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,57 @@
|
||||
package retrieval
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/geo-platform/tenant-api/internal/shared/config"
|
||||
)
|
||||
|
||||
var ErrNotConfigured = errors.New("retrieval provider is not configured")
|
||||
|
||||
type Provider interface {
|
||||
Validate() error
|
||||
Embed(ctx context.Context, inputs []string) ([][]float64, error)
|
||||
Rerank(ctx context.Context, query string, documents []string, topN int) ([]RerankResult, error)
|
||||
}
|
||||
|
||||
type RerankResult struct {
|
||||
Index int
|
||||
RelevanceScore float64
|
||||
Text string
|
||||
}
|
||||
|
||||
func NewProvider(cfg config.RetrievalConfig, logger *zap.Logger) Provider {
|
||||
provider := strings.ToLower(strings.TrimSpace(cfg.Provider))
|
||||
switch provider {
|
||||
case "", "disabled":
|
||||
return disabledProvider{reason: "retrieval provider is disabled"}
|
||||
case "siliconflow":
|
||||
return NewSiliconFlowProvider(cfg, logger)
|
||||
default:
|
||||
return disabledProvider{reason: fmt.Sprintf("unsupported retrieval provider %q", cfg.Provider)}
|
||||
}
|
||||
}
|
||||
|
||||
type disabledProvider struct {
|
||||
reason string
|
||||
}
|
||||
|
||||
func (p disabledProvider) Validate() error {
|
||||
if p.reason == "" {
|
||||
p.reason = "missing retrieval configuration"
|
||||
}
|
||||
return fmt.Errorf("%w: %s", ErrNotConfigured, p.reason)
|
||||
}
|
||||
|
||||
func (p disabledProvider) Embed(context.Context, []string) ([][]float64, error) {
|
||||
return nil, p.Validate()
|
||||
}
|
||||
|
||||
func (p disabledProvider) Rerank(context.Context, string, []string, int) ([]RerankResult, error) {
|
||||
return nil, p.Validate()
|
||||
}
|
||||
@@ -0,0 +1,426 @@
|
||||
package retrieval
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
sdk "github.com/qdrant/go-client/qdrant"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/geo-platform/tenant-api/internal/shared/config"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/middleware"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultQdrantHost = "localhost"
|
||||
defaultQdrantGRPCPort = 6334
|
||||
defaultQdrantCollection = "geo_knowledge_chunks"
|
||||
)
|
||||
|
||||
type VectorStore interface {
|
||||
Validate() error
|
||||
EnsureCollection(ctx context.Context, vectorSize int) error
|
||||
Upsert(ctx context.Context, points []VectorPoint) error
|
||||
Search(ctx context.Context, req SearchRequest) ([]SearchPoint, error)
|
||||
DeletePoints(ctx context.Context, pointIDs []string) error
|
||||
}
|
||||
|
||||
type VectorPoint struct {
|
||||
ID string
|
||||
Vector []float64
|
||||
Payload map[string]any
|
||||
}
|
||||
|
||||
type SearchRequest struct {
|
||||
Vector []float64
|
||||
TenantID int64
|
||||
GroupIDs []int64
|
||||
Limit int
|
||||
}
|
||||
|
||||
type SearchPoint struct {
|
||||
ID string
|
||||
Score float64
|
||||
Payload map[string]any
|
||||
}
|
||||
|
||||
type qdrantStore struct {
|
||||
client *sdk.Client
|
||||
collection string
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
type disabledVectorStore struct {
|
||||
reason string
|
||||
}
|
||||
|
||||
func NewQdrantStore(cfg config.QdrantConfig, logger *zap.Logger) VectorStore {
|
||||
host, port, useTLS, err := parseQdrantEndpoint(cfg.URL)
|
||||
if err != nil {
|
||||
return disabledVectorStore{reason: fmt.Sprintf("invalid qdrant.url: %v", err)}
|
||||
}
|
||||
|
||||
collection := strings.TrimSpace(cfg.Collection)
|
||||
if collection == "" {
|
||||
collection = defaultQdrantCollection
|
||||
}
|
||||
|
||||
client, err := sdk.NewClient(&sdk.Config{
|
||||
Host: host,
|
||||
Port: port,
|
||||
APIKey: strings.TrimSpace(cfg.APIKey),
|
||||
UseTLS: useTLS,
|
||||
SkipCompatibilityCheck: true,
|
||||
})
|
||||
if err != nil {
|
||||
return disabledVectorStore{reason: fmt.Sprintf("init qdrant client failed: %v", err)}
|
||||
}
|
||||
|
||||
return &qdrantStore{
|
||||
client: client,
|
||||
collection: collection,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *qdrantStore) Validate() error {
|
||||
if s == nil || s.client == nil {
|
||||
return fmt.Errorf("%w: qdrant client is not initialized", ErrNotConfigured)
|
||||
}
|
||||
if strings.TrimSpace(s.collection) == "" {
|
||||
return fmt.Errorf("%w: qdrant collection is empty", ErrNotConfigured)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *qdrantStore) EnsureCollection(ctx context.Context, vectorSize int) error {
|
||||
if err := s.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if vectorSize <= 0 {
|
||||
return fmt.Errorf("invalid qdrant vector size %d", vectorSize)
|
||||
}
|
||||
|
||||
exists, err := s.client.CollectionExists(ctx, s.collection)
|
||||
if err != nil {
|
||||
s.logError(ctx, "qdrant collection exists check failed",
|
||||
zap.String("collection", s.collection),
|
||||
zap.Error(err),
|
||||
)
|
||||
return fmt.Errorf("query qdrant collection: %w", err)
|
||||
}
|
||||
if exists {
|
||||
return nil
|
||||
}
|
||||
|
||||
err = s.client.CreateCollection(ctx, &sdk.CreateCollection{
|
||||
CollectionName: s.collection,
|
||||
VectorsConfig: sdk.NewVectorsConfig(&sdk.VectorParams{
|
||||
Size: uint64(vectorSize),
|
||||
Distance: sdk.Distance_Cosine,
|
||||
}),
|
||||
})
|
||||
if err != nil {
|
||||
s.logError(ctx, "qdrant create collection failed",
|
||||
zap.String("collection", s.collection),
|
||||
zap.Int("vector_size", vectorSize),
|
||||
zap.Error(err),
|
||||
)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *qdrantStore) Upsert(ctx context.Context, points []VectorPoint) error {
|
||||
if err := s.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(points) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
converted := make([]*sdk.PointStruct, 0, len(points))
|
||||
for _, point := range points {
|
||||
if strings.TrimSpace(point.ID) == "" || len(point.Vector) == 0 {
|
||||
continue
|
||||
}
|
||||
converted = append(converted, &sdk.PointStruct{
|
||||
Id: sdk.NewID(point.ID),
|
||||
Vectors: sdk.NewVectors(float64SliceToFloat32(point.Vector)...),
|
||||
Payload: sdk.NewValueMap(point.Payload),
|
||||
})
|
||||
}
|
||||
if len(converted) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
wait := true
|
||||
_, err := s.client.Upsert(ctx, &sdk.UpsertPoints{
|
||||
CollectionName: s.collection,
|
||||
Wait: &wait,
|
||||
Points: converted,
|
||||
})
|
||||
if err != nil {
|
||||
s.logError(ctx, "qdrant upsert failed",
|
||||
zap.String("collection", s.collection),
|
||||
zap.Int("point_count", len(converted)),
|
||||
zap.Error(err),
|
||||
)
|
||||
return fmt.Errorf("upsert qdrant points: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *qdrantStore) Search(ctx context.Context, req SearchRequest) ([]SearchPoint, error) {
|
||||
if err := s.Validate(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(req.Vector) == 0 || req.TenantID <= 0 {
|
||||
return []SearchPoint{}, nil
|
||||
}
|
||||
limit := req.Limit
|
||||
if limit <= 0 {
|
||||
limit = 10
|
||||
}
|
||||
limitValue := uint64(limit)
|
||||
|
||||
filter := &sdk.Filter{
|
||||
Must: []*sdk.Condition{
|
||||
sdk.NewMatchInt("tenant_id", req.TenantID),
|
||||
sdk.NewMatchKeyword("status", "active"),
|
||||
},
|
||||
}
|
||||
groupIDs := normalizeGroupIDs(req.GroupIDs)
|
||||
if len(groupIDs) > 0 {
|
||||
filter.Should = []*sdk.Condition{
|
||||
sdk.NewMatchInts("group_id", groupIDs...),
|
||||
}
|
||||
}
|
||||
|
||||
points, err := s.client.Query(ctx, &sdk.QueryPoints{
|
||||
CollectionName: s.collection,
|
||||
Query: sdk.NewQueryDense(float64SliceToFloat32(req.Vector)),
|
||||
Filter: filter,
|
||||
Limit: &limitValue,
|
||||
WithPayload: sdk.NewWithPayload(true),
|
||||
})
|
||||
if err != nil {
|
||||
s.logError(ctx, "qdrant search failed",
|
||||
zap.String("collection", s.collection),
|
||||
zap.Int64("tenant_id", req.TenantID),
|
||||
zap.Int("group_count", len(groupIDs)),
|
||||
zap.Int("limit", limit),
|
||||
zap.Error(err),
|
||||
)
|
||||
return nil, fmt.Errorf("query qdrant points: %w", err)
|
||||
}
|
||||
|
||||
results := make([]SearchPoint, 0, len(points))
|
||||
for _, point := range points {
|
||||
if point == nil {
|
||||
continue
|
||||
}
|
||||
results = append(results, SearchPoint{
|
||||
ID: pointIDToString(point.GetId()),
|
||||
Score: float64(point.GetScore()),
|
||||
Payload: valueMapToAny(point.GetPayload()),
|
||||
})
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func (s *qdrantStore) DeletePoints(ctx context.Context, pointIDs []string) error {
|
||||
if err := s.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(pointIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
ids := make([]*sdk.PointId, 0, len(pointIDs))
|
||||
for _, pointID := range pointIDs {
|
||||
pointID = strings.TrimSpace(pointID)
|
||||
if pointID == "" {
|
||||
continue
|
||||
}
|
||||
ids = append(ids, sdk.NewID(pointID))
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
wait := true
|
||||
_, err := s.client.Delete(ctx, &sdk.DeletePoints{
|
||||
CollectionName: s.collection,
|
||||
Wait: &wait,
|
||||
Points: sdk.NewPointsSelector(ids...),
|
||||
})
|
||||
if err != nil {
|
||||
s.logError(ctx, "qdrant delete points failed",
|
||||
zap.String("collection", s.collection),
|
||||
zap.Int("point_count", len(ids)),
|
||||
zap.Error(err),
|
||||
)
|
||||
return fmt.Errorf("delete qdrant points: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *qdrantStore) logError(ctx context.Context, msg string, fields ...zap.Field) {
|
||||
if s == nil || s.logger == nil {
|
||||
return
|
||||
}
|
||||
if requestID := middleware.RequestIDFromContext(ctx); requestID != "" {
|
||||
fields = append(fields, zap.String("request_id", requestID))
|
||||
}
|
||||
s.logger.Error(msg, fields...)
|
||||
}
|
||||
|
||||
func parseQdrantEndpoint(raw string) (string, int, bool, error) {
|
||||
value := strings.TrimSpace(raw)
|
||||
if value == "" {
|
||||
return defaultQdrantHost, defaultQdrantGRPCPort, false, nil
|
||||
}
|
||||
|
||||
if strings.Contains(value, "://") {
|
||||
parsed, err := url.Parse(value)
|
||||
if err != nil {
|
||||
return "", 0, false, err
|
||||
}
|
||||
port := defaultQdrantGRPCPort
|
||||
if parsed.Port() != "" {
|
||||
port, err = strconv.Atoi(parsed.Port())
|
||||
if err != nil {
|
||||
return "", 0, false, err
|
||||
}
|
||||
}
|
||||
host := parsed.Hostname()
|
||||
if host == "" {
|
||||
host = defaultQdrantHost
|
||||
}
|
||||
return host, port, parsed.Scheme == "https", nil
|
||||
}
|
||||
|
||||
host, portText, err := net.SplitHostPort(value)
|
||||
if err == nil {
|
||||
port, convErr := strconv.Atoi(portText)
|
||||
if convErr != nil {
|
||||
return "", 0, false, convErr
|
||||
}
|
||||
return host, port, false, nil
|
||||
}
|
||||
|
||||
if strings.Contains(err.Error(), "missing port in address") {
|
||||
return value, defaultQdrantGRPCPort, false, nil
|
||||
}
|
||||
|
||||
return "", 0, false, err
|
||||
}
|
||||
|
||||
func float64SliceToFloat32(values []float64) []float32 {
|
||||
result := make([]float32, 0, len(values))
|
||||
for _, value := range values {
|
||||
result = append(result, float32(value))
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func pointIDToString(id *sdk.PointId) string {
|
||||
if id == nil {
|
||||
return ""
|
||||
}
|
||||
if uuid := strings.TrimSpace(id.GetUuid()); uuid != "" {
|
||||
return uuid
|
||||
}
|
||||
if num := id.GetNum(); num > 0 {
|
||||
return strconv.FormatUint(num, 10)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func valueMapToAny(values map[string]*sdk.Value) map[string]any {
|
||||
if len(values) == 0 {
|
||||
return map[string]any{}
|
||||
}
|
||||
result := make(map[string]any, len(values))
|
||||
for key, value := range values {
|
||||
result[key] = valueToAny(value)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func valueToAny(value *sdk.Value) any {
|
||||
if value == nil {
|
||||
return nil
|
||||
}
|
||||
switch value.GetKind().(type) {
|
||||
case *sdk.Value_StringValue:
|
||||
return value.GetStringValue()
|
||||
case *sdk.Value_IntegerValue:
|
||||
return value.GetIntegerValue()
|
||||
case *sdk.Value_DoubleValue:
|
||||
return value.GetDoubleValue()
|
||||
case *sdk.Value_BoolValue:
|
||||
return value.GetBoolValue()
|
||||
case *sdk.Value_ListValue:
|
||||
list := value.GetListValue()
|
||||
if list == nil {
|
||||
return []any{}
|
||||
}
|
||||
items := make([]any, 0, len(list.GetValues()))
|
||||
for _, item := range list.GetValues() {
|
||||
items = append(items, valueToAny(item))
|
||||
}
|
||||
return items
|
||||
case *sdk.Value_StructValue:
|
||||
structValue := value.GetStructValue()
|
||||
if structValue == nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
return valueMapToAny(structValue.GetFields())
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeGroupIDs(ids []int64) []int64 {
|
||||
result := make([]int64, 0, len(ids))
|
||||
seen := make(map[int64]struct{}, len(ids))
|
||||
for _, id := range ids {
|
||||
if id <= 0 {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[id]; ok {
|
||||
continue
|
||||
}
|
||||
seen[id] = struct{}{}
|
||||
result = append(result, id)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func (s disabledVectorStore) Validate() error {
|
||||
if s.reason == "" {
|
||||
s.reason = "missing qdrant configuration"
|
||||
}
|
||||
return fmt.Errorf("%w: %s", ErrNotConfigured, s.reason)
|
||||
}
|
||||
|
||||
func (s disabledVectorStore) EnsureCollection(context.Context, int) error {
|
||||
return s.Validate()
|
||||
}
|
||||
|
||||
func (s disabledVectorStore) Upsert(context.Context, []VectorPoint) error {
|
||||
return s.Validate()
|
||||
}
|
||||
|
||||
func (s disabledVectorStore) Search(context.Context, SearchRequest) ([]SearchPoint, error) {
|
||||
return nil, s.Validate()
|
||||
}
|
||||
|
||||
func (s disabledVectorStore) DeletePoints(context.Context, []string) error {
|
||||
return s.Validate()
|
||||
}
|
||||
@@ -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