feat(libsy): expose typed Python protocol bindings#88
Conversation
Signed-off-by: nachiketb <nachiketb@nvidia.com>
b383288 to
b7c5507
Compare
WalkthroughChangesAdds Rust/PyO3 bindings for the libsy protocol, typed Python exports, aggregate and streaming response wrappers, error registration, and tests covering serialization, variants, namespace behavior, and stream cleanup. Libsy Python Protocol
Estimated code review effort: 4 (Complex) | ~60 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (1)
tests/test_libsy_protocol_bindings.py (1)
181-199: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider using
AsyncGeneratorinstead ofAsyncIteratorfor_chunks.
self._chunks.aclose()is called on line 199, butAsyncIteratordoesn't declareaclose()— that method belongs toAsyncGenerator. Using the more precise type avoids type-checker errors and correctly documents the capability.♻️ Proposed type annotation fix
from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator class _ChunkSource: def __init__(self) -> None: - self._chunks: AsyncIterator[libsy.protocol.LlmResponseChunkValue] = self._iterate() + self._chunks: AsyncGenerator[libsy.protocol.LlmResponseChunkValue, None] = self._iterate() self.closed = False - async def _iterate(self) -> AsyncIterator[libsy.protocol.LlmResponseChunkValue]: + async def _iterate(self) -> AsyncGenerator[libsy.protocol.LlmResponseChunkValue, None]: yield libsy.protocol.LlmResponseChunk.MessageStart("response-1", "selected-model")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_libsy_protocol_bindings.py` around lines 181 - 199, Update the _ChunkSource._chunks annotation from AsyncIterator to AsyncGenerator, matching the _iterate() implementation and its use of aclose() in aclose(). Preserve the existing iteration behavior and generator element type.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/switchyard-py/src/libsy_bindings/protocol.rs`:
- Around line 38-53: Document the complete Python-visible Rust API: in
crates/switchyard-py/src/libsy_bindings/protocol.rs:38-53, add concise
item-level documentation for every exported protocol PyO3 class, including
PyWireFormat; in crates/switchyard-py/src/libsy_bindings/values.rs:29-38,
document every exported value wrapper; in
crates/switchyard-py/src/libsy_bindings/values.rs:375-440, document stream
single-consumption, closure, and concurrency invariants around the relevant
stream types and methods; and in crates/switchyard-py/src/errors.rs:13-13,
document the exception’s purpose and hierarchy.
In `@crates/switchyard-py/src/libsy_bindings/values.rs`:
- Around line 275-294: Update the is_streaming and selected_model getters to
detect the consumed response state after the stored LlmResponse has been taken
by stream, and return the existing “already consumed” PyRuntimeError instead of
false or None. Reuse the established consumed-state check and preserve the
current lock-poisoning behavior and normal response results.
- Around line 571-579: Update SourceClosingLlmStream to capture and store a
Tokio runtime Handle during construction, then use that stored handle to spawn
the stream-close future instead of calling Handle::try_current() in the
drop/completion path. Mirror core_bindings/response.rs and add the requested
warning fallback when no runtime handle was captured, ensuring the aclose
coroutine is polled whenever a handle is available.
In `@switchyard_rust/libsy_protocol.py`:
- Around line 12-44: Keep the lazy facade and type stub exports synchronized: in
switchyard_rust/libsy_protocol.py (lines 12-44), expose the five aliases
ImageSourceValue, FileSourceValue, MediaSourceValue, ContentBlockValue, and
LlmResponseChunkValue, or explicitly remove them from the public stub; in
switchyard_rust/libsy_protocol.pyi, align the declarations at lines 60, 79, 99,
162-173, and 415-423 with that choice so explicit imports remain type- and
runtime-correct.
In `@switchyard/__init__.py`:
- Line 14: Update switchyard/__init__.py to re-export every newly introduced
public protocol class alongside libsy, and add each class to the module’s
__all__ so they are available from the switchyard package root.
---
Nitpick comments:
In `@tests/test_libsy_protocol_bindings.py`:
- Around line 181-199: Update the _ChunkSource._chunks annotation from
AsyncIterator to AsyncGenerator, matching the _iterate() implementation and its
use of aclose() in aclose(). Preserve the existing iteration behavior and
generator element type.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 03b97e98-5b92-4b4b-a6cd-82f51ab3c070
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock,!Cargo.lock
📒 Files selected for processing (12)
crates/switchyard-py/Cargo.tomlcrates/switchyard-py/src/errors.rscrates/switchyard-py/src/lib.rscrates/switchyard-py/src/libsy_bindings.rscrates/switchyard-py/src/libsy_bindings/protocol.rscrates/switchyard-py/src/libsy_bindings/values.rsswitchyard/__init__.pyswitchyard/libsy/__init__.pyswitchyard/libsy/protocol.pyswitchyard_rust/libsy_protocol.pyswitchyard_rust/libsy_protocol.pyitests/test_libsy_protocol_bindings.py
| #[pyclass( | ||
| name = "WireFormat", | ||
| module = "switchyard.libsy.protocol", | ||
| eq, | ||
| frozen, | ||
| from_py_object | ||
| )] | ||
| #[derive(Clone, Copy, Debug, Eq, PartialEq)] | ||
| pub(crate) enum PyWireFormat { | ||
| #[pyo3(name = "OPENAI_CHAT")] | ||
| OpenAiChat, | ||
| #[pyo3(name = "ANTHROPIC_MESSAGES")] | ||
| AnthropicMessages, | ||
| #[pyo3(name = "OPENAI_RESPONSES")] | ||
| OpenAiResponses, | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
Document the complete Python-visible Rust API. The new binding layer has module documentation but lacks item-level behavior and invariant documentation.
crates/switchyard-py/src/libsy_bindings/protocol.rs#L38-L53: document every exported protocol PyO3 class.crates/switchyard-py/src/libsy_bindings/values.rs#L29-L38: document every exported value wrapper.crates/switchyard-py/src/libsy_bindings/values.rs#L375-L440: document stream single-consumption, closure, and concurrency invariants.crates/switchyard-py/src/errors.rs#L13-L13: document the exception’s purpose and hierarchy.
As per coding guidelines, public Rust items and complex async/lifecycle logic require concise documentation.
📍 Affects 3 files
crates/switchyard-py/src/libsy_bindings/protocol.rs#L38-L53(this comment)crates/switchyard-py/src/libsy_bindings/values.rs#L29-L38crates/switchyard-py/src/libsy_bindings/values.rs#L375-L440crates/switchyard-py/src/errors.rs#L13-L13
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/switchyard-py/src/libsy_bindings/protocol.rs` around lines 38 - 53,
Document the complete Python-visible Rust API: in
crates/switchyard-py/src/libsy_bindings/protocol.rs:38-53, add concise
item-level documentation for every exported protocol PyO3 class, including
PyWireFormat; in crates/switchyard-py/src/libsy_bindings/values.rs:29-38,
document every exported value wrapper; in
crates/switchyard-py/src/libsy_bindings/values.rs:375-440, document stream
single-consumption, closure, and concurrency invariants around the relevant
stream types and methods; and in crates/switchyard-py/src/errors.rs:13-13,
document the exception’s purpose and hierarchy.
Source: Coding guidelines
| #[getter] | ||
| fn is_streaming(&self) -> PyResult<bool> { | ||
| let response = self | ||
| .llm_response | ||
| .lock() | ||
| .map_err(|_| PyRuntimeError::new_err("Response lock is poisoned"))?; | ||
| Ok(matches!(response.as_ref(), Some(LlmResponse::Stream(_)))) | ||
| } | ||
|
|
||
| #[getter] | ||
| fn selected_model(&self) -> PyResult<Option<String>> { | ||
| let response = self | ||
| .llm_response | ||
| .lock() | ||
| .map_err(|_| PyRuntimeError::new_err("Response lock is poisoned"))?; | ||
| Ok(response | ||
| .as_ref() | ||
| .and_then(LlmResponse::selected_model) | ||
| .map(str::to_owned)) | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Do not silently misclassify a consumed response.
After .stream takes the value, is_streaming returns False and selected_model returns None. Both look valid but represent consumed state; raise the existing “already consumed” error instead.
Proposed fix
- Ok(matches!(response.as_ref(), Some(LlmResponse::Stream(_))))
+ match response.as_ref() {
+ Some(LlmResponse::Stream(_)) => Ok(true),
+ Some(LlmResponse::Agg(_)) => Ok(false),
+ None => Err(PyRuntimeError::new_err(
+ "Response has already been consumed",
+ )),
+ }Apply the equivalent consumed-state check before calling selected_model.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/switchyard-py/src/libsy_bindings/values.rs` around lines 275 - 294,
Update the is_streaming and selected_model getters to detect the consumed
response state after the stored LlmResponse has been taken by stream, and return
the existing “already consumed” PyRuntimeError instead of false or None. Reuse
the established consumed-state check and preserve the current lock-poisoning
behavior and normal response results.
| Ok(Some(future)) => { | ||
| if let Ok(runtime) = tokio::runtime::Handle::try_current() { | ||
| runtime.spawn(async move { | ||
| if let Err(error) = future.await { | ||
| tracing::warn!(error = %error, "libsy response stream close failed"); | ||
| } | ||
| }); | ||
| } | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm dependency versions and locate established configured-runtime patterns.
rg -n -C3 'pyo3-async-runtimes|tokio' Cargo.lock crates/switchyard-py/Cargo.toml
rg -n -C4 'get_runtime|Handle::try_current|future_into_py|into_future' crates/switchyard-py/srcRepository: NVIDIA-NeMo/Switchyard
Length of output: 45802
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline crates/switchyard-py/src/libsy_bindings/values.rs --view expanded
printf '\n--- relevant range ---\n'
sed -n '520,610p' crates/switchyard-py/src/libsy_bindings/values.rs
printf '\n--- surrounding types ---\n'
sed -n '1,220p' crates/switchyard-py/src/libsy_bindings/values.rsRepository: NVIDIA-NeMo/Switchyard
Length of output: 12602
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n 'runtime: Option<tokio::runtime::Handle>|Handle::try_current|runtime\.spawn|self\.runtime|try_current\(\)\.ok\(\)' crates/switchyard-py/src/libsy_bindings/values.rsRepository: NVIDIA-NeMo/Switchyard
Length of output: 284
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '250,290p' crates/switchyard-py/src/libsy_bindings/values.rsRepository: NVIDIA-NeMo/Switchyard
Length of output: 1497
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '220,290p' crates/switchyard-py/src/libsy_bindings/values.rsRepository: NVIDIA-NeMo/Switchyard
Length of output: 2315
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '618,710p' crates/switchyard-py/src/core_bindings/response.rsRepository: NVIDIA-NeMo/Switchyard
Length of output: 4010
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '440,500p' crates/switchyard-py/src/libsy_bindings/values.rsRepository: NVIDIA-NeMo/Switchyard
Length of output: 1734
Use a captured runtime handle for stream-source close. Handle::try_current() can fail when SourceClosingLlmStream is dropped or completes on a non-Tokio Python thread, so the aclose coroutine is created and then never polled, leaving the source open. Mirror core_bindings/response.rs by storing the runtime handle at construction and spawning there, with a warning fallback when no handle exists.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/switchyard-py/src/libsy_bindings/values.rs` around lines 571 - 579,
Update SourceClosingLlmStream to capture and store a Tokio runtime Handle during
construction, then use that stored handle to spawn the stream-close future
instead of calling Handle::try_current() in the drop/completion path. Mirror
core_bindings/response.rs and add the requested warning fallback when no runtime
handle was captured, ensuring the aclose coroutine is polled whenever a handle
is available.
| _EXPORTS = frozenset( | ||
| { | ||
| "AggLlmResponse", | ||
| "ContentBlock", | ||
| "Context", | ||
| "FileSource", | ||
| "FormatId", | ||
| "ImageSource", | ||
| "InstructionBlock", | ||
| "LlmRequest", | ||
| "LlmResponseChunk", | ||
| "LlmResponseStream", | ||
| "MediaSource", | ||
| "Message", | ||
| "Metadata", | ||
| "OutputParams", | ||
| "PreservationMetadata", | ||
| "ProviderExtensions", | ||
| "ReasoningParams", | ||
| "Request", | ||
| "Response", | ||
| "ResponseOutput", | ||
| "Role", | ||
| "SamplingParams", | ||
| "StopReason", | ||
| "ToolCall", | ||
| "ToolChoice", | ||
| "ToolDefinition", | ||
| "ToolResult", | ||
| "Usage", | ||
| "WireFormat", | ||
| } | ||
| ) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Keep the lazy facade and stub exports identical. Five aliases exist only in the stub, so type-correct explicit imports fail at runtime.
switchyard_rust/libsy_protocol.py#L12-L44: lazily expose the five aliases or intentionally exclude them from the public stub.switchyard_rust/libsy_protocol.pyi#L60-L60: alignImageSourceValue.switchyard_rust/libsy_protocol.pyi#L79-L79: alignFileSourceValue.switchyard_rust/libsy_protocol.pyi#L99-L99: alignMediaSourceValue.switchyard_rust/libsy_protocol.pyi#L162-L173: alignContentBlockValue.switchyard_rust/libsy_protocol.pyi#L415-L423: alignLlmResponseChunkValue.
📍 Affects 2 files
switchyard_rust/libsy_protocol.py#L12-L44(this comment)switchyard_rust/libsy_protocol.pyi#L60-L60switchyard_rust/libsy_protocol.pyi#L79-L79switchyard_rust/libsy_protocol.pyi#L99-L99switchyard_rust/libsy_protocol.pyi#L162-L173switchyard_rust/libsy_protocol.pyi#L415-L423
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@switchyard_rust/libsy_protocol.py` around lines 12 - 44, Keep the lazy facade
and type stub exports synchronized: in switchyard_rust/libsy_protocol.py (lines
12-44), expose the five aliases ImageSourceValue, FileSourceValue,
MediaSourceValue, ContentBlockValue, and LlmResponseChunkValue, or explicitly
remove them from the public stub; in switchyard_rust/libsy_protocol.pyi, align
the declarations at lines 60, 79, 99, 162-173, and 415-423 with that choice so
explicit imports remain type- and runtime-correct.
|
|
||
| from typing import TYPE_CHECKING, Any | ||
|
|
||
| from switchyard import libsy as libsy |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
Re-export the new public protocol classes.
Exporting only libsy leaves all newly introduced public classes absent from switchyard and its __all__.
As per coding guidelines, switchyard/__init__.py must export every new public class and include it in __all__.
Also applies to: 143-144
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@switchyard/__init__.py` at line 14, Update switchyard/__init__.py to
re-export every newly introduced public protocol class alongside libsy, and add
each class to the module’s __all__ so they are available from the switchyard
package root.
Source: Coding guidelines
Stack
Review and merge this PR first.
Why
Python hosts need typed access to the neutral request and response values already owned by switchyard-protocol. Untyped dictionaries would duplicate the Rust model and lose structural matching at the language boundary.
What
How
PyO3 wrappers own or clone the Rust values and convert only raw JSON extension fields. Complex enums use nested Python variants for structural pattern matching.
What to review
Validation
Summary by CodeRabbit
New Features
LibsyErrorexception for orchestration failures.Tests