diff --git a/.env.example b/.env.example index cd343a8437..2b21b915f6 100644 --- a/.env.example +++ b/.env.example @@ -71,8 +71,12 @@ # ── LLM-backed workers (inherit from LLM defaults when unset) ───────── # PROBO_AGENT_PROVIDER=anthropic # PROBO_AGENT_MODEL_NAME=claude-sonnet-4-6 -# EVIDENCE_DESCRIBER_PROVIDER=openai -# EVIDENCE_DESCRIBER_MODEL_NAME=gpt-4o-mini +# AGENT_EVIDENCE_ASSESSOR_PROVIDER=openai # inherits AGENT_DEFAULT_PROVIDER when unset +# AGENT_EVIDENCE_ASSESSOR_MODEL_NAME=gpt-5.4-mini # default; vision-capable model for evidence assessment +# AGENT_EVIDENCE_ASSESSOR_THINKING=2048 # extended-thinking budget; pair with a thinking-capable model (e.g. claude-sonnet-4-6) +# EVIDENCE_ASSESSOR_INTERVAL=10 +# EVIDENCE_ASSESSOR_STALE_AFTER=300 +# EVIDENCE_ASSESSOR_MAX_CONCURRENCY=10 # AGENT_THIRD_PARTY_VETTER_PROVIDER=openai # inherits AGENT_DEFAULT_PROVIDER when unset # AGENT_THIRD_PARTY_VETTER_MODEL_NAME=gpt-4o # inherits AGENT_DEFAULT_MODEL_NAME when unset # THIRD_PARTY_VETTING_INTERVAL=10 diff --git a/contrib/claude/file-naming.md b/contrib/claude/file-naming.md index 1dadcf969a..a2456dc24e 100644 --- a/contrib/claude/file-naming.md +++ b/contrib/claude/file-naming.md @@ -76,7 +76,7 @@ auth). In that case the file keeps its service name and the agent is built inline: ``` -pkg/evidencedescriber/evidencedescriber.go -- single-file describer service +pkg/evidenceassessor/assessment.go -- single-file evidence assessment service pkg/vetting/assessment.go -- third-party assessment service ``` diff --git a/pkg/agent/output_type.go b/pkg/agent/output_type.go index d99c132ed0..df36272ebc 100644 --- a/pkg/agent/output_type.go +++ b/pkg/agent/output_type.go @@ -50,3 +50,35 @@ func (o *OutputType) responseFormat() *llm.ResponseFormat { }, } } + +// DecorateEnum injects an explicit `enum` constraint on a single +// top-level property of the schema. jsonschema-go reads struct tags as +// free-form descriptions only, so enums cannot be encoded on the tag +// itself; callers chain one call per enum field after NewOutputType. +func (o *OutputType) DecorateEnum(field string, values []string) error { + var schema map[string]any + if err := json.Unmarshal(o.Schema, &schema); err != nil { + return fmt.Errorf("cannot unmarshal output type schema: %w", err) + } + + properties, ok := schema["properties"].(map[string]any) + if !ok { + return fmt.Errorf("output type schema has no properties") + } + + prop, ok := properties[field].(map[string]any) + if !ok { + return fmt.Errorf("output type schema has no %q property", field) + } + + prop["enum"] = values + + decorated, err := json.Marshal(schema) + if err != nil { + return fmt.Errorf("cannot marshal decorated output type schema: %w", err) + } + + o.Schema = decorated + + return nil +} diff --git a/pkg/bootstrap/builder.go b/pkg/bootstrap/builder.go index 42a729e386..1704901cde 100644 --- a/pkg/bootstrap/builder.go +++ b/pkg/bootstrap/builder.go @@ -193,6 +193,7 @@ func (b *Builder) Build() (*probodconfig.FullConfig, error) { 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)), + Thinking: b.getEnvIntPtr("AGENT_DEFAULT_THINKING"), }, Probo: probodconfig.LLMAgentConfig{ Provider: b.getEnvOrDefault("AGENT_PROBO_PROVIDER", ""), @@ -200,11 +201,15 @@ func (b *Builder) Build() (*probodconfig.FullConfig, error) { Temperature: b.getEnvFloatPtr("AGENT_PROBO_TEMPERATURE"), MaxTokens: b.getEnvIntPtr("AGENT_PROBO_MAX_TOKENS"), }, - EvidenceDescriber: probodconfig.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"), + EvidenceAssessor: probodconfig.LLMAgentConfig{ + Provider: b.getEnvOrDefault("AGENT_EVIDENCE_ASSESSOR_PROVIDER", ""), + // Evidence assessment is vision-first (it reads screenshots / + // PDFs / console exports), so it pins a current vision-capable + // model rather than inheriting the generic AGENT_DEFAULT model. + ModelName: b.getEnvOrDefault("AGENT_EVIDENCE_ASSESSOR_MODEL_NAME", "gpt-5.4-mini"), + Temperature: b.getEnvFloatPtr("AGENT_EVIDENCE_ASSESSOR_TEMPERATURE"), + MaxTokens: b.getEnvIntPtr("AGENT_EVIDENCE_ASSESSOR_MAX_TOKENS"), + Thinking: b.getEnvIntPtr("AGENT_EVIDENCE_ASSESSOR_THINKING"), }, ThirdPartyVetter: probodconfig.LLMAgentConfig{ Provider: b.getEnvOrDefault("AGENT_THIRD_PARTY_VETTER_PROVIDER", ""), @@ -247,10 +252,10 @@ func (b *Builder) Build() (*probodconfig.FullConfig, error) { ESign: probodconfig.ESignConfig{ TSAURL: b.getEnvOrDefault("ESIGN_TSA_URL", "http://timestamp.digicert.com"), }, - EvidenceDescriber: probodconfig.EvidenceDescriberConfig{ - Interval: b.getEnvIntOrDefault("EVIDENCE_DESCRIBER_INTERVAL", 10), - StaleAfter: b.getEnvIntOrDefault("EVIDENCE_DESCRIBER_STALE_AFTER", 300), - MaxConcurrency: b.getEnvIntOrDefault("EVIDENCE_DESCRIBER_MAX_CONCURRENCY", 10), + EvidenceAssessor: probodconfig.EvidenceAssessmentConfig{ + Interval: b.getEnvIntOrDefault("EVIDENCE_ASSESSOR_INTERVAL", 10), + StaleAfter: b.getEnvIntOrDefault("EVIDENCE_ASSESSOR_STALE_AFTER", 300), + MaxConcurrency: b.getEnvIntOrDefault("EVIDENCE_ASSESSOR_MAX_CONCURRENCY", 10), }, ThirdPartyVetting: probodconfig.ThirdPartyVettingWorkerConfig{ Interval: b.getEnvIntOrDefault("THIRD_PARTY_VETTING_INTERVAL", 10), diff --git a/pkg/bootstrap/builder_test.go b/pkg/bootstrap/builder_test.go index 3529489425..7ddf5c8b44 100644 --- a/pkg/bootstrap/builder_test.go +++ b/pkg/bootstrap/builder_test.go @@ -211,10 +211,11 @@ func TestBuilder_Build_Defaults(t *testing.T) { 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) + assert.Empty(t, cfg.Probod.Agents.EvidenceAssessor.Provider) + assert.Equal(t, "gpt-5.4-mini", cfg.Probod.Agents.EvidenceAssessor.ModelName) + assert.Nil(t, cfg.Probod.Agents.EvidenceAssessor.Temperature) + assert.Nil(t, cfg.Probod.Agents.EvidenceAssessor.MaxTokens) + assert.Nil(t, cfg.Probod.Agents.EvidenceAssessor.Thinking) assert.Empty(t, cfg.Probod.Agents.ThirdPartyVetter.Provider) assert.Empty(t, cfg.Probod.Agents.ThirdPartyVetter.ModelName) assert.Nil(t, cfg.Probod.Agents.ThirdPartyVetter.Temperature) @@ -238,6 +239,9 @@ func TestBuilder_Build_Defaults(t *testing.T) { assert.Equal(t, 10, cfg.Probod.ThirdPartyVetting.Interval) assert.Equal(t, 1500, cfg.Probod.ThirdPartyVetting.StaleAfter) assert.Equal(t, 1, cfg.Probod.ThirdPartyVetting.MaxConcurrency) + assert.Equal(t, 10, cfg.Probod.EvidenceAssessor.Interval) + assert.Equal(t, 300, cfg.Probod.EvidenceAssessor.StaleAfter) + assert.Equal(t, 10, cfg.Probod.EvidenceAssessor.MaxConcurrency) // Custom domains config assert.Equal(t, 3600, cfg.Probod.CustomDomains.RenewalInterval) @@ -322,11 +326,12 @@ func TestBuilder_Build_CustomValues(t *testing.T) { 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" + // Agents — evidence-assessor override + env["AGENT_EVIDENCE_ASSESSOR_PROVIDER"] = "anthropic" + env["AGENT_EVIDENCE_ASSESSOR_MODEL_NAME"] = "claude-sonnet-4-6" + env["AGENT_EVIDENCE_ASSESSOR_TEMPERATURE"] = "0.2" + env["AGENT_EVIDENCE_ASSESSOR_MAX_TOKENS"] = "4096" + env["AGENT_EVIDENCE_ASSESSOR_THINKING"] = "2048" // Agents — third-party-vetter override env["AGENT_THIRD_PARTY_VETTER_PROVIDER"] = "openai" env["AGENT_THIRD_PARTY_VETTER_MODEL_NAME"] = "gpt-4o" @@ -351,6 +356,9 @@ func TestBuilder_Build_CustomValues(t *testing.T) { env["THIRD_PARTY_VETTING_INTERVAL"] = "15" env["THIRD_PARTY_VETTING_STALE_AFTER"] = "1800" env["THIRD_PARTY_VETTING_MAX_CONCURRENCY"] = "2" + env["EVIDENCE_ASSESSOR_INTERVAL"] = "20" + env["EVIDENCE_ASSESSOR_STALE_AFTER"] = "600" + env["EVIDENCE_ASSESSOR_MAX_CONCURRENCY"] = "5" // 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-----" @@ -432,11 +440,12 @@ func TestBuilder_Build_CustomValues(t *testing.T) { // 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) + // Agents — evidence-assessor overrides + assert.Equal(t, "anthropic", cfg.Probod.Agents.EvidenceAssessor.Provider) + assert.Equal(t, "claude-sonnet-4-6", cfg.Probod.Agents.EvidenceAssessor.ModelName) + assert.Equal(t, new(0.2), cfg.Probod.Agents.EvidenceAssessor.Temperature) + assert.Equal(t, new(4096), cfg.Probod.Agents.EvidenceAssessor.MaxTokens) + assert.Equal(t, new(2048), cfg.Probod.Agents.EvidenceAssessor.Thinking) // Agents — third-party-vetter overrides assert.Equal(t, "openai", cfg.Probod.Agents.ThirdPartyVetter.Provider) assert.Equal(t, "gpt-4o", cfg.Probod.Agents.ThirdPartyVetter.ModelName) @@ -461,6 +470,9 @@ func TestBuilder_Build_CustomValues(t *testing.T) { assert.Equal(t, 15, cfg.Probod.ThirdPartyVetting.Interval) assert.Equal(t, 1800, cfg.Probod.ThirdPartyVetting.StaleAfter) assert.Equal(t, 2, cfg.Probod.ThirdPartyVetting.MaxConcurrency) + assert.Equal(t, 20, cfg.Probod.EvidenceAssessor.Interval) + assert.Equal(t, 600, cfg.Probod.EvidenceAssessor.StaleAfter) + assert.Equal(t, 5, cfg.Probod.EvidenceAssessor.MaxConcurrency) // Custom domains assert.Equal(t, "1.1.1.1:53", cfg.Probod.CustomDomains.ResolverAddr) assert.Equal(t, "-----BEGIN EC PRIVATE KEY-----\ntest\n-----END EC PRIVATE KEY-----", cfg.Probod.CustomDomains.ACME.AccountKey) diff --git a/pkg/coredata/connector.go b/pkg/coredata/connector.go index 403577f68f..3f0b452815 100644 --- a/pkg/coredata/connector.go +++ b/pkg/coredata/connector.go @@ -16,6 +16,7 @@ package coredata import ( "context" + "database/sql/driver" "encoding/json" "errors" "fmt" @@ -59,6 +60,17 @@ func (j *jsonRawMessageOrNull) Scan(src any) error { } } +// Value implements database/sql/driver.Valuer so pgx binds the column +// without an explicit helper at every call site. pgx rejects empty byte +// slices as invalid JSON, so an empty/nil value is sent as SQL NULL. +func (j jsonRawMessageOrNull) Value() (driver.Value, error) { + if len(j) == 0 { + return nil, nil + } + + return []byte(j), nil +} + type ( Connector struct { ID gid.GID `db:"id"` @@ -389,18 +401,13 @@ INSERT INTO connectors ( return fmt.Errorf("cannot encrypt connection: %w", err) } - var settingsArg any - if len(c.RawSettings) > 0 { - settingsArg = []byte(c.RawSettings) - } - args := pgx.StrictNamedArgs{ "id": c.ID, "tenant_id": scope.GetTenantID(), "organization_id": c.OrganizationID, "provider": c.Provider, "protocol": c.Protocol, - "settings": settingsArg, + "settings": c.RawSettings, "encrypted_connection": encryptedConnection, "created_at": c.CreatedAt, "updated_at": c.UpdatedAt, @@ -613,14 +620,9 @@ WHERE return fmt.Errorf("cannot encrypt connection: %w", err) } - var settingsArg any - if len(c.RawSettings) > 0 { - settingsArg = []byte(c.RawSettings) - } - args := pgx.StrictNamedArgs{ "id": c.ID, - "settings": settingsArg, + "settings": c.RawSettings, "encrypted_connection": encryptedConnection, "updated_at": c.UpdatedAt, } diff --git a/pkg/coredata/evidence.go b/pkg/coredata/evidence.go index 47a44ac542..c033496a75 100644 --- a/pkg/coredata/evidence.go +++ b/pkg/coredata/evidence.go @@ -31,20 +31,21 @@ 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"` - 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"` + 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"` + Assessment jsonRawMessageOrNull `db:"assessment"` + AssessmentStatus EvidenceAssessmentStatus `db:"assessment_status"` + AssessmentProcessingStartedAt *time.Time `db:"assessment_processing_started_at"` + CreatedAt time.Time `db:"created_at"` + UpdatedAt time.Time `db:"updated_at"` } Evidences []*Evidence @@ -117,8 +118,9 @@ INSERT INTO url, evidence_file_id, description, - description_status, - description_processing_started_at, + assessment, + assessment_status, + assessment_processing_started_at, created_at, updated_at ) @@ -133,8 +135,9 @@ VALUES ( @url, @evidence_file_id, @description, - @description_status, - @description_processing_started_at, + @assessment, + @assessment_status, + @assessment_processing_started_at, @created_at, @updated_at ) @@ -146,24 +149,29 @@ 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, - "description_status": e.DescriptionStatus, - "description_processing_started_at": e.DescriptionProcessingStartedAt, + "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, + "assessment": e.Assessment, + "assessment_status": e.AssessmentStatus, + "assessment_processing_started_at": e.AssessmentProcessingStartedAt, } + _, err := conn.Exec(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot upsert evidence: %w", err) + } - return err + return nil } func (e Evidence) Insert( @@ -185,8 +193,9 @@ INSERT INTO url, evidence_file_id, description, - description_status, - description_processing_started_at, + assessment, + assessment_status, + assessment_processing_started_at, created_at, updated_at ) @@ -202,29 +211,31 @@ VALUES ( @url, @evidence_file_id, @description, - @description_status, - @description_processing_started_at, + @assessment, + @assessment_status, + @assessment_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, - "description_status": e.DescriptionStatus, - "description_processing_started_at": e.DescriptionProcessingStartedAt, + "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, + "assessment": e.Assessment, + "assessment_status": e.AssessmentStatus, + "assessment_processing_started_at": e.AssessmentProcessingStartedAt, } _, err := conn.Exec(ctx, q, args) @@ -259,8 +270,9 @@ SELECT url, evidence_file_id, description, - description_status, - description_processing_started_at, + assessment, + assessment_status, + assessment_processing_started_at, created_at, updated_at FROM @@ -343,8 +355,9 @@ SELECT url, evidence_file_id, description, - description_status, - description_processing_started_at, + assessment, + assessment_status, + assessment_processing_started_at, created_at, updated_at FROM @@ -428,8 +441,9 @@ SELECT url, evidence_file_id, description, - description_status, - description_processing_started_at, + assessment, + assessment_status, + assessment_processing_started_at, created_at, updated_at FROM @@ -475,8 +489,9 @@ SET evidence_file_id = @evidence_file_id, url = @url, description = @description, - description_status = @description_status, - description_processing_started_at = @description_processing_started_at, + assessment = @assessment, + assessment_status = @assessment_status, + assessment_processing_started_at = @assessment_processing_started_at, updated_at = @updated_at WHERE %s @@ -486,21 +501,139 @@ 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, - "description_status": e.DescriptionStatus, - "description_processing_started_at": e.DescriptionProcessingStartedAt, - "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, + "assessment": e.Assessment, + "assessment_status": e.AssessmentStatus, + "assessment_processing_started_at": e.AssessmentProcessingStartedAt, + "updated_at": e.UpdatedAt, } maps.Copy(args, scope.SQLArguments()) _, err := conn.Exec(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot update evidence: %w", err) + } - return err + return nil +} + +// SetAssessmentFailed transitions the row to FAILED, touching only the +// two status columns so it never clobbers concurrent edits to +// description/assessment/url/etc. made while the LLM call was running. +// +// It guards on BOTH the PROCESSING status AND the claim's +// assessment_processing_started_at. The started_at predicate is what makes +// a late failure safe: if stale-recovery recycled this claim and another +// worker has since re-claimed the row with a fresh timestamp, the receiver +// still carries the original timestamp, so the UPDATE matches no row and is +// a no-op rather than flipping the new owner's in-flight claim to FAILED. +// It returns whether a row was transitioned so the caller can distinguish a +// real failure from a superseded one. The terminal status guard +// additionally prevents resurrecting an already COMPLETED/FAILED row. +// +// Note FAILED is terminal: no path re-queues it (the worker claims only +// PENDING and stale-recovery resets only PROCESSING), so a transient error +// permanently parks the evidence until it is re-triggered externally. This +// matches the prior describer behaviour and is intentional. +// +// Callers must carry the claim's AssessmentProcessingStartedAt on the +// receiver (Claim sets it). +func (e Evidence) SetAssessmentFailed( + ctx context.Context, + conn pg.Tx, + scope Scoper, +) (bool, error) { + q := ` +UPDATE + evidences +SET + assessment_status = @assessment_status, + assessment_processing_started_at = NULL, + updated_at = @updated_at +WHERE + %s + AND id = @evidence_id + AND assessment_status = @expected_status + AND assessment_processing_started_at = @expected_started_at +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{ + "evidence_id": e.ID, + "assessment_status": EvidenceAssessmentStatusFailed, + "expected_status": EvidenceAssessmentStatusProcessing, + "expected_started_at": e.AssessmentProcessingStartedAt, + "updated_at": time.Now(), + } + maps.Copy(args, scope.SQLArguments()) + + tag, err := conn.Exec(ctx, q, args) + if err != nil { + return false, fmt.Errorf("cannot mark evidence assessment failed: %w", err) + } + + return tag.RowsAffected() > 0, nil +} + +// SetAssessmentCompleted records a successful assessment, writing only +// the assessment-owned columns (assessment, description) and +// transitioning PROCESSING -> COMPLETED. Like SetAssessmentFailed it +// leaves user-editable columns (type/state/url/...) untouched so an edit +// made while the assessment was running is preserved, and it guards on +// both the PROCESSING status and the claim's +// assessment_processing_started_at. The guard makes the write idempotent +// and claim-scoped: if stale-recovery recycled the claim and another +// worker re-claimed (fresh started_at) or already finished it, no row +// matches and the method reports updated=false so the caller can drop the +// superseded result instead of clobbering the fresh one. Callers must +// populate Assessment and Description and carry the claim's +// AssessmentProcessingStartedAt on the receiver before calling. +func (e Evidence) SetAssessmentCompleted( + ctx context.Context, + conn pg.Tx, + scope Scoper, +) (bool, error) { + q := ` +UPDATE + evidences +SET + assessment = @assessment, + description = @description, + assessment_status = @assessment_status, + assessment_processing_started_at = NULL, + updated_at = @updated_at +WHERE + %s + AND id = @evidence_id + AND assessment_status = @expected_status + AND assessment_processing_started_at = @expected_started_at +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{ + "evidence_id": e.ID, + "assessment": e.Assessment, + "description": e.Description, + "assessment_status": EvidenceAssessmentStatusCompleted, + "expected_status": EvidenceAssessmentStatusProcessing, + "expected_started_at": e.AssessmentProcessingStartedAt, + "updated_at": time.Now(), + } + maps.Copy(args, scope.SQLArguments()) + + tag, err := conn.Exec(ctx, q, args) + if err != nil { + return false, fmt.Errorf("cannot mark evidence assessment completed: %w", err) + } + + return tag.RowsAffected() > 0, nil } func (e Evidence) Delete( @@ -529,7 +662,7 @@ WHERE return nil } -func (e *Evidence) LoadNextPendingDescriptionForUpdateSkipLocked( +func (e *Evidence) LoadNextPendingAssessmentForUpdateSkipLocked( ctx context.Context, conn pg.Tx, ) error { @@ -545,14 +678,15 @@ SELECT url, evidence_file_id, description, - description_status, - description_processing_started_at, + assessment, + assessment_status, + assessment_processing_started_at, created_at, updated_at FROM evidences WHERE - description_status = 'PENDING' + assessment_status = 'PENDING' AND evidence_file_id IS NOT NULL ORDER BY created_at ASC @@ -579,7 +713,7 @@ FOR UPDATE SKIP LOCKED; return nil } -func ResetStaleDescriptionProcessing( +func ResetStaleAssessmentProcessing( ctx context.Context, conn pg.Querier, staleAfter time.Duration, @@ -587,14 +721,17 @@ func ResetStaleDescriptionProcessing( q := ` UPDATE evidences SET - description_status = 'PENDING', - description_processing_started_at = NULL + assessment_status = 'PENDING', + assessment_processing_started_at = NULL WHERE - description_status = 'PROCESSING' - AND description_processing_started_at < $1; + assessment_status = 'PROCESSING' + AND assessment_processing_started_at < $1; ` _, err := conn.Exec(ctx, q, time.Now().Add(-staleAfter)) + if err != nil { + return fmt.Errorf("cannot reset stale assessment processing: %w", err) + } - return err + return nil } diff --git a/pkg/coredata/evidence_assessment.go b/pkg/coredata/evidence_assessment.go new file mode 100644 index 0000000000..0210efdc35 --- /dev/null +++ b/pkg/coredata/evidence_assessment.go @@ -0,0 +1,62 @@ +// Copyright (c) 2026 Probo Inc . +// +// 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 ( + "encoding/json" + "fmt" +) + +// SetAssessment marshals a typed assessment into Evidence.Assessment. +// Passing nil — or any value that marshals to JSON null, including a +// typed-nil pointer — clears the field rather than persisting JSON +// null. The shape of the assessment is defined by its producer (see +// pkg/evidenceassessor.EvidenceAssessment); this package is +// intentionally agnostic about the schema and only owns the raw JSONB +// round-trip. +func (e *Evidence) SetAssessment(v any) error { + if v == nil { + e.Assessment = nil + return nil + } + + data, err := json.Marshal(v) + if err != nil { + return fmt.Errorf("cannot marshal evidence assessment: %w", err) + } + + if string(data) == "null" { + e.Assessment = nil + return nil + } + + e.Assessment = data + + return nil +} + +// GetAssessment unmarshals Evidence.Assessment into dst. It is a no-op +// when the column is NULL/empty, leaving dst untouched. +func (e *Evidence) GetAssessment(dst any) error { + if len(e.Assessment) == 0 { + return nil + } + + if err := json.Unmarshal(e.Assessment, dst); err != nil { + return fmt.Errorf("cannot unmarshal evidence assessment: %w", err) + } + + return nil +} diff --git a/pkg/coredata/evidence_assessment_state_test.go b/pkg/coredata/evidence_assessment_state_test.go new file mode 100644 index 0000000000..036b2d6516 --- /dev/null +++ b/pkg/coredata/evidence_assessment_state_test.go @@ -0,0 +1,361 @@ +// Copyright (c) 2026 Probo Inc . +// +// 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_test + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.gearno.de/kit/pg" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/gid" +) + +// These integration tests exercise the assessment status state machine — the +// concurrency-sensitive guards the evidence-assessment worker relies on. They +// reuse newTestPgClient (access_entry_upsert_test.go) and skip when +// PROBO_TEST_PG_URL is unset, so `make test` stays a pure unit run. +// +// They lock in particular the claim-ownership guard on the terminal +// transitions: a worker whose claim was recycled by stale recovery and +// re-claimed by another worker must NOT clobber the new owner's PROCESSING +// row when its late SetAssessmentCompleted/SetAssessmentFailed fires. +// +// LoadNextPendingAssessmentForUpdateSkipLocked is intentionally not asserted +// here: it scans globally (no tenant scope), so it cannot be isolated against +// other rows in a shared parallel test database; the worker integration path +// covers its selection. + +type assessmentEvidenceFixture struct { + scope *coredata.Scope + organizationID gid.GID + measureID gid.GID +} + +// seedAssessmentFixture bootstraps the organization and measure that an +// evidence row's FKs require. +func seedAssessmentFixture(t *testing.T, ctx context.Context, client *pg.Client) assessmentEvidenceFixture { + t.Helper() + + tenantID := gid.NewTenantID() + scope := coredata.NewScope(tenantID) + organizationID := gid.New(tenantID, coredata.OrganizationEntityType) + measureID := gid.New(tenantID, coredata.MeasureEntityType) + now := time.Now().UTC().Truncate(time.Microsecond) + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + org := &coredata.Organization{ + ID: organizationID, + TenantID: tenantID, + Name: "Evidence Assessment Test Org", + CreatedAt: now, + UpdatedAt: now, + } + if err := org.Insert(ctx, tx); err != nil { + return err + } + + measure := &coredata.Measure{ + ID: measureID, + OrganizationID: organizationID, + Category: "Test", + Name: "Evidence Assessment Test Measure", + State: coredata.MeasureStateNotImplemented, + ReferenceID: "measure-assessment-" + measureID.String(), + CreatedAt: now, + UpdatedAt: now, + } + + return measure.Insert(ctx, tx, scope) + })) + + t.Cleanup(func() { + _ = client.WithTx(context.Background(), func(ctx context.Context, tx pg.Tx) error { + if _, err := tx.Exec(ctx, `DELETE FROM evidences WHERE measure_id = $1`, measureID); err != nil { + return err + } + + if _, err := tx.Exec(ctx, `DELETE FROM measures WHERE id = $1`, measureID); err != nil { + return err + } + + if _, err := tx.Exec(ctx, `DELETE FROM organizations WHERE id = $1`, organizationID); err != nil { + return err + } + + return nil + }) + }) + + return assessmentEvidenceFixture{ + scope: scope, + organizationID: organizationID, + measureID: measureID, + } +} + +// insertPendingEvidence inserts a PENDING evidence with a file attached. +func (fx assessmentEvidenceFixture) insertPendingEvidence( + t *testing.T, + ctx context.Context, + client *pg.Client, + referenceID string, + createdAt time.Time, +) gid.GID { + t.Helper() + + tenantID := fx.scope.GetTenantID() + evidenceID := gid.New(tenantID, coredata.EvidenceEntityType) + fileID := gid.New(tenantID, coredata.FileEntityType) + + evidence := coredata.Evidence{ + ID: evidenceID, + OrganizationID: fx.organizationID, + MeasureID: fx.measureID, + State: coredata.EvidenceStateFulfilled, + ReferenceID: referenceID, + Type: coredata.EvidenceTypeFile, + EvidenceFileID: &fileID, + AssessmentStatus: coredata.EvidenceAssessmentStatusPending, + CreatedAt: createdAt, + UpdatedAt: createdAt, + } + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + return evidence.Insert(ctx, tx, fx.scope) + })) + + return evidenceID +} + +func loadEvidence( + t *testing.T, + ctx context.Context, + client *pg.Client, + scope coredata.Scoper, + id gid.GID, +) coredata.Evidence { + t.Helper() + + var e coredata.Evidence + + require.NoError(t, client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { + return e.LoadByID(ctx, conn, scope, id) + })) + + return e +} + +// claimEvidence transitions a row to PROCESSING with the given started_at, +// mirroring what the worker's Claim does, and returns the claimed snapshot +// (carrying that started_at) so the caller can drive a terminal transition. +func claimEvidence( + t *testing.T, + ctx context.Context, + client *pg.Client, + scope coredata.Scoper, + id gid.GID, + startedAt time.Time, +) coredata.Evidence { + t.Helper() + + e := loadEvidence(t, ctx, client, scope, id) + e.AssessmentStatus = coredata.EvidenceAssessmentStatusProcessing + e.AssessmentProcessingStartedAt = &startedAt + e.UpdatedAt = startedAt + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + return e.Update(ctx, tx, scope) + })) + + return e +} + +func TestEvidence_SetAssessmentCompleted_TransitionsAndPreservesUserColumns(t *testing.T) { + t.Parallel() + + client := newTestPgClient(t) + ctx := context.Background() + fx := seedAssessmentFixture(t, ctx, client) + + t0 := time.Now().UTC().Truncate(time.Microsecond) + id := fx.insertPendingEvidence(t, ctx, client, "ev-complete-"+t0.Format(time.RFC3339Nano), t0) + + claimedAt := t0.Add(time.Minute) + claimed := claimEvidence(t, ctx, client, fx.scope, id, claimedAt) + + // Simulate a concurrent user edit to a user-owned column while the + // assessment ran; the terminal write must not clobber it. + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + _, err := tx.Exec(ctx, `UPDATE evidences SET url = $1 WHERE id = $2`, "https://edited.example", id) + return err + })) + + summary := "Admin console shows enforced MFA." + claimed.Description = &summary + require.NoError(t, claimed.SetAssessment(map[string]any{"summary": summary, "confidence": "HIGH"})) + + var updated bool + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + var err error + + updated, err = claimed.SetAssessmentCompleted(ctx, tx, fx.scope) + + return err + })) + assert.True(t, updated, "completing a PROCESSING row owned by this claim must update it") + + loaded := loadEvidence(t, ctx, client, fx.scope, id) + assert.Equal(t, coredata.EvidenceAssessmentStatusCompleted, loaded.AssessmentStatus) + require.NotNil(t, loaded.Description) + assert.Equal(t, summary, *loaded.Description) + assert.NotEmpty(t, loaded.Assessment) + assert.Nil(t, loaded.AssessmentProcessingStartedAt) + assert.Equal(t, "https://edited.example", loaded.URL, "a user edit during the run must be preserved") +} + +func TestEvidence_SetAssessmentCompleted_NoOpOnSupersededClaim(t *testing.T) { + t.Parallel() + + client := newTestPgClient(t) + ctx := context.Background() + fx := seedAssessmentFixture(t, ctx, client) + + t0 := time.Now().UTC().Truncate(time.Microsecond) + id := fx.insertPendingEvidence(t, ctx, client, "ev-superseded-complete-"+t0.Format(time.RFC3339Nano), t0) + + // Worker A claims the row. + firstClaim := claimEvidence(t, ctx, client, fx.scope, id, t0.Add(time.Minute)) + + // Stale recovery recycled A's claim and worker B re-claimed it with a + // fresh started_at — the row is now owned by B. + secondClaimAt := t0.Add(3 * time.Minute) + _ = claimEvidence(t, ctx, client, fx.scope, id, secondClaimAt) + + // Worker A finishes late and tries to commit its result. + summary := "stale result from A" + firstClaim.Description = &summary + require.NoError(t, firstClaim.SetAssessment(map[string]any{"summary": summary})) + + var updated bool + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + var err error + + updated, err = firstClaim.SetAssessmentCompleted(ctx, tx, fx.scope) + + return err + })) + assert.False(t, updated, "a superseded claim must not overwrite the re-claimed row") + + loaded := loadEvidence(t, ctx, client, fx.scope, id) + assert.Equal(t, coredata.EvidenceAssessmentStatusProcessing, loaded.AssessmentStatus, "row stays owned by the re-claim") + assert.Empty(t, loaded.Assessment, "the stale result must not be persisted") + require.NotNil(t, loaded.AssessmentProcessingStartedAt) + assert.WithinDuration(t, secondClaimAt, *loaded.AssessmentProcessingStartedAt, time.Second) +} + +func TestEvidence_SetAssessmentFailed_TransitionsOwningClaim(t *testing.T) { + t.Parallel() + + client := newTestPgClient(t) + ctx := context.Background() + fx := seedAssessmentFixture(t, ctx, client) + + t0 := time.Now().UTC().Truncate(time.Microsecond) + id := fx.insertPendingEvidence(t, ctx, client, "ev-fail-"+t0.Format(time.RFC3339Nano), t0) + + claimed := claimEvidence(t, ctx, client, fx.scope, id, t0.Add(time.Minute)) + + var updated bool + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + var err error + + updated, err = claimed.SetAssessmentFailed(ctx, tx, fx.scope) + + return err + })) + assert.True(t, updated, "failing a PROCESSING row owned by this claim must update it") + + loaded := loadEvidence(t, ctx, client, fx.scope, id) + assert.Equal(t, coredata.EvidenceAssessmentStatusFailed, loaded.AssessmentStatus) + assert.Nil(t, loaded.AssessmentProcessingStartedAt) +} + +func TestEvidence_SetAssessmentFailed_NoOpOnSupersededClaim(t *testing.T) { + t.Parallel() + + client := newTestPgClient(t) + ctx := context.Background() + fx := seedAssessmentFixture(t, ctx, client) + + t0 := time.Now().UTC().Truncate(time.Microsecond) + id := fx.insertPendingEvidence(t, ctx, client, "ev-superseded-fail-"+t0.Format(time.RFC3339Nano), t0) + + // Worker A claims, then B re-claims with a fresh started_at. + firstClaim := claimEvidence(t, ctx, client, fx.scope, id, t0.Add(time.Minute)) + secondClaimAt := t0.Add(3 * time.Minute) + _ = claimEvidence(t, ctx, client, fx.scope, id, secondClaimAt) + + // Worker A's late failure must not flip B's live claim to FAILED. + var updated bool + + require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + var err error + + updated, err = firstClaim.SetAssessmentFailed(ctx, tx, fx.scope) + + return err + })) + assert.False(t, updated, "a stale worker's failure must not clobber the re-claimed row") + + loaded := loadEvidence(t, ctx, client, fx.scope, id) + assert.Equal(t, coredata.EvidenceAssessmentStatusProcessing, loaded.AssessmentStatus, "row must remain PROCESSING under the new owner") + require.NotNil(t, loaded.AssessmentProcessingStartedAt) + assert.WithinDuration(t, secondClaimAt, *loaded.AssessmentProcessingStartedAt, time.Second) +} + +func TestResetStaleAssessmentProcessing_ResetsOnlyStaleClaims(t *testing.T) { + t.Parallel() + + client := newTestPgClient(t) + ctx := context.Background() + fx := seedAssessmentFixture(t, ctx, client) + + now := time.Now().UTC().Truncate(time.Microsecond) + staleID := fx.insertPendingEvidence(t, ctx, client, "ev-stale-"+now.Format(time.RFC3339Nano), now.Add(-time.Hour)) + freshID := fx.insertPendingEvidence(t, ctx, client, "ev-fresh-"+now.Format(time.RFC3339Nano), now.Add(-time.Hour)) + + // One claim started long ago (stale), one started now (fresh). + _ = claimEvidence(t, ctx, client, fx.scope, staleID, now.Add(-10*time.Minute)) + _ = claimEvidence(t, ctx, client, fx.scope, freshID, now) + + require.NoError(t, client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { + return coredata.ResetStaleAssessmentProcessing(ctx, conn, 5*time.Minute) + })) + + stale := loadEvidence(t, ctx, client, fx.scope, staleID) + assert.Equal(t, coredata.EvidenceAssessmentStatusPending, stale.AssessmentStatus, "a stale claim is recycled to PENDING") + assert.Nil(t, stale.AssessmentProcessingStartedAt) + + fresh := loadEvidence(t, ctx, client, fx.scope, freshID) + assert.Equal(t, coredata.EvidenceAssessmentStatusProcessing, fresh.AssessmentStatus, "a fresh claim is left running") + require.NotNil(t, fresh.AssessmentProcessingStartedAt) +} diff --git a/pkg/coredata/evidence_assessment_status.go b/pkg/coredata/evidence_assessment_status.go new file mode 100644 index 0000000000..4bc1d4082c --- /dev/null +++ b/pkg/coredata/evidence_assessment_status.go @@ -0,0 +1,78 @@ +// Copyright (c) 2026 Probo Inc . +// +// 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 ( + "encoding" + "fmt" +) + +type ( + EvidenceAssessmentStatus string +) + +const ( + EvidenceAssessmentStatusPending EvidenceAssessmentStatus = "PENDING" + EvidenceAssessmentStatusProcessing EvidenceAssessmentStatus = "PROCESSING" + EvidenceAssessmentStatusCompleted EvidenceAssessmentStatus = "COMPLETED" + EvidenceAssessmentStatusFailed EvidenceAssessmentStatus = "FAILED" +) + +var ( + _ fmt.Stringer = EvidenceAssessmentStatus("") + _ encoding.TextMarshaler = EvidenceAssessmentStatus("") + _ encoding.TextUnmarshaler = (*EvidenceAssessmentStatus)(nil) +) + +func EvidenceAssessmentStatuses() []EvidenceAssessmentStatus { + return []EvidenceAssessmentStatus{ + EvidenceAssessmentStatusPending, + EvidenceAssessmentStatusProcessing, + EvidenceAssessmentStatusCompleted, + EvidenceAssessmentStatusFailed, + } +} + +func (v EvidenceAssessmentStatus) IsValid() bool { + switch v { + case + EvidenceAssessmentStatusPending, + EvidenceAssessmentStatusProcessing, + EvidenceAssessmentStatusCompleted, + EvidenceAssessmentStatusFailed: + return true + } + + return false +} + +func (v EvidenceAssessmentStatus) String() string { + return string(v) +} + +func (v EvidenceAssessmentStatus) MarshalText() ([]byte, error) { + return []byte(v.String()), nil +} + +func (v *EvidenceAssessmentStatus) UnmarshalText(text []byte) error { + val := EvidenceAssessmentStatus(text) + if !val.IsValid() { + return fmt.Errorf("invalid EvidenceAssessmentStatus value: %q", string(text)) + } + + *v = val + + return nil +} diff --git a/pkg/coredata/evidence_assessment_status_test.go b/pkg/coredata/evidence_assessment_status_test.go new file mode 100644 index 0000000000..5279b3efc7 --- /dev/null +++ b/pkg/coredata/evidence_assessment_status_test.go @@ -0,0 +1,78 @@ +// Copyright (c) 2026 Probo Inc . +// +// 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 "testing" + +func TestEvidenceAssessmentStatusIsValid(t *testing.T) { + t.Parallel() + + for _, value := range EvidenceAssessmentStatuses() { + if !value.IsValid() { + t.Fatalf("IsValid() = false for %q", value) + } + } + + if EvidenceAssessmentStatus("BOGUS").IsValid() { + t.Fatal("IsValid() = true for invalid value") + } +} + +func TestEvidenceAssessmentStatusUnmarshalText(t *testing.T) { + t.Parallel() + + for _, value := range EvidenceAssessmentStatuses() { + t.Run(string(value), func(t *testing.T) { + t.Parallel() + + var got EvidenceAssessmentStatus + if err := got.UnmarshalText([]byte(value)); err != nil { + t.Fatalf("UnmarshalText(%q) returned error: %v", value, err) + } + + if got != value { + t.Fatalf("UnmarshalText(%q) = %q, want %q", value, got, value) + } + }) + } + + t.Run("invalid", func(t *testing.T) { + t.Parallel() + + var got EvidenceAssessmentStatus + if err := got.UnmarshalText([]byte("BOGUS")); err == nil { + t.Fatal("UnmarshalText(BOGUS) expected error") + } + }) +} + +func TestEvidenceAssessmentStatusMarshalText(t *testing.T) { + t.Parallel() + + for _, value := range EvidenceAssessmentStatuses() { + t.Run(string(value), func(t *testing.T) { + t.Parallel() + + got, err := value.MarshalText() + if err != nil { + t.Fatalf("MarshalText() returned error: %v", err) + } + + if string(got) != value.String() { + t.Fatalf("MarshalText() = %q, want %q", string(got), value.String()) + } + }) + } +} diff --git a/pkg/coredata/evidence_assessment_test.go b/pkg/coredata/evidence_assessment_test.go new file mode 100644 index 0000000000..82865c9b94 --- /dev/null +++ b/pkg/coredata/evidence_assessment_test.go @@ -0,0 +1,81 @@ +// Copyright (c) 2026 Probo Inc . +// +// 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 ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// A representative payload shape. The coredata layer is schema-agnostic +// so any JSON-serialisable struct should round-trip. +type testAssessmentPayload struct { + Summary string `json:"summary"` + Confidence string `json:"confidence"` + Frameworks []string `json:"frameworks"` +} + +func TestEvidence_SetAssessment_ReadBack(t *testing.T) { + t.Parallel() + + in := &testAssessmentPayload{ + Summary: "Google Workspace admin console showing enforced 2SV for all users.", + Confidence: "HIGH", + Frameworks: []string{"SOC2", "ISO27001"}, + } + + var e Evidence + require.NoError(t, e.SetAssessment(in)) + require.NotEmpty(t, e.Assessment, "Assessment raw bytes should be populated") + + var out testAssessmentPayload + require.NoError(t, e.GetAssessment(&out)) + assert.Equal(t, *in, out) +} + +func TestEvidence_SetAssessment_Nil_ClearsField(t *testing.T) { + t.Parallel() + + e := Evidence{} + require.NoError(t, e.SetAssessment(&testAssessmentPayload{Summary: "stub"})) + require.NotEmpty(t, e.Assessment) + + require.NoError(t, e.SetAssessment(nil)) + assert.Empty(t, e.Assessment, "untyped nil should clear the raw bytes") +} + +func TestEvidence_SetAssessment_TypedNil_ClearsField(t *testing.T) { + t.Parallel() + + e := Evidence{} + require.NoError(t, e.SetAssessment(&testAssessmentPayload{Summary: "stub"})) + require.NotEmpty(t, e.Assessment) + + var typedNil *testAssessmentPayload + require.NoError(t, e.SetAssessment(typedNil)) + assert.Empty(t, e.Assessment, "typed-nil pointer should clear, not persist JSON null") +} + +func TestEvidence_GetAssessment_EmptyIsNoOp(t *testing.T) { + t.Parallel() + + var e Evidence + + out := testAssessmentPayload{Summary: "unchanged"} + require.NoError(t, e.GetAssessment(&out)) + assert.Equal(t, "unchanged", out.Summary, "empty column should leave dst untouched") +} diff --git a/pkg/coredata/evidence_description_status.go b/pkg/coredata/evidence_description_status.go deleted file mode 100644 index 64c3d10a47..0000000000 --- a/pkg/coredata/evidence_description_status.go +++ /dev/null @@ -1,78 +0,0 @@ -// Copyright (c) 2026 Probo Inc . -// -// 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 ( - "encoding" - "fmt" -) - -type ( - EvidenceDescriptionStatus string -) - -const ( - EvidenceDescriptionStatusPending EvidenceDescriptionStatus = "PENDING" - EvidenceDescriptionStatusProcessing EvidenceDescriptionStatus = "PROCESSING" - EvidenceDescriptionStatusCompleted EvidenceDescriptionStatus = "COMPLETED" - EvidenceDescriptionStatusFailed EvidenceDescriptionStatus = "FAILED" -) - -var ( - _ fmt.Stringer = EvidenceDescriptionStatus("") - _ encoding.TextMarshaler = EvidenceDescriptionStatus("") - _ encoding.TextUnmarshaler = (*EvidenceDescriptionStatus)(nil) -) - -func EvidenceDescriptionStatuses() []EvidenceDescriptionStatus { - return []EvidenceDescriptionStatus{ - EvidenceDescriptionStatusPending, - EvidenceDescriptionStatusProcessing, - EvidenceDescriptionStatusCompleted, - EvidenceDescriptionStatusFailed, - } -} - -func (v EvidenceDescriptionStatus) IsValid() bool { - switch v { - case - EvidenceDescriptionStatusPending, - EvidenceDescriptionStatusProcessing, - EvidenceDescriptionStatusCompleted, - EvidenceDescriptionStatusFailed: - return true - } - - return false -} - -func (v EvidenceDescriptionStatus) String() string { - return string(v) -} - -func (v EvidenceDescriptionStatus) MarshalText() ([]byte, error) { - return []byte(v.String()), nil -} - -func (v *EvidenceDescriptionStatus) UnmarshalText(text []byte) error { - val := EvidenceDescriptionStatus(text) - if !val.IsValid() { - return fmt.Errorf("invalid EvidenceDescriptionStatus value: %q", string(text)) - } - - *v = val - - return nil -} diff --git a/pkg/coredata/migrations/20260424T418273Z.sql b/pkg/coredata/migrations/20260424T418273Z.sql new file mode 100644 index 0000000000..fe1e0bea9e --- /dev/null +++ b/pkg/coredata/migrations/20260424T418273Z.sql @@ -0,0 +1,22 @@ +-- Copyright (c) 2026 Probo Inc . +-- +-- 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. + +-- Nullable with no default: the assessor writes structured output here on +-- completion. Evidences that were already COMPLETED before this column +-- existed keep a NULL assessment and are not re-enqueued (the worker claims +-- only PENDING rows); this is intentional. GetAssessment is a no-op on NULL, +-- so readers must treat a NULL assessment as "not assessed", independent of +-- assessment_status. +ALTER TABLE evidences + ADD COLUMN assessment JSONB NULL; diff --git a/pkg/coredata/migrations/20260424T527194Z.sql b/pkg/coredata/migrations/20260424T527194Z.sql new file mode 100644 index 0000000000..eb14f9337d --- /dev/null +++ b/pkg/coredata/migrations/20260424T527194Z.sql @@ -0,0 +1,19 @@ +-- Copyright (c) 2026 Probo Inc . +-- +-- 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. + +ALTER TABLE evidences + RENAME COLUMN description_status TO assessment_status; + +ALTER TABLE evidences + RENAME COLUMN description_processing_started_at TO assessment_processing_started_at; diff --git a/pkg/coredata/migrations/20260424T603849Z.sql b/pkg/coredata/migrations/20260424T603849Z.sql new file mode 100644 index 0000000000..291e1b2f81 --- /dev/null +++ b/pkg/coredata/migrations/20260424T603849Z.sql @@ -0,0 +1,15 @@ +-- Copyright (c) 2026 Probo Inc . +-- +-- 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. + +ALTER TYPE evidence_description_status RENAME TO evidence_assessment_status; diff --git a/pkg/evidenceassessor/assessment.go b/pkg/evidenceassessor/assessment.go new file mode 100644 index 0000000000..741f88664c --- /dev/null +++ b/pkg/evidenceassessor/assessment.go @@ -0,0 +1,178 @@ +// Copyright (c) 2026 Probo Inc . +// +// 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 evidenceassessor + +import ( + "context" + _ "embed" + "encoding/json" + "fmt" + "time" + + "go.gearno.de/kit/log" + "go.probo.inc/probo/pkg/agent" + "go.probo.inc/probo/pkg/llm" +) + +// AssessmentTimeout is the hard upper bound on the assessor's single LLM +// turn. The assessor performs one tool-less structured-output turn, so +// this is far tighter than pkg/vetting's multi-turn budget. It exists so a +// stalled provider connection cannot pin a worker concurrency slot forever +// (the kit worker holds the slot until Process returns). Note it bounds +// only the LLM call, not the surrounding file download or commit, so it +// does not by itself guarantee Process finishes before the worker's +// stale-recovery window (EvidenceAssessmentConfig.StaleAfter, default 300s) +// elapses; correctness under a stale re-claim is enforced by the +// claim-ownership guard on the terminal transitions (see +// coredata.Evidence.SetAssessmentCompleted / SetAssessmentFailed), not by +// this timeout. +const AssessmentTimeout = 4 * time.Minute + +//go:embed prompt.txt +var systemPrompt string + +// confidenceEnum is the canonical set of allowed values for +// EvidenceAssessment.Confidence. Injected into the generated JSON +// schema via agent.OutputType.DecorateEnum because jsonschema struct +// tags cannot encode enum constraints directly. +var confidenceEnum = []string{"HIGH", "MEDIUM", "LOW"} + +type ( + // EvidenceAssessment is the structured output produced by the + // evidence assessor. The worker persists it as JSONB on the + // evidences table and derives Evidence.Description from Summary. + // + // Only Summary is surfaced to users today (through Description); the + // remaining structured fields (confidence, readable, issues, + // frameworks, ...) are stored as deliberate groundwork for a later + // PR that exposes them on the GraphQL/MCP Evidence type. They are + // intentionally write-only at the API boundary for now. + EvidenceAssessment struct { + Summary string `json:"summary" jsonschema:"One plain-text sentence (two at most) summarising what the evidence shows. When readable is false this field must restate the rejection reason."` + System string `json:"system" jsonschema:"Tool or platform visible on the file; empty string if not clearly identifiable."` + Setting string `json:"setting" jsonschema:"Specific configuration or state demonstrated; empty string if not clearly identifiable."` + Scope string `json:"scope" jsonschema:"Who or what the setting applies to; empty string if not stated."` + CapturedAt string `json:"captured_at" jsonschema:"ISO-8601 date or datetime visible on the file; empty string when no date is shown."` + Language string `json:"language" jsonschema:"BCP-47 language tag of the visible text; empty when unclear."` + Frameworks []string `json:"frameworks" jsonschema:"Compliance frameworks explicitly named on the file; empty when unclear."` + Issues []string `json:"issues" jsonschema:"Quality problems on the file itself; empty when the file is clean."` + Confidence string `json:"confidence" jsonschema:"One of HIGH, MEDIUM, LOW."` + Readable bool `json:"readable" jsonschema:"True if the file is a usable piece of compliance evidence."` + RejectionReason string `json:"rejection_reason" jsonschema:"Populated only when readable is false; empty otherwise."` + } + + Config struct { + Client *llm.Client + Model string + Temp float64 + MaxTokens int + // Thinking is the extended-thinking budget in tokens. 0 disables + // extended thinking entirely. + Thinking int + Logger *log.Logger + } + + Assessor struct { + cfg Config + outputType *agent.OutputType + } +) + +// New builds an Assessor. The structured output type is decorated once +// (enum on confidence) and cached so every call reuses the same schema. +func New(cfg Config) (*Assessor, error) { + outputType, err := agent.NewOutputType[EvidenceAssessment]("evidence_assessment") + if err != nil { + return nil, fmt.Errorf("cannot create evidence assessment output type: %w", err) + } + + if err := outputType.DecorateEnum("confidence", confidenceEnum); err != nil { + return nil, fmt.Errorf("cannot decorate evidence assessment schema: %w", err) + } + + return &Assessor{cfg: cfg, outputType: outputType}, nil +} + +// Assess runs the evidence assessor agent against a single uploaded +// file and returns a structured assessment. The worker persists the +// result as JSONB and derives Evidence.Description from Summary. +func (a *Assessor) Assess( + ctx context.Context, + filename string, + mimeType string, + fileBase64 string, +) (*EvidenceAssessment, error) { + // Detach from the caller so an unrelated cancellation cannot abort a + // committed assessment, but impose a hard deadline so a stalled + // provider response cannot run forever and pin the worker slot. + ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), AssessmentTimeout) + defer cancel() + + opts := []agent.Option{ + agent.WithInstructions(systemPrompt), + agent.WithModel(a.cfg.Model), + agent.WithMaxTokens(a.cfg.MaxTokens), + agent.WithOutputType(a.outputType), + } + if a.cfg.Thinking > 0 { + // Extended thinking and a custom temperature are mutually + // exclusive on Anthropic (the Messages API rejects any temperature + // other than 1 when thinking is enabled), and OpenAI reasoning + // models ignore temperature anyway, so only send a temperature + // when thinking is off. Mirrors pkg/vetting, which sets thinking + // without a temperature for the same reason. + opts = append(opts, agent.WithThinking(a.cfg.Thinking)) + } else { + opts = append(opts, agent.WithTemperature(a.cfg.Temp)) + } + + if a.cfg.Logger != nil { + opts = append(opts, agent.WithLogger(a.cfg.Logger)) + } + + assessmentAgent := agent.New("evidence_assessor", a.cfg.Client, opts...) + + result, err := assessmentAgent.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 assess evidence: %w", err) + } + + text := result.FinalMessage().Text() + + var out EvidenceAssessment + if err := json.Unmarshal([]byte(text), &out); err != nil { + // Surface the output size (not the content, which may carry + // sensitive evidence detail) so a MaxTokens-truncated response is + // distinguishable from other parse failures in logs. + return nil, fmt.Errorf("cannot parse evidence assessment (%d bytes of output): %w", len(text), err) + } + + return &out, nil +} diff --git a/pkg/evidenceassessor/assessment_test.go b/pkg/evidenceassessor/assessment_test.go new file mode 100644 index 0000000000..9c4498d435 --- /dev/null +++ b/pkg/evidenceassessor/assessment_test.go @@ -0,0 +1,55 @@ +// Copyright (c) 2026 Probo Inc . +// +// 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 evidenceassessor + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestNew_DecoratesConfidenceEnum guards the schema-mutation trick used +// to inject an enum into the generated JSON schema. If jsonschema-go +// ever reshapes the "properties" block or stops preserving the +// confidence field path, this test will fail loudly instead of silently +// shipping an un-constrained schema. +func TestNew_DecoratesConfidenceEnum(t *testing.T) { + t.Parallel() + + assessor, err := New(Config{}) + require.NoError(t, err) + require.NotNil(t, assessor.outputType) + + var schema map[string]any + require.NoError(t, json.Unmarshal(assessor.outputType.Schema, &schema)) + + properties, ok := schema["properties"].(map[string]any) + require.True(t, ok, "schema has no properties block") + + confidence, ok := properties["confidence"].(map[string]any) + require.True(t, ok, "schema has no confidence property") + + enumRaw, ok := confidence["enum"].([]any) + require.True(t, ok, "confidence has no enum array") + + actual := make([]string, len(enumRaw)) + for i, v := range enumRaw { + actual[i] = v.(string) + } + + assert.Equal(t, confidenceEnum, actual) +} diff --git a/pkg/evidenceassessor/prompt.txt b/pkg/evidenceassessor/prompt.txt new file mode 100644 index 0000000000..e04c352222 --- /dev/null +++ b/pkg/evidenceassessor/prompt.txt @@ -0,0 +1,21 @@ +You are an ISO/SOC auditor producing a structured record for one piece of compliance evidence (a screenshot, PDF, or document export). Your reader is another auditor; write for a compliance review, not a user. + +Output a single JSON object that conforms exactly to the provided schema. Emit nothing outside the object — no prose, no markdown, no code fences, no preamble. + +Rules that go beyond the schema: + +- Do not speculate. If a field cannot be determined from the file alone, leave it empty rather than inferring it from the filename, from the surrounding metadata, or from general knowledge about the platform. This applies especially to `frameworks`: only list a framework when it is explicitly named on the file, or strongly implied by a recognisable domain artifact (e.g. an AWS Config dashboard for SOC2). + +- `summary` is the only free-text field a downstream system is guaranteed to display, so it must be self-sufficient: one sentence (two at most), plain text, no markdown, no line breaks, no filenames, no image-quality comments, no phrases like "this screenshot shows". Use the language of the document. + +- When `readable` is false, `summary` must restate the rejection reason in a human-readable sentence. `rejection_reason` is populated in that case and must be empty when `readable` is true. + +- `confidence` = HIGH only when all three of `system`, `setting`, and `scope` are unambiguous on the file. MEDIUM when at most one is inferred or partially visible. LOW otherwise. + +- `readable` = false when the file is unreadable (blank, corrupt, mis-OCRed), off-topic (an invoice, a marketing page, a cat picture), or when the content is compliance-adjacent but too ambiguous to describe. + +Example — readable (no frameworks named on the file, so the array is empty): +{"summary":"Google Workspace admin console showing enforced 2-step verification for all users in the organization, with no exceptions permitted.","system":"Google Workspace","setting":"enforced 2-step verification","scope":"organization-wide","captured_at":"","language":"en","frameworks":[],"issues":[],"confidence":"HIGH","readable":true,"rejection_reason":""} + +Example — unreadable: +{"summary":"Unable to describe: image is blank and contains no recognisable compliance evidence.","system":"","setting":"","scope":"","captured_at":"","language":"","frameworks":[],"issues":[],"confidence":"LOW","readable":false,"rejection_reason":"Image is blank and contains no recognisable compliance evidence."} diff --git a/pkg/evidencedescriber/evidencedescriber.go b/pkg/evidencedescriber/evidencedescriber.go deleted file mode 100644 index db7d552169..0000000000 --- a/pkg/evidencedescriber/evidencedescriber.go +++ /dev/null @@ -1,82 +0,0 @@ -// Copyright (c) 2026 Probo Inc . -// -// 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 -} diff --git a/pkg/evidencedescriber/prompt.txt b/pkg/evidencedescriber/prompt.txt deleted file mode 100644 index 049b0da72a..0000000000 --- a/pkg/evidencedescriber/prompt.txt +++ /dev/null @@ -1,17 +0,0 @@ -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." \ No newline at end of file diff --git a/pkg/llm/registry_gen.go b/pkg/llm/registry_gen.go index b74bcd7ada..38ddc76307 100644 --- a/pkg/llm/registry_gen.go +++ b/pkg/llm/registry_gen.go @@ -14,11 +14,311 @@ // Code generated by genmodels; DO NOT EDIT. // Source: https://openrouter.ai/api/v1/models -// Generated: 2026-04-15T06:30:41Z +// Generated: 2026-06-04T20:22:07Z package llm var generatedModels = map[string]ModelDefinition{ + "anthropic/claude-opus-4.8-fast": { + Name: "Anthropic: Claude Opus 4.8 (Fast)", + ContextLength: 1000000, + MaxOutputTokens: 128000, + Supports: SupportedParameters{ + Temperature: false, + TopP: false, + TopK: false, + FrequencyPenalty: false, + PresencePenalty: false, + Stop: true, + Seed: false, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, + "anthropic/claude-opus-4.8": { + Name: "Anthropic: Claude Opus 4.8", + ContextLength: 1000000, + MaxOutputTokens: 128000, + Supports: SupportedParameters{ + Temperature: false, + TopP: false, + TopK: false, + FrequencyPenalty: false, + PresencePenalty: false, + Stop: true, + Seed: false, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, + "x-ai/grok-build-0.1": { + Name: "xAI: Grok Build 0.1", + ContextLength: 256000, + MaxOutputTokens: 0, + Supports: SupportedParameters{ + Temperature: true, + TopP: true, + TopK: false, + FrequencyPenalty: true, + PresencePenalty: true, + Stop: true, + Seed: true, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, + "google/gemini-3.5-flash": { + Name: "Google: Gemini 3.5 Flash", + ContextLength: 1048576, + MaxOutputTokens: 65536, + Supports: SupportedParameters{ + Temperature: true, + TopP: true, + TopK: false, + FrequencyPenalty: false, + PresencePenalty: false, + Stop: true, + Seed: true, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, + "anthropic/claude-opus-4.7-fast": { + Name: "Anthropic: Claude Opus 4.7 (Fast)", + ContextLength: 1000000, + MaxOutputTokens: 128000, + Supports: SupportedParameters{ + Temperature: false, + TopP: false, + TopK: false, + FrequencyPenalty: false, + PresencePenalty: false, + Stop: true, + Seed: false, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, + "google/gemini-3.1-flash-lite": { + Name: "Google: Gemini 3.1 Flash Lite", + ContextLength: 1048576, + MaxOutputTokens: 65536, + Supports: SupportedParameters{ + Temperature: true, + TopP: true, + TopK: false, + FrequencyPenalty: false, + PresencePenalty: false, + Stop: true, + Seed: true, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, + "openai/gpt-chat-latest": { + Name: "OpenAI: GPT Chat Latest", + ContextLength: 400000, + MaxOutputTokens: 128000, + Supports: SupportedParameters{ + Temperature: false, + TopP: false, + TopK: false, + FrequencyPenalty: true, + PresencePenalty: true, + Stop: true, + Seed: true, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: false, + }, + }, + "x-ai/grok-4.3": { + Name: "xAI: Grok 4.3", + ContextLength: 1000000, + MaxOutputTokens: 0, + Supports: SupportedParameters{ + Temperature: true, + TopP: true, + TopK: false, + FrequencyPenalty: true, + PresencePenalty: true, + Stop: true, + Seed: true, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, + "mistralai/mistral-medium-3-5": { + Name: "Mistral: Mistral Medium 3.5", + ContextLength: 262144, + MaxOutputTokens: 0, + Supports: SupportedParameters{ + Temperature: true, + TopP: true, + TopK: false, + FrequencyPenalty: true, + PresencePenalty: true, + Stop: true, + Seed: true, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, + "openai/gpt-5.5-pro": { + Name: "OpenAI: GPT-5.5 Pro", + ContextLength: 1050000, + MaxOutputTokens: 128000, + Supports: SupportedParameters{ + Temperature: false, + TopP: false, + TopK: false, + FrequencyPenalty: false, + PresencePenalty: false, + Stop: false, + Seed: true, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, + "openai/gpt-5.5": { + Name: "OpenAI: GPT-5.5", + ContextLength: 1050000, + MaxOutputTokens: 128000, + Supports: SupportedParameters{ + Temperature: false, + TopP: false, + TopK: false, + FrequencyPenalty: false, + PresencePenalty: false, + Stop: false, + Seed: true, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, + "deepseek/deepseek-v4-pro": { + Name: "DeepSeek: DeepSeek V4 Pro", + ContextLength: 1048576, + MaxOutputTokens: 384000, + Supports: SupportedParameters{ + Temperature: true, + TopP: true, + TopK: true, + FrequencyPenalty: true, + PresencePenalty: true, + Stop: true, + Seed: true, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, + "deepseek/deepseek-v4-flash": { + Name: "DeepSeek: DeepSeek V4 Flash", + ContextLength: 1048576, + MaxOutputTokens: 131072, + Supports: SupportedParameters{ + Temperature: true, + TopP: true, + TopK: true, + FrequencyPenalty: true, + PresencePenalty: true, + Stop: true, + Seed: true, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, + "openai/gpt-5.4-image-2": { + Name: "OpenAI: GPT-5.4 Image 2", + ContextLength: 272000, + MaxOutputTokens: 128000, + Supports: SupportedParameters{ + Temperature: false, + TopP: false, + TopK: false, + FrequencyPenalty: true, + PresencePenalty: true, + Stop: true, + Seed: true, + MaxTokens: true, + ToolChoice: false, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, + "anthropic/claude-opus-4.7": { + Name: "Anthropic: Claude Opus 4.7", + ContextLength: 1000000, + MaxOutputTokens: 128000, + Supports: SupportedParameters{ + Temperature: false, + TopP: false, + TopK: false, + FrequencyPenalty: false, + PresencePenalty: false, + Stop: true, + Seed: false, + MaxTokens: true, + ToolChoice: true, + ParallelToolCalls: false, + ResponseFormat: true, + StructuredOutputs: true, + Reasoning: true, + }, + }, "anthropic/claude-opus-4.6-fast": { Name: "Anthropic: Claude Opus 4.6 (Fast)", ContextLength: 1000000, @@ -89,7 +389,7 @@ var generatedModels = map[string]ModelDefinition{ TopK: false, FrequencyPenalty: false, PresencePenalty: false, - Stop: false, + Stop: true, Seed: true, MaxTokens: true, ToolChoice: true, @@ -102,7 +402,7 @@ var generatedModels = map[string]ModelDefinition{ "google/gemma-4-31b-it": { Name: "Google: Gemma 4 31B", ContextLength: 262144, - MaxOutputTokens: 0, + MaxOutputTokens: 16384, Supports: SupportedParameters{ Temperature: true, TopP: true, @@ -341,7 +641,7 @@ var generatedModels = map[string]ModelDefinition{ }, "google/gemini-3.1-flash-image-preview": { Name: "Google: Nano Banana 2 (Gemini 3.1 Flash Image Preview)", - ContextLength: 65536, + ContextLength: 131072, MaxOutputTokens: 65536, Supports: SupportedParameters{ Temperature: true, @@ -361,7 +661,7 @@ var generatedModels = map[string]ModelDefinition{ }, "google/gemini-3.1-pro-preview-customtools": { Name: "Google: Gemini 3.1 Pro Preview Custom Tools", - ContextLength: 1048576, + ContextLength: 1048756, MaxOutputTokens: 65536, Supports: SupportedParameters{ Temperature: true, @@ -539,30 +839,10 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: true, }, }, - "mistralai/mistral-small-creative": { - Name: "Mistral: Mistral Small Creative", - ContextLength: 32768, - MaxOutputTokens: 0, - Supports: SupportedParameters{ - Temperature: false, - TopP: false, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: false, - Seed: false, - MaxTokens: false, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: false, - StructuredOutputs: false, - Reasoning: false, - }, - }, "openai/gpt-5.2-chat": { Name: "OpenAI: GPT-5.2 Chat", ContextLength: 128000, - MaxOutputTokens: 32000, + MaxOutputTokens: 16384, Supports: SupportedParameters{ Temperature: false, TopP: false, @@ -759,30 +1039,10 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: false, }, }, - "deepseek/deepseek-v3.2-speciale": { - Name: "DeepSeek: DeepSeek V3.2 Speciale", - ContextLength: 163840, - MaxOutputTokens: 163840, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: true, - FrequencyPenalty: true, - PresencePenalty: true, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: false, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: true, - }, - }, "deepseek/deepseek-v3.2": { Name: "DeepSeek: DeepSeek V3.2", - ContextLength: 163840, - MaxOutputTokens: 0, + ContextLength: 131072, + MaxOutputTokens: 64000, Supports: SupportedParameters{ Temperature: true, TopP: true, @@ -839,26 +1099,6 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: true, }, }, - "x-ai/grok-4.1-fast": { - Name: "xAI: Grok 4.1 Fast", - ContextLength: 2000000, - MaxOutputTokens: 30000, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: false, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: true, - }, - }, "openai/gpt-5.1": { Name: "OpenAI: GPT-5.1", ContextLength: 400000, @@ -882,7 +1122,7 @@ var generatedModels = map[string]ModelDefinition{ "openai/gpt-5.1-chat": { Name: "OpenAI: GPT-5.1 Chat", ContextLength: 128000, - MaxOutputTokens: 16384, + MaxOutputTokens: 32000, Supports: SupportedParameters{ Temperature: false, TopP: false, @@ -922,7 +1162,7 @@ var generatedModels = map[string]ModelDefinition{ "openai/gpt-5.1-codex-mini": { Name: "OpenAI: GPT-5.1-Codex-Mini", ContextLength: 400000, - MaxOutputTokens: 128000, + MaxOutputTokens: 100000, Supports: SupportedParameters{ Temperature: false, TopP: false, @@ -1032,7 +1272,7 @@ var generatedModels = map[string]ModelDefinition{ Stop: true, Seed: true, MaxTokens: true, - ToolChoice: true, + ToolChoice: false, ParallelToolCalls: false, ResponseFormat: true, StructuredOutputs: true, @@ -1072,7 +1312,7 @@ var generatedModels = map[string]ModelDefinition{ Stop: true, Seed: true, MaxTokens: true, - ToolChoice: true, + ToolChoice: false, ParallelToolCalls: false, ResponseFormat: true, StructuredOutputs: true, @@ -1182,67 +1422,7 @@ var generatedModels = map[string]ModelDefinition{ "deepseek/deepseek-v3.2-exp": { Name: "DeepSeek: DeepSeek V3.2 Exp", ContextLength: 163840, - MaxOutputTokens: 65536, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: true, - FrequencyPenalty: true, - PresencePenalty: true, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: true, - }, - }, - "google/gemini-2.5-flash-lite-preview-09-2025": { - Name: "Google: Gemini 2.5 Flash Lite Preview 09-2025", - ContextLength: 1048576, - MaxOutputTokens: 65535, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: true, - }, - }, - "openai/gpt-5-codex": { - Name: "OpenAI: GPT-5 Codex", - ContextLength: 400000, - MaxOutputTokens: 128000, - Supports: SupportedParameters{ - Temperature: false, - TopP: false, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: false, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: true, - }, - }, - "deepseek/deepseek-v3.1-terminus": { - Name: "DeepSeek: DeepSeek V3.1 Terminus", - ContextLength: 163840, - MaxOutputTokens: 0, + MaxOutputTokens: 65536, Supports: SupportedParameters{ Temperature: true, TopP: true, @@ -1259,17 +1439,17 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: true, }, }, - "x-ai/grok-4-fast": { - Name: "xAI: Grok 4 Fast", - ContextLength: 2000000, - MaxOutputTokens: 30000, + "google/gemini-2.5-flash-lite-preview-09-2025": { + Name: "Google: Gemini 2.5 Flash Lite Preview 09-2025", + ContextLength: 1048576, + MaxOutputTokens: 65535, Supports: SupportedParameters{ Temperature: true, TopP: true, TopK: false, FrequencyPenalty: false, PresencePenalty: false, - Stop: false, + Stop: true, Seed: true, MaxTokens: true, ToolChoice: true, @@ -1279,17 +1459,17 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: true, }, }, - "x-ai/grok-code-fast-1": { - Name: "xAI: Grok Code Fast 1", - ContextLength: 256000, - MaxOutputTokens: 10000, + "openai/gpt-5-codex": { + Name: "OpenAI: GPT-5 Codex", + ContextLength: 400000, + MaxOutputTokens: 128000, Supports: SupportedParameters{ - Temperature: true, - TopP: true, + Temperature: false, + TopP: false, TopK: false, FrequencyPenalty: false, PresencePenalty: false, - Stop: true, + Stop: false, Seed: true, MaxTokens: true, ToolChoice: true, @@ -1299,10 +1479,10 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: true, }, }, - "deepseek/deepseek-chat-v3.1": { - Name: "DeepSeek: DeepSeek V3.1", - ContextLength: 32768, - MaxOutputTokens: 7168, + "deepseek/deepseek-v3.1-terminus": { + Name: "DeepSeek: DeepSeek V3.1 Terminus", + ContextLength: 163840, + MaxOutputTokens: 32768, Supports: SupportedParameters{ Temperature: true, TopP: true, @@ -1319,14 +1499,14 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: true, }, }, - "openai/gpt-4o-audio-preview": { - Name: "OpenAI: GPT-4o Audio", - ContextLength: 128000, - MaxOutputTokens: 16384, + "deepseek/deepseek-chat-v3.1": { + Name: "DeepSeek: DeepSeek V3.1", + ContextLength: 163840, + MaxOutputTokens: 32768, Supports: SupportedParameters{ Temperature: true, TopP: true, - TopK: false, + TopK: true, FrequencyPenalty: true, PresencePenalty: true, Stop: true, @@ -1336,7 +1516,7 @@ var generatedModels = map[string]ModelDefinition{ ParallelToolCalls: false, ResponseFormat: true, StructuredOutputs: true, - Reasoning: false, + Reasoning: true, }, }, "mistralai/mistral-medium-3.1": { @@ -1579,90 +1759,10 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: true, }, }, - "mistralai/devstral-medium": { - Name: "Mistral: Devstral Medium", - ContextLength: 131072, - MaxOutputTokens: 0, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: true, - PresencePenalty: true, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: false, - }, - }, - "mistralai/devstral-small": { - Name: "Mistral: Devstral Small 1.1", - ContextLength: 131072, - MaxOutputTokens: 0, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: true, - PresencePenalty: true, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: false, - }, - }, - "x-ai/grok-4": { - Name: "xAI: Grok 4", - ContextLength: 256000, - MaxOutputTokens: 0, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: false, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: true, - }, - }, - "google/gemma-3n-e2b-it:free": { - Name: "Google: Gemma 3n 2B (free)", - ContextLength: 8192, - MaxOutputTokens: 2048, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: false, - Seed: true, - MaxTokens: true, - ToolChoice: false, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: false, - Reasoning: false, - }, - }, "mistralai/mistral-small-3.2-24b-instruct": { Name: "Mistral: Mistral Small 3.2 24B", ContextLength: 128000, - MaxOutputTokens: 0, + MaxOutputTokens: 16384, Supports: SupportedParameters{ Temperature: true, TopP: true, @@ -1739,46 +1839,6 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: true, }, }, - "x-ai/grok-3-mini": { - Name: "xAI: Grok 3 Mini", - ContextLength: 131072, - MaxOutputTokens: 0, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: true, - }, - }, - "x-ai/grok-3": { - Name: "xAI: Grok 3", - ContextLength: 131072, - MaxOutputTokens: 0, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: true, - PresencePenalty: true, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: false, - }, - }, "google/gemini-2.5-pro-preview": { Name: "Google: Gemini 2.5 Pro Preview 06-05", ContextLength: 1048576, @@ -1802,7 +1862,7 @@ var generatedModels = map[string]ModelDefinition{ "deepseek/deepseek-r1-0528": { Name: "DeepSeek: R1 0528", ContextLength: 163840, - MaxOutputTokens: 0, + MaxOutputTokens: 32768, Supports: SupportedParameters{ Temperature: true, TopP: true, @@ -1826,7 +1886,7 @@ var generatedModels = map[string]ModelDefinition{ Supports: SupportedParameters{ Temperature: true, TopP: true, - TopK: true, + TopK: false, FrequencyPenalty: false, PresencePenalty: false, Stop: true, @@ -1859,26 +1919,6 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: true, }, }, - "google/gemma-3n-e4b-it:free": { - Name: "Google: Gemma 3n 4B (free)", - ContextLength: 8192, - MaxOutputTokens: 2048, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: false, - Seed: true, - MaxTokens: true, - ToolChoice: false, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: false, - Reasoning: false, - }, - }, "google/gemma-3n-e4b-it": { Name: "Google: Gemma 3n 4B", ContextLength: 32768, @@ -1942,7 +1982,7 @@ var generatedModels = map[string]ModelDefinition{ "meta-llama/llama-guard-4-12b": { Name: "Meta: Llama Guard 4 12B", ContextLength: 163840, - MaxOutputTokens: 0, + MaxOutputTokens: 16384, Supports: SupportedParameters{ Temperature: true, TopP: true, @@ -2079,46 +2119,6 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: false, }, }, - "x-ai/grok-3-mini-beta": { - Name: "xAI: Grok 3 Mini Beta", - ContextLength: 131072, - MaxOutputTokens: 0, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: false, - Reasoning: true, - }, - }, - "x-ai/grok-3-beta": { - Name: "xAI: Grok 3 Beta", - ContextLength: 131072, - MaxOutputTokens: 0, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: true, - PresencePenalty: true, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: false, - Reasoning: false, - }, - }, "meta-llama/llama-4-maverick": { Name: "Meta: Llama 4 Maverick", ContextLength: 1048576, @@ -2141,7 +2141,7 @@ var generatedModels = map[string]ModelDefinition{ }, "meta-llama/llama-4-scout": { Name: "Meta: Llama 4 Scout", - ContextLength: 327680, + ContextLength: 10000000, MaxOutputTokens: 16384, Supports: SupportedParameters{ Temperature: true, @@ -2162,7 +2162,7 @@ var generatedModels = map[string]ModelDefinition{ "deepseek/deepseek-chat-v3-0324": { Name: "DeepSeek: DeepSeek V3 0324", ContextLength: 163840, - MaxOutputTokens: 0, + MaxOutputTokens: 16384, Supports: SupportedParameters{ Temperature: true, TopP: true, @@ -2176,7 +2176,7 @@ var generatedModels = map[string]ModelDefinition{ ParallelToolCalls: false, ResponseFormat: true, StructuredOutputs: true, - Reasoning: true, + Reasoning: false, }, }, "openai/o1-pro": { @@ -2202,39 +2202,19 @@ var generatedModels = map[string]ModelDefinition{ "mistralai/mistral-small-3.1-24b-instruct": { Name: "Mistral: Mistral Small 3.1 24B", ContextLength: 128000, - MaxOutputTokens: 0, + MaxOutputTokens: 128000, Supports: SupportedParameters{ Temperature: true, TopP: true, TopK: true, FrequencyPenalty: true, PresencePenalty: true, - Stop: false, - Seed: true, - MaxTokens: true, - ToolChoice: false, - ParallelToolCalls: false, - ResponseFormat: false, - StructuredOutputs: false, - Reasoning: false, - }, - }, - "google/gemma-3-4b-it:free": { - Name: "Google: Gemma 3 4B (free)", - ContextLength: 32768, - MaxOutputTokens: 8192, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, Stop: true, Seed: true, MaxTokens: true, ToolChoice: false, ParallelToolCalls: false, - ResponseFormat: true, + ResponseFormat: false, StructuredOutputs: false, Reasoning: false, }, @@ -2242,47 +2222,27 @@ var generatedModels = map[string]ModelDefinition{ "google/gemma-3-4b-it": { Name: "Google: Gemma 3 4B", ContextLength: 131072, - MaxOutputTokens: 0, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: true, - FrequencyPenalty: true, - PresencePenalty: true, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: false, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: false, - Reasoning: false, - }, - }, - "google/gemma-3-12b-it:free": { - Name: "Google: Gemma 3 12B (free)", - ContextLength: 32768, - MaxOutputTokens: 8192, + MaxOutputTokens: 16384, Supports: SupportedParameters{ Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, + TopP: true, + TopK: true, + FrequencyPenalty: true, + PresencePenalty: true, Stop: true, Seed: true, MaxTokens: true, ToolChoice: false, ParallelToolCalls: false, - ResponseFormat: false, - StructuredOutputs: false, + ResponseFormat: true, + StructuredOutputs: true, Reasoning: false, }, }, "google/gemma-3-12b-it": { Name: "Google: Gemma 3 12B", ContextLength: 131072, - MaxOutputTokens: 0, + MaxOutputTokens: 16384, Supports: SupportedParameters{ Temperature: true, TopP: true, @@ -2292,7 +2252,7 @@ var generatedModels = map[string]ModelDefinition{ Stop: true, Seed: true, MaxTokens: true, - ToolChoice: false, + ToolChoice: true, ParallelToolCalls: false, ResponseFormat: true, StructuredOutputs: true, @@ -2339,26 +2299,6 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: false, }, }, - "google/gemma-3-27b-it:free": { - Name: "Google: Gemma 3 27B (free)", - ContextLength: 131072, - MaxOutputTokens: 8192, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: false, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: false, - Reasoning: false, - }, - }, "google/gemma-3-27b-it": { Name: "Google: Gemma 3 27B", ContextLength: 131072, @@ -2372,7 +2312,7 @@ var generatedModels = map[string]ModelDefinition{ Stop: true, Seed: true, MaxTokens: true, - ToolChoice: false, + ToolChoice: true, ParallelToolCalls: false, ResponseFormat: true, StructuredOutputs: true, @@ -2439,66 +2379,6 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: true, }, }, - "google/gemini-2.0-flash-lite-001": { - Name: "Google: Gemini 2.0 Flash Lite", - ContextLength: 1048576, - MaxOutputTokens: 8192, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: false, - }, - }, - "anthropic/claude-3.7-sonnet": { - Name: "Anthropic: Claude 3.7 Sonnet", - ContextLength: 200000, - MaxOutputTokens: 128000, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: true, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: true, - Seed: false, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: false, - StructuredOutputs: false, - Reasoning: true, - }, - }, - "anthropic/claude-3.7-sonnet:thinking": { - Name: "Anthropic: Claude 3.7 Sonnet (thinking)", - ContextLength: 200000, - MaxOutputTokens: 64000, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: true, - Seed: false, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: false, - StructuredOutputs: false, - Reasoning: true, - }, - }, "mistralai/mistral-saba": { Name: "Mistral: Saba", ContextLength: 32768, @@ -2522,14 +2402,14 @@ var generatedModels = map[string]ModelDefinition{ "meta-llama/llama-guard-3-8b": { Name: "Llama Guard 3 8B", ContextLength: 131072, - MaxOutputTokens: 0, + MaxOutputTokens: 131072, Supports: SupportedParameters{ Temperature: true, TopP: true, TopK: true, FrequencyPenalty: true, PresencePenalty: true, - Stop: false, + Stop: true, Seed: true, MaxTokens: true, ToolChoice: false, @@ -2559,26 +2439,6 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: true, }, }, - "google/gemini-2.0-flash-001": { - Name: "Google: Gemini 2.0 Flash", - ContextLength: 1048576, - MaxOutputTokens: 8192, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: false, - PresencePenalty: false, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: false, - }, - }, "openai/o3-mini": { Name: "OpenAI: o3 Mini", ContextLength: 200000, @@ -2615,13 +2475,13 @@ var generatedModels = map[string]ModelDefinition{ ToolChoice: false, ParallelToolCalls: false, ResponseFormat: true, - StructuredOutputs: false, + StructuredOutputs: true, Reasoning: false, }, }, "deepseek/deepseek-r1-distill-qwen-32b": { Name: "DeepSeek: R1 Distill Qwen 32B", - ContextLength: 32768, + ContextLength: 128000, MaxOutputTokens: 32768, Supports: SupportedParameters{ Temperature: true, @@ -2630,7 +2490,7 @@ var generatedModels = map[string]ModelDefinition{ FrequencyPenalty: true, PresencePenalty: true, Stop: true, - Seed: false, + Seed: true, MaxTokens: true, ToolChoice: false, ParallelToolCalls: false, @@ -2681,7 +2541,7 @@ var generatedModels = map[string]ModelDefinition{ }, "deepseek/deepseek-r1": { Name: "DeepSeek: R1", - ContextLength: 64000, + ContextLength: 163840, MaxOutputTokens: 16000, Supports: SupportedParameters{ Temperature: true, @@ -2694,15 +2554,15 @@ var generatedModels = map[string]ModelDefinition{ MaxTokens: true, ToolChoice: true, ParallelToolCalls: false, - ResponseFormat: false, - StructuredOutputs: false, + ResponseFormat: true, + StructuredOutputs: true, Reasoning: true, }, }, "deepseek/deepseek-chat": { Name: "DeepSeek: DeepSeek V3", - ContextLength: 163840, - MaxOutputTokens: 163840, + ContextLength: 131072, + MaxOutputTokens: 16000, Supports: SupportedParameters{ Temperature: true, TopP: true, @@ -2715,7 +2575,7 @@ var generatedModels = map[string]ModelDefinition{ ToolChoice: true, ParallelToolCalls: false, ResponseFormat: true, - StructuredOutputs: false, + StructuredOutputs: true, Reasoning: false, }, }, @@ -2741,7 +2601,7 @@ var generatedModels = map[string]ModelDefinition{ }, "meta-llama/llama-3.3-70b-instruct:free": { Name: "Meta: Llama 3.3 70B Instruct (free)", - ContextLength: 65536, + ContextLength: 131072, MaxOutputTokens: 0, Supports: SupportedParameters{ Temperature: true, @@ -2859,26 +2719,6 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: false, }, }, - "mistralai/mistral-large-2411": { - Name: "Mistral Large 2411", - ContextLength: 131072, - MaxOutputTokens: 0, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: true, - PresencePenalty: true, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: false, - }, - }, "mistralai/mistral-large-2407": { Name: "Mistral Large 2407", ContextLength: 131072, @@ -2899,26 +2739,6 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: false, }, }, - "mistralai/pixtral-large-2411": { - Name: "Mistral: Pixtral Large 2411", - ContextLength: 131072, - MaxOutputTokens: 0, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: true, - PresencePenalty: true, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: false, - }, - }, "anthropic/claude-3.5-haiku": { Name: "Anthropic: Claude 3.5 Haiku", ContextLength: 200000, @@ -2961,15 +2781,15 @@ var generatedModels = map[string]ModelDefinition{ }, "meta-llama/llama-3.2-3b-instruct": { Name: "Meta: Llama 3.2 3B Instruct", - ContextLength: 80000, - MaxOutputTokens: 0, + ContextLength: 131072, + MaxOutputTokens: 80000, Supports: SupportedParameters{ Temperature: true, TopP: true, TopK: true, FrequencyPenalty: true, PresencePenalty: true, - Stop: false, + Stop: true, Seed: true, MaxTokens: true, ToolChoice: false, @@ -2981,15 +2801,15 @@ var generatedModels = map[string]ModelDefinition{ }, "meta-llama/llama-3.2-1b-instruct": { Name: "Meta: Llama 3.2 1B Instruct", - ContextLength: 60000, - MaxOutputTokens: 0, + ContextLength: 131072, + MaxOutputTokens: 60000, Supports: SupportedParameters{ Temperature: true, TopP: true, TopK: true, FrequencyPenalty: true, PresencePenalty: true, - Stop: false, + Stop: true, Seed: true, MaxTokens: true, ToolChoice: false, @@ -3041,7 +2861,7 @@ var generatedModels = map[string]ModelDefinition{ }, "meta-llama/llama-3.1-8b-instruct": { Name: "Meta: Llama 3.1 8B Instruct", - ContextLength: 16384, + ContextLength: 131072, MaxOutputTokens: 16384, Supports: SupportedParameters{ Temperature: true, @@ -3062,7 +2882,7 @@ var generatedModels = map[string]ModelDefinition{ "meta-llama/llama-3.1-70b-instruct": { Name: "Meta: Llama 3.1 70B Instruct", ContextLength: 131072, - MaxOutputTokens: 0, + MaxOutputTokens: 16384, Supports: SupportedParameters{ Temperature: true, TopP: true, @@ -3075,14 +2895,14 @@ var generatedModels = map[string]ModelDefinition{ ToolChoice: true, ParallelToolCalls: false, ResponseFormat: true, - StructuredOutputs: false, + StructuredOutputs: true, Reasoning: false, }, }, "mistralai/mistral-nemo": { Name: "Mistral: Mistral Nemo", ContextLength: 131072, - MaxOutputTokens: 16384, + MaxOutputTokens: 0, Supports: SupportedParameters{ Temperature: true, TopP: true, @@ -3150,7 +2970,7 @@ var generatedModels = map[string]ModelDefinition{ FrequencyPenalty: true, PresencePenalty: true, Stop: true, - Seed: false, + Seed: true, MaxTokens: true, ToolChoice: false, ParallelToolCalls: false, @@ -3199,30 +3019,10 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: false, }, }, - "openai/gpt-4o:extended": { - Name: "OpenAI: GPT-4o (extended)", - ContextLength: 128000, - MaxOutputTokens: 64000, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: true, - PresencePenalty: true, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: false, - }, - }, "meta-llama/llama-3-8b-instruct": { Name: "Meta: Llama 3 8B Instruct", ContextLength: 8192, - MaxOutputTokens: 16384, + MaxOutputTokens: 8192, Supports: SupportedParameters{ Temperature: true, TopP: true, @@ -3232,7 +3032,7 @@ var generatedModels = map[string]ModelDefinition{ Stop: true, Seed: true, MaxTokens: true, - ToolChoice: true, + ToolChoice: false, ParallelToolCalls: false, ResponseFormat: true, StructuredOutputs: false, @@ -3255,7 +3055,7 @@ var generatedModels = map[string]ModelDefinition{ ToolChoice: false, ParallelToolCalls: false, ResponseFormat: false, - StructuredOutputs: false, + StructuredOutputs: true, Reasoning: false, }, }, @@ -3379,26 +3179,6 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: false, }, }, - "mistralai/mixtral-8x7b-instruct": { - Name: "Mistral: Mixtral 8x7B Instruct", - ContextLength: 32768, - MaxOutputTokens: 16384, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: true, - FrequencyPenalty: true, - PresencePenalty: true, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: false, - Reasoning: false, - }, - }, "openai/gpt-4-1106-preview": { Name: "OpenAI: GPT-4 Turbo (older v1106)", ContextLength: 128000, @@ -3439,26 +3219,6 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: false, }, }, - "mistralai/mistral-7b-instruct-v0.1": { - Name: "Mistral: Mistral 7B Instruct v0.1", - ContextLength: 2824, - MaxOutputTokens: 0, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: true, - FrequencyPenalty: true, - PresencePenalty: true, - Stop: false, - Seed: true, - MaxTokens: true, - ToolChoice: false, - ParallelToolCalls: false, - ResponseFormat: false, - StructuredOutputs: false, - Reasoning: false, - }, - }, "openai/gpt-3.5-turbo-16k": { Name: "OpenAI: GPT-3.5 Turbo 16k", ContextLength: 16385, @@ -3479,26 +3239,6 @@ var generatedModels = map[string]ModelDefinition{ Reasoning: false, }, }, - "openai/gpt-4-0314": { - Name: "OpenAI: GPT-4 (older v0314)", - ContextLength: 8191, - MaxOutputTokens: 4096, - Supports: SupportedParameters{ - Temperature: true, - TopP: true, - TopK: false, - FrequencyPenalty: true, - PresencePenalty: true, - Stop: true, - Seed: true, - MaxTokens: true, - ToolChoice: true, - ParallelToolCalls: false, - ResponseFormat: true, - StructuredOutputs: true, - Reasoning: false, - }, - }, "openai/gpt-4": { Name: "OpenAI: GPT-4", ContextLength: 8191, diff --git a/pkg/probo/evidence_assessment_worker.go b/pkg/probo/evidence_assessment_worker.go new file mode 100644 index 0000000000..a7bc1bf041 --- /dev/null +++ b/pkg/probo/evidence_assessment_worker.go @@ -0,0 +1,226 @@ +// Copyright (c) 2026 Probo Inc . +// +// 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" + "time" + + "go.gearno.de/kit/log" + "go.gearno.de/kit/pg" + "go.gearno.de/kit/worker" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/evidenceassessor" + "go.probo.inc/probo/pkg/filemanager" +) + +type ( + evidenceAssessmentHandler struct { + pg *pg.Client + fileManager *filemanager.Service + assessor *evidenceassessor.Assessor + logger *log.Logger + staleAfter time.Duration + } + + EvidenceAssessmentWorkerConfig struct { + StaleAfter time.Duration + } +) + +func NewEvidenceAssessmentWorker( + pgClient *pg.Client, + fileManager *filemanager.Service, + assessor *evidenceassessor.Assessor, + logger *log.Logger, + cfg EvidenceAssessmentWorkerConfig, + opts ...worker.Option, +) *worker.Worker[coredata.Evidence] { + staleAfter := cfg.StaleAfter + if staleAfter == 0 { + staleAfter = 5 * time.Minute + } + + h := &evidenceAssessmentHandler{ + pg: pgClient, + fileManager: fileManager, + assessor: assessor, + logger: logger, + staleAfter: staleAfter, + } + + return worker.New( + "evidence-assessment-worker", + h, + logger, + opts..., + ) +} + +func (h *evidenceAssessmentHandler) Claim(ctx context.Context) (coredata.Evidence, error) { + var evidence coredata.Evidence + + if err := h.pg.WithTx( + ctx, + func(ctx context.Context, tx pg.Tx) error { + if err := evidence.LoadNextPendingAssessmentForUpdateSkipLocked(ctx, tx); err != nil { + return err + } + + now := time.Now() + evidence.AssessmentStatus = coredata.EvidenceAssessmentStatusProcessing + evidence.AssessmentProcessingStartedAt = &now + + evidence.UpdatedAt = now + if err := evidence.Update(ctx, tx, coredata.NewScopeFromObjectID(evidence.ID)); err != nil { + return fmt.Errorf("cannot update evidence: %w", err) + } + + return nil + }, + ); err != nil { + if errors.Is(err, coredata.ErrResourceNotFound) { + return coredata.Evidence{}, worker.ErrNoTask + } + + return coredata.Evidence{}, err + } + + return evidence, nil +} + +func (h *evidenceAssessmentHandler) Process(ctx context.Context, evidence coredata.Evidence) error { + if err := h.assessAndCommit(ctx, evidence); err != nil { + h.logger.ErrorCtx( + ctx, + "evidence assessment worker failure", + log.Error(err), + log.String("evidence_id", evidence.ID.String()), + ) + + if err := h.failEvidence(ctx, evidence); err != nil { + h.logger.ErrorCtx(ctx, "cannot mark evidence assessment as failed", log.Error(err)) + } + + return err + } + + return nil +} + +func (h *evidenceAssessmentHandler) RecoverStale(ctx context.Context) error { + return h.pg.WithConn( + ctx, + func(ctx context.Context, conn pg.Querier) error { + if err := coredata.ResetStaleAssessmentProcessing(ctx, conn, h.staleAfter); err != nil { + return fmt.Errorf("cannot reset stale assessment processing: %w", err) + } + + return nil + }, + ) +} + +// assessAndCommit deliberately takes evidence by value; mutations made +// inside the transaction stay local, so a failed commit cannot leak +// partial state to the subsequent failEvidence call. +func (h *evidenceAssessmentHandler) assessAndCommit( + ctx context.Context, + evidence coredata.Evidence, +) error { + if evidence.EvidenceFileID == nil { + return fmt.Errorf("cannot assess evidence %s: no file attached", evidence.ID) + } + + scope := coredata.NewScopeFromObjectID(evidence.ID) + + var file coredata.File + + if err := h.pg.WithConn( + ctx, + func(ctx context.Context, conn pg.Querier) error { + return file.LoadByID(ctx, conn, scope, *evidence.EvidenceFileID) + }, + ); err != nil { + return fmt.Errorf("cannot load evidence file: %w", err) + } + + base64Data, mimeType, err := h.fileManager.GetFileBase64(ctx, &file) + if err != nil { + return fmt.Errorf("cannot download file: %w", err) + } + + assessment, err := h.assessor.Assess(ctx, file.FileName, mimeType, base64Data) + if err != nil { + return fmt.Errorf("cannot assess evidence: %w", err) + } + + return h.pg.WithTx( + ctx, + func(ctx context.Context, tx pg.Tx) error { + if err := evidence.SetAssessment(assessment); err != nil { + return err + } + + summary := assessment.Summary + evidence.Description = &summary + + updated, err := evidence.SetAssessmentCompleted(ctx, tx, scope) + if err != nil { + return err + } + + if !updated { + // The claim was recycled by stale recovery and finished by + // another worker; drop this result rather than clobber it. + h.logger.WarnCtx( + ctx, + "evidence assessment claim superseded, dropping result", + log.String("evidence_id", evidence.ID.String()), + ) + } + + return nil + }, + ) +} + +func (h *evidenceAssessmentHandler) failEvidence(ctx context.Context, evidence coredata.Evidence) error { + return h.pg.WithTx( + ctx, + func(ctx context.Context, tx pg.Tx) error { + updated, err := evidence.SetAssessmentFailed(ctx, tx, coredata.NewScopeFromObjectID(evidence.ID)) + if err != nil { + return err + } + + if !updated { + // The claim was recycled by stale recovery and re-claimed (or + // already finished) by another worker; the claim-scoped guard + // matched no row, so we leave the status to the current owner + // rather than forcing FAILED over its in-flight claim. + h.logger.WarnCtx( + ctx, + "evidence assessment failure superseded, leaving status to current owner", + log.String("evidence_id", evidence.ID.String()), + ) + } + + return nil + }, + ) +} diff --git a/pkg/probo/evidence_description_worker.go b/pkg/probo/evidence_description_worker.go deleted file mode 100644 index 85db1a001d..0000000000 --- a/pkg/probo/evidence_description_worker.go +++ /dev/null @@ -1,210 +0,0 @@ -// Copyright (c) 2026 Probo Inc . -// -// 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" - "time" - - "go.gearno.de/kit/log" - "go.gearno.de/kit/pg" - "go.gearno.de/kit/worker" - "go.probo.inc/probo/pkg/coredata" - "go.probo.inc/probo/pkg/evidencedescriber" - "go.probo.inc/probo/pkg/filemanager" -) - -type ( - evidenceDescriptionHandler struct { - pg *pg.Client - fileManager *filemanager.Service - describer *evidencedescriber.Describer - logger *log.Logger - staleAfter time.Duration - } - - EvidenceDescriptionWorkerConfig struct { - StaleAfter time.Duration - } -) - -func NewEvidenceDescriptionWorker( - pgClient *pg.Client, - fileManager *filemanager.Service, - describer *evidencedescriber.Describer, - logger *log.Logger, - cfg EvidenceDescriptionWorkerConfig, - opts ...worker.Option, -) *worker.Worker[coredata.Evidence] { - staleAfter := cfg.StaleAfter - if staleAfter == 0 { - staleAfter = 5 * time.Minute - } - - h := &evidenceDescriptionHandler{ - pg: pgClient, - fileManager: fileManager, - describer: describer, - logger: logger, - staleAfter: staleAfter, - } - - return worker.New( - "evidence-description-worker", - h, - logger, - opts..., - ) -} - -func (h *evidenceDescriptionHandler) Claim(ctx context.Context) (coredata.Evidence, error) { - var evidence coredata.Evidence - - if err := h.pg.WithTx( - ctx, - func(ctx context.Context, tx pg.Tx) error { - if err := evidence.LoadNextPendingDescriptionForUpdateSkipLocked(ctx, tx); err != nil { - return err - } - - now := time.Now() - evidence.DescriptionStatus = coredata.EvidenceDescriptionStatusProcessing - evidence.DescriptionProcessingStartedAt = &now - - evidence.UpdatedAt = now - if err := evidence.Update(ctx, tx, coredata.NewNoScope()); err != nil { - return fmt.Errorf("cannot update evidence: %w", err) - } - - return nil - }, - ); err != nil { - if errors.Is(err, coredata.ErrResourceNotFound) { - return coredata.Evidence{}, worker.ErrNoTask - } - - return coredata.Evidence{}, err - } - - return evidence, nil -} - -func (h *evidenceDescriptionHandler) Process(ctx context.Context, evidence coredata.Evidence) error { - if err := h.describeAndCommit(ctx, &evidence); err != nil { - h.logger.ErrorCtx( - ctx, - "evidence description worker failure", - log.Error(err), - log.String("evidence_id", evidence.ID.String()), - ) - - if err := h.failEvidence(ctx, &evidence); err != nil { - h.logger.ErrorCtx(ctx, "cannot mark evidence description as failed", log.Error(err)) - } - - return err - } - - return nil -} - -func (h *evidenceDescriptionHandler) RecoverStale(ctx context.Context) error { - return h.pg.WithConn( - ctx, - func(ctx context.Context, conn pg.Querier) error { - if err := coredata.ResetStaleDescriptionProcessing(ctx, conn, h.staleAfter); err != nil { - return fmt.Errorf("cannot reset stale description processing: %w", err) - } - - return nil - }, - ) -} - -func (h *evidenceDescriptionHandler) 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 := h.pg.WithConn( - ctx, - func(ctx context.Context, conn pg.Querier) 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 := h.fileManager.GetFileBase64(ctx, &file) - if err != nil { - return fmt.Errorf("cannot download file: %w", err) - } - - description, err := h.describer.Describe(ctx, file.FileName, mimeType, base64Data) - if err != nil { - return fmt.Errorf("cannot describe evidence: %w", err) - } - - return h.pg.WithTx( - ctx, - func(ctx context.Context, tx pg.Tx) 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 (h *evidenceDescriptionHandler) failEvidence( - ctx context.Context, - evidence *coredata.Evidence, -) error { - scope := coredata.NewScopeFromObjectID(evidence.ID) - - return h.pg.WithTx( - ctx, - func(ctx context.Context, tx pg.Tx) error { - evidence.DescriptionStatus = coredata.EvidenceDescriptionStatusFailed - 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 - }, - ) -} diff --git a/pkg/probo/evidence_service.go b/pkg/probo/evidence_service.go index ce345b364d..060de1e45b 100644 --- a/pkg/probo/evidence_service.go +++ b/pkg/probo/evidence_service.go @@ -91,14 +91,14 @@ func (s EvidenceService) UploadMeasureEvidence( } evidence := &coredata.Evidence{ - ID: evidenceID, - MeasureID: req.MeasureID, - State: coredata.EvidenceStateFulfilled, - ReferenceID: "custom-evidence-" + referenceID.String(), - Type: coredata.EvidenceTypeFile, - DescriptionStatus: coredata.EvidenceDescriptionStatusPending, - CreatedAt: now, - UpdatedAt: now, + ID: evidenceID, + MeasureID: req.MeasureID, + State: coredata.EvidenceStateFulfilled, + ReferenceID: "custom-evidence-" + referenceID.String(), + Type: coredata.EvidenceTypeFile, + AssessmentStatus: coredata.EvidenceAssessmentStatusPending, + CreatedAt: now, + UpdatedAt: now, } err = s.svc.pg.WithTx( @@ -130,7 +130,7 @@ func (s EvidenceService) UploadMeasureEvidence( } evidence.OrganizationID = measure.OrganizationID - evidence.EvidenceFileId = &file.ID + evidence.EvidenceFileID = &file.ID evidence.MeasureID = req.MeasureID if err := evidence.Insert(ctx, conn, scope); err != nil { diff --git a/pkg/probo/framework_service.go b/pkg/probo/framework_service.go index f949d6f740..3096a25515 100644 --- a/pkg/probo/framework_service.go +++ b/pkg/probo/framework_service.go @@ -248,12 +248,12 @@ func (s FrameworkService) Export( for _, evidence := range evidences { if evidence.Type != coredata.EvidenceTypeFile || evidence.State != coredata.EvidenceStateFulfilled || - evidence.EvidenceFileId == nil { + evidence.EvidenceFileID == nil { continue } evidence_file := &coredata.File{} - if err := evidence_file.LoadByID(ctx, conn, scope, *evidence.EvidenceFileId); err != nil { + if err := evidence_file.LoadByID(ctx, conn, scope, *evidence.EvidenceFileID); err != nil { return fmt.Errorf("cannot load evidence file: %w", err) } diff --git a/pkg/probo/measure_service.go b/pkg/probo/measure_service.go index 8c75ede4b1..bf7f64d61b 100644 --- a/pkg/probo/measure_service.go +++ b/pkg/probo/measure_service.go @@ -426,15 +426,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, - DescriptionStatus: coredata.EvidenceDescriptionStatusPending, - 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, + AssessmentStatus: coredata.EvidenceAssessmentStatusPending, + CreatedAt: now, + UpdatedAt: now, } if err := evidence.Upsert(ctx, tx, scope); err != nil { diff --git a/pkg/probod/aliases.go b/pkg/probod/aliases.go index d2a30d28df..c1f14994d4 100644 --- a/pkg/probod/aliases.go +++ b/pkg/probod/aliases.go @@ -39,7 +39,7 @@ type ( ACMEConfig = probodconfig.ACMEConfig LLMProviderConfig = probodconfig.LLMProviderConfig LLMAgentConfig = probodconfig.LLMAgentConfig - EvidenceDescriberConfig = probodconfig.EvidenceDescriberConfig + EvidenceAssessmentConfig = probodconfig.EvidenceAssessmentConfig ThirdPartyVettingWorkerConfig = probodconfig.ThirdPartyVettingWorkerConfig AgentsConfig = probodconfig.AgentsConfig diff --git a/pkg/probod/probod.go b/pkg/probod/probod.go index ad90f24a4c..7264f7aa8f 100644 --- a/pkg/probod/probod.go +++ b/pkg/probod/probod.go @@ -55,7 +55,7 @@ import ( "go.probo.inc/probo/pkg/crypto/passwdhash" pemutil "go.probo.inc/probo/pkg/crypto/pem" "go.probo.inc/probo/pkg/esign" - "go.probo.inc/probo/pkg/evidencedescriber" + "go.probo.inc/probo/pkg/evidenceassessor" "go.probo.inc/probo/pkg/file" "go.probo.inc/probo/pkg/filemanager" "go.probo.inc/probo/pkg/filesign" @@ -173,7 +173,7 @@ func New() *Implm { TSAURL: "http://timestamp.digicert.com", }, Branding: true, - EvidenceDescriber: EvidenceDescriberConfig{ + EvidenceAssessor: EvidenceAssessmentConfig{ Interval: 10, StaleAfter: 300, MaxConcurrency: 10, @@ -310,7 +310,7 @@ func (impl *Implm) Run( return err } - evidenceDescriberAgentCfg, evidenceDescriberLLMClient, err := impl.resolveAgentClient("evidence-describer", impl.cfg.Agents.EvidenceDescriber, l, tp, r) + evidenceAssessorAgentCfg, evidenceAssessorLLMClient, err := impl.resolveAgentClient("evidence-assessor", impl.cfg.Agents.EvidenceAssessor, l, tp, r) if err != nil { return err } @@ -632,6 +632,25 @@ func (impl *Implm) Run( return fmt.Errorf("cannot create server: %w", err) } + // Build the evidence assessor before any worker cancel func is created + // so its fallible construction can return early without tripping govet's + // lostcancel on the context cancels defined further below. + evidenceAssessorCfg := evidenceassessor.Config{ + Client: evidenceAssessorLLMClient, + Model: evidenceAssessorAgentCfg.ModelName, + Temp: ref.UnrefOrZero(evidenceAssessorAgentCfg.Temperature), + MaxTokens: ref.UnrefOrZero(evidenceAssessorAgentCfg.MaxTokens), + Logger: l.Named("evidence-assessor"), + } + if evidenceAssessorAgentCfg.Thinking != nil { + evidenceAssessorCfg.Thinking = *evidenceAssessorAgentCfg.Thinking + } + + evidenceAssessor, err := evidenceassessor.New(evidenceAssessorCfg) + if err != nil { + return fmt.Errorf("cannot build evidence assessor: %w", err) + } + apiServerCtx, stopApiServer := context.WithCancel(context.Background()) defer stopApiServer() @@ -838,31 +857,23 @@ func (impl *Implm) Run( }, ) - evidenceDescriber := evidencedescriber.New( - evidenceDescriberLLMClient, - evidencedescriber.Config{ - Model: evidenceDescriberAgentCfg.ModelName, - Temp: ref.UnrefOrZero(evidenceDescriberAgentCfg.Temperature), - MaxTokens: ref.UnrefOrZero(evidenceDescriberAgentCfg.MaxTokens), - }, - ) - evidenceDescriptionWorker := probo.NewEvidenceDescriptionWorker( + evidenceAssessmentWorker := probo.NewEvidenceAssessmentWorker( pgClient, fileManagerService, - evidenceDescriber, - l.Named("evidence-description-worker"), - probo.EvidenceDescriptionWorkerConfig{ - StaleAfter: time.Duration(impl.cfg.EvidenceDescriber.StaleAfter) * time.Second, + evidenceAssessor, + l.Named("evidence-assessment-worker"), + probo.EvidenceAssessmentWorkerConfig{ + StaleAfter: time.Duration(impl.cfg.EvidenceAssessor.StaleAfter) * time.Second, }, - worker.WithInterval(time.Duration(impl.cfg.EvidenceDescriber.Interval)*time.Second), - worker.WithMaxConcurrency(impl.cfg.EvidenceDescriber.MaxConcurrency), + worker.WithInterval(time.Duration(impl.cfg.EvidenceAssessor.Interval)*time.Second), + worker.WithMaxConcurrency(impl.cfg.EvidenceAssessor.MaxConcurrency), ) - evidenceDescriptionWorkerCtx, stopEvidenceDescriptionWorker := context.WithCancel(context.Background()) + evidenceAssessmentWorkerCtx, stopEvidenceAssessmentWorker := context.WithCancel(context.Background()) wg.Go( func() { - if err := evidenceDescriptionWorker.Run(evidenceDescriptionWorkerCtx); err != nil { - cancel(fmt.Errorf("evidence description worker crashed: %w", err)) + if err := evidenceAssessmentWorker.Run(evidenceAssessmentWorkerCtx); err != nil { + cancel(fmt.Errorf("evidence assessment worker crashed: %w", err)) } }, ) @@ -920,7 +931,7 @@ func (impl *Implm) Run( stopCommonPatternEnrichmentWorker() stopMailingListWorker() stopVettingWorker() - stopEvidenceDescriptionWorker() + stopEvidenceAssessmentWorker() stopDocumentPDFWorker() stopExportJobExporter() stopAccessReviewWorker() diff --git a/pkg/probod/third_party_vetter.go b/pkg/probod/third_party_vetter.go index ba041676a6..85f5c43a00 100644 --- a/pkg/probod/third_party_vetter.go +++ b/pkg/probod/third_party_vetter.go @@ -24,7 +24,7 @@ import ( // buildThirdPartyVetter wires the third-party vetting agent. Unset // third-party-vetter fields inherit from the default agent config -// (AGENT_DEFAULT_*), same as evidence-describer and probo. +// (AGENT_DEFAULT_*), same as evidence-assessor and probo. func (impl *Implm) buildThirdPartyVetter( l *log.Logger, tp trace.TracerProvider, diff --git a/pkg/probodconfig/config.go b/pkg/probodconfig/config.go index 2334d8d6fa..6cb9d7545d 100644 --- a/pkg/probodconfig/config.go +++ b/pkg/probodconfig/config.go @@ -59,7 +59,7 @@ type ( Notifications NotificationsConfig `json:"notifications"` Connectors []ConnectorConfig `json:"connectors"` Agents AgentsConfig `json:"llm"` - EvidenceDescriber EvidenceDescriberConfig `json:"evidence-describer"` + EvidenceAssessor EvidenceAssessmentConfig `json:"evidence-assessor"` ThirdPartyVetting ThirdPartyVettingWorkerConfig `json:"third-party-vetting-worker"` TrackerMappingWorker TrackerMappingWorkerConfig `json:"tracker-mapping-worker"` diff --git a/pkg/probodconfig/llm_config.go b/pkg/probodconfig/llm_config.go index 2dda1df55b..9969d0cf7c 100644 --- a/pkg/probodconfig/llm_config.go +++ b/pkg/probodconfig/llm_config.go @@ -29,12 +29,17 @@ type ( ModelName string `json:"model-name"` Temperature *float64 `json:"temperature"` MaxTokens *int `json:"max-tokens"` + // Thinking is the extended-thinking budget in tokens for agents + // that opt in. Leave nil to disable extended thinking; set to 0 + // to explicitly disable via config. Only a few providers and + // models support this; see pkg/agent/WithThinking. + Thinking *int `json:"thinking"` } - // EvidenceDescriberConfig holds worker-side tuning for the evidence - // description background worker. LLM parameters for the same worker - // live under AgentsConfig.EvidenceDescriber. - EvidenceDescriberConfig struct { + // EvidenceAssessmentConfig holds worker-side tuning for the evidence + // assessment background worker. LLM parameters for the same worker + // live under AgentsConfig.EvidenceAssessor. + EvidenceAssessmentConfig struct { Interval int `json:"interval"` // seconds between polls StaleAfter int `json:"stale-after"` // seconds before a claim is recycled MaxConcurrency int `json:"max-concurrency"` @@ -84,13 +89,13 @@ type ( // 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:"defaults"` - Probo LLMAgentConfig `json:"probo"` - EvidenceDescriber LLMAgentConfig `json:"evidence-describer"` - ThirdPartyVetter LLMAgentConfig `json:"third-party-vetter"` - TrackerMapping LLMAgentConfig `json:"tracker-mapping"` - Tools AgentToolsConfig `json:"tools"` + Providers map[string]LLMProviderConfig `json:"providers"` + Default LLMAgentConfig `json:"defaults"` + Probo LLMAgentConfig `json:"probo"` + EvidenceAssessor LLMAgentConfig `json:"evidence-assessor"` + ThirdPartyVetter LLMAgentConfig `json:"third-party-vetter"` + TrackerMapping LLMAgentConfig `json:"tracker-mapping"` + Tools AgentToolsConfig `json:"tools"` } ) @@ -109,5 +114,9 @@ func (c *AgentsConfig) ResolveAgent(agent LLMAgentConfig) LLMAgentConfig { agent.MaxTokens = new(*c.Default.MaxTokens) } + if agent.Thinking == nil { + agent.Thinking = c.Default.Thinking + } + return agent } diff --git a/pkg/server/api/console/v1/types/evidence.go b/pkg/server/api/console/v1/types/evidence.go index ea4c8f9842..20e4396238 100644 --- a/pkg/server/api/console/v1/types/evidence.go +++ b/pkg/server/api/console/v1/types/evidence.go @@ -81,9 +81,9 @@ func NewEvidence(e *coredata.Evidence) *Evidence { UpdatedAt: e.UpdatedAt, } - if e.EvidenceFileId != nil { + if e.EvidenceFileID != nil { evidence.File = &File{ - ID: *e.EvidenceFileId, + ID: *e.EvidenceFileID, } } diff --git a/pkg/vetting/assessment.go b/pkg/vetting/assessment.go index a9413f06f2..f04e113b1c 100644 --- a/pkg/vetting/assessment.go +++ b/pkg/vetting/assessment.go @@ -314,35 +314,15 @@ func thirdPartyInfoOutputType() (*agent.OutputType, error) { return nil, fmt.Errorf("cannot create thirdParty info output type: %w", err) } - var schema map[string]any - if err := json.Unmarshal(outputType.Schema, &schema); err != nil { - return nil, fmt.Errorf("cannot unmarshal thirdParty info schema: %w", err) + if err := outputType.DecorateEnum("category", thirdPartyCategoryEnum); err != nil { + return nil, fmt.Errorf("cannot decorate thirdParty info schema: %w", err) } - properties, ok := schema["properties"].(map[string]any) - if !ok { - return nil, fmt.Errorf("thirdParty info schema has no properties") + if err := outputType.DecorateEnum("third_party_type", thirdPartyTypeEnum); err != nil { + return nil, fmt.Errorf("cannot decorate thirdParty info schema: %w", err) } - enums := map[string][]string{ - "category": thirdPartyCategoryEnum, - "third_party_type": thirdPartyTypeEnum, - } - for field, values := range enums { - prop, ok := properties[field].(map[string]any) - if !ok { - return nil, fmt.Errorf("thirdParty info schema has no %q property", field) - } - - prop["enum"] = values - } - - decorated, err := json.Marshal(schema) - if err != nil { - return nil, fmt.Errorf("cannot marshal decorated thirdParty info schema: %w", err) - } - - strict, err := enforceStrictJSONSchema(decorated) + strict, err := enforceStrictJSONSchema(outputType.Schema) if err != nil { return nil, fmt.Errorf("cannot enforce strict thirdParty info schema: %w", err) } diff --git a/pkg/vetting/openai_schema.go b/pkg/vetting/openai_schema.go index a27046d0bc..8dae55c79e 100644 --- a/pkg/vetting/openai_schema.go +++ b/pkg/vetting/openai_schema.go @@ -34,8 +34,13 @@ type strictFunctionTool[P any] struct { } // jsonSchemaForTool builds an OpenAI strict-mode JSON schema for vetting tools -// and structured outputs. OpenAI requires every property in required and -// additionalProperties=false; the shared agent schema generator does not. +// and structured outputs. The shared agent schema generator already emits +// additionalProperties=false and a complete required list on every struct and +// array-of-struct object, so enforceStrictJSONSchema only adds the remaining +// strict-mode normalisation the generator skips: it sorts required and locks +// down map[string]T fields (which serialise to a non-false additionalProperties). +// Output types that are flat or map-free (e.g. evidenceassessor.EvidenceAssessment) +// are already strict-valid straight from agent.NewOutputType and do not need this. func jsonSchemaForTool[T any]() (json.RawMessage, error) { outputType, err := agent.NewOutputType[T]("_") if err != nil {