Refactor LLM config into top-level settings

Replace the monolithic agents config with a cleaner structure: - llm: holds provider credentials and default model settings - probo-agent: LLM overrides for the probo agent - evidence-describer: worker config (interval, stale-after, max-concurrency) alongside LLM overrides This makes worker tuning configurable via YAML and env vars instead of being hardcoded in Go, and separates provider credentials from per-consumer model settings. Signed-off-by: Bryan Frimin <bryan@getprobo.com>
2026-03-27 07:48:43 +01:00
parent 0926a8828a
commit 050154ab6a
6 changed files with 179 additions and 113 deletions
--- a/pkg/probod/evidence_describer_config.go
+++ b/pkg/probod/evidence_describer_config.go
@@ -0,0 +1,38 @@
+// Copyright (c) 2025-2026 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package probod
+
+// EvidenceDescriberConfig holds both the worker settings and LLM overrides
+// for the evidence description worker.
+type EvidenceDescriberConfig struct {
+	Interval       int `json:"interval"`    // seconds
+	StaleAfter     int `json:"stale-after"` // seconds
+	MaxConcurrency int `json:"max-concurrency"`
+
+	Provider    string   `json:"provider"`
+	ModelName   string   `json:"model-name"`
+	Temperature *float64 `json:"temperature"`
+	MaxTokens   *int     `json:"max-tokens"`
+}
+
+// LLMConfig extracts the LLM-specific fields as an LLMConfig.
+func (c *EvidenceDescriberConfig) LLMConfig() LLMConfig {
+	return LLMConfig{
+		Provider:    c.Provider,
+		ModelName:   c.ModelName,
+		Temperature: c.Temperature,
+		MaxTokens:   c.MaxTokens,
+	}
+}
--- a/pkg/probod/openai_config.go
+++ b/pkg/probod/openai_config.go
@@ -22,40 +22,38 @@ type (
 		APIKey string `json:"api-key"` // for OpenAI and Anthropic
 	}

