| 1 | from helpers.tool import Tool, Response |
| 2 | from helpers import parallel_tools |
| 3 | from helpers.strings import sanitize_string |
| 4 | |
| 5 | |
| 6 | class ParallelTool(Tool): |
| 7 | async def before_execution(self, **kwargs): |
| 8 | self.log = None |
| 9 | |
| 10 | async def after_execution(self, response: Response, **kwargs): |
| 11 | text = sanitize_string(response.message.strip()) |
| 12 | self.agent.hist_add_tool_result( |
| 13 | self.name, |
| 14 | text, |
| 15 | **(response.additional or {}), |
| 16 | ) |
| 17 | await parallel_tools.collect_parallel_jobs( |
| 18 | self.agent, |
| 19 | getattr(self, "_collect_job_ids", []), |
| 20 | promote_parent_history=True, |
| 21 | ) |
| 22 | |
| 23 | async def execute(self, **kwargs) -> Response: |
| 24 | self._collect_job_ids = [] |
| 25 | args = {**self.args, **kwargs} |
| 26 | action = str(args.get("action") or "").strip().lower() |
| 27 | |
| 28 | try: |
| 29 | timeout = parallel_tools.coerce_timeout(args.get("timeout")) |
| 30 | job_ids = parallel_tools.normalize_job_ids(args.get("job_ids")) |
| 31 | |
| 32 | if action == "cancel": |
| 33 | results = await parallel_tools.cancel_parallel_jobs(self.agent, job_ids) |
| 34 | return Response( |
| 35 | message=parallel_tools.format_parallel_results(results), |
| 36 | break_loop=False, |
| 37 | ) |
| 38 | |
| 39 | raw_calls = parallel_tools.extract_tool_calls(args) |
| 40 | started_jobs = [] |
| 41 | if raw_calls is not None: |
| 42 | calls = parallel_tools.normalize_parallel_tool_calls(raw_calls) |
| 43 | started_jobs = await parallel_tools.start_parallel_jobs(self.agent, calls) |
| 44 | |
| 45 | started_job_ids = [job.id for job in started_jobs] |
| 46 | all_job_ids = [*job_ids, *started_job_ids] |
| 47 | |
| 48 | if not all_job_ids: |
| 49 | return Response( |
| 50 | message=( |
| 51 | "Error: provide `tool_calls` to start parallel jobs, " |
| 52 | "or `job_ids` to await/cancel existing jobs." |
| 53 | ), |
| 54 | break_loop=False, |
| 55 | ) |
| 56 | |
| 57 | wait_default = action not in {"start", "background", "collect"} |
| 58 | wait = parallel_tools.coerce_bool(args.get("wait"), wait_default) |
| 59 | if action in {"await", "wait"}: |
| 60 | wait = True |
| 61 | |
| 62 | if not wait: |
| 63 | if not started_jobs and not job_ids: |
| 64 | return Response( |
| 65 | message="Error: `wait: false` requires `tool_calls` to start new jobs.", |
| 66 | break_loop=False, |
| 67 | ) |
| 68 | if not job_ids: |
| 69 | return Response( |
| 70 | message=parallel_tools.format_started_jobs(started_jobs), |
| 71 | break_loop=False, |
| 72 | ) |
| 73 | results = await parallel_tools.await_parallel_jobs( |
| 74 | self.agent, |
| 75 | all_job_ids, |
| 76 | timeout=timeout, |
| 77 | collect=False, |
| 78 | wait=False, |
| 79 | ) |
| 80 | self._collect_job_ids = [result["job_id"] for result in results] |
| 81 | return Response( |
| 82 | message=parallel_tools.format_parallel_results(results), |
| 83 | break_loop=False, |
| 84 | ) |
| 85 | |
| 86 | results = await parallel_tools.await_parallel_jobs( |
| 87 | self.agent, |
| 88 | all_job_ids, |
| 89 | timeout=timeout, |
| 90 | collect=False, |
| 91 | wait=True, |
| 92 | ) |
| 93 | self._collect_job_ids = [result["job_id"] for result in results] |
| 94 | return Response( |
| 95 | message=parallel_tools.format_parallel_results(results), |
| 96 | break_loop=False, |
| 97 | ) |
| 98 | except ValueError as exc: |
| 99 | return Response(message=f"Error: {exc}", break_loop=False) |