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

1
apps/console/src/hooks/graph/EvidenceGraph.ts
View File

@@ -7,6 +7,7 @@ export const evidenceFileQuery = graphql`
node(id: $evidenceId) {
... on Evidence {
id
description
file {
mimeType
fileName

43
apps/console/src/pages/organizations/measures/dialog/EvidencePreviewDialog.tsx
View File

@@ -125,39 +125,48 @@ function EvidencePreviewContent({
);
}
let preview;
if (evidence.file.mimeType?.startsWith("image/")) {
return (
preview = (
<img
src={evidence.file.downloadUrl}
alt={evidence.file.fileName}
className="max-h-[70vh] object-contain"
/>
);
}
if (evidence.file.mimeType?.includes("pdf")) {
return (
} else if (evidence.file.mimeType?.includes("pdf")) {
preview = (
<iframe
src={evidence.file.downloadUrl}
className="w-full h-[70vh]"
title={evidence.file.fileName}
/>
);
} else {
preview = (
<div className="flex flex-col items-center gap-2 justify-center">
<IconWarning size={20} />
<p className="text-txt-secondary text-center">
{__("Preview not available for this file type")
+ " "
+ evidence.file.mimeType}
</p>
<Button asChild variant="secondary" icon={IconArrowInbox}>
<a href={evidence.file.downloadUrl} target="_blank" rel="noreferrer">
{__("Download File")}
</a>
</Button>
</div>
);
}
return (
<div className="flex flex-col items-center gap-2 justify-center">
<IconWarning size={20} />
<p className="text-txt-secondary text-center">
{__("Preview not available for this file type")
+ " "
+ evidence.file.mimeType}
</p>
<Button asChild variant="secondary" icon={IconArrowInbox}>
<a href={evidence.file.downloadUrl} target="_blank" rel="noreferrer">
{__("Download File")}
</a>
</Button>
<div className="space-y-4">
{preview}
{evidence.description && (
<p className="text-txt-secondary text-sm">{evidence.description}</p>
)}
</div>
);
}

17
apps/console/src/pages/organizations/measures/tabs/MeasureEvidencesTab.tsx
View File

@@ -1,4 +1,4 @@
import { fileSize, fileType, formatDate, sprintf } from "@probo/helpers";
import { fileSize, formatDate, sprintf } from "@probo/helpers";
import { usePageTitle } from "@probo/hooks";
import { useTranslate } from "@probo/i18n";
import {
@@ -79,7 +79,7 @@ export const evidenceFragment = graphql`
mimeType
size
}
type
description
createdAt
canDelete: permission(action: "core:evidence:delete")
}
@@ -123,8 +123,8 @@ export default function MeasureEvidencesTab() {
<SortableTable {...pagination}>
<Thead>
<Tr>
<Th>{__("Evidence name")}</Th>
<Th>{__("Type")}</Th>
<Th>{__("Description")}</Th>
<Th>{__("File type")}</Th>
<Th>{__("File size")}</Th>
<Th>{__("Created at")}</Th>
<Th width={50}></Th>
@@ -248,13 +248,12 @@ function EvidenceRow(props: {
/>
)}
<Tr to={evidenceUrl}>
<Td>{evidence.file?.fileName}</Td>
<Td>
{fileType(__, {
type: evidence.type,
mimeType: evidence.file?.mimeType || "",
})}
<span className="text-txt-secondary text-sm line-clamp-2">
{evidence.description || "—"}
</span>
</Td>
<Td>{evidence.file?.mimeType || "—"}</Td>
<Td>{fileSize(__, evidence.file?.size || 0)}</Td>
<Td>{formatDate(evidence.createdAt)}</Td>
<Td noLink>

53
pkg/bootstrap/builder.go
View File

@@ -157,11 +157,35 @@ func (b *Builder) Build() (*probod.FullConfig, error) {
CacheTTL: b.getEnvIntOrDefault("WEBHOOK_CACHE_TTL", 86400),
},
},
OpenAI: probod.OpenAIConfig{
APIKey: b.getEnv("OPENAI_API_KEY"),
Temperature: b.getEnvFloatOrDefault("OPENAI_TEMPERATURE", 0.1),
ModelName: b.getEnvOrDefault("OPENAI_MODEL_NAME", "gpt-4o"),
MaxTokens: b.getEnvIntOrDefault("OPENAI_MAX_TOKENS", 4096),
Agents: probod.AgentsConfig{
Providers: map[string]probod.LLMProviderConfig{
"openai": {
Type: "openai",
APIKey: b.getEnv("OPENAI_API_KEY"),
},
"anthropic": {
Type: "anthropic",
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"),
},
},
CustomDomains: probod.CustomDomainsConfig{
RenewalInterval: b.getEnvIntOrDefault("CUSTOM_DOMAINS_RENEWAL_INTERVAL", 3600),
@@ -302,6 +326,25 @@ func (b *Builder) getEnvFloatOrDefault(key string, defaultValue float64) float64
return defaultValue
}
func (b *Builder) getEnvFloatPtr(key string) *float64 {
if value := b.getEnv(key); value != "" {
if floatValue, err := strconv.ParseFloat(value, 64); err == nil {
return &floatValue
}
}
return nil
}
func (b *Builder) getEnvIntPtr(key string) *int {
if value := b.getEnv(key); value != "" {
if intValue, err := strconv.ParseInt(value, 10, 32); err == nil {
v := int(intValue)
return &v
}
}
return nil
}
func (b *Builder) getEnvBoolOrDefault(key string, defaultValue bool) bool {
if value := b.getEnv(key); value != "" {
if boolValue, err := strconv.ParseBool(value); err == nil {

57
pkg/bootstrap/builder_test.go
View File

@@ -161,10 +161,20 @@ 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)
// OpenAI config
assert.Equal(t, 0.1, cfg.Probod.OpenAI.Temperature)
assert.Equal(t, "gpt-4o", cfg.Probod.OpenAI.ModelName)
assert.Equal(t, 4096, cfg.Probod.OpenAI.MaxTokens)
// 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)
// Custom domains config
assert.Equal(t, 3600, cfg.Probod.CustomDomains.RenewalInterval)
@@ -231,11 +241,19 @@ 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"
// OpenAI
// Agents — providers
env["OPENAI_API_KEY"] = "sk-test-key"
env["OPENAI_TEMPERATURE"] = "0.5"
env["OPENAI_MODEL_NAME"] = "gpt-4-turbo"
env["OPENAI_MAX_TOKENS"] = "8192"
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"
// 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-----"
@@ -295,11 +313,24 @@ 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)
// OpenAI
assert.Equal(t, "sk-test-key", cfg.Probod.OpenAI.APIKey)
assert.Equal(t, 0.5, cfg.Probod.OpenAI.Temperature)
assert.Equal(t, "gpt-4-turbo", cfg.Probod.OpenAI.ModelName)
assert.Equal(t, 8192, cfg.Probod.OpenAI.MaxTokens)
// 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)
// 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)

106
pkg/cmd/evidence/delete/delete.go Normal file
View File

@@ -0,0 +1,106 @@
// 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 delete
import (
"encoding/json"
"fmt"
"github.com/charmbracelet/huh"
"github.com/spf13/cobra"
"go.probo.inc/probo/pkg/cli/api"
"go.probo.inc/probo/pkg/cmd/cmdutil"
)
const deleteMutation = `
mutation($input: DeleteEvidenceInput!) {
deleteEvidence(input: $input) {
deletedEvidenceId
}
}
`
type deleteResponse struct {
DeleteEvidence struct {
DeletedEvidenceID string `json:"deletedEvidenceId"`
} `json:"deleteEvidence"`
}
func NewCmdDelete(f *cmdutil.Factory) *cobra.Command {
var flagYes bool
cmd := &cobra.Command{
Use: "delete <id>",
Short: "Delete an evidence",
Args: cobra.ExactArgs(1),
RunE: func(cmd *cobra.Command, args []string) error {
if !flagYes {
if !f.IOStreams.IsInteractive() {
return fmt.Errorf("cannot delete evidence: confirmation required, use --yes to confirm")
}
var confirmed bool
err := huh.NewConfirm().Title(fmt.Sprintf("Delete evidence %s?", args[0])).Value(&confirmed).Run()
if err != nil {
return err
}
if !confirmed {
return nil
}
}
cfg, err := f.Config()
if err != nil {
return err
}
host, hc, err := cfg.DefaultHost()
if err != nil {
return err
}
client := api.NewClient(
host,
hc.Token,
"/api/console/v1/graphql",
cfg.HTTPTimeoutDuration(),
)
data, err := client.Do(
deleteMutation,
map[string]any{
"input": map[string]any{
"evidenceId": args[0],
},
},
)
if err != nil {
return err
}
var resp deleteResponse
if err := json.Unmarshal(data, &resp); err != nil {
return fmt.Errorf("cannot parse response: %w", err)
}
_, _ = fmt.Fprintf(f.IOStreams.Out, "Deleted evidence %s\n", resp.DeleteEvidence.DeletedEvidenceID)
return nil
},
}
cmd.Flags().BoolVarP(&flagYes, "yes", "y", false, "Skip confirmation prompt")
return cmd
}

36
pkg/cmd/evidence/evidence.go Normal file
View File

@@ -0,0 +1,36 @@
// 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 evidence
import (
"github.com/spf13/cobra"
"go.probo.inc/probo/pkg/cmd/cmdutil"
"go.probo.inc/probo/pkg/cmd/evidence/delete"
"go.probo.inc/probo/pkg/cmd/evidence/list"
"go.probo.inc/probo/pkg/cmd/evidence/view"
)
func NewCmdEvidence(f *cmdutil.Factory) *cobra.Command {
cmd := &cobra.Command{
Use: "evidence <command>",
Short: "Manage evidences",
}
cmd.AddCommand(list.NewCmdList(f))
cmd.AddCommand(view.NewCmdView(f))
cmd.AddCommand(delete.NewCmdDelete(f))
return cmd
}

199
pkg/cmd/evidence/list/list.go Normal file
View File

@@ -0,0 +1,199 @@
// 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 list
import (
"encoding/json"
"fmt"
"github.com/spf13/cobra"
"go.probo.inc/probo/pkg/cli/api"
"go.probo.inc/probo/pkg/cmd/cmdutil"
)
const listQuery = `
query($id: ID!, $first: Int, $after: CursorKey, $orderBy: EvidenceOrder) {
node(id: $id) {
__typename
... on Measure {
evidences(first: $first, after: $after, orderBy: $orderBy) {
totalCount
edges {
node {
id
state
type
url
description
createdAt
}
}
pageInfo {
hasNextPage
endCursor
}
}
}
}
}
`
type evidence struct {
ID string `json:"id"`
State string `json:"state"`
Type string `json:"type"`
URL string `json:"url"`
Description *string `json:"description"`
CreatedAt string `json:"createdAt"`
}
func NewCmdList(f *cmdutil.Factory) *cobra.Command {
var (
flagMeasure string
flagLimit int
flagOrderBy string
flagOrderDir string
flagOutput *string
)
cmd := &cobra.Command{
Use: "list",
Short: "List evidences for a measure",
Aliases: []string{"ls"},
Example: ` # List evidences for a measure
prb evidence list --measure <measure-id>
# Output as JSON
prb evidence ls --measure <measure-id> --output json`,
Args: cobra.NoArgs,
RunE: func(cmd *cobra.Command, args []string) error {
if err := cmdutil.ValidateOutputFlag(flagOutput); err != nil {
return err
}
cfg, err := f.Config()
if err != nil {
return err
}
host, hc, err := cfg.DefaultHost()
if err != nil {
return err
}
client := api.NewClient(
host,
hc.Token,
"/api/console/v1/graphql",
cfg.HTTPTimeoutDuration(),
)
variables := map[string]any{
"id": flagMeasure,
}
if flagOrderBy != "" {
if err := cmdutil.ValidateEnum("order-by", flagOrderBy, []string{"CREATED_AT"}); err != nil {
return err
}
variables["orderBy"] = map[string]any{
"field": flagOrderBy,
"direction": flagOrderDir,
}
}
evidences, totalCount, err := api.Paginate(
client,
listQuery,
variables,
flagLimit,
func(data json.RawMessage) (*api.Connection[evidence], error) {
var resp struct {
Node *struct {
Typename string `json:"__typename"`
Evidences api.Connection[evidence] `json:"evidences"`
} `json:"node"`
}
if err := json.Unmarshal(data, &resp); err != nil {
return nil, err
}
if resp.Node == nil {
return nil, fmt.Errorf("measure %s not found", flagMeasure)
}
if resp.Node.Typename != "Measure" {
return nil, fmt.Errorf("expected Measure node, got %s", resp.Node.Typename)
}
return &resp.Node.Evidences, nil
},
)
if err != nil {
return err
}
if *flagOutput == cmdutil.OutputJSON {
return cmdutil.PrintJSON(f.IOStreams.Out, evidences)
}
if len(evidences) == 0 {
_, _ = fmt.Fprintln(f.IOStreams.Out, "No evidences found.")
return nil
}
rows := make([][]string, 0, len(evidences))
for _, e := range evidences {
desc := "-"
if e.Description != nil && *e.Description != "" {
d := *e.Description
if len(d) > 60 {
d = d[:57] + "..."
}
desc = d
}
rows = append(rows, []string{
e.ID,
e.Type,
e.State,
desc,
cmdutil.FormatTime(e.CreatedAt),
})
}
t := cmdutil.NewTable("ID", "TYPE", "STATE", "DESCRIPTION", "CREATED").Rows(rows...)
_, _ = fmt.Fprintln(f.IOStreams.Out, t)
if totalCount > len(evidences) {
_, _ = fmt.Fprintf(
f.IOStreams.ErrOut,
"\nShowing %d of %d evidences\n",
len(evidences),
totalCount,
)
}
return nil
},
}
cmd.Flags().StringVar(&flagMeasure, "measure", "", "Measure ID (required)")
cmd.Flags().IntVarP(&flagLimit, "limit", "L", 30, "Maximum number of evidences to list")
cmd.Flags().StringVar(&flagOrderBy, "order-by", "", "Order by field (CREATED_AT)")
cmd.Flags().StringVar(&flagOrderDir, "order-direction", "DESC", "Sort direction (ASC, DESC)")
flagOutput = cmdutil.AddOutputFlag(cmd)
_ = cmd.MarkFlagRequired("measure")
return cmd
}

171
pkg/cmd/evidence/view/view.go Normal file
View File

@@ -0,0 +1,171 @@
// 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 view
import (
"encoding/json"
"fmt"
"github.com/charmbracelet/lipgloss"
"github.com/spf13/cobra"
"go.probo.inc/probo/pkg/cli/api"
"go.probo.inc/probo/pkg/cmd/cmdutil"
)
const viewQuery = `
query($id: ID!) {
node(id: $id) {
__typename
... on Evidence {
id
size
state
type
url
description
file {
id
filename
contentType
}
measure {
id
}
task {
id
}
createdAt
updatedAt
}
}
}
`
type viewResponse struct {
Node *struct {
Typename string `json:"__typename"`
ID string `json:"id"`
Size int `json:"size"`
State string `json:"state"`
Type string `json:"type"`
URL string `json:"url"`
Description *string `json:"description"`
File *struct {
ID string `json:"id"`
Filename string `json:"filename"`
ContentType string `json:"contentType"`
} `json:"file"`
Measure struct {
ID string `json:"id"`
} `json:"measure"`
Task *struct {
ID string `json:"id"`
} `json:"task"`
CreatedAt string `json:"createdAt"`
UpdatedAt string `json:"updatedAt"`
} `json:"node"`
}
func NewCmdView(f *cmdutil.Factory) *cobra.Command {
var flagOutput *string
cmd := &cobra.Command{
Use: "view <id>",
Short: "View an evidence",
Args: cobra.ExactArgs(1),
RunE: func(cmd *cobra.Command, args []string) error {
if err := cmdutil.ValidateOutputFlag(flagOutput); err != nil {
return err
}
cfg, err := f.Config()
if err != nil {
return err
}
host, hc, err := cfg.DefaultHost()
if err != nil {
return err
}
client := api.NewClient(
host,
hc.Token,
"/api/console/v1/graphql",
cfg.HTTPTimeoutDuration(),
)
data, err := client.Do(
viewQuery,
map[string]any{"id": args[0]},
)
if err != nil {
return err
}
var resp viewResponse
if err := json.Unmarshal(data, &resp); err != nil {
return fmt.Errorf("cannot parse response: %w", err)
}
if resp.Node == nil {
return fmt.Errorf("evidence %s not found", args[0])
}
if resp.Node.Typename != "Evidence" {
return fmt.Errorf("expected Evidence node, got %s", resp.Node.Typename)
}
if *flagOutput == cmdutil.OutputJSON {
return cmdutil.PrintJSON(f.IOStreams.Out, resp.Node)
}
n := resp.Node
out := f.IOStreams.Out
bold := lipgloss.NewStyle().Bold(true)
label := lipgloss.NewStyle().Foreground(lipgloss.Color("242")).Width(22)
_, _ = fmt.Fprintf(out, "%s\n\n", bold.Render(n.ID))
_, _ = fmt.Fprintf(out, "%s%s\n", label.Render("Type:"), n.Type)
_, _ = fmt.Fprintf(out, "%s%s\n", label.Render("State:"), n.State)
_, _ = fmt.Fprintf(out, "%s%s\n", label.Render("Measure:"), n.Measure.ID)
if n.File != nil {
_, _ = fmt.Fprintf(out, "%s%s (%s)\n", label.Render("File:"), n.File.Filename, n.File.ContentType)
}
if n.URL != "" {
_, _ = fmt.Fprintf(out, "%s%s\n", label.Render("URL:"), n.URL)
}
if n.Task != nil {
_, _ = fmt.Fprintf(out, "%s%s\n", label.Render("Task:"), n.Task.ID)
}
if n.Description != nil && *n.Description != "" {
_, _ = fmt.Fprintf(out, "%s%s\n", label.Render("Description:"), *n.Description)
}
_, _ = fmt.Fprintln(out)
_, _ = fmt.Fprintf(out, "%s%s\n", label.Render("Created:"), cmdutil.FormatTime(n.CreatedAt))
_, _ = fmt.Fprintf(out, "%s%s\n", label.Render("Updated:"), cmdutil.FormatTime(n.UpdatedAt))
return nil
},
}
flagOutput = cmdutil.AddOutputFlag(cmd)
return cmd
}

2
pkg/cmd/root/root.go
View File

@@ -25,6 +25,7 @@ import (
cmdconfig "go.probo.inc/probo/pkg/cmd/config"
cmdcontext "go.probo.inc/probo/pkg/cmd/context"
"go.probo.inc/probo/pkg/cmd/control"
"go.probo.inc/probo/pkg/cmd/evidence"
"go.probo.inc/probo/pkg/cmd/finding"
"go.probo.inc/probo/pkg/cmd/framework"
"go.probo.inc/probo/pkg/cmd/org"
@@ -73,6 +74,7 @@ func NewCmdRoot(f *cmdutil.Factory) *cobra.Command {
cmd.AddCommand(cmdconfig.NewCmdConfig(f))
cmd.AddCommand(cmdcontext.NewCmdContext(f))
cmd.AddCommand(control.NewCmdControl(f))
cmd.AddCommand(evidence.NewCmdEvidence(f))
cmd.AddCommand(finding.NewCmdFinding(f))
cmd.AddCommand(framework.NewCmdFramework(f))
cmd.AddCommand(org.NewCmdOrg(f))

180
pkg/coredata/evidence.go
View File

@@ -30,18 +30,20 @@ import (
type (
Evidence struct {
ID gid.GID `db:"id"`
OrganizationID gid.GID `db:"organization_id"`
MeasureID gid.GID `db:"measure_id"`
TaskID *gid.GID `db:"task_id"`
State EvidenceState `db:"state"`
ReferenceID string `db:"reference_id"`
Type EvidenceType `db:"type"`
URL string `db:"url"`
EvidenceFileId *gid.GID `db:"evidence_file_id"`
Description *string `db:"description"`
CreatedAt time.Time `db:"created_at"`
UpdatedAt time.Time `db:"updated_at"`
ID gid.GID `db:"id"`
OrganizationID gid.GID `db:"organization_id"`
MeasureID gid.GID `db:"measure_id"`
TaskID *gid.GID `db:"task_id"`
State EvidenceState `db:"state"`
ReferenceID string `db:"reference_id"`
Type EvidenceType `db:"type"`
URL string `db:"url"`
EvidenceFileId *gid.GID `db:"evidence_file_id"`
Description *string `db:"description"`
DescriptionStatus EvidenceDescriptionStatus `db:"description_status"`
DescriptionProcessingStartedAt *time.Time `db:"description_processing_started_at"`
CreatedAt time.Time `db:"created_at"`
UpdatedAt time.Time `db:"updated_at"`
}
Evidences []*Evidence
@@ -89,6 +91,8 @@ INSERT INTO
url,
evidence_file_id,
description,
description_status,
description_processing_started_at,
created_at,
updated_at
)
@@ -103,6 +107,8 @@ VALUES (
@url,
@evidence_file_id,
@description,
@description_status,
@description_processing_started_at,
@created_at,
@updated_at
)
@@ -114,18 +120,20 @@ WHERE evidences.state = 'REQUESTED';
`
args := pgx.StrictNamedArgs{
"tenant_id": scope.GetTenantID(),
"evidence_id": e.ID,
"measure_id": e.MeasureID,
"task_id": e.TaskID,
"reference_id": e.ReferenceID,
"evidence_file_id": e.EvidenceFileId,
"created_at": e.CreatedAt,
"updated_at": e.UpdatedAt,
"state": e.State,
"type": e.Type,
"url": e.URL,
"description": e.Description,
"tenant_id": scope.GetTenantID(),
"evidence_id": e.ID,
"measure_id": e.MeasureID,
"task_id": e.TaskID,
"reference_id": e.ReferenceID,
"evidence_file_id": e.EvidenceFileId,
"created_at": e.CreatedAt,
"updated_at": e.UpdatedAt,
"state": e.State,
"type": e.Type,
"url": e.URL,
"description": e.Description,
"description_status": e.DescriptionStatus,
"description_processing_started_at": e.DescriptionProcessingStartedAt,
}
_, err := conn.Exec(ctx, q, args)
return err
@@ -150,6 +158,8 @@ INSERT INTO
url,
evidence_file_id,
description,
description_status,
description_processing_started_at,
created_at,
updated_at
)
@@ -165,25 +175,29 @@ VALUES (
@url,
@evidence_file_id,
@description,
@description_status,
@description_processing_started_at,
@created_at,
@updated_at
)
`
args := pgx.StrictNamedArgs{
"tenant_id": scope.GetTenantID(),
"evidence_id": e.ID,
"organization_id": e.OrganizationID,
"measure_id": e.MeasureID,
"task_id": e.TaskID,
"reference_id": e.ReferenceID,
"evidence_file_id": e.EvidenceFileId,
"created_at": e.CreatedAt,
"updated_at": e.UpdatedAt,
"state": e.State,
"type": e.Type,
"url": e.URL,
"description": e.Description,
"tenant_id": scope.GetTenantID(),
"evidence_id": e.ID,
"organization_id": e.OrganizationID,
"measure_id": e.MeasureID,
"task_id": e.TaskID,
"reference_id": e.ReferenceID,
"evidence_file_id": e.EvidenceFileId,
"created_at": e.CreatedAt,
"updated_at": e.UpdatedAt,
"state": e.State,
"type": e.Type,
"url": e.URL,
"description": e.Description,
"description_status": e.DescriptionStatus,
"description_processing_started_at": e.DescriptionProcessingStartedAt,
}
_, err := conn.Exec(ctx, q, args)
@@ -218,6 +232,8 @@ SELECT
url,
evidence_file_id,
description,
description_status,
description_processing_started_at,
created_at,
updated_at
FROM
@@ -299,6 +315,8 @@ SELECT
url,
evidence_file_id,
description,
description_status,
description_processing_started_at,
created_at,
updated_at
FROM
@@ -381,6 +399,8 @@ SELECT
url,
evidence_file_id,
description,
description_status,
description_processing_started_at,
created_at,
updated_at
FROM
@@ -426,6 +446,8 @@ SET
evidence_file_id = @evidence_file_id,
url = @url,
description = @description,
description_status = @description_status,
description_processing_started_at = @description_processing_started_at,
updated_at = @updated_at
WHERE
%s
@@ -435,13 +457,15 @@ WHERE
q = fmt.Sprintf(q, scope.SQLFragment())
args := pgx.StrictNamedArgs{
"evidence_id": e.ID,
"type": e.Type,
"state": e.State,
"evidence_file_id": e.EvidenceFileId,
"url": e.URL,
"description": e.Description,
"updated_at": e.UpdatedAt,
"evidence_id": e.ID,
"type": e.Type,
"state": e.State,
"evidence_file_id": e.EvidenceFileId,
"url": e.URL,
"description": e.Description,
"description_status": e.DescriptionStatus,
"description_processing_started_at": e.DescriptionProcessingStartedAt,
"updated_at": e.UpdatedAt,
}
maps.Copy(args, scope.SQLArguments())
@@ -474,3 +498,71 @@ WHERE
return nil
}
func (e *Evidence) LoadNextPendingDescriptionForUpdateSkipLocked(
ctx context.Context,
conn pg.Conn,
) error {
q := `
SELECT
id,
organization_id,
task_id,
measure_id,
reference_id,
state,
type,
url,
evidence_file_id,
description,
description_status,
description_processing_started_at,
created_at,
updated_at
FROM
evidences
WHERE
description_status = 'PENDING'
AND evidence_file_id IS NOT NULL
ORDER BY
created_at ASC
LIMIT 1
FOR UPDATE SKIP LOCKED;
`
rows, err := conn.Query(ctx, q)
if err != nil {
return fmt.Errorf("cannot query evidence: %w", err)
}
evidence, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[Evidence])
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return ErrResourceNotFound
}
return fmt.Errorf("cannot collect evidence: %w", err)
}
*e = evidence
return nil
}
func ResetStaleDescriptionProcessing(
ctx context.Context,
conn pg.Conn,
staleAfter time.Duration,
) error {
q := `
UPDATE evidences
SET
description_status = 'PENDING',
description_processing_started_at = NULL
WHERE
description_status = 'PROCESSING'
AND description_processing_started_at < $1;
`
_, err := conn.Exec(ctx, q, time.Now().Add(-staleAfter))
return err
}

68
pkg/coredata/evidence_description_status.go Normal file
View File

@@ -0,0 +1,68 @@
// 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 coredata
import (
"database/sql/driver"
"fmt"
)
type (
EvidenceDescriptionStatus string
)
const (
EvidenceDescriptionStatusPending EvidenceDescriptionStatus = "PENDING"
EvidenceDescriptionStatusProcessing EvidenceDescriptionStatus = "PROCESSING"
EvidenceDescriptionStatusCompleted EvidenceDescriptionStatus = "COMPLETED"
)
func (s EvidenceDescriptionStatus) MarshalText() ([]byte, error) {
return []byte(s.String()), nil
}
func (s *EvidenceDescriptionStatus) UnmarshalText(data []byte) error {
val := string(data)
switch val {
case EvidenceDescriptionStatusPending.String():
*s = EvidenceDescriptionStatusPending
case EvidenceDescriptionStatusProcessing.String():
*s = EvidenceDescriptionStatusProcessing
case EvidenceDescriptionStatusCompleted.String():
*s = EvidenceDescriptionStatusCompleted
default:
return fmt.Errorf("invalid EvidenceDescriptionStatus value: %q", val)
}
return nil
}
func (s EvidenceDescriptionStatus) String() string {
return string(s)
}
func (s *EvidenceDescriptionStatus) Scan(value any) error {
val, ok := value.(string)
if !ok {
return fmt.Errorf("invalid scan source for EvidenceDescriptionStatus, expected string got %T", value)
}
return s.UnmarshalText([]byte(val))
}
func (s EvidenceDescriptionStatus) Value() (driver.Value, error) {
return s.String(), nil
}

10
pkg/coredata/migrations/20260326T131816Z.sql Normal file
View File

@@ -0,0 +1,10 @@
CREATE TYPE evidence_description_status AS ENUM ('PENDING', 'PROCESSING', 'COMPLETED');
ALTER TABLE evidences
ADD COLUMN description_status evidence_description_status NOT NULL DEFAULT 'PENDING',
ADD COLUMN description_processing_started_at TIMESTAMPTZ;
UPDATE evidences SET description_status = 'COMPLETED' WHERE evidence_file_id IS NULL;
ALTER TABLE evidences
ALTER COLUMN description_status DROP DEFAULT;

81
pkg/evidencedescriber/evidencedescriber.go Normal file
View File

@@ -0,0 +1,81 @@
// 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 evidencedescriber
import (
"context"
_ "embed"
"fmt"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/llm"
)
//go:embed prompt.txt
var systemPrompt string
type (
Config struct {
Model string
Temp float64
MaxTokens int
}
Describer struct {
client *llm.Client
config Config
}
)
func New(client *llm.Client, cfg Config) *Describer {
return &Describer{
client: client,
config: cfg,
}
}
func (d *Describer) Describe(ctx context.Context, filename string, mimeType string, fileBase64 string) (*string, error) {
ag := agent.New(
"evidence_describer",
d.client,
agent.WithInstructions(systemPrompt),
agent.WithModel(d.config.Model),
agent.WithTemperature(d.config.Temp),
agent.WithMaxTokens(d.config.MaxTokens),
)
result, err := ag.Run(
ctx,
[]llm.Message{
{
Role: llm.RoleUser,
Parts: []llm.Part{
llm.TextPart{Text: fmt.Sprintf("Filename: %s", filename)},
llm.FilePart{
Data: fileBase64,
MimeType: mimeType,
Filename: filename,
},
},
},
},
)
if err != nil {
return nil, fmt.Errorf("cannot describe evidence: %w", err)
}
text := result.FinalMessage().Text()
return &text, nil
}

17
pkg/evidencedescriber/prompt.txt Normal file
View File

@@ -0,0 +1,17 @@
You are an ISO/SOC auditor writing evidence descriptions for a compliance review.
Input: a single image of an evidence file (screenshot or document export).
Output: plain text, 1–2 sentences. No markdown, no line breaks, no labels, no preamble.
Include these elements if clearly present; omit any that are not:
— System: the tool or platform shown (e.g. GitHub, Google Workspace, AWS).
— Setting: the specific configuration, feature, or state demonstrated.
— Scope: who or what it applies to (e.g. organization-wide, all users, a specific repository).
Use the language of the document. Do not include greetings, caveats, file names, image quality comments, or phrases like "This screenshot shows." Never guess — if something is not clearly visible, leave it out.
If the image is unreadable or is not compliance evidence, respond with exactly: "Unable to describe: unreadable or not recognized as compliance evidence."
Example — good: "Google Workspace admin console showing enforced 2-step verification for all users in the organization, with no exceptions permitted."
Example — bad: "This shows Google security settings."

25
pkg/llm/anthropic/provider.go
View File

@@ -16,11 +16,13 @@ package anthropic
import (
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"net/http"
"strconv"
"strings"
"time"
"github.com/anthropics/anthropic-sdk-go"
@@ -185,6 +187,8 @@ func buildMessages(messages []llm.Message) []anthropic.MessageParam {
},
),
)
case llm.FilePart:
blocks = append(blocks, buildFilePart(p))
}
}
out = append(out, anthropic.NewUserMessage(blocks...))
@@ -460,3 +464,24 @@ func (s *anthropicStream) mapStreamEvent(event *anthropic.MessageStreamEventUnio
return llm.ChatCompletionStreamEvent{}, false
}
}
func buildFilePart(p llm.FilePart) anthropic.ContentBlockParamUnion {
switch {
case strings.HasPrefix(p.MimeType, "image/"):
return anthropic.NewImageBlockBase64(p.MimeType, p.Data)
case p.MimeType == "application/pdf":
return anthropic.NewDocumentBlock(anthropic.Base64PDFSourceParam{
Data: p.Data,
})
case strings.HasPrefix(p.MimeType, "text/"):
decoded, err := base64.StdEncoding.DecodeString(p.Data)
if err != nil {
return anthropic.NewTextBlock(fmt.Sprintf("[file: %s, type: %s, error decoding content]", p.Filename, p.MimeType))
}
return anthropic.NewDocumentBlock(anthropic.PlainTextSourceParam{
Data: string(decoded),
})
default:
return anthropic.NewTextBlock(fmt.Sprintf("[file: %s, type: %s, unsupported format]", p.Filename, p.MimeType))
}
}

