diff --git a/pkg/cookiebanner/common_pattern_enrichment_agent.go b/pkg/cookiebanner/common_pattern_enrichment_agent.go new file mode 100644 index 000000000..a1d216637 --- /dev/null +++ b/pkg/cookiebanner/common_pattern_enrichment_agent.go @@ -0,0 +1,91 @@ +// 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 ( + _ "embed" + "fmt" + "strings" + + "go.gearno.de/kit/log" + "go.gearno.de/kit/pg" + "go.probo.inc/probo/pkg/agent" + "go.probo.inc/probo/pkg/agent/tools/search" + "go.probo.inc/probo/pkg/coredata" +) + +//go:embed prompts/tracker_enrichment.txt.tmpl +var trackerEnrichmentPrompt string + +// CommonPatternEnrichmentResult is the structured output the +// common-pattern enrichment agent returns. +type CommonPatternEnrichmentResult struct { + Description string `json:"description" jsonschema:"A concise, factual, compliance-grade description of what this tracker stores or does and its purpose. One or two sentences. Name the operating company when known."` +} + +func buildCommonPatternEnrichmentAgent( + cfg TrackerAgentsConfig, + pgClient *pg.Client, + logger *log.Logger, +) *agent.Agent { + tools := []agent.Tool{ + searchThirdPartiesTool(pgClient), + } + + if cfg.FirecrawlAPIKey != "" { + tools = append(tools, search.FirecrawlSearchTool(cfg.FirecrawlAPIKey)) + } + + outputType, err := agent.NewOutputType[CommonPatternEnrichmentResult]("tracker_enrichment") + if err != nil { + panic(fmt.Sprintf("cookiebanner: cannot build tracker enrichment output type: %s", err)) + } + + return agent.New( + "common-pattern-enrichment", + cfg.LLMClient, + agent.WithInstructions(trackerEnrichmentPrompt), + agent.WithModel(cfg.Model), + agent.WithTools(tools...), + agent.WithOutputType(outputType), + agent.WithMaxTurns(agentMaxTurns), + agent.WithLogger(logger), + ) +} + +func buildEnrichmentPrompt(cp coredata.CommonTrackerPattern, thirdPartyName string) string { + maxAge := "session" + if cp.MaxAgeSeconds != nil { + maxAge = fmt.Sprintf("%d seconds", *cp.MaxAgeSeconds) + } + + prompt := fmt.Sprintf( + "Describe the following tracker:\n\n"+ + " %s \n"+ + " %s \n"+ + " %s \n"+ + " %s \n", + cp.Pattern, + cp.TrackerType, + cp.MatchType, + maxAge, + ) + + if name := strings.TrimSpace(thirdPartyName); name != "" { + prompt += fmt.Sprintf(" %s \n", name) + } + + return prompt +} diff --git a/pkg/cookiebanner/common_pattern_enrichment_worker.go b/pkg/cookiebanner/common_pattern_enrichment_worker.go new file mode 100644 index 000000000..317d69343 --- /dev/null +++ b/pkg/cookiebanner/common_pattern_enrichment_worker.go @@ -0,0 +1,208 @@ +// 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" + "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/llm" +) + +const enrichmentStaleAfter = 10 * time.Minute + +type commonPatternEnrichmentHandler struct { + pg *pg.Client + logger *log.Logger + enrichmentAgent *agent.Agent + staleAfter time.Duration +} + +// NewCommonPatternEnrichmentWorker builds the worker that fills +// descriptions on common_tracker_patterns using an agent with web +// search, then fans the result out to every linked tracker pattern. It is +// a global system worker: common_tracker_patterns is not tenant-scoped, +// so a single enrichment benefits all tenants. The worker no-ops when no +// LLM client is configured; callers should gate registration on config +// presence. +func NewCommonPatternEnrichmentWorker( + pgClient *pg.Client, + logger *log.Logger, + cfg TrackerAgentsConfig, + opts ...worker.Option, +) *worker.Worker[coredata.CommonTrackerPattern] { + h := &commonPatternEnrichmentHandler{ + pg: pgClient, + logger: logger, + staleAfter: enrichmentStaleAfter, + } + + if cfg.LLMClient != nil { + h.enrichmentAgent = buildCommonPatternEnrichmentAgent(cfg, pgClient, logger) + } + + return worker.New( + "common-pattern-enrichment-worker", + h, + logger, + opts..., + ) +} + +func (h *commonPatternEnrichmentHandler) Claim(ctx context.Context) (coredata.CommonTrackerPattern, error) { + var cp coredata.CommonTrackerPattern + + if err := h.pg.WithTx( + ctx, + func(ctx context.Context, tx pg.Tx) error { + if err := cp.LoadNextForEnrichmentForUpdateSkipLocked(ctx, tx); err != nil { + return err + } + + return cp.ClearEnrichmentRequestedAt(ctx, tx) + }, + ); err != nil { + if errors.Is(err, coredata.ErrResourceNotFound) { + return coredata.CommonTrackerPattern{}, worker.ErrNoTask + } + + return coredata.CommonTrackerPattern{}, fmt.Errorf("cannot claim common tracker pattern enrichment task: %w", err) + } + + return cp, nil +} + +func (h *commonPatternEnrichmentHandler) Process(ctx context.Context, cp coredata.CommonTrackerPattern) error { + if h.enrichmentAgent == nil { + return nil + } + + thirdPartyName, err := h.loadThirdPartyName(ctx, cp) + if err != nil { + return err + } + + description, err := h.research(ctx, cp, thirdPartyName) + if err != nil { + return fmt.Errorf("cannot research tracker description: %w", err) + } + + if description == "" { + return fmt.Errorf("enrichment produced empty description for pattern %q", cp.Pattern) + } + + return h.pg.WithTx( + ctx, + func(ctx context.Context, tx pg.Tx) error { + if err := cp.SetEnriched(ctx, tx, description); err != nil { + return fmt.Errorf("cannot set common tracker pattern enriched: %w", err) + } + + var patterns coredata.TrackerPatterns + + count, 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.Int64("backfilled_tracker_patterns", count), + ) + + return nil + }, + ) +} + +func (h *commonPatternEnrichmentHandler) RecoverStale(ctx context.Context) error { + return h.pg.WithConn( + ctx, + func(ctx context.Context, conn pg.Querier) error { + if err := coredata.ResetStaleEnrichments(ctx, conn, h.staleAfter); err != nil { + return fmt.Errorf("cannot reset stale common tracker pattern enrichments: %w", err) + } + + return nil + }, + ) +} + +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, 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 +} diff --git a/pkg/cookiebanner/prompts/tracker_enrichment.txt.tmpl b/pkg/cookiebanner/prompts/tracker_enrichment.txt.tmpl new file mode 100644 index 000000000..ad30767b8 --- /dev/null +++ b/pkg/cookiebanner/prompts/tracker_enrichment.txt.tmpl @@ -0,0 +1,27 @@ + +You are a privacy and web-tracking compliance expert. Your job is to write an accurate, source-grounded description of what a given cookie or web tracker does, for use in a privacy/compliance register. + + + +Given a tracker pattern (cookie name, local storage key, etc.), its type, max-age, and the third party that operates it when known, produce a concise factual description of the tracker's purpose. + +Return a structured JSON response with: +- description: one or two sentences describing what this tracker stores or does and the purpose it serves (e.g. analytics, advertising, session management, security). When the operating company is known, name it. + + + +1. Use the search_third_parties tool to confirm details about the operating company when one is associated with this tracker. + +2. Use web_search to find authoritative information about the tracker's purpose. Try up to 3 targeted queries, adapting to the available signals: + - With a recognizable prefix or name: "[name] cookie purpose" (e.g. "_ga cookie purpose"). + - For localStorage keys: "[name] localStorage purpose tracking". + - Broaden if needed: "[name] cookie what is it used for". + - Stop once you have a confident, well-sourced answer; do not exhaust all queries if the first succeeds. + - Verify that any result discusses a tracker whose name shares a meaningful prefix with the pattern being described. Discard results about a differently-named tracker. + +3. Be factual and conservative. Describe only what the evidence supports. Do not speculate about data flows or purposes you cannot substantiate. + +4. Keep the description concise (one to two sentences) and free of marketing language. It should read as a neutral, compliance-grade statement of purpose. + +5. If you genuinely cannot determine the tracker's purpose, write a minimal factual description based on its type and name (e.g. "Cookie set by an unidentified third party; purpose could not be determined.") rather than inventing details. + diff --git a/pkg/cookiebanner/prompts/tracker_identification.txt.tmpl b/pkg/cookiebanner/prompts/tracker_identification.txt.tmpl index ed9f75690..6b9df3453 100644 --- a/pkg/cookiebanner/prompts/tracker_identification.txt.tmpl +++ b/pkg/cookiebanner/prompts/tracker_identification.txt.tmpl @@ -8,7 +8,6 @@ Given a tracker pattern (cookie name, local storage key, etc.), its type, max-ag Return a structured JSON response with: - third_party_name: the canonical company/service name (e.g. "Google Analytics", not "google" or "GA") - category: the business category of the third party -- description: a one-sentence description of what this tracker does - confidence: how confident you are in the identification (0.0 to 1.0) diff --git a/pkg/cookiebanner/tracker_agents_config.go b/pkg/cookiebanner/tracker_agents_config.go new file mode 100644 index 000000000..ee1fbdff9 --- /dev/null +++ b/pkg/cookiebanner/tracker_agents_config.go @@ -0,0 +1,28 @@ +// 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 "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. +type TrackerAgentsConfig struct { + LLMClient *llm.Client + Model string + FirecrawlAPIKey string +} diff --git a/pkg/cookiebanner/tracker_mapping_agent.go b/pkg/cookiebanner/tracker_mapping_agent.go index 13e02fe39..a0eea3e6a 100644 --- a/pkg/cookiebanner/tracker_mapping_agent.go +++ b/pkg/cookiebanner/tracker_mapping_agent.go @@ -26,7 +26,6 @@ import ( "go.probo.inc/probo/pkg/agent" "go.probo.inc/probo/pkg/agent/tools/search" "go.probo.inc/probo/pkg/coredata" - "go.probo.inc/probo/pkg/llm" ) const ( @@ -44,21 +43,11 @@ var trackerIdentificationPrompt string type TrackerMappingAgentResult struct { ThirdPartyName string `json:"third_party_name" jsonschema:"Name of the company or service that sets this tracker (e.g. 'Google Analytics', 'Meta Pixel'). Empty string if truly unknown."` Category coredata.ThirdPartyCategory `json:"category" jsonschema:"Third party category"` - Description string `json:"description" jsonschema:"What this tracker does in one sentence"` Confidence float64 `json:"confidence" jsonschema:"Confidence level from 0.0 to 1.0. Set below 0.5 if unsure."` } -// TrackerMappingConfig configures the tracker-mapping agent (catalog -// identification). The agent uses DB-backed search tools and may also -// use Firecrawl for web search when an API key is supplied. -type TrackerMappingConfig struct { - LLMClient *llm.Client - Model string - FirecrawlAPIKey string -} - func buildTrackerMappingAgent( - cfg TrackerMappingConfig, + cfg TrackerAgentsConfig, pgClient *pg.Client, logger *log.Logger, ) *agent.Agent { diff --git a/pkg/cookiebanner/tracker_mapping_worker.go b/pkg/cookiebanner/tracker_mapping_worker.go index f75b2a875..02c0ffd1a 100644 --- a/pkg/cookiebanner/tracker_mapping_worker.go +++ b/pkg/cookiebanner/tracker_mapping_worker.go @@ -42,7 +42,7 @@ type trackerMappingHandler struct { func NewTrackerMappingWorker( pgClient *pg.Client, logger *log.Logger, - mappingCfg TrackerMappingConfig, + mappingCfg TrackerAgentsConfig, disambiguationCfg thirdparty.DisambiguationConfig, opts ...worker.Option, ) *worker.Worker[coredata.TrackerPattern] { @@ -260,6 +260,11 @@ func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.Tracker tp.ThirdPartyID = thirdPartyID tp.UpdatedAt = time.Now() + // Descriptions are owned by the common-pattern enrichment + // worker. Here we only propagate: if the linked catalog row + // is already enriched, copy its description onto this + // pattern. A pattern linked before enrichment is filled + // later by the enrichment worker's fan-out instead. if tp.Description == "" && commonPatternID != nil { var commonPattern coredata.CommonTrackerPattern if err := commonPattern.LoadByID(ctx, tx, *commonPatternID); err == nil && commonPattern.Description != "" { @@ -451,7 +456,6 @@ func (h *trackerMappingHandler) matchByDomain( TrackerType: tp.TrackerType, Pattern: tp.Pattern, MatchType: tp.MatchType, - Description: tp.Description, MaxAgeSeconds: tp.MaxAgeSeconds, Confidence: 0.7, CreatedAt: now, @@ -547,7 +551,6 @@ func (h *trackerMappingHandler) identifyWithAgent( TrackerType: tp.TrackerType, Pattern: tp.Pattern, MatchType: tp.MatchType, - Description: identification.Description, MaxAgeSeconds: tp.MaxAgeSeconds, Confidence: confidence, CreatedAt: now, @@ -693,7 +696,6 @@ func (h *trackerMappingHandler) matchBySiblingOrigin( TrackerType: tp.TrackerType, Pattern: tp.Pattern, MatchType: tp.MatchType, - Description: tp.Description, MaxAgeSeconds: tp.MaxAgeSeconds, Confidence: 0.7, CreatedAt: now, @@ -822,7 +824,6 @@ func (h *trackerMappingHandler) createUnmatchedPattern( TrackerType: tp.TrackerType, Pattern: tp.Pattern, MatchType: tp.MatchType, - Description: tp.Description, MaxAgeSeconds: tp.MaxAgeSeconds, Confidence: 0.5, CreatedAt: now, diff --git a/pkg/coredata/common_tracker_pattern.go b/pkg/coredata/common_tracker_pattern.go index 5b88baa60..575fa0613 100644 --- a/pkg/coredata/common_tracker_pattern.go +++ b/pkg/coredata/common_tracker_pattern.go @@ -27,16 +27,18 @@ import ( type ( CommonTrackerPattern struct { - ID gid.GID `db:"id"` - CommonThirdPartyID *gid.GID `db:"common_third_party_id"` - TrackerType TrackerType `db:"tracker_type"` - Pattern string `db:"pattern"` - MatchType TrackerPatternMatchType `db:"match_type"` - Description string `db:"description"` - MaxAgeSeconds *int `db:"max_age_seconds"` - Confidence float32 `db:"confidence"` - CreatedAt time.Time `db:"created_at"` - UpdatedAt time.Time `db:"updated_at"` + ID gid.GID `db:"id"` + CommonThirdPartyID *gid.GID `db:"common_third_party_id"` + TrackerType TrackerType `db:"tracker_type"` + Pattern string `db:"pattern"` + MatchType TrackerPatternMatchType `db:"match_type"` + Description string `db:"description"` + MaxAgeSeconds *int `db:"max_age_seconds"` + Confidence float32 `db:"confidence"` + EnrichmentRequestedAt *time.Time `db:"enrichment_requested_at"` + EnrichedAt *time.Time `db:"enriched_at"` + CreatedAt time.Time `db:"created_at"` + UpdatedAt time.Time `db:"updated_at"` } CommonTrackerPatterns []*CommonTrackerPattern @@ -57,6 +59,8 @@ SELECT description, max_age_seconds, confidence, + enrichment_requested_at, + enriched_at, created_at, updated_at FROM @@ -104,6 +108,8 @@ SELECT description, max_age_seconds, confidence, + enrichment_requested_at, + enriched_at, created_at, updated_at FROM @@ -154,6 +160,8 @@ INSERT INTO common_tracker_patterns ( description, max_age_seconds, confidence, + enrichment_requested_at, + enriched_at, created_at, updated_at ) VALUES ( @@ -165,22 +173,26 @@ INSERT INTO common_tracker_patterns ( @description, @max_age_seconds, @confidence, + @enrichment_requested_at, + @enriched_at, @created_at, @updated_at ) ` args := pgx.StrictNamedArgs{ - "id": p.ID, - "common_third_party_id": p.CommonThirdPartyID, - "tracker_type": p.TrackerType, - "pattern": p.Pattern, - "match_type": p.MatchType, - "description": p.Description, - "max_age_seconds": p.MaxAgeSeconds, - "confidence": p.Confidence, - "created_at": p.CreatedAt, - "updated_at": p.UpdatedAt, + "id": p.ID, + "common_third_party_id": p.CommonThirdPartyID, + "tracker_type": p.TrackerType, + "pattern": p.Pattern, + "match_type": p.MatchType, + "description": p.Description, + "max_age_seconds": p.MaxAgeSeconds, + "confidence": p.Confidence, + "enrichment_requested_at": p.EnrichmentRequestedAt, + "enriched_at": p.EnrichedAt, + "created_at": p.CreatedAt, + "updated_at": p.UpdatedAt, } _, err := conn.Exec(ctx, q, args) @@ -195,6 +207,12 @@ func (p *CommonTrackerPattern) Upsert( ctx context.Context, conn pg.Tx, ) (inserted bool, err error) { + // On insert, a description-less row is immediately queued for the + // enrichment worker (enrichment_requested_at = NOW()). On conflict the + // enrichment columns are left untouched, and an empty incoming + // description never overwrites an existing one — descriptions are owned + // by the enrichment worker, so mapping-side upserts must not clobber a + // researched description with an empty string. q := ` INSERT INTO common_tracker_patterns ( id, @@ -205,6 +223,8 @@ INSERT INTO common_tracker_patterns ( description, max_age_seconds, confidence, + enrichment_requested_at, + enriched_at, created_at, updated_at ) VALUES ( @@ -216,6 +236,8 @@ INSERT INTO common_tracker_patterns ( @description, @max_age_seconds, @confidence, + CASE WHEN @description = '' THEN NOW() ELSE NULL END, + NULL, @created_at, @updated_at ) @@ -223,7 +245,10 @@ ON CONFLICT (tracker_type, pattern, COALESCE(max_age_seconds, -1)) DO UPDATE SET common_third_party_id = EXCLUDED.common_third_party_id, match_type = EXCLUDED.match_type, - description = EXCLUDED.description, + description = CASE + WHEN EXCLUDED.description = '' THEN common_tracker_patterns.description + ELSE EXCLUDED.description + END, confidence = EXCLUDED.confidence, updated_at = EXCLUDED.updated_at RETURNING @@ -235,6 +260,8 @@ RETURNING description, max_age_seconds, confidence, + enrichment_requested_at, + enriched_at, created_at, updated_at ` @@ -303,6 +330,8 @@ SELECT description, max_age_seconds, confidence, + enrichment_requested_at, + enriched_at, created_at, updated_at FROM @@ -418,6 +447,8 @@ SELECT description, max_age_seconds, confidence, + enrichment_requested_at, + enriched_at, created_at, updated_at FROM @@ -444,6 +475,151 @@ ORDER BY pattern ASC; return nil } +// LoadNextForEnrichmentForUpdateSkipLocked claims the next common tracker +// pattern queued for description enrichment, oldest request first. It +// mirrors the mapping worker's claim pattern: the row is locked FOR +// UPDATE SKIP LOCKED so concurrent enrichment workers never pick the same +// row. +func (p *CommonTrackerPattern) LoadNextForEnrichmentForUpdateSkipLocked( + ctx context.Context, + tx pg.Tx, +) 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 + enrichment_requested_at IS NOT NULL +ORDER BY + enrichment_requested_at ASC +FOR UPDATE SKIP LOCKED +LIMIT 1; +` + + rows, err := tx.Query(ctx, q) + if err != nil { + return fmt.Errorf("cannot query common tracker pattern for enrichment: %w", err) + } + + pattern, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[CommonTrackerPattern]) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return ErrResourceNotFound + } + + return fmt.Errorf("cannot collect common tracker pattern for enrichment: %w", err) + } + + *p = pattern + + return nil +} + +// ClearEnrichmentRequestedAt removes the row from the enrichment queue. It +// bumps updated_at so the stale-recovery clock starts at claim time. +func (p *CommonTrackerPattern) ClearEnrichmentRequestedAt( + ctx context.Context, + tx pg.Tx, +) error { + q := ` +UPDATE common_tracker_patterns +SET + enrichment_requested_at = NULL, + updated_at = NOW() +WHERE id = @id +` + + args := pgx.StrictNamedArgs{"id": p.ID} + + _, err := tx.Exec(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot clear enrichment requested at: %w", err) + } + + p.EnrichmentRequestedAt = nil + + return nil +} + +// SetEnriched records the researched description and marks the row +// terminally enriched so it is never re-queued. +func (p *CommonTrackerPattern) SetEnriched( + ctx context.Context, + tx pg.Tx, + description string, +) error { + q := ` +UPDATE common_tracker_patterns +SET + description = @description, + enriched_at = NOW(), + enrichment_requested_at = NULL, + updated_at = NOW() +WHERE id = @id +` + + args := pgx.StrictNamedArgs{ + "id": p.ID, + "description": description, + } + + result, err := tx.Exec(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot mark common tracker pattern enriched: %w", err) + } + + if result.RowsAffected() == 0 { + return ErrResourceNotFound + } + + p.Description = description + + return nil +} + +// ResetStaleEnrichments re-queues rows whose enrichment was claimed but +// never completed (no enriched_at, still description-less) and have been +// idle longer than staleAfter, so a crashed or timed-out enrichment is +// retried. +func ResetStaleEnrichments( + ctx context.Context, + conn pg.Querier, + staleAfter time.Duration, +) error { + q := ` +UPDATE common_tracker_patterns +SET + enrichment_requested_at = NOW(), + updated_at = NOW() +WHERE + enrichment_requested_at IS NULL + AND enriched_at IS NULL + AND description = '' + AND updated_at < @stale_before +` + + args := pgx.StrictNamedArgs{"stale_before": time.Now().Add(-staleAfter)} + + _, err := conn.Exec(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot reset stale common tracker pattern enrichments: %w", err) + } + + return nil +} + func (ps *CommonTrackerPatterns) LoadByIDs( ctx context.Context, conn pg.Querier, @@ -459,6 +635,8 @@ SELECT description, max_age_seconds, confidence, + enrichment_requested_at, + enriched_at, created_at, updated_at FROM diff --git a/pkg/coredata/migrations/20260529T091300Z.sql b/pkg/coredata/migrations/20260529T091300Z.sql new file mode 100644 index 000000000..4cd49501b --- /dev/null +++ b/pkg/coredata/migrations/20260529T091300Z.sql @@ -0,0 +1,31 @@ +-- 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. + +-- The common-pattern enrichment worker fills descriptions on +-- common_tracker_patterns using an agent with web search, then fans the +-- result out to every linked tracker pattern. enrichment_requested_at is +-- the work queue (claimed FOR UPDATE SKIP LOCKED); enriched_at marks a +-- row as terminally enriched so it is never re-enqueued. +ALTER TABLE common_tracker_patterns + ADD COLUMN enrichment_requested_at TIMESTAMP WITH TIME ZONE, + ADD COLUMN enriched_at TIMESTAMP WITH TIME ZONE; + +CREATE INDEX idx_common_tracker_patterns_enrichment + ON common_tracker_patterns (enrichment_requested_at) + WHERE enrichment_requested_at IS NOT NULL; + +-- Enqueue existing description-less rows for a first enrichment pass. +UPDATE common_tracker_patterns +SET enrichment_requested_at = NOW() +WHERE description = ''; diff --git a/pkg/coredata/tracker_pattern.go b/pkg/coredata/tracker_pattern.go index 52aa7505e..f806c1797 100644 --- a/pkg/coredata/tracker_pattern.go +++ b/pkg/coredata/tracker_pattern.go @@ -1303,3 +1303,40 @@ WHERE return result.RowsAffected(), nil } + +// BackfillDescriptionByCommonTrackerPatternID copies an enriched +// description onto every tracker pattern linked to the given common +// pattern that does not yet have one. It is invoked by the common-pattern +// enrichment worker, a global system process, so it is intentionally not +// tenant-scoped: a single catalog enrichment fans out to all tenants' +// linked patterns. The description = ” guard guarantees a pattern that +// already carries a description is never overwritten. Returns the number +// of patterns backfilled. +func (tps *TrackerPatterns) BackfillDescriptionByCommonTrackerPatternID( + ctx context.Context, + tx pg.Tx, + commonTrackerPatternID gid.GID, + description string, +) (int64, error) { + q := ` +UPDATE tracker_patterns +SET + description = @description, + updated_at = NOW() +WHERE + common_tracker_pattern_id = @common_tracker_pattern_id + AND description = '' +` + + args := pgx.StrictNamedArgs{ + "common_tracker_pattern_id": commonTrackerPatternID, + "description": description, + } + + result, err := tx.Exec(ctx, q, args) + if err != nil { + return 0, fmt.Errorf("cannot backfill tracker pattern descriptions: %w", err) + } + + return result.RowsAffected(), nil +} diff --git a/pkg/probod/probod.go b/pkg/probod/probod.go index ce912a45f..b912c4525 100644 --- a/pkg/probod/probod.go +++ b/pkg/probod/probod.go @@ -313,7 +313,7 @@ func (impl *Implm) Run( return err } - trackerMappingCfg, thirdPartyDisambiguationCfg, err := impl.buildTrackerMappingConfig(l, tp, r) + trackerAgentsCfg, thirdPartyDisambiguationCfg, err := impl.buildTrackerAgentsConfig(l, tp, r) if err != nil { return err } @@ -725,7 +725,7 @@ func (impl *Implm) Run( }, ) - trackerMappingWorker := cookiebanner.NewTrackerMappingWorker(pgClient, l, trackerMappingCfg, thirdPartyDisambiguationCfg) + trackerMappingWorker := cookiebanner.NewTrackerMappingWorker(pgClient, l, trackerAgentsCfg, thirdPartyDisambiguationCfg) trackerMappingWorkerCtx, stopTrackerMappingWorker := context.WithCancel(context.Background()) wg.Go( @@ -736,6 +736,26 @@ func (impl *Implm) Run( }, ) + // The common-pattern enrichment worker needs an LLM client (it + // researches descriptions via the agent), so it is only started when + // the tracker agents are configured. + stopCommonPatternEnrichmentWorker := func() {} + + if trackerAgentsCfg.LLMClient != nil { + commonPatternEnrichmentWorker := cookiebanner.NewCommonPatternEnrichmentWorker(pgClient, l, trackerAgentsCfg) + + var commonPatternEnrichmentWorkerCtx context.Context + commonPatternEnrichmentWorkerCtx, stopCommonPatternEnrichmentWorker = context.WithCancel(context.Background()) + + wg.Go( + func() { + if err := commonPatternEnrichmentWorker.Run(commonPatternEnrichmentWorkerCtx); err != nil { + cancel(fmt.Errorf("common pattern enrichment worker crashed: %w", err)) + } + }, + ) + } + mailingListWorker := mailman.NewMailingListWorker(mailmanService, pgClient, l.Named("mailing-list-worker")) mailingListWorkerCtx, stopMailingListWorker := context.WithCancel(context.Background()) @@ -805,6 +825,7 @@ func (impl *Implm) Run( stopESignService() stopTrackerPatternAnalysisWorker() stopTrackerMappingWorker() + stopCommonPatternEnrichmentWorker() stopMailingListWorker() stopEvidenceDescriptionWorker() stopDocumentPDFWorker() diff --git a/pkg/probod/tracker_mapping.go b/pkg/probod/tracker_agents.go similarity index 63% rename from pkg/probod/tracker_mapping.go rename to pkg/probod/tracker_agents.go index 7414b0d57..dd7048bf1 100644 --- a/pkg/probod/tracker_mapping.go +++ b/pkg/probod/tracker_agents.go @@ -24,25 +24,26 @@ import ( "go.probo.inc/probo/pkg/thirdparty" ) -// buildTrackerMappingConfig wires the tracker-mapping agent (catalog -// identification) and the third-party disambiguation agent that the -// tracker-mapping worker uses to promote patterns to org ThirdParties. -// Both are opt-in: deployments that do not set -// `llm.tracker-mapping.provider` get zero configs (nil LLM client) so -// the worker runs without agent fallback. +// 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. // -// Both agents are sourced from the same `tracker-mapping` config slot +// 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. -func (impl *Implm) buildTrackerMappingConfig( +func (impl *Implm) buildTrackerAgentsConfig( l *log.Logger, tp trace.TracerProvider, r prometheus.Registerer, -) (cookiebanner.TrackerMappingConfig, thirdparty.DisambiguationConfig, error) { +) (cookiebanner.TrackerAgentsConfig, thirdparty.DisambiguationConfig, error) { if impl.cfg.Agents.TrackerMapping.Provider == "" { - return cookiebanner.TrackerMappingConfig{}, thirdparty.DisambiguationConfig{}, nil + return cookiebanner.TrackerAgentsConfig{}, thirdparty.DisambiguationConfig{}, nil } agentCfg, llmClient, err := impl.resolveAgentClient( @@ -53,10 +54,10 @@ func (impl *Implm) buildTrackerMappingConfig( r, ) if err != nil { - return cookiebanner.TrackerMappingConfig{}, thirdparty.DisambiguationConfig{}, fmt.Errorf("cannot resolve tracker mapping agent client: %w", err) + return cookiebanner.TrackerAgentsConfig{}, thirdparty.DisambiguationConfig{}, fmt.Errorf("cannot resolve tracker mapping agent client: %w", err) } - mappingCfg := cookiebanner.TrackerMappingConfig{ + trackerAgentsCfg := cookiebanner.TrackerAgentsConfig{ LLMClient: llmClient, Model: agentCfg.ModelName, FirecrawlAPIKey: impl.cfg.Agents.Tools.FirecrawlAPIKey, @@ -67,5 +68,5 @@ func (impl *Implm) buildTrackerMappingConfig( Model: agentCfg.ModelName, } - return mappingCfg, disambiguationCfg, nil + return trackerAgentsCfg, disambiguationCfg, nil }