Coverage for src/alprina_cli/workflows.py: 0%

222 statements  

« prev     ^ index     » next       coverage.py v7.11.3, created at 2025-11-14 11:27 +0100

1""" 

2Alprina Agent Workflows - Structured orchestration patterns. 

3 

4Implements AI SDK workflow patterns for coordinating security agents: 

5- Sequential (Chain): Step-by-step security analysis 

6- Routing: Intelligent agent selection 

7- Parallel: Run multiple agents simultaneously 

8- Orchestrator-Worker: Main Agent coordinates specialized agents 

9- Evaluator-Optimizer: Quality control and error recovery 

10 

11Reference: https://ai-sdk.dev/docs/agents/workflows 

12""" 

13 

14from typing import Dict, Any, List, Optional, Callable 

15from enum import Enum 

16from loguru import logger 

17from datetime import datetime 

18import asyncio 

19 

20 

21class WorkflowType(Enum): 

22 """Types of agent workflows.""" 

23 SEQUENTIAL = "sequential" # Chain: A → B → C 

24 ROUTING = "routing" # Route to best agent 

25 PARALLEL = "parallel" # Run agents concurrently 

26 ORCHESTRATOR_WORKER = "orchestrator_worker" # Main coordinates workers 

27 EVALUATOR_OPTIMIZER = "evaluator_optimizer" # Quality control loop 

28 

29 

30class WorkflowResult: 

31 """Result from workflow execution.""" 

32 

33 def __init__(self, success: bool = True): 

34 self.success = success 

35 self.steps = [] 

36 self.errors = [] 

37 self.start_time = datetime.now() 

38 self.end_time = None 

39 self.final_output = None 

40 self.metadata = {} 

41 

42 def add_step(self, step_name: str, output: Any, agent: str = None): 

43 """Add a completed workflow step.""" 

44 self.steps.append({ 

45 "name": step_name, 

46 "agent": agent, 

47 "output": output, 

48 "timestamp": datetime.now().isoformat() 

49 }) 

50 

51 def add_error(self, error: str, step: str = None): 

52 """Add an error that occurred during workflow.""" 

53 self.errors.append({ 

54 "error": error, 

55 "step": step, 

56 "timestamp": datetime.now().isoformat() 

57 }) 

58 self.success = False 

59 

60 def complete(self, final_output: Any): 

61 """Mark workflow as complete.""" 

62 self.end_time = datetime.now() 

63 self.final_output = final_output 

64 

65 def duration(self) -> float: 

66 """Get workflow duration in seconds.""" 

67 if self.end_time: 

68 return (self.end_time - self.start_time).total_seconds() 

69 return 0.0 

70 

71 def to_dict(self) -> Dict[str, Any]: 

72 """Convert to dictionary.""" 

73 return { 

74 "success": self.success, 

75 "steps": self.steps, 

76 "errors": self.errors, 

77 "duration": self.duration(), 

78 "final_output": self.final_output, 

79 "metadata": self.metadata 

80 } 

81 

82 

83class AlprinaWorkflow: 

84 """ 

85 Base workflow orchestrator for Alprina agents. 

86 

87 Implements the Orchestrator-Worker pattern where the Main Alprina Agent 

88 coordinates specialized security agents through structured workflows. 

89 """ 

90 

91 def __init__(self, workflow_type: WorkflowType = WorkflowType.ORCHESTRATOR_WORKER): 

92 """ 

93 Initialize workflow. 

94 

95 Args: 

96 workflow_type: Type of workflow pattern to use 

97 """ 

98 self.workflow_type = workflow_type 

99 self.result = WorkflowResult() 

100 logger.info(f"Initialized {workflow_type.value} workflow") 

101 

102 async def execute_sequential( 

103 self, 

104 steps: List[Dict[str, Any]], 

105 context: Dict[str, Any] = None 

106 ) -> WorkflowResult: 

107 """ 

108 Execute sequential workflow (Chain pattern). 

109 

110 Steps executed in order: A → B → C 

111 Each step's output becomes next step's input. 

112 

113 Args: 

114 steps: List of steps [{agent, task, params}] 

115 context: Shared context across steps 

116 

117 Returns: 

118 WorkflowResult with final output 

119 

120 Example: 

121 steps = [ 

122 {"agent": "secret_detection", "task": "scan", "params": {"target": "./"}}, 

123 {"agent": "codeagent", "task": "analyze_secrets", "params": {}}, 

124 {"agent": "report_generator", "task": "create_report", "params": {}} 

125 ] 

126 """ 