25
pkg/llm/openai/provider.go
View File

@@ -16,10 +16,13 @@ package openai
import (
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"net/http"
"strconv"
"strings"
"time"
"github.com/openai/openai-go"
@@ -188,6 +191,8 @@ func buildMessages(messages []llm.Message) []openai.ChatCompletionMessageParamUn
},
),
)
case llm.FilePart:
parts = append(parts, buildFilePart(p))
}
}
out = append(out, openai.UserMessage(parts))
@@ -450,3 +455,23 @@ func mapChunkToEvent(chunk *openai.ChatCompletionChunk) llm.ChatCompletionStream
return event
}
func buildFilePart(p llm.FilePart) openai.ChatCompletionContentPartUnionParam {
switch {
case strings.HasPrefix(p.MimeType, "image/"):
return openai.ImageContentPart(openai.ChatCompletionContentPartImageImageURLParam{
URL: fmt.Sprintf("data:%s;base64,%s", p.MimeType, p.Data),
})
case strings.HasPrefix(p.MimeType, "text/"):
decoded, err := base64.StdEncoding.DecodeString(p.Data)
if err != nil {
return openai.TextContentPart(fmt.Sprintf("[file: %s, type: %s, error decoding content]", p.Filename, p.MimeType))
}
return openai.TextContentPart(fmt.Sprintf("File: %s\n\n%s", p.Filename, string(decoded)))
default:
return openai.FileContentPart(openai.ChatCompletionContentPartFileFileParam{
FileData: param.NewOpt(fmt.Sprintf("data:%s;base64,%s", p.MimeType, p.Data)),
Filename: param.NewOpt(p.Filename),
})
}
}

