eric/notarius
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Cleanup and complete the pipeline refactor

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-07 15:32:35 -05:00
parent a9d8505cdb
commit 582c5dceed
14 changed files with 164 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
internal/modules/extract/dnd/spells/extractor.go
View File

@@ -1,6 +1,7 @@
package spells
import (
"bytes"
"context"
"encoding/json"
"fmt"
@@ -89,6 +90,10 @@ func (e *Extractor) Extract(ctx context.Context, req contracts.ExtractionRequest
if req.LLMClient == nil {
return contracts.ExtractionResult{}, extractorErrorf("LLM client must not be nil")
}
sourceInput, err := chunkSourceInput(req)
if err != nil {
return contracts.ExtractionResult{}, err
}
var response extractionResponse
completion, err := req.LLMClient.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
@@ -97,7 +102,7 @@ func (e *Extractor) Extract(ctx context.Context, req contracts.ExtractionRequest
PromptVersion: SchemaVersion,
ProfileID: req.LLMProfile,
SessionID: req.SessionID,
Inputs: dnd.PromptInputs(req.SourceInput, req.References),
Inputs: dnd.PromptInputs(sourceInput, req.References),
}, &response)
if err != nil {
return contracts.ExtractionResult{}, extractorErrorf("complete structured output: %w", err)
@@ -128,6 +133,26 @@ func (e *Extractor) Extract(ctx context.Context, req contracts.ExtractionRequest
}, nil
}
func chunkSourceInput(req contracts.ExtractionRequest) (contracts.LLMInputMaterial, error) {
material := req.SourceInput.Clone()
if len(material.Content) == 0 {
material = contracts.NewLLMInputMaterial("source", req.Chunk.MediaType, req.Chunk.Content, "", "")
}
if !bytes.Equal(material.Content, req.Chunk.Content) {
return contracts.LLMInputMaterial{}, extractorErrorf("source input must match chunk %q content", req.Chunk.ID)
}
if material.Name == "" {
material.Name = "source"
}
if material.MediaType == "" {
material.MediaType = req.Chunk.MediaType
}
if material.SizeBytes == 0 {
material.SizeBytes = int64(len(material.Content))
}
return material, nil
}
func ModuleSpec() pipeline.ModuleSpec {
return pipeline.ModuleSpec{
Key: Key,

37
internal/modules/extract/dnd/spells/extractor_test.go
View File

@@ -27,7 +27,8 @@ func TestExtractReturnsRawOutputFromStructuredResponse(t *testing.T) {
content: []byte(`{"spell_casts":[{"caster":" Aria ","spell":" Cure Wounds ","effect":" Heals an injured ally. ","narrative_description":" Aria restores the fighter after the fight. ","source_refs":[{"source_id":"session-alpha","start_unit_id":1,"end_unit_id":2}]}],"raw_marker":true}`),
}
result, err := New().Extract(context.Background(), extractionRequestWithClient(client))
extractReq := extractionRequestWithClient(client)
result, err := New().Extract(context.Background(), extractReq)
if err != nil {
t.Fatalf("Extract() error = %v, want nil", err)
}
@@ -35,22 +36,22 @@ func TestExtractReturnsRawOutputFromStructuredResponse(t *testing.T) {
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)
llmReq := client.requests[0]
if llmReq.StageName != Key {
t.Fatalf("StageName = %q, want %q", llmReq.StageName, Key)
}
if req.PromptID != PromptID || req.PromptVersion != SchemaVersion {
t.Fatalf("prompt = %q/%q, want %q/%q", req.PromptID, req.PromptVersion, PromptID, SchemaVersion)
if llmReq.PromptID != PromptID || llmReq.PromptVersion != SchemaVersion {
t.Fatalf("prompt = %q/%q, want %q/%q", llmReq.PromptID, llmReq.PromptVersion, PromptID, SchemaVersion)
}
if req.SessionID != "session-123" || req.ProfileID != "profile-spells" {
t.Fatalf("session/profile = %q/%q, want session-123/profile-spells", req.SessionID, req.ProfileID)
if llmReq.SessionID != "session-123" || llmReq.ProfileID != "profile-spells" {
t.Fatalf("session/profile = %q/%q, want session-123/profile-spells", llmReq.SessionID, llmReq.ProfileID)
}
transcript := req.Inputs["transcript"]
if transcript.Name != "transcript" || transcript.MediaType != "application/json" || transcript.Digest != "sha256:transcript" || transcript.OriginURI != "file:///session-alpha.json" {
transcript := llmReq.Inputs["transcript"]
if transcript.Name != "transcript" || transcript.MediaType != "application/json" || transcript.Digest != "sha256:chunk" || transcript.OriginURI != "file:///session-alpha.json" {
t.Fatalf("transcript metadata = %#v", transcript)
}
if got := string(transcript.Content); got != spellTranscriptJSON {
t.Fatalf("transcript content = %q, want original source input", got)
if got := string(transcript.Content); got != string(extractReq.Chunk.Content) {
t.Fatalf("transcript content = %q, want chunk content %q", got, extractReq.Chunk.Content)
}
if result.Output.Payload.MediaType != "application/json" {
@@ -227,6 +228,7 @@ func TestExtractRejectsInvalidRequests(t *testing.T) {
{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"},
{name: "source input mismatches chunk", extractor: New(), ctx: context.Background(), req: mismatchedSourceInputRequest(validReq), want: "must match chunk"},
}
for _, tt := range tests {
@@ -310,7 +312,7 @@ func TestExtractDefensivelyCopiesRawContent(t *testing.T) {
func extractionRequestWithClient(client contracts.StructuredLLMClient) contracts.ExtractionRequest {
req := promptExtractionRequest()
req.LLMClient = client
req.SourceInput = spellSourceInput()
req.SourceInput = spellChunkInput(req.Chunk)
req.SessionID = "session-123"
req.LLMProfile = "profile-spells"
return req
@@ -322,6 +324,10 @@ func spellSourceInput() contracts.LLMInputMaterial {
return contracts.NewLLMInputMaterial("source", "application/json", []byte(spellTranscriptJSON), "sha256:transcript", "file:///session-alpha.json")
}
func spellChunkInput(chunk *contracts.SourceChunk) contracts.LLMInputMaterial {
return contracts.NewLLMInputMaterial("source", chunk.MediaType, chunk.Content, "sha256:chunk", "file:///session-alpha.json")
}
func emptyChunkRequest(req contracts.ExtractionRequest) contracts.ExtractionRequest {
req.Chunk = &contracts.SourceChunk{
ID: req.Chunk.ID,
@@ -331,6 +337,11 @@ func emptyChunkRequest(req contracts.ExtractionRequest) contracts.ExtractionRequ
return req
}
func mismatchedSourceInputRequest(req contracts.ExtractionRequest) contracts.ExtractionRequest {
req.SourceInput = spellSourceInput()
return req
}
type fakeSpellsLLMClient struct {
response extractionResponse
content []byte