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

package scenes

import (
	"context"
	"fmt"

	"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",
}

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 := promptAssetMetadata()
	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)
	}

	plan, err := planFromResponse(req.Source, response)
	if err != nil {
		return contracts.ChunkPlanResult{}, chunkerErrorf("malformed structured output: %w", err)
	}
	return contracts.ChunkPlanResult{Plan: plan}, 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) (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")
	}

	index := source.NewDocumentIndex(doc)

	ranges := make([]source.ChunkRange, 0, len(response.Scenes))
	previousEnd := -1
	for i, scene := range response.Scenes {
		startUnitID, err := shared.ResolveUnitID(index, "start_unit_id", scene.StartUnitID)
		if err != nil {
			return source.ChunkPlan{}, fmt.Errorf("scene[%d] %w", i, err)
		}
		endUnitID, err := shared.ResolveUnitID(index, "end_unit_id", scene.EndUnitID)
		if err != nil {
			return source.ChunkPlan{}, fmt.Errorf("scene[%d] %w", i, err)
		}

		startIndex, ok := index.Position(startUnitID)
		if !ok {
			return source.ChunkPlan{}, fmt.Errorf("scene[%d] start_unit_id %d was not found", i, startUnitID)
		}
		endIndex, ok := index.Position(endUnitID)
		if !ok {
			return source.ChunkPlan{}, fmt.Errorf("scene[%d] end_unit_id %d was not found", i, endUnitID)
		}
		if startIndex > endIndex {
			return source.ChunkPlan{}, fmt.Errorf("scene[%d] start_unit_id %d appears after end_unit_id %d", i, startUnitID, 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

		ranges = append(ranges, source.ChunkRange{
			StartUnitID: startUnitID,
			EndUnitID:   endUnitID,
		})
	}

	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)
	}
	return source.ChunkPlan{
		SourceDigest: doc.Digest,
		Ranges:       ranges,
	}, nil
}

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