7
pkg/llm/part.go
View File

@@ -26,7 +26,14 @@ type (
ImagePart struct {
URL string
}
FilePart struct {
Data string // base64-encoded content
MimeType string // e.g. "application/pdf", "text/csv", "image/png"
Filename string
}
)
func (TextPart) part() {}
func (ImagePart) part() {}
func (FilePart) part() {}

255
pkg/probo/evidence_description_worker.go Normal file
View File

@@ -0,0 +1,255 @@
// 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 probo
import (
"context"
"errors"
"fmt"
"sync"
"time"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/evidencedescriber"
"go.probo.inc/probo/pkg/filemanager"
)
type (
EvidenceDescriptionWorker struct {
pg *pg.Client
fileManager *filemanager.Service
describer *evidencedescriber.Describer
logger *log.Logger
interval time.Duration
staleAfter time.Duration
maxConcurrency int
}
EvidenceDescriptionWorkerOption func(*EvidenceDescriptionWorker)
)
func WithEvidenceDescriptionWorkerInterval(d time.Duration) EvidenceDescriptionWorkerOption {
return func(w *EvidenceDescriptionWorker) { w.interval = d }
}
func WithEvidenceDescriptionWorkerStaleAfter(d time.Duration) EvidenceDescriptionWorkerOption {
return func(w *EvidenceDescriptionWorker) { w.staleAfter = d }
}
func WithEvidenceDescriptionWorkerMaxConcurrency(n int) EvidenceDescriptionWorkerOption {
return func(w *EvidenceDescriptionWorker) {
if n > 0 {
w.maxConcurrency = n
}
}
}
func NewEvidenceDescriptionWorker(
pgClient *pg.Client,
fileManager *filemanager.Service,
describer *evidencedescriber.Describer,
logger *log.Logger,
opts ...EvidenceDescriptionWorkerOption,
) *EvidenceDescriptionWorker {
w := &EvidenceDescriptionWorker{
pg: pgClient,
fileManager: fileManager,
describer: describer,
logger: logger,
interval: 10 * time.Second,
staleAfter: 5 * time.Minute,
maxConcurrency: 10,
}
for _, opt := range opts {
opt(w)
}
return w
}
func (w *EvidenceDescriptionWorker) Run(ctx context.Context) error {
var (
wg sync.WaitGroup
sem = make(chan struct{}, w.maxConcurrency)
ticker = time.NewTicker(w.interval)
)
defer ticker.Stop()
defer wg.Wait()
for {
select {
case <-ctx.Done():
return ctx.Err()
case <-ticker.C:
nonCancelableCtx := context.WithoutCancel(ctx)
w.recoverStaleRows(nonCancelableCtx)
for {
if err := w.processNext(ctx, sem, &wg); err != nil {
if !errors.Is(err, coredata.ErrResourceNotFound) {
w.logger.ErrorCtx(nonCancelableCtx, "cannot claim evidence for description", log.Error(err))
}
break
}
}
}
}
}
func (w *EvidenceDescriptionWorker) processNext(ctx context.Context, sem chan struct{}, wg *sync.WaitGroup) error {
select {
case sem <- struct{}{}:
case <-ctx.Done():
return ctx.Err()
}
var (
evidence = coredata.Evidence{}
now = time.Now()
nonCancelableCtx = context.WithoutCancel(ctx)
)
if err := w.pg.WithTx(
nonCancelableCtx,
func(tx pg.Conn) error {
if err := evidence.LoadNextPendingDescriptionForUpdateSkipLocked(
nonCancelableCtx,
tx,
); err != nil {
return err
}
evidence.DescriptionStatus = coredata.EvidenceDescriptionStatusProcessing
evidence.DescriptionProcessingStartedAt = &now
evidence.UpdatedAt = now
if err := evidence.Update(nonCancelableCtx, tx, coredata.NewNoScope()); err != nil {
return fmt.Errorf("cannot update evidence: %w", err)
}
return nil
},
); err != nil {
<-sem
return err
}
wg.Add(1)
go func(evidence coredata.Evidence) {
defer wg.Done()
defer func() { <-sem }()
if err := w.describeAndCommit(nonCancelableCtx, &evidence); err != nil {
w.logger.ErrorCtx(
nonCancelableCtx,
"evidence description worker failure",
log.Error(err),
log.String("evidence_id", evidence.ID.String()),
)
if err := w.resetEvidence(nonCancelableCtx, &evidence); err != nil {
w.logger.ErrorCtx(nonCancelableCtx, "cannot reset evidence description status", log.Error(err))
}
}
}(evidence)
return nil
}
func (w *EvidenceDescriptionWorker) describeAndCommit(
ctx context.Context,
evidence *coredata.Evidence,
) error {
if evidence.EvidenceFileId == nil {
return fmt.Errorf("evidence %s has no file", evidence.ID)
}
scope := coredata.NewScopeFromObjectID(evidence.ID)
var file coredata.File
if err := w.pg.WithConn(
ctx,
func(conn pg.Conn) error {
if err := file.LoadByID(ctx, conn, scope, *evidence.EvidenceFileId); err != nil {
return fmt.Errorf("cannot load file: %w", err)
}
return nil
},
); err != nil {
return fmt.Errorf("cannot load file: %w", err)
}
base64Data, mimeType, err := w.fileManager.GetFileBase64(ctx, &file)
if err != nil {
return fmt.Errorf("cannot download file: %w", err)
}
description, err := w.describer.Describe(ctx, file.FileName, mimeType, base64Data)
if err != nil {
return fmt.Errorf("cannot describe evidence: %w", err)
}
return w.pg.WithTx(
ctx,
func(tx pg.Conn) error {
evidence.Description = description
evidence.DescriptionStatus = coredata.EvidenceDescriptionStatusCompleted
evidence.DescriptionProcessingStartedAt = nil
evidence.UpdatedAt = time.Now()
if err := evidence.Update(ctx, tx, scope); err != nil {
return fmt.Errorf("cannot update evidence: %w", err)
}
return nil
},
)
}
func (w *EvidenceDescriptionWorker) resetEvidence(
ctx context.Context,
evidence *coredata.Evidence,
) error {
scope := coredata.NewScopeFromObjectID(evidence.ID)
return w.pg.WithTx(
ctx,
func(tx pg.Conn) error {
evidence.DescriptionStatus = coredata.EvidenceDescriptionStatusPending
evidence.DescriptionProcessingStartedAt = nil
evidence.UpdatedAt = time.Now()
if err := evidence.Update(ctx, tx, scope); err != nil {
return fmt.Errorf("cannot update evidence: %w", err)
}
return nil
},
)
}
func (w *EvidenceDescriptionWorker) recoverStaleRows(ctx context.Context) {
if err := w.pg.WithConn(
ctx,
func(conn pg.Conn) error {
if err := coredata.ResetStaleDescriptionProcessing(ctx, conn, w.staleAfter); err != nil {
return fmt.Errorf("cannot reset stale description processing: %w", err)
}
return nil
},
); err != nil {
w.logger.ErrorCtx(ctx, "cannot recover stale evidence descriptions", log.Error(err))
}
}