127 logger.info(f"Starting sequential workflow with {len(steps)} steps") 

128 

129 current_output = context or {} 

130 

131 for i, step in enumerate(steps, 1): 

132 try: 

133 logger.info(f"Step {i}/{len(steps)}: {step.get('task')} via {step.get('agent')}") 

134 

135 # Execute step with previous output as input 

136 step_output = await self._execute_step( 

137 agent=step["agent"], 

138 task=step["task"], 

139 params={**step.get("params", {}), "input": current_output} 

140 ) 

141 

142 self.result.add_step( 

143 step_name=step["task"], 

144 output=step_output, 

145 agent=step["agent"] 

146 ) 

147 

148 # Output becomes input for next step 

149 current_output = step_output 

150 

151 except Exception as e: 

152 logger.error(f"Step {i} failed: {e}") 

153 self.result.add_error(str(e), step=step["task"]) 

154 break 

155 

156 self.result.complete(current_output) 

157 return self.result 

158 

159 async def execute_routing( 

160 self, 

161 user_request: str, 

162 available_agents: List[Dict[str, Any]], 

163 router_fn: Callable 

164 ) -> WorkflowResult: 

165 """ 

166 Execute routing workflow. 

167 

168 Router intelligently selects best agent based on request. 

169 

170 Args: 

171 user_request: User's natural language request 

172 available_agents: List of available agents 

173 router_fn: Function that routes request to agent 

174 

175 Returns: 

176 WorkflowResult with selected agent's output 

177 

178 Example: 

179 router_fn decides: "scan code" → CodeAgent 

180 "check API" → Web Scanner Agent 

181 """ 

182 logger.info("Starting routing workflow") 

183 

184 try: 

185 # Route to best agent 

186 selected_agent = router_fn(user_request, available_agents) 

187 logger.info(f"Routed to: {selected_agent['name']}") 

188 

189 # Execute with selected agent 

190 output = await self._execute_step( 

191 agent=selected_agent["id"], 

192 task=selected_agent["task"], 

193 params=selected_agent.get("params", {}) 

194 ) 

195 

196 self.result.add_step( 

197 step_name="routing_decision", 

198 output={"selected_agent": selected_agent["name"]}, 

199 agent="router" 

200 ) 

201 

202 self.result.add_step( 

203 step_name=selected_agent["task"], 

204 output=output, 

205 agent=selected_agent["id"] 

206 ) 

207 

208 self.result.complete(output) 

209 

210 except Exception as e: 

211 logger.error(f"Routing workflow failed: {e}") 

212 self.result.add_error(str(e)) 

213 self.result.complete(None) 

214 

215 return self.result 

216 

217 async def execute_parallel( 

218 self, 

219 tasks: List[Dict[str, Any]], 

220 context: Dict[str, Any] = None 

221 ) -> WorkflowResult: 

222 """ 

223 Execute parallel workflow. 

224 

225 Multiple agents run simultaneously for efficiency. 

226 

227 Args: 

228 tasks: List of independent tasks to run in parallel 

229 context: Shared context 

230 

231 Returns: 

232 WorkflowResult with all outputs 

233 

234 Example: 

235 Scan code + Check secrets + Audit config simultaneously 

236 """ 

237 logger.info(f"Starting parallel workflow with {len(tasks)} tasks") 

238 

239 try: 

240 # Create async tasks 

241 async_tasks = [] 

242 for task in tasks: 

243 async_task = self._execute_step( 

244 agent=task["agent"], 

245 task=task["task"], 

246 params={**task.get("params", {}), **(context or {})} 

247 ) 

248 async_tasks.append(async_task) 

249 

250 # Execute all in parallel 

251 outputs = await asyncio.gather(*async_tasks, return_exceptions=True) 

252 

253 # Process results 

254 for i, (task, output) in enumerate(zip(tasks, outputs)): 

255 if isinstance(output, Exception): 

256 logger.error(f"Task {i+1} failed: {output}") 

