105 lines
3.9 KiB
Go
105 lines
3.9 KiB
Go
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
|
|
}
|