opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor(memory): update RL schema, delegate, sqlc queries

Browse source
This commit is contained in:
ZanzyTHEbar 2026-03-05 19:43:10 +00:00
parent 05183881b0
commit 7efb45854d
6 changed files with 427 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

166
pkg/memory/delegate/rl_store_test.go
View file

@ -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))
}
}

90
pkg/memory/delegate/sqlite.go
View file

@ -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 ---

35
pkg/memory/migrations/016_rl_schema.go
View file

@ -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)

22
pkg/memory/sqlc/querier.go
View file

@ -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

22
pkg/memory/sqlc/queries/rl.sql
View file

@ -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);

109
pkg/memory/sqlc/rl.sql.go
View file

@ -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,