From d35d417818f04ada0d452775933941d92d04cfb8 Mon Sep 17 00:00:00 2001 From: nachiketb Date: Wed, 12 Aug 2026 14:01:11 -0700 Subject: [PATCH] feat(python): expose libsy run stream Signed-off-by: nachiketb --- Cargo.lock | 3 +- crates/switchyard-py/Cargo.toml | 3 +- crates/switchyard-py/src/libsy_bindings.rs | 434 +++++++++++---------- examples/libsy.py | 45 ++- switchyard/libsy/__init__.py | 10 +- switchyard_rust/libsy.py | 90 +++-- tests/test_libsy_minimal_bindings.py | 158 +++++--- 7 files changed, 440 insertions(+), 303 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 39871f0e9..3d0ea85f6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2327,7 +2327,7 @@ dependencies = [ name = "switchyard-py" version = "0.2.0" dependencies = [ - "async-trait", + "futures", "http", "pyo3", "pyo3-async-runtimes", @@ -2335,7 +2335,6 @@ dependencies = [ "serde", "serde_json", "switchyard-libsy", - "switchyard-llm-client", "switchyard-protocol", "switchyard-server", "tokio", diff --git a/crates/switchyard-py/Cargo.toml b/crates/switchyard-py/Cargo.toml index beb4e1b9f..6fe4fdbb4 100644 --- a/crates/switchyard-py/Cargo.toml +++ b/crates/switchyard-py/Cargo.toml @@ -16,10 +16,9 @@ name = "_switchyard_rust" crate-type = ["cdylib", "rlib"] [dependencies] -async-trait.workspace = true +futures.workspace = true http.workspace = true switchyard-libsy.workspace = true -switchyard-llm-client.workspace = true pyo3 = { version = "0.28.3", features = ["abi3-py310", "extension-module"] } pyo3-async-runtimes = { version = "0.28", features = ["tokio-runtime"] } pythonize = "0.28.0" diff --git a/crates/switchyard-py/src/libsy_bindings.rs b/crates/switchyard-py/src/libsy_bindings.rs index d7ca40399..ee2b70fd9 100644 --- a/crates/switchyard-py/src/libsy_bindings.rs +++ b/crates/switchyard-py/src/libsy_bindings.rs @@ -6,21 +6,20 @@ use std::collections::HashMap; use std::sync::Arc; -use async_trait::async_trait; +use futures::StreamExt; use http::header::{HeaderName, HeaderValue}; -use pyo3::exceptions::{PyTypeError, PyValueError}; +use pyo3::exceptions::{PyBaseException, PyStopAsyncIteration, PyTypeError, PyValueError}; use pyo3::prelude::*; -use serde_json::{Value, json}; use switchyard_libsy::{ - Algorithm, ClassifierContractConfig, HandoffNoteConfig, LibsyError as RustLibsyError, - LlmClassifierConfig, LlmFallback, LlmTaskClassifier, Noop, PickerMode, Random, StageRouter, - StageRouterConfig, TaskClassifierConfig, + Algorithm, CallModel, ClassifierContractConfig, HandoffNoteConfig, + LibsyError as RustLibsyError, LlmClassifierConfig, LlmFallback, LlmTaskClassifier, Noop, + PickerMode, Random, StageRouter, StageRouterConfig, Step as RustStep, StepStream, + TaskClassifierConfig, }; -use switchyard_llm_client::ClientRouter; use switchyard_protocol::{ - AggLlmResponse, LlmClientError, LlmResponse, Metadata, ModelId, Request, Response, - RoutedLlmClient, + AggLlmResponse, Decision, LlmClientError, LlmResponse, Metadata, ModelId, Request, Response, }; +use tokio::sync::Mutex; use crate::errors::py_libsy_error; use crate::py_serde::{from_python, to_python}; @@ -40,84 +39,6 @@ fn header_map_from_python(headers: &HashMap) -> PyResult, -} - -#[async_trait] -impl RoutedLlmClient for PythonLlmClient { - async fn call(&self, request: Request) -> Result { - let metadata = request.metadata; - let future = Python::attach(|py| { - let request = to_python(py, &request.llm_request)?; - let awaitable = self.inner.bind(py).call_method1("call", (request,))?; - pyo3_async_runtimes::tokio::into_future(awaitable) - }) - .map_err(other_python_error)?; - - let response = future.await.map_err(other_python_error)?; - let aggregate = Python::attach(|py| from_python::(response.bind(py))) - .map_err(invalid_python_response)?; - Ok(Response { - llm_response: LlmResponse::Agg(aggregate), - metadata, - }) - } -} - -/// A required-client routing target used by Python-created algorithms. -#[pyclass(name = "LlmTarget", module = "switchyard.libsy", frozen)] -struct PyLlmTarget { - name: String, - client: Py, -} - -impl PyLlmTarget { - /// The bare model id libsy routes by; the client behind it stays with the bindings. - fn clone_core(&self, _py: Python<'_>) -> ModelId { - ModelId::new(self.name.clone()) - } - - /// The `selected_model -> client` entry this target contributes to the algorithm's - /// [`ClientRouter`]. libsy no longer carries the client, so the bindings keep the - /// mapping and serve the calls themselves. - fn client_entry(&self, py: Python<'_>) -> ClientEntry { - ( - ModelId::new(self.name.clone()), - Arc::new(PythonLlmClient { - inner: self.client.clone_ref(py), - }), - ) - } -} - -#[pymethods] -impl PyLlmTarget { - #[new] - fn new(py: Python<'_>, name: String, client: Py) -> PyResult { - let call = client - .bind(py) - .getattr("call") - .map_err(|_| PyTypeError::new_err("client must define async call(request)"))?; - if !call.is_callable() { - return Err(PyTypeError::new_err( - "client.call must be callable as async call(request)", - )); - } - Ok(Self { name, client }) - } - - #[getter] - fn name(&self) -> &str { - &self.name - } - - fn __repr__(&self) -> String { - format!("LlmTarget(name={:?})", self.name) - } -} - /// Classifier settings shared by standalone and stage-router classifiers. #[pyclass( name = "TaskClassifierConfig", @@ -185,19 +106,14 @@ impl PyTaskClassifierConfig { skip_from_py_object )] struct PyLlmFallback { - judge_target: Py, + judge_target: String, config: Py, } impl PyLlmFallback { - /// The judge's client entry, so the caller can register it with the algorithm's router. - fn judge_client_entry(&self, py: Python<'_>) -> PyResult { - Ok(self.judge_target.bind(py).try_borrow()?.client_entry(py)) - } - fn clone_core(&self, py: Python<'_>) -> PyResult { Ok(LlmFallback { - judge_target: self.judge_target.bind(py).try_borrow()?.clone_core(py), + judge_target: ModelId::new(self.judge_target.clone()), config: self.config.bind(py).try_borrow()?.clone_core(), }) } @@ -207,7 +123,7 @@ impl PyLlmFallback { impl PyLlmFallback { #[new] #[pyo3(signature = (judge_target, *, config))] - fn new(judge_target: Py, config: Py) -> Self { + fn new(judge_target: String, config: Py) -> Self { Self { judge_target, config, @@ -215,117 +131,260 @@ impl PyLlmFallback { } } -/// Opaque handle shared by every Rust-owned algorithm exposed to Python. -#[pyclass(name = "Algorithm", module = "switchyard.libsy", frozen)] -struct PyAlgorithm { - inner: Arc, - /// Resolves the calls `inner` offloads to each target's Python client. - client_router: ClientRouter, +/// A routing choice produced by an algorithm. +#[pyclass(name = "Decision", module = "switchyard.libsy", frozen)] +struct PyDecision { + inner: Decision, } -/// One target's `selected_model -> client` mapping for an algorithm's router. -type ClientEntry = (ModelId, Arc); +impl From for PyDecision { + fn from(inner: Decision) -> Self { + Self { inner } + } +} -impl PyAlgorithm { - fn new(inner: Arc, clients: impl IntoIterator) -> Self { - Self { - inner, - client_router: clients.into_iter().collect(), +#[pymethods] +impl PyDecision { + /// The semantic model id selected for the call. + #[getter] + fn selected_model_id(&self) -> &str { + self.inner.selected_model_id().as_str() + } + + /// Why the algorithm selected this model, when supplied. + #[getter] + fn reasoning(&self) -> Option<&str> { + self.inner.reasoning() + } + + /// Whether this call produces the answer rather than a routing verdict. + #[getter] + fn is_answer_call(&self) -> bool { + self.inner.is_answer_call() + } + + fn __repr__(&self) -> String { + format!( + "Decision(selected_model_id={:?}, reasoning={:?}, is_answer_call={})", + self.inner.selected_model_id(), + self.inner.reasoning(), + self.inner.is_answer_call() + ) + } +} + +/// One model call yielded by [`PyAlgorithm::run_stream`]. +#[pyclass(name = "ModelCall", module = "switchyard.libsy")] +struct PyModelCall { + inner: Option, + algorithm: String, + request: Py, + decision: Py, +} + +impl PyModelCall { + fn new(py: Python<'_>, call: CallModel) -> PyResult { + let request = to_python(py, &call.request.llm_request)?; + let decision = Py::new(py, PyDecision::from(call.decision.clone()))?; + Ok(Self { + algorithm: call.algorithm.clone(), + inner: Some(call), + request, + decision, + }) + } + + fn take(&mut self) -> PyResult { + self.inner + .take() + .ok_or_else(|| py_libsy_error("model call has already been completed")) + } +} + +#[pymethods] +impl PyModelCall { + /// The algorithm that produced this call. + #[getter] + fn algorithm(&self) -> &str { + &self.algorithm + } + + /// The normalized LLM request to serve as a Python dictionary. + #[getter] + fn request(&self, py: Python<'_>) -> Py { + self.request.clone_ref(py) + } + + /// The routing decision behind this call. + #[getter] + fn decision(&self, py: Python<'_>) -> Py { + self.decision.clone_ref(py) + } + + /// Consume the answer call without serving it and return its rewritten request and decision. + #[pyo3(name = "into_parts")] + fn take_parts(&mut self, py: Python<'_>) -> PyResult<(Py, Py)> { + let (request, decision) = self.take()?.into_parts(); + Ok(( + to_python(py, &request.llm_request)?, + Py::new(py, PyDecision::from(decision))?, + )) + } + + /// Fulfill this call with an aggregate normalized response dictionary. + fn respond(&mut self, response: &Bound<'_, PyAny>) -> PyResult<()> { + let aggregate = from_python::(response)?; + let call = self.take()?; + let metadata = call.request.metadata.clone(); + call.respond(Ok(Response { + llm_response: LlmResponse::Agg(aggregate), + metadata, + })) + .map_err(py_libsy_error) + } + + /// Fulfill this call with a Python client failure. + fn fail(&mut self, error: &Bound<'_, PyAny>) -> PyResult<()> { + if !error.is_instance_of::() { + return Err(PyTypeError::new_err("error must derive from BaseException")); } + let call = self.take()?; + let target = call.decision.selected_model_id().clone(); + let source = LlmClientError::Ffi { + source: Box::new(PyErr::from_value(error.clone())), + }; + call.respond(Err(RustLibsyError::client_call(target, source))) + .map_err(py_libsy_error) } } +/// One item yielded by a Python algorithm stream. +#[pyclass(name = "Step", module = "switchyard.libsy", frozen)] +enum PyStep { + /// The host must serve the model call before the algorithm can continue. + CallModel { call: Py }, + /// A routing decision emitted by the algorithm. + Decision { decision: Py }, + /// The terminal aggregate response. + Done { response: Py }, +} + +/// Async Python iterator over one Rust algorithm run. +#[pyclass(name = "_RunStream", module = "switchyard.libsy", frozen)] +struct PyRunStream { + inner: Arc>, +} + +#[pymethods] +impl PyRunStream { + fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn __anext__<'py>(&self, py: Python<'py>) -> PyResult> { + let stream = Arc::clone(&self.inner); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let step = stream.lock().await.next().await; + match step { + Some(Ok(step)) => step_to_python(step).await, + Some(Err(error)) => Err(py_libsy_error(error)), + None => Err(PyStopAsyncIteration::new_err(())), + } + }) + } +} + +/// Opaque handle shared by every Rust-owned algorithm exposed to Python. +#[pyclass(name = "Algorithm", module = "switchyard.libsy", frozen)] +struct PyAlgorithm { + inner: Arc, +} + #[pymethods] impl PyAlgorithm { - /// Run to completion using the clients configured on the algorithm's targets. + /// Run the algorithm as a stream of model calls, decisions, and one terminal response. /// /// `headers`, when given, is normalized into the request's correlation /// [`Metadata`] exactly as an HTTP host would (`Metadata::from_headers`), /// so metadata-driven algorithms see the same signals in Python as when /// served over HTTP. #[pyo3(signature = (request, headers=None))] - fn run<'py>( + fn run_stream( &self, - py: Python<'py>, request: &Bound<'_, PyAny>, headers: Option>, - ) -> PyResult> { - let algorithm = Arc::clone(&self.inner); - let client_router = self.client_router.clone(); + ) -> PyResult { let headers = headers.as_ref().map(header_map_from_python).transpose()?; - let request = Request { llm_request: from_python(request)?, raw_request: None, metadata: headers.map(|headers| Metadata::from_headers(&headers)), }; - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let (decisions, response) = - switchyard_llm_client::run(algorithm, client_router, request, None) - .await - .map_err(py_libsy_error)?; + let stream = { + let _guard = pyo3_async_runtimes::tokio::get_runtime().enter(); + Arc::clone(&self.inner).run_stream(request) + }; + Ok(PyRunStream { + inner: Arc::new(Mutex::new(stream)), + }) + } + + fn __repr__(&self) -> &'static str { + "Algorithm()" + } +} + +async fn step_to_python(step: RustStep) -> PyResult { + match step { + RustStep::CallModel(call) => Python::attach(|py| { + Ok(PyStep::CallModel { + call: Py::new(py, PyModelCall::new(py, *call)?)?, + }) + }), + RustStep::Decision(decision) => Python::attach(|py| { + Ok(PyStep::Decision { + decision: Py::new(py, PyDecision::from(decision))?, + }) + }), + RustStep::Done(response) => { let response = response .llm_response .into_agg() .await .map_err(py_libsy_error)?; - let decisions = decisions - .iter() - .map(|decision| { - json!({ - "selected_model_id": decision.selected_model_id(), - "reasoning": decision.reasoning(), - "is_answer_call": decision.is_answer_call(), - }) + Python::attach(|py| { + Ok(PyStep::Done { + response: to_python(py, &response)?, }) - .collect::>(); - Python::attach(|py| Ok((to_python(py, &decisions)?, to_python(py, &response)?))) - }) - } - - fn __repr__(&self) -> &'static str { - "Algorithm()" + }) + } } } /// Construct the no-op reference algorithm. #[pyfunction(name = "noop")] fn noop_algorithm() -> PyAlgorithm { - // `Noop` synthesizes its own response and never offloads a call, so it needs no clients. - PyAlgorithm::new(Arc::new(Noop {}), []) + PyAlgorithm { + inner: Arc::new(Noop {}), + } } /// Construct random routing over targets with optional relative weights and seed. #[pyfunction(name = "random")] #[pyo3(signature = (targets, *, weights=None, seed=None))] fn random_algorithm( - py: Python<'_>, - targets: Vec>, + targets: Vec, weights: Option>, seed: Option, ) -> PyResult { - let (cores, clients) = target_cores(py, &targets)?; - let algorithm = Random::new(cores, weights, seed).map_err(|error| match error { + let model_ids = targets.into_iter().map(ModelId::new).collect(); + let algorithm = Random::new(model_ids, weights, seed).map_err(|error| match error { RustLibsyError::NoTargets => PyValueError::new_err("random requires at least one target"), other => PyValueError::new_err(other.to_string()), })?; - Ok(PyAlgorithm::new(Arc::new(algorithm), clients)) -} - -/// Splits a list of Python targets into libsy's client-free targets and the client entries -/// the bindings keep for the algorithm's router. -fn target_cores( - py: Python<'_>, - targets: &[Py], -) -> PyResult<(Vec, Vec)> { - let mut cores = Vec::with_capacity(targets.len()); - let mut clients = Vec::with_capacity(targets.len()); - for target in targets { - let target = target.bind(py).try_borrow()?; - cores.push(target.clone_core(py)); - clients.push(target.client_entry(py)); - } - Ok((cores, clients)) + Ok(PyAlgorithm { + inner: Arc::new(algorithm), + }) } /// Construct task-level LLM classifier routing. @@ -339,26 +398,21 @@ fn target_cores( ))] fn llm_task_classifier_algorithm( py: Python<'_>, - judge_target: Py, - efficient_target: Py, - capable_target: Py, + judge_target: String, + efficient_target: String, + capable_target: String, config: Py, ) -> PyResult { - let (cores, clients) = target_cores( - py, - &[judge_target.clone_ref(py), efficient_target, capable_target], - )?; - let [judge, efficient, capable] = cores - .try_into() - .map_err(|_| PyValueError::new_err("expected three targets"))?; let algorithm = LlmTaskClassifier::new(LlmClassifierConfig::Capability { - judge_target: judge, - efficient_target: efficient, - capable_target: capable, + judge_target: ModelId::new(judge_target), + efficient_target: ModelId::new(efficient_target), + capable_target: ModelId::new(capable_target), config: config.bind(py).try_borrow()?.clone_core(), }) .map_err(|error| PyValueError::new_err(error.to_string()))?; - Ok(PyAlgorithm::new(Arc::new(algorithm), clients)) + Ok(PyAlgorithm { + inner: Arc::new(algorithm), + }) } /// Construct signal-driven stage routing with an optional LLM classifier fallback. @@ -380,8 +434,8 @@ fn llm_task_classifier_algorithm( #[allow(clippy::too_many_arguments)] fn stage_router_algorithm( py: Python<'_>, - capable_target: Py, - efficient_target: Py, + capable_target: String, + efficient_target: String, picker: &str, confidence_threshold: f64, recent_window: Option, @@ -401,10 +455,8 @@ fn stage_router_algorithm( ))); } }; - let (cores, mut clients) = target_cores(py, &[capable_target, efficient_target])?; - let [capable, efficient] = cores - .try_into() - .map_err(|_| PyValueError::new_err("expected two targets"))?; + let capable = ModelId::new(capable_target); + let efficient = ModelId::new(efficient_target); let mut config = StageRouterConfig::new(mode, confidence_threshold); config.recent_window = recent_window; config.handoff_notes = match (escalation_note, deescalation_note) { @@ -426,37 +478,25 @@ fn stage_router_algorithm( if let Some(prompt) = efficient_system_prompt { config.tier_prompts = config.tier_prompts.with(efficient.clone(), prompt); } - // The judge is only reachable through the optional classifier fallback, so its client - // joins the router only when a fallback is configured. - if let Some(classifier) = &classifier { - clients.push(classifier.bind(py).try_borrow()?.judge_client_entry(py)?); - } config.llm_fallback = classifier .map(|classifier| classifier.bind(py).try_borrow()?.clone_core(py)) .transpose()?; let algorithm = StageRouter::new(capable, efficient, config) .map_err(|error| PyValueError::new_err(error.to_string()))?; - Ok(PyAlgorithm::new(Arc::new(algorithm), clients)) -} - -fn other_python_error(error: PyErr) -> LlmClientError { - LlmClientError::Ffi { - source: Box::new(error), - } -} - -fn invalid_python_response(error: PyErr) -> LlmClientError { - LlmClientError::InvalidResponse { - source: Box::new(error), - } + Ok(PyAlgorithm { + inner: Arc::new(algorithm), + }) } pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { let libsy_module = PyModule::new(module.py(), "libsy")?; libsy_module.add_class::()?; + libsy_module.add_class::()?; libsy_module.add_class::()?; - libsy_module.add_class::()?; + libsy_module.add_class::()?; + libsy_module.add_class::()?; + libsy_module.add_class::()?; libsy_module.add_class::()?; libsy_module.add_function(wrap_pyfunction!(noop_algorithm, &libsy_module)?)?; libsy_module.add_function(wrap_pyfunction!(random_algorithm, &libsy_module)?)?; diff --git a/examples/libsy.py b/examples/libsy.py index 5da1d2530..3a1536331 100644 --- a/examples/libsy.py +++ b/examples/libsy.py @@ -2,49 +2,52 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Run libsy's no-op and random-routing algorithms from Python.""" +"""Drive a libsy algorithm stream from Python.""" import asyncio from collections.abc import Mapping -from switchyard.libsy import LlmTarget, algorithms +from switchyard.libsy import Step, algorithms class EchoClient: - """Return its configured model as the completion.""" + """Return a fixed completion for any selected target.""" - def __init__(self, model: str) -> None: - self.model = model - - async def call(self, request: Mapping[str, object]) -> Mapping[str, object]: + async def call( + self, + request: Mapping[str, object], + model: str, + ) -> Mapping[str, object]: return { - "model": self.model, + "model": model, "outputs": [ - {"role": "assistant", "content": [{"type": "text", "text": self.model}]} + {"role": "assistant", "content": [{"type": "text", "text": "Hello"}]} ], } async def main() -> None: - """Run both algorithms and print their aggregate results.""" + """Run random routing and serve its selected target.""" request = { "model": "auto", "messages": [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}], } - - noop_decisions, noop_response = await algorithms.noop().run(request) - print("No-op:", noop_decisions, noop_response) - - random = algorithms.random( - [ - LlmTarget("fast", EchoClient("fast")), - LlmTarget("quality", EchoClient("quality")), - ], + client = EchoClient() + algorithm = algorithms.random( + ["fast", "quality"], weights=[1, 3], seed=42, ) - random_decisions, random_response = await random.run(request) - print("Random:", random_decisions, random_response) + + async for step in algorithm.run_stream(request): + match step: + case Step.Decision(decision): + print("Decision:", decision.selected_model_id, decision.reasoning) + case Step.CallModel(call): + model = call.decision.selected_model_id + call.respond(await client.call(call.request, model)) + case Step.Done(response): + print("Response:", response) if __name__ == "__main__": diff --git a/switchyard/libsy/__init__.py b/switchyard/libsy/__init__.py index 44ca4cea3..6d2b5ba74 100644 --- a/switchyard/libsy/__init__.py +++ b/switchyard/libsy/__init__.py @@ -5,10 +5,11 @@ from switchyard_rust.libsy import ( Algorithm, + Decision, LibsyError, - LlmClient, LlmFallback, - LlmTarget, + ModelCall, + Step, TaskClassifierConfig, ) @@ -16,10 +17,11 @@ __all__ = [ "Algorithm", + "Decision", "LibsyError", - "LlmClient", "LlmFallback", - "LlmTarget", + "ModelCall", + "Step", "TaskClassifierConfig", "algorithms", ] diff --git a/switchyard_rust/libsy.py b/switchyard_rust/libsy.py index a7eb5c85a..244b1dd02 100644 --- a/switchyard_rust/libsy.py +++ b/switchyard_rust/libsy.py @@ -6,16 +6,18 @@ from __future__ import annotations from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Protocol +from typing import TYPE_CHECKING, Any from switchyard_rust._native import load_native _EXPORTS = frozenset( { "Algorithm", + "Decision", "LibsyError", "LlmFallback", - "LlmTarget", + "ModelCall", + "Step", "TaskClassifierConfig", "llm_task_classifier", "noop", @@ -24,30 +26,59 @@ } ) +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Sequence + from typing import ClassVar, Literal, final -class LlmClient(Protocol): - """Structural interface for a Python-hosted model client.""" + class LibsyError(RuntimeError): ... - async def call( - self, - request: Mapping[str, object], - ) -> Mapping[str, object]: - """Call the configured target and return an aggregate neutral response.""" - ... + @final + class Decision: + """A semantic routing choice produced by an algorithm.""" + @property + def selected_model_id(self) -> str: ... -if TYPE_CHECKING: - from collections.abc import Sequence - from typing import final + @property + def reasoning(self) -> str | None: ... - class LibsyError(RuntimeError): ... + @property + def is_answer_call(self) -> bool: ... + + _RoutingDecision = Decision @final - class LlmTarget: - def __init__(self, name: str, client: LlmClient) -> None: ... + class ModelCall: + @property + def algorithm(self) -> str: ... @property - def name(self) -> str: ... + def request(self) -> dict[str, object]: ... + + @property + def decision(self) -> Decision: ... + + def into_parts(self) -> tuple[dict[str, object], Decision]: ... + + def respond(self, response: Mapping[str, object]) -> None: ... + + def fail(self, error: BaseException) -> None: ... + + class Step: + @final + class CallModel: + __match_args__: ClassVar[tuple[Literal["call"]]] = ("call",) + call: ModelCall + + @final + class Decision: + __match_args__: ClassVar[tuple[Literal["decision"]]] = ("decision",) + decision: _RoutingDecision + + @final + class Done: + __match_args__: ClassVar[tuple[Literal["response"]]] = ("response",) + response: dict[str, object] @final class TaskClassifierConfig: @@ -67,34 +98,39 @@ def __init__( class LlmFallback: def __init__( self, - judge_target: LlmTarget, + judge_target: str, *, config: TaskClassifierConfig, ) -> None: ... @final class Algorithm: - async def run( + def run_stream( self, request: Mapping[str, object], headers: Mapping[str, str] | None = None, - ) -> tuple[list[dict[str, object]], dict[str, object]]: ... + ) -> AsyncIterator[Step.CallModel | Step.Decision | Step.Done]: ... def noop() -> Algorithm: ... - def random(targets: Sequence[LlmTarget]) -> Algorithm: ... + def random( + targets: Sequence[str], + *, + weights: Sequence[float] | None = None, + seed: int | None = None, + ) -> Algorithm: ... def llm_task_classifier( - judge_target: LlmTarget, - efficient_target: LlmTarget, - capable_target: LlmTarget, + judge_target: str, + efficient_target: str, + capable_target: str, *, config: TaskClassifierConfig, ) -> Algorithm: ... def stage_router( - capable_target: LlmTarget, - efficient_target: LlmTarget, + capable_target: str, + efficient_target: str, *, picker: str, confidence_threshold: float, @@ -115,4 +151,4 @@ def __getattr__(name: str) -> object: raise AttributeError(f"module {__name__!r} has no attribute {name!r}") -__all__ = [*sorted(_EXPORTS), "LlmClient"] +__all__ = sorted(_EXPORTS) diff --git a/tests/test_libsy_minimal_bindings.py b/tests/test_libsy_minimal_bindings.py index 059ccf0aa..eb1d6c8ce 100644 --- a/tests/test_libsy_minimal_bindings.py +++ b/tests/test_libsy_minimal_bindings.py @@ -7,7 +7,14 @@ import pytest -from switchyard.libsy import LibsyError, LlmTarget, TaskClassifierConfig, algorithms +from switchyard.libsy import ( + Algorithm, + Decision, + LibsyError, + Step, + TaskClassifierConfig, + algorithms, +) def request_body() -> dict[str, Any]: @@ -41,27 +48,84 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: } -async def test_random_runs_with_a_python_client() -> None: +async def run_algorithm( + algorithm: Algorithm, + clients: dict[str, Any] | None = None, + *, + request: dict[str, Any] | None = None, + headers: dict[str, str] | None = None, +) -> tuple[list[Decision], dict[str, Any]]: + decisions: list[Decision] = [] + async for step in algorithm.run_stream(request or request_body(), headers=headers): + match step: + case Step.CallModel(call): + target = call.decision.selected_model_id + client = (clients or {})[target] + try: + response = await client.call(call.request) + except Exception as error: + call.fail(error) + else: + call.respond(response) + case Step.Decision(decision): + decisions.append(decision) + case Step.Done(response): + return decisions, response + raise AssertionError("algorithm stream ended without a response") + + +async def test_random_streams_complex_steps_and_accepts_a_dictionary_response() -> None: client = EchoClient("fast") - algorithm = algorithms.random([LlmTarget("fast", client)]) - - decisions, response = await algorithm.run(request_body()) - - assert decisions == [ - { - "selected_model_id": "fast", - "reasoning": "random routing selected target 'fast'", - "is_answer_call": True, - } - ] + algorithm = algorithms.random(["fast"]) + decisions: list[Decision] = [] + response: dict[str, Any] | None = None + variants: list[str] = [] + + async for step in algorithm.run_stream(request_body()): + match step: + case Step.CallModel(call): + variants.append("call_model") + client_response = await client.call(call.request) + call.respond(client_response) + with pytest.raises(LibsyError, match="already been completed"): + call.respond(client_response) + case Step.Decision(decision): + variants.append("decision") + decisions.append(decision) + case Step.Done(done): + variants.append("done") + response = done + + assert variants == ["decision", "call_model", "done"] + assert len(decisions) == 1 + assert decisions[0].selected_model_id == "fast" + assert decisions[0].reasoning == "random routing selected target 'fast'" + assert decisions[0].is_answer_call is True assert client.calls[0]["model"] == "fast" assert client.calls[0]["messages"][0]["content"] == [ {"type": "text", "text": "hello"} ] + assert response is not None assert response["model"] == "fast" assert response["outputs"][0]["content"] == [{"type": "text", "text": "fast"}] +async def test_into_parts_supports_decision_only_routing() -> None: + algorithm = algorithms.random(["fast"]) + + async for step in algorithm.run_stream(request_body()): + match step: + case Step.CallModel(call) if call.decision.is_answer_call: + request, decision = call.into_parts() + assert call.algorithm == "random" + assert request["messages"] == request_body()["messages"] + assert decision.selected_model_id == "fast" + assert decision.is_answer_call is True + with pytest.raises(LibsyError, match="already been completed"): + call.into_parts() + break + + async def test_classifier_config_accepts_a_prompt_override() -> None: """Verify that a configured classifier prompt is rendered for the judge.""" @@ -90,9 +154,9 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: judge = JudgeClient("judge") weak = EchoClient("weak") algorithm = algorithms.llm_task_classifier( - LlmTarget("judge", judge), - LlmTarget("weak", weak), - LlmTarget("strong", EchoClient("strong")), + "judge", + "weak", + "strong", config=TaskClassifierConfig( 0.5, threshold_step=0.1, @@ -100,7 +164,14 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: ), ) - _, response = await algorithm.run(request_body()) + _, response = await run_algorithm( + algorithm, + { + "judge": judge, + "weak": weak, + "strong": EchoClient("strong"), + }, + ) prompt = judge.calls[0]["instructions"][0]["content"][0]["text"] assert prompt == "Custom capability rubric." @@ -113,37 +184,32 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: async def test_random_weights_and_seed_are_reproducible() -> None: def algorithm(): return algorithms.random( - [ - LlmTarget("fast", EchoClient("fast")), - LlmTarget("capable", EchoClient("capable")), - ], + ["fast", "capable"], weights=[1, 3], seed=42, ) first_router = algorithm() second_router = algorithm() - first = [(await first_router.run(request_body()))[1]["model"] for _ in range(100)] - second = [(await second_router.run(request_body()))[1]["model"] for _ in range(100)] + clients = {"fast": EchoClient("fast"), "capable": EchoClient("capable")} + first = [(await run_algorithm(first_router, clients))[1]["model"] for _ in range(100)] + second = [(await run_algorithm(second_router, clients))[1]["model"] for _ in range(100)] assert first == second assert 65 <= second.count("capable") <= 85 def test_random_rejects_invalid_weights() -> None: - targets = [ - LlmTarget("fast", EchoClient("fast")), - LlmTarget("capable", EchoClient("capable")), - ] + targets = ["fast", "capable"] with pytest.raises(ValueError, match="expected 2 weights, got 1"): algorithms.random(targets, weights=[1]) async def test_noop_needs_no_client() -> None: - decisions, response = await algorithms.noop().run(request_body()) + decisions, response = await run_algorithm(algorithms.noop()) - assert decisions[0]["selected_model_id"] == "auto" + assert decisions[0].selected_model_id == "auto" assert response["outputs"][0]["content"] == [{"type": "text", "text": "OK"}] @@ -156,37 +222,29 @@ async def test_noop_needs_no_client() -> None: ) def test_algorithm_rejects_invalid_headers(headers: dict[str, str], message: str) -> None: with pytest.raises(ValueError, match=message): - algorithms.noop().run(request_body(), headers=headers) + algorithms.noop().run_stream(request_body(), headers=headers) async def test_algorithm_accepts_case_insensitive_duplicate_names() -> None: - decisions, _ = await algorithms.noop().run( - request_body(), headers={"X-Unused": "first", "x-unused": "second"} + decisions, _ = await run_algorithm( + algorithms.noop(), headers={"X-Unused": "first", "x-unused": "second"} ) - assert decisions[0]["selected_model_id"] == "auto" + assert decisions[0].selected_model_id == "auto" def test_algorithm_rejects_header_map_capacity_overflow() -> None: headers = {f"x-header-{index}": "value" for index in range(32_769)} with pytest.raises(ValueError, match="max size reached"): - algorithms.noop().run(request_body(), headers=headers) + algorithms.noop().run_stream(request_body(), headers=headers) -def test_algorithm_exposes_only_managed_execution() -> None: +def test_algorithm_exposes_only_streaming_execution() -> None: algorithm = algorithms.noop() - assert callable(algorithm.run) - assert not hasattr(algorithm, "run_stream") - - -def test_target_requires_a_callable_client() -> None: - with pytest.raises(TypeError, match="client must define async call"): - LlmTarget("fast", object()) - - with pytest.raises(TypeError, match="client.call must be callable"): - LlmTarget("fast", type("Client", (), {"call": None})()) + assert callable(algorithm.run_stream) + assert not hasattr(algorithm, "run") def test_random_requires_a_target() -> None: @@ -194,11 +252,11 @@ def test_random_requires_a_target() -> None: algorithms.random([]) -async def test_invalid_request_is_rejected_at_the_boundary() -> None: - algorithm = algorithms.random([LlmTarget("fast", EchoClient("fast"))]) +def test_invalid_request_is_rejected_at_the_boundary() -> None: + algorithm = algorithms.random(["fast"]) with pytest.raises(ValueError, match="unknown variant"): - await algorithm.run( + algorithm.run_stream( { "model": "auto", "messages": [{"role": "invalid", "content": []}], @@ -211,7 +269,7 @@ class FailingClient: async def call(self, request: dict[str, Any]) -> dict[str, Any]: raise RuntimeError("client failed") - algorithm = algorithms.random([LlmTarget("broken", FailingClient())]) + algorithm = algorithms.random(["broken"]) with pytest.raises(LibsyError, match="client failed"): - await algorithm.run(request_body()) + await run_algorithm(algorithm, {"broken": FailingClient()})