feat(rust-bridge): add a typed stub for the _native module

basedpyright-code-budget.json

@@ -9,7 +9,7 @@
     "limit": 319
   },
   "reportAttributeAccessIssue": {
-    "limit": 480
+    "limit": 479
   },
   "reportCallIssue": {
     "limit": 112
@@ -111,7 +111,7 @@
     "limit": 19624
   },
   "reportUnknownVariableType": {
-    "limit": 29861
+    "limit": 29858
   },
   "reportUnnecessaryCast": {
     "limit": 111

litellm/rust_bridge/_native.pyi

@@ -0,0 +1,133 @@
+"""Types for the compiled ``litellm.rust_bridge._native`` PyO3 extension.
+
+Keep in sync with the surface-pinning test in
+litellm-rust/crates/python-bridge/src/lib.rs and the ``__text_signature__``
+contract tests in litellm-rust/crates/python-bridge/src/routes/definition.rs.
+
+``_panic_for_test`` is only built with the ``panic-test`` cargo feature and is
+absent from default wheels, so it is intentionally not stubbed.
+"""
+
+from collections.abc import Coroutine, Mapping, Sequence
+
+__version__: str
+
+class RustBridgeDeclined(Exception):
+    """The route declined before calling the provider; the host may retry on its own path."""
+
+    def __init__(self, message: str) -> None: ...
+
+class RustUpstreamError(Exception):
+    """The provider call was issued and failed.
+
+    Args are ``(status, message)``; ``status`` is 0 when there was no HTTP response.
+    """
+
+    def __init__(self, status: int, message: str) -> None: ...
+
+class ResponsesWebSocketConnection:
+    @classmethod
+    def connect(
+        cls,
+        url: str,
+        headers: Mapping[str, str] | None = ...,
+        timeout_seconds: float | None = ...,
+    ) -> Coroutine[object, object, ResponsesWebSocketConnection]: ...
+    def send_text(self, text: str) -> Coroutine[object, object, None]: ...
+    def recv_text(self) -> Coroutine[object, object, str | None]: ...
+    def close(self) -> Coroutine[object, object, None]: ...
+
+def ocr(
+    model: str,
+    document: dict[str, object],
+    api_key: str | None = ...,
+    api_base: str | None = ...,
+    custom_llm_provider: str | None = ...,
+    extra_headers: Mapping[str, object] | None = ...,
+    optional_params: Mapping[str, object] | None = ...,
+    timeout_seconds: float | None = ...,
+    trace: bool = ...,
+) -> dict[str, object]: ...
+def aocr(
+    model: str,
+    document: dict[str, object],
+    api_key: str | None = ...,
+    api_base: str | None = ...,
+    custom_llm_provider: str | None = ...,
+    extra_headers: Mapping[str, object] | None = ...,
+    optional_params: Mapping[str, object] | None = ...,
+    timeout_seconds: float | None = ...,
+    trace: bool = ...,
+) -> Coroutine[object, object, dict[str, object]]: ...
+def transcription(
+    model: str,
+    audio: dict[str, object],
+    api_key: str | None = ...,
+    api_base: str | None = ...,
+    custom_llm_provider: str | None = ...,
+    extra_headers: Mapping[str, object] | None = ...,
+    optional_params: Mapping[str, object] | None = ...,
+    timeout_seconds: float | None = ...,
+    trace: bool = ...,
+) -> dict[str, object]: ...
+def atranscription(
+    model: str,
+    audio: dict[str, object],
+    api_key: str | None = ...,
+    api_base: str | None = ...,
+    custom_llm_provider: str | None = ...,
+    extra_headers: Mapping[str, object] | None = ...,
+    optional_params: Mapping[str, object] | None = ...,
+    timeout_seconds: float | None = ...,
+    trace: bool = ...,
+) -> Coroutine[object, object, dict[str, object]]: ...
+def messages(
+    model: str,
+    body: dict[str, object],
+    api_key: str | None = ...,
+    api_base: str | None = ...,
+    custom_llm_provider: str | None = ...,
+    extra_headers: Mapping[str, object] | None = ...,
+    timeout_seconds: float | None = ...,
+    trace: bool = ...,
+) -> dict[str, object]: ...
+def amessages(
+    model: str,
+    body: dict[str, object],
+    api_key: str | None = ...,
+    api_base: str | None = ...,
+    custom_llm_provider: str | None = ...,
+    extra_headers: Mapping[str, object] | None = ...,
+    timeout_seconds: float | None = ...,
+    trace: bool = ...,
+) -> Coroutine[object, object, dict[str, object]]: ...
+def chat_completions_decline(
+    model: str,
+    messages: Sequence[object],
+    optional_params: Mapping[str, object] | None = ...,
+    custom_llm_provider: str | None = ...,
+) -> str | None: ...
+def chat_completions(
+    model: str,
+    messages: Sequence[object],
+    optional_params: Mapping[str, object] | None = ...,
+    api_key: str | None = ...,
+    api_base: str | None = ...,
+    custom_llm_provider: str | None = ...,
+    extra_headers: Mapping[str, object] | None = ...,
+    timeout_seconds: float | None = ...,
+    trace: bool = ...,
+) -> dict[str, object]: ...
+def achat_completions(
+    model: str,
+    messages: Sequence[object],
+    optional_params: Mapping[str, object] | None = ...,
+    api_key: str | None = ...,
+    api_base: str | None = ...,
+    custom_llm_provider: str | None = ...,
+    extra_headers: Mapping[str, object] | None = ...,
+    timeout_seconds: float | None = ...,
+    trace: bool = ...,
+) -> Coroutine[object, object, dict[str, object]]: ...
+def gil_stats() -> dict[str, object]: ...
+def build_info() -> dict[str, object]: ...

tests/test_litellm/rust_bridge/test_native_stub.py

@@ -0,0 +1,74 @@
+"""Pin the ``_native.pyi`` stub to the compiled module's public surface."""
+
+from __future__ import annotations
+
+import ast
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Final
+
+import pytest
+
+import litellm.rust_bridge
+
+# keep in sync with crates/python-bridge/src/lib.rs surface test
+PINNED_SURFACE: Final = (
+    "RustBridgeDeclined",
+    "RustUpstreamError",
+    "ocr",
+    "aocr",
+    "transcription",
+    "atranscription",
+    "messages",
+    "amessages",
+    "chat_completions_decline",
+    "chat_completions",
+    "achat_completions",
+    "ResponsesWebSocketConnection",
+    "gil_stats",
+    "build_info",
+)
+
+
+def _stub_path() -> Path:
+    module_file: Final = litellm.rust_bridge.__file__
+    assert module_file is not None
+    return Path(module_file).with_name("_native.pyi")
+
+
+def _parse_stub() -> ast.Module:
+    return ast.parse(_stub_path().read_text(encoding="utf-8"), filename=str(_stub_path()))
+
+
+def _module_level_names(module: ast.Module) -> frozenset[str]:
+    def names(body: list[ast.stmt]) -> Iterator[str]:
+        for node in body:
+            if isinstance(node, ast.ClassDef | ast.FunctionDef | ast.AsyncFunctionDef):
+                yield node.name
+            elif isinstance(node, ast.Assign):
+                for target in node.targets:
+                    if isinstance(target, ast.Name):
+                        yield target.id
+            elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
+                yield node.target.id
+
+    return frozenset(names(module.body))
+
+
+def test_stub_pins_the_native_surface() -> None:
+    stub_names: Final = _module_level_names(_parse_stub())
+    public_names: Final = {name for name in stub_names if not name.startswith("__")}
+
+    assert public_names == set(PINNED_SURFACE)
+    assert "__version__" in stub_names
+
+
+def test_stub_names_exist_on_the_compiled_module() -> None:
+    try:
+        from litellm.rust_bridge import _native
+    except ImportError:
+        pytest.skip("compiled _native extension is not importable in this environment")
+
+    runtime_names: Final = frozenset(dir(_native))
+
+    assert _module_level_names(_parse_stub()) <= runtime_names