notarius/internal/modules/dnd/extract/combatturns/extractor.go

package combatturns

import (
	"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"
	sceneregistry "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/scenedescriptions/registry"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
)

const (
	Key             = "dnd/combat-turns"
	mappingPolicy   = "dnd.combat_turns.extract_mapping.v3"
	sceneGatePolicy = "dnd.combat_turns.scene_gate.v1"
)

const (
	NPCRegistryReferenceSlot      = npcregistry.ReferenceSlot
	NPCRegistryMaxBytes           = npcregistry.MaxBytes
	SceneDescriptionReferenceSlot = sceneregistry.ReferenceSlot
	SceneDescriptionMaxBytes      = sceneregistry.MaxBytes
)

var requiredCapabilities = []string{
	"chunks",
	"source.transcript",
}

var providedCapabilities = []string{
	"dnd.combat_turns",
}

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:                  NPCRegistryReferenceSlot,
		Description:           "Optional normalized NPC registry used for canonical actor grounding.",
		AcceptedMediaTypes:    []string{"application/json"},
		AcceptedArtifactKinds: []contracts.ArtifactKind{dnd.NPCRegistryKind},
		MaxBytes:              NPCRegistryMaxBytes,
	})
	slots = append(slots, contracts.ReferenceSlot{
		Name:                  SceneDescriptionReferenceSlot,
		Description:           "Required scene descriptions that gate combat eligibility and are not combat evidence.",
		Required:              true,
		AcceptedMediaTypes:    []string{"application/json"},
		AcceptedArtifactKinds: []contracts.ArtifactKind{dnd.SceneDescriptionListKind},
		MaxBytes:              SceneDescriptionMaxBytes,
	})
	sort.Slice(slots, func(i, j int) bool { return slots[i].Name < slots[j].Name })
	return slots
}

var _ contracts.Extractor[dnd.CombatTurnList] = (*Extractor)(nil)
var _ contracts.ManifestMetadataProvider = (*Extractor)(nil)
var _ pipeline.CheckpointFingerprintProvider = (*Extractor)(nil)

type Options struct{}

type Extractor struct {
	llm               contracts.StructuredLLMClient
	npcResolver       *npcregistry.Resolver
	sceneResolver     *sceneregistry.Resolver
	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]
	}
	npcResolver, err := npcregistry.NewResolver(referenceSet)
	if err != nil {
		return nil, extractorErrorf("prepare NPC registry prompt input: %w", err)
	}
	sceneResolver, err := sceneregistry.NewResolver(referenceSet)
	if err != nil {
		return nil, extractorErrorf("prepare scene eligibility: %w", err)
	}
	promptSHA, err := promptAssetMetadata()
	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,
		npcResolver:       npcResolver,
		sceneResolver:     sceneResolver,
		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 {
	if e == nil {
		return nil
	}
	metadata := map[string]any{
		"prompt_id":               PromptID,
		"prompt_version":          SchemaVersion,
		"prompt_sha256":           e.promptSHA,
		"mapping_policy":          mappingPolicy,
		"scene_gate_policy":       sceneGatePolicy,
		"response_schema_key":     string(ResponseSchemaKey),
		"response_schema_id":      ResponseSchemaID,
		"response_schema_name":    ResponseSchemaName,
		"response_schema_version": SchemaVersion,
		"response_schema_sha256":  e.responseSchemaSHA,
	}
	seeded := e.npcResolver.Seeded()
	if seeded.Bound() {
		metadata["npc_registry_digest"] = seeded.Digest()
		metadata["npc_count"] = seeded.Count()
	}
	seededScenes := e.sceneResolver.Seeded()
	if seededScenes.Bound() {
		metadata["scene_eligibility_digest"] = seededScenes.EligibilityDigest()
		metadata["scene_description_count"] = seededScenes.Count()
	}
	return metadata
}

func (e *Extractor) CheckpointFingerprints() []pipeline.CheckpointFingerprint {
	if e == nil {
		return nil
	}
	fingerprints := []pipeline.CheckpointFingerprint{
		{Name: "prompt", Value: e.promptSHA},
		{Name: "response_schema", Value: e.responseSchemaSHA},
		{Name: "mapping_policy", Value: mappingPolicy},
		{Name: "scene_gate_policy", Value: sceneGatePolicy},
	}
	seeded := e.npcResolver.Seeded()
	fingerprints = append(fingerprints, pipeline.CheckpointFingerprint{Name: "npc_registry", Value: seeded.ProjectionDigest()})
	seededScenes := e.sceneResolver.Seeded()
	fingerprints = append(fingerprints, pipeline.CheckpointFingerprint{Name: "scene_eligibility", Value: seededScenes.EligibilityDigest()})
	return fingerprints
}

