diff --git a/contrib/helm/charts/probo/CHANGELOG.md b/contrib/helm/charts/probo/CHANGELOG.md index 607c60acc..315e6e925 100644 --- a/contrib/helm/charts/probo/CHANGELOG.md +++ b/contrib/helm/charts/probo/CHANGELOG.md @@ -4,10 +4,6 @@ All notable changes to the Probo Helm chart will be documented in this file. ## Unreleased -### Removed - -- `thirdPartyDisambiguation` agent config slot and `trackerMappingWorker.disambiguationAgentTimeout`; the tracker-mapping worker no longer runs a third-party disambiguation agent - ## [0.5.0] - 2026-06-09 ### Added diff --git a/contrib/helm/charts/probo/templates/deployment.yaml b/contrib/helm/charts/probo/templates/deployment.yaml index db8f724fe..795fc1bbd 100644 --- a/contrib/helm/charts/probo/templates/deployment.yaml +++ b/contrib/helm/charts/probo/templates/deployment.yaml @@ -281,6 +281,23 @@ spec: - name: THIRD_PARTY_VETTING_STALE_AFTER value: {{ .Values.probo.thirdPartyVettingWorker.staleAfter | quote }} {{- end }} + # Third-party Disambiguation Agent + {{- if .Values.probo.thirdPartyDisambiguation.provider }} + - name: AGENT_THIRD_PARTY_DISAMBIGUATION_PROVIDER + value: {{ .Values.probo.thirdPartyDisambiguation.provider | quote }} + {{- end }} + {{- if .Values.probo.thirdPartyDisambiguation.modelName }} + - name: AGENT_THIRD_PARTY_DISAMBIGUATION_MODEL_NAME + value: {{ .Values.probo.thirdPartyDisambiguation.modelName | quote }} + {{- end }} + {{- if .Values.probo.thirdPartyDisambiguation.temperature }} + - name: AGENT_THIRD_PARTY_DISAMBIGUATION_TEMPERATURE + value: {{ .Values.probo.thirdPartyDisambiguation.temperature | quote }} + {{- end }} + {{- if .Values.probo.thirdPartyDisambiguation.maxTokens }} + - name: AGENT_THIRD_PARTY_DISAMBIGUATION_MAX_TOKENS + value: {{ .Values.probo.thirdPartyDisambiguation.maxTokens | quote }} + {{- end }} # Tracker Mapping Agent {{- if .Values.probo.trackerMapping.provider }} - name: AGENT_TRACKER_MAPPING_PROVIDER @@ -336,6 +353,10 @@ spec: - name: TRACKER_MAPPING_AGENT_MAX_TURNS value: {{ .Values.probo.trackerMappingWorker.agentMaxTurns | quote }} {{- end }} + {{- if .Values.probo.trackerMappingWorker.disambiguationAgentTimeout }} + - name: TRACKER_MAPPING_DISAMBIGUATION_AGENT_TIMEOUT + value: {{ .Values.probo.trackerMappingWorker.disambiguationAgentTimeout | quote }} + {{- end }} # Common Pattern Enrichment Worker {{- if .Values.probo.commonPatternEnrichmentWorker.interval }} - name: COMMON_PATTERN_ENRICHMENT_INTERVAL diff --git a/contrib/helm/charts/probo/values-production.yaml.example b/contrib/helm/charts/probo/values-production.yaml.example index f302ad9e5..4de046539 100644 --- a/contrib/helm/charts/probo/values-production.yaml.example +++ b/contrib/helm/charts/probo/values-production.yaml.example @@ -185,6 +185,15 @@ probo: # maxConcurrency: 1 # staleAfter: 1500 + # Third-party disambiguation agent (optional; runs inside the tracker + # mapping worker to pick the best matching org third party). Falls back + # to trackerMapping when its provider is unset. + # thirdPartyDisambiguation: + # provider: "openai" + # modelName: "gpt-4o" + # temperature: "0.4" + # maxTokens: "4096" + # Tracker mapping agent (optional, auto-links tracker patterns to vendors) # trackerMapping: # provider: "openai" @@ -201,7 +210,7 @@ probo: # maxTokens: "4096" # Tracker mapping worker tuning (optional; seconds for - # interval/staleAfter/agentTimeout). + # interval/staleAfter/agentTimeout/disambiguationAgentTimeout). # Keep concurrency modest to stay under OpenAI/Firecrawl limits and the DB pool. # trackerMappingWorker: # interval: 10 @@ -209,6 +218,7 @@ probo: # staleAfter: 600 # agentTimeout: 45 # agentMaxTurns: 10 + # disambiguationAgentTimeout: 45 # Common-pattern enrichment worker tuning (optional; seconds for # interval/staleAfter/agentTimeout). diff --git a/contrib/helm/charts/probo/values.yaml b/contrib/helm/charts/probo/values.yaml index 9724e2381..27e2d8656 100644 --- a/contrib/helm/charts/probo/values.yaml +++ b/contrib/helm/charts/probo/values.yaml @@ -286,6 +286,15 @@ probo: maxConcurrency: 1 staleAfter: 1500 + # Third-party disambiguation agent (optional; runs inside the tracker + # mapping worker to pick the best matching org third party). Falls back + # to trackerMapping when its provider is unset. + thirdPartyDisambiguation: + provider: "" + modelName: "" + temperature: "" + maxTokens: "" + # Tracker mapping agent (optional, requires openai.apiKey or anthropic key) trackerMapping: provider: "" @@ -302,15 +311,16 @@ probo: maxTokens: "" # Tracker mapping background worker tuning (optional). interval, - # staleAfter, and agentTimeout are in seconds. Keep concurrency modest - # to stay under OpenAI/Firecrawl rate limits and the database - # connection pool. + # staleAfter, agentTimeout, and disambiguationAgentTimeout are in + # seconds. Keep concurrency modest to stay under OpenAI/Firecrawl rate + # limits and the database connection pool. trackerMappingWorker: interval: 10 maxConcurrency: 3 staleAfter: 600 agentTimeout: 45 agentMaxTurns: 10 + disambiguationAgentTimeout: 45 # Common-pattern enrichment background worker tuning (optional). # interval, staleAfter, and agentTimeout are in seconds. diff --git a/pkg/bootstrap/builder.go b/pkg/bootstrap/builder.go index b065293cd..2eac4449d 100644 --- a/pkg/bootstrap/builder.go +++ b/pkg/bootstrap/builder.go @@ -213,6 +213,17 @@ func (b *Builder) Build() (*probodconfig.FullConfig, error) { Temperature: b.getEnvFloatPtr("AGENT_THIRD_PARTY_VETTER_TEMPERATURE"), MaxTokens: b.getEnvIntPtr("AGENT_THIRD_PARTY_VETTER_MAX_TOKENS"), }, + ThirdPartyDisambiguation: probodconfig.LLMAgentConfig{ + Provider: b.getEnvOrDefault("AGENT_THIRD_PARTY_DISAMBIGUATION_PROVIDER", ""), + ModelName: b.getEnvOrDefault("AGENT_THIRD_PARTY_DISAMBIGUATION_MODEL_NAME", ""), + // The disambiguation agent emits a single id plus a + // short rationale, but the budget must leave headroom + // for reasoning models whose reasoning tokens count + // against max_tokens; too small a budget truncates the + // JSON. + Temperature: b.getEnvFloatPtr("AGENT_THIRD_PARTY_DISAMBIGUATION_TEMPERATURE"), + MaxTokens: new(b.getEnvIntOrDefault("AGENT_THIRD_PARTY_DISAMBIGUATION_MAX_TOKENS", 4096)), + }, TrackerMapping: probodconfig.LLMAgentConfig{ Provider: b.getEnvOrDefault("AGENT_TRACKER_MAPPING_PROVIDER", ""), ModelName: b.getEnvOrDefault("AGENT_TRACKER_MAPPING_MODEL_NAME", ""), @@ -267,11 +278,12 @@ func (b *Builder) Build() (*probodconfig.FullConfig, error) { MaxConcurrency: b.getEnvIntOrDefault("THIRD_PARTY_VETTING_MAX_CONCURRENCY", 1), }, TrackerMappingWorker: probodconfig.TrackerMappingWorkerConfig{ - Interval: b.getEnvIntOrDefault("TRACKER_MAPPING_INTERVAL", 10), - MaxConcurrency: b.getEnvIntOrDefault("TRACKER_MAPPING_MAX_CONCURRENCY", 3), - StaleAfter: b.getEnvIntOrDefault("TRACKER_MAPPING_STALE_AFTER", 600), - AgentTimeout: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_TIMEOUT", 45), - AgentMaxTurns: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_MAX_TURNS", 10), + Interval: b.getEnvIntOrDefault("TRACKER_MAPPING_INTERVAL", 10), + MaxConcurrency: b.getEnvIntOrDefault("TRACKER_MAPPING_MAX_CONCURRENCY", 3), + StaleAfter: b.getEnvIntOrDefault("TRACKER_MAPPING_STALE_AFTER", 600), + AgentTimeout: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_TIMEOUT", 45), + AgentMaxTurns: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_MAX_TURNS", 10), + DisambiguationAgentTimeout: b.getEnvIntOrDefault("TRACKER_MAPPING_DISAMBIGUATION_AGENT_TIMEOUT", 45), }, CommonPatternEnrichmentWorker: probodconfig.CommonPatternEnrichmentWorkerConfig{ Interval: b.getEnvIntOrDefault("COMMON_PATTERN_ENRICHMENT_INTERVAL", 10), diff --git a/pkg/bootstrap/builder_test.go b/pkg/bootstrap/builder_test.go index bbd3a296e..b8889ee6b 100644 --- a/pkg/bootstrap/builder_test.go +++ b/pkg/bootstrap/builder_test.go @@ -220,6 +220,10 @@ func TestBuilder_Build_Defaults(t *testing.T) { assert.Empty(t, cfg.Probod.Agents.ThirdPartyVetter.ModelName) assert.Nil(t, cfg.Probod.Agents.ThirdPartyVetter.Temperature) assert.Nil(t, cfg.Probod.Agents.ThirdPartyVetter.MaxTokens) + assert.Empty(t, cfg.Probod.Agents.ThirdPartyDisambiguation.Provider) + assert.Empty(t, cfg.Probod.Agents.ThirdPartyDisambiguation.ModelName) + assert.Nil(t, cfg.Probod.Agents.ThirdPartyDisambiguation.Temperature) + assert.Equal(t, new(4096), cfg.Probod.Agents.ThirdPartyDisambiguation.MaxTokens) assert.Empty(t, cfg.Probod.Agents.TrackerMapping.Provider) assert.Empty(t, cfg.Probod.Agents.TrackerMapping.ModelName) assert.Nil(t, cfg.Probod.Agents.TrackerMapping.Temperature) @@ -235,6 +239,7 @@ func TestBuilder_Build_Defaults(t *testing.T) { assert.Equal(t, 600, cfg.Probod.TrackerMappingWorker.StaleAfter) assert.Equal(t, 45, cfg.Probod.TrackerMappingWorker.AgentTimeout) assert.Equal(t, 10, cfg.Probod.TrackerMappingWorker.AgentMaxTurns) + assert.Equal(t, 45, cfg.Probod.TrackerMappingWorker.DisambiguationAgentTimeout) assert.Equal(t, 10, cfg.Probod.CommonPatternEnrichmentWorker.Interval) assert.Equal(t, 2, cfg.Probod.CommonPatternEnrichmentWorker.MaxConcurrency) assert.Equal(t, 600, cfg.Probod.CommonPatternEnrichmentWorker.StaleAfter) @@ -337,6 +342,11 @@ func TestBuilder_Build_CustomValues(t *testing.T) { env["AGENT_THIRD_PARTY_VETTER_MODEL_NAME"] = "gpt-4o" env["AGENT_THIRD_PARTY_VETTER_TEMPERATURE"] = "0.3" env["AGENT_THIRD_PARTY_VETTER_MAX_TOKENS"] = "8192" + // Agents — third-party-disambiguation override + env["AGENT_THIRD_PARTY_DISAMBIGUATION_PROVIDER"] = "anthropic" + env["AGENT_THIRD_PARTY_DISAMBIGUATION_MODEL_NAME"] = "claude-sonnet-4-20250514" + env["AGENT_THIRD_PARTY_DISAMBIGUATION_TEMPERATURE"] = "0.4" + env["AGENT_THIRD_PARTY_DISAMBIGUATION_MAX_TOKENS"] = "2048" // Agents — tracker-mapping override env["AGENT_TRACKER_MAPPING_PROVIDER"] = "openai" env["AGENT_TRACKER_MAPPING_MODEL_NAME"] = "gpt-4o-mini" @@ -353,6 +363,7 @@ func TestBuilder_Build_CustomValues(t *testing.T) { env["TRACKER_MAPPING_STALE_AFTER"] = "1200" env["TRACKER_MAPPING_AGENT_TIMEOUT"] = "30" env["TRACKER_MAPPING_AGENT_MAX_TURNS"] = "6" + env["TRACKER_MAPPING_DISAMBIGUATION_AGENT_TIMEOUT"] = "35" env["COMMON_PATTERN_ENRICHMENT_INTERVAL"] = "15" env["COMMON_PATTERN_ENRICHMENT_MAX_CONCURRENCY"] = "4" env["COMMON_PATTERN_ENRICHMENT_STALE_AFTER"] = "900" @@ -452,6 +463,11 @@ func TestBuilder_Build_CustomValues(t *testing.T) { assert.Equal(t, "gpt-4o", cfg.Probod.Agents.ThirdPartyVetter.ModelName) assert.Equal(t, new(0.3), cfg.Probod.Agents.ThirdPartyVetter.Temperature) assert.Equal(t, new(8192), cfg.Probod.Agents.ThirdPartyVetter.MaxTokens) + // Agents — third-party-disambiguation overrides + assert.Equal(t, "anthropic", cfg.Probod.Agents.ThirdPartyDisambiguation.Provider) + assert.Equal(t, "claude-sonnet-4-20250514", cfg.Probod.Agents.ThirdPartyDisambiguation.ModelName) + assert.Equal(t, new(0.4), cfg.Probod.Agents.ThirdPartyDisambiguation.Temperature) + assert.Equal(t, new(2048), cfg.Probod.Agents.ThirdPartyDisambiguation.MaxTokens) // Agents — tracker-mapping overrides assert.Equal(t, "openai", cfg.Probod.Agents.TrackerMapping.Provider) assert.Equal(t, "gpt-4o-mini", cfg.Probod.Agents.TrackerMapping.ModelName) @@ -468,6 +484,7 @@ func TestBuilder_Build_CustomValues(t *testing.T) { assert.Equal(t, 1200, cfg.Probod.TrackerMappingWorker.StaleAfter) assert.Equal(t, 30, cfg.Probod.TrackerMappingWorker.AgentTimeout) assert.Equal(t, 6, cfg.Probod.TrackerMappingWorker.AgentMaxTurns) + assert.Equal(t, 35, cfg.Probod.TrackerMappingWorker.DisambiguationAgentTimeout) assert.Equal(t, 15, cfg.Probod.CommonPatternEnrichmentWorker.Interval) assert.Equal(t, 4, cfg.Probod.CommonPatternEnrichmentWorker.MaxConcurrency) assert.Equal(t, 900, cfg.Probod.CommonPatternEnrichmentWorker.StaleAfter) diff --git a/pkg/cookiebanner/pattern_analysis_worker.go b/pkg/cookiebanner/pattern_analysis_worker.go index 4085b25e9..3fadbae0a 100644 --- a/pkg/cookiebanner/pattern_analysis_worker.go +++ b/pkg/cookiebanner/pattern_analysis_worker.go @@ -230,9 +230,10 @@ func (h *patternAnalysisHandler) Process(ctx context.Context, banner coredata.Co return fmt.Errorf("cannot promote source on glob pattern %q: %w", key.template, err) } - // A stronger source can unblock mapping (e.g. - // EXTENSION->SCRIPT lifts the creationAllowed - // gate), so re-arm mapping on the existing glob. + // A stronger source can unblock mapping (a fresh + // initiator domain lets matchByDomain/ + // matchBySiblingOrigin resolve a vendor), so + // re-arm mapping on the existing glob. if err := globPattern.SetMappingRequested(ctx, tx); err != nil { return fmt.Errorf("cannot request mapping after source promotion on glob pattern %q: %w", key.template, err) } @@ -894,9 +895,9 @@ func (h *patternAnalysisHandler) adoptUncategorisedPatterns( return false, fmt.Errorf("cannot promote source on glob pattern %q: %w", match.Pattern, err) } - // A stronger source can unblock mapping (e.g. - // EXTENSION->SCRIPT lifts the creationAllowed gate), so - // re-arm mapping on the adopted glob. + // A stronger source can unblock mapping (a fresh + // initiator domain lets matchByDomain/matchBySiblingOrigin + // resolve a vendor), so re-arm mapping on the adopted glob. if err := match.SetMappingRequested(ctx, tx); err != nil { return false, fmt.Errorf("cannot request mapping after source promotion on glob pattern %q: %w", match.Pattern, err) } diff --git a/pkg/cookiebanner/service.go b/pkg/cookiebanner/service.go index 934cee2fc..99d11218d 100644 --- a/pkg/cookiebanner/service.go +++ b/pkg/cookiebanner/service.go @@ -2284,9 +2284,8 @@ func (s *Service) reportDetectedTracker( // A stronger source can unblock mapping: the detection // upserted below carries a fresh initiator domain that - // matchByDomain/matchBySiblingOrigin can now use, and an - // EXTENSION->SCRIPT promotion lifts the creationAllowed - // gate. Re-arm mapping so the worker revisits the pattern. + // matchByDomain/matchBySiblingOrigin can now use. Re-arm + // mapping so the worker revisits the pattern. if err := matchedPattern.SetMappingRequested(ctx, tx); err != nil { return fmt.Errorf("cannot request mapping after source promotion on tracker pattern %q: %w", matchedPattern.Pattern, err) } diff --git a/pkg/cookiebanner/tracker_mapping_worker.go b/pkg/cookiebanner/tracker_mapping_worker.go index ea8e2f775..f38e3429c 100644 --- a/pkg/cookiebanner/tracker_mapping_worker.go +++ b/pkg/cookiebanner/tracker_mapping_worker.go @@ -39,17 +39,20 @@ import ( const defaultMappingStaleAfter = 10 * time.Minute type trackerMappingHandler struct { - pg *pg.Client - logger *log.Logger - mappingAgent *agent.Agent - agentTimeout time.Duration - staleAfter time.Duration + pg *pg.Client + logger *log.Logger + mappingAgent *agent.Agent + disambiguationAgent *agent.Agent + agentTimeout time.Duration + disambiguationTimeout time.Duration + staleAfter time.Duration } func NewTrackerMappingWorker( pgClient *pg.Client, logger *log.Logger, mappingCfg TrackerMappingAgentConfig, + disambiguationCfg thirdparty.DisambiguationAgentConfig, staleAfter time.Duration, opts ...worker.Option, ) *worker.Worker[coredata.TrackerPattern] { @@ -63,16 +66,21 @@ func NewTrackerMappingWorker( } h := &trackerMappingHandler{ - pg: pgClient, - logger: logger, - agentTimeout: agentTimeout, - staleAfter: staleAfter, + pg: pgClient, + logger: logger, + agentTimeout: agentTimeout, + disambiguationTimeout: disambiguationCfg.Timeout, + staleAfter: staleAfter, } if mappingCfg.LLMClient != nil { h.mappingAgent = buildTrackerMappingAgent(mappingCfg, pgClient, logger) } + if disambiguationCfg.LLMClient != nil { + h.disambiguationAgent = thirdparty.BuildDisambiguationAgent(disambiguationCfg, logger) + } + return worker.New( "tracker-mapping-worker", h, @@ -124,11 +132,14 @@ func (h *trackerMappingHandler) RecoverStale(ctx context.Context) error { // catalogMatch is the result of a single catalog signal. commonPatternID // is the catalog row the signal resolved (or backfilled); commonThirdPartyID -// is the catalog third party the signal discovered, when any. A nil -// *catalogMatch means the signal produced nothing. +// is the catalog third party the signal discovered, when any; thirdPartyID +// is an existing org ThirdParty the signal knows directly (e.g. a sibling +// pattern already promoted in the same organization). A nil *catalogMatch +// means the signal produced nothing. type catalogMatch struct { commonPatternID *gid.GID commonThirdPartyID *gid.GID + thirdPartyID *gid.GID } // Process resolves the catalog mapping for a tracker pattern and links it @@ -145,11 +156,10 @@ type catalogMatch struct { // common_tracker_pattern_id but its catalog row has no common third // party yet. // -// The worker no longer materializes per-org ThirdParty rows: it resolves -// the shared catalog link only. An org ThirdParty is created exclusively -// through the explicit per-vendor import action, which also backfills -// tracker_patterns.third_party_id; an already-set third_party_id is -// preserved here untouched. +// Org ThirdParty resolution only links to an existing party (even for +// uncategorised or extension-sourced patterns); it never creates a brand +// new org ThirdParty. Creating an org ThirdParty from a catalog vendor is +// done exclusively through the explicit ImportFromCommon action. func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.TrackerPattern) error { scope := coredata.NewScopeFromObjectID(tp.ID) @@ -174,6 +184,7 @@ func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.Tracker commonPatternID := det.commonPatternID commonThirdPartyID := det.commonThirdPartyID + directThirdPartyID := det.directThirdPartyID // Phase 2: tracker-mapping agent (no transaction). It runs only when // the deterministic signals could not resolve a catalog third party. @@ -205,6 +216,25 @@ func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.Tracker } } + // Phase 3: org ThirdParty resolution. The heuristic ranking and the + // disambiguation agent run without a transaction; only the final link + // touches the database (in a short transaction). + thirdPartyID := tp.ThirdPartyID + + if thirdPartyID == nil { + switch { + case directThirdPartyID != nil: + thirdPartyID = directThirdPartyID + case commonThirdPartyID != nil: + resolved, err := h.resolveOrgThirdParty(ctx, tp, *commonThirdPartyID) + if err != nil { + return fmt.Errorf("cannot resolve org third party: %w", err) + } + + thirdPartyID = resolved + } + } + // Phase 4: persist the pattern mapping in a short transaction. The // unmatched fallback keeps catalog coverage complete even when no // vendor was resolved. @@ -221,6 +251,7 @@ func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.Tracker } tp.CommonTrackerPatternID = commonPatternID + tp.ThirdPartyID = thirdPartyID tp.UpdatedAt = time.Now() // Descriptions are owned by the common-pattern enrichment @@ -290,6 +321,7 @@ type deterministicResult struct { origin string commonPatternID *gid.GID commonThirdPartyID *gid.GID + directThirdPartyID *gid.GID domains []string commonThirdPartyPreexisted bool } @@ -371,6 +403,7 @@ func (h *trackerMappingHandler) resolveDeterministic( if siblingMatch != nil { res.commonPatternID = firstNonNil(res.commonPatternID, siblingMatch.commonPatternID) res.commonThirdPartyID = siblingMatch.commonThirdPartyID + res.directThirdPartyID = siblingMatch.thirdPartyID } if res.commonThirdPartyID != nil { @@ -805,8 +838,10 @@ func (h *trackerMappingHandler) persistAgentIdentification( // matchBySiblingOrigin finds other tracker patterns on the same banner // that share initiator domains with the current pattern. Sharing an // origin across multiple detected patterns is a strong indicator of the -// same third party, so the common third party the siblings resolve to is -// upserted onto the catalog row. +// same third party. When the siblings resolve to a single existing org +// ThirdParty, that id is returned directly so promotion can link to it +// without re-running heuristics; otherwise the resolved common third +// party is upserted onto the catalog row. func (h *trackerMappingHandler) matchBySiblingOrigin( ctx context.Context, tx pg.Tx, @@ -837,14 +872,19 @@ func (h *trackerMappingHandler) matchBySiblingOrigin( scope := coredata.NewScopeFromObjectID(tp.ID) - commonThirdPartyID, err := h.resolveThirdPartyFromSiblings(ctx, tx, scope, siblingIDs) + commonThirdPartyID, thirdPartyID, err := h.resolveThirdPartyFromSiblings(ctx, tx, scope, siblingIDs) if err != nil { return nil, fmt.Errorf("cannot resolve third party from siblings: %w", err) } - // No catalog third party to record: leave catalog creation to a later - // signal or the unmatched fallback. + // No catalog third party to record: surface a directly-known org + // third party (if any) so promotion can still link to it, and leave + // catalog creation to a later signal or the unmatched fallback. if commonThirdPartyID == nil { + if thirdPartyID != nil { + return &catalogMatch{thirdPartyID: thirdPartyID}, nil + } + return nil, nil } @@ -876,25 +916,37 @@ func (h *trackerMappingHandler) matchBySiblingOrigin( return &catalogMatch{ commonPatternID: &commonPattern.ID, commonThirdPartyID: commonPattern.CommonThirdPartyID, + thirdPartyID: thirdPartyID, }, nil } // resolveThirdPartyFromSiblings inspects sibling patterns to resolve a -// single unambiguous catalog third party for backfill. It is resolved -// first from the siblings' org ThirdParties, then, when those carry none, -// from siblings' common_tracker_pattern rows. It returns nil when the -// siblings carry no catalog third party or disagree on one. +// third party. It returns two independent signals: a direct org +// ThirdParty (set only when the siblings share a single one — the +// strongest, same-org signal), and a single unambiguous catalog third +// party for backfill. The catalog third party is resolved first from the +// siblings' org ThirdParties, then, when those carry none, from siblings' +// common_tracker_pattern rows. Either signal may be nil; siblings that +// disagree on the catalog third party resolve it to nothing. func (h *trackerMappingHandler) resolveThirdPartyFromSiblings( ctx context.Context, conn pg.Querier, scope coredata.Scoper, siblingIDs []gid.GID, -) (commonThirdPartyID *gid.GID, err error) { +) (commonThirdPartyID *gid.GID, thirdPartyID *gid.GID, err error) { var patterns coredata.TrackerPatterns thirdPartyIDs, err := patterns.LoadDistinctThirdPartyIDsByIDs(ctx, conn, scope, siblingIDs) if err != nil { - return nil, fmt.Errorf("cannot load distinct third party ids from siblings: %w", err) + return nil, nil, fmt.Errorf("cannot load distinct third party ids from siblings: %w", err) + } + + // A single org third party shared across the siblings is the + // strongest, same-org signal: link to it directly. This is resolved + // independently from the catalog third party used for backfill. + if len(thirdPartyIDs) == 1 { + directID := thirdPartyIDs[0] + thirdPartyID = &directID } if len(thirdPartyIDs) > 0 { @@ -913,26 +965,29 @@ func (h *trackerMappingHandler) resolveThirdPartyFromSiblings( if len(commonIDs) == 1 { for id := range commonIDs { - return &id, nil + return &id, thirdPartyID, nil } } - // Siblings are linked to several different catalog third - // parties: do not guess one. + // Siblings are promoted to several different catalog third + // parties: do not guess one. A single shared org third party (if + // any) is still a safe direct link. if len(commonIDs) > 1 { - return nil, nil + return nil, thirdPartyID, nil } } // Fall back to siblings carrying only a common_tracker_pattern_id, or - // whose org ThirdParty is not itself linked to the catalog. + // whose org ThirdParty is not itself linked to the catalog. This is + // reached when the org-third-party scan above found no catalog third + // party, so it must not be short-circuited by a direct match. commonPatternIDs, err := patterns.LoadDistinctCommonTrackerPatternIDsByIDs(ctx, conn, scope, siblingIDs) if err != nil { - return nil, fmt.Errorf("cannot load distinct common tracker pattern ids from siblings: %w", err) + return nil, nil, fmt.Errorf("cannot load distinct common tracker pattern ids from siblings: %w", err) } if len(commonPatternIDs) == 0 { - return nil, nil + return nil, thirdPartyID, nil } commonIDs := make(map[gid.GID]struct{}) @@ -950,11 +1005,11 @@ func (h *trackerMappingHandler) resolveThirdPartyFromSiblings( if len(commonIDs) == 1 { for id := range commonIDs { - return &id, nil + return &id, thirdPartyID, nil } } - return nil, nil + return nil, thirdPartyID, nil } func (h *trackerMappingHandler) createUnmatchedPattern( @@ -980,3 +1035,217 @@ func (h *trackerMappingHandler) createUnmatchedPattern( return &commonPattern.ID, nil } + +// resolveOrgThirdParty resolves an org ThirdParty for the given pattern +// from a known catalog third party by linking to an existing party. The +// resolution order is: +// +// 1. Exact link by common_third_party_id (O(1)). +// 2. Heuristic match against the org's existing ThirdParty rows +// (lowercased name, suffix-stripped name, slug, website host, +// CommonThirdPartyDomain overlap). +// 3. Agent disambiguation when the heuristic is ambiguous. +// +// When none of these resolve an existing party, the function returns +// (nil, nil); it never creates a brand new org ThirdParty (that happens +// only through the explicit ImportFromCommon action). A confident +// heuristic/agent match is auto-tagged with common_third_party_id so +// subsequent resolutions hit the exact-link path in O(1). +func (h *trackerMappingHandler) resolveOrgThirdParty( + ctx context.Context, + tp coredata.TrackerPattern, + commonThirdPartyID gid.GID, +) (*gid.GID, error) { + scope := coredata.NewScopeFromObjectID(tp.ID) + + // Read phase: exact link, candidate ranking, and eligibility. No + // write or LLM call happens here. + var prep orgThirdPartyPrep + + if err := h.pg.WithConn( + ctx, + func(ctx context.Context, conn pg.Querier) error { + var err error + + prep, err = h.prepareOrgThirdParty(ctx, conn, scope, tp, commonThirdPartyID) + + return err + }, + ); err != nil { + return nil, err + } + + if prep.existingID != nil { + return prep.existingID, nil + } + + picked := prep.highConfidence + viaAgent := false + + // Agent phase (no transaction): disambiguate among the heuristic + // candidates when none scored high enough on its own. + if picked == nil && prep.eligibleForAgent && h.disambiguationAgent != nil { + matchedID, err := thirdparty.Disambiguate( + ctx, + h.disambiguationAgent, + h.logger, + prep.commonParty, + prep.commonDomains, + prep.agentSet, + h.disambiguationTimeout, + ) + if err != nil { + h.logger.WarnCtx( + ctx, + "third-party disambiguation agent failed", + log.Error(err), + log.String("tracker_pattern_id", tp.ID.String()), + ) + } + + if matchedID != nil { + for _, c := range prep.agentSet { + if c.ThirdParty.ID == *matchedID { + picked = c.ThirdParty + viaAgent = true + + break + } + } + } + } + + // Nothing to link: leave the pattern without an org third party. An + // org ThirdParty is created only through the explicit ImportFromCommon + // action, never here. + if picked == nil { + return nil, nil + } + + // Write phase: link the picked candidate to the catalog entry in a + // short transaction. + if err := h.pg.WithTx( + ctx, + func(ctx context.Context, tx pg.Tx) error { + if err := thirdparty.LinkToCommon(ctx, tx, scope, picked, commonThirdPartyID); err != nil { + return fmt.Errorf("cannot link third party to common: %w", err) + } + + if viaAgent { + h.logger.InfoCtx( + ctx, + "promoted tracker pattern via disambiguation agent", + log.String("tracker_pattern_id", tp.ID.String()), + log.String("third_party_id", picked.ID.String()), + ) + } else { + h.logger.InfoCtx( + ctx, + "promoted tracker pattern via heuristic match", + log.String("tracker_pattern_id", tp.ID.String()), + log.String("third_party_id", picked.ID.String()), + log.Float64("score", prep.highScore), + ) + } + + return nil + }, + ); err != nil { + return nil, err + } + + return &picked.ID, nil +} + +// orgThirdPartyPrep is the read-phase outcome for org ThirdParty +// resolution. existingID is set when an exact common-id link already +// exists (the other fields are then unused). Otherwise highConfidence +// holds a heuristic match at or above HighConfidenceScore (with +// highScore), or agentSet/eligibleForAgent describe the disambiguation +// candidates. +type orgThirdPartyPrep struct { + existingID *gid.GID + commonParty coredata.CommonThirdParty + commonDomains coredata.CommonThirdPartyDomains + agentSet []thirdparty.ScoredCandidate + highConfidence *coredata.ThirdParty + highScore float64 + eligibleForAgent bool +} + +// prepareOrgThirdParty performs the read-only work for org ThirdParty +// resolution: it checks for an exact common-id link, loads the catalog +// entry and the org's existing third parties, and ranks the candidates. +// It makes no writes and no LLM call. +func (h *trackerMappingHandler) prepareOrgThirdParty( + ctx context.Context, + conn pg.Querier, + scope coredata.Scoper, + tp coredata.TrackerPattern, + commonThirdPartyID gid.GID, +) (orgThirdPartyPrep, error) { + var prep orgThirdPartyPrep + + var existing coredata.ThirdParty + + err := existing.LoadByOrganizationIDAndCommonThirdPartyID( + ctx, + conn, + scope, + tp.OrganizationID, + commonThirdPartyID, + ) + if err == nil { + id := existing.ID + prep.existingID = &id + + return prep, nil + } + + if !errors.Is(err, coredata.ErrResourceNotFound) { + return prep, fmt.Errorf("cannot load org third party by common id: %w", err) + } + + if err := prep.commonParty.LoadByID(ctx, conn, commonThirdPartyID); err != nil { + return prep, fmt.Errorf("cannot load common third party: %w", err) + } + + if err := prep.commonDomains.LoadByCommonThirdPartyID(ctx, conn, commonThirdPartyID); err != nil { + return prep, fmt.Errorf("cannot load common third party domains: %w", err) + } + + firstLevel := 1 + + var orgThirdParties coredata.ThirdParties + if err := orgThirdParties.LoadAllByOrganizationID( + ctx, + conn, + scope, + tp.OrganizationID, + coredata.NewThirdPartyFilter(nil, &firstLevel, nil), + ); err != nil { + return prep, fmt.Errorf("cannot load org third parties: %w", err) + } + + ranked := thirdparty.RankCandidates(prep.commonParty, prep.commonDomains, orgThirdParties) + + if len(ranked) > 0 && ranked[0].Score >= thirdparty.HighConfidenceScore { + prep.highConfidence = ranked[0].ThirdParty + prep.highScore = ranked[0].Score + } else { + prep.agentSet = ranked + if len(prep.agentSet) > thirdparty.MaxAgentCandidates { + prep.agentSet = prep.agentSet[:thirdparty.MaxAgentCandidates] + } + + for _, c := range prep.agentSet { + if c.Score >= thirdparty.MinAgentScore { + prep.eligibleForAgent = true + + break + } + } + } + + return prep, nil +} diff --git a/pkg/cookiebanner/tracker_mapping_worker_test.go b/pkg/cookiebanner/tracker_mapping_worker_test.go index 363a9e4e5..d95af0e56 100644 --- a/pkg/cookiebanner/tracker_mapping_worker_test.go +++ b/pkg/cookiebanner/tracker_mapping_worker_test.go @@ -31,10 +31,10 @@ import ( // promotionFixture extends workerFixture with a CommonThirdParty and a // CommonTrackerPattern linking the catalog to the test pattern. It is -// the minimum scaffolding the catalog-resolution paths need to run -// end-to-end. +// the minimum scaffolding resolveOrgThirdParty needs to run end-to-end. type promotionFixture struct { workerFixture + commonThirdParty coredata.CommonThirdParty commonPatternID gid.GID trackerPattern coredata.TrackerPattern commonThirdPartyID gid.GID @@ -130,6 +130,7 @@ func seedPromotionFixture(t *testing.T, ctx context.Context, client *pg.Client) return promotionFixture{ workerFixture: fx, + commonThirdParty: commonThirdParty, commonPatternID: commonPattern.ID, commonThirdPartyID: commonThirdPartyID, trackerPattern: pattern, @@ -143,6 +144,112 @@ func newMappingHandler(client *pg.Client) *trackerMappingHandler { } } +// promote runs resolveOrgThirdParty, which manages its own short +// transactions internally and only links to an existing org ThirdParty +// (it never creates one). +func promote( + t *testing.T, + ctx context.Context, + h *trackerMappingHandler, + tp coredata.TrackerPattern, + commonThirdPartyID gid.GID, +) *gid.GID { + t.Helper() + + got, err := h.resolveOrgThirdParty(ctx, tp, commonThirdPartyID) + require.NoError(t, err) + + return got +} + +func TestPromoteThirdParty_ExactCommonLink(t *testing.T) { + t.Parallel() + + client := test.PGClient(t) + ctx := context.Background() + fx := seedPromotionFixture(t, ctx, client) + + now := time.Now().UTC().Truncate(time.Microsecond) + existing := coredata.ThirdParty{ + ID: gid.New(fx.scope.GetTenantID(), coredata.ThirdPartyEntityType), + OrganizationID: fx.organizationID, + CommonThirdPartyID: &fx.commonThirdPartyID, + Name: "Google LLC", + Category: coredata.ThirdPartyCategoryAnalytics, + Certifications: []string{}, + Countries: coredata.CountryCodes{}, + CreatedAt: now, + UpdatedAt: now, + } + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + return existing.Insert(ctx, tx, fx.scope) + })) + + got := promote(t, ctx, newMappingHandler(client), fx.trackerPattern, fx.commonThirdPartyID) + + require.NotNil(t, got) + assert.Equal(t, existing.ID, *got, "should return the existing org ThirdParty linked by common id") +} + +func TestPromoteThirdParty_HeuristicMatch(t *testing.T) { + t.Parallel() + + client := test.PGClient(t) + ctx := context.Background() + fx := seedPromotionFixture(t, ctx, client) + + now := time.Now().UTC().Truncate(time.Microsecond) + // Append a corporate suffix to the catalog name so the heuristic + // matches on the suffix-stripped name (score 0.9) rather than an + // exact link. + manualEntry := coredata.ThirdParty{ + ID: gid.New(fx.scope.GetTenantID(), coredata.ThirdPartyEntityType), + OrganizationID: fx.organizationID, + Name: fx.commonThirdParty.Name + " LLC", + Category: coredata.ThirdPartyCategoryAnalytics, + Certifications: []string{}, + Countries: coredata.CountryCodes{}, + Level: 1, + CreatedAt: now, + UpdatedAt: now, + } + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + return manualEntry.Insert(ctx, tx, fx.scope) + })) + + got := promote(t, ctx, newMappingHandler(client), fx.trackerPattern, fx.commonThirdPartyID) + + require.NotNil(t, got) + assert.Equal(t, manualEntry.ID, *got, "heuristic match should return the manually-entered ThirdParty") + + var reloaded coredata.ThirdParty + + require.NoError(t, client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { + return reloaded.LoadByID(ctx, conn, fx.scope, manualEntry.ID) + })) + + require.NotNil(t, reloaded.CommonThirdPartyID, "matched row must be tagged with common_third_party_id") + assert.Equal(t, fx.commonThirdPartyID, *reloaded.CommonThirdPartyID) +} + +// TestPromoteThirdParty_NoCreateWithoutMatch asserts that when no +// existing org ThirdParty matches the catalog third party, resolution +// returns nothing: the worker never creates a brand new org ThirdParty +// (that is done only through the explicit ImportFromCommon action). +func TestPromoteThirdParty_NoCreateWithoutMatch(t *testing.T) { + t.Parallel() + + client := test.PGClient(t) + ctx := context.Background() + fx := seedPromotionFixture(t, ctx, client) + + got := promote(t, ctx, newMappingHandler(client), fx.trackerPattern, fx.commonThirdPartyID) + + assert.Nil(t, got, "resolution must not create a new org ThirdParty") +} + // TestProcess_PreservesCatalogMappingOnReTrigger asserts that when // Process is called for a pattern that already carries a // common_tracker_pattern_id, the catalog pipeline is skipped and the @@ -169,7 +276,7 @@ func TestProcess_PreservesCatalogMappingOnReTrigger(t *testing.T) { require.NotNil(t, reloaded.CommonTrackerPatternID, "common tracker pattern link must be preserved") assert.Equal(t, fx.commonPatternID, *reloaded.CommonTrackerPatternID) - assert.Nil(t, reloaded.ThirdPartyID, "the worker must not auto-create or link an org ThirdParty") + assert.Nil(t, reloaded.ThirdPartyID, "no org ThirdParty exists to link, so third_party_id stays unset") } // TestProcess_UncategorisedPatternIsNotPromoted asserts that a pattern @@ -442,6 +549,8 @@ func TestMatchBySiblingOrigin_SiblingWithThirdPartyID(t *testing.T) { require.NotNil(t, got, "sibling origin match should return a catalog match") require.NotNil(t, got.commonPatternID, "sibling origin match should return a common tracker pattern ID") + require.NotNil(t, got.thirdPartyID, "sibling origin match should surface the sibling's org third party directly") + assert.Equal(t, orgThirdParty.ID, *got.thirdPartyID) var commonPattern coredata.CommonTrackerPattern @@ -874,11 +983,56 @@ func TestMatchBySiblingOrigin_ConvergentSiblings(t *testing.T) { assert.Equal(t, fx.commonThirdPartyID, *commonPattern.CommonThirdPartyID) } +func TestPromoteThirdParty_ExactCommonLinkIgnoresSimilarUnlinked(t *testing.T) { + t.Parallel() + + client := test.PGClient(t) + ctx := context.Background() + fx := seedPromotionFixture(t, ctx, client) + + now := time.Now().UTC().Truncate(time.Microsecond) + + manualEntry := coredata.ThirdParty{ + ID: gid.New(fx.scope.GetTenantID(), coredata.ThirdPartyEntityType), + OrganizationID: fx.organizationID, + Name: "Google LLC", + Category: coredata.ThirdPartyCategoryAnalytics, + Certifications: []string{}, + Countries: coredata.CountryCodes{}, + CreatedAt: now, + UpdatedAt: now, + } + + linked := coredata.ThirdParty{ + ID: gid.New(fx.scope.GetTenantID(), coredata.ThirdPartyEntityType), + OrganizationID: fx.organizationID, + CommonThirdPartyID: &fx.commonThirdPartyID, + Name: "Google", + Category: coredata.ThirdPartyCategoryAnalytics, + Certifications: []string{}, + Countries: coredata.CountryCodes{}, + CreatedAt: now, + UpdatedAt: now, + } + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + if err := manualEntry.Insert(ctx, tx, fx.scope); err != nil { + return err + } + + return linked.Insert(ctx, tx, fx.scope) + })) + + got := promote(t, ctx, newMappingHandler(client), fx.trackerPattern, fx.commonThirdPartyID) + + require.NotNil(t, got) + assert.Equal(t, linked.ID, *got, "exact-link path must short-circuit before the heuristic fires") +} + // TestProcess_BackfillsCommonThirdPartyFromSibling asserts that a pattern // linked to an unlinked catalog row (no common_third_party_id) gets its -// catalog row backfilled from a sibling signal. The worker resolves the -// catalog link only; it must not promote the pattern to an org -// ThirdParty. +// catalog row backfilled from a sibling signal, and is promoted directly +// to the sibling's existing org ThirdParty. func TestProcess_BackfillsCommonThirdPartyFromSibling(t *testing.T) { t.Parallel() @@ -1022,18 +1176,81 @@ func TestProcess_BackfillsCommonThirdPartyFromSibling(t *testing.T) { return reloadedTarget.LoadByID(ctx, conn, fx.scope, target.ID) })) - assert.Nil(t, reloadedTarget.ThirdPartyID, "target must not be auto-promoted to an org third party") + require.NotNil(t, reloadedTarget.ThirdPartyID, "target must be promoted to the sibling's org third party") + assert.Equal(t, orgThirdParty.ID, *reloadedTarget.ThirdPartyID) require.NotNil(t, reloadedTarget.CommonTrackerPatternID) assert.Equal(t, unlinkedCommon.ID, *reloadedTarget.CommonTrackerPatternID, "the existing catalog link must be preserved") } -// TestProcess_SiblingCatalogResolutionOnFirstPartyOrigin asserts that a -// pattern detected on the banner's own (first-party) origin is still -// grouped with its siblings sharing that origin for catalog resolution. -// Sibling matching is an org-local co-occurrence signal and must not be -// defeated by the first-party domain filter that only protects the global -// catalog (domain) match. -func TestProcess_SiblingCatalogResolutionOnFirstPartyOrigin(t *testing.T) { +// TestProcess_UncategorisedLinksExistingThirdParty asserts that an +// uncategorised pattern is still linked to an already-existing matching +// org ThirdParty (linking to an existing party is ungated); only the +// creation of a new party stays gated, as covered by +// TestProcess_UncategorisedPatternIsNotPromoted. +func TestProcess_UncategorisedLinksExistingThirdParty(t *testing.T) { + t.Parallel() + + client := test.PGClient(t) + ctx := context.Background() + fx := seedPromotionFixture(t, ctx, client) + + now := time.Now().UTC().Truncate(time.Microsecond) + existing := coredata.ThirdParty{ + ID: gid.New(fx.scope.GetTenantID(), coredata.ThirdPartyEntityType), + OrganizationID: fx.organizationID, + CommonThirdPartyID: &fx.commonThirdPartyID, + Name: "Google LLC", + Category: coredata.ThirdPartyCategoryAnalytics, + Certifications: []string{}, + Countries: coredata.CountryCodes{}, + CreatedAt: now, + UpdatedAt: now, + } + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + if err := existing.Insert(ctx, tx, fx.scope); err != nil { + return err + } + + _, err := tx.Exec( + ctx, + `UPDATE tracker_patterns + SET cookie_category_id = $1, + mapping_requested_at = $2 + WHERE id = $3`, + fx.uncategorisedID, + now, + fx.trackerPattern.ID, + ) + + return err + })) + + var reloadedBefore coredata.TrackerPattern + + require.NoError(t, client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { + return reloadedBefore.LoadByID(ctx, conn, fx.scope, fx.trackerPattern.ID) + })) + + h := newMappingHandler(client) + require.NoError(t, h.Process(ctx, reloadedBefore)) + + var reloaded coredata.TrackerPattern + + require.NoError(t, client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { + return reloaded.LoadByID(ctx, conn, fx.scope, fx.trackerPattern.ID) + })) + + require.NotNil(t, reloaded.ThirdPartyID, "uncategorised pattern must still link to an existing org third party") + assert.Equal(t, existing.ID, *reloaded.ThirdPartyID) +} + +// TestProcess_SiblingPromotionOnFirstPartyOrigin asserts that a pattern +// detected on the banner's own (first-party) origin is still grouped with +// its siblings sharing that origin. Sibling matching is an org-local +// co-occurrence signal and must not be defeated by the first-party domain +// filter that only protects the global catalog (domain) match. +func TestProcess_SiblingPromotionOnFirstPartyOrigin(t *testing.T) { t.Parallel() client := test.PGClient(t) @@ -1177,7 +1394,8 @@ func TestProcess_SiblingCatalogResolutionOnFirstPartyOrigin(t *testing.T) { return reloadedTarget.LoadByID(ctx, conn, fx.scope, target.ID) })) - assert.Nil(t, reloadedTarget.ThirdPartyID, "target must not be auto-promoted to an org third party") + require.NotNil(t, reloadedTarget.ThirdPartyID, "target sharing a first-party origin must be promoted via its sibling") + assert.Equal(t, orgThirdParty.ID, *reloadedTarget.ThirdPartyID) } // TestProcess_ReenqueuesUnmappedSiblingOnResolve asserts that when a @@ -1295,8 +1513,7 @@ func TestProcess_ReenqueuesUnmappedSiblingOnResolve(t *testing.T) { return reloadedTarget.LoadByID(ctx, conn, fx.scope, target.ID) })) - require.NotNil(t, reloadedTarget.CommonTrackerPatternID, "target must resolve a catalog link via its sibling") - assert.Nil(t, reloadedTarget.ThirdPartyID, "target must not be auto-promoted to an org third party") + require.NotNil(t, reloadedTarget.ThirdPartyID, "target must resolve via its promoted sibling") var reloadedUnmapped coredata.TrackerPattern @@ -1450,8 +1667,7 @@ func TestProcess_DoesNotReenqueuePromotedOrExtensionSiblings(t *testing.T) { return p } - require.NotNil(t, reload(target.ID).CommonTrackerPatternID, "target must resolve a catalog link via its sibling") - assert.Nil(t, reload(target.ID).ThirdPartyID, "target must not be auto-promoted to an org third party") + require.NotNil(t, reload(target.ID).ThirdPartyID, "target must resolve via its promoted sibling") require.NotNil(t, reload(plainSibling.ID).MappingRequestedAt, "plain unmapped sibling must be re-enqueued") assert.Nil(t, reload(mappedSibling.ID).MappingRequestedAt, "promoted sibling must not be re-enqueued") assert.Nil(t, reload(extensionSibling.ID).MappingRequestedAt, "EXTENSION-sourced sibling must not be re-enqueued") diff --git a/pkg/coredata/migrations/20260611T120000Z.sql b/pkg/coredata/migrations/20260611T120000Z.sql deleted file mode 100644 index 924528923..000000000 --- a/pkg/coredata/migrations/20260611T120000Z.sql +++ /dev/null @@ -1,134 +0,0 @@ --- Copyright (c) 2026 Probo Inc . --- --- Permission to use, copy, modify, and/or distribute this software for any --- purpose with or without fee is hereby granted, provided that the above --- copyright notice and this permission notice appear in all copies. --- --- THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH --- REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY --- AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, --- INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM --- LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR --- OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR --- PERFORMANCE OF THIS SOFTWARE. - --- Enforce a single org ThirdParty per catalog vendor per organization. --- --- The tracker-mapping worker used to materialize org third parties from --- the catalog and could race the load-then-create check, leaving several --- rows that share the same (organization_id, common_third_party_id). --- Org third parties are now created only through the explicit import --- action; this migration both cleans up the historical duplicates and --- adds the partial unique index that makes the invariant enforceable. --- --- DESTRUCTIVE: the DO block below merges duplicate org third parties onto --- the earliest-created survivor by repointing every foreign key that --- references third_parties(id) and then deleting the extras. It is driven --- by catalog introspection (pg_constraint / pg_index) so it covers every --- referencing table without hard-coding the list, and it removes link --- rows that would collide on a referencing table's unique key before --- repointing. Validate it against a production dump before deploying. - -DO $$ -DECLARE - fk RECORD; - uq RECORD; - has_dupes boolean; -BEGIN - -- Map each duplicate row to the survivor it should be merged into: - -- the earliest-created row for the (organization_id, - -- common_third_party_id) pair, ties broken by id. - CREATE TEMP TABLE _tp_dedupe ON COMMIT DROP AS - SELECT t.id AS dup_id, k.survivor_id - FROM third_parties t - JOIN ( - SELECT - organization_id, - common_third_party_id, - (array_agg(id ORDER BY created_at ASC, id ASC))[1] AS survivor_id - FROM third_parties - WHERE common_third_party_id IS NOT NULL - GROUP BY organization_id, common_third_party_id - HAVING count(*) > 1 - ) k - ON t.organization_id = k.organization_id - AND t.common_third_party_id = k.common_third_party_id - WHERE t.id <> k.survivor_id; - - SELECT EXISTS (SELECT 1 FROM _tp_dedupe) INTO has_dupes; - - IF has_dupes THEN - -- For every single-column foreign key that references - -- third_parties(id) ... - FOR fk IN - SELECT c.conrelid::regclass::text AS tbl, - a.attname::text AS col - FROM pg_constraint c - JOIN pg_attribute a - ON a.attrelid = c.conrelid - AND a.attnum = c.conkey[1] - WHERE c.contype = 'f' - AND c.confrelid = 'third_parties'::regclass - AND array_length(c.conkey, 1) = 1 - LOOP - -- ... drop duplicate-side rows that would collide with a - -- survivor-side row on any unique key that includes the FK - -- column (the survivor's row wins; this realizes the union of - -- the two vendors' links). - FOR uq IN - SELECT ( - SELECT string_agg( - format('s.%I IS NOT DISTINCT FROM d.%I', - att.attname, att.attname), - ' AND ') - FROM pg_attribute att - WHERE att.attrelid = i.indrelid - AND att.attnum = ANY (i.indkey) - AND att.attname <> fk.col - ) AS match_pred - FROM pg_index i - WHERE i.indrelid = fk.tbl::regclass - AND (i.indisunique OR i.indisprimary) - AND EXISTS ( - SELECT 1 - FROM pg_attribute att - WHERE att.attrelid = i.indrelid - AND att.attnum = ANY (i.indkey) - AND att.attname = fk.col - ) - LOOP - IF uq.match_pred IS NOT NULL AND length(uq.match_pred) > 0 THEN - EXECUTE format( - 'DELETE FROM %1$s d - USING _tp_dedupe m - WHERE d.%2$I = m.dup_id - AND EXISTS ( - SELECT 1 FROM %1$s s - WHERE s.%2$I = m.survivor_id - AND %3$s - )', - fk.tbl, fk.col, uq.match_pred - ); - END IF; - END LOOP; - - -- Repoint the remaining references onto the survivor. - EXECUTE format( - 'UPDATE %1$s d - SET %2$I = m.survivor_id - FROM _tp_dedupe m - WHERE d.%2$I = m.dup_id', - fk.tbl, fk.col - ); - END LOOP; - - -- Drop the now-unreferenced duplicate rows. - DELETE FROM third_parties t - USING _tp_dedupe m - WHERE t.id = m.dup_id; - END IF; -END $$; - -CREATE UNIQUE INDEX third_parties_org_common_key - ON third_parties (organization_id, common_third_party_id) - WHERE common_third_party_id IS NOT NULL; diff --git a/pkg/probod/probod.go b/pkg/probod/probod.go index 2e50aafea..3ab0f2305 100644 --- a/pkg/probod/probod.go +++ b/pkg/probod/probod.go @@ -321,7 +321,7 @@ func (impl *Implm) Run( return err } - trackerMappingCfg, trackerEnrichmentCfg, err := impl.buildTrackerAgents(l, tp, r) + trackerMappingCfg, trackerEnrichmentCfg, thirdPartyDisambiguationCfg, err := impl.buildTrackerAgents(l, tp, r) if err != nil { return err } @@ -775,6 +775,7 @@ func (impl *Implm) Run( pgClient, l, trackerMappingCfg, + thirdPartyDisambiguationCfg, time.Duration(impl.cfg.TrackerMappingWorker.StaleAfter)*time.Second, worker.WithInterval(time.Duration(impl.cfg.TrackerMappingWorker.Interval)*time.Second), worker.WithMaxConcurrency(impl.cfg.TrackerMappingWorker.MaxConcurrency), diff --git a/pkg/probod/tracker_agents.go b/pkg/probod/tracker_agents.go index 154faa6ac..ce8b1ab63 100644 --- a/pkg/probod/tracker_agents.go +++ b/pkg/probod/tracker_agents.go @@ -22,25 +22,27 @@ import ( "go.gearno.de/kit/log" "go.opentelemetry.io/otel/trace" "go.probo.inc/probo/pkg/cookiebanner" + "go.probo.inc/probo/pkg/thirdparty" ) -// buildTrackerAgents wires the two tracker agents from the probod +// buildTrackerAgents wires the three tracker agents from the probod // config, each with its own LLM client and tuning: the tracker-mapping -// agent (catalog identification) and the common-pattern enrichment agent -// (description research). Both are opt-in: when -// `llm.tracker-mapping.provider` is empty it returns zero configs (nil -// LLM clients) so callers run without agent fallback. +// agent (catalog identification), the common-pattern enrichment agent +// (description research), and the third-party disambiguation agent. All +// are opt-in: when `llm.tracker-mapping.provider` is empty it returns +// zero configs (nil LLM clients) so callers run without agent fallback. // -// The enrichment agent falls back to the tracker-mapping config when its -// own provider slot is empty, so a deployment that configures only -// `tracker-mapping` keeps wiring both agents. +// The enrichment and disambiguation agents fall back to the +// tracker-mapping config when their own provider slot is empty, so a +// deployment that configures only `tracker-mapping` keeps wiring all +// three agents. func (impl *Implm) buildTrackerAgents( l *log.Logger, tp trace.TracerProvider, r prometheus.Registerer, -) (cookiebanner.TrackerMappingAgentConfig, cookiebanner.TrackerEnrichmentAgentConfig, error) { +) (cookiebanner.TrackerMappingAgentConfig, cookiebanner.TrackerEnrichmentAgentConfig, thirdparty.DisambiguationAgentConfig, error) { if impl.cfg.Agents.TrackerMapping.Provider == "" { - return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, nil + return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, thirdparty.DisambiguationAgentConfig{}, nil } firecrawlAPIKey := impl.cfg.Agents.Tools.FirecrawlAPIKey @@ -53,7 +55,7 @@ func (impl *Implm) buildTrackerAgents( r, ) if err != nil { - return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, fmt.Errorf("cannot resolve tracker mapping agent client: %w", err) + return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, thirdparty.DisambiguationAgentConfig{}, fmt.Errorf("cannot resolve tracker mapping agent client: %w", err) } mappingCfg := cookiebanner.TrackerMappingAgentConfig{ @@ -79,7 +81,7 @@ func (impl *Implm) buildTrackerAgents( r, ) if err != nil { - return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, fmt.Errorf("cannot resolve tracker enrichment agent client: %w", err) + return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, thirdparty.DisambiguationAgentConfig{}, fmt.Errorf("cannot resolve tracker enrichment agent client: %w", err) } enrichmentCfg := cookiebanner.TrackerEnrichmentAgentConfig{ @@ -92,5 +94,29 @@ func (impl *Implm) buildTrackerAgents( MaxTurns: impl.cfg.CommonPatternEnrichmentWorker.AgentMaxTurns, } - return mappingCfg, enrichmentCfg, nil + disambiguationSlot := impl.cfg.Agents.ThirdPartyDisambiguation + if disambiguationSlot.Provider == "" { + disambiguationSlot = impl.cfg.Agents.TrackerMapping + } + + disambiguationAgentCfg, disambiguationClient, err := impl.resolveAgentClient( + "third-party-disambiguation", + disambiguationSlot, + l, + tp, + r, + ) + if err != nil { + return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, thirdparty.DisambiguationAgentConfig{}, fmt.Errorf("cannot resolve third party disambiguation agent client: %w", err) + } + + disambiguationCfg := thirdparty.DisambiguationAgentConfig{ + LLMClient: disambiguationClient, + Model: disambiguationAgentCfg.ModelName, + MaxTokens: disambiguationAgentCfg.MaxTokens, + Temperature: disambiguationAgentCfg.Temperature, + Timeout: time.Duration(impl.cfg.TrackerMappingWorker.DisambiguationAgentTimeout) * time.Second, + } + + return mappingCfg, enrichmentCfg, disambiguationCfg, nil } diff --git a/pkg/probodconfig/llm_config.go b/pkg/probodconfig/llm_config.go index 63a447a74..589c042a6 100644 --- a/pkg/probodconfig/llm_config.go +++ b/pkg/probodconfig/llm_config.go @@ -53,12 +53,16 @@ type ( // tracker-mapping background worker. LLM parameters for the mapping // agent it runs live under AgentsConfig.TrackerMapping. AgentTimeout // and AgentMaxTurns bound a single mapping agent run. + // DisambiguationAgentTimeout caps a single third-party + // disambiguation agent run; that agent runs inside this worker but + // uses its own LLM parameters from AgentsConfig.ThirdPartyDisambiguation. TrackerMappingWorkerConfig struct { - Interval int `json:"interval"` // seconds between polls - MaxConcurrency int `json:"max-concurrency"` - StaleAfter int `json:"stale-after"` // seconds before a claim is recycled - AgentTimeout int `json:"agent-timeout"` // seconds, single agent run - AgentMaxTurns int `json:"agent-max-turns"` + Interval int `json:"interval"` // seconds between polls + MaxConcurrency int `json:"max-concurrency"` + StaleAfter int `json:"stale-after"` // seconds before a claim is recycled + AgentTimeout int `json:"agent-timeout"` // seconds, single agent run + AgentMaxTurns int `json:"agent-max-turns"` + DisambiguationAgentTimeout int `json:"disambiguation-agent-timeout"` // seconds, single disambiguation run } // CommonPatternEnrichmentWorkerConfig holds worker-side tuning for @@ -82,14 +86,15 @@ type ( // settings. Default is used as a fallback when an agent-specific field // is zero-valued. AgentsConfig struct { - Providers map[string]LLMProviderConfig `json:"providers"` - Default LLMAgentConfig `json:"defaults"` - Probo LLMAgentConfig `json:"probo"` - EvidenceDescriber LLMAgentConfig `json:"evidence-describer"` - ThirdPartyVetter LLMAgentConfig `json:"third-party-vetter"` - TrackerMapping LLMAgentConfig `json:"tracker-mapping"` - TrackerEnrichment LLMAgentConfig `json:"tracker-enrichment"` - Tools AgentToolsConfig `json:"tools"` + Providers map[string]LLMProviderConfig `json:"providers"` + Default LLMAgentConfig `json:"defaults"` + Probo LLMAgentConfig `json:"probo"` + EvidenceDescriber LLMAgentConfig `json:"evidence-describer"` + ThirdPartyVetter LLMAgentConfig `json:"third-party-vetter"` + ThirdPartyDisambiguation LLMAgentConfig `json:"third-party-disambiguation"` + TrackerMapping LLMAgentConfig `json:"tracker-mapping"` + TrackerEnrichment LLMAgentConfig `json:"tracker-enrichment"` + Tools AgentToolsConfig `json:"tools"` } ) diff --git a/pkg/thirdparty/common_match.go b/pkg/thirdparty/common_match.go new file mode 100644 index 000000000..e48ba6977 --- /dev/null +++ b/pkg/thirdparty/common_match.go @@ -0,0 +1,217 @@ +// 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 thirdparty + +import ( + "context" + "fmt" + "sort" + "strings" + + "go.gearno.de/kit/pg" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/gid" + "go.probo.inc/probo/pkg/slug" + "go.probo.inc/probo/pkg/uri" +) + +// Heuristic thresholds for matching a CommonThirdParty to an existing +// org ThirdParty. Exported so callers can short-circuit explicitly +// (skip the agent, leave the pattern unlinked, etc.) instead of +// duplicating magic numbers. +const ( + // HighConfidenceScore is the floor at which a heuristic match + // is treated as obvious (exact name, suffix-stripped name, slug + // equality). Callers typically link without consulting the + // agent at this score. + HighConfidenceScore = 0.85 + + // MinAgentScore is the floor below which a candidate is + // statistical noise and should not be shown to the + // disambiguation agent. Below this, callers typically prefer + // to create a fresh row over asking the model to pick among + // weak candidates. + MinAgentScore = 0.6 + + // MaxAgentCandidates caps the candidate list shown to the + // disambiguation agent. The list is heuristic-ranked, so the + // top few are the only ones worth the agent's tokens. + MaxAgentCandidates = 5 +) + +// ScoredCandidate is a scored heuristic-match candidate. It is the +// unified currency between the heuristic ranker (RankCandidates) and +// the disambiguation agent (Disambiguate): the agent renders the +// `ThirdParty` fields plus the score directly into its prompt, with +// no intermediate DTO. +type ScoredCandidate struct { + ThirdParty *coredata.ThirdParty + Score float64 +} + +// corporateSuffixes are the legal-form noise words stripped when +// comparing third-party names heuristically. The list is intentionally +// short and conservative: matching "Foo Inc" to "Foo" is safe, but +// stripping "Group" or "Services" would over-match unrelated entries. +// +// Order matters: stripCorporateSuffixes returns on the first match, +// so longer / comma-prefixed forms must come before their shorter +// siblings (", inc." before " inc.", which itself comes before " inc"). +var corporateSuffixes = []string{ + " incorporated", + " corporation", + ", inc.", + ", inc", + " l.l.c.", + " s.a.s.", + " inc.", + " inc", + " llc", + " ltd.", + " ltd", + " limited", + " gmbh", + " s.a.", + " sas", + " sa", + " ag", + " plc", + " corp.", + " corp", + " co.", + " co", + " b.v.", + " bv", +} + +// RankCandidates ranks org ThirdParty rows by how likely each is to +// represent the given CommonThirdParty. Returned slice is sorted by +// descending score; only candidates with score > 0 are kept. Pure +// function: no I/O, deterministic on its inputs. +// +// Scoring (highest match wins; website-host overlap can lift a name +// miss to 0.8): +// +// - exact lowercase name = 1.0 +// - lowercase name with corporate suffix stripped, equal = 0.9 +// - slug equality (slug.Make on the org's name) = 0.85 +// - website host (eTLD+1) overlap with the catalog domain set = 0.8 +func RankCandidates( + commonParty coredata.CommonThirdParty, + commonDomains coredata.CommonThirdPartyDomains, + candidates coredata.ThirdParties, +) []ScoredCandidate { + commonName := strings.ToLower(strings.TrimSpace(commonParty.Name)) + commonStripped := stripCorporateSuffixes(commonName) + commonSlug := commonParty.Slug + + commonHost := "" + if commonParty.WebsiteURL != nil { + commonHost = uri.ExtractDomain(*commonParty.WebsiteURL) + } + + commonDomainSet := make(map[string]struct{}, len(commonDomains)) + for _, d := range commonDomains { + commonDomainSet[strings.ToLower(d.Domain)] = struct{}{} + } + + if commonHost != "" { + commonDomainSet[commonHost] = struct{}{} + } + + scored := make([]ScoredCandidate, 0, len(candidates)) + + for _, tp := range candidates { + score := 0.0 + + orgName := strings.ToLower(strings.TrimSpace(tp.Name)) + orgStripped := stripCorporateSuffixes(orgName) + + switch { + case orgName != "" && orgName == commonName: + score = 1.0 + case orgStripped != "" && orgStripped == commonStripped: + score = 0.9 + case commonSlug != "" && slug.Make(tp.Name) == commonSlug: + score = 0.85 + } + + if tp.WebsiteURL != nil { + orgHost := uri.ExtractDomain(*tp.WebsiteURL) + if orgHost != "" { + if _, hit := commonDomainSet[orgHost]; hit { + if score < 0.8 { + score = 0.8 + } + } + } + } + + if score == 0 { + continue + } + + scored = append(scored, ScoredCandidate{ + ThirdParty: tp, + Score: score, + }) + } + + sort.SliceStable(scored, func(i, j int) bool { + return scored[i].Score > scored[j].Score + }) + + return scored +} + +// stripCorporateSuffixes removes a single trailing legal-form suffix +// from a lowercased name. Only one suffix is stripped to avoid +// mangling names that happen to end in two stop-words (e.g. "Foo Inc +// LLC" → "Foo Inc", not "Foo"). +func stripCorporateSuffixes(lowerName string) string { + for _, s := range corporateSuffixes { + if before, ok := strings.CutSuffix(lowerName, s); ok { + return strings.TrimSpace(before) + } + } + + return lowerName +} + +// LinkToCommon writes common_third_party_id onto an org ThirdParty so +// future matches against the same CommonThirdParty can short-circuit +// to the exact-link path in O(1). No-op when the field is already set +// (to any value) — we never overwrite an existing catalog link because +// a heuristic or agent false-positive must not corrupt a previous, +// possibly more accurate, association. +func LinkToCommon( + ctx context.Context, + tx pg.Tx, + scope coredata.Scoper, + orgThirdParty *coredata.ThirdParty, + commonID gid.GID, +) error { + if orgThirdParty.CommonThirdPartyID != nil { + return nil + } + + orgThirdParty.CommonThirdPartyID = &commonID + + if err := orgThirdParty.Update(ctx, tx, scope); err != nil { + return fmt.Errorf("cannot update third party with common id: %w", err) + } + + return nil +} diff --git a/pkg/thirdparty/common_match_test.go b/pkg/thirdparty/common_match_test.go new file mode 100644 index 000000000..53bfd7d75 --- /dev/null +++ b/pkg/thirdparty/common_match_test.go @@ -0,0 +1,178 @@ +// 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 thirdparty + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/gid" +) + +func TestStripCorporateSuffixes(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + in string + want string + }{ + {name: "llc suffix", in: "google llc", want: "google"}, + {name: "comma inc", in: "stripe, inc", want: "stripe"}, + {name: "inc dot", in: "meta inc.", want: "meta"}, + {name: "ltd", in: "deepmind ltd", want: "deepmind"}, + {name: "gmbh", in: "n8n gmbh", want: "n8n"}, + {name: "no suffix", in: "cloudflare", want: "cloudflare"}, + {name: "trailing space", in: "github inc", want: "github"}, + {name: "only one suffix stripped", in: "foo inc llc", want: "foo inc"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + assert.Equal(t, tt.want, stripCorporateSuffixes(tt.in)) + }) + } +} + +func TestRankCandidates(t *testing.T) { + t.Parallel() + + tenantID := gid.NewTenantID() + + mkTP := func(name, website string) *coredata.ThirdParty { + tp := &coredata.ThirdParty{ + ID: gid.New(tenantID, coredata.ThirdPartyEntityType), + Name: name, + } + + if website != "" { + tp.WebsiteURL = new(website) + } + + return tp + } + + t.Run("exact name match scores 1.0", func(t *testing.T) { + t.Parallel() + + common := coredata.CommonThirdParty{Name: "Google", Slug: "google"} + got := RankCandidates(common, nil, coredata.ThirdParties{ + mkTP("Google", ""), + mkTP("Stripe", ""), + }) + + require.Len(t, got, 1) + assert.Equal(t, 1.0, got[0].Score) + assert.Equal(t, "Google", got[0].ThirdParty.Name) + }) + + t.Run("suffix-stripped name scores 0.9", func(t *testing.T) { + t.Parallel() + + common := coredata.CommonThirdParty{Name: "Google", Slug: "google"} + got := RankCandidates(common, nil, coredata.ThirdParties{ + mkTP("Google LLC", ""), + }) + + require.Len(t, got, 1) + assert.Equal(t, 0.9, got[0].Score) + }) + + t.Run("slug equality scores 0.85", func(t *testing.T) { + t.Parallel() + + common := coredata.CommonThirdParty{Name: "Google", Slug: "google"} + got := RankCandidates(common, nil, coredata.ThirdParties{ + mkTP("google!", ""), + }) + + require.Len(t, got, 1) + assert.Equal(t, 0.85, got[0].Score) + }) + + t.Run("website host overlap scores 0.8 when name does not match", func(t *testing.T) { + t.Parallel() + + common := coredata.CommonThirdParty{ + Name: "Google Analytics", + Slug: "google-analytics", + WebsiteURL: new("https://google.com"), + } + + got := RankCandidates(common, nil, coredata.ThirdParties{ + mkTP("Sundar's Search Co", "https://www.google.com/about"), + }) + + require.Len(t, got, 1) + assert.Equal(t, 0.8, got[0].Score) + }) + + t.Run("domain set overlap scores 0.8", func(t *testing.T) { + t.Parallel() + + common := coredata.CommonThirdParty{Name: "Stripe", Slug: "stripe"} + domains := coredata.CommonThirdPartyDomains{ + {Domain: "stripe.com"}, + {Domain: "stripe.network"}, + } + + got := RankCandidates(common, domains, coredata.ThirdParties{ + mkTP("Payment Processor", "https://api.stripe.com/v1"), + }) + + require.Len(t, got, 1) + assert.Equal(t, 0.8, got[0].Score) + }) + + t.Run("no match returns empty", func(t *testing.T) { + t.Parallel() + + common := coredata.CommonThirdParty{Name: "Stripe", Slug: "stripe"} + got := RankCandidates(common, nil, coredata.ThirdParties{ + mkTP("Acme", "https://acme.example"), + mkTP("Widgets Inc", "https://widgets.example"), + }) + + assert.Empty(t, got) + }) + + t.Run("ranks descending by score", func(t *testing.T) { + t.Parallel() + + common := coredata.CommonThirdParty{ + Name: "Google", + Slug: "google", + WebsiteURL: new("https://google.com"), + } + + got := RankCandidates(common, nil, coredata.ThirdParties{ + mkTP("Random", "https://google.com"), + mkTP("Google", ""), + mkTP("Google LLC", ""), + }) + + require.Len(t, got, 3) + assert.Equal(t, "Google", got[0].ThirdParty.Name) + assert.Equal(t, 1.0, got[0].Score) + assert.Equal(t, "Google LLC", got[1].ThirdParty.Name) + assert.Equal(t, 0.9, got[1].Score) + assert.Equal(t, "Random", got[2].ThirdParty.Name) + assert.Equal(t, 0.8, got[2].Score) + }) +} diff --git a/pkg/thirdparty/disambiguation_agent.go b/pkg/thirdparty/disambiguation_agent.go new file mode 100644 index 000000000..e7487c8f7 --- /dev/null +++ b/pkg/thirdparty/disambiguation_agent.go @@ -0,0 +1,241 @@ +// 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 thirdparty + +import ( + "context" + _ "embed" + "fmt" + "strings" + "time" + + "go.gearno.de/kit/log" + "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:embed prompts/disambiguation.txt.tmpl +var disambiguationPrompt string + +const ( + // disambiguationConfidenceThreshold is the floor below which we + // treat the agent's pick as "no confident match" even when it + // returned a non-nil matched_id. Mirrors the conservative bias + // described in the prompt. + disambiguationConfidenceThreshold = 0.6 + + // defaultDisambiguationTimeout caps a single disambiguation run + // when the config supplies none. The agent has no tools and a + // single turn, so this is mostly a guard against a hung LLM + // provider, not a real budget. + defaultDisambiguationTimeout = 45 * time.Second + + // defaultDisambiguationMaxTokens caps the agent's output when the + // config carries no max-tokens budget. The final output is tiny (a + // single id plus a one-sentence rationale), but the budget must + // leave ample headroom for reasoning models (e.g. the GPT-5 + // family): their reasoning tokens count against max_tokens, so too + // small a budget gets consumed by reasoning and truncates the JSON, + // surfacing as "unexpected end of JSON input". + defaultDisambiguationMaxTokens = 4096 +) + +// DisambiguationAgentConfig configures the third-party disambiguation +// agent. The agent has no DB tools and no web-search tools: the +// candidate list is supplied entirely in the prompt and the agent +// only picks among it. +// +// MaxTokens and Temperature bound and steer the single LLM call, and +// Timeout caps a single run. Zero-valued fields fall back to package +// defaults. +type DisambiguationAgentConfig struct { + LLMClient *llm.Client + Model string + MaxTokens *int + Temperature *float64 + Timeout time.Duration +} + +// DisambiguationResult is the structured output the disambiguation +// agent returns when picking the best existing org ThirdParty for a +// catalog entry. +type DisambiguationResult struct { + MatchedID *string `json:"matched_id" jsonschema:"GID of the org third party that best matches, or null if none of the candidates is a confident match."` + Confidence float64 `json:"confidence" jsonschema:"Confidence level from 0.0 to 1.0. Below 0.6 means 'no confident match' and matched_id MUST be null."` + Reasoning string `json:"reasoning" jsonschema:"One short sentence describing the rationale."` +} + +// BuildDisambiguationAgent wires the agent that picks the best +// existing org ThirdParty for a catalog entry. It deliberately has +// no tools: the candidate list is supplied in the prompt and the +// agent must only choose among it. +func BuildDisambiguationAgent( + cfg DisambiguationAgentConfig, + logger *log.Logger, +) *agent.Agent { + outputType, err := agent.NewOutputType[DisambiguationResult]("third_party_disambiguation") + if err != nil { + panic(fmt.Sprintf("thirdparty: cannot build disambiguation output type: %s", err)) + } + + maxTokens := defaultDisambiguationMaxTokens + if cfg.MaxTokens != nil && *cfg.MaxTokens > 0 { + maxTokens = *cfg.MaxTokens + } + + opts := []agent.Option{ + agent.WithInstructions(disambiguationPrompt), + agent.WithModel(cfg.Model), + agent.WithOutputType(outputType), + agent.WithMaxTurns(1), + agent.WithMaxTokens(maxTokens), + agent.WithLogger(logger), + } + + if cfg.Temperature != nil { + opts = append(opts, agent.WithTemperature(*cfg.Temperature)) + } + + return agent.New("third-party-disambiguation", cfg.LLMClient, opts...) +} + +// Disambiguate runs the agent against the given catalog third party +// and candidate list, and returns the matched candidate's ID — or +// nil when the agent picks "none", returns a confidence below the +// threshold, or fails. Errors from the agent itself are returned; +// "no confident match" is not an error. +// +// The matched candidate is identified by string equality against the +// IDs supplied in `candidates`; we never invent IDs from the agent's +// output, so a model that hallucinates an ID is treated as "none". +func Disambiguate( + ctx context.Context, + a *agent.Agent, + logger *log.Logger, + commonParty coredata.CommonThirdParty, + commonDomains coredata.CommonThirdPartyDomains, + candidates []ScoredCandidate, + timeout time.Duration, +) (*gid.GID, error) { + if a == nil || len(candidates) == 0 { + return nil, nil + } + + if timeout <= 0 { + timeout = defaultDisambiguationTimeout + } + + prompt := buildDisambiguationPrompt(commonParty, commonDomains, candidates) + + agentCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + result, err := agent.RunTyped[DisambiguationResult]( + agentCtx, + a, + []llm.Message{ + { + Role: llm.RoleUser, + Parts: []llm.Part{llm.TextPart{Text: prompt}}, + }, + }, + ) + if err != nil { + return nil, fmt.Errorf("cannot run disambiguation agent: %w", err) + } + + out := result.Output + + if out.MatchedID == nil || *out.MatchedID == "" { + return nil, nil + } + + if out.Confidence < disambiguationConfidenceThreshold { + logger.InfoCtx( + ctx, + "disambiguation agent below confidence threshold", + log.String("matched_id", *out.MatchedID), + log.Float64("confidence", out.Confidence), + ) + + return nil, nil + } + + for _, c := range candidates { + if c.ThirdParty.ID.String() == *out.MatchedID { + id := c.ThirdParty.ID + + return &id, nil + } + } + + logger.WarnCtx( + ctx, + "disambiguation agent returned id not in candidate list", + log.String("matched_id", *out.MatchedID), + ) + + return nil, nil +} + +// buildDisambiguationPrompt formats the catalog third party and the +// heuristic-ranked candidate list into the user message for the +// disambiguation agent. The prompt is intentionally compact: the +// agent only needs ids, names, websites, and the heuristic score to +// decide. +func buildDisambiguationPrompt( + commonParty coredata.CommonThirdParty, + commonDomains coredata.CommonThirdPartyDomains, + candidates []ScoredCandidate, +) string { + var b strings.Builder + + b.WriteString("Catalog third party:\n") + fmt.Fprintf(&b, " name: %s\n", commonParty.Name) + + if commonParty.WebsiteURL != nil && *commonParty.WebsiteURL != "" { + fmt.Fprintf(&b, " website: %s\n", *commonParty.WebsiteURL) + } + + if len(commonDomains) > 0 { + domains := make([]string, len(commonDomains)) + for i, d := range commonDomains { + domains[i] = d.Domain + } + + fmt.Fprintf(&b, " domains: %s\n", strings.Join(domains, ", ")) + } + + b.WriteString("\nCandidate organisation third parties (heuristic-ranked):\n") + + for i, c := range candidates { + fmt.Fprintf(&b, "- id: %s\n", c.ThirdParty.ID.String()) + fmt.Fprintf(&b, " name: %s\n", c.ThirdParty.Name) + + if c.ThirdParty.WebsiteURL != nil && *c.ThirdParty.WebsiteURL != "" { + fmt.Fprintf(&b, " website: %s\n", *c.ThirdParty.WebsiteURL) + } + + fmt.Fprintf(&b, " heuristic_score: %.2f\n", c.Score) + + if i < len(candidates)-1 { + b.WriteString("\n") + } + } + + return b.String() +}