Cleanup and complete the pipeline refactor

This commit is contained in:
2026-07-07 15:32:35 -05:00
parent a9d8505cdb
commit 582c5dceed
14 changed files with 164 additions and 160 deletions

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