zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-15 20:27:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

Implement retrieval_test in GO (#14231)

Browse Source 
### What problem does this PR solve?

Implement retrieval_test in GO

### Type of change

- [x] Refactoring
This commit is contained in:
qinling0210
2026-04-24 15:30:14 +08:00
committed by GitHub
parent aadd9a333f
commit 1473000135
42 changed files with 4822 additions and 1609 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
api/apps/chunk_app.py
View File

@@ -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)

26
conf/models/siliconflow.json Normal file
View File

@@ -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": {}
}
]
}

4
go.mod
View File

@@ -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

4
go.sum
View File

@@ -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=

3
internal/cli/user_parser.go
View File

@@ -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

8
internal/common/constants.go Normal file
View File

@@ -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"
)

68
internal/dao/tenant_llm.go
View File

@@ -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
}

31
internal/engine/elasticsearch/get.go
View File

@@ -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
}

278
internal/engine/elasticsearch/search.go
View File

@@ -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
}

16
internal/engine/engine.go
View File

@@ -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

9
internal/engine/global.go
View File

@@ -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

59
internal/engine/infinity/common.go
View File

@@ -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
}

2
internal/engine/infinity/dataset.go
View File

@@ -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" {

200
internal/engine/infinity/get.go
View File

@@ -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)
}

1578
internal/engine/infinity/search.go
View File

File diff suppressed because it is too large Load Diff

99
internal/engine/types/types.go
View File

@@ -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"})
}

1
internal/entity/kb.go
View File

@@ -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"`

5
internal/entity/models/deepseek.go
View File

@@ -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

5
internal/entity/models/dummy.go
View File

@@ -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")

5
internal/entity/models/minimax.go
View File

@@ -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())

5
internal/entity/models/moonshot.go
View File

@@ -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")

8
internal/entity/models/types.go
View File

@@ -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

100
internal/entity/models/zhipu-ai.go
View File

@@ -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"

7
internal/entity/types.go
View File

@@ -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
}

5
internal/logger/logger.go
View File

@@ -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()

743
internal/service/chunk.go
View File

@@ -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 "<nil>"
}
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.

167
internal/service/generator.go Normal file
View File

@@ -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
}

160
internal/service/load_prompt.go Normal file
View File

@@ -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(`<think>[\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
}

223
internal/service/metadata.go
View File

@@ -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

563
internal/service/metadata_filter.go Normal file
View File

@@ -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
}

193
internal/service/model_service.go
View File

@@ -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
}

60
internal/service/models/factory.go
View File

@@ -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
}

258
internal/service/models/siliconflow_model.go
View File

@@ -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,
}
})
}

45
internal/service/nlp/query_builder.go
View File

@@ -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,

254
internal/service/nlp/reranker.go
View File

@@ -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
}

787
internal/service/nlp/retrieval.go Normal file
View File

@@ -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
}

27
internal/service/search.go
View File

@@ -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
}

358
internal/service/tag.go Normal file
View File

@@ -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
}

11
internal/tokenizer/tokenizer.go
View File

@@ -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)
})

29
internal/utility/convert.go
View File

@@ -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
}

3
rag/llm/rerank_model.py
View File

@@ -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"]:

23
rag/nlp/search.py
View File

@@ -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(",")