15
pkg/probo/evidence_service.go
View File

@@ -92,13 +92,14 @@ func (s EvidenceService) UploadMeasureEvidence(
}
evidence := &coredata.Evidence{
ID: evidenceID,
MeasureID: req.MeasureID,
State: coredata.EvidenceStateFulfilled,
ReferenceID: "custom-evidence-" + referenceID.String(),
Type: coredata.EvidenceTypeFile,
CreatedAt: now,
UpdatedAt: now,
ID: evidenceID,
MeasureID: req.MeasureID,
State: coredata.EvidenceStateFulfilled,
ReferenceID: "custom-evidence-" + referenceID.String(),
Type: coredata.EvidenceTypeFile,
DescriptionStatus: coredata.EvidenceDescriptionStatusPending,
CreatedAt: now,
UpdatedAt: now,
}
err = s.svc.pg.WithTx(

17
pkg/probo/measure_service.go
View File

@@ -424,14 +424,15 @@ func (s MeasureService) Import(
evidenceDescription := req.Measures[i].Tasks[j].RequestedEvidences[k].Name
evidence := &coredata.Evidence{
State: coredata.EvidenceStateRequested,
ID: evidenceID,
TaskID: &task.ID,
ReferenceID: req.Measures[i].Tasks[j].RequestedEvidences[k].ReferenceID,
Type: req.Measures[i].Tasks[j].RequestedEvidences[k].Type,
Description: &evidenceDescription,
CreatedAt: now,
UpdatedAt: now,
State: coredata.EvidenceStateRequested,
ID: evidenceID,
TaskID: &task.ID,
ReferenceID: req.Measures[i].Tasks[j].RequestedEvidences[k].ReferenceID,
Type: req.Measures[i].Tasks[j].RequestedEvidences[k].Type,
Description: &evidenceDescription,
DescriptionStatus: coredata.EvidenceDescriptionStatusPending,
CreatedAt: now,
UpdatedAt: now,
}
if err := evidence.Upsert(ctx, tx, s.svc.scope); err != nil {

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