From f53518c11008af5fa62783c15974b009b6dc7288 Mon Sep 17 00:00:00 2001 From: Jin Hai Date: Mon, 27 Jul 2026 10:20:16 +0800 Subject: [PATCH] Go: add context, part5 (#17392) Signed-off-by: Jin Hai --- cmd/ragflow_server.go | 10 +- internal/admin/service.go | 2 +- internal/agent/canvas/node_body.go | 3 +- .../canvas/node_body_override_params_test.go | 4 +- ...body_per_class_timeout_integration_test.go | 4 +- .../agent/canvas/node_body_timeout_test.go | 6 +- internal/agent/component/agent.go | 34 +++-- internal/agent/component/agent_test.go | 70 ++++----- internal/agent/component/base.go | 6 +- internal/agent/component/begin.go | 8 +- internal/agent/component/begin_test.go | 10 +- internal/agent/component/browser.go | 6 +- internal/agent/component/browser_test.go | 12 +- .../component/canvas_var_resolution_test.go | 6 +- internal/agent/component/categorize.go | 9 +- internal/agent/component/categorize_test.go | 36 ++--- .../component/chat_template_kwargs_test.go | 2 +- internal/agent/component/citation_test.go | 4 +- internal/agent/component/data_operations.go | 8 +- .../agent/component/data_operations_test.go | 14 +- internal/agent/component/docs_generator.go | 7 +- .../agent/component/docs_generator_test.go | 10 +- .../component/downstream_message_test.go | 2 +- internal/agent/component/excel_processor.go | 7 +- .../agent/component/excel_processor_test.go | 12 +- internal/agent/component/fillup.go | 8 +- internal/agent/component/fillup_test.go | 8 +- internal/agent/component/fixture_stubs.go | 48 +++--- internal/agent/component/grounding_test.go | 15 +- .../agent/component/history_window_test.go | 6 +- internal/agent/component/invoke.go | 7 +- internal/agent/component/invoke_test.go | 22 +-- internal/agent/component/list_operations.go | 8 +- .../agent/component/list_operations_test.go | 32 ++-- internal/agent/component/llm.go | 15 +- internal/agent/component/llm_retry.go | 6 +- internal/agent/component/llm_retry_test.go | 28 ++-- internal/agent/component/llm_test.go | 17 ++- internal/agent/component/loop.go | 8 +- internal/agent/component/loop_test.go | 4 +- internal/agent/component/message.go | 8 +- .../agent/component/message_phase8b_test.go | 22 +-- internal/agent/component/message_test.go | 30 ++-- internal/agent/component/parallel.go | 8 +- internal/agent/component/parallel_test.go | 4 +- .../component/production_chain_fixes_test.go | 32 ++-- .../agent/component/retrieval_swap_test.go | 4 +- .../agent/component/sampler_params_test.go | 10 +- .../stagehand_runtime_integration_test.go | 2 +- internal/agent/component/streaming_test.go | 8 +- internal/agent/component/string_transform.go | 8 +- .../agent/component/string_transform_test.go | 12 +- .../agent/component/structured_output_test.go | 10 +- internal/agent/component/switch.go | 8 +- internal/agent/component/switch_test.go | 28 ++-- .../agent/component/tool_call_memory_test.go | 6 +- internal/agent/component/tool_component.go | 6 +- .../agent/component/tool_component_test.go | 18 +-- .../agent/component/tool_dispatch_test.go | 2 +- .../agent/component/universe_a_wrappers.go | 28 ++-- internal/agent/component/userfillup.go | 8 +- internal/agent/component/userfillup_test.go | 8 +- .../agent/component/variable_aggregator.go | 8 +- .../component/variable_aggregator_test.go | 8 +- internal/agent/component/variable_assigner.go | 8 +- .../agent/component/variable_assigner_test.go | 16 +- internal/agent/component/vision_test.go | 6 +- internal/agent/runtime/component.go | 4 +- internal/agent/runtime/registry_test.go | 4 +- internal/agent/tool/code_exec.go | 4 +- internal/agent/tool/exesql.go | 7 +- internal/agent/tool/mcp.go | 4 +- internal/agent/tool/registry_test.go | 2 +- internal/agent/tool/retrieval.go | 3 +- internal/agent/tool/retrieval_nlp.go | 12 +- internal/agent/tool/retrieval_nlp_test.go | 12 +- internal/agent/tool/retrieval_service.go | 12 +- internal/agent/tool/retrieval_service_test.go | 13 +- internal/agent/tool/retrieval_test.go | 6 +- internal/agent/tool/retrieval_wiring_test.go | 14 +- internal/dao/kb.go | 143 +++++++++--------- internal/dao/kb_test.go | 11 +- internal/engine/elasticsearch/metadata.go | 6 +- internal/engine/engine.go | 4 +- internal/engine/infinity/metadata.go | 11 +- internal/handler/agent_component.go | 2 +- internal/handler/agent_component_test.go | 3 +- internal/handler/dataset.go | 64 +++++--- internal/handler/dataset_search_test.go | 5 +- internal/handler/dify_retrieval_handler.go | 30 ++-- .../handler/dify_retrieval_handler_test.go | 52 +++---- internal/handler/document.go | 47 +++--- internal/handler/file_commit.go | 2 +- internal/handler/mcp_server.go | 25 +-- internal/handler/search.go | 4 +- internal/handler/tenant.go | 10 +- internal/ingestion/component/chunker/group.go | 3 +- .../ingestion/component/chunker/group_test.go | 13 +- .../ingestion/component/chunker/hierarchy.go | 3 +- .../component/chunker/hierarchy_test.go | 25 ++- .../component/chunker/image_upload_test.go | 2 +- internal/ingestion/component/chunker/one.go | 4 +- .../ingestion/component/chunker/one_test.go | 2 +- .../component/chunker/pdfcrop_cgo.go | 6 +- .../component/chunker/pdfcrop_nocgo.go | 2 +- .../component/chunker/presentation.go | 4 +- .../component/chunker/presentation_test.go | 2 +- .../component/chunker/python_parity_test.go | 2 +- internal/ingestion/component/chunker/qa.go | 3 +- .../component/chunker/qa_batch2_test.go | 2 +- .../ingestion/component/chunker/qa_test.go | 22 +-- .../ingestion/component/chunker/register.go | 6 +- internal/ingestion/component/chunker/table.go | 4 +- .../ingestion/component/chunker/table_test.go | 2 +- internal/ingestion/component/chunker/tag.go | 4 +- .../ingestion/component/chunker/tag_test.go | 2 +- internal/ingestion/component/chunker/title.go | 4 +- .../ingestion/component/chunker/title_test.go | 10 +- internal/ingestion/component/chunker/token.go | 10 +- .../ingestion/component/chunker/token_test.go | 16 +- .../ingestion/component/document_storage.go | 10 +- internal/ingestion/component/extractor.go | 3 +- .../ingestion/component/extractor_tag_test.go | 6 +- .../ingestion/component/extractor_test.go | 27 ++-- internal/ingestion/component/file.go | 4 +- internal/ingestion/component/file_test.go | 13 +- internal/ingestion/component/parser.go | 13 +- .../parser_dispatch_pdf_vision_cgo_test.go | 2 +- .../component/parser_dispatch_test.go | 30 ++-- internal/ingestion/component/parser_test.go | 18 +-- .../component/slice3_integration_test.go | 4 +- internal/ingestion/component/tokenizer.go | 7 +- .../ingestion/component/tokenizer_test.go | 42 ++--- .../component/tokenizer_unit_test.go | 16 +- internal/ingestion/pipeline/pipeline_test.go | 11 +- .../pipeline/template_integration_test.go | 2 +- internal/ingestion/task/embedder.go | 14 +- internal/ingestion/task/embedder_test.go | 14 +- .../ingestion/task/pipeline_executor_test.go | 2 +- internal/ingestion/task/task_context.go | 2 +- internal/mcp/connector.go | 25 +-- internal/mcp/server.go | 39 ++--- internal/service/agent_run_e2e_test.go | 3 +- internal/service/chat.go | 36 ++--- internal/service/chat_pipeline.go | 8 +- internal/service/chunk/chunk.go | 44 +++--- internal/service/chunk/chunk_test.go | 8 +- internal/service/chunk_types.go | 14 +- internal/service/dataset/create_test.go | 21 ++- internal/service/dataset/crud.go | 26 ++-- .../service/dataset/fake_doc_engine_test.go | 4 +- internal/service/dataset/index.go | 32 ++-- internal/service/dataset/index_delete_test.go | 15 +- internal/service/dataset/ingestion.go | 12 +- internal/service/dataset/metadata.go | 16 +- internal/service/dataset/permission.go | 18 ++- internal/service/dataset/search.go | 16 +- internal/service/dataset/tags.go | 18 +-- .../service/dataset/tags_aggregate_test.go | 18 ++- internal/service/dataset/tags_list_test.go | 15 +- internal/service/dataset/tags_rename_test.go | 12 +- internal/service/dataset/update.go | 2 +- internal/service/dataset/update_test.go | 30 ++-- internal/service/document/document_crud.go | 6 +- .../document/document_dataset_update.go | 4 +- internal/service/document/document_ingest.go | 4 +- .../service/document/document_metadata.go | 14 +- internal/service/document/document_test.go | 22 +-- internal/service/document/file2document.go | 4 +- internal/service/file_permission.go | 2 +- internal/service/memory_message_test.go | 2 +- internal/service/metadata.go | 18 +-- internal/service/metadata_filter.go | 3 +- internal/service/model_chat.go | 4 +- internal/service/model_service.go | 4 +- internal/service/openai_chat.go | 2 +- internal/service/search.go | 5 +- internal/service/skill_indexer.go | 6 +- internal/service/skill_search.go | 2 +- internal/service/tag.go | 18 +-- internal/service/tenant.go | 12 +- internal/service/test_helpers_test.go | 2 +- 182 files changed, 1249 insertions(+), 1066 deletions(-) diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index dc8e020f1d..2cb42e7cdf 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -742,14 +742,14 @@ func startServer(ctx context.Context, config *server.Config) { // (ragflow_retrieval, ragflow_list_datasets, ragflow_list_chats) to // external AI clients via JSON-RPC over HTTP. mcpServerHandler := handler.NewMCPServerHandler( - func(userID string, page, pageSize int, orderBy string, desc bool) ([]map[string]interface{}, int64, error) { - return handler.MCPListDatasets(datasetsService, userID, page, pageSize, orderBy, desc) + func(ctx context.Context, userID string, page, pageSize int, orderBy string, desc bool) ([]map[string]interface{}, int64, error) { + return handler.MCPListDatasets(ctx, datasetsService, userID, page, pageSize, orderBy, desc) }, - func(userID string, page, pageSize int, orderBy string, desc bool) ([]map[string]interface{}, int64, error) { + func(ctx context.Context, userID string, page, pageSize int, orderBy string, desc bool) ([]map[string]interface{}, int64, error) { return handler.MCPListChats(ctx, chatService, userID, page, pageSize, orderBy, desc) }, - func(userID string, req mcp.RetrievalRequest) (string, error) { - return handler.MCPRetrieval(datasetsService, userID, req) + func(ctx context.Context, userID string, req mcp.RetrievalRequest) (string, error) { + return handler.MCPRetrieval(ctx, datasetsService, userID, req) }, ) skillSearchHandler := handler.NewSkillSearchHandler(docEngine, documentService) diff --git a/internal/admin/service.go b/internal/admin/service.go index a3d968e461..f06c18ba69 100644 --- a/internal/admin/service.go +++ b/internal/admin/service.go @@ -552,7 +552,7 @@ func (s *Service) DeleteUser(ctx context.Context, username string) (*DeleteUserR // Delete owned tenant data if ownedTenantID != "" { // 1. Get knowledge base IDs - kbIDs, err := s.kbDAO.GetKBIDsByTenantIDSimple(ownedTenantID) + kbIDs, err := s.kbDAO.GetKBIDsByTenantIDSimple(ctx, tx, ownedTenantID) if err != nil { common.Warn("failed to get knowledge base IDs", zap.Error(err)) } diff --git a/internal/agent/canvas/node_body.go b/internal/agent/canvas/node_body.go index 8e78ec3135..d01ec37b9f 100644 --- a/internal/agent/canvas/node_body.go +++ b/internal/agent/canvas/node_body.go @@ -34,6 +34,7 @@ import ( "errors" "fmt" "ragflow/internal/common" + "ragflow/internal/dao" "strconv" "strings" "time" @@ -255,7 +256,7 @@ func realComponentBodyWithOptions(cpnID, componentClass string, comp runtime.Com invokeErr := runtime.TrackProgress(cpnID, runtime.ProgressCallbackFromContext(ctx), func() error { var e error out, e = runtime.TrackElapsed(componentClass, func() (map[string]any, error) { - return comp.Invoke(runtime.WithComponentExecutionOptions(cctx, opts), in) + return comp.Invoke(runtime.WithComponentExecutionOptions(cctx, opts), dao.DB, in) }) return e }) diff --git a/internal/agent/canvas/node_body_override_params_test.go b/internal/agent/canvas/node_body_override_params_test.go index 4462437fb3..3c27b92918 100644 --- a/internal/agent/canvas/node_body_override_params_test.go +++ b/internal/agent/canvas/node_body_override_params_test.go @@ -5,13 +5,15 @@ import ( "testing" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) type stubComponent struct { params map[string]any } -func (s *stubComponent) Invoke(_ context.Context, in map[string]any) (map[string]any, error) { +func (s *stubComponent) Invoke(_ context.Context, _ *gorm.DB, in map[string]any) (map[string]any, error) { return in, nil } diff --git a/internal/agent/canvas/node_body_per_class_timeout_integration_test.go b/internal/agent/canvas/node_body_per_class_timeout_integration_test.go index 7e2c48045a..1479785096 100644 --- a/internal/agent/canvas/node_body_per_class_timeout_integration_test.go +++ b/internal/agent/canvas/node_body_per_class_timeout_integration_test.go @@ -35,6 +35,8 @@ import ( "time" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) // captureCtxComponent is a runtime.Component that records the @@ -51,7 +53,7 @@ type captureCtxComponent struct { } func (c *captureCtxComponent) Name() string { return c.name } -func (c *captureCtxComponent) Invoke(ctx context.Context, _ map[string]any) (map[string]any, error) { +func (c *captureCtxComponent) Invoke(ctx context.Context, _ *gorm.DB, _ map[string]any) (map[string]any, error) { if dl, ok := ctx.Deadline(); ok { c.deadline.Store(dl.UnixNano()) c.timeoutOK.Store(true) diff --git a/internal/agent/canvas/node_body_timeout_test.go b/internal/agent/canvas/node_body_timeout_test.go index 424ad06893..34a74db353 100644 --- a/internal/agent/canvas/node_body_timeout_test.go +++ b/internal/agent/canvas/node_body_timeout_test.go @@ -23,6 +23,8 @@ import ( "time" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) // blockingComponent is a runtime.Component whose Invoke blocks until ctx @@ -32,7 +34,7 @@ type blockingComponent struct{} func (b *blockingComponent) Name() string { return "blocking" } -func (b *blockingComponent) Invoke(ctx context.Context, _ map[string]any) (map[string]any, error) { +func (b *blockingComponent) Invoke(ctx context.Context, _ *gorm.DB, _ map[string]any) (map[string]any, error) { <-ctx.Done() return nil, ctx.Err() } @@ -160,7 +162,7 @@ type echoComponent struct{} func (e *echoComponent) Name() string { return "echo" } -func (e *echoComponent) Invoke(_ context.Context, in map[string]any) (map[string]any, error) { +func (e *echoComponent) Invoke(_ context.Context, _ *gorm.DB, in map[string]any) (map[string]any, error) { out := make(map[string]any, len(in)) for k, v := range in { out[k] = v diff --git a/internal/agent/component/agent.go b/internal/agent/component/agent.go index e9844d9c12..37b9e7ec3f 100644 --- a/internal/agent/component/agent.go +++ b/internal/agent/component/agent.go @@ -17,6 +17,7 @@ import ( "errors" "fmt" "io" + "ragflow/internal/dao" "sort" "strings" "time" @@ -25,6 +26,7 @@ import ( "github.com/cloudwego/eino/compose" "github.com/cloudwego/eino/flow/agent/react" "github.com/cloudwego/eino/schema" + "gorm.io/gorm" "ragflow/internal/agent/component/prompts" "ragflow/internal/agent/runtime" @@ -304,7 +306,7 @@ func emitAgentModelStreams(ctx context.Context, future react.MessageFuture) <-ch // // When the LLM call fails or there are no tool calls, the function // returns ("", nil) and the caller skips appending to history. -func addToolCallMemory(ctx context.Context, p AgentParam, msg *schema.Message) (string, error) { +func addToolCallMemory(ctx context.Context, db *gorm.DB, p AgentParam, msg *schema.Message) (string, error) { calls := extractToolCalls(msg) if len(calls) == 0 { return "", nil @@ -320,7 +322,7 @@ func addToolCallMemory(ctx context.Context, p AgentParam, msg *schema.Message) ( system := "You are a memory summarizer. Given a list of tool calls the assistant just made, output ONE short sentence (max 30 words) describing what the assistant did, suitable for a future-turn conversation history. Output ONLY the sentence, no preamble, no quotes." user := "Tool calls: " + callsDesc.String() inv := getDefaultChatInvoker() - resp, err := inv.Invoke(ctx, ChatInvokeRequest{ + resp, err := inv.Invoke(ctx, db, ChatInvokeRequest{ Driver: p.Driver, ModelName: p.ModelID, APIKey: p.APIKey, @@ -347,7 +349,7 @@ func addToolCallMemory(ctx context.Context, p AgentParam, msg *schema.Message) ( // Returns the grounded content on success, the original content // unchanged when no chunks are available or the call fails. Mirrors // Python's `cite_letter` / `generate_with_citation` flow. -func applyCitationGrounding(ctx context.Context, p AgentParam, content string, chunks []prompts.CitationSource) (string, error) { +func applyCitationGrounding(ctx context.Context, db *gorm.DB, p AgentParam, content string, chunks []prompts.CitationSource) (string, error) { if !p.Cite { return content, nil } @@ -359,7 +361,7 @@ func applyCitationGrounding(ctx context.Context, p AgentParam, content string, c } systemPrompt, _ := prompts.CitationPlusPrompt(chunks) inv := getDefaultChatInvoker() - resp, err := inv.Invoke(ctx, ChatInvokeRequest{ + resp, err := inv.Invoke(ctx, db, ChatInvokeRequest{ Driver: p.Driver, ModelName: p.ModelID, APIKey: p.APIKey, @@ -499,7 +501,7 @@ func (c *AgentComponent) Reset() { // - the rephrase LLM call fails // // Window defaults to AgentParam.OptimizeHistoryWindow (3) when zero. -func optimizeMultiTurnQuestion(ctx context.Context, p AgentParam, history []map[string]any) (string, error) { +func optimizeMultiTurnQuestion(ctx context.Context, db *gorm.DB, p AgentParam, history []map[string]any) (string, error) { window := p.OptimizeHistoryWindow if window <= 0 { window = 3 @@ -527,7 +529,7 @@ func optimizeMultiTurnQuestion(ctx context.Context, p AgentParam, history []map[ system := "You are a question rephraser. Given conversation history and the user's latest input, rewrite the latest input as a self-contained question that does not require the history to understand. Output ONLY the rephrased question, no preamble, no quotes." user := "Conversation history:\n" + histBuf.String() + "\n\nUser's latest input:\n" + p.UserPrompt inv := getDefaultChatInvoker() - resp, err := inv.Invoke(ctx, ChatInvokeRequest{ + resp, err := inv.Invoke(ctx, db, ChatInvokeRequest{ Driver: p.Driver, ModelName: p.ModelID, APIKey: p.APIKey, @@ -629,7 +631,7 @@ func (t *subAgentTool) InvokableRun(ctx context.Context, argsJSON string, _ ...e } } - out, err := NewAgentComponent(t.spec.Param).Invoke(ctx, inputs) + out, err := NewAgentComponent(t.spec.Param).Invoke(ctx, dao.DB, inputs) if err != nil { return "", fmt.Errorf("sub-agent tool %q: %w", normalizeAgentToolName(t.spec.Name), err) } @@ -715,7 +717,7 @@ func (c *AgentComponent) Name() string { return "Agent" } // Invoke either returns a lazy Agent stream for a direct downstream Message, // or executes the Agent eagerly for all other graph shapes. The mode is a // compile-time canvas decision carried through context, not a DSL parameter. -func (c *AgentComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *AgentComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { if runtime.ComponentExecutionOptionsFromContext(ctx).DeferAgentToMessage { // Preserve the Agent-class duration installed by the node wrapper, then // start a fresh deadline when Message opens the lazy stream. @@ -729,18 +731,18 @@ func (c *AgentComponent) Invoke(ctx context.Context, inputs map[string]any) (map Open: func(openCtx context.Context, sink runtime.AgentDeltaSink) (map[string]any, error) { agentCtx, cancel := context.WithTimeout(openCtx, timeout) defer cancel() - return c.invokeNow(runtime.WithAgentDeltaSink(agentCtx, sink), inputs) + return c.invokeNow(runtime.WithAgentDeltaSink(agentCtx, sink), db, inputs) }, } return map[string]any{"content": deferred}, nil } - return c.invokeNow(ctx, inputs) + return c.invokeNow(ctx, db, inputs) } // invokeNow contains the original eager Agent execution path. Deferred // Message consumption calls this function later with an invocation-local // delta sink, so the same ReAct/citation/tool behavior is reused once. -func (c *AgentComponent) invokeNow(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *AgentComponent) invokeNow(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { runtime.ResetAgentMessageEmission(ctx) defer runtime.FinalizeAgentMessage(ctx) @@ -800,7 +802,7 @@ func (c *AgentComponent) invokeNow(ctx context.Context, inputs map[string]any) ( // actually consumes. if p.OptimizeMultiTurn { if state, _, sErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); sErr == nil && state != nil { - if rephrased, err := optimizeMultiTurnQuestion(ctx, p, state.SnapshotPriorHistory()); err == nil && rephrased != "" { + if rephrased, err := optimizeMultiTurnQuestion(ctx, db, p, state.SnapshotPriorHistory()); err == nil && rephrased != "" { p.UserPrompt = rephrased } } @@ -813,7 +815,7 @@ func (c *AgentComponent) invokeNow(ctx context.Context, inputs map[string]any) ( // actual user/assistant turns maintained by the canvas service. if err == nil && msg != nil { if state, _, sErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); sErr == nil && state != nil { - if summary, sErr2 := addToolCallMemory(ctx, p, msg); sErr2 == nil && summary != "" { + if summary, sErr2 := addToolCallMemory(ctx, db, p, msg); sErr2 == nil && summary != "" { state.AppendMemory(p.UserPrompt, msg.Content, summary) } } @@ -860,7 +862,7 @@ func (c *AgentComponent) invokeNow(ctx context.Context, inputs map[string]any) ( if len(chunks) == 0 { groundingStatus = "no_chunks" } else { - grounded, gErr := applyCitationGrounding(ctx, p, content, chunks) + grounded, gErr := applyCitationGrounding(ctx, db, p, content, chunks) if gErr == nil && grounded != content { content = grounded groundingStatus = "applied" @@ -901,11 +903,11 @@ func isAgentGraphRunError(err error) bool { // Stream implements Component.Stream. Mirrors Invoke then pushes the // single payload through the channel. -func (c *AgentComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { +func (c *AgentComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { out := make(chan map[string]any, 1) go func() { defer close(out) - result, err := c.Invoke(ctx, inputs) + result, err := c.Invoke(ctx, db, inputs) if err != nil { out <- map[string]any{"error": err.Error()} return diff --git a/internal/agent/component/agent_test.go b/internal/agent/component/agent_test.go index e7c67ccd62..dc90c2e6db 100644 --- a/internal/agent/component/agent_test.go +++ b/internal/agent/component/agent_test.go @@ -49,7 +49,7 @@ func TestAgent_NoToolsReAct(t *testing.T) { }) c := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 3}) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "user_prompt": "what is 6*7?", }) if err != nil { @@ -80,7 +80,7 @@ func TestAgent_EmitsThinking(t *testing.T) { }) c := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "user_prompt": "hello", }) if err != nil { @@ -106,7 +106,7 @@ func TestAgent_MessageEmissionIsScopedPerInvocation(t *testing.T) { }) state := runtime.NewCanvasState("run-1", "task-1") - ctx := runtime.WithState(context.Background(), state) + ctx := runtime.WithState(t.Context(), state) var contents []string ctx = runtime.WithAgentMessageEmitterControl(ctx, func(contentDelta, _ string) { @@ -119,11 +119,11 @@ func TestAgent_MessageEmissionIsScopedPerInvocation(t *testing.T) { ) first := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}) - if _, err := first.Invoke(ctx, map[string]any{"user_prompt": "first"}); err != nil { + if _, err := first.Invoke(ctx, nil, map[string]any{"user_prompt": "first"}); err != nil { t.Fatalf("first Invoke: %v", err) } second := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}) - if _, err := second.Invoke(ctx, map[string]any{"user_prompt": "second"}); err != nil { + if _, err := second.Invoke(ctx, nil, map[string]any{"user_prompt": "second"}); err != nil { t.Fatalf("second Invoke: %v", err) } @@ -139,12 +139,12 @@ func TestAgent_DefersExecutionForDownstreamMessage(t *testing.T) { return &schema.Message{Role: schema.Assistant, Content: "lazy answer"}, nil }) - ctx := runtime.WithComponentExecutionOptions(context.Background(), runtime.ComponentExecutionOptions{ + ctx := runtime.WithComponentExecutionOptions(t.Context(), runtime.ComponentExecutionOptions{ DeferAgentToMessage: true, }) ctx = runtime.WithAgentMessageEmitter(ctx, func(string, string) {}) agent := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}) - out, err := agent.Invoke(ctx, map[string]any{"user_prompt": "hello"}) + out, err := agent.Invoke(ctx, nil, map[string]any{"user_prompt": "hello"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -178,10 +178,10 @@ func TestAgent_DeferredStreamAppliesFreshTimeout(t *testing.T) { return &schema.Message{Role: schema.Assistant, Content: "answer"}, nil }) - ctx := runtime.WithComponentExecutionOptions(context.Background(), runtime.ComponentExecutionOptions{ + ctx := runtime.WithComponentExecutionOptions(t.Context(), runtime.ComponentExecutionOptions{ DeferAgentToMessage: true, }) - out, err := NewAgentComponent(AgentParam{ModelID: "stub"}).Invoke(ctx, map[string]any{"user_prompt": "hello"}) + out, err := NewAgentComponent(AgentParam{ModelID: "stub"}).Invoke(ctx, nil, map[string]any{"user_prompt": "hello"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -189,7 +189,7 @@ func TestAgent_DeferredStreamAppliesFreshTimeout(t *testing.T) { if !ok || deferred == nil { t.Fatalf("content=%T, want *runtime.DeferredStream", out["content"]) } - if _, err := deferred.Open(context.Background(), func(string, string) {}); err != nil { + if _, err := deferred.Open(t.Context(), func(string, string) {}); err != nil { t.Fatalf("Open: %v", err) } } @@ -201,11 +201,11 @@ func TestAgent_DeferredStreamDoesNotAppendFinalAnswerAfterDeltas(t *testing.T) { return &schema.Message{Role: schema.Assistant, Content: "lazy answer"}, nil }) - ctx := runtime.WithComponentExecutionOptions(context.Background(), runtime.ComponentExecutionOptions{ + ctx := runtime.WithComponentExecutionOptions(t.Context(), runtime.ComponentExecutionOptions{ DeferAgentToMessage: true, }) ctx = runtime.WithAgentMessageEmitter(ctx, func(string, string) {}) - out, err := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}).Invoke(ctx, map[string]any{ + out, err := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}).Invoke(ctx, nil, map[string]any{ "user_prompt": "hello", }) if err != nil { @@ -236,7 +236,7 @@ func TestAgent_SuppressesVisibleEventsWithoutMessageDownstream(t *testing.T) { return &schema.Message{Role: schema.Assistant, Content: "timeline answer"}, nil }) var emitted []string - ctx := runtime.WithAgentMessageEmitter(context.Background(), func(contentDelta, _ string) { + ctx := runtime.WithAgentMessageEmitter(t.Context(), func(contentDelta, _ string) { if contentDelta != "" { emitted = append(emitted, contentDelta) } @@ -244,7 +244,7 @@ func TestAgent_SuppressesVisibleEventsWithoutMessageDownstream(t *testing.T) { ctx = runtime.WithComponentExecutionOptions(ctx, runtime.ComponentExecutionOptions{ SuppressAgentMessageEvents: true, }) - if _, err := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}).Invoke(ctx, map[string]any{ + if _, err := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}).Invoke(ctx, nil, map[string]any{ "user_prompt": "hello", }); err != nil { t.Fatalf("Invoke: %v", err) @@ -272,7 +272,7 @@ func TestAgent_ForwardsThinkingParam(t *testing.T) { if err != nil { t.Fatalf("New(Agent): %v", err) } - if _, err := cmp.Invoke(context.Background(), nil); err != nil { + if _, err = cmp.Invoke(t.Context(), nil, nil); err != nil { t.Fatalf("Invoke: %v", err) } if gotThinking != "enabled" { @@ -289,7 +289,7 @@ func TestAgent_ResolvesUserPromptFromCanvasState(t *testing.T) { state := runtime.NewCanvasState("run-1", "task-1") state.Sys["query"] = "what is marigold" - ctx := runtime.WithState(context.Background(), state) + ctx := runtime.WithState(t.Context(), state) c := NewAgentComponent(AgentParam{ ModelID: "stub", @@ -297,7 +297,7 @@ func TestAgent_ResolvesUserPromptFromCanvasState(t *testing.T) { UserPrompt: "Question: {sys.query}", MaxRounds: 1, }) - if _, err := c.Invoke(ctx, nil); err != nil { + if _, err := c.Invoke(ctx, nil, nil); err != nil { t.Fatalf("Invoke: %v", err) } if gotPrompt != "Question: what is marigold" { @@ -327,9 +327,9 @@ func TestAgent_UsesPromptsListForSysQuery(t *testing.T) { state := runtime.NewCanvasState("run-1", "task-1") state.Sys["query"] = "用户真正的问题" - ctx := runtime.WithState(context.Background(), state) + ctx := runtime.WithState(t.Context(), state) - if _, err := cmp.Invoke(ctx, nil); err != nil { + if _, err = cmp.Invoke(ctx, nil, nil); err != nil { t.Fatalf("Invoke: %v", err) } if gotPrompt != "用户真正的问题" { @@ -355,9 +355,9 @@ func TestAgent_NewTreatsSchemaDefaultUserPromptAsSysQueryPlaceholder(t *testing. state := runtime.NewCanvasState("run-1", "task-1") state.Sys["query"] = "用户真正的问题" - ctx := runtime.WithState(context.Background(), state) + ctx := runtime.WithState(t.Context(), state) - if _, err := cmp.Invoke(ctx, nil); err != nil { + if _, err = cmp.Invoke(ctx, nil, nil); err != nil { t.Fatalf("Invoke: %v", err) } if gotPrompt != "用户真正的问题" { @@ -388,9 +388,9 @@ func TestAgent_EmptyConfiguredUserPromptDoesNotFallbackToSysQuery(t *testing.T) state := runtime.NewCanvasState("run-1", "task-1") state.Sys["query"] = "1" state.SetVar("UserFillUp:TwelveBadgersRescue", "key", "21") - ctx := runtime.WithState(context.Background(), state) + ctx := runtime.WithState(t.Context(), state) - if _, err := cmp.Invoke(ctx, nil); err != nil { + if _, err = cmp.Invoke(ctx, nil, nil); err != nil { t.Fatalf("Invoke: %v", err) } if gotSystemPrompt != "User answer: 21" { @@ -416,7 +416,7 @@ func TestAgent_FormatsRuntimePromptLikePython(t *testing.T) { APIKey: "test-key", MaxRounds: 1, }) - if _, err := c.Invoke(context.Background(), map[string]any{ + if _, err := c.Invoke(t.Context(), nil, map[string]any{ "user_prompt": "write answer", "reasoning": "selected because it can answer", "context": "known facts", @@ -451,7 +451,7 @@ func TestAgent_ToolCallRound(t *testing.T) { }) c := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 3}) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "user_prompt": "find out about ragflow", }) if err != nil { @@ -481,7 +481,7 @@ func TestAgent_ExhaustRoundsError(t *testing.T) { }) c := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 2}) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "user_prompt": "x", }) if err != nil { @@ -494,7 +494,7 @@ func TestAgent_ExhaustRoundsError(t *testing.T) { func TestAgent_MissingModelID(t *testing.T) { c := NewAgentComponent(AgentParam{MaxRounds: 1}) - _, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "x"}) + _, err := c.Invoke(t.Context(), nil, map[string]any{"user_prompt": "x"}) if err == nil { t.Fatal("expected ParamError for missing model_id") } @@ -521,7 +521,7 @@ func TestAgent_Invoke_RespectsParentCancellation(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() // pre-cancel - _, err := c.Invoke(ctx, map[string]any{"user_prompt": "hi"}) + _, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "hi"}) if err == nil { t.Fatal("expected error from pre-cancelled context") } @@ -536,7 +536,7 @@ func TestAgent_UnknownToolName(t *testing.T) { MaxRounds: 1, Tools: []string{"does_not_exist"}, }) - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(t.Context(), nil, map[string]any{ "user_prompt": "x", }) if err == nil { @@ -555,7 +555,7 @@ func TestAgent_AllRegisteredToolsConfigPassesToRunner(t *testing.T) { }) c := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}) - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(t.Context(), nil, map[string]any{ "user_prompt": "x", "tools": []any{ "akshare", "arxiv", "code_exec", "crawler", "deepl", "duckduckgo", @@ -594,7 +594,7 @@ func TestAgent_AcceptsCanvasToolObjects(t *testing.T) { }) c := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}) - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(t.Context(), nil, map[string]any{ "user_prompt": "x", "tools": []any{ map[string]any{ @@ -705,7 +705,7 @@ func TestAgent_CanCreateReactAgentWithAllRegisteredTools(t *testing.T) { if len(tools) != len(p.Tools) { t.Fatalf("len(tools) = %d, want %d", len(tools), len(p.Tools)) } - _, err = react.NewAgent(context.Background(), &react.AgentConfig{ + _, err = react.NewAgent(t.Context(), &react.AgentConfig{ ToolCallingModel: &fakeToolCallingChatModel{}, ToolsConfig: compose.ToolsNodeConfig{ Tools: tools, @@ -846,7 +846,7 @@ func TestAgent_SubAgentToolInvokableRunCallsChildAgent(t *testing.T) { }, }} - out, err := tool.InvokableRun(context.Background(), `{"user_prompt":"ask child","reasoning":"because","context":"facts"}`) + out, err := tool.InvokableRun(t.Context(), `{"user_prompt":"ask child","reasoning":"because","context":"facts"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -968,7 +968,7 @@ func TestAgent_ReActExhaustsSteps(t *testing.T) { toolArgs: `{"sql": "SELECT 1"}`, } - agent, err := react.NewAgent(context.Background(), &react.AgentConfig{ + agent, err := react.NewAgent(t.Context(), &react.AgentConfig{ ToolCallingModel: mdl, ToolsConfig: compose.ToolsNodeConfig{ Tools: []einotool.BaseTool{realTool}, @@ -979,7 +979,7 @@ func TestAgent_ReActExhaustsSteps(t *testing.T) { t.Fatalf("react.NewAgent: %v", err) } - out, err := agent.Generate(context.Background(), []*schema.Message{ + out, err := agent.Generate(t.Context(), []*schema.Message{ schema.UserMessage("loop forever"), }) if err == nil { diff --git a/internal/agent/component/base.go b/internal/agent/component/base.go index 9b585c162e..b58c8af33e 100644 --- a/internal/agent/component/base.go +++ b/internal/agent/component/base.go @@ -19,6 +19,8 @@ import ( "context" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) // Component is the runtime contract every RAGFlow component implements. @@ -42,13 +44,13 @@ type Component interface { // parameter map (variable references already substituted by the canvas // engine). Returns the output map; components should put their public // outputs at top-level keys. - Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) + Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) // Stream is the streaming variant. The default implementation may // return a buffered channel that emits the same payload as Invoke, then // closes — components that natively stream (LLM, Message) override. // May return (nil, nil) for non-streaming components. - Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) + Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) // Inputs returns parameter metadata: param_name → description. Inputs() map[string]string diff --git a/internal/agent/component/begin.go b/internal/agent/component/begin.go index 3a255ceaa1..03f2dec47b 100644 --- a/internal/agent/component/begin.go +++ b/internal/agent/component/begin.go @@ -29,6 +29,8 @@ import ( "maps" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) // mapsCopy is a thin alias for the stdlib maps.Copy to keep the call @@ -62,7 +64,7 @@ func (b *BeginComponent) Name() string { return b.name } // the shared *CanvasState.Sys namespace, then returns the input map as // outputs unchanged. The input map is shallow-copied to avoid aliasing // surprises across concurrent goroutines that share an inputs map. -func (b *BeginComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (b *BeginComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) if err != nil { return nil, fmt.Errorf("Begin: %w", err) @@ -101,8 +103,8 @@ func (b *BeginComponent) Invoke(ctx context.Context, inputs map[string]any) (map // Stream is a synchronous facade over Invoke for P0. SSE streaming of // Begin output is not meaningful (Begin has no I/O), so the channel // receives a single payload and closes — same shape as Invoke's return. -func (b *BeginComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := b.Invoke(ctx, inputs) +func (b *BeginComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := b.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/begin_test.go b/internal/agent/component/begin_test.go index 7f056f3f7b..efab2e2fa6 100644 --- a/internal/agent/component/begin_test.go +++ b/internal/agent/component/begin_test.go @@ -35,7 +35,7 @@ func TestBegin_InjectsSys(t *testing.T) { state := canvas.NewCanvasState("run-1", "task-1") ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{"query": "hello"}) + out, err := c.Invoke(ctx, nil, map[string]any{"query": "hello"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -67,7 +67,7 @@ func TestBegin_PassesThroughInputs(t *testing.T) { "inputs": map[string]any{"k": "v"}, "extra": 42, } - out, err := c.Invoke(ctx, inputs) + out, err := c.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -110,7 +110,7 @@ func TestBegin_InjectsWebhookPayload(t *testing.T) { "query": "", "webhook_payload": payload, } - out, err := c.Invoke(ctx, inputs) + out, err := c.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -135,7 +135,7 @@ func TestBegin_AbsentWebhookPayload(t *testing.T) { state := canvas.NewCanvasState("run-4", "task-4") ctx := canvas.WithState(context.Background(), state) - if _, err := c.Invoke(ctx, map[string]any{"query": "plain chat"}); err != nil { + if _, err := c.Invoke(ctx, nil, map[string]any{"query": "plain chat"}); err != nil { t.Fatalf("Invoke: %v", err) } if _, ok := state.Sys["webhook_payload"]; ok { @@ -150,7 +150,7 @@ func TestBegin_EmptyWebhookPayload(t *testing.T) { state := canvas.NewCanvasState("run-5", "task-5") ctx := canvas.WithState(context.Background(), state) - if _, err := c.Invoke(ctx, map[string]any{ + if _, err := c.Invoke(ctx, nil, map[string]any{ "query": "", "webhook_payload": map[string]any{}, }); err != nil { diff --git a/internal/agent/component/browser.go b/internal/agent/component/browser.go index 7d7372aed9..21e730a8e2 100644 --- a/internal/agent/component/browser.go +++ b/internal/agent/component/browser.go @@ -282,7 +282,7 @@ func (b *BrowserComponent) Name() string { return b.name } // // File upload/download and session persistence are not supported // in this component; they are v1-deferred. -func (b *BrowserComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (b *BrowserComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) if err != nil { return nil, fmt.Errorf("Browser: %w", err) @@ -385,8 +385,8 @@ func stagehandModelName(providerName, modelName string) string { } // Stream mirrors Invoke; Browser is a single-shot generator. -func (b *BrowserComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := b.Invoke(ctx, inputs) +func (b *BrowserComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := b.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/browser_test.go b/internal/agent/component/browser_test.go index ebc9bc6125..17f685cfed 100644 --- a/internal/agent/component/browser_test.go +++ b/internal/agent/component/browser_test.go @@ -199,7 +199,7 @@ func TestBrowser_DispatchesToRuntime(t *testing.T) { }) ctx := stateWith(t, map[string]any{"tenant_id": "tenant-1"}) - _, err := c.Invoke(ctx, nil) + _, err := c.Invoke(ctx, nil, nil) if err == nil { t.Fatal("expected tenant LLM lookup error, got nil") } @@ -281,7 +281,7 @@ func TestBrowser_MissingTenant(t *testing.T) { state := canvas.NewCanvasState("run-1", "task-1") ctx := canvas.WithState(context.Background(), state) - _, err := c.Invoke(ctx, nil) + _, err := c.Invoke(ctx, nil, nil) if err == nil || !strings.Contains(err.Error(), "tenant_id") { t.Errorf("expected tenant_id error, got %v", err) } @@ -313,7 +313,7 @@ func TestBrowser_PropagatesRuntimeError(t *testing.T) { }) ctx := stateWith(t, map[string]any{"tenant_id": "tenant-1"}) - _, err := c.Invoke(ctx, nil) + _, err := c.Invoke(ctx, nil, nil) if err == nil { t.Fatal("expected runtime error, got nil") } @@ -431,7 +431,7 @@ func TestBrowser_RunExtractRequestShape(t *testing.T) { }) ctx := stateWith(t, map[string]any{"tenant_id": "tenant-1"}) - if _, err := c.Invoke(ctx, nil); err != nil { + if _, err := c.Invoke(ctx, nil, nil); err != nil { t.Fatalf("Invoke: %v", err) } req := mock.lastRequest(t) @@ -474,7 +474,7 @@ func TestBrowser_HeadlessPropagates(t *testing.T) { }) ctx := stateWith(t, map[string]any{"tenant_id": "tenant-1"}) - if _, err := c.Invoke(ctx, nil); err != nil { + if _, err := c.Invoke(ctx, nil, nil); err != nil { t.Fatalf("Invoke: %v", err) } req := mock.lastRequest(t) @@ -505,7 +505,7 @@ func TestBrowser_OutputsShape(t *testing.T) { }) ctx := stateWith(t, map[string]any{"tenant_id": "tenant-1"}) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } diff --git a/internal/agent/component/canvas_var_resolution_test.go b/internal/agent/component/canvas_var_resolution_test.go index 8a065c6044..e4747a5c4b 100644 --- a/internal/agent/component/canvas_var_resolution_test.go +++ b/internal/agent/component/canvas_var_resolution_test.go @@ -39,7 +39,7 @@ func TestLLM_Invoke_ResolvesTemplateRefs(t *testing.T) { UserPrompt: "What does {{retrieval:0@content}} say?", }) ctx := runtime.WithState(context.Background(), state) - _, err := c.Invoke(ctx, map[string]any{}) + _, err := c.Invoke(ctx, nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -64,7 +64,7 @@ func TestLLM_Invoke_NoState_LeavesPromptsUnchanged(t *testing.T) { ModelID: "echo", UserPrompt: "What does {{retrieval:0@content}} say?", }) - _, err := c.Invoke(context.Background(), map[string]any{}) + _, err := c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -90,7 +90,7 @@ func TestLLM_Invoke_UnresolvedRef_LeavesPromptIntact(t *testing.T) { UserPrompt: "Use {{retrieval:0@content}} please", }) ctx := runtime.WithState(context.Background(), state) - _, err := c.Invoke(ctx, map[string]any{}) + _, err := c.Invoke(ctx, nil, map[string]any{}) if err != nil { t.Fatalf("Invoke should not error on unresolved ref: %v", err) } diff --git a/internal/agent/component/categorize.go b/internal/agent/component/categorize.go index 9c8168c9e9..e482c7a15f 100644 --- a/internal/agent/component/categorize.go +++ b/internal/agent/component/categorize.go @@ -15,6 +15,7 @@ import ( "strings" "github.com/cloudwego/eino/schema" + "gorm.io/gorm" "ragflow/internal/agent/runtime" ) @@ -64,7 +65,7 @@ func (c *CategorizeComponent) Name() string { return "Categorize" } // Invoke calls the chat model, parses the response for a category, and // returns the chosen category (or the default if the model returned // something outside the configured set). -func (c *CategorizeComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *CategorizeComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { p := mergeCategorizeParam(c.param, inputs) var err error p.ModelID, p.Driver, p.APIKey, p.BaseURL, err = resolveChatModelRef(ctx, p.ModelID, p.Driver, p.APIKey, p.BaseURL) @@ -97,7 +98,7 @@ func (c *CategorizeComponent) Invoke(ctx context.Context, inputs map[string]any) {Role: schema.System, Content: sysPrompt}, {Role: schema.User, Content: userPrompt}, } - resp, err := inv.Invoke(ctx, ChatInvokeRequest{ + resp, err := inv.Invoke(ctx, db, ChatInvokeRequest{ Driver: p.Driver, ModelName: p.ModelID, APIKey: p.APIKey, @@ -122,11 +123,11 @@ func (c *CategorizeComponent) Invoke(ctx context.Context, inputs map[string]any) } // Stream mirrors Invoke as a single chunk. -func (c *CategorizeComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { +func (c *CategorizeComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { out := make(chan map[string]any, 1) go func() { defer close(out) - result, err := c.Invoke(ctx, inputs) + result, err := c.Invoke(ctx, db, inputs) if err != nil { out <- map[string]any{"error": err.Error()} return diff --git a/internal/agent/component/categorize_test.go b/internal/agent/component/categorize_test.go index 943e76d73b..cdbbcfcafc 100644 --- a/internal/agent/component/categorize_test.go +++ b/internal/agent/component/categorize_test.go @@ -23,7 +23,7 @@ func TestCategorize_ChosenCategory(t *testing.T) { Categories: []string{"sales", "support", "billing"}, DefaultCategory: "support", }) - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -61,7 +61,7 @@ func TestCategorize_FallbackToDefault(t *testing.T) { Categories: []string{"a", "b", "c"}, DefaultCategory: "b", }) - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -79,7 +79,7 @@ func TestCategorize_DefaultDefaultsToFirstCategory(t *testing.T) { Categories: []string{"alpha", "beta", "gamma"}, // no default_category }) - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -97,7 +97,7 @@ func TestCategorize_CaseInsensitive(t *testing.T) { Categories: []string{"sales", "support", "billing"}, DefaultCategory: "sales", }) - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -118,7 +118,7 @@ func TestCategorize_PromptListsCategories(t *testing.T) { DefaultCategory: "x", Items: []string{"foo", "bar"}, }) - _, err := c.Invoke(context.Background(), map[string]any{}) + _, err := c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -159,7 +159,7 @@ func TestCategorize_PromptIncludesRuntimeQuery(t *testing.T) { "English": "this query has english letter", }, }) - _, err := c.Invoke(ctx, map[string]any{}) + _, err := c.Invoke(ctx, nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -189,7 +189,7 @@ func TestCategorize_PromptUsesInputQueryValue(t *testing.T) { Query: "sys.query", Categories: []string{"Number", "chinese", "English"}, }) - _, err := c.Invoke(context.Background(), map[string]any{"query": "测试"}) + _, err := c.Invoke(context.Background(), nil, map[string]any{"query": "测试"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -225,7 +225,7 @@ func TestCategorize_HistoryWindowRealData(t *testing.T) { Categories: []string{"Number", "chinese", "English"}, MessageHistoryWindowSize: 2, }) - _, err := c.Invoke(ctx, map[string]any{"query": "current question"}) + _, err := c.Invoke(ctx, nil, map[string]any{"query": "current question"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -262,7 +262,7 @@ func TestCategorizeRegistered_DefaultHistoryWindow(t *testing.T) { if err != nil { t.Fatalf("New(Categorize): %v", err) } - _, err = c.Invoke(ctx, map[string]any{"query": "current question"}) + _, err = c.Invoke(ctx, nil, map[string]any{"query": "current question"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -286,7 +286,7 @@ func TestCategorize_HistoryWindowRejectsNegative(t *testing.T) { Categories: []string{"Number", "chinese", "English"}, MessageHistoryWindowSize: -1, }) - _, err := c.Invoke(context.Background(), map[string]any{"query": "current question"}) + _, err := c.Invoke(context.Background(), nil, map[string]any{"query": "current question"}) if err == nil { t.Fatal("expected negative message_history_window_size error") } @@ -334,7 +334,7 @@ func TestCategorize_SplitsCompositeLLMIDIntoDriverAndModel(t *testing.T) { Categories: []string{"sales", "support"}, DefaultCategory: "support", }) - _, err := c.Invoke(context.Background(), map[string]any{}) + _, err := c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -398,7 +398,7 @@ func TestCategorize_ResolvesTenantModelInstanceCredentials(t *testing.T) { Categories: []string{"sales", "support"}, DefaultCategory: "support", }) - _, err := c.Invoke(stateWithTenant("tenant-1"), map[string]any{}) + _, err := c.Invoke(stateWithTenant("tenant-1"), db, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -452,7 +452,7 @@ func TestCategorize_ResolvesTenantModelID(t *testing.T) { Categories: []string{"sales", "support"}, DefaultCategory: "support", }) - _, err := c.Invoke(stateWithTenant("tenant-1"), map[string]any{}) + _, err := c.Invoke(stateWithTenant("tenant-1"), db, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -502,7 +502,7 @@ func TestCategorize_ResolvesSoleActiveInstanceWhenDefaultMissing(t *testing.T) { Categories: []string{"sales", "support"}, DefaultCategory: "support", }) - _, err := c.Invoke(stateWithTenant("tenant-1"), map[string]any{}) + _, err := c.Invoke(stateWithTenant("tenant-1"), db, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -527,7 +527,7 @@ func TestCategorize_RoutesToSelectedCategoryHandle(t *testing.T) { CategoryRoutes: map[string]string{"打招呼": "a111", "Retrieval": "b222", "Other": "c333"}, DefaultCategory: "Other", }) - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -545,7 +545,7 @@ func TestCategorize_RoutesFromCategoryDescriptionToList(t *testing.T) { withStubInvoker(t, stub) c := NewCategorizeComponent(CategorizeParam{ModelID: "stub"}) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "category_description": map[string]any{ "打招呼": map[string]any{"description": "hello", "to": []any{"Message:CateLoop"}}, "Retrieval": map[string]any{"description": "rag", "to": []any{"Message:CateRetrieval"}}, @@ -569,7 +569,7 @@ func TestCategorize_ExplicitCategoriesKeepCategoryDescriptionMetadata(t *testing withStubInvoker(t, stub) c := NewCategorizeComponent(CategorizeParam{ModelID: "stub"}) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "categories": []any{"Number", "chinese", "English"}, "category_description": map[string]any{ "Number": map[string]any{"description": "This query has only a number", "examples": []any{"4321"}, "to": []any{"Message:Number"}}, @@ -617,7 +617,7 @@ func TestCategorizeRegistered_ExplicitCategoriesKeepCategoryDescriptionMetadata( if err != nil { t.Fatalf("New(Categorize): %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{"query": "hello"}) + out, err := c.Invoke(context.Background(), nil, map[string]any{"query": "hello"}) if err != nil { t.Fatalf("Invoke: %v", err) } diff --git a/internal/agent/component/chat_template_kwargs_test.go b/internal/agent/component/chat_template_kwargs_test.go index bedc90ae16..74b95e0d4b 100644 --- a/internal/agent/component/chat_template_kwargs_test.go +++ b/internal/agent/component/chat_template_kwargs_test.go @@ -76,7 +76,7 @@ func TestLLM_Invoke_ChatTemplateKwargsDoesNotBreak(t *testing.T) { ModelID: "echo", ChatTemplateKwargs: map[string]any{"seed": 1}, }) - _, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "hi"}) + _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "hi"}) if err != nil { t.Fatalf("Invoke: %v", err) } diff --git a/internal/agent/component/citation_test.go b/internal/agent/component/citation_test.go index ce25ca9f19..cec92ad073 100644 --- a/internal/agent/component/citation_test.go +++ b/internal/agent/component/citation_test.go @@ -132,7 +132,7 @@ func TestLLM_Invoke_AppendsCitationPrompt(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo", SystemPrompt: "You are a bot.", Cite: true}) - if _, err := c.Invoke(context.Background(), map[string]any{ + if _, err := c.Invoke(context.Background(), nil, map[string]any{ "user_prompt": "hi", }); err != nil { t.Fatalf("Invoke: %v", err) @@ -159,7 +159,7 @@ func TestLLM_Invoke_CiteFalseDisablesInjection(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo", SystemPrompt: "sys"}) - _, _ = c.Invoke(context.Background(), map[string]any{ + _, _ = c.Invoke(context.Background(), nil, map[string]any{ "user_prompt": "hi", "cite": false, }) diff --git a/internal/agent/component/data_operations.go b/internal/agent/component/data_operations.go index a740dfdbe0..12c35344c7 100644 --- a/internal/agent/component/data_operations.go +++ b/internal/agent/component/data_operations.go @@ -44,6 +44,8 @@ import ( "strings" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) const componentNameDataOperations = "DataOperations" @@ -182,7 +184,7 @@ func (d *DataOperationsComponent) Name() string { return d.name } // Invoke loads input_objects from the configured query refs, then // dispatches to the operation-specific helper. -func (d *DataOperationsComponent) Invoke(ctx context.Context, _ map[string]any) (map[string]any, error) { +func (d *DataOperationsComponent) Invoke(ctx context.Context, db *gorm.DB, _ map[string]any) (map[string]any, error) { state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) if err != nil { return nil, fmt.Errorf("DataOperations: %w", err) @@ -245,8 +247,8 @@ func (d *DataOperationsComponent) Invoke(ctx context.Context, _ map[string]any) } // Stream mirrors Invoke; DataOperations is a single-shot transform. -func (d *DataOperationsComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := d.Invoke(ctx, inputs) +func (d *DataOperationsComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := d.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/data_operations_test.go b/internal/agent/component/data_operations_test.go index 663c7311a9..329abb4848 100644 --- a/internal/agent/component/data_operations_test.go +++ b/internal/agent/component/data_operations_test.go @@ -68,7 +68,7 @@ func TestDataOperations_SelectKeys(t *testing.T) { }} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -100,7 +100,7 @@ func TestDataOperations_Combine(t *testing.T) { state.Outputs["cpn_1"] = map[string]any{"d2": map[string]any{"k": []any{2, 3}}} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -130,7 +130,7 @@ func TestDataOperations_RemoveKeys(t *testing.T) { }} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -171,7 +171,7 @@ func TestDataOperations_LiteralEval(t *testing.T) { }} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -212,7 +212,7 @@ func TestDataOperations_FilterValues(t *testing.T) { }} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -237,7 +237,7 @@ func TestDataOperations_AppendOrUpdate(t *testing.T) { }} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -264,7 +264,7 @@ func TestDataOperations_RenameKeys(t *testing.T) { }} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } diff --git a/internal/agent/component/docs_generator.go b/internal/agent/component/docs_generator.go index 768eda54b0..5ad6d9762a 100644 --- a/internal/agent/component/docs_generator.go +++ b/internal/agent/component/docs_generator.go @@ -39,6 +39,7 @@ import ( "time" "github.com/google/uuid" + "gorm.io/gorm" iow "ragflow/internal/agent/component/io" "ragflow/internal/agent/runtime" @@ -196,7 +197,7 @@ func (d *DocsGenerator) Name() string { return d.name } // Invoke dispatches to the appropriate writer. Input overrides for // content / filename are honored. -func (d *DocsGenerator) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (d *DocsGenerator) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { param := d.param content, err := resolveDocsGeneratorContent(ctx, param.Content, inputs) if err != nil { @@ -412,8 +413,8 @@ func agentAttachmentDownloadPath(attachmentID, ext, mimeType, filename string) s } // Stream mirrors Invoke; DocsGenerator is a single-shot generator. -func (d *DocsGenerator) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := d.Invoke(ctx, inputs) +func (d *DocsGenerator) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := d.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/docs_generator_test.go b/internal/agent/component/docs_generator_test.go index 80b6108f9b..f1320e886a 100644 --- a/internal/agent/component/docs_generator_test.go +++ b/internal/agent/component/docs_generator_test.go @@ -96,7 +96,7 @@ func TestDocsGenerator_Invoke_HappyPath(t *testing.T) { if err != nil { t.Fatalf("New(DocsGenerator): %v", err) } - _, _ = c.Invoke(context.Background(), map[string]any{}) + _, _ = c.Invoke(context.Background(), nil, map[string]any{}) // We do not assert err == nil here because txt output requires // the internal writer (which may not be available in this // checkout). The test pins that the call doesn't panic. @@ -110,7 +110,7 @@ func TestDocsGenerator_Invoke_UsesQueryWhenContentEmpty(t *testing.T) { if err != nil { t.Fatalf("New(DocsGenerator): %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "query": "Hello from begin query", }) if err != nil { @@ -137,7 +137,7 @@ func TestDocsGenerator_Invoke_AcceptsDecorationParamNames(t *testing.T) { if err != nil { t.Fatalf("New(DocsGenerator): %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke canonical params: %v", err) } @@ -161,7 +161,7 @@ func TestDocsGenerator_Invoke_AcceptsDecorationParamNames(t *testing.T) { if err != nil { t.Fatalf("New(DocsGenerator alias params): %v", err) } - out, err = c.Invoke(context.Background(), map[string]any{}) + out, err = c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke alias params: %v", err) } @@ -194,7 +194,7 @@ func TestDocsGenerator_Invoke_StoresAgentAttachment(t *testing.T) { state.Sys["tenant_id"] = "tenant-1" ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{}) + out, err := c.Invoke(ctx, nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } diff --git a/internal/agent/component/downstream_message_test.go b/internal/agent/component/downstream_message_test.go index a71c3bf245..d00e23b609 100644 --- a/internal/agent/component/downstream_message_test.go +++ b/internal/agent/component/downstream_message_test.go @@ -48,7 +48,7 @@ func TestLLM_Stream_ExposesThinkingAndContentForDownstreamMessage(t *testing.T) withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo"}) - ch, err := c.Stream(context.Background(), map[string]any{"user_prompt": "go"}) + ch, err := c.Stream(context.Background(), nil, map[string]any{"user_prompt": "go"}) if err != nil { t.Fatalf("Stream: %v", err) } diff --git a/internal/agent/component/excel_processor.go b/internal/agent/component/excel_processor.go index ffdd109bd8..b348c004bc 100644 --- a/internal/agent/component/excel_processor.go +++ b/internal/agent/component/excel_processor.go @@ -45,6 +45,7 @@ import ( "strings" "github.com/xuri/excelize/v2" + "gorm.io/gorm" "ragflow/internal/agent/runtime" ) @@ -154,7 +155,7 @@ func (e *ExcelProcessorComponent) Name() string { return e.name } // write — {"rows": [][]any, "sheet_names": []string, "size": , // "bytes": <[]byte>} // merge — {"rows": [][]any, "sheet_names": []string, "size": } -func (e *ExcelProcessorComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (e *ExcelProcessorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { // ExcelProcessor does not currently read from canvas state for // binary blobs, but we still pull state so a nil-state error is // surfaced early. @@ -194,8 +195,8 @@ func (e *ExcelProcessorComponent) Invoke(ctx context.Context, inputs map[string] } // Stream mirrors Invoke; ExcelProcessor is a single-shot transform. -func (e *ExcelProcessorComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := e.Invoke(ctx, inputs) +func (e *ExcelProcessorComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := e.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/excel_processor_test.go b/internal/agent/component/excel_processor_test.go index 348946fe4c..d4f4664815 100644 --- a/internal/agent/component/excel_processor_test.go +++ b/internal/agent/component/excel_processor_test.go @@ -49,7 +49,7 @@ func TestExcelProcessor_WriteThenRead(t *testing.T) { if err != nil { t.Fatalf("NewExcelProcessorComponent (write): %v", err) } - out, err := w.Invoke(excelCtx(t), map[string]any{}) + out, err := w.Invoke(excelCtx(t), nil, map[string]any{}) if err != nil { t.Fatalf("write Invoke: %v", err) } @@ -73,7 +73,7 @@ func TestExcelProcessor_WriteThenRead(t *testing.T) { if err != nil { t.Fatalf("NewExcelProcessorComponent (read): %v", err) } - rout, err := r.Invoke(excelCtx(t), map[string]any{"bytes": raw}) + rout, err := r.Invoke(excelCtx(t), nil, map[string]any{"bytes": raw}) if err != nil { t.Fatalf("read Invoke: %v", err) } @@ -116,7 +116,7 @@ func TestExcelProcessor_ReadSheetNames(t *testing.T) { } r, _ := NewExcelProcessorComponent(map[string]any{"operation": "read"}) - out, err := r.Invoke(excelCtx(t), map[string]any{"bytes": buf.Bytes()}) + out, err := r.Invoke(excelCtx(t), nil, map[string]any{"bytes": buf.Bytes()}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -134,7 +134,7 @@ func TestExcelProcessor_EmptyFile(t *testing.T) { "operation": "write", "output_data": [][]any{}, }) - out, err := w.Invoke(excelCtx(t), map[string]any{}) + out, err := w.Invoke(excelCtx(t), nil, map[string]any{}) if err != nil { t.Fatalf("write Invoke: %v", err) } @@ -147,7 +147,7 @@ func TestExcelProcessor_EmptyFile(t *testing.T) { } r, _ := NewExcelProcessorComponent(map[string]any{"operation": "read"}) - rout, err := r.Invoke(excelCtx(t), map[string]any{"bytes": raw}) + rout, err := r.Invoke(excelCtx(t), nil, map[string]any{"bytes": raw}) if err != nil { t.Fatalf("read Invoke: %v", err) } @@ -172,7 +172,7 @@ func TestExcelProcessor_ParamCheck(t *testing.T) { // surfaces a ParamError. func TestExcelProcessor_ReadMissingBytes(t *testing.T) { r, _ := NewExcelProcessorComponent(map[string]any{"operation": "read"}) - _, err := r.Invoke(excelCtx(t), map[string]any{}) + _, err := r.Invoke(excelCtx(t), nil, map[string]any{}) if err == nil { t.Fatal("expected error for missing bytes, got nil") } diff --git a/internal/agent/component/fillup.go b/internal/agent/component/fillup.go index df9aabd988..7cd7111db8 100644 --- a/internal/agent/component/fillup.go +++ b/internal/agent/component/fillup.go @@ -29,6 +29,8 @@ import ( "fmt" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) const componentNameFillup = "Fillup" @@ -80,7 +82,7 @@ func (f *FillupComponent) Name() string { return f.name } // Invoke emits one output per form field, with file-typed fields // stubbed as "". No "tips" key is added — that is the // defining difference from UserFillUp. -func (f *FillupComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (f *FillupComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { // State is required by the engine contract; we don't read from it // here, but we still extract it to fail loudly if the engine forgot // to wire it (consistent with UserFillUp's behavior). @@ -98,8 +100,8 @@ func (f *FillupComponent) Invoke(ctx context.Context, inputs map[string]any) (ma // Stream is the synchronous facade over Invoke: a single payload, then // close. Mirrors the pattern used by UserFillUp and the P0 components. -func (f *FillupComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := f.Invoke(ctx, inputs) +func (f *FillupComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := f.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/fillup_test.go b/internal/agent/component/fillup_test.go index 204f02f39d..c6c97c5c7d 100644 --- a/internal/agent/component/fillup_test.go +++ b/internal/agent/component/fillup_test.go @@ -32,7 +32,7 @@ func TestFillup_PassesThroughInputs(t *testing.T) { state := canvas.NewCanvasState("run-1", "task-1") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "inputs": map[string]any{ "a": map[string]any{"value": 1}, "b": map[string]any{"value": "x"}, @@ -62,7 +62,7 @@ func TestFillup_NoTipsKey(t *testing.T) { state := canvas.NewCanvasState("run-2", "task-2") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "inputs": map[string]any{ "x": map[string]any{"value": "y"}, }, @@ -88,7 +88,7 @@ func TestFillup_NonDictInput(t *testing.T) { state := canvas.NewCanvasState("run-3", "task-3") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "inputs": map[string]any{ "plain_str": "just a string", "plain_int": 42, @@ -119,7 +119,7 @@ func TestFillup_FileInputStub(t *testing.T) { state := canvas.NewCanvasState("run-4", "task-4") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "inputs": map[string]any{ "cv": map[string]any{ "value": []any{"file-1"}, diff --git a/internal/agent/component/fixture_stubs.go b/internal/agent/component/fixture_stubs.go index d46f59cffa..aa0a6df81c 100644 --- a/internal/agent/component/fixture_stubs.go +++ b/internal/agent/component/fixture_stubs.go @@ -41,6 +41,8 @@ import ( "fmt" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) // ----- Retrieval ----- @@ -69,13 +71,13 @@ func (r *RetrievalStub) Name() string { return componentNameRetrieval } // Invoke returns a stub result that downstream templates can // resolve. `formalized_content` is the field the test fixtures // reference; empty string is the safe fixture value. -func (r *RetrievalStub) Invoke(_ context.Context, _ map[string]any) (map[string]any, error) { +func (r *RetrievalStub) Invoke(_ context.Context, _ *gorm.DB, _ map[string]any) (map[string]any, error) { return map[string]any{"formalized_content": ""}, nil } // Stream mirrors Invoke as a single-chunk SSE stream. -func (r *RetrievalStub) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := r.Invoke(ctx, inputs) +func (r *RetrievalStub) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := r.Invoke(ctx, db, inputs) if err != nil { return nil, err } @@ -125,13 +127,13 @@ func (t *TavilySearchStub) Name() string { return componentNameTavilySearch } // Invoke returns an empty `formalized_content` so downstream // templates resolve. -func (t *TavilySearchStub) Invoke(_ context.Context, _ map[string]any) (map[string]any, error) { +func (t *TavilySearchStub) Invoke(_ context.Context, _ *gorm.DB, _ map[string]any) (map[string]any, error) { return map[string]any{"formalized_content": ""}, nil } // Stream mirrors Invoke. -func (t *TavilySearchStub) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := t.Invoke(ctx, inputs) +func (t *TavilySearchStub) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := t.Invoke(ctx, db, inputs) if err != nil { return nil, err } @@ -187,7 +189,7 @@ func (e *ExeSQLStub) Name() string { return componentNameExeSQL } // Invoke returns a stable two-column stub result. Downstream // templates that render SQL output will see headers + an empty row // — enough for the message surface to format a string. -func (e *ExeSQLStub) Invoke(_ context.Context, _ map[string]any) (map[string]any, error) { +func (e *ExeSQLStub) Invoke(_ context.Context, _ *gorm.DB, _ map[string]any) (map[string]any, error) { return map[string]any{ "columns": []string{"col1", "col2"}, "rows": [][]any{{"", ""}}, @@ -196,8 +198,8 @@ func (e *ExeSQLStub) Invoke(_ context.Context, _ map[string]any) (map[string]any } // Stream mirrors Invoke. -func (e *ExeSQLStub) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := e.Invoke(ctx, inputs) +func (e *ExeSQLStub) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := e.Invoke(ctx, db, inputs) if err != nil { return nil, err } @@ -258,13 +260,13 @@ func NewGenerateStub(params map[string]any) (Component, error) { func (g *GenerateStub) Name() string { return componentNameGenerate } // Invoke delegates to the LLM component. -func (g *GenerateStub) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - return g.inner.Invoke(ctx, inputs) +func (g *GenerateStub) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { + return g.inner.Invoke(ctx, db, inputs) } // Stream delegates to the LLM component. -func (g *GenerateStub) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - return g.inner.Stream(ctx, inputs) +func (g *GenerateStub) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + return g.inner.Stream(ctx, db, inputs) } // Inputs returns the DSL param surface. Matches LLM's surface @@ -342,7 +344,7 @@ func (a *AnswerStub) Name() string { return componentNameAnswer } // Invoke returns an empty answer. Real implementation will block // until the user provides input; the stub is fire-and-forget so // the e2e flow doesn't deadlock. -func (a *AnswerStub) Invoke(ctx context.Context, _ map[string]any) (map[string]any, error) { +func (a *AnswerStub) Invoke(ctx context.Context, db *gorm.DB, _ map[string]any) (map[string]any, error) { // Mirror the no-state-check pattern of Message/Retrieval: we // don't read state, but the signature must match. if _, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx); err != nil { @@ -352,8 +354,8 @@ func (a *AnswerStub) Invoke(ctx context.Context, _ map[string]any) (map[string]a } // Stream mirrors Invoke. -func (a *AnswerStub) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := a.Invoke(ctx, inputs) +func (a *AnswerStub) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := a.Invoke(ctx, db, inputs) if err != nil { return nil, err } @@ -418,13 +420,13 @@ func NewIterationStub(params map[string]any) (Component, error) { func (i *IterationStub) Name() string { return componentNameIteration } // Invoke delegates to the inner Parallel component. -func (i *IterationStub) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - return i.inner.Invoke(ctx, inputs) +func (i *IterationStub) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { + return i.inner.Invoke(ctx, db, inputs) } // Stream delegates to the inner Parallel component. -func (i *IterationStub) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - return i.inner.Stream(ctx, inputs) +func (i *IterationStub) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + return i.inner.Stream(ctx, db, inputs) } // Inputs mirrors Parallel's surface for introspection. @@ -451,13 +453,13 @@ func NewIterationItemStub(_ map[string]any) (Component, error) { func (it *IterationItemStub) Name() string { return componentNameIterationItem } // Invoke returns a passthrough empty map. -func (it *IterationItemStub) Invoke(_ context.Context, _ map[string]any) (map[string]any, error) { +func (it *IterationItemStub) Invoke(_ context.Context, _ *gorm.DB, _ map[string]any) (map[string]any, error) { return map[string]any{"result": ""}, nil } // Stream mirrors Invoke. -func (it *IterationItemStub) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := it.Invoke(ctx, inputs) +func (it *IterationItemStub) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := it.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/grounding_test.go b/internal/agent/component/grounding_test.go index 36fad6ea3f..204c3baa73 100644 --- a/internal/agent/component/grounding_test.go +++ b/internal/agent/component/grounding_test.go @@ -29,6 +29,7 @@ import ( "testing" "github.com/cloudwego/eino/schema" + "gorm.io/gorm" "ragflow/internal/agent/canvas" "ragflow/internal/agent/runtime" @@ -45,7 +46,7 @@ type groundingTestInvoker struct { calls int } -func (g *groundingTestInvoker) Invoke(_ context.Context, req ChatInvokeRequest) (*ChatInvokeResponse, error) { +func (g *groundingTestInvoker) Invoke(_ context.Context, _ *gorm.DB, req ChatInvokeRequest) (*ChatInvokeResponse, error) { g.mu.Lock() defer g.mu.Unlock() g.lastReq = req @@ -79,7 +80,7 @@ func TestGrounding_Applied(t *testing.T) { ctx := runtime.WithState(context.Background(), state) c := NewAgentComponent(AgentParam{ModelID: "stub", Cite: true}) - out, err := c.Invoke(ctx, map[string]any{"user_prompt": "q"}) + out, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "q"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -118,7 +119,7 @@ func TestGrounding_NoChunks(t *testing.T) { ctx := runtime.WithState(context.Background(), state) c := NewAgentComponent(AgentParam{ModelID: "stub", Cite: true}) - out, err := c.Invoke(ctx, map[string]any{"user_prompt": "q"}) + out, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "q"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -146,7 +147,7 @@ func TestGrounding_CiteFalse(t *testing.T) { }) c := NewAgentComponent(AgentParam{ModelID: "stub", Cite: false}) - out, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "q"}) + out, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "q"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -175,7 +176,7 @@ func TestGrounding_LLMError(t *testing.T) { ctx := runtime.WithState(context.Background(), state) c := NewAgentComponent(AgentParam{ModelID: "stub", Cite: true}) - out, err := c.Invoke(ctx, map[string]any{"user_prompt": "q"}) + out, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "q"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -205,7 +206,7 @@ func TestGrounding_EmptyContent(t *testing.T) { ctx := runtime.WithState(context.Background(), state) c := NewAgentComponent(AgentParam{ModelID: "stub", Cite: true}) - out, err := c.Invoke(ctx, map[string]any{"user_prompt": "q"}) + out, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "q"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -218,7 +219,7 @@ type errInvoker struct { err error } -func (e *errInvoker) Invoke(_ context.Context, _ ChatInvokeRequest) (*ChatInvokeResponse, error) { +func (e *errInvoker) Invoke(_ context.Context, _ *gorm.DB, _ ChatInvokeRequest) (*ChatInvokeResponse, error) { return nil, e.err } diff --git a/internal/agent/component/history_window_test.go b/internal/agent/component/history_window_test.go index 9d831f5a59..aa4c13e528 100644 --- a/internal/agent/component/history_window_test.go +++ b/internal/agent/component/history_window_test.go @@ -131,7 +131,7 @@ func TestLLM_Invoke_HistoryWindow_PrependsFromState(t *testing.T) { MessageHistoryWindowSize: window, }) ctx := runtime.WithState(context.Background(), state) - _, err := c.Invoke(ctx, map[string]any{}) + _, err := c.Invoke(ctx, nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -167,7 +167,7 @@ func TestLLM_Invoke_HistoryWindow_DoesNotDuplicateCurrentUser(t *testing.T) { MessageHistoryWindowSize: 5, }) ctx := runtime.WithState(context.Background(), state) - if _, err := c.Invoke(ctx, map[string]any{}); err != nil { + if _, err := c.Invoke(ctx, nil, map[string]any{}); err != nil { t.Fatalf("Invoke: %v", err) } @@ -201,7 +201,7 @@ func TestLLM_Invoke_HistoryWindow_ZeroIsNoop(t *testing.T) { // MessageHistoryWindowSize: 0 (default) }) ctx := runtime.WithState(context.Background(), state) - _, _ = c.Invoke(ctx, map[string]any{}) + _, _ = c.Invoke(ctx, nil, map[string]any{}) if len(stub.captured.Messages) != 1 { t.Errorf("expected only 1 message (no history), got %d", len(stub.captured.Messages)) } diff --git a/internal/agent/component/invoke.go b/internal/agent/component/invoke.go index 655c114d16..4b14d19015 100644 --- a/internal/agent/component/invoke.go +++ b/internal/agent/component/invoke.go @@ -49,6 +49,7 @@ import ( "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp" "go.uber.org/zap" + "gorm.io/gorm" "ragflow/internal/utility" ) @@ -101,7 +102,7 @@ func (i *InvokeComponent) Name() string { return i.name } // On any of those checks failing the function returns an `_ERROR` // output (no Go error) so the canvas can route around the failure // the same way the Python fix does, instead of crashing the node. -func (i *InvokeComponent) Invoke(ctx context.Context, inputs map[string]any) (output map[string]any, invokeErr error) { +func (i *InvokeComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (output map[string]any, invokeErr error) { startedAt := time.Now() defer func() { if output == nil { @@ -404,8 +405,8 @@ func sanitizeLogURL(raw string) string { // Stream is a synchronous facade over Invoke. Real streaming // (chunked transfer as it arrives) is a future enhancement. -func (i *InvokeComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := i.Invoke(ctx, inputs) +func (i *InvokeComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := i.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/invoke_test.go b/internal/agent/component/invoke_test.go index f50c08390b..13a8c3d956 100644 --- a/internal/agent/component/invoke_test.go +++ b/internal/agent/component/invoke_test.go @@ -63,7 +63,7 @@ func TestInvoke_GET(t *testing.T) { defer srv.Close() c, _ := NewInvokeComponent(nil) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "method": "GET", "url": srv.URL, }) @@ -101,7 +101,7 @@ func TestInvoke_POST(t *testing.T) { defer srv.Close() c, _ := NewInvokeComponent(nil) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "method": "POST", "url": srv.URL, "body": `{"k":"v"}`, @@ -139,7 +139,7 @@ func TestInvoke_UsesNodeParams(t *testing.T) { if err != nil { t.Fatalf("NewInvokeComponent: %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -176,7 +176,7 @@ func TestInvoke_BadMethod(t *testing.T) { setupAllowAnyHost(t, true) c, _ := NewInvokeComponent(nil) - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(context.Background(), nil, map[string]any{ "method": "PATCH", "url": "http://localhost:1", }) @@ -193,7 +193,7 @@ func TestInvoke_MissingURL(t *testing.T) { setupAllowAnyHost(t, true) c, _ := NewInvokeComponent(nil) - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(context.Background(), nil, map[string]any{ "method": "GET", }) if err == nil { @@ -222,7 +222,7 @@ func TestInvoke_SSRFGuard_BlocksLoopback(t *testing.T) { defer srv.Close() c, _ := NewInvokeComponent(nil) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "method": "GET", "url": srv.URL, }) @@ -252,7 +252,7 @@ func TestInvoke_SSRFGuard_BlocksMetadataIP(t *testing.T) { defer srv.Close() c, _ := NewInvokeComponent(nil) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "method": "GET", "url": "http://169.254.169.254/latest/meta-data/", }) @@ -288,7 +288,7 @@ func TestInvoke_SSRFGuard_BlocksProxy(t *testing.T) { // (without a side-effect on the local server) is still // caught. PR review round 5 / duplicate fix from the // serverHit-only assertion in the earlier review. - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "method": "GET", "url": "https://example.com/api", "proxy": "http://127.0.0.1:8080", @@ -320,7 +320,7 @@ func TestInvoke_NoRedirects_NotFollowed(t *testing.T) { defer srv.Close() c, _ := NewInvokeComponent(nil) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "method": "GET", "url": srv.URL, }) @@ -403,7 +403,7 @@ func TestInvoke_ProxyDNSPin(t *testing.T) { c, _ := NewInvokeComponent(nil) ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - _, err = c.Invoke(ctx, map[string]any{ + _, err = c.Invoke(ctx, nil, map[string]any{ "method": "GET", "url": "http://8.8.8.8/api", "proxy": "http://proxy.test.invalid:" + proxyPort, @@ -429,7 +429,7 @@ func TestInvoke_ProxyRejectsHostnameTarget(t *testing.T) { setupAllowAnyHost(t, false) c, _ := NewInvokeComponent(nil) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "method": "GET", "url": "http://example.com/api", "proxy": "http://proxy.example.invalid:9999", diff --git a/internal/agent/component/list_operations.go b/internal/agent/component/list_operations.go index 74bb4f02e3..831a4edc2a 100644 --- a/internal/agent/component/list_operations.go +++ b/internal/agent/component/list_operations.go @@ -41,6 +41,8 @@ import ( "strings" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) const componentNameListOperations = "ListOperations" @@ -269,7 +271,7 @@ func (l *ListOperationsComponent) Name() string { return l.name } // errors) into a returned error. Any other panic is re-raised so a // real bug in the operator code is not masked as a "ListOperations: // ..." error. -func (l *ListOperationsComponent) Invoke(ctx context.Context, _ map[string]any) (result map[string]any, err error) { +func (l *ListOperationsComponent) Invoke(ctx context.Context, db *gorm.DB, _ map[string]any) (result map[string]any, err error) { defer func() { r := recover() if r == nil { @@ -330,8 +332,8 @@ func (l *ListOperationsComponent) Invoke(ctx context.Context, _ map[string]any) } // Stream mirrors Invoke; ListOperations is a single-shot transform. -func (l *ListOperationsComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := l.Invoke(ctx, inputs) +func (l *ListOperationsComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := l.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/list_operations_test.go b/internal/agent/component/list_operations_test.go index cbb93d4310..b9c13376bf 100644 --- a/internal/agent/component/list_operations_test.go +++ b/internal/agent/component/list_operations_test.go @@ -40,7 +40,7 @@ func TestListOperations_Head(t *testing.T) { state.Outputs["cpn_0"] = map[string]any{"xs": []any{1, 2, 3, 4, 5}} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -73,7 +73,7 @@ func TestListOperations_TopNLegacyAlias(t *testing.T) { state.Outputs["cpn_0"] = map[string]any{"xs": []any{1, 2, 3, 4}} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -96,7 +96,7 @@ func TestListOperations_Filter(t *testing.T) { state.Outputs["cpn_0"] = map[string]any{"xs": []any{"foo", "bar", "foobar"}} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -121,7 +121,7 @@ func TestListOperations_DropDuplicates(t *testing.T) { }} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -148,7 +148,7 @@ func TestListOperations_Tail(t *testing.T) { state.Outputs["cpn_0"] = map[string]any{"xs": []any{1, 2, 3, 4, 5}} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -170,7 +170,7 @@ func TestListOperations_NthPositive(t *testing.T) { state.Outputs["cpn_0"] = map[string]any{"xs": []any{"a", "b", "c", "d"}} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -192,7 +192,7 @@ func TestListOperations_SortDesc(t *testing.T) { state.Outputs["cpn_0"] = map[string]any{"xs": []any{3, 1, 4, 1, 5, 9, 2, 6}} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -239,7 +239,7 @@ func TestListOperations_SortByFieldList(t *testing.T) { if err != nil { t.Fatalf("NewListOperationsComponent: %v", err) } - out, err := cSortDesc.Invoke(ctx, nil) + out, err := cSortDesc.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke sort_by=score desc: %v", err) } @@ -264,7 +264,7 @@ func TestListOperations_SortByFieldList(t *testing.T) { if err != nil { t.Fatalf("NewListOperationsComponent (legacy): %v", err) } - out, err = cLegacy.Invoke(ctx, nil) + out, err = cLegacy.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke sort_by=\"\" desc: %v", err) } @@ -288,7 +288,7 @@ func TestListOperations_SortByFieldList(t *testing.T) { if err != nil { t.Fatalf("NewListOperationsComponent (tiebreak): %v", err) } - out, err = cTiebreak.Invoke(ctx, nil) + out, err = cTiebreak.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke sort_by=score,title asc: %v", err) } @@ -314,7 +314,7 @@ func TestListOperations_NotAList(t *testing.T) { state.Outputs["cpn_0"] = map[string]any{"x": "not-a-list"} ctx := canvas.WithState(context.Background(), state) - _, err := c.Invoke(ctx, nil) + _, err := c.Invoke(ctx, nil, nil) if err == nil { t.Fatal("expected error for non-list input, got nil") } @@ -364,7 +364,7 @@ func TestListOperations_StrictMode_ReturnsError(t *testing.T) { state.Outputs["cpn_0"] = map[string]any{"xs": []any{1, 2, 3}} ctx := canvas.WithState(context.Background(), state) - _, err = c.Invoke(ctx, nil) + _, err = c.Invoke(ctx, nil, nil) if err == nil { t.Fatal("expected error from strict-mode nth with n=0, got nil") } @@ -391,7 +391,7 @@ func TestListOperations_StrictStringCoercion(t *testing.T) { state.Outputs["cpn_0"] = map[string]any{"xs": []any{1, 2, 3}} ctx := canvas.WithState(context.Background(), state) - _, err = c.Invoke(ctx, nil) + _, err = c.Invoke(ctx, nil, nil) if err == nil { t.Errorf("[strict=%q] expected error from strict-mode nth with n=0, got nil", v) } @@ -416,7 +416,7 @@ func TestListOperations_StrictFalseStringsIgnored(t *testing.T) { state.Outputs["cpn_0"] = map[string]any{"xs": []any{1, 2, 3}} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Errorf("[strict=%q] expected non-error (coerced to false), got %v", v, err) } @@ -454,7 +454,7 @@ func TestListOperations_FilterEqBool(t *testing.T) { state.Outputs["cpn_0"] = map[string]any{"xs": []any{true, false, true, "True"}} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -482,7 +482,7 @@ func TestListOperations_UnknownOp_ReturnsError(t *testing.T) { state.Outputs["cpn_0"] = map[string]any{"xs": []any{1, 2, 3}} ctx := canvas.WithState(context.Background(), state) - _, err := c.Invoke(ctx, nil) + _, err := c.Invoke(ctx, nil, nil) if err == nil { t.Fatal("expected error for unknown operation, got nil") } diff --git a/internal/agent/component/llm.go b/internal/agent/component/llm.go index eb9ebdaa30..5101221ed9 100644 --- a/internal/agent/component/llm.go +++ b/internal/agent/component/llm.go @@ -22,6 +22,7 @@ import ( "time" "github.com/cloudwego/eino/schema" + "gorm.io/gorm" "ragflow/internal/agent/component/prompts" "ragflow/internal/agent/runtime" @@ -133,7 +134,7 @@ type LLMOutput struct { // chat model. The default implementation lives in this file; tests can // override the package-level defaultChatInvoker to inject a stub. type ChatInvoker interface { - Invoke(ctx context.Context, req ChatInvokeRequest) (*ChatInvokeResponse, error) + Invoke(ctx context.Context, db *gorm.DB, req ChatInvokeRequest) (*ChatInvokeResponse, error) } // ChatInvokeRequest is the minimal surface the LLM component needs to @@ -204,7 +205,7 @@ func GetDefaultChatInvokerForTest() ChatInvoker { type einoChatInvoker struct{} // Invoke satisfies ChatInvoker. -func (e *einoChatInvoker) Invoke(ctx context.Context, req ChatInvokeRequest) (*ChatInvokeResponse, error) { +func (e *einoChatInvoker) Invoke(ctx context.Context, db *gorm.DB, req ChatInvokeRequest) (*ChatInvokeResponse, error) { if req.ModelName == "" { return nil, fmt.Errorf("component: LLM: model_id is required") } @@ -312,7 +313,7 @@ func NewLLMComponent(p LLMParam) *LLMComponent { func (c *LLMComponent) Name() string { return "LLM" } // Invoke runs the LLM and returns the output map. -func (c *LLMComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *LLMComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { p := mergeLLMParam(c.param, inputs) // Resolve tenant-scoped custom models (and fill missing driver/credentials) @@ -493,7 +494,7 @@ func (c *LLMComponent) Invoke(ctx context.Context, inputs map[string]any) (map[s // MaxRetries is an absolute count, not a stacked one. inv = newRetryInvoker(unwrapChatInvoker(inv), maxRetries, delay) } - resp, err := inv.Invoke(ctx, ChatInvokeRequest{ + resp, err := inv.Invoke(ctx, db, ChatInvokeRequest{ Driver: p.Driver, ModelName: p.ModelID, APIKey: p.APIKey, @@ -544,7 +545,7 @@ func (c *LLMComponent) Invoke(ctx context.Context, inputs map[string]any) (map[s // deferred to a future phase. parsed, ok := matchOutputStructure(resp.Content, p.OutputStructure) if !ok { - retryResp, err := inv.Invoke(ctx, ChatInvokeRequest{ + retryResp, err := inv.Invoke(ctx, db, ChatInvokeRequest{ Driver: p.Driver, ModelName: p.ModelID, APIKey: p.APIKey, @@ -599,7 +600,7 @@ func (c *LLMComponent) Invoke(ctx context.Context, inputs map[string]any) (map[s // is deferred — the public surface here is correct, only the data // source needs to be swapped to a real StreamReader consumer in a // follow-up. -func (c *LLMComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { +func (c *LLMComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { out := make(chan map[string]any, 16) go func() { defer close(out) @@ -610,7 +611,7 @@ func (c *LLMComponent) Stream(ctx context.Context, inputs map[string]any) (<-cha if err := ctx.Err(); err != nil { return } - result, err := c.Invoke(ctx, inputs) + result, err := c.Invoke(ctx, db, inputs) if err != nil { select { case out <- map[string]any{"error": err.Error()}: diff --git a/internal/agent/component/llm_retry.go b/internal/agent/component/llm_retry.go index edaae0d2b8..b654f111f9 100644 --- a/internal/agent/component/llm_retry.go +++ b/internal/agent/component/llm_retry.go @@ -25,6 +25,8 @@ import ( "context" "fmt" "time" + + "gorm.io/gorm" ) // retryInvoker decorates a ChatInvoker with exponential-backoff @@ -113,14 +115,14 @@ func unwrapChatInvoker(inv ChatInvoker) ChatInvoker { // sleeping initialDelay * 2^attempt between failures. The sleep // honours ctx cancellation: a cancelled context aborts the backoff // and returns ctx.Err() immediately. -func (r *retryInvoker) Invoke(ctx context.Context, req ChatInvokeRequest) (*ChatInvokeResponse, error) { +func (r *retryInvoker) Invoke(ctx context.Context, db *gorm.DB, req ChatInvokeRequest) (*ChatInvokeResponse, error) { if r.inner == nil { return nil, fmt.Errorf("component: retryInvoker: nil inner") } delay := r.initialDelay var lastErr error for attempt := 0; attempt <= r.maxRetries; attempt++ { - resp, err := r.inner.Invoke(ctx, req) + resp, err := r.inner.Invoke(ctx, db, req) if err == nil { return resp, nil } diff --git a/internal/agent/component/llm_retry_test.go b/internal/agent/component/llm_retry_test.go index 2b9fdc1047..c55298e969 100644 --- a/internal/agent/component/llm_retry_test.go +++ b/internal/agent/component/llm_retry_test.go @@ -16,6 +16,8 @@ import ( "sync/atomic" "testing" "time" + + "gorm.io/gorm" ) // scriptedInvoker fails the first failTimes calls then succeeds. @@ -27,7 +29,7 @@ type scriptedInvoker struct { calls int32 } -func (s *scriptedInvoker) Invoke(_ context.Context, _ ChatInvokeRequest) (*ChatInvokeResponse, error) { +func (s *scriptedInvoker) Invoke(_ context.Context, _ *gorm.DB, _ ChatInvokeRequest) (*ChatInvokeResponse, error) { n := atomic.AddInt32(&s.calls, 1) if n <= atomic.LoadInt32(&s.failTimes) { return nil, s.err @@ -44,7 +46,7 @@ type alwaysFailInvoker struct { calls int32 } -func (a *alwaysFailInvoker) Invoke(_ context.Context, _ ChatInvokeRequest) (*ChatInvokeResponse, error) { +func (a *alwaysFailInvoker) Invoke(_ context.Context, _ *gorm.DB, _ ChatInvokeRequest) (*ChatInvokeResponse, error) { atomic.AddInt32(&a.calls, 1) return nil, a.err } @@ -59,7 +61,7 @@ func TestRetryInvoker_SucceedsOnSecondAttempt(t *testing.T) { inner := &scriptedInvoker{failTimes: 1, err: errors.New("transient"), resp: want} r := newRetryInvoker(inner, 3, time.Millisecond) - resp, err := r.Invoke(context.Background(), ChatInvokeRequest{ModelName: "m"}) + resp, err := r.Invoke(context.Background(), nil, ChatInvokeRequest{ModelName: "m"}) if err != nil { t.Fatalf("Invoke: unexpected err: %v", err) } @@ -79,7 +81,7 @@ func TestRetryInvoker_ExhaustsRetries(t *testing.T) { inner := &alwaysFailInvoker{err: sentinel} r := newRetryInvoker(inner, 3, time.Millisecond) - _, err := r.Invoke(context.Background(), ChatInvokeRequest{ModelName: "m"}) + _, err := r.Invoke(context.Background(), nil, ChatInvokeRequest{ModelName: "m"}) if err == nil { t.Fatal("expected error after exhaustion") } @@ -127,7 +129,7 @@ func TestRetryInvoker_HonorsContextCancellation(t *testing.T) { cancel() }() - _, err := r.Invoke(ctx, ChatInvokeRequest{ModelName: "m"}) + _, err := r.Invoke(ctx, nil, ChatInvokeRequest{ModelName: "m"}) if err == nil { t.Fatal("expected error from cancelled context") @@ -162,7 +164,7 @@ func TestRetryInvoker_ExponentialBackoff(t *testing.T) { r := newRetryInvoker(inner, 3, initial) start := time.Now() - _, _ = r.Invoke(context.Background(), ChatInvokeRequest{ModelName: "m"}) + _, _ = r.Invoke(context.Background(), nil, ChatInvokeRequest{ModelName: "m"}) elapsed := time.Since(start) // 20 + 40 + 80 = 140ms of backoff (3 retries, 4 attempts total). @@ -183,7 +185,7 @@ func TestRetryInvoker_NoRetriesWhenZero(t *testing.T) { inner := &alwaysFailInvoker{err: errors.New("nope")} r := newRetryInvoker(inner, 0, 50*time.Millisecond) - _, err := r.Invoke(context.Background(), ChatInvokeRequest{ModelName: "m"}) + _, err := r.Invoke(context.Background(), nil, ChatInvokeRequest{ModelName: "m"}) if err == nil { t.Fatal("expected error") } @@ -196,7 +198,7 @@ func TestRetryInvoker_NoRetriesWhenZero(t *testing.T) { // should not panic when constructed with nil inner. func TestRetryInvoker_NilInner(t *testing.T) { r := newRetryInvoker(nil, 3, time.Millisecond) - _, err := r.Invoke(context.Background(), ChatInvokeRequest{ModelName: "m"}) + _, err := r.Invoke(context.Background(), nil, ChatInvokeRequest{ModelName: "m"}) if err == nil { t.Fatal("expected error for nil inner") } @@ -218,7 +220,7 @@ func TestLLMParam_RespectsMaxRetries(t *testing.T) { // zero-value default is 2s, which would make this test slow. c.param.DelayAfterError = time.Millisecond - _, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "x"}) + _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "x"}) if err == nil { t.Fatal("expected error from exhausted retries") } @@ -247,7 +249,7 @@ func TestLLMParam_ZeroRetriesMeansOneAttempt(t *testing.T) { // with maxRetries=0. c.param.DelayAfterError = time.Millisecond - _, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "x"}) + _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "x"}) if err == nil { t.Fatal("expected error") } @@ -278,7 +280,7 @@ func TestLLMParam_DefaultRetries(t *testing.T) { SetDefaultChatInvoker(newRetryInvoker(fastInner, 2, time.Millisecond)) t.Cleanup(func() { SetDefaultChatInvoker(nil) }) - _, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "x"}) + _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "x"}) if err == nil { t.Fatal("expected error after default retries") } @@ -362,7 +364,7 @@ func TestLLM_ParamOverride_AbsoluteCount_NotStacked(t *testing.T) { // Force a tiny delay so the test runs fast. c.param.DelayAfterError = time.Millisecond - _, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "x"}) + _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "x"}) if err == nil { t.Fatal("expected error from exhausted retries") } @@ -392,7 +394,7 @@ func TestLLM_NoParamOverride_StackingPreserved(t *testing.T) { // No param override: MaxRetries=0 AND DelayAfterError=0. c := NewLLMComponent(LLMParam{ModelID: "m"}) - _, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "x"}) + _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "x"}) if err == nil { t.Fatal("expected error from exhausted retries") } diff --git a/internal/agent/component/llm_test.go b/internal/agent/component/llm_test.go index dcab7bc5c6..9f0625a949 100644 --- a/internal/agent/component/llm_test.go +++ b/internal/agent/component/llm_test.go @@ -21,6 +21,7 @@ import ( "ragflow/internal/entity" "github.com/cloudwego/eino/schema" + "gorm.io/gorm" ) // stubInvoker is a programmable ChatInvoker used by these tests. @@ -31,7 +32,7 @@ type stubInvoker struct { calls int } -func (s *stubInvoker) Invoke(_ context.Context, req ChatInvokeRequest) (*ChatInvokeResponse, error) { +func (s *stubInvoker) Invoke(_ context.Context, _ *gorm.DB, req ChatInvokeRequest) (*ChatInvokeResponse, error) { s.calls++ cp := req s.captured = &cp @@ -54,7 +55,7 @@ func TestLLM_Invoke_HappyPath(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo-model"}) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "user_prompt": "hi", }) if err != nil { @@ -85,7 +86,7 @@ func TestLLM_Invoke_JSONOutput(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo"}) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "user_prompt": "give me json", "json_output": true, }) @@ -109,7 +110,7 @@ func TestLLM_Invoke_SystemAndUser(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo"}) - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(t.Context(), nil, map[string]any{ "system_prompt": "you are helpful", "user_prompt": "say hi", }) @@ -132,7 +133,7 @@ func TestLLM_Stream(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo"}) - ch, err := c.Stream(context.Background(), map[string]any{"user_prompt": "go"}) + ch, err := c.Stream(t.Context(), nil, map[string]any{"user_prompt": "go"}) if err != nil { t.Fatalf("Stream: %v", err) } @@ -156,7 +157,7 @@ func TestLLM_Stream(t *testing.T) { func TestLLM_Invoke_MissingModelID(t *testing.T) { withStubInvoker(t, &stubInvoker{resp: &ChatInvokeResponse{Content: "should not be called"}}) c := NewLLMComponent(LLMParam{}) // no model_id - _, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "x"}) + _, err := c.Invoke(t.Context(), nil, map[string]any{"user_prompt": "x"}) if err == nil { t.Fatal("expected ParamError for missing model_id") } @@ -170,7 +171,7 @@ func TestLLM_Invoke_InvokerError(t *testing.T) { stub := &stubInvoker{err: errors.New("upstream blew up")} withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo"}) - _, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "x"}) + _, err := c.Invoke(t.Context(), nil, map[string]any{"user_prompt": "x"}) if err == nil { t.Fatal("expected error to propagate") } @@ -298,7 +299,7 @@ func TestLLM_ResolvesTenantModelID(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "3d2d824e7e5d11f1a845455b140cef90"}) - _, err := c.Invoke(stateWithTenant("tenant-1"), map[string]any{"user_prompt": "hi"}) + _, err := c.Invoke(stateWithTenant("tenant-1"), db, map[string]any{"user_prompt": "hi"}) if err != nil { t.Fatalf("Invoke: %v", err) } diff --git a/internal/agent/component/loop.go b/internal/agent/component/loop.go index def6d0a5c8..d6bc771a06 100644 --- a/internal/agent/component/loop.go +++ b/internal/agent/component/loop.go @@ -46,6 +46,8 @@ package component import ( "context" + + "gorm.io/gorm" ) const componentNameLoop = "Loop" @@ -171,13 +173,13 @@ func (c *LoopComponent) Outputs() map[string]string { // The returned map is empty. State writes from this method would be // silently dropped by the eino graph, because LoopComponent is not // registered as an eino node when the macro expansion fires. -func (c *LoopComponent) Invoke(_ context.Context, _ map[string]any) (map[string]any, error) { +func (c *LoopComponent) Invoke(_ context.Context, _ *gorm.DB, _ map[string]any) (map[string]any, error) { return map[string]any{}, nil } // Stream mirrors Invoke and emits an empty map as a single chunk. -func (c *LoopComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := c.Invoke(ctx, inputs) +func (c *LoopComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := c.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/loop_test.go b/internal/agent/component/loop_test.go index ab453bb964..2154493661 100644 --- a/internal/agent/component/loop_test.go +++ b/internal/agent/component/loop_test.go @@ -84,7 +84,7 @@ func TestLoop_InvokeIsNoOp(t *testing.T) { {"variable": "counter", "input_mode": "constant", "value": 7, "type": "number"}, }, }) - out, err := c.Invoke(context.Background(), map[string]any{"in": 1}) + out, err := c.Invoke(context.Background(), nil, map[string]any{"in": 1}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -97,7 +97,7 @@ func TestLoop_InvokeIsNoOp(t *testing.T) { // empty-map chunk and closes. func TestLoop_StreamMirrorsInvoke(t *testing.T) { c := NewLoopComponent(loopParam{}) - ch, err := c.Stream(context.Background(), nil) + ch, err := c.Stream(context.Background(), nil, nil) if err != nil { t.Fatalf("Stream: %v", err) } diff --git a/internal/agent/component/message.go b/internal/agent/component/message.go index 58e4e542f6..8b21ea146c 100644 --- a/internal/agent/component/message.go +++ b/internal/agent/component/message.go @@ -38,6 +38,8 @@ import ( "ragflow/internal/agent/audio" "ragflow/internal/agent/runtime" "ragflow/internal/common" + + "gorm.io/gorm" ) const componentNameMessage = "Message" @@ -169,7 +171,7 @@ func (m *MessageComponent) Name() string { return m.name } // inputs["text"] takes precedence over the per-instance text so the // same node can be reused with different templates at run time when // the orchestrator wants to override the DSL-declared value. -func (m *MessageComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (m *MessageComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) if err != nil { return nil, fmt.Errorf("Message: %w", err) @@ -506,11 +508,11 @@ func stringFromStateSys(state *runtime.CanvasState, key string) string { // Stream resolves the message and emits the content chunk. The outer // Agent SSE handler owns the final [DONE] frame, matching Python's // agent_api.py rather than leaking a component-local done marker. -func (m *MessageComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { +func (m *MessageComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { ch := make(chan map[string]any, 16) go func() { defer close(ch) - result, err := m.Invoke(ctx, inputs) + result, err := m.Invoke(ctx, db, inputs) if err != nil { select { case ch <- map[string]any{"error": err.Error()}: diff --git a/internal/agent/component/message_phase8b_test.go b/internal/agent/component/message_phase8b_test.go index 2064b5bcfa..cadda9ac29 100644 --- a/internal/agent/component/message_phase8b_test.go +++ b/internal/agent/component/message_phase8b_test.go @@ -51,7 +51,7 @@ func TestMessage_OutputFormatParam(t *testing.T) { state := canvas.NewCanvasState("r1", "t1") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{"text": "hello", "stream": false}) + out, err := c.Invoke(ctx, nil, map[string]any{"text": "hello", "stream": false}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -71,7 +71,7 @@ func TestMessage_OutputFormatInputOverride(t *testing.T) { state := canvas.NewCanvasState("r1", "t1") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "text": "hi", "output_format": "plain", }) @@ -97,7 +97,7 @@ func TestMessage_DownloadsExtraction(t *testing.T) { "mime_type": "text/csv", "url": "/dl/d-1", } - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "text": "see attachment", "stream": false, "attach": dl, @@ -127,7 +127,7 @@ func TestMessage_DownloadJSONStringSuppressesContent(t *testing.T) { ctx := withStateForTest(context.Background(), state) downloadJSON := `{"doc_id":"d-1","filename":"report.md","mime_type":"text/markdown","url":"/api/v1/agents/attachments/d-1/download","include_download_info_in_content":true}` - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "text": downloadJSON, "stream": false, }) @@ -161,7 +161,7 @@ func TestMessage_AutoPlay_NoEngine(t *testing.T) { state := canvas.NewCanvasState("r1", "t1") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{"text": "hi", "stream": false}) + out, err := c.Invoke(ctx, nil, map[string]any{"text": "hi", "stream": false}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -194,7 +194,7 @@ func TestMessage_AutoPlay_Success(t *testing.T) { state := canvas.NewCanvasState("r1", "t1") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{"text": "hi", "stream": false}) + out, err := c.Invoke(ctx, nil, map[string]any{"text": "hi", "stream": false}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -226,7 +226,7 @@ func TestMessage_MemorySave_NoService(t *testing.T) { state.Sys["query"] = "what?" ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "text": "hi", "stream": false, "memory_save": true, @@ -259,7 +259,7 @@ func TestMessage_MemorySave_Success(t *testing.T) { state.Sys["query"] = "what?" ctx := withStateForTest(context.Background(), state) - _, err := c.Invoke(ctx, map[string]any{ + _, err := c.Invoke(ctx, nil, map[string]any{ "text": "hi", "stream": false, "memory_save": true, @@ -308,7 +308,7 @@ func TestMessage_MemorySave_FromDSLParams(t *testing.T) { // Inputs simulate what the pipeline actually provides: only upstream // outputs, NO memory_ids or memory_save keys. - _, err := c.Invoke(ctx, map[string]any{ + _, err := c.Invoke(ctx, nil, map[string]any{ "text": "hi", "stream": false, }) @@ -345,7 +345,7 @@ func TestMessage_MemorySave_UserIDVariable(t *testing.T) { state.SetVar("begin", "user_id", "resolved-user-123") ctx := withStateForTest(context.Background(), state) - _, err := c.Invoke(ctx, map[string]any{ + _, err := c.Invoke(ctx, nil, map[string]any{ "text": "hi", "stream": false, "memory_save": true, @@ -378,7 +378,7 @@ func TestMessage_MemorySave_UserIDLiteral(t *testing.T) { state.Sys["query"] = "hello?" ctx := withStateForTest(context.Background(), state) - _, err := c.Invoke(ctx, map[string]any{ + _, err := c.Invoke(ctx, nil, map[string]any{ "text": "hi", "stream": false, "memory_save": true, diff --git a/internal/agent/component/message_test.go b/internal/agent/component/message_test.go index 0fe16d26e3..d22bb744ee 100644 --- a/internal/agent/component/message_test.go +++ b/internal/agent/component/message_test.go @@ -34,7 +34,7 @@ func TestMessage_ResolveTemplate(t *testing.T) { state.Sys["query"] = "world" ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "text": "hello {{sys.query}}", "stream": false, }) @@ -59,7 +59,7 @@ func TestMessage_ResolveListReferenceAsJSON(t *testing.T) { state.SetVar("list_0", "result", []any{"user: 1"}) ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "text": "{{list_0@result}}", "stream": false, }) @@ -80,7 +80,7 @@ func TestMessage_Stream(t *testing.T) { state.Sys["query"] = "alice" ctx := withStateForTest(context.Background(), state) - ch, err := c.Stream(ctx, map[string]any{ + ch, err := c.Stream(ctx, nil, map[string]any{ "text": "hi", "stream": true, }) @@ -109,7 +109,7 @@ func TestMessage_NoTemplate(t *testing.T) { state := canvas.NewCanvasState("run-3", "task-3") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{"text": "no refs here", "stream": false}) + out, err := c.Invoke(ctx, nil, map[string]any{"text": "no refs here", "stream": false}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -123,7 +123,7 @@ func TestMessage_RuntimeContentInput(t *testing.T) { state := canvas.NewCanvasState("run-4", "task-4") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{"content": "from upstream", "stream": false}) + out, err := c.Invoke(ctx, nil, map[string]any{"content": "from upstream", "stream": false}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -146,7 +146,7 @@ func TestMessage_EmitsAgentMessage(t *testing.T) { } }) - out, err := c.Invoke(ctx, map[string]any{"content": "visible message"}) + out, err := c.Invoke(ctx, nil, map[string]any{"content": "visible message"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -176,7 +176,7 @@ func TestMessage_EmitsDirectCanvasMessage(t *testing.T) { direct = append(direct, content) }) - out, err := c.Invoke(ctx, map[string]any{"content": "direct message"}) + out, err := c.Invoke(ctx, nil, map[string]any{"content": "direct message"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -204,7 +204,7 @@ func TestMessage_NormalTemplateEmitsOnlyRenderedMessage(t *testing.T) { direct = append(direct, content) }) - out, err := c.Invoke(ctx, map[string]any{"text": "hello {{sys.query}}"}) + out, err := c.Invoke(ctx, nil, map[string]any{"text": "hello {{sys.query}}"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -241,7 +241,7 @@ func TestMessage_SkipsEmissionWhenAgentAlreadyStreamed(t *testing.T) { runtime.EmitAgentMessage(ctx, "streamed this", "") emitted = nil // clear setup emission so we only capture Message's output - out, err := c.Invoke(ctx, map[string]any{"content": "agent streamed this"}) + out, err := c.Invoke(ctx, nil, map[string]any{"content": "agent streamed this"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -276,7 +276,7 @@ func TestMessage_EmitsContentDifferentFromAgentStream(t *testing.T) { runtime.EmitAgentMessage(ctx, "draft answer", "") emitted = nil - out, err := c.Invoke(ctx, map[string]any{"content": "final answer"}) + out, err := c.Invoke(ctx, nil, map[string]any{"content": "final answer"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -309,7 +309,7 @@ func TestMessage_ConsumesDeferredAgentStream(t *testing.T) { } }) - out, err := c.Invoke(ctx, map[string]any{"text": "answer: {{agent_0@content}}"}) + out, err := c.Invoke(ctx, nil, map[string]any{"text": "answer: {{agent_0@content}}"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -344,7 +344,7 @@ func TestMessage_DeferredStreamThinkingEvents(t *testing.T) { } }) - if _, err := c.Invoke(ctx, map[string]any{"text": "{{agent_0@content}}"}); err != nil { + if _, err := c.Invoke(ctx, nil, map[string]any{"text": "{{agent_0@content}}"}); err != nil { t.Fatalf("Invoke: %v", err) } if got, want := strings.Join(events, "|"), "start|plan|end|answer"; got != want { @@ -367,7 +367,7 @@ func TestMessage_DeferredStreamUsesCompletedContent(t *testing.T) { streamed = append(streamed, content) }) - out, err := c.Invoke(ctx, map[string]any{"text": "{{agent_0@content}}"}) + out, err := c.Invoke(ctx, nil, map[string]any{"text": "{{agent_0@content}}"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -387,7 +387,7 @@ func TestMessage_FormalizedContentFallback(t *testing.T) { state := canvas.NewCanvasState("run-5", "task-5") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "formalized_content": "retrieved answer", "_created_time": "2026-07-15T00:00:00Z", "_elapsed_time": 0.01, @@ -405,7 +405,7 @@ func TestMessage_SingleStringFallback(t *testing.T) { state := canvas.NewCanvasState("run-6", "task-6") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "value": "single upstream text", "_elapsed_time": 0.01, }) diff --git a/internal/agent/component/parallel.go b/internal/agent/component/parallel.go index 89cc2386f5..f934dd1abd 100644 --- a/internal/agent/component/parallel.go +++ b/internal/agent/component/parallel.go @@ -50,6 +50,8 @@ package component import ( "context" + + "gorm.io/gorm" ) const componentNameParallel = "Parallel" @@ -148,13 +150,13 @@ func (c *ParallelComponent) Outputs() map[string]string { // The returned map is empty. State writes from this method would be // silently dropped by the eino graph, because ParallelComponent is not // registered as an eino node when the macro expansion fires. -func (c *ParallelComponent) Invoke(_ context.Context, _ map[string]any) (map[string]any, error) { +func (c *ParallelComponent) Invoke(_ context.Context, _ *gorm.DB, _ map[string]any) (map[string]any, error) { return map[string]any{}, nil } // Stream mirrors Invoke and emits an empty map as a single chunk. -func (c *ParallelComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := c.Invoke(ctx, inputs) +func (c *ParallelComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := c.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/parallel_test.go b/internal/agent/component/parallel_test.go index 31ec3e1d2b..507966854a 100644 --- a/internal/agent/component/parallel_test.go +++ b/internal/agent/component/parallel_test.go @@ -75,7 +75,7 @@ func TestParallel_InvokeIsNoOp(t *testing.T) { ItemsRef: "sys.arr", MaxConcurrency: 3, }) - out, err := c.Invoke(context.Background(), map[string]any{"in": 1}) + out, err := c.Invoke(context.Background(), nil, map[string]any{"in": 1}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -88,7 +88,7 @@ func TestParallel_InvokeIsNoOp(t *testing.T) { // empty-map chunk and closes. func TestParallel_StreamMirrorsInvoke(t *testing.T) { c := NewParallelComponent(ParallelParam{}) - ch, err := c.Stream(context.Background(), nil) + ch, err := c.Stream(context.Background(), nil, nil) if err != nil { t.Fatalf("Stream: %v", err) } diff --git a/internal/agent/component/production_chain_fixes_test.go b/internal/agent/component/production_chain_fixes_test.go index 6f4ac1fc25..8336a84253 100644 --- a/internal/agent/component/production_chain_fixes_test.go +++ b/internal/agent/component/production_chain_fixes_test.go @@ -214,7 +214,7 @@ func TestRetrieval_LegacyQueryStringNormalized(t *testing.T) { }) state := runtime.NewCanvasState("run-1", "task-1") state.Sys["user_id"] = "user-1" - normalizeLegacyRetrievalInputs(runtime.WithState(context.Background(), state), merged) + normalizeLegacyRetrievalInputs(runtime.WithState(context.Background(), state), db, merged) if got, _ := merged["query"].(string); got != "diamond necklace" { t.Fatalf("query = %q, want diamond necklace", got) @@ -235,7 +235,7 @@ func TestRetrieval_StructuredUserFillInputNormalized(t *testing.T) { t.Fatalf("failed to unwrap sql db: %v", err) } sqlDB.SetMaxOpenConns(1) - if err := db.AutoMigrate(&entity.Knowledgebase{}, &entity.UserTenant{}); err != nil { + if err = db.AutoMigrate(&entity.Knowledgebase{}, &entity.UserTenant{}); err != nil { t.Fatalf("failed to migrate tables: %v", err) } origDB := dao.DB @@ -243,7 +243,7 @@ func TestRetrieval_StructuredUserFillInputNormalized(t *testing.T) { t.Cleanup(func() { dao.DB = origDB }) activeStatus := "1" - if err := db.Create(&entity.UserTenant{ + if err = db.Create(&entity.UserTenant{ ID: "ut-1", UserID: "user-1", TenantID: "tenant-1", @@ -253,7 +253,7 @@ func TestRetrieval_StructuredUserFillInputNormalized(t *testing.T) { }).Error; err != nil { t.Fatalf("failed to seed user_tenant: %v", err) } - if err := db.Create(&entity.Knowledgebase{ + if err = db.Create(&entity.Knowledgebase{ ID: "kb-da1", Name: "da1", TenantID: "tenant-1", @@ -280,7 +280,7 @@ func TestRetrieval_StructuredUserFillInputNormalized(t *testing.T) { }) state := runtime.NewCanvasState("run-1", "task-1") state.Sys["user_id"] = "user-1" - normalizeLegacyRetrievalInputs(runtime.WithState(context.Background(), state), merged) + normalizeLegacyRetrievalInputs(runtime.WithState(context.Background(), state), db, merged) if got, _ := merged["query"].(string); got != "合同" { t.Fatalf("query = %q, want 合同", got) @@ -324,7 +324,7 @@ func TestRetrieval_ResolveDatasetIDByTenantName(t *testing.T) { state.Sys["tenant_id"] = "tenant-1" ctx := runtime.WithState(context.Background(), state) - if got := resolveRetrievalDatasetID(ctx, "da1"); got != "kb-da1" { + if got := resolveRetrievalDatasetID(ctx, db, "da1"); got != "kb-da1" { t.Fatalf("resolveRetrievalDatasetID = %q, want kb-da1", got) } } @@ -345,7 +345,7 @@ func TestRetrieval_StructuredInputPreservesQueryWhenDatasetIDsAlreadyPresent(t * }, }) - consumed := normalizeStructuredRetrievalInputs(context.Background(), merged) + consumed := normalizeStructuredRetrievalInputs(context.Background(), nil, merged) if !consumed { t.Fatal("normalizeStructuredRetrievalInputs should consume structured query") } @@ -387,7 +387,7 @@ func TestRetrieval_KbIDsEndToEndThroughTool(t *testing.T) { if err != nil { t.Fatalf("New(Retrieval): %v", err) } - out, err := wrapper.Invoke(context.Background(), map[string]any{"query": "ragflow"}) + out, err := wrapper.Invoke(context.Background(), nil, map[string]any{"query": "ragflow"}) if err != nil { t.Fatalf("Retrieval.Invoke: %v", err) } @@ -401,7 +401,7 @@ func TestRetrieval_KbIDsEndToEndThroughTool(t *testing.T) { // variants must also resolve to the Universe B tool. The Universe // B tool registry has always // accepted all three spellings (search_my_dataset / -// search_my_dateset) plus PascalCase; Universe A now mirrors that +// search_my_dataset) plus PascalCase; Universe A now mirrors that // surface so older DSLs don't fail with "unknown component" at // buildNodeBody time. // @@ -501,7 +501,7 @@ func TestCodeExec_LegacyDSLWrapperBridgesParamsAndOutputs(t *testing.T) { t.Fatalf("New(CodeExec): %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "arguments": map[string]any{ "x": 7, }, @@ -573,7 +573,7 @@ func TestCodeExec_LegacyDSLWrapperResolvesArgumentRefsFromState(t *testing.T) { t.Fatalf("New(CodeExec): %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "state": map[string]map[string]any{ "UserFillUp:CodeInput": { "x": "8", @@ -625,7 +625,7 @@ func TestCodeExec_LegacyDSLWrapperResolvesSysArgumentRefsFromCanvasState(t *test state := runtime.NewCanvasState("run-codeexec", "task-codeexec") state.Sys["query"] = "532" ctx := runtime.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("CodeExec.Invoke: %v", err) } @@ -664,7 +664,7 @@ func TestCodeExec_LegacyDSLWrapperContractMismatchSetsError(t *testing.T) { t.Fatalf("New(CodeExec): %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("CodeExec.Invoke: %v", err) } @@ -704,7 +704,7 @@ func TestCodeExec_LegacyDSLWrapperPreservesExecutionError(t *testing.T) { t.Fatalf("New(CodeExec): %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(context.Background(), nil, map[string]any{}) if err == nil { t.Fatal("CodeExec.Invoke: want wrapped execution error, got nil") } @@ -723,7 +723,7 @@ type countingInvoker struct { calls int } -func (c *countingInvoker) Invoke(_ context.Context, _ ChatInvokeRequest) (*ChatInvokeResponse, error) { +func (c *countingInvoker) Invoke(_ context.Context, _ *gorm.DB, _ ChatInvokeRequest) (*ChatInvokeResponse, error) { c.calls++ return nil, errLLMRetryTestAlwaysFail } @@ -796,7 +796,7 @@ func TestLLM_RetryStackingSemantics(t *testing.T) { } // The retry chain always returns errLLMRetryTestAlwaysFail // so we can count attempts deterministically. - _, _ = comp.Invoke(context.Background(), nil) + _, _ = comp.Invoke(context.Background(), nil, nil) if counter.calls < tc.wantMinAttempts { t.Errorf("counter.calls = %d, want >= %d (MaxRetries=%d stacking regression?)", counter.calls, tc.wantMinAttempts, tc.maxRetries) diff --git a/internal/agent/component/retrieval_swap_test.go b/internal/agent/component/retrieval_swap_test.go index 3417a48409..a3691ee194 100644 --- a/internal/agent/component/retrieval_swap_test.go +++ b/internal/agent/component/retrieval_swap_test.go @@ -55,7 +55,7 @@ func TestRetrieval_DelegatesToRealWrapper(t *testing.T) { t.Errorf("Retrieval c.Name() = %q, want %q", got, componentNameRetrieval) } - out, err := c.Invoke(context.Background(), map[string]any{"query": "ragflow"}) + out, err := c.Invoke(context.Background(), nil, map[string]any{"query": "ragflow"}) if err != nil { t.Fatalf("Retrieval Invoke errored: %v", err) } @@ -87,7 +87,7 @@ func TestSearchMyDataset_AliasDelegatesToRealWrapper(t *testing.T) { t.Errorf("SearchMyDataset alias c.Name() = %q, want %q (canonical)", got, componentNameRetrieval) } - out, err := c.Invoke(context.Background(), map[string]any{"query": "kb"}) + out, err := c.Invoke(context.Background(), nil, map[string]any{"query": "kb"}) if err != nil { t.Fatalf("SearchMyDataset Invoke errored: %v", err) } diff --git a/internal/agent/component/sampler_params_test.go b/internal/agent/component/sampler_params_test.go index 1110000acc..90284b6233 100644 --- a/internal/agent/component/sampler_params_test.go +++ b/internal/agent/component/sampler_params_test.go @@ -35,7 +35,7 @@ func TestLLM_ForwardsTopP(t *testing.T) { ModelID: "echo", TopP: &topP, }) - if _, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "hi"}); err != nil { + if _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "hi"}); err != nil { t.Fatalf("Invoke: %v", err) } if stub.calls != 1 { @@ -59,7 +59,7 @@ func TestLLM_TopPFromInputs(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo"}) - if _, err := c.Invoke(context.Background(), map[string]any{ + if _, err := c.Invoke(context.Background(), nil, map[string]any{ "user_prompt": "hi", "top_p": 0.7, }); err != nil { @@ -80,7 +80,7 @@ func TestLLM_NoTopPByDefault(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo"}) - if _, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "hi"}); err != nil { + if _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "hi"}); err != nil { t.Fatalf("Invoke: %v", err) } if stub.captured == nil { @@ -135,7 +135,7 @@ func TestAgentParam_ForwardsTopP(t *testing.T) { TopP: &topP, MaxRounds: 1, }) - if _, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "hi"}); err != nil { + if _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "hi"}); err != nil { t.Fatalf("Invoke: %v", err) } } @@ -150,7 +150,7 @@ func TestAgent_TopPFromInputs(t *testing.T) { }) c := NewAgentComponent(AgentParam{ModelID: "echo", MaxRounds: 1}) - if _, err := c.Invoke(context.Background(), map[string]any{ + if _, err := c.Invoke(context.Background(), nil, map[string]any{ "user_prompt": "hi", "top_p": 0.42, }); err != nil { diff --git a/internal/agent/component/stagehand_runtime_integration_test.go b/internal/agent/component/stagehand_runtime_integration_test.go index 01c532d381..f9e18ba2fb 100644 --- a/internal/agent/component/stagehand_runtime_integration_test.go +++ b/internal/agent/component/stagehand_runtime_integration_test.go @@ -242,7 +242,7 @@ func TestBrowser_E2E_Extract(t *testing.T) { invokeCtx, cancel := context.WithTimeout(ctx, 3*time.Minute) defer cancel() t.Logf("starting browser.Invoke (RunExtract) against %s (timeout 3m)", srv.URL) - out, err := c.Invoke(invokeCtx, nil) + out, err := c.Invoke(invokeCtx, nil, nil) if err != nil { t.Logf("extraction failed (best-effort): %v", err) return diff --git a/internal/agent/component/streaming_test.go b/internal/agent/component/streaming_test.go index 116e896062..3128c673d9 100644 --- a/internal/agent/component/streaming_test.go +++ b/internal/agent/component/streaming_test.go @@ -29,7 +29,7 @@ func TestLLM_Stream_HappyPath(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo"}) - ch, err := c.Stream(context.Background(), map[string]any{"user_prompt": "hi"}) + ch, err := c.Stream(context.Background(), nil, map[string]any{"user_prompt": "hi"}) if err != nil { t.Fatalf("Stream: %v", err) } @@ -63,7 +63,7 @@ func TestLLM_Stream_Error(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo"}) - ch, err := c.Stream(context.Background(), map[string]any{"user_prompt": "hi"}) + ch, err := c.Stream(context.Background(), nil, map[string]any{"user_prompt": "hi"}) if err != nil { t.Fatalf("Stream: %v", err) } @@ -89,7 +89,7 @@ func TestLLM_Stream_RespectsCancellation(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() // pre-cancel - ch, err := c.Stream(ctx, map[string]any{"user_prompt": "hi"}) + ch, err := c.Stream(ctx, nil, map[string]any{"user_prompt": "hi"}) if err != nil { t.Fatalf("Stream: %v", err) } @@ -113,7 +113,7 @@ func TestLLM_Stream_BufferDoesNotBlock(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo"}) - ch, err := c.Stream(context.Background(), map[string]any{"user_prompt": "hi"}) + ch, err := c.Stream(context.Background(), nil, map[string]any{"user_prompt": "hi"}) if err != nil { t.Fatalf("Stream: %v", err) } diff --git a/internal/agent/component/string_transform.go b/internal/agent/component/string_transform.go index 94a72284cf..a9503bbf37 100644 --- a/internal/agent/component/string_transform.go +++ b/internal/agent/component/string_transform.go @@ -34,6 +34,8 @@ import ( "strings" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) const componentNameStringTransform = "StringTransform" @@ -137,7 +139,7 @@ func (s *StringTransformComponent) Name() string { return s.name } // Invoke runs the configured method (split or merge) and returns // outputs["result"] with the transformed payload. -func (s *StringTransformComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (s *StringTransformComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) if err != nil { return nil, fmt.Errorf("StringTransform: %w", err) @@ -153,8 +155,8 @@ func (s *StringTransformComponent) Invoke(ctx context.Context, inputs map[string } // Stream mirrors Invoke; StringTransform is a single-shot transform. -func (s *StringTransformComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := s.Invoke(ctx, inputs) +func (s *StringTransformComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := s.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/string_transform_test.go b/internal/agent/component/string_transform_test.go index 43cc039e39..3f2cde17a8 100644 --- a/internal/agent/component/string_transform_test.go +++ b/internal/agent/component/string_transform_test.go @@ -36,7 +36,7 @@ func TestStringTransform_SplitBasic(t *testing.T) { state := canvas.NewCanvasState("run-1", "task-1") ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{"line": "a,b;c"}) + out, err := c.Invoke(ctx, nil, map[string]any{"line": "a,b;c"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -56,7 +56,7 @@ func TestStringTransform_SplitNoDelim(t *testing.T) { state := canvas.NewCanvasState("run-2", "task-2") ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{"line": "abc"}) + out, err := c.Invoke(ctx, nil, map[string]any{"line": "abc"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -77,7 +77,7 @@ func TestStringTransform_Merge(t *testing.T) { state := canvas.NewCanvasState("run-3", "task-3") ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{"x": "foo", "y": "bar"}) + out, err := c.Invoke(ctx, nil, map[string]any{"x": "foo", "y": "bar"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -97,7 +97,7 @@ func TestStringTransform_MergeIterationAliases(t *testing.T) { state.Globals["__index__"] = 1 ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{}) + out, err := c.Invoke(ctx, nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -118,7 +118,7 @@ func TestStringTransform_SplitFromStateRef(t *testing.T) { state.Outputs["cpn_0"] = map[string]any{"x": "alpha,beta,gamma"} ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -140,7 +140,7 @@ func TestStringTransform_MergeMissingPlaceholder(t *testing.T) { state := canvas.NewCanvasState("run-5", "task-5") ctx := canvas.WithState(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{}) + out, err := c.Invoke(ctx, nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } diff --git a/internal/agent/component/structured_output_test.go b/internal/agent/component/structured_output_test.go index f47de3f121..1e30575158 100644 --- a/internal/agent/component/structured_output_test.go +++ b/internal/agent/component/structured_output_test.go @@ -20,6 +20,8 @@ import ( "context" "strings" "testing" + + "gorm.io/gorm" ) // TestMatchOutputStructure_ValidJSON_AllKeysPresent: response has all @@ -105,7 +107,7 @@ func TestLLM_Invoke_OutputStructure_ValidFirstTry(t *testing.T) { ModelID: "echo", OutputStructure: map[string]any{"name": "", "age": 0}, }) - out, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "who?"}) + out, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "who?"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -139,7 +141,7 @@ func TestLLM_Invoke_OutputStructure_RetryOnInvalid(t *testing.T) { ModelID: "echo", OutputStructure: map[string]any{"name": ""}, }) - out, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "who?"}) + out, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "who?"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -176,7 +178,7 @@ func TestLLM_Invoke_OutputStructure_RetryStillFails(t *testing.T) { ModelID: "echo", OutputStructure: map[string]any{"x": 0}, }) - out, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "go"}) + out, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "go"}) if err != nil { t.Fatalf("Invoke should not error on parse failure: %v", err) } @@ -200,7 +202,7 @@ type callCountingInvoker struct { calls int } -func (c *callCountingInvoker) Invoke(_ context.Context, _ ChatInvokeRequest) (*ChatInvokeResponse, error) { +func (c *callCountingInvoker) Invoke(_ context.Context, _ *gorm.DB, _ ChatInvokeRequest) (*ChatInvokeResponse, error) { if c.onCall != nil { c.onCall() } diff --git a/internal/agent/component/switch.go b/internal/agent/component/switch.go index d594e5cc1e..e804938c12 100644 --- a/internal/agent/component/switch.go +++ b/internal/agent/component/switch.go @@ -40,6 +40,8 @@ import ( "strings" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) const componentNameSwitch = "Switch" @@ -78,7 +80,7 @@ func (s *SwitchComponent) Name() string { return s.name } // returning "_next" as a []any (list of cpn_ids). The canvas // scheduler's MultiBranch condition consumes this list via // NewGraphMultiBranch so every declared target fires. -func (s *SwitchComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (s *SwitchComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) if err != nil { return nil, fmt.Errorf("Switch: %w", err) @@ -135,8 +137,8 @@ func (s *SwitchComponent) Invoke(ctx context.Context, inputs map[string]any) (ma // Stream is a synchronous facade over Invoke for P0. Switch is a // routing decision, not a stream of partial results; the channel // receives one payload and closes. -func (s *SwitchComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := s.Invoke(ctx, inputs) +func (s *SwitchComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := s.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/switch_test.go b/internal/agent/component/switch_test.go index f526f30a50..773d5e3bc6 100644 --- a/internal/agent/component/switch_test.go +++ b/internal/agent/component/switch_test.go @@ -69,7 +69,7 @@ func TestSwitch_AndMatches(t *testing.T) { }, "default": "fallback", } - out, err := s.Invoke(ctx, inputs) + out, err := s.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -113,7 +113,7 @@ func TestSwitch_OrMatches(t *testing.T) { }, "default": "fallback", } - out, err := s.Invoke(ctx, inputs) + out, err := s.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -142,7 +142,7 @@ func TestSwitch_DefaultFallback(t *testing.T) { }, "default": "fallback_0", } - out, err := s.Invoke(ctx, inputs) + out, err := s.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -169,7 +169,7 @@ func TestSwitch_LegacyEndCpnIDsFallback(t *testing.T) { }, "end_cpn_ids": []any{"legacy_fallback"}, } - out, err := s.Invoke(ctx, inputs) + out, err := s.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -200,7 +200,7 @@ func TestSwitch_ContainsAndEmpty(t *testing.T) { }, "default": "x", } - out, err := s.Invoke(ctx, inputs) + out, err := s.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -231,7 +231,7 @@ func TestSwitch_LegacyConditionsAndArrayTo(t *testing.T) { state.SetVar("UserFillUp:Menu", "demo", "loop") ctx := withStateForTest(context.Background(), state) - out, err := s.Invoke(ctx, nil) + out, err := s.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -267,7 +267,7 @@ func TestSwitch_NilUpstreamContainsEmptyNeedleMatches(t *testing.T) { }, "default": "else_target", } - out, err := s.Invoke(ctx, inputs) + out, err := s.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -298,7 +298,7 @@ func TestSwitch_NilUpstreamContainsNonEmptyDoesNotMatch(t *testing.T) { }, "default": "else_target", } - out, err := s.Invoke(ctx, inputs) + out, err := s.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -330,7 +330,7 @@ func TestSwitch_NilValueContainsDoesNotRaise(t *testing.T) { }, "default": "else_target", } - out, err := s.Invoke(ctx, inputs) + out, err := s.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -372,7 +372,7 @@ func TestSwitch_NilUpstreamStartWithEndWithDoNotCrash(t *testing.T) { }, "default": "else_target", } - out, err := s.Invoke(ctx, inputs) + out, err := s.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -411,7 +411,7 @@ func TestSwitch_MultiTargetTo(t *testing.T) { }, "default": "Message:Help", } - out, err := s.Invoke(ctx, inputs) + out, err := s.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -443,7 +443,7 @@ func TestSwitch_EmptyAndConditionFallsThrough(t *testing.T) { }, "default": "DEFAULT", } - out, err := s.Invoke(ctx, inputs) + out, err := s.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -476,7 +476,7 @@ func TestSwitch_LegacyEmptyItemsFallsThrough(t *testing.T) { }, "default": "DEFAULT", } - out, err := s.Invoke(ctx, inputs) + out, err := s.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -510,7 +510,7 @@ func TestSwitch_SatisfiedAndConditionStillRoutes(t *testing.T) { }, "default": "DEFAULT", } - out, err := s.Invoke(ctx, inputs) + out, err := s.Invoke(ctx, nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } diff --git a/internal/agent/component/tool_call_memory_test.go b/internal/agent/component/tool_call_memory_test.go index 0d9a808efe..30783db5e1 100644 --- a/internal/agent/component/tool_call_memory_test.go +++ b/internal/agent/component/tool_call_memory_test.go @@ -32,7 +32,7 @@ func TestAddToolCallMemory_NoToolCalls(t *testing.T) { stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}} withStubInvoker(t, stub) - got, err := addToolCallMemory(context.Background(), AgentParam{ModelID: "echo"}, &schema.Message{Content: "no tools"}) + got, err := addToolCallMemory(context.Background(), nil, AgentParam{ModelID: "echo"}, &schema.Message{Content: "no tools"}) if err != nil { t.Fatalf("err: %v", err) } @@ -71,7 +71,7 @@ func TestAddToolCallMemory_SummarizesAndAppendsToState(t *testing.T) { state := runtime.NewCanvasState("rid", "tid") c := NewAgentComponent(AgentParam{ModelID: "echo", MaxRounds: 1}) ctx := runtime.WithState(context.Background(), state) - _, err := c.Invoke(ctx, map[string]any{"user_prompt": "do it"}) + _, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "do it"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -113,7 +113,7 @@ func TestAddToolCallMemory_LLMFailure(t *testing.T) { state := runtime.NewCanvasState("rid", "tid") c := NewAgentComponent(AgentParam{ModelID: "echo", MaxRounds: 1}) ctx := runtime.WithState(context.Background(), state) - _, err := c.Invoke(ctx, map[string]any{"user_prompt": "do it"}) + _, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "do it"}) if err != nil { t.Fatalf("Invoke should not error when memory summary fails: %v", err) } diff --git a/internal/agent/component/tool_component.go b/internal/agent/component/tool_component.go index 79e6be2beb..f4b5919a36 100644 --- a/internal/agent/component/tool_component.go +++ b/internal/agent/component/tool_component.go @@ -24,6 +24,8 @@ import ( "ragflow/internal/agent/runtime" agenttool "ragflow/internal/agent/tool" + + "gorm.io/gorm" ) // ToolBackedComponent is the single Canvas adapter for tools that implement @@ -60,7 +62,7 @@ func (c *ToolBackedComponent) Outputs() map[string]string { return c.spec.Output func (c *ToolBackedComponent) GetInputForm() map[string]any { return c.spec.InputForm } -func (c *ToolBackedComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *ToolBackedComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { argsJSON, err := json.Marshal(inputs) if err != nil { return nil, fmt.Errorf("canvas: %s: encode inputs: %w", c.name, err) @@ -95,7 +97,7 @@ func (c *ToolBackedComponent) Invoke(ctx context.Context, inputs map[string]any) return c.tool.BuildComponentOutputs(decoded), nil } -func (c *ToolBackedComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { +func (c *ToolBackedComponent) Stream(_ context.Context, _ *gorm.DB, _ map[string]any) (<-chan map[string]any, error) { return nil, nil } diff --git a/internal/agent/component/tool_component_test.go b/internal/agent/component/tool_component_test.go index 221607766f..c2b0b9caf4 100644 --- a/internal/agent/component/tool_component_test.go +++ b/internal/agent/component/tool_component_test.go @@ -262,7 +262,7 @@ func TestToolBackedComponentWenCaiInvoke(t *testing.T) { if err != nil { t.Fatalf("New(WenCai): %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "query": "商业航天", "unrelated": "ignored by the tool parameter struct", }) @@ -320,7 +320,7 @@ func TestToolBackedComponentInvokeOrdersReferencesBeforeRendering(t *testing.T) fake := &fakeToolAdapter{out: `{"results":[{"title":"RAGFlow","content":"RAG engine"}],"tool_metadata":{"request_id":"request-1"}}`} c := &ToolBackedComponent{name: "Search", tool: fake, spec: fake.ComponentSpec()} state := runtime.NewCanvasState("run", "task") - out, err := c.Invoke(runtime.WithState(context.Background(), state), map[string]any{ + out, err := c.Invoke(runtime.WithState(context.Background(), state), nil, map[string]any{ "query": "ragflow", "top_n": float64(10), "outputs": map[string]any{"json": map[string]any{}}, @@ -361,7 +361,7 @@ func TestToolBackedComponentReturnsErrorEnvelopeWithoutReferences(t *testing.T) } c := &ToolBackedComponent{name: "Search", tool: fake, spec: fake.ComponentSpec()} state := runtime.NewCanvasState("run-error", "task-error") - out, err := c.Invoke(runtime.WithState(context.Background(), state), map[string]any{"query": "ragflow"}) + out, err := c.Invoke(runtime.WithState(context.Background(), state), nil, map[string]any{"query": "ragflow"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -397,14 +397,14 @@ func TestToolBackedComponentGitHubIntegration(t *testing.T) { github := agenttool.NewGitHubToolWith(helper) component := &ToolBackedComponent{name: "GitHub", tool: github, spec: github.ComponentSpec()} state := runtime.NewCanvasState("run-github", "task-github") - empty, err := component.Invoke(context.Background(), map[string]any{"query": ""}) + empty, err := component.Invoke(context.Background(), nil, map[string]any{"query": ""}) if err != nil { t.Fatalf("Invoke(empty query): %v", err) } if serverCalls != 0 || len(empty["json"].([]any)) != 0 || empty["formalized_content"] != "" { t.Fatalf("empty query result = %#v, server calls = %d", empty, serverCalls) } - out, err := component.Invoke(runtime.WithState(context.Background(), state), map[string]any{"query": "ragflow"}) + out, err := component.Invoke(runtime.WithState(context.Background(), state), nil, map[string]any{"query": "ragflow"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -443,14 +443,14 @@ func TestToolBackedComponentTavilyIntegration(t *testing.T) { tavily := agenttool.NewTavilyToolWith(helper) component := &ToolBackedComponent{name: "TavilySearch", tool: tavily, spec: tavily.ComponentSpec()} state := runtime.NewCanvasState("run-tavily", "task-tavily") - empty, err := component.Invoke(context.Background(), map[string]any{"query": ""}) + empty, err := component.Invoke(context.Background(), nil, map[string]any{"query": ""}) if err != nil { t.Fatalf("Invoke(empty query): %v", err) } if serverCalls != 0 || len(empty["json"].([]any)) != 0 || empty["formalized_content"] != "" { t.Fatalf("empty query result = %#v, server calls = %d", empty, serverCalls) } - out, err := component.Invoke(runtime.WithState(context.Background(), state), map[string]any{"query": "ragflow", "api_key": "key"}) + out, err := component.Invoke(runtime.WithState(context.Background(), state), nil, map[string]any{"query": "ragflow", "api_key": "key"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -508,7 +508,7 @@ func TestToolBackedComponentYahooFinanceIntegration(t *testing.T) { yahoo := agenttool.NewYahooFinanceToolWith(helper) component := &ToolBackedComponent{name: "YahooFinance", tool: yahoo, spec: yahoo.ComponentSpec()} - empty, err := component.Invoke(context.Background(), map[string]any{"stock_code": ""}) + empty, err := component.Invoke(context.Background(), nil, map[string]any{"stock_code": ""}) if err != nil { t.Fatalf("Invoke(empty stock_code): %v", err) } @@ -516,7 +516,7 @@ func TestToolBackedComponentYahooFinanceIntegration(t *testing.T) { t.Fatalf("empty result = %#v, server calls = %d", empty, serverCalls) } - out, err := component.Invoke(context.Background(), map[string]any{"stock_code": "AAPL"}) + out, err := component.Invoke(context.Background(), nil, map[string]any{"stock_code": "AAPL"}) if err != nil { t.Fatalf("Invoke: %v", err) } diff --git a/internal/agent/component/tool_dispatch_test.go b/internal/agent/component/tool_dispatch_test.go index 165d265368..2886611a57 100644 --- a/internal/agent/component/tool_dispatch_test.go +++ b/internal/agent/component/tool_dispatch_test.go @@ -69,7 +69,7 @@ func TestPhase3_6_ToolDSLLoading(t *testing.T) { MaxRounds: 1, Tools: []string{"retrieval"}, // known tool }) - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(context.Background(), nil, map[string]any{ "user_prompt": "test", }) if err != nil { diff --git a/internal/agent/component/universe_a_wrappers.go b/internal/agent/component/universe_a_wrappers.go index a2d016e033..9413620450 100644 --- a/internal/agent/component/universe_a_wrappers.go +++ b/internal/agent/component/universe_a_wrappers.go @@ -164,9 +164,9 @@ func (c *retrievalComponent) Outputs() map[string]string { } } -func (c *retrievalComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *retrievalComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { merged := c.applyDefaults(inputs) - normalizeLegacyRetrievalInputs(ctx, merged) + normalizeLegacyRetrievalInputs(ctx, db, merged) common.Debug("agent retrieval component: invoke", zap.Any("inputs", inputs), zap.Any("merged", merged), @@ -182,7 +182,7 @@ func (c *retrievalComponent) Invoke(ctx context.Context, inputs map[string]any) return parseToolEnvelope(out), nil } -func (c *retrievalComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { +func (c *retrievalComponent) Stream(_ context.Context, _ *gorm.DB, _ map[string]any) (<-chan map[string]any, error) { // V1: retrieval is a non-streaming node (the Python // Retrieval component also blocks on Dealer.search). A // streaming retrieval lands with the streaming-dealer @@ -255,8 +255,8 @@ func (c *retrievalComponent) applyDefaults(inputs map[string]any) map[string]any return out } -func normalizeLegacyRetrievalInputs(ctx context.Context, out map[string]any) { - if normalizeStructuredRetrievalInputs(ctx, out) { +func normalizeLegacyRetrievalInputs(ctx context.Context, db *gorm.DB, out map[string]any) { + if normalizeStructuredRetrievalInputs(ctx, db, out) { return } rawQuery, _ := out["query"].(string) @@ -279,12 +279,12 @@ func normalizeLegacyRetrievalInputs(ctx context.Context, out map[string]any) { if kbName == "" { return } - if datasetID := resolveRetrievalDatasetID(ctx, kbName); datasetID != "" { + if datasetID := resolveRetrievalDatasetID(ctx, db, kbName); datasetID != "" { out["dataset_ids"] = []string{datasetID} } } -func normalizeStructuredRetrievalInputs(ctx context.Context, out map[string]any) bool { +func normalizeStructuredRetrievalInputs(ctx context.Context, db *gorm.DB, out map[string]any) bool { _, hasDatasetIDs := out["dataset_ids"] candidateMaps := []map[string]any{} if stateMap, ok := out["state"].(map[string]any); ok { @@ -309,7 +309,7 @@ func normalizeStructuredRetrievalInputs(ctx context.Context, out map[string]any) out["query"] = queryText } if kbName != "" && !hasDatasetIDs { - if datasetID := resolveRetrievalDatasetID(ctx, strings.TrimSpace(kbName)); datasetID != "" { + if datasetID := resolveRetrievalDatasetID(ctx, db, strings.TrimSpace(kbName)); datasetID != "" { out["dataset_ids"] = []string{datasetID} common.Debug("agent retrieval component: resolved dataset id") } @@ -324,11 +324,11 @@ func normalizeStructuredRetrievalInputs(ctx context.Context, out map[string]any) return consumed } -func resolveRetrievalDatasetID(ctx context.Context, kbName string) string { +func resolveRetrievalDatasetID(ctx context.Context, db *gorm.DB, kbName string) string { if kbName == "" { return "" } - if kb, err := dao.NewKnowledgebaseDAO().GetByID(kbName); err == nil && kb != nil { + if kb, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, kbName); err == nil && kb != nil { common.Debug("agent retrieval component: resolved dataset id by direct id") return kb.ID } else if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { @@ -338,7 +338,7 @@ func resolveRetrievalDatasetID(ctx context.Context, kbName string) string { if state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx); err == nil && state != nil { common.Debug("agent retrieval component: resolve dataset id context") if tenantID, _ := state.Sys["tenant_id"].(string); tenantID != "" { - if kb, lookupErr := dao.NewKnowledgebaseDAO().GetByName(kbName, tenantID); lookupErr == nil && kb != nil { + if kb, lookupErr := dao.NewKnowledgebaseDAO().GetByName(ctx, db, kbName, tenantID); lookupErr == nil && kb != nil { common.Debug("agent retrieval component: resolved dataset id by tenant") return kb.ID } else if lookupErr != nil && !errors.Is(lookupErr, gorm.ErrRecordNotFound) { @@ -349,7 +349,7 @@ func resolveRetrievalDatasetID(ctx context.Context, kbName string) string { } } if userID, _ := state.Sys["user_id"].(string); userID != "" { - if kbs, lookupErr := dao.NewKnowledgebaseDAO().GetKBByNameAndUserID(kbName, userID); lookupErr == nil && len(kbs) > 0 { + if kbs, lookupErr := dao.NewKnowledgebaseDAO().GetKBByNameAndUserID(ctx, db, kbName, userID); lookupErr == nil && len(kbs) > 0 { for _, kb := range kbs { if kb == nil || kb.Status == nil || *kb.Status != string(entity.StatusValid) { continue @@ -429,7 +429,7 @@ func (c *codeExecComponent) Outputs() map[string]string { } } -func (c *codeExecComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *codeExecComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { merged := make(map[string]any, len(c.params)+len(inputs)) for k, v := range c.params { merged[k] = v @@ -468,7 +468,7 @@ func (c *codeExecComponent) Invoke(ctx context.Context, inputs map[string]any) ( return decoded, nil } -func (c *codeExecComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { +func (c *codeExecComponent) Stream(_ context.Context, _ *gorm.DB, _ map[string]any) (<-chan map[string]any, error) { return nil, nil } diff --git a/internal/agent/component/userfillup.go b/internal/agent/component/userfillup.go index dc03ca7fa2..7ef3d5bafe 100644 --- a/internal/agent/component/userfillup.go +++ b/internal/agent/component/userfillup.go @@ -33,6 +33,8 @@ import ( "strings" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) const componentNameUserFillUp = "UserFillUp" @@ -115,7 +117,7 @@ func (u *UserFillUpComponent) Name() string { return u.name } // Invoke renders the tips template (when enable_tips) and emits one // output per form field. Inputs are expected under the top-level // "inputs" key. -func (u *UserFillUpComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (u *UserFillUpComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { // State is required for the canvas ref grammar, but UserFillUp's // tips substitution uses simple {{key}} placeholders resolved // against the form input map. We still extract state so a @@ -141,8 +143,8 @@ func (u *UserFillUpComponent) Invoke(ctx context.Context, inputs map[string]any) // Stream is the synchronous facade over Invoke: a single payload, then // close. SSE streaming of the rendered tips is not meaningful for the // P3 port (form-fill is a one-shot interaction in the DSL). -func (u *UserFillUpComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := u.Invoke(ctx, inputs) +func (u *UserFillUpComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := u.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/userfillup_test.go b/internal/agent/component/userfillup_test.go index c743552dc3..66ee37d23b 100644 --- a/internal/agent/component/userfillup_test.go +++ b/internal/agent/component/userfillup_test.go @@ -35,7 +35,7 @@ func TestUserFillUp_RendersTips(t *testing.T) { state := canvas.NewCanvasState("run-1", "task-1") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "inputs": map[string]any{ "name": map[string]any{"value": "World"}, }, @@ -59,7 +59,7 @@ func TestUserFillUp_DisableTips(t *testing.T) { state := canvas.NewCanvasState("run-2", "task-2") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "inputs": map[string]any{ "name": map[string]any{"value": "World"}, }, @@ -85,7 +85,7 @@ func TestUserFillUp_PassesThroughInputs(t *testing.T) { state := canvas.NewCanvasState("run-3", "task-3") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "inputs": map[string]any{ "q": map[string]any{"value": "What is RAGFlow?"}, "top_k": map[string]any{"value": 5}, @@ -118,7 +118,7 @@ func TestUserFillUp_FileInputStub(t *testing.T) { state := canvas.NewCanvasState("run-4", "task-4") ctx := withStateForTest(context.Background(), state) - out, err := c.Invoke(ctx, map[string]any{ + out, err := c.Invoke(ctx, nil, map[string]any{ "inputs": map[string]any{ "cv": map[string]any{ "value": []any{"file-1", "file-2"}, diff --git a/internal/agent/component/variable_aggregator.go b/internal/agent/component/variable_aggregator.go index 9955998fae..263529eda5 100644 --- a/internal/agent/component/variable_aggregator.go +++ b/internal/agent/component/variable_aggregator.go @@ -32,6 +32,8 @@ import ( "fmt" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) const componentNameVariableAggregator = "VariableAggregator" @@ -158,7 +160,7 @@ func (v *VariableAggregatorComponent) Name() string { return v.name } // contract: the engine is allowed to pass the resolved variable list // per-invocation. When inputs["variables"] is absent the static param // config is used unchanged. -func (v *VariableAggregatorComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (v *VariableAggregatorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) if err != nil { return nil, fmt.Errorf("VariableAggregator: %w", err) @@ -230,8 +232,8 @@ func (v *VariableAggregatorComponent) Invoke(ctx context.Context, inputs map[str } // Stream mirrors Invoke; VariableAggregator is a single-shot reduce. -func (v *VariableAggregatorComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := v.Invoke(ctx, inputs) +func (v *VariableAggregatorComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := v.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/variable_aggregator_test.go b/internal/agent/component/variable_aggregator_test.go index 4c21ca834b..d7e90bcf62 100644 --- a/internal/agent/component/variable_aggregator_test.go +++ b/internal/agent/component/variable_aggregator_test.go @@ -60,7 +60,7 @@ func TestVariableAggregator_FirstNonEmpty(t *testing.T) { if err != nil { t.Fatalf("NewVariableAggregatorComponent: %v", err) } - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -99,7 +99,7 @@ func TestVariableAggregator_SkipsEmptyString(t *testing.T) { if err != nil { t.Fatalf("NewVariableAggregatorComponent: %v", err) } - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -143,7 +143,7 @@ func TestVariableAggregator_MultipleGroups(t *testing.T) { if err != nil { t.Fatalf("NewVariableAggregatorComponent: %v", err) } - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -177,7 +177,7 @@ func TestVariableAggregator_AllEmpty(t *testing.T) { if err != nil { t.Fatalf("NewVariableAggregatorComponent: %v", err) } - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } diff --git a/internal/agent/component/variable_assigner.go b/internal/agent/component/variable_assigner.go index 8ed5684bb8..43f232bb40 100644 --- a/internal/agent/component/variable_assigner.go +++ b/internal/agent/component/variable_assigner.go @@ -43,6 +43,8 @@ import ( "strings" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) const componentNameVariableAssigner = "VariableAssigner" @@ -144,7 +146,7 @@ func (v *VariableAssignerComponent) GetInputForm() map[string]any { // returned an "ERROR:..." sentinel. The list of refs that were // assigned is returned at outputs["assignments"]; per-item errors (if // any) are returned at outputs["errors"]. -func (v *VariableAssignerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (v *VariableAssignerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) if err != nil { return nil, fmt.Errorf("VariableAssigner: %w", err) @@ -200,8 +202,8 @@ func (v *VariableAssignerComponent) Invoke(ctx context.Context, inputs map[strin } // Stream mirrors Invoke; VariableAssigner is a single-shot apply. -func (v *VariableAssignerComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := v.Invoke(ctx, inputs) +func (v *VariableAssignerComponent) Stream(ctx context.Context, db *gorm.DB, inputs map[string]any) (<-chan map[string]any, error) { + out, err := v.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/agent/component/variable_assigner_test.go b/internal/agent/component/variable_assigner_test.go index a21df59f52..27137b8b9e 100644 --- a/internal/agent/component/variable_assigner_test.go +++ b/internal/agent/component/variable_assigner_test.go @@ -41,7 +41,7 @@ func TestVariableAssigner_Append(t *testing.T) { if err != nil { t.Fatalf("NewVariableAssignerComponent: %v", err) } - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -75,7 +75,7 @@ func TestVariableAssigner_Overwrite(t *testing.T) { if err != nil { t.Fatalf("NewVariableAssignerComponent: %v", err) } - if _, err := c.Invoke(ctx, nil); err != nil { + if _, err = c.Invoke(ctx, nil, nil); err != nil { t.Fatalf("Invoke: %v", err) } if got, want := state.Outputs["cpn_0"]["x"], "fresh"; got != want { @@ -101,7 +101,7 @@ func TestVariableAssigner_DivideByZero(t *testing.T) { if err != nil { t.Fatalf("NewVariableAssignerComponent: %v", err) } - out, err := c.Invoke(ctx, nil) + out, err := c.Invoke(ctx, nil, nil) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -145,7 +145,7 @@ func TestVariableAssigner_Clear(t *testing.T) { if err != nil { t.Fatalf("NewVariableAssignerComponent: %v", err) } - if _, err := c.Invoke(ctx, nil); err != nil { + if _, err = c.Invoke(ctx, nil, nil); err != nil { t.Fatalf("Invoke: %v", err) } if got := state.Outputs["cpn_0"]["a"]; reflect.DeepEqual(got, []any{1, 2}) { @@ -178,7 +178,7 @@ func TestVariableAssigner_Arithmetic(t *testing.T) { if err != nil { t.Fatalf("NewVariableAssignerComponent: %v", err) } - if _, err := c.Invoke(ctx, nil); err != nil { + if _, err = c.Invoke(ctx, nil, nil); err != nil { t.Fatalf("Invoke: %v", err) } // 10 + 5 = 15, 15 - 3 = 12, 12 * 2 = 24, 24 / 4 = 6 @@ -197,7 +197,7 @@ func TestVariableAssigner_RemoveFirstLast(t *testing.T) { {"variable": "cpn_0@xs", "operator": "remove_first", "parameter": "x"}, } c, _ := NewVariableAssignerComponent(map[string]any{"variables": vars}) - if _, err := c.Invoke(ctx, nil); err != nil { + if _, err := c.Invoke(ctx, nil, nil); err != nil { t.Fatalf("Invoke: %v", err) } if got, want := state.Outputs["cpn_0"]["xs"], []any{"b", "c", "d"}; !reflect.DeepEqual(got, want) { @@ -208,7 +208,7 @@ func TestVariableAssigner_RemoveFirstLast(t *testing.T) { {"variable": "cpn_0@xs", "operator": "remove_last", "parameter": "x"}, } c, _ = NewVariableAssignerComponent(map[string]any{"variables": vars}) - if _, err := c.Invoke(ctx, nil); err != nil { + if _, err := c.Invoke(ctx, nil, nil); err != nil { t.Fatalf("Invoke: %v", err) } if got, want := state.Outputs["cpn_0"]["xs"], []any{"b", "c"}; !reflect.DeepEqual(got, want) { @@ -225,7 +225,7 @@ func TestVariableAssigner_SysTarget(t *testing.T) { {"variable": "sys.x", "operator": "set", "parameter": "hello"}, } c, _ := NewVariableAssignerComponent(map[string]any{"variables": vars}) - if _, err := c.Invoke(ctx, nil); err != nil { + if _, err := c.Invoke(ctx, nil, nil); err != nil { t.Fatalf("Invoke: %v", err) } if got, want := state.Sys["x"], "hello"; got != want { diff --git a/internal/agent/component/vision_test.go b/internal/agent/component/vision_test.go index 3997ebc432..1ee411f04a 100644 --- a/internal/agent/component/vision_test.go +++ b/internal/agent/component/vision_test.go @@ -288,7 +288,7 @@ func TestLLM_Invoke_ForwardsImagesToInvoker(t *testing.T) { uri := "data:image/png;base64,iVBORw0KGgo=" c := NewLLMComponent(LLMParam{ModelID: "echo"}) - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(context.Background(), nil, map[string]any{ "user_prompt": "what is this?", "visual_files": []string{uri}, }) @@ -320,7 +320,7 @@ func TestLLM_Invoke_NoVisualFiles_BackwardCompat(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo"}) - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(context.Background(), nil, map[string]any{ "user_prompt": "hi", }) if err != nil { @@ -350,7 +350,7 @@ func TestLLM_Invoke_VisualFilesAsString(t *testing.T) { uri := "data:image/jpeg;base64,/9j/4AAQ" c := NewLLMComponent(LLMParam{ModelID: "echo"}) - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(context.Background(), nil, map[string]any{ "user_prompt": "describe", "visual_files": "see " + uri, }) diff --git a/internal/agent/runtime/component.go b/internal/agent/runtime/component.go index 73c802d1ea..a6bedb4a71 100644 --- a/internal/agent/runtime/component.go +++ b/internal/agent/runtime/component.go @@ -33,6 +33,8 @@ import ( "context" "fmt" "sync" + + "gorm.io/gorm" ) // Component is the minimal interface the canvas builder needs at @@ -40,7 +42,7 @@ import ( // Component type has more methods (Name / Stream / Inputs / Outputs); // it satisfies this smaller interface by Go's structural typing. type Component interface { - Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) + Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) } // ComponentFactory builds a Component from a DSL name + params map. diff --git a/internal/agent/runtime/registry_test.go b/internal/agent/runtime/registry_test.go index eb33afe389..039b9bfaf8 100644 --- a/internal/agent/runtime/registry_test.go +++ b/internal/agent/runtime/registry_test.go @@ -23,6 +23,8 @@ import ( "sort" "sync" "testing" + + "gorm.io/gorm" ) // stubComponent is a minimal Component impl used as the factory's @@ -33,7 +35,7 @@ type stubComponent struct { params map[string]any } -func (s *stubComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (s *stubComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { return map[string]any{"name": s.name, "params": s.params}, nil } diff --git a/internal/agent/tool/code_exec.go b/internal/agent/tool/code_exec.go index a14effaf35..02f2365c3d 100644 --- a/internal/agent/tool/code_exec.go +++ b/internal/agent/tool/code_exec.go @@ -24,11 +24,11 @@ import ( "os" "strings" + "ragflow/internal/common" + "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" "go.uber.org/zap" - - "ragflow/internal/common" ) // ErrCodeExecSandboxMissing is returned when no sandbox client is diff --git a/internal/agent/tool/exesql.go b/internal/agent/tool/exesql.go index b30344afb2..c9f5478a7e 100644 --- a/internal/agent/tool/exesql.go +++ b/internal/agent/tool/exesql.go @@ -53,14 +53,13 @@ import ( "strings" "time" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" // SQL drivers — registered via their init() side effects. _ "github.com/denisenkom/go-mssqldb" _ "github.com/go-sql-driver/mysql" _ "github.com/lib/pq" - "github.com/cloudwego/eino/components/tool" - "github.com/cloudwego/eino/schema" - "ragflow/internal/agent/runtime" ) @@ -347,7 +346,7 @@ func (e *ExeSQLTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ defer cancel() } - if err := db.PingContext(ctx); err != nil { + if err = db.PingContext(ctx); err != nil { return exesqlErrorResult(fmt.Errorf("exesql: ping: %w", err)), fmt.Errorf("exesql: ping: %w", err) } diff --git a/internal/agent/tool/mcp.go b/internal/agent/tool/mcp.go index 9c64d5ecf5..f7deabc820 100644 --- a/internal/agent/tool/mcp.go +++ b/internal/agent/tool/mcp.go @@ -30,10 +30,10 @@ import ( "net/http" "time" + mcpclient "ragflow/internal/utility" + "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" - - mcpclient "ragflow/internal/utility" ) // MCPToolAdapter wraps a single MCP-discovered tool descriptor as an diff --git a/internal/agent/tool/registry_test.go b/internal/agent/tool/registry_test.go index 60dcf83578..03b5fb8118 100644 --- a/internal/agent/tool/registry_test.go +++ b/internal/agent/tool/registry_test.go @@ -188,7 +188,7 @@ func TestToolRegistry_SchemasAreComplete(t *testing.T) { // Alias consistency: execute_sql and exesql must surface the // same canonical Info().Name; same for retrieval/search_my_dataset/ - // search_my_dateset and crawler/web_crawler. A bug here would mean + // search_my_dataset and crawler/web_crawler. A bug here would mean // an alias was accidentally pointed at a different tool. canonicalByAlias := map[string]string{ "execute_sql": "execute_sql", diff --git a/internal/agent/tool/retrieval.go b/internal/agent/tool/retrieval.go index 9f38959836..c50ea199ed 100644 --- a/internal/agent/tool/retrieval.go +++ b/internal/agent/tool/retrieval.go @@ -21,6 +21,7 @@ import ( "encoding/json" "errors" "fmt" + "ragflow/internal/dao" "strings" "github.com/cloudwego/eino/components/tool" @@ -161,7 +162,7 @@ func (r *RetrievalTool) InvokableRun(ctx context.Context, argumentsInJSON string // via SetRetrievalService (or SetSimpleRetrievalService for // dev), the chunks flow through normally. svc := GetRetrievalService() - chunks, err := svc.Search(ctx, RetrievalRequest{ + chunks, err := svc.Search(ctx, dao.DB, RetrievalRequest{ Query: args.Query, DatasetIDs: args.DatasetIDs, TopN: args.TopN, diff --git a/internal/agent/tool/retrieval_nlp.go b/internal/agent/tool/retrieval_nlp.go index 8b1abe2ae2..88f9e14d72 100644 --- a/internal/agent/tool/retrieval_nlp.go +++ b/internal/agent/tool/retrieval_nlp.go @@ -74,6 +74,8 @@ import ( "ragflow/internal/engine" "ragflow/internal/entity" "ragflow/internal/service/nlp" + + "gorm.io/gorm" ) // NLPRetrievalAdapter wraps *nlp.RetrievalService behind the @@ -87,7 +89,7 @@ type NLPRetrievalAdapter struct { } type knowledgebaseLookup interface { - GetByIDs(ids []string) ([]*entity.Knowledgebase, error) + GetByIDs(ctx context.Context, sqlDB *gorm.DB, ids []string) ([]*entity.Knowledgebase, error) } // NewNLPRetrievalAdapter wraps an already-constructed @@ -115,7 +117,7 @@ func NewNLPRetrievalAdapterFromDeps(docEngine engine.DocEngine, documentDAO *dao // Search implements RetrievalService. The translation rules live // at the top of this file. -func (a *NLPRetrievalAdapter) Search(ctx context.Context, req RetrievalRequest) ([]RetrievalChunk, error) { +func (a *NLPRetrievalAdapter) Search(ctx context.Context, db *gorm.DB, req RetrievalRequest) ([]RetrievalChunk, error) { if a == nil || a.svc == nil { return nil, ErrRetrievalServiceMissing } @@ -135,7 +137,7 @@ func (a *NLPRetrievalAdapter) Search(ctx context.Context, req RetrievalRequest) topN = 8 } - tenantIDs, err := a.resolveTenantIDs(req) + tenantIDs, err := a.resolveTenantIDs(ctx, db, req) if err != nil { return nil, err } @@ -181,7 +183,7 @@ func nlpRequestFromRetrieval(req RetrievalRequest, tenantIDs []string, topN int) return nlpReq } -func (a *NLPRetrievalAdapter) resolveTenantIDs(req RetrievalRequest) ([]string, error) { +func (a *NLPRetrievalAdapter) resolveTenantIDs(ctx context.Context, db *gorm.DB, req RetrievalRequest) ([]string, error) { seen := map[string]struct{}{} tenantIDs := make([]string, 0, 1) appendTenantID := func(tenantID string) { @@ -202,7 +204,7 @@ func (a *NLPRetrievalAdapter) resolveTenantIDs(req RetrievalRequest) ([]string, return tenantIDs, nil } - kbs, err := a.kbDAO.GetByIDs(datasetIDs) + kbs, err := a.kbDAO.GetByIDs(ctx, db, datasetIDs) if err != nil { return nil, fmt.Errorf("retrieval: resolve dataset tenants: %w", err) } diff --git a/internal/agent/tool/retrieval_nlp_test.go b/internal/agent/tool/retrieval_nlp_test.go index 7771073781..d98f50d613 100644 --- a/internal/agent/tool/retrieval_nlp_test.go +++ b/internal/agent/tool/retrieval_nlp_test.go @@ -32,6 +32,8 @@ import ( "testing" "ragflow/internal/entity" + + "gorm.io/gorm" ) // floatEqual compares two floats with a small epsilon so @@ -221,7 +223,7 @@ func TestTranslateChunk_MissingAllScores(t *testing.T) { // error, not a panic. func TestNewNLPRetrievalAdapter_NilService(t *testing.T) { a := NewNLPRetrievalAdapter(nil) - _, err := a.Search(context.TODO(), RetrievalRequest{Query: "hi"}) + _, err := a.Search(context.TODO(), nil, RetrievalRequest{Query: "hi"}) if err == nil { t.Fatal("expected error from nil-service adapter") } @@ -281,7 +283,7 @@ func TestNLPRequestFromRetrieval_FallsBackToTopNHeadroom(t *testing.T) { func TestNLPRetrievalAdapter_ResolveTenantIDsStaysWithinRequestTenant(t *testing.T) { a := &NLPRetrievalAdapter{} - got, err := a.resolveTenantIDs(RetrievalRequest{ + got, err := a.resolveTenantIDs(nil, nil, RetrievalRequest{ TenantID: "tenant-a", DatasetIDs: []string{"kb-1", "kb-2", "kb-missing"}, }) @@ -307,7 +309,7 @@ func TestNLPRetrievalAdapter_ResolveTenantIDsFromDatasetIDs(t *testing.T) { }, }, } - got, err := a.resolveTenantIDs(RetrievalRequest{ + got, err := a.resolveTenantIDs(nil, nil, RetrievalRequest{ DatasetIDs: []string{"kb-1", "kb-2", "kb-3", " "}, }) if err != nil { @@ -326,7 +328,7 @@ func TestNLPRetrievalAdapter_ResolveTenantIDsKeepsRequestTenantFirst(t *testing. }, }, } - got, err := a.resolveTenantIDs(RetrievalRequest{ + got, err := a.resolveTenantIDs(nil, nil, RetrievalRequest{ TenantID: "tenant-a", DatasetIDs: []string{"kb-1"}, }) @@ -343,6 +345,6 @@ type fakeKnowledgebaseLookup struct { err error } -func (f fakeKnowledgebaseLookup) GetByIDs(_ []string) ([]*entity.Knowledgebase, error) { +func (f fakeKnowledgebaseLookup) GetByIDs(ctx context.Context, db *gorm.DB, ids []string) ([]*entity.Knowledgebase, error) { return f.kbs, f.err } diff --git a/internal/agent/tool/retrieval_service.go b/internal/agent/tool/retrieval_service.go index 86d40e26f5..ecc0615987 100644 --- a/internal/agent/tool/retrieval_service.go +++ b/internal/agent/tool/retrieval_service.go @@ -25,6 +25,8 @@ import ( "errors" "fmt" "sync" + + "gorm.io/gorm" ) // RetrievalChunk is the minimal shape RetrievalService returns. The @@ -65,7 +67,7 @@ type RetrievalRequest struct { // Today only the stub impl exists; production code can register // a real impl via SetRetrievalService during boot. type RetrievalService interface { - Search(ctx context.Context, req RetrievalRequest) ([]RetrievalChunk, error) + Search(ctx context.Context, db *gorm.DB, req RetrievalRequest) ([]RetrievalChunk, error) } // KGRetrievalService is the GraphRAG retrieval surface. The @@ -77,7 +79,7 @@ type RetrievalService interface { // adapter can be tested in isolation and wired independently // at boot. type KGRetrievalService interface { - Search(ctx context.Context, req RetrievalRequest) ([]RetrievalChunk, error) + Search(ctx context.Context, db *gorm.DB, req RetrievalRequest) ([]RetrievalChunk, error) } // ErrRetrievalServiceMissing is declared in retrieval.go (kept @@ -108,7 +110,7 @@ func GetRetrievalService() RetrievalService { type stubRetrievalService struct{} -func (stubRetrievalService) Search(_ context.Context, _ RetrievalRequest) ([]RetrievalChunk, error) { +func (stubRetrievalService) Search(_ context.Context, _ *gorm.DB, _ RetrievalRequest) ([]RetrievalChunk, error) { return nil, ErrRetrievalServiceMissing } @@ -119,7 +121,7 @@ func (stubRetrievalService) Search(_ context.Context, _ RetrievalRequest) ([]Ret // SetRetrievalService. type simpleRetrievalService struct{} -func (simpleRetrievalService) Search(_ context.Context, req RetrievalRequest) ([]RetrievalChunk, error) { +func (simpleRetrievalService) Search(_ context.Context, _ *gorm.DB, req RetrievalRequest) ([]RetrievalChunk, error) { if req.Query == "" { return nil, nil } @@ -199,6 +201,6 @@ func GetKGRetrievalService() KGRetrievalService { type stubKGRetrievalService struct{} -func (stubKGRetrievalService) Search(_ context.Context, _ RetrievalRequest) ([]RetrievalChunk, error) { +func (stubKGRetrievalService) Search(_ context.Context, _ *gorm.DB, _ RetrievalRequest) ([]RetrievalChunk, error) { return nil, ErrKGRetrievalServiceMissing } diff --git a/internal/agent/tool/retrieval_service_test.go b/internal/agent/tool/retrieval_service_test.go index 5d98e3048d..3bf14e0920 100644 --- a/internal/agent/tool/retrieval_service_test.go +++ b/internal/agent/tool/retrieval_service_test.go @@ -20,6 +20,8 @@ import ( "context" "errors" "testing" + + "gorm.io/gorm" ) // TestRetrievalService_DefaultStub: with no SetRetrievalService @@ -29,7 +31,9 @@ func TestRetrievalService_DefaultStub(t *testing.T) { if svc == nil { t.Fatal("GetRetrievalService returned nil") } - _, err := svc.Search(context.Background(), RetrievalRequest{Query: "hi"}) + + ctx := t.Context() + _, err := svc.Search(ctx, nil, RetrievalRequest{Query: "hi"}) if !errors.Is(err, ErrRetrievalServiceMissing) { t.Errorf("err=%v, want ErrRetrievalServiceMissing", err) } @@ -41,9 +45,10 @@ func TestRetrievalService_RegisterAndRestore(t *testing.T) { prev := GetRetrievalService() t.Cleanup(func() { SetRetrievalService(prev) }) + ctx := t.Context() // Install a fake. SetRetrievalService(fakeRetrievalService{chunks: []RetrievalChunk{{ID: "1", Content: "fake"}}}) - got, err := GetRetrievalService().Search(context.Background(), RetrievalRequest{Query: "x"}) + got, err := GetRetrievalService().Search(ctx, nil, RetrievalRequest{Query: "x"}) if err != nil { t.Fatalf("Search: %v", err) } @@ -53,7 +58,7 @@ func TestRetrievalService_RegisterAndRestore(t *testing.T) { // Restore via nil. SetRetrievalService(nil) - if _, err := GetRetrievalService().Search(context.Background(), RetrievalRequest{}); !errors.Is(err, ErrRetrievalServiceMissing) { + if _, err := GetRetrievalService().Search(ctx, nil, RetrievalRequest{}); !errors.Is(err, ErrRetrievalServiceMissing) { t.Errorf("after nil-reset: err=%v, want ErrRetrievalServiceMissing", err) } } @@ -64,7 +69,7 @@ type fakeRetrievalService struct { err error } -func (f fakeRetrievalService) Search(_ context.Context, _ RetrievalRequest) ([]RetrievalChunk, error) { +func (f fakeRetrievalService) Search(_ context.Context, _ *gorm.DB, _ RetrievalRequest) ([]RetrievalChunk, error) { return f.chunks, f.err } diff --git a/internal/agent/tool/retrieval_test.go b/internal/agent/tool/retrieval_test.go index 16072a5c4f..89f79f5e18 100644 --- a/internal/agent/tool/retrieval_test.go +++ b/internal/agent/tool/retrieval_test.go @@ -24,6 +24,8 @@ import ( "testing" "ragflow/internal/agent/runtime" + + "gorm.io/gorm" ) func TestRetrieval_StubsErrorWhenServiceMissing(t *testing.T) { @@ -307,7 +309,7 @@ type capturingRetrievalService struct { req RetrievalRequest } -func (s *capturingRetrievalService) Search(_ context.Context, req RetrievalRequest) ([]RetrievalChunk, error) { +func (s *capturingRetrievalService) Search(_ context.Context, _ *gorm.DB, req RetrievalRequest) ([]RetrievalChunk, error) { s.req = req return []RetrievalChunk{{ID: "ck-1", Content: "answer"}}, nil } @@ -316,6 +318,6 @@ type staticRetrievalService struct { chunks []RetrievalChunk } -func (s staticRetrievalService) Search(_ context.Context, _ RetrievalRequest) ([]RetrievalChunk, error) { +func (s staticRetrievalService) Search(_ context.Context, _ *gorm.DB, _ RetrievalRequest) ([]RetrievalChunk, error) { return s.chunks, nil } diff --git a/internal/agent/tool/retrieval_wiring_test.go b/internal/agent/tool/retrieval_wiring_test.go index f41bdde0e8..fdc86829e3 100644 --- a/internal/agent/tool/retrieval_wiring_test.go +++ b/internal/agent/tool/retrieval_wiring_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "errors" "strings" "testing" @@ -31,7 +30,8 @@ func TestRetrievalTool_StubReturnsServiceMissing(t *testing.T) { // Ensure stub is the active service. SetRetrievalService(nil) tool := NewRetrievalTool() - _, err := tool.InvokableRun(context.Background(), `{"query":"hi"}`) + ctx := t.Context() + _, err := tool.InvokableRun(ctx, `{"query":"hi"}`) if !errors.Is(err, ErrRetrievalServiceMissing) { t.Errorf("err=%v, want ErrRetrievalServiceMissing", err) } @@ -45,8 +45,8 @@ func TestRetrievalTool_SimpleServiceReturnsChunks(t *testing.T) { SetSimpleRetrievalService() tool := NewRetrievalTool() - out, err := tool.InvokableRun(context.Background(), - `{"query":"hello world","top_n":2}`) + ctx := t.Context() + out, err := tool.InvokableRun(ctx, `{"query":"hello world","top_n":2}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -64,7 +64,8 @@ func TestRetrievalTool_SimpleServiceReturnsChunks(t *testing.T) { // TestSimpleRetrievalService_EmptyQuery: empty query → no chunks, // no error. func TestSimpleRetrievalService_EmptyQuery(t *testing.T) { - chunks, err := simpleRetrievalService{}.Search(context.Background(), RetrievalRequest{Query: ""}) + ctx := t.Context() + chunks, err := simpleRetrievalService{}.Search(ctx, nil, RetrievalRequest{Query: ""}) if err != nil { t.Errorf("err: %v", err) } @@ -76,7 +77,8 @@ func TestSimpleRetrievalService_EmptyQuery(t *testing.T) { // TestSimpleRetrievalService_RespectsTopN: top_n=1 → exactly 1 // chunk. func TestSimpleRetrievalService_RespectsTopN(t *testing.T) { - chunks, err := simpleRetrievalService{}.Search(context.Background(), RetrievalRequest{ + ctx := t.Context() + chunks, err := simpleRetrievalService{}.Search(ctx, nil, RetrievalRequest{ Query: "x", TopN: 1, }) diff --git a/internal/dao/kb.go b/internal/dao/kb.go index d9dafcc2c1..1dabf86997 100644 --- a/internal/dao/kb.go +++ b/internal/dao/kb.go @@ -17,6 +17,7 @@ package dao import ( + "context" "errors" "fmt" "path" @@ -31,11 +32,11 @@ import ( // GetTenantIDByKBID is a convenience function that retrieves the tenant ID // for a given knowledge base ID. It is a package-level helper so both the // service and engine layers can use it without circular imports. -func GetTenantIDByKBID(kbID string) (string, error) { +func GetTenantIDByKBID(ctx context.Context, db *gorm.DB, kbID string) (string, error) { kbDAO := NewKnowledgebaseDAO() - kb, err := kbDAO.GetByID(kbID) + kb, err := kbDAO.GetByID(ctx, db, kbID) if err != nil { - return "", fmt.Errorf("knowledgebase not found: %w", err) + return "", fmt.Errorf("dataset not found: %w", err) } return kb.TenantID, nil } @@ -59,29 +60,29 @@ func NewKnowledgebaseDAO() *KnowledgebaseDAO { } // Create creates a new knowledge base record -func (dao *KnowledgebaseDAO) Create(kb *entity.Knowledgebase) error { - return DB.Create(kb).Error +func (dao *KnowledgebaseDAO) Create(ctx context.Context, db *gorm.DB, kb *entity.Knowledgebase) error { + return db.WithContext(ctx).Create(kb).Error } // Update updates a knowledge base record -func (dao *KnowledgebaseDAO) Update(kb *entity.Knowledgebase) error { - return DB.Save(kb).Error +func (dao *KnowledgebaseDAO) Update(ctx context.Context, db *gorm.DB, kb *entity.Knowledgebase) error { + return db.WithContext(ctx).Save(kb).Error } // UpdateByID updates a knowledge base by ID with the given fields -func (dao *KnowledgebaseDAO) UpdateByID(id string, updates map[string]interface{}) error { - return DB.Model(&entity.Knowledgebase{}).Where("id = ?", id).Updates(updates).Error +func (dao *KnowledgebaseDAO) UpdateByID(ctx context.Context, db *gorm.DB, id string, updates map[string]interface{}) error { + return db.WithContext(ctx).Model(&entity.Knowledgebase{}).Where("id = ?", id).Updates(updates).Error } // Delete soft deletes a knowledge base by setting status to invalid -func (dao *KnowledgebaseDAO) Delete(id string) error { - return DB.Model(&entity.Knowledgebase{}).Where("id = ?", id).Update("status", string(entity.StatusInvalid)).Error +func (dao *KnowledgebaseDAO) Delete(ctx context.Context, db *gorm.DB, id string) error { + return db.WithContext(ctx).Model(&entity.Knowledgebase{}).Where("id = ?", id).Update("status", string(entity.StatusInvalid)).Error } // GetByID retrieves a knowledge base by ID -func (dao *KnowledgebaseDAO) GetByID(id string) (*entity.Knowledgebase, error) { +func (dao *KnowledgebaseDAO) GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.Knowledgebase, error) { var kb entity.Knowledgebase - err := DB.Where("id = ? AND status = ?", id, string(entity.StatusValid)).First(&kb).Error + err := db.WithContext(ctx).Where("id = ? AND status = ?", id, string(entity.StatusValid)).First(&kb).Error if err != nil { return nil, err } @@ -89,9 +90,9 @@ func (dao *KnowledgebaseDAO) GetByID(id string) (*entity.Knowledgebase, error) { } // GetByIDAndTenantID retrieves a knowledge base by ID and tenant ID -func (dao *KnowledgebaseDAO) GetByIDAndTenantID(id, tenantID string) (*entity.Knowledgebase, error) { +func (dao *KnowledgebaseDAO) GetByIDAndTenantID(ctx context.Context, db *gorm.DB, id, tenantID string) (*entity.Knowledgebase, error) { var kb entity.Knowledgebase - err := DB.Where("id = ? AND tenant_id = ? AND status = ?", id, tenantID, string(entity.StatusValid)).First(&kb).Error + err := db.WithContext(ctx).Where("id = ? AND tenant_id = ? AND status = ?", id, tenantID, string(entity.StatusValid)).First(&kb).Error if err != nil { return nil, err } @@ -99,16 +100,16 @@ func (dao *KnowledgebaseDAO) GetByIDAndTenantID(id, tenantID string) (*entity.Kn } // GetByIDs retrieves multiple knowledge bases by IDs -func (dao *KnowledgebaseDAO) GetByIDs(ids []string) ([]*entity.Knowledgebase, error) { +func (dao *KnowledgebaseDAO) GetByIDs(ctx context.Context, db *gorm.DB, ids []string) ([]*entity.Knowledgebase, error) { var kbs []*entity.Knowledgebase - err := DB.Where("id IN ? AND status = ?", ids, string(entity.StatusValid)).Find(&kbs).Error + err := db.WithContext(ctx).Where("id IN ? AND status = ?", ids, string(entity.StatusValid)).Find(&kbs).Error return kbs, err } // GetByName retrieves a knowledge base by name and tenant ID -func (dao *KnowledgebaseDAO) GetByName(name, tenantID string) (*entity.Knowledgebase, error) { +func (dao *KnowledgebaseDAO) GetByName(ctx context.Context, db *gorm.DB, name, tenantID string) (*entity.Knowledgebase, error) { var kb entity.Knowledgebase - err := DB.Where("LOWER(name) = LOWER(?) AND tenant_id = ? AND status = ?", name, tenantID, string(entity.StatusValid)).First(&kb).Error + err := db.WithContext(ctx).Where("LOWER(name) = LOWER(?) AND tenant_id = ? AND status = ?", name, tenantID, string(entity.StatusValid)).First(&kb).Error if err != nil { return nil, err } @@ -116,16 +117,16 @@ func (dao *KnowledgebaseDAO) GetByName(name, tenantID string) (*entity.Knowledge } // GetByCreatedBy retrieves knowledge bases created by a specific user -func (dao *KnowledgebaseDAO) GetByCreatedBy(createdBy string) ([]*entity.Knowledgebase, error) { +func (dao *KnowledgebaseDAO) GetByCreatedBy(ctx context.Context, db *gorm.DB, createdBy string) ([]*entity.Knowledgebase, error) { var kbs []*entity.Knowledgebase - err := DB.Where("created_by = ? AND status = ?", createdBy, string(entity.StatusValid)).Find(&kbs).Error + err := db.WithContext(ctx).Where("created_by = ? AND status = ?", createdBy, string(entity.StatusValid)).Find(&kbs).Error return kbs, err } // Query retrieves knowledge bases with filters -func (dao *KnowledgebaseDAO) Query(filters map[string]interface{}) ([]*entity.Knowledgebase, error) { +func (dao *KnowledgebaseDAO) Query(ctx context.Context, db *gorm.DB, filters map[string]interface{}) ([]*entity.Knowledgebase, error) { var kbs []*entity.Knowledgebase - query := DB.Where("status = ?", string(entity.StatusValid)) + query := db.WithContext(ctx).Where("status = ?", string(entity.StatusValid)) for key, value := range filters { if value != nil && value != "" { @@ -138,9 +139,9 @@ func (dao *KnowledgebaseDAO) Query(filters map[string]interface{}) ([]*entity.Kn } // QueryOne retrieves a single knowledge base with filters -func (dao *KnowledgebaseDAO) QueryOne(filters map[string]interface{}) (*entity.Knowledgebase, error) { +func (dao *KnowledgebaseDAO) QueryOne(ctx context.Context, db *gorm.DB, filters map[string]interface{}) (*entity.Knowledgebase, error) { var kb entity.Knowledgebase - query := DB.Where("status = ?", string(entity.StatusValid)) + query := db.WithContext(ctx).Where("status = ?", string(entity.StatusValid)) for key, value := range filters { if value != nil && value != "" { @@ -156,9 +157,9 @@ func (dao *KnowledgebaseDAO) QueryOne(filters map[string]interface{}) (*entity.K } // Count returns the count of knowledge bases matching the filters -func (dao *KnowledgebaseDAO) Count(filters map[string]interface{}) (int64, error) { +func (dao *KnowledgebaseDAO) Count(ctx context.Context, db *gorm.DB, filters map[string]interface{}) (int64, error) { var count int64 - query := DB.Model(&entity.Knowledgebase{}).Where("status = ?", string(entity.StatusValid)) + query := db.WithContext(ctx).Model(&entity.Knowledgebase{}).Where("status = ?", string(entity.StatusValid)) for key, value := range filters { if value != nil && value != "" { @@ -172,11 +173,11 @@ func (dao *KnowledgebaseDAO) Count(filters map[string]interface{}) (int64, error // GetByTenantIDs retrieves knowledge bases by tenant IDs with pagination // This matches the Python get_by_tenant_ids method -func (dao *KnowledgebaseDAO) GetByTenantIDs(tenantIDs []string, userID string, pageNumber, itemsPerPage int, orderby string, desc bool, keywords, parserID, id, name string) ([]*entity.KnowledgebaseListItem, int64, error) { +func (dao *KnowledgebaseDAO) GetByTenantIDs(ctx context.Context, db *gorm.DB, tenantIDs []string, userID string, pageNumber, itemsPerPage int, orderby string, desc bool, keywords, parserID, id, name string) ([]*entity.KnowledgebaseListItem, int64, error) { var kbs []*entity.KnowledgebaseListItem var total int64 - query := DB.Model(&entity.Knowledgebase{}). + query := db.WithContext(ctx).Model(&entity.Knowledgebase{}). Select(`knowledgebase.id, knowledgebase.avatar, knowledgebase.name, knowledgebase.language, knowledgebase.description, knowledgebase.tenant_id, knowledgebase.permission, knowledgebase.doc_num, knowledgebase.token_num, @@ -228,10 +229,10 @@ func (dao *KnowledgebaseDAO) GetByTenantIDs(tenantIDs []string, userID string, p // GetAllByTenantIDs retrieves all permitted knowledge bases by tenant IDs // This matches the Python get_all_kb_by_tenant_ids method -func (dao *KnowledgebaseDAO) GetAllByTenantIDs(tenantIDs []string, userID string) ([]*entity.Knowledgebase, error) { +func (dao *KnowledgebaseDAO) GetAllByTenantIDs(ctx context.Context, db *gorm.DB, tenantIDs []string, userID string) ([]*entity.Knowledgebase, error) { var kbs []*entity.Knowledgebase - err := DB.Where( + err := db.WithContext(ctx).Where( "(tenant_id IN ? AND permission = ?) OR tenant_id = ?", tenantIDs, string(entity.TenantPermissionTeam), userID, ).Order("create_time ASC").Find(&kbs).Error @@ -241,10 +242,10 @@ func (dao *KnowledgebaseDAO) GetAllByTenantIDs(tenantIDs []string, userID string // GetDetail retrieves detailed knowledge base information with joined pipeline data // This matches the Python get_detail method -func (dao *KnowledgebaseDAO) GetDetail(kbID string) (*entity.KnowledgebaseDetail, error) { +func (dao *KnowledgebaseDAO) GetDetail(ctx context.Context, db *gorm.DB, kbID string) (*entity.KnowledgebaseDetail, error) { var detail entity.KnowledgebaseDetail - err := DB.Table("knowledgebase"). + err := db.WithContext(ctx).Table("knowledgebase"). Select(`knowledgebase.id, knowledgebase.embd_id, knowledgebase.avatar, knowledgebase.name, knowledgebase.language, knowledgebase.description, knowledgebase.permission, knowledgebase.doc_num, knowledgebase.token_num, knowledgebase.chunk_num, @@ -272,9 +273,9 @@ func (dao *KnowledgebaseDAO) GetDetail(kbID string) (*entity.KnowledgebaseDetail // 2. If user is the owner tenant, return true // 3. If permission is "me", only owner tenant can access // 4. If permission is "team", user must be a member of the tenant -func (dao *KnowledgebaseDAO) Accessible(datasetID, userID string) bool { +func (dao *KnowledgebaseDAO) Accessible(ctx context.Context, db *gorm.DB, datasetID, userID string) bool { var kb entity.Knowledgebase - err := DB.Where("id = ? AND status = ?", datasetID, string(entity.StatusValid)).First(&kb).Error + err := db.WithContext(ctx).Where("id = ? AND status = ?", datasetID, string(entity.StatusValid)).First(&kb).Error if err != nil { return false } @@ -290,7 +291,7 @@ func (dao *KnowledgebaseDAO) Accessible(datasetID, userID string) bool { } var count int64 - err = DB.Table("user_tenant"). + err = db.WithContext(ctx).Table("user_tenant"). Where("tenant_id = ? AND user_id = ? AND status = ?", kb.TenantID, userID, "1"). Count(&count).Error @@ -302,9 +303,9 @@ func (dao *KnowledgebaseDAO) Accessible(datasetID, userID string) bool { // Accessible4Deletion checks if a knowledge base can be deleted by a user // This matches the Python accessible4deletion method -func (dao *KnowledgebaseDAO) Accessible4Deletion(kbID, userID string) bool { +func (dao *KnowledgebaseDAO) Accessible4Deletion(ctx context.Context, db *gorm.DB, kbID, userID string) bool { var count int64 - err := DB.Model(&entity.Knowledgebase{}). + err := db.WithContext(ctx).Model(&entity.Knowledgebase{}). Where("id = ? AND created_by = ? AND status = ?", kbID, userID, string(entity.StatusValid)). Count(&count).Error @@ -316,13 +317,13 @@ func (dao *KnowledgebaseDAO) Accessible4Deletion(kbID, userID string) bool { // DuplicateName generates a unique name by appending parentheses if name already exists // This matches the Python duplicate_name function behavior -func (dao *KnowledgebaseDAO) DuplicateName(name, tenantID string) string { +func (dao *KnowledgebaseDAO) DuplicateName(ctx context.Context, db *gorm.DB, name, tenantID string) string { const maxRetries = 1000 currentName := name for retries := 0; retries < maxRetries; retries++ { var count int64 - err := DB.Model(&entity.Knowledgebase{}). + err := db.WithContext(ctx).Model(&entity.Knowledgebase{}). Where("LOWER(name) = ? AND tenant_id = ? AND status = ?", strings.ToLower(currentName), tenantID, string(entity.StatusValid)). Count(&count).Error if err != nil || count == 0 { @@ -364,31 +365,31 @@ func splitNameCounter(name string) (string, int) { // AtomicIncreaseDocNumByID atomically increments the document count // This matches the Python atomic_increase_doc_num_by_id method -func (dao *KnowledgebaseDAO) AtomicIncreaseDocNumByID(kbID string) error { - return DB.Model(&entity.Knowledgebase{}). +func (dao *KnowledgebaseDAO) AtomicIncreaseDocNumByID(ctx context.Context, db *gorm.DB, kbID string) error { + return db.WithContext(ctx).Model(&entity.Knowledgebase{}). Where("id = ?", kbID). Updates(map[string]interface{}{ - "doc_num": DB.Raw("doc_num + 1"), + "doc_num": db.Raw("doc_num + 1"), }).Error } // DecreaseDocumentNum decreases document, chunk, and token counts // This matches the Python decrease_document_num_in_delete method -func (dao *KnowledgebaseDAO) DecreaseDocumentNum(kbID string, docNum, chunkNum, tokenNum int64) error { - return DB.Model(&entity.Knowledgebase{}). +func (dao *KnowledgebaseDAO) DecreaseDocumentNum(ctx context.Context, db *gorm.DB, kbID string, docNum, chunkNum, tokenNum int64) error { + return db.WithContext(ctx).Model(&entity.Knowledgebase{}). Where("id = ?", kbID). Updates(map[string]interface{}{ - "doc_num": DB.Raw("doc_num - ?", docNum), - "chunk_num": DB.Raw("chunk_num - ?", chunkNum), - "token_num": DB.Raw("token_num - ?", tokenNum), + "doc_num": db.Raw("doc_num - ?", docNum), + "chunk_num": db.Raw("chunk_num - ?", chunkNum), + "token_num": db.Raw("token_num - ?", tokenNum), }).Error } // GetKBIDsByTenantID retrieves all knowledge base IDs for a tenant // This matches the Python get_kb_ids method -func (dao *KnowledgebaseDAO) GetKBIDsByTenantID(tenantID string) ([]string, error) { +func (dao *KnowledgebaseDAO) GetKBIDsByTenantID(ctx context.Context, db *gorm.DB, tenantID string) ([]string, error) { var kbIDs []string - err := DB.Model(&entity.Knowledgebase{}). + err := db.WithContext(ctx).Model(&entity.Knowledgebase{}). Where("tenant_id = ? AND status = ?", tenantID, string(entity.StatusValid)). Pluck("id", &kbIDs).Error return kbIDs, err @@ -396,9 +397,9 @@ func (dao *KnowledgebaseDAO) GetKBIDsByTenantID(tenantID string) ([]string, erro // GetAllIDs retrieves all knowledge base IDs // This matches the Python get_all_ids method -func (dao *KnowledgebaseDAO) GetAllIDs() ([]string, error) { +func (dao *KnowledgebaseDAO) GetAllIDs(ctx context.Context, db *gorm.DB) ([]string, error) { var kbIDs []string - err := DB.Model(&entity.Knowledgebase{}). + err := db.WithContext(ctx).Model(&entity.Knowledgebase{}). Where("status = ?", string(entity.StatusValid)). Pluck("id", &kbIDs).Error return kbIDs, err @@ -406,29 +407,29 @@ func (dao *KnowledgebaseDAO) GetAllIDs() ([]string, error) { // UpdateParserConfig updates the parser configuration with deep merge // This matches the Python update_parser_config method -func (dao *KnowledgebaseDAO) UpdateParserConfig(id string, config map[string]interface{}) error { +func (dao *KnowledgebaseDAO) UpdateParserConfig(ctx context.Context, db *gorm.DB, id string, config map[string]interface{}) error { var kb entity.Knowledgebase - if err := DB.Where("id = ? AND status = ?", id, string(entity.StatusValid)).First(&kb).Error; err != nil { + if err := db.WithContext(ctx).Where("id = ? AND status = ?", id, string(entity.StatusValid)).First(&kb).Error; err != nil { return err } mergedConfig := mergeConfig(kb.ParserConfig, config) - return DB.Model(&entity.Knowledgebase{}). + return db.WithContext(ctx).Model(&entity.Knowledgebase{}). Where("id = ?", id). Update("parser_config", mergedConfig).Error } // DeleteFieldMap removes the field_map from parser_config // This matches the Python delete_field_map method -func (dao *KnowledgebaseDAO) DeleteFieldMap(id string) error { +func (dao *KnowledgebaseDAO) DeleteFieldMap(ctx context.Context, db *gorm.DB, id string) error { var kb entity.Knowledgebase - if err := DB.Where("id = ? AND status = ?", id, string(entity.StatusValid)).First(&kb).Error; err != nil { + if err := db.WithContext(ctx).Where("id = ? AND status = ?", id, string(entity.StatusValid)).First(&kb).Error; err != nil { return err } if kb.ParserConfig != nil { delete(kb.ParserConfig, "field_map") - return DB.Model(&entity.Knowledgebase{}). + return db.WithContext(ctx).Model(&entity.Knowledgebase{}). Where("id = ?", id). Update("parser_config", kb.ParserConfig).Error } @@ -437,9 +438,9 @@ func (dao *KnowledgebaseDAO) DeleteFieldMap(id string) error { // GetFieldMap retrieves field mappings from multiple knowledge bases // This matches the Python get_field_map method -func (dao *KnowledgebaseDAO) GetFieldMap(ids []string) (map[string]interface{}, error) { +func (dao *KnowledgebaseDAO) GetFieldMap(ctx context.Context, db *gorm.DB, ids []string) (map[string]interface{}, error) { conf := make(map[string]interface{}) - kbs, err := dao.GetByIDs(ids) + kbs, err := dao.GetByIDs(ctx, db, ids) if err != nil { return nil, err } @@ -460,9 +461,9 @@ func (dao *KnowledgebaseDAO) GetFieldMap(ids []string) (map[string]interface{}, // GetKBByIDAndUserID retrieves a knowledge base by ID and user ID with tenant join // This matches the Python get_kb_by_id method -func (dao *KnowledgebaseDAO) GetKBByIDAndUserID(kbID, userID string) ([]*entity.Knowledgebase, error) { +func (dao *KnowledgebaseDAO) GetKBByIDAndUserID(ctx context.Context, db *gorm.DB, kbID, userID string) ([]*entity.Knowledgebase, error) { var kbs []*entity.Knowledgebase - err := DB.Model(&entity.Knowledgebase{}). + err := db.WithContext(ctx).Model(&entity.Knowledgebase{}). Joins("JOIN user_tenant ON user_tenant.tenant_id = knowledgebase.tenant_id"). Where("knowledgebase.id = ? AND user_tenant.user_id = ?", kbID, userID). Limit(1). @@ -472,9 +473,9 @@ func (dao *KnowledgebaseDAO) GetKBByIDAndUserID(kbID, userID string) ([]*entity. // GetKBByNameAndUserID retrieves a knowledge base by name and user ID with tenant join // This matches the Python get_kb_by_name method -func (dao *KnowledgebaseDAO) GetKBByNameAndUserID(kbName, userID string) ([]*entity.Knowledgebase, error) { +func (dao *KnowledgebaseDAO) GetKBByNameAndUserID(ctx context.Context, db *gorm.DB, kbName, userID string) ([]*entity.Knowledgebase, error) { var kbs []*entity.Knowledgebase - err := DB.Model(&entity.Knowledgebase{}). + err := db.WithContext(ctx).Model(&entity.Knowledgebase{}). Joins("JOIN user_tenant ON user_tenant.tenant_id = knowledgebase.tenant_id"). Where("knowledgebase.name = ? AND user_tenant.user_id = ?", kbName, userID). Limit(1). @@ -484,11 +485,11 @@ func (dao *KnowledgebaseDAO) GetKBByNameAndUserID(kbName, userID string) ([]*ent // GetList retrieves knowledge bases with filtering by ID and name // This matches the Python get_list method -func (dao *KnowledgebaseDAO) GetList(tenantIDs []string, userID string, pageNumber, itemsPerPage int, orderby string, desc bool, id, name string) ([]*entity.Knowledgebase, int64, error) { +func (dao *KnowledgebaseDAO) GetList(ctx context.Context, db *gorm.DB, tenantIDs []string, userID string, pageNumber, itemsPerPage int, orderby string, desc bool, id, name string) ([]*entity.Knowledgebase, int64, error) { var kbs []*entity.Knowledgebase var total int64 - query := DB.Model(&entity.Knowledgebase{}). + query := db.WithContext(ctx).Model(&entity.Knowledgebase{}). Where("((tenant_id IN ? AND permission = ?) OR tenant_id = ?) AND status = ?", tenantIDs, string(entity.TenantPermissionTeam), userID, string(entity.StatusValid)) @@ -561,15 +562,15 @@ func mergeConfig(old, new map[string]interface{}) map[string]interface{} { } // DeleteByTenantID deletes all knowledge bases by tenant ID (hard delete) -func (dao *KnowledgebaseDAO) DeleteByTenantID(tenantID string) (int64, error) { - result := DB.Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.Knowledgebase{}) +func (dao *KnowledgebaseDAO) DeleteByTenantID(ctx context.Context, db *gorm.DB, tenantID string) (int64, error) { + result := db.WithContext(ctx).Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.Knowledgebase{}) return result.RowsAffected, result.Error } -// GetKBIDsByTenantID gets all knowledge base IDs by tenant ID -func (dao *KnowledgebaseDAO) GetKBIDsByTenantIDSimple(tenantID string) ([]string, error) { +// GetKBIDsByTenantIDSimple GetKBIDsByTenantID gets all knowledge base IDs by tenant ID +func (dao *KnowledgebaseDAO) GetKBIDsByTenantIDSimple(ctx context.Context, db *gorm.DB, tenantID string) ([]string, error) { var kbIDs []string - err := DB.Model(&entity.Knowledgebase{}). + err := db.WithContext(ctx).Model(&entity.Knowledgebase{}). Where("tenant_id = ?", tenantID). Pluck("id", &kbIDs).Error return kbIDs, err diff --git a/internal/dao/kb_test.go b/internal/dao/kb_test.go index 47b0704e25..79924179cb 100644 --- a/internal/dao/kb_test.go +++ b/internal/dao/kb_test.go @@ -73,12 +73,13 @@ func TestKnowledgebaseDAO_DecreaseDocumentNum(t *testing.T) { testKnowledgebase(t, db, "kb-1", 5, 100, 50) - err := dao.DecreaseDocumentNum("kb-1", 1, 20, 10) + ctx := t.Context() + err := dao.DecreaseDocumentNum(ctx, db, "kb-1", 1, 20, 10) if err != nil { t.Fatalf("DecreaseDocumentNum failed: %v", err) } - kb, err := dao.GetByID("kb-1") + kb, err := dao.GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("GetByID failed: %v", err) } @@ -98,14 +99,16 @@ func TestKnowledgebaseDAO_DecreaseDocumentNum_ZeroDecrement(t *testing.T) { pushDB(t, db) dao := NewKnowledgebaseDAO() + ctx := t.Context() + testKnowledgebase(t, db, "kb-2", 3, 60, 15) - err := dao.DecreaseDocumentNum("kb-2", 0, 0, 0) + err := dao.DecreaseDocumentNum(ctx, db, "kb-2", 0, 0, 0) if err != nil { t.Fatalf("DecreaseDocumentNum failed: %v", err) } - kb, err := dao.GetByID("kb-2") + kb, err := dao.GetByID(ctx, db, "kb-2") if err != nil { t.Fatalf("GetByID failed: %v", err) } diff --git a/internal/engine/elasticsearch/metadata.go b/internal/engine/elasticsearch/metadata.go index 2d231fbe96..4cf462aed4 100644 --- a/internal/engine/elasticsearch/metadata.go +++ b/internal/engine/elasticsearch/metadata.go @@ -25,6 +25,8 @@ import ( "strings" "github.com/elastic/go-elasticsearch/v8/esapi" + "gorm.io/gorm" + "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/engine/types" @@ -720,7 +722,7 @@ const metaPushdownMaxSize = 10000 // nil -> push-down was not viable / errored / result overflowed the // push-down cap (caller should fall back to in-memory) // []string{} -> push-down succeeded but found 0 matching docs (empty result is definitive) -func (e *elasticsearchEngine) FilterDocIdsByMetaPushdown(ctx context.Context, kbIDs []string, conditions []map[string]interface{}, logic string) []string { +func (e *elasticsearchEngine) FilterDocIdsByMetaPushdown(ctx context.Context, sqlDB *gorm.DB, kbIDs []string, conditions []map[string]interface{}, logic string) []string { if len(conditions) == 0 || len(kbIDs) == 0 { return nil } @@ -733,7 +735,7 @@ func (e *elasticsearchEngine) FilterDocIdsByMetaPushdown(ctx context.Context, kb // Execute search for each tenant (use first KB to get tenant) kb := kbIDs[0] - tenantID, err := dao.GetTenantIDByKBID(kb) + tenantID, err := dao.GetTenantIDByKBID(ctx, sqlDB, kb) if err != nil { common.Warn("FilterDocIdsByMetaPushdown: failed to get tenant for KB", zap.String("kbID", kb), zap.Error(err)) return nil diff --git a/internal/engine/engine.go b/internal/engine/engine.go index 0e4956d2bf..9040bbaa29 100644 --- a/internal/engine/engine.go +++ b/internal/engine/engine.go @@ -20,6 +20,8 @@ import ( "context" "ragflow/internal/common" "ragflow/internal/engine/types" + + "gorm.io/gorm" ) // EngineType document engine type @@ -83,7 +85,7 @@ type DocEngine interface { // the doc metadata index, returning matching doc IDs or nil if push-down // is not supported (caller should fall back to in-memory filtering). // conditions is a list of filter objects with keys: key, op, value - FilterDocIdsByMetaPushdown(ctx context.Context, kbIDs []string, conditions []map[string]interface{}, logic string) []string + FilterDocIdsByMetaPushdown(ctx context.Context, sqlDB *gorm.DB, kbIDs []string, conditions []map[string]interface{}, logic string) []string } // Type returns the engine type (helper method for runtime type checking) diff --git a/internal/engine/infinity/metadata.go b/internal/engine/infinity/metadata.go index 774612f560..2d5b8691c4 100644 --- a/internal/engine/infinity/metadata.go +++ b/internal/engine/infinity/metadata.go @@ -29,6 +29,7 @@ import ( "ragflow/internal/utility" infinity "github.com/infiniflow/infinity-go-sdk" + "gorm.io/gorm" "go.uber.org/zap" ) @@ -794,7 +795,7 @@ const metaPushdownMaxSize = 10000 // nil -> push-down was not viable / errored / result overflowed the // push-down cap (caller should fall back to in-memory) // []string{} -> push-down succeeded but found 0 matching docs (empty result is definitive) -func (e *infinityEngine) FilterDocIdsByMetaPushdown(ctx context.Context, kbIDs []string, conditions []map[string]interface{}, logic string) []string { +func (e *infinityEngine) FilterDocIdsByMetaPushdown(ctx context.Context, sqlDB *gorm.DB, kbIDs []string, conditions []map[string]interface{}, logic string) []string { if len(conditions) == 0 || len(kbIDs) == 0 { return nil } @@ -806,7 +807,7 @@ func (e *infinityEngine) FilterDocIdsByMetaPushdown(ctx context.Context, kbIDs [ } // Get tenant ID from first KB - tenantID, err := dao.GetTenantIDByKBID(kbIDs[0]) + tenantID, err := dao.GetTenantIDByKBID(ctx, sqlDB, kbIDs[0]) if err != nil { common.Warn("FilterDocIdsByMetaPushdown: failed to get tenant for KB", zap.String("kbID", kbIDs[0]), zap.Error(err)) return nil @@ -836,13 +837,13 @@ func (e *infinityEngine) FilterDocIdsByMetaPushdown(ctx context.Context, kbIDs [ whereClause = kbFilter + " AND (" + whereClause + ")" // Use Infinity connection to execute query - db, release, err := e.client.checkoutDatabase(ctx, "metadata.go") - if err != nil || db == nil { + infinityDB, release, err := e.client.checkoutDatabase(ctx, "metadata.go") + if err != nil || infinityDB == nil { return nil } defer release() - table, err := db.GetTable(tableName) + table, err := infinityDB.GetTable(tableName) if err != nil || table == nil { return nil } diff --git a/internal/handler/agent_component.go b/internal/handler/agent_component.go index 9cb0ae16fa..6ee2181e02 100644 --- a/internal/handler/agent_component.go +++ b/internal/handler/agent_component.go @@ -218,7 +218,7 @@ func (h *AgentHandler) DebugComponent(c *gin.Context) { invokeCtx := runtime.WithState(c.Request.Context(), debugState) outputs, err := runtime.TrackElapsed(name, func() (map[string]any, error) { - return comp.Invoke(invokeCtx, inputs) + return comp.Invoke(invokeCtx, dao.DB, inputs) }) if err != nil { common.ResponseWithCodeData(c, common.CodeServerError, nil, "invoke: "+err.Error()) diff --git a/internal/handler/agent_component_test.go b/internal/handler/agent_component_test.go index 7980b63fc5..d9dc21ee0f 100644 --- a/internal/handler/agent_component_test.go +++ b/internal/handler/agent_component_test.go @@ -25,6 +25,7 @@ import ( "testing" "github.com/gin-gonic/gin" + "gorm.io/gorm" _ "ragflow/internal/agent/component" // registers the production factory agentruntime "ragflow/internal/agent/runtime" @@ -329,7 +330,7 @@ func TestDebugComponent_HappyPath_Begin(t *testing.T) { type sysEchoComponent struct{} -func (s *sysEchoComponent) Invoke(ctx context.Context, _ map[string]any) (map[string]any, error) { +func (s *sysEchoComponent) Invoke(ctx context.Context, _ *gorm.DB, _ map[string]any) (map[string]any, error) { state, _, err := agentruntime.GetStateFromContext[*agentruntime.CanvasState](ctx) if err != nil { return nil, err diff --git a/internal/handler/dataset.go b/internal/handler/dataset.go index b21511daeb..9f2ecdc717 100644 --- a/internal/handler/dataset.go +++ b/internal/handler/dataset.go @@ -17,6 +17,7 @@ package handler import ( + "context" "encoding/json" "fmt" "net/http" @@ -42,11 +43,11 @@ type DatasetsHandler struct { } type searchDatasetsService interface { - SearchDatasets(req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error) + SearchDatasets(ctx context.Context, req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error) } type searchDatasetService interface { - SearchDataset(datasetID, userID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error) + SearchDataset(ctx context.Context, datasetID, userID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error) } type listDatasetsExt struct { @@ -116,7 +117,9 @@ func (h *DatasetsHandler) ListDatasets(c *gin.Context) { ownerIDs = ext.OwnerIDs } + ctx := c.Request.Context() data, total, code, err := h.datasetsService.ListDatasets( + ctx, c.Query("id"), c.Query("name"), page, @@ -154,7 +157,9 @@ func (h *DatasetsHandler) CreateDataset(c *gin.Context) { return } - result, code, err := h.datasetsService.CreateDataset(&req, user.ID) + ctx := c.Request.Context() + + result, code, err := h.datasetsService.CreateDataset(ctx, &req, user.ID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -244,8 +249,10 @@ func (h *DatasetsHandler) GetMetadataConfig(c *gin.Context) { return } + ctx := c.Request.Context() + datasetID := c.Param("dataset_id") - result, code, err := h.datasetsService.GetMetadataConfig(datasetID, user.ID) + result, code, err := h.datasetsService.GetMetadataConfig(ctx, datasetID, user.ID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -268,7 +275,9 @@ func (h *DatasetsHandler) UpdateMetadataConfig(c *gin.Context) { return } - result, code, err := h.datasetsService.UpdateMetadataConfig(datasetID, user.ID, &req) + ctx := c.Request.Context() + + result, code, err := h.datasetsService.UpdateMetadataConfig(ctx, datasetID, user.ID, &req) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -336,7 +345,9 @@ func (h *DatasetsHandler) ListIngestionLogs(c *gin.Context) { logType := c.DefaultQuery("log_type", "dataset") keywords := c.Query("keywords") - result, code, err := h.datasetsService.ListIngestionLogs(datasetID, user.ID, page, pageSize, orderby, desc, operationStatus, createDateFrom, createDateTo, logType, keywords) + ctx := c.Request.Context() + + result, code, err := h.datasetsService.ListIngestionLogs(ctx, datasetID, user.ID, page, pageSize, orderby, desc, operationStatus, createDateFrom, createDateTo, logType, keywords) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -353,9 +364,11 @@ func (h *DatasetsHandler) GetIngestionLog(c *gin.Context) { return } + ctx := c.Request.Context() + datasetID := c.Param("dataset_id") logID := c.Param("log_id") - result, code, err := h.datasetsService.GetIngestionLog(datasetID, user.ID, logID) + result, code, err := h.datasetsService.GetIngestionLog(ctx, datasetID, user.ID, logID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -388,7 +401,9 @@ func (h *DatasetsHandler) DeleteDatasets(c *gin.Context) { ids = *req.IDs } - result, code, err := h.datasetsService.DeleteDatasets(ids, req.DeleteAll, user.ID) + ctx := c.Request.Context() + + result, code, err := h.datasetsService.DeleteDatasets(ctx, ids, req.DeleteAll, user.ID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -513,8 +528,10 @@ func (h *DatasetsHandler) ListTags(c *gin.Context) { return } + ctx := c.Request.Context() + datasetID := strings.TrimSpace(c.Param("dataset_id")) - result, code, err := h.datasetsService.ListTags(datasetID, user.ID) + result, code, err := h.datasetsService.ListTags(ctx, datasetID, user.ID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -559,7 +576,9 @@ func (h *DatasetsHandler) RenameTag(c *gin.Context) { return } - result, code, err := h.datasetsService.RenameTag(datasetID, user.ID, req.FromTag, req.ToTag) + ctx := c.Request.Context() + + result, code, err := h.datasetsService.RenameTag(ctx, datasetID, user.ID, req.FromTag, req.ToTag) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -756,7 +775,9 @@ func (h *DatasetsHandler) AggregateTags(c *gin.Context) { return } - result, code, err := h.datasetsService.AggregateTags(datasetIDs, user.ID) + ctx := c.Request.Context() + + result, code, err := h.datasetsService.AggregateTags(ctx, datasetIDs, user.ID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -815,8 +836,10 @@ func (h *DatasetsHandler) TraceIndex(c *gin.Context) { return } + ctx := c.Request.Context() + indexType := strings.ToLower(strings.TrimSpace(c.Query("type"))) - result, code, err := h.datasetsService.TraceIndex(datasetID, userID, indexType) + result, code, err := h.datasetsService.TraceIndex(ctx, datasetID, userID, indexType) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -861,7 +884,9 @@ func (h *DatasetsHandler) DeleteIndex(c *gin.Context) { wipe = false } - code, err := h.datasetsService.DeleteIndex(userID, datasetID, indexType, wipe) + ctx := c.Request.Context() + + code, err := h.datasetsService.DeleteIndex(ctx, userID, datasetID, indexType, wipe) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -905,15 +930,16 @@ func (h *DatasetsHandler) ListMetadataFlattened(c *gin.Context) { return } + ctx := c.Request.Context() // Check access for each dataset for _, datasetID := range datasetIDs { - if !h.datasetsService.Accessible(datasetID, user.ID) { + if !h.datasetsService.Accessible(ctx, datasetID, user.ID) { common.ResponseWithCodeData(c, common.CodeAuthenticationError, nil, "No authorization for dataset: "+datasetID) return } } - flattenedMeta, err := h.metadataService.GetFlattedMetaByKBs(datasetIDs) + flattenedMeta, err := h.metadataService.GetFlattedMetaByKBs(ctx, datasetIDs) if err != nil { common.ResponseWithCodeData(c, common.CodeServerError, nil, "Failed to get metadata: "+err.Error()) return @@ -1026,7 +1052,9 @@ func (h *DatasetsHandler) SearchDatasets(c *gin.Context) { return } - resp, err := searchService.SearchDatasets(&req, user.ID) + ctx := c.Request.Context() + + resp, err := searchService.SearchDatasets(ctx, &req, user.ID) if err != nil { common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error()) return @@ -1080,7 +1108,9 @@ func (h *DatasetsHandler) SearchDataset(c *gin.Context) { return } - resp, err := searchService.SearchDataset(datasetID, user.ID, &req) + ctx := c.Request.Context() + + resp, err := searchService.SearchDataset(ctx, datasetID, user.ID, &req) if err != nil { common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error()) return diff --git a/internal/handler/dataset_search_test.go b/internal/handler/dataset_search_test.go index da430ec298..9eaf94773c 100644 --- a/internal/handler/dataset_search_test.go +++ b/internal/handler/dataset_search_test.go @@ -1,6 +1,7 @@ package handler import ( + "context" "encoding/json" "errors" "net/http" @@ -23,7 +24,7 @@ type fakeSearchDatasetService struct { err error } -func (f *fakeSearchDatasetService) SearchDataset(datasetID, userID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error) { +func (f *fakeSearchDatasetService) SearchDataset(ctx context.Context, datasetID, userID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error) { f.datasetID = datasetID f.userID = userID f.req = req @@ -37,7 +38,7 @@ type fakeSearchDatasetsService struct { err error } -func (f *fakeSearchDatasetsService) SearchDatasets(req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error) { +func (f *fakeSearchDatasetsService) SearchDatasets(ctx context.Context, req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error) { f.userID = userID f.req = req return f.resp, f.err diff --git a/internal/handler/dify_retrieval_handler.go b/internal/handler/dify_retrieval_handler.go index 948fb4e52e..d022303987 100644 --- a/internal/handler/dify_retrieval_handler.go +++ b/internal/handler/dify_retrieval_handler.go @@ -43,21 +43,21 @@ import ( // KBServiceIface abstracts KnowledgebaseService for the Dify handler. type KBServiceIface interface { - GetByID(kbID string) (*entity.Knowledgebase, error) - Accessible(kbID, userID string) bool + GetByID(ctx context.Context, kbID string) (*entity.Knowledgebase, error) + Accessible(ctx context.Context, kbID, userID string) bool } // ModelServiceIface abstracts ModelProviderService for the Dify handler. type ModelServiceIface interface { - GetEmbeddingModel(tenantID, embdID string) (*modelModule.EmbeddingModel, error) - GetChatModel(tenantID, compositeModelName string) (*modelModule.ChatModel, error) + GetEmbeddingModel(ctx context.Context, tenantID, embdID string) (*modelModule.EmbeddingModel, error) + GetChatModel(ctx context.Context, tenantID, compositeModelName string) (*modelModule.ChatModel, error) } // MetadataServiceIface abstracts MetadataService for the Dify handler. type MetadataServiceIface interface { - GetFlattedMetaByKBs(kbIDs []string) (common.MetaData, error) - SearchMetadataByKBs(kbIDs []string, size int) (*service.SearchMetadataResponse, error) - LabelQuestion(question string, kbs []*entity.Knowledgebase) map[string]float64 + GetFlattedMetaByKBs(ctx context.Context, kbIDs []string) (common.MetaData, error) + SearchMetadataByKBs(ctx context.Context, kbIDs []string, size int) (*service.SearchMetadataResponse, error) + LabelQuestion(ctx context.Context, question string, kbs []*entity.Knowledgebase) map[string]float64 } // RetrievalServiceIface abstracts RetrievalService for the Dify handler. @@ -204,7 +204,8 @@ func (h *DifyRetrievalHandler) Retrieval(c *gin.Context) { return } - kb, err := h.kbSvc.GetByID(req.KnowledgeID) + ctx := c.Request.Context() + kb, err := h.kbSvc.GetByID(ctx, req.KnowledgeID) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { common.ResponseWithHttpCodeData(c, http.StatusNotFound, common.CodeNotFound, nil, "Knowledge base not found!") @@ -214,7 +215,7 @@ func (h *DifyRetrievalHandler) Retrieval(c *gin.Context) { return } - if !h.kbSvc.Accessible(req.KnowledgeID, user.ID) { + if !h.kbSvc.Accessible(ctx, req.KnowledgeID, user.ID) { common.ResponseWithHttpCodeData(c, http.StatusUnauthorized, common.CodeAuthenticationError, nil, "No authorization") return } @@ -234,14 +235,14 @@ func (h *DifyRetrievalHandler) Retrieval(c *gin.Context) { } // Get embedding model - embModel, err := h.modelSvc.GetEmbeddingModel(kb.TenantID, kb.EmbdID) + embModel, err := h.modelSvc.GetEmbeddingModel(ctx, kb.TenantID, kb.EmbdID) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, fmt.Sprintf("failed to get embedding model: %v", err)) return } // Metadata filter - metas, metaErr := h.metadataSvc.GetFlattedMetaByKBs([]string{req.KnowledgeID}) + metas, metaErr := h.metadataSvc.GetFlattedMetaByKBs(ctx, []string{req.KnowledgeID}) docIDs := make([]string, 0) if metaErr == nil && req.MetadataCondition != nil { logic := req.MetadataCondition.Logic @@ -257,7 +258,7 @@ func (h *DifyRetrievalHandler) Retrieval(c *gin.Context) { // Label question for rank features kbs := []*entity.Knowledgebase{kb} - rankFeature := h.metadataSvc.LabelQuestion(req.Query, kbs) + rankFeature := h.metadataSvc.LabelQuestion(ctx, req.Query, kbs) // Chunk retrieval sr := &nlp.RetrievalRequest{ @@ -290,7 +291,7 @@ func (h *DifyRetrievalHandler) Retrieval(c *gin.Context) { // KG retrieval (optional) if req.UseKG { - chatModel, kgErr := h.modelSvc.GetChatModel(kb.TenantID, "") + chatModel, kgErr := h.modelSvc.GetChatModel(ctx, kb.TenantID, "") if kgErr != nil { common.Warn("KG retrieval: failed to get chat model", zap.String("kbID", req.KnowledgeID), zap.Error(kgErr)) } else if chatModel != nil { @@ -322,7 +323,6 @@ func (h *DifyRetrievalHandler) Retrieval(c *gin.Context) { allDocIDs = append(allDocIDs, id) } - ctx := c.Request.Context() docMap := make(map[string]*entity.Document) if len(allDocIDs) > 0 { var docs []*entity.Document @@ -337,7 +337,7 @@ func (h *DifyRetrievalHandler) Retrieval(c *gin.Context) { } metaMap := make(map[string]map[string]interface{}) - metaResult, err := h.metadataSvc.SearchMetadataByKBs([]string{kb.ID}, 10000) + metaResult, err := h.metadataSvc.SearchMetadataByKBs(ctx, []string{kb.ID}, 10000) if err == nil { for _, metadata := range metaResult.MetadataRecords { docID, ok := service.ExtractDocumentID(metadata) diff --git a/internal/handler/dify_retrieval_handler_test.go b/internal/handler/dify_retrieval_handler_test.go index 55c5666db6..383df813af 100644 --- a/internal/handler/dify_retrieval_handler_test.go +++ b/internal/handler/dify_retrieval_handler_test.go @@ -42,63 +42,63 @@ import ( type mockKBService struct { KBServiceIface - getByIDFn func(kbID string) (*entity.Knowledgebase, error) - accessibleFn func(kbID, userID string) bool + getByIDFn func(ctx context.Context, kbID string) (*entity.Knowledgebase, error) + accessibleFn func(ctx context.Context, kbID, userID string) bool } -func (m *mockKBService) GetByID(kbID string) (*entity.Knowledgebase, error) { +func (m *mockKBService) GetByID(ctx context.Context, kbID string) (*entity.Knowledgebase, error) { if m.getByIDFn != nil { - return m.getByIDFn(kbID) + return m.getByIDFn(ctx, kbID) } return &entity.Knowledgebase{ ID: kbID, TenantID: "tenant1", EmbdID: "text-embedding", }, nil } -func (m *mockKBService) Accessible(kbID, userID string) bool { +func (m *mockKBService) Accessible(ctx context.Context, kbID, userID string) bool { if m.accessibleFn != nil { - return m.accessibleFn(kbID, userID) + return m.accessibleFn(ctx, kbID, userID) } return true } type mockModelService struct { ModelServiceIface - getEmbeddingFn func(tenantID, embdID string) (*modelModule.EmbeddingModel, error) - getChatModelFn func(tenantID, llmID string) (*modelModule.ChatModel, error) + getEmbeddingFn func(ctx context.Context, tenantID, embdID string) (*modelModule.EmbeddingModel, error) + getChatModelFn func(ctx context.Context, tenantID, llmID string) (*modelModule.ChatModel, error) } -func (m *mockModelService) GetEmbeddingModel(tenantID, embdID string) (*modelModule.EmbeddingModel, error) { +func (m *mockModelService) GetEmbeddingModel(ctx context.Context, tenantID, embdID string) (*modelModule.EmbeddingModel, error) { if m.getEmbeddingFn != nil { - return m.getEmbeddingFn(tenantID, embdID) + return m.getEmbeddingFn(ctx, tenantID, embdID) } return &modelModule.EmbeddingModel{}, nil } -func (m *mockModelService) GetChatModel(tenantID, llmID string) (*modelModule.ChatModel, error) { +func (m *mockModelService) GetChatModel(ctx context.Context, tenantID, llmID string) (*modelModule.ChatModel, error) { if m.getChatModelFn != nil { - return m.getChatModelFn(tenantID, llmID) + return m.getChatModelFn(ctx, tenantID, llmID) } return &modelModule.ChatModel{}, nil } type mockMetadataService struct { MetadataServiceIface - getFlattedMetaFn func(kbIDs []string) (common.MetaData, error) - searchMetadataByKBs func(kbIDs []string, size int) (*service.SearchMetadataResponse, error) - labelQuestionFn func(question string, kbs []*entity.Knowledgebase) map[string]float64 + getFlattedMetaFn func(ctx context.Context, kbIDs []string) (common.MetaData, error) + searchMetadataByKBs func(ctx context.Context, kbIDs []string, size int) (*service.SearchMetadataResponse, error) + labelQuestionFn func(ctx context.Context, question string, kbs []*entity.Knowledgebase) map[string]float64 } -func (m *mockMetadataService) GetFlattedMetaByKBs(kbIDs []string) (common.MetaData, error) { +func (m *mockMetadataService) GetFlattedMetaByKBs(ctx context.Context, kbIDs []string) (common.MetaData, error) { if m.getFlattedMetaFn != nil { - return m.getFlattedMetaFn(kbIDs) + return m.getFlattedMetaFn(ctx, kbIDs) } return common.MetaData{}, nil } -func (m *mockMetadataService) SearchMetadataByKBs(kbIDs []string, size int) (*service.SearchMetadataResponse, error) { +func (m *mockMetadataService) SearchMetadataByKBs(ctx context.Context, kbIDs []string, size int) (*service.SearchMetadataResponse, error) { if m.searchMetadataByKBs != nil { - return m.searchMetadataByKBs(kbIDs, size) + return m.searchMetadataByKBs(ctx, kbIDs, size) } return &service.SearchMetadataResponse{ MetadataRecords: []map[string]interface{}{ @@ -107,9 +107,9 @@ func (m *mockMetadataService) SearchMetadataByKBs(kbIDs []string, size int) (*se }, nil } -func (m *mockMetadataService) LabelQuestion(question string, kbs []*entity.Knowledgebase) map[string]float64 { +func (m *mockMetadataService) LabelQuestion(ctx context.Context, question string, kbs []*entity.Knowledgebase) map[string]float64 { if m.labelQuestionFn != nil { - return m.labelQuestionFn(question, kbs) + return m.labelQuestionFn(ctx, question, kbs) } return nil } @@ -275,7 +275,7 @@ func TestDifyRetrieval_MissingArgs(t *testing.T) { func TestDifyRetrieval_KBNotFound(t *testing.T) { h, r := setupDifyTest("user1") h.kbSvc = &mockKBService{ - getByIDFn: func(kbID string) (*entity.Knowledgebase, error) { + getByIDFn: func(ctx context.Context, kbID string) (*entity.Knowledgebase, error) { return nil, gorm.ErrRecordNotFound }, } @@ -304,7 +304,7 @@ func TestDifyRetrieval_NoAuth(t *testing.T) { func TestDifyRetrieval_Unauthorized(t *testing.T) { h, r := setupDifyTest("user1") h.kbSvc = &mockKBService{ - accessibleFn: func(kbID, userID string) bool { return false }, + accessibleFn: func(ctx context.Context, kbID, userID string) bool { return false }, } w := httptest.NewRecorder() body := `{"knowledge_id": "kb1", "query": "test"}` @@ -319,7 +319,7 @@ func TestDifyRetrieval_Unauthorized(t *testing.T) { func TestDifyRetrieval_WithMetadataFilter(t *testing.T) { h, r := setupDifyTest("user1") h.metadataSvc = &mockMetadataService{ - getFlattedMetaFn: func(kbIDs []string) (common.MetaData, error) { + getFlattedMetaFn: func(ctx context.Context, kbIDs []string) (common.MetaData, error) { return common.MetaData{}, nil }, } @@ -347,7 +347,7 @@ func TestDifyRetrieval_InvalidJSON(t *testing.T) { func TestDifyRetrieval_UseKG(t *testing.T) { h, r := setupDifyTest("user1") h.metadataSvc = &mockMetadataService{ - labelQuestionFn: func(question string, kbs []*entity.Knowledgebase) map[string]float64 { + labelQuestionFn: func(ctx context.Context, question string, kbs []*entity.Knowledgebase) map[string]float64 { return map[string]float64{"tag_1": 0.8} }, } @@ -366,7 +366,7 @@ func strPtr(s string) *string { return &s } func TestDifyRetrieval_KBDBError(t *testing.T) { h, r := setupDifyTest("user1") h.kbSvc = &mockKBService{ - getByIDFn: func(kbID string) (*entity.Knowledgebase, error) { + getByIDFn: func(ctx context.Context, kbID string) (*entity.Knowledgebase, error) { return nil, errors.New("connection refused") }, } diff --git a/internal/handler/document.go b/internal/handler/document.go index 4e67d7579d..5db4f9b30a 100644 --- a/internal/handler/document.go +++ b/internal/handler/document.go @@ -308,7 +308,7 @@ func (h *DocumentHandler) UpdateDocument(c *gin.Context) { common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 1, nil, "document not found!") return } - if !h.datasetService.Accessible(doc.KbID, user.ID) { + if !h.datasetService.Accessible(ctx, doc.KbID, user.ID) { common.ResponseWithCodeData(c, common.CodeAuthenticationError, nil, "No authorization.") return } @@ -361,7 +361,7 @@ func (h *DocumentHandler) DeleteDocument(c *gin.Context) { common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 1, nil, "document not found!") return } - if !h.datasetService.Accessible(doc.KbID, user.ID) { + if !h.datasetService.Accessible(ctx, doc.KbID, user.ID) { common.ResponseWithCodeData(c, common.CodeAuthenticationError, nil, "No authorization.") return } @@ -505,7 +505,8 @@ func (h *DocumentHandler) ListDocuments(c *gin.Context) { userID := c.GetString("user_id") - if !h.datasetService.Accessible(datasetID, userID) { + ctx := c.Request.Context() + if !h.datasetService.Accessible(ctx, datasetID, userID) { common.ResponseWithCodeData(c, common.CodeDataError, nil, fmt.Sprintf("You don't own the dataset %s.", datasetID)) return } @@ -528,7 +529,6 @@ func (h *DocumentHandler) ListDocuments(c *gin.Context) { return } - ctx := c.Request.Context() if c.Query("type") == "filter" { filters, total, err := h.documentService.GetDocumentFiltersByDatasetID(ctx, opts) if err != nil { @@ -830,7 +830,8 @@ func (h *DocumentHandler) UploadDocuments(c *gin.Context) { datasetID := c.Param("dataset_id") uploadType := strings.ToLower(c.DefaultQuery("type", "local")) - kb, err := h.datasetService.GetKnowledgebaseByID(datasetID) + ctx := c.Request.Context() + kb, err := h.datasetService.GetKnowledgebaseByID(ctx, datasetID) if err != nil || kb == nil { common.ResponseWithCodeData(c, common.CodeDataError, nil, fmt.Sprintf("Can't find the dataset with ID %s!", datasetID)) return @@ -1236,7 +1237,7 @@ func (h *DocumentHandler) SetMeta(c *gin.Context) { common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 1, nil, "document not found") return } - if !h.datasetService.Accessible(doc.KbID, user.ID) { + if !h.datasetService.Accessible(ctx, doc.KbID, user.ID) { common.ResponseWithCodeData(c, common.CodeAuthenticationError, nil, "No authorization.") return } @@ -1331,7 +1332,7 @@ func (h *DocumentHandler) DeleteMeta(c *gin.Context) { common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 1, nil, "document not found") return } - if !h.datasetService.Accessible(doc.KbID, user.ID) { + if !h.datasetService.Accessible(ctx, doc.KbID, user.ID) { common.ResponseWithCodeData(c, common.CodeAuthenticationError, nil, "No authorization.") return } @@ -1392,13 +1393,13 @@ func (h *DocumentHandler) ListIngestionTasks(c *gin.Context) { var parseResult []*entity.IngestionTask var err error + ctx := c.Request.Context() if req.DatasetID != nil { - if !h.datasetService.Accessible(*req.DatasetID, userID) { + if !h.datasetService.Accessible(ctx, *req.DatasetID, userID) { common.ResponseWithCodeData(c, common.CodeAuthenticationError, nil, "No authorization to access the dataset.") return } } - ctx := c.Request.Context() parseResult, err = h.documentService.ListIngestionTasks(ctx, userID, req.DatasetID, 0, 0) if err != nil { common.ResponseWithCodeData(c, IngestionTaskErrorCode(err), nil, err.Error()) @@ -1422,12 +1423,12 @@ func (h *DocumentHandler) StartIngestionTask(c *gin.Context) { } userID := c.GetString("user_id") - - if !h.datasetService.Accessible(datasetID, userID) { + ctx := c.Request.Context() + if !h.datasetService.Accessible(ctx, datasetID, userID) { common.ResponseWithCodeData(c, common.CodeDataError, nil, fmt.Sprintf("You don't own the dataset %s.", datasetID)) return } - ctx := c.Request.Context() + parseResult, err := h.documentService.IngestDocuments(ctx, datasetID, userID, req.DocumentIDs) if err != nil { common.ResponseWithCodeData(c, common.CodeExceptionError, nil, err.Error()) @@ -1505,12 +1506,12 @@ func (h *DocumentHandler) ParseDocuments(c *gin.Context) { } userID := c.GetString("user_id") - - if !h.datasetService.Accessible(datasetID, userID) { + ctx := c.Request.Context() + if !h.datasetService.Accessible(ctx, datasetID, userID) { common.ResponseWithCodeData(c, common.CodeAuthenticationError, nil, "No authorization to access the dataset.") return } - ctx := c.Request.Context() + parseResult, err := h.documentService.ParseDocuments(ctx, datasetID, userID, req.Documents) if err != nil { common.ResponseWithCodeData(c, common.CodeExceptionError, nil, err.Error()) @@ -1538,12 +1539,12 @@ func (h *DocumentHandler) StopParseDocuments(c *gin.Context) { } userID := c.GetString("user_id") - - if !h.datasetService.Accessible(datasetID, userID) { + ctx := c.Request.Context() + if !h.datasetService.Accessible(ctx, datasetID, userID) { common.ResponseWithCodeData(c, common.CodeDataError, nil, fmt.Sprintf("You don't own the dataset %s.", datasetID)) return } - ctx := c.Request.Context() + result, err := h.documentService.StopParseDocuments(ctx, datasetID, req.DocumentIDs) if err != nil { common.ResponseWithCodeData(c, common.CodeExceptionError, nil, err.Error()) @@ -1564,7 +1565,8 @@ func (h *DocumentHandler) MetadataSummaryByDataset(c *gin.Context) { common.ErrorWithCode(c, common.CodeServerError, "dataset_id is required") return } - if !h.datasetService.Accessible(datasetID, user.ID) { + ctx := c.Request.Context() + if !h.datasetService.Accessible(ctx, datasetID, user.ID) { common.ErrorWithCode(c, common.CodeServerError, "You don't own the dataset "+datasetID) return } @@ -1573,7 +1575,7 @@ func (h *DocumentHandler) MetadataSummaryByDataset(c *gin.Context) { if docIDsParam := c.Query("doc_ids"); docIDsParam != "" { docIDS = strings.Split(docIDsParam, ",") } - ctx := c.Request.Context() + summary, err := h.documentService.GetMetadataSummary(ctx, datasetID, docIDS) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, "Failed to get metadata summary"+err.Error()) @@ -1710,7 +1712,8 @@ func (h *DocumentHandler) handleBatchUpdateDocumentMetadatas(c *gin.Context) { common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "dataset_id is required") return } - if !h.datasetService.Accessible(datasetID, user.ID) { + ctx := c.Request.Context() + if !h.datasetService.Accessible(ctx, datasetID, user.ID) { common.ResponseWithCodeData(c, common.CodeDataError, nil, "You don't own the dataset "+datasetID+".") return } @@ -1742,7 +1745,7 @@ func (h *DocumentHandler) handleBatchUpdateDocumentMetadatas(c *gin.Context) { Updates: updates, Deletes: deletes, } - ctx := c.Request.Context() + resp, code, err := h.documentService.BatchUpdateDocumentMetadatas(ctx, datasetID, req.Selector, req.Updates, req.Deletes) if err != nil { common.ErrorWithCode(c, code, err.Error()) diff --git a/internal/handler/file_commit.go b/internal/handler/file_commit.go index 071fec237e..2da5aef295 100644 --- a/internal/handler/file_commit.go +++ b/internal/handler/file_commit.go @@ -94,7 +94,7 @@ func CommitFolderResolver(h *FileCommitHandler, entityType, urlParam string) gin } func (h *FileCommitHandler) resolveDatasetFolderID(ctx context.Context, datasetID string) (string, error) { - kb, err := h.kbDAO.GetByID(datasetID) + kb, err := h.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil { return "", err } diff --git a/internal/handler/mcp_server.go b/internal/handler/mcp_server.go index 10da9c86ef..fe0334ce80 100644 --- a/internal/handler/mcp_server.go +++ b/internal/handler/mcp_server.go @@ -41,18 +41,18 @@ type MCPRetrievalService interface { // MCPServerHandler handles MCP protocol requests (JSON-RPC over HTTP). // It exposes RAGFlow capabilities as MCP tools to external AI clients. type MCPServerHandler struct { - listDatasetsFunc func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) - listChatsFunc func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) - retrievalFunc func(userID string, req mcp.RetrievalRequest) (string, error) + listDatasetsFunc func(ctx context.Context, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) + listChatsFunc func(ctx context.Context, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) + retrievalFunc func(ctx context.Context, userID string, req mcp.RetrievalRequest) (string, error) } // NewMCPServerHandler creates a new MCPServerHandler. // The service functions are passed as closures to avoid importing the service // package directly from the handler layer. func NewMCPServerHandler( - listDatasetsFunc func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error), - listChatsFunc func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error), - retrievalFunc func(userID string, req mcp.RetrievalRequest) (string, error), + listDatasetsFunc func(ctx context.Context, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error), + listChatsFunc func(ctx context.Context, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error), + retrievalFunc func(ctx context.Context, userID string, req mcp.RetrievalRequest) (string, error), ) *MCPServerHandler { return &MCPServerHandler{ listDatasetsFunc: listDatasetsFunc, @@ -95,8 +95,9 @@ func (h *MCPServerHandler) HandleMCP(c *gin.Context) { h.retrievalFunc, ) + ctx := c.Request.Context() server := mcp.NewServer(connector) - respBody, hasResponse, err := server.HandleRequest(body) + respBody, hasResponse, err := server.HandleRequest(ctx, body) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeBadRequest, nil, "MCP server error: "+err.Error()) return @@ -114,8 +115,8 @@ func (h *MCPServerHandler) HandleMCP(c *gin.Context) { // MCPListDatasets wraps DatasetService.ListDatasets for the MCP tool handler, // filling in default values for parameters that the MCP tool does not expose. -func MCPListDatasets(ds *dataset.DatasetService, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) { - data, total, _, err := ds.ListDatasets( +func MCPListDatasets(ctx context.Context, ds *dataset.DatasetService, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) { + data, total, _, err := ds.ListDatasets(ctx, "", "", page, pageSize, orderby, desc, "", nil, "", userID, ) @@ -143,7 +144,7 @@ func MCPListChats(ctx context.Context, chatService *service.ChatService, userID // MCPRetrieval executes a retrieval request on behalf of the MCP tool handler. // It translates the mcp.RetrievalRequest into a service.SearchDatasetsRequest // and calls DatasetService.SearchDatasets. The result is serialized as JSON. -func MCPRetrieval(ds *dataset.DatasetService, userID string, req mcp.RetrievalRequest) (string, error) { +func MCPRetrieval(ctx context.Context, ds *dataset.DatasetService, userID string, req mcp.RetrievalRequest) (string, error) { // Resolve dataset IDs: if none provided, fetch ALL accessible datasets // across all pages (matching Python _fetch_all_datasets behaviour). datasetIDs := req.DatasetIDs @@ -151,7 +152,7 @@ func MCPRetrieval(ds *dataset.DatasetService, userID string, req mcp.RetrievalRe const maxPageSize = 100 page := 1 for { - data, _, _, err := ds.ListDatasets( + data, _, _, err := ds.ListDatasets(ctx, "", "", page, maxPageSize, "create_time", true, "", nil, "", userID, ) @@ -213,7 +214,7 @@ func MCPRetrieval(ds *dataset.DatasetService, userID string, req mcp.RetrievalRe searchReq.Keyword = &v } - resp, err := ds.SearchDatasets(searchReq, userID) + resp, err := ds.SearchDatasets(ctx, searchReq, userID) if err != nil { return "", err } diff --git a/internal/handler/search.go b/internal/handler/search.go index 14abd3f2a2..9616e48b9d 100644 --- a/internal/handler/search.go +++ b/internal/handler/search.go @@ -382,8 +382,8 @@ func (h *SearchHandler) Completion(c *gin.Context) { if searchSvc == nil { searchSvc = service.NewSearchService() } - - plan, code, err := searchSvc.PrepareCompletion(user.ID, c.Param("search_id"), &req) + ctx := c.Request.Context() + plan, code, err := searchSvc.PrepareCompletion(ctx, user.ID, c.Param("search_id"), &req) if err != nil { if code == common.CodeAuthenticationError { common.ResponseWithCodeData(c, code, false, err.Error()) diff --git a/internal/handler/tenant.go b/internal/handler/tenant.go index d4bb0835fd..5849757da6 100644 --- a/internal/handler/tenant.go +++ b/internal/handler/tenant.go @@ -257,8 +257,9 @@ func (h *TenantHandler) CreateChunkStore(c *gin.Context) { return } + ctx := c.Request.Context() // Check authorization - user must have access to this kb - if !h.datasetService.Accessible(req.KBID, user.ID) { + if !h.datasetService.Accessible(ctx, req.KBID, user.ID) { common.ResponseWithCodeData(c, common.CodeAuthenticationError, nil, "No authorization.") return } @@ -267,7 +268,7 @@ func (h *TenantHandler) CreateChunkStore(c *gin.Context) { KBID: req.KBID, VectorSize: req.VectorSize, } - result, code, err := h.tenantService.CreateChunkStore(serviceReq) + result, code, err := h.tenantService.CreateChunkStore(ctx, serviceReq) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -298,6 +299,7 @@ func (h *TenantHandler) DeleteChunkStore(c *gin.Context) { return } + ctx := c.Request.Context() var req DeleteChunkTableRequest if err := c.ShouldBindJSON(&req); err != nil { common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error()) @@ -305,12 +307,12 @@ func (h *TenantHandler) DeleteChunkStore(c *gin.Context) { } // Check authorization - if !h.datasetService.Accessible(req.KBID, user.ID) { + if !h.datasetService.Accessible(ctx, req.KBID, user.ID) { common.ResponseWithCodeData(c, common.CodeAuthenticationError, nil, "No authorization.") return } - code, err := h.tenantService.DeleteChunkStore(req.KBID) + code, err := h.tenantService.DeleteChunkStore(ctx, req.KBID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return diff --git a/internal/ingestion/component/chunker/group.go b/internal/ingestion/component/chunker/group.go index 21fb0082e0..992d4d9825 100644 --- a/internal/ingestion/component/chunker/group.go +++ b/internal/ingestion/component/chunker/group.go @@ -41,6 +41,7 @@ import ( "strings" "go.uber.org/zap" + "gorm.io/gorm" "ragflow/internal/agent/runtime" "ragflow/internal/common" @@ -468,7 +469,7 @@ func (c *GroupTitleChunkerComponent) Outputs() map[string]string { return ChunkerOutputs } -func (c *GroupTitleChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *GroupTitleChunkerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { if inputs == nil { inputs = map[string]any{} } diff --git a/internal/ingestion/component/chunker/group_test.go b/internal/ingestion/component/chunker/group_test.go index 3fe0c3b835..bf97caabc5 100644 --- a/internal/ingestion/component/chunker/group_test.go +++ b/internal/ingestion/component/chunker/group_test.go @@ -17,7 +17,6 @@ package chunker import ( - "context" "strings" "testing" @@ -114,7 +113,7 @@ func TestGroupTitleChunker_InvokeEmptyInput(t *testing.T) { if err != nil { t.Fatalf("NewGroupTitleChunker: %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(t.Context(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -138,7 +137,7 @@ func TestGroupTitleChunker_Headings_ASCII(t *testing.T) { t.Fatalf("NewGroupTitleChunker: %v", err) } input := "# Heading One\nBody line under H1.\nAnother body line.\n# Heading Two\nBody under second H1." - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "name": "doc.md", "text": input, }) @@ -161,7 +160,7 @@ func TestGroupTitleChunker_RootChunkAsHeading_StillSingleGroup(t *testing.T) { if err != nil { t.Fatalf("NewGroupTitleChunker: %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "name": "doc.md", "text": "Body without any heading here.\nMore body.", }) @@ -208,7 +207,7 @@ func TestGroupChunker_StructuredMetadata(t *testing.T) { {"text": "body line", "doc_type_kwd": "text"}, {"text": "an image caption", "doc_type_kwd": "image", "img_id": "img-9"}, } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "name": "doc", "output_format": "chunks", "chunks": items, @@ -249,7 +248,7 @@ func TestGroupChunker_MergesPDFPositionsAndRemovesTags(t *testing.T) { {"text": "body one @@1\t10.0\t20.0\t30.0\t40.0## tail", "doc_type_kwd": "text", "positions": [][]float64{{1, 10, 20, 30, 40}}}, {"text": "body two @@2\t15.0\t25.0\t35.0\t45.0## tail", "doc_type_kwd": "text", "positions": [][]float64{{2, 15, 25, 35, 45}}}, } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "name": "doc", "output_format": "chunks", "chunks": items, @@ -301,7 +300,7 @@ func TestGroupTitleChunker_InvokeDeterministic(t *testing.T) { var firstLen int var firstTexts []string for run := 0; run < 10; run++ { - out, err := c.Invoke(context.Background(), inputs) + out, err := c.Invoke(t.Context(), nil, inputs) if err != nil { t.Fatalf("Invoke run %d: %v", run, err) } diff --git a/internal/ingestion/component/chunker/hierarchy.go b/internal/ingestion/component/chunker/hierarchy.go index a659b8b34c..107e3b1233 100644 --- a/internal/ingestion/component/chunker/hierarchy.go +++ b/internal/ingestion/component/chunker/hierarchy.go @@ -43,6 +43,7 @@ import ( "fmt" "go.uber.org/zap" + "gorm.io/gorm" "ragflow/internal/agent/runtime" "ragflow/internal/common" @@ -286,7 +287,7 @@ func (c *HierarchyTitleChunkerComponent) Outputs() map[string]string { return ChunkerOutputs } -func (c *HierarchyTitleChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *HierarchyTitleChunkerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { if inputs == nil { inputs = map[string]any{} } diff --git a/internal/ingestion/component/chunker/hierarchy_test.go b/internal/ingestion/component/chunker/hierarchy_test.go index b1537e7a03..5ed21ccc7a 100644 --- a/internal/ingestion/component/chunker/hierarchy_test.go +++ b/internal/ingestion/component/chunker/hierarchy_test.go @@ -17,7 +17,6 @@ package chunker import ( - "context" "strings" "testing" @@ -81,7 +80,7 @@ func TestHierarchyTitleChunker_InvokeEmptyInput(t *testing.T) { if err != nil { t.Fatalf("NewHierarchyTitleChunker: %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(t.Context(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -104,7 +103,7 @@ func TestHierarchyTitleChunker_NestedHeadings(t *testing.T) { t.Fatalf("NewHierarchyTitleChunker: %v", err) } input := "# H1\nbody1a\nbody1b\n## H2a\nbody2a1\nbody2a2\n## H2b\nbody2b1\n# H1-2\nbody-last" - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "name": "doc.md", "text": input, }) @@ -133,7 +132,7 @@ func TestHierarchyTitleChunker_LeafOnlyDefault(t *testing.T) { if err != nil { t.Fatalf("NewHierarchyTitleChunker: %v", err) } - _, err = c.Invoke(context.Background(), map[string]any{ + _, err = c.Invoke(t.Context(), nil, map[string]any{ "name": "doc.md", "text": "# A\nbody a\n## A1\nbody a1\n# B\nbody b", }) @@ -158,7 +157,7 @@ func TestHierarchyChunker_StructuredMetadata(t *testing.T) { {"text": "body one", "doc_type_kwd": "text"}, {"text": "imgA caption", "doc_type_kwd": "image", "img_id": "a"}, } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "name": "doc", "output_format": "chunks", "chunks": items, @@ -196,7 +195,7 @@ func TestHierarchyTitleChunker_InvokeDeterministic(t *testing.T) { var firstLen int var firstTexts []string for run := 0; run < 10; run++ { - out, err := c.Invoke(context.Background(), inputs) + out, err := c.Invoke(t.Context(), nil, inputs) if err != nil { t.Fatalf("Invoke run %d: %v", run, err) } @@ -246,7 +245,7 @@ func TestHierarchyTitleChunker_ColonTitlePromotion(t *testing.T) { {"text": "Another section continuing the discussion. The key provisions are outlined below:", "doc_type_kwd": "text"}, {"text": "Body two.", "doc_type_kwd": "text"}, } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "name": "test.txt", "chunks": items, }) @@ -278,7 +277,7 @@ func TestHierarchyTitleChunker_ColonTitleShortLine_Negative(t *testing.T) { {"text": "Note:", "doc_type_kwd": "text"}, {"text": "Body text.", "doc_type_kwd": "text"}, } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "name": "test.txt", "chunks": items, }) @@ -313,7 +312,7 @@ func TestHierarchyTitleChunker_ShortNumericLineFilter(t *testing.T) { {"text": "2", "doc_type_kwd": "text"}, {"text": "Methodology paragraph.", "doc_type_kwd": "text"}, } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "name": "test.txt", "chunks": items, }) @@ -389,7 +388,7 @@ func TestHierarchyTitleChunker_ColonTitlePromotion_CJK_EdgeCase(t *testing.T) { {"text": "def。1234567890123456789012345678901：", "doc_type_kwd": "text"}, {"text": "Body two.", "doc_type_kwd": "text"}, } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "name": "test.txt", "chunks": items, }) @@ -427,7 +426,7 @@ func TestHierarchyTitleChunker_ShortSingleCJKLineFilter(t *testing.T) { {"text": "例", "doc_type_kwd": "text"}, {"text": "Body two.", "doc_type_kwd": "text"}, } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "name": "test.txt", "chunks": items, }) @@ -466,7 +465,7 @@ func TestHierarchyTitleChunker_CKTypeHeadingFallback(t *testing.T) { {"text": "Conclusion", "doc_type_kwd": "text", "ck_type": "heading"}, {"text": "Final words.", "doc_type_kwd": "text"}, } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "name": "test.docx", "chunks": items, }) @@ -507,7 +506,7 @@ func TestHierarchyTitleChunker_ConsecutiveNonTextOrder(t *testing.T) { {"text": "# H2", "doc_type_kwd": "text"}, {"text": "body two", "doc_type_kwd": "text"}, } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "name": "doc", "chunks": items, }) diff --git a/internal/ingestion/component/chunker/image_upload_test.go b/internal/ingestion/component/chunker/image_upload_test.go index 0384439103..2a11e05e65 100644 --- a/internal/ingestion/component/chunker/image_upload_test.go +++ b/internal/ingestion/component/chunker/image_upload_test.go @@ -228,7 +228,7 @@ func TestImageUploadDecorator_EndToEnd(t *testing.T) { {"content_with_weight": "a cropped figure", "image": "data:image/png;base64," + pngBase64}, }, } - out, err := decorated.Invoke(context.Background(), inputs) + out, err := decorated.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("decorated Invoke: %v", err) } diff --git a/internal/ingestion/component/chunker/one.go b/internal/ingestion/component/chunker/one.go index 8abc084f04..daff2a4c1d 100644 --- a/internal/ingestion/component/chunker/one.go +++ b/internal/ingestion/component/chunker/one.go @@ -34,6 +34,8 @@ import ( "ragflow/internal/agent/runtime" "ragflow/internal/ingestion/component/schema" + + "gorm.io/gorm" ) const ComponentNameOneChunker = "OneChunker" @@ -66,7 +68,7 @@ func (c *OneChunkerComponent) Inputs() map[string]string { return ChunkerInputs func (c *OneChunkerComponent) Outputs() map[string]string { return ChunkerOutputs } -func (c *OneChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *OneChunkerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { return c.invoke(ctx, inputs) } diff --git a/internal/ingestion/component/chunker/one_test.go b/internal/ingestion/component/chunker/one_test.go index 0cd371da8c..031774dfaf 100644 --- a/internal/ingestion/component/chunker/one_test.go +++ b/internal/ingestion/component/chunker/one_test.go @@ -28,7 +28,7 @@ func oneChunksOf(t *testing.T, inputs map[string]any) []map[string]any { if err != nil { t.Fatalf("NewOneChunker: %v", err) } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("OneChunker.Invoke: %v", err) } diff --git a/internal/ingestion/component/chunker/pdfcrop_cgo.go b/internal/ingestion/component/chunker/pdfcrop_cgo.go index 9a72b77a6d..04647bb042 100644 --- a/internal/ingestion/component/chunker/pdfcrop_cgo.go +++ b/internal/ingestion/component/chunker/pdfcrop_cgo.go @@ -21,13 +21,15 @@ import ( deepdoctype "ragflow/internal/deepdoc/parser/type" "ragflow/internal/ingestion/component" "ragflow/internal/ingestion/component/schema" + + "gorm.io/gorm" ) // newPDFEngineFromUpstream re-acquires the source PDF from storage using the // same resolution the Parser uses, then opens a native engine. It returns // (nil, nil) when no storage reference is present or the bytes are not a PDF, // so callers can treat a nil engine as "no cropping". -func newPDFEngineFromUpstream(ctx context.Context, up schema.ChunkerFromUpstream) (deepdoctype.PDFEngine, error) { +func newPDFEngineFromUpstream(ctx context.Context, db *gorm.DB, up schema.ChunkerFromUpstream) (deepdoctype.PDFEngine, error) { var data []byte var err error switch { @@ -35,7 +37,7 @@ func newPDFEngineFromUpstream(ctx context.Context, up schema.ChunkerFromUpstream data, err = component.FetchBinary(ctx, up.Bucket, up.Path) case up.DocID != "": var ref *component.DocumentStorageRef - ref, err = component.ResolveDocumentStorage(ctx, up.DocID) + ref, err = component.ResolveDocumentStorage(ctx, db, up.DocID) if err == nil && ref != nil { data, err = component.FetchBinary(ctx, ref.Bucket, ref.Path) } diff --git a/internal/ingestion/component/chunker/pdfcrop_nocgo.go b/internal/ingestion/component/chunker/pdfcrop_nocgo.go index 51778ff0f1..179e199c33 100644 --- a/internal/ingestion/component/chunker/pdfcrop_nocgo.go +++ b/internal/ingestion/component/chunker/pdfcrop_nocgo.go @@ -14,7 +14,7 @@ import ( "ragflow/internal/ingestion/component/schema" ) -func newPDFEngineFromUpstream(_ context.Context, _ schema.ChunkerFromUpstream) (deepdoctype.PDFEngine, error) { +func newPDFEngineFromUpstream(_ context.Context, _ *gorm.DB, _ schema.ChunkerFromUpstream) (deepdoctype.PDFEngine, error) { return nil, nil } diff --git a/internal/ingestion/component/chunker/presentation.go b/internal/ingestion/component/chunker/presentation.go index de1a3ffda2..63b51113fb 100644 --- a/internal/ingestion/component/chunker/presentation.go +++ b/internal/ingestion/component/chunker/presentation.go @@ -32,6 +32,8 @@ import ( "ragflow/internal/agent/runtime" "ragflow/internal/ingestion/component/schema" + + "gorm.io/gorm" ) const ComponentNamePresentationChunker = "PresentationChunker" @@ -66,7 +68,7 @@ func (c *PresentationChunkerComponent) Inputs() map[string]string { return Chunk func (c *PresentationChunkerComponent) Outputs() map[string]string { return ChunkerOutputs } -func (c *PresentationChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *PresentationChunkerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { return c.invoke(ctx, inputs) } diff --git a/internal/ingestion/component/chunker/presentation_test.go b/internal/ingestion/component/chunker/presentation_test.go index 9a695b4907..17a3d518ac 100644 --- a/internal/ingestion/component/chunker/presentation_test.go +++ b/internal/ingestion/component/chunker/presentation_test.go @@ -28,7 +28,7 @@ func slideChunksOf(t *testing.T, inputs map[string]any) []map[string]any { if err != nil { t.Fatalf("NewPresentationChunker: %v", err) } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("PresentationChunker.Invoke: %v", err) } diff --git a/internal/ingestion/component/chunker/python_parity_test.go b/internal/ingestion/component/chunker/python_parity_test.go index a6296a3a74..6b15ee6938 100644 --- a/internal/ingestion/component/chunker/python_parity_test.go +++ b/internal/ingestion/component/chunker/python_parity_test.go @@ -46,7 +46,7 @@ func invokeAsTokenChunker(t *testing.T, items []map[string]any) map[string]any { if err != nil { t.Fatalf("NewTokenChunker: %v", err) } - out, err := comp.Invoke(context.Background(), map[string]any{ + out, err := comp.Invoke(context.Background(), nil, map[string]any{ "name": "test", "output_format": "json", "json": items, diff --git a/internal/ingestion/component/chunker/qa.go b/internal/ingestion/component/chunker/qa.go index 4dce3f7573..33c4359d99 100644 --- a/internal/ingestion/component/chunker/qa.go +++ b/internal/ingestion/component/chunker/qa.go @@ -37,6 +37,7 @@ import ( "github.com/gomarkdown/markdown" "github.com/gomarkdown/markdown/parser" + "gorm.io/gorm" "ragflow/internal/agent/runtime" "ragflow/internal/ingestion/component/schema" @@ -81,7 +82,7 @@ func (c *QAChunkerComponent) Inputs() map[string]string { return ChunkerInputs } func (c *QAChunkerComponent) Outputs() map[string]string { return ChunkerOutputs } -func (c *QAChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *QAChunkerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { return c.invoke(ctx, inputs) } diff --git a/internal/ingestion/component/chunker/qa_batch2_test.go b/internal/ingestion/component/chunker/qa_batch2_test.go index 6ad8329d5f..84b334ac94 100644 --- a/internal/ingestion/component/chunker/qa_batch2_test.go +++ b/internal/ingestion/component/chunker/qa_batch2_test.go @@ -30,7 +30,7 @@ func qaInvoke(t *testing.T, inputs map[string]any) []map[string]any { if err != nil { t.Fatalf("NewQAChunker: %v", err) } - out, err := c.Invoke(context.Background(), inputs) + out, err := c.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke: %v", err) } diff --git a/internal/ingestion/component/chunker/qa_test.go b/internal/ingestion/component/chunker/qa_test.go index 4502cd00ed..ad85e13486 100644 --- a/internal/ingestion/component/chunker/qa_test.go +++ b/internal/ingestion/component/chunker/qa_test.go @@ -48,7 +48,7 @@ func TestQAChunker_DelimiterTab(t *testing.T) { "output_format": "text", "text": "What is Go?\tGo is a programming language.", } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke failed: %v", err) } @@ -73,7 +73,7 @@ func TestQAChunker_DelimiterComma(t *testing.T) { "output_format": "text", "text": "What is Rust?,Rust is a systems language.", } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke failed: %v", err) } @@ -98,7 +98,7 @@ func TestQAChunker_Markdown(t *testing.T) { "output_format": "markdown", "markdown": "# What is Go?\nGo is a programming language.\n\n# What is Rust?\nRust is a systems language.", } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke failed: %v", err) } @@ -118,7 +118,7 @@ func TestQAChunker_HTMLTable(t *testing.T) { "output_format": "html", "html": "
Q1A1
Q2A2
", } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke failed: %v", err) } @@ -138,7 +138,7 @@ func TestQAChunker_RmQAPrefix(t *testing.T) { "output_format": "text", "text": "Question: What is Go?\tAnswer: Go is a language.", } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke failed: %v", err) } @@ -159,7 +159,7 @@ func TestQAChunker_Empty(t *testing.T) { "output_format": "text", "text": "", } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke failed: %v", err) } @@ -179,7 +179,7 @@ func TestQAChunker_CaseInsensitivePrefix(t *testing.T) { "output_format": "text", "text": "QUESTION: Hello\tANSWER: World", } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke failed: %v", err) } @@ -203,7 +203,7 @@ func TestQAChunker_PrefixSpaceSeparatorStrips(t *testing.T) { "output_format": "text", "text": "A language model is useful\tQ How does it work", } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke failed: %v", err) } @@ -229,7 +229,7 @@ func TestQAChunker_HeadingNoTrailingSpace(t *testing.T) { "output_format": "markdown", "markdown": "#Hello\nWorld\n", } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke failed: %v", err) } @@ -249,7 +249,7 @@ func TestQAChunker_ChineseLang(t *testing.T) { "output_format": "text", "text": "什么是Go？\tGo是一种编程语言。", } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke failed: %v", err) } @@ -273,7 +273,7 @@ func TestQAChunker_MarkdownRendersHTML(t *testing.T) { "output_format": "markdown", "markdown": "# Title\nThis is **bold** text.\n", } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke failed: %v", err) } diff --git a/internal/ingestion/component/chunker/register.go b/internal/ingestion/component/chunker/register.go index 04498bdff6..6b716fa76c 100644 --- a/internal/ingestion/component/chunker/register.go +++ b/internal/ingestion/component/chunker/register.go @@ -30,6 +30,8 @@ import ( "ragflow/internal/agent/runtime" "ragflow/internal/common" + + "gorm.io/gorm" ) // MustRegisterChunker registers a single chunker component under @@ -60,8 +62,8 @@ type imageUploadDecorator struct { inner runtime.Component } -func (d *imageUploadDecorator) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - out, err := d.inner.Invoke(ctx, inputs) +func (d *imageUploadDecorator) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { + out, err := d.inner.Invoke(ctx, db, inputs) if err != nil { return nil, err } diff --git a/internal/ingestion/component/chunker/table.go b/internal/ingestion/component/chunker/table.go index 3e552321bc..74fb1bee18 100644 --- a/internal/ingestion/component/chunker/table.go +++ b/internal/ingestion/component/chunker/table.go @@ -33,6 +33,8 @@ import ( "ragflow/internal/agent/runtime" "ragflow/internal/ingestion/component/schema" + + "gorm.io/gorm" ) const ComponentNameTableChunker = "TableChunker" @@ -65,7 +67,7 @@ func (c *TableChunkerComponent) Inputs() map[string]string { return ChunkerInput func (c *TableChunkerComponent) Outputs() map[string]string { return ChunkerOutputs } -func (c *TableChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *TableChunkerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { return c.invoke(ctx, inputs) } diff --git a/internal/ingestion/component/chunker/table_test.go b/internal/ingestion/component/chunker/table_test.go index 91df228933..b7b0275efc 100644 --- a/internal/ingestion/component/chunker/table_test.go +++ b/internal/ingestion/component/chunker/table_test.go @@ -28,7 +28,7 @@ func tableChunksOf(t *testing.T, inputs map[string]any) []map[string]any { if err != nil { t.Fatalf("NewTableChunker: %v", err) } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("TableChunker.Invoke: %v", err) } diff --git a/internal/ingestion/component/chunker/tag.go b/internal/ingestion/component/chunker/tag.go index 864d6ea5bc..e4a1e8496f 100644 --- a/internal/ingestion/component/chunker/tag.go +++ b/internal/ingestion/component/chunker/tag.go @@ -46,6 +46,8 @@ import ( "ragflow/internal/agent/runtime" "ragflow/internal/ingestion/component/schema" "ragflow/internal/tokenizer" + + "gorm.io/gorm" ) const ComponentNameTagChunker = "TagChunker" @@ -78,7 +80,7 @@ func (c *TagChunkerComponent) Inputs() map[string]string { return ChunkerInputs func (c *TagChunkerComponent) Outputs() map[string]string { return ChunkerOutputs } -func (c *TagChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *TagChunkerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { return c.invoke(ctx, inputs) } diff --git a/internal/ingestion/component/chunker/tag_test.go b/internal/ingestion/component/chunker/tag_test.go index 246637b3e7..7fb503f4ac 100644 --- a/internal/ingestion/component/chunker/tag_test.go +++ b/internal/ingestion/component/chunker/tag_test.go @@ -28,7 +28,7 @@ func tagChunksOf(t *testing.T, inputs map[string]any) []map[string]any { if err != nil { t.Fatalf("NewTagChunker: %v", err) } - out, err := comp.Invoke(context.Background(), inputs) + out, err := comp.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("TagChunker.Invoke: %v", err) } diff --git a/internal/ingestion/component/chunker/title.go b/internal/ingestion/component/chunker/title.go index 791a962dad..3403dd3ada 100644 --- a/internal/ingestion/component/chunker/title.go +++ b/internal/ingestion/component/chunker/title.go @@ -63,6 +63,8 @@ import ( "ragflow/internal/agent/runtime" "ragflow/internal/ingestion/component/globals" "ragflow/internal/ingestion/component/schema" + + "gorm.io/gorm" ) const ComponentNameTitleChunker = "TitleChunker" @@ -629,7 +631,7 @@ func (c *TitleChunkerComponent) Inputs() map[string]string { return ChunkerInput func (c *TitleChunkerComponent) Outputs() map[string]string { return ChunkerOutputs } // Invoke delegates to the chosen strategy (group or hierarchy). -func (c *TitleChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *TitleChunkerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { if inputs == nil { inputs = map[string]any{} } diff --git a/internal/ingestion/component/chunker/title_test.go b/internal/ingestion/component/chunker/title_test.go index 9a8523d531..4cd660af3f 100644 --- a/internal/ingestion/component/chunker/title_test.go +++ b/internal/ingestion/component/chunker/title_test.go @@ -92,7 +92,7 @@ func TestTitleChunker_InvokeEmptyInput(t *testing.T) { if err != nil { t.Fatalf("NewTitleChunker: %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -127,7 +127,7 @@ func TestTitleChunker_Headings_ASCII(t *testing.T) { } input := "# Top\nFirst body line under Top.\nSecond body.\n## Sub\nBody under Sub heading.\n# TopTwo\nBody under TopTwo." - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.md", "text": input, }) @@ -159,7 +159,7 @@ func TestTitleChunker_NoHeadings_FallsBack(t *testing.T) { if err != nil { t.Fatalf("NewTitleChunker: %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.txt", "text": "alpha line one\nalpha line two", }) @@ -186,7 +186,7 @@ func TestTitleChunker_DispatcherHierarchy(t *testing.T) { if _, ok := c.(*TitleChunkerComponent); !ok { t.Fatalf("NewTitleChunker returned %T", c) } - _, err = c.Invoke(context.Background(), map[string]any{ + _, err = c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.md", "text": "# Top\nFirst body.\n# TopTwo\nBody under TopTwo.", }) @@ -215,7 +215,7 @@ func TestTitleChunker_InvokeDeterministic(t *testing.T) { var firstLen int var firstTexts []string for run := 0; run < 10; run++ { - out, err := c.Invoke(context.Background(), inputs) + out, err := c.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke run %d: %v", run, err) } diff --git a/internal/ingestion/component/chunker/token.go b/internal/ingestion/component/chunker/token.go index e08af56daf..0d6e0ebc81 100644 --- a/internal/ingestion/component/chunker/token.go +++ b/internal/ingestion/component/chunker/token.go @@ -65,6 +65,8 @@ import ( deepdoctype "ragflow/internal/deepdoc/parser/type" "ragflow/internal/ingestion/component/globals" "ragflow/internal/ingestion/component/schema" + + "gorm.io/gorm" ) const ComponentNameTokenChunker = "TokenChunker" @@ -146,11 +148,11 @@ func (c *TokenChunkerComponent) Outputs() map[string]string { return ChunkerOutp // // Timeout: honours ctx cancellation only — there is no inner @timeout // decorator equivalent (plan §8 R1). -func (c *TokenChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - return c.invoke(ctx, inputs) +func (c *TokenChunkerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { + return c.invoke(ctx, db, inputs) } -func (c *TokenChunkerComponent) invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *TokenChunkerComponent) invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { if inputs == nil { return emptyOutputs(), nil } @@ -211,7 +213,7 @@ func (c *TokenChunkerComponent) invoke(ctx context.Context, inputs map[string]an // refs) so image/table sections are cropped on demand rather // than carried through the wire. Best-effort: a nil engine // simply skips cropping. - engine, engErr := newPDFEngineFromUpstream(ctx, upstream) + engine, engErr := newPDFEngineFromUpstream(ctx, db, upstream) if engErr != nil { slog.Warn("TokenChunker: could not open PDF for on-demand cropping", "err", engErr) } diff --git a/internal/ingestion/component/chunker/token_test.go b/internal/ingestion/component/chunker/token_test.go index f37d3a4484..342ec8c0ad 100644 --- a/internal/ingestion/component/chunker/token_test.go +++ b/internal/ingestion/component/chunker/token_test.go @@ -52,7 +52,7 @@ func TestTokenChunker_InvokeEmptyInput(t *testing.T) { if err != nil { t.Fatalf("NewTokenChunker: %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(context.Background(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -76,7 +76,7 @@ func TestTokenChunker_InvokeDelimMode_BasicChunking(t *testing.T) { if err != nil { t.Fatalf("NewTokenChunker: %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.txt", "output_format": "text", "text": "alpha\n\nbeta\n\ngamma", @@ -110,7 +110,7 @@ func TestTokenChunker_DelimNeverStandaloneChunk(t *testing.T) { if err != nil { t.Fatalf("NewTokenChunker: %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.txt", "output_format": "text", "text": "alpha section\n666\nbeta section", @@ -150,7 +150,7 @@ func TestTokenChunker_InvokeTokenSize_FallbackToMerge(t *testing.T) { } // Input without any \n\n so the delimiter miss branch triggers // the token_size merge. - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.txt", "output_format": "text", "text": "First sentence. Second sentence. Third sentence. Fourth.", @@ -179,7 +179,7 @@ func TestTokenChunker_InvokeChildrenDelim(t *testing.T) { if err != nil { t.Fatalf("NewTokenChunker: %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.txt", "output_format": "text", "text": "alpha line\nbeta line", @@ -209,7 +209,7 @@ func TestTokenChunker_InvokeJSONPayload(t *testing.T) { {"text": "Alpha text\nBeta text", "doc_type_kwd": "text"}, {"text": "Gamma text\nDelta text", "doc_type_kwd": "text"}, } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.md", "output_format": "json", "json": items, @@ -251,7 +251,7 @@ func TestTokenChunker_InvokeDeterministic(t *testing.T) { } var firstfp fingerprint for run := 0; run < 10; run++ { - out, err := c.Invoke(context.Background(), inputs) + out, err := c.Invoke(context.Background(), nil, inputs) if err != nil { t.Fatalf("Invoke run %d: %v", run, err) } @@ -336,7 +336,7 @@ func TestTokenChunker_PrefersUpstreamChunks(t *testing.T) { if err != nil { t.Fatalf("NewTokenChunker: %v", err) } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.md", "output_format": "chunks", "chunks": []map[string]any{{"text": "CHAPTER-AWARE", "doc_type_kwd": "text"}}, diff --git a/internal/ingestion/component/document_storage.go b/internal/ingestion/component/document_storage.go index 1bbd31cd31..2b9789aa14 100644 --- a/internal/ingestion/component/document_storage.go +++ b/internal/ingestion/component/document_storage.go @@ -23,6 +23,8 @@ import ( "ragflow/internal/dao" "ragflow/internal/entity" "ragflow/internal/storage" + + "gorm.io/gorm" ) // DocumentStorageRef is the resolved backing storage location for a document. @@ -81,23 +83,23 @@ func resolveStorage() storage.Storage { // location. It is exported so downstream components can re-acquire the // source PDF without threading the raw bytes across the component // boundary. -func ResolveDocumentStorage(ctx context.Context, docID string) (*DocumentStorageRef, error) { +func ResolveDocumentStorage(ctx context.Context, db *gorm.DB, docID string) (*DocumentStorageRef, error) { if ResolveDocumentStorageOverride != nil { return ResolveDocumentStorageOverride(docID) } - doc, err := dao.NewDocumentDAO().GetByID(ctx, dao.DB, docID) + doc, err := dao.NewDocumentDAO().GetByID(ctx, db, docID) if err != nil { return nil, err } ref := &DocumentStorageRef{Name: documentNameOrID(doc)} - mappings, err := dao.NewFile2DocumentDAO().GetByDocumentID(ctx, dao.DB, doc.ID) + mappings, err := dao.NewFile2DocumentDAO().GetByDocumentID(ctx, db, doc.ID) if err != nil { return nil, err } if len(mappings) > 0 && mappings[0].FileID != nil && *mappings[0].FileID != "" { - file, err := dao.NewFileDAO().GetByID(ctx, dao.DB, *mappings[0].FileID) + file, err := dao.NewFileDAO().GetByID(ctx, db, *mappings[0].FileID) if err != nil { return nil, err } diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index ac4ae73e5d..588c272a41 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -72,6 +72,7 @@ import ( eschema "github.com/cloudwego/eino/schema" "go.uber.org/zap" + "gorm.io/gorm" "ragflow/internal/agent/runtime" "ragflow/internal/common" @@ -523,7 +524,7 @@ func extractorChunkList(v any) ([]map[string]any, bool) { // short-circuits with an error. // _created_time, _elapsed_time — stamped by the canvas framework // (realComponentBody), not here. -func (c *ExtractorComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *ExtractorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { if err := c.Param.Validate(); err != nil { return nil, fmt.Errorf("extractor: %w", err) } diff --git a/internal/ingestion/component/extractor_tag_test.go b/internal/ingestion/component/extractor_tag_test.go index 7266c7e5ce..42c56e632a 100644 --- a/internal/ingestion/component/extractor_tag_test.go +++ b/internal/ingestion/component/extractor_tag_test.go @@ -45,7 +45,7 @@ func pushExtractorTagTargetResolverStub(t *testing.T) { func TestExtractorTags_NoTagFileID(t *testing.T) { pushExtractorTagChatStub(t, nil) comp, _ := NewExtractorComponent(map[string]any{"auto_tags": 3}) - out, err := comp.Invoke(context.Background(), map[string]any{ + out, err := comp.Invoke(t.Context(), nil, map[string]any{ "chunks": []map[string]any{ {"content_with_weight": "test"}, }, @@ -65,7 +65,7 @@ func TestExtractorTags_NoTagFileID(t *testing.T) { func TestExtractorTags_NoLLMID(t *testing.T) { pushExtractorTagChatStub(t, nil) comp, _ := NewExtractorComponent(map[string]any{"auto_tags": 3}) - out, err := comp.Invoke(context.Background(), map[string]any{ + out, err := comp.Invoke(t.Context(), nil, map[string]any{ "chunks": []map[string]any{ {"content_with_weight": "some unrelated text"}, }, @@ -88,7 +88,7 @@ func TestExtractorTags_WithKeywords(t *testing.T) { "auto_tags": 3, "auto_keywords": 3, }) - out, err := comp.Invoke(context.Background(), map[string]any{ + out, err := comp.Invoke(t.Context(), nil, map[string]any{ "chunks": []map[string]any{ {"content_with_weight": "some unrelated textxyz"}, }, diff --git a/internal/ingestion/component/extractor_test.go b/internal/ingestion/component/extractor_test.go index 79ccb1df0e..9394c9ae6d 100644 --- a/internal/ingestion/component/extractor_test.go +++ b/internal/ingestion/component/extractor_test.go @@ -127,7 +127,7 @@ func TestExtractorComponent_Invoke_HappyPath(t *testing.T) { LLMID: "gpt-4o-mini", Prompt: "Summarize:", }} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "chunks": []map[string]any{ {"text": "first text"}, {"text": "second text"}, @@ -167,7 +167,7 @@ func TestExtractorComponent_Invoke_LLMError(t *testing.T) { FieldName: "summary", LLMID: "gpt-4o-mini", }} - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(t.Context(), nil, map[string]any{ "chunks": []map[string]any{{"text": "x"}}, }) if err == nil { @@ -247,7 +247,7 @@ func TestExtractorComponent_Invoke_ParsesJSON(t *testing.T) { FieldName: "extraction", Prompt: "extract:", }} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "chunks": []map[string]any{{"text": "doc"}}}, ) if err != nil { @@ -279,7 +279,7 @@ func TestExtractorComponent_Invoke_ParsesJSONInFence(t *testing.T) { c := &ExtractorComponent{Param: schema.ExtractorParam{ FieldName: "out", }} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "chunks": []map[string]any{{"text": "x"}}}, ) if err != nil { @@ -307,7 +307,7 @@ func TestExtractorComponent_Invoke_HandlesMalformedJSON(t *testing.T) { c := &ExtractorComponent{Param: schema.ExtractorParam{ FieldName: "raw", }} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "chunks": []map[string]any{{"text": "x"}}}, ) if err != nil { @@ -327,7 +327,7 @@ func TestExtractorComponent_Invoke_TOCNotPorted(t *testing.T) { c := &ExtractorComponent{Param: schema.ExtractorParam{ FieldName: "toc", }} - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(t.Context(), nil, map[string]any{ "chunks": []map[string]any{{"text": "x"}}}, ) if err == nil { @@ -352,7 +352,7 @@ func TestExtractorComponent_Invoke_NoChunksFastPath(t *testing.T) { c := &ExtractorComponent{Param: schema.ExtractorParam{ FieldName: "answer", }} - out, err := c.Invoke(context.Background(), map[string]any{}) + out, err := c.Invoke(t.Context(), nil, map[string]any{}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -376,7 +376,7 @@ func TestExtractorComponent_Invoke_JSONListInput(t *testing.T) { c := &ExtractorComponent{Param: schema.ExtractorParam{ FieldName: "answer", }} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "json": []map[string]any{{"text": "json payload chunk"}}, }) if err != nil { @@ -404,7 +404,7 @@ func TestExtractorComponent_Invoke_PerCallLLMIDOverride(t *testing.T) { FieldName: "out", LLMID: "static-llm", }} - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(t.Context(), nil, map[string]any{ "llm_id": "override-llm", }) if err != nil { @@ -430,7 +430,7 @@ func TestExtractorComponent_Invoke_CompositeLLMID(t *testing.T) { FieldName: "out", LLMID: "gpt-4o-mini@openai", }} - if _, err := c.Invoke(context.Background(), map[string]any{}); err != nil { + if _, err := c.Invoke(t.Context(), nil, map[string]any{}); err != nil { t.Fatalf("Invoke: %v", err) } stub.mu.Lock() @@ -455,7 +455,7 @@ func TestExtractorComponent_Invoke_ChunkIndexInError(t *testing.T) { c := &ExtractorComponent{Param: schema.ExtractorParam{ FieldName: "out", }} - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(t.Context(), nil, map[string]any{ "chunks": []map[string]any{ {"text": "first"}, {"text": "second"}, @@ -499,7 +499,7 @@ func TestExtractorComponent_NewExtractorComponent_Happy(t *testing.T) { if err != nil { t.Fatalf("NewExtractorComponent: %v", err) } - if _, err := c.Invoke(context.Background(), map[string]any{ + if _, err = c.Invoke(t.Context(), nil, map[string]any{ "chunks": []map[string]any{{"text": "x"}}}, ); err != nil { t.Fatalf("Invoke: %v", err) @@ -737,11 +737,10 @@ func TestExtractorComponent_ConcurrentInvoke(t *testing.T) { var wg sync.WaitGroup errs := make(chan error, len(chunks)) for _, ck := range chunks { - ck := ck wg.Add(1) go func() { defer wg.Done() - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(t.Context(), nil, map[string]any{ "chunks": []map[string]any{ck}, }) if err != nil { diff --git a/internal/ingestion/component/file.go b/internal/ingestion/component/file.go index 8e80027189..bd037079cb 100644 --- a/internal/ingestion/component/file.go +++ b/internal/ingestion/component/file.go @@ -46,6 +46,8 @@ import ( "ragflow/internal/ingestion/component/globals" "ragflow/internal/ingestion/component/schema" "ragflow/internal/storage" + + "gorm.io/gorm" ) const ComponentNameFile = "File" @@ -128,7 +130,7 @@ func (c *FileComponent) Outputs() map[string]string { // the document name and emit metadata only. // 2. doc_id is empty — pull the first file descriptor out of // `file` and use its `name`/`id` directly. -func (c *FileComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *FileComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { // Parse the wire input through the schema type so the // validation errors match the package convention. in, err := parseFileInputs(ctx, inputs) diff --git a/internal/ingestion/component/file_test.go b/internal/ingestion/component/file_test.go index 70ce40d191..2de65cba67 100644 --- a/internal/ingestion/component/file_test.go +++ b/internal/ingestion/component/file_test.go @@ -17,7 +17,6 @@ package component import ( - "context" "strings" "testing" @@ -89,7 +88,7 @@ func TestFileComponent_Invoke_HappyPath(t *testing.T) { } c := &FileComponent{} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "file": []map[string]any{{"name": "file.txt"}}, "bucket": "bucketA", "path": "path/to/file.txt", @@ -135,7 +134,7 @@ func TestFileComponent_Invoke_ResolvesDocIDViaDocumentLocation(t *testing.T) { } c := &FileComponent{} - out, err := c.Invoke(context.Background(), map[string]any{"doc_id": "doc-loc"}) + out, err := c.Invoke(t.Context(), db, map[string]any{"doc_id": "doc-loc"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -192,7 +191,7 @@ func TestFileComponent_Invoke_ResolvesDocIDViaFileMapping(t *testing.T) { } c := &FileComponent{} - out, err := c.Invoke(context.Background(), map[string]any{"doc_id": "doc-file"}) + out, err := c.Invoke(t.Context(), db, map[string]any{"doc_id": "doc-file"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -223,7 +222,7 @@ func TestFileComponent_Invoke_DocIDWithoutStorageLocationStillSucceeds(t *testin } c := &FileComponent{} - out, err := c.Invoke(context.Background(), map[string]any{"doc_id": "doc-empty"}) + out, err := c.Invoke(t.Context(), db, map[string]any{"doc_id": "doc-empty"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -237,7 +236,7 @@ func TestFileComponent_Invoke_DocIDWithoutStorageLocationStillSucceeds(t *testin func TestFileComponent_Invoke_MissingDoc(t *testing.T) { withMemoryStorage(t) c := &FileComponent{} - _, err := c.Invoke(context.Background(), map[string]any{}) + _, err := c.Invoke(t.Context(), nil, map[string]any{}) if err == nil { t.Fatal("expected error for empty inputs, got nil") } @@ -255,7 +254,7 @@ func TestFileComponent_Invoke_IncludesCheckpointPath(t *testing.T) { t.Fatalf("seed: %v", err) } c := &FileComponent{} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "file": []map[string]any{{"name": "checkpoint.bin"}}, "bucket": "b", "path": wantPath, diff --git a/internal/ingestion/component/parser.go b/internal/ingestion/component/parser.go index 1b6f867302..8724424ab1 100644 --- a/internal/ingestion/component/parser.go +++ b/internal/ingestion/component/parser.go @@ -84,6 +84,7 @@ import ( "unicode/utf8" "go.uber.org/zap" + "gorm.io/gorm" "ragflow/internal/agent/runtime" "ragflow/internal/common" @@ -392,9 +393,9 @@ func (c *ParserComponent) Outputs() map[string]string { // that downstream Chunker / Tokenizer rely on for stable chunk // IDs (chunks that span pages must reference adjacent PageNumbers // in input order). -func (c *ParserComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *ParserComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { // 1. Decode the binary input. - binary, err := readParserBinary(ctx, inputs) + binary, err := readParserBinary(ctx, db, inputs) if err != nil { return nil, err } @@ -618,7 +619,7 @@ func buildPagesFromBytes(ctx context.Context, pages [][]byte, docType string) ([ // A non-UTF-8 string is rejected with a clear error so a caller // that mistakenly hands a base64 string sees the failure // immediately (mirrors pipeline_chunker's "no try-base64" rule). -func readParserBinary(ctx context.Context, inputs map[string]any) ([]byte, error) { +func readParserBinary(ctx context.Context, db *gorm.DB, inputs map[string]any) ([]byte, error) { if inputs == nil { return nil, nil } @@ -628,8 +629,8 @@ func readParserBinary(ctx context.Context, inputs map[string]any) ([]byte, error if s, ok := inputs["binary"].(string); ok { if !utf8.ValidString(s) { return nil, errors.New( - "Parser: binary string is not valid UTF-8. " + - "Text-page mode only accepts UTF-8 text input.") + "parser: binary string is not valid UTF-8. " + + "Text-page mode only accepts UTF-8 text input") } return []byte(s), nil } @@ -639,7 +640,7 @@ func readParserBinary(ctx context.Context, inputs map[string]any) ([]byte, error return FetchBinary(ctx, bucket, path) } if docID, ok := getString(inputs, "doc_id"); ok && docID != "" { - ref, err := ResolveDocumentStorage(ctx, docID) + ref, err := ResolveDocumentStorage(ctx, db, docID) if err != nil { return nil, fmt.Errorf("Parser: resolve doc_id %q: %w", docID, err) } diff --git a/internal/ingestion/component/parser_dispatch_pdf_vision_cgo_test.go b/internal/ingestion/component/parser_dispatch_pdf_vision_cgo_test.go index b697bec671..62d6bae540 100644 --- a/internal/ingestion/component/parser_dispatch_pdf_vision_cgo_test.go +++ b/internal/ingestion/component/parser_dispatch_pdf_vision_cgo_test.go @@ -65,7 +65,7 @@ func TestDispatch_PDFVisionJSON_RealPDFFixture(t *testing.T) { setups["pdf"]["output_format"] = "json" c := &ParserComponent{Param: param, Setups: setups} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": data, "file_type": "pdf", "name": "Doc1.pdf", diff --git a/internal/ingestion/component/parser_dispatch_test.go b/internal/ingestion/component/parser_dispatch_test.go index dbd2cba282..3506b37cec 100644 --- a/internal/ingestion/component/parser_dispatch_test.go +++ b/internal/ingestion/component/parser_dispatch_test.go @@ -62,7 +62,7 @@ func TestDispatch_OutputFormatValidation_Allowed(t *testing.T) { // Defaults already include markdown → {text, json}. c := &ParserComponent{Param: param, Setups: setups} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": []byte("# Title\n\nbody\n"), "doc_id": "doc.md", "file_type": "md", @@ -116,7 +116,7 @@ func TestDispatch_OutputFormatValidation_Rejection(t *testing.T) { // returns a FileType whose string form matches the setup key. c := &ParserComponent{Param: param, Setups: setups} - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": []byte("# Title\n"), "file_type": "md", }) @@ -142,7 +142,7 @@ func TestDispatch_TextPageMode_NoFileType(t *testing.T) { setups := defaultSetups() c := &ParserComponent{Param: param, Setups: setups} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": []byte("plain content\n"), "doc_id": "unknown", }) @@ -167,7 +167,7 @@ func TestDispatch_SupportedFamilyFailure_HardErrors(t *testing.T) { setups := defaultSetups() c := &ParserComponent{Param: param, Setups: setups} - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": []byte("PDF payload as bytes (not a real PDF — stub test)\n"), "file_type": "pdf", }) @@ -330,7 +330,7 @@ func TestDispatch_PDFMarkdown_UsesConfiguredOutputFormat(t *testing.T) { setups["pdf"]["output_format"] = "markdown" c := &ParserComponent{Param: param, Setups: setups} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": data, "file_type": "pdf", "name": "Doc1.pdf", @@ -363,7 +363,7 @@ func TestDispatch_PDFPlainText_UsesConfiguredBackend(t *testing.T) { setups["pdf"]["output_format"] = "json" c := &ParserComponent{Param: param, Setups: setups} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": data, "file_type": "pdf", "name": "Doc1.pdf", @@ -386,7 +386,7 @@ func TestDispatch_PDFUnsupportedParseMethod_HardErrors(t *testing.T) { setups["pdf"]["parse_method"] = "CustomVLM" c := &ParserComponent{Param: param, Setups: setups} - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": []byte("%PDF-1.4"), "file_type": "pdf", "name": "bad.pdf", @@ -457,7 +457,7 @@ func TestDispatch_PDFVisionJSON_UsesTenantAwareModel(t *testing.T) { setups["pdf"]["output_format"] = "json" c := &ParserComponent{Param: param, Setups: setups} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": []byte("%PDF-1.4"), "file_type": "pdf", "name": "vision.pdf", @@ -523,7 +523,7 @@ func TestDispatch_PDFVisionJSON_PreservesEmptyPages(t *testing.T) { setups["pdf"]["output_format"] = "json" c := &ParserComponent{Param: param, Setups: setups} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": []byte("%PDF-1.4"), "file_type": "pdf", "name": "vision.pdf", @@ -572,7 +572,7 @@ func TestDispatch_PDFMinerUMarkdown_UsesConfiguredBackend(t *testing.T) { setups["pdf"]["output_format"] = "markdown" c := &ParserComponent{Param: param, Setups: setups} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": []byte("%PDF-1.4"), "file_type": "pdf", "name": "sample.pdf", @@ -664,7 +664,7 @@ func TestDispatch_PDFPaddleOCRMarkdown_UsesConfiguredBackend(t *testing.T) { setups["pdf"]["paddleocr_api_key"] = "paddle-secret" c := &ParserComponent{Param: param, Setups: setups} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": []byte("%PDF-1.4"), "file_type": "pdf", "name": "sample.pdf", @@ -713,7 +713,7 @@ func TestDispatch_PDFDoclingMarkdown_UsesConfiguredBackend(t *testing.T) { setups["pdf"]["docling_api_key"] = "doc-secret" c := &ParserComponent{Param: param, Setups: setups} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": []byte("%PDF-1.4"), "file_type": "pdf", "name": "sample.pdf", @@ -751,7 +751,7 @@ func TestDispatch_PDFOpenDataLoaderMarkdown_UsesConfiguredBackend(t *testing.T) setups["pdf"]["opendataloader_apiserver"] = server.URL c := &ParserComponent{Param: param, Setups: setups} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": []byte("%PDF-1.4"), "file_type": "pdf", "name": "sample.pdf", @@ -785,7 +785,7 @@ func TestDispatch_PDFSoMarkMarkdown_UsesConfiguredBackend(t *testing.T) { setups["pdf"]["somark_base_url"] = server.URL c := &ParserComponent{Param: param, Setups: setups} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": []byte("%PDF-1.4"), "file_type": "pdf", "name": "sample.pdf", @@ -821,7 +821,7 @@ func TestDispatch_PDFTCADPMarkdown_UsesConfiguredBackend(t *testing.T) { setups["pdf"]["tcadp_apiserver"] = server.URL c := &ParserComponent{Param: param, Setups: setups} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(t.Context(), nil, map[string]any{ "binary": []byte("%PDF-1.4"), "file_type": "pdf", "name": "sample.pdf", diff --git a/internal/ingestion/component/parser_test.go b/internal/ingestion/component/parser_test.go index 8c7a0218c4..0924b14674 100644 --- a/internal/ingestion/component/parser_test.go +++ b/internal/ingestion/component/parser_test.go @@ -84,7 +84,7 @@ func TestParserComponent_InputsOutputs_NonEmpty(t *testing.T) { // not asserted here. func TestParserComponent_Invoke_TextInput(t *testing.T) { c := &ParserComponent{Param: schema.ParserParam{}.Defaults()} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "binary": "hello world", }) if err != nil { @@ -107,7 +107,7 @@ func TestParserComponent_Invoke_TextInput(t *testing.T) { // pages, in input order, with text intact. func TestParserComponent_Invoke_PageRangeFilter(t *testing.T) { c := &ParserComponent{Param: schema.ParserParam{}.Defaults()} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "binary": "pageA\fpageB\fpageC", }) if err != nil { @@ -146,7 +146,7 @@ func TestParserComponent_Invoke_DeterministicMerge(t *testing.T) { input := "p1\fp2\fp3\fp4\fp5\fp6\fp7\fp8" // First call: produce the canonical bytes. - first, err := c.Invoke(context.Background(), map[string]any{ + first, err := c.Invoke(context.Background(), nil, map[string]any{ "binary": input, }) if err != nil { @@ -159,7 +159,7 @@ func TestParserComponent_Invoke_DeterministicMerge(t *testing.T) { // Subsequent calls: must produce the same bytes. for i := 0; i < 5; i++ { - got, err := c.Invoke(context.Background(), map[string]any{ + got, err := c.Invoke(context.Background(), nil, map[string]any{ "binary": input, }) if err != nil { @@ -228,7 +228,7 @@ func TestParserComponent_New_Overrides(t *testing.T) { // doc_id input flows through to the "name" output. func TestParserComponent_Invoke_DocIDCarried(t *testing.T) { c := &ParserComponent{Param: schema.ParserParam{}.Defaults()} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "binary": "x", "doc_id": "doc-123", }) @@ -263,7 +263,7 @@ func TestParserComponent_Invoke_ResolvesBinaryFromDocID(t *testing.T) { } c := &ParserComponent{Param: schema.ParserParam{}.Defaults()} - out, err := c.Invoke(context.Background(), map[string]any{"doc_id": "doc-parser"}) + out, err := c.Invoke(context.Background(), db, map[string]any{"doc_id": "doc-parser"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -286,7 +286,7 @@ func TestParserComponent_Invoke_ResolvesBinaryFromBucketPath(t *testing.T) { } c := &ParserComponent{Param: schema.ParserParam{}.Defaults()} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "bucket": "bucket-1", "path": "docs/explicit.txt", }) @@ -309,7 +309,7 @@ func TestParserComponent_Invoke_ResolvesBinaryFromBucketPath(t *testing.T) { // without decoding it). func TestParserComponent_Invoke_RejectsInvalidUTF8(t *testing.T) { c := &ParserComponent{Param: schema.ParserParam{}.Defaults()} - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(context.Background(), nil, map[string]any{ // 0xFF alone is not valid UTF-8 start byte. "binary": string([]byte{0xFF, 0xFE, 0xFD}), }) @@ -326,7 +326,7 @@ func TestParserComponent_Invoke_RejectsInvalidUTF8(t *testing.T) { // string. func TestParserComponent_Invoke_AcceptsBytes(t *testing.T) { c := &ParserComponent{Param: schema.ParserParam{}.Defaults()} - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "binary": []byte("alpha\fbeta"), }) if err != nil { diff --git a/internal/ingestion/component/slice3_integration_test.go b/internal/ingestion/component/slice3_integration_test.go index c2e494f500..9d94fbaacf 100644 --- a/internal/ingestion/component/slice3_integration_test.go +++ b/internal/ingestion/component/slice3_integration_test.go @@ -35,7 +35,7 @@ func TestTokenizer_FallsBackToContentWithWeight(t *testing.T) { c.param.Fields = []string{"text"} c.param.FilenameEmbdWeight = 0 - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc", "output_format": "chunks", "chunks": []map[string]any{ @@ -66,7 +66,7 @@ func TestTokenizer_DoesNotChangeChunkText(t *testing.T) { c.param.Fields = []string{"text"} c.param.FilenameEmbdWeight = 0 - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc", "output_format": "chunks", "chunks": []map[string]any{ diff --git a/internal/ingestion/component/tokenizer.go b/internal/ingestion/component/tokenizer.go index 5029772032..524801a2e9 100644 --- a/internal/ingestion/component/tokenizer.go +++ b/internal/ingestion/component/tokenizer.go @@ -102,6 +102,7 @@ import ( "time" "go.uber.org/zap" + "gorm.io/gorm" "ragflow/internal/agent/runtime" "ragflow/internal/common" @@ -160,7 +161,7 @@ type Embedder interface { // embeddingModel is the Tokenizer-scoped embedding-model identifier (from the // component's setups); an empty value tells the resolver to fall back to the // dataset's configured model. -type EmbedderResolver func(tenantID, kbID, embeddingModel string) (Embedder, error) +type EmbedderResolver func(ctx context.Context, tenantID, kbID, embeddingModel string) (Embedder, error) // DefaultEmbedderResolver is the production embedder resolver. It is nil in // this leaf package — which must not import internal/service (see the @@ -323,7 +324,7 @@ func (c *TokenizerComponent) Outputs() map[string]string { // embedding batch (python tokenizer.py:80-82 `if not cleaned_txt: // continue`), but the chunk still carries tokenized fields if // `full_text` is in `search_method`. -func (c *TokenizerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (c *TokenizerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { // Run-level metadata lives in the workflow-wide CanvasState.Globals // bag (seeded at pipeline start, published by the File component), // not in the upstream output map — see GlobalOrInput. @@ -401,7 +402,7 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, em if resolver == nil { return nil, 0, fmt.Errorf("Tokenizer: embedding requested but no embedder resolver configured") } - embedder, err := resolver(tenantID, kbID, embeddingModel) + embedder, err := resolver(ctx, tenantID, kbID, embeddingModel) if err != nil { return nil, 0, fmt.Errorf("Tokenizer: resolve embedder: %w", err) } diff --git a/internal/ingestion/component/tokenizer_test.go b/internal/ingestion/component/tokenizer_test.go index 4883504aab..afdebcb587 100644 --- a/internal/ingestion/component/tokenizer_test.go +++ b/internal/ingestion/component/tokenizer_test.go @@ -67,7 +67,7 @@ func TestTokenizerComponent_Invoke_HappyPath(t *testing.T) { {"text": "bravo chunk text"}, {"text": "charlie chunk text"}, } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "tenant_id": "t1", "model_id": "embd-1", "name": "doc.pdf", @@ -137,7 +137,7 @@ func TestTokenizerComponent_Invoke_Unicode(t *testing.T) { chunks[i] = map[string]any{"text": txt} } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "output_format": "chunks", "chunks": chunks, }) @@ -166,7 +166,7 @@ func TestTokenizerComponent_Invoke_TextPayload(t *testing.T) { c, _ := NewTokenizerComponent(map[string]any{ "search_method": []any{"full_text"}, }) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "note.txt", "output_format": "text", "text": "plain payload", @@ -192,7 +192,7 @@ func TestTokenizerComponent_Invoke_JSONPayload(t *testing.T) { c, _ := NewTokenizerComponent(map[string]any{ "search_method": []any{"full_text"}, }) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "note.pdf", "output_format": "json", "json": []map[string]any{{"text": "row one"}, {"text": "row two"}}, @@ -220,7 +220,7 @@ func TestTokenizerComponent_Invoke_BatchedEmbedding(t *testing.T) { {"text": "two"}, {"text": "three"}, } - if _, err := c.Invoke(context.Background(), map[string]any{ + if _, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.txt", "output_format": "chunks", "chunks": chunks, @@ -241,7 +241,7 @@ func TestTokenizerComponent_Invoke_FullTextOnly(t *testing.T) { c, _ := NewTokenizerComponent(map[string]any{ "search_method": []any{"full_text"}, }) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "output_format": "chunks", "chunks": []map[string]any{{"text": "alpha bravo"}}, }) @@ -270,7 +270,7 @@ func TestTokenizerComponent_Invoke_KeywordSplitCJK(t *testing.T) { c, _ := NewTokenizerComponent(map[string]any{ "search_method": []any{"full_text"}, }) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "output_format": "chunks", "chunks": []map[string]any{{"text": "alpha", "keywords": "kw1，kw2；kw3"}}, }) @@ -296,7 +296,7 @@ func TestTokenizerComponent_Invoke_KeywordSplitCJK(t *testing.T) { func TestTokenizerComponent_Invoke_FullTextAndEmbedding(t *testing.T) { requireTokenizerPool(t) c, _ := withStubEmbedder(t, 4) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.pdf", "output_format": "chunks", "chunks": []map[string]any{{"text": "alpha bravo"}}, @@ -326,7 +326,7 @@ func TestTokenizerComponent_Invoke_FullTextAndEmbedding(t *testing.T) { func TestTokenizerComponent_Invoke_EmbedNoResolver(t *testing.T) { requireTokenizerPool(t) c, _ := NewTokenizerComponent(nil) - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(context.Background(), nil, map[string]any{ "output_format": "chunks", "chunks": []map[string]any{{"text": "alpha"}}, }) @@ -345,7 +345,7 @@ func TestTokenizerComponent_Invoke_EmbedderError(t *testing.T) { c, stub := withStubEmbedder(t, 4) stub.err = errors.New("simulated upstream error") - _, err := c.Invoke(context.Background(), map[string]any{ + _, err := c.Invoke(context.Background(), nil, map[string]any{ "output_format": "chunks", "chunks": []map[string]any{{"text": "alpha"}}, }) @@ -371,7 +371,7 @@ func TestTokenizerComponent_Invoke_EncoderCountMismatch(t *testing.T) { } c := cIntf.(*TokenizerComponent) _ = stub - _, err = c.Invoke(context.Background(), map[string]any{ + _, err = c.Invoke(context.Background(), nil, map[string]any{ "output_format": "chunks", "chunks": []map[string]any{{"text": "a"}, {"text": "b"}, {"text": "c"}}, }) @@ -408,7 +408,7 @@ func TestTokenizerComponent_Invoke_HonorsTimeout(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - _, err := c.Invoke(ctx, map[string]any{ + _, err := c.Invoke(ctx, nil, map[string]any{ "output_format": "chunks", "chunks": []map[string]any{{"text": "alpha"}}, }) @@ -454,7 +454,7 @@ func TestTokenizerComponent_Smoke_EndToEnd(t *testing.T) { } start := time.Now() - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "tenant_id": "tenant-smoke", "model_id": "embd-smoke", "name": "smoke.pdf", @@ -502,7 +502,7 @@ func TestTokenizerComponent_Embedding_MergesTitleAndContentVectors(t *testing.T) {vectors: [][]float64{{10, 20}}, tokenCount: 7}, {vectors: [][]float64{{1, 2}, {3, 4}}, tokenCount: 11}, } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.pdf", "output_format": "chunks", "chunks": []map[string]any{{"text": "alpha"}, {"text": "beta"}}, @@ -537,7 +537,7 @@ func TestTokenizerComponent_Embedding_UsesFilenameWeight(t *testing.T) { t.Fatalf("NewTokenizerComponentWithResolver: %v", err) } c := cIntf.(*TokenizerComponent) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.pdf", "output_format": "chunks", "chunks": []map[string]any{{"text": "alpha"}}, @@ -567,7 +567,7 @@ func TestTokenizerComponent_Embedding_EmptyNameWarnsAndUsesContentVector(t *test log.SetFlags(prevFlags) }) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": " ", "output_format": "chunks", "chunks": []map[string]any{{"text": "alpha"}}, @@ -600,7 +600,7 @@ func TestTokenizerComponent_Embedding_UsesRawNameNotTrimmed(t *testing.T) { requireTokenizerPool(t) c, stub := withStubEmbedder(t, 2) - if _, err := c.Invoke(context.Background(), map[string]any{ + if _, err := c.Invoke(context.Background(), nil, map[string]any{ "name": " report.pdf ", "output_format": "chunks", "chunks": []map[string]any{{"text": "alpha"}}, @@ -622,7 +622,7 @@ func TestTokenizerComponent_Embedding_TruncatesByMaxTokensMinus10(t *testing.T) c, stub := withStubEmbedder(t, 2) stub.maxTokens = 12 longText := strings.Repeat("hello world ", 20) - if _, err := c.Invoke(context.Background(), map[string]any{ + if _, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.pdf", "output_format": "chunks", "chunks": []map[string]any{{"text": longText}}, @@ -643,7 +643,7 @@ func TestTokenizerComponent_Embedding_TruncatesByMaxTokensMinus10(t *testing.T) func TestTokenizerComponent_Embedding_SkipsEmptyCleanedTextsButReturnsZeroWhenAllSkipped(t *testing.T) { requireTokenizerPool(t) c, stub := withStubEmbedder(t, 2) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.pdf", "output_format": "chunks", "chunks": []map[string]any{ @@ -676,7 +676,7 @@ func TestTokenizerComponent_Embedding_SetsTokenConsumptionIncludingTitleCall(t * tokenizerEmbeddingBatchSize = 1 t.Cleanup(func() { tokenizerEmbeddingBatchSize = prevBatchSize }) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.pdf", "output_format": "chunks", "chunks": []map[string]any{{"text": "alpha"}, {"text": "beta"}}, @@ -696,7 +696,7 @@ func TestTokenizerComponent_Embedding_BatchesByConfiguredBatchSize(t *testing.T) tokenizerEmbeddingBatchSize = 2 t.Cleanup(func() { tokenizerEmbeddingBatchSize = prevBatchSize }) - if _, err := c.Invoke(context.Background(), map[string]any{ + if _, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.pdf", "output_format": "chunks", "chunks": []map[string]any{ diff --git a/internal/ingestion/component/tokenizer_unit_test.go b/internal/ingestion/component/tokenizer_unit_test.go index 4710044a08..b50ac430c4 100644 --- a/internal/ingestion/component/tokenizer_unit_test.go +++ b/internal/ingestion/component/tokenizer_unit_test.go @@ -101,7 +101,7 @@ func newStubEmbedder(dim int) *stubEmbedder { func withStubEmbedder(t *testing.T, dim int) (*TokenizerComponent, *stubEmbedder) { t.Helper() stub := newStubEmbedder(dim) - comp, err := NewTokenizerComponentWithResolver(nil, func(_, _, _ string) (Embedder, error) { return stub, nil }) + comp, err := NewTokenizerComponentWithResolver(nil, func(ctx context.Context, _, _, _ string) (Embedder, error) { return stub, nil }) if err != nil { t.Fatalf("NewTokenizerComponentWithResolver: %v", err) } @@ -135,12 +135,8 @@ func TestTokenizerComponent_Registered(t *testing.T) { func TestTokenizerComponent_Invoke_EmptyChunks(t *testing.T) { c, stub := withStubEmbedder(t, 4) _ = stub - var err error - if err != nil { - t.Fatalf("NewTokenizerComponent: %v", err) - } - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "output_format": "chunks", "chunks": []map[string]any{}, }) @@ -168,7 +164,7 @@ func TestTokenizerComponent_Invoke_EmptyChunks(t *testing.T) { func TestTokenizerComponent_Invoke_NilChunks(t *testing.T) { c, stub := withStubEmbedder(t, 4) _ = stub - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "output_format": "chunks", }) if err != nil { @@ -183,13 +179,13 @@ func TestTokenizerComponent_Invoke_NilChunks(t *testing.T) { func TestTokenizerComponent_Invoke_EmbeddingOnly(t *testing.T) { cIntf, err := NewTokenizerComponentWithResolver(map[string]any{ "search_method": []any{"embedding"}, - }, func(_, _, _ string) (Embedder, error) { + }, func(ctx context.Context, _, _, _ string) (Embedder, error) { return newStubEmbedder(4), nil }) if err != nil { t.Fatalf("NewTokenizerComponentWithResolver: %v", err) } - out, err := cIntf.(*TokenizerComponent).Invoke(context.Background(), map[string]any{ + out, err := cIntf.(*TokenizerComponent).Invoke(context.Background(), nil, map[string]any{ "name": "doc.pdf", "output_format": "chunks", "chunks": []map[string]any{{"text": "alpha bravo"}}, @@ -216,7 +212,7 @@ func TestTokenizerComponent_Invoke_EmbeddingOnly(t *testing.T) { // empty chunk list, so tokenizeChunks is a no-op and the C++ pool is not needed. func TestTokenizerComponent_Embedding_ZeroChunksStillEmitsConsumptionZero(t *testing.T) { c, stub := withStubEmbedder(t, 2) - out, err := c.Invoke(context.Background(), map[string]any{ + out, err := c.Invoke(context.Background(), nil, map[string]any{ "name": "doc.pdf", "output_format": "chunks", "chunks": []map[string]any{}, diff --git a/internal/ingestion/pipeline/pipeline_test.go b/internal/ingestion/pipeline/pipeline_test.go index ce1f93cb4b..1aa9970b22 100644 --- a/internal/ingestion/pipeline/pipeline_test.go +++ b/internal/ingestion/pipeline/pipeline_test.go @@ -29,6 +29,7 @@ import ( "github.com/alicebob/miniredis/v2" "github.com/redis/go-redis/v9" + "gorm.io/gorm" ) type mockCanvasStage struct { @@ -37,7 +38,7 @@ type mockCanvasStage struct { calls int } -func (m *mockCanvasStage) Invoke(_ context.Context, inputs map[string]any) (map[string]any, error) { +func (m *mockCanvasStage) Invoke(_ context.Context, _ *gorm.DB, inputs map[string]any) (map[string]any, error) { m.called = true m.calls++ out := cloneMapOrEmpty(inputs) @@ -58,12 +59,12 @@ type oneShotErrStage struct { n int } -func (s *oneShotErrStage) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { +func (s *oneShotErrStage) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { s.n++ if s.n == 1 { return nil, errors.New("simulated crash") } - return s.mockCanvasStage.Invoke(ctx, inputs) + return s.mockCanvasStage.Invoke(ctx, db, inputs) } func (s *oneShotErrStage) Inputs() map[string]string { return s.mockCanvasStage.Inputs() } func (s *oneShotErrStage) Outputs() map[string]string { return s.mockCanvasStage.Outputs() } @@ -171,7 +172,7 @@ func TestNewPipelineFromDSLUnwrapsTemplateDSL(t *testing.T) { type errCanvasStage struct{} -func (e *errCanvasStage) Invoke(_ context.Context, _ map[string]any) (map[string]any, error) { +func (e *errCanvasStage) Invoke(_ context.Context, _ *gorm.DB, _ map[string]any) (map[string]any, error) { return nil, &stageError{Stage: "p.RunErrStage", Reason: "intentional"} } func (e *errCanvasStage) Inputs() map[string]string { return nil } @@ -181,7 +182,7 @@ type factorySentinelStage struct { marker string } -func (s *factorySentinelStage) Invoke(_ context.Context, inputs map[string]any) (map[string]any, error) { +func (s *factorySentinelStage) Invoke(_ context.Context, _ *gorm.DB, inputs map[string]any) (map[string]any, error) { out := cloneMapOrEmpty(inputs) out["marker"] = s.marker return out, nil diff --git a/internal/ingestion/pipeline/template_integration_test.go b/internal/ingestion/pipeline/template_integration_test.go index 0b25a3d0a5..8d22b96ff4 100644 --- a/internal/ingestion/pipeline/template_integration_test.go +++ b/internal/ingestion/pipeline/template_integration_test.go @@ -842,7 +842,7 @@ func attachFixedEmbedderFactory(t *testing.T, pipe *Pipeline) { t.Helper() pipe.WithComponentFactory(func(name string, params map[string]any) (runtime.Component, error) { if name == componentpkg.ComponentNameTokenizer { - return componentpkg.NewTokenizerComponentWithResolver(params, func(_, _, _ string) (componentpkg.Embedder, error) { + return componentpkg.NewTokenizerComponentWithResolver(params, func(ctx context.Context, _, _, _ string) (componentpkg.Embedder, error) { return fixedEmbedder{}, nil }) } diff --git a/internal/ingestion/task/embedder.go b/internal/ingestion/task/embedder.go index 80ca34f202..6e0c1ec97a 100644 --- a/internal/ingestion/task/embedder.go +++ b/internal/ingestion/task/embedder.go @@ -61,11 +61,11 @@ func (e *embedder) Encode(ctx context.Context, texts []string) ([]componentpkg.E // injectable deps so the resolution logic stays unit-testable without a live // model provider / DB. func newEmbedderResolver( - getKBEmbdID func(kbID string) (string, error), - getEmbeddingModel func(tenantID, embdID string) (*models.EmbeddingModel, error), + getKBEmbdID func(ctx context.Context, kbID string) (string, error), + getEmbeddingModel func(ctx context.Context, tenantID, embdID string) (*models.EmbeddingModel, error), ) componentpkg.EmbedderResolver { - return func(tenantID, kbID, _ string) (componentpkg.Embedder, error) { - embdID, err := getKBEmbdID(kbID) + return func(ctx context.Context, tenantID, kbID, _ string) (componentpkg.Embedder, error) { + embdID, err := getKBEmbdID(ctx, kbID) if err != nil { return nil, fmt.Errorf("embedder: resolve kb embd_id for kb_id=%s: %w", kbID, err) } @@ -73,7 +73,7 @@ func newEmbedderResolver( if embdID == "" { return nil, nil } - model, err := getEmbeddingModel(tenantID, embdID) + model, err := getEmbeddingModel(ctx, tenantID, embdID) if err != nil { return nil, err } @@ -90,8 +90,8 @@ func newEmbedderResolver( // composition root for ingestion runs. func init() { componentpkg.DefaultEmbedderResolver = newEmbedderResolver( - func(kbID string) (string, error) { - kb, err := dao.NewKnowledgebaseDAO().GetByID(kbID) + func(ctx context.Context, kbID string) (string, error) { + kb, err := dao.NewKnowledgebaseDAO().GetByID(ctx, dao.DB, kbID) if err != nil { return "", err } diff --git a/internal/ingestion/task/embedder_test.go b/internal/ingestion/task/embedder_test.go index df2803e789..7df8a57bcc 100644 --- a/internal/ingestion/task/embedder_test.go +++ b/internal/ingestion/task/embedder_test.go @@ -31,15 +31,16 @@ func makeEmbeddingModelForResolver() *models.EmbeddingModel { func TestEmbedderResolver_UsesKBEmbdID(t *testing.T) { var gotTenantID, gotEmbdID string resolver := newEmbedderResolver( - func(kbID string) (string, error) { + func(ctx context.Context, kbID string) (string, error) { return "kb-embd-1", nil }, - func(tenantID, embdID string) (*models.EmbeddingModel, error) { + func(ctx context.Context, tenantID, embdID string) (*models.EmbeddingModel, error) { gotTenantID, gotEmbdID = tenantID, embdID return makeEmbeddingModelForResolver(), nil }, ) - emb, err := resolver("tenant-1", "kb-1", "should-be-ignored") + ctx := t.Context() + emb, err := resolver(ctx, "tenant-1", "kb-1", "should-be-ignored") if err != nil { t.Fatalf("resolver: %v", err) } @@ -53,15 +54,16 @@ func TestEmbedderResolver_UsesKBEmbdID(t *testing.T) { func TestEmbedderResolver_EmptyKBEmbdIDReturnsNil(t *testing.T) { resolver := newEmbedderResolver( - func(kbID string) (string, error) { + func(ctx context.Context, kbID string) (string, error) { return "", nil }, - func(string, string) (*models.EmbeddingModel, error) { + func(context.Context, string, string) (*models.EmbeddingModel, error) { t.Fatal("model resolver should not be called when kb embd_id is empty") return nil, nil }, ) - emb, err := resolver("tenant-1", "kb-1", "ignored") + ctx := t.Context() + emb, err := resolver(ctx, "tenant-1", "kb-1", "ignored") if err != nil { t.Fatalf("resolver: %v", err) } diff --git a/internal/ingestion/task/pipeline_executor_test.go b/internal/ingestion/task/pipeline_executor_test.go index 3c63367963..a32a173c01 100644 --- a/internal/ingestion/task/pipeline_executor_test.go +++ b/internal/ingestion/task/pipeline_executor_test.go @@ -432,7 +432,7 @@ func (r *recordingProgressSink) OnComponentProgress(ctx context.Context, ev pipe type sinkPassthroughStage struct{} -func (sinkPassthroughStage) Invoke(_ context.Context, inputs map[string]any) (map[string]any, error) { +func (sinkPassthroughStage) Invoke(_ context.Context, _ *gorm.DB, inputs map[string]any) (map[string]any, error) { return inputs, nil } diff --git a/internal/ingestion/task/task_context.go b/internal/ingestion/task/task_context.go index e3759694ac..1acd1f3a14 100644 --- a/internal/ingestion/task/task_context.go +++ b/internal/ingestion/task/task_context.go @@ -66,7 +66,7 @@ func LoadFromIngestionTask(ctx context.Context, ingestionTask *entity.IngestionT return nil, fmt.Errorf("document %s not found", ingestionTask.DocumentID) } - kb, err := dao.NewKnowledgebaseDAO().GetByID(doc.KbID) + kb, err := dao.NewKnowledgebaseDAO().GetByID(ctx, dao.DB, doc.KbID) if err != nil || kb == nil { return nil, fmt.Errorf("error when load knowledgebase %s: %w", doc.KbID, err) } diff --git a/internal/mcp/connector.go b/internal/mcp/connector.go index fbdfd23c88..a03e713920 100644 --- a/internal/mcp/connector.go +++ b/internal/mcp/connector.go @@ -17,6 +17,7 @@ package mcp import ( + "context" "encoding/json" "fmt" "strings" @@ -26,9 +27,9 @@ import ( // service layer calls, avoiding HTTP round-trips to self. type ServiceConnector struct { userID string - listDatasets func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) - listChats func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) - retrieval func(userID string, req RetrievalRequest) (string, error) + listDatasets func(ctx context.Context, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) + listChats func(ctx context.Context, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) + retrieval func(ctx context.Context, userID string, req RetrievalRequest) (string, error) } // NewServiceConnector creates a ServiceConnector. @@ -36,9 +37,9 @@ type ServiceConnector struct { // does not import the service layer directly. func NewServiceConnector( userID string, - listDatasetsFunc func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error), - listChatsFunc func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error), - retrievalFunc func(userID string, req RetrievalRequest) (string, error), + listDatasetsFunc func(ctx context.Context, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error), + listChatsFunc func(ctx context.Context, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error), + retrievalFunc func(ctx context.Context, userID string, req RetrievalRequest) (string, error), ) *ServiceConnector { return &ServiceConnector{ userID: userID, @@ -50,8 +51,8 @@ func NewServiceConnector( // ListDatasets returns newline-delimited JSON, each line being // {"id": "...", "description": "..."} for a dataset. -func (c *ServiceConnector) ListDatasets(page, pageSize int, orderby string, desc bool) (string, error) { - data, _, err := c.listDatasets(c.userID, page, pageSize, orderby, desc) +func (c *ServiceConnector) ListDatasets(ctx context.Context, page, pageSize int, orderby string, desc bool) (string, error) { + data, _, err := c.listDatasets(ctx, c.userID, page, pageSize, orderby, desc) if err != nil { return "", fmt.Errorf("list datasets: %w", err) } @@ -88,8 +89,8 @@ func (c *ServiceConnector) ListDatasets(page, pageSize int, orderby string, desc // ListChats returns newline-delimited JSON, each line being // {"id": "...", "name": "...", "description": "..."} for a chat assistant. -func (c *ServiceConnector) ListChats(page, pageSize int, orderby string, desc bool) (string, error) { - data, _, err := c.listChats(c.userID, page, pageSize, orderby, desc) +func (c *ServiceConnector) ListChats(ctx context.Context, page, pageSize int, orderby string, desc bool) (string, error) { + data, _, err := c.listChats(ctx, c.userID, page, pageSize, orderby, desc) if err != nil { return "", fmt.Errorf("list chats: %w", err) } @@ -125,6 +126,6 @@ func (c *ServiceConnector) ListChats(page, pageSize int, orderby string, desc bo // Retrieval executes a retrieval via the in-process service and returns // the result as a JSON string. -func (c *ServiceConnector) Retrieval(req RetrievalRequest) (string, error) { - return c.retrieval(c.userID, req) +func (c *ServiceConnector) Retrieval(ctx context.Context, req RetrievalRequest) (string, error) { + return c.retrieval(ctx, c.userID, req) } diff --git a/internal/mcp/server.go b/internal/mcp/server.go index 12b1a25aea..dd73e4ee96 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -17,6 +17,7 @@ package mcp import ( + "context" "encoding/json" "fmt" "strings" @@ -28,15 +29,15 @@ import ( type Connector interface { // ListDatasets returns newline-delimited JSON lines, each containing // at minimum {"id": "...", "description": "..."}. - ListDatasets(page, pageSize int, orderby string, desc bool) (string, error) + ListDatasets(ctx context.Context, page, pageSize int, orderby string, desc bool) (string, error) // ListChats returns newline-delimited JSON lines, each containing // at minimum {"id": "...", "name": "...", "description": "..."}. - ListChats(page, pageSize int, orderby string, desc bool) (string, error) + ListChats(ctx context.Context, page, pageSize int, orderby string, desc bool) (string, error) // Retrieval executes a retrieval request and returns the result as // a JSON string. - Retrieval(req RetrievalRequest) (string, error) + Retrieval(ctx context.Context, req RetrievalRequest) (string, error) } // RetrievalRequest carries all parameters for a retrieval query. @@ -71,7 +72,7 @@ func NewServer(connector Connector) *Server { // HandleRequest dispatches a raw JSON-RPC request body and returns the // serialized JSON-RPC response. Returns nil if the request is a // notification (no id) and requires no response. -func (s *Server) HandleRequest(body []byte) ([]byte, bool, error) { +func (s *Server) HandleRequest(ctx context.Context, body []byte) ([]byte, bool, error) { // Try to decode as a request (with an id) first. var req JSONRPCRequest if err := json.Unmarshal(body, &req); err != nil { @@ -98,9 +99,9 @@ func (s *Server) HandleRequest(body []byte) ([]byte, bool, error) { case "initialize": resp = s.handleInitialize(req.ID) case "tools/list": - resp = s.handleListTools(req.ID) + resp = s.handleListTools(ctx, req.ID) case "tools/call": - resp = s.handleCallTool(req.ID, req.Params) + resp = s.handleCallTool(ctx, req.ID, req.Params) case "ping": resp = s.handlePing(req.ID) default: @@ -133,14 +134,14 @@ func (s *Server) handlePing(id json.RawMessage) JSONRPCResponse { return NewSuccessResponse(id, struct{}{}) } -func (s *Server) handleListTools(id json.RawMessage) JSONRPCResponse { +func (s *Server) handleListTools(ctx context.Context, id json.RawMessage) JSONRPCResponse { // Fetch dataset and chat descriptions for embedding into tool descriptions, // matching the Python MCP server behavior. - datasetDescription, err := s.connector.ListDatasets(1, 100, "create_time", true) + datasetDescription, err := s.connector.ListDatasets(ctx, 1, 100, "create_time", true) if err != nil { datasetDescription = "" } - chatDescription, err := s.connector.ListChats(1, 30, "create_time", true) + chatDescription, err := s.connector.ListChats(ctx, 1, 30, "create_time", true) if err != nil { chatDescription = "" } @@ -267,7 +268,7 @@ func (s *Server) handleListTools(id json.RawMessage) JSONRPCResponse { return NewSuccessResponse(id, ListToolsResult{Tools: tools}) } -func (s *Server) handleCallTool(id json.RawMessage, rawParams json.RawMessage) JSONRPCResponse { +func (s *Server) handleCallTool(ctx context.Context, id json.RawMessage, rawParams json.RawMessage) JSONRPCResponse { var params CallToolParams if err := json.Unmarshal(rawParams, ¶ms); err != nil { return NewErrorResponse(id, ErrCodeInvalidParams, @@ -280,18 +281,18 @@ func (s *Server) handleCallTool(id json.RawMessage, rawParams json.RawMessage) J switch params.Name { case "ragflow_retrieval": - return s.callRagflowRetrieval(id, params.Arguments) + return s.callRagflowRetrieval(ctx, id, params.Arguments) case "ragflow_list_datasets": - return s.callListDatasets(id, params.Arguments) + return s.callListDatasets(ctx, id, params.Arguments) case "ragflow_list_chats": - return s.callListChats(id, params.Arguments) + return s.callListChats(ctx, id, params.Arguments) default: return NewErrorResponse(id, ErrCodeMethodNotFound, fmt.Sprintf("Tool not found: %s", params.Name)) } } -func (s *Server) callRagflowRetrieval(id json.RawMessage, args map[string]interface{}) JSONRPCResponse { +func (s *Server) callRagflowRetrieval(ctx context.Context, id json.RawMessage, args map[string]interface{}) JSONRPCResponse { req := RetrievalRequest{ Page: getBoundedIntArg(args, "page", 1, 1, 1_000_000), PageSize: getBoundedIntArg(args, "page_size", 10, 1, 100), @@ -315,29 +316,29 @@ func (s *Server) callRagflowRetrieval(id json.RawMessage, args map[string]interf return NewSuccessResponse(id, NewErrorResult("question is required")) } - result, err := s.connector.Retrieval(req) + result, err := s.connector.Retrieval(ctx, req) if err != nil { return NewSuccessResponse(id, NewErrorResult(err.Error())) } return NewSuccessResponse(id, NewTextResult(result)) } -func (s *Server) callListDatasets(id json.RawMessage, args map[string]interface{}) JSONRPCResponse { +func (s *Server) callListDatasets(ctx context.Context, id json.RawMessage, args map[string]interface{}) JSONRPCResponse { page := getBoundedIntArg(args, "page", 1, 1, 1_000_000) pageSize := getBoundedIntArg(args, "page_size", 100, 1, 1000) - result, err := s.connector.ListDatasets(page, pageSize, "create_time", true) + result, err := s.connector.ListDatasets(ctx, page, pageSize, "create_time", true) if err != nil { return NewSuccessResponse(id, NewErrorResult(err.Error())) } return NewSuccessResponse(id, NewTextResult(result)) } -func (s *Server) callListChats(id json.RawMessage, args map[string]interface{}) JSONRPCResponse { +func (s *Server) callListChats(ctx context.Context, id json.RawMessage, args map[string]interface{}) JSONRPCResponse { page := getBoundedIntArg(args, "page", 1, 1, 1_000_000) pageSize := getBoundedIntArg(args, "page_size", 30, 1, 100) - result, err := s.connector.ListChats(page, pageSize, "create_time", true) + result, err := s.connector.ListChats(ctx, page, pageSize, "create_time", true) if err != nil { return NewSuccessResponse(id, NewErrorResult(err.Error())) } diff --git a/internal/service/agent_run_e2e_test.go b/internal/service/agent_run_e2e_test.go index e2141eac5b..2bb04e42c9 100644 --- a/internal/service/agent_run_e2e_test.go +++ b/internal/service/agent_run_e2e_test.go @@ -56,6 +56,7 @@ import ( "github.com/alicebob/miniredis/v2" "github.com/redis/go-redis/v9" + "gorm.io/gorm" ) // makeCanvasWithDSL inserts a canvas + tenant + user + a published @@ -1689,7 +1690,7 @@ func TestRunAgent_AllFixture_CategorizeResume(t *testing.T) { type categorizeResumeInvoker struct{} -func (i *categorizeResumeInvoker) Invoke(_ context.Context, req component.ChatInvokeRequest) (*component.ChatInvokeResponse, error) { +func (i *categorizeResumeInvoker) Invoke(_ context.Context, _ *gorm.DB, req component.ChatInvokeRequest) (*component.ChatInvokeResponse, error) { return &component.ChatInvokeResponse{ Content: "Retrieval", Model: req.ModelName, diff --git a/internal/service/chat.go b/internal/service/chat.go index 939ed6f94e..869a78f8de 100644 --- a/internal/service/chat.go +++ b/internal/service/chat.go @@ -134,7 +134,7 @@ func (s *ChatService) ListChats(ctx context.Context, userID, status, keywords st // Enrich with knowledge base names chatsWithKBNames := make([]*ChatWithKBNames, 0, len(chats)) for _, chat := range chats { - kbNames, datasetIDs := s.getDatasetNamesAndIDs(chat.KBIDs) + kbNames, datasetIDs := s.getDatasetNamesAndIDs(ctx, chat.KBIDs) chatsWithKBNames = append(chatsWithKBNames, &ChatWithKBNames{ Chat: &chat.Chat, KBNames: kbNames, @@ -224,7 +224,7 @@ func (s *ChatService) Create(ctx context.Context, userID string, req map[string] req["name"] = name if datasetIDsValue, ok := req["dataset_ids"]; ok { - kbIDs, err := s.validateCreateDatasetIDs(datasetIDsValue, userID) + kbIDs, err := s.validateCreateDatasetIDs(ctx, datasetIDsValue, userID) if err != nil { return nil, common.CodeDataError, err } @@ -339,7 +339,7 @@ func (s *ChatService) Create(ctx context.Context, userID string, req map[string] return nil, common.CodeDataError, errors.New("failed to retrieve created chat") } - response, err := s.buildCreateChatResponse(chat) + response, err := s.buildCreateChatResponse(ctx, chat) if err != nil { return nil, common.CodeServerError, err } @@ -364,7 +364,7 @@ func validateCreateChatName(value interface{}) (string, error) { return name, nil } -func (s *ChatService) validateCreateDatasetIDs(value interface{}, tenantID string) ([]string, error) { +func (s *ChatService) validateCreateDatasetIDs(ctx context.Context, value interface{}, tenantID string) ([]string, error) { if value == nil { return []string{}, nil } @@ -384,10 +384,10 @@ func (s *ChatService) validateCreateDatasetIDs(value interface{}, tenantID strin } for _, datasetID := range normalizedIDs { - if !s.kbDAO.Accessible(datasetID, tenantID) { + if !s.kbDAO.Accessible(ctx, dao.DB, datasetID, tenantID) { return nil, fmt.Errorf("you don't own the dataset %s", datasetID) } - kb, err := s.kbDAO.GetByID(datasetID) + kb, err := s.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil { return nil, fmt.Errorf("you don't own the dataset %s", datasetID) } @@ -578,12 +578,12 @@ func stringPtrIfNotEmpty(value string) *string { return &value } -func (s *ChatService) buildCreateChatResponse(chat *entity.Chat) (map[string]interface{}, error) { +func (s *ChatService) buildCreateChatResponse(ctx context.Context, chat *entity.Chat) (map[string]interface{}, error) { data, err := structToMap(chat) if err != nil { return nil, err } - kbNames, datasetIDs := s.getDatasetNamesAndIDs(chat.KBIDs) + kbNames, datasetIDs := s.getDatasetNamesAndIDs(ctx, chat.KBIDs) data["dataset_ids"] = datasetIDs delete(data, "kb_ids") data["kb_names"] = kbNames @@ -735,7 +735,7 @@ func isTruthy(value interface{}) bool { } // getDatasetNamesAndIDs gets knowledge base names by IDs -func (s *ChatService) getDatasetNamesAndIDs(kbIDs entity.JSONSlice) ([]string, []string) { +func (s *ChatService) getDatasetNamesAndIDs(ctx context.Context, kbIDs entity.JSONSlice) ([]string, []string) { var names = make([]string, 0, len(kbIDs)) var ids = make([]string, 0, len(kbIDs)) for _, kbID := range kbIDs { @@ -743,7 +743,7 @@ func (s *ChatService) getDatasetNamesAndIDs(kbIDs entity.JSONSlice) ([]string, [ if !ok { continue } - kb, err := s.kbDAO.GetByID(kbIDStr) + kb, err := s.kbDAO.GetByID(ctx, dao.DB, kbIDStr) if err != nil || kb == nil { continue } @@ -869,7 +869,7 @@ func (s *ChatService) updateChatREST(ctx context.Context, userID, chatID string, } if value, ok := req["dataset_ids"]; ok { - kbIDs, err := s.validateRESTDatasetIDs(value, userID) + kbIDs, err := s.validateRESTDatasetIDs(ctx, value, userID) if err != nil { return nil, err } @@ -977,7 +977,7 @@ func (s *ChatService) updateChatREST(ctx context.Context, userID, chatID string, if err != nil { return nil, errors.New("failed to retrieve updated chat") } - return s.buildRESTChatResponse(updatedChat), nil + return s.buildRESTChatResponse(ctx, updatedChat), nil } func validateRESTChatName(value interface{}, required bool) (string, bool, error) { @@ -1004,7 +1004,7 @@ func validateRESTChatName(value interface{}, required bool) (string, bool, error return name, true, nil } -func (s *ChatService) validateRESTDatasetIDs(value interface{}, userID string) (entity.JSONSlice, error) { +func (s *ChatService) validateRESTDatasetIDs(ctx context.Context, value interface{}, userID string) (entity.JSONSlice, error) { if value == nil { return entity.JSONSlice{}, nil } @@ -1020,10 +1020,10 @@ func (s *ChatService) validateRESTDatasetIDs(value interface{}, userID string) ( continue } datasetID := fmt.Sprint(item) - if !s.kbDAO.Accessible(datasetID, userID) { + if !s.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return nil, fmt.Errorf("you don't own the dataset %s", datasetID) } - kb, err := s.kbDAO.GetByID(datasetID) + kb, err := s.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil || kb == nil { return nil, fmt.Errorf("you don't own the dataset %s", datasetID) } @@ -1121,8 +1121,8 @@ func mergeJSONMap(base entity.JSONMap, patch map[string]interface{}) entity.JSON return merged } -func (s *ChatService) buildRESTChatResponse(chat *entity.Chat) map[string]interface{} { - kbNames, datasetIDs := s.getDatasetNamesAndIDs(chat.KBIDs) +func (s *ChatService) buildRESTChatResponse(ctx context.Context, chat *entity.Chat) map[string]interface{} { + kbNames, datasetIDs := s.getDatasetNamesAndIDs(ctx, chat.KBIDs) return map[string]interface{}{ "id": chat.ID, "tenant_id": chat.TenantID, @@ -1298,7 +1298,7 @@ func (s *ChatService) GetChat(ctx context.Context, userID string, chatID string) // Step 4: Build response with kb_names (same as Python _build_chat_response) // Resolve kb_ids to kb_names - kbNames, datasetIDs := s.getDatasetNamesAndIDs(chat.KBIDs) + kbNames, datasetIDs := s.getDatasetNamesAndIDs(ctx, chat.KBIDs) return &GetChatResponse{ Chat: chat, diff --git a/internal/service/chat_pipeline.go b/internal/service/chat_pipeline.go index 636696b231..ce5ed22b9c 100644 --- a/internal/service/chat_pipeline.go +++ b/internal/service/chat_pipeline.go @@ -348,7 +348,7 @@ func (s *ChatPipelineService) AsyncChat( // === Phase 6: SQL Retrieval === // Retrieve field_map for SQL retrieval (preferred over vector search) promptConfig := chat.PromptConfig - fieldMap, fmErr := s.kbDAO.GetFieldMap(kbIDStrings(kbs)) + fieldMap, fmErr := s.kbDAO.GetFieldMap(ctx, dao.DB, kbIDStrings(kbs)) if fmErr != nil { common.Warn("get_field_map failed; proceeding without field_map", zap.Error(fmErr)) fieldMap = nil @@ -552,7 +552,7 @@ func (s *ChatPipelineService) AsyncChat( var flattedMeta common.MetaData var mErr error if s.MetadataSvc != nil { - flattedMeta, mErr = s.MetadataSvc.GetFlattedMetaByKBs(kbIDs) + flattedMeta, mErr = s.MetadataSvc.GetFlattedMetaByKBs(ctx, kbIDs) } if mErr == nil { if filtered, ok := ApplyMetaDataFilter( @@ -689,7 +689,7 @@ func (s *ChatPipelineService) AsyncChat( searchQuestion := strings.Join(questions, " ") if embModel != nil { // Retrieval - rankFeature := s.MetadataSvc.LabelQuestion(searchQuestion, kbs) + rankFeature := s.MetadataSvc.LabelQuestion(ctx, searchQuestion, kbs) { tenantIDs := make([]string, 0) kbIDs := make([]string, 0) @@ -1962,7 +1962,7 @@ func (s *ChatPipelineService) getModels(ctx context.Context, chat *entity.Chat) var kbs []*entity.Knowledgebase if len(kbIDs) > 0 { var err error - kbs, err = kbDAO.GetByIDs(kbIDs) + kbs, err = kbDAO.GetByIDs(ctx, dao.DB, kbIDs) if err != nil { common.Warn("Failed to get KBs by IDs; retrieval may be incomplete", zap.Strings("kbIDs", kbIDs), zap.Error(err)) diff --git a/internal/service/chunk/chunk.go b/internal/service/chunk/chunk.go index c9d6c29e1f..1820ecec22 100644 --- a/internal/service/chunk/chunk.go +++ b/internal/service/chunk/chunk.go @@ -174,7 +174,7 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval for _, datasetID := range req.Datasets { found := false for _, tenant := range tenants { - kb, err := s.kbDAO.GetByIDAndTenantID(datasetID, tenant.TenantID) + kb, err := s.kbDAO.GetByIDAndTenantID(ctx, dao.DB, datasetID, tenant.TenantID) if err == nil && kb != nil { common.Debug("Found knowledge base in database", zap.String("datasetID", datasetID), @@ -269,7 +269,7 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval if filter != nil { // Get flattened metadata metadataSvc := service.NewMetadataService() - flattedMeta, err := metadataSvc.GetFlattedMetaByKBs([]string(req.Datasets)) + flattedMeta, err := metadataSvc.GetFlattedMetaByKBs(ctx, []string(req.Datasets)) if err != nil { common.Warn("Failed to get flatted metadata", zap.Error(err)) } else { @@ -336,7 +336,7 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval // Get tag-based rank features via LabelQuestion metadataSvc := service.NewMetadataService() - labels := metadataSvc.LabelQuestion(modifiedQuestion, kbRecords) + labels := metadataSvc.LabelQuestion(ctx, modifiedQuestion, kbRecords) common.Debug("LabelQuestion result", zap.Any("labels", labels)) // Determine embedding model. @@ -524,7 +524,7 @@ func (s *ChunkService) Get(ctx context.Context, req *service.GetChunkRequest, us var chunk map[string]interface{} for _, tenant := range tenants { // Get kbIDs for this tenant - kbIDs, err := s.kbDAO.GetKBIDsByTenantID(tenant.TenantID) + kbIDs, err := s.kbDAO.GetKBIDsByTenantID(ctx, dao.DB, tenant.TenantID) if err != nil { continue } @@ -604,7 +604,7 @@ func (s *ChunkService) cancelAllTasksOfDoc(ctx context.Context, doc *entity.Docu } func (s *ChunkService) StopParsing(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) { - if !s.kbDAO.Accessible(datasetID, userID) { + if !s.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return nil, common.CodeDataError, fmt.Errorf("you don't own the dataset %s", datasetID) } @@ -612,7 +612,7 @@ func (s *ChunkService) StopParsing(ctx context.Context, userID, datasetID string return nil, common.CodeDataError, fmt.Errorf("`document_ids` is required") } - _, err := s.kbDAO.GetByID(datasetID) + _, err := s.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil { return nil, common.CodeDataError, fmt.Errorf("you don't own the dataset %s", datasetID) } @@ -681,18 +681,18 @@ func checkDuplicateIDs(documentIDs []string, idTypes string) ([]string, []string return uniqueDocIDs, duplicateMessages } -func (s *ChunkService) accessible(datasetID, userID string) bool { +func (s *ChunkService) accessible(ctx context.Context, datasetID, userID string) bool { if s.accessibleFunc != nil { return s.accessibleFunc(datasetID, userID) } - return s.kbDAO.Accessible(datasetID, userID) + return s.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) } -func (s *ChunkService) getKnowledgebaseByID(datasetID string) (*entity.Knowledgebase, error) { +func (s *ChunkService) getKnowledgebaseByID(ctx context.Context, datasetID string) (*entity.Knowledgebase, error) { if s.getKnowledgebaseByIDFunc != nil { return s.getKnowledgebaseByIDFunc(datasetID) } - return s.kbDAO.GetByID(datasetID) + return s.kbDAO.GetByID(ctx, dao.DB, datasetID) } func (s *ChunkService) getDocumentsByIDs(ctx context.Context, docIDs []string) ([]*entity.Document, error) { @@ -703,14 +703,14 @@ func (s *ChunkService) getDocumentsByIDs(ctx context.Context, docIDs []string) ( } func (s *ChunkService) Parse(ctx context.Context, userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) { - if !s.accessible(datasetID, userID) { + if !s.accessible(ctx, datasetID, userID) { return nil, common.CodeOperatingError, fmt.Errorf("you don't own the dataset %s", datasetID) } if req == nil || len(req.DocumentIDs) == 0 { return nil, common.CodeDataError, fmt.Errorf("`document_ids` is required") } - kb, err := s.getKnowledgebaseByID(datasetID) + kb, err := s.getKnowledgebaseByID(ctx, datasetID) if err != nil || kb == nil { return nil, common.CodeDataError, fmt.Errorf("dataset not found") } @@ -806,7 +806,7 @@ func (s *ChunkService) List(ctx context.Context, req *service.ListChunksRequest, } // Get knowledge base to find tenant - kb, err := s.kbDAO.GetByID(doc.KbID) + kb, err := s.kbDAO.GetByID(ctx, dao.DB, doc.KbID) if err != nil || kb == nil { return nil, fmt.Errorf("knowledge base not found") } @@ -824,7 +824,7 @@ func (s *ChunkService) List(ctx context.Context, req *service.ListChunksRequest, } // Get kbIDs for this tenant - kbIDs, err := s.kbDAO.GetKBIDsByTenantID(targetTenantID) + kbIDs, err := s.kbDAO.GetKBIDsByTenantID(ctx, dao.DB, targetTenantID) if err != nil { return nil, fmt.Errorf("failed to get kb ids: %w", err) } @@ -1024,7 +1024,7 @@ func (s *ChunkService) SwitchChunks(ctx context.Context, userID, datasetID, docu // Find the tenant that owns this dataset var targetTenantID string for _, tenant := range tenants { - kb, err := s.kbDAO.GetByIDAndTenantID(datasetID, tenant.TenantID) + kb, err := s.kbDAO.GetByIDAndTenantID(ctx, dao.DB, datasetID, tenant.TenantID) if err == nil && kb != nil { targetTenantID = tenant.TenantID break @@ -1081,7 +1081,7 @@ func (s *ChunkService) UpdateChunk(ctx context.Context, req *service.UpdateChunk // Find the tenant that owns this dataset var targetTenantID string for _, tenant := range tenants { - kb, err := s.kbDAO.GetByIDAndTenantID(req.DatasetID, tenant.TenantID) + kb, err := s.kbDAO.GetByIDAndTenantID(ctx, dao.DB, req.DatasetID, tenant.TenantID) if err == nil && kb != nil { targetTenantID = tenant.TenantID break @@ -1227,7 +1227,7 @@ func (s *ChunkService) RemoveChunks(ctx context.Context, req *service.RemoveChun // Find the tenant that owns this document var targetTenantID string for _, tenant := range tenants { - kb, err := s.kbDAO.GetByIDAndTenantID(doc.KbID, tenant.TenantID) + kb, err := s.kbDAO.GetByIDAndTenantID(ctx, dao.DB, doc.KbID, tenant.TenantID) if err == nil && kb != nil { targetTenantID = tenant.TenantID break @@ -1280,11 +1280,11 @@ func (s *ChunkService) AddChunk(ctx context.Context, req *service.AddChunkReques if req == nil { return nil, addChunkError{code: common.CodeDataError, message: "invalid request payload"} } - if !s.accessible(req.DatasetID, userID) { + if !s.accessible(ctx, req.DatasetID, userID) { return nil, addChunkError{code: common.CodeDataError, message: fmt.Sprintf("You don't own the dataset %s.", req.DatasetID)} } - kb, err := s.getKnowledgebaseByID(req.DatasetID) + kb, err := s.getKnowledgebaseByID(ctx, req.DatasetID) if err != nil || kb == nil { return nil, addChunkError{code: common.CodeDataError, message: fmt.Sprintf("You don't own the dataset %s.", req.DatasetID)} } @@ -1371,7 +1371,7 @@ func (s *ChunkService) AddChunk(ctx context.Context, req *service.AddChunkReques chunkData["doc_type_kwd"] = "image" } - embeddingModel, err := s.getEmbeddingModel(kb.TenantID, kb.EmbdID) + embeddingModel, err := s.getEmbeddingModel(ctx, kb.TenantID, kb.EmbdID) if err != nil { return nil, addChunkError{code: common.CodeServerError, message: fmt.Sprintf("get embedding model: %v", err)} } @@ -1561,11 +1561,11 @@ func (s *ChunkService) numTokens(text string) int { return tokenizer.NumTokensFromString(text) } -func (s *ChunkService) getEmbeddingModel(tenantID, embdID string) (*models.EmbeddingModel, error) { +func (s *ChunkService) getEmbeddingModel(ctx context.Context, tenantID, embdID string) (*models.EmbeddingModel, error) { if s.getEmbeddingModelFunc != nil { return s.getEmbeddingModelFunc(tenantID, embdID) } - return service.NewModelProviderService().GetEmbeddingModel(tenantID, embdID) + return service.NewModelProviderService().GetEmbeddingModel(ctx, tenantID, embdID) } func (s *ChunkService) incrementChunkStats(docID, kbID string, tokenNum, chunkNum int64, duration float64) error { diff --git a/internal/service/chunk/chunk_test.go b/internal/service/chunk/chunk_test.go index b23b661825..02d3a01540 100644 --- a/internal/service/chunk/chunk_test.go +++ b/internal/service/chunk/chunk_test.go @@ -756,7 +756,7 @@ func TestRemoveChunksDecrementsStatsAfterDelete(t *testing.T) { t.Fatalf("document chunk_num = %d, want 4", doc.ChunkNum) } - kb, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + kb, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get kb: %v", err) } @@ -858,7 +858,7 @@ func TestDecrementChunkStatsClampsCounters(t *testing.T) { t.Fatalf("document stats token=%d chunk=%d duration=%v, want zeros", doc.TokenNum, doc.ChunkNum, doc.ProcessDuration) } - kb, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + kb, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get kb: %v", err) } @@ -1301,7 +1301,7 @@ func (e *parseTestDocEngine) GetType() string { } func (e *parseTestDocEngine) SupportsPageRank() bool { return false } -func (e *parseTestDocEngine) FilterDocIdsByMetaPushdown(context.Context, []string, []map[string]interface{}, string) []string { +func (e *parseTestDocEngine) FilterDocIdsByMetaPushdown(context.Context, *gorm.DB, []string, []map[string]interface{}, string) []string { return nil } @@ -1475,7 +1475,7 @@ func (m *switchChunksEngineMock) Ping(context.Context) error func (m *switchChunksEngineMock) Close() error { return nil } func (m *switchChunksEngineMock) GetType() string { return "elasticsearch" } func (m *switchChunksEngineMock) SupportsPageRank() bool { return true } -func (m *switchChunksEngineMock) FilterDocIdsByMetaPushdown(context.Context, []string, []map[string]interface{}, string) []string { +func (m *switchChunksEngineMock) FilterDocIdsByMetaPushdown(context.Context, *gorm.DB, []string, []map[string]interface{}, string) []string { return nil } diff --git a/internal/service/chunk_types.go b/internal/service/chunk_types.go index bd112ad36c..6229a26f16 100644 --- a/internal/service/chunk_types.go +++ b/internal/service/chunk_types.go @@ -147,7 +147,7 @@ func (s *ChunkService) Get(req *GetChunkRequest, userID string) (*GetChunkRespon var chunk map[string]interface{} for _, tenant := range tenants { // Get kbIDs for this tenant - kbIDs, err := s.kbDAO.GetKBIDsByTenantID(tenant.TenantID) + kbIDs, err := s.kbDAO.GetKBIDsByTenantID(ctx, dao.DB, tenant.TenantID) if err != nil { continue } @@ -254,7 +254,7 @@ func IndexName(uid string) string { } func (s *ChunkService) StopParsing(ctx context.Context, userID, datasetID string, req StopParsingRequest) (map[string]interface{}, common.ErrorCode, error) { - if !s.kbDAO.Accessible(datasetID, userID) { + if !s.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return nil, common.CodeAuthenticationError, fmt.Errorf("You don't own the dataset %s", datasetID) } @@ -264,7 +264,7 @@ func (s *ChunkService) StopParsing(ctx context.Context, userID, datasetID string docList, duplicateMessages := CheckDuplicateIDs(req.DocumentIDs, "document") - kb, err := s.kbDAO.GetByID(datasetID) + kb, err := s.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil { return nil, common.CodeDataError, fmt.Errorf("You don't own the dataset %s", datasetID) } @@ -373,7 +373,7 @@ func (s *ChunkService) List(ctx context.Context, req *ListChunksRequest, userID } // Get knowledge base to find tenant - kb, err := s.kbDAO.GetByID(doc.KbID) + kb, err := s.kbDAO.GetByID(ctx, dao.DB, doc.KbID) if err != nil || kb == nil { return nil, fmt.Errorf("knowledge base not found") } @@ -391,7 +391,7 @@ func (s *ChunkService) List(ctx context.Context, req *ListChunksRequest, userID } // Get kbIDs for this tenant - kbIDs, err := s.kbDAO.GetKBIDsByTenantID(targetTenantID) + kbIDs, err := s.kbDAO.GetKBIDsByTenantID(ctx, dao.DB, targetTenantID) if err != nil { return nil, fmt.Errorf("failed to get kb ids: %w", err) } @@ -536,7 +536,7 @@ func (s *ChunkService) UpdateChunk(ctx context.Context, req *UpdateChunkRequest, // Find the tenant that owns this dataset var targetTenantID string for _, tenant := range tenants { - kb, err := s.kbDAO.GetByIDAndTenantID(req.DatasetID, tenant.TenantID) + kb, err := s.kbDAO.GetByIDAndTenantID(ctx, dao.DB, req.DatasetID, tenant.TenantID) if err == nil && kb != nil { targetTenantID = tenant.TenantID break @@ -744,7 +744,7 @@ func (s *ChunkService) RemoveChunks(ctx context.Context, req *RemoveChunksReques // Find the tenant that owns this document var targetTenantID string for _, tenant := range tenants { - kb, err := s.kbDAO.GetByIDAndTenantID(doc.KbID, tenant.TenantID) + kb, err := s.kbDAO.GetByIDAndTenantID(ctx, dao.DB, doc.KbID, tenant.TenantID) if err == nil && kb != nil { targetTenantID = tenant.TenantID break diff --git a/internal/service/dataset/create_test.go b/internal/service/dataset/create_test.go index 897e121785..947ab1ec84 100644 --- a/internal/service/dataset/create_test.go +++ b/internal/service/dataset/create_test.go @@ -42,10 +42,11 @@ func TestCreateDataset_NoComponentParams(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) insertCreateDatasetTenant(t, "tenant-1") + ctx := t.Context() chunkMethod := "naive" parseType := 1 - result, code, err := testDatasetCreateService(t).CreateDataset(&service.CreateDatasetRequest{ + result, code, err := testDatasetCreateService(t).CreateDataset(ctx, &service.CreateDatasetRequest{ Name: "ds-no-cp", ParserID: &chunkMethod, ParseType: &parseType, @@ -65,10 +66,11 @@ func TestCreateDataset_ComponentParamsPopulated(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) insertCreateDatasetTenant(t, "tenant-1") + ctx := t.Context() chunkMethod := "general" parseType := 1 - result, code, err := testDatasetCreateService(t).CreateDataset(&service.CreateDatasetRequest{ + result, code, err := testDatasetCreateService(t).CreateDataset(ctx, &service.CreateDatasetRequest{ Name: "ds-with-cp", ParserID: &chunkMethod, ParseType: &parseType, @@ -89,11 +91,12 @@ func TestCreateDataset_ParseTypeBuiltinClearsPipelineID(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) insertCreateDatasetTenant(t, "tenant-1") + ctx := t.Context() pipelineID := "0123456789abcdef0123456789abcdef" parseTypeBuiltin := 1 chunkMethod := "naive" - result, code, err := testDatasetCreateService(t).CreateDataset(&service.CreateDatasetRequest{ + result, code, err := testDatasetCreateService(t).CreateDataset(ctx, &service.CreateDatasetRequest{ Name: "ds-parse-builtin", ParserID: &chunkMethod, PipelineID: &pipelineID, @@ -118,11 +121,12 @@ func TestCreateDataset_ParseTypePipelineIgnoresParserID(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) insertCreateDatasetTenant(t, "tenant-1") + ctx := t.Context() pipelineID := "0123456789abcdef0123456789abcdef" parseTypePipeline := 2 chunkMethod := "naive" - result, code, err := testDatasetCreateService(t).CreateDataset(&service.CreateDatasetRequest{ + result, code, err := testDatasetCreateService(t).CreateDataset(ctx, &service.CreateDatasetRequest{ Name: "ds-parse-pipeline", ParserID: &chunkMethod, PipelineID: &pipelineID, @@ -144,8 +148,9 @@ func TestCreateDataset_ValidatesName(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) insertCreateDatasetTenant(t, "tenant-1") + ctx := t.Context() - _, code, err := testDatasetCreateService(t).CreateDataset(&service.CreateDatasetRequest{Name: " "}, "tenant-1") + _, code, err := testDatasetCreateService(t).CreateDataset(ctx, &service.CreateDatasetRequest{Name: " "}, "tenant-1") if err == nil { t.Fatal("expected name validation error") } @@ -173,7 +178,8 @@ func TestCreateDataset_RejectsDuplicateName(t *testing.T) { t.Fatalf("failed to create existing kb: %v", err) } - _, code, err := testDatasetCreateService(t).CreateDataset(&service.CreateDatasetRequest{Name: "Existing"}, "tenant-1") + ctx := t.Context() + _, code, err := testDatasetCreateService(t).CreateDataset(ctx, &service.CreateDatasetRequest{Name: "Existing"}, "tenant-1") if err == nil { t.Fatal("expected duplicate name error") } @@ -200,13 +206,14 @@ func TestCreateDataset_RejectsInvalidEmbeddingModel(t *testing.T) { {"whitespace_provider", "BAAI/bge-small-en-v1.5@ ", "Both model_name and provider must be non-empty strings"}, } + ctx := t.Context() for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) insertCreateDatasetTenant(t, "tenant-1") - _, code, err := testDatasetCreateService(t).CreateDataset(&service.CreateDatasetRequest{ + _, code, err := testDatasetCreateService(t).CreateDataset(ctx, &service.CreateDatasetRequest{ Name: "ds-embd-" + tc.name, EmbeddingModel: &tc.embeddingModel, }, "tenant-1") diff --git a/internal/service/dataset/crud.go b/internal/service/dataset/crud.go index 4347f51839..44d1fc353e 100644 --- a/internal/service/dataset/crud.go +++ b/internal/service/dataset/crud.go @@ -16,7 +16,7 @@ import ( "gorm.io/gorm" ) -func (d *DatasetService) CreateDataset(req *service.CreateDatasetRequest, tenantID string) (map[string]interface{}, common.ErrorCode, error) { +func (d *DatasetService) CreateDataset(ctx context.Context, req *service.CreateDatasetRequest, tenantID string) (map[string]interface{}, common.ErrorCode, error) { if !common.IsValidString(req.Name) { return nil, common.CodeDataError, errors.New("dataset name must be string") } @@ -118,7 +118,7 @@ func (d *DatasetService) CreateDataset(req *service.CreateDatasetRequest, tenant kbID := utility.GenerateToken() status := string(entity.StatusValid) // Reject duplicate name within tenant to match the established API contract. - existing, err := d.kbDAO.GetByName(name, tenantID) + existing, err := d.kbDAO.GetByName(ctx, dao.DB, name, tenantID) if err != nil && !dao.IsNotFoundErr(err) { return nil, common.CodeServerError, errors.New("database operation failed") } @@ -140,14 +140,14 @@ func (d *DatasetService) CreateDataset(req *service.CreateDatasetRequest, tenant Status: &status, } - if err = d.kbDAO.Create(kb); err != nil { + if err = d.kbDAO.Create(ctx, dao.DB, kb); err != nil { if dao.IsDuplicateKeyErr(err) { return nil, common.CodeDataError, fmt.Errorf("dataset name '%s' already exists", name) } return nil, common.CodeServerError, errors.New("failed to save dataset") } - createdKB, err := d.kbDAO.GetByID(kbID) + createdKB, err := d.kbDAO.GetByID(ctx, dao.DB, kbID) if err != nil || createdKB == nil { return nil, common.CodeServerError, errors.New("dataset created failed") } @@ -167,11 +167,11 @@ func (d *DatasetService) GetDataset(ctx context.Context, datasetID, userID strin } datasetID = normalizedID - if !d.kbDAO.Accessible(datasetID, userID) { + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return nil, common.CodeDataError, fmt.Errorf("user '%s' lacks permission for dataset '%s'", userID, datasetID) } - kb, err := d.kbDAO.GetByID(datasetID) + kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil || kb == nil { return nil, common.CodeDataError, errors.New("invalid Dataset ID") } @@ -193,7 +193,7 @@ func (d *DatasetService) GetDataset(ctx context.Context, datasetID, userID strin return data, common.CodeSuccess, nil } -func (d *DatasetService) DeleteDatasets(ids []string, deleteAll bool, tenantID string) (map[string]interface{}, common.ErrorCode, error) { +func (d *DatasetService) DeleteDatasets(ctx context.Context, ids []string, deleteAll bool, tenantID string) (map[string]interface{}, common.ErrorCode, error) { normalizedIDs := make([]string, 0, len(ids)) seenIDs := make(map[string]struct{}, len(ids)) for _, id := range ids { @@ -213,7 +213,7 @@ func (d *DatasetService) DeleteDatasets(ids []string, deleteAll bool, tenantID s if !deleteAll { return map[string]interface{}{"deleted": []string{}}, common.CodeSuccess, nil } - kbs, err := d.kbDAO.Query(map[string]interface{}{"tenant_id": tenantID}) + kbs, err := d.kbDAO.Query(ctx, dao.DB, map[string]interface{}{"tenant_id": tenantID}) if err != nil { return nil, common.CodeServerError, errors.New("database operation failed") } @@ -226,7 +226,7 @@ func (d *DatasetService) DeleteDatasets(ids []string, deleteAll bool, tenantID s kbs := make([]*entity.Knowledgebase, 0, len(normalizedIDs)) unauthorizedIDs := make([]string, 0) for _, id := range normalizedIDs { - kb, err := d.kbDAO.GetByIDAndTenantID(id, tenantID) + kb, err := d.kbDAO.GetByIDAndTenantID(ctx, dao.DB, id, tenantID) if err != nil || kb == nil { unauthorizedIDs = append(unauthorizedIDs, id) continue @@ -313,7 +313,7 @@ func (d *DatasetService) deleteDataset(tenantID string, kb *entity.Knowledgebase }) } -func (d *DatasetService) ListDatasets(id, name string, page, pageSize int, orderby string, desc bool, keywords string, ownerIDs []string, parserID, userID string) ([]map[string]interface{}, int64, common.ErrorCode, error) { +func (d *DatasetService) ListDatasets(ctx context.Context, id, name string, page, pageSize int, orderby string, desc bool, keywords string, ownerIDs []string, parserID, userID string) ([]map[string]interface{}, int64, common.ErrorCode, error) { id = strings.TrimSpace(id) if id != "" { normalizedID, err := normalizeDatasetID(id) @@ -322,7 +322,7 @@ func (d *DatasetService) ListDatasets(id, name string, page, pageSize int, order } id = normalizedID - kbs, err := d.kbDAO.GetKBByIDAndUserID(id, userID) + kbs, err := d.kbDAO.GetKBByIDAndUserID(ctx, dao.DB, id, userID) if err != nil { return nil, 0, common.CodeServerError, errors.New("database operation failed") } @@ -333,7 +333,7 @@ func (d *DatasetService) ListDatasets(id, name string, page, pageSize int, order name = strings.TrimSpace(name) if name != "" { - kbs, err := d.kbDAO.GetKBByNameAndUserID(name, userID) + kbs, err := d.kbDAO.GetKBByNameAndUserID(ctx, dao.DB, name, userID) if err != nil { return nil, 0, common.CodeServerError, errors.New("database operation failed") } @@ -377,7 +377,7 @@ func (d *DatasetService) ListDatasets(id, name string, page, pageSize int, order } } - kbs, total, err := d.kbDAO.GetByTenantIDs(tenantIDs, userID, page, pageSize, orderby, desc, keywords, parserID, id, name) + kbs, total, err := d.kbDAO.GetByTenantIDs(ctx, dao.DB, tenantIDs, userID, page, pageSize, orderby, desc, keywords, parserID, id, name) if err != nil { return nil, 0, common.CodeServerError, errors.New("database operation failed") } diff --git a/internal/service/dataset/fake_doc_engine_test.go b/internal/service/dataset/fake_doc_engine_test.go index 0358cc901e..7fa80346a8 100644 --- a/internal/service/dataset/fake_doc_engine_test.go +++ b/internal/service/dataset/fake_doc_engine_test.go @@ -4,6 +4,8 @@ import ( "context" "ragflow/internal/engine/types" + + "gorm.io/gorm" ) // fakeChatDocEngine is a no-op DocEngine implementation for tests. @@ -77,6 +79,6 @@ func (fakeChatDocEngine) Ping(context.Context) error { func (fakeChatDocEngine) Close() error { return nil } func (fakeChatDocEngine) GetType() string { return "fake" } func (fakeChatDocEngine) SupportsPageRank() bool { return false } -func (fakeChatDocEngine) FilterDocIdsByMetaPushdown(context.Context, []string, []map[string]interface{}, string) []string { +func (fakeChatDocEngine) FilterDocIdsByMetaPushdown(context.Context, *gorm.DB, []string, []map[string]interface{}, string) []string { return nil } diff --git a/internal/service/dataset/index.go b/internal/service/dataset/index.go index 47ce1aca86..3f7ac1a9de 100644 --- a/internal/service/dataset/index.go +++ b/internal/service/dataset/index.go @@ -270,11 +270,11 @@ func (d *DatasetService) RunIndex(ctx context.Context, userID, datasetID, indexT if datasetID == "" { return nil, common.CodeDataError, errors.New(`lack of "Dataset ID"`) } - if !d.kbDAO.Accessible(datasetID, userID) { + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return nil, common.CodeDataError, errors.New("no authorization") } - kb, err := d.kbDAO.GetByID(datasetID) + kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, errors.New("invalid Dataset ID") @@ -306,7 +306,7 @@ func (d *DatasetService) RunIndex(ctx context.Context, userID, datasetID, indexT var updatedDocument *entity.Document var dataErr error - err = dao.DB.Transaction(func(tx *gorm.DB) error { + err = dao.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { var lockedKB entity.Knowledgebase if err = tx.Clauses(clause.Locking{Strength: "UPDATE"}). Where("id = ? AND status = ?", kb.ID, string(entity.StatusValid)). @@ -366,7 +366,7 @@ func (d *DatasetService) getDocumentsByDatasetForIndex(ctx context.Context, data return documents, common.CodeSuccess, nil } -func (d *DatasetService) TraceIndex(datasetID, userID, indexType string) (*entity.Task, common.ErrorCode, error) { +func (d *DatasetService) TraceIndex(ctx context.Context, datasetID, userID, indexType string) (*entity.Task, common.ErrorCode, error) { if !checkType(indexType) { return nil, common.CodeDataError, fmt.Errorf("Invalid index type '%s'. Must be one of %v", indexType, validIndexTypes) } @@ -374,11 +374,11 @@ func (d *DatasetService) TraceIndex(datasetID, userID, indexType string) (*entit if datasetID == "" { return nil, common.CodeDataError, errors.New(`Lack of "Dataset ID"`) } - if !d.kbDAO.Accessible(datasetID, userID) { + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return nil, common.CodeDataError, errors.New("No authorization.") } - kb, err := d.kbDAO.GetByID(datasetID) + kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, errors.New("Invalid Dataset ID") @@ -423,11 +423,11 @@ func (d *DatasetService) CheckEmbedding(ctx context.Context, userID, datasetID s if datasetID == "" { return nil, common.CodeDataError, errors.New(`Lack of "Dataset ID"`) } - if !d.kbDAO.Accessible(datasetID, userID) { + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return nil, common.CodeDataError, errors.New("No authorization.") } - kb, err := d.kbDAO.GetByID(datasetID) + kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, errors.New("Invalid Dataset ID") @@ -460,7 +460,7 @@ func (d *DatasetService) CheckEmbedding(ctx context.Context, userID, datasetID s checkNum = defaultEmbeddingCheckNum } - samples, err := d.sampleRandomChunksWithVectors(context.Background(), kb.TenantID, datasetID, checkNum) + samples, err := d.sampleRandomChunksWithVectors(ctx, kb.TenantID, datasetID, checkNum) if err != nil { return nil, common.CodeServerError, err } @@ -479,7 +479,7 @@ func (d *DatasetService) CheckEmbedding(ctx context.Context, userID, datasetID s continue } - rawChunk, err := d.docEngine.GetChunk(context.Background(), fmt.Sprintf("ragflow_%s", kb.TenantID), sample.ChunkID, []string{datasetID}) + rawChunk, err := d.docEngine.GetChunk(ctx, fmt.Sprintf("ragflow_%s", kb.TenantID), sample.ChunkID, []string{datasetID}) if err != nil { continue } @@ -662,7 +662,7 @@ func (d *DatasetService) verifyEmbeddingAvailability(embdID string, tenantID str return true, "" } -func (d *DatasetService) DeleteIndex(userID, datasetID, indexType string, wipe bool) (common.ErrorCode, error) { +func (d *DatasetService) DeleteIndex(ctx context.Context, userID, datasetID, indexType string, wipe bool) (common.ErrorCode, error) { if !checkType(indexType) { return common.CodeArgumentError, fmt.Errorf("Invalid index type '%s'", indexType) } @@ -671,11 +671,11 @@ func (d *DatasetService) DeleteIndex(userID, datasetID, indexType string, wipe b return common.CodeDataError, errors.New(`Lack of "Dataset ID"`) } - if !d.kbDAO.Accessible(datasetID, userID) { + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return common.CodeDataError, errors.New("No authorization.") } - kb, err := d.kbDAO.GetByID(datasetID) + kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil { if dao.IsNotFoundErr(err) { return common.CodeDataError, errors.New("Invalid Dataset ID") @@ -704,7 +704,7 @@ func (d *DatasetService) DeleteIndex(userID, datasetID, indexType string, wipe b return common.CodeServerError, errors.New("Document engine is not initialized") } indexName := fmt.Sprintf("ragflow_%s", kb.TenantID) - _, err = d.docEngine.DeleteChunks(context.Background(), map[string]interface{}{ + _, err = d.docEngine.DeleteChunks(ctx, map[string]interface{}{ "knowledge_graph_kwd": []interface{}{"graph", "subgraph", "entity", "relation", "community_report"}, "kb_id": datasetID, }, indexName, datasetID) @@ -719,7 +719,7 @@ func (d *DatasetService) DeleteIndex(userID, datasetID, indexType string, wipe b return common.CodeServerError, errors.New("Document engine is not initialized") } indexName := fmt.Sprintf("ragflow_%s", kb.TenantID) - _, err = d.docEngine.DeleteChunks(context.Background(), map[string]interface{}{ + _, err = d.docEngine.DeleteChunks(ctx, map[string]interface{}{ "raptor_kwd": []interface{}{"raptor"}, "kb_id": datasetID, }, indexName, datasetID) @@ -734,7 +734,7 @@ func (d *DatasetService) DeleteIndex(userID, datasetID, indexType string, wipe b updates[taskFinishAtField] = nil } if len(updates) > 0 { - if err := d.kbDAO.UpdateByID(kb.ID, updates); err != nil { + if err = d.kbDAO.UpdateByID(ctx, dao.DB, kb.ID, updates); err != nil { common.Warn("Failed to clear KB index task refs", zap.String("dataset_id", datasetID), zap.Error(err)) } } diff --git a/internal/service/dataset/index_delete_test.go b/internal/service/dataset/index_delete_test.go index e63b1b36f0..a4bb1cbdae 100644 --- a/internal/service/dataset/index_delete_test.go +++ b/internal/service/dataset/index_delete_test.go @@ -99,9 +99,10 @@ func TestDatasetServiceDeleteIndexGraphWipeFalseOnlyCancelsTask(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) insertDeleteIndexKB(t, "graph", "graph-task") + ctx := t.Context() docEngine := &deleteIndexDocEngine{} - code, err := testDatasetServiceForDeleteIndex(docEngine).DeleteIndex("user-1", "kb-1", "graph", false) + code, err := testDatasetServiceForDeleteIndex(docEngine).DeleteIndex(ctx, "user-1", "kb-1", "graph", false) if err != nil { t.Fatalf("DeleteIndex failed: %v", err) } @@ -126,9 +127,10 @@ func TestDatasetServiceDeleteIndexGraphWipeTrueDeletesArtefacts(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) insertDeleteIndexKB(t, "graph", "graph-task") + ctx := t.Context() docEngine := &deleteIndexDocEngine{} - code, err := testDatasetServiceForDeleteIndex(docEngine).DeleteIndex("user-1", "kb-1", "graph", true) + code, err := testDatasetServiceForDeleteIndex(docEngine).DeleteIndex(ctx, "user-1", "kb-1", "graph", true) if err != nil { t.Fatalf("DeleteIndex failed: %v", err) } @@ -154,9 +156,10 @@ func TestDatasetServiceDeleteIndexRaptorWipeTrueDeletesRaptorArtefacts(t *testin db := setupServiceTestDB(t) pushServiceDB(t, db) insertDeleteIndexKB(t, "raptor", "raptor-task") + ctx := t.Context() docEngine := &deleteIndexDocEngine{} - code, err := testDatasetServiceForDeleteIndex(docEngine).DeleteIndex("user-1", "kb-1", "raptor", true) + code, err := testDatasetServiceForDeleteIndex(docEngine).DeleteIndex(ctx, "user-1", "kb-1", "raptor", true) if err != nil { t.Fatalf("DeleteIndex failed: %v", err) } @@ -186,9 +189,10 @@ func TestDatasetServiceDeleteIndexMindmapDoesNotDeleteDocStore(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) insertDeleteIndexKB(t, "mindmap", "mindmap-task") + ctx := t.Context() docEngine := &deleteIndexDocEngine{} - code, err := testDatasetServiceForDeleteIndex(docEngine).DeleteIndex("user-1", "kb-1", "mindmap", true) + code, err := testDatasetServiceForDeleteIndex(docEngine).DeleteIndex(ctx, "user-1", "kb-1", "mindmap", true) if err != nil { t.Fatalf("DeleteIndex failed: %v", err) } @@ -212,8 +216,9 @@ func TestDatasetServiceDeleteIndexMindmapDoesNotDeleteDocStore(t *testing.T) { func TestDatasetServiceDeleteIndexRejectsInvalidType(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) + ctx := t.Context() - code, err := testDatasetServiceForDeleteIndex(&deleteIndexDocEngine{}).DeleteIndex("user-1", "kb-1", "invalid", true) + code, err := testDatasetServiceForDeleteIndex(&deleteIndexDocEngine{}).DeleteIndex(ctx, "user-1", "kb-1", "invalid", true) if err == nil { t.Fatal("expected invalid index type error") } diff --git a/internal/service/dataset/ingestion.go b/internal/service/dataset/ingestion.go index 23a4cb3e03..792095cf6b 100644 --- a/internal/service/dataset/ingestion.go +++ b/internal/service/dataset/ingestion.go @@ -14,11 +14,11 @@ func (d *DatasetService) GetIngestionSummary(ctx context.Context, datasetID, use if datasetID == "" { return nil, common.CodeDataError, errors.New(`Lack of "Dataset ID"`) } - if !d.kbDAO.Accessible(datasetID, userID) { + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return nil, common.CodeDataError, errors.New("No authorization.") } - kb, err := d.kbDAO.GetByID(datasetID) + kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, fmt.Errorf("Invalid Dataset ID '%s'", datasetID) @@ -39,11 +39,11 @@ func (d *DatasetService) GetIngestionSummary(ctx context.Context, datasetID, use }, common.CodeSuccess, nil } -func (d *DatasetService) ListIngestionLogs(datasetID, userID string, page, pageSize int, orderby string, desc bool, operationStatus []string, createDateFrom, createDateTo, logType, keywords string) (map[string]interface{}, common.ErrorCode, error) { +func (d *DatasetService) ListIngestionLogs(ctx context.Context, datasetID, userID string, page, pageSize int, orderby string, desc bool, operationStatus []string, createDateFrom, createDateTo, logType, keywords string) (map[string]interface{}, common.ErrorCode, error) { if datasetID == "" { return nil, common.CodeDataError, errors.New(`Lack of "Dataset ID"`) } - if !d.kbDAO.Accessible(datasetID, userID) { + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return nil, common.CodeDataError, errors.New("No authorization.") } @@ -89,14 +89,14 @@ func (d *DatasetService) ListIngestionLogs(datasetID, userID string, page, pageS }, common.CodeSuccess, nil } -func (d *DatasetService) GetIngestionLog(datasetID, userID, logID string) (map[string]interface{}, common.ErrorCode, error) { +func (d *DatasetService) GetIngestionLog(ctx context.Context, datasetID, userID, logID string) (map[string]interface{}, common.ErrorCode, error) { if datasetID == "" { return nil, common.CodeDataError, errors.New(`Lack of "Dataset ID"`) } if logID == "" { return nil, common.CodeDataError, errors.New(`Lack of "Log ID"`) } - if !d.kbDAO.Accessible(datasetID, userID) { + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return nil, common.CodeDataError, errors.New("No authorization.") } diff --git a/internal/service/dataset/metadata.go b/internal/service/dataset/metadata.go index a18b7c68c7..88e0bbd351 100644 --- a/internal/service/dataset/metadata.go +++ b/internal/service/dataset/metadata.go @@ -16,7 +16,7 @@ import ( func (d *DatasetService) UpdateDocumentMetadataConfig(ctx context.Context, userID, datasetID, documentID string, req map[string]interface{}) (*entity.Document, common.ErrorCode, error) { userID = strings.TrimSpace(userID) datasetID = strings.TrimSpace(datasetID) - if !d.Accessible(datasetID, userID) { + if !d.Accessible(ctx, datasetID, userID) { return nil, common.CodeDataError, errors.New("You don't own the dataset.") } @@ -51,14 +51,14 @@ func (d *DatasetService) UpdateDocumentMetadataConfig(ctx context.Context, userI } // GetMetadataConfig gets the auto-metadata configuration for a dataset. -func (d *DatasetService) GetMetadataConfig(datasetID, tenantID string) (map[string]interface{}, common.ErrorCode, error) { +func (d *DatasetService) GetMetadataConfig(ctx context.Context, datasetID, tenantID string) (map[string]interface{}, common.ErrorCode, error) { datasetID = strings.TrimSpace(datasetID) tenantID = strings.TrimSpace(tenantID) - if !d.Accessible(datasetID, tenantID) { + if !d.Accessible(ctx, datasetID, tenantID) { return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID) } - kb, err := d.kbDAO.GetByID(datasetID) + kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, errors.New("Dataset not found") @@ -76,15 +76,15 @@ func (d *DatasetService) GetMetadataConfig(datasetID, tenantID string) (map[stri } // UpdateMetadataConfig updates the auto-metadata configuration for a dataset. -func (d *DatasetService) UpdateMetadataConfig(datasetID, tenantID string, req *service.MetadataConfigRequest) (map[string]interface{}, common.ErrorCode, error) { +func (d *DatasetService) UpdateMetadataConfig(ctx context.Context, datasetID, tenantID string, req *service.MetadataConfigRequest) (map[string]interface{}, common.ErrorCode, error) { datasetID = strings.TrimSpace(datasetID) tenantID = strings.TrimSpace(tenantID) - if !d.Accessible(datasetID, tenantID) { + if !d.Accessible(ctx, datasetID, tenantID) { return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID) } - kb, err := d.kbDAO.GetByID(datasetID) + kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, errors.New("Dataset not found") @@ -115,7 +115,7 @@ func (d *DatasetService) UpdateMetadataConfig(datasetID, tenantID string, req *s parserConfig["metadata"] = metadata parserConfig["built_in_metadata"] = builtInMetadata - if err = d.kbDAO.UpdateByID(kb.ID, map[string]interface{}{"parser_config": parserConfig}); err != nil { + if err = d.kbDAO.UpdateByID(ctx, dao.DB, kb.ID, map[string]interface{}{"parser_config": parserConfig}); err != nil { return nil, common.CodeServerError, errors.New("Update auto-metadata error.(Database error)") } diff --git a/internal/service/dataset/permission.go b/internal/service/dataset/permission.go index 55dab64063..31b920f0a0 100644 --- a/internal/service/dataset/permission.go +++ b/internal/service/dataset/permission.go @@ -1,25 +1,27 @@ package dataset import ( + "context" "errors" + "ragflow/internal/dao" "strings" "ragflow/internal/entity" ) // Accessible checks if a user has access to a dataset. -func (d *DatasetService) Accessible(kbID, userID string) bool { - return d.kbDAO.Accessible(kbID, userID) +func (d *DatasetService) Accessible(ctx context.Context, kbID, userID string) bool { + return d.kbDAO.Accessible(ctx, dao.DB, kbID, userID) } // GetByID retrieves a knowledge base by ID. -func (d *DatasetService) GetByID(kbID string) (*entity.Knowledgebase, error) { - return d.kbDAO.GetByID(kbID) +func (d *DatasetService) GetByID(ctx context.Context, kbID string) (*entity.Knowledgebase, error) { + return d.kbDAO.GetByID(ctx, dao.DB, kbID) } // GetKnowledgebaseByID resolves a dataset entity without applying permission // checks. Upload needs the same existence-then-auth ordering as Python. -func (d *DatasetService) GetKnowledgebaseByID(datasetID string) (*entity.Knowledgebase, error) { +func (d *DatasetService) GetKnowledgebaseByID(ctx context.Context, datasetID string) (*entity.Knowledgebase, error) { datasetID = strings.TrimSpace(datasetID) if datasetID == "" { return nil, errors.New("Lack of \"Dataset ID\"") @@ -28,7 +30,7 @@ func (d *DatasetService) GetKnowledgebaseByID(datasetID string) (*entity.Knowled if err != nil { return nil, err } - return d.kbDAO.GetByID(normalizedID) + return d.kbDAO.GetByID(ctx, dao.DB, normalizedID) } // CheckKBTeamPermission checks if a user has team-level permission for the KB. @@ -55,6 +57,6 @@ func (d *DatasetService) CheckKBTeamPermission(kb *entity.Knowledgebase, userID } // GetFieldMap returns the field map for the given knowledge base IDs. -func (d *DatasetService) GetFieldMap(ids []string) (map[string]interface{}, error) { - return d.kbDAO.GetFieldMap(ids) +func (d *DatasetService) GetFieldMap(ctx context.Context, ids []string) (map[string]interface{}, error) { + return d.kbDAO.GetFieldMap(ctx, dao.DB, ids) } diff --git a/internal/service/dataset/search.go b/internal/service/dataset/search.go index 5811b8a58e..0e1c3d361e 100644 --- a/internal/service/dataset/search.go +++ b/internal/service/dataset/search.go @@ -3,6 +3,7 @@ package dataset import ( "context" "fmt" + "ragflow/internal/dao" "go.uber.org/zap" @@ -13,14 +14,14 @@ import ( "ragflow/internal/service/nlp" ) -func (d *DatasetService) SearchDataset(datasetID, userID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error) { +func (d *DatasetService) SearchDataset(ctx context.Context, datasetID, userID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error) { if datasetID == "" { return nil, fmt.Errorf("dataset_id is required") } - return d.SearchDatasets(req.ToSearchDatasetsRequest(datasetID), userID) + return d.SearchDatasets(ctx, req.ToSearchDatasetsRequest(datasetID), userID) } -func (d *DatasetService) SearchDatasets(req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error) { +func (d *DatasetService) SearchDatasets(ctx context.Context, req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error) { if req.Question == "" { return nil, fmt.Errorf("question is required") } @@ -76,7 +77,6 @@ func (d *DatasetService) SearchDatasets(req *service.SearchDatasetsRequest, user metadataFilter := req.MetadataFilter crossLanguages := req.CrossLanguages - ctx := context.Background() modelProviderSvc := service.NewModelProviderService() // Access check for all datasets @@ -84,12 +84,12 @@ func (d *DatasetService) SearchDatasets(req *service.SearchDatasetsRequest, user var kbRecords []*entity.Knowledgebase seenTenants := make(map[string]bool) for _, datasetID := range datasetIDs { - if !d.kbDAO.Accessible(datasetID, userID) { + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { common.Warn("SearchDatasets access denied", zap.String("datasetID", datasetID), zap.String("userID", userID)) return nil, fmt.Errorf("only owner of dataset %s is authorized for this operation", datasetID) } - kb, err := d.kbDAO.GetByID(datasetID) + kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil || kb == nil { common.Warn("SearchDatasets dataset not found", zap.String("datasetID", datasetID)) return nil, fmt.Errorf("dataset %s not found", datasetID) @@ -189,7 +189,7 @@ func (d *DatasetService) SearchDatasets(req *service.SearchDatasetsRequest, user copy(docIDs, req.DocIDs) if len(metadataFilter) > 0 { metadataSvc := service.NewMetadataService() - flattedMeta, err := metadataSvc.GetFlattedMetaByKBs(datasetIDs) + flattedMeta, err := metadataSvc.GetFlattedMetaByKBs(ctx, datasetIDs) if err != nil { common.Warn("Failed to get flatted metadata, using empty metadata for filter", zap.Error(err)) flattedMeta = make(common.MetaData) @@ -225,7 +225,7 @@ func (d *DatasetService) SearchDatasets(req *service.SearchDatasetsRequest, user // Get tag-based rank features via LabelQuestion metadataSvc := service.NewMetadataService() - labels := metadataSvc.LabelQuestion(modifiedQuestion, kbRecords) + labels := metadataSvc.LabelQuestion(ctx, modifiedQuestion, kbRecords) // Determine embedding model var embeddingModel *modelModule.EmbeddingModel diff --git a/internal/service/dataset/tags.go b/internal/service/dataset/tags.go index 2caee6dc73..77bb70487f 100644 --- a/internal/service/dataset/tags.go +++ b/internal/service/dataset/tags.go @@ -13,7 +13,7 @@ import ( enginetypes "ragflow/internal/engine/types" ) -func (d *DatasetService) AggregateTags(datasetIDs []string, userID string) ([]map[string]interface{}, common.ErrorCode, error) { +func (d *DatasetService) AggregateTags(ctx context.Context, datasetIDs []string, userID string) ([]map[string]interface{}, common.ErrorCode, error) { if len(datasetIDs) == 0 { return nil, common.CodeDataError, errors.New("Lack of dataset_ids in query parameters") } @@ -31,10 +31,10 @@ func (d *DatasetService) AggregateTags(datasetIDs []string, userID string) ([]ma if err != nil { return nil, common.CodeDataError, err } - if !d.kbDAO.Accessible(datasetID, userID) { + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return nil, common.CodeDataError, fmt.Errorf("No authorization for dataset '%s'", datasetID) } - kb, err := d.kbDAO.GetByID(datasetID) + kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, fmt.Errorf("Invalid Dataset ID '%s'", datasetID) @@ -96,7 +96,7 @@ func (d *DatasetService) AggregateTags(datasetIDs []string, userID string) ([]ma return result, common.CodeSuccess, nil } -func (d *DatasetService) ListTags(datasetID, userID string) ([]map[string]interface{}, common.ErrorCode, error) { +func (d *DatasetService) ListTags(ctx context.Context, datasetID, userID string) ([]map[string]interface{}, common.ErrorCode, error) { datasetID = strings.TrimSpace(datasetID) if datasetID == "" { return nil, common.CodeDataError, errors.New("Lack of \"Dataset ID\"") @@ -106,13 +106,13 @@ func (d *DatasetService) ListTags(datasetID, userID string) ([]map[string]interf return nil, common.CodeDataError, err } datasetID = normalizedID - if !d.kbDAO.Accessible(datasetID, userID) { + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return nil, common.CodeDataError, errors.New("No authorization.") } if d.docEngine == nil { return nil, common.CodeServerError, errors.New("Document engine is not initialized") } - kb, err := d.kbDAO.GetByID(datasetID) + kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil || kb == nil { return nil, common.CodeDataError, errors.New("Invalid Dataset ID") } @@ -189,7 +189,7 @@ func (d *DatasetService) ListTags(datasetID, userID string) ([]map[string]interf return result, common.CodeSuccess, nil } -func (d *DatasetService) RenameTag(datasetID, userID, fromTag, toTag string) (map[string]interface{}, common.ErrorCode, error) { +func (d *DatasetService) RenameTag(ctx context.Context, datasetID, userID, fromTag, toTag string) (map[string]interface{}, common.ErrorCode, error) { fromTag = strings.TrimSpace(fromTag) toTag = strings.TrimSpace(toTag) datasetID, err := normalizeDatasetID(datasetID) @@ -199,13 +199,13 @@ func (d *DatasetService) RenameTag(datasetID, userID, fromTag, toTag string) (ma if strings.TrimSpace(datasetID) == "" { return nil, common.CodeDataError, errors.New("Lack of \"Dataset ID\"") } - if !d.kbDAO.Accessible(datasetID, userID) { + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return nil, common.CodeDataError, errors.New("No authorization.") } if d.docEngine == nil { return nil, common.CodeServerError, errors.New("Document engine is not initialized") } - kb, err := d.kbDAO.GetByID(datasetID) + kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil || kb == nil { return nil, common.CodeDataError, errors.New("Invalid Dataset ID") } diff --git a/internal/service/dataset/tags_aggregate_test.go b/internal/service/dataset/tags_aggregate_test.go index 1a5f1b7c81..e560a06e87 100644 --- a/internal/service/dataset/tags_aggregate_test.go +++ b/internal/service/dataset/tags_aggregate_test.go @@ -169,8 +169,9 @@ func TestDatasetServiceAggregateTagsMergesAcrossTenants(t *testing.T) { }, }, } + ctx := t.Context() - result, code, err := testDatasetServiceForAggregateTags(t, docEngine).AggregateTags([]string{kb1Input, kb2Input}, "user-1") + result, code, err := testDatasetServiceForAggregateTags(t, docEngine).AggregateTags(ctx, []string{kb1Input, kb2Input}, "user-1") if err != nil { t.Fatalf("AggregateTags failed: %v", err) } @@ -252,8 +253,9 @@ func TestDatasetServiceAggregateTagsPagesThroughAllChunks(t *testing.T) { }, }, } + ctx := t.Context() - result, code, err := testDatasetServiceForAggregateTags(t, docEngine).AggregateTags([]string{kbInput}, "user-1") + result, code, err := testDatasetServiceForAggregateTags(t, docEngine).AggregateTags(ctx, []string{kbInput}, "user-1") if err != nil { t.Fatalf("AggregateTags failed: %v", err) } @@ -287,8 +289,9 @@ func TestDatasetServiceAggregateTagsRejectsUnauthorizedDataset(t *testing.T) { kbInput := "323e4567-e89b-12d3-a456-426614174002" kbID := strings.ReplaceAll(kbInput, "-", "") insertAggregateTagsKB(t, kbID, "tenant-9", string(entity.TenantPermissionMe), 1) + ctx := t.Context() - _, code, err := testDatasetServiceForAggregateTags(t, &aggregateTagsMockEngine{}).AggregateTags([]string{kbInput}, "user-1") + _, code, err := testDatasetServiceForAggregateTags(t, &aggregateTagsMockEngine{}).AggregateTags(ctx, []string{kbInput}, "user-1") if err == nil { t.Fatal("expected authorization error") } @@ -301,7 +304,8 @@ func TestDatasetServiceAggregateTagsRejectsUnauthorizedDataset(t *testing.T) { } func TestDatasetServiceAggregateTagsRequiresDocumentEngine(t *testing.T) { - _, code, err := testDatasetServiceForAggregateTags(t, nil).AggregateTags([]string{"123e4567-e89b-12d3-a456-426614174000"}, "user-1") + ctx := t.Context() + _, code, err := testDatasetServiceForAggregateTags(t, nil).AggregateTags(ctx, []string{"123e4567-e89b-12d3-a456-426614174000"}, "user-1") if err == nil { t.Fatal("expected missing doc engine error") } @@ -334,8 +338,9 @@ func TestDatasetServiceAggregateTagsSkipsDatasetsWithoutDocuments(t *testing.T) }, }, } + ctx := t.Context() - result, code, err := testDatasetServiceForAggregateTags(t, docEngine).AggregateTags([]string{emptyInput, liveInput}, "user-1") + result, code, err := testDatasetServiceForAggregateTags(t, docEngine).AggregateTags(ctx, []string{emptyInput, liveInput}, "user-1") if err != nil { t.Fatalf("AggregateTags failed: %v", err) } @@ -362,9 +367,10 @@ func TestDatasetServiceAggregateTagsReturnsSearchError(t *testing.T) { kbInput := "623e4567-e89b-12d3-a456-426614174005" kbID := strings.ReplaceAll(kbInput, "-", "") insertAggregateTagsKB(t, kbID, "user-1", string(entity.TenantPermissionMe), 1) + ctx := t.Context() docEngine := &aggregateTagsMockEngine{searchErr: errors.New("boom")} - _, code, err := testDatasetServiceForAggregateTags(t, docEngine).AggregateTags([]string{kbInput}, "user-1") + _, code, err := testDatasetServiceForAggregateTags(t, docEngine).AggregateTags(ctx, []string{kbInput}, "user-1") if err == nil { t.Fatal("expected search error") } diff --git a/internal/service/dataset/tags_list_test.go b/internal/service/dataset/tags_list_test.go index 7304bf2b0e..da36503fb1 100644 --- a/internal/service/dataset/tags_list_test.go +++ b/internal/service/dataset/tags_list_test.go @@ -137,8 +137,9 @@ func TestDatasetServiceListTagsSuccess(t *testing.T) { }, }, } + ctx := t.Context() - result, code, err := testDatasetServiceForListTags(t, docEngine).ListTags(kbInput, "user-1") + result, code, err := testDatasetServiceForListTags(t, docEngine).ListTags(ctx, kbInput, "user-1") if err != nil { t.Fatalf("ListTags failed: %v", err) } @@ -176,7 +177,9 @@ func TestDatasetServiceListTagsReturnsEmptyWhenChunkStoreMissing(t *testing.T) { docEngine := &listTagsMockEngine{chunkStoreExists: false} - result, code, err := testDatasetServiceForListTags(t, docEngine).ListTags(kbInput, "user-1") + ctx := t.Context() + + result, code, err := testDatasetServiceForListTags(t, docEngine).ListTags(ctx, kbInput, "user-1") if err != nil { t.Fatalf("ListTags failed: %v", err) } @@ -201,7 +204,9 @@ func TestDatasetServiceListTagsRejectsUnauthorizedDataset(t *testing.T) { docEngine := &listTagsMockEngine{chunkStoreExists: true} - _, code, err := testDatasetServiceForListTags(t, docEngine).ListTags(kbInput, "user-1") + ctx := t.Context() + + _, code, err := testDatasetServiceForListTags(t, docEngine).ListTags(ctx, kbInput, "user-1") if err == nil { t.Fatal("expected unauthorized error") } @@ -225,7 +230,9 @@ func TestDatasetServiceListTagsReturnsChunkStoreError(t *testing.T) { chunkStoreErr: errors.New("boom"), } - _, code, err := testDatasetServiceForListTags(t, docEngine).ListTags(kbInput, "user-1") + ctx := t.Context() + + _, code, err := testDatasetServiceForListTags(t, docEngine).ListTags(ctx, kbInput, "user-1") if err == nil { t.Fatal("expected chunk store error") } diff --git a/internal/service/dataset/tags_rename_test.go b/internal/service/dataset/tags_rename_test.go index 567a6c7f22..47199b7d20 100644 --- a/internal/service/dataset/tags_rename_test.go +++ b/internal/service/dataset/tags_rename_test.go @@ -74,7 +74,9 @@ func TestDatasetServiceRenameTagSuccess(t *testing.T) { insertRenameTagKB(t, kbID, "user-1", string(entity.TenantPermissionMe)) docEngine := &renameTagMockEngine{} - result, code, err := testDatasetServiceForRenameTag(t, docEngine).RenameTag(kbInput, "user-1", "old-tag ", "new-tag") + ctx := t.Context() + + result, code, err := testDatasetServiceForRenameTag(t, docEngine).RenameTag(ctx, kbInput, "user-1", "old-tag ", "new-tag") if err != nil { t.Fatalf("RenameTag failed: %v", err) } @@ -120,7 +122,9 @@ func TestDatasetServiceRenameTagUnauthorized(t *testing.T) { kbID := strings.ReplaceAll(kbInput, "-", "") insertRenameTagKB(t, kbID, "tenant-9", string(entity.TenantPermissionMe)) - _, code, err := testDatasetServiceForRenameTag(t, &renameTagMockEngine{}).RenameTag(kbInput, "user-1", "old", "new") + ctx := t.Context() + + _, code, err := testDatasetServiceForRenameTag(t, &renameTagMockEngine{}).RenameTag(ctx, kbInput, "user-1", "old", "new") if err == nil { t.Fatal("expected unauthorized error") } @@ -141,7 +145,9 @@ func TestDatasetServiceRenameTagUpdateError(t *testing.T) { insertRenameTagKB(t, kbID, "user-1", string(entity.TenantPermissionMe)) docEngine := &renameTagMockEngine{updateErr: errors.New("boom")} - _, code, err := testDatasetServiceForRenameTag(t, docEngine).RenameTag(kbInput, "user-1", "old", "new") + ctx := t.Context() + + _, code, err := testDatasetServiceForRenameTag(t, docEngine).RenameTag(ctx, kbInput, "user-1", "old", "new") if err == nil { t.Fatal("expected update error") } diff --git a/internal/service/dataset/update.go b/internal/service/dataset/update.go index f5e01f969d..d7b2d820c2 100644 --- a/internal/service/dataset/update.go +++ b/internal/service/dataset/update.go @@ -29,7 +29,7 @@ type datasetPagerankUpdate struct { func (d *DatasetService) UpdateDataset(ctx context.Context, datasetID, tenantID string, req service.UpdateDatasetRequest) (map[string]interface{}, common.ErrorCode, error) { datasetID = strings.TrimSpace(datasetID) tenantID = strings.TrimSpace(tenantID) - if _, err := d.kbDAO.GetByID(datasetID); err != nil { + if _, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID); err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, errors.New("dataset not found") } diff --git a/internal/service/dataset/update_test.go b/internal/service/dataset/update_test.go index 4efb9443f7..caa14ec33b 100644 --- a/internal/service/dataset/update_test.go +++ b/internal/service/dataset/update_test.go @@ -73,7 +73,7 @@ func TestDatasetServiceUpdateDatasetUpdatesFields(t *testing.T) { t.Fatalf("expected empty connector list, got %#v", result["connectors"]) } - persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get updated kb: %v", err) } @@ -359,7 +359,7 @@ func TestDatasetServiceUpdateDatasetRejectsTeamMemberPermissionChange(t *testing t.Fatalf("unexpected error: %v", err) } - persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get dataset: %v", err) } @@ -517,7 +517,7 @@ func TestDatasetServiceUpdateDatasetAcceptsProviderInstanceEmbedding(t *testing. t.Fatalf("expected embedding model %q, got %#v", embeddingModel, result["embedding_model"]) } - persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get updated kb: %v", err) } @@ -549,7 +549,7 @@ func TestDatasetServiceUpdateDatasetAcceptsEmbeddingModelID(t *testing.T) { t.Fatalf("expected embedding model %q, got %#v", embeddingModelID, result["embedding_model"]) } - persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get updated kb: %v", err) } @@ -678,7 +678,7 @@ func TestDatasetServiceUpdateDatasetPreservesUnmodifiedFields(t *testing.T) { t.Fatalf("expected embedding_model preserved, got %#v", result["embedding_model"]) } - persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get updated kb: %v", err) } @@ -711,7 +711,7 @@ func TestDatasetServiceUpdateDatasetPreservesParserConfigOnEmptyUpdate(t *testin t.Fatalf("expected success code, got %d", code) } - persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get updated kb: %v", err) } @@ -728,9 +728,10 @@ func TestDatasetServiceDeleteDatasetsRejectsUnauthorizedID(t *testing.T) { pushServiceDB(t, db) insertDatasetUpdateKB(t, "11111111111141118111111111111111", "tenant-1", "Test") + ctx := t.Context() svc := NewDatasetService() normalizedID := "11111111111141118111111111111111" - _, code, err := svc.DeleteDatasets([]string{normalizedID}, false, "tenant-2") + _, code, err := svc.DeleteDatasets(ctx, []string{normalizedID}, false, "tenant-2") if err == nil { t.Fatal("expected unauthorized error") } @@ -746,8 +747,9 @@ func TestDatasetServiceDeleteDatasetsRejectsAllUnauthorized(t *testing.T) { db := setupDatasetUpdateTestDB(t) pushServiceDB(t, db) + ctx := t.Context() svc := NewDatasetService() - _, code, err := svc.DeleteDatasets([]string{"d94a8dc02c9711f0930f7fbc369eab6d"}, false, "tenant-1") + _, code, err := svc.DeleteDatasets(ctx, []string{"d94a8dc02c9711f0930f7fbc369eab6d"}, false, "tenant-1") if err == nil { t.Fatal("expected unauthorized error") } @@ -986,7 +988,7 @@ func TestUpdateDataset_StripsUnknownParam_Builtin(t *testing.T) { t.Fatalf("expected success code, got %d", code) } - persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get updated kb: %v", err) } @@ -1026,7 +1028,7 @@ func TestUpdateDataset_AcceptsValidComponentParams_Builtin(t *testing.T) { t.Fatalf("expected parser_id preserved, got %#v", result["parser_id"]) } - persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get updated kb: %v", err) } @@ -1074,7 +1076,8 @@ func TestUpdateDataset_PreservesIncomingMetadataWhenCleaningParserConfig(t *test t.Fatalf("UpdateDataset failed: code=%d err=%v", code, err) } - persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + ctx := t.Context() + persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get updated kb: %v", err) } @@ -1124,7 +1127,8 @@ func TestUpdateDataset_PreservesExistingMetadataWhenParserConfigOmitsIt(t *testi t.Fatalf("UpdateDataset failed: code=%d err=%v", code, err) } - persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + ctx := t.Context() + persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get updated kb: %v", err) } @@ -1161,7 +1165,7 @@ func TestUpdateDataset_StripsCanvasUnknownParam(t *testing.T) { t.Fatalf("expected success code, got %d", code) } - persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get updated kb: %v", err) } diff --git a/internal/service/document/document_crud.go b/internal/service/document/document_crud.go index ff71bd1606..b44c6a5baa 100644 --- a/internal/service/document/document_crud.go +++ b/internal/service/document/document_crud.go @@ -26,7 +26,7 @@ func (s *DocumentService) Accessible(ctx context.Context, docID, userID string) if err != nil || doc == nil { return false } - return s.kbDAO.Accessible(doc.KbID, userID) + return s.kbDAO.Accessible(ctx, dao.DB, doc.KbID, userID) } func (s *DocumentService) GetDocumentStorageAddress(ctx context.Context, doc *entity.Document) (string, string, error) { @@ -190,7 +190,7 @@ func (s *DocumentService) DeleteDocument(ctx context.Context, id string) error { // Returns the number of successfully deleted documents. func (s *DocumentService) DeleteDocuments(ctx context.Context, ids []string, deleteAll bool, datasetID, userID string) (int, error) { // 1. Check dataset is accessible by the user - if !s.kbDAO.Accessible(datasetID, userID) { + if !s.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { return 0, fmt.Errorf("you don't own the dataset %s", datasetID) } @@ -333,7 +333,7 @@ func (s *DocumentService) resolveDocAndKB(ctx context.Context, docID string) (*e if err != nil { return nil, nil, fmt.Errorf("document not found: %w", err) } - kb, err := s.kbDAO.GetByID(doc.KbID) + kb, err := s.kbDAO.GetByID(ctx, dao.DB, doc.KbID) if err != nil { return nil, nil, fmt.Errorf("knowledgebase not found: %w", err) } diff --git a/internal/service/document/document_dataset_update.go b/internal/service/document/document_dataset_update.go index baeec31845..69df41b799 100644 --- a/internal/service/document/document_dataset_update.go +++ b/internal/service/document/document_dataset_update.go @@ -20,7 +20,7 @@ import ( ) func (s *DocumentService) BatchUpdateDocumentStatus(ctx context.Context, userID, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) { - kb, err := s.kbDAO.GetByIDAndTenantID(datasetID, userID) + kb, err := s.kbDAO.GetByIDAndTenantID(ctx, dao.DB, datasetID, userID) if err != nil { return nil, common.CodeDataError, fmt.Errorf("You don't own the dataset.") } @@ -110,7 +110,7 @@ func (s *DocumentService) BatchUpdateDocumentStatus(ctx context.Context, userID, func (s *DocumentService) UpdateDatasetDocument(ctx context.Context, userID, datasetID, documentID string, req *UpdateDatasetDocumentRequest, present map[string]bool) (*UpdateDatasetDocumentResponse, common.ErrorCode, error) { tenantID := userID - kb, err := s.kbDAO.GetByIDAndTenantID(datasetID, tenantID) + kb, err := s.kbDAO.GetByIDAndTenantID(ctx, dao.DB, datasetID, tenantID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, errors.New("You don't own the dataset.") diff --git a/internal/service/document/document_ingest.go b/internal/service/document/document_ingest.go index ddcca010b6..eb398e4981 100644 --- a/internal/service/document/document_ingest.go +++ b/internal/service/document/document_ingest.go @@ -59,11 +59,11 @@ func (s *DocumentService) Ingest(ctx context.Context, userID string, req *Ingest if doc == nil { return common.CodeDataError, fmt.Errorf("document not found") } - kb, err := s.kbDAO.GetByID(doc.KbID) + kb, err := s.kbDAO.GetByID(ctx, dao.DB, doc.KbID) if err != nil { return common.CodeDataError, fmt.Errorf("dataset not found") } - if !s.kbDAO.Accessible(kb.ID, userID) { + if !s.kbDAO.Accessible(ctx, dao.DB, kb.ID, userID) { return common.CodeAuthenticationError, fmt.Errorf("no authorization") } validated = append(validated, validatedDoc{doc, kb}) diff --git a/internal/service/document/document_metadata.go b/internal/service/document/document_metadata.go index 14fa8cdbd7..fc4b3fda62 100644 --- a/internal/service/document/document_metadata.go +++ b/internal/service/document/document_metadata.go @@ -20,7 +20,7 @@ import ( // GetMetadataSummary get metadata summary for documents func (s *DocumentService) GetMetadataSummary(ctx context.Context, kbID string, docIDs []string) (map[string]interface{}, error) { - tenantID, err := s.metadataSvc.GetTenantIDByKBID(kbID) + tenantID, err := s.metadataSvc.GetTenantIDByKBID(ctx, kbID) if err != nil { return nil, err } @@ -43,7 +43,7 @@ func (s *DocumentService) SetDocumentMetadata(ctx context.Context, docID string, } // Get tenant ID - tenantID, err := s.metadataSvc.GetTenantIDByKBID(doc.KbID) + tenantID, err := s.metadataSvc.GetTenantIDByKBID(ctx, doc.KbID) if err != nil { return fmt.Errorf("failed to get tenant ID: %w", err) } @@ -64,7 +64,7 @@ func (s *DocumentService) DeleteDocumentMetadata(ctx context.Context, docID stri } // Get tenant ID - tenantID, err := s.metadataSvc.GetTenantIDByKBID(doc.KbID) + tenantID, err := s.metadataSvc.GetTenantIDByKBID(ctx, doc.KbID) if err != nil { return fmt.Errorf("failed to get tenant ID: %w", err) } @@ -87,7 +87,7 @@ func (s *DocumentService) DeleteDocumentAllMetadata(ctx context.Context, docID s } // Get tenant ID - tenantID, err := s.metadataSvc.GetTenantIDByKBID(doc.KbID) + tenantID, err := s.metadataSvc.GetTenantIDByKBID(ctx, doc.KbID) if err != nil { return fmt.Errorf("failed to get tenant ID: %w", err) } @@ -115,7 +115,7 @@ func (s *DocumentService) GetDocumentMetadataByID(ctx context.Context, docID str return nil, fmt.Errorf("document not found: %w", err) } - tenantID, err := s.metadataSvc.GetTenantIDByKBID(doc.KbID) + tenantID, err := s.metadataSvc.GetTenantIDByKBID(ctx, doc.KbID) if err != nil { return nil, err } @@ -140,7 +140,7 @@ func (s *DocumentService) GetMetadataByKBs(ctx context.Context, kbIDs []string) return make(map[string]interface{}), nil } - searchResult, err := s.metadataSvc.SearchMetadataByKBs(kbIDs, 10000) + searchResult, err := s.metadataSvc.SearchMetadataByKBs(ctx, kbIDs, 10000) if err != nil { return nil, err } @@ -586,7 +586,7 @@ func (s *DocumentService) BatchUpdateDocumentMetadatas( // Apply metadata_condition filter. if len(selector.MetadataCondition) > 0 { - flattedMeta, err := s.metadataSvc.GetFlattedMetaByKBs([]string{datasetID}) + flattedMeta, err := s.metadataSvc.GetFlattedMetaByKBs(ctx, []string{datasetID}) if err != nil { return nil, common.CodeServerError, fmt.Errorf("failed to get flattened metadata: %w", err) } diff --git a/internal/service/document/document_test.go b/internal/service/document/document_test.go index 05138455bc..5b62bedfa8 100644 --- a/internal/service/document/document_test.go +++ b/internal/service/document/document_test.go @@ -198,7 +198,7 @@ func (fakeChatDocEngine) GetType() string { return "fake" } func (fakeChatDocEngine) SupportsPageRank() bool { return false } -func (fakeChatDocEngine) FilterDocIdsByMetaPushdown(context.Context, []string, []map[string]interface{}, string) []string { +func (fakeChatDocEngine) FilterDocIdsByMetaPushdown(context.Context, *gorm.DB, []string, []map[string]interface{}, string) []string { return nil } @@ -609,7 +609,7 @@ func TestDeleteDocumentFull_Basic(t *testing.T) { } // Verify KB counters decremented - kb, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + kb, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("kb not found: %v", err) } @@ -917,7 +917,7 @@ func TestDeleteDocuments_DeleteAll(t *testing.T) { } // KB counters: doc_num 3→0, token_num 100→0, chunk_num 50→0 - kb, _ := dao.NewKnowledgebaseDAO().GetByID("kb-1") + kb, _ := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if kb.DocNum != 0 { t.Fatalf("doc_num: expected 0, got %d", kb.DocNum) } @@ -1077,7 +1077,7 @@ func TestStartParseDocuments_EnqueuesIngestionTask(t *testing.T) { svc.ingestionTaskSvc.SetTaskPublisher(publisher) ctx := t.Context() - kb, err := svc.kbDAO.GetByID("kb-1") + kb, err := svc.kbDAO.GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("load kb: %v", err) } @@ -1407,7 +1407,7 @@ func TestDeleteDocRecordWithCounters_Success(t *testing.T) { } // Counters decremented - kb, _ := dao.NewKnowledgebaseDAO().GetByID("kb-1") + kb, _ := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if kb.DocNum != 2 { t.Fatalf("doc_num: expected 2, got %d", kb.DocNum) } @@ -1441,7 +1441,7 @@ func TestDeleteDocRecordWithCounters_DocAlreadyDeleted(t *testing.T) { } // KB counters should be decremented exactly once: 1→0 for doc_num - kb, _ := dao.NewKnowledgebaseDAO().GetByID("kb-1") + kb, _ := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if kb.DocNum != 0 { t.Fatalf("doc_num: expected 0 (decremented once), got %d", kb.DocNum) } @@ -1752,7 +1752,7 @@ func TestUpdateDatasetDocumentParserIDResetsForReparse(t *testing.T) { if doc.TokenNum != 0 || doc.ChunkNum != 0 || doc.Progress != 0 { t.Fatalf("doc counters/progress = token:%d chunk:%d progress:%f, want zero", doc.TokenNum, doc.ChunkNum, doc.Progress) } - kb, _ := dao.NewKnowledgebaseDAO().GetByID("kb-1") + kb, _ := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if kb.TokenNum != 0 || kb.ChunkNum != 0 { t.Fatalf("kb counters = token:%d chunk:%d, want zero", kb.TokenNum, kb.ChunkNum) } @@ -1783,7 +1783,7 @@ func TestResetDocumentForReparseSkipsSecondCounterDecrement(t *testing.T) { if doc.TokenNum != 0 || doc.ChunkNum != 0 { t.Fatalf("doc counters = token:%d chunk:%d, want zero", doc.TokenNum, doc.ChunkNum) } - kb, _ := dao.NewKnowledgebaseDAO().GetByID("kb-1") + kb, _ := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if kb.TokenNum != 0 || kb.ChunkNum != 0 { t.Fatalf("kb counters = token:%d chunk:%d, want zero after duplicate reset", kb.TokenNum, kb.ChunkNum) } @@ -1814,7 +1814,7 @@ func TestClearDocumentParseResultsClearsCountersTasksAndChunks(t *testing.T) { if updatedDoc.TokenNum != 0 || updatedDoc.ChunkNum != 0 { t.Fatalf("doc counters = token:%d chunk:%d, want zero", updatedDoc.TokenNum, updatedDoc.ChunkNum) } - kb, _ := dao.NewKnowledgebaseDAO().GetByID("kb-1") + kb, _ := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if kb.TokenNum != 0 || kb.ChunkNum != 0 { t.Fatalf("kb counters = token:%d chunk:%d, want zero", kb.TokenNum, kb.ChunkNum) } @@ -1855,7 +1855,7 @@ func TestClearDocumentParseResultsIsIdempotentForStaleDocSnapshot(t *testing.T) if doc.TokenNum != 0 || doc.ChunkNum != 0 { t.Fatalf("doc counters = token:%d chunk:%d, want zero", doc.TokenNum, doc.ChunkNum) } - kb, _ := dao.NewKnowledgebaseDAO().GetByID("kb-1") + kb, _ := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if kb.TokenNum != 0 || kb.ChunkNum != 0 { t.Fatalf("kb counters = token:%d chunk:%d, want zero after duplicate prepare", kb.TokenNum, kb.ChunkNum) } @@ -2786,7 +2786,7 @@ func TestStartParseDocuments_FailsBeforeClearing(t *testing.T) { if err != nil { t.Fatalf("get doc: %v", err) } - kb, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + kb, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") if err != nil { t.Fatalf("get kb: %v", err) } diff --git a/internal/service/document/file2document.go b/internal/service/document/file2document.go index 9e587082de..bacb3586bf 100644 --- a/internal/service/document/file2document.go +++ b/internal/service/document/file2document.go @@ -102,7 +102,7 @@ func (s *File2DocumentService) LinkToDatasets(ctx context.Context, userID string // ── 2. Validate KBs exist ──────────────────────────────────────────────── kbMap := make(map[string]*entity.Knowledgebase, len(req.KbIDs)) for _, kbID := range req.KbIDs { - kb, err := s.kbDAO.GetByID(kbID) + kb, err := s.kbDAO.GetByID(ctx, dao.DB, kbID) if err != nil || kb == nil { return ErrLinkDatasetNotFound } @@ -222,7 +222,7 @@ func (s *File2DocumentService) convertFiles(ctx context.Context, fileIDs, kbIDs if _, exists := existingKBIDs[kbID]; exists { continue } - kb, err := s.kbDAO.GetByID(kbID) + kb, err := s.kbDAO.GetByID(ctx, dao.DB, kbID) if err != nil || kb == nil { continue } diff --git a/internal/service/file_permission.go b/internal/service/file_permission.go index a91bf12887..0cb26be104 100644 --- a/internal/service/file_permission.go +++ b/internal/service/file_permission.go @@ -23,7 +23,7 @@ func CheckFileTeamPermission(ctx context.Context, fileDAO *dao.FileDAO, file *en kbDAO := dao.NewKnowledgebaseDAO() for _, datasetID := range datasetIDs { - kb, err := kbDAO.GetByID(datasetID) + kb, err := kbDAO.GetByID(ctx, dao.DB, datasetID) if err != nil || kb == nil { continue } diff --git a/internal/service/memory_message_test.go b/internal/service/memory_message_test.go index 6ff427f3c5..685c35dcff 100644 --- a/internal/service/memory_message_test.go +++ b/internal/service/memory_message_test.go @@ -62,7 +62,7 @@ func (e *memoryMessageDocEngine) UpdateChunks(ctx context.Context, condition map return nil } -func (e *memoryMessageDocEngine) FilterDocIdsByMetaPushdown(_ context.Context, _ []string, _ []map[string]interface{}, _ string) []string { +func (e *memoryMessageDocEngine) FilterDocIdsByMetaPushdown(_ context.Context, _ *gorm.DB, _ []string, _ []map[string]interface{}, _ string) []string { return nil } diff --git a/internal/service/metadata.go b/internal/service/metadata.go index dee0dae91c..3e8abd5e52 100644 --- a/internal/service/metadata.go +++ b/internal/service/metadata.go @@ -65,16 +65,16 @@ func BuildMetadataIndexName(tenantID string) string { } // GetTenantIDByKBID retrieves tenant ID from knowledge base ID -func (s *MetadataService) GetTenantIDByKBID(kbID string) (string, error) { - return dao.GetTenantIDByKBID(kbID) +func (s *MetadataService) GetTenantIDByKBID(ctx context.Context, kbID string) (string, error) { + return dao.GetTenantIDByKBID(ctx, dao.DB, kbID) } // GetTenantIDByKBIDs retrieves tenant ID from the first knowledge base ID in the list -func (s *MetadataService) GetTenantIDByKBIDs(kbIDs []string) (string, error) { +func (s *MetadataService) GetTenantIDByKBIDs(ctx context.Context, kbIDs []string) (string, error) { if len(kbIDs) == 0 { return "", fmt.Errorf("no kb_ids provided") } - return dao.GetTenantIDByKBID(kbIDs[0]) + return dao.GetTenantIDByKBID(ctx, dao.DB, kbIDs[0]) } // SearchMetadataResponse holds the result of a metadata search @@ -107,12 +107,12 @@ func (s *MetadataService) SearchMetadata(kbID, tenantID string, docIDs []string, } // SearchMetadataByKBs searches the metadata index for multiple knowledge bases -func (s *MetadataService) SearchMetadataByKBs(kbIDs []string, size int) (*SearchMetadataResponse, error) { +func (s *MetadataService) SearchMetadataByKBs(ctx context.Context, kbIDs []string, size int) (*SearchMetadataResponse, error) { if len(kbIDs) == 0 { return &SearchMetadataResponse{MetadataRecords: []map[string]interface{}{}}, nil } - tenantID, err := s.GetTenantIDByKBIDs(kbIDs) + tenantID, err := s.GetTenantIDByKBIDs(ctx, kbIDs) if err != nil { return nil, err } @@ -126,7 +126,7 @@ func (s *MetadataService) SearchMetadataByKBs(kbIDs []string, size int) (*Search }, } - searchResult, err := s.docEngine.SearchMetadata(context.Background(), searchReq) + searchResult, err := s.docEngine.SearchMetadata(ctx, searchReq) if err != nil { return nil, fmt.Errorf("search failed: %w", err) } @@ -139,13 +139,13 @@ func (s *MetadataService) SearchMetadataByKBs(kbIDs []string, size int) (*Search // GetFlattedMetaByKBs returns flattened metadata in the format: // {field_name: {value: [doc_ids]}} -func (s *MetadataService) GetFlattedMetaByKBs(kbIDs []string) (common.MetaData, error) { +func (s *MetadataService) GetFlattedMetaByKBs(ctx context.Context, kbIDs []string) (common.MetaData, error) { if len(kbIDs) == 0 { return make(common.MetaData), nil } // Get metadata for all docs in KBs (use large limit like Python's 10000) - result, err := s.SearchMetadataByKBs(kbIDs, 10000) + result, err := s.SearchMetadataByKBs(ctx, kbIDs, 10000) if err != nil { return nil, err } diff --git a/internal/service/metadata_filter.go b/internal/service/metadata_filter.go index f27e7975f1..cda02729ac 100644 --- a/internal/service/metadata_filter.go +++ b/internal/service/metadata_filter.go @@ -22,6 +22,7 @@ import ( "fmt" "os" "ragflow/internal/common" + "ragflow/internal/dao" "ragflow/internal/engine" "regexp" @@ -647,7 +648,7 @@ func ApplyMetaDataFilter( "value": c.Value, } } - pushdownIDs := docEngine.FilterDocIdsByMetaPushdown(ctx, kbIDs, condMaps, logic) + pushdownIDs := docEngine.FilterDocIdsByMetaPushdown(ctx, dao.DB, kbIDs, condMaps, logic) // nil = push-down not viable / errored -> fall back to in-memory // non-nil (including empty slice) = push-down definitive -> use as-is if pushdownIDs != nil { diff --git a/internal/service/model_chat.go b/internal/service/model_chat.go index ac1946e01d..eba4bd549b 100644 --- a/internal/service/model_chat.go +++ b/internal/service/model_chat.go @@ -35,7 +35,7 @@ const streamDoneSentinel = "[DONE]" var errStreamDone = errors.New("chat stream done") func (m *ModelProviderService) Chat(ctx context.Context, tenantID, modelID string, messages []modelModule.Message, config *modelModule.ChatConfig) (*modelModule.ChatResponse, error) { - chatModel, err := m.GetChatModel(tenantID, modelID) + chatModel, err := m.GetChatModel(ctx, tenantID, modelID) if err != nil { return nil, err } @@ -43,7 +43,7 @@ func (m *ModelProviderService) Chat(ctx context.Context, tenantID, modelID strin } func (m *ModelProviderService) ChatStream(ctx context.Context, tenantID, modelID string, messages []modelModule.Message, config *modelModule.ChatConfig) (<-chan string, error) { - chatModel, err := m.GetChatModel(tenantID, modelID) + chatModel, err := m.GetChatModel(ctx, tenantID, modelID) if err != nil { return nil, err } diff --git a/internal/service/model_service.go b/internal/service/model_service.go index bb54e6fa1e..a5ddd32fc4 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -3189,7 +3189,7 @@ func (m *ModelProviderService) ParseFile(ctx context.Context, providerName, inst } // GetEmbeddingModel returns an EmbeddingModel wrapper for the given tenant -func (m *ModelProviderService) GetEmbeddingModel(tenantID, compositeModelName string) (*modelModule.EmbeddingModel, error) { +func (m *ModelProviderService) GetEmbeddingModel(ctx context.Context, tenantID, compositeModelName string) (*modelModule.EmbeddingModel, error) { driver, modelName, apiConfig, maxTokens, err := m.ResolveModelConfig(tenantID, entity.ModelTypeEmbedding, compositeModelName) if err != nil { return nil, err @@ -3198,7 +3198,7 @@ func (m *ModelProviderService) GetEmbeddingModel(tenantID, compositeModelName st } // GetChatModel returns a ChatModel wrapper for the given tenant -func (m *ModelProviderService) GetChatModel(tenantID, compositeModelName string) (*modelModule.ChatModel, error) { +func (m *ModelProviderService) GetChatModel(ctx context.Context, tenantID, compositeModelName string) (*modelModule.ChatModel, error) { driver, modelName, apiConfig, _, err := m.ResolveModelConfig(tenantID, entity.ModelTypeChat, compositeModelName) if err != nil { return nil, err diff --git a/internal/service/openai_chat.go b/internal/service/openai_chat.go index 78bfca4c41..ef98986187 100644 --- a/internal/service/openai_chat.go +++ b/internal/service/openai_chat.go @@ -286,7 +286,7 @@ func (s *OpenAIChatService) OpenAIChatCompletions(c *gin.Context, userID, chatID kbIDs = append(kbIDs, id) } } - metas, mdErr := s.pipeline.MetadataSvc.GetFlattedMetaByKBs(kbIDs) + metas, mdErr := s.pipeline.MetadataSvc.GetFlattedMetaByKBs(ctx, kbIDs) if mdErr != nil { s.writeDataError(c, fmt.Errorf("metadata_condition: load metadata: %w", mdErr).Error()) return diff --git a/internal/service/search.go b/internal/service/search.go index dd926707de..797706f968 100644 --- a/internal/service/search.go +++ b/internal/service/search.go @@ -17,6 +17,7 @@ package service import ( + "context" "fmt" "ragflow/internal/common" "ragflow/internal/dao" @@ -349,7 +350,7 @@ type SearchCompletionPlan struct { Options AskStreamOptions } -func (s *SearchService) PrepareCompletion(userID, searchID string, req *SearchCompletionsRequest) (*SearchCompletionPlan, common.ErrorCode, error) { +func (s *SearchService) PrepareCompletion(ctx context.Context, userID, searchID string, req *SearchCompletionsRequest) (*SearchCompletionPlan, common.ErrorCode, error) { userID = strings.TrimSpace(userID) if userID == "" { return nil, common.CodeBadRequest, fmt.Errorf("user id is required") @@ -389,7 +390,7 @@ func (s *SearchService) PrepareCompletion(userID, searchID string, req *SearchCo } for _, datasetID := range datasetIDs { - accessible = s.datasetDAO.Accessible(datasetID, userID) + accessible = s.datasetDAO.Accessible(ctx, dao.DB, datasetID, userID) if !accessible { return nil, common.CodeAuthenticationError, fmt.Errorf("no authorization for dataset %s", datasetID) } diff --git a/internal/service/skill_indexer.go b/internal/service/skill_indexer.go index 68bb2fa1cf..8bf3c4ce6e 100644 --- a/internal/service/skill_indexer.go +++ b/internal/service/skill_indexer.go @@ -922,7 +922,7 @@ func (s *SkillIndexerService) generateEmbedding(ctx context.Context, text, embdI return nil, fmt.Errorf("embedding model ID not configured") } - embeddingModel, err := s.modelProvider.GetEmbeddingModel(tenantID, embdID) + embeddingModel, err := s.modelProvider.GetEmbeddingModel(ctx, tenantID, embdID) if err != nil { return nil, fmt.Errorf("failed to get embedding model: %w", err) } @@ -960,7 +960,7 @@ func (s *SkillIndexerService) generateEmbeddings(ctx context.Context, texts []st } common.Info(fmt.Sprintf("Getting embedding model for %s", embdID)) - embeddingModel, err := s.modelProvider.GetEmbeddingModel(tenantID, embdID) + embeddingModel, err := s.modelProvider.GetEmbeddingModel(ctx, tenantID, embdID) if err != nil { common.Error(fmt.Sprintf("Failed to get embedding model: %v", err), err) return nil, fmt.Errorf("failed to get embedding model: %w", err) @@ -1018,7 +1018,7 @@ func (s *SkillIndexerService) getEmbeddingDimension(ctx context.Context, tenantI return 0, fmt.Errorf("embedding model ID not configured") } - embeddingModel, err := s.modelProvider.GetEmbeddingModel(tenantID, embdID) + embeddingModel, err := s.modelProvider.GetEmbeddingModel(ctx, tenantID, embdID) if err != nil { return 0, fmt.Errorf("failed to get embedding model: %w", err) } diff --git a/internal/service/skill_search.go b/internal/service/skill_search.go index 5a946219aa..d5c071a4a5 100644 --- a/internal/service/skill_search.go +++ b/internal/service/skill_search.go @@ -667,7 +667,7 @@ func (s *SkillSearchService) getEmbedding(ctx context.Context, text, embdID, ten return nil, fmt.Errorf("embedding model ID not configured") } - embeddingModel, err := s.modelProvider.GetEmbeddingModel(tenantID, embdID) + embeddingModel, err := s.modelProvider.GetEmbeddingModel(ctx, tenantID, embdID) if err != nil { return nil, fmt.Errorf("failed to get embedding model: %w", err) } diff --git a/internal/service/tag.go b/internal/service/tag.go index e5be8a2a29..16d7d7029e 100644 --- a/internal/service/tag.go +++ b/internal/service/tag.go @@ -116,7 +116,7 @@ type Knowledgebase = entity.Knowledgebase // GetAllTagsInPortion returns all tag_kwd values and their occurrence counts // for documents belonging to the given kbIDs. -func (s *MetadataService) GetAllTagsInPortion(tenantID string, kbIDs []string) (map[string]float64, error) { +func (s *MetadataService) GetAllTagsInPortion(ctx context.Context, tenantID string, kbIDs []string) (map[string]float64, error) { if len(kbIDs) == 0 { return make(map[string]float64), nil } @@ -133,7 +133,7 @@ func (s *MetadataService) GetAllTagsInPortion(tenantID string, kbIDs []string) ( SelectFields: []string{"tag_kwd"}, } - searchResp, err := s.docEngine.Search(context.Background(), searchReq) + searchResp, err := s.docEngine.Search(ctx, searchReq) if err != nil { return nil, err } @@ -164,7 +164,7 @@ func (s *MetadataService) GetAllTagsInPortion(tenantID string, kbIDs []string) ( } // 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) { +func (s *MetadataService) TagQuery(ctx context.Context, 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 } @@ -195,7 +195,7 @@ func (s *MetadataService) TagQuery(question string, tenantIDs []string, kbIDs [] MatchExprs: []interface{}{matchTextExpr}, } - searchResp, err := s.docEngine.Search(context.Background(), searchReq) + searchResp, err := s.docEngine.Search(ctx, searchReq) if err != nil { return nil, err } @@ -261,7 +261,7 @@ func (s *MetadataService) TagQuery(question string, tenantIDs []string, kbIDs [] // 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 { +func (s *MetadataService) LabelQuestion(ctx context.Context, question string, kbs []*Knowledgebase) map[string]float64 { if len(kbs) == 0 || question == "" { return nil } @@ -296,20 +296,20 @@ func (s *MetadataService) LabelQuestion(question string, kbs []*Knowledgebase) m } if allTags == nil { // Cache miss - compute all_tags_in_portion - allTags, err = s.GetAllTagsInPortion(lastKB.TenantID, tagKBIDs) + allTags, err = s.GetAllTagsInPortion(ctx, lastKB.TenantID, tagKBIDs) if err != nil { common.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 { + if err = SetTagsToCache(tagKBIDs, allTags); err != nil { common.Warn("Failed to set tags cache", zap.Error(err)) } } // Get tag_kbs by IDs kbDAO := dao.NewKnowledgebaseDAO() - tagKBs, err := kbDAO.GetByIDs(tagKBIDs) + tagKBs, err := kbDAO.GetByIDs(ctx, dao.DB, tagKBIDs) if err != nil || len(tagKBs) == 0 { // Return nil if no tag_kbs found return nil @@ -347,7 +347,7 @@ func (s *MetadataService) LabelQuestion(question string, kbs []*Knowledgebase) m } // Query tags for the question using unique tenant IDs - tagFeatures, err := s.TagQuery(question, uniqueTenantIDs, tagKBIDs, allTags, topnTags) + tagFeatures, err := s.TagQuery(ctx, question, uniqueTenantIDs, tagKBIDs, allTags, topnTags) if err != nil { return nil } diff --git a/internal/service/tenant.go b/internal/service/tenant.go index 7e38ac3e8d..5b3fa12b0a 100644 --- a/internal/service/tenant.go +++ b/internal/service/tenant.go @@ -356,12 +356,12 @@ type CreateChunkStoreResponse struct { } // CreateChunkStore creates a chunk store in the document engine for a knowledge base -func (s *TenantService) CreateChunkStore(req *CreateDatasetTableRequest) (*CreateChunkStoreResponse, common.ErrorCode, error) { +func (s *TenantService) CreateChunkStore(ctx context.Context, req *CreateDatasetTableRequest) (*CreateChunkStoreResponse, common.ErrorCode, error) { if req == nil { return nil, common.CodeDataError, fmt.Errorf("request is required") } // Get KB to find tenant_id for building table name - kb, err := s.kbDAO.GetByID(req.KBID) + kb, err := s.kbDAO.GetByID(ctx, dao.DB, req.KBID) if err != nil { if dao.IsNotFoundErr(err) { return nil, common.CodeDataError, fmt.Errorf("knowledge base not found: %s", req.KBID) @@ -380,7 +380,7 @@ func (s *TenantService) CreateChunkStore(req *CreateDatasetTableRequest) (*Creat // Call document engine to create table // Full table name will be built as "{tableName}_{kb_id}" - err = s.docEngine.CreateChunkStore(context.Background(), tableName, req.KBID, vecSize, req.ParserID) + err = s.docEngine.CreateChunkStore(ctx, tableName, req.KBID, vecSize, req.ParserID) if err != nil { return nil, common.CodeServerError, fmt.Errorf("failed to create dataset: %w", err) } @@ -393,9 +393,9 @@ func (s *TenantService) CreateChunkStore(req *CreateDatasetTableRequest) (*Creat } // DeleteChunkStore deletes the chunk store in the document engine for a knowledge base -func (s *TenantService) DeleteChunkStore(kbID string) (common.ErrorCode, error) { +func (s *TenantService) DeleteChunkStore(ctx context.Context, kbID string) (common.ErrorCode, error) { // Get KB to find tenant_id for building table name - kb, err := s.kbDAO.GetByID(kbID) + kb, err := s.kbDAO.GetByID(ctx, dao.DB, kbID) if err != nil { if dao.IsNotFoundErr(err) { return common.CodeDataError, fmt.Errorf("knowledge base not found: %s", kbID) @@ -404,7 +404,7 @@ func (s *TenantService) DeleteChunkStore(kbID string) (common.ErrorCode, error) } // Call document engine to delete table - err = s.docEngine.DropChunkStore(context.Background(), fmt.Sprintf("ragflow_%s", kb.TenantID), kbID) + err = s.docEngine.DropChunkStore(ctx, fmt.Sprintf("ragflow_%s", kb.TenantID), kbID) if err != nil { return common.CodeServerError, fmt.Errorf("failed to delete table: %w", err) } diff --git a/internal/service/test_helpers_test.go b/internal/service/test_helpers_test.go index 099d3ec03b..e35bb3c8e6 100644 --- a/internal/service/test_helpers_test.go +++ b/internal/service/test_helpers_test.go @@ -144,7 +144,7 @@ func (fakeChatDocEngine) GetScores(map[string]interface{}) map[string]float64 { func (fakeChatDocEngine) Ping(context.Context) error { return nil } func (fakeChatDocEngine) Close() error { return nil } func (fakeChatDocEngine) CheckStatus() error { return nil } -func (fakeChatDocEngine) FilterDocIdsByMetaPushdown(context.Context, []string, []map[string]interface{}, string) []string { +func (fakeChatDocEngine) FilterDocIdsByMetaPushdown(context.Context, *gorm.DB, []string, []map[string]interface{}, string) []string { return nil } func (fakeChatDocEngine) GetMessages(context.Context, string, int, int) ([]interface{}, error) {