| 1495 | /// without one it is the plain list of outcome dicts, unchanged. |
| 1496 | #[pyo3(signature = (specs, budget_tokens=None))] |
| 1497 | fn parallel( |
| 1498 | &self, |
| 1499 | py: Python<'_>, |
| 1500 | specs: Vec<Bound<'_, PyAny>>, |
| 1501 | budget_tokens: Option<u64>, |
| 1502 | ) -> PyResult<PyObject> { |
| 1503 | let rust_specs = specs |
| 1504 | .iter() |
| 1505 | .map(|s| py_to_step_spec(py, s)) |
| 1506 | .collect::<PyResult<Vec<_>>>()?; |
| 1507 | let session = self.inner.clone(); |
| 1508 | |
| 1509 | // No budget → unchanged behavior: a plain list of outcome dicts. |
| 1510 | let Some(limit) = budget_tokens else { |
| 1511 | let outcomes = py.allow_threads(move || { |
| 1512 | get_runtime().block_on(async move { |
| 1513 | let executor = session.agent_executor(); |
| 1514 | execute_steps_parallel(executor, rust_specs, None).await |
| 1515 | }) |
| 1516 | }); |
| 1517 | let items = outcomes |
| 1518 | .iter() |
| 1519 | .map(|o| step_outcome_to_py(py, o)) |
| 1520 | .collect::<PyResult<Vec<_>>>()?; |
| 1521 | return Ok(PyList::new(py, items)?.into_any().unbind()); |
| 1522 | }; |
| 1523 | |
| 1524 | // Budget → shared ledger across the fan-out; return {"outcomes", "budget"}. |
| 1525 | let (outcomes, snapshot) = py.allow_threads(move || { |
| 1526 | get_runtime().block_on(async move { |
| 1527 | let wf = session.workflow_with_token_budget(Some(limit)); |
| 1528 | let outcomes = wf.parallel(rust_specs).await; |
| 1529 | (outcomes, wf.budget_snapshot()) |
| 1530 | }) |
| 1531 | }); |
| 1532 | let outcomes_py = outcomes |
| 1533 | .iter() |
| 1534 | .map(|o| step_outcome_to_py(py, o)) |
| 1535 | .collect::<PyResult<Vec<_>>>()?; |
| 1536 | let budget = PyDict::new(py); |
| 1537 | budget.set_item( |
| 1538 | "consumed_tokens", |
| 1539 | snapshot.as_ref().map(|b| b.consumed_tokens).unwrap_or(0), |
| 1540 | )?; |
| 1541 | budget.set_item( |
| 1542 | "limit_tokens", |
| 1543 | snapshot |
| 1544 | .as_ref() |
| 1545 | .and_then(|b| b.limit_tokens) |
| 1546 | .or(Some(limit)), |
| 1547 | )?; |
| 1548 | let result = PyDict::new(py); |
| 1549 | result.set_item("outcomes", outcomes_py)?; |
| 1550 | result.set_item("budget", budget)?; |
| 1551 | Ok(result.into_any().unbind()) |
| 1552 | } |
| 1553 | |
| 1554 | /// Like `parallel`, but resumable: progress is journaled under |