Use contextual descriptors for entity reconciliation

internal/modules/dnd/normalize/locationregistry/test_helpers_test.go

@@ -3,11 +3,14 @@ package locationregistry
 import (
 	"context"
 	"encoding/json"
+	"strconv"
+	"strings"
 	"testing"
 
 	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
 	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
 	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
+	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared/entityreconcile"
 )
 
 type recordingLocationNormalizerClient struct {
@@ -25,10 +28,51 @@ func (c *recordingLocationNormalizerClient) CompleteStructured(_ context.Context
 	if response == "" {
 		response = `{"duplicate_groups":[]}`
 	}
-	if err := json.Unmarshal([]byte(response), output); err != nil {
+	content, err := contextualProposalResponse(response, request.Inputs["candidates"].Content)
+	if err != nil {
 		return contracts.StructuredCompletionResponse{}, err
 	}
-	return contracts.StructuredCompletionResponse{Content: json.RawMessage(response)}, nil
+	if err := json.Unmarshal(content, output); err != nil {
+		return contracts.StructuredCompletionResponse{}, err
+	}
+	return contracts.StructuredCompletionResponse{Content: content}, nil
+}
+
+func contextualProposalResponse(response string, candidateContent []byte) ([]byte, error) {
+	if !strings.Contains(response, "candidate-") {
+		return []byte(response), nil
+	}
+	var selection struct {
+		DuplicateGroups []struct {
+			Members   []string `json:"members"`
+			Canonical string   `json:"canonical"`
+		} `json:"duplicate_groups"`
+	}
+	if err := json.Unmarshal([]byte(response), &selection); err != nil {
+		return nil, err
+	}
+	var candidates struct {
+		Candidates []entityreconcile.Selector `json:"candidates"`
+	}
+	if err := json.Unmarshal(candidateContent, &candidates); err != nil {
+		return nil, err
+	}
+	selector := func(key string) entityreconcile.Selector {
+		index, err := strconv.Atoi(strings.TrimPrefix(key, "candidate-"))
+		if err != nil || index < 1 || index > len(candidates.Candidates) {
+			return entityreconcile.Selector{Name: key, SourceRefs: []entityreconcile.SourceRange{}}
+		}
+		return candidates.Candidates[index-1].Clone()
+	}
+	proposal := entityreconcile.ProposalResponse{DuplicateGroups: make([]entityreconcile.DuplicateGroup, len(selection.DuplicateGroups))}
+	for index, group := range selection.DuplicateGroups {
+		members := make([]entityreconcile.Selector, len(group.Members))
+		for memberIndex, key := range group.Members {
+			members[memberIndex] = selector(key)
+		}
+		proposal.DuplicateGroups[index] = entityreconcile.DuplicateGroup{Members: members, Canonical: selector(group.Canonical)}
+	}
+	return json.Marshal(proposal)
 }
 
 func newNormalizer(t *testing.T, client contracts.StructuredLLMClient) *Normalizer {