Reuse mapping agent to attribute trackers in enricher
The enrichment worker no longer invents a description when a tracker's purpose cannot be substantiated; it records an empty description and marks the row enriched so the stale-recovery loop does not retry it. Vendor identification is the mapping pipeline's job, so the enricher reuses the existing tracker-mapping agent to attribute a third party for an unlinked common pattern before describing it. A confident catalog match seeds the enrichment prompt and links the pattern, but the enricher never creates or overrides an attribution. When a blank, unlinked catalog row later gains a third party through the mapping pipeline's upsert, enrichment is re-armed so the now-known vendor gets a second, better-informed description attempt. Signed-off-by: Émile Ré <emile@probo.com>
This commit is contained in:
@@ -32,7 +32,7 @@ var trackerEnrichmentPrompt string
|
||||
// CommonPatternEnrichmentResult is the structured output the
|
||||
// common-pattern enrichment agent returns.
|
||||
type CommonPatternEnrichmentResult struct {
|
||||
Description string `json:"description" jsonschema:"A concise, factual, compliance-grade description of what this tracker stores or does and its purpose. One or two sentences. Name the operating company when known."`
|
||||
Description string `json:"description" jsonschema:"A concise, factual, compliance-grade description of what this tracker stores or does and its purpose. One or two sentences. Name the operating company when known. Empty when the purpose cannot be substantiated from evidence."`
|
||||
}
|
||||
|
||||
func buildCommonPatternEnrichmentAgent(
|
||||
|
||||
@@ -26,6 +26,7 @@ import (
|
||||
"go.gearno.de/kit/worker"
|
||||
"go.probo.inc/probo/pkg/agent"
|
||||
"go.probo.inc/probo/pkg/coredata"
|
||||
"go.probo.inc/probo/pkg/gid"
|
||||
"go.probo.inc/probo/pkg/llm"
|
||||
)
|
||||
|
||||
@@ -35,6 +36,7 @@ type commonPatternEnrichmentHandler struct {
|
||||
pg *pg.Client
|
||||
logger *log.Logger
|
||||
enrichmentAgent *agent.Agent
|
||||
mappingAgent *agent.Agent
|
||||
staleAfter time.Duration
|
||||
agentTimeout time.Duration
|
||||
}
|
||||
@@ -71,6 +73,7 @@ func NewCommonPatternEnrichmentWorker(
|
||||
|
||||
if cfg.LLMClient != nil {
|
||||
h.enrichmentAgent = buildCommonPatternEnrichmentAgent(cfg, pgClient, logger)
|
||||
h.mappingAgent = buildTrackerMappingAgent(cfg, pgClient, logger)
|
||||
}
|
||||
|
||||
return worker.New(
|
||||
@@ -114,27 +117,49 @@ func (h *commonPatternEnrichmentHandler) Process(ctx context.Context, cp coredat
|
||||
return err
|
||||
}
|
||||
|
||||
// Map before enriching: an unlinked pattern is run through the
|
||||
// mapping agent first so a resolved vendor both seeds the enrichment
|
||||
// prompt and gets linked. Attribution stays the mapping pipeline's
|
||||
// job; the enricher only reuses it.
|
||||
var thirdPartyID *gid.GID
|
||||
|
||||
if cp.CommonThirdPartyID == nil {
|
||||
id, name, err := h.identifyThirdParty(ctx, cp)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
thirdPartyID = id
|
||||
if name != "" {
|
||||
thirdPartyName = name
|
||||
}
|
||||
}
|
||||
|
||||
description, err := h.research(ctx, cp, thirdPartyName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot research tracker description: %w", err)
|
||||
}
|
||||
|
||||
if description == "" {
|
||||
return fmt.Errorf("enrichment produced empty description for pattern %q", cp.Pattern)
|
||||
}
|
||||
|
||||
return h.pg.WithTx(
|
||||
ctx,
|
||||
func(ctx context.Context, tx pg.Tx) error {
|
||||
if err := cp.SetEnriched(ctx, tx, description); err != nil {
|
||||
// A blank description is recorded as a terminal-for-now state:
|
||||
// the row is marked enriched so the stale-recovery loop never
|
||||
// re-queues it, but a later third-party link (mapping worker)
|
||||
// re-arms enrichment for a vendor-informed second attempt.
|
||||
if err := cp.SetEnriched(ctx, tx, description, thirdPartyID); err != nil {
|
||||
return fmt.Errorf("cannot set common tracker pattern enriched: %w", err)
|
||||
}
|
||||
|
||||
var patterns coredata.TrackerPatterns
|
||||
var backfilled int64
|
||||
|
||||
count, err := patterns.BackfillDescriptionByCommonTrackerPatternID(ctx, tx, cp.ID, description)
|
||||
if err != nil {
|
||||
return err
|
||||
if description != "" {
|
||||
var patterns coredata.TrackerPatterns
|
||||
|
||||
backfilled, err = patterns.BackfillDescriptionByCommonTrackerPatternID(ctx, tx, cp.ID, description)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
h.logger.InfoCtx(
|
||||
@@ -142,7 +167,9 @@ func (h *commonPatternEnrichmentHandler) Process(ctx context.Context, cp coredat
|
||||
"enriched common tracker pattern",
|
||||
log.String("common_tracker_pattern_id", cp.ID.String()),
|
||||
log.String("pattern", cp.Pattern),
|
||||
log.Int64("backfilled_tracker_patterns", count),
|
||||
log.Bool("described", description != ""),
|
||||
log.Bool("third_party_linked", thirdPartyID != nil),
|
||||
log.Int64("backfilled_tracker_patterns", backfilled),
|
||||
)
|
||||
|
||||
return nil
|
||||
@@ -150,6 +177,45 @@ func (h *commonPatternEnrichmentHandler) Process(ctx context.Context, cp coredat
|
||||
)
|
||||
}
|
||||
|
||||
// resolveThirdPartyID maps the agent's returned company name to an
|
||||
// existing catalog third party, but only when the pattern has none yet:
|
||||
// the enrichment worker links, it never overrides an attribution the
|
||||
// mapping pipeline already resolved. A name that matches no catalog row
|
||||
// resolves to nil, so the worker never invents a third party.
|
||||
func (h *commonPatternEnrichmentHandler) resolveThirdPartyID(
|
||||
ctx context.Context,
|
||||
cp coredata.CommonTrackerPattern,
|
||||
thirdPartyName string,
|
||||
) (*gid.GID, error) {
|
||||
if cp.CommonThirdPartyID != nil || thirdPartyName == "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var id *gid.GID
|
||||
|
||||
if err := h.pg.WithConn(
|
||||
ctx,
|
||||
func(ctx context.Context, conn pg.Querier) error {
|
||||
var party coredata.CommonThirdParty
|
||||
if err := party.LoadByName(ctx, conn, thirdPartyName); err != nil {
|
||||
if errors.Is(err, coredata.ErrResourceNotFound) {
|
||||
return nil
|
||||
}
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
id = &party.ID
|
||||
|
||||
return nil
|
||||
},
|
||||
); err != nil {
|
||||
return nil, fmt.Errorf("cannot resolve common third party for enrichment: %w", err)
|
||||
}
|
||||
|
||||
return id, nil
|
||||
}
|
||||
|
||||
func (h *commonPatternEnrichmentHandler) RecoverStale(ctx context.Context) error {
|
||||
return h.pg.WithConn(
|
||||
ctx,
|
||||
@@ -218,3 +284,61 @@ func (h *commonPatternEnrichmentHandler) research(
|
||||
|
||||
return strings.TrimSpace(result.Output.Description), nil
|
||||
}
|
||||
|
||||
// identifyThirdParty reuses the tracker-mapping agent to attribute a
|
||||
// vendor to an unlinked catalog pattern. It returns the resolved
|
||||
// existing third party id and its name (for the enrichment prompt) only
|
||||
// when the agent is confident and the name matches a catalog row;
|
||||
// otherwise it returns nils so enrichment proceeds without a vendor. A
|
||||
// failed agent run is best-effort and non-fatal, mirroring the mapping
|
||||
// worker's identifyWithAgent.
|
||||
func (h *commonPatternEnrichmentHandler) identifyThirdParty(
|
||||
ctx context.Context,
|
||||
cp coredata.CommonTrackerPattern,
|
||||
) (*gid.GID, string, error) {
|
||||
if h.mappingAgent == nil {
|
||||
return nil, "", nil
|
||||
}
|
||||
|
||||
prompt := buildCommonPatternIdentificationPrompt(cp)
|
||||
|
||||
agentCtx, cancel := context.WithTimeout(ctx, h.agentTimeout)
|
||||
defer cancel()
|
||||
|
||||
result, err := agent.RunTyped[TrackerMappingAgentResult](
|
||||
agentCtx,
|
||||
h.mappingAgent,
|
||||
[]llm.Message{
|
||||
{
|
||||
Role: llm.RoleUser,
|
||||
Parts: []llm.Part{llm.TextPart{Text: prompt}},
|
||||
},
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
h.logger.WarnCtx(
|
||||
ctx,
|
||||
"mapping agent identification failed during enrichment",
|
||||
log.Error(err),
|
||||
log.String("pattern", cp.Pattern),
|
||||
)
|
||||
|
||||
return nil, "", nil
|
||||
}
|
||||
|
||||
name := strings.TrimSpace(result.Output.ThirdPartyName)
|
||||
if name == "" || result.Output.ThirdPartyConfidence < agentThirdPartyConfidenceThreshold {
|
||||
return nil, "", nil
|
||||
}
|
||||
|
||||
id, err := h.resolveThirdPartyID(ctx, cp, name)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
if id == nil {
|
||||
return nil, "", nil
|
||||
}
|
||||
|
||||
return id, name, nil
|
||||
}
|
||||
|
||||
126
pkg/cookiebanner/common_pattern_enrichment_worker_test.go
Normal file
126
pkg/cookiebanner/common_pattern_enrichment_worker_test.go
Normal file
@@ -0,0 +1,126 @@
|
||||
// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
|
||||
//
|
||||
// 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"
|
||||
)
|
||||
|
||||
func newEnrichmentHandler(client *pg.Client) *commonPatternEnrichmentHandler {
|
||||
return &commonPatternEnrichmentHandler{
|
||||
pg: client,
|
||||
logger: log.NewLogger(log.WithOutput(io.Discard)),
|
||||
}
|
||||
}
|
||||
|
||||
// seedEnrichmentThirdParty inserts a collision-free catalog third party
|
||||
// for the resolver to match against.
|
||||
func seedEnrichmentThirdParty(t *testing.T, ctx context.Context, client *pg.Client) coredata.CommonThirdParty {
|
||||
t.Helper()
|
||||
|
||||
now := time.Now().UTC().Truncate(time.Microsecond)
|
||||
id := gid.New(gid.NilTenant, coredata.CommonThirdPartyEntityType)
|
||||
suffix := id.String()
|
||||
|
||||
party := coredata.CommonThirdParty{
|
||||
ID: id,
|
||||
Name: "Hotjar " + suffix,
|
||||
Slug: "hotjar-" + suffix,
|
||||
Category: coredata.ThirdPartyCategoryAnalytics,
|
||||
Certifications: []string{},
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
|
||||
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
|
||||
return party.Insert(ctx, tx)
|
||||
}))
|
||||
|
||||
t.Cleanup(func() {
|
||||
_ = client.WithTx(context.Background(), func(ctx context.Context, tx pg.Tx) error {
|
||||
_, err := tx.Exec(ctx, `DELETE FROM common_third_parties WHERE id = $1`, id)
|
||||
return err
|
||||
})
|
||||
})
|
||||
|
||||
return party
|
||||
}
|
||||
|
||||
// TestResolveThirdPartyID pins the enrichment worker's third-party
|
||||
// resolution: it links the agent's returned company to an existing
|
||||
// catalog row by name, but only when the pattern has no third party yet,
|
||||
// and it never invents one for a name absent from the catalog.
|
||||
func TestResolveThirdPartyID(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
client := newTestPgClient(t)
|
||||
ctx := context.Background()
|
||||
h := newEnrichmentHandler(client)
|
||||
|
||||
party := seedEnrichmentThirdParty(t, ctx, client)
|
||||
existingID := gid.New(gid.NilTenant, coredata.CommonThirdPartyEntityType)
|
||||
|
||||
t.Run("links existing catalog third party when unset", func(t *testing.T) {
|
||||
cp := coredata.CommonTrackerPattern{}
|
||||
|
||||
got, err := h.resolveThirdPartyID(ctx, cp, party.Name)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, got)
|
||||
assert.Equal(t, party.ID, *got)
|
||||
})
|
||||
|
||||
t.Run("matches catalog name case-insensitively", func(t *testing.T) {
|
||||
cp := coredata.CommonTrackerPattern{}
|
||||
|
||||
got, err := h.resolveThirdPartyID(ctx, cp, "hOtJaR "+party.ID.String())
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, got)
|
||||
assert.Equal(t, party.ID, *got)
|
||||
})
|
||||
|
||||
t.Run("does not resolve when pattern already linked", func(t *testing.T) {
|
||||
cp := coredata.CommonTrackerPattern{CommonThirdPartyID: &existingID}
|
||||
|
||||
got, err := h.resolveThirdPartyID(ctx, cp, party.Name)
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, got, "must not override an existing third-party link")
|
||||
})
|
||||
|
||||
t.Run("returns nil for a name absent from the catalog", func(t *testing.T) {
|
||||
cp := coredata.CommonTrackerPattern{}
|
||||
|
||||
got, err := h.resolveThirdPartyID(ctx, cp, "Nonexistent Vendor "+party.ID.String())
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, got, "must not invent a third party")
|
||||
})
|
||||
|
||||
t.Run("returns nil for an empty name", func(t *testing.T) {
|
||||
cp := coredata.CommonTrackerPattern{}
|
||||
|
||||
got, err := h.resolveThirdPartyID(ctx, cp, "")
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, got)
|
||||
})
|
||||
}
|
||||
@@ -6,7 +6,7 @@ You are a privacy and web-tracking compliance expert. Your job is to write an ac
|
||||
Given a tracker pattern (cookie name, local storage key, etc.), its type, max-age, and the third party that operates it when known, produce a concise factual description of the tracker's purpose.
|
||||
|
||||
Return a structured JSON response with:
|
||||
- description: one or two sentences describing what this tracker stores or does and the purpose it serves (e.g. analytics, advertising, session management, security). When the operating company is known, name it.
|
||||
- description: one or two sentences describing what this tracker stores or does and the purpose it serves (e.g. analytics, advertising, session management, security). When the operating company is known, name it. Leave this empty when you cannot substantiate the purpose from evidence.
|
||||
</task>
|
||||
|
||||
<instructions>
|
||||
@@ -23,5 +23,5 @@ Return a structured JSON response with:
|
||||
|
||||
4. Keep the description concise (one to two sentences) and free of marketing language. It should read as a neutral, compliance-grade statement of purpose.
|
||||
|
||||
5. If you genuinely cannot determine the tracker's purpose, write a minimal factual description based on its type and name (e.g. "Cookie set by an unidentified third party; purpose could not be determined.") rather than inventing details.
|
||||
5. If you genuinely cannot substantiate the tracker's purpose from evidence, return an empty description. Do not write a fallback such as "purpose could not be determined", and do not guess a purpose from the name or max-age alone (e.g. do not claim a key is "used for session" just because it has no expiry). An empty description is preferable to an unverified one.
|
||||
</instructions>
|
||||
|
||||
@@ -161,3 +161,27 @@ func buildAgentPrompt(tp coredata.TrackerPattern, domains []string) string {
|
||||
|
||||
return prompt
|
||||
}
|
||||
|
||||
// buildCommonPatternIdentificationPrompt builds the mapping-agent input
|
||||
// for a global catalog pattern. Catalog rows carry no observed domains,
|
||||
// so the prompt omits the <observed_domains> signal and relies on the
|
||||
// pattern name, type, and naming conventions. It lets the enrichment
|
||||
// worker reuse the mapping agent to attribute a vendor before describing.
|
||||
func buildCommonPatternIdentificationPrompt(cp coredata.CommonTrackerPattern) string {
|
||||
maxAge := "session"
|
||||
if cp.MaxAgeSeconds != nil {
|
||||
maxAge = fmt.Sprintf("%d seconds", *cp.MaxAgeSeconds)
|
||||
}
|
||||
|
||||
return fmt.Sprintf(
|
||||
"Identify the following tracker:\n\n"+
|
||||
"<pattern> %s </pattern>\n"+
|
||||
"<type> %s </type>\n"+
|
||||
"<match_type> %s </match_type>\n"+
|
||||
"<max_age> %s </max_age>\n",
|
||||
cp.Pattern,
|
||||
cp.TrackerType,
|
||||
cp.MatchType,
|
||||
maxAge,
|
||||
)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user