From a99a4dde14f786cea95bfca6bb31e3b17911f0f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=89mile=20R=C3=A9?= Date: Thu, 28 May 2026 17:20:36 +0200 Subject: [PATCH] Promote tracker patterns to org third parties via worker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Manual moves of a non-extension TrackerPattern lacking a ThirdPartyID now request mapping, which the tracker-mapping worker resolves with a four-stage pipeline: exact common_third_party_id link, heuristic ranking, agent disambiguation, and finally CreateFromCommon. Existing fuzzy-matched org rows are tagged with common_third_party_id so the next promotion takes the O(1) exact-link path. The matching primitives live in pkg/thirdparty (RankCandidates, LinkToCommon, CreateFromCommon, ScoredCandidate, threshold constants) so the disambiguation agent and the heuristic share one candidate type. Cookiebanner orchestrates them; cookie-banner-specific concerns (pattern -> common-pattern -> common-party navigation, the EXTENSION gate, and structured logs) stay in the worker. Signed-off-by: Émile Ré --- contrib/claude/file-naming.md | 45 +- pkg/cookiebanner/service.go | 14 + pkg/cookiebanner/tracker_mapping_agent.go | 3 + pkg/cookiebanner/tracker_mapping_worker.go | 300 ++++++++--- .../tracker_mapping_worker_test.go | 495 ++++++++++++++++++ pkg/coredata/third_party.go | 2 + pkg/probod/probod.go | 4 +- pkg/probod/tracker_mapping.go | 33 +- pkg/thirdparty/disambiguation_agent.go | 209 ++++++++ pkg/thirdparty/match.go | 274 ++++++++++ pkg/thirdparty/match_test.go | 181 +++++++ .../prompts/disambiguation.txt.tmpl | 25 + 12 files changed, 1491 insertions(+), 94 deletions(-) create mode 100644 pkg/cookiebanner/tracker_mapping_worker_test.go create mode 100644 pkg/thirdparty/disambiguation_agent.go create mode 100644 pkg/thirdparty/match.go create mode 100644 pkg/thirdparty/match_test.go create mode 100644 pkg/thirdparty/prompts/disambiguation.txt.tmpl diff --git a/contrib/claude/file-naming.md b/contrib/claude/file-naming.md index 2e9ecb9f0..1dadcf969 100644 --- a/contrib/claude/file-naming.md +++ b/contrib/claude/file-naming.md @@ -42,17 +42,48 @@ pkg/cookiebanner/pattern_analysis_worker.go ## Agent files -When a package has a worker that uses an agent, the agent construction logic -goes in `_agent.go` alongside `_worker.go`. The -worker file stays focused on `Claim`/`Process` and handler methods; the agent -file owns agent construction, prompt building, constants, config, and the -`//go:embed` directive for prompt templates. +A file whose **sole purpose** is to construct and operate an agent uses the +`_agent.go` suffix. The agent file owns agent construction, prompt +building, the `//go:embed` directive for prompt templates, the typed result, +the agent-specific config, and any agent-only constants (timeout, confidence +threshold). Callers (workers, services) hold a `*agent.Agent` field and import +the file's `Build…Agent` constructor. + +This applies in two shapes: + +1. **Paired with a worker** (most common). The worker file + `_worker.go` stays focused on `Claim`/`Process`, and the + agent it uses lives in `_agent.go` next to it. + + ``` + pkg/cookiebanner/tracker_mapping_worker.go -- worker handler + pkg/cookiebanner/tracker_mapping_agent.go -- agent construction + prompts + ``` + +2. **Standalone, called from elsewhere.** When the agent is consumed by a + different package (or by multiple packages — e.g. an agent that operates on + a domain entity, used by several feature workers), it lives in the package + that owns the domain, named `_agent.go`. + + ``` + pkg/thirdparty/disambiguation_agent.go -- catalog→org ThirdParty matcher + pkg/vetting/sub_agent.go -- generic vetting sub-agent + ``` + +A file is NOT renamed to `_agent.go` when the agent is incidental to a service +that does substantially more than agent orchestration (e.g. CRUD, caching, +auth). In that case the file keeps its service name and the agent is built +inline: ``` -pkg/cookiebanner/tracker_mapping_worker.go -- worker handler -pkg/cookiebanner/tracker_mapping_agent.go -- agent construction + prompts +pkg/evidencedescriber/evidencedescriber.go -- single-file describer service +pkg/vetting/assessment.go -- third-party assessment service ``` +If the agent construction grows past a few dozen lines or sprouts its own +prompt embed / typed result / config struct, extract it into a sibling +`_agent.go`. + ## Tool files Each agent tool lives in its own `_tool.go` file, named after the diff --git a/pkg/cookiebanner/service.go b/pkg/cookiebanner/service.go index 680851005..20a79428a 100644 --- a/pkg/cookiebanner/service.go +++ b/pkg/cookiebanner/service.go @@ -2628,6 +2628,20 @@ func (s *Service) MoveTrackerPatternToCategory( return fmt.Errorf("cannot update tracker pattern: %w", err) } + // A manual move is the user's signal that this is a + // real tracker. Enqueue the tracker-mapping worker so + // it can promote the pattern to an org ThirdParty (or + // link an existing one) — never EXTENSION-sourced + // patterns, and never patterns we already promoted. + // SetMappingRequested is idempotent: it short-circuits + // when mapping_requested_at is already non-NULL. + if pattern.ThirdPartyID == nil && + (pattern.Source == nil || *pattern.Source != coredata.CookieSourceExtension) { + if err := pattern.SetMappingRequested(ctx, tx); err != nil { + return fmt.Errorf("cannot enqueue tracker mapping after move: %w", err) + } + } + var banner coredata.CookieBanner if err := banner.LoadByID(ctx, tx, scope, pattern.CookieBannerID); err != nil { return fmt.Errorf("cannot load cookie banner: %w", err) diff --git a/pkg/cookiebanner/tracker_mapping_agent.go b/pkg/cookiebanner/tracker_mapping_agent.go index 133f95526..13e02fe39 100644 --- a/pkg/cookiebanner/tracker_mapping_agent.go +++ b/pkg/cookiebanner/tracker_mapping_agent.go @@ -48,6 +48,9 @@ type TrackerMappingAgentResult struct { 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 diff --git a/pkg/cookiebanner/tracker_mapping_worker.go b/pkg/cookiebanner/tracker_mapping_worker.go index 053e9055b..0daf5f342 100644 --- a/pkg/cookiebanner/tracker_mapping_worker.go +++ b/pkg/cookiebanner/tracker_mapping_worker.go @@ -28,18 +28,21 @@ import ( "go.probo.inc/probo/pkg/gid" "go.probo.inc/probo/pkg/llm" "go.probo.inc/probo/pkg/slug" + "go.probo.inc/probo/pkg/thirdparty" ) type trackerMappingHandler struct { - pg *pg.Client - logger *log.Logger - agent *agent.Agent + pg *pg.Client + logger *log.Logger + mappingAgent *agent.Agent + disambiguationAgent *agent.Agent } func NewTrackerMappingWorker( pgClient *pg.Client, logger *log.Logger, - cfg TrackerMappingConfig, + mappingCfg TrackerMappingConfig, + disambiguationCfg thirdparty.DisambiguationConfig, opts ...worker.Option, ) *worker.Worker[coredata.TrackerPattern] { h := &trackerMappingHandler{ @@ -47,8 +50,12 @@ func NewTrackerMappingWorker( logger: logger, } - if cfg.LLMClient != nil { - h.agent = buildTrackerMappingAgent(cfg, pgClient, logger) + if mappingCfg.LLMClient != nil { + h.mappingAgent = buildTrackerMappingAgent(mappingCfg, pgClient, logger) + } + + if disambiguationCfg.LLMClient != nil { + h.disambiguationAgent = thirdparty.BuildDisambiguationAgent(disambiguationCfg, logger) } return worker.New( @@ -82,40 +89,62 @@ func (h *trackerMappingHandler) Claim(ctx context.Context) (coredata.TrackerPatt return tp, nil } +// Process resolves the catalog mapping (when missing) and then promotes +// the pattern to an org ThirdParty (when eligible). When a pattern is +// re-triggered by a manual move (it already carries a +// common_tracker_pattern_id), we MUST NOT re-resolve the catalog: the +// existing link is preserved and we jump straight to third-party +// promotion. func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.TrackerPattern) error { return h.pg.WithTx( ctx, func(ctx context.Context, tx pg.Tx) error { var ( commonPatternID *gid.GID - thirdPartyID *gid.GID err error ) - commonPatternID, thirdPartyID, err = h.matchByPattern(ctx, tx, tp) - if err != nil { - return fmt.Errorf("cannot match by pattern: %w", err) - } - - if commonPatternID == nil { - commonPatternID, thirdPartyID, err = h.matchByDomain(ctx, tx, tp) + if tp.CommonTrackerPatternID != nil { + commonPatternID = tp.CommonTrackerPatternID + } else { + commonPatternID, err = h.matchByPattern(ctx, tx, tp) if err != nil { - return fmt.Errorf("cannot match by domain: %w", err) + return fmt.Errorf("cannot match by pattern: %w", err) + } + + if commonPatternID == nil { + commonPatternID, err = h.matchByDomain(ctx, tx, tp) + if err != nil { + return fmt.Errorf("cannot match by domain: %w", err) + } + } + + if commonPatternID == nil && h.mappingAgent != nil { + commonPatternID, err = h.identifyWithAgent(ctx, tx, tp) + if err != nil { + return fmt.Errorf("cannot identify with agent: %w", err) + } + } + + if commonPatternID == nil { + commonPatternID, err = h.createUnmatchedPattern(ctx, tx, tp) + if err != nil { + return fmt.Errorf("cannot create unmatched pattern: %w", err) + } } } - if commonPatternID == nil && h.agent != nil { - commonPatternID, thirdPartyID, err = h.identifyWithAgent(ctx, tx, tp) - if err != nil { - return fmt.Errorf("cannot identify with agent: %w", err) - } - } + thirdPartyID := tp.ThirdPartyID - if commonPatternID == nil { - commonPatternID, err = h.createUnmatchedPattern(ctx, tx, tp) + if thirdPartyID == nil && + commonPatternID != nil && + (tp.Source == nil || *tp.Source != coredata.CookieSourceExtension) { + promoted, err := h.promoteThirdParty(ctx, tx, tp, *commonPatternID) if err != nil { - return fmt.Errorf("cannot create unmatched pattern: %w", err) + return fmt.Errorf("cannot promote third party: %w", err) } + + thirdPartyID = promoted } if commonPatternID != nil || thirdPartyID != nil { @@ -136,59 +165,55 @@ func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.Tracker ) } +// matchByPattern looks for a catalog row with the same pattern. It now +// only returns the catalog ID; third-party resolution happens later in +// promoteThirdParty. func (h *trackerMappingHandler) matchByPattern( ctx context.Context, conn pg.Querier, tp coredata.TrackerPattern, -) (*gid.GID, *gid.GID, error) { +) (*gid.GID, error) { var commonPattern coredata.CommonTrackerPattern if err := commonPattern.LoadByPattern(ctx, conn, tp.TrackerType, tp.Pattern, tp.MaxAgeSeconds); err != nil { if errors.Is(err, coredata.ErrResourceNotFound) { - return nil, nil, nil + return nil, nil } - return nil, nil, fmt.Errorf("cannot load common tracker pattern: %w", err) + return nil, fmt.Errorf("cannot load common tracker pattern: %w", err) } - var thirdPartyID *gid.GID - - if commonPattern.CommonThirdPartyID != nil { - var err error - - thirdPartyID, err = h.resolveThirdParty(ctx, conn, tp, &commonPattern) - if err != nil { - return nil, nil, fmt.Errorf("cannot resolve third party from pattern match: %w", err) - } - } - - return &commonPattern.ID, thirdPartyID, nil + return &commonPattern.ID, nil } +// matchByDomain finds a CommonThirdParty whose registered domains +// overlap the pattern's observed initiator domains, and upserts a +// CommonTrackerPattern linking the two. As with matchByPattern, +// third-party resolution is deferred to promoteThirdParty. func (h *trackerMappingHandler) matchByDomain( ctx context.Context, tx pg.Tx, tp coredata.TrackerPattern, -) (*gid.GID, *gid.GID, error) { +) (*gid.GID, error) { var trackers coredata.DetectedTrackers domains, err := trackers.LoadInitiatorDomainsByTrackerPatternID(ctx, tx, tp.ID, 10) if err != nil { - return nil, nil, fmt.Errorf("cannot load initiator domains: %w", err) + return nil, fmt.Errorf("cannot load initiator domains: %w", err) } if len(domains) == 0 { - return nil, nil, nil + return nil, nil } filter := coredata.NewCommonThirdPartyDomainFilter(domains) var matchedDomains coredata.CommonThirdPartyDomains if err := matchedDomains.Load(ctx, tx, 1, filter); err != nil { - return nil, nil, fmt.Errorf("cannot load common third party domain by domain match: %w", err) + return nil, fmt.Errorf("cannot load common third party domain by domain match: %w", err) } if len(matchedDomains) == 0 { - return nil, nil, nil + return nil, nil } commonThirdPartyID := matchedDomains[0].CommonThirdPartyID @@ -208,22 +233,17 @@ func (h *trackerMappingHandler) matchByDomain( } if _, err := commonPattern.Upsert(ctx, tx); err != nil { - return nil, nil, fmt.Errorf("cannot upsert common tracker pattern from domain match: %w", err) + return nil, fmt.Errorf("cannot upsert common tracker pattern from domain match: %w", err) } - thirdPartyID, err := h.resolveThirdParty(ctx, tx, tp, &commonPattern) - if err != nil { - return nil, nil, fmt.Errorf("cannot resolve third party from domain match: %w", err) - } - - return &commonPattern.ID, thirdPartyID, nil + return &commonPattern.ID, nil } func (h *trackerMappingHandler) identifyWithAgent( ctx context.Context, tx pg.Tx, tp coredata.TrackerPattern, -) (*gid.GID, *gid.GID, error) { +) (*gid.GID, error) { var trackers coredata.DetectedTrackers domains, err := trackers.LoadInitiatorDomainsByTrackerPatternID(ctx, tx, tp.ID, 5) @@ -238,7 +258,7 @@ func (h *trackerMappingHandler) identifyWithAgent( result, err := agent.RunTyped[TrackerMappingAgentResult]( agentCtx, - h.agent, + h.mappingAgent, []llm.Message{ { Role: llm.RoleUser, @@ -254,7 +274,7 @@ func (h *trackerMappingHandler) identifyWithAgent( log.String("pattern", tp.Pattern), ) - return nil, nil, nil + return nil, nil } identification := result.Output @@ -267,7 +287,7 @@ func (h *trackerMappingHandler) identifyWithAgent( log.Float64("confidence", identification.Confidence), ) - return nil, nil, nil + return nil, nil } confidence := float32(identification.Confidence) @@ -284,7 +304,7 @@ func (h *trackerMappingHandler) identifyWithAgent( domains, ) if err != nil { - return nil, nil, fmt.Errorf("cannot resolve or create common third party: %w", err) + return nil, fmt.Errorf("cannot resolve or create common third party: %w", err) } } @@ -303,12 +323,7 @@ func (h *trackerMappingHandler) identifyWithAgent( } if _, err := commonPattern.Upsert(ctx, tx); err != nil { - return nil, nil, fmt.Errorf("cannot upsert common tracker pattern from agent: %w", err) - } - - thirdPartyID, err := h.resolveThirdParty(ctx, tx, tp, &commonPattern) - if err != nil { - return nil, nil, fmt.Errorf("cannot resolve third party from agent match: %w", err) + return nil, fmt.Errorf("cannot upsert common tracker pattern from agent: %w", err) } h.logger.InfoCtx( @@ -319,7 +334,7 @@ func (h *trackerMappingHandler) identifyWithAgent( log.Float64("confidence", identification.Confidence), ) - return &commonPattern.ID, thirdPartyID, nil + return &commonPattern.ID, nil } func (h *trackerMappingHandler) resolveOrCreateCommonThirdParty( @@ -408,32 +423,163 @@ func (h *trackerMappingHandler) createUnmatchedPattern( return &commonPattern.ID, nil } -func (h *trackerMappingHandler) resolveThirdParty( +// promoteThirdParty resolves an org ThirdParty for the given pattern +// once the catalog mapping is known. 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. +// 4. Fallback create from CommonThirdParty. +// +// A confident heuristic/agent match is auto-tagged with +// common_third_party_id so subsequent promotions hit the exact-link +// path in O(1). Returns (nil, nil) when the catalog row has no +// CommonThirdPartyID — there is nothing to promote to. +func (h *trackerMappingHandler) promoteThirdParty( ctx context.Context, - conn pg.Querier, + tx pg.Tx, tp coredata.TrackerPattern, - commonPattern *coredata.CommonTrackerPattern, + commonPatternID gid.GID, ) (*gid.GID, error) { + var commonPattern coredata.CommonTrackerPattern + if err := commonPattern.LoadByID(ctx, tx, commonPatternID); err != nil { + return nil, fmt.Errorf("cannot load common tracker pattern: %w", err) + } + if commonPattern.CommonThirdPartyID == nil { return nil, nil } + commonThirdPartyID := *commonPattern.CommonThirdPartyID scope := coredata.NewScopeFromObjectID(tp.ID) - var t coredata.ThirdParty - if err := t.LoadByOrganizationIDAndCommonThirdPartyID( + var existing coredata.ThirdParty + + err := existing.LoadByOrganizationIDAndCommonThirdPartyID( ctx, - conn, + tx, scope, tp.OrganizationID, - *commonPattern.CommonThirdPartyID, - ); err != nil { - if errors.Is(err, coredata.ErrResourceNotFound) { - return nil, nil - } - - return nil, fmt.Errorf("cannot resolve third party: %w", err) + commonThirdPartyID, + ) + if err == nil { + return &existing.ID, nil } - return &t.ID, nil + if !errors.Is(err, coredata.ErrResourceNotFound) { + return nil, fmt.Errorf("cannot load org third party by common id: %w", err) + } + + var commonParty coredata.CommonThirdParty + if err := commonParty.LoadByID(ctx, tx, commonThirdPartyID); err != nil { + return nil, fmt.Errorf("cannot load common third party: %w", err) + } + + var commonDomains coredata.CommonThirdPartyDomains + if err := commonDomains.LoadByCommonThirdPartyID(ctx, tx, commonThirdPartyID); err != nil { + return nil, fmt.Errorf("cannot load common third party domains: %w", err) + } + + var orgThirdParties coredata.ThirdParties + if err := orgThirdParties.LoadAllByOrganizationID(ctx, tx, scope, tp.OrganizationID); err != nil { + return nil, fmt.Errorf("cannot load org third parties: %w", err) + } + + ranked := thirdparty.RankCandidates(commonParty, commonDomains, orgThirdParties) + + if len(ranked) > 0 && ranked[0].Score >= thirdparty.HighConfidenceScore { + picked := ranked[0].ThirdParty + + if err := thirdparty.LinkToCommon(ctx, tx, scope, picked, commonThirdPartyID); err != nil { + return nil, fmt.Errorf("cannot link fuzzy-matched third party to common: %w", err) + } + + 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", ranked[0].Score), + ) + + return &picked.ID, nil + } + + agentSet := ranked + if len(agentSet) > thirdparty.MaxAgentCandidates { + agentSet = agentSet[:thirdparty.MaxAgentCandidates] + } + + eligibleForAgent := false + + for _, c := range agentSet { + if c.Score >= thirdparty.MinAgentScore { + eligibleForAgent = true + + break + } + } + + if eligibleForAgent && h.disambiguationAgent != nil { + matchedID, err := thirdparty.Disambiguate( + ctx, + h.disambiguationAgent, + h.logger, + commonParty, + commonDomains, + agentSet, + ) + 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 { + var picked *coredata.ThirdParty + + for _, c := range agentSet { + if c.ThirdParty.ID == *matchedID { + picked = c.ThirdParty + + break + } + } + + if picked != nil { + if err := thirdparty.LinkToCommon(ctx, tx, scope, picked, commonThirdPartyID); err != nil { + return nil, fmt.Errorf("cannot link agent-matched third party to common: %w", err) + } + + 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()), + ) + + return &picked.ID, nil + } + } + } + + created, err := thirdparty.CreateFromCommon(ctx, tx, scope, tp.OrganizationID, commonParty) + if err != nil { + return nil, fmt.Errorf("cannot create third party from common: %w", err) + } + + h.logger.InfoCtx( + ctx, + "promoted tracker pattern by creating org third party from catalog", + log.String("tracker_pattern_id", tp.ID.String()), + log.String("third_party_id", created.ID.String()), + log.String("common_third_party_id", commonThirdPartyID.String()), + ) + + return &created.ID, nil } diff --git a/pkg/cookiebanner/tracker_mapping_worker_test.go b/pkg/cookiebanner/tracker_mapping_worker_test.go new file mode 100644 index 000000000..c07c504e8 --- /dev/null +++ b/pkg/cookiebanner/tracker_mapping_worker_test.go @@ -0,0 +1,495 @@ +// 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" + "io" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.gearno.de/kit/log" + "go.gearno.de/kit/pg" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/gid" +) + +//go:fix inline +func ptr[T any](v T) *T { return new(v) } + +// promotionFixture extends workerFixture with a CommonThirdParty and a +// CommonTrackerPattern linking the catalog to the test pattern. It is +// the minimum scaffolding promoteThirdParty needs to run end-to-end. +type promotionFixture struct { + workerFixture + commonThirdParty coredata.CommonThirdParty + commonPatternID gid.GID + trackerPattern coredata.TrackerPattern + commonThirdPartyID gid.GID +} + +func seedPromotionFixture(t *testing.T, ctx context.Context, client *pg.Client) promotionFixture { + t.Helper() + + fx := seedWorkerFixture(t, ctx, client) + now := time.Now().UTC().Truncate(time.Microsecond) + + commonThirdPartyID := gid.New(gid.NilTenant, coredata.CommonThirdPartyEntityType) + commonThirdParty := coredata.CommonThirdParty{ + ID: commonThirdPartyID, + Name: "Google", + Slug: "google", + Category: coredata.ThirdPartyCategoryAnalytics, + WebsiteURL: new("https://google.com"), + Certifications: []string{}, + CreatedAt: now, + UpdatedAt: now, + } + + commonPattern := coredata.CommonTrackerPattern{ + ID: gid.New(gid.NilTenant, coredata.CommonTrackerPatternEntityType), + CommonThirdPartyID: &commonThirdPartyID, + TrackerType: coredata.TrackerTypeCookie, + Pattern: "_ga", + MatchType: coredata.TrackerPatternMatchTypeExact, + Description: "", + Confidence: 0.9, + CreatedAt: now, + UpdatedAt: now, + } + + pattern := coredata.TrackerPattern{ + ID: gid.New(fx.scope.GetTenantID(), coredata.TrackerPatternEntityType), + OrganizationID: fx.organizationID, + CookieBannerID: fx.banner.ID, + CookieCategoryID: fx.normalCategoryID, + CommonTrackerPatternID: &commonPattern.ID, + TrackerType: coredata.TrackerTypeCookie, + Pattern: "_ga", + MatchType: coredata.TrackerPatternMatchTypeExact, + DisplayName: "_ga", + Description: "", + CreatedAt: now, + UpdatedAt: now, + } + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + if err := commonThirdParty.Insert(ctx, tx); err != nil { + return err + } + + if _, err := commonPattern.Upsert(ctx, tx); err != nil { + return err + } + + return pattern.Insert(ctx, tx, fx.scope) + })) + + t.Cleanup(func() { + _ = client.WithTx(context.Background(), func(ctx context.Context, tx pg.Tx) error { + if _, err := tx.Exec(ctx, `DELETE FROM common_third_party_domains WHERE common_third_party_id = $1`, commonThirdPartyID); err != nil { + return err + } + + if _, err := tx.Exec(ctx, `DELETE FROM common_tracker_patterns WHERE id = $1`, commonPattern.ID); err != nil { + return err + } + + if _, err := tx.Exec(ctx, `DELETE FROM common_third_parties WHERE id = $1`, commonThirdPartyID); err != nil { + return err + } + + if _, err := tx.Exec(ctx, `DELETE FROM third_parties WHERE organization_id = $1`, fx.organizationID); err != nil { + return err + } + + return nil + }) + }) + + return promotionFixture{ + workerFixture: fx, + commonThirdParty: commonThirdParty, + commonPatternID: commonPattern.ID, + commonThirdPartyID: commonThirdPartyID, + trackerPattern: pattern, + } +} + +func newMappingHandler(client *pg.Client) *trackerMappingHandler { + return &trackerMappingHandler{ + pg: client, + logger: log.NewLogger(log.WithOutput(io.Discard)), + } +} + +// promote runs promoteThirdParty inside its own transaction so each +// test case starts from a clean state. +func promote( + t *testing.T, + ctx context.Context, + h *trackerMappingHandler, + client *pg.Client, + tp coredata.TrackerPattern, + commonPatternID gid.GID, +) *gid.GID { + t.Helper() + + var got *gid.GID + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + var err error + + got, err = h.promoteThirdParty(ctx, tx, tp, commonPatternID) + + return err + })) + + return got +} + +func TestPromoteThirdParty_ExactCommonLink(t *testing.T) { + t.Parallel() + + client := newTestPgClient(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), client, fx.trackerPattern, fx.commonPatternID) + + 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 := newTestPgClient(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, + } + + 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), client, fx.trackerPattern, fx.commonPatternID) + + 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) +} + +func TestPromoteThirdParty_FallbackCreate(t *testing.T) { + t.Parallel() + + client := newTestPgClient(t) + ctx := context.Background() + fx := seedPromotionFixture(t, ctx, client) + + got := promote(t, ctx, newMappingHandler(client), client, fx.trackerPattern, fx.commonPatternID) + + require.NotNil(t, got, "fallback should create a new 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, *got) + })) + + assert.Equal(t, fx.organizationID, reloaded.OrganizationID) + assert.Equal(t, "Google", reloaded.Name) + require.NotNil(t, reloaded.CommonThirdPartyID) + assert.Equal(t, fx.commonThirdPartyID, *reloaded.CommonThirdPartyID) + assert.Equal(t, coredata.ThirdPartyCategoryAnalytics, reloaded.Category) + assert.False(t, reloaded.FirstLevel) + assert.False(t, reloaded.ShowOnTrustCenter) +} + +func TestPromoteThirdParty_NoCommonThirdPartyOnPattern(t *testing.T) { + t.Parallel() + + client := newTestPgClient(t) + ctx := context.Background() + fx := seedWorkerFixture(t, ctx, client) + + now := time.Now().UTC().Truncate(time.Microsecond) + commonPattern := coredata.CommonTrackerPattern{ + ID: gid.New(gid.NilTenant, coredata.CommonTrackerPatternEntityType), + TrackerType: coredata.TrackerTypeCookie, + Pattern: "unknown_xyz", + MatchType: coredata.TrackerPatternMatchTypeExact, + Description: "", + Confidence: 0.5, + CreatedAt: now, + UpdatedAt: now, + } + + pattern := coredata.TrackerPattern{ + ID: gid.New(fx.scope.GetTenantID(), coredata.TrackerPatternEntityType), + OrganizationID: fx.organizationID, + CookieBannerID: fx.banner.ID, + CookieCategoryID: fx.normalCategoryID, + CommonTrackerPatternID: &commonPattern.ID, + TrackerType: coredata.TrackerTypeCookie, + Pattern: "unknown_xyz", + MatchType: coredata.TrackerPatternMatchTypeExact, + DisplayName: "unknown_xyz", + CreatedAt: now, + UpdatedAt: now, + } + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + if _, err := commonPattern.Upsert(ctx, tx); err != nil { + return err + } + + return pattern.Insert(ctx, tx, fx.scope) + })) + + t.Cleanup(func() { + _ = client.WithTx(context.Background(), func(ctx context.Context, tx pg.Tx) error { + _, err := tx.Exec(ctx, `DELETE FROM common_tracker_patterns WHERE id = $1`, commonPattern.ID) + + return err + }) + }) + + got := promote(t, ctx, newMappingHandler(client), client, pattern, commonPattern.ID) + + assert.Nil(t, got, "patterns whose catalog row has no CommonThirdPartyID should not be promoted") +} + +// 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 +// existing catalog link is preserved verbatim. +func TestProcess_PreservesCatalogMappingOnReTrigger(t *testing.T) { + t.Parallel() + + client := newTestPgClient(t) + ctx := context.Background() + fx := seedPromotionFixture(t, ctx, client) + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + return fx.trackerPattern.SetMappingRequested(ctx, tx) + })) + + h := newMappingHandler(client) + require.NoError(t, h.Process(ctx, fx.trackerPattern)) + + 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.CommonTrackerPatternID, "common tracker pattern link must be preserved") + assert.Equal(t, fx.commonPatternID, *reloaded.CommonTrackerPatternID) + require.NotNil(t, reloaded.ThirdPartyID, "the worker should have promoted to an org ThirdParty") +} + +// TestProcess_ExtensionPatternIsNotPromoted asserts that even when a +// pattern has a catalog link, a Source=EXTENSION pattern stays +// un-promoted. +func TestProcess_ExtensionPatternIsNotPromoted(t *testing.T) { + t.Parallel() + + client := newTestPgClient(t) + ctx := context.Background() + fx := seedPromotionFixture(t, ctx, client) + + now := time.Now().UTC().Truncate(time.Microsecond) + source := coredata.CookieSourceExtension + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + _, err := tx.Exec( + ctx, + `UPDATE tracker_patterns + SET source = $1, + mapping_requested_at = $2 + WHERE id = $3`, + source, + 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) + })) + + assert.Nil(t, reloaded.ThirdPartyID, "EXTENSION-sourced pattern must not be promoted") +} + +// TestProcess_NoOpWhenAlreadyPromoted asserts that re-running the +// worker on a pattern that already has a third_party_id leaves the +// row alone (the guard in Process). +func TestProcess_NoOpWhenAlreadyPromoted(t *testing.T) { + t.Parallel() + + client := newTestPgClient(t) + ctx := context.Background() + fx := seedPromotionFixture(t, ctx, client) + + now := time.Now().UTC().Truncate(time.Microsecond) + preExisting := 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 := preExisting.Insert(ctx, tx, fx.scope); err != nil { + return err + } + + fx.trackerPattern.ThirdPartyID = &preExisting.ID + + _, err := tx.Exec( + ctx, + `UPDATE tracker_patterns + SET third_party_id = $1, + mapping_requested_at = $2 + WHERE id = $3`, + preExisting.ID, + 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) + assert.Equal(t, preExisting.ID, *reloaded.ThirdPartyID, "third_party_id must not be overwritten") +} + +func TestPromoteThirdParty_ExactCommonLinkIgnoresSimilarUnlinked(t *testing.T) { + t.Parallel() + + client := newTestPgClient(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), client, fx.trackerPattern, fx.commonPatternID) + + require.NotNil(t, got) + assert.Equal(t, linked.ID, *got, "exact-link path must short-circuit before the heuristic fires") +} diff --git a/pkg/coredata/third_party.go b/pkg/coredata/third_party.go index 6a51df4f9..c8d814c14 100644 --- a/pkg/coredata/third_party.go +++ b/pkg/coredata/third_party.go @@ -645,6 +645,7 @@ func (v *ThirdParty) Update( q := ` UPDATE third_parties SET + common_third_party_id = @common_third_party_id, name = @name, description = @description, category = @category, @@ -675,6 +676,7 @@ WHERE %s args := pgx.StrictNamedArgs{ "third_party_id": v.ID, "updated_at": time.Now(), + "common_third_party_id": v.CommonThirdPartyID, "name": v.Name, "description": v.Description, "category": v.Category, diff --git a/pkg/probod/probod.go b/pkg/probod/probod.go index 3fd557952..ce912a45f 100644 --- a/pkg/probod/probod.go +++ b/pkg/probod/probod.go @@ -313,7 +313,7 @@ func (impl *Implm) Run( return err } - trackerMappingCfg, err := impl.buildTrackerMappingConfig(l, tp, r) + trackerMappingCfg, thirdPartyDisambiguationCfg, err := impl.buildTrackerMappingConfig(l, tp, r) if err != nil { return err } @@ -725,7 +725,7 @@ func (impl *Implm) Run( }, ) - trackerMappingWorker := cookiebanner.NewTrackerMappingWorker(pgClient, l, trackerMappingCfg) + trackerMappingWorker := cookiebanner.NewTrackerMappingWorker(pgClient, l, trackerMappingCfg, thirdPartyDisambiguationCfg) trackerMappingWorkerCtx, stopTrackerMappingWorker := context.WithCancel(context.Background()) wg.Go( diff --git a/pkg/probod/tracker_mapping.go b/pkg/probod/tracker_mapping.go index 0b13a25b1..7414b0d57 100644 --- a/pkg/probod/tracker_mapping.go +++ b/pkg/probod/tracker_mapping.go @@ -21,18 +21,28 @@ import ( "go.gearno.de/kit/log" "go.opentelemetry.io/otel/trace" "go.probo.inc/probo/pkg/cookiebanner" + "go.probo.inc/probo/pkg/thirdparty" ) -// buildTrackerMappingConfig wires the tracker-mapping agent. It is opt-in: -// deployments that do not set `llm.tracker-mapping.provider` get a zero -// config (nil LLM client) so the worker runs without agent fallback. +// 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. +// +// Both 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( l *log.Logger, tp trace.TracerProvider, r prometheus.Registerer, -) (cookiebanner.TrackerMappingConfig, error) { +) (cookiebanner.TrackerMappingConfig, thirdparty.DisambiguationConfig, error) { if impl.cfg.Agents.TrackerMapping.Provider == "" { - return cookiebanner.TrackerMappingConfig{}, nil + return cookiebanner.TrackerMappingConfig{}, thirdparty.DisambiguationConfig{}, nil } agentCfg, llmClient, err := impl.resolveAgentClient( @@ -43,12 +53,19 @@ func (impl *Implm) buildTrackerMappingConfig( r, ) if err != nil { - return cookiebanner.TrackerMappingConfig{}, fmt.Errorf("cannot resolve tracker mapping agent client: %w", err) + return cookiebanner.TrackerMappingConfig{}, thirdparty.DisambiguationConfig{}, fmt.Errorf("cannot resolve tracker mapping agent client: %w", err) } - return cookiebanner.TrackerMappingConfig{ + mappingCfg := cookiebanner.TrackerMappingConfig{ LLMClient: llmClient, Model: agentCfg.ModelName, FirecrawlAPIKey: impl.cfg.Agents.Tools.FirecrawlAPIKey, - }, nil + } + + disambiguationCfg := thirdparty.DisambiguationConfig{ + LLMClient: llmClient, + Model: agentCfg.ModelName, + } + + return mappingCfg, disambiguationCfg, nil } diff --git a/pkg/thirdparty/disambiguation_agent.go b/pkg/thirdparty/disambiguation_agent.go new file mode 100644 index 000000000..93ecaec41 --- /dev/null +++ b/pkg/thirdparty/disambiguation_agent.go @@ -0,0 +1,209 @@ +// 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 + + // disambiguationTimeout caps a single disambiguation run. The + // agent has no tools and a single turn, so this is mostly a + // guard against a hung LLM provider, not a real budget. + disambiguationTimeout = 60 * time.Second +) + +// DisambiguationConfig 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. +type DisambiguationConfig struct { + LLMClient *llm.Client + Model string +} + +// 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 DisambiguationConfig, + 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)) + } + + return agent.New( + "third-party-disambiguation", + cfg.LLMClient, + agent.WithInstructions(disambiguationPrompt), + agent.WithModel(cfg.Model), + agent.WithOutputType(outputType), + agent.WithMaxTurns(1), + agent.WithLogger(logger), + ) +} + +// 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, +) (*gid.GID, error) { + if a == nil || len(candidates) == 0 { + return nil, nil + } + + prompt := buildDisambiguationPrompt(commonParty, commonDomains, candidates) + + agentCtx, cancel := context.WithTimeout(ctx, disambiguationTimeout) + 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() +} diff --git a/pkg/thirdparty/match.go b/pkg/thirdparty/match.go new file mode 100644 index 000000000..787d301c1 --- /dev/null +++ b/pkg/thirdparty/match.go @@ -0,0 +1,274 @@ +// 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" + "time" + + "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, fall back to creating, 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 commonID; otherwise writes the field via ThirdParty.Update and +// updates the receiver in place. +func LinkToCommon( + ctx context.Context, + tx pg.Tx, + scope coredata.Scoper, + orgThirdParty *coredata.ThirdParty, + commonID gid.GID, +) error { + if orgThirdParty.CommonThirdPartyID != nil && *orgThirdParty.CommonThirdPartyID == commonID { + 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 +} + +// CreateFromCommon inserts a new org ThirdParty seeded from the catalog +// row (name, category, addresses, URLs, certifications, …). The new row +// has common_third_party_id pointed at commonParty, an empty Countries +// list, and ShowOnTrustCenter / FirstLevel both false — mirroring the +// front-end CreateThirdPartyDialog's "pick from catalog" seeding shape +// so the result is indistinguishable from a manual creation. +// +// Deliberately bypasses any service-level webhook emission: callers +// that need a webhook for the implicit creation should emit it +// themselves. +func CreateFromCommon( + ctx context.Context, + tx pg.Tx, + scope coredata.Scoper, + organizationID gid.GID, + commonParty coredata.CommonThirdParty, +) (*coredata.ThirdParty, error) { + commonID := commonParty.ID + now := time.Now() + + tp := &coredata.ThirdParty{ + ID: gid.New(scope.GetTenantID(), coredata.ThirdPartyEntityType), + OrganizationID: organizationID, + CommonThirdPartyID: &commonID, + Name: commonParty.Name, + Category: commonParty.Category, + HeadquarterAddress: commonParty.HeadquarterAddress, + LegalName: commonParty.LegalName, + WebsiteURL: commonParty.WebsiteURL, + PrivacyPolicyURL: commonParty.PrivacyPolicyURL, + ServiceLevelAgreementURL: commonParty.ServiceLevelAgreementURL, + DataProcessingAgreementURL: commonParty.DataProcessingAgreementURL, + BusinessAssociateAgreementURL: commonParty.BusinessAssociateAgreementURL, + SubprocessorsListURL: commonParty.SubprocessorsListURL, + Certifications: commonParty.Certifications, + Countries: coredata.CountryCodes{}, + StatusPageURL: commonParty.StatusPageURL, + TermsOfServiceURL: commonParty.TermsOfServiceURL, + SecurityPageURL: commonParty.SecurityPageURL, + TrustPageURL: commonParty.TrustPageURL, + ShowOnTrustCenter: false, + FirstLevel: false, + CreatedAt: now, + UpdatedAt: now, + } + + if tp.Certifications == nil { + tp.Certifications = []string{} + } + + if err := tp.Insert(ctx, tx, scope); err != nil { + return nil, fmt.Errorf("cannot insert org third party: %w", err) + } + + return tp, nil +} diff --git a/pkg/thirdparty/match_test.go b/pkg/thirdparty/match_test.go new file mode 100644 index 000000000..8eb075aa0 --- /dev/null +++ b/pkg/thirdparty/match_test.go @@ -0,0 +1,181 @@ +// 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" +) + +//go:fix inline +func ptr[T any](v T) *T { return new(v) } + +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/prompts/disambiguation.txt.tmpl b/pkg/thirdparty/prompts/disambiguation.txt.tmpl new file mode 100644 index 000000000..7879a5216 --- /dev/null +++ b/pkg/thirdparty/prompts/disambiguation.txt.tmpl @@ -0,0 +1,25 @@ + +You are a third-party catalog matcher. The product groups web trackers under a global catalog of "common" third parties (Google Analytics, Stripe, Meta Pixel, …). Each customer organisation also maintains its own list of "third parties" — sometimes seeded from the catalog, sometimes typed manually. Your only job is to decide whether one of the organisation's existing third parties already represents a given catalog entry, so we don't create a duplicate. + + + +You are given: +- A catalog third party: name, website, and known domains. +- A small list of candidate organisation third parties: each with a stable id, a name, and (optionally) a website. + +Pick the candidate that best represents the catalog third party, or none. + +Return a structured JSON response with: +- matched_id: the candidate id, or null if none of them is a confident match. +- confidence: 0.0 to 1.0; below 0.6 means "no confident match" (set matched_id to null in that case). +- reasoning: one short sentence describing the rationale. + + + +1. The candidates have already been ranked by a heuristic; the list is small (usually a handful). Use that as a hint, but do not trust it blindly. +2. Treat corporate suffixes (LLC, Inc, Ltd, GmbH, SA, …) as noise. "Google" and "Google LLC" are the same company. +3. Treat brand/product names as the same when the parent company is obvious: "Google Analytics" maps to "Google" if the org only has the parent. Only do this when the catalog domains plainly match the parent's domains. +4. Website hostnames and known domains are the strongest signal. If the catalog domain (or its eTLD+1) matches a candidate's website host, they are almost certainly the same. +5. Be conservative. If two candidates look plausible and you cannot rule one out, return matched_id=null with confidence < 0.6 — we will create a fresh org third party from the catalog rather than risk a wrong link. +6. Do not invent ids. Return only ids that appear verbatim in the candidate list. +