func (e *Extractor) Extract(ctx context.Context, req contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[dnd.CombatTurnList], error) {
	if e == nil {
		return contracts.TypedExtractionResult[dnd.CombatTurnList]{}, extractorErrorf("extractor must not be nil")
	}
	if e.llm == nil {
		return contracts.TypedExtractionResult[dnd.CombatTurnList]{}, extractorErrorf("LLM client must not be nil")
	}
	sourceInput, err := shared.PrepareChunkExtraction(ctx, req)
	if err != nil {
		return contracts.TypedExtractionResult[dnd.CombatTurnList]{}, extractorErrorf("%w", err)
	}
	if _, ok := req.References.Slots[SceneDescriptionReferenceSlot]; !ok {
		return contracts.TypedExtractionResult[dnd.CombatTurnList]{}, extractorErrorf("required scene descriptions reference is missing")
	}
	scenes, err := e.sceneResolver.Resolve(req.References)
	if err != nil {
		return contracts.TypedExtractionResult[dnd.CombatTurnList]{}, extractorErrorf("resolve scene eligibility: %w", err)
	}
	match := scenes.Match(req.Chunk)
	switch match.State {
	case sceneregistry.MatchExact:
		if match.Kind != dnd.SceneKindCombat {
			return contracts.TypedExtractionResult[dnd.CombatTurnList]{Value: dnd.CombatTurnList{CombatTurns: []dnd.CombatTurn{}}}, nil
		}
	case sceneregistry.MatchMissing, sceneregistry.MatchMismatched:
		return contracts.TypedExtractionResult[dnd.CombatTurnList]{
			Value: dnd.CombatTurnList{CombatTurns: []dnd.CombatTurn{}},
			Diagnostics: []contracts.ProducerDiagnostic{{
				Disposition:     contracts.DiagnosticDispositionWarning,
				Category:        contracts.DiagnosticCategoryDegradation,
				ReasonCode:      "scene_classification_unavailable",
				OccurrenceCount: 1,
				Samples:         []contracts.DiagnosticSample{{Scope: SceneDescriptionReferenceSlot, Message: "No exact scene classification was available; combat extraction was skipped."}},
			}},
		}, nil
	default:
		return contracts.TypedExtractionResult[dnd.CombatTurnList]{}, extractorErrorf("unsupported scene eligibility match state %q", match.State)
	}
	order := shared.NewSourceRefOrder(req.Source)
	npcRegistry, err := e.npcResolver.Resolve(req.References)
	if err != nil {
		return contracts.TypedExtractionResult[dnd.CombatTurnList]{}, extractorErrorf("resolve NPC registry: %w", err)
	}

	var response extractionResponse
	inputs := shared.PromptInputs(sourceInput, req.References)
	inputs[NPCRegistryReferenceSlot] = npcRegistry.PromptInput()
	completion, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
		StageName:                      Key,
		PromptID:                       PromptID,
		PromptVersion:                  SchemaVersion,
		ProfileID:                      req.LLMProfile,
		SessionID:                      req.SessionID,
		StructuredOutputRepairAttempts: req.StructuredOutputRepairAttempts,
		Correction:                     req.Correction,
		Inputs:                         inputs,
	}, &response)
	if err != nil {
		return contracts.TypedExtractionResult[dnd.CombatTurnList]{}, extractorErrorf("complete structured output: %w", err)
	}
	candidate, err := shared.ModelCandidateFromResponse(completion)
	if err != nil {
		return contracts.TypedExtractionResult[dnd.CombatTurnList]{}, extractorErrorf("capture model candidate: %w", err)
	}
	canonicalizeResponse(&response, order, req.Source.ID)
	return contracts.TypedExtractionResult[dnd.CombatTurnList]{Value: canonicalCombatTurnList(response, req.Source.ID), ModelCandidate: candidate}, nil
}

func ModuleSpec() pipeline.ModuleSpec {
	return pipeline.ModuleSpec{
		Key:                Key,
		Stage:              pipeline.StageExtract,
		ExecutionClass:     contracts.ExecutionClassLLMBacked,
		CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1,
		Requires:           append([]string(nil), requiredCapabilities...),
		Provides:           append([]string(nil), providedCapabilities...),
		ArtifactKind:       dnd.CombatTurnListKind,
		ReferenceSlots:     referenceSlots(),
	}
}

func Register(registry *pipeline.ExtractorRegistry) error {
	return pipeline.RegisterExtractorBuilder(registry, ModuleSpec(), validateOptions, func(request pipeline.BuildRequest) (contracts.Extractor[dnd.CombatTurnList], 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 combat turns extractor: "+format, args...)
}