notarius/internal/cli/catalog.go

package cli

import (
	"context"
	"fmt"

	"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
	"gitea.maximumdirect.net/eric/notarius/internal/core/config"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/chunk/dnd/scenes"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/chunk/generic"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/promptassets"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/extract/dnd/spells"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/input/seriatim"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/merge/appendorder"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/normalize/noop"
	jsonoutput "gitea.maximumdirect.net/eric/notarius/internal/modules/output/json"
)

func productionRegistries() (pipeline.Registries, error) {
	registries := pipeline.Registries{
		Inputs:      pipeline.NewInputAdapterRegistry(),
		Chunkers:    pipeline.NewChunkerRegistry(),
		Extractors:  pipeline.NewExtractorRegistry(),
		Mergers:     pipeline.NewMergerRegistry(),
		Normalizers: pipeline.NewNormalizerRegistry(),
		Validators:  pipeline.NewValidatorRegistry(),
		Outputs:     pipeline.NewOutputEncoderRegistry(),
	}
	if err := seriatim.Register(registries.Inputs); err != nil {
		return pipeline.Registries{}, fmt.Errorf("register seriatim input: %w", err)
	}
	if err := generic.Register(registries.Chunkers); err != nil {
		return pipeline.Registries{}, fmt.Errorf("register generic chunker: %w", err)
	}
	if err := scenes.Register(registries.Chunkers); err != nil {
		return pipeline.Registries{}, fmt.Errorf("register dnd scenes chunker: %w", err)
	}
	if err := spells.Register(registries.Extractors); err != nil {
		return pipeline.Registries{}, fmt.Errorf("register dnd spells extractor: %w", err)
	}
	if err := appendorder.Register(registries.Mergers); err != nil {
		return pipeline.Registries{}, fmt.Errorf("register appendorder merger: %w", err)
	}
	if err := noop.Register(registries.Normalizers); err != nil {
		return pipeline.Registries{}, fmt.Errorf("register noop normalizer: %w", err)
	}
	if err := jsonoutput.Register(registries.Outputs); err != nil {
		return pipeline.Registries{}, fmt.Errorf("register json output encoder: %w", err)
	}
	return registries, nil
}

func productionCatalog() (pipeline.ModuleCatalog, error) {
	registries, err := productionRegistries()
	if err != nil {
		return pipeline.ModuleCatalog{}, err
	}
	return catalogFromRegistries(registries), nil
}

func productionPromptAssets() (*llm.AssetRegistry, error) {
	registry := llm.NewAssetRegistry()
	if err := promptassets.Register(registry); err != nil {
		return nil, fmt.Errorf("register shared dnd prompt assets: %w", err)
	}
	if err := scenes.RegisterPromptAssets(registry); err != nil {
		return nil, fmt.Errorf("register dnd scenes prompt assets: %w", err)
	}
	if err := spells.RegisterPromptAssets(registry); err != nil {
		return nil, fmt.Errorf("register dnd spells prompt assets: %w", err)
	}
	return registry, nil
}

func effectiveCatalog(opts Options) (pipeline.ModuleCatalog, error) {
	if !isEmptyCatalog(opts.Catalog) {
		return opts.Catalog, nil
	}
	if !isEmptyRegistries(opts.Registries) {
		return catalogFromRegistries(opts.Registries), nil
	}
	return productionCatalog()
}

func effectiveRegistries(opts Options) (pipeline.Registries, error) {
	if !isEmptyRegistries(opts.Registries) {
		return opts.Registries, nil
	}
	if !isEmptyCatalog(opts.Catalog) {
		return registriesFromCatalog(opts.Catalog), nil
	}
	return productionRegistries()
}

func catalogFromRegistries(registries pipeline.Registries) pipeline.ModuleCatalog {
	return pipeline.ModuleCatalog{
		Inputs:      registries.Inputs,
		Chunkers:    registries.Chunkers,
		Extractors:  registries.Extractors,
		Mergers:     registries.Mergers,
		Normalizers: registries.Normalizers,
		Validators:  registries.Validators,
		Outputs:     registries.Outputs,
	}
}

func registriesFromCatalog(catalog pipeline.ModuleCatalog) pipeline.Registries {
	return pipeline.Registries{
		Inputs:      catalog.Inputs,
		Chunkers:    catalog.Chunkers,
		Extractors:  catalog.Extractors,
		Mergers:     catalog.Mergers,
		Normalizers: catalog.Normalizers,
		Validators:  catalog.Validators,
		Outputs:     catalog.Outputs,
	}
}

func isEmptyCatalog(catalog pipeline.ModuleCatalog) bool {
	return catalog.Inputs == nil &&
		catalog.Chunkers == nil &&
		catalog.Extractors == nil &&
		catalog.Mergers == nil &&
		catalog.Normalizers == nil &&
		catalog.Validators == nil &&
		catalog.Outputs == nil
}

func isEmptyRegistries(registries pipeline.Registries) bool {
	return registries.Inputs == nil &&
		registries.Chunkers == nil &&
		registries.Extractors == nil &&
		registries.Mergers == nil &&
		registries.Normalizers == nil &&
		registries.Validators == nil &&
		registries.Outputs == nil
}

func productionLLMClientFactory(ctx context.Context, cfg config.Config, profileID string) (contracts.StructuredLLMClient, []artifacts.LLMProfileManifest, error) {
	if err := ctx.Err(); err != nil {
		return nil, nil, err
	}
	assets, err := productionPromptAssets()
	if err != nil {
		return nil, nil, err
	}
	recorder := llm.NewLLMProfileRecorder()
	client, err := llm.NewScriptoriumClient(llm.ScriptoriumClientConfig{
		ProfileDir:  cfg.Scriptorium.ProfileDir,
		ProfileFile: cfg.Scriptorium.ProfileFile,
		Assets:      assets,
		Recorder:    recorder,
	})
	if err != nil {
		return nil, nil, fmt.Errorf("create Scriptorium-backed LLM client: %w", err)
	}
	scheduler, err := llm.NewScheduler(cfg.Concurrency.TotalLLM)
	if err != nil {
		return nil, nil, fmt.Errorf("create LLM scheduler: %w", err)
	}
	return llm.NewScheduledClient(client, scheduler), nil, nil
}