notarius/internal/modules/dnd/chunk/scenes/chunker.go

package scenes

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
)

const Key = "dnd/scenes"

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

var providedCapabilities = []string{
	"chunks",
}

const annotationNamespace = "dnd/scenes"

var referenceSlotDescriptions = shared.ReferenceSlotDescriptions{
	Glossary: "Optional campaign glossary reference material used only for scene disambiguation.",
	Party:    "Optional party roster reference material used only for scene disambiguation.",
	Players:  "Optional player list reference material used only for scene disambiguation.",
	Roster:   "Deprecated alias for party roster reference material used only for scene disambiguation.",
}

var _ contracts.Chunker = (*Chunker)(nil)
var _ contracts.ManifestMetadataProvider = (*Chunker)(nil)

type Options struct{}

type Chunker struct {
	llm contracts.StructuredLLMClient
}

func New(llmClient contracts.StructuredLLMClient, _ Options) (*Chunker, error) {
	if llmClient == nil {
		return nil, chunkerErrorf("LLM client must not be nil")
	}
	return &Chunker{llm: llmClient}, nil
}

func (c *Chunker) Key() string {
	return Key
}

func (*Chunker) ExecutionClass() contracts.ExecutionClass {
	return contracts.ExecutionClassLLMBacked
}

func (c *Chunker) ReferenceSlots() []contracts.ReferenceSlot {
	return shared.ReferenceSlots(referenceSlotDescriptions)
}

func (c *Chunker) ManifestMetadata() map[string]any {
	promptSHA, err := scriptoriumPromptMetadata()
	if err != nil {
		promptSHA = ""
	}
	metadata := map[string]any{
		"prompt_id":            PromptID,
		"prompt_version":       ResponseSchemaVersion,
		"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 (c *Chunker) Plan(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
	if c == nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("chunker must not be nil")
	}
	if c.llm == nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("LLM client must not be nil")
	}
	if ctx == nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("context must not be nil")
	}
	if err := ctx.Err(); err != nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("context error before chunking: %w", err)
	}
	if req.Source == nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("source must not be nil")
	}
	if len(req.Source.Units) == 0 {
		return contracts.ChunkPlanResult{}, chunkerErrorf("source units must not be empty")
	}
	if err := source.ValidateDocument(req.Source); err != nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("validate source document: %w", err)
	}
	var response chunkResponse
	if _, err := c.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
		StageName:     Key,
		PromptID:      PromptID,
		PromptVersion: ResponseSchemaVersion,
		ProfileID:     req.LLMProfile,
		SessionID:     req.SessionID,
		Inputs:        shared.PromptInputs(req.SourceInput, req.References),
	}, &response); err != nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("complete structured output: %w", err)
	}

	warnings, err := warningsFromCaveats(response.BoundaryCaveats)
	if err != nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("malformed structured output: %w", err)
	}
	plan, err := planFromResponse(req.Source, response, warnings)
	if err != nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("malformed structured output: %w", err)
	}
	return contracts.ChunkPlanResult{
		Plan:     plan,
		Warnings: warnings,
	}, nil
}

func ModuleSpec() pipeline.ModuleSpec {
	return pipeline.ModuleSpec{
		Key:            Key,
		Stage:          pipeline.StageChunk,
		Requires:       append([]string(nil), requiredCapabilities...),
		Provides:       append([]string(nil), providedCapabilities...),
		ReferenceSlots: shared.ReferenceSlots(referenceSlotDescriptions),
	}
}

func Register(registry *pipeline.ChunkerRegistry) error {
	return registry.RegisterBuilderWithSpec(ModuleSpec(), validateOptions, func(request pipeline.BuildRequest) (contracts.Chunker, error) {
		options, err := DecodeOptions(request.Options)
		if err != nil {
			return nil, err
		}
		return New(request.Dependencies.LLM, options)
	})
}

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{}, chunkerErrorf("%w", err)
	}
	return Options{}, nil
}

