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 { response string err error requests []contracts.StructuredCompletionRequest } func (c *recordingLocationNormalizerClient) CompleteStructured(_ context.Context, request contracts.StructuredCompletionRequest, output any) (contracts.StructuredCompletionResponse, error) { c.requests = append(c.requests, request) if c.err != nil { return contracts.StructuredCompletionResponse{}, c.err } response := c.response if response == "" { response = `{"duplicate_groups":[]}` } content, err := contextualProposalResponse(response, request.Inputs["candidates"].Content) if err != nil { return contracts.StructuredCompletionResponse{}, err } 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 { t.Helper() normalizer, err := New(client, Options{}) if err != nil { t.Fatalf("New() error = %v", err) } return normalizer } func normalizeRequest(value dnd.LocationRegistry) contracts.TypedNormalizeRequest[dnd.LocationRegistry] { return contracts.TypedNormalizeRequest[dnd.LocationRegistry]{MergeOutput: contracts.MergeArtifact[dnd.LocationRegistry]{Value: value}} } func normalizeRequestWithSource(value dnd.LocationRegistry, doc *source.SourceDocument) contracts.TypedNormalizeRequest[dnd.LocationRegistry] { request := normalizeRequest(value) request.Source = doc return request } func semanticDocument() *source.SourceDocument { return &source.SourceDocument{ID: "location-session", Units: []source.SourceUnit{{ID: 10, Kind: "speech", Text: "The old mill is the Greencloak's refuge."}, {ID: 20, Kind: "speech", Text: "The mill stands on the northern road."}, {ID: 30, Kind: "speech", Text: "The tavern is beside the mill."}, {ID: 40, Kind: "speech", Text: "The mill's cellar is flooded."}}} } func hasWarning(warnings []contracts.Warning, reason string) bool { for _, warning := range warnings { if warning.ReasonCode == reason { return true } } return false }