mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
105 lines
4 KiB
Python
105 lines
4 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Awaitable, Callable, Mapping
|
|
from dataclasses import dataclass
|
|
from typing import Final, Generic, TypeAlias, TypeVar
|
|
|
|
from litellm.rust_bridge import catalog, runtime
|
|
from litellm.rust_bridge.bindings import NativeBinding
|
|
from litellm.rust_bridge.catalog import Route, RouteContext, Rules
|
|
|
|
RequestT: Final = TypeVar("RequestT")
|
|
NativeT: Final = TypeVar("NativeT")
|
|
ResultT: Final = TypeVar("ResultT")
|
|
|
|
NativeHook: TypeAlias = Callable[[RequestT, tuple[object, ...], Mapping[str, object]], ResultT]
|
|
|
|
|
|
def call_hook(
|
|
hook: NativeHook[RequestT, ResultT],
|
|
request: RequestT,
|
|
args: tuple[object, ...],
|
|
kwargs: Mapping[str, object],
|
|
) -> ResultT:
|
|
return hook(request, args, kwargs)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class PublicDispatch(Generic[RequestT]):
|
|
route: Route
|
|
request: Callable[[tuple[object, ...], Mapping[str, object]], RequestT | None]
|
|
context: Callable[[RequestT], RouteContext]
|
|
|
|
def _native_request(self, args: tuple[object, ...], kwargs: Mapping[str, object]) -> RequestT:
|
|
request: Final = self.request(args, kwargs)
|
|
if request is None:
|
|
raise runtime.NoPythonImplementationError(
|
|
f"{self.route.value} has no Python implementation, so every call must project to a native request"
|
|
)
|
|
return request
|
|
|
|
def run(
|
|
self,
|
|
args: tuple[object, ...],
|
|
kwargs: Mapping[str, object],
|
|
*,
|
|
python: Callable[..., ResultT] | runtime.NoPythonImplementation,
|
|
binding: NativeBinding[NativeT],
|
|
native: Callable[[NativeT, RequestT, tuple[object, ...], Mapping[str, object]], ResultT],
|
|
rules: Rules | None = None,
|
|
) -> ResultT:
|
|
selected_rules: Final = catalog.RULES if rules is None else rules
|
|
if isinstance(python, runtime.NoPythonImplementation):
|
|
native_request: Final = self._native_request(args, kwargs)
|
|
return runtime.run(
|
|
self.context(native_request),
|
|
binding=binding,
|
|
native=lambda hook: native(hook, native_request, args, kwargs),
|
|
python=python,
|
|
rules=selected_rules,
|
|
)
|
|
if not catalog.may_leave_python(self.route, selected_rules):
|
|
return python(*args, **kwargs)
|
|
request: Final = self.request(args, kwargs)
|
|
if request is None:
|
|
return python(*args, **kwargs)
|
|
return runtime.run(
|
|
self.context(request),
|
|
binding=binding,
|
|
native=lambda hook: native(hook, request, args, kwargs),
|
|
python=lambda: python(*args, **kwargs),
|
|
rules=selected_rules,
|
|
)
|
|
|
|
async def arun(
|
|
self,
|
|
args: tuple[object, ...],
|
|
kwargs: Mapping[str, object],
|
|
*,
|
|
python: Callable[..., Awaitable[ResultT]] | runtime.NoPythonImplementation,
|
|
binding: NativeBinding[NativeT],
|
|
native: Callable[[NativeT, RequestT, tuple[object, ...], Mapping[str, object]], Awaitable[ResultT]],
|
|
rules: Rules | None = None,
|
|
) -> ResultT:
|
|
selected_rules: Final = catalog.RULES if rules is None else rules
|
|
if isinstance(python, runtime.NoPythonImplementation):
|
|
native_request: Final = self._native_request(args, kwargs)
|
|
return await runtime.arun(
|
|
self.context(native_request),
|
|
binding=binding,
|
|
native=lambda hook: native(hook, native_request, args, kwargs),
|
|
python=python,
|
|
rules=selected_rules,
|
|
)
|
|
if not catalog.may_leave_python(self.route, selected_rules):
|
|
return await python(*args, **kwargs)
|
|
request: Final = self.request(args, kwargs)
|
|
if request is None:
|
|
return await python(*args, **kwargs)
|
|
return await runtime.arun(
|
|
self.context(request),
|
|
binding=binding,
|
|
native=lambda hook: native(hook, request, args, kwargs),
|
|
python=lambda: python(*args, **kwargs),
|
|
rules=selected_rules,
|
|
)
|