73 lines
2.3 KiB
Go
73 lines
2.3 KiB
Go
package integration_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
|
|
)
|
|
|
|
type extractionResponse struct {
|
|
SpellCasts []spellCastResponse `json:"spell_casts"`
|
|
}
|
|
|
|
type spellCastResponse struct {
|
|
Caster string `json:"caster"`
|
|
Spell string `json:"spell"`
|
|
Effect string `json:"effect"`
|
|
NarrativeDescription string `json:"narrative_description"`
|
|
SourceRefs []shared.SourceRefResponse `json:"source_refs"`
|
|
}
|
|
|
|
type fakeSpellsLLMClient struct {
|
|
response extractionResponse
|
|
responses []extractionResponse
|
|
requests []contracts.StructuredCompletionRequest
|
|
}
|
|
|
|
func (client *fakeSpellsLLMClient) CompleteStructured(_ context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
|
|
client.requests = append(client.requests, cloneStructuredCompletionRequest(req))
|
|
response := client.response
|
|
if client.responses != nil {
|
|
index := len(client.requests) - 1
|
|
if index >= len(client.responses) {
|
|
return contracts.StructuredCompletionResponse{}, fmt.Errorf("missing fake response %d", index)
|
|
}
|
|
response = client.responses[index]
|
|
}
|
|
content, err := json.Marshal(response)
|
|
if err != nil {
|
|
return contracts.StructuredCompletionResponse{}, err
|
|
}
|
|
if err := json.Unmarshal(content, out); err != nil {
|
|
return contracts.StructuredCompletionResponse{}, fmt.Errorf("populate structured target: %w", err)
|
|
}
|
|
return contracts.StructuredCompletionResponse{Content: content}, nil
|
|
}
|
|
|
|
func responseSourceRefs(sourceID string, startUnitID int, endUnitID int) []shared.SourceRefResponse {
|
|
return []shared.SourceRefResponse{
|
|
{
|
|
SourceID: sourceID,
|
|
StartUnitID: shared.UnitRefFromInt(startUnitID),
|
|
EndUnitID: shared.UnitRefFromInt(endUnitID),
|
|
},
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|