mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-23 00:41:40 +00:00
64f2a3d098
Move each route's bridge module under litellm/rust_bridge/<route>/ so a folder means a Rust implementation exists while the catalog row says whether it is used. OCR now keeps the Python implementation in litellm/ocr/main.py and the Rust selection in litellm/ocr/rust.py, removing litellm/ocr/legacy.py Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
88 lines
2.4 KiB
Python
88 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from litellm.rust_bridge import configuration
|
|
from litellm.rust_bridge.responses import websocket as responses_websocket
|
|
|
|
|
|
class _FakeNativeConnection:
|
|
def __init__(self) -> None:
|
|
self.sent: list[str] = []
|
|
self.closed = False
|
|
|
|
async def send_text(self, text: str) -> None:
|
|
self.sent.append(text)
|
|
|
|
async def recv_text(self) -> str:
|
|
return "response.completed"
|
|
|
|
async def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class _ClosedNativeConnection:
|
|
async def recv_text(self) -> None:
|
|
return None
|
|
|
|
|
|
class _FakeNativeBridge:
|
|
@classmethod
|
|
async def connect(
|
|
cls,
|
|
*,
|
|
url: str,
|
|
headers: dict[str, str],
|
|
timeout_seconds: float | None,
|
|
) -> _FakeNativeConnection:
|
|
return _FakeNativeConnection()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_responses_websocket():
|
|
responses_websocket.set_rust_responses_websocket(connection=None)
|
|
configuration.reset_rust_configuration()
|
|
yield
|
|
responses_websocket.set_rust_responses_websocket(connection=None)
|
|
configuration.reset_rust_configuration()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None:
|
|
adapter = responses_websocket._ConnectionAdapter(_ClosedNativeConnection())
|
|
|
|
with pytest.raises(responses_websocket.ConnectionClosedOK):
|
|
await adapter.recv()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(responses_websocket, "_STATE", responses_websocket._RustResponsesWebSocketState())
|
|
monkeypatch.setattr(responses_websocket, "get_native_bridge", lambda: None)
|
|
|
|
assert (
|
|
await responses_websocket.connect(
|
|
url="wss://example.test/responses",
|
|
headers={},
|
|
timeout=None,
|
|
)
|
|
is None
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enabled_bridge_connects_and_adapts_socket(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
responses_websocket.set_rust_responses_websocket(connection=_FakeNativeBridge)
|
|
|
|
connection = await responses_websocket.connect(
|
|
url="wss://example.test/responses",
|
|
headers={"Authorization": "Bearer key"},
|
|
timeout=1.0,
|
|
)
|
|
|
|
assert connection is not None
|
|
await connection.send("response.create")
|
|
assert await connection.recv() == "response.completed"
|
|
await connection.close()
|