diff --git a/crates/switchyard-py/src/libsy_bindings/run.rs b/crates/switchyard-py/src/libsy_bindings/run.rs index 0b227bbd..4953965f 100644 --- a/crates/switchyard-py/src/libsy_bindings/run.rs +++ b/crates/switchyard-py/src/libsy_bindings/run.rs @@ -1,18 +1,205 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Python execution for Rust-owned libsy algorithms. +//! Host-driven libsy run streams and their Python step variants. -use std::sync::Arc; +use std::sync::{Arc, Mutex}; -use libsy::{Algorithm, NoopAlgo, RandomAlgo}; +use futures_util::{stream, StreamExt}; +use libsy::{Algorithm, CallLlmRequest, NoopAlgo, RandomAlgo, Step, StepStream}; +use pyo3::exceptions::{PyRuntimeError, PyStopAsyncIteration}; use pyo3::prelude::*; +use switchyard_protocol::{Context, Decision, Request}; use super::protocol::PyDecision; use super::target::PyLlmTargetSet; use super::values::{PyContext, PyRequest, PyResponse}; use crate::errors::py_libsy_error; +#[pyclass(name = "LlmCall", module = "switchyard.libsy")] +pub(crate) struct PyLlmCall { + inner: Mutex>, + context: Context, + request: Request, + decision: Arc, +} + +impl PyLlmCall { + fn new(inner: CallLlmRequest) -> Self { + let context = inner.get_routed().ctx.clone(); + let request = inner.get_request().clone(); + let decision = inner.get_routed().decision.clone(); + Self { + inner: Mutex::new(Some(inner)), + context, + request, + decision, + } + } +} + +#[pymethods] +impl PyLlmCall { + #[getter] + fn context(&self, py: Python<'_>) -> PyResult> { + Py::new(py, PyContext::from_core(self.context.clone())) + } + + #[getter] + fn request(&self, py: Python<'_>) -> PyResult> { + Py::new(py, PyRequest::from_core(self.request.clone())) + } + + #[getter] + fn decision(&self, py: Python<'_>) -> PyResult> { + Py::new(py, PyDecision::new(self.decision.clone())) + } + + fn respond(&self, response: PyRef<'_, PyResponse>) -> PyResult<()> { + let mut call = self + .inner + .lock() + .map_err(|_| PyRuntimeError::new_err("LlmCall lock is poisoned"))?; + if call.is_none() { + return Err(PyRuntimeError::new_err( + "LlmCall has already been fulfilled", + )); + } + let response = response.take_core()?; + let Some(call) = call.take() else { + return Err(PyRuntimeError::new_err( + "LlmCall has already been fulfilled", + )); + }; + call.respond(Ok(response)).map_err(py_libsy_error) + } + + fn fail(&self, message: String) -> PyResult<()> { + let mut call = self + .inner + .lock() + .map_err(|_| PyRuntimeError::new_err("LlmCall lock is poisoned"))?; + let Some(call) = call.take() else { + return Err(PyRuntimeError::new_err( + "LlmCall has already been fulfilled", + )); + }; + call.respond(Err(std::io::Error::other(message).into())) + .map_err(py_libsy_error) + } + + #[getter] + fn is_pending(&self) -> PyResult { + self.inner + .lock() + .map(|call| call.is_some()) + .map_err(|_| PyRuntimeError::new_err("LlmCall lock is poisoned")) + } + + fn __repr__(&self) -> String { + format!( + "LlmCall(selected_model={:?})", + self.decision.selected_model() + ) + } +} + +/// Python's discriminated-union view of [`libsy::Step`]. +#[pyclass(name = "Step", module = "switchyard.libsy", frozen)] +pub(crate) enum PyStep { + CallLlm { call: Py }, + Decision { decision: Py }, + ReturnToAgent { response: Py }, +} + +fn step_to_python(py: Python<'_>, step: Step) -> PyResult> { + let step = match step { + Step::CallLlm(call) => PyStep::CallLlm { + call: Py::new(py, PyLlmCall::new(*call))?, + }, + Step::Decision(decision) => PyStep::Decision { + decision: Py::new(py, PyDecision::new(decision))?, + }, + Step::ReturnToAgent(response) => PyStep::ReturnToAgent { + response: Py::new(py, PyResponse::from_core(*response))?, + }, + }; + step.into_pyobject(py) + .map(|value| value.unbind().into_any()) +} + +#[pyclass(name = "RunStream", module = "switchyard.libsy")] +pub(crate) struct PyRunStream { + stream: Arc>>, + consumed: bool, +} + +impl PyRunStream { + fn new(algorithm: Arc, context: Context, request: Request) -> Self { + let stream = stream::once(async move { algorithm.run_stream(context, request) }).flatten(); + Self { + stream: Arc::new(tokio::sync::Mutex::new(Some(Box::pin(stream)))), + consumed: false, + } + } +} + +#[pymethods] +impl PyRunStream { + fn __aiter__(mut slf: PyRefMut<'_, Self>) -> PyResult> { + if slf.consumed { + return Err(PyRuntimeError::new_err( + "RunStream has already been consumed", + )); + } + slf.consumed = true; + let py = slf.py(); + Ok(slf.into_pyobject(py)?.unbind().into_any()) + } + + fn __anext__<'py>(&self, py: Python<'py>) -> PyResult> { + let stream = Arc::clone(&self.stream); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let mut guard = stream.lock().await; + let Some(inner) = guard.as_mut() else { + return Err(PyStopAsyncIteration::new_err(())); + }; + match inner.next().await { + Some(Ok(step)) => { + let is_terminal = matches!(&step, Step::ReturnToAgent(_)); + let step = Python::attach(|py| step_to_python(py, step))?; + if is_terminal { + guard.take(); + } + drop(guard); + Ok(step) + } + Some(Err(error)) => { + guard.take(); + Err(py_libsy_error(error)) + } + None => { + guard.take(); + Err(PyStopAsyncIteration::new_err(())) + } + } + }) + } + + fn aclose<'py>(&mut self, py: Python<'py>) -> PyResult> { + self.consumed = true; + let stream = Arc::clone(&self.stream); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + stream.lock().await.take(); + Ok(()) + }) + } + + fn __repr__(&self) -> &'static str { + "RunStream()" + } +} + /// Opaque Python handle to a Rust-owned routing algorithm. #[pyclass(name = "Algorithm", module = "switchyard.libsy", frozen)] pub(crate) struct PyAlgorithm { @@ -56,6 +243,19 @@ impl PyAlgorithm { }) } + /// Return a stream that lets the Python host serve model calls. + #[pyo3(signature = (request, *, context=None))] + fn run_stream( + &self, + request: PyRef<'_, PyRequest>, + context: Option>, + ) -> PyRunStream { + let context = context + .map(|context| context.clone_core()) + .unwrap_or_default(); + PyRunStream::new(Arc::clone(&self.inner), context, request.clone_core()) + } + fn __repr__(&self) -> &'static str { "Algorithm()" } @@ -76,6 +276,9 @@ fn random_algorithm(py: Python<'_>, targets: PyRef<'_, PyLlmTargetSet>) -> PyRes } pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_class::()?; + module.add_class::()?; + module.add_class::()?; module.add_class::()?; module.add_function(wrap_pyfunction!(noop_algorithm, module)?)?; module.add_function(wrap_pyfunction!(random_algorithm, module)?)?; diff --git a/examples/libsy.py b/examples/libsy.py new file mode 100644 index 00000000..9ee6f99c --- /dev/null +++ b/examples/libsy.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run libsy algorithms to completion and as a host-driven step stream.""" + +import asyncio + +from switchyard import libsy + + +class ExampleRoutedLlmClient: + """Return a local response for the model selected by an algorithm.""" + + async def call( + self, + context: libsy.protocol.Context, + request: libsy.protocol.Request, + decision: libsy.protocol.Decision, + ) -> libsy.protocol.Response: + return libsy.protocol.Response( + libsy.protocol.AggLlmResponse( + model=decision.selected_model, + outputs=[ + libsy.protocol.ResponseOutput( + libsy.protocol.Role.ASSISTANT, + [ + libsy.protocol.ContentBlock.Text( + f"Response from {decision.selected_model}" + ) + ], + ) + ], + ) + ) + + +def response_text(response: libsy.protocol.Response) -> str: + """Read the first text block from an aggregate response.""" + block = response.aggregate.outputs[0].content[0] + match block: + case libsy.protocol.ContentBlock.Text(text=text): + return text + case _: + raise TypeError("expected a text response") + + +def user_message(text: str) -> libsy.protocol.Message: + """Build a typed user message.""" + content: libsy.protocol.ContentBlockValue = libsy.protocol.ContentBlock.Text(text) + return libsy.protocol.Message(libsy.protocol.Role.USER, [content]) + + +async def main() -> None: + """Exercise no-op and random routing through the libsy execution APIs.""" + algorithm = libsy.algorithms.noop() + request = libsy.protocol.Request( + libsy.protocol.LlmRequest( + model="example-model", + messages=[user_message("Hello from Python")], + ) + ) + + decisions, response = await algorithm.run(request) + print("run") + for decision in decisions: + print(f" decision: {decision.selected_model}") + print(f" response: {response_text(response)}") + + print("run_stream") + async for step in algorithm.run_stream(request): + match step: + case libsy.Step.Decision(decision=decision): + print(f" decision: {decision.selected_model}") + case libsy.Step.CallLlm(call=call): + raise RuntimeError(f"unexpected LLM call for {call.decision.selected_model}") + case libsy.Step.ReturnToAgent(response=response): + print(f" response: {response_text(response)}") + + client = ExampleRoutedLlmClient() + random_algorithm = libsy.algorithms.random( + libsy.LlmTargetSet( + [ + libsy.LlmTarget("fast-model", llm_client=client), + libsy.LlmTarget("smart-model", llm_client=client), + ] + ) + ) + decisions, response = await random_algorithm.run(request) + + print("random routing") + for decision in decisions: + print(f" decision: {decision.selected_model}") + print(f" response: {response_text(response)}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/switchyard/libsy/__init__.py b/switchyard/libsy/__init__.py index ea8dd62b..3de61bff 100644 --- a/switchyard/libsy/__init__.py +++ b/switchyard/libsy/__init__.py @@ -1,9 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Public Python API for running Rust-owned libsy algorithms.""" +"""Public Python API for driving Rust-owned libsy algorithms.""" -from switchyard_rust.libsy import Algorithm, LibsyError, LlmTarget, LlmTargetSet +from switchyard_rust.libsy import ( + Algorithm, + LibsyError, + LlmCall, + LlmTarget, + LlmTargetSet, + RunStream, + Step, +) from . import algorithms as algorithms from . import protocol as protocol @@ -11,8 +19,11 @@ __all__ = [ "Algorithm", "LibsyError", + "LlmCall", "LlmTarget", "LlmTargetSet", + "RunStream", + "Step", "algorithms", "protocol", ] diff --git a/switchyard_rust/libsy.py b/switchyard_rust/libsy.py index 12063955..52f6422f 100644 --- a/switchyard_rust/libsy.py +++ b/switchyard_rust/libsy.py @@ -13,8 +13,11 @@ { "Algorithm", "LibsyError", + "LlmCall", "LlmTarget", "LlmTargetSet", + "RunStream", + "Step", "noop", "random", } @@ -23,8 +26,11 @@ if TYPE_CHECKING: Algorithm: type[Any] LibsyError: type[RuntimeError] + LlmCall: type[Any] LlmTarget: type[Any] LlmTargetSet: type[Any] + RunStream: type[Any] + Step: type[Any] def __getattr__(name: str) -> object: diff --git a/switchyard_rust/libsy.pyi b/switchyard_rust/libsy.pyi index 617213b8..1f8f2609 100644 --- a/switchyard_rust/libsy.pyi +++ b/switchyard_rust/libsy.pyi @@ -1,12 +1,21 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import AsyncIterator from typing import final -from switchyard_rust.libsy_protocol import Context, Decision, Request, Response, RoutedLlmClient +from switchyard_rust.libsy_protocol import ( + Context, + Decision, + Request, + Response, + RoutedLlmClient, +) class LibsyError(RuntimeError): ... +_DecisionValue = Decision + class LlmTarget: def __init__( self, @@ -29,6 +38,30 @@ class LlmTargetSet: def get_target(self, name: str) -> LlmTarget: ... def __len__(self) -> int: ... +class LlmCall: + context: Context + request: Request + decision: Decision + is_pending: bool + + def respond(self, response: Response) -> None: ... + def fail(self, message: str) -> None: ... + +class Step: + class CallLlm: + call: LlmCall + + class Decision: + decision: _DecisionValue + + class ReturnToAgent: + response: Response + +class RunStream(AsyncIterator[Step]): + def __aiter__(self) -> RunStream: ... + async def __anext__(self) -> Step: ... + async def aclose(self) -> None: ... + @final class Algorithm: async def run( @@ -37,6 +70,12 @@ class Algorithm: *, context: Context | None = None, ) -> tuple[list[Decision], Response]: ... + def run_stream( + self, + request: Request, + *, + context: Context | None = None, + ) -> RunStream: ... def noop() -> Algorithm: ... def random(targets: LlmTargetSet) -> Algorithm: ... diff --git a/tests/test_libsy_protocol_bindings.py b/tests/test_libsy_protocol_bindings.py index 15ac656e..b67ac792 100644 --- a/tests/test_libsy_protocol_bindings.py +++ b/tests/test_libsy_protocol_bindings.py @@ -15,7 +15,6 @@ def test_libsy_uses_an_explicit_namespace() -> None: assert switchyard.libsy is libsy assert libsy.protocol.__name__ == "switchyard.libsy.protocol" - assert not hasattr(libsy, "RunStream") def test_protocol_request_and_response_fields_stay_typed() -> None: diff --git a/tests/test_libsy_run_stream_bindings.py b/tests/test_libsy_run_stream_bindings.py new file mode 100644 index 00000000..7da6c0e2 --- /dev/null +++ b/tests/test_libsy_run_stream_bindings.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio + +import pytest + +from switchyard import libsy + + +async def test_noop_algorithm_is_driven_with_structural_step_matching() -> None: + run = libsy.algorithms.noop().run_stream( + libsy.protocol.Request(libsy.protocol.LlmRequest(model="requested-model")) + ) + observed: list[tuple[str, str | None]] = [] + + async for step in run: + match step: + case libsy.Step.Decision(decision=decision): + observed.append(("decision", decision.selected_model)) + case libsy.Step.ReturnToAgent(response=response): + observed.append(("return", response.selected_model)) + match response.aggregate.outputs[0].content[0]: + case libsy.protocol.ContentBlock.Text(text="OK"): + pass + case block: + pytest.fail(f"unexpected content block: {block!r}") + case libsy.Step.CallLlm(): + pytest.fail("the no-op algorithm must not request an LLM call") + + assert observed == [ + ("decision", "requested-model"), + ("return", "requested-model"), + ] + with pytest.raises(RuntimeError, match="RunStream has already been consumed"): + async for _ in run: + pass + + +def test_run_stream_starts_lazily_without_an_event_loop() -> None: + run = libsy.algorithms.noop().run_stream( + libsy.protocol.Request(libsy.protocol.LlmRequest(model="requested-model")) + ) + + async def consume() -> list[str | None]: + selected_models: list[str | None] = [] + async for step in run: + match step: + case libsy.Step.Decision(decision=decision): + selected_models.append(decision.selected_model) + case _: + pass + return selected_models + + assert asyncio.run(consume()) == ["requested-model"] + + +async def test_random_algorithm_is_driven_through_the_same_handle() -> None: + algorithm = libsy.algorithms.random(libsy.LlmTargetSet([libsy.LlmTarget("only/model")])) + request = libsy.protocol.Request(libsy.protocol.LlmRequest(model="auto")) + context = libsy.protocol.Context(values={"request_id": "request-1"}) + observed: list[tuple[str, str | None]] = [] + + async for step in algorithm.run_stream(request, context=context): + match step: + case libsy.Step.Decision(decision=decision): + observed.append(("decision", decision.selected_model)) + case libsy.Step.CallLlm(call=call): + observed.append(("call", call.decision.selected_model)) + assert call.context.values == {"request_id": "request-1"} + call.respond( + libsy.protocol.Response( + libsy.protocol.AggLlmResponse(model=call.decision.selected_model) + ) + ) + case libsy.Step.ReturnToAgent(response=response): + observed.append(("return", response.selected_model)) + + assert observed == [ + ("decision", "only/model"), + ("call", "only/model"), + ("return", "only/model"), + ]