eric/notarius
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add prompt input materials and session IDs

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-05 17:51:36 +00:00
parent 291298cf7b
commit 49d94cc2e9
7 changed files with 423 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
internal/framework/contracts/contracts.go
View File

@@ -15,6 +15,12 @@ type LLMMessage struct {
type StructuredCompletionRequest struct {
StageName string `json:"stage_name"`
PromptID string `json:"prompt_id,omitempty"`
PromptVersion string `json:"prompt_version,omitempty"`
ProfileID string `json:"profile_id,omitempty"`
SessionID string `json:"session_id,omitempty"`
Inputs LLMInputSet `json:"inputs,omitempty"`
Vars map[string]any `json:"vars,omitempty"`
Messages []LLMMessage `json:"messages"`
Model string `json:"model,omitempty"`
ResponseSchemaName string `json:"response_schema_name,omitempty"`
@@ -34,6 +40,44 @@ type StructuredLLMClient interface {
CompleteStructured(ctx context.Context, req StructuredCompletionRequest, out any) (StructuredCompletionResponse, error)
}
type LLMInputMaterial struct {
Name string `json:"name"`
MediaType string `json:"media_type,omitempty"`
Content []byte `json:"-"`
Digest string `json:"digest,omitempty"`
OriginURI string `json:"origin_uri,omitempty"`
SizeBytes int64 `json:"size_bytes,omitempty"`
}
func NewLLMInputMaterial(name string, mediaType string, content []byte, digest string, originURI string) LLMInputMaterial {
return LLMInputMaterial{
Name: name,
MediaType: mediaType,
Content: append([]byte(nil), content...),
Digest: digest,
OriginURI: originURI,
SizeBytes: int64(len(content)),
}
}
func (material LLMInputMaterial) Clone() LLMInputMaterial {
material.Content = append([]byte(nil), material.Content...)
return material
}
type LLMInputSet map[string]LLMInputMaterial
func (set LLMInputSet) Clone() LLMInputSet {
if len(set) == 0 {
return nil
}
out := make(LLMInputSet, len(set))
for key, material := range set {
out[key] = material.Clone()
}
return out
}
type ParseRequest struct {
SourceID string `json:"source_id,omitempty"`
Path string `json:"path,omitempty"`
@@ -57,12 +101,14 @@ type SourceChunk struct {
}
type ChunkRequest struct {
Source *source.SourceDocument `json:"-"`
References ReferenceSet `json:"references,omitempty"`
LLMClient StructuredLLMClient `json:"-"`
LLMProfile string `json:"llm_profile,omitempty"`
Options map[string]any `json:"options,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
Source *source.SourceDocument `json:"-"`
SourceInput LLMInputMaterial `json:"source_input,omitempty"`
SessionID string `json:"session_id,omitempty"`
References ReferenceSet `json:"references,omitempty"`
LLMClient StructuredLLMClient `json:"-"`
LLMProfile string `json:"llm_profile,omitempty"`
Options map[string]any `json:"options,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type ChunkResult struct {
@@ -118,6 +164,8 @@ type ExtractionRequest struct {
Source *source.SourceDocument `json:"-"`
Chunk *SourceChunk `json:"chunk,omitempty"`
AmbientContext map[string]any `json:"ambient_context,omitempty"`
SourceInput LLMInputMaterial `json:"source_input,omitempty"`
SessionID string `json:"session_id,omitempty"`
References ReferenceSet `json:"references,omitempty"`
LLMClient StructuredLLMClient `json:"-"`
LLMProfile string `json:"llm_profile,omitempty"`
@@ -164,14 +212,16 @@ type Merger interface {
}
type NormalizeRequest struct {
Source *source.SourceDocument `json:"-"`
LaneID string `json:"lane_id"`
Candidates []artifacts.ArtifactCandidate `json:"candidates"`
References ReferenceSet `json:"references,omitempty"`
LLMClient StructuredLLMClient `json:"-"`
LLMProfile string `json:"llm_profile,omitempty"`
Options map[string]any `json:"options,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
Source *source.SourceDocument `json:"-"`
LaneID string `json:"lane_id"`
Candidates []artifacts.ArtifactCandidate `json:"candidates"`
SourceInput LLMInputMaterial `json:"source_input,omitempty"`
SessionID string `json:"session_id,omitempty"`
References ReferenceSet `json:"references,omitempty"`
LLMClient StructuredLLMClient `json:"-"`
LLMProfile string `json:"llm_profile,omitempty"`
Options map[string]any `json:"options,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type NormalizeResult struct {

47
internal/framework/contracts/contracts_test.go
View File

@@ -251,6 +251,53 @@ func TestReferenceItemJSONOmitsContent(t *testing.T) {
}
}
func TestLLMInputMaterialCopiesContentAndOmitsContentFromJSON(t *testing.T) {
content := []byte("raw source bytes")
material := NewLLMInputMaterial("transcript", "application/json", content, "sha256:source", "file:///tmp/source.json")
content[0] = 'R'
if got := string(material.Content); got != "raw source bytes" {
t.Fatalf("material content = %q, want defensive copy", got)
}
if material.SizeBytes != int64(len("raw source bytes")) {
t.Fatalf("SizeBytes = %d, want content length", material.SizeBytes)
}
clone := material.Clone()
clone.Content[0] = 'X'
if got := string(material.Content); got != "raw source bytes" {
t.Fatalf("cloned material content aliased original: %q", got)
}
encoded, err := json.Marshal(material)
if err != nil {
t.Fatalf("json.Marshal() error = %v, want nil", err)
}
var got map[string]any
if err := json.Unmarshal(encoded, &got); err != nil {
t.Fatalf("json.Unmarshal() error = %v, want nil", err)
}
if _, ok := got["content"]; ok {
t.Fatalf("encoded material leaked content: %s", encoded)
}
if _, ok := got["Content"]; ok {
t.Fatalf("encoded material leaked Content: %s", encoded)
}
if got["digest"] != "sha256:source" || got["origin_uri"] != "file:///tmp/source.json" {
t.Fatalf("encoded material = %#v, want non-secret provenance", got)
}
}
func TestLLMInputSetCloneCopiesContent(t *testing.T) {
set := LLMInputSet{
"transcript": NewLLMInputMaterial("transcript", "application/json", []byte("source"), "sha256:source", "file:///tmp/source.json"),
}
clone := set.Clone()
clone["transcript"].Content[0] = 'S'
if got := string(set["transcript"].Content); got != "source" {
t.Fatalf("input set clone aliased content: %q", got)
}
}
func TestFakeMergeNormalizeAndOutputContracts(t *testing.T) {
candidate := artifacts.ArtifactCandidate{
Index: 0,

113
internal/framework/pipeline/runner.go
View File

@@ -2,8 +2,12 @@ package pipeline
import (
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"mime"
"path"
"path/filepath"
"strings"
"time"
@@ -37,6 +41,7 @@ type RunInput struct {
Path string
RawInput []byte
LLMClient contracts.StructuredLLMClient
SessionID string
RunID string
StartedAt time.Time
LLMProfiles []artifacts.LLMProfileManifest
@@ -86,6 +91,9 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
if err := source.ValidateDocument(doc); err != nil {
return failOutput(output), fmt.Errorf("validate source document: %w", err)
}
sourceInput := sourceInputMaterial(input.Path, input.RawInput)
sessionID := resolvedSessionID(input.SessionID, doc.ID)
output.Manifest.Metadata = manifestMetadataWithSessionID(output.Manifest.Metadata, sessionID)
output.Manifest.SourceDigests = []string{doc.Digest}
chunker, err := r.registries.Chunkers.Build(input.Pipeline.Chunk.Module)
@@ -94,12 +102,14 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
}
attachModuleManifestMetadata(&output, "chunker", chunker)
chunkResult, err := chunker.Chunk(ctx, contracts.ChunkRequest{
Source: doc,
References: CloneReferenceSet(input.Pipeline.ChunkReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: input.Pipeline.Chunk.LLMProfile,
Options: cloneOptions(input.Pipeline.Chunk.Options),
Metadata: input.Metadata,
Source: doc,
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(input.Pipeline.ChunkReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: input.Pipeline.Chunk.LLMProfile,
Options: cloneOptions(input.Pipeline.Chunk.Options),
Metadata: input.Metadata,
})
output.Warnings = append(output.Warnings, chunkResult.Warnings...)
if err != nil {
@@ -115,7 +125,7 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
nextCandidateIndex := 0
for _, lane := range input.Pipeline.ArtifactLanes {
if err := r.runLane(ctx, input, doc, canonicalChunks, lane, &output, &nextCandidateIndex); err != nil {
if err := r.runLane(ctx, input, doc, sourceInput, sessionID, canonicalChunks, lane, &output, &nextCandidateIndex); err != nil {
return failOutput(output), err
}
}
@@ -154,7 +164,7 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
return output, nil
}
func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.SourceDocument, chunks []contracts.SourceChunk, lane ResolvedArtifactLane, output *RunOutput, nextCandidateIndex *int) error {
func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, chunks []contracts.SourceChunk, lane ResolvedArtifactLane, output *RunOutput, nextCandidateIndex *int) error {
extractor, err := r.registries.Extractors.Build(lane.Extract.Module)
if err != nil {
return fmt.Errorf("build extractor %q for lane %q: %w", lane.Extract.Module, lane.ID, err)
@@ -185,13 +195,15 @@ func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.Source
for index := range chunks {
chunk := chunks[index]
result, err := extractor.Extract(ctx, contracts.ExtractionRequest{
Source: doc,
Chunk: &chunk,
References: CloneReferenceSet(lane.ExtractReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Extract.LLMProfile,
Options: cloneOptions(lane.Extract.Options),
Metadata: input.Metadata,
Source: doc,
Chunk: &chunk,
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(lane.ExtractReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Extract.LLMProfile,
Options: cloneOptions(lane.Extract.Options),
Metadata: input.Metadata,
})
output.Warnings = append(output.Warnings, result.Warnings...)
if err != nil {
@@ -222,14 +234,16 @@ func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.Source
}
normalizeResult, err := normalizer.Normalize(ctx, contracts.NormalizeRequest{
Source: doc,
LaneID: lane.ID,
Candidates: mergeResult.Candidates,
References: CloneReferenceSet(lane.NormalizeReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Normalize.LLMProfile,
Options: cloneOptions(lane.Normalize.Options),
Metadata: input.Metadata,
Source: doc,
LaneID: lane.ID,
Candidates: mergeResult.Candidates,
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(lane.NormalizeReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Normalize.LLMProfile,
Options: cloneOptions(lane.Normalize.Options),
Metadata: input.Metadata,
})
output.Warnings = append(output.Warnings, normalizeResult.Warnings...)
if err != nil {
@@ -510,6 +524,59 @@ func cloneLLMProfiles(profiles []artifacts.LLMProfileManifest) []artifacts.LLMPr
return append([]artifacts.LLMProfileManifest(nil), profiles...)
}
func sourceInputMaterial(inputPath string, content []byte) contracts.LLMInputMaterial {
return contracts.NewLLMInputMaterial(
"source",
sourceInputMediaType(inputPath),
content,
sourceInputDigest(content),
sourceInputOriginURI(inputPath),
)
}
func sourceInputMediaType(inputPath string) string {
extension := strings.ToLower(filepath.Ext(strings.TrimSpace(inputPath)))
if extension == ".json" {
return "application/json"
}
mediaType := mime.TypeByExtension(extension)
if strings.TrimSpace(mediaType) == "" {
return unknownMediaType
}
return canonicalMediaType(mediaType)
}
func sourceInputDigest(content []byte) string {
sum := sha256.Sum256(content)
return "sha256:" + hex.EncodeToString(sum[:])
}
func sourceInputOriginURI(inputPath string) string {
if strings.TrimSpace(inputPath) == "" {
return ""
}
return fileURI(inputPath)
}
func resolvedSessionID(explicit string, sourceDocumentID string) string {
if trimmed := strings.TrimSpace(explicit); trimmed != "" {
return trimmed
}
return strings.TrimSpace(sourceDocumentID)
}
func manifestMetadataWithSessionID(metadata map[string]any, sessionID string) map[string]any {
out := cloneMetadata(metadata)
if strings.TrimSpace(sessionID) == "" {
return out
}
if out == nil {
out = make(map[string]any)
}
out["session_id"] = sessionID
return out
}
func cloneWarnings(warnings []contracts.Warning) []contracts.Warning {
if len(warnings) == 0 {
return nil

78
internal/framework/pipeline/runner_test.go
View File

@@ -476,6 +476,84 @@ func TestRunExecutesChunksAndPassesChunkAndLLMClient(t *testing.T) {
}
}
func TestRunPassesSourceInputAndSessionIDToPromptCapableStages(t *testing.T) {
modules := defaultRunnerModules()
rawInput := []byte("{\"source\":\"exact bytes\"}")
output, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{
Pipeline: resolvedPipeline(),
Path: "session.json",
RawInput: rawInput,
SessionID: " explicit-session ",
})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
if got := output.Manifest.Metadata["session_id"]; got != "explicit-session" {
t.Fatalf("manifest metadata = %#v, want session_id", output.Manifest.Metadata)
}
requests := []struct {
name string
material contracts.LLMInputMaterial
sessionID string
}{
{name: "chunk", material: modules.chunker.requests[0].SourceInput, sessionID: modules.chunker.requests[0].SessionID},
{name: "extract first", material: modules.extractors["extract-alpha"].requests[0].SourceInput, sessionID: modules.extractors["extract-alpha"].requests[0].SessionID},
{name: "extract second", material: modules.extractors["extract-alpha"].requests[1].SourceInput, sessionID: modules.extractors["extract-alpha"].requests[1].SessionID},
{name: "normalize", material: modules.normalizers["normalize"].requests[0].SourceInput, sessionID: modules.normalizers["normalize"].requests[0].SessionID},
}
for _, req := range requests {
if req.sessionID != "explicit-session" {
t.Fatalf("%s session ID = %q, want explicit-session", req.name, req.sessionID)
}
if got := string(req.material.Content); got != string(rawInput) {
t.Fatalf("%s source input content = %q, want exact raw input", req.name, got)
}
if req.material.Name != "source" || req.material.MediaType != "application/json" || req.material.SizeBytes != int64(len(rawInput)) {
t.Fatalf("%s source input = %#v, want source metadata", req.name, req.material)
}
if req.material.Digest != sourceInputDigest(rawInput) {
t.Fatalf("%s digest = %q, want %q", req.name, req.material.Digest, sourceInputDigest(rawInput))
}
if !strings.HasPrefix(req.material.OriginURI, "file://") || !strings.HasSuffix(req.material.OriginURI, "/session.json") {
t.Fatalf("%s origin URI = %q, want file URI ending in session.json", req.name, req.material.OriginURI)
}
}
modules.chunker.requests[0].SourceInput.Content[0] = 'X'
if got := string(modules.extractors["extract-alpha"].requests[0].SourceInput.Content); got != string(rawInput) {
t.Fatalf("source input content aliased across requests: %q", got)
}
if got := string(rawInput); got != "{\"source\":\"exact bytes\"}" {
t.Fatalf("raw input mutated through request material: %q", got)
}
}
func TestRunDefaultsSessionIDFromParsedSourceDocumentID(t *testing.T) {
modules := defaultRunnerModules()
output, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{
Pipeline: resolvedPipeline(),
Path: "notes.unknown",
RawInput: []byte("notes"),
})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
if got := modules.chunker.requests[0].SessionID; got != "source-1" {
t.Fatalf("chunk session ID = %q, want parsed source document ID", got)
}
if got := output.Manifest.Metadata["session_id"]; got != "source-1" {
t.Fatalf("manifest metadata = %#v, want default session id", output.Manifest.Metadata)
}
if got := modules.chunker.requests[0].SourceInput.MediaType; got != unknownMediaType {
t.Fatalf("source input media type = %q, want fallback %q", got, unknownMediaType)
}
}
func TestRunPassesInputRequestFields(t *testing.T) {
modules := defaultRunnerModules()
metadata := map[string]any{"request": "test"}