rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add common catalog query layer and enricher service

Browse Source 
Introduce an API-style data layer for the global common tracker pattern
and common third party catalogs: typed filters, order fields, CursorKey,
cursor-paginated Load and CountAll, plus by-id enrichment re-queue and a
scoped reset/remap helper for a banner's tracker patterns. These reuse
the same page.Cursor/filter/order types the GraphQL API consumes, so a
future proboctl API can back them unchanged.

Extract the common-pattern enrichment logic out of the worker into a
CommonPatternEnricher service so it can run either from the background
queue or synchronously over a known set of ids; the worker becomes a
thin poller that delegates to it.

Extract the LLM client and tracker-agents config wiring into
pkg/agentsbuild so probod and other binaries build agents identically;
probod now delegates to it.

Signed-off-by: Émile Ré <emile@probo.com>
This commit is contained in:
Émile Ré
2026-06-03 13:37:36 +02:00
parent aebb2a1ed0
commit 662c0ae428
13 changed files with 1452 additions and 335 deletions
Download Patch File Download Diff File Expand all files Collapse all files

222
pkg/cookiebanner/common_pattern_enrichment_worker.go
View File

@@ -18,28 +18,25 @@ import (
"context"
"errors"
"fmt"
"strings"
"time"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.gearno.de/kit/worker"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/llm"
"go.probo.inc/probo/pkg/thirdparty"
)
const defaultEnrichmentStaleAfter = 10 * time.Minute
// commonPatternEnrichmentHandler is the queue poller for common tracker
// pattern enrichment. It owns only the claim/dequeue and stale-recovery
// mechanics; the enrichment work itself lives in CommonPatternEnricher so
// it can also run synchronously from operator tooling.
type commonPatternEnrichmentHandler struct {
pg *pg.Client
logger *log.Logger
enrichmentAgent *agent.Agent
mappingAgent *agent.Agent
staleAfter time.Duration
agentTimeout time.Duration
pg *pg.Client
logger *log.Logger
enricher *CommonPatternEnricher
staleAfter time.Duration
}
// NewCommonPatternEnrichmentWorker builds the worker that fills
@@ -60,21 +57,11 @@ func NewCommonPatternEnrichmentWorker(
staleAfter = defaultEnrichmentStaleAfter
}
agentTimeout := cfg.AgentTimeout
if agentTimeout <= 0 {
agentTimeout = defaultAgentTimeout
}
h := &commonPatternEnrichmentHandler{
pg: pgClient,
logger: logger,
staleAfter: staleAfter,
agentTimeout: agentTimeout,
}
if cfg.LLMClient != nil {
h.enrichmentAgent = buildCommonPatternEnrichmentAgent(cfg, pgClient, logger)
h.mappingAgent = buildTrackerMappingAgent(cfg, pgClient, logger)
pg: pgClient,
logger: logger,
enricher: NewCommonPatternEnricher(pgClient, logger, cfg),
staleAfter: staleAfter,
}
return worker.New(
@@ -109,86 +96,11 @@ func (h *commonPatternEnrichmentHandler) Claim(ctx context.Context) (coredata.Co
}
func (h *commonPatternEnrichmentHandler) Process(ctx context.Context, cp coredata.CommonTrackerPattern) error {
if h.enrichmentAgent == nil {
if !h.enricher.Enabled() {
return nil
}
thirdPartyName, err := h.loadThirdPartyName(ctx, cp)
if err != nil {
return err
}
// Map before enriching: an unlinked pattern is run through the
// mapping agent first so a confident vendor both seeds the enrichment
// prompt and gets linked. Attribution stays the mapping pipeline's
// job; the enricher only reuses it. An already-linked pattern skips
// this entirely.
var attribution *TrackerMappingAgentResult
if cp.CommonThirdPartyID == nil {
attribution, err = h.identifyThirdParty(ctx, cp)
if err != nil {
return err
}
if attribution != nil {
thirdPartyName = attribution.ThirdPartyName
}
}
description, err := h.research(ctx, cp, thirdPartyName)
if err != nil {
return fmt.Errorf("cannot research tracker description: %w", err)
}
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
// Resolve or create the catalog vendor only for an unlinked
// pattern; the mapping pipeline owns creation, so we reuse its
// name+slug dedup and never create a duplicate or override an
// existing link.
var thirdPartyID *gid.GID
if attribution != nil && cp.CommonThirdPartyID == nil {
thirdPartyID, err = thirdparty.ResolveOrCreateCommonThirdParty(ctx, tx, h.logger, attribution.ThirdPartyName, attribution.Category)
if err != nil {
return fmt.Errorf("cannot resolve or create common third party: %w", err)
}
}
// 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 backfilled int64
if description != "" {
var patterns coredata.TrackerPatterns
backfilled, err = patterns.BackfillDescriptionByCommonTrackerPatternID(ctx, tx, cp.ID, description)
if err != nil {
return err
}
}
h.logger.InfoCtx(
ctx,
"enriched common tracker pattern",
log.String("common_tracker_pattern_id", cp.ID.String()),
log.String("pattern", cp.Pattern),
log.Bool("described", description != ""),
log.Bool("third_party_linked", thirdPartyID != nil),
log.Int64("backfilled_tracker_patterns", backfilled),
)
return nil
},
)
return h.enricher.EnrichPattern(ctx, cp)
}
func (h *commonPatternEnrichmentHandler) RecoverStale(ctx context.Context) error {
@@ -203,109 +115,3 @@ func (h *commonPatternEnrichmentHandler) RecoverStale(ctx context.Context) error
},
)
}
func (h *commonPatternEnrichmentHandler) loadThirdPartyName(
ctx context.Context,
cp coredata.CommonTrackerPattern,
) (string, error) {
if cp.CommonThirdPartyID == nil {
return "", nil
}
var name string
if err := h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
var party coredata.CommonThirdParty
if err := party.LoadByID(ctx, conn, *cp.CommonThirdPartyID); err != nil {
return err
}
name = party.Name
return nil
},
); err != nil {
return "", fmt.Errorf("cannot load common third party for enrichment: %w", err)
}
return name, nil
}
func (h *commonPatternEnrichmentHandler) research(
ctx context.Context,
cp coredata.CommonTrackerPattern,
thirdPartyName string,
) (string, error) {
prompt := buildEnrichmentPrompt(cp, thirdPartyName)
agentCtx, cancel := context.WithTimeout(ctx, h.agentTimeout)
defer cancel()
result, err := agent.RunTyped[CommonPatternEnrichmentResult](
agentCtx,
h.enrichmentAgent,
[]llm.Message{
{
Role: llm.RoleUser,
Parts: []llm.Part{llm.TextPart{Text: prompt}},
},
},
)
if err != nil {
return "", fmt.Errorf("enrichment agent run failed: %w", err)
}
return strings.TrimSpace(result.Output.Description), nil
}
// identifyThirdParty reuses the tracker-mapping agent to attribute a
// vendor to an unlinked catalog pattern. It performs no DB writes: it
// returns the confident attribution (name, category, confidence) or nil
// when the agent is unsure, leaving the caller to resolve or create the
// catalog row. 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,
) (*TrackerMappingAgentResult, 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
}
out := result.Output
out.ThirdPartyName = strings.TrimSpace(out.ThirdPartyName)
if out.ThirdPartyName == "" || out.ThirdPartyConfidence < agentThirdPartyConfidenceThreshold {
return nil, nil
}
return &out, nil
}