feat: add knowledge management functionality with CRUD operations and database schema
- 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.
This commit is contained in:
@@ -0,0 +1,312 @@
|
||||
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)"
|
||||
}
|
||||
Reference in New Issue
Block a user