Blizzard
84d1a1dd3a
mcp-go 接通向量 RAG：embedding(OpenAI 兼容 provider 抽象) + Milvus 真实连接， kb_ingest 入库、wiki_search 真检索。retriever 节点一行不改即从桩变真。 - mcp-go internal/rag: embed.go(OpenAI 兼容 /embeddings 客户端) + milvus.go(milvus-sdk-go 真连，集合按首次 embedding 维度懒建+AUTOINDEX/COSINE索引+加载，insert/向量search) + rag.go(Engine: 切块→embed→insert / embed query→search；embedding 或 Milvus 缺则降级) - mcp-go gateway: 新工具 kb_ingest，wiki_search 换真(RAG 向量检索，kb 过滤 topK) - mcp-go main: rag.Open 读 MILVUS_ADDR/EMBED_BASE_URL/EMBED_API_KEY/EMBED_MODEL 环境变量 - gateway: POST /api/v1/kb/ingest → kb_ingest（供知识库页/脚本） - scripts/mock_embeddings.py: 确定性词法向量(字+bigram 哈希)，无真 key 验证检索 - 开发期 embedding 接在线 API(无真 key 用 mock)，见 llm-provider-strategy - 验证: 全模块 build✓ + e2e PASS; live——入库5条→Milvus；retriever 节点查'向量数据库' →召回 Milvus 那条→DeepSeek 答'Milvus'；查'知识图谱'→Neo4j（向量检索区分正确） 注: 当前向量单路；Bleve/Neo4j 融合 + rerank + 真实语义 embedding 为后续。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
98 lines
2.5 KiB
Go
98 lines
2.5 KiB
Go
// Package rag 实现 RAG 核心链：embedding(provider 抽象) + Milvus 向量库 + 入库/检索。
|
||
// 是 LLM Wiki 混合检索的向量路；Bleve/Neo4j 融合为后续扩展。
|
||
package rag
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"log"
|
||
"strings"
|
||
)
|
||
|
||
// Engine 聚合 embedding 与 Milvus，对外提供入库/检索。
|
||
type Engine struct {
|
||
emb *embedClient
|
||
mv *milvusStore
|
||
}
|
||
|
||
// Open 建立 RAG 引擎。embedding 未配 / Milvus 连不上 → 降级（检索返回空，不阻断工具服务）。
|
||
func Open(ctx context.Context, milvusAddr, embBase, embKey, embModel string) *Engine {
|
||
e := &Engine{}
|
||
if embBase != "" && embModel != "" {
|
||
e.emb = newEmbedClient(embBase, embKey, embModel)
|
||
log.Printf("[rag] embedding: %s model=%s", embBase, embModel)
|
||
} else {
|
||
log.Println("[rag] embedding 未配置，向量检索降级")
|
||
}
|
||
if milvusAddr != "" {
|
||
mv, err := openMilvus(ctx, milvusAddr)
|
||
if err != nil {
|
||
log.Printf("[rag] Milvus 不可用，向量检索降级: %v", err)
|
||
} else {
|
||
e.mv = mv
|
||
log.Printf("[rag] Milvus connected %s", milvusAddr)
|
||
}
|
||
}
|
||
return e
|
||
}
|
||
|
||
// Ready 报告 RAG 是否可用（embedding + Milvus 均就绪）。
|
||
func (e *Engine) Ready() bool { return e.emb.ready() && e.mv != nil }
|
||
|
||
// Ingest 把一段文本切块 → 向量化 → 写入 Milvus，返回块数。
|
||
func (e *Engine) Ingest(ctx context.Context, kb, text string) (int, error) {
|
||
if !e.Ready() {
|
||
return 0, errors.New("rag 未配置（需 embedding + Milvus）")
|
||
}
|
||
chunks := chunk(text)
|
||
if len(chunks) == 0 {
|
||
return 0, nil
|
||
}
|
||
vecs, err := e.emb.Embed(ctx, chunks)
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
if err := e.mv.insert(ctx, kb, chunks, vecs); err != nil {
|
||
return 0, err
|
||
}
|
||
return len(chunks), nil
|
||
}
|
||
|
||
// Search 向量化查询 → Milvus topK 检索。降级时返回空。
|
||
func (e *Engine) Search(ctx context.Context, kb, query string, topK int) ([]Hit, error) {
|
||
if !e.Ready() {
|
||
return nil, nil
|
||
}
|
||
if topK <= 0 {
|
||
topK = 5
|
||
}
|
||
vecs, err := e.emb.Embed(ctx, []string{query})
|
||
if err != nil || len(vecs) == 0 {
|
||
return nil, err
|
||
}
|
||
return e.mv.search(ctx, kb, vecs[0], topK)
|
||
}
|
||
|
||
func (e *Engine) Close() {
|
||
if e.mv != nil {
|
||
e.mv.close()
|
||
}
|
||
}
|
||
|
||
// chunk 朴素切块：按行切，去空白；过长再按长度切。真实系统应做版面/语义切块。
|
||
func chunk(text string) []string {
|
||
var out []string
|
||
for _, line := range strings.Split(text, "\n") {
|
||
s := strings.TrimSpace(line)
|
||
if s == "" {
|
||
continue
|
||
}
|
||
for len(s) > 2000 {
|
||
out = append(out, s[:2000])
|
||
s = s[2000:]
|
||
}
|
||
out = append(out, s)
|
||
}
|
||
return out
|
||
}
|