111 lines
3.5 KiB
Go
111 lines
3.5 KiB
Go
package npcs
|
|
|
|
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-alpha:chunk:0",
|
|
SourceID: doc.ID,
|
|
Index: 0,
|
|
Ref: source.SourceRef{SourceID: doc.ID, StartUnitID: 1, EndUnitID: 3},
|
|
Content: []byte(`{"units":[1,2,3]}`),
|
|
MediaType: "application/json",
|
|
Units: append([]source.SourceUnit(nil), doc.Units...),
|
|
Metadata: map[string]any{"ignored": "chunk metadata"},
|
|
}
|
|
return contracts.TypedExtractionRequest{
|
|
Source: doc,
|
|
Chunk: chunk,
|
|
SourceInput: contracts.NewLLMInputMaterial("source", chunk.MediaType, chunk.Content, "sha256:chunk", "file:///session-alpha.json"),
|
|
SessionID: "session-123",
|
|
LLMProfile: "profile-npcs",
|
|
}
|
|
}
|
|
|
|
func sourceDocument() *source.SourceDocument {
|
|
return &source.SourceDocument{
|
|
ID: "session-alpha",
|
|
Kind: "transcript",
|
|
Format: "application/vnd.seriatim.minimal+json",
|
|
Digest: "sha256:test",
|
|
Units: []source.SourceUnit{
|
|
{ID: 1, Kind: "transcript_segment", Text: "Mira Thorn greets the party."},
|
|
{ID: 2, Kind: "transcript_segment", Text: "The Greencloak watches the northern road."},
|
|
{ID: 3, Kind: "transcript_segment", Text: "Captain Vale reports to Mira Thorn."},
|
|
},
|
|
}
|
|
}
|
|
|
|
func responseSourceRefs(startUnitID, endUnitID int) []npcSourceRefResponse {
|
|
return []npcSourceRefResponse{{StartUnitID: startUnitID, EndUnitID: endUnitID}}
|
|
}
|
|
|
|
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, want nil", err)
|
|
}
|
|
return extractor
|
|
}
|
|
|
|
func mismatchedSourceInputRequest(req contracts.TypedExtractionRequest) contracts.TypedExtractionRequest {
|
|
req.SourceInput = contracts.NewLLMInputMaterial("source", "application/json", []byte(`{"different":true}`), "sha256:other", "file:///other.json")
|
|
return req
|
|
}
|
|
|
|
func cloneStructuredCompletionRequest(req contracts.StructuredCompletionRequest) contracts.StructuredCompletionRequest {
|
|
req.Inputs = req.Inputs.Clone()
|
|
if len(req.Vars) == 0 {
|
|
req.Vars = nil
|
|
return req
|
|
}
|
|
vars := make(map[string]any, len(req.Vars))
|
|
for key, value := range req.Vars {
|
|
vars[key] = value
|
|
}
|
|
req.Vars = vars
|
|
return req
|
|
}
|
|
|
|
type fakeNPCsLLMClient struct {
|
|
response extractionResponse
|
|
content []byte
|
|
err error
|
|
requests []contracts.StructuredCompletionRequest
|
|
}
|
|
|
|
func (client *fakeNPCsLLMClient) 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 := append([]byte(nil), client.content...)
|
|
if len(content) != 0 {
|
|
if err := json.Unmarshal(content, target); err != nil {
|
|
return contracts.StructuredCompletionResponse{}, err
|
|
}
|
|
} else {
|
|
*target = client.response
|
|
var err error
|
|
content, err = json.Marshal(client.response)
|
|
if err != nil {
|
|
return contracts.StructuredCompletionResponse{}, err
|
|
}
|
|
}
|
|
return contracts.StructuredCompletionResponse{Content: content}, nil
|
|
}
|