diff --git a/go/internal/models/memory.go b/go/internal/models/memory.go new file mode 100644 index 0000000..e10c86d --- /dev/null +++ b/go/internal/models/memory.go @@ -0,0 +1,55 @@ +// Package models 定义织忆 MemoryWeave 的核心数据模型。 +package models + +import "time" + +// MemoryRecord 织忆中的单条记忆(存储在 LanceDB memories 表中)。 +type MemoryRecord struct { + ID string `json:"id"` + AgentID string `json:"agent_id"` + Namespace string `json:"namespace"` + Content string `json:"content"` + Category string `json:"category"` // system_fact, user_pref, proj_context, tool_usage, code_snippet + Vector []float32 `json:"vector"` // bge-m3 1024维,L2归一化 + Tier string `json:"tier"` // normal, core(core永不衰减) + QualityScore float64 `json:"quality_score"` + RecallCount int `json:"recall_count"` + Freshness string `json:"freshness"` // fresh, stale, verified + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` + IsDeleted bool `json:"is_deleted"` +} + +// EpisodeRecord 原始对话/任务日志(存储在 LanceDB episodes 表中)。 +type EpisodeRecord struct { + ID string `json:"id"` + AgentID string `json:"agent_id"` + Namespace string `json:"namespace"` + Content string `json:"content"` + Category string `json:"category"` + CreatedAt time.Time `json:"created_at"` +} + +// TombstoneRecord 已删除/废弃记忆的墓碑记录。 +type TombstoneRecord struct { + ID string `json:"id"` + OriginalID string `json:"original_id"` + Reason string `json:"reason"` // deprecated, merged, corrected, evicted + CreatedAt time.Time `json:"created_at"` +} + +// RecallResult recall 接口返回的单条结果。 +type RecallResult struct { + ID string `json:"id"` + Content string `json:"content"` + Category string `json:"category"` + Score float64 `json:"score"` + Timestamp string `json:"timestamp"` +} + +// RerankResult rerank 重排后的单条结果。 +type RerankResult struct { + Index int `json:"index"` + Score float64 `json:"score"` + Text string `json:"text"` +} diff --git a/go/internal/storage/embedder.go b/go/internal/storage/embedder.go new file mode 100644 index 0000000..cbaaa0c --- /dev/null +++ b/go/internal/storage/embedder.go @@ -0,0 +1,150 @@ +package storage + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "math" + "net/http" + "os" + "sync" +) + +// Embedder bge-m3 编码客户端。优先使用本地 vLLM(端口 8000),fallback 到模力方舟 API。 +type Embedder struct { + endpoint string // 本地 vLLM: http://localhost:8000/v1/embeddings + apiKey string // 模力方舟 API key(fallback 用) + httpClient *http.Client + cache sync.Map // string → []float32 + dim int +} + +// NewEmbedder 创建编码客户端。endpoint 默认为本地 vLLM。 +func NewEmbedder(endpoint string) *Embedder { + if endpoint == "" { + endpoint = os.Getenv("VLLM_ENDPOINT") + if endpoint == "" { + endpoint = "http://localhost:8000/v1/embeddings" + } + } + return &Embedder{ + endpoint: endpoint, + apiKey: os.Getenv("MOLIFANG_API_KEY"), + httpClient: &http.Client{}, + dim: 1024, + } +} + +// Encode 批量编码文本,返回 1024 维归一化向量。 +func (e *Embedder) Encode(texts []string) ([][]float32, error) { + // 检查缓存 + result := make([][]float32, len(texts)) + uncached := make([]int, 0) + for i, t := range texts { + if v, ok := e.cache.Load(t); ok { + result[i] = v.([]float32) + } else { + uncached = append(uncached, i) + } + } + if len(uncached) == 0 { + return result, nil + } + + // 收集未缓存的文本 + uncachedTexts := make([]string, len(uncached)) + for idx, i := range uncached { + uncachedTexts[idx] = texts[i] + } + + // 优先本地 vLLM + vectors, err := e.encodeRemote(e.endpoint, uncachedTexts) + if err != nil { + // fallback:模力方舟 + if e.apiKey != "" { + vectors, err = e.encodeRemote("https://ai.gitee.com/v1/embeddings", uncachedTexts) + } + if err != nil { + return nil, fmt.Errorf("encode: all endpoints failed: %w", err) + } + } + + // L2 归一化 + 缓存 + for idx, i := range uncached { + normalized := l2Normalize(vectors[idx]) + result[i] = normalized + e.cache.Store(texts[i], normalized) + } + return result, nil +} + +// EncodeSingle 编码单条文本。 +func (e *Embedder) EncodeSingle(text string) ([]float32, error) { + vecs, err := e.Encode([]string{text}) + if err != nil { + return nil, err + } + return vecs[0], nil +} + +// encodeRemote 调用远端 OpenAI 兼容 embeddings API。 +func (e *Embedder) encodeRemote(endpoint string, texts []string) ([][]float32, error) { + reqBody := map[string]any{ + "model": "bge-m3", + "input": texts, + } + body, _ := json.Marshal(reqBody) + + req, _ := http.NewRequest("POST", endpoint, bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + if e.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+e.apiKey) + } + + resp, err := e.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + rbody, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("embedding API: status %d: %s", resp.StatusCode, string(rbody)) + } + + var result struct { + Data []struct { + Embedding []float64 `json:"embedding"` + } `json:"data"` + } + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("decode embedding response: %w", err) + } + + vectors := make([][]float32, len(result.Data)) + for i, d := range result.Data { + vectors[i] = make([]float32, len(d.Embedding)) + for j, v := range d.Embedding { + vectors[i][j] = float32(v) + } + } + return vectors, nil +} + +// l2Normalize L2 归一化。 +func l2Normalize(v []float32) []float32 { + var sum float64 + for _, x := range v { + sum += float64(x) * float64(x) + } + norm := float32(math.Sqrt(sum)) + if norm == 0 { + return v + } + result := make([]float32, len(v)) + for i, x := range v { + result[i] = x / norm + } + return result +} diff --git a/go/internal/storage/lancedb.go b/go/internal/storage/lancedb.go new file mode 100644 index 0000000..85e2ee7 --- /dev/null +++ b/go/internal/storage/lancedb.go @@ -0,0 +1,175 @@ +// Package storage 提供 LanceDB 向量数据库的 HTTP REST 客户端封装。 +// LanceDB 本身是 Rust 编写的,当 Go 绑定不可用时通过 HTTP 与 LanceDB REST API 通信。 +// 生产部署时 LanceDB 由 zhiyi-consolidate (Rust) 原生管理。 +package storage + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + + "github.com/xiaoxue/memoryweave/internal/models" +) + +// LanceClient LanceDB HTTP REST 客户端。 +// 默认连接 http://localhost:8080(LanceDB REST 服务),将来 Rust sidecar 内嵌原生数据库后不再需要此层。 +type LanceClient struct { + baseURL string + httpClient *http.Client +} + +// NewLanceClient 创建 LanceDB 客户端,从环境变量 LANCEDB_URL 读取地址(默认 http://localhost:8080)。 +func NewLanceClient() *LanceClient { + url := os.Getenv("LANCEDB_URL") + if url == "" { + url = "http://localhost:8080" + } + return &LanceClient{ + baseURL: url, + httpClient: &http.Client{}, + } +} + +// CreateTables 初始化三张表(仅当不存在时创建)。 +func (c *LanceClient) CreateTables() error { + tables := []struct { + name string + schema any + }{ + {"memories", models.MemoryRecord{}}, + {"episodes", models.EpisodeRecord{}}, + {"tombstones", models.TombstoneRecord{}}, + } + for _, t := range tables { + if err := c.createTableIfNotExists(t.name); err != nil { + return fmt.Errorf("create table %s: %w", t.name, err) + } + } + return nil +} + +func (c *LanceClient) createTableIfNotExists(name string) error { + req, _ := http.NewRequest("POST", fmt.Sprintf("%s/v1/table/%s/create", c.baseURL, name), nil) + req.Header.Set("Content-Type", "application/json") + resp, err := c.httpClient.Do(req) + if err != nil { + return err + } + resp.Body.Close() + // 如果已存在则忽略错误 + if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusConflict { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("lanceDB create table %s: status %d, body: %s", name, resp.StatusCode, body) + } + return nil +} + +// Insert 向指定表插入一条记录。 +func (c *LanceClient) Insert(table string, record any) error { + body, _ := json.Marshal(record) + req, _ := http.NewRequest("POST", fmt.Sprintf("%s/v1/table/%s/insert", c.baseURL, table), + bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + resp, err := c.httpClient.Do(req) + if err != nil { + return fmt.Errorf("insert: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated { + rbody, _ := io.ReadAll(resp.Body) + return fmt.Errorf("insert: status %d: %s", resp.StatusCode, string(rbody)) + } + return nil +} + +// Search 向量搜索,返回 top_k 条最相似记录。可选按 namespace 过滤。 +func (c *LanceClient) Search(table string, vector []float32, topK int, namespaceFilter string) ([]models.MemoryRecord, error) { + reqBody := map[string]any{ + "vector": vector, + "top_k": topK, + "metric": "cosine", + "nprobes": 10, + "refine_factor": 2, + } + if namespaceFilter != "" { + reqBody["filter"] = fmt.Sprintf("namespace = '%s' AND is_deleted = false", namespaceFilter) + } else { + reqBody["filter"] = "is_deleted = false" + } + + body, _ := json.Marshal(reqBody) + req, _ := http.NewRequest("POST", fmt.Sprintf("%s/v1/table/%s/query", c.baseURL, table), + bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("search: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + rbody, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("search: status %d: %s", resp.StatusCode, string(rbody)) + } + + var results []models.MemoryRecord + if err := json.NewDecoder(resp.Body).Decode(&results); err != nil { + return nil, fmt.Errorf("decode search results: %w", err) + } + return results, nil +} + +// Update 更新指定记录的字段。 +func (c *LanceClient) Update(table, id string, fields map[string]any) error { + reqBody := map[string]any{ + "id": id, + "fields": fields, + } + body, _ := json.Marshal(reqBody) + req, _ := http.NewRequest("PUT", fmt.Sprintf("%s/v1/table/%s/update", c.baseURL, table), + bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + resp, err := c.httpClient.Do(req) + if err != nil { + return fmt.Errorf("update: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + rbody, _ := io.ReadAll(resp.Body) + return fmt.Errorf("update: status %d: %s", resp.StatusCode, string(rbody)) + } + return nil +} + +// SoftDelete 软删除(标记 is_deleted=true 并写入 tombstones)。 +func (c *LanceClient) SoftDelete(id, reason string) error { + // 1. 标记删除 + if err := c.Update("memories", id, map[string]any{ + "is_deleted": true, + }); err != nil { + return err + } + // 2. 写墓碑 + return c.Insert("tombstones", models.TombstoneRecord{ + OriginalID: id, + Reason: reason, + }) +} + +// Stats 返回各表记录数。 +func (c *LanceClient) Stats() (map[string]int, error) { + req, _ := http.NewRequest("GET", fmt.Sprintf("%s/v1/stats", c.baseURL), nil) + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("stats: %w", err) + } + defer resp.Body.Close() + + var stats map[string]int + if err := json.NewDecoder(resp.Body).Decode(&stats); err != nil { + return nil, fmt.Errorf("decode stats: %w", err) + } + return stats, nil +} diff --git a/go/internal/storage/recall.go b/go/internal/storage/recall.go new file mode 100644 index 0000000..349fa76 --- /dev/null +++ b/go/internal/storage/recall.go @@ -0,0 +1,168 @@ +// Package storage 提供记忆召回完整管线:编码 → 向量搜索 → 重排 → MMR 多样性。 +package storage + +import ( + "fmt" + "math" + "time" + + "github.com/xiaoxue/memoryweave/internal/models" +) + +// RecallPipeline 组合 Embedder、LanceClient、Reranker 构成完整召回管线。 +type RecallPipeline struct { + embedder *Embedder + lancedb *LanceClient + reranker *Reranker +} + +// NewRecallPipeline 创建召回管线。 +func NewRecallPipeline(embedder *Embedder, lancedb *LanceClient, reranker *Reranker) *RecallPipeline { + return &RecallPipeline{ + embedder: embedder, + lancedb: lancedb, + reranker: reranker, + } +} + +// Recall 完整召回流程: +// 1. query → bge-m3 编码 → 1024d 向量 +// 2. LanceDB ANN 搜索 → top 50 候选(粗排) +// 3. bge-reranker-v2-m3 重排 → top_k(精排) +// 4. MMR 多样性 → 最终结果 +func (p *RecallPipeline) Recall(query, namespace string, topK int, diversity float64) ([]models.RecallResult, error) { + if topK <= 0 { + topK = 10 + } + + // Step 1: Encode query + queryVec, err := p.embedder.EncodeSingle(query) + if err != nil { + return nil, fmt.Errorf("recall encode: %w", err) + } + + // Step 2: Coarse search — top 50 + candidates, err := p.lancedb.Search("memories", queryVec, 50, namespace) + if err != nil { + return nil, fmt.Errorf("recall search: %w", err) + } + if len(candidates) == 0 { + return []models.RecallResult{}, nil + } + + // Step 3: Rerank to top_k + docs := make([]string, len(candidates)) + for i, c := range candidates { + docs[i] = c.Content + } + reranked, err := p.reranker.Rerank(query, docs, topK) + if err != nil { + return nil, fmt.Errorf("recall rerank: %w", err) + } + + // Step 4: MMR diversity + finalResults := mmrSelect(reranked, candidates, topK, diversity) + + // Convert to RecallResult + results := make([]models.RecallResult, len(finalResults)) + for i, idx := range finalResults { + if idx < len(candidates) { + c := candidates[idx] + score := 0.0 + if idx < len(reranked) { + for _, r := range reranked { + if r.Index == idx { + score = r.Score + break + } + } + } + results[i] = models.RecallResult{ + ID: c.ID, + Content: c.Content, + Category: c.Category, + Score: score, + Timestamp: c.CreatedAt.Format(time.RFC3339), + } + } + } + + // 异步更新 recall_count(非阻塞) + go p.incrementRecallCount(results) + + return results, nil +} + +// mmrSelect 最大边际相关性选择:平衡相关性与多样性。 +// diversity=0 → 纯相关性排序,diversity=1 → 最大多样性。 +func mmrSelect(reranked []models.RerankResult, candidates []models.MemoryRecord, k int, lambda float64) []int { + if len(reranked) == 0 { + return nil + } + + selected := make([]int, 0, k) + candidateIndices := make([]int, len(reranked)) + for i := range reranked { + candidateIndices[i] = reranked[i].Index + } + + for len(selected) < k && len(candidateIndices) > 0 { + bestIdx := 0 + bestScore := -math.MaxFloat64 + + for i, ci := range candidateIndices { + relevance := 0.0 + for _, r := range reranked { + if r.Index == ci { + relevance = r.Score + break + } + } + + // Max similarity to already selected + maxSim := 0.0 + for _, si := range selected { + sim := cosineSimilarity(candidates[ci].Vector, candidates[si].Vector) + if sim > maxSim { + maxSim = sim + } + } + + mmr := (1-lambda)*relevance - lambda*maxSim + if mmr > bestScore { + bestScore = mmr + bestIdx = i + } + } + + selected = append(selected, candidateIndices[bestIdx]) + // Remove from candidates + candidateIndices = append(candidateIndices[:bestIdx], candidateIndices[bestIdx+1:]...) + } + + return selected +} + +// cosineSimilarity 计算两个向量的余弦相似度(向量已 L2 归一化时等于内积)。 +func cosineSimilarity(a, b []float32) float64 { + var sum float64 + minLen := len(a) + if len(b) < minLen { + minLen = len(b) + } + for i := 0; i < minLen; i++ { + sum += float64(a[i]) * float64(b[i]) + } + return sum +} + +// incrementRecallCount 异步递增被召回记忆的 recall_count。 +func (p *RecallPipeline) incrementRecallCount(results []models.RecallResult) { + for _, r := range results { + if r.ID != "" { + _ = p.lancedb.Update("memories", r.ID, map[string]any{ + "recall_count": map[string]string{"$inc": "1"}, + }) + } + } +} diff --git a/go/internal/storage/reranker.go b/go/internal/storage/reranker.go new file mode 100644 index 0000000..2b0d316 --- /dev/null +++ b/go/internal/storage/reranker.go @@ -0,0 +1,85 @@ +package storage + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + + "github.com/xiaoxue/memoryweave/internal/models" +) + +// Reranker bge-reranker-v2-m3 重排客户端,默认使用模力方舟 API。 +type Reranker struct { + endpoint string + apiKey string + httpClient *http.Client +} + +// NewReranker 创建重排客户端。 +func NewReranker(endpoint string) *Reranker { + if endpoint == "" { + endpoint = os.Getenv("RERANK_ENDPOINT") + if endpoint == "" { + endpoint = "https://ai.gitee.com/v1/rerank" + } + } + return &Reranker{ + endpoint: endpoint, + apiKey: os.Getenv("MOLIFANG_API_KEY"), + httpClient: &http.Client{}, + } +} + +// Rerank 对候选文档重排,返回 top_n 条最相关结果。 +func (r *Reranker) Rerank(query string, documents []string, topN int) ([]models.RerankResult, error) { + reqBody := map[string]any{ + "model": "bge-reranker-v2-m3", + "query": query, + "documents": documents, + "top_n": topN, + } + body, _ := json.Marshal(reqBody) + + req, _ := http.NewRequest("POST", r.endpoint, bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + if r.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+r.apiKey) + } + + resp, err := r.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("rerank: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + rbody, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("rerank API: status %d: %s", resp.StatusCode, string(rbody)) + } + + var result struct { + Results []struct { + Index int `json:"index"` + RelevanceScore float64 `json:"relevance_score"` + Document struct { + Text string `json:"text"` + } `json:"document"` + } `json:"results"` + } + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("decode rerank response: %w", err) + } + + results := make([]models.RerankResult, len(result.Results)) + for i, r := range result.Results { + results[i] = models.RerankResult{ + Index: r.Index, + Score: r.RelevanceScore, + Text: r.Document.Text, + } + } + return results, nil +} diff --git a/go/proto/consolidate.proto b/go/proto/consolidate.proto index 9cf6743..7a397e0 100644 --- a/go/proto/consolidate.proto +++ b/go/proto/consolidate.proto @@ -4,7 +4,7 @@ syntax = "proto3"; package consolidate; -option go_package = "github.com/xiaoxue/memoryweave/go/proto/consolidate"; +option go_package = "github.com/xiaoxue/memoryweave/proto/consolidate"; // Consolidation 服务 — 由 Rust sidecar 实现,Go 调用 service Consolidation {