Agent-System-Backend/internal/rag/embed.go

package rag

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"

	core "agent-system-core"
)

// 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"`
}

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]) + "…"
}