244 lines
12 KiB
Go
244 lines
12 KiB
Go
package locationoccurrences
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
|
|
locationcodec "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/codec/locationregistry"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/locations/identity"
|
|
)
|
|
|
|
func TestExtractMapsKindsOrdersOccurrencesAndPreservesIndependentFacts(t *testing.T) {
|
|
locations := locationRegistry(t, "The Tavern", "The Tavern")
|
|
first, second := locations.Locations[0], locations.Locations[1]
|
|
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{
|
|
{LocationID: second.ID, Name: second.Name, Kind: "mentioned", SourceRefs: occurrenceRefs(30, 30)},
|
|
{LocationID: first.ID, Name: first.Name, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)},
|
|
{LocationID: first.ID, Name: first.Name, Kind: "recalled", SourceRefs: occurrenceRefs(10, 10)},
|
|
{LocationID: first.ID, Name: first.Name, Kind: "planned", SourceRefs: occurrenceRefs(10, 10)},
|
|
{LocationID: first.ID, Name: first.Name, Kind: "visited", SourceRefs: append(occurrenceRefs(10, 10), occurrenceRefs(10, 10)...)},
|
|
{LocationID: first.ID, Name: first.Name, Kind: "visited", SourceRefs: occurrenceRefs(20, 20)},
|
|
{LocationID: first.ID, Name: first.Name, Kind: "visited", SourceRefs: occurrenceRefs(10, 10)},
|
|
}}}
|
|
references := registryReferences(t, locations)
|
|
req := extractionRequest()
|
|
req.References = references
|
|
result, err := newExtractor(t, client, references).Extract(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(result.Value.Occurrences) != 6 {
|
|
t.Fatalf("occurrences = %#v, want exact duplicate removed", result.Value.Occurrences)
|
|
}
|
|
got := result.Value.Occurrences
|
|
if kinds := []dnd.LocationOccurrenceKind{got[0].Kind, got[1].Kind, got[2].Kind, got[3].Kind}; !reflect.DeepEqual(kinds, []dnd.LocationOccurrenceKind{dnd.LocationOccurrenceKindVisited, dnd.LocationOccurrenceKindPlanned, dnd.LocationOccurrenceKindRecalled, dnd.LocationOccurrenceKindMentioned}) {
|
|
t.Fatalf("same-evidence kind order = %#v", kinds)
|
|
}
|
|
if got[4].Kind != dnd.LocationOccurrenceKindVisited || got[4].SourceRefs[0].StartUnitID != 20 || got[5].LocationID != second.ID || got[5].SourceRefs[0].StartUnitID != 30 {
|
|
t.Fatalf("occurrence order = %#v", got)
|
|
}
|
|
if !reflect.DeepEqual(got[0].SourceRefs, []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 10, EndUnitID: 10}}) {
|
|
t.Fatalf("canonical evidence = %#v", got[0].SourceRefs)
|
|
}
|
|
}
|
|
|
|
func TestExtractUsesIDsNamesAndCurrentTranscriptEvidenceOnly(t *testing.T) {
|
|
locations := locationRegistry(t, "The Tavern", "The Tavern")
|
|
first, second := locations.Locations[0], locations.Locations[1]
|
|
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{{
|
|
LocationID: second.ID, Name: second.Name, Kind: "visited", SourceRefs: occurrenceRefs(10, 10),
|
|
}}}}
|
|
references := registryReferences(t, locations)
|
|
req := extractionRequest()
|
|
req.References = references
|
|
result, err := newExtractor(t, client, references).Extract(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if occurrence := result.Value.Occurrences[0]; occurrence.LocationID != second.ID || occurrence.Name != second.Name || occurrence.SourceRefs[0].SourceID != req.Source.ID {
|
|
t.Fatalf("occurrence = %#v", occurrence)
|
|
}
|
|
input := client.requests[0].Inputs[LocationRegistryReferenceSlot]
|
|
if input.Name != LocationRegistryReferenceSlot || !strings.Contains(string(input.Content), first.ID) || !strings.Contains(string(input.Content), second.ID) {
|
|
t.Fatalf("location prompt input = %#v", input)
|
|
}
|
|
for _, forbidden := range []string{"source_refs", "source_id", "other-session"} {
|
|
if strings.Contains(string(input.Content), forbidden) {
|
|
t.Fatalf("location prompt leaked %q: %s", forbidden, input.Content)
|
|
}
|
|
}
|
|
if strings.Contains(string(client.requests[0].Inputs["transcript"].Content), "other-session") {
|
|
t.Fatal("transcript input contains registry evidence")
|
|
}
|
|
metadata, err := json.Marshal(newExtractor(t, &fakeOccurrencesLLMClient{}, references).ManifestMetadata())
|
|
if err != nil || strings.Contains(string(metadata), "other-session") || strings.Contains(string(metadata), first.ID) {
|
|
t.Fatalf("manifest metadata = %s, %v", metadata, err)
|
|
}
|
|
}
|
|
|
|
func TestExtractPreservesUnknownOrMismatchedGroundingForValidators(t *testing.T) {
|
|
locations := locationRegistry(t, "The Mill")
|
|
known := locations.Locations[0]
|
|
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{
|
|
{LocationID: "location:sha256:unknown", Name: "The Mill", Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)},
|
|
{LocationID: known.ID, Name: "A Different Mill", Kind: "mentioned", SourceRefs: occurrenceRefs(20, 20)},
|
|
}}}
|
|
references := registryReferences(t, locations)
|
|
req := extractionRequest()
|
|
req.References = references
|
|
result, err := newExtractor(t, client, references).Extract(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if result.Value.Occurrences[0].LocationID != "location:sha256:unknown" || result.Value.Occurrences[1].Name != "A Different Mill" {
|
|
t.Fatalf("extractor repaired validator-owned grounding errors: %#v", result.Value.Occurrences)
|
|
}
|
|
}
|
|
|
|
func TestExtractRequiresRegistryAndAcceptsEmptyRegistryWithNoOccurrences(t *testing.T) {
|
|
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{}}}
|
|
if _, err := newExtractor(t, client).Extract(context.Background(), extractionRequest()); err == nil || !strings.Contains(err.Error(), "location registry reference is required") {
|
|
t.Fatalf("Extract() error = %v", err)
|
|
}
|
|
if len(client.requests) != 0 {
|
|
t.Fatalf("LLM calls = %d", len(client.requests))
|
|
}
|
|
empty := dnd.LocationRegistry{Locations: []dnd.Location{}}
|
|
references := registryReferences(t, empty)
|
|
req := extractionRequest()
|
|
req.References = references
|
|
result, err := newExtractor(t, client, references).Extract(context.Background(), req)
|
|
if err != nil || result.Value.Occurrences == nil || len(result.Value.Occurrences) != 0 {
|
|
t.Fatalf("empty registry result = %#v, %v", result, err)
|
|
}
|
|
}
|
|
|
|
func TestExtractResolvesGeneratedRegistryAtOperationTimeAndDoesNotMutateResponse(t *testing.T) {
|
|
locations := locationRegistry(t, "The Mill")
|
|
location := locations.Locations[0]
|
|
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{{
|
|
LocationID: location.ID, Name: location.Name, Kind: "mentioned", SourceRefs: occurrenceRefs(30, 30),
|
|
}}}}
|
|
references := registryReferences(t, locations)
|
|
req := extractionRequest()
|
|
req.References = references
|
|
extractor := newExtractor(t, client)
|
|
result, err := extractor.Extract(context.Background(), req)
|
|
if err != nil || result.Value.Occurrences[0].Name != "The Mill" {
|
|
t.Fatalf("Extract() = %#v, %v", result, err)
|
|
}
|
|
if input := client.requests[0].Inputs[LocationRegistryReferenceSlot]; !strings.Contains(string(input.Content), location.ID) || input.OriginURI != "" {
|
|
t.Fatalf("generated registry prompt input = %#v", input)
|
|
}
|
|
if _, ok := extractor.ManifestMetadata()["location_registry_digest"]; ok {
|
|
t.Fatalf("operation registry leaked into static metadata: %#v", extractor.ManifestMetadata())
|
|
}
|
|
if client.response.Occurrences[0].SourceRefs[0].StartUnitID != 30 {
|
|
t.Fatalf("model response mutated: %#v", client.response)
|
|
}
|
|
}
|
|
|
|
func TestExtractorContractsMetadataAndFailures(t *testing.T) {
|
|
if _, err := New(nil, Options{}); err == nil || !strings.Contains(err.Error(), "LLM client") {
|
|
t.Fatalf("New(nil) error = %v", err)
|
|
}
|
|
if _, err := New(&fakeOccurrencesLLMClient{}, Options{}, contracts.ReferenceSet{}, contracts.ReferenceSet{}); err == nil || !strings.Contains(err.Error(), "at most one reference set") {
|
|
t.Fatalf("New() error = %v", err)
|
|
}
|
|
malformed := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{LocationRegistryReferenceSlot: {
|
|
Items: []contracts.ReferenceItem{{SlotName: LocationRegistryReferenceSlot, MediaType: "application/json", Content: []byte(`{"secret":"registry evidence"}`)}},
|
|
}}}
|
|
if _, err := New(&fakeOccurrencesLLMClient{}, Options{}, malformed); err == nil || !strings.Contains(err.Error(), "prepare location registry") || strings.Contains(err.Error(), "registry evidence") {
|
|
t.Fatalf("New() error = %v", err)
|
|
}
|
|
|
|
locations := locationRegistry(t, "The Mill")
|
|
references := registryReferences(t, locations)
|
|
req := extractionRequest()
|
|
req.References = references
|
|
extractor := newExtractor(t, &fakeOccurrencesLLMClient{}, references)
|
|
var nilExtractor *Extractor
|
|
for _, test := range []struct {
|
|
name string
|
|
extractor *Extractor
|
|
req contracts.TypedExtractionRequest
|
|
want string
|
|
}{
|
|
{"nil extractor", nilExtractor, req, "extractor"},
|
|
{"nil client", &Extractor{}, req, "LLM client"},
|
|
{"invalid request", extractor, mismatchedSourceInputRequest(req), "must match chunk"},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
if _, err := test.extractor.Extract(context.Background(), test.req); err == nil || !strings.Contains(err.Error(), test.want) {
|
|
t.Fatalf("Extract() error = %v", err)
|
|
}
|
|
})
|
|
}
|
|
if _, err := newExtractor(t, &fakeOccurrencesLLMClient{err: errors.New("provider unavailable")}, references).Extract(context.Background(), req); err == nil || !strings.Contains(err.Error(), "provider unavailable") {
|
|
t.Fatalf("provider error = %v", err)
|
|
}
|
|
|
|
spec := ModuleSpec()
|
|
if spec.Key != Key || spec.Stage != pipeline.StageExtract || spec.ExecutionClass != contracts.ExecutionClassLLMBacked || spec.ArtifactKind != dnd.LocationOccurrenceListKind {
|
|
t.Fatalf("ModuleSpec() = %#v", spec)
|
|
}
|
|
var slot contracts.ReferenceSlot
|
|
for _, candidate := range spec.ReferenceSlots {
|
|
if candidate.Name == LocationRegistryReferenceSlot {
|
|
slot = candidate
|
|
}
|
|
}
|
|
if !slot.Required || !reflect.DeepEqual(slot.AcceptedArtifactKinds, []contracts.ArtifactKind{dnd.LocationRegistryKind}) || slot.MaxBytes != LocationRegistryMaxBytes {
|
|
t.Fatalf("location registry slot = %#v", slot)
|
|
}
|
|
registry := pipeline.NewExtractorRegistry()
|
|
if err := Register(registry); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, ok := registry.Spec(Key); !ok {
|
|
t.Fatalf("registration missing %q", Key)
|
|
}
|
|
if _, err := DecodeOptions(map[string]any{"unexpected": true}); err == nil {
|
|
t.Fatal("DecodeOptions() accepted unknown options")
|
|
}
|
|
metadata := newExtractor(t, &fakeOccurrencesLLMClient{}, references).ManifestMetadata()
|
|
for _, key := range []string{"prompt_sha256", "response_schema_sha256", "location_registry_digest"} {
|
|
if value, ok := metadata[key].(string); !ok || !strings.HasPrefix(value, "sha256:") {
|
|
t.Fatalf("metadata[%q] = %#v", key, metadata[key])
|
|
}
|
|
}
|
|
if got := newExtractor(t, &fakeOccurrencesLLMClient{}, references).CheckpointFingerprints(); len(got) != 4 || got[3].Name != "location_registry" {
|
|
t.Fatalf("fingerprints = %#v", got)
|
|
}
|
|
}
|
|
|
|
func locationRegistry(t *testing.T, names ...string) dnd.LocationRegistry {
|
|
t.Helper()
|
|
locations := make([]dnd.Location, len(names))
|
|
for index, name := range names {
|
|
refs := []source.SourceRef{{SourceID: "other-session", StartUnitID: index + 1, EndUnitID: index + 1}}
|
|
locations[index] = dnd.Location{ID: identity.DeriveID(name, refs), Name: name, SourceRefs: refs}
|
|
}
|
|
return dnd.LocationRegistry{Locations: locations}
|
|
}
|
|
|
|
func registryReferences(t *testing.T, locations dnd.LocationRegistry) contracts.ReferenceSet {
|
|
t.Helper()
|
|
content, err := locationcodec.New().Encode(locations)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{LocationRegistryReferenceSlot: {
|
|
Slot: contracts.ReferenceSlot{Name: LocationRegistryReferenceSlot},
|
|
Items: []contracts.ReferenceItem{{SlotName: LocationRegistryReferenceSlot, MediaType: locationcodec.MediaType, Content: content, Origin: contracts.ReferenceOrigin{Type: "generated"}}},
|
|
}}}
|
|
}
|