gooseek/backend/internal/db/memory_repo.go

package db

import (
	"context"
	"encoding/json"
	"time"
)

type UserMemory struct {
	ID          string                 `json:"id"`
	UserID      string                 `json:"userId"`
	MemoryType  string                 `json:"memoryType"`
	Key         string                 `json:"key"`
	Value       string                 `json:"value"`
	Metadata    map[string]interface{} `json:"metadata"`
	Importance  int                    `json:"importance"`
	LastUsed    time.Time              `json:"lastUsed"`
	UseCount    int                    `json:"useCount"`
	CreatedAt   time.Time              `json:"createdAt"`
	UpdatedAt   time.Time              `json:"updatedAt"`
}

type MemoryRepository struct {
	db *PostgresDB
}

func NewMemoryRepository(db *PostgresDB) *MemoryRepository {
	return &MemoryRepository{db: db}
}

func (r *MemoryRepository) RunMigrations(ctx context.Context) error {
	migrations := []string{
		`CREATE TABLE IF NOT EXISTS user_memories (
			id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
			user_id UUID NOT NULL,
			memory_type VARCHAR(50) NOT NULL,
			key VARCHAR(255) NOT NULL,
			value TEXT NOT NULL,
			metadata JSONB DEFAULT '{}',
			importance INT DEFAULT 5,
			last_used TIMESTAMPTZ DEFAULT NOW(),
			use_count INT DEFAULT 0,
			created_at TIMESTAMPTZ DEFAULT NOW(),
			updated_at TIMESTAMPTZ DEFAULT NOW(),
			UNIQUE(user_id, memory_type, key)
		)`,
		`CREATE INDEX IF NOT EXISTS idx_user_memories_user ON user_memories(user_id)`,
		`CREATE INDEX IF NOT EXISTS idx_user_memories_type ON user_memories(user_id, memory_type)`,
		`CREATE INDEX IF NOT EXISTS idx_user_memories_importance ON user_memories(user_id, importance DESC)`,
	}

	for _, m := range migrations {
		if _, err := r.db.db.ExecContext(ctx, m); err != nil {
			return err
		}
	}
	return nil
}

func (r *MemoryRepository) Save(ctx context.Context, mem *UserMemory) error {
	metadataJSON, _ := json.Marshal(mem.Metadata)

	query := `
		INSERT INTO user_memories (user_id, memory_type, key, value, metadata, importance)
		VALUES ($1, $2, $3, $4, $5, $6)
		ON CONFLICT (user_id, memory_type, key) 
		DO UPDATE SET 
			value = EXCLUDED.value,
			metadata = EXCLUDED.metadata,
			importance = EXCLUDED.importance,
			updated_at = NOW()
		RETURNING id, created_at, updated_at
	`
	return r.db.db.QueryRowContext(ctx, query,
		mem.UserID, mem.MemoryType, mem.Key, mem.Value, metadataJSON, mem.Importance,
	).Scan(&mem.ID, &mem.CreatedAt, &mem.UpdatedAt)
}

func (r *MemoryRepository) GetByUserID(ctx context.Context, userID string, memoryType string, limit int) ([]*UserMemory, error) {
	query := `
		SELECT id, user_id, memory_type, key, value, metadata, importance, last_used, use_count, created_at, updated_at
		FROM user_memories 
		WHERE user_id = $1
	`
	args := []interface{}{userID}

	if memoryType != "" {
		query += " AND memory_type = $2"
		args = append(args, memoryType)
	}

	query += " ORDER BY importance DESC, last_used DESC"

	if limit > 0 {
		query += " LIMIT $" + string(rune('0'+len(args)+1))
		args = append(args, limit)
	}

	rows, err := r.db.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var memories []*UserMemory
	for rows.Next() {
		var mem UserMemory
		var metadataJSON []byte

		if err := rows.Scan(
			&mem.ID, &mem.UserID, &mem.MemoryType, &mem.Key, &mem.Value,
			&metadataJSON, &mem.Importance, &mem.LastUsed, &mem.UseCount,
			&mem.CreatedAt, &mem.UpdatedAt,
		); err != nil {
			return nil, err
		}

		json.Unmarshal(metadataJSON, &mem.Metadata)
		memories = append(memories, &mem)
	}

	return memories, nil
}

func (r *MemoryRepository) GetContextForUser(ctx context.Context, userID string) (string, error) {
	memories, err := r.GetByUserID(ctx, userID, "", 20)
	if err != nil {
		return "", err
	}

	var context string
	for _, mem := range memories {
		switch mem.MemoryType {
		case "preference":
			context += "User preference: " + mem.Key + " = " + mem.Value + "\n"
		case "fact":
			context += "Known fact about user: " + mem.Value + "\n"
		case "instruction":
			context += "User instruction: " + mem.Value + "\n"
		case "interest":
			context += "User interest: " + mem.Value + "\n"
		default:
			context += mem.Key + ": " + mem.Value + "\n"
		}
	}

	return context, nil
}

func (r *MemoryRepository) IncrementUseCount(ctx context.Context, id, userID string) error {
	result, err := r.db.db.ExecContext(ctx,
		"UPDATE user_memories SET use_count = use_count + 1, last_used = NOW() WHERE id = $1 AND user_id = $2",
		id, userID,
	)
	if err != nil {
		return err
	}
	rows, _ := result.RowsAffected()
	if rows == 0 {
		return ErrNotFound
	}
	return nil
}

func (r *MemoryRepository) Delete(ctx context.Context, id, userID string) error {
	result, err := r.db.db.ExecContext(ctx, "DELETE FROM user_memories WHERE id = $1 AND user_id = $2", id, userID)
	if err != nil {
		return err
	}
	rows, _ := result.RowsAffected()
	if rows == 0 {
		return ErrNotFound
	}
	return nil
}

func (r *MemoryRepository) DeleteByUserID(ctx context.Context, userID string) error {
	_, err := r.db.db.ExecContext(ctx, "DELETE FROM user_memories WHERE user_id = $1", userID)
	return err
}

func ExtractMemoriesFromConversation(ctx context.Context, llmClient interface{}, conversation, answer string) ([]UserMemory, error) {
	return nil, nil
}