diff --git a/internal/memory/block.go b/internal/memory/block.go new file mode 100644 index 0000000..415f3b1 --- /dev/null +++ b/internal/memory/block.go @@ -0,0 +1,219 @@ +package memory + +import ( + "database/sql" + "encoding/json" + "fmt" + "time" +) + +// BlockModality 是一等记忆块的原生模态。 +type BlockModality string + +const ( + BlockText BlockModality = "text" + BlockImage BlockModality = "image" + BlockVideo BlockModality = "video" + BlockAudio BlockModality = "audio" +) + +// MemoryBlock 是 Context、Document、Graph 三层共同使用的记忆块值。 +// +// 它不携带 Layer、Owner 或 RefCount：块当前由哪个层的容器持有，哪个层就是 +// 唯一事实源。Context→Document→Graph 迁移的是这个值本身，不建立平行保活账本。 +// PayloadDigest 仅用于定位内容寻址的原始字节，不表示另一条逻辑记忆。 +type MemoryBlock struct { + ID string `json:"id"` + Modality BlockModality `json:"modality"` + Text string `json:"text,omitempty"` + PayloadDigest string `json:"payload_digest,omitempty"` + MIME string `json:"mime,omitempty"` + Size int64 `json:"size,omitempty"` + Width int `json:"width,omitempty"` + Height int `json:"height,omitempty"` + Vector []float64 `json:"vector,omitempty"` + Fingerprint string `json:"fingerprint,omitempty"` + Source string `json:"source,omitempty"` + Tool string `json:"tool,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +// MemoryBlockEdge 是 L3 中连接一等记忆节点的结构化语义边。 +// source/target kind 当前允许 block、entity、sentence。 +type MemoryBlockEdge struct { + ID int64 `json:"id"` + SourceKind string `json:"source_kind"` + SourceID string `json:"source_id"` + TargetKind string `json:"target_kind"` + TargetID string `json:"target_id"` + Type string `json:"type"` + CreatedAt time.Time `json:"created_at"` +} + +func validBlockModality(modality BlockModality) bool { + return modality == BlockText || modality == BlockImage || modality == BlockVideo || modality == BlockAudio +} + +// PutMemoryBlocks 将完成 L2→L3 迁移的块写成 GraphDB 原生节点。 +// 调用方只有在本事务成功后才能从 Document 删除这些块。 +func (g *GraphDB) PutMemoryBlocks(blocks []MemoryBlock) error { + if len(blocks) == 0 { + return nil + } + g.mu.Lock() + defer g.mu.Unlock() + tx, err := g.db.Begin() + if err != nil { + return err + } + defer tx.Rollback() + + for _, block := range blocks { + if block.ID == "" { + return fmt.Errorf("memory block id is required") + } + if !validBlockModality(block.Modality) { + return fmt.Errorf("memory block %s has invalid modality %q", block.ID, block.Modality) + } + vectorJSON, err := json.Marshal(block.Vector) + if err != nil { + return fmt.Errorf("marshal memory block %s vector: %w", block.ID, err) + } + now := time.Now() + if block.CreatedAt.IsZero() { + block.CreatedAt = now + } + _, err = tx.Exec(`INSERT INTO memory_blocks ( + id, modality, text_content, payload_digest, mime, size, width, height, + vector, fingerprint, source, tool, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + modality = excluded.modality, + text_content = excluded.text_content, + payload_digest = excluded.payload_digest, + mime = excluded.mime, + size = excluded.size, + width = excluded.width, + height = excluded.height, + vector = excluded.vector, + fingerprint = excluded.fingerprint, + source = excluded.source, + tool = excluded.tool, + updated_at = excluded.updated_at`, + block.ID, block.Modality, block.Text, block.PayloadDigest, block.MIME, + block.Size, block.Width, block.Height, string(vectorJSON), block.Fingerprint, + block.Source, block.Tool, block.CreatedAt, now) + if err != nil { + return fmt.Errorf("put memory block %s: %w", block.ID, err) + } + } + return tx.Commit() +} + +// MemoryBlocks 查询 Graph 层实际持有的一等记忆节点。 +func (g *GraphDB) MemoryBlocks() ([]MemoryBlock, error) { + g.mu.RLock() + defer g.mu.RUnlock() + rows, err := g.db.Query(`SELECT id, modality, text_content, payload_digest, mime, + size, width, height, vector, fingerprint, source, tool, created_at, updated_at + FROM memory_blocks ORDER BY created_at, id`) + if err != nil { + return nil, err + } + defer rows.Close() + + var blocks []MemoryBlock + for rows.Next() { + var block MemoryBlock + var vectorJSON string + if err := rows.Scan(&block.ID, &block.Modality, &block.Text, &block.PayloadDigest, + &block.MIME, &block.Size, &block.Width, &block.Height, &vectorJSON, + &block.Fingerprint, &block.Source, &block.Tool, &block.CreatedAt, + &block.UpdatedAt); err != nil { + return nil, err + } + if vectorJSON != "" && vectorJSON != "null" { + if err := json.Unmarshal([]byte(vectorJSON), &block.Vector); err != nil { + return nil, fmt.Errorf("decode memory block %s vector: %w", block.ID, err) + } + } + blocks = append(blocks, block) + } + return blocks, rows.Err() +} + +func validGraphNodeKind(kind string) bool { + return kind == "block" || kind == "entity" || kind == "sentence" +} + +func graphNodeExists(tx *sql.Tx, kind, id string) (bool, error) { + var n int + var err error + switch kind { + case "block": + err = tx.QueryRow(`SELECT COUNT(*) FROM memory_blocks WHERE id = ?`, id).Scan(&n) + case "entity": + err = tx.QueryRow(`SELECT COUNT(*) FROM entities WHERE CAST(id AS TEXT) = ?`, id).Scan(&n) + case "sentence": + err = tx.QueryRow(`SELECT COUNT(*) FROM sentences WHERE CAST(id AS TEXT) = ?`, id).Scan(&n) + default: + return false, fmt.Errorf("invalid graph node kind %q", kind) + } + return n == 1, err +} + +// AddMemoryBlockEdge 建立 contains、depicts、derived_from 等原生图边。 +// 端点必须是真实 Graph 节点，不能用 owner 字符串伪装关系。 +func (g *GraphDB) AddMemoryBlockEdge(sourceKind, sourceID, targetKind, targetID, edgeType string) error { + if !validGraphNodeKind(sourceKind) || !validGraphNodeKind(targetKind) { + return fmt.Errorf("invalid memory block edge kinds %q -> %q", sourceKind, targetKind) + } + if sourceID == "" || targetID == "" || edgeType == "" { + return fmt.Errorf("memory block edge endpoints and type are required") + } + g.mu.Lock() + defer g.mu.Unlock() + tx, err := g.db.Begin() + if err != nil { + return err + } + defer tx.Rollback() + for _, endpoint := range []struct{ kind, id string }{{sourceKind, sourceID}, {targetKind, targetID}} { + exists, err := graphNodeExists(tx, endpoint.kind, endpoint.id) + if err != nil { + return err + } + if !exists { + return fmt.Errorf("%s graph node %s does not exist", endpoint.kind, endpoint.id) + } + } + _, err = tx.Exec(`INSERT OR IGNORE INTO memory_block_edges + (source_kind, source_id, target_kind, target_id, edge_type) + VALUES (?, ?, ?, ?, ?)`, sourceKind, sourceID, targetKind, targetID, edgeType) + if err != nil { + return err + } + return tx.Commit() +} + +func (g *GraphDB) MemoryBlockEdges() ([]MemoryBlockEdge, error) { + g.mu.RLock() + defer g.mu.RUnlock() + rows, err := g.db.Query(`SELECT id, source_kind, source_id, target_kind, target_id, + edge_type, created_at FROM memory_block_edges ORDER BY id`) + if err != nil { + return nil, err + } + defer rows.Close() + var edges []MemoryBlockEdge + for rows.Next() { + var edge MemoryBlockEdge + if err := rows.Scan(&edge.ID, &edge.SourceKind, &edge.SourceID, &edge.TargetKind, + &edge.TargetID, &edge.Type, &edge.CreatedAt); err != nil { + return nil, err + } + edges = append(edges, edge) + } + return edges, rows.Err() +} diff --git a/internal/memory/graph.go b/internal/memory/graph.go index 3167de2..26668b5 100644 --- a/internal/memory/graph.go +++ b/internal/memory/graph.go @@ -2,6 +2,7 @@ package memory import ( "database/sql" + "encoding/json" "fmt" "strings" "sync" @@ -107,6 +108,36 @@ func (g *GraphDB) initSchema() error { FOREIGN KEY (target_id) REFERENCES entities(id), UNIQUE(source_id, target_id, relation_type, session_id) )`, + `CREATE TABLE IF NOT EXISTS memory_blocks ( + id TEXT PRIMARY KEY, + modality TEXT NOT NULL, + text_content TEXT DEFAULT '', + payload_digest TEXT DEFAULT '', + mime TEXT DEFAULT '', + size INTEGER DEFAULT 0, + width INTEGER DEFAULT 0, + height INTEGER DEFAULT 0, + vector TEXT DEFAULT '', + fingerprint TEXT DEFAULT '', + source TEXT DEFAULT '', + tool TEXT DEFAULT '', + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + )`, + `CREATE TABLE IF NOT EXISTS memory_block_edges ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + source_kind TEXT NOT NULL, + source_id TEXT NOT NULL, + target_kind TEXT NOT NULL, + target_id TEXT NOT NULL, + edge_type TEXT NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + UNIQUE(source_kind, source_id, target_kind, target_id, edge_type) + )`, + `CREATE INDEX IF NOT EXISTS idx_memory_blocks_modality ON memory_blocks(modality)`, + `CREATE INDEX IF NOT EXISTS idx_memory_blocks_digest ON memory_blocks(payload_digest)`, + `CREATE INDEX IF NOT EXISTS idx_memory_block_edges_source ON memory_block_edges(source_kind, source_id)`, + `CREATE INDEX IF NOT EXISTS idx_memory_block_edges_target ON memory_block_edges(target_kind, target_id)`, `CREATE INDEX IF NOT EXISTS idx_entity_name ON entities(name)`, `CREATE INDEX IF NOT EXISTS idx_entity_type ON entities(type)`, `CREATE INDEX IF NOT EXISTS idx_relation_source ON relations(source_id)`, @@ -714,9 +745,58 @@ func (g *GraphDB) GraphData() (map[string]interface{}, error) { return nil, err } + brows, err := g.db.Query(`SELECT id, modality, text_content, payload_digest, mime, + size, width, height, vector, fingerprint, source, tool, created_at, updated_at + FROM memory_blocks ORDER BY created_at, id`) + if err != nil { + return nil, err + } + defer brows.Close() + var blocks []MemoryBlock + for brows.Next() { + var block MemoryBlock + var vectorJSON string + if err := brows.Scan(&block.ID, &block.Modality, &block.Text, &block.PayloadDigest, + &block.MIME, &block.Size, &block.Width, &block.Height, &vectorJSON, + &block.Fingerprint, &block.Source, &block.Tool, &block.CreatedAt, + &block.UpdatedAt); err != nil { + return nil, err + } + if vectorJSON != "" && vectorJSON != "null" { + if err := json.Unmarshal([]byte(vectorJSON), &block.Vector); err != nil { + return nil, fmt.Errorf("decode memory block %s vector: %w", block.ID, err) + } + } + blocks = append(blocks, block) + } + if err := brows.Err(); err != nil { + return nil, err + } + + berows, err := g.db.Query(`SELECT id, source_kind, source_id, target_kind, target_id, + edge_type, created_at FROM memory_block_edges ORDER BY id`) + if err != nil { + return nil, err + } + defer berows.Close() + var blockEdges []MemoryBlockEdge + for berows.Next() { + var edge MemoryBlockEdge + if err := berows.Scan(&edge.ID, &edge.SourceKind, &edge.SourceID, + &edge.TargetKind, &edge.TargetID, &edge.Type, &edge.CreatedAt); err != nil { + return nil, err + } + blockEdges = append(blockEdges, edge) + } + if err := berows.Err(); err != nil { + return nil, err + } + return map[string]interface{}{ - "nodes": entities, - "edges": relations, + "nodes": entities, + "edges": relations, + "memory_blocks": blocks, + "memory_block_edges": blockEdges, }, nil } diff --git a/internal/memory/graph_test.go b/internal/memory/graph_test.go index 79a9281..ee32ed5 100644 --- a/internal/memory/graph_test.go +++ b/internal/memory/graph_test.go @@ -2,8 +2,8 @@ package memory import ( "database/sql" - "path/filepath" "os" + "path/filepath" "testing" ) @@ -434,6 +434,84 @@ func TestMergeEntitiesNonexistent(t *testing.T) { } } +func TestMemoryBlocksAreFirstClassGraphNodes(t *testing.T) { + g := newTestGraph(t) + defer os.Remove(g.dbPath) + defer g.Close() + + image := MemoryBlock{ + ID: "block_image_1", Modality: BlockImage, + PayloadDigest: "0123456789abcdef", MIME: "image/png", Size: 1234, + Width: 768, Height: 512, Vector: []float64{0.1, 0.2, 0.3}, + Fingerprint: "qwen:test", Source: "qq", Tool: "upload", + } + text := MemoryBlock{ID: "block_text_1", Modality: BlockText, Text: "用户上传了一张架构图"} + if err := g.PutMemoryBlocks([]MemoryBlock{image, text}); err != nil { + t.Fatalf("PutMemoryBlocks: %v", err) + } + if err := g.AddMemoryBlockEdge("block", text.ID, "block", image.ID, "contains"); err != nil { + t.Fatalf("AddMemoryBlockEdge: %v", err) + } + + blocks, err := g.MemoryBlocks() + if err != nil { + t.Fatal(err) + } + if len(blocks) != 2 { + t.Fatalf("memory blocks=%d, want 2", len(blocks)) + } + var gotImage *MemoryBlock + for i := range blocks { + if blocks[i].ID == image.ID { + gotImage = &blocks[i] + } + } + if gotImage == nil || gotImage.Modality != BlockImage || gotImage.PayloadDigest != image.PayloadDigest || gotImage.Fingerprint != image.Fingerprint || len(gotImage.Vector) != 3 { + t.Fatalf("image block not round-tripped: %+v", gotImage) + } + + edges, err := g.MemoryBlockEdges() + if err != nil { + t.Fatal(err) + } + if len(edges) != 1 || edges[0].Type != "contains" || edges[0].SourceID != text.ID || edges[0].TargetID != image.ID { + t.Fatalf("memory block edges=%+v", edges) + } + + graph, err := g.GraphData() + if err != nil { + t.Fatal(err) + } + graphBlocks, ok := graph["memory_blocks"].([]MemoryBlock) + if !ok || len(graphBlocks) != 2 { + t.Fatalf("GraphData memory_blocks=%T %+v", graph["memory_blocks"], graph["memory_blocks"]) + } + graphEdges, ok := graph["memory_block_edges"].([]MemoryBlockEdge) + if !ok || len(graphEdges) != 1 { + t.Fatalf("GraphData memory_block_edges=%T %+v", graph["memory_block_edges"], graph["memory_block_edges"]) + } +} + +func TestMemoryBlockEdgeRejectsMissingEndpoint(t *testing.T) { + g := newTestGraph(t) + defer os.Remove(g.dbPath) + defer g.Close() + + if err := g.PutMemoryBlocks([]MemoryBlock{{ID: "known", Modality: BlockText, Text: "known"}}); err != nil { + t.Fatal(err) + } + if err := g.AddMemoryBlockEdge("block", "known", "block", "missing", "derived_from"); err == nil { + t.Fatal("edge to missing node must fail") + } + edges, err := g.MemoryBlockEdges() + if err != nil { + t.Fatal(err) + } + if len(edges) != 0 { + t.Fatalf("failed transaction left edges: %+v", edges) + } +} + func TestPlaceholders(t *testing.T) { if placeholders(0) != "NULL" { t.Errorf("expected NULL for n=0, got %s", placeholders(0))