257 self.result.add_error(str(output), step=task["task"]) 

258 else: 

259 self.result.add_step( 

260 step_name=task["task"], 

261 output=output, 

262 agent=task["agent"] 

263 ) 

264 

265 # Aggregate all outputs 

266 final_output = { 

267 "parallel_results": [ 

268 {"task": t["task"], "output": o} 

269 for t, o in zip(tasks, outputs) 

270 if not isinstance(o, Exception) 

271 ] 

272 } 

273 

274 self.result.complete(final_output) 

275 

276 except Exception as e: 

277 logger.error(f"Parallel workflow failed: {e}") 

278 self.result.add_error(str(e)) 

279 self.result.complete(None) 

280 

281 return self.result 

282 

283 async def execute_orchestrator_worker( 

284 self, 

285 orchestrator_agent: str, 

286 workers: List[Dict[str, Any]], 

287 task: str, 

288 params: Dict[str, Any] 

289 ) -> WorkflowResult: 

290 """ 

291 Execute Orchestrator-Worker workflow. 

292 

293 Main Agent (orchestrator) coordinates specialized workers. 

294 This is the PRIMARY pattern for Alprina! 

295 

296 Args: 

297 orchestrator_agent: Main agent coordinating work 

298 workers: List of worker agents [{agent, task, params}] 

299 task: Overall task to accomplish 

300 params: Parameters for orchestration 

301 

302 Returns: 

303 WorkflowResult with coordinated output 

304 

305 Example: 

306 Orchestrator: Main Alprina Agent 

307 Workers: [CodeAgent, Web Scanner, Secret Detection] 

308 Task: "Comprehensive security scan" 

309 """ 

310 logger.info(f"Starting orchestrator-worker workflow: {orchestrator_agent} coordinating {len(workers)} workers") 

311 

312 try: 

313 # Step 1: Orchestrator plans work 

314 plan = await self._execute_step( 

315 agent=orchestrator_agent, 

316 task="plan_work", 

317 params={"task": task, "workers": workers, **params} 

318 ) 

319 

320 self.result.add_step( 

321 step_name="orchestration_planning", 

322 output=plan, 

323 agent=orchestrator_agent 

324 ) 

325 

326 # Step 2: Execute worker tasks 

327 worker_results = [] 

328 for worker in workers: 

329 try: 

330 result = await self._execute_step( 

331 agent=worker["agent"], 

332 task=worker["task"], 

333 params=worker.get("params", {}) 

334 ) 

335 

336 self.result.add_step( 

337 step_name=worker["task"], 

338 output=result, 

339 agent=worker["agent"] 

340 ) 

341 

342 worker_results.append({ 

343 "agent": worker["agent"], 

344 "output": result 

345 }) 

346 

347 except Exception as e: 

348 logger.error(f"Worker {worker['agent']} failed: {e}") 

349 self.result.add_error(str(e), step=worker["task"]) 

350 

351 # Step 3: Orchestrator aggregates results 

352 final_output = await self._execute_step( 

353 agent=orchestrator_agent, 

354 task="aggregate_results", 

355 params={ 

356 "worker_results": worker_results, 

357 "original_task": task 

358 } 

359 ) 

360 

361 self.result.add_step( 

362 step_name="result_aggregation", 

363 output=final_output, 

364 agent=orchestrator_agent 

365 ) 

366 

367 self.result.complete(final_output) 

368 

369 except Exception as e: 

370 logger.error(f"Orchestrator-worker workflow failed: {e}") 

371 self.result.add_error(str(e)) 

372 self.result.complete(None) 

373 

374 return self.result 

375 

376 async def execute_evaluator_optimizer( 

377 self, 

378 agent: str, 

379 task: str, 

380 params: Dict[str, Any], 

381 evaluator_fn: Callable, 

382 max_iterations: int = 3 

383 ) -> WorkflowResult: 

384 """ 

385 Execute Evaluator-Optimizer workflow. 

386 

387 Quality control loop: Execute → Evaluate → Improve → Repeat 

388 

389 Args: 

390 agent: Agent to execute task 

391 task: Task to perform 

392 params: Task parameters 

393 evaluator_fn: Function to evaluate output quality 

394 max_iterations: Maximum improvement iterations 

395 

396 Returns: 

397 WorkflowResult with optimized output 

398 

399 Example: 

400 1. CodeAgent scans code 

401 2. Evaluator checks for false positives 

402 3. If quality < threshold, re-scan with refined params 

403 4. Repeat until quality acceptable or max iterations 

404 """ 

