Add LLM agent fallback for unmapped tracker patterns
When both pattern matching and domain matching fail to identify a tracker, an opt-in LLM agent can now attempt identification using internal database searches and optional web search. The agent returns structured output (third party name, category, description, confidence) and the worker auto-creates CommonThirdParty records when needed. The feature is gated behind the `llm.tracker-mapping.provider` config field; when unset the worker behaves exactly as before. Signed-off-by: Émile Ré <emile@probo.com>
This commit is contained in:
Émile Ré
@@ -16,25 +16,50 @@ package cookiebanner
|
||||
|
||||
import (
|
||||
"context"
|
||||
_ "embed"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.gearno.de/kit/log"
|
||||
"go.gearno.de/kit/pg"
|
||||
"go.gearno.de/kit/worker"
|
||||
"go.probo.inc/probo/pkg/agent"
|
||||
"go.probo.inc/probo/pkg/agent/tools/search"
|
||||
"go.probo.inc/probo/pkg/coredata"
|
||||
"go.probo.inc/probo/pkg/gid"
|
||||
"go.probo.inc/probo/pkg/llm"
|
||||
"go.probo.inc/probo/pkg/slug"
|
||||
)
|
||||
|
||||
const (
|
||||
agentTimeout = 60 * time.Second
|
||||
agentMaxTurns = 5
|
||||
agentConfidenceThreshold = 0.6
|
||||
agentMaxPatternConfidence = 0.8
|
||||
)
|
||||
|
||||
//go:embed prompts/tracker_identification.txt
|
||||
var trackerIdentificationPrompt string
|
||||
|
||||
type trackerMappingHandler struct {
|
||||
pg *pg.Client
|
||||
logger *log.Logger
|
||||
agent *agent.Agent
|
||||
}
|
||||
|
||||
type TrackerMappingConfig struct {
|
||||
LLMClient *llm.Client
|
||||
Model string
|
||||
SearchEndpoint string
|
||||
}
|
||||
|
||||
func NewTrackerMappingWorker(
|
||||
pgClient *pg.Client,
|
||||
logger *log.Logger,
|
||||
cfg TrackerMappingConfig,
|
||||
opts ...worker.Option,
|
||||
) *worker.Worker[coredata.TrackerPattern] {
|
||||
h := &trackerMappingHandler{
|
||||
@@ -42,6 +67,16 @@ func NewTrackerMappingWorker(
|
||||
logger: logger,
|
||||
}
|
||||
|
||||
if cfg.LLMClient != nil {
|
||||
h.agent = buildTrackerMappingAgent(
|
||||
cfg.LLMClient,
|
||||
cfg.Model,
|
||||
cfg.SearchEndpoint,
|
||||
pgClient,
|
||||
logger,
|
||||
)
|
||||
}
|
||||
|
||||
return worker.New(
|
||||
"tracker-mapping-worker",
|
||||
h,
|
||||
@@ -50,6 +85,39 @@ func NewTrackerMappingWorker(
|
||||
)
|
||||
}
|
||||
|
||||
func buildTrackerMappingAgent(
|
||||
llmClient *llm.Client,
|
||||
model string,
|
||||
searchEndpoint string,
|
||||
pgClient *pg.Client,
|
||||
logger *log.Logger,
|
||||
) *agent.Agent {
|
||||
tools := []agent.Tool{
|
||||
searchTrackerPatternsTool(pgClient),
|
||||
searchThirdPartiesTool(pgClient),
|
||||
}
|
||||
|
||||
if searchEndpoint != "" {
|
||||
tools = append(tools, search.WebSearchTool(searchEndpoint))
|
||||
}
|
||||
|
||||
outputType, err := agent.NewOutputType[TrackerIdentification]("tracker_identification")
|
||||
if err != nil {
|
||||
panic(fmt.Sprintf("cookiebanner: cannot build tracker identification output type: %s", err))
|
||||
}
|
||||
|
||||
return agent.New(
|
||||
"tracker-mapping",
|
||||
llmClient,
|
||||
agent.WithInstructions(trackerIdentificationPrompt),
|
||||
agent.WithModel(model),
|
||||
agent.WithTools(tools...),
|
||||
agent.WithOutputType(outputType),
|
||||
agent.WithMaxTurns(agentMaxTurns),
|
||||
agent.WithLogger(logger),
|
||||
)
|
||||
}
|
||||
|
||||
func (h *trackerMappingHandler) Claim(ctx context.Context) (coredata.TrackerPattern, error) {
|
||||
var tp coredata.TrackerPattern
|
||||
|
||||
@@ -85,6 +153,10 @@ func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.Tracker
|
||||
commonPatternID, thirdPartyID = h.matchByDomain(ctx, tx, tp)
|
||||
}
|
||||
|
||||
if commonPatternID == nil && h.agent != nil {
|
||||
commonPatternID, thirdPartyID = h.identifyWithAgent(ctx, tx, tp)
|
||||
}
|
||||
|
||||
if commonPatternID == nil {
|
||||
commonPatternID = h.createUnmatchedPattern(ctx, tx, tp)
|
||||
}
|
||||
@@ -183,6 +255,195 @@ func (h *trackerMappingHandler) matchByDomain(
|
||||
return &commonPattern.ID, thirdPartyID
|
||||
}
|
||||
|
||||
func (h *trackerMappingHandler) identifyWithAgent(
|
||||
ctx context.Context,
|
||||
tx pg.Tx,
|
||||
tp coredata.TrackerPattern,
|
||||
) (*gid.GID, *gid.GID) {
|
||||
var trackers coredata.DetectedTrackers
|
||||
domains, err := trackers.LoadInitiatorDomainsByTrackerPatternID(ctx, tx, tp.ID)
|
||||
if err != nil {
|
||||
h.logger.WarnCtx(ctx, "cannot load initiator domains for agent", log.Error(err))
|
||||
}
|
||||
|
||||
prompt := buildAgentPrompt(tp, domains)
|
||||
|
||||
agentCtx, cancel := context.WithTimeout(ctx, agentTimeout)
|
||||
defer cancel()
|
||||
|
||||
result, err := h.agent.Run(
|
||||
agentCtx,
|
||||
[]llm.Message{
|
||||
{
|
||||
Role: llm.RoleUser,
|
||||
Parts: []llm.Part{llm.TextPart{Text: prompt}},
|
||||
},
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
h.logger.WarnCtx(
|
||||
ctx,
|
||||
"agent identification failed",
|
||||
log.Error(err),
|
||||
log.String("pattern", tp.Pattern),
|
||||
)
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var identification TrackerIdentification
|
||||
if err := json.Unmarshal([]byte(result.FinalMessage().Text()), &identification); err != nil {
|
||||
h.logger.WarnCtx(
|
||||
ctx,
|
||||
"cannot parse agent identification output",
|
||||
log.Error(err),
|
||||
log.String("pattern", tp.Pattern),
|
||||
)
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
if identification.Confidence < agentConfidenceThreshold {
|
||||
h.logger.InfoCtx(
|
||||
ctx,
|
||||
"agent identification below confidence threshold",
|
||||
log.String("pattern", tp.Pattern),
|
||||
log.Float64("confidence", identification.Confidence),
|
||||
)
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
confidence := float32(identification.Confidence)
|
||||
if confidence > agentMaxPatternConfidence {
|
||||
confidence = agentMaxPatternConfidence
|
||||
}
|
||||
|
||||
var commonThirdPartyID *gid.GID
|
||||
if identification.ThirdPartyName != "" {
|
||||
commonThirdPartyID = h.resolveOrCreateCommonThirdParty(
|
||||
ctx,
|
||||
tx,
|
||||
identification,
|
||||
domains,
|
||||
)
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
commonPattern := coredata.CommonTrackerPattern{
|
||||
ID: gid.New(gid.NilTenant, coredata.CommonTrackerPatternEntityType),
|
||||
CommonThirdPartyID: commonThirdPartyID,
|
||||
TrackerType: tp.TrackerType,
|
||||
Pattern: tp.Pattern,
|
||||
MatchType: tp.MatchType,
|
||||
Description: identification.Description,
|
||||
MaxAgeSeconds: tp.MaxAgeSeconds,
|
||||
Confidence: confidence,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
|
||||
actualID, _, err := commonPattern.Upsert(ctx, tx)
|
||||
if err != nil {
|
||||
h.logger.ErrorCtx(
|
||||
ctx,
|
||||
"cannot upsert common tracker pattern from agent",
|
||||
log.Error(err),
|
||||
log.String("pattern", tp.Pattern),
|
||||
)
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
commonPattern.ID = actualID
|
||||
thirdPartyID, err := h.resolveThirdParty(ctx, tx, tp, &commonPattern)
|
||||
if err != nil {
|
||||
h.logger.ErrorCtx(ctx, "cannot resolve third party from agent match", log.Error(err))
|
||||
return &commonPattern.ID, nil
|
||||
}
|
||||
|
||||
h.logger.InfoCtx(
|
||||
ctx,
|
||||
"agent identified tracker pattern",
|
||||
log.String("pattern", tp.Pattern),
|
||||
log.String("third_party", identification.ThirdPartyName),
|
||||
log.Float64("confidence", identification.Confidence),
|
||||
)
|
||||
|
||||
return &commonPattern.ID, thirdPartyID
|
||||
}
|
||||
|
||||
func (h *trackerMappingHandler) resolveOrCreateCommonThirdParty(
|
||||
ctx context.Context,
|
||||
tx pg.Tx,
|
||||
identification TrackerIdentification,
|
||||
domains []string,
|
||||
) *gid.GID {
|
||||
var party coredata.CommonThirdParty
|
||||
if err := party.LoadByName(ctx, tx, identification.ThirdPartyName); err == nil {
|
||||
return &party.ID
|
||||
}
|
||||
|
||||
partySlug := slug.Make(identification.ThirdPartyName)
|
||||
if partySlug == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := party.LoadBySlug(ctx, tx, partySlug); err == nil {
|
||||
return &party.ID
|
||||
}
|
||||
|
||||
category := coredata.ThirdPartyCategoryOther
|
||||
if parsed := parseThirdPartyCategory(identification.Category); parsed != "" {
|
||||
category = parsed
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
party = coredata.CommonThirdParty{
|
||||
ID: gid.New(gid.NilTenant, coredata.CommonThirdPartyEntityType),
|
||||
Name: identification.ThirdPartyName,
|
||||
Slug: partySlug,
|
||||
Category: category,
|
||||
Certifications: []string{},
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
|
||||
if err := party.Insert(ctx, tx); err != nil {
|
||||
h.logger.WarnCtx(
|
||||
ctx,
|
||||
"cannot create common third party from agent",
|
||||
log.Error(err),
|
||||
log.String("name", identification.ThirdPartyName),
|
||||
)
|
||||
return nil
|
||||
}
|
||||
|
||||
for _, domain := range domains {
|
||||
domainRecord := coredata.CommonThirdPartyDomain{
|
||||
ID: gid.New(gid.NilTenant, coredata.CommonThirdPartyDomainEntityType),
|
||||
CommonThirdPartyID: party.ID,
|
||||
Domain: domain,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
|
||||
if _, err := domainRecord.Upsert(ctx, tx); err != nil {
|
||||
h.logger.WarnCtx(
|
||||
ctx,
|
||||
"cannot create common third party domain from agent",
|
||||
log.Error(err),
|
||||
log.String("domain", domain),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
h.logger.InfoCtx(
|
||||
ctx,
|
||||
"created common third party from agent identification",
|
||||
log.String("name", identification.ThirdPartyName),
|
||||
log.String("category", string(category)),
|
||||
)
|
||||
|
||||
return &party.ID
|
||||
}
|
||||
|
||||
func (h *trackerMappingHandler) createUnmatchedPattern(
|
||||
ctx context.Context,
|
||||
tx pg.Tx,
|
||||
@@ -242,3 +503,77 @@ func (h *trackerMappingHandler) resolveThirdParty(
|
||||
|
||||
return &t.ID, nil
|
||||
}
|
||||
|
||||
func buildAgentPrompt(tp coredata.TrackerPattern, domains []string) string {
|
||||
var b strings.Builder
|
||||
|
||||
fmt.Fprintf(&b, "Identify the following tracker:\n\n")
|
||||
fmt.Fprintf(&b, "- Pattern: %s\n", tp.Pattern)
|
||||
fmt.Fprintf(&b, "- Type: %s\n", tp.TrackerType)
|
||||
fmt.Fprintf(&b, "- Match type: %s\n", tp.MatchType)
|
||||
|
||||
if tp.MaxAgeSeconds != nil {
|
||||
fmt.Fprintf(&b, "- Max age: %d seconds\n", *tp.MaxAgeSeconds)
|
||||
} else {
|
||||
fmt.Fprintf(&b, "- Max age: session\n")
|
||||
}
|
||||
|
||||
if len(domains) > 0 {
|
||||
fmt.Fprintf(&b, "- Observed on domains: %s\n", strings.Join(domains, ", "))
|
||||
}
|
||||
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func parseThirdPartyCategory(s string) coredata.ThirdPartyCategory {
|
||||
switch s {
|
||||
case "ANALYTICS":
|
||||
return coredata.ThirdPartyCategoryAnalytics
|
||||
case "ADVERTISING":
|
||||
return coredata.ThirdPartyCategoryMarketing
|
||||
case "CLOUD_MONITORING":
|
||||
return coredata.ThirdPartyCategoryCloudMonitoring
|
||||
case "CLOUD_PROVIDER":
|
||||
return coredata.ThirdPartyCategoryCloudProvider
|
||||
case "COLLABORATION":
|
||||
return coredata.ThirdPartyCategoryCollaboration
|
||||
case "CUSTOMER_SUPPORT":
|
||||
return coredata.ThirdPartyCategoryCustomerSupport
|
||||
case "DATA_STORAGE_AND_PROCESSING":
|
||||
return coredata.ThirdPartyCategoryDataStorageAndProcessing
|
||||
case "DOCUMENT_MANAGEMENT":
|
||||
return coredata.ThirdPartyCategoryDocumentManagement
|
||||
case "EMPLOYEE_MANAGEMENT":
|
||||
return coredata.ThirdPartyCategoryEmployeeManagement
|
||||
case "ENGINEERING":
|
||||
return coredata.ThirdPartyCategoryEngineering
|
||||
case "FINANCE":
|
||||
return coredata.ThirdPartyCategoryFinance
|
||||
case "IDENTITY_PROVIDER":
|
||||
return coredata.ThirdPartyCategoryIdentityProvider
|
||||
case "IT":
|
||||
return coredata.ThirdPartyCategoryIT
|
||||
case "MARKETING":
|
||||
return coredata.ThirdPartyCategoryMarketing
|
||||
case "OFFICE_OPERATIONS":
|
||||
return coredata.ThirdPartyCategoryOfficeOperations
|
||||
case "OTHER":
|
||||
return coredata.ThirdPartyCategoryOther
|
||||
case "PASSWORD_MANAGEMENT":
|
||||
return coredata.ThirdPartyCategoryPasswordManagement
|
||||
case "PRODUCT_AND_DESIGN":
|
||||
return coredata.ThirdPartyCategoryProductAndDesign
|
||||
case "PROFESSIONAL_SERVICES":
|
||||
return coredata.ThirdPartyCategoryProfessionalServices
|
||||
case "RECRUITING":
|
||||
return coredata.ThirdPartyCategoryRecruiting
|
||||
case "SALES":
|
||||
return coredata.ThirdPartyCategorySales
|
||||
case "SECURITY":
|
||||
return coredata.ThirdPartyCategorySecurity
|
||||
case "VERSION_CONTROL":
|
||||
return coredata.ThirdPartyCategoryVersionControl
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user