Extract shared HTTP client for search tools
The userAgentTransport and 15s timeout were inconsistently applied: government_db.go lost its timeout, wayback.go used a bare http.Client without the pooled transport or user-agent header. A new httpclient.go centralizes the setup so all search tools share the same configuration. Signed-off-by: Émile Ré <emile@probo.com>
This commit is contained in:
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
39
pkg/agent/tools/search/httpclient.go
Normal file
39
pkg/agent/tools/search/httpclient.go
Normal file
@@ -0,0 +1,39 @@
|
||||
// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
|
||||
//
|
||||
// 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
|
||||
}
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user