Construct universal modules with decoded options
This commit is contained in:
@@ -17,7 +17,6 @@ var _ contracts.LegacyRawExtractor = compositionExtractor{}
|
||||
var _ contracts.LegacyRawMerger = compositionMerger{}
|
||||
var _ contracts.LegacyRawNormalizer = compositionNormalizer{}
|
||||
var _ contracts.LegacyRawValidator = compositionValidator{}
|
||||
var _ contracts.StructuredLLMClient = compositionLLMClient{}
|
||||
var _ contracts.OutputEncoder = compositionOutputEncoder{}
|
||||
|
||||
func TestContractsComposeAcrossPackages(t *testing.T) {
|
||||
@@ -38,9 +37,8 @@ func TestContractsComposeAcrossPackages(t *testing.T) {
|
||||
}
|
||||
|
||||
chunking, err := chunker.Chunk(ctx, contracts.ChunkRequest{
|
||||
Source: doc,
|
||||
LLMClient: compositionLLMClient{},
|
||||
Metadata: map[string]any{"max_units": 2},
|
||||
Source: doc,
|
||||
Metadata: map[string]any{"request": "test"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Chunk() error = %v, want nil", err)
|
||||
@@ -136,10 +134,6 @@ func (chunker compositionChunker) Chunk(ctx context.Context, req contracts.Chunk
|
||||
if req.Source == nil {
|
||||
return contracts.ChunkResult{}, errors.New("source document is required")
|
||||
}
|
||||
if req.LLMClient == nil {
|
||||
return contracts.ChunkResult{}, errors.New("structured llm client is required")
|
||||
}
|
||||
|
||||
return contracts.ChunkResult{
|
||||
Chunks: []source.Chunk{
|
||||
{
|
||||
@@ -160,12 +154,6 @@ func (chunker compositionChunker) Chunk(ctx context.Context, req contracts.Chunk
|
||||
}, nil
|
||||
}
|
||||
|
||||
type compositionLLMClient struct{}
|
||||
|
||||
func (client compositionLLMClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
|
||||
return contracts.StructuredCompletionResponse{}, nil
|
||||
}
|
||||
|
||||
type compositionExtractor struct{}
|
||||
|
||||
func (extractor compositionExtractor) Key() string {
|
||||
|
||||
@@ -129,7 +129,6 @@ type ParseRequest struct {
|
||||
Path string `json:"path,omitempty"`
|
||||
Raw []byte `json:"-"`
|
||||
LLMProfile string `json:"llm_profile,omitempty"`
|
||||
Options map[string]any `json:"options,omitempty"`
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
@@ -143,9 +142,7 @@ type ChunkRequest struct {
|
||||
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"`
|
||||
}
|
||||
|
||||
@@ -377,7 +374,6 @@ type OutputRequest struct {
|
||||
Rejected []RejectedOutput `json:"rejected,omitempty"`
|
||||
Warnings []Warning `json:"warnings,omitempty"`
|
||||
LLMProfile string `json:"llm_profile,omitempty"`
|
||||
Options map[string]any `json:"options,omitempty"`
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
|
||||
@@ -97,7 +97,7 @@ func TestFakeChunkerReturnsSourceChunks(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestFakeChunkerReceivesLLMClient(t *testing.T) {
|
||||
func TestFakeChunkerReceivesPerRunContext(t *testing.T) {
|
||||
doc := &source.SourceDocument{
|
||||
ID: "source-1",
|
||||
Kind: "document",
|
||||
@@ -107,14 +107,13 @@ func TestFakeChunkerReceivesLLMClient(t *testing.T) {
|
||||
{ID: 1, Kind: "section", Text: "Source text."},
|
||||
},
|
||||
}
|
||||
client := fakeLLMClient{}
|
||||
chunker := &recordingChunker{key: "llm-chunker"}
|
||||
|
||||
if _, err := chunker.Chunk(context.Background(), ChunkRequest{Source: doc, LLMClient: client}); err != nil {
|
||||
if _, err := chunker.Chunk(context.Background(), ChunkRequest{Source: doc, SessionID: "session", LLMProfile: "profile"}); err != nil {
|
||||
t.Fatalf("Chunk() error = %v, want nil", err)
|
||||
}
|
||||
if chunker.request.LLMClient == nil {
|
||||
t.Fatal("ChunkRequest.LLMClient = nil, want structured LLM client")
|
||||
if chunker.request.SessionID != "session" || chunker.request.LLMProfile != "profile" {
|
||||
t.Fatalf("ChunkRequest = %#v, want per-run session and profile", chunker.request)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -238,9 +238,18 @@ func wrapDebugLLMClient(client contracts.StructuredLLMClient, recorder DebugReco
|
||||
if client == nil || recorder == nil || !recorder.Enabled() {
|
||||
return client
|
||||
}
|
||||
if _, ok := client.(*debugLLMClient); ok {
|
||||
return client
|
||||
}
|
||||
return &debugLLMClient{inner: client, recorder: recorder}
|
||||
}
|
||||
|
||||
// WithDebugLLMRecording decorates a shared LLM client so calls made by
|
||||
// construction-injected modules participate in the run's debug recording.
|
||||
func WithDebugLLMRecording(client contracts.StructuredLLMClient, recorder DebugRecorder) contracts.StructuredLLMClient {
|
||||
return wrapDebugLLMClient(client, recorder)
|
||||
}
|
||||
|
||||
func (client *debugLLMClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
|
||||
client.mu.Lock()
|
||||
client.counter++
|
||||
|
||||
@@ -133,7 +133,6 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
|
||||
Path: input.Path,
|
||||
Raw: input.RawInput,
|
||||
LLMProfile: input.pipeline.Input.LLMProfile,
|
||||
Options: cloneOptions(input.pipeline.Input.Options),
|
||||
Metadata: input.Metadata,
|
||||
})
|
||||
if err != nil {
|
||||
@@ -206,9 +205,7 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
|
||||
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,
|
||||
})
|
||||
if err != nil {
|
||||
@@ -352,7 +349,6 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
|
||||
Rejected: cloneRejectedOutputs(output.Rejected),
|
||||
Warnings: output.Warnings,
|
||||
LLMProfile: input.pipeline.Output.LLMProfile,
|
||||
Options: cloneOptions(input.pipeline.Output.Options),
|
||||
Metadata: input.Metadata,
|
||||
})
|
||||
output.Warnings = append(output.Warnings, encoded.Warnings...)
|
||||
|
||||
@@ -55,6 +55,7 @@ func newPreparedRunner(t *testing.T, registries Registries) preparedRunnerHarnes
|
||||
|
||||
func (h preparedRunnerHarness) Run(ctx context.Context, input RunInput) (RunOutput, error) {
|
||||
h.t.Helper()
|
||||
input.llmClient = WithDebugLLMRecording(input.llmClient, input.Debug)
|
||||
prepared, err := Prepare(input.pipeline, h.registries, ModuleDependencies{LLM: input.llmClient})
|
||||
if err != nil {
|
||||
return RunOutput{}, err
|
||||
@@ -541,7 +542,7 @@ func TestRunPreservesChunkMetadataDuringCanonicalization(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunExecutesChunksAndPassesChunkAndLLMClient(t *testing.T) {
|
||||
func TestRunExecutesChunksAndUsesConstructedLLMClient(t *testing.T) {
|
||||
modules := defaultRunnerModules()
|
||||
llmClient := fakeLLMClient{}
|
||||
|
||||
@@ -558,8 +559,8 @@ func TestRunExecutesChunksAndPassesChunkAndLLMClient(t *testing.T) {
|
||||
if !reflect.DeepEqual(extractor.seenChunkIDs, []string{"chunk-0", "chunk-1"}) {
|
||||
t.Fatalf("seen chunks = %#v, want both chunks", extractor.seenChunkIDs)
|
||||
}
|
||||
if len(modules.chunker.requests) != 1 || modules.chunker.requests[0].LLMClient == nil {
|
||||
t.Fatalf("chunker LLM client = %#v, want client on chunk request", modules.chunker.requests)
|
||||
if len(modules.chunker.requests) != 1 || modules.chunker.llmClient == nil {
|
||||
t.Fatalf("chunker LLM client = %#v, want injected client", modules.chunker.llmClient)
|
||||
}
|
||||
if len(extractor.seenLLMClients) != 2 || extractor.seenLLMClients[0] == nil || extractor.seenLLMClients[1] == nil {
|
||||
t.Fatalf("seen LLM clients = %#v, want client for each chunk", extractor.seenLLMClients)
|
||||
@@ -721,14 +722,14 @@ func TestRunPassesModuleBindingConfigToStageRequests(t *testing.T) {
|
||||
if got := modules.input.requests[0].LLMProfile; got != "input-profile" {
|
||||
t.Fatalf("input LLMProfile = %q, want input-profile", got)
|
||||
}
|
||||
if got := modules.input.requests[0].Options["input_option"]; got != "input-value" {
|
||||
t.Fatalf("input Options = %#v, want input option", modules.input.requests[0].Options)
|
||||
if got := modules.input.buildOptions["input_option"]; got != "input-value" {
|
||||
t.Fatalf("input construction options = %#v, want input option", modules.input.buildOptions)
|
||||
}
|
||||
if got := modules.chunker.requests[0].LLMProfile; got != "chunk-profile" {
|
||||
t.Fatalf("chunk LLMProfile = %q, want chunk-profile", got)
|
||||
}
|
||||
if got := modules.chunker.requests[0].Options["chunk_option"]; got != "chunk-value" {
|
||||
t.Fatalf("chunk Options = %#v, want chunk option", modules.chunker.requests[0].Options)
|
||||
if got := modules.chunker.buildOptions["chunk_option"]; got != "chunk-value" {
|
||||
t.Fatalf("chunk construction options = %#v, want chunk option", modules.chunker.buildOptions)
|
||||
}
|
||||
if got := modules.extractors["extract-alpha"].requests[0].LLMProfile; got != "extract-profile" {
|
||||
t.Fatalf("extract LLMProfile = %q, want extract-profile", got)
|
||||
@@ -751,8 +752,8 @@ func TestRunPassesModuleBindingConfigToStageRequests(t *testing.T) {
|
||||
if got := modules.output.requests[0].LLMProfile; got != "output-profile" {
|
||||
t.Fatalf("output LLMProfile = %q, want output-profile", got)
|
||||
}
|
||||
if got := modules.output.requests[0].Options["output_option"]; got != "output-value" {
|
||||
t.Fatalf("output Options = %#v, want output option", modules.output.requests[0].Options)
|
||||
if got := modules.output.buildOptions["output_option"]; got != "output-value" {
|
||||
t.Fatalf("output construction options = %#v, want output option", modules.output.buildOptions)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2034,18 +2035,22 @@ func newRunnerRegistries(t *testing.T, modules *runnerModules) Registries {
|
||||
ValidatorChains: NewValidatorChainRegistry(),
|
||||
Outputs: NewOutputEncoderRegistry(),
|
||||
}
|
||||
if err := registries.Inputs.Register("input", func() (contracts.InputAdapter, error) {
|
||||
allowOptions := func(map[string]any) error { return nil }
|
||||
if err := registries.Inputs.RegisterBuilderWithSpec(defaultModuleSpec("input", StageInput), allowOptions, func(request BuildRequest) (contracts.InputAdapter, error) {
|
||||
if modules.inputBuildErr != nil {
|
||||
return nil, modules.inputBuildErr
|
||||
}
|
||||
modules.input.buildOptions = cloneOptions(request.Options)
|
||||
return modules.input, nil
|
||||
}); err != nil {
|
||||
t.Fatalf("register input: %v", err)
|
||||
}
|
||||
if err := registries.Chunkers.Register("chunk", func() (contracts.Chunker, error) {
|
||||
if err := registries.Chunkers.RegisterBuilderWithSpec(defaultModuleSpec("chunk", StageChunk), allowOptions, func(request BuildRequest) (contracts.Chunker, error) {
|
||||
if modules.chunkerBuildErr != nil {
|
||||
return nil, modules.chunkerBuildErr
|
||||
}
|
||||
modules.chunker.buildOptions = cloneOptions(request.Options)
|
||||
modules.chunker.llmClient = request.Dependencies.LLM
|
||||
return modules.chunker, nil
|
||||
}); err != nil {
|
||||
t.Fatalf("register chunker: %v", err)
|
||||
@@ -2075,7 +2080,10 @@ func newRunnerRegistries(t *testing.T, modules *runnerModules) Registries {
|
||||
t.Fatalf("register validator %q: %v", key, err)
|
||||
}
|
||||
}
|
||||
if err := registries.Outputs.Register("output", func() (contracts.OutputEncoder, error) { return modules.output, nil }); err != nil {
|
||||
if err := registries.Outputs.RegisterBuilderWithSpec(defaultModuleSpec("output", StageOutput), allowOptions, func(request BuildRequest) (contracts.OutputEncoder, error) {
|
||||
modules.output.buildOptions = cloneOptions(request.Options)
|
||||
return modules.output, nil
|
||||
}); err != nil {
|
||||
t.Fatalf("register output: %v", err)
|
||||
}
|
||||
return registries
|
||||
@@ -2087,6 +2095,7 @@ type runnerInputAdapter struct {
|
||||
err error
|
||||
manifestMetadata map[string]any
|
||||
requests []contracts.ParseRequest
|
||||
buildOptions map[string]any
|
||||
}
|
||||
|
||||
func (adapter *runnerInputAdapter) Key() string {
|
||||
@@ -2113,6 +2122,8 @@ type runnerChunker struct {
|
||||
llmPromptID string
|
||||
manifestMetadata map[string]any
|
||||
requests []contracts.ChunkRequest
|
||||
buildOptions map[string]any
|
||||
llmClient contracts.StructuredLLMClient
|
||||
}
|
||||
|
||||
func (chunker *runnerChunker) Key() string {
|
||||
@@ -2133,13 +2144,13 @@ func (chunker *runnerChunker) Chunk(ctx context.Context, req contracts.ChunkRequ
|
||||
}
|
||||
return contracts.ChunkResult{}, err
|
||||
}
|
||||
if chunker.callLLM && req.LLMClient != nil {
|
||||
if chunker.callLLM && chunker.llmClient != nil {
|
||||
promptID := strings.TrimSpace(chunker.llmPromptID)
|
||||
if promptID == "" {
|
||||
promptID = "runner.chunk"
|
||||
}
|
||||
var out map[string]any
|
||||
if _, err := req.LLMClient.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
||||
if _, err := chunker.llmClient.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
||||
StageName: chunker.key,
|
||||
PromptID: promptID,
|
||||
ProfileID: req.LLMProfile,
|
||||
@@ -2402,6 +2413,7 @@ type runnerOutputEncoder struct {
|
||||
err error
|
||||
manifestMetadata map[string]any
|
||||
requests []contracts.OutputRequest
|
||||
buildOptions map[string]any
|
||||
}
|
||||
|
||||
func (encoder *runnerOutputEncoder) Key() string {
|
||||
|
||||
Reference in New Issue
Block a user