From 1a433dc59628d0391ea9964627874f3120198025 Mon Sep 17 00:00:00 2001 From: ZanzyTHEbar Date: Sat, 21 Feb 2026 18:57:46 +0000 Subject: [PATCH] feat(memory): add kernel contract, DAG snapshot persistence, and map_ops schema pkg/memory/kernel_contract.go - Defines the MemoryKernelContract interface for active context projection: ActiveContext(), ProjectSession(), SyncProjection(), RecoverSession() - Provides a deterministic contract for lossless session continuity with DAG-based compression and projection pointer tracking pkg/memory/dag/store.go - DAGStore: persistent snapshot storage for DAG nodes and edges in SQLite - Stores compression snapshots with provenance metadata (origin session, message range, token counts) - Supports node/edge insertion, snapshot listing, and node lookup by ID pkg/memory/dag/backfill.go + backfill_test.go - BackfillDAGFromSession(): builds a DAG snapshot from an existing session history for sessions that pre-date DAG compression - Idempotent: skips sessions that already have a snapshot pkg/memory/migrations/011_dag_tables.go - SQL migration adding dag_snapshots, dag_nodes, and dag_edges tables - Includes indexes for snapshot_id, session_key, and node type lookups pkg/memory/migrations/012_map_operator_runs.go - SQL migration adding map_operator_runs and map_operator_items tables - Tracks batch map operation runs with per-item status and FlatBuffers serialized payloads pkg/memory/sqlc/queries/dag.sql + dag.sql.go - SQLC queries for DAG snapshot CRUD operations pkg/memory/sqlc/queries/map_ops.sql + map_ops.sql.go - SQLC queries for map operator run and item tracking pkg/memory/store/retrieval_policy.go - RetrievalPolicy: configures token budget, recency bias, and semantic score thresholds for memory retrieval operations - Used by the kernel contract to bound context window usage pkg/memory/delegate/sqlite_dag_test.go - Integration tests for DAG table creation, snapshot insertion, and node/edge persistence through the SQLite delegate --- pkg/memory/dag/backfill.go | 236 +++++ pkg/memory/dag/backfill_test.go | 79 ++ pkg/memory/dag/store.go | 166 ++++ pkg/memory/delegate/sqlite_dag_test.go | 50 ++ pkg/memory/kernel_contract.go | 95 ++ pkg/memory/kernel_contract_test.go | 47 + pkg/memory/migrations/011_dag_tables.go | 86 ++ .../migrations/012_map_operator_runs.go | 77 ++ pkg/memory/sqlc/dag.sql.go | 638 ++++++++++++++ pkg/memory/sqlc/map_ops.sql.go | 815 ++++++++++++++++++ pkg/memory/sqlc/queries/dag.sql | 172 ++++ pkg/memory/sqlc/queries/map_ops.sql | 143 +++ pkg/memory/store/retrieval_policy.go | 404 +++++++++ pkg/memory/store/token_test_helpers_test.go | 18 + 14 files changed, 3026 insertions(+) create mode 100644 pkg/memory/dag/backfill.go create mode 100644 pkg/memory/dag/backfill_test.go create mode 100644 pkg/memory/dag/store.go create mode 100644 pkg/memory/delegate/sqlite_dag_test.go create mode 100644 pkg/memory/kernel_contract.go create mode 100644 pkg/memory/kernel_contract_test.go create mode 100644 pkg/memory/migrations/011_dag_tables.go create mode 100644 pkg/memory/migrations/012_map_operator_runs.go create mode 100644 pkg/memory/sqlc/dag.sql.go create mode 100644 pkg/memory/sqlc/map_ops.sql.go create mode 100644 pkg/memory/sqlc/queries/dag.sql create mode 100644 pkg/memory/sqlc/queries/map_ops.sql create mode 100644 pkg/memory/store/retrieval_policy.go create mode 100644 pkg/memory/store/token_test_helpers_test.go diff --git a/pkg/memory/dag/backfill.go b/pkg/memory/dag/backfill.go new file mode 100644 index 000000000..e78477aa0 --- /dev/null +++ b/pkg/memory/dag/backfill.go @@ -0,0 +1,236 @@ +package dag + +import ( + "context" + "database/sql" + "errors" + "sort" + "strings" + "time" + + jsonv2 "github.com/go-json-experiment/json" + + "github.com/ZanzyTHEbar/dragonscale/pkg/memory" + memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc" +) + +const dagBackfillStatusKVKey = "migration:dag_backfill:v1" + +// BackfillOptions controls DAG backfill behavior. +type BackfillOptions struct { + PageSize int + MaxSessions int + Force bool +} + +func DefaultBackfillOptions() BackfillOptions { + return BackfillOptions{ + PageSize: 500, + MaxSessions: 0, + Force: false, + } +} + +// BackfillStatus tracks one DAG backfill pass over existing sessions. +type BackfillStatus struct { + Version int `json:"version"` + SessionsScanned int `json:"sessions_scanned"` + SnapshotsCreated int `json:"snapshots_created"` + SkippedExisting int `json:"skipped_existing"` + Failures int `json:"failures"` + CompletedAt time.Time `json:"completed_at"` + Skipped bool `json:"skipped,omitempty"` +} + +// BackfillMissingSessionDAGs creates DAG snapshots for sessions that have +// session-message history but no persisted DAG snapshot yet. +// +// The pass is one-shot by default and records status in KV under +// dagBackfillStatusKVKey. Use opts.Force=true to run again. +func BackfillMissingSessionDAGs( + ctx context.Context, + delegate memory.MemoryDelegate, + queries *memsqlc.Queries, + agentID string, + opts BackfillOptions, +) (*BackfillStatus, error) { + if delegate == nil || queries == nil { + return nil, nil + } + persister, ok := delegate.(DAGPersister) + if !ok { + return nil, nil + } + + if opts.PageSize <= 0 { + opts.PageSize = 500 + } + + if !opts.Force { + if raw, err := delegate.GetKV(ctx, agentID, dagBackfillStatusKVKey); err == nil && strings.TrimSpace(raw) != "" { + var status BackfillStatus + if uErr := jsonv2.Unmarshal([]byte(raw), &status); uErr == nil { + status.Skipped = true + return &status, nil + } + return &BackfillStatus{Version: 1, Skipped: true}, nil + } + } + + sessionKeys, err := collectSessionKeysForBackfill(ctx, delegate, agentID, opts.PageSize) + if err != nil { + return nil, err + } + if opts.MaxSessions > 0 && len(sessionKeys) > opts.MaxSessions { + sessionKeys = sessionKeys[:opts.MaxSessions] + } + + status := &BackfillStatus{ + Version: 1, + SessionsScanned: len(sessionKeys), + } + compressor := NewCompressor(DefaultCompressorConfig()) + + for _, sessionKey := range sessionKeys { + if ctx.Err() != nil { + return status, ctx.Err() + } + + _, err := queries.GetLatestDAGSnapshotBySession(ctx, memsqlc.GetLatestDAGSnapshotBySessionParams{ + AgentID: agentID, + SessionKey: sessionKey, + }) + if err == nil { + status.SkippedExisting++ + continue + } + if !errors.Is(err, sql.ErrNoRows) { + status.Failures++ + continue + } + + msgs, err := loadSessionMessagesForBackfill(ctx, delegate, agentID, sessionKey, opts.PageSize) + if err != nil { + status.Failures++ + continue + } + if len(msgs) == 0 { + continue + } + + d := compressor.Compress(msgs) + if d == nil || len(d.Nodes) == 0 { + continue + } + + if err := persister.PersistDAG(ctx, agentID, sessionKey, &PersistSnapshot{ + FromMsgIdx: 0, + ToMsgIdx: len(msgs), + MsgCount: len(msgs), + DAG: d, + }); err != nil { + status.Failures++ + continue + } + status.SnapshotsCreated++ + } + + status.CompletedAt = time.Now().UTC() + data, err := jsonv2.Marshal(status) + if err != nil { + return status, err + } + if err := delegate.UpsertKV(ctx, agentID, dagBackfillStatusKVKey, string(data)); err != nil { + return status, err + } + return status, nil +} + +func collectSessionKeysForBackfill(ctx context.Context, delegate memory.MemoryDelegate, agentID string, pageSize int) ([]string, error) { + keys := make(map[string]struct{}) + offset := 0 + for { + items, err := delegate.ListRecallItems(ctx, agentID, "", pageSize, offset) + if err != nil { + return nil, err + } + if len(items) == 0 { + break + } + for _, item := range items { + if !strings.Contains(item.Tags, "session-message") { + continue + } + keys[item.SessionKey] = struct{}{} + } + if len(items) < pageSize { + break + } + offset += len(items) + } + + out := make([]string, 0, len(keys)) + for k := range keys { + out = append(out, k) + } + sort.Strings(out) + return out, nil +} + +func loadSessionMessagesForBackfill( + ctx context.Context, + delegate memory.MemoryDelegate, + agentID string, + sessionKey string, + pageSize int, +) ([]Message, error) { + type sessionLister interface { + ListSessionMessages(ctx context.Context, agentID, sessionKey, role string, limit int) ([]*memory.RecallItem, error) + } + + if lister, ok := delegate.(sessionLister); ok { + count, err := delegate.CountRecallItems(ctx, agentID, sessionKey) + if err != nil { + count = 0 + } + limit := count + 32 + if limit < 32 { + limit = 32 + } + rows, err := lister.ListSessionMessages(ctx, agentID, sessionKey, "", limit) + if err == nil { + msgs := make([]Message, 0, len(rows)) + for _, row := range rows { + msgs = append(msgs, Message{Role: row.Role, Content: row.Content}) + } + return msgs, nil + } + } + + var all []*memory.RecallItem + offset := 0 + for { + items, err := delegate.ListRecallItems(ctx, agentID, sessionKey, pageSize, offset) + if err != nil { + return nil, err + } + if len(items) == 0 { + break + } + all = append(all, items...) + if len(items) < pageSize { + break + } + offset += len(items) + } + + msgs := make([]Message, 0, len(all)) + for i := len(all) - 1; i >= 0; i-- { + item := all[i] + if !strings.Contains(item.Tags, "session-message") { + continue + } + msgs = append(msgs, Message{Role: item.Role, Content: item.Content}) + } + return msgs, nil +} diff --git a/pkg/memory/dag/backfill_test.go b/pkg/memory/dag/backfill_test.go new file mode 100644 index 000000000..d5e40ca81 --- /dev/null +++ b/pkg/memory/dag/backfill_test.go @@ -0,0 +1,79 @@ +package dag_test + +import ( + "context" + "fmt" + "testing" + "time" + + jsonv2 "github.com/go-json-experiment/json" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/ZanzyTHEbar/dragonscale/pkg/ids" + "github.com/ZanzyTHEbar/dragonscale/pkg/memory" + "github.com/ZanzyTHEbar/dragonscale/pkg/memory/dag" + "github.com/ZanzyTHEbar/dragonscale/pkg/memory/delegate" + memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc" +) + +func TestBackfillMissingSessionDAGs_CreatesSnapshotsAndPersistsStatus(t *testing.T) { + ctx := context.Background() + d, err := delegate.NewLibSQLInMemory() + require.NoError(t, err) + require.NoError(t, d.Init(ctx)) + defer d.Close() + + agentID := "agent-backfill" + sessionKey := "legacy-session-1" + for i := 0; i < 12; i++ { + role := "user" + if i%2 == 1 { + role = "assistant" + } + require.NoError(t, d.InsertRecallItem(ctx, &memory.RecallItem{ + ID: ids.New(), + AgentID: agentID, + SessionKey: sessionKey, + Role: role, + Sector: memory.SectorEpisodic, + Importance: 0.5, + Salience: 0.5, + DecayRate: 0.01, + Content: fmt.Sprintf("legacy message %d", i), + Tags: "session-message", + CreatedAt: time.Now().UTC(), + UpdatedAt: time.Now().UTC(), + })) + } + + status, err := dag.BackfillMissingSessionDAGs(ctx, d, d.Queries(), agentID, dag.DefaultBackfillOptions()) + require.NoError(t, err) + require.NotNil(t, status) + assert.Equal(t, 1, status.SnapshotsCreated) + assert.Equal(t, 0, status.Failures) + + row, err := d.Queries().GetLatestDAGSnapshotBySession(ctx, memsqlc.GetLatestDAGSnapshotBySessionParams{ + AgentID: agentID, + SessionKey: sessionKey, + }) + require.NoError(t, err) + assert.Equal(t, int64(12), row.MsgCount) + + kv, err := d.ListKVByPrefix(ctx, agentID, "migration:dag_backfill", 10) + require.NoError(t, err) + require.NotEmpty(t, kv) + + var stored dag.BackfillStatus + for _, raw := range kv { + require.NoError(t, jsonv2.Unmarshal([]byte(raw), &stored)) + break + } + assert.Equal(t, status.SnapshotsCreated, stored.SnapshotsCreated) + assert.False(t, stored.CompletedAt.IsZero()) + + status2, err := dag.BackfillMissingSessionDAGs(ctx, d, d.Queries(), agentID, dag.DefaultBackfillOptions()) + require.NoError(t, err) + require.NotNil(t, status2) + assert.True(t, status2.Skipped) +} diff --git a/pkg/memory/dag/store.go b/pkg/memory/dag/store.go new file mode 100644 index 000000000..efbb49b25 --- /dev/null +++ b/pkg/memory/dag/store.go @@ -0,0 +1,166 @@ +package dag + +import ( + "context" + "crypto/sha256" + "database/sql" + "encoding/hex" + "encoding/json" + "fmt" + "sort" + + "github.com/ZanzyTHEbar/dragonscale/pkg/ids" + memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc" +) + +// PersistSnapshot is the serializable shape for persisting a DAG snapshot. +type PersistSnapshot struct { + FromMsgIdx int + ToMsgIdx int + MsgCount int + DAG *DAG +} + +// contentHash returns a deterministic hash of the DAG structure for deduplication. +func contentHash(d *DAG) string { + if d == nil || len(d.Nodes) == 0 { + return "" + } + h := sha256.New() + roots, _ := json.Marshal(d.Roots) + h.Write(roots) + + nodeIDs := make([]string, 0, len(d.Nodes)) + for id := range d.Nodes { + nodeIDs = append(nodeIDs, id) + } + sort.Strings(nodeIDs) + + for _, id := range nodeIDs { + n := d.Nodes[id] + h.Write([]byte(id)) + h.Write([]byte(n.Summary)) + h.Write([]byte(fmt.Sprintf("%d:%d:%d", n.StartIdx, n.EndIdx, n.Tokens))) + if len(n.Children) > 0 { + children := append([]string(nil), n.Children...) + sort.Strings(children) + for _, child := range children { + h.Write([]byte(child)) + } + } + } + return hex.EncodeToString(h.Sum(nil)) +} + +// DAGPersister persists DAG snapshots to storage. +// Implemented by delegates that support DAG persistence. +type DAGPersister interface { + PersistDAG(ctx context.Context, agentID, sessionKey string, snap *PersistSnapshot) error +} + +// PersistDAG persists a DAG snapshot via sqlc in a transaction. +func PersistDAG(ctx context.Context, db *sql.DB, q *memsqlc.Queries, agentID, sessionKey string, snap *PersistSnapshot) error { + if snap == nil || snap.DAG == nil || len(snap.DAG.Nodes) == 0 { + return nil + } + + hash := contentHash(snap.DAG) + snapshotID := ids.New() + + rootsJSON, err := json.Marshal(snap.DAG.Roots) + if err != nil { + return fmt.Errorf("marshal roots: %w", err) + } + + tx, err := db.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("begin tx: %w", err) + } + defer tx.Rollback() + + qTx := q.WithTx(tx) + + _, err = qTx.InsertDAGSnapshot(ctx, memsqlc.InsertDAGSnapshotParams{ + ID: snapshotID, + AgentID: agentID, + SessionKey: sessionKey, + FromMsgIdx: int64(snap.FromMsgIdx), + ToMsgIdx: int64(snap.ToMsgIdx), + MsgCount: int64(snap.MsgCount), + RootsJson: string(rootsJSON), + ContentHash: hash, + }) + if err != nil { + return fmt.Errorf("insert snapshot: %w", err) + } + + nodeIDsOrdered := make([]string, 0, len(snap.DAG.Nodes)) + for nodeID := range snap.DAG.Nodes { + nodeIDsOrdered = append(nodeIDsOrdered, nodeID) + } + sort.Strings(nodeIDsOrdered) + + nodeIDs := make(map[string]ids.UUID) + for _, nodeID := range nodeIDsOrdered { + node := snap.DAG.Nodes[nodeID] + nodeUUID := ids.New() + nodeIDs[nodeID] = nodeUUID + + nodeHash := sha256.Sum256([]byte(fmt.Sprintf("%s|%s|%d|%d|%d", node.ID, node.Summary, node.StartIdx, node.EndIdx, node.Tokens))) + + metricsJSONBytes, _ := json.Marshal(map[string]any{ + "tokens": node.Tokens, + "span": node.Span(), + "children_count": len(node.Children), + }) + metadataJSONBytes, _ := json.Marshal(map[string]any{ + "level": node.Level.String(), + "children": node.Children, + }) + _, err = qTx.InsertDAGNode(ctx, memsqlc.InsertDAGNodeParams{ + ID: nodeUUID, + SnapshotID: snapshotID, + NodeID: node.ID, + Level: int64(node.Level), + Summary: node.Summary, + Tokens: int64(node.Tokens), + StartIdx: int64(node.StartIdx), + EndIdx: int64(node.EndIdx), + Span: int64(node.Span()), + ContentHash: hex.EncodeToString(nodeHash[:]), + MetricsJson: metricsJSONBytes, + MetadataJson: metadataJSONBytes, + }) + if err != nil { + return fmt.Errorf("insert node %s: %w", node.ID, err) + } + } + + for _, nodeID := range nodeIDsOrdered { + node := snap.DAG.Nodes[nodeID] + for edgeIdx, childID := range node.Children { + childUUID, ok := nodeIDs[childID] + if !ok { + continue + } + parentUUID := nodeIDs[node.ID] + edgeMetadata, _ := json.Marshal(map[string]any{ + "parent_node_id": node.ID, + "child_node_id": childID, + }) + err = qTx.InsertDAGEdge(ctx, memsqlc.InsertDAGEdgeParams{ + ID: ids.New(), + SnapshotID: snapshotID, + ParentNodeID: parentUUID, + ChildNodeID: childUUID, + EdgeIndex: int64(edgeIdx), + MetadataJson: edgeMetadata, + }) + if err != nil { + return fmt.Errorf("insert edge %s->%s: %w", node.ID, childID, err) + } + } + } + + return tx.Commit() +} diff --git a/pkg/memory/delegate/sqlite_dag_test.go b/pkg/memory/delegate/sqlite_dag_test.go new file mode 100644 index 000000000..d2b59553a --- /dev/null +++ b/pkg/memory/delegate/sqlite_dag_test.go @@ -0,0 +1,50 @@ +package delegate + +import ( + "context" + "testing" + + "github.com/ZanzyTHEbar/dragonscale/pkg/memory/dag" + memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc" + "github.com/stretchr/testify/require" +) + +func TestLibSQLDelegate_PersistDAG(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + + compressor := dag.NewCompressor(dag.DefaultCompressorConfig()) + msgs := make([]dag.Message, 16) + for i := range msgs { + role := "user" + if i%2 == 1 { + role = "assistant" + } + msgs[i] = dag.Message{Role: role, Content: "msg " + string(rune('A'+i%26))} + } + dagOut := compressor.Compress(msgs) + require.NotEmpty(t, dagOut.Nodes) + + snap := &dag.PersistSnapshot{ + FromMsgIdx: 0, + ToMsgIdx: 16, + MsgCount: 16, + DAG: dagOut, + } + require.NoError(t, d.PersistDAG(ctx, "agent1", "session1", snap)) + + row, err := d.Queries().GetLatestDAGSnapshotBySession(ctx, memsqlc.GetLatestDAGSnapshotBySessionParams{ + AgentID: "agent1", + SessionKey: "session1", + }) + require.NoError(t, err) + require.Equal(t, "agent1", row.AgentID) + require.Equal(t, "session1", row.SessionKey) + require.Equal(t, int64(16), row.MsgCount) + + nodes, err := d.Queries().ListDAGNodesBySnapshotID(ctx, memsqlc.ListDAGNodesBySnapshotIDParams{ + SnapshotID: row.ID, + }) + require.NoError(t, err) + require.NotEmpty(t, nodes) +} diff --git a/pkg/memory/kernel_contract.go b/pkg/memory/kernel_contract.go new file mode 100644 index 000000000..b790e2e5d --- /dev/null +++ b/pkg/memory/kernel_contract.go @@ -0,0 +1,95 @@ +package memory + +import ( + "context" + "time" + + "github.com/ZanzyTHEbar/dragonscale/pkg/ids" +) + +// ProjectionSegmentKind classifies what role a segment plays in active context. +type ProjectionSegmentKind string + +const ( + ProjectionSegmentSystem ProjectionSegmentKind = "system" + ProjectionSegmentRecent ProjectionSegmentKind = "recent" + ProjectionSegmentDAG ProjectionSegmentKind = "dag" + ProjectionSegmentRecall ProjectionSegmentKind = "recall" + ProjectionSegmentArchival ProjectionSegmentKind = "archival" + ProjectionSegmentTool ProjectionSegmentKind = "tool_result" +) + +// ImmutableSpanRef is a lossless reference back into immutable session history. +type ImmutableSpanRef struct { + SessionKey string `json:"session_key"` + StartIdx int `json:"start_idx"` // inclusive + EndIdx int `json:"end_idx"` // exclusive + FirstID ids.UUID `json:"first_id"` + LastID ids.UUID `json:"last_id"` + FromTime time.Time `json:"from_time"` + ToTime time.Time `json:"to_time"` +} + +// ProjectionSegment is one unit included in active context. +// Every segment must carry a lossless reference back to immutable history. +type ProjectionSegment struct { + Kind ProjectionSegmentKind `json:"kind"` + Source string `json:"source"` + Text string `json:"text"` + Tokens int `json:"tokens"` + Ref ImmutableSpanRef `json:"ref"` +} + +// ActiveContextProjection is the assembled view injected into the model for a turn. +// It is a materialized projection, not a source of truth. +type ActiveContextProjection struct { + AgentID string `json:"agent_id"` + SessionKey string `json:"session_key"` + BudgetTokens int `json:"budget_tokens"` + GeneratedAt time.Time `json:"generated_at"` + ProjectionRef string `json:"projection_ref"` + Segments []ProjectionSegment `json:"segments"` +} + +func (p *ActiveContextProjection) TotalTokens() int { + if p == nil { + return 0 + } + total := 0 + for _, seg := range p.Segments { + total += seg.Tokens + } + return total +} + +// HasLosslessRefs verifies the projection preserves deterministic pointers to +// immutable history for every segment. +func (p *ActiveContextProjection) HasLosslessRefs() bool { + if p == nil { + return false + } + for _, seg := range p.Segments { + if seg.Ref.SessionKey == "" || seg.Ref.EndIdx < seg.Ref.StartIdx { + return false + } + } + return true +} + +type ProjectionRequest struct { + AgentID string + SessionKey string + MaxTokens int + IncludeTools bool +} + +// ImmutableHistoryReader resolves lossless references to original messages. +type ImmutableHistoryReader interface { + ListRecallItems(ctx context.Context, agentID, sessionKey string, limit, offset int) ([]*RecallItem, error) +} + +// ActiveContextBuilder materializes the turn-time projection from immutable +// storage, retrieval tiers, and DAG summaries. +type ActiveContextBuilder interface { + BuildActiveContext(ctx context.Context, req ProjectionRequest) (*ActiveContextProjection, error) +} diff --git a/pkg/memory/kernel_contract_test.go b/pkg/memory/kernel_contract_test.go new file mode 100644 index 000000000..0ef4ce8ae --- /dev/null +++ b/pkg/memory/kernel_contract_test.go @@ -0,0 +1,47 @@ +package memory + +import ( + "testing" + "time" + + "github.com/ZanzyTHEbar/dragonscale/pkg/ids" + "github.com/stretchr/testify/assert" +) + +func TestActiveContextProjection_TotalTokens(t *testing.T) { + p := &ActiveContextProjection{ + Segments: []ProjectionSegment{ + {Tokens: 120}, + {Tokens: 80}, + {Tokens: 35}, + }, + } + assert.Equal(t, 235, p.TotalTokens()) +} + +func TestActiveContextProjection_HasLosslessRefs(t *testing.T) { + now := time.Now().UTC() + p := &ActiveContextProjection{ + Segments: []ProjectionSegment{ + { + Kind: ProjectionSegmentRecent, + Source: "session-tail", + Text: "latest messages", + Tokens: 64, + Ref: ImmutableSpanRef{ + SessionKey: "s1", + StartIdx: 10, + EndIdx: 15, + FirstID: ids.New(), + LastID: ids.New(), + FromTime: now, + ToTime: now, + }, + }, + }, + } + assert.True(t, p.HasLosslessRefs()) + + p.Segments[0].Ref.SessionKey = "" + assert.False(t, p.HasLosslessRefs()) +} diff --git a/pkg/memory/migrations/011_dag_tables.go b/pkg/memory/migrations/011_dag_tables.go new file mode 100644 index 000000000..9bed9a8d5 --- /dev/null +++ b/pkg/memory/migrations/011_dag_tables.go @@ -0,0 +1,86 @@ +package migrations + +import ( + "context" + "database/sql" + "fmt" + + "github.com/pressly/goose/v3" +) + +func init() { + goose.AddMigrationContext(up011DAGTables, down011DAGTables) +} + +func up011DAGTables(ctx context.Context, tx *sql.Tx) error { + stmts := []string{ + `CREATE TABLE IF NOT EXISTS dag_snapshots ( + id BLOB PRIMARY KEY, + agent_id TEXT NOT NULL, + session_key TEXT NOT NULL DEFAULT '', + from_msg_idx INTEGER NOT NULL DEFAULT 0, + to_msg_idx INTEGER NOT NULL DEFAULT 0, + msg_count INTEGER NOT NULL DEFAULT 0, + roots_json TEXT NOT NULL DEFAULT '[]', + content_hash TEXT NOT NULL DEFAULT '', + created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + updated_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')) +)`, + `CREATE INDEX IF NOT EXISTS idx_dag_snapshots_agent_session ON dag_snapshots(agent_id, session_key)`, + `CREATE INDEX IF NOT EXISTS idx_dag_snapshots_created_at ON dag_snapshots(created_at DESC)`, + + `CREATE TABLE IF NOT EXISTS dag_nodes ( + id BLOB PRIMARY KEY, + snapshot_id BLOB NOT NULL REFERENCES dag_snapshots(id) ON DELETE CASCADE, + node_id TEXT NOT NULL, + level INTEGER NOT NULL DEFAULT 1, + summary TEXT NOT NULL DEFAULT '', + tokens INTEGER NOT NULL DEFAULT 0, + start_idx INTEGER NOT NULL DEFAULT 0, + end_idx INTEGER NOT NULL DEFAULT 0, + span INTEGER NOT NULL DEFAULT 0, + content_hash TEXT NOT NULL DEFAULT '', + metrics_json JSON NOT NULL DEFAULT '{}', + metadata_json JSON NOT NULL DEFAULT '{}', + created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + updated_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')) +)`, + `CREATE INDEX IF NOT EXISTS idx_dag_nodes_snapshot ON dag_nodes(snapshot_id)`, + `CREATE INDEX IF NOT EXISTS idx_dag_nodes_node_id ON dag_nodes(snapshot_id, node_id)`, + `CREATE INDEX IF NOT EXISTS idx_dag_nodes_level_start ON dag_nodes(snapshot_id, level, start_idx)`, + + `CREATE TABLE IF NOT EXISTS dag_edges ( + id BLOB PRIMARY KEY, + snapshot_id BLOB NOT NULL REFERENCES dag_snapshots(id) ON DELETE CASCADE, + parent_node_id BLOB NOT NULL REFERENCES dag_nodes(id) ON DELETE CASCADE, + child_node_id BLOB NOT NULL REFERENCES dag_nodes(id) ON DELETE CASCADE, + edge_index INTEGER NOT NULL DEFAULT 0, + metadata_json JSON NOT NULL DEFAULT '{}', + created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + updated_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')) +)`, + `CREATE INDEX IF NOT EXISTS idx_dag_edges_snapshot ON dag_edges(snapshot_id)`, + `CREATE INDEX IF NOT EXISTS idx_dag_edges_parent ON dag_edges(parent_node_id)`, + `CREATE INDEX IF NOT EXISTS idx_dag_edges_child ON dag_edges(child_node_id)`, + } + for _, s := range stmts { + if _, err := tx.ExecContext(ctx, s); err != nil { + return fmt.Errorf("011_dag_tables up: %w\nSQL: %s", err, s) + } + } + return nil +} + +func down011DAGTables(ctx context.Context, tx *sql.Tx) error { + stmts := []string{ + `DROP TABLE IF EXISTS dag_edges`, + `DROP TABLE IF EXISTS dag_nodes`, + `DROP TABLE IF EXISTS dag_snapshots`, + } + for _, s := range stmts { + if _, err := tx.ExecContext(ctx, s); err != nil { + return fmt.Errorf("011_dag_tables down: %w\nSQL: %s", err, s) + } + } + return nil +} diff --git a/pkg/memory/migrations/012_map_operator_runs.go b/pkg/memory/migrations/012_map_operator_runs.go new file mode 100644 index 000000000..cc98cb1c4 --- /dev/null +++ b/pkg/memory/migrations/012_map_operator_runs.go @@ -0,0 +1,77 @@ +package migrations + +import ( + "context" + "database/sql" + "fmt" + + "github.com/pressly/goose/v3" +) + +func init() { + goose.AddMigrationContext(up012MapOperatorRuns, down012MapOperatorRuns) +} + +func up012MapOperatorRuns(ctx context.Context, tx *sql.Tx) error { + stmts := []string{ + `CREATE TABLE IF NOT EXISTS map_runs ( + id BLOB PRIMARY KEY, + agent_id TEXT NOT NULL, + session_key TEXT NOT NULL, + operator_kind TEXT NOT NULL, + idempotency_key TEXT, + status TEXT NOT NULL DEFAULT 'queued', + total_items INTEGER NOT NULL DEFAULT 0, + queued_items INTEGER NOT NULL DEFAULT 0, + running_items INTEGER NOT NULL DEFAULT 0, + succeeded_items INTEGER NOT NULL DEFAULT 0, + failed_items INTEGER NOT NULL DEFAULT 0, + spec_fb BLOB NOT NULL, + last_error TEXT, + created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + updated_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + completed_at DATETIME +)`, + `CREATE INDEX IF NOT EXISTS idx_map_runs_agent_session_created_at ON map_runs(agent_id, session_key, created_at DESC)`, + `CREATE INDEX IF NOT EXISTS idx_map_runs_status_updated_at ON map_runs(status, updated_at DESC)`, + `CREATE UNIQUE INDEX IF NOT EXISTS idx_map_runs_dedupe ON map_runs(agent_id, session_key, operator_kind, idempotency_key) WHERE idempotency_key IS NOT NULL`, + `CREATE TABLE IF NOT EXISTS map_items ( + id BLOB PRIMARY KEY, + run_id BLOB NOT NULL REFERENCES map_runs(id) ON DELETE CASCADE, + item_index INTEGER NOT NULL, + status TEXT NOT NULL DEFAULT 'queued', + attempts INTEGER NOT NULL DEFAULT 0, + last_error TEXT, + input_fb BLOB NOT NULL, + output_fb BLOB, + input_hash TEXT, + output_hash TEXT, + created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + updated_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + completed_at DATETIME, + UNIQUE(run_id, item_index) +)`, + `CREATE INDEX IF NOT EXISTS idx_map_items_run_id_status_item_index ON map_items(run_id, status, item_index)`, + `CREATE INDEX IF NOT EXISTS idx_map_items_run_id_item_index ON map_items(run_id, item_index)`, + } + + for _, s := range stmts { + if _, err := tx.ExecContext(ctx, s); err != nil { + return fmt.Errorf("012_map_operator_runs up: %w\nSQL: %s", err, s) + } + } + return nil +} + +func down012MapOperatorRuns(ctx context.Context, tx *sql.Tx) error { + stmts := []string{ + `DROP TABLE IF EXISTS map_items`, + `DROP TABLE IF EXISTS map_runs`, + } + for _, s := range stmts { + if _, err := tx.ExecContext(ctx, s); err != nil { + return fmt.Errorf("012_map_operator_runs down: %w\nSQL: %s", err, s) + } + } + return nil +} diff --git a/pkg/memory/sqlc/dag.sql.go b/pkg/memory/sqlc/dag.sql.go new file mode 100644 index 000000000..848e31f6a --- /dev/null +++ b/pkg/memory/sqlc/dag.sql.go @@ -0,0 +1,638 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: dag.sql + +package sqlc + +import ( + "context" + "encoding/json" + + "github.com/ZanzyTHEbar/dragonscale/pkg/ids" +) + +const GetDAGNodeBySnapshotAndNodeID = `-- name: GetDAGNodeBySnapshotAndNodeID :one +SELECT id, + snapshot_id, + node_id, + level, + summary, + tokens, + start_idx, + end_idx, + span, + content_hash, + metrics_json, + metadata_json, + created_at, + updated_at +FROM dag_nodes +WHERE snapshot_id = ?1 + AND node_id = ?2 +LIMIT 1 +` + +type GetDAGNodeBySnapshotAndNodeIDParams struct { + SnapshotID ids.UUID `db:"snapshot_id" json:"snapshot_id"` + NodeID string `db:"node_id" json:"node_id"` +} + +// GetDAGNodeBySnapshotAndNodeID +// +// SELECT id, +// snapshot_id, +// node_id, +// level, +// summary, +// tokens, +// start_idx, +// end_idx, +// span, +// content_hash, +// metrics_json, +// metadata_json, +// created_at, +// updated_at +// FROM dag_nodes +// WHERE snapshot_id = ?1 +// AND node_id = ?2 +// LIMIT 1 +func (q *Queries) GetDAGNodeBySnapshotAndNodeID(ctx context.Context, arg GetDAGNodeBySnapshotAndNodeIDParams) (DagNode, error) { + row := q.db.QueryRowContext(ctx, GetDAGNodeBySnapshotAndNodeID, arg.SnapshotID, arg.NodeID) + var i DagNode + err := row.Scan( + &i.ID, + &i.SnapshotID, + &i.NodeID, + &i.Level, + &i.Summary, + &i.Tokens, + &i.StartIdx, + &i.EndIdx, + &i.Span, + &i.ContentHash, + &i.MetricsJson, + &i.MetadataJson, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} + +const GetDAGSnapshotByID = `-- name: GetDAGSnapshotByID :one +SELECT id, + agent_id, + session_key, + from_msg_idx, + to_msg_idx, + msg_count, + roots_json, + content_hash, + created_at, + updated_at +FROM dag_snapshots +WHERE id = ?1 +LIMIT 1 +` + +type GetDAGSnapshotByIDParams struct { + ID ids.UUID `db:"id" json:"id"` +} + +// GetDAGSnapshotByID +// +// SELECT id, +// agent_id, +// session_key, +// from_msg_idx, +// to_msg_idx, +// msg_count, +// roots_json, +// content_hash, +// created_at, +// updated_at +// FROM dag_snapshots +// WHERE id = ?1 +// LIMIT 1 +func (q *Queries) GetDAGSnapshotByID(ctx context.Context, arg GetDAGSnapshotByIDParams) (DagSnapshot, error) { + row := q.db.QueryRowContext(ctx, GetDAGSnapshotByID, arg.ID) + var i DagSnapshot + err := row.Scan( + &i.ID, + &i.AgentID, + &i.SessionKey, + &i.FromMsgIdx, + &i.ToMsgIdx, + &i.MsgCount, + &i.RootsJson, + &i.ContentHash, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} + +const GetLatestDAGSnapshotBySession = `-- name: GetLatestDAGSnapshotBySession :one +SELECT id, + agent_id, + session_key, + from_msg_idx, + to_msg_idx, + msg_count, + roots_json, + content_hash, + created_at, + updated_at +FROM dag_snapshots +WHERE agent_id = ?1 + AND session_key = ?2 +ORDER BY created_at DESC +LIMIT 1 +` + +type GetLatestDAGSnapshotBySessionParams struct { + AgentID string `db:"agent_id" json:"agent_id"` + SessionKey string `db:"session_key" json:"session_key"` +} + +// GetLatestDAGSnapshotBySession +// +// SELECT id, +// agent_id, +// session_key, +// from_msg_idx, +// to_msg_idx, +// msg_count, +// roots_json, +// content_hash, +// created_at, +// updated_at +// FROM dag_snapshots +// WHERE agent_id = ?1 +// AND session_key = ?2 +// ORDER BY created_at DESC +// LIMIT 1 +func (q *Queries) GetLatestDAGSnapshotBySession(ctx context.Context, arg GetLatestDAGSnapshotBySessionParams) (DagSnapshot, error) { + row := q.db.QueryRowContext(ctx, GetLatestDAGSnapshotBySession, arg.AgentID, arg.SessionKey) + var i DagSnapshot + err := row.Scan( + &i.ID, + &i.AgentID, + &i.SessionKey, + &i.FromMsgIdx, + &i.ToMsgIdx, + &i.MsgCount, + &i.RootsJson, + &i.ContentHash, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} + +const InsertDAGEdge = `-- name: InsertDAGEdge :exec +INSERT INTO dag_edges ( + id, + snapshot_id, + parent_node_id, + child_node_id, + edge_index, + metadata_json + ) +VALUES ( + ?1, + ?2, + ?3, + ?4, + ?5, + ?6 + ) +` + +type InsertDAGEdgeParams struct { + ID ids.UUID `db:"id" json:"id"` + SnapshotID ids.UUID `db:"snapshot_id" json:"snapshot_id"` + ParentNodeID ids.UUID `db:"parent_node_id" json:"parent_node_id"` + ChildNodeID ids.UUID `db:"child_node_id" json:"child_node_id"` + EdgeIndex int64 `db:"edge_index" json:"edge_index"` + MetadataJson json.RawMessage `db:"metadata_json" json:"metadata_json"` +} + +// InsertDAGEdge +// +// INSERT INTO dag_edges ( +// id, +// snapshot_id, +// parent_node_id, +// child_node_id, +// edge_index, +// metadata_json +// ) +// VALUES ( +// ?1, +// ?2, +// ?3, +// ?4, +// ?5, +// ?6 +// ) +func (q *Queries) InsertDAGEdge(ctx context.Context, arg InsertDAGEdgeParams) error { + _, err := q.db.ExecContext(ctx, InsertDAGEdge, + arg.ID, + arg.SnapshotID, + arg.ParentNodeID, + arg.ChildNodeID, + arg.EdgeIndex, + arg.MetadataJson, + ) + return err +} + +const InsertDAGNode = `-- name: InsertDAGNode :one +INSERT INTO dag_nodes ( + id, + snapshot_id, + node_id, + level, + summary, + tokens, + start_idx, + end_idx, + span, + content_hash, + metrics_json, + metadata_json + ) +VALUES ( + ?1, + ?2, + ?3, + ?4, + ?5, + ?6, + ?7, + ?8, + ?9, + ?10, + ?11, + ?12 + ) +RETURNING id, + snapshot_id, + node_id, + level, + summary, + tokens, + start_idx, + end_idx, + span, + content_hash, + metrics_json, + metadata_json, + created_at, + updated_at +` + +type InsertDAGNodeParams struct { + ID ids.UUID `db:"id" json:"id"` + SnapshotID ids.UUID `db:"snapshot_id" json:"snapshot_id"` + NodeID string `db:"node_id" json:"node_id"` + Level int64 `db:"level" json:"level"` + Summary string `db:"summary" json:"summary"` + Tokens int64 `db:"tokens" json:"tokens"` + StartIdx int64 `db:"start_idx" json:"start_idx"` + EndIdx int64 `db:"end_idx" json:"end_idx"` + Span int64 `db:"span" json:"span"` + ContentHash string `db:"content_hash" json:"content_hash"` + MetricsJson json.RawMessage `db:"metrics_json" json:"metrics_json"` + MetadataJson json.RawMessage `db:"metadata_json" json:"metadata_json"` +} + +// InsertDAGNode +// +// INSERT INTO dag_nodes ( +// id, +// snapshot_id, +// node_id, +// level, +// summary, +// tokens, +// start_idx, +// end_idx, +// span, +// content_hash, +// metrics_json, +// metadata_json +// ) +// VALUES ( +// ?1, +// ?2, +// ?3, +// ?4, +// ?5, +// ?6, +// ?7, +// ?8, +// ?9, +// ?10, +// ?11, +// ?12 +// ) +// RETURNING id, +// snapshot_id, +// node_id, +// level, +// summary, +// tokens, +// start_idx, +// end_idx, +// span, +// content_hash, +// metrics_json, +// metadata_json, +// created_at, +// updated_at +func (q *Queries) InsertDAGNode(ctx context.Context, arg InsertDAGNodeParams) (DagNode, error) { + row := q.db.QueryRowContext(ctx, InsertDAGNode, + arg.ID, + arg.SnapshotID, + arg.NodeID, + arg.Level, + arg.Summary, + arg.Tokens, + arg.StartIdx, + arg.EndIdx, + arg.Span, + arg.ContentHash, + arg.MetricsJson, + arg.MetadataJson, + ) + var i DagNode + err := row.Scan( + &i.ID, + &i.SnapshotID, + &i.NodeID, + &i.Level, + &i.Summary, + &i.Tokens, + &i.StartIdx, + &i.EndIdx, + &i.Span, + &i.ContentHash, + &i.MetricsJson, + &i.MetadataJson, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} + +const InsertDAGSnapshot = `-- name: InsertDAGSnapshot :one +INSERT INTO dag_snapshots ( + id, + agent_id, + session_key, + from_msg_idx, + to_msg_idx, + msg_count, + roots_json, + content_hash + ) +VALUES ( + ?1, + ?2, + ?3, + ?4, + ?5, + ?6, + ?7, + ?8 + ) +RETURNING id, + agent_id, + session_key, + from_msg_idx, + to_msg_idx, + msg_count, + roots_json, + content_hash, + created_at, + updated_at +` + +type InsertDAGSnapshotParams struct { + ID ids.UUID `db:"id" json:"id"` + AgentID string `db:"agent_id" json:"agent_id"` + SessionKey string `db:"session_key" json:"session_key"` + FromMsgIdx int64 `db:"from_msg_idx" json:"from_msg_idx"` + ToMsgIdx int64 `db:"to_msg_idx" json:"to_msg_idx"` + MsgCount int64 `db:"msg_count" json:"msg_count"` + RootsJson string `db:"roots_json" json:"roots_json"` + ContentHash string `db:"content_hash" json:"content_hash"` +} + +// DAG persistence queries +// +// INSERT INTO dag_snapshots ( +// id, +// agent_id, +// session_key, +// from_msg_idx, +// to_msg_idx, +// msg_count, +// roots_json, +// content_hash +// ) +// VALUES ( +// ?1, +// ?2, +// ?3, +// ?4, +// ?5, +// ?6, +// ?7, +// ?8 +// ) +// RETURNING id, +// agent_id, +// session_key, +// from_msg_idx, +// to_msg_idx, +// msg_count, +// roots_json, +// content_hash, +// created_at, +// updated_at +func (q *Queries) InsertDAGSnapshot(ctx context.Context, arg InsertDAGSnapshotParams) (DagSnapshot, error) { + row := q.db.QueryRowContext(ctx, InsertDAGSnapshot, + arg.ID, + arg.AgentID, + arg.SessionKey, + arg.FromMsgIdx, + arg.ToMsgIdx, + arg.MsgCount, + arg.RootsJson, + arg.ContentHash, + ) + var i DagSnapshot + err := row.Scan( + &i.ID, + &i.AgentID, + &i.SessionKey, + &i.FromMsgIdx, + &i.ToMsgIdx, + &i.MsgCount, + &i.RootsJson, + &i.ContentHash, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} + +const ListDAGEdgesBySnapshotID = `-- name: ListDAGEdgesBySnapshotID :many +SELECT id, + snapshot_id, + parent_node_id, + child_node_id, + edge_index, + metadata_json, + created_at, + updated_at +FROM dag_edges +WHERE snapshot_id = ?1 +ORDER BY edge_index ASC, + created_at ASC +` + +type ListDAGEdgesBySnapshotIDParams struct { + SnapshotID ids.UUID `db:"snapshot_id" json:"snapshot_id"` +} + +// ListDAGEdgesBySnapshotID +// +// SELECT id, +// snapshot_id, +// parent_node_id, +// child_node_id, +// edge_index, +// metadata_json, +// created_at, +// updated_at +// FROM dag_edges +// WHERE snapshot_id = ?1 +// ORDER BY edge_index ASC, +// created_at ASC +func (q *Queries) ListDAGEdgesBySnapshotID(ctx context.Context, arg ListDAGEdgesBySnapshotIDParams) ([]DagEdge, error) { + rows, err := q.db.QueryContext(ctx, ListDAGEdgesBySnapshotID, arg.SnapshotID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []DagEdge{} + for rows.Next() { + var i DagEdge + if err := rows.Scan( + &i.ID, + &i.SnapshotID, + &i.ParentNodeID, + &i.ChildNodeID, + &i.EdgeIndex, + &i.MetadataJson, + &i.CreatedAt, + &i.UpdatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const ListDAGNodesBySnapshotID = `-- name: ListDAGNodesBySnapshotID :many +SELECT id, + snapshot_id, + node_id, + level, + summary, + tokens, + start_idx, + end_idx, + span, + content_hash, + metrics_json, + metadata_json, + created_at, + updated_at +FROM dag_nodes +WHERE snapshot_id = ?1 +ORDER BY start_idx ASC +` + +type ListDAGNodesBySnapshotIDParams struct { + SnapshotID ids.UUID `db:"snapshot_id" json:"snapshot_id"` +} + +// ListDAGNodesBySnapshotID +// +// SELECT id, +// snapshot_id, +// node_id, +// level, +// summary, +// tokens, +// start_idx, +// end_idx, +// span, +// content_hash, +// metrics_json, +// metadata_json, +// created_at, +// updated_at +// FROM dag_nodes +// WHERE snapshot_id = ?1 +// ORDER BY start_idx ASC +func (q *Queries) ListDAGNodesBySnapshotID(ctx context.Context, arg ListDAGNodesBySnapshotIDParams) ([]DagNode, error) { + rows, err := q.db.QueryContext(ctx, ListDAGNodesBySnapshotID, arg.SnapshotID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []DagNode{} + for rows.Next() { + var i DagNode + if err := rows.Scan( + &i.ID, + &i.SnapshotID, + &i.NodeID, + &i.Level, + &i.Summary, + &i.Tokens, + &i.StartIdx, + &i.EndIdx, + &i.Span, + &i.ContentHash, + &i.MetricsJson, + &i.MetadataJson, + &i.CreatedAt, + &i.UpdatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} diff --git a/pkg/memory/sqlc/map_ops.sql.go b/pkg/memory/sqlc/map_ops.sql.go new file mode 100644 index 000000000..292c27f65 --- /dev/null +++ b/pkg/memory/sqlc/map_ops.sql.go @@ -0,0 +1,815 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: map_ops.sql + +package sqlc + +import ( + "context" + "time" + + "github.com/ZanzyTHEbar/dragonscale/pkg/ids" +) + +const CountMapItemsByRun = `-- name: CountMapItemsByRun :one +SELECT count(*) AS count +FROM map_items +WHERE run_id = ?1 +` + +type CountMapItemsByRunParams struct { + RunID ids.UUID `db:"run_id" json:"run_id"` +} + +// CountMapItemsByRun +// +// SELECT count(*) AS count +// FROM map_items +// WHERE run_id = ?1 +func (q *Queries) CountMapItemsByRun(ctx context.Context, arg CountMapItemsByRunParams) (int64, error) { + row := q.db.QueryRowContext(ctx, CountMapItemsByRun, arg.RunID) + var count int64 + err := row.Scan(&count) + return count, err +} + +const CountMapItemsByRunAndStatus = `-- name: CountMapItemsByRunAndStatus :many +SELECT status, + count(*) AS count +FROM map_items +WHERE run_id = ?1 +GROUP BY status +` + +type CountMapItemsByRunAndStatusParams struct { + RunID ids.UUID `db:"run_id" json:"run_id"` +} + +type CountMapItemsByRunAndStatusRow struct { + Status string `db:"status" json:"status"` + Count int64 `db:"count" json:"count"` +} + +// CountMapItemsByRunAndStatus +// +// SELECT status, +// count(*) AS count +// FROM map_items +// WHERE run_id = ?1 +// GROUP BY status +func (q *Queries) CountMapItemsByRunAndStatus(ctx context.Context, arg CountMapItemsByRunAndStatusParams) ([]CountMapItemsByRunAndStatusRow, error) { + rows, err := q.db.QueryContext(ctx, CountMapItemsByRunAndStatus, arg.RunID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []CountMapItemsByRunAndStatusRow{} + for rows.Next() { + var i CountMapItemsByRunAndStatusRow + if err := rows.Scan(&i.Status, &i.Count); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const GetMapItemByID = `-- name: GetMapItemByID :one +SELECT id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +FROM map_items +WHERE id = ?1 +LIMIT 1 +` + +type GetMapItemByIDParams struct { + ID ids.UUID `db:"id" json:"id"` +} + +// GetMapItemByID +// +// SELECT id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +// FROM map_items +// WHERE id = ?1 +// LIMIT 1 +func (q *Queries) GetMapItemByID(ctx context.Context, arg GetMapItemByIDParams) (MapItem, error) { + row := q.db.QueryRowContext(ctx, GetMapItemByID, arg.ID) + var i MapItem + err := row.Scan( + &i.ID, + &i.RunID, + &i.ItemIndex, + &i.Status, + &i.Attempts, + &i.LastError, + &i.InputFb, + &i.OutputFb, + &i.InputHash, + &i.OutputHash, + &i.CreatedAt, + &i.UpdatedAt, + &i.CompletedAt, + ) + return i, err +} + +const GetMapItemByRunAndIndex = `-- name: GetMapItemByRunAndIndex :one +SELECT id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +FROM map_items +WHERE run_id = ?1 + AND item_index = ?2 +LIMIT 1 +` + +type GetMapItemByRunAndIndexParams struct { + RunID ids.UUID `db:"run_id" json:"run_id"` + ItemIndex int64 `db:"item_index" json:"item_index"` +} + +// GetMapItemByRunAndIndex +// +// SELECT id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +// FROM map_items +// WHERE run_id = ?1 +// AND item_index = ?2 +// LIMIT 1 +func (q *Queries) GetMapItemByRunAndIndex(ctx context.Context, arg GetMapItemByRunAndIndexParams) (MapItem, error) { + row := q.db.QueryRowContext(ctx, GetMapItemByRunAndIndex, arg.RunID, arg.ItemIndex) + var i MapItem + err := row.Scan( + &i.ID, + &i.RunID, + &i.ItemIndex, + &i.Status, + &i.Attempts, + &i.LastError, + &i.InputFb, + &i.OutputFb, + &i.InputHash, + &i.OutputHash, + &i.CreatedAt, + &i.UpdatedAt, + &i.CompletedAt, + ) + return i, err +} + +const GetMapRunByID = `-- name: GetMapRunByID :one +SELECT id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at +FROM map_runs +WHERE id = ?1 +LIMIT 1 +` + +type GetMapRunByIDParams struct { + ID ids.UUID `db:"id" json:"id"` +} + +// GetMapRunByID +// +// SELECT id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at +// FROM map_runs +// WHERE id = ?1 +// LIMIT 1 +func (q *Queries) GetMapRunByID(ctx context.Context, arg GetMapRunByIDParams) (MapRun, error) { + row := q.db.QueryRowContext(ctx, GetMapRunByID, arg.ID) + var i MapRun + err := row.Scan( + &i.ID, + &i.AgentID, + &i.SessionKey, + &i.OperatorKind, + &i.IdempotencyKey, + &i.Status, + &i.TotalItems, + &i.QueuedItems, + &i.RunningItems, + &i.SucceededItems, + &i.FailedItems, + &i.SpecFb, + &i.LastError, + &i.CreatedAt, + &i.UpdatedAt, + &i.CompletedAt, + ) + return i, err +} + +const GetMapRunByIdempotencyKey = `-- name: GetMapRunByIdempotencyKey :one +SELECT id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at +FROM map_runs +WHERE agent_id = ?1 + AND session_key = ?2 + AND operator_kind = ?3 + AND idempotency_key = ?4 +LIMIT 1 +` + +type GetMapRunByIdempotencyKeyParams struct { + AgentID string `db:"agent_id" json:"agent_id"` + SessionKey string `db:"session_key" json:"session_key"` + OperatorKind string `db:"operator_kind" json:"operator_kind"` + IdempotencyKey *string `db:"idempotency_key" json:"idempotency_key"` +} + +// GetMapRunByIdempotencyKey +// +// SELECT id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at +// FROM map_runs +// WHERE agent_id = ?1 +// AND session_key = ?2 +// AND operator_kind = ?3 +// AND idempotency_key = ?4 +// LIMIT 1 +func (q *Queries) GetMapRunByIdempotencyKey(ctx context.Context, arg GetMapRunByIdempotencyKeyParams) (MapRun, error) { + row := q.db.QueryRowContext(ctx, GetMapRunByIdempotencyKey, + arg.AgentID, + arg.SessionKey, + arg.OperatorKind, + arg.IdempotencyKey, + ) + var i MapRun + err := row.Scan( + &i.ID, + &i.AgentID, + &i.SessionKey, + &i.OperatorKind, + &i.IdempotencyKey, + &i.Status, + &i.TotalItems, + &i.QueuedItems, + &i.RunningItems, + &i.SucceededItems, + &i.FailedItems, + &i.SpecFb, + &i.LastError, + &i.CreatedAt, + &i.UpdatedAt, + &i.CompletedAt, + ) + return i, err +} + +const InsertMapItem = `-- name: InsertMapItem :one +INSERT INTO map_items ( + id, + run_id, + item_index, + status, + attempts, + last_error, + input_fb, + output_fb, + input_hash, + output_hash + ) +VALUES ( + ?1, + ?2, + ?3, + ?4, + ?5, + ?6, + ?7, + ?8, + ?9, + ?10 + ) +RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +` + +type InsertMapItemParams struct { + ID ids.UUID `db:"id" json:"id"` + RunID ids.UUID `db:"run_id" json:"run_id"` + ItemIndex int64 `db:"item_index" json:"item_index"` + Status string `db:"status" json:"status"` + Attempts int64 `db:"attempts" json:"attempts"` + LastError *string `db:"last_error" json:"last_error"` + InputFb []byte `db:"input_fb" json:"input_fb"` + OutputFb []byte `db:"output_fb" json:"output_fb"` + InputHash *string `db:"input_hash" json:"input_hash"` + OutputHash *string `db:"output_hash" json:"output_hash"` +} + +// InsertMapItem +// +// INSERT INTO map_items ( +// id, +// run_id, +// item_index, +// status, +// attempts, +// last_error, +// input_fb, +// output_fb, +// input_hash, +// output_hash +// ) +// VALUES ( +// ?1, +// ?2, +// ?3, +// ?4, +// ?5, +// ?6, +// ?7, +// ?8, +// ?9, +// ?10 +// ) +// RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +func (q *Queries) InsertMapItem(ctx context.Context, arg InsertMapItemParams) (MapItem, error) { + row := q.db.QueryRowContext(ctx, InsertMapItem, + arg.ID, + arg.RunID, + arg.ItemIndex, + arg.Status, + arg.Attempts, + arg.LastError, + arg.InputFb, + arg.OutputFb, + arg.InputHash, + arg.OutputHash, + ) + var i MapItem + err := row.Scan( + &i.ID, + &i.RunID, + &i.ItemIndex, + &i.Status, + &i.Attempts, + &i.LastError, + &i.InputFb, + &i.OutputFb, + &i.InputHash, + &i.OutputHash, + &i.CreatedAt, + &i.UpdatedAt, + &i.CompletedAt, + ) + return i, err +} + +const InsertMapRun = `-- name: InsertMapRun :one +INSERT INTO map_runs ( + id, + agent_id, + session_key, + operator_kind, + idempotency_key, + status, + total_items, + queued_items, + running_items, + succeeded_items, + failed_items, + spec_fb, + last_error + ) +VALUES ( + ?1, + ?2, + ?3, + ?4, + ?5, + ?6, + ?7, + ?8, + ?9, + ?10, + ?11, + ?12, + ?13 + ) +RETURNING id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at +` + +type InsertMapRunParams struct { + ID ids.UUID `db:"id" json:"id"` + AgentID string `db:"agent_id" json:"agent_id"` + SessionKey string `db:"session_key" json:"session_key"` + OperatorKind string `db:"operator_kind" json:"operator_kind"` + IdempotencyKey *string `db:"idempotency_key" json:"idempotency_key"` + Status string `db:"status" json:"status"` + TotalItems int64 `db:"total_items" json:"total_items"` + QueuedItems int64 `db:"queued_items" json:"queued_items"` + RunningItems int64 `db:"running_items" json:"running_items"` + SucceededItems int64 `db:"succeeded_items" json:"succeeded_items"` + FailedItems int64 `db:"failed_items" json:"failed_items"` + SpecFb []byte `db:"spec_fb" json:"spec_fb"` + LastError *string `db:"last_error" json:"last_error"` +} + +// InsertMapRun +// +// INSERT INTO map_runs ( +// id, +// agent_id, +// session_key, +// operator_kind, +// idempotency_key, +// status, +// total_items, +// queued_items, +// running_items, +// succeeded_items, +// failed_items, +// spec_fb, +// last_error +// ) +// VALUES ( +// ?1, +// ?2, +// ?3, +// ?4, +// ?5, +// ?6, +// ?7, +// ?8, +// ?9, +// ?10, +// ?11, +// ?12, +// ?13 +// ) +// RETURNING id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at +func (q *Queries) InsertMapRun(ctx context.Context, arg InsertMapRunParams) (MapRun, error) { + row := q.db.QueryRowContext(ctx, InsertMapRun, + arg.ID, + arg.AgentID, + arg.SessionKey, + arg.OperatorKind, + arg.IdempotencyKey, + arg.Status, + arg.TotalItems, + arg.QueuedItems, + arg.RunningItems, + arg.SucceededItems, + arg.FailedItems, + arg.SpecFb, + arg.LastError, + ) + var i MapRun + err := row.Scan( + &i.ID, + &i.AgentID, + &i.SessionKey, + &i.OperatorKind, + &i.IdempotencyKey, + &i.Status, + &i.TotalItems, + &i.QueuedItems, + &i.RunningItems, + &i.SucceededItems, + &i.FailedItems, + &i.SpecFb, + &i.LastError, + &i.CreatedAt, + &i.UpdatedAt, + &i.CompletedAt, + ) + return i, err +} + +const ListMapItemsByRunPaged = `-- name: ListMapItemsByRunPaged :many +SELECT id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +FROM map_items +WHERE run_id = ?1 +ORDER BY item_index ASC +LIMIT ?3 OFFSET ?2 +` + +type ListMapItemsByRunPagedParams struct { + RunID ids.UUID `db:"run_id" json:"run_id"` + Off int64 `db:"off" json:"off"` + Lim int64 `db:"lim" json:"lim"` +} + +// ListMapItemsByRunPaged +// +// SELECT id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +// FROM map_items +// WHERE run_id = ?1 +// ORDER BY item_index ASC +// LIMIT ?3 OFFSET ?2 +func (q *Queries) ListMapItemsByRunPaged(ctx context.Context, arg ListMapItemsByRunPagedParams) ([]MapItem, error) { + rows, err := q.db.QueryContext(ctx, ListMapItemsByRunPaged, arg.RunID, arg.Off, arg.Lim) + if err != nil { + return nil, err + } + defer rows.Close() + items := []MapItem{} + for rows.Next() { + var i MapItem + if err := rows.Scan( + &i.ID, + &i.RunID, + &i.ItemIndex, + &i.Status, + &i.Attempts, + &i.LastError, + &i.InputFb, + &i.OutputFb, + &i.InputHash, + &i.OutputHash, + &i.CreatedAt, + &i.UpdatedAt, + &i.CompletedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const ListMapRunsBySession = `-- name: ListMapRunsBySession :many +SELECT id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at +FROM map_runs +WHERE agent_id = ?1 + AND session_key = ?2 +ORDER BY created_at DESC +LIMIT ?4 OFFSET ?3 +` + +type ListMapRunsBySessionParams struct { + AgentID string `db:"agent_id" json:"agent_id"` + SessionKey string `db:"session_key" json:"session_key"` + Off int64 `db:"off" json:"off"` + Lim int64 `db:"lim" json:"lim"` +} + +// ListMapRunsBySession +// +// SELECT id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at +// FROM map_runs +// WHERE agent_id = ?1 +// AND session_key = ?2 +// ORDER BY created_at DESC +// LIMIT ?4 OFFSET ?3 +func (q *Queries) ListMapRunsBySession(ctx context.Context, arg ListMapRunsBySessionParams) ([]MapRun, error) { + rows, err := q.db.QueryContext(ctx, ListMapRunsBySession, + arg.AgentID, + arg.SessionKey, + arg.Off, + arg.Lim, + ) + if err != nil { + return nil, err + } + defer rows.Close() + items := []MapRun{} + for rows.Next() { + var i MapRun + if err := rows.Scan( + &i.ID, + &i.AgentID, + &i.SessionKey, + &i.OperatorKind, + &i.IdempotencyKey, + &i.Status, + &i.TotalItems, + &i.QueuedItems, + &i.RunningItems, + &i.SucceededItems, + &i.FailedItems, + &i.SpecFb, + &i.LastError, + &i.CreatedAt, + &i.UpdatedAt, + &i.CompletedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const MarkMapItemFailed = `-- name: MarkMapItemFailed :one +UPDATE map_items +SET status = 'failed', + last_error = ?1, + updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'), + completed_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') +WHERE id = ?2 +RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +` + +type MarkMapItemFailedParams struct { + LastError *string `db:"last_error" json:"last_error"` + ID ids.UUID `db:"id" json:"id"` +} + +// MarkMapItemFailed +// +// UPDATE map_items +// SET status = 'failed', +// last_error = ?1, +// updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'), +// completed_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') +// WHERE id = ?2 +// RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +func (q *Queries) MarkMapItemFailed(ctx context.Context, arg MarkMapItemFailedParams) (MapItem, error) { + row := q.db.QueryRowContext(ctx, MarkMapItemFailed, arg.LastError, arg.ID) + var i MapItem + err := row.Scan( + &i.ID, + &i.RunID, + &i.ItemIndex, + &i.Status, + &i.Attempts, + &i.LastError, + &i.InputFb, + &i.OutputFb, + &i.InputHash, + &i.OutputHash, + &i.CreatedAt, + &i.UpdatedAt, + &i.CompletedAt, + ) + return i, err +} + +const MarkMapItemRunning = `-- name: MarkMapItemRunning :one +UPDATE map_items +SET status = 'running', + attempts = attempts + 1, + last_error = NULL, + updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') +WHERE id = ?1 +RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +` + +type MarkMapItemRunningParams struct { + ID ids.UUID `db:"id" json:"id"` +} + +// MarkMapItemRunning +// +// UPDATE map_items +// SET status = 'running', +// attempts = attempts + 1, +// last_error = NULL, +// updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') +// WHERE id = ?1 +// RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +func (q *Queries) MarkMapItemRunning(ctx context.Context, arg MarkMapItemRunningParams) (MapItem, error) { + row := q.db.QueryRowContext(ctx, MarkMapItemRunning, arg.ID) + var i MapItem + err := row.Scan( + &i.ID, + &i.RunID, + &i.ItemIndex, + &i.Status, + &i.Attempts, + &i.LastError, + &i.InputFb, + &i.OutputFb, + &i.InputHash, + &i.OutputHash, + &i.CreatedAt, + &i.UpdatedAt, + &i.CompletedAt, + ) + return i, err +} + +const MarkMapItemSucceeded = `-- name: MarkMapItemSucceeded :one +UPDATE map_items +SET status = 'succeeded', + output_fb = ?1, + output_hash = ?2, + last_error = NULL, + updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'), + completed_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') +WHERE id = ?3 +RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +` + +type MarkMapItemSucceededParams struct { + OutputFb []byte `db:"output_fb" json:"output_fb"` + OutputHash *string `db:"output_hash" json:"output_hash"` + ID ids.UUID `db:"id" json:"id"` +} + +// MarkMapItemSucceeded +// +// UPDATE map_items +// SET status = 'succeeded', +// output_fb = ?1, +// output_hash = ?2, +// last_error = NULL, +// updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'), +// completed_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') +// WHERE id = ?3 +// RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at +func (q *Queries) MarkMapItemSucceeded(ctx context.Context, arg MarkMapItemSucceededParams) (MapItem, error) { + row := q.db.QueryRowContext(ctx, MarkMapItemSucceeded, arg.OutputFb, arg.OutputHash, arg.ID) + var i MapItem + err := row.Scan( + &i.ID, + &i.RunID, + &i.ItemIndex, + &i.Status, + &i.Attempts, + &i.LastError, + &i.InputFb, + &i.OutputFb, + &i.InputHash, + &i.OutputHash, + &i.CreatedAt, + &i.UpdatedAt, + &i.CompletedAt, + ) + return i, err +} + +const UpdateMapRunProgress = `-- name: UpdateMapRunProgress :one +UPDATE map_runs +SET status = ?1, + queued_items = ?2, + running_items = ?3, + succeeded_items = ?4, + failed_items = ?5, + last_error = ?6, + updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'), + completed_at = ?7 +WHERE id = ?8 +RETURNING id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at +` + +type UpdateMapRunProgressParams struct { + Status string `db:"status" json:"status"` + QueuedItems int64 `db:"queued_items" json:"queued_items"` + RunningItems int64 `db:"running_items" json:"running_items"` + SucceededItems int64 `db:"succeeded_items" json:"succeeded_items"` + FailedItems int64 `db:"failed_items" json:"failed_items"` + LastError *string `db:"last_error" json:"last_error"` + CompletedAt *time.Time `db:"completed_at" json:"completed_at"` + ID ids.UUID `db:"id" json:"id"` +} + +// UpdateMapRunProgress +// +// UPDATE map_runs +// SET status = ?1, +// queued_items = ?2, +// running_items = ?3, +// succeeded_items = ?4, +// failed_items = ?5, +// last_error = ?6, +// updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'), +// completed_at = ?7 +// WHERE id = ?8 +// RETURNING id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at +func (q *Queries) UpdateMapRunProgress(ctx context.Context, arg UpdateMapRunProgressParams) (MapRun, error) { + row := q.db.QueryRowContext(ctx, UpdateMapRunProgress, + arg.Status, + arg.QueuedItems, + arg.RunningItems, + arg.SucceededItems, + arg.FailedItems, + arg.LastError, + arg.CompletedAt, + arg.ID, + ) + var i MapRun + err := row.Scan( + &i.ID, + &i.AgentID, + &i.SessionKey, + &i.OperatorKind, + &i.IdempotencyKey, + &i.Status, + &i.TotalItems, + &i.QueuedItems, + &i.RunningItems, + &i.SucceededItems, + &i.FailedItems, + &i.SpecFb, + &i.LastError, + &i.CreatedAt, + &i.UpdatedAt, + &i.CompletedAt, + ) + return i, err +} diff --git a/pkg/memory/sqlc/queries/dag.sql b/pkg/memory/sqlc/queries/dag.sql new file mode 100644 index 000000000..998b3975b --- /dev/null +++ b/pkg/memory/sqlc/queries/dag.sql @@ -0,0 +1,172 @@ +-- DAG persistence queries +-- name: InsertDAGSnapshot :one +INSERT INTO dag_snapshots ( + id, + agent_id, + session_key, + from_msg_idx, + to_msg_idx, + msg_count, + roots_json, + content_hash + ) +VALUES ( + sqlc.arg(id), + sqlc.arg(agent_id), + sqlc.arg(session_key), + sqlc.arg(from_msg_idx), + sqlc.arg(to_msg_idx), + sqlc.arg(msg_count), + sqlc.arg(roots_json), + sqlc.arg(content_hash) + ) +RETURNING id, + agent_id, + session_key, + from_msg_idx, + to_msg_idx, + msg_count, + roots_json, + content_hash, + created_at, + updated_at; +-- name: InsertDAGNode :one +INSERT INTO dag_nodes ( + id, + snapshot_id, + node_id, + level, + summary, + tokens, + start_idx, + end_idx, + span, + content_hash, + metrics_json, + metadata_json + ) +VALUES ( + sqlc.arg(id), + sqlc.arg(snapshot_id), + sqlc.arg(node_id), + sqlc.arg(level), + sqlc.arg(summary), + sqlc.arg(tokens), + sqlc.arg(start_idx), + sqlc.arg(end_idx), + sqlc.arg(span), + sqlc.arg(content_hash), + sqlc.arg(metrics_json), + sqlc.arg(metadata_json) + ) +RETURNING id, + snapshot_id, + node_id, + level, + summary, + tokens, + start_idx, + end_idx, + span, + content_hash, + metrics_json, + metadata_json, + created_at, + updated_at; +-- name: InsertDAGEdge :exec +INSERT INTO dag_edges ( + id, + snapshot_id, + parent_node_id, + child_node_id, + edge_index, + metadata_json + ) +VALUES ( + sqlc.arg(id), + sqlc.arg(snapshot_id), + sqlc.arg(parent_node_id), + sqlc.arg(child_node_id), + sqlc.arg(edge_index), + sqlc.arg(metadata_json) + ); +-- name: GetLatestDAGSnapshotBySession :one +SELECT id, + agent_id, + session_key, + from_msg_idx, + to_msg_idx, + msg_count, + roots_json, + content_hash, + created_at, + updated_at +FROM dag_snapshots +WHERE agent_id = sqlc.arg(agent_id) + AND session_key = sqlc.arg(session_key) +ORDER BY created_at DESC +LIMIT 1; +-- name: GetDAGSnapshotByID :one +SELECT id, + agent_id, + session_key, + from_msg_idx, + to_msg_idx, + msg_count, + roots_json, + content_hash, + created_at, + updated_at +FROM dag_snapshots +WHERE id = sqlc.arg(id) +LIMIT 1; +-- name: ListDAGNodesBySnapshotID :many +SELECT id, + snapshot_id, + node_id, + level, + summary, + tokens, + start_idx, + end_idx, + span, + content_hash, + metrics_json, + metadata_json, + created_at, + updated_at +FROM dag_nodes +WHERE snapshot_id = sqlc.arg(snapshot_id) +ORDER BY start_idx ASC; +-- name: GetDAGNodeBySnapshotAndNodeID :one +SELECT id, + snapshot_id, + node_id, + level, + summary, + tokens, + start_idx, + end_idx, + span, + content_hash, + metrics_json, + metadata_json, + created_at, + updated_at +FROM dag_nodes +WHERE snapshot_id = sqlc.arg(snapshot_id) + AND node_id = sqlc.arg(node_id) +LIMIT 1; +-- name: ListDAGEdgesBySnapshotID :many +SELECT id, + snapshot_id, + parent_node_id, + child_node_id, + edge_index, + metadata_json, + created_at, + updated_at +FROM dag_edges +WHERE snapshot_id = sqlc.arg(snapshot_id) +ORDER BY edge_index ASC, + created_at ASC; \ No newline at end of file diff --git a/pkg/memory/sqlc/queries/map_ops.sql b/pkg/memory/sqlc/queries/map_ops.sql new file mode 100644 index 000000000..024fd04a2 --- /dev/null +++ b/pkg/memory/sqlc/queries/map_ops.sql @@ -0,0 +1,143 @@ +-- name: InsertMapRun :one +INSERT INTO map_runs ( + id, + agent_id, + session_key, + operator_kind, + idempotency_key, + status, + total_items, + queued_items, + running_items, + succeeded_items, + failed_items, + spec_fb, + last_error + ) +VALUES ( + sqlc.arg(id), + sqlc.arg(agent_id), + sqlc.arg(session_key), + sqlc.arg(operator_kind), + sqlc.arg(idempotency_key), + sqlc.arg(status), + sqlc.arg(total_items), + sqlc.arg(queued_items), + sqlc.arg(running_items), + sqlc.arg(succeeded_items), + sqlc.arg(failed_items), + sqlc.arg(spec_fb), + sqlc.arg(last_error) + ) +RETURNING *; +-- name: GetMapRunByIdempotencyKey :one +SELECT * +FROM map_runs +WHERE agent_id = sqlc.arg(agent_id) + AND session_key = sqlc.arg(session_key) + AND operator_kind = sqlc.arg(operator_kind) + AND idempotency_key = sqlc.arg(idempotency_key) +LIMIT 1; +-- name: GetMapRunByID :one +SELECT * +FROM map_runs +WHERE id = sqlc.arg(id) +LIMIT 1; +-- name: ListMapRunsBySession :many +SELECT * +FROM map_runs +WHERE agent_id = sqlc.arg(agent_id) + AND session_key = sqlc.arg(session_key) +ORDER BY created_at DESC +LIMIT sqlc.arg(lim) OFFSET sqlc.arg(off); +-- name: UpdateMapRunProgress :one +UPDATE map_runs +SET status = sqlc.arg(status), + queued_items = sqlc.arg(queued_items), + running_items = sqlc.arg(running_items), + succeeded_items = sqlc.arg(succeeded_items), + failed_items = sqlc.arg(failed_items), + last_error = sqlc.arg(last_error), + updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'), + completed_at = sqlc.arg(completed_at) +WHERE id = sqlc.arg(id) +RETURNING *; +-- name: InsertMapItem :one +INSERT INTO map_items ( + id, + run_id, + item_index, + status, + attempts, + last_error, + input_fb, + output_fb, + input_hash, + output_hash + ) +VALUES ( + sqlc.arg(id), + sqlc.arg(run_id), + sqlc.arg(item_index), + sqlc.arg(status), + sqlc.arg(attempts), + sqlc.arg(last_error), + sqlc.arg(input_fb), + sqlc.arg(output_fb), + sqlc.arg(input_hash), + sqlc.arg(output_hash) + ) +RETURNING *; +-- name: GetMapItemByID :one +SELECT * +FROM map_items +WHERE id = sqlc.arg(id) +LIMIT 1; +-- name: GetMapItemByRunAndIndex :one +SELECT * +FROM map_items +WHERE run_id = sqlc.arg(run_id) + AND item_index = sqlc.arg(item_index) +LIMIT 1; +-- name: ListMapItemsByRunPaged :many +SELECT * +FROM map_items +WHERE run_id = sqlc.arg(run_id) +ORDER BY item_index ASC +LIMIT sqlc.arg(lim) OFFSET sqlc.arg(off); +-- name: MarkMapItemRunning :one +UPDATE map_items +SET status = 'running', + attempts = attempts + 1, + last_error = NULL, + updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') +WHERE id = sqlc.arg(id) +RETURNING *; +-- name: MarkMapItemSucceeded :one +UPDATE map_items +SET status = 'succeeded', + output_fb = sqlc.arg(output_fb), + output_hash = sqlc.arg(output_hash), + last_error = NULL, + updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'), + completed_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') +WHERE id = sqlc.arg(id) +RETURNING *; +-- name: MarkMapItemFailed :one +UPDATE map_items +SET status = 'failed', + last_error = sqlc.arg(last_error), + updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'), + completed_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') +WHERE id = sqlc.arg(id) +RETURNING *; +-- name: CountMapItemsByRun :one +SELECT count(*) AS count +FROM map_items +WHERE run_id = sqlc.arg(run_id); +-- name: CountMapItemsByRunAndStatus :many +SELECT status, + count(*) AS count +FROM map_items +WHERE run_id = sqlc.arg(run_id) +GROUP BY status; \ No newline at end of file diff --git a/pkg/memory/store/retrieval_policy.go b/pkg/memory/store/retrieval_policy.go new file mode 100644 index 000000000..b5e79b8da --- /dev/null +++ b/pkg/memory/store/retrieval_policy.go @@ -0,0 +1,404 @@ +package store + +import ( + "context" + "encoding/json" + "fmt" + "time" + + "github.com/ZanzyTHEbar/dragonscale/pkg/ids" + "github.com/ZanzyTHEbar/dragonscale/pkg/logger" + "github.com/ZanzyTHEbar/dragonscale/pkg/memory" +) + +type retrievalMode string + +const ( + retrievalModeShadow retrievalMode = "shadow" + retrievalModePromoted retrievalMode = "promoted" + retrievalModeRollback retrievalMode = "rollback" +) + +const ( + retrievalPolicyStateKey = "memory:retrieval:policy_state" + retrievalPolicyGatesKey = "memory:retrieval:promotion_gates" + retrievalPolicyMetricsKey = "memory:retrieval:shadow_metrics" +) + +type retrievalPolicyState struct { + Mode retrievalMode `json:"mode"` + Reason string `json:"reason,omitempty"` + UpdatedAt time.Time `json:"updated_at"` +} + +type retrievalPromotionGates struct { + MinSamples int `json:"min_samples"` + MinAugmentedSamples int `json:"min_augmented_samples"` + MinTop1Parity float64 `json:"min_top1_parity"` + MinOverlapAtK float64 `json:"min_overlap_at_k"` + MinPromotedSamples int `json:"min_promoted_samples"` + RollbackTop1Parity float64 `json:"rollback_top1_parity"` + RollbackOverlapAtK float64 `json:"rollback_overlap_at_k"` + MaxNoResultRate float64 `json:"max_no_result_rate"` + PromotedNoResultRate float64 `json:"promoted_no_result_rate"` +} + +type retrievalShadowMetrics struct { + TotalQueries int `json:"total_queries"` + NoResultQueries int `json:"no_result_queries"` + AugmentedQueries int `json:"augmented_queries"` + Top1ParityHits int `json:"top1_parity_hits"` + OverlapAtKTotal float64 `json:"overlap_at_k_total"` + PromotedQueries int `json:"promoted_queries"` + PromotedTop1Hits int `json:"promoted_top1_hits"` + PromotedOverlap float64 `json:"promoted_overlap_total"` + PromotedNoResultQueries int `json:"promoted_no_result_queries"` + UpdatedAt time.Time `json:"updated_at"` +} + +type retrievalParity struct { + Top1Match bool + OverlapAtK float64 + NoResultPair bool +} + +func defaultRetrievalPolicyState() retrievalPolicyState { + return retrievalPolicyState{ + Mode: retrievalModeShadow, + Reason: "default_shadow_bootstrap", + UpdatedAt: time.Now().UTC(), + } +} + +func defaultRetrievalPromotionGates() retrievalPromotionGates { + return retrievalPromotionGates{ + MinSamples: 25, + MinAugmentedSamples: 10, + MinTop1Parity: 0.65, + MinOverlapAtK: 0.60, + MinPromotedSamples: 10, + RollbackTop1Parity: 0.45, + RollbackOverlapAtK: 0.35, + MaxNoResultRate: 0.90, + PromotedNoResultRate: 0.95, + } +} + +func defaultRetrievalShadowMetrics() retrievalShadowMetrics { + return retrievalShadowMetrics{ + UpdatedAt: time.Now().UTC(), + } +} + +func (m *MemoryStore) loadRetrievalPolicyState(ctx context.Context) retrievalPolicyState { + raw, err := m.delegate.GetKV(ctx, m.agentID, retrievalPolicyStateKey) + if err != nil || raw == "" { + state := defaultRetrievalPolicyState() + if perr := m.persistRetrievalPolicyState(ctx, state); perr != nil { + logger.WarnCF("memory", "failed to persist default retrieval policy state", + map[string]interface{}{"error": perr.Error()}) + } + return state + } + + var state retrievalPolicyState + if err := json.Unmarshal([]byte(raw), &state); err != nil { + state = defaultRetrievalPolicyState() + if perr := m.persistRetrievalPolicyState(ctx, state); perr != nil { + logger.WarnCF("memory", "failed to persist repaired retrieval policy state", + map[string]interface{}{"error": perr.Error()}) + } + return state + } + if !isValidRetrievalMode(state.Mode) { + state = defaultRetrievalPolicyState() + if perr := m.persistRetrievalPolicyState(ctx, state); perr != nil { + logger.WarnCF("memory", "failed to persist healed retrieval policy state", + map[string]interface{}{"error": perr.Error()}) + } + } + return state +} + +func (m *MemoryStore) loadRetrievalPromotionGates(ctx context.Context) retrievalPromotionGates { + raw, err := m.delegate.GetKV(ctx, m.agentID, retrievalPolicyGatesKey) + if err != nil || raw == "" { + gates := defaultRetrievalPromotionGates() + if perr := m.persistRetrievalPromotionGates(ctx, gates); perr != nil { + logger.WarnCF("memory", "failed to persist default retrieval promotion gates", + map[string]interface{}{"error": perr.Error()}) + } + return gates + } + + var gates retrievalPromotionGates + if err := json.Unmarshal([]byte(raw), &gates); err != nil { + gates = defaultRetrievalPromotionGates() + if perr := m.persistRetrievalPromotionGates(ctx, gates); perr != nil { + logger.WarnCF("memory", "failed to persist repaired retrieval promotion gates", + map[string]interface{}{"error": perr.Error()}) + } + return gates + } + + // Guard-clauses for malformed/partial gate records. + if gates.MinSamples <= 0 { + gates.MinSamples = defaultRetrievalPromotionGates().MinSamples + } + if gates.MinPromotedSamples <= 0 { + gates.MinPromotedSamples = defaultRetrievalPromotionGates().MinPromotedSamples + } + if gates.MinAugmentedSamples <= 0 { + gates.MinAugmentedSamples = defaultRetrievalPromotionGates().MinAugmentedSamples + } + if gates.MinTop1Parity <= 0 { + gates.MinTop1Parity = defaultRetrievalPromotionGates().MinTop1Parity + } + if gates.MinOverlapAtK <= 0 { + gates.MinOverlapAtK = defaultRetrievalPromotionGates().MinOverlapAtK + } + if gates.RollbackTop1Parity <= 0 { + gates.RollbackTop1Parity = defaultRetrievalPromotionGates().RollbackTop1Parity + } + if gates.RollbackOverlapAtK <= 0 { + gates.RollbackOverlapAtK = defaultRetrievalPromotionGates().RollbackOverlapAtK + } + if gates.MaxNoResultRate <= 0 { + gates.MaxNoResultRate = defaultRetrievalPromotionGates().MaxNoResultRate + } + if gates.PromotedNoResultRate <= 0 { + gates.PromotedNoResultRate = defaultRetrievalPromotionGates().PromotedNoResultRate + } + return gates +} + +func (m *MemoryStore) loadRetrievalShadowMetrics(ctx context.Context) retrievalShadowMetrics { + raw, err := m.delegate.GetKV(ctx, m.agentID, retrievalPolicyMetricsKey) + if err != nil || raw == "" { + metrics := defaultRetrievalShadowMetrics() + if perr := m.persistRetrievalShadowMetrics(ctx, metrics); perr != nil { + logger.WarnCF("memory", "failed to persist default retrieval shadow metrics", + map[string]interface{}{"error": perr.Error()}) + } + return metrics + } + + var metrics retrievalShadowMetrics + if err := json.Unmarshal([]byte(raw), &metrics); err != nil { + metrics = defaultRetrievalShadowMetrics() + if perr := m.persistRetrievalShadowMetrics(ctx, metrics); perr != nil { + logger.WarnCF("memory", "failed to persist repaired retrieval shadow metrics", + map[string]interface{}{"error": perr.Error()}) + } + return metrics + } + return metrics +} + +func (m *MemoryStore) persistRetrievalPolicyState(ctx context.Context, state retrievalPolicyState) error { + payload, err := json.Marshal(state) + if err != nil { + return fmt.Errorf("marshal retrieval policy state: %w", err) + } + if err := m.delegate.UpsertKV(ctx, m.agentID, retrievalPolicyStateKey, string(payload)); err != nil { + return fmt.Errorf("upsert retrieval policy state: %w", err) + } + return nil +} + +func (m *MemoryStore) persistRetrievalPromotionGates(ctx context.Context, gates retrievalPromotionGates) error { + payload, err := json.Marshal(gates) + if err != nil { + return fmt.Errorf("marshal retrieval promotion gates: %w", err) + } + if err := m.delegate.UpsertKV(ctx, m.agentID, retrievalPolicyGatesKey, string(payload)); err != nil { + return fmt.Errorf("upsert retrieval promotion gates: %w", err) + } + return nil +} + +func (m *MemoryStore) persistRetrievalShadowMetrics(ctx context.Context, metrics retrievalShadowMetrics) error { + payload, err := json.Marshal(metrics) + if err != nil { + return fmt.Errorf("marshal retrieval shadow metrics: %w", err) + } + if err := m.delegate.UpsertKV(ctx, m.agentID, retrievalPolicyMetricsKey, string(payload)); err != nil { + return fmt.Errorf("upsert retrieval shadow metrics: %w", err) + } + return nil +} + +func computeRetrievalParity(baseline, augmented []memory.SearchResult, k int) retrievalParity { + if k <= 0 { + k = 5 + } + parity := retrievalParity{ + NoResultPair: len(baseline) == 0 && len(augmented) == 0, + } + if len(baseline) > 0 && len(augmented) > 0 { + parity.Top1Match = baseline[0].ID == augmented[0].ID + } + + baseTop := topKResultIDs(baseline, k) + augTop := topKResultIDs(augmented, k) + if len(baseTop) == 0 && len(augTop) == 0 { + parity.OverlapAtK = 1.0 + return parity + } + + union := make(map[ids.UUID]bool, len(baseTop)+len(augTop)) + intersection := 0 + for id := range baseTop { + union[id] = true + if augTop[id] { + intersection++ + } + } + for id := range augTop { + union[id] = true + } + if len(union) > 0 { + parity.OverlapAtK = float64(intersection) / float64(len(union)) + } + return parity +} + +func topKResultIDs(results []memory.SearchResult, k int) map[ids.UUID]bool { + if k <= 0 { + return nil + } + if len(results) < k { + k = len(results) + } + out := make(map[ids.UUID]bool, k) + for i := 0; i < k; i++ { + out[results[i].ID] = true + } + return out +} + +func (m *MemoryStore) updateRetrievalPolicy( + ctx context.Context, + state retrievalPolicyState, + gates retrievalPromotionGates, + metrics retrievalShadowMetrics, + parity retrievalParity, + augmentedUsed bool, + baseline []memory.SearchResult, +) retrievalPolicyState { + prevMode := state.Mode + + metrics.TotalQueries++ + if len(baseline) == 0 { + metrics.NoResultQueries++ + } + if augmentedUsed { + metrics.AugmentedQueries++ + } + if parity.Top1Match { + metrics.Top1ParityHits++ + } + metrics.OverlapAtKTotal += parity.OverlapAtK + metrics.UpdatedAt = time.Now().UTC() + + switch state.Mode { + case retrievalModeShadow: + top1Rate := safeRate(metrics.Top1ParityHits, metrics.TotalQueries) + overlapAvg := safeAvg(metrics.OverlapAtKTotal, metrics.TotalQueries) + noResultRate := safeRate(metrics.NoResultQueries, metrics.TotalQueries) + + if metrics.TotalQueries >= gates.MinSamples && + metrics.AugmentedQueries >= gates.MinAugmentedSamples && + top1Rate >= gates.MinTop1Parity && + overlapAvg >= gates.MinOverlapAtK && + noResultRate <= gates.MaxNoResultRate { + state.Mode = retrievalModePromoted + state.Reason = fmt.Sprintf( + "promotion_gates_passed(top1=%.3f overlap=%.3f samples=%d)", + top1Rate, + overlapAvg, + metrics.TotalQueries, + ) + state.UpdatedAt = time.Now().UTC() + metrics.PromotedQueries = 0 + metrics.PromotedTop1Hits = 0 + metrics.PromotedOverlap = 0 + metrics.PromotedNoResultQueries = 0 + } + + case retrievalModePromoted: + metrics.PromotedQueries++ + if len(baseline) == 0 { + metrics.PromotedNoResultQueries++ + } + if parity.Top1Match { + metrics.PromotedTop1Hits++ + } + metrics.PromotedOverlap += parity.OverlapAtK + + if metrics.PromotedQueries >= gates.MinPromotedSamples { + promotedTop1 := safeRate(metrics.PromotedTop1Hits, metrics.PromotedQueries) + promotedOverlap := safeAvg(metrics.PromotedOverlap, metrics.PromotedQueries) + promotedNoResultRate := safeRate(metrics.PromotedNoResultQueries, metrics.PromotedQueries) + + if promotedTop1 < gates.RollbackTop1Parity || + promotedOverlap < gates.RollbackOverlapAtK || + promotedNoResultRate > gates.PromotedNoResultRate { + state.Mode = retrievalModeRollback + state.Reason = fmt.Sprintf( + "rollback_triggered(top1=%.3f overlap=%.3f no_result_rate=%.3f promoted_samples=%d)", + promotedTop1, + promotedOverlap, + promotedNoResultRate, + metrics.PromotedQueries, + ) + state.UpdatedAt = time.Now().UTC() + } + } + } + + if prevMode != state.Mode { + logger.InfoCF("memory", "retrieval policy mode transition", + map[string]interface{}{ + "from": prevMode, + "to": state.Mode, + "reason": state.Reason, + "total_queries": metrics.TotalQueries, + "augmented_queries": metrics.AugmentedQueries, + "promoted_queries": metrics.PromotedQueries, + }) + } + + if err := m.persistRetrievalShadowMetrics(ctx, metrics); err != nil { + logger.WarnCF("memory", "failed to persist retrieval shadow metrics", + map[string]interface{}{"mode": state.Mode, "error": err.Error()}) + } + if err := m.persistRetrievalPolicyState(ctx, state); err != nil { + logger.WarnCF("memory", "failed to persist retrieval policy state", + map[string]interface{}{"mode": state.Mode, "error": err.Error()}) + } + return state +} + +func safeRate(numerator, denominator int) float64 { + if denominator <= 0 { + return 0 + } + return float64(numerator) / float64(denominator) +} + +func safeAvg(total float64, count int) float64 { + if count <= 0 { + return 0 + } + return total / float64(count) +} + +func isValidRetrievalMode(mode retrievalMode) bool { + switch mode { + case retrievalModeShadow, retrievalModePromoted, retrievalModeRollback: + return true + default: + return false + } +} diff --git a/pkg/memory/store/token_test_helpers_test.go b/pkg/memory/store/token_test_helpers_test.go new file mode 100644 index 000000000..259c072de --- /dev/null +++ b/pkg/memory/store/token_test_helpers_test.go @@ -0,0 +1,18 @@ +package store + +import ( + "fmt" + "strings" +) + +func contentAtLeastTokens(target int) string { + if target <= 0 { + return "" + } + + var b strings.Builder + for i := 0; estimateTokens(b.String()) < target; i++ { + b.WriteString(fmt.Sprintf("token-%d ", i)) + } + return b.String() +}