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

Remove sync common tracker pattern reenriching

Browse Source 
Signed-off-by: Émile Ré <emile@probo.com>
This commit is contained in:
Émile Ré
2026-06-09 11:47:18 +02:00
parent 83d36d345f
commit 4f6fcb42f9
5 changed files with 15 additions and 297 deletions
Download Patch File Download Diff File Expand all files Collapse all files

92
pkg/proboctl/cmdutil/cmdutil.go
View File

@@ -16,23 +16,16 @@ package cmdutil
import (
"fmt"
"os"
"time"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/pkg/cmd/iostreams"
"go.probo.inc/probo/pkg/cookiebanner"
"go.probo.inc/probo/pkg/proboctl/pgconn"
"go.probo.inc/probo/pkg/probodconfig"
"sigs.k8s.io/yaml"
)
type Factory struct {
IOStreams *iostreams.IOStreams
Version string
PgDSN string
CfgFile string
pgClient *pg.Client
}
@@ -58,88 +51,3 @@ func (f *Factory) PgClient() (*pg.Client, error) {
return f.pgClient, nil
}
// ProbodConfig loads the shared probod configuration file (--cfg-file).
// It reuses the exact file, struct, and json-tagged (un)marshaling probod
// uses, so proboctl and probod stay consistent.
func (f *Factory) ProbodConfig() (probodconfig.Config, error) {
if f.CfgFile == "" {
return probodconfig.Config{}, fmt.Errorf("set --cfg-file to the probod config file")
}
data, err := os.ReadFile(f.CfgFile)
if err != nil {
return probodconfig.Config{}, fmt.Errorf("cannot read config file %q: %w", f.CfgFile, err)
}
var full probodconfig.FullConfig
if err := yaml.Unmarshal(data, &full); err != nil {
return probodconfig.Config{}, fmt.Errorf("cannot parse config file %q: %w", f.CfgFile, err)
}
return full.Probod, nil
}
// 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.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(
log.WithName("proboctl"),
log.WithOutput(f.IOStreams.ErrOut),
)
firecrawlAPIKey := cfg.Agents.Tools.FirecrawlAPIKey
mappingAgentCfg, mappingClient, err := resolveAgentClient(cfg.Agents, "tracker-mapping", cfg.Agents.TrackerMapping, logger)
if err != nil {
return cookiebanner.TrackerEnrichmentAgentConfig{}, cookiebanner.TrackerMappingAgentConfig{}, fmt.Errorf("cannot build tracker mapping agent: %w", err)
}
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,
}
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
View File

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

61
pkg/proboctl/commontrackerpattern/reenrich.go
View File

@@ -20,9 +20,7 @@ import (
"io"
"github.com/spf13/cobra"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/pkg/cookiebanner"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/proboctl/cmdutil"
@@ -38,19 +36,16 @@ func newCmdReenrich(f *cmdutil.Factory) *cobra.Command {
flagKeyword string
flagState string
flagWithoutDescription bool
flagConcurrency int
flagDryRun bool
flagYes bool
flagEnqueue bool
)
cmd := &cobra.Command{
Use: "reenrich",
Short: "Re-describe common tracker patterns by running the enrichment agent",
Long: "Re-describe selected common tracker patterns. By default the enrichment " +
"agent runs in-process and the command returns only when the work is done " +
"(requires --cfg-file with an LLM provider). Use --enqueue to instead arm the " +
"async enrichment worker. Re-describe a banner's catalog rows with --linked-banner " +
Short: "Re-describe common tracker patterns via the enrichment worker",
Long: "Re-describe selected common tracker patterns by arming the async " +
"enrichment worker, which fills descriptions and fans them out to linked " +
"tracker patterns. Re-describe a banner's catalog rows with --linked-banner " +
"before running 'cookie-banner reset-trackers' so fresh descriptions copy down.",
Args: cobra.NoArgs,
}
@@ -63,10 +58,8 @@ func newCmdReenrich(f *cmdutil.Factory) *cobra.Command {
cmd.Flags().StringVar(&flagKeyword, "keyword", "", "Filter selected patterns by a pattern/description substring")
cmd.Flags().StringVar(&flagState, "state", "", "Filter selected patterns by enrichment state (queued, enriched, unenriched)")
cmd.Flags().BoolVar(&flagWithoutDescription, "without-description", false, "Only patterns with a blank description")
cmd.Flags().IntVar(&flagConcurrency, "concurrency", 4, "Number of patterns to enrich in parallel (sync mode)")
cmd.Flags().BoolVar(&flagDryRun, "dry-run", false, "Print the selected patterns without enriching")
cmd.Flags().BoolVar(&flagYes, "yes", false, "Skip confirmation")
cmd.Flags().BoolVar(&flagEnqueue, "enqueue", false, "Arm the async enrichment worker instead of running the agent in-process")
cmd.RunE = func(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()
@@ -110,46 +103,22 @@ func newCmdReenrich(f *cmdutil.Factory) *cobra.Command {
return fmt.Errorf("about to re-enrich %d pattern(s); pass --yes to proceed or --dry-run to preview", len(ids))
}
if flagEnqueue {
var requeued int64
var requeued int64
if err := pgClient.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
var ps coredata.CommonTrackerPatterns
if err := pgClient.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
var ps coredata.CommonTrackerPatterns
requeued, err = ps.RequestEnrichmentByIDs(ctx, tx, ids)
requeued, err = ps.RequestEnrichmentByIDs(ctx, tx, ids)
return err
},
); err != nil {
return fmt.Errorf("cannot enqueue enrichment: %w", err)
}
_, _ = fmt.Fprintf(out, "Queued %d common tracker pattern(s) for the enrichment worker.\n", requeued)
return nil
return err
},
); err != nil {
return fmt.Errorf("cannot enqueue enrichment: %w", err)
}
enrichmentCfg, mappingCfg, err := f.TrackerAgentsConfig()
if err != nil {
return err
}
logger := log.NewLogger(
log.WithName("proboctl"),
log.WithOutput(f.IOStreams.ErrOut),
)
enricher := cookiebanner.NewCommonPatternEnricher(pgClient, logger, enrichmentCfg, mappingCfg)
enriched, err := enricher.EnrichByIDs(ctx, ids, flagConcurrency)
_, _ = fmt.Fprintf(out, "Enriched %d of %d common tracker pattern(s).\n", enriched, len(ids))
if err != nil {
return fmt.Errorf("enrichment did not complete for all patterns: %w", err)
}
_, _ = fmt.Fprintf(out, "Queued %d common tracker pattern(s) for the enrichment worker.\n", requeued)
return nil
}

7
pkg/proboctl/root/root.go
View File

@@ -41,13 +41,6 @@ func NewCmdRoot(f *cmdutil.Factory) *cobra.Command {
"PostgreSQL connection URL (default: DATABASE_URL env)",
)
cmd.PersistentFlags().StringVar(
&f.CfgFile,
"cfg-file",
os.Getenv("PROBOD_CFG_FILE"),
"Path to the probod config file (default: PROBOD_CFG_FILE env); required for agent-backed commands",
)
cmd.AddCommand(seed.NewCmdSeed(f))
cmd.AddCommand(commontrackerpattern.NewCmdCommonTrackerPattern(f))
cmd.AddCommand(commonthirdparty.NewCmdCommonThirdParty(f))