108 lines
2.5 KiB
Go
108 lines
2.5 KiB
Go
package spells
|
|
|
|
import (
|
|
"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/prompt"
|
|
)
|
|
|
|
type promptData struct {
|
|
SourceID string
|
|
HasChunk bool
|
|
ChunkID string
|
|
ChunkIndex int
|
|
Units []promptUnit
|
|
}
|
|
|
|
type promptUnit struct {
|
|
ID string
|
|
Text string
|
|
Metadata []promptMetadata
|
|
}
|
|
|
|
type promptMetadata struct {
|
|
Key string
|
|
Value string
|
|
}
|
|
|
|
var spellsPromptBundle = mustLoadPromptBundle()
|
|
|
|
func mustLoadPromptBundle() *prompt.Bundle {
|
|
bundle, err := prompt.LoadBundle(embeddedAssets, prompt.Definition{
|
|
PromptID: PromptID,
|
|
Version: SchemaVersion,
|
|
EmbeddedPath: "assets/prompts",
|
|
SystemPath: "assets/prompts/system.md",
|
|
UserPath: "assets/prompts/user.md",
|
|
ReferenceSlots: cloneReferenceSlots(referenceSlots),
|
|
})
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return bundle
|
|
}
|
|
|
|
func buildPromptData(req contracts.ExtractionRequest) (promptData, error) {
|
|
if req.Source == nil {
|
|
return promptData{}, fmt.Errorf("dnd spells prompt: source must not be nil")
|
|
}
|
|
if req.Chunk == nil {
|
|
return promptData{}, fmt.Errorf("dnd spells prompt: chunk must not be nil")
|
|
}
|
|
|
|
data := promptData{
|
|
SourceID: req.Source.ID,
|
|
HasChunk: true,
|
|
ChunkID: req.Chunk.ID,
|
|
ChunkIndex: req.Chunk.Index,
|
|
Units: make([]promptUnit, 0, len(req.Chunk.Units)),
|
|
}
|
|
for _, unit := range req.Chunk.Units {
|
|
data.Units = append(data.Units, promptUnit{
|
|
ID: unit.ID,
|
|
Text: unit.Text,
|
|
Metadata: selectedMetadata(unit),
|
|
})
|
|
}
|
|
return data, nil
|
|
}
|
|
|
|
func renderPrompt(req contracts.ExtractionRequest) (system string, user string, metadata prompt.Metadata, err error) {
|
|
data, err := buildPromptData(req)
|
|
if err != nil {
|
|
return "", "", prompt.Metadata{}, err
|
|
}
|
|
system, user, metadata, err = spellsPromptBundle.RenderUserSystemWithReferences(data, req.References)
|
|
if err != nil {
|
|
return "", "", prompt.Metadata{}, fmt.Errorf("dnd spells prompt: %w", err)
|
|
}
|
|
return system, user, metadata, nil
|
|
}
|
|
|
|
func selectedMetadata(unit source.SourceUnit) []promptMetadata {
|
|
if len(unit.Metadata) == 0 {
|
|
return nil
|
|
}
|
|
|
|
keys := []string{"speaker", "start", "end"}
|
|
metadata := make([]promptMetadata, 0, len(keys))
|
|
for _, key := range keys {
|
|
value, ok := unit.Metadata[key]
|
|
if !ok {
|
|
continue
|
|
}
|
|
rendered := strings.TrimSpace(fmt.Sprint(value))
|
|
if rendered == "" {
|
|
continue
|
|
}
|
|
metadata = append(metadata, promptMetadata{
|
|
Key: key,
|
|
Value: rendered,
|
|
})
|
|
}
|
|
return metadata
|
|
}
|