286 lines
15 KiB
Go
286 lines
15 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"
|
|
locationregistry "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/locations/registry"
|
|
)
|
|
|
|
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{
|
|
{Name: second.Name, RegistryRefs: registryRefs(second), Kind: "mentioned", SourceRefs: occurrenceRefs(30, 30)},
|
|
{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)},
|
|
{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "recalled", SourceRefs: occurrenceRefs(10, 10)},
|
|
{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "planned", SourceRefs: occurrenceRefs(10, 10)},
|
|
{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "visited", SourceRefs: append(occurrenceRefs(10, 10), occurrenceRefs(10, 10)...)},
|
|
{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "visited", SourceRefs: occurrenceRefs(20, 20)},
|
|
{Name: first.Name, RegistryRefs: registryRefs(first), 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 TestExtractResolvesContextualSelectorsAndUsesCurrentTranscriptEvidenceOnly(t *testing.T) {
|
|
locations := locationRegistry(t, "The Tavern", "The Tavern")
|
|
first, second := locations.Locations[0], locations.Locations[1]
|
|
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{{
|
|
Name: second.Name, RegistryRefs: registryRefs(second), 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), `"registry_refs":[{"start_unit_id":20,"end_unit_id":20}]`) {
|
|
t.Fatalf("location prompt input = %#v", input)
|
|
}
|
|
for _, forbidden := range []string{"source_refs", "source_id", first.ID, second.ID} {
|
|
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 TestExtractRejectsUnknownMalformedOrMismatchedSelectorsAtomically(t *testing.T) {
|
|
locations := locationRegistry(t, "The Mill")
|
|
known := locations.Locations[0]
|
|
references := registryReferences(t, locations)
|
|
req := extractionRequest()
|
|
req.References = references
|
|
for _, occurrences := range [][]occurrenceResponse{
|
|
{{Name: "Unknown", RegistryRefs: []locationregistry.RegistryRef{}, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}},
|
|
{{Name: known.Name, RegistryRefs: []locationregistry.RegistryRef{{StartUnitID: 10, EndUnitID: 0}}, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}},
|
|
{{Name: known.Name, RegistryRefs: []locationregistry.RegistryRef{{StartUnitID: 10, EndUnitID: 10}}, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}},
|
|
{{Name: known.Name, RegistryRefs: []locationregistry.RegistryRef{}, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}, {Name: "Unknown", RegistryRefs: []locationregistry.RegistryRef{}, Kind: "mentioned", SourceRefs: occurrenceRefs(20, 20)}},
|
|
} {
|
|
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: occurrences}}
|
|
if result, err := newExtractor(t, client, references).Extract(context.Background(), req); err == nil || result.Value.Occurrences != nil || !strings.Contains(err.Error(), "location selector") {
|
|
t.Fatalf("Extract() = %#v, %v", result, err)
|
|
}
|
|
}
|
|
|
|
sharedName := "The Tavern"
|
|
firstRefs := []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 10, EndUnitID: 10}, {SourceID: req.Source.ID, StartUnitID: 20, EndUnitID: 20}}
|
|
secondRefs := []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 30, EndUnitID: 30}}
|
|
duplicateLocations := dnd.LocationRegistry{Locations: []dnd.Location{
|
|
{ID: identity.DeriveID(sharedName, firstRefs), Name: sharedName, SourceRefs: firstRefs},
|
|
{ID: identity.DeriveID(sharedName, secondRefs), Name: sharedName, SourceRefs: secondRefs},
|
|
}}
|
|
duplicateReferences := registryReferences(t, duplicateLocations)
|
|
for _, selector := range []struct {
|
|
name string
|
|
refs []locationregistry.RegistryRef
|
|
}{
|
|
{name: sharedName, refs: []locationregistry.RegistryRef{}},
|
|
{name: sharedName, refs: []locationregistry.RegistryRef{{StartUnitID: 10, EndUnitID: 10}}},
|
|
{name: sharedName, refs: []locationregistry.RegistryRef{{StartUnitID: 20, EndUnitID: 20}, {StartUnitID: 10, EndUnitID: 10}}},
|
|
} {
|
|
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{{Name: selector.name, RegistryRefs: selector.refs, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}}}}
|
|
duplicateRequest := extractionRequest()
|
|
duplicateRequest.References = duplicateReferences
|
|
if _, err := newExtractor(t, client, duplicateReferences).Extract(context.Background(), duplicateRequest); err == nil || !strings.Contains(err.Error(), "location selector") {
|
|
t.Fatalf("Extract(%#v) error = %v", selector, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
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{{
|
|
Name: location.Name, RegistryRefs: []locationregistry.RegistryRef{}, 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" || got[3].Value != locationRegistryIdentityDigest(t, references) {
|
|
t.Fatalf("fingerprints = %#v", got)
|
|
}
|
|
}
|
|
|
|
func locationRegistryIdentityDigest(t *testing.T, references contracts.ReferenceSet) string {
|
|
t.Helper()
|
|
resolver, err := locationregistry.NewResolver(references)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return resolver.Seeded().IdentityDigest()
|
|
}
|
|
|
|
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: sourceDocument().ID, StartUnitID: (index + 1) * 10, EndUnitID: (index + 1) * 10}}
|
|
locations[index] = dnd.Location{ID: identity.DeriveID(name, refs), Name: name, SourceRefs: refs}
|
|
}
|
|
return dnd.LocationRegistry{Locations: locations}
|
|
}
|
|
|
|
func registryRefs(location dnd.Location) []locationregistry.RegistryRef {
|
|
refs := make([]locationregistry.RegistryRef, len(location.SourceRefs))
|
|
for index, ref := range location.SourceRefs {
|
|
refs[index] = locationregistry.RegistryRef{StartUnitID: ref.StartUnitID, EndUnitID: ref.EndUnitID}
|
|
}
|
|
return refs
|
|
}
|
|
|
|
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"}}},
|
|
}}}
|
|
}
|