405 logger.info(f"Starting evaluator-optimizer workflow (max {max_iterations} iterations)") 

406 

407 best_output = None 

408 best_score = 0.0 

409 

410 for iteration in range(1, max_iterations + 1): 

411 try: 

412 logger.info(f"Iteration {iteration}/{max_iterations}") 

413 

414 # Execute task 

415 output = await self._execute_step( 

416 agent=agent, 

417 task=task, 

418 params=params 

419 ) 

420 

421 # Evaluate quality 

422 evaluation = evaluator_fn(output) 

423 score = evaluation.get("score", 0.0) 

424 feedback = evaluation.get("feedback", "") 

425 

426 logger.info(f"Quality score: {score:.2f} - {feedback}") 

427 

428 self.result.add_step( 

429 step_name=f"iteration_{iteration}", 

430 output={ 

431 "result": output, 

432 "score": score, 

433 "feedback": feedback 

434 }, 

435 agent=agent 

436 ) 

437 

438 # Track best result 

439 if score > best_score: 

440 best_score = score 

441 best_output = output 

442 

443 # Check if quality acceptable 

444 if evaluation.get("acceptable", False): 

445 logger.info(f"Quality acceptable after {iteration} iterations") 

446 break 

447 

448 # Optimize parameters for next iteration 

449 if iteration < max_iterations: 

450 params = self._optimize_params(params, evaluation) 

451 

452 except Exception as e: 

453 logger.error(f"Iteration {iteration} failed: {e}") 

454 self.result.add_error(str(e), step=f"iteration_{iteration}") 

455 break 

456 

457 self.result.metadata["iterations"] = iteration 

458 self.result.metadata["best_score"] = best_score 

459 self.result.complete(best_output) 

460 

461 return self.result 

462 

463 async def _execute_step( 

464 self, 

465 agent: str, 

466 task: str, 

467 params: Dict[str, Any] 

468 ) -> Any: 

469 """ 

470 Execute a single workflow step with actual Alprina agents. 

471 

472 Args: 

473 agent: Agent identifier 

474 task: Task to perform 

475 params: Task parameters 

476 

477 Returns: 

478 Step output 

479 """ 

480 logger.debug(f"Executing: {agent}.{task}()") 

481 

482 # Import here to avoid circular dependencies 

483 from .security_engine import run_agent, run_local_scan, run_remote_scan 

484 from .report_generator import generate_security_reports 

485 from pathlib import Path 

486 

487 try: 

488 # Route to appropriate Alprina agent 

489 if task in ["scan", "code_audit", "web_recon", "vuln_scan", "secret_detection", "config_audit"]: 

490 target = params.get("target") 

491 

492 if not target: 

493 return {"error": "No target specified", "agent": agent} 

494 

495 # Determine if local or remote scan 

496 is_local = Path(target).exists() if target else False 

497 

498 if is_local: 

499 # Local file/directory scan 

500 result = await asyncio.get_event_loop().run_in_executor( 

501 None, 

502 run_local_scan, 

503 target, 

504 task if task != "scan" else "code-audit", 

505 params.get("safe_only", True) 

506 ) 

507 else: 

508 # Remote URL/IP scan 

509 result = await asyncio.get_event_loop().run_in_executor( 

510 None, 

511 run_remote_scan, 

512 target, 

513 task if task != "scan" else "web-recon", 

514 params.get("safe_only", True) 

515 ) 

516 

517 return result 

518 

519 elif task == "create_reports" or task == "generate_reports": 

520 # Generate markdown reports 

521 scan_results = params.get("input") or params.get("results") 

522 target = params.get("target") 

523 

524 if scan_results and target: 

525 report_path = await asyncio.get_event_loop().run_in_executor( 

526 None, 

527 generate_security_reports, 

528 scan_results, 

529 target 

530 ) 

531 return {"report_path": report_path, "status": "success"} 

532 

533 return {"error": "Missing scan results or target", "agent": agent} 

534 

535 elif task == "plan_work": 

