diff --git a/agent.py b/agent.py index d1607d0..d30ef29 100644 --- a/agent.py +++ b/agent.py @@ -213,10 +213,6 @@ class GenerateRemediationRequest(StrictModel): AgentStatus = Literal["ok", "partial", "unreachable", "setup_required"] -class AgentComplianceAuditorResult(dict[str, Any]): - """Compact JSON-shaped result alias for a smaller Agent Card output schema.""" - - class AgentComplianceAuditorConfig(StrictModel): @@ -366,7 +362,7 @@ class AgentComplianceAuditor(A2AAgent[AgentComplianceAuditorConfig, NoAuth]): return FindingResult(audit_id=audit_id, target=request.target, status="ok", output_dir=output_dir, generated_files=files, findings=findings, evidence=evidence, warnings=warnings) @a2a.tool(description="Deduplicate findings, map evidence to controls, and score severity, confidence, and aggregate risk", timeout_seconds=120, idempotent=True, cost_class="read-only") - async def assess_risk(self, ctx: RunContext[NoAuth], request: AssessRiskRequest) -> dict[str, Any]: + async def assess_risk(self, ctx: RunContext[NoAuth], request: AssessRiskRequest) -> RiskSummaryResult: audit_id = _audit_id(request.target, request.options.audit_id) await _event(ctx, "audit_started", audit_id=audit_id, skill="assess_risk", target=request.target.agent_name) all_findings = _dedupe_findings([*request.configuration_findings, *request.runtime_findings])