Aska
8603839711
- orchestrator: Tracer interface with spans for intent analysis, persona, subsession dispatch, context build, LLM streaming and each tool call; traces finish when the stream completes — the admin Trace waterfall finally receives data - background: Thinker generates real LLM thoughts (post-chat + silence triggers, recent history as context, empty/overlong outputs discarded); thoughts broadcast to web clients as 'thinking' WS messages; pipeline postChatHook triggers post-chat thinking after each round - memory: ApplyDecayAll + ListUserIDs with a 6h decay/purge cron in main (was implemented but never scheduled)
304 lines
10 KiB
Go
304 lines
10 KiB
Go
package memory
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/lib/pq"
|
|
)
|
|
|
|
// Repository handles database operations for memories.
|
|
type Repository struct {
|
|
db *sql.DB
|
|
vectorEnabled bool // true when the pgvector embedding column exists
|
|
}
|
|
|
|
// NewRepository creates a new memory repository.
|
|
func NewRepository(db *sql.DB) *Repository {
|
|
return &Repository{db: db}
|
|
}
|
|
|
|
// EnableVector marks the embedding column as available (after migration).
|
|
func (r *Repository) EnableVector() { r.vectorEnabled = true }
|
|
|
|
// VectorEnabled reports whether vector search is usable.
|
|
func (r *Repository) VectorEnabled() bool { return r.vectorEnabled }
|
|
|
|
// VectorMigrationSQL returns the statements to enable pgvector and add the
|
|
// embedding column with the given dimension.
|
|
func VectorMigrationSQL(dims int) string {
|
|
return fmt.Sprintf(`
|
|
CREATE EXTENSION IF NOT EXISTS vector;
|
|
ALTER TABLE memories ADD COLUMN IF NOT EXISTS embedding vector(%d);
|
|
CREATE INDEX IF NOT EXISTS idx_memories_embedding ON memories USING hnsw (embedding vector_cosine_ops);
|
|
`, dims)
|
|
}
|
|
|
|
// vectorString formats an embedding for pgvector input.
|
|
func vectorString(vec []float64) string {
|
|
var sb strings.Builder
|
|
sb.WriteByte('[')
|
|
for i, v := range vec {
|
|
if i > 0 {
|
|
sb.WriteByte(',')
|
|
}
|
|
fmt.Fprintf(&sb, "%g", v)
|
|
}
|
|
sb.WriteByte(']')
|
|
return sb.String()
|
|
}
|
|
|
|
// Create inserts a new memory and returns its ID.
|
|
func (r *Repository) Create(ctx context.Context, m *Memory) (string, error) {
|
|
query := `
|
|
INSERT INTO memories (user_id, source_id, content, summary, keywords, category, priority, decay, access_count, last_accessed_at, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
|
|
RETURNING id
|
|
`
|
|
now := time.Now()
|
|
m.CreatedAt = now
|
|
m.UpdatedAt = now
|
|
if m.LastAcc.IsZero() {
|
|
m.LastAcc = now
|
|
}
|
|
if m.Decay == 0 {
|
|
m.Decay = 1.0
|
|
}
|
|
if m.Keywords == nil {
|
|
m.Keywords = []string{}
|
|
}
|
|
|
|
var id string
|
|
if r.vectorEnabled && len(m.Embedding) > 0 {
|
|
query = `
|
|
INSERT INTO memories (user_id, source_id, content, summary, keywords, category, priority, decay, access_count, last_accessed_at, created_at, updated_at, embedding)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13::vector)
|
|
RETURNING id
|
|
`
|
|
err := r.db.QueryRowContext(ctx, query,
|
|
m.UserID, m.SourceID, m.Content, m.Summary,
|
|
pq.Array(m.Keywords), m.Category, m.Priority, m.Decay,
|
|
m.AccessCnt, m.LastAcc, m.CreatedAt, m.UpdatedAt, vectorString(m.Embedding),
|
|
).Scan(&id)
|
|
if err != nil {
|
|
return "", fmt.Errorf("insert memory: %w", err)
|
|
}
|
|
m.ID = id
|
|
return id, nil
|
|
}
|
|
|
|
err := r.db.QueryRowContext(ctx, query,
|
|
m.UserID, m.SourceID, m.Content, m.Summary,
|
|
pq.Array(m.Keywords), m.Category, m.Priority, m.Decay,
|
|
m.AccessCnt, m.LastAcc, m.CreatedAt, m.UpdatedAt,
|
|
).Scan(&id)
|
|
if err != nil {
|
|
return "", fmt.Errorf("insert memory: %w", err)
|
|
}
|
|
m.ID = id
|
|
return id, nil
|
|
}
|
|
|
|
// GetByID retrieves a memory by ID.
|
|
func (r *Repository) GetByID(ctx context.Context, id string) (*Memory, error) {
|
|
query := `
|
|
SELECT id, user_id, source_id, content, summary, keywords, category, priority, decay, access_count, last_accessed_at, created_at, updated_at
|
|
FROM memories WHERE id = $1
|
|
`
|
|
m := &Memory{}
|
|
err := r.db.QueryRowContext(ctx, query, id).Scan(
|
|
&m.ID, &m.UserID, &m.SourceID, &m.Content, &m.Summary,
|
|
pq.Array(&m.Keywords), &m.Category, &m.Priority, &m.Decay,
|
|
&m.AccessCnt, &m.LastAcc, &m.CreatedAt, &m.UpdatedAt,
|
|
)
|
|
if err == sql.ErrNoRows {
|
|
return nil, fmt.Errorf("memory %s not found", id)
|
|
}
|
|
if err != nil {
|
|
return nil, fmt.Errorf("get memory: %w", err)
|
|
}
|
|
return m, nil
|
|
}
|
|
|
|
// ListByUser retrieves all memories for a user, ordered by priority descending.
|
|
func (r *Repository) ListByUser(ctx context.Context, userID string, limit int) ([]Memory, error) {
|
|
if limit <= 0 {
|
|
limit = 50
|
|
}
|
|
query := `
|
|
SELECT id, user_id, source_id, content, summary, keywords, category, priority, decay, access_count, last_accessed_at, created_at, updated_at
|
|
FROM memories WHERE user_id = $1 AND decay > $2
|
|
ORDER BY priority DESC, created_at DESC
|
|
LIMIT $3
|
|
`
|
|
return r.queryMemories(ctx, query, userID, DecayThreshold, limit)
|
|
}
|
|
|
|
// SearchByKeywords performs full-text search on content and keywords.
|
|
func (r *Repository) SearchByKeywords(ctx context.Context, userID, q string, limit int) ([]Memory, error) {
|
|
if limit <= 0 { limit = 20 }
|
|
tokens := strings.Fields(q)
|
|
var conditions []string
|
|
var args []any
|
|
args = append(args, userID, DecayThreshold)
|
|
for _, token := range tokens {
|
|
paramIdx := len(args) + 1
|
|
conditions = append(conditions, fmt.Sprintf("(content ILIKE $%d OR $%d = ANY(keywords))", paramIdx, paramIdx))
|
|
args = append(args, "%"+token+"%")
|
|
}
|
|
query := fmt.Sprintf(`
|
|
SELECT id, user_id, source_id, content, summary, keywords, category, priority, decay, access_count, last_accessed_at, created_at, updated_at
|
|
FROM memories WHERE user_id = $1 AND decay > $2 AND (%s) ORDER BY priority DESC LIMIT $%d
|
|
`, strings.Join(conditions, " OR "), len(args)+1)
|
|
args = append(args, limit)
|
|
return r.queryMemories(ctx, query, args...)
|
|
}
|
|
|
|
// SearchByNGrams searches memories by n-gram tokens matching on keywords array.
|
|
// More effective for Chinese text than ILIKE substring matching.
|
|
func (r *Repository) SearchByNGrams(ctx context.Context, userID string, nGrams []string, limit int) ([]Memory, error) {
|
|
if limit <= 0 { limit = 10 }
|
|
if len(nGrams) == 0 {
|
|
return r.ListByUser(ctx, userID, limit)
|
|
}
|
|
|
|
// Search where any n-gram matches any keyword in the keywords array
|
|
// Use postgres array overlap operator: keywords && ARRAY[$1, $2, ...]
|
|
query := fmt.Sprintf(`
|
|
SELECT id, user_id, source_id, content, summary, keywords, category, priority, decay, access_count, last_accessed_at, created_at, updated_at
|
|
FROM memories WHERE user_id = $1 AND decay > $2 AND keywords && $3
|
|
ORDER BY priority DESC LIMIT $4
|
|
`)
|
|
|
|
return r.queryMemories(ctx, query, userID, DecayThreshold, pq.Array(nGrams), limit)
|
|
}
|
|
|
|
// SearchByVector returns the nearest memories by cosine distance.
|
|
func (r *Repository) SearchByVector(ctx context.Context, userID string, vec []float64, limit int) ([]Memory, error) {
|
|
if !r.vectorEnabled || len(vec) == 0 {
|
|
return nil, nil
|
|
}
|
|
if limit <= 0 {
|
|
limit = 10
|
|
}
|
|
query := `
|
|
SELECT id, user_id, source_id, content, summary, keywords, category, priority, decay, access_count, last_accessed_at, created_at, updated_at
|
|
FROM memories WHERE user_id = $1 AND decay > $2 AND embedding IS NOT NULL
|
|
ORDER BY embedding <=> $3::vector LIMIT $4
|
|
`
|
|
return r.queryMemories(ctx, query, userID, DecayThreshold, vectorString(vec), limit)
|
|
}
|
|
|
|
// Update updates an existing memory.
|
|
func (r *Repository) Update(ctx context.Context, m *Memory) error {
|
|
m.UpdatedAt = time.Now()
|
|
query := `
|
|
UPDATE memories SET content=$1, summary=$2, keywords=$3, category=$4, priority=$5, decay=$6, access_count=$7, last_accessed_at=$8, updated_at=$9
|
|
WHERE id=$10
|
|
`
|
|
_, err := r.db.ExecContext(ctx, query,
|
|
m.Content, m.Summary, pq.Array(m.Keywords), m.Category,
|
|
m.Priority, m.Decay, m.AccessCnt, m.LastAcc, m.UpdatedAt, m.ID,
|
|
)
|
|
return err
|
|
}
|
|
|
|
// Delete removes a memory by ID.
|
|
func (r *Repository) Delete(ctx context.Context, id string) error {
|
|
_, err := r.db.ExecContext(ctx, "DELETE FROM memories WHERE id=$1", id)
|
|
return err
|
|
}
|
|
|
|
// RecordAccess increments the access count and updates last_accessed_at.
|
|
func (r *Repository) RecordAccess(ctx context.Context, id string) error {
|
|
_, err := r.db.ExecContext(ctx,
|
|
"UPDATE memories SET access_count = access_count + 1, last_accessed_at = $1 WHERE id = $2",
|
|
time.Now(), id,
|
|
)
|
|
return err
|
|
}
|
|
|
|
// ListUserIDs returns distinct user IDs that have memories.
|
|
func (r *Repository) ListUserIDs(ctx context.Context) ([]string, error) {
|
|
rows, err := r.db.QueryContext(ctx, `SELECT DISTINCT user_id FROM memories`)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list user ids: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
var ids []string
|
|
for rows.Next() {
|
|
var id string
|
|
if err := rows.Scan(&id); err != nil {
|
|
return nil, err
|
|
}
|
|
ids = append(ids, id)
|
|
}
|
|
return ids, rows.Err()
|
|
}
|
|
|
|
// PurgeDecayed removes memories below the decay threshold.
|
|
func (r *Repository) PurgeDecayed(ctx context.Context, userID string) (int64, error) {
|
|
result, err := r.db.ExecContext(ctx,
|
|
"DELETE FROM memories WHERE user_id = $1 AND decay < $2",
|
|
userID, DecayThreshold,
|
|
)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
n, _ := result.RowsAffected()
|
|
return n, nil
|
|
}
|
|
|
|
// CreateTable returns the SQL to create the memories table.
|
|
// This is used for migrations, not called at runtime.
|
|
func CreateTableSQL() string {
|
|
return `
|
|
CREATE TABLE IF NOT EXISTS memories (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
user_id TEXT NOT NULL,
|
|
source_id TEXT NOT NULL DEFAULT '',
|
|
content TEXT NOT NULL,
|
|
summary TEXT NOT NULL DEFAULT '',
|
|
keywords TEXT[] NOT NULL DEFAULT '{}',
|
|
category TEXT NOT NULL DEFAULT 'fact',
|
|
priority DOUBLE PRECISION NOT NULL DEFAULT 0.5,
|
|
decay DOUBLE PRECISION NOT NULL DEFAULT 1.0,
|
|
access_count INTEGER NOT NULL DEFAULT 0,
|
|
last_accessed_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
|
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_memories_user_id ON memories(user_id);
|
|
CREATE INDEX IF NOT EXISTS idx_memories_source_id ON memories(source_id);
|
|
CREATE INDEX IF NOT EXISTS idx_memories_category ON memories(category);
|
|
CREATE INDEX IF NOT EXISTS idx_memories_decay ON memories(decay) WHERE decay < 0.1;
|
|
CREATE INDEX IF NOT EXISTS idx_memories_keywords ON memories USING gin(keywords);
|
|
`
|
|
}
|
|
|
|
// queryMemories is a helper that scans rows into Memory structs.
|
|
func (r *Repository) queryMemories(ctx context.Context, query string, args ...any) ([]Memory, error) {
|
|
rows, err := r.db.QueryContext(ctx, query, args...)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("query memories: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
|
|
var memories []Memory
|
|
for rows.Next() {
|
|
var m Memory
|
|
if err := rows.Scan(
|
|
&m.ID, &m.UserID, &m.SourceID, &m.Content, &m.Summary,
|
|
pq.Array(&m.Keywords), &m.Category, &m.Priority, &m.Decay,
|
|
&m.AccessCnt, &m.LastAcc, &m.CreatedAt, &m.UpdatedAt,
|
|
); err != nil {
|
|
return nil, fmt.Errorf("scan memory: %w", err)
|
|
}
|
|
memories = append(memories, m)
|
|
}
|
|
return memories, rows.Err()
|
|
}
|