From c4147e68010af83d1fdae9be3f7add014097dfdd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=89mile=20R=C3=A9?= Date: Tue, 19 May 2026 12:06:15 +0400 Subject: [PATCH] Extract agent and tools from tracker mapping worker into dedicated files MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Split tracker_mapping_worker.go: agent construction, prompts, and structured output type move to tracker_mapping_agent.go; each tool gets its own *_tool.go file. Update naming conventions (worker, agent, tool file patterns, AgentResult suffix, RunTyped preference). Signed-off-by: Émile Ré --- contrib/claude/agent.md | 12 +- contrib/claude/file-naming.md | 37 +++++ pkg/cookiebanner/search_third_parties_tool.go | 80 +++++++++++ ...ing.go => search_tracker_patterns_tool.go} | 55 -------- pkg/cookiebanner/tracker_identification.go | 26 ---- pkg/cookiebanner/tracker_mapping_agent.go | 126 ++++++++++++++++++ pkg/cookiebanner/tracker_mapping_worker.go | 94 +------------ 7 files changed, 256 insertions(+), 174 deletions(-) create mode 100644 pkg/cookiebanner/search_third_parties_tool.go rename pkg/cookiebanner/{tools_tracker_mapping.go => search_tracker_patterns_tool.go} (62%) delete mode 100644 pkg/cookiebanner/tracker_identification.go create mode 100644 pkg/cookiebanner/tracker_mapping_agent.go diff --git a/contrib/claude/agent.md b/contrib/claude/agent.md index cbd6fcde0..0f294b9c7 100644 --- a/contrib/claude/agent.md +++ b/contrib/claude/agent.md @@ -24,7 +24,17 @@ result.FinalMessage().Text() // final output result.LastAgent // agent that produced the result ``` -Typed output via `RunTyped[T](ctx, agent, messages)` — validates against JSON Schema. +Prefer `RunTyped[T]` over `Run` + manual unmarshalling when the agent declares a +structured output type via `WithOutputType`. `RunTyped` validates the response +against the JSON Schema and returns the typed value directly: + +```go +result, err := agent.RunTyped[TrackerIdentification](ctx, ag, messages) +identification := result.Output // already typed, no json.Unmarshal needed +``` + +Only fall back to `agent.Run` when the agent produces free-form text with no +output schema. ## Tool interface diff --git a/contrib/claude/file-naming.md b/contrib/claude/file-naming.md index 0c4214434..2e9ecb9f0 100644 --- a/contrib/claude/file-naming.md +++ b/contrib/claude/file-naming.md @@ -27,3 +27,40 @@ Use one of: {{.Categories}}. For single-placeholder templates, `strings.Replace` is sufficient. For multiple placeholders, use `text/template`. + +## Worker files + +Background worker files use the `_worker.go` naming pattern. The file +contains the handler struct, `NewXxxWorker` constructor, and the `Claim`/`Process` +methods. Keep domain-specific helper methods (match, resolve, etc.) in the same +file. + +``` +pkg/cookiebanner/tracker_mapping_worker.go +pkg/cookiebanner/pattern_analysis_worker.go +``` + +## Agent files + +When a package has a worker that uses an agent, the agent construction logic +goes in `_agent.go` alongside `_worker.go`. The +worker file stays focused on `Claim`/`Process` and handler methods; the agent +file owns agent construction, prompt building, constants, config, and the +`//go:embed` directive for prompt templates. + +``` +pkg/cookiebanner/tracker_mapping_worker.go -- worker handler +pkg/cookiebanner/tracker_mapping_agent.go -- agent construction + prompts +``` + +## Tool files + +Each agent tool lives in its own `_tool.go` file, named after the +tool constructor function (without the `Tool` suffix). The file contains the +constructor function plus its param/result types. Avoid grouping multiple tools +in a single file. + +``` +pkg/cookiebanner/search_tracker_patterns_tool.go -- searchTrackerPatternsTool +pkg/cookiebanner/search_third_parties_tool.go -- searchThirdPartiesTool +``` diff --git a/pkg/cookiebanner/search_third_parties_tool.go b/pkg/cookiebanner/search_third_parties_tool.go new file mode 100644 index 000000000..669f2226a --- /dev/null +++ b/pkg/cookiebanner/search_third_parties_tool.go @@ -0,0 +1,80 @@ +// 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 cookiebanner + +import ( + "context" + + "go.gearno.de/kit/pg" + "go.probo.inc/probo/pkg/agent" + "go.probo.inc/probo/pkg/coredata" +) + +type ( + searchThirdPartiesParams struct { + Query string `json:"query" jsonschema:"Search fragment to match against known third party names (e.g. 'Google', 'Meta', 'Hotjar')"` + } + + searchThirdPartiesResult struct { + Name string `json:"name"` + Category string `json:"category"` + WebsiteURL string `json:"website_url,omitempty"` + } +) + +func searchThirdPartiesTool(pgClient *pg.Client) agent.Tool { + return agent.FunctionTool( + "search_third_parties", + "Search the internal database of known third parties (companies/services) by name fragment. Returns matching third party names, categories, and website URLs. Use this to find the exact name of a known third party to link the tracker to.", + func(ctx context.Context, p searchThirdPartiesParams) (agent.ToolResult, error) { + if p.Query == "" { + return agent.ResultError("query is required"), nil + } + + var out []searchThirdPartiesResult + + if err := pgClient.WithConn( + ctx, + func(ctx context.Context, conn pg.Querier) error { + var parties coredata.CommonThirdParties + if err := parties.LoadAll( + ctx, + conn, + coredata.NewCommonThirdPartyFilter(&p.Query), + ); err != nil { + return err + } + + out = make([]searchThirdPartiesResult, len(parties)) + for i, tp := range parties { + out[i] = searchThirdPartiesResult{ + Name: tp.Name, + Category: string(tp.Category), + } + if tp.WebsiteURL != nil { + out[i].WebsiteURL = *tp.WebsiteURL + } + } + + return nil + }, + ); err != nil { + return agent.ResultErrorf("search failed: %s", err), nil + } + + return agent.ResultJSON(out), nil + }, + ) +} diff --git a/pkg/cookiebanner/tools_tracker_mapping.go b/pkg/cookiebanner/search_tracker_patterns_tool.go similarity index 62% rename from pkg/cookiebanner/tools_tracker_mapping.go rename to pkg/cookiebanner/search_tracker_patterns_tool.go index cbba74170..355554052 100644 --- a/pkg/cookiebanner/tools_tracker_mapping.go +++ b/pkg/cookiebanner/search_tracker_patterns_tool.go @@ -34,16 +34,6 @@ type ( ThirdPartyName string `json:"third_party_name,omitempty"` Confidence float32 `json:"confidence"` } - - searchThirdPartiesParams struct { - Query string `json:"query" jsonschema:"Search fragment to match against known third party names (e.g. 'Google', 'Meta', 'Hotjar')"` - } - - searchThirdPartiesResult struct { - Name string `json:"name"` - Category string `json:"category"` - WebsiteURL string `json:"website_url,omitempty"` - } ) func searchTrackerPatternsTool(pgClient *pg.Client) agent.Tool { @@ -89,48 +79,3 @@ func searchTrackerPatternsTool(pgClient *pg.Client) agent.Tool { }, ) } - -func searchThirdPartiesTool(pgClient *pg.Client) agent.Tool { - return agent.FunctionTool( - "search_third_parties", - "Search the internal database of known third parties (companies/services) by name fragment. Returns matching third party names, categories, and website URLs. Use this to find the exact name of a known third party to link the tracker to.", - func(ctx context.Context, p searchThirdPartiesParams) (agent.ToolResult, error) { - if p.Query == "" { - return agent.ResultError("query is required"), nil - } - - var out []searchThirdPartiesResult - - if err := pgClient.WithConn( - ctx, - func(ctx context.Context, conn pg.Querier) error { - var parties coredata.CommonThirdParties - if err := parties.LoadAll( - ctx, - conn, - coredata.NewCommonThirdPartyFilter(&p.Query), - ); err != nil { - return err - } - - out = make([]searchThirdPartiesResult, len(parties)) - for i, tp := range parties { - out[i] = searchThirdPartiesResult{ - Name: tp.Name, - Category: string(tp.Category), - } - if tp.WebsiteURL != nil { - out[i].WebsiteURL = *tp.WebsiteURL - } - } - - return nil - }, - ); err != nil { - return agent.ResultErrorf("search failed: %s", err), nil - } - - return agent.ResultJSON(out), nil - }, - ) -} diff --git a/pkg/cookiebanner/tracker_identification.go b/pkg/cookiebanner/tracker_identification.go deleted file mode 100644 index 83629440c..000000000 --- a/pkg/cookiebanner/tracker_identification.go +++ /dev/null @@ -1,26 +0,0 @@ -// 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 cookiebanner - -import "go.probo.inc/probo/pkg/coredata" - -// TrackerIdentification is the structured output the tracker-mapping -// agent returns. -type TrackerIdentification struct { - ThirdPartyName string `json:"third_party_name" jsonschema:"Name of the company or service that sets this tracker (e.g. 'Google Analytics', 'Meta Pixel'). Empty string if truly unknown."` - Category coredata.ThirdPartyCategory `json:"category" jsonschema:"Third party category"` - Description string `json:"description" jsonschema:"What this tracker does in one sentence"` - Confidence float64 `json:"confidence" jsonschema:"Confidence level from 0.0 to 1.0. Set below 0.5 if unsure."` -} diff --git a/pkg/cookiebanner/tracker_mapping_agent.go b/pkg/cookiebanner/tracker_mapping_agent.go new file mode 100644 index 000000000..765c98d80 --- /dev/null +++ b/pkg/cookiebanner/tracker_mapping_agent.go @@ -0,0 +1,126 @@ +// 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 cookiebanner + +import ( + "context" + _ "embed" + "fmt" + "strings" + "time" + + "go.gearno.de/kit/log" + "go.gearno.de/kit/pg" + "go.probo.inc/probo/pkg/agent" + "go.probo.inc/probo/pkg/agent/tools/search" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/llm" +) + +const ( + agentTimeout = 60 * time.Second + agentMaxTurns = 5 + agentConfidenceThreshold = 0.6 + agentMaxPatternConfidence = 0.8 +) + +//go:embed prompts/tracker_identification.txt.tmpl +var trackerIdentificationPrompt string + +// TrackerMappingAgentResult is the structured output the tracker-mapping +// agent returns. +type TrackerMappingAgentResult struct { + ThirdPartyName string `json:"third_party_name" jsonschema:"Name of the company or service that sets this tracker (e.g. 'Google Analytics', 'Meta Pixel'). Empty string if truly unknown."` + Category coredata.ThirdPartyCategory `json:"category" jsonschema:"Third party category"` + Description string `json:"description" jsonschema:"What this tracker does in one sentence"` + Confidence float64 `json:"confidence" jsonschema:"Confidence level from 0.0 to 1.0. Set below 0.5 if unsure."` +} + +type TrackerMappingConfig struct { + LLMClient *llm.Client + Model string + FirecrawlAPIKey string +} + +func buildTrackerMappingAgent( + cfg TrackerMappingConfig, + pgClient *pg.Client, + logger *log.Logger, +) *agent.Agent { + tools := []agent.Tool{ + searchTrackerPatternsTool(pgClient), + searchThirdPartiesTool(pgClient), + } + + if cfg.FirecrawlAPIKey != "" { + tools = append(tools, search.FirecrawlSearchTool(cfg.FirecrawlAPIKey)) + } + + outputType, err := agent.NewOutputType[TrackerMappingAgentResult]("tracker_identification") + if err != nil { + panic(fmt.Sprintf("cookiebanner: cannot build tracker identification output type: %s", err)) + } + + return agent.New( + "tracker-mapping", + cfg.LLMClient, + agent.WithInstructionsFunc(trackerMappingInstructions), + agent.WithModel(cfg.Model), + agent.WithTools(tools...), + agent.WithOutputType(outputType), + agent.WithMaxTurns(agentMaxTurns), + agent.WithLogger(logger), + ) +} + +func trackerMappingInstructions(_ context.Context, _ *agent.Agent) string { + categories := coredata.ThirdPartyCategories() + parts := make([]string, len(categories)) + for i, c := range categories { + parts[i] = string(c) + } + + return strings.Replace( + trackerIdentificationPrompt, + "{{.Categories}}", + strings.Join(parts, ", "), + 1, + ) +} + +func buildAgentPrompt(tp coredata.TrackerPattern, domains []string) string { + maxAge := "session" + if tp.MaxAgeSeconds != nil { + maxAge = fmt.Sprintf("%d seconds", *tp.MaxAgeSeconds) + } + + prompt := fmt.Sprintf( + "Identify the following tracker:\n\n"+ + " %s \n"+ + " %s \n"+ + " %s \n"+ + " %s \n", + tp.Pattern, + tp.TrackerType, + tp.MatchType, + maxAge, + ) + + if len(domains) > 0 { + prompt += fmt.Sprintf(" %s \n", strings.Join(domains, ", ")) + } + + return prompt +} diff --git a/pkg/cookiebanner/tracker_mapping_worker.go b/pkg/cookiebanner/tracker_mapping_worker.go index 140579fdb..7d4acb6a3 100644 --- a/pkg/cookiebanner/tracker_mapping_worker.go +++ b/pkg/cookiebanner/tracker_mapping_worker.go @@ -16,45 +16,26 @@ package cookiebanner import ( "context" - _ "embed" "errors" "fmt" - "strings" "time" "go.gearno.de/kit/log" "go.gearno.de/kit/pg" "go.gearno.de/kit/worker" "go.probo.inc/probo/pkg/agent" - "go.probo.inc/probo/pkg/agent/tools/search" "go.probo.inc/probo/pkg/coredata" "go.probo.inc/probo/pkg/gid" "go.probo.inc/probo/pkg/llm" "go.probo.inc/probo/pkg/slug" ) -const ( - agentTimeout = 60 * time.Second - agentMaxTurns = 5 - agentConfidenceThreshold = 0.6 - agentMaxPatternConfidence = 0.8 -) - -//go:embed prompts/tracker_identification.txt.tmpl -var trackerIdentificationPrompt string - type trackerMappingHandler struct { pg *pg.Client logger *log.Logger agent *agent.Agent } -type TrackerMappingConfig struct { - LLMClient *llm.Client - Model string - FirecrawlAPIKey string -} - func NewTrackerMappingWorker( pgClient *pg.Client, logger *log.Logger, @@ -78,52 +59,6 @@ func NewTrackerMappingWorker( ) } -func buildTrackerMappingAgent( - cfg TrackerMappingConfig, - pgClient *pg.Client, - logger *log.Logger, -) *agent.Agent { - tools := []agent.Tool{ - searchTrackerPatternsTool(pgClient), - searchThirdPartiesTool(pgClient), - } - - if cfg.FirecrawlAPIKey != "" { - tools = append(tools, search.FirecrawlSearchTool(cfg.FirecrawlAPIKey)) - } - - outputType, err := agent.NewOutputType[TrackerIdentification]("tracker_identification") - if err != nil { - panic(fmt.Sprintf("cookiebanner: cannot build tracker identification output type: %s", err)) - } - - return agent.New( - "tracker-mapping", - cfg.LLMClient, - agent.WithInstructionsFunc(trackerMappingInstructions), - agent.WithModel(cfg.Model), - agent.WithTools(tools...), - agent.WithOutputType(outputType), - agent.WithMaxTurns(agentMaxTurns), - agent.WithLogger(logger), - ) -} - -func trackerMappingInstructions(_ context.Context, _ *agent.Agent) string { - categories := coredata.ThirdPartyCategories() - parts := make([]string, len(categories)) - for i, c := range categories { - parts[i] = string(c) - } - - return strings.Replace( - trackerIdentificationPrompt, - "{{.Categories}}", - strings.Join(parts, ", "), - 1, - ) -} - func (h *trackerMappingHandler) Claim(ctx context.Context) (coredata.TrackerPattern, error) { var tp coredata.TrackerPattern @@ -292,7 +227,7 @@ func (h *trackerMappingHandler) identifyWithAgent( agentCtx, cancel := context.WithTimeout(ctx, agentTimeout) defer cancel() - result, err := agent.RunTyped[TrackerIdentification]( + result, err := agent.RunTyped[TrackerMappingAgentResult]( agentCtx, h.agent, []llm.Message{ @@ -379,7 +314,7 @@ func (h *trackerMappingHandler) identifyWithAgent( func (h *trackerMappingHandler) resolveOrCreateCommonThirdParty( ctx context.Context, tx pg.Tx, - identification TrackerIdentification, + identification TrackerMappingAgentResult, domains []string, ) (*gid.GID, error) { var party coredata.CommonThirdParty @@ -490,28 +425,3 @@ func (h *trackerMappingHandler) resolveThirdParty( return &t.ID, nil } - -func buildAgentPrompt(tp coredata.TrackerPattern, domains []string) string { - maxAge := "session" - if tp.MaxAgeSeconds != nil { - maxAge = fmt.Sprintf("%d seconds", *tp.MaxAgeSeconds) - } - - prompt := fmt.Sprintf( - "Identify the following tracker:\n\n"+ - " %s \n"+ - " %s \n"+ - " %s \n"+ - " %s \n", - tp.Pattern, - tp.TrackerType, - tp.MatchType, - maxAge, - ) - - if len(domains) > 0 { - prompt += fmt.Sprintf(" %s \n", strings.Join(domains, ", ")) - } - - return prompt -}