rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Extract agent and tools from tracker mapping worker into dedicated files

Browse Source 
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é <emile@probo.com>
This commit is contained in:
Émile Ré
2026-05-19 12:06:15 +04:00
parent 10adf2bd4b
commit c4147e6801
7 changed files with 256 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
contrib/claude/agent.md
View File

@@ -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

37
contrib/claude/file-naming.md
View File

@@ -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 `<name>_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 `<worker_prefix>_agent.go` alongside `<worker_prefix>_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_name>_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
```

80
pkg/cookiebanner/search_third_parties_tool.go Normal file
View File

@@ -0,0 +1,80 @@
// 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 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
},
)
}

55
pkg/cookiebanner/tools_tracker_mapping.go → pkg/cookiebanner/search_tracker_patterns_tool.go
View File

@@ -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
},
)
}

26
pkg/cookiebanner/tracker_identification.go
View File

@@ -1,26 +0,0 @@
// 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 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."`
}

126
pkg/cookiebanner/tracker_mapping_agent.go Normal file
View File

@@ -0,0 +1,126 @@
// 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 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"+
"<pattern> %s </pattern>\n"+
"<type> %s </type>\n"+
"<match_type> %s </match_type>\n"+
"<max_age> %s </max_age>\n",
tp.Pattern,
tp.TrackerType,
tp.MatchType,
maxAge,
)
if len(domains) > 0 {
prompt += fmt.Sprintf("<observed_domains> %s </observed_domains>\n", strings.Join(domains, ", "))
}
return prompt
}

94
pkg/cookiebanner/tracker_mapping_worker.go
View File

@@ -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"+
"<pattern> %s </pattern>\n"+
"<type> %s </type>\n"+
"<match_type> %s </match_type>\n"+
"<max_age> %s </max_age>\n",
tp.Pattern,
tp.TrackerType,
tp.MatchType,
maxAge,
)
if len(domains) > 0 {
prompt += fmt.Sprintf("<observed_domains> %s </observed_domains>\n", strings.Join(domains, ", "))
}
return prompt
}