diff --git a/pkg/memory/delegate/rl_store_test.go b/pkg/memory/delegate/rl_store_test.go index 1085c8a62..f37ba145d 100644 --- a/pkg/memory/delegate/rl_store_test.go +++ b/pkg/memory/delegate/rl_store_test.go @@ -364,13 +364,13 @@ func TestSQLiteDelegate_GetCompletedTasks(t *testing.T) { d := setupRLTest(t) ctx := t.Context() - // This is a placeholder implementation that returns empty list - tasks, err := d.GetCompletedTasks(ctx, time.Time{}) + tasks, err := d.GetCompletedTasks(ctx, "test-agent", time.Time{}) if err != nil { t.Fatalf("GetCompletedTasks: %v", err) } + // Should return empty list when no tasks exist if len(tasks) != 0 { - t.Errorf("expected 0 tasks (placeholder), got %d", len(tasks)) + t.Errorf("expected 0 tasks, got %d", len(tasks)) } } @@ -379,13 +379,167 @@ func TestSQLiteDelegate_GetRetrievedMemories(t *testing.T) { d := setupRLTest(t) ctx := t.Context() - // This is a placeholder implementation that returns empty list - memories, err := d.GetRetrievedMemories(ctx, "task-123") + // Drop and recreate tables with relaxed constraints for testing + // (the real schema from Init() has FK constraints requiring valid conversations/runs) + _, _ = d.db.ExecContext(ctx, `DROP TABLE IF EXISTS task_retrievals`) + _, _ = d.db.ExecContext(ctx, `DROP TABLE IF EXISTS task_completions`) + + // Create test tables with relaxed constraints + // Use WITHOUT ROWID for BLOB PRIMARY KEY to match schema.sql + _, err := d.db.ExecContext(ctx, ` + CREATE TABLE task_completions ( + id BLOB PRIMARY KEY, + agent_id TEXT NOT NULL, + conversation_id BLOB NOT NULL DEFAULT (x'00000000000000000000000000000000'), + run_id BLOB NOT NULL DEFAULT (x'00000000000000000000000000000000'), + description TEXT, + tokens_used INTEGER DEFAULT 0, + tool_calls INTEGER DEFAULT 0, + errors INTEGER DEFAULT 0, + user_corrections INTEGER DEFAULT 0, + completed BOOLEAN NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')) + ) WITHOUT ROWID + `) + if err != nil { + t.Fatalf("create task_completions table: %v", err) + } + + _, err = d.db.ExecContext(ctx, ` + CREATE TABLE task_retrievals ( + id BLOB PRIMARY KEY, + task_id BLOB NOT NULL, + memory_id BLOB NOT NULL, + similarity REAL NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + UNIQUE(task_id, memory_id) + ) WITHOUT ROWID + `) + if err != nil { + t.Fatalf("create task_retrievals table: %v", err) + } + + // Create a valid task ID + taskID := ids.New() + + // Insert a task completion first (required for FK) + _, err = d.db.ExecContext(ctx, ` + INSERT INTO task_completions (id, agent_id, description, completed) + VALUES (?, ?, ?, ?) + `, taskID, "test-agent", "test task", true) + if err != nil { + t.Fatalf("insert task completion: %v", err) + } + + // Insert a recall item to retrieve + memoryID := insertTestRecallItem(ctx, t, d, "test-agent") + + // Store a task retrieval record + err = d.queries.StoreTaskRetrieval(ctx, memsqlc.StoreTaskRetrievalParams{ + ID: ids.New(), + TaskID: taskID, + MemoryID: memoryID, + Similarity: 0.95, + }) + if err != nil { + t.Fatalf("StoreTaskRetrieval: %v", err) + } + + // Update self-report score on the recall item + selfReportScore := int64(3) + err = d.queries.UpdateMemorySelfReportScore(ctx, memsqlc.UpdateMemorySelfReportScoreParams{ + SelfReportScore: &selfReportScore, + ID: memoryID, + AgentID: "test-agent", + }) + if err != nil { + t.Fatalf("UpdateMemorySelfReportScore: %v", err) + } + + // Retrieve memories for the task + memories, err := d.GetRetrievedMemories(ctx, taskID.String()) + if err != nil { + t.Fatalf("GetRetrievedMemories: %v", err) + } + if len(memories) != 1 { + t.Fatalf("expected 1 memory, got %d", len(memories)) + } + + // Verify the retrieved memory + if memories[0].MemoryID != memoryID { + t.Errorf("expected memory ID %s, got %s", memoryID, memories[0].MemoryID) + } + if memories[0].Similarity != 0.95 { + t.Errorf("expected similarity 0.95, got %f", memories[0].Similarity) + } + if memories[0].SelfReportScore == nil || *memories[0].SelfReportScore != 3 { + t.Errorf("expected self-report score 3, got %v", memories[0].SelfReportScore) + } +} + +func TestSQLiteDelegate_GetRetrievedMemories_InvalidTaskID(t *testing.T) { + t.Parallel() + d := setupRLTest(t) + ctx := t.Context() + + // Test with invalid task ID format + _, err := d.GetRetrievedMemories(ctx, "invalid-task-id") + if err == nil { + t.Error("expected error for invalid task ID, got nil") + } +} + +func TestSQLiteDelegate_GetRetrievedMemories_Empty(t *testing.T) { + t.Parallel() + d := setupRLTest(t) + ctx := t.Context() + + // Drop and recreate tables with relaxed constraints for testing + _, _ = d.db.ExecContext(ctx, `DROP TABLE IF EXISTS task_retrievals`) + _, _ = d.db.ExecContext(ctx, `DROP TABLE IF EXISTS task_completions`) + + // Create test tables with relaxed constraints + _, err := d.db.ExecContext(ctx, ` + CREATE TABLE task_completions ( + id BLOB PRIMARY KEY, + agent_id TEXT NOT NULL, + conversation_id BLOB NOT NULL DEFAULT (x'00000000000000000000000000000000'), + run_id BLOB NOT NULL DEFAULT (x'00000000000000000000000000000000'), + description TEXT, + tokens_used INTEGER DEFAULT 0, + tool_calls INTEGER DEFAULT 0, + errors INTEGER DEFAULT 0, + user_corrections INTEGER DEFAULT 0, + completed BOOLEAN NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')) + ) WITHOUT ROWID + `) + if err != nil { + t.Fatalf("create task_completions table: %v", err) + } + + _, err = d.db.ExecContext(ctx, ` + CREATE TABLE task_retrievals ( + id BLOB PRIMARY KEY, + task_id BLOB NOT NULL, + memory_id BLOB NOT NULL, + similarity REAL NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + UNIQUE(task_id, memory_id) + ) WITHOUT ROWID + `) + if err != nil { + t.Fatalf("create task_retrievals table: %v", err) + } + + // Test with valid UUID but no retrievals + taskID := ids.New() + memories, err := d.GetRetrievedMemories(ctx, taskID.String()) if err != nil { t.Fatalf("GetRetrievedMemories: %v", err) } if len(memories) != 0 { - t.Errorf("expected 0 memories (placeholder), got %d", len(memories)) + t.Errorf("expected 0 memories for task with no retrievals, got %d", len(memories)) } } diff --git a/pkg/memory/delegate/sqlite.go b/pkg/memory/delegate/sqlite.go index 1b3428a61..c09dad96f 100644 --- a/pkg/memory/delegate/sqlite.go +++ b/pkg/memory/delegate/sqlite.go @@ -1114,22 +1114,92 @@ func (d *LibSQLDelegate) UpdateMemorySelfReport(ctx context.Context, memoryID id }) } -// GetCompletedTasks returns tasks completed since the given time. +// GetCompletedTasks returns tasks completed since the given time for a specific agent. // Implements cortex.RLStore interface. -// Note: This is a placeholder implementation - actual task storage needs to be defined. -func (d *LibSQLDelegate) GetCompletedTasks(ctx context.Context, since time.Time) ([]TaskRecord, error) { - // TODO: Implement actual task retrieval from jobs or runs tables - // For now, return empty list - return []TaskRecord{}, nil +func (d *LibSQLDelegate) GetCompletedTasks(ctx context.Context, agentID string, since time.Time) ([]TaskRecord, error) { + rows, err := d.queries.GetCompletedTasks(ctx, memsqlc.GetCompletedTasksParams{ + AgentID: agentID, + Since: since, + }) + if err != nil { + logger.WarnCF("memory", "Failed to get completed tasks", map[string]interface{}{ + "agent_id": agentID, + "since": since.String(), + "error": err.Error(), + }) + return nil, fmt.Errorf("get completed tasks: %w", err) + } + + tasks := make([]TaskRecord, 0, len(rows)) + for _, task := range rows { + record := TaskRecord{ + ID: task.ID.String(), + Description: task.Description, + Completed: task.Completed, + CreatedAt: task.CreatedAt, + } + if task.TokensUsed != nil { + record.TokensUsed = int(*task.TokensUsed) + } + if task.ToolCalls != nil { + record.ToolCalls = int(*task.ToolCalls) + } + if task.Errors != nil { + record.Errors = int(*task.Errors) + } + if task.UserCorrections != nil { + record.UserCorrections = int(*task.UserCorrections) + } + tasks = append(tasks, record) + } + + return tasks, nil } // GetRetrievedMemories returns memories retrieved during a task. // Implements cortex.RLStore interface. -// Note: This is a placeholder implementation - actual retrieval tracking needs to be defined. func (d *LibSQLDelegate) GetRetrievedMemories(ctx context.Context, taskID string) ([]RetrievedMemoryRecord, error) { - // TODO: Implement actual retrieved memory tracking - // For now, return empty list - return []RetrievedMemoryRecord{}, nil + parsedID, err := ids.Parse(taskID) + if err != nil { + return nil, fmt.Errorf("invalid task id: %w", err) + } + + rows, err := d.queries.GetRetrievedMemories(ctx, memsqlc.GetRetrievedMemoriesParams{ + TaskID: parsedID, + }) + if err != nil { + return nil, err + } + + records := make([]RetrievedMemoryRecord, len(rows)) + for i, row := range rows { + records[i] = RetrievedMemoryRecord{ + MemoryID: row.MemoryID, + Similarity: row.Similarity, + } + if row.SelfReportScore != nil { + score := int(*row.SelfReportScore) + records[i].SelfReportScore = &score + } + } + + return records, nil +} + +// ListActiveAgents returns all agent IDs that have completed tasks since the given time. +// Implements cortex.RLStore interface for multi-agent support. +func (d *LibSQLDelegate) ListActiveAgents(ctx context.Context, since time.Time) ([]string, error) { + rows, err := d.queries.ListActiveAgents(ctx, memsqlc.ListActiveAgentsParams{ + Since: since, + }) + if err != nil { + logger.WarnCF("memory", "Failed to list active agents", map[string]interface{}{ + "since": since, + "error": err.Error(), + }) + return nil, fmt.Errorf("list active agents: %w", err) + } + return rows, nil } // --- Audit Analysis Store Methods --- diff --git a/pkg/memory/migrations/016_rl_schema.go b/pkg/memory/migrations/016_rl_schema.go index c759764d5..e2a8d8c8b 100644 --- a/pkg/memory/migrations/016_rl_schema.go +++ b/pkg/memory/migrations/016_rl_schema.go @@ -26,6 +26,33 @@ func up016RLSchema(ctx context.Context, tx *sql.Tx) error { m2_user_corrections REAL DEFAULT 0, updated_at DATETIME )`, + // Create task_completions table for recording completed agent runs + `CREATE TABLE IF NOT EXISTS task_completions ( + id BLOB PRIMARY KEY, + agent_id TEXT NOT NULL, + conversation_id BLOB NOT NULL REFERENCES agent_conversations(id) ON DELETE CASCADE, + run_id BLOB NOT NULL REFERENCES agent_runs(id) ON DELETE CASCADE, + description TEXT NOT NULL DEFAULT '', + tokens_used INTEGER DEFAULT 0, + tool_calls INTEGER DEFAULT 0, + errors INTEGER DEFAULT 0, + user_corrections INTEGER DEFAULT 0, + completed BOOLEAN NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')) + )`, + `CREATE INDEX IF NOT EXISTS idx_task_completions_agent_created ON task_completions(agent_id, created_at)`, + `CREATE INDEX IF NOT EXISTS idx_task_completions_run ON task_completions(run_id)`, + // Create task_retrievals table for RL credit assignment tracking + `CREATE TABLE IF NOT EXISTS task_retrievals ( + id BLOB PRIMARY KEY, + task_id BLOB NOT NULL REFERENCES task_completions(id) ON DELETE CASCADE, + memory_id BLOB NOT NULL REFERENCES recall_items(id) ON DELETE CASCADE, + similarity REAL NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + UNIQUE(task_id, memory_id) + )`, + `CREATE INDEX IF NOT EXISTS idx_task_retrievals_task ON task_retrievals(task_id)`, + `CREATE INDEX IF NOT EXISTS idx_task_retrievals_memory ON task_retrievals(memory_id)`, // Add RL-related columns to recall_items `ALTER TABLE recall_items ADD COLUMN rl_weight REAL DEFAULT 1.0`, `ALTER TABLE recall_items ADD COLUMN rl_credit REAL`, @@ -47,9 +74,15 @@ func up016RLSchema(ctx context.Context, tx *sql.Tx) error { func down016RLSchema(ctx context.Context, tx *sql.Tx) error { stmts := []string{ // Drop indexes first + `DROP INDEX IF EXISTS idx_task_retrievals_memory`, + `DROP INDEX IF EXISTS idx_task_retrievals_task`, + `DROP INDEX IF EXISTS idx_task_completions_run`, + `DROP INDEX IF EXISTS idx_task_completions_agent_created`, `DROP INDEX IF EXISTS idx_recall_task_retrieval`, `DROP INDEX IF EXISTS idx_recall_rl_weight`, - // Drop task_baselines table + // Drop tables in reverse order of creation (respect foreign keys) + `DROP TABLE IF EXISTS task_retrievals`, + `DROP TABLE IF EXISTS task_completions`, `DROP TABLE IF EXISTS task_baselines`, // Note: SQLite doesn't support DROP COLUMN directly // The columns (rl_weight, rl_credit, self_report_score, task_retrieval_count) diff --git a/pkg/memory/sqlc/querier.go b/pkg/memory/sqlc/querier.go index 15bc34d90..3dc226b88 100644 --- a/pkg/memory/sqlc/querier.go +++ b/pkg/memory/sqlc/querier.go @@ -504,6 +504,21 @@ type Querier interface { // AND name = ?2 // LIMIT 1 GetDocument(ctx context.Context, arg GetDocumentParams) (AgentDocument, error) + // Get sessions with high token usage grouped by conversation/agent + // + // SELECT + // conversation_id as session_id, + // agent_id, + // SUM(tokens_used) as total_tokens, + // COUNT(*) as task_count + // FROM task_completions + // WHERE tokens_used > ?1 + // AND created_at > datetime('now', '-24 hours') + // GROUP BY conversation_id, agent_id + // HAVING total_tokens > ?1 + // ORDER BY total_tokens DESC + // LIMIT ?2 + GetHighTokenSessions(ctx context.Context, arg GetHighTokenSessionsParams) ([]GetHighTokenSessionsRow, error) //GetImmutableMessage // // SELECT id, @@ -1092,6 +1107,13 @@ type Querier interface { // ) // RETURNING id, agent_id, session_key, content, from_msg_idx, to_msg_idx, created_at InsertSummary(ctx context.Context, arg InsertSummaryParams) (MemorySummary, error) + // Get all unique agent IDs that have completed tasks (for multi-agent processing) + // + // SELECT DISTINCT agent_id + // FROM task_completions + // WHERE created_at > ?1 + // ORDER BY agent_id + ListActiveAgents(ctx context.Context, arg ListActiveAgentsParams) ([]string, error) //ListAgentCheckpointsByConversationID // // SELECT id, conversation_id, name, run_state_id, metadata_json, created_at, updated_at diff --git a/pkg/memory/sqlc/queries/rl.sql b/pkg/memory/sqlc/queries/rl.sql index 03932e5cf..a97d60ae7 100644 --- a/pkg/memory/sqlc/queries/rl.sql +++ b/pkg/memory/sqlc/queries/rl.sql @@ -200,3 +200,25 @@ SELECT tr.memory_id, FROM task_retrievals tr JOIN recall_items ri ON tr.memory_id = ri.id WHERE tr.task_id = sqlc.arg(task_id); + +-- name: ListActiveAgents :many +-- Get all unique agent IDs that have completed tasks (for multi-agent processing) +SELECT DISTINCT agent_id +FROM task_completions +WHERE created_at > sqlc.arg(since) +ORDER BY agent_id; + +-- name: GetHighTokenSessions :many +-- Get sessions with high token usage grouped by conversation/agent +SELECT + conversation_id as session_id, + agent_id, + SUM(tokens_used) as total_tokens, + COUNT(*) as task_count +FROM task_completions +WHERE tokens_used > sqlc.arg(min_tokens) + AND created_at > datetime('now', '-24 hours') +GROUP BY conversation_id, agent_id +HAVING total_tokens > sqlc.arg(min_tokens) +ORDER BY total_tokens DESC +LIMIT sqlc.arg(lim); diff --git a/pkg/memory/sqlc/rl.sql.go b/pkg/memory/sqlc/rl.sql.go index 21a446808..4161a9596 100644 --- a/pkg/memory/sqlc/rl.sql.go +++ b/pkg/memory/sqlc/rl.sql.go @@ -90,6 +90,75 @@ func (q *Queries) GetCompletedTasks(ctx context.Context, arg GetCompletedTasksPa return items, nil } +const GetHighTokenSessions = `-- name: GetHighTokenSessions :many +SELECT + conversation_id as session_id, + agent_id, + SUM(tokens_used) as total_tokens, + COUNT(*) as task_count +FROM task_completions +WHERE tokens_used > ?1 + AND created_at > datetime('now', '-24 hours') +GROUP BY conversation_id, agent_id +HAVING total_tokens > ?1 +ORDER BY total_tokens DESC +LIMIT ?2 +` + +type GetHighTokenSessionsParams struct { + MinTokens *int64 `db:"min_tokens" json:"min_tokens"` + Lim int64 `db:"lim" json:"lim"` +} + +type GetHighTokenSessionsRow struct { + SessionID ids.UUID `db:"session_id" json:"session_id"` + AgentID string `db:"agent_id" json:"agent_id"` + TotalTokens *float64 `db:"total_tokens" json:"total_tokens"` + TaskCount int64 `db:"task_count" json:"task_count"` +} + +// Get sessions with high token usage grouped by conversation/agent +// +// SELECT +// conversation_id as session_id, +// agent_id, +// SUM(tokens_used) as total_tokens, +// COUNT(*) as task_count +// FROM task_completions +// WHERE tokens_used > ?1 +// AND created_at > datetime('now', '-24 hours') +// GROUP BY conversation_id, agent_id +// HAVING total_tokens > ?1 +// ORDER BY total_tokens DESC +// LIMIT ?2 +func (q *Queries) GetHighTokenSessions(ctx context.Context, arg GetHighTokenSessionsParams) ([]GetHighTokenSessionsRow, error) { + rows, err := q.db.QueryContext(ctx, GetHighTokenSessions, arg.MinTokens, arg.Lim) + if err != nil { + return nil, err + } + defer rows.Close() + items := []GetHighTokenSessionsRow{} + for rows.Next() { + var i GetHighTokenSessionsRow + if err := rows.Scan( + &i.SessionID, + &i.AgentID, + &i.TotalTokens, + &i.TaskCount, + ); 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 GetMemoriesByRetrievalCount = `-- name: GetMemoriesByRetrievalCount :many SELECT id, agent_id, @@ -329,6 +398,46 @@ func (q *Queries) IncrementTaskRetrievalCount(ctx context.Context, arg Increment return err } +const ListActiveAgents = `-- name: ListActiveAgents :many +SELECT DISTINCT agent_id +FROM task_completions +WHERE created_at > ?1 +ORDER BY agent_id +` + +type ListActiveAgentsParams struct { + Since time.Time `db:"since" json:"since"` +} + +// Get all unique agent IDs that have completed tasks (for multi-agent processing) +// +// SELECT DISTINCT agent_id +// FROM task_completions +// WHERE created_at > ?1 +// ORDER BY agent_id +func (q *Queries) ListActiveAgents(ctx context.Context, arg ListActiveAgentsParams) ([]string, error) { + rows, err := q.db.QueryContext(ctx, ListActiveAgents, arg.Since) + if err != nil { + return nil, err + } + defer rows.Close() + items := []string{} + for rows.Next() { + var agent_id string + if err := rows.Scan(&agent_id); err != nil { + return nil, err + } + items = append(items, agent_id) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const ListHighValueMemories = `-- name: ListHighValueMemories :many SELECT id, agent_id,