root
446f865cdf
- Implemented KnowledgeHandler for managing knowledge groups and items, including listing, creating, updating, and deleting operations. - Added database migration scripts to create necessary tables for knowledge management, including knowledge_groups, knowledge_items, knowledge_parse_tasks, and knowledge_chunks_meta. - Introduced prompt_rule_knowledge_groups table to associate prompt rules with knowledge groups.
313 lines
8.0 KiB
Go
313 lines
8.0 KiB
Go
package retrieval
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"go.uber.org/zap"
|
|
|
|
"github.com/geo-platform/tenant-api/internal/shared/config"
|
|
"github.com/geo-platform/tenant-api/internal/shared/middleware"
|
|
)
|
|
|
|
const (
|
|
defaultSiliconFlowBaseURL = "https://api.siliconflow.cn/v1"
|
|
defaultSiliconFlowEmbeddingModel = "BAAI/bge-m3"
|
|
defaultSiliconFlowRerankerModel = "BAAI/bge-reranker-v2-m3"
|
|
defaultSiliconFlowTimeout = 30 * time.Second
|
|
)
|
|
|
|
type siliconFlowProvider struct {
|
|
client *http.Client
|
|
logger *zap.Logger
|
|
baseURL string
|
|
apiKey string
|
|
embeddingModel string
|
|
rerankerModel string
|
|
maxChunksPerDoc int
|
|
overlapTokens int
|
|
}
|
|
|
|
func NewSiliconFlowProvider(cfg config.RetrievalConfig, logger *zap.Logger) Provider {
|
|
if strings.TrimSpace(cfg.APIKey) == "" {
|
|
return disabledProvider{reason: "missing retrieval.api_key or SiliconFlow api key env (SILICONFLOW_API_KEY / RETRIEVAL_API_KEY)"}
|
|
}
|
|
|
|
baseURL := strings.TrimSpace(cfg.BaseURL)
|
|
if baseURL == "" {
|
|
baseURL = defaultSiliconFlowBaseURL
|
|
}
|
|
|
|
timeout := cfg.Timeout
|
|
if timeout <= 0 {
|
|
timeout = defaultSiliconFlowTimeout
|
|
}
|
|
|
|
embeddingModel := strings.TrimSpace(cfg.EmbeddingModel)
|
|
if embeddingModel == "" {
|
|
embeddingModel = defaultSiliconFlowEmbeddingModel
|
|
}
|
|
|
|
rerankerModel := strings.TrimSpace(cfg.RerankerModel)
|
|
if rerankerModel == "" {
|
|
rerankerModel = defaultSiliconFlowRerankerModel
|
|
}
|
|
|
|
maxChunksPerDoc := cfg.MaxChunksPerDoc
|
|
if maxChunksPerDoc <= 0 {
|
|
maxChunksPerDoc = 6
|
|
}
|
|
|
|
overlapTokens := cfg.OverlapTokens
|
|
if overlapTokens <= 0 {
|
|
overlapTokens = 60
|
|
}
|
|
|
|
return &siliconFlowProvider{
|
|
client: &http.Client{Timeout: timeout},
|
|
logger: logger,
|
|
baseURL: strings.TrimRight(baseURL, "/"),
|
|
apiKey: strings.TrimSpace(cfg.APIKey),
|
|
embeddingModel: embeddingModel,
|
|
rerankerModel: rerankerModel,
|
|
maxChunksPerDoc: maxChunksPerDoc,
|
|
overlapTokens: overlapTokens,
|
|
}
|
|
}
|
|
|
|
func (p *siliconFlowProvider) Validate() error {
|
|
if p == nil || p.client == nil {
|
|
return fmt.Errorf("%w: siliconflow client is not initialized", ErrNotConfigured)
|
|
}
|
|
if strings.TrimSpace(p.apiKey) == "" {
|
|
return fmt.Errorf("%w: siliconflow api key is empty", ErrNotConfigured)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (p *siliconFlowProvider) Embed(ctx context.Context, inputs []string) ([][]float64, error) {
|
|
if err := p.Validate(); err != nil {
|
|
return nil, err
|
|
}
|
|
if len(inputs) == 0 {
|
|
return [][]float64{}, nil
|
|
}
|
|
|
|
type embeddingRequest struct {
|
|
Model string `json:"model"`
|
|
Input []string `json:"input"`
|
|
}
|
|
type embeddingItem struct {
|
|
Embedding []float64 `json:"embedding"`
|
|
Index int `json:"index"`
|
|
}
|
|
type embeddingResponse struct {
|
|
Data []embeddingItem `json:"data"`
|
|
Error *struct {
|
|
Message string `json:"message"`
|
|
} `json:"error"`
|
|
}
|
|
|
|
payload := embeddingRequest{
|
|
Model: p.embeddingModel,
|
|
Input: inputs,
|
|
}
|
|
|
|
var resp embeddingResponse
|
|
if err := p.doJSON(ctx, http.MethodPost, "/embeddings", payload, &resp); err != nil {
|
|
return nil, err
|
|
}
|
|
if resp.Error != nil && strings.TrimSpace(resp.Error.Message) != "" {
|
|
return nil, errors.New(resp.Error.Message)
|
|
}
|
|
|
|
vectors := make([][]float64, len(inputs))
|
|
for _, item := range resp.Data {
|
|
if item.Index < 0 || item.Index >= len(vectors) {
|
|
continue
|
|
}
|
|
vectors[item.Index] = item.Embedding
|
|
}
|
|
for index, vector := range vectors {
|
|
if len(vector) == 0 {
|
|
return nil, fmt.Errorf("embedding vector missing at index %d", index)
|
|
}
|
|
}
|
|
|
|
return vectors, nil
|
|
}
|
|
|
|
func (p *siliconFlowProvider) Rerank(
|
|
ctx context.Context,
|
|
query string,
|
|
documents []string,
|
|
topN int,
|
|
) ([]RerankResult, error) {
|
|
if err := p.Validate(); err != nil {
|
|
return nil, err
|
|
}
|
|
if strings.TrimSpace(query) == "" || len(documents) == 0 {
|
|
return []RerankResult{}, nil
|
|
}
|
|
if topN <= 0 || topN > len(documents) {
|
|
topN = len(documents)
|
|
}
|
|
|
|
type rerankRequest struct {
|
|
Model string `json:"model"`
|
|
Query string `json:"query"`
|
|
Documents []string `json:"documents"`
|
|
TopN int `json:"top_n,omitempty"`
|
|
ReturnDocuments bool `json:"return_documents"`
|
|
MaxChunksPerDoc int `json:"max_chunks_per_doc,omitempty"`
|
|
OverlapTokens int `json:"overlap_tokens,omitempty"`
|
|
}
|
|
type rerankResponse struct {
|
|
Results []struct {
|
|
Index int `json:"index"`
|
|
RelevanceScore float64 `json:"relevance_score"`
|
|
Document struct {
|
|
Text string `json:"text"`
|
|
} `json:"document"`
|
|
} `json:"results"`
|
|
Error *struct {
|
|
Message string `json:"message"`
|
|
} `json:"error"`
|
|
}
|
|
|
|
payload := rerankRequest{
|
|
Model: p.rerankerModel,
|
|
Query: strings.TrimSpace(query),
|
|
Documents: documents,
|
|
TopN: topN,
|
|
ReturnDocuments: true,
|
|
MaxChunksPerDoc: p.maxChunksPerDoc,
|
|
OverlapTokens: p.overlapTokens,
|
|
}
|
|
|
|
var resp rerankResponse
|
|
if err := p.doJSON(ctx, http.MethodPost, "/rerank", payload, &resp); err != nil {
|
|
return nil, err
|
|
}
|
|
if resp.Error != nil && strings.TrimSpace(resp.Error.Message) != "" {
|
|
return nil, errors.New(resp.Error.Message)
|
|
}
|
|
|
|
results := make([]RerankResult, 0, len(resp.Results))
|
|
for _, item := range resp.Results {
|
|
text := item.Document.Text
|
|
if text == "" && item.Index >= 0 && item.Index < len(documents) {
|
|
text = documents[item.Index]
|
|
}
|
|
results = append(results, RerankResult{
|
|
Index: item.Index,
|
|
RelevanceScore: item.RelevanceScore,
|
|
Text: text,
|
|
})
|
|
}
|
|
|
|
return results, nil
|
|
}
|
|
|
|
func (p *siliconFlowProvider) doJSON(
|
|
ctx context.Context,
|
|
method, path string,
|
|
payload any,
|
|
target any,
|
|
) error {
|
|
body, err := json.Marshal(payload)
|
|
if err != nil {
|
|
p.logError(ctx, "siliconflow marshal request failed",
|
|
zap.String("method", method),
|
|
zap.String("path", path),
|
|
zap.Error(err),
|
|
)
|
|
return fmt.Errorf("marshal request: %w", err)
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, method, p.baseURL+path, bytes.NewReader(body))
|
|
if err != nil {
|
|
p.logError(ctx, "siliconflow create request failed",
|
|
zap.String("method", method),
|
|
zap.String("url", p.baseURL+path),
|
|
zap.Error(err),
|
|
)
|
|
return fmt.Errorf("create request: %w", err)
|
|
}
|
|
req.Header.Set("Authorization", "Bearer "+p.apiKey)
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
resp, err := p.client.Do(req)
|
|
if err != nil {
|
|
p.logError(ctx, "siliconflow request failed",
|
|
zap.String("method", method),
|
|
zap.String("url", req.URL.String()),
|
|
zap.Error(err),
|
|
)
|
|
return fmt.Errorf("request siliconflow: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
data, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
p.logError(ctx, "siliconflow read response failed",
|
|
zap.String("method", method),
|
|
zap.String("url", req.URL.String()),
|
|
zap.Int("status", resp.StatusCode),
|
|
zap.Error(err),
|
|
)
|
|
return fmt.Errorf("read response body: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode >= http.StatusBadRequest {
|
|
p.logError(ctx, "siliconflow non-success response",
|
|
zap.String("method", method),
|
|
zap.String("url", req.URL.String()),
|
|
zap.Int("status", resp.StatusCode),
|
|
zap.String("response_body", truncateForLog(string(data), 512)),
|
|
)
|
|
return fmt.Errorf("siliconflow request failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(data)))
|
|
}
|
|
|
|
if target == nil || len(data) == 0 {
|
|
return nil
|
|
}
|
|
if err := json.Unmarshal(data, target); err != nil {
|
|
p.logError(ctx, "siliconflow decode response failed",
|
|
zap.String("method", method),
|
|
zap.String("url", req.URL.String()),
|
|
zap.Int("status", resp.StatusCode),
|
|
zap.String("response_body", truncateForLog(string(data), 512)),
|
|
zap.Error(err),
|
|
)
|
|
return fmt.Errorf("decode response: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (p *siliconFlowProvider) logError(ctx context.Context, msg string, fields ...zap.Field) {
|
|
if p == nil || p.logger == nil {
|
|
return
|
|
}
|
|
if requestID := middleware.RequestIDFromContext(ctx); requestID != "" {
|
|
fields = append(fields, zap.String("request_id", requestID))
|
|
}
|
|
p.logger.Error(msg, fields...)
|
|
}
|
|
|
|
func truncateForLog(value string, limit int) string {
|
|
value = strings.TrimSpace(value)
|
|
if limit <= 0 || len(value) <= limit {
|
|
return value
|
|
}
|
|
return value[:limit] + "...(truncated)"
|
|
}
|