func planFromResponse(doc *source.SourceDocument, response chunkResponse, warnings []contracts.Warning) (source.ChunkPlan, error) {
	if response.Scenes == nil {
		return source.ChunkPlan{}, fmt.Errorf("scenes must be present")
	}
	if len(response.Scenes) == 0 {
		return source.ChunkPlan{}, fmt.Errorf("scenes must not be empty")
	}

	unitIndexes := make(map[int]int, len(doc.Units))
	for i, unit := range doc.Units {
		unitIndexes[unit.ID] = i
	}

	ranges := make([]source.ChunkRange, 0, len(response.Scenes))
	previousEnd := -1
	for i, scene := range response.Scenes {
		normalized, err := normalizeScene(doc, i, scene)
		if err != nil {
			return source.ChunkPlan{}, err
		}

		startIndex, ok := unitIndexes[normalized.StartUnitID]
		if !ok {
			return source.ChunkPlan{}, fmt.Errorf("scene[%d] start_unit_id %d was not found", i, normalized.StartUnitID)
		}
		endIndex, ok := unitIndexes[normalized.EndUnitID]
		if !ok {
			return source.ChunkPlan{}, fmt.Errorf("scene[%d] end_unit_id %d was not found", i, normalized.EndUnitID)
		}
		if startIndex > endIndex {
			return source.ChunkPlan{}, fmt.Errorf("scene[%d] start_unit_id %d appears after end_unit_id %d", i, normalized.StartUnitID, normalized.EndUnitID)
		}

		if i == 0 && startIndex != 0 {
			return source.ChunkPlan{}, fmt.Errorf("first scene must start at first source unit %d", doc.Units[0].ID)
		}
		if i > 0 {
			if startIndex <= previousEnd {
				return source.ChunkPlan{}, fmt.Errorf("scene[%d] overlaps previous scene", i)
			}
			if startIndex > previousEnd+1 {
				return source.ChunkPlan{}, fmt.Errorf("scene[%d] leaves a gap after previous scene", i)
			}
		}
		previousEnd = endIndex

		annotation, err := json.Marshal(struct {
			ShortTitle         string   `json:"short_title"`
			PrimaryMode        string   `json:"primary_mode"`
			MainParticipants   []string `json:"main_participants"`
			Summary            string   `json:"summary"`
			BoundaryNote       string   `json:"boundary_note"`
			BoundaryConfidence string   `json:"boundary_confidence"`
		}{normalized.ShortTitle, normalized.PrimaryMode, normalized.MainParticipants, normalized.Summary, normalized.BoundaryNote, normalized.BoundaryConfidence})
		if err != nil {
			return source.ChunkPlan{}, fmt.Errorf("encode scene[%d] annotation: %w", i, err)
		}
		ranges = append(ranges, source.ChunkRange{
			StartUnitID: normalized.StartUnitID,
			EndUnitID:   normalized.EndUnitID,
			Annotations: source.ChunkAnnotations{annotationNamespace: annotation},
		})
	}

	if previousEnd != len(doc.Units)-1 {
		return source.ChunkPlan{}, fmt.Errorf("final scene must end at final source unit %d", doc.Units[len(doc.Units)-1].ID)
	}
	caveats := make([]string, 0, len(warnings))
	for _, warning := range warnings {
		caveats = append(caveats, warning.Message)
	}
	annotation, err := json.Marshal(struct {
		BoundaryCaveats []string `json:"boundary_caveats"`
	}{BoundaryCaveats: caveats})
	if err != nil {
		return source.ChunkPlan{}, fmt.Errorf("encode plan annotation: %w", err)
	}
	return source.ChunkPlan{
		SourceDigest: doc.Digest,
		Ranges:       ranges,
		Annotations:  source.ChunkAnnotations{annotationNamespace: annotation},
	}, nil
}

func normalizeScene(doc *source.SourceDocument, index int, scene sceneResponse) (normalizedScene, error) {
	startUnitID, err := shared.ResolveUnitID(doc, "start_unit_id", scene.StartUnitID)
	if err != nil {
		return normalizedScene{}, fmt.Errorf("scene[%d] %w", index, err)
	}
	endUnitID, err := shared.ResolveUnitID(doc, "end_unit_id", scene.EndUnitID)
	if err != nil {
		return normalizedScene{}, fmt.Errorf("scene[%d] %w", index, err)
	}

	out := normalizedScene{
		StartUnitID:        startUnitID,
		EndUnitID:          endUnitID,
		ShortTitle:         strings.TrimSpace(scene.ShortTitle),
		PrimaryMode:        strings.TrimSpace(scene.PrimaryMode),
		Summary:            strings.TrimSpace(scene.Summary),
		BoundaryNote:       strings.TrimSpace(scene.BoundaryNote),
		BoundaryConfidence: strings.TrimSpace(scene.BoundaryConfidence),
	}

	requiredInts := map[string]int{
		"start_unit_id": out.StartUnitID,
		"end_unit_id":   out.EndUnitID,
	}
	for field, value := range requiredInts {
		if value <= 0 {
			return normalizedScene{}, fmt.Errorf("scene[%d] %s must be positive", index, field)
		}
	}
	required := map[string]string{
		"short_title":         out.ShortTitle,
		"primary_mode":        out.PrimaryMode,
		"summary":             out.Summary,
		"boundary_note":       out.BoundaryNote,
		"boundary_confidence": out.BoundaryConfidence,
	}
	for field, value := range required {
		if value == "" {
			return normalizedScene{}, fmt.Errorf("scene[%d] %s must not be empty", index, field)
		}
	}
	if !validPrimaryMode(out.PrimaryMode) {
		return normalizedScene{}, fmt.Errorf("scene[%d] primary_mode %q is not supported", index, out.PrimaryMode)
	}
	if !validBoundaryConfidence(out.BoundaryConfidence) {
		return normalizedScene{}, fmt.Errorf("scene[%d] boundary_confidence %q is not supported", index, out.BoundaryConfidence)
	}
	if len(scene.MainParticipants) == 0 {
		return normalizedScene{}, fmt.Errorf("scene[%d] main_participants must not be empty", index)
	}
	out.MainParticipants = make([]string, 0, len(scene.MainParticipants))
	for participantIndex, participant := range scene.MainParticipants {
		trimmed := strings.TrimSpace(participant)
		if trimmed == "" {
			return normalizedScene{}, fmt.Errorf("scene[%d] main_participants[%d] must not be empty", index, participantIndex)
		}
		out.MainParticipants = append(out.MainParticipants, trimmed)
	}
	return out, nil
}

func validPrimaryMode(value string) bool {
	switch value {
	case "Recap", "Discussion", "Combat", "Narrative":
		return true
	default:
		return false
	}
}

func validBoundaryConfidence(value string) bool {
	switch value {
	case "High", "Medium", "Low":
		return true
	default:
		return false
	}
}

func warningsFromCaveats(caveats []string) ([]contracts.Warning, error) {
	if len(caveats) == 0 {
		return nil, nil
	}
	warnings := make([]contracts.Warning, 0, len(caveats))
	for i, caveat := range caveats {
		trimmed := strings.TrimSpace(caveat)
		if trimmed == "" {
			return nil, fmt.Errorf("boundary_caveats[%d] must not be empty after trimming", i)
		}
		warnings = append(warnings, contracts.Warning{
			Scope:      Key,
			ReasonCode: "scene_boundary_caveat",
			Message:    trimmed,
		})
	}
	return warnings, nil
}

func chunkerErrorf(format string, args ...any) error {
	return fmt.Errorf("dnd scenes chunker: "+format, args...)
}