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

Add AI-powered evidence description generation

Browse Source 
Introduce a background worker that automatically generates
compliance-focused descriptions for uploaded evidence files
using configurable LLM providers. Descriptions are surfaced
across all interfaces: GraphQL API, MCP API, CLI, and the
console UI.

Key changes:
- Multi-provider LLM config with per-agent settings (pointer
  types for Temperature/MaxTokens to preserve zero values)
- Evidence description worker with bounded concurrency
- EvidenceDescriptionStatus typed enum with PostgreSQL enum type
- New `prb evidence` CLI commands (list, view, delete)
- Evidence description displayed in console table and preview
- Migration only marks evidences without files as completed

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
This commit is contained in:
Bryan Frimin
2026-03-26 14:30:04 +01:00
parent 8bbed534a9
commit 619ec7b882
24 changed files with 1390 additions and 125 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
pkg/probod/llm.go
View File

@@ -22,13 +22,14 @@ import (
"go.gearno.de/kit/log"
"go.opentelemetry.io/otel/trace"
"go.probo.inc/probo/pkg/llm"
llmanthropic "go.probo.inc/probo/pkg/llm/anthropic"
llmopenai "go.probo.inc/probo/pkg/llm/openai"
)
func buildLLMClient(cfg LLMConfig, l *log.Logger, tp trace.TracerProvider, r prometheus.Registerer) (*llm.Client, error) {
provider := cfg.Provider
if provider == "" {
provider = "openai"
func buildLLMClient(cfg LLMProviderConfig, l *log.Logger, tp trace.TracerProvider, r prometheus.Registerer) (*llm.Client, error) {
providerType := cfg.Type
if providerType == "" {
providerType = "openai"
}
httpClient := httpclient.DefaultPooledClient(
@@ -37,7 +38,7 @@ func buildLLMClient(cfg LLMConfig, l *log.Logger, tp trace.TracerProvider, r pro
httpclient.WithRegisterer(r),
)
switch provider {
switch providerType {
case "openai":
p := llmopenai.NewProvider(
cfg.APIKey,
@@ -50,10 +51,19 @@ func buildLLMClient(cfg LLMConfig, l *log.Logger, tp trace.TracerProvider, r pro
llm.WithTracerProvider(tp),
), nil
case "anthropic":
return nil, fmt.Errorf("anthropic provider not yet wired; add import and construct here")
p := llmanthropic.NewProvider(
cfg.APIKey,
llmanthropic.WithHTTPClient(httpClient),
)
return llm.NewClient(
p,
"anthropic",
llm.WithLogger(l),
llm.WithTracerProvider(tp),
), nil
case "bedrock":
return nil, fmt.Errorf("bedrock provider not yet wired; requires aws.Config")
default:
return nil, fmt.Errorf("unsupported LLM provider: %q", provider)
return nil, fmt.Errorf("unsupported LLM provider type: %q", providerType)
}
}

52
pkg/probod/openai_config.go
View File

@@ -14,12 +14,48 @@
package probod
type LLMConfig struct {
Provider string `json:"provider"` // "openai", "anthropic", "bedrock"
APIKey string `json:"api-key"` // for OpenAI and Anthropic
ModelName string `json:"model-name"`
Temperature float64 `json:"temperature"`
MaxTokens int `json:"max-tokens"`
}
type (
// LLMProviderConfig holds authentication and connection settings for an
// LLM provider (e.g. OpenAI, Anthropic).
LLMProviderConfig struct {
Type string `json:"type"` // "openai", "anthropic", "bedrock"
APIKey string `json:"api-key"` // for OpenAI and Anthropic
}
type OpenAIConfig = LLMConfig
// 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
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"`
}
)
// 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
}
if agent.ModelName == "" {
agent.ModelName = c.Default.ModelName
}
if agent.Temperature == nil {
agent.Temperature = c.Default.Temperature
}
if agent.MaxTokens == nil {
agent.MaxTokens = c.Default.MaxTokens
}
return agent
}

54
pkg/probod/probod.go
View File

@@ -51,6 +51,7 @@ import (
"go.probo.inc/probo/pkg/crypto/keys"
"go.probo.inc/probo/pkg/crypto/passwdhash"
"go.probo.inc/probo/pkg/esign"
"go.probo.inc/probo/pkg/evidencedescriber"
"go.probo.inc/probo/pkg/file"
"go.probo.inc/probo/pkg/filemanager"
"go.probo.inc/probo/pkg/html2pdf"
@@ -115,7 +116,7 @@ type (
AWS AWSConfig `json:"aws"`
Notifications NotificationsConfig `json:"notifications"`
Connectors []ConnectorConfig `json:"connectors"`
OpenAI OpenAIConfig `json:"openai"`
Agents AgentsConfig `json:"agents"`
ChromeDPAddr string `json:"chrome-dp-addr"`
CustomDomains CustomDomainsConfig `json:"custom-domains"`
SCIMBridge SCIMBridgeConfig `json:"scim-bridge"`
@@ -316,9 +317,24 @@ func (impl *Implm) Run(
}
}
llmClient, err := buildLLMClient(impl.cfg.OpenAI, l.Named("llm"), tp, r)
proboAgentCfg := impl.cfg.Agents.ResolveAgent(impl.cfg.Agents.Probo)
proboProviderCfg, ok := impl.cfg.Agents.Providers[proboAgentCfg.Provider]
if !ok {
return fmt.Errorf("unknown LLM provider %q for probo agent", proboAgentCfg.Provider)
}
proboLLMClient, err := buildLLMClient(proboProviderCfg, l.Named("llm.probo"), tp, r)
if err != nil {
return fmt.Errorf("cannot create LLM client: %w", err)
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]
if !ok {
return fmt.Errorf("unknown LLM provider %q for evidence-describer agent", evidenceDescriberAgentCfg.Provider)
}
evidenceDescriberLLMClient, err := buildLLMClient(evidenceDescriberProviderCfg, l.Named("llm.evidence-describer"), tp, r)
if err != nil {
return fmt.Errorf("cannot create evidence describer LLM client: %w", err)
}
fileManagerService := filemanager.NewService(s3Client)
@@ -454,10 +470,10 @@ func (impl *Implm) Run(
impl.cfg.AWS.Bucket,
baseURL.String(),
impl.cfg.Auth.Cookie.Secret,
llmClient,
impl.cfg.OpenAI.ModelName,
impl.cfg.OpenAI.Temperature,
impl.cfg.OpenAI.MaxTokens,
proboLLMClient,
proboAgentCfg.ModelName,
*proboAgentCfg.Temperature,
*proboAgentCfg.MaxTokens,
html2pdfConverter,
acmeService,
fileManagerService,
@@ -617,6 +633,29 @@ func (impl *Implm) Run(
},
)
evidenceDescriber := evidencedescriber.New(
evidenceDescriberLLMClient,
evidencedescriber.Config{
Model: evidenceDescriberAgentCfg.ModelName,
Temp: *evidenceDescriberAgentCfg.Temperature,
MaxTokens: *evidenceDescriberAgentCfg.MaxTokens,
},
)
evidenceDescriptionWorker := probo.NewEvidenceDescriptionWorker(
pgClient,
fileManagerService,
evidenceDescriber,
l.Named("evidence-description-worker"),
)
evidenceDescriptionWorkerCtx, stopEvidenceDescriptionWorker := context.WithCancel(context.Background())
wg.Go(
func() {
if err := evidenceDescriptionWorker.Run(evidenceDescriptionWorkerCtx); err != nil {
cancel(fmt.Errorf("evidence description worker crashed: %w", err))
}
},
)
trustCenterServerCtx, stopTrustCenterServer := context.WithCancel(context.Background())
defer stopTrustCenterServer()
wg.Go(
@@ -644,6 +683,7 @@ func (impl *Implm) Run(
stopWebhookSender()
stopESignService()
stopMailingListWorker()
stopEvidenceDescriptionWorker()
stopExportJobExporter()
stopIAMService()
stopMailer()