diff --git a/crates/switchyard-py/src/libsy_bindings.rs b/crates/switchyard-py/src/libsy_bindings.rs index 8be1bede0..cf4bce07a 100644 --- a/crates/switchyard-py/src/libsy_bindings.rs +++ b/crates/switchyard-py/src/libsy_bindings.rs @@ -12,9 +12,10 @@ use pyo3::exceptions::{PyTypeError, PyValueError}; use pyo3::prelude::*; use serde_json::{Value, json}; use switchyard_libsy::{ - Algorithm, ClassifierContractConfig, HandoffNoteConfig, LibsyError as RustLibsyError, - LlmClassifierConfig, LlmFallback, LlmTarget, LlmTargetSet, LlmTaskClassifier, Noop, PickerMode, - Random, StageRouter, StageRouterConfig, TaskClassifierConfig, + Algorithm, ClassifierContractConfig, CustomClassifierConfig, CustomClassifierPolicy, + HandoffNoteConfig, LibsyError as RustLibsyError, LlmClassifierConfig, LlmFallback, LlmTarget, + LlmTargetSet, LlmTaskClassifier, Noop, PickerMode, Random, StageRouter, StageRouterConfig, + TaskClassifierConfig, }; use switchyard_llm_client::ClientRouter; use switchyard_protocol::{ @@ -182,6 +183,74 @@ impl PyTaskClassifierConfig { } } +/// Settings for a custom-schema classifier: a user-supplied prompt, an inner JSON +/// Schema for the verdict, and a JSON Pointer selecting the target label from it. +#[pyclass( + name = "CustomClassifierConfig", + module = "switchyard.libsy", + frozen, + skip_from_py_object +)] +struct PyCustomClassifierConfig { + prompt: String, + response_schema: Value, + selector: String, + session_affinity: bool, + message_hash_fallback: bool, + recent_turn_window: Option, + max_output_tokens: u64, +} + +impl PyCustomClassifierConfig { + fn clone_core(&self) -> CustomClassifierConfig { + CustomClassifierConfig { + prompt: self.prompt.clone(), + response_schema: self.response_schema.clone(), + policy: CustomClassifierPolicy::target_selector(self.selector.clone()), + session_affinity: self.session_affinity, + message_hash_fallback: self.message_hash_fallback, + recent_turn_window: self.recent_turn_window, + max_output_tokens: self.max_output_tokens, + } + } +} + +#[pymethods] +impl PyCustomClassifierConfig { + #[new] + #[pyo3(signature = ( + prompt, + response_schema, + selector, + *, + session_affinity=false, + message_hash_fallback=false, + recent_turn_window=None, + max_output_tokens=4096 + ))] + #[allow(clippy::too_many_arguments)] + fn new( + prompt: String, + response_schema: &Bound<'_, PyAny>, + selector: String, + session_affinity: bool, + message_hash_fallback: bool, + recent_turn_window: Option, + max_output_tokens: u64, + ) -> PyResult { + let response_schema: Value = from_python(response_schema)?; + Ok(Self { + prompt, + response_schema, + selector, + session_affinity, + message_hash_fallback, + recent_turn_window, + max_output_tokens, + }) + } +} + /// Judge target and policy used when stage-router signals are inconclusive. #[pyclass( name = "LlmFallback", @@ -369,6 +438,45 @@ fn llm_task_classifier_algorithm( Ok(PyAlgorithm::new(Arc::new(algorithm), clients)) } +/// Construct schema-driven classifier routing across two or more named targets. +/// +/// `targets` pairs each user-facing label with its routing target; the judge's +/// schema-validated verdict selects a label through the config's JSON Pointer, +/// and `default_target` is used when the judge does not produce a usable verdict. +#[pyfunction(name = "custom_classifier")] +#[pyo3(signature = ( + judge_target, + targets, + *, + default_target, + config +))] +fn custom_classifier_algorithm( + py: Python<'_>, + judge_target: Py, + targets: Vec<(String, Py)>, + default_target: String, + config: Py, +) -> PyResult { + let (labels, target_handles): (Vec, Vec>) = + targets.into_iter().unzip(); + let mut handles = vec![judge_target]; + handles.extend(target_handles); + let (cores, clients) = target_cores(py, &handles)?; + let mut cores = cores.into_iter(); + let judge = cores + .next() + .ok_or_else(|| PyValueError::new_err("expected a judge target"))?; + let algorithm = LlmTaskClassifier::new(LlmClassifierConfig::Custom { + judge_target: judge, + targets: labels.into_iter().zip(cores).collect(), + default_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)) +} + /// Construct signal-driven stage routing with an optional LLM classifier fallback. #[pyfunction(name = "stage_router")] #[pyo3(signature = ( @@ -467,6 +575,7 @@ fn invalid_python_response(error: PyErr) -> LlmClientError { 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::()?; @@ -476,6 +585,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { llm_task_classifier_algorithm, &libsy_module )?)?; + libsy_module.add_function(wrap_pyfunction!(custom_classifier_algorithm, &libsy_module)?)?; libsy_module.add_function(wrap_pyfunction!(stage_router_algorithm, &libsy_module)?)?; libsy_module.add("LibsyError", module.getattr("LibsyError")?)?; module.add_submodule(&libsy_module)?; diff --git a/switchyard/libsy/__init__.py b/switchyard/libsy/__init__.py index 44ca4cea3..d32a3513c 100644 --- a/switchyard/libsy/__init__.py +++ b/switchyard/libsy/__init__.py @@ -5,6 +5,7 @@ from switchyard_rust.libsy import ( Algorithm, + CustomClassifierConfig, LibsyError, LlmClient, LlmFallback, @@ -16,6 +17,7 @@ __all__ = [ "Algorithm", + "CustomClassifierConfig", "LibsyError", "LlmClient", "LlmFallback", diff --git a/switchyard/libsy/algorithms.py b/switchyard/libsy/algorithms.py index 92e9e747e..5f637bebd 100644 --- a/switchyard/libsy/algorithms.py +++ b/switchyard/libsy/algorithms.py @@ -3,9 +3,10 @@ """Factories for Rust-owned libsy algorithms.""" +from switchyard_rust.libsy import custom_classifier as custom_classifier from switchyard_rust.libsy import llm_task_classifier as llm_task_classifier from switchyard_rust.libsy import noop as noop from switchyard_rust.libsy import random as random from switchyard_rust.libsy import stage_router as stage_router -__all__ = ["llm_task_classifier", "noop", "random", "stage_router"] +__all__ = ["custom_classifier", "llm_task_classifier", "noop", "random", "stage_router"] diff --git a/switchyard_rust/libsy.py b/switchyard_rust/libsy.py index a7eb5c85a..13dc3badd 100644 --- a/switchyard_rust/libsy.py +++ b/switchyard_rust/libsy.py @@ -13,10 +13,12 @@ _EXPORTS = frozenset( { "Algorithm", + "CustomClassifierConfig", "LibsyError", "LlmFallback", "LlmTarget", "TaskClassifierConfig", + "custom_classifier", "llm_task_classifier", "noop", "random", @@ -63,6 +65,20 @@ def __init__( prompt: str | None = None, ) -> None: ... + @final + class CustomClassifierConfig: + def __init__( + self, + prompt: str, + response_schema: Mapping[str, object], + selector: str, + *, + session_affinity: bool = False, + message_hash_fallback: bool = False, + recent_turn_window: int | None = None, + max_output_tokens: int = 4096, + ) -> None: ... + @final class LlmFallback: def __init__( @@ -92,6 +108,14 @@ def llm_task_classifier( config: TaskClassifierConfig, ) -> Algorithm: ... + def custom_classifier( + judge_target: LlmTarget, + targets: Sequence[tuple[str, LlmTarget]], + *, + default_target: str, + config: CustomClassifierConfig, + ) -> Algorithm: ... + def stage_router( capable_target: LlmTarget, efficient_target: LlmTarget, diff --git a/tests/test_libsy_custom_classifier.py b/tests/test_libsy_custom_classifier.py new file mode 100644 index 000000000..602eb0a3a --- /dev/null +++ b/tests/test_libsy_custom_classifier.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the custom N-target classifier Python binding.""" + +from typing import Any + +import pytest + +from switchyard.libsy import CustomClassifierConfig, LlmTarget, algorithms + +LANES = ("grok", "luna", "flash", "sol", "opus") + +SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "required": ["model"], + "properties": {"model": {"type": "string", "enum": list(LANES)}}, +} + + +def request_body() -> dict[str, Any]: + return { + "model": "auto", + "messages": [ + { + "role": "user", + "content": [{"type": "text", "text": "summarize this document"}], + } + ], + } + + +class EchoClient: + def __init__(self, model: str) -> None: + self.model = model + self.calls: list[dict[str, Any]] = [] + + async def call(self, request: dict[str, Any]) -> dict[str, Any]: + self.calls.append(request) + return { + "model": self.model, + "outputs": [ + { + "role": "assistant", + "content": [{"type": "text", "text": self.model}], + "stop_reason": "end_turn", + } + ], + } + + +class VerdictClient(EchoClient): + def __init__(self, model: str, verdict: str) -> None: + super().__init__(model) + self.verdict = verdict + + async def call(self, request: dict[str, Any]) -> dict[str, Any]: + self.calls.append(request) + return { + "model": self.model, + "outputs": [ + { + "role": "assistant", + "content": [{"type": "text", "text": self.verdict}], + "stop_reason": "end_turn", + } + ], + } + + +class FailingClient: + async def call(self, request: dict[str, Any]) -> dict[str, Any]: + raise RuntimeError("judge unavailable") + + +def build(judge_client: Any, clients: dict[str, EchoClient]): + return algorithms.custom_classifier( + LlmTarget("judge", judge_client), + [(lane, LlmTarget(lane, clients[lane])) for lane in LANES], + default_target="grok", + config=CustomClassifierConfig( + "Pick the best lane.", + SCHEMA, + "/model", + ), + ) + + +def lane_clients() -> dict[str, EchoClient]: + return {lane: EchoClient(lane) for lane in LANES} + + +@pytest.mark.parametrize("lane", LANES) +async def test_verdict_routes_each_lane(lane: str) -> None: + clients = lane_clients() + judge = VerdictClient("judge", f'{{"model":"{lane}"}}') + algorithm = build(judge, clients) + + decisions, response = await algorithm.run(request_body()) + + assert response["model"] == lane + assert clients[lane].calls, "selected lane's client must serve the answer call" + assert decisions[-1]["selected_model"] == lane + + +async def test_judge_receives_prompt_and_inner_schema() -> None: + clients = lane_clients() + judge = VerdictClient("judge", '{"model":"luna"}') + algorithm = build(judge, clients) + + await algorithm.run(request_body()) + + judge_request = judge.calls[0] + assert judge_request["instructions"][0]["content"][0]["text"] == "Pick the best lane." + schema = judge_request["output"]["response_format"]["json_schema"]["schema"] + assert schema["properties"]["model"]["enum"] == list(LANES) + + +async def test_judge_failure_falls_open_to_the_default_target() -> None: + clients = lane_clients() + algorithm = build(FailingClient(), clients) + + _, response = await algorithm.run(request_body()) + + assert response["model"] == "grok" + + +async def test_unusable_verdict_falls_open_to_the_default_target() -> None: + clients = lane_clients() + judge = VerdictClient("judge", '{"model":"a-lane-that-does-not-exist"}') + algorithm = build(judge, clients) + + _, response = await algorithm.run(request_body()) + + assert response["model"] == "grok" + + +def test_default_target_must_be_a_configured_label() -> None: + clients = lane_clients() + with pytest.raises(ValueError, match="default_target"): + algorithms.custom_classifier( + LlmTarget("judge", EchoClient("judge")), + [(lane, LlmTarget(lane, clients[lane])) for lane in LANES], + default_target="not-a-lane", + config=CustomClassifierConfig("Pick.", SCHEMA, "/model"), + ) + + +def test_requires_at_least_two_targets() -> None: + with pytest.raises(ValueError, match="at least two targets"): + algorithms.custom_classifier( + LlmTarget("judge", EchoClient("judge")), + [("grok", LlmTarget("grok", EchoClient("grok")))], + default_target="grok", + config=CustomClassifierConfig("Pick.", SCHEMA, "/model"), + )