feat(M2): 存储层完成
This commit is contained in:
xiaowei
parent
6667d42e14
commit
ad671bda34
|
|
@ -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"`
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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"},
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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 {
|
||||
|
|
|
|||
Loading…
Reference in New Issue