selfheal-lab-apps/ai/internal/rag/store.go

// Package rag is the AI VM state store (experiment design 6.4): runbook chunks with pgvector embeddings, prompt
// versions, diagnosis sessions and the Discord HITL message ledger. Schema is created at start-up; idempotent.
package rag

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
	pgvector "github.com/pgvector/pgvector-go"
)

type Store struct{ pool *pgxpool.Pool }

const schema = `
CREATE EXTENSION IF NOT EXISTS vector;
CREATE SCHEMA IF NOT EXISTS ai;
CREATE TABLE IF NOT EXISTS ai.rag_docs (
  id         bigserial PRIMARY KEY,
  source     text NOT NULL,
  chunk      text NOT NULL,
  chunk_hash text NOT NULL UNIQUE,
  embedding  vector(1536) NOT NULL,
  created_at timestamptz NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS rag_docs_embedding_idx ON ai.rag_docs USING hnsw (embedding vector_cosine_ops);
CREATE TABLE IF NOT EXISTS ai.prompt_versions (
  version    text PRIMARY KEY,
  body       text NOT NULL,
  created_at timestamptz NOT NULL DEFAULT now()
);
CREATE TABLE IF NOT EXISTS ai.sessions (
  session_id     text PRIMARY KEY,
  run_id         text,
  scenario_id    text,
  model          text NOT NULL,
  prompt_version text NOT NULL,
  input          jsonb NOT NULL,
  output         jsonb,
  tool_calls     jsonb NOT NULL DEFAULT '[]'::jsonb,
  usage          jsonb,
  error          text,
  duration_ms    bigint,
  created_at     timestamptz NOT NULL DEFAULT now()
);
CREATE TABLE IF NOT EXISTS ai.hitl_messages (
  token       text PRIMARY KEY,
  channel_id  text NOT NULL,
  message_id  text NOT NULL,
  action_id   text NOT NULL,
  target      jsonb NOT NULL DEFAULT '{}'::jsonb,
  run_id      text,
  scenario_id text,
  session_id  text,
  expires_at  timestamptz NOT NULL,
  decided_at  timestamptz,
  decision    text,
  decided_by  text,
  result      jsonb,
  created_at  timestamptz NOT NULL DEFAULT now()
);`

func New(ctx context.Context, dsn string) (*Store, error) {
	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		return nil, err
	}
	var lastErr error
	for i := 0; i < 30; i++ { // rag-db container may still be starting
		if lastErr = pool.Ping(ctx); lastErr == nil {
			break
		}
		time.Sleep(2 * time.Second)
	}
	if lastErr != nil {
		return nil, fmt.Errorf("rag db: %w", lastErr)
	}
	if _, err := pool.Exec(ctx, schema); err != nil {
		return nil, fmt.Errorf("rag schema: %w", err)
	}
	return &Store{pool: pool}, nil
}

func (s *Store) Ping(ctx context.Context) error { return s.pool.Ping(ctx) }

// ---- prompt versions ----

// PromptVersion is the first 12 hex characters of sha256(body); the same body always maps to the same version.
func PromptVersion(body string) string {
	h := sha256.Sum256([]byte(body))
	return hex.EncodeToString(h[:])[:12]
}

func (s *Store) UpsertPrompt(ctx context.Context, body string) (string, error) {
	v := PromptVersion(body)
	_, err := s.pool.Exec(ctx, `INSERT INTO ai.prompt_versions (version, body) VALUES ($1,$2) ON CONFLICT (version) DO NOTHING`, v, body)
	return v, err
}

// ---- runbook chunks ----

// Chunk splits a document into pieces of about maxChars characters on paragraph boundaries (design 6.4: ~500 tokens).
func Chunk(text string, maxChars int) []string {
	paras := strings.Split(strings.ReplaceAll(text, "\r\n", "\n"), "\n\n")
	var out []string
	var cur strings.Builder
	flush := func() {
		if t := strings.TrimSpace(cur.String()); t != "" {
			out = append(out, t)
		}
		cur.Reset()
	}
	for _, p := range paras {
		p = strings.TrimSpace(p)
		if p == "" {
			continue
		}
		if cur.Len() > 0 && cur.Len()+len(p)+2 > maxChars {
			flush()
		}
		for len(p) > maxChars { // a single oversized paragraph is cut hard
			cut := strings.LastIndex(p[:maxChars], "\n")
			if cut < maxChars/2 {
				cut = maxChars
			}
			if cur.Len() > 0 {
				flush()
			}
			out = append(out, strings.TrimSpace(p[:cut]))
			p = strings.TrimSpace(p[cut:])
		}
		if cur.Len() > 0 {
			cur.WriteString("\n\n")
		}
		cur.WriteString(p)
	}
	flush()
	return out
}

