Add AI-powered evidence description generation
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:
@@ -7,6 +7,7 @@ export const evidenceFileQuery = graphql`
|
|||||||
node(id: $evidenceId) {
|
node(id: $evidenceId) {
|
||||||
... on Evidence {
|
... on Evidence {
|
||||||
id
|
id
|
||||||
|
description
|
||||||
file {
|
file {
|
||||||
mimeType
|
mimeType
|
||||||
fileName
|
fileName
|
||||||
|
|||||||
@@ -125,39 +125,48 @@ function EvidencePreviewContent({
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let preview;
|
||||||
|
|
||||||
if (evidence.file.mimeType?.startsWith("image/")) {
|
if (evidence.file.mimeType?.startsWith("image/")) {
|
||||||
return (
|
preview = (
|
||||||
<img
|
<img
|
||||||
src={evidence.file.downloadUrl}
|
src={evidence.file.downloadUrl}
|
||||||
alt={evidence.file.fileName}
|
alt={evidence.file.fileName}
|
||||||
className="max-h-[70vh] object-contain"
|
className="max-h-[70vh] object-contain"
|
||||||
/>
|
/>
|
||||||
);
|
);
|
||||||
}
|
} else if (evidence.file.mimeType?.includes("pdf")) {
|
||||||
|
preview = (
|
||||||
if (evidence.file.mimeType?.includes("pdf")) {
|
|
||||||
return (
|
|
||||||
<iframe
|
<iframe
|
||||||
src={evidence.file.downloadUrl}
|
src={evidence.file.downloadUrl}
|
||||||
className="w-full h-[70vh]"
|
className="w-full h-[70vh]"
|
||||||
title={evidence.file.fileName}
|
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 (
|
return (
|
||||||
<div className="flex flex-col items-center gap-2 justify-center">
|
<div className="space-y-4">
|
||||||
<IconWarning size={20} />
|
{preview}
|
||||||
<p className="text-txt-secondary text-center">
|
{evidence.description && (
|
||||||
{__("Preview not available for this file type")
|
<p className="text-txt-secondary text-sm">{evidence.description}</p>
|
||||||
+ " "
|
)}
|
||||||
+ evidence.file.mimeType}
|
|
||||||
</p>
|
|
||||||
<Button asChild variant="secondary" icon={IconArrowInbox}>
|
|
||||||
<a href={evidence.file.downloadUrl} target="_blank" rel="noreferrer">
|
|
||||||
{__("Download File")}
|
|
||||||
</a>
|
|
||||||
</Button>
|
|
||||||
</div>
|
</div>
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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 { usePageTitle } from "@probo/hooks";
|
||||||
import { useTranslate } from "@probo/i18n";
|
import { useTranslate } from "@probo/i18n";
|
||||||
import {
|
import {
|
||||||
@@ -79,7 +79,7 @@ export const evidenceFragment = graphql`
|
|||||||
mimeType
|
mimeType
|
||||||
size
|
size
|
||||||
}
|
}
|
||||||
type
|
description
|
||||||
createdAt
|
createdAt
|
||||||
canDelete: permission(action: "core:evidence:delete")
|
canDelete: permission(action: "core:evidence:delete")
|
||||||
}
|
}
|
||||||
@@ -123,8 +123,8 @@ export default function MeasureEvidencesTab() {
|
|||||||
<SortableTable {...pagination}>
|
<SortableTable {...pagination}>
|
||||||
<Thead>
|
<Thead>
|
||||||
<Tr>
|
<Tr>
|
||||||
<Th>{__("Evidence name")}</Th>
|
<Th>{__("Description")}</Th>
|
||||||
<Th>{__("Type")}</Th>
|
<Th>{__("File type")}</Th>
|
||||||
<Th>{__("File size")}</Th>
|
<Th>{__("File size")}</Th>
|
||||||
<Th>{__("Created at")}</Th>
|
<Th>{__("Created at")}</Th>
|
||||||
<Th width={50}></Th>
|
<Th width={50}></Th>
|
||||||
@@ -248,13 +248,12 @@ function EvidenceRow(props: {
|
|||||||
/>
|
/>
|
||||||
)}
|
)}
|
||||||
<Tr to={evidenceUrl}>
|
<Tr to={evidenceUrl}>
|
||||||
<Td>{evidence.file?.fileName}</Td>
|
|
||||||
<Td>
|
<Td>
|
||||||
{fileType(__, {
|
<span className="text-txt-secondary text-sm line-clamp-2">
|
||||||
type: evidence.type,
|
{evidence.description || "—"}
|
||||||
mimeType: evidence.file?.mimeType || "",
|
</span>
|
||||||
})}
|
|
||||||
</Td>
|
</Td>
|
||||||
|
<Td>{evidence.file?.mimeType || "—"}</Td>
|
||||||
<Td>{fileSize(__, evidence.file?.size || 0)}</Td>
|
<Td>{fileSize(__, evidence.file?.size || 0)}</Td>
|
||||||
<Td>{formatDate(evidence.createdAt)}</Td>
|
<Td>{formatDate(evidence.createdAt)}</Td>
|
||||||
<Td noLink>
|
<Td noLink>
|
||||||
|
|||||||
@@ -157,11 +157,35 @@ func (b *Builder) Build() (*probod.FullConfig, error) {
|
|||||||
CacheTTL: b.getEnvIntOrDefault("WEBHOOK_CACHE_TTL", 86400),
|
CacheTTL: b.getEnvIntOrDefault("WEBHOOK_CACHE_TTL", 86400),
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
OpenAI: probod.OpenAIConfig{
|
Agents: probod.AgentsConfig{
|
||||||
APIKey: b.getEnv("OPENAI_API_KEY"),
|
Providers: map[string]probod.LLMProviderConfig{
|
||||||
Temperature: b.getEnvFloatOrDefault("OPENAI_TEMPERATURE", 0.1),
|
"openai": {
|
||||||
ModelName: b.getEnvOrDefault("OPENAI_MODEL_NAME", "gpt-4o"),
|
Type: "openai",
|
||||||
MaxTokens: b.getEnvIntOrDefault("OPENAI_MAX_TOKENS", 4096),
|
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{
|
CustomDomains: probod.CustomDomainsConfig{
|
||||||
RenewalInterval: b.getEnvIntOrDefault("CUSTOM_DOMAINS_RENEWAL_INTERVAL", 3600),
|
RenewalInterval: b.getEnvIntOrDefault("CUSTOM_DOMAINS_RENEWAL_INTERVAL", 3600),
|
||||||
@@ -302,6 +326,25 @@ func (b *Builder) getEnvFloatOrDefault(key string, defaultValue float64) float64
|
|||||||
return defaultValue
|
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 {
|
func (b *Builder) getEnvBoolOrDefault(key string, defaultValue bool) bool {
|
||||||
if value := b.getEnv(key); value != "" {
|
if value := b.getEnv(key); value != "" {
|
||||||
if boolValue, err := strconv.ParseBool(value); err == nil {
|
if boolValue, err := strconv.ParseBool(value); err == nil {
|
||||||
|
|||||||
@@ -161,10 +161,20 @@ func TestBuilder_Build_Defaults(t *testing.T) {
|
|||||||
assert.Equal(t, 5, cfg.Probod.Notifications.Webhook.SenderInterval)
|
assert.Equal(t, 5, cfg.Probod.Notifications.Webhook.SenderInterval)
|
||||||
assert.Equal(t, 86400, cfg.Probod.Notifications.Webhook.CacheTTL)
|
assert.Equal(t, 86400, cfg.Probod.Notifications.Webhook.CacheTTL)
|
||||||
|
|
||||||
// OpenAI config
|
// Agents config — default
|
||||||
assert.Equal(t, 0.1, cfg.Probod.OpenAI.Temperature)
|
assert.Equal(t, "openai", cfg.Probod.Agents.Default.Provider)
|
||||||
assert.Equal(t, "gpt-4o", cfg.Probod.OpenAI.ModelName)
|
assert.Equal(t, "gpt-4o", cfg.Probod.Agents.Default.ModelName)
|
||||||
assert.Equal(t, 4096, cfg.Probod.OpenAI.MaxTokens)
|
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
|
// Custom domains config
|
||||||
assert.Equal(t, 3600, cfg.Probod.CustomDomains.RenewalInterval)
|
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_SENDER_INTERVAL"] = "10"
|
||||||
env["WEBHOOK_CACHE_TTL"] = "3600"
|
env["WEBHOOK_CACHE_TTL"] = "3600"
|
||||||
env["CONNECTOR_SLACK_SIGNING_SECRET"] = "slack-signing-secret"
|
env["CONNECTOR_SLACK_SIGNING_SECRET"] = "slack-signing-secret"
|
||||||
// OpenAI
|
// Agents — providers
|
||||||
env["OPENAI_API_KEY"] = "sk-test-key"
|
env["OPENAI_API_KEY"] = "sk-test-key"
|
||||||
env["OPENAI_TEMPERATURE"] = "0.5"
|
env["ANTHROPIC_API_KEY"] = "sk-ant-test-key"
|
||||||
env["OPENAI_MODEL_NAME"] = "gpt-4-turbo"
|
// Agents — default
|
||||||
env["OPENAI_MAX_TOKENS"] = "8192"
|
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
|
// Custom domains
|
||||||
env["CUSTOM_DOMAINS_RESOLVER_ADDR"] = "1.1.1.1:53"
|
env["CUSTOM_DOMAINS_RESOLVER_ADDR"] = "1.1.1.1:53"
|
||||||
env["ACME_ACCOUNT_KEY"] = "-----BEGIN EC PRIVATE KEY-----\ntest\n-----END EC PRIVATE KEY-----"
|
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, "slack-signing-secret", cfg.Probod.Notifications.Slack.SigningSecret)
|
||||||
assert.Equal(t, 10, cfg.Probod.Notifications.Webhook.SenderInterval)
|
assert.Equal(t, 10, cfg.Probod.Notifications.Webhook.SenderInterval)
|
||||||
assert.Equal(t, 3600, cfg.Probod.Notifications.Webhook.CacheTTL)
|
assert.Equal(t, 3600, cfg.Probod.Notifications.Webhook.CacheTTL)
|
||||||
// OpenAI
|
// Agents — providers
|
||||||
assert.Equal(t, "sk-test-key", cfg.Probod.OpenAI.APIKey)
|
assert.Equal(t, "openai", cfg.Probod.Agents.Providers["openai"].Type)
|
||||||
assert.Equal(t, 0.5, cfg.Probod.OpenAI.Temperature)
|
assert.Equal(t, "sk-test-key", cfg.Probod.Agents.Providers["openai"].APIKey)
|
||||||
assert.Equal(t, "gpt-4-turbo", cfg.Probod.OpenAI.ModelName)
|
assert.Equal(t, "anthropic", cfg.Probod.Agents.Providers["anthropic"].Type)
|
||||||
assert.Equal(t, 8192, cfg.Probod.OpenAI.MaxTokens)
|
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
|
// Custom domains
|
||||||
assert.Equal(t, "1.1.1.1:53", cfg.Probod.CustomDomains.ResolverAddr)
|
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)
|
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
106
pkg/cmd/evidence/delete/delete.go
Normal 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
36
pkg/cmd/evidence/evidence.go
Normal 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
199
pkg/cmd/evidence/list/list.go
Normal 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
171
pkg/cmd/evidence/view/view.go
Normal 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
|
||||||
|
}
|
||||||
@@ -25,6 +25,7 @@ import (
|
|||||||
cmdconfig "go.probo.inc/probo/pkg/cmd/config"
|
cmdconfig "go.probo.inc/probo/pkg/cmd/config"
|
||||||
cmdcontext "go.probo.inc/probo/pkg/cmd/context"
|
cmdcontext "go.probo.inc/probo/pkg/cmd/context"
|
||||||
"go.probo.inc/probo/pkg/cmd/control"
|
"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/finding"
|
||||||
"go.probo.inc/probo/pkg/cmd/framework"
|
"go.probo.inc/probo/pkg/cmd/framework"
|
||||||
"go.probo.inc/probo/pkg/cmd/org"
|
"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(cmdconfig.NewCmdConfig(f))
|
||||||
cmd.AddCommand(cmdcontext.NewCmdContext(f))
|
cmd.AddCommand(cmdcontext.NewCmdContext(f))
|
||||||
cmd.AddCommand(control.NewCmdControl(f))
|
cmd.AddCommand(control.NewCmdControl(f))
|
||||||
|
cmd.AddCommand(evidence.NewCmdEvidence(f))
|
||||||
cmd.AddCommand(finding.NewCmdFinding(f))
|
cmd.AddCommand(finding.NewCmdFinding(f))
|
||||||
cmd.AddCommand(framework.NewCmdFramework(f))
|
cmd.AddCommand(framework.NewCmdFramework(f))
|
||||||
cmd.AddCommand(org.NewCmdOrg(f))
|
cmd.AddCommand(org.NewCmdOrg(f))
|
||||||
|
|||||||
@@ -30,18 +30,20 @@ import (
|
|||||||
|
|
||||||
type (
|
type (
|
||||||
Evidence struct {
|
Evidence struct {
|
||||||
ID gid.GID `db:"id"`
|
ID gid.GID `db:"id"`
|
||||||
OrganizationID gid.GID `db:"organization_id"`
|
OrganizationID gid.GID `db:"organization_id"`
|
||||||
MeasureID gid.GID `db:"measure_id"`
|
MeasureID gid.GID `db:"measure_id"`
|
||||||
TaskID *gid.GID `db:"task_id"`
|
TaskID *gid.GID `db:"task_id"`
|
||||||
State EvidenceState `db:"state"`
|
State EvidenceState `db:"state"`
|
||||||
ReferenceID string `db:"reference_id"`
|
ReferenceID string `db:"reference_id"`
|
||||||
Type EvidenceType `db:"type"`
|
Type EvidenceType `db:"type"`
|
||||||
URL string `db:"url"`
|
URL string `db:"url"`
|
||||||
EvidenceFileId *gid.GID `db:"evidence_file_id"`
|
EvidenceFileId *gid.GID `db:"evidence_file_id"`
|
||||||
Description *string `db:"description"`
|
Description *string `db:"description"`
|
||||||
CreatedAt time.Time `db:"created_at"`
|
DescriptionStatus EvidenceDescriptionStatus `db:"description_status"`
|
||||||
UpdatedAt time.Time `db:"updated_at"`
|
DescriptionProcessingStartedAt *time.Time `db:"description_processing_started_at"`
|
||||||
|
CreatedAt time.Time `db:"created_at"`
|
||||||
|
UpdatedAt time.Time `db:"updated_at"`
|
||||||
}
|
}
|
||||||
|
|
||||||
Evidences []*Evidence
|
Evidences []*Evidence
|
||||||
@@ -89,6 +91,8 @@ INSERT INTO
|
|||||||
url,
|
url,
|
||||||
evidence_file_id,
|
evidence_file_id,
|
||||||
description,
|
description,
|
||||||
|
description_status,
|
||||||
|
description_processing_started_at,
|
||||||
created_at,
|
created_at,
|
||||||
updated_at
|
updated_at
|
||||||
)
|
)
|
||||||
@@ -103,6 +107,8 @@ VALUES (
|
|||||||
@url,
|
@url,
|
||||||
@evidence_file_id,
|
@evidence_file_id,
|
||||||
@description,
|
@description,
|
||||||
|
@description_status,
|
||||||
|
@description_processing_started_at,
|
||||||
@created_at,
|
@created_at,
|
||||||
@updated_at
|
@updated_at
|
||||||
)
|
)
|
||||||
@@ -114,18 +120,20 @@ WHERE evidences.state = 'REQUESTED';
|
|||||||
`
|
`
|
||||||
|
|
||||||
args := pgx.StrictNamedArgs{
|
args := pgx.StrictNamedArgs{
|
||||||
"tenant_id": scope.GetTenantID(),
|
"tenant_id": scope.GetTenantID(),
|
||||||
"evidence_id": e.ID,
|
"evidence_id": e.ID,
|
||||||
"measure_id": e.MeasureID,
|
"measure_id": e.MeasureID,
|
||||||
"task_id": e.TaskID,
|
"task_id": e.TaskID,
|
||||||
"reference_id": e.ReferenceID,
|
"reference_id": e.ReferenceID,
|
||||||
"evidence_file_id": e.EvidenceFileId,
|
"evidence_file_id": e.EvidenceFileId,
|
||||||
"created_at": e.CreatedAt,
|
"created_at": e.CreatedAt,
|
||||||
"updated_at": e.UpdatedAt,
|
"updated_at": e.UpdatedAt,
|
||||||
"state": e.State,
|
"state": e.State,
|
||||||
"type": e.Type,
|
"type": e.Type,
|
||||||
"url": e.URL,
|
"url": e.URL,
|
||||||
"description": e.Description,
|
"description": e.Description,
|
||||||
|
"description_status": e.DescriptionStatus,
|
||||||
|
"description_processing_started_at": e.DescriptionProcessingStartedAt,
|
||||||
}
|
}
|
||||||
_, err := conn.Exec(ctx, q, args)
|
_, err := conn.Exec(ctx, q, args)
|
||||||
return err
|
return err
|
||||||
@@ -150,6 +158,8 @@ INSERT INTO
|
|||||||
url,
|
url,
|
||||||
evidence_file_id,
|
evidence_file_id,
|
||||||
description,
|
description,
|
||||||
|
description_status,
|
||||||
|
description_processing_started_at,
|
||||||
created_at,
|
created_at,
|
||||||
updated_at
|
updated_at
|
||||||
)
|
)
|
||||||
@@ -165,25 +175,29 @@ VALUES (
|
|||||||
@url,
|
@url,
|
||||||
@evidence_file_id,
|
@evidence_file_id,
|
||||||
@description,
|
@description,
|
||||||
|
@description_status,
|
||||||
|
@description_processing_started_at,
|
||||||
@created_at,
|
@created_at,
|
||||||
@updated_at
|
@updated_at
|
||||||
)
|
)
|
||||||
`
|
`
|
||||||
|
|
||||||
args := pgx.StrictNamedArgs{
|
args := pgx.StrictNamedArgs{
|
||||||
"tenant_id": scope.GetTenantID(),
|
"tenant_id": scope.GetTenantID(),
|
||||||
"evidence_id": e.ID,
|
"evidence_id": e.ID,
|
||||||
"organization_id": e.OrganizationID,
|
"organization_id": e.OrganizationID,
|
||||||
"measure_id": e.MeasureID,
|
"measure_id": e.MeasureID,
|
||||||
"task_id": e.TaskID,
|
"task_id": e.TaskID,
|
||||||
"reference_id": e.ReferenceID,
|
"reference_id": e.ReferenceID,
|
||||||
"evidence_file_id": e.EvidenceFileId,
|
"evidence_file_id": e.EvidenceFileId,
|
||||||
"created_at": e.CreatedAt,
|
"created_at": e.CreatedAt,
|
||||||
"updated_at": e.UpdatedAt,
|
"updated_at": e.UpdatedAt,
|
||||||
"state": e.State,
|
"state": e.State,
|
||||||
"type": e.Type,
|
"type": e.Type,
|
||||||
"url": e.URL,
|
"url": e.URL,
|
||||||
"description": e.Description,
|
"description": e.Description,
|
||||||
|
"description_status": e.DescriptionStatus,
|
||||||
|
"description_processing_started_at": e.DescriptionProcessingStartedAt,
|
||||||
}
|
}
|
||||||
_, err := conn.Exec(ctx, q, args)
|
_, err := conn.Exec(ctx, q, args)
|
||||||
|
|
||||||
@@ -218,6 +232,8 @@ SELECT
|
|||||||
url,
|
url,
|
||||||
evidence_file_id,
|
evidence_file_id,
|
||||||
description,
|
description,
|
||||||
|
description_status,
|
||||||
|
description_processing_started_at,
|
||||||
created_at,
|
created_at,
|
||||||
updated_at
|
updated_at
|
||||||
FROM
|
FROM
|
||||||
@@ -299,6 +315,8 @@ SELECT
|
|||||||
url,
|
url,
|
||||||
evidence_file_id,
|
evidence_file_id,
|
||||||
description,
|
description,
|
||||||
|
description_status,
|
||||||
|
description_processing_started_at,
|
||||||
created_at,
|
created_at,
|
||||||
updated_at
|
updated_at
|
||||||
FROM
|
FROM
|
||||||
@@ -381,6 +399,8 @@ SELECT
|
|||||||
url,
|
url,
|
||||||
evidence_file_id,
|
evidence_file_id,
|
||||||
description,
|
description,
|
||||||
|
description_status,
|
||||||
|
description_processing_started_at,
|
||||||
created_at,
|
created_at,
|
||||||
updated_at
|
updated_at
|
||||||
FROM
|
FROM
|
||||||
@@ -426,6 +446,8 @@ SET
|
|||||||
evidence_file_id = @evidence_file_id,
|
evidence_file_id = @evidence_file_id,
|
||||||
url = @url,
|
url = @url,
|
||||||
description = @description,
|
description = @description,
|
||||||
|
description_status = @description_status,
|
||||||
|
description_processing_started_at = @description_processing_started_at,
|
||||||
updated_at = @updated_at
|
updated_at = @updated_at
|
||||||
WHERE
|
WHERE
|
||||||
%s
|
%s
|
||||||
@@ -435,13 +457,15 @@ WHERE
|
|||||||
q = fmt.Sprintf(q, scope.SQLFragment())
|
q = fmt.Sprintf(q, scope.SQLFragment())
|
||||||
|
|
||||||
args := pgx.StrictNamedArgs{
|
args := pgx.StrictNamedArgs{
|
||||||
"evidence_id": e.ID,
|
"evidence_id": e.ID,
|
||||||
"type": e.Type,
|
"type": e.Type,
|
||||||
"state": e.State,
|
"state": e.State,
|
||||||
"evidence_file_id": e.EvidenceFileId,
|
"evidence_file_id": e.EvidenceFileId,
|
||||||
"url": e.URL,
|
"url": e.URL,
|
||||||
"description": e.Description,
|
"description": e.Description,
|
||||||
"updated_at": e.UpdatedAt,
|
"description_status": e.DescriptionStatus,
|
||||||
|
"description_processing_started_at": e.DescriptionProcessingStartedAt,
|
||||||
|
"updated_at": e.UpdatedAt,
|
||||||
}
|
}
|
||||||
maps.Copy(args, scope.SQLArguments())
|
maps.Copy(args, scope.SQLArguments())
|
||||||
|
|
||||||
@@ -474,3 +498,71 @@ WHERE
|
|||||||
|
|
||||||
return nil
|
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
68
pkg/coredata/evidence_description_status.go
Normal 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
10
pkg/coredata/migrations/20260326T131816Z.sql
Normal 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
81
pkg/evidencedescriber/evidencedescriber.go
Normal 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
17
pkg/evidencedescriber/prompt.txt
Normal 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."
|
||||||
@@ -16,11 +16,13 @@ package anthropic
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"encoding/base64"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"net/http"
|
"net/http"
|
||||||
"strconv"
|
"strconv"
|
||||||
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/anthropics/anthropic-sdk-go"
|
"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...))
|
out = append(out, anthropic.NewUserMessage(blocks...))
|
||||||
@@ -460,3 +464,24 @@ func (s *anthropicStream) mapStreamEvent(event *anthropic.MessageStreamEventUnio
|
|||||||
return llm.ChatCompletionStreamEvent{}, false
|
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))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -16,10 +16,13 @@ package openai
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"encoding/base64"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
|
"fmt"
|
||||||
"net/http"
|
"net/http"
|
||||||
"strconv"
|
"strconv"
|
||||||
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/openai/openai-go"
|
"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))
|
out = append(out, openai.UserMessage(parts))
|
||||||
@@ -450,3 +455,23 @@ func mapChunkToEvent(chunk *openai.ChatCompletionChunk) llm.ChatCompletionStream
|
|||||||
|
|
||||||
return event
|
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),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -26,7 +26,14 @@ type (
|
|||||||
ImagePart struct {
|
ImagePart struct {
|
||||||
URL string
|
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 (TextPart) part() {}
|
||||||
func (ImagePart) part() {}
|
func (ImagePart) part() {}
|
||||||
|
func (FilePart) part() {}
|
||||||
|
|||||||
255
pkg/probo/evidence_description_worker.go
Normal file
255
pkg/probo/evidence_description_worker.go
Normal 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))
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -92,13 +92,14 @@ func (s EvidenceService) UploadMeasureEvidence(
|
|||||||
}
|
}
|
||||||
|
|
||||||
evidence := &coredata.Evidence{
|
evidence := &coredata.Evidence{
|
||||||
ID: evidenceID,
|
ID: evidenceID,
|
||||||
MeasureID: req.MeasureID,
|
MeasureID: req.MeasureID,
|
||||||
State: coredata.EvidenceStateFulfilled,
|
State: coredata.EvidenceStateFulfilled,
|
||||||
ReferenceID: "custom-evidence-" + referenceID.String(),
|
ReferenceID: "custom-evidence-" + referenceID.String(),
|
||||||
Type: coredata.EvidenceTypeFile,
|
Type: coredata.EvidenceTypeFile,
|
||||||
CreatedAt: now,
|
DescriptionStatus: coredata.EvidenceDescriptionStatusPending,
|
||||||
UpdatedAt: now,
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
}
|
}
|
||||||
|
|
||||||
err = s.svc.pg.WithTx(
|
err = s.svc.pg.WithTx(
|
||||||
|
|||||||
@@ -424,14 +424,15 @@ func (s MeasureService) Import(
|
|||||||
|
|
||||||
evidenceDescription := req.Measures[i].Tasks[j].RequestedEvidences[k].Name
|
evidenceDescription := req.Measures[i].Tasks[j].RequestedEvidences[k].Name
|
||||||
evidence := &coredata.Evidence{
|
evidence := &coredata.Evidence{
|
||||||
State: coredata.EvidenceStateRequested,
|
State: coredata.EvidenceStateRequested,
|
||||||
ID: evidenceID,
|
ID: evidenceID,
|
||||||
TaskID: &task.ID,
|
TaskID: &task.ID,
|
||||||
ReferenceID: req.Measures[i].Tasks[j].RequestedEvidences[k].ReferenceID,
|
ReferenceID: req.Measures[i].Tasks[j].RequestedEvidences[k].ReferenceID,
|
||||||
Type: req.Measures[i].Tasks[j].RequestedEvidences[k].Type,
|
Type: req.Measures[i].Tasks[j].RequestedEvidences[k].Type,
|
||||||
Description: &evidenceDescription,
|
Description: &evidenceDescription,
|
||||||
CreatedAt: now,
|
DescriptionStatus: coredata.EvidenceDescriptionStatusPending,
|
||||||
UpdatedAt: now,
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := evidence.Upsert(ctx, tx, s.svc.scope); err != nil {
|
if err := evidence.Upsert(ctx, tx, s.svc.scope); err != nil {
|
||||||
|
|||||||
@@ -22,13 +22,14 @@ import (
|
|||||||
"go.gearno.de/kit/log"
|
"go.gearno.de/kit/log"
|
||||||
"go.opentelemetry.io/otel/trace"
|
"go.opentelemetry.io/otel/trace"
|
||||||
"go.probo.inc/probo/pkg/llm"
|
"go.probo.inc/probo/pkg/llm"
|
||||||
|
llmanthropic "go.probo.inc/probo/pkg/llm/anthropic"
|
||||||
llmopenai "go.probo.inc/probo/pkg/llm/openai"
|
llmopenai "go.probo.inc/probo/pkg/llm/openai"
|
||||||
)
|
)
|
||||||
|
|
||||||
func buildLLMClient(cfg LLMConfig, l *log.Logger, tp trace.TracerProvider, r prometheus.Registerer) (*llm.Client, error) {
|
func buildLLMClient(cfg LLMProviderConfig, l *log.Logger, tp trace.TracerProvider, r prometheus.Registerer) (*llm.Client, error) {
|
||||||
provider := cfg.Provider
|
providerType := cfg.Type
|
||||||
if provider == "" {
|
if providerType == "" {
|
||||||
provider = "openai"
|
providerType = "openai"
|
||||||
}
|
}
|
||||||
|
|
||||||
httpClient := httpclient.DefaultPooledClient(
|
httpClient := httpclient.DefaultPooledClient(
|
||||||
@@ -37,7 +38,7 @@ func buildLLMClient(cfg LLMConfig, l *log.Logger, tp trace.TracerProvider, r pro
|
|||||||
httpclient.WithRegisterer(r),
|
httpclient.WithRegisterer(r),
|
||||||
)
|
)
|
||||||
|
|
||||||
switch provider {
|
switch providerType {
|
||||||
case "openai":
|
case "openai":
|
||||||
p := llmopenai.NewProvider(
|
p := llmopenai.NewProvider(
|
||||||
cfg.APIKey,
|
cfg.APIKey,
|
||||||
@@ -50,10 +51,19 @@ func buildLLMClient(cfg LLMConfig, l *log.Logger, tp trace.TracerProvider, r pro
|
|||||||
llm.WithTracerProvider(tp),
|
llm.WithTracerProvider(tp),
|
||||||
), nil
|
), nil
|
||||||
case "anthropic":
|
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":
|
case "bedrock":
|
||||||
return nil, fmt.Errorf("bedrock provider not yet wired; requires aws.Config")
|
return nil, fmt.Errorf("bedrock provider not yet wired; requires aws.Config")
|
||||||
default:
|
default:
|
||||||
return nil, fmt.Errorf("unsupported LLM provider: %q", provider)
|
return nil, fmt.Errorf("unsupported LLM provider type: %q", providerType)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -14,12 +14,48 @@
|
|||||||
|
|
||||||
package probod
|
package probod
|
||||||
|
|
||||||
type LLMConfig struct {
|
type (
|
||||||
Provider string `json:"provider"` // "openai", "anthropic", "bedrock"
|
// LLMProviderConfig holds authentication and connection settings for an
|
||||||
APIKey string `json:"api-key"` // for OpenAI and Anthropic
|
// LLM provider (e.g. OpenAI, Anthropic).
|
||||||
ModelName string `json:"model-name"`
|
LLMProviderConfig struct {
|
||||||
Temperature float64 `json:"temperature"`
|
Type string `json:"type"` // "openai", "anthropic", "bedrock"
|
||||||
MaxTokens int `json:"max-tokens"`
|
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
|
||||||
|
}
|
||||||
|
|||||||
@@ -51,6 +51,7 @@ import (
|
|||||||
"go.probo.inc/probo/pkg/crypto/keys"
|
"go.probo.inc/probo/pkg/crypto/keys"
|
||||||
"go.probo.inc/probo/pkg/crypto/passwdhash"
|
"go.probo.inc/probo/pkg/crypto/passwdhash"
|
||||||
"go.probo.inc/probo/pkg/esign"
|
"go.probo.inc/probo/pkg/esign"
|
||||||
|
"go.probo.inc/probo/pkg/evidencedescriber"
|
||||||
"go.probo.inc/probo/pkg/file"
|
"go.probo.inc/probo/pkg/file"
|
||||||
"go.probo.inc/probo/pkg/filemanager"
|
"go.probo.inc/probo/pkg/filemanager"
|
||||||
"go.probo.inc/probo/pkg/html2pdf"
|
"go.probo.inc/probo/pkg/html2pdf"
|
||||||
@@ -115,7 +116,7 @@ type (
|
|||||||
AWS AWSConfig `json:"aws"`
|
AWS AWSConfig `json:"aws"`
|
||||||
Notifications NotificationsConfig `json:"notifications"`
|
Notifications NotificationsConfig `json:"notifications"`
|
||||||
Connectors []ConnectorConfig `json:"connectors"`
|
Connectors []ConnectorConfig `json:"connectors"`
|
||||||
OpenAI OpenAIConfig `json:"openai"`
|
Agents AgentsConfig `json:"agents"`
|
||||||
ChromeDPAddr string `json:"chrome-dp-addr"`
|
ChromeDPAddr string `json:"chrome-dp-addr"`
|
||||||
CustomDomains CustomDomainsConfig `json:"custom-domains"`
|
CustomDomains CustomDomainsConfig `json:"custom-domains"`
|
||||||
SCIMBridge SCIMBridgeConfig `json:"scim-bridge"`
|
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 {
|
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)
|
fileManagerService := filemanager.NewService(s3Client)
|
||||||
@@ -454,10 +470,10 @@ func (impl *Implm) Run(
|
|||||||
impl.cfg.AWS.Bucket,
|
impl.cfg.AWS.Bucket,
|
||||||
baseURL.String(),
|
baseURL.String(),
|
||||||
impl.cfg.Auth.Cookie.Secret,
|
impl.cfg.Auth.Cookie.Secret,
|
||||||
llmClient,
|
proboLLMClient,
|
||||||
impl.cfg.OpenAI.ModelName,
|
proboAgentCfg.ModelName,
|
||||||
impl.cfg.OpenAI.Temperature,
|
*proboAgentCfg.Temperature,
|
||||||
impl.cfg.OpenAI.MaxTokens,
|
*proboAgentCfg.MaxTokens,
|
||||||
html2pdfConverter,
|
html2pdfConverter,
|
||||||
acmeService,
|
acmeService,
|
||||||
fileManagerService,
|
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())
|
trustCenterServerCtx, stopTrustCenterServer := context.WithCancel(context.Background())
|
||||||
defer stopTrustCenterServer()
|
defer stopTrustCenterServer()
|
||||||
wg.Go(
|
wg.Go(
|
||||||
@@ -644,6 +683,7 @@ func (impl *Implm) Run(
|
|||||||
stopWebhookSender()
|
stopWebhookSender()
|
||||||
stopESignService()
|
stopESignService()
|
||||||
stopMailingListWorker()
|
stopMailingListWorker()
|
||||||
|
stopEvidenceDescriptionWorker()
|
||||||
stopExportJobExporter()
|
stopExportJobExporter()
|
||||||
stopIAMService()
|
stopIAMService()
|
||||||
stopMailer()
|
stopMailer()
|
||||||
|
|||||||
Reference in New Issue
Block a user