Wrap a Python callable `(name, LlmRequest, next) -> dict` for LLM execution intercepts.
(
py_fn: Py<PyAny>,
)
| 462 | |
| 463 | /// Wrap a Python callable `(name, LlmRequest, next) -> dict` for LLM execution intercepts. |
| 464 | pub fn wrap_py_llm_exec_intercept_fn( |
| 465 | py_fn: Py<PyAny>, |
| 466 | ) -> Arc< |
| 467 | dyn Fn( |
| 468 | &str, |
| 469 | LlmRequest, |
| 470 | LlmExecutionNextFn, |
| 471 | ) -> Pin<Box<dyn Future<Output = FlowResult<Json>> + Send>> |
| 472 | + Send |
| 473 | + Sync, |
| 474 | > { |
| 475 | let py_fn = Arc::new(py_fn); |
| 476 | Arc::new( |
| 477 | move |name: &str, request: LlmRequest, next: LlmExecutionNextFn| { |
| 478 | let py_fn = py_fn.clone(); |
| 479 | let name = name.to_string(); |
| 480 | Box::pin(async move { |
| 481 | let outcome: FlowResult< |
| 482 | Result<Json, Pin<Box<dyn Future<Output = PyResult<Py<PyAny>>> + Send>>>, |
| 483 | > = Python::attach(|py| { |
| 484 | let py_req = PyLLMRequest { inner: request }; |
| 485 | let py_next = PyLlmNextFn { inner: next }; |
| 486 | let result = py_fn |
| 487 | .call1( |
| 488 | py, |
| 489 | ( |
| 490 | &name, |
| 491 | py_req |
| 492 | .into_pyobject(py) |
| 493 | .map_err(|e| FlowError::Internal(e.to_string()))? |
| 494 | .into_any(), |
| 495 | py_next |
| 496 | .into_pyobject(py) |
| 497 | .map_err(|e| FlowError::Internal(e.to_string()))? |
| 498 | .into_any(), |
| 499 | ), |
| 500 | ) |
| 501 | .map_err(|e: PyErr| FlowError::Internal(e.to_string()))?; |
| 502 | |
| 503 | let bound = result.bind(py); |
| 504 | if bound.getattr("__await__").is_ok() { |
| 505 | let future = pyo3_async_runtimes::tokio::into_future(result.into_bound(py)) |
| 506 | .map_err(|e| FlowError::Internal(e.to_string()))?; |
| 507 | Ok(Err(Box::pin(future) |
| 508 | as Pin< |
| 509 | Box<dyn Future<Output = PyResult<Py<PyAny>>> + Send>, |
| 510 | >)) |
| 511 | } else { |
| 512 | let json = py_to_json(bound) |
| 513 | .map_err(|e: PyErr| FlowError::Internal(e.to_string()))?; |
| 514 | Ok(Ok(json)) |
| 515 | } |
| 516 | }); |
| 517 | |
| 518 | match outcome? { |
| 519 | Ok(json) => Ok(json), |
| 520 | Err(future) => { |
| 521 | let py_result = future |