From 662c0ae4286b12938b21ddbe4b1bba43da6e87e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=89mile=20R=C3=A9?= Date: Wed, 3 Jun 2026 13:37:36 +0200 Subject: [PATCH] Add common catalog query layer and enricher service MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce an API-style data layer for the global common tracker pattern and common third party catalogs: typed filters, order fields, CursorKey, cursor-paginated Load and CountAll, plus by-id enrichment re-queue and a scoped reset/remap helper for a banner's tracker patterns. These reuse the same page.Cursor/filter/order types the GraphQL API consumes, so a future proboctl API can back them unchanged. Extract the common-pattern enrichment logic out of the worker into a CommonPatternEnricher service so it can run either from the background queue or synchronously over a known set of ids; the worker becomes a thin poller that delegates to it. Extract the LLM client and tracker-agents config wiring into pkg/agentsbuild so probod and other binaries build agents identically; probod now delegates to it. Signed-off-by: Émile Ré --- pkg/agentsbuild/agentsbuild.go | 172 ++++++++++ pkg/cookiebanner/common_pattern_enricher.go | 322 ++++++++++++++++++ .../common_pattern_enrichment_worker.go | 222 +----------- pkg/coredata/common_third_party.go | 106 ++++++ pkg/coredata/common_third_party_filter.go | 42 ++- .../common_third_party_order_field.go | 89 +++++ pkg/coredata/common_tracker_pattern.go | 175 ++++++++++ pkg/coredata/common_tracker_pattern_filter.go | 213 ++++++++++++ .../common_tracker_pattern_order_field.go | 99 ++++++ pkg/coredata/detected_tracker.go | 93 +++++ pkg/coredata/tracker_pattern.go | 119 +++++++ pkg/probod/llm.go | 67 +--- pkg/probod/tracker_agents.go | 68 +--- 13 files changed, 1452 insertions(+), 335 deletions(-) create mode 100644 pkg/agentsbuild/agentsbuild.go create mode 100644 pkg/cookiebanner/common_pattern_enricher.go create mode 100644 pkg/coredata/common_third_party_order_field.go create mode 100644 pkg/coredata/common_tracker_pattern_filter.go create mode 100644 pkg/coredata/common_tracker_pattern_order_field.go diff --git a/pkg/agentsbuild/agentsbuild.go b/pkg/agentsbuild/agentsbuild.go new file mode 100644 index 000000000..a36ae8b91 --- /dev/null +++ b/pkg/agentsbuild/agentsbuild.go @@ -0,0 +1,172 @@ +// 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/cookiebanner/common_pattern_enricher.go b/pkg/cookiebanner/common_pattern_enricher.go new file mode 100644 index 000000000..b64ec320a --- /dev/null +++ b/pkg/cookiebanner/common_pattern_enricher.go @@ -0,0 +1,322 @@ +// 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 cookiebanner + +import ( + "context" + "fmt" + "strings" + "sync/atomic" + "time" + + "go.gearno.de/kit/log" + "go.gearno.de/kit/pg" + "go.probo.inc/probo/pkg/agent" + "go.probo.inc/probo/pkg/coredata" + "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 +// using an agent with web search, optionally attributing a vendor first, +// then fans the result out to every linked tracker pattern. It holds the +// agent dependencies so the enrichment logic can run from either the +// background worker (one claimed row at a time) or synchronously over a +// 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 +} + +// 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. +func NewCommonPatternEnricher( + pgClient *pg.Client, + logger *log.Logger, + cfg TrackerAgentsConfig, +) *CommonPatternEnricher { + agentTimeout := cfg.AgentTimeout + if agentTimeout <= 0 { + agentTimeout = defaultAgentTimeout + } + + e := &CommonPatternEnricher{ + pg: pgClient, + logger: logger, + agentTimeout: agentTimeout, + } + + if cfg.LLMClient != nil { + e.enrichmentAgent = buildCommonPatternEnrichmentAgent(cfg, pgClient, logger) + e.mappingAgent = buildTrackerMappingAgent(cfg, pgClient, logger) + } + + return e +} + +// Enabled reports whether an LLM-backed enrichment agent is configured. +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: +// the row is marked enriched so stale recovery never re-queues it, while a +// later third-party link re-arms a vendor-informed second attempt. +func (e *CommonPatternEnricher) EnrichPattern(ctx context.Context, cp coredata.CommonTrackerPattern) error { + if !e.Enabled() { + return nil + } + + thirdPartyName, err := e.loadThirdPartyName(ctx, cp) + if err != nil { + return err + } + + // Map before enriching: an unlinked pattern is run through the + // mapping agent first so a confident vendor both seeds the enrichment + // prompt and gets linked. Attribution stays the mapping pipeline's + // job; the enricher only reuses it. An already-linked pattern skips + // this entirely. + var attribution *TrackerMappingAgentResult + + if cp.CommonThirdPartyID == nil { + attribution, err = e.identifyThirdParty(ctx, cp) + if err != nil { + return err + } + + if attribution != nil { + thirdPartyName = attribution.ThirdPartyName + } + } + + description, err := e.research(ctx, cp, thirdPartyName) + if err != nil { + return fmt.Errorf("cannot research tracker description: %w", err) + } + + return e.pg.WithTx( + ctx, + func(ctx context.Context, tx pg.Tx) error { + // Resolve or create the catalog vendor only for an unlinked + // pattern; the mapping pipeline owns creation, so we reuse its + // name+slug dedup and never create a duplicate or override an + // existing link. + var thirdPartyID *gid.GID + + if attribution != nil && cp.CommonThirdPartyID == nil { + thirdPartyID, err = thirdparty.ResolveOrCreateCommonThirdParty(ctx, tx, e.logger, attribution.ThirdPartyName, attribution.Category) + if err != nil { + return fmt.Errorf("cannot resolve or create common third party: %w", err) + } + } + + if err := cp.SetEnriched(ctx, tx, description, thirdPartyID); err != nil { + return fmt.Errorf("cannot set common tracker pattern enriched: %w", err) + } + + var backfilled int64 + + if description != "" { + var patterns coredata.TrackerPatterns + + backfilled, err = patterns.BackfillDescriptionByCommonTrackerPatternID(ctx, tx, cp.ID, description) + if err != nil { + return err + } + } + + e.logger.InfoCtx( + ctx, + "enriched common tracker pattern", + log.String("common_tracker_pattern_id", cp.ID.String()), + log.String("pattern", cp.Pattern), + log.Bool("described", description != ""), + log.Bool("third_party_linked", thirdPartyID != nil), + log.Int64("backfilled_tracker_patterns", backfilled), + ) + + return nil + }, + ) +} + +func (e *CommonPatternEnricher) loadThirdPartyName( + ctx context.Context, + cp coredata.CommonTrackerPattern, +) (string, error) { + if cp.CommonThirdPartyID == nil { + return "", nil + } + + var name string + + if err := e.pg.WithConn( + ctx, + func(ctx context.Context, conn pg.Querier) error { + var party coredata.CommonThirdParty + if err := party.LoadByID(ctx, conn, *cp.CommonThirdPartyID); err != nil { + return err + } + + name = party.Name + + return nil + }, + ); err != nil { + return "", fmt.Errorf("cannot load common third party for enrichment: %w", err) + } + + return name, nil +} + +func (e *CommonPatternEnricher) research( + ctx context.Context, + cp coredata.CommonTrackerPattern, + thirdPartyName string, +) (string, error) { + prompt := buildEnrichmentPrompt(cp, thirdPartyName) + + agentCtx, cancel := context.WithTimeout(ctx, e.agentTimeout) + defer cancel() + + result, err := agent.RunTyped[CommonPatternEnrichmentResult]( + agentCtx, + e.enrichmentAgent, + []llm.Message{ + { + Role: llm.RoleUser, + Parts: []llm.Part{llm.TextPart{Text: prompt}}, + }, + }, + ) + if err != nil { + return "", fmt.Errorf("enrichment agent run failed: %w", err) + } + + return strings.TrimSpace(result.Output.Description), nil +} + +// identifyThirdParty reuses the tracker-mapping agent to attribute a +// vendor to an unlinked catalog pattern. It performs no DB writes: it +// returns the confident attribution (name, category, confidence) or nil +// when the agent is unsure, leaving the caller to resolve or create the +// catalog row. A failed agent run is best-effort and non-fatal, +// mirroring the mapping worker's identifyWithAgent. +func (e *CommonPatternEnricher) identifyThirdParty( + ctx context.Context, + cp coredata.CommonTrackerPattern, +) (*TrackerMappingAgentResult, error) { + if e.mappingAgent == nil { + return nil, nil + } + + prompt := buildCommonPatternIdentificationPrompt(cp) + + agentCtx, cancel := context.WithTimeout(ctx, e.agentTimeout) + defer cancel() + + result, err := agent.RunTyped[TrackerMappingAgentResult]( + agentCtx, + e.mappingAgent, + []llm.Message{ + { + Role: llm.RoleUser, + Parts: []llm.Part{llm.TextPart{Text: prompt}}, + }, + }, + ) + if err != nil { + e.logger.WarnCtx( + ctx, + "mapping agent identification failed during enrichment", + log.Error(err), + log.String("pattern", cp.Pattern), + ) + + return nil, nil + } + + out := result.Output + out.ThirdPartyName = strings.TrimSpace(out.ThirdPartyName) + + if out.ThirdPartyName == "" || out.ThirdPartyConfidence < agentThirdPartyConfidenceThreshold { + return nil, nil + } + + return &out, nil +} diff --git a/pkg/cookiebanner/common_pattern_enrichment_worker.go b/pkg/cookiebanner/common_pattern_enrichment_worker.go index 7429291d2..e31b4bc24 100644 --- a/pkg/cookiebanner/common_pattern_enrichment_worker.go +++ b/pkg/cookiebanner/common_pattern_enrichment_worker.go @@ -18,28 +18,25 @@ import ( "context" "errors" "fmt" - "strings" "time" "go.gearno.de/kit/log" "go.gearno.de/kit/pg" "go.gearno.de/kit/worker" - "go.probo.inc/probo/pkg/agent" "go.probo.inc/probo/pkg/coredata" - "go.probo.inc/probo/pkg/gid" - "go.probo.inc/probo/pkg/llm" - "go.probo.inc/probo/pkg/thirdparty" ) const defaultEnrichmentStaleAfter = 10 * time.Minute +// commonPatternEnrichmentHandler is the queue poller for common tracker +// pattern enrichment. It owns only the claim/dequeue and stale-recovery +// mechanics; the enrichment work itself lives in CommonPatternEnricher so +// it can also run synchronously from operator tooling. type commonPatternEnrichmentHandler struct { - pg *pg.Client - logger *log.Logger - enrichmentAgent *agent.Agent - mappingAgent *agent.Agent - staleAfter time.Duration - agentTimeout time.Duration + pg *pg.Client + logger *log.Logger + enricher *CommonPatternEnricher + staleAfter time.Duration } // NewCommonPatternEnrichmentWorker builds the worker that fills @@ -60,21 +57,11 @@ func NewCommonPatternEnrichmentWorker( staleAfter = defaultEnrichmentStaleAfter } - agentTimeout := cfg.AgentTimeout - if agentTimeout <= 0 { - agentTimeout = defaultAgentTimeout - } - h := &commonPatternEnrichmentHandler{ - pg: pgClient, - logger: logger, - staleAfter: staleAfter, - agentTimeout: agentTimeout, - } - - if cfg.LLMClient != nil { - h.enrichmentAgent = buildCommonPatternEnrichmentAgent(cfg, pgClient, logger) - h.mappingAgent = buildTrackerMappingAgent(cfg, pgClient, logger) + pg: pgClient, + logger: logger, + enricher: NewCommonPatternEnricher(pgClient, logger, cfg), + staleAfter: staleAfter, } return worker.New( @@ -109,86 +96,11 @@ func (h *commonPatternEnrichmentHandler) Claim(ctx context.Context) (coredata.Co } func (h *commonPatternEnrichmentHandler) Process(ctx context.Context, cp coredata.CommonTrackerPattern) error { - if h.enrichmentAgent == nil { + if !h.enricher.Enabled() { return nil } - thirdPartyName, err := h.loadThirdPartyName(ctx, cp) - if err != nil { - return err - } - - // Map before enriching: an unlinked pattern is run through the - // mapping agent first so a confident vendor both seeds the enrichment - // prompt and gets linked. Attribution stays the mapping pipeline's - // job; the enricher only reuses it. An already-linked pattern skips - // this entirely. - var attribution *TrackerMappingAgentResult - - if cp.CommonThirdPartyID == nil { - attribution, err = h.identifyThirdParty(ctx, cp) - if err != nil { - return err - } - - if attribution != nil { - thirdPartyName = attribution.ThirdPartyName - } - } - - description, err := h.research(ctx, cp, thirdPartyName) - if err != nil { - return fmt.Errorf("cannot research tracker description: %w", err) - } - - return h.pg.WithTx( - ctx, - func(ctx context.Context, tx pg.Tx) error { - // Resolve or create the catalog vendor only for an unlinked - // pattern; the mapping pipeline owns creation, so we reuse its - // name+slug dedup and never create a duplicate or override an - // existing link. - var thirdPartyID *gid.GID - - if attribution != nil && cp.CommonThirdPartyID == nil { - thirdPartyID, err = thirdparty.ResolveOrCreateCommonThirdParty(ctx, tx, h.logger, attribution.ThirdPartyName, attribution.Category) - if err != nil { - return fmt.Errorf("cannot resolve or create common third party: %w", err) - } - } - - // A blank description is recorded as a terminal-for-now state: - // the row is marked enriched so the stale-recovery loop never - // re-queues it, but a later third-party link (mapping worker) - // re-arms enrichment for a vendor-informed second attempt. - if err := cp.SetEnriched(ctx, tx, description, thirdPartyID); err != nil { - return fmt.Errorf("cannot set common tracker pattern enriched: %w", err) - } - - var backfilled int64 - - if description != "" { - var patterns coredata.TrackerPatterns - - backfilled, err = patterns.BackfillDescriptionByCommonTrackerPatternID(ctx, tx, cp.ID, description) - if err != nil { - return err - } - } - - h.logger.InfoCtx( - ctx, - "enriched common tracker pattern", - log.String("common_tracker_pattern_id", cp.ID.String()), - log.String("pattern", cp.Pattern), - log.Bool("described", description != ""), - log.Bool("third_party_linked", thirdPartyID != nil), - log.Int64("backfilled_tracker_patterns", backfilled), - ) - - return nil - }, - ) + return h.enricher.EnrichPattern(ctx, cp) } func (h *commonPatternEnrichmentHandler) RecoverStale(ctx context.Context) error { @@ -203,109 +115,3 @@ func (h *commonPatternEnrichmentHandler) RecoverStale(ctx context.Context) error }, ) } - -func (h *commonPatternEnrichmentHandler) loadThirdPartyName( - ctx context.Context, - cp coredata.CommonTrackerPattern, -) (string, error) { - if cp.CommonThirdPartyID == nil { - return "", nil - } - - var name string - - if err := h.pg.WithConn( - ctx, - func(ctx context.Context, conn pg.Querier) error { - var party coredata.CommonThirdParty - if err := party.LoadByID(ctx, conn, *cp.CommonThirdPartyID); err != nil { - return err - } - - name = party.Name - - return nil - }, - ); err != nil { - return "", fmt.Errorf("cannot load common third party for enrichment: %w", err) - } - - return name, nil -} - -func (h *commonPatternEnrichmentHandler) research( - ctx context.Context, - cp coredata.CommonTrackerPattern, - thirdPartyName string, -) (string, error) { - prompt := buildEnrichmentPrompt(cp, thirdPartyName) - - agentCtx, cancel := context.WithTimeout(ctx, h.agentTimeout) - defer cancel() - - result, err := agent.RunTyped[CommonPatternEnrichmentResult]( - agentCtx, - h.enrichmentAgent, - []llm.Message{ - { - Role: llm.RoleUser, - Parts: []llm.Part{llm.TextPart{Text: prompt}}, - }, - }, - ) - if err != nil { - return "", fmt.Errorf("enrichment agent run failed: %w", err) - } - - return strings.TrimSpace(result.Output.Description), nil -} - -// identifyThirdParty reuses the tracker-mapping agent to attribute a -// vendor to an unlinked catalog pattern. It performs no DB writes: it -// returns the confident attribution (name, category, confidence) or nil -// when the agent is unsure, leaving the caller to resolve or create the -// catalog row. A failed agent run is best-effort and non-fatal, -// mirroring the mapping worker's identifyWithAgent. -func (h *commonPatternEnrichmentHandler) identifyThirdParty( - ctx context.Context, - cp coredata.CommonTrackerPattern, -) (*TrackerMappingAgentResult, error) { - if h.mappingAgent == nil { - return nil, nil - } - - prompt := buildCommonPatternIdentificationPrompt(cp) - - agentCtx, cancel := context.WithTimeout(ctx, h.agentTimeout) - defer cancel() - - result, err := agent.RunTyped[TrackerMappingAgentResult]( - agentCtx, - h.mappingAgent, - []llm.Message{ - { - Role: llm.RoleUser, - Parts: []llm.Part{llm.TextPart{Text: prompt}}, - }, - }, - ) - if err != nil { - h.logger.WarnCtx( - ctx, - "mapping agent identification failed during enrichment", - log.Error(err), - log.String("pattern", cp.Pattern), - ) - - return nil, nil - } - - out := result.Output - out.ThirdPartyName = strings.TrimSpace(out.ThirdPartyName) - - if out.ThirdPartyName == "" || out.ThirdPartyConfidence < agentThirdPartyConfidenceThreshold { - return nil, nil - } - - return &out, nil -} diff --git a/pkg/coredata/common_third_party.go b/pkg/coredata/common_third_party.go index 17d743e92..b078b5525 100644 --- a/pkg/coredata/common_third_party.go +++ b/pkg/coredata/common_third_party.go @@ -25,6 +25,7 @@ import ( "go.gearno.de/kit/pg" "go.probo.inc/probo/pkg/gid" "go.probo.inc/probo/pkg/iam/policy" + "go.probo.inc/probo/pkg/page" ) type ( @@ -644,3 +645,108 @@ WHERE return nil } + +func (t *CommonThirdParty) CursorKey(field CommonThirdPartyOrderField) page.CursorKey { + switch field { + case CommonThirdPartyOrderFieldName: + return page.NewCursorKey(t.ID, t.Name) + case CommonThirdPartyOrderFieldCreatedAt: + return page.NewCursorKey(t.ID, t.CreatedAt) + case CommonThirdPartyOrderFieldUpdatedAt: + return page.NewCursorKey(t.ID, t.UpdatedAt) + } + + panic(fmt.Sprintf("unsupported order by: %s", field)) +} + +// Load returns a cursor-paginated, filtered page of common third +// parties. The catalog is global (no tenant scope); the cursor supplies +// the limit and ordering. Unlike LoadAll (capped at 20, name only), this +// is the listing entry point a future API/CLI consumes. +func (t *CommonThirdParties) Load( + ctx context.Context, + conn pg.Querier, + cursor *page.Cursor[CommonThirdPartyOrderField], + filter *CommonThirdPartyFilter, +) error { + q := ` +SELECT + id, + name, + slug, + category, + headquarter_address, + legal_name, + website_url, + privacy_policy_url, + service_level_agreement_url, + service_software_agreement_url, + data_processing_agreement_url, + business_associate_agreement_url, + subprocessors_list_url, + certifications, + status_page_url, + terms_of_service_url, + security_page_url, + trust_page_url, + logo_file_id, + created_at, + updated_at +FROM + common_third_parties +WHERE + %s + AND %s +` + + q = fmt.Sprintf(q, filter.SQLFragment(), cursor.SQLFragment()) + + args := pgx.StrictNamedArgs{} + maps.Copy(args, filter.SQLArguments()) + maps.Copy(args, cursor.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query common third parties: %w", err) + } + + parties, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[CommonThirdParty]) + if err != nil { + return fmt.Errorf("cannot collect common third parties: %w", err) + } + + *t = parties + + return nil +} + +// CountAll returns the number of common third parties matching the +// filter, ignoring pagination. +func (t *CommonThirdParties) CountAll( + ctx context.Context, + conn pg.Querier, + filter *CommonThirdPartyFilter, +) (int, error) { + q := ` +SELECT + COUNT(id) +FROM + common_third_parties +WHERE + %s +` + + q = fmt.Sprintf(q, filter.SQLFragment()) + + args := pgx.StrictNamedArgs{} + maps.Copy(args, filter.SQLArguments()) + + row := conn.QueryRow(ctx, q, args) + + var count int + if err := row.Scan(&count); err != nil { + return 0, fmt.Errorf("cannot count common third parties: %w", err) + } + + return count, nil +} diff --git a/pkg/coredata/common_third_party_filter.go b/pkg/coredata/common_third_party_filter.go index c747ab77d..ac78e4947 100644 --- a/pkg/coredata/common_third_party_filter.go +++ b/pkg/coredata/common_third_party_filter.go @@ -19,13 +19,25 @@ import ( ) type CommonThirdPartyFilter struct { - name *string + name *string + category *ThirdPartyCategory + keyword *string } func NewCommonThirdPartyFilter(name *string) *CommonThirdPartyFilter { return &CommonThirdPartyFilter{name: name} } +func (f *CommonThirdPartyFilter) WithCategory(category *ThirdPartyCategory) *CommonThirdPartyFilter { + f.category = category + return f +} + +func (f *CommonThirdPartyFilter) WithKeyword(keyword *string) *CommonThirdPartyFilter { + f.keyword = keyword + return f +} + func (f *CommonThirdPartyFilter) SQLFragment() string { return `( CASE @@ -33,14 +45,40 @@ func (f *CommonThirdPartyFilter) SQLFragment() string { name ILIKE '%' || @filter_name || '%' ELSE TRUE END + AND + CASE + WHEN @filter_category::text IS NOT NULL THEN + category = @filter_category::third_party_category + ELSE TRUE + END + AND + CASE + WHEN @filter_keyword::text IS NOT NULL AND @filter_keyword::text != '' THEN + (name ILIKE '%' || @filter_keyword || '%' + OR slug ILIKE '%' || @filter_keyword || '%') + ELSE TRUE + END )` } func (f *CommonThirdPartyFilter) SQLArguments() pgx.StrictNamedArgs { - args := pgx.StrictNamedArgs{"filter_name": nil} + args := pgx.StrictNamedArgs{ + "filter_name": nil, + "filter_category": nil, + "filter_keyword": nil, + } + if f.name != nil { args["filter_name"] = *f.name } + if f.category != nil { + args["filter_category"] = string(*f.category) + } + + if f.keyword != nil { + args["filter_keyword"] = *f.keyword + } + return args } diff --git a/pkg/coredata/common_third_party_order_field.go b/pkg/coredata/common_third_party_order_field.go new file mode 100644 index 000000000..70cbd2907 --- /dev/null +++ b/pkg/coredata/common_third_party_order_field.go @@ -0,0 +1,89 @@ +// 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 coredata + +import ( + "encoding" + "fmt" + + "go.probo.inc/probo/pkg/page" +) + +type CommonThirdPartyOrderField string + +const ( + CommonThirdPartyOrderFieldName CommonThirdPartyOrderField = "NAME" + CommonThirdPartyOrderFieldCreatedAt CommonThirdPartyOrderField = "CREATED_AT" + CommonThirdPartyOrderFieldUpdatedAt CommonThirdPartyOrderField = "UPDATED_AT" +) + +var ( + _ page.OrderField = CommonThirdPartyOrderField("") + _ fmt.Stringer = CommonThirdPartyOrderField("") + _ encoding.TextMarshaler = CommonThirdPartyOrderField("") + _ encoding.TextUnmarshaler = (*CommonThirdPartyOrderField)(nil) +) + +func CommonThirdPartyOrderFields() []CommonThirdPartyOrderField { + return []CommonThirdPartyOrderField{ + CommonThirdPartyOrderFieldName, + CommonThirdPartyOrderFieldCreatedAt, + CommonThirdPartyOrderFieldUpdatedAt, + } +} + +func (v CommonThirdPartyOrderField) IsValid() bool { + switch v { + case + CommonThirdPartyOrderFieldName, + CommonThirdPartyOrderFieldCreatedAt, + CommonThirdPartyOrderFieldUpdatedAt: + return true + } + + return false +} + +func (v CommonThirdPartyOrderField) String() string { + return string(v) +} + +func (v CommonThirdPartyOrderField) MarshalText() ([]byte, error) { + return []byte(v.String()), nil +} + +func (v *CommonThirdPartyOrderField) UnmarshalText(text []byte) error { + val := CommonThirdPartyOrderField(text) + if !val.IsValid() { + return fmt.Errorf("invalid CommonThirdPartyOrderField value: %q", string(text)) + } + + *v = val + + return nil +} + +func (v CommonThirdPartyOrderField) Column() string { + switch v { + case CommonThirdPartyOrderFieldName: + return "name" + case CommonThirdPartyOrderFieldCreatedAt: + return "created_at" + case CommonThirdPartyOrderFieldUpdatedAt: + return "updated_at" + } + + panic(fmt.Sprintf("unsupported order by: %s", v)) +} diff --git a/pkg/coredata/common_tracker_pattern.go b/pkg/coredata/common_tracker_pattern.go index 07a768399..43d19c127 100644 --- a/pkg/coredata/common_tracker_pattern.go +++ b/pkg/coredata/common_tracker_pattern.go @@ -18,11 +18,13 @@ import ( "context" "errors" "fmt" + "maps" "time" "github.com/jackc/pgx/v5" "go.gearno.de/kit/pg" "go.probo.inc/probo/pkg/gid" + "go.probo.inc/probo/pkg/page" ) type ( @@ -693,3 +695,176 @@ WHERE return nil } + +func (p *CommonTrackerPattern) CursorKey(field CommonTrackerPatternOrderField) page.CursorKey { + switch field { + case CommonTrackerPatternOrderFieldPattern: + return page.NewCursorKey(p.ID, p.Pattern) + case CommonTrackerPatternOrderFieldConfidence: + return page.NewCursorKey(p.ID, p.Confidence) + case CommonTrackerPatternOrderFieldCreatedAt: + return page.NewCursorKey(p.ID, p.CreatedAt) + case CommonTrackerPatternOrderFieldUpdatedAt: + return page.NewCursorKey(p.ID, p.UpdatedAt) + case CommonTrackerPatternOrderFieldEnrichedAt: + if p.EnrichedAt == nil { + return page.NewCursorKey(p.ID, time.Time{}) + } + + return page.NewCursorKey(p.ID, *p.EnrichedAt) + } + + panic(fmt.Sprintf("unsupported order by: %s", field)) +} + +// Load returns a cursor-paginated, filtered page of common tracker +// patterns. The catalog is global (no tenant scope). The cursor supplies +// the limit and ordering; callers wrap the result with page.NewPage. +func (ps *CommonTrackerPatterns) Load( + ctx context.Context, + conn pg.Querier, + cursor *page.Cursor[CommonTrackerPatternOrderField], + filter *CommonTrackerPatternFilter, +) error { + q := ` +SELECT + id, + common_third_party_id, + tracker_type, + pattern, + match_type, + description, + max_age_seconds, + confidence, + enrichment_requested_at, + enriched_at, + created_at, + updated_at +FROM + common_tracker_patterns +WHERE + %s + AND %s +` + + q = fmt.Sprintf(q, filter.SQLFragment(), cursor.SQLFragment()) + + args := pgx.StrictNamedArgs{} + maps.Copy(args, filter.SQLArguments()) + maps.Copy(args, cursor.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query common tracker patterns: %w", err) + } + + patterns, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[CommonTrackerPattern]) + if err != nil { + return fmt.Errorf("cannot collect common tracker patterns: %w", err) + } + + *ps = patterns + + return nil +} + +// CountAll returns the number of common tracker patterns matching the +// filter, ignoring pagination. +func (ps *CommonTrackerPatterns) CountAll( + ctx context.Context, + conn pg.Querier, + filter *CommonTrackerPatternFilter, +) (int, error) { + q := ` +SELECT + COUNT(id) +FROM + common_tracker_patterns +WHERE + %s +` + + q = fmt.Sprintf(q, filter.SQLFragment()) + + args := pgx.StrictNamedArgs{} + maps.Copy(args, filter.SQLArguments()) + + row := conn.QueryRow(ctx, q, args) + + var count int + if err := row.Scan(&count); err != nil { + return 0, fmt.Errorf("cannot count common tracker patterns: %w", err) + } + + return count, nil +} + +// LoadAllIDs returns every common tracker pattern id matching the filter, +// with no pagination. It backs bulk operations (e.g. operator-driven +// re-enrichment) that act on the entire matching set. +func (ps *CommonTrackerPatterns) LoadAllIDs( + ctx context.Context, + conn pg.Querier, + filter *CommonTrackerPatternFilter, +) ([]gid.GID, error) { + q := ` +SELECT + id +FROM + common_tracker_patterns +WHERE + %s +ORDER BY pattern ASC +` + + q = fmt.Sprintf(q, filter.SQLFragment()) + + args := pgx.StrictNamedArgs{} + maps.Copy(args, filter.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return nil, fmt.Errorf("cannot query common tracker pattern ids: %w", err) + } + + ids, err := pgx.CollectRows(rows, pgx.RowTo[gid.GID]) + if err != nil { + return nil, fmt.Errorf("cannot collect common tracker pattern ids: %w", err) + } + + return ids, nil +} + +// RequestEnrichmentByIDs arms enrichment on the given common tracker +// patterns. When resetEnriched is true it also clears enriched_at so rows +// that previously reached a terminal state are re-processed. Returns the +// number of rows re-queued. This is the async fallback path; the +// synchronous enricher service is preferred. +func (ps *CommonTrackerPatterns) RequestEnrichmentByIDs( + ctx context.Context, + tx pg.Tx, + ids []gid.GID, + resetEnriched bool, +) (int64, error) { + q := ` +UPDATE common_tracker_patterns +SET + enrichment_requested_at = NOW(), + enriched_at = CASE WHEN @reset_enriched THEN NULL ELSE enriched_at END, + updated_at = NOW() +WHERE + id = ANY(@ids) +` + + args := pgx.StrictNamedArgs{ + "ids": ids, + "reset_enriched": resetEnriched, + } + + result, err := tx.Exec(ctx, q, args) + if err != nil { + return 0, fmt.Errorf("cannot request common tracker pattern enrichment: %w", err) + } + + return result.RowsAffected(), nil +} diff --git a/pkg/coredata/common_tracker_pattern_filter.go b/pkg/coredata/common_tracker_pattern_filter.go new file mode 100644 index 000000000..dbfd4f611 --- /dev/null +++ b/pkg/coredata/common_tracker_pattern_filter.go @@ -0,0 +1,213 @@ +// 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 coredata + +import ( + "fmt" + + "github.com/jackc/pgx/v5" + "go.probo.inc/probo/pkg/gid" +) + +// CommonTrackerPatternEnrichmentState is a synthetic filter over the +// enrichment_requested_at / enriched_at columns. It is not a stored +// column; it classifies a row's position in the enrichment lifecycle. +type CommonTrackerPatternEnrichmentState string + +const ( + // CommonTrackerPatternEnrichmentStateQueued: a row armed for the + // enrichment worker (enrichment_requested_at IS NOT NULL). + CommonTrackerPatternEnrichmentStateQueued CommonTrackerPatternEnrichmentState = "QUEUED" + // CommonTrackerPatternEnrichmentStateEnriched: a row whose + // enrichment has completed (enriched_at IS NOT NULL) and is not + // re-queued. + CommonTrackerPatternEnrichmentStateEnriched CommonTrackerPatternEnrichmentState = "ENRICHED" + // CommonTrackerPatternEnrichmentStateUnenriched: a row never enriched + // and not currently queued. + CommonTrackerPatternEnrichmentStateUnenriched CommonTrackerPatternEnrichmentState = "UNENRICHED" +) + +func (s CommonTrackerPatternEnrichmentState) IsValid() bool { + switch s { + case + CommonTrackerPatternEnrichmentStateQueued, + CommonTrackerPatternEnrichmentStateEnriched, + CommonTrackerPatternEnrichmentStateUnenriched: + return true + } + + return false +} + +func (s CommonTrackerPatternEnrichmentState) String() string { + return string(s) +} + +func (s CommonTrackerPatternEnrichmentState) MarshalText() ([]byte, error) { + return []byte(s.String()), nil +} + +func (s *CommonTrackerPatternEnrichmentState) UnmarshalText(text []byte) error { + val := CommonTrackerPatternEnrichmentState(text) + if !val.IsValid() { + return fmt.Errorf("invalid CommonTrackerPatternEnrichmentState value: %q", string(text)) + } + + *s = val + + return nil +} + +type CommonTrackerPatternFilter struct { + trackerType *TrackerType + matchType *TrackerPatternMatchType + commonThirdPartyID *gid.GID + keyword *string + linked *bool + state *CommonTrackerPatternEnrichmentState +} + +func NewCommonTrackerPatternFilter() *CommonTrackerPatternFilter { + return &CommonTrackerPatternFilter{} +} + +func (f *CommonTrackerPatternFilter) WithTrackerType(trackerType *TrackerType) *CommonTrackerPatternFilter { + f.trackerType = trackerType + return f +} + +func (f *CommonTrackerPatternFilter) WithMatchType(matchType *TrackerPatternMatchType) *CommonTrackerPatternFilter { + f.matchType = matchType + return f +} + +func (f *CommonTrackerPatternFilter) WithCommonThirdPartyID(id *gid.GID) *CommonTrackerPatternFilter { + f.commonThirdPartyID = id + return f +} + +func (f *CommonTrackerPatternFilter) WithKeyword(keyword *string) *CommonTrackerPatternFilter { + f.keyword = keyword + return f +} + +func (f *CommonTrackerPatternFilter) WithLinked(linked *bool) *CommonTrackerPatternFilter { + f.linked = linked + return f +} + +func (f *CommonTrackerPatternFilter) WithState(state *CommonTrackerPatternEnrichmentState) *CommonTrackerPatternFilter { + f.state = state + return f +} + +func (f *CommonTrackerPatternFilter) SQLFragment() string { + if f == nil { + return "TRUE" + } + + return ` +( + CASE + WHEN @filter_tracker_type::text IS NOT NULL THEN + tracker_type = @filter_tracker_type::tracker_type + ELSE TRUE + END + AND + CASE + WHEN @filter_match_type::text IS NOT NULL THEN + match_type = @filter_match_type::cookie_pattern_match_type + ELSE TRUE + END + AND + CASE + WHEN @filter_common_third_party_id::text IS NOT NULL THEN + common_third_party_id = @filter_common_third_party_id::text + ELSE TRUE + END + AND + CASE + WHEN @filter_keyword::text IS NOT NULL AND @filter_keyword::text != '' THEN + (pattern ILIKE '%' || @filter_keyword || '%' + OR description ILIKE '%' || @filter_keyword || '%') + ELSE TRUE + END + AND + CASE + WHEN @filter_linked::boolean IS NULL THEN TRUE + WHEN @filter_linked::boolean THEN common_third_party_id IS NOT NULL + ELSE common_third_party_id IS NULL + END + AND + CASE + WHEN @filter_state_queued::boolean THEN enrichment_requested_at IS NOT NULL + WHEN @filter_state_enriched::boolean THEN + enrichment_requested_at IS NULL AND enriched_at IS NOT NULL + WHEN @filter_state_unenriched::boolean THEN + enrichment_requested_at IS NULL AND enriched_at IS NULL + ELSE TRUE + END +)` +} + +func (f *CommonTrackerPatternFilter) SQLArguments() pgx.StrictNamedArgs { + args := pgx.StrictNamedArgs{ + "filter_tracker_type": nil, + "filter_match_type": nil, + "filter_common_third_party_id": nil, + "filter_keyword": nil, + "filter_linked": nil, + "filter_state_queued": false, + "filter_state_enriched": false, + "filter_state_unenriched": false, + } + + if f == nil { + return args + } + + if f.trackerType != nil { + args["filter_tracker_type"] = string(*f.trackerType) + } + + if f.matchType != nil { + args["filter_match_type"] = string(*f.matchType) + } + + if f.commonThirdPartyID != nil { + args["filter_common_third_party_id"] = *f.commonThirdPartyID + } + + if f.keyword != nil { + args["filter_keyword"] = *f.keyword + } + + if f.linked != nil { + args["filter_linked"] = *f.linked + } + + if f.state != nil { + switch *f.state { + case CommonTrackerPatternEnrichmentStateQueued: + args["filter_state_queued"] = true + case CommonTrackerPatternEnrichmentStateEnriched: + args["filter_state_enriched"] = true + case CommonTrackerPatternEnrichmentStateUnenriched: + args["filter_state_unenriched"] = true + } + } + + return args +} diff --git a/pkg/coredata/common_tracker_pattern_order_field.go b/pkg/coredata/common_tracker_pattern_order_field.go new file mode 100644 index 000000000..a52eb72fa --- /dev/null +++ b/pkg/coredata/common_tracker_pattern_order_field.go @@ -0,0 +1,99 @@ +// 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 coredata + +import ( + "encoding" + "fmt" + + "go.probo.inc/probo/pkg/page" +) + +type CommonTrackerPatternOrderField string + +const ( + CommonTrackerPatternOrderFieldPattern CommonTrackerPatternOrderField = "PATTERN" + CommonTrackerPatternOrderFieldConfidence CommonTrackerPatternOrderField = "CONFIDENCE" + CommonTrackerPatternOrderFieldCreatedAt CommonTrackerPatternOrderField = "CREATED_AT" + CommonTrackerPatternOrderFieldUpdatedAt CommonTrackerPatternOrderField = "UPDATED_AT" + CommonTrackerPatternOrderFieldEnrichedAt CommonTrackerPatternOrderField = "ENRICHED_AT" +) + +var ( + _ page.OrderField = CommonTrackerPatternOrderField("") + _ fmt.Stringer = CommonTrackerPatternOrderField("") + _ encoding.TextMarshaler = CommonTrackerPatternOrderField("") + _ encoding.TextUnmarshaler = (*CommonTrackerPatternOrderField)(nil) +) + +func CommonTrackerPatternOrderFields() []CommonTrackerPatternOrderField { + return []CommonTrackerPatternOrderField{ + CommonTrackerPatternOrderFieldPattern, + CommonTrackerPatternOrderFieldConfidence, + CommonTrackerPatternOrderFieldCreatedAt, + CommonTrackerPatternOrderFieldUpdatedAt, + CommonTrackerPatternOrderFieldEnrichedAt, + } +} + +func (v CommonTrackerPatternOrderField) IsValid() bool { + switch v { + case + CommonTrackerPatternOrderFieldPattern, + CommonTrackerPatternOrderFieldConfidence, + CommonTrackerPatternOrderFieldCreatedAt, + CommonTrackerPatternOrderFieldUpdatedAt, + CommonTrackerPatternOrderFieldEnrichedAt: + return true + } + + return false +} + +func (v CommonTrackerPatternOrderField) String() string { + return string(v) +} + +func (v CommonTrackerPatternOrderField) MarshalText() ([]byte, error) { + return []byte(v.String()), nil +} + +func (v *CommonTrackerPatternOrderField) UnmarshalText(text []byte) error { + val := CommonTrackerPatternOrderField(text) + if !val.IsValid() { + return fmt.Errorf("invalid CommonTrackerPatternOrderField value: %q", string(text)) + } + + *v = val + + return nil +} + +func (v CommonTrackerPatternOrderField) Column() string { + switch v { + case CommonTrackerPatternOrderFieldPattern: + return "pattern" + case CommonTrackerPatternOrderFieldConfidence: + return "confidence" + case CommonTrackerPatternOrderFieldCreatedAt: + return "created_at" + case CommonTrackerPatternOrderFieldUpdatedAt: + return "updated_at" + case CommonTrackerPatternOrderFieldEnrichedAt: + return "COALESCE(enriched_at, '0001-01-01T00:00:00Z'::timestamptz)" + } + + panic(fmt.Sprintf("unsupported order by: %s", v)) +} diff --git a/pkg/coredata/detected_tracker.go b/pkg/coredata/detected_tracker.go index ba9d12b49..80aecf241 100644 --- a/pkg/coredata/detected_tracker.go +++ b/pkg/coredata/detected_tracker.go @@ -295,6 +295,99 @@ LIMIT @limit; return ids, nil } +// LoadAllByTrackerPatternID returns every detected tracker linked to the +// pattern, with no pagination. It backs the banner-reset rebuild, which +// recreates exact patterns from a glob's detections. +func (dts *DetectedTrackers) LoadAllByTrackerPatternID( + ctx context.Context, + conn pg.Querier, + scope Scoper, + trackerPatternID gid.GID, +) error { + q := ` +SELECT + id, + cookie_banner_id, + tracker_pattern_id, + tracker_type, + identifier, + max_age_seconds, + source, + value_size, + initiator_url, + initiator_domain, + last_detected_at, + created_at, + updated_at +FROM + detected_trackers +WHERE + %s + AND tracker_pattern_id = @tracker_pattern_id +ORDER BY + identifier ASC, id ASC +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"tracker_pattern_id": trackerPatternID} + maps.Copy(args, scope.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query detected trackers: %w", err) + } + + trackers, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[DetectedTracker]) + if err != nil { + return fmt.Errorf("cannot collect detected trackers: %w", err) + } + + *dts = trackers + + return nil +} + +// UpdateTrackerPatternID repoints a single detected tracker at another +// pattern. It is the per-row counterpart of RelinkByTrackerPatternID, +// used by the banner-reset rebuild where each detection of a glob moves +// to its own recreated exact pattern. +func (dt *DetectedTracker) UpdateTrackerPatternID( + ctx context.Context, + tx pg.Tx, + scope Scoper, +) error { + q := ` +UPDATE detected_trackers +SET + tracker_pattern_id = @tracker_pattern_id, + updated_at = @updated_at +WHERE + %s + AND id = @id +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{ + "id": dt.ID, + "tracker_pattern_id": dt.TrackerPatternID, + "updated_at": time.Now(), + } + maps.Copy(args, scope.SQLArguments()) + + result, err := tx.Exec(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot update detected tracker pattern: %w", err) + } + + if result.RowsAffected() == 0 { + return ErrResourceNotFound + } + + return nil +} + func (dts *DetectedTrackers) RelinkByTrackerPatternID( ctx context.Context, tx pg.Tx, diff --git a/pkg/coredata/tracker_pattern.go b/pkg/coredata/tracker_pattern.go index 39a5ee4fc..1bf103468 100644 --- a/pkg/coredata/tracker_pattern.go +++ b/pkg/coredata/tracker_pattern.go @@ -1382,3 +1382,122 @@ WHERE return result.RowsAffected(), nil } + +// ResetAndRequestMappingByCookieCategoryID detaches every pattern in the +// given category from its catalog row, org third party, and copied +// description, then re-arms mapping. Operators run this (via proboctl) on +// a banner's uncategorised category to force a clean re-map when +// iterating on the mapping agent. Excluded patterns are left untouched - +// exclusion is a deliberate suppression. The cookie_category_id key +// scopes the reset to the uncategorised category the caller resolves; +// the Scoper keeps it tenant-isolated. Returns the number of patterns +// reset. +func (tps *TrackerPatterns) ResetAndRequestMappingByCookieCategoryID( + ctx context.Context, + tx pg.Tx, + scope Scoper, + cookieCategoryID gid.GID, +) (int64, error) { + q := ` +UPDATE tracker_patterns +SET + common_tracker_pattern_id = NULL, + third_party_id = NULL, + description = '', + mapping_requested_at = NOW(), + updated_at = NOW() +WHERE + %s + AND cookie_category_id = @cookie_category_id + AND excluded = false +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"cookie_category_id": cookieCategoryID} + maps.Copy(args, scope.SQLArguments()) + + result, err := tx.Exec(ctx, q, args) + if err != nil { + return 0, fmt.Errorf("cannot reset and request mapping by cookie category: %w", err) + } + + return result.RowsAffected(), nil +} + +// LoadAllLinkedCommonTrackerPatternIDsByCookieBannerID returns every +// distinct common_tracker_pattern_id referenced by the banner's patterns, +// regardless of mapping state. Unlike +// LoadDistinctCommonTrackerPatternIDsByCookieBannerID (which restricts to +// unmapped patterns for the mapping pipeline), this returns the full set +// of catalog rows the banner depends on, so an operator can re-describe +// exactly those before a reset. +func (tps *TrackerPatterns) LoadAllLinkedCommonTrackerPatternIDsByCookieBannerID( + ctx context.Context, + conn pg.Querier, + scope Scoper, + cookieBannerID gid.GID, +) ([]gid.GID, error) { + q := ` +SELECT DISTINCT common_tracker_pattern_id +FROM tracker_patterns +WHERE + %s + AND cookie_banner_id = @cookie_banner_id + AND common_tracker_pattern_id IS NOT NULL +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"cookie_banner_id": cookieBannerID} + maps.Copy(args, scope.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return nil, fmt.Errorf("cannot query linked common tracker pattern ids: %w", err) + } + + ids, err := pgx.CollectRows(rows, pgx.RowTo[gid.GID]) + if err != nil { + return nil, fmt.Errorf("cannot collect linked common tracker pattern ids: %w", err) + } + + return ids, nil +} + +// LoadAllLinkedCommonTrackerPatternIDsByOrganizationID is the org-wide +// counterpart of LoadAllLinkedCommonTrackerPatternIDsByCookieBannerID: +// every distinct catalog row the organization's tracker patterns depend +// on, regardless of mapping state. +func (tps *TrackerPatterns) LoadAllLinkedCommonTrackerPatternIDsByOrganizationID( + ctx context.Context, + conn pg.Querier, + scope Scoper, + organizationID gid.GID, +) ([]gid.GID, error) { + q := ` +SELECT DISTINCT common_tracker_pattern_id +FROM tracker_patterns +WHERE + %s + AND organization_id = @organization_id + AND common_tracker_pattern_id IS NOT NULL +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"organization_id": organizationID} + maps.Copy(args, scope.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return nil, fmt.Errorf("cannot query linked common tracker pattern ids: %w", err) + } + + ids, err := pgx.CollectRows(rows, pgx.RowTo[gid.GID]) + if err != nil { + return nil, fmt.Errorf("cannot collect linked common tracker pattern ids: %w", err) + } + + return ids, nil +} diff --git a/pkg/probod/llm.go b/pkg/probod/llm.go index a53a04082..003e204fe 100644 --- a/pkg/probod/llm.go +++ b/pkg/probod/llm.go @@ -15,20 +15,17 @@ 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" ) // 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. +// in error messages. It delegates to pkg/agentsbuild so probod and proboctl +// wire LLM clients identically. func (impl *Implm) resolveAgentClient( name string, agent LLMAgentConfig, @@ -36,61 +33,5 @@ func (impl *Implm) resolveAgentClient( tp trace.TracerProvider, r prometheus.Registerer, ) (LLMAgentConfig, *llm.Client, error) { - 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 := 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 -} - -func 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) - } + return agentsbuild.ResolveAgentClient(impl.cfg.Agents, name, agent, l, tp, r) } diff --git a/pkg/probod/tracker_agents.go b/pkg/probod/tracker_agents.go index 7341ef908..f93011d76 100644 --- a/pkg/probod/tracker_agents.go +++ b/pkg/probod/tracker_agents.go @@ -15,78 +15,22 @@ 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 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 that the tracker-mapping worker uses -// to promote patterns to org ThirdParties. All are opt-in: deployments -// that do not set `llm.tracker-mapping.provider` get zero configs (nil -// LLM client) so the workers run without agent fallback. -// -// The agents are sourced from the same `tracker-mapping` config slot -// because they share the LLM client, model, and lifecycle. The -// disambiguation agent has no Firecrawl/DB tools, so its config -// surface is narrower and it lives in the cross-domain pkg/thirdparty -// package. +// 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( l *log.Logger, tp trace.TracerProvider, r prometheus.Registerer, ) (cookiebanner.TrackerAgentsConfig, thirdparty.DisambiguationConfig, error) { - if impl.cfg.Agents.TrackerMapping.Provider == "" { - return cookiebanner.TrackerAgentsConfig{}, thirdparty.DisambiguationConfig{}, nil - } - - agentCfg, llmClient, err := impl.resolveAgentClient( - "tracker-mapping", - impl.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 := impl.cfg.TrackerMappingWorker - enrichmentWorkerCfg := impl.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: impl.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 + return agentsbuild.BuildTrackerAgentsConfig(impl.cfg, l, tp, r) }