Implement retrieval_test in GO (infiniflow#14231)

### What problem does this PR solve?

Implement retrieval_test in GO

### Type of change

- [x] Refactoring

Author: qinling0210

api/apps/chunk_app.py

@@ -157,6 +157,7 @@ async def _retrieval():
             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)

conf/models/siliconflow.json

@@ -0,0 +1,26 @@
+{
+  "name": "SILICONFLOW",
+  "tags": "LLM,TEXT EMBEDDING,TEXT RE-RANK,IMAGE2TEXT",
+  "url": {
+    "default": "https://api.siliconflow.cn/v1"
+  },
+  "url_suffix": {
+    "chat": "chat/completions",
+    "async_chat": "async/chat/completions",
+    "async_result": "async-result",
+    "embedding": "embedding",
+    "rerank": "rerank"
+  },
+  "models": [
+    {
+      "name": "BAAI/bge-reranker-v2-m3",
+      "max_tokens": 8192,
+      "model_types": [
+        "rerank"
+      ],
+      "features": {}
+    }
+  ]
+}
+
+

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/credentials v1.19.11
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.96.4
 	github.com/aws/smithy-go v1.24.2
+	github.com/cespare/xxhash/v2 v2.3.0
 	github.com/elastic/go-elasticsearch/v8 v8.19.1
 	github.com/gin-gonic/gin v1.9.1
 	github.com/google/uuid v1.6.0
@@ -43,7 +44,6 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.16 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sts v1.41.8 // indirect
 	github.com/bytedance/sonic v1.9.1 // indirect
-	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect
@@ -106,4 +106,4 @@ require (
 	gopkg.in/ini.v1 v1.67.0 // indirect
 )
 
-replace github.com/infiniflow/infinity-go-sdk => github.com/infiniflow/infinity/go v0.0.0-20260331112649-9bcd52a3d364
+replace github.com/infiniflow/infinity-go-sdk => github.com/infiniflow/infinity/go v0.0.0-20260424025959-72028e662929

go.sum

@@ -98,8 +98,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
-github.com/infiniflow/infinity/go v0.0.0-20260331112649-9bcd52a3d364 h1:0v5TjSirmCAUX3oaIV8Rd9d5B+kHPdymveETUU8OcC0=
-github.com/infiniflow/infinity/go v0.0.0-20260331112649-9bcd52a3d364/go.mod h1:hw3z5AwNFsGy1cdrE0Mfjot2y9jqVHTxBufUx9VzZ+0=
+github.com/infiniflow/infinity/go v0.0.0-20260424025959-72028e662929 h1:0M1BNouFVpnF12XEmF/42aR8CRU0bt/rMEVEsRUtSfQ=
+github.com/infiniflow/infinity/go v0.0.0-20260424025959-72028e662929/go.mod h1:hw3z5AwNFsGy1cdrE0Mfjot2y9jqVHTxBufUx9VzZ+0=
 github.com/iromli/go-itsdangerous v0.0.0-20220223194502-9c8bef8dac6a h1:Inib12UR9HAfBubrGNraPjKt/Cu8xPbTJbC50+0wP5U=
 github.com/iromli/go-itsdangerous v0.0.0-20220223194502-9c8bef8dac6a/go.mod h1:8N0Hlye5Lzw+H/yHWpZMkT0QLA+iOHG7KLdvAm95DZg=
 github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=

internal/cli/user_parser.go

@@ -1907,7 +1907,7 @@ func (p *Parser) parseInsertDatasetFromFile() (*Command, error) {
 }
 
 // Internal CLI for GO
