179 lines
5.3 KiB
Bash
Executable file
179 lines
5.3 KiB
Bash
Executable file
#!/usr/bin/env bash
|
|
# memory-pg.sh — PostgreSQL-backed memory operations
|
|
#
|
|
# Usage:
|
|
# ./memory-pg.sh store "summary text" [topics] [importance]
|
|
# ./memory-pg.sh search "query text" [limit]
|
|
# ./memory-pg.sh recent [limit]
|
|
# ./memory-pg.sh important [limit]
|
|
# ./memory-pg.sh count
|
|
|
|
. "$(dirname "$0")/common.sh"
|
|
|
|
DISPLAY_DATETIME_SQL=$(pg_display_datetime_expr "created_at")
|
|
|
|
# Sanitize an integer argument (for LIMIT clauses which don't support psql variables)
|
|
safe_int() {
|
|
local val="$1"
|
|
if ! [[ "$val" =~ ^[0-9]+$ ]]; then
|
|
echo "Error: expected integer, got '$val'" >&2
|
|
exit 1
|
|
fi
|
|
echo "$val"
|
|
}
|
|
|
|
# ── Store a new memory with automatic chunking + embedding ──
|
|
cmd_store() {
|
|
local summary="$1"
|
|
local topics="${2:-}"
|
|
local importance="${3:-2}"
|
|
|
|
# Parse topics into PostgreSQL array
|
|
local pg_topics="{}"
|
|
if [ -n "$topics" ]; then
|
|
# Convert comma-separated to PostgreSQL array: "a,b,c" → "{a,b,c}"
|
|
pg_topics="{$topics}"
|
|
fi
|
|
|
|
# Escape summary for SQL
|
|
local escaped_summary
|
|
escaped_summary=$(python3 -c "import sys; print(sys.argv[1].replace(\"'\", \"''\"))" "$summary")
|
|
local safe_importance
|
|
safe_importance=$(safe_int "$importance")
|
|
|
|
# Insert memory
|
|
local memory_id
|
|
memory_id=$(pg -c "
|
|
INSERT INTO memories (summary, topics, importance)
|
|
VALUES ('$escaped_summary', '$pg_topics'::text[], $safe_importance)
|
|
RETURNING id;
|
|
")
|
|
|
|
if [ -z "$memory_id" ]; then
|
|
echo "Error: failed to insert memory" >&2
|
|
exit 1
|
|
fi
|
|
|
|
echo "Stored memory: $memory_id"
|
|
|
|
# Chunk the summary
|
|
local chunk_order=0
|
|
while IFS= read -r chunk; do
|
|
[ -z "$chunk" ] && continue
|
|
|
|
local content_hash
|
|
content_hash=$(echo -n "$chunk" | md5)
|
|
|
|
# Escape chunk text for SQL
|
|
local escaped_chunk
|
|
escaped_chunk=$(python3 -c "import sys; print(sys.argv[1].replace(\"'\", \"''\"))" "$chunk")
|
|
|
|
# Insert chunk
|
|
local chunk_id
|
|
chunk_id=$(pg -c "
|
|
INSERT INTO memory_chunks (memory_id, chunk_order, chunk_text, content_hash)
|
|
VALUES ('$memory_id', $chunk_order, '$escaped_chunk', '$content_hash')
|
|
RETURNING id;
|
|
")
|
|
|
|
# Generate and store embedding
|
|
local embedding
|
|
embedding=$("$SCRIPT_DIR/embed.sh" "$chunk")
|
|
|
|
pg -c "
|
|
INSERT INTO memory_embeddings (chunk_id, embedding, embedding_provider, embedding_model)
|
|
VALUES ('$chunk_id', '$embedding'::vector, '$EMBED_PROVIDER', '$EMBED_MODEL');
|
|
" >/dev/null
|
|
|
|
echo " chunk $chunk_order embedded (${#chunk} chars)"
|
|
chunk_order=$((chunk_order + 1))
|
|
done < <("$SCRIPT_DIR/chunk.sh" "$summary")
|
|
|
|
echo "Done: $chunk_order chunks embedded"
|
|
}
|
|
|
|
# ── Hybrid search ──
|
|
cmd_search() {
|
|
local query="$1"
|
|
local safe_limit
|
|
safe_limit=$(safe_int "${2:-5}")
|
|
|
|
# Generate query embedding
|
|
local query_embedding
|
|
query_embedding=$("$SCRIPT_DIR/embed.sh" "$query")
|
|
|
|
# Escape query for SQL (single quotes → '')
|
|
local escaped_query
|
|
escaped_query=$(python3 -c "import sys; print(sys.argv[1].replace(\"'\", \"''\"))" "$query")
|
|
|
|
pg_pretty -c "
|
|
SELECT
|
|
left(chunk_text, 80) AS match,
|
|
importance,
|
|
topics,
|
|
round(combined_score::numeric, 6) AS score,
|
|
${DISPLAY_DATETIME_SQL} AS created
|
|
FROM search_memories('$escaped_query', '$query_embedding'::vector, $safe_limit);
|
|
"
|
|
}
|
|
|
|
# ── Recent memories ──
|
|
cmd_recent() {
|
|
local safe_limit
|
|
safe_limit=$(safe_int "${1:-5}")
|
|
|
|
pg_pretty -c "
|
|
SELECT
|
|
${DISPLAY_DATETIME_SQL} AS created,
|
|
importance AS imp,
|
|
left(summary, 80) AS summary,
|
|
topics
|
|
FROM memories
|
|
ORDER BY created_at DESC
|
|
LIMIT $safe_limit;
|
|
"
|
|
}
|
|
|
|
# ── High-importance memories ──
|
|
cmd_important() {
|
|
local safe_limit
|
|
safe_limit=$(safe_int "${1:-10}")
|
|
|
|
pg_pretty -c "
|
|
SELECT
|
|
${DISPLAY_DATETIME_SQL} AS created,
|
|
importance AS imp,
|
|
left(summary, 80) AS summary,
|
|
topics
|
|
FROM memories
|
|
WHERE importance >= 4
|
|
ORDER BY importance DESC, created_at DESC
|
|
LIMIT $safe_limit;
|
|
"
|
|
}
|
|
|
|
# ── Count ──
|
|
cmd_count() {
|
|
echo "memories: $(pg -c "SELECT count(*) FROM memories;")"
|
|
echo "chunks: $(pg -c "SELECT count(*) FROM memory_chunks;")"
|
|
echo "embeddings:$(pg -c "SELECT count(*) FROM memory_embeddings;")"
|
|
}
|
|
|
|
# ── Main dispatcher ──
|
|
case "${1:-help}" in
|
|
store) cmd_store "${2:?Summary required}" "${3:-}" "${4:-2}" ;;
|
|
search) cmd_search "${2:?Query required}" "${3:-5}" ;;
|
|
recent) cmd_recent "${2:-5}" ;;
|
|
important) cmd_important "${2:-10}" ;;
|
|
count) cmd_count ;;
|
|
*)
|
|
echo "Usage: memory-pg.sh <command> [args]"
|
|
echo ""
|
|
echo "Commands:"
|
|
echo " store <summary> [topics] [importance] Store a memory with auto-chunking + embedding"
|
|
echo " search <query> [limit] Hybrid search (full-text + vector)"
|
|
echo " recent [limit] Show recent memories"
|
|
echo " important [limit] Show high-importance memories"
|
|
echo " count Show table counts"
|
|
;;
|
|
esac
|