From 0e672782ac83ee20b9cdaa5ada99d7d96760696e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=89mile=20R=C3=A9?= Date: Fri, 15 May 2026 17:43:42 +0400 Subject: [PATCH] Add Firecrawl web search tool for tracker mapping MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Firecrawl provides higher quality search results than SearXNG. When configured (firecrawl-endpoint + firecrawl-api-key), the tracker-mapping agent and search toolset prefer it over the SearXNG backend. Also improves the tracker identification prompt with multi-strategy search queries that leverage domain signals and adapt to tracker type. Signed-off-by: Émile Ré --- pkg/agent/tools/search/firecrawl.go | 139 ++++++++++++++++++ pkg/agent/tools/search/government_db.go | 6 +- pkg/agent/tools/search/search.go | 45 +++++- pkg/agent/tools/search/web_search.go | 16 +- .../prompts/tracker_identification.txt | 7 +- pkg/cookiebanner/tracker_mapping_worker.go | 30 ++-- pkg/probod/tracker_mapping.go | 8 +- pkg/probodconfig/config.go | 2 + 8 files changed, 222 insertions(+), 31 deletions(-) create mode 100644 pkg/agent/tools/search/firecrawl.go diff --git a/pkg/agent/tools/search/firecrawl.go b/pkg/agent/tools/search/firecrawl.go new file mode 100644 index 000000000..83b9a4a64 --- /dev/null +++ b/pkg/agent/tools/search/firecrawl.go @@ -0,0 +1,139 @@ +// 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 ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + + "go.gearno.de/kit/httpclient" + + "go.probo.inc/probo/pkg/agent" +) + +type ( + 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)"` + } + + firecrawlRequest struct { + Query string `json:"query"` + Limit int `json:"limit"` + } + + firecrawlResponse struct { + Success bool `json:"success"` + Data struct { + Web []firecrawlWebResult `json:"web"` + } `json:"data"` + } + + firecrawlWebResult struct { + Title string `json:"title"` + Description string `json:"description"` + URL string `json:"url"` + } +) + +// 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() + + return agent.FunctionTool( + "web_search", + "Search the web for information about a topic. Returns a list of results with title, URL, and snippet. Use this to find news, reviews, breach reports, regulatory actions, and other external information about a vendor.", + func(ctx context.Context, p firecrawlParams) (agent.ToolResult, error) { + maxResults := p.MaxResults + if maxResults <= 0 { + maxResults = 5 + } + if maxResults > 10 { + maxResults = 10 + } + + results, err := firecrawlSearch(ctx, client, endpoint, apiKey, p.Query, maxResults) + if err != nil { + return agent.ResultErrorf("search request failed: %s", err), nil + } + + return agent.ResultJSON(results), nil + }, + ) +} + +func firecrawlSearch(ctx context.Context, client *http.Client, endpoint, apiKey, query string, maxResults int) ([]searchResult, error) { + u, err := url.JoinPath(endpoint, "search") + if err != nil { + return nil, fmt.Errorf("cannot build search URL: %w", err) + } + + body, err := json.Marshal(firecrawlRequest{ + Query: query, + Limit: maxResults, + }) + if err != nil { + return nil, fmt.Errorf("cannot marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, u, bytes.NewReader(body)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+apiKey) + + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer func() { _ = resp.Body.Close() }() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("search returned status %d", resp.StatusCode) + } + + var fcResp firecrawlResponse + if err := json.Unmarshal(respBody, &fcResp); err != nil { + return nil, err + } + + if !fcResp.Success { + return nil, fmt.Errorf("search returned success=false") + } + + results := make([]searchResult, 0, len(fcResp.Data.Web)) + for _, r := range fcResp.Data.Web { + results = append(results, searchResult{ + Title: r.Title, + URL: r.URL, + Snippet: r.Description, + }) + } + + return results, nil +} diff --git a/pkg/agent/tools/search/government_db.go b/pkg/agent/tools/search/government_db.go index b4d40dc60..02df23f10 100644 --- a/pkg/agent/tools/search/government_db.go +++ b/pkg/agent/tools/search/government_db.go @@ -21,7 +21,8 @@ import ( "io" "net/http" "net/url" - "time" + + "go.gearno.de/kit/httpclient" "go.probo.inc/probo/pkg/agent" ) @@ -49,7 +50,8 @@ type ( ) func CheckGovernmentDBTool(searchEndpoint string) agent.Tool { - client := &http.Client{Timeout: 15 * time.Second} + client := httpclient.DefaultPooledClient() + client.Transport = &userAgentTransport{next: client.Transport} return agent.FunctionTool( "check_government_databases", diff --git a/pkg/agent/tools/search/search.go b/pkg/agent/tools/search/search.go index a5e5eed5b..fb582cb10 100644 --- a/pkg/agent/tools/search/search.go +++ b/pkg/agent/tools/search/search.go @@ -20,19 +20,52 @@ import ( // Toolset provides web search tools. type Toolset struct { - endpoint string + endpoint string + firecrawlEndpoint string + firecrawlAPIKey string +} + +// Option configures a Toolset. +type Option func(*Toolset) + +// WithFirecrawl configures the Firecrawl search backend. The endpoint +// should be the base URL of the Firecrawl instance (e.g. +// "https://api.firecrawl.dev/v2"). +func WithFirecrawl(endpoint, apiKey string) Option { + return func(t *Toolset) { + t.firecrawlEndpoint = endpoint + t.firecrawlAPIKey = apiKey + } } // NewToolset creates a search toolset with the given SearXNG endpoint. -func NewToolset(endpoint string) *Toolset { - return &Toolset{endpoint: endpoint} +func NewToolset(endpoint string, opts ...Option) *Toolset { + t := &Toolset{endpoint: endpoint} + for _, opt := range opts { + opt(t) + } + return t } func (t *Toolset) Tools() []agent.Tool { - return []agent.Tool{ - WebSearchTool(t.endpoint), - CheckGovernmentDBTool(t.endpoint), + tools := []agent.Tool{ CheckWaybackTool(), DiffDocumentsTool(), } + + if t.firecrawlEndpoint != "" && t.firecrawlAPIKey != "" { + tools = append( + tools, + FirecrawlSearchTool(t.firecrawlEndpoint, t.firecrawlAPIKey), + CheckGovernmentDBTool(t.endpoint), + ) + } else if t.endpoint != "" { + tools = append( + tools, + WebSearchTool(t.endpoint), + CheckGovernmentDBTool(t.endpoint), + ) + } + + return tools } diff --git a/pkg/agent/tools/search/web_search.go b/pkg/agent/tools/search/web_search.go index d9cabbc7a..536d888a4 100644 --- a/pkg/agent/tools/search/web_search.go +++ b/pkg/agent/tools/search/web_search.go @@ -17,7 +17,8 @@ package search import ( "context" "net/http" - "time" + + "go.gearno.de/kit/httpclient" "go.probo.inc/probo/pkg/agent" ) @@ -43,13 +44,24 @@ type ( URL string `json:"url"` Content string `json:"content"` } + + 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) +} + // WebSearchTool creates a tool that searches the web using a SearXNG instance. // The endpoint should be the base URL of the SearXNG instance (e.g. // "http://localhost:8888"). func WebSearchTool(endpoint string) agent.Tool { - client := &http.Client{Timeout: 15 * time.Second} + client := httpclient.DefaultPooledClient() + client.Transport = &userAgentTransport{next: client.Transport} return agent.FunctionTool( "web_search", diff --git a/pkg/cookiebanner/prompts/tracker_identification.txt b/pkg/cookiebanner/prompts/tracker_identification.txt index b0a36b0db..ace804bc5 100644 --- a/pkg/cookiebanner/prompts/tracker_identification.txt +++ b/pkg/cookiebanner/prompts/tracker_identification.txt @@ -17,7 +17,12 @@ Return a structured JSON response with: 2. If search_tracker_patterns returns results with a third party name, use the search_third_parties tool to confirm the exact name in the database and get its category. -3. Only use web_search as a last resort if the internal tools return nothing useful. Search for "what is cookie [name]" or "[name] cookie tracker". +3. Only use web_search as a last resort if the internal tools return nothing useful. Try up to 3 different queries, adapting your strategy based on the available signals: + - If the tracker has a recognizable prefix or name, start with: "[prefix] cookie tracking" (e.g. "_ce.s cookie tracking"). + - If observed domains are available and the name is opaque, search by domain instead: "[domain] cookies tracking privacy" (e.g. "clarity.ms cookies tracking privacy"). + - For localStorage keys, include the type: "[name] localStorage tracking script". + - If the first query returns nothing useful, broaden: "[name] web tracker" or "site:[domain] cookie documentation". + - Stop searching once you get a confident match; do not exhaust all query slots if the first one succeeds. 4. Common cookie naming conventions to recognize: - _ga*, _gid, _gat*: Google Analytics diff --git a/pkg/cookiebanner/tracker_mapping_worker.go b/pkg/cookiebanner/tracker_mapping_worker.go index 33ca05c29..15db3ea89 100644 --- a/pkg/cookiebanner/tracker_mapping_worker.go +++ b/pkg/cookiebanner/tracker_mapping_worker.go @@ -51,9 +51,11 @@ type trackerMappingHandler struct { } type TrackerMappingConfig struct { - LLMClient *llm.Client - Model string - SearchEndpoint string + LLMClient *llm.Client + Model string + SearchEndpoint string + FirecrawlEndpoint string + FirecrawlAPIKey string } func NewTrackerMappingWorker( @@ -68,13 +70,7 @@ func NewTrackerMappingWorker( } if cfg.LLMClient != nil { - h.agent = buildTrackerMappingAgent( - cfg.LLMClient, - cfg.Model, - cfg.SearchEndpoint, - pgClient, - logger, - ) + h.agent = buildTrackerMappingAgent(cfg, pgClient, logger) } return worker.New( @@ -86,9 +82,7 @@ func NewTrackerMappingWorker( } func buildTrackerMappingAgent( - llmClient *llm.Client, - model string, - searchEndpoint string, + cfg TrackerMappingConfig, pgClient *pg.Client, logger *log.Logger, ) *agent.Agent { @@ -97,8 +91,10 @@ func buildTrackerMappingAgent( searchThirdPartiesTool(pgClient), } - if searchEndpoint != "" { - tools = append(tools, search.WebSearchTool(searchEndpoint)) + if cfg.FirecrawlEndpoint != "" && cfg.FirecrawlAPIKey != "" { + tools = append(tools, search.FirecrawlSearchTool(cfg.FirecrawlEndpoint, cfg.FirecrawlAPIKey)) + } else if cfg.SearchEndpoint != "" { + tools = append(tools, search.WebSearchTool(cfg.SearchEndpoint)) } outputType, err := agent.NewOutputType[TrackerIdentification]("tracker_identification") @@ -108,9 +104,9 @@ func buildTrackerMappingAgent( return agent.New( "tracker-mapping", - llmClient, + cfg.LLMClient, agent.WithInstructions(trackerIdentificationPrompt), - agent.WithModel(model), + agent.WithModel(cfg.Model), agent.WithTools(tools...), agent.WithOutputType(outputType), agent.WithMaxTurns(agentMaxTurns), diff --git a/pkg/probod/tracker_mapping.go b/pkg/probod/tracker_mapping.go index d0960825d..f9c2a3756 100644 --- a/pkg/probod/tracker_mapping.go +++ b/pkg/probod/tracker_mapping.go @@ -45,8 +45,10 @@ func (impl *Implm) buildTrackerMappingConfig( } return cookiebanner.TrackerMappingConfig{ - LLMClient: llmClient, - Model: agentCfg.ModelName, - SearchEndpoint: impl.cfg.SearchEndpoint, + LLMClient: llmClient, + Model: agentCfg.ModelName, + SearchEndpoint: impl.cfg.SearchEndpoint, + FirecrawlEndpoint: impl.cfg.FirecrawlEndpoint, + FirecrawlAPIKey: impl.cfg.FirecrawlAPIKey, }, nil } diff --git a/pkg/probodconfig/config.go b/pkg/probodconfig/config.go index f63696d1c..5ed9ad9e0 100644 --- a/pkg/probodconfig/config.go +++ b/pkg/probodconfig/config.go @@ -62,6 +62,8 @@ type ( EvidenceDescriber EvidenceDescriberConfig `json:"evidence-describer"` ChromeDPAddr string `json:"chrome-dp-addr"` SearchEndpoint string `json:"search-endpoint"` + FirecrawlEndpoint string `json:"firecrawl-endpoint"` + FirecrawlAPIKey string `json:"firecrawl-api-key"` CustomDomains CustomDomainsConfig `json:"custom-domains"` SCIMBridge SCIMBridgeConfig `json:"scim-bridge"` ESign ESignConfig `json:"esign"`