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",
	"chunks.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 (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) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
	if c == nil {
		return contracts.ChunkResult{}, chunkerErrorf("chunker must not be nil")
	}
	if c.llm == nil {
		return contracts.ChunkResult{}, chunkerErrorf("LLM client must not be nil")
	}
	if ctx == nil {
		return contracts.ChunkResult{}, chunkerErrorf("context must not be nil")
	}
	if err := ctx.Err(); err != nil {
		return contracts.ChunkResult{}, chunkerErrorf("context error before chunking: %w", err)
	}
	if req.Source == nil {
		return contracts.ChunkResult{}, chunkerErrorf("source must not be nil")
	}
	if len(req.Source.Units) == 0 {
		return contracts.ChunkResult{}, chunkerErrorf("source units must not be empty")
	}
	if err := source.ValidateDocument(req.Source); err != nil {
		return contracts.ChunkResult{}, 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.ChunkResult{}, chunkerErrorf("complete structured output: %w", err)
	}

	warnings, err := warningsFromCaveats(response.BoundaryCaveats)
	if err != nil {
		return contracts.ChunkResult{}, chunkerErrorf("malformed structured output: %w", err)
	}
	chunks, err := chunksFromResponse(req.Source, response)
	if err != nil {
		return contracts.ChunkResult{}, chunkerErrorf("malformed structured output: %w", err)
	}
	return contracts.ChunkResult{
		Chunks:   chunks,
		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 chunksFromResponse(doc *source.SourceDocument, response chunkResponse) ([]source.Chunk, error) {
	if response.Scenes == nil {
		return nil, fmt.Errorf("scenes must be present")
	}
	if len(response.Scenes) == 0 {
		return nil, 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
	}

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

		startIndex, ok := unitIndexes[normalized.StartUnitID]
		if !ok {
			return nil, fmt.Errorf("scene[%d] start_unit_id %d was not found", i, normalized.StartUnitID)
		}
		endIndex, ok := unitIndexes[normalized.EndUnitID]
		if !ok {
			return nil, fmt.Errorf("scene[%d] end_unit_id %d was not found", i, normalized.EndUnitID)
		}
		if startIndex > endIndex {
			return nil, 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 nil, fmt.Errorf("first scene must start at first source unit %d", doc.Units[0].ID)
		}
		if i > 0 {
			if startIndex <= previousEnd {
				return nil, fmt.Errorf("scene[%d] overlaps previous scene", i)
			}
			if startIndex > previousEnd+1 {
				return nil, fmt.Errorf("scene[%d] leaves a gap after previous scene", i)
			}
		}
		previousEnd = endIndex

		units := cloneUnits(doc.Units[startIndex : endIndex+1])
		content, err := chunkContent(units)
		if err != nil {
			return nil, err
		}
		chunks = append(chunks, source.Chunk{
			ID:       fmt.Sprintf("scene-%06d", i+1),
			SourceID: doc.ID,
			Index:    i,
			Ref: source.SourceRef{
				SourceID:    doc.ID,
				StartUnitID: units[0].Ref.StartUnitID,
				EndUnitID:   units[len(units)-1].Ref.EndUnitID,
			},
			Content:   content,
			MediaType: "application/json",
			Units:     units,
			Metadata: map[string]any{
				"scene_title":         normalized.ShortTitle,
				"primary_mode":        normalized.PrimaryMode,
				"main_participants":   append([]string(nil), normalized.MainParticipants...),
				"summary":             normalized.Summary,
				"boundary_note":       normalized.BoundaryNote,
				"boundary_confidence": normalized.BoundaryConfidence,
				"start_unit_id":       normalized.StartUnitID,
				"end_unit_id":         normalized.EndUnitID,
				"unit_count":          len(units),
			},
		})
	}

	if previousEnd != len(doc.Units)-1 {
		return nil, fmt.Errorf("final scene must end at final source unit %d", doc.Units[len(doc.Units)-1].ID)
	}
	return chunks, nil
}

func chunkContent(units []source.SourceUnit) ([]byte, error) {
	content, err := json.Marshal(struct {
		Units []source.SourceUnit `json:"units"`
	}{
		Units: units,
	})
	if err != nil {
		return nil, fmt.Errorf("encode chunk content: %w", err)
	}
	return content, 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 cloneUnits(units []source.SourceUnit) []source.SourceUnit {
	out := make([]source.SourceUnit, 0, len(units))
	for _, unit := range units {
		out = append(out, source.SourceUnit{
			ID:       unit.ID,
			Kind:     unit.Kind,
			Text:     unit.Text,
			Ref:      unit.Ref,
			Metadata: cloneMetadata(unit.Metadata),
		})
	}
	return out
}

func cloneMetadata(metadata map[string]any) map[string]any {
	if len(metadata) == 0 {
		return nil
	}
	out := make(map[string]any, len(metadata))
	for key, value := range metadata {
		out[key] = value
	}
	return out
}

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