Holds mutable context shared across all workflow nodes during execution.
| 28 | |
| 29 | @dataclass(slots=True) |
| 30 | class WorkflowContext: |
| 31 | """Holds mutable context shared across all workflow nodes during execution.""" |
| 32 | benchmark_name: str |
| 33 | config: AppConfig | None = None |
| 34 | problem_template: str = DEFAULT_PROBLEM_TEMPLATE |
| 35 | problem: str = DEFAULT_PROBLEM_TEMPLATE |
| 36 | history: list[str] = field(default_factory=list) |
| 37 | commands: list[str] = field(default_factory=list) |
| 38 | benchmark: "BenchmarkItem | None" = None |
| 39 | ip_addr: str = "" |
| 40 | target: str = "" |
| 41 | |
| 42 | def bind_target(self, ip_addr: str, target: str) -> None: |
| 43 | self.ip_addr = ip_addr |
| 44 | self.target = target |
| 45 | self.problem = self.problem_template.format(ip_addr=ip_addr, vul_target=target) |
| 46 | benchmark_context = self._benchmark_context() |
| 47 | if benchmark_context: |
| 48 | self.problem += ( |
| 49 | "Known benchmark context:\n" |
| 50 | f"{benchmark_context}\n" |
| 51 | "Prefer this local benchmark context before fetching external pages.\n" |
| 52 | ) |
| 53 | |
| 54 | def append_information(self, information: str) -> None: |
| 55 | self.problem += f"Information: {information}\n" |
| 56 | |
| 57 | def max_iterations_for(self, sender: str) -> int: |
| 58 | workflow_config = self._workflow_config() |
| 59 | if sender == "Exploit": |
| 60 | return workflow_config.exp_iterations |
| 61 | if sender == "Inquire": |
| 62 | return workflow_config.query_iterations |
| 63 | return workflow_config.scan_iterations |
| 64 | |
| 65 | def recursion_limit(self) -> int: |
| 66 | return self._workflow_config().sys_iterations |
| 67 | |
| 68 | def debug_enabled(self) -> bool: |
| 69 | return self._workflow_config().debug |
| 70 | |
| 71 | def draw_graph_enabled(self) -> bool: |
| 72 | return self._workflow_config().draw_graph |
| 73 | |
| 74 | def reset_runtime(self) -> None: |
| 75 | self.problem = self.problem_template |
| 76 | self.history.clear() |
| 77 | self.commands.clear() |
| 78 | if self.ip_addr or self.target: |
| 79 | self.bind_target(self.ip_addr, self.target) |
| 80 | |
| 81 | def initial_state(self) -> WorkflowState: |
| 82 | try: |
| 83 | from langchain_core.messages import HumanMessage |
| 84 | |
| 85 | initial_message: Any = HumanMessage(content=self.problem) |
| 86 | except ImportError: |
| 87 | initial_message = self.problem |
no outgoing calls
no test coverage detected