Blizzard
3b54e59ecf
embedding 从 env 改为控制面驱动(持久化+可视化)，复用 chat 模型同套范式： 配置控制面泛化为按 kind(chat/embedding)，加 embedding kind。 - shared: 配置 subjects 泛化 sundynix.config.<kind>.get/.updated；bus 方法改 kind 参数 (RequestConfig/ServeConfig/PublishConfigUpdated/SubscribeConfigUpdated) - gateway: sundynix_model 加 kind 列(每 kind 唯一激活)+旧行回填 chat；admin 按 kind 增删改/激活/列表，测试连接 embedding 走 POST /embeddings；main 按 kind ServeConfig； 变更广播各 kind - dispatcher: 取 chat 配置(kind 化) - mcp-go: rag.Engine.SetEmbedding 热更新(RWMutex)；main 取/订阅 embedding 控制面配置 (覆盖 env) - admin 控制台: api 按 kind；抽出复用 ModelManager；ModelsPage(chat)+新 DatasourcesPage (embedding + 向量/图库占位)；routes 数据源页就绪 - 验证: 全模块 build✓ + e2e PASS + 控制台 npm build✓；live 全跑通——chat(DeepSeek 回填 kind 仍工作)；mcp-go 不带 EMBED env 启动→控制台配 embedding(百炼)→测试连接✓→激活 →NATS 热更新 mcp-go→入库+语义检索'存向量的数据库'→Milvus；浏览器数据源页拉到激活配置 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
117 lines
3.0 KiB
Go
117 lines
3.0 KiB
Go
// Package rag 实现 RAG 核心链：embedding(provider 抽象) + Milvus 向量库 + 入库/检索。
|
||
// 是 LLM Wiki 混合检索的向量路；Bleve/Neo4j 融合为后续扩展。
|
||
package rag
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"log"
|
||
"strings"
|
||
"sync"
|
||
)
|
||
|
||
// Engine 聚合 embedding 与 Milvus，对外提供入库/检索。embedding 可热更新（控制面下发）。
|
||
type Engine struct {
|
||
mu sync.RWMutex
|
||
emb *embedClient
|
||
mv *milvusStore
|
||
}
|
||
|
||
// SetEmbedding 热更新 embedding 配置（控制面变更时调用）。空配置=关闭向量检索。
|
||
func (e *Engine) SetEmbedding(base, key, model string) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
if base == "" || model == "" {
|
||
e.emb = nil
|
||
return
|
||
}
|
||
e.emb = newEmbedClient(base, key, model)
|
||
log.Printf("[rag] embedding 配置: %s model=%s", base, model)
|
||
}
|
||
|
||
func (e *Engine) embed() *embedClient {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
return e.emb
|
||
}
|
||
|
||
// Open 建立 RAG 引擎。embedding 未配 / Milvus 连不上 → 降级（检索返回空，不阻断工具服务）。
|
||
func Open(ctx context.Context, milvusAddr, embBase, embKey, embModel string) *Engine {
|
||
e := &Engine{}
|
||
if embBase != "" && embModel != "" {
|
||
e.SetEmbedding(embBase, embKey, embModel) // env 初值（控制面会覆盖）
|
||
} 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.embed().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.embed().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.embed().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
|
||
}
|