243 lines
6.6 KiB
Go
243 lines
6.6 KiB
Go
package config
|
|
|
|
import (
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
|
)
|
|
|
|
const SupportedFileConfigVersion = 4
|
|
|
|
type Config struct {
|
|
PromptKit PromptKitConfig `json:"promptkit,omitempty"`
|
|
Pipelines map[string]pipeline.PipelineProfile `json:"pipelines"`
|
|
Concurrency ConcurrencyConfig `json:"concurrency"`
|
|
Output OutputConfig `json:"output"`
|
|
Cache CacheConfig `json:"cache"`
|
|
Debug DebugConfig `json:"debug"`
|
|
}
|
|
|
|
type PromptKitConfig struct {
|
|
ProfileDir string `json:"profile_dir,omitempty"`
|
|
ProfileFile string `json:"profile_file,omitempty"`
|
|
LocalBackend *PromptKitLocalBackendConfig `json:"local_backend,omitempty"`
|
|
}
|
|
|
|
type PromptKitLocalBackendConfig struct {
|
|
Endpoint string `json:"endpoint"`
|
|
ConcurrencyLimit int `json:"concurrency_limit"`
|
|
}
|
|
|
|
type ConcurrencyConfig struct {
|
|
TotalLLM int `json:"total_llm"`
|
|
StageWorkers map[string]int `json:"stage_workers"`
|
|
|
|
extractWorkersConfigured bool
|
|
defaultedExtractWorkers int
|
|
}
|
|
|
|
const defaultLLMConcurrency = 16
|
|
|
|
type OutputConfig struct {
|
|
Directory string `json:"directory"`
|
|
}
|
|
|
|
type CacheConfig struct {
|
|
ChunkPlans ChunkPlanCacheConfig `json:"chunk_plans"`
|
|
Checkpoints CheckpointCacheConfig `json:"checkpoints"`
|
|
}
|
|
|
|
type ChunkPlanCacheConfig struct {
|
|
Directory string `json:"directory,omitempty"`
|
|
Mode pipeline.ChunkCacheMode `json:"mode"`
|
|
}
|
|
|
|
type CheckpointCacheConfig struct {
|
|
Enabled bool `json:"enabled"`
|
|
Directory string `json:"directory,omitempty"`
|
|
}
|
|
type DebugConfig struct {
|
|
Directory string `json:"directory"`
|
|
}
|
|
|
|
func Default() Config {
|
|
return Config{
|
|
Pipelines: map[string]pipeline.PipelineProfile{},
|
|
Concurrency: ConcurrencyConfig{
|
|
TotalLLM: defaultLLMConcurrency,
|
|
StageWorkers: map[string]int{"extract": defaultLLMConcurrency},
|
|
defaultedExtractWorkers: defaultLLMConcurrency,
|
|
},
|
|
Output: OutputConfig{Directory: "./notarius-output"},
|
|
Cache: CacheConfig{ChunkPlans: ChunkPlanCacheConfig{Mode: pipeline.ChunkCacheAuto}},
|
|
Debug: DebugConfig{Directory: "./notarius-debug"},
|
|
}
|
|
}
|
|
|
|
func cloneConfig(in Config) Config {
|
|
out := in
|
|
if in.PromptKit.LocalBackend != nil {
|
|
localBackend := *in.PromptKit.LocalBackend
|
|
out.PromptKit.LocalBackend = &localBackend
|
|
}
|
|
out.Concurrency.StageWorkers = cloneIntMap(in.Concurrency.StageWorkers)
|
|
out.Pipelines = make(map[string]pipeline.PipelineProfile, len(in.Pipelines))
|
|
for key, profile := range in.Pipelines {
|
|
out.Pipelines[key] = clonePipelineProfile(profile)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func cloneIntMap(in map[string]int) map[string]int {
|
|
if len(in) == 0 {
|
|
return nil
|
|
}
|
|
out := make(map[string]int, len(in))
|
|
for key, value := range in {
|
|
out[key] = value
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (c *ConcurrencyConfig) recomputeStageWorkerDefaults() {
|
|
if c == nil {
|
|
return
|
|
}
|
|
if c.StageWorkers == nil {
|
|
c.StageWorkers = make(map[string]int)
|
|
}
|
|
if !c.extractWorkersConfigured {
|
|
if value, ok := c.StageWorkers["extract"]; ok && (c.defaultedExtractWorkers == 0 || value != c.defaultedExtractWorkers) {
|
|
c.extractWorkersConfigured = true
|
|
return
|
|
}
|
|
c.StageWorkers["extract"] = c.TotalLLM
|
|
c.defaultedExtractWorkers = c.TotalLLM
|
|
}
|
|
}
|
|
|
|
func clonePipelineProfile(in pipeline.PipelineProfile) pipeline.PipelineProfile {
|
|
out := in
|
|
out.Input = cloneModuleBinding(in.Input)
|
|
out.Chunk = cloneModuleBinding(in.Chunk)
|
|
out.Output = cloneModuleBinding(in.Output)
|
|
out.References = cloneReferenceSourceMap(in.References)
|
|
if len(in.Artifacts) > 0 {
|
|
out.Artifacts = make(map[string]pipeline.ArtifactLaneProfile, len(in.Artifacts))
|
|
for key, lane := range in.Artifacts {
|
|
out.Artifacts[key] = cloneArtifactLaneProfile(lane)
|
|
}
|
|
}
|
|
if in.Steps != nil {
|
|
out.Steps = make([]pipeline.PipelineStepProfile, len(in.Steps))
|
|
for i, step := range in.Steps {
|
|
out.Steps[i] = clonePipelineStepProfile(step)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func clonePipelineStepProfile(in pipeline.PipelineStepProfile) pipeline.PipelineStepProfile {
|
|
out := in
|
|
out.ID = in.ID
|
|
out.References = cloneReferenceSourceMap(in.References)
|
|
if len(in.Artifacts) > 0 {
|
|
out.Artifacts = make(map[string]pipeline.ArtifactLaneProfile, len(in.Artifacts))
|
|
for key, lane := range in.Artifacts {
|
|
out.Artifacts[key] = cloneArtifactLaneProfile(lane)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func cloneArtifactLaneProfile(in pipeline.ArtifactLaneProfile) pipeline.ArtifactLaneProfile {
|
|
out := in
|
|
out.Extract = cloneModuleBinding(in.Extract)
|
|
out.Merge = cloneModuleBinding(in.Merge)
|
|
out.Normalize = cloneModuleBinding(in.Normalize)
|
|
out.References = cloneReferenceSourceMap(in.References)
|
|
if len(in.Validators) > 0 {
|
|
out.Validators = make([]pipeline.ModuleBinding, len(in.Validators))
|
|
for i, binding := range in.Validators {
|
|
out.Validators[i] = cloneModuleBinding(binding)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func cloneStringMap(in map[string]string) map[string]string {
|
|
if len(in) == 0 {
|
|
return nil
|
|
}
|
|
out := make(map[string]string, len(in))
|
|
for key, value := range in {
|
|
out[key] = value
|
|
}
|
|
return out
|
|
}
|
|
|
|
func cloneReferenceSourceMap(in map[string]pipeline.ReferenceSource) map[string]pipeline.ReferenceSource {
|
|
if len(in) == 0 {
|
|
return nil
|
|
}
|
|
out := make(map[string]pipeline.ReferenceSource, len(in))
|
|
for key, source := range in {
|
|
out[key] = cloneReferenceSource(source)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func cloneReferenceSource(in pipeline.ReferenceSource) pipeline.ReferenceSource {
|
|
out := in
|
|
if in.Artifact != nil {
|
|
artifact := *in.Artifact
|
|
out.Artifact = &artifact
|
|
}
|
|
return out
|
|
}
|
|
|
|
func cloneModuleBinding(in pipeline.ModuleBinding) pipeline.ModuleBinding {
|
|
out := in
|
|
if len(in.Options) > 0 {
|
|
out.Options = cloneOptions(in.Options)
|
|
}
|
|
out.References = cloneReferenceSourceMap(in.References)
|
|
out.Validators = cloneValidatorOverride(in.Validators)
|
|
return out
|
|
}
|
|
|
|
func cloneValidatorOverride(in pipeline.ValidatorOverride) pipeline.ValidatorOverride {
|
|
out := pipeline.ValidatorOverride{Set: in.Set}
|
|
if len(in.Validators) > 0 {
|
|
out.Validators = make([]pipeline.ModuleBinding, len(in.Validators))
|
|
for i, binding := range in.Validators {
|
|
out.Validators[i] = cloneModuleBinding(binding)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func cloneOptions(in map[string]any) map[string]any {
|
|
if len(in) == 0 {
|
|
return nil
|
|
}
|
|
out := make(map[string]any, len(in))
|
|
for key, value := range in {
|
|
out[key] = cloneOptionValue(value)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func cloneOptionValue(value any) any {
|
|
switch typed := value.(type) {
|
|
case map[string]any:
|
|
return cloneOptions(typed)
|
|
case []any:
|
|
out := make([]any, len(typed))
|
|
for i, item := range typed {
|
|
out[i] = cloneOptionValue(item)
|
|
}
|
|
return out
|
|
default:
|
|
return typed
|
|
}
|
|
}
|