355 lines
11 KiB
Go
355 lines
11 KiB
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/sharedassets/dnd"
|
|
)
|
|
|
|
const Key = "dnd/scenes"
|
|
|
|
var requiredCapabilities = []string{
|
|
"source.transcript",
|
|
}
|
|
|
|
var providedCapabilities = []string{
|
|
"chunks",
|
|
"chunks.scenes",
|
|
}
|
|
|
|
var referenceSlotDescriptions = dnd.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 Chunker struct{}
|
|
|
|
func New() *Chunker {
|
|
return &Chunker{}
|
|
}
|
|
|
|
func (c *Chunker) Key() string {
|
|
return Key
|
|
}
|
|
|
|
func (c *Chunker) ReferenceSlots() []contracts.ReferenceSlot {
|
|
return dnd.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 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)
|
|
}
|
|
if req.LLMClient == nil {
|
|
return contracts.ChunkResult{}, chunkerErrorf("LLM client must not be nil")
|
|
}
|
|
if len(req.Options) > 0 {
|
|
return contracts.ChunkResult{}, chunkerErrorf("options are not supported")
|
|
}
|
|
|
|
var response chunkResponse
|
|
if _, err := req.LLMClient.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
|
StageName: Key,
|
|
PromptID: PromptID,
|
|
PromptVersion: ResponseSchemaVersion,
|
|
ProfileID: req.LLMProfile,
|
|
SessionID: req.SessionID,
|
|
Inputs: dnd.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: dnd.ReferenceSlots(referenceSlotDescriptions),
|
|
}
|
|
}
|
|
|
|
func Register(registry *pipeline.ChunkerRegistry) error {
|
|
return registry.RegisterWithSpec(ModuleSpec(), func() (contracts.Chunker, error) {
|
|
return New(), nil
|
|
})
|
|
}
|
|
|
|
func chunksFromResponse(doc *source.SourceDocument, response chunkResponse) ([]contracts.SourceChunk, 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([]contracts.SourceChunk, 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, contracts.SourceChunk{
|
|
ID: fmt.Sprintf("scene-%06d", i+1),
|
|
SourceID: doc.ID,
|
|
Index: i,
|
|
StartUnitID: units[0].ID,
|
|
EndUnitID: units[len(units)-1].ID,
|
|
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 := dnd.ResolveUnitID(doc, "start_unit_id", scene.StartUnitID)
|
|
if err != nil {
|
|
return normalizedScene{}, fmt.Errorf("scene[%d] %w", index, err)
|
|
}
|
|
endUnitID, err := dnd.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,
|
|
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...)
|
|
}
|