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
« prev ^ index » next coverage.py v7.11.3, created at 2025-11-14 11:27 +0100
1"""
2Alprina Agent Workflows - Structured orchestration patterns.
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
11Reference: https://ai-sdk.dev/docs/agents/workflows
12"""
14from typing import Dict, Any, List, Optional, Callable
15from enum import Enum
16from loguru import logger
17from datetime import datetime
18import asyncio
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
30class WorkflowResult:
31 """Result from workflow execution."""
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 = {}
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 })
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
60 def complete(self, final_output: Any):
61 """Mark workflow as complete."""
62 self.end_time = datetime.now()
63 self.final_output = final_output
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
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 }
83class AlprinaWorkflow:
84 """
85 Base workflow orchestrator for Alprina agents.
87 Implements the Orchestrator-Worker pattern where the Main Alprina Agent
88 coordinates specialized security agents through structured workflows.
89 """
91 def __init__(self, workflow_type: WorkflowType = WorkflowType.ORCHESTRATOR_WORKER):
92 """
93 Initialize workflow.
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")
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).
110 Steps executed in order: A → B → C
111 Each step's output becomes next step's input.
113 Args:
114 steps: List of steps [{agent, task, params}]
115 context: Shared context across steps
117 Returns:
118 WorkflowResult with final output
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")
129 current_output = context or {}
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')}")
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 )
142 self.result.add_step(
143 step_name=step["task"],
144 output=step_output,
145 agent=step["agent"]
146 )
148 # Output becomes input for next step
149 current_output = step_output
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
156 self.result.complete(current_output)
157 return self.result
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.
168 Router intelligently selects best agent based on request.
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
175 Returns:
176 WorkflowResult with selected agent's output
178 Example:
179 router_fn decides: "scan code" → CodeAgent
180 "check API" → Web Scanner Agent
181 """
182 logger.info("Starting routing workflow")
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']}")
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 )
196 self.result.add_step(
197 step_name="routing_decision",
198 output={"selected_agent": selected_agent["name"]},
199 agent="router"
200 )
202 self.result.add_step(
203 step_name=selected_agent["task"],
204 output=output,
205 agent=selected_agent["id"]
206 )
208 self.result.complete(output)
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)
215 return self.result
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.
225 Multiple agents run simultaneously for efficiency.
227 Args:
228 tasks: List of independent tasks to run in parallel
229 context: Shared context
231 Returns:
232 WorkflowResult with all outputs
234 Example:
235 Scan code + Check secrets + Audit config simultaneously
236 """
237 logger.info(f"Starting parallel workflow with {len(tasks)} tasks")
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)
250 # Execute all in parallel
251 outputs = await asyncio.gather(*async_tasks, return_exceptions=True)
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 )
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 }
274 self.result.complete(final_output)
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)
281 return self.result
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.
293 Main Agent (orchestrator) coordinates specialized workers.
294 This is the PRIMARY pattern for Alprina!
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
302 Returns:
303 WorkflowResult with coordinated output
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")
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 )
320 self.result.add_step(
321 step_name="orchestration_planning",
322 output=plan,
323 agent=orchestrator_agent
324 )
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 )
336 self.result.add_step(
337 step_name=worker["task"],
338 output=result,
339 agent=worker["agent"]
340 )
342 worker_results.append({
343 "agent": worker["agent"],
344 "output": result
345 })
347 except Exception as e:
348 logger.error(f"Worker {worker['agent']} failed: {e}")
349 self.result.add_error(str(e), step=worker["task"])
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 )
361 self.result.add_step(
362 step_name="result_aggregation",
363 output=final_output,
364 agent=orchestrator_agent
365 )
367 self.result.complete(final_output)
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)
374 return self.result
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.
387 Quality control loop: Execute → Evaluate → Improve → Repeat
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
396 Returns:
397 WorkflowResult with optimized output
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)")
407 best_output = None
408 best_score = 0.0
410 for iteration in range(1, max_iterations + 1):
411 try:
412 logger.info(f"Iteration {iteration}/{max_iterations}")
414 # Execute task
415 output = await self._execute_step(
416 agent=agent,
417 task=task,
418 params=params
419 )
421 # Evaluate quality
422 evaluation = evaluator_fn(output)
423 score = evaluation.get("score", 0.0)
424 feedback = evaluation.get("feedback", "")
426 logger.info(f"Quality score: {score:.2f} - {feedback}")
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 )
438 # Track best result
439 if score > best_score:
440 best_score = score
441 best_output = output
443 # Check if quality acceptable
444 if evaluation.get("acceptable", False):
445 logger.info(f"Quality acceptable after {iteration} iterations")
446 break
448 # Optimize parameters for next iteration
449 if iteration < max_iterations:
450 params = self._optimize_params(params, evaluation)
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
457 self.result.metadata["iterations"] = iteration
458 self.result.metadata["best_score"] = best_score
459 self.result.complete(best_output)
461 return self.result
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.
472 Args:
473 agent: Agent identifier
474 task: Task to perform
475 params: Task parameters
477 Returns:
478 Step output
479 """
480 logger.debug(f"Executing: {agent}.{task}()")
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
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")
492 if not target:
493 return {"error": "No target specified", "agent": agent}
495 # Determine if local or remote scan
496 is_local = Path(target).exists() if target else False
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 )
517 return result
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")
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"}
533 return {"error": "Missing scan results or target", "agent": agent}
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 }
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 }
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
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 }
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.
585 Args:
586 params: Current parameters
587 evaluation: Evaluation result with feedback
589 Returns:
590 Optimized parameters
591 """
592 # Example optimization logic
593 optimized = params.copy()
595 feedback = evaluation.get("feedback", "").lower()
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
601 if "missed vulnerabilities" in feedback:
602 optimized["sensitivity"] = optimized.get("sensitivity", 0.5) + 0.1
604 return optimized
607# Quality evaluator functions
609def evaluate_scan_quality(output: Any) -> Dict[str, Any]:
610 """
611 Evaluate quality of scan results.
613 Checks for:
614 - False positive indicators
615 - Coverage completeness
616 - Confidence scores
617 - Result consistency
619 Args:
620 output: Scan results to evaluate
622 Returns:
623 Evaluation dict with score and feedback
624 """
625 findings = output.get("findings", []) if isinstance(output, dict) else []
627 if not findings:
628 return {
629 "score": 1.0,
630 "acceptable": True,
631 "feedback": "No findings - scan complete"
632 }
634 # Count findings by severity
635 severity_counts = {}
636 low_confidence_count = 0
637 total_findings = len(findings)
639 for finding in findings:
640 severity = finding.get("severity", "UNKNOWN")
641 severity_counts[severity] = severity_counts.get(severity, 0) + 1
643 # Check confidence if available
644 confidence = finding.get("confidence", 1.0)
645 if confidence < 0.6:
646 low_confidence_count += 1
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
652 # Acceptable if score >= 0.7 and not too many critical with low confidence
653 acceptable = quality_score >= 0.7
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")
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 }
674def evaluate_comprehensive_scan(output: Any) -> Dict[str, Any]:
675 """
676 Evaluate comprehensive security scan results from multiple agents.
678 Args:
679 output: Workflow results from parallel/orchestrator-worker execution
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 )
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}
696 coverage = len(agents_run & expected_agents) / len(expected_agents)
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 }
709 # Single scan result
710 return evaluate_scan_quality(output)
713# Convenience functions for common workflows
715async def comprehensive_security_scan(target: str) -> WorkflowResult:
716 """
717 Run comprehensive security scan using orchestrator-worker pattern.
719 Orchestrator: Main Alprina Agent
720 Workers: CodeAgent, Secret Detection, Config Audit
722 Args:
723 target: Path to scan
725 Returns:
726 WorkflowResult with all findings
727 """
728 workflow = AlprinaWorkflow(WorkflowType.ORCHESTRATOR_WORKER)
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 )
754async def parallel_multi_target_scan(targets: List[str]) -> WorkflowResult:
755 """
756 Scan multiple targets in parallel.
758 Args:
759 targets: List of paths/URLs to scan
761 Returns:
762 WorkflowResult with all scan results
763 """
764 workflow = AlprinaWorkflow(WorkflowType.PARALLEL)
766 tasks = [
767 {
768 "agent": "codeagent",
769 "task": "scan",
770 "params": {"target": target}
771 }
772 for target in targets
773 ]
775 return await workflow.execute_parallel(tasks)
778async def sequential_scan_and_report(target: str) -> WorkflowResult:
779 """
780 Sequential workflow: Scan → Analyze → Generate Report
782 Args:
783 target: Path to scan
785 Returns:
786 WorkflowResult with final report
787 """
788 workflow = AlprinaWorkflow(WorkflowType.SEQUENTIAL)
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 ]
808 return await workflow.execute_sequential(steps)