From 4f6fcb42f9eb47208ef31973beee4f84651058aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=89mile=20R=C3=A9?= Date: Tue, 9 Jun 2026 11:47:18 +0200 Subject: [PATCH] Remove sync common tracker pattern reenriching MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Émile Ré --- pkg/cookiebanner/common_pattern_enricher.go | 55 ----------- pkg/proboctl/cmdutil/cmdutil.go | 92 ------------------ pkg/proboctl/cmdutil/llm.go | 97 ------------------- pkg/proboctl/commontrackerpattern/reenrich.go | 61 +++--------- pkg/proboctl/root/root.go | 7 -- 5 files changed, 15 insertions(+), 297 deletions(-) delete mode 100644 pkg/proboctl/cmdutil/llm.go diff --git a/pkg/cookiebanner/common_pattern_enricher.go b/pkg/cookiebanner/common_pattern_enricher.go index 2ccd17756..cae826a04 100644 --- a/pkg/cookiebanner/common_pattern_enricher.go +++ b/pkg/cookiebanner/common_pattern_enricher.go @@ -18,7 +18,6 @@ import ( "context" "fmt" "strings" - "sync/atomic" "time" "go.gearno.de/kit/log" @@ -28,7 +27,6 @@ import ( "go.probo.inc/probo/pkg/gid" "go.probo.inc/probo/pkg/llm" "go.probo.inc/probo/pkg/thirdparty" - "golang.org/x/sync/errgroup" ) // CommonPatternEnricher fills descriptions on common_tracker_patterns @@ -92,59 +90,6 @@ func (e *CommonPatternEnricher) Enabled() bool { return e.enrichmentAgent != nil } -// EnrichByIDs enriches each common tracker pattern id, with bounded -// concurrency, and returns the number successfully enriched. It is the -// synchronous entry point used by operator tooling: it completes when the -// work is done rather than arming the async queue. A concurrency <= 0 is -// treated as 1. -func (e *CommonPatternEnricher) EnrichByIDs( - ctx context.Context, - ids []gid.GID, - concurrency int, -) (int, error) { - if !e.Enabled() { - return 0, fmt.Errorf("common pattern enricher is disabled: no LLM client configured") - } - - if concurrency <= 0 { - concurrency = 1 - } - - var enriched atomic.Int64 - - g, gctx := errgroup.WithContext(ctx) - g.SetLimit(concurrency) - - for _, id := range ids { - g.Go(func() error { - var cp coredata.CommonTrackerPattern - - if err := e.pg.WithConn( - gctx, - func(ctx context.Context, conn pg.Querier) error { - return cp.LoadByID(ctx, conn, id) - }, - ); err != nil { - return fmt.Errorf("cannot load common tracker pattern %s: %w", id, err) - } - - if err := e.EnrichPattern(gctx, cp); err != nil { - return fmt.Errorf("cannot enrich common tracker pattern %s: %w", id, err) - } - - enriched.Add(1) - - return nil - }) - } - - if err := g.Wait(); err != nil { - return int(enriched.Load()), err - } - - return int(enriched.Load()), nil -} - // EnrichPattern researches a description for one common tracker pattern // (attributing a vendor first when unlinked), records it, and fans it out // to linked org patterns. A blank description is a terminal-for-now state: diff --git a/pkg/proboctl/cmdutil/cmdutil.go b/pkg/proboctl/cmdutil/cmdutil.go index 482ce1a57..4099b9ef0 100644 --- a/pkg/proboctl/cmdutil/cmdutil.go +++ b/pkg/proboctl/cmdutil/cmdutil.go @@ -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 -} diff --git a/pkg/proboctl/cmdutil/llm.go b/pkg/proboctl/cmdutil/llm.go deleted file mode 100644 index f6096c481..000000000 --- a/pkg/proboctl/cmdutil/llm.go +++ /dev/null @@ -1,97 +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 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 2542e3691..4fe88417d 100644 --- a/pkg/proboctl/commontrackerpattern/reenrich.go +++ b/pkg/proboctl/commontrackerpattern/reenrich.go @@ -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 } diff --git a/pkg/proboctl/root/root.go b/pkg/proboctl/root/root.go index 73a3d624b..4de1d0806 100644 --- a/pkg/proboctl/root/root.go +++ b/pkg/proboctl/root/root.go @@ -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))