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

Wire persistent chunk plan caching into the CLI

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-18 00:31:53 +00:00
parent 51a36efb6b
commit 8ba5228c01
5 changed files with 601 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
internal/cli/run.go
View File

@@ -20,6 +20,7 @@ import (
"gitea.maximumdirect.net/eric/notarius/internal/core/diagnostics"
"gitea.maximumdirect.net/eric/notarius/internal/core/workspace"
"gitea.maximumdirect.net/eric/notarius/internal/framework/checkpoint"
"gitea.maximumdirect.net/eric/notarius/internal/framework/chunkplan"
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
frameworkdebug "gitea.maximumdirect.net/eric/notarius/internal/framework/debug"
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
@@ -30,17 +31,19 @@ const defaultOutputRoot = "./notarius-output"
const usage = `Usage:
notarius help
notarius run <pipeline-id> --input path/to/source.json [--config path/to/config.yml] [--only lane-a,lane-b] [--resume] [--session-id id] [--reference selector=path] [--without-reference selector]
notarius run <pipeline-id> --input path/to/source.json [--config path/to/config.yml] [--only lane-a,lane-b] [--chunk_cache auto|bypass|refresh] [--resume] [--session-id id] [--reference selector=path] [--without-reference selector]
notarius config validate --config path/to/config.yml [--pipeline pipeline-id] [--only lane-a,lane-b]
notarius pipelines list --config path/to/config.yml [--json]
`
type Options struct {
Catalog pipeline.ModuleCatalog
Registries pipeline.Registries
LLMClientFactory LLMClientFactory
LookupEnv func(string) (string, bool)
Now func() time.Time
Catalog pipeline.ModuleCatalog
Registries pipeline.Registries
LLMClientFactory LLMClientFactory
LookupEnv func(string) (string, bool)
Now func() time.Time
UserCacheDir func() (string, error)
ChunkPlanStoreFactory pipeline.ChunkPlanStoreFactory
}
type LLMClientFactory func(ctx context.Context, cfg config.Config, profileID string) (contracts.StructuredLLMClient, []artifacts.LLMProfileManifest, error)
@@ -90,6 +93,12 @@ func normalizeOptions(opts Options) (Options, error) {
if opts.Now == nil {
opts.Now = time.Now
}
if opts.UserCacheDir == nil {
opts.UserCacheDir = os.UserCacheDir
}
if opts.ChunkPlanStoreFactory == nil {
opts.ChunkPlanStoreFactory = chunkplan.NewFilesystemStore
}
if isEmptyCatalog(opts.Catalog) && isEmptyRegistries(opts.Registries) {
components, err := newProductionComponents()
if err != nil {
@@ -117,10 +126,12 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
diagnosticsDir := fs.String("diagnostics-dir", "", "diagnostics directory")
llmProfile := fs.String("llm-profile", "", "LLM profile override")
resume := fs.Bool("resume", false, "reuse valid workspace checkpoints")
chunkCache := chunkCacheFlag{}
sessionID := sessionIDFlag{}
referenceFlags := stringListFlag{}
withoutReferenceFlags := stringListFlag{}
fs.Var(&sessionID, "session-id", "prompt session identifier")
fs.Var(&chunkCache, "chunk_cache", "chunk plan cache mode: auto, bypass, or refresh")
fs.Var(&referenceFlags, "reference", "reference binding, as slot=path, chunk.slot=path, merge.slot=path, lane.slot=path, lane.extract.slot=path, lane.merge.slot=path, or lane.normalize.slot=path")
fs.Var(&withoutReferenceFlags, "without-reference", "unbind a reference, using the same selector forms as --reference")
if err := validateRunFlagValues(args); err != nil {
@@ -173,6 +184,9 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
fmt.Fprintf(stderr, "notarius: %v\n", err)
return 1
}
if chunkCache.set {
cfg.Workspace.ChunkCache.Mode = chunkCache.value
}
workspaceSettings := workspace.FromConfig(cfg)
if dir := strings.TrimSpace(*diagnosticsDir); dir != "" {
workspaceSettings.DiagnosticsRoot = dir
@@ -191,15 +205,16 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
runID = runDir.RunID()
}
invocation := diagnostics.InvocationMetadata{
Operation: "run",
PipelineID: pipelineID,
InputPath: strings.TrimSpace(*inputPath),
ConfigPath: loadedConfigPath,
ConfigSource: configSource(*configPath),
OnlyLanes: append([]string(nil), only...),
Resume: *resume,
RunID: runID,
StartedAt: startedAt,
Operation: "run",
PipelineID: pipelineID,
InputPath: strings.TrimSpace(*inputPath),
ConfigPath: loadedConfigPath,
ConfigSource: configSource(*configPath),
OnlyLanes: append([]string(nil), only...),
ChunkCacheOverride: chunkCache.explicitValue(),
Resume: *resume,
RunID: runID,
StartedAt: startedAt,
}
if err := writeDiagnostics(runDir, func() error { return runDir.WriteInvocationMetadata(invocation) }); err != nil {
return failPipelineCommand(stderr, runDir, cfg.Diagnostics.Retention, fmt.Errorf("write diagnostics invocation metadata: %w", err))
@@ -287,6 +302,10 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
if err != nil {
return failPipelineCommand(stderr, runDir, cfg.Diagnostics.Retention, fmt.Errorf("read input %q: %w", strings.TrimSpace(*inputPath), err))
}
chunkPlans, err := chunkPlanStoreForRun(effective.Config.Workspace.ChunkCache, opts)
if err != nil {
return failPipelineCommand(stderr, runDir, cfg.Diagnostics.Retention, err)
}
checkpointRecorder, checkpointLoader, err := checkpointHandlersForRun(workspaceSettings, effective.ResolvedPipeline, rawInput, only, llmProfiles, strings.TrimSpace(*llmProfile), strings.TrimSpace(sessionID.value), *resume)
if err != nil {
return failPipelineCommand(stderr, runDir, cfg.Diagnostics.Retention, err)
@@ -302,6 +321,8 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
LLMProfiles: llmProfiles,
Metadata: runMetadata(*outputDir, *diagnosticsDir),
Warnings: referenceWarnings,
ChunkCacheMode: effective.Config.Workspace.ChunkCache.Mode,
ChunkPlans: chunkPlans,
Checkpoints: checkpointRecorder,
Checkpoint: checkpointLoader,
Debug: debugRecorder,
@@ -598,13 +619,64 @@ func reorderRunArgs(args []string) []string {
func runFlagTakesValue(arg string) bool {
switch arg {
case "--config", "--input", "--only", "--output-dir", "--diagnostics-dir", "--llm-profile", "--session-id", "--reference", "--without-reference":
case "--config", "--input", "--only", "--output-dir", "--diagnostics-dir", "--llm-profile", "--session-id", "--chunk_cache", "--reference", "--without-reference":
return true
default:
return false
}
}
type chunkCacheFlag struct {
value pipeline.ChunkCacheMode
set bool
}
func (f *chunkCacheFlag) String() string {
if f == nil {
return ""
}
return string(f.value)
}
func (f *chunkCacheFlag) Set(raw string) error {
mode, err := pipeline.ParseChunkCacheMode(raw)
if err != nil {
return err
}
f.value = mode
f.set = true
return nil
}
func (f chunkCacheFlag) explicitValue() string {
if !f.set {
return ""
}
return string(f.value)
}
func chunkPlanStoreForRun(cfg config.WorkspaceChunkCacheConfig, opts Options) (pipeline.ChunkPlanStore, error) {
if cfg.Mode == pipeline.ChunkCacheBypass {
return nil, nil
}
root := strings.TrimSpace(cfg.Directory)
if root == "" {
var err error
root, err = workspace.DefaultChunkPlanRoot(opts.UserCacheDir)
if err != nil {
return nil, fmt.Errorf("resolve chunk plan root: %w", err)
}
}
store, err := opts.ChunkPlanStoreFactory(root)
if err != nil {
return nil, fmt.Errorf("create chunk plan store at %q: %w", root, err)
}
if store == nil {
return nil, fmt.Errorf("create chunk plan store at %q: factory returned nil", root)
}
return store, nil
}
func validateRunFlagValues(args []string) error {
for i, arg := range args {
if arg != "--session-id" {