536 # Orchestrator planning step 

537 return { 

538 "plan": "analyzed_request", 

539 "workers_assigned": params.get("workers", []), 

540 "status": "planned" 

541 } 

542 

543 elif task == "aggregate_results": 

544 # Orchestrator aggregation step 

545 worker_results = params.get("worker_results", []) 

546 return { 

547 "aggregated": True, 

548 "total_findings": sum( 

549 len(r.get("output", {}).get("findings", [])) 

550 for r in worker_results 

551 ), 

552 "worker_results": worker_results, 

553 "status": "aggregated" 

554 } 

555 

556 else: 

557 # Generic agent execution 

558 logger.warning(f"Unknown task type: {task}, using run_agent fallback") 

559 result = await asyncio.get_event_loop().run_in_executor( 

560 None, 

561 run_agent, 

562 task, 

563 params.get("input_data", ""), 

564 params.get("metadata", {}) 

565 ) 

566 return result 

567 

568 except Exception as e: 

569 logger.error(f"Step execution failed: {e}", exc_info=True) 

570 return { 

571 "error": str(e), 

572 "agent": agent, 

573 "task": task, 

574 "status": "failed" 

575 } 

576 

577 def _optimize_params( 

578 self, 

579 params: Dict[str, Any], 

580 evaluation: Dict[str, Any] 

581 ) -> Dict[str, Any]: 

582 """ 

583 Optimize parameters based on evaluation feedback. 

584 

585 Args: 

586 params: Current parameters 

587 evaluation: Evaluation result with feedback 

588 

589 Returns: 

590 Optimized parameters 

591 """ 

592 # Example optimization logic 

593 optimized = params.copy() 

594 

595 feedback = evaluation.get("feedback", "").lower() 

596 

597 # Adjust based on feedback 

598 if "too many false positives" in feedback: 

599 optimized["confidence_threshold"] = optimized.get("confidence_threshold", 0.5) + 0.1 

600 

601 if "missed vulnerabilities" in feedback: 

602 optimized["sensitivity"] = optimized.get("sensitivity", 0.5) + 0.1 

603 

604 return optimized 

605 

606 

607# Quality evaluator functions 

608 

609def evaluate_scan_quality(output: Any) -> Dict[str, Any]: 

610 """ 

611 Evaluate quality of scan results. 

612 

613 Checks for: 

614 - False positive indicators 

615 - Coverage completeness 

616 - Confidence scores 

617 - Result consistency 

618 

619 Args: 

620 output: Scan results to evaluate 

621 

622 Returns: 

623 Evaluation dict with score and feedback 

624 """ 

625 findings = output.get("findings", []) if isinstance(output, dict) else [] 

626 

627 if not findings: 

628 return { 

629 "score": 1.0, 

630 "acceptable": True, 

631 "feedback": "No findings - scan complete" 

632 } 

633 

634 # Count findings by severity 

635 severity_counts = {} 

636 low_confidence_count = 0 

637 total_findings = len(findings) 

638 

639 for finding in findings: 

640 severity = finding.get("severity", "UNKNOWN") 

641 severity_counts[severity] = severity_counts.get(severity, 0) + 1 

642 

643 # Check confidence if available 

644 confidence = finding.get("confidence", 1.0) 

645 if confidence < 0.6: 

646 low_confidence_count += 1 

647 

648 # Calculate quality score 

649 false_positive_ratio = low_confidence_count / total_findings if total_findings > 0 else 0 

650 quality_score = 1.0 - (false_positive_ratio * 0.5) # Penalize for low confidence findings 

651 

652 # Acceptable if score >= 0.7 and not too many critical with low confidence 

653 acceptable = quality_score >= 0.7 

654 

655 feedback = [] 

656 if false_positive_ratio > 0.3: 

657 feedback.append(f"High false positive ratio ({false_positive_ratio:.1%})") 

658 if severity_counts.get("CRITICAL", 0) > 10: 

659 feedback.append("Unusually high critical findings - verify accuracy") 

660 

661 return { 

662 "score": quality_score, 

663 "acceptable": acceptable, 

664 "feedback": "; ".join(feedback) if feedback else "Quality acceptable", 

665 "metrics": { 

666 "total_findings": total_findings, 

667 "low_confidence": low_confidence_count, 

668 "false_positive_ratio": false_positive_ratio, 

669 "severity_distribution": severity_counts 

670 } 

671 } 

