Add D&D location tracking to complete example
This commit is contained in:
268
internal/modules/integration/dnd_locations_runner_test.go
Normal file
268
internal/modules/integration/dnd_locations_runner_test.go
Normal file
@@ -0,0 +1,268 @@
|
||||
package integration_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/core/config"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/evidencecontext"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
|
||||
locationoccurrencecodec "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/codec/locationoccurrences"
|
||||
locationcodec "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/codec/locations"
|
||||
locationoccurrences "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/locationoccurrences"
|
||||
locations "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/locations"
|
||||
locationnormalize "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/normalize/locations"
|
||||
)
|
||||
|
||||
func TestLocationRegistryHandoffProducesOccurrencesAndEvidence(t *testing.T) {
|
||||
registries := productionNPCRegistries(t)
|
||||
resolved := resolveLocationPipeline(t, registries)
|
||||
client := &locationHandoffLLMClient{}
|
||||
output, err := runPreparedPipeline(t, registries, resolved, client, pipeline.RunInput{RawInput: readNPCFixture(t)})
|
||||
if err != nil {
|
||||
t.Fatalf("Run() error = %v", err)
|
||||
}
|
||||
if len(output.Rejected) != 0 || len(output.NormalizeOutputs) != 2 {
|
||||
t.Fatalf("run outputs = %#v rejected = %#v, want accepted registry and occurrences", output.NormalizeOutputs, output.Rejected)
|
||||
}
|
||||
|
||||
registryOutput := normalizedLane(t, output, "locations")
|
||||
if registryOutput.Artifact.Kind != dnd.LocationListKind || registryOutput.Artifact.Schema.ID != locationcodec.SchemaID {
|
||||
t.Fatalf("registry envelope = %#v, want durable location artifact", registryOutput)
|
||||
}
|
||||
registry, err := locationcodec.New().Decode(registryOutput.Artifact.Content)
|
||||
if err != nil {
|
||||
t.Fatalf("Decode(registry) error = %v", err)
|
||||
}
|
||||
if len(registry.Locations) != 2 || registry.Locations[0].Name != "Moon Gate" || registry.Locations[1].Name != "Moon Gate" || registry.Locations[0].ID == registry.Locations[1].ID {
|
||||
t.Fatalf("registry = %#v, want same-name locations distinguished by source-derived IDs", registry)
|
||||
}
|
||||
|
||||
occurrenceOutput := normalizedLane(t, output, "occurrences")
|
||||
if occurrenceOutput.Artifact.Kind != dnd.LocationOccurrenceListKind || occurrenceOutput.Artifact.Schema.ID != locationoccurrencecodec.SchemaID {
|
||||
t.Fatalf("occurrence envelope = %#v, want durable occurrence artifact", occurrenceOutput)
|
||||
}
|
||||
occurrences, err := locationoccurrencecodec.New().Decode(occurrenceOutput.Artifact.Content)
|
||||
if err != nil {
|
||||
t.Fatalf("Decode(occurrences) error = %v", err)
|
||||
}
|
||||
if len(occurrences.Occurrences) != 2 || occurrences.Occurrences[0].LocationID == occurrences.Occurrences[1].LocationID || occurrences.Occurrences[0].LocationID != registry.Locations[0].ID || occurrences.Occurrences[1].LocationID != registry.Locations[1].ID {
|
||||
t.Fatalf("occurrences = %#v, want exact registry ID grounding", occurrences)
|
||||
}
|
||||
for _, occurrence := range occurrences.Occurrences {
|
||||
if occurrence.SourceRefs[0].SourceID != "npc-session" {
|
||||
t.Fatalf("occurrence evidence = %#v, want current source evidence only", occurrence.SourceRefs)
|
||||
}
|
||||
}
|
||||
|
||||
request := client.requestFor(t, locationoccurrences.PromptID)
|
||||
registryInput := request.Inputs["locations"]
|
||||
if registryInput.MediaType != locationcodec.MediaType || strings.Contains(string(registryInput.Content), "source_refs") || !strings.Contains(string(registryInput.Content), registry.Locations[0].ID) || !strings.Contains(string(registryInput.Content), registry.Locations[1].ID) {
|
||||
t.Fatalf("occurrence registry input = %#v, want source-free generated ID projection", registryInput)
|
||||
}
|
||||
contextArtifact := outputFileContent(t, output.OutputFiles, "evidence-context.json")
|
||||
evidence, err := evidencecontext.New().Decode(contextArtifact)
|
||||
if err != nil {
|
||||
t.Fatalf("Decode(evidence context) error = %v", err)
|
||||
}
|
||||
if !locationEvidenceHasLane(evidence, "locations") || !locationEvidenceHasLane(evidence, "occurrences") {
|
||||
t.Fatalf("evidence context = %#v, want registry and occurrence evidence from their own artifacts", evidence)
|
||||
}
|
||||
}
|
||||
|
||||
func locationEvidenceHasLane(document evidencecontext.Document, laneID string) bool {
|
||||
for _, context := range document.Contexts {
|
||||
for _, reference := range context.EvidenceRefs {
|
||||
if reference.LaneID == laneID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func TestLocationOccurrenceConsumerDoesNotRunAfterRejectedRegistry(t *testing.T) {
|
||||
registries := productionNPCRegistries(t)
|
||||
client := &locationHandoffLLMClient{rejectRegistry: true}
|
||||
_, err := runPreparedPipeline(t, registries, resolveLocationPipeline(t, registries), client, pipeline.RunInput{RawInput: readNPCFixture(t)})
|
||||
if err == nil {
|
||||
t.Fatal("Run() error = nil, want missing accepted location registry")
|
||||
}
|
||||
if client.requestCount(locationoccurrences.PromptID) != 0 {
|
||||
t.Fatalf("occurrence requests = %d, want none after rejected registry", client.requestCount(locationoccurrences.PromptID))
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocationOccurrenceConfigurationRequiresEarlierRegistry(t *testing.T) {
|
||||
registries := productionNPCRegistries(t)
|
||||
cfg, err := configFromYAML(strings.Replace(locationPipelineConfig, "step: identify-locations", "step: track-location-occurrences", 1))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = cfg.Resolve(config.ResolveInput{PipelineID: "dnd-locations-fixture", Catalog: moduleCatalog(registries)})
|
||||
if err == nil || !strings.Contains(err.Error(), "earlier step") {
|
||||
t.Fatalf("Resolve() error = %v, want an earlier generated registry requirement", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocationOccurrenceCheckpointTracksGeneratedRegistry(t *testing.T) {
|
||||
registries := productionNPCRegistries(t)
|
||||
checkpoint := newGeneratedReferenceCheckpointLoader()
|
||||
_, err := runPreparedPipeline(t, registries, resolveLocationPipeline(t, registries), &locationHandoffLLMClient{}, pipeline.RunInput{
|
||||
RawInput: readNPCFixture(t),
|
||||
Checkpoint: checkpoint,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Run() error = %v", err)
|
||||
}
|
||||
for _, dependency := range checkpoint.extractDependencies("occurrences") {
|
||||
if dependency.Name == "generated-reference:locations:0" && strings.HasPrefix(dependency.Value, "sha256:") {
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Fatalf("occurrence checkpoint dependencies = %#v, want generated location registry fingerprint", checkpoint.extractDependencies("occurrences"))
|
||||
}
|
||||
|
||||
func resolveLocationPipeline(t *testing.T, registries pipeline.Registries) pipeline.ResolvedPipeline {
|
||||
t.Helper()
|
||||
cfg, err := configFromYAML(locationPipelineConfig)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
effective, err := cfg.Resolve(config.ResolveInput{PipelineID: "dnd-locations-fixture", Catalog: moduleCatalog(registries)})
|
||||
if err != nil {
|
||||
t.Fatalf("Resolve() error = %v", err)
|
||||
}
|
||||
resolved, warnings, err := pipeline.MaterializeReferences(effective.ResolvedPipeline, moduleCatalog(registries), pipeline.ReferenceMaterializationOptions{})
|
||||
if err != nil || len(warnings) != 0 {
|
||||
t.Fatalf("MaterializeReferences() error = %v warnings = %#v", err, warnings)
|
||||
}
|
||||
return resolved
|
||||
}
|
||||
|
||||
func configFromYAML(content string) (config.Config, error) {
|
||||
fileConfig, err := config.ParseFileConfigYAML([]byte(content))
|
||||
if err != nil {
|
||||
return config.Config{}, err
|
||||
}
|
||||
cfg := config.Default()
|
||||
if err := cfg.ApplyFileConfig(fileConfig); err != nil {
|
||||
return config.Config{}, err
|
||||
}
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
const locationPipelineConfig = `version: 4
|
||||
cache:
|
||||
chunk_plans:
|
||||
mode: bypass
|
||||
checkpoints: {}
|
||||
pipelines:
|
||||
dnd-locations-fixture:
|
||||
input: seriatim
|
||||
output:
|
||||
module: json
|
||||
options:
|
||||
evidence_context:
|
||||
enabled: true
|
||||
lanes: [locations, occurrences]
|
||||
steps:
|
||||
- id: identify-locations
|
||||
artifacts:
|
||||
locations:
|
||||
extract: dnd/locations
|
||||
merge: appendorder
|
||||
normalize: dnd/locations
|
||||
- id: track-location-occurrences
|
||||
references:
|
||||
locations:
|
||||
artifact:
|
||||
step: identify-locations
|
||||
lane: locations
|
||||
artifacts:
|
||||
occurrences:
|
||||
extract: dnd/location-occurrences
|
||||
merge: appendorder
|
||||
normalize: dnd/location-occurrences
|
||||
`
|
||||
|
||||
type locationHandoffLLMClient struct {
|
||||
mu sync.Mutex
|
||||
requests []contracts.StructuredCompletionRequest
|
||||
rejectRegistry bool
|
||||
}
|
||||
|
||||
func (client *locationHandoffLLMClient) CompleteStructured(ctx context.Context, request contracts.StructuredCompletionRequest, output any) (contracts.StructuredCompletionResponse, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return contracts.StructuredCompletionResponse{}, err
|
||||
}
|
||||
client.mu.Lock()
|
||||
client.requests = append(client.requests, cloneStructuredCompletionRequest(request))
|
||||
client.mu.Unlock()
|
||||
var payload any
|
||||
switch request.PromptID {
|
||||
case locations.PromptID:
|
||||
name := "Moon Gate"
|
||||
if client.rejectRegistry {
|
||||
name = ""
|
||||
}
|
||||
payload = map[string]any{"locations": []any{map[string]any{"name": name, "source_refs": []any{map[string]int{"start_unit_id": 1, "end_unit_id": 1}}}, map[string]any{"name": name, "source_refs": []any{map[string]int{"start_unit_id": 3, "end_unit_id": 3}}}}}
|
||||
case locationnormalize.PromptID:
|
||||
payload = map[string]any{"duplicate_groups": []any{}}
|
||||
case locationoccurrences.PromptID:
|
||||
var projection struct {
|
||||
Locations []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"locations"`
|
||||
}
|
||||
if err := json.Unmarshal(request.Inputs["locations"].Content, &projection); err != nil {
|
||||
return contracts.StructuredCompletionResponse{}, fmt.Errorf("decode location registry: %w", err)
|
||||
}
|
||||
if len(projection.Locations) != 2 {
|
||||
return contracts.StructuredCompletionResponse{}, fmt.Errorf("generated location registry has %d locations, want 2", len(projection.Locations))
|
||||
}
|
||||
payload = map[string]any{"occurrences": []any{map[string]any{"location_id": projection.Locations[0].ID, "name": "Moon Gate", "kind": "visited", "source_refs": []any{map[string]int{"start_unit_id": 1, "end_unit_id": 1}}}, map[string]any{"location_id": projection.Locations[1].ID, "name": "Moon Gate", "kind": "mentioned", "source_refs": []any{map[string]int{"start_unit_id": 3, "end_unit_id": 3}}}}}
|
||||
default:
|
||||
return contracts.StructuredCompletionResponse{}, fmt.Errorf("unexpected prompt %q", request.PromptID)
|
||||
}
|
||||
content, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return contracts.StructuredCompletionResponse{}, err
|
||||
}
|
||||
if err := json.Unmarshal(content, output); err != nil {
|
||||
return contracts.StructuredCompletionResponse{}, err
|
||||
}
|
||||
return contracts.StructuredCompletionResponse{Content: content, Provider: "test", Model: "location-handoff"}, nil
|
||||
}
|
||||
|
||||
func (client *locationHandoffLLMClient) requestFor(t *testing.T, promptID string) contracts.StructuredCompletionRequest {
|
||||
t.Helper()
|
||||
client.mu.Lock()
|
||||
defer client.mu.Unlock()
|
||||
for _, request := range client.requests {
|
||||
if request.PromptID == promptID {
|
||||
return request
|
||||
}
|
||||
}
|
||||
t.Fatalf("requests = %#v, missing %q", client.requests, promptID)
|
||||
return contracts.StructuredCompletionRequest{}
|
||||
}
|
||||
func (client *locationHandoffLLMClient) requestCount(promptID string) int {
|
||||
client.mu.Lock()
|
||||
defer client.mu.Unlock()
|
||||
count := 0
|
||||
for _, request := range client.requests {
|
||||
if request.PromptID == promptID {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
var _ contracts.StructuredLLMClient = (*locationHandoffLLMClient)(nil)
|
||||
Reference in New Issue
Block a user