package rag import ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "time" core "git-ssh.u1680832.nyat.app/weichen/Agent-System-Core.git" ) // Route 解析一次模型路由(设置层;组合根/管线调用方负责取有效设置)。 type Route struct { Base string Key string Model string } // OK 表示该 kind 已配置可用。 func (r Route) OK() bool { return r.Base != "" && r.Model != "" } // HTTPEmbedder 是 OpenAI 兼容 /embeddings 客户端(core.EmbeddingProvider)。 // resolve 每次调用解析当前路由(设置可热更)。 type HTTPEmbedder struct { resolve func(ctx context.Context) Route client *http.Client } func NewHTTPEmbedder(resolve func(ctx context.Context) Route) *HTTPEmbedder { return &HTTPEmbedder{resolve: resolve, client: &http.Client{Timeout: 60 * time.Second}} } type embeddingsRequest struct { Model string `json:"model"` Input []string `json:"input"` } type embeddingsResponse struct { Data []struct { Index int `json:"index"` Embedding []float32 `json:"embedding"` } `json:"data"` Error *struct { Message string `json:"message"` } `json:"error,omitempty"` } // Model 当前 embedding 模型身份(U6-3:写入与检索同源;未配置返回空)。 func (e *HTTPEmbedder) Model() string { route := e.resolve(context.Background()) return route.Model } func (e *HTTPEmbedder) Embed(ctx context.Context, texts []string) ([][]float32, error) { if len(texts) == 0 { return nil, nil } route := e.resolve(ctx) if !route.OK() { return nil, fmt.Errorf("rag: embedding 模型未配置(settings.models.embedding)") } body, err := json.Marshal(embeddingsRequest{Model: route.Model, Input: texts}) if err != nil { return nil, err } req, err := http.NewRequestWithContext(ctx, http.MethodPost, route.Base+"/embeddings", bytes.NewReader(body)) if err != nil { return nil, err } req.Header.Set("Content-Type", "application/json") if route.Key != "" { req.Header.Set("Authorization", "Bearer "+route.Key) } resp, err := e.client.Do(req) if err != nil { return nil, err } defer resp.Body.Close() raw, _ := io.ReadAll(io.LimitReader(resp.Body, 32<<20)) if resp.StatusCode != http.StatusOK { return nil, fmt.Errorf("rag: embeddings %d: %s", resp.StatusCode, truncateStr(string(raw), 200)) } var out embeddingsResponse if err := json.Unmarshal(raw, &out); err != nil { return nil, err } if out.Error != nil { return nil, fmt.Errorf("rag: embeddings: %s", out.Error.Message) } if len(out.Data) != len(texts) { return nil, fmt.Errorf("rag: embeddings 返回 %d 条(期望 %d)", len(out.Data), len(texts)) } vecs := make([][]float32, len(texts)) for _, d := range out.Data { if d.Index < 0 || d.Index >= len(vecs) { return nil, fmt.Errorf("rag: embeddings index 越界 %d", d.Index) } vecs[d.Index] = d.Embedding } return vecs, nil } // HTTPReranker 是通用 /rerank 客户端(Cohere/Jina 兼容形状;core.Reranker)。 type HTTPReranker struct { resolve func(ctx context.Context) Route client *http.Client } func NewHTTPReranker(resolve func(ctx context.Context) Route) *HTTPReranker { return &HTTPReranker{resolve: resolve, client: &http.Client{Timeout: 30 * time.Second}} } type rerankRequest struct { Model string `json:"model"` Query string `json:"query"` Documents []struct { Text string `json:"text"` } `json:"documents"` TopN int `json:"top_n"` } type rerankResponse struct { Results []struct { Index int `json:"index"` Relevance float64 `json:"relevance_score"` } `json:"results"` } func (r *HTTPReranker) Rerank(ctx context.Context, query string, docs []core.KnowledgeDoc, topK int) ([]core.KnowledgeDoc, error) { route := r.resolve(ctx) if !route.OK() || len(docs) == 0 { return docs, nil // 未配置 = 显式跳过(排序退化为混合分) } reqBody := rerankRequest{Model: route.Model, Query: query, TopN: topK} for _, d := range docs { reqBody.Documents = append(reqBody.Documents, struct { Text string `json:"text"` }{Text: d.Chunk}) } body, err := json.Marshal(reqBody) if err != nil { return docs, err } req, err := http.NewRequestWithContext(ctx, http.MethodPost, route.Base+"/rerank", bytes.NewReader(body)) if err != nil { return docs, err } req.Header.Set("Content-Type", "application/json") if route.Key != "" { req.Header.Set("Authorization", "Bearer "+route.Key) } resp, err := r.client.Do(req) if err != nil { return docs, err // rerank 失败不阻断检索:返回原序 } defer resp.Body.Close() raw, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20)) if resp.StatusCode != http.StatusOK { return docs, fmt.Errorf("rag: rerank %d: %s", resp.StatusCode, truncateStr(string(raw), 200)) } var out rerankResponse if err := json.Unmarshal(raw, &out); err != nil { return docs, err } sorted := make([]core.KnowledgeDoc, 0, len(out.Results)) for _, res := range out.Results { if res.Index >= 0 && res.Index < len(docs) { sorted = append(sorted, docs[res.Index]) } } if len(sorted) == 0 { return docs, nil } return sorted, nil } func truncateStr(s string, n int) string { r := []rune(s) if len(r) <= n { return s } return string(r[:n]) + "…" }