mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
import asyncio
|
|
import copy
|
|
from collections.abc import Awaitable, Callable, Generator, Mapping
|
|
from contextlib import contextmanager
|
|
from typing import Final
|
|
|
|
from litellm.rust_bridge import catalog
|
|
from litellm.rust_bridge.catalog import Route, RouteRule
|
|
from litellm.rust_bridge.configuration import Rollout
|
|
from tests.test_litellm_rust.support.isolation import rebound
|
|
from tests.test_litellm_rust.support.recording_server import RecordingServer
|
|
from tests.test_litellm_rust.support.served import assert_served_natively
|
|
|
|
__all__ = ["CALL_ID", "assert_served_natively", "call_public", "native_route", "routed"]
|
|
|
|
CALL_ID: Final = "callback-contract"
|
|
|
|
|
|
@contextmanager
|
|
def native_route(*routes: Route) -> Generator[None]:
|
|
with rebound(catalog, "RULES", (*(RouteRule(route, Rollout.RUST_REQUIRED) for route in routes), *catalog.RULES)):
|
|
yield
|
|
|
|
|
|
@contextmanager
|
|
def routed(*rules: RouteRule) -> Generator[None]:
|
|
with rebound(catalog, "RULES", rules):
|
|
yield
|
|
|
|
|
|
async def call_public(
|
|
create: Callable[..., object],
|
|
acreate: Callable[..., Awaitable[object]],
|
|
request: Mapping[str, object],
|
|
server: RecordingServer,
|
|
*,
|
|
asynchronous: bool,
|
|
**kwargs: object,
|
|
) -> object:
|
|
arguments: Final = {
|
|
**copy.deepcopy(dict(request)),
|
|
"api_key": "test-key",
|
|
"api_base": server.base_url,
|
|
"litellm_call_id": CALL_ID,
|
|
**kwargs,
|
|
}
|
|
if asynchronous:
|
|
return await acreate(**arguments)
|
|
return await asyncio.to_thread(create, **arguments)
|