145 lines
4.3 KiB
Go
145 lines
4.3 KiB
Go
package spells
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
|
)
|
|
|
|
const Key = "dnd/spells"
|
|
const ArtifactType = "dnd.spell_cast"
|
|
const SchemaVersion = "v1"
|
|
|
|
var requiredCapabilities = []string{
|
|
"chunks",
|
|
"source.transcript",
|
|
}
|
|
|
|
var providedCapabilities = []string{
|
|
"dnd.spell_casts",
|
|
}
|
|
|
|
var _ contracts.Extractor = (*Extractor)(nil)
|
|
|
|
type Extractor struct{}
|
|
|
|
func New() *Extractor {
|
|
return &Extractor{}
|
|
}
|
|
|
|
func (e *Extractor) Key() string {
|
|
return Key
|
|
}
|
|
|
|
func (e *Extractor) ArtifactType() string {
|
|
return ArtifactType
|
|
}
|
|
|
|
func (e *Extractor) SchemaVersion() string {
|
|
return SchemaVersion
|
|
}
|
|
|
|
func (e *Extractor) Validators() []contracts.Validator {
|
|
return nil
|
|
}
|
|
|
|
func (e *Extractor) Extract(ctx context.Context, req contracts.ExtractionRequest) (contracts.ExtractionResult, error) {
|
|
if e == nil {
|
|
return contracts.ExtractionResult{}, extractorErrorf("extractor must not be nil")
|
|
}
|
|
if ctx == nil {
|
|
return contracts.ExtractionResult{}, extractorErrorf("context must not be nil")
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return contracts.ExtractionResult{}, extractorErrorf("context error before extraction: %w", err)
|
|
}
|
|
if req.Source == nil {
|
|
return contracts.ExtractionResult{}, extractorErrorf("source must not be nil")
|
|
}
|
|
if req.Chunk == nil {
|
|
return contracts.ExtractionResult{}, extractorErrorf("chunk must not be nil")
|
|
}
|
|
if len(req.Chunk.Units) == 0 {
|
|
return contracts.ExtractionResult{}, extractorErrorf("chunk %q units must not be empty", req.Chunk.ID)
|
|
}
|
|
if req.LLMClient == nil {
|
|
return contracts.ExtractionResult{}, extractorErrorf("LLM client must not be nil")
|
|
}
|
|
|
|
system, user, _, err := renderPrompt(req)
|
|
if err != nil {
|
|
return contracts.ExtractionResult{}, extractorErrorf("render prompt: %w", err)
|
|
}
|
|
schema, ok := llm.LookupResponseSchema(llm.DNDSpellsSchemaKey)
|
|
if !ok {
|
|
return contracts.ExtractionResult{}, extractorErrorf("lookup response schema %q", llm.DNDSpellsSchemaKey)
|
|
}
|
|
|
|
var response extractionResponse
|
|
if _, err := req.LLMClient.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
|
StageName: Key,
|
|
Messages: []contracts.LLMMessage{
|
|
{Role: "system", Content: system},
|
|
{Role: "user", Content: user},
|
|
},
|
|
ResponseSchemaName: schema.Name,
|
|
ResponseSchema: schema.JSONSchema,
|
|
}, &response); err != nil {
|
|
return contracts.ExtractionResult{}, extractorErrorf("complete structured output: %w", err)
|
|
}
|
|
if response.SpellCasts == nil {
|
|
return contracts.ExtractionResult{}, extractorErrorf("malformed structured output: spell_casts must be present")
|
|
}
|
|
if len(response.SpellCasts) == 0 {
|
|
return contracts.ExtractionResult{}, nil
|
|
}
|
|
|
|
candidates := make([]artifacts.ArtifactCandidate, 0, len(response.SpellCasts))
|
|
for i, spellCast := range response.SpellCasts {
|
|
payload, err := spellCastPayload(spellCast)
|
|
if err != nil {
|
|
return contracts.ExtractionResult{}, extractorErrorf("marshal spell cast[%d]: %w", i, err)
|
|
}
|
|
candidates = append(candidates, artifacts.ArtifactCandidate{
|
|
Payload: payload,
|
|
SourceRefs: append([]source.SourceRef(nil), spellCast.SourceRefs...),
|
|
})
|
|
}
|
|
return contracts.ExtractionResult{Candidates: candidates}, nil
|
|
}
|
|
|
|
func ModuleSpec() pipeline.ModuleSpec {
|
|
return pipeline.ModuleSpec{
|
|
Key: Key,
|
|
Stage: pipeline.StageExtract,
|
|
Requires: append([]string(nil), requiredCapabilities...),
|
|
Provides: append([]string(nil), providedCapabilities...),
|
|
}
|
|
}
|
|
|
|
func Register(registry *pipeline.ExtractorRegistry) error {
|
|
return registry.RegisterWithSpec(ModuleSpec(), func() (contracts.Extractor, error) {
|
|
return New(), nil
|
|
})
|
|
}
|
|
|
|
func spellCastPayload(spellCast spellCastResponse) (json.RawMessage, error) {
|
|
return json.Marshal(SpellCast{
|
|
Caster: strings.TrimSpace(spellCast.Caster),
|
|
Spell: strings.TrimSpace(spellCast.Spell),
|
|
Effect: strings.TrimSpace(spellCast.Effect),
|
|
NarrativeDescription: strings.TrimSpace(spellCast.NarrativeDescription),
|
|
})
|
|
}
|
|
|
|
func extractorErrorf(format string, args ...any) error {
|
|
return fmt.Errorf("dnd spells extractor: "+format, args...)
|
|
}
|