Add D&D location occurrence extractor
This commit is contained in:
@@ -0,0 +1,6 @@
|
||||
package locationoccurrences
|
||||
|
||||
import "embed"
|
||||
|
||||
//go:embed assets/prompts/*.yaml assets/prompts/*.md assets/schemas/*.json
|
||||
var embeddedAssets embed.FS
|
||||
@@ -0,0 +1,47 @@
|
||||
id: dnd.location_occurrences
|
||||
version: "v1"
|
||||
default_profile: dnd-extraction
|
||||
inputs:
|
||||
- name: transcript
|
||||
required: true
|
||||
content_type: application/json
|
||||
- name: players
|
||||
required: false
|
||||
content_type: text/plain
|
||||
- name: party
|
||||
required: false
|
||||
content_type: text/plain
|
||||
- name: glossary
|
||||
required: false
|
||||
content_type: text/plain
|
||||
- name: locations
|
||||
required: true
|
||||
content_type: application/json
|
||||
messages:
|
||||
- role: system
|
||||
content_file: ./sharedassets/common-dnd-system.md
|
||||
- role: user
|
||||
content_file: ./sharedassets/common-dnd-identity.md
|
||||
- role: user
|
||||
content_file: ./sharedassets/common-dnd-references.md
|
||||
cache_control:
|
||||
type: ephemeral
|
||||
- role: user
|
||||
content_file: ./sharedassets/common-dnd-transcript.md
|
||||
cache_control:
|
||||
type: ephemeral
|
||||
- role: user
|
||||
content_file: ./sharedassets/common-dnd-extraction-evidence.md
|
||||
- role: user
|
||||
content_file: ./locations.md
|
||||
- role: user
|
||||
content_file: ./task.md
|
||||
- role: user
|
||||
content_file: ./instructions.md
|
||||
cache_control:
|
||||
type: ephemeral
|
||||
output:
|
||||
format: json
|
||||
validation_mode: json_schema
|
||||
schema_path: dnd_location_occurrences_llm.v1.json
|
||||
repair_attempts: 0
|
||||
@@ -0,0 +1,8 @@
|
||||
Return the occurrences array even when no occurrence is established. Every
|
||||
record must contain location_id, name, kind, and source_refs. Copy location_id
|
||||
and name from one supplied registry record, and cite only narrow transcript
|
||||
ranges that support both that location and its classified occurrence.
|
||||
|
||||
Do not summarize location descriptions, infer a missing registry record, or
|
||||
use registry context as evidence. Omit source_id; Notarius assigns the current
|
||||
transcript source identity.
|
||||
@@ -0,0 +1,9 @@
|
||||
A normalized location registry is provided below for identity grounding. It may
|
||||
be empty. Each record contains the exact location ID and canonical display name
|
||||
to copy when the transcript establishes an occurrence of that place.
|
||||
|
||||
Registry content is context, not occurrence evidence. Do not derive an
|
||||
occurrence or a source range from the registry, and do not infer a location
|
||||
that is absent from it.
|
||||
|
||||
{{ input "locations" }}
|
||||
@@ -0,0 +1,20 @@
|
||||
Extract Dungeons & Dragons location occurrences from the supplied transcript.
|
||||
|
||||
Include an occurrence only when the transcript establishes one supplied
|
||||
location, one occurrence kind, and a coherent passage supporting both. Use
|
||||
only the exact ID and name pair from the supplied location registry. Return an
|
||||
empty occurrences array when no supplied location has an evidenced occurrence
|
||||
in this transcript passage.
|
||||
|
||||
Use exactly one kind per occurrence:
|
||||
|
||||
- visited: party members are physically present, arrive, remain, or depart;
|
||||
- planned: the party explicitly proposes, intends, or agrees to future travel;
|
||||
- recalled: the transcript explicitly recounts an earlier party visit; or
|
||||
- mentioned: the location is explicitly referenced without stronger support.
|
||||
|
||||
For overlapping support, visited outranks planned, recalled, and mentioned;
|
||||
planned outranks recalled and mentioned; recalled outranks mentioned. A passage
|
||||
may produce multiple records when it independently establishes separate facts,
|
||||
such as recalling an earlier visit while planning a return. Omit inferred,
|
||||
uncertain, hypothetical, or unsupported places and occurrences.
|
||||
@@ -0,0 +1,34 @@
|
||||
{
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"$id": "notarius.dnd.location_occurrences.llm",
|
||||
"type": "object",
|
||||
"additionalProperties": false,
|
||||
"required": ["occurrences"],
|
||||
"properties": {
|
||||
"occurrences": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"additionalProperties": false,
|
||||
"required": ["location_id", "name", "kind", "source_refs"],
|
||||
"properties": {
|
||||
"location_id": {"type": "string"},
|
||||
"name": {"type": "string"},
|
||||
"kind": {"enum": ["visited", "planned", "recalled", "mentioned"]},
|
||||
"source_refs": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"additionalProperties": false,
|
||||
"required": ["start_unit_id", "end_unit_id"],
|
||||
"properties": {
|
||||
"start_unit_id": {"type": "integer"},
|
||||
"end_unit_id": {"type": "integer"}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
106
internal/modules/dnd/extract/locationoccurrences/canonicalize.go
Normal file
106
internal/modules/dnd/extract/locationoccurrences/canonicalize.go
Normal file
@@ -0,0 +1,106 @@
|
||||
package locationoccurrences
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"sort"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
|
||||
)
|
||||
|
||||
type orderedOccurrence struct {
|
||||
value dnd.LocationOccurrence
|
||||
earliest int
|
||||
hasEvidence bool
|
||||
}
|
||||
|
||||
func canonicalOccurrenceList(response extractionResponse, order shared.SourceRefOrder, sourceID string) dnd.LocationOccurrenceList {
|
||||
if response.Occurrences == nil {
|
||||
return dnd.LocationOccurrenceList{}
|
||||
}
|
||||
ordered := make([]orderedOccurrence, len(response.Occurrences))
|
||||
for index, occurrence := range response.Occurrences {
|
||||
refs := order.Canonicalize(canonicalSourceRefs(occurrence.SourceRefs, sourceID))
|
||||
earliest, hasEvidence := order.EarliestValid(refs)
|
||||
ordered[index] = orderedOccurrence{value: dnd.LocationOccurrence{
|
||||
LocationID: occurrence.LocationID,
|
||||
Name: occurrence.Name,
|
||||
Kind: dnd.LocationOccurrenceKind(occurrence.Kind),
|
||||
SourceRefs: refs,
|
||||
}, earliest: earliest, hasEvidence: hasEvidence}
|
||||
}
|
||||
sort.SliceStable(ordered, func(left, right int) bool {
|
||||
return lessOccurrence(ordered[left], ordered[right], order)
|
||||
})
|
||||
occurrences := make([]dnd.LocationOccurrence, 0, len(ordered))
|
||||
for _, occurrence := range ordered {
|
||||
if len(occurrences) == 0 || !sameOccurrence(occurrences[len(occurrences)-1], occurrence.value) {
|
||||
occurrences = append(occurrences, occurrence.value)
|
||||
}
|
||||
}
|
||||
return dnd.LocationOccurrenceList{Occurrences: occurrences}
|
||||
}
|
||||
|
||||
func lessOccurrence(left, right orderedOccurrence, order shared.SourceRefOrder) bool {
|
||||
if left.hasEvidence != right.hasEvidence {
|
||||
return left.hasEvidence
|
||||
}
|
||||
if left.hasEvidence && left.earliest != right.earliest {
|
||||
return left.earliest < right.earliest
|
||||
}
|
||||
if left.value.LocationID != right.value.LocationID {
|
||||
return left.value.LocationID < right.value.LocationID
|
||||
}
|
||||
if left.value.Name != right.value.Name {
|
||||
return left.value.Name < right.value.Name
|
||||
}
|
||||
if kindOrder(left.value.Kind) != kindOrder(right.value.Kind) {
|
||||
return kindOrder(left.value.Kind) < kindOrder(right.value.Kind)
|
||||
}
|
||||
return lessReferences(left.value.SourceRefs, right.value.SourceRefs, order)
|
||||
}
|
||||
|
||||
func kindOrder(kind dnd.LocationOccurrenceKind) int {
|
||||
switch kind {
|
||||
case dnd.LocationOccurrenceKindVisited:
|
||||
return 0
|
||||
case dnd.LocationOccurrenceKindPlanned:
|
||||
return 1
|
||||
case dnd.LocationOccurrenceKindRecalled:
|
||||
return 2
|
||||
case dnd.LocationOccurrenceKindMentioned:
|
||||
return 3
|
||||
default:
|
||||
return 4
|
||||
}
|
||||
}
|
||||
|
||||
func lessReferences(left, right []source.SourceRef, order shared.SourceRefOrder) bool {
|
||||
limit := len(left)
|
||||
if len(right) < limit {
|
||||
limit = len(right)
|
||||
}
|
||||
for index := 0; index < limit; index++ {
|
||||
if left[index] == right[index] {
|
||||
continue
|
||||
}
|
||||
return order.Less(left[index], right[index])
|
||||
}
|
||||
return len(left) < len(right)
|
||||
}
|
||||
|
||||
func sameOccurrence(left, right dnd.LocationOccurrence) bool {
|
||||
return left.LocationID == right.LocationID && left.Name == right.Name && left.Kind == right.Kind && reflect.DeepEqual(left.SourceRefs, right.SourceRefs)
|
||||
}
|
||||
|
||||
func canonicalSourceRefs(values []occurrenceSourceRefResponse, sourceID string) []source.SourceRef {
|
||||
if values == nil {
|
||||
return nil
|
||||
}
|
||||
refs := make([]source.SourceRef, len(values))
|
||||
for index, value := range values {
|
||||
refs[index] = source.SourceRef{SourceID: sourceID, StartUnitID: value.StartUnitID, EndUnitID: value.EndUnitID}
|
||||
}
|
||||
return refs
|
||||
}
|
||||
188
internal/modules/dnd/extract/locationoccurrences/extractor.go
Normal file
188
internal/modules/dnd/extract/locationoccurrences/extractor.go
Normal file
@@ -0,0 +1,188 @@
|
||||
// Package locationoccurrences extracts source-grounded D&D location occurrences.
|
||||
package locationoccurrences
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
|
||||
locationregistry "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/locations/registry"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
|
||||
)
|
||||
|
||||
const (
|
||||
Key = "dnd/location-occurrences"
|
||||
mappingPolicy = "dnd.location_occurrences.extract_mapping.v1"
|
||||
)
|
||||
|
||||
const (
|
||||
LocationRegistryReferenceSlot = locationregistry.ReferenceSlot
|
||||
LocationRegistryMaxBytes = locationregistry.MaxBytes
|
||||
)
|
||||
|
||||
var requiredCapabilities = []string{"chunks", "source.transcript"}
|
||||
var providedCapabilities = []string{"dnd.location_occurrences"}
|
||||
|
||||
var referenceSlotDescriptions = shared.ReferenceSlotDescriptions{
|
||||
Glossary: "Optional campaign glossary reference material used only for location-occurrence disambiguation.",
|
||||
Party: "Optional party roster reference material used only for location-occurrence disambiguation.",
|
||||
Players: "Optional player list reference material used only for location-occurrence disambiguation.",
|
||||
Roster: "Deprecated alias for party roster reference material used only for location-occurrence disambiguation.",
|
||||
}
|
||||
|
||||
func referenceSlots() []contracts.ReferenceSlot {
|
||||
slots := shared.ReferenceSlots(referenceSlotDescriptions)
|
||||
slots = append(slots, contracts.ReferenceSlot{
|
||||
Name: LocationRegistryReferenceSlot,
|
||||
Description: "Required normalized location registry used only for location identity grounding, never as occurrence evidence.",
|
||||
Required: true,
|
||||
AcceptedMediaTypes: []string{"application/json"},
|
||||
AcceptedArtifactKinds: []contracts.ArtifactKind{dnd.LocationListKind},
|
||||
MaxBytes: LocationRegistryMaxBytes,
|
||||
})
|
||||
sort.Slice(slots, func(left, right int) bool { return slots[left].Name < slots[right].Name })
|
||||
return slots
|
||||
}
|
||||
|
||||
var _ contracts.Extractor[dnd.LocationOccurrenceList] = (*Extractor)(nil)
|
||||
var _ contracts.ManifestMetadataProvider = (*Extractor)(nil)
|
||||
var _ pipeline.CheckpointFingerprintProvider = (*Extractor)(nil)
|
||||
|
||||
type Options struct{}
|
||||
|
||||
type Extractor struct {
|
||||
llm contracts.StructuredLLMClient
|
||||
locationResolver *locationregistry.Resolver
|
||||
promptSHA string
|
||||
responseSchemaSHA string
|
||||
}
|
||||
|
||||
func New(llmClient contracts.StructuredLLMClient, _ Options, references ...contracts.ReferenceSet) (*Extractor, error) {
|
||||
if llmClient == nil {
|
||||
return nil, extractorErrorf("LLM client must not be nil")
|
||||
}
|
||||
if len(references) > 1 {
|
||||
return nil, extractorErrorf("at most one reference set may be supplied")
|
||||
}
|
||||
var referenceSet contracts.ReferenceSet
|
||||
if len(references) == 1 {
|
||||
referenceSet = references[0]
|
||||
}
|
||||
locationResolver, err := locationregistry.NewResolver(referenceSet)
|
||||
if err != nil {
|
||||
return nil, extractorErrorf("prepare location registry prompt input: %w", err)
|
||||
}
|
||||
promptSHA, err := promptAssetMetadata()
|
||||
if err != nil {
|
||||
return nil, extractorErrorf("load prompt metadata: %w", err)
|
||||
}
|
||||
responseSchema, err := loadResponseSchema()
|
||||
if err != nil {
|
||||
return nil, extractorErrorf("load response schema: %w", err)
|
||||
}
|
||||
return &Extractor{llm: llmClient, locationResolver: locationResolver, promptSHA: promptSHA, responseSchemaSHA: responseSchema.SHA256}, nil
|
||||
}
|
||||
|
||||
func (e *Extractor) Key() string { return Key }
|
||||
|
||||
func (e *Extractor) ReferenceSlots() []contracts.ReferenceSlot { return referenceSlots() }
|
||||
|
||||
func (e *Extractor) ManifestMetadata() map[string]any {
|
||||
if e == nil {
|
||||
return nil
|
||||
}
|
||||
metadata := map[string]any{
|
||||
"prompt_id": PromptID,
|
||||
"prompt_version": SchemaVersion,
|
||||
"prompt_sha256": e.promptSHA,
|
||||
"mapping_policy": mappingPolicy,
|
||||
"response_schema_key": string(ResponseSchemaKey),
|
||||
"response_schema_id": ResponseSchemaID,
|
||||
"response_schema_name": ResponseSchemaName,
|
||||
"response_schema_version": SchemaVersion,
|
||||
"response_schema_sha256": e.responseSchemaSHA,
|
||||
}
|
||||
seeded := e.locationResolver.Seeded()
|
||||
if seeded.Bound() {
|
||||
metadata["location_registry_digest"] = seeded.Digest()
|
||||
metadata["location_count"] = seeded.Count()
|
||||
}
|
||||
return metadata
|
||||
}
|
||||
|
||||
func (e *Extractor) CheckpointFingerprints() []pipeline.CheckpointFingerprint {
|
||||
if e == nil {
|
||||
return nil
|
||||
}
|
||||
return []pipeline.CheckpointFingerprint{
|
||||
{Name: "prompt", Value: e.promptSHA},
|
||||
{Name: "response_schema", Value: e.responseSchemaSHA},
|
||||
{Name: "mapping_policy", Value: mappingPolicy},
|
||||
{Name: "location_registry", Value: e.locationResolver.Seeded().ProjectionDigest()},
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Extractor) Extract(ctx context.Context, req contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[dnd.LocationOccurrenceList], error) {
|
||||
if e == nil {
|
||||
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{}, extractorErrorf("extractor must not be nil")
|
||||
}
|
||||
if e.llm == nil {
|
||||
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{}, extractorErrorf("LLM client must not be nil")
|
||||
}
|
||||
sourceInput, err := shared.PrepareChunkExtraction(ctx, req)
|
||||
if err != nil {
|
||||
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{}, extractorErrorf("%w", err)
|
||||
}
|
||||
registry, err := e.locationResolver.Resolve(req.References)
|
||||
if err != nil {
|
||||
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{}, extractorErrorf("resolve location registry: %w", err)
|
||||
}
|
||||
if !registry.Bound() {
|
||||
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{}, extractorErrorf("location registry reference is required")
|
||||
}
|
||||
|
||||
var response extractionResponse
|
||||
inputs := shared.PromptInputs(sourceInput, req.References)
|
||||
inputs[LocationRegistryReferenceSlot] = registry.PromptInput()
|
||||
if _, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
||||
StageName: Key, PromptID: PromptID, PromptVersion: SchemaVersion,
|
||||
ProfileID: req.LLMProfile, SessionID: req.SessionID, Inputs: inputs,
|
||||
}, &response); err != nil {
|
||||
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{}, extractorErrorf("complete structured output: %w", err)
|
||||
}
|
||||
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{Value: canonicalOccurrenceList(response, shared.NewSourceRefOrder(req.Source), req.Source.ID)}, nil
|
||||
}
|
||||
|
||||
func ModuleSpec() pipeline.ModuleSpec {
|
||||
return pipeline.ModuleSpec{
|
||||
Key: Key, Stage: pipeline.StageExtract, ExecutionClass: contracts.ExecutionClassLLMBacked,
|
||||
Requires: append([]string(nil), requiredCapabilities...), Provides: append([]string(nil), providedCapabilities...),
|
||||
ArtifactKind: dnd.LocationOccurrenceListKind, ReferenceSlots: referenceSlots(),
|
||||
}
|
||||
}
|
||||
|
||||
func Register(registry *pipeline.ExtractorRegistry) error {
|
||||
return pipeline.RegisterExtractorBuilder(registry, ModuleSpec(), validateOptions, func(request pipeline.BuildRequest) (contracts.Extractor[dnd.LocationOccurrenceList], error) {
|
||||
options, err := DecodeOptions(request.Options)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return New(request.Dependencies.LLM, options, request.References)
|
||||
})
|
||||
}
|
||||
|
||||
func validateOptions(options map[string]any) error { _, err := DecodeOptions(options); return err }
|
||||
|
||||
func DecodeOptions(options map[string]any) (Options, error) {
|
||||
if err := pipeline.RejectUnknownOptions(options); err != nil {
|
||||
return Options{}, extractorErrorf("%w", err)
|
||||
}
|
||||
return Options{}, nil
|
||||
}
|
||||
|
||||
func extractorErrorf(format string, args ...any) error {
|
||||
return fmt.Errorf("dnd location occurrences extractor: "+format, args...)
|
||||
}
|
||||
@@ -0,0 +1,243 @@
|
||||
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/locations"
|
||||
"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.LocationList{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.LocationListKind}) || 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.LocationList {
|
||||
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.LocationList{Locations: locations}
|
||||
}
|
||||
|
||||
func registryReferences(t *testing.T, locations dnd.LocationList) 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"}}},
|
||||
}}}
|
||||
}
|
||||
17
internal/modules/dnd/extract/locationoccurrences/model.go
Normal file
17
internal/modules/dnd/extract/locationoccurrences/model.go
Normal file
@@ -0,0 +1,17 @@
|
||||
package locationoccurrences
|
||||
|
||||
type extractionResponse struct {
|
||||
Occurrences []occurrenceResponse `json:"occurrences"`
|
||||
}
|
||||
|
||||
type occurrenceResponse struct {
|
||||
LocationID string `json:"location_id"`
|
||||
Name string `json:"name"`
|
||||
Kind string `json:"kind"`
|
||||
SourceRefs []occurrenceSourceRefResponse `json:"source_refs"`
|
||||
}
|
||||
|
||||
type occurrenceSourceRefResponse struct {
|
||||
StartUnitID int `json:"start_unit_id"`
|
||||
EndUnitID int `json:"end_unit_id"`
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
package locationoccurrences
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/promptfs"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
|
||||
)
|
||||
|
||||
const promptAssetRoot = "assets/prompts"
|
||||
|
||||
var promptAssetManifest = shared.PromptAssetManifest{
|
||||
ModuleDir: PromptID,
|
||||
ModuleFiles: []promptfs.ModulePromptFile{
|
||||
{Name: "dnd.location_occurrences.yaml", Path: "assets/prompts/dnd.location_occurrences.yaml"},
|
||||
{Name: "locations.md", Path: "assets/prompts/locations.md"},
|
||||
{Name: "task.md", Path: "assets/prompts/task.md"},
|
||||
{Name: "instructions.md", Path: "assets/prompts/instructions.md"},
|
||||
},
|
||||
SharedFiles: []string{
|
||||
"common-dnd-system.md",
|
||||
"common-dnd-extraction-evidence.md",
|
||||
"common-dnd-identity.md",
|
||||
"common-dnd-transcript.md",
|
||||
"common-dnd-references.md",
|
||||
},
|
||||
}
|
||||
|
||||
func RegisterPromptAssets(registry *llm.AssetRegistry) error {
|
||||
promptFS, err := promptAssetManifest.PromptFS(embeddedAssets)
|
||||
if err != nil {
|
||||
return fmt.Errorf("prepare location-occurrence prompt assets: %w", err)
|
||||
}
|
||||
if err := registry.RegisterPromptFS(promptFS, promptAssetRoot); err != nil {
|
||||
return err
|
||||
}
|
||||
return registry.RegisterSchemaFS(embeddedAssets, "assets/schemas")
|
||||
}
|
||||
|
||||
func promptAssetMetadata() (string, error) {
|
||||
promptAssetHashOnce.Do(func() {
|
||||
promptAssetHash, promptAssetHashErr = promptAssetManifest.Hash(embeddedAssets)
|
||||
})
|
||||
return promptAssetHash, promptAssetHashErr
|
||||
}
|
||||
|
||||
var (
|
||||
promptAssetHashOnce sync.Once
|
||||
promptAssetHash string
|
||||
promptAssetHashErr error
|
||||
)
|
||||
@@ -0,0 +1,64 @@
|
||||
package locationoccurrences
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
|
||||
"gitea.maximumdirect.net/eric/promptkit"
|
||||
)
|
||||
|
||||
func TestRegisterPromptAssetsPreparesLocationOccurrencePrompt(t *testing.T) {
|
||||
registry := llm.NewAssetRegistry()
|
||||
if err := RegisterPromptAssets(registry); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
options, err := registry.PromptKitOptions()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
options = append(options, promptkit.WithProfiles(promptkit.OpenAICompatibleProfile(promptkit.OpenAICompatibleProfileConfig{
|
||||
ID: "location-occurrences-test", Endpoint: "http://127.0.0.1:1/v1", Model: "test",
|
||||
})))
|
||||
engine, err := promptkit.NewEngine(promptkit.Config{Timeout: time.Second}, options...)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{
|
||||
PromptID: PromptID, PromptVersion: SchemaVersion, ProfileID: "location-occurrences-test",
|
||||
Inputs: map[string]promptkit.ArtifactRef{
|
||||
"transcript": promptkit.Inline(`{"units":[1]}`), "players": promptkit.Inline(" "), "party": promptkit.Inline(" "), "glossary": promptkit.Inline(" "),
|
||||
"locations": promptkit.Inline(`{"locations":[{"id":"location:sha256:test","name":"The Mill"}]}`),
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if prepared.PromptID != PromptID || prepared.OutputContract.SchemaPath != "dnd_location_occurrences_llm.v1.json" {
|
||||
t.Fatalf("prepared prompt = %#v", prepared)
|
||||
}
|
||||
var registryMessage string
|
||||
registryIndex := -1
|
||||
evidenceIndex := -1
|
||||
taskIndex := -1
|
||||
for index, message := range prepared.Messages {
|
||||
if strings.Contains(message.Content, "normalized location registry") {
|
||||
registryMessage = message.Content
|
||||
registryIndex = index
|
||||
}
|
||||
if strings.Contains(message.Content, "Transcript units are the only evidence") {
|
||||
evidenceIndex = index
|
||||
}
|
||||
if strings.Contains(message.Content, "Extract Dungeons & Dragons location occurrences") {
|
||||
taskIndex = index
|
||||
}
|
||||
}
|
||||
if !strings.Contains(registryMessage, "location:sha256:test") || !strings.Contains(registryMessage, "The Mill") || strings.Contains(registryMessage, "source_refs") {
|
||||
t.Fatalf("rendered prompt did not preserve source-free registry grounding: %s", registryMessage)
|
||||
}
|
||||
if evidenceIndex < 0 || taskIndex < 0 || registryIndex <= evidenceIndex || registryIndex >= taskIndex {
|
||||
t.Fatalf("registry prompt placement = evidence %d, registry %d, task %d", evidenceIndex, registryIndex, taskIndex)
|
||||
}
|
||||
}
|
||||
21
internal/modules/dnd/extract/locationoccurrences/schema.go
Normal file
21
internal/modules/dnd/extract/locationoccurrences/schema.go
Normal file
@@ -0,0 +1,21 @@
|
||||
package locationoccurrences
|
||||
|
||||
import "gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
|
||||
|
||||
const (
|
||||
PromptID = "dnd.location_occurrences"
|
||||
ResponseSchemaKey = llm.ResponseSchemaKey("dnd_location_occurrences_llm")
|
||||
ResponseSchemaID = "notarius.dnd.location_occurrences.llm"
|
||||
ResponseSchemaName = "notarius_dnd_location_occurrences_llm_v1"
|
||||
SchemaVersion = "v1"
|
||||
)
|
||||
|
||||
func loadResponseSchema() (llm.ResponseSchema, error) {
|
||||
return llm.LoadResponseSchema(embeddedAssets, llm.ResponseSchemaDefinition{
|
||||
Key: ResponseSchemaKey,
|
||||
ID: ResponseSchemaID,
|
||||
Version: SchemaVersion,
|
||||
Name: ResponseSchemaName,
|
||||
AssetPath: "assets/schemas/dnd_location_occurrences_llm.v1.json",
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package locationoccurrences
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/santhosh-tekuri/jsonschema/v6"
|
||||
)
|
||||
|
||||
func TestResponseSchemaRestrictsPrivateOccurrenceStructureAndKinds(t *testing.T) {
|
||||
schema, err := loadResponseSchema()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if schema.Key != ResponseSchemaKey || schema.ID != ResponseSchemaID || schema.Name != ResponseSchemaName || schema.Version != SchemaVersion || !strings.HasPrefix(schema.SHA256, "sha256:") || !json.Valid(schema.JSONSchema) {
|
||||
t.Fatalf("schema = %#v", schema)
|
||||
}
|
||||
valid := map[string]any{"occurrences": []any{map[string]any{
|
||||
"location_id": "location:sha256:test", "name": "The Mill", "kind": "visited",
|
||||
"source_refs": []any{map[string]any{"start_unit_id": 1, "end_unit_id": 2}},
|
||||
}}}
|
||||
if err := validateSchema(t, valid, schema.JSONSchema); err != nil {
|
||||
t.Fatalf("valid response rejected: %v", err)
|
||||
}
|
||||
for _, mutate := range []func(map[string]any){
|
||||
func(value map[string]any) { delete(value, "location_id") },
|
||||
func(value map[string]any) { value["kind"] = "other" },
|
||||
func(value map[string]any) { value["unexpected"] = true },
|
||||
func(value map[string]any) {
|
||||
value["source_refs"].([]any)[0].(map[string]any)["source_id"] = "assigned later"
|
||||
},
|
||||
} {
|
||||
candidate := cloneCandidate(t, valid)
|
||||
mutate(candidate["occurrences"].([]any)[0].(map[string]any))
|
||||
if err := validateSchema(t, candidate, schema.JSONSchema); err == nil {
|
||||
t.Fatal("schema accepted structurally invalid response")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponseSchemaIsDefensiveAndContentSafe(t *testing.T) {
|
||||
first, err := loadResponseSchema()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
first.JSONSchema[0] = '['
|
||||
second, err := loadResponseSchema()
|
||||
if err != nil || !json.Valid(second.JSONSchema) || bytes.Equal(first.JSONSchema, second.JSONSchema) {
|
||||
t.Fatalf("schema = %s, %v", second.JSONSchema, err)
|
||||
}
|
||||
if diagnostics := second.DiagnosticsMap(); diagnostics["key"] != ResponseSchemaKey || diagnostics["id"] != ResponseSchemaID {
|
||||
t.Fatalf("diagnostics = %#v", diagnostics)
|
||||
} else if _, ok := diagnostics["json_schema"]; ok {
|
||||
t.Fatalf("diagnostics leaked schema content: %#v", diagnostics)
|
||||
}
|
||||
}
|
||||
|
||||
func cloneCandidate(t *testing.T, value map[string]any) map[string]any {
|
||||
t.Helper()
|
||||
content, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var clone map[string]any
|
||||
if err := json.Unmarshal(content, &clone); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return clone
|
||||
}
|
||||
|
||||
func validateSchema(t *testing.T, value map[string]any, schemaContent []byte) error {
|
||||
t.Helper()
|
||||
content, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
instance, err := jsonschema.UnmarshalJSON(bytes.NewReader(content))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
document, err := jsonschema.UnmarshalJSON(bytes.NewReader(schemaContent))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
compiler := jsonschema.NewCompiler()
|
||||
if err := compiler.AddResource("schema.json", document); err != nil {
|
||||
return err
|
||||
}
|
||||
schema, err := compiler.Compile("schema.json")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return schema.Validate(instance)
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
package locationoccurrences
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||
)
|
||||
|
||||
func extractionRequest() contracts.TypedExtractionRequest {
|
||||
doc := sourceDocument()
|
||||
chunk := &source.Chunk{
|
||||
ID: "session-occurrences:chunk:0", SourceID: doc.ID, Index: 0,
|
||||
Ref: source.SourceRef{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 30},
|
||||
Content: []byte(`{"units":[10,20,30]}`), MediaType: "application/json",
|
||||
Units: append([]source.SourceUnit(nil), doc.Units...),
|
||||
}
|
||||
return contracts.TypedExtractionRequest{
|
||||
Source: doc, Chunk: chunk,
|
||||
SourceInput: contracts.NewLLMInputMaterial("source", chunk.MediaType, chunk.Content, "sha256:chunk", "file:///session-occurrences.json"),
|
||||
SessionID: "occurrence-session", LLMProfile: "occurrence-profile",
|
||||
}
|
||||
}
|
||||
|
||||
func sourceDocument() *source.SourceDocument {
|
||||
return &source.SourceDocument{ID: "session-occurrences", Kind: "transcript", Format: "application/json", Digest: "sha256:test", Units: []source.SourceUnit{
|
||||
{ID: 10, Kind: "transcript_segment", Text: "The party returns to the tavern."},
|
||||
{ID: 20, Kind: "transcript_segment", Text: "They plan to travel to the tavern tomorrow."},
|
||||
{ID: 30, Kind: "transcript_segment", Text: "They recall their first visit to the tavern."},
|
||||
}}
|
||||
}
|
||||
|
||||
func occurrenceRefs(start, end int) []occurrenceSourceRefResponse {
|
||||
return []occurrenceSourceRefResponse{{StartUnitID: start, EndUnitID: end}}
|
||||
}
|
||||
|
||||
func newExtractor(t *testing.T, client contracts.StructuredLLMClient, references ...contracts.ReferenceSet) *Extractor {
|
||||
t.Helper()
|
||||
extractor, err := New(client, Options{}, references...)
|
||||
if err != nil {
|
||||
t.Fatalf("New() error = %v", err)
|
||||
}
|
||||
return extractor
|
||||
}
|
||||
|
||||
func mismatchedSourceInputRequest(req contracts.TypedExtractionRequest) contracts.TypedExtractionRequest {
|
||||
req.SourceInput = contracts.NewLLMInputMaterial("source", "application/json", []byte(`{"different":true}`), "sha256:other", "")
|
||||
return req
|
||||
}
|
||||
|
||||
type fakeOccurrencesLLMClient struct {
|
||||
response extractionResponse
|
||||
err error
|
||||
requests []contracts.StructuredCompletionRequest
|
||||
}
|
||||
|
||||
func (client *fakeOccurrencesLLMClient) CompleteStructured(_ context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
|
||||
client.requests = append(client.requests, cloneStructuredCompletionRequest(req))
|
||||
if client.err != nil {
|
||||
return contracts.StructuredCompletionResponse{}, client.err
|
||||
}
|
||||
target, ok := out.(*extractionResponse)
|
||||
if !ok {
|
||||
return contracts.StructuredCompletionResponse{}, errors.New("unexpected output target")
|
||||
}
|
||||
content, err := json.Marshal(client.response)
|
||||
if err != nil {
|
||||
return contracts.StructuredCompletionResponse{}, err
|
||||
}
|
||||
if err := json.Unmarshal(content, target); err != nil {
|
||||
return contracts.StructuredCompletionResponse{}, err
|
||||
}
|
||||
return contracts.StructuredCompletionResponse{Content: content}, nil
|
||||
}
|
||||
|
||||
func cloneStructuredCompletionRequest(req contracts.StructuredCompletionRequest) contracts.StructuredCompletionRequest {
|
||||
req.Inputs = req.Inputs.Clone()
|
||||
return req
|
||||
}
|
||||
Reference in New Issue
Block a user