259 lines
9.0 KiB
Go
259 lines
9.0 KiB
Go
package spells
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"sort"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
|
|
npcregistry "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/npcs/registry"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
|
|
spellcatalog "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/spells/catalog"
|
|
)
|
|
|
|
const Key = "dnd/spells"
|
|
const ArtifactType = "dnd.spell_cast"
|
|
const SchemaVersion = "v1"
|
|
|
|
const (
|
|
NPCRegistryReferenceSlot = npcregistry.ReferenceSlot
|
|
NPCRegistryMaxBytes = npcregistry.MaxBytes
|
|
)
|
|
|
|
var requiredCapabilities = []string{
|
|
"chunks",
|
|
"source.transcript",
|
|
}
|
|
|
|
var providedCapabilities = []string{
|
|
"dnd.spell_casts",
|
|
}
|
|
|
|
var referenceSlotDescriptions = shared.ReferenceSlotDescriptions{
|
|
Glossary: "Optional campaign glossary reference material used only for disambiguation.",
|
|
Party: "Optional party roster reference material used only for disambiguation.",
|
|
Players: "Optional player list reference material used only for disambiguation.",
|
|
Roster: "Deprecated alias for party roster reference material used only for disambiguation.",
|
|
}
|
|
|
|
func referenceSlots() []contracts.ReferenceSlot {
|
|
slots := shared.ReferenceSlots(referenceSlotDescriptions)
|
|
slots = append(slots, contracts.ReferenceSlot{
|
|
Name: spellcatalog.SpellCatalogReferenceSlot,
|
|
Description: "Optional canonical spell-name catalog used for extraction grounding.",
|
|
AcceptedMediaTypes: []string{"application/json"},
|
|
MaxBytes: 1048576,
|
|
}, contracts.ReferenceSlot{
|
|
Name: NPCRegistryReferenceSlot,
|
|
Description: "Optional normalized NPC registry used for canonical caster-name grounding.",
|
|
AcceptedMediaTypes: []string{"application/json"},
|
|
MaxBytes: NPCRegistryMaxBytes,
|
|
})
|
|
sort.Slice(slots, func(i, j int) bool { return slots[i].Name < slots[j].Name })
|
|
return slots
|
|
}
|
|
|
|
var _ contracts.Extractor[dnd.SpellList] = (*Extractor)(nil)
|
|
var _ pipeline.CheckpointFingerprintProvider = (*Extractor)(nil)
|
|
|
|
type Options struct{}
|
|
|
|
type Extractor struct {
|
|
llm contracts.StructuredLLMClient
|
|
effectiveCatalog spellcatalog.EffectiveCatalog
|
|
catalogPromptInput contracts.LLMInputMaterial
|
|
npcRegistry *npcregistry.Registry
|
|
promptSHA string
|
|
responseSchemaSHA string
|
|
}
|
|
|
|
func New(llmClient contracts.StructuredLLMClient, _ Options, references ...contracts.ReferenceSet) (*Extractor, error) {
|
|
if llmClient == nil {
|
|
return nil, extractorErrorf("LLM client must not be nil")
|
|
}
|
|
if len(references) > 1 {
|
|
return nil, extractorErrorf("at most one reference set may be supplied")
|
|
}
|
|
var referenceSet contracts.ReferenceSet
|
|
if len(references) == 1 {
|
|
referenceSet = references[0]
|
|
}
|
|
effectiveCatalog, err := spellcatalog.ResolveEffectiveCatalog(referenceSet)
|
|
if err != nil {
|
|
return nil, extractorErrorf("resolve effective spell catalog: %w", err)
|
|
}
|
|
catalogPromptInput, err := newCatalogPromptInput(effectiveCatalog)
|
|
if err != nil {
|
|
return nil, extractorErrorf("prepare spell catalog prompt input: %w", err)
|
|
}
|
|
npcRegistry, err := npcregistry.Resolve(referenceSet)
|
|
if err != nil {
|
|
return nil, extractorErrorf("prepare NPC registry prompt input: %w", err)
|
|
}
|
|
promptSHA, err := scriptoriumPromptMetadata()
|
|
if err != nil {
|
|
return nil, extractorErrorf("load prompt metadata: %w", err)
|
|
}
|
|
responseSchema, err := loadResponseSchema()
|
|
if err != nil {
|
|
return nil, extractorErrorf("load response schema: %w", err)
|
|
}
|
|
return &Extractor{
|
|
llm: llmClient,
|
|
effectiveCatalog: effectiveCatalog,
|
|
catalogPromptInput: catalogPromptInput,
|
|
npcRegistry: npcRegistry,
|
|
promptSHA: promptSHA,
|
|
responseSchemaSHA: responseSchema.SHA256,
|
|
}, nil
|
|
}
|
|
|
|
func (e *Extractor) Key() string {
|
|
return Key
|
|
}
|
|
|
|
func (e *Extractor) ReferenceSlots() []contracts.ReferenceSlot {
|
|
return referenceSlots()
|
|
}
|
|
|
|
func (e *Extractor) ManifestMetadata() map[string]any {
|
|
metadata := map[string]any{
|
|
"prompt_id": PromptID,
|
|
"prompt_version": SchemaVersion,
|
|
"prompt_sha256": e.promptSHA,
|
|
"catalog_base_id": e.effectiveCatalog.BaseID(),
|
|
"catalog_digest": e.effectiveCatalog.Digest(),
|
|
"catalog_overlay_ids": e.effectiveCatalog.OverlayIDs(),
|
|
"response_schema_key": string(ResponseSchemaKey),
|
|
"response_schema_id": ResponseSchemaID,
|
|
"response_schema_name": ResponseSchemaName,
|
|
"response_schema_version": SchemaVersion,
|
|
"response_schema_sha256": e.responseSchemaSHA,
|
|
}
|
|
if e.npcRegistry.Bound() {
|
|
metadata["npc_registry_digest"] = e.npcRegistry.Digest()
|
|
metadata["npc_count"] = e.npcRegistry.Count()
|
|
}
|
|
return metadata
|
|
}
|
|
|
|
func (e *Extractor) CheckpointFingerprints() []pipeline.CheckpointFingerprint {
|
|
if e == nil {
|
|
return nil
|
|
}
|
|
fingerprints := []pipeline.CheckpointFingerprint{
|
|
{Name: "effective_catalog", Value: e.effectiveCatalog.Digest()},
|
|
{Name: "prompt", Value: e.promptSHA},
|
|
{Name: "response_schema", Value: e.responseSchemaSHA},
|
|
}
|
|
if e.npcRegistry.Bound() {
|
|
fingerprints = append(fingerprints, pipeline.CheckpointFingerprint{Name: "npc_registry", Value: e.npcRegistry.Digest()})
|
|
}
|
|
return fingerprints
|
|
}
|
|
|
|
func (e *Extractor) Extract(ctx context.Context, req contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[dnd.SpellList], error) {
|
|
if e == nil {
|
|
return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("extractor must not be nil")
|
|
}
|
|
if e.llm == nil {
|
|
return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("LLM client must not be nil")
|
|
}
|
|
if ctx == nil {
|
|
return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("context must not be nil")
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("context error before extraction: %w", err)
|
|
}
|
|
if req.Source == nil {
|
|
return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("source must not be nil")
|
|
}
|
|
if req.Chunk == nil {
|
|
return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("chunk must not be nil")
|
|
}
|
|
if len(req.Chunk.Units) == 0 {
|
|
return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("chunk %q units must not be empty", req.Chunk.ID)
|
|
}
|
|
sourceInput, err := chunkSourceInput(req)
|
|
if err != nil {
|
|
return contracts.TypedExtractionResult[dnd.SpellList]{}, err
|
|
}
|
|
|
|
var response extractionResponse
|
|
inputs := shared.PromptInputs(sourceInput, req.References)
|
|
inputs[spellcatalog.SpellCatalogReferenceSlot] = e.catalogPromptInput.Clone()
|
|
inputs[NPCRegistryReferenceSlot] = e.npcRegistry.PromptInput()
|
|
if _, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
|
StageName: Key,
|
|
PromptID: PromptID,
|
|
PromptVersion: SchemaVersion,
|
|
ProfileID: req.LLMProfile,
|
|
SessionID: req.SessionID,
|
|
Inputs: inputs,
|
|
}, &response); err != nil {
|
|
return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("complete structured output: %w", err)
|
|
}
|
|
canonicalizeResponse(&response)
|
|
return contracts.TypedExtractionResult[dnd.SpellList]{Value: canonicalSpellList(response, req.Source.ID)}, nil
|
|
}
|
|
|
|
func chunkSourceInput(req contracts.TypedExtractionRequest) (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,
|
|
Stage: pipeline.StageExtract,
|
|
Requires: append([]string(nil), requiredCapabilities...),
|
|
Provides: append([]string(nil), providedCapabilities...),
|
|
ArtifactKind: dnd.SpellListKind,
|
|
ReferenceSlots: referenceSlots(),
|
|
}
|
|
}
|
|
|
|
func Register(registry *pipeline.ExtractorRegistry) error {
|
|
return pipeline.RegisterExtractorBuilder(registry, ModuleSpec(), validateOptions, func(request pipeline.BuildRequest) (contracts.Extractor[dnd.SpellList], error) {
|
|
options, err := DecodeOptions(request.Options)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return New(request.Dependencies.LLM, options, request.References)
|
|
})
|
|
}
|
|
|
|
func validateOptions(options map[string]any) error {
|
|
_, err := DecodeOptions(options)
|
|
return err
|
|
}
|
|
|
|
func DecodeOptions(options map[string]any) (Options, error) {
|
|
if err := pipeline.RejectUnknownOptions(options); err != nil {
|
|
return Options{}, extractorErrorf("%w", err)
|
|
}
|
|
return Options{}, nil
|
|
}
|
|
|
|
func extractorErrorf(format string, args ...any) error {
|
|
return fmt.Errorf("dnd spells extractor: "+format, args...)
|
|
}
|