-// parseInsertMetadataFromFile parses: INSERT INTO METADATA FROM FILE "file_path"
+// parseInsertMetadataFromFile parses: INSERT METADATA FROM FILE "file_path"
 func (p *Parser) parseInsertMetadataFromFile() (*Command, error) {
 	p.nextToken() // consume METADATA
 
@@ -2617,6 +2617,7 @@ func (p *Parser) parseUpdateCommand() (*Command, error) {
 	return nil, fmt.Errorf("unknown UPDATE target: %s", p.curToken.Value)
 }
 
+// Internal CLI for GO
 // parseUpdateChunk parses: UPDATE CHUNK 'chunk_id' OF DATASET 'dataset_name' SET '{"content": "..."}'
 func (p *Parser) parseUpdateChunk() (*Command, error) {
 	p.nextToken() // consume CHUNK

internal/common/constants.go

@@ -0,0 +1,8 @@
+package common
+
+const (
+	// PAGERANK_FLD is the field name for pagerank score
+	PAGERANK_FLD = "pagerank_fea"
+	// TAG_FLD is the field name for tag features
+	TAG_FLD = "tag_feas"
+)

internal/dao/tenant_llm.go

@@ -17,6 +17,7 @@
 package dao
 
 import (
+	"fmt"
 	"ragflow/internal/entity"
 )
 
@@ -28,6 +29,16 @@ func NewTenantLLMDAO() *TenantLLMDAO {
 	return &TenantLLMDAO{}
 }
 
+// GetByID get tenant LLM by primary key ID
+func (dao *TenantLLMDAO) GetByID(id int64) (*entity.TenantLLM, error) {
+	var tenantLLM entity.TenantLLM
+	err := DB.Where("id = ?", id).First(&tenantLLM).Error
+	if err != nil {
+		return nil, err
+	}
+	return &tenantLLM, nil
+}
+
 // GetByTenantAndModelName get tenant LLM by tenant ID and model name
 func (dao *TenantLLMDAO) GetByTenantAndModelName(tenantID, providerName string, modelName string) (*entity.TenantLLM, error) {
 	var tenantLLM entity.TenantLLM
@@ -38,6 +49,16 @@ func (dao *TenantLLMDAO) GetByTenantAndModelName(tenantID, providerName string,
 	return &tenantLLM, nil
 }
 
+// GetByTenantNameAndType get tenant LLM by tenant ID, model name, and model type
+func (dao *TenantLLMDAO) GetByTenantNameAndType(tenantID, modelName string, modelType entity.ModelType) (*entity.TenantLLM, error) {
+	var tenantLLM entity.TenantLLM
+	err := DB.Where("tenant_id = ? AND llm_name = ? AND model_type = ?", tenantID, modelName, modelType).First(&tenantLLM).Error
+	if err != nil {
+		return nil, err
+	}
+	return &tenantLLM, nil
+}
+
 // GetByTenantAndType get tenant LLM by tenant ID and model type
 func (dao *TenantLLMDAO) GetByTenantAndType(tenantID string, modelType entity.ModelType) (*entity.TenantLLM, error) {
 	var tenantLLM entity.TenantLLM
@@ -268,3 +289,50 @@ func (dao *TenantLLMDAO) GetByTenantIDLLMNameAndFactory(tenantID, llmName, facto
 	}
 	return &tenantLLM, nil
 }
+
+// LookupTenantLLMByID looks up a TenantLLM record by ID and returns the record plus composite model name.
+func LookupTenantLLMByID(tenantLLMDao *TenantLLMDAO, id int64) (*entity.TenantLLM, string, error) {
+	tenantLLM, err := tenantLLMDao.GetByID(id)
+	if err != nil {
+		return nil, "", fmt.Errorf("failed to get tenant_llm by id %d: %w", id, err)
+	}
+	if tenantLLM == nil || tenantLLM.LLMName == nil || *tenantLLM.LLMName == "" {
+		return nil, "", fmt.Errorf("tenant_llm record not found for id %d", id)
+	}
+	compositeName := fmt.Sprintf("%s@%s", *tenantLLM.LLMName, tenantLLM.LLMFactory)
+	return tenantLLM, compositeName, nil
+}
+
+// LookupTenantLLMByName looks up a TenantLLM record by tenant name and model type.
+func LookupTenantLLMByName(tenantLLMDao *TenantLLMDAO, tenantID, name string, modelType entity.ModelType) (*entity.TenantLLM, string, error) {
+	// Parse factory from name if present (e.g., "model@Factory")
+	modelName, factory := splitModelNameAndFactory(name)
+
+	// If factory is found, use factory-based lookup
+	if factory != "" {
+		return LookupTenantLLMByFactory(tenantLLMDao, tenantID, factory, modelName, modelType)
+	}
+
+	tenantLLM, err := tenantLLMDao.GetByTenantNameAndType(tenantID, modelName, modelType)
+	if err != nil {
+		return nil, "", fmt.Errorf("failed to get tenant_llm by name %s: %w", name, err)
+	}
+	if tenantLLM == nil || tenantLLM.LLMName == nil || *tenantLLM.LLMName == "" {
+		return nil, "", fmt.Errorf("tenant_llm record not found for name %s", name)
+	}
+	compositeName := fmt.Sprintf("%s@%s", *tenantLLM.LLMName, tenantLLM.LLMFactory)
+	return tenantLLM, compositeName, nil
+}
+
+// LookupTenantLLMByFactory looks up a TenantLLM record by tenant, factory, and model name.
+func LookupTenantLLMByFactory(tenantLLMDao *TenantLLMDAO, tenantID, factory, name string, modelType entity.ModelType) (*entity.TenantLLM, string, error) {
+	tenantLLM, err := tenantLLMDao.GetByTenantFactoryAndModelName(tenantID, factory, name)
+	if err != nil {
+		return nil, "", fmt.Errorf("failed to get tenant_llm by factory %s and name %s: %w", factory, name, err)
+	}
+	if tenantLLM == nil || tenantLLM.LLMName == nil || *tenantLLM.LLMName == "" {
+		return nil, "", fmt.Errorf("tenant_llm record not found for factory %s and name %s", factory, name)
+	}
+	compositeName := fmt.Sprintf("%s@%s", *tenantLLM.LLMName, tenantLLM.LLMFactory)
+	return tenantLLM, compositeName, nil
+}

internal/engine/elasticsearch/get.go

@@ -19,38 +19,31 @@ package elasticsearch
 import (
 	"context"
 	"fmt"
+
+	"ragflow/internal/engine/types"
 )
 
 // GetChunk gets a chunk by ID
 func (e *elasticsearchEngine) GetChunk(ctx context.Context, indexName, chunkID string, kbIDs []string) (interface{}, error) {
-	// Build query to get the chunk by ID
-	query := map[string]interface{}{
-		"term": map[string]interface{}{
+	// Build unified search request to get the chunk by ID
+	searchReq := &types.SearchRequest{
+		IndexNames: []string{indexName},
+		Limit:      1,
+		Offset:     0,
+		Filter: map[string]interface{}{
 			"id": chunkID,
 		},
 	}
 
-	searchReq := &SearchRequest{
-		IndexNames: []string{indexName},
-		Query:     query,
-		Size:      1,
-		From:      0,
-	}
-
 	// Execute search
-	result, err := e.Search(ctx, searchReq)
+	searchResp, err := e.Search(ctx, searchReq)
 	if err != nil {
 		return nil, fmt.Errorf("failed to search: %w", err)
 	}
 
-	esResp, ok := result.(*SearchResponse)
-	if !ok {
-		return nil, fmt.Errorf("invalid search response type")
-	}
-
-	if len(esResp.Hits.Hits) == 0 {
+	if len(searchResp.Chunks) == 0 {
 		return nil, nil
 	}
 
-	return esResp.Hits.Hits[0].Source, nil
-}
+	return searchResp.Chunks[0], nil
+}