302 lines
11 KiB
Go
302 lines
11 KiB
Go
package spells
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"strings"
|
|
"testing"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
)
|
|
|
|
func TestExtractReturnsSpellCandidateFromStructuredOutput(t *testing.T) {
|
|
client := &fakeSpellsLLMClient{
|
|
response: extractionResponse{
|
|
SpellCasts: []spellCastResponse{
|
|
{
|
|
Caster: " Aria ",
|
|
Spell: " Cure Wounds ",
|
|
Effect: " Heals an injured ally. ",
|
|
NarrativeDescription: " Aria restores the fighter after the fight. ",
|
|
SourceRefs: []source.SourceRef{
|
|
{SourceID: "session-alpha", StartUnitID: "seg-001", EndUnitID: "seg-002"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
result, err := New().Extract(context.Background(), extractionRequestWithClient(client))
|
|
if err != nil {
|
|
t.Fatalf("Extract() error = %v, want nil", err)
|
|
}
|
|
|
|
if len(client.requests) != 1 {
|
|
t.Fatalf("LLM calls = %d, want 1", len(client.requests))
|
|
}
|
|
req := client.requests[0]
|
|
if req.StageName != Key {
|
|
t.Fatalf("StageName = %q, want %q", req.StageName, Key)
|
|
}
|
|
schema, err := loadResponseSchema()
|
|
if err != nil {
|
|
t.Fatalf("loadResponseSchema() error = %v, want nil", err)
|
|
}
|
|
if req.ResponseSchemaName != schema.Name {
|
|
t.Fatalf("ResponseSchemaName = %q, want %q", req.ResponseSchemaName, schema.Name)
|
|
}
|
|
if !bytes.Equal(req.ResponseSchema, schema.JSONSchema) {
|
|
t.Fatal("ResponseSchema does not match registered D&D spells schema")
|
|
}
|
|
if len(req.Messages) != 2 {
|
|
t.Fatalf("len(Messages) = %d, want 2", len(req.Messages))
|
|
}
|
|
if req.Messages[0].Role != "system" || req.Messages[1].Role != "user" {
|
|
t.Fatalf("Messages roles = %#v, want system then user", req.Messages)
|
|
}
|
|
if !strings.Contains(req.Messages[0].Content, "D&D spell-cast") {
|
|
t.Fatalf("system message = %q, want D&D spell context", req.Messages[0].Content)
|
|
}
|
|
for _, want := range []string{"session-alpha", "session-alpha:chunk:0", "seg-001", "Cure Wounds"} {
|
|
if !strings.Contains(req.Messages[1].Content, want) {
|
|
t.Fatalf("user message = %q, want substring %q", req.Messages[1].Content, want)
|
|
}
|
|
}
|
|
|
|
if len(result.Candidates) != 1 {
|
|
t.Fatalf("len(Candidates) = %d, want 1", len(result.Candidates))
|
|
}
|
|
candidate := result.Candidates[0]
|
|
if candidate.Index != 0 || candidate.ExtractorKey != "" || candidate.ArtifactType != "" || candidate.SchemaVersion != "" {
|
|
t.Fatalf("candidate envelope fields = %#v, want runner-normalized zero values", candidate)
|
|
}
|
|
var payload SpellCast
|
|
if err := json.Unmarshal(candidate.Payload, &payload); err != nil {
|
|
t.Fatalf("Unmarshal(Payload) error = %v, want nil", err)
|
|
}
|
|
wantPayload := SpellCast{
|
|
Caster: "Aria",
|
|
Spell: "Cure Wounds",
|
|
Effect: "Heals an injured ally.",
|
|
NarrativeDescription: "Aria restores the fighter after the fight.",
|
|
}
|
|
if payload != wantPayload {
|
|
t.Fatalf("payload = %#v, want %#v", payload, wantPayload)
|
|
}
|
|
wantRef := source.SourceRef{SourceID: "session-alpha", StartUnitID: "seg-001", EndUnitID: "seg-002"}
|
|
if len(candidate.SourceRefs) != 1 || candidate.SourceRefs[0] != wantRef {
|
|
t.Fatalf("SourceRefs = %#v, want %#v", candidate.SourceRefs, []source.SourceRef{wantRef})
|
|
}
|
|
}
|
|
|
|
func TestExtractorManifestMetadataIncludesPromptAndSchemaProvenance(t *testing.T) {
|
|
metadata := New().ManifestMetadata()
|
|
|
|
tests := map[string]string{
|
|
"prompt_id": PromptID,
|
|
"prompt_version": SchemaVersion,
|
|
"response_schema_key": string(ResponseSchemaKey),
|
|
"response_schema_id": ResponseSchemaID,
|
|
"response_schema_name": ResponseSchemaName,
|
|
"response_schema_version": SchemaVersion,
|
|
}
|
|
for key, want := range tests {
|
|
if metadata[key] != want {
|
|
t.Fatalf("metadata[%q] = %#v, want %q", key, metadata[key], want)
|
|
}
|
|
}
|
|
for _, key := range []string{"prompt_sha256", "response_schema_sha256"} {
|
|
value, ok := metadata[key].(string)
|
|
if !ok || !strings.HasPrefix(value, "sha256:") {
|
|
t.Fatalf("metadata[%q] = %#v, want sha256 value", key, metadata[key])
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestExtractReturnsNoCandidatesForEmptyResponse(t *testing.T) {
|
|
client := &fakeSpellsLLMClient{response: extractionResponse{SpellCasts: []spellCastResponse{}}}
|
|
|
|
result, err := New().Extract(context.Background(), extractionRequestWithClient(client))
|
|
if err != nil {
|
|
t.Fatalf("Extract() error = %v, want nil", err)
|
|
}
|
|
if len(result.Candidates) != 0 {
|
|
t.Fatalf("Candidates = %#v, want none", result.Candidates)
|
|
}
|
|
}
|
|
|
|
func TestExtractRejectsMissingSpellCasts(t *testing.T) {
|
|
client := &fakeSpellsLLMClient{response: extractionResponse{}}
|
|
|
|
_, err := New().Extract(context.Background(), extractionRequestWithClient(client))
|
|
if err == nil {
|
|
t.Fatal("Extract() error = nil, want malformed output error")
|
|
}
|
|
if !strings.Contains(err.Error(), "dnd spells") || !strings.Contains(err.Error(), "spell_casts") {
|
|
t.Fatalf("Extract() error = %q, want spell_casts context", err.Error())
|
|
}
|
|
}
|
|
|
|
func TestExtractWrapsLLMClientError(t *testing.T) {
|
|
client := &fakeSpellsLLMClient{err: errors.New("provider unavailable")}
|
|
|
|
_, err := New().Extract(context.Background(), extractionRequestWithClient(client))
|
|
if err == nil {
|
|
t.Fatal("Extract() error = nil, want LLM error")
|
|
}
|
|
if !strings.Contains(err.Error(), "dnd spells") || !strings.Contains(err.Error(), "provider unavailable") {
|
|
t.Fatalf("Extract() error = %q, want wrapped LLM context", err.Error())
|
|
}
|
|
}
|
|
|
|
func TestExtractRejectsInvalidRequests(t *testing.T) {
|
|
validClient := &fakeSpellsLLMClient{response: extractionResponse{SpellCasts: []spellCastResponse{}}}
|
|
validReq := extractionRequestWithClient(validClient)
|
|
canceledCtx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
extractor *Extractor
|
|
ctx context.Context
|
|
req contracts.ExtractionRequest
|
|
want string
|
|
}{
|
|
{name: "nil extractor", extractor: nil, ctx: context.Background(), req: validReq, want: "extractor"},
|
|
{name: "nil context", extractor: New(), ctx: nil, req: validReq, want: "context"},
|
|
{name: "canceled context", extractor: New(), ctx: canceledCtx, req: validReq, want: "context"},
|
|
{name: "nil source", extractor: New(), ctx: context.Background(), req: contracts.ExtractionRequest{Chunk: validReq.Chunk, LLMClient: validReq.LLMClient}, want: "source"},
|
|
{name: "nil chunk", extractor: New(), ctx: context.Background(), req: contracts.ExtractionRequest{Source: validReq.Source, LLMClient: validReq.LLMClient}, want: "chunk"},
|
|
{name: "empty chunk units", extractor: New(), ctx: context.Background(), req: emptyChunkRequest(validReq), want: "units"},
|
|
{name: "nil LLM client", extractor: New(), ctx: context.Background(), req: contracts.ExtractionRequest{Source: validReq.Source, Chunk: validReq.Chunk}, want: "LLM client"},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
_, err := tt.extractor.Extract(tt.ctx, tt.req)
|
|
if err == nil {
|
|
t.Fatal("Extract() error = nil, want error")
|
|
}
|
|
if !strings.Contains(err.Error(), "dnd spells") || !strings.Contains(err.Error(), tt.want) {
|
|
t.Fatalf("Extract() error = %q, want %q context", err.Error(), tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestExtractPreservesResponseOrder(t *testing.T) {
|
|
client := &fakeSpellsLLMClient{
|
|
response: extractionResponse{
|
|
SpellCasts: []spellCastResponse{
|
|
{
|
|
Caster: "Aria",
|
|
Spell: "Cure Wounds",
|
|
Effect: "Heals.",
|
|
NarrativeDescription: "First spell.",
|
|
SourceRefs: []source.SourceRef{{SourceID: "session-alpha", StartUnitID: "seg-001", EndUnitID: "seg-001"}},
|
|
},
|
|
{
|
|
Caster: "Bandit Shaman",
|
|
Spell: "Fire Bolt",
|
|
Effect: "Burns.",
|
|
NarrativeDescription: "Second spell.",
|
|
SourceRefs: []source.SourceRef{{SourceID: "session-alpha", StartUnitID: "seg-002", EndUnitID: "seg-002"}},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
result, err := New().Extract(context.Background(), extractionRequestWithClient(client))
|
|
if err != nil {
|
|
t.Fatalf("Extract() error = %v, want nil", err)
|
|
}
|
|
if len(result.Candidates) != 2 {
|
|
t.Fatalf("len(Candidates) = %d, want 2", len(result.Candidates))
|
|
}
|
|
|
|
var first, second SpellCast
|
|
if err := json.Unmarshal(result.Candidates[0].Payload, &first); err != nil {
|
|
t.Fatalf("Unmarshal(first) error = %v, want nil", err)
|
|
}
|
|
if err := json.Unmarshal(result.Candidates[1].Payload, &second); err != nil {
|
|
t.Fatalf("Unmarshal(second) error = %v, want nil", err)
|
|
}
|
|
if first.Spell != "Cure Wounds" || second.Spell != "Fire Bolt" {
|
|
t.Fatalf("candidate order = %q, %q; want response order", first.Spell, second.Spell)
|
|
}
|
|
}
|
|
|
|
func TestExtractCopiesCandidateSourceRefs(t *testing.T) {
|
|
client := &fakeSpellsLLMClient{
|
|
response: extractionResponse{
|
|
SpellCasts: []spellCastResponse{
|
|
{
|
|
Caster: "Aria",
|
|
Spell: "Cure Wounds",
|
|
Effect: "Heals.",
|
|
NarrativeDescription: "Aria heals.",
|
|
SourceRefs: []source.SourceRef{{SourceID: "session-alpha", StartUnitID: "seg-001", EndUnitID: "seg-002"}},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
result, err := New().Extract(context.Background(), extractionRequestWithClient(client))
|
|
if err != nil {
|
|
t.Fatalf("Extract() error = %v, want nil", err)
|
|
}
|
|
client.response.SpellCasts[0].SourceRefs[0].StartUnitID = "mutated"
|
|
|
|
if got := result.Candidates[0].SourceRefs[0].StartUnitID; got != "seg-001" {
|
|
t.Fatalf("candidate source ref start = %q, want copied seg-001", got)
|
|
}
|
|
}
|
|
|
|
func extractionRequestWithClient(client contracts.StructuredLLMClient) contracts.ExtractionRequest {
|
|
req := promptExtractionRequest()
|
|
req.LLMClient = client
|
|
return req
|
|
}
|
|
|
|
func emptyChunkRequest(req contracts.ExtractionRequest) contracts.ExtractionRequest {
|
|
req.Chunk = &contracts.SourceChunk{
|
|
ID: req.Chunk.ID,
|
|
SourceID: req.Chunk.SourceID,
|
|
Index: req.Chunk.Index,
|
|
}
|
|
return req
|
|
}
|
|
|
|
type fakeSpellsLLMClient struct {
|
|
response extractionResponse
|
|
err error
|
|
requests []contracts.StructuredCompletionRequest
|
|
}
|
|
|
|
func (client *fakeSpellsLLMClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
|
|
client.requests = append(client.requests, contracts.StructuredCompletionRequest{
|
|
StageName: req.StageName,
|
|
Messages: append([]contracts.LLMMessage(nil), req.Messages...),
|
|
Model: req.Model,
|
|
ResponseSchemaName: req.ResponseSchemaName,
|
|
ResponseSchema: append(json.RawMessage(nil), req.ResponseSchema...),
|
|
})
|
|
if client.err != nil {
|
|
return contracts.StructuredCompletionResponse{}, client.err
|
|
}
|
|
|
|
target, ok := out.(*extractionResponse)
|
|
if !ok {
|
|
return contracts.StructuredCompletionResponse{}, errors.New("unexpected output target")
|
|
}
|
|
*target = client.response
|
|
content, err := json.Marshal(client.response)
|
|
if err != nil {
|
|
return contracts.StructuredCompletionResponse{}, err
|
|
}
|
|
return contracts.StructuredCompletionResponse{Content: content}, nil
|
|
}
|