247 lines
8.3 KiB
Go
247 lines
8.3 KiB
Go
package spells
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
|
|
)
|
|
|
|
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 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.",
|
|
}
|
|
|
|
var _ contracts.Extractor[dnd.SpellList] = (*Extractor)(nil)
|
|
|
|
type Options struct{}
|
|
|
|
type Extractor struct {
|
|
llm contracts.StructuredLLMClient
|
|
}
|
|
|
|
type rawAdapter struct {
|
|
extractor *Extractor
|
|
codec RawAdapterCodec
|
|
}
|
|
|
|
type RawAdapterCodec interface {
|
|
contracts.ArtifactCodec[dnd.SpellList]
|
|
EncodeCandidate(dnd.SpellList) ([]byte, error)
|
|
}
|
|
|
|
var _ contracts.LegacyRawExtractor = (*rawAdapter)(nil)
|
|
|
|
func New(llmClient contracts.StructuredLLMClient, _ Options) (*Extractor, error) {
|
|
if llmClient == nil {
|
|
return nil, extractorErrorf("LLM client must not be nil")
|
|
}
|
|
return &Extractor{llm: llmClient}, nil
|
|
}
|
|
|
|
func (e *Extractor) Key() string {
|
|
return Key
|
|
}
|
|
|
|
func (e *Extractor) ReferenceSlots() []contracts.ReferenceSlot {
|
|
return shared.ReferenceSlots(referenceSlotDescriptions)
|
|
}
|
|
|
|
func (e *Extractor) ManifestMetadata() map[string]any {
|
|
promptSHA, err := scriptoriumPromptMetadata()
|
|
if err != nil {
|
|
promptSHA = ""
|
|
}
|
|
metadata := map[string]any{
|
|
"prompt_id": PromptID,
|
|
"prompt_version": SchemaVersion,
|
|
"prompt_sha256": promptSHA,
|
|
"response_schema_key": string(ResponseSchemaKey),
|
|
"response_schema_id": ResponseSchemaID,
|
|
"response_schema_name": ResponseSchemaName,
|
|
}
|
|
if schema, err := loadResponseSchema(); err == nil {
|
|
metadata["response_schema_version"] = schema.Version
|
|
metadata["response_schema_sha256"] = schema.SHA256
|
|
}
|
|
return metadata
|
|
}
|
|
|
|
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
|
|
if _, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
|
StageName: Key,
|
|
PromptID: PromptID,
|
|
PromptVersion: SchemaVersion,
|
|
ProfileID: req.LLMProfile,
|
|
SessionID: req.SessionID,
|
|
Inputs: shared.PromptInputs(sourceInput, req.References),
|
|
}, &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: shared.ReferenceSlots(referenceSlotDescriptions),
|
|
}
|
|
}
|
|
|
|
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)
|
|
})
|
|
}
|
|
|
|
// RegisterWithRawAdapter keeps existing raw downstream implementations usable
|
|
// while the extractor itself produces the canonical typed artifact.
|
|
func RegisterWithRawAdapter(registry *pipeline.ExtractorRegistry, codec RawAdapterCodec) error {
|
|
if codec == nil {
|
|
return extractorErrorf("artifact codec must not be nil")
|
|
}
|
|
build := func(request pipeline.BuildRequest) (*Extractor, error) {
|
|
options, err := DecodeOptions(request.Options)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return New(request.Dependencies.LLM, options)
|
|
}
|
|
return pipeline.RegisterExtractorBuilderWithRawAdapter(registry, ModuleSpec(), validateOptions,
|
|
func(request pipeline.BuildRequest) (contracts.Extractor[dnd.SpellList], error) {
|
|
return build(request)
|
|
},
|
|
func(request pipeline.BuildRequest) (contracts.LegacyRawExtractor, error) {
|
|
extractor, err := build(request)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &rawAdapter{extractor: extractor, codec: codec}, nil
|
|
},
|
|
)
|
|
}
|
|
|
|
func (adapter *rawAdapter) Key() string { return Key }
|
|
|
|
func (adapter *rawAdapter) ReferenceSlots() []contracts.ReferenceSlot {
|
|
return adapter.extractor.ReferenceSlots()
|
|
}
|
|
|
|
func (adapter *rawAdapter) ManifestMetadata() map[string]any {
|
|
return adapter.extractor.ManifestMetadata()
|
|
}
|
|
|
|
func (adapter *rawAdapter) Extract(ctx context.Context, req contracts.ExtractionRequest) (contracts.ExtractionResult, error) {
|
|
result, err := adapter.extractor.Extract(ctx, contracts.TypedExtractionRequest{
|
|
Source: req.Source, Chunk: req.Chunk, AmbientContext: req.AmbientContext,
|
|
SourceInput: req.SourceInput, SessionID: req.SessionID, References: req.References,
|
|
LLMProfile: req.LLMProfile, Metadata: req.Metadata,
|
|
})
|
|
if err != nil {
|
|
return contracts.ExtractionResult{}, err
|
|
}
|
|
content, err := adapter.codec.EncodeCandidate(result.Value)
|
|
if err != nil {
|
|
return contracts.ExtractionResult{}, extractorErrorf("encode canonical output: %w", err)
|
|
}
|
|
schema := adapter.codec.Schema()
|
|
return contracts.ExtractionResult{
|
|
Output: contracts.ExtractOutput{
|
|
Schema: contracts.ResponseSchema{ID: schema.ID, Name: schema.Name, Version: schema.Version, JSONSchema: append([]byte(nil), schema.JSONSchema...)},
|
|
Payload: contracts.RawPayload{Content: content, MediaType: adapter.codec.MediaType(), Metadata: map[string]any{"spell_cast_count": len(result.Value.SpellCasts)}},
|
|
},
|
|
Warnings: result.Warnings,
|
|
}, nil
|
|
}
|
|
|
|
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...)
|
|
}
|