Add D&D spells prompt assets
This commit is contained in:
90
internal/modules/extract/dnd/spells/prompt.go
Normal file
90
internal/modules/extract/dnd/spells/prompt.go
Normal file
@@ -0,0 +1,90 @@
|
||||
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
|
||||
}
|
||||
|
||||
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 = prompt.RenderUserSystem(prompt.DNDSpellsPromptID, data)
|
||||
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
|
||||
}
|
||||
159
internal/modules/extract/dnd/spells/prompt_test.go
Normal file
159
internal/modules/extract/dnd/spells/prompt_test.go
Normal file
@@ -0,0 +1,159 @@
|
||||
package spells
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/prompt"
|
||||
)
|
||||
|
||||
func TestBuildPromptDataFromGenericSourceChunk(t *testing.T) {
|
||||
req := promptExtractionRequest()
|
||||
|
||||
data, err := buildPromptData(req)
|
||||
if err != nil {
|
||||
t.Fatalf("buildPromptData() error = %v, want nil", err)
|
||||
}
|
||||
|
||||
if data.SourceID != "session-alpha" {
|
||||
t.Fatalf("SourceID = %q, want session-alpha", data.SourceID)
|
||||
}
|
||||
if !data.HasChunk || data.ChunkID != "session-alpha:chunk:0" || data.ChunkIndex != 0 {
|
||||
t.Fatalf("chunk data = %#v, want fixture chunk", data)
|
||||
}
|
||||
if len(data.Units) != 2 {
|
||||
t.Fatalf("len(Units) = %d, want 2", len(data.Units))
|
||||
}
|
||||
first := data.Units[0]
|
||||
if first.ID != "seg-001" || first.Text != "Aria raises her hand and casts Cure Wounds." {
|
||||
t.Fatalf("first unit = %#v, want source unit data", first)
|
||||
}
|
||||
wantMetadata := []promptMetadata{
|
||||
{Key: "speaker", Value: "Alice"},
|
||||
{Key: "start", Value: "1.25"},
|
||||
{Key: "end", Value: "3.5"},
|
||||
}
|
||||
if !reflect.DeepEqual(first.Metadata, wantMetadata) {
|
||||
t.Fatalf("first.Metadata = %#v, want %#v", first.Metadata, wantMetadata)
|
||||
}
|
||||
if len(data.Units[1].Metadata) != 0 {
|
||||
t.Fatalf("second.Metadata = %#v, want no selected metadata", data.Units[1].Metadata)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildPromptDataDoesNotMutateRequest(t *testing.T) {
|
||||
req := promptExtractionRequest()
|
||||
beforeSource := mustJSON(t, req.Source)
|
||||
beforeChunk := mustJSON(t, req.Chunk)
|
||||
beforeRequest := mustJSON(t, req)
|
||||
|
||||
if _, err := buildPromptData(req); err != nil {
|
||||
t.Fatalf("buildPromptData() error = %v, want nil", err)
|
||||
}
|
||||
afterSource := mustJSON(t, req.Source)
|
||||
afterChunk := mustJSON(t, req.Chunk)
|
||||
afterRequest := mustJSON(t, req)
|
||||
if beforeSource != afterSource || beforeChunk != afterChunk || beforeRequest != afterRequest {
|
||||
t.Fatalf(
|
||||
"request mutated:\nsource before: %s\nsource after: %s\nchunk before: %s\nchunk after: %s\nrequest before: %s\nrequest after: %s",
|
||||
beforeSource,
|
||||
afterSource,
|
||||
beforeChunk,
|
||||
afterChunk,
|
||||
beforeRequest,
|
||||
afterRequest,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderPromptIncludesSourceContext(t *testing.T) {
|
||||
system, user, metadata, err := renderPrompt(promptExtractionRequest())
|
||||
if err != nil {
|
||||
t.Fatalf("renderPrompt() error = %v, want nil", err)
|
||||
}
|
||||
|
||||
if !strings.Contains(system, prompt.HardeningText()) {
|
||||
t.Fatalf("system prompt = %q, want hardening text", system)
|
||||
}
|
||||
for _, want := range []string{
|
||||
"session-alpha",
|
||||
"session-alpha:chunk:0",
|
||||
"seg-001",
|
||||
"Aria raises her hand and casts Cure Wounds.",
|
||||
"speaker: Alice",
|
||||
"start: 1.25",
|
||||
"end: 3.5",
|
||||
} {
|
||||
if !strings.Contains(user, want) {
|
||||
t.Fatalf("user prompt = %q, want substring %q", user, want)
|
||||
}
|
||||
}
|
||||
if metadata.PromptID != prompt.DNDSpellsPromptID {
|
||||
t.Fatalf("metadata.PromptID = %q, want %q", metadata.PromptID, prompt.DNDSpellsPromptID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildPromptDataRejectsMissingSourceContext(t *testing.T) {
|
||||
if _, err := buildPromptData(contracts.ExtractionRequest{}); err == nil || !strings.Contains(err.Error(), "source") {
|
||||
t.Fatalf("buildPromptData() error = %v, want source error", err)
|
||||
}
|
||||
if _, err := buildPromptData(contracts.ExtractionRequest{Source: promptSourceDocument()}); err == nil || !strings.Contains(err.Error(), "chunk") {
|
||||
t.Fatalf("buildPromptData() error = %v, want chunk error", err)
|
||||
}
|
||||
}
|
||||
|
||||
func promptExtractionRequest() contracts.ExtractionRequest {
|
||||
doc := promptSourceDocument()
|
||||
chunk := &contracts.SourceChunk{
|
||||
ID: "session-alpha:chunk:0",
|
||||
SourceID: doc.ID,
|
||||
Index: 0,
|
||||
Units: append([]source.SourceUnit(nil), doc.Units...),
|
||||
Metadata: map[string]any{"ignored": "chunk metadata"},
|
||||
}
|
||||
return contracts.ExtractionRequest{
|
||||
Source: doc,
|
||||
Chunk: chunk,
|
||||
}
|
||||
}
|
||||
|
||||
func promptSourceDocument() *source.SourceDocument {
|
||||
return &source.SourceDocument{
|
||||
ID: "session-alpha",
|
||||
Kind: "transcript",
|
||||
Format: "application/vnd.seriatim.minimal+json",
|
||||
Digest: "sha256:test",
|
||||
Units: []source.SourceUnit{
|
||||
{
|
||||
ID: "seg-001",
|
||||
Kind: "transcript_segment",
|
||||
Text: "Aria raises her hand and casts Cure Wounds.",
|
||||
Metadata: map[string]any{
|
||||
"speaker": "Alice",
|
||||
"start": json.Number("1.25"),
|
||||
"end": json.Number("3.5"),
|
||||
"ignored": "not rendered",
|
||||
},
|
||||
},
|
||||
{
|
||||
ID: "seg-002",
|
||||
Kind: "transcript_segment",
|
||||
Text: "The fighter's wounds begin to close.",
|
||||
Metadata: map[string]any{"ignored": "not rendered"},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func mustJSON(t *testing.T, value any) string {
|
||||
t.Helper()
|
||||
encoded, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal() error = %v, want nil", err)
|
||||
}
|
||||
return string(encoded)
|
||||
}
|
||||
Reference in New Issue
Block a user