-	// LLMAgentConfig holds model parameters for a single agent. Provider
-	// references one of the keys in AgentsConfig.Providers.
-	LLMAgentConfig struct {
-		Provider    string   `json:"provider"` // key into AgentsConfig.Providers
+	// LLMConfig holds model parameters for a single LLM consumer. Provider
+	// references one of the keys in LLMSettings.Providers.
+	LLMConfig struct {
+		Provider    string   `json:"provider"` // key into LLMSettings.Providers
 		ModelName   string   `json:"model-name"`
 		Temperature *float64 `json:"temperature"`
 		MaxTokens   *int     `json:"max-tokens"`
 	}

-	// AgentsConfig groups LLM provider credentials and per-agent model
-	// settings. Default is used as a fallback when an agent-specific field
-	// is zero-valued.
-	AgentsConfig struct {
-		Providers         map[string]LLMProviderConfig `json:"providers"`
-		Default           LLMAgentConfig               `json:"default"`
-		Probo             LLMAgentConfig               `json:"probo"`
-		EvidenceDescriber LLMAgentConfig               `json:"evidence-describer"`
+	// LLMSettings groups LLM provider credentials and default model
+	// settings. Defaults is used as a fallback when a consumer-specific
+	// field is zero-valued.
+	LLMSettings struct {
+		Providers map[string]LLMProviderConfig `json:"providers"`
+		Defaults  LLMConfig                    `json:"defaults"`
 	}
 )

-// ResolveAgent returns a fully populated LLMAgentConfig by filling in
-// zero-valued fields from the default config.
-func (c *AgentsConfig) ResolveAgent(agent LLMAgentConfig) LLMAgentConfig {
-	if agent.Provider == "" {
-		agent.Provider = c.Default.Provider
+// ResolveLLMConfig returns a fully populated LLMConfig by filling in
+// zero-valued fields from the defaults.
+func (s *LLMSettings) ResolveLLMConfig(cfg LLMConfig) LLMConfig {
+	if cfg.Provider == "" {
+		cfg.Provider = s.Defaults.Provider
 	}
-	if agent.ModelName == "" {
-		agent.ModelName = c.Default.ModelName
+	if cfg.ModelName == "" {
+		cfg.ModelName = s.Defaults.ModelName
 	}
-	if agent.Temperature == nil {
-		agent.Temperature = c.Default.Temperature
+	if cfg.Temperature == nil {
+		cfg.Temperature = s.Defaults.Temperature
 	}
-	if agent.MaxTokens == nil {
-		agent.MaxTokens = c.Default.MaxTokens
+	if cfg.MaxTokens == nil {
+		cfg.MaxTokens = s.Defaults.MaxTokens
 	}
-	return agent
+	return cfg
 }
--- a/pkg/probod/probod.go
+++ b/pkg/probod/probod.go
@@ -108,20 +108,22 @@ type (

 	// Config represents the probod application configuration.
 	Config struct {
-		BaseURL       string              `json:"base-url"`
-		EncryptionKey string              `json:"encryption-key"`
-		Pg            PgConfig            `json:"pg"`
-		Api           APIConfig           `json:"api"`
-		Auth          AuthConfig          `json:"auth"`
-		TrustCenter   TrustCenterConfig   `json:"trust-center"`
-		AWS           AWSConfig           `json:"aws"`
-		Notifications NotificationsConfig `json:"notifications"`
-		Connectors    []ConnectorConfig   `json:"connectors"`
-		Agents        AgentsConfig        `json:"agents"`
-		ChromeDPAddr  string              `json:"chrome-dp-addr"`
-		CustomDomains CustomDomainsConfig `json:"custom-domains"`
-		SCIMBridge    SCIMBridgeConfig    `json:"scim-bridge"`
-		ESign         ESignConfig         `json:"esign"`
+		BaseURL           string                  `json:"base-url"`
+		EncryptionKey     string                  `json:"encryption-key"`
+		Pg                PgConfig                `json:"pg"`
+		Api               APIConfig               `json:"api"`
+		Auth              AuthConfig              `json:"auth"`
+		TrustCenter       TrustCenterConfig       `json:"trust-center"`
+		AWS               AWSConfig               `json:"aws"`
+		Notifications     NotificationsConfig     `json:"notifications"`
+		Connectors        []ConnectorConfig       `json:"connectors"`
+		LLM               LLMSettings             `json:"llm"`
+		ProboAgent        LLMConfig               `json:"probo-agent"`
+		EvidenceDescriber EvidenceDescriberConfig `json:"evidence-describer"`
+		ChromeDPAddr      string                  `json:"chrome-dp-addr"`
+		CustomDomains     CustomDomainsConfig     `json:"custom-domains"`
+		SCIMBridge        SCIMBridgeConfig        `json:"scim-bridge"`
+		ESign             ESignConfig             `json:"esign"`
 	}

 	// TrustCenterConfig contains trust center server configuration.
@@ -217,6 +219,11 @@ func New() *Implm {
 			ESign: ESignConfig{
 				TSAURL: "http://timestamp.digicert.com",
 			},
+			EvidenceDescriber: EvidenceDescriberConfig{
+				Interval:       10,
+				StaleAfter:     300,
+				MaxConcurrency: 10,
+			},
 		},
 	}
 }
@@ -318,8 +325,8 @@ func (impl *Implm) Run(
 		}
 	}

-	proboAgentCfg := impl.cfg.Agents.ResolveAgent(impl.cfg.Agents.Probo)
-	proboProviderCfg, ok := impl.cfg.Agents.Providers[proboAgentCfg.Provider]
+	proboAgentCfg := impl.cfg.LLM.ResolveLLMConfig(impl.cfg.ProboAgent)
+	proboProviderCfg, ok := impl.cfg.LLM.Providers[proboAgentCfg.Provider]
 	if !ok {
 		return fmt.Errorf("unknown LLM provider %q for probo agent", proboAgentCfg.Provider)
 	}
@@ -328,12 +335,12 @@ func (impl *Implm) Run(
 		return fmt.Errorf("cannot create probo LLM client: %w", err)
 	}

-	evidenceDescriberAgentCfg := impl.cfg.Agents.ResolveAgent(impl.cfg.Agents.EvidenceDescriber)
-	evidenceDescriberProviderCfg, ok := impl.cfg.Agents.Providers[evidenceDescriberAgentCfg.Provider]
+	edLLMCfg := impl.cfg.LLM.ResolveLLMConfig(impl.cfg.EvidenceDescriber.LLMConfig())
+	edProviderCfg, ok := impl.cfg.LLM.Providers[edLLMCfg.Provider]
 	if !ok {
-		return fmt.Errorf("unknown LLM provider %q for evidence-describer agent", evidenceDescriberAgentCfg.Provider)
+		return fmt.Errorf("unknown LLM provider %q for evidence-describer agent", edLLMCfg.Provider)
 	}
-	evidenceDescriberLLMClient, err := buildLLMClient(evidenceDescriberProviderCfg, l.Named("llm.evidence-describer"), tp, r)
+	evidenceDescriberLLMClient, err := buildLLMClient(edProviderCfg, l.Named("llm.evidence-describer"), tp, r)
 	if err != nil {
 		return fmt.Errorf("cannot create evidence describer LLM client: %w", err)
 	}
@@ -655,9 +662,9 @@ func (impl *Implm) Run(
 	evidenceDescriber := evidencedescriber.New(
 		evidenceDescriberLLMClient,
 		evidencedescriber.Config{
-			Model:     evidenceDescriberAgentCfg.ModelName,
-			Temp:      *evidenceDescriberAgentCfg.Temperature,
-			MaxTokens: *evidenceDescriberAgentCfg.MaxTokens,
+			Model:     edLLMCfg.ModelName,
+			Temp:      *edLLMCfg.Temperature,
+			MaxTokens: *edLLMCfg.MaxTokens,
 		},
 	)
 	evidenceDescriptionWorker := probo.NewEvidenceDescriptionWorker(
@@ -665,6 +672,9 @@ func (impl *Implm) Run(
 		fileManagerService,
 		evidenceDescriber,
 		l.Named("evidence-description-worker"),
+		probo.WithEvidenceDescriptionWorkerInterval(time.Duration(impl.cfg.EvidenceDescriber.Interval)*time.Second),
+		probo.WithEvidenceDescriptionWorkerStaleAfter(time.Duration(impl.cfg.EvidenceDescriber.StaleAfter)*time.Second),
+		probo.WithEvidenceDescriptionWorkerMaxConcurrency(impl.cfg.EvidenceDescriber.MaxConcurrency),
 	)
 	evidenceDescriptionWorkerCtx, stopEvidenceDescriptionWorker := context.WithCancel(context.Background())
 	wg.Go(