diff --git a/pkg/agent/tools/search/firecrawl.go b/pkg/agent/tools/search/firecrawl.go index dcc473ca2..d3545ea53 100644 --- a/pkg/agent/tools/search/firecrawl.go +++ b/pkg/agent/tools/search/firecrawl.go @@ -22,9 +22,7 @@ import ( "io" "net/http" "net/url" - "time" - "go.gearno.de/kit/httpclient" "go.probo.inc/probo/pkg/agent" ) @@ -35,10 +33,6 @@ type ( Snippet string `json:"snippet"` } - userAgentTransport struct { - next http.RoundTripper - } - firecrawlParams struct { Query string `json:"query" jsonschema:"The search query to execute"` MaxResults int `json:"max_results" jsonschema:"Maximum number of results to return (default 5, max 10)"` @@ -63,19 +57,11 @@ type ( } ) -func (t *userAgentTransport) RoundTrip(r *http.Request) (*http.Response, error) { - r2 := r.Clone(r.Context()) - r2.Header.Set("User-Agent", "Probo-Agent/1.0") - return t.next.RoundTrip(r2) -} - // FirecrawlSearchTool creates a tool that searches the web using the Firecrawl // API. The endpoint should be the base URL of the Firecrawl instance (e.g. // "https://api.firecrawl.dev/v2"). The apiKey is used for Bearer authentication. func FirecrawlSearchTool(endpoint, apiKey string) agent.Tool { - client := httpclient.DefaultPooledClient() - client.Timeout = 15 * time.Second - client.Transport = &userAgentTransport{next: client.Transport} + client := newHTTPClient() return agent.FunctionTool( "web_search", diff --git a/pkg/agent/tools/search/government_db.go b/pkg/agent/tools/search/government_db.go index 0bdab10c8..132771d2d 100644 --- a/pkg/agent/tools/search/government_db.go +++ b/pkg/agent/tools/search/government_db.go @@ -18,7 +18,6 @@ import ( "context" "fmt" - "go.gearno.de/kit/httpclient" "go.probo.inc/probo/pkg/agent" ) @@ -45,8 +44,7 @@ type ( ) func CheckGovernmentDBTool(endpoint, apiKey string) agent.Tool { - client := httpclient.DefaultPooledClient() - client.Transport = &userAgentTransport{next: client.Transport} + client := newHTTPClient() return agent.FunctionTool( "check_government_databases", diff --git a/pkg/agent/tools/search/httpclient.go b/pkg/agent/tools/search/httpclient.go new file mode 100644 index 000000000..bcc0a22a9 --- /dev/null +++ b/pkg/agent/tools/search/httpclient.go @@ -0,0 +1,39 @@ +// Copyright (c) 2026 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package search + +import ( + "net/http" + "time" + + "go.gearno.de/kit/httpclient" +) + +type userAgentTransport struct { + next http.RoundTripper +} + +func (t *userAgentTransport) RoundTrip(r *http.Request) (*http.Response, error) { + r2 := r.Clone(r.Context()) + r2.Header.Set("User-Agent", "Probo-Agent/1.0") + return t.next.RoundTrip(r2) +} + +func newHTTPClient() *http.Client { + client := httpclient.DefaultPooledClient() + client.Timeout = 15 * time.Second + client.Transport = &userAgentTransport{next: client.Transport} + return client +} diff --git a/pkg/agent/tools/search/wayback.go b/pkg/agent/tools/search/wayback.go index 6c6f08798..3b26a84b6 100644 --- a/pkg/agent/tools/search/wayback.go +++ b/pkg/agent/tools/search/wayback.go @@ -21,7 +21,6 @@ import ( "io" "net/http" "net/url" - "time" "go.probo.inc/probo/pkg/agent" ) @@ -57,7 +56,7 @@ type ( ) func CheckWaybackTool() agent.Tool { - client := &http.Client{Timeout: 15 * time.Second} + client := newHTTPClient() return agent.FunctionTool( "check_wayback",