From b29a4a61eb7175839e2b80d493839dee27b0d925 Mon Sep 17 00:00:00 2001 From: Hz_ Date: Thu, 9 Jul 2026 16:38:42 +0800 Subject: [PATCH] fix(ao-agent): add Tavily input_form, wire real search component, and port TavilyExtract to Go (#16693) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary - TavilySearch now stores api_key from component params and injects it into tool calls when runtime inputs omit it. - TavilyExtract and BGPT now follow the same stored api_key behavior. - Canvas decoding now recovers api_key from graph.nodes[].data.form when components[].obj.params.api_key is empty, matching frontend payload behavior without changing frontend data. - Added regression tests for graph form key recovery and stored key injection / caller key precedence. Tests: build.sh --test ./internal/agent/component ./internal/service — all pass. image image --- .../agent/component/bgpt_component_test.go | 31 ++ internal/agent/component/fixture_stubs.go | 20 +- internal/agent/component/invoke_test.go | 69 +++-- .../agent/component/tavily_component_test.go | 102 +++++++ .../agent/component/universe_a_wrappers.go | 154 +++++++++- internal/agent/component/verify_p1_test.go | 4 +- internal/agent/tool/registry.go | 1 + internal/agent/tool/registry_test.go | 8 +- internal/agent/tool/tavily.go | 275 +++++++++++++++++- internal/agent/tool/tavily_test.go | 106 +++++++ 10 files changed, 700 insertions(+), 70 deletions(-) create mode 100644 internal/agent/component/tavily_component_test.go diff --git a/internal/agent/component/bgpt_component_test.go b/internal/agent/component/bgpt_component_test.go index 8b7265a8c..ac7242cdb 100644 --- a/internal/agent/component/bgpt_component_test.go +++ b/internal/agent/component/bgpt_component_test.go @@ -106,3 +106,34 @@ func TestBGPT_InvokeAdaptsCanvasInputsAndOutputs(t *testing.T) { 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/fixture_stubs.go b/internal/agent/component/fixture_stubs.go index c8d52a909..a160cc7f9 100644 --- a/internal/agent/component/fixture_stubs.go +++ b/internal/agent/component/fixture_stubs.go @@ -17,9 +17,11 @@ // Package component — e2e fixture stubs and compat shims. // // The test fixtures under internal/agent/dsl/testdata reference -// seven component names that are registered here: Retrieval, +// fixture-backed component names that are registered here: Retrieval, // TavilySearch, ExeSQL, Generate, Answer, Iteration, -// IterationItem. Their bodies are deliberately trivial — they +// 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 // 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 @@ -27,7 +29,7 @@ // universe_a_wrappers.go and the real production bodies in // their own .go files replace these stubs in production paths. // -// The seven names were chosen by enumerating the component_name +// The fixture names were chosen by enumerating the component_name // values in the testdata fixtures (see the `examples` var in // internal/agent/canvas/dsl_examples_test.go). Keeping the list // in sync with the fixture set is a single-source-of-truth @@ -156,6 +158,15 @@ func (t *TavilySearchStub) Outputs() map[string]string { } } +func (t *TavilySearchStub) GetInputForm() map[string]any { + return map[string]any{ + "query": map[string]any{ + "name": "Query", + "type": "line", + }, + } +} + // ----- ExeSQL ----- const componentNameExeSQL = "ExeSQL" @@ -498,7 +509,8 @@ func init() { Register("SearchMyDataset", newRetrievalComponent) Register("search_my_dataset", newRetrievalComponent) Register("search_my_dateset", newRetrievalComponent) - Register(componentNameTavilySearch, NewTavilySearchStub) + Register(componentNameTavilySearch, newTavilySearchComponent) + Register("TavilyExtract", newTavilyExtractComponent) Register(componentNameExeSQL, newExeSQLComponent) Register(componentNameCodeExec, newCodeExecComponent) Register(componentNameGenerate, NewGenerateStub) diff --git a/internal/agent/component/invoke_test.go b/internal/agent/component/invoke_test.go index eceebd256..649592171 100644 --- a/internal/agent/component/invoke_test.go +++ b/internal/agent/component/invoke_test.go @@ -19,8 +19,10 @@ package component import ( "context" "io" + "net" "net/http" "net/http/httptest" + "net/url" "strings" "testing" "time" @@ -317,17 +319,26 @@ func TestInvoke_NoRedirects_NotFollowed(t *testing.T) { // would re-open the rebinding window. PR review round 5, // Major #3. func TestInvoke_ProxyDNSPin(t *testing.T) { - setupAllowAnyHost(t, false) - // The validated proxy IP we will pin the dial to. - // 192.88.99.1 is the 6to4 anycast prefix — a real public - // IP that the SSRF guard accepts (not in any block-list - // range) but is highly unlikely to be listening on port - // 9999 in the test environment, so the dial fails fast - // with "connection refused" carrying the IP we pinned. - // Earlier versions used 1.1.1.1, but Cloudflare's edge - // network has started responding on 9999, so we switched - // to a less-likely-to-listen anycast prefix. - const pinnedProxyIP = "192.88.99.1" + setupAllowAnyHost(t, true) + + proxyHit := make(chan struct{}, 1) + proxySrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + proxyHit <- struct{}{} + if r.RequestURI != "http://8.8.8.8/api" { + t.Errorf("proxy RequestURI = %q, want absolute-form target", r.RequestURI) + } + w.WriteHeader(http.StatusNoContent) + })) + defer proxySrv.Close() + + proxyURL, err := url.Parse(proxySrv.URL) + if err != nil { + t.Fatalf("parse proxy server URL: %v", err) + } + pinnedProxyIP, proxyPort, err := net.SplitHostPort(proxyURL.Host) + if err != nil { + t.Fatalf("split proxy server host: %v", err) + } // Build a small Invoke call with a proxy URL whose // hostname will resolve via SSRF (we override the @@ -339,43 +350,29 @@ func TestInvoke_ProxyDNSPin(t *testing.T) { // to return a different IP on a second call. originalLookup := utility.LookupHost utility.LookupHost = func(host string) ([]string, error) { - // Always return a public IP so SSRF accepts. + // Always return the already-running fake proxy. If the + // Invoke transport re-resolves proxy.test.invalid instead + // of using the pinned IP, the request will never hit it. return []string{pinnedProxyIP}, nil } t.Cleanup(func() { utility.LookupHost = originalLookup }) - // We expect the dial to fail (no proxy server at - // 1.1.1.1:9999). The error message tells us whether - // the pin was active: with the fix, the dial targets - // 1.1.1.1:9999; without the fix, the Go transport - // would have re-resolved the hostname and dialed a - // different IP (or refused to dial because the - // stubbed hostname is unreachable). c, _ := NewInvokeComponent(nil) - // Tight per-call timeout so the test fails fast — the request - // will hang waiting for the unreachable proxy (1.1.1.1:9999) - // to respond, and the default 30s client timeout would make - // this test take 30s on every run. ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - _, err := c.Invoke(ctx, map[string]any{ + _, err = c.Invoke(ctx, map[string]any{ "method": "GET", "url": "http://8.8.8.8/api", - "proxy": "http://proxy.test.invalid:9999", + "proxy": "http://proxy.test.invalid:" + proxyPort, "timeout": 2, }) - if err == nil { - t.Fatal("expected dial error (no proxy listening), got nil") + if err != nil { + t.Fatalf("Invoke: %v", err) } - // The dial error must reference 1.1.1.1:9999 (the - // pinned IP) — NOT the unresolved proxy.test.invalid - // hostname, which would prove the dialer fell through - // to the default resolver. - if !strings.Contains(err.Error(), pinnedProxyIP+":9999") { - t.Fatalf("dial error = %v; want pinned proxy IP %s:9999 "+ - "(connection-refused is acceptable; an absent IP means "+ - "the dialer fell through to the default resolver and the "+ - "pinning regression went undetected)", err, pinnedProxyIP) + select { + case <-proxyHit: + case <-time.After(time.Second): + t.Fatal("fake proxy was not hit; proxy dial was not pinned to validated IP") } } diff --git a/internal/agent/component/tavily_component_test.go b/internal/agent/component/tavily_component_test.go new file mode 100644 index 000000000..efa310136 --- /dev/null +++ b/internal/agent/component/tavily_component_test.go @@ -0,0 +1,102 @@ +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/universe_a_wrappers.go b/internal/agent/component/universe_a_wrappers.go index 6f292ea39..df27eff46 100644 --- a/internal/agent/component/universe_a_wrappers.go +++ b/internal/agent/component/universe_a_wrappers.go @@ -53,11 +53,15 @@ import ( // The underlying tool makes a real HTTP call; the wrapper is the // canvas-facing surface. type tavilySearchComponent struct { - inner *agenttool.TavilyTool + inner *agenttool.TavilyTool + apiKey string } -func newTavilySearchComponent(_ map[string]any) (Component, error) { - return &tavilySearchComponent{inner: agenttool.NewTavilyTool()}, nil +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" } @@ -74,23 +78,110 @@ func (c *tavilySearchComponent) Inputs() map[string]string { func (c *tavilySearchComponent) Outputs() map[string]string { return map[string]string{ "formalized_content": "Rendered search results for downstream LLM prompts.", - "results": "Raw result list (url, title, content).", + "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) } - return parseToolEnvelope(out), nil + 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 } +// 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) @@ -99,15 +190,20 @@ type bgptInvoker interface { // bgptComponent delegates to internal/agent/tool/BGPTTool and adapts // the tool envelope to the BGPT canvas output contract. type bgptComponent struct { - inner bgptInvoker + inner bgptInvoker + apiKey string } -func newBGPTComponent(_ map[string]any) (Component, error) { - return newBGPTComponentWithInvoker(agenttool.NewBGPTTool()), nil +func newBGPTComponent(params map[string]any) (Component, error) { + return newBGPTComponentWithInvoker(agenttool.NewBGPTTool(), stringParam(params["api_key"])), nil } -func newBGPTComponentWithInvoker(inner bgptInvoker) Component { - return &bgptComponent{inner: inner} +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" } @@ -138,6 +234,10 @@ func (c *bgptComponent) Outputs() map[string]string { } 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 @@ -202,6 +302,38 @@ 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 "" @@ -1049,6 +1181,7 @@ func (c *yahooFinanceComponent) Stream(_ context.Context, _ map[string]any) (<-c var ( _ Component = (*retrievalComponent)(nil) _ Component = (*tavilySearchComponent)(nil) + _ Component = (*tavilyExtractComponent)(nil) _ Component = (*exesqlComponent)(nil) _ Component = (*codeExecComponent)(nil) _ Component = (*yahooFinanceComponent)(nil) @@ -1057,4 +1190,5 @@ var ( // 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.TavilyExtractTool)(nil) var _ einotool.InvokableTool = (*agenttool.YahooFinanceTool)(nil) diff --git a/internal/agent/component/verify_p1_test.go b/internal/agent/component/verify_p1_test.go index 62244e0c1..4fa8af1a6 100644 --- a/internal/agent/component/verify_p1_test.go +++ b/internal/agent/component/verify_p1_test.go @@ -44,8 +44,8 @@ func TestVerifyRegistration_P1(t *testing.T) { if len(missing) > 0 { t.Fatalf("missing P0/P1 components: %v (have %d: %v)", missing, len(names), names) } - if got := len(names); got < 12 || got > 32 { - t.Errorf("expected 12-32 registered (current plan scope + v1 stubs), got %d: %v", got, names) + if got := len(names); got < 12 || got > 33 { + t.Errorf("expected 12-33 registered (current plan scope + v1 wrappers/stubs), got %d: %v", got, names) } // ExitLoop must NOT be in the registry (legacy compat lives at diff --git a/internal/agent/tool/registry.go b/internal/agent/tool/registry.go index 0d1de6c6f..3b802977b 100644 --- a/internal/agent/tool/registry.go +++ b/internal/agent/tool/registry.go @@ -51,6 +51,7 @@ var registry = map[string]Factory{ "search_my_dateset": noConfig("search_my_dateset", func() einotool.BaseTool { return NewRetrievalTool() }), "searxng": noConfig("searxng", func() einotool.BaseTool { return NewSearXNGTool() }), "tavily": noConfig("tavily", func() einotool.BaseTool { return NewTavilyTool() }), + "tavily_extract": noConfig("tavily_extract", func() einotool.BaseTool { return NewTavilyExtractTool() }), "tushare": noConfig("tushare", func() einotool.BaseTool { return NewTushareTool() }), "wencai": noConfig("wencai", func() einotool.BaseTool { return NewWencaiTool() }), "web_crawler": noConfig("web_crawler", func() einotool.BaseTool { return NewCrawlerTool() }), diff --git a/internal/agent/tool/registry_test.go b/internal/agent/tool/registry_test.go index de889a7fa..b23318aa0 100644 --- a/internal/agent/tool/registry_test.go +++ b/internal/agent/tool/registry_test.go @@ -41,13 +41,13 @@ func TestBuildAll_UnknownTool(t *testing.T) { } func TestBuildAll_AllRegisteredTools(t *testing.T) { - // Every key in registry (27 entries, 23 unique canonical tools). + // Every key in registry (28 entries, 24 unique canonical tools). names := []string{ "akshare", "arxiv", "bgpt", "code_exec", "crawler", "deepl", "duckduckgo", "email", "exesql", "execute_sql", "github", "google", "google_scholar", "jin10", "keenable", "pubmed", "qweather", "retrieval", "search_my_dataset", "search_my_dateset", "searxng", - "tavily", "tushare", "web_crawler", "wencai", "wikipedia", + "tavily", "tavily_extract", "tushare", "web_crawler", "wencai", "wikipedia", "yahoo_finance", } params := map[string]map[string]any{ @@ -116,7 +116,7 @@ func TestBuildAll_KeenableRejectsEmptyNodeAPIKey(t *testing.T) { func TestToolRegistry_SchemasAreComplete(t *testing.T) { t.Parallel() - // Every entry the registry advertises. 27 names, 23 unique + // Every entry the registry advertises. 28 names, 24 unique // canonical tools (execute_sql == exesql, retrieval == // search_my_dataset == search_my_dateset, crawler == web_crawler). names := []string{ @@ -124,7 +124,7 @@ func TestToolRegistry_SchemasAreComplete(t *testing.T) { "duckduckgo", "email", "execute_sql", "exesql", "github", "google", "google_scholar", "jin10", "keenable", "pubmed", "qweather", "retrieval", "search_my_dataset", "search_my_dateset", "searxng", - "tavily", "tushare", "web_crawler", "wencai", "wikipedia", + "tavily", "tavily_extract", "tushare", "web_crawler", "wencai", "wikipedia", "yahoo_finance", } params := map[string]map[string]any{ diff --git a/internal/agent/tool/tavily.go b/internal/agent/tool/tavily.go index 5ac385dc7..b25b79c2b 100644 --- a/internal/agent/tool/tavily.go +++ b/internal/agent/tool/tavily.go @@ -22,6 +22,7 @@ import ( "fmt" "net/http" "os" + "strings" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" @@ -29,32 +30,52 @@ import ( const tavilyToolName = "tavily" +const tavilyExtractToolName = "tavily_extract" + const tavilyToolDescription = "Search the web via the Tavily API. Returns a list of {url, title, content} results." // 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. type tavilyParams struct { - APIKey string `json:"api_key"` - Query string `json:"query"` - MaxResults int `json:"max_results"` - SearchDepth string `json:"search_depth"` + APIKey string `json:"api_key"` + Query string `json:"query"` + MaxResults int `json:"max_results"` + SearchDepth string `json:"search_depth"` + Topic string `json:"topic"` + Days int `json:"days"` + IncludeAnswer bool `json:"include_answer"` + IncludeRawContent bool `json:"include_raw_content"` + IncludeImages bool `json:"include_images"` + IncludeImageDescriptions bool `json:"include_image_descriptions"` + IncludeDomains []string `json:"include_domains"` + ExcludeDomains []string `json:"exclude_domains"` } // tavilyRequestBody is the JSON body POSTed to the Tavily /search endpoint. // The struct matches the upstream API (https://docs.tavily.com). type tavilyRequestBody struct { - Query string `json:"query"` - MaxResults int `json:"max_results"` - SearchDepth string `json:"search_depth"` + Query string `json:"query"` + MaxResults int `json:"max_results"` + SearchDepth string `json:"search_depth"` + Topic string `json:"topic,omitempty"` + Days int `json:"days"` + IncludeAnswer bool `json:"include_answer"` + IncludeRawContent bool `json:"include_raw_content"` + IncludeImages bool `json:"include_images"` + IncludeImageDescriptions bool `json:"include_image_descriptions"` + IncludeDomains []string `json:"include_domains,omitempty"` + 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"` + 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 @@ -70,6 +91,36 @@ type tavilyEnvelope struct { Error string `json:"_ERROR,omitempty"` } +type tavilyExtractParams struct { + APIKey string `json:"api_key"` + URLs any `json:"urls"` + ExtractDepth string `json:"extract_depth"` + Format string `json:"format"` +} + +type tavilyExtractRequestBody struct { + URLs []string `json:"urls"` + ExtractDepth string `json:"extract_depth"` + Format string `json:"format"` + 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"` +} + +type tavilyExtractEnvelope struct { + Results []tavilyExtractResult `json:"results"` + Error string `json:"_ERROR,omitempty"` +} + // TavilyTool is the Tavily search // tool. It POSTs a search request // to https://api.tavily.com/search using the shared HTTPHelper and returns @@ -79,6 +130,13 @@ type TavilyTool struct { envKey func() string } +// 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 +} + // NewTavilyTool returns a TavilyTool using the default HTTPHelper and // the TAVILY_API_KEY env var for credential resolution. func NewTavilyTool() *TavilyTool { @@ -98,12 +156,39 @@ func NewTavilyToolWith(h *HTTPHelper) *TavilyTool { // resolver. Useful for tests that want to inject a fake credential // without mutating process state. func NewTavilyToolWithEnvKey(h *HTTPHelper, envKey func() string) *TavilyTool { + if h == nil { + h = NewHTTPHelper() + } if envKey == nil { envKey = defaultTavilyEnvKey } return &TavilyTool{helper: h, envKey: envKey} } +// NewTavilyExtractTool returns a TavilyExtractTool using the default HTTPHelper. +func NewTavilyExtractTool() *TavilyExtractTool { + return NewTavilyExtractToolWith(NewHTTPHelper()) +} + +// NewTavilyExtractToolWith returns a TavilyExtractTool using the provided helper. +func NewTavilyExtractToolWith(h *HTTPHelper) *TavilyExtractTool { + if h == nil { + h = NewHTTPHelper() + } + return &TavilyExtractTool{helper: h, envKey: defaultTavilyEnvKey} +} + +// NewTavilyExtractToolWithEnvKey returns a TavilyExtractTool with a custom env-key resolver. +func NewTavilyExtractToolWithEnvKey(h *HTTPHelper, envKey func() string) *TavilyExtractTool { + if h == nil { + h = NewHTTPHelper() + } + if envKey == nil { + envKey = defaultTavilyEnvKey + } + return &TavilyExtractTool{helper: h, envKey: envKey} +} + // defaultTavilyEnvKey is the production env-key resolver. Pulled out // as a named function (not a var) so tests cannot accidentally // mutate it via package-var assignment. @@ -135,6 +220,46 @@ func (t *TavilyTool) Info(_ context.Context) (*schema.ToolInfo, error) { Desc: `Tavily search depth: "basic" (default) or "advanced".`, Required: false, }, + "topic": { + Type: schema.String, + Desc: `Search topic: "general" (default) or "news".`, + Required: false, + }, + "include_domains": { + Type: schema.Array, + Desc: "Domains that search results must include.", + Required: false, + }, + "exclude_domains": { + Type: schema.Array, + Desc: "Domains that search results must exclude.", + Required: false, + }, + }), + }, nil +} + +// Info returns the Tavily Extract tool metadata for the chat model. +func (t *TavilyExtractTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: tavilyExtractToolName, + Desc: "Extract web page content from one or more specified URLs using Tavily Extract.", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "urls": { + Type: schema.Array, + Desc: "The URLs to extract content from.", + Required: true, + }, + "extract_depth": { + Type: schema.String, + Desc: `Extraction depth: "basic" (default) or "advanced".`, + Required: false, + }, + "format": { + Type: schema.String, + Desc: `Output format: "markdown" (default) or "text".`, + Required: false, + }, }), }, nil } @@ -144,6 +269,10 @@ func (t *TavilyTool) Info(_ context.Context) (*schema.ToolInfo, error) { // access with tavilyEndpointMu if running in parallel. var tavilyEndpoint = "https://api.tavily.com/search" +// tavilyExtractEndpoint is the Tavily /extract URL. Exposed as a package var +// so tests can substitute it through rewriteHostTransport. +var tavilyExtractEndpoint = "https://api.tavily.com/extract" + // InvokableRun performs the Tavily search. The api_key may come from the // argument or the TAVILY_API_KEY env var. func (t *TavilyTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) { @@ -157,7 +286,7 @@ func (t *TavilyTool) InvokableRun(ctx context.Context, argsJSON string, _ ...too fmt.Errorf("tavily: query is required") } if p.MaxResults <= 0 { - p.MaxResults = 5 + p.MaxResults = 6 } if p.SearchDepth == "" { p.SearchDepth = "basic" @@ -173,9 +302,17 @@ func (t *TavilyTool) InvokableRun(ctx context.Context, argsJSON string, _ ...too } body, _ := json.Marshal(tavilyRequestBody{ - Query: p.Query, - MaxResults: p.MaxResults, - SearchDepth: p.SearchDepth, + Query: p.Query, + MaxResults: p.MaxResults, + SearchDepth: p.SearchDepth, + Topic: defaultString(p.Topic, "general"), + Days: p.Days, + IncludeAnswer: p.IncludeAnswer, + IncludeRawContent: false, + IncludeImages: false, + IncludeImageDescriptions: p.IncludeImageDescriptions, + IncludeDomains: p.IncludeDomains, + ExcludeDomains: p.ExcludeDomains, }) resp, err := t.helper.Do(ctx, @@ -200,6 +337,63 @@ func (t *TavilyTool) InvokableRun(ctx context.Context, argsJSON string, _ ...too return tavilyJSON(tavilyEnvelope{Results: raw.Results}), nil } +// 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 { + return tavilyExtractErrJSON(fmt.Errorf("tavily_extract: parse arguments: %w", err)), + fmt.Errorf("tavily_extract: parse arguments: %w", err) + } + urls := normalizeTavilyURLs(p.URLs) + if len(urls) == 0 { + return tavilyExtractErrJSON(fmt.Errorf("urls is required")), + fmt.Errorf("tavily_extract: urls is required") + } + if p.ExtractDepth == "" { + p.ExtractDepth = "basic" + } + if p.Format == "" { + p.Format = "markdown" + } + + apiKey := p.APIKey + if apiKey == "" { + apiKey = t.envKey() + } + if apiKey == "" { + return tavilyExtractErrJSON(fmt.Errorf("tavily_extract: api_key is required (or set TAVILY_API_KEY)")), + fmt.Errorf("tavily_extract: api_key is required (or set TAVILY_API_KEY)") + } + + body, _ := json.Marshal(tavilyExtractRequestBody{ + URLs: urls, + ExtractDepth: p.ExtractDepth, + Format: p.Format, + IncludeImages: false, + }) + resp, err := t.helper.Do(ctx, + http.MethodPost, tavilyExtractEndpoint, string(body), "application/json", + map[string]string{"Authorization": "Bearer " + apiKey}, + ) + if err != nil { + return tavilyExtractErrJSON(err), err + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return tavilyExtractErrJSON(fmt.Errorf("tavily_extract: upstream returned %d", resp.StatusCode)), + fmt.Errorf("tavily_extract: upstream returned %d", resp.StatusCode) + } + + var raw tavilyExtractResponse + if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil { + return tavilyExtractErrJSON(fmt.Errorf("tavily_extract: decode response: %w", err)), + fmt.Errorf("tavily_extract: decode response: %w", err) + } + return tavilyExtractJSON(tavilyExtractEnvelope{Results: raw.Results}), nil +} + // tavilyJSON marshals the envelope to a JSON string for the model. func tavilyJSON(env tavilyEnvelope) string { b, err := json.Marshal(env) @@ -213,3 +407,56 @@ func tavilyJSON(env tavilyEnvelope) string { func tavilyErrJSON(err error) string { return tavilyJSON(tavilyEnvelope{Error: err.Error()}) } + +func tavilyExtractJSON(env tavilyExtractEnvelope) string { + b, err := json.Marshal(env) + if err != nil { + return fmt.Sprintf(`{"_ERROR":"tavily_extract: marshal result: %s"}`, err) + } + return string(b) +} + +func tavilyExtractErrJSON(err error) string { + return tavilyExtractJSON(tavilyExtractEnvelope{Error: err.Error()}) +} + +func normalizeTavilyURLs(raw any) []string { + switch v := raw.(type) { + case string: + return splitTavilyURLs(v) + case []string: + return compactStrings(v) + case []any: + out := make([]string, 0, len(v)) + for _, item := range v { + if s, ok := item.(string); ok { + out = append(out, s) + } + } + return compactStrings(out) + default: + return nil + } +} + +func splitTavilyURLs(s string) []string { + parts := strings.Split(s, ",") + return compactStrings(parts) +} + +func compactStrings(in []string) []string { + out := make([]string, 0, len(in)) + for _, s := range in { + if trimmed := strings.TrimSpace(s); trimmed != "" { + out = append(out, trimmed) + } + } + return out +} + +func defaultString(v, fallback string) string { + if v == "" { + return fallback + } + return v +} diff --git a/internal/agent/tool/tavily_test.go b/internal/agent/tool/tavily_test.go index 2d2b2d93e..7442b7453 100644 --- a/internal/agent/tool/tavily_test.go +++ b/internal/agent/tool/tavily_test.go @@ -176,3 +176,109 @@ func TestTavily_Info(t *testing.T) { t.Errorf("Desc = %q, want to mention Tavily", info.Desc) } } + +func TestTavilyExtract_BuildRequest(t *testing.T) { + t.Parallel() + + var gotPath, gotAuth, gotMethod string + var gotBody map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotMethod = r.Method + gotAuth = r.Header.Get("Authorization") + _ = json.NewDecoder(r.Body).Decode(&gotBody) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"results":[{"url":"https://example.com","raw_content":"hello"}]}`)) + })) + defer srv.Close() + + helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)}) + tool := NewTavilyExtractToolWith(helper) + out, err := tool.InvokableRun(context.Background(), + `{"urls":"https://a.example, https://b.example","api_key":"key-xyz","extract_depth":"advanced","format":"text"}`) + if err != nil { + t.Fatalf("InvokableRun: %v", err) + } + if out == "" { + t.Fatal("InvokableRun returned empty string") + } + if gotMethod != http.MethodPost { + t.Errorf("method = %q, want POST", gotMethod) + } + if !strings.HasSuffix(gotPath, "extract") { + t.Errorf("path = %q, want .../extract", gotPath) + } + if gotAuth != "Bearer key-xyz" { + t.Errorf("Authorization = %q, want %q", gotAuth, "Bearer key-xyz") + } + urls, ok := gotBody["urls"].([]any) + if !ok || len(urls) != 2 || urls[0] != "https://a.example" || urls[1] != "https://b.example" { + t.Errorf("body.urls = %#v, want two trimmed URLs", gotBody["urls"]) + } + if gotBody["extract_depth"] != "advanced" { + t.Errorf("body.extract_depth = %v, want advanced", gotBody["extract_depth"]) + } + if gotBody["format"] != "text" { + t.Errorf("body.format = %v, want text", gotBody["format"]) + } +} + +func TestTavilyExtract_ParseResponse(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/","raw_content":"alpha"}]}`)) + })) + defer srv.Close() + + helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)}) + tool := NewTavilyExtractToolWith(helper) + out, err := tool.InvokableRun(context.Background(), `{"urls":["https://a.example/"],"api_key":"k"}`) + if err != nil { + t.Fatalf("InvokableRun: %v", err) + } + + var env tavilyExtractEnvelope + 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) != 1 { + t.Fatalf("Results len = %d, want 1", len(env.Results)) + } + if env.Results[0].URL != "https://a.example/" || env.Results[0].RawContent != "alpha" { + t.Errorf("Results[0] = %+v, want url and raw_content", env.Results[0]) + } +} + +func TestTavilyExtract_RequiresAPIKey(t *testing.T) { + t.Parallel() + + tool := NewTavilyExtractToolWithEnvKey(NewHTTPHelper(), func() string { return "" }) + _, err := tool.InvokableRun(context.Background(), `{"urls":["https://a.example/"]}`) + if err == nil { + t.Fatal("expected error for missing api_key") + } + if !strings.Contains(err.Error(), "api_key") { + t.Errorf("err = %v, want to mention api_key", err) + } +} + +func TestTavilyExtract_Info(t *testing.T) { + t.Parallel() + + tool := NewTavilyExtractTool() + info, err := tool.Info(context.Background()) + if err != nil { + t.Fatalf("Info: %v", err) + } + if info.Name != "tavily_extract" { + t.Errorf("Name = %q, want tavily_extract", info.Name) + } + if !strings.Contains(info.Desc, "Tavily Extract") { + t.Errorf("Desc = %q, want to mention Tavily Extract", info.Desc) + } +}