672 

673 

674def evaluate_comprehensive_scan(output: Any) -> Dict[str, Any]: 

675 """ 

676 Evaluate comprehensive security scan results from multiple agents. 

677 

678 Args: 

679 output: Workflow results from parallel/orchestrator-worker execution 

680 

681 Returns: 

682 Evaluation dict 

683 """ 

684 if isinstance(output, dict) and "parallel_results" in output: 

685 # Parallel execution results 

686 results = output["parallel_results"] 

687 total_findings = sum( 

688 len(r.get("output", {}).get("findings", [])) 

689 for r in results 

690 ) 

691 

692 # Check coverage - did all agents run? 

693 expected_agents = {"codeagent", "secret_detection", "config_audit"} 

694 agents_run = {r.get("task") for r in results} 

695 

696 coverage = len(agents_run & expected_agents) / len(expected_agents) 

697 

698 return { 

699 "score": coverage, 

700 "acceptable": coverage >= 0.66, # At least 2/3 agents ran 

701 "feedback": f"Coverage: {coverage:.1%} ({len(agents_run)}/{len(expected_agents)} agents)", 

702 "metrics": { 

703 "total_findings": total_findings, 

704 "agents_run": list(agents_run), 

705 "coverage": coverage 

706 } 

707 } 

708 

709 # Single scan result 

710 return evaluate_scan_quality(output) 

711 

712 

713# Convenience functions for common workflows 

714 

715async def comprehensive_security_scan(target: str) -> WorkflowResult: 

716 """ 

717 Run comprehensive security scan using orchestrator-worker pattern. 

718 

719 Orchestrator: Main Alprina Agent 

720 Workers: CodeAgent, Secret Detection, Config Audit 

721 

722 Args: 

723 target: Path to scan 

724 

725 Returns: 

726 WorkflowResult with all findings 

727 """ 

728 workflow = AlprinaWorkflow(WorkflowType.ORCHESTRATOR_WORKER) 

729 

730 return await workflow.execute_orchestrator_worker( 

731 orchestrator_agent="main_alprina_agent", 

732 workers=[ 

733 { 

734 "agent": "codeagent", 

735 "task": "code_audit", 

736 "params": {"target": target} 

737 }, 

738 { 

739 "agent": "secret_detection", 

740 "task": "find_secrets", 

741 "params": {"target": target} 

742 }, 

743 { 

744 "agent": "config_audit", 

745 "task": "audit_configs", 

746 "params": {"target": target} 

747 } 

748 ], 

749 task="comprehensive_scan", 

750 params={"target": target} 

751 ) 

752 

753 

754async def parallel_multi_target_scan(targets: List[str]) -> WorkflowResult: 

755 """ 

756 Scan multiple targets in parallel. 

757 

758 Args: 

759 targets: List of paths/URLs to scan 

760 

761 Returns: 

762 WorkflowResult with all scan results 

763 """ 

764 workflow = AlprinaWorkflow(WorkflowType.PARALLEL) 

765 

766 tasks = [ 

767 { 

768 "agent": "codeagent", 

769 "task": "scan", 

770 "params": {"target": target} 

771 } 

772 for target in targets 

773 ] 

774 

775 return await workflow.execute_parallel(tasks) 

776 

777 

778async def sequential_scan_and_report(target: str) -> WorkflowResult: 

779 """ 

780 Sequential workflow: Scan → Analyze → Generate Report 

781 

782 Args: 

783 target: Path to scan 

784 

785 Returns: 

786 WorkflowResult with final report 

787 """ 

788 workflow = AlprinaWorkflow(WorkflowType.SEQUENTIAL) 

789 

790 steps = [ 

791 { 

792 "agent": "codeagent", 

793 "task": "scan", 

794 "params": {"target": target} 

795 }, 

796 { 

797 "agent": "main_alprina_agent", 

798 "task": "analyze_findings", 

799 "params": {} 

800 }, 

801 { 

802 "agent": "report_generator", 

803 "task": "create_markdown_reports", 

804 "params": {} 

805 } 

806 ] 

807 

808 return await workflow.execute_sequential(steps)