diff --git a/litellm/extensions/__init__.py b/litellm/extensions/__init__.py new file mode 100644 index 00000000000..aaed9ea1250 --- /dev/null +++ b/litellm/extensions/__init__.py @@ -0,0 +1,5 @@ +"""Client-side adapters for the external Python extension host.""" + +from .runtime import configure_extension_runtime, get_extension_runtime + +__all__ = ("configure_extension_runtime", "get_extension_runtime") diff --git a/litellm/extensions/adapters.py b/litellm/extensions/adapters.py new file mode 100644 index 00000000000..c44b2122e4c --- /dev/null +++ b/litellm/extensions/adapters.py @@ -0,0 +1,467 @@ +# pyright: reportAny=false, reportUnknownArgumentType=false, reportUnknownMemberType=false +# pyright: reportUnknownParameterType=false, reportUnknownVariableType=false +# pyright: reportMissingParameterType=false +from __future__ import annotations + +import hashlib +import time +from collections.abc import AsyncGenerator, AsyncIterator, Mapping +from datetime import datetime +from typing import Final + +from litellm.caching import DualCache +from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.python_extension.generated.v1 import extension_host_pb2 as pb +from litellm.python_extension_host.compatibility import decode_json, encode_json, response_from_json +from litellm.types.guardrails import GuardrailEventHooks, Mode +from litellm.types.utils import CallTypesLiteral, ModelResponseStream + +from .cache_gateway import InvocationCacheRegistry +from .client import PythonExtensionClient + + +class _RemoteHooks: + def __init__( + self, + client: PythonExtensionClient, + plugin_id: str, + cache_registry: InvocationCacheRegistry, + ) -> None: + self._extension_client: Final = client + self._extension_plugin_id: Final = plugin_id + self._extension_cache_registry: Final = cache_registry + + async def execute( + self, + phase: pb.HookPhase, + data: dict[str, object], # mutable-ok: LiteLLM compatibility payload + auth: UserAPIKeyAuth, + call_type: str, + cache: object | None = None, + response: object | None = None, + ) -> object | None: + context = _invocation_context( # rebind-ok: invocation-scoped RPC state + data, self._extension_client.manifest.revision_id, call_type + ) # rebind-ok: invocation-scoped RPC state + cache_ref: Final = self._extension_cache_registry.register(context.invocation_id, cache) + invocation: Final = pb.GuardrailInvocation( + context=context, + plugin_id=self._extension_plugin_id, + hook_phase=phase, + request_json=encode_json(data), + auth=_auth_context(auth), + ) + if response is not None: + invocation.response_json = encode_json(response) + if cache_ref is not None: + invocation.cache.CopyFrom(cache_ref) + try: + result = await self._extension_client.execute_guardrail( # rebind-ok: invocation-scoped RPC state + invocation + ) # rebind-ok: invocation-scoped RPC state + finally: + self._extension_cache_registry.revoke(cache_ref) + if result.decision in ( + pb.GUARDRAIL_DECISION_ALLOW, + pb.GUARDRAIL_DECISION_ERROR, + pb.GUARDRAIL_DECISION_UNSPECIFIED, + ): + return response if phase == pb.HOOK_PHASE_POST_CALL else None + if result.decision == pb.GUARDRAIL_DECISION_BLOCK: + public_error: Final = result.public_error if result.HasField("public_error") else None + raise GuardrailRaisedException( + guardrail_name=self._extension_plugin_id, + message=public_error.message if public_error is not None else "request blocked by extension", + should_wrap_with_default_message=False, + status_code=( + public_error.status_code + if public_error is not None and public_error.HasField("status_code") + else 400 + ), + blocked_content=True, + ) + if result.decision == pb.GUARDRAIL_DECISION_REPLACE_REQUEST and result.HasField("request_json"): + replacement: Final = decode_json(result.request_json, "request_json") + if isinstance(replacement, dict): + data.clear() + data.update(replacement) + return data + if result.decision == pb.GUARDRAIL_DECISION_REPLACE_RESPONSE and result.HasField("response_json"): + return response_from_json(result.response_json) + return response if phase == pb.HOOK_PHASE_POST_CALL else None + + def enqueue_event( + self, + kind: pb.CallbackEventKind, + kwargs: dict[str, object], # mutable-ok: LiteLLM compatibility payload + response: object, + start_time: object, + end_time: object, + streaming: bool = False, + ) -> None: + context = _invocation_context( # rebind-ok: invocation-scoped RPC state + kwargs, + self._extension_client.manifest.revision_id, + str(kwargs.get("call_type", "unknown")), + ) + payload: Final = kwargs.get("standard_logging_payload", kwargs) + event: Final = pb.CallbackEvent( + context=context, + plugin_id=self._extension_plugin_id, + kind=kind, + standard_logging_payload_json=encode_json(payload), + start_time_seconds=_timestamp(start_time), + end_time_seconds=_timestamp(end_time), + auth=_auth_context_from_kwargs(kwargs), + streaming=streaming, + ) + if response is not None and not isinstance(response, Exception): + event.response_json = encode_json(response) + if isinstance(response, Exception): + event.error_json = encode_json( + {"type": type(response).__name__, "message": str(response)} # mutable-ok: LiteLLM compatibility payload + ) # mutable-ok: LiteLLM compatibility payload + self._extension_client.enqueue_callback(event) + + async def transform_iterator( + self, + auth: UserAPIKeyAuth, + response: AsyncIterator[ModelResponseStream], + request_data: dict[str, object], # mutable-ok: LiteLLM compatibility payload + ) -> AsyncGenerator[ModelResponseStream, None]: + context = _invocation_context( # rebind-ok: invocation-scoped RPC state + request_data, + self._extension_client.manifest.revision_id, + str(request_data.get("call_type", "stream")), + ) + pending: list[ # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + ModelResponseStream + ] = [] # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + + async def frames() -> AsyncIterator[pb.StreamFrame]: + yield pb.StreamFrame( + kind=pb.STREAM_FRAME_KIND_OPEN, + stream_id=context.invocation_id, + open=pb.StreamOpen( + context=context, + plugin_id=self._extension_plugin_id, + request_json=encode_json(request_data), + auth=_auth_context(auth), + iterator_hook=( + "async_post_call_streaming_iterator_hook" + in self._extension_client.descriptor_hooks(self._extension_plugin_id) + ), + ), + ) + async for chunk in response: + pending.append(chunk) + yield pb.StreamFrame( + kind=pb.STREAM_FRAME_KIND_INPUT_CHUNK, + stream_id=context.invocation_id, + chunk_json=encode_json(chunk), + ) + yield pb.StreamFrame(kind=pb.STREAM_FRAME_KIND_END, stream_id=context.invocation_id) + + completed = False # rebind-ok: invocation-scoped RPC state + async for frame in self._extension_client.transform_stream(frames()): + if frame.kind == pb.STREAM_FRAME_KIND_OUTPUT_CHUNK and frame.HasField("chunk_json"): + transformed = response_from_json(frame.chunk_json) + if isinstance(transformed, ModelResponseStream): + yield transformed + elif frame.kind == pb.STREAM_FRAME_KIND_END: + completed = True + elif frame.kind == pb.STREAM_FRAME_KIND_ERROR: + break + if not completed: + for chunk in pending: + yield chunk + async for chunk in response: + yield chunk + + +class RemoteCustomLogger(CustomLogger): + def __init__( + self, + client: PythonExtensionClient, + plugin_id: str, + cache_registry: InvocationCacheRegistry, + ) -> None: + super().__init__() + self._remote = _RemoteHooks(client, plugin_id, cache_registry) + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], # mutable-ok: LiteLLM compatibility payload + call_type: CallTypesLiteral, + ) -> Exception | str | dict[str, object] | None: # mutable-ok: LiteLLM compatibility payload + result = await self._remote.execute( # rebind-ok: invocation-scoped RPC state + pb.HOOK_PHASE_PRE_CALL, + data, + user_api_key_dict, + call_type, + cache=cache, + ) + return result if isinstance(result, Exception | str | dict) else None + + async def async_moderation_hook( + self, + data: dict[str, object], # mutable-ok: LiteLLM compatibility payload + user_api_key_dict: UserAPIKeyAuth, + call_type: CallTypesLiteral, + ) -> object | None: + return await self._remote.execute(pb.HOOK_PHASE_DURING_CALL, data, user_api_key_dict, call_type) + + async def async_post_call_success_hook( + self, + data: dict[str, object], # mutable-ok: LiteLLM compatibility payload + user_api_key_dict: UserAPIKeyAuth, + response: object, + ) -> object | None: + return await self._remote.execute( + pb.HOOK_PHASE_POST_CALL, + data, + user_api_key_dict, + str(data.get("call_type", "unknown")), + response=response, + ) + + async def async_log_success_event( + self, + kwargs: dict[str, object], # mutable-ok: LiteLLM compatibility payload + response_obj: object, + start_time: object, + end_time: object, + ) -> None: + self._remote.enqueue_event(pb.CALLBACK_EVENT_KIND_SUCCESS, kwargs, response_obj, start_time, end_time) + + async def async_log_failure_event( + self, + kwargs: dict[str, object], # mutable-ok: LiteLLM compatibility payload + response_obj: object, + start_time: object, + end_time: object, + ) -> None: + self._remote.enqueue_event(pb.CALLBACK_EVENT_KIND_FAILURE, kwargs, response_obj, start_time, end_time) + + async def async_log_stream_event( + self, + kwargs: dict[str, object], # mutable-ok: LiteLLM compatibility payload + response_obj: object, + start_time: object, + end_time: object, + ) -> None: + self._remote.enqueue_event( + pb.CALLBACK_EVENT_KIND_SUCCESS, + kwargs, + response_obj, + start_time, + end_time, + streaming=True, + ) + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: AsyncIterator[ModelResponseStream], + request_data: dict[str, object], # mutable-ok: LiteLLM compatibility payload + ) -> AsyncGenerator[ModelResponseStream, None]: + async for chunk in self._remote.transform_iterator(user_api_key_dict, response, request_data): + yield chunk + + +class RemoteCustomGuardrail(CustomGuardrail): + def __init__( + self, + client: PythonExtensionClient, + plugin_id: str, + cache_registry: InvocationCacheRegistry, + guardrail_name: str | None = None, + event_hook: GuardrailEventHooks # mutable-ok: LiteLLM compatibility payload + | list[GuardrailEventHooks] + | Mode + | None = None, # mutable-ok: LiteLLM compatibility payload + default_on: bool = False, + extra_params: Mapping[str, object] | None = None, + ) -> None: + params = ( # rebind-ok: invocation-scoped RPC state + extra_params or {} # mutable-ok: LiteLLM compatibility payload + ) # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + super().__init__( + guardrail_name=guardrail_name, + event_hook=event_hook, + default_on=default_on, + mask_request_content=_bool_param(params, "mask_request_content", False), + mask_response_content=_bool_param(params, "mask_response_content", False), + violation_message_template=_str_param(params, "violation_message_template"), + end_session_after_n_fails=_int_param(params, "end_session_after_n_fails"), + on_violation=_str_param(params, "on_violation"), + realtime_violation_message=_str_param(params, "realtime_violation_message"), + on_sensitive_data=_str_param(params, "on_sensitive_data"), + sensitive_data_route_to_model=_str_param(params, "sensitive_data_route_to_model"), + sticky_session_routing=_bool_param(params, "sticky_session_routing", True), + run_in_parallel=_bool_param(params, "run_in_parallel", False), + scan_raw_request=_bool_param(params, "scan_raw_request", False), + only_scan_new_messages=_bool_param(params, "only_scan_new_messages", False), + turn_off_message_logging=_bool_param(params, "turn_off_message_logging", False), + message_logging=_bool_param(params, "message_logging", True), + ) + self._remote = _RemoteHooks(client, plugin_id, cache_registry) + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], # mutable-ok: LiteLLM compatibility payload + call_type: CallTypesLiteral, + ) -> Exception | str | dict[str, object] | None: # mutable-ok: LiteLLM compatibility payload + result = await self._remote.execute( # rebind-ok: invocation-scoped RPC state + pb.HOOK_PHASE_PRE_CALL, + data, + user_api_key_dict, + call_type, + cache=cache, + ) + return result if isinstance(result, Exception | str | dict) else None + + async def async_moderation_hook( + self, + data: dict[str, object], # mutable-ok: LiteLLM compatibility payload + user_api_key_dict: UserAPIKeyAuth, + call_type: CallTypesLiteral, + ) -> object | None: + return await self._remote.execute(pb.HOOK_PHASE_DURING_CALL, data, user_api_key_dict, call_type) + + async def async_post_call_success_hook( + self, + data: dict[str, object], # mutable-ok: LiteLLM compatibility payload + user_api_key_dict: UserAPIKeyAuth, + response: object, + ) -> object | None: + return await self._remote.execute( + pb.HOOK_PHASE_POST_CALL, + data, + user_api_key_dict, + str(data.get("call_type", "unknown")), + response=response, + ) + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: AsyncIterator[ModelResponseStream], + request_data: dict[str, object], # mutable-ok: LiteLLM compatibility payload + ) -> AsyncGenerator[ModelResponseStream, None]: + async for chunk in self._remote.transform_iterator(user_api_key_dict, response, request_data): + yield chunk + + +def _invocation_context(data: Mapping[str, object], revision: str, call_type: str) -> pb.InvocationContext: + metadata = data.get("metadata") # rebind-ok: invocation-scoped RPC state + metadata_mapping = ( # rebind-ok: invocation-scoped RPC state + metadata if isinstance(metadata, Mapping) else {} # mutable-ok: LiteLLM compatibility payload + ) # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + request_id: Final = str( + data.get("request_id") or data.get("litellm_call_id") or metadata_mapping.get("request_id") or "" + ) + invocation_id: Final = str( + data.get("litellm_call_id") or request_id or hashlib.sha256(str(time.time_ns()).encode()).hexdigest()[:24] + ) + return pb.InvocationContext( + request_id=request_id, + invocation_id=invocation_id, + active_revision=revision, + api_surface=str(data.get("api_surface", data.get("call_type", "unknown"))), + call_type=call_type, + trace_context=_trace_context(metadata_mapping), + ) + + +def _auth_context(auth: UserAPIKeyAuth) -> pb.AuthContext: + key = getattr(auth, "api_key", None) or getattr(auth, "token", None) or "" # rebind-ok: invocation-scoped RPC state + metadata = getattr(auth, "metadata", None) # rebind-ok: invocation-scoped RPC state + return pb.AuthContext( + key_hash=_hash_key(str(key)) if key else "", + user_id=str(getattr(auth, "user_id", None) or ""), + team_id=str(getattr(auth, "team_id", None) or ""), + request_metadata=_safe_metadata( + metadata if isinstance(metadata, Mapping) else {} # mutable-ok: LiteLLM compatibility payload + ), # mutable-ok: LiteLLM compatibility payload + ) + + +def _auth_context_from_kwargs(kwargs: Mapping[str, object]) -> pb.AuthContext: + litellm_params: Final = kwargs.get("litellm_params") + params = ( # rebind-ok: invocation-scoped RPC state + litellm_params if isinstance(litellm_params, Mapping) else {} # mutable-ok: LiteLLM compatibility payload + ) # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + metadata = params.get("metadata") # rebind-ok: invocation-scoped RPC state + metadata_mapping = ( # rebind-ok: invocation-scoped RPC state + metadata if isinstance(metadata, Mapping) else {} # mutable-ok: LiteLLM compatibility payload + ) # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + key = params.get("api_key") or metadata_mapping.get("user_api_key") or "" # rebind-ok: invocation-scoped RPC state + return pb.AuthContext( + key_hash=_hash_key(str(key)) if key else "", + user_id=str(metadata_mapping.get("user_api_key_user_id") or ""), + team_id=str(metadata_mapping.get("user_api_key_team_id") or ""), + request_metadata=_safe_metadata(metadata_mapping), + ) + + +def _hash_key(value: str) -> str: + return hashlib.sha256(value.encode()).hexdigest() + + +def _safe_metadata( + metadata: Mapping[object, object], +) -> dict[str, str]: # mutable-ok: LiteLLM compatibility payload + denied: Final = ( + "authorization", + "api_key", + "token", + "cookie", + "secret", + "password", + ) + return { # mutable-ok: LiteLLM compatibility payload + str(key): str(value) + for key, value in metadata.items() + if not any(part in str(key).lower() for part in denied) and isinstance(value, str | int | float | bool) + } + + +def _trace_context( + metadata: Mapping[object, object], +) -> dict[str, str]: # mutable-ok: LiteLLM compatibility payload + return { # mutable-ok: LiteLLM compatibility payload + name: str(metadata[name]) + for name in ("traceparent", "tracestate") + if name in metadata and isinstance(metadata[name], str) + } + + +def _timestamp(value: object) -> float: + if isinstance(value, datetime): + return float(value.timestamp()) + if isinstance(value, int | float): + return float(value) + return time.time() + + +def _bool_param(params: Mapping[str, object], key: str, default: bool) -> bool: + value = params.get(key) # rebind-ok: invocation-scoped RPC state + return value if isinstance(value, bool) else default + + +def _str_param(params: Mapping[str, object], key: str) -> str | None: + value = params.get(key) # rebind-ok: invocation-scoped RPC state + return value if isinstance(value, str) else None + + +def _int_param(params: Mapping[str, object], key: str) -> int | None: + value = params.get(key) # rebind-ok: invocation-scoped RPC state + return value if isinstance(value, int) and not isinstance(value, bool) else None diff --git a/litellm/extensions/cache_gateway.py b/litellm/extensions/cache_gateway.py new file mode 100644 index 00000000000..0d99a27c928 --- /dev/null +++ b/litellm/extensions/cache_gateway.py @@ -0,0 +1,115 @@ +# pyright: reportMissingModuleSource=false, reportMissingTypeArgument=false +# pyright: reportUnknownMemberType=false, reportUnknownParameterType=false, reportUnknownVariableType=false +from __future__ import annotations + +import json +import secrets +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Final, Protocol, cast # noqa: TID251 # validated cache/protobuf boundaries below + +import grpc + +from litellm.python_extension.generated.v1 import extension_host_pb2 as pb +from litellm.python_extension.generated.v1 import extension_host_pb2_grpc as pb_grpc +from litellm.python_extension_host.constants import TOKEN_METADATA_KEY + + +class _AsyncCache(Protocol): + async def async_get_cache( + self, + key: str, + local_only: bool = False, + **kwargs: object, # kwargs-ok: LiteLLM callback compatibility + ) -> object | None: ... + + async def async_set_cache( + self, + key: str, + value: object, + **kwargs: object, # kwargs-ok: LiteLLM callback compatibility + ) -> None: ... + + +@dataclass(frozen=True, slots=True) +class _CacheBinding: + invocation_id: str + cache: _AsyncCache + + +class InvocationCacheRegistry: + def __init__(self) -> None: + self._bindings: dict[str, _CacheBinding] = {} # mutable-ok: LiteLLM compatibility payload + + def register(self, invocation_id: str, cache: object | None) -> pb.CacheRef | None: + if cache is None: + return None + handle: Final = secrets.token_urlsafe(24) + self._bindings[handle] = _CacheBinding( + invocation_id, + cast(_AsyncCache, cache), # cast-ok: validated protobuf boundary + ) + return pb.CacheRef(invocation_id=invocation_id, opaque_handle=handle) + + def resolve(self, cache_ref: pb.CacheRef) -> _AsyncCache | None: + binding: Final = self._bindings.get(cache_ref.opaque_handle) + if binding is None or binding.invocation_id != cache_ref.invocation_id: + return None + return binding.cache + + def revoke(self, cache_ref: pb.CacheRef | None) -> None: + if cache_ref is not None: + self._bindings.pop(cache_ref.opaque_handle, None) + + +class GatewayServices(pb_grpc.GatewayServicesServicer): + def __init__(self, token: str, registry: InvocationCacheRegistry) -> None: + self._token: Final = token + self._registry: Final = registry + + async def CacheGet(self, request: pb.CacheGetRequest, context: grpc.aio.ServicerContext) -> pb.CacheGetResponse: + await self._authenticate(context) + cache = self._registry.resolve(request.cache) # rebind-ok: invocation-scoped RPC state + if cache is None: + return pb.CacheGetResponse(operation=_error(pb.ERROR_CODE_NOT_FOUND, "cache reference is invalid")) + try: + value = await cache.async_get_cache( # rebind-ok: invocation-scoped RPC state + key=request.key, local_only=request.local_only + ) + response: Final = pb.CacheGetResponse(operation=pb.OperationResult(ok=True)) + if value is not None: + response.value_json = json.dumps(value, separators=(",", ":"), default=str).encode() + except Exception as error: # noqa: BLE001 # cache backends may raise arbitrary exceptions + return pb.CacheGetResponse(operation=_error(pb.ERROR_CODE_EXTENSION_FAILED, str(error))) + else: + return response + + async def CacheSet(self, request: pb.CacheSetRequest, context: grpc.aio.ServicerContext) -> pb.OperationResult: + await self._authenticate(context) + cache = self._registry.resolve(request.cache) # rebind-ok: invocation-scoped RPC state + if cache is None: + return _error(pb.ERROR_CODE_NOT_FOUND, "cache reference is invalid") + try: + kwargs: dict[str, object] = { # mutable-ok: cache kwargs # rebind-ok: cache kwargs + "local_only": request.local_only + } + if request.HasField("ttl_seconds"): + kwargs["ttl"] = request.ttl_seconds + value = cast( # cast-ok: protobuf value # rebind-ok: decoded value + object, json.loads(request.value_json) + ) + await cache.async_set_cache(request.key, value, **kwargs) + return pb.OperationResult(ok=True) + except Exception as error: # noqa: BLE001 # cache backends may raise arbitrary exceptions + return _error(pb.ERROR_CODE_EXTENSION_FAILED, str(error)) + + async def _authenticate(self, context: grpc.aio.ServicerContext) -> None: + metadata: Final = cast( # cast-ok: validated protobuf boundary + Iterable[tuple[str, str]], context.invocation_metadata() or () + ) # cast-ok: validated protobuf boundary + if dict(metadata).get(TOKEN_METADATA_KEY) != self._token: # mutable-ok: LiteLLM compatibility payload + await context.abort(grpc.StatusCode.UNAUTHENTICATED, "invalid extension host token") + + +def _error(code: pb.ErrorCode, message: str) -> pb.OperationResult: + return pb.OperationResult(ok=False, error_code=code, error_message=message) diff --git a/litellm/extensions/client.py b/litellm/extensions/client.py new file mode 100644 index 00000000000..d40d5f5d562 --- /dev/null +++ b/litellm/extensions/client.py @@ -0,0 +1,283 @@ +# pyright: reportMissingModuleSource=false, reportUnknownMemberType=false, reportUnknownVariableType=false +from __future__ import annotations + +import asyncio +import contextlib +from collections.abc import AsyncIterator, Iterable +from dataclasses import dataclass +from typing import Final, cast # noqa: TID251 # protobuf repeated fields lack precise runtime typing + +import grpc + +from litellm._logging import verbose_proxy_logger +from litellm.python_extension.generated.v1 import extension_host_pb2 as pb +from litellm.python_extension.generated.v1 import extension_host_pb2_grpc as pb_grpc +from litellm.python_extension_host.constants import PROTOCOL_MAJOR, PROTOCOL_MINOR, TOKEN_METADATA_KEY + +from .config import ExtensionHostSettings +from .manifest import ExtensionManifest + + +@dataclass(frozen=True, slots=True) +class ExtensionHostHealth: + healthy: bool + reason: str | None = None + + +class PythonExtensionClient: + def __init__(self, settings: ExtensionHostSettings, manifest: ExtensionManifest) -> None: + self.settings: Final = settings + self.manifest: Final = manifest + target: Final = settings.endpoint.removeprefix("http://").removeprefix("https://") + self._channel: Final = grpc.aio.insecure_channel(target) + self._stub: Final = pb_grpc.PythonExtensionHostStub(self._channel) + self._metadata: Final = ((TOKEN_METADATA_KEY, settings.token),) + self._queue: Final[asyncio.Queue[pb.CallbackEvent]] = asyncio.Queue(settings.callback_queue_size) + self._worker: asyncio.Task[None] | None = None + self._recovery: asyncio.Task[None] | None = None + self._recovery_lock: Final = asyncio.Lock() + self._health = ExtensionHostHealth(False, "not connected") + self._bypass_counts: dict[tuple[str, str], int] = {} # mutable-ok: LiteLLM compatibility payload + self._descriptor_hooks: dict[str, frozenset[str]] = {} # mutable-ok: active revision descriptors + self._closed = False + + @property + def health(self) -> ExtensionHostHealth: + return self._health + + @property + def bypass_counts( + self, + ) -> dict[tuple[str, str], int]: # mutable-ok: LiteLLM compatibility payload + return dict(self._bypass_counts) # mutable-ok: LiteLLM compatibility payload + + def descriptor_hooks(self, plugin_id: str) -> frozenset[str]: + return self._descriptor_hooks.get(plugin_id, frozenset()) + + async def start(self) -> tuple[pb.ExtensionDescriptor, ...]: + descriptors: tuple[pb.ExtensionDescriptor, ...] = () # rebind-ok: invocation-scoped RPC state + try: + await asyncio.wait_for(self._channel.channel_ready(), self.settings.connect_timeout_seconds) + descriptors = await self._activate() # rebind-ok: invocation-scoped RPC state + except TimeoutError: + self._mark_unhealthy("connect timeout") + self._schedule_recovery() + except grpc.aio.AioRpcError as error: + if error.code() in (grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.DEADLINE_EXCEEDED): + self._mark_unhealthy(error.code().name) + self._schedule_recovery() + else: + raise RuntimeError(f"extension host startup failed: {error.details()}") from error + self._worker = asyncio.create_task(self._callback_worker(), name="python-extension-callbacks") + return descriptors + + async def close(self) -> None: + self._closed = True + if self._worker is not None: + self._worker.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._worker + if self._recovery is not None: + self._recovery.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._recovery + await self._channel.close() + + async def retire(self, revision_id: str) -> None: + try: + await self._stub.RetireRevision( + pb.RetireRevisionRequest(revision_id=revision_id), + timeout=self.settings.hook_timeout_seconds, + metadata=self._metadata, + ) + except grpc.aio.AioRpcError: + return + + async def execute_guardrail(self, invocation: pb.GuardrailInvocation) -> pb.GuardrailResult: + try: + result: Final = await self._stub.ExecuteGuardrail( + invocation, + timeout=self.settings.hook_timeout_seconds, + metadata=self._metadata, + ) + except (grpc.aio.AioRpcError, TimeoutError) as error: + reason: Final = _rpc_reason(error) + self._record_bypass(invocation.plugin_id, pb.HookPhase.Name(invocation.hook_phase), reason) + self._schedule_recovery() + return pb.GuardrailResult( + operation=pb.OperationResult(ok=True), + decision=pb.GUARDRAIL_DECISION_ALLOW, + ) + else: + if result.operation.ok: + self._health = ExtensionHostHealth(True) + else: + self._record_bypass( + invocation.plugin_id, + pb.HookPhase.Name(invocation.hook_phase), + "extension_error", + ) + return result + + def enqueue_callback(self, event: pb.CallbackEvent) -> bool: + try: + self._queue.put_nowait(event) + except asyncio.QueueFull: + self._record_bypass(event.plugin_id, "callback", "queue_full") + return False + else: + return True + + async def transform_stream(self, frames: AsyncIterator[pb.StreamFrame]) -> AsyncIterator[pb.StreamFrame]: + call: Final = self._stub.TransformStream(frames, metadata=self._metadata) + try: + async for frame in call: + yield frame + except grpc.aio.AioRpcError as error: + self._record_bypass("stream", "transform", error.code().name) + self._schedule_recovery() + finally: + if not call.done(): + call.cancel() + + async def _activate(self) -> tuple[pb.ExtensionDescriptor, ...]: + capabilities: Final = await self._stub.GetCapabilities( + pb.GetCapabilitiesRequest(protocol_major=PROTOCOL_MAJOR, protocol_minor=PROTOCOL_MINOR), + timeout=self.settings.connect_timeout_seconds, + metadata=self._metadata, + ) + if capabilities.protocol_major != PROTOCOL_MAJOR: + raise RuntimeError( + f"extension host protocol major {capabilities.protocol_major} does not match {PROTOCOL_MAJOR}" + ) + has_callbacks: Final = any(spec.kind == pb.EXTENSION_KIND_CALLBACK for spec in self.manifest.specs) + if has_callbacks and not capabilities.supports_callback_batching: + raise RuntimeError("extension host does not support callback batching") + if ( + capabilities.max_callback_batch_size + and self.settings.callback_batch_size > capabilities.max_callback_batch_size + ): + raise RuntimeError("python_extension_host.callback_batch_size exceeds the host capability") + if self.settings.gateway_listen is not None and not capabilities.supports_cache: + raise RuntimeError("extension host was not configured for reverse cache access") + response: pb.PrepareRevisionResponse = ( # rebind-ok: invocation-scoped RPC state + await self._stub.PrepareRevision( # rebind-ok: invocation-scoped RPC state + pb.PrepareRevisionRequest( + revision_id=self.manifest.revision_id, + extensions=self.manifest.specs, + ), + timeout=self.settings.hook_timeout_seconds, + metadata=self._metadata, + ) + ) + if not response.operation.ok and response.operation.error_code != pb.ERROR_CODE_ALREADY_EXISTS: + raise RuntimeError(f"extension manifest rejected: {response.operation.error_message}") + streaming_hooks: Final = { # mutable-ok: LiteLLM compatibility payload + "async_post_call_streaming_hook", + "async_post_call_streaming_iterator_hook", + } + if not capabilities.supports_duplex_streaming and any( + streaming_hooks.intersection(cast(Iterable[str], descriptor.hooks)) # cast-ok: validated protobuf boundary + for descriptor in cast( # cast-ok: validated protobuf boundary + Iterable[pb.ExtensionDescriptor], response.extensions + ) # cast-ok: validated protobuf boundary + ): + raise RuntimeError("extension host does not support required duplex streaming hooks") + commit: Final = await self._stub.CommitRevision( + pb.CommitRevisionRequest(revision_id=self.manifest.revision_id), + timeout=self.settings.hook_timeout_seconds, + metadata=self._metadata, + ) + if not commit.ok: + raise RuntimeError(f"extension manifest commit failed: {commit.error_message}") + self._health = ExtensionHostHealth(True) + extensions: Final = cast( # cast-ok: validated protobuf boundary + Iterable[pb.ExtensionDescriptor], response.extensions + ) # cast-ok: validated protobuf boundary + descriptors: Final = tuple(extensions) + self._descriptor_hooks = { # mutable-ok: active revision descriptors + descriptor.id: frozenset(descriptor.hooks) for descriptor in descriptors + } + return descriptors + + async def _callback_worker(self) -> None: + while True: + first = await self._queue.get() + batch = [first] # mutable-ok: LiteLLM compatibility payload + await asyncio.sleep(0.01) + while len(batch) < self.settings.callback_batch_size: + try: + batch.append(self._queue.get_nowait()) + except asyncio.QueueEmpty: + break + try: + response = await self._stub.PublishCallbackEvents( # rebind-ok: callback batch response + pb.PublishCallbackEventsRequest(events=batch), + timeout=self.settings.hook_timeout_seconds, + metadata=self._metadata, + ) + except (grpc.aio.AioRpcError, TimeoutError) as error: + reason = _rpc_reason(error) + for event in batch: + self._record_bypass(event.plugin_id, "callback", reason) + self._schedule_recovery() + else: + operations = tuple( # cast-ok: validated protobuf repeated field # rebind-ok: callback batch operations + cast(Iterable[pb.OperationResult], response.operations) # cast-ok: protobuf repeated field + ) + all_ok = len(operations) == len(batch) + for index, event in enumerate(batch): + if index >= len(operations): + self._record_bypass(event.plugin_id, "callback", "missing_operation") + continue + operation = operations[index] # rebind-ok: each callback has one operation + if not operation.ok: + all_ok = False + self._record_bypass( + event.plugin_id, + "callback", + _operation_reason(operation), + ) + if all_ok: + self._health = ExtensionHostHealth(True) + finally: + for _ in batch: + self._queue.task_done() + + def _schedule_recovery(self) -> None: + if self._closed or (self._recovery is not None and not self._recovery.done()): + return + self._recovery = asyncio.create_task(self._recover(), name="python-extension-recovery") + + async def _recover(self) -> None: + async with self._recovery_lock: + delay = 0.25 # rebind-ok: invocation-scoped RPC state + while not self._closed: + try: + await asyncio.wait_for(self._channel.channel_ready(), self.settings.connect_timeout_seconds) + await self._activate() + except (grpc.aio.AioRpcError, RuntimeError, TimeoutError) as error: + self._mark_unhealthy(_rpc_reason(error)) + await asyncio.sleep(delay) + delay = min(delay * 2, 5.0) + else: + return + + def _record_bypass(self, plugin_id: str, hook: str, reason: str) -> None: + key: Final = (plugin_id, f"{hook}:{reason}") + self._bypass_counts[key] = self._bypass_counts.get(key, 0) + 1 + self._mark_unhealthy(reason) + verbose_proxy_logger.warning( + "python extension host bypass plugin=%s hook=%s reason=%s", plugin_id, hook, reason + ) + + def _mark_unhealthy(self, reason: str) -> None: + self._health = ExtensionHostHealth(False, reason) + + +def _rpc_reason(error: BaseException) -> str: + return error.code().name if isinstance(error, grpc.aio.AioRpcError) else str(error) + + +def _operation_reason(operation: pb.OperationResult) -> str: + return operation.error_message or "extension_error" diff --git a/litellm/extensions/config.py b/litellm/extensions/config.py new file mode 100644 index 00000000000..1d1a9af2799 --- /dev/null +++ b/litellm/extensions/config.py @@ -0,0 +1,56 @@ +# pyright: reportUnknownArgumentType=false, reportUnknownMemberType=false, reportUnknownVariableType=false +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Final + + +@dataclass(frozen=True, slots=True) +class ExtensionHostSettings: + endpoint: str + token: str + connect_timeout_seconds: float = 5.0 + hook_timeout_seconds: float = 30.0 + callback_queue_size: int = 1_000 + callback_batch_size: int = 50 + gateway_listen: str | None = None + + +def settings_from_config( + config: dict[str, object], # mutable-ok: LiteLLM compatibility payload +) -> ExtensionHostSettings | None: + general_settings: Final = config.get("general_settings") + if not isinstance(general_settings, dict): + return None + raw: Final = general_settings.get("python_extension_host") + if raw is None: + return None + if not isinstance(raw, dict): + raise TypeError("general_settings.python_extension_host must be an object") + endpoint: Final = raw.get("endpoint") + token: Final = raw.get("token") + if not isinstance(endpoint, str) or not endpoint: + raise ValueError("python_extension_host.endpoint is required") + if not isinstance(token, str) or not token: + raise ValueError("python_extension_host.token is required") + resolved_token: Final = _resolve_env(token) + return ExtensionHostSettings( + endpoint=endpoint, + token=resolved_token, + connect_timeout_seconds=float(raw.get("connect_timeout_seconds", 5)), + hook_timeout_seconds=float(raw.get("hook_timeout_seconds", 30)), + callback_queue_size=int(raw.get("callback_queue_size", 1_000)), + callback_batch_size=int(raw.get("callback_batch_size", 50)), + gateway_listen=(str(raw["gateway_listen"]) if raw.get("gateway_listen") else None), + ) + + +def _resolve_env(value: str) -> str: + if not value.startswith("os.environ/"): + return value + name: Final = value.removeprefix("os.environ/") + resolved: Final = os.environ.get(name) + if not resolved: + raise ValueError(f"environment variable {name!r} is required for python_extension_host.token") + return resolved diff --git a/litellm/extensions/manifest.py b/litellm/extensions/manifest.py new file mode 100644 index 00000000000..d0913317e57 --- /dev/null +++ b/litellm/extensions/manifest.py @@ -0,0 +1,149 @@ +# pyright: reportUnknownArgumentType=false, reportUnknownMemberType=false, reportUnknownVariableType=false +from __future__ import annotations + +import hashlib +import json +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from typing import Final + +from litellm.python_extension.generated.v1 import extension_host_pb2 as pb + + +@dataclass(frozen=True, slots=True) +class ExtensionManifest: + revision_id: str + specs: tuple[pb.ExtensionSpec, ...] + callback_ids: Mapping[str, str] + guardrail_ids: Mapping[tuple[str, str], str] + + +def build_manifest(config: Mapping[str, object]) -> ExtensionManifest: + specs: list[ # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + pb.ExtensionSpec + ] = [] # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + callback_ids: dict[ # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + str, str + ] = {} # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + guardrail_ids: dict[ # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + tuple[str, str], str + ] = {} # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + callback_events: dict[ # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + str, set[str] + ] = {} # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + litellm_settings: Final = config.get("litellm_settings") + settings: Final = ( + litellm_settings if isinstance(litellm_settings, Mapping) else {} # mutable-ok: LiteLLM compatibility payload + ) # mutable-ok: LiteLLM compatibility payload + for setting, event in ( + ("callbacks", None), + ("success_callback", "success"), + ("failure_callback", "failure"), + ): + for entrypoint in _string_entries(settings.get(setting)): + if _is_customer_entrypoint(entrypoint): + events = callback_events.setdefault( + entrypoint, + set(), # mutable-ok: LiteLLM compatibility payload + ) + events.update(("success", "failure") if event is None else (event,)) + for entrypoint, events in sorted(callback_events.items()): + extension_id = _stable_id("callback", entrypoint) + callback_ids[entrypoint] = extension_id + specs.append( + pb.ExtensionSpec( + id=extension_id, + kind=pb.EXTENSION_KIND_CALLBACK, + entrypoint=entrypoint, + constructor_json=_canonical_json( + {"callback_events": sorted(events)} # mutable-ok: LiteLLM compatibility payload + ), # mutable-ok: LiteLLM compatibility payload + ) + ) + guardrail_configs: Final = _guardrail_configs(config, settings) + for guardrail in guardrail_configs: + name = guardrail.get("guardrail_name") + params = guardrail.get("litellm_params") + if not isinstance(name, str) or not isinstance(params, Mapping): + continue + entrypoint = params.get("guardrail") + if not isinstance(entrypoint, str) or not _is_customer_entrypoint(entrypoint): + continue + extension_id = _stable_id("guardrail", entrypoint, name) + guardrail_ids[(entrypoint, name)] = extension_id + kwargs = dict(params) # mutable-ok: LiteLLM compatibility payload + kwargs.pop("guardrail", None) + mode = kwargs.pop("mode", None) + default_on = kwargs.pop("default_on", False) + kwargs.update(guardrail_name=name, event_hook=mode, default_on=default_on) + specs.append( + pb.ExtensionSpec( + id=extension_id, + kind=pb.EXTENSION_KIND_GUARDRAIL, + entrypoint=entrypoint, + constructor_json=_canonical_json({"kwargs": kwargs}), # mutable-ok: LiteLLM compatibility payload + ) + ) + canonical_specs: Final = b"".join(spec.SerializeToString(deterministic=True) for spec in specs) + revision_id: Final = hashlib.sha256(canonical_specs).hexdigest()[:24] + return ExtensionManifest(revision_id, tuple(specs), callback_ids, guardrail_ids) + + +def manifest_json_from_config_path(config_path: str) -> str: + import yaml + + with open(config_path, encoding="utf-8") as config_file: + raw: Final = yaml.safe_load(config_file) or {} # mutable-ok: LiteLLM compatibility payload + if not isinstance(raw, dict): + raise TypeError("LiteLLM config must contain an object") + manifest: Final = build_manifest(raw) + extensions: Final = [ # mutable-ok: LiteLLM compatibility payload + { # mutable-ok: LiteLLM compatibility payload + "id": spec.id, + "kind": "guardrail" if spec.kind == pb.EXTENSION_KIND_GUARDRAIL else "callback", + "entrypoint": spec.entrypoint, + "constructor": json.loads(spec.constructor_json or b"{}"), + } + for spec in manifest.specs + ] + return json.dumps( + { # mutable-ok: LiteLLM compatibility payload + "revision_id": manifest.revision_id, + "extensions": extensions, + }, + separators=(",", ":"), + sort_keys=True, + ) + + +def _guardrail_configs( + config: Mapping[str, object], settings: Mapping[str, object] +) -> tuple[Mapping[str, object], ...]: + values: list[ # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + Mapping[str, object] + ] = [] # mutable-ok: LiteLLM compatibility payload # rebind-ok: invocation-scoped RPC state + for candidate in (config.get("guardrails"), settings.get("guardrails")): + if isinstance(candidate, list): + values.extend(item for item in candidate if isinstance(item, Mapping)) + return tuple(values) + + +def _string_entries(value: object) -> Iterable[str]: + if isinstance(value, str): + return (value,) + if isinstance(value, list): + return (item for item in value if isinstance(item, str)) + return () + + +def _is_customer_entrypoint(value: str) -> bool: + return ("." in value or ":" in value) and not value.startswith(("http://", "https://")) + + +def _stable_id(kind: str, entrypoint: str, name: str = "") -> str: + digest: Final = hashlib.sha256(f"{kind}\0{entrypoint}\0{name}".encode()).hexdigest()[:16] + return f"{kind}-{digest}" + + +def _canonical_json(value: object) -> bytes: + return json.dumps(value, separators=(",", ":"), sort_keys=True, default=str).encode() diff --git a/litellm/extensions/runtime.py b/litellm/extensions/runtime.py new file mode 100644 index 00000000000..36e54e4465f --- /dev/null +++ b/litellm/extensions/runtime.py @@ -0,0 +1,138 @@ +# pyright: reportMissingModuleSource=false, reportUnknownMemberType=false +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +import grpc + +from litellm.python_extension.generated.v1 import extension_host_pb2_grpc as pb_grpc +from litellm.types.guardrails import GuardrailEventHooks, Mode + +from .adapters import RemoteCustomGuardrail, RemoteCustomLogger +from .cache_gateway import GatewayServices, InvocationCacheRegistry +from .client import PythonExtensionClient +from .config import ExtensionHostSettings, settings_from_config +from .manifest import ExtensionManifest, build_manifest + + +@dataclass(slots=True) +class ExtensionRuntime: + settings: ExtensionHostSettings + manifest: ExtensionManifest + client: PythonExtensionClient + cache_registry: InvocationCacheRegistry + gateway_server: grpc.aio.Server | None = None + + def callback(self, entrypoint: str) -> RemoteCustomLogger | None: + plugin_id = self.manifest.callback_ids.get(entrypoint) # rebind-ok: invocation-scoped RPC state + if plugin_id is None: + return None + return RemoteCustomLogger(self.client, plugin_id, self.cache_registry) + + def guardrail( + self, + entrypoint: str, + name: str, + event_hook: GuardrailEventHooks # mutable-ok: LiteLLM compatibility payload + | list[GuardrailEventHooks] + | Mode + | None = None, # mutable-ok: LiteLLM compatibility payload + default_on: bool = False, + **kwargs: object, # kwargs-ok: LiteLLM callback compatibility + ) -> RemoteCustomGuardrail | None: + plugin_id = self.manifest.guardrail_ids.get((entrypoint, name)) # rebind-ok: invocation-scoped RPC state + if plugin_id is None: + return None + return RemoteCustomGuardrail( + self.client, + plugin_id, + self.cache_registry, + guardrail_name=name, + event_hook=event_hook, + default_on=default_on, + extra_params=kwargs, + ) + + async def close(self) -> None: + await self.client.close() + if self.gateway_server is not None: + await self.gateway_server.stop(grace=5) + + +_runtime: ExtensionRuntime | None = None # rebind-ok: invocation-scoped RPC state + + +def get_extension_runtime() -> ExtensionRuntime | None: + return _runtime + + +async def configure_extension_runtime( + config: dict[str, object], # mutable-ok: LiteLLM compatibility payload +) -> ExtensionRuntime | None: + global _runtime # noqa: PLW0603 # process-wide extension lifecycle singleton + settings: Final = settings_from_config(config) + if settings is None: + if _runtime is not None: + await _runtime.close() + _runtime = None # rebind-ok: invocation-scoped RPC state + return None + manifest: Final = build_manifest(config) + cache_registry: Final = InvocationCacheRegistry() + gateway_server: Final = await _start_gateway_services(settings, cache_registry) + client: Final = PythonExtensionClient(settings, manifest) + try: + await client.start() + except Exception: + await client.close() + if gateway_server is not None: + await gateway_server.stop(grace=0) + raise + previous: Final = _runtime + _runtime = ExtensionRuntime( # rebind-ok: invocation-scoped RPC state + settings, manifest, client, cache_registry, gateway_server + ) # rebind-ok: invocation-scoped RPC state + if previous is not None: + if previous.manifest.revision_id != manifest.revision_id: + await previous.client.retire(previous.manifest.revision_id) + await previous.close() + return _runtime + + +def remote_callback(entrypoint: str) -> RemoteCustomLogger | None: + return _runtime.callback(entrypoint) if _runtime is not None else None + + +def remote_guardrail( + entrypoint: str, + name: str, + event_hook: GuardrailEventHooks # mutable-ok: LiteLLM compatibility payload + | list[GuardrailEventHooks] + | Mode + | None = None, # mutable-ok: LiteLLM compatibility payload + default_on: bool = False, + **kwargs: object, # kwargs-ok: LiteLLM callback compatibility +) -> RemoteCustomGuardrail | None: + if _runtime is None: + return None + return _runtime.guardrail( + entrypoint, + name, + event_hook=event_hook, + default_on=default_on, + **kwargs, + ) + + +async def _start_gateway_services( + settings: ExtensionHostSettings, registry: InvocationCacheRegistry +) -> grpc.aio.Server | None: + if settings.gateway_listen is None: + return None + server: Final = grpc.aio.server() + pb_grpc.add_GatewayServicesServicer_to_server(GatewayServices(settings.token, registry), server) + target: Final = settings.gateway_listen.removeprefix("http://").removeprefix("https://") + if server.add_insecure_port(target) == 0: + raise RuntimeError(f"failed to bind GatewayServices to {settings.gateway_listen}") + await server.start() + return server diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 39e74d2c8bd..e26dd9be1e4 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -369,15 +369,25 @@ def initialize_callbacks_on_proxy( verbose_proxy_logger.debug( "%s attempting to import custom calback=%s %s", blue_color_code, callback, reset_color_code ) - imported_list.append( - _loaded_callback_or_raise( - entry=callback, - loaded=get_instance_fn( - value=callback, - config_file_path=config_file_path, - ), - ) + from litellm.extensions.runtime import get_extension_runtime, remote_callback + + remote_for_callback = remote_callback( # rebind-ok: each configured callback is independent + callback ) + if remote_for_callback is not None: + imported_list.append(remote_for_callback) + elif get_extension_runtime() is not None: + raise ValueError(f"callback {callback!r} was not present in the extension host manifest") + else: + imported_list.append( + _loaded_callback_or_raise( + entry=callback, + loaded=get_instance_fn( + value=callback, + config_file_path=config_file_path, + ), + ) + ) if isinstance(litellm.callbacks, list): litellm.callbacks.extend(imported_list) else: @@ -388,15 +398,23 @@ def initialize_callbacks_on_proxy( PrometheusLogger._mount_metrics_endpoint() else: - litellm.callbacks = [ - _loaded_callback_or_raise( - entry=value, - loaded=get_instance_fn( - value=value, - config_file_path=config_file_path, - ), - ) - ] + from litellm.extensions.runtime import get_extension_runtime, remote_callback + + remote_single: Final = remote_callback(value) + if remote_single is not None: + litellm.callbacks = [remote_single] # mutable-ok: global callback registry requires a list + elif get_extension_runtime() is not None: + raise ValueError(f"callback {value!r} was not present in the extension host manifest") + else: + litellm.callbacks = [ # mutable-ok: global callback registry requires a list + _loaded_callback_or_raise( + entry=value, + loaded=get_instance_fn( + value=value, + config_file_path=config_file_path, + ), + ) + ] verbose_proxy_logger.debug("%s Initialized Callbacks - %s %s", blue_color_code, litellm.callbacks, reset_color_code) diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index dc13c09dd38..09ce9023998 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -47,6 +47,7 @@ from litellm.types.guardrails import ( GuardrailEventHooks, LakeraCategoryThresholds, LitellmParams, + Mode, SupportedGuardrailIntegrations, ) @@ -424,6 +425,16 @@ def _apply_configured_bool_overrides(instance: CustomGuardrail, litellm_params: instance.scan_raw_request = bool(litellm_params.scan_raw_request) +def _remote_guardrail_event_hook( + mode: str | list[str] | Mode, # mutable-ok: guardrail config accepts a hook list +) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode: # mutable-ok: CustomGuardrail requires a list + if isinstance(mode, Mode): + return mode + if isinstance(mode, list): + return [GuardrailEventHooks(hook) for hook in mode] # mutable-ok: CustomGuardrail requires a list + return GuardrailEventHooks(mode) + + class InMemoryGuardrailHandler: """ Class that handles initializing guardrails and adding them to the CallbackManager @@ -581,10 +592,6 @@ class InMemoryGuardrailHandler: guardrail_type, ) - _guardrail_class: Final[Callable[..., CustomGuardrail]] = get_instance_fn( - guardrail_type, config_file_path=config_file_path - ) - mode: Final = litellm_params.mode if mode is None: raise ValueError( @@ -605,6 +612,25 @@ class InMemoryGuardrailHandler: for key in ["guardrail", "mode", "default_on"]: extra_params.pop(key, None) + from litellm.extensions.runtime import get_extension_runtime, remote_guardrail + + remote: Final = remote_guardrail( + guardrail_type, + guardrail["guardrail_name"], + event_hook=_remote_guardrail_event_hook(mode), + default_on=bool(default_on), + **extra_params, + ) + if remote is not None: + litellm.logging_callback_manager.add_litellm_callback(remote) + return remote + if get_extension_runtime() is not None: + raise ValueError(f"guardrail {guardrail_type!r} was not present in the extension host manifest") + + _guardrail_class: Final[Callable[..., CustomGuardrail]] = get_instance_fn( + guardrail_type, config_file_path=config_file_path + ) + _guardrail_callback: Final = _guardrail_class( guardrail_name=guardrail["guardrail_name"], event_hook=mode, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 77a80ea0052..4d03d1df38a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5044,6 +5044,10 @@ class ProxyConfig: self._load_environment_variables(config=config) + from litellm.extensions.runtime import configure_extension_runtime + + await configure_extension_runtime(config) + ## Coordination Redis (before cache init, so the explicit block wins) coordination_redis_cache: Final = self._init_coordination_redis(config=config) if coordination_redis_cache is not None: @@ -5246,8 +5250,16 @@ class ProxyConfig: for callback in value: # user passed custom_callbacks.async_on_succes_logger. They need us to import a function if "." in callback: + from litellm.extensions.runtime import get_extension_runtime, remote_callback + + callback_instance = remote_callback(callback) + if callback_instance is None and get_extension_runtime() is not None: + raise ValueError( + f"callback {callback!r} was not present in the extension host manifest" + ) litellm.logging_callback_manager.add_litellm_success_callback( - get_instance_fn( + callback_instance + or get_instance_fn( value=callback, config_file_path=config_file_path, ) @@ -5273,8 +5285,16 @@ class ProxyConfig: for callback in value: # user passed custom_callbacks.async_on_succes_logger. They need us to import a function if "." in callback: + from litellm.extensions.runtime import get_extension_runtime, remote_callback + + callback_instance = remote_callback(callback) + if callback_instance is None and get_extension_runtime() is not None: + raise ValueError( + f"callback {callback!r} was not present in the extension host manifest" + ) litellm.logging_callback_manager.add_litellm_failure_callback( - get_instance_fn( + callback_instance + or get_instance_fn( value=callback, config_file_path=config_file_path, ) diff --git a/tests/test_litellm/python_extension_host/__init__.py b/tests/test_litellm/python_extension_host/__init__.py new file mode 100644 index 00000000000..aacd20e6eac --- /dev/null +++ b/tests/test_litellm/python_extension_host/__init__.py @@ -0,0 +1 @@ +"""Python extension host tests.""" diff --git a/tests/test_litellm/python_extension_host/test_rpc_sidecar.py b/tests/test_litellm/python_extension_host/test_rpc_sidecar.py new file mode 100644 index 00000000000..01fe2cfa850 --- /dev/null +++ b/tests/test_litellm/python_extension_host/test_rpc_sidecar.py @@ -0,0 +1,537 @@ +from __future__ import annotations + +import asyncio +import json +import os +import socket +import sys +from collections.abc import AsyncIterator +from datetime import datetime +from pathlib import Path + +import grpc +import pytest +import pytest_asyncio + +from litellm.exceptions import GuardrailRaisedException +from litellm.extensions.cache_gateway import GatewayServices, InvocationCacheRegistry +from litellm.extensions.client import PythonExtensionClient +from litellm.extensions.config import ExtensionHostSettings, settings_from_config +from litellm.extensions.manifest import build_manifest +from litellm.extensions.runtime import configure_extension_runtime +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.python_extension.generated.v1 import extension_host_pb2 as pb +from litellm.python_extension.generated.v1 import extension_host_pb2_grpc as pb_grpc +from litellm.python_extension_host.constants import TOKEN_METADATA_KEY +from litellm.python_extension_host.service import PythonExtensionHostService + +MODULE = "tests.test_litellm.python_extension_host.test_rpc_sidecar" +TOKEN = "test-extension-token" +METADATA = ((TOKEN_METADATA_KEY, TOKEN),) + + +class GuardrailFixture(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + data["hosted"] = True + data["received_plaintext_key"] = bool(user_api_key_dict.api_key) + return data + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + response["post_call"] = data["model"] + return response + + +class BlockingGuardrailFixture(CustomGuardrail): + async def async_moderation_hook(self, data, user_api_key_dict, call_type): + raise GuardrailRaisedException( + guardrail_name="blocking", + message="blocked by fixture", + blocked_content=True, + ) + + +class CacheGuardrailFixture(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + data["cached"] = await cache.async_get_cache("seed") + await cache.async_set_cache("written", {"ok": True}, ttl=30) + return data + + +class FakeCache: + def __init__(self): + self.values: dict[str, object] = {"seed": {"value": 7}} + + async def async_get_cache(self, key: str, **kwargs: object) -> object | None: + return self.values.get(key) + + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: + self.values[key] = value + + +class UnsupportedFixture(CustomLogger): + async def async_pre_request_hook(self, model, messages, kwargs): + return kwargs + + +class CallbackFixture(CustomLogger): + events: list[tuple[str, str]] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.events.append((kwargs["model"], response_obj["id"])) + + +function_events: list[str] = [] + + +async def callback_function(kwargs, response_obj, start_time, end_time): + function_events.append(kwargs["model"]) + + +class StreamFixture(CustomLogger): + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict, response, request_data + ) -> AsyncIterator[dict[str, str]]: + async for chunk in response: + yield {"value": chunk["value"].upper()} + yield {"value": "!"} + + +class ChunkStreamFixture(CustomLogger): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + return {"value": response["value"].upper()} + + +@pytest_asyncio.fixture(loop_scope="function") +async def host_stub(): + server = grpc.aio.server() + pb_grpc.add_PythonExtensionHostServicer_to_server(PythonExtensionHostService(TOKEN), server) + port = server.add_insecure_port("127.0.0.1:0") + await server.start() + channel = grpc.aio.insecure_channel(f"127.0.0.1:{port}") + await channel.channel_ready() + try: + yield pb_grpc.PythonExtensionHostStub(channel) + finally: + await channel.close() + await server.stop(grace=0) + + +@pytest.mark.asyncio +async def test_capabilities_require_auth_and_negotiate_version(host_stub): + with pytest.raises(grpc.aio.AioRpcError) as error: + await host_stub.GetCapabilities(pb.GetCapabilitiesRequest(protocol_major=1, protocol_minor=0)) + assert error.value.code() == grpc.StatusCode.UNAUTHENTICATED + + capabilities = await host_stub.GetCapabilities( + pb.GetCapabilitiesRequest(protocol_major=1, protocol_minor=3), metadata=METADATA + ) + assert capabilities.protocol_major == 1 + assert capabilities.protocol_minor == 0 + assert capabilities.supports_duplex_streaming is True + assert "async_pre_call_hook" in capabilities.supported_hooks + + +@pytest.mark.asyncio +async def test_prepare_commit_execute_replace_and_retire(host_stub): + prepare = await host_stub.PrepareRevision( + pb.PrepareRevisionRequest( + revision_id="revision-1", + extensions=[_spec("guardrail", pb.EXTENSION_KIND_GUARDRAIL, "GuardrailFixture")], + ), + metadata=METADATA, + ) + assert prepare.operation.ok is True + assert prepare.extensions[0].hooks == ["async_post_call_success_hook", "async_pre_call_hook"] + + prepared_again = await host_stub.PrepareRevision( + pb.PrepareRevisionRequest( + revision_id="revision-1", + extensions=[_spec("guardrail", pb.EXTENSION_KIND_GUARDRAIL, "GuardrailFixture")], + ), + metadata=METADATA, + ) + assert prepared_again.operation.error_code == pb.ERROR_CODE_ALREADY_EXISTS + assert prepared_again.extensions[0].hooks == prepare.extensions[0].hooks + + inactive = await host_stub.ExecuteGuardrail(_guardrail_request("revision-1"), metadata=METADATA) + assert inactive.decision == pb.GUARDRAIL_DECISION_ERROR + + committed = await host_stub.CommitRevision(pb.CommitRevisionRequest(revision_id="revision-1"), metadata=METADATA) + assert committed.ok is True + replaced = await host_stub.ExecuteGuardrail(_guardrail_request("revision-1"), metadata=METADATA) + assert replaced.decision == pb.GUARDRAIL_DECISION_REPLACE_REQUEST + replaced_request = json.loads(replaced.request_json) + assert replaced_request["hosted"] is True + assert replaced_request["received_plaintext_key"] is False + + active_retire = await host_stub.RetireRevision( + pb.RetireRevisionRequest(revision_id="revision-1"), metadata=METADATA + ) + assert active_retire.error_code == pb.ERROR_CODE_INVALID_ARGUMENT + + await host_stub.PrepareRevision( + pb.PrepareRevisionRequest(revision_id="revision-2", extensions=[]), metadata=METADATA + ) + await host_stub.CommitRevision(pb.CommitRevisionRequest(revision_id="revision-2"), metadata=METADATA) + retired = await host_stub.RetireRevision(pb.RetireRevisionRequest(revision_id="revision-1"), metadata=METADATA) + assert retired.ok is True + + +@pytest.mark.asyncio +async def test_recognized_guardrail_exception_becomes_block(host_stub): + await _activate( + host_stub, + "block-revision", + [_spec("blocking", pb.EXTENSION_KIND_GUARDRAIL, "BlockingGuardrailFixture")], + ) + request = _guardrail_request("block-revision", plugin_id="blocking") + request.hook_phase = pb.HOOK_PHASE_DURING_CALL + result = await host_stub.ExecuteGuardrail(request, metadata=METADATA) + assert result.operation.ok is True + assert result.decision == pb.GUARDRAIL_DECISION_BLOCK + assert result.public_error.status_code == 400 + assert "blocked by fixture" in result.public_error.message + + +@pytest.mark.asyncio +async def test_prepare_rejects_any_unsupported_override(host_stub): + response = await host_stub.PrepareRevision( + pb.PrepareRevisionRequest( + revision_id="unsupported", + extensions=[_spec("unsupported", pb.EXTENSION_KIND_CALLBACK, "UnsupportedFixture")], + ), + metadata=METADATA, + ) + assert response.operation.ok is False + assert response.operation.error_code == pb.ERROR_CODE_LOAD_FAILED + assert "async_pre_request_hook" in response.operation.error_message + + +@pytest.mark.asyncio +async def test_callback_batch_supports_logger_and_function(host_stub): + CallbackFixture.events.clear() + function_events.clear() + specs = [ + _spec("logger", pb.EXTENSION_KIND_CALLBACK, "CallbackFixture"), + _spec( + "function", + pb.EXTENSION_KIND_CALLBACK, + "callback_function", + {"callback_events": ["success"]}, + ), + ] + await _activate(host_stub, "callbacks", specs) + events = [_callback_event("callbacks", plugin_id) for plugin_id in ("logger", "function")] + response = await host_stub.PublishCallbackEvents(pb.PublishCallbackEventsRequest(events=events), metadata=METADATA) + assert all(operation.ok for operation in response.operations) + assert CallbackFixture.events == [("test-model", "response-1")] + assert function_events == ["test-model"] + + +@pytest.mark.asyncio +async def test_duplex_iterator_can_emit_multiple_chunks_per_input(host_stub): + await _activate( + host_stub, + "streaming", + [_spec("stream", pb.EXTENSION_KIND_CALLBACK, "StreamFixture")], + ) + + async def frames() -> AsyncIterator[pb.StreamFrame]: + yield pb.StreamFrame( + kind=pb.STREAM_FRAME_KIND_OPEN, + stream_id="stream-1", + open=pb.StreamOpen( + context=_context("streaming"), + plugin_id="stream", + request_json=b'{"model":"test-model"}', + auth=pb.AuthContext(key_hash="hashed"), + iterator_hook=True, + ), + ) + yield pb.StreamFrame( + kind=pb.STREAM_FRAME_KIND_INPUT_CHUNK, + stream_id="stream-1", + chunk_json=b'{"value":"hello"}', + ) + yield pb.StreamFrame(kind=pb.STREAM_FRAME_KIND_END, stream_id="stream-1") + + output = [frame async for frame in host_stub.TransformStream(frames(), metadata=METADATA)] + assert [json.loads(frame.chunk_json) for frame in output[:-1]] == [ + {"value": "HELLO"}, + {"value": "!"}, + ], [(frame.kind, frame.error.message) for frame in output] + assert output[-1].kind == pb.STREAM_FRAME_KIND_END + + +@pytest.mark.asyncio +async def test_duplex_chunk_hook_transforms_each_input(host_stub): + await _activate( + host_stub, + "chunk-streaming", + [_spec("chunk-stream", pb.EXTENSION_KIND_CALLBACK, "ChunkStreamFixture")], + ) + + async def frames() -> AsyncIterator[pb.StreamFrame]: + yield pb.StreamFrame( + kind=pb.STREAM_FRAME_KIND_OPEN, + stream_id="stream-2", + open=pb.StreamOpen( + context=_context("chunk-streaming"), + plugin_id="chunk-stream", + request_json=b'{"model":"test-model"}', + auth=pb.AuthContext(key_hash="hashed"), + iterator_hook=False, + ), + ) + yield pb.StreamFrame( + kind=pb.STREAM_FRAME_KIND_INPUT_CHUNK, + stream_id="stream-2", + chunk_json=b'{"value":"hello"}', + ) + yield pb.StreamFrame(kind=pb.STREAM_FRAME_KIND_END, stream_id="stream-2") + + output = [frame async for frame in host_stub.TransformStream(frames(), metadata=METADATA)] + assert json.loads(output[0].chunk_json) == {"value": "HELLO"} + assert output[1].kind == pb.STREAM_FRAME_KIND_END + + +@pytest.mark.asyncio +async def test_invocation_scoped_reverse_cache_access_is_revoked(): + registry = InvocationCacheRegistry() + cache = FakeCache() + cache_ref = registry.register("invocation-1", cache) + assert cache_ref is not None + + gateway_server = grpc.aio.server() + pb_grpc.add_GatewayServicesServicer_to_server(GatewayServices(TOKEN, registry), gateway_server) + gateway_port = gateway_server.add_insecure_port("127.0.0.1:0") + await gateway_server.start() + gateway_channel = grpc.aio.insecure_channel(f"127.0.0.1:{gateway_port}") + gateway_stub = pb_grpc.GatewayServicesStub(gateway_channel) + + host_server = grpc.aio.server() + pb_grpc.add_PythonExtensionHostServicer_to_server( + PythonExtensionHostService(TOKEN, gateway_stub=gateway_stub), host_server + ) + host_port = host_server.add_insecure_port("127.0.0.1:0") + await host_server.start() + host_channel = grpc.aio.insecure_channel(f"127.0.0.1:{host_port}") + host = pb_grpc.PythonExtensionHostStub(host_channel) + try: + await _activate( + host, + "cache-revision", + [_spec("cache", pb.EXTENSION_KIND_GUARDRAIL, "CacheGuardrailFixture")], + ) + request = _guardrail_request("cache-revision", plugin_id="cache") + request.cache.CopyFrom(cache_ref) + result = await host.ExecuteGuardrail(request, metadata=METADATA) + assert json.loads(result.request_json)["cached"] == {"value": 7} + assert cache.values["written"] == {"ok": True} + + registry.revoke(cache_ref) + revoked = await gateway_stub.CacheGet(pb.CacheGetRequest(cache=cache_ref, key="seed"), metadata=METADATA) + assert revoked.operation.error_code == pb.ERROR_CODE_NOT_FOUND + finally: + await host_channel.close() + await host_server.stop(grace=0) + await gateway_channel.close() + await gateway_server.stop(grace=0) + + +def test_manifest_build_does_not_import_customer_modules(): + module_name = "customer_extension_that_proxy_must_not_import" + config = { + "litellm_settings": { + "callbacks": [f"{module_name}.logger"], + "success_callback": [f"{module_name}.success"], + "failure_callback": [f"{module_name}.failure"], + }, + "guardrails": [ + { + "guardrail_name": "customer-guardrail", + "litellm_params": { + "guardrail": f"{module_name}.Guardrail", + "mode": "pre_call", + }, + } + ], + } + + manifest = build_manifest(config) + + assert len(manifest.specs) == 4 + assert module_name not in sys.modules + assert settings_from_config(config) is None + + +@pytest.mark.asyncio +async def test_python_client_fails_open_when_host_is_unavailable(): + port = _unused_port() + manifest = build_manifest({"litellm_settings": {"callbacks": [f"{MODULE}.CallbackFixture"]}}) + client = PythonExtensionClient( + ExtensionHostSettings( + endpoint=f"http://127.0.0.1:{port}", + token=TOKEN, + connect_timeout_seconds=0.05, + hook_timeout_seconds=0.05, + ), + manifest, + ) + try: + assert await client.start() == () + result = await client.execute_guardrail(_guardrail_request(manifest.revision_id)) + assert result.decision == pb.GUARDRAIL_DECISION_ALLOW + assert client.health.healthy is False + assert client.bypass_counts + finally: + await client.close() + + +@pytest.mark.asyncio +async def test_proxy_and_sidecar_share_contract_without_proxy_import(tmp_path: Path): + module_name = "customer_sidecar_only_plugin" + import_marker = tmp_path / "import.pid" + event_marker = tmp_path / "event.pid" + (tmp_path / f"{module_name}.py").write_text( + """ +import os +from pathlib import Path +from litellm.integrations.custom_logger import CustomLogger + +Path(os.environ["EXTENSION_IMPORT_MARKER"]).write_text(str(os.getpid())) + +class HostedLogger(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + Path(os.environ["EXTENSION_EVENT_MARKER"]).write_text(str(os.getpid())) +""".lstrip(), + encoding="utf-8", + ) + port = _unused_port() + environment = os.environ.copy() + environment["PYTHONPATH"] = os.pathsep.join((str(tmp_path), os.getcwd())) + environment["LITELLM_EXTENSION_HOST_TOKEN"] = TOKEN + environment["EXTENSION_IMPORT_MARKER"] = str(import_marker) + environment["EXTENSION_EVENT_MARKER"] = str(event_marker) + process = await asyncio.create_subprocess_exec( + sys.executable, + "-m", + "litellm.python_extension_host.server", + "--listen", + f"127.0.0.1:{port}", + env=environment, + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + ) + channel = grpc.aio.insecure_channel(f"127.0.0.1:{port}") + try: + await asyncio.wait_for(channel.channel_ready(), timeout=10) + config = { + "general_settings": { + "python_extension_host": { + "endpoint": f"http://127.0.0.1:{port}", + "token": TOKEN, + "connect_timeout_seconds": 2, + "hook_timeout_seconds": 2, + } + }, + "litellm_settings": { + "success_callback": [f"{module_name}.HostedLogger"], + }, + } + runtime = await configure_extension_runtime(config) + assert runtime is not None + assert module_name not in sys.modules + await _wait_for_file(import_marker) + assert int(import_marker.read_text()) == process.pid + + callback = runtime.callback(f"{module_name}.HostedLogger") + assert callback is not None + await callback.async_log_success_event( + {"model": "hosted-model"}, + {"id": "response-1"}, + datetime.now(), + datetime.now(), + ) + await _wait_for_file(event_marker) + assert int(event_marker.read_text()) == process.pid + assert process.pid != os.getpid() + finally: + await configure_extension_runtime({}) + await channel.close() + process.terminate() + try: + await asyncio.wait_for(process.wait(), timeout=5) + except TimeoutError: + process.kill() + await process.wait() + + +def _unused_port() -> int: + with socket.socket() as server_socket: + server_socket.bind(("127.0.0.1", 0)) + return int(server_socket.getsockname()[1]) + + +async def _wait_for_file(path: Path) -> None: + for _ in range(200): + if path.exists(): + return + await asyncio.sleep(0.025) + raise AssertionError(f"timed out waiting for {path}") + + +def _spec( + plugin_id: str, + kind: int, + target: str, + constructor: dict[str, object] | None = None, +) -> pb.ExtensionSpec: + return pb.ExtensionSpec( + id=plugin_id, + kind=kind, + entrypoint=f"{MODULE}.{target}", + constructor_json=json.dumps(constructor or {}).encode(), + ) + + +async def _activate(host_stub, revision: str, specs: list[pb.ExtensionSpec]) -> None: + prepared = await host_stub.PrepareRevision( + pb.PrepareRevisionRequest(revision_id=revision, extensions=specs), metadata=METADATA + ) + assert prepared.operation.ok, prepared.operation.error_message + committed = await host_stub.CommitRevision(pb.CommitRevisionRequest(revision_id=revision), metadata=METADATA) + assert committed.ok + + +def _context(revision: str) -> pb.InvocationContext: + return pb.InvocationContext( + request_id="request-1", + invocation_id="invocation-1", + active_revision=revision, + api_surface="chat.completions", + call_type="acompletion", + ) + + +def _guardrail_request(revision: str, plugin_id: str = "guardrail") -> pb.GuardrailInvocation: + return pb.GuardrailInvocation( + context=_context(revision), + plugin_id=plugin_id, + hook_phase=pb.HOOK_PHASE_PRE_CALL, + request_json=b'{"model":"test-model"}', + auth=pb.AuthContext(key_hash="sha256-only", user_id="user-1", team_id="team-1"), + ) + + +def _callback_event(revision: str, plugin_id: str) -> pb.CallbackEvent: + return pb.CallbackEvent( + context=_context(revision), + plugin_id=plugin_id, + kind=pb.CALLBACK_EVENT_KIND_SUCCESS, + standard_logging_payload_json=b'{"model":"test-model"}', + response_json=b'{"id":"response-1"}', + start_time_seconds=datetime(2026, 1, 1).timestamp(), + end_time_seconds=datetime(2026, 1, 1, 0, 0, 1).timestamp(), + )