From 92a73e93026edb6c2e65aee823827e30946519be Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=89mile=20R=C3=A9?= Date: Mon, 8 Jun 2026 11:24:37 +0200 Subject: [PATCH] Give each tracker agent its own config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The tracker-mapping, common-pattern enrichment, and third-party disambiguation agents were all built from one shared TrackerAgentsConfig fed by a single tracker-mapping config slot. That forced a single AgentTimeout to be reused and patched per worker, and two unrelated max-turns fields to share one struct. Split the in-code config into TrackerMappingAgentConfig, TrackerEnrichmentAgentConfig, and DisambiguationAgentConfig, each with its own timeout and max-turns, and add dedicated tracker-enrichment and third-party-disambiguation provider slots (the latter resolving next to third-party-vetter). Enrichment and disambiguation fall back to the tracker-mapping slot when their own provider is unset, preserving single-config deployments. Drop the shared pkg/agentsbuild package and duplicate its small wiring into probod and proboctl so the two executables stay decoupled. Wire the new env vars, builder test coverage, and Helm values. Signed-off-by: Émile Ré --- .../charts/probo/templates/deployment.yaml | 38 ++++ .../probo/values-production.yaml.example | 20 +- contrib/helm/charts/probo/values.yaml | 24 ++- pkg/agentsbuild/agentsbuild.go | 172 ------------------ pkg/bootstrap/builder.go | 30 ++- pkg/bootstrap/builder_test.go | 31 ++++ pkg/cookiebanner/common_pattern_enricher.go | 54 ++++-- .../common_pattern_enrichment_agent.go | 4 +- .../common_pattern_enrichment_worker.go | 5 +- pkg/cookiebanner/tracker_agents_config.go | 51 ++++-- pkg/cookiebanner/tracker_mapping_agent.go | 4 +- pkg/cookiebanner/tracker_mapping_worker.go | 6 +- pkg/proboctl/cmdutil/cmdutil.go | 68 +++++-- pkg/proboctl/cmdutil/llm.go | 97 ++++++++++ pkg/proboctl/commontrackerpattern/reenrich.go | 4 +- pkg/probod/llm.go | 73 +++++++- pkg/probod/probod.go | 12 +- pkg/probod/tracker_agents.go | 102 ++++++++++- pkg/probodconfig/llm_config.go | 40 ++-- pkg/thirdparty/disambiguation_agent.go | 6 +- 20 files changed, 552 insertions(+), 289 deletions(-) delete mode 100644 pkg/agentsbuild/agentsbuild.go create mode 100644 pkg/proboctl/cmdutil/llm.go diff --git a/contrib/helm/charts/probo/templates/deployment.yaml b/contrib/helm/charts/probo/templates/deployment.yaml index 57310f03e..795fc1bbd 100644 --- a/contrib/helm/charts/probo/templates/deployment.yaml +++ b/contrib/helm/charts/probo/templates/deployment.yaml @@ -281,6 +281,23 @@ spec: - name: THIRD_PARTY_VETTING_STALE_AFTER value: {{ .Values.probo.thirdPartyVettingWorker.staleAfter | quote }} {{- end }} + # Third-party Disambiguation Agent + {{- if .Values.probo.thirdPartyDisambiguation.provider }} + - name: AGENT_THIRD_PARTY_DISAMBIGUATION_PROVIDER + value: {{ .Values.probo.thirdPartyDisambiguation.provider | quote }} + {{- end }} + {{- if .Values.probo.thirdPartyDisambiguation.modelName }} + - name: AGENT_THIRD_PARTY_DISAMBIGUATION_MODEL_NAME + value: {{ .Values.probo.thirdPartyDisambiguation.modelName | quote }} + {{- end }} + {{- if .Values.probo.thirdPartyDisambiguation.temperature }} + - name: AGENT_THIRD_PARTY_DISAMBIGUATION_TEMPERATURE + value: {{ .Values.probo.thirdPartyDisambiguation.temperature | quote }} + {{- end }} + {{- if .Values.probo.thirdPartyDisambiguation.maxTokens }} + - name: AGENT_THIRD_PARTY_DISAMBIGUATION_MAX_TOKENS + value: {{ .Values.probo.thirdPartyDisambiguation.maxTokens | quote }} + {{- end }} # Tracker Mapping Agent {{- if .Values.probo.trackerMapping.provider }} - name: AGENT_TRACKER_MAPPING_PROVIDER @@ -298,6 +315,23 @@ spec: - name: AGENT_TRACKER_MAPPING_MAX_TOKENS value: {{ .Values.probo.trackerMapping.maxTokens | quote }} {{- end }} + # Tracker Enrichment Agent + {{- if .Values.probo.trackerEnrichment.provider }} + - name: AGENT_TRACKER_ENRICHMENT_PROVIDER + value: {{ .Values.probo.trackerEnrichment.provider | quote }} + {{- end }} + {{- if .Values.probo.trackerEnrichment.modelName }} + - name: AGENT_TRACKER_ENRICHMENT_MODEL_NAME + value: {{ .Values.probo.trackerEnrichment.modelName | quote }} + {{- end }} + {{- if .Values.probo.trackerEnrichment.temperature }} + - name: AGENT_TRACKER_ENRICHMENT_TEMPERATURE + value: {{ .Values.probo.trackerEnrichment.temperature | quote }} + {{- end }} + {{- if .Values.probo.trackerEnrichment.maxTokens }} + - name: AGENT_TRACKER_ENRICHMENT_MAX_TOKENS + value: {{ .Values.probo.trackerEnrichment.maxTokens | quote }} + {{- end }} # Tracker Mapping Worker {{- if .Values.probo.trackerMappingWorker.interval }} - name: TRACKER_MAPPING_INTERVAL @@ -319,6 +353,10 @@ spec: - name: TRACKER_MAPPING_AGENT_MAX_TURNS value: {{ .Values.probo.trackerMappingWorker.agentMaxTurns | quote }} {{- end }} + {{- if .Values.probo.trackerMappingWorker.disambiguationAgentTimeout }} + - name: TRACKER_MAPPING_DISAMBIGUATION_AGENT_TIMEOUT + value: {{ .Values.probo.trackerMappingWorker.disambiguationAgentTimeout | quote }} + {{- end }} # Common Pattern Enrichment Worker {{- if .Values.probo.commonPatternEnrichmentWorker.interval }} - name: COMMON_PATTERN_ENRICHMENT_INTERVAL diff --git a/contrib/helm/charts/probo/values-production.yaml.example b/contrib/helm/charts/probo/values-production.yaml.example index 909806a55..4de046539 100644 --- a/contrib/helm/charts/probo/values-production.yaml.example +++ b/contrib/helm/charts/probo/values-production.yaml.example @@ -185,6 +185,15 @@ probo: # maxConcurrency: 1 # staleAfter: 1500 + # Third-party disambiguation agent (optional; runs inside the tracker + # mapping worker to pick the best matching org third party). Falls back + # to trackerMapping when its provider is unset. + # thirdPartyDisambiguation: + # provider: "openai" + # modelName: "gpt-4o" + # temperature: "0.4" + # maxTokens: "4096" + # Tracker mapping agent (optional, auto-links tracker patterns to vendors) # trackerMapping: # provider: "openai" @@ -192,8 +201,16 @@ probo: # temperature: "0.1" # maxTokens: "4096" + # Tracker enrichment agent (optional; researches catalog descriptions). + # Falls back to trackerMapping when its provider is unset. + # trackerEnrichment: + # provider: "openai" + # modelName: "gpt-4o" + # temperature: "0.2" + # maxTokens: "4096" + # Tracker mapping worker tuning (optional; seconds for - # interval/staleAfter/agentTimeout). + # interval/staleAfter/agentTimeout/disambiguationAgentTimeout). # Keep concurrency modest to stay under OpenAI/Firecrawl limits and the DB pool. # trackerMappingWorker: # interval: 10 @@ -201,6 +218,7 @@ probo: # staleAfter: 600 # agentTimeout: 45 # agentMaxTurns: 10 + # disambiguationAgentTimeout: 45 # Common-pattern enrichment worker tuning (optional; seconds for # interval/staleAfter/agentTimeout). diff --git a/contrib/helm/charts/probo/values.yaml b/contrib/helm/charts/probo/values.yaml index 94dcede64..27e2d8656 100644 --- a/contrib/helm/charts/probo/values.yaml +++ b/contrib/helm/charts/probo/values.yaml @@ -286,6 +286,15 @@ probo: maxConcurrency: 1 staleAfter: 1500 + # Third-party disambiguation agent (optional; runs inside the tracker + # mapping worker to pick the best matching org third party). Falls back + # to trackerMapping when its provider is unset. + thirdPartyDisambiguation: + provider: "" + modelName: "" + temperature: "" + maxTokens: "" + # Tracker mapping agent (optional, requires openai.apiKey or anthropic key) trackerMapping: provider: "" @@ -293,16 +302,25 @@ probo: temperature: "" maxTokens: "" + # Tracker enrichment agent (optional; researches catalog descriptions). + # Falls back to trackerMapping when its provider is unset. + trackerEnrichment: + provider: "" + modelName: "" + temperature: "" + maxTokens: "" + # Tracker mapping background worker tuning (optional). interval, - # staleAfter, and agentTimeout are in seconds. Keep concurrency modest - # to stay under OpenAI/Firecrawl rate limits and the database - # connection pool. + # staleAfter, agentTimeout, and disambiguationAgentTimeout are in + # seconds. Keep concurrency modest to stay under OpenAI/Firecrawl rate + # limits and the database connection pool. trackerMappingWorker: interval: 10 maxConcurrency: 3 staleAfter: 600 agentTimeout: 45 agentMaxTurns: 10 + disambiguationAgentTimeout: 45 # Common-pattern enrichment background worker tuning (optional). # interval, staleAfter, and agentTimeout are in seconds. diff --git a/pkg/agentsbuild/agentsbuild.go b/pkg/agentsbuild/agentsbuild.go deleted file mode 100644 index a36ae8b91..000000000 --- a/pkg/agentsbuild/agentsbuild.go +++ /dev/null @@ -1,172 +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 agentsbuild builds LLM clients and tracker-agent configuration -// from the shared probodconfig types. It is the single wiring used by -// both probod (background workers) and proboctl (synchronous operator -// commands) so the two executables build agents identically. -package agentsbuild - -import ( - "fmt" - "time" - - "github.com/prometheus/client_golang/prometheus" - "go.gearno.de/kit/httpclient" - "go.gearno.de/kit/log" - "go.opentelemetry.io/otel/trace" - "go.probo.inc/probo/pkg/cookiebanner" - "go.probo.inc/probo/pkg/llm" - llmanthropic "go.probo.inc/probo/pkg/llm/anthropic" - llmopenai "go.probo.inc/probo/pkg/llm/openai" - "go.probo.inc/probo/pkg/probodconfig" - "go.probo.inc/probo/pkg/thirdparty" -) - -// BuildLLMClient creates an LLM client for the given provider config. -func BuildLLMClient( - cfg probodconfig.LLMProviderConfig, - l *log.Logger, - tp trace.TracerProvider, - r prometheus.Registerer, -) (*llm.Client, error) { - providerType := cfg.Type - if providerType == "" { - providerType = "openai" - } - - httpClient := httpclient.DefaultPooledClient( - httpclient.WithLogger(l), - httpclient.WithTracerProvider(tp), - httpclient.WithRegisterer(r), - ) - - switch providerType { - case "openai": - p := llmopenai.NewProvider( - cfg.APIKey, - llmopenai.WithHTTPClient(httpClient), - ) - - return llm.NewClient( - p, - "openai", - llm.WithLogger(l), - llm.WithTracerProvider(tp), - ), nil - case "anthropic": - p := llmanthropic.NewProvider( - cfg.APIKey, - llmanthropic.WithHTTPClient(httpClient), - ) - - return llm.NewClient( - p, - "anthropic", - llm.WithLogger(l), - llm.WithTracerProvider(tp), - ), nil - case "bedrock": - return nil, fmt.Errorf("bedrock provider not yet wired; requires aws.Config") - default: - return nil, fmt.Errorf("unsupported LLM provider type: %q", providerType) - } -} - -// ResolveAgentClient resolves an agent's effective config from defaults -// and builds an LLM client for it. The name is used in the logger name -// and error messages. -func ResolveAgentClient( - agents probodconfig.AgentsConfig, - name string, - agent probodconfig.LLMAgentConfig, - l *log.Logger, - tp trace.TracerProvider, - r prometheus.Registerer, -) (probodconfig.LLMAgentConfig, *llm.Client, error) { - resolved := agents.ResolveAgent(agent) - - providerCfg, ok := agents.Providers[resolved.Provider] - if !ok { - return probodconfig.LLMAgentConfig{}, nil, fmt.Errorf("unknown LLM provider %q for %s agent", resolved.Provider, name) - } - - client, err := BuildLLMClient(providerCfg, l.Named("llm."+name), tp, r) - if err != nil { - return probodconfig.LLMAgentConfig{}, nil, fmt.Errorf("cannot create %s LLM client: %w", name, err) - } - - return resolved, client, nil -} - -// BuildTrackerAgentsConfig wires the tracker agents that share one LLM -// client and model: the tracker-mapping agent (catalog identification), -// the common-pattern enrichment agent (description research), and the -// third-party disambiguation agent. All are opt-in: when -// `llm.tracker-mapping.provider` is empty it returns zero configs (nil -// LLM client) so callers run without agent fallback. -func BuildTrackerAgentsConfig( - cfg probodconfig.Config, - l *log.Logger, - tp trace.TracerProvider, - r prometheus.Registerer, -) (cookiebanner.TrackerAgentsConfig, thirdparty.DisambiguationConfig, error) { - if cfg.Agents.TrackerMapping.Provider == "" { - return cookiebanner.TrackerAgentsConfig{}, thirdparty.DisambiguationConfig{}, nil - } - - agentCfg, llmClient, err := ResolveAgentClient( - cfg.Agents, - "tracker-mapping", - cfg.Agents.TrackerMapping, - l, - tp, - r, - ) - if err != nil { - return cookiebanner.TrackerAgentsConfig{}, thirdparty.DisambiguationConfig{}, fmt.Errorf("cannot resolve tracker mapping agent client: %w", err) - } - - mappingWorkerCfg := cfg.TrackerMappingWorker - enrichmentWorkerCfg := cfg.CommonPatternEnrichmentWorker - - // The mapping and enrichment agents share one config slot but run - // from separate workers with separate max-turns. AgentTimeout here - // carries the mapping worker's value (also reused by the - // disambiguation agent); the enrichment worker overrides it on its - // own copy at registration. - trackerAgentsCfg := cookiebanner.TrackerAgentsConfig{ - LLMClient: llmClient, - Model: agentCfg.ModelName, - FirecrawlAPIKey: cfg.Agents.Tools.FirecrawlAPIKey, - MaxTokens: agentCfg.MaxTokens, - Temperature: agentCfg.Temperature, - AgentTimeout: time.Duration(mappingWorkerCfg.AgentTimeout) * time.Second, - MappingMaxTurns: mappingWorkerCfg.AgentMaxTurns, - EnrichmentMaxTurns: enrichmentWorkerCfg.AgentMaxTurns, - } - - // The disambiguation agent emits a single id plus a short rationale, - // so it keeps its own smaller token budget (left unset here) rather - // than inheriting the mapping agent's. It shares the mapping worker's - // timeout. - disambiguationCfg := thirdparty.DisambiguationConfig{ - LLMClient: llmClient, - Model: agentCfg.ModelName, - Temperature: agentCfg.Temperature, - Timeout: time.Duration(mappingWorkerCfg.AgentTimeout) * time.Second, - } - - return trackerAgentsCfg, disambiguationCfg, nil -} diff --git a/pkg/bootstrap/builder.go b/pkg/bootstrap/builder.go index cf1d65142..2eac4449d 100644 --- a/pkg/bootstrap/builder.go +++ b/pkg/bootstrap/builder.go @@ -213,6 +213,17 @@ func (b *Builder) Build() (*probodconfig.FullConfig, error) { Temperature: b.getEnvFloatPtr("AGENT_THIRD_PARTY_VETTER_TEMPERATURE"), MaxTokens: b.getEnvIntPtr("AGENT_THIRD_PARTY_VETTER_MAX_TOKENS"), }, + ThirdPartyDisambiguation: probodconfig.LLMAgentConfig{ + Provider: b.getEnvOrDefault("AGENT_THIRD_PARTY_DISAMBIGUATION_PROVIDER", ""), + ModelName: b.getEnvOrDefault("AGENT_THIRD_PARTY_DISAMBIGUATION_MODEL_NAME", ""), + // The disambiguation agent emits a single id plus a + // short rationale, but the budget must leave headroom + // for reasoning models whose reasoning tokens count + // against max_tokens; too small a budget truncates the + // JSON. + Temperature: b.getEnvFloatPtr("AGENT_THIRD_PARTY_DISAMBIGUATION_TEMPERATURE"), + MaxTokens: new(b.getEnvIntOrDefault("AGENT_THIRD_PARTY_DISAMBIGUATION_MAX_TOKENS", 4096)), + }, TrackerMapping: probodconfig.LLMAgentConfig{ Provider: b.getEnvOrDefault("AGENT_TRACKER_MAPPING_PROVIDER", ""), ModelName: b.getEnvOrDefault("AGENT_TRACKER_MAPPING_MODEL_NAME", ""), @@ -223,6 +234,14 @@ func (b *Builder) Build() (*probodconfig.FullConfig, error) { Temperature: b.getEnvFloatPtr("AGENT_TRACKER_MAPPING_TEMPERATURE"), MaxTokens: new(b.getEnvIntOrDefault("AGENT_TRACKER_MAPPING_MAX_TOKENS", 4096)), }, + TrackerEnrichment: probodconfig.LLMAgentConfig{ + Provider: b.getEnvOrDefault("AGENT_TRACKER_ENRICHMENT_PROVIDER", ""), + ModelName: b.getEnvOrDefault("AGENT_TRACKER_ENRICHMENT_MODEL_NAME", ""), + // See the tracker-mapping note: keep ample headroom so + // reasoning models do not truncate the structured JSON. + Temperature: b.getEnvFloatPtr("AGENT_TRACKER_ENRICHMENT_TEMPERATURE"), + MaxTokens: new(b.getEnvIntOrDefault("AGENT_TRACKER_ENRICHMENT_MAX_TOKENS", 4096)), + }, Tools: probodconfig.AgentToolsConfig{ FirecrawlAPIKey: b.getEnv("FIRECRAWL_API_KEY"), }, @@ -259,11 +278,12 @@ func (b *Builder) Build() (*probodconfig.FullConfig, error) { MaxConcurrency: b.getEnvIntOrDefault("THIRD_PARTY_VETTING_MAX_CONCURRENCY", 1), }, TrackerMappingWorker: probodconfig.TrackerMappingWorkerConfig{ - Interval: b.getEnvIntOrDefault("TRACKER_MAPPING_INTERVAL", 10), - MaxConcurrency: b.getEnvIntOrDefault("TRACKER_MAPPING_MAX_CONCURRENCY", 3), - StaleAfter: b.getEnvIntOrDefault("TRACKER_MAPPING_STALE_AFTER", 600), - AgentTimeout: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_TIMEOUT", 45), - AgentMaxTurns: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_MAX_TURNS", 10), + Interval: b.getEnvIntOrDefault("TRACKER_MAPPING_INTERVAL", 10), + MaxConcurrency: b.getEnvIntOrDefault("TRACKER_MAPPING_MAX_CONCURRENCY", 3), + StaleAfter: b.getEnvIntOrDefault("TRACKER_MAPPING_STALE_AFTER", 600), + AgentTimeout: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_TIMEOUT", 45), + AgentMaxTurns: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_MAX_TURNS", 10), + DisambiguationAgentTimeout: b.getEnvIntOrDefault("TRACKER_MAPPING_DISAMBIGUATION_AGENT_TIMEOUT", 45), }, CommonPatternEnrichmentWorker: probodconfig.CommonPatternEnrichmentWorkerConfig{ Interval: b.getEnvIntOrDefault("COMMON_PATTERN_ENRICHMENT_INTERVAL", 10), diff --git a/pkg/bootstrap/builder_test.go b/pkg/bootstrap/builder_test.go index f876ea80e..b8889ee6b 100644 --- a/pkg/bootstrap/builder_test.go +++ b/pkg/bootstrap/builder_test.go @@ -220,10 +220,18 @@ func TestBuilder_Build_Defaults(t *testing.T) { assert.Empty(t, cfg.Probod.Agents.ThirdPartyVetter.ModelName) assert.Nil(t, cfg.Probod.Agents.ThirdPartyVetter.Temperature) assert.Nil(t, cfg.Probod.Agents.ThirdPartyVetter.MaxTokens) + assert.Empty(t, cfg.Probod.Agents.ThirdPartyDisambiguation.Provider) + assert.Empty(t, cfg.Probod.Agents.ThirdPartyDisambiguation.ModelName) + assert.Nil(t, cfg.Probod.Agents.ThirdPartyDisambiguation.Temperature) + assert.Equal(t, new(4096), cfg.Probod.Agents.ThirdPartyDisambiguation.MaxTokens) assert.Empty(t, cfg.Probod.Agents.TrackerMapping.Provider) assert.Empty(t, cfg.Probod.Agents.TrackerMapping.ModelName) assert.Nil(t, cfg.Probod.Agents.TrackerMapping.Temperature) assert.Equal(t, new(4096), cfg.Probod.Agents.TrackerMapping.MaxTokens) + assert.Empty(t, cfg.Probod.Agents.TrackerEnrichment.Provider) + assert.Empty(t, cfg.Probod.Agents.TrackerEnrichment.ModelName) + assert.Nil(t, cfg.Probod.Agents.TrackerEnrichment.Temperature) + assert.Equal(t, new(4096), cfg.Probod.Agents.TrackerEnrichment.MaxTokens) // Tracker worker tuning — defaults assert.Equal(t, 10, cfg.Probod.TrackerMappingWorker.Interval) @@ -231,6 +239,7 @@ func TestBuilder_Build_Defaults(t *testing.T) { assert.Equal(t, 600, cfg.Probod.TrackerMappingWorker.StaleAfter) assert.Equal(t, 45, cfg.Probod.TrackerMappingWorker.AgentTimeout) assert.Equal(t, 10, cfg.Probod.TrackerMappingWorker.AgentMaxTurns) + assert.Equal(t, 45, cfg.Probod.TrackerMappingWorker.DisambiguationAgentTimeout) assert.Equal(t, 10, cfg.Probod.CommonPatternEnrichmentWorker.Interval) assert.Equal(t, 2, cfg.Probod.CommonPatternEnrichmentWorker.MaxConcurrency) assert.Equal(t, 600, cfg.Probod.CommonPatternEnrichmentWorker.StaleAfter) @@ -333,17 +342,28 @@ func TestBuilder_Build_CustomValues(t *testing.T) { env["AGENT_THIRD_PARTY_VETTER_MODEL_NAME"] = "gpt-4o" env["AGENT_THIRD_PARTY_VETTER_TEMPERATURE"] = "0.3" env["AGENT_THIRD_PARTY_VETTER_MAX_TOKENS"] = "8192" + // Agents — third-party-disambiguation override + env["AGENT_THIRD_PARTY_DISAMBIGUATION_PROVIDER"] = "anthropic" + env["AGENT_THIRD_PARTY_DISAMBIGUATION_MODEL_NAME"] = "claude-sonnet-4-20250514" + env["AGENT_THIRD_PARTY_DISAMBIGUATION_TEMPERATURE"] = "0.4" + env["AGENT_THIRD_PARTY_DISAMBIGUATION_MAX_TOKENS"] = "2048" // Agents — tracker-mapping override env["AGENT_TRACKER_MAPPING_PROVIDER"] = "openai" env["AGENT_TRACKER_MAPPING_MODEL_NAME"] = "gpt-4o-mini" env["AGENT_TRACKER_MAPPING_TEMPERATURE"] = "0.1" env["AGENT_TRACKER_MAPPING_MAX_TOKENS"] = "1024" + // Agents — tracker-enrichment override + env["AGENT_TRACKER_ENRICHMENT_PROVIDER"] = "openai" + env["AGENT_TRACKER_ENRICHMENT_MODEL_NAME"] = "gpt-4o" + env["AGENT_TRACKER_ENRICHMENT_TEMPERATURE"] = "0.2" + env["AGENT_TRACKER_ENRICHMENT_MAX_TOKENS"] = "2048" // Tracker worker tuning override env["TRACKER_MAPPING_INTERVAL"] = "20" env["TRACKER_MAPPING_MAX_CONCURRENCY"] = "5" env["TRACKER_MAPPING_STALE_AFTER"] = "1200" env["TRACKER_MAPPING_AGENT_TIMEOUT"] = "30" env["TRACKER_MAPPING_AGENT_MAX_TURNS"] = "6" + env["TRACKER_MAPPING_DISAMBIGUATION_AGENT_TIMEOUT"] = "35" env["COMMON_PATTERN_ENRICHMENT_INTERVAL"] = "15" env["COMMON_PATTERN_ENRICHMENT_MAX_CONCURRENCY"] = "4" env["COMMON_PATTERN_ENRICHMENT_STALE_AFTER"] = "900" @@ -443,17 +463,28 @@ func TestBuilder_Build_CustomValues(t *testing.T) { assert.Equal(t, "gpt-4o", cfg.Probod.Agents.ThirdPartyVetter.ModelName) assert.Equal(t, new(0.3), cfg.Probod.Agents.ThirdPartyVetter.Temperature) assert.Equal(t, new(8192), cfg.Probod.Agents.ThirdPartyVetter.MaxTokens) + // Agents — third-party-disambiguation overrides + assert.Equal(t, "anthropic", cfg.Probod.Agents.ThirdPartyDisambiguation.Provider) + assert.Equal(t, "claude-sonnet-4-20250514", cfg.Probod.Agents.ThirdPartyDisambiguation.ModelName) + assert.Equal(t, new(0.4), cfg.Probod.Agents.ThirdPartyDisambiguation.Temperature) + assert.Equal(t, new(2048), cfg.Probod.Agents.ThirdPartyDisambiguation.MaxTokens) // Agents — tracker-mapping overrides assert.Equal(t, "openai", cfg.Probod.Agents.TrackerMapping.Provider) assert.Equal(t, "gpt-4o-mini", cfg.Probod.Agents.TrackerMapping.ModelName) assert.Equal(t, new(0.1), cfg.Probod.Agents.TrackerMapping.Temperature) assert.Equal(t, new(1024), cfg.Probod.Agents.TrackerMapping.MaxTokens) + // Agents — tracker-enrichment overrides + assert.Equal(t, "openai", cfg.Probod.Agents.TrackerEnrichment.Provider) + assert.Equal(t, "gpt-4o", cfg.Probod.Agents.TrackerEnrichment.ModelName) + assert.Equal(t, new(0.2), cfg.Probod.Agents.TrackerEnrichment.Temperature) + assert.Equal(t, new(2048), cfg.Probod.Agents.TrackerEnrichment.MaxTokens) // Tracker worker tuning — overrides assert.Equal(t, 20, cfg.Probod.TrackerMappingWorker.Interval) assert.Equal(t, 5, cfg.Probod.TrackerMappingWorker.MaxConcurrency) assert.Equal(t, 1200, cfg.Probod.TrackerMappingWorker.StaleAfter) assert.Equal(t, 30, cfg.Probod.TrackerMappingWorker.AgentTimeout) assert.Equal(t, 6, cfg.Probod.TrackerMappingWorker.AgentMaxTurns) + assert.Equal(t, 35, cfg.Probod.TrackerMappingWorker.DisambiguationAgentTimeout) assert.Equal(t, 15, cfg.Probod.CommonPatternEnrichmentWorker.Interval) assert.Equal(t, 4, cfg.Probod.CommonPatternEnrichmentWorker.MaxConcurrency) assert.Equal(t, 900, cfg.Probod.CommonPatternEnrichmentWorker.StaleAfter) diff --git a/pkg/cookiebanner/common_pattern_enricher.go b/pkg/cookiebanner/common_pattern_enricher.go index b64ec320a..2ccd17756 100644 --- a/pkg/cookiebanner/common_pattern_enricher.go +++ b/pkg/cookiebanner/common_pattern_enricher.go @@ -39,35 +39,49 @@ import ( // known set of ids (e.g. proboctl). It is enrichment's single source of // truth - the worker is a thin queue poller that delegates here. type CommonPatternEnricher struct { - pg *pg.Client - logger *log.Logger - enrichmentAgent *agent.Agent - mappingAgent *agent.Agent - agentTimeout time.Duration + pg *pg.Client + logger *log.Logger + enrichmentAgent *agent.Agent + mappingAgent *agent.Agent + enrichmentTimeout time.Duration + mappingTimeout time.Duration } -// NewCommonPatternEnricher builds the enricher from the shared tracker -// agents config. When no LLM client is configured the agents are left nil -// and Enabled reports false; callers must gate on Enabled before running. +// NewCommonPatternEnricher builds the enricher from the enrichment and +// mapping agent configs. It runs the enrichment agent to research a +// description and reuses the mapping agent to attribute a vendor first, +// so it needs both configs. When the enrichment config has no LLM client +// the agents are left nil and Enabled reports false; callers must gate on +// Enabled before running. func NewCommonPatternEnricher( pgClient *pg.Client, logger *log.Logger, - cfg TrackerAgentsConfig, + enrichmentCfg TrackerEnrichmentAgentConfig, + mappingCfg TrackerMappingAgentConfig, ) *CommonPatternEnricher { - agentTimeout := cfg.AgentTimeout - if agentTimeout <= 0 { - agentTimeout = defaultAgentTimeout + enrichmentTimeout := enrichmentCfg.Timeout + if enrichmentTimeout <= 0 { + enrichmentTimeout = defaultAgentTimeout + } + + mappingTimeout := mappingCfg.Timeout + if mappingTimeout <= 0 { + mappingTimeout = defaultAgentTimeout } e := &CommonPatternEnricher{ - pg: pgClient, - logger: logger, - agentTimeout: agentTimeout, + pg: pgClient, + logger: logger, + enrichmentTimeout: enrichmentTimeout, + mappingTimeout: mappingTimeout, } - if cfg.LLMClient != nil { - e.enrichmentAgent = buildCommonPatternEnrichmentAgent(cfg, pgClient, logger) - e.mappingAgent = buildTrackerMappingAgent(cfg, pgClient, logger) + if enrichmentCfg.LLMClient != nil { + e.enrichmentAgent = buildCommonPatternEnrichmentAgent(enrichmentCfg, pgClient, logger) + } + + if mappingCfg.LLMClient != nil { + e.mappingAgent = buildTrackerMappingAgent(mappingCfg, pgClient, logger) } return e @@ -251,7 +265,7 @@ func (e *CommonPatternEnricher) research( ) (string, error) { prompt := buildEnrichmentPrompt(cp, thirdPartyName) - agentCtx, cancel := context.WithTimeout(ctx, e.agentTimeout) + agentCtx, cancel := context.WithTimeout(ctx, e.enrichmentTimeout) defer cancel() result, err := agent.RunTyped[CommonPatternEnrichmentResult]( @@ -287,7 +301,7 @@ func (e *CommonPatternEnricher) identifyThirdParty( prompt := buildCommonPatternIdentificationPrompt(cp) - agentCtx, cancel := context.WithTimeout(ctx, e.agentTimeout) + agentCtx, cancel := context.WithTimeout(ctx, e.mappingTimeout) defer cancel() result, err := agent.RunTyped[TrackerMappingAgentResult]( diff --git a/pkg/cookiebanner/common_pattern_enrichment_agent.go b/pkg/cookiebanner/common_pattern_enrichment_agent.go index 0d96e6e11..3c61c2c3b 100644 --- a/pkg/cookiebanner/common_pattern_enrichment_agent.go +++ b/pkg/cookiebanner/common_pattern_enrichment_agent.go @@ -36,7 +36,7 @@ type CommonPatternEnrichmentResult struct { } func buildCommonPatternEnrichmentAgent( - cfg TrackerAgentsConfig, + cfg TrackerEnrichmentAgentConfig, pgClient *pg.Client, logger *log.Logger, ) *agent.Agent { @@ -53,7 +53,7 @@ func buildCommonPatternEnrichmentAgent( panic(fmt.Sprintf("cookiebanner: cannot build tracker enrichment output type: %s", err)) } - maxTurns := cfg.EnrichmentMaxTurns + maxTurns := cfg.MaxTurns if maxTurns < 1 { maxTurns = defaultEnrichmentMaxTurns } diff --git a/pkg/cookiebanner/common_pattern_enrichment_worker.go b/pkg/cookiebanner/common_pattern_enrichment_worker.go index e31b4bc24..d65b15d88 100644 --- a/pkg/cookiebanner/common_pattern_enrichment_worker.go +++ b/pkg/cookiebanner/common_pattern_enrichment_worker.go @@ -49,7 +49,8 @@ type commonPatternEnrichmentHandler struct { func NewCommonPatternEnrichmentWorker( pgClient *pg.Client, logger *log.Logger, - cfg TrackerAgentsConfig, + enrichmentCfg TrackerEnrichmentAgentConfig, + mappingCfg TrackerMappingAgentConfig, staleAfter time.Duration, opts ...worker.Option, ) *worker.Worker[coredata.CommonTrackerPattern] { @@ -60,7 +61,7 @@ func NewCommonPatternEnrichmentWorker( h := &commonPatternEnrichmentHandler{ pg: pgClient, logger: logger, - enricher: NewCommonPatternEnricher(pgClient, logger, cfg), + enricher: NewCommonPatternEnricher(pgClient, logger, enrichmentCfg, mappingCfg), staleAfter: staleAfter, } diff --git a/pkg/cookiebanner/tracker_agents_config.go b/pkg/cookiebanner/tracker_agents_config.go index 31b32ec51..22e00f582 100644 --- a/pkg/cookiebanner/tracker_agents_config.go +++ b/pkg/cookiebanner/tracker_agents_config.go @@ -20,23 +20,38 @@ import ( "go.probo.inc/probo/pkg/llm" ) -// TrackerAgentsConfig configures the tracker agents that share one LLM -// client, model, and tool surface: the tracker-mapping agent (catalog -// identification) and the common-pattern enrichment agent (description -// research). Both use DB-backed search tools and may also use Firecrawl -// for web search when an API key is supplied. +// TrackerMappingAgentConfig configures the tracker-mapping agent +// (catalog identification). It uses DB-backed search tools and may also +// use Firecrawl for web search when an API key is supplied. // -// MaxTokens and Temperature bound and steer each LLM call (both -// outputs are tiny structured JSON). AgentTimeout caps a single agent -// run, and the per-worker max-turns bound the agent reasoning loop. -// Zero-valued tuning fields fall back to package defaults. -type TrackerAgentsConfig struct { - LLMClient *llm.Client - Model string - FirecrawlAPIKey string - MaxTokens *int - Temperature *float64 - AgentTimeout time.Duration - MappingMaxTurns int - EnrichmentMaxTurns int +// MaxTokens and Temperature bound and steer the LLM call (the output is +// tiny structured JSON). Timeout caps a single agent run and MaxTurns +// bounds the agent reasoning loop. Zero-valued tuning fields fall back +// to package defaults. +type TrackerMappingAgentConfig struct { + LLMClient *llm.Client + Model string + FirecrawlAPIKey string + MaxTokens *int + Temperature *float64 + Timeout time.Duration + MaxTurns int +} + +// TrackerEnrichmentAgentConfig configures the common-pattern enrichment +// agent (description research). It uses DB-backed search tools and may +// also use Firecrawl for web search when an API key is supplied. +// +// MaxTokens and Temperature bound and steer the LLM call (the output is +// tiny structured JSON). Timeout caps a single agent run and MaxTurns +// bounds the agent reasoning loop. Zero-valued tuning fields fall back +// to package defaults. +type TrackerEnrichmentAgentConfig struct { + LLMClient *llm.Client + Model string + FirecrawlAPIKey string + MaxTokens *int + Temperature *float64 + Timeout time.Duration + MaxTurns int } diff --git a/pkg/cookiebanner/tracker_mapping_agent.go b/pkg/cookiebanner/tracker_mapping_agent.go index efbab2a45..2873c38c0 100644 --- a/pkg/cookiebanner/tracker_mapping_agent.go +++ b/pkg/cookiebanner/tracker_mapping_agent.go @@ -76,7 +76,7 @@ type TrackerMappingAgentResult struct { } func buildTrackerMappingAgent( - cfg TrackerAgentsConfig, + cfg TrackerMappingAgentConfig, pgClient *pg.Client, logger *log.Logger, ) *agent.Agent { @@ -94,7 +94,7 @@ func buildTrackerMappingAgent( panic(fmt.Sprintf("cookiebanner: cannot build tracker identification output type: %s", err)) } - maxTurns := cfg.MappingMaxTurns + maxTurns := cfg.MaxTurns if maxTurns < 1 { maxTurns = defaultMappingMaxTurns } diff --git a/pkg/cookiebanner/tracker_mapping_worker.go b/pkg/cookiebanner/tracker_mapping_worker.go index f75c46a5e..66b67e967 100644 --- a/pkg/cookiebanner/tracker_mapping_worker.go +++ b/pkg/cookiebanner/tracker_mapping_worker.go @@ -50,12 +50,12 @@ type trackerMappingHandler struct { func NewTrackerMappingWorker( pgClient *pg.Client, logger *log.Logger, - mappingCfg TrackerAgentsConfig, - disambiguationCfg thirdparty.DisambiguationConfig, + mappingCfg TrackerMappingAgentConfig, + disambiguationCfg thirdparty.DisambiguationAgentConfig, staleAfter time.Duration, opts ...worker.Option, ) *worker.Worker[coredata.TrackerPattern] { - agentTimeout := mappingCfg.AgentTimeout + agentTimeout := mappingCfg.Timeout if agentTimeout <= 0 { agentTimeout = defaultAgentTimeout } diff --git a/pkg/proboctl/cmdutil/cmdutil.go b/pkg/proboctl/cmdutil/cmdutil.go index 3a8d64138..2b54df6b6 100644 --- a/pkg/proboctl/cmdutil/cmdutil.go +++ b/pkg/proboctl/cmdutil/cmdutil.go @@ -17,12 +17,10 @@ package cmdutil import ( "fmt" "os" + "time" - "github.com/prometheus/client_golang/prometheus" "go.gearno.de/kit/log" "go.gearno.de/kit/pg" - "go.opentelemetry.io/otel/trace/noop" - "go.probo.inc/probo/pkg/agentsbuild" "go.probo.inc/probo/pkg/cmd/iostreams" "go.probo.inc/probo/pkg/cookiebanner" "go.probo.inc/probo/pkg/proboctl/pgconn" @@ -66,14 +64,23 @@ func (f *Factory) ProbodConfig() (probodconfig.Config, error) { return full.Probod, nil } -// TrackerAgentsConfig builds the tracker-agents config (LLM client + -// Firecrawl key) from the shared probod config for in-process agent -// execution, e.g. synchronous common-pattern re-enrichment. It errors -// when no LLM provider is configured. -func (f *Factory) TrackerAgentsConfig() (cookiebanner.TrackerAgentsConfig, error) { +// TrackerAgentsConfig builds the enrichment and mapping agent configs +// (LLM clients + Firecrawl key) from the shared probod config for +// in-process agent execution, e.g. synchronous common-pattern +// re-enrichment. The enricher runs the enrichment agent and reuses the +// mapping agent to attribute a vendor first, so both configs are +// returned. It errors when no LLM provider is configured. +// +// This duplicates the small wiring probod does in buildTrackerAgents +// rather than sharing a package, keeping the two executables decoupled. +func (f *Factory) TrackerAgentsConfig() (cookiebanner.TrackerEnrichmentAgentConfig, cookiebanner.TrackerMappingAgentConfig, error) { cfg, err := f.ProbodConfig() if err != nil { - return cookiebanner.TrackerAgentsConfig{}, err + return cookiebanner.TrackerEnrichmentAgentConfig{}, cookiebanner.TrackerMappingAgentConfig{}, err + } + + if cfg.Agents.TrackerMapping.Provider == "" { + return cookiebanner.TrackerEnrichmentAgentConfig{}, cookiebanner.TrackerMappingAgentConfig{}, fmt.Errorf("no LLM provider configured; set llm.tracker-mapping.provider in %q", f.CfgFile) } logger := log.NewLogger( @@ -81,19 +88,42 @@ func (f *Factory) TrackerAgentsConfig() (cookiebanner.TrackerAgentsConfig, error log.WithOutput(f.IOStreams.ErrOut), ) - trackerCfg, _, err := agentsbuild.BuildTrackerAgentsConfig( - cfg, - logger, - noop.NewTracerProvider(), - prometheus.NewRegistry(), - ) + firecrawlAPIKey := cfg.Agents.Tools.FirecrawlAPIKey + + mappingAgentCfg, mappingClient, err := resolveAgentClient(cfg.Agents, "tracker-mapping", cfg.Agents.TrackerMapping, logger) if err != nil { - return cookiebanner.TrackerAgentsConfig{}, fmt.Errorf("cannot build tracker agents config: %w", err) + return cookiebanner.TrackerEnrichmentAgentConfig{}, cookiebanner.TrackerMappingAgentConfig{}, fmt.Errorf("cannot build tracker mapping agent: %w", err) } - if trackerCfg.LLMClient == nil { - return cookiebanner.TrackerAgentsConfig{}, fmt.Errorf("no LLM provider configured; set llm.tracker-mapping.provider in %q", f.CfgFile) + mappingCfg := cookiebanner.TrackerMappingAgentConfig{ + LLMClient: mappingClient, + Model: mappingAgentCfg.ModelName, + FirecrawlAPIKey: firecrawlAPIKey, + MaxTokens: mappingAgentCfg.MaxTokens, + Temperature: mappingAgentCfg.Temperature, + Timeout: time.Duration(cfg.TrackerMappingWorker.AgentTimeout) * time.Second, + MaxTurns: cfg.TrackerMappingWorker.AgentMaxTurns, } - return trackerCfg, nil + enrichmentSlot := cfg.Agents.TrackerEnrichment + if enrichmentSlot.Provider == "" { + enrichmentSlot = cfg.Agents.TrackerMapping + } + + enrichmentAgentCfg, enrichmentClient, err := resolveAgentClient(cfg.Agents, "tracker-enrichment", enrichmentSlot, logger) + if err != nil { + return cookiebanner.TrackerEnrichmentAgentConfig{}, cookiebanner.TrackerMappingAgentConfig{}, fmt.Errorf("cannot build tracker enrichment agent: %w", err) + } + + enrichmentCfg := cookiebanner.TrackerEnrichmentAgentConfig{ + LLMClient: enrichmentClient, + Model: enrichmentAgentCfg.ModelName, + FirecrawlAPIKey: firecrawlAPIKey, + MaxTokens: enrichmentAgentCfg.MaxTokens, + Temperature: enrichmentAgentCfg.Temperature, + Timeout: time.Duration(cfg.CommonPatternEnrichmentWorker.AgentTimeout) * time.Second, + MaxTurns: cfg.CommonPatternEnrichmentWorker.AgentMaxTurns, + } + + return enrichmentCfg, mappingCfg, nil } diff --git a/pkg/proboctl/cmdutil/llm.go b/pkg/proboctl/cmdutil/llm.go new file mode 100644 index 000000000..f6096c481 --- /dev/null +++ b/pkg/proboctl/cmdutil/llm.go @@ -0,0 +1,97 @@ +// 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 cmdutil + +import ( + "fmt" + + "github.com/prometheus/client_golang/prometheus" + "go.gearno.de/kit/httpclient" + "go.gearno.de/kit/log" + "go.opentelemetry.io/otel/trace/noop" + "go.probo.inc/probo/pkg/llm" + llmanthropic "go.probo.inc/probo/pkg/llm/anthropic" + llmopenai "go.probo.inc/probo/pkg/llm/openai" + "go.probo.inc/probo/pkg/probodconfig" +) + +// resolveAgentClient resolves an agent's effective config from defaults +// and builds an LLM client for it. The name is used in the logger name +// and error messages. proboctl runs synchronously with no tracing or +// metrics, so it duplicates probod's wiring with a no-op tracer and a +// throwaway registry rather than sharing a package. +func resolveAgentClient( + agents probodconfig.AgentsConfig, + name string, + agent probodconfig.LLMAgentConfig, + l *log.Logger, +) (probodconfig.LLMAgentConfig, *llm.Client, error) { + resolved := agents.ResolveAgent(agent) + + providerCfg, ok := agents.Providers[resolved.Provider] + if !ok { + return probodconfig.LLMAgentConfig{}, nil, fmt.Errorf("unknown LLM provider %q for %s agent", resolved.Provider, name) + } + + client, err := buildLLMClient(providerCfg, l.Named("llm."+name)) + if err != nil { + return probodconfig.LLMAgentConfig{}, nil, fmt.Errorf("cannot create %s LLM client: %w", name, err) + } + + return resolved, client, nil +} + +// buildLLMClient creates an LLM client for the given provider config. +func buildLLMClient(cfg probodconfig.LLMProviderConfig, l *log.Logger) (*llm.Client, error) { + providerType := cfg.Type + if providerType == "" { + providerType = "openai" + } + + httpClient := httpclient.DefaultPooledClient( + httpclient.WithLogger(l), + httpclient.WithTracerProvider(noop.NewTracerProvider()), + httpclient.WithRegisterer(prometheus.NewRegistry()), + ) + + switch providerType { + case "openai": + p := llmopenai.NewProvider( + cfg.APIKey, + llmopenai.WithHTTPClient(httpClient), + ) + + return llm.NewClient( + p, + "openai", + llm.WithLogger(l), + ), nil + case "anthropic": + p := llmanthropic.NewProvider( + cfg.APIKey, + llmanthropic.WithHTTPClient(httpClient), + ) + + return llm.NewClient( + p, + "anthropic", + llm.WithLogger(l), + ), nil + case "bedrock": + return nil, fmt.Errorf("bedrock provider not yet wired; requires aws.Config") + default: + return nil, fmt.Errorf("unsupported LLM provider type: %q", providerType) + } +} diff --git a/pkg/proboctl/commontrackerpattern/reenrich.go b/pkg/proboctl/commontrackerpattern/reenrich.go index 305cca5c1..bade09859 100644 --- a/pkg/proboctl/commontrackerpattern/reenrich.go +++ b/pkg/proboctl/commontrackerpattern/reenrich.go @@ -132,7 +132,7 @@ func newCmdReenrich(f *cmdutil.Factory) *cobra.Command { return nil } - cfg, err := f.TrackerAgentsConfig() + enrichmentCfg, mappingCfg, err := f.TrackerAgentsConfig() if err != nil { return err } @@ -142,7 +142,7 @@ func newCmdReenrich(f *cmdutil.Factory) *cobra.Command { log.WithOutput(f.IOStreams.ErrOut), ) - enricher := cookiebanner.NewCommonPatternEnricher(pgClient, logger, cfg) + enricher := cookiebanner.NewCommonPatternEnricher(pgClient, logger, enrichmentCfg, mappingCfg) enriched, err := enricher.EnrichByIDs(ctx, ids, flagConcurrency) diff --git a/pkg/probod/llm.go b/pkg/probod/llm.go index 003e204fe..5ecf827b6 100644 --- a/pkg/probod/llm.go +++ b/pkg/probod/llm.go @@ -15,17 +15,70 @@ package probod import ( + "fmt" + "github.com/prometheus/client_golang/prometheus" + "go.gearno.de/kit/httpclient" "go.gearno.de/kit/log" "go.opentelemetry.io/otel/trace" - "go.probo.inc/probo/pkg/agentsbuild" "go.probo.inc/probo/pkg/llm" + llmanthropic "go.probo.inc/probo/pkg/llm/anthropic" + llmopenai "go.probo.inc/probo/pkg/llm/openai" ) +// buildLLMClient creates an LLM client for the given provider config. +func (impl *Implm) buildLLMClient( + cfg LLMProviderConfig, + l *log.Logger, + tp trace.TracerProvider, + r prometheus.Registerer, +) (*llm.Client, error) { + providerType := cfg.Type + if providerType == "" { + providerType = "openai" + } + + httpClient := httpclient.DefaultPooledClient( + httpclient.WithLogger(l), + httpclient.WithTracerProvider(tp), + httpclient.WithRegisterer(r), + ) + + switch providerType { + case "openai": + p := llmopenai.NewProvider( + cfg.APIKey, + llmopenai.WithHTTPClient(httpClient), + ) + + return llm.NewClient( + p, + "openai", + llm.WithLogger(l), + llm.WithTracerProvider(tp), + ), nil + case "anthropic": + p := llmanthropic.NewProvider( + cfg.APIKey, + llmanthropic.WithHTTPClient(httpClient), + ) + + return llm.NewClient( + p, + "anthropic", + llm.WithLogger(l), + llm.WithTracerProvider(tp), + ), nil + case "bedrock": + return nil, fmt.Errorf("bedrock provider not yet wired; requires aws.Config") + default: + return nil, fmt.Errorf("unsupported LLM provider type: %q", providerType) + } +} + // resolveAgentClient resolves the agent's effective config from defaults and // builds an LLM client for it. The name parameter is used in the logger and -// in error messages. It delegates to pkg/agentsbuild so probod and proboctl -// wire LLM clients identically. +// in error messages. func (impl *Implm) resolveAgentClient( name string, agent LLMAgentConfig, @@ -33,5 +86,17 @@ func (impl *Implm) resolveAgentClient( tp trace.TracerProvider, r prometheus.Registerer, ) (LLMAgentConfig, *llm.Client, error) { - return agentsbuild.ResolveAgentClient(impl.cfg.Agents, name, agent, l, tp, r) + resolved := impl.cfg.Agents.ResolveAgent(agent) + + providerCfg, ok := impl.cfg.Agents.Providers[resolved.Provider] + if !ok { + return LLMAgentConfig{}, nil, fmt.Errorf("unknown LLM provider %q for %s agent", resolved.Provider, name) + } + + client, err := impl.buildLLMClient(providerCfg, l.Named("llm."+name), tp, r) + if err != nil { + return LLMAgentConfig{}, nil, fmt.Errorf("cannot create %s LLM client: %w", name, err) + } + + return resolved, client, nil } diff --git a/pkg/probod/probod.go b/pkg/probod/probod.go index d16f17e3e..4435050d3 100644 --- a/pkg/probod/probod.go +++ b/pkg/probod/probod.go @@ -324,7 +324,7 @@ func (impl *Implm) Run( return err } - trackerAgentsCfg, thirdPartyDisambiguationCfg, err := impl.buildTrackerAgentsConfig(l, tp, r) + trackerMappingCfg, trackerEnrichmentCfg, thirdPartyDisambiguationCfg, err := impl.buildTrackerAgents(l, tp, r) if err != nil { return err } @@ -790,7 +790,7 @@ func (impl *Implm) Run( trackerMappingWorker := cookiebanner.NewTrackerMappingWorker( pgClient, l, - trackerAgentsCfg, + trackerMappingCfg, thirdPartyDisambiguationCfg, time.Duration(impl.cfg.TrackerMappingWorker.StaleAfter)*time.Second, worker.WithInterval(time.Duration(impl.cfg.TrackerMappingWorker.Interval)*time.Second), @@ -811,14 +811,12 @@ func (impl *Implm) Run( // the tracker agents are configured. stopCommonPatternEnrichmentWorker := func() {} - if trackerAgentsCfg.LLMClient != nil { - enrichmentCfg := trackerAgentsCfg - enrichmentCfg.AgentTimeout = time.Duration(impl.cfg.CommonPatternEnrichmentWorker.AgentTimeout) * time.Second - + if trackerEnrichmentCfg.LLMClient != nil { commonPatternEnrichmentWorker := cookiebanner.NewCommonPatternEnrichmentWorker( pgClient, l, - enrichmentCfg, + trackerEnrichmentCfg, + trackerMappingCfg, time.Duration(impl.cfg.CommonPatternEnrichmentWorker.StaleAfter)*time.Second, worker.WithInterval(time.Duration(impl.cfg.CommonPatternEnrichmentWorker.Interval)*time.Second), worker.WithMaxConcurrency(impl.cfg.CommonPatternEnrichmentWorker.MaxConcurrency), diff --git a/pkg/probod/tracker_agents.go b/pkg/probod/tracker_agents.go index f93011d76..ce8b1ab63 100644 --- a/pkg/probod/tracker_agents.go +++ b/pkg/probod/tracker_agents.go @@ -15,22 +15,108 @@ package probod import ( + "fmt" + "time" + "github.com/prometheus/client_golang/prometheus" "go.gearno.de/kit/log" "go.opentelemetry.io/otel/trace" - "go.probo.inc/probo/pkg/agentsbuild" "go.probo.inc/probo/pkg/cookiebanner" "go.probo.inc/probo/pkg/thirdparty" ) -// buildTrackerAgentsConfig wires the tracker agents (mapping, enrichment, -// disambiguation) from the probod config. It delegates to pkg/agentsbuild -// so probod and proboctl build the same agent configuration; see that -// package for the wiring rationale. -func (impl *Implm) buildTrackerAgentsConfig( +// buildTrackerAgents wires the three tracker agents from the probod +// config, each with its own LLM client and tuning: the tracker-mapping +// agent (catalog identification), the common-pattern enrichment agent +// (description research), and the third-party disambiguation agent. All +// are opt-in: when `llm.tracker-mapping.provider` is empty it returns +// zero configs (nil LLM clients) so callers run without agent fallback. +// +// The enrichment and disambiguation agents fall back to the +// tracker-mapping config when their own provider slot is empty, so a +// deployment that configures only `tracker-mapping` keeps wiring all +// three agents. +func (impl *Implm) buildTrackerAgents( l *log.Logger, tp trace.TracerProvider, r prometheus.Registerer, -) (cookiebanner.TrackerAgentsConfig, thirdparty.DisambiguationConfig, error) { - return agentsbuild.BuildTrackerAgentsConfig(impl.cfg, l, tp, r) +) (cookiebanner.TrackerMappingAgentConfig, cookiebanner.TrackerEnrichmentAgentConfig, thirdparty.DisambiguationAgentConfig, error) { + if impl.cfg.Agents.TrackerMapping.Provider == "" { + return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, thirdparty.DisambiguationAgentConfig{}, nil + } + + firecrawlAPIKey := impl.cfg.Agents.Tools.FirecrawlAPIKey + + mappingAgentCfg, mappingClient, err := impl.resolveAgentClient( + "tracker-mapping", + impl.cfg.Agents.TrackerMapping, + l, + tp, + r, + ) + if err != nil { + return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, thirdparty.DisambiguationAgentConfig{}, fmt.Errorf("cannot resolve tracker mapping agent client: %w", err) + } + + mappingCfg := cookiebanner.TrackerMappingAgentConfig{ + LLMClient: mappingClient, + Model: mappingAgentCfg.ModelName, + FirecrawlAPIKey: firecrawlAPIKey, + MaxTokens: mappingAgentCfg.MaxTokens, + Temperature: mappingAgentCfg.Temperature, + Timeout: time.Duration(impl.cfg.TrackerMappingWorker.AgentTimeout) * time.Second, + MaxTurns: impl.cfg.TrackerMappingWorker.AgentMaxTurns, + } + + enrichmentSlot := impl.cfg.Agents.TrackerEnrichment + if enrichmentSlot.Provider == "" { + enrichmentSlot = impl.cfg.Agents.TrackerMapping + } + + enrichmentAgentCfg, enrichmentClient, err := impl.resolveAgentClient( + "tracker-enrichment", + enrichmentSlot, + l, + tp, + r, + ) + if err != nil { + return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, thirdparty.DisambiguationAgentConfig{}, fmt.Errorf("cannot resolve tracker enrichment agent client: %w", err) + } + + enrichmentCfg := cookiebanner.TrackerEnrichmentAgentConfig{ + LLMClient: enrichmentClient, + Model: enrichmentAgentCfg.ModelName, + FirecrawlAPIKey: firecrawlAPIKey, + MaxTokens: enrichmentAgentCfg.MaxTokens, + Temperature: enrichmentAgentCfg.Temperature, + Timeout: time.Duration(impl.cfg.CommonPatternEnrichmentWorker.AgentTimeout) * time.Second, + MaxTurns: impl.cfg.CommonPatternEnrichmentWorker.AgentMaxTurns, + } + + disambiguationSlot := impl.cfg.Agents.ThirdPartyDisambiguation + if disambiguationSlot.Provider == "" { + disambiguationSlot = impl.cfg.Agents.TrackerMapping + } + + disambiguationAgentCfg, disambiguationClient, err := impl.resolveAgentClient( + "third-party-disambiguation", + disambiguationSlot, + l, + tp, + r, + ) + if err != nil { + return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, thirdparty.DisambiguationAgentConfig{}, fmt.Errorf("cannot resolve third party disambiguation agent client: %w", err) + } + + disambiguationCfg := thirdparty.DisambiguationAgentConfig{ + LLMClient: disambiguationClient, + Model: disambiguationAgentCfg.ModelName, + MaxTokens: disambiguationAgentCfg.MaxTokens, + Temperature: disambiguationAgentCfg.Temperature, + Timeout: time.Duration(impl.cfg.TrackerMappingWorker.DisambiguationAgentTimeout) * time.Second, + } + + return mappingCfg, enrichmentCfg, disambiguationCfg, nil } diff --git a/pkg/probodconfig/llm_config.go b/pkg/probodconfig/llm_config.go index d19f9d4aa..589c042a6 100644 --- a/pkg/probodconfig/llm_config.go +++ b/pkg/probodconfig/llm_config.go @@ -50,22 +50,24 @@ type ( } // TrackerMappingWorkerConfig holds worker-side tuning for the - // tracker-mapping background worker. LLM parameters for the agents - // it runs live under AgentsConfig.TrackerMapping. AgentTimeout and - // AgentMaxTurns bound a single agent run (the identification and - // disambiguation agents). + // tracker-mapping background worker. LLM parameters for the mapping + // agent it runs live under AgentsConfig.TrackerMapping. AgentTimeout + // and AgentMaxTurns bound a single mapping agent run. + // DisambiguationAgentTimeout caps a single third-party + // disambiguation agent run; that agent runs inside this worker but + // uses its own LLM parameters from AgentsConfig.ThirdPartyDisambiguation. TrackerMappingWorkerConfig struct { - Interval int `json:"interval"` // seconds between polls - MaxConcurrency int `json:"max-concurrency"` - StaleAfter int `json:"stale-after"` // seconds before a claim is recycled - AgentTimeout int `json:"agent-timeout"` // seconds, single agent run - AgentMaxTurns int `json:"agent-max-turns"` + Interval int `json:"interval"` // seconds between polls + MaxConcurrency int `json:"max-concurrency"` + StaleAfter int `json:"stale-after"` // seconds before a claim is recycled + AgentTimeout int `json:"agent-timeout"` // seconds, single agent run + AgentMaxTurns int `json:"agent-max-turns"` + DisambiguationAgentTimeout int `json:"disambiguation-agent-timeout"` // seconds, single disambiguation run } // CommonPatternEnrichmentWorkerConfig holds worker-side tuning for // the common-pattern enrichment background worker. LLM parameters - // for the enrichment agent live under AgentsConfig.TrackerMapping - // (the agents share one config slot). + // for the enrichment agent live under AgentsConfig.TrackerEnrichment. CommonPatternEnrichmentWorkerConfig struct { Interval int `json:"interval"` // seconds between polls MaxConcurrency int `json:"max-concurrency"` @@ -84,13 +86,15 @@ type ( // settings. Default is used as a fallback when an agent-specific field // is zero-valued. AgentsConfig struct { - Providers map[string]LLMProviderConfig `json:"providers"` - Default LLMAgentConfig `json:"defaults"` - Probo LLMAgentConfig `json:"probo"` - EvidenceDescriber LLMAgentConfig `json:"evidence-describer"` - ThirdPartyVetter LLMAgentConfig `json:"third-party-vetter"` - TrackerMapping LLMAgentConfig `json:"tracker-mapping"` - Tools AgentToolsConfig `json:"tools"` + Providers map[string]LLMProviderConfig `json:"providers"` + Default LLMAgentConfig `json:"defaults"` + Probo LLMAgentConfig `json:"probo"` + EvidenceDescriber LLMAgentConfig `json:"evidence-describer"` + ThirdPartyVetter LLMAgentConfig `json:"third-party-vetter"` + ThirdPartyDisambiguation LLMAgentConfig `json:"third-party-disambiguation"` + TrackerMapping LLMAgentConfig `json:"tracker-mapping"` + TrackerEnrichment LLMAgentConfig `json:"tracker-enrichment"` + Tools AgentToolsConfig `json:"tools"` } ) diff --git a/pkg/thirdparty/disambiguation_agent.go b/pkg/thirdparty/disambiguation_agent.go index 53079808d..e7487c8f7 100644 --- a/pkg/thirdparty/disambiguation_agent.go +++ b/pkg/thirdparty/disambiguation_agent.go @@ -54,7 +54,7 @@ const ( defaultDisambiguationMaxTokens = 4096 ) -// DisambiguationConfig configures the third-party disambiguation +// DisambiguationAgentConfig configures the third-party disambiguation // agent. The agent has no DB tools and no web-search tools: the // candidate list is supplied entirely in the prompt and the agent // only picks among it. @@ -62,7 +62,7 @@ const ( // MaxTokens and Temperature bound and steer the single LLM call, and // Timeout caps a single run. Zero-valued fields fall back to package // defaults. -type DisambiguationConfig struct { +type DisambiguationAgentConfig struct { LLMClient *llm.Client Model string MaxTokens *int @@ -84,7 +84,7 @@ type DisambiguationResult struct { // no tools: the candidate list is supplied in the prompt and the // agent must only choose among it. func BuildDisambiguationAgent( - cfg DisambiguationConfig, + cfg DisambiguationAgentConfig, logger *log.Logger, ) *agent.Agent { outputType, err := agent.NewOutputType[DisambiguationResult]("third_party_disambiguation")