From a7da78d0d7d6f3e30c5d44f68242c6be36df33b9 Mon Sep 17 00:00:00 2001 From: Hz_ Date: Wed, 15 Jul 2026 21:42:08 +0800 Subject: [PATCH] refactor(go-agent): unify tool-backed canvas components (#16912) ## Summary This PR consolidates Eino-backed Agent tools behind the shared ToolBackedComponent implementation and aligns their Canvas configuration, runtime inputs, output conversion, validation, and registration. ## What changed - Migrated these Canvas components to the unified tool-backed path: - Tavily Search and Extract - Execute SQL - Google - Yahoo Finance - Email - DuckDuckGo - Wikipedia - Google Scholar - ArXiv - PubMed - BGPT - GitHub - WenCai - SearXNG - Keenable Search - Removed superseded component wrappers and their duplicate tests. - Added dedicated registry builders for node-level configuration and validation. - Kept model-emitted runtime inputs separate from Canvas node configuration. - Moved Email defaults, template resolution, recipient parsing, SMTP execution, and output conversion into the owning tool. - Added complete ToolComponent specifications for Canvas inputs, outputs, and input forms. - Preserved raw upstream fields where downstream workflows may depend on them. - Added workflow registration coverage for all migrated component names. - Kept HTTP Request, Docs Generator, and Browser as standalone components because they are not Eino-backed tools. ## Testing Passed: bash build.sh --test ./internal/agent/tool/... bash build.sh --test ./internal/agent/component/... bash build.sh --test ./internal/agent/runtime/... image image image --- .../agent/component/arxiv_component_test.go | 151 -- .../agent/component/bgpt_component_test.go | 139 -- .../component/duckduckgo_component_test.go | 132 -- internal/agent/component/email.go | 251 --- internal/agent/component/email_test.go | 223 --- .../agent/component/exesql_component_test.go | 100 -- internal/agent/component/fixture_stubs.go | 37 +- internal/agent/component/github.go | 230 --- internal/agent/component/github_test.go | 165 -- .../agent/component/google_component_test.go | 71 - .../google_scholar_component_test.go | 131 -- .../component/keenable_component_test.go | 169 -- .../component/production_chain_fixes_test.go | 43 +- .../agent/component/pubmed_component_test.go | 150 -- internal/agent/component/searxng.go | 252 --- internal/agent/component/searxng_test.go | 250 --- .../agent/component/tavily_component_test.go | 102 -- internal/agent/component/tool_component.go | 131 ++ .../agent/component/tool_component_test.go | 514 ++++++ .../agent/component/tool_dispatch_test.go | 2 +- .../agent/component/universe_a_wrappers.go | 1572 +---------------- internal/agent/component/wencai.go | 109 -- internal/agent/component/wencai_test.go | 190 -- .../component/wikipedia_component_test.go | 132 -- internal/agent/tool/arxiv.go | 124 +- internal/agent/tool/arxiv_test.go | 56 +- internal/agent/tool/bgpt.go | 315 ++-- internal/agent/tool/bgpt_test.go | 160 ++ internal/agent/tool/duckduckgo.go | 157 +- internal/agent/tool/duckduckgo_test.go | 66 +- internal/agent/tool/email.go | 196 +- internal/agent/tool/email_test.go | 278 +-- internal/agent/tool/exesql.go | 96 + internal/agent/tool/exesql_test.go | 52 + internal/agent/tool/github.go | 155 +- internal/agent/tool/github_test.go | 93 +- internal/agent/tool/google.go | 155 +- internal/agent/tool/google_scholar.go | 137 +- internal/agent/tool/google_scholar_test.go | 85 +- internal/agent/tool/google_test.go | 95 +- internal/agent/tool/keenable.go | 220 ++- internal/agent/tool/keenable_test.go | 113 +- internal/agent/tool/pubmed.go | 122 +- internal/agent/tool/pubmed_test.go | 58 +- internal/agent/tool/registry.go | 336 +++- internal/agent/tool/registry_test.go | 16 +- internal/agent/tool/searxng.go | 168 +- internal/agent/tool/searxng_test.go | 86 +- internal/agent/tool/tavily.go | 351 +++- internal/agent/tool/tavily_test.go | 319 +++- internal/agent/tool/tool2component.go | 68 + internal/agent/tool/wencai.go | 27 + internal/agent/tool/wencai_test.go | 42 +- internal/agent/tool/wikipedia.go | 127 +- internal/agent/tool/wikipedia_test.go | 85 +- internal/agent/tool/yahoo_finance.go | 153 +- internal/agent/tool/yahoo_finance_test.go | 347 ++-- 57 files changed, 4779 insertions(+), 5325 deletions(-) delete mode 100644 internal/agent/component/arxiv_component_test.go delete mode 100644 internal/agent/component/bgpt_component_test.go delete mode 100644 internal/agent/component/duckduckgo_component_test.go delete mode 100644 internal/agent/component/email.go delete mode 100644 internal/agent/component/email_test.go delete mode 100644 internal/agent/component/exesql_component_test.go delete mode 100644 internal/agent/component/github.go delete mode 100644 internal/agent/component/github_test.go delete mode 100644 internal/agent/component/google_component_test.go delete mode 100644 internal/agent/component/google_scholar_component_test.go delete mode 100644 internal/agent/component/keenable_component_test.go delete mode 100644 internal/agent/component/pubmed_component_test.go delete mode 100644 internal/agent/component/searxng.go delete mode 100644 internal/agent/component/searxng_test.go delete mode 100644 internal/agent/component/tavily_component_test.go create mode 100644 internal/agent/component/tool_component.go create mode 100644 internal/agent/component/tool_component_test.go delete mode 100644 internal/agent/component/wencai.go delete mode 100644 internal/agent/component/wencai_test.go delete mode 100644 internal/agent/component/wikipedia_component_test.go create mode 100644 internal/agent/tool/bgpt_test.go create mode 100644 internal/agent/tool/tool2component.go diff --git a/internal/agent/component/arxiv_component_test.go b/internal/agent/component/arxiv_component_test.go deleted file mode 100644 index 6baa4a4bd5..0000000000 --- a/internal/agent/component/arxiv_component_test.go +++ /dev/null @@ -1,151 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "net/http" - "net/http/httptest" - "net/url" - "strings" - "testing" - - agenttool "ragflow/internal/agent/tool" -) - -type arxivRoundTripper func(*http.Request) (*http.Response, error) - -func (fn arxivRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) { - return fn(request) -} - -func TestArXiv_RegisteredFactoryAndInputForm(t *testing.T) { - t.Parallel() - - c, err := New("ArXiv", map[string]any{ - "top_n": float64(7), - "sort_by": "relevance", - }) - if err != nil { - t.Fatalf("New(ArXiv): %v", err) - } - if got := c.Name(); got != "ArXiv" { - t.Fatalf("Name() = %q, want ArXiv", got) - } - formGetter, ok := c.(interface{ GetInputForm() map[string]any }) - if !ok { - t.Fatal("ArXiv component does not expose GetInputForm") - } - query, ok := formGetter.GetInputForm()["query"].(map[string]any) - if !ok || query["type"] != "line" { - t.Fatalf("GetInputForm()[query] = %#v, want Query line input", query) - } - if _, ok := c.Outputs()["formalized_content"]; !ok { - t.Fatal("Outputs() missing formalized_content") - } - if _, ok := c.Outputs()["json"]; !ok { - t.Fatal("Outputs() missing json") - } -} - -func TestArXiv_InvalidNodeParams(t *testing.T) { - t.Parallel() - - for _, params := range []map[string]any{ - {"top_n": 0}, - {"sort_by": "newest"}, - } { - if _, err := New("ArXiv", params); err == nil { - t.Fatalf("New(ArXiv, %#v) succeeded, want validation error", params) - } - } -} - -func TestArXiv_InvokeSendsOnlyQueryAndFormatsPythonFields(t *testing.T) { - t.Parallel() - - const atom = ` - - - http://arxiv.org/abs/2501.12345v1 - Test Paper - Paper summary. - Author One - - -` - - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - query := r.URL.Query() - if got := query.Get("search_query"); got != "all:retrieval augmented generation" { - t.Errorf("search_query = %q", got) - } - if got := query.Get("max_results"); got != "7" { - t.Errorf("max_results = %q, want 7", got) - } - if got := query.Get("sortBy"); got != "relevance" { - t.Errorf("sortBy = %q, want relevance", got) - } - _, _ = w.Write([]byte(atom)) - })) - defer server.Close() - - serverURL, err := url.Parse(server.URL) - if err != nil { - t.Fatalf("url.Parse(server.URL): %v", err) - } - component := &arxivComponent{inner: agenttool.NewArxivToolWithParams(agenttool.NewHTTPHelper().WithClient(&http.Client{ - Transport: arxivRoundTripper(func(request *http.Request) (*http.Response, error) { - request.URL.Scheme = serverURL.Scheme - request.URL.Host = serverURL.Host - return http.DefaultTransport.RoundTrip(request) - }), - }), 7, "relevance")} - - out, err := component.Invoke(context.Background(), map[string]any{"query": " retrieval augmented generation "}) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - formalized, _ := out["formalized_content"].(string) - for _, want := range []string{"Test Paper", "http://arxiv.org/pdf/2501.12345v1", "Paper summary."} { - if !strings.Contains(formalized, want) { - t.Errorf("formalized_content missing %q: %s", want, formalized) - } - } - if results, ok := out["json"].([]any); !ok || len(results) != 1 { - t.Fatalf("json output = %#v, want one paper", out["json"]) - } -} - -func TestArXiv_InvokeEmptyQueryReturnsEmptyPayload(t *testing.T) { - t.Parallel() - - c, err := newArxivComponent(nil) - if err != nil { - t.Fatalf("newArxivComponent: %v", err) - } - out, err := c.Invoke(context.Background(), map[string]any{"query": " "}) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if got := out["formalized_content"]; got != "" { - t.Errorf("formalized_content = %v, want empty string", got) - } - if results, ok := out["json"].([]any); !ok || len(results) != 0 { - t.Fatalf("json output = %#v, want empty []any", out["json"]) - } -} diff --git a/internal/agent/component/bgpt_component_test.go b/internal/agent/component/bgpt_component_test.go deleted file mode 100644 index ac7242cdb0..0000000000 --- a/internal/agent/component/bgpt_component_test.go +++ /dev/null @@ -1,139 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "encoding/json" - "strings" - "testing" - - einotool "github.com/cloudwego/eino/components/tool" -) - -type fakeBGPTInvoker struct { - args map[string]any -} - -func (f *fakeBGPTInvoker) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) { - if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil { - return "", err - } - return `{"results":[{"title":"Paper A","authors":"Lee, Kim","journal":"Science","year":"2026","doi":"10.1/a","abstract":"Abstract A","methods":"RCT","sample_size":"120","results":"Improved outcomes","limitations":"Small cohort","conflict_of_interest":"None","data_availability":"Available","blind_spots":"Long term effects","falsify":"Run a larger trial"}]}`, nil -} - -func TestBGPT_RegisteredFactory(t *testing.T) { - c, err := New("BGPT", nil) - if err != nil { - t.Fatalf("New(BGPT) errored: %v", err) - } - if got := c.Name(); got != "BGPT" { - t.Fatalf("Name() = %q, want BGPT", got) - } - formGetter, ok := c.(interface{ GetInputForm() map[string]any }) - if !ok { - t.Fatal("BGPT component does not expose GetInputForm") - } - form := formGetter.GetInputForm() - query, ok := form["query"].(map[string]any) - if !ok { - t.Fatalf("GetInputForm()[query] has type %T, want map", form["query"]) - } - if query["type"] != "line" { - t.Fatalf("GetInputForm()[query][type] = %v, want line", query["type"]) - } - if _, ok := c.Outputs()["formalized_content"]; !ok { - t.Fatal("Outputs() missing formalized_content") - } - if _, ok := c.Outputs()["json"]; !ok { - t.Fatal("Outputs() missing json") - } -} - -func TestBGPT_InvokeAdaptsCanvasInputsAndOutputs(t *testing.T) { - fake := &fakeBGPTInvoker{} - c := newBGPTComponentWithInvoker(fake) - - out, err := c.Invoke(context.Background(), map[string]any{ - "query": " cancer therapy ", - "api_key": "key-1", - "days_back": float64(30), - "top_n": float64(3), - }) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - - if got := fake.args["query"]; got != "cancer therapy" { - t.Errorf("query arg = %v, want trimmed query", got) - } - if got := fake.args["api_key"]; got != "key-1" { - t.Errorf("api_key arg = %v, want key-1", got) - } - if got := fake.args["days_back"]; got != float64(30) { - t.Errorf("days_back arg = %v, want 30", got) - } - if got := fake.args["num_results"]; got != float64(3) { - t.Errorf("num_results arg = %v, want 3 from top_n", got) - } - - formalized, _ := out["formalized_content"].(string) - for _, want := range []string{"Paper A", "Lee, Kim", "Improved outcomes", "Run a larger trial"} { - if !strings.Contains(formalized, want) { - t.Errorf("formalized_content missing %q: %s", want, formalized) - } - } - - results, ok := out["json"].([]any) - if !ok { - t.Fatalf("json output has type %T, want []any", out["json"]) - } - if len(results) != 1 { - t.Fatalf("json output length = %d, want 1", len(results)) - } -} - -func TestBGPT_InvokeUsesStoredAPIKeyWhenInputOmitsIt(t *testing.T) { - fake := &fakeBGPTInvoker{} - c := newBGPTComponentWithInvoker(fake, "stored-key") - - _, err := c.Invoke(context.Background(), map[string]any{ - "query": "cancer therapy", - }) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - if got := fake.args["api_key"]; got != "stored-key" { - t.Fatalf("api_key arg = %v, want stored-key", got) - } -} - -func TestBGPT_InvokeDoesNotOverrideCallerAPIKey(t *testing.T) { - fake := &fakeBGPTInvoker{} - c := newBGPTComponentWithInvoker(fake, "stored-key") - - _, err := c.Invoke(context.Background(), map[string]any{ - "query": "cancer therapy", - "api_key": "call-key", - }) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - if got := fake.args["api_key"]; got != "call-key" { - t.Fatalf("api_key arg = %v, want call-key", got) - } -} diff --git a/internal/agent/component/duckduckgo_component_test.go b/internal/agent/component/duckduckgo_component_test.go deleted file mode 100644 index 2729a7d3b0..0000000000 --- a/internal/agent/component/duckduckgo_component_test.go +++ /dev/null @@ -1,132 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "encoding/json" - "strings" - "testing" - - einotool "github.com/cloudwego/eino/components/tool" -) - -type fakeDuckDuckGoInvoker struct { - args map[string]any -} - -func (f *fakeDuckDuckGoInvoker) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) { - if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil { - return "", err - } - return `{"results":[{"title":"RAGFlow","url":"https://ragflow.io","body":"Open source RAG engine"}]}`, nil -} - -func TestDuckDuckGo_RegisteredFactory(t *testing.T) { - t.Parallel() - - c, err := New("DuckDuckGo", nil) - if err != nil { - t.Fatalf("New(DuckDuckGo) errored: %v", err) - } - if got := c.Name(); got != "DuckDuckGo" { - t.Fatalf("Name() = %q, want DuckDuckGo", got) - } - formGetter, ok := c.(interface{ GetInputForm() map[string]any }) - if !ok { - t.Fatal("DuckDuckGo component does not expose GetInputForm") - } - form := formGetter.GetInputForm() - query, ok := form["query"].(map[string]any) - if !ok { - t.Fatalf("GetInputForm()[query] has type %T, want map", form["query"]) - } - if query["type"] != "line" { - t.Fatalf("GetInputForm()[query][type] = %v, want line", query["type"]) - } - channel, ok := form["channel"].(map[string]any) - if !ok { - t.Fatalf("GetInputForm()[channel] has type %T, want map", form["channel"]) - } - if channel["value"] != "general" { - t.Fatalf("GetInputForm()[channel][value] = %v, want general", channel["value"]) - } - if _, ok := c.Outputs()["formalized_content"]; !ok { - t.Fatal("Outputs() missing formalized_content") - } - if _, ok := c.Outputs()["json"]; !ok { - t.Fatal("Outputs() missing json") - } -} - -func TestDuckDuckGo_InvokeAdaptsCanvasInputsAndOutputs(t *testing.T) { - t.Parallel() - - fake := &fakeDuckDuckGoInvoker{} - c := newDuckDuckGoComponentWithInvoker(fake) - - out, err := c.Invoke(context.Background(), map[string]any{ - "query": " privacy search ", - "channel": "news", - "top_n": float64(3), - }) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - - if got := fake.args["query"]; got != "privacy search" { - t.Errorf("query arg = %v, want trimmed query", got) - } - if got := fake.args["channel"]; got != "news" { - t.Errorf("channel arg = %v, want news", got) - } - if got := fake.args["top_n"]; got != float64(3) { - t.Errorf("top_n arg = %v, want 3", got) - } - - formalized, _ := out["formalized_content"].(string) - for _, want := range []string{"RAGFlow", "https://ragflow.io", "Open source RAG engine"} { - if !strings.Contains(formalized, want) { - t.Errorf("formalized_content missing %q: %s", want, formalized) - } - } - - results, ok := out["json"].([]any) - if !ok { - t.Fatalf("json output has type %T, want []any", out["json"]) - } - if len(results) != 1 { - t.Fatalf("json output length = %d, want 1", len(results)) - } -} - -func TestDuckDuckGo_InvokeEmptyQueryReturnsEmptyPayload(t *testing.T) { - t.Parallel() - - c := newDuckDuckGoComponentWithInvoker(&fakeDuckDuckGoInvoker{}) - out, err := c.Invoke(context.Background(), map[string]any{"query": " "}) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - if got := out["formalized_content"]; got != "" { - t.Errorf("formalized_content = %v, want empty string", got) - } - results, ok := out["json"].([]any) - if !ok || len(results) != 0 { - t.Fatalf("json output = %#v, want empty []any", out["json"]) - } -} diff --git a/internal/agent/component/email.go b/internal/agent/component/email.go deleted file mode 100644 index 0e36395a32..0000000000 --- a/internal/agent/component/email.go +++ /dev/null @@ -1,251 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "encoding/json" - "fmt" - "ragflow/internal/agent/runtime" - "strconv" - "strings" - - agenttool "ragflow/internal/agent/tool" -) - -const componentNameEmail = "Email" - -// EmailComponent sends a canvas email node through the Go SMTP tool. -// It accepts the same DSL fields as the Python Email tool: -// smtp_server, smtp_port, email, smtp_username, password, sender_name, -// to_email, cc_email, content, and subject. -type EmailComponent struct { - name string - smtpServer string - smtpPort int - email string - smtpUsername string - password string - senderName string - toEmail string - ccEmail string - content string - subject string - tool *agenttool.EmailTool -} - -// NewEmailComponent constructs an Email component from DSL params. -func NewEmailComponent(params map[string]any) (Component, error) { - smtpPort, err := emailIntParam(params, "smtp_port", 465) - if err != nil { - return nil, err - } - return &EmailComponent{ - name: componentNameEmail, - smtpServer: emailStringParam(params, "smtp_server"), - smtpPort: smtpPort, - email: emailStringParam(params, "email"), - smtpUsername: emailStringParam(params, "smtp_username"), - password: emailStringParam(params, "password"), - senderName: emailStringParam(params, "sender_name"), - toEmail: emailStringParam(params, "to_email"), - ccEmail: emailStringParam(params, "cc_email"), - content: emailStringParam(params, "content"), - subject: emailStringParam(params, "subject"), - tool: agenttool.NewEmailTool(), - }, nil -} - -// Name returns the registered component name. -func (e *EmailComponent) Name() string { return e.name } - -// Invoke sends the email and returns success/_ERROR fields, matching the -// Python node's soft-failure behaviour. -func (e *EmailComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - toEmail := firstString(inputs, "to_email", e.toEmail) - content := firstString(inputs, "content", e.content) - subject := firstString(inputs, "subject", e.subject) - ccEmail := firstString(inputs, "cc_email", e.ccEmail) - - state, _, _ := runtime.GetStateFromContext[*runtime.CanvasState](ctx) - toEmail = runtime.ResolveTemplateForDisplay(toEmail, state) - ccEmail = runtime.ResolveTemplateForDisplay(ccEmail, state) - subject = runtime.ResolveTemplateForDisplay(subject, state) - subject = stripEmailHeaderLineBreaks(subject) - content = runtime.ResolveTemplateForDisplay(content, state) - - args := map[string]any{ - "smtp_host": e.smtpServer, - "smtp_port": e.smtpPort, - "username": firstNonEmpty(e.smtpUsername, e.email), - "password": e.password, - "from_addr": e.email, - "to_addrs": emailRecipients(toEmail, ccEmail), - "subject": subject, - "body": content, - } - argsJSON, err := json.Marshal(args) - if err != nil { - return map[string]any{"success": false, "_ERROR": fmt.Sprintf("email: marshal arguments: %v", err)}, nil - } - - out, err := e.tool.InvokableRun(ctx, string(argsJSON)) - if err != nil { - return map[string]any{"success": false, "_ERROR": err.Error()}, nil - } - var env struct { - OK bool `json:"ok"` - Error string `json:"_ERROR"` - } - if err := json.Unmarshal([]byte(out), &env); err != nil { - return map[string]any{"success": false, "_ERROR": fmt.Sprintf("email: parse result: %v", err)}, nil - } - if env.Error != "" { - return map[string]any{"success": false, "_ERROR": env.Error}, nil - } - return map[string]any{"success": env.OK}, nil -} - -// Stream mirrors Invoke as a single-chunk stream. -func (e *EmailComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { - out, err := e.Invoke(ctx, inputs) - if err != nil { - return nil, err - } - ch := make(chan map[string]any, 1) - ch <- out - close(ch) - return ch, nil -} - -func (e *EmailComponent) GetInputForm() map[string]any { - return map[string]any{ - "to_email": map[string]any{ - "name": "To", - "type": "line", - }, - "subject": map[string]any{ - "name": "Subject", - "type": "line", - }, - "cc_email": map[string]any{ - "name": "CC To", - "type": "line", - }, - } -} - -// Inputs returns the DSL input surface. -func (e *EmailComponent) Inputs() map[string]string { - return map[string]string{ - "to_email": "Recipient email address list.", - "cc_email": "Optional CC recipient list.", - "content": "Email body.", - "subject": "Email subject.", - "smtp_server": "SMTP server hostname.", - "smtp_port": "SMTP server port.", - "email": "Sender email address.", - "smtp_username": "SMTP username; defaults to sender email.", - "password": "SMTP password or authorization code.", - "sender_name": "Sender display name.", - } -} - -// Outputs returns the public output surface. -func (e *EmailComponent) Outputs() map[string]string { - return map[string]string{ - "success": "Whether the email was sent successfully.", - "_ERROR": "SMTP error message when sending fails.", - } -} - -func emailStringParam(params map[string]any, key string) string { - if v, ok := params[key].(string); ok { - return v - } - return "" -} - -func emailIntParam(params map[string]any, key string, fallback int) (int, error) { - switch v := params[key].(type) { - case nil: - return fallback, nil - case int: - return v, nil - case int64: - return int(v), nil - case float64: - return int(v), nil - case json.Number: - n, err := v.Int64() - return int(n), err - case string: - if strings.TrimSpace(v) == "" { - return fallback, nil - } - n, err := strconv.Atoi(strings.TrimSpace(v)) - if err != nil { - return 0, &ParamError{Field: key, Reason: "must be an integer"} - } - return n, nil - default: - return 0, &ParamError{Field: key, Reason: "must be an integer"} - } -} - -func firstString(inputs map[string]any, key, fallback string) string { - if v, ok := inputs[key].(string); ok { - return v - } - return fallback -} - -func firstNonEmpty(values ...string) string { - for _, v := range values { - if strings.TrimSpace(v) != "" { - return v - } - } - return "" -} - -func emailRecipients(toEmail, ccEmail string) []string { - recipients := splitEmailList(toEmail) - recipients = append(recipients, splitEmailList(ccEmail)...) - return recipients -} - -func splitEmailList(value string) []string { - parts := strings.FieldsFunc(value, func(r rune) bool { - return r == ',' || r == ';' || r == '\n' || r == '\r' - }) - out := make([]string, 0, len(parts)) - for _, part := range parts { - if s := strings.TrimSpace(part); s != "" { - out = append(out, s) - } - } - return out -} - -func stripEmailHeaderLineBreaks(value string) string { - return strings.NewReplacer("\r", "", "\n", "").Replace(value) -} - -func init() { - Register(componentNameEmail, NewEmailComponent) -} diff --git a/internal/agent/component/email_test.go b/internal/agent/component/email_test.go deleted file mode 100644 index a81a04a384..0000000000 --- a/internal/agent/component/email_test.go +++ /dev/null @@ -1,223 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "bufio" - "context" - "fmt" - "net" - "strings" - "testing" - "time" - - "ragflow/internal/agent/runtime" -) - -func TestEmailComponentRegisteredAndSends(t *testing.T) { - t.Parallel() - - ln, err := net.Listen("tcp", "127.0.0.1:0") - if err != nil { - t.Fatalf("Listen: %v", err) - } - defer ln.Close() - - var receivedData strings.Builder - done := make(chan struct{}) - go runEmailMockSMTP(t, ln, &receivedData, done) - - _, port, err := net.SplitHostPort(ln.Addr().String()) - if err != nil { - t.Fatalf("SplitHostPort: %v", err) - } - var portInt int - _, _ = fmt.Sscanf(port, "%d", &portInt) - - c, err := New(componentNameEmail, map[string]any{ - "smtp_server": "127.0.0.1", - "smtp_port": portInt, - "email": "alice@example.com", - "smtp_username": "", - "password": "", - "sender_name": "Alice", - "to_email": "bob@example.com", - "subject": "Build status", - "content": "Build succeeded.", - }) - if err != nil { - t.Fatalf("New Email: %v", err) - } - - state := runtime.NewCanvasState("run-email", "task-email") - state.Sys["date"] = "2026-07-14 03:04:05" - ctx := runtime.WithState(context.Background(), state) - - out, err := c.Invoke(ctx, nil) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if out["success"] != true { - t.Fatalf("success = %v, want true; out=%v", out["success"], out) - } - - select { - case <-done: - case <-time.After(2 * time.Second): - t.Fatal("mock SMTP server did not close in time") - } - - data := receivedData.String() - for _, want := range []string{"Subject: Build status", "bob@example.com", "Build succeeded."} { - if !strings.Contains(data, want) { - t.Fatalf("mock SMTP payload missing %q\n--- data ---\n%s\n---", want, data) - } - } -} - -func TestEmailComponentResolvesSysDateInSubject(t *testing.T) { - t.Parallel() - - ln, err := net.Listen("tcp", "127.0.0.1:0") - if err != nil { - t.Fatalf("Listen: %v", err) - } - defer ln.Close() - - var receivedData strings.Builder - done := make(chan struct{}) - go runEmailMockSMTP(t, ln, &receivedData, done) - - _, port, err := net.SplitHostPort(ln.Addr().String()) - if err != nil { - t.Fatalf("SplitHostPort: %v", err) - } - var portInt int - _, _ = fmt.Sscanf(port, "%d", &portInt) - - c, err := New(componentNameEmail, map[string]any{ - "smtp_server": "127.0.0.1", - "smtp_port": portInt, - "email": "alice@example.com", - "to_email": "bob@example.com", - "subject": "[noreply]{sys.date}", - "content": "body", - }) - if err != nil { - t.Fatalf("New Email: %v", err) - } - - state := runtime.NewCanvasState("run-email", "task-email") - state.Sys["date"] = "2026-07-14 03:04:05" - out, err := c.Invoke(runtime.WithState(context.Background(), state), nil) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if out["success"] != true { - t.Fatalf("success = %v, want true; out=%v", out["success"], out) - } - - select { - case <-done: - case <-time.After(2 * time.Second): - t.Fatal("mock SMTP server did not close in time") - } - - if strings.Contains(receivedData.String(), "{sys.date}") { - t.Fatalf("subject still contains unresolved sys.date\n--- data ---\n%s\n---", receivedData.String()) - } - if !strings.Contains(receivedData.String(), "Subject: [noreply]2026-07-14 03:04:05") { - t.Fatalf("subject did not resolve sys.date\n--- data ---\n%s\n---", receivedData.String()) - } -} - -func TestEmailComponentSoftFails(t *testing.T) { - t.Parallel() - - c, err := New(componentNameEmail, map[string]any{ - "smtp_port": 465, - "email": "alice@example.com", - "to_email": "bob@example.com", - "subject": "Subject", - }) - if err != nil { - t.Fatalf("New Email: %v", err) - } - out, err := c.Invoke(context.Background(), nil) - if err != nil { - t.Fatalf("Invoke returned hard error: %v", err) - } - if out["success"] != false { - t.Fatalf("success = %v, want false; out=%v", out["success"], out) - } - if _, ok := out["_ERROR"].(string); !ok { - t.Fatalf("_ERROR missing or not string: %v", out) - } -} - -func runEmailMockSMTP(t *testing.T, ln net.Listener, receivedData *strings.Builder, done chan<- struct{}) { - t.Helper() - defer close(done) - - conn, err := ln.Accept() - if err != nil { - return - } - defer conn.Close() - _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) - reader := bufio.NewReader(conn) - writer := bufio.NewWriter(conn) - _, _ = writer.WriteString("220 mock-smtp ready\r\n") - _ = writer.Flush() - - inData := false - for { - line, err := reader.ReadString('\n') - if err != nil { - return - } - up := strings.ToUpper(strings.TrimSpace(line)) - switch { - case strings.HasPrefix(up, "EHLO"), strings.HasPrefix(up, "HELO"): - _, _ = writer.WriteString("250-mock-smtp\r\n250-AUTH PLAIN\r\n250 OK\r\n") - _ = writer.Flush() - case strings.HasPrefix(up, "AUTH PLAIN"): - _, _ = writer.WriteString("235 Authentication successful\r\n") - _ = writer.Flush() - case strings.HasPrefix(up, "MAIL FROM:"), strings.HasPrefix(up, "RCPT TO:"): - _, _ = writer.WriteString("250 OK\r\n") - _ = writer.Flush() - case strings.HasPrefix(up, "DATA"): - _, _ = writer.WriteString("354 End data with .\r\n") - _ = writer.Flush() - inData = true - case inData && strings.TrimSpace(line) == ".": - _, _ = writer.WriteString("250 Queued\r\n") - _ = writer.Flush() - inData = false - case inData: - receivedData.WriteString(line) - case strings.HasPrefix(up, "QUIT"): - _, _ = writer.WriteString("221 Bye\r\n") - _ = writer.Flush() - return - default: - _, _ = writer.WriteString("250 OK\r\n") - _ = writer.Flush() - } - } -} diff --git a/internal/agent/component/exesql_component_test.go b/internal/agent/component/exesql_component_test.go deleted file mode 100644 index a069a863c2..0000000000 --- a/internal/agent/component/exesql_component_test.go +++ /dev/null @@ -1,100 +0,0 @@ -package component - -import ( - "context" - "encoding/json" - "errors" - "strings" - "testing" - - einotool "github.com/cloudwego/eino/components/tool" - - "ragflow/internal/agent/runtime" -) - -type exesqlInvokerStub struct { - arguments string - result string - err error -} - -func (s *exesqlInvokerStub) InvokableRun(_ context.Context, arguments string, _ ...einotool.Option) (string, error) { - s.arguments = arguments - return s.result, s.err -} - -func TestExeSQLComponentResolvesConfiguredSQL(t *testing.T) { - stub := &exesqlInvokerStub{ - result: `{"columns":["id","status"],"rows":[{"id":1,"status":"Completed"}]}`, - } - state := runtime.NewCanvasState("run", "task") - state.SetVar("Agent:SparklyMooseDivide", "content", "SELECT id FROM orders WHERE status = 'Completed'") - c := &exesqlComponent{ - inner: stub, - sql: "{Agent:SparklyMooseDivide@content}", - } - - out, err := c.Invoke(runtime.WithState(context.Background(), state), map[string]any{ - "content": "this must not be used as SQL", - }) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - var arguments map[string]any - if err := json.Unmarshal([]byte(stub.arguments), &arguments); err != nil { - t.Fatalf("decode tool arguments: %v", err) - } - if got := arguments["sql"]; got != "SELECT id FROM orders WHERE status = 'Completed'" { - t.Fatalf("sql argument = %#v, want resolved SQL", got) - } - if _, exists := arguments["content"]; exists { - t.Fatalf("tool arguments unexpectedly contain upstream content: %s", stub.arguments) - } - formalized, ok := out["formalized_content"].(string) - if !ok || !strings.Contains(formalized, "Completed") { - t.Fatalf("formalized_content = %#v, want rendered row", out["formalized_content"]) - } - jsonResult, ok := out["json"].([]any) - if !ok || len(jsonResult) != 1 { - t.Fatalf("json = %#v, want one statement result", out["json"]) - } -} - -func TestNewExeSQLComponentKeepsConfiguredSQL(t *testing.T) { - component, err := newExeSQLComponent(map[string]any{ - "db_type": "mysql", - "database": "demo", - "username": "root", - "host": "db.example.com", - "port": 3306, - "password": "secret", - "max_records": 100, - "sql": "{Agent:SparklyMooseDivide@content}", - }) - if err != nil { - t.Fatalf("newExeSQLComponent: %v", err) - } - exeSQL, ok := component.(*exesqlComponent) - if !ok { - t.Fatalf("component type = %T, want *exesqlComponent", component) - } - if exeSQL.sql != "{Agent:SparklyMooseDivide@content}" { - t.Fatalf("configured SQL = %q", exeSQL.sql) - } -} - -func TestExeSQLComponentPreservesToolErrorAsCanvasOutput(t *testing.T) { - stub := &exesqlInvokerStub{ - result: `{"_ERROR":"exesql: empty sql"}`, - err: errors.New("exesql: empty sql"), - } - c := &exesqlComponent{inner: stub, sql: ""} - - out, err := c.Invoke(context.Background(), map[string]any{"sql": ""}) - if err != nil { - t.Fatalf("Invoke returned a hard error: %v", err) - } - if got := out["_ERROR"]; got != "exesql: empty sql" { - t.Fatalf("_ERROR = %#v, want tool error", got) - } -} diff --git a/internal/agent/component/fixture_stubs.go b/internal/agent/component/fixture_stubs.go index c60d49ace5..d46f59cffa 100644 --- a/internal/agent/component/fixture_stubs.go +++ b/internal/agent/component/fixture_stubs.go @@ -14,20 +14,18 @@ // limitations under the License. // -// Package component — e2e fixture stubs and compat shims. +// Package component contains e2e fixture stubs used directly by tests. // // The test fixtures under internal/agent/dsl/testdata reference -// fixture-backed component names that are registered here: Retrieval, -// TavilySearch, ExeSQL, Generate, Answer, Iteration, -// IterationItem. Some names (for example TavilySearch) now route to -// production wrappers while their stubs remain available as direct test -// constructors. The fixture stub bodies are deliberately trivial — they +// fixture-backed component names that are registered here: Retrieval and its +// aliases, CodeExec, Generate, Answer, Iteration, and IterationItem. +// Production TavilySearch and ExeSQL nodes are registered through +// ToolBackedComponent; their fixture stubs are direct-only test constructors. The +// fixture stub bodies are deliberately trivial — they // echo a stable, template-friendly output shape and never call // the network or DB. The contract is "registered, non-panicking, // and produces outputs downstream templates can resolve", not -// "do something useful". The Universe A wrappers in -// universe_a_wrappers.go and the real production bodies in -// their own .go files replace these stubs in production paths. +// "do something useful". // // The fixture names were chosen by enumerating the component_name // values in the testdata fixtures (see the `examples` var in @@ -490,13 +488,9 @@ func (it *IterationItemStub) Outputs() map[string]string { // uniqueness), so accidental double-registration in a later refactor // surfaces as a panic at init time, not as a silent override. func init() { - // Primary registration: Retrieval and ExeSQL go through the - // Universe A delegation wrappers in universe_a_wrappers.go - // (real eino tool plumbing). The stubs remain available for - // unit tests that want to assert the "no service wired" path - // via a direct constructor. + // Retrieval still requires its specialized adapter. The stub remains a + // direct test constructor for the "no service wired" path. Register(componentNameRetrieval, newRetrievalComponent) - // The Python-side // The agent canvas uses both a PascalCase "SearchMyDataset" // and the original snake_case typo "search_my_dateset"; an // intermediate "search_my_dataset" form also exists in some @@ -509,22 +503,9 @@ func init() { Register("SearchMyDataset", newRetrievalComponent) Register("search_my_dataset", newRetrievalComponent) Register("search_my_dateset", newRetrievalComponent) - Register(componentNameTavilySearch, newTavilySearchComponent) - Register("TavilyExtract", newTavilyExtractComponent) - Register(componentNameExeSQL, newExeSQLComponent) Register(componentNameCodeExec, newCodeExecComponent) Register(componentNameGenerate, NewGenerateStub) Register(componentNameAnswer, NewAnswerStub) Register(componentNameIteration, NewIterationStub) Register(componentNameIterationItem, NewIterationItemStub) - Register("BGPT", newBGPTComponent) - Register("GitHub", newGitHubComponent) - Register("Wikipedia", newWikipediaComponent) - Register("DuckDuckGo", newDuckDuckGoComponent) - Register("KeenableSearch", newKeenableSearchComponent) - Register("Google", newGoogleComponent) - Register("GoogleScholar", newGoogleScholarComponent) - Register("ArXiv", newArxivComponent) - Register("PubMed", newPubMedComponent) - Register("YahooFinance", newYahooFinanceComponent) } diff --git a/internal/agent/component/github.go b/internal/agent/component/github.go deleted file mode 100644 index b338370426..0000000000 --- a/internal/agent/component/github.go +++ /dev/null @@ -1,230 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "crypto/sha1" - "encoding/json" - "fmt" - "math/big" - "strconv" - "strings" - "unicode/utf8" - - einotool "github.com/cloudwego/eino/components/tool" - - "ragflow/internal/agent/runtime" - agenttool "ragflow/internal/agent/tool" - "ragflow/internal/tokenizer" -) - -const githubPromptMaxTokens = 200000 - -type githubInvoker interface { - InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error) -} - -// githubComponent is the Canvas-facing GitHub repository search component. -// It mirrors agent/tools/github.py: query is a runtime input, while top_n is -// validated once from node parameters and defaults to ten. -type githubComponent struct { - inner githubInvoker -} - -func newGitHubComponent(params map[string]any) (Component, error) { - toolParams := make(map[string]any, 1) - for _, key := range []string{"top_n"} { - if value, ok := params[key]; ok { - toolParams[key] = value - } - } - inner, err := agenttool.BuildByName("github", toolParams) - if err != nil { - return nil, err - } - invoker, ok := inner.(githubInvoker) - if !ok { - return nil, fmt.Errorf("GitHub: tool does not implement InvokableRun") - } - return newGitHubComponentWithInvoker(invoker), nil -} - -func newGitHubComponentWithInvoker(inner githubInvoker) Component { - return &githubComponent{inner: inner} -} - -func (c *githubComponent) Name() string { return "GitHub" } - -func (c *githubComponent) Inputs() map[string]string { - return map[string]string{ - "query": "The search keywords to execute with GitHub. Use the most important terms and synonyms from the original request.", - } -} - -func (c *githubComponent) GetInputForm() map[string]any { - return map[string]any{ - "query": map[string]any{ - "type": "line", - "name": "Query", - }, - } -} - -func (c *githubComponent) Outputs() map[string]string { - return map[string]string{ - "formalized_content": "GitHub repositories formatted for downstream prompts.", - "json": "Raw GitHub repository items.", - } -} - -func (c *githubComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - query := stringParam(inputs["query"]) - if query == "" { - return map[string]any{"formalized_content": "", "json": []any{}}, nil - } - argsJSON, err := json.Marshal(map[string]any{"query": query}) - if err != nil { - return nil, fmt.Errorf("canvas: GitHub: encode query: %w", err) - } - out, err := c.inner.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - items := anySlice(decoded["results"]) - if existing, _ := decoded["_ERROR"].(string); strings.TrimSpace(existing) != "" { - return map[string]any{"formalized_content": "", "json": items, "_ERROR": existing}, nil - } - if err != nil { - if len(decoded) > 0 { - return map[string]any{"formalized_content": "", "json": items, "_ERROR": decoded["_ERROR"]}, nil - } - return nil, fmt.Errorf("canvas: GitHub: %w", err) - } - chunks, docAggs := buildGitHubReferences(items) - if state, _, stateErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); stateErr == nil && state != nil { - state.SetRetrievalReferences(chunks, docAggs) - } - return map[string]any{ - "formalized_content": renderGitHubReferences(chunks), - "json": items, - }, nil -} - -func (c *githubComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -func buildGitHubReferences(items []any) ([]map[string]any, []map[string]any) { - chunks := make([]map[string]any, 0, len(items)) - docAggs := make([]map[string]any, 0, len(items)) - for _, item := range items { - repository, ok := item.(map[string]any) - if !ok { - continue - } - content := truncateRunes(githubValueString(repository["description"])+"\n stars:"+githubValueString(repository["watchers"]), 10000) - if content == "" { - continue - } - documentID := strconv.FormatInt(githubHashInt(content, 100000000), 10) - title := githubValueString(repository["name"]) - url := githubValueString(repository["html_url"]) - displayID := strconv.FormatInt(githubHashInt(documentID, 500), 10) - chunks = append(chunks, map[string]any{ - "id": displayID, - "chunk_id": documentID, - "content": content, - "doc_id": documentID, - "document_id": documentID, - "docnm_kwd": title, - "document_name": title, - "similarity": 1, - "score": 1, - "url": url, - }) - docAggs = append(docAggs, map[string]any{ - "doc_name": title, - "doc_id": documentID, - "count": 1, - "url": url, - }) - } - return chunks, docAggs -} - -// renderGitHubReferences mirrors the Python -// "\n".join(kb_prompt({"chunks": chunks, "doc_aggs": aggs}, 200000, True)) -// layout for the GitHub chunks built above. -func renderGitHubReferences(chunks []map[string]any) string { - chunks = limitGitHubReferences(chunks, githubPromptMaxTokens) - blocks := make([]string, 0, len(chunks)) - for _, chunk := range chunks { - blocks = append(blocks, strings.Join([]string{ - "\nID: " + githubValueString(chunk["id"]), - "├── Title: " + githubValueString(chunk["docnm_kwd"]), - "├── URL: " + githubValueString(chunk["url"]), - "└── Content:\n" + githubValueString(chunk["content"]), - }, "\n")) - } - return strings.Join(blocks, "\n") -} - -// limitGitHubReferences mirrors kb_prompt's 200000-token guard. References -// are recorded before this step, as Python calls add_reference before -// formatting the prompt. -func limitGitHubReferences(chunks []map[string]any, maxTokens int) []map[string]any { - if maxTokens <= 0 { - return nil - } - usedTokens := 0 - for index, chunk := range chunks { - content := githubValueString(chunk["content"]) - if content == "" { - continue - } - usedTokens += tokenizer.NumTokensFromString(content) - if float64(maxTokens)*0.97 < float64(usedTokens) { - return chunks[:index+1] - } - } - return chunks -} - -func githubValueString(value any) string { - if value == nil { - return "None" - } - if boolean, ok := value.(bool); ok { - if boolean { - return "True" - } - return "False" - } - return fmt.Sprint(value) -} - -func githubHashInt(value string, modulus int64) int64 { - sum := sha1.Sum([]byte(value)) - number := new(big.Int).SetBytes(sum[:]) - return new(big.Int).Mod(number, big.NewInt(modulus)).Int64() -} - -func truncateRunes(value string, limit int) string { - if utf8.RuneCountInString(value) <= limit { - return value - } - return string([]rune(value)[:limit]) -} diff --git a/internal/agent/component/github_test.go b/internal/agent/component/github_test.go deleted file mode 100644 index 00863cd62c..0000000000 --- a/internal/agent/component/github_test.go +++ /dev/null @@ -1,165 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "encoding/json" - "strings" - "testing" - - einotool "github.com/cloudwego/eino/components/tool" - - "ragflow/internal/agent/canvas" - "ragflow/internal/agent/runtime" -) - -type fakeGitHubInvoker struct { - args map[string]any - out string - err error -} - -func (f *fakeGitHubInvoker) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) { - if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil { - return "", err - } - return f.out, f.err -} - -func TestGitHub_RegisteredFactory(t *testing.T) { - c, err := New("GitHub", map[string]any{"top_n": float64(10)}) - if err != nil { - t.Fatalf("New(GitHub) errored: %v", err) - } - if got := c.Name(); got != "GitHub" { - t.Fatalf("Name() = %q, want GitHub", got) - } - formGetter, ok := c.(interface{ GetInputForm() map[string]any }) - if !ok { - t.Fatal("GitHub component does not expose GetInputForm") - } - query, ok := formGetter.GetInputForm()["query"].(map[string]any) - if !ok || query["type"] != "line" { - t.Fatalf("GetInputForm query = %#v, want line input", query) - } - if _, ok := c.Outputs()["formalized_content"]; !ok { - t.Fatal("Outputs() missing formalized_content") - } - if _, ok := c.Outputs()["json"]; !ok { - t.Fatal("Outputs() missing json") - } -} - -func TestGitHub_CanvasBuildWorkflow(t *testing.T) { - c := &canvas.Canvas{ - Components: map[string]canvas.CanvasComponent{ - "begin_0": { - Obj: canvas.CanvasComponentObj{ComponentName: "Begin", Params: map[string]any{}}, - Downstream: []string{"github_0"}, - }, - "github_0": { - Obj: canvas.CanvasComponentObj{ComponentName: "GitHub", Params: map[string]any{"top_n": float64(10)}}, - Upstream: []string{"begin_0"}, - Downstream: []string{}, - }, - }, - Path: []string{"begin_0", "github_0"}, - } - if _, err := canvas.BuildWorkflow(context.Background(), c); err != nil { - t.Fatalf("BuildWorkflow with GitHub component: %v", err) - } -} - -func TestGitHub_InvokeMatchesPythonOutputsAndReferences(t *testing.T) { - fake := &fakeGitHubInvoker{out: `{"results":[{"name":"ragflow","html_url":"https://github.com/infiniflow/ragflow","description":"RAG engine","watchers":12000}]}`} - c := newGitHubComponentWithInvoker(fake) - state := runtime.NewCanvasState("run-github", "task-github") - ctx := runtime.WithState(context.Background(), state) - - out, err := c.Invoke(ctx, map[string]any{"query": "ragflow"}) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - if got := fake.args["query"]; got != "ragflow" { - t.Errorf("query = %v, want ragflow", got) - } - content, _ := out["formalized_content"].(string) - for _, want := range []string{"Title: ragflow", "URL: https://github.com/infiniflow/ragflow", "RAG engine\n stars:12000"} { - if !strings.Contains(content, want) { - t.Errorf("formalized_content missing %q: %q", want, content) - } - } - items := anySlice(out["json"]) - if len(items) != 1 { - t.Fatalf("json length = %d, want 1", len(items)) - } - chunks := state.GetRetrievalChunks() - if len(chunks) != 1 { - t.Fatalf("recorded reference chunks = %d, want 1", len(chunks)) - } - if chunks[0]["document_name"] != "ragflow" || chunks[0]["url"] != "https://github.com/infiniflow/ragflow" { - t.Errorf("reference chunk metadata = %#v", chunks[0]) - } - if chunks[0]["similarity"] != 1 { - t.Errorf("reference similarity = %v, want 1", chunks[0]["similarity"]) - } - encodedState, err := json.Marshal(state) - if err != nil { - t.Fatalf("marshal state: %v", err) - } - var statePayload struct { - Retrieval struct { - DocAggs map[string]any `json:"doc_aggs"` - } `json:"retrieval"` - } - if err := json.Unmarshal(encodedState, &statePayload); err != nil { - t.Fatalf("unmarshal state: %v", err) - } - if _, ok := statePayload.Retrieval.DocAggs["ragflow"]; !ok { - t.Fatalf("doc_aggs missing ragflow: %#v", statePayload.Retrieval.DocAggs) - } -} - -func TestGitHub_InvokeEmptyQueryMatchesPython(t *testing.T) { - fake := &fakeGitHubInvoker{} - c := newGitHubComponentWithInvoker(fake) - out, err := c.Invoke(context.Background(), map[string]any{"query": ""}) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - if got := out["formalized_content"]; got != "" { - t.Errorf("formalized_content = %v, want empty", got) - } - if items := anySlice(out["json"]); len(items) != 0 { - t.Errorf("json = %#v, want empty", items) - } - if fake.args != nil { - t.Fatal("GitHub tool was called for an empty query") - } -} - -func TestGitHub_LimitReferencesKeepsBoundaryChunk(t *testing.T) { - chunks := []map[string]any{ - {"content": "first repository description"}, - {"content": "second repository description"}, - } - limited := limitGitHubReferences(chunks, 1) - if len(limited) != 1 { - t.Fatalf("limited chunks = %d, want 1", len(limited)) - } -} diff --git a/internal/agent/component/google_component_test.go b/internal/agent/component/google_component_test.go deleted file mode 100644 index ecaa3e922b..0000000000 --- a/internal/agent/component/google_component_test.go +++ /dev/null @@ -1,71 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "strings" - "testing" -) - -func TestGoogleComponent_RegisteredAndInputForm(t *testing.T) { - c, err := New("Google", map[string]any{ - "api_key": "", - "country": "us", - "language": "en", - }) - if err != nil { - t.Fatalf("New(Google): %v", err) - } - if got := c.Name(); got != "Google" { - t.Fatalf("Name() = %q, want Google", got) - } - formGetter, ok := c.(interface{ GetInputForm() map[string]any }) - if !ok { - t.Fatal("Google component does not expose GetInputForm") - } - form := formGetter.GetInputForm() - if _, ok := form["q"]; !ok { - t.Fatalf("GetInputForm missing q: %+v", form) - } - if _, ok := c.Outputs()["formalized_content"]; !ok { - t.Fatal("Outputs() missing formalized_content") - } - if _, ok := c.Outputs()["json"]; !ok { - t.Fatal("Outputs() missing json") - } -} - -func TestGoogleComponent_MissingAPIKeyMatchesToolError(t *testing.T) { - c, err := New("Google", map[string]any{}) - if err != nil { - t.Fatalf("New(Google): %v", err) - } - out, err := c.Invoke(context.Background(), map[string]any{"q": "ragflow"}) - if err != nil { - t.Fatalf("Invoke returned error: %v", err) - } - if got, _ := out["_ERROR"].(string); !strings.Contains(got, "api_key") { - t.Fatalf("_ERROR = %q, want api_key error (out=%+v)", got, out) - } - if got, ok := out["formalized_content"].(string); !ok || got != "" { - t.Fatalf("formalized_content = %#v, want empty string", out["formalized_content"]) - } - if got := anySlice(out["json"]); len(got) != 0 { - t.Fatalf("json len = %d, want 0", len(got)) - } -} diff --git a/internal/agent/component/google_scholar_component_test.go b/internal/agent/component/google_scholar_component_test.go deleted file mode 100644 index fbb91df323..0000000000 --- a/internal/agent/component/google_scholar_component_test.go +++ /dev/null @@ -1,131 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "encoding/json" - "testing" - - einotool "github.com/cloudwego/eino/components/tool" -) - -type fakeGoogleScholarInvoker struct { - args map[string]any -} - -func (f *fakeGoogleScholarInvoker) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) { - if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil { - return "", err - } - return `{"results":[{"title":"Paper","link":"https://example.com","authors":"A Author","year":"2024","snippet":"Abstract"}]}`, nil -} - -func TestGoogleScholar_InvokePassesCanvasParams(t *testing.T) { - t.Parallel() - - fake := &fakeGoogleScholarInvoker{} - c := newGoogleScholarComponentWithInvoker(fake, nil) - - _, err := c.Invoke(context.Background(), map[string]any{ - "query": " retrieval augmented generation ", - "top_n": float64(7), - "sort_by": "date", - "year_low": float64(2020), - "year_high": float64(2024), - "patents": false, - }) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - - if got := fake.args["query"]; got != "retrieval augmented generation" { - t.Errorf("query arg = %v, want trimmed query", got) - } - if got := fake.args["top_n"]; got != float64(7) { - t.Errorf("top_n arg = %v, want 7", got) - } - if _, ok := fake.args["max_results"]; ok { - t.Fatalf("max_results should not be sent for GoogleScholar args: %#v", fake.args) - } - if got := fake.args["sort_by"]; got != "date" { - t.Errorf("sort_by arg = %v, want date", got) - } - if got := fake.args["year_low"]; got != float64(2020) { - t.Errorf("year_low arg = %v, want 2020", got) - } - if got := fake.args["year_high"]; got != float64(2024) { - t.Errorf("year_high arg = %v, want 2024", got) - } - if got := fake.args["patents"]; got != false { - t.Errorf("patents arg = %v, want false", got) - } -} - -func TestGoogleScholar_InvokeMergesNodeParams(t *testing.T) { - t.Parallel() - - fake := &fakeGoogleScholarInvoker{} - c := newGoogleScholarComponentWithInvoker(fake, map[string]any{ - "top_n": float64(20), - "sort_by": "date", - "patents": false, - }) - - _, err := c.Invoke(context.Background(), map[string]any{ - "query": "machine learning", - }) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - - if got := fake.args["top_n"]; got != float64(20) { - t.Errorf("top_n arg = %v, want 20 (from node params)", got) - } - if got := fake.args["sort_by"]; got != "date" { - t.Errorf("sort_by arg = %v, want date (from node params)", got) - } - if got := fake.args["patents"]; got != false { - t.Errorf("patents arg = %v, want false (from node params)", got) - } -} - -func TestGoogleScholar_InvokeInputsOverrideNodeParams(t *testing.T) { - t.Parallel() - - fake := &fakeGoogleScholarInvoker{} - c := newGoogleScholarComponentWithInvoker(fake, map[string]any{ - "top_n": float64(20), - "sort_by": "relevance", - }) - - _, err := c.Invoke(context.Background(), map[string]any{ - "query": "deep learning", - "top_n": float64(5), - "sort_by": "date", - }) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - - if got := fake.args["top_n"]; got != float64(5) { - t.Errorf("top_n arg = %v, want 5 (inputs override node params)", got) - } - if got := fake.args["sort_by"]; got != "date" { - t.Errorf("sort_by arg = %v, want date (inputs override node params)", got) - } -} diff --git a/internal/agent/component/keenable_component_test.go b/internal/agent/component/keenable_component_test.go deleted file mode 100644 index 03886bb28f..0000000000 --- a/internal/agent/component/keenable_component_test.go +++ /dev/null @@ -1,169 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "encoding/json" - "strings" - "testing" - - einotool "github.com/cloudwego/eino/components/tool" - - "ragflow/internal/agent/runtime" -) - -type fakeKeenableInvoker struct { - args map[string]any -} - -func (f *fakeKeenableInvoker) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) { - if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil { - return "", err - } - return `{"results":[{"title":"Keenable result","url":"https://example.com/item","description":"Fresh search result"}]}`, nil -} - -func TestKeenableSearch_RegisteredFactory(t *testing.T) { - t.Parallel() - - c, err := New("KeenableSearch", map[string]any{"api_key": "key-1", "mode": "realtime", "top_n": float64(3)}) - if err != nil { - t.Fatalf("New(KeenableSearch) errored: %v", err) - } - kc, ok := c.(*keenableSearchComponent) - if !ok { - t.Fatalf("New(KeenableSearch) returned %T, want *keenableSearchComponent", c) - } - if kc.apiKey != "key-1" { - t.Fatalf("apiKey = %q, want key-1", kc.apiKey) - } - if kc.mode != "realtime" { - t.Fatalf("mode = %q, want realtime", kc.mode) - } - if kc.topN != 3 { - t.Fatalf("topN = %d, want 3", kc.topN) - } - - formGetter, ok := c.(interface{ GetInputForm() map[string]any }) - if !ok { - t.Fatal("KeenableSearch component does not expose GetInputForm") - } - form := formGetter.GetInputForm() - query, ok := form["query"].(map[string]any) - if !ok { - t.Fatalf("GetInputForm()[query] has type %T, want map", form["query"]) - } - if query["type"] != "line" { - t.Fatalf("GetInputForm()[query][type] = %v, want line", query["type"]) - } - site, ok := form["site"].(map[string]any) - if !ok { - t.Fatalf("GetInputForm()[site] has type %T, want map", form["site"]) - } - if site["type"] != "line" { - t.Fatalf("GetInputForm()[site][type] = %v, want line", site["type"]) - } - if _, ok := c.Outputs()["formalized_content"]; !ok { - t.Fatal("Outputs() missing formalized_content") - } - if _, ok := c.Outputs()["json"]; !ok { - t.Fatal("Outputs() missing json") - } -} - -func TestKeenableSearch_InvokeAdaptsCanvasInputsAndOutputs(t *testing.T) { - t.Parallel() - - fake := &fakeKeenableInvoker{} - c := newKeenableSearchComponentWithInvoker(fake, map[string]any{ - "mode": "pro", - "top_n": 5, - "site": "example.com", - }) - - state := runtime.NewCanvasState("run-keenable", "task-keenable") - ctx := runtime.WithState(context.Background(), state) - - out, err := c.Invoke(ctx, map[string]any{ - "query": " agent search ", - "mode": "realtime", - "top_n": float64(2), - }) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - - if got := fake.args["query"]; got != "agent search" { - t.Errorf("query arg = %v, want trimmed query", got) - } - if got := fake.args["mode"]; got != "realtime" { - t.Errorf("mode arg = %v, want realtime", got) - } - if got := fake.args["top_n"]; got != float64(2) && got != 2 { - t.Errorf("top_n arg = %v, want 2", got) - } - if got := fake.args["site"]; got != "example.com" { - t.Errorf("site arg = %v, want default site", got) - } - - formalized, _ := out["formalized_content"].(string) - for _, want := range []string{"Keenable result", "https://example.com/item", "Fresh search result"} { - if !strings.Contains(formalized, want) { - t.Errorf("formalized_content missing %q: %s", want, formalized) - } - } - - results, ok := out["json"].([]any) - if !ok { - t.Fatalf("json output has type %T, want []any", out["json"]) - } - if len(results) != 1 { - t.Fatalf("json output length = %d, want 1", len(results)) - } - - reference := state.GetRetrievalReference() - chunks, _ := reference["chunks"].([]any) - if len(chunks) != 1 { - t.Fatalf("reference chunks length = %d, want 1", len(chunks)) - } - chunk, _ := chunks[0].(map[string]any) - if chunk["document_name"] != "Keenable result" || chunk["url"] != "https://example.com/item" { - t.Fatalf("reference chunk metadata = %#v", chunk) - } - docAggs, _ := reference["doc_aggs"].([]any) - if len(docAggs) != 1 { - t.Fatalf("reference doc_aggs length = %d, want 1", len(docAggs)) - } -} - -func TestKeenableSearch_InvokeEmptyQueryReturnsEmptyPayload(t *testing.T) { - t.Parallel() - - c := newKeenableSearchComponentWithInvoker(&fakeKeenableInvoker{}, nil) - out, err := c.Invoke(context.Background(), map[string]any{"query": " "}) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - if got := out["formalized_content"]; got != "" { - t.Errorf("formalized_content = %v, want empty string", got) - } - results, ok := out["json"].([]any) - if !ok || len(results) != 0 { - t.Fatalf("json output = %#v, want empty []any", out["json"]) - } -} diff --git a/internal/agent/component/production_chain_fixes_test.go b/internal/agent/component/production_chain_fixes_test.go index 4b2796e206..6f4ac1fc25 100644 --- a/internal/agent/component/production_chain_fixes_test.go +++ b/internal/agent/component/production_chain_fixes_test.go @@ -59,12 +59,8 @@ func (s *codeExecSandboxRecorder) ExecuteCode(_ context.Context, req agenttool.S // host/port/password/top_n, no db_type) into the tool's required shape // (db_type/database/username/host/port/password/max_records). // -// Without the translator, NewExeSQLConnParams would reject the v1 -// shape with "missing required connection params (db_type/host/ -// database/username)" and every legacy v1 ExeSQL canvas would fail -// at buildNodeBody time. With the translator in place, the v1 shape -// compiles cleanly (db_type defaults to "mysql"; port is coerced from -// float64; top_n is mapped to max_records). +// The tool factory accepts the v1 shape directly: db_type defaults to mysql, +// JSON numeric ports are accepted, and top_n becomes max_records. func TestExeSQL_V1DSLParamsAccepted(t *testing.T) { t.Parallel() @@ -90,39 +86,8 @@ func TestExeSQL_V1DSLParamsAccepted(t *testing.T) { t.Errorf("ExeSQL c.Name() = %q, want %q", got, componentNameExeSQL) } - // translateExeSQLParamsToToolShape should also be directly - // testable as a pure function: the same shape, the same result. - got := translateExeSQLParamsToToolShape(v1Params) - if got["db_type"] != "mysql" { - t.Errorf("translated db_type = %v, want %q", got["db_type"], "mysql") - } - if v, ok := got["port"].(int); !ok || v != 3306 { - t.Errorf("translated port = %v (%T), want int 3306", got["port"], got["port"]) - } - if v, ok := got["max_records"].(int); !ok || v != 50 { - t.Errorf("translated max_records = %v (%T), want int 50", got["max_records"], got["max_records"]) - } - if _, ok := got["top_n"]; ok { - t.Errorf("translated map should drop top_n (mapped to max_records)") - } - - // Idempotency: a second pass must not double-default. - got2 := translateExeSQLParamsToToolShape(got) - if got2["db_type"] != "mysql" { - t.Errorf("idempotent db_type = %v, want %q", got2["db_type"], "mysql") - } - if v, ok := got2["port"].(int); !ok || v != 3306 { - t.Errorf("idempotent port = %v (%T), want int 3306", got2["port"], got2["port"]) - } - - // Explicit override wins: passing db_type=postgres must be - // preserved through the translator. - override := translateExeSQLParamsToToolShape(map[string]any{ - "db_type": "postgres", - "host": "10.0.0.1", - }) - if override["db_type"] != "postgres" { - t.Errorf("override db_type = %v, want %q", override["db_type"], "postgres") + if _, ok := c.(*ToolBackedComponent); !ok { + t.Fatalf("New(ExeSQL) returned %T, want *ToolBackedComponent", c) } } diff --git a/internal/agent/component/pubmed_component_test.go b/internal/agent/component/pubmed_component_test.go deleted file mode 100644 index fdad06faf9..0000000000 --- a/internal/agent/component/pubmed_component_test.go +++ /dev/null @@ -1,150 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "encoding/json" - "errors" - "strings" - "testing" - - einotool "github.com/cloudwego/eino/components/tool" -) - -type fakePubMedInvoker struct { - args map[string]any - err error - out string -} - -func (f *fakePubMedInvoker) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) { - if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil { - return "", err - } - if f.out != "" || f.err != nil { - return f.out, f.err - } - return `{"results":[{"title":"Deep learning for retrieval augmented generation","url":"https://pubmed.ncbi.nlm.nih.gov/12345678","content":"Title: Deep learning for retrieval augmented generation\nAuthors: Furqan Khan, Jane Smith\nJournal: Nature Machine Intelligence\nVolume: 10\nIssue: 2\nPages: 101-110\nDOI: 10.1000/example.doi\nAbstract: A short abstract."}]}`, nil -} - -func TestPubMed_RegisteredFactory(t *testing.T) { - t.Parallel() - - c, err := New("PubMed", map[string]any{ - "top_n": 8, - "email": "node@example.com", - "outputs": map[string]any{"formalized_content": map[string]any{}}, - "setups": map[string]any{"query": "configured query"}, - }) - if err != nil { - t.Fatalf("New(PubMed) errored: %v", err) - } - if got := c.Name(); got != "PubMed" { - t.Fatalf("Name() = %q, want PubMed", got) - } - formGetter, ok := c.(interface{ GetInputForm() map[string]any }) - if !ok { - t.Fatal("PubMed component does not expose GetInputForm") - } - form := formGetter.GetInputForm() - if len(form) != 1 { - t.Fatalf("GetInputForm size = %d, want 1", len(form)) - } - query, ok := form["query"].(map[string]any) - if !ok { - t.Fatalf("GetInputForm()[query] has type %T, want map", form["query"]) - } - if query["type"] != "line" { - t.Fatalf("GetInputForm()[query][type] = %v, want line", query["type"]) - } - if _, ok := c.Outputs()["formalized_content"]; !ok { - t.Fatal("Outputs() missing formalized_content") - } - if _, ok := c.Outputs()["json"]; !ok { - t.Fatal("Outputs() missing json") - } -} - -func TestPubMed_InvokeOnlyPassesQuery(t *testing.T) { - t.Parallel() - - fake := &fakePubMedInvoker{} - c := newPubMedComponentWithInvoker(fake) - out, err := c.Invoke(context.Background(), map[string]any{ - "query": " retrieval augmented generation ", - "top_n": float64(8), - "email": "ignored@example.com", - "unused": true, - }) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - if got := fake.args["query"]; got != "retrieval augmented generation" { - t.Fatalf("query arg = %v, want trimmed query", got) - } - if len(fake.args) != 1 { - t.Fatalf("runtime args = %#v, want only query", fake.args) - } - formalized, _ := out["formalized_content"].(string) - for _, want := range []string{"ID: 0", "Title: Deep learning for retrieval augmented generation", "URL: https://pubmed.ncbi.nlm.nih.gov/12345678", "Content:", "Abstract: A short abstract."} { - if !strings.Contains(formalized, want) { - t.Fatalf("formalized_content missing %q: %s", want, formalized) - } - } - results, ok := out["json"].([]any) - if !ok || len(results) != 1 { - t.Fatalf("json output = %#v, want one result", out["json"]) - } -} - -func TestPubMed_InvokeEmptyQueryReturnsEmptyPayload(t *testing.T) { - t.Parallel() - - c := newPubMedComponentWithInvoker(&fakePubMedInvoker{}) - out, err := c.Invoke(context.Background(), map[string]any{"query": " "}) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - if got := out["formalized_content"]; got != "" { - t.Fatalf("formalized_content = %v, want empty string", got) - } - results, ok := out["json"].([]any) - if !ok || len(results) != 0 { - t.Fatalf("json output = %#v, want empty []any", out["json"]) - } -} - -func TestPubMed_InvokeSurfacesToolErrorEnvelope(t *testing.T) { - t.Parallel() - - fake := &fakePubMedInvoker{ - out: `{"results":[],"_ERROR":"upstream down"}`, - err: errors.New("boom"), - } - c := newPubMedComponentWithInvoker(fake) - out, err := c.Invoke(context.Background(), map[string]any{"query": "pubmed"}) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - if got := out["_ERROR"]; got != "upstream down" { - t.Fatalf("_ERROR = %v, want upstream down", got) - } - if got := out["formalized_content"]; got != "" { - t.Fatalf("formalized_content = %v, want empty string", got) - } -} diff --git a/internal/agent/component/searxng.go b/internal/agent/component/searxng.go deleted file mode 100644 index 8c8855535e..0000000000 --- a/internal/agent/component/searxng.go +++ /dev/null @@ -1,252 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "crypto/sha1" - "encoding/json" - "fmt" - "regexp" - "strconv" - "strings" - - einotool "github.com/cloudwego/eino/components/tool" - - "ragflow/internal/agent/runtime" - agenttool "ragflow/internal/agent/tool" - "ragflow/internal/tokenizer" -) - -const searxngPromptTokenLimit = 200000 - -var searxngDataImagePattern = regexp.MustCompile(`!?\[[a-z]+\]\(data:image/png;base64,[ 0-9A-Za-z/_=+\-]+\)`) - -var searxngNewlinePattern = regexp.MustCompile(`\n+`) - -type searxngInvoker interface { - InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error) -} - -type searxngComponent struct { - inner searxngInvoker -} - -func newSearXNGComponent(params map[string]any) (Component, error) { - toolParams := make(map[string]any, 2) - for _, key := range []string{"top_n", "searxng_url"} { - if value, ok := params[key]; ok { - toolParams[key] = value - } - } - inner, err := agenttool.BuildByName("searxng", toolParams) - if err != nil { - return nil, err - } - invoker, ok := inner.(searxngInvoker) - if !ok { - return nil, fmt.Errorf("SearXNG: tool does not implement InvokableRun") - } - return newSearXNGComponentWithInvoker(invoker), nil -} - -func newSearXNGComponentWithInvoker(inner searxngInvoker) Component { - return &searxngComponent{inner: inner} -} - -func (c *searxngComponent) Name() string { return "SearXNG" } - -func (c *searxngComponent) Inputs() map[string]string { - return map[string]string{ - "query": "The search keywords to execute with SearXNG.", - "searxng_url": "The base URL of the SearXNG instance.", - } -} - -func (c *searxngComponent) Outputs() map[string]string { - return map[string]string{ - "formalized_content": "Rendered SearXNG references for downstream LLM prompts.", - "json": "Raw SearXNG result list.", - } -} - -func (c *searxngComponent) GetInputForm() map[string]any { - return map[string]any{ - "query": map[string]any{ - "name": "Query", - "type": "line", - }, - "searxng_url": map[string]any{ - "name": "SearXNG URL", - "type": "line", - "placeholder": "http://localhost:4000", - }, - } -} - -func (c *searxngComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - query := stringParam(inputs["query"]) - if strings.TrimSpace(query) == "" { - return map[string]any{"formalized_content": "", "json": []any{}}, nil - } - args := map[string]any{"query": query} - if searxngURL, ok := inputs["searxng_url"].(string); ok { - args["searxng_url"] = searxngURL - } - argsJSON, err := json.Marshal(args) - if err != nil { - return nil, fmt.Errorf("canvas: SearXNG: encode inputs: %w", err) - } - - out, invokeErr := c.inner.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - results := anySlice(decoded["results"]) - if message, _ := decoded["_ERROR"].(string); strings.TrimSpace(message) != "" { - return map[string]any{ - "formalized_content": "", - "json": results, - "_ERROR": message, - }, nil - } - if invokeErr != nil { - return nil, fmt.Errorf("canvas: SearXNG: %w", invokeErr) - } - - chunks, docAggs := buildSearXNGReferences(results) - if state, _, stateErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); stateErr == nil && state != nil { - state.SetRetrievalReferences(chunks, docAggs) - } - return map[string]any{ - "formalized_content": renderSearXNGReferences(chunks, searxngPromptTokenLimit), - "json": results, - }, nil -} - -func (c *searxngComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -func buildSearXNGReferences(results []any) ([]map[string]any, []map[string]any) { - chunks := make([]map[string]any, 0, len(results)) - docAggs := make([]map[string]any, 0, len(results)) - for _, result := range results { - item, ok := result.(map[string]any) - if !ok { - continue - } - content, _ := item["content"].(string) - if content == "" { - continue - } - content = searxngDataImagePattern.ReplaceAllString(content, "") - runes := []rune(content) - if len(runes) > 10000 { - content = string(runes[:10000]) - } - if content == "" { - continue - } - - documentID := strconv.Itoa(hashSearXNGString(content, 100000000)) - displayID := strconv.Itoa(hashSearXNGString(documentID, 500)) - title := searxngText(item["title"]) - resultURL := searxngText(item["url"]) - chunks = append(chunks, map[string]any{ - "id": displayID, - "chunk_id": documentID, - "content": content, - "doc_id": documentID, - "docnm_kwd": title, - "document_id": documentID, - "document_name": title, - "dataset_id": nil, - "image_id": nil, - "positions": nil, - "url": resultURL, - "similarity": 1, - "vector_similarity": nil, - "term_similarity": nil, - "row_id": nil, - "doc_type": nil, - "document_metadata": nil, - }) - docAggs = append(docAggs, map[string]any{ - "doc_name": title, - "doc_id": documentID, - "count": 1, - "url": resultURL, - }) - } - return chunks, docAggs -} - -func renderSearXNGReferences(chunks []map[string]any, maxTokens int) string { - if len(chunks) == 0 { - return "" - } - usedTokens := 0 - blocks := make([]string, 0, len(chunks)) - for _, chunk := range chunks { - content := searxngText(chunk["content"]) - if content == "" { - continue - } - usedTokens += tokenizer.NumTokensFromString(content) - var block strings.Builder - fmt.Fprintf(&block, "\nID: %s", searxngText(chunk["id"])) - if title := searxngPromptField(chunk["document_name"]); title != "" { - fmt.Fprintf(&block, "\n├── Title: %s", title) - } - if resultURL := searxngPromptField(chunk["url"]); resultURL != "" { - fmt.Fprintf(&block, "\n├── URL: %s", resultURL) - } - block.WriteString("\n└── Content:\n") - block.WriteString(content) - blocks = append(blocks, block.String()) - if maxTokens > 0 && float64(maxTokens)*0.97 < float64(usedTokens) { - break - } - } - return strings.Join(blocks, "\n") -} - -func searxngText(value any) string { - if value == nil { - return "" - } - if text, ok := value.(string); ok { - return text - } - return fmt.Sprint(value) -} - -func searxngPromptField(value any) string { - return searxngNewlinePattern.ReplaceAllString(searxngText(value), " ") -} - -func hashSearXNGString(value string, modulus int) int { - digest := sha1.Sum([]byte(value)) - result := 0 - for _, part := range digest { - result = (result*256 + int(part)) % modulus - } - return result -} - -func init() { - Register("SearXNG", newSearXNGComponent) -} diff --git a/internal/agent/component/searxng_test.go b/internal/agent/component/searxng_test.go deleted file mode 100644 index 9c029089ec..0000000000 --- a/internal/agent/component/searxng_test.go +++ /dev/null @@ -1,250 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "encoding/json" - "errors" - "math" - "strconv" - "strings" - "testing" - - einotool "github.com/cloudwego/eino/components/tool" - - agentcanvas "ragflow/internal/agent/canvas" - "ragflow/internal/agent/runtime" - "ragflow/internal/tokenizer" -) - -type fakeSearXNGInvoker struct { - args map[string]any - calls int - out string - err error -} - -func (f *fakeSearXNGInvoker) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) { - f.calls++ - if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil { - return "", err - } - return f.out, f.err -} - -func TestSearXNGRegisteredFactoryMatchesPythonSurface(t *testing.T) { - t.Parallel() - - component, err := New("SearXNG", map[string]any{ - "top_n": "10", - "searxng_url": "http://localhost:4000", - "outputs": map[string]any{"formalized_content": map[string]any{}}, - "setups": map[string]any{"query": "configured query"}, - }) - if err != nil { - t.Fatalf("New(SearXNG): %v", err) - } - if component.Name() != "SearXNG" { - t.Fatalf("Name = %q, want SearXNG", component.Name()) - } - if _, ok := component.Inputs()["query"]; !ok { - t.Fatal("Inputs missing query") - } - if _, ok := component.Inputs()["searxng_url"]; !ok { - t.Fatal("Inputs missing searxng_url") - } - for _, output := range []string{"formalized_content", "json"} { - if _, ok := component.Outputs()[output]; !ok { - t.Fatalf("Outputs missing %s", output) - } - } - form := component.(*searxngComponent).GetInputForm() - if len(form) != 2 { - t.Fatalf("GetInputForm size = %d, want 2", len(form)) - } - query := form["query"].(map[string]any) - if query["name"] != "Query" || query["type"] != "line" { - t.Fatalf("query form = %#v", query) - } - serverURL := form["searxng_url"].(map[string]any) - if serverURL["name"] != "SearXNG URL" || serverURL["type"] != "line" || serverURL["placeholder"] != "http://localhost:4000" { - t.Fatalf("searxng_url form = %#v", serverURL) - } -} - -func TestSearXNGInvokePreservesRawJSONPromptAndReferences(t *testing.T) { - t.Parallel() - - content := "RAGFlow content ![img](data:image/png;base64,AAAA) remains" - fake := &fakeSearXNGInvoker{out: `{"results":[{"title":"RAGFlow\nDocs","url":"https://ragflow.io","content":` + mustJSONText(t, content) + `,"engine":"bing","score":0.9},{"title":"Empty","url":"https://example.com","content":""}]}`} - component := newSearXNGComponentWithInvoker(fake) - state := runtime.NewCanvasState("run", "task") - ctx := runtime.WithState(context.Background(), state) - out, err := component.Invoke(ctx, map[string]any{ - "query": " ragflow ", - "searxng_url": "http://localhost:4000", - }) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if fake.args["query"] != " ragflow " || fake.args["searxng_url"] != "http://localhost:4000" { - t.Fatalf("runtime args = %#v", fake.args) - } - results, ok := out["json"].([]any) - if !ok || len(results) != 2 { - t.Fatalf("json = %#v, want two raw results", out["json"]) - } - first := results[0].(map[string]any) - if first["engine"] != "bing" || first["score"] != float64(0.9) { - t.Fatalf("raw result lost fields: %#v", first) - } - - cleaned := "RAGFlow content remains" - documentID := hashSearXNGString(cleaned, 100000000) - referenceID := hashSearXNGString(strconv.Itoa(documentID), 500) - if documentID != 93760153 || referenceID != 491 { - t.Fatalf("hash parity = %d/%d, want Python 93760153/491", documentID, referenceID) - } - formalized := out["formalized_content"].(string) - for _, want := range []string{ - "ID: " + strconv.Itoa(referenceID), - "Title: RAGFlow Docs", - "URL: https://ragflow.io", - "Content:\n" + cleaned, - } { - if !strings.Contains(formalized, want) { - t.Fatalf("formalized_content missing %q: %s", want, formalized) - } - } - - chunks := state.GetRetrievalChunks() - if len(chunks) != 1 { - t.Fatalf("retrieval chunks = %#v, want one non-empty-content chunk", chunks) - } - if chunks[0]["document_id"] != strconv.Itoa(documentID) || chunks[0]["content"] != cleaned { - t.Fatalf("retrieval chunk = %#v", chunks[0]) - } - if chunks[0]["id"] != strconv.Itoa(referenceID) { - t.Fatalf("chunk id = %#v, want displayed reference ID", chunks[0]["id"]) - } - if chunks[0]["chunk_id"] != strconv.Itoa(documentID) || chunks[0]["doc_id"] != strconv.Itoa(documentID) { - t.Fatalf("raw document IDs = %#v, want Python content hash", chunks[0]) - } - if chunks[0]["similarity"] != 1 { - t.Fatalf("similarity = %#v, want 1", chunks[0]["similarity"]) - } - aggs := state.GetRetrievalDocAggs() - if len(aggs) != 1 || aggs["RAGFlow\nDocs"]["doc_id"] != strconv.Itoa(documentID) { - t.Fatalf("doc_aggs = %#v", aggs) - } -} - -func TestSearXNGInvokeEmptyQuerySkipsTool(t *testing.T) { - t.Parallel() - - fake := &fakeSearXNGInvoker{} - component := newSearXNGComponentWithInvoker(fake) - out, err := component.Invoke(context.Background(), map[string]any{"query": " "}) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if fake.calls != 0 { - t.Fatalf("tool calls = %d, want 0", fake.calls) - } - if out["formalized_content"] != "" { - t.Fatalf("formalized_content = %#v", out["formalized_content"]) - } - if results, ok := out["json"].([]any); !ok || len(results) != 0 { - t.Fatalf("json = %#v, want empty []any", out["json"]) - } -} - -func TestSearXNGInvokePreservesErrorEnvelope(t *testing.T) { - t.Parallel() - - fake := &fakeSearXNGInvoker{ - out: `{"results":[],"_ERROR":"Network error: upstream down"}`, - err: errors.New("upstream down"), - } - component := newSearXNGComponentWithInvoker(fake) - out, err := component.Invoke(context.Background(), map[string]any{"query": "ragflow"}) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if out["_ERROR"] != "Network error: upstream down" || out["formalized_content"] != "" { - t.Fatalf("output = %#v", out) - } -} - -func TestSearXNGPromptBoundaryMatchesPython(t *testing.T) { - t.Parallel() - - first := "small content" - second := strings.Repeat("cross the token budget ", 50) - third := "must not render" - chunks := []map[string]any{ - {"id": "1", "content": first}, - {"id": "2", "content": second}, - {"id": "3", "content": third}, - } - firstTokens := tokenizer.NumTokensFromString(first) - maxTokens := int(math.Ceil(float64(firstTokens)/0.97)) + 1 - rendered := renderSearXNGReferences(chunks, maxTokens) - if !strings.Contains(rendered, "ID: 1") || !strings.Contains(rendered, "ID: 2") { - t.Fatalf("crossing chunk must be included like Python: %s", rendered) - } - if strings.Contains(rendered, "ID: 3") { - t.Fatalf("chunk after budget crossing must be excluded: %s", rendered) - } -} - -func TestSearXNGBuildWorkflowUsesPythonComponentName(t *testing.T) { - canvas := &agentcanvas.Canvas{ - Components: map[string]agentcanvas.CanvasComponent{ - "begin_0": { - Obj: agentcanvas.CanvasComponentObj{ComponentName: "Begin", Params: map[string]any{}}, - Downstream: []string{"searxng_0"}, - }, - "searxng_0": { - Obj: agentcanvas.CanvasComponentObj{ComponentName: "SearXNG", Params: map[string]any{ - "top_n": "10", - "searxng_url": "http://localhost:4000", - }}, - Upstream: []string{"begin_0"}, - Downstream: []string{"message_0"}, - }, - "message_0": { - Obj: agentcanvas.CanvasComponentObj{ComponentName: "Message", Params: map[string]any{}}, - Upstream: []string{"searxng_0"}, - }, - }, - Path: []string{"begin_0", "searxng_0", "message_0"}, - } - if _, err := agentcanvas.BuildWorkflow(context.Background(), canvas); err != nil { - t.Fatalf("BuildWorkflow with SearXNG: %v", err) - } -} - -func mustJSONText(t *testing.T, value string) string { - t.Helper() - raw, err := json.Marshal(value) - if err != nil { - t.Fatalf("json.Marshal: %v", err) - } - return string(raw) -} diff --git a/internal/agent/component/tavily_component_test.go b/internal/agent/component/tavily_component_test.go deleted file mode 100644 index efa3101365..0000000000 --- a/internal/agent/component/tavily_component_test.go +++ /dev/null @@ -1,102 +0,0 @@ -package component - -import ( - "context" - "testing" -) - -func TestTavilySearch_RegisteredRealComponentWithInputForm(t *testing.T) { - c, err := New("TavilySearch", nil) - if err != nil { - t.Fatalf("New(TavilySearch): %v", err) - } - if _, ok := c.(*tavilySearchComponent); !ok { - t.Fatalf("New(TavilySearch) returned %T, want *tavilySearchComponent", c) - } - formGetter, ok := c.(interface{ GetInputForm() map[string]any }) - if !ok { - t.Fatal("TavilySearch component does not expose GetInputForm") - } - form := formGetter.GetInputForm() - query, ok := form["query"].(map[string]any) - if !ok { - t.Fatalf("GetInputForm()[query] has type %T, want map", form["query"]) - } - if query["type"] != "line" || query["name"] != "Query" { - t.Fatalf("GetInputForm()[query] = %#v, want Query line input", query) - } -} - -func TestTavilyExtract_RegisteredWithInputForm(t *testing.T) { - c, err := New("TavilyExtract", nil) - if err != nil { - t.Fatalf("New(TavilyExtract): %v", err) - } - if _, ok := c.(*tavilyExtractComponent); !ok { - t.Fatalf("New(TavilyExtract) returned %T, want *tavilyExtractComponent", c) - } - formGetter, ok := c.(interface{ GetInputForm() map[string]any }) - if !ok { - t.Fatal("TavilyExtract component does not expose GetInputForm") - } - form := formGetter.GetInputForm() - urls, ok := form["urls"].(map[string]any) - if !ok { - t.Fatalf("GetInputForm()[urls] has type %T, want map", form["urls"]) - } - if urls["type"] != "line" || urls["name"] != "URLs" { - t.Fatalf("GetInputForm()[urls] = %#v, want URLs line input", urls) - } -} - -func TestTavilySearch_StoresAPIKeyAndInjectsWhenInputOmitsIt(t *testing.T) { - c, err := New("TavilySearch", map[string]any{"api_key": "tvly-stored"}) - if err != nil { - t.Fatalf("New(TavilySearch): %v", err) - } - tc, ok := c.(*tavilySearchComponent) - if !ok { - t.Fatalf("New(TavilySearch) returned %T, want *tavilySearchComponent", c) - } - if tc.apiKey != "tvly-stored" { - t.Fatalf("stored apiKey = %q, want tvly-stored", tc.apiKey) - } - - inputs := map[string]any{"query": ""} - if _, err := tc.Invoke(context.Background(), inputs); err != nil { - t.Fatalf("Invoke with empty query errored: %v", err) - } - if got := inputs["api_key"]; got != "tvly-stored" { - t.Fatalf("injected api_key = %v, want tvly-stored", got) - } -} - -func TestTavilySearch_DoesNotOverrideCallerAPIKey(t *testing.T) { - c, err := New("TavilySearch", map[string]any{"api_key": "tvly-stored"}) - if err != nil { - t.Fatalf("New(TavilySearch): %v", err) - } - tc := c.(*tavilySearchComponent) - - inputs := map[string]any{"query": "", "api_key": "tvly-call"} - if _, err := tc.Invoke(context.Background(), inputs); err != nil { - t.Fatalf("Invoke with empty query errored: %v", err) - } - if got := inputs["api_key"]; got != "tvly-call" { - t.Fatalf("api_key = %v, want caller key", got) - } -} - -func TestTavilyExtract_StoresAPIKey(t *testing.T) { - c, err := New("TavilyExtract", map[string]any{"api_key": "tvly-extract"}) - if err != nil { - t.Fatalf("New(TavilyExtract): %v", err) - } - tc, ok := c.(*tavilyExtractComponent) - if !ok { - t.Fatalf("New(TavilyExtract) returned %T, want *tavilyExtractComponent", c) - } - if tc.apiKey != "tvly-extract" { - t.Fatalf("stored apiKey = %q, want tvly-extract", tc.apiKey) - } -} diff --git a/internal/agent/component/tool_component.go b/internal/agent/component/tool_component.go new file mode 100644 index 0000000000..79e6be2beb --- /dev/null +++ b/internal/agent/component/tool_component.go @@ -0,0 +1,131 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package component + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "ragflow/internal/agent/runtime" + agenttool "ragflow/internal/agent/tool" +) + +// ToolBackedComponent is the single Canvas adapter for tools that implement +// agenttool.ToolComponent. +type ToolBackedComponent struct { + name string + tool agenttool.ToolComponent + spec agenttool.ComponentSpec +} + +func newToolComponentFactory(componentName, toolName string) Factory { + return func(params map[string]any) (Component, error) { + base, err := agenttool.BuildByName(toolName, params) + if err != nil { + return nil, err + } + componentTool, ok := base.(agenttool.ToolComponent) + if !ok { + return nil, fmt.Errorf("%s: tool %q does not implement ToolComponent", componentName, toolName) + } + return &ToolBackedComponent{ + name: componentName, + tool: componentTool, + spec: componentTool.ComponentSpec(), + }, nil + } +} + +func (c *ToolBackedComponent) Name() string { return c.name } + +func (c *ToolBackedComponent) Inputs() map[string]string { return c.spec.Inputs } + +func (c *ToolBackedComponent) Outputs() map[string]string { return c.spec.Outputs } + +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) { + argsJSON, err := json.Marshal(inputs) + if err != nil { + return nil, fmt.Errorf("canvas: %s: encode inputs: %w", c.name, err) + } + + raw, invokeErr := c.tool.InvokableRun(ctx, string(argsJSON)) + decoded := parseToolEnvelope(raw) + if rawValue, invalid := decoded["_raw"]; invalid { + if invokeErr != nil { + return nil, fmt.Errorf("canvas: %s: %w", c.name, invokeErr) + } + return nil, fmt.Errorf("canvas: %s: invalid tool result: %v", c.name, rawValue) + } + if existing, _ := decoded["_ERROR"].(string); strings.TrimSpace(existing) != "" { + outputs := c.tool.BuildComponentOutputs(decoded) + if outputs == nil { + outputs = make(map[string]any, 1) + } + outputs["_ERROR"] = existing + return outputs, nil + } + if invokeErr != nil { + return nil, fmt.Errorf("canvas: %s: %w", c.name, invokeErr) + } + + if builder, ok := c.tool.(agenttool.ReferenceBuilder); ok { + chunks, docAggs := builder.BuildReferences(ctx, decoded) + if state, _, stateErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); stateErr == nil && state != nil { + state.SetRetrievalReferences(chunks, docAggs) + } + } + return c.tool.BuildComponentOutputs(decoded), nil +} + +func (c *ToolBackedComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { + return nil, nil +} + +var toolComponentRegistrations = []struct { + componentName string + toolName string +}{ + {componentName: "GitHub", toolName: "github"}, + {componentName: "BGPT", toolName: "bgpt"}, + {componentName: "ArXiv", toolName: "arxiv"}, + {componentName: "DuckDuckGo", toolName: "duckduckgo"}, + {componentName: "Email", toolName: "email"}, + {componentName: "ExeSQL", toolName: "execute_sql"}, + {componentName: "Google", toolName: "google"}, + {componentName: "GoogleScholar", toolName: "google_scholar"}, + {componentName: "KeenableSearch", toolName: "keenable"}, + {componentName: "PubMed", toolName: "pubmed"}, + {componentName: "SearXNG", toolName: "searxng"}, + {componentName: "TavilySearch", toolName: "tavily"}, + {componentName: "TavilyExtract", toolName: "tavily_extract"}, + {componentName: "WenCai", toolName: "wencai"}, + {componentName: "Wikipedia", toolName: "wikipedia"}, + {componentName: "YahooFinance", toolName: "yahoo_finance"}, +} + +func init() { + for _, registration := range toolComponentRegistrations { + Register( + registration.componentName, + newToolComponentFactory(registration.componentName, registration.toolName), + ) + } +} diff --git a/internal/agent/component/tool_component_test.go b/internal/agent/component/tool_component_test.go new file mode 100644 index 0000000000..c50736fd85 --- /dev/null +++ b/internal/agent/component/tool_component_test.go @@ -0,0 +1,514 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package component + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "net/url" + "reflect" + "strings" + "testing" + + einotool "github.com/cloudwego/eino/components/tool" + + "ragflow/internal/agent/canvas" + "ragflow/internal/agent/runtime" + agenttool "ragflow/internal/agent/tool" +) + +type fakeToolAdapter struct { + args map[string]any + referenceEnvelope map[string]any + outputEnvelope map[string]any + calls int + events []string + out string + err error +} + +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (f roundTripperFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return f(request) +} + +func (f *fakeToolAdapter) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) { + f.calls++ + if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil { + return "", err + } + if f.out != "" || f.err != nil { + return f.out, f.err + } + return `{"results":[{"title":"RAGFlow","url":"https://ragflow.io","content":"RAG engine"}]}`, nil +} + +func (f *fakeToolAdapter) ComponentSpec() agenttool.ComponentSpec { + return agenttool.ComponentSpec{ + Inputs: map[string]string{"query": "Search query."}, + Outputs: map[string]string{"formalized_content": "Rendered results.", "json": "Raw results."}, + InputForm: map[string]any{"query": map[string]any{"name": "Query", "type": "line"}}, + } +} + +func (f *fakeToolAdapter) BuildReferences(_ context.Context, envelope map[string]any) ([]map[string]any, []map[string]any) { + f.referenceEnvelope = envelope + f.events = append(f.events, "references") + return []map[string]any{{"chunk_id": "1", "content": "RAG engine", "docnm_kwd": "RAGFlow"}}, + []map[string]any{{"doc_name": "RAGFlow", "doc_id": "1", "count": 1}} +} + +func (f *fakeToolAdapter) BuildComponentOutputs(envelope map[string]any) map[string]any { + f.outputEnvelope = envelope + f.events = append(f.events, "render") + formalizedContent := "" + results := anySlice(envelope["results"]) + if len(results) > 0 { + if result, ok := results[0].(map[string]any); ok { + formalizedContent, _ = result["content"].(string) + } + } + return map[string]any{ + "json": results, + "formalized_content": formalizedContent, + "tool_metadata": envelope["tool_metadata"], + } +} + +func TestToolBackedComponentRegisteredGitHubFactory(t *testing.T) { + c, err := New("GitHub", map[string]any{ + "top_n": float64(10), + "query": "runtime query", + "outputs": map[string]any{"json": map[string]any{}}, + "setups": map[string]any{"query": "configured query"}, + }) + if err != nil { + t.Fatalf("New(GitHub): %v", err) + } + if _, ok := c.(*ToolBackedComponent); !ok { + t.Fatalf("New(GitHub) returned %T, want *ToolBackedComponent", c) + } + if c.Name() != "GitHub" { + t.Fatalf("Name = %q, want GitHub", c.Name()) + } + form := c.(interface{ GetInputForm() map[string]any }).GetInputForm() + if query, ok := form["query"].(map[string]any); !ok || query["type"] != "line" { + t.Fatalf("query input form = %#v, want line", form["query"]) + } +} + +func TestToolBackedComponentRegisteredFactories(t *testing.T) { + tests := []struct { + name string + toolName string + params map[string]any + outputKey string + inputKey string + }{ + { + name: "ArXiv", + toolName: "ArXiv", + params: map[string]any{"top_n": float64(3), "sort_by": "relevance", "outputs": map[string]any{"json": map[string]any{}}}, + outputKey: "json", + inputKey: "query", + }, + { + name: "BGPT", + toolName: "BGPT", + params: map[string]any{"api_key": "stored-key", "top_n": float64(3), "outputs": map[string]any{"json": map[string]any{}}}, + outputKey: "json", + inputKey: "query", + }, + { + name: "KeenableSearch", + toolName: "KeenableSearch", + params: map[string]any{"api_key": "stored-key", "mode": "realtime", "top_n": float64(3), "outputs": map[string]any{"json": map[string]any{}}}, + outputKey: "json", + inputKey: "query", + }, + { + name: "PubMed", + toolName: "PubMed", + params: map[string]any{"top_n": float64(3), "email": "node@example.com", "outputs": map[string]any{"json": map[string]any{}}}, + outputKey: "json", + inputKey: "query", + }, + { + name: "Google", + toolName: "Google", + params: map[string]any{"api_key": "stored-key", "country": "cn", "language": "en", "outputs": map[string]any{"json": map[string]any{}}}, + outputKey: "json", + inputKey: "q", + }, + { + name: "ExeSQL", + toolName: "ExeSQL", + params: map[string]any{ + "database": "demo", "username": "root", "host": "db.example.com", "port": float64(3306), "password": "secret", + "top_n": float64(50), "outputs": map[string]any{"json": map[string]any{}}, + }, + outputKey: "json", + inputKey: "sql", + }, + { + name: "GoogleScholar", + toolName: "GoogleScholar", + params: map[string]any{"top_n": float64(3), "sort_by": "relevance", "outputs": map[string]any{"json": map[string]any{}}}, + outputKey: "json", + inputKey: "query", + }, + { + name: "DuckDuckGo", + toolName: "DuckDuckGo", + params: map[string]any{"top_n": float64(3), "channel": "news", "outputs": map[string]any{"json": map[string]any{}}}, + outputKey: "json", + inputKey: "query", + }, + { + name: "Email", + toolName: "Email", + params: map[string]any{ + "smtp_server": "smtp.example.com", "smtp_port": float64(465), "email": "sender@example.com", + "password": "secret", "sender_name": "Sender", "outputs": map[string]any{"success": map[string]any{}}, + }, + outputKey: "success", + inputKey: "to_email", + }, + { + name: "SearXNG", + toolName: "SearXNG", + params: map[string]any{"top_n": "10", "searxng_url": "https://searx.example.com", "outputs": map[string]any{"json": map[string]any{}}}, + outputKey: "json", + inputKey: "query", + }, + { + name: "WenCai", + toolName: "WenCai", + params: map[string]any{"top_n": float64(20), "query_type": "stock", "outputs": map[string]any{"report": map[string]any{}}}, + outputKey: "report", + inputKey: "query", + }, + { + name: "Wikipedia", + toolName: "Wikipedia", + params: map[string]any{"top_n": float64(3), "language": "en", "outputs": map[string]any{"json": map[string]any{}}}, + outputKey: "json", + inputKey: "query", + }, + { + name: "YahooFinance", + toolName: "YahooFinance", + params: map[string]any{"outputs": map[string]any{"report": map[string]any{}}}, + outputKey: "report", + inputKey: "stock_code", + }, + { + name: "TavilyExtract", + toolName: "TavilyExtract", + params: map[string]any{"api_key": "stored-key", "extract_depth": "advanced", "format": "text", "outputs": map[string]any{"json": map[string]any{}}}, + outputKey: "json", + inputKey: "urls", + }, + { + name: "TavilySearch", + toolName: "TavilySearch", + params: map[string]any{"api_key": "stored-key", "search_depth": "advanced", "max_results": float64(3), "outputs": map[string]any{"json": map[string]any{}}}, + outputKey: "json", + inputKey: "query", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c, err := New(tt.toolName, tt.params) + if err != nil { + t.Fatalf("New(%s): %v", tt.toolName, err) + } + if _, ok := c.(*ToolBackedComponent); !ok { + t.Fatalf("New(%s) returned %T, want *ToolBackedComponent", tt.toolName, c) + } + if c.Name() != tt.toolName { + t.Fatalf("Name = %q, want %q", c.Name(), tt.toolName) + } + if _, ok := c.Inputs()[tt.inputKey]; !ok { + t.Fatalf("Inputs missing %q: %#v", tt.inputKey, c.Inputs()) + } + if _, ok := c.Outputs()[tt.outputKey]; !ok { + t.Fatalf("Outputs missing %q: %#v", tt.outputKey, c.Outputs()) + } + }) + } +} + +func TestToolBackedComponentWenCaiInvoke(t *testing.T) { + c, err := New("WenCai", map[string]any{"top_n": float64(20), "query_type": "stock"}) + if err != nil { + t.Fatalf("New(WenCai): %v", err) + } + out, err := c.Invoke(context.Background(), map[string]any{ + "query": "商业航天", + "unrelated": "ignored by the tool parameter struct", + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if out["report"] != "" { + t.Fatalf("outputs = %#v, want empty report", out) + } +} + +func TestToolBackedComponentRegisteredBuildWorkflow(t *testing.T) { + for _, componentName := range []string{"ArXiv", "BGPT", "DuckDuckGo", "Email", "Google", "GoogleScholar", "KeenableSearch", "PubMed", "SearXNG", "WenCai", "TavilyExtract", "TavilySearch", "Wikipedia", "YahooFinance"} { + t.Run(componentName, func(t *testing.T) { + c := &canvas.Canvas{ + Components: map[string]canvas.CanvasComponent{ + "begin_0": { + Obj: canvas.CanvasComponentObj{ComponentName: "Begin", Params: map[string]any{}}, + Downstream: []string{"tool_0"}, + }, + "tool_0": { + Obj: canvas.CanvasComponentObj{ComponentName: componentName, Params: map[string]any{}}, + Upstream: []string{"begin_0"}, + }, + }, + Path: []string{"begin_0", "tool_0"}, + } + if _, err := canvas.BuildWorkflow(context.Background(), c); err != nil { + t.Fatalf("BuildWorkflow with %s: %v", componentName, err) + } + }) + } +} + +func TestToolBackedComponentCanvasBuildWorkflow(t *testing.T) { + c := &canvas.Canvas{ + Components: map[string]canvas.CanvasComponent{ + "begin_0": { + Obj: canvas.CanvasComponentObj{ComponentName: "Begin", Params: map[string]any{}}, + Downstream: []string{"github_0"}, + }, + "github_0": { + Obj: canvas.CanvasComponentObj{ComponentName: "GitHub", Params: map[string]any{"top_n": float64(10)}}, + Upstream: []string{"begin_0"}, + }, + }, + Path: []string{"begin_0", "github_0"}, + } + if _, err := canvas.BuildWorkflow(context.Background(), c); err != nil { + t.Fatalf("BuildWorkflow with GitHub: %v", err) + } +} + +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{ + "query": "ragflow", + "top_n": float64(10), + "outputs": map[string]any{"json": map[string]any{}}, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if fake.args["query"] != "ragflow" { + t.Fatalf("runtime args = %#v", fake.args) + } + if fake.args["top_n"] != float64(10) || fake.args["outputs"] == nil { + t.Fatalf("generic component did not pass all inputs: %#v", fake.args) + } + if !reflect.DeepEqual(fake.events, []string{"references", "render"}) { + t.Fatalf("post-process order = %#v, want references then render", fake.events) + } + if fake.referenceEnvelope["tool_metadata"] == nil || fake.outputEnvelope["tool_metadata"] == nil { + t.Fatalf("complete tool envelope was not passed to post-processors: references=%#v outputs=%#v", fake.referenceEnvelope, fake.outputEnvelope) + } + if _, exists := fake.outputEnvelope["chunks"]; exists { + t.Fatalf("generic component injected references into the tool envelope: %#v", fake.outputEnvelope) + } + if metadata, ok := out["tool_metadata"].(map[string]any); !ok || metadata["request_id"] != "request-1" { + t.Fatalf("tool-specific envelope fields were lost: %#v", out) + } + if out["formalized_content"] != "RAG engine" { + t.Fatalf("formalized_content = %#v", out["formalized_content"]) + } + if len(state.GetRetrievalChunks()) != 1 { + t.Fatalf("retrieval chunks = %#v", state.GetRetrievalChunks()) + } +} + +func TestToolBackedComponentReturnsErrorEnvelopeWithoutReferences(t *testing.T) { + fake := &fakeToolAdapter{ + out: `{"results":[],"_ERROR":"rate limited"}`, + err: errors.New("rate limited"), + } + 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"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if out["_ERROR"] != "rate limited" || out["formalized_content"] != "" { + t.Fatalf("error outputs = %#v", out) + } + if results, ok := out["json"].([]any); !ok || len(results) != 0 { + t.Fatalf("error json output = %#v, want an empty array", out["json"]) + } + if len(state.GetRetrievalChunks()) != 0 { + t.Fatalf("error path recorded references: %#v", state.GetRetrievalChunks()) + } +} + +func TestToolBackedComponentGitHubIntegration(t *testing.T) { + serverCalls := 0 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + serverCalls++ + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"items":[{"name":"ragflow","html_url":"https://github.com/infiniflow/ragflow","description":"RAG engine","watchers":12000,"private":false}]}`)) + })) + defer server.Close() + target, err := url.Parse(server.URL) + if err != nil { + t.Fatalf("parse test server URL: %v", err) + } + helper := agenttool.NewHTTPHelper().WithClient(&http.Client{Transport: roundTripperFunc(func(request *http.Request) (*http.Response, error) { + cloned := request.Clone(request.Context()) + cloned.URL.Scheme = target.Scheme + cloned.URL.Host = target.Host + return http.DefaultTransport.RoundTrip(cloned) + })}) + 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": ""}) + 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"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + rendered, _ := out["formalized_content"].(string) + if !strings.Contains(rendered, "Title: ragflow") || !strings.Contains(rendered, "RAG engine\n stars:12000") { + t.Fatalf("formalized_content = %q", rendered) + } + results, ok := out["json"].([]any) + if !ok || len(results) != 1 || results[0].(map[string]any)["private"] != false { + t.Fatalf("raw json results = %#v", out["json"]) + } + chunks := state.GetRetrievalChunks() + if len(chunks) != 1 || chunks[0]["document_name"] != "ragflow" { + t.Fatalf("recorded references = %#v", chunks) + } +} + +func TestToolBackedComponentTavilyIntegration(t *testing.T) { + serverCalls := 0 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + serverCalls++ + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"results":[{"title":"RAGFlow","url":"https://ragflow.io","raw_content":"RAG article","content":"fallback","score":0.8,"custom":"preserved"}]}`)) + })) + defer server.Close() + target, err := url.Parse(server.URL) + if err != nil { + t.Fatalf("parse test server URL: %v", err) + } + helper := agenttool.NewHTTPHelper().WithClient(&http.Client{Transport: roundTripperFunc(func(request *http.Request) (*http.Response, error) { + cloned := request.Clone(request.Context()) + cloned.URL.Scheme = target.Scheme + cloned.URL.Host = target.Host + return http.DefaultTransport.RoundTrip(cloned) + })}) + 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": ""}) + 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"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + rendered, _ := out["formalized_content"].(string) + if !strings.Contains(rendered, "Title: RAGFlow") || !strings.Contains(rendered, "RAG article") { + t.Fatalf("formalized_content = %q", rendered) + } + results, ok := out["json"].([]any) + if !ok || len(results) != 1 || results[0].(map[string]any)["custom"] != "preserved" { + t.Fatalf("raw json results = %#v", out["json"]) + } + chunks := state.GetRetrievalChunks() + if len(chunks) != 1 || chunks[0]["document_name"] != "RAGFlow" || chunks[0]["similarity"] != float64(0.8) { + t.Fatalf("recorded references = %#v", chunks) + } +} + +func TestToolBackedComponentYahooFinanceIntegration(t *testing.T) { + serverCalls := 0 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + serverCalls++ + if symbols := request.URL.Query().Get("symbols"); symbols != "AAPL" { + t.Errorf("symbols = %q", symbols) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"quoteResponse":{"result":[{"symbol":"AAPL","regularMarketPrice":189.5,"currency":"USD"}],"error":null}}`)) + })) + defer server.Close() + target, err := url.Parse(server.URL) + if err != nil { + t.Fatalf("parse test server URL: %v", err) + } + helper := agenttool.NewHTTPHelper().WithClient(&http.Client{Transport: roundTripperFunc(func(request *http.Request) (*http.Response, error) { + cloned := request.Clone(request.Context()) + cloned.URL.Scheme = target.Scheme + cloned.URL.Host = target.Host + return http.DefaultTransport.RoundTrip(cloned) + })}) + yahoo := agenttool.NewYahooFinanceToolWith(helper) + component := &ToolBackedComponent{name: "YahooFinance", tool: yahoo, spec: yahoo.ComponentSpec()} + + empty, err := component.Invoke(context.Background(), map[string]any{"stock_code": ""}) + if err != nil { + t.Fatalf("Invoke(empty stock_code): %v", err) + } + if serverCalls != 0 || empty["report"] != "" { + t.Fatalf("empty result = %#v, server calls = %d", empty, serverCalls) + } + + out, err := component.Invoke(context.Background(), map[string]any{"stock_code": "AAPL"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + report, ok := out["report"].(string) + if !ok || !strings.Contains(report, "# Information:") || !strings.Contains(report, "| symbol | AAPL |") { + t.Fatalf("report = %#v", out["report"]) + } + if serverCalls != 1 { + t.Fatalf("server calls = %d, want 1", serverCalls) + } +} diff --git a/internal/agent/component/tool_dispatch_test.go b/internal/agent/component/tool_dispatch_test.go index 49fd735ee5..165d265368 100644 --- a/internal/agent/component/tool_dispatch_test.go +++ b/internal/agent/component/tool_dispatch_test.go @@ -94,7 +94,7 @@ func TestAgent_GoogleToolDSLParamsLoading(t *testing.T) { }, }) form := c.GetInputForm() - googleForm, ok := form["google"].(map[string]any) + googleForm, ok := form["google_search"].(map[string]any) if !ok { t.Fatalf("GetInputForm missing google tool form: %+v", form) } diff --git a/internal/agent/component/universe_a_wrappers.go b/internal/agent/component/universe_a_wrappers.go index 15c3bac30a..2d3a272342 100644 --- a/internal/agent/component/universe_a_wrappers.go +++ b/internal/agent/component/universe_a_wrappers.go @@ -14,18 +14,9 @@ // limitations under the License. // -// Package component Universe A delegation wrappers. Canvas-facing components that -// delegate to their corresponding Universe B eino tool -// implementations. The delegation pattern keeps the canvas -// scheduler's Component contract thin and the eino tool's -// InvokableRun interface as the actual implementation seam. -// -// Primary registration: TavilySearch, Retrieval (incl. the -// Python-typo SearchMyDataset alias), and ExeSQL all delegate to -// the real Universe B tools. fixture_stubs.go's init() wires the -// registry to these wrappers; the legacy stub-only path is -// preserved as NewRetrievalStub / NewExeSQLStub for unit tests -// that want to assert the "no service wired" state directly. +// Package component contains the remaining specialized Canvas adapters for +// Retrieval and CodeExec. Tools with a standard Canvas surface are registered +// through ToolBackedComponent instead. package component import ( @@ -34,12 +25,9 @@ import ( "errors" "fmt" "regexp" - "sort" "strconv" "strings" - einotool "github.com/cloudwego/eino/components/tool" - "ragflow/internal/agent/runtime" agenttool "ragflow/internal/agent/tool" "ragflow/internal/common" @@ -50,758 +38,6 @@ import ( "gorm.io/gorm" ) -// tavilySearchComponent delegates to internal/agent/tool/TavilyTool. -// The underlying tool makes a real HTTP call; the wrapper is the -// canvas-facing surface. -type tavilySearchComponent struct { - inner *agenttool.TavilyTool - apiKey string -} - -func newTavilySearchComponent(params map[string]any) (Component, error) { - return &tavilySearchComponent{ - inner: agenttool.NewTavilyTool(), - apiKey: stringParam(params["api_key"]), - }, nil -} - -func (c *tavilySearchComponent) Name() string { return "TavilySearch" } - -func (c *tavilySearchComponent) Inputs() map[string]string { - return map[string]string{ - "query": "Search query.", - "api_key": "Tavily API key (overrides TAVILY_API_KEY env var).", - "max_results": "Maximum results to return (default 5).", - "search_depth": "\"basic\" (default) or \"advanced\".", - } -} - -func (c *tavilySearchComponent) Outputs() map[string]string { - return map[string]string{ - "formalized_content": "Rendered search results for downstream LLM prompts.", - "json": "Raw result list (url, title, content, score).", - } -} - -func (c *tavilySearchComponent) GetInputForm() map[string]any { - return map[string]any{ - "query": map[string]any{ - "name": "Query", - "type": "line", - }, - } -} - -func (c *tavilySearchComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - if c.apiKey != "" && stringParam(inputs["api_key"]) == "" { - inputs["api_key"] = c.apiKey - } - if strings.TrimSpace(stringParam(inputs["query"])) == "" { - return map[string]any{"formalized_content": "", "json": []any{}}, nil - } - argsJSON, _ := json.Marshal(inputs) - out, err := c.inner.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - if err != nil { - if len(decoded) > 0 { - return map[string]any{ - "formalized_content": "", - "json": []any{}, - "_ERROR": decoded["_ERROR"], - }, nil - } - return nil, fmt.Errorf("canvas: TavilySearch: %w", err) - } - results := anySlice(decoded["results"]) - return map[string]any{ - "formalized_content": renderTavilySearchResults(results), - "json": results, - }, nil -} - -func (c *tavilySearchComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -// googleComponent wraps internal/agent/tool/GoogleTool for canvas execution and -// adapts the tool envelope to the Google component outputs. -type googleComponent struct { - inner *agenttool.GoogleTool - params map[string]any -} - -func newGoogleComponent(params map[string]any) (Component, error) { - cloned := make(map[string]any, len(params)) - for k, v := range params { - cloned[k] = v - } - return &googleComponent{inner: agenttool.NewGoogleTool(), params: cloned}, nil -} - -func (c *googleComponent) Name() string { return "Google" } - -func (c *googleComponent) Inputs() map[string]string { - return map[string]string{ - "q": "Search query.", - "api_key": "SerpApi API key.", - "start": "Result offset.", - "num": "Maximum number of results.", - "country": "Google country code.", - "language": "Google language code.", - } -} - -func (c *googleComponent) GetInputForm() map[string]any { - return agenttool.NewGoogleTool().InputForm() -} - -func (c *googleComponent) Outputs() map[string]string { - return map[string]string{ - "formalized_content": "Rendered search results for downstream LLM prompts.", - "json": "Raw Google organic result list.", - } -} - -func (c *googleComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - merged := make(map[string]any, len(c.params)+len(inputs)+2) - for k, v := range c.params { - merged[k] = v - } - for k, v := range inputs { - merged[k] = v - } - if _, ok := merged["q"]; !ok { - if query, ok := merged["query"]; ok { - merged["q"] = query - } - } - if _, ok := merged["num"]; !ok { - if maxResults, ok := merged["max_results"]; ok { - merged["num"] = maxResults - } - } - - argsJSON, _ := json.Marshal(merged) - out, err := c.inner.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - results := anySlice(decoded["organic_results"]) - if len(results) == 0 { - results = anySlice(decoded["results"]) - } - formalized := renderGoogleResults(results) - if existing, _ := decoded["_ERROR"].(string); strings.TrimSpace(existing) != "" { - return map[string]any{"formalized_content": formalized, "json": results, "_ERROR": existing}, nil - } - if err != nil { - if len(decoded) > 0 { - return map[string]any{"formalized_content": formalized, "json": results, "_ERROR": decoded["_ERROR"]}, nil - } - return nil, fmt.Errorf("canvas: Google: %w", err) - } - return map[string]any{"formalized_content": formalized, "json": results}, nil -} - -func (c *googleComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -func renderGoogleResults(results []any) string { - if len(results) == 0 { - return "" - } - blocks := make([]string, 0, len(results)) - for _, item := range results { - m, ok := item.(map[string]any) - if !ok { - continue - } - title := strings.TrimSpace(stringParam(m["title"])) - link := strings.TrimSpace(stringParam(m["link"])) - content := strings.TrimSpace(stringParam(m["snippet"])) - if content == "" { - content = strings.TrimSpace(googleAboutDescription(m["about_this_result"])) - } - if content == "" { - continue - } - lines := []string{} - if title != "" { - lines = append(lines, "Title: "+title) - } - if link != "" { - lines = append(lines, "URL: "+link) - } - lines = append(lines, "Content: "+content) - blocks = append(blocks, strings.Join(lines, "\n")) - } - return strings.Join(blocks, "\n\n") -} - -func googleAboutDescription(v any) string { - about, ok := v.(map[string]any) - if !ok { - return "" - } - source, ok := about["source"].(map[string]any) - if !ok { - return "" - } - return stringParam(source["description"]) -} - -// tavilyExtractComponent delegates to internal/agent/tool/TavilyExtractTool. -type tavilyExtractComponent struct { - inner *agenttool.TavilyExtractTool - apiKey string -} - -func newTavilyExtractComponent(params map[string]any) (Component, error) { - return &tavilyExtractComponent{ - inner: agenttool.NewTavilyExtractTool(), - apiKey: stringParam(params["api_key"]), - }, nil -} - -func (c *tavilyExtractComponent) Name() string { return "TavilyExtract" } - -func (c *tavilyExtractComponent) Inputs() map[string]string { - return map[string]string{ - "urls": "URLs to extract content from. Accepts a comma-separated string or array.", - "api_key": "Tavily API key (overrides TAVILY_API_KEY env var).", - "extract_depth": "\"basic\" (default) or \"advanced\".", - "format": "\"markdown\" (default) or \"text\".", - } -} - -func (c *tavilyExtractComponent) GetInputForm() map[string]any { - return map[string]any{ - "urls": map[string]any{ - "name": "URLs", - "type": "line", - }, - } -} - -func (c *tavilyExtractComponent) Outputs() map[string]string { - return map[string]string{ - "json": "Raw Tavily Extract results.", - } -} - -func (c *tavilyExtractComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - if c.apiKey != "" && stringParam(inputs["api_key"]) == "" { - inputs["api_key"] = c.apiKey - } - - argsJSON, _ := json.Marshal(inputs) - out, err := c.inner.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - if err != nil { - if len(decoded) > 0 { - return map[string]any{"json": []any{}, "_ERROR": decoded["_ERROR"]}, nil - } - return nil, fmt.Errorf("canvas: TavilyExtract: %w", err) - } - return map[string]any{"json": anySlice(decoded["results"])}, nil -} - -func (c *tavilyExtractComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -// bgptInvoker is the subset of BGPTTool used by the canvas wrapper. -type bgptInvoker interface { - InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error) -} - -type duckDuckGoInvoker interface { - InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error) -} - -type keenableInvoker interface { - InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error) -} - -// bgptComponent delegates to internal/agent/tool/BGPTTool and adapts -// the tool envelope to the BGPT canvas output contract. -type bgptComponent struct { - inner bgptInvoker - apiKey string -} - -func newBGPTComponent(params map[string]any) (Component, error) { - return newBGPTComponentWithInvoker(agenttool.NewBGPTTool(), stringParam(params["api_key"])), nil -} - -func newBGPTComponentWithInvoker(inner bgptInvoker, apiKey ...string) Component { - c := &bgptComponent{inner: inner} - if len(apiKey) > 0 { - c.apiKey = apiKey[0] - } - return c -} - -func (c *bgptComponent) Name() string { return "BGPT" } - -func (c *bgptComponent) Inputs() map[string]string { - return map[string]string{ - "query": "Scientific search query.", - "api_key": "Optional BGPT API key.", - "days_back": "Optional recency filter in days.", - "top_n": "Maximum number of results.", - } -} - -func (c *bgptComponent) GetInputForm() map[string]any { - return map[string]any{ - "query": map[string]any{ - "name": "Query", - "type": "line", - }, - } -} - -func (c *bgptComponent) Outputs() map[string]string { - return map[string]string{ - "formalized_content": "Rendered scientific paper evidence for downstream LLM prompts.", - "json": "Raw BGPT result list.", - } -} - -func (c *bgptComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - if c.apiKey != "" && stringParam(inputs["api_key"]) == "" { - inputs["api_key"] = c.apiKey - } - - query := strings.TrimSpace(stringParam(inputs["query"])) - if query == "" { - return map[string]any{"formalized_content": "", "json": []any{}}, nil - } - args := map[string]any{ - "query": query, - } - if apiKey := strings.TrimSpace(stringParam(inputs["api_key"])); apiKey != "" { - args["api_key"] = apiKey - } - if daysBack := toIntParam(inputs["days_back"]); daysBack > 0 { - args["days_back"] = daysBack - } - if topN := toIntParam(inputs["top_n"]); topN > 0 { - args["num_results"] = topN - } - - argsJSON, _ := json.Marshal(args) - out, err := c.inner.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - if err != nil { - if len(decoded) > 0 { - return map[string]any{ - "formalized_content": "", - "json": []any{}, - "_ERROR": decoded["_ERROR"], - }, nil - } - return nil, fmt.Errorf("canvas: BGPT: %w", err) - } - - results := anySlice(decoded["results"]) - return map[string]any{ - "formalized_content": renderBGPTResults(results), - "json": results, - }, nil -} - -func (c *bgptComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -// wikipediaComponent delegates to internal/agent/tool/WikipediaTool. -// Python's canvas component is named "Wikipedia" and stores top_n/language -// on the node params while accepting query at runtime. -type wikipediaComponent struct { - inner *agenttool.WikipediaTool -} - -func newWikipediaComponent(params map[string]any) (Component, error) { - topN := 10 - if v, ok := params["top_n"]; ok { - topN = toIntParam(v) - } - if topN <= 0 { - return nil, fmt.Errorf("canvas: Wikipedia: top_n must be a positive integer") - } - language := "en" - if v, ok := params["language"].(string); ok && strings.TrimSpace(v) != "" { - language = strings.TrimSpace(v) - } - if !agenttool.WikipediaLanguageSupported(language) { - return nil, fmt.Errorf("canvas: Wikipedia: unsupported language %q", language) - } - return &wikipediaComponent{inner: agenttool.NewWikipediaToolWithParams(nil, topN, language)}, nil -} - -func (c *wikipediaComponent) Name() string { return "Wikipedia" } - -func (c *wikipediaComponent) Inputs() map[string]string { - return map[string]string{ - "query": "The search keyword to execute with wikipedia. The keyword MUST be a specific subject that can match the title.", - } -} - -func (c *wikipediaComponent) GetInputForm() map[string]any { - return map[string]any{ - "query": map[string]any{ - "name": "Query", - "type": "line", - }, - } -} - -func (c *wikipediaComponent) Outputs() map[string]string { - return map[string]string{ - "formalized_content": "Rendered Wikipedia article summaries for downstream LLM prompts.", - "json": "Raw Wikipedia result list.", - } -} - -func (c *wikipediaComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - query := strings.TrimSpace(stringParam(inputs["query"])) - if query == "" { - return map[string]any{"formalized_content": "", "json": []any{}}, nil - } - argsJSON, _ := json.Marshal(map[string]any{"query": query}) - out, err := c.inner.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - if results, ok := decoded["results"]; ok { - decoded["json"] = results - } - if err != nil { - if len(decoded) > 0 { - if _, ok := decoded["formalized_content"]; !ok { - decoded["formalized_content"] = "" - } - if _, ok := decoded["json"]; !ok { - decoded["json"] = []any{} - } - return decoded, nil - } - return nil, fmt.Errorf("canvas: Wikipedia: %w", err) - } - return decoded, nil -} - -func (c *wikipediaComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -// duckDuckGoComponent delegates to internal/agent/tool/DuckDuckGoTool. -type duckDuckGoComponent struct { - inner duckDuckGoInvoker -} - -func newDuckDuckGoComponent(_ map[string]any) (Component, error) { - return newDuckDuckGoComponentWithInvoker(agenttool.NewDuckDuckGoTool()), nil -} - -func newDuckDuckGoComponentWithInvoker(inner duckDuckGoInvoker) Component { - return &duckDuckGoComponent{inner: inner} -} - -func (c *duckDuckGoComponent) Name() string { return "DuckDuckGo" } - -func (c *duckDuckGoComponent) Inputs() map[string]string { - return map[string]string{ - "query": "Search query.", - "channel": "Search channel: general or news.", - "top_n": "Maximum number of results.", - } -} - -func (c *duckDuckGoComponent) GetInputForm() map[string]any { - return map[string]any{ - "query": map[string]any{ - "name": "Query", - "type": "line", - }, - "channel": map[string]any{ - "name": "Channel", - "type": "options", - "value": "general", - "options": []string{"general", "news"}, - }, - } -} - -func (c *duckDuckGoComponent) Outputs() map[string]string { - return map[string]string{ - "formalized_content": "Rendered search results for downstream LLM prompts.", - "json": "Raw result list.", - } -} - -func (c *duckDuckGoComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - query := strings.TrimSpace(stringParam(inputs["query"])) - if query == "" { - return map[string]any{"formalized_content": "", "json": []any{}}, nil - } - args := map[string]any{ - "query": query, - } - if channel := strings.TrimSpace(stringParam(inputs["channel"])); channel != "" { - args["channel"] = channel - } - if topN := toIntParam(inputs["top_n"]); topN > 0 { - args["top_n"] = topN - } - - argsJSON, _ := json.Marshal(args) - out, err := c.inner.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - if err != nil { - if len(decoded) > 0 { - return map[string]any{ - "formalized_content": "", - "json": []any{}, - "_ERROR": decoded["_ERROR"], - }, nil - } - return nil, fmt.Errorf("canvas: DuckDuckGo: %w", err) - } - - results := anySlice(decoded["results"]) - return map[string]any{ - "formalized_content": renderDuckDuckGoResults(results), - "json": results, - }, nil -} - -func (c *duckDuckGoComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -// keenableSearchComponent delegates to internal/agent/tool/KeenableTool. -type keenableSearchComponent struct { - inner keenableInvoker - apiKey string - mode string - topN int - site string -} - -func newKeenableSearchComponent(params map[string]any) (Component, error) { - apiKey := stringParam(params["api_key"]) - return newKeenableSearchComponentWithInvoker(agenttool.NewKeenableToolWithAPIKey(nil, apiKey), params), nil -} - -func newKeenableSearchComponentWithInvoker(inner keenableInvoker, params map[string]any) Component { - if inner == nil { - inner = agenttool.NewKeenableTool() - } - mode := strings.TrimSpace(stringParam(params["mode"])) - if mode == "" { - mode = "pro" - } - topN := toIntParam(params["top_n"]) - if topN <= 0 { - topN = 10 - } - return &keenableSearchComponent{ - inner: inner, - apiKey: stringParam(params["api_key"]), - mode: mode, - topN: topN, - site: stringParam(params["site"]), - } -} - -func (c *keenableSearchComponent) Name() string { return "KeenableSearch" } - -func (c *keenableSearchComponent) Inputs() map[string]string { - return map[string]string{ - "query": "Search query.", - "site": "Optional single-domain filter.", - "api_key": "Optional Keenable API key.", - "mode": "Search mode: pro or realtime.", - "top_n": "Maximum number of results.", - } -} - -func (c *keenableSearchComponent) GetInputForm() map[string]any { - return map[string]any{ - "query": map[string]any{ - "name": "Query", - "type": "line", - }, - "site": map[string]any{ - "name": "Site", - "type": "line", - }, - } -} - -func (c *keenableSearchComponent) Outputs() map[string]string { - return map[string]string{ - "formalized_content": "Rendered search results for downstream LLM prompts.", - "json": "Raw Keenable result list.", - } -} - -func (c *keenableSearchComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - query := strings.TrimSpace(stringParam(inputs["query"])) - if query == "" { - return map[string]any{"formalized_content": "", "json": []any{}}, nil - } - - args := map[string]any{ - "query": query, - "mode": c.mode, - "top_n": c.topN, - } - if mode := strings.TrimSpace(stringParam(inputs["mode"])); mode != "" { - args["mode"] = mode - } - if topN := toIntParam(inputs["top_n"]); topN > 0 { - args["top_n"] = topN - } - site := strings.TrimSpace(stringParam(inputs["site"])) - if site == "" { - site = strings.TrimSpace(c.site) - } - if site != "" { - args["site"] = site - } - - invoker := c.inner - if apiKey := strings.TrimSpace(stringParam(inputs["api_key"])); apiKey != "" && apiKey != strings.TrimSpace(c.apiKey) { - invoker = agenttool.NewKeenableToolWithAPIKey(nil, apiKey) - } - - argsJSON, _ := json.Marshal(args) - out, err := invoker.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - results := anySlice(decoded["results"]) - if existing, _ := decoded["_ERROR"].(string); strings.TrimSpace(existing) != "" { - return map[string]any{ - "formalized_content": "", - "json": results, - "_ERROR": existing, - }, nil - } - if err != nil { - if len(decoded) > 0 { - return map[string]any{ - "formalized_content": "", - "json": results, - "_ERROR": decoded["_ERROR"], - }, nil - } - return nil, fmt.Errorf("canvas: KeenableSearch: %w", err) - } - chunks, docAggs := buildKeenableReferences(results) - if state, _, stateErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); stateErr == nil && state != nil { - state.SetRetrievalReferences(chunks, docAggs) - } - return map[string]any{ - "formalized_content": renderKeenableReferences(chunks), - "json": results, - }, nil -} - -func (c *keenableSearchComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -func buildKeenableReferences(results []any) ([]map[string]any, []map[string]any) { - chunks := make([]map[string]any, 0, len(results)) - docAggs := make([]map[string]any, 0, len(results)) - for _, item := range results { - m, ok := item.(map[string]any) - if !ok { - continue - } - content := truncateRunes(strings.TrimSpace(keenableValueString(m["description"])), 10000) - if content == "" || content == "None" { - continue - } - documentID := strconv.FormatInt(githubHashInt(content, 100000000), 10) - title := keenableValueString(m["title"]) - url := keenableValueString(m["url"]) - displayID := strconv.FormatInt(githubHashInt(documentID, 500), 10) - chunks = append(chunks, map[string]any{ - "id": displayID, - "chunk_id": documentID, - "content": content, - "doc_id": documentID, - "document_id": documentID, - "docnm_kwd": title, - "document_name": title, - "similarity": 1, - "score": 1, - "url": url, - }) - docAggs = append(docAggs, map[string]any{ - "doc_name": title, - "doc_id": documentID, - "count": 1, - "url": url, - }) - } - return chunks, docAggs -} - -func renderKeenableReferences(chunks []map[string]any) string { - if len(chunks) == 0 { - return "" - } - blocks := make([]string, 0, len(chunks)) - for _, chunk := range chunks { - blocks = append(blocks, strings.Join([]string{ - "\nID: " + keenableValueString(chunk["id"]), - "├── Title: " + keenableValueString(chunk["docnm_kwd"]), - "├── URL: " + keenableValueString(chunk["url"]), - "└── Content:\n" + keenableValueString(chunk["content"]), - }, "\n")) - } - return strings.Join(blocks, "\n") -} - -func keenableValueString(value any) string { - if value == nil { - return "None" - } - return strings.TrimSpace(fmt.Sprint(value)) -} - -func renderDuckDuckGoResults(results []any) string { - if len(results) == 0 { - return "" - } - blocks := make([]string, 0, len(results)) - for _, item := range results { - m, ok := item.(map[string]any) - if !ok { - continue - } - field := func(key string) string { - v, ok := m[key] - if !ok || v == nil { - return "-" - } - text := strings.TrimSpace(fmt.Sprintf("%v", v)) - if text == "" { - return "-" - } - return text - } - blocks = append(blocks, strings.Join([]string{ - fmt.Sprintf("Title: %s", field("title")), - fmt.Sprintf("URL: %s", field("url")), - fmt.Sprintf("Body: %s", field("body")), - }, "\n")) - } - return strings.Join(blocks, "\n\n") -} - func stringParam(v any) string { if s, ok := v.(string); ok { return s @@ -824,86 +60,6 @@ func anySlice(v any) []any { } } -func renderTavilySearchResults(results []any) string { - if len(results) == 0 { - return "" - } - blocks := make([]string, 0, len(results)) - for _, item := range results { - m, ok := item.(map[string]any) - if !ok { - continue - } - title := strings.TrimSpace(fmt.Sprint(m["title"])) - url := strings.TrimSpace(fmt.Sprint(m["url"])) - content := strings.TrimSpace(fmt.Sprint(m["raw_content"])) - if content == "" || content == "" { - content = strings.TrimSpace(fmt.Sprint(m["content"])) - } - if content == "" || content == "" { - continue - } - lines := []string{} - if title != "" && title != "" { - lines = append(lines, fmt.Sprintf("Title: %s", title)) - } - if url != "" && url != "" { - lines = append(lines, fmt.Sprintf("URL: %s", url)) - } - lines = append(lines, content) - blocks = append(blocks, strings.Join(lines, "\n")) - } - return strings.Join(blocks, "\n\n") -} - -func renderBGPTResults(results []any) string { - if len(results) == 0 { - return "" - } - blocks := make([]string, 0, len(results)) - for _, item := range results { - m, ok := item.(map[string]any) - if !ok { - continue - } - bgptField := func(key string) string { - v, ok := m[key] - if !ok || v == nil { - return "-" - } - switch vv := v.(type) { - case string: - if text := strings.TrimSpace(vv); text != "" { - return text - } - default: - if text := strings.TrimSpace(fmt.Sprintf("%v", vv)); text != "" { - return text - } - } - return "-" - } - lines := []string{ - fmt.Sprintf("Title: %s", bgptField("title")), - fmt.Sprintf("Authors: %s", bgptField("authors")), - fmt.Sprintf("Journal: %s", bgptField("journal")), - fmt.Sprintf("Year: %s", bgptField("year")), - fmt.Sprintf("DOI: %s", bgptField("doi")), - fmt.Sprintf("Abstract: %s", bgptField("abstract")), - fmt.Sprintf("Methods: %s", bgptField("methods")), - fmt.Sprintf("Sample size / population: %s", bgptField("sample_size")), - fmt.Sprintf("Results: %s", bgptField("results")), - fmt.Sprintf("Limitations: %s", bgptField("limitations")), - fmt.Sprintf("Conflicts of interest: %s", bgptField("conflict_of_interest")), - fmt.Sprintf("Data availability: %s", bgptField("data_availability")), - fmt.Sprintf("Blind spots: %s", bgptField("blind_spots")), - fmt.Sprintf("How to falsify: %s", bgptField("falsify")), - } - blocks = append(blocks, strings.Join(lines, "\n")) - } - return strings.Join(blocks, "\n\n") -} - // retrievalParams mirrors the Python RetrievalParam shape: the // values the canvas node declares at build time, applied as // defaults to the per-invocation RetrievalRequest. The fields are @@ -1146,9 +302,7 @@ func normalizeStructuredRetrievalInputs(ctx context.Context, out map[string]any) if kbName != "" && !hasDatasetIDs { if datasetID := resolveRetrievalDatasetID(ctx, strings.TrimSpace(kbName)); datasetID != "" { out["dataset_ids"] = []string{datasetID} - common.Debug("agent retrieval component: resolved dataset id", - zap.String("kb", strings.TrimSpace(kbName)), - zap.String("dataset_id", datasetID)) + common.Debug("agent retrieval component: resolved dataset id") } } if queryText != "" { @@ -1166,37 +320,23 @@ func resolveRetrievalDatasetID(ctx context.Context, kbName string) string { return "" } if kb, err := dao.NewKnowledgebaseDAO().GetByID(kbName); err == nil && kb != nil { - common.Debug("agent retrieval component: resolved dataset id by direct id", - zap.String("kb", kbName), - zap.String("dataset_id", kb.ID)) + common.Debug("agent retrieval component: resolved dataset id by direct id") return kb.ID } else if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { common.Warn("agent retrieval component: resolve dataset id by id failed", - zap.String("kb", kbName), zap.Error(err)) } if state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx); err == nil && state != nil { - common.Debug("agent retrieval component: resolve dataset id context", - zap.String("kb", kbName), - zap.Any("sys_query", state.Sys["query"]), - zap.Any("tenant_id", state.Sys["tenant_id"]), - zap.Any("user_id", state.Sys["user_id"])) + 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 { - common.Debug("agent retrieval component: resolved dataset id by tenant", - zap.String("kb", kbName), - zap.String("tenant_id", tenantID), - zap.String("dataset_id", kb.ID)) + common.Debug("agent retrieval component: resolved dataset id by tenant") return kb.ID } else if lookupErr != nil && !errors.Is(lookupErr, gorm.ErrRecordNotFound) { common.Warn("agent retrieval component: resolve dataset id by tenant failed", - zap.String("kb", kbName), - zap.String("tenant_id", tenantID), zap.Error(lookupErr)) } else { - common.Debug("agent retrieval component: tenant lookup missed", - zap.String("kb", kbName), - zap.String("tenant_id", tenantID)) + common.Debug("agent retrieval component: tenant lookup missed") } } if userID, _ := state.Sys["user_id"].(string); userID != "" { @@ -1205,263 +345,24 @@ func resolveRetrievalDatasetID(ctx context.Context, kbName string) string { if kb == nil || kb.Status == nil || *kb.Status != string(entity.StatusValid) { continue } - common.Debug("agent retrieval component: resolved dataset id by user visibility", - zap.String("kb", kbName), - zap.String("user_id", userID), - zap.String("dataset_id", kb.ID)) + common.Debug("agent retrieval component: resolved dataset id by user visibility") return kb.ID } } else if lookupErr != nil { common.Warn("agent retrieval component: resolve dataset id by name failed", - zap.String("kb", kbName), - zap.String("user_id", userID), zap.Error(lookupErr)) } else { - common.Debug("agent retrieval component: user visibility lookup missed", - zap.String("kb", kbName), - zap.String("user_id", userID)) + common.Debug("agent retrieval component: user visibility lookup missed") } } } else { common.Debug("agent retrieval component: resolve dataset id missing canvas state", - zap.String("kb", kbName), zap.Error(err)) } - common.Debug("agent retrieval component: dataset id unresolved", - zap.String("kb", kbName)) + common.Debug("agent retrieval component: dataset id unresolved") return "" } -// exesqlComponent delegates to internal/agent/tool/ExeSQLTool. The -// connection params (db_type, host, port, database, username, -// password) are passed via the canvas node's params map at build -// time, matching Python's ExeSQLParam semantics. -// -// v1 → tool param translation: the legacy v1 ExeSQL canvas node -// surface used (database, username, host, port, password, top_n) -// and did NOT declare db_type. The tool, by contrast, REQUIRES -// db_type (and uses max_records for the row cap, not top_n). A -// naive passthrough would turn every v1 canvas into a build-time -// error (NewExeSQLConnParams returns "missing required connection -// params (db_type/host/database/username)"). The adapter below -// bridges the two surfaces so existing v1 DSLs keep compiling. -// -// Defaults applied: db_type defaults to "mysql" (matches the v1 -// Python default); top_n is mapped to max_records; port is coerced -// from JSON-decoded float64 to int. See TestExeSQL_V1DSLParamsAccepted. -func newExeSQLComponent(params map[string]any) (Component, error) { - toolParams := translateExeSQLParamsToToolShape(params) - conn, err := agenttool.NewExeSQLConnParams(toolParams) - if err != nil { - return nil, fmt.Errorf("canvas: ExeSQL: %w", err) - } - return &exesqlComponent{ - inner: agenttool.NewExeSQLTool(conn), - sql: conn.SQL, - }, nil -} - -// translateExeSQLParamsToToolShape adapts a v1 DSL ExeSQL params -// map into the tool's expected param surface. Idempotent: callers -// that already supply db_type / max_records / int-typed port pass -// through unchanged. -// -// Field map: -// -// v1 surface → tool surface -// ------------------- -------------- -// db_type (optional) → db_type (defaults to "mysql") -// database → database -// username → username -// host → host -// port (float64) → port (coerced to int) -// password → password -// top_n (numeric) → max_records (and dropped from out) -// -// Returns a fresh map; the input is not mutated. -func translateExeSQLParamsToToolShape(v1Params map[string]any) map[string]any { - out := make(map[string]any, len(v1Params)+2) - for k, v := range v1Params { - out[k] = v - } - // db_type: required by the tool, absent in v1 DSL — default - // to mysql to match the v1 Python default and most legacy - // canvases. Operators wanting a different engine can set - // db_type explicitly in the params map. - if _, ok := out["db_type"]; !ok { - out["db_type"] = "mysql" - } - // port: JSON-decoded numeric comes through as float64, but - // NewExeSQLConnParams asserts on int via type-switch. Coerce. - if v, ok := out["port"]; ok { - switch x := v.(type) { - case float64: - out["port"] = int(x) - case int64: - out["port"] = int(x) - } - } - // top_n: v1's row-limit param. Map to max_records (the tool's - // equivalent). If both keys are present, max_records wins — the - // tool's name is the canonical one. - if v, ok := out["top_n"]; ok { - if _, hasMaxRecords := out["max_records"]; !hasMaxRecords { - switch x := v.(type) { - case float64: - out["max_records"] = int(x) - case int: - out["max_records"] = x - case int64: - out["max_records"] = int(x) - } - } - delete(out, "top_n") - } - return out -} - -type exeSQLInvoker interface { - InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error) -} - -type exesqlComponent struct { - inner exeSQLInvoker - sql string -} - -func (c *exesqlComponent) Name() string { return "ExeSQL" } - -func (c *exesqlComponent) Inputs() map[string]string { - return map[string]string{ - "sql": "SQL statement to execute (SELECT-only; DML/DDL rejected).", - } -} - -func (c *exesqlComponent) GetInputForm() map[string]any { - return map[string]any{ - "sql": map[string]any{ - "name": "SQL", - "type": "line", - }, - } -} - -func (c *exesqlComponent) Outputs() map[string]string { - return map[string]string{ - "formalized_content": "SQL result rendered as Markdown.", - "json": "Raw SQL statement results.", - } -} - -func (c *exesqlComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - sqlText := c.sql - if value, ok := inputs["sql"].(string); ok && strings.TrimSpace(value) != "" { - sqlText = value - } - if state, _, stateErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); stateErr == nil && state != nil { - resolved, resolveErr := runtime.ResolveTemplate(sqlText, state) - if resolveErr != nil { - return map[string]any{ - "formalized_content": "", - "json": []any{}, - "_ERROR": resolveErr.Error(), - }, nil - } - sqlText = resolved - } - - argsJSON, err := json.Marshal(map[string]any{"sql": sqlText}) - if err != nil { - return nil, fmt.Errorf("canvas: ExeSQL: encode SQL: %w", err) - } - out, invokeErr := c.inner.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - result := formatExeSQLCanvasOutput(decoded) - if invokeErr != nil { - if message := stringParam(decoded["_ERROR"]); message != "" { - result["_ERROR"] = message - } else { - result["_ERROR"] = invokeErr.Error() - } - } - return result, nil -} - -func formatExeSQLCanvasOutput(decoded map[string]any) map[string]any { - rows := anySlice(decoded["rows"]) - columns := anySlice(decoded["columns"]) - jsonResult := make([]any, 0, 1) - if len(rows) == 1 { - if row, ok := rows[0].(map[string]any); ok && len(row) == 1 { - if _, hasContent := row["content"]; hasContent { - jsonResult = append(jsonResult, row) - } - } - } - if len(jsonResult) == 0 && len(rows) > 0 { - jsonResult = append(jsonResult, rows) - } - result := map[string]any{ - "formalized_content": renderExeSQLMarkdown(columns, rows), - "json": jsonResult, - } - if message := stringParam(decoded["_ERROR"]); message != "" { - result["_ERROR"] = message - } - return result -} - -func renderExeSQLMarkdown(columns, rows []any) string { - if len(rows) == 0 { - return "" - } - for _, value := range rows { - if row, ok := value.(map[string]any); ok && len(row) == 1 { - if message, exists := row["content"]; exists { - return stringParam(message) - } - } - } - columnNames := make([]string, 0, len(columns)) - for _, column := range columns { - columnNames = append(columnNames, fmt.Sprint(column)) - } - if len(columnNames) == 0 { - if first, ok := rows[0].(map[string]any); ok { - for column := range first { - columnNames = append(columnNames, column) - } - sort.Strings(columnNames) - } - } - if len(columnNames) == 0 { - return "" - } - var builder strings.Builder - fmt.Fprintf(&builder, "| %s |\n", strings.Join(columnNames, " | ")) - separators := make([]string, len(columnNames)) - for i := range separators { - separators[i] = "---" - } - fmt.Fprintf(&builder, "| %s |\n", strings.Join(separators, " | ")) - for _, value := range rows { - row, ok := value.(map[string]any) - if !ok { - continue - } - cells := make([]string, len(columnNames)) - for i, column := range columnNames { - cells[i] = strings.ReplaceAll(strings.ReplaceAll(fmt.Sprint(row[column]), "|", "\\|"), "\n", "
") - } - fmt.Fprintf(&builder, "| %s |\n", strings.Join(cells, " | ")) - } - return strings.TrimSuffix(builder.String(), "\n") -} - -func (c *exesqlComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - // codeExecComponent delegates to internal/agent/tool/CodeExecTool. // The node-level params map carries the legacy v1 DSL surface // (`lang`, `script`, `arguments`, optional `timeout`). Per-call inputs @@ -1762,462 +663,11 @@ func toFloatParam(v any) float64 { return 0 } -// yahooFinanceComponent delegates to internal/agent/tool/YahooFinanceTool. -type yahooFinanceComponent struct { - inner *agenttool.YahooFinanceTool -} - -func newYahooFinanceComponent(_ map[string]any) (Component, error) { - return &yahooFinanceComponent{inner: agenttool.NewYahooFinanceTool()}, nil -} - -func (c *yahooFinanceComponent) Name() string { return "YahooFinance" } - -func (c *yahooFinanceComponent) Inputs() map[string]string { - return map[string]string{ - "stock_code": "Stock symbol to look up (e.g. AAPL, MSFT, 0005.HK).", - } -} - -func (c *yahooFinanceComponent) Outputs() map[string]string { - return map[string]string{ - "report": "Stock quote data (JSON).", - } -} - -func (c *yahooFinanceComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - stockCode, _ := inputs["stock_code"].(string) - if strings.TrimSpace(stockCode) == "" { - return map[string]any{"_ERROR": "stock_code is required"}, nil - } - toolInput := map[string]any{ - "symbols": []string{stockCode}, - } - argsJSON, _ := json.Marshal(toolInput) - out, err := c.inner.InvokableRun(ctx, string(argsJSON)) - if err != nil { - if out != "" { - return parseToolEnvelope(out), nil - } - return nil, fmt.Errorf("canvas: YahooFinance: %w", err) - } - result := parseToolEnvelope(out) - return map[string]any{"report": result["results"]}, nil -} - -func (c *yahooFinanceComponent) GetInputForm() map[string]any { - return map[string]any{ - "stock_code": map[string]any{ - "type": "line", - "name": "Stock code/Company name", - }, - } -} - -func (c *yahooFinanceComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -// arxivComponent delegates to internal/agent/tool/ArxivTool. Query is the -// runtime input; top_n and sort_by are validated node parameters. -type arxivComponent struct { - inner *agenttool.ArxivTool -} - -func newArxivComponent(params map[string]any) (Component, error) { - topN := 12 - if v, ok := params["top_n"]; ok { - topN = toIntParam(v) - } - if topN <= 0 { - return nil, fmt.Errorf("canvas: ArXiv: top_n must be a positive integer") - } - sortBy := "submittedDate" - if v, ok := params["sort_by"].(string); ok && strings.TrimSpace(v) != "" { - sortBy = strings.TrimSpace(v) - } - if !agenttool.ArxivSortBySupported(sortBy) { - return nil, fmt.Errorf("canvas: ArXiv: unsupported sort_by %q", sortBy) - } - return &arxivComponent{inner: agenttool.NewArxivToolWithParams(nil, topN, sortBy)}, nil -} - -func (c *arxivComponent) Name() string { return "ArXiv" } - -func (c *arxivComponent) Inputs() map[string]string { - return map[string]string{ - "query": "Search query.", - } -} - -func (c *arxivComponent) Outputs() map[string]string { - return map[string]string{ - "formalized_content": "Rendered arXiv papers for downstream LLM prompts.", - "json": "Raw arXiv paper list.", - } -} - -func (c *arxivComponent) GetInputForm() map[string]any { - return map[string]any{ - "query": map[string]any{ - "name": "Query", - "type": "line", - }, - } -} - -func (c *arxivComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - query := strings.TrimSpace(stringParam(inputs["query"])) - if query == "" { - return map[string]any{"formalized_content": "", "json": []any{}}, nil - } - argsJSON, _ := json.Marshal(map[string]any{"query": query}) - out, err := c.inner.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - if err != nil { - if len(decoded) > 0 { - return map[string]any{ - "formalized_content": "", - "json": []any{}, - "_ERROR": decoded["_ERROR"], - }, nil - } - return nil, fmt.Errorf("canvas: ArXiv: %w", err) - } - results := anySlice(decoded["results"]) - return map[string]any{ - "formalized_content": renderArxivResults(results), - "json": results, - }, nil -} - -func (c *arxivComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -func renderArxivResults(results []any) string { - if len(results) == 0 { - return "" - } - blocks := make([]string, 0, len(results)) - for _, item := range results { - paper, ok := item.(map[string]any) - if !ok { - continue - } - summary := strings.TrimSpace(stringParam(paper["summary"])) - if summary == "" { - continue - } - title := strings.TrimSpace(stringParam(paper["title"])) - url := strings.TrimSpace(stringParam(paper["pdf_url"])) - lines := make([]string, 0, 3) - if title != "" { - lines = append(lines, fmt.Sprintf("Title: %s", title)) - } - if url != "" { - lines = append(lines, fmt.Sprintf("URL: %s", url)) - } - lines = append(lines, summary) - blocks = append(blocks, strings.Join(lines, "\n")) - } - return strings.Join(blocks, "\n\n") -} - -// googleScholarComponent delegates to internal/agent/tool/GoogleScholarTool. -type googleScholarComponent struct { - inner googleScholarInvoker - params map[string]any -} - -type googleScholarInvoker interface { - InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error) -} - -type pubmedInvoker interface { - InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error) -} - -func newGoogleScholarComponent(params map[string]any) (Component, error) { - cloned := make(map[string]any, len(params)) - for k, v := range params { - cloned[k] = v - } - return &googleScholarComponent{ - inner: agenttool.NewGoogleScholarTool(), - params: cloned, - }, nil -} - -func newGoogleScholarComponentWithInvoker(inner googleScholarInvoker, params map[string]any) Component { - cloned := make(map[string]any, len(params)) - for k, v := range params { - cloned[k] = v - } - return &googleScholarComponent{inner: inner, params: cloned} -} - -func (c *googleScholarComponent) Name() string { return "GoogleScholar" } - -func (c *googleScholarComponent) Inputs() map[string]string { - return map[string]string{ - "query": "Search query.", - "top_n": "Maximum number of results (default 12).", - "sort_by": "Sort order: relevance or date.", - "year_low": "Earliest publication year to include.", - "year_high": "Latest publication year to include.", - "patents": "Whether to include patents, defaults to true.", - } -} - -func (c *googleScholarComponent) Outputs() map[string]string { - return map[string]string{ - "formalized_content": "Rendered search results for downstream LLM prompts.", - "json": "Raw result list.", - } -} - -func (c *googleScholarComponent) GetInputForm() map[string]any { - return map[string]any{ - "query": map[string]any{ - "name": "Query", - "type": "line", - }, - } -} - -func (c *googleScholarComponent) Invoke(ctx context.Context, 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 - } - for k, v := range inputs { - merged[k] = v - } - - query := strings.TrimSpace(stringParam(merged["query"])) - if query == "" { - return map[string]any{"formalized_content": "", "json": []any{}}, nil - } - args := map[string]any{ - "query": query, - } - if topN := toIntParam(merged["top_n"]); topN > 0 { - args["top_n"] = topN - } - if sortBy := strings.TrimSpace(stringParam(merged["sort_by"])); sortBy != "" { - args["sort_by"] = sortBy - } - if yearLow := toIntParam(merged["year_low"]); yearLow > 0 { - args["year_low"] = yearLow - } - if yearHigh := toIntParam(merged["year_high"]); yearHigh > 0 { - args["year_high"] = yearHigh - } - if patents, ok := merged["patents"].(bool); ok { - args["patents"] = patents - } else { - args["patents"] = true - } - - argsJSON, _ := json.Marshal(args) - out, err := c.inner.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - if err != nil { - if len(decoded) > 0 { - return map[string]any{ - "formalized_content": "", - "json": []any{}, - "_ERROR": decoded["_ERROR"], - }, nil - } - return nil, fmt.Errorf("canvas: GoogleScholar: %w", err) - } - - results := anySlice(decoded["results"]) - return map[string]any{ - "formalized_content": renderGoogleScholarResults(results), - "json": results, - }, nil -} - -func (c *googleScholarComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -func renderGoogleScholarResults(results []any) string { - if len(results) == 0 { - return "" - } - blocks := make([]string, 0, len(results)) - for _, item := range results { - m, ok := item.(map[string]any) - if !ok { - continue - } - field := func(key string) string { - v, ok := m[key] - if !ok || v == nil { - return "-" - } - text := strings.TrimSpace(fmt.Sprintf("%v", v)) - if text == "" { - return "-" - } - return text - } - blocks = append(blocks, strings.Join([]string{ - fmt.Sprintf("Title: %s", field("title")), - fmt.Sprintf("URL: %s", field("link")), - fmt.Sprintf("Authors: %s", field("authors")), - fmt.Sprintf("Year: %s", field("year")), - fmt.Sprintf("Snippet: %s", field("snippet")), - }, "\n")) - } - return strings.Join(blocks, "\n\n") -} - -// pubMedComponent delegates to the PubMed tool. Its node parameters are -// consumed at construction time, leaving query as the sole runtime input. -type pubMedComponent struct { - inner pubmedInvoker -} - -func newPubMedComponent(params map[string]any) (Component, error) { - toolParams := make(map[string]any, 2) - for _, key := range []string{"top_n", "email"} { - if value, ok := params[key]; ok { - toolParams[key] = value - } - } - inner, err := agenttool.BuildByName("pubmed", toolParams) - if err != nil { - return nil, err - } - invoker, ok := inner.(pubmedInvoker) - if !ok { - return nil, fmt.Errorf("PubMed: tool does not implement InvokableRun") - } - return newPubMedComponentWithInvoker(invoker), nil -} - -func newPubMedComponentWithInvoker(inner pubmedInvoker) Component { - return &pubMedComponent{inner: inner} -} - -func (c *pubMedComponent) Name() string { return "PubMed" } - -func (c *pubMedComponent) Inputs() map[string]string { - return map[string]string{ - "query": "PubMed search query.", - } -} - -func (c *pubMedComponent) Outputs() map[string]string { - return map[string]string{ - "formalized_content": "Rendered PubMed references for downstream LLM prompts.", - "json": "Raw PubMed result list.", - } -} - -func (c *pubMedComponent) GetInputForm() map[string]any { - return map[string]any{ - "query": map[string]any{ - "name": "Query", - "type": "line", - }, - } -} - -func (c *pubMedComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - query := strings.TrimSpace(stringParam(inputs["query"])) - if query == "" { - return map[string]any{"formalized_content": "", "json": []any{}}, nil - } - argsJSON, err := json.Marshal(map[string]any{"query": query}) - if err != nil { - return nil, fmt.Errorf("canvas: PubMed: encode query: %w", err) - } - out, err := c.inner.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - results := anySlice(decoded["results"]) - if existing, _ := decoded["_ERROR"].(string); strings.TrimSpace(existing) != "" { - return map[string]any{ - "formalized_content": "", - "json": results, - "_ERROR": existing, - }, nil - } - if err != nil { - if len(decoded) > 0 { - return map[string]any{ - "formalized_content": "", - "json": results, - "_ERROR": decoded["_ERROR"], - }, nil - } - return nil, fmt.Errorf("canvas: PubMed: %w", err) - } - return map[string]any{ - "formalized_content": renderPubMedResults(results), - "json": results, - }, nil -} - -func (c *pubMedComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -func renderPubMedResults(results []any) string { - if len(results) == 0 { - return "" - } - blocks := make([]string, 0, len(results)) - for i, item := range results { - result, ok := item.(map[string]any) - if !ok { - continue - } - content := strings.TrimSpace(stringParam(result["content"])) - if content == "" { - continue - } - lines := []string{fmt.Sprintf("ID: %d", i)} - if title := strings.TrimSpace(stringParam(result["title"])); title != "" { - lines = append(lines, "Title: "+title) - } - if link := strings.TrimSpace(stringParam(result["url"])); link != "" { - lines = append(lines, "URL: "+link) - } - lines = append(lines, "Content:", content) - blocks = append(blocks, strings.Join(lines, "\n")) - } - return strings.Join(blocks, "\n\n") -} - // Compile-time interface checks. var ( _ Component = (*retrievalComponent)(nil) - _ Component = (*tavilySearchComponent)(nil) - _ Component = (*googleComponent)(nil) - _ Component = (*tavilyExtractComponent)(nil) - _ Component = (*duckDuckGoComponent)(nil) - _ Component = (*exesqlComponent)(nil) _ Component = (*codeExecComponent)(nil) - _ Component = (*arxivComponent)(nil) - _ Component = (*wikipediaComponent)(nil) - _ Component = (*googleScholarComponent)(nil) - _ Component = (*pubMedComponent)(nil) - _ Component = (*yahooFinanceComponent)(nil) ) // Compile-time check that the eino InvokableTool methods we call // are reachable (catches a future refactor that renames them). -var _ einotool.InvokableTool = (*agenttool.TavilyTool)(nil) -var _ einotool.InvokableTool = (*agenttool.WikipediaTool)(nil) -var _ einotool.InvokableTool = (*agenttool.GoogleTool)(nil) -var _ einotool.InvokableTool = (*agenttool.TavilyExtractTool)(nil) -var _ einotool.InvokableTool = (*agenttool.DuckDuckGoTool)(nil) -var _ einotool.InvokableTool = (*agenttool.YahooFinanceTool)(nil) -var _ einotool.InvokableTool = (*agenttool.GoogleScholarTool)(nil) -var _ einotool.InvokableTool = (*agenttool.PubMedTool)(nil) diff --git a/internal/agent/component/wencai.go b/internal/agent/component/wencai.go deleted file mode 100644 index 979749213c..0000000000 --- a/internal/agent/component/wencai.go +++ /dev/null @@ -1,109 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "encoding/json" - "fmt" - - einotool "github.com/cloudwego/eino/components/tool" - - agenttool "ragflow/internal/agent/tool" -) - -type wencaiInvoker interface { - InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error) -} - -type wencaiComponent struct { - inner wencaiInvoker -} - -func newWencaiComponent(params map[string]any) (Component, error) { - toolParams := make(map[string]any, 2) - for _, key := range []string{"top_n", "query_type"} { - if value, ok := params[key]; ok { - toolParams[key] = value - } - } - inner, err := agenttool.BuildByName("wencai", toolParams) - if err != nil { - return nil, err - } - invoker, ok := inner.(wencaiInvoker) - if !ok { - return nil, fmt.Errorf("WenCai: tool does not implement InvokableRun") - } - return newWencaiComponentWithInvoker(invoker), nil -} - -func newWencaiComponentWithInvoker(inner wencaiInvoker) Component { - return &wencaiComponent{inner: inner} -} - -func (c *wencaiComponent) Name() string { return "WenCai" } - -func (c *wencaiComponent) Inputs() map[string]string { - return map[string]string{ - "query": "The question/conditions to select stocks.", - } -} - -func (c *wencaiComponent) Outputs() map[string]string { - return map[string]string{ - "report": "WenCai query report.", - } -} - -func (c *wencaiComponent) GetInputForm() map[string]any { - return map[string]any{ - "query": map[string]any{ - "name": "Query", - "type": "line", - }, - } -} - -func (c *wencaiComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - query := stringParam(inputs["query"]) - if query == "" { - return map[string]any{"report": ""}, nil - } - argsJSON, err := json.Marshal(map[string]any{"query": query}) - if err != nil { - return nil, fmt.Errorf("canvas: WenCai: encode query: %w", err) - } - out, err := c.inner.InvokableRun(ctx, string(argsJSON)) - decoded := parseToolEnvelope(out) - report, _ := decoded["report"].(string) - if message, _ := decoded["_ERROR"].(string); message != "" { - return map[string]any{"report": report, "_ERROR": message}, nil - } - if err != nil { - return nil, fmt.Errorf("canvas: WenCai: %w", err) - } - return map[string]any{"report": report}, nil -} - -func (c *wencaiComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { - return nil, nil -} - -func init() { - Register("WenCai", newWencaiComponent) -} diff --git a/internal/agent/component/wencai_test.go b/internal/agent/component/wencai_test.go deleted file mode 100644 index f849f0827a..0000000000 --- a/internal/agent/component/wencai_test.go +++ /dev/null @@ -1,190 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "encoding/json" - "errors" - "testing" - - einotool "github.com/cloudwego/eino/components/tool" - - agentcanvas "ragflow/internal/agent/canvas" -) - -type fakeWencaiInvoker struct { - args map[string]any - calls int - out string - err error -} - -func (f *fakeWencaiInvoker) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) { - f.calls++ - if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil { - return "", err - } - if f.out == "" && f.err == nil { - return `{"report":""}`, nil - } - return f.out, f.err -} - -func TestWencai_RegisteredFactoryMatchesPythonSurface(t *testing.T) { - t.Parallel() - - c, err := New("WenCai", map[string]any{ - "top_n": float64(20), - "query_type": "stock", - "outputs": map[string]any{"report": map[string]any{}}, - "setups": map[string]any{"query": "configured query"}, - }) - if err != nil { - t.Fatalf("New(WenCai): %v", err) - } - if got := c.Name(); got != "WenCai" { - t.Fatalf("Name() = %q, want WenCai", got) - } - if _, ok := c.Inputs()["query"]; !ok { - t.Fatal("Inputs() missing query") - } - if _, ok := c.Outputs()["report"]; !ok { - t.Fatal("Outputs() missing report") - } - formGetter, ok := c.(interface{ GetInputForm() map[string]any }) - if !ok { - t.Fatal("WenCai component does not expose GetInputForm") - } - query, ok := formGetter.GetInputForm()["query"].(map[string]any) - if !ok { - t.Fatalf("query form has type %T, want map", formGetter.GetInputForm()["query"]) - } - if query["name"] != "Query" || query["type"] != "line" { - t.Fatalf("query form = %#v, want name=Query type=line", query) - } -} - -func TestWencai_InvokeReturnsEmptyReportWithoutError(t *testing.T) { - t.Parallel() - - fake := &fakeWencaiInvoker{} - c := newWencaiComponentWithInvoker(fake) - out, err := c.Invoke(context.Background(), map[string]any{ - "query": "商业航天", - "top_n": 20, - "query_type": "stock", - }) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if got := out["report"]; got != "" { - t.Fatalf("report = %v, want empty string", got) - } - if fake.calls != 1 { - t.Fatalf("calls = %d, want 1", fake.calls) - } - if got := fake.args["query"]; got != "商业航天" { - t.Fatalf("runtime query = %v, want 商业航天", got) - } - if len(fake.args) != 1 { - t.Fatalf("runtime args = %#v, want query only", fake.args) - } -} - -func TestWencai_InvokeEmptyQuerySkipsTool(t *testing.T) { - t.Parallel() - - fake := &fakeWencaiInvoker{} - c := newWencaiComponentWithInvoker(fake) - out, err := c.Invoke(context.Background(), map[string]any{"query": ""}) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if got := out["report"]; got != "" { - t.Fatalf("report = %v, want empty string", got) - } - if fake.calls != 0 { - t.Fatalf("calls = %d, want 0", fake.calls) - } -} - -func TestWencai_InvokePreservesErrorEnvelope(t *testing.T) { - t.Parallel() - - fake := &fakeWencaiInvoker{ - out: `{"report":"","_ERROR":"upstream failed"}`, - err: errors.New("upstream failed"), - } - c := newWencaiComponentWithInvoker(fake) - out, err := c.Invoke(context.Background(), map[string]any{"query": "商业航天"}) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if got := out["report"]; got != "" { - t.Fatalf("report = %v, want empty string", got) - } - if got := out["_ERROR"]; got != "upstream failed" { - t.Fatalf("_ERROR = %v, want upstream failed", got) - } -} - -func TestWencai_InvokePreservesErrorEnvelopeWithoutGoError(t *testing.T) { - t.Parallel() - - fake := &fakeWencaiInvoker{ - out: `{"report":"","_ERROR":"business error"}`, - } - c := newWencaiComponentWithInvoker(fake) - out, err := c.Invoke(context.Background(), map[string]any{"query": "商业航天"}) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if got := out["report"]; got != "" { - t.Fatalf("report = %v, want empty string", got) - } - if got := out["_ERROR"]; got != "business error" { - t.Fatalf("_ERROR = %v, want business error", got) - } -} - -func TestWencai_BuildWorkflowUsesPythonComponentName(t *testing.T) { - c := &agentcanvas.Canvas{ - Components: map[string]agentcanvas.CanvasComponent{ - "begin_0": { - Obj: agentcanvas.CanvasComponentObj{ComponentName: "Begin", Params: map[string]any{}}, - Downstream: []string{"wencai_0"}, - }, - "wencai_0": { - Obj: agentcanvas.CanvasComponentObj{ComponentName: "WenCai", Params: map[string]any{ - "top_n": float64(20), - "query_type": "stock", - }}, - Upstream: []string{"begin_0"}, - Downstream: []string{"message_0"}, - }, - "message_0": { - Obj: agentcanvas.CanvasComponentObj{ComponentName: "Message", Params: map[string]any{}}, - Upstream: []string{"wencai_0"}, - }, - }, - Path: []string{"begin_0", "wencai_0", "message_0"}, - } - if _, err := agentcanvas.BuildWorkflow(context.Background(), c); err != nil { - t.Fatalf("BuildWorkflow with WenCai: %v", err) - } -} diff --git a/internal/agent/component/wikipedia_component_test.go b/internal/agent/component/wikipedia_component_test.go deleted file mode 100644 index eb6f0df9a8..0000000000 --- a/internal/agent/component/wikipedia_component_test.go +++ /dev/null @@ -1,132 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "net/http" - "net/http/httptest" - "net/url" - "strings" - "testing" - - agenttool "ragflow/internal/agent/tool" -) - -func TestWikipedia_RegisteredFactory(t *testing.T) { - c, err := New("Wikipedia", map[string]any{"top_n": float64(3), "language": "en"}) - if err != nil { - t.Fatalf("New(Wikipedia) errored: %v", err) - } - if got := c.Name(); got != "Wikipedia" { - t.Fatalf("Name() = %q, want Wikipedia", got) - } - formGetter, ok := c.(interface{ GetInputForm() map[string]any }) - if !ok { - t.Fatal("Wikipedia component does not expose GetInputForm") - } - query, ok := formGetter.GetInputForm()["query"].(map[string]any) - if !ok { - t.Fatalf("GetInputForm()[query] has type %T, want map", formGetter.GetInputForm()["query"]) - } - if query["type"] != "line" { - t.Fatalf("GetInputForm()[query][type] = %v, want line", query["type"]) - } - if _, ok := c.Outputs()["formalized_content"]; !ok { - t.Fatal("Outputs() missing formalized_content") - } - if _, ok := c.Outputs()["json"]; !ok { - t.Fatal("Outputs() missing json") - } -} - -func TestWikipedia_InvokeEmptyQueryMatchesPython(t *testing.T) { - c, err := New("Wikipedia", nil) - if err != nil { - t.Fatalf("New(Wikipedia) errored: %v", err) - } - out, err := c.Invoke(context.Background(), map[string]any{"query": " "}) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - if got, _ := out["formalized_content"].(string); got != "" { - t.Fatalf("formalized_content = %q, want empty", got) - } - if _, ok := out["json"].([]any); !ok { - t.Fatalf("json output has type %T, want []any", out["json"]) - } -} - -func TestWikipedia_InvokeParamsBakedAtConstruct(t *testing.T) { - t.Parallel() - - var gotLimit string - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - gotLimit = r.URL.Query().Get("gsrlimit") - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{ - "query": { - "pages": { - "10": {"index":1,"title":"RAG","extract":"RAG is an acronym."} - } - } - }`)) - })) - defer srv.Close() - - helper := agenttool.NewHTTPHelper().WithClient(&http.Client{ - Transport: componentRewriteHostTransport(srv.URL), - }) - inner := agenttool.NewWikipediaToolWithParams(helper, 2, "en") - c := &wikipediaComponent{inner: inner} - out, err := c.Invoke(context.Background(), map[string]any{"query": "rag"}) - if err != nil { - t.Fatalf("Invoke errored: %v", err) - } - if gotLimit != "2" { - t.Fatalf("gsrlimit = %q, want 2", gotLimit) - } - if got, _ := out["formalized_content"].(string); !strings.Contains(got, "RAG is an acronym.") { - t.Fatalf("formalized_content = %q, want rendered result", got) - } - results := anySlice(out["json"]) - if len(results) != 1 { - t.Fatalf("json len = %d, want 1", len(results)) - } -} - -func componentRewriteHostTransport(srvURL string) http.RoundTripper { - u, err := url.Parse(srvURL) - if err != nil { - panic("componentRewriteHostTransport: bad srvURL: " + err.Error()) - } - return &componentHostSwapRT{inner: http.DefaultTransport, host: u.Host, scheme: u.Scheme} -} - -type componentHostSwapRT struct { - inner http.RoundTripper - host string - scheme string -} - -func (t *componentHostSwapRT) RoundTrip(req *http.Request) (*http.Response, error) { - r2 := req.Clone(req.Context()) - r2.URL.Scheme = t.scheme - r2.URL.Host = t.host - r2.Host = t.host - return t.inner.RoundTrip(r2) -} diff --git a/internal/agent/tool/arxiv.go b/internal/agent/tool/arxiv.go index 6071c6d1bf..cbd00ce0bb 100644 --- a/internal/agent/tool/arxiv.go +++ b/internal/agent/tool/arxiv.go @@ -19,19 +19,26 @@ package tool import ( "bytes" "context" + "crypto/sha1" "encoding/json" "encoding/xml" "fmt" "io" + "math/big" "net/http" "net/url" + "regexp" + "strconv" "strings" + "unicode/utf8" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" + + "ragflow/internal/tokenizer" ) -const arxivToolName = "arxiv" +const arxivToolName = "arxiv_search" const arxivToolDescription = "Search arXiv and return matching preprints as {title, authors, summary, pdf_url, entry_id}." @@ -39,6 +46,12 @@ const defaultArxivTopN = 12 const defaultArxivSortBy = "submittedDate" +const arxivPromptMaxTokens = 200000 + +var arxivDataImagePattern = regexp.MustCompile(`!?\[[a-z]+\]\(data:image/png;base64,[ 0-9A-Za-z/_=+\-]+\)`) + +var arxivNewlinePattern = regexp.MustCompile(`\n+`) + // arxivParams carries the query and ArXiv search settings. Info exposes only // query to match Python's tool meta; top_n and sort_by come from node params. type arxivParams struct { @@ -98,6 +111,9 @@ type ArxivTool struct { defaults arxivParams } +var _ ToolComponent = (*ArxivTool)(nil) +var _ ReferenceBuilder = (*ArxivTool)(nil) + // NewArxivTool returns an ArxivTool using the default HTTPHelper. func NewArxivTool() *ArxivTool { return NewArxivToolWith(NewHTTPHelper()) @@ -139,6 +155,19 @@ func (a *ArxivTool) Info(_ context.Context) (*schema.ToolInfo, error) { }, nil } +func (a *ArxivTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{"query": "Search query."}, + Outputs: map[string]string{ + "formalized_content": "Rendered arXiv references for downstream prompts.", + "json": "arXiv paper list.", + }, + InputForm: map[string]any{ + "query": map[string]any{"name": "Query", "type": "line"}, + }, + } +} + // buildArxivURL constructs the ArXiv /api/query URL. func buildArxivURL(query string, topN int, sortBy string) string { if topN <= 0 { @@ -229,8 +258,7 @@ func (a *ArxivTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool } p = mergeArxivDefaults(a.defaults, p) if strings.TrimSpace(p.Query) == "" { - return arxivErrJSON(fmt.Errorf("query is required")), - fmt.Errorf("arxiv: query is required") + return arxivJSON(arxivEnvelope{Results: []arxivResult{}}), nil } if p.TopN <= 0 { return arxivErrJSON(fmt.Errorf("top_n must be a positive integer")), @@ -266,6 +294,96 @@ func (a *ArxivTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool return arxivJSON(arxivEnvelope{Results: results}), nil } +func (a *ArxivTool) BuildReferences(_ context.Context, envelope map[string]any) ([]map[string]any, []map[string]any) { + return buildArxivReferences(envelope) +} + +func (a *ArxivTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + results := envelopeSlice(envelope, "results") + chunks, _ := buildArxivReferences(envelope) + return map[string]any{ + "formalized_content": renderArxivReferences(chunks, arxivPromptMaxTokens), + "json": results, + } +} + +func buildArxivReferences(envelope map[string]any) ([]map[string]any, []map[string]any) { + results := envelopeSlice(envelope, "results") + chunks := make([]map[string]any, 0, len(results)) + docAggs := make([]map[string]any, 0, len(results)) + for _, result := range results { + paper, ok := result.(map[string]any) + if !ok { + continue + } + content := arxivDataImagePattern.ReplaceAllString(arxivText(paper["summary"]), "") + content = truncateArxivRunes(content, 10000) + if content == "" { + continue + } + documentID := strconv.FormatInt(arxivHashInt(content, 100000000), 10) + displayID := strconv.FormatInt(arxivHashInt(documentID, 500), 10) + title := arxivText(paper["title"]) + resultURL := arxivText(paper["pdf_url"]) + chunks = append(chunks, map[string]any{ + "id": displayID, + "chunk_id": documentID, + "content": content, + "doc_id": documentID, + "document_id": documentID, + "docnm_kwd": title, + "document_name": title, + "similarity": 1, + "score": 1, + "url": resultURL, + }) + docAggs = append(docAggs, map[string]any{"doc_name": title, "doc_id": documentID, "count": 1, "url": resultURL}) + } + return chunks, docAggs +} + +func renderArxivReferences(chunks []map[string]any, maxTokens int) string { + usedTokens := 0 + blocks := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + content := arxivText(chunk["content"]) + usedTokens += tokenizer.NumTokensFromString(content) + blocks = append(blocks, strings.Join([]string{ + "\nID: " + arxivText(chunk["id"]), + "├── Title: " + arxivNewlinePattern.ReplaceAllString(arxivText(chunk["document_name"]), " "), + "├── URL: " + arxivNewlinePattern.ReplaceAllString(arxivText(chunk["url"]), " "), + "└── Content:\n" + content, + }, "\n")) + if maxTokens > 0 && float64(maxTokens)*0.97 < float64(usedTokens) { + break + } + } + return strings.Join(blocks, "\n") +} + +func arxivText(value any) string { + if value == nil { + return "" + } + if text, ok := value.(string); ok { + return text + } + return fmt.Sprint(value) +} + +func arxivHashInt(value string, modulus int64) int64 { + digest := sha1.Sum([]byte(value)) + number := new(big.Int).SetBytes(digest[:]) + return new(big.Int).Mod(number, big.NewInt(modulus)).Int64() +} + +func truncateArxivRunes(value string, limit int) string { + if utf8.RuneCountInString(value) <= limit { + return value + } + return string([]rune(value)[:limit]) +} + func mergeArxivDefaults(defaults, p arxivParams) arxivParams { if p.TopN == 0 { p.TopN = defaults.TopN diff --git a/internal/agent/tool/arxiv_test.go b/internal/agent/tool/arxiv_test.go index 67f19acc44..5b15622bb2 100644 --- a/internal/agent/tool/arxiv_test.go +++ b/internal/agent/tool/arxiv_test.go @@ -131,8 +131,8 @@ func TestArxiv_Info(t *testing.T) { if err != nil { t.Fatalf("Info: %v", err) } - if info.Name != "arxiv" { - t.Errorf("Name = %q, want arxiv", info.Name) + if info.Name != "arxiv_search" { + t.Errorf("Name = %q, want arxiv_search", info.Name) } if !strings.Contains(info.Desc, "arXiv") { t.Errorf("Desc = %q, want to mention arXiv", info.Desc) @@ -146,16 +146,17 @@ func TestArxiv_Info(t *testing.T) { } } -func TestArxiv_RequiresQuery(t *testing.T) { +func TestArxiv_EmptyQuery(t *testing.T) { t.Parallel() tool := NewArxivTool() - _, err := tool.InvokableRun(context.Background(), `{"query":""}`) - if err == nil { - t.Fatal("expected error for empty query") + out, err := tool.InvokableRun(context.Background(), `{"query":""}`) + if err != nil { + t.Fatalf("InvokableRun(empty): %v", err) } - if !strings.Contains(err.Error(), "query") { - t.Errorf("err = %v, want to mention query", err) + var envelope arxivEnvelope + if err := json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 { + t.Fatalf("empty result = %s / %v", out, err) } } @@ -208,3 +209,42 @@ func TestArxiv_FullRoundtrip(t *testing.T) { t.Errorf("PDFURL = %q, want http://arxiv.org/pdf/2501.12345v1", env.Results[0].PDFURL) } } + +func TestArxiv_ComponentReferencesAndDefaults(t *testing.T) { + t.Parallel() + + built, err := BuildByName("arxiv", map[string]any{ + "top_n": float64(7), + "sort_by": "relevance", + "outputs": map[string]any{"json": map[string]any{}}, + }) + if err != nil { + t.Fatalf("BuildByName: %v", err) + } + arxiv := built.(*ArxivTool) + if arxiv.defaults.TopN != 7 || arxiv.defaults.SortBy != "relevance" { + t.Fatalf("defaults = %+v", arxiv.defaults) + } + spec := arxiv.ComponentSpec() + if query, ok := spec.InputForm["query"].(map[string]any); !ok || query["type"] != "line" { + t.Fatalf("query input form = %#v", spec.InputForm["query"]) + } + envelope := map[string]any{"results": []any{map[string]any{ + "title": "Paper", "summary": "Paper summary.", "pdf_url": "https://arxiv.org/pdf/1", "entry_id": "kept", + }}} + chunks, docAggs := arxiv.BuildReferences(context.Background(), envelope) + if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["content"] != "Paper summary." { + t.Fatalf("references = %#v / %#v", chunks, docAggs) + } + outputs := arxiv.BuildComponentOutputs(envelope) + results := outputs["json"].([]any) + if results[0].(map[string]any)["entry_id"] != "kept" { + t.Fatalf("json output = %#v", results) + } + if !strings.Contains(outputs["formalized_content"].(string), "Paper summary.") { + t.Fatalf("formalized_content = %q", outputs["formalized_content"]) + } + if _, exists := envelope["chunks"]; exists { + t.Fatalf("output conversion mutated envelope: %#v", envelope) + } +} diff --git a/internal/agent/tool/bgpt.go b/internal/agent/tool/bgpt.go index 8dc69276c9..c6737b4263 100644 --- a/internal/agent/tool/bgpt.go +++ b/internal/agent/tool/bgpt.go @@ -17,17 +17,22 @@ package tool import ( - "bytes" "context" + "crypto/sha1" "encoding/json" "fmt" - "io" + "math/big" "net/http" + "regexp" + "strconv" "strings" "time" + "unicode/utf8" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" + + "ragflow/internal/tokenizer" ) const bgptToolName = "bgpt_search" @@ -39,37 +44,26 @@ const bgptToolDescription = "Search scientific papers via BGPT (bgpt.pro) and re // bgptEndpoint is the BGPT search URL. Exposed as a package var for tests. var bgptEndpoint = "https://bgpt.pro/api/mcp-search" -// bgptParams is the JSON shape the model sends into InvokableRun. -type bgptParams struct { - Query string `json:"query"` - MaxResults int `json:"num_results"` - APIKey string `json:"api_key,omitempty"` - DaysBack int `json:"days_back,omitempty"` -} +const bgptPromptMaxTokens = 200000 -// bgptResult is one paper in the result list. -type bgptResult struct { - Title string `json:"title"` - Authors string `json:"authors"` - Journal string `json:"journal"` - Year string `json:"year"` - DOI string `json:"doi"` - URL string `json:"url"` - Abstract string `json:"abstract"` - Methods string `json:"methods"` - SampleSize string `json:"sample_size"` - Results string `json:"results"` - Limitations string `json:"limitations"` - ConflictOfInterest string `json:"conflict_of_interest"` - DataAvailability string `json:"data_availability"` - BlindSpots string `json:"blind_spots"` - Falsify string `json:"falsify"` +var bgptDataImagePattern = regexp.MustCompile(`!?\[[a-z]+\]\(data:image/png;base64,[ 0-9A-Za-z/_=+\-]+\)`) + +var bgptNewlinePattern = regexp.MustCompile(`\n+`) + +// bgptParams holds the Canvas node configuration and the model-emitted +// runtime input. Info exposes only query; api_key / top_n / days_back +// come from canvas DSL params and live in BGPTTool.defaults. +type bgptParams struct { + Query string `json:"query"` + TopN int `json:"top_n"` + APIKey string `json:"api_key,omitempty"` + DaysBack int `json:"days_back,omitempty"` } // bgptEnv is what the model sees. type bgptEnv struct { - Results []bgptResult `json:"results"` - Error string `json:"_ERROR,omitempty"` + Results []map[string]any `json:"results"` + Error string `json:"_ERROR,omitempty"` } // bgptResponse is the upstream BGPT API envelope. @@ -80,20 +74,32 @@ type bgptResponse struct { // BGPTTool searches scientific papers via bgpt.pro. type BGPTTool struct { - helper *HTTPHelper + helper *HTTPHelper + defaults bgptParams } +var _ ToolComponent = (*BGPTTool)(nil) +var _ ReferenceBuilder = (*BGPTTool)(nil) + // NewBGPTTool returns a BGPTTool using the default HTTPHelper. func NewBGPTTool() *BGPTTool { - return NewBGPTToolWith(NewHTTPHelper()) + return newBGPTTool(nil, bgptParams{}) } // NewBGPTToolWith returns a BGPTTool with the given helper. func NewBGPTToolWith(h *HTTPHelper) *BGPTTool { + return newBGPTTool(h, bgptParams{}) +} + +func newBGPTTool(h *HTTPHelper, defaults bgptParams) *BGPTTool { if h == nil { - h = NewHTTPHelper() + h = NewHTTPHelperWithRetry(RetryConfig{MaxAttempts: 1}) + h.client.Timeout = 25 * time.Second } - return &BGPTTool{helper: h} + if defaults.TopN == 0 { + defaults.TopN = 10 + } + return &BGPTTool{helper: h, defaults: defaults} } // Info returns the tool's metadata for the chat model. @@ -107,25 +113,28 @@ func (b *BGPTTool) Info(_ context.Context) (*schema.ToolInfo, error) { Desc: "Natural-language scientific search query.", Required: true, }, - "num_results": { - Type: schema.Integer, - Desc: "Maximum number of results. Defaults to 10.", - Required: false, - }, - "api_key": { - Type: schema.String, - Desc: "Optional BGPT API key. Leave blank for the free tier.", - Required: false, - }, - "days_back": { - Type: schema.Integer, - Desc: "Optional recency filter (e.g. 365 for last year).", - Required: false, - }, }), }, nil } +func (b *BGPTTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "query": "Scientific search query.", + "api_key": "Optional BGPT API key.", + "days_back": "Optional recency filter in days.", + "top_n": "Maximum number of results.", + }, + Outputs: map[string]string{ + "formalized_content": "Rendered scientific paper references for downstream prompts.", + "json": "Raw BGPT result list.", + }, + InputForm: map[string]any{ + "query": map[string]any{"name": "Query", "type": "line"}, + }, + } +} + // InvokableRun performs the BGPT search. func (b *BGPTTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) { var p bgptParams @@ -134,16 +143,16 @@ func (b *BGPTTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool. fmt.Errorf("bgpt: parse arguments: %w", err) } if strings.TrimSpace(p.Query) == "" { - return bgptErrJSON(fmt.Errorf("query is required")), - fmt.Errorf("bgpt: query is required") + return bgptJSON(bgptEnv{Results: []map[string]any{}}), nil } - if p.MaxResults <= 0 { - p.MaxResults = 10 + p = mergeBGPTParams(b.defaults, p) + if p.TopN <= 0 { + p.TopN = 10 } reqBody := map[string]interface{}{ "query": strings.TrimSpace(p.Query), - "num_results": p.MaxResults, + "num_results": p.TopN, } if p.APIKey != "" { reqBody["api_key"] = p.APIKey @@ -157,34 +166,20 @@ func (b *BGPTTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool. return bgptErrJSON(err), err } - req, err := http.NewRequestWithContext(ctx, http.MethodPost, bgptEndpoint, bytes.NewReader(bodyBytes)) - if err != nil { - return bgptErrJSON(err), err - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Accept", "application/json") - - client := &http.Client{Timeout: 25 * time.Second} - resp, err := client.Do(req) + resp, err := b.helper.Do(ctx, http.MethodPost, bgptEndpoint, string(bodyBytes), "application/json", map[string]string{"Accept": "application/json"}) if err != nil { return bgptErrJSON(fmt.Errorf("bgpt: request failed: %w", err)), fmt.Errorf("bgpt: request failed: %w", err) } defer resp.Body.Close() - body, err := io.ReadAll(resp.Body) - if err != nil { - return bgptErrJSON(fmt.Errorf("bgpt: read response: %w", err)), - fmt.Errorf("bgpt: read response: %w", err) - } - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - return bgptErrJSON(fmt.Errorf("bgpt: upstream returned %d: %s", resp.StatusCode, string(body))), - fmt.Errorf("bgpt: upstream returned %d: %s", resp.StatusCode, string(body)) + return bgptErrJSON(fmt.Errorf("bgpt: upstream returned %d", resp.StatusCode)), + fmt.Errorf("bgpt: upstream returned %d", resp.StatusCode) } var raw bgptResponse - if err := json.Unmarshal(body, &raw); err != nil { + if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil { return bgptErrJSON(fmt.Errorf("bgpt: parse response: %w", err)), fmt.Errorf("bgpt: parse response: %w", err) } @@ -193,51 +188,149 @@ func (b *BGPTTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool. return bgptErrJSON(err), err } - results := make([]bgptResult, 0, len(raw.Results)) - for _, r := range raw.Results { - results = append(results, bgptResult{ - Title: strVal(r["title"]), - Authors: strVal(r["authors"]), - Journal: strVal(r["journal"]), - Year: strVal(r["year"]), - DOI: strVal(r["doi"]), - URL: strVal(r["url"]), - Abstract: strVal(r["abstract"]), - Methods: firstStr(r, "methods_and_experimental_techniques", "methods"), - SampleSize: firstStr(r, "sample_size_and_population_characteristics", "sample_size_and_population"), - Results: firstStr(r, "results_and_conclusions", "results"), - Limitations: firstStr(r, "paper_limitations_and_biases", "limitations"), - ConflictOfInterest: firstStr(r, "conflict_of_interest_statements", "conflict_of_interest"), - DataAvailability: firstStr(r, "data_availability_statements", "data_availability"), - BlindSpots: strVal(r["study_blindspots"]), - Falsify: strVal(r["how_to_falsify"]), - }) - } - - env := bgptEnv{Results: results} - out, _ := json.Marshal(env) - return string(out), nil + return bgptJSON(bgptEnv{Results: raw.Results}), nil } -func bgptErrJSON(err error) string { - env := bgptEnv{Error: err.Error()} +func mergeBGPTParams(defaults, params bgptParams) bgptParams { + if params.APIKey == "" { + params.APIKey = defaults.APIKey + } + if params.DaysBack == 0 { + params.DaysBack = defaults.DaysBack + } + if params.TopN == 0 { + params.TopN = defaults.TopN + } + return params +} + +func (b *BGPTTool) BuildReferences(_ context.Context, envelope map[string]any) ([]map[string]any, []map[string]any) { + return buildBGPTReferences(envelope) +} + +func (b *BGPTTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + results := envelopeSlice(envelope, "results") + chunks, _ := buildBGPTReferences(envelope) + return map[string]any{ + "formalized_content": renderBGPTReferences(chunks, bgptPromptMaxTokens), + "json": results, + } +} + +func buildBGPTReferences(envelope map[string]any) ([]map[string]any, []map[string]any) { + results := envelopeSlice(envelope, "results") + chunks := make([]map[string]any, 0, len(results)) + docAggs := make([]map[string]any, 0, len(results)) + for _, result := range results { + paper, ok := result.(map[string]any) + if !ok { + continue + } + content := bgptDataImagePattern.ReplaceAllString(formatBGPTPaper(paper), "") + content = truncateBGPTRunes(content, 10000) + if content == "" { + continue + } + documentID := strconv.FormatInt(bgptHashInt(content, 100000000), 10) + displayID := strconv.FormatInt(bgptHashInt(documentID, 500), 10) + title := bgptFirstField(paper, "title") + if title == "-" { + title = "Untitled" + } + resultURL := bgptFirstField(paper, "url", "doi") + if resultURL == "-" { + resultURL = "" + } + chunks = append(chunks, map[string]any{ + "id": displayID, + "chunk_id": documentID, + "content": content, + "doc_id": documentID, + "document_id": documentID, + "docnm_kwd": title, + "document_name": title, + "similarity": 1, + "score": 1, + "url": resultURL, + }) + docAggs = append(docAggs, map[string]any{"doc_name": title, "doc_id": documentID, "count": 1, "url": resultURL}) + } + return chunks, docAggs +} + +func formatBGPTPaper(paper map[string]any) string { + lines := []string{ + "Title: " + bgptFirstField(paper, "title"), + "Authors: " + bgptFirstField(paper, "authors"), + "Journal: " + bgptFirstField(paper, "journal"), + "Year: " + bgptFirstField(paper, "year"), + "DOI: " + bgptFirstField(paper, "doi"), + "Abstract: " + bgptFirstField(paper, "abstract"), + "Methods: " + bgptFirstField(paper, "methods_and_experimental_techniques", "methods"), + "Sample size / population: " + bgptFirstField(paper, "sample_size_and_population_characteristics", "sample_size_and_population", "sample_size"), + "Results: " + bgptFirstField(paper, "results_and_conclusions", "results"), + "Limitations: " + bgptFirstField(paper, "paper_limitations_and_biases", "limitations"), + "Conflicts of interest: " + bgptFirstField(paper, "conflict_of_interest_statements", "conflict_of_interest"), + "Data availability: " + bgptFirstField(paper, "data_availability_statements", "data_availability"), + "Blind spots: " + bgptFirstField(paper, "study_blindspots", "blind_spots"), + "How to falsify: " + bgptFirstField(paper, "how_to_falsify", "falsify"), + } + return strings.Join(lines, "\n") +} + +func bgptFirstField(paper map[string]any, keys ...string) string { + for _, key := range keys { + value, ok := paper[key] + if !ok || value == nil { + continue + } + text := strings.TrimSpace(fmt.Sprint(value)) + if text != "" { + return text + } + } + return "-" +} + +func renderBGPTReferences(chunks []map[string]any, maxTokens int) string { + usedTokens := 0 + blocks := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + content := fmt.Sprint(chunk["content"]) + block := strings.Join([]string{ + "\nID: " + fmt.Sprint(chunk["id"]), + "├── Title: " + bgptNewlinePattern.ReplaceAllString(fmt.Sprint(chunk["document_name"]), " "), + "├── URL: " + bgptNewlinePattern.ReplaceAllString(fmt.Sprint(chunk["url"]), " "), + "└── Content:\n" + content, + }, "\n") + blockTokens := tokenizer.NumTokensFromString(block) + if maxTokens > 0 && float64(usedTokens+blockTokens) > float64(maxTokens)*0.97 { + break + } + usedTokens += blockTokens + blocks = append(blocks, block) + } + return strings.Join(blocks, "\n") +} + +func bgptHashInt(value string, modulus int64) int64 { + digest := sha1.Sum([]byte(value)) + number := new(big.Int).SetBytes(digest[:]) + return new(big.Int).Mod(number, big.NewInt(modulus)).Int64() +} + +func truncateBGPTRunes(value string, limit int) string { + if utf8.RuneCountInString(value) <= limit { + return value + } + return string([]rune(value)[:limit]) +} + +func bgptJSON(env bgptEnv) string { b, _ := json.Marshal(env) return string(b) } -func strVal(v interface{}) string { - if v == nil { - return "" - } - s, _ := v.(string) - return s -} - -func firstStr(m map[string]interface{}, keys ...string) string { - for _, k := range keys { - if s := strVal(m[k]); s != "" { - return s - } - } - return "" +func bgptErrJSON(err error) string { + return bgptJSON(bgptEnv{Error: err.Error()}) } diff --git a/internal/agent/tool/bgpt_test.go b/internal/agent/tool/bgpt_test.go new file mode 100644 index 0000000000..f50ae8f0bf --- /dev/null +++ b/internal/agent/tool/bgpt_test.go @@ -0,0 +1,160 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package tool + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "ragflow/internal/tokenizer" +) + +func TestBGPT_RequestAndRawResults(t *testing.T) { + t.Parallel() + + var body map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = json.NewDecoder(r.Body).Decode(&body) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"results":[{"title":"Paper","url":"https://paper.example","methods_and_experimental_techniques":"RCT","quality_score":0.95}]}`)) + })) + defer srv.Close() + helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)}) + bgpt := NewBGPTToolWith(helper) + out, err := bgpt.InvokableRun(context.Background(), `{"query":" cancer ","top_n":3,"api_key":"key","days_back":30}`) + if err != nil { + t.Fatalf("InvokableRun: %v", err) + } + if body["query"] != "cancer" || body["num_results"] != float64(3) || body["api_key"] != "key" || body["days_back"] != float64(30) { + t.Fatalf("request body = %#v", body) + } + var envelope bgptEnv + if err := json.Unmarshal([]byte(out), &envelope); err != nil { + t.Fatalf("decode output: %v", err) + } + if len(envelope.Results) != 1 || envelope.Results[0]["quality_score"] != float64(0.95) { + t.Fatalf("raw upstream result was narrowed: %#v", envelope.Results) + } +} + +func TestBGPT_InfoAndNodeDefaults(t *testing.T) { + t.Parallel() + + info, err := NewBGPTTool().Info(context.Background()) + if err != nil { + t.Fatalf("Info: %v", err) + } + if info.Name != "bgpt_search" { + t.Fatalf("Info.Name = %q", info.Name) + } + built, err := BuildByName("bgpt", map[string]any{ + "api_key": "stored", + "top_n": "7", + "days_back": float64(14), + "outputs": map[string]any{"json": map[string]any{}}, + }) + if err != nil { + t.Fatalf("BuildByName: %v", err) + } + bgpt := built.(*BGPTTool) + if bgpt.defaults.APIKey != "stored" || bgpt.defaults.TopN != 7 || bgpt.defaults.DaysBack != 14 { + t.Fatalf("defaults = %+v", bgpt.defaults) + } + for _, params := range []map[string]any{{"top_n": 0}, {"days_back": 1.5}, {"api_key": 1}} { + if _, err := BuildByName("bgpt", params); err == nil { + t.Fatalf("BuildByName(%#v) succeeded", params) + } + } +} + +func TestBGPT_EmptyQuery(t *testing.T) { + t.Parallel() + + out, err := NewBGPTTool().InvokableRun(context.Background(), `{"query":""}`) + if err != nil { + t.Fatalf("InvokableRun(empty): %v", err) + } + var envelope bgptEnv + if err := json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 { + t.Fatalf("empty result = %s / %v", out, err) + } +} + +func TestBGPT_ComponentReferencesAndOutputs(t *testing.T) { + t.Parallel() + + bgpt := NewBGPTTool() + spec := bgpt.ComponentSpec() + if query, ok := spec.InputForm["query"].(map[string]any); !ok || query["type"] != "line" { + t.Fatalf("query input form = %#v", spec.InputForm["query"]) + } + envelope := map[string]any{"results": []any{map[string]any{ + "title": "Paper A", + "url": "https://paper.example", + "authors": []any{"Lee", "Kim"}, + "methods_and_experimental_techniques": "RCT", + "sample_size_and_population_characteristics": "120 adults", + "results_and_conclusions": "Improved outcomes", + "quality_score": float64(0.9), + }}} + chunks, docAggs := bgpt.BuildReferences(context.Background(), envelope) + if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["document_name"] != "Paper A" { + t.Fatalf("references = %#v / %#v", chunks, docAggs) + } + content := chunks[0]["content"].(string) + for _, want := range []string{"Methods: RCT", "Sample size / population: 120 adults", "Results: Improved outcomes"} { + if !strings.Contains(content, want) { + t.Fatalf("content missing %q: %q", want, content) + } + } + outputs := bgpt.BuildComponentOutputs(envelope) + results := outputs["json"].([]any) + if results[0].(map[string]any)["quality_score"] != float64(0.9) { + t.Fatalf("raw json was narrowed: %#v", results) + } + if !strings.Contains(outputs["formalized_content"].(string), "Title: Paper A") { + t.Fatalf("formalized_content = %q", outputs["formalized_content"]) + } + if _, exists := envelope["chunks"]; exists { + t.Fatalf("output conversion mutated envelope: %#v", envelope) + } +} + +func TestRenderBGPTReferencesStopsBeforeOverBudgetBlock(t *testing.T) { + t.Parallel() + + chunks := []map[string]any{ + {"id": "1", "document_name": "First", "url": "https://first.example", "content": "first reference content"}, + {"id": "2", "document_name": "Second", "url": "https://second.example", "content": "second reference content"}, + } + firstBlock := renderBGPTReferences(chunks[:1], 0) + firstTokens := tokenizer.NumTokensFromString(firstBlock) + maxTokens := (firstTokens*100 + 96) / 97 + if got := renderBGPTReferences(chunks, maxTokens); got != firstBlock { + t.Fatalf("rendered = %q, want only first block %q", got, firstBlock) + } + if got := renderBGPTReferences(chunks, 1); got != "" { + t.Fatalf("over-budget first block was appended: %q", got) + } + if got := renderBGPTReferences(chunks, 0); !strings.Contains(got, "Title: First") || !strings.Contains(got, "Title: Second") { + t.Fatalf("unlimited rendering dropped blocks: %q", got) + } +} diff --git a/internal/agent/tool/duckduckgo.go b/internal/agent/tool/duckduckgo.go index a6e3a56e64..efc0af0dc6 100644 --- a/internal/agent/tool/duckduckgo.go +++ b/internal/agent/tool/duckduckgo.go @@ -18,34 +18,45 @@ package tool import ( "context" + "crypto/sha1" "encoding/json" "fmt" "html" "io" + "math/big" "net/http" "net/url" "regexp" "strconv" "strings" + "unicode/utf8" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" xhtml "golang.org/x/net/html" + + "ragflow/internal/tokenizer" ) -const duckduckgoToolName = "duckduckgo" +const duckduckgoToolName = "duckduckgo_search" const duckduckgoToolDescription = "Search DuckDuckGo web or news results. Returns results[].{title, url, body}." const duckduckgoChannelGeneral = "general" const duckduckgoChannelNews = "news" +const duckduckgoPromptMaxTokens = 200000 + var duckduckgoSearchEndpoint = "https://duckduckgo.com/html/" var duckduckgoNewsEndpoint = "https://duckduckgo.com/news.js" var duckduckgoNewsBootstrapEndpoint = "https://duckduckgo.com/" var duckduckgoVQDPattern = regexp.MustCompile(`vqd="([^"]+)"`) +var duckduckgoDataImagePattern = regexp.MustCompile(`!?\[[a-z]+\]\(data:image/png;base64,[ 0-9A-Za-z/_=+\-]+\)`) + +var duckduckgoNewlinePattern = regexp.MustCompile(`\n+`) + type duckduckgoParams struct { Query string `json:"query"` Channel string `json:"channel"` @@ -74,18 +85,32 @@ type duckduckgoNewsItem struct { } type DuckDuckGoTool struct { - helper *HTTPHelper + helper *HTTPHelper + defaults duckduckgoParams } +var _ ToolComponent = (*DuckDuckGoTool)(nil) +var _ ReferenceBuilder = (*DuckDuckGoTool)(nil) + func NewDuckDuckGoTool() *DuckDuckGoTool { - return NewDuckDuckGoToolWith(NewHTTPHelper()) + return newDuckDuckGoTool(nil, duckduckgoParams{}) } func NewDuckDuckGoToolWith(h *HTTPHelper) *DuckDuckGoTool { + return newDuckDuckGoTool(h, duckduckgoParams{}) +} + +func newDuckDuckGoTool(h *HTTPHelper, defaults duckduckgoParams) *DuckDuckGoTool { if h == nil { h = NewHTTPHelper() } - return &DuckDuckGoTool{helper: h} + if defaults.TopN == 0 { + defaults.TopN = 10 + } + if defaults.Channel == "" { + defaults.Channel = duckduckgoChannelGeneral + } + return &DuckDuckGoTool{helper: h, defaults: defaults} } func (d *DuckDuckGoTool) Info(_ context.Context) (*schema.ToolInfo, error) { @@ -107,6 +132,26 @@ func (d *DuckDuckGoTool) Info(_ context.Context) (*schema.ToolInfo, error) { }, nil } +func (d *DuckDuckGoTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "query": "Search query.", + "channel": "Search channel: general or news.", + "top_n": "Maximum number of results.", + }, + Outputs: map[string]string{ + "formalized_content": "Rendered DuckDuckGo references for downstream prompts.", + "json": "DuckDuckGo result list.", + }, + InputForm: map[string]any{ + "query": map[string]any{"name": "Query", "type": "line"}, + "channel": map[string]any{ + "name": "Channel", "type": "options", "value": "general", "options": []string{"general", "news"}, + }, + }, + } +} + func buildDuckDuckGoSearchURL(query string, topN int) string { if topN <= 0 { topN = 10 @@ -156,9 +201,9 @@ func (d *DuckDuckGoTool) InvokableRun(ctx context.Context, argsJSON string, _ .. fmt.Errorf("duckduckgo: parse arguments: %w", err) } if strings.TrimSpace(p.Query) == "" { - return duckduckgoErrJSON(fmt.Errorf("query is required")), - fmt.Errorf("duckduckgo: query is required") + return duckduckgoJSON(duckduckgoEnvelope{Results: []duckduckgoResult{}}), nil } + p = mergeDuckDuckGoParams(d.defaults, p) channel := normalizeDuckDuckGoChannel(p.Channel) topN := p.TopN @@ -195,6 +240,106 @@ func (d *DuckDuckGoTool) InvokableRun(ctx context.Context, argsJSON string, _ .. return duckduckgoJSON(duckduckgoEnvelope{Results: results}), nil } +func mergeDuckDuckGoParams(defaults, params duckduckgoParams) duckduckgoParams { + if params.Channel == "" { + params.Channel = defaults.Channel + } + if params.TopN == 0 { + params.TopN = defaults.TopN + } + return params +} + +func (d *DuckDuckGoTool) BuildReferences(_ context.Context, envelope map[string]any) ([]map[string]any, []map[string]any) { + return buildDuckDuckGoReferences(envelope) +} + +func (d *DuckDuckGoTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + results := envelopeSlice(envelope, "results") + chunks, _ := buildDuckDuckGoReferences(envelope) + return map[string]any{ + "formalized_content": renderDuckDuckGoReferences(chunks, duckduckgoPromptMaxTokens), + "json": results, + } +} + +func buildDuckDuckGoReferences(envelope map[string]any) ([]map[string]any, []map[string]any) { + results := envelopeSlice(envelope, "results") + chunks := make([]map[string]any, 0, len(results)) + docAggs := make([]map[string]any, 0, len(results)) + for _, result := range results { + item, ok := result.(map[string]any) + if !ok { + continue + } + content := duckduckgoDataImagePattern.ReplaceAllString(duckduckgoText(item["body"]), "") + content = truncateDuckDuckGoRunes(content, 10000) + if content == "" { + continue + } + documentID := strconv.FormatInt(duckduckgoHashInt(content, 100000000), 10) + displayID := strconv.FormatInt(duckduckgoHashInt(documentID, 500), 10) + title := duckduckgoText(item["title"]) + resultURL := duckduckgoText(item["url"]) + chunks = append(chunks, map[string]any{ + "id": displayID, + "chunk_id": documentID, + "content": content, + "doc_id": documentID, + "document_id": documentID, + "docnm_kwd": title, + "document_name": title, + "similarity": 1, + "score": 1, + "url": resultURL, + }) + docAggs = append(docAggs, map[string]any{"doc_name": title, "doc_id": documentID, "count": 1, "url": resultURL}) + } + return chunks, docAggs +} + +func renderDuckDuckGoReferences(chunks []map[string]any, maxTokens int) string { + usedTokens := 0 + blocks := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + content := duckduckgoText(chunk["content"]) + usedTokens += tokenizer.NumTokensFromString(content) + blocks = append(blocks, strings.Join([]string{ + "\nID: " + duckduckgoText(chunk["id"]), + "├── Title: " + duckduckgoNewlinePattern.ReplaceAllString(duckduckgoText(chunk["document_name"]), " "), + "├── URL: " + duckduckgoNewlinePattern.ReplaceAllString(duckduckgoText(chunk["url"]), " "), + "└── Content:\n" + content, + }, "\n")) + if maxTokens > 0 && float64(maxTokens)*0.97 < float64(usedTokens) { + break + } + } + return strings.Join(blocks, "\n") +} + +func duckduckgoText(value any) string { + if value == nil { + return "" + } + if text, ok := value.(string); ok { + return text + } + return fmt.Sprint(value) +} + +func duckduckgoHashInt(value string, modulus int64) int64 { + digest := sha1.Sum([]byte(value)) + number := new(big.Int).SetBytes(digest[:]) + return new(big.Int).Mod(number, big.NewInt(modulus)).Int64() +} + +func truncateDuckDuckGoRunes(value string, limit int) string { + if utf8.RuneCountInString(value) <= limit { + return value + } + return string([]rune(value)[:limit]) +} + func (d *DuckDuckGoTool) runNewsSearch(ctx context.Context, query string, topN int) (string, error) { vqd, err := d.fetchDuckDuckGoNewsVQD(ctx, query) if err != nil { diff --git a/internal/agent/tool/duckduckgo_test.go b/internal/agent/tool/duckduckgo_test.go index de6ead6dce..435a80158a 100644 --- a/internal/agent/tool/duckduckgo_test.go +++ b/internal/agent/tool/duckduckgo_test.go @@ -226,8 +226,8 @@ func TestDuckDuckGo_Info(t *testing.T) { if err != nil { t.Fatalf("Info: %v", err) } - if info.Name != "duckduckgo" { - t.Errorf("Name = %q, want duckduckgo", info.Name) + if info.Name != "duckduckgo_search" { + t.Errorf("Name = %q, want duckduckgo_search", info.Name) } if !strings.Contains(info.Desc, "DuckDuckGo") { t.Errorf("Desc = %q, want to mention DuckDuckGo", info.Desc) @@ -252,14 +252,15 @@ func TestDuckDuckGo_Info(t *testing.T) { } } -func TestDuckDuckGo_RequiresQuery(t *testing.T) { +func TestDuckDuckGo_EmptyQuery(t *testing.T) { tool := NewDuckDuckGoTool() - _, err := tool.InvokableRun(context.Background(), `{"query":""}`) - if err == nil { - t.Fatal("expected error for empty query") + out, err := tool.InvokableRun(context.Background(), `{"query":""}`) + if err != nil { + t.Fatalf("InvokableRun(empty): %v", err) } - if !strings.Contains(err.Error(), "query") { - t.Errorf("err = %v, want to mention query", err) + var envelope duckduckgoEnvelope + if err := json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 { + t.Fatalf("empty result = %s / %v", out, err) } } @@ -284,7 +285,7 @@ func TestDuckDuckGo_RealReactAgent_ExecutesTool(t *testing.T) { realTool := NewDuckDuckGoTool() mdl := newReactScriptedModel( - "duckduckgo", + "duckduckgo_search", `{"query":"ragflow"}`, "RAGFlow is an open-source RAG engine.", ) @@ -312,12 +313,12 @@ func TestDuckDuckGo_RealReactAgent_ExecutesTool(t *testing.T) { if mdl.turn != 2 { t.Errorf("Generate called %d times, want 2 (tool_call + final)", mdl.turn) } - if len(mdl.boundTools) != 1 || mdl.boundTools[0].Name != "duckduckgo" { + if len(mdl.boundTools) != 1 || mdl.boundTools[0].Name != "duckduckgo_search" { names := make([]string, 0, len(mdl.boundTools)) for _, ti := range mdl.boundTools { names = append(names, ti.Name) } - t.Errorf("tools bound to model = %v, want [duckduckgo]", names) + t.Errorf("tools bound to model = %v, want [duckduckgo_search]", names) } if len(mdl.rounds) < 2 { t.Fatalf("only %d rounds captured, want >= 2", len(mdl.rounds)) @@ -336,3 +337,46 @@ func TestDuckDuckGo_RealReactAgent_ExecutesTool(t *testing.T) { t.Error("test server was never hit; the tool did not actually call the upstream") } } + +func TestDuckDuckGo_ComponentReferencesAndDefaults(t *testing.T) { + t.Parallel() + + built, err := BuildByName("duckduckgo", map[string]any{ + "top_n": float64(4), + "channel": "news", + "outputs": map[string]any{"json": map[string]any{}}, + }) + if err != nil { + t.Fatalf("BuildByName: %v", err) + } + duck := built.(*DuckDuckGoTool) + if duck.defaults.TopN != 4 || duck.defaults.Channel != "news" { + t.Fatalf("defaults = %+v", duck.defaults) + } + for _, params := range []map[string]any{{"top_n": 0}, {"top_n": 1.5}, {"channel": "images"}} { + if _, err := BuildByName("duckduckgo", params); err == nil { + t.Fatalf("BuildByName(%#v) succeeded", params) + } + } + spec := duck.ComponentSpec() + if channel, ok := spec.InputForm["channel"].(map[string]any); !ok || channel["value"] != "general" { + t.Fatalf("channel input form = %#v", spec.InputForm["channel"]) + } + envelope := map[string]any{"results": []any{map[string]any{ + "title": "Story", "url": "https://news.example/story", "body": "Breaking update", + }}} + chunks, docAggs := duck.BuildReferences(context.Background(), envelope) + if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["content"] != "Breaking update" { + t.Fatalf("references = %#v / %#v", chunks, docAggs) + } + outputs := duck.BuildComponentOutputs(envelope) + if results, ok := outputs["json"].([]any); !ok || len(results) != 1 { + t.Fatalf("json output = %#v", outputs["json"]) + } + if !strings.Contains(outputs["formalized_content"].(string), "Breaking update") { + t.Fatalf("formalized_content = %q", outputs["formalized_content"]) + } + if _, exists := envelope["chunks"]; exists { + t.Fatalf("output conversion mutated envelope: %#v", envelope) + } +} diff --git a/internal/agent/tool/email.go b/internal/agent/tool/email.go index f7b19a93ca..ea64304762 100644 --- a/internal/agent/tool/email.go +++ b/internal/agent/tool/email.go @@ -23,12 +23,15 @@ import ( "fmt" "io" "net" + "net/mail" "net/smtp" "strings" "time" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" + + "ragflow/internal/agent/runtime" ) const emailToolName = "email" @@ -40,16 +43,19 @@ const ( emailSessionTimeout = 30 * time.Second ) -// emailParams is the JSON shape the model sends into InvokableRun. +// emailParams contains both the Python model-call inputs and the Canvas node +// configuration used to deliver the message. Info exposes only the former. type emailParams struct { - SMTPHost string `json:"smtp_host"` - SMTPPort int `json:"smtp_port"` - Username string `json:"username"` - Password string `json:"password"` - FromAddr string `json:"from_addr"` - ToAddrs []string `json:"to_addrs"` - Subject string `json:"subject"` - Body string `json:"body"` + SMTPServer string `json:"smtp_server"` + SMTPPort int `json:"smtp_port"` + Email string `json:"email"` + SMTPUsername string `json:"smtp_username"` + Password string `json:"password"` + SenderName string `json:"sender_name"` + ToEmail string `json:"to_email"` + CCEmail string `json:"cc_email"` + Content string `json:"content"` + Subject string `json:"subject"` } // emailEnvelope is what the model sees. @@ -63,12 +69,23 @@ type emailEnvelope struct { // RFC 822 message and submits it via the stdlib net/smtp client. All // authentication modes supported by net/smtp.Auth are available // (PLAIN, LOGIN, CRAM-MD5) by selecting the appropriate creds. -type EmailTool struct{} +type EmailTool struct { + defaults emailParams +} + +var _ ToolComponent = (*EmailTool)(nil) // NewEmailTool returns an EmailTool. There is no shared HTTPHelper // (SMTP is not HTTP), so the constructor is the simplest possible. func NewEmailTool() *EmailTool { - return &EmailTool{} + return newEmailTool(emailParams{SMTPPort: 465}) +} + +func newEmailTool(defaults emailParams) *EmailTool { + if defaults.SMTPPort == 0 { + defaults.SMTPPort = 465 + } + return &EmailTool{defaults: defaults} } // Info returns the tool's metadata for the chat model. @@ -77,57 +94,63 @@ func (e *EmailTool) Info(_ context.Context) (*schema.ToolInfo, error) { Name: emailToolName, Desc: emailToolDescription, ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ - "smtp_host": { + "to_email": { Type: schema.String, - Desc: "SMTP server hostname (e.g. smtp.gmail.com).", + Desc: "The target email address.", Required: true, }, - "smtp_port": { - Type: schema.Integer, - Desc: "SMTP server port (e.g. 587 for STARTTLS, 465 for implicit TLS).", - Required: true, - }, - "username": { + "cc_email": { Type: schema.String, - Desc: "SMTP authentication username. Empty for unauthenticated relay.", + Desc: "Other email addresses to send to, separated by commas.", Required: false, }, - "password": { - Type: schema.String, - Desc: "SMTP authentication password (or app password for Gmail/Yahoo).", - Required: false, - }, - "from_addr": { - Type: schema.String, - Desc: "Sender email address (RFC 5322).", - Required: true, - }, - "to_addrs": { - Type: schema.Array, - Desc: "Recipient email addresses.", - Required: true, - }, "subject": { Type: schema.String, - Desc: "Email subject line.", - Required: true, + Desc: "The subject/title of the email.", + Required: false, }, - "body": { + "content": { Type: schema.String, - Desc: "Email body (plain text).", - Required: true, + Desc: "The content of the email.", + Required: false, }, }), }, nil } +func (e *EmailTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "to_email": "Recipient email address list.", + "cc_email": "Optional CC recipient list.", + "content": "Email body.", + "subject": "Email subject.", + }, + Outputs: map[string]string{ + "success": "Whether the email was sent successfully.", + }, + InputForm: map[string]any{ + "to_email": map[string]any{"name": "To ", "type": "line"}, + "subject": map[string]any{"name": "Subject", "type": "line", "optional": true}, + "cc_email": map[string]any{"name": "CC To", "type": "line", "optional": true}, + }, + } +} + // buildEmailMessage composes the RFC 822 wire format: headers + blank // line + body. Extracted so tests can verify subject / recipient // inclusion without opening a real socket. -func buildEmailMessage(from string, to []string, subject, body string) []byte { +func buildEmailMessage(from, senderName string, to, cc []string, subject, body string) []byte { var b strings.Builder - b.WriteString("From: " + from + "\r\n") + fromHeader := (&mail.Address{ + Name: stripEmailHeaderLineBreaks(senderName), + Address: stripEmailHeaderLineBreaks(from), + }).String() + b.WriteString("From: " + fromHeader + "\r\n") b.WriteString("To: " + strings.Join(to, ", ") + "\r\n") + if len(cc) > 0 { + b.WriteString("Cc: " + strings.Join(cc, ", ") + "\r\n") + } b.WriteString("Subject: " + stripEmailHeaderLineBreaks(subject) + "\r\n") b.WriteString("MIME-Version: 1.0\r\n") b.WriteString("Content-Type: text/plain; charset=UTF-8\r\n") @@ -147,16 +170,35 @@ var sendEmail = sendEmailSMTP // InvokableRun sends the email. func (e *EmailTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) { - var p emailParams + p := e.defaults if err := json.Unmarshal([]byte(argsJSON), &p); err != nil { return emailErrJSON(fmt.Errorf("email: parse arguments: %w", err)), fmt.Errorf("email: parse arguments: %w", err) } + state, _, _ := runtime.GetStateFromContext[*runtime.CanvasState](ctx) + p.ToEmail = runtime.ResolveTemplateForDisplay(p.ToEmail, state) + p.CCEmail = runtime.ResolveTemplateForDisplay(p.CCEmail, state) + p.Subject = stripEmailHeaderLineBreaks(runtime.ResolveTemplateForDisplay(p.Subject, state)) + p.Content = runtime.ResolveTemplateForDisplay(p.Content, state) if err := validateEmailParams(&p); err != nil { return emailErrJSON(err), err } - msg := buildEmailMessage(p.FromAddr, p.ToAddrs, p.Subject, p.Body) + toRecipients := splitEmailList(p.ToEmail) + ccRecipients := splitEmailList(p.CCEmail) + if len(toRecipients) == 0 { + err := fmt.Errorf("email: to_email is required") + return emailErrJSON(err), err + } + subject := p.Subject + if subject == "" { + subject = "No Subject" + } + content := p.Content + if content == "" { + content = "No content provided" + } + msg := buildEmailMessage(p.Email, p.SenderName, toRecipients, ccRecipients, subject, content) if err := sendEmail(ctx, p, msg); err != nil { return emailErrJSON(fmt.Errorf("email: send: %w", err)), fmt.Errorf("email: send: %w", err) @@ -168,6 +210,11 @@ func (e *EmailTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool return emailJSON(emailEnvelope{OK: true}), nil } +func (e *EmailTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + ok, _ := envelope["ok"].(bool) + return map[string]any{"success": ok} +} + func sendEmailSMTP(ctx context.Context, p emailParams, msg []byte) error { if p.SMTPPort == 465 { return sendEmailSMTPS(ctx, p, msg) @@ -176,7 +223,7 @@ func sendEmailSMTP(ctx context.Context, p emailParams, msg []byte) error { } func sendEmailSTARTTLS(ctx context.Context, p emailParams, msg []byte) error { - addr := fmt.Sprintf("%s:%d", p.SMTPHost, p.SMTPPort) + addr := fmt.Sprintf("%s:%d", p.SMTPServer, p.SMTPPort) dialer := &net.Dialer{Timeout: emailDialTimeout} conn, err := dialer.DialContext(ctx, "tcp", addr) if err != nil { @@ -187,7 +234,7 @@ func sendEmailSTARTTLS(ctx context.Context, p emailParams, msg []byte) error { stopWatch := watchEmailContext(ctx, conn) defer stopWatch() - client, err := smtp.NewClient(conn, p.SMTPHost) + client, err := smtp.NewClient(conn, p.SMTPServer) if err != nil { return err } @@ -196,16 +243,17 @@ func sendEmailSTARTTLS(ctx context.Context, p emailParams, msg []byte) error { if err := client.Hello("localhost"); err != nil { return err } - if ok, _ := client.Extension("STARTTLS"); ok { - if err := client.StartTLS(&tls.Config{ServerName: p.SMTPHost, MinVersion: tls.VersionTLS12}); err != nil { - return err - } + if ok, _ := client.Extension("STARTTLS"); !ok { + return fmt.Errorf("email: SMTP server does not advertise STARTTLS") + } + if err := client.StartTLS(&tls.Config{ServerName: p.SMTPServer, MinVersion: tls.VersionTLS12}); err != nil { + return err } return submitEmail(ctx, client, p, msg) } func sendEmailSMTPS(ctx context.Context, p emailParams, msg []byte) error { - addr := fmt.Sprintf("%s:%d", p.SMTPHost, p.SMTPPort) + addr := fmt.Sprintf("%s:%d", p.SMTPServer, p.SMTPPort) dialer := &net.Dialer{Timeout: emailDialTimeout} rawConn, err := dialer.DialContext(ctx, "tcp", addr) if err != nil { @@ -215,14 +263,14 @@ func sendEmailSMTPS(ctx context.Context, p emailParams, msg []byte) error { stopWatch := watchEmailContext(ctx, rawConn) defer stopWatch() - conn := tls.Client(rawConn, &tls.Config{ServerName: p.SMTPHost, MinVersion: tls.VersionTLS12}) + conn := tls.Client(rawConn, &tls.Config{ServerName: p.SMTPServer, MinVersion: tls.VersionTLS12}) if err := conn.HandshakeContext(ctx); err != nil { _ = rawConn.Close() return err } defer conn.Close() - client, err := smtp.NewClient(conn, p.SMTPHost) + client, err := smtp.NewClient(conn, p.SMTPServer) if err != nil { return err } @@ -235,15 +283,19 @@ func sendEmailSMTPS(ctx context.Context, p emailParams, msg []byte) error { } func submitEmail(ctx context.Context, client *smtp.Client, p emailParams, msg []byte) error { - if p.Username != "" { - if err := client.Auth(smtp.PlainAuth("", p.Username, p.Password, p.SMTPHost)); err != nil { + username := p.SMTPUsername + if username == "" { + username = p.Email + } + if username != "" && p.Password != "" { + if err := client.Auth(smtp.PlainAuth("", username, p.Password, p.SMTPServer)); err != nil { return err } } - if err := client.Mail(p.FromAddr); err != nil { + if err := client.Mail(p.Email); err != nil { return err } - for _, addr := range p.ToAddrs { + for _, addr := range emailRecipients(p.ToEmail, p.CCEmail) { if err := client.Rcpt(addr); err != nil { return err } @@ -293,20 +345,34 @@ func watchEmailContext(ctx context.Context, conn net.Conn) func() { // addresses, but the common case (empty / missing) is caught here. func validateEmailParams(p *emailParams) error { switch { - case p.SMTPHost == "": - return fmt.Errorf("email: smtp_host is required") + case p.SMTPServer == "": + return fmt.Errorf("email: smtp_server is required") case p.SMTPPort <= 0 || p.SMTPPort > 65535: return fmt.Errorf("email: smtp_port must be in [1, 65535]") - case p.FromAddr == "": - return fmt.Errorf("email: from_addr is required") - case len(p.ToAddrs) == 0: - return fmt.Errorf("email: to_addrs is required and must be non-empty") - case p.Subject == "": - return fmt.Errorf("email: subject is required") + case p.Email == "": + return fmt.Errorf("email: email is required") } return nil } +func emailRecipients(toEmail, ccEmail string) []string { + recipients := splitEmailList(toEmail) + return append(recipients, splitEmailList(ccEmail)...) +} + +func splitEmailList(value string) []string { + parts := strings.FieldsFunc(value, func(r rune) bool { + return r == ',' || r == ';' || r == '\n' || r == '\r' + }) + out := make([]string, 0, len(parts)) + for _, part := range parts { + if recipient := strings.TrimSpace(part); recipient != "" { + out = append(out, recipient) + } + } + return out +} + func emailJSON(env emailEnvelope) string { b, err := json.Marshal(env) if err != nil { diff --git a/internal/agent/tool/email_test.go b/internal/agent/tool/email_test.go index 4ae1e6d157..1e2142d395 100644 --- a/internal/agent/tool/email_test.go +++ b/internal/agent/tool/email_test.go @@ -25,6 +25,8 @@ import ( "strings" "testing" "time" + + "ragflow/internal/agent/runtime" ) func TestEmail_BuildMessage(t *testing.T) { @@ -32,15 +34,18 @@ func TestEmail_BuildMessage(t *testing.T) { msg := buildEmailMessage( "alice@example.com", - []string{"bob@example.com", "carol@example.com"}, + "Alice Sender", + []string{"bob@example.com"}, + []string{"carol@example.com"}, "Hello, world", "Body of the message.", ) s := string(msg) for _, want := range []string{ - "From: alice@example.com", - "To: bob@example.com, carol@example.com", + `From: "Alice Sender" `, + "To: bob@example.com\r\n", + "Cc: carol@example.com\r\n", "Subject: Hello, world", "Content-Type: text/plain; charset=UTF-8", "Body of the message.", @@ -49,159 +54,165 @@ func TestEmail_BuildMessage(t *testing.T) { t.Errorf("message missing %q\n--- message ---\n%s\n---", want, s) } } + if strings.Contains(s, "To: bob@example.com, carol@example.com") { + t.Fatalf("CC recipient leaked into To header:\n%s", s) + } // RFC 822 mandates a blank line between headers and body. if !strings.Contains(s, "\r\n\r\n") { t.Errorf("message missing blank line between headers and body\n%s", s) } } -func TestEmail_SendAgainstMockSMTP(t *testing.T) { +func TestEmail_SendBuildsDistinctHeadersAndEnvelopeRecipients(t *testing.T) { + originalSendEmail := sendEmail + t.Cleanup(func() { sendEmail = originalSendEmail }) + var sentParams emailParams + var sentMessage []byte + sendEmail = func(_ context.Context, p emailParams, msg []byte) error { + sentParams = p + sentMessage = append([]byte(nil), msg...) + return nil + } + + built, err := BuildByName("email", map[string]any{ + "smtp_server": "smtp.example.com", + "smtp_port": 587, + "email": "alice@example.com", + "sender_name": "Alice Sender", + }) + if err != nil { + t.Fatalf("BuildByName(email): %v", err) + } + args := map[string]any{ + "to_email": "bob@example.com", + "cc_email": "carol@example.com, dave@example.com", + "subject": "Test {sys.date}", + "content": "Test body content.", + } + argsJSON, _ := json.Marshal(args) + state := runtime.NewCanvasState("run-email", "task-email") + state.Sys["date"] = "2026-07-15" + out, err := built.(*EmailTool).InvokableRun(runtime.WithState(context.Background(), state), string(argsJSON)) + if err != nil { + t.Fatalf("InvokableRun: %v", err) + } + var env emailEnvelope + if err := json.Unmarshal([]byte(out), &env); err != nil || !env.OK || env.Error != "" { + t.Fatalf("output = %s, decode error = %v", out, err) + } + + message := string(sentMessage) + for _, want := range []string{ + `From: "Alice Sender" `, + "To: bob@example.com\r\n", + "Cc: carol@example.com, dave@example.com\r\n", + "Subject: Test 2026-07-15", + "Test body content.", + } { + if !strings.Contains(message, want) { + t.Fatalf("message missing %q:\n%s", want, message) + } + } + if got := emailRecipients(sentParams.ToEmail, sentParams.CCEmail); strings.Join(got, ",") != "bob@example.com,carol@example.com,dave@example.com" { + t.Fatalf("SMTP envelope recipients = %#v", got) + } +} + +func TestEmail_STARTTLSRequiredBeforeSubmission(t *testing.T) { t.Parallel() - // Spin up a minimal SMTP server: read commands, respond 250 to - // everything, and copy the DATA payload bytes so the test can - // inspect them. - ln, err := net.Listen("tcp", "127.0.0.1:0") + listener, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatalf("Listen: %v", err) } - defer ln.Close() - - var receivedData strings.Builder - done := make(chan struct{}) + defer listener.Close() + commands := make(chan []string, 1) go func() { - defer close(done) - conn, err := ln.Accept() - if err != nil { + conn, acceptErr := listener.Accept() + if acceptErr != nil { + commands <- nil return } defer conn.Close() - _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) + _ = conn.SetDeadline(time.Now().Add(2 * time.Second)) reader := bufio.NewReader(conn) writer := bufio.NewWriter(conn) - // Greeting _, _ = writer.WriteString("220 mock-smtp ready\r\n") _ = writer.Flush() - inData := false + var received []string for { - line, err := reader.ReadString('\n') - if err != nil { + line, readErr := reader.ReadString('\n') + if readErr != nil { + commands <- received return } - up := strings.ToUpper(strings.TrimSpace(line)) - switch { - case strings.HasPrefix(up, "EHLO"), strings.HasPrefix(up, "HELO"): - _, _ = writer.WriteString("250-mock-smtp\r\n250 OK\r\n") - _ = writer.Flush() - case strings.HasPrefix(up, "MAIL FROM:"), strings.HasPrefix(up, "RCPT TO:"): - _, _ = writer.WriteString("250 OK\r\n") - _ = writer.Flush() - case strings.HasPrefix(up, "DATA"): - _, _ = writer.WriteString("354 End data with .\r\n") - _ = writer.Flush() - inData = true - case inData && strings.TrimSpace(line) == ".": - _, _ = writer.WriteString("250 Queued\r\n") - _ = writer.Flush() - inData = false - case inData: - receivedData.WriteString(line) - case strings.HasPrefix(up, "QUIT"): - _, _ = writer.WriteString("221 Bye\r\n") - _ = writer.Flush() - return - default: - _, _ = writer.WriteString("250 OK\r\n") + command := strings.ToUpper(strings.TrimSpace(line)) + received = append(received, command) + if strings.HasPrefix(command, "EHLO") || strings.HasPrefix(command, "HELO") { + _, _ = writer.WriteString("250 mock-smtp\r\n") _ = writer.Flush() } } }() - host, port, err := net.SplitHostPort(ln.Addr().String()) + host, port, err := net.SplitHostPort(listener.Addr().String()) if err != nil { t.Fatalf("SplitHostPort: %v", err) } - _ = host - var portInt int - _, _ = fmt.Sscanf(port, "%d", &portInt) - - tool := NewEmailTool() - args := map[string]any{ - "smtp_host": "127.0.0.1", - "smtp_port": portInt, - "from_addr": "alice@example.com", - "to_addrs": []string{"bob@example.com"}, - "subject": "Test Subject", - "body": "Test body content.", - } - argsJSON, _ := json.Marshal(args) - out, err := tool.InvokableRun(context.Background(), string(argsJSON)) - if err != nil { - t.Fatalf("InvokableRun: %v", err) + var portNumber int + _, _ = fmt.Sscanf(port, "%d", &portNumber) + err = sendEmailSTARTTLS(context.Background(), emailParams{ + SMTPServer: host, SMTPPort: portNumber, Email: "alice@example.com", + ToEmail: "bob@example.com", + }, []byte("message")) + if err == nil || !strings.Contains(err.Error(), "does not advertise STARTTLS") { + t.Fatalf("err = %v", err) } - var env emailEnvelope - if jerr := json.Unmarshal([]byte(out), &env); jerr != nil { - t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out) - } - if env.Error != "" { - t.Errorf("Error = %q, want empty", env.Error) - } - if !env.OK { - t.Errorf("OK = false, want true") - } - - // Wait for the mock server to finish. select { - case <-done: - case <-time.After(2 * time.Second): - t.Fatal("mock SMTP server did not close in time") - } - - if !strings.Contains(receivedData.String(), "Subject: Test Subject") { - t.Errorf("mock server did not receive subject\n--- data ---\n%s\n---", - receivedData.String()) - } - if !strings.Contains(receivedData.String(), "bob@example.com") { - t.Errorf("mock server did not receive recipient\n--- data ---\n%s\n---", - receivedData.String()) - } - if !strings.Contains(receivedData.String(), "Test body content.") { - t.Errorf("mock server did not receive body\n--- data ---\n%s\n---", - receivedData.String()) + case received := <-commands: + for _, command := range received { + if strings.HasPrefix(command, "MAIL FROM") || strings.HasPrefix(command, "RCPT TO") || command == "DATA" { + t.Fatalf("message submission started without STARTTLS: %#v", received) + } + } + case <-time.After(3 * time.Second): + t.Fatal("mock SMTP server did not finish") } } func TestEmail_RequiresFields(t *testing.T) { t.Parallel() - tool := NewEmailTool() - cases := []struct { name string + tool *EmailTool args string wantErr string }{ { - name: "missing smtp_host", - args: `{"smtp_port":587,"from_addr":"a@b","to_addrs":["c@d"],"subject":"s","body":"b"}`, - wantErr: "smtp_host", + name: "missing smtp_server", + tool: newEmailTool(emailParams{SMTPPort: 587, Email: "a@b"}), + args: `{"to_email":"c@d","subject":"s","content":"b"}`, + wantErr: "smtp_server", }, { - name: "missing to_addrs", - args: `{"smtp_host":"x","smtp_port":587,"from_addr":"a@b","subject":"s","body":"b"}`, - wantErr: "to_addrs", + name: "missing to_email", + tool: newEmailTool(emailParams{SMTPServer: "x", SMTPPort: 587, Email: "a@b"}), + args: `{"subject":"s","content":"b"}`, + wantErr: "to_email", }, { name: "bad smtp_port", - args: `{"smtp_host":"x","smtp_port":0,"from_addr":"a@b","to_addrs":["c@d"],"subject":"s","body":"b"}`, + tool: newEmailTool(emailParams{SMTPServer: "x", SMTPPort: -1, Email: "a@b"}), + args: `{"to_email":"c@d","subject":"s","content":"b"}`, wantErr: "smtp_port", }, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { t.Parallel() - _, err := tool.InvokableRun(context.Background(), tc.args) + _, err := tc.tool.InvokableRun(context.Background(), tc.args) if err == nil { t.Fatalf("expected error for %s", tc.name) } @@ -226,4 +237,71 @@ func TestEmail_Info(t *testing.T) { if !strings.Contains(info.Desc, "SMTP") { t.Errorf("Desc = %q, want to mention SMTP", info.Desc) } + schemaJSON, err := info.ToJSONSchema() + if err != nil { + t.Fatalf("ToJSONSchema: %v", err) + } + raw, err := json.Marshal(schemaJSON) + if err != nil { + t.Fatalf("marshal schema: %v", err) + } + for _, runtimeField := range []string{"to_email", "cc_email", "content", "subject"} { + if !strings.Contains(string(raw), `"`+runtimeField+`"`) { + t.Errorf("schema missing runtime field %q: %s", runtimeField, raw) + } + } + for _, configField := range []string{"smtp_server", "smtp_port", "email", "password", "sender_name"} { + if strings.Contains(string(raw), `"`+configField+`"`) { + t.Errorf("schema leaked node config %q: %s", configField, raw) + } + } +} + +func TestEmail_ComponentContractAndFactory(t *testing.T) { + t.Parallel() + + built, err := BuildByName("email", map[string]any{ + "smtp_server": "smtp.example.com", + "smtp_port": "587", + "email": "sender@example.com", + "password": "secret", + "sender_name": "Sender", + "outputs": map[string]any{"success": map[string]any{}}, + "setups": map[string]any{"to_email": "configured@example.com"}, + }) + if err != nil { + t.Fatalf("BuildByName(email): %v", err) + } + emailTool := built.(*EmailTool) + if emailTool.defaults.SMTPPort != 587 || emailTool.defaults.SMTPServer != "smtp.example.com" { + t.Fatalf("node defaults = %#v", emailTool.defaults) + } + spec := emailTool.ComponentSpec() + if _, ok := spec.Inputs["to_email"]; !ok { + t.Fatalf("component inputs = %#v", spec.Inputs) + } + if _, ok := spec.Outputs["success"]; !ok { + t.Fatalf("component outputs = %#v", spec.Outputs) + } + toEmail, ok := spec.InputForm["to_email"].(map[string]any) + if !ok || toEmail["name"] != "To " || toEmail["type"] != "line" { + t.Fatalf("to_email input form = %#v", spec.InputForm["to_email"]) + } + if outputs := emailTool.BuildComponentOutputs(map[string]any{"ok": true, "provider": "smtp"}); outputs["success"] != true { + t.Fatalf("component outputs = %#v", outputs) + } +} + +func TestEmail_BuildByNameRejectsInvalidNodeParams(t *testing.T) { + t.Parallel() + + for _, params := range []map[string]any{ + {"smtp_port": float64(1.5)}, + {"smtp_port": 70000}, + {"smtp_server": 1}, + } { + if _, err := BuildByName("email", params); err == nil { + t.Fatalf("BuildByName(email, %#v) succeeded, want error", params) + } + } } diff --git a/internal/agent/tool/exesql.go b/internal/agent/tool/exesql.go index 851b3770a3..b30344afb2 100644 --- a/internal/agent/tool/exesql.go +++ b/internal/agent/tool/exesql.go @@ -48,6 +48,7 @@ import ( "fmt" "net" "regexp" + "sort" "strconv" "strings" "time" @@ -59,6 +60,8 @@ import ( "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" + + "ragflow/internal/agent/runtime" ) // ExeSQL-specific errors. ErrExeSQLDAOMissing is surfaced when @@ -155,6 +158,8 @@ func NewExeSQLConnParams(params map[string]any) (ExeSQLConnParams, error) { } if v, ok := intParam(params, "max_records"); ok { conn.MaxRecords = v + } else if v, ok := intParam(params, "top_n"); ok { + conn.MaxRecords = v } if conn.DBType == "" || conn.Host == "" || conn.Username == "" || conn.Database == "" { return conn, fmt.Errorf("ExeSQL: missing required connection params (db_type/host/database/username)") @@ -209,6 +214,8 @@ type ExeSQLTool struct { dialer exesqlDialer } +var _ ToolComponent = (*ExeSQLTool)(nil) + // NewExeSQLTool returns an ExeSQLTool wired to the given connection // params. The dialer defaults to `sql.Open`; tests can pass a // sqlmock-backed dialer via WithExeSQLDialer. @@ -261,6 +268,21 @@ func (e *ExeSQLTool) Info(_ context.Context) (*schema.ToolInfo, error) { }, nil } +func (e *ExeSQLTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "sql": "SQL statement to execute (SELECT-only; DML/DDL rejected).", + }, + Outputs: map[string]string{ + "formalized_content": "SQL result rendered as Markdown.", + "json": "Raw SQL statement results.", + }, + InputForm: map[string]any{ + "sql": map[string]any{"name": "SQL", "type": "line"}, + }, + } +} + // InvokableRun validates the SQL, opens a fresh connection scoped to // the tool's params, executes each semicolon-separated statement, and // returns the rows. Per-statement errors do not abort the node: they @@ -275,6 +297,13 @@ func (e *ExeSQLTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ return exesqlErrorResult(fmt.Errorf("exesql: parse arguments: %w", err)), fmt.Errorf("exesql: parse arguments: %w", err) } + if state, _, stateErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); stateErr == nil && state != nil { + resolved, resolveErr := runtime.ResolveTemplate(args.SQL, state) + if resolveErr != nil { + return exesqlErrorResult(resolveErr), resolveErr + } + args.SQL = resolved + } if strings.TrimSpace(args.SQL) == "" { return exesqlErrorResult(errors.New("exesql: empty sql")), errors.New("exesql: empty sql") } @@ -330,6 +359,73 @@ func (e *ExeSQLTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ return exesqlMarshalResult(res) } +func (e *ExeSQLTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + rows := envelopeSlice(envelope, "rows") + columns := envelopeSlice(envelope, "columns") + jsonResult := make([]any, 0, 1) + if len(rows) == 1 { + if row, ok := rows[0].(map[string]any); ok && len(row) == 1 { + if _, hasContent := row["content"]; hasContent { + jsonResult = append(jsonResult, row) + } + } + } + if len(jsonResult) == 0 && len(rows) > 0 { + jsonResult = append(jsonResult, rows) + } + return map[string]any{ + "formalized_content": renderExeSQLMarkdown(columns, rows), + "json": jsonResult, + } +} + +func renderExeSQLMarkdown(columns, rows []any) string { + if len(rows) == 0 { + return "" + } + for _, value := range rows { + if row, ok := value.(map[string]any); ok && len(row) == 1 { + if message, exists := row["content"]; exists { + return fmt.Sprint(message) + } + } + } + columnNames := make([]string, 0, len(columns)) + for _, column := range columns { + columnNames = append(columnNames, fmt.Sprint(column)) + } + if len(columnNames) == 0 { + if first, ok := rows[0].(map[string]any); ok { + for column := range first { + columnNames = append(columnNames, column) + } + sort.Strings(columnNames) + } + } + if len(columnNames) == 0 { + return "" + } + var builder strings.Builder + fmt.Fprintf(&builder, "| %s |\n", strings.Join(columnNames, " | ")) + separators := make([]string, len(columnNames)) + for i := range separators { + separators[i] = "---" + } + fmt.Fprintf(&builder, "| %s |\n", strings.Join(separators, " | ")) + for _, value := range rows { + row, ok := value.(map[string]any) + if !ok { + continue + } + cells := make([]string, len(columnNames)) + for i, column := range columnNames { + cells[i] = strings.ReplaceAll(strings.ReplaceAll(fmt.Sprint(row[column]), "|", "\\|"), "\n", "
") + } + fmt.Fprintf(&builder, "| %s |\n", strings.Join(cells, " | ")) + } + return strings.TrimSuffix(builder.String(), "\n") +} + // exesqlExecute splits the SQL on statement-delimiting semicolons and runs // each statement independently. A failing statement is recorded as an // error entry but does not abort subsequent statements — this is diff --git a/internal/agent/tool/exesql_test.go b/internal/agent/tool/exesql_test.go index 7890d86b38..c3f5f7ce0e 100644 --- a/internal/agent/tool/exesql_test.go +++ b/internal/agent/tool/exesql_test.go @@ -30,6 +30,8 @@ import ( "github.com/cloudwego/eino/compose" "github.com/cloudwego/eino/flow/agent/react" "github.com/cloudwego/eino/schema" + + "ragflow/internal/agent/runtime" ) // testConn is a fully-populated connection params struct used by @@ -357,6 +359,56 @@ func TestExeSQL_UsesConfiguredSQLDefault(t *testing.T) { } } +func TestExeSQL_ComponentContractAndTemplateResolution(t *testing.T) { + dialer, mock, cleanup := sqlmockDialer(t) + defer cleanup() + mock.ExpectPing() + mock.ExpectQuery("SELECT id FROM orders WHERE status = 'Completed'"). + WillReturnRows(sqlmock.NewRows([]string{"id", "status"}).AddRow(1, "Completed")) + conn := testConn() + conn.SQL = "{Agent:Result@content}" + exesql := NewExeSQLTool(conn).WithExeSQLDialer(dialer) + state := runtime.NewCanvasState("run", "task") + state.SetVar("Agent:Result", "content", "SELECT id FROM orders WHERE status = 'Completed'") + out, err := exesql.InvokableRun(runtime.WithState(context.Background(), state), `{}`) + if err != nil { + t.Fatalf("InvokableRun: %v", err) + } + var envelope map[string]any + if err := json.Unmarshal([]byte(out), &envelope); err != nil { + t.Fatalf("decode output: %v", err) + } + outputs := exesql.BuildComponentOutputs(envelope) + if !strings.Contains(outputs["formalized_content"].(string), "Completed") { + t.Fatalf("formalized_content = %q", outputs["formalized_content"]) + } + jsonRows, ok := outputs["json"].([]any) + if !ok || len(jsonRows) != 1 { + t.Fatalf("json output = %#v", outputs["json"]) + } + spec := exesql.ComponentSpec() + if sqlInput, ok := spec.InputForm["sql"].(map[string]any); !ok || sqlInput["type"] != "line" { + t.Fatalf("sql input form = %#v", spec.InputForm["sql"]) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatalf("sql expectations: %v", err) + } +} + +func TestExeSQL_BuildByNameAcceptsCanvasShape(t *testing.T) { + built, err := BuildByName("execute_sql", map[string]any{ + "database": "demo", "username": "root", "host": "db.example.com", "port": float64(3306), "password": "secret", + "top_n": float64(50), "sql": "SELECT 1", "outputs": map[string]any{"json": map[string]any{}}, + }) + if err != nil { + t.Fatalf("BuildByName: %v", err) + } + exesql := built.(*ExeSQLTool) + if exesql.conn.DBType != "mysql" || exesql.conn.Port != 3306 || exesql.conn.MaxRecords != 50 || exesql.conn.SQL != "SELECT 1" { + t.Fatalf("connection defaults = %+v", exesql.conn) + } +} + func TestExeSQL_ExecuteSelect_ReturnsRows(t *testing.T) { t.Parallel() diff --git a/internal/agent/tool/github.go b/internal/agent/tool/github.go index 7589f67270..4c6a3c800b 100644 --- a/internal/agent/tool/github.go +++ b/internal/agent/tool/github.go @@ -18,14 +18,21 @@ package tool import ( "context" + "crypto/sha1" "encoding/json" "fmt" + "math/big" "net/http" "net/url" + "strconv" + "strings" "time" + "unicode/utf8" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" + + "ragflow/internal/tokenizer" ) const githubToolName = "github_search" @@ -33,10 +40,13 @@ const githubToolName = "github_search" const githubToolDescription = "GitHub repository search finds repositories, projects, and codebases hosted on GitHub." const ( - defaultGitHubTopN = 10 - maxGitHubTopN = 100 + defaultGitHubTopN = 10 + maxGitHubTopN = 100 + githubPromptMaxTokens = 200000 ) +const githubQueryDescription = "The search keywords to execute with GitHub. Use the most important terms and synonyms from the original request." + // githubParams mirrors Python GitHubParam. Info() exposes only Query to the // model, while TopN is a canvas-side configuration value merged with defaults. type githubParams struct { @@ -66,6 +76,10 @@ type GitHubTool struct { defaults githubParams } +var _ ToolInvoker = (*GitHubTool)(nil) +var _ ToolComponent = (*GitHubTool)(nil) +var _ ReferenceBuilder = (*GitHubTool)(nil) + // NewGitHubTool returns a GitHubTool using the default HTTPHelper. func NewGitHubTool() *GitHubTool { return NewGitHubToolWithDefaults(nil, githubParams{}) @@ -101,7 +115,7 @@ func (g *GitHubTool) Info(_ context.Context) (*schema.ToolInfo, error) { ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ "query": { Type: schema.String, - Desc: "The search keywords to execute with GitHub. Use the most important terms and synonyms from the original request.", + Desc: githubQueryDescription, Required: true, }, }), @@ -130,8 +144,7 @@ func (g *GitHubTool) InvokableRun(ctx context.Context, argsJSON string, _ ...too fmt.Errorf("github: parse arguments: %w", err) } if p.Query == "" { - return githubErrJSON(fmt.Errorf("query is required")), - fmt.Errorf("github: query is required") + return githubJSON(githubEnvelope{Results: []map[string]any{}}), nil } p = mergeGitHubDefaults(g.defaults, p) @@ -178,3 +191,135 @@ func githubJSON(env githubEnvelope) string { func githubErrJSON(err error) string { return githubJSON(githubEnvelope{Error: err.Error()}) } + +// ComponentSpec returns the Python-compatible GitHub Canvas surface. +func (g *GitHubTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "query": githubQueryDescription, + }, + Outputs: map[string]string{ + "formalized_content": "GitHub repositories formatted for downstream prompts.", + "json": "Raw GitHub repository items.", + }, + InputForm: map[string]any{ + "query": map[string]any{ + "type": "line", + "name": "Query", + }, + }, + } +} + +// BuildReferences creates the chunks and document aggregates Python's +// ToolBase._retrieve_chunks records for GitHub results. +func (g *GitHubTool) BuildReferences(_ context.Context, envelope map[string]any) ([]map[string]any, []map[string]any) { + return buildGitHubReferences(envelope) +} + +func buildGitHubReferences(envelope map[string]any) ([]map[string]any, []map[string]any) { + results := envelopeSlice(envelope, "results") + chunks := make([]map[string]any, 0, len(results)) + docAggs := make([]map[string]any, 0, len(results)) + for _, result := range results { + repository, ok := result.(map[string]any) + if !ok { + continue + } + content := truncateGitHubRunes(githubValueString(repository["description"])+"\n stars:"+githubValueString(repository["watchers"]), 10000) + if content == "" { + continue + } + documentID := strconv.FormatInt(githubHashInt(content, 100000000), 10) + title := githubValueString(repository["name"]) + resultURL := githubValueString(repository["html_url"]) + displayID := strconv.FormatInt(githubHashInt(documentID, 500), 10) + chunks = append(chunks, map[string]any{ + "id": displayID, + "chunk_id": documentID, + "content": content, + "doc_id": documentID, + "document_id": documentID, + "docnm_kwd": title, + "document_name": title, + "similarity": 1, + "score": 1, + "url": resultURL, + }) + docAggs = append(docAggs, map[string]any{ + "doc_name": title, + "doc_id": documentID, + "count": 1, + "url": resultURL, + }) + } + return chunks, docAggs +} + +// BuildComponentOutputs constructs GitHub's complete Canvas output map. +func (g *GitHubTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + results := envelopeSlice(envelope, "results") + chunks, _ := buildGitHubReferences(envelope) + return map[string]any{ + "json": results, + "formalized_content": renderGitHubReferences(chunks), + } +} + +func renderGitHubReferences(chunks []map[string]any) string { + chunks = limitGitHubReferences(chunks, githubPromptMaxTokens) + blocks := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + blocks = append(blocks, strings.Join([]string{ + "\nID: " + githubValueString(chunk["id"]), + "├── Title: " + githubValueString(chunk["docnm_kwd"]), + "├── URL: " + githubValueString(chunk["url"]), + "└── Content:\n" + githubValueString(chunk["content"]), + }, "\n")) + } + return strings.Join(blocks, "\n") +} + +func limitGitHubReferences(chunks []map[string]any, maxTokens int) []map[string]any { + if maxTokens <= 0 { + return nil + } + usedTokens := 0 + for index, chunk := range chunks { + content := githubValueString(chunk["content"]) + if content == "" { + continue + } + usedTokens += tokenizer.NumTokensFromString(content) + if float64(maxTokens)*0.97 < float64(usedTokens) { + return chunks[:index+1] + } + } + return chunks +} + +func githubValueString(value any) string { + if value == nil { + return "None" + } + if boolean, ok := value.(bool); ok { + if boolean { + return "True" + } + return "False" + } + return fmt.Sprint(value) +} + +func githubHashInt(value string, modulus int64) int64 { + sum := sha1.Sum([]byte(value)) + number := new(big.Int).SetBytes(sum[:]) + return new(big.Int).Mod(number, big.NewInt(modulus)).Int64() +} + +func truncateGitHubRunes(value string, limit int) string { + if utf8.RuneCountInString(value) <= limit { + return value + } + return string([]rune(value)[:limit]) +} diff --git a/internal/agent/tool/github_test.go b/internal/agent/tool/github_test.go index e18cd81009..281e5c6f5f 100644 --- a/internal/agent/tool/github_test.go +++ b/internal/agent/tool/github_test.go @@ -147,21 +147,30 @@ func TestGitHub_ParseResponse(t *testing.T) { } } -func TestGitHub_RequiresQuery(t *testing.T) { +func TestGitHub_EmptyQueryReturnsEmptyResults(t *testing.T) { t.Parallel() tool := NewGitHubTool() - _, err := tool.InvokableRun(context.Background(), `{"query":""}`) - if err == nil { - t.Fatal("expected error for empty query") + out, err := tool.InvokableRun(context.Background(), `{"query":""}`) + if err != nil { + t.Fatalf("InvokableRun(empty query): %v", err) } - if !strings.Contains(err.Error(), "query") { - t.Errorf("err = %v, want to mention query", err) + var envelope githubEnvelope + if err := json.Unmarshal([]byte(out), &envelope); err != nil { + t.Fatalf("decode empty result: %v", err) + } + if len(envelope.Results) != 0 || envelope.Error != "" { + t.Fatalf("empty query result = %#v", envelope) } } func TestGitHub_BuildByNameUsesPythonNodeParams(t *testing.T) { - built, err := BuildByName("github", map[string]any{"top_n": float64(17)}) + built, err := BuildByName("github", map[string]any{ + "top_n": float64(17), + "query": "runtime query", + "outputs": map[string]any{"json": map[string]any{}}, + "setups": map[string]any{"query": "configured query"}, + }) if err != nil { t.Fatalf("BuildByName(github): %v", err) } @@ -187,8 +196,74 @@ func TestGitHub_BuildByNameUsesPythonNodeParams(t *testing.T) { if _, err := BuildByName("github", map[string]any{"top_n": 101}); err == nil { t.Fatal("BuildByName(github) accepted top_n above GitHub's per_page limit") } - if _, err := BuildByName("github", map[string]any{"max_results": 5}); err == nil { - t.Fatal("BuildByName(github) accepted removed max_results parameter") + ignored, err := BuildByName("github", map[string]any{"max_results": 5}) + if err != nil { + t.Fatalf("BuildByName(github) rejected unrelated Canvas params: %v", err) + } + if ignored.(*GitHubTool).defaults.TopN != defaultGitHubTopN { + t.Fatalf("unrelated params changed top_n: %d", ignored.(*GitHubTool).defaults.TopN) + } +} + +func TestGitHub_ComponentContractMatchesPython(t *testing.T) { + github := NewGitHubTool() + spec := github.ComponentSpec() + if _, ok := spec.Outputs["json"]; !ok { + t.Fatalf("component outputs missing json: %#v", spec.Outputs) + } + if _, ok := spec.Outputs["formalized_content"]; !ok { + t.Fatalf("component outputs missing formalized_content: %#v", spec.Outputs) + } + if query, ok := spec.InputForm["query"].(map[string]any); !ok || query["name"] != "Query" || query["type"] != "line" { + t.Fatalf("query input form = %#v", spec.InputForm["query"]) + } +} + +func TestGitHub_ReferencesAndOutputsPreserveRawResults(t *testing.T) { + github := NewGitHubTool() + results := []any{map[string]any{ + "name": "ragflow", + "html_url": "https://github.com/infiniflow/ragflow", + "description": "RAG engine", + "watchers": float64(12000), + "private": false, + }} + repository := results[0].(map[string]any) + if repository["private"] != false { + t.Fatalf("raw repository fields were lost: %#v", repository) + } + envelope := map[string]any{"results": results} + + chunks, docAggs := github.BuildReferences(context.Background(), envelope) + if len(chunks) != 1 || len(docAggs) != 1 { + t.Fatalf("references = %#v / %#v", chunks, docAggs) + } + if chunks[0]["document_name"] != "ragflow" || chunks[0]["similarity"] != 1 { + t.Fatalf("reference metadata = %#v", chunks[0]) + } + outputs := github.BuildComponentOutputs(envelope) + if _, exists := envelope["chunks"]; exists { + t.Fatalf("component output conversion mutated the tool envelope: %#v", envelope) + } + if results, ok := outputs["json"].([]any); !ok || len(results) != 1 { + t.Fatalf("component json output = %#v", outputs["json"]) + } + rendered, _ := outputs["formalized_content"].(string) + for _, want := range []string{"Title: ragflow", "URL: https://github.com/infiniflow/ragflow", "RAG engine\n stars:12000"} { + if !strings.Contains(rendered, want) { + t.Fatalf("rendered results missing %q: %q", want, rendered) + } + } +} + +func TestGitHub_LimitReferencesKeepsBoundaryChunk(t *testing.T) { + chunks := []map[string]any{ + {"content": "first repository description"}, + {"content": "second repository description"}, + } + limited := limitGitHubReferences(chunks, 1) + if len(limited) != 1 { + t.Fatalf("limited chunks = %d, want 1", len(limited)) } } diff --git a/internal/agent/tool/google.go b/internal/agent/tool/google.go index a62c002216..be9827c8a5 100644 --- a/internal/agent/tool/google.go +++ b/internal/agent/tool/google.go @@ -18,21 +18,33 @@ package tool import ( "context" + "crypto/sha1" "encoding/json" "fmt" + "math/big" "net/http" "net/url" + "regexp" "strconv" "strings" + "unicode/utf8" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" + + "ragflow/internal/tokenizer" ) -const googleToolName = "google" +const googleToolName = "google_search" const googleToolDescription = "Search the web via Google using SerpApi. Returns organic_results[].{title, link, snippet}." +const googlePromptMaxTokens = 200000 + +var googleDataImagePattern = regexp.MustCompile(`!?\[[a-z]+\]\(data:image/png;base64,[ 0-9A-Za-z/_=+\-]+\)`) + +var googleNewlinePattern = regexp.MustCompile(`\n+`) + // googleParams is the JSON shape the model or canvas sends into InvokableRun. type googleParams struct { APIKey string `json:"api_key"` @@ -59,6 +71,9 @@ type GoogleTool struct { defaults googleParams } +var _ ToolComponent = (*GoogleTool)(nil) +var _ ReferenceBuilder = (*GoogleTool)(nil) + func NewGoogleTool() *GoogleTool { return NewGoogleToolWith(NewHTTPHelper()) } @@ -121,6 +136,24 @@ func (g *GoogleTool) InputForm() map[string]any { } } +func (g *GoogleTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "q": "Search query.", + "start": "Result offset.", + "num": "Maximum number of results.", + "api_key": "SerpApi API key.", + "country": "Google country code.", + "language": "Google language code.", + }, + Outputs: map[string]string{ + "formalized_content": "Rendered Google references for downstream prompts.", + "json": "Raw Google organic result list.", + }, + InputForm: g.InputForm(), + } +} + var googleEndpoint = "https://serpapi.com/search.json" func buildGoogleURL(p googleParams) string { @@ -207,6 +240,126 @@ func mergeGoogleDefaults(defaults, p googleParams) googleParams { return p } +func (g *GoogleTool) BuildReferences(_ context.Context, envelope map[string]any) ([]map[string]any, []map[string]any) { + return buildGoogleReferences(envelope) +} + +func (g *GoogleTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + results := envelopeSlice(envelope, "results") + chunks, _ := buildGoogleReferences(envelope) + return map[string]any{ + "formalized_content": renderGoogleReferences(chunks, googlePromptMaxTokens), + "json": results, + } +} + +func buildGoogleReferences(envelope map[string]any) ([]map[string]any, []map[string]any) { + results := envelopeSlice(envelope, "results") + chunks := make([]map[string]any, 0, len(results)) + docAggs := make([]map[string]any, 0, len(results)) + for _, result := range results { + item, ok := result.(map[string]any) + if !ok { + continue + } + content := strings.TrimSpace(googleText(item["snippet"])) + if content == "" { + content = strings.TrimSpace(googleAboutDescription(item["about_this_result"])) + } + content = googleDataImagePattern.ReplaceAllString(content, "") + content = truncateGoogleRunes(content, 10000) + if content == "" { + continue + } + documentID := strconv.FormatInt(googleHashInt(content, 100000000), 10) + displayID := strconv.FormatInt(googleHashInt(documentID, 500), 10) + title := strings.TrimSpace(googleText(item["title"])) + resultURL := strings.TrimSpace(googleText(item["link"])) + chunks = append(chunks, map[string]any{ + "id": displayID, + "chunk_id": documentID, + "content": content, + "doc_id": documentID, + "document_id": documentID, + "docnm_kwd": title, + "document_name": title, + "similarity": 1, + "score": 1, + "url": resultURL, + }) + docAggs = append(docAggs, map[string]any{ + "doc_name": title, + "doc_id": documentID, + "count": 1, + "url": resultURL, + }) + } + return chunks, docAggs +} + +func renderGoogleReferences(chunks []map[string]any, maxTokens int) string { + usedTokens := 0 + blocks := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + content := googleText(chunk["content"]) + if content == "" { + continue + } + block := strings.Join([]string{ + "\nID: " + googleText(chunk["id"]), + "├── Title: " + googlePromptField(chunk["document_name"]), + "├── URL: " + googlePromptField(chunk["url"]), + "└── Content:\n" + content, + }, "\n") + blockTokens := tokenizer.NumTokensFromString(block) + if maxTokens > 0 && float64(usedTokens+blockTokens) > float64(maxTokens)*0.97 { + break + } + usedTokens += blockTokens + blocks = append(blocks, block) + } + return strings.Join(blocks, "\n") +} + +func googleAboutDescription(value any) string { + about, ok := value.(map[string]any) + if !ok { + return "" + } + source, ok := about["source"].(map[string]any) + if !ok { + return "" + } + return googleText(source["description"]) +} + +func googleText(value any) string { + if value == nil { + return "" + } + if text, ok := value.(string); ok { + return text + } + return fmt.Sprint(value) +} + +func googlePromptField(value any) string { + return googleNewlinePattern.ReplaceAllString(googleText(value), " ") +} + +func googleHashInt(value string, modulus int64) int64 { + digest := sha1.Sum([]byte(value)) + number := new(big.Int).SetBytes(digest[:]) + return new(big.Int).Mod(number, big.NewInt(modulus)).Int64() +} + +func truncateGoogleRunes(value string, limit int) string { + if utf8.RuneCountInString(value) <= limit { + return value + } + return string([]rune(value)[:limit]) +} + func googleJSON(env googleEnvelope) string { b, err := json.Marshal(env) if err != nil { diff --git a/internal/agent/tool/google_scholar.go b/internal/agent/tool/google_scholar.go index fc3036709d..462e0e99fc 100644 --- a/internal/agent/tool/google_scholar.go +++ b/internal/agent/tool/google_scholar.go @@ -18,24 +18,36 @@ package tool import ( "context" + "crypto/sha1" "encoding/json" "fmt" + "math/big" "net/http" "net/url" + "regexp" "strconv" "strings" + "unicode/utf8" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" "golang.org/x/net/html" + + "ragflow/internal/tokenizer" ) -const googleScholarToolName = "google_scholar" +const googleScholarToolName = "google_scholar_search" const googleScholarToolDescription = "Google Scholar provides a simple way to broadly search for scholarly literature. From one place, you can search across many disciplines and sources: articles, theses, books, abstracts and court opinions, from academic publishers, professional societies, online repositories, universities and other web sites. Google Scholar helps you find relevant work across the world of scholarly research." const defaultGoogleScholarTopN = 12 +const googleScholarPromptMaxTokens = 200000 + +var googleScholarDataImagePattern = regexp.MustCompile(`!?\[[a-z]+\]\(data:image/png;base64,[ 0-9A-Za-z/_=+\-]+\)`) + +var googleScholarNewlinePattern = regexp.MustCompile(`\n+`) + // googleScholarParams is the JSON shape the model sends into InvokableRun. // All fields come from the canvas node form; only query is exposed to // the LLM via Info() (matching Python's meta — the LLM only sees query). @@ -78,6 +90,9 @@ type GoogleScholarTool struct { defaults googleScholarParams } +var _ ToolComponent = (*GoogleScholarTool)(nil) +var _ ReferenceBuilder = (*GoogleScholarTool)(nil) + // NewGoogleScholarTool returns a GoogleScholarTool using the default // HTTPHelper. func NewGoogleScholarTool() *GoogleScholarTool { @@ -120,6 +135,26 @@ func (g *GoogleScholarTool) Info(_ context.Context) (*schema.ToolInfo, error) { }, nil } +func (g *GoogleScholarTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "query": "Search query.", + "top_n": "Maximum number of results.", + "sort_by": "Sort order: relevance or date.", + "year_low": "Earliest publication year to include.", + "year_high": "Latest publication year to include.", + "patents": "Whether to include patents.", + }, + Outputs: map[string]string{ + "formalized_content": "Rendered Google Scholar references for downstream prompts.", + "json": "Google Scholar result list.", + }, + InputForm: map[string]any{ + "query": map[string]any{"name": "Query", "type": "line"}, + }, + } +} + // buildGoogleScholarURL composes the Scholar query URL. Centralized // for testability. sortBy: "relevance" (default) or "date". // yearLow / yearHigh: 0 means no filter. patents: nil or true includes @@ -158,8 +193,7 @@ func (g *GoogleScholarTool) InvokableRun(ctx context.Context, argsJSON string, _ } p = mergeGoogleScholarDefaults(g.defaults, p) if strings.TrimSpace(p.Query) == "" { - return googleScholarErrJSON(fmt.Errorf("query is required")), - fmt.Errorf("google_scholar: query is required") + return googleScholarJSON(googleScholarEnvelope{Results: []googleScholarResult{}}), nil } endpoint := buildGoogleScholarURL(p.Query, p.TopN, p.SortBy, p.YearLow, p.YearHigh, p.Patents) @@ -188,6 +222,103 @@ func (g *GoogleScholarTool) InvokableRun(ctx context.Context, argsJSON string, _ return googleScholarJSON(googleScholarEnvelope{Results: results}), nil } +func (g *GoogleScholarTool) BuildReferences(_ context.Context, envelope map[string]any) ([]map[string]any, []map[string]any) { + return buildGoogleScholarReferences(envelope) +} + +func (g *GoogleScholarTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + results := envelopeSlice(envelope, "results") + chunks, _ := buildGoogleScholarReferences(envelope) + return map[string]any{ + "formalized_content": renderGoogleScholarReferences(chunks, googleScholarPromptMaxTokens), + "json": results, + } +} + +func buildGoogleScholarReferences(envelope map[string]any) ([]map[string]any, []map[string]any) { + results := envelopeSlice(envelope, "results") + chunks := make([]map[string]any, 0, len(results)) + docAggs := make([]map[string]any, 0, len(results)) + for _, result := range results { + paper, ok := result.(map[string]any) + if !ok { + continue + } + content := strings.Join([]string{ + "Authors: " + googleScholarText(paper["authors"]), + "Year: " + googleScholarText(paper["year"]), + "Snippet: " + googleScholarText(paper["snippet"]), + }, "\n") + content = googleScholarDataImagePattern.ReplaceAllString(content, "") + content = truncateGoogleScholarRunes(content, 10000) + if content == "" { + continue + } + documentID := strconv.FormatInt(googleScholarHashInt(content, 100000000), 10) + displayID := strconv.FormatInt(googleScholarHashInt(documentID, 500), 10) + title := googleScholarText(paper["title"]) + resultURL := googleScholarText(paper["link"]) + chunks = append(chunks, map[string]any{ + "id": displayID, + "chunk_id": documentID, + "content": content, + "doc_id": documentID, + "document_id": documentID, + "docnm_kwd": title, + "document_name": title, + "similarity": 1, + "score": 1, + "url": resultURL, + }) + docAggs = append(docAggs, map[string]any{"doc_name": title, "doc_id": documentID, "count": 1, "url": resultURL}) + } + return chunks, docAggs +} + +func renderGoogleScholarReferences(chunks []map[string]any, maxTokens int) string { + usedTokens := 0 + blocks := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + content := googleScholarText(chunk["content"]) + block := strings.Join([]string{ + "\nID: " + googleScholarText(chunk["id"]), + "├── Title: " + googleScholarNewlinePattern.ReplaceAllString(googleScholarText(chunk["document_name"]), " "), + "├── URL: " + googleScholarNewlinePattern.ReplaceAllString(googleScholarText(chunk["url"]), " "), + "└── Content:\n" + content, + }, "\n") + blockTokens := tokenizer.NumTokensFromString(block) + if maxTokens > 0 && float64(usedTokens+blockTokens) > float64(maxTokens)*0.97 { + break + } + usedTokens += blockTokens + blocks = append(blocks, block) + } + return strings.Join(blocks, "\n") +} + +func googleScholarText(value any) string { + if value == nil { + return "" + } + if text, ok := value.(string); ok { + return text + } + return fmt.Sprint(value) +} + +func googleScholarHashInt(value string, modulus int64) int64 { + digest := sha1.Sum([]byte(value)) + number := new(big.Int).SetBytes(digest[:]) + return new(big.Int).Mod(number, big.NewInt(modulus)).Int64() +} + +func truncateGoogleScholarRunes(value string, limit int) string { + if utf8.RuneCountInString(value) <= limit { + return value + } + return string([]rune(value)[:limit]) +} + func mergeGoogleScholarDefaults(defaults, p googleScholarParams) googleScholarParams { if p.Query == "" { p.Query = defaults.Query diff --git a/internal/agent/tool/google_scholar_test.go b/internal/agent/tool/google_scholar_test.go index 91342d13a5..a13215c531 100644 --- a/internal/agent/tool/google_scholar_test.go +++ b/internal/agent/tool/google_scholar_test.go @@ -24,6 +24,8 @@ import ( "net/url" "strings" "testing" + + "ragflow/internal/tokenizer" ) const cannedScholarHTML = ` @@ -225,16 +227,17 @@ func TestGoogleScholar_ParseResults(t *testing.T) { } } -func TestGoogleScholar_RequiresQuery(t *testing.T) { +func TestGoogleScholar_EmptyQuery(t *testing.T) { t.Parallel() tool := NewGoogleScholarTool() - _, err := tool.InvokableRun(context.Background(), `{"query":""}`) - if err == nil { - t.Fatal("expected error for empty query") + out, err := tool.InvokableRun(context.Background(), `{"query":""}`) + if err != nil { + t.Fatalf("InvokableRun(empty): %v", err) } - if !strings.Contains(err.Error(), "query") { - t.Errorf("err = %v, want to mention query", err) + var envelope googleScholarEnvelope + if err := json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 { + t.Fatalf("empty result = %s / %v", out, err) } } @@ -246,14 +249,80 @@ func TestGoogleScholar_Info(t *testing.T) { if err != nil { t.Fatalf("Info: %v", err) } - if info.Name != "google_scholar" { - t.Errorf("Name = %q, want google_scholar", info.Name) + if info.Name != "google_scholar_search" { + t.Errorf("Name = %q, want google_scholar_search", info.Name) } if !strings.Contains(info.Desc, "Scholar") { t.Errorf("Desc = %q, want to mention Scholar", info.Desc) } } +func TestGoogleScholar_ComponentReferencesAndValidation(t *testing.T) { + t.Parallel() + + built, err := BuildByName("google_scholar", map[string]any{ + "top_n": float64(7), + "sort_by": "date", + "year_low": float64(2020), + "patents": false, + "outputs": map[string]any{"json": map[string]any{}}, + }) + if err != nil { + t.Fatalf("BuildByName: %v", err) + } + scholar := built.(*GoogleScholarTool) + if scholar.defaults.TopN != 7 || scholar.defaults.SortBy != "date" || scholar.defaults.YearLow != 2020 || scholar.defaults.Patents == nil || *scholar.defaults.Patents { + t.Fatalf("defaults = %+v", scholar.defaults) + } + for _, params := range []map[string]any{{"top_n": 0}, {"top_n": 1.5}, {"sort_by": "newest"}, {"patents": "yes"}} { + if _, err := BuildByName("google_scholar", params); err == nil { + t.Fatalf("BuildByName(%#v) succeeded", params) + } + } + spec := scholar.ComponentSpec() + if query, ok := spec.InputForm["query"].(map[string]any); !ok || query["type"] != "line" { + t.Fatalf("query input form = %#v", spec.InputForm["query"]) + } + envelope := map[string]any{"results": []any{map[string]any{ + "title": "Paper", "link": "https://paper.example", "authors": "A Author", "year": "2024", "snippet": "Abstract", + }}} + chunks, docAggs := scholar.BuildReferences(context.Background(), envelope) + if len(chunks) != 1 || len(docAggs) != 1 || !strings.Contains(chunks[0]["content"].(string), "Authors: A Author") { + t.Fatalf("references = %#v / %#v", chunks, docAggs) + } + outputs := scholar.BuildComponentOutputs(envelope) + if results, ok := outputs["json"].([]any); !ok || len(results) != 1 { + t.Fatalf("json output = %#v", outputs["json"]) + } + if !strings.Contains(outputs["formalized_content"].(string), "Snippet: Abstract") { + t.Fatalf("formalized_content = %q", outputs["formalized_content"]) + } + if _, exists := envelope["chunks"]; exists { + t.Fatalf("output conversion mutated envelope: %#v", envelope) + } +} + +func TestRenderGoogleScholarReferencesStopsBeforeOverBudgetBlock(t *testing.T) { + t.Parallel() + + chunks := []map[string]any{ + {"id": "1", "document_name": "First", "url": "https://first.example", "content": "first reference content"}, + {"id": "2", "document_name": "Second", "url": "https://second.example", "content": "second reference content"}, + } + firstBlock := renderGoogleScholarReferences(chunks[:1], 0) + firstTokens := tokenizer.NumTokensFromString(firstBlock) + maxTokens := (firstTokens*100 + 96) / 97 + if got := renderGoogleScholarReferences(chunks, maxTokens); got != firstBlock { + t.Fatalf("rendered = %q, want only first block %q", got, firstBlock) + } + if got := renderGoogleScholarReferences(chunks, 1); got != "" { + t.Fatalf("over-budget first block was appended: %q", got) + } + if got := renderGoogleScholarReferences(chunks, 0); !strings.Contains(got, "Title: First") || !strings.Contains(got, "Title: Second") { + t.Fatalf("unlimited rendering dropped blocks: %q", got) + } +} + func TestGoogleScholar_MergesNodeLevelDefaults(t *testing.T) { t.Parallel() diff --git a/internal/agent/tool/google_test.go b/internal/agent/tool/google_test.go index 895e3d1a8c..242a8323c0 100644 --- a/internal/agent/tool/google_test.go +++ b/internal/agent/tool/google_test.go @@ -24,6 +24,8 @@ import ( "net/url" "strings" "testing" + + "ragflow/internal/tokenizer" ) func TestGoogle_BuildURL(t *testing.T) { @@ -145,8 +147,8 @@ func TestGoogle_InfoAndInputForm(t *testing.T) { if err != nil { t.Fatalf("Info: %v", err) } - if info.Name != "google" { - t.Errorf("Name = %q, want google", info.Name) + if info.Name != "google_search" { + t.Errorf("Name = %q, want google_search", info.Name) } if !strings.Contains(info.Desc, "Google") { t.Errorf("Desc = %q, want to mention Google", info.Desc) @@ -161,6 +163,11 @@ func TestGoogle_InfoAndInputForm(t *testing.T) { if _, ok := form["num"]; !ok { t.Fatalf("InputForm missing num: %+v", form) } + for _, configField := range []string{"api_key", "country", "language"} { + if _, exists := form[configField]; exists { + t.Fatalf("InputForm leaked node configuration %q: %+v", configField, form) + } + } } func TestGoogle_MergeDefaultsPrefersExplicitInputs(t *testing.T) { @@ -220,17 +227,89 @@ func TestGoogle_BuildByNameAcceptsNodeParams(t *testing.T) { } } -func TestGoogle_BuildByNameRejectsRemovedAliasNodeParams(t *testing.T) { +func TestGoogle_BuildByNameIgnoresUnrelatedCanvasParams(t *testing.T) { t.Parallel() - _, err := BuildByName("google", map[string]any{ + built, err := BuildByName("google", map[string]any{ "query": "ragflow", "max_results": 5, + "outputs": map[string]any{"json": map[string]any{}}, }) - if err == nil { - t.Fatal("expected error for removed google alias node params") + if err != nil { + t.Fatalf("BuildByName rejected unrelated Canvas params: %v", err) } - if !strings.Contains(err.Error(), "does not accept node-level param") { - t.Fatalf("err = %q, want unsupported node-level param error", err.Error()) + google := built.(*GoogleTool) + if google.defaults.Q != "" || google.defaults.Num != 0 { + t.Fatalf("unrelated params changed defaults: %+v", google.defaults) + } +} + +func TestGoogle_ComponentReferencesAndOutputs(t *testing.T) { + t.Parallel() + + google := NewGoogleTool() + spec := google.ComponentSpec() + if query, ok := spec.InputForm["q"].(map[string]any); !ok || query["name"] != "Query" || query["type"] != "line" { + t.Fatalf("query input form = %#v", spec.InputForm["q"]) + } + envelope := map[string]any{"results": []any{ + map[string]any{ + "title": "RAGFlow", + "link": "https://ragflow.io", + "snippet": "Go-first snippet", + "about_this_result": map[string]any{"source": map[string]any{ + "description": "fallback description", + }}, + "custom": "preserved", + }, + map[string]any{ + "title": "Fallback", + "link": "https://example.com", + "about_this_result": map[string]any{"source": map[string]any{ + "description": "about content", + }}, + }, + }} + chunks, docAggs := google.BuildReferences(context.Background(), envelope) + if len(chunks) != 2 || len(docAggs) != 2 { + t.Fatalf("references = %#v / %#v", chunks, docAggs) + } + if chunks[0]["content"] != "Go-first snippet" || chunks[1]["content"] != "about content" || chunks[0]["similarity"] != 1 { + t.Fatalf("reference content = %#v", chunks) + } + outputs := google.BuildComponentOutputs(envelope) + results, ok := outputs["json"].([]any) + if !ok || len(results) != 2 || results[0].(map[string]any)["custom"] != "preserved" { + t.Fatalf("json output = %#v", outputs["json"]) + } + rendered, _ := outputs["formalized_content"].(string) + for _, want := range []string{"Title: RAGFlow", "Go-first snippet", "about content"} { + if !strings.Contains(rendered, want) { + t.Fatalf("formalized_content missing %q: %q", want, rendered) + } + } + if _, exists := envelope["chunks"]; exists { + t.Fatalf("output conversion mutated envelope: %#v", envelope) + } +} + +func TestRenderGoogleReferencesStopsBeforeOverBudgetBlock(t *testing.T) { + t.Parallel() + + chunks := []map[string]any{ + {"id": "1", "document_name": "First", "url": "https://first.example", "content": "first reference content"}, + {"id": "2", "document_name": "Second", "url": "https://second.example", "content": "second reference content"}, + } + firstBlock := renderGoogleReferences(chunks[:1], 0) + firstTokens := tokenizer.NumTokensFromString(firstBlock) + maxTokens := (firstTokens*100 + 96) / 97 + if got := renderGoogleReferences(chunks, maxTokens); got != firstBlock { + t.Fatalf("rendered = %q, want only first block %q", got, firstBlock) + } + if got := renderGoogleReferences(chunks, 1); got != "" { + t.Fatalf("over-budget first block was appended: %q", got) + } + if got := renderGoogleReferences(chunks, 0); !strings.Contains(got, "Title: First") || !strings.Contains(got, "Title: Second") { + t.Fatalf("unlimited rendering dropped blocks: %q", got) } } diff --git a/internal/agent/tool/keenable.go b/internal/agent/tool/keenable.go index aeecfe8f2b..43c8eb0731 100644 --- a/internal/agent/tool/keenable.go +++ b/internal/agent/tool/keenable.go @@ -18,18 +18,23 @@ package tool import ( "context" + "crypto/sha1" "encoding/json" "fmt" + "math/big" "net/http" neturl "net/url" - "ragflow/internal/common" + "strconv" "strings" + "unicode/utf8" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" + + "ragflow/internal/common" ) -const keenableToolName = "keenable" +const keenableToolName = "keenable_search" // keenableToolDescription follows the upstream Python tool's description, // trimmed for the chat model. The "no API key required" line is the @@ -58,27 +63,16 @@ type keenableRequestBody struct { Site string `json:"site,omitempty"` } -// keenableResult mirrors one element of the upstream `results` array. -// The Python tool's _retrieve_chunks reads `title`, `url`, `description`, -// so we model those fields and pass everything else through verbatim -// when serializing to the model. -type keenableResult struct { - Title string `json:"title"` - URL string `json:"url"` - Description string `json:"description"` -} - -// keenableResponse is the envelope returned by Keenable. We only model -// the fields we care about; the upstream API has more, but they are -// ignored. +// keenableResponse preserves complete upstream result objects because Python +// exposes them unchanged through the Canvas json output. type keenableResponse struct { - Results []keenableResult `json:"results"` + Results []map[string]any `json:"results"` } // keenableEnvelope is the shape the model actually sees, identical to // the Python tool's output convention. type keenableEnvelope struct { - Results []keenableResult `json:"results"` + Results []map[string]any `json:"results"` Error string `json:"_ERROR,omitempty"` } @@ -87,8 +81,9 @@ type keenableEnvelope struct { // endpoint (with X-API-Key) when an API key is provided. The upstream // `results` array is returned as JSON. type KeenableTool struct { - helper *HTTPHelper - apiKey string + helper *HTTPHelper + apiKey string + defaults keenableParams // envBaseURL resolves the Keenable API base URL from the // KEENABLE_API_URL env var (HTTPS enforced). Exposed as a @@ -97,40 +92,53 @@ type KeenableTool struct { envBaseURL func() string } +var _ ToolComponent = (*KeenableTool)(nil) +var _ ReferenceBuilder = (*KeenableTool)(nil) + // NewKeenableTool returns a KeenableTool using the default HTTPHelper // and the KEENABLE_API_URL env var for base-URL resolution. func NewKeenableTool() *KeenableTool { - return NewKeenableToolWith(NewHTTPHelper()) + return newKeenableTool(nil, nil, "", keenableParams{}) } // NewKeenableToolWithAPIKey returns a KeenableTool that uses a // server-provided API key instead of model-visible runtime args. func NewKeenableToolWithAPIKey(h *HTTPHelper, apiKey string) *KeenableTool { - t := NewKeenableToolWith(h) - t.apiKey = strings.TrimSpace(apiKey) - return t + return newKeenableTool(h, nil, apiKey, keenableParams{}) } // NewKeenableToolWith returns a KeenableTool that uses the provided // HTTPHelper. Useful for tests that want to inject a custom transport. func NewKeenableToolWith(h *HTTPHelper) *KeenableTool { - if h == nil { - h = NewHTTPHelper() - } - return &KeenableTool{helper: h, envBaseURL: defaultKeenableEnvBaseURL} + return newKeenableTool(h, nil, "", keenableParams{}) } // NewKeenableToolWithEnvBaseURL returns a KeenableTool with a custom // base-URL resolver. Useful for tests that want to inject a fake env // without mutating process state. func NewKeenableToolWithEnvBaseURL(h *HTTPHelper, envBaseURL func() string) *KeenableTool { + return newKeenableTool(h, envBaseURL, "", keenableParams{}) +} + +func newKeenableTool(h *HTTPHelper, envBaseURL func() string, apiKey string, defaults keenableParams) *KeenableTool { if h == nil { h = NewHTTPHelper() } if envBaseURL == nil { envBaseURL = defaultKeenableEnvBaseURL } - return &KeenableTool{helper: h, envBaseURL: envBaseURL} + if strings.TrimSpace(defaults.Mode) == "" { + defaults.Mode = "pro" + } + if defaults.TopN == 0 { + defaults.TopN = 10 + } + return &KeenableTool{ + helper: h, + apiKey: strings.TrimSpace(apiKey), + defaults: defaults, + envBaseURL: envBaseURL, + } } // defaultKeenableEnvBaseURL is the production base-URL resolver. @@ -195,43 +203,20 @@ func (k *KeenableTool) Info(_ context.Context) (*schema.ToolInfo, error) { Desc: "Optional. Restrict results to a single domain, e.g. 'techcrunch.com'. Defaults to '' (no filter).", Required: false, }, - "mode": { - Type: schema.String, - Desc: `Search mode: "pro" (default, deeper) or "realtime" (low latency; requires a server-configured API key).`, - Required: false, - }, - "top_n": { - Type: schema.Integer, - Desc: "Maximum number of results to return. Defaults to 10.", - Required: false, - }, }), }, nil } -func (k *KeenableTool) GetInputForm() map[string]any { - return map[string]any{ - "query": map[string]any{ - "name": "Query", - "type": "line", - }, - "site": map[string]any{ - "name": "Site", - "type": "line", - }, - } -} - // InvokableRun performs the Keenable search. func (k *KeenableTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) { - var p keenableParams - if err := json.Unmarshal([]byte(argsJSON), &p); err != nil { + var runtimeParams keenableParams + if err := json.Unmarshal([]byte(argsJSON), &runtimeParams); err != nil { return keenableErrJSON(fmt.Errorf("keenable: parse arguments: %w", err)), fmt.Errorf("keenable: parse arguments: %w", err) } + p := mergeKeenableParams(k.defaults, runtimeParams) if strings.TrimSpace(p.Query) == "" { - return keenableErrJSON(fmt.Errorf("query is required")), - fmt.Errorf("keenable: query is required") + return keenableJSON(keenableEnvelope{Results: []map[string]any{}}), nil } mode := strings.TrimSpace(p.Mode) @@ -310,6 +295,131 @@ func (k *KeenableTool) InvokableRun(ctx context.Context, argsJSON string, _ ...t return keenableJSON(keenableEnvelope{Results: results}), nil } +func mergeKeenableParams(defaults, params keenableParams) keenableParams { + if strings.TrimSpace(params.Site) == "" { + params.Site = defaults.Site + } + if strings.TrimSpace(params.Mode) == "" { + params.Mode = defaults.Mode + } + if params.TopN == 0 { + params.TopN = defaults.TopN + } + return params +} + +// ComponentSpec returns the Python-compatible KeenableSearch Canvas surface. +func (k *KeenableTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "query": "The search keywords to execute with Keenable.", + "site": "Optional single-domain filter.", + }, + Outputs: map[string]string{ + "formalized_content": "Rendered search results for downstream LLM prompts.", + "json": "Raw Keenable result list.", + }, + InputForm: map[string]any{ + "query": map[string]any{ + "name": "Query", + "type": "line", + }, + "site": map[string]any{ + "name": "Site", + "type": "line", + }, + }, + } +} + +// BuildReferences builds the same references as Python ToolBase._retrieve_chunks. +func (k *KeenableTool) BuildReferences(_ context.Context, envelope map[string]any) ([]map[string]any, []map[string]any) { + return buildKeenableReferences(envelope) +} + +// BuildComponentOutputs converts Keenable's complete tool envelope into its +// public Canvas outputs. +func (k *KeenableTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + results := envelopeSlice(envelope, "results") + chunks, _ := buildKeenableReferences(envelope) + return map[string]any{ + "formalized_content": renderKeenableReferences(chunks), + "json": results, + } +} + +func buildKeenableReferences(envelope map[string]any) ([]map[string]any, []map[string]any) { + results := envelopeSlice(envelope, "results") + chunks := make([]map[string]any, 0, len(results)) + docAggs := make([]map[string]any, 0, len(results)) + for _, item := range results { + result, ok := item.(map[string]any) + if !ok { + continue + } + content := truncateKeenableRunes(strings.TrimSpace(keenableValueString(result["description"])), 10000) + if content == "" || content == "None" { + continue + } + documentID := strconv.FormatInt(keenableHashInt(content, 100000000), 10) + title := keenableValueString(result["title"]) + resultURL := keenableValueString(result["url"]) + displayID := strconv.FormatInt(keenableHashInt(documentID, 500), 10) + chunks = append(chunks, map[string]any{ + "id": displayID, + "chunk_id": documentID, + "content": content, + "doc_id": documentID, + "document_id": documentID, + "docnm_kwd": title, + "document_name": title, + "similarity": 1, + "score": 1, + "url": resultURL, + }) + docAggs = append(docAggs, map[string]any{ + "doc_name": title, + "doc_id": documentID, + "count": 1, + "url": resultURL, + }) + } + return chunks, docAggs +} + +func renderKeenableReferences(chunks []map[string]any) string { + blocks := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + blocks = append(blocks, strings.Join([]string{ + "\nID: " + keenableValueString(chunk["id"]), + "├── Title: " + keenableValueString(chunk["docnm_kwd"]), + "├── URL: " + keenableValueString(chunk["url"]), + "└── Content:\n" + keenableValueString(chunk["content"]), + }, "\n")) + } + return strings.Join(blocks, "\n") +} + +func keenableValueString(value any) string { + if value == nil { + return "None" + } + return strings.TrimSpace(fmt.Sprint(value)) +} + +func keenableHashInt(value string, modulus int64) int64 { + sum := sha1.Sum([]byte(value)) + number := new(big.Int).SetBytes(sum[:]) + return new(big.Int).Mod(number, big.NewInt(modulus)).Int64() +} + +func truncateKeenableRunes(value string, limit int) string { + if utf8.RuneCountInString(value) <= limit { + return value + } + return string([]rune(value)[:limit]) +} + // keenableJSON marshals the envelope to a JSON string for the model. func keenableJSON(env keenableEnvelope) string { b, err := json.Marshal(env) diff --git a/internal/agent/tool/keenable_test.go b/internal/agent/tool/keenable_test.go index 82a12213d3..ca8d21e0c2 100644 --- a/internal/agent/tool/keenable_test.go +++ b/internal/agent/tool/keenable_test.go @@ -126,7 +126,7 @@ func TestKeenable_SiteAndTopN(t *testing.T) { _ = json.NewDecoder(r.Body).Decode(&gotBody) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"results":[ - {"title":"A","url":"https://a","description":"alpha"}, + {"title":"A","url":"https://a","description":"alpha","custom":"preserved"}, {"title":"B","url":"https://b","description":"beta"}, {"title":"C","url":"https://c","description":"gamma"}, {"title":"D","url":"https://d","description":"delta"} @@ -159,9 +159,12 @@ func TestKeenable_SiteAndTopN(t *testing.T) { if len(env.Results) != 2 { t.Fatalf("Results len = %d, want 2 (capped by top_n)", len(env.Results)) } - if env.Results[0].Title != "A" || env.Results[1].Title != "B" { + if env.Results[0]["title"] != "A" || env.Results[1]["title"] != "B" { t.Errorf("Results = %+v, want first 2 upstream items", env.Results) } + if env.Results[0]["custom"] != "preserved" { + t.Fatalf("raw upstream fields were lost: %#v", env.Results[0]) + } } // TestKeenable_DefaultTopN verifies that omitting top_n keeps up to 10 @@ -204,18 +207,20 @@ func TestKeenable_DefaultTopN(t *testing.T) { } } -// TestKeenable_MissingQuery verifies that an empty query is rejected -// before any HTTP request is made. func TestKeenable_MissingQuery(t *testing.T) { t.Parallel() tool := NewKeenableTool() - _, err := tool.InvokableRun(context.Background(), `{}`) - if err == nil { - t.Fatal("expected error for missing query") + out, err := tool.InvokableRun(context.Background(), `{}`) + if err != nil { + t.Fatalf("InvokableRun: %v", err) } - if !strings.Contains(err.Error(), "query") { - t.Errorf("err = %v, want to mention query", err) + var envelope keenableEnvelope + if err := json.Unmarshal([]byte(out), &envelope); err != nil { + t.Fatalf("output is not valid JSON: %v", err) + } + if len(envelope.Results) != 0 || envelope.Error != "" { + t.Fatalf("envelope = %+v, want empty results without error", envelope) } } @@ -373,8 +378,8 @@ func TestKeenable_Info(t *testing.T) { if err != nil { t.Fatalf("Info: %v", err) } - if info.Name != "keenable" { - t.Errorf("Name = %q, want keenable", info.Name) + if info.Name != "keenable_search" { + t.Errorf("Name = %q, want keenable_search", info.Name) } if !strings.Contains(info.Desc, "Keenable") { t.Errorf("Desc = %q, want to mention Keenable", info.Desc) @@ -389,7 +394,10 @@ func TestKeenable_Info(t *testing.T) { if strings.Contains(string(paramsJSON), "api_key") { t.Fatalf("Info ParamsOneOf unexpectedly exposes api_key: %s", string(paramsJSON)) } - form := tool.GetInputForm() + if strings.Contains(string(paramsJSON), "mode") || strings.Contains(string(paramsJSON), "top_n") { + t.Fatalf("Info ParamsOneOf leaked node configuration: %s", string(paramsJSON)) + } + form := tool.ComponentSpec().InputForm for _, key := range []string{"query", "site"} { field, ok := form[key].(map[string]any) if !ok { @@ -400,3 +408,84 @@ func TestKeenable_Info(t *testing.T) { } } } + +func TestKeenable_ComponentContractReferencesAndOutputs(t *testing.T) { + t.Parallel() + + tool := NewKeenableTool() + spec := tool.ComponentSpec() + for _, input := range []string{"query", "site"} { + if _, ok := spec.Inputs[input]; !ok { + t.Fatalf("component inputs missing %s: %#v", input, spec.Inputs) + } + } + for _, output := range []string{"formalized_content", "json"} { + if _, ok := spec.Outputs[output]; !ok { + t.Fatalf("component outputs missing %s: %#v", output, spec.Outputs) + } + } + + envelope := map[string]any{"results": []any{map[string]any{ + "title": "Keenable result", + "url": "https://example.com/item", + "description": "Fresh search result", + }}} + chunks, docAggs := tool.BuildReferences(context.Background(), envelope) + if len(chunks) != 1 || len(docAggs) != 1 { + t.Fatalf("references = %#v / %#v", chunks, docAggs) + } + if chunks[0]["document_name"] != "Keenable result" || chunks[0]["url"] != "https://example.com/item" { + t.Fatalf("reference metadata = %#v", chunks[0]) + } + outputs := tool.BuildComponentOutputs(envelope) + formalized, _ := outputs["formalized_content"].(string) + for _, want := range []string{"Keenable result", "https://example.com/item", "Fresh search result"} { + if !strings.Contains(formalized, want) { + t.Fatalf("formalized_content missing %q: %s", want, formalized) + } + } + results, ok := outputs["json"].([]any) + if !ok || len(results) != 1 { + t.Fatalf("json output = %#v", outputs["json"]) + } + if _, exists := envelope["chunks"]; exists { + t.Fatalf("component conversion mutated envelope: %#v", envelope) + } +} + +func TestKeenable_BuildByNameAcceptsCanvasParams(t *testing.T) { + t.Parallel() + + built, err := BuildByName("keenable", map[string]any{ + "api_key": "stored-key", + "mode": "realtime", + "top_n": float64(3), + "site": "example.com", + "outputs": map[string]any{"json": map[string]any{}}, + }) + if err != nil { + t.Fatalf("BuildByName: %v", err) + } + tool := built.(*KeenableTool) + if tool.apiKey != "stored-key" || tool.defaults.Mode != "realtime" || tool.defaults.TopN != 3 || tool.defaults.Site != "example.com" { + t.Fatalf("tool config = apiKey=%q defaults=%+v", tool.apiKey, tool.defaults) + } +} + +func TestKeenable_BuildByNameRejectsInvalidCanvasParams(t *testing.T) { + t.Parallel() + + invalid := []map[string]any{ + {"api_key": 1}, + {"mode": "fast"}, + {"mode": "realtime"}, + {"top_n": 0}, + {"top_n": 1.5}, + {"site": 1}, + } + for _, params := range invalid { + if _, err := BuildByName("keenable", params); err == nil { + t.Fatalf("BuildByName(%#v) succeeded, want validation error", params) + } + } +} diff --git a/internal/agent/tool/pubmed.go b/internal/agent/tool/pubmed.go index 5ae3cbcb97..612552e9d1 100644 --- a/internal/agent/tool/pubmed.go +++ b/internal/agent/tool/pubmed.go @@ -18,25 +18,36 @@ package tool import ( "context" + "crypto/sha1" "encoding/json" "encoding/xml" "fmt" + "math/big" "net/http" "net/url" + "regexp" "strconv" "strings" + "unicode/utf8" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" + + "ragflow/internal/tokenizer" ) const ( - pubmedToolName = "pubmed" + pubmedToolName = "pubmed_search" pubmedToolDescription = "Search PubMed for life sciences and biomedical references." defaultPubMedTopN = 12 defaultPubMedEmail = "A.N.Other@example.com" + pubmedPromptMaxTokens = 200000 ) +var pubmedDataImagePattern = regexp.MustCompile(`!?\[[a-z]+\]\(data:image/png;base64,[ 0-9A-Za-z/_=+\-]+\)`) + +var pubmedNewlinePattern = regexp.MustCompile(`\n+`) + // pubmedParams mirrors Python PubMedParam. Info() still only exposes query to // the LLM; top_n and email are canvas-side params merged with constructor // defaults at runtime. @@ -72,6 +83,9 @@ type PubMedTool struct { defaults pubmedParams } +var _ ToolComponent = (*PubMedTool)(nil) +var _ ReferenceBuilder = (*PubMedTool)(nil) + func NewPubMedTool() *PubMedTool { return NewPubMedToolWith(NewHTTPHelper()) } @@ -109,6 +123,19 @@ func (p *PubMedTool) Info(_ context.Context) (*schema.ToolInfo, error) { }, nil } +func (p *PubMedTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{"query": "PubMed search query."}, + Outputs: map[string]string{ + "formalized_content": "Rendered PubMed references for downstream prompts.", + "json": "PubMed result list.", + }, + InputForm: map[string]any{ + "query": map[string]any{"name": "Query", "type": "line"}, + }, + } +} + func buildPubMedESearchURL(query string, topN int, email string) string { q := url.Values{} q.Set("db", "pubmed") @@ -179,8 +206,7 @@ func (p *PubMedTool) InvokableRun(ctx context.Context, argsJSON string, _ ...too params = mergePubMedDefaults(p.defaults, params) params.Query = strings.TrimSpace(params.Query) if params.Query == "" { - err := fmt.Errorf("pubmed: query is required") - return pubmedErrJSON(err), err + return pubmedJSON(pubmedEnvelope{Results: []pubmedResult{}}), nil } headers := map[string]string{ @@ -229,6 +255,96 @@ func (p *PubMedTool) InvokableRun(ctx context.Context, argsJSON string, _ ...too return pubmedJSON(pubmedEnvelope{Results: results}), nil } +func (p *PubMedTool) BuildReferences(_ context.Context, envelope map[string]any) ([]map[string]any, []map[string]any) { + return buildPubMedReferences(envelope) +} + +func (p *PubMedTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + results := envelopeSlice(envelope, "results") + chunks, _ := buildPubMedReferences(envelope) + return map[string]any{ + "formalized_content": renderPubMedReferences(chunks, pubmedPromptMaxTokens), + "json": results, + } +} + +func buildPubMedReferences(envelope map[string]any) ([]map[string]any, []map[string]any) { + results := envelopeSlice(envelope, "results") + chunks := make([]map[string]any, 0, len(results)) + docAggs := make([]map[string]any, 0, len(results)) + for _, result := range results { + item, ok := result.(map[string]any) + if !ok { + continue + } + content := pubmedDataImagePattern.ReplaceAllString(pubmedText(item["content"]), "") + content = truncatePubMedRunes(content, 10000) + if content == "" { + continue + } + documentID := strconv.FormatInt(pubmedHashInt(content, 100000000), 10) + displayID := strconv.FormatInt(pubmedHashInt(documentID, 500), 10) + title := pubmedText(item["title"]) + resultURL := pubmedText(item["url"]) + chunks = append(chunks, map[string]any{ + "id": displayID, + "chunk_id": documentID, + "content": content, + "doc_id": documentID, + "document_id": documentID, + "docnm_kwd": title, + "document_name": title, + "similarity": 1, + "score": 1, + "url": resultURL, + }) + docAggs = append(docAggs, map[string]any{"doc_name": title, "doc_id": documentID, "count": 1, "url": resultURL}) + } + return chunks, docAggs +} + +func renderPubMedReferences(chunks []map[string]any, maxTokens int) string { + usedTokens := 0 + blocks := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + content := pubmedText(chunk["content"]) + usedTokens += tokenizer.NumTokensFromString(content) + blocks = append(blocks, strings.Join([]string{ + "\nID: " + pubmedText(chunk["id"]), + "├── Title: " + pubmedNewlinePattern.ReplaceAllString(pubmedText(chunk["document_name"]), " "), + "├── URL: " + pubmedNewlinePattern.ReplaceAllString(pubmedText(chunk["url"]), " "), + "└── Content:\n" + content, + }, "\n")) + if maxTokens > 0 && float64(maxTokens)*0.97 < float64(usedTokens) { + break + } + } + return strings.Join(blocks, "\n") +} + +func pubmedText(value any) string { + if value == nil { + return "" + } + if text, ok := value.(string); ok { + return text + } + return fmt.Sprint(value) +} + +func pubmedHashInt(value string, modulus int64) int64 { + digest := sha1.Sum([]byte(value)) + number := new(big.Int).SetBytes(digest[:]) + return new(big.Int).Mod(number, big.NewInt(modulus)).Int64() +} + +func truncatePubMedRunes(value string, limit int) string { + if utf8.RuneCountInString(value) <= limit { + return value + } + return string([]rune(value)[:limit]) +} + func formatPubMedResult(article pubmedXMLArticle) pubmedResult { title := fallbackPubMedField(article.Article.Title, "No title") abstract := fallbackPubMedField(strings.Join(article.Article.Abstract.Text, " "), "No abstract available") diff --git a/internal/agent/tool/pubmed_test.go b/internal/agent/tool/pubmed_test.go index efb8102b3d..071116c017 100644 --- a/internal/agent/tool/pubmed_test.go +++ b/internal/agent/tool/pubmed_test.go @@ -191,16 +191,17 @@ func TestPubMed_InvokableRunEmptyResults(t *testing.T) { } } -func TestPubMed_InvokableRunRequiresQuery(t *testing.T) { +func TestPubMed_InvokableRunEmptyQuery(t *testing.T) { t.Parallel() tool := NewPubMedTool() - _, err := tool.InvokableRun(context.Background(), `{"query":""}`) - if err == nil { - t.Fatal("expected error for empty query") + out, err := tool.InvokableRun(context.Background(), `{"query":""}`) + if err != nil { + t.Fatalf("InvokableRun(empty): %v", err) } - if !strings.Contains(err.Error(), "query") { - t.Fatalf("err = %q, want query validation", err.Error()) + var envelope pubmedEnvelope + if err := json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 { + t.Fatalf("empty result = %s / %v", out, err) } } @@ -212,8 +213,8 @@ func TestPubMed_InfoOnlyExposesQuery(t *testing.T) { if err != nil { t.Fatalf("Info: %v", err) } - if info.Name != "pubmed" { - t.Fatalf("Name = %q, want pubmed", info.Name) + if info.Name != "pubmed_search" { + t.Fatalf("Name = %q, want pubmed_search", info.Name) } schema, err := info.ParamsOneOf.ToJSONSchema() if err != nil { @@ -238,7 +239,9 @@ func TestPubMed_InfoOnlyExposesQuery(t *testing.T) { func TestPubMed_BuildByNameAcceptsNodeParams(t *testing.T) { t.Parallel() - built, err := BuildByName("pubmed", map[string]any{"top_n": 8, "email": "node@example.com"}) + built, err := BuildByName("pubmed", map[string]any{ + "top_n": 8, "email": "node@example.com", "outputs": map[string]any{"json": map[string]any{}}, + }) if err != nil { t.Fatalf("BuildByName: %v", err) } @@ -254,6 +257,33 @@ func TestPubMed_BuildByNameAcceptsNodeParams(t *testing.T) { } } +func TestPubMed_ComponentReferencesAndOutputs(t *testing.T) { + t.Parallel() + + pubmed := NewPubMedTool() + spec := pubmed.ComponentSpec() + if query, ok := spec.InputForm["query"].(map[string]any); !ok || query["type"] != "line" { + t.Fatalf("query input form = %#v", spec.InputForm["query"]) + } + envelope := map[string]any{"results": []any{map[string]any{ + "title": "Paper", "url": "https://pubmed.ncbi.nlm.nih.gov/1", "content": "Title: Paper\nAbstract: Evidence.", + }}} + chunks, docAggs := pubmed.BuildReferences(context.Background(), envelope) + if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["document_name"] != "Paper" { + t.Fatalf("references = %#v / %#v", chunks, docAggs) + } + outputs := pubmed.BuildComponentOutputs(envelope) + if results, ok := outputs["json"].([]any); !ok || len(results) != 1 { + t.Fatalf("json output = %#v", outputs["json"]) + } + if !strings.Contains(outputs["formalized_content"].(string), "Abstract: Evidence.") { + t.Fatalf("formalized_content = %q", outputs["formalized_content"]) + } + if _, exists := envelope["chunks"]; exists { + t.Fatalf("output conversion mutated envelope: %#v", envelope) + } +} + func TestPubMed_MergeDefaults(t *testing.T) { t.Parallel() @@ -277,3 +307,13 @@ func TestPubMed_BuildByNameRejectsInvalidTopN(t *testing.T) { t.Fatalf("err = %q, want positive integer validation", err.Error()) } } + +func TestPubMed_BuildByNameRejectsInvalidNodeTypes(t *testing.T) { + t.Parallel() + + for _, params := range []map[string]any{{"top_n": 1.5}, {"email": 1}, {"email": ""}} { + if _, err := BuildByName("pubmed", params); err == nil { + t.Fatalf("BuildByName(%#v) succeeded", params) + } + } +} diff --git a/internal/agent/tool/registry.go b/internal/agent/tool/registry.go index 7ece5ff0ca..76cc6edc0b 100644 --- a/internal/agent/tool/registry.go +++ b/internal/agent/tool/registry.go @@ -19,6 +19,7 @@ package tool import ( "fmt" "math" + "strconv" "strings" einotool "github.com/cloudwego/eino/components/tool" @@ -32,12 +33,12 @@ type Factory func(params map[string]any) (einotool.BaseTool, error) var registry = map[string]Factory{ "akshare": buildAkShareTool, "arxiv": buildArxivTool, - "bgpt": noConfig("bgpt", func() einotool.BaseTool { return NewBGPTTool() }), + "bgpt": buildBGPTTool, "code_exec": noConfig("code_exec", func() einotool.BaseTool { return NewCodeExecTool() }), "crawler": noConfig("crawler", func() einotool.BaseTool { return NewCrawlerTool() }), "deepl": noConfig("deepl", func() einotool.BaseTool { return NewDeepLTool() }), - "duckduckgo": noConfig("duckduckgo", func() einotool.BaseTool { return NewDuckDuckGoTool() }), - "email": noConfig("email", func() einotool.BaseTool { return NewEmailTool() }), + "duckduckgo": buildDuckDuckGoTool, + "email": buildEmailTool, "execute_sql": buildExeSQLTool, "exesql": buildExeSQLTool, "github": buildGitHubTool, @@ -52,14 +53,14 @@ var registry = map[string]Factory{ "search_my_dataset": noConfig("search_my_dataset", func() einotool.BaseTool { return NewRetrievalTool() }), "search_my_dateset": noConfig("search_my_dateset", func() einotool.BaseTool { return NewRetrievalTool() }), "searxng": buildSearXNGTool, - "tavily": noConfig("tavily", func() einotool.BaseTool { return NewTavilyTool() }), - "tavily_extract": noConfig("tavily_extract", func() einotool.BaseTool { return NewTavilyExtractTool() }), + "tavily": buildTavilyTool, + "tavily_extract": buildTavilyExtractTool, "tushare": noConfig("tushare", func() einotool.BaseTool { return NewTushareTool() }), "wencai": buildWencaiTool, "web_crawler": noConfig("web_crawler", func() einotool.BaseTool { return NewCrawlerTool() }), "wikipedia": buildWikipediaTool, "wikipedia_search": buildWikipediaTool, - "yahoo_finance": noConfig("yahoo_finance", func() einotool.BaseTool { return NewYahooFinanceTool() }), + "yahoo_finance": buildYahooFinanceTool, } func noConfig(name string, fn func() einotool.BaseTool) Factory { @@ -132,13 +133,6 @@ func buildAkShareTool(params map[string]any) (einotool.BaseTool, error) { func buildArxivTool(params map[string]any) (einotool.BaseTool, error) { topN := defaultArxivTopN sortBy := defaultArxivSortBy - for key := range params { - switch key { - case "top_n", "sort_by": - default: - return nil, fmt.Errorf("agent tool: tool %q does not accept node-level param %s", "arxiv", key) - } - } if v, ok := intParam(params, "top_n"); ok { topN = v } @@ -154,6 +148,36 @@ func buildArxivTool(params map[string]any) (einotool.BaseTool, error) { return NewArxivToolWithParams(nil, topN, sortBy), nil } +func buildBGPTTool(params map[string]any) (einotool.BaseTool, error) { + defaults := bgptParams{} + if value, ok := params["api_key"]; ok { + apiKey, valid := value.(string) + if !valid { + return nil, fmt.Errorf("agent tool: tool %q requires string node-level param api_key", "bgpt") + } + defaults.APIKey = apiKey + } + if value, ok := params["top_n"]; ok { + topN, valid := strictInt(value) + if text, isString := value.(string); isString { + parsed, err := strconv.Atoi(strings.TrimSpace(text)) + topN, valid = parsed, err == nil + } + if !valid || topN <= 0 { + return nil, fmt.Errorf("agent tool: tool %q requires positive integer node-level param top_n", "bgpt") + } + defaults.TopN = topN + } + if value, ok := params["days_back"]; ok && value != nil && value != "" { + daysBack, valid := strictInt(value) + if !valid || daysBack <= 0 { + return nil, fmt.Errorf("agent tool: tool %q requires positive integer node-level param days_back", "bgpt") + } + defaults.DaysBack = daysBack + } + return newBGPTTool(nil, defaults), nil +} + func buildExeSQLTool(params map[string]any) (einotool.BaseTool, error) { conn, err := decodeExeSQLConnParams(params) if err != nil { @@ -162,17 +186,64 @@ func buildExeSQLTool(params map[string]any) (einotool.BaseTool, error) { return NewExeSQLTool(conn), nil } -func buildGoogleTool(params map[string]any) (einotool.BaseTool, error) { - if len(params) == 0 { - return NewGoogleTool(), nil +func buildEmailTool(params map[string]any) (einotool.BaseTool, error) { + defaults := emailParams{SMTPPort: 465} + stringFields := map[string]*string{ + "smtp_server": &defaults.SMTPServer, + "email": &defaults.Email, + "smtp_username": &defaults.SMTPUsername, + "password": &defaults.Password, + "sender_name": &defaults.SenderName, + "to_email": &defaults.ToEmail, + "cc_email": &defaults.CCEmail, + "content": &defaults.Content, + "subject": &defaults.Subject, } - for key := range params { - switch key { - case "api_key", "country", "language", "q", "start", "num": - default: - return nil, fmt.Errorf("agent tool: tool %q does not accept node-level param %s", "google", key) + for key, destination := range stringFields { + value, exists := params[key] + if !exists { + continue } + text, valid := value.(string) + if !valid { + return nil, fmt.Errorf("agent tool: tool %q requires string node-level param %s", "email", key) + } + *destination = text } + if value, exists := params["smtp_port"]; exists { + port, valid := strictInt(value) + if text, isString := value.(string); isString { + parsed, err := strconv.Atoi(strings.TrimSpace(text)) + port, valid = parsed, err == nil + } + if !valid || port <= 0 || port > 65535 { + return nil, fmt.Errorf("agent tool: tool %q requires integer node-level param smtp_port in [1, 65535]", "email") + } + defaults.SMTPPort = port + } + return newEmailTool(defaults), nil +} + +func buildDuckDuckGoTool(params map[string]any) (einotool.BaseTool, error) { + defaults := duckduckgoParams{} + if value, ok := params["top_n"]; ok { + topN, valid := strictInt(value) + if !valid || topN <= 0 { + return nil, fmt.Errorf("agent tool: tool %q requires positive integer node-level param top_n", "duckduckgo") + } + defaults.TopN = topN + } + if value, ok := params["channel"]; ok { + channel, valid := value.(string) + if !valid || (channel != "text" && channel != "general" && channel != "news") { + return nil, fmt.Errorf("agent tool: tool %q has unsupported channel %q", "duckduckgo", channel) + } + defaults.Channel = normalizeDuckDuckGoChannel(channel) + } + return newDuckDuckGoTool(nil, defaults), nil +} + +func buildGoogleTool(params map[string]any) (einotool.BaseTool, error) { defaults := googleParams{} if v, ok := stringParam(params, "api_key"); ok { defaults.APIKey = v @@ -197,11 +268,6 @@ func buildGoogleTool(params map[string]any) (einotool.BaseTool, error) { func buildGitHubTool(params map[string]any) (einotool.BaseTool, error) { topN := defaultGitHubTopN - for key := range params { - if key != "top_n" { - return nil, fmt.Errorf("agent tool: tool %q only accepts node-level param top_n", "github") - } - } if raw, exists := params["top_n"]; exists { value, ok := intParam(params, "top_n") if !ok { @@ -222,55 +288,61 @@ func buildGitHubTool(params map[string]any) (einotool.BaseTool, error) { } func buildGoogleScholarTool(params map[string]any) (einotool.BaseTool, error) { - if len(params) == 0 { - return NewGoogleScholarTool(), nil - } - for key := range params { - switch key { - case "query", "top_n", "sort_by", "year_low", "year_high", "patents": - default: - return nil, fmt.Errorf("agent tool: tool %q does not accept node-level param %s", "google_scholar", key) - } - } defaults := googleScholarParams{} if v, ok := stringParam(params, "query"); ok { defaults.Query = v } - if v, ok := intParam(params, "top_n"); ok { - defaults.TopN = v + if value, ok := params["top_n"]; ok { + topN, valid := strictInt(value) + if !valid || topN <= 0 { + return nil, fmt.Errorf("agent tool: tool %q requires positive integer node-level param top_n", "google_scholar") + } + defaults.TopN = topN } - if v, ok := stringParam(params, "sort_by"); ok { - defaults.SortBy = v + if value, ok := params["sort_by"]; ok { + sortBy, valid := value.(string) + if !valid || (sortBy != "date" && sortBy != "relevance") { + return nil, fmt.Errorf("agent tool: tool %q has unsupported sort_by %q", "google_scholar", sortBy) + } + defaults.SortBy = sortBy } - if v, ok := intParam(params, "year_low"); ok { - defaults.YearLow = v + if value, ok := params["year_low"]; ok && value != nil { + yearLow, valid := strictInt(value) + if !valid || yearLow <= 0 { + return nil, fmt.Errorf("agent tool: tool %q requires positive integer node-level param year_low", "google_scholar") + } + defaults.YearLow = yearLow } - if v, ok := intParam(params, "year_high"); ok { - defaults.YearHigh = v + if value, ok := params["year_high"]; ok && value != nil { + yearHigh, valid := strictInt(value) + if !valid || yearHigh <= 0 { + return nil, fmt.Errorf("agent tool: tool %q requires positive integer node-level param year_high", "google_scholar") + } + defaults.YearHigh = yearHigh } if v, ok := boolParam(params, "patents"); ok { defaults.Patents = &v } + if value, ok := params["patents"]; ok { + if _, valid := value.(bool); !valid { + return nil, fmt.Errorf("agent tool: tool %q requires boolean node-level param patents", "google_scholar") + } + } return NewGoogleScholarToolWithDefaults(nil, defaults), nil } func buildPubMedTool(params map[string]any) (einotool.BaseTool, error) { defaults := pubmedParams{} - for key := range params { - switch key { - case "top_n", "email": - default: - return nil, fmt.Errorf("agent tool: tool %q does not accept node-level param %s", "pubmed", key) - } - } - if topN, ok := intParam(params, "top_n"); ok { - if topN <= 0 { + if value, ok := params["top_n"]; ok { + topN, valid := strictInt(value) + if !valid || topN <= 0 { return nil, fmt.Errorf("agent tool: tool %q requires positive integer node-level param top_n", "pubmed") } defaults.TopN = topN } - if email, ok := stringParam(params, "email"); ok { - if strings.TrimSpace(email) == "" { + if value, ok := params["email"]; ok { + email, valid := value.(string) + if !valid || strings.TrimSpace(email) == "" { return nil, fmt.Errorf("agent tool: tool %q requires non-empty string node-level param email", "pubmed") } defaults.Email = email @@ -280,13 +352,6 @@ func buildPubMedTool(params map[string]any) (einotool.BaseTool, error) { func buildSearXNGTool(params map[string]any) (einotool.BaseTool, error) { defaults := defaultSearXNGParams() - for key := range params { - switch key { - case "top_n", "searxng_url": - default: - return nil, fmt.Errorf("agent tool: tool %q does not accept node-level param %s", "searxng", key) - } - } if value, ok := params["top_n"]; ok { topN, valid := parseSearXNGTopN(value) if !valid || topN <= 0 { @@ -306,13 +371,6 @@ func buildSearXNGTool(params map[string]any) (einotool.BaseTool, error) { func buildWencaiTool(params map[string]any) (einotool.BaseTool, error) { defaults := wencaiParams{} - for key := range params { - switch key { - case "top_n", "query_type": - default: - return nil, fmt.Errorf("agent tool: tool %q does not accept node-level param %s", "wencai", key) - } - } if value, ok := params["top_n"]; ok { topN, valid := strictInt(value) if !valid || topN <= 0 { @@ -330,30 +388,128 @@ func buildWencaiTool(params map[string]any) (einotool.BaseTool, error) { return newWencaiTool(defaults), nil } -func buildKeenableTool(params map[string]any) (einotool.BaseTool, error) { - if len(params) == 0 { - return NewKeenableTool(), nil +func buildTavilyExtractTool(params map[string]any) (einotool.BaseTool, error) { + defaults := tavilyExtractParams{} + if value, ok := params["api_key"]; ok { + apiKey, valid := value.(string) + if !valid { + return nil, fmt.Errorf("agent tool: tool %q requires string node-level param api_key", "tavily_extract") + } + defaults.APIKey = apiKey } - for key := range params { - if key != "api_key" { - return nil, fmt.Errorf("agent tool: tool %q only accepts node-level param api_key", "keenable") + if value, ok := params["urls"]; ok { + defaults.URLs = value + } + if value, ok := params["extract_depth"]; ok { + extractDepth, valid := value.(string) + if !valid || (extractDepth != "basic" && extractDepth != "advanced") { + return nil, fmt.Errorf("agent tool: tool %q has unsupported extract_depth %q", "tavily_extract", extractDepth) + } + defaults.ExtractDepth = extractDepth + } + if value, ok := params["format"]; ok { + format, valid := value.(string) + if !valid || (format != "markdown" && format != "text") { + return nil, fmt.Errorf("agent tool: tool %q has unsupported format %q", "tavily_extract", format) + } + defaults.Format = format + } + return newTavilyExtractTool(nil, nil, defaults), nil +} + +func buildTavilyTool(params map[string]any) (einotool.BaseTool, error) { + defaults := tavilyParams{} + if value, ok := params["api_key"]; ok { + apiKey, valid := value.(string) + if !valid { + return nil, fmt.Errorf("agent tool: tool %q requires string node-level param api_key", "tavily") + } + defaults.APIKey = apiKey + } + if value, ok := params["search_depth"]; ok { + searchDepth, valid := value.(string) + if !valid || (searchDepth != "basic" && searchDepth != "advanced") { + return nil, fmt.Errorf("agent tool: tool %q has unsupported search_depth %q", "tavily", searchDepth) + } + defaults.SearchDepth = searchDepth + } + if value, ok := params["max_results"]; ok { + maxResults, valid := strictInt(value) + if !valid || maxResults <= 0 || maxResults > 20 { + return nil, fmt.Errorf("agent tool: tool %q requires integer node-level param max_results within [1, 20]", "tavily") + } + defaults.MaxResults = maxResults + } + if value, ok := params["days"]; ok { + days, valid := strictInt(value) + if !valid || days <= 0 { + return nil, fmt.Errorf("agent tool: tool %q requires positive integer node-level param days", "tavily") + } + defaults.Days = days + } + for _, key := range []string{"include_answer", "include_raw_content", "include_images", "include_image_descriptions"} { + value, ok := params[key] + if !ok { + continue + } + flag, valid := value.(bool) + if !valid { + return nil, fmt.Errorf("agent tool: tool %q requires boolean node-level param %s", "tavily", key) + } + switch key { + case "include_answer": + defaults.IncludeAnswer = flag + case "include_raw_content": + defaults.IncludeRawContent = flag + case "include_images": + defaults.IncludeImages = flag + case "include_image_descriptions": + defaults.IncludeImageDescriptions = flag } } - apiKey, ok := params["api_key"].(string) - if !ok || strings.TrimSpace(apiKey) == "" { - return nil, fmt.Errorf("agent tool: tool %q requires non-empty string node-level param api_key", "keenable") + return newTavilyTool(nil, nil, defaults), nil +} + +func buildKeenableTool(params map[string]any) (einotool.BaseTool, error) { + defaults := keenableParams{} + apiKey := "" + if value, ok := params["api_key"]; ok { + var valid bool + apiKey, valid = value.(string) + if !valid { + return nil, fmt.Errorf("agent tool: tool %q requires string node-level param api_key", "keenable") + } } - return NewKeenableToolWithAPIKey(nil, apiKey), nil + if value, ok := params["mode"]; ok { + mode, valid := value.(string) + if !valid || (mode != "pro" && mode != "realtime") { + return nil, fmt.Errorf("agent tool: tool %q has unsupported mode %q", "keenable", mode) + } + defaults.Mode = mode + } + if value, ok := params["top_n"]; ok { + topN, valid := strictInt(value) + if !valid || topN <= 0 { + return nil, fmt.Errorf("agent tool: tool %q requires positive integer node-level param top_n", "keenable") + } + defaults.TopN = topN + } + if value, ok := params["site"]; ok { + site, valid := value.(string) + if !valid { + return nil, fmt.Errorf("agent tool: tool %q requires string node-level param site", "keenable") + } + defaults.Site = site + } + if defaults.Mode == "realtime" && strings.TrimSpace(apiKey) == "" { + return nil, fmt.Errorf("agent tool: tool %q requires api_key for realtime mode", "keenable") + } + return newKeenableTool(nil, nil, apiKey, defaults), nil } func buildWikipediaTool(params map[string]any) (einotool.BaseTool, error) { topN := defaultWikipediaTopN language := defaultWikipediaLanguage - for key := range params { - if key != "top_n" && key != "language" { - return nil, fmt.Errorf("agent tool: tool %q only accepts node-level params top_n/language", "wikipedia") - } - } if v, ok := intParam(params, "top_n"); ok { topN = v } @@ -372,6 +528,18 @@ func buildWikipediaTool(params map[string]any) (einotool.BaseTool, error) { return NewWikipediaToolWithParams(nil, topN, language), nil } +func buildYahooFinanceTool(params map[string]any) (einotool.BaseTool, error) { + defaults := defaultYahooFinanceParams() + if value, exists := params["info"]; exists { + flag, valid := value.(bool) + if !valid { + return nil, fmt.Errorf("agent tool: tool %q requires boolean node-level param info", "yahoo_finance") + } + defaults.Info = flag + } + return NewYahooFinanceToolWithDefaults(nil, defaults), nil +} + func decodeExeSQLConnParams(params map[string]any) (exesqlConnParams, error) { if len(params) == 0 { return exesqlConnParams{}, fmt.Errorf( diff --git a/internal/agent/tool/registry_test.go b/internal/agent/tool/registry_test.go index dba1f761ca..a01a318f8b 100644 --- a/internal/agent/tool/registry_test.go +++ b/internal/agent/tool/registry_test.go @@ -92,18 +92,6 @@ func TestBuildAll_ExeSQLRequiresNodeParams(t *testing.T) { } } -func TestBuildAll_KeenableRejectsEmptyNodeAPIKey(t *testing.T) { - _, err := BuildAll([]string{"keenable"}, map[string]map[string]any{ - "keenable": {"api_key": ""}, - }) - if err == nil { - t.Fatal("expected keenable config error") - } - if !strings.Contains(err.Error(), "requires non-empty string node-level param api_key") { - t.Fatalf("err = %q, want keenable api_key validation error", err.Error()) - } -} - // TestToolRegistry_SchemasAreComplete sweeps every name the public // registry advertises (including the execute_sql/exesql and // retrieval/search_my_dateset alias pairs), builds the tool, and @@ -181,8 +169,8 @@ func TestToolRegistry_SchemasAreComplete(t *testing.T) { canonicalByAlias := map[string]string{ "execute_sql": "execute_sql", "exesql": "execute_sql", - "google_scholar": "google_scholar", - "google_scholar_search": "google_scholar", + "google_scholar": "google_scholar_search", + "google_scholar_search": "google_scholar_search", "retrieval": "search_my_dateset", "search_my_dataset": "search_my_dateset", "search_my_dateset": "search_my_dateset", diff --git a/internal/agent/tool/searxng.go b/internal/agent/tool/searxng.go index f6f82972bc..64f365442e 100644 --- a/internal/agent/tool/searxng.go +++ b/internal/agent/tool/searxng.go @@ -18,26 +18,35 @@ package tool import ( "context" + "crypto/sha1" "encoding/json" "fmt" "net" "net/http" "net/url" + "regexp" "strconv" "strings" "time" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" + + "ragflow/internal/tokenizer" ) const ( - searxngToolName = "searxng_search" - searxngToolDescription = "SearXNG is a privacy-focused metasearch engine that aggregates results from multiple search engines without tracking users. It provides comprehensive web search capabilities." - defaultSearXNGTopN = 10 - searxngRequestTimeout = 10 * time.Second + searxngToolName = "searxng_search" + searxngToolDescription = "SearXNG is a privacy-focused metasearch engine that aggregates results from multiple search engines without tracking users. It provides comprehensive web search capabilities." + defaultSearXNGTopN = 10 + searxngRequestTimeout = 10 * time.Second + searxngPromptTokenLimit = 200000 ) +var searxngDataImagePattern = regexp.MustCompile(`!?\[[a-z]+\]\(data:image/png;base64,[ 0-9A-Za-z/_=+\-]+\)`) + +var searxngNewlinePattern = regexp.MustCompile(`\n+`) + // searxngParams mirrors the SearXNG-specific Python parameters. Query and // searxng_url are model inputs; searxng_url may also be node configuration, // while top_n is Canvas-only configuration. @@ -61,6 +70,9 @@ type SearXNGTool struct { resolve searxngResolver } +var _ ToolComponent = (*SearXNGTool)(nil) +var _ ReferenceBuilder = (*SearXNGTool)(nil) + func defaultSearXNGParams() searxngParams { return searxngParams{TopN: defaultSearXNGTopN} } @@ -212,6 +224,154 @@ func parseSearXNGTopN(value any) (int, bool) { return strictInt(value) } +// ComponentSpec returns the Python-compatible SearXNG Canvas surface. +func (s *SearXNGTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "query": "The search keywords to execute with SearXNG.", + "searxng_url": "The base URL of the SearXNG instance.", + }, + Outputs: map[string]string{ + "formalized_content": "Rendered SearXNG references for downstream LLM prompts.", + "json": "Raw SearXNG result list.", + }, + InputForm: map[string]any{ + "query": map[string]any{ + "name": "Query", + "type": "line", + }, + "searxng_url": map[string]any{ + "name": "SearXNG URL", + "type": "line", + "placeholder": "http://localhost:4000", + }, + }, + } +} + +// BuildReferences builds the same references as Python ToolBase._retrieve_chunks. +func (s *SearXNGTool) BuildReferences(_ context.Context, envelope map[string]any) ([]map[string]any, []map[string]any) { + return buildSearXNGReferences(envelope) +} + +// BuildComponentOutputs converts SearXNG's complete tool envelope into its +// public Canvas outputs. +func (s *SearXNGTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + results := envelopeSlice(envelope, "results") + chunks, _ := buildSearXNGReferences(envelope) + return map[string]any{ + "formalized_content": renderSearXNGReferences(chunks, searxngPromptTokenLimit), + "json": results, + } +} + +func buildSearXNGReferences(envelope map[string]any) ([]map[string]any, []map[string]any) { + results := envelopeSlice(envelope, "results") + chunks := make([]map[string]any, 0, len(results)) + docAggs := make([]map[string]any, 0, len(results)) + for _, result := range results { + item, ok := result.(map[string]any) + if !ok { + continue + } + content, _ := item["content"].(string) + if content == "" { + continue + } + content = searxngDataImagePattern.ReplaceAllString(content, "") + runes := []rune(content) + if len(runes) > 10000 { + content = string(runes[:10000]) + } + if content == "" { + continue + } + + documentID := strconv.Itoa(hashSearXNGString(content, 100000000)) + displayID := strconv.Itoa(hashSearXNGString(documentID, 500)) + title := searxngText(item["title"]) + resultURL := searxngText(item["url"]) + chunks = append(chunks, map[string]any{ + "id": displayID, + "chunk_id": documentID, + "content": content, + "doc_id": documentID, + "docnm_kwd": title, + "document_id": documentID, + "document_name": title, + "dataset_id": nil, + "image_id": nil, + "positions": nil, + "url": resultURL, + "similarity": 1, + "vector_similarity": nil, + "term_similarity": nil, + "row_id": nil, + "doc_type": nil, + "document_metadata": nil, + }) + docAggs = append(docAggs, map[string]any{ + "doc_name": title, + "doc_id": documentID, + "count": 1, + "url": resultURL, + }) + } + return chunks, docAggs +} + +func renderSearXNGReferences(chunks []map[string]any, maxTokens int) string { + usedTokens := 0 + blocks := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + content := searxngText(chunk["content"]) + if content == "" { + continue + } + var block strings.Builder + fmt.Fprintf(&block, "\nID: %s", searxngText(chunk["id"])) + if title := searxngPromptField(chunk["document_name"]); title != "" { + fmt.Fprintf(&block, "\n├── Title: %s", title) + } + if resultURL := searxngPromptField(chunk["url"]); resultURL != "" { + fmt.Fprintf(&block, "\n├── URL: %s", resultURL) + } + block.WriteString("\n└── Content:\n") + block.WriteString(content) + completeBlock := block.String() + blockTokens := tokenizer.NumTokensFromString(completeBlock) + if maxTokens > 0 && float64(usedTokens+blockTokens) > float64(maxTokens)*0.97 { + break + } + usedTokens += blockTokens + blocks = append(blocks, completeBlock) + } + return strings.Join(blocks, "\n") +} + +func searxngText(value any) string { + if value == nil { + return "" + } + if text, ok := value.(string); ok { + return text + } + return fmt.Sprint(value) +} + +func searxngPromptField(value any) string { + return searxngNewlinePattern.ReplaceAllString(searxngText(value), " ") +} + +func hashSearXNGString(value string, modulus int) int { + digest := sha1.Sum([]byte(value)) + result := 0 + for _, part := range digest { + result = (result*256 + int(part)) % modulus + } + return result +} + func searxngJSON(env searxngEnvelope) string { data, err := json.Marshal(env) if err != nil { diff --git a/internal/agent/tool/searxng_test.go b/internal/agent/tool/searxng_test.go index 8d9210f1fa..22af5f3f0a 100644 --- a/internal/agent/tool/searxng_test.go +++ b/internal/agent/tool/searxng_test.go @@ -24,9 +24,12 @@ import ( "net/http" "net/http/httptest" "net/url" + "strconv" "strings" "sync/atomic" "testing" + + "ragflow/internal/tokenizer" ) func TestSearXNGBuildURLMatchesPythonQuery(t *testing.T) { @@ -168,6 +171,7 @@ func TestSearXNGBuildByNameAcceptsPythonNodeParams(t *testing.T) { built, err := BuildByName("searxng", map[string]any{ "top_n": "8", "searxng_url": "https://searx.example.com", + "outputs": map[string]any{"json": map[string]any{}}, }) if err != nil { t.Fatalf("BuildByName: %v", err) @@ -188,7 +192,6 @@ func TestSearXNGBuildByNameRejectsInvalidNodeParams(t *testing.T) { {"top_n": "abc"}, {"top_n": 0}, {"top_n": 1.5}, - {"unknown": true}, } for _, params := range invalid { if _, err := BuildByName("searxng", params); err == nil { @@ -197,6 +200,87 @@ func TestSearXNGBuildByNameRejectsInvalidNodeParams(t *testing.T) { } } +func TestSearXNGComponentContractReferencesAndOutputs(t *testing.T) { + t.Parallel() + + tool := NewSearXNGTool() + spec := tool.ComponentSpec() + for _, input := range []string{"query", "searxng_url"} { + if _, ok := spec.Inputs[input]; !ok { + t.Fatalf("component inputs missing %s: %#v", input, spec.Inputs) + } + } + for _, output := range []string{"formalized_content", "json"} { + if _, ok := spec.Outputs[output]; !ok { + t.Fatalf("component outputs missing %s: %#v", output, spec.Outputs) + } + } + serverURL := spec.InputForm["searxng_url"].(map[string]any) + if serverURL["placeholder"] != "http://localhost:4000" { + t.Fatalf("searxng_url input form = %#v", serverURL) + } + + content := "RAGFlow content ![img](data:image/png;base64,AAAA) remains" + envelope := map[string]any{"results": []any{ + map[string]any{"title": "RAGFlow\nDocs", "url": "https://ragflow.io", "content": content, "engine": "bing", "score": 0.9}, + map[string]any{"title": "Empty", "url": "https://example.com", "content": ""}, + }} + chunks, docAggs := tool.BuildReferences(context.Background(), envelope) + if len(chunks) != 1 || len(docAggs) != 1 { + t.Fatalf("references = %#v / %#v", chunks, docAggs) + } + cleaned := "RAGFlow content remains" + documentID := hashSearXNGString(cleaned, 100000000) + referenceID := hashSearXNGString(strconv.Itoa(documentID), 500) + if documentID != 93760153 || referenceID != 491 { + t.Fatalf("hash parity = %d/%d, want Python 93760153/491", documentID, referenceID) + } + if chunks[0]["document_id"] != strconv.Itoa(documentID) || chunks[0]["id"] != strconv.Itoa(referenceID) || chunks[0]["content"] != cleaned { + t.Fatalf("reference chunk = %#v", chunks[0]) + } + + outputs := tool.BuildComponentOutputs(envelope) + results, ok := outputs["json"].([]any) + if !ok || len(results) != 2 || results[0].(map[string]any)["engine"] != "bing" { + t.Fatalf("json output lost raw fields: %#v", outputs["json"]) + } + formalized := outputs["formalized_content"].(string) + for _, want := range []string{ + "ID: " + strconv.Itoa(referenceID), + "Title: RAGFlow Docs", + "URL: https://ragflow.io", + "Content:\n" + cleaned, + } { + if !strings.Contains(formalized, want) { + t.Fatalf("formalized_content missing %q: %s", want, formalized) + } + } + if _, exists := envelope["chunks"]; exists { + t.Fatalf("component conversion mutated envelope: %#v", envelope) + } +} + +func TestRenderSearXNGReferencesStopsBeforeOverBudgetBlock(t *testing.T) { + t.Parallel() + + chunks := []map[string]any{ + {"id": "1", "document_name": "First", "url": "https://first.example", "content": "first reference content"}, + {"id": "2", "document_name": "Second", "url": "https://second.example", "content": "second reference content"}, + } + firstBlock := renderSearXNGReferences(chunks[:1], 0) + firstTokens := tokenizer.NumTokensFromString(firstBlock) + maxTokens := (firstTokens*100 + 96) / 97 + if got := renderSearXNGReferences(chunks, maxTokens); got != firstBlock { + t.Fatalf("rendered = %q, want only first block %q", got, firstBlock) + } + if got := renderSearXNGReferences(chunks, 1); got != "" { + t.Fatalf("over-budget first block was appended: %q", got) + } + if got := renderSearXNGReferences(chunks, 0); !strings.Contains(got, "Title: First") || !strings.Contains(got, "Title: Second") { + t.Fatalf("unlimited rendering dropped blocks: %q", got) + } +} + func TestSearXNGDoesNotRetryFailedRequest(t *testing.T) { t.Parallel() diff --git a/internal/agent/tool/tavily.go b/internal/agent/tool/tavily.go index d692172643..aab4a49a3f 100644 --- a/internal/agent/tool/tavily.go +++ b/internal/agent/tool/tavily.go @@ -18,22 +18,35 @@ package tool import ( "context" + "crypto/sha1" "encoding/json" "fmt" + "math/big" "net/http" - "ragflow/internal/common" + "regexp" + "strconv" "strings" + "unicode/utf8" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" + + "ragflow/internal/common" + "ragflow/internal/tokenizer" ) -const tavilyToolName = "tavily" +const tavilyToolName = "tavily_search" const tavilyExtractToolName = "tavily_extract" const tavilyToolDescription = "Search the web via the Tavily API. Returns a list of {url, title, content} results." +const tavilyPromptMaxTokens = 200000 + +var tavilyDataImagePattern = regexp.MustCompile(`!?\[[a-z]+\]\(data:image/png;base64,[ 0-9A-Za-z/_=+\-]+\)`) + +var tavilyNewlinePattern = regexp.MustCompile(`\n+`) + // tavilyParams is the JSON shape the model sends into InvokableRun. The // api_key may be omitted when the env var TAVILY_API_KEY is set; the tool // resolves it from the environment in that case. @@ -68,27 +81,17 @@ type tavilyRequestBody struct { ExcludeDomains []string `json:"exclude_domains,omitempty"` } -// tavilyResult mirrors one element of the upstream `results` array. We -// return these verbatim to the model. -type tavilyResult struct { - URL string `json:"url"` - Title string `json:"title"` - Content string `json:"content"` - RawContent string `json:"raw_content,omitempty"` - Score float64 `json:"score,omitempty"` -} - -// tavilyResponse is the envelope returned by Tavily. We only model the -// fields we care about; the upstream API has more, but they are ignored. +// tavilyResponse preserves every upstream result field because Python exposes +// response["results"] directly through the Canvas json output. type tavilyResponse struct { - Results []tavilyResult `json:"results"` + Results []map[string]any `json:"results"` } // tavilyEnvelope is the shape the model actually sees, identical to the // Python tool's output convention. type tavilyEnvelope struct { - Results []tavilyResult `json:"results"` - Error string `json:"_ERROR,omitempty"` + Results []map[string]any `json:"results"` + Error string `json:"_ERROR,omitempty"` } type tavilyExtractParams struct { @@ -105,20 +108,13 @@ type tavilyExtractRequestBody struct { IncludeImages bool `json:"include_images"` } -type tavilyExtractResult struct { - URL string `json:"url"` - RawContent string `json:"raw_content,omitempty"` - Content string `json:"content,omitempty"` - Error string `json:"error,omitempty"` -} - type tavilyExtractResponse struct { - Results []tavilyExtractResult `json:"results"` + Results []map[string]any `json:"results"` } type tavilyExtractEnvelope struct { - Results []tavilyExtractResult `json:"results"` - Error string `json:"_ERROR,omitempty"` + Results []map[string]any `json:"results"` + Error string `json:"_ERROR,omitempty"` } // TavilyTool is the Tavily search @@ -126,67 +122,93 @@ type tavilyExtractEnvelope struct { // to https://api.tavily.com/search using the shared HTTPHelper and returns // the upstream `results` array as JSON. type TavilyTool struct { - helper *HTTPHelper - envKey func() string + helper *HTTPHelper + envKey func() string + defaults tavilyParams } // TavilyExtractTool is the Tavily Extract tool. It POSTs URLs to // https://api.tavily.com/extract and returns the upstream results array. type TavilyExtractTool struct { - helper *HTTPHelper - envKey func() string + helper *HTTPHelper + envKey func() string + defaults tavilyExtractParams } +var _ ToolComponent = (*TavilyTool)(nil) +var _ ReferenceBuilder = (*TavilyTool)(nil) +var _ ToolComponent = (*TavilyExtractTool)(nil) + // NewTavilyTool returns a TavilyTool using the default HTTPHelper and // the TAVILY_API_KEY env var for credential resolution. func NewTavilyTool() *TavilyTool { - return NewTavilyToolWith(NewHTTPHelper()) + return newTavilyTool(nil, nil, tavilyParams{}) } // NewTavilyToolWith returns a TavilyTool that uses the provided // HTTPHelper. Useful for tests that want to inject a custom transport. func NewTavilyToolWith(h *HTTPHelper) *TavilyTool { - if h == nil { - h = NewHTTPHelper() - } - return &TavilyTool{helper: h, envKey: defaultTavilyEnvKey} + return newTavilyTool(h, nil, tavilyParams{}) } // NewTavilyToolWithEnvKey returns a TavilyTool with a custom env-key // resolver. Useful for tests that want to inject a fake credential // without mutating process state. func NewTavilyToolWithEnvKey(h *HTTPHelper, envKey func() string) *TavilyTool { + return newTavilyTool(h, envKey, tavilyParams{}) +} + +func newTavilyTool(h *HTTPHelper, envKey func() string, defaults tavilyParams) *TavilyTool { if h == nil { h = NewHTTPHelper() } if envKey == nil { envKey = defaultTavilyEnvKey } - return &TavilyTool{helper: h, envKey: envKey} + if defaults.MaxResults == 0 { + defaults.MaxResults = 6 + } + if defaults.SearchDepth == "" { + defaults.SearchDepth = "basic" + } + if defaults.Topic == "" { + defaults.Topic = "general" + } + if defaults.Days == 0 { + defaults.Days = 14 + } + return &TavilyTool{helper: h, envKey: envKey, defaults: defaults} } // NewTavilyExtractTool returns a TavilyExtractTool using the default HTTPHelper. func NewTavilyExtractTool() *TavilyExtractTool { - return NewTavilyExtractToolWith(NewHTTPHelper()) + return newTavilyExtractTool(nil, nil, tavilyExtractParams{}) } // NewTavilyExtractToolWith returns a TavilyExtractTool using the provided helper. func NewTavilyExtractToolWith(h *HTTPHelper) *TavilyExtractTool { - if h == nil { - h = NewHTTPHelper() - } - return &TavilyExtractTool{helper: h, envKey: defaultTavilyEnvKey} + return newTavilyExtractTool(h, nil, tavilyExtractParams{}) } // NewTavilyExtractToolWithEnvKey returns a TavilyExtractTool with a custom env-key resolver. func NewTavilyExtractToolWithEnvKey(h *HTTPHelper, envKey func() string) *TavilyExtractTool { + return newTavilyExtractTool(h, envKey, tavilyExtractParams{}) +} + +func newTavilyExtractTool(h *HTTPHelper, envKey func() string, defaults tavilyExtractParams) *TavilyExtractTool { if h == nil { h = NewHTTPHelper() } if envKey == nil { envKey = defaultTavilyEnvKey } - return &TavilyExtractTool{helper: h, envKey: envKey} + if defaults.ExtractDepth == "" { + defaults.ExtractDepth = "basic" + } + if defaults.Format == "" { + defaults.Format = "markdown" + } + return &TavilyExtractTool{helper: h, envKey: envKey, defaults: defaults} } // defaultTavilyEnvKey is the production env-key resolver. Pulled out @@ -205,21 +227,6 @@ func (t *TavilyTool) Info(_ context.Context) (*schema.ToolInfo, error) { Desc: "Search query", Required: true, }, - "api_key": { - Type: schema.String, - Desc: "Tavily API key. Falls back to TAVILY_API_KEY env var.", - Required: false, - }, - "max_results": { - Type: schema.Integer, - Desc: "Maximum number of results to return. Defaults to 5.", - Required: false, - }, - "search_depth": { - Type: schema.String, - Desc: `Tavily search depth: "basic" (default) or "advanced".`, - Required: false, - }, "topic": { Type: schema.String, Desc: `Search topic: "general" (default) or "news".`, @@ -282,15 +289,11 @@ func (t *TavilyTool) InvokableRun(ctx context.Context, argsJSON string, _ ...too fmt.Errorf("tavily: parse arguments: %w", err) } if p.Query == "" { - return tavilyErrJSON(fmt.Errorf("query is required")), - fmt.Errorf("tavily: query is required") - } - if p.MaxResults <= 0 { - p.MaxResults = 6 - } - if p.SearchDepth == "" { - p.SearchDepth = "basic" + return tavilyJSON(tavilyEnvelope{Results: []map[string]any{}}), nil } + var provided map[string]json.RawMessage + _ = json.Unmarshal([]byte(argsJSON), &provided) + p = mergeTavilyParams(t.defaults, p, provided) apiKey := p.APIKey if apiKey == "" { @@ -308,8 +311,8 @@ func (t *TavilyTool) InvokableRun(ctx context.Context, argsJSON string, _ ...too Topic: defaultString(p.Topic, "general"), Days: p.Days, IncludeAnswer: p.IncludeAnswer, - IncludeRawContent: false, - IncludeImages: false, + IncludeRawContent: p.IncludeRawContent, + IncludeImages: p.IncludeImages, IncludeImageDescriptions: p.IncludeImageDescriptions, IncludeDomains: p.IncludeDomains, ExcludeDomains: p.ExcludeDomains, @@ -337,14 +340,181 @@ func (t *TavilyTool) InvokableRun(ctx context.Context, argsJSON string, _ ...too return tavilyJSON(tavilyEnvelope{Results: raw.Results}), nil } +func mergeTavilyParams(defaults, params tavilyParams, provided map[string]json.RawMessage) tavilyParams { + if params.APIKey == "" { + params.APIKey = defaults.APIKey + } + if params.MaxResults == 0 { + params.MaxResults = defaults.MaxResults + } + if params.SearchDepth == "" { + params.SearchDepth = defaults.SearchDepth + } + if params.Topic == "" { + params.Topic = defaults.Topic + } + if params.Days == 0 { + params.Days = defaults.Days + } + if params.IncludeDomains == nil { + params.IncludeDomains = defaults.IncludeDomains + } + if params.ExcludeDomains == nil { + params.ExcludeDomains = defaults.ExcludeDomains + } + if _, ok := provided["include_answer"]; !ok { + params.IncludeAnswer = defaults.IncludeAnswer + } + if _, ok := provided["include_raw_content"]; !ok { + params.IncludeRawContent = defaults.IncludeRawContent + } + if _, ok := provided["include_images"]; !ok { + params.IncludeImages = defaults.IncludeImages + } + if _, ok := provided["include_image_descriptions"]; !ok { + params.IncludeImageDescriptions = defaults.IncludeImageDescriptions + } + return params +} + +// ComponentSpec returns TavilySearch's Canvas-facing metadata. +func (t *TavilyTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "query": "Search query.", + "topic": `Search topic: "general" or "news".`, + "include_domains": "Domains that search results must include.", + "exclude_domains": "Domains that search results must exclude.", + }, + Outputs: map[string]string{ + "formalized_content": "Rendered Tavily references for downstream prompts.", + "json": "Raw Tavily result list.", + }, + InputForm: map[string]any{ + "query": map[string]any{"name": "Query", "type": "line"}, + "topic": map[string]any{"name": "Topic", "type": "line"}, + "include_domains": map[string]any{"name": "Include domains", "type": "line"}, + "exclude_domains": map[string]any{"name": "Exclude domains", "type": "line"}, + }, + } +} + +func (t *TavilyTool) BuildReferences(_ context.Context, envelope map[string]any) ([]map[string]any, []map[string]any) { + return buildTavilyReferences(envelope) +} + +func (t *TavilyTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + results := envelopeSlice(envelope, "results") + chunks, _ := buildTavilyReferences(envelope) + return map[string]any{ + "formalized_content": renderTavilyReferences(chunks, tavilyPromptMaxTokens), + "json": results, + } +} + +func buildTavilyReferences(envelope map[string]any) ([]map[string]any, []map[string]any) { + results := envelopeSlice(envelope, "results") + chunks := make([]map[string]any, 0, len(results)) + docAggs := make([]map[string]any, 0, len(results)) + for _, result := range results { + item, ok := result.(map[string]any) + if !ok { + continue + } + content := tavilyText(item["raw_content"]) + if content == "" { + content = tavilyText(item["content"]) + } + content = tavilyDataImagePattern.ReplaceAllString(content, "") + content = truncateTavilyRunes(content, 10000) + if content == "" { + continue + } + documentID := strconv.FormatInt(tavilyHashInt(content, 100000000), 10) + displayID := strconv.FormatInt(tavilyHashInt(documentID, 500), 10) + title := tavilyText(item["title"]) + resultURL := tavilyText(item["url"]) + score := item["score"] + chunks = append(chunks, map[string]any{ + "id": displayID, + "chunk_id": documentID, + "content": content, + "doc_id": documentID, + "document_id": documentID, + "docnm_kwd": title, + "document_name": title, + "similarity": score, + "score": score, + "url": resultURL, + }) + docAggs = append(docAggs, map[string]any{ + "doc_name": title, + "doc_id": documentID, + "count": 1, + "url": resultURL, + }) + } + return chunks, docAggs +} + +func renderTavilyReferences(chunks []map[string]any, maxTokens int) string { + usedTokens := 0 + blocks := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + content := tavilyText(chunk["content"]) + if content == "" { + continue + } + usedTokens += tokenizer.NumTokensFromString(content) + blocks = append(blocks, strings.Join([]string{ + "\nID: " + tavilyText(chunk["id"]), + "├── Title: " + tavilyPromptField(chunk["document_name"]), + "├── URL: " + tavilyPromptField(chunk["url"]), + "└── Content:\n" + content, + }, "\n")) + if maxTokens > 0 && float64(maxTokens)*0.97 < float64(usedTokens) { + break + } + } + return strings.Join(blocks, "\n") +} + +func tavilyText(value any) string { + if value == nil { + return "" + } + if text, ok := value.(string); ok { + return text + } + return fmt.Sprint(value) +} + +func tavilyPromptField(value any) string { + return tavilyNewlinePattern.ReplaceAllString(tavilyText(value), " ") +} + +func tavilyHashInt(value string, modulus int64) int64 { + digest := sha1.Sum([]byte(value)) + number := new(big.Int).SetBytes(digest[:]) + return new(big.Int).Mod(number, big.NewInt(modulus)).Int64() +} + +func truncateTavilyRunes(value string, limit int) string { + if utf8.RuneCountInString(value) <= limit { + return value + } + return string([]rune(value)[:limit]) +} + // InvokableRun performs the Tavily Extract request. The api_key may come from // the argument or the TAVILY_API_KEY env var. func (t *TavilyExtractTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) { - var p tavilyExtractParams - if err := json.Unmarshal([]byte(argsJSON), &p); err != nil { + var runtimeParams tavilyExtractParams + if err := json.Unmarshal([]byte(argsJSON), &runtimeParams); err != nil { return tavilyExtractErrJSON(fmt.Errorf("tavily_extract: parse arguments: %w", err)), fmt.Errorf("tavily_extract: parse arguments: %w", err) } + p := mergeTavilyExtractParams(t.defaults, runtimeParams) urls := normalizeTavilyURLs(p.URLs) if len(urls) == 0 { return tavilyExtractErrJSON(fmt.Errorf("urls is required")), @@ -394,6 +564,47 @@ func (t *TavilyExtractTool) InvokableRun(ctx context.Context, argsJSON string, _ return tavilyExtractJSON(tavilyExtractEnvelope{Results: raw.Results}), nil } +func mergeTavilyExtractParams(defaults, params tavilyExtractParams) tavilyExtractParams { + if params.APIKey == "" { + params.APIKey = defaults.APIKey + } + if params.URLs == nil { + params.URLs = defaults.URLs + } + if params.ExtractDepth == "" { + params.ExtractDepth = defaults.ExtractDepth + } + if params.Format == "" { + params.Format = defaults.Format + } + return params +} + +// ComponentSpec returns the Python-compatible TavilyExtract Canvas surface. +func (t *TavilyExtractTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "urls": "The URLs to extract content from.", + "extract_depth": `Extraction depth: "basic" or "advanced".`, + "format": `Output format: "markdown" or "text".`, + }, + Outputs: map[string]string{ + "json": "Raw Tavily Extract results.", + }, + InputForm: map[string]any{ + "urls": map[string]any{"name": "URLs", "type": "line"}, + "extract_depth": map[string]any{"name": "Extract depth", "type": "line"}, + "format": map[string]any{"name": "Format", "type": "line"}, + }, + } +} + +// BuildComponentOutputs converts TavilyExtract's complete tool envelope into +// its public Canvas outputs. +func (t *TavilyExtractTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + return map[string]any{"json": envelopeSlice(envelope, "results")} +} + // tavilyJSON marshals the envelope to a JSON string for the model. func tavilyJSON(env tavilyEnvelope) string { b, err := json.Marshal(env) diff --git a/internal/agent/tool/tavily_test.go b/internal/agent/tool/tavily_test.go index 7442b7453b..1113f56b42 100644 --- a/internal/agent/tool/tavily_test.go +++ b/internal/agent/tool/tavily_test.go @@ -50,7 +50,7 @@ func TestTavily_BuildRequest(t *testing.T) { }) tool := NewTavilyToolWith(helper) out, err := tool.InvokableRun(context.Background(), - `{"query":"ragflow","api_key":"key-xyz","max_results":3,"search_depth":"advanced"}`) + `{"query":"ragflow","api_key":"key-xyz","max_results":3,"search_depth":"advanced","include_raw_content":true,"include_images":true}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -79,6 +79,9 @@ func TestTavily_BuildRequest(t *testing.T) { if gotBody["search_depth"] != "advanced" { t.Errorf("body.search_depth = %v, want advanced", gotBody["search_depth"]) } + if gotBody["include_raw_content"] != true || gotBody["include_images"] != true { + t.Errorf("include flags = raw:%v images:%v, want true/true", gotBody["include_raw_content"], gotBody["include_images"]) + } } func TestTavily_ParseResponse(t *testing.T) { @@ -116,11 +119,11 @@ func TestTavily_ParseResponse(t *testing.T) { if len(env.Results) != 2 { t.Fatalf("Results len = %d, want 2", len(env.Results)) } - if env.Results[0].URL != "https://a.example/" || env.Results[0].Title != "A" { + if env.Results[0]["url"] != "https://a.example/" || env.Results[0]["title"] != "A" { t.Errorf("Results[0] = %+v, want url=https://a.example/ title=A", env.Results[0]) } - if env.Results[1].Content != "beta" { - t.Errorf("Results[1].Content = %q, want beta", env.Results[1].Content) + if env.Results[1]["content"] != "beta" { + t.Errorf("Results[1].content = %q, want beta", env.Results[1]["content"]) } } @@ -169,14 +172,176 @@ func TestTavily_Info(t *testing.T) { if err != nil { t.Fatalf("Info: %v", err) } - if info.Name != "tavily" { - t.Errorf("Name = %q, want tavily", info.Name) + if info.Name != "tavily_search" { + t.Errorf("Name = %q, want tavily_search", info.Name) } if !strings.Contains(info.Desc, "Tavily") { t.Errorf("Desc = %q, want to mention Tavily", info.Desc) } } +func TestTavily_EmptyQueryReturnsEmptyResults(t *testing.T) { + t.Parallel() + + tavily := NewTavilyToolWithEnvKey(NewHTTPHelper(), func() string { return "" }) + out, err := tavily.InvokableRun(context.Background(), `{"query":""}`) + if err != nil { + t.Fatalf("InvokableRun(empty query): %v", err) + } + var envelope tavilyEnvelope + if err := json.Unmarshal([]byte(out), &envelope); err != nil { + t.Fatalf("decode empty result: %v", err) + } + if len(envelope.Results) != 0 || envelope.Error != "" { + t.Fatalf("empty query result = %#v", envelope) + } +} + +func TestTavily_PreservesRawResults(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"results":[{"url":"https://a.example","title":"A","content":"alpha","score":0.9,"custom":{"request_id":"kept"}}]}`)) + })) + defer srv.Close() + helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)}) + out, err := NewTavilyToolWith(helper).InvokableRun(context.Background(), `{"query":"x","api_key":"k"}`) + if err != nil { + t.Fatalf("InvokableRun: %v", err) + } + var envelope tavilyEnvelope + if err := json.Unmarshal([]byte(out), &envelope); err != nil { + t.Fatalf("decode response: %v", err) + } + custom, ok := envelope.Results[0]["custom"].(map[string]any) + if !ok || custom["request_id"] != "kept" { + t.Fatalf("raw result fields were lost: %#v", envelope.Results[0]) + } +} + +func TestTavily_BuildByNameUsesNodeDefaults(t *testing.T) { + t.Parallel() + + built, err := BuildByName("tavily", map[string]any{ + "api_key": "stored-key", + "search_depth": "advanced", + "max_results": float64(12), + "days": float64(7), + "include_answer": true, + "include_raw_content": true, + "include_images": true, + "include_image_descriptions": true, + "query": "ignored runtime input", + "outputs": map[string]any{"json": map[string]any{}}, + }) + if err != nil { + t.Fatalf("BuildByName: %v", err) + } + tavily, ok := built.(*TavilyTool) + if !ok { + t.Fatalf("built type = %T, want *TavilyTool", built) + } + if tavily.defaults.APIKey != "stored-key" || tavily.defaults.SearchDepth != "advanced" || tavily.defaults.MaxResults != 12 || tavily.defaults.Days != 7 { + t.Fatalf("defaults = %+v", tavily.defaults) + } + if !tavily.defaults.IncludeAnswer || !tavily.defaults.IncludeRawContent || !tavily.defaults.IncludeImages || !tavily.defaults.IncludeImageDescriptions { + t.Fatalf("boolean defaults = %+v", tavily.defaults) + } +} + +func TestTavily_ExplicitFlagsOverrideNodeDefaults(t *testing.T) { + t.Parallel() + + var gotBody map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = json.NewDecoder(r.Body).Decode(&gotBody) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"results":[]}`)) + })) + defer srv.Close() + + helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)}) + tavily := newTavilyTool(helper, func() string { return "" }, tavilyParams{ + APIKey: "stored-key", IncludeRawContent: true, IncludeImages: true, + }) + if _, err := tavily.InvokableRun(context.Background(), `{"query":"ragflow"}`); err != nil { + t.Fatalf("InvokableRun(node defaults): %v", err) + } + if gotBody["include_raw_content"] != true || gotBody["include_images"] != true { + t.Fatalf("node default flags were not sent: %#v", gotBody) + } + if _, err := tavily.InvokableRun(context.Background(), `{"query":"ragflow","include_raw_content":false,"include_images":false}`); err != nil { + t.Fatalf("InvokableRun: %v", err) + } + if gotBody["include_raw_content"] != false || gotBody["include_images"] != false { + t.Fatalf("explicit false flags were not preserved: %#v", gotBody) + } +} + +func TestTavily_BuildByNameRejectsInvalidNodeDefaults(t *testing.T) { + t.Parallel() + + cases := []map[string]any{ + {"api_key": 1}, + {"search_depth": "deep"}, + {"max_results": 0}, + {"max_results": 21}, + {"max_results": 1.5}, + {"days": 0}, + {"include_answer": "yes"}, + } + for _, params := range cases { + if _, err := BuildByName("tavily", params); err == nil { + t.Fatalf("BuildByName(%#v) succeeded, want validation error", params) + } + } +} + +func TestTavily_ComponentReferencesAndOutputs(t *testing.T) { + t.Parallel() + + tavily := NewTavilyTool() + spec := tavily.ComponentSpec() + for key, name := range map[string]string{ + "query": "Query", "topic": "Topic", "include_domains": "Include domains", "exclude_domains": "Exclude domains", + } { + field, ok := spec.InputForm[key].(map[string]any) + if !ok || field["name"] != name || field["type"] != "line" { + t.Fatalf("%s input form = %#v", key, spec.InputForm[key]) + } + } + envelope := map[string]any{"results": []any{map[string]any{ + "title": "RAGFlow", + "url": "https://ragflow.io", + "raw_content": "raw article", + "content": "fallback article", + "score": float64(0.75), + "custom": "preserved", + }}} + chunks, docAggs := tavily.BuildReferences(context.Background(), envelope) + if len(chunks) != 1 || len(docAggs) != 1 { + t.Fatalf("references = %#v / %#v", chunks, docAggs) + } + if chunks[0]["content"] != "raw article" || chunks[0]["similarity"] != float64(0.75) || docAggs[0]["doc_name"] != "RAGFlow" { + t.Fatalf("reference metadata = %#v / %#v", chunks[0], docAggs[0]) + } + outputs := tavily.BuildComponentOutputs(envelope) + results, ok := outputs["json"].([]any) + if !ok || len(results) != 1 || results[0].(map[string]any)["custom"] != "preserved" { + t.Fatalf("json output = %#v", outputs["json"]) + } + rendered, _ := outputs["formalized_content"].(string) + for _, want := range []string{"Title: RAGFlow", "URL: https://ragflow.io", "raw article"} { + if !strings.Contains(rendered, want) { + t.Fatalf("formalized_content missing %q: %q", want, rendered) + } + } + if _, exists := envelope["chunks"]; exists { + t.Fatalf("output conversion mutated envelope: %#v", envelope) + } +} + func TestTavilyExtract_BuildRequest(t *testing.T) { t.Parallel() @@ -228,7 +393,7 @@ func TestTavilyExtract_ParseResponse(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"results":[{"url":"https://a.example/","raw_content":"alpha"}]}`)) + _, _ = w.Write([]byte(`{"results":[{"url":"https://a.example/","raw_content":"alpha","custom":"preserved"}]}`)) })) defer srv.Close() @@ -249,9 +414,12 @@ func TestTavilyExtract_ParseResponse(t *testing.T) { if len(env.Results) != 1 { t.Fatalf("Results len = %d, want 1", len(env.Results)) } - if env.Results[0].URL != "https://a.example/" || env.Results[0].RawContent != "alpha" { + if env.Results[0]["url"] != "https://a.example/" || env.Results[0]["raw_content"] != "alpha" { t.Errorf("Results[0] = %+v, want url and raw_content", env.Results[0]) } + if env.Results[0]["custom"] != "preserved" { + t.Fatalf("raw upstream fields were lost: %#v", env.Results[0]) + } } func TestTavilyExtract_RequiresAPIKey(t *testing.T) { @@ -282,3 +450,138 @@ func TestTavilyExtract_Info(t *testing.T) { t.Errorf("Desc = %q, want to mention Tavily Extract", info.Desc) } } + +func TestTavilyExtract_ComponentContract(t *testing.T) { + t.Parallel() + + tavily := NewTavilyExtractTool() + spec := tavily.ComponentSpec() + for _, input := range []string{"urls", "extract_depth", "format"} { + if _, ok := spec.Inputs[input]; !ok { + t.Fatalf("component inputs missing %s: %#v", input, spec.Inputs) + } + } + if _, ok := spec.Outputs["json"]; !ok { + t.Fatalf("component outputs missing json: %#v", spec.Outputs) + } + for key, name := range map[string]string{ + "urls": "URLs", "extract_depth": "Extract depth", "format": "Format", + } { + field, ok := spec.InputForm[key].(map[string]any) + if !ok || field["name"] != name || field["type"] != "line" { + t.Fatalf("%s input form = %#v", key, spec.InputForm[key]) + } + } + + envelope := map[string]any{ + "results": []any{map[string]any{ + "url": "https://example.com", + "raw_content": "content", + "custom": "preserved", + }}, + } + outputs := tavily.BuildComponentOutputs(envelope) + results, ok := outputs["json"].([]any) + if !ok || len(results) != 1 || results[0].(map[string]any)["custom"] != "preserved" { + t.Fatalf("component outputs = %#v", outputs) + } + if _, exists := envelope["chunks"]; exists { + t.Fatalf("component output conversion mutated envelope: %#v", envelope) + } +} + +func TestTavilyExtract_BuildByNameAcceptsNodeDefaults(t *testing.T) { + t.Parallel() + + built, err := BuildByName("tavily_extract", map[string]any{ + "api_key": "stored-key", + "urls": "https://example.com", + "extract_depth": "advanced", + "format": "text", + "outputs": map[string]any{"json": map[string]any{}}, + }) + if err != nil { + t.Fatalf("BuildByName: %v", err) + } + tavily, ok := built.(*TavilyExtractTool) + if !ok { + t.Fatalf("built type = %T, want *TavilyExtractTool", built) + } + if tavily.defaults.APIKey != "stored-key" || tavily.defaults.ExtractDepth != "advanced" || tavily.defaults.Format != "text" { + t.Fatalf("defaults = %+v", tavily.defaults) + } + if tavily.defaults.URLs != "https://example.com" { + t.Fatalf("defaults.URLs = %#v", tavily.defaults.URLs) + } +} + +func TestTavilyExtract_UsesNodeDefaults(t *testing.T) { + t.Parallel() + + var gotAuth string + var gotBody map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + _ = json.NewDecoder(r.Body).Decode(&gotBody) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"results":[]}`)) + })) + defer srv.Close() + + helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)}) + tavily := newTavilyExtractTool(helper, func() string { return "" }, tavilyExtractParams{ + APIKey: "stored-key", + URLs: "https://stored.example", + ExtractDepth: "advanced", + Format: "text", + }) + if _, err := tavily.InvokableRun(context.Background(), `{"unrelated":"ignored"}`); err != nil { + t.Fatalf("InvokableRun: %v", err) + } + if gotAuth != "Bearer stored-key" { + t.Fatalf("Authorization = %q, want stored node API key", gotAuth) + } + urls, ok := gotBody["urls"].([]any) + if !ok || len(urls) != 1 || urls[0] != "https://stored.example" { + t.Fatalf("body.urls = %#v", gotBody["urls"]) + } + if gotBody["extract_depth"] != "advanced" || gotBody["format"] != "text" { + t.Fatalf("request body = %#v", gotBody) + } +} + +func TestTavilyExtract_BuildByNameRejectsInvalidNodeDefaults(t *testing.T) { + t.Parallel() + + cases := []map[string]any{ + {"api_key": 1}, + {"extract_depth": "deep"}, + {"format": "html"}, + } + for _, params := range cases { + if _, err := BuildByName("tavily_extract", params); err == nil { + t.Fatalf("BuildByName(%#v) succeeded, want validation error", params) + } + } +} + +func TestTavilyExtract_MergeDefaults(t *testing.T) { + t.Parallel() + + got := mergeTavilyExtractParams( + tavilyExtractParams{ + APIKey: "stored-key", + URLs: []string{"https://stored.example"}, + ExtractDepth: "advanced", + Format: "text", + }, + tavilyExtractParams{URLs: []string{"https://runtime.example"}}, + ) + if got.APIKey != "stored-key" || got.ExtractDepth != "advanced" || got.Format != "text" { + t.Fatalf("merged params = %+v", got) + } + urls, ok := got.URLs.([]string) + if !ok || len(urls) != 1 || urls[0] != "https://runtime.example" { + t.Fatalf("merged URLs = %#v", got.URLs) + } +} diff --git a/internal/agent/tool/tool2component.go b/internal/agent/tool/tool2component.go new file mode 100644 index 0000000000..df487921d7 --- /dev/null +++ b/internal/agent/tool/tool2component.go @@ -0,0 +1,68 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package tool + +import ( + "context" + + einotool "github.com/cloudwego/eino/components/tool" +) + +// ToolInvoker is the invocation seam shared by Eino tools and Canvas +// components backed by those tools. +type ToolInvoker interface { + InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error) +} + +// ComponentSpec describes the Canvas-facing surface of a tool. It is kept +// separate from Info(), whose schema contains only model-emitted arguments. +type ComponentSpec struct { + Inputs map[string]string + Outputs map[string]string + InputForm map[string]any +} + +// ToolComponent is the required Canvas adaptation contract implemented by a +// tool that can back a Canvas component. +type ToolComponent interface { + ToolInvoker + ComponentSpec() ComponentSpec + // BuildComponentOutputs converts the complete decoded tool envelope into + // the component's public Canvas outputs. + BuildComponentOutputs(envelope map[string]any) map[string]any +} + +// ReferenceBuilder is an optional capability for tools that add retrieval +// references to Canvas state. +type ReferenceBuilder interface { + BuildReferences(ctx context.Context, envelope map[string]any) (chunks []map[string]any, docAggs []map[string]any) +} + +func envelopeSlice(envelope map[string]any, key string) []any { + switch values := envelope[key].(type) { + case []any: + return values + case []map[string]any: + result := make([]any, 0, len(values)) + for _, value := range values { + result = append(result, value) + } + return result + default: + return []any{} + } +} diff --git a/internal/agent/tool/wencai.go b/internal/agent/tool/wencai.go index dcdc5708f7..b28a3f9d0e 100644 --- a/internal/agent/tool/wencai.go +++ b/internal/agent/tool/wencai.go @@ -73,6 +73,8 @@ type WencaiTool struct { defaults wencaiParams } +var _ ToolComponent = (*WencaiTool)(nil) + func NewWencaiTool() *WencaiTool { return newWencaiTool(wencaiParams{}) } @@ -139,6 +141,31 @@ func mergeWencaiParams(defaults, params wencaiParams) wencaiParams { return params } +// ComponentSpec returns the Python-compatible WenCai Canvas surface. +func (w *WencaiTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "query": "The question/conditions to select stocks.", + }, + Outputs: map[string]string{ + "report": "WenCai query report.", + }, + InputForm: map[string]any{ + "query": map[string]any{ + "name": "Query", + "type": "line", + }, + }, + } +} + +// BuildComponentOutputs converts WenCai's complete tool envelope into its +// public Canvas outputs. +func (w *WencaiTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + report, _ := envelope["report"].(string) + return map[string]any{"report": report} +} + func isWencaiQueryTypeSupported(queryType string) bool { _, ok := wencaiQueryTypes[queryType] return ok diff --git a/internal/agent/tool/wencai_test.go b/internal/agent/tool/wencai_test.go index a105a18fb3..bc326baa80 100644 --- a/internal/agent/tool/wencai_test.go +++ b/internal/agent/tool/wencai_test.go @@ -171,7 +171,6 @@ func TestWencai_BuildByNameRejectsInvalidNodeParams(t *testing.T) { {name: "string top_n", params: map[string]any{"top_n": "10"}}, {name: "invalid query_type", params: map[string]any{"query_type": "crypto"}}, {name: "non-string query_type", params: map[string]any{"query_type": 1}}, - {name: "unknown key", params: map[string]any{"page": 1}}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { @@ -183,6 +182,47 @@ func TestWencai_BuildByNameRejectsInvalidNodeParams(t *testing.T) { } } +func TestWencai_BuildByNameIgnoresUnrelatedCanvasParams(t *testing.T) { + t.Parallel() + + built, err := BuildByName("wencai", map[string]any{ + "top_n": float64(20), + "outputs": map[string]any{"report": map[string]any{}}, + "setups": map[string]any{"query": "configured query"}, + }) + if err != nil { + t.Fatalf("BuildByName: %v", err) + } + wencai := built.(*WencaiTool) + if wencai.defaults.TopN != 20 { + t.Fatalf("defaults.TopN = %d, want 20", wencai.defaults.TopN) + } +} + +func TestWencai_ComponentContract(t *testing.T) { + t.Parallel() + + wencai := NewWencaiTool() + spec := wencai.ComponentSpec() + if _, ok := spec.Inputs["query"]; !ok { + t.Fatalf("component inputs missing query: %#v", spec.Inputs) + } + if _, ok := spec.Outputs["report"]; !ok { + t.Fatalf("component outputs missing report: %#v", spec.Outputs) + } + query, ok := spec.InputForm["query"].(map[string]any) + if !ok || query["name"] != "Query" || query["type"] != "line" { + t.Fatalf("query input form = %#v", spec.InputForm["query"]) + } + outputs := wencai.BuildComponentOutputs(map[string]any{ + "report": "market report", + "tool_metadata": map[string]any{"request_id": "request-1"}, + }) + if outputs["report"] != "market report" { + t.Fatalf("component outputs = %#v", outputs) + } +} + func TestWencai_MergeDefaults(t *testing.T) { t.Parallel() diff --git a/internal/agent/tool/wikipedia.go b/internal/agent/tool/wikipedia.go index 05705a44ba..20651d0633 100644 --- a/internal/agent/tool/wikipedia.go +++ b/internal/agent/tool/wikipedia.go @@ -18,15 +18,22 @@ package tool import ( "context" + "crypto/sha1" "encoding/json" "fmt" + "math/big" "net/http" "net/url" + "regexp" "sort" + "strconv" "strings" + "unicode/utf8" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" + + "ragflow/internal/tokenizer" ) const wikipediaToolName = "wikipedia_search" @@ -38,8 +45,13 @@ const wikipediaUserAgent = "Mozilla/5.0 (compatible; ragflow/1.0; +https://githu const ( defaultWikipediaTopN = 10 defaultWikipediaLanguage = "en" + wikipediaPromptMaxTokens = 200000 ) +var wikipediaDataImagePattern = regexp.MustCompile(`!?\[[a-z]+\]\(data:image/png;base64,[ 0-9A-Za-z/_=+\-]+\)`) + +var wikipediaNewlinePattern = regexp.MustCompile(`\n+`) + var wikipediaLanguages = map[string]struct{}{ "af": {}, "pl": {}, "ar": {}, "ast": {}, "az": {}, "bg": {}, "nan": {}, "bn": {}, "be": {}, "ca": {}, "cs": {}, "cy": {}, "da": {}, "de": {}, "et": {}, "el": {}, "en": {}, "es": {}, "eo": {}, "eu": {}, @@ -91,7 +103,7 @@ type wikipediaResponse struct { // is the canvas-facing equivalent of Python ToolBase._retrieve_chunks(). type wikipediaEnvelope struct { FormalizedContent string `json:"formalized_content,omitempty"` - Results []wikipediaResult `json:"results,omitempty"` + Results []wikipediaResult `json:"results"` Error string `json:"_ERROR,omitempty"` } @@ -105,6 +117,9 @@ type WikipediaTool struct { lang string } +var _ ToolComponent = (*WikipediaTool)(nil) +var _ ReferenceBuilder = (*WikipediaTool)(nil) + // NewWikipediaTool returns a WikipediaTool using the default HTTPHelper. func NewWikipediaTool() *WikipediaTool { return NewWikipediaToolWith(NewHTTPHelper()) @@ -146,6 +161,21 @@ func (w *WikipediaTool) Info(_ context.Context) (*schema.ToolInfo, error) { }, nil } +func (w *WikipediaTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "query": "The search keyword to execute with wikipedia. The keyword MUST be a specific subject that can match the title.", + }, + Outputs: map[string]string{ + "formalized_content": "Rendered Wikipedia references for downstream prompts.", + "json": "Wikipedia result list.", + }, + InputForm: map[string]any{ + "query": map[string]any{"name": "Query", "type": "line"}, + }, + } +} + // buildWikipediaURL constructs a MediaWiki generator=search URL that returns // the same page fields Python reads from wikipedia.page(): title, url, // and summary-like introductory extract. @@ -172,8 +202,7 @@ func (w *WikipediaTool) InvokableRun(ctx context.Context, argsJSON string, _ ... fmt.Errorf("wikipedia: parse arguments: %w", err) } if p.Query == "" { - return wikipediaErrJSON(fmt.Errorf("query is required")), - fmt.Errorf("wikipedia: query is required") + return wikipediaJSON(wikipediaEnvelope{Results: []wikipediaResult{}}), nil } lang := p.Lang @@ -254,6 +283,98 @@ func (w *WikipediaTool) InvokableRun(ctx context.Context, argsJSON string, _ ... }), nil } +func (w *WikipediaTool) BuildReferences(_ context.Context, envelope map[string]any) ([]map[string]any, []map[string]any) { + return buildWikipediaReferences(envelope) +} + +func (w *WikipediaTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + results := envelopeSlice(envelope, "results") + chunks, _ := buildWikipediaReferences(envelope) + return map[string]any{ + "formalized_content": renderWikipediaReferences(chunks, wikipediaPromptMaxTokens), + "json": results, + } +} + +func buildWikipediaReferences(envelope map[string]any) ([]map[string]any, []map[string]any) { + results := envelopeSlice(envelope, "results") + chunks := make([]map[string]any, 0, len(results)) + docAggs := make([]map[string]any, 0, len(results)) + for _, result := range results { + item, ok := result.(map[string]any) + if !ok { + continue + } + content := wikipediaDataImagePattern.ReplaceAllString(wikipediaText(item["content"]), "") + content = truncateWikipediaRunes(content, 10000) + if content == "" { + continue + } + documentID := strconv.FormatInt(wikipediaHashInt(content, 100000000), 10) + displayID := strconv.FormatInt(wikipediaHashInt(documentID, 500), 10) + title := wikipediaText(item["title"]) + resultURL := wikipediaText(item["url"]) + chunks = append(chunks, map[string]any{ + "id": displayID, + "chunk_id": documentID, + "content": content, + "doc_id": documentID, + "document_id": documentID, + "docnm_kwd": title, + "document_name": title, + "similarity": 1, + "score": 1, + "url": resultURL, + }) + docAggs = append(docAggs, map[string]any{"doc_name": title, "doc_id": documentID, "count": 1, "url": resultURL}) + } + return chunks, docAggs +} + +func renderWikipediaReferences(chunks []map[string]any, maxTokens int) string { + usedTokens := 0 + blocks := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + content := wikipediaText(chunk["content"]) + block := strings.Join([]string{ + "\nID: " + wikipediaText(chunk["id"]), + "├── Title: " + wikipediaNewlinePattern.ReplaceAllString(wikipediaText(chunk["document_name"]), " "), + "├── URL: " + wikipediaNewlinePattern.ReplaceAllString(wikipediaText(chunk["url"]), " "), + "└── Content:\n" + content, + }, "\n") + blockTokens := tokenizer.NumTokensFromString(block) + if maxTokens > 0 && float64(usedTokens+blockTokens) > float64(maxTokens)*0.97 { + break + } + usedTokens += blockTokens + blocks = append(blocks, block) + } + return strings.Join(blocks, "\n") +} + +func wikipediaText(value any) string { + if value == nil { + return "" + } + if text, ok := value.(string); ok { + return text + } + return fmt.Sprint(value) +} + +func wikipediaHashInt(value string, modulus int64) int64 { + digest := sha1.Sum([]byte(value)) + number := new(big.Int).SetBytes(digest[:]) + return new(big.Int).Mod(number, big.NewInt(modulus)).Int64() +} + +func truncateWikipediaRunes(value string, limit int) string { + if utf8.RuneCountInString(value) <= limit { + return value + } + return string([]rune(value)[:limit]) +} + func renderWikipediaResults(results []wikipediaResult) string { if len(results) == 0 { return "" diff --git a/internal/agent/tool/wikipedia_test.go b/internal/agent/tool/wikipedia_test.go index 862fb571e5..e29cf03100 100644 --- a/internal/agent/tool/wikipedia_test.go +++ b/internal/agent/tool/wikipedia_test.go @@ -24,6 +24,8 @@ import ( "net/url" "strings" "testing" + + "ragflow/internal/tokenizer" ) func TestWikipedia_BuildURL(t *testing.T) { @@ -191,15 +193,86 @@ func TestWikipedia_Info(t *testing.T) { } } -func TestWikipedia_RequiresQuery(t *testing.T) { +func TestWikipedia_EmptyQuery(t *testing.T) { t.Parallel() tool := NewWikipediaTool() - _, err := tool.InvokableRun(context.Background(), `{"query":""}`) - if err == nil { - t.Fatal("expected error for empty query") + out, err := tool.InvokableRun(context.Background(), `{"query":""}`) + if err != nil { + t.Fatalf("InvokableRun(empty): %v", err) } - if !strings.Contains(err.Error(), "query") { - t.Errorf("err = %v, want to mention query", err) + var envelope wikipediaEnvelope + if err := json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 { + t.Fatalf("empty result = %s / %v", out, err) + } + if !strings.Contains(out, `"results":[]`) { + t.Fatalf("empty result omitted results key: %s", out) + } +} + +func TestWikipedia_ComponentReferencesAndOutputs(t *testing.T) { + t.Parallel() + + wikipedia := NewWikipediaTool() + spec := wikipedia.ComponentSpec() + if query, ok := spec.InputForm["query"].(map[string]any); !ok || query["name"] != "Query" || query["type"] != "line" { + t.Fatalf("query input form = %#v", spec.InputForm["query"]) + } + envelope := map[string]any{"results": []any{map[string]any{ + "title": "RAG", + "url": "https://en.wikipedia.org/wiki/RAG", + "content": "RAG is an acronym.", + }}} + chunks, docAggs := wikipedia.BuildReferences(context.Background(), envelope) + if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["document_name"] != "RAG" || chunks[0]["similarity"] != 1 { + t.Fatalf("references = %#v / %#v", chunks, docAggs) + } + outputs := wikipedia.BuildComponentOutputs(envelope) + if results, ok := outputs["json"].([]any); !ok || len(results) != 1 { + t.Fatalf("json output = %#v", outputs["json"]) + } + if !strings.Contains(outputs["formalized_content"].(string), "RAG is an acronym.") { + t.Fatalf("formalized_content = %q", outputs["formalized_content"]) + } + if _, exists := envelope["chunks"]; exists { + t.Fatalf("output conversion mutated envelope: %#v", envelope) + } +} + +func TestRenderWikipediaReferencesStopsBeforeOverBudgetBlock(t *testing.T) { + t.Parallel() + + chunks := []map[string]any{ + {"id": "1", "document_name": "First", "url": "https://first.example", "content": "first reference content"}, + {"id": "2", "document_name": "Second", "url": "https://second.example", "content": "second reference content"}, + } + firstBlock := renderWikipediaReferences(chunks[:1], 0) + firstTokens := tokenizer.NumTokensFromString(firstBlock) + maxTokens := (firstTokens*100 + 96) / 97 + if got := renderWikipediaReferences(chunks, maxTokens); got != firstBlock { + t.Fatalf("rendered = %q, want only first block %q", got, firstBlock) + } + if got := renderWikipediaReferences(chunks, 1); got != "" { + t.Fatalf("over-budget first block was appended: %q", got) + } + if got := renderWikipediaReferences(chunks, 0); !strings.Contains(got, "Title: First") || !strings.Contains(got, "Title: Second") { + t.Fatalf("unlimited rendering dropped blocks: %q", got) + } +} + +func TestWikipedia_BuildByNameIgnoresCanvasParams(t *testing.T) { + t.Parallel() + + built, err := BuildByName("wikipedia", map[string]any{ + "top_n": float64(3), + "language": "en", + "outputs": map[string]any{"json": map[string]any{}}, + }) + if err != nil { + t.Fatalf("BuildByName: %v", err) + } + wikipedia := built.(*WikipediaTool) + if wikipedia.topN != 3 || wikipedia.lang != "en" { + t.Fatalf("node defaults = %d/%q", wikipedia.topN, wikipedia.lang) } } diff --git a/internal/agent/tool/yahoo_finance.go b/internal/agent/tool/yahoo_finance.go index 98920cc696..c6aa219701 100644 --- a/internal/agent/tool/yahoo_finance.go +++ b/internal/agent/tool/yahoo_finance.go @@ -22,6 +22,7 @@ import ( "fmt" "net/http" "net/url" + "sort" "strings" "github.com/cloudwego/eino/components/tool" @@ -30,35 +31,30 @@ import ( const yahooFinanceToolName = "yahoo_finance" -const yahooFinanceToolDescription = "Fetch stock quote snapshots from Yahoo Finance. Returns quoteResponse.result[].{symbol, regularMarketPrice, currency, regularMarketChangePercent}." +const yahooFinanceToolDescription = "The Yahoo Finance service provides access to real-time and historical stock market data, company profiles, and financial news." -// yahooFinanceParams is the JSON shape the model sends into InvokableRun. +const yahooFinanceStockCodeDescription = "The stock code or company name." + +// yahooFinanceParams contains the model-emitted stock code and the supported +// Canvas-side information switch. Info exposes only StockCode. type yahooFinanceParams struct { - Symbols []string `json:"symbols"` - Fields []string `json:"fields"` -} - -// yahooFinanceQuote is one element of the upstream result array. -type yahooFinanceQuote struct { - Symbol string `json:"symbol"` - RegularMarketPrice float64 `json:"regularMarketPrice"` - Currency string `json:"currency"` - RegularMarketChangePercent float64 `json:"regularMarketChangePercent"` + StockCode string `json:"stock_code"` + Info bool `json:"info"` } // yahooFinanceResponse is the upstream Yahoo Finance /v7/finance/quote // envelope. type yahooFinanceResponse struct { QuoteResponse struct { - Result []yahooFinanceQuote `json:"result"` - Error any `json:"error,omitempty"` + Result []map[string]any `json:"result"` + Error any `json:"error,omitempty"` } `json:"quoteResponse"` } -// yahooFinanceEnvelope is what the model sees. +// yahooFinanceEnvelope is the tool-to-component transport shape. type yahooFinanceEnvelope struct { - Results []yahooFinanceQuote `json:"results"` - Error string `json:"_ERROR,omitempty"` + Report string `json:"report"` + Error string `json:"_ERROR,omitempty"` } // yahooFinanceEndpoint is the Yahoo Finance quote URL. Exposed as a @@ -70,22 +66,37 @@ var yahooFinanceEndpoint = "https://query1.finance.yahoo.com/v7/finance/quote" // It performs an unauthenticated GET against the public quote API // via the shared HTTPHelper and returns the parsed quote records. type YahooFinanceTool struct { - helper *HTTPHelper + helper *HTTPHelper + defaults yahooFinanceParams } +var _ ToolComponent = (*YahooFinanceTool)(nil) + // NewYahooFinanceTool returns a YahooFinanceTool using the default // HTTPHelper. func NewYahooFinanceTool() *YahooFinanceTool { - return NewYahooFinanceToolWith(NewHTTPHelper()) + return NewYahooFinanceToolWithDefaults(nil, defaultYahooFinanceParams()) } // NewYahooFinanceToolWith returns a YahooFinanceTool that uses the // provided HTTPHelper. Useful for tests. func NewYahooFinanceToolWith(h *HTTPHelper) *YahooFinanceTool { + return NewYahooFinanceToolWithDefaults(h, defaultYahooFinanceParams()) +} + +// NewYahooFinanceToolWithDefaults returns a YahooFinanceTool with Canvas-side +// defaults. This follows the same constructor pattern as GitHubTool. +func NewYahooFinanceToolWithDefaults(h *HTTPHelper, defaults yahooFinanceParams) *YahooFinanceTool { if h == nil { - h = NewHTTPHelper() + // ToolParamBase defaults to max_retries=0, so the tool itself performs + // one request unless a caller injects a differently configured helper. + h = NewHTTPHelperWithRetry(RetryConfig{MaxAttempts: 1}) } - return &YahooFinanceTool{helper: h} + return &YahooFinanceTool{helper: h, defaults: defaults} +} + +func defaultYahooFinanceParams() yahooFinanceParams { + return yahooFinanceParams{Info: true} } // Info returns the tool's metadata for the chat model. @@ -94,28 +105,32 @@ func (y *YahooFinanceTool) Info(_ context.Context) (*schema.ToolInfo, error) { Name: yahooFinanceToolName, Desc: yahooFinanceToolDescription, ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ - "symbols": { - Type: schema.Array, - Desc: "Stock symbols to look up (e.g. AAPL, MSFT, 0005.HK).", + "stock_code": { + Type: schema.String, + Desc: yahooFinanceStockCodeDescription, Required: true, }, - "fields": { - Type: schema.Array, - Desc: "Optional list of fields to request via the `fields` query parameter.", - Required: false, - }, }), }, nil } -// buildYahooFinanceURL composes the quote URL with the symbol list -// and an optional `fields` parameter. Centralized for testability. -func buildYahooFinanceURL(symbols []string, fields []string) string { - q := url.Values{} - q.Set("symbols", strings.Join(symbols, ",")) - if len(fields) > 0 { - q.Set("fields", strings.Join(fields, ",")) +func (y *YahooFinanceTool) ComponentSpec() ComponentSpec { + return ComponentSpec{ + Inputs: map[string]string{ + "stock_code": yahooFinanceStockCodeDescription, + }, + Outputs: map[string]string{"report": "Yahoo Finance data formatted as Markdown."}, + InputForm: map[string]any{ + "stock_code": map[string]any{"type": "line", "name": "Stock code/Company name"}, + }, } +} + +// buildYahooFinanceURL composes the quote URL for one Python-compatible +// stock_code input. Centralized for testability. +func buildYahooFinanceURL(stockCode string) string { + q := url.Values{} + q.Set("symbols", stockCode) return yahooFinanceEndpoint + "?" + q.Encode() } @@ -126,12 +141,13 @@ func (y *YahooFinanceTool) InvokableRun(ctx context.Context, argsJSON string, _ return yahooFinanceErrJSON(fmt.Errorf("yahoo_finance: parse arguments: %w", err)), fmt.Errorf("yahoo_finance: parse arguments: %w", err) } - if len(p.Symbols) == 0 { - return yahooFinanceErrJSON(fmt.Errorf("symbols is required and must be non-empty")), - fmt.Errorf("yahoo_finance: symbols is required and must be non-empty") + p = mergeYahooFinanceParams(y.defaults, p) + p.StockCode = strings.TrimSpace(p.StockCode) + if p.StockCode == "" || !p.Info { + return yahooFinanceJSON(yahooFinanceEnvelope{Report: ""}), nil } - endpoint := buildYahooFinanceURL(p.Symbols, p.Fields) + endpoint := buildYahooFinanceURL(p.StockCode) // Yahoo Finance returns 401 unless we send a User-Agent that // looks like a real browser. curl-style UA is the conventional // workaround for the public (unauthenticated) endpoint. @@ -156,7 +172,62 @@ func (y *YahooFinanceTool) InvokableRun(ctx context.Context, argsJSON string, _ return yahooFinanceErrJSON(fmt.Errorf("yahoo_finance: decode response: %w", err)), fmt.Errorf("yahoo_finance: decode response: %w", err) } - return yahooFinanceJSON(yahooFinanceEnvelope{Results: raw.QuoteResponse.Result}), nil + if raw.QuoteResponse.Error != nil { + err := fmt.Errorf("yahoo_finance: upstream error: %v", raw.QuoteResponse.Error) + return yahooFinanceErrJSON(err), err + } + return yahooFinanceJSON(yahooFinanceEnvelope{Report: renderYahooFinanceReport(raw.QuoteResponse.Result)}), nil +} + +func (y *YahooFinanceTool) BuildComponentOutputs(envelope map[string]any) map[string]any { + report, _ := envelope["report"].(string) + return map[string]any{"report": report} +} + +func mergeYahooFinanceParams(defaults, params yahooFinanceParams) yahooFinanceParams { + params.Info = defaults.Info + return params +} + +// renderYahooFinanceReport keeps the existing public quote endpoint while +// returning the string report expected by the Canvas component. +func renderYahooFinanceReport(quotes []map[string]any) string { + if len(quotes) == 0 { + return "" + } + sections := make([]string, 0, len(quotes)) + for _, quote := range quotes { + keys := make([]string, 0, len(quote)) + for key := range quote { + keys = append(keys, key) + } + sort.Strings(keys) + + rows := []string{"# Information:", "| | 0 |", "|:---|:---|"} + for _, key := range keys { + rows = append(rows, fmt.Sprintf("| %s | %s |", markdownCell(key), markdownCell(yahooFinanceValue(quote[key])))) + } + sections = append(sections, strings.Join(rows, "\n")) + } + return strings.Join(sections, "\n\n") +} + +func yahooFinanceValue(value any) string { + if value == nil { + return "None" + } + if text, ok := value.(string); ok { + return text + } + if encoded, err := json.Marshal(value); err == nil { + return string(encoded) + } + return fmt.Sprint(value) +} + +func markdownCell(value string) string { + value = strings.ReplaceAll(value, "|", `\|`) + return strings.ReplaceAll(value, "\n", "
") } func yahooFinanceJSON(env yahooFinanceEnvelope) string { diff --git a/internal/agent/tool/yahoo_finance_test.go b/internal/agent/tool/yahoo_finance_test.go index 595909e9cc..b24dd64726 100644 --- a/internal/agent/tool/yahoo_finance_test.go +++ b/internal/agent/tool/yahoo_finance_test.go @@ -26,135 +26,282 @@ import ( "testing" ) -func TestYahooFinance_BuildURL(t *testing.T) { +func TestYahooFinanceBuildURL(t *testing.T) { t.Parallel() - cases := []struct { - name string - symbols []string - fields []string - wantSymbols string - wantFields string - wantHost string - }{ - { - name: "single symbol, no fields", - symbols: []string{"AAPL"}, - fields: nil, - wantSymbols: "AAPL", - wantHost: "query1.finance.yahoo.com", - }, - { - name: "multi symbol, with fields", - symbols: []string{"AAPL", "MSFT", "0005.HK"}, - fields: []string{"symbol", "regularMarketPrice"}, - wantSymbols: "AAPL,MSFT,0005.HK", - wantFields: "symbol,regularMarketPrice", - wantHost: "query1.finance.yahoo.com", - }, + got := buildYahooFinanceURL("0005.HK") + parsed, err := url.Parse(got) + if err != nil { + t.Fatalf("url.Parse(%q): %v", got, err) } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - got := buildYahooFinanceURL(tc.symbols, tc.fields) - u, err := url.Parse(got) - if err != nil { - t.Fatalf("url.Parse(%q): %v", got, err) + if parsed.Host != "query1.finance.yahoo.com" { + t.Fatalf("host = %q", parsed.Host) + } + if parsed.Path != "/v7/finance/quote" { + t.Fatalf("path = %q", parsed.Path) + } + if symbols := parsed.Query().Get("symbols"); symbols != "0005.HK" { + t.Fatalf("symbols = %q", symbols) + } + if _, exists := parsed.Query()["fields"]; exists { + t.Fatalf("unexpected legacy fields query: %s", parsed.RawQuery) + } +} + +func TestYahooFinanceInvokableRunBuildsMarkdownReport(t *testing.T) { + t.Parallel() + + var gotSymbols, gotUserAgent string + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + gotSymbols = request.URL.Query().Get("symbols") + gotUserAgent = request.Header.Get("User-Agent") + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{ + "quoteResponse": { + "result": [{ + "symbol":"AAPL", + "regularMarketPrice":189.5, + "currency":"USD", + "marketState":"REGULAR", + "note":"left|right" + }], + "error": null } - if u.Host != tc.wantHost { - t.Errorf("host = %q, want %q", u.Host, tc.wantHost) + }`)) + })) + defer server.Close() + + helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)}) + raw, err := NewYahooFinanceToolWith(helper).InvokableRun(context.Background(), `{"stock_code":" AAPL "}`) + if err != nil { + t.Fatalf("InvokableRun: %v", err) + } + if gotSymbols != "AAPL" { + t.Fatalf("symbols = %q", gotSymbols) + } + if !strings.Contains(gotUserAgent, "ragflow") { + t.Fatalf("User-Agent = %q", gotUserAgent) + } + + var envelope yahooFinanceEnvelope + if err := json.Unmarshal([]byte(raw), &envelope); err != nil { + t.Fatalf("unmarshal output %q: %v", raw, err) + } + for _, expected := range []string{ + "# Information:", + "| currency | USD |", + "| marketState | REGULAR |", + `| note | left\|right |`, + "| regularMarketPrice | 189.5 |", + "| symbol | AAPL |", + } { + if !strings.Contains(envelope.Report, expected) { + t.Fatalf("report missing %q:\n%s", expected, envelope.Report) + } + } + if envelope.Error != "" { + t.Fatalf("unexpected error = %q", envelope.Error) + } +} + +func TestYahooFinanceEmptyStockCodeSkipsRequest(t *testing.T) { + t.Parallel() + + calls := 0 + server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + calls++ + })) + defer server.Close() + helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)}) + + for _, args := range []string{`{"stock_code":""}`, `{"stock_code":" "}`} { + raw, err := NewYahooFinanceToolWith(helper).InvokableRun(context.Background(), args) + if err != nil { + t.Fatalf("InvokableRun(%s): %v", args, err) + } + var envelope yahooFinanceEnvelope + if err := json.Unmarshal([]byte(raw), &envelope); err != nil { + t.Fatalf("unmarshal output: %v", err) + } + if envelope.Report != "" || envelope.Error != "" { + t.Fatalf("empty input output = %#v", envelope) + } + } + if calls != 0 { + t.Fatalf("server calls = %d, want 0", calls) + } +} + +func TestYahooFinanceAllSectionsDisabledSkipsRequest(t *testing.T) { + t.Parallel() + + calls := 0 + server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + calls++ + })) + defer server.Close() + helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)}) + yahoo := NewYahooFinanceToolWithDefaults(helper, yahooFinanceParams{}) + + raw, err := yahoo.InvokableRun(context.Background(), `{"stock_code":"AAPL"}`) + if err != nil { + t.Fatalf("InvokableRun: %v", err) + } + var envelope yahooFinanceEnvelope + if err := json.Unmarshal([]byte(raw), &envelope); err != nil { + t.Fatalf("unmarshal output: %v", err) + } + if envelope.Report != "" || calls != 0 { + t.Fatalf("output = %#v, server calls = %d", envelope, calls) + } +} + +func TestMergeYahooFinanceParamsKeepsStockCodeAndUsesNodeConfig(t *testing.T) { + t.Parallel() + + defaults := yahooFinanceParams{Info: true} + params := yahooFinanceParams{ + StockCode: "AAPL", + Info: false, + } + + got := mergeYahooFinanceParams(defaults, params) + want := defaults + want.StockCode = "AAPL" + if got != want { + t.Fatalf("merged params = %#v, want %#v", got, want) + } +} + +func TestYahooFinanceErrorsReturnEnvelope(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + statusCode int + body string + wantError string + }{ + {name: "http status", statusCode: http.StatusUnauthorized, body: `denied`, wantError: "upstream returned 401"}, + {name: "invalid json", statusCode: http.StatusOK, body: `{`, wantError: "decode response"}, + {name: "upstream envelope", statusCode: http.StatusOK, body: `{"quoteResponse":{"result":[],"error":{"code":"Not Found"}}}`, wantError: "upstream error"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.WriteHeader(test.statusCode) + _, _ = writer.Write([]byte(test.body)) + })) + defer server.Close() + helper := NewHTTPHelperWithRetry(RetryConfig{MaxAttempts: 1}).WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)}) + + raw, err := NewYahooFinanceToolWith(helper).InvokableRun(context.Background(), `{"stock_code":"AAPL"}`) + if err == nil || !strings.Contains(err.Error(), test.wantError) { + t.Fatalf("err = %v, want %q", err, test.wantError) } - if u.Path != "/v7/finance/quote" { - t.Errorf("path = %q, want /v7/finance/quote", u.Path) + var envelope yahooFinanceEnvelope + if decodeErr := json.Unmarshal([]byte(raw), &envelope); decodeErr != nil { + t.Fatalf("error result is not JSON: %s: %v", raw, decodeErr) } - q := u.Query() - if q.Get("symbols") != tc.wantSymbols { - t.Errorf("symbols = %q, want %q", q.Get("symbols"), tc.wantSymbols) - } - if tc.wantFields != "" && q.Get("fields") != tc.wantFields { - t.Errorf("fields = %q, want %q", q.Get("fields"), tc.wantFields) + if !strings.Contains(envelope.Error, test.wantError) || envelope.Report != "" { + t.Fatalf("envelope = %#v", envelope) } }) } } -func TestYahooFinance_ParseQuote(t *testing.T) { +func TestYahooFinanceMalformedArguments(t *testing.T) { t.Parallel() - var gotUA string - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - gotUA = r.Header.Get("User-Agent") - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{ - "quoteResponse": { - "result": [ - {"symbol":"AAPL","regularMarketPrice":189.5,"currency":"USD","regularMarketChangePercent":1.23}, - {"symbol":"MSFT","regularMarketPrice":421.0,"currency":"USD","regularMarketChangePercent":-0.5} - ] - } - }`)) - })) - defer srv.Close() - - helper := NewHTTPHelper().WithClient(&http.Client{ - Transport: rewriteHostTransport(srv.URL), - }) - tool := NewYahooFinanceToolWith(helper) - out, err := tool.InvokableRun(context.Background(), - `{"symbols":["AAPL","MSFT"]}`) - if err != nil { - t.Fatalf("InvokableRun: %v", err) + raw, err := NewYahooFinanceTool().InvokableRun(context.Background(), `{`) + if err == nil || !strings.Contains(err.Error(), "parse arguments") { + t.Fatalf("err = %v", err) } - if !strings.Contains(gotUA, "ragflow") { - t.Errorf("User-Agent = %q, want to contain ragflow", gotUA) + var envelope yahooFinanceEnvelope + if decodeErr := json.Unmarshal([]byte(raw), &envelope); decodeErr != nil { + t.Fatalf("error result is not JSON: %s: %v", raw, decodeErr) } - - var env yahooFinanceEnvelope - if jerr := json.Unmarshal([]byte(out), &env); jerr != nil { - t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out) - } - if env.Error != "" { - t.Errorf("Error = %q, want empty", env.Error) - } - if len(env.Results) != 2 { - t.Fatalf("Results len = %d, want 2", len(env.Results)) - } - if env.Results[0].Symbol != "AAPL" { - t.Errorf("Results[0].Symbol = %q, want AAPL", env.Results[0].Symbol) - } - if env.Results[0].RegularMarketPrice != 189.5 { - t.Errorf("Results[0].Price = %v, want 189.5", env.Results[0].RegularMarketPrice) - } - if env.Results[1].RegularMarketChangePercent != -0.5 { - t.Errorf("Results[1].ChangePct = %v, want -0.5", env.Results[1].RegularMarketChangePercent) + if !strings.Contains(envelope.Error, "parse arguments") { + t.Fatalf("envelope = %#v", envelope) } } -func TestYahooFinance_RequiresSymbols(t *testing.T) { +func TestYahooFinanceComponentContract(t *testing.T) { t.Parallel() - tool := NewYahooFinanceTool() - _, err := tool.InvokableRun(context.Background(), `{"symbols":[]}`) - if err == nil { - t.Fatal("expected error for empty symbols") + yahoo := NewYahooFinanceTool() + spec := yahoo.ComponentSpec() + stockCode, ok := spec.InputForm["stock_code"].(map[string]any) + if !ok || stockCode["type"] != "line" || stockCode["name"] != "Stock code/Company name" { + t.Fatalf("stock_code input form = %#v", spec.InputForm["stock_code"]) } - if !strings.Contains(err.Error(), "symbols") { - t.Errorf("err = %v, want to mention symbols", err) + if _, exists := spec.Outputs["report"]; !exists { + t.Fatalf("outputs = %#v", spec.Outputs) + } + outputs := yahoo.BuildComponentOutputs(map[string]any{"report": "# Information:\nAAPL"}) + if report, ok := outputs["report"].(string); !ok || report != "# Information:\nAAPL" { + t.Fatalf("report = %#v", outputs["report"]) } } -func TestYahooFinance_Info(t *testing.T) { +func TestYahooFinanceInfoOnlyExposesStockCode(t *testing.T) { t.Parallel() - tool := NewYahooFinanceTool() - info, err := tool.Info(context.Background()) + info, err := NewYahooFinanceTool().Info(context.Background()) if err != nil { t.Fatalf("Info: %v", err) } - if info.Name != "yahoo_finance" { - t.Errorf("Name = %q, want yahoo_finance", info.Name) + if info.Name != "yahoo_finance" || info.Desc != yahooFinanceToolDescription { + t.Fatalf("Info = name %q, desc %q", info.Name, info.Desc) } - if !strings.Contains(info.Desc, "Yahoo") { - t.Errorf("Desc = %q, want to mention Yahoo", info.Desc) + jsonSchema, err := info.ParamsOneOf.ToJSONSchema() + if err != nil { + t.Fatalf("ToJSONSchema: %v", err) + } + raw, err := json.Marshal(jsonSchema) + if err != nil { + t.Fatalf("marshal schema: %v", err) + } + schemaText := string(raw) + for _, expected := range []string{`"stock_code"`, `"required":["stock_code"]`} { + if !strings.Contains(schemaText, expected) { + t.Fatalf("schema missing %s: %s", expected, schemaText) + } + } + for _, leaked := range []string{`"symbols"`, `"fields"`, `"info"`, `"news"`} { + if strings.Contains(schemaText, leaked) { + t.Fatalf("schema leaked node config %s: %s", leaked, schemaText) + } + } +} + +func TestBuildYahooFinanceToolUsesNodeDefaults(t *testing.T) { + t.Parallel() + + built, err := BuildByName("yahoo_finance", map[string]any{ + "info": false, "history": true, "balance_sheet": true, "news": true, + "stock_code": "sys.query", "outputs": map[string]any{"report": map[string]any{}}, + }) + if err != nil { + t.Fatalf("BuildByName: %v", err) + } + yahoo, ok := built.(*YahooFinanceTool) + if !ok { + t.Fatalf("tool type = %T", built) + } + if yahoo.defaults.Info { + t.Fatalf("defaults = %#v", yahoo.defaults) + } + if yahoo.defaults.StockCode != "" { + t.Fatalf("runtime stock_code leaked into defaults: %#v", yahoo.defaults) + } +} + +func TestBuildYahooFinanceToolRejectsInvalidInfoParam(t *testing.T) { + t.Parallel() + + _, err := BuildByName("yahoo_finance", map[string]any{"info": "true"}) + if err == nil || !strings.Contains(err.Error(), "requires boolean node-level param info") { + t.Fatalf("err = %v", err) } }