feat(memory): add KV store, documents, and audit log subsystems
Extend the memory delegate with three new storage subsystems: - agent_kv: key-value store for agent state and preferences - agent_documents: structured document storage with metadata - agent_audit_log: append-only audit trail for agent actions Includes goose migrations (004-006), sqlc query generation, delegate methods, and comprehensive test coverage.
This commit is contained in:
ZanzyTHEbar
parent
ca56eb2fe1
commit
aaaafb723d
24 changed files with 3676 additions and 2 deletions
|
|
@ -411,6 +411,235 @@ func (d *LibSQLDelegate) CountArchivalChunks(ctx context.Context) (int, error) {
|
||||||
return int(count), err
|
return int(count), err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// --- Key-Value Store ---
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) GetKV(ctx context.Context, agentID, key string) (string, error) {
|
||||||
|
row, err := d.queries.GetKV(ctx, memsqlc.GetKVParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
Key: key,
|
||||||
|
})
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return row.Value, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) UpsertKV(ctx context.Context, agentID, key, value string) error {
|
||||||
|
return d.queries.UpsertKV(ctx, memsqlc.UpsertKVParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
Key: key,
|
||||||
|
Value: value,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) DeleteKV(ctx context.Context, agentID, key string) error {
|
||||||
|
return d.queries.DeleteKV(ctx, memsqlc.DeleteKVParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
Key: key,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) ListKVByPrefix(ctx context.Context, agentID, prefix string, limit int) (map[string]string, error) {
|
||||||
|
rows, err := d.queries.ListKVByPrefix(ctx, memsqlc.ListKVByPrefixParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
Prefix: &prefix,
|
||||||
|
Lim: int64(limit),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
result := make(map[string]string, len(rows))
|
||||||
|
for _, row := range rows {
|
||||||
|
result[row.Key] = row.Value
|
||||||
|
}
|
||||||
|
return result, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Documents ---
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) GetDocument(ctx context.Context, agentID, name string) (*memory.AgentDocument, error) {
|
||||||
|
row, err := d.queries.GetDocument(ctx, memsqlc.GetDocumentParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
Name: name,
|
||||||
|
})
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return sqlcDocToMemory(row), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) UpsertDocument(ctx context.Context, doc *memory.AgentDocument) error {
|
||||||
|
return d.queries.UpsertDocument(ctx, memsqlc.UpsertDocumentParams{
|
||||||
|
ID: doc.ID,
|
||||||
|
AgentID: doc.AgentID,
|
||||||
|
Name: doc.Name,
|
||||||
|
Category: doc.Category,
|
||||||
|
Content: doc.Content,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) DeleteDocument(ctx context.Context, agentID, name string) error {
|
||||||
|
return d.queries.DeleteDocument(ctx, memsqlc.DeleteDocumentParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
Name: name,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) ListDocumentsByCategory(ctx context.Context, agentID, category string) ([]*memory.AgentDocument, error) {
|
||||||
|
rows, err := d.queries.ListDocumentsByCategory(ctx, memsqlc.ListDocumentsByCategoryParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
Category: category,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
docs := make([]*memory.AgentDocument, len(rows))
|
||||||
|
for i, row := range rows {
|
||||||
|
docs[i] = sqlcDocToMemory(row)
|
||||||
|
}
|
||||||
|
return docs, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) ListAllDocuments(ctx context.Context, agentID string) ([]*memory.AgentDocument, error) {
|
||||||
|
rows, err := d.queries.ListAllDocuments(ctx, memsqlc.ListAllDocumentsParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
docs := make([]*memory.AgentDocument, len(rows))
|
||||||
|
for i, row := range rows {
|
||||||
|
docs[i] = sqlcDocToMemory(row)
|
||||||
|
}
|
||||||
|
return docs, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Session Messages ---
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) InsertSessionMessage(ctx context.Context, agentID, sessionKey, role, content string) error {
|
||||||
|
return d.queries.InsertSessionMessage(ctx, memsqlc.InsertSessionMessageParams{
|
||||||
|
ID: ids.New(),
|
||||||
|
AgentID: agentID,
|
||||||
|
SessionKey: sessionKey,
|
||||||
|
Role: role,
|
||||||
|
Content: content,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) ListSessionMessages(ctx context.Context, agentID, sessionKey, role string, limit int) ([]*memory.RecallItem, error) {
|
||||||
|
rows, err := d.queries.ListSessionMessages(ctx, memsqlc.ListSessionMessagesParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
SessionKey: sessionKey,
|
||||||
|
Role: role,
|
||||||
|
Lim: int64(limit),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
items := make([]*memory.RecallItem, len(rows))
|
||||||
|
for i, row := range rows {
|
||||||
|
items[i] = sqlcRecallToMemory(row)
|
||||||
|
}
|
||||||
|
return items, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) CountSessionMessages(ctx context.Context, agentID, sessionKey string) (int64, error) {
|
||||||
|
return d.queries.CountSessionMessages(ctx, memsqlc.CountSessionMessagesParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
SessionKey: sessionKey,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Audit Log ---
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) InsertAuditEntry(ctx context.Context, entry *memory.AuditEntry) error {
|
||||||
|
return d.queries.InsertAuditEntry(ctx, memsqlc.InsertAuditEntryParams{
|
||||||
|
ID: entry.ID,
|
||||||
|
AgentID: entry.AgentID,
|
||||||
|
SessionKey: entry.SessionKey,
|
||||||
|
Action: entry.Action,
|
||||||
|
Target: entry.Target,
|
||||||
|
Input: &entry.Input,
|
||||||
|
Output: &entry.Output,
|
||||||
|
DurationMs: ptrInt64(int64(entry.DurationMS)),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) ListAuditEntries(ctx context.Context, agentID string, limit int) ([]*memory.AuditEntry, error) {
|
||||||
|
rows, err := d.queries.ListAuditEntries(ctx, memsqlc.ListAuditEntriesParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
Lim: int64(limit),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
entries := make([]*memory.AuditEntry, len(rows))
|
||||||
|
for i, row := range rows {
|
||||||
|
entries[i] = sqlcAuditToMemory(row)
|
||||||
|
}
|
||||||
|
return entries, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) ListAuditEntriesByAction(ctx context.Context, agentID, action string, limit int) ([]*memory.AuditEntry, error) {
|
||||||
|
rows, err := d.queries.ListAuditEntriesByAction(ctx, memsqlc.ListAuditEntriesByActionParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
Action: action,
|
||||||
|
Lim: int64(limit),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
entries := make([]*memory.AuditEntry, len(rows))
|
||||||
|
for i, row := range rows {
|
||||||
|
entries[i] = sqlcAuditToMemory(row)
|
||||||
|
}
|
||||||
|
return entries, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) CountAuditEntries(ctx context.Context, agentID string) (int, error) {
|
||||||
|
count, err := d.queries.CountAuditEntries(ctx, memsqlc.CountAuditEntriesParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
})
|
||||||
|
return int(count), err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) ListAuditEntriesBySession(ctx context.Context, agentID, sessionKey string, limit int) ([]*memory.AuditEntry, error) {
|
||||||
|
rows, err := d.queries.ListAuditEntriesBySession(ctx, memsqlc.ListAuditEntriesBySessionParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
SessionKey: sessionKey,
|
||||||
|
Lim: int64(limit),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
entries := make([]*memory.AuditEntry, len(rows))
|
||||||
|
for i, row := range rows {
|
||||||
|
entries[i] = sqlcAuditToMemory(row)
|
||||||
|
}
|
||||||
|
return entries, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) PruneOldAuditEntries(ctx context.Context, agentID string, before time.Time) error {
|
||||||
|
return d.queries.PruneOldAuditEntries(ctx, memsqlc.PruneOldAuditEntriesParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
Before: before,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *LibSQLDelegate) CountAuditEntriesByAction(ctx context.Context, agentID, action string) (int, error) {
|
||||||
|
count, err := d.queries.CountAuditEntriesByAction(ctx, memsqlc.CountAuditEntriesByActionParams{
|
||||||
|
AgentID: agentID,
|
||||||
|
Action: action,
|
||||||
|
})
|
||||||
|
return int(count), err
|
||||||
|
}
|
||||||
|
|
||||||
// --- Conversion helpers ---
|
// --- Conversion helpers ---
|
||||||
|
|
||||||
func sqlcRecallToMemory(row memsqlc.RecallItem) *memory.RecallItem {
|
func sqlcRecallToMemory(row memsqlc.RecallItem) *memory.RecallItem {
|
||||||
|
|
@ -442,3 +671,40 @@ func sqlcChunkToMemory(row memsqlc.ArchivalChunk) *memory.ArchivalChunk {
|
||||||
CreatedAt: row.CreatedAt,
|
CreatedAt: row.CreatedAt,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func sqlcDocToMemory(row memsqlc.AgentDocument) *memory.AgentDocument {
|
||||||
|
return &memory.AgentDocument{
|
||||||
|
ID: row.ID,
|
||||||
|
AgentID: row.AgentID,
|
||||||
|
Name: row.Name,
|
||||||
|
Category: row.Category,
|
||||||
|
Content: row.Content,
|
||||||
|
Version: int(row.Version),
|
||||||
|
IsActive: row.IsActive,
|
||||||
|
CreatedAt: row.CreatedAt,
|
||||||
|
UpdatedAt: row.UpdatedAt,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func sqlcAuditToMemory(row memsqlc.AgentAuditLog) *memory.AuditEntry {
|
||||||
|
entry := &memory.AuditEntry{
|
||||||
|
ID: row.ID,
|
||||||
|
AgentID: row.AgentID,
|
||||||
|
SessionKey: row.SessionKey,
|
||||||
|
Action: row.Action,
|
||||||
|
Target: row.Target,
|
||||||
|
CreatedAt: row.CreatedAt,
|
||||||
|
}
|
||||||
|
if row.Input != nil {
|
||||||
|
entry.Input = *row.Input
|
||||||
|
}
|
||||||
|
if row.Output != nil {
|
||||||
|
entry.Output = *row.Output
|
||||||
|
}
|
||||||
|
if row.DurationMs != nil {
|
||||||
|
entry.DurationMS = int(*row.DurationMs)
|
||||||
|
}
|
||||||
|
return entry
|
||||||
|
}
|
||||||
|
|
||||||
|
func ptrInt64(v int64) *int64 { return &v }
|
||||||
|
|
|
||||||
281
pkg/memory/delegate/sqlite_audit_test.go
Normal file
281
pkg/memory/delegate/sqlite_audit_test.go
Normal file
|
|
@ -0,0 +1,281 @@
|
||||||
|
package delegate
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/ids"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/memory"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func makeAuditEntry(agentID, sessionKey, action, target string) *memory.AuditEntry {
|
||||||
|
return &memory.AuditEntry{
|
||||||
|
ID: ids.New(),
|
||||||
|
AgentID: agentID,
|
||||||
|
SessionKey: sessionKey,
|
||||||
|
Action: action,
|
||||||
|
Target: target,
|
||||||
|
Input: `{"arg":"val"}`,
|
||||||
|
Output: `{"result":"ok"}`,
|
||||||
|
DurationMS: 42,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_InsertAuditEntry(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
entry *memory.AuditEntry
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "insert tool_call entry",
|
||||||
|
entry: makeAuditEntry("a1", "sess-1", "tool_call", "exec"),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "insert memory_write entry",
|
||||||
|
entry: makeAuditEntry("a1", "sess-1", "memory_write", "working_context"),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "insert entry with empty optional fields",
|
||||||
|
entry: &memory.AuditEntry{
|
||||||
|
ID: ids.New(),
|
||||||
|
AgentID: "a1",
|
||||||
|
SessionKey: "sess-1",
|
||||||
|
Action: "state_change",
|
||||||
|
Target: "",
|
||||||
|
Input: "",
|
||||||
|
Output: "",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, tt.entry))
|
||||||
|
|
||||||
|
count, err := d.CountAuditEntries(ctx, tt.entry.AgentID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, 1, count)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_ListAuditEntries(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context)
|
||||||
|
agentID string
|
||||||
|
limit int
|
||||||
|
wantLen int
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "empty returns empty",
|
||||||
|
agentID: "a1",
|
||||||
|
limit: 10,
|
||||||
|
wantLen: 0,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "returns all for agent",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "exec")))
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "memory_write", "kv")))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
limit: 10,
|
||||||
|
wantLen: 2,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "agent isolation",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "t1")))
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a2", "s1", "tool_call", "t2")))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
limit: 10,
|
||||||
|
wantLen: 1,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "respects limit",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", kvKey("t", i))))
|
||||||
|
}
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
limit: 3,
|
||||||
|
wantLen: 3,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
if tt.setup != nil {
|
||||||
|
tt.setup(t, d, ctx)
|
||||||
|
}
|
||||||
|
entries, err := d.ListAuditEntries(ctx, tt.agentID, tt.limit)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Len(t, entries, tt.wantLen)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_ListAuditEntriesByAction(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context)
|
||||||
|
agentID string
|
||||||
|
action string
|
||||||
|
limit int
|
||||||
|
wantLen int
|
||||||
|
wantAction string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "no matches returns empty",
|
||||||
|
agentID: "a1",
|
||||||
|
action: "tool_call",
|
||||||
|
limit: 10,
|
||||||
|
wantLen: 0,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "filters by action",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "exec")))
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "memory_write", "kv")))
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "read")))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
action: "tool_call",
|
||||||
|
limit: 10,
|
||||||
|
wantLen: 2,
|
||||||
|
wantAction: "tool_call",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
if tt.setup != nil {
|
||||||
|
tt.setup(t, d, ctx)
|
||||||
|
}
|
||||||
|
entries, err := d.ListAuditEntriesByAction(ctx, tt.agentID, tt.action, tt.limit)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Len(t, entries, tt.wantLen)
|
||||||
|
if tt.wantAction != "" {
|
||||||
|
for _, e := range entries {
|
||||||
|
assert.Equal(t, tt.wantAction, e.Action)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_ListAuditEntriesBySession(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context)
|
||||||
|
agentID string
|
||||||
|
sessionKey string
|
||||||
|
limit int
|
||||||
|
wantLen int
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "empty returns empty",
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-1",
|
||||||
|
limit: 10,
|
||||||
|
wantLen: 0,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "filters by session",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "sess-A", "tool_call", "t1")))
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "sess-A", "tool_call", "t2")))
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "sess-B", "tool_call", "t3")))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-A",
|
||||||
|
limit: 10,
|
||||||
|
wantLen: 2,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
if tt.setup != nil {
|
||||||
|
tt.setup(t, d, ctx)
|
||||||
|
}
|
||||||
|
entries, err := d.ListAuditEntriesBySession(ctx, tt.agentID, tt.sessionKey, tt.limit)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Len(t, entries, tt.wantLen)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_PruneOldAuditEntries(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "t1")))
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "t2")))
|
||||||
|
|
||||||
|
count, err := d.CountAuditEntries(ctx, "a1")
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, 2, count)
|
||||||
|
|
||||||
|
// Prune entries created before "now + 1 minute" (should remove all)
|
||||||
|
require.NoError(t, d.PruneOldAuditEntries(ctx, "a1", time.Now().Add(time.Minute)))
|
||||||
|
|
||||||
|
count, err = d.CountAuditEntries(ctx, "a1")
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, 0, count)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_CountAuditEntriesByAction(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context)
|
||||||
|
agentID string
|
||||||
|
action string
|
||||||
|
want int
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "no entries returns zero",
|
||||||
|
agentID: "a1",
|
||||||
|
action: "tool_call",
|
||||||
|
want: 0,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "counts only matching action",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "t1")))
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "memory_write", "kv")))
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "t2")))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
action: "tool_call",
|
||||||
|
want: 2,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
if tt.setup != nil {
|
||||||
|
tt.setup(t, d, ctx)
|
||||||
|
}
|
||||||
|
count, err := d.CountAuditEntriesByAction(ctx, tt.agentID, tt.action)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, tt.want, count)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
330
pkg/memory/delegate/sqlite_doc_test.go
Normal file
330
pkg/memory/delegate/sqlite_doc_test.go
Normal file
|
|
@ -0,0 +1,330 @@
|
||||||
|
package delegate
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/ids"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/memory"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func makeDoc(agentID, name, category, content string) *memory.AgentDocument {
|
||||||
|
return &memory.AgentDocument{
|
||||||
|
ID: ids.New(),
|
||||||
|
AgentID: agentID,
|
||||||
|
Name: name,
|
||||||
|
Category: category,
|
||||||
|
Content: content,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_GetDocument(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context)
|
||||||
|
agentID string
|
||||||
|
docName string
|
||||||
|
wantNil bool
|
||||||
|
wantContent string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "missing document returns nil",
|
||||||
|
agentID: "a1",
|
||||||
|
docName: "nonexistent",
|
||||||
|
wantNil: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "returns stored document",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "README", "core", "hello world")))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
docName: "README",
|
||||||
|
wantContent: "hello world",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "agent isolation — different agent sees nil",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "SECRET", "core", "mine")))
|
||||||
|
},
|
||||||
|
agentID: "a2",
|
||||||
|
docName: "SECRET",
|
||||||
|
wantNil: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "document with special characters in content",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "config.json", "config", `{"key":"val","nested":{"a":1}}`)))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
docName: "config.json",
|
||||||
|
wantContent: `{"key":"val","nested":{"a":1}}`,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
if tt.setup != nil {
|
||||||
|
tt.setup(t, d, ctx)
|
||||||
|
}
|
||||||
|
doc, err := d.GetDocument(ctx, tt.agentID, tt.docName)
|
||||||
|
require.NoError(t, err)
|
||||||
|
if tt.wantNil {
|
||||||
|
assert.Nil(t, doc)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
require.NotNil(t, doc)
|
||||||
|
assert.Equal(t, tt.wantContent, doc.Content)
|
||||||
|
assert.Equal(t, tt.agentID, doc.AgentID)
|
||||||
|
assert.Equal(t, tt.docName, doc.Name)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_UpsertDocument(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
ops []*memory.AgentDocument
|
||||||
|
agentID string
|
||||||
|
docName string
|
||||||
|
wantContent string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "insert new document",
|
||||||
|
ops: []*memory.AgentDocument{makeDoc("a1", "doc1", "cat", "first")},
|
||||||
|
agentID: "a1",
|
||||||
|
docName: "doc1",
|
||||||
|
wantContent: "first",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "overwrite existing document content",
|
||||||
|
ops: []*memory.AgentDocument{
|
||||||
|
makeDoc("a1", "doc1", "cat", "original"),
|
||||||
|
{ID: ids.New(), AgentID: "a1", Name: "doc1", Category: "cat", Content: "updated"},
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
docName: "doc1",
|
||||||
|
wantContent: "updated",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "same name different agents are independent",
|
||||||
|
ops: []*memory.AgentDocument{
|
||||||
|
makeDoc("a1", "shared", "cat", "from-a1"),
|
||||||
|
makeDoc("a2", "shared", "cat", "from-a2"),
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
docName: "shared",
|
||||||
|
wantContent: "from-a1",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "empty content is valid",
|
||||||
|
ops: []*memory.AgentDocument{makeDoc("a1", "empty", "cat", "")},
|
||||||
|
agentID: "a1",
|
||||||
|
docName: "empty",
|
||||||
|
wantContent: "",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "large content roundtrip",
|
||||||
|
ops: []*memory.AgentDocument{makeDoc("a1", "big", "cat", longValue(8192))},
|
||||||
|
agentID: "a1",
|
||||||
|
docName: "big",
|
||||||
|
wantContent: longValue(8192),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
for _, doc := range tt.ops {
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, doc))
|
||||||
|
}
|
||||||
|
got, err := d.GetDocument(ctx, tt.agentID, tt.docName)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NotNil(t, got)
|
||||||
|
assert.Equal(t, tt.wantContent, got.Content)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_DeleteDocument(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context)
|
||||||
|
agentID string
|
||||||
|
docName string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "delete nonexistent is idempotent",
|
||||||
|
agentID: "a1",
|
||||||
|
docName: "nope",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "delete existing document",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "doomed", "cat", "bye")))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
docName: "doomed",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "delete only affects target agent",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "shared", "cat", "a1-doc")))
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a2", "shared", "cat", "a2-doc")))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
docName: "shared",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
if tt.setup != nil {
|
||||||
|
tt.setup(t, d, ctx)
|
||||||
|
}
|
||||||
|
require.NoError(t, d.DeleteDocument(ctx, tt.agentID, tt.docName))
|
||||||
|
|
||||||
|
got, err := d.GetDocument(ctx, tt.agentID, tt.docName)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Nil(t, got, "document should be gone after delete")
|
||||||
|
|
||||||
|
if tt.name == "delete only affects target agent" {
|
||||||
|
other, err := d.GetDocument(ctx, "a2", "shared")
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NotNil(t, other, "other agent's document must survive")
|
||||||
|
assert.Equal(t, "a2-doc", other.Content)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_ListDocumentsByCategory(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context)
|
||||||
|
agentID string
|
||||||
|
category string
|
||||||
|
wantLen int
|
||||||
|
wantName string // first result name, if any
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "empty store returns empty slice",
|
||||||
|
agentID: "a1",
|
||||||
|
category: "core",
|
||||||
|
wantLen: 0,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "filters by category",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d1", "core", "one")))
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d2", "core", "two")))
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d3", "config", "three")))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
category: "core",
|
||||||
|
wantLen: 2,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "agent isolation in category listing",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d1", "core", "a1-val")))
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a2", "d2", "core", "a2-val")))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
category: "core",
|
||||||
|
wantLen: 1,
|
||||||
|
wantName: "d1",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "no match returns empty slice",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d1", "core", "val")))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
category: "nonexistent",
|
||||||
|
wantLen: 0,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
if tt.setup != nil {
|
||||||
|
tt.setup(t, d, ctx)
|
||||||
|
}
|
||||||
|
docs, err := d.ListDocumentsByCategory(ctx, tt.agentID, tt.category)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Len(t, docs, tt.wantLen)
|
||||||
|
if tt.wantName != "" && len(docs) > 0 {
|
||||||
|
assert.Equal(t, tt.wantName, docs[0].Name)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_ListAllDocuments(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context)
|
||||||
|
agentID string
|
||||||
|
wantLen int
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "empty store returns empty slice",
|
||||||
|
agentID: "a1",
|
||||||
|
wantLen: 0,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "returns all documents for agent",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d1", "core", "one")))
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d2", "config", "two")))
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d3", "skills", "three")))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
wantLen: 3,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "agent isolation — counts only own docs",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d1", "core", "a1")))
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a2", "d2", "core", "a2")))
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a2", "d3", "core", "a2b")))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
wantLen: 1,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "multiple categories all returned",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
cats := []string{"core", "config", "skills", "templates"}
|
||||||
|
for i, c := range cats {
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", kvKey("doc-", i), c, "val")))
|
||||||
|
}
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
wantLen: 4,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
if tt.setup != nil {
|
||||||
|
tt.setup(t, d, ctx)
|
||||||
|
}
|
||||||
|
docs, err := d.ListAllDocuments(ctx, tt.agentID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Len(t, docs, tt.wantLen)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
227
pkg/memory/delegate/sqlite_integration_test.go
Normal file
227
pkg/memory/delegate/sqlite_integration_test.go
Normal file
|
|
@ -0,0 +1,227 @@
|
||||||
|
package delegate
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/ids"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/memory"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestCronKVBackend_Roundtrip(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
agentID := "picoclaw"
|
||||||
|
kvKey := "cron:store"
|
||||||
|
|
||||||
|
type cronStore struct {
|
||||||
|
Version int `json:"version"`
|
||||||
|
Jobs []struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
} `json:"jobs"`
|
||||||
|
}
|
||||||
|
|
||||||
|
store := cronStore{
|
||||||
|
Version: 1,
|
||||||
|
Jobs: []struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
}{
|
||||||
|
{ID: "job-1", Name: "daily report", Enabled: true},
|
||||||
|
{ID: "job-2", Name: "weekly backup", Enabled: false},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
data, err := json.Marshal(store)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, agentID, kvKey, string(data)))
|
||||||
|
|
||||||
|
raw, err := d.GetKV(ctx, agentID, kvKey)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NotEmpty(t, raw)
|
||||||
|
|
||||||
|
var loaded cronStore
|
||||||
|
require.NoError(t, json.Unmarshal([]byte(raw), &loaded))
|
||||||
|
|
||||||
|
assert.Equal(t, 1, loaded.Version)
|
||||||
|
assert.Len(t, loaded.Jobs, 2)
|
||||||
|
assert.Equal(t, "daily report", loaded.Jobs[0].Name)
|
||||||
|
assert.True(t, loaded.Jobs[0].Enabled)
|
||||||
|
assert.Equal(t, "weekly backup", loaded.Jobs[1].Name)
|
||||||
|
assert.False(t, loaded.Jobs[1].Enabled)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCronKVBackend_UpdatePreservesShape(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
agentID := "picoclaw"
|
||||||
|
kvKey := "cron:store"
|
||||||
|
|
||||||
|
v1 := `{"version":1,"jobs":[{"id":"j1","name":"test","enabled":true}]}`
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, agentID, kvKey, v1))
|
||||||
|
|
||||||
|
v2 := `{"version":1,"jobs":[{"id":"j1","name":"test","enabled":true},{"id":"j2","name":"new","enabled":true}]}`
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, agentID, kvKey, v2))
|
||||||
|
|
||||||
|
raw, err := d.GetKV(ctx, agentID, kvKey)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, v2, raw)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCronKVBackend_PrefixScan(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
agentID := "picoclaw"
|
||||||
|
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, agentID, "cron:store", "{}"))
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, agentID, "cron:lock", "held"))
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, agentID, "focus:sess1", "{}"))
|
||||||
|
|
||||||
|
result, err := d.ListKVByPrefix(ctx, agentID, "cron:", 10)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Len(t, result, 2)
|
||||||
|
assert.Contains(t, result, "cron:store")
|
||||||
|
assert.Contains(t, result, "cron:lock")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestEndToEnd_SessionAndAuditFlow(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
agentID := "a1"
|
||||||
|
sessionKey := "sess-integration"
|
||||||
|
|
||||||
|
// 1. Insert session messages
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, agentID, sessionKey, "user", "What is the weather?"))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, agentID, sessionKey, "assistant", "Let me check..."))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, agentID, sessionKey, "tool", `{"temp":72,"unit":"F"}`))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, agentID, sessionKey, "assistant", "It's 72F."))
|
||||||
|
|
||||||
|
// 2. Verify session message count
|
||||||
|
count, err := d.CountSessionMessages(ctx, agentID, sessionKey)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, int64(4), count)
|
||||||
|
|
||||||
|
// 3. Verify message ordering (ASC)
|
||||||
|
msgs, err := d.ListSessionMessages(ctx, agentID, sessionKey, "", 50)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Len(t, msgs, 4)
|
||||||
|
assert.Equal(t, "user", msgs[0].Role)
|
||||||
|
assert.Equal(t, "What is the weather?", msgs[0].Content)
|
||||||
|
assert.Equal(t, "assistant", msgs[3].Role)
|
||||||
|
|
||||||
|
// 4. Insert audit entries for the tool call
|
||||||
|
auditEntry := &memory.AuditEntry{
|
||||||
|
ID: ids.New(),
|
||||||
|
AgentID: agentID,
|
||||||
|
SessionKey: sessionKey,
|
||||||
|
Action: "tool_call",
|
||||||
|
Target: "weather_api",
|
||||||
|
Input: `{"location":"here"}`,
|
||||||
|
Output: `{"temp":72}`,
|
||||||
|
DurationMS: 150,
|
||||||
|
}
|
||||||
|
require.NoError(t, d.InsertAuditEntry(ctx, auditEntry))
|
||||||
|
|
||||||
|
// 5. Verify audit log by session
|
||||||
|
auditEntries, err := d.ListAuditEntriesBySession(ctx, agentID, sessionKey, 10)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Len(t, auditEntries, 1)
|
||||||
|
assert.Equal(t, "tool_call", auditEntries[0].Action)
|
||||||
|
assert.Equal(t, "weather_api", auditEntries[0].Target)
|
||||||
|
|
||||||
|
// 6. Store KV state (e.g. focus checkpoint)
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, agentID, "focus:"+sessionKey, `{"topic":"weather query","checkpoint_index":2}`))
|
||||||
|
|
||||||
|
kvVal, err := d.GetKV(ctx, agentID, "focus:"+sessionKey)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Contains(t, kvVal, "weather query")
|
||||||
|
|
||||||
|
// 7. Store a document
|
||||||
|
doc := &memory.AgentDocument{
|
||||||
|
ID: ids.New(),
|
||||||
|
AgentID: agentID,
|
||||||
|
Name: "AGENT.md",
|
||||||
|
Category: "core",
|
||||||
|
Content: "# Agent Identity\nI am picoclaw.",
|
||||||
|
}
|
||||||
|
require.NoError(t, d.UpsertDocument(ctx, doc))
|
||||||
|
|
||||||
|
loadedDoc, err := d.GetDocument(ctx, agentID, "AGENT.md")
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NotNil(t, loadedDoc)
|
||||||
|
assert.Equal(t, "# Agent Identity\nI am picoclaw.", loadedDoc.Content)
|
||||||
|
|
||||||
|
// 8. Verify cross-table isolation: different session sees nothing
|
||||||
|
otherMsgs, err := d.ListSessionMessages(ctx, agentID, "sess-other", "", 50)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Empty(t, otherMsgs)
|
||||||
|
|
||||||
|
otherAudit, err := d.ListAuditEntriesBySession(ctx, agentID, "sess-other", 10)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Empty(t, otherAudit)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestEndToEnd_WorkingContextAndRecallRoundtrip(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
agentID := "a1"
|
||||||
|
sessionKey := "sess-wc"
|
||||||
|
|
||||||
|
// 1. Upsert working context
|
||||||
|
require.NoError(t, d.UpsertWorkingContext(ctx, agentID, sessionKey, "Initial system prompt state"))
|
||||||
|
|
||||||
|
wc, err := d.GetWorkingContext(ctx, agentID, sessionKey)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NotNil(t, wc)
|
||||||
|
assert.Equal(t, "Initial system prompt state", wc.Content)
|
||||||
|
|
||||||
|
// 2. Insert recall items
|
||||||
|
item := &memory.RecallItem{
|
||||||
|
ID: ids.New(),
|
||||||
|
AgentID: agentID,
|
||||||
|
SessionKey: sessionKey,
|
||||||
|
Role: "user",
|
||||||
|
Sector: memory.SectorEpisodic,
|
||||||
|
Importance: 0.8,
|
||||||
|
Salience: 0.7,
|
||||||
|
DecayRate: 0.01,
|
||||||
|
Content: "Important user preference",
|
||||||
|
Tags: "preference",
|
||||||
|
}
|
||||||
|
require.NoError(t, d.InsertRecallItem(ctx, item))
|
||||||
|
|
||||||
|
// 3. Count recall items
|
||||||
|
recallCount, err := d.CountRecallItems(ctx, agentID, sessionKey)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, 1, recallCount)
|
||||||
|
|
||||||
|
// 4. Update working context
|
||||||
|
require.NoError(t, d.UpsertWorkingContext(ctx, agentID, sessionKey, "Updated with preference awareness"))
|
||||||
|
|
||||||
|
wc, err = d.GetWorkingContext(ctx, agentID, sessionKey)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, "Updated with preference awareness", wc.Content)
|
||||||
|
|
||||||
|
// 5. Insert a summary
|
||||||
|
summary := &memory.MemorySummary{
|
||||||
|
ID: ids.New(),
|
||||||
|
AgentID: agentID,
|
||||||
|
SessionKey: sessionKey,
|
||||||
|
Content: "User discussed preferences. Key info captured.",
|
||||||
|
FromMsgIdx: 0,
|
||||||
|
ToMsgIdx: 5,
|
||||||
|
}
|
||||||
|
require.NoError(t, d.InsertSummary(ctx, summary))
|
||||||
|
|
||||||
|
summaries, err := d.ListSummaries(ctx, agentID, sessionKey, 10)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Len(t, summaries, 1)
|
||||||
|
assert.Equal(t, "User discussed preferences. Key info captured.", summaries[0].Content)
|
||||||
|
}
|
||||||
330
pkg/memory/delegate/sqlite_kv_test.go
Normal file
330
pkg/memory/delegate/sqlite_kv_test.go
Normal file
|
|
@ -0,0 +1,330 @@
|
||||||
|
package delegate
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_GetKV(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context)
|
||||||
|
agentID string
|
||||||
|
key string
|
||||||
|
wantValue string
|
||||||
|
wantErr bool
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "missing key returns empty string",
|
||||||
|
agentID: "agent-1",
|
||||||
|
key: "nonexistent",
|
||||||
|
wantValue: "",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "returns stored value",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "agent-1", "greeting", "hello"))
|
||||||
|
},
|
||||||
|
agentID: "agent-1",
|
||||||
|
key: "greeting",
|
||||||
|
wantValue: "hello",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "agent isolation — different agent sees empty",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "agent-A", "secret", "mine"))
|
||||||
|
},
|
||||||
|
agentID: "agent-B",
|
||||||
|
key: "secret",
|
||||||
|
wantValue: "",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "empty value is valid",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "agent-1", "blank", ""))
|
||||||
|
},
|
||||||
|
agentID: "agent-1",
|
||||||
|
key: "blank",
|
||||||
|
wantValue: "",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "value with special characters",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "agent-1", "json", `{"key":"val","n":42}`))
|
||||||
|
},
|
||||||
|
agentID: "agent-1",
|
||||||
|
key: "json",
|
||||||
|
wantValue: `{"key":"val","n":42}`,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "key with colons and slashes",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "agent-1", "focus:session/abc", "data"))
|
||||||
|
},
|
||||||
|
agentID: "agent-1",
|
||||||
|
key: "focus:session/abc",
|
||||||
|
wantValue: "data",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
if tt.setup != nil {
|
||||||
|
tt.setup(t, d, ctx)
|
||||||
|
}
|
||||||
|
got, err := d.GetKV(ctx, tt.agentID, tt.key)
|
||||||
|
if tt.wantErr {
|
||||||
|
assert.Error(t, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, tt.wantValue, got)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_UpsertKV(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
ops []kvOp
|
||||||
|
agentID string
|
||||||
|
key string
|
||||||
|
want string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "insert new key",
|
||||||
|
ops: []kvOp{{agent: "a1", key: "k1", val: "v1"}},
|
||||||
|
agentID: "a1",
|
||||||
|
key: "k1",
|
||||||
|
want: "v1",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "overwrite existing key",
|
||||||
|
ops: []kvOp{
|
||||||
|
{agent: "a1", key: "k1", val: "first"},
|
||||||
|
{agent: "a1", key: "k1", val: "second"},
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
key: "k1",
|
||||||
|
want: "second",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "overwrite with empty",
|
||||||
|
ops: []kvOp{
|
||||||
|
{agent: "a1", key: "k1", val: "nonempty"},
|
||||||
|
{agent: "a1", key: "k1", val: ""},
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
key: "k1",
|
||||||
|
want: "",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "same key different agents are independent",
|
||||||
|
ops: []kvOp{
|
||||||
|
{agent: "a1", key: "shared", val: "from-a1"},
|
||||||
|
{agent: "a2", key: "shared", val: "from-a2"},
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
key: "shared",
|
||||||
|
want: "from-a1",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "large value roundtrip",
|
||||||
|
ops: []kvOp{
|
||||||
|
{agent: "a1", key: "big", val: longValue(4096)},
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
key: "big",
|
||||||
|
want: longValue(4096),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
for _, op := range tt.ops {
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, op.agent, op.key, op.val))
|
||||||
|
}
|
||||||
|
got, err := d.GetKV(ctx, tt.agentID, tt.key)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, tt.want, got)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_DeleteKV(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context)
|
||||||
|
agentID string
|
||||||
|
key string
|
||||||
|
wantErr bool
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "delete nonexistent key is idempotent",
|
||||||
|
agentID: "a1",
|
||||||
|
key: "nope",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "delete existing key",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "a1", "doomed", "bye"))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
key: "doomed",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "delete only affects target agent",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "a1", "shared", "a1-val"))
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "a2", "shared", "a2-val"))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
key: "shared",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
if tt.setup != nil {
|
||||||
|
tt.setup(t, d, ctx)
|
||||||
|
}
|
||||||
|
err := d.DeleteKV(ctx, tt.agentID, tt.key)
|
||||||
|
if tt.wantErr {
|
||||||
|
assert.Error(t, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
got, err := d.GetKV(ctx, tt.agentID, tt.key)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Empty(t, got, "key should be gone after delete")
|
||||||
|
|
||||||
|
if tt.name == "delete only affects target agent" {
|
||||||
|
other, err := d.GetKV(ctx, "a2", "shared")
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, "a2-val", other, "other agent's key must survive")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_ListKVByPrefix(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context)
|
||||||
|
agentID string
|
||||||
|
prefix string
|
||||||
|
limit int
|
||||||
|
want map[string]string
|
||||||
|
wantErr bool
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "empty store returns empty map",
|
||||||
|
agentID: "a1",
|
||||||
|
prefix: "obs:",
|
||||||
|
limit: 10,
|
||||||
|
want: map[string]string{},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "matches prefix exactly",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "a1", "obs:1", "one"))
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "a1", "obs:2", "two"))
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "a1", "focus:x", "other"))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
prefix: "obs:",
|
||||||
|
limit: 10,
|
||||||
|
want: map[string]string{"obs:1": "one", "obs:2": "two"},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "respects limit",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
for i := 0; i < 5; i++ {
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "a1", kvKey("item:", i), kvVal(i)))
|
||||||
|
}
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
prefix: "item:",
|
||||||
|
limit: 3,
|
||||||
|
want: nil, // checked via length only
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "agent isolation in prefix scan",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "a1", "ns:k1", "a1-val"))
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "a2", "ns:k2", "a2-val"))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
prefix: "ns:",
|
||||||
|
limit: 10,
|
||||||
|
want: map[string]string{"ns:k1": "a1-val"},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "prefix with no trailing separator",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "a1", "abc", "one"))
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "a1", "abd", "two"))
|
||||||
|
require.NoError(t, d.UpsertKV(ctx, "a1", "xyz", "three"))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
prefix: "ab",
|
||||||
|
limit: 10,
|
||||||
|
want: map[string]string{"abc": "one", "abd": "two"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
if tt.setup != nil {
|
||||||
|
tt.setup(t, d, ctx)
|
||||||
|
}
|
||||||
|
got, err := d.ListKVByPrefix(ctx, tt.agentID, tt.prefix, tt.limit)
|
||||||
|
if tt.wantErr {
|
||||||
|
assert.Error(t, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
if tt.name == "respects limit" {
|
||||||
|
assert.LessOrEqual(t, len(got), tt.limit)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
assert.Equal(t, tt.want, got)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- helpers ---
|
||||||
|
|
||||||
|
type kvOp struct {
|
||||||
|
agent, key, val string
|
||||||
|
}
|
||||||
|
|
||||||
|
func longValue(n int) string {
|
||||||
|
b := make([]byte, n)
|
||||||
|
for i := range b {
|
||||||
|
b[i] = 'A' + byte(i%26)
|
||||||
|
}
|
||||||
|
return string(b)
|
||||||
|
}
|
||||||
|
|
||||||
|
func kvKey(prefix string, i int) string {
|
||||||
|
return prefix + string(rune('0'+i))
|
||||||
|
}
|
||||||
|
|
||||||
|
func kvVal(i int) string {
|
||||||
|
return string(rune('a'+i)) + "-val"
|
||||||
|
}
|
||||||
269
pkg/memory/delegate/sqlite_session_msg_test.go
Normal file
269
pkg/memory/delegate/sqlite_session_msg_test.go
Normal file
|
|
@ -0,0 +1,269 @@
|
||||||
|
package delegate
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/ids"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/memory"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func makeRecallItem(agentID, sessionKey, role, content, tags string) *memory.RecallItem {
|
||||||
|
now := time.Now()
|
||||||
|
return &memory.RecallItem{
|
||||||
|
ID: ids.New(),
|
||||||
|
AgentID: agentID,
|
||||||
|
SessionKey: sessionKey,
|
||||||
|
Role: role,
|
||||||
|
Sector: memory.SectorEpisodic,
|
||||||
|
Importance: 0.5,
|
||||||
|
Salience: 0.5,
|
||||||
|
DecayRate: 0.01,
|
||||||
|
Content: content,
|
||||||
|
Tags: tags,
|
||||||
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_InsertSessionMessage(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
agentID string
|
||||||
|
sessionKey string
|
||||||
|
role string
|
||||||
|
content string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "insert user message",
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-1",
|
||||||
|
role: "user",
|
||||||
|
content: "Hello, agent!",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "insert assistant message",
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-1",
|
||||||
|
role: "assistant",
|
||||||
|
content: "Hello, human!",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "insert tool message",
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-1",
|
||||||
|
role: "tool",
|
||||||
|
content: `{"result":"ok"}`,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "empty content is valid",
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-1",
|
||||||
|
role: "system",
|
||||||
|
content: "",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, tt.agentID, tt.sessionKey, tt.role, tt.content))
|
||||||
|
|
||||||
|
count, err := d.CountSessionMessages(ctx, tt.agentID, tt.sessionKey)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, int64(1), count)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_ListSessionMessages(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context)
|
||||||
|
agentID string
|
||||||
|
sessionKey string
|
||||||
|
role string
|
||||||
|
limit int
|
||||||
|
wantLen int
|
||||||
|
wantRole string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "empty session returns empty",
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-empty",
|
||||||
|
role: "",
|
||||||
|
limit: 50,
|
||||||
|
wantLen: 0,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "list all roles (role=empty string)",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "msg1"))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "assistant", "msg2"))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "tool", "msg3"))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-1",
|
||||||
|
role: "",
|
||||||
|
limit: 50,
|
||||||
|
wantLen: 3,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "filter by role=user",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "u1"))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "assistant", "a1"))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "u2"))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-1",
|
||||||
|
role: "user",
|
||||||
|
limit: 50,
|
||||||
|
wantLen: 2,
|
||||||
|
wantRole: "user",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "session isolation",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-A", "user", "msgA"))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-B", "user", "msgB"))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-A",
|
||||||
|
role: "",
|
||||||
|
limit: 50,
|
||||||
|
wantLen: 1,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "agent isolation",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "from-a1"))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a2", "sess-1", "user", "from-a2"))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-1",
|
||||||
|
role: "",
|
||||||
|
limit: 50,
|
||||||
|
wantLen: 1,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "respects limit",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", kvVal(i)))
|
||||||
|
}
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-1",
|
||||||
|
role: "",
|
||||||
|
limit: 5,
|
||||||
|
wantLen: 5,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "ordered by created_at ASC",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "first"))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "assistant", "second"))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "third"))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-1",
|
||||||
|
role: "",
|
||||||
|
limit: 50,
|
||||||
|
wantLen: 3,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
if tt.setup != nil {
|
||||||
|
tt.setup(t, d, ctx)
|
||||||
|
}
|
||||||
|
msgs, err := d.ListSessionMessages(ctx, tt.agentID, tt.sessionKey, tt.role, tt.limit)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Len(t, msgs, tt.wantLen)
|
||||||
|
|
||||||
|
if tt.wantRole != "" {
|
||||||
|
for _, m := range msgs {
|
||||||
|
assert.Equal(t, tt.wantRole, m.Role)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if tt.name == "ordered by created_at ASC" && len(msgs) == 3 {
|
||||||
|
assert.Equal(t, "first", msgs[0].Content)
|
||||||
|
assert.Equal(t, "second", msgs[1].Content)
|
||||||
|
assert.Equal(t, "third", msgs[2].Content)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLibSQLDelegate_CountSessionMessages(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context)
|
||||||
|
agentID string
|
||||||
|
sessionKey string
|
||||||
|
want int64
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "empty session returns zero",
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-empty",
|
||||||
|
want: 0,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "counts session messages only",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "m1"))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "assistant", "m2"))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-1",
|
||||||
|
want: 2,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "excludes non-session recall items",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "session-msg"))
|
||||||
|
// Insert a generic recall item via the standard method
|
||||||
|
item := makeRecallItem("a1", "sess-1", "user", "generic-item", "other-tag")
|
||||||
|
require.NoError(t, d.InsertRecallItem(ctx, item))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-1",
|
||||||
|
want: 1,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "session isolation in count",
|
||||||
|
setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) {
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-A", "user", "m1"))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-A", "user", "m2"))
|
||||||
|
require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-B", "user", "m3"))
|
||||||
|
},
|
||||||
|
agentID: "a1",
|
||||||
|
sessionKey: "sess-A",
|
||||||
|
want: 2,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
if tt.setup != nil {
|
||||||
|
tt.setup(t, d, ctx)
|
||||||
|
}
|
||||||
|
count, err := d.CountSessionMessages(ctx, tt.agentID, tt.sessionKey)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, tt.want, count)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -564,3 +564,141 @@ func TestEmbeddingValueScanRoundTrip(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestIntegration_FullStackNoDisk(t *testing.T) {
|
||||||
|
d := newTestDelegate(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
agentID := "integration-agent"
|
||||||
|
|
||||||
|
t.Run("KV_Store", func(t *testing.T) {
|
||||||
|
if err := d.UpsertKV(ctx, agentID, "last_channel", "telegram:42"); err != nil {
|
||||||
|
t.Fatalf("UpsertKV: %v", err)
|
||||||
|
}
|
||||||
|
val, err := d.GetKV(ctx, agentID, "last_channel")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetKV: %v", err)
|
||||||
|
}
|
||||||
|
if val != "telegram:42" {
|
||||||
|
t.Errorf("expected telegram:42, got %q", val)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := d.UpsertKV(ctx, agentID, "cron:store", `{"version":1,"jobs":[]}`); err != nil {
|
||||||
|
t.Fatalf("UpsertKV cron: %v", err)
|
||||||
|
}
|
||||||
|
kvs, err := d.ListKVByPrefix(ctx, agentID, "cron:", 10)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ListKVByPrefix: %v", err)
|
||||||
|
}
|
||||||
|
if len(kvs) != 1 {
|
||||||
|
t.Errorf("expected 1 cron KV, got %d", len(kvs))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("Documents", func(t *testing.T) {
|
||||||
|
doc := &memory.AgentDocument{
|
||||||
|
ID: ids.New(),
|
||||||
|
AgentID: agentID,
|
||||||
|
Name: "AGENTS.md",
|
||||||
|
Category: "bootstrap",
|
||||||
|
Content: "# Agent Config\nBootstrap content",
|
||||||
|
Version: 1,
|
||||||
|
IsActive: true,
|
||||||
|
}
|
||||||
|
if err := d.UpsertDocument(ctx, doc); err != nil {
|
||||||
|
t.Fatalf("UpsertDocument: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
got, err := d.GetDocument(ctx, agentID, "AGENTS.md")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetDocument: %v", err)
|
||||||
|
}
|
||||||
|
if got.Content != doc.Content {
|
||||||
|
t.Errorf("content mismatch: %q vs %q", got.Content, doc.Content)
|
||||||
|
}
|
||||||
|
|
||||||
|
docs, err := d.ListDocumentsByCategory(ctx, agentID, "bootstrap")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ListDocumentsByCategory: %v", err)
|
||||||
|
}
|
||||||
|
if len(docs) != 1 {
|
||||||
|
t.Errorf("expected 1 bootstrap doc, got %d", len(docs))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("Sessions_via_RecallItems", func(t *testing.T) {
|
||||||
|
for i, msg := range []struct{ role, content string }{
|
||||||
|
{"user", "hello agent"},
|
||||||
|
{"assistant", "hi there!"},
|
||||||
|
{"user", "how are you?"},
|
||||||
|
} {
|
||||||
|
item := &memory.RecallItem{
|
||||||
|
ID: ids.New(),
|
||||||
|
AgentID: agentID,
|
||||||
|
SessionKey: "test-session",
|
||||||
|
Role: msg.role,
|
||||||
|
Sector: memory.SectorEpisodic,
|
||||||
|
Importance: 0.5,
|
||||||
|
Salience: 0.5,
|
||||||
|
Content: msg.content,
|
||||||
|
Tags: "session-message",
|
||||||
|
}
|
||||||
|
if err := d.InsertRecallItem(ctx, item); err != nil {
|
||||||
|
t.Fatalf("InsertRecallItem[%d]: %v", i, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
items, err := d.ListRecallItems(ctx, agentID, "test-session", 100, 0)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ListRecallItems: %v", err)
|
||||||
|
}
|
||||||
|
if len(items) != 3 {
|
||||||
|
t.Errorf("expected 3 session items, got %d", len(items))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("AuditLog", func(t *testing.T) {
|
||||||
|
entry := &memory.AuditEntry{
|
||||||
|
ID: ids.New(),
|
||||||
|
AgentID: agentID,
|
||||||
|
SessionKey: "test-session",
|
||||||
|
Action: "tool_call",
|
||||||
|
Target: "exec",
|
||||||
|
Input: `{"command":"ls"}`,
|
||||||
|
}
|
||||||
|
if err := d.InsertAuditEntry(ctx, entry); err != nil {
|
||||||
|
t.Fatalf("InsertAuditEntry: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
entries, err := d.ListAuditEntries(ctx, agentID, 10)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ListAuditEntries: %v", err)
|
||||||
|
}
|
||||||
|
if len(entries) != 1 {
|
||||||
|
t.Errorf("expected 1 audit entry, got %d", len(entries))
|
||||||
|
}
|
||||||
|
if entries[0].Target != "exec" {
|
||||||
|
t.Errorf("expected target 'exec', got %q", entries[0].Target)
|
||||||
|
}
|
||||||
|
|
||||||
|
count, err := d.CountAuditEntries(ctx, agentID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CountAuditEntries: %v", err)
|
||||||
|
}
|
||||||
|
if count != 1 {
|
||||||
|
t.Errorf("expected count 1, got %d", count)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("WorkingContext", func(t *testing.T) {
|
||||||
|
if err := d.UpsertWorkingContext(ctx, agentID, "sess-1", "agent memory contents"); err != nil {
|
||||||
|
t.Fatalf("UpsertWorkingContext: %v", err)
|
||||||
|
}
|
||||||
|
wc, err := d.GetWorkingContext(ctx, agentID, "sess-1")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetWorkingContext: %v", err)
|
||||||
|
}
|
||||||
|
if wc == nil || wc.Content != "agent memory contents" {
|
||||||
|
t.Errorf("unexpected working context: %v", wc)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -211,6 +211,34 @@ const (
|
||||||
PressureFlush PressureLevel = "flush" // > 85%
|
PressureFlush PressureLevel = "flush" // > 85%
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// --- Additional domain types ---
|
||||||
|
|
||||||
|
// AgentDocument is a versioned named document stored in the database.
|
||||||
|
type AgentDocument struct {
|
||||||
|
ID ids.UUID
|
||||||
|
AgentID string
|
||||||
|
Name string
|
||||||
|
Category string
|
||||||
|
Content string
|
||||||
|
Version int
|
||||||
|
IsActive bool
|
||||||
|
CreatedAt time.Time
|
||||||
|
UpdatedAt time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
// AuditEntry records a single agent action for the audit log.
|
||||||
|
type AuditEntry struct {
|
||||||
|
ID ids.UUID
|
||||||
|
AgentID string
|
||||||
|
SessionKey string
|
||||||
|
Action string // "tool_call", "memory_write", "doc_update", "state_change"
|
||||||
|
Target string // tool name, doc name, key name
|
||||||
|
Input string
|
||||||
|
Output string
|
||||||
|
DurationMS int
|
||||||
|
CreatedAt time.Time
|
||||||
|
}
|
||||||
|
|
||||||
// --- Delegate interface (backend) ---
|
// --- Delegate interface (backend) ---
|
||||||
|
|
||||||
// MemoryDelegate is the pure storage backend for the memory system.
|
// MemoryDelegate is the pure storage backend for the memory system.
|
||||||
|
|
@ -260,6 +288,25 @@ type MemoryDelegate interface {
|
||||||
CountRecallItems(ctx context.Context, agentID, sessionKey string) (int, error)
|
CountRecallItems(ctx context.Context, agentID, sessionKey string) (int, error)
|
||||||
CountArchivalChunks(ctx context.Context) (int, error)
|
CountArchivalChunks(ctx context.Context) (int, error)
|
||||||
|
|
||||||
|
// --- Key-Value Store ---
|
||||||
|
GetKV(ctx context.Context, agentID, key string) (string, error)
|
||||||
|
UpsertKV(ctx context.Context, agentID, key, value string) error
|
||||||
|
DeleteKV(ctx context.Context, agentID, key string) error
|
||||||
|
ListKVByPrefix(ctx context.Context, agentID, prefix string, limit int) (map[string]string, error)
|
||||||
|
|
||||||
|
// --- Documents ---
|
||||||
|
GetDocument(ctx context.Context, agentID, name string) (*AgentDocument, error)
|
||||||
|
UpsertDocument(ctx context.Context, doc *AgentDocument) error
|
||||||
|
DeleteDocument(ctx context.Context, agentID, name string) error
|
||||||
|
ListDocumentsByCategory(ctx context.Context, agentID, category string) ([]*AgentDocument, error)
|
||||||
|
ListAllDocuments(ctx context.Context, agentID string) ([]*AgentDocument, error)
|
||||||
|
|
||||||
|
// --- Audit Log ---
|
||||||
|
InsertAuditEntry(ctx context.Context, entry *AuditEntry) error
|
||||||
|
ListAuditEntries(ctx context.Context, agentID string, limit int) ([]*AuditEntry, error)
|
||||||
|
ListAuditEntriesByAction(ctx context.Context, agentID, action string, limit int) ([]*AuditEntry, error)
|
||||||
|
CountAuditEntries(ctx context.Context, agentID string) (int, error)
|
||||||
|
|
||||||
// --- Capability Detection ---
|
// --- Capability Detection ---
|
||||||
HasVectorSearch() bool
|
HasVectorSearch() bool
|
||||||
HasFTS() bool
|
HasFTS() bool
|
||||||
|
|
|
||||||
35
pkg/memory/migrations/004_agent_kv.go
Normal file
35
pkg/memory/migrations/004_agent_kv.go
Normal file
|
|
@ -0,0 +1,35 @@
|
||||||
|
package migrations
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"github.com/pressly/goose/v3"
|
||||||
|
)
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
goose.AddMigrationContext(up004AgentKV, down004AgentKV)
|
||||||
|
}
|
||||||
|
|
||||||
|
func up004AgentKV(ctx context.Context, tx *sql.Tx) error {
|
||||||
|
_, err := tx.ExecContext(ctx, `CREATE TABLE IF NOT EXISTS agent_kv (
|
||||||
|
agent_id TEXT NOT NULL,
|
||||||
|
key TEXT NOT NULL,
|
||||||
|
value TEXT NOT NULL DEFAULT '',
|
||||||
|
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
PRIMARY KEY (agent_id, key)
|
||||||
|
)`)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("004_agent_kv up: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func down004AgentKV(ctx context.Context, tx *sql.Tx) error {
|
||||||
|
_, err := tx.ExecContext(ctx, `DROP TABLE IF EXISTS agent_kv`)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("004_agent_kv down: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
47
pkg/memory/migrations/005_agent_documents.go
Normal file
47
pkg/memory/migrations/005_agent_documents.go
Normal file
|
|
@ -0,0 +1,47 @@
|
||||||
|
package migrations
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"github.com/pressly/goose/v3"
|
||||||
|
)
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
goose.AddMigrationContext(up005AgentDocuments, down005AgentDocuments)
|
||||||
|
}
|
||||||
|
|
||||||
|
func up005AgentDocuments(ctx context.Context, tx *sql.Tx) error {
|
||||||
|
stmts := []string{
|
||||||
|
`CREATE TABLE IF NOT EXISTS agent_documents (
|
||||||
|
id BLOB PRIMARY KEY,
|
||||||
|
agent_id TEXT NOT NULL,
|
||||||
|
name TEXT NOT NULL,
|
||||||
|
category TEXT NOT NULL DEFAULT 'bootstrap',
|
||||||
|
content TEXT NOT NULL,
|
||||||
|
version INTEGER NOT NULL DEFAULT 1,
|
||||||
|
is_active INTEGER NOT NULL DEFAULT 1,
|
||||||
|
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
UNIQUE(agent_id, name)
|
||||||
|
)`,
|
||||||
|
`CREATE INDEX IF NOT EXISTS idx_docs_agent_cat ON agent_documents(agent_id, category)`,
|
||||||
|
`CREATE INDEX IF NOT EXISTS idx_docs_name ON agent_documents(name)`,
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, s := range stmts {
|
||||||
|
if _, err := tx.ExecContext(ctx, s); err != nil {
|
||||||
|
return fmt.Errorf("005_agent_documents up: %w\nSQL: %s", err, s)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func down005AgentDocuments(ctx context.Context, tx *sql.Tx) error {
|
||||||
|
_, err := tx.ExecContext(ctx, `DROP TABLE IF EXISTS agent_documents`)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("005_agent_documents down: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
46
pkg/memory/migrations/006_agent_audit_log.go
Normal file
46
pkg/memory/migrations/006_agent_audit_log.go
Normal file
|
|
@ -0,0 +1,46 @@
|
||||||
|
package migrations
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"github.com/pressly/goose/v3"
|
||||||
|
)
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
goose.AddMigrationContext(up006AgentAuditLog, down006AgentAuditLog)
|
||||||
|
}
|
||||||
|
|
||||||
|
func up006AgentAuditLog(ctx context.Context, tx *sql.Tx) error {
|
||||||
|
stmts := []string{
|
||||||
|
`CREATE TABLE IF NOT EXISTS agent_audit_log (
|
||||||
|
id BLOB PRIMARY KEY,
|
||||||
|
agent_id TEXT NOT NULL,
|
||||||
|
session_key TEXT NOT NULL DEFAULT '',
|
||||||
|
action TEXT NOT NULL,
|
||||||
|
target TEXT NOT NULL DEFAULT '',
|
||||||
|
input TEXT,
|
||||||
|
output TEXT,
|
||||||
|
duration_ms INTEGER,
|
||||||
|
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||||
|
)`,
|
||||||
|
`CREATE INDEX IF NOT EXISTS idx_audit_agent_time ON agent_audit_log(agent_id, created_at DESC)`,
|
||||||
|
`CREATE INDEX IF NOT EXISTS idx_audit_action ON agent_audit_log(action)`,
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, s := range stmts {
|
||||||
|
if _, err := tx.ExecContext(ctx, s); err != nil {
|
||||||
|
return fmt.Errorf("006_agent_audit_log up: %w\nSQL: %s", err, s)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func down006AgentAuditLog(ctx context.Context, tx *sql.Tx) error {
|
||||||
|
_, err := tx.ExecContext(ctx, `DROP TABLE IF EXISTS agent_audit_log`)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("006_agent_audit_log down: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
368
pkg/memory/sqlc/agent_audit_log.sql.go
Normal file
368
pkg/memory/sqlc/agent_audit_log.sql.go
Normal file
|
|
@ -0,0 +1,368 @@
|
||||||
|
// Code generated by sqlc. DO NOT EDIT.
|
||||||
|
// versions:
|
||||||
|
// sqlc v1.30.0
|
||||||
|
// source: agent_audit_log.sql
|
||||||
|
|
||||||
|
package sqlc
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/ids"
|
||||||
|
)
|
||||||
|
|
||||||
|
const CountAuditEntries = `-- name: CountAuditEntries :one
|
||||||
|
SELECT COUNT(*)
|
||||||
|
FROM agent_audit_log
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
`
|
||||||
|
|
||||||
|
type CountAuditEntriesParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// CountAuditEntries
|
||||||
|
//
|
||||||
|
// SELECT COUNT(*)
|
||||||
|
// FROM agent_audit_log
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
func (q *Queries) CountAuditEntries(ctx context.Context, arg CountAuditEntriesParams) (int64, error) {
|
||||||
|
row := q.db.QueryRowContext(ctx, CountAuditEntries, arg.AgentID)
|
||||||
|
var count int64
|
||||||
|
err := row.Scan(&count)
|
||||||
|
return count, err
|
||||||
|
}
|
||||||
|
|
||||||
|
const CountAuditEntriesByAction = `-- name: CountAuditEntriesByAction :one
|
||||||
|
SELECT COUNT(*)
|
||||||
|
FROM agent_audit_log
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
AND action = ?2
|
||||||
|
`
|
||||||
|
|
||||||
|
type CountAuditEntriesByActionParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Action string `json:"action"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// CountAuditEntriesByAction
|
||||||
|
//
|
||||||
|
// SELECT COUNT(*)
|
||||||
|
// FROM agent_audit_log
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND action = ?2
|
||||||
|
func (q *Queries) CountAuditEntriesByAction(ctx context.Context, arg CountAuditEntriesByActionParams) (int64, error) {
|
||||||
|
row := q.db.QueryRowContext(ctx, CountAuditEntriesByAction, arg.AgentID, arg.Action)
|
||||||
|
var count int64
|
||||||
|
err := row.Scan(&count)
|
||||||
|
return count, err
|
||||||
|
}
|
||||||
|
|
||||||
|
const InsertAuditEntry = `-- name: InsertAuditEntry :exec
|
||||||
|
INSERT INTO agent_audit_log (
|
||||||
|
id,
|
||||||
|
agent_id,
|
||||||
|
session_key,
|
||||||
|
action,
|
||||||
|
target,
|
||||||
|
input,
|
||||||
|
output,
|
||||||
|
duration_ms,
|
||||||
|
created_at
|
||||||
|
)
|
||||||
|
VALUES (
|
||||||
|
?1,
|
||||||
|
?2,
|
||||||
|
?3,
|
||||||
|
?4,
|
||||||
|
?5,
|
||||||
|
?6,
|
||||||
|
?7,
|
||||||
|
?8,
|
||||||
|
datetime('now')
|
||||||
|
)
|
||||||
|
`
|
||||||
|
|
||||||
|
type InsertAuditEntryParams struct {
|
||||||
|
ID ids.UUID `json:"id"`
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
SessionKey string `json:"session_key"`
|
||||||
|
Action string `json:"action"`
|
||||||
|
Target string `json:"target"`
|
||||||
|
Input *string `json:"input"`
|
||||||
|
Output *string `json:"output"`
|
||||||
|
DurationMs *int64 `json:"duration_ms"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Agent Audit Log queries
|
||||||
|
//
|
||||||
|
// INSERT INTO agent_audit_log (
|
||||||
|
// id,
|
||||||
|
// agent_id,
|
||||||
|
// session_key,
|
||||||
|
// action,
|
||||||
|
// target,
|
||||||
|
// input,
|
||||||
|
// output,
|
||||||
|
// duration_ms,
|
||||||
|
// created_at
|
||||||
|
// )
|
||||||
|
// VALUES (
|
||||||
|
// ?1,
|
||||||
|
// ?2,
|
||||||
|
// ?3,
|
||||||
|
// ?4,
|
||||||
|
// ?5,
|
||||||
|
// ?6,
|
||||||
|
// ?7,
|
||||||
|
// ?8,
|
||||||
|
// datetime('now')
|
||||||
|
// )
|
||||||
|
func (q *Queries) InsertAuditEntry(ctx context.Context, arg InsertAuditEntryParams) error {
|
||||||
|
_, err := q.db.ExecContext(ctx, InsertAuditEntry,
|
||||||
|
arg.ID,
|
||||||
|
arg.AgentID,
|
||||||
|
arg.SessionKey,
|
||||||
|
arg.Action,
|
||||||
|
arg.Target,
|
||||||
|
arg.Input,
|
||||||
|
arg.Output,
|
||||||
|
arg.DurationMs,
|
||||||
|
)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
const ListAuditEntries = `-- name: ListAuditEntries :many
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
session_key,
|
||||||
|
action,
|
||||||
|
target,
|
||||||
|
input,
|
||||||
|
output,
|
||||||
|
duration_ms,
|
||||||
|
created_at
|
||||||
|
FROM agent_audit_log
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
ORDER BY created_at DESC
|
||||||
|
LIMIT ?2
|
||||||
|
`
|
||||||
|
|
||||||
|
type ListAuditEntriesParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Lim int64 `json:"lim"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListAuditEntries
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// session_key,
|
||||||
|
// action,
|
||||||
|
// target,
|
||||||
|
// input,
|
||||||
|
// output,
|
||||||
|
// duration_ms,
|
||||||
|
// created_at
|
||||||
|
// FROM agent_audit_log
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// ORDER BY created_at DESC
|
||||||
|
// LIMIT ?2
|
||||||
|
func (q *Queries) ListAuditEntries(ctx context.Context, arg ListAuditEntriesParams) ([]AgentAuditLog, error) {
|
||||||
|
rows, err := q.db.QueryContext(ctx, ListAuditEntries, arg.AgentID, arg.Lim)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
items := []AgentAuditLog{}
|
||||||
|
for rows.Next() {
|
||||||
|
var i AgentAuditLog
|
||||||
|
if err := rows.Scan(
|
||||||
|
&i.ID,
|
||||||
|
&i.AgentID,
|
||||||
|
&i.SessionKey,
|
||||||
|
&i.Action,
|
||||||
|
&i.Target,
|
||||||
|
&i.Input,
|
||||||
|
&i.Output,
|
||||||
|
&i.DurationMs,
|
||||||
|
&i.CreatedAt,
|
||||||
|
); 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 ListAuditEntriesByAction = `-- name: ListAuditEntriesByAction :many
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
session_key,
|
||||||
|
action,
|
||||||
|
target,
|
||||||
|
input,
|
||||||
|
output,
|
||||||
|
duration_ms,
|
||||||
|
created_at
|
||||||
|
FROM agent_audit_log
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
AND action = ?2
|
||||||
|
ORDER BY created_at DESC
|
||||||
|
LIMIT ?3
|
||||||
|
`
|
||||||
|
|
||||||
|
type ListAuditEntriesByActionParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Action string `json:"action"`
|
||||||
|
Lim int64 `json:"lim"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListAuditEntriesByAction
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// session_key,
|
||||||
|
// action,
|
||||||
|
// target,
|
||||||
|
// input,
|
||||||
|
// output,
|
||||||
|
// duration_ms,
|
||||||
|
// created_at
|
||||||
|
// FROM agent_audit_log
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND action = ?2
|
||||||
|
// ORDER BY created_at DESC
|
||||||
|
// LIMIT ?3
|
||||||
|
func (q *Queries) ListAuditEntriesByAction(ctx context.Context, arg ListAuditEntriesByActionParams) ([]AgentAuditLog, error) {
|
||||||
|
rows, err := q.db.QueryContext(ctx, ListAuditEntriesByAction, arg.AgentID, arg.Action, arg.Lim)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
items := []AgentAuditLog{}
|
||||||
|
for rows.Next() {
|
||||||
|
var i AgentAuditLog
|
||||||
|
if err := rows.Scan(
|
||||||
|
&i.ID,
|
||||||
|
&i.AgentID,
|
||||||
|
&i.SessionKey,
|
||||||
|
&i.Action,
|
||||||
|
&i.Target,
|
||||||
|
&i.Input,
|
||||||
|
&i.Output,
|
||||||
|
&i.DurationMs,
|
||||||
|
&i.CreatedAt,
|
||||||
|
); 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 ListAuditEntriesBySession = `-- name: ListAuditEntriesBySession :many
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
session_key,
|
||||||
|
action,
|
||||||
|
target,
|
||||||
|
input,
|
||||||
|
output,
|
||||||
|
duration_ms,
|
||||||
|
created_at
|
||||||
|
FROM agent_audit_log
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
AND session_key = ?2
|
||||||
|
ORDER BY created_at DESC
|
||||||
|
LIMIT ?3
|
||||||
|
`
|
||||||
|
|
||||||
|
type ListAuditEntriesBySessionParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
SessionKey string `json:"session_key"`
|
||||||
|
Lim int64 `json:"lim"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListAuditEntriesBySession
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// session_key,
|
||||||
|
// action,
|
||||||
|
// target,
|
||||||
|
// input,
|
||||||
|
// output,
|
||||||
|
// duration_ms,
|
||||||
|
// created_at
|
||||||
|
// FROM agent_audit_log
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND session_key = ?2
|
||||||
|
// ORDER BY created_at DESC
|
||||||
|
// LIMIT ?3
|
||||||
|
func (q *Queries) ListAuditEntriesBySession(ctx context.Context, arg ListAuditEntriesBySessionParams) ([]AgentAuditLog, error) {
|
||||||
|
rows, err := q.db.QueryContext(ctx, ListAuditEntriesBySession, arg.AgentID, arg.SessionKey, arg.Lim)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
items := []AgentAuditLog{}
|
||||||
|
for rows.Next() {
|
||||||
|
var i AgentAuditLog
|
||||||
|
if err := rows.Scan(
|
||||||
|
&i.ID,
|
||||||
|
&i.AgentID,
|
||||||
|
&i.SessionKey,
|
||||||
|
&i.Action,
|
||||||
|
&i.Target,
|
||||||
|
&i.Input,
|
||||||
|
&i.Output,
|
||||||
|
&i.DurationMs,
|
||||||
|
&i.CreatedAt,
|
||||||
|
); 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 PruneOldAuditEntries = `-- name: PruneOldAuditEntries :exec
|
||||||
|
DELETE FROM agent_audit_log
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
AND created_at < ?2
|
||||||
|
`
|
||||||
|
|
||||||
|
type PruneOldAuditEntriesParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Before time.Time `json:"before"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// PruneOldAuditEntries
|
||||||
|
//
|
||||||
|
// DELETE FROM agent_audit_log
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND created_at < ?2
|
||||||
|
func (q *Queries) PruneOldAuditEntries(ctx context.Context, arg PruneOldAuditEntriesParams) error {
|
||||||
|
_, err := q.db.ExecContext(ctx, PruneOldAuditEntries, arg.AgentID, arg.Before)
|
||||||
|
return err
|
||||||
|
}
|
||||||
305
pkg/memory/sqlc/agent_documents.sql.go
Normal file
305
pkg/memory/sqlc/agent_documents.sql.go
Normal file
|
|
@ -0,0 +1,305 @@
|
||||||
|
// Code generated by sqlc. DO NOT EDIT.
|
||||||
|
// versions:
|
||||||
|
// sqlc v1.30.0
|
||||||
|
// source: agent_documents.sql
|
||||||
|
|
||||||
|
package sqlc
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/ids"
|
||||||
|
)
|
||||||
|
|
||||||
|
const DeleteDocument = `-- name: DeleteDocument :exec
|
||||||
|
DELETE FROM agent_documents
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
AND name = ?2
|
||||||
|
`
|
||||||
|
|
||||||
|
type DeleteDocumentParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteDocument
|
||||||
|
//
|
||||||
|
// DELETE FROM agent_documents
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND name = ?2
|
||||||
|
func (q *Queries) DeleteDocument(ctx context.Context, arg DeleteDocumentParams) error {
|
||||||
|
_, err := q.db.ExecContext(ctx, DeleteDocument, arg.AgentID, arg.Name)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
const GetDocument = `-- name: GetDocument :one
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
name,
|
||||||
|
category,
|
||||||
|
content,
|
||||||
|
version,
|
||||||
|
is_active,
|
||||||
|
created_at,
|
||||||
|
updated_at
|
||||||
|
FROM agent_documents
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
AND name = ?2
|
||||||
|
`
|
||||||
|
|
||||||
|
type GetDocumentParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Agent Documents queries
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// name,
|
||||||
|
// category,
|
||||||
|
// content,
|
||||||
|
// version,
|
||||||
|
// is_active,
|
||||||
|
// created_at,
|
||||||
|
// updated_at
|
||||||
|
// FROM agent_documents
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND name = ?2
|
||||||
|
func (q *Queries) GetDocument(ctx context.Context, arg GetDocumentParams) (AgentDocument, error) {
|
||||||
|
row := q.db.QueryRowContext(ctx, GetDocument, arg.AgentID, arg.Name)
|
||||||
|
var i AgentDocument
|
||||||
|
err := row.Scan(
|
||||||
|
&i.ID,
|
||||||
|
&i.AgentID,
|
||||||
|
&i.Name,
|
||||||
|
&i.Category,
|
||||||
|
&i.Content,
|
||||||
|
&i.Version,
|
||||||
|
&i.IsActive,
|
||||||
|
&i.CreatedAt,
|
||||||
|
&i.UpdatedAt,
|
||||||
|
)
|
||||||
|
return i, err
|
||||||
|
}
|
||||||
|
|
||||||
|
const ListAllDocuments = `-- name: ListAllDocuments :many
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
name,
|
||||||
|
category,
|
||||||
|
content,
|
||||||
|
version,
|
||||||
|
is_active,
|
||||||
|
created_at,
|
||||||
|
updated_at
|
||||||
|
FROM agent_documents
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
AND is_active = 1
|
||||||
|
ORDER BY category,
|
||||||
|
name
|
||||||
|
`
|
||||||
|
|
||||||
|
type ListAllDocumentsParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListAllDocuments
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// name,
|
||||||
|
// category,
|
||||||
|
// content,
|
||||||
|
// version,
|
||||||
|
// is_active,
|
||||||
|
// created_at,
|
||||||
|
// updated_at
|
||||||
|
// FROM agent_documents
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND is_active = 1
|
||||||
|
// ORDER BY category,
|
||||||
|
// name
|
||||||
|
func (q *Queries) ListAllDocuments(ctx context.Context, arg ListAllDocumentsParams) ([]AgentDocument, error) {
|
||||||
|
rows, err := q.db.QueryContext(ctx, ListAllDocuments, arg.AgentID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
items := []AgentDocument{}
|
||||||
|
for rows.Next() {
|
||||||
|
var i AgentDocument
|
||||||
|
if err := rows.Scan(
|
||||||
|
&i.ID,
|
||||||
|
&i.AgentID,
|
||||||
|
&i.Name,
|
||||||
|
&i.Category,
|
||||||
|
&i.Content,
|
||||||
|
&i.Version,
|
||||||
|
&i.IsActive,
|
||||||
|
&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 ListDocumentsByCategory = `-- name: ListDocumentsByCategory :many
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
name,
|
||||||
|
category,
|
||||||
|
content,
|
||||||
|
version,
|
||||||
|
is_active,
|
||||||
|
created_at,
|
||||||
|
updated_at
|
||||||
|
FROM agent_documents
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
AND category = ?2
|
||||||
|
AND is_active = 1
|
||||||
|
ORDER BY name
|
||||||
|
`
|
||||||
|
|
||||||
|
type ListDocumentsByCategoryParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Category string `json:"category"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListDocumentsByCategory
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// name,
|
||||||
|
// category,
|
||||||
|
// content,
|
||||||
|
// version,
|
||||||
|
// is_active,
|
||||||
|
// created_at,
|
||||||
|
// updated_at
|
||||||
|
// FROM agent_documents
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND category = ?2
|
||||||
|
// AND is_active = 1
|
||||||
|
// ORDER BY name
|
||||||
|
func (q *Queries) ListDocumentsByCategory(ctx context.Context, arg ListDocumentsByCategoryParams) ([]AgentDocument, error) {
|
||||||
|
rows, err := q.db.QueryContext(ctx, ListDocumentsByCategory, arg.AgentID, arg.Category)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
items := []AgentDocument{}
|
||||||
|
for rows.Next() {
|
||||||
|
var i AgentDocument
|
||||||
|
if err := rows.Scan(
|
||||||
|
&i.ID,
|
||||||
|
&i.AgentID,
|
||||||
|
&i.Name,
|
||||||
|
&i.Category,
|
||||||
|
&i.Content,
|
||||||
|
&i.Version,
|
||||||
|
&i.IsActive,
|
||||||
|
&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 UpsertDocument = `-- name: UpsertDocument :exec
|
||||||
|
INSERT INTO agent_documents (
|
||||||
|
id,
|
||||||
|
agent_id,
|
||||||
|
name,
|
||||||
|
category,
|
||||||
|
content,
|
||||||
|
version,
|
||||||
|
is_active,
|
||||||
|
created_at,
|
||||||
|
updated_at
|
||||||
|
)
|
||||||
|
VALUES (
|
||||||
|
?1,
|
||||||
|
?2,
|
||||||
|
?3,
|
||||||
|
?4,
|
||||||
|
?5,
|
||||||
|
1,
|
||||||
|
1,
|
||||||
|
datetime('now'),
|
||||||
|
datetime('now')
|
||||||
|
) ON CONFLICT (agent_id, name) DO
|
||||||
|
UPDATE
|
||||||
|
SET content = excluded.content,
|
||||||
|
category = excluded.category,
|
||||||
|
version = agent_documents.version + 1,
|
||||||
|
is_active = 1,
|
||||||
|
updated_at = datetime('now')
|
||||||
|
`
|
||||||
|
|
||||||
|
type UpsertDocumentParams struct {
|
||||||
|
ID ids.UUID `json:"id"`
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Category string `json:"category"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpsertDocument
|
||||||
|
//
|
||||||
|
// INSERT INTO agent_documents (
|
||||||
|
// id,
|
||||||
|
// agent_id,
|
||||||
|
// name,
|
||||||
|
// category,
|
||||||
|
// content,
|
||||||
|
// version,
|
||||||
|
// is_active,
|
||||||
|
// created_at,
|
||||||
|
// updated_at
|
||||||
|
// )
|
||||||
|
// VALUES (
|
||||||
|
// ?1,
|
||||||
|
// ?2,
|
||||||
|
// ?3,
|
||||||
|
// ?4,
|
||||||
|
// ?5,
|
||||||
|
// 1,
|
||||||
|
// 1,
|
||||||
|
// datetime('now'),
|
||||||
|
// datetime('now')
|
||||||
|
// ) ON CONFLICT (agent_id, name) DO
|
||||||
|
// UPDATE
|
||||||
|
// SET content = excluded.content,
|
||||||
|
// category = excluded.category,
|
||||||
|
// version = agent_documents.version + 1,
|
||||||
|
// is_active = 1,
|
||||||
|
// updated_at = datetime('now')
|
||||||
|
func (q *Queries) UpsertDocument(ctx context.Context, arg UpsertDocumentParams) error {
|
||||||
|
_, err := q.db.ExecContext(ctx, UpsertDocument,
|
||||||
|
arg.ID,
|
||||||
|
arg.AgentID,
|
||||||
|
arg.Name,
|
||||||
|
arg.Category,
|
||||||
|
arg.Content,
|
||||||
|
)
|
||||||
|
return err
|
||||||
|
}
|
||||||
160
pkg/memory/sqlc/agent_kv.sql.go
Normal file
160
pkg/memory/sqlc/agent_kv.sql.go
Normal file
|
|
@ -0,0 +1,160 @@
|
||||||
|
// Code generated by sqlc. DO NOT EDIT.
|
||||||
|
// versions:
|
||||||
|
// sqlc v1.30.0
|
||||||
|
// source: agent_kv.sql
|
||||||
|
|
||||||
|
package sqlc
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
)
|
||||||
|
|
||||||
|
const DeleteKV = `-- name: DeleteKV :exec
|
||||||
|
DELETE FROM agent_kv
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
AND key = ?2
|
||||||
|
`
|
||||||
|
|
||||||
|
type DeleteKVParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Key string `json:"key"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteKV
|
||||||
|
//
|
||||||
|
// DELETE FROM agent_kv
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND key = ?2
|
||||||
|
func (q *Queries) DeleteKV(ctx context.Context, arg DeleteKVParams) error {
|
||||||
|
_, err := q.db.ExecContext(ctx, DeleteKV, arg.AgentID, arg.Key)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
const GetKV = `-- name: GetKV :one
|
||||||
|
SELECT agent_id,
|
||||||
|
key,
|
||||||
|
value,
|
||||||
|
updated_at
|
||||||
|
FROM agent_kv
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
AND key = ?2
|
||||||
|
`
|
||||||
|
|
||||||
|
type GetKVParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Key string `json:"key"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Agent KV Store queries
|
||||||
|
//
|
||||||
|
// SELECT agent_id,
|
||||||
|
// key,
|
||||||
|
// value,
|
||||||
|
// updated_at
|
||||||
|
// FROM agent_kv
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND key = ?2
|
||||||
|
func (q *Queries) GetKV(ctx context.Context, arg GetKVParams) (AgentKv, error) {
|
||||||
|
row := q.db.QueryRowContext(ctx, GetKV, arg.AgentID, arg.Key)
|
||||||
|
var i AgentKv
|
||||||
|
err := row.Scan(
|
||||||
|
&i.AgentID,
|
||||||
|
&i.Key,
|
||||||
|
&i.Value,
|
||||||
|
&i.UpdatedAt,
|
||||||
|
)
|
||||||
|
return i, err
|
||||||
|
}
|
||||||
|
|
||||||
|
const ListKVByPrefix = `-- name: ListKVByPrefix :many
|
||||||
|
SELECT agent_id,
|
||||||
|
key,
|
||||||
|
value,
|
||||||
|
updated_at
|
||||||
|
FROM agent_kv
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
AND key LIKE ?2 || '%'
|
||||||
|
ORDER BY key
|
||||||
|
LIMIT ?3
|
||||||
|
`
|
||||||
|
|
||||||
|
type ListKVByPrefixParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Prefix *string `json:"prefix"`
|
||||||
|
Lim int64 `json:"lim"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListKVByPrefix
|
||||||
|
//
|
||||||
|
// SELECT agent_id,
|
||||||
|
// key,
|
||||||
|
// value,
|
||||||
|
// updated_at
|
||||||
|
// FROM agent_kv
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND key LIKE ?2 || '%'
|
||||||
|
// ORDER BY key
|
||||||
|
// LIMIT ?3
|
||||||
|
func (q *Queries) ListKVByPrefix(ctx context.Context, arg ListKVByPrefixParams) ([]AgentKv, error) {
|
||||||
|
rows, err := q.db.QueryContext(ctx, ListKVByPrefix, arg.AgentID, arg.Prefix, arg.Lim)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
items := []AgentKv{}
|
||||||
|
for rows.Next() {
|
||||||
|
var i AgentKv
|
||||||
|
if err := rows.Scan(
|
||||||
|
&i.AgentID,
|
||||||
|
&i.Key,
|
||||||
|
&i.Value,
|
||||||
|
&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 UpsertKV = `-- name: UpsertKV :exec
|
||||||
|
INSERT INTO agent_kv (agent_id, key, value, updated_at)
|
||||||
|
VALUES (
|
||||||
|
?1,
|
||||||
|
?2,
|
||||||
|
?3,
|
||||||
|
datetime('now')
|
||||||
|
) ON CONFLICT (agent_id, key) DO
|
||||||
|
UPDATE
|
||||||
|
SET value = excluded.value,
|
||||||
|
updated_at = excluded.updated_at
|
||||||
|
`
|
||||||
|
|
||||||
|
type UpsertKVParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Key string `json:"key"`
|
||||||
|
Value string `json:"value"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpsertKV
|
||||||
|
//
|
||||||
|
// INSERT INTO agent_kv (agent_id, key, value, updated_at)
|
||||||
|
// VALUES (
|
||||||
|
// ?1,
|
||||||
|
// ?2,
|
||||||
|
// ?3,
|
||||||
|
// datetime('now')
|
||||||
|
// ) ON CONFLICT (agent_id, key) DO
|
||||||
|
// UPDATE
|
||||||
|
// SET value = excluded.value,
|
||||||
|
// updated_at = excluded.updated_at
|
||||||
|
func (q *Queries) UpsertKV(ctx context.Context, arg UpsertKVParams) error {
|
||||||
|
_, err := q.db.ExecContext(ctx, UpsertKV, arg.AgentID, arg.Key, arg.Value)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
@ -11,6 +11,37 @@ import (
|
||||||
"github.com/sipeed/picoclaw/pkg/memory"
|
"github.com/sipeed/picoclaw/pkg/memory"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
type AgentAuditLog struct {
|
||||||
|
ID ids.UUID `json:"id"`
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
SessionKey string `json:"session_key"`
|
||||||
|
Action string `json:"action"`
|
||||||
|
Target string `json:"target"`
|
||||||
|
Input *string `json:"input"`
|
||||||
|
Output *string `json:"output"`
|
||||||
|
DurationMs *int64 `json:"duration_ms"`
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type AgentDocument struct {
|
||||||
|
ID ids.UUID `json:"id"`
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Category string `json:"category"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
Version int64 `json:"version"`
|
||||||
|
IsActive bool `json:"is_active"`
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type AgentKv struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
Key string `json:"key"`
|
||||||
|
Value string `json:"value"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
|
}
|
||||||
|
|
||||||
type ArchivalChunk struct {
|
type ArchivalChunk struct {
|
||||||
ID ids.UUID `json:"id"`
|
ID ids.UUID `json:"id"`
|
||||||
RecallID ids.UUID `json:"recall_id"`
|
RecallID ids.UUID `json:"recall_id"`
|
||||||
|
|
|
||||||
|
|
@ -14,6 +14,19 @@ type Querier interface {
|
||||||
// SELECT COUNT(*)
|
// SELECT COUNT(*)
|
||||||
// FROM archival_chunks
|
// FROM archival_chunks
|
||||||
CountArchivalChunks(ctx context.Context) (int64, error)
|
CountArchivalChunks(ctx context.Context) (int64, error)
|
||||||
|
//CountAuditEntries
|
||||||
|
//
|
||||||
|
// SELECT COUNT(*)
|
||||||
|
// FROM agent_audit_log
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
CountAuditEntries(ctx context.Context, arg CountAuditEntriesParams) (int64, error)
|
||||||
|
//CountAuditEntriesByAction
|
||||||
|
//
|
||||||
|
// SELECT COUNT(*)
|
||||||
|
// FROM agent_audit_log
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND action = ?2
|
||||||
|
CountAuditEntriesByAction(ctx context.Context, arg CountAuditEntriesByActionParams) (int64, error)
|
||||||
//CountRecallItems
|
//CountRecallItems
|
||||||
//
|
//
|
||||||
// SELECT COUNT(*)
|
// SELECT COUNT(*)
|
||||||
|
|
@ -24,11 +37,31 @@ type Querier interface {
|
||||||
// OR ?2 = ''
|
// OR ?2 = ''
|
||||||
// )
|
// )
|
||||||
CountRecallItems(ctx context.Context, arg CountRecallItemsParams) (int64, error)
|
CountRecallItems(ctx context.Context, arg CountRecallItemsParams) (int64, error)
|
||||||
|
//CountSessionMessages
|
||||||
|
//
|
||||||
|
// SELECT COUNT(*)
|
||||||
|
// FROM recall_items
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND session_key = ?2
|
||||||
|
// AND tags = 'session-message'
|
||||||
|
CountSessionMessages(ctx context.Context, arg CountSessionMessagesParams) (int64, error)
|
||||||
//DeleteArchivalChunksByRecall
|
//DeleteArchivalChunksByRecall
|
||||||
//
|
//
|
||||||
// DELETE FROM archival_chunks
|
// DELETE FROM archival_chunks
|
||||||
// WHERE recall_id = ?1
|
// WHERE recall_id = ?1
|
||||||
DeleteArchivalChunksByRecall(ctx context.Context, arg DeleteArchivalChunksByRecallParams) error
|
DeleteArchivalChunksByRecall(ctx context.Context, arg DeleteArchivalChunksByRecallParams) error
|
||||||
|
//DeleteDocument
|
||||||
|
//
|
||||||
|
// DELETE FROM agent_documents
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND name = ?2
|
||||||
|
DeleteDocument(ctx context.Context, arg DeleteDocumentParams) error
|
||||||
|
//DeleteKV
|
||||||
|
//
|
||||||
|
// DELETE FROM agent_kv
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND key = ?2
|
||||||
|
DeleteKV(ctx context.Context, arg DeleteKVParams) error
|
||||||
//DeleteRecallItem
|
//DeleteRecallItem
|
||||||
//
|
//
|
||||||
// DELETE FROM recall_items
|
// DELETE FROM recall_items
|
||||||
|
|
@ -60,6 +93,31 @@ type Querier interface {
|
||||||
// FROM archival_chunks
|
// FROM archival_chunks
|
||||||
// WHERE id IN (/*SLICE:ids*/?)
|
// WHERE id IN (/*SLICE:ids*/?)
|
||||||
GetArchivalChunksByIDs(ctx context.Context, arg GetArchivalChunksByIDsParams) ([]ArchivalChunk, error)
|
GetArchivalChunksByIDs(ctx context.Context, arg GetArchivalChunksByIDsParams) ([]ArchivalChunk, error)
|
||||||
|
// Agent Documents queries
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// name,
|
||||||
|
// category,
|
||||||
|
// content,
|
||||||
|
// version,
|
||||||
|
// is_active,
|
||||||
|
// created_at,
|
||||||
|
// updated_at
|
||||||
|
// FROM agent_documents
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND name = ?2
|
||||||
|
GetDocument(ctx context.Context, arg GetDocumentParams) (AgentDocument, error)
|
||||||
|
// Agent KV Store queries
|
||||||
|
//
|
||||||
|
// SELECT agent_id,
|
||||||
|
// key,
|
||||||
|
// value,
|
||||||
|
// updated_at
|
||||||
|
// FROM agent_kv
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND key = ?2
|
||||||
|
GetKV(ctx context.Context, arg GetKVParams) (AgentKv, error)
|
||||||
//GetRecallItem
|
//GetRecallItem
|
||||||
//
|
//
|
||||||
// SELECT id,
|
// SELECT id,
|
||||||
|
|
@ -127,6 +185,31 @@ type Querier interface {
|
||||||
// datetime('now')
|
// datetime('now')
|
||||||
// )
|
// )
|
||||||
InsertArchivalChunk(ctx context.Context, arg InsertArchivalChunkParams) error
|
InsertArchivalChunk(ctx context.Context, arg InsertArchivalChunkParams) error
|
||||||
|
// Agent Audit Log queries
|
||||||
|
//
|
||||||
|
// INSERT INTO agent_audit_log (
|
||||||
|
// id,
|
||||||
|
// agent_id,
|
||||||
|
// session_key,
|
||||||
|
// action,
|
||||||
|
// target,
|
||||||
|
// input,
|
||||||
|
// output,
|
||||||
|
// duration_ms,
|
||||||
|
// created_at
|
||||||
|
// )
|
||||||
|
// VALUES (
|
||||||
|
// ?1,
|
||||||
|
// ?2,
|
||||||
|
// ?3,
|
||||||
|
// ?4,
|
||||||
|
// ?5,
|
||||||
|
// ?6,
|
||||||
|
// ?7,
|
||||||
|
// ?8,
|
||||||
|
// datetime('now')
|
||||||
|
// )
|
||||||
|
InsertAuditEntry(ctx context.Context, arg InsertAuditEntryParams) error
|
||||||
// Recall Item queries
|
// Recall Item queries
|
||||||
//
|
//
|
||||||
// INSERT INTO recall_items (
|
// INSERT INTO recall_items (
|
||||||
|
|
@ -158,6 +241,37 @@ type Querier interface {
|
||||||
// datetime('now')
|
// datetime('now')
|
||||||
// )
|
// )
|
||||||
InsertRecallItem(ctx context.Context, arg InsertRecallItemParams) error
|
InsertRecallItem(ctx context.Context, arg InsertRecallItemParams) error
|
||||||
|
//InsertSessionMessage
|
||||||
|
//
|
||||||
|
// INSERT INTO recall_items (
|
||||||
|
// id,
|
||||||
|
// agent_id,
|
||||||
|
// session_key,
|
||||||
|
// role,
|
||||||
|
// sector,
|
||||||
|
// importance,
|
||||||
|
// salience,
|
||||||
|
// decay_rate,
|
||||||
|
// content,
|
||||||
|
// tags,
|
||||||
|
// created_at,
|
||||||
|
// updated_at
|
||||||
|
// )
|
||||||
|
// VALUES (
|
||||||
|
// ?1,
|
||||||
|
// ?2,
|
||||||
|
// ?3,
|
||||||
|
// ?4,
|
||||||
|
// 'episodic',
|
||||||
|
// 0.5,
|
||||||
|
// 0.5,
|
||||||
|
// 0.01,
|
||||||
|
// ?5,
|
||||||
|
// 'session-message',
|
||||||
|
// datetime('now'),
|
||||||
|
// datetime('now')
|
||||||
|
// )
|
||||||
|
InsertSessionMessage(ctx context.Context, arg InsertSessionMessageParams) error
|
||||||
// Memory Summary queries
|
// Memory Summary queries
|
||||||
//
|
//
|
||||||
// INSERT INTO memory_summaries (
|
// INSERT INTO memory_summaries (
|
||||||
|
|
@ -193,6 +307,23 @@ type Querier interface {
|
||||||
// ORDER BY created_at DESC
|
// ORDER BY created_at DESC
|
||||||
// LIMIT ?2 OFFSET ?1
|
// LIMIT ?2 OFFSET ?1
|
||||||
ListAllArchivalChunks(ctx context.Context, arg ListAllArchivalChunksParams) ([]ArchivalChunk, error)
|
ListAllArchivalChunks(ctx context.Context, arg ListAllArchivalChunksParams) ([]ArchivalChunk, error)
|
||||||
|
//ListAllDocuments
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// name,
|
||||||
|
// category,
|
||||||
|
// content,
|
||||||
|
// version,
|
||||||
|
// is_active,
|
||||||
|
// created_at,
|
||||||
|
// updated_at
|
||||||
|
// FROM agent_documents
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND is_active = 1
|
||||||
|
// ORDER BY category,
|
||||||
|
// name
|
||||||
|
ListAllDocuments(ctx context.Context, arg ListAllDocumentsParams) ([]AgentDocument, error)
|
||||||
//ListArchivalChunks
|
//ListArchivalChunks
|
||||||
//
|
//
|
||||||
// SELECT id,
|
// SELECT id,
|
||||||
|
|
@ -207,6 +338,85 @@ type Querier interface {
|
||||||
// WHERE recall_id = ?1
|
// WHERE recall_id = ?1
|
||||||
// ORDER BY chunk_index
|
// ORDER BY chunk_index
|
||||||
ListArchivalChunks(ctx context.Context, arg ListArchivalChunksParams) ([]ArchivalChunk, error)
|
ListArchivalChunks(ctx context.Context, arg ListArchivalChunksParams) ([]ArchivalChunk, error)
|
||||||
|
//ListAuditEntries
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// session_key,
|
||||||
|
// action,
|
||||||
|
// target,
|
||||||
|
// input,
|
||||||
|
// output,
|
||||||
|
// duration_ms,
|
||||||
|
// created_at
|
||||||
|
// FROM agent_audit_log
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// ORDER BY created_at DESC
|
||||||
|
// LIMIT ?2
|
||||||
|
ListAuditEntries(ctx context.Context, arg ListAuditEntriesParams) ([]AgentAuditLog, error)
|
||||||
|
//ListAuditEntriesByAction
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// session_key,
|
||||||
|
// action,
|
||||||
|
// target,
|
||||||
|
// input,
|
||||||
|
// output,
|
||||||
|
// duration_ms,
|
||||||
|
// created_at
|
||||||
|
// FROM agent_audit_log
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND action = ?2
|
||||||
|
// ORDER BY created_at DESC
|
||||||
|
// LIMIT ?3
|
||||||
|
ListAuditEntriesByAction(ctx context.Context, arg ListAuditEntriesByActionParams) ([]AgentAuditLog, error)
|
||||||
|
//ListAuditEntriesBySession
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// session_key,
|
||||||
|
// action,
|
||||||
|
// target,
|
||||||
|
// input,
|
||||||
|
// output,
|
||||||
|
// duration_ms,
|
||||||
|
// created_at
|
||||||
|
// FROM agent_audit_log
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND session_key = ?2
|
||||||
|
// ORDER BY created_at DESC
|
||||||
|
// LIMIT ?3
|
||||||
|
ListAuditEntriesBySession(ctx context.Context, arg ListAuditEntriesBySessionParams) ([]AgentAuditLog, error)
|
||||||
|
//ListDocumentsByCategory
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// name,
|
||||||
|
// category,
|
||||||
|
// content,
|
||||||
|
// version,
|
||||||
|
// is_active,
|
||||||
|
// created_at,
|
||||||
|
// updated_at
|
||||||
|
// FROM agent_documents
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND category = ?2
|
||||||
|
// AND is_active = 1
|
||||||
|
// ORDER BY name
|
||||||
|
ListDocumentsByCategory(ctx context.Context, arg ListDocumentsByCategoryParams) ([]AgentDocument, error)
|
||||||
|
//ListKVByPrefix
|
||||||
|
//
|
||||||
|
// SELECT agent_id,
|
||||||
|
// key,
|
||||||
|
// value,
|
||||||
|
// updated_at
|
||||||
|
// FROM agent_kv
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND key LIKE ?2 || '%'
|
||||||
|
// ORDER BY key
|
||||||
|
// LIMIT ?3
|
||||||
|
ListKVByPrefix(ctx context.Context, arg ListKVByPrefixParams) ([]AgentKv, error)
|
||||||
//ListRecallItems
|
//ListRecallItems
|
||||||
//
|
//
|
||||||
// SELECT id,
|
// SELECT id,
|
||||||
|
|
@ -230,6 +440,31 @@ type Querier interface {
|
||||||
// ORDER BY created_at DESC
|
// ORDER BY created_at DESC
|
||||||
// LIMIT ?4 OFFSET ?3
|
// LIMIT ?4 OFFSET ?3
|
||||||
ListRecallItems(ctx context.Context, arg ListRecallItemsParams) ([]RecallItem, error)
|
ListRecallItems(ctx context.Context, arg ListRecallItemsParams) ([]RecallItem, error)
|
||||||
|
//ListSessionMessages
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// session_key,
|
||||||
|
// role,
|
||||||
|
// sector,
|
||||||
|
// importance,
|
||||||
|
// salience,
|
||||||
|
// decay_rate,
|
||||||
|
// content,
|
||||||
|
// tags,
|
||||||
|
// created_at,
|
||||||
|
// updated_at
|
||||||
|
// FROM recall_items
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND session_key = ?2
|
||||||
|
// AND tags = 'session-message'
|
||||||
|
// AND (
|
||||||
|
// role = ?3
|
||||||
|
// OR ?3 = ''
|
||||||
|
// )
|
||||||
|
// ORDER BY created_at ASC
|
||||||
|
// LIMIT ?4
|
||||||
|
ListSessionMessages(ctx context.Context, arg ListSessionMessagesParams) ([]RecallItem, error)
|
||||||
//ListSummaries
|
//ListSummaries
|
||||||
//
|
//
|
||||||
// SELECT id,
|
// SELECT id,
|
||||||
|
|
@ -248,6 +483,12 @@ type Querier interface {
|
||||||
// ORDER BY created_at DESC
|
// ORDER BY created_at DESC
|
||||||
// LIMIT ?3
|
// LIMIT ?3
|
||||||
ListSummaries(ctx context.Context, arg ListSummariesParams) ([]MemorySummary, error)
|
ListSummaries(ctx context.Context, arg ListSummariesParams) ([]MemorySummary, error)
|
||||||
|
//PruneOldAuditEntries
|
||||||
|
//
|
||||||
|
// DELETE FROM agent_audit_log
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND created_at < ?2
|
||||||
|
PruneOldAuditEntries(ctx context.Context, arg PruneOldAuditEntriesParams) error
|
||||||
//SearchRecallByKeyword
|
//SearchRecallByKeyword
|
||||||
//
|
//
|
||||||
// SELECT ri.id,
|
// SELECT ri.id,
|
||||||
|
|
@ -281,6 +522,50 @@ type Querier interface {
|
||||||
// updated_at = datetime('now')
|
// updated_at = datetime('now')
|
||||||
// WHERE id = ?8
|
// WHERE id = ?8
|
||||||
UpdateRecallItem(ctx context.Context, arg UpdateRecallItemParams) error
|
UpdateRecallItem(ctx context.Context, arg UpdateRecallItemParams) error
|
||||||
|
//UpsertDocument
|
||||||
|
//
|
||||||
|
// INSERT INTO agent_documents (
|
||||||
|
// id,
|
||||||
|
// agent_id,
|
||||||
|
// name,
|
||||||
|
// category,
|
||||||
|
// content,
|
||||||
|
// version,
|
||||||
|
// is_active,
|
||||||
|
// created_at,
|
||||||
|
// updated_at
|
||||||
|
// )
|
||||||
|
// VALUES (
|
||||||
|
// ?1,
|
||||||
|
// ?2,
|
||||||
|
// ?3,
|
||||||
|
// ?4,
|
||||||
|
// ?5,
|
||||||
|
// 1,
|
||||||
|
// 1,
|
||||||
|
// datetime('now'),
|
||||||
|
// datetime('now')
|
||||||
|
// ) ON CONFLICT (agent_id, name) DO
|
||||||
|
// UPDATE
|
||||||
|
// SET content = excluded.content,
|
||||||
|
// category = excluded.category,
|
||||||
|
// version = agent_documents.version + 1,
|
||||||
|
// is_active = 1,
|
||||||
|
// updated_at = datetime('now')
|
||||||
|
UpsertDocument(ctx context.Context, arg UpsertDocumentParams) error
|
||||||
|
//UpsertKV
|
||||||
|
//
|
||||||
|
// INSERT INTO agent_kv (agent_id, key, value, updated_at)
|
||||||
|
// VALUES (
|
||||||
|
// ?1,
|
||||||
|
// ?2,
|
||||||
|
// ?3,
|
||||||
|
// datetime('now')
|
||||||
|
// ) ON CONFLICT (agent_id, key) DO
|
||||||
|
// UPDATE
|
||||||
|
// SET value = excluded.value,
|
||||||
|
// updated_at = excluded.updated_at
|
||||||
|
UpsertKV(ctx context.Context, arg UpsertKVParams) error
|
||||||
//UpsertWorkingContext
|
//UpsertWorkingContext
|
||||||
//
|
//
|
||||||
// INSERT INTO working_context (agent_id, session_key, content, updated_at)
|
// INSERT INTO working_context (agent_id, session_key, content, updated_at)
|
||||||
|
|
|
||||||
81
pkg/memory/sqlc/queries/agent_audit_log.sql
Normal file
81
pkg/memory/sqlc/queries/agent_audit_log.sql
Normal file
|
|
@ -0,0 +1,81 @@
|
||||||
|
-- Agent Audit Log queries
|
||||||
|
-- name: InsertAuditEntry :exec
|
||||||
|
INSERT INTO agent_audit_log (
|
||||||
|
id,
|
||||||
|
agent_id,
|
||||||
|
session_key,
|
||||||
|
action,
|
||||||
|
target,
|
||||||
|
input,
|
||||||
|
output,
|
||||||
|
duration_ms,
|
||||||
|
created_at
|
||||||
|
)
|
||||||
|
VALUES (
|
||||||
|
sqlc.arg(id),
|
||||||
|
sqlc.arg(agent_id),
|
||||||
|
sqlc.arg(session_key),
|
||||||
|
sqlc.arg(action),
|
||||||
|
sqlc.arg(target),
|
||||||
|
sqlc.arg(input),
|
||||||
|
sqlc.arg(output),
|
||||||
|
sqlc.arg(duration_ms),
|
||||||
|
datetime('now')
|
||||||
|
);
|
||||||
|
-- name: ListAuditEntries :many
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
session_key,
|
||||||
|
action,
|
||||||
|
target,
|
||||||
|
input,
|
||||||
|
output,
|
||||||
|
duration_ms,
|
||||||
|
created_at
|
||||||
|
FROM agent_audit_log
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
ORDER BY created_at DESC
|
||||||
|
LIMIT sqlc.arg(lim);
|
||||||
|
-- name: ListAuditEntriesByAction :many
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
session_key,
|
||||||
|
action,
|
||||||
|
target,
|
||||||
|
input,
|
||||||
|
output,
|
||||||
|
duration_ms,
|
||||||
|
created_at
|
||||||
|
FROM agent_audit_log
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
AND action = sqlc.arg(action)
|
||||||
|
ORDER BY created_at DESC
|
||||||
|
LIMIT sqlc.arg(lim);
|
||||||
|
-- name: CountAuditEntries :one
|
||||||
|
SELECT COUNT(*)
|
||||||
|
FROM agent_audit_log
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id);
|
||||||
|
-- name: ListAuditEntriesBySession :many
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
session_key,
|
||||||
|
action,
|
||||||
|
target,
|
||||||
|
input,
|
||||||
|
output,
|
||||||
|
duration_ms,
|
||||||
|
created_at
|
||||||
|
FROM agent_audit_log
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
AND session_key = sqlc.arg(session_key)
|
||||||
|
ORDER BY created_at DESC
|
||||||
|
LIMIT sqlc.arg(lim);
|
||||||
|
-- name: PruneOldAuditEntries :exec
|
||||||
|
DELETE FROM agent_audit_log
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
AND created_at < sqlc.arg(before);
|
||||||
|
-- name: CountAuditEntriesByAction :one
|
||||||
|
SELECT COUNT(*)
|
||||||
|
FROM agent_audit_log
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
AND action = sqlc.arg(action);
|
||||||
77
pkg/memory/sqlc/queries/agent_documents.sql
Normal file
77
pkg/memory/sqlc/queries/agent_documents.sql
Normal file
|
|
@ -0,0 +1,77 @@
|
||||||
|
-- Agent Documents queries
|
||||||
|
-- name: GetDocument :one
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
name,
|
||||||
|
category,
|
||||||
|
content,
|
||||||
|
version,
|
||||||
|
is_active,
|
||||||
|
created_at,
|
||||||
|
updated_at
|
||||||
|
FROM agent_documents
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
AND name = sqlc.arg(name);
|
||||||
|
-- name: UpsertDocument :exec
|
||||||
|
INSERT INTO agent_documents (
|
||||||
|
id,
|
||||||
|
agent_id,
|
||||||
|
name,
|
||||||
|
category,
|
||||||
|
content,
|
||||||
|
version,
|
||||||
|
is_active,
|
||||||
|
created_at,
|
||||||
|
updated_at
|
||||||
|
)
|
||||||
|
VALUES (
|
||||||
|
sqlc.arg(id),
|
||||||
|
sqlc.arg(agent_id),
|
||||||
|
sqlc.arg(name),
|
||||||
|
sqlc.arg(category),
|
||||||
|
sqlc.arg(content),
|
||||||
|
1,
|
||||||
|
1,
|
||||||
|
datetime('now'),
|
||||||
|
datetime('now')
|
||||||
|
) ON CONFLICT (agent_id, name) DO
|
||||||
|
UPDATE
|
||||||
|
SET content = excluded.content,
|
||||||
|
category = excluded.category,
|
||||||
|
version = agent_documents.version + 1,
|
||||||
|
is_active = 1,
|
||||||
|
updated_at = datetime('now');
|
||||||
|
-- name: ListDocumentsByCategory :many
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
name,
|
||||||
|
category,
|
||||||
|
content,
|
||||||
|
version,
|
||||||
|
is_active,
|
||||||
|
created_at,
|
||||||
|
updated_at
|
||||||
|
FROM agent_documents
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
AND category = sqlc.arg(category)
|
||||||
|
AND is_active = 1
|
||||||
|
ORDER BY name;
|
||||||
|
-- name: DeleteDocument :exec
|
||||||
|
DELETE FROM agent_documents
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
AND name = sqlc.arg(name);
|
||||||
|
-- name: ListAllDocuments :many
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
name,
|
||||||
|
category,
|
||||||
|
content,
|
||||||
|
version,
|
||||||
|
is_active,
|
||||||
|
created_at,
|
||||||
|
updated_at
|
||||||
|
FROM agent_documents
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
AND is_active = 1
|
||||||
|
ORDER BY category,
|
||||||
|
name;
|
||||||
34
pkg/memory/sqlc/queries/agent_kv.sql
Normal file
34
pkg/memory/sqlc/queries/agent_kv.sql
Normal file
|
|
@ -0,0 +1,34 @@
|
||||||
|
-- Agent KV Store queries
|
||||||
|
-- name: GetKV :one
|
||||||
|
SELECT agent_id,
|
||||||
|
key,
|
||||||
|
value,
|
||||||
|
updated_at
|
||||||
|
FROM agent_kv
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
AND key = sqlc.arg(key);
|
||||||
|
-- name: UpsertKV :exec
|
||||||
|
INSERT INTO agent_kv (agent_id, key, value, updated_at)
|
||||||
|
VALUES (
|
||||||
|
sqlc.arg(agent_id),
|
||||||
|
sqlc.arg(key),
|
||||||
|
sqlc.arg(value),
|
||||||
|
datetime('now')
|
||||||
|
) ON CONFLICT (agent_id, key) DO
|
||||||
|
UPDATE
|
||||||
|
SET value = excluded.value,
|
||||||
|
updated_at = excluded.updated_at;
|
||||||
|
-- name: DeleteKV :exec
|
||||||
|
DELETE FROM agent_kv
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
AND key = sqlc.arg(key);
|
||||||
|
-- name: ListKVByPrefix :many
|
||||||
|
SELECT agent_id,
|
||||||
|
key,
|
||||||
|
value,
|
||||||
|
updated_at
|
||||||
|
FROM agent_kv
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
AND key LIKE sqlc.arg(prefix) || '%'
|
||||||
|
ORDER BY key
|
||||||
|
LIMIT sqlc.arg(lim);
|
||||||
|
|
@ -118,4 +118,62 @@ SELECT id,
|
||||||
created_at,
|
created_at,
|
||||||
updated_at
|
updated_at
|
||||||
FROM recall_items
|
FROM recall_items
|
||||||
WHERE id IN (sqlc.slice('ids'));
|
WHERE id IN (sqlc.slice('ids'));
|
||||||
|
-- name: InsertSessionMessage :exec
|
||||||
|
INSERT INTO recall_items (
|
||||||
|
id,
|
||||||
|
agent_id,
|
||||||
|
session_key,
|
||||||
|
role,
|
||||||
|
sector,
|
||||||
|
importance,
|
||||||
|
salience,
|
||||||
|
decay_rate,
|
||||||
|
content,
|
||||||
|
tags,
|
||||||
|
created_at,
|
||||||
|
updated_at
|
||||||
|
)
|
||||||
|
VALUES (
|
||||||
|
sqlc.arg(id),
|
||||||
|
sqlc.arg(agent_id),
|
||||||
|
sqlc.arg(session_key),
|
||||||
|
sqlc.arg(role),
|
||||||
|
'episodic',
|
||||||
|
0.5,
|
||||||
|
0.5,
|
||||||
|
0.01,
|
||||||
|
sqlc.arg(content),
|
||||||
|
'session-message',
|
||||||
|
datetime('now'),
|
||||||
|
datetime('now')
|
||||||
|
);
|
||||||
|
-- name: ListSessionMessages :many
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
session_key,
|
||||||
|
role,
|
||||||
|
sector,
|
||||||
|
importance,
|
||||||
|
salience,
|
||||||
|
decay_rate,
|
||||||
|
content,
|
||||||
|
tags,
|
||||||
|
created_at,
|
||||||
|
updated_at
|
||||||
|
FROM recall_items
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
AND session_key = sqlc.arg(session_key)
|
||||||
|
AND tags = 'session-message'
|
||||||
|
AND (
|
||||||
|
role = sqlc.arg(role)
|
||||||
|
OR sqlc.arg(role) = ''
|
||||||
|
)
|
||||||
|
ORDER BY created_at ASC
|
||||||
|
LIMIT sqlc.arg(lim);
|
||||||
|
-- name: CountSessionMessages :one
|
||||||
|
SELECT COUNT(*)
|
||||||
|
FROM recall_items
|
||||||
|
WHERE agent_id = sqlc.arg(agent_id)
|
||||||
|
AND session_key = sqlc.arg(session_key)
|
||||||
|
AND tags = 'session-message';
|
||||||
|
|
@ -44,6 +44,33 @@ func (q *Queries) CountRecallItems(ctx context.Context, arg CountRecallItemsPara
|
||||||
return count, err
|
return count, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const CountSessionMessages = `-- name: CountSessionMessages :one
|
||||||
|
SELECT COUNT(*)
|
||||||
|
FROM recall_items
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
AND session_key = ?2
|
||||||
|
AND tags = 'session-message'
|
||||||
|
`
|
||||||
|
|
||||||
|
type CountSessionMessagesParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
SessionKey string `json:"session_key"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// CountSessionMessages
|
||||||
|
//
|
||||||
|
// SELECT COUNT(*)
|
||||||
|
// FROM recall_items
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND session_key = ?2
|
||||||
|
// AND tags = 'session-message'
|
||||||
|
func (q *Queries) CountSessionMessages(ctx context.Context, arg CountSessionMessagesParams) (int64, error) {
|
||||||
|
row := q.db.QueryRowContext(ctx, CountSessionMessages, arg.AgentID, arg.SessionKey)
|
||||||
|
var count int64
|
||||||
|
err := row.Scan(&count)
|
||||||
|
return count, err
|
||||||
|
}
|
||||||
|
|
||||||
const DeleteRecallItem = `-- name: DeleteRecallItem :exec
|
const DeleteRecallItem = `-- name: DeleteRecallItem :exec
|
||||||
DELETE FROM recall_items
|
DELETE FROM recall_items
|
||||||
WHERE id = ?1
|
WHERE id = ?1
|
||||||
|
|
@ -292,6 +319,86 @@ func (q *Queries) InsertRecallItem(ctx context.Context, arg InsertRecallItemPara
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const InsertSessionMessage = `-- name: InsertSessionMessage :exec
|
||||||
|
INSERT INTO recall_items (
|
||||||
|
id,
|
||||||
|
agent_id,
|
||||||
|
session_key,
|
||||||
|
role,
|
||||||
|
sector,
|
||||||
|
importance,
|
||||||
|
salience,
|
||||||
|
decay_rate,
|
||||||
|
content,
|
||||||
|
tags,
|
||||||
|
created_at,
|
||||||
|
updated_at
|
||||||
|
)
|
||||||
|
VALUES (
|
||||||
|
?1,
|
||||||
|
?2,
|
||||||
|
?3,
|
||||||
|
?4,
|
||||||
|
'episodic',
|
||||||
|
0.5,
|
||||||
|
0.5,
|
||||||
|
0.01,
|
||||||
|
?5,
|
||||||
|
'session-message',
|
||||||
|
datetime('now'),
|
||||||
|
datetime('now')
|
||||||
|
)
|
||||||
|
`
|
||||||
|
|
||||||
|
type InsertSessionMessageParams struct {
|
||||||
|
ID ids.UUID `json:"id"`
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
SessionKey string `json:"session_key"`
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// InsertSessionMessage
|
||||||
|
//
|
||||||
|
// INSERT INTO recall_items (
|
||||||
|
// id,
|
||||||
|
// agent_id,
|
||||||
|
// session_key,
|
||||||
|
// role,
|
||||||
|
// sector,
|
||||||
|
// importance,
|
||||||
|
// salience,
|
||||||
|
// decay_rate,
|
||||||
|
// content,
|
||||||
|
// tags,
|
||||||
|
// created_at,
|
||||||
|
// updated_at
|
||||||
|
// )
|
||||||
|
// VALUES (
|
||||||
|
// ?1,
|
||||||
|
// ?2,
|
||||||
|
// ?3,
|
||||||
|
// ?4,
|
||||||
|
// 'episodic',
|
||||||
|
// 0.5,
|
||||||
|
// 0.5,
|
||||||
|
// 0.01,
|
||||||
|
// ?5,
|
||||||
|
// 'session-message',
|
||||||
|
// datetime('now'),
|
||||||
|
// datetime('now')
|
||||||
|
// )
|
||||||
|
func (q *Queries) InsertSessionMessage(ctx context.Context, arg InsertSessionMessageParams) error {
|
||||||
|
_, err := q.db.ExecContext(ctx, InsertSessionMessage,
|
||||||
|
arg.ID,
|
||||||
|
arg.AgentID,
|
||||||
|
arg.SessionKey,
|
||||||
|
arg.Role,
|
||||||
|
arg.Content,
|
||||||
|
)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
const ListRecallItems = `-- name: ListRecallItems :many
|
const ListRecallItems = `-- name: ListRecallItems :many
|
||||||
SELECT id,
|
SELECT id,
|
||||||
agent_id,
|
agent_id,
|
||||||
|
|
@ -385,6 +492,103 @@ func (q *Queries) ListRecallItems(ctx context.Context, arg ListRecallItemsParams
|
||||||
return items, nil
|
return items, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const ListSessionMessages = `-- name: ListSessionMessages :many
|
||||||
|
SELECT id,
|
||||||
|
agent_id,
|
||||||
|
session_key,
|
||||||
|
role,
|
||||||
|
sector,
|
||||||
|
importance,
|
||||||
|
salience,
|
||||||
|
decay_rate,
|
||||||
|
content,
|
||||||
|
tags,
|
||||||
|
created_at,
|
||||||
|
updated_at
|
||||||
|
FROM recall_items
|
||||||
|
WHERE agent_id = ?1
|
||||||
|
AND session_key = ?2
|
||||||
|
AND tags = 'session-message'
|
||||||
|
AND (
|
||||||
|
role = ?3
|
||||||
|
OR ?3 = ''
|
||||||
|
)
|
||||||
|
ORDER BY created_at ASC
|
||||||
|
LIMIT ?4
|
||||||
|
`
|
||||||
|
|
||||||
|
type ListSessionMessagesParams struct {
|
||||||
|
AgentID string `json:"agent_id"`
|
||||||
|
SessionKey string `json:"session_key"`
|
||||||
|
Role string `json:"role"`
|
||||||
|
Lim int64 `json:"lim"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListSessionMessages
|
||||||
|
//
|
||||||
|
// SELECT id,
|
||||||
|
// agent_id,
|
||||||
|
// session_key,
|
||||||
|
// role,
|
||||||
|
// sector,
|
||||||
|
// importance,
|
||||||
|
// salience,
|
||||||
|
// decay_rate,
|
||||||
|
// content,
|
||||||
|
// tags,
|
||||||
|
// created_at,
|
||||||
|
// updated_at
|
||||||
|
// FROM recall_items
|
||||||
|
// WHERE agent_id = ?1
|
||||||
|
// AND session_key = ?2
|
||||||
|
// AND tags = 'session-message'
|
||||||
|
// AND (
|
||||||
|
// role = ?3
|
||||||
|
// OR ?3 = ''
|
||||||
|
// )
|
||||||
|
// ORDER BY created_at ASC
|
||||||
|
// LIMIT ?4
|
||||||
|
func (q *Queries) ListSessionMessages(ctx context.Context, arg ListSessionMessagesParams) ([]RecallItem, error) {
|
||||||
|
rows, err := q.db.QueryContext(ctx, ListSessionMessages,
|
||||||
|
arg.AgentID,
|
||||||
|
arg.SessionKey,
|
||||||
|
arg.Role,
|
||||||
|
arg.Lim,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
items := []RecallItem{}
|
||||||
|
for rows.Next() {
|
||||||
|
var i RecallItem
|
||||||
|
if err := rows.Scan(
|
||||||
|
&i.ID,
|
||||||
|
&i.AgentID,
|
||||||
|
&i.SessionKey,
|
||||||
|
&i.Role,
|
||||||
|
&i.Sector,
|
||||||
|
&i.Importance,
|
||||||
|
&i.Salience,
|
||||||
|
&i.DecayRate,
|
||||||
|
&i.Content,
|
||||||
|
&i.Tags,
|
||||||
|
&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 SearchRecallByKeyword = `-- name: SearchRecallByKeyword :many
|
const SearchRecallByKeyword = `-- name: SearchRecallByKeyword :many
|
||||||
SELECT ri.id,
|
SELECT ri.id,
|
||||||
ri.agent_id,
|
ri.agent_id,
|
||||||
|
|
|
||||||
|
|
@ -58,4 +58,41 @@ CREATE TABLE IF NOT EXISTS memory_summaries (
|
||||||
to_msg_idx INTEGER NOT NULL DEFAULT 0,
|
to_msg_idx INTEGER NOT NULL DEFAULT 0,
|
||||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||||
);
|
);
|
||||||
CREATE INDEX IF NOT EXISTS idx_summaries_agent_session ON memory_summaries(agent_id, session_key);
|
CREATE INDEX IF NOT EXISTS idx_summaries_agent_session ON memory_summaries(agent_id, session_key);
|
||||||
|
-- Agent KV store: generic key-value pairs for agent state, preferences, config
|
||||||
|
CREATE TABLE IF NOT EXISTS agent_kv (
|
||||||
|
agent_id TEXT NOT NULL,
|
||||||
|
key TEXT NOT NULL,
|
||||||
|
value TEXT NOT NULL DEFAULT '',
|
||||||
|
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
PRIMARY KEY (agent_id, key)
|
||||||
|
);
|
||||||
|
-- Agent documents: versioned named documents (bootstrap files, identity, etc.)
|
||||||
|
CREATE TABLE IF NOT EXISTS agent_documents (
|
||||||
|
id BLOB PRIMARY KEY,
|
||||||
|
agent_id TEXT NOT NULL,
|
||||||
|
name TEXT NOT NULL,
|
||||||
|
category TEXT NOT NULL DEFAULT 'bootstrap',
|
||||||
|
content TEXT NOT NULL,
|
||||||
|
version INTEGER NOT NULL DEFAULT 1,
|
||||||
|
is_active BOOLEAN NOT NULL DEFAULT 1,
|
||||||
|
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
UNIQUE(agent_id, name)
|
||||||
|
);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_docs_agent_cat ON agent_documents(agent_id, category);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_docs_name ON agent_documents(name);
|
||||||
|
-- Agent audit log: append-only record of tool calls, state changes, etc.
|
||||||
|
CREATE TABLE IF NOT EXISTS agent_audit_log (
|
||||||
|
id BLOB PRIMARY KEY,
|
||||||
|
agent_id TEXT NOT NULL,
|
||||||
|
session_key TEXT NOT NULL DEFAULT '',
|
||||||
|
action TEXT NOT NULL,
|
||||||
|
target TEXT NOT NULL DEFAULT '',
|
||||||
|
input TEXT,
|
||||||
|
output TEXT,
|
||||||
|
duration_ms INTEGER,
|
||||||
|
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||||
|
);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_audit_agent_time ON agent_audit_log(agent_id, created_at DESC);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_audit_action ON agent_audit_log(action);
|
||||||
|
|
@ -40,6 +40,14 @@ sql:
|
||||||
go_type:
|
go_type:
|
||||||
import: "github.com/sipeed/picoclaw/pkg/ids"
|
import: "github.com/sipeed/picoclaw/pkg/ids"
|
||||||
type: "UUID"
|
type: "UUID"
|
||||||
|
- column: "agent_documents.id"
|
||||||
|
go_type:
|
||||||
|
import: "github.com/sipeed/picoclaw/pkg/ids"
|
||||||
|
type: "UUID"
|
||||||
|
- column: "agent_audit_log.id"
|
||||||
|
go_type:
|
||||||
|
import: "github.com/sipeed/picoclaw/pkg/ids"
|
||||||
|
type: "UUID"
|
||||||
# Domain type: recall_items.sector → memory.Sector
|
# Domain type: recall_items.sector → memory.Sector
|
||||||
- column: "recall_items.sector"
|
- column: "recall_items.sector"
|
||||||
go_type:
|
go_type:
|
||||||
|
|
|
||||||
|
|
@ -204,6 +204,16 @@ func (m *MemoryStore) RetrieveArchival(ctx context.Context, id ids.UUID) (string
|
||||||
return string(buf), nil
|
return string(buf), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ReadByID loads the content of a memory entry by its UUID string.
|
||||||
|
// Tries archival chunks first, then falls back to recall item.
|
||||||
|
func (m *MemoryStore) ReadByID(ctx context.Context, agentID, idStr string) (string, error) {
|
||||||
|
id, err := ids.Parse(idStr)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("invalid ID: %w", err)
|
||||||
|
}
|
||||||
|
return m.RetrieveArchival(ctx, id)
|
||||||
|
}
|
||||||
|
|
||||||
// --- Retrieval pipeline ---
|
// --- Retrieval pipeline ---
|
||||||
|
|
||||||
func (m *MemoryStore) Search(ctx context.Context, query string, opts memory.SearchOptions) ([]memory.SearchResult, error) {
|
func (m *MemoryStore) Search(ctx context.Context, query string, opts memory.SearchOptions) ([]memory.SearchResult, error) {
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue