Build tracker identification prompt categories dynamically

Generate the category list from coredata.ThirdPartyCategories() at
runtime instead of hardcoding it in the prompt text. Type the
TrackerIdentification.Category field as coredata.ThirdPartyCategory so
JSON unmarshaling validates values automatically.

Also documents the .txt.tmpl template file naming convention.

Signed-off-by: Émile Ré <emile@probo.com>

pkg/cookiebanner/prompts/tracker_identification.txt.tmpl

@@ -45,6 +45,6 @@ Return a structured JSON response with:
 
 7. If you truly cannot identify the tracker, set third_party_name to an empty string and confidence below 0.3.
 
-8. For the category field, use one of: ANALYTICS, ADVERTISING, CLOUD_MONITORING, CLOUD_PROVIDER, COLLABORATION, CUSTOMER_SUPPORT, DATA_STORAGE_AND_PROCESSING, DOCUMENT_MANAGEMENT, EMPLOYEE_MANAGEMENT, ENGINEERING, FINANCE, IDENTITY_PROVIDER, IT, MARKETING, OFFICE_OPERATIONS, OTHER, PASSWORD_MANAGEMENT, PRODUCT_AND_DESIGN, PROFESSIONAL_SERVICES, RECRUITING, SALES, SECURITY, VERSION_CONTROL.
-   Most cookies fall under ANALYTICS, ADVERTISING, or MARKETING.
+8. For the category field, use one of: {{.Categories}}.
+   Most cookies fall under ANALYTICS or MARKETING.
 </instructions>

pkg/cookiebanner/tracker_identification.go

@@ -14,13 +14,13 @@
 
 package cookiebanner
 
+import "go.probo.inc/probo/pkg/coredata"
+
 // TrackerIdentification is the structured output the tracker-mapping
-// agent returns. The Category field uses the same values as the
-// third_party_category PostgreSQL enum so auto-created CommonThirdParty
-// rows get a valid category without mapping.
+// agent returns.
 type TrackerIdentification struct {
-	ThirdPartyName string  `json:"third_party_name" jsonschema:"Name of the company or service that sets this tracker (e.g. 'Google Analytics', 'Meta Pixel'). Empty string if truly unknown."`
-	Category       string  `json:"category" jsonschema:"Third party category. One of: ANALYTICS, ADVERTISING, CLOUD_MONITORING, CLOUD_PROVIDER, COLLABORATION, CUSTOMER_SUPPORT, DATA_STORAGE_AND_PROCESSING, DOCUMENT_MANAGEMENT, EMPLOYEE_MANAGEMENT, ENGINEERING, FINANCE, IDENTITY_PROVIDER, IT, MARKETING, OFFICE_OPERATIONS, OTHER, PASSWORD_MANAGEMENT, PRODUCT_AND_DESIGN, PROFESSIONAL_SERVICES, RECRUITING, SALES, SECURITY, VERSION_CONTROL"`
-	Description    string  `json:"description" jsonschema:"What this tracker does in one sentence"`
-	Confidence     float64 `json:"confidence" jsonschema:"Confidence level from 0.0 to 1.0. Set below 0.5 if unsure."`
+	ThirdPartyName string                      `json:"third_party_name" jsonschema:"Name of the company or service that sets this tracker (e.g. 'Google Analytics', 'Meta Pixel'). Empty string if truly unknown."`
+	Category       coredata.ThirdPartyCategory `json:"category" jsonschema:"Third party category"`
+	Description    string                      `json:"description" jsonschema:"What this tracker does in one sentence"`
+	Confidence     float64                     `json:"confidence" jsonschema:"Confidence level from 0.0 to 1.0. Set below 0.5 if unsure."`
 }

pkg/cookiebanner/tracker_mapping_worker.go

@@ -40,7 +40,7 @@ const (
 	agentMaxPatternConfidence = 0.8
 )
 
-//go:embed prompts/tracker_identification.txt
+//go:embed prompts/tracker_identification.txt.tmpl
 var trackerIdentificationPrompt string
 
 type trackerMappingHandler struct {
@@ -100,7 +100,7 @@ func buildTrackerMappingAgent(
 	return agent.New(
 		"tracker-mapping",
 		cfg.LLMClient,
-		agent.WithInstructions(trackerIdentificationPrompt),
+		agent.WithInstructionsFunc(trackerMappingInstructions),
 		agent.WithModel(cfg.Model),
 		agent.WithTools(tools...),
 		agent.WithOutputType(outputType),
@@ -109,6 +109,21 @@ func buildTrackerMappingAgent(
 	)
 }
 
+func trackerMappingInstructions(_ context.Context, _ *agent.Agent) string {
+	categories := coredata.ThirdPartyCategories()
+	parts := make([]string, len(categories))
+	for i, c := range categories {
+		parts[i] = string(c)
+	}
+
+	return strings.Replace(
+		trackerIdentificationPrompt,
+		"{{.Categories}}",
+		strings.Join(parts, ", "),
+		1,
+	)
+}
+
 func (h *trackerMappingHandler) Claim(ctx context.Context) (coredata.TrackerPattern, error) {
 	var tp coredata.TrackerPattern
 
@@ -373,10 +388,7 @@ func (h *trackerMappingHandler) resolveOrCreateCommonThirdParty(
 		return &party.ID, nil
 	}
 
-	category := coredata.ThirdPartyCategoryOther
-	if parsed := parseThirdPartyCategory(identification.Category); parsed != "" {
-		category = parsed
-	}
+	category := identification.Category
 
 	now := time.Now()
 	party = coredata.CommonThirdParty{
@@ -411,7 +423,7 @@ func (h *trackerMappingHandler) resolveOrCreateCommonThirdParty(
 		ctx,
 		"created common third party from agent identification",
 		log.String("name", identification.ThirdPartyName),
-		log.String("category", string(category)),
+		log.String("category", category.String()),
 	)
 
 	return &party.ID, nil
@@ -473,75 +485,26 @@ func (h *trackerMappingHandler) resolveThirdParty(
 }
 
 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)
-
+	maxAge := "session"
 	if tp.MaxAgeSeconds != nil {
-		fmt.Fprintf(&b, "- Max age: %d seconds\n", *tp.MaxAgeSeconds)
-	} else {
-		fmt.Fprintf(&b, "- Max age: session\n")
+		maxAge = fmt.Sprintf("%d seconds", *tp.MaxAgeSeconds)
 	}
 
+	prompt := 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",
+		tp.Pattern,
+		tp.TrackerType,
+		tp.MatchType,
+		maxAge,
+	)
+
 	if len(domains) > 0 {
-		fmt.Fprintf(&b, "- Observed on domains: %s\n", strings.Join(domains, ", "))
+		prompt += fmt.Sprintf("<observed_domains> %s </observed_domains>\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 ""
-	}
+	return prompt
 }

pkg/coredata/common_third_party.go

@@ -214,6 +214,7 @@ LIMIT 1;
 		if errors.Is(err, pgx.ErrNoRows) {
 			return ErrResourceNotFound
 		}
+
 		return fmt.Errorf("cannot collect common third party by slug: %w", err)
 	}