rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor LLM config into top-level settings

Browse Source 
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>
This commit is contained in:
Bryan Frimin
2026-03-27 07:48:43 +01:00
committed by Sacha Al Himdani
parent 0926a8828a
commit 050154ab6a
6 changed files with 179 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
cfg/dev.yaml
View File

@@ -85,17 +85,22 @@ probod:
slack:
sender-interval: 60
agents:
llm:
providers:
openai:
type: "openai"
type: openai
api-key: "thisisnotasecret"
default:
provider: "openai"
model-name: "gpt-4o"
defaults:
provider: openai
model-name: gpt-4o
temperature: 0.1
max-tokens: 4096
evidence-describer:
interval: 10
stale-after: 300
max-concurrency: 10
custom-domains:
renewal-interval: 3600
provision-interval: 30

39
pkg/bootstrap/builder.go
View File

@@ -157,7 +157,7 @@ func (b *Builder) Build() (*probod.FullConfig, error) {
CacheTTL: b.getEnvIntOrDefault("WEBHOOK_CACHE_TTL", 86400),
},
},
Agents: probod.AgentsConfig{
LLM: probod.LLMSettings{
Providers: map[string]probod.LLMProviderConfig{
"openai": {
Type: "openai",
@@ -168,25 +168,28 @@ func (b *Builder) Build() (*probod.FullConfig, error) {
APIKey: b.getEnv("ANTHROPIC_API_KEY"),
},
},
Default: probod.LLMAgentConfig{
Provider: b.getEnvOrDefault("AGENT_DEFAULT_PROVIDER", "openai"),
ModelName: b.getEnvOrDefault("AGENT_DEFAULT_MODEL_NAME", "gpt-4o"),
Temperature: new(b.getEnvFloatOrDefault("AGENT_DEFAULT_TEMPERATURE", 0.1)),
MaxTokens: new(b.getEnvIntOrDefault("AGENT_DEFAULT_MAX_TOKENS", 4096)),
},
Probo: probod.LLMAgentConfig{
Provider: b.getEnvOrDefault("AGENT_PROBO_PROVIDER", ""),
ModelName: b.getEnvOrDefault("AGENT_PROBO_MODEL_NAME", ""),
Temperature: b.getEnvFloatPtr("AGENT_PROBO_TEMPERATURE"),
MaxTokens: b.getEnvIntPtr("AGENT_PROBO_MAX_TOKENS"),
},
EvidenceDescriber: probod.LLMAgentConfig{
Provider: b.getEnvOrDefault("AGENT_EVIDENCE_DESCRIBER_PROVIDER", ""),
ModelName: b.getEnvOrDefault("AGENT_EVIDENCE_DESCRIBER_MODEL_NAME", ""),
Temperature: b.getEnvFloatPtr("AGENT_EVIDENCE_DESCRIBER_TEMPERATURE"),
MaxTokens: b.getEnvIntPtr("AGENT_EVIDENCE_DESCRIBER_MAX_TOKENS"),
Defaults: probod.LLMConfig{
Provider: b.getEnvOrDefault("LLM_DEFAULT_PROVIDER", "openai"),
ModelName: b.getEnvOrDefault("LLM_DEFAULT_MODEL_NAME", "gpt-4o"),
Temperature: new(b.getEnvFloatOrDefault("LLM_DEFAULT_TEMPERATURE", 0.1)),
MaxTokens: new(b.getEnvIntOrDefault("LLM_DEFAULT_MAX_TOKENS", 4096)),
},
},
ProboAgent: probod.LLMConfig{
Provider: b.getEnvOrDefault("PROBO_AGENT_PROVIDER", ""),
ModelName: b.getEnvOrDefault("PROBO_AGENT_MODEL_NAME", ""),
Temperature: b.getEnvFloatPtr("PROBO_AGENT_TEMPERATURE"),
MaxTokens: b.getEnvIntPtr("PROBO_AGENT_MAX_TOKENS"),
},
EvidenceDescriber: probod.EvidenceDescriberConfig{
Interval: b.getEnvIntOrDefault("EVIDENCE_DESCRIBER_INTERVAL", 10),
StaleAfter: b.getEnvIntOrDefault("EVIDENCE_DESCRIBER_STALE_AFTER", 300),
MaxConcurrency: b.getEnvIntOrDefault("EVIDENCE_DESCRIBER_MAX_CONCURRENCY", 10),
Provider: b.getEnvOrDefault("EVIDENCE_DESCRIBER_PROVIDER", ""),
ModelName: b.getEnvOrDefault("EVIDENCE_DESCRIBER_MODEL_NAME", ""),
Temperature: b.getEnvFloatPtr("EVIDENCE_DESCRIBER_TEMPERATURE"),
MaxTokens: b.getEnvIntPtr("EVIDENCE_DESCRIBER_MAX_TOKENS"),
},
CustomDomains: probod.CustomDomainsConfig{
RenewalInterval: b.getEnvIntOrDefault("CUSTOM_DOMAINS_RENEWAL_INTERVAL", 3600),
ProvisionInterval: b.getEnvIntOrDefault("CUSTOM_DOMAINS_PROVISION_INTERVAL", 30),

98
pkg/bootstrap/builder_test.go
View File

@@ -161,20 +161,25 @@ func TestBuilder_Build_Defaults(t *testing.T) {
assert.Equal(t, 5, cfg.Probod.Notifications.Webhook.SenderInterval)
assert.Equal(t, 86400, cfg.Probod.Notifications.Webhook.CacheTTL)
// Agents config — default
assert.Equal(t, "openai", cfg.Probod.Agents.Default.Provider)
assert.Equal(t, "gpt-4o", cfg.Probod.Agents.Default.ModelName)
assert.Equal(t, new(0.1), cfg.Probod.Agents.Default.Temperature)
assert.Equal(t, new(4096), cfg.Probod.Agents.Default.MaxTokens)
// Agents config — per-agent overrides are empty (inherit from default)
assert.Empty(t, cfg.Probod.Agents.Probo.Provider)
assert.Empty(t, cfg.Probod.Agents.Probo.ModelName)
assert.Nil(t, cfg.Probod.Agents.Probo.Temperature)
assert.Nil(t, cfg.Probod.Agents.Probo.MaxTokens)
assert.Empty(t, cfg.Probod.Agents.EvidenceDescriber.Provider)
assert.Empty(t, cfg.Probod.Agents.EvidenceDescriber.ModelName)
assert.Nil(t, cfg.Probod.Agents.EvidenceDescriber.Temperature)
assert.Nil(t, cfg.Probod.Agents.EvidenceDescriber.MaxTokens)
// LLM config — defaults
assert.Equal(t, "openai", cfg.Probod.LLM.Defaults.Provider)
assert.Equal(t, "gpt-4o", cfg.Probod.LLM.Defaults.ModelName)
assert.Equal(t, new(0.1), cfg.Probod.LLM.Defaults.Temperature)
assert.Equal(t, new(4096), cfg.Probod.LLM.Defaults.MaxTokens)
// Probo agent — empty (inherits from defaults)
assert.Empty(t, cfg.Probod.ProboAgent.Provider)
assert.Empty(t, cfg.Probod.ProboAgent.ModelName)
assert.Nil(t, cfg.Probod.ProboAgent.Temperature)
assert.Nil(t, cfg.Probod.ProboAgent.MaxTokens)
// Evidence describer — LLM fields empty (inherits from defaults)
assert.Empty(t, cfg.Probod.EvidenceDescriber.Provider)
assert.Empty(t, cfg.Probod.EvidenceDescriber.ModelName)
assert.Nil(t, cfg.Probod.EvidenceDescriber.Temperature)
assert.Nil(t, cfg.Probod.EvidenceDescriber.MaxTokens)
// Evidence describer — worker defaults
assert.Equal(t, 10, cfg.Probod.EvidenceDescriber.Interval)
assert.Equal(t, 300, cfg.Probod.EvidenceDescriber.StaleAfter)
assert.Equal(t, 10, cfg.Probod.EvidenceDescriber.MaxConcurrency)
// Custom domains config
assert.Equal(t, 3600, cfg.Probod.CustomDomains.RenewalInterval)
@@ -241,19 +246,22 @@ func TestBuilder_Build_CustomValues(t *testing.T) {
env["WEBHOOK_SENDER_INTERVAL"] = "10"
env["WEBHOOK_CACHE_TTL"] = "3600"
env["CONNECTOR_SLACK_SIGNING_SECRET"] = "slack-signing-secret"
// Agents — providers
// LLM — providers
env["OPENAI_API_KEY"] = "sk-test-key"
env["ANTHROPIC_API_KEY"] = "sk-ant-test-key"
// Agents — default
env["AGENT_DEFAULT_PROVIDER"] = "openai"
env["AGENT_DEFAULT_MODEL_NAME"] = "gpt-4-turbo"
env["AGENT_DEFAULT_TEMPERATURE"] = "0.5"
env["AGENT_DEFAULT_MAX_TOKENS"] = "8192"
// Agents — evidence-describer override
env["AGENT_EVIDENCE_DESCRIBER_PROVIDER"] = "anthropic"
env["AGENT_EVIDENCE_DESCRIBER_MODEL_NAME"] = "claude-sonnet-4-20250514"
env["AGENT_EVIDENCE_DESCRIBER_TEMPERATURE"] = "0.2"
env["AGENT_EVIDENCE_DESCRIBER_MAX_TOKENS"] = "4096"
// LLM — defaults
env["LLM_DEFAULT_PROVIDER"] = "openai"
env["LLM_DEFAULT_MODEL_NAME"] = "gpt-4-turbo"
env["LLM_DEFAULT_TEMPERATURE"] = "0.5"
env["LLM_DEFAULT_MAX_TOKENS"] = "8192"
// Evidence describer
env["EVIDENCE_DESCRIBER_PROVIDER"] = "anthropic"
env["EVIDENCE_DESCRIBER_MODEL_NAME"] = "claude-sonnet-4-20250514"
env["EVIDENCE_DESCRIBER_TEMPERATURE"] = "0.2"
env["EVIDENCE_DESCRIBER_MAX_TOKENS"] = "4096"
env["EVIDENCE_DESCRIBER_INTERVAL"] = "15"
env["EVIDENCE_DESCRIBER_STALE_AFTER"] = "600"
env["EVIDENCE_DESCRIBER_MAX_CONCURRENCY"] = "20"
// Custom domains
env["CUSTOM_DOMAINS_RESOLVER_ADDR"] = "1.1.1.1:53"
env["ACME_ACCOUNT_KEY"] = "-----BEGIN EC PRIVATE KEY-----\ntest\n-----END EC PRIVATE KEY-----"
@@ -313,24 +321,28 @@ func TestBuilder_Build_CustomValues(t *testing.T) {
assert.Equal(t, "slack-signing-secret", cfg.Probod.Notifications.Slack.SigningSecret)
assert.Equal(t, 10, cfg.Probod.Notifications.Webhook.SenderInterval)
assert.Equal(t, 3600, cfg.Probod.Notifications.Webhook.CacheTTL)
// Agents — providers
assert.Equal(t, "openai", cfg.Probod.Agents.Providers["openai"].Type)
assert.Equal(t, "sk-test-key", cfg.Probod.Agents.Providers["openai"].APIKey)
assert.Equal(t, "anthropic", cfg.Probod.Agents.Providers["anthropic"].Type)
assert.Equal(t, "sk-ant-test-key", cfg.Probod.Agents.Providers["anthropic"].APIKey)
// Agents — default
assert.Equal(t, "openai", cfg.Probod.Agents.Default.Provider)
assert.Equal(t, "gpt-4-turbo", cfg.Probod.Agents.Default.ModelName)
assert.Equal(t, new(0.5), cfg.Probod.Agents.Default.Temperature)
assert.Equal(t, new(8192), cfg.Probod.Agents.Default.MaxTokens)
// Agents — probo inherits default (no overrides set)
assert.Empty(t, cfg.Probod.Agents.Probo.Provider)
assert.Empty(t, cfg.Probod.Agents.Probo.ModelName)
// Agents — evidence-describer overrides
assert.Equal(t, "anthropic", cfg.Probod.Agents.EvidenceDescriber.Provider)
assert.Equal(t, "claude-sonnet-4-20250514", cfg.Probod.Agents.EvidenceDescriber.ModelName)
assert.Equal(t, new(0.2), cfg.Probod.Agents.EvidenceDescriber.Temperature)
assert.Equal(t, new(4096), cfg.Probod.Agents.EvidenceDescriber.MaxTokens)
// LLM — providers
assert.Equal(t, "openai", cfg.Probod.LLM.Providers["openai"].Type)
assert.Equal(t, "sk-test-key", cfg.Probod.LLM.Providers["openai"].APIKey)
assert.Equal(t, "anthropic", cfg.Probod.LLM.Providers["anthropic"].Type)
assert.Equal(t, "sk-ant-test-key", cfg.Probod.LLM.Providers["anthropic"].APIKey)
// LLM — defaults
assert.Equal(t, "openai", cfg.Probod.LLM.Defaults.Provider)
assert.Equal(t, "gpt-4-turbo", cfg.Probod.LLM.Defaults.ModelName)
assert.Equal(t, new(0.5), cfg.Probod.LLM.Defaults.Temperature)
assert.Equal(t, new(8192), cfg.Probod.LLM.Defaults.MaxTokens)
// Probo agent — inherits defaults (no overrides set)
assert.Empty(t, cfg.Probod.ProboAgent.Provider)
assert.Empty(t, cfg.Probod.ProboAgent.ModelName)
// Evidence describer — LLM overrides
assert.Equal(t, "anthropic", cfg.Probod.EvidenceDescriber.Provider)
assert.Equal(t, "claude-sonnet-4-20250514", cfg.Probod.EvidenceDescriber.ModelName)
assert.Equal(t, new(0.2), cfg.Probod.EvidenceDescriber.Temperature)
assert.Equal(t, new(4096), cfg.Probod.EvidenceDescriber.MaxTokens)
// Evidence describer — worker config
assert.Equal(t, 15, cfg.Probod.EvidenceDescriber.Interval)
assert.Equal(t, 600, cfg.Probod.EvidenceDescriber.StaleAfter)
assert.Equal(t, 20, cfg.Probod.EvidenceDescriber.MaxConcurrency)
// Custom domains
assert.Equal(t, "1.1.1.1:53", cfg.Probod.CustomDomains.ResolverAddr)
assert.Equal(t, "-----BEGIN EC PRIVATE KEY-----\ntest\n-----END EC PRIVATE KEY-----", cfg.Probod.CustomDomains.ACME.AccountKey)

38
pkg/probod/evidence_describer_config.go Normal file
View File

@@ -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,
}
}

46
pkg/probod/openai_config.go → pkg/probod/llm_config.go
View File

@@ -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
}

56
pkg/probod/probod.go
View File

@@ -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(