rony-llm-agent/pkg/rag/memory.go

package rag

import (
	"context"
	"fmt"
	"time"

	"github.com/google/uuid"
)

// MemoryType classifies a fragment within the three-tier taxonomy from the
// Phase 2 spec (docs/phase2.md §3.1). Working memory (the current session's
// messages) lives in the consuming product's own state, not here.
type MemoryType string

const (
	// MemoryEpisodic records past events: "what happened / what I did on
	// <date>". Typically auto-captured at the end of successful turns (see
	// EpisodeCapture) rather than saved deliberately.
	MemoryEpisodic MemoryType = "episodic"
	// MemorySemantic records consolidated knowledge and facts: "how the
	// architecture works", "the API returns X". Curated — saved when
	// something is worth knowing independent of when it was learned.
	MemorySemantic MemoryType = "semantic"
	// MemoryProcedural records how to do things: workflows, procedures,
	// user preferences about process. This is also what every fragment
	// saved before the taxonomy existed is treated as — the pre-taxonomy
	// tools (save_process et al.) only ever stored procedures.
	MemoryProcedural MemoryType = "procedural"
)

// metaTypeKey is the metadata key the fragment's MemoryType round-trips
// through, so backends need no schema change to support the taxonomy.
const metaTypeKey = "memory_type"

// Fragment represents a piece of content stored in the RAG system.
type Fragment struct {
	ID        string
	Content   string
	Type      MemoryType // defaults to MemoryProcedural when empty (pre-taxonomy compatibility)
	Vector    []float32
	Metadata  map[string]string
	Timestamp time.Time
	ProjectID string
}

// Memory provides persistent memory and semantic search over agent content.
type Memory interface {
	Add(ctx context.Context, fragment Fragment) error
	Search(ctx context.Context, query string, topK int) ([]Fragment, error)
	// SearchByType is Search restricted to the given memory types. No types
	// means no restriction (same as Search). Fragments stored before the
	// taxonomy existed match MemoryProcedural.
	SearchByType(ctx context.Context, query string, topK int, types ...MemoryType) ([]Fragment, error)
	Forget(ctx context.Context, id string) error
	ForgetAll(ctx context.Context) error
}

// Config holds the settings for creating a Memory.
type Config struct {
	Backend  Backend
	Embedder Embedder
}

// Backend is the interface for storage backends. queryVector is nil when no
// embedder produced one (e.g. it's unavailable or failed); backends that
// can't search without a vector (e.g. a pure vector database like Chroma)
// should return an error in that case, while backends capable of lexical
// search (e.g. SQLite FTS5) can fall back to matching on query instead.
type Backend interface {
	Upsert(ctx context.Context, id string, vector []float32, content string, metadata map[string]string) error
	Search(ctx context.Context, query string, queryVector []float32, topK int) ([]SearchResult, error)
	Forget(ctx context.Context, id string) error
	ForgetAll(ctx context.Context) error
}

// SearchResult represents a matched fragment from a search.
type SearchResult struct {
	ID       string
	Content  string
	Score    float32
	Metadata map[string]string
}

// Embedder generates embeddings for text.
type Embedder interface {
	Embed(ctx context.Context, text string) ([]float32, error)
	Dimensions() int
}

// memory implements Memory using a Backend and Embedder.
type memory struct {
	backend  Backend
	embedder Embedder
}

// New creates a new Memory with the given config.
func New(cfg Config) (Memory, error) {
	if cfg.Backend == nil {
		return nil, fmt.Errorf("backend is required")
	}
	if cfg.Embedder == nil {
		return nil, fmt.Errorf("embedder is required")
	}
	return &memory{
		backend:  cfg.Backend,
		embedder: cfg.Embedder,
	}, nil
}

func (m *memory) Add(ctx context.Context, fragment Fragment) error {
	if fragment.ID == "" {
		fragment.ID = uuid.New().String()
	}
	if fragment.Metadata == nil {
		fragment.Metadata = make(map[string]string)
	}
	if fragment.Type == "" {
		fragment.Type = MemoryProcedural
	}
	fragment.Metadata[metaTypeKey] = string(fragment.Type)
	fragment.Metadata["project_id"] = fragment.ProjectID
	fragment.Timestamp = time.Now()

	// A failed embedding doesn't block saving: the backend still has the
	// raw content and can index it for lexical search (e.g. FTS5), so the
	// fragment just won't be reachable by vector similarity later.
	vector, err := m.embedder.Embed(ctx, fragment.Content)
	if err != nil {
		vector = nil
	}
	fragment.Vector = vector

	return m.backend.Upsert(ctx, fragment.ID, fragment.Vector, fragment.Content, fragment.Metadata)
}

func (m *memory) Search(ctx context.Context, query string, topK int) ([]Fragment, error) {
	return m.SearchByType(ctx, query, topK)
}

// typeOf resolves a stored result's memory type; fragments saved before the
// taxonomy existed carry no memory_type metadata and were all procedures.
func typeOf(metadata map[string]string) MemoryType {
	if t := MemoryType(metadata[metaTypeKey]); t != "" {
		return t
	}
	return MemoryProcedural
}

// typeFilterOverfetch is how many times topK gets requested from the backend
// when SearchByType has to post-filter by memory type: the Backend interface
// has no type predicate (deliberately — backends stay schema-agnostic), so
// filtering happens here and the extra headroom keeps a type-restricted
// search from coming back near-empty just because the top raw matches
// happened to be of other types.
const typeFilterOverfetch = 4

func (m *memory) SearchByType(ctx context.Context, query string, topK int, types ...MemoryType) ([]Fragment, error) {
	if topK <= 0 {
		topK = 5
	}
	fetchK := topK
	if len(types) > 0 {
		fetchK = topK * typeFilterOverfetch
	}

	// Same fallback as Add: if embedding the query fails, search proceeds
	// with no vector so the backend can fall back to lexical matching.
	queryVector, err := m.embedder.Embed(ctx, query)
	if err != nil {
		queryVector = nil
	}

	results, err := m.backend.Search(ctx, query, queryVector, fetchK)
	if err != nil {
		return nil, fmt.Errorf("search: %w", err)
	}

	wanted := make(map[MemoryType]bool, len(types))
	for _, t := range types {
		wanted[t] = true
	}

	fragments := make([]Fragment, 0, topK)
	for _, r := range results {
		fragType := typeOf(r.Metadata)
		if len(wanted) > 0 && !wanted[fragType] {
			continue
		}
		fragments = append(fragments, Fragment{
			ID:        r.ID,
			Content:   r.Content,
			Type:      fragType,
			Metadata:  r.Metadata,
			ProjectID: r.Metadata["project_id"],
		})
		if len(fragments) == topK {
			break
		}
	}
	return fragments, nil
}

func (m *memory) Forget(ctx context.Context, id string) error {
	return m.backend.Forget(ctx, id)
}

func (m *memory) ForgetAll(ctx context.Context) error {
	return m.backend.ForgetAll(ctx)
}