type Embedder interface {
	Embed(ctx context.Context, model string, inputs []string) ([][]float32, error)
}

// IngestDir embeds every *.md under dir whose chunk is not stored yet and removes chunks of files that disappeared.
// Returns (new chunks, removed chunks). Unchanged files cost nothing (hash lookup only).
func (s *Store) IngestDir(ctx context.Context, dir string, emb Embedder, model string) (int, int, error) {
	files, err := filepath.Glob(filepath.Join(dir, "*.md"))
	if err != nil {
		return 0, 0, err
	}
	sort.Strings(files)
	type pending struct{ source, chunk, hash string }
	var todo []pending
	keep := map[string]bool{}
	sources := map[string]bool{}
	for _, f := range files {
		b, err := os.ReadFile(f)
		if err != nil {
			return 0, 0, err
		}
		src := filepath.Base(f)
		sources[src] = true
		for _, c := range Chunk(string(b), 2000) {
			h := sha256.Sum256([]byte(src + "\x00" + c))
			hash := hex.EncodeToString(h[:])
			keep[hash] = true
			var exists bool
			if err := s.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM ai.rag_docs WHERE chunk_hash=$1)`, hash).Scan(&exists); err != nil {
				return 0, 0, err
			}
			if !exists {
				todo = append(todo, pending{src, c, hash})
			}
		}
	}
	added := 0
	for i := 0; i < len(todo); i += 32 {
		j := i + 32
		if j > len(todo) {
			j = len(todo)
		}
		batch := todo[i:j]
		texts := make([]string, len(batch))
		for k, p := range batch {
			texts[k] = p.chunk
		}
		vecs, err := emb.Embed(ctx, model, texts)
		if err != nil {
			return added, 0, err
		}
		for k, p := range batch {
			if len(vecs[k]) == 0 {
				return added, 0, fmt.Errorf("empty embedding for chunk of %s", p.source)
			}
			if _, err := s.pool.Exec(ctx, `INSERT INTO ai.rag_docs (source, chunk, chunk_hash, embedding) VALUES ($1,$2,$3,$4) ON CONFLICT (chunk_hash) DO NOTHING`,
				p.source, p.chunk, p.hash, pgvector.NewVector(vecs[k])); err != nil {
				return added, 0, err
			}
			added++
		}
	}
	// drop stale chunks (file removed or content changed)
	rows, err := s.pool.Query(ctx, `SELECT chunk_hash, source FROM ai.rag_docs`)
	if err != nil {
		return added, 0, err
	}
	var stale []string
	for rows.Next() {
		var h, src string
		if err := rows.Scan(&h, &src); err != nil {
			rows.Close()
			return added, 0, err
		}
		if !keep[h] {
			stale = append(stale, h)
		}
	}
	rows.Close()
	for _, h := range stale {
		if _, err := s.pool.Exec(ctx, `DELETE FROM ai.rag_docs WHERE chunk_hash=$1`, h); err != nil {
			return added, len(stale), err
		}
	}
	return added, len(stale), nil
}

type Hit struct {
	Source   string  `json:"source"`
	Chunk    string  `json:"chunk"`
	Distance float64 `json:"distance"`
}

// Search returns the k nearest chunks by cosine distance (design 6.4: ORDER BY embedding <=> $1 LIMIT k).
func (s *Store) Search(ctx context.Context, vec []float32, k int) ([]Hit, error) {
	if k <= 0 || k > 20 {
		k = 5
	}
	rows, err := s.pool.Query(ctx, `SELECT source, chunk, embedding <=> $1 AS d FROM ai.rag_docs ORDER BY d LIMIT $2`, pgvector.NewVector(vec), k)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Hit
	for rows.Next() {
		var h Hit
		if err := rows.Scan(&h.Source, &h.Chunk, &h.Distance); err != nil {
			return nil, err
		}
		out = append(out, h)
	}
	return out, rows.Err()
}

func (s *Store) DocCount(ctx context.Context) (int, error) {
	var n int
	err := s.pool.QueryRow(ctx, `SELECT count(*) FROM ai.rag_docs`).Scan(&n)
	return n, err
}

// ---- sessions ----

type Session struct {
	SessionID     string
	RunID         *string
	ScenarioID    *string
	Model         string
	PromptVersion string
	Input         any
	Output        any
	ToolCalls     any
	Usage         any
	Error         string
	Duration      time.Duration
}

func (s *Store) SaveSession(ctx context.Context, se Session) error {
	in, _ := json.Marshal(se.Input)
	out, _ := json.Marshal(se.Output)
	tc, _ := json.Marshal(se.ToolCalls)
	us, _ := json.Marshal(se.Usage)
	_, err := s.pool.Exec(ctx, `INSERT INTO ai.sessions (session_id,run_id,scenario_id,model,prompt_version,input,output,tool_calls,usage,error,duration_ms)
	  VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,NULLIF($10,''),$11)
	  ON CONFLICT (session_id) DO UPDATE SET output=EXCLUDED.output, tool_calls=EXCLUDED.tool_calls, usage=EXCLUDED.usage, error=EXCLUDED.error, duration_ms=EXCLUDED.duration_ms`,
		se.SessionID, se.RunID, se.ScenarioID, se.Model, se.PromptVersion, in, out, tc, us, se.Error, se.Duration.Milliseconds())
	return err
}

// ---- HITL ledger ----

type HitlMessage struct {
	Token      string
	ChannelID  string
	MessageID  string
	ActionID   string
	Target     map[string]any
	RunID      *string
	ScenarioID *string
	SessionID  string
	ExpiresAt  time.Time
	DecidedAt  *time.Time
	Decision   *string
	DecidedBy  *string
}

func (s *Store) InsertHitl(ctx context.Context, m HitlMessage) error {
	t, _ := json.Marshal(m.Target)
	_, err := s.pool.Exec(ctx, `INSERT INTO ai.hitl_messages (token,channel_id,message_id,action_id,target,run_id,scenario_id,session_id,expires_at)
	  VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9)`, m.Token, m.ChannelID, m.MessageID, m.ActionID, t, m.RunID, m.ScenarioID, m.SessionID, m.ExpiresAt)
	return err
}

func (s *Store) GetHitl(ctx context.Context, token string) (*HitlMessage, error) {
	var m HitlMessage
	var t []byte
	err := s.pool.QueryRow(ctx, `SELECT token,channel_id,message_id,action_id,target,run_id,scenario_id,COALESCE(session_id,''),expires_at,decided_at,decision,decided_by FROM ai.hitl_messages WHERE token=$1`, token).
		Scan(&m.Token, &m.ChannelID, &m.MessageID, &m.ActionID, &t, &m.RunID, &m.ScenarioID, &m.SessionID, &m.ExpiresAt, &m.DecidedAt, &m.Decision, &m.DecidedBy)
	if err != nil {
		return nil, err
	}
	_ = json.Unmarshal(t, &m.Target)
	return &m, nil
}

// Decide records the first decision only; returns false when one was already recorded (double click / race).
func (s *Store) Decide(ctx context.Context, token, decision, by string, result any) (bool, error) {
	r, _ := json.Marshal(result)
	tag, err := s.pool.Exec(ctx, `UPDATE ai.hitl_messages SET decided_at=now(), decision=$2, decided_by=$3, result=$4 WHERE token=$1 AND decided_at IS NULL`, token, decision, by, r)
	if err != nil {
		return false, err
	}
	return tag.RowsAffected() == 1, nil
}

func (s *Store) SetHitlResult(ctx context.Context, token string, result any) error {
	r, _ := json.Marshal(result)
	_, err := s.pool.Exec(ctx, `UPDATE ai.hitl_messages SET result=$2 WHERE token=$1`, token, r)
	return err
}

// Pending returns undecided messages (restart safety: expiry timers are re-armed from this list).
func (s *Store) Pending(ctx context.Context) ([]HitlMessage, error) {
	rows, err := s.pool.Query(ctx, `SELECT token,channel_id,message_id,action_id,target,run_id,scenario_id,COALESCE(session_id,''),expires_at FROM ai.hitl_messages WHERE decided_at IS NULL`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []HitlMessage
	for rows.Next() {
		var m HitlMessage
		var t []byte
		if err := rows.Scan(&m.Token, &m.ChannelID, &m.MessageID, &m.ActionID, &t, &m.RunID, &m.ScenarioID, &m.SessionID, &m.ExpiresAt); err != nil {
			return nil, err
		}
		_ = json.Unmarshal(t, &m.Target)
		out = append(out, m)
	}
	return out, rows.Err()
}