From 1473000135cf0ddca79f61f36babd4a714ee5279 Mon Sep 17 00:00:00 2001 From: qinling0210 <88864212+qinling0210@users.noreply.github.com> Date: Fri, 24 Apr 2026 15:30:14 +0800 Subject: [PATCH] Implement retrieval_test in GO (#14231) ### What problem does this PR solve? Implement retrieval_test in GO ### Type of change - [x] Refactoring --- api/apps/chunk_app.py | 1 + conf/models/siliconflow.json | 26 + go.mod | 4 +- go.sum | 4 +- internal/cli/user_parser.go | 3 +- internal/common/constants.go | 8 + internal/dao/tenant_llm.go | 68 + internal/engine/elasticsearch/get.go | 31 +- internal/engine/elasticsearch/search.go | 278 ++- internal/engine/engine.go | 16 +- internal/engine/global.go | 9 +- internal/engine/infinity/common.go | 59 +- internal/engine/infinity/dataset.go | 2 +- internal/engine/infinity/get.go | 200 ++- internal/engine/infinity/search.go | 1578 +++++++++--------- internal/engine/types/types.go | 99 +- internal/entity/kb.go | 1 + internal/entity/models/deepseek.go | 5 + internal/entity/models/dummy.go | 5 + internal/entity/models/minimax.go | 5 + internal/entity/models/moonshot.go | 5 + internal/entity/models/types.go | 8 + internal/entity/models/zhipu-ai.go | 100 ++ internal/entity/types.go | 7 + internal/logger/logger.go | 5 + internal/service/chunk.go | 743 ++++----- internal/service/generator.go | 167 ++ internal/service/load_prompt.go | 160 ++ internal/service/metadata.go | 223 ++- internal/service/metadata_filter.go | 563 +++++++ internal/service/model_service.go | 193 ++- internal/service/models/factory.go | 60 + internal/service/models/siliconflow_model.go | 258 ++- internal/service/nlp/query_builder.go | 45 +- internal/service/nlp/reranker.go | 254 ++- internal/service/nlp/retrieval.go | 787 +++++++++ internal/service/search.go | 27 + internal/service/tag.go | 358 ++++ internal/tokenizer/tokenizer.go | 11 + internal/utility/convert.go | 29 + rag/llm/rerank_model.py | 3 +- rag/nlp/search.py | 23 +- 42 files changed, 4822 insertions(+), 1609 deletions(-) create mode 100644 conf/models/siliconflow.json create mode 100644 internal/common/constants.go create mode 100644 internal/service/generator.go create mode 100644 internal/service/load_prompt.go create mode 100644 internal/service/metadata_filter.go create mode 100644 internal/service/nlp/retrieval.go create mode 100644 internal/service/tag.go diff --git a/api/apps/chunk_app.py b/api/apps/chunk_app.py index c7dc45b004..99159c878d 100644 --- a/api/apps/chunk_app.py +++ b/api/apps/chunk_app.py @@ -157,6 +157,7 @@ async def retrieval_test(): if ck["content_with_weight"]: ranks["chunks"].insert(0, ck) ranks["chunks"] = settings.retriever.retrieval_by_children(ranks["chunks"], tenant_ids) + ranks["total"] = len(ranks["chunks"]) for c in ranks["chunks"]: c.pop("vector", None) diff --git a/conf/models/siliconflow.json b/conf/models/siliconflow.json new file mode 100644 index 0000000000..80acb6c8ba --- /dev/null +++ b/conf/models/siliconflow.json @@ -0,0 +1,26 @@ +{ + "name": "SILICONFLOW", + "tags": "LLM,TEXT EMBEDDING,TEXT RE-RANK,IMAGE2TEXT", + "url": { + "default": "https://api.siliconflow.cn/v1" + }, + "url_suffix": { + "chat": "chat/completions", + "async_chat": "async/chat/completions", + "async_result": "async-result", + "embedding": "embedding", + "rerank": "rerank" + }, + "models": [ + { + "name": "BAAI/bge-reranker-v2-m3", + "max_tokens": 8192, + "model_types": [ + "rerank" + ], + "features": {} + } + ] +} + + diff --git a/go.mod b/go.mod index 9f06faffc6..f3c1021708 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,7 @@ require ( github.com/aws/aws-sdk-go-v2/credentials v1.19.11 github.com/aws/aws-sdk-go-v2/service/s3 v1.96.4 github.com/aws/smithy-go v1.24.2 + github.com/cespare/xxhash/v2 v2.3.0 github.com/elastic/go-elasticsearch/v8 v8.19.1 github.com/gin-gonic/gin v1.9.1 github.com/google/uuid v1.6.0 @@ -43,7 +44,6 @@ require ( github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.16 // indirect github.com/aws/aws-sdk-go-v2/service/sts v1.41.8 // indirect github.com/bytedance/sonic v1.9.1 // indirect - github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/dustin/go-humanize v1.0.1 // indirect @@ -106,4 +106,4 @@ require ( gopkg.in/ini.v1 v1.67.0 // indirect ) -replace github.com/infiniflow/infinity-go-sdk => github.com/infiniflow/infinity/go v0.0.0-20260331112649-9bcd52a3d364 +replace github.com/infiniflow/infinity-go-sdk => github.com/infiniflow/infinity/go v0.0.0-20260424025959-72028e662929 diff --git a/go.sum b/go.sum index fe150a81b9..5e9818e0e7 100644 --- a/go.sum +++ b/go.sum @@ -98,8 +98,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4= github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ= -github.com/infiniflow/infinity/go v0.0.0-20260331112649-9bcd52a3d364 h1:0v5TjSirmCAUX3oaIV8Rd9d5B+kHPdymveETUU8OcC0= -github.com/infiniflow/infinity/go v0.0.0-20260331112649-9bcd52a3d364/go.mod h1:hw3z5AwNFsGy1cdrE0Mfjot2y9jqVHTxBufUx9VzZ+0= +github.com/infiniflow/infinity/go v0.0.0-20260424025959-72028e662929 h1:0M1BNouFVpnF12XEmF/42aR8CRU0bt/rMEVEsRUtSfQ= +github.com/infiniflow/infinity/go v0.0.0-20260424025959-72028e662929/go.mod h1:hw3z5AwNFsGy1cdrE0Mfjot2y9jqVHTxBufUx9VzZ+0= github.com/iromli/go-itsdangerous v0.0.0-20220223194502-9c8bef8dac6a h1:Inib12UR9HAfBubrGNraPjKt/Cu8xPbTJbC50+0wP5U= github.com/iromli/go-itsdangerous v0.0.0-20220223194502-9c8bef8dac6a/go.mod h1:8N0Hlye5Lzw+H/yHWpZMkT0QLA+iOHG7KLdvAm95DZg= github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= diff --git a/internal/cli/user_parser.go b/internal/cli/user_parser.go index a597ac64cf..951c389326 100644 --- a/internal/cli/user_parser.go +++ b/internal/cli/user_parser.go @@ -1907,7 +1907,7 @@ func (p *Parser) parseInsertDatasetFromFile() (*Command, error) { } // Internal CLI for GO -// parseInsertMetadataFromFile parses: INSERT INTO METADATA FROM FILE "file_path" +// parseInsertMetadataFromFile parses: INSERT METADATA FROM FILE "file_path" func (p *Parser) parseInsertMetadataFromFile() (*Command, error) { p.nextToken() // consume METADATA @@ -2617,6 +2617,7 @@ func (p *Parser) parseUpdateCommand() (*Command, error) { return nil, fmt.Errorf("unknown UPDATE target: %s", p.curToken.Value) } +// Internal CLI for GO // parseUpdateChunk parses: UPDATE CHUNK 'chunk_id' OF DATASET 'dataset_name' SET '{"content": "..."}' func (p *Parser) parseUpdateChunk() (*Command, error) { p.nextToken() // consume CHUNK diff --git a/internal/common/constants.go b/internal/common/constants.go new file mode 100644 index 0000000000..c9d730727a --- /dev/null +++ b/internal/common/constants.go @@ -0,0 +1,8 @@ +package common + +const ( + // PAGERANK_FLD is the field name for pagerank score + PAGERANK_FLD = "pagerank_fea" + // TAG_FLD is the field name for tag features + TAG_FLD = "tag_feas" +) diff --git a/internal/dao/tenant_llm.go b/internal/dao/tenant_llm.go index c57ca6f32d..63ef0eecef 100644 --- a/internal/dao/tenant_llm.go +++ b/internal/dao/tenant_llm.go @@ -17,6 +17,7 @@ package dao import ( + "fmt" "ragflow/internal/entity" ) @@ -28,6 +29,16 @@ func NewTenantLLMDAO() *TenantLLMDAO { return &TenantLLMDAO{} } +// GetByID get tenant LLM by primary key ID +func (dao *TenantLLMDAO) GetByID(id int64) (*entity.TenantLLM, error) { + var tenantLLM entity.TenantLLM + err := DB.Where("id = ?", id).First(&tenantLLM).Error + if err != nil { + return nil, err + } + return &tenantLLM, nil +} + // GetByTenantAndModelName get tenant LLM by tenant ID and model name func (dao *TenantLLMDAO) GetByTenantAndModelName(tenantID, providerName string, modelName string) (*entity.TenantLLM, error) { var tenantLLM entity.TenantLLM @@ -38,6 +49,16 @@ func (dao *TenantLLMDAO) GetByTenantAndModelName(tenantID, providerName string, return &tenantLLM, nil } +// GetByTenantNameAndType get tenant LLM by tenant ID, model name, and model type +func (dao *TenantLLMDAO) GetByTenantNameAndType(tenantID, modelName string, modelType entity.ModelType) (*entity.TenantLLM, error) { + var tenantLLM entity.TenantLLM + err := DB.Where("tenant_id = ? AND llm_name = ? AND model_type = ?", tenantID, modelName, modelType).First(&tenantLLM).Error + if err != nil { + return nil, err + } + return &tenantLLM, nil +} + // GetByTenantAndType get tenant LLM by tenant ID and model type func (dao *TenantLLMDAO) GetByTenantAndType(tenantID string, modelType entity.ModelType) (*entity.TenantLLM, error) { var tenantLLM entity.TenantLLM @@ -268,3 +289,50 @@ func (dao *TenantLLMDAO) GetByTenantIDLLMNameAndFactory(tenantID, llmName, facto } return &tenantLLM, nil } + +// LookupTenantLLMByID looks up a TenantLLM record by ID and returns the record plus composite model name. +func LookupTenantLLMByID(tenantLLMDao *TenantLLMDAO, id int64) (*entity.TenantLLM, string, error) { + tenantLLM, err := tenantLLMDao.GetByID(id) + if err != nil { + return nil, "", fmt.Errorf("failed to get tenant_llm by id %d: %w", id, err) + } + if tenantLLM == nil || tenantLLM.LLMName == nil || *tenantLLM.LLMName == "" { + return nil, "", fmt.Errorf("tenant_llm record not found for id %d", id) + } + compositeName := fmt.Sprintf("%s@%s", *tenantLLM.LLMName, tenantLLM.LLMFactory) + return tenantLLM, compositeName, nil +} + +// LookupTenantLLMByName looks up a TenantLLM record by tenant name and model type. +func LookupTenantLLMByName(tenantLLMDao *TenantLLMDAO, tenantID, name string, modelType entity.ModelType) (*entity.TenantLLM, string, error) { + // Parse factory from name if present (e.g., "model@Factory") + modelName, factory := splitModelNameAndFactory(name) + + // If factory is found, use factory-based lookup + if factory != "" { + return LookupTenantLLMByFactory(tenantLLMDao, tenantID, factory, modelName, modelType) + } + + tenantLLM, err := tenantLLMDao.GetByTenantNameAndType(tenantID, modelName, modelType) + if err != nil { + return nil, "", fmt.Errorf("failed to get tenant_llm by name %s: %w", name, err) + } + if tenantLLM == nil || tenantLLM.LLMName == nil || *tenantLLM.LLMName == "" { + return nil, "", fmt.Errorf("tenant_llm record not found for name %s", name) + } + compositeName := fmt.Sprintf("%s@%s", *tenantLLM.LLMName, tenantLLM.LLMFactory) + return tenantLLM, compositeName, nil +} + +// LookupTenantLLMByFactory looks up a TenantLLM record by tenant, factory, and model name. +func LookupTenantLLMByFactory(tenantLLMDao *TenantLLMDAO, tenantID, factory, name string, modelType entity.ModelType) (*entity.TenantLLM, string, error) { + tenantLLM, err := tenantLLMDao.GetByTenantFactoryAndModelName(tenantID, factory, name) + if err != nil { + return nil, "", fmt.Errorf("failed to get tenant_llm by factory %s and name %s: %w", factory, name, err) + } + if tenantLLM == nil || tenantLLM.LLMName == nil || *tenantLLM.LLMName == "" { + return nil, "", fmt.Errorf("tenant_llm record not found for factory %s and name %s", factory, name) + } + compositeName := fmt.Sprintf("%s@%s", *tenantLLM.LLMName, tenantLLM.LLMFactory) + return tenantLLM, compositeName, nil +} diff --git a/internal/engine/elasticsearch/get.go b/internal/engine/elasticsearch/get.go index a2a4071260..625bacdda7 100644 --- a/internal/engine/elasticsearch/get.go +++ b/internal/engine/elasticsearch/get.go @@ -19,38 +19,31 @@ package elasticsearch import ( "context" "fmt" + + "ragflow/internal/engine/types" ) // GetChunk gets a chunk by ID func (e *elasticsearchEngine) GetChunk(ctx context.Context, indexName, chunkID string, kbIDs []string) (interface{}, error) { - // Build query to get the chunk by ID - query := map[string]interface{}{ - "term": map[string]interface{}{ + // Build unified search request to get the chunk by ID + searchReq := &types.SearchRequest{ + IndexNames: []string{indexName}, + Limit: 1, + Offset: 0, + Filter: map[string]interface{}{ "id": chunkID, }, } - searchReq := &SearchRequest{ - IndexNames: []string{indexName}, - Query: query, - Size: 1, - From: 0, - } - // Execute search - result, err := e.Search(ctx, searchReq) + searchResp, err := e.Search(ctx, searchReq) if err != nil { return nil, fmt.Errorf("failed to search: %w", err) } - esResp, ok := result.(*SearchResponse) - if !ok { - return nil, fmt.Errorf("invalid search response type") - } - - if len(esResp.Hits.Hits) == 0 { + if len(searchResp.Chunks) == 0 { return nil, nil } - return esResp.Hits.Hits[0].Source, nil -} + return searchResp.Chunks[0], nil +} \ No newline at end of file diff --git a/internal/engine/elasticsearch/search.go b/internal/engine/elasticsearch/search.go index c433829520..1f3935b069 100644 --- a/internal/engine/elasticsearch/search.go +++ b/internal/engine/elasticsearch/search.go @@ -22,8 +22,6 @@ import ( "encoding/json" "fmt" "io" - "strconv" - "strings" "github.com/elastic/go-elasticsearch/v8/esapi" "go.uber.org/zap" @@ -32,18 +30,6 @@ import ( "ragflow/internal/logger" ) -// SearchRequest Elasticsearch search request (legacy, kept for backward compatibility) -type SearchRequest struct { - IndexNames []string - Query map[string]interface{} - Filters map[string]interface{} // Filter conditions (e.g., kb_id, doc_id, available_int) - Size int - From int - Highlight map[string]interface{} - Source []string - Sort []interface{} -} - // SearchResponse Elasticsearch search response type SearchResponse struct { Hits struct { @@ -59,49 +45,59 @@ type SearchResponse struct { Aggregations map[string]interface{} `json:"aggregations"` } -// Search executes search (supports both unified engine.SearchRequest and legacy SearchRequest) -func (e *elasticsearchEngine) Search(ctx context.Context, req interface{}) (interface{}, error) { - - switch searchReq := req.(type) { - case *types.SearchRequest: - return e.searchUnified(ctx, searchReq) - case *SearchRequest: - return e.searchLegacy(ctx, searchReq) - default: - return nil, fmt.Errorf("invalid search request type: %T", req) - } +// Search executes search with unified types.SearchRequest +func (e *elasticsearchEngine) Search(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error) { + return e.searchUnified(ctx, req) } -// searchUnified handles the unified engine.SearchRequest -func (e *elasticsearchEngine) searchUnified(ctx context.Context, req *types.SearchRequest) (*types.SearchResponse, error) { +// searchUnified handles the unified types.SearchRequest +func (e *elasticsearchEngine) searchUnified(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error) { if len(req.IndexNames) == 0 { return nil, fmt.Errorf("index names cannot be empty") } // Build pagination parameters - offset, limit := calculatePagination(req.Page, req.Size, req.TopK) + offset := req.Offset + limit := req.Limit + if limit <= 0 { + limit = 30 // default ES size + } // Build filter clauses (default: available=1, meaning available_int >= 1) // Reference: rag/utils/es_conn.py L60-L78 - filterClauses := buildFilterClauses(req.KbIDs, req.DocIDs, 1) + filterClauses := buildFilterClauses(req.KbIDs, 1) // Build search query body queryBody := make(map[string]interface{}) - // Use MatchText if available (from QueryBuilder), otherwise use original Question - matchText := req.MatchText - if matchText == "" { - matchText = req.Question + // Determine search type from MatchExprs + var matchText string + var matchDense interface{} + var textWeight float64 = 1.0 + var hasVectorMatch bool + + for _, expr := range req.MatchExprs { + if expr == nil { + continue + } + switch e := expr.(type) { + case string: + matchText = e + case *types.MatchDenseExpr: + hasVectorMatch = true + matchDense = e + textWeight = 0.3 // default, should be passed via SimilarityThreshold + } } var vectorFieldName string - if req.KeywordOnly || len(req.Vector) == 0 { + if !hasVectorMatch { // Keyword-only search queryBody["query"] = buildESKeywordQuery(matchText, filterClauses, 1.0) } else { // Hybrid search: keyword + vector - // Calculate text weight - textWeight := 1.0 - req.VectorSimilarityWeight + // Calculate text weight (use SimilarityThreshold as text weight if provided) + // Build boolean query for text match and filters boolQuery := buildESKeywordQuery(matchText, filterClauses, 1.0) // Add boost to the bool query (as in Python code) @@ -109,30 +105,49 @@ func (e *elasticsearchEngine) searchUnified(ctx context.Context, req *types.Sear boolMap["boost"] = textWeight } // Build kNN query - dimension := len(req.Vector) - var fieldBuilder strings.Builder - fieldBuilder.WriteString("q_") - fieldBuilder.WriteString(strconv.Itoa(dimension)) - fieldBuilder.WriteString("_vec") - vectorFieldName = fieldBuilder.String() + var vectorData []float64 + if md, ok := matchDense.(*types.MatchDenseExpr); ok { + vectorData = md.EmbeddingData + vectorFieldName = md.VectorColumnName + k := md.TopN + if k <= 0 { + k = req.Limit + } + if k <= 0 { + k = 1024 + } + numCandidates := k * 2 - k := req.TopK - if k <= 0 { - k = 1024 - } - numCandidates := k * 2 + knnQuery := map[string]interface{}{ + "field": vectorFieldName, + "query_vector": vectorData, + "k": k, + "num_candidates": numCandidates, + "filter": boolQuery, + "similarity": 0.0, + } - knnQuery := map[string]interface{}{ - "field": vectorFieldName, - "query_vector": req.Vector, - "k": k, - "num_candidates": numCandidates, - "filter": boolQuery, - "similarity": req.SimilarityThreshold, + queryBody["knn"] = knnQuery + queryBody["query"] = boolQuery } - queryBody["knn"] = knnQuery - queryBody["query"] = boolQuery + // Add vector column to Source fields (matching Python ES: src.append(f"q_{len(q_vec)}_vec")) + // Only modify Source if it was explicitly set by the caller + if vectorFieldName != "" && len(req.SelectFields) > 0 { + sourceFields := req.SelectFields + // Check if vector column already in source + found := false + for _, f := range sourceFields { + if f == vectorFieldName { + found = true + break + } + } + if !found { + sourceFields = append(sourceFields, vectorFieldName) + } + req.SelectFields = sourceFields + } } queryBody["size"] = limit @@ -179,129 +194,12 @@ func (e *elasticsearchEngine) searchUnified(ctx context.Context, req *types.Sear // Convert to unified response chunks := convertESResponse(&esResp, vectorFieldName) - return &types.SearchResponse{ + return &types.SearchResult{ Chunks: chunks, Total: esResp.Hits.Total.Value, }, nil } -// searchLegacy handles the legacy elasticsearch.SearchRequest (backward compatibility) -func (e *elasticsearchEngine) searchLegacy(ctx context.Context, searchReq *SearchRequest) (*SearchResponse, error) { - if len(searchReq.IndexNames) == 0 { - return nil, fmt.Errorf("index names cannot be empty") - } - - // Build search query - queryBody := make(map[string]interface{}) - - // Process Filters first - convert to Elasticsearch filter clauses - var filterClauses []map[string]interface{} - if searchReq.Filters != nil && len(searchReq.Filters) > 0 { - for field, value := range searchReq.Filters { - switch v := value.(type) { - case map[string]interface{}: - filterClauses = append(filterClauses, map[string]interface{}{ - field: v, - }) - default: - filterClauses = append(filterClauses, map[string]interface{}{ - "term": map[string]interface{}{ - field: v, - }, - }) - } - } - } - - if searchReq.Query != nil { - queryCopy := make(map[string]interface{}) - for k, v := range searchReq.Query { - queryCopy[k] = v - } - - if knnValue, ok := queryCopy["knn"]; ok { - queryBody["knn"] = knnValue - delete(queryCopy, "knn") - } - - if len(queryCopy) > 0 { - if len(filterClauses) > 0 { - queryBody["query"] = map[string]interface{}{ - "bool": map[string]interface{}{ - "must": queryCopy, - "filter": filterClauses, - }, - } - } else { - queryBody["query"] = queryCopy - } - } else if len(filterClauses) > 0 { - queryBody["query"] = map[string]interface{}{ - "bool": map[string]interface{}{ - "filter": filterClauses, - }, - } - } - } else if len(filterClauses) > 0 { - queryBody["query"] = map[string]interface{}{ - "bool": map[string]interface{}{ - "filter": filterClauses, - }, - } - } - if searchReq.Size > 0 { - queryBody["size"] = searchReq.Size - } - if searchReq.From > 0 { - queryBody["from"] = searchReq.From - } - if searchReq.Highlight != nil { - queryBody["highlight"] = searchReq.Highlight - } - if len(searchReq.Source) > 0 { - queryBody["_source"] = searchReq.Source - } - if len(searchReq.Sort) > 0 { - queryBody["sort"] = searchReq.Sort - } - - var buf bytes.Buffer - if err := json.NewEncoder(&buf).Encode(queryBody); err != nil { - return nil, fmt.Errorf("error encoding query: %w", err) - } - - logger.Debug("Elasticsearch searching indices", zap.Strings("indices", searchReq.IndexNames)) - logger.Debug("Elasticsearch DSL", zap.Any("dsl", queryBody)) - - reqES := esapi.SearchRequest{ - Index: searchReq.IndexNames, - Body: &buf, - } - - res, err := reqES.Do(ctx, e.client) - if err != nil { - return nil, fmt.Errorf("search failed: %w", err) - } - defer res.Body.Close() - - if res.IsError() { - bodyBytes, err := io.ReadAll(res.Body) - if err != nil { - logger.Error("Elasticsearch failed to read error response body", err) - } else { - logger.Warn("Elasticsearch error response", zap.String("body", string(bodyBytes))) - } - return nil, fmt.Errorf("Elasticsearch returned error: %s", res.Status()) - } - - var response SearchResponse - if err := json.NewDecoder(res.Body).Decode(&response); err != nil { - return nil, fmt.Errorf("error parsing response: %w", err) - } - - return &response, nil -} - // calculatePagination calculates offset and limit based on page, size and topK func calculatePagination(page, size, topK int) (int, int) { if page < 1 { @@ -334,7 +232,7 @@ func calculatePagination(page, size, topK int) (int, int) { // Reference: rag/utils/es_conn.py L60-L78 // When available=0: available_int < 1 // When available!=0: NOT (available_int < 1) -func buildFilterClauses(kbIDs, docIDs []string, available int) []map[string]interface{} { +func buildFilterClauses(kbIDs []string, available int) []map[string]interface{} { var filters []map[string]interface{} if len(kbIDs) > 0 { @@ -343,12 +241,6 @@ func buildFilterClauses(kbIDs, docIDs []string, available int) []map[string]inte }) } - if len(docIDs) > 0 { - filters = append(filters, map[string]interface{}{ - "terms": map[string]interface{}{"doc_id": docIDs}, - }) - } - // Add available_int filter // Reference: rag/utils/es_conn.py L63-L68 if available == 0 { @@ -526,3 +418,27 @@ func AddMustNot(query map[string]interface{}, clauses ...map[string]interface{}) } } } + +// GetFields is not implemented for Elasticsearch +func (e *elasticsearchEngine) GetFields(chunks []map[string]interface{}, fields []string) map[string]map[string]interface{} { + logger.Warn("GetFields not implemented for Elasticsearch") + return nil +} + +// GetAggregation is not implemented for Elasticsearch +func (e *elasticsearchEngine) GetAggregation(chunks []map[string]interface{}, fieldName string) []map[string]interface{} { + logger.Warn("GetAggregation not implemented for Elasticsearch") + return nil +} + +// GetHighlight is not implemented for Elasticsearch +func (e *elasticsearchEngine) GetHighlight(chunks []map[string]interface{}, keywords []string, fieldName string) map[string]string { + logger.Warn("GetHighlight not implemented for Elasticsearch") + return nil +} + +// GetDocIDs is not implemented for Elasticsearch +func (e *elasticsearchEngine) GetDocIDs(chunks []map[string]interface{}) []string { + logger.Warn("GetDocIDs not implemented for Elasticsearch") + return nil +} diff --git a/internal/engine/engine.go b/internal/engine/engine.go index 6ea188f8db..149f96ed00 100644 --- a/internal/engine/engine.go +++ b/internal/engine/engine.go @@ -30,16 +30,10 @@ const ( EngineInfinity EngineType = "infinity" ) -// SearchRequest is an alias for types.SearchRequest -type SearchRequest = types.SearchRequest - -// SearchResponse is an alias for types.SearchResponse -type SearchResponse = types.SearchResponse - // DocEngine document storage engine interface type DocEngine interface { // Search - Search(ctx context.Context, req interface{}) (interface{}, error) + Search(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error) // Dataset operations CreateDataset(ctx context.Context, indexName, datasetID string, vectorSize int, parserID string) error @@ -56,9 +50,15 @@ type DocEngine interface { // Operations for both dataset and metadata tables Delete(ctx context.Context, condition map[string]interface{}, indexName string, datasetID string) (int64, error) - DropTable(ctx context.Context, indexName string) error + DropTable(ctx context.Context, indexName string) error TableExists(ctx context.Context, indexName string) (bool, error) + // Utility functions for search result processing + GetFields(chunks []map[string]interface{}, fields []string) map[string]map[string]interface{} + GetAggregation(chunks []map[string]interface{}, fieldName string) []map[string]interface{} + GetHighlight(chunks []map[string]interface{}, keywords []string, fieldName string) map[string]string + GetDocIDs(chunks []map[string]interface{}) []string + // Health check Ping(ctx context.Context) error Close() error diff --git a/internal/engine/global.go b/internal/engine/global.go index 315dfb4baa..fb213e65f6 100644 --- a/internal/engine/global.go +++ b/internal/engine/global.go @@ -30,6 +30,7 @@ import ( var ( globalEngine DocEngine + engineType EngineType once sync.Once ) @@ -37,8 +38,9 @@ var ( func Init(cfg *server.DocEngineConfig) error { var initErr error once.Do(func() { + engineType = EngineType(cfg.Type) var err error - switch EngineType(cfg.Type) { + switch engineType { case EngineElasticsearch: globalEngine, err = elasticsearch.NewEngine(cfg.ES) case EngineInfinity: @@ -56,6 +58,11 @@ func Init(cfg *server.DocEngineConfig) error { return initErr } +// GetEngineType returns the document engine type +func GetEngineType() EngineType { + return engineType +} + // Get gets global document engine instance func Get() DocEngine { return globalEngine diff --git a/internal/engine/infinity/common.go b/internal/engine/infinity/common.go index 0837fe080d..663d50c744 100644 --- a/internal/engine/infinity/common.go +++ b/internal/engine/infinity/common.go @@ -23,8 +23,9 @@ import ( "fmt" "strings" - infinity "github.com/infiniflow/infinity-go-sdk" "ragflow/internal/logger" + + infinity "github.com/infiniflow/infinity-go-sdk" ) // Delete deletes rows from either a dataset table or metadata table. @@ -127,10 +128,10 @@ func (e *infinityEngine) TableExists(ctx context.Context, indexName string) (boo // fieldInfo represents a field in the infinity mapping schema type fieldInfo struct { Type string `json:"type"` - Default interface{} `json:"default"` - Analyzer interface{} `json:"analyzer"` // string or []string + Default interface{} `json:"default"` + Analyzer interface{} `json:"analyzer"` // string or []string IndexType interface{} `json:"index_type"` // string or map - Comment string `json:"comment"` + Comment string `json:"comment"` } // orderedFields preserves the order of fields as defined in JSON @@ -176,7 +177,22 @@ func (o *orderedFields) UnmarshalJSON(data []byte) error { return nil } -// existsCondition builds a NOT EXISTS or field!='' condition +// fieldKeyword checks if field is a keyword field +func fieldKeyword(fieldName string) bool { + if fieldName == "source_id" { + return true + } + if strings.HasSuffix(fieldName, "_kwd") && + fieldName != "knowledge_graph_kwd" && + fieldName != "docnm_kwd" && + fieldName != "important_kwd" && + fieldName != "question_kwd" { + return true + } + return false +} + +// existsCondition builds a NOT EXISTS or field!=" condition func existsCondition(field string, tableColumns map[string]struct { Type string Default interface{} @@ -228,20 +244,29 @@ func buildFilterFromCondition(condition map[string]interface{}, tableColumns map // Handle keyword fields -> filter_fulltext with converted field name if fieldKeyword(k) { - if listVal, ok := v.([]interface{}); ok { - var orConds []string - for _, item := range listVal { - if strItem, ok := item.(string); ok { - strItem = strings.ReplaceAll(strItem, "'", "''") - orConds = append(orConds, fmt.Sprintf("filter_fulltext('%s', '%s')", convertMatchingField(k), strItem)) - } + var orConds []string + addFullText := func(item string) { + item = strings.ReplaceAll(item, "'", "''") + orConds = append(orConds, fmt.Sprintf("filter_fulltext('%s', '%s')", convertMatchingField(k), item)) + } + + switch val := v.(type) { + case []string: + for _, item := range val { + addFullText(item) } - if len(orConds) > 0 { - conditions = append(conditions, "("+strings.Join(orConds, " OR ")+")") + case []interface{}: + for _, item := range val { + addFullText(fmt.Sprintf("%v", item)) } - } else if strVal, ok := v.(string); ok { - strVal = strings.ReplaceAll(strVal, "'", "''") - conditions = append(conditions, fmt.Sprintf("filter_fulltext('%s', '%s')", convertMatchingField(k), strVal)) + case string: + addFullText(val) + default: + addFullText(fmt.Sprintf("%v", val)) + } + + if len(orConds) > 0 { + conditions = append(conditions, "("+strings.Join(orConds, " OR ")+")") } continue } diff --git a/internal/engine/infinity/dataset.go b/internal/engine/infinity/dataset.go index c671ddab32..2043c6145e 100644 --- a/internal/engine/infinity/dataset.go +++ b/internal/engine/infinity/dataset.go @@ -403,7 +403,7 @@ func (e *infinityEngine) UpdateDataset(ctx context.Context, condition map[string if ok && len(qr.Data) > 0 { // Get the id column and columns to remove idCol := qr.Data["id"] - removeOpt := make(map[string]map[string][]string); // column -> value -> [ids] + removeOpt := make(map[string]map[string][]string) // column -> value -> [ids] for colName, colData := range qr.Data { if colName == "id" { diff --git a/internal/engine/infinity/get.go b/internal/engine/infinity/get.go index a8f8b58135..fe42f92837 100644 --- a/internal/engine/infinity/get.go +++ b/internal/engine/infinity/get.go @@ -21,10 +21,11 @@ import ( "fmt" "strings" - infinity "github.com/infiniflow/infinity-go-sdk" "ragflow/internal/logger" "ragflow/internal/utility" + infinity "github.com/infiniflow/infinity-go-sdk" + "go.uber.org/zap" ) @@ -114,16 +115,9 @@ func (e *infinityEngine) GetChunk(ctx context.Context, tableName, chunkID string return nil, nil } - getFields(chunk) - logger.Debug("infinity get chunk", zap.String("chunkID", chunkID), zap.Any("tables", tableNames)) - return chunk, nil -} - -// getFields applies field mappings to a chunk, similar to Python's get_fields function. -func getFields(chunk map[string]interface{}) { - // Field mappings + // Apply field mappings (same as in GetFields) // docnm -> docnm_kwd, title_tks, title_sm_tks if val, ok := chunk["docnm"].(string); ok { chunk["docnm_kwd"] = val @@ -131,6 +125,13 @@ func getFields(chunk map[string]interface{}) { chunk["title_sm_tks"] = val } + // content -> content_with_weight, content_ltks, content_sm_ltks + if val, ok := chunk["content"].(string); ok { + chunk["content_with_weight"] = val + chunk["content_ltks"] = val + chunk["content_sm_ltks"] = val + } + // important_keywords -> important_kwd (split by comma), important_tks if val, ok := chunk["important_keywords"].(string); ok { if val == "" { @@ -159,61 +160,144 @@ func getFields(chunk map[string]interface{}) { chunk["question_tks"] = []interface{}{} } - // content -> content_with_weight, content_ltks, content_sm_ltks - if val, ok := chunk["content"].(string); ok { - chunk["content_with_weight"] = val - chunk["content_ltks"] = val - chunk["content_sm_ltks"] = val - } - - // authors -> authors_tks, authors_sm_tks - if val, ok := chunk["authors"].(string); ok { - chunk["authors_tks"] = val - chunk["authors_sm_tks"] = val - } - - // position_int: convert from hex string to array format (grouped by 5) - if val, ok := chunk["position_int"].(string); ok { - chunk["position_int"] = utility.ConvertHexToPositionIntArray(val) + if posVal, ok := chunk["position_int"].(string); ok { + chunk["position_int"] = utility.ConvertHexToPositionIntArray(posVal) } else { chunk["position_int"] = []interface{}{} } - // Convert page_num_int and top_int from hex string to array - for _, colName := range []string{"page_num_int", "top_int"} { - if val, ok := chunk[colName].(string); ok && val != "" { - chunk[colName] = utility.ConvertHexToIntArray(val) + return chunk, nil +} + +// GetFields applies field mappings to chunks and returns a dict keyed by chunk ID. +// Equivalent to Python's get_fields() in infinity_conn.py. +// When fields is nil/empty, returns all fields from chunks. +func GetFields(chunks []map[string]interface{}, fields []string) map[string]map[string]interface{} { + result := make(map[string]map[string]interface{}) + if len(chunks) == 0 { + return result + } + + // If fields is provided, create a set for lookup + fieldSet := make(map[string]bool) + for _, f := range fields { + fieldSet[f] = true + } + + for _, chunk := range chunks { + // Apply field mappings + // docnm -> docnm_kwd, title_tks, title_sm_tks + if val, ok := chunk["docnm"].(string); ok { + chunk["docnm_kwd"] = val + chunk["title_tks"] = val + chunk["title_sm_tks"] = val + } + + // important_keywords -> important_kwd (split by comma), important_tks + if val, ok := chunk["important_keywords"].(string); ok { + if val == "" { + chunk["important_kwd"] = []interface{}{} + } else { + parts := strings.Split(val, ",") + chunk["important_kwd"] = parts + } + chunk["important_tks"] = val } else { - chunk[colName] = []int{} + chunk["important_kwd"] = []interface{}{} + chunk["important_tks"] = []interface{}{} + } + + // questions -> question_kwd (split by newline), question_tks + if val, ok := chunk["questions"].(string); ok { + if val == "" { + chunk["question_kwd"] = []interface{}{} + } else { + parts := strings.Split(val, "\n") + chunk["question_kwd"] = parts + } + chunk["question_tks"] = val + } else { + chunk["question_kwd"] = []interface{}{} + chunk["question_tks"] = []interface{}{} + } + + // content -> content_with_weight, content_ltks, content_sm_ltks + if val, ok := chunk["content"].(string); ok { + chunk["content_with_weight"] = val + chunk["content_ltks"] = val + chunk["content_sm_ltks"] = val + } + + // authors -> authors_tks, authors_sm_tks + if val, ok := chunk["authors"].(string); ok { + chunk["authors_tks"] = val + chunk["authors_sm_tks"] = val + } + + // position_int: convert from hex string to array format (grouped by 5) + if val, ok := chunk["position_int"].(string); ok { + chunk["position_int"] = utility.ConvertHexToPositionIntArray(val) + } + + // Convert page_num_int and top_int from hex string to array + for _, colName := range []string{"page_num_int", "top_int"} { + if val, ok := chunk[colName].(string); ok && val != "" { + chunk[colName] = utility.ConvertHexToIntArray(val) + } + } + + // Post-process: convert nil/empty values to empty slices for array-like fields + // and split _kwd fields by "###" (except knowledge_graph_kwd, docnm_kwd, important_kwd, question_kwd) + kwdNoSplit := map[string]bool{ + "knowledge_graph_kwd": true, "docnm_kwd": true, + "important_kwd": true, "question_kwd": true, + } + arrayFields := []string{ + "doc_type_kwd", "important_kwd", "important_tks", "question_tks", + "question_kwd", "authors_tks", "authors_sm_tks", "title_tks", + "title_sm_tks", "content_ltks", "content_sm_ltks", "tag_kwd", + } + for _, colName := range arrayFields { + val, ok := chunk[colName] + if !ok || val == nil || val == "" { + chunk[colName] = []interface{}{} + } else if !kwdNoSplit[colName] { + // Split by "###" for _kwd fields + if strVal, ok := val.(string); ok && strings.Contains(strVal, "###") { + parts := strings.Split(strVal, "###") + var filtered []interface{} + for _, p := range parts { + if p != "" { + filtered = append(filtered, p) + } + } + chunk[colName] = filtered + } + } + } + + // Handle row_id mapping - Infinity returns "ROW_ID" but we use "row_id()" + if val, ok := chunk["ROW_ID"]; ok { + chunk["row_id()"] = val + delete(chunk, "ROW_ID") + } + + // Build result map keyed by id + if id, ok := chunk["id"].(string); ok { + fieldMap := make(map[string]interface{}) + for field, value := range chunk { + if len(fieldSet) == 0 || fieldSet[field] { + fieldMap[field] = value + } + } + result[id] = fieldMap } } - // Post-process: convert nil/empty values to empty slices for array-like fields - // and split _kwd fields by "###" (except knowledge_graph_kwd, docnm_kwd, important_kwd, question_kwd) - kwdNoSplit := map[string]bool{ - "knowledge_graph_kwd": true, "docnm_kwd": true, - "important_kwd": true, "question_kwd": true, - } - arrayFields := []string{ - "doc_type_kwd", "important_kwd", "important_tks", "question_tks", - "question_kwd", "authors_tks", "authors_sm_tks", "title_tks", - "title_sm_tks", "content_ltks", "content_sm_ltks", - } - for _, colName := range arrayFields { - if val, ok := chunk[colName]; !ok || val == nil || val == "" { - chunk[colName] = []interface{}{} - } else if !kwdNoSplit[colName] { - // Split by "###" for _kwd fields - if strVal, ok := val.(string); ok && strings.Contains(strVal, "###") { - parts := strings.Split(strVal, "###") - var filtered []interface{} - for _, p := range parts { - if p != "" { - filtered = append(filtered, p) - } - } - chunk[colName] = filtered - } - } - } + return result +} + +// GetFields is a method wrapper for infinityEngine to satisfy DocEngine interface +func (e *infinityEngine) GetFields(chunks []map[string]interface{}, fields []string) map[string]map[string]interface{} { + return GetFields(chunks, fields) } diff --git a/internal/engine/infinity/search.go b/internal/engine/infinity/search.go index a196b4e223..e82ba35223 100644 --- a/internal/engine/infinity/search.go +++ b/internal/engine/infinity/search.go @@ -18,195 +18,473 @@ package infinity import ( "context" + "encoding/json" "fmt" + "ragflow/internal/common" "ragflow/internal/engine/types" "ragflow/internal/utility" + "regexp" + "slices" + "sort" + "strconv" "strings" - "unicode/utf8" + "unicode" + + "ragflow/internal/logger" infinity "github.com/infiniflow/infinity-go-sdk" + "go.uber.org/zap" ) -const ( - PAGERANK_FLD = "pagerank_fea" - TAG_FLD = "tag_feas" -) - -type SortType int - -const ( - SortAsc SortType = 0 - SortDesc SortType = 1 -) - -type OrderByExpr struct { - Fields []OrderByField -} - -type OrderByField struct { - Field string - Type SortType -} - -// fieldKeyword checks if field is a keyword field -func fieldKeyword(fieldName string) bool { - // Treat "*_kwd" tag-like columns as keyword lists except knowledge_graph_kwd - if fieldName == "source_id" { - return true - } - if strings.HasSuffix(fieldName, "_kwd") && - fieldName != "knowledge_graph_kwd" && - fieldName != "docnm_kwd" && - fieldName != "important_kwd" && - fieldName != "question_kwd" { - return true - } - return false -} - -// equivalentConditionToStr converts condition dict to filter string -func equivalentConditionToStr(condition map[string]interface{}, tableColumns map[string]struct { - Type string - Default interface{} -}) string { - if len(condition) == 0 { - return "" - } - - var conditions []string - - for k, v := range condition { - if !strings.HasPrefix(k, "_") { - continue - } - if v == nil || v == "" { - continue +// Search searches the Infinity engine for matching chunks. +// It supports three matching types: MatchTextExpr (full-text), MatchDenseExpr (vector), and FusionExpr (combined). +// If no match expressions are provided, Search relies solely on filter (e.g., doc_id, available_int) to find results. +func (e *infinityEngine) Search(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error) { + logger.Info("Search in Infinity started", zap.Any("indexNames", req.IndexNames)) + if logger.IsDebugEnabled() { + // Format match expressions for logging + var matchExprsStr string + for i, expr := range req.MatchExprs { + switch e := expr.(type) { + case *types.MatchTextExpr: + matchExprsStr += fmt.Sprintf(" [%d] MatchTextExpr: fields=%v, matchingText=%s, topN=%d, extraOptions=%v\n", i, e.Fields, e.MatchingText, e.TopN, e.ExtraOptions) + case *types.MatchDenseExpr: + matchExprsStr += fmt.Sprintf(" [%d] MatchDenseExpr: vectorColumn=%s, vectorSize=%d, topN=%d, extraOptions=%v\n", i, e.VectorColumnName, len(e.EmbeddingData), e.TopN, e.ExtraOptions) + case *types.FusionExpr: + matchExprsStr += fmt.Sprintf(" [%d] FusionExpr: method=%s, topN=%d, fusionParams=%v\n", i, e.Method, e.TopN, e.FusionParams) + default: + matchExprsStr += fmt.Sprintf(" [%d] unknown type\n", i) + } } + logger.Debug(fmt.Sprintf("Search request:\n"+ + " indexNames=%v\n"+ + " KbIDs=%v\n"+ + " offset=%d, limit=%d\n"+ + " SelectFields=%v\n"+ + " Filter=%v\n"+ + " MatchExprs:\n%s orderBy=%v\n"+ + " RankFeature=%v", + req.IndexNames, req.KbIDs, req.Offset, req.Limit, req.SelectFields, req.Filter, matchExprsStr, req.OrderBy, req.RankFeature)) + } - // Handle keyword fields with filter_fulltext - if fieldKeyword(k) { - if listVal, isList := v.([]interface{}); isList { - var orConds []string - for _, item := range listVal { - if strItem, ok := item.(string); ok { - strItem = strings.ReplaceAll(strItem, "'", "''") - orConds = append(orConds, fmt.Sprintf("filter_fulltext('%s', '%s')", convertMatchingField(k), strItem)) - } - } - if len(orConds) > 0 { - conditions = append(conditions, "("+strings.Join(orConds, " OR ")+")") - } - } else if strVal, ok := v.(string); ok { - strVal = strings.ReplaceAll(strVal, "'", "''") - conditions = append(conditions, fmt.Sprintf("filter_fulltext('%s', '%s')", convertMatchingField(k), strVal)) + if len(req.IndexNames) == 0 { + return nil, fmt.Errorf("index names cannot be empty") + } + + // Get retrieval parameters with defaults + pageSize := req.Limit + if pageSize <= 0 { + pageSize = 30 + } + + offset := req.Offset + if offset < 0 { + offset = 0 + } + + db, err := e.client.conn.GetDatabase(e.client.dbName) + if err != nil { + return nil, fmt.Errorf("failed to get database: %w", err) + } + + isMetadataTable := false + for _, idx := range req.IndexNames { + if strings.HasPrefix(idx, "ragflow_doc_meta_") { + isMetadataTable = true + break + } + } + + var outputColumns []string + if isMetadataTable { + outputColumns = []string{"id", "kb_id", "meta_fields"} + } else { + outputColumns = []string{ + "id", "doc_id", "kb_id", "content_ltks", "content_with_weight", + "title_tks", "docnm_kwd", "img_id", "available_int", "important_kwd", + "position_int", "page_num_int", "top_int", "chunk_order_int", + "create_timestamp_flt", "knowledge_graph_kwd", "question_kwd", "question_tks", + "doc_type_kwd", "mom_id", "tag_kwd", "pagerank_fea", "tag_feas", + } + outputColumns = convertSelectFields(outputColumns) + } + + hasTextMatch := false + hasVectorMatch := false + var matchText *types.MatchTextExpr + var matchDense *types.MatchDenseExpr + if req.MatchExprs != nil && len(req.MatchExprs) > 0 { + for _, expr := range req.MatchExprs { + if expr == nil { + continue } - } else if listVal, isList := v.([]interface{}); isList { - // Handle IN conditions - var inVals []string - for _, item := range listVal { - if strItem, ok := item.(string); ok { - strItem = strings.ReplaceAll(strItem, "'", "''") - inVals = append(inVals, fmt.Sprintf("'%s'", strItem)) - } else { - inVals = append(inVals, fmt.Sprintf("%v", item)) - } + switch e := expr.(type) { + case *types.MatchTextExpr: + hasTextMatch = true + matchText = e + case *types.MatchDenseExpr: + hasVectorMatch = true + matchDense = e } - if len(inVals) > 0 { - conditions = append(conditions, fmt.Sprintf("%s IN (%s)", k, strings.Join(inVals, ", "))) - } - } else if k == "must_not" { - // Handle must_not conditions - if mustNotMap, ok := v.(map[string]interface{}); ok { - if existsVal, ok := mustNotMap["exists"]; ok { - if existsField, ok := existsVal.(string); ok { - col, colOk := tableColumns[existsField] - if colOk && strings.Contains(strings.ToLower(col.Type), "char") { - conditions = append(conditions, fmt.Sprintf(" %s!='' ", existsField)) - } else { - conditions = append(conditions, fmt.Sprintf("%s!=null", existsField)) - } - } - } - } - } else if strVal, ok := v.(string); ok { - strVal = strings.ReplaceAll(strVal, "'", "''") - conditions = append(conditions, fmt.Sprintf("%s='%s'", k, strVal)) - } else if k == "exists" { - if existsField, ok := v.(string); ok { - col, colOk := tableColumns[existsField] - if colOk && strings.Contains(strings.ToLower(col.Type), "char") { - conditions = append(conditions, fmt.Sprintf(" %s!='' ", existsField)) - } else { - conditions = append(conditions, fmt.Sprintf("%s!=null", existsField)) - } + } + } + + if hasTextMatch || hasVectorMatch { + if hasTextMatch { + outputColumns = append(outputColumns, "score()") + } else if hasVectorMatch { + outputColumns = append(outputColumns, "similarity()") + } + if !slices.Contains(outputColumns, common.PAGERANK_FLD) { + outputColumns = append(outputColumns, common.PAGERANK_FLD) + } + if !slices.Contains(outputColumns, common.TAG_FLD) { + outputColumns = append(outputColumns, common.TAG_FLD) + } + } + + if !slices.Contains(outputColumns, "row_id") && !slices.Contains(outputColumns, "row_id()") { + outputColumns = append(outputColumns, "row_id()") + } + + outputColumns = convertSelectFields(outputColumns) + if hasVectorMatch && matchDense != nil && matchDense.VectorColumnName != "" { + outputColumns = append(outputColumns, matchDense.VectorColumnName) + } + + var filterParts []string + if isMetadataTable && len(req.KbIDs) > 0 && req.KbIDs[0] != "" { + kbIDs := req.KbIDs + if len(kbIDs) == 1 { + filterParts = append(filterParts, fmt.Sprintf("kb_id = '%s'", kbIDs[0])) + } else { + kbIDStr := strings.Join(kbIDs, "', '") + filterParts = append(filterParts, fmt.Sprintf("kb_id IN ('%s')", kbIDStr)) + } + } + + if !isMetadataTable && (hasTextMatch || hasVectorMatch) { + if req.Filter != nil { + if availInt, ok := req.Filter["available_int"]; ok { + filterParts = append(filterParts, fmt.Sprintf("available_int=%v", availInt)) + } else { + filterParts = append(filterParts, "available_int=1") } } else { - conditions = append(conditions, fmt.Sprintf("%s=%v", k, v)) + filterParts = append(filterParts, "available_int=1") } } - if len(conditions) == 0 { - return "" + // Build filter string from req.Filter + if req.Filter != nil { + filterCopy := req.Filter + if !isMetadataTable { + filterCopy = make(map[string]interface{}) + for k, v := range req.Filter { + if k != "kb_id" { + filterCopy[k] = v + } + } + } + + condStr := equivalentConditionToStr(filterCopy) + if condStr != "" { + filterParts = append(filterParts, condStr) + } } - return strings.Join(conditions, " AND ") -} + filterStr := strings.Join(filterParts, " AND ") -// SearchRequest Infinity search request (legacy, kept for backward compatibility) -type SearchRequest struct { - TableName string - ColumnNames []string - MatchText *MatchTextExpr - MatchDense *MatchDenseExpr - Fusion *FusionExpr - Offset int - Limit int - Filter map[string]interface{} - OrderBy *OrderByExpr -} - -// SearchResponse Infinity search response -type SearchResponse struct { - Rows []map[string]interface{} - Total int64 -} - -// MatchTextExpr text match expression -type MatchTextExpr struct { - Fields []string - MatchingText string - TopN int - ExtraOptions map[string]interface{} -} - -// MatchDenseExpr vector match expression -type MatchDenseExpr struct { - VectorColumnName string - EmbeddingData []float64 - EmbeddingDataType string - DistanceType string - TopN int - ExtraOptions map[string]interface{} -} - -// FusionExpr fusion expression -type FusionExpr struct { - Method string - TopN int - Weights []float64 - FusionParams map[string]interface{} -} - -// Search executes search (supports unified engine.SearchRequest only) -func (e *infinityEngine) Search(ctx context.Context, req interface{}) (interface{}, error) { - switch searchReq := req.(type) { - case *types.SearchRequest: - return e.searchUnified(ctx, searchReq) - default: - return nil, fmt.Errorf("invalid search request type: %T", req) + orderBy := req.OrderBy + var rankFeature map[string]float64 + if req.RankFeature != nil { + rankFeature = req.RankFeature } + + var fusionExpr *types.FusionExpr + if len(req.MatchExprs) > 2 { + if fe, ok := req.MatchExprs[2].(*types.FusionExpr); ok { + fusionExpr = fe + } + } + + var allResults []map[string]interface{} + totalHits := int64(0) + + for _, indexName := range req.IndexNames { + var tableNames []string + if strings.HasPrefix(indexName, "ragflow_doc_meta_") { + tableNames = []string{indexName} + } else { + kbIDs := req.KbIDs + if len(kbIDs) == 0 { + kbIDs = []string{""} + } + for _, kbID := range kbIDs { + if kbID == "" { + tableNames = append(tableNames, indexName) + } else { + tableNames = append(tableNames, fmt.Sprintf("%s_%s", indexName, kbID)) + } + } + } + + minMatch := 0.3 + + var questionText string + var vectorData []float64 + textTopN := pageSize + var originalQuery string + if matchText != nil { + questionText = matchText.MatchingText + textTopN = int(matchText.TopN) + if matchText.ExtraOptions != nil { + if oq, ok := matchText.ExtraOptions["original_query"].(string); ok { + originalQuery = oq + } + } + } + if matchDense != nil { + vectorData = matchDense.EmbeddingData + } + + for _, tableName := range tableNames { + tbl, err := db.GetTable(tableName) + if err != nil { + continue + } + table := tbl.Output(outputColumns) + + var textFields []string + if matchText != nil && len(matchText.Fields) > 0 { + textFields = matchText.Fields + } else { + textFields = []string{ + "title_tks^10", + "title_sm_tks^5", + "important_kwd^30", + "important_tks^20", + "question_tks^20", + "content_ltks^2", + "content_sm_ltks", + } + } + + // Convert field names for Infinity + var convertedFields []string + for _, f := range textFields { + cf := convertMatchingField(f) + convertedFields = append(convertedFields, cf) + } + fields := strings.Join(convertedFields, ",") + + hasTextMatch := questionText != "" + hasVectorMatch := len(vectorData) > 0 + // Add text match if question is provided + if hasTextMatch { + extraOptions := map[string]string{ + "minimum_should_match": fmt.Sprintf("%d%%", int(minMatch*100)), + } + + if filterStr != "" { + extraOptions["filter"] = filterStr + } + + if rankFeature != nil { + var rankFeaturesList []string + for featureName, weight := range rankFeature { + rankFeaturesList = append(rankFeaturesList, fmt.Sprintf("%s^%s^%.0f", common.TAG_FLD, featureName, weight)) + } + if len(rankFeaturesList) > 0 { + extraOptions["rank_features"] = strings.Join(rankFeaturesList, ",") + } + } + + if originalQuery != "" { + extraOptions["original_query"] = originalQuery + } + + table = table.MatchText(fields, questionText, textTopN, extraOptions) + + logger.Debug(fmt.Sprintf( + "MatchTextExpr:\n"+ + " fields=%s\n"+ + " matching_text=%s\n"+ + " topn=%d\n"+ + " extra_options=%v", + fields, questionText, textTopN, extraOptions, + )) + } + + // Add vector match if provided + if hasVectorMatch { + vectorSize := len(vectorData) + fieldName := fmt.Sprintf("q_%d_vec", vectorSize) + dataType := "float" + distanceType := "cosine" + + if matchDense != nil { + if matchDense.VectorColumnName != "" { + fieldName = matchDense.VectorColumnName + } + if matchDense.EmbeddingDataType != "" { + dataType = matchDense.EmbeddingDataType + } + if matchDense.DistanceType != "" { + distanceType = matchDense.DistanceType + } + } + + vectorTopN := pageSize + if matchDense != nil && matchDense.TopN > 0 { + vectorTopN = int(matchDense.TopN) + } + + denseFilterStr := filterStr + if denseFilterStr == "" { + denseFilterStr = "available_int=1" + } + + if hasTextMatch { + fieldsStr := strings.Join(convertedFields, ",") + filterFulltext := fmt.Sprintf("filter_fulltext('%s', '%s')", fieldsStr, questionText) + denseFilterStr = fmt.Sprintf("(%s) AND %s", denseFilterStr, filterFulltext) + } + extraOptions := map[string]string{ + "threshold": utility.FloatToString(0.0), + "filter": denseFilterStr, + } + + logger.Debug(fmt.Sprintf( + "MatchDenseExpr:\n"+ + " field=%s\n"+ + " topn=%d\n"+ + " extra_options=%v", + fieldName, vectorTopN, extraOptions, + )) + + table = table.MatchDense(fieldName, vectorData, dataType, distanceType, vectorTopN, extraOptions) + } + + // Add fusion (for text + vector combination) + if hasTextMatch && hasVectorMatch && fusionExpr != nil { + fusionMethod := fusionExpr.Method + fusionTopK := fusionExpr.TopN + if fusionTopK == 0 { + fusionTopK = pageSize + } + fusionParams := map[string]interface{}{ + "normalize": "atan", + } + if fusionExpr.FusionParams != nil { + for k, v := range fusionExpr.FusionParams { + fusionParams[k] = v + } + } + logger.Debug(fmt.Sprintf( + "FusionExpr:\n"+ + " method=%s\n"+ + " topn=%d\n"+ + " fusion_params=%v", + fusionMethod, fusionTopK, fusionParams, + )) + + table = table.Fusion(fusionMethod, fusionTopK, fusionParams) + } + + // Add order_by if provided + if orderBy != nil && len(orderBy.Fields) > 0 { + var sortFields [][2]interface{} + for _, orderField := range orderBy.Fields { + sortType := infinity.SortTypeAsc + if orderField.Type == types.SortDesc { + sortType = infinity.SortTypeDesc + } + sortFields = append(sortFields, [2]interface{}{orderField.Field, sortType}) + } + table = table.Sort(sortFields) + } + + // Add filter when there's no text/vector match (like metadata queries) + if !hasTextMatch && !hasVectorMatch && filterStr != "" { + logger.Debug(fmt.Sprintf("Adding filter for no-match query: %s", filterStr)) + table = table.Filter(filterStr) + } + + // Set limit and offset + table = table.Limit(pageSize) + if offset > 0 { + table = table.Offset(offset) + } + + // Request total_hits_count from Infinity + table = table.Option(map[string]interface{}{"total_hits_count": true}) + + // Execute query + df, err := table.ToDataFrame() + if err != nil { + continue + } + + // Convert DataFrame to chunks format (column-oriented to row-oriented) + chunks := make([]map[string]interface{}, 0) + for colName, colData := range df.ColumnData { + for i, val := range colData { + for len(chunks) <= i { + chunks = append(chunks, make(map[string]interface{})) + } + chunks[i][colName] = val + } + } + + // Apply field name mapping and row_id handling + GetFields(chunks, nil) + + // Parse total_hits_count from ExtraInfo + var tableTotal int64 + if df.ExtraInfo != "" { + var extraResult map[string]interface{} + if err := json.Unmarshal([]byte(df.ExtraInfo), &extraResult); err == nil { + if count, ok := extraResult["total_hits_count"].(float64); ok { + tableTotal = int64(count) + } + } + } + + searchResult := &types.SearchResult{ + Chunks: chunks, + Total: tableTotal, + } + + allResults = append(allResults, searchResult.Chunks...) + totalHits += searchResult.Total + } + } + + if hasTextMatch || hasVectorMatch { + scoreColumn := "" + if hasTextMatch { + scoreColumn = "SCORE" + } else if hasVectorMatch { + scoreColumn = "SIMILARITY" + } + allResults = calculateScores(allResults, scoreColumn) + allResults = sortByScore(allResults, len(allResults)) + } + + if len(allResults) > pageSize { + allResults = allResults[:pageSize] + } + + logger.Info("Search in Infinity completed", zap.Any("indexNames", req.IndexNames), zap.Int("returnedRows", len(allResults)), zap.Int64("totalHits", totalHits)) + + return &types.SearchResult{ + Chunks: allResults, + Total: totalHits, + }, nil } -// convertSelectFields converts field names to Infinity format +// convertSelectFields converts RAG field names to Infinity column names for SELECT (output_columns). +// Example: docnm_kwd → docnm, content_ltks → content func convertSelectFields(output []string) []string { fieldMapping := map[string]string{ "docnm_kwd": "docnm", @@ -262,69 +540,8 @@ func convertSelectFields(output []string) []string { return result } -// isChinese checks if a string contains Chinese characters -func isChinese(s string) bool { - for _, r := range s { - if '\u4e00' <= r && r <= '\u9fff' { - return true - } - } - return false -} - -// hasSubTokens checks if the text has sub-tokens after fine-grained tokenization -// - Returns False if len < 3 -// - Returns False if text is only ASCII alphanumeric -// - Returns True otherwise (meaning there are sub-tokens) -func hasSubTokens(s string) bool { - if utf8.RuneCountInString(s) < 3 { - return false - } - isASCIIOnly := true - for _, r := range s { - if r > 127 { - isASCIIOnly = false - break - } - } - if isASCIIOnly { - // Check if it's only alphanumeric and allowed special chars - for _, r := range s { - if !((r >= '0' && r <= '9') || (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || r == '.' || r == '+' || r == '#' || r == '_' || r == '*' || r == '-') { - isASCIIOnly = false - break - } - } - if isASCIIOnly { - return false - } - } - // Has sub-tokens if it's Chinese and length >= 3 - return isChinese(s) -} - -// formatQuestion formats the question -// - If len < 3: returns ((query)^1.0) -// - If has sub-tokens: adds fuzzy search ((query OR "query" OR ("query"~2)^0.5)^1.0) -// - Otherwise: returns ((query)^1.0) -func formatQuestion(question string) string { - // Trim whitespace - question = strings.TrimSpace(question) - fmt.Printf("[DEBUG formatQuestion] input: %q, len: %d, hasSubTokens: %v\n", question, len(question), hasSubTokens(question)) - - // If no sub-tokens, use simple format - if !hasSubTokens(question) { - result := fmt.Sprintf("((%s)^1.0)", question) - fmt.Printf("[DEBUG formatQuestion] simple: %s\n", result) - return result - } - - result := fmt.Sprintf("((%s OR \"%s\" OR (\"%s\"~2)^0.5)^1.0)", question, question, question) - fmt.Printf("[DEBUG formatQuestion] fuzzy: %s\n", result) - return result -} - -// convertMatchingField converts field names for matching +// convertMatchingField converts RAG field names to Infinity full-text index names for MATCH expressions. +// Example: docnm_kwd → docnm@ft_docnm_rag_coarse, content_ltks → content@ft_content_rag_coarse func convertMatchingField(fieldWeightStr string) string { // Split on ^ to get field name parts := strings.Split(fieldWeightStr, "^") @@ -354,309 +571,178 @@ func convertMatchingField(fieldWeightStr string) string { return strings.Join(parts, "^") } -// searchUnified handles the unified engine.SearchRequest -func (e *infinityEngine) searchUnified(ctx context.Context, req *types.SearchRequest) (*types.SearchResponse, error) { - if len(req.IndexNames) == 0 { - return nil, fmt.Errorf("index names cannot be empty") - } - - // Get retrieval parameters with defaults - topK := req.TopK - if topK <= 0 { - topK = 1024 - } - - pageSize := req.Size - if pageSize <= 0 { - pageSize = 30 - } - - offset := (req.Page - 1) * pageSize - if offset < 0 { - offset = 0 - } - - // Get database - db, err := e.client.conn.GetDatabase(e.client.dbName) - if err != nil { - return nil, fmt.Errorf("failed to get database: %w", err) - } - - // Determine if this is a metadata table - isMetadataTable := false - for _, idx := range req.IndexNames { - if strings.HasPrefix(idx, "ragflow_doc_meta_") { - isMetadataTable = true - break - } - } - - // Build output columns - // For metadata tables, only use: id, kb_id, meta_fields - // For chunk tables, use all the standard fields - var outputColumns []string - if isMetadataTable { - outputColumns = []string{"id", "kb_id", "meta_fields"} - } else { - outputColumns = []string{ - "id", - "doc_id", - "kb_id", - "content", - "content_ltks", - "content_with_weight", - "title_tks", - "docnm_kwd", - "img_id", - "available_int", - "important_kwd", - "position_int", - "page_num_int", - "doc_type_kwd", - "mom_id", - "question_tks", - } - } - outputColumns = convertSelectFields(outputColumns) - - // Determine if text or vector search - hasTextMatch := req.Question != "" - hasVectorMatch := !req.KeywordOnly && len(req.Vector) > 0 - - // Determine score column - scoreColumn := "" - if hasTextMatch { - scoreColumn = "SCORE" - } else if hasVectorMatch { - scoreColumn = "SIMILARITY" - } - - // Add score column if needed - if hasTextMatch || hasVectorMatch { - if hasTextMatch { - outputColumns = append(outputColumns, "score()") - } else if hasVectorMatch { - outputColumns = append(outputColumns, "similarity()") - } - // Add pagerank field - outputColumns = append(outputColumns, PAGERANK_FLD) - } - - // Remove duplicates - outputColumns = convertSelectFields(outputColumns) - - // Build filter string - var filterParts []string - - // For metadata tables, add kb_id filter if provided - if isMetadataTable && len(req.KbIDs) > 0 && req.KbIDs[0] != "" { - kbIDs := req.KbIDs - if len(kbIDs) == 1 { - filterParts = append(filterParts, fmt.Sprintf("kb_id = '%s'", kbIDs[0])) - } else { - kbIDStr := strings.Join(kbIDs, "', '") - filterParts = append(filterParts, fmt.Sprintf("kb_id IN ('%s')", kbIDStr)) - } - } - - // DocIDs filters by doc_id (document ID) to find all chunks belonging to a document - // This is used by ChunkService.List() to list all chunks for a document - if len(req.DocIDs) > 0 { - if len(req.DocIDs) == 1 { - filterParts = append(filterParts, fmt.Sprintf("doc_id = '%s'", req.DocIDs[0])) - } else { - docIDs := strings.Join(req.DocIDs, "', '") - filterParts = append(filterParts, fmt.Sprintf("doc_id IN ('%s')", docIDs)) - } - } - - // Only add available_int filter when there's text/vector match or AvailableInt is explicitly set - // This matches Python's behavior where chunk_list doesn't filter by available_int - if !isMetadataTable && (hasTextMatch || hasVectorMatch || req.AvailableInt != nil) { - if req.AvailableInt != nil { - filterParts = append(filterParts, fmt.Sprintf("available_int=%d", *req.AvailableInt)) - } else { - filterParts = append(filterParts, "available_int=1") - } - } - - filterStr := strings.Join(filterParts, " AND ") - - // Build order_by - var orderBy *OrderByExpr - if req.OrderBy != "" { - orderBy = &OrderByExpr{Fields: []OrderByField{}} - // Parse order_by field and direction - fields := strings.Split(req.OrderBy, ",") - for _, field := range fields { - field = strings.TrimSpace(field) - if strings.HasSuffix(field, " desc") || strings.HasSuffix(field, " DESC") { - fieldName := strings.TrimSuffix(field, " desc") - fieldName = strings.TrimSuffix(fieldName, " DESC") - orderBy.Fields = append(orderBy.Fields, OrderByField{Field: fieldName, Type: SortDesc}) - } else { - orderBy.Fields = append(orderBy.Fields, OrderByField{Field: field, Type: SortAsc}) - } - } - } - - // rank_feature support - var rankFeature map[string]float64 - if req.RankFeature != nil { - rankFeature = req.RankFeature - } - - // Results from all tables - var allResults []map[string]interface{} - totalHits := int64(0) - - // Search across all tables - for _, indexName := range req.IndexNames { - // Determine table names to search - var tableNames []string - if strings.HasPrefix(indexName, "ragflow_doc_meta_") { - tableNames = []string{indexName} - } else { - // For each KB ID, create a table name - kbIDs := req.KbIDs - if len(kbIDs) == 0 { - // If no KB IDs, use the index name directly - kbIDs = []string{""} - } - for _, kbID := range kbIDs { - if kbID == "" { - tableNames = append(tableNames, indexName) - } else { - tableNames = append(tableNames, fmt.Sprintf("%s_%s", indexName, kbID)) - } - } - } - - // Search each table - // 1. First try with min_match=0.3 (30%) - // 2. If no results and has doc_id filter: search without match - // 3. If no results and no doc_id filter: retry with min_match=0.1 (10%) and lower similarity - minMatch := 0.3 - hasDocIDFilter := len(req.DocIDs) > 0 - - for _, tableName := range tableNames { - fmt.Printf("[DEBUG] Searching table: %s\n", tableName) - // Try to get table - _, err := db.GetTable(tableName) - if err != nil { - // Table doesn't exist, skip - continue - } - - // Build query for this table - result, err := e.executeTableSearch(db, tableName, outputColumns, req.Question, req.Vector, filterStr, topK, pageSize, offset, orderBy, rankFeature, req.SimilarityThreshold, minMatch) - if err != nil { - // Skip this table on error - continue - } - - allResults = append(allResults, result.Chunks...) - totalHits += result.Total - } - - // If no results, try fallback strategies - if totalHits == 0 && (hasTextMatch || hasVectorMatch) { - fmt.Printf("[DEBUG] No results, trying fallback strategies\n") - allResults = nil - totalHits = 0 - - if hasDocIDFilter { - // If has doc_id filter, search without match - fmt.Printf("[DEBUG] Retry with no match (has doc_id filter)\n") - for _, tableName := range tableNames { - _, err := db.GetTable(tableName) - if err != nil { - continue - } - // Search without match - pass empty question - result, err := e.executeTableSearch(db, tableName, outputColumns, "", req.Vector, filterStr, topK, pageSize, offset, orderBy, rankFeature, req.SimilarityThreshold, 0.0) - if err != nil { - continue - } - allResults = append(allResults, result.Chunks...) - totalHits += result.Total - } - } else { - // Retry with lower min_match and similarity - fmt.Printf("[DEBUG] Retry with min_match=0.1, similarity=0.17\n") - lowerThreshold := 0.17 - for _, tableName := range tableNames { - _, err := db.GetTable(tableName) - if err != nil { - continue - } - result, err := e.executeTableSearch(db, tableName, outputColumns, req.Question, req.Vector, filterStr, topK, pageSize, offset, orderBy, rankFeature, lowerThreshold, 0.1) - if err != nil { - continue - } - allResults = append(allResults, result.Chunks...) - totalHits += result.Total - } - } - } - } - - if hasTextMatch || hasVectorMatch { - allResults = calculateScores(allResults, scoreColumn, PAGERANK_FLD) - } - - if hasTextMatch || hasVectorMatch { - allResults = sortByScore(allResults, len(allResults)) - } - - // Apply threshold filter to combined results - fmt.Printf("[DEBUG] Threshold check: SimilarityThreshold=%f, hasVectorMatch=%v, hasTextMatch=%v\n", req.SimilarityThreshold, hasVectorMatch, hasTextMatch) - if req.SimilarityThreshold > 0 && hasVectorMatch { - var filteredResults []map[string]interface{} - for _, chunk := range allResults { - score := getScore(chunk) - chunkID := "" - if id, ok := chunk["id"]; ok { - chunkID = fmt.Sprintf("%v", id) - } - fmt.Printf("[DEBUG] Threshold filter: id=%s, score=%f, threshold=%f, pass=%v\n", chunkID, score, req.SimilarityThreshold, score >= req.SimilarityThreshold) - if score >= req.SimilarityThreshold { - filteredResults = append(filteredResults, chunk) - } - } - fmt.Printf("[DEBUG] After threshold filter (combined): %d -> %d chunks\n", len(allResults), len(filteredResults)) - allResults = filteredResults - } - - // Limit to pageSize - if len(allResults) > pageSize { - allResults = allResults[:pageSize] - } - - return &types.SearchResponse{ - Chunks: allResults, - Total: totalHits, - }, nil +// escapeFilterValue escapes single quotes for filter values +func escapeFilterValue(s string) string { + return strings.ReplaceAll(s, "'", "''") } -// calculateScores calculates _score = score_column + pagerank -func calculateScores(chunks []map[string]interface{}, scoreColumn, pagerankField string) []map[string]interface{} { - fmt.Printf("[DEBUG] calculateScores: scoreColumn=%s, pagerankField=%s\n", scoreColumn, pagerankField) +// equivalentConditionToStr converts a condition map to an Infinity filter string +func equivalentConditionToStr(condition map[string]interface{}) string { + if len(condition) == 0 { + return "" + } + + var cond []string + + for k, v := range condition { + if k == "_id" || utility.IsEmpty(v) { + continue + } + + // Handle must_not specially + if k == "must_not" { + if m, ok := v.(map[string]interface{}); ok { + for kk, vv := range m { + if kk == "exists" { + // For must_not exists, use !='' since we don't have table schema + cond = append(cond, fmt.Sprintf("NOT (%v!='')", vv)) + } + } + } + continue + } + + // Handle exists specially (without table schema, use string comparison) + if k == "exists" { + cond = append(cond, fmt.Sprintf("%v!=''", v)) + continue + } + + // Handle keyword fields (using full-text filter) + if fieldKeyword(k) { + // For keyword fields, values are always treated as strings for filter_fulltext + switch val := v.(type) { + case []string: + var inCond []string + for _, item := range val { + inCond = append(inCond, fmt.Sprintf("filter_fulltext('%s', '%s')", + convertMatchingField(k), escapeFilterValue(item))) + } + if len(inCond) > 0 { + cond = append(cond, "("+strings.Join(inCond, " or ")+")") + } + case []interface{}: + var inCond []string + for _, item := range val { + if s, ok := item.(string); ok { + inCond = append(inCond, fmt.Sprintf("filter_fulltext('%s', '%s')", + convertMatchingField(k), escapeFilterValue(s))) + } else { + inCond = append(inCond, fmt.Sprintf("filter_fulltext('%s', '%s')", + convertMatchingField(k), escapeFilterValue(fmt.Sprintf("%v", item)))) + } + } + if len(inCond) > 0 { + cond = append(cond, "("+strings.Join(inCond, " or ")+")") + } + case string: + cond = append(cond, fmt.Sprintf("filter_fulltext('%s', '%s')", + convertMatchingField(k), escapeFilterValue(val))) + default: + cond = append(cond, fmt.Sprintf("filter_fulltext('%s', '%s')", + convertMatchingField(k), escapeFilterValue(fmt.Sprintf("%v", v)))) + } + continue + } + + // Handle list values (mixed types - strings get quotes, numbers don't) + if list, ok := v.([]interface{}); ok && len(list) > 0 { + var strItems, numItems []string + for _, item := range list { + if s, ok := item.(string); ok { + strItems = append(strItems, fmt.Sprintf("'%s'", escapeFilterValue(s))) + } else if n, ok := item.(int); ok { + numItems = append(numItems, strconv.Itoa(n)) + } else if n, ok := item.(int64); ok { + numItems = append(numItems, strconv.FormatInt(n, 10)) + } else if f, ok := item.(float64); ok { + numItems = append(numItems, strconv.FormatFloat(f, 'f', -1, 64)) + } else if s, ok := item.(fmt.Stringer); ok { + strItems = append(strItems, fmt.Sprintf("'%s'", escapeFilterValue(s.String()))) + } else { + strItems = append(strItems, fmt.Sprintf("'%s'", escapeFilterValue(fmt.Sprintf("%v", item)))) + } + } + if len(strItems) > 0 { + if len(strItems) == 1 { + cond = append(cond, fmt.Sprintf("%s=%s", k, strItems[0])) + } else { + cond = append(cond, fmt.Sprintf("%s IN (%s)", k, strings.Join(strItems, ", "))) + } + } + if len(numItems) > 0 { + if len(numItems) == 1 { + cond = append(cond, fmt.Sprintf("%s=%s", k, numItems[0])) + } else { + cond = append(cond, fmt.Sprintf("%s IN (%s)", k, strings.Join(numItems, ", "))) + } + } + continue + } + + if list, ok := v.([]string); ok && len(list) > 0 { + if len(list) == 1 { + cond = append(cond, fmt.Sprintf("%s='%s'", k, escapeFilterValue(list[0]))) + } else { + var items []string + for _, item := range list { + items = append(items, fmt.Sprintf("'%s'", escapeFilterValue(item))) + } + cond = append(cond, fmt.Sprintf("%s IN (%s)", k, strings.Join(items, ", "))) + } + continue + } + + if list, ok := v.([]int); ok && len(list) > 0 { + if len(list) == 1 { + cond = append(cond, fmt.Sprintf("%s=%d", k, list[0])) + } else { + var strs []string + for _, n := range list { + strs = append(strs, strconv.Itoa(n)) + } + cond = append(cond, fmt.Sprintf("%s IN (%s)", k, strings.Join(strs, ", "))) + } + continue + } + + // Handle numeric values (no quotes) + if utility.IsNumericValue(v) { + cond = append(cond, fmt.Sprintf("%s=%v", k, v)) + continue + } + + // Handle string values (with quotes and escaping) + if str, ok := v.(string); ok { + cond = append(cond, fmt.Sprintf("%s='%s'", k, escapeFilterValue(str))) + continue + } + + // Fallback: treat as string + cond = append(cond, fmt.Sprintf("%s='%s'", k, escapeFilterValue(fmt.Sprintf("%v", v)))) + } + + if len(cond) == 0 { + return "" + } + return strings.Join(cond, " AND ") +} + +// calculateScores calculates _score = score_column + pagerank_fld +func calculateScores(chunks []map[string]interface{}, scoreColumn string) []map[string]interface{} { for i := range chunks { score := 0.0 if scoreVal, ok := chunks[i][scoreColumn]; ok { if f, ok := utility.ToFloat64(scoreVal); ok { score += f - fmt.Printf("[DEBUG] chunk[%d]: %s=%f\n", i, scoreColumn, f) } } - if pagerankVal, ok := chunks[i][pagerankField]; ok { - if f, ok := utility.ToFloat64(pagerankVal); ok { + if prVal, ok := chunks[i][common.PAGERANK_FLD]; ok { + if f, ok := utility.ToFloat64(prVal); ok { score += f } } chunks[i]["_score"] = score - fmt.Printf("[DEBUG] chunk[%d]: _score=%f\n", i, score) } return chunks } @@ -668,15 +754,11 @@ func sortByScore(chunks []map[string]interface{}, limit int) []map[string]interf } // Sort by _score descending - for i := 0; i < len(chunks)-1; i++ { - for j := i + 1; j < len(chunks); j++ { - scoreI := getScore(chunks[i]) - scoreJ := getScore(chunks[j]) - if scoreI < scoreJ { - chunks[i], chunks[j] = chunks[j], chunks[i] - } - } - } + sort.Slice(chunks, func(i, j int) bool { + scoreI := getChunkScore(chunks[i]) + scoreJ := getChunkScore(chunks[j]) + return scoreI > scoreJ + }) // Limit if len(chunks) > limit && limit > 0 { @@ -686,270 +768,244 @@ func sortByScore(chunks []map[string]interface{}, limit int) []map[string]interf return chunks } -func getScore(chunk map[string]interface{}) float64 { - // Check _score first - if score, ok := chunk["_score"].(float64); ok { - return score +// getChunkScore extracts the score from a chunk +func getChunkScore(chunk map[string]interface{}) float64 { + if v, ok := chunk["_score"].(float64); ok { + return v } - if score, ok := chunk["_score"].(int); ok { - return float64(score) + if v, ok := chunk["SCORE"].(float64); ok { + return v } - if score, ok := chunk["_score"].(int64); ok { - return float64(score) - } - // Fallback to SCORE (for fusion) or SIMILARITY (for vector-only) - if score, ok := chunk["SCORE"].(float64); ok { - return score - } - if score, ok := chunk["SIMILARITY"].(float64); ok { - return score + if v, ok := chunk["SIMILARITY"].(float64); ok { + return v } return 0.0 } -// executeTableSearch executes search on a single table -func (e *infinityEngine) executeTableSearch(db *infinity.Database, tableName string, outputColumns []string, question string, vector []float64, filterStr string, topK, pageSize, offset int, orderBy *OrderByExpr, rankFeature map[string]float64, similarityThreshold float64, minMatch float64) (*types.SearchResponse, error) { - // Debug logging - fmt.Printf("[DEBUG] executeTableSearch: question=%s, topK=%d, pageSize=%d, similarityThreshold=%f, filterStr=%s\n", question, topK, pageSize, similarityThreshold, filterStr) - - // Get table - table, err := db.GetTable(tableName) - if err != nil { - return nil, err +// GetAggregation aggregates field values from search results. +// +// Example: +// input chunks: +// +// [{"docnm_kwd": "docA"}, {"docnm_kwd": "docA"}, {"docnm_kwd": "docB"}] +// +// GetAggregation(chunks, "docnm_kwd") returns: +// +// [{"key": "docA", "count": 2}, {"key": "docB", "count": 1}] +// +// For tag_kwd field, splits values by "###" separator. +// For other fields, uses comma separation. +func (e *infinityEngine) GetAggregation(chunks []map[string]interface{}, fieldName string) []map[string]interface{} { + if len(chunks) == 0 { + return []map[string]interface{}{} } - // Build query using Table's chainable methods - hasTextMatch := question != "" - hasVectorMatch := len(vector) > 0 - - table = table.Output(outputColumns) - - // Define text fields - textFields := []string{ - "title_tks^10", - "title_sm_tks^5", - "important_kwd^30", - "important_tks^20", - "question_tks^20", - "content_ltks^2", - "content_sm_ltks", + // Check if field exists in first chunk + hasField := false + for _, chunk := range chunks { + if _, ok := chunk[fieldName]; ok { + hasField = true + break + } + } + if !hasField { + return []map[string]interface{}{} } - // Convert field names for Infinity - var convertedFields []string - for _, f := range textFields { - cf := convertMatchingField(f) - convertedFields = append(convertedFields, cf) - } - fields := strings.Join(convertedFields, ",") - - // Format question - formattedQuestion := formatQuestion(question) - - // Compute full filter with filter_fulltext for MatchDense extra_options - var fullFilterWithFulltext string - if filterStr != "" && fields != "" { - fullFilterWithFulltext = fmt.Sprintf("(%s) AND FILTER_FULLTEXT('%s', '%s')", filterStr, fields, formattedQuestion) - } - - // Add text match if question is provided - if hasTextMatch { - extraOptions := map[string]string{ - "topn": fmt.Sprintf("%d", topK), - "minimum_should_match": fmt.Sprintf("%d%%", int(minMatch*100)), + // Count occurrences + tagCounts := make(map[string]int) + for _, chunk := range chunks { + value, ok := chunk[fieldName] + if !ok || value == nil { + continue } - // Add rank_features support - if rankFeature != nil { - var rankFeaturesList []string - for featureName, weight := range rankFeature { - rankFeaturesList = append(rankFeaturesList, fmt.Sprintf("%s^%s^%f", TAG_FLD, featureName, weight)) + // Handle string value + if valueStr, ok := value.(string); ok { + if valueStr == "" { + continue } - if len(rankFeaturesList) > 0 { - extraOptions["rank_features"] = strings.Join(rankFeaturesList, ",") + + var tags []string + // Split by "###" for tag_kwd field + if fieldName == "tag_kwd" && strings.Contains(valueStr, "###") { + for _, tag := range strings.Split(valueStr, "###") { + tag = strings.TrimSpace(tag) + if tag != "" { + tags = append(tags, tag) + } + } + } else { + // Fallback to comma separation + for _, tag := range strings.Split(valueStr, ",") { + tag = strings.TrimSpace(tag) + if tag != "" { + tags = append(tags, tag) + } + } + } + + for _, tag := range tags { + tagCounts[tag]++ + } + continue + } + + // Handle list value + if valueList, ok := value.([]interface{}); ok { + for _, item := range valueList { + if itemStr, ok := item.(string); ok { + tag := strings.TrimSpace(itemStr) + if tag != "" { + tagCounts[tag]++ + } + } } } - - table = table.MatchText(fields, formattedQuestion, topK, extraOptions) - fmt.Printf("[DEBUG] MatchTextExpr: fields=%s, matching_text=%s, topn=%d, extra_options=%v\n", fields, formattedQuestion, topK, extraOptions) } - // Add vector match if provided - if hasVectorMatch { - vectorSize := len(vector) - fieldName := fmt.Sprintf("q_%d_vec", vectorSize) - threshold := similarityThreshold - if threshold <= 0 { - threshold = 0.1 // default - } - extraOptions := map[string]string{ - // Add threshold - "threshold": fmt.Sprintf("%f", threshold), - } - - // Add filter with filter_fulltext, add to MatchDense extra_options - // This is the full filter that includes both available_int=1 AND filter_fulltext - if fullFilterWithFulltext != "" { - extraOptions["filter"] = fullFilterWithFulltext - fmt.Printf("[DEBUG] filterStr=%s, fullFilterWithFulltext=%s\n", filterStr, fullFilterWithFulltext) - } - - fmt.Printf("[DEBUG] MatchDenseExpr: field=%s, topn=%d, extra_options=%v\n", fieldName, topK, extraOptions) - - table = table.MatchDense(fieldName, vector, "float", "cosine", topK, extraOptions) + if len(tagCounts) == 0 { + return []map[string]interface{}{} } - // Add fusion (for text+vector combination) - if hasTextMatch && hasVectorMatch { - fusionParams := map[string]interface{}{ - "normalize": "atan", - "weights": "0.05,0.95", - } - fmt.Printf("[DEBUG] FusionExpr: method=weighted_sum, topn=%d, fusion_params=%v\n", topK, fusionParams) - fmt.Printf("[DEBUG] Before Fusion - table has MatchText=%v, MatchDense=%v\n", hasTextMatch, hasVectorMatch) - table = table.Fusion("weighted_sum", topK, fusionParams) + // Convert to slice and sort by count descending + type tagCountPair struct { + tag string + count int + } + pairs := make([]tagCountPair, 0, len(tagCounts)) + for tag, count := range tagCounts { + pairs = append(pairs, tagCountPair{tag, count}) + } + sort.Slice(pairs, func(i, j int) bool { + return pairs[i].count > pairs[j].count + }) + + // Convert to []map[string]interface{} directly + result := make([]map[string]interface{}, len(pairs)) + for i, p := range pairs { + result[i] = map[string]interface{}{"key": p.tag, "count": p.count} } - // Add order_by if provided - if orderBy != nil && len(orderBy.Fields) > 0 { - var sortFields [][2]interface{} - for _, field := range orderBy.Fields { - sortType := infinity.SortTypeAsc - if field.Type == SortDesc { - sortType = infinity.SortTypeDesc - } - sortFields = append(sortFields, [2]interface{}{field.Field, sortType}) - } - table = table.Sort(sortFields) - } - - // Add filter when there's no text/vector match (like metadata queries) - if !hasTextMatch && !hasVectorMatch && filterStr != "" { - fmt.Printf("[DEBUG] Adding filter for no-match query: %s\n", filterStr) - table = table.Filter(filterStr) - } - - // Set limit and offset - // Use topK to get more results from Infinity, then filter/sort in Go - table = table.Limit(topK) - if offset > 0 { - table = table.Offset(offset) - } - - // Execute query - get the raw query and execute via SDK - result, err := e.executeQuery(table) - if err != nil { - return nil, err - } - - // Debug logging - show returned chunks - scoreColumn := "SIMILARITY" - if hasTextMatch { - scoreColumn = "SCORE" - } - fmt.Printf("[DEBUG] executeTableSearch returned %d chunks\n", len(result.Chunks)) - - result.Chunks = calculateScores(result.Chunks, scoreColumn, PAGERANK_FLD) - - // Debug after calculateScores - for i, chunk := range result.Chunks { - chunkID := "" - if id, ok := chunk["id"]; ok { - chunkID = fmt.Sprintf("%v", id) - } - score := getScore(chunk) - fmt.Printf("[DEBUG] chunk[%d]: id=%s, _score=%f\n", i, chunkID, score) - } - - // Sort by score - result.Chunks = sortByScore(result.Chunks, len(result.Chunks)) - - if len(result.Chunks) > pageSize { - result.Chunks = result.Chunks[:pageSize] - } - result.Total = int64(len(result.Chunks)) - - return result, nil + return result } -// executeQuery executes the query and returns results -func (e *infinityEngine) executeQuery(table *infinity.Table) (*types.SearchResponse, error) { - // Use ToResult() to execute query - result, err := table.ToResult() - if err != nil { - return nil, fmt.Errorf("Infinity query failed: %w", err) +// GetDocIDs extracts document IDs from search results. +// Extracts "id" field from each chunk and returns as a list. +func (e *infinityEngine) GetDocIDs(chunks []map[string]interface{}) []string { + if len(chunks) == 0 { + return nil } - - // Debug: print raw result info - // fmt.Printf("[DEBUG] Infinity raw result: %+v\n", result) - - // Convert result to SearchResponse format - // The SDK returns QueryResult with Data as map[string][]interface{} - qr, ok := result.(*infinity.QueryResult) - if !ok { - return &types.SearchResponse{ - Chunks: []map[string]interface{}{}, - Total: 0, - }, nil - } - - // Convert to chunks format - chunks := make([]map[string]interface{}, 0) - for colName, colData := range qr.Data { - for i, val := range colData { - // Ensure we have a row for this index - for len(chunks) <= i { - chunks = append(chunks, make(map[string]interface{})) - } - chunks[i][colName] = val + ids := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + if id, ok := chunk["id"].(string); ok { + ids = append(ids, id) } } + return ids +} - // Post-process: convert nil/empty values to empty slices for array-like fields - arrayFields := map[string]bool{ - "doc_type_kwd": true, - "important_kwd": true, - "important_tks": true, - "question_tks": true, - "authors_tks": true, - "authors_sm_tks": true, - "title_tks": true, - "title_sm_tks": true, - "content_ltks": true, - "content_sm_ltks": true, +// GetHighlight generates highlighted text snippets for search results. +// Matches keywords in text and wraps them with tags. +func (e *infinityEngine) GetHighlight(chunks []map[string]interface{}, keywords []string, fieldName string) map[string]string { + result := make(map[string]string) + if len(chunks) == 0 || len(keywords) == 0 { + return result } - for i := range chunks { - for colName := range arrayFields { - if val, ok := chunks[i][colName]; !ok || val == nil || val == "" { - chunks[i][colName] = []interface{}{} + + // Check if field exists + hasField := false + for _, chunk := range chunks { + if _, ok := chunk[fieldName]; ok { + hasField = true + break + } + } + if !hasField { + // Try alternative field names + if fieldName == "content_with_weight" { + if _, ok := chunks[0]["content"]; ok { + fieldName = "content" + hasField = true } } - // Convert position_int from hex string to array format - if posVal, ok := chunks[i]["position_int"].(string); ok { - chunks[i]["position_int"] = utility.ConvertHexToPositionIntArray(posVal) + } + if !hasField { + return result + } + + emTag := regexp.MustCompile(`[^<>]+`) + + for _, chunk := range chunks { + id := "" + if idVal, ok := chunk["id"].(string); ok { + id = idVal + } + + txt, ok := chunk[fieldName].(string) + if !ok || txt == "" { + continue + } + + // Check if already highlighted + if emTag.MatchString(txt) { + result[id] = txt + continue + } + + // Replace newlines with spaces + txt = regexp.MustCompile(`[\r\n]`).ReplaceAllString(txt, " ") + + // Split by sentence delimiters + delimiters := regexp.MustCompile(`[.?!;\n]`) + segments := delimiters.Split(txt, -1) + + var highlightedSegments []string + for _, segment := range segments { + // Check if segment is English or contains keywords + englishCount := 0 + totalCount := 0 + for _, r := range segment { + if unicode.IsLetter(r) { + totalCount++ + if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') { + englishCount++ + } + } + } + isEnglish := totalCount > 0 && float64(englishCount)/float64(totalCount) > 0.5 + segmentToCheck := segment + if isEnglish { + // For English: match whole words with boundaries + for _, kw := range keywords { + re := regexp.MustCompile(`(^|[ .?/'\"\(\)!,:;-])` + regexp.QuoteMeta(kw) + `([ .?/'\"\(\)!,:;-]|$)`) + segmentToCheck = re.ReplaceAllString(segmentToCheck, "$1"+kw+"$2") + } + } else { + // For non-English: simple keyword replacement (sorted by length desc for longer matches first) + sortedKeywords := make([]string, len(keywords)) + copy(sortedKeywords, keywords) + sort.Slice(sortedKeywords, func(i, j int) bool { + return len(sortedKeywords[i]) > len(sortedKeywords[j]) + }) + for _, kw := range sortedKeywords { + re := regexp.MustCompile(regexp.QuoteMeta(kw)) + segmentToCheck = re.ReplaceAllString(segmentToCheck, ""+kw+"") + } + } + + // Check if any keywords were highlighted + if emTag.MatchString(segmentToCheck) { + highlightedSegments = append(highlightedSegments, segmentToCheck) + } + } + + if len(highlightedSegments) > 0 { + result[id] = "..." + strings.Join(highlightedSegments, "...") + "..." } else { - chunks[i]["position_int"] = []interface{}{} - } - // Convert page_num_int and top_int from hex string to array - for _, colName := range []string{"page_num_int", "top_int"} { - if val, ok := chunks[i][colName].(string); ok { - chunks[i][colName] = utility.ConvertHexToIntArray(val) - } + result[id] = txt } } - return &types.SearchResponse{ - Chunks: chunks, - Total: int64(len(chunks)), - }, nil -} - -// contains checks if slice contains string -func contains(slice []string, item string) bool { - for _, s := range slice { - if s == item { - return true - } - } - return false + return result } diff --git a/internal/engine/types/types.go b/internal/engine/types/types.go index 5556774121..a7413c018c 100644 --- a/internal/engine/types/types.go +++ b/internal/engine/types/types.go @@ -18,42 +18,87 @@ package types // SearchRequest unified search request for all engines type SearchRequest struct { - // Common fields - IndexNames []string // For ES: index names; For Infinity: treated as table names - Question string // Search query text - Vector []float64 // Embedding vector (optional, for hybrid search) - - // Query analysis results (from QueryBuilder.Question) - MatchText string // Processed match text for ES query_string - Keywords []string // Extracted keywords from question - - // Filters - KbIDs []string // Knowledge base IDs filter - DocIDs []string // Document IDs filter - AvailableInt *int // Available_int filter (1 = available, 0 = unavailable) + // Search target + IndexNames []string // For ES: index names; For Infinity: treated as table name prefixes + KbIDs []string // Knowledge base IDs filter // Pagination - Page int // Page number (1-based) - Size int // Page size - TopK int // Number of candidates for retrieval + Offset int // Offset for pagination (0-based) + Limit int // Limit for pagination - // Search mode - KeywordOnly bool // If true, only do keyword search (no vector search) + // Source fields (for ES: fields to return) + SelectFields []string // List of field names to return - // Scoring parameters - SimilarityThreshold float64 // Minimum similarity score (default: 0.1) - VectorSimilarityWeight float64 // Weight for vector vs keyword (default: 0.3) + // Filtering + Filter map[string]interface{} // Filters for search + + // Match expressions + MatchExprs []interface{} // List of match expressions: [matchText, matchDense, fusionExpr] // Sorting and ranking - OrderBy string // Order by field (e.g., "field1 desc, field2 asc") + OrderBy *OrderByExpr // Order by expression (asc/desc on fields) RankFeature map[string]float64 // Rank features for learning to rank - - // Engine-specific options (optional, for advanced use) - Options map[string]interface{} } -// SearchResponse unified search response for all engines -type SearchResponse struct { +// SearchResult unified search result for all engines +type SearchResult struct { Chunks []map[string]interface{} // Search results Total int64 // Total number of matches } + +type OrderByExpr struct { + Fields []OrderByField +} + +// OrderByField represents a single field ordering. +type OrderByField struct { + Field string + Type OrderByType +} + +// OrderByType represents ascending or descending order. +type OrderByType int + +const ( + // SortAsc represents ascending order. + SortAsc OrderByType = 0 + // SortDesc represents descending order. + SortDesc OrderByType = 1 +) + +// Asc adds an ascending order field. +func (o *OrderByExpr) Asc(field string) *OrderByExpr { + o.Fields = append(o.Fields, OrderByField{Field: field, Type: SortAsc}) + return o +} + +// Desc adds a descending order field. +func (o *OrderByExpr) Desc(field string) *OrderByExpr { + o.Fields = append(o.Fields, OrderByField{Field: field, Type: SortDesc}) + return o +} + +// MatchTextExpr represents a text match expression +type MatchTextExpr struct { + Fields []string // Field names to search (with optional boost, e.g., "title_tks^10") + MatchingText string // Text to match + TopN int // Number of results to return + ExtraOptions map[string]interface{} // Additional options (e.g., minimum_should_match, filter) +} + +// MatchDenseExpr represents a dense vector match expression +type MatchDenseExpr struct { + VectorColumnName string + EmbeddingData []float64 + EmbeddingDataType string + DistanceType string + TopN int + ExtraOptions map[string]interface{} +} + +// FusionExpr represents a fusion expression for hybrid search +type FusionExpr struct { + Method string // Fusion method (e.g., "weighted_sum") + TopN int // TopK for fusion + FusionParams map[string]interface{} // Fusion parameters (e.g., {"weights": "0.05,0.95"}) +} diff --git a/internal/entity/kb.go b/internal/entity/kb.go index 7e4ccb16f9..9424e85812 100644 --- a/internal/entity/kb.go +++ b/internal/entity/kb.go @@ -104,6 +104,7 @@ type Knowledgebase struct { Language *string `gorm:"column:language;size:32;index" json:"language,omitempty"` Description *string `gorm:"column:description;type:longtext" json:"description,omitempty"` EmbdID string `gorm:"column:embd_id;size:128;not null;index" json:"embd_id"` + TenantEmbdID *int64 `gorm:"column:tenant_embd_id;index" json:"tenant_embd_id,omitempty"` Permission string `gorm:"column:permission;size:16;not null;default:me;index" json:"permission"` CreatedBy string `gorm:"column:created_by;size:32;not null;index" json:"created_by"` DocNum int64 `gorm:"column:doc_num;default:0;index" json:"doc_num"` diff --git a/internal/entity/models/deepseek.go b/internal/entity/models/deepseek.go index f215df7b1c..5b7a43d905 100644 --- a/internal/entity/models/deepseek.go +++ b/internal/entity/models/deepseek.go @@ -58,6 +58,11 @@ func (z *DeepSeekModel) Chat(modelName, message *string, apiConfig *APIConfig, c return nil, fmt.Errorf("%s, no such method", z.Name()) } +// ChatWithMessages sends multiple messages with roles and returns response +func (z *DeepSeekModel) ChatWithMessages(modelName string, apiKey *string, messages []Message, chatModelConfig *ChatConfig) (string, error) { + return "", fmt.Errorf("%s, ChatWithMessages not implemented", z.Name()) +} + // ChatStreamlyWithSender sends a message and streams response via sender function (best performance, no channel) func (z *DeepSeekModel) ChatStreamlyWithSender(modelName, message *string, apiConfig *APIConfig, chatModelConfig *ChatConfig, sender func(*string, *string) error) error { return nil diff --git a/internal/entity/models/dummy.go b/internal/entity/models/dummy.go index 4d81c62bdc..e7be91543c 100644 --- a/internal/entity/models/dummy.go +++ b/internal/entity/models/dummy.go @@ -43,6 +43,11 @@ func (z *DummyModel) Chat(modelName, message *string, apiConfig *APIConfig, mode return nil, fmt.Errorf("not implemented") } +// ChatWithMessages sends multiple messages with roles and returns response +func (z *DummyModel) ChatWithMessages(modelName string, apiKey *string, messages []Message, modelConfig *ChatConfig) (string, error) { + return "", fmt.Errorf("not implemented") +} + // ChatStreamlyWithSender sends a message and streams response via sender function (best performance, no channel) func (z *DummyModel) ChatStreamlyWithSender(modelName, message *string, apiConfig *APIConfig, modelConfig *ChatConfig, sender func(*string, *string) error) error { return fmt.Errorf("not implemented") diff --git a/internal/entity/models/minimax.go b/internal/entity/models/minimax.go index f090a2b58b..836e639b02 100644 --- a/internal/entity/models/minimax.go +++ b/internal/entity/models/minimax.go @@ -56,6 +56,11 @@ func (z *MinimaxModel) Chat(modelName, message *string, apiConfig *APIConfig, mo return nil, fmt.Errorf("%s, no such method", z.Name()) } +// ChatWithMessages sends multiple messages with roles and returns response +func (z *MinimaxModel) ChatWithMessages(modelName string, apiKey *string, messages []Message, chatModelConfig *ChatConfig) (string, error) { + return "", fmt.Errorf("%s, ChatWithMessages not implemented", z.Name()) +} + // ChatStreamlyWithSender sends a message and streams response via sender function (best performance, no channel) func (z *MinimaxModel) ChatStreamlyWithSender(modelName, message *string, apiConfig *APIConfig, modelConfig *ChatConfig, sender func(*string, *string) error) error { return fmt.Errorf("%s, no such method", z.Name()) diff --git a/internal/entity/models/moonshot.go b/internal/entity/models/moonshot.go index 7117874e52..ab7ba2aeaf 100644 --- a/internal/entity/models/moonshot.go +++ b/internal/entity/models/moonshot.go @@ -58,6 +58,11 @@ func (z *MoonshotModel) Chat(modelName, message *string, apiConfig *APIConfig, c return nil, fmt.Errorf("not implemented") } +// ChatWithMessages sends multiple messages with roles and returns response +func (z *MoonshotModel) ChatWithMessages(modelName string, apiKey *string, messages []Message, chatModelConfig *ChatConfig) (string, error) { + return "", fmt.Errorf("%s, ChatWithMessages not implemented", z.Name()) +} + // ChatStreamlyWithSender sends a message and streams response via sender function (best performance, no channel) func (z *MoonshotModel) ChatStreamlyWithSender(modelName, message *string, apiConfig *APIConfig, chatModelConfig *ChatConfig, sender func(*string, *string) error) error { return fmt.Errorf("not implemented") diff --git a/internal/entity/models/types.go b/internal/entity/models/types.go index 705dc92595..3a398f01f7 100644 --- a/internal/entity/models/types.go +++ b/internal/entity/models/types.go @@ -1,11 +1,19 @@ package models +// Message represents a chat message with role +type Message struct { + Role string + Content string +} + // EmbeddingModel interface for embedding models type ModelDriver interface { Name() string // Chat sends a message and returns response Chat(modelName, message *string, apiConfig *APIConfig, modelConfig *ChatConfig) (*ChatResponse, error) + // ChatWithMessages sends multiple messages with roles (system, user, etc.) and returns response + ChatWithMessages(modelName string, apiKey *string, messages []Message, modelConfig *ChatConfig) (string, error) // ChatStreamlyWithSender sends a message and streams response via sender function (best performance, no channel) ChatStreamlyWithSender(modelName, message *string, apiConfig *APIConfig, modelConfig *ChatConfig, sender func(*string, *string) error) error // Encode encodes a list of texts into embeddings diff --git a/internal/entity/models/zhipu-ai.go b/internal/entity/models/zhipu-ai.go index e30a4aeac5..ce9eb4c481 100644 --- a/internal/entity/models/zhipu-ai.go +++ b/internal/entity/models/zhipu-ai.go @@ -185,6 +185,106 @@ func (z *ZhipuAIModel) Chat(modelName, message *string, apiConfig *APIConfig, ch return chatResponse, nil } +// ChatWithMessages sends multiple messages with roles and returns response +func (z *ZhipuAIModel) ChatWithMessages(modelName string, apiKey *string, messages []Message, chatModelConfig *ChatConfig) (string, error) { + if apiKey == nil || *apiKey == "" { + return "", fmt.Errorf("api key is nil or empty") + } + + if len(messages) == 0 { + return "", fmt.Errorf("messages is empty") + } + + url := fmt.Sprintf("%s/%s", z.BaseURL["default"], z.URLSuffix.Chat) + + // Convert messages to the format expected by API + apiMessages := make([]map[string]string, len(messages)) + for i, msg := range messages { + apiMessages[i] = map[string]string{ + "role": msg.Role, + "content": msg.Content, + } + } + + // Build request body + reqBody := map[string]interface{}{ + "model": modelName, + "messages": apiMessages, + "stream": false, + "temperature": 1, + } + + if chatModelConfig != nil { + if chatModelConfig.MaxTokens != nil { + reqBody["max_tokens"] = *chatModelConfig.MaxTokens + } + + if chatModelConfig.Temperature != nil { + reqBody["temperature"] = *chatModelConfig.Temperature + } + + if chatModelConfig.TopP != nil { + reqBody["top_p"] = *chatModelConfig.TopP + } + } + + jsonData, err := json.Marshal(reqBody) + if err != nil { + return "", fmt.Errorf("failed to marshal request: %w", err) + } + + req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData)) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiKey)) + + resp, err := z.httpClient.Do(req) + if err != nil { + return "", fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return "", fmt.Errorf("failed to read response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) + } + + // Parse response + var result map[string]interface{} + if err := json.Unmarshal(body, &result); err != nil { + return "", fmt.Errorf("failed to parse response: %w", err) + } + + choices, ok := result["choices"].([]interface{}) + if !ok || len(choices) == 0 { + return "", fmt.Errorf("no choices in response") + } + + firstChoice, ok := choices[0].(map[string]interface{}) + if !ok { + return "", fmt.Errorf("invalid choice format") + } + + messageMap, ok := firstChoice["message"].(map[string]interface{}) + if !ok { + return "", fmt.Errorf("invalid message format") + } + + content, ok := messageMap["content"].(string) + if !ok { + return "", fmt.Errorf("invalid content format") + } + + return content, nil +} + // ChatStreamlyWithSender sends a message and streams response via sender function (best performance, no channel) func (z *ZhipuAIModel) ChatStreamlyWithSender(modelName, message *string, apiConfig *APIConfig, chatModelConfig *ChatConfig, sender func(*string, *string) error) error { var region = "default" diff --git a/internal/entity/types.go b/internal/entity/types.go index 1812a5aa69..b2f2df2958 100644 --- a/internal/entity/types.go +++ b/internal/entity/types.go @@ -69,3 +69,10 @@ type ModelConfig struct { MaxTokens int64 `json:"max_tokens"` IsTools bool `json:"is_tools"` } + +// ModelCredentials holds the credentials for a model +type ModelCredentials struct { + ProviderName string + ModelName string + APIKey string +} diff --git a/internal/logger/logger.go b/internal/logger/logger.go index 65ac2c7f20..f01f613ecf 100644 --- a/internal/logger/logger.go +++ b/internal/logger/logger.go @@ -143,6 +143,11 @@ func Warn(msg string, fields ...zap.Field) { Logger.Warn(msg, fields...) } +// IsDebugEnabled returns true if debug logging is enabled +func IsDebugEnabled() bool { + return atomicLevel.Enabled(zapcore.DebugLevel) +} + // GetLevel returns the current log level func GetLevel() string { levelMu.RLock() diff --git a/internal/service/chunk.go b/internal/service/chunk.go index 4cc98cf6a8..53f8d7db74 100644 --- a/internal/service/chunk.go +++ b/internal/service/chunk.go @@ -21,12 +21,14 @@ import ( "fmt" "ragflow/internal/entity" "ragflow/internal/server" + "strconv" "strings" "go.uber.org/zap" "ragflow/internal/dao" "ragflow/internal/engine" + "ragflow/internal/engine/types" "ragflow/internal/logger" "ragflow/internal/service/nlp" @@ -42,6 +44,7 @@ type ChunkService struct { embeddingCache *utility.EmbeddingLRU kbDAO *dao.KnowledgebaseDAO userTenantDAO *dao.UserTenantDAO + searchService *SearchService } // NewChunkService creates chunk service @@ -54,6 +57,7 @@ func NewChunkService() *ChunkService { embeddingCache: utility.NewEmbeddingLRU(1000), // default capacity kbDAO: dao.NewKnowledgebaseDAO(), userTenantDAO: dao.NewUserTenantDAO(), + searchService: NewSearchService(), } } @@ -68,36 +72,82 @@ type RetrievalTestRequest struct { TopK *int `json:"top_k,omitempty"` CrossLanguages []string `json:"cross_languages,omitempty"` SearchID *string `json:"search_id,omitempty"` - MetaDataFilter map[string]interface{} `json:"meta_data_filter,omitempty"` + Filter map[string]interface{} `json:"meta_data_filter,omitempty"` + TenantRerankID *string `json:"tenant_rerank_id,omitempty"` RerankID *string `json:"rerank_id,omitempty"` Keyword *bool `json:"keyword,omitempty"` SimilarityThreshold *float64 `json:"similarity_threshold,omitempty"` VectorSimilarityWeight *float64 `json:"vector_similarity_weight,omitempty"` - TenantIDs []string `json:"tenant_ids,omitempty"` } // RetrievalTestResponse retrieval test response type RetrievalTestResponse struct { - Chunks []map[string]interface{} `json:"chunks"` - DocAggs []map[string]interface{} `json:"doc_aggs"` - Labels *[]map[string]interface{} `json:"labels"` - Total int64 `json:"total,omitempty"` + Chunks []map[string]interface{} `json:"chunks"` + DocAggs []map[string]interface{} `json:"doc_aggs"` + Labels *map[string]float64 `json:"labels"` + Total int64 `json:"total"` } -// RetrievalTest performs retrieval test +// RetrievalTest performs retrieval test for a given question against specified knowledge bases. +// Corresponds to Python's api/apps/chunk_app.py:retrieval_test() +// +// Flow: +// 1. Validate kbs permissions and embedding model +// 2. Apply metadata filter if specified (auto/semi_auto uses LLM, manual uses provided conditions) +// 3. Apply cross_languages transformation if requested (translate question) +// 4. Apply keyword extraction if requested (append keywords to question) +// 5. Get rank features via LabelQuestion() - tag-based weights or pagerank_fld fallback +// 6. Call RetrievalService.Retrieval() which: +// - Computes query embedding +// - Performs hybrid search (text + vector) with rank features +// - Reranks results +// - Builds doc_aggs by aggregating chunks per document +// 7. knowledge graph retrieval (not implemented) +// 8. Apply retrieval by children to group child chunks under parent chunks func (s *ChunkService) RetrievalTest(req *RetrievalTestRequest, userID string) (*RetrievalTestResponse, error) { - if s.docEngine == nil { - return nil, fmt.Errorf("doc engine not initialized") - } + logger.Info("RetrievalTest started", zap.String("userID", userID), zap.Any("kbID", req.KbID), zap.String("question", req.Question)) + + logger.Debug(fmt.Sprintf("RetrievalTest request:\n"+ + " kbID=%v\n"+ + " question=%s\n"+ + " page=%v, size=%v\n"+ + " docIDs=%v\n"+ + " useKG=%v, topK=%v\n"+ + " crossLanguages=%v\n"+ + " searchID=%v\n"+ + " filter=%v\n"+ + " tenantRerankID=%v\n"+ + " rerankID=%v\n"+ + " keyword=%v\n"+ + " similarityThreshold=%v, vectorSimilarityWeight=%v", + req.KbID, req.Question, + ptrString(req.Page), ptrString(req.Size), req.DocIDs, + ptrString(req.UseKG), ptrString(req.TopK), req.CrossLanguages, ptrString(req.SearchID), + req.Filter, + ptrString(req.TenantRerankID), ptrString(req.RerankID), + ptrString(req.Keyword), + ptrString(req.SimilarityThreshold), ptrString(req.VectorSimilarityWeight))) - // Validate question is required if req.Question == "" { return nil, fmt.Errorf("question is required") } ctx := context.Background() - // Get user's tenants + // Determine kb_id list and check permission for each kb_id + var kbIDs []string + switch v := req.KbID.(type) { + case string: + kbIDs = []string{v} + case []string: + kbIDs = v + default: + return nil, fmt.Errorf("kb_id must be string or array of strings") + } + if len(kbIDs) == 0 { + return nil, fmt.Errorf("kb_id cannot be empty") + } + tenants, err := s.userTenantDAO.GetByUserID(userID) if err != nil { return nil, fmt.Errorf("failed to get user tenants: %w", err) @@ -107,39 +157,14 @@ func (s *ChunkService) RetrievalTest(req *RetrievalTestRequest, userID string) ( } logger.Debug("Retrieved user tenants from database", zap.String("userID", userID), zap.Int("tenantCount", len(tenants))) - // Determine kb_id list - var kbIDs []string - switch v := req.KbID.(type) { - case string: - kbIDs = []string{v} - case []interface{}: - for _, item := range v { - if str, ok := item.(string); ok { - kbIDs = append(kbIDs, str) - } else { - return nil, fmt.Errorf("kb_id array must contain strings") - } - } - case []string: - kbIDs = v - default: - return nil, fmt.Errorf("kb_id must be string or array of strings") - } - - if len(kbIDs) == 0 { - return nil, fmt.Errorf("kb_id cannot be empty") - } - - // Check permission for each kb_id var tenantIDs []string var kbRecords []*entity.Knowledgebase - for _, kbID := range kbIDs { found := false for _, tenant := range tenants { kb, err := s.kbDAO.GetByIDAndTenantID(kbID, tenant.TenantID) if err == nil && kb != nil { - logger.Debug("Found knowledge base record in database", + logger.Debug("Found knowledge base in database", zap.String("kbID", kbID), zap.String("tenantID", tenant.TenantID), zap.String("kbName", kb.Name), @@ -155,7 +180,7 @@ func (s *ChunkService) RetrievalTest(req *RetrievalTestRequest, userID string) ( } } - // Check if all kb records have the same embedding model + // Check if all kbs have the same embedding model if len(kbRecords) > 1 { firstEmbdID := kbRecords[0].EmbdID for i := 1; i < len(kbRecords); i++ { @@ -165,391 +190,268 @@ func (s *ChunkService) RetrievalTest(req *RetrievalTestRequest, userID string) ( } } - // Get user's owner tenants to prioritize - ownerTenants, err := s.userTenantDAO.GetByUserIDAndRole(userID, "owner") - if err != nil { - return nil, fmt.Errorf("failed to get user owner tenants: %w", err) + // Determine meta_data_filter + var chatID string + var creds *entity.ModelCredentials + filter := req.Filter + + if req.SearchID != nil && *req.SearchID != "" { + // If search_id is set, get meta_data_filter and chat_id from search_config + searchDetail, err := s.searchService.GetDetail(*req.SearchID) + if err != nil { + logger.Warn("Failed to get search detail for search_id, proceeding without it", zap.String("searchID", *req.SearchID), zap.Error(err)) + } else if searchConfig, ok := searchDetail["search_config"].(entity.JSONMap); ok && searchConfig != nil { + if searchMetaFilter, ok := searchConfig["meta_data_filter"].(map[string]interface{}); ok { + filter = searchMetaFilter + } + chatID, _ = searchConfig["chat_id"].(string) + } else { + logger.Warn("No search_config found in search detail", zap.String("searchID", *req.SearchID)) + } } - logger.Debug("Retrieved owner tenants from database", - zap.String("userID", userID), - zap.Int("ownerTenantCount", len(ownerTenants))) - req.TenantIDs = tenantIDs - // Choose target tenant: prioritize owner tenant if available in tenantIDs - targetTenantID := tenantIDs[0] + // If meta_data_filter method is auto/semi_auto, get chat model + if filter != nil { + method, _ := filter["method"].(string) + if method == "auto" || method == "semi_auto" { + modelProviderSvc := NewModelProviderService() + if chatID != "" { + // Use chat_id from search_config + creds, err = modelProviderSvc.GetModelByName(chatID, tenantIDs[0]) + if err != nil { + logger.Warn("Failed to get chat model from search_config chat_id, using tenant default", zap.String("chatID", chatID), zap.Error(err)) + } else { + logger.Info("Fetched chat model (from search_config) for metadata filter", + zap.String("chatID", chatID), + zap.String("tenantID", tenantIDs[0]), + zap.String("providerName", creds.ProviderName), + zap.String("modelName", creds.ModelName)) + } + } + // If no chatID from search_config, or creds not found, use tenant default + if creds == nil { + creds, err = modelProviderSvc.GetDefaultModel(entity.ModelTypeChat, tenantIDs[0]) + if err != nil { + logger.Warn("Failed to get tenant default chat model for meta_data_filter", zap.Error(err)) + } else { + logger.Info("Fetched chat model (tenant default) for metadata filter", + zap.String("tenantID", tenantIDs[0]), + zap.String("providerName", creds.ProviderName), + zap.String("modelName", creds.ModelName)) + } + } + } + } - // Get embedding model for the target tenant - embeddingModel, err := s.modelProvider.GetEmbeddingModel(ctx, targetTenantID, kbRecords[0].EmbdID) + // Apply meta_data_filter to get filtered doc_ids (filter by metadata before retrieval) + docIDs := make([]string, len(req.DocIDs)) + copy(docIDs, req.DocIDs) + if filter != nil { + // Get flattened metadata + metadataSvc := NewMetadataService() + flattedMeta, err := metadataSvc.GetFlattedMetaByKBs(kbIDs) + if err != nil { + logger.Warn("Failed to get flatted metadata", zap.Error(err)) + } else { + logger.Info("metadata filter conditions", zap.Any("filter", filter)) + filteredDocIDs, _ := ApplyMetaDataFilter(ctx, filter, flattedMeta, req.Question, creds, req.DocIDs) + docIDs = filteredDocIDs + logger.Info("ApplyMetaDataFilter result", zap.Strings("docIDs", docIDs)) + } + } + + // Apply cross_languages and keyword extraction with tenant default chat model + modifiedQuestion := req.Question + + // Get chat model for cross_languages and keyword_extraction + if len(req.CrossLanguages) > 0 || (req.Keyword != nil && *req.Keyword) { + modelProviderSvc := NewModelProviderService() + creds, err = modelProviderSvc.GetDefaultModel(entity.ModelTypeChat, tenantIDs[0]) + if err != nil { + logger.Warn("Failed to get default chat model for LLM transformations", zap.Error(err)) + } else { + logger.Info("Fetched chat model (tenant default) for cross_languages/keyword_extraction", + zap.String("tenantID", tenantIDs[0]), + zap.String("providerName", creds.ProviderName), + zap.String("modelName", creds.ModelName)) + } + } + + // Apply cross_languages on the question (translate question) + if creds != nil && len(req.CrossLanguages) > 0 { + translated, err := CrossLanguages(ctx, creds, req.Question, req.CrossLanguages) + if err != nil { + logger.Warn("Failed to translate question", zap.Error(err)) + } else { + modifiedQuestion = translated + } + } + + // Apply keyword extraction on the question (append keywords to question) + if creds != nil && req.Keyword != nil && *req.Keyword { + extractedKeywords, err := KeywordExtraction(ctx, creds, modifiedQuestion, 3) + if err != nil { + logger.Warn("Failed to extract keywords from question", zap.Error(err)) + } else if extractedKeywords != "" { + modifiedQuestion = modifiedQuestion + " " + extractedKeywords + } + } + + if modifiedQuestion != req.Question { + logger.Info("Modified question after transformations", + zap.String("originalQuestion", req.Question), + zap.String("modifiedQuestion", modifiedQuestion), + zap.Strings("crossLanguages", req.CrossLanguages), + zap.Bool("keywordExtraction", req.Keyword != nil && *req.Keyword)) + } + + // Get tag-based rank features via LabelQuestion + metadataSvc := NewMetadataService() + labels := metadataSvc.LabelQuestion(modifiedQuestion, kbRecords) + logger.Debug("LabelQuestion result", zap.Any("labels", labels)) + + // Determine embedding model + var embdID string + var tenantLLM *entity.TenantLLM + if kbRecords[0].TenantEmbdID != nil && *kbRecords[0].TenantEmbdID > 0 { + tenantLLM, embdID, err = dao.LookupTenantLLMByID(dao.NewTenantLLMDAO(), *kbRecords[0].TenantEmbdID) + if err != nil { + return nil, fmt.Errorf("failed to get embedding model by tenant_embd_id: %w", err) + } + } else if kbRecords[0].EmbdID != "" { + parts := strings.Split(kbRecords[0].EmbdID, "@") + if len(parts) == 2 && parts[1] != "" { + tenantLLM, embdID, err = dao.LookupTenantLLMByFactory(dao.NewTenantLLMDAO(), tenantIDs[0], parts[1], parts[0], entity.ModelTypeEmbedding) + } else { + tenantLLM, embdID, err = dao.LookupTenantLLMByName(dao.NewTenantLLMDAO(), tenantIDs[0], kbRecords[0].EmbdID, entity.ModelTypeEmbedding) + } + if err != nil { + return nil, fmt.Errorf("failed to get embedding model by embd_id: %w", err) + } + } else { + tenantLLM, err = dao.NewTenantLLMDAO().GetByTenantAndType(tenantIDs[0], entity.ModelTypeEmbedding) + if err != nil { + return nil, fmt.Errorf("failed to get tenant default embedding model: %w", err) + } + if tenantLLM == nil || tenantLLM.LLMName == nil || *tenantLLM.LLMName == "" { + return nil, fmt.Errorf("no default embedding model found for tenant %s", tenantIDs[0]) + } + embdID = fmt.Sprintf("%s@%s", *tenantLLM.LLMName, tenantLLM.LLMFactory) + } + + // Get embedding model for the tenant + var embeddingModel entity.EmbeddingModel + embeddingModel, err = s.modelProvider.GetEmbeddingModel(ctx, tenantIDs[0], embdID) if err != nil { return nil, fmt.Errorf("failed to get embedding model: %w", err) } - logger.Debug("Retrieved embedding model from database", - zap.String("targetTenantID", targetTenantID), - zap.String("embdID", kbRecords[0].EmbdID)) + logger.Info("Fetched embedding model for retrieval", + zap.String("tenantID", tenantIDs[0]), + zap.String("embdID", embdID)) - // Try to get embedding from cache first - embdID := kbRecords[0].EmbdID - var questionVector []float64 - - if s.embeddingCache != nil { - if cachedVector, ok := s.embeddingCache.Get(req.Question, embdID); ok { - logger.Debug("Embedding cache hit", - zap.String("question", req.Question), - zap.String("embdID", embdID), - zap.Int("cacheSize", s.embeddingCache.Len())) - questionVector = cachedVector - } else { - // Cache miss, encode and store - questionVector, err = embeddingModel.EncodeQuery(req.Question) - if err != nil { - return nil, fmt.Errorf("failed to encode query: %w", err) - } - s.embeddingCache.Put(req.Question, embdID, questionVector) - logger.Debug("Embedding cache miss, stored", - zap.String("question", req.Question), - zap.String("embdID", embdID), - zap.Int("vectorDim", len(questionVector)), - zap.Int("cacheSize", s.embeddingCache.Len())) - } - } else { - // No cache, just encode - questionVector, err = embeddingModel.EncodeQuery(req.Question) - if err != nil { - return nil, fmt.Errorf("failed to encode query: %w", err) - } - } - - // Use global QueryBuilder to process question and get matchText and keywords - // Reference: rag/nlp/search.py L115 - queryBuilder := nlp.GetQueryBuilder() - if queryBuilder == nil { - return nil, fmt.Errorf("query builder not initialized") - } - matchTextExpr, keywords := queryBuilder.Question(req.Question, "qa", 0.6) - - //if matchTextExpr == nil { - // return nil, fmt.Errorf("failed to process question") - //} - logger.Debug("QueryBuilder processed question", - zap.String("original", req.Question), - zap.String("matchingText", matchTextExpr.MatchingText), - zap.Strings("keywords", keywords)) - - // Build unified search request - searchReq := &engine.SearchRequest{ - IndexNames: buildIndexNames(tenantIDs), - Question: req.Question, - MatchText: matchTextExpr.MatchingText, - Keywords: keywords, - Vector: questionVector, - KbIDs: kbIDs, - DocIDs: req.DocIDs, - Page: getPageNum(req.Page), - Size: getPageSize(req.Size), - TopK: getTopK(req.TopK), - KeywordOnly: req.Keyword != nil && *req.Keyword, - SimilarityThreshold: getSimilarityThreshold(req.SimilarityThreshold), - VectorSimilarityWeight: getVectorSimilarityWeight(req.VectorSimilarityWeight), - } - - // Execute search through unified engine interface - result, err := s.docEngine.Search(ctx, searchReq) - if err != nil { - return nil, fmt.Errorf("search failed: %w", err) - } - - // Convert result to unified response - searchResp, ok := result.(*engine.SearchResponse) - if !ok { - return nil, fmt.Errorf("invalid search response type") - } - - //return &RetrievalTestResponse{ - // Chunks: searchResp.Chunks, - // Labels: []map[string]interface{}{}, // Empty labels for now - // Total: searchResp.Total, - //}, nil - - //// Build SearchResult for reranker - //sres := buildSearchResult(searchResp, questionVector) - // - // Get rerank model if RerankID is specified (can be nil) + // Get rerank model if RerankID is specified var rerankModel nlp.RerankModel - if req.RerankID != nil && *req.RerankID != "" { - rerankModel, err = s.modelProvider.GetRerankModel(ctx, targetTenantID, *req.RerankID) + var rerankCompositeName string + if req.TenantRerankID != nil && *req.TenantRerankID != "" { + tenantRerankIDInt, parseErr := strconv.ParseInt(*req.TenantRerankID, 10, 64) + if parseErr != nil { + return nil, fmt.Errorf("invalid tenant_rerank_id: %w", parseErr) + } + _, rerankCompositeName, err = dao.LookupTenantLLMByID(dao.NewTenantLLMDAO(), tenantRerankIDInt) if err != nil { - logger.Warn("Failed to get rerank model, falling back to standard reranking", zap.Error(err)) - rerankModel = nil + return nil, fmt.Errorf("failed to get rerank model by tenant_rerank_id: %w", err) + } + rerankModel, err = s.modelProvider.GetRerankModel(ctx, tenantIDs[0], rerankCompositeName) + if err != nil { + return nil, fmt.Errorf("failed to get rerank model by tenant_rerank_id: %w", err) + } + } else if req.RerankID != nil && *req.RerankID != "" { + var err error + _, rerankCompositeName, err = dao.LookupTenantLLMByName(dao.NewTenantLLMDAO(), tenantIDs[0], *req.RerankID, entity.ModelTypeRerank) + if err != nil { + return nil, fmt.Errorf("failed to get rerank model by rerank_id: %w", err) + } + rerankModel, err = s.modelProvider.GetRerankModel(ctx, tenantIDs[0], rerankCompositeName) + if err != nil { + return nil, fmt.Errorf("failed to get rerank model by rerank_id: %w", err) } } - // Perform reranking - // Reference: rag/nlp/search.py L404-L429 - vtWeight := getVectorSimilarityWeight(req.VectorSimilarityWeight) - tkWeight := 1.0 - vtWeight - useInfinity := s.engineType == server.EngineInfinity - - sim, term_similarity, vector_similarity := nlp.Rerank( - rerankModel, - searchResp, - keywords, - questionVector, - nil, - req.Question, - tkWeight, - vtWeight, - useInfinity, - "content_ltks", - queryBuilder, - ) - // - // Apply similarity threshold and sort chunks - similarityThreshold := getSimilarityThreshold(req.SimilarityThreshold) - filteredChunks := applyRerankResults(searchResp.Chunks, sim, similarityThreshold) - for idx, _ := range filteredChunks { - filteredChunks[idx]["similarity"] = sim[idx] - filteredChunks[idx]["term_similarity"] = term_similarity[idx] - filteredChunks[idx]["vector_similarity"] = vector_similarity[idx] + if rerankModel != nil { + logger.Info("Fetched rerank model", + zap.String("tenantID", tenantIDs[0]), + zap.String("rerankCompositeName", rerankCompositeName)) } - convertedChunks := buildRetrievalTestResults(filteredChunks) - - // Build doc_aggs by aggregating chunks by docnm - docAggsMap := make(map[string]struct { - docID string - count int - }) - docNameOrder := []string{} // Track insertion order of doc names - for _, chunk := range filteredChunks { - docName := "" - docID := "" - if v, ok := chunk["docnm"].(string); ok { - docName = v - } - if v, ok := chunk["doc_id"].(string); ok { - docID = v - } - if docName == "" { - continue - } - if entry, exists := docAggsMap[docName]; exists { - entry.count++ - docAggsMap[docName] = entry - } else { - docAggsMap[docName] = struct { - docID string - count int - }{docID: docID, count: 1} - docNameOrder = append(docNameOrder, docName) - } + retrievalReq := &nlp.RetrievalRequest{ + TenantIDs: tenantIDs, + Question: modifiedQuestion, + KbIDs: kbIDs, + DocIDs: docIDs, + Page: getPageNum(req.Page, 1), + PageSize: getPageSize(req.Size, 30), + Top: req.TopK, + SimilarityThreshold: req.SimilarityThreshold, + VectorSimilarityWeight: req.VectorSimilarityWeight, + RerankModel: rerankModel, + RankFeature: &labels, + EmbeddingModel: embeddingModel, } - // Convert to list maintaining insertion order - type docAggEntry struct { - docName string - docID string - count int - order int + // Call RetrievalService to perform retrieval + retrievalResult, err := nlp.NewRetrievalService(s.docEngine).Retrieval(ctx, retrievalReq) + if err != nil { + return nil, fmt.Errorf("retrieval search failed: %w", err) } - docAggsList := make([]docAggEntry, 0, len(docAggsMap)) - for order, docName := range docNameOrder { - entry := docAggsMap[docName] - docAggsList = append(docAggsList, docAggEntry{docName: docName, docID: entry.docID, count: entry.count, order: order}) + + filteredChunks := retrievalResult.Chunks + + // Handle knowledge graph retrieval + // TODO: KG retrieval requires GraphRAG infrastructure which is not yet implemented in Go + if req.UseKG != nil && *req.UseKG { + logger.Warn("use_kg is not yet implemented in Go - skipping KG retrieval") } - // Sort by count descending, then by order ascending (for tie-breaking) - for i := 0; i < len(docAggsList)-1; i++ { - for j := i + 1; j < len(docAggsList); j++ { - if docAggsList[j].count > docAggsList[i].count || - (docAggsList[j].count == docAggsList[i].count && docAggsList[j].order < docAggsList[i].order) { - docAggsList[i], docAggsList[j] = docAggsList[j], docAggsList[i] - } - } - } - docAggs := make([]map[string]interface{}, 0, len(docAggsList)) - for _, entry := range docAggsList { - docAggs = append(docAggs, map[string]interface{}{ - "doc_name": entry.docName, - "doc_id": entry.docID, - "count": entry.count, - }) + + // Apply retrieval_by_children - aggregate child chunks into parent chunks + filteredChunks = nlp.RetrievalByChildren(filteredChunks, tenantIDs, s.docEngine, ctx) + + // Remove vector field from each chunk + for i := range filteredChunks { + delete(filteredChunks[i], "vector") } + logger.Info("RetrievalTest completed", zap.String("userID", userID), zap.Any("kbID", req.KbID), zap.String("question", req.Question), zap.Int64("chunkCount", int64(len(filteredChunks)))) + return &RetrievalTestResponse{ - Chunks: convertedChunks, - DocAggs: docAggs, - Labels: nil, - Total: int64(len(convertedChunks)), + Chunks: filteredChunks, + DocAggs: retrievalResult.DocAggs, + Labels: &labels, + Total: int64(len(filteredChunks)), }, nil } // Helper functions -func getPageNum(page *int) int { +// ptrString converts a pointer to a formatted string +func ptrString[T any](p *T) string { + if p == nil { + return "" + } + return fmt.Sprintf("%v", *p) +} + +func getPageNum(page *int, defaultVal int) int { if page != nil && *page > 0 { return *page } - return 1 + return defaultVal } -func getPageSize(size *int) int { +func getPageSize(size *int, defaultVal int) int { if size != nil && *size > 0 { return *size } - return 30 -} - -func getTopK(topk *int) int { - if topk != nil && *topk > 0 { - return *topk - } - return 1024 -} - -func getSimilarityThreshold(threshold *float64) float64 { - if threshold != nil && *threshold >= 0 { - return *threshold - } - return 0.1 -} - -func getVectorSimilarityWeight(weight *float64) float64 { - if weight != nil && *weight >= 0 && *weight <= 1 { - return *weight - } - return 0.3 -} - -func buildIndexNames(tenantIDs []string) []string { - indexNames := make([]string, len(tenantIDs)) - for i, tenantID := range tenantIDs { - indexNames[i] = fmt.Sprintf("ragflow_%s", tenantID) - } - return indexNames -} - -// buildSearchResult converts engine.SearchResponse to nlp.SearchResult for reranking -func buildSearchResult(resp *engine.SearchResponse, queryVector []float64) *nlp.SearchResult { - field := make(map[string]map[string]interface{}) - ids := make([]string, 0, len(resp.Chunks)) - - for i, chunk := range resp.Chunks { - // Extract ID from chunk - id := "" - if idVal, ok := chunk["_id"].(string); ok { - id = idVal - } else { - id = fmt.Sprintf("chunk_%d", i) - } - ids = append(ids, id) - - // Store fields by id - field[id] = chunk - } - - return &nlp.SearchResult{ - Total: len(resp.Chunks), - IDs: ids, - QueryVector: queryVector, - Field: field, - } -} - -// applyRerankResults sorts and filters chunks based on reranking results -// Reference: rag/nlp/search.py L430-L439 -func applyRerankResults(chunks []map[string]interface{}, sim []float64, threshold float64) []map[string]interface{} { - if len(chunks) == 0 || len(sim) == 0 { - return chunks - } - - // Get sorted indices (descending by similarity) - sortedIndices := nlp.ArgsortDescending(sim) - - // Sort and filter chunks based on reranking results - var filteredChunks []map[string]interface{} - for _, idx := range sortedIndices { - if idx < 0 || idx >= len(chunks) { - continue - } - if sim[idx] >= threshold { - chunk := chunks[idx] - // Add similarity score to chunk - chunk["_score"] = sim[idx] - filteredChunks = append(filteredChunks, chunk) - } - } - - return filteredChunks -} - -// buildRetrievalTestResults converts filtered chunks to retrieval test results with renamed keys -func buildRetrievalTestResults(filteredChunks []map[string]interface{}) []map[string]interface{} { - results := make([]map[string]interface{}, 0, len(filteredChunks)) - - for _, chunk := range filteredChunks { - result := make(map[string]interface{}) - - // Key mappings - if v, ok := chunk["id"]; ok { - result["chunk_id"] = v - } else if v, ok := chunk["_id"]; ok { - result["chunk_id"] = v - } - if v, ok := chunk["content"]; ok { - result["content_ltks"] = v - result["content_with_weight"] = v - } else { - if v, ok := chunk["content_ltks"]; ok { - result["content_ltks"] = v - } - if v, ok := chunk["content_with_weight"]; ok { - result["content_with_weight"] = v - } - } - if v, ok := chunk["doc_id"]; ok { - result["doc_id"] = v - } - if v, ok := chunk["docnm"]; ok { - result["docnm_kwd"] = v - } else if v, ok := chunk["docnm_kwd"]; ok { - result["docnm_kwd"] = v - } - if v, ok := chunk["img_id"]; ok { - result["image_id"] = v - } - if v, ok := chunk["kb_id"]; ok { - result["kb_id"] = v - } - if v, ok := chunk["position_int"]; ok { - result["positions"] = v - } - if v, ok := chunk["doc_type_kwd"]; ok { - result["doc_type_kwd"] = v - } - if v, ok := chunk["mom_id"]; ok { - result["mom_id"] = v - } - if v, ok := chunk["important_kwd"]; ok { - result["important_kwd"] = v - } else if v, ok := chunk["important_keywords"]; ok { - result["important_kwd"] = v - } - if v, ok := chunk["tag_kwd"]; ok { - result["tag_kwd"] = v - } - if v, ok := chunk["similarity"]; ok { - result["similarity"] = v - } - if v, ok := chunk["term_similarity"]; ok { - result["term_similarity"] = v - } - if v, ok := chunk["vector_similarity"]; ok { - result["vector_similarity"] = v - } - - results = append(results, result) - } - - return results + return defaultVal } // GetChunkRequest request for getting a chunk by ID @@ -602,7 +504,6 @@ func (s *ChunkService) Get(req *GetChunkRequest, userID string) (*GetChunkRespon if doc != nil { chunk, ok := doc.(map[string]interface{}) if ok { - // Format to match Python output result := make(map[string]interface{}) skipFields := map[string]bool{ "id": true, "authors": true, "_score": true, "SCORE": true, @@ -724,39 +625,33 @@ func (s *ChunkService) List(req *ListChunksRequest, userID string) (*ListChunksR indexName := fmt.Sprintf("ragflow_%s", targetTenantID) - page := getPageNum(req.Page) - size := getPageSize(req.Size) + page := getPageNum(req.Page, 1) + size := getPageSize(req.Size, 30) keywords := req.Keywords // Build search request - same as retrieval test but filtered by doc_id - searchReq := &engine.SearchRequest{ + searchReq := &types.SearchRequest{ IndexNames: []string{indexName}, - Question: keywords, + MatchExprs: []interface{}{keywords}, KbIDs: kbIDs, - DocIDs: []string{req.DocID}, - Page: page, - Size: size, - TopK: size, + Offset: (page - 1) * size, + Limit: size, + Filter: map[string]interface{}{ + "doc_id": req.DocID, + }, } // Add available_int filter if specified if req.AvailableInt != nil { - searchReq.AvailableInt = req.AvailableInt + searchReq.Filter["available_int"] = *req.AvailableInt } // Execute search through unified engine interface - result, err := s.docEngine.Search(ctx, searchReq) + searchResp, err := s.docEngine.Search(ctx, searchReq) if err != nil { return nil, fmt.Errorf("search failed: %w", err) } - // Convert result to unified response - searchResp, ok := result.(*engine.SearchResponse) - if !ok { - return nil, fmt.Errorf("invalid search response type") - } - - // Format output to match Python chunks := make([]map[string]interface{}, 0, len(searchResp.Chunks)) for _, chunk := range searchResp.Chunks { // Inline formatChunkForList @@ -819,7 +714,7 @@ func (s *ChunkService) List(req *ListChunksRequest, userID string) (*ListChunksR chunks = append(chunks, result) } - // Build document info (matching Python doc.to_dict()) + // Build document info timeFormat := "2006-01-02T15:04:05" docInfo := map[string]interface{}{ "id": doc.ID, @@ -859,16 +754,16 @@ func (s *ChunkService) List(req *ListChunksRequest, userID string) (*ListChunksR // UpdateChunkRequest request for updating a chunk type UpdateChunkRequest struct { - DatasetID string `json:"dataset_id"` - DocumentID string `json:"document_id"` - ChunkID string `json:"chunk_id"` - Content *string `json:"content,omitempty"` - ImportantKwd []string `json:"important_keywords,omitempty"` - Questions []string `json:"questions,omitempty"` - Available *bool `json:"available,omitempty"` - Positions []interface{} `json:"positions,omitempty"` - TagKwd []string `json:"tag_kwd,omitempty"` - TagFeas interface{} `json:"tag_feas,omitempty"` + DatasetID string `json:"dataset_id"` + DocumentID string `json:"document_id"` + ChunkID string `json:"chunk_id"` + Content *string `json:"content,omitempty"` + ImportantKwd []string `json:"important_keywords,omitempty"` + Questions []string `json:"questions,omitempty"` + Available *bool `json:"available,omitempty"` + Positions []interface{} `json:"positions,omitempty"` + TagKwd []string `json:"tag_kwd,omitempty"` + TagFeas interface{} `json:"tag_feas,omitempty"` } // UpdateChunk updates a chunk fields @@ -915,7 +810,7 @@ func (s *ChunkService) UpdateChunk(req *UpdateChunkRequest, userID string) error return fmt.Errorf("document does not belong to this dataset") } - // Fetch existing chunk first (like Python does) + // Fetch existing chunk first indexName := fmt.Sprintf("ragflow_%s", targetTenantID) existingChunk, err := s.docEngine.GetChunk(ctx, indexName, req.ChunkID, []string{req.DatasetID}) if err != nil { @@ -927,7 +822,7 @@ func (s *ChunkService) UpdateChunk(req *UpdateChunkRequest, userID string) error return fmt.Errorf("invalid chunk format") } - // Build update dict like Python does (doc.py:1476-1523) + // Build update dict d := make(map[string]interface{}) // Content - use new value or existing @@ -1012,9 +907,9 @@ func (s *ChunkService) UpdateChunk(req *UpdateChunkRequest, userID string) error // RemoveChunksRequest request for removing chunks type RemoveChunksRequest struct { - DocID string `json:"doc_id"` - ChunkIDs []string `json:"chunk_ids,omitempty"` - DeleteAll bool `json:"delete_all,omitempty"` + DocID string `json:"doc_id"` + ChunkIDs []string `json:"chunk_ids,omitempty"` + DeleteAll bool `json:"delete_all,omitempty"` } // RemoveChunks removes chunks from the dataset table. diff --git a/internal/service/generator.go b/internal/service/generator.go new file mode 100644 index 0000000000..901a486790 --- /dev/null +++ b/internal/service/generator.go @@ -0,0 +1,167 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package service + +import ( + "context" + "fmt" + "regexp" + "strings" + + "go.uber.org/zap" + + "ragflow/internal/entity" + modelModule "ragflow/internal/entity/models" + "ragflow/internal/logger" +) + +// KeywordExtraction extracts keywords from content using LLM. +// Corresponds to rag/prompts/generator.py:keyword_extraction(). +// +// Uses ChatToModelByApiKey via ModelCredentials to call the LLM with a keyword extraction prompt. +// Returns comma-separated top N important keywords/phrases from the content. +func KeywordExtraction(ctx context.Context, creds *entity.ModelCredentials, content string, topN int) (string, error) { + if creds == nil { + return "", fmt.Errorf("model credentials is nil") + } + + if content == "" { + return "", nil + } + + if topN <= 0 { + topN = 3 + } + + // Load keyword prompt template from file + keywordPromptTemplate, err := LoadPrompt("keyword_prompt") + if err != nil { + return "", fmt.Errorf("failed to load keyword prompt: %w", err) + } + + // Render template with content and topn + renderedPrompt := RenderPrompt(keywordPromptTemplate, map[string]interface{}{ + "content": content, + "topn": topN, + }) + + // Build messages: system prompt + user "Output:" + messages := []modelModule.Message{ + {Role: "system", Content: renderedPrompt}, + {Role: "user", Content: "Output: "}, + } + + // Call LLM using ChatWithMessagesToModelByApiKey + modelProviderSvc := NewModelProviderService() + responsePtr, code, err := modelProviderSvc.ChatWithMessagesToModelByApiKey(creds.ProviderName, creds.ModelName, creds.APIKey, messages) + if err != nil { + return "", fmt.Errorf("failed to extract keywords: code=%d, err=%w", int(code), err) + } + + response := *responsePtr + logger.Info("KeywordExtraction result", zap.String("response", response)) + + // Clean up response - remove thinking tags if present + response = strings.TrimSpace(response) + response = thinkBlockRE.ReplaceAllString(response, "") + response = strings.TrimSpace(response) + + if strings.Contains(response, "**ERROR**") { + return "", fmt.Errorf("error in keyword extraction response") + } + + return response, nil +} + +// CrossLanguages translates a question into multiple languages using LLM. +func CrossLanguages(ctx context.Context, creds *entity.ModelCredentials, query string, languages []string) (string, error) { + if creds == nil { + return "", fmt.Errorf("model credentials is nil") + } + + if query == "" { + return query, nil + } + + if len(languages) == 0 { + return query, nil + } + + // Load system prompt from embedded file + systemPrompt, err := LoadPrompt("cross_languages_sys_prompt") + if err != nil { + return query, fmt.Errorf("failed to load system prompt: %w", err) + } + + // Load user prompt template from file + userPromptTemplate, err := LoadPrompt("cross_languages_user_prompt") + if err != nil { + return query, fmt.Errorf("failed to load user prompt: %w", err) + } + + // Render user prompt with query and languages + userPrompt := RenderPrompt(userPromptTemplate, map[string]interface{}{ + "query": query, + "languages": languages, + }) + + // Build messages: system prompt + user prompt + messages := []modelModule.Message{ + {Role: "system", Content: systemPrompt}, + {Role: "user", Content: userPrompt}, + } + + // Call LLM using ChatWithMessagesToModelByApiKey + modelProviderSvc := NewModelProviderService() + responsePtr, code, err := modelProviderSvc.ChatWithMessagesToModelByApiKey(creds.ProviderName, creds.ModelName, creds.APIKey, messages) + if err != nil { + return query, fmt.Errorf("failed to translate question: code=%d, err=%w", int(code), err) + } + + response := *responsePtr + + // Clean up response - remove think tags and trim + response = strings.TrimSpace(response) + response = thinkBlockRE.ReplaceAllString(response, "") + response = strings.TrimSpace(response) + + if strings.Contains(response, "**ERROR**") { + return query, nil + } + + // Parse response + response = strings.TrimPrefix(response, "Output:") + response = strings.TrimPrefix(response, "output:") + response = regexp.MustCompile(`(?i)^output:\s*`).ReplaceAllString(response, "") + response = regexp.MustCompile(`\n+`).ReplaceAllString(response, "") + response = strings.TrimSpace(response) + + parts := strings.Split(response, "===") + var translations []string + for _, part := range parts { + trimmed := strings.TrimSpace(part) + if trimmed != "" { + translations = append(translations, trimmed) + } + } + + if len(translations) > 0 { + return strings.Join(translations, "\n"), nil + } + + return query, nil +} diff --git a/internal/service/load_prompt.go b/internal/service/load_prompt.go new file mode 100644 index 0000000000..138a88822e --- /dev/null +++ b/internal/service/load_prompt.go @@ -0,0 +1,160 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package service + +import ( + "fmt" + "os" + "path/filepath" + "regexp" + "strings" + "sync" +) + +var ( + promptCache = make(map[string]string) + promptMu sync.RWMutex + promptsBaseDir string +) + +// thinkBlockRE is used to strip think blocks from LLM responses +var thinkBlockRE = regexp.MustCompile(`[\s\S]*?`) + +func init() { + // Strategy 1: Check working directory first (most reliable during development/tests) + cwd, err := os.Getwd() + if err == nil { + // Check if CWD has rag/prompts directly + if _, err := os.Stat(filepath.Join(cwd, "rag", "prompts")); err == nil { + promptsBaseDir = cwd + return + } + // Walk up from CWD looking for rag/prompts + dir := cwd + for dir != "/" && dir != "" { + if _, err := os.Stat(filepath.Join(dir, "rag", "prompts")); err == nil { + promptsBaseDir = dir + return + } + dir = filepath.Dir(dir) + } + } + + // Strategy 2: Walk up from executable (for production Docker where binary is in /ragflow/bin/) + exe, err := os.Executable() + if err == nil { + dir := filepath.Dir(exe) + for dir != "/" && dir != "" { + if _, err := os.Stat(filepath.Join(dir, "rag", "prompts")); err == nil { + promptsBaseDir = dir + return + } + dir = filepath.Dir(dir) + } + } + + // Final fallback + promptsBaseDir = "/ragflow" +} + +// LoadPrompt loads a prompt by name from the rag/prompts/ directory. +// It caches loaded prompts for subsequent calls. +// Corresponds to rag/prompts/template.py:load_prompt() +func LoadPrompt(name string) (string, error) { + promptMu.RLock() + if cached, ok := promptCache[name]; ok { + promptMu.RUnlock() + return cached, nil + } + promptMu.RUnlock() + + promptPath := filepath.Join(promptsBaseDir, "rag", "prompts", fmt.Sprintf("%s.md", name)) + content, err := os.ReadFile(promptPath) + if err != nil { + return "", fmt.Errorf("prompt file '%s.md' not found in rag/prompts/: %w", name, err) + } + + cached := strings.TrimSpace(string(content)) + promptMu.Lock() + promptCache[name] = cached + promptMu.Unlock() + + return cached, nil +} + +// RenderPrompt renders a prompt template with the given variables. +// Supports {{ variable }} and {{ variable | filter(args) }} syntax. +// Corresponds to rag/prompts/generator.py template rendering (Jinja2). +func RenderPrompt(template string, data map[string]interface{}) string { + // Handle {{ variable | filter(args) }} syntax - capture filter arguments too + filterPattern := regexp.MustCompile(`\{\{\s*(\w+)\s*\|\s*(\w+)\s*\(\s*([^)]*)\s*\)\s*\}\}`) + result := filterPattern.ReplaceAllStringFunc(template, func(match string) string { + matches := filterPattern.FindStringSubmatch(match) + if len(matches) < 4 { + return match + } + key := matches[1] + filter := matches[2] + args := matches[3] + value := data[key] + return applyFilter(value, filter, args) + }) + + // Handle simple {{ variable }} syntax + varPattern := regexp.MustCompile(`\{\{\s*(\w+)\s*\}\}`) + result = varPattern.ReplaceAllStringFunc(result, func(match string) string { + matches := varPattern.FindStringSubmatch(match) + if len(matches) < 2 { + return match + } + key := matches[1] + if value, ok := data[key]; ok { + return fmt.Sprintf("%v", value) + } + return match + }) + + return result +} + +// applyFilter applies a filter to a value with optional arguments. +func applyFilter(value interface{}, filter string, args string) string { + switch filter { + case "join": + // {{ variable | join(', ') }} - expects value to be a slice, args is the separator + if slice, ok := value.([]string); ok { + sep := stripQuotes(strings.TrimSpace(args)) + if sep == "" { + sep = ", " + } + return strings.Join(slice, sep) + } + return fmt.Sprintf("%v", value) + default: + return fmt.Sprintf("%v", value) + } +} + +// stripQuotes removes matching surrounding single or double quotes. +func stripQuotes(s string) string { + if len(s) >= 2 { + if (s[0] == '\'' && s[len(s)-1] == '\'') || (s[0] == '"' && s[len(s)-1] == '"') { + return s[1 : len(s)-1] + } + } + return s +} diff --git a/internal/service/metadata.go b/internal/service/metadata.go index 7f21775a13..a4be1412e3 100644 --- a/internal/service/metadata.go +++ b/internal/service/metadata.go @@ -20,6 +20,7 @@ import ( "context" "encoding/json" "fmt" + "strconv" "ragflow/internal/dao" "ragflow/internal/engine" @@ -77,27 +78,23 @@ func (s *MetadataService) SearchMetadata(kbID, tenantID string, docIDs []string, indexName := BuildMetadataIndexName(tenantID) searchReq := &types.SearchRequest{ - IndexNames: []string{indexName}, - KbIDs: []string{kbID}, - DocIDs: docIDs, - Page: 1, - Size: size, - KeywordOnly: true, + IndexNames: []string{indexName}, + KbIDs: []string{kbID}, + Offset: 0, + Limit: size, + Filter: map[string]interface{}{ + "doc_id": docIDs, + }, } - result, err := s.docEngine.Search(context.Background(), searchReq) + searchResult, err := s.docEngine.Search(context.Background(), searchReq) if err != nil { return nil, fmt.Errorf("search failed: %w", err) } - searchResp, ok := result.(*types.SearchResponse) - if !ok { - return nil, fmt.Errorf("invalid search response type") - } - return &SearchMetadataResult{ IndexName: indexName, - Chunks: searchResp.Chunks, + Chunks: searchResult.Chunks, }, nil } @@ -115,29 +112,135 @@ func (s *MetadataService) SearchMetadataByKBs(kbIDs []string, size int) (*Search indexName := BuildMetadataIndexName(tenantID) searchReq := &types.SearchRequest{ - IndexNames: []string{indexName}, - KbIDs: kbIDs, - Page: 1, - Size: size, - KeywordOnly: true, + IndexNames: []string{indexName}, + KbIDs: kbIDs, + Offset: 0, + Limit: size, } - result, err := s.docEngine.Search(context.Background(), searchReq) + searchResult, err := s.docEngine.Search(context.Background(), searchReq) if err != nil { return nil, fmt.Errorf("search failed: %w", err) } - searchResp, ok := result.(*types.SearchResponse) - if !ok { - return nil, fmt.Errorf("invalid search response type") - } - return &SearchMetadataResult{ IndexName: indexName, - Chunks: searchResp.Chunks, + Chunks: searchResult.Chunks, }, nil } +// GetFlattedMetaByKBs returns flattened metadata in the format: +// {field_name: {value: [doc_ids]}} +func (s *MetadataService) GetFlattedMetaByKBs(kbIDs []string) (map[string]interface{}, error) { + if len(kbIDs) == 0 { + return make(map[string]interface{}), nil + } + + // Get metadata for all docs in KBs (use large limit like Python's 10000) + result, err := s.SearchMetadataByKBs(kbIDs, 10000) + if err != nil { + return nil, err + } + + flattedMeta := make(map[string]interface{}) + + for _, chunk := range result.Chunks { + // Extract doc_id from chunk + docID := "" + if id, ok := chunk["id"].(string); ok { + docID = id + } else if id, ok := chunk["doc_id"].(string); ok { + docID = id + } + + if docID == "" { + continue + } + + // Extract metadata fields + metaFields, err := ExtractMetaFields(chunk) + if err != nil || len(metaFields) == 0 { + continue + } + + // Flatten each field + for fieldName, fieldValue := range metaFields { + if fieldValue == nil { + continue + } + + // Initialize field map if not exists + if _, exists := flattedMeta[fieldName]; !exists { + flattedMeta[fieldName] = make(map[string]interface{}) + } + + valueMap, ok := flattedMeta[fieldName].(map[string]interface{}) + if !ok { + continue + } + + // Handle string, number (float64/int), and list of string/number + switch v := fieldValue.(type) { + case string: + // Single string value (including time strings) + if v != "" { + if _, exists := valueMap[v]; !exists { + valueMap[v] = []string{docID} + } else { + valueMap[v] = appendDocID(valueMap[v], docID) + } + } + case float64: + // Numeric value - convert to string (matching Python's str()) + strVal := strconv.FormatFloat(v, 'f', -1, 64) + if _, exists := valueMap[strVal]; !exists { + valueMap[strVal] = []string{docID} + } else { + valueMap[strVal] = appendDocID(valueMap[strVal], docID) + } + case int: + // Integer value - convert to string + strVal := fmt.Sprintf("%d", v) + if _, exists := valueMap[strVal]; !exists { + valueMap[strVal] = []string{docID} + } else { + valueMap[strVal] = appendDocID(valueMap[strVal], docID) + } + case []interface{}: + // List of values (string, number, or time) + for _, item := range v { + switch itemVal := item.(type) { + case string: + if itemVal != "" { + if _, exists := valueMap[itemVal]; !exists { + valueMap[itemVal] = []string{docID} + } else { + valueMap[itemVal] = appendDocID(valueMap[itemVal], docID) + } + } + case float64: + strVal := strconv.FormatFloat(itemVal, 'f', -1, 64) + if _, exists := valueMap[strVal]; !exists { + valueMap[strVal] = []string{docID} + } else { + valueMap[strVal] = appendDocID(valueMap[strVal], docID) + } + case int: + strVal := fmt.Sprintf("%d", itemVal) + if _, exists := valueMap[strVal]; !exists { + valueMap[strVal] = []string{docID} + } else { + valueMap[strVal] = appendDocID(valueMap[strVal], docID) + } + } + } + } + } + } + + return flattedMeta, nil +} + // ExtractDocumentID extracts the document ID from a chunk func ExtractDocumentID(chunk map[string]interface{}) (string, bool) { docID, ok := chunk["id"].(string) @@ -160,11 +263,22 @@ func ExtractMetaFields(chunk map[string]interface{}) (map[string]interface{}, er return make(map[string]interface{}), nil } case []byte: - metaFields = ParseLengthPrefixedJSON(v) - if metaFields == nil { - if err := json.Unmarshal(v, &metaFields); err != nil { - return make(map[string]interface{}), nil + allResults := ParseAllLengthPrefixedJSON(v) + if len(allResults) > 0 { + // Merge all JSON objects - when same key appears with different values, collect all + metaFields = make(map[string]interface{}) + for _, result := range allResults { + for k, val := range result { + if existing, exists := metaFields[k]; exists { + // Key already exists - merge values + metaFields[k] = mergeFieldValues(existing, val) + } else { + metaFields[k] = val + } + } } + } else if err := json.Unmarshal(v, &metaFields); err != nil { + return make(map[string]interface{}), nil } default: return make(map[string]interface{}), nil @@ -173,6 +287,57 @@ func ExtractMetaFields(chunk map[string]interface{}) (map[string]interface{}, er return metaFields, nil } +// mergeFieldValues merges two field values when the same key appears multiple times +// If both are arrays, append all elements. If one is array and other is string, append string to array. +// Returns []interface{} with all merged values (flattened). +func mergeFieldValues(existing, new interface{}) []interface{} { + result := []interface{}{} + + var addValue func(v interface{}) + addValue = func(v interface{}) { + if v == nil { + return + } + switch val := v.(type) { + case string: + if val != "" { + result = append(result, val) + } + case []interface{}: + for _, item := range val { + addValue(item) + } + } + } + + addValue(existing) + addValue(new) + + return result +} + +// appendDocID appends a docID to an existing value that may be []string or []interface{} +func appendDocID(existing interface{}, docID string) []string { + result := []string{docID} + if existing == nil { + return result + } + switch v := existing.(type) { + case []string: + return append(v, docID) + case []interface{}: + for _, item := range v { + if s, ok := item.(string); ok { + result = append(result, s) + } + } + return result + case string: + return append(result, v) + } + return result +} + // ParseLengthPrefixedJSON parses Infinity's length-prefixed JSON format // Format: [4-byte length (little-endian)][JSON][4-byte length][JSON]... // Returns the FIRST valid JSON object found diff --git a/internal/service/metadata_filter.go b/internal/service/metadata_filter.go new file mode 100644 index 0000000000..5e445cf347 --- /dev/null +++ b/internal/service/metadata_filter.go @@ -0,0 +1,563 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package service + +import ( + "context" + "encoding/json" + "fmt" + "os" + "regexp" + "strings" + "time" + + "go.uber.org/zap" + + "ragflow/internal/entity" + modelModule "ragflow/internal/entity/models" + "ragflow/internal/logger" +) + +// MetaFilterCondition represents a single filter condition +type MetaFilterCondition struct { + Key string `json:"key"` + Value string `json:"value"` + Op string `json:"op"` +} + +// MetaFilterResult represents the result of LLM-generated filter +type MetaFilterResult struct { + Conditions []MetaFilterCondition `json:"conditions"` + Logic string `json:"logic"` +} + +// ManualValueResolver is a callback function to transform manual filter values +type ManualValueResolver func(map[string]interface{}) map[string]interface{} + +// metaFilterTemplateCache caches the template content +var metaFilterTemplateCache string + +// getMetaFilterTemplate loads and caches the meta_filter.md template +func getMetaFilterTemplate() (string, error) { + if metaFilterTemplateCache != "" { + return metaFilterTemplateCache, nil + } + + // Try to find meta_filter.md relative to the rag module + // Look for it in rag/prompts/ directory + possiblePaths := []string{ + "rag/prompts/meta_filter.md", + "../rag/prompts/meta_filter.md", + "../../rag/prompts/meta_filter.md", + } + + var templateContent string + for _, path := range possiblePaths { + content, err := os.ReadFile(path) + if err == nil { + templateContent = string(content) + break + } + } + + if templateContent == "" { + // Fallback: return error + return "", fmt.Errorf("could not find meta_filter.md template") + } + + metaFilterTemplateCache = templateContent + return templateContent, nil +} + +// renderMetaFilterTemplate renders the Jinja2-like template from meta_filter.md +func renderMetaFilterTemplate(currentDate, metadataKeys, question, constraints string) (string, error) { + templateContent, err := getMetaFilterTemplate() + if err != nil { + return "", err + } + + // Replace variables + result := strings.ReplaceAll(templateContent, "{{ current_date }}", currentDate) + result = strings.ReplaceAll(result, "{{ metadata_keys }}", metadataKeys) + result = strings.ReplaceAll(result, "{{ user_question }}", question) + + // Handle {% if constraints %}...{% endif %} + constraintRegex := regexp.MustCompile(`(?s)\{%\s*if\s+constraints\s*%\}(.+?)\{%\s*endif\s*%\}`) + if constraints != "" { + // Replace with the content inside the if block + result = constraintRegex.ReplaceAllString(result, "$1") + } else { + // Remove the entire if block + result = constraintRegex.ReplaceAllString(result, "") + } + + // Clean up any extra newlines from removed blocks + result = regexp.MustCompile(`\n{3,}`).ReplaceAllString(result, "\n\n") + + return strings.TrimSpace(result), nil +} + +// genMetaFilterPrompt builds the prompt for LLM-based metadata filter generation +func genMetaFilterPrompt(metaDataJSON, question, constraintsJSON, currentDate string) string { + prompt, err := renderMetaFilterTemplate(currentDate, metaDataJSON, question, constraintsJSON) + if err != nil { + logger.Warn("Failed to render meta filter template, using fallback", zap.Error(err)) + // Fallback to empty prompt + return "" + } + return prompt +} + +// GenMetaFilter generates filter conditions using LLM based on metadata and question. +func GenMetaFilter(ctx context.Context, creds *entity.ModelCredentials, metaData map[string]interface{}, question string, constraints map[string]string) (*MetaFilterResult, error) { + if creds == nil { + return nil, fmt.Errorf("model credentials is nil") + } + + if len(metaData) == 0 { + return &MetaFilterResult{Conditions: []MetaFilterCondition{}, Logic: "and"}, nil + } + + // Build metadata structure for prompt + metaDataStructure := make(map[string][]string) + for key, values := range metaData { + if valueMap, ok := values.(map[string]interface{}); ok { + keys := make([]string, 0, len(valueMap)) + for k := range valueMap { + keys = append(keys, k) + } + metaDataStructure[key] = keys + } + } + + metaDataJSON, _ := json.Marshal(metaDataStructure) + constraintsJSON := "" + if constraints != nil { + constraintsBytes, _ := json.Marshal(constraints) + constraintsJSON = string(constraintsBytes) + } + + // Build the prompt + currentDate := time.Now().Format("2006-01-02") + systemPrompt := genMetaFilterPrompt(string(metaDataJSON), question, constraintsJSON, currentDate) + + // Build user message + userMessage := "Generate filters:" + + // Build messages: system prompt + user message + messages := []modelModule.Message{ + {Role: "system", Content: systemPrompt}, + {Role: "user", Content: userMessage}, + } + + // Call LLM using ChatWithMessagesToModelByApiKey + modelProviderSvc := NewModelProviderService() + response, code, err := modelProviderSvc.ChatWithMessagesToModelByApiKey(creds.ProviderName, creds.ModelName, creds.APIKey, messages) + if err != nil { + logger.Warn("ChatWithMessagesToModelByApiKey failed for GenMetaFilter", + zap.String("provider", creds.ProviderName), + zap.String("model", creds.ModelName), + zap.Int("code", int(code)), + zap.Error(err)) + return nil, fmt.Errorf("failed to generate meta filter: %w", err) + } + + // Clean up response + responseStr := strings.TrimSpace(*response) + responseStr = thinkBlockRE.ReplaceAllString(responseStr, "") + responseStr = strings.TrimSpace(responseStr) + + // Remove markdown code blocks if present + responseStr = strings.TrimPrefix(responseStr, "```json") + responseStr = strings.TrimPrefix(responseStr, "```") + responseStr = strings.TrimSuffix(responseStr, "```") + responseStr = strings.TrimSpace(responseStr) + + // Parse JSON + var result MetaFilterResult + if err := json.Unmarshal([]byte(responseStr), &result); err != nil { + logger.Warn("Failed to parse meta filter response, returning empty conditions", zap.Error(err)) + return &MetaFilterResult{Conditions: []MetaFilterCondition{}, Logic: "and"}, nil + } + + logger.Info("GenMetaFilter result", zap.Any("conditions", result.Conditions), zap.String("logic", result.Logic)) + + return &result, nil +} + +// ApplyMetaFilter applies filter conditions to metadata and returns matching doc IDs +func ApplyMetaFilter(metaData map[string]interface{}, filters []MetaFilterCondition, logic string) []string { + if len(filters) == 0 { + return []string{} + } + + docIDSet := make(map[string]bool) + + for i, condition := range filters { + matchingIDs := applySingleCondition(metaData, condition) + if i == 0 { + for _, id := range matchingIDs { + docIDSet[id] = true + } + } else { + if logic == "or" { + // Union + for _, id := range matchingIDs { + docIDSet[id] = true + } + } else { + // AND - intersection + newSet := make(map[string]bool) + for _, id := range matchingIDs { + if docIDSet[id] { + newSet[id] = true + } + } + docIDSet = newSet + } + } + } + + // Convert to list + result := make([]string, 0, len(docIDSet)) + for id := range docIDSet { + result = append(result, id) + } + return result +} + +// applySingleCondition applies a single filter condition and returns matching doc IDs +func applySingleCondition(metaData map[string]interface{}, condition MetaFilterCondition) []string { + key := condition.Key + value := condition.Value + op := condition.Op + + valueMap, ok := metaData[key].(map[string]interface{}) + if !ok { + return []string{} + } + + var result []string + + switch op { + case "=", "==": + if docIDs, exists := valueMap[value]; exists { + switch v := docIDs.(type) { + case []interface{}: + for _, id := range v { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + case []string: + result = append(result, v...) + } + } + case "!=", "≠": + for val, docIDs := range valueMap { + if val != value { + if ids, ok := docIDs.([]interface{}); ok { + for _, id := range ids { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + } + } + } + case "contains": + for val, docIDs := range valueMap { + if strings.Contains(strings.ToLower(val), strings.ToLower(value)) { + if ids, ok := docIDs.([]interface{}); ok { + for _, id := range ids { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + } + } + } + case "not contains": + for val, docIDs := range valueMap { + if !strings.Contains(strings.ToLower(val), strings.ToLower(value)) { + if ids, ok := docIDs.([]interface{}); ok { + for _, id := range ids { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + } + } + } + case "in": + values := strings.Split(value, ",") + for _, v := range values { + v = strings.TrimSpace(v) + if docIDs, exists := valueMap[v]; exists { + if ids, ok := docIDs.([]interface{}); ok { + for _, id := range ids { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + } + } + } + case "not in": + excludeValues := make(map[string]bool) + for _, v := range strings.Split(value, ",") { + excludeValues[strings.TrimSpace(strings.ToLower(v))] = true + } + for val, docIDs := range valueMap { + if !excludeValues[strings.ToLower(val)] { + if ids, ok := docIDs.([]interface{}); ok { + for _, id := range ids { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + } + } + } + case "start with": + for val, docIDs := range valueMap { + if strings.HasPrefix(strings.ToLower(val), strings.ToLower(value)) { + if ids, ok := docIDs.([]interface{}); ok { + for _, id := range ids { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + } + } + } + case "end with": + for val, docIDs := range valueMap { + if strings.HasSuffix(strings.ToLower(val), strings.ToLower(value)) { + if ids, ok := docIDs.([]interface{}); ok { + for _, id := range ids { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + } + } + } + case "empty": + if len(valueMap) == 0 { + return []string{} + } + case "not empty": + if len(valueMap) > 0 { + for _, docIDs := range valueMap { + if ids, ok := docIDs.([]interface{}); ok { + for _, id := range ids { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + } + } + } + case ">": + for val, docIDs := range valueMap { + if val > value { + if ids, ok := docIDs.([]interface{}); ok { + for _, id := range ids { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + } + } + } + case "<": + for val, docIDs := range valueMap { + if val < value { + if ids, ok := docIDs.([]interface{}); ok { + for _, id := range ids { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + } + } + } + case ">=": + for val, docIDs := range valueMap { + if val >= value { + if ids, ok := docIDs.([]interface{}); ok { + for _, id := range ids { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + } + } + } + case "<=": + for val, docIDs := range valueMap { + if val <= value { + if ids, ok := docIDs.([]interface{}); ok { + for _, id := range ids { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + } + } + } + default: + // Default to equality check + if docIDs, exists := valueMap[value]; exists { + if ids, ok := docIDs.([]interface{}); ok { + for _, id := range ids { + if idStr, ok := id.(string); ok { + result = append(result, idStr) + } + } + } + } + } + + return result +} + +// ApplyMetaDataFilter applies metadata filtering rules and returns filtered doc_ids +// Supports three modes: +// - auto: generate filter conditions via LLM +// - semi_auto: generate conditions using selected metadata keys only via LLM +// - manual: directly filter based on provided conditions +func ApplyMetaDataFilter( + ctx context.Context, + metaDataFilter map[string]interface{}, + metaData map[string]interface{}, + question string, + creds *entity.ModelCredentials, + baseDocIDs []string, + manualValueResolver ...ManualValueResolver, +) ([]string, bool) { + if metaDataFilter == nil { + return baseDocIDs, false + } + + docIDs := make([]string, len(baseDocIDs)) + copy(docIDs, baseDocIDs) + + method, _ := metaDataFilter["method"].(string) + + switch method { + case "auto": + filters, err := GenMetaFilter(ctx, creds, metaData, question, nil) + if err != nil { + logger.Warn("Failed to generate meta filter", zap.Error(err)) + return docIDs, false + } + filteredIDs := ApplyMetaFilter(metaData, filters.Conditions, filters.Logic) + docIDs = append(docIDs, filteredIDs...) + if len(docIDs) == 0 { + return nil, true // Return nil to indicate auto filter returned empty + } + + case "semi_auto": + selectedKeys := []string{} + constraints := make(map[string]string) + + if semiAuto, ok := metaDataFilter["semi_auto"].([]interface{}); ok { + for _, item := range semiAuto { + switch v := item.(type) { + case string: + selectedKeys = append(selectedKeys, v) + case map[string]interface{}: + if key, ok := v["key"].(string); ok { + selectedKeys = append(selectedKeys, key) + if op, ok := v["op"].(string); ok { + constraints[key] = op + } + } + } + } + } + + if len(selectedKeys) > 0 { + // Filter metadata to only selected keys + filteredMeta := make(map[string]interface{}) + for _, key := range selectedKeys { + if val, exists := metaData[key]; exists { + filteredMeta[key] = val + } + } + + if len(filteredMeta) > 0 { + filters, err := GenMetaFilter(ctx, creds, filteredMeta, question, constraints) + if err != nil { + logger.Warn("Failed to generate meta filter", zap.Error(err)) + return docIDs, false + } + filteredIDs := ApplyMetaFilter(metaData, filters.Conditions, filters.Logic) + docIDs = append(docIDs, filteredIDs...) + if len(docIDs) == 0 { + return nil, true + } + } + } + + case "manual": + manualFilters, _ := metaDataFilter["manual"].([]interface{}) + logic := "and" + if logicVal, ok := metaDataFilter["logic"].(string); ok { + logic = logicVal + } + + // Apply manual_value_resolver callback if provided + if len(manualValueResolver) > 0 && manualValueResolver[0] != nil { + resolver := manualValueResolver[0] + resolvedFilters := make([]interface{}, 0, len(manualFilters)) + for _, item := range manualFilters { + if cond, ok := item.(map[string]interface{}); ok { + resolvedFilters = append(resolvedFilters, resolver(cond)) + } + } + manualFilters = resolvedFilters + } + + conditions := make([]MetaFilterCondition, 0, len(manualFilters)) + for _, item := range manualFilters { + if cond, ok := item.(map[string]interface{}); ok { + condition := MetaFilterCondition{} + if key, ok := cond["key"].(string); ok { + condition.Key = key + } + if value, ok := cond["value"].(string); ok { + condition.Value = value + } + if op, ok := cond["op"].(string); ok { + condition.Op = op + } + conditions = append(conditions, condition) + } + } + + filteredIDs := ApplyMetaFilter(metaData, conditions, logic) + docIDs = append(docIDs, filteredIDs...) + if len(manualFilters) > 0 && len(docIDs) == 0 { + return []string{"-999"}, false + } + } + + return docIDs, false +} diff --git a/internal/service/model_service.go b/internal/service/model_service.go index 1eb71a1432..3862bd4e2f 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -87,13 +87,19 @@ func (p *ModelProviderImpl) GetEmbeddingModel(ctx context.Context, tenantID stri if apiKey == nil || *apiKey == "" { return nil, fmt.Errorf("no API key found for tenant %s and model %s", tenantID, compositeModelName) } - // Always get API base from model provider configuration - providerDAO := dao.NewModelProviderDAO() - providerConfig := providerDAO.GetProviderByName(provider) - if providerConfig == nil || providerConfig.DefaultURL == "" { - return nil, fmt.Errorf("no API base found for provider %s", provider) + + // Get API base from TenantLLM if set, otherwise from model provider configuration + apiBase := "" + if embeddingModel.APIBase != nil && *embeddingModel.APIBase != "" { + apiBase = *embeddingModel.APIBase + } else { + providerDAO := dao.NewModelProviderDAO() + providerConfig := providerDAO.GetProviderByName(provider) + if providerConfig == nil || providerConfig.DefaultURL == "" { + return nil, fmt.Errorf("no API base found for provider %s", provider) + } + apiBase = providerConfig.DefaultURL } - apiBase := fmt.Sprintf("%sembeddings/", providerConfig.DefaultURL) return models.CreateEmbeddingModel(provider, *apiKey, apiBase, modelName, p.httpClient) } @@ -101,23 +107,71 @@ func (p *ModelProviderImpl) GetEmbeddingModel(ctx context.Context, tenantID stri // GetChatModel returns a chat model for the given tenant func (p *ModelProviderImpl) GetChatModel(ctx context.Context, tenantID string, compositeModelName string) (entity.ChatModel, error) { // Parse composite model name to extract model name and provider - _, _, err := parseModelName(compositeModelName) + modelName, provider, err := parseModelName(compositeModelName) if err != nil { return nil, err } - // TODO: implement chat model creation - return nil, fmt.Errorf("chat model not implemented yet for model: %s", compositeModelName) + + // Get chat model from database + chatModel, err := dao.NewTenantLLMDAO().GetByTenantFactoryAndModelName(tenantID, provider, modelName) + if err != nil { + return nil, fmt.Errorf("no chat model found for tenant %s and model %s: %w", tenantID, compositeModelName, err) + } + + apiKey := chatModel.APIKey + if apiKey == nil || *apiKey == "" { + return nil, fmt.Errorf("no API key found for tenant %s and model %s", tenantID, compositeModelName) + } + + // Get API base from TenantLLM if set, otherwise from model provider configuration + apiBase := "" + if chatModel.APIBase != nil && *chatModel.APIBase != "" { + apiBase = *chatModel.APIBase + } else { + providerDAO := dao.NewModelProviderDAO() + providerConfig := providerDAO.GetProviderByName(provider) + if providerConfig == nil || providerConfig.DefaultURL == "" { + return nil, fmt.Errorf("no API base found for provider %s", provider) + } + apiBase = providerConfig.DefaultURL + } + + return models.CreateChatModel(provider, *apiKey, apiBase, modelName, p.httpClient) } // GetRerankModel returns a rerank model for the given tenant func (p *ModelProviderImpl) GetRerankModel(ctx context.Context, tenantID string, compositeModelName string) (entity.RerankModel, error) { // Parse composite model name to extract model name and provider - _, _, err := parseModelName(compositeModelName) + modelName, provider, err := parseModelName(compositeModelName) if err != nil { return nil, err } - // TODO: implement rerank model creation - return nil, fmt.Errorf("rerank model not implemented yet for model: %s", compositeModelName) + + // Get rerank model from database + rerankModel, err := dao.NewTenantLLMDAO().GetByTenantFactoryAndModelName(tenantID, provider, modelName) + if err != nil { + return nil, fmt.Errorf("no rerank model found for tenant %s and model %s: %w", tenantID, compositeModelName, err) + } + + apiKey := rerankModel.APIKey + if apiKey == nil || *apiKey == "" { + return nil, fmt.Errorf("no API key found for tenant %s and model %s", tenantID, compositeModelName) + } + + // Get API base from TenantLLM if set, otherwise from model provider configuration + apiBase := "" + if rerankModel.APIBase != nil && *rerankModel.APIBase != "" { + apiBase = *rerankModel.APIBase + } else { + providerDAO := dao.NewModelProviderDAO() + providerConfig := providerDAO.GetProviderByName(provider) + if providerConfig == nil || providerConfig.DefaultURL == "" { + return nil, fmt.Errorf("no API base found for provider %s", provider) + } + apiBase = providerConfig.DefaultURL + } + + return models.CreateRerankModel(provider, *apiKey, apiBase, modelName, p.httpClient) } func NewModelProviderService() *ModelProviderService { @@ -743,6 +797,49 @@ func (m *ModelProviderService) ChatToModel(providerName, instanceName, modelName return nil, common.CodeServerError, errors.New("model is disabled") } +func (m *ModelProviderService) ChatToModelByApiKey(providerName, modelName, apiKey, message string) (*string, common.ErrorCode, error) { + providerInfo := dao.GetModelProviderManager().FindProvider(providerName) + if providerInfo == nil { + return nil, common.CodeNotFound, errors.New("provider not found") + } + + _, err := dao.GetModelProviderManager().GetModelByName(providerName, modelName) + if err != nil { + return nil, common.CodeNotFound, errors.New(fmt.Sprintf("provider %s model %s not found", providerName, modelName)) + } + + var apiConfig = &modelModule.APIConfig{} + apiConfig.ApiKey = &apiKey + var response *modelModule.ChatResponse + response, err = providerInfo.ModelDriver.Chat(&modelName, &message, apiConfig, nil) + if err != nil { + return nil, common.CodeServerError, err + } + + return response.Answer, common.CodeSuccess, nil +} + +// ChatWithMessagesToModelByApiKey sends multiple messages with roles and returns response +func (m *ModelProviderService) ChatWithMessagesToModelByApiKey(providerName, modelName, apiKey string, messages []modelModule.Message) (*string, common.ErrorCode, error) { + providerInfo := dao.GetModelProviderManager().FindProvider(providerName) + if providerInfo == nil { + return nil, common.CodeNotFound, errors.New("provider not found") + } + + _, err := dao.GetModelProviderManager().GetModelByName(providerName, modelName) + if err != nil { + return nil, common.CodeNotFound, errors.New(fmt.Sprintf("provider %s model %s not found", providerName, modelName)) + } + + var response string + response, err = providerInfo.ModelDriver.ChatWithMessages(modelName, &apiKey, messages, nil) + if err != nil { + return nil, common.CodeServerError, err + } + + return &response, common.CodeSuccess, nil +} + // ChatToModelStreamWithSender streams chat response directly via sender function (best performance, no channel) func (m *ModelProviderService) ChatToModelStreamWithSender(providerName, instanceName, modelName, userID, message string, apiConfig *modelModule.APIConfig, modelConfig *modelModule.ChatConfig, sender func(*string, *string) error) common.ErrorCode { // Get tenant ID from user @@ -801,3 +898,75 @@ func (m *ModelProviderService) ChatToModelStreamWithSender(providerName, instanc return common.CodeServerError } + +func (m *ModelProviderService) GetDefaultModel(modelType entity.ModelType, tenantID string) (*entity.ModelCredentials, error) { + // Get tenant record to find default model name + tenant, err := dao.NewTenantDAO().GetByID(tenantID) + if err != nil { + return nil, fmt.Errorf("tenant not found: %w", err) + } + + // Determine model name based on model type + var defaultModelName string + switch modelType { + case entity.ModelTypeChat: + defaultModelName = tenant.LLMID + case entity.ModelTypeEmbedding: + defaultModelName = tenant.EmbdID + case entity.ModelTypeSpeech2Text: + defaultModelName = tenant.ASRID + case entity.ModelTypeImage2Text: + defaultModelName = tenant.Img2TxtID + case entity.ModelTypeRerank: + defaultModelName = tenant.RerankID + case entity.ModelTypeTTS: + if tenant.TTSID != nil { + defaultModelName = *tenant.TTSID + } + case entity.ModelTypeOCR: + return nil, errors.New("OCR model name is required") + default: + return nil, fmt.Errorf("unknown model type: %s", modelType) + } + + if defaultModelName == "" { + return nil, fmt.Errorf("no default %s model is set", modelType) + } + + // Look up the TenantLLM record to get provider name and API key + // Use GetByTenantIDAndLLMName which handles splitting model name and factory + tenantLLM, err := dao.NewTenantLLMDAO().GetByTenantIDAndLLMName(tenantID, defaultModelName) + if err != nil { + return nil, fmt.Errorf("failed to get tenant default model: %w", err) + } + + if tenantLLM == nil { + return nil, fmt.Errorf("no default %s model found for tenant", modelType) + } + + if tenantLLM.LLMName == nil || tenantLLM.APIKey == nil { + return nil, fmt.Errorf("tenant model %q has missing name or api key", defaultModelName) + } + return &entity.ModelCredentials{ + ProviderName: tenantLLM.LLMFactory, + ModelName: *tenantLLM.LLMName, + APIKey: *tenantLLM.APIKey, + }, nil +} + +// GetModelByName gets model credentials by model name (chat_id from search_config) +func (m *ModelProviderService) GetModelByName(modelName string, tenantID string) (*entity.ModelCredentials, error) { + tenantLLM, err := dao.NewTenantLLMDAO().GetByTenantIDAndLLMName(tenantID, modelName) + if err != nil { + return nil, fmt.Errorf("failed to get model by name: %w", err) + } + if tenantLLM == nil { + return nil, fmt.Errorf("model not found: %s", modelName) + } + + return &entity.ModelCredentials{ + ProviderName: tenantLLM.LLMFactory, + ModelName: *tenantLLM.LLMName, + APIKey: *tenantLLM.APIKey, + }, nil +} diff --git a/internal/service/models/factory.go b/internal/service/models/factory.go index 6a148e4417..b3ed9c5c76 100644 --- a/internal/service/models/factory.go +++ b/internal/service/models/factory.go @@ -27,8 +27,16 @@ import ( // EmbeddingModelFactory creates an EmbeddingModel instance type EmbeddingModelFactory func(apiKey, apiBase, modelName string, httpClient *http.Client) entity.EmbeddingModel +// ChatModelFactory creates a ChatModel instance +type ChatModelFactory func(apiKey, apiBase, modelName string, httpClient *http.Client) entity.ChatModel + +// RerankModelFactory creates a RerankModel instance +type RerankModelFactory func(apiKey, apiBase, modelName string, httpClient *http.Client) entity.RerankModel + var ( embeddingModelFactories = make(map[string]EmbeddingModelFactory) + chatModelFactories = make(map[string]ChatModelFactory) + rerankModelFactories = make(map[string]RerankModelFactory) factoryMu sync.RWMutex ) @@ -40,6 +48,22 @@ func RegisterEmbeddingModelFactory(providerName string, factory EmbeddingModelFa embeddingModelFactories[providerName] = factory } +// RegisterChatModelFactory registers a factory for a chat provider name. +// Should be called from init() functions of provider implementations. +func RegisterChatModelFactory(providerName string, factory ChatModelFactory) { + factoryMu.Lock() + defer factoryMu.Unlock() + chatModelFactories[providerName] = factory +} + +// RegisterRerankModelFactory registers a factory for a rerank provider name. +// Should be called from init() functions of provider implementations. +func RegisterRerankModelFactory(providerName string, factory RerankModelFactory) { + factoryMu.Lock() + defer factoryMu.Unlock() + rerankModelFactories[providerName] = factory +} + // GetEmbeddingModelFactory returns the factory for the given provider name. // Returns nil if not found. func GetEmbeddingModelFactory(providerName string) EmbeddingModelFactory { @@ -48,6 +72,22 @@ func GetEmbeddingModelFactory(providerName string) EmbeddingModelFactory { return embeddingModelFactories[providerName] } +// GetChatModelFactory returns the factory for the given chat provider name. +// Returns nil if not found. +func GetChatModelFactory(providerName string) ChatModelFactory { + factoryMu.RLock() + defer factoryMu.RUnlock() + return chatModelFactories[providerName] +} + +// GetRerankModelFactory returns the factory for the given rerank provider name. +// Returns nil if not found. +func GetRerankModelFactory(providerName string) RerankModelFactory { + factoryMu.RLock() + defer factoryMu.RUnlock() + return rerankModelFactories[providerName] +} + // CreateEmbeddingModel creates an EmbeddingModel instance for the given provider. // Returns error if provider not registered. func CreateEmbeddingModel(providerName, apiKey, apiBase, modelName string, httpClient *http.Client) (entity.EmbeddingModel, error) { @@ -57,3 +97,23 @@ func CreateEmbeddingModel(providerName, apiKey, apiBase, modelName string, httpC } return factory(apiKey, apiBase, modelName, httpClient), nil } + +// CreateChatModel creates a ChatModel instance for the given provider. +// Returns error if provider not registered. +func CreateChatModel(providerName, apiKey, apiBase, modelName string, httpClient *http.Client) (entity.ChatModel, error) { + factory := GetChatModelFactory(providerName) + if factory == nil { + return nil, fmt.Errorf("no chat model factory registered for provider %s", providerName) + } + return factory(apiKey, apiBase, modelName, httpClient), nil +} + +// CreateRerankModel creates a RerankModel instance for the given provider. +// Returns error if provider not registered. +func CreateRerankModel(providerName, apiKey, apiBase, modelName string, httpClient *http.Client) (entity.RerankModel, error) { + factory := GetRerankModelFactory(providerName) + if factory == nil { + return nil, fmt.Errorf("no rerank model factory registered for provider %s", providerName) + } + return factory(apiKey, apiBase, modelName, httpClient), nil +} diff --git a/internal/service/models/siliconflow_model.go b/internal/service/models/siliconflow_model.go index 0333da2d07..75f89f3525 100644 --- a/internal/service/models/siliconflow_model.go +++ b/internal/service/models/siliconflow_model.go @@ -34,6 +34,22 @@ type siliconflowEmbeddingModel struct { httpClient *http.Client } +// siliconflowChatModel implements ChatModel for SILICONFLOW API +type siliconflowChatModel struct { + apiKey string + apiBase string + model string + httpClient *http.Client +} + +// siliconflowRerankModel implements RerankModel for SILICONFLOW API +type siliconflowRerankModel struct { + apiKey string + apiBase string + model string + httpClient *http.Client +} + // SiliconflowEmbeddingRequest represents SILICONFLOW embedding request type SiliconflowEmbeddingRequest struct { Model string `json:"model"` @@ -48,6 +64,54 @@ type SiliconflowEmbeddingResponse struct { } `json:"data"` } +// SiliconflowChatRequest represents SILICONFLOW chat request +type SiliconflowChatRequest struct { + Model string `json:"model"` + Messages []ChatMessage `json:"messages"` + Temperature float64 `json:"temperature,omitempty"` + MaxTokens int `json:"max_tokens,omitempty"` + Stream bool `json:"stream,omitempty"` +} + +// SiliconflowChatResponse represents SILICONFLOW chat response +type SiliconflowChatResponse struct { + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + FinishReason string `json:"finish_reason"` + } `json:"choices"` + Error struct { + Message string `json:"message"` + Code string `json:"code"` + } `json:"error,omitempty"` +} + +// ChatMessage represents a chat message +type ChatMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +// SiliconflowRerankRequest represents SILICONFLOW rerank request +type SiliconflowRerankRequest struct { + Model string `json:"model"` + Query string `json:"query"` + Documents []string `json:"documents"` + TopN int `json:"top_n"` + ReturnDocuments bool `json:"return_documents"` + MaxChunksPerDoc int `json:"max_chunks_per_doc"` + OverlapTokens int `json:"overlap_tokens"` +} + +// SiliconflowRerankResponse represents SILICONFLOW rerank response +type SiliconflowRerankResponse struct { + Results []struct { + Index int `json:"index"` + RelevanceScore float64 `json:"relevance_score"` + } `json:"results"` +} + // Encode encodes a list of texts into embeddings using SILICONFLOW API func (m *siliconflowEmbeddingModel) Encode(texts []string) ([][]float64, error) { if len(texts) == 0 { @@ -111,7 +175,181 @@ func (m *siliconflowEmbeddingModel) EncodeQuery(query string) ([]float64, error) return embeddings[0], nil } -// init registers the SILICONFLOW embedding model factory +// Chat sends a chat message and returns response +func (m *siliconflowChatModel) Chat(system string, history []map[string]string, genConf map[string]interface{}) (string, error) { + // Build messages array + var messages []ChatMessage + + // Add system message if provided + if system != "" { + messages = append(messages, ChatMessage{Role: "system", Content: system}) + } + + // Add history messages + for _, msg := range history { + role := msg["role"] + content := msg["content"] + if role != "" && content != "" { + messages = append(messages, ChatMessage{Role: role, Content: content}) + } + } + + // Extract generation config + temperature := 0.7 + if temp, ok := genConf["temperature"].(float64); ok { + temperature = temp + } + maxTokens := 1024 + if mt, ok := genConf["max_tokens"].(int); ok { + maxTokens = mt + } + + // Build request + reqBody := SiliconflowChatRequest{ + Model: m.model, + Messages: messages, + Temperature: temperature, + MaxTokens: maxTokens, + } + + jsonData, err := json.Marshal(reqBody) + if err != nil { + return "", fmt.Errorf("failed to marshal request: %w", err) + } + + // Build URL - append /chat/completions if not already present + url := m.apiBase + if !strings.HasSuffix(url, "/chat/completions") { + if !strings.HasSuffix(url, "/") { + url += "/" + } + url += "chat/completions" + } + + req, err := http.NewRequest("POST", url, strings.NewReader(string(jsonData))) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+m.apiKey) + + resp, err := m.httpClient.Do(req) + if err != nil { + return "", fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return "", fmt.Errorf("failed to read response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("SILICONFLOW API error: %s, body: %s", resp.Status, string(body)) + } + + var chatResp SiliconflowChatResponse + if err := json.Unmarshal(body, &chatResp); err != nil { + return "", fmt.Errorf("failed to decode response: %w", err) + } + + if chatResp.Error.Message != "" { + return "", fmt.Errorf("chat error: %s", chatResp.Error.Message) + } + + if len(chatResp.Choices) == 0 { + return "", fmt.Errorf("no response choices returned") + } + + return chatResp.Choices[0].Message.Content, nil +} + +// ChatStreamly sends a chat message and streams response +func (m *siliconflowChatModel) ChatStreamly(system string, history []map[string]string, genConf map[string]interface{}) (<-chan string, error) { + // For now, return a simple non-streaming implementation + // Streaming can be implemented later with SSE support + responseChan := make(chan string) + + go func() { + defer close(responseChan) + response, err := m.Chat(system, history, genConf) + if err != nil { + responseChan <- "**ERROR**: " + err.Error() + return + } + responseChan <- response + }() + + return responseChan, nil +} + +// Similarity calculates similarity scores between query and texts using SiliconFlow API +func (m *siliconflowRerankModel) Similarity(query string, texts []string) ([]float64, error) { + if len(texts) == 0 { + return []float64{}, nil + } + + reqBody := SiliconflowRerankRequest{ + Model: m.model, + Query: query, + Documents: texts, + TopN: len(texts), + ReturnDocuments: false, + MaxChunksPerDoc: 1024, + OverlapTokens: 80, + } + + jsonData, err := json.Marshal(reqBody) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + reqURL := m.apiBase + if !strings.Contains(reqURL, "/rerank") { + if !strings.HasSuffix(reqURL, "/") { + reqURL += "/" + } + reqURL += "rerank" + } + + req, err := http.NewRequest("POST", reqURL, strings.NewReader(string(jsonData))) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+m.apiKey) + + resp, err := m.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("SiliconFlow Rerank API error: %s, body: %s", resp.Status, string(body)) + } + + body, _ := io.ReadAll(resp.Body) + + var rerankResp SiliconflowRerankResponse + if err := json.Unmarshal(body, &rerankResp); err != nil { + return nil, fmt.Errorf("failed to decode response: %w", err) + } + + scores := make([]float64, len(texts)) + for _, result := range rerankResp.Results { + if result.Index >= 0 && result.Index < len(texts) { + scores[result.Index] = result.RelevanceScore + } + } + + return scores, nil +} + +// init registers the SILICONFLOW model factories func init() { RegisterEmbeddingModelFactory("SILICONFLOW", func(apiKey, apiBase, modelName string, httpClient *http.Client) entity.EmbeddingModel { return &siliconflowEmbeddingModel{ @@ -121,4 +359,22 @@ func init() { httpClient: httpClient, } }) + + RegisterChatModelFactory("SILICONFLOW", func(apiKey, apiBase, modelName string, httpClient *http.Client) entity.ChatModel { + return &siliconflowChatModel{ + apiKey: apiKey, + apiBase: apiBase, + model: modelName, + httpClient: httpClient, + } + }) + + RegisterRerankModelFactory("SILICONFLOW", func(apiKey, apiBase, modelName string, httpClient *http.Client) entity.RerankModel { + return &siliconflowRerankModel{ + apiKey: apiKey, + apiBase: apiBase, + model: modelName, + httpClient: httpClient, + } + }) } diff --git a/internal/service/nlp/query_builder.go b/internal/service/nlp/query_builder.go index 1a4cdf37b3..991bcdb53d 100644 --- a/internal/service/nlp/query_builder.go +++ b/internal/service/nlp/query_builder.go @@ -21,8 +21,9 @@ import ( "sort" "strings" "sync" + "unicode/utf8" - "ragflow/internal/engine/infinity" + "ragflow/internal/engine/types" "ragflow/internal/tokenizer" "github.com/siongui/gojianfan" @@ -198,7 +199,7 @@ func (qb *QueryBuilder) Traditional2Simplified(line string) string { // NeedFineGrainedTokenize determines if fine-grained tokenization is needed for a token. // Reference: rag/nlp/query.py L88-93 func (qb *QueryBuilder) NeedFineGrainedTokenize(tk string) bool { - if len(tk) < 3 { + if utf8.RuneCountInString(tk) < 3 { return false } if matched, _ := regexp.MatchString(`^[0-9a-z\.\+#_\*-]+$`, tk); matched { @@ -209,8 +210,7 @@ func (qb *QueryBuilder) NeedFineGrainedTokenize(tk string) bool { // Question builds a full-text query expression based on input text. // References Python FulltextQueryer.question method. -// Currently, a simplified version, returns basic MatchTextExpr; future integration of term weight and synonyms. -func (qb *QueryBuilder) Question(txt string, tbl string, minMatch float64) (*infinity.MatchTextExpr, []string) { +func (qb *QueryBuilder) Question(txt string, tbl string, minMatch float64) (*types.MatchTextExpr, []string) { // originalQuery stores the original input text for later use in query expression. originalQuery := txt @@ -299,10 +299,27 @@ func (qb *QueryBuilder) Question(txt string, tbl string, minMatch float64) (*inf tksW = tksW[:256] } - // TODO: Synonym expansion (reference L61-67) - // For now, use empty synonyms - // syns is a placeholder for synonym expansion (currently empty). + // Synonym expansion + // Look up synonyms for each token syns := make([]string, len(tksW)) + for i, tw := range tksW { + tk := tw.tk + // Lookup synonyms (limit to 8 per Python) + tkSyns := qb.synonym.Lookup(tk, 8) + if len(tkSyns) > 0 { + // Format synonyms with weight boost: term^weight + var synParts []string + for _, syn := range tkSyns { + syn = strings.TrimSpace(syn) + if syn != "" { + synParts = append(synParts, fmt.Sprintf(`"%s"^%.1f`, syn, tw.w/4.0)) + } + } + syns[i] = strings.Join(synParts, " ") + } else { + syns[i] = "" + } + } // Build query parts // Reference: rag/nlp/query.py L69-70 @@ -316,7 +333,7 @@ func (qb *QueryBuilder) Question(txt string, tbl string, minMatch float64) (*inf continue } // Format: (token^weight synonym) - q = append(q, fmt.Sprintf("(%s^%.4f %s)", tk, w, syns[i])) + q = append(q, fmt.Sprintf("(%s^%.1f %s)", tk, w, syns[i])) } // Add phrase queries for adjacent tokens @@ -332,7 +349,7 @@ func (qb *QueryBuilder) Question(txt string, tbl string, minMatch float64) (*inf if tksW[i].w > maxW { maxW = tksW[i].w } - q = append(q, fmt.Sprintf(`"%s %s"^%.4f`, left, right, maxW*2)) + q = append(q, fmt.Sprintf(`"%s %s"^%.1f`, left, right, maxW*2)) } if len(q) == 0 { @@ -341,7 +358,7 @@ func (qb *QueryBuilder) Question(txt string, tbl string, minMatch float64) (*inf // query is the final query string built from all query parts. query := strings.Join(q, " ") - return &infinity.MatchTextExpr{ + return &types.MatchTextExpr{ Fields: qb.queryFields, MatchingText: query, TopN: 100, @@ -504,7 +521,7 @@ func (qb *QueryBuilder) Question(txt string, tbl string, minMatch float64) (*inf // termParts collects query parts for each term in the segment. var termParts []string for _, termWeight := range terms { - termParts = append(termParts, fmt.Sprintf("(%s)^%.4f", termWeight.term, termWeight.weight)) + termParts = append(termParts, fmt.Sprintf("(%s)^%.1f", termWeight.term, termWeight.weight)) } // tmsStr is the query string for the current segment. tmsStr := strings.Join(termParts, " ") @@ -557,7 +574,7 @@ func (qb *QueryBuilder) Question(txt string, tbl string, minMatch float64) (*inf if query == "" { query = otxt } - return &infinity.MatchTextExpr{ + return &types.MatchTextExpr{ Fields: qb.queryFields, MatchingText: query, TopN: 100, @@ -573,7 +590,7 @@ func (qb *QueryBuilder) Question(txt string, tbl string, minMatch float64) (*inf // Paragraph builds a query expression based on content terms and keywords. // References Python FulltextQueryer.paragraph method. -func (qb *QueryBuilder) Paragraph(contentTks string, keywords []string, keywordsTopN int) *infinity.MatchTextExpr { +func (qb *QueryBuilder) Paragraph(contentTks string, keywords []string, keywordsTopN int) *types.MatchTextExpr { // Simplified implementation: merge keywords and content terms allTerms := make([]string, 0, len(keywords)) for _, k := range keywords { @@ -598,7 +615,7 @@ func (qb *QueryBuilder) Paragraph(contentTks string, keywords []string, keywords } _ = calc } - return &infinity.MatchTextExpr{ + return &types.MatchTextExpr{ Fields: qb.queryFields, MatchingText: query, TopN: 100, diff --git a/internal/service/nlp/reranker.go b/internal/service/nlp/reranker.go index 7ac1a2a31a..0ab4d1c5c8 100644 --- a/internal/service/nlp/reranker.go +++ b/internal/service/nlp/reranker.go @@ -15,11 +15,17 @@ package nlp import ( + "encoding/json" "math" - "ragflow/internal/engine" + "regexp" "sort" "strconv" "strings" + + "ragflow/internal/common" + "ragflow/internal/logger" + + "go.uber.org/zap" ) // RerankModel defines the interface for reranker models @@ -55,69 +61,70 @@ type SearchResult struct { // - vsim: vector similarity scores func Rerank( rerankModel RerankModel, - resp *engine.SearchResponse, + chunks []map[string]interface{}, + total int, keywords []string, questionVector []float64, - sres *SearchResult, query string, tkWeight, vtWeight float64, useInfinity bool, cfield string, qb *QueryBuilder, + rankFeature map[string]float64, ) (sim []float64, tsim []float64, vsim []float64) { // If reranker model is provided and there are results, use model reranking - if rerankModel != nil && resp.Total > 0 { - return RerankByModel(rerankModel, nil, query, tkWeight, vtWeight, cfield, qb) + if rerankModel != nil && total > 0 { + return RerankByModel(rerankModel, chunks, query, tkWeight, vtWeight, cfield, qb, rankFeature) } // Otherwise, use fallback logic based on engine type if useInfinity { // For Infinity: scores are already normalized before fusion // Just extract the scores from results - // Check if there are results to rerank - if resp == nil || resp.Total == 0 || len(resp.Chunks) == 0 { + if chunks == nil || total == 0 || len(chunks) == 0 { return []float64{}, []float64{}, []float64{} } - return RerankInfinityFallback(resp) + return RerankInfinityFallback(chunks) } - // For Elasticsearch: need to perform reranking - return RerankStandard(resp, keywords, questionVector, nil, query, tkWeight, vtWeight, cfield, qb) + // For Elasticsearch: need to perform reranking and apply rank features + return RerankStandard(chunks, keywords, questionVector, query, tkWeight, vtWeight, cfield, qb, rankFeature) } // RerankByModel performs reranking using a reranker model -// Reference: rag/nlp/search.py L333-L354 func RerankByModel( rerankModel RerankModel, - sres *SearchResult, + chunks []map[string]interface{}, query string, tkWeight, vtWeight float64, cfield string, qb *QueryBuilder, + rankFeature map[string]float64, ) (sim []float64, tsim []float64, vsim []float64) { - if sres.Total == 0 || len(sres.IDs) == 0 { + if chunks == nil || len(chunks) == 0 { return []float64{}, []float64{}, []float64{} } + chunkCount := len(chunks) + + logger.Info("RerankByModel started", zap.String("query", query), zap.Int("chunkCount", chunkCount), zap.Float64("tkWeight", tkWeight), zap.Float64("vtWeight", vtWeight)) + // Extract keywords from query - _, keywords := qb.Question(query, "qa", 0.6) + keywords := []string{} + if qb != nil { + _, keywords = qb.Question(query, "qa", 0.6) + } + logger.Info("RerankByModel keywords extracted", zap.Any("keywords", keywords)) // Build token lists and document texts for each chunk - insTw := make([][]string, 0, len(sres.IDs)) - docs := make([]string, 0, len(sres.IDs)) + insTw := make([][]string, 0, chunkCount) + docs := make([]string, 0, chunkCount) - for _, id := range sres.IDs { - fields := sres.Field[id] - if fields == nil { - insTw = append(insTw, []string{}) - docs = append(docs, "") - continue - } - - contentLtks := extractContentTokens(fields, cfield) - titleTks := extractTitleTokens(fields) - importantKwd := extractImportantKeywords(fields) + for _, chunk := range chunks { + contentLtks := extractContentTokens(chunk, cfield) + titleTks := extractTitleTokens(chunk) + importantKwd := extractImportantKeywords(chunk) // Combine tokens without repetition (simpler version for model reranking) tks := make([]string, 0, len(contentLtks)+len(titleTks)+len(importantKwd)) @@ -127,7 +134,7 @@ func RerankByModel( insTw = append(insTw, tks) // Build document text for model reranking - docText := removeRedundantSpaces(strings.Join(tks, " ")) + docText := RemoveRedundantSpaces(strings.Join(tks, " ")) docs = append(docs, docText) } @@ -137,38 +144,57 @@ func RerankByModel( // Get similarity scores from reranker model modelSim, err := rerankModel.Similarity(query, docs) if err != nil { + logger.Error("RerankByModel: rerankModel.Similarity failed; falling back to token-only similarity", err) // If model fails, fall back to token similarity only modelSim = make([]float64, len(tsim)) } - + if len(modelSim) != chunkCount { + logger.Warn("reranker returned mismatched score length; padding/truncating", + zap.Int("got", len(modelSim)), zap.Int("want", chunkCount)) + fixed := make([]float64, chunkCount) + copy(fixed, modelSim) + modelSim = fixed + } // Combine token similarity with model similarity // Model similarity is treated as vector similarity component - sim = make([]float64, len(tsim)) + sim = make([]float64, chunkCount) for i := range tsim { sim[i] = tkWeight*tsim[i] + vtWeight*modelSim[i] } + // Apply rank feature scores (tag_score * 10 + pagerank) + // Always apply pageranks, even when rankFeature is nil/empty + sim = applyRankFeatureScores(chunks, sim, rankFeature) + + logger.Info("RerankByModel completed") return sim, tsim, modelSim } // RerankStandard performs standard reranking without a reranker model // Used for Elasticsearch when no reranker model is provided -// Reference: rag/nlp/search.py L294-L331 func RerankStandard( - resp *engine.SearchResponse, + chunks []map[string]interface{}, keywords []string, questionVector []float64, - sres *SearchResult, query string, tkWeight, vtWeight float64, cfield string, qb *QueryBuilder, + rankFeature map[string]float64, ) (sim []float64, tsim []float64, vsim []float64) { - chunkCount := len(resp.Chunks) - if resp.Total == 0 || chunkCount == 0 { + chunkCount := len(chunks) + if chunkCount == 0 { return []float64{}, []float64{}, []float64{} } + logger.Info("RerankStandard started", zap.Int("chunkCount", chunkCount), zap.Float64("tkWeight", tkWeight), zap.Float64("vtWeight", vtWeight)) + + // Compute keywords fresh from query + if qb != nil && len(keywords) == 0 { + _, keywords = qb.Question(query, "qa", 0.6) + } + logger.Info("RerankStandard keywords", zap.Any("keywords", keywords)) + // Get vector information vectorSize := len(questionVector) vectorColumn := getVectorColumnName(vectorSize) @@ -178,9 +204,9 @@ func RerankStandard( insEmbd := make([][]float64, 0, chunkCount) insTw := make([][]string, 0, chunkCount) - for index := range resp.Chunks { + for index := range chunks { // Extract vector - chunk := resp.Chunks[index] + chunk := chunks[index] chunkVector := extractVector(chunk, vectorColumn, zeroVector) insEmbd = append(insEmbd, chunkVector) @@ -210,16 +236,25 @@ func RerankStandard( } // Calculate hybrid similarity - return HybridSimilarity(questionVector, insEmbd, keywords, insTw, tkWeight, vtWeight, qb) + sim, tsim, vsim = HybridSimilarity(questionVector, insEmbd, keywords, insTw, tkWeight, vtWeight, qb) + + // Apply rank feature scores (tag_score * 10 + pagerank) + // Always apply pageranks, even when rankFeature is nil/empty + sim = applyRankFeatureScores(chunks, sim, rankFeature) + + logger.Info("RerankStandard completed") + return sim, tsim, vsim } // RerankInfinityFallback is used as a fallback when no reranker model is provided for Infinity engine. // Infinity can return scores in various field names (SCORE, score, SIMILARITY, etc.), // so we check multiple possible field names. If no score is found, we default to 1.0 // to ensure the chunk passes through any similarity threshold filters. -func RerankInfinityFallback(resp *engine.SearchResponse) (sim []float64, tsim []float64, vsim []float64) { - sim = make([]float64, len(resp.Chunks)) - for i, chunk := range resp.Chunks { +func RerankInfinityFallback(chunks []map[string]interface{}) (sim []float64, tsim []float64, vsim []float64) { + logger.Info("RerankInfinityFallback started", zap.Int("chunkCount", len(chunks))) + + sim = make([]float64, len(chunks)) + for i, chunk := range chunks { scoreFound := false scoreFields := []string{"SCORE", "score", "SIMILARITY", "similarity", "_score", "score()", "similarity()"} for _, field := range scoreFields { @@ -233,11 +268,11 @@ func RerankInfinityFallback(resp *engine.SearchResponse) (sim []float64, tsim [] sim[i] = 1.0 } } + logger.Info("RerankInfinityFallback completed") return sim, sim, sim } // HybridSimilarity calculates hybrid similarity between query and documents -// Reference: rag/nlp/query.py L174-L182 func HybridSimilarity( avec []float64, bvecs [][]float64, @@ -277,7 +312,6 @@ func HybridSimilarity( } // TokenSimilarity calculates token-based similarity -// Reference: rag/nlp/query.py L184-L199 func TokenSimilarity(atks []string, btkss [][]string, qb *QueryBuilder) []float64 { atksDict := tokensToDict(atks, qb) btkssDicts := make([]map[string]float64, len(btkss)) @@ -294,9 +328,11 @@ func TokenSimilarity(atks []string, btkss [][]string, qb *QueryBuilder) []float6 } // tokensToDict converts tokens to a weighted dictionary -// Reference: rag/nlp/query.py L185-L195 func tokensToDict(tks []string, qb *QueryBuilder) map[string]float64 { d := make(map[string]float64) + if qb == nil || qb.termWeight == nil { + return d + } wts := qb.termWeight.Weights(tks, false) for i, tw := range wts { @@ -314,7 +350,6 @@ func tokensToDict(tks []string, qb *QueryBuilder) map[string]float64 { } // tokenDictSimilarity calculates similarity between two token dictionaries -// Reference: rag/nlp/query.py L201-L213 func tokenDictSimilarity(qtwt, dtwt map[string]float64) float64 { if len(qtwt) == 0 || len(dtwt) == 0 { return 0.0 @@ -386,7 +421,10 @@ func extractContentTokens(fields map[string]interface{}, cfield string) []string return []string{} } - // Remove duplicates while preserving order + // Remove redundant spaces first to handle irregular spacing in Chinese text + v = RemoveRedundantSpaces(v) + + // Now split by whitespace to get individual tokens seen := make(map[string]bool) var result []string for _, t := range strings.Fields(v) { @@ -404,6 +442,8 @@ func extractTitleTokens(fields map[string]interface{}) []string { if !ok { return []string{} } + // Remove redundant spaces first + v = RemoveRedundantSpaces(v) var result []string for _, t := range strings.Fields(v) { if t != "" { @@ -473,12 +513,128 @@ func cosineSimilarity(a, b []float64) float64 { return dot / (math.Sqrt(normA) * math.Sqrt(normB)) } -// removeRedundantSpaces removes redundant spaces from text -func removeRedundantSpaces(s string) string { - return strings.Join(strings.Fields(s), " ") +// RemoveRedundantSpaces removes redundant spaces from text +// First pass: remove spaces after left-boundary characters +// Second pass: remove spaces before right-boundary characters +func RemoveRedundantSpaces(s string) string { + // First pass: remove spaces after left-boundary characters (opening brackets, etc.) + // e.g., "（ text" -> "（text", "【 text" -> "【text" + s = regexp.MustCompile(`([^\sa-z0-9.,\)>]) +([^\s])`).ReplaceAllString(s, "$1$2") + + // Second pass: remove spaces before right-boundary characters (closing brackets, punctuation) + // e.g., "text ！" -> "text！" + s = regexp.MustCompile(`([^\s]) +([^\sa-z0-9.,\(])`).ReplaceAllString(s, "$1$2") + + return s } // parseFloat parses a string to float64 func parseFloat(s string) (float64, error) { return strconv.ParseFloat(strings.TrimSpace(s), 64) } + +// applyRankFeatureScores applies rank feature scores to similarity +// Formula: tag_score * 10 + pagerank (per document) +func applyRankFeatureScores(chunks []map[string]interface{}, sim []float64, rankFeature map[string]float64) []float64 { + if len(chunks) == 0 || len(sim) == 0 { + return sim + } + + // Collect pageranks from each chunk + pageranks := make([]float64, len(chunks)) + for i, chunk := range chunks { + if pr, ok := chunk[common.PAGERANK_FLD]; ok { + if f, ok := toFloat64(pr); ok { + pageranks[i] = f + } + } + } + + // If no query rank features (no tag features), just add pageranks to sim + if len(rankFeature) == 0 { + for i := range sim { + sim[i] += pageranks[i] + } + return sim + } + + // Compute query denominator: sqrt(sum of squares of query rank feature weights, excluding pagerank) + qDenor := 0.0 + for t, s := range rankFeature { + if t != common.PAGERANK_FLD { + qDenor += s * s + } + } + qDenor = math.Sqrt(qDenor) + + // Compute tag score for each chunk + tagScores := make([]float64, len(chunks)) + for i, chunk := range chunks { + tagFeaStr, ok := chunk[common.TAG_FLD].(string) + if !ok || tagFeaStr == "" { + tagScores[i] = 0 + continue + } + + // Parse tag_feas JSON string: {"tag1": 0.5, "tag2": 0.3} + nor, denor := 0.0, 0.0 + tagFeaMap := parseTagFeasRerank(tagFeaStr) + for t, sc := range tagFeaMap { + if weight, exists := rankFeature[t]; exists { + nor += weight * sc + } + denor += sc * sc + } + if denor == 0 { + tagScores[i] = 0 + } else { + tagScores[i] = nor / math.Sqrt(denor) / qDenor + } + } + + // Final score: tag_score * 10 + pagerank + for i := range sim { + sim[i] += tagScores[i]*10 + pageranks[i] + } + + return sim +} + +// toFloat64 converts various numeric types to float64 +func toFloat64(v interface{}) (float64, bool) { + switch val := v.(type) { + case float64: + return val, true + case float32: + return float64(val), true + case int: + return float64(val), true + case int64: + return float64(val), true + case int32: + return float64(val), true + default: + return 0, false + } +} + +// parseTagFeasRerank parses a tag_feas JSON string into a map +// Format: {"tag1": 0.5, "tag2": 0.3} +func parseTagFeasRerank(tagFeasStr string) map[string]float64 { + result := make(map[string]float64) + if tagFeasStr == "" || tagFeasStr == "{}" { + return result + } + + // Parse JSON string + var m map[string]interface{} + if err := json.Unmarshal([]byte(tagFeasStr), &m); err != nil { + return result + } + for k, v := range m { + if f, ok := toFloat64(v); ok { + result[k] = f + } + } + return result +} diff --git a/internal/service/nlp/retrieval.go b/internal/service/nlp/retrieval.go new file mode 100644 index 0000000000..5f6bb8185f --- /dev/null +++ b/internal/service/nlp/retrieval.go @@ -0,0 +1,787 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package nlp + +import ( + "context" + "fmt" + "math" + "ragflow/internal/logger" + "sort" + "strings" + + "ragflow/internal/engine" + "ragflow/internal/engine/types" + "ragflow/internal/entity" + "ragflow/internal/tokenizer" + + "go.uber.org/zap" +) + +// RetrievalService provides retrieval search functionality +type RetrievalService struct { + docEngine engine.DocEngine +} + +// NewRetrievalService creates a new RetrievalService with the given doc engine +func NewRetrievalService(docEngine engine.DocEngine) *RetrievalService { + return &RetrievalService{docEngine: docEngine} +} + +// RetrievalRequest request for retrieval search +type RetrievalRequest struct { + Question string + TenantIDs []string + KbIDs []string + DocIDs []string + Page int + PageSize int + Top *int + SimilarityThreshold *float64 + VectorSimilarityWeight *float64 + RankFeature *map[string]float64 + RerankModel RerankModel + EmbeddingModel entity.EmbeddingModel + Aggs *bool + Highlight *bool +} + +// RetrievalResult result from retrieval search +type RetrievalResult struct { + Chunks []map[string]interface{} + DocAggs []map[string]interface{} // Aggregated document counts, sorted by count desc +} + +// Retrieval performs hybrid search + reranking + pagination +// - Calculate rerank limit and call Search() to fetch rerankLimit candidates for reranking +// - Perform reranking via Rerank() +// - Sort indices by score descending and filter by threshold +// - Calculate pagination to extract actual page returned from reranked results +// - Build chunks +// - Build document aggregation if specified +func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) (*RetrievalResult, error) { + if req.Question == "" { + return &RetrievalResult{Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}}, nil + } + + // Apply default values + if req.Top == nil { + req.Top = func() *int { v := 1024; return &v }() + } + if req.SimilarityThreshold == nil { + req.SimilarityThreshold = func() *float64 { v := 0.0; return &v }() + } + if req.VectorSimilarityWeight == nil { + req.VectorSimilarityWeight = func() *float64 { v := 0.3; return &v }() + } + if req.RankFeature == nil { + req.RankFeature = &map[string]float64{"pagerank_fea": 10.0} + } + if req.Aggs == nil { + req.Aggs = func() *bool { v := true; return &v }() + } + + if req.Page <= 0 { + req.Page = 1 + } + if req.PageSize <= 0 { + req.PageSize = 1 + } + + // Calculate rerank limit to ensure we get enough results for proper pagination + pageSize := req.PageSize + rerankLimit := pageSize + if pageSize > 1 { + rerankLimit = int(math.Ceil(64.0/float64(pageSize))) * pageSize + } else { + rerankLimit = 1 + } + if rerankLimit < 30 { + rerankLimit = 30 + } + // Cap rerank limit when external rerank model is used + if req.RerankModel != nil && *req.Top > 0 { + if rerankLimit > *req.Top { + rerankLimit = *req.Top + } + if rerankLimit > 64 { + rerankLimit = 64 + } + } + + page := req.Page + globalOffset := (page - 1) * pageSize + searchPage := globalOffset/rerankLimit + 1 + logger.Debug("Retrieval rerank params", zap.Int("page", req.Page), zap.Int("pageSize", pageSize), + zap.Int("searchPage", searchPage), zap.Int("rerankLimit", rerankLimit), zap.Int("globalOffset", globalOffset)) + + // Execute search via Search() + searchReq := &RetrievalSearchRequest{ + TenantIDs: req.TenantIDs, + Question: req.Question, + KbIDs: req.KbIDs, + DocIDs: req.DocIDs, + Page: searchPage, + PageSize: rerankLimit, + Top: *req.Top, + RankFeature: *req.RankFeature, + EmbeddingModel: req.EmbeddingModel, + } + searchResult, err := s.Search(ctx, searchReq) + if err != nil { + return nil, fmt.Errorf("Search failed: %w", err) + } + + // Perform reranking + vtWeight := *req.VectorSimilarityWeight + tkWeight := 1.0 - vtWeight + qb := GetQueryBuilder() + useInfinity := engine.GetEngineType() != engine.EngineElasticsearch + sim, term_similarity, vector_similarity := Rerank( + req.RerankModel, + searchResult.Chunks, + int(searchResult.Total), + nil, + searchResult.QueryVector, + req.Question, + tkWeight, + vtWeight, + useInfinity, + "content_ltks", + qb, + *req.RankFeature, + ) + if len(sim) == 0 { + return &RetrievalResult{Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}}, nil + } + + // Sort indices (positions into search results) by score descending + // After sorting by score descending, we process chunks in relevance order + type idxScore struct { + idx int + score float64 + } + idxScores := make([]idxScore, 0, len(sim)) + for i, s := range sim { + idxScores = append(idxScores, idxScore{idx: i, score: s}) + } + sort.Slice(idxScores, func(i, j int) bool { + return idxScores[i].score > idxScores[j].score + }) + + // When vector_similarity_weight is 0, similarity_threshold is not meaningful for term-only scores + // When doc_ids is explicitly provided (metadata or document filtering), bypass threshold + // User wants those specific documents regardless of their relevance score + postThreshold := *req.SimilarityThreshold + if *req.VectorSimilarityWeight <= 0 || len(req.DocIDs) > 0 { + postThreshold = 0.0 + } + + // Get valid indices where score >= postThreshold + validIdx := make([]int, 0) + for _, is := range idxScores { + if is.score >= postThreshold { + validIdx = append(validIdx, is.idx) + } + } + if len(validIdx) == 0 { + return &RetrievalResult{Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}}, nil + } + + // Calculate pagination + // begin and end define which of validIdx to return as the page + begin := globalOffset % rerankLimit + end := begin + pageSize + + // Get page indices + var pageIdx []int + if begin < len(validIdx) { + if end > len(validIdx) { + end = len(validIdx) + } + pageIdx = validIdx[begin:end] + } + logger.Debug("Pagination result info", zap.Int("totalValid", len(validIdx)), zap.Int("begin", begin), + zap.Int("end", end), zap.Int("chunkCount", len(pageIdx))) + + // Build chunks for pageIdx, transforms raw search results into the API response format + var filteredChunks []map[string]interface{} + dim := 0 + if searchResult.QueryVector != nil { + dim = len(searchResult.QueryVector) + } + zeroVector := make([]float64, dim) + for j := 0; j < dim; j++ { + zeroVector[j] = 0.0 + } + + for _, i := range pageIdx { + if i < 0 || i >= len(searchResult.IDs) { + continue + } + chunkID := searchResult.IDs[i] + chunk, exists := searchResult.Field[chunkID] + if !exists { + continue + } + + resultChunk := make(map[string]interface{}) + resultChunk["chunk_id"] = chunkID + if v, ok := chunk["content_ltks"]; ok { + resultChunk["content_ltks"] = v + } + if v, ok := chunk["content_with_weight"]; ok { + resultChunk["content_with_weight"] = v + } + if v, ok := chunk["doc_id"]; ok { + resultChunk["doc_id"] = v + } + if v, ok := chunk["docnm_kwd"]; ok { + resultChunk["docnm_kwd"] = v + } + if v, ok := chunk["kb_id"]; ok { + resultChunk["kb_id"] = v + } + if v, ok := chunk["important_kwd"]; ok { + resultChunk["important_kwd"] = v + } + if v, ok := chunk["tag_kwd"]; ok { + resultChunk["tag_kwd"] = v + } + if v, ok := chunk["img_id"]; ok { + resultChunk["image_id"] = v + } + if v, ok := chunk["position_int"]; ok { + resultChunk["positions"] = v + } + if v, ok := chunk["doc_type_kwd"]; ok { + resultChunk["doc_type_kwd"] = v + } + if v, ok := chunk["mom_id"]; ok { + resultChunk["mom_id"] = v + } + // row_id: row identifier (for structured data like tables) + if v, ok := chunk["row_id()"]; ok { + resultChunk["row_id"] = v + } + resultChunk["similarity"] = sim[i] + resultChunk["term_similarity"] = term_similarity[i] + resultChunk["vector_similarity"] = vector_similarity[i] + vectorColumn := fmt.Sprintf("q_%d_vec", dim) + if v, ok := chunk[vectorColumn]; ok { + resultChunk["vector"] = v + } else { + resultChunk["vector"] = zeroVector + } + + highlightEnabled := false + if req.Highlight != nil && *req.Highlight { + highlightEnabled = true + } + if highlightEnabled && searchResult.Highlight != nil { + if highlightText, ok := searchResult.Highlight[chunkID]; ok { + resultChunk["highlight"] = RemoveRedundantSpaces(highlightText) + } else if contentWithWeight, ok := chunk["content_with_weight"].(string); ok { + resultChunk["highlight"] = RemoveRedundantSpaces(contentWithWeight) + } + } + filteredChunks = append(filteredChunks, resultChunk) + } + + // Build document aggregation, aggregates document-level statistics across all valid chunks + // This is useful for showing users which documents are most relevant to their query. + var docAggs []map[string]interface{} + if req.Aggs != nil && *req.Aggs { + docAggsMap := make(map[string]struct { + docID string + count int + }) + for _, i := range validIdx { + if i < 0 || i >= len(searchResult.IDs) { + continue + } + chunkID := searchResult.IDs[i] + chunk, exists := searchResult.Field[chunkID] + if !exists { + continue + } + docName := "" + docID := "" + if v, ok := chunk["docnm_kwd"].(string); ok { + docName = v + } + if v, ok := chunk["doc_id"].(string); ok { + docID = v + } + if entry, exists := docAggsMap[docName]; exists { + entry.count++ + docAggsMap[docName] = entry + } else { + docAggsMap[docName] = struct { + docID string + count int + }{docID: docID, count: 1} + } + } + + // Sort by count descending + type docAggEntry struct { + docName string + docID string + count int + } + docAggsList := make([]docAggEntry, 0, len(docAggsMap)) + for docName, entry := range docAggsMap { + docAggsList = append(docAggsList, docAggEntry{docName: docName, docID: entry.docID, count: entry.count}) + } + sort.Slice(docAggsList, func(i, j int) bool { + return docAggsList[i].count > docAggsList[j].count + }) + + docAggs = make([]map[string]interface{}, 0, len(docAggsList)) + for _, entry := range docAggsList { + docAggs = append(docAggs, map[string]interface{}{ + "doc_name": entry.docName, + "doc_id": entry.docID, + "count": entry.count, + }) + } + } else { + docAggs = []map[string]interface{}{} + } + + return &RetrievalResult{ + Chunks: filteredChunks, + DocAggs: docAggs, + }, nil +} + +// RetrievalSearchRequest is the request struct for RetrievalService.Search() +type RetrievalSearchRequest struct { + Question string + TenantIDs []string + KbIDs []string + DocIDs []string + Top int + Page int + PageSize int + Sort bool + Highlight *bool + SimilarityThreshold float64 + RankFeature map[string]float64 + Filter map[string]interface{} + EmbeddingModel interface{} +} + +type RetrievalSearchResult struct { + Chunks []map[string]interface{} // Search results + Total int64 // Total number of matches + QueryVector []float64 // Query vector (for hybrid search, used in reranking) + Highlight map[string]string // Highlighted snippets (chunk_id -> highlighted text) + Field map[string]map[string]interface{} // ID -> chunk mapping + IDs []string // Ordered list of chunk IDs + Keywords []string // Keywords from query + Aggregation []map[string]interface{} // Doc aggregation by field + Options map[string]interface{} // Engine-specific options (e.g., total from get_total) +} + +// Search performs search based on question and EmbeddingModel: +// - Empty question: list data matching filters, optionally sorted +// - Non-empty question, no EmbeddingModel: fulltext search only +// - Non-empty question, with EmbeddingModel: hybrid search (fulltext + vector + fusion) +// +// Hybrid search path retries with lower thresholds if no results found. +func (s *RetrievalService) Search(ctx context.Context, req *RetrievalSearchRequest) (*RetrievalSearchResult, error) { + if req.Highlight == nil { + req.Highlight = func() *bool { v := false; return &v }() + } + filters := req.GetFilters() + pg := req.Page - 1 + if pg < 0 { + pg = 0 + } + topk := req.Top + if topk <= 0 { + topk = 1024 + } + pageSize := req.PageSize + if pageSize <= 0 { + pageSize = topk + } + limit := pageSize + + // Build Source field list + src := []string{ + "docnm_kwd", "content_ltks", "kb_id", "img_id", "title_tks", "important_kwd", "position_int", + "doc_id", "chunk_order_int", "page_num_int", "top_int", "create_timestamp_flt", "knowledge_graph_kwd", + "question_kwd", "question_tks", "doc_type_kwd", + "available_int", "content_with_weight", "mom_id", "pagerank_fea", "tag_feas", "row_id()", + } + + kwds := make(map[string]struct{}) + + // Build base engine request with common fields + // Note: RankFeature is NOT set here, it's set per-call where needed + searchRequest := &types.SearchRequest{ + IndexNames: buildIndexNames(req.TenantIDs), + KbIDs: req.KbIDs, + Offset: pg * pageSize, + Limit: limit, + Filter: filters, + SelectFields: src, + } + + // engineResult holds the result from docEngine.Search() (types.SearchResult) + // queryVector tracks the query vector for reranking + var engineResult *types.SearchResult + var queryVector []float64 + var err error + + if req.Question == "" { + // Empty question + if req.Sort { + searchRequest.OrderBy = &types.OrderByExpr{} + searchRequest.OrderBy.Asc("chunk_order_int").Asc("page_num_int").Asc("top_int").Desc("create_timestamp_flt") + } + searchRequest.MatchExprs = []interface{}{} + engineResult, err = s.docEngine.Search(ctx, searchRequest) + if err != nil { + return nil, fmt.Errorf("Search failed: %w", err) + } + } else { + // Non-empty question + + // Compute keywords via QueryBuilder + matchText, keywords := GetQueryBuilder().Question(req.Question, "", 0.3) + for _, k := range keywords { + kwds[k] = struct{}{} + } + + // Check if EmbeddingModel is available + if req.EmbeddingModel == nil { + // Keyword-only search + searchRequestWithRank := *searchRequest + searchRequestWithRank.MatchExprs = []interface{}{matchText} + searchRequestWithRank.RankFeature = req.RankFeature + + engineResult, err = s.docEngine.Search(ctx, &searchRequestWithRank) + if err != nil { + return nil, fmt.Errorf("Search failed: %w", err) + } + queryVector = nil + } else { + // Compute question vector via GetVector + similarityForGetVector := req.SimilarityThreshold + if similarityForGetVector <= 0 { + similarityForGetVector = 0.1 + } + matchDense, err := s.GetVector(req.Question, req.EmbeddingModel.(entity.EmbeddingModel), topk, similarityForGetVector) + if err != nil { + return nil, fmt.Errorf("GetVector failed: %w", err) + } + + // Execute search with fusion + fusionExpr := &types.FusionExpr{ + Method: "weighted_sum", + TopN: topk, + FusionParams: map[string]interface{}{"weights": "0.05,0.95"}, + } + + // Build source with vector column for ES + searchSrc := make([]string, len(searchRequest.SelectFields)) + copy(searchSrc, searchRequest.SelectFields) + if engine.GetEngineType() == engine.EngineElasticsearch { + searchSrc = append(searchSrc, matchDense.VectorColumnName) + } + + searchRequest.SelectFields = searchSrc + searchRequest.MatchExprs = []interface{}{matchText, matchDense, fusionExpr} + searchRequest.RankFeature = req.RankFeature + + engineResult, err = s.docEngine.Search(ctx, searchRequest) + if err != nil { + return nil, fmt.Errorf("Search failed: %w", err) + } + // If result is empty, retry with lower min_match + if engineResult.Total == 0 { + _, hasDocIDFilter := filters["doc_id"] + if hasDocIDFilter { + // Fallback without vector query when doc_id filter is present + searchRequest.SelectFields = src + searchRequest.MatchExprs = []interface{}{} + searchRequest.RankFeature = nil + + engineResult, err = s.docEngine.Search(ctx, searchRequest) + if err != nil { + return nil, fmt.Errorf("Search retry failed: %w", err) + } + } else { + // Retry with lower min_match via QueryBuilder + matchText, _ := GetQueryBuilder().Question(req.Question, "qa", 0.1) + matchDense.ExtraOptions["similarity"] = 0.17 + searchRequest.MatchExprs = []interface{}{matchText, matchDense, fusionExpr} + searchRequest.RankFeature = req.RankFeature + + engineResult, err = s.docEngine.Search(ctx, searchRequest) + if err != nil { + return nil, fmt.Errorf("Search retry failed: %w", err) + } + } + } + + queryVector = matchDense.EmbeddingData + } + + // Build kwds from keywords with fine-grained tokenization + for _, k := range keywords { + kwds[k] = struct{}{} + fgToken, _ := tokenizer.FineGrainedTokenize(k) + for _, kk := range strings.Fields(fgToken) { + if len(kk) < 2 { + continue + } + if _, ok := kwds[kk]; ok { + continue + } + kwds[kk] = struct{}{} + } + } + } + + searchResult := engineResult + ids := s.docEngine.GetDocIDs(searchResult.Chunks) + + // Build Keywords list from kwds set + keywordsList := make([]string, 0, len(kwds)) + for k := range kwds { + keywordsList = append(keywordsList, k) + } + + // Build Field map + fieldMap := s.docEngine.GetFields(searchResult.Chunks, nil) + + // Build Aggregation + aggregation := s.docEngine.GetAggregation(searchResult.Chunks, "docnm_kwd") + + // Build Highlight using GetHighlight + var highlight map[string]string + if len(keywordsList) > 0 { + highlight = s.docEngine.GetHighlight(searchResult.Chunks, keywordsList, "content_with_weight") + } + + return &RetrievalSearchResult{ + Chunks: searchResult.Chunks, + Total: searchResult.Total, + QueryVector: queryVector, + Highlight: highlight, + Field: fieldMap, + IDs: ids, + Keywords: keywordsList, + Aggregation: aggregation, + }, nil +} + +// GetVector computes query vector and returns MatchDenseExpr for hybrid search +func (s *RetrievalService) GetVector(txt string, embModel entity.EmbeddingModel, topk int, similarity float64) (*types.MatchDenseExpr, error) { + vector, err := embModel.EncodeQuery(txt) + if err != nil { + return nil, err + } + + vectorSize := len(vector) + vectorColumnName := fmt.Sprintf("q_%d_vec", vectorSize) + + return &types.MatchDenseExpr{ + VectorColumnName: vectorColumnName, + EmbeddingData: vector, + EmbeddingDataType: "float", + DistanceType: "cosine", + TopN: topk, + ExtraOptions: map[string]interface{}{"similarity": similarity}, + }, nil +} + +// GetFilters builds metadata filter map from RetrievalSearchRequest +func (r *RetrievalSearchRequest) GetFilters() map[string]interface{} { + filters := make(map[string]interface{}) + + if len(r.KbIDs) > 0 { + filters["kb_id"] = r.KbIDs + } + if len(r.DocIDs) > 0 { + filters["doc_id"] = r.DocIDs + } + for _, key := range []string{"knowledge_graph_kwd", "available_int", "entity_kwd", "from_entity_kwd", "to_entity_kwd", "removed_kwd"} { + if val, ok := r.Filter[key]; ok && val != nil { + filters[key] = val + } + } + for key, val := range r.Filter { + if _, exists := filters[key]; !exists && val != nil { + filters[key] = val + } + } + return filters +} + +// RetrievalByChildren aggregates child chunks into parent chunks +func RetrievalByChildren(chunks []map[string]interface{}, tenantIDs []string, docEngine engine.DocEngine, ctx context.Context) []map[string]interface{} { + logger.Info("RetrievalByChildren started", zap.Int("chunks", len(chunks)), zap.Strings("tenantIDs", tenantIDs)) + + indexNames := buildIndexNames(tenantIDs) + if len(chunks) == 0 || len(indexNames) == 0 { + return chunks + } + + // Group child chunks by mom_id + type childChunk struct { + chunk map[string]interface{} + kbID string + } + momChunks := make(map[string][]childChunk) + remainingChunks := make([]map[string]interface{}, 0, len(chunks)) + + for _, ck := range chunks { + momID, ok := ck["mom_id"].(string) + if !ok || momID == "" { + remainingChunks = append(remainingChunks, ck) + continue + } + kbID, _ := ck["kb_id"].(string) + momChunks[momID] = append(momChunks[momID], childChunk{chunk: ck, kbID: kbID}) + } + + if len(momChunks) == 0 { + logger.Info("RetrievalByChildren finished", zap.Int("momChunks", len(momChunks)), zap.Int("resultChunks", len(chunks))) + return chunks + } + + // Fetch parent chunks and aggregate + vectorSize := 1024 + for momID, childList := range momChunks { + kbIDs := make([]string, 0, len(childList)) + for _, c := range childList { + if c.kbID != "" { + kbIDs = append(kbIDs, c.kbID) + } + } + if len(kbIDs) == 0 { + kbIDs = append(kbIDs, "") + } + + parent, err := docEngine.GetChunk(ctx, indexNames[0], momID, kbIDs) + if err != nil { + logger.Warn("Failed to get parent chunk", zap.String("momID", momID), zap.Error(err)) + continue + } + parentMap, ok := parent.(map[string]interface{}) + if !ok { + continue + } + + // Calculate average similarity + var totalSim float64 + for _, c := range childList { + if sim, ok := c.chunk["similarity"].(float64); ok { + totalSim += sim + } + } + avgSim := totalSim / float64(len(childList)) + + // Collect content_ltks from children + var contentParts []string + for _, c := range childList { + if ltks, ok := c.chunk["content_ltks"].(string); ok { + contentParts = append(contentParts, ltks) + } + } + contentLTKS := strings.Join(contentParts, " ") + + // Collect important_kwd from children + allImportantKwd := []string{} + for _, c := range childList { + if kwd, ok := c.chunk["important_kwd"].([]interface{}); ok { + for _, k := range kwd { + if ks, ok := k.(string); ok { + allImportantKwd = append(allImportantKwd, ks) + } + } + } + } + + // Build aggregated chunk + docTypeKwd := parentMap["doc_type_kwd"] + if v, ok := docTypeKwd.(string); ok && v == "" { + docTypeKwd = []interface{}{} + } + aggregated := map[string]interface{}{ + "chunk_id": momID, + "content_ltks": contentLTKS, + "content_with_weight": parentMap["content_with_weight"], + "doc_id": parentMap["doc_id"], + "docnm_kwd": parentMap["docnm_kwd"], + "kb_id": parentMap["kb_id"], + "important_kwd": allImportantKwd, + "image_id": parentMap["img_id"], + "similarity": avgSim, + "vector_similarity": avgSim, + "term_similarity": avgSim, + "vector": make([]float64, vectorSize), + "positions": parentMap["position_int"], + "doc_type_kwd": docTypeKwd, + } + + // Get vector from first child if available + childVecLoop: + for _, c := range childList { + for k := range c.chunk { + if strings.HasSuffix(k, "_vec") { + if vec, ok := c.chunk[k].([]float64); ok { + aggregated["vector"] = vec + vectorSize = len(vec) + break childVecLoop + } + } + } + } + + remainingChunks = append(remainingChunks, aggregated) + } + + // Sort by similarity descending + for i := 0; i < len(remainingChunks); i++ { + for j := i + 1; j < len(remainingChunks); j++ { + simI, _ := remainingChunks[i]["similarity"].(float64) + simJ, _ := remainingChunks[j]["similarity"].(float64) + if simJ > simI { + remainingChunks[i], remainingChunks[j] = remainingChunks[j], remainingChunks[i] + } + } + } + + logger.Info("RetrievalByChildren finished", zap.Int("momChunks", len(momChunks)), zap.Int("resultChunks", len(remainingChunks))) + return remainingChunks +} + +// buildIndexNames creates index names for the given tenant IDs +func buildIndexNames(tenantIDs []string) []string { + indexNames := make([]string, len(tenantIDs)) + for i, tenantID := range tenantIDs { + indexNames[i] = fmt.Sprintf("ragflow_%s", tenantID) + } + return indexNames +} diff --git a/internal/service/search.go b/internal/service/search.go index cc2c0f38e5..901cebb423 100644 --- a/internal/service/search.go +++ b/internal/service/search.go @@ -330,3 +330,30 @@ func (s *SearchService) UpdateSearch(userID string, searchID string, req *Update return updatedSearch, nil } + +// GetDetail gets search details by ID including search_config +func (s *SearchService) GetDetail(searchID string) (map[string]interface{}, error) { + search, err := s.searchDAO.GetByID(searchID) + + if err != nil { + return nil, err + } + + result := map[string]interface{}{ + "id": search.ID, + "tenant_id": search.TenantID, + "name": search.Name, + "description": search.Description, + "created_by": search.CreatedBy, + "status": search.Status, + "create_time": search.CreateTime, + "update_time": search.UpdateTime, + "search_config": search.SearchConfig, + } + + if search.Avatar != nil { + result["avatar"] = *search.Avatar + } + + return result, nil +} diff --git a/internal/service/tag.go b/internal/service/tag.go new file mode 100644 index 0000000000..edb6a88e24 --- /dev/null +++ b/internal/service/tag.go @@ -0,0 +1,358 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package service + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "strings" + "time" + + "go.uber.org/zap" + + "ragflow/internal/cache" + "ragflow/internal/dao" + "ragflow/internal/engine/types" + "ragflow/internal/entity" + "ragflow/internal/logger" + "ragflow/internal/service/nlp" + + "github.com/cespare/xxhash/v2" +) + +// getTagsCacheKey generates a cache key from kb_ids using xxhash64 +func getTagsCacheKey(kbIDs []string) string { + // Normalize: unique + sorted so the key is set-stable regardless of caller order. + seen := make(map[string]struct{}, len(kbIDs)) + norm := make([]string, 0, len(kbIDs)) + for _, id := range kbIDs { + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + norm = append(norm, id) + } + sort.Strings(norm) + hasher := xxhash.New() + hasher.Write([]byte(strings.Join(norm, "\x00"))) + return fmt.Sprintf("%x", hasher.Sum64()) +} + +// GetTagsFromCache retrieves cached tags for given kb_ids +// Returns nil if not found (cache miss) +func GetTagsFromCache(kbIDs []string) (map[string]float64, error) { + if len(kbIDs) == 0 { + return nil, nil + } + + redisClient := cache.Get() + if redisClient == nil { + logger.Warn("Redis client not available, skipping cache lookup") + return nil, nil + } + + key := getTagsCacheKey(kbIDs) + data, err := redisClient.Get(key) + if err != nil || data == "" { + // Cache miss or error + return nil, nil + } + + var tags map[string]float64 + if err := json.Unmarshal([]byte(data), &tags); err != nil { + logger.Warn("Failed to unmarshal cached tags", zap.Error(err)) + return nil, nil + } + + return tags, nil +} + +// SetTagsToCache stores tags in cache for given kb_ids with 10 minute expiry +func SetTagsToCache(kbIDs []string, tags map[string]float64) error { + if len(kbIDs) == 0 || tags == nil { + return nil + } + + redisClient := cache.Get() + if redisClient == nil { + logger.Warn("Redis client not available, skipping cache store") + return nil + } + + key := getTagsCacheKey(kbIDs) + data, err := json.Marshal(tags) + if err != nil { + return fmt.Errorf("failed to marshal tags for cache: %w", err) + } + + // Cache for 10 minutes (600 seconds) + ok := redisClient.Set(key, string(data), 10*time.Minute) + if !ok { + logger.Warn("Failed to set tags cache") + return fmt.Errorf("failed to set tags cache") + } + + return nil +} + +// Knowledgebase type alias for entity.Knowledgebase +type Knowledgebase = entity.Knowledgebase + +// GetAllTagsInPortion returns the tag distribution for given KBs +func (s *MetadataService) GetAllTagsInPortion(tenantID string, kbIDs []string) (map[string]float64, error) { + if len(kbIDs) == 0 { + return make(map[string]float64), nil + } + + indexName := fmt.Sprintf("ragflow_%s", tenantID) + + // Search with large limit to get all tag_kwd values + searchReq := &types.SearchRequest{ + IndexNames: []string{indexName}, + KbIDs: kbIDs, + Offset: 0, + Limit: 10000, // Large limit to get all docs + } + + searchResp, err := s.docEngine.Search(context.Background(), searchReq) + if err != nil { + return nil, err + } + + // Use GetAggregation for tag counting + tagAgg := s.docEngine.GetAggregation(searchResp.Chunks, "tag_kwd") + if len(tagAgg) == 0 { + return make(map[string]float64), nil + } + + // Calculate total count for proportion calculation + total := 0 + for _, tc := range tagAgg { + total += tc["count"].(int) + } + if total == 0 { + return make(map[string]float64), nil + } + + // Calculate tag proportions: (count + 1) / (total + 1000) + S := 1000.0 + allTags := make(map[string]float64) + for _, tc := range tagAgg { + allTags[tc["key"].(string)] = float64(tc["count"].(int)+1) / (float64(total) + S) + } + + return allTags, nil +} + +// TagQuery returns weighted tag features for a question +func (s *MetadataService) TagQuery(question string, tenantIDs []string, kbIDs []string, allTags map[string]float64, topnTags int) (map[string]float64, error) { + if len(kbIDs) == 0 || len(allTags) == 0 || len(tenantIDs) == 0 { + return make(map[string]float64), nil + } + + // Build index names for all tenant IDs + indexNames := make([]string, len(tenantIDs)) + for i, tenantID := range tenantIDs { + indexNames[i] = fmt.Sprintf("ragflow_%s", tenantID) + } + + // Process question to get match text + queryBuilder := nlp.GetQueryBuilder() + matchTextExpr, warns := queryBuilder.Question(question, "qa", 0.0) // min_match=0.0 + if len(warns) > 0 { + logger.Warn("TagQuery: failed to build match text", zap.Any("warnings", warns)) + return make(map[string]float64), nil + } + matchText := matchTextExpr.MatchingText + + logger.Debug("TagQuery match_text", zap.String("match_text", matchText)) + + // Search with match text to get relevant docs + searchReq := &types.SearchRequest{ + IndexNames: indexNames, + KbIDs: kbIDs, + Offset: 0, + Limit: 1000, + MatchExprs: []interface{}{matchTextExpr}, + } + + searchResp, err := s.docEngine.Search(context.Background(), searchReq) + if err != nil { + return nil, err + } + + // Use GetAggregation for tag counting + aggs := s.docEngine.GetAggregation(searchResp.Chunks, "tag_kwd") + if len(aggs) == 0 { + return make(map[string]float64), nil + } + + // Calculate total count + cnt := 0 + for _, agg := range aggs { + cnt += agg["count"].(int) + } + if cnt == 0 { + return make(map[string]float64), nil + } + + // Calculate weighted tag features + // Formula: 0.1 * (c + 1) / (cnt + S) / max(1e-6, all_tags.get(a, 0.0001)) + S := 1000.0 + type tagScore struct { + tag string + score float64 + } + scoredTags := make([]tagScore, 0, len(aggs)) + + for _, agg := range aggs { + tag := agg["key"].(string) + c := agg["count"].(int) + allTagValue := allTags[tag] + if allTagValue <= 0 { + allTagValue = 0.0001 + } + score := 0.1 * float64(c+1) / (float64(cnt) + S) / max(1e-6, allTagValue) + scoredTags = append(scoredTags, tagScore{tag: tag, score: score}) + } + + // Sort by score descending + sort.Slice(scoredTags, func(i, j int) bool { + return scoredTags[i].score > scoredTags[j].score + }) + + // Take top N tags and normalize dot notation + resultTags := make(map[string]float64) + for i := 0; i < topnTags && i < len(scoredTags); i++ { + normalizedTag := strings.ReplaceAll(scoredTags[i].tag, ".", "_") + score := max(1.0, scoredTags[i].score) + if existing, ok := resultTags[normalizedTag]; !ok || score > existing { + resultTags[normalizedTag] = score + } + } + + return resultTags, nil +} + +// LabelQuestion returns rank features for a question based on KB's tag configuration. +// +// Flow: +// 1. Collect tag_kb_ids from KBs' parser_config +// 2. Try to get all_tags from cache (via GetTagsFromCache) +// 3. If cache miss, call GetAllTagsInPortion and cache the result (via SetTagsToCache) +// 4. Get tag KBs by IDs +// 5. Call TagQuery to get weighted tag features for the question +func (s *MetadataService) LabelQuestion(question string, kbs []*Knowledgebase) map[string]float64 { + if len(kbs) == 0 { + return nil + } + + // Collect tag_kb_ids from KBs' parser_config and track last KB + var tagKBIDs []string + var lastKB *Knowledgebase + for _, kb := range kbs { + if kb.ParserConfig == nil { + continue + } + lastKB = kb + if rawTagKBIDs, ok := kb.ParserConfig["tag_kb_ids"].([]interface{}); ok { + for _, id := range rawTagKBIDs { + if idStr, ok := id.(string); ok { + tagKBIDs = append(tagKBIDs, idStr) + } + } + } + } + + if len(tagKBIDs) == 0 { + return nil + } + + logger.Debug("tag_kb_ids found in parser_config", zap.Strings("tag_kb_ids", tagKBIDs)) + + // Get all tags from cache or compute and cache + allTags, err := GetTagsFromCache(tagKBIDs) + if err != nil { + logger.Warn("Failed to get tags from cache", zap.Error(err)) + } + if allTags == nil { + // Cache miss - compute all_tags_in_portion + allTags, err = s.GetAllTagsInPortion(lastKB.TenantID, tagKBIDs) + if err != nil { + logger.Warn("Failed to get all tags in portion", zap.Error(err)) + return nil + } + // Store in cache for future lookups + if err := SetTagsToCache(tagKBIDs, allTags); err != nil { + logger.Warn("Failed to set tags cache", zap.Error(err)) + } + } + + // Get tag_kbs by IDs + kbDAO := dao.NewKnowledgebaseDAO() + tagKBs, err := kbDAO.GetByIDs(tagKBIDs) + if err != nil || len(tagKBs) == 0 { + // Return nil if no tag_kbs found + return nil + } + + // Get unique tenant IDs from tag_kbs + tenantIDSet := make(map[string]bool) + for _, kb := range tagKBs { + tenantIDSet[kb.TenantID] = true + } + var uniqueTenantIDs []string + for tid := range tenantIDSet { + uniqueTenantIDs = append(uniqueTenantIDs, tid) + } + if len(uniqueTenantIDs) == 0 { + return nil + } + + // Get topn_tags from last KB's parser_config + // JSON-decoded numbers arrive as float64; also tolerate int/int64/json.Number for safety + topnTags := 3 + if lastKB != nil && lastKB.ParserConfig != nil { + switch v := lastKB.ParserConfig["topn_tags"].(type) { + case float64: + topnTags = int(v) + case int: + topnTags = v + case int64: + topnTags = int(v) + case json.Number: + if n, err := v.Int64(); err == nil { + topnTags = int(n) + } + } + } + + // Query tags for the question using unique tenant IDs + tagFeatures, err := s.TagQuery(question, uniqueTenantIDs, tagKBIDs, allTags, topnTags) + if err != nil { + return nil + } + if len(tagFeatures) == 0 { + // Tag kb exists but returned no matching tags - return empty map (not nil) + // so caller knows tag kb was configured vs not configured at all + return make(map[string]float64) + } + + return tagFeatures +} diff --git a/internal/tokenizer/tokenizer.go b/internal/tokenizer/tokenizer.go index d3dd867abd..8355f7b2e0 100644 --- a/internal/tokenizer/tokenizer.go +++ b/internal/tokenizer/tokenizer.go @@ -19,6 +19,7 @@ package tokenizer import ( "context" "fmt" + "ragflow/internal/engine" "runtime" "sync" "sync/atomic" @@ -408,7 +409,12 @@ func withAnalyzerResult[T any](fn func(*rag.Analyzer) (T, error)) (T, error) { // Tokenize tokenizes the text and returns a space-separated string of tokens // Example: "hello world" -> "hello world" +// +// NOTE: For Infinity engine, returns input unchanged to match python's behavior func Tokenize(text string) (string, error) { + if engine.GetEngineType() == "infinity" { + return text, nil + } return withAnalyzerResult(func(a *rag.Analyzer) (string, error) { return a.Tokenize(text) }) @@ -440,7 +446,12 @@ func SetFineGrained(fineGrained bool) { // FineGrainedTokenize performs fine-grained tokenization on space-separated tokens // Input: space-separated tokens (e.g., "hello world 测试") // Output: space-separated fine-grained tokens (e.g., "hello world 测 试") +// +// NOTE: For Infinity engine, returns input unchanged to match python's behavior func FineGrainedTokenize(tokens string) (string, error) { + if engine.GetEngineType() == "infinity" { + return tokens, nil + } return withAnalyzerResult(func(a *rag.Analyzer) (string, error) { return a.FineGrainedTokenize(tokens) }) diff --git a/internal/utility/convert.go b/internal/utility/convert.go index 5d88969d18..a13041a212 100644 --- a/internal/utility/convert.go +++ b/internal/utility/convert.go @@ -224,6 +224,26 @@ func IsEmpty(v interface{}) bool { return false } +// IsNumericValue checks if a value is numeric (int, uint, float, or numeric string) +func IsNumericValue(v interface{}) bool { + if v == nil { + return false + } + switch val := v.(type) { + case int, int8, int16, int32, int64: + return true + case uint, uint8, uint16, uint32, uint64: + return true + case float32, float64: + return true + case string: + _, err := strconv.ParseFloat(val, 64) + return err == nil + default: + return false + } +} + // SetFieldArray copies value to dest key, or sets empty array if value is empty func SetFieldArray(result map[string]interface{}, destKey string, v interface{}) { if IsEmpty(v) { @@ -321,4 +341,13 @@ func ConvertMapToJSONString(v interface{}) interface{} { return string(jsonBytes) } return v +} + +// FloatToString formats a float like Python's str() - adds ".0" if needed +func FloatToString(f float64) string { + s := strconv.FormatFloat(f, 'f', -1, 64) + if !strings.Contains(s, ".") && !strings.Contains(s, "e") { + s = s + ".0" + } + return s } \ No newline at end of file diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py index 6730261ea7..3a07e60067 100644 --- a/rag/llm/rerank_model.py +++ b/rag/llm/rerank_model.py @@ -297,7 +297,8 @@ class SILICONFLOWRerank(Base): "max_chunks_per_doc": 1024, "overlap_tokens": 80, } - response = requests.post(self.base_url, json=payload, headers=self.headers).json() + response_raw = requests.post(self.base_url, json=payload, headers=self.headers) + response = response_raw.json() rank = np.zeros(len(texts), dtype=float) try: for d in response["results"]: diff --git a/rag/nlp/search.py b/rag/nlp/search.py index 7ad19fe7c4..f37ce24572 100644 --- a/rag/nlp/search.py +++ b/rag/nlp/search.py @@ -343,7 +343,9 @@ class Dealer: def rerank_by_model(self, rerank_mdl, sres, query, tkweight=0.3, vtweight=0.7, cfield="content_ltks", rank_feature: dict | None = None): + print(f"[DEBUG rerank_by_model] query={query}, tkweight={tkweight}, vtweight={vtweight}") _, keywords = self.qryr.question(query) + print(f"[DEBUG rerank_by_model] keywords={keywords}") for i in sres.ids: if isinstance(sres.field[i].get("important_kwd", []), str): @@ -355,11 +357,29 @@ class Dealer: important_kwd = sres.field[i].get("important_kwd", []) tks = content_ltks + title_tks + important_kwd ins_tw.append(tks) + print(f"[DEBUG rerank_by_model] chunk id={i}, content_ltks={len(content_ltks)}, title_tks={len(title_tks)}, important_kwd={len(important_kwd)}") + doc_text = remove_redundant_spaces(" ".join(tks)) + if len(doc_text) > 100: + print(f"[DEBUG rerank_by_model] chunk id={i}, doc_text (first 100)={doc_text[:100]}...") + else: + print(f"[DEBUG rerank_by_model] chunk id={i}, doc_text={doc_text}") + + docs = [remove_redundant_spaces(" ".join(tks)) for tks in ins_tw] + print(f"[DEBUG rerank_by_model] docs sent to reranker: {len(docs)} docs") + for idx, doc in enumerate(docs[:2]): # Print first 2 + print(f"[DEBUG rerank_by_model] doc[{idx}] len={len(doc)}, full={doc}") + if len(doc) > 100: + print(f"[DEBUG rerank_by_model] doc[{idx}] (first 100)={doc[:100]}...") + else: + print(f"[DEBUG rerank_by_model] doc[{idx}]={doc}") tksim = self.qryr.token_similarity(keywords, ins_tw) - vtsim, _ = rerank_mdl.similarity(query, [remove_redundant_spaces(" ".join(tks)) for tks in ins_tw]) + print(f"[DEBUG rerank_by_model] tksim={tksim}") + vtsim, _ = rerank_mdl.similarity(query, docs) + print(f"[DEBUG rerank_by_model] vtsim from reranker={vtsim}") ## For rank feature(tag_fea) scores. rank_fea = self._rank_feature_scores(rank_feature, sres) + print(f"[DEBUG rerank_by_model] rank_fea={rank_fea}") return tkweight * np.array(tksim) + vtweight * vtsim + rank_fea, tksim, vtsim @@ -409,6 +429,7 @@ class Dealer: "similarity": similarity_threshold, "available_int": 1, } + logging.debug(f"[Search] global_offset={global_offset}, rerank_limit={RERANK_LIMIT}, page_size={page_size}, page={page}") if isinstance(tenant_ids, str): tenant_ids = tenant_ids.split(",")