From 1155db7b39a76a81f8629adafc9f681ae85c7253 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 3 Sep 2026 14:32:14 -0700 Subject: [PATCH] feat(rust-bridge): add a typed stub for the _native module --- basedpyright-code-budget.json | 4 +- litellm/rust_bridge/_native.pyi | 133 ++++++++++++++++++ .../rust_bridge/test_native_stub.py | 74 ++++++++++ 3 files changed, 209 insertions(+), 2 deletions(-) create mode 100644 litellm/rust_bridge/_native.pyi create mode 100644 tests/test_litellm/rust_bridge/test_native_stub.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index a37cb194757..a4424e42e21 100644 --- a/basedpyright-code-budget.json +++ b/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 diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi new file mode 100644 index 00000000000..14bdb83d007 --- /dev/null +++ b/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]: ... diff --git a/tests/test_litellm/rust_bridge/test_native_stub.py b/tests/test_litellm/rust_bridge/test_native_stub.py new file mode 100644 index 00000000000..2e1c6a7e413 --- /dev/null +++ b/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