Give each tracker agent its own config
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é <emile@probo.com>
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -1,172 +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 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
|
||||
}
|
||||
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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](
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
97
pkg/proboctl/cmdutil/llm.go
Normal file
97
pkg/proboctl/cmdutil/llm.go
Normal file
@@ -0,0 +1,97 @@
|
||||
// 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 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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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"`
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
6
pkg/thirdparty/disambiguation_agent.go
vendored
6
pkg/thirdparty/disambiguation_agent.go
vendored
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user