diff --git a/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs b/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs index 9036deb9871..e247c650fad 100644 --- a/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs +++ b/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs @@ -5,6 +5,7 @@ use serde_json::Value; #[derive(Deserialize)] struct Input { + path: String, model_alias: String, provider_model: String, api_base: String, @@ -21,7 +22,8 @@ async fn main() { Ok(input) => input, Err(error) => fail(error), }; - let result = litellm_ai_gateway::trace_parity::traced_messages_request( + let result = litellm_ai_gateway::trace_parity::traced_request( + input.path, input.model_alias, input.provider_model, input.api_base, diff --git a/litellm-rust/crates/ai-gateway/src/trace_parity.rs b/litellm-rust/crates/ai-gateway/src/trace_parity.rs index 00c9b53e691..7540a71fb12 100644 --- a/litellm-rust/crates/ai-gateway/src/trace_parity.rs +++ b/litellm-rust/crates/ai-gateway/src/trace_parity.rs @@ -29,14 +29,15 @@ pub struct TracedGatewayResponse { pub trace: Vec, } -pub async fn traced_messages_request( +pub async fn traced_request( + path: String, model_alias: String, provider_model: String, api_base: String, body: Value, ) -> TracedGatewayResponse { let trace = litellm_core::observability::FunctionTrace::default(); - let result = messages_request(model_alias, provider_model, api_base, body) + let result = request(path, model_alias, provider_model, api_base, body) .with_subscriber(trace.dispatcher()) .await; let events = trace.events(); @@ -54,7 +55,8 @@ pub async fn traced_messages_request( } } -pub async fn messages_request( +pub async fn request( + path: String, model_alias: String, provider_model: String, api_base: String, @@ -75,7 +77,7 @@ pub async fn messages_request( }; let request = Request::builder() .method("POST") - .uri("/v1/messages") + .uri(path) .header(AUTHORIZATION, "Bearer trace-master-key") .header(CONTENT_TYPE, "application/json") .body(Body::from(body.to_string())) diff --git a/litellm/constants.py b/litellm/constants.py index e8aafda1797..c106be688e4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -317,6 +317,8 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure" +REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" +REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged" # SSL/TLS cipher configuration for faster handshakes # Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones diff --git a/litellm/integrations/arize/arize_phoenix_prompt_manager.py b/litellm/integrations/arize/arize_phoenix_prompt_manager.py index 0c9e868c146..fbf25ea87fd 100644 --- a/litellm/integrations/arize/arize_phoenix_prompt_manager.py +++ b/litellm/integrations/arize/arize_phoenix_prompt_manager.py @@ -379,7 +379,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement): """Reload prompts from Arize Phoenix.""" if self.prompt_id: self._prompt_manager = None # Reset to force reload - self.prompt_manager # This will trigger reload + _ = self.prompt_manager # access triggers lazy reload def should_run_prompt_management( self, diff --git a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py index ff34bd91e31..e98fa77a562 100644 --- a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py +++ b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py @@ -406,7 +406,7 @@ class BitBucketPromptManager(CustomPromptManagement): """Reload prompts from BitBucket.""" if self.prompt_id: self._prompt_manager = None # Reset to force reload - self.prompt_manager # This will trigger reload + _ = self.prompt_manager # access triggers lazy reload def should_run_prompt_management( self, diff --git a/litellm/integrations/weights_biases.py b/litellm/integrations/weights_biases.py index 97a2acbac08..d1a8ec098cf 100644 --- a/litellm/integrations/weights_biases.py +++ b/litellm/integrations/weights_biases.py @@ -4,18 +4,10 @@ imported_openAIResponse = True try: import io import logging - import sys - from typing import Any, TypeVar + from typing import Any, Literal, Protocol, TypeVar from wandb.sdk.data_types import trace_tree - if sys.version_info >= (3, 8): - from typing import Literal, Protocol - else: - from typing import Literal - - from typing_extensions import Protocol - logger: Final = logging.getLogger(__name__) K = TypeVar("K", bound=str) diff --git a/litellm/litellm_core_utils/coroutine_checker.py b/litellm/litellm_core_utils/coroutine_checker.py index 52fc44ba8dc..7b9a650c66b 100644 --- a/litellm/litellm_core_utils/coroutine_checker.py +++ b/litellm/litellm_core_utils/coroutine_checker.py @@ -36,7 +36,7 @@ class CoroutineChecker: target = callback if not inspect.isfunction(target) and not inspect.ismethod(target): try: - call_attr: Final = getattr(target, "__call__", None) + call_attr: Final = getattr(target, "__call__", None) # noqa: B004 # value unwrap so iscoroutinefunction sees through functors if call_attr is not None: target = call_attr except Exception: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index ece619e3883..21ae8b001dd 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1757,7 +1757,7 @@ def convert_to_anthropic_tool_invoke( anthropic_tool_invoke: Final[list[AnthropicMessagesToolUseParam | dict[str, object]]] = [] for tool in tool_calls: - if not get_attribute_or_key(tool, "type") == "function": + if get_attribute_or_key(tool, "type") != "function": continue tool_id = cast(str, get_attribute_or_key(tool, "id")) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index e3f8786a39a..fa567bdf4c9 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -10,6 +10,7 @@ from typing_extensions import ReadOnly import litellm from litellm._logging import redact_internal_details_from_client_message, verbose_logger +from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.types.llms.openai import ( @@ -35,9 +36,6 @@ else: CLIENT_CONNECTION_CLASS = Any -REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" - - @dataclass(frozen=True, slots=True) class BackendClose: code: int @@ -1153,6 +1151,7 @@ class RealTimeStreaming: self._logging_worker.ensure_initialized_and_enqueue( self.logging_obj.dispatch_failure_handlers(error, traceback.format_exc(), prefer_async_handlers=True) ) + self.logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True @staticmethod def _detect_beta_header(websocket: ScopedWebSocket) -> bool: diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 3b128899f45..4c61fac82bb 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -454,7 +454,7 @@ def token_counter( params: Final = _MessageCountParams(model, custom_tokenizer) num_tokens = _count_messages(params, new_messages, use_default_image_token_count, default_token_count) if count_response_tokens is False: - includes_system_message: Final = any([message.get("role", None) == "system" for message in new_messages]) + includes_system_message: Final = any(message.get("role", None) == "system" for message in new_messages) num_tokens += _count_extra(params.count_function, tools, tool_choice, includes_system_message) else: diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 0665b3f64c5..53a864a880a 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -144,10 +144,13 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): def get_api_key(api_key: str | None = None) -> str | None: return api_key or litellm.api_key or get_secret_str("AZURE_AI_API_KEY") + @staticmethod + def get_api_version(api_version: str | None = None) -> str | None: + return api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") + @property - def api_version(self, api_version: str | None = None) -> str | None: - api_version = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") - return api_version + def api_version(self) -> str | None: + return AzureFoundryModelInfo.get_api_version() def get_token_counter(self) -> BaseTokenCounter | None: """ diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 26d47330859..93e8813fe7e 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -20,6 +20,7 @@ from litellm.constants import ( BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY, BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY, + REALTIME_SESSION_SUCCESS_LOGGED_KEY, ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -346,6 +347,7 @@ class BedrockRealtime(BaseAWSLLM): prefer_async_handlers=True, ) ) + logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True if outcome.provider_failure is None: return @@ -382,7 +384,7 @@ class BedrockRealtime(BaseAWSLLM): ) bedrock_task: Final = asyncio.create_task(collect_logged_events()) - await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_EXCEPTION) + await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_COMPLETED) client_disconnected: Final = ( client_task.done() and not client_task.cancelled() and client_task.exception() is None ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d80a39a41da..708a94123de 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4418,6 +4418,9 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): access_group_mcp_server_ids: list[str] | None = None access_group_agent_ids: list[str] | None = None access_group_details: tuple[TeamAccessGroupModelGrant, ...] | None = None + # Parent org's model ceiling, reported only to callers who can manage the team. + # None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling. + organization_models: list[str] | None = None class TeamInfoResponseObject(TypedDict): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8abd13ce02d..e58c8437899 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -328,9 +328,23 @@ _safe_json_loads_obj: Final = _typed_json_loads(safe_json_loads) last_db_access_time: Final = LimitedSizeOrderedDict(max_size=100) db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s +_TEAM_MEMBERSHIP_INFLIGHT_MAX: Final = 10000 +_team_membership_inflight: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX) + + +class _TeamMembershipCacheMiss: + __slots__ = () + + +_TEAM_MEMBERSHIP_CACHE_MISS: Final = _TeamMembershipCacheMiss() + all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value +def _membership_from_shared_load(result: object) -> LiteLLM_TeamMembership | None: + return result if isinstance(result, LiteLLM_TeamMembership) else None + + def _log_budget_lookup_failure(entity: str, error: Exception) -> None: """ Log a warning when budget lookup fails; cache will not be populated. @@ -888,6 +902,22 @@ async def common_checks( and (route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route)) ) + membership_user_id: Final = ( + valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None + ) + team_membership_loaded: Final = team_object is not None and membership_user_id is not None + loaded_team_membership: Final = ( + await get_team_membership( + user_id=membership_user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if team_object is not None and membership_user_id is not None + else None + ) + unpriced_models: Final = ( _unpriced_models_in_request(model=_model, llm_router=llm_router) if litellm.block_requests_for_models_without_pricing and RouteChecks.is_llm_api_route(route=route) @@ -937,6 +967,8 @@ async def common_checks( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=loaded_team_membership, + team_membership_loaded=team_membership_loaded, ) # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent @@ -988,6 +1020,8 @@ async def common_checks( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=loaded_team_membership, + team_membership_loaded=team_membership_loaded, ) # Run before apply_key_tags_pre_auth injects key metadata.tags into request_body. @@ -1097,6 +1131,8 @@ async def common_checks( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=loaded_team_membership, + team_membership_loaded=team_membership_loaded, ), _check_end_user_budget(end_user_obj=end_user_object, route=route) if end_user_object is not None and end_user_object.litellm_budget_table is not None @@ -2142,7 +2178,76 @@ async def get_tag_object( return tag_objects.get(tag_name) +def _membership_from_cached_payload( + cached: object, +) -> LiteLLM_TeamMembership | None | _TeamMembershipCacheMiss: + if cached is None: + return _TEAM_MEMBERSHIP_CACHE_MISS + if cached == NO_TEAM_MEMBERSHIP_SENTINEL: + return None + cached_membership: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership) + return cached_membership if cached_membership is not None else _TEAM_MEMBERSHIP_CACHE_MISS + + @log_db_metrics +async def _fetch_team_membership_from_db( + user_id: str, + team_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, +) -> LiteLLM_TeamMembership | None: + _ = parent_otel_span, proxy_logging_obj + response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, + include={"litellm_budget_table": True}, + ) + membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict()) + _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) + if membership is None: + await user_api_key_cache.async_set_cache( + key=_key, + value=NO_TEAM_MEMBERSHIP_SENTINEL, + ttl=get_management_object_ttl(user_api_key_cache), + ) + else: + await user_api_key_cache.async_set_cache( + key=_key, + value=membership, + model_type=LiteLLM_TeamMembership, + ) + return membership + + +async def _load_team_membership_on_cache_miss( + user_id: str, + team_id: str, + cache_key: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, +) -> LiteLLM_TeamMembership | None: + try: + redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) + redis_membership: Final = _membership_from_cached_payload(redis_cached) + if not isinstance(redis_membership, _TeamMembershipCacheMiss): + return redis_membership + + return await _fetch_team_membership_from_db( + user_id=user_id, + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception: + verbose_proxy_logger.exception("Error getting team membership") + return None + + async def get_team_membership( user_id: str, team_id: str, @@ -2156,54 +2261,42 @@ async def get_team_membership( Do a isolated check for team membership vs. doing a combined key + team + user + team-membership check, as key might come in frequently for different users/teams. Larger call will slowdown query time. This way we get to cache the constant (key/team/user info) and only update based on the changing value (team membership). """ - from litellm.proxy._types import LiteLLM_TeamMembership - - if prisma_client is None: - raise Exception("No db connected") - if user_id is None or team_id is None: return None _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) - # check if in cache - cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key) - if cached == NO_TEAM_MEMBERSHIP_SENTINEL: - return None - cached_membership_obj: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership) - if cached_membership_obj is not None: - return cached_membership_obj + l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True) + l1_membership: Final = _membership_from_cached_payload(l1_cached) + if not isinstance(l1_membership, _TeamMembershipCacheMiss): + return l1_membership - # else, check db - try: - response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique( - where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, - include={"litellm_budget_table": True}, + inflight: Final[object] = _team_membership_inflight.get(_key) + if isinstance(inflight, asyncio.Task): + return _membership_from_shared_load(await asyncio.shield(inflight)) + + if prisma_client is None: + raise Exception("No db connected") + + task: Final = asyncio.ensure_future( + _load_team_membership_on_cache_miss( + user_id=user_id, + team_id=team_id, + cache_key=_key, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) + ) + _team_membership_inflight[_key] = task - if response is None: - await user_api_key_cache.async_set_cache( - key=_key, - value=NO_TEAM_MEMBERSHIP_SENTINEL, - ttl=get_management_object_ttl(user_api_key_cache), - ) - return None + def _clear_inflight(_done: object) -> None: + if _team_membership_inflight.get(_key) is task: + _team_membership_inflight.pop(_key, None) - _response: Final = LiteLLM_TeamMembership.model_validate(response.dict()) - await user_api_key_cache.async_set_cache( - key=_key, - value=_response, - model_type=LiteLLM_TeamMembership, - ) - - return _response - except Exception: - verbose_proxy_logger.exception( - "Error getting team membership for user_id: %s, team_id: %s", - user_id, - team_id, - ) - return None + task.add_done_callback(_clear_inflight) + return _membership_from_shared_load(await asyncio.shield(task)) def model_in_access_group(model: str, team_models: list[str] | None, llm_router: Router | None) -> bool: @@ -2662,6 +2755,12 @@ async def invalidate_team_member_spend_state( publish_auth_cache_invalidation, ) + inflight: Final[object] = _team_membership_inflight.pop( + team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None + ) + if isinstance(inflight, asyncio.Task) and inflight is not asyncio.current_task(): + await asyncio.wait((inflight,)) + if new_spend is not None: from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache @@ -4116,18 +4215,21 @@ async def _team_member_granted_models( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> Sequence[str]: """The member's own ``allowed_models`` scope; empty when the member is not narrowed below the team.""" if team_object is None or valid_token.user_id is None: return () - team_membership: Final = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + if not team_membership_loaded: + team_membership = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) return () if team_membership is None else _member_allowed_models(team_membership) @@ -4163,6 +4265,8 @@ async def _granted_model_lists( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> tuple[Sequence[str], ...]: """One model allowlist per level that participates in authorizing the request.""" return ( @@ -4174,6 +4278,8 @@ async def _granted_model_lists( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=team_membership_loaded, ), project_object.models if project_object is not None else (), await _org_granted_models( @@ -4268,6 +4374,8 @@ async def collect_matched_model_access_groups( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> tuple[str, ...]: """ The budgeted model access groups that authorized this request, sorted and deduplicated. @@ -4313,6 +4421,8 @@ async def collect_matched_model_access_groups( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=team_membership_loaded, ) for granted_model in granted_models ) @@ -4328,6 +4438,8 @@ async def stamp_matched_model_access_groups( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> tuple[str, ...]: """Record the groups that authorized this request on its auth object, for the post-call spend writer and the reservation counters, and hand them back for the budget check.""" @@ -4344,6 +4456,8 @@ async def stamp_matched_model_access_groups( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=team_membership_loaded, ) except Exception as e: # noqa: BLE001 # fail-safe: attribution is spend telemetry, it must never break auth verbose_proxy_logger.debug("model access group attribution failed: %s", e) @@ -5146,6 +5260,8 @@ async def _check_team_member_budget( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ): """Check if team member is over their max budget within the team.""" if ( @@ -5154,23 +5270,25 @@ async def _check_team_member_budget( and valid_token is not None and valid_token.user_id is not None ): - team_membership: Final = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + if not team_membership_loaded: + team_membership = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + loaded_membership = team_membership # Per-member override wins; otherwise fall back to the team-level # default configured via team.metadata["team_member_budget_id"]. team_member_budget: float | None = None if ( - team_membership is not None - and team_membership.litellm_budget_table is not None - and team_membership.litellm_budget_table.max_budget is not None + loaded_membership is not None + and loaded_membership.litellm_budget_table is not None + and loaded_membership.litellm_budget_table.max_budget is not None ): - team_member_budget = team_membership.litellm_budget_table.max_budget + team_member_budget = loaded_membership.litellm_budget_table.max_budget else: default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): @@ -5189,7 +5307,7 @@ async def _check_team_member_budget( team_member_budget = default_budget.max_budget if team_member_budget is not None: - team_member_spend = (team_membership.spend if team_membership is not None else 0.0) or 0.0 + team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 # Read from cross-pod counter (Redis-first) if available from litellm.proxy.proxy_server import get_current_spend @@ -5218,6 +5336,8 @@ async def _check_team_member_model_access( prisma_client: Optional["PrismaClient"], user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> None: """ Check if a team member's per-member model scope allows access to the requested model. @@ -5228,22 +5348,24 @@ async def _check_team_member_model_access( if valid_token.user_id is None or team_object.team_id is None: return - team_membership: Final = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + if not team_membership_loaded: + team_membership = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + loaded_membership = team_membership if ( - team_membership is None - or team_membership.litellm_budget_table is None - or not team_membership.litellm_budget_table.allowed_models + loaded_membership is None + or loaded_membership.litellm_budget_table is None + or not loaded_membership.litellm_budget_table.allowed_models ): return # no per-member restriction — inherit team-level check - member_allowed_models: Final[list[str]] = team_membership.litellm_budget_table.allowed_models + member_allowed_models: Final[list[str]] = loaded_membership.litellm_budget_table.allowed_models try: _can_object_call_model( model=model, diff --git a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py index 3e5d8fb311d..e54e07b6a1b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py +++ b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py @@ -44,7 +44,7 @@ class JavelinGuardrail(CustomGuardrail): application: str | None = None, **kwargs, ): - f""" + """ Initialize the JavelinGuardrail class. This calls: {api_base}/{api_version}/guardrail/{guardrail_name}/apply diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py index a911c78ddc3..95db4bd4f77 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py @@ -15,7 +15,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" # We check the raw guardrail dict because LitellmParams normalizes None → False, # making it impossible to distinguish "not set" from "explicitly false" via litellm_params. _raw_default_on: Final = cast(dict[str, Any], guardrail).get("litellm_params", {}).get("default_on") - _default_on: Final = False if _raw_default_on is False else True + _default_on: Final = _raw_default_on is not False _callback: Final = MCPEndUserPermissionGuardrail( guardrail_name=guardrail.get("guardrail_name", ""), diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 2ca92d2c37e..ebf7c8f6dce 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -252,7 +252,7 @@ class _RawTeamRow(_TeamIdRow, _ModelDumpRow, _ObjectPermissionRow, _TeamBudgetRo @property def members_with_roles( self, - ) -> Sequence[dict[str, object]] | None: ... + ) -> Sequence[dict[str, object]] | None: ... # mutable-ok: prisma deserializes this JSON column into plain dicts @property def organization_id(self) -> str | None: ... @@ -431,27 +431,26 @@ async def _refresh_cached_team( ) +async def _can_manage_team( + team_obj: LiteLLM_TeamTable, + user_api_key_dict: UserAPIKeyAuth, +) -> bool: + """True for a proxy admin, an admin of this team, or an org admin for the team's organization.""" + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return True + + if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj): + return True + + return await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj) + + async def _verify_team_access( team_obj: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth, ) -> None: - """ - Verify the caller is authorized to manage the given team. - - Access is granted if: - - Caller is a proxy admin, OR - - Caller is an org admin for the team's organization, OR - - Caller is a team admin of this team - - Raises HTTPException(403) otherwise. - """ - if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: - return - - if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj): - return - - if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj): + """Raise HTTPException(403) unless the caller can manage the given team.""" + if await _can_manage_team(team_obj=team_obj, user_api_key_dict=user_api_key_dict): return raise HTTPException( @@ -2197,7 +2196,7 @@ async def update_team( if "metadata" in updated_kv: stored_metadata: Final[Mapping[str, JsonValue] | None] = ( - { + { # mutable-ok: the validator payload's isinstance guard requires a plain dict key: value for key, value in existing_team_row.metadata.items() if key not in TeamMemberBudgetHandler.SYSTEM_MANAGED_METADATA_KEYS @@ -2881,7 +2880,11 @@ async def _resolve_existing_member_user_ids( return frozenset() found: Final = await _user_id_rows_db(UserRepository(prisma_client)).find_many( - where={"user_id": {"in": sorted(requested_user_ids)}} + where={ # mutable-ok: Prisma query filters are dict-shaped + "user_id": { # mutable-ok: Prisma query filters are dict-shaped + "in": sorted(requested_user_ids) + } + } ) return frozenset(user.user_id for user in found or () if user.user_id is not None) @@ -2935,7 +2938,7 @@ def _validate_member_user_id_provisioning( remaining: Final = len(unknown_user_ids) - _MAX_REPORTED_UNKNOWN_USER_IDS raise HTTPException( status_code=403, - detail={ + detail={ # mutable-ok: HTTPException detail must be a plain mapping to keep this route's {"error": ...} response shape "error": ( "Only proxy admins can add a user_id that does not exist yet: {}{}. " "Add the member by user_email to invite a new user, or ask a proxy admin " @@ -2951,7 +2954,11 @@ def _members_audit_value(members: Sequence[Member]) -> str: The audit-log columns hold a JSON object, so the member list is nested under a key rather than serialized as a top-level array. """ - return safe_dumps({"members_with_roles": tuple(member.model_dump() for member in members)}) + return safe_dumps( + { # mutable-ok: the audit-log JSON column rejects a top-level array, so this value must be an object + "members_with_roles": tuple(member.model_dump() for member in members) + } + ) async def _create_team_member_add_audit_logs( @@ -3638,7 +3645,7 @@ def _check_not_resetting_own_spend(user_id: str, user_api_key_dict: UserAPIKeyAu def _raise_reset_spend_error(status_code: int, message: str) -> NoReturn: - detail: Final = {"error": message} + detail: Final = {"error": message} # mutable-ok: HTTPException.detail takes a dict raise HTTPException(status_code=status_code, detail=detail) @@ -3672,7 +3679,7 @@ def _validate_team_member_reset_spend_value( @router.post( "/team/{team_id}/member/{user_id}/reset_spend", - tags=["team management"], + tags=["team management"], # mutable-ok: FastAPI's `tags` param is typed as list[str], not Sequence dependencies=(Depends(user_api_key_auth),), ) @management_endpoint_wrapper @@ -3708,10 +3715,12 @@ async def reset_team_member_spend_fn( await _verify_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict) _check_not_resetting_own_spend(user_id=user_id, user_api_key_dict=user_api_key_dict) - membership_where: Final = {"user_id_team_id": {"user_id": user_id, "team_id": team_id}} + membership_where: Final = { # mutable-ok: prisma client requires a plain dict where= argument + "user_id_team_id": {"user_id": user_id, "team_id": team_id} # mutable-ok: same prisma where= argument + } _membership_row: Final = await _team_membership_db(prisma_client).find_unique( where=membership_where, - include={"litellm_budget_table": True}, + include={"litellm_budget_table": True}, # mutable-ok: prisma client requires a plain dict include= argument ) if _membership_row is None: _raise_reset_spend_error(status.HTTP_404_NOT_FOUND, f"User {user_id} is not a member of team {team_id}.") @@ -3722,7 +3731,7 @@ async def reset_team_member_spend_fn( await _team_membership_db(prisma_client).update( where=membership_where, - data={"spend": reset_to}, + data={"spend": reset_to}, # mutable-ok: prisma client requires a plain dict data= argument ) await invalidate_team_member_spend_state( @@ -3732,7 +3741,7 @@ async def reset_team_member_spend_fn( new_spend=reset_to, ) - return { + return { # mutable-ok: matches this router's established untyped-response-dict convention "team_id": team_id, "user_id": user_id, "spend": reset_to, @@ -4336,7 +4345,15 @@ async def _hydrate_member_user_details( """Attach ``user_alias`` and fill in a missing ``user_email`` from ``LiteLLM_UserTable`` in one query.""" user_ids: Final = frozenset(m.user_id for m in members if m.user_id is not None) user_rows: Final[Sequence[prisma_models.LiteLLM_UserTable]] = ( - await _user_db(prisma_client).find_many(where={"user_id": {"in": sorted(user_ids)}}) if user_ids else () + await _user_db(prisma_client).find_many( + where={ # mutable-ok: Prisma query filters are dict-shaped + "user_id": { # mutable-ok: Prisma query filters are dict-shaped + "in": sorted(user_ids) + } + } + ) + if user_ids + else () ) user_by_id: Final = MappingProxyType({u.user_id: u for u in user_rows}) @@ -4352,6 +4369,20 @@ async def _hydrate_member_user_details( return tuple(hydrate(m) for m in members) +class _OrganizationModelsRow(BaseModel): + models: list[str] = [] # mutable-ok: pydantic field default + + +class _TeamRowWithOrganization(BaseModel): + litellm_organization_table: _OrganizationModelsRow | None = None + + +def _parent_organization_models(team_row: BaseModel) -> list[str] | None: + """Return the parent org's model allow-list, or None when the team has no org.""" + organization: Final = _TeamRowWithOrganization.model_validate(team_row.model_dump()).litellm_organization_table + return organization.models if organization is not None else None + + async def _resolve_team_access_group_resources( _team_info: TeamInfoResponseObjectTeamTable, ) -> TeamInfoResponseObjectTeamTable: @@ -4423,7 +4454,11 @@ async def team_info( try: team_info: BaseModel | None = await _team_db(prisma_client).find_unique( where={"team_id": team_id}, - include={"litellm_model_table": True, "object_permission": True}, + include={ + "litellm_model_table": True, + "object_permission": True, + "litellm_organization_table": True, + }, ) if team_info is None: raise Exception @@ -4432,9 +4467,12 @@ async def team_info( status_code=status.HTTP_404_NOT_FOUND, detail={"message": f"Team not found, passed team id: {team_id}."}, ) - await validate_membership( - user_api_key_dict=user_api_key_dict, - team_table=LiteLLM_TeamTable.model_validate(team_info.model_dump()), + team_table: Final = LiteLLM_TeamTable.model_validate(team_info.model_dump()) + await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_table) + organization_models: Final[list[str] | None] = ( + _parent_organization_models(team_info) + if await _can_manage_team(team_obj=team_table, user_api_key_dict=user_api_key_dict) + else None ) ## GET ALL KEYS ## @@ -4493,7 +4531,12 @@ async def team_info( prisma_client=prisma_client, members=resolved_team_info.members_with_roles, ) - hydrated_team_info: Final = resolved_team_info.model_copy(update={"members_with_roles": hydrated_members}) + hydrated_team_info: Final = resolved_team_info.model_copy( + update={ + "members_with_roles": hydrated_members, + "organization_models": organization_models, + } + ) response_object: Final = TeamInfoResponseObject( team_id=team_id, @@ -4746,7 +4789,7 @@ async def unblock_team( @router.get( "/team/metadata_schema", - tags=["team management"], + tags=["team management"], # mutable-ok: fastapi's decorator signature types tags as a list dependencies=(Depends(user_api_key_auth),), response_model=TeamMetadataSchemaResponse, ) @@ -6028,13 +6071,13 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi def _daily_activity_error(*, status_code: int, message: str) -> HTTPException: """Single construction site for the `{"error": ...}` detail shape the /team/daily/activity endpoints have always returned.""" - return HTTPException(status_code=status_code, detail={"error": message}) + return HTTPException(status_code=status_code, detail={"error": message}) # mutable-ok: FastAPI JSON detail class _TeamDailyActivityScope(NamedTuple): team_ids: list[str] | None # mutable-ok: downstream daily-activity signatures take str | list unions exclude_team_ids: list[str] | None # mutable-ok: downstream daily-activity signatures take str | list unions - team_alias_metadata: dict[str, dict[str, object]] + team_alias_metadata: dict[str, dict[str, object]] # mutable-ok: entity_metadata_field shape api_key_filter: str | list[str] | None # mutable-ok: downstream daily-activity signatures take str | list unions @@ -6337,7 +6380,7 @@ class _TeamUserSpendDbRow(TypedDict): @router.get( "/team/spend/by_user", response_model=TeamUserSpendResponse, - tags=["team management"], + tags=["team management"], # mutable-ok: fastapi route tags must be a list ) async def get_team_spend_by_user( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], diff --git a/litellm/proxy/management_helpers/team_metadata_validation.py b/litellm/proxy/management_helpers/team_metadata_validation.py index 8755522e206..8bb32696857 100644 --- a/litellm/proxy/management_helpers/team_metadata_validation.py +++ b/litellm/proxy/management_helpers/team_metadata_validation.py @@ -115,13 +115,13 @@ async def run_team_metadata_validation( "error": f"custom_team_metadata_validate is an Enterprise feature. {CommonProxyErrors.not_premium_user.value}" }, ) - if not ( - inspect.iscoroutinefunction(validator) or inspect.iscoroutinefunction(getattr(validator, "__call__", None)) - ): - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": "custom_team_metadata_validate must be an async function"}, - ) + if not inspect.iscoroutinefunction(validator): + validator_call: Final = getattr(validator, "__call__", None) # noqa: B004 # value unwrap for the functor check + if not inspect.iscoroutinefunction(validator_call): + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": "custom_team_metadata_validate must be an async function"}, + ) try: raw_result: Final = await asyncio.wait_for(validator(payload), timeout=timeout_seconds) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 41f51cdb769..2a54cc4ced3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -274,6 +274,8 @@ from litellm.constants import ( PROXY_BUDGET_RESCHEDULER_MAX_TIME, PROXY_BUDGET_RESCHEDULER_MIN_TIME, PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + REALTIME_SESSION_FAILURE_LOGGED_KEY, + REALTIME_SESSION_SUCCESS_LOGGED_KEY, ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG, USER_SPEND_ALERTS_JOB_ID, WEEKLY_SPEND_REPORT_JOB_ID, @@ -11895,6 +11897,13 @@ async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth ) +async def _release_realtime_max_parallel_slot(user_api_key_dict: UserAPIKeyAuth) -> None: + release_like_http_disconnect: Final = ( + proxy_logging_obj._arelease_max_parallel_requests_on_disconnect # pyright: ignore[reportPrivateUsage] # shared + ) + await release_like_http_disconnect(user_api_key_dict) + + async def _reject_realtime_session( websocket: WebSocket, user_api_key_dict: UserAPIKeyAuth, @@ -11914,6 +11923,7 @@ async def _reject_realtime_session( await websocket.close(code=code, reason=reason) finally: await _release_realtime_budget_reservation(user_api_key_dict) + await _release_realtime_max_parallel_slot(user_api_key_dict) @app.websocket("/openai/v1/realtime") @@ -12017,6 +12027,9 @@ async def realtime_websocket_endpoint( websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e) ) return + except BaseException: + await _release_realtime_max_parallel_slot(user_api_key_dict) + raise # Phase 2: route to upstream LLM. try: @@ -12046,12 +12059,10 @@ async def realtime_websocket_endpoint( except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone") finally: - from litellm.litellm_core_utils.realtime_streaming import ( - REALTIME_SESSION_SUCCESS_LOGGED_KEY, - ) - if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): await _release_realtime_budget_reservation(user_api_key_dict) + if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY): + await _release_realtime_max_parallel_slot(user_api_key_dict) ###################################################################### diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 3ef21996b9c..4bcdf6aad22 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -159,7 +159,7 @@ def _get_spend_logs_metadata( requester_ip_address=None, additional_usage_values=None, applied_guardrails=None, - status=None or "success", + status="success", error_information=None, proxy_server_request=None, batch_models=None, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8ac20929337..67bd215329b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4,6 +4,7 @@ import copy import hashlib import inspect import json +import math import os import smtplib import ssl @@ -6384,7 +6385,7 @@ class PrismaClient: return None try: value: Final = float(response_time_ms) - return value if value == value and value not in (float("inf"), float("-inf")) else None + return value if math.isfinite(value) else None except (ValueError, TypeError): verbose_proxy_logger.warning("Invalid response_time_ms value: %s", response_time_ms) return None diff --git a/litellm/router.py b/litellm/router.py index 86a276c96d7..7e78fb608d4 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1532,6 +1532,18 @@ class Router: return False return sum(len(self.model_name_to_deployment_indices.get(member) or ()) for member in group.models) > 1 + def team_model_has_alternatives(self, deployment_id: str) -> bool: + deployment: Final = self.get_deployment(model_id=deployment_id) + if deployment is None: + return False + team_id: Final = deployment.model_info.team_id + public_model_name: Final = deployment.model_info.team_public_model_name + if team_id is None or public_model_name is None: + return False + sibling_indices: Final = self.team_model_to_deployment_indices.get((team_id, public_model_name)) or () + routable_siblings: Final = self._filter_blocked_deployments([self.model_list[idx] for idx in sibling_indices]) + return len(routable_siblings) > 1 + _OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY}) def _get_request_routing_strategy_override(self, request_kwargs: dict | None) -> str | None: @@ -4087,16 +4099,16 @@ class Router: models: Final = [m.strip() for m in model.split(",")] async def _async_completion_no_exceptions( - model: str, messages: list[dict[str, str]], stream: bool, **kwargs: Any + model_name: str, messages: list[dict[str, str]], stream: bool, **kwargs: Any ) -> ModelResponse | CustomStreamWrapper | Exception: """ Wrapper around self.acompletion that catches exceptions and returns them as a result """ try: - result = await self.acompletion(model=model, messages=messages, stream=stream, **kwargs) + result = await self.acompletion(model=model_name, messages=messages, stream=stream, **kwargs) return result except asyncio.CancelledError: - verbose_router_logger.debug("Received 'task.cancel'. Cancelling call w/ model=%s.", model) + verbose_router_logger.debug("Received 'task.cancel'. Cancelling call w/ model=%s.", model_name) raise except Exception as e: return e @@ -4123,9 +4135,9 @@ class Router: except KeyError: pass - for model in models: + for model_name in models: task = asyncio.create_task( - _async_completion_no_exceptions(model=model, messages=messages, stream=stream, **kwargs) + _async_completion_no_exceptions(model_name=model_name, messages=messages, stream=stream, **kwargs) ) pending_tasks.append(task) diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index f722b6fd20c..027f0a9ca05 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -343,8 +343,9 @@ def _should_cooldown_deployment( model_group: Final = litellm_router_instance.get_model_group(id=deployment) is_single_deployment_model_group = False if model_group is not None and len(model_group) == 1: - is_single_deployment_model_group = not litellm_router_instance.routing_group_has_alternatives( - requested_model_group + is_single_deployment_model_group = not ( + litellm_router_instance.routing_group_has_alternatives(requested_model_group) + or litellm_router_instance.team_model_has_alternatives(deployment) ) ## CHECK DEPLOYMENT-LEVEL POLICY FIRST (overrides router-level) diff --git a/litellm/utils.py b/litellm/utils.py index 3a0fec1d027..c9f26a92516 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6260,7 +6260,7 @@ def function_to_dict(input_function) -> dict: "enum": param_enum, } - parameters[param_name] = dict([(k, v) for k, v in param_dict.items() if isinstance(v, str)]) + parameters[param_name] = {k: v for k, v in param_dict.items() if isinstance(v, str)} # Check if the parameter has no default value (i.e., it's required) if param.default == param.empty: diff --git a/ruff.toml b/ruff.toml index 3ac4c1fc94d..fab3fe27aed 100644 --- a/ruff.toml +++ b/ruff.toml @@ -4,11 +4,12 @@ lint.ignore = ["F405", "E402", "F403"] # That gives editors and `ruff check --fix` the diagnostic, which the gate script cannot. lint.extend-select = [ "T20", "PGH004", "RUF008", "RUF009", "RUF100", - "B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", "PIE800", "PLC0208", - "PLR0402", "PLR1711", "PLR1730", "PLR2044", "PLW0133", "PYI030", "PYI041", "PYI064", "RET501", - "RUF010", "RUF022", "RUF023", "RUF051", "S113", "SIM114", "SIM118", "TC005", "UP006", "UP007", - "UP008", - "UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045", + "B004", "B018", "B021", "B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", + "PIE800", "PLC0208", "PLR0124", "PLR0402", "PLR0206", "PLR1704", "PLR1711", "PLR1730", "PLR2044", + "PLW0133", "PYI030", "PYI041", "PYI064", "RET501", "RUF010", "RUF022", "RUF023", "RUF051", "S113", + "SIM114", "SIM118", "SIM201", "SIM211", "SIM222", "TC005", "UP006", "UP007", "UP008", + "UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045", + "C404", "C419", ] # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip # `# noqa` directives that protect rules enforced elsewhere. List those codes as external diff --git a/tests/proxy_behavior/auth/test_auth_object_prefetch.py b/tests/proxy_behavior/auth/test_auth_object_prefetch.py index 59e9585a296..cfa958500af 100644 --- a/tests/proxy_behavior/auth/test_auth_object_prefetch.py +++ b/tests/proxy_behavior/auth/test_auth_object_prefetch.py @@ -66,6 +66,7 @@ async def test_join_binds_the_membership_to_the_requested_team(prisma): cache = _frozen_cache() refs = AuthObjectRefs(user_id=user_id, team_id=team_a, membership_user_id=user_id, organization_id=org_id) await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma) + assert cache.in_memory_cache.get_cache(f"org_id:{org_id}") is not None dead_db = _dead_db() membership = await get_team_membership( diff --git a/tests/rust-python-harness/AGENTS.md b/tests/rust-python-harness/AGENTS.md index 017668d4289..71e17541fd2 100644 --- a/tests/rust-python-harness/AGENTS.md +++ b/tests/rust-python-harness/AGENTS.md @@ -63,7 +63,7 @@ tests/rust-python-harness/ - Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr` - `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases - `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses -- `trace_parity/` compares mapped operations, call counts, and required execution ordering; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`) +- `trace_parity/` prints every collected Python call under `litellm/` and every Rust span without comparing them; mappings only filter the separate unit-test mapping strategy. Before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`) - E2E and trace strategies load their registered module cases and run surface-specific execution from their folders - `unit_tests_mapping/contracts.py` owns typed harness-side mapping contracts, per-function contracts live below `cases/`, and `mappings.py` exports the registry; live test discovery derives unmapped Python and Rust-only tests without an exhaustive manifest - `unit_tests_mapping/runner.py` validates confirmed mappings against the live Python and Rust inventories and attaches the derived status report diff --git a/tests/rust-python-harness/cli/__init__.py b/tests/rust-python-harness/cli/__init__.py index 13b995825dd..d2bfdc55b19 100644 --- a/tests/rust-python-harness/cli/__init__.py +++ b/tests/rust-python-harness/cli/__init__.py @@ -58,16 +58,29 @@ def _strategy_command(strategy: Strategy) -> click.Command: help=runner_argument.help, ) ) + for runner_option in strategy.definition.runner_options: + name: Final = runner_option.option.removeprefix("--").replace("-", "_") + params.append( + click.Option( + (runner_option.option, name), + type=click.Choice(runner_option.choices), + help=runner_option.help, + ) + ) def run_strategy( sdk_functions: tuple[str, ...], surface: str | None = None, runner_args: tuple[str, ...] = (), + **runner_options: str | None, ) -> int: selected_functions: Final = cast(frozenset[SdkFunction], frozenset(sdk_functions)) selected_surface: Final = cast(Surface | None, surface) cases: Final = select_cases((strategy,), selected_functions, selected_surface) - return run_command((strategy,), cases, runner_args) + option_args: Final = tuple( + f"--{name.replace('_', '-')}={value}" for name, value in runner_options.items() if value is not None + ) + return run_command((strategy,), cases, (*runner_args, *option_args)) return click.Command( strategy.id, diff --git a/tests/rust-python-harness/cli/catalog.py b/tests/rust-python-harness/cli/catalog.py index 03eb032d9c6..073873d0121 100644 --- a/tests/rust-python-harness/cli/catalog.py +++ b/tests/rust-python-harness/cli/catalog.py @@ -21,9 +21,7 @@ def _load_strategy_module(name: str, folder: Path, prefix: str | None) -> Module if prefix is not None: return importlib.import_module(f"{prefix}.{name}") module_name: Final = _synthetic_module_name(folder) - spec: Final = importlib.util.spec_from_file_location( - module_name, folder / "__init__.py" - ) + spec: Final = importlib.util.spec_from_file_location(module_name, folder / "__init__.py") if spec is None or spec.loader is None: raise ValueError(f"{folder}: cannot load strategy package") module: Final = importlib.util.module_from_spec(spec) @@ -59,9 +57,7 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy: if duplicates: raise ValueError(f"{folder}: duplicate strategy cases: {duplicates}") expected: Final = frozenset( - (surface, function) - for surface in (definition.surfaces or (None,)) - for function in SDK_FUNCTIONS + (surface, function) for surface in (definition.surfaces or (None,)) for function in SDK_FUNCTIONS ) actual: Final = frozenset(keys) if actual != expected: @@ -73,8 +69,7 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy: incompatible: Final = tuple( (case.surface, case.sdk_function) for case in definition.cases - if case.spec.disposition is CaseDisposition.RUNNABLE - and not isinstance(case.spec, definition.runnable_spec) + if case.spec.disposition is CaseDisposition.RUNNABLE and not isinstance(case.spec, definition.runnable_spec) ) if incompatible: raise ValueError(f"{folder}: runnable cases do not match {definition.runnable_spec.__name__}: {incompatible}") @@ -102,14 +97,10 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy: def load_catalog(root: Path | None = None) -> tuple[Strategy, ...]: resolved: Final = STRATEGIES_ROOT if root is None else root prefix: Final = _STRATEGIES_PACKAGE.__name__ if resolved == STRATEGIES_ROOT else None - folders: Final = tuple( - info.name for info in pkgutil.iter_modules([str(resolved)]) if info.ispkg - ) + folders: Final = tuple(info.name for info in pkgutil.iter_modules([str(resolved)]) if info.ispkg) if not folders: raise ValueError(f"No strategy packages found below {resolved}") - strategies: Final = tuple( - _load_strategy(name, resolved / name, prefix) for name in sorted(folders) - ) + strategies: Final = tuple(_load_strategy(name, resolved / name, prefix) for name in sorted(folders)) ids: Final = [strategy.id for strategy in strategies] if len(set(ids)) != len(ids): raise ValueError(f"Duplicate strategy id in {resolved}") diff --git a/tests/rust-python-harness/cli/commands.py b/tests/rust-python-harness/cli/commands.py index f94c3277dc2..e51bbb5d966 100644 --- a/tests/rust-python-harness/cli/commands.py +++ b/tests/rust-python-harness/cli/commands.py @@ -21,8 +21,7 @@ def select_cases( case for strategy in strategies for case in strategy.cases - if (not sdk_functions or case.sdk_function in sdk_functions) - and (surface is None or case.surface == surface) + if (not sdk_functions or case.sdk_function in sdk_functions) and (surface is None or case.surface == surface) ) @@ -32,8 +31,7 @@ def run_command( runner_args: Sequence[str] = (), ) -> int: grouped: Final = { - strategy.id: tuple(case for case in cases if case.strategy_id == strategy.id) - for strategy in strategies + strategy.id: tuple(case for case in cases if case.strategy_id == strategy.id) for strategy in strategies } visible: Final = tuple(strategy for strategy in strategies if grouped[strategy.id]) runners: Final = tuple(replace(strategy, cases=grouped[strategy.id]) for strategy in visible) diff --git a/tests/rust-python-harness/cli/test_cli.py b/tests/rust-python-harness/cli/test_cli.py index 226c89843d0..5641aa8a539 100644 --- a/tests/rust-python-harness/cli/test_cli.py +++ b/tests/rust-python-harness/cli/test_cli.py @@ -242,7 +242,7 @@ def _assert_unavailable_cell(strategy: Strategy, case: HarnessCase, section_titl def test_every_unavailable_case_finishes_and_explains_itself() -> None: section_titles: Final = { "e2e_parity": "End-to-end parity outcomes", - "trace_parity": "trace comparisons", + "trace_parity": "traces", "unit_tests_mapping": "Python/Rust unit-test mappings", "unit_tests_parity": "Python backend parity outcomes", "unit_tests_rust": "Native Rust unit-test outcomes", @@ -359,6 +359,25 @@ def test_strategy_command_forwards_repeated_filters_and_runner_arguments( ] +def test_trace_command_forwards_engine_and_scenario(monkeypatch: pytest.MonkeyPatch) -> None: + cli: Final = importlib.import_module("tests.rust-python-harness.cli") + captured: list[tuple[str, ...]] = [] + + def capture_run( + strategies: Sequence[Strategy], + cases: Sequence[HarnessCase], + runner_args: Sequence[str] = (), + ) -> int: + del strategies, cases + captured.append(tuple(runner_args)) + return 0 + + monkeypatch.setattr(cli, "run_command", capture_run) + + assert main(["run", "trace_parity", "--scenario", "async-mistral", "--engine", "python"]) == 0 + assert captured == [("async-mistral", "--engine=python")] + + def test_omitted_surface_selects_every_strategy_surface(monkeypatch: pytest.MonkeyPatch) -> None: cli: Final = importlib.import_module("tests.rust-python-harness.cli") selected: list[str] = [] diff --git a/tests/rust-python-harness/conftest.py b/tests/rust-python-harness/conftest.py index d50d0fa4204..c487b25a9cc 100644 --- a/tests/rust-python-harness/conftest.py +++ b/tests/rust-python-harness/conftest.py @@ -19,9 +19,7 @@ def subprocess_test_environment(monkeypatch: pytest.MonkeyPatch) -> None: def cargo_project(tmp_path: Path) -> Callable[[str, str], Path]: def create(package: str, source: str) -> Path: manifest: Final = tmp_path / "Cargo.toml" - manifest.write_text( - f'[package]\nname = "{package}"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n' - ) + manifest.write_text(f'[package]\nname = "{package}"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n') (tmp_path / "src").mkdir() (tmp_path / "src/lib.rs").write_text(source) return manifest diff --git a/tests/rust-python-harness/shared/native_build.py b/tests/rust-python-harness/shared/native_build.py index 8693cf3bac2..f67488cecb4 100644 --- a/tests/rust-python-harness/shared/native_build.py +++ b/tests/rust-python-harness/shared/native_build.py @@ -16,6 +16,11 @@ _RUST_ROOT: Final = "litellm-rust" _LOCKFILE: Final = "Cargo.lock" _SOURCE_SUFFIXES: Final = frozenset({".rs", ".toml"}) _FAILURE_OUTPUT_LINES: Final = 15 +_TRACE_CHECK: Final = ( + "from litellm.rust_bridge import get_native_bridge; " + "bridge = get_native_bridge(); " + "raise SystemExit(0 if bridge is not None and getattr(bridge, '_trace', None) is not None else 1)" +) def needs_rebuild(native_mtime: float | None, newest_source_mtime: float | None) -> bool: @@ -73,6 +78,17 @@ def _rebuild(repo_root: Path) -> tuple[bool, str]: return completed.returncode == 0, "\n".join(lines[-_FAILURE_OUTPUT_LINES:]) +def _installed_bridge_has_trace(repo_root: Path) -> bool: + completed: Final = subprocess.run( + (sys.executable, "-c", _TRACE_CHECK), + cwd=repo_root, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=False, + ) + return completed.returncode == 0 + + def trace_bridge_error() -> str | None: bridge: Final = get_native_bridge() if bridge is None: @@ -85,7 +101,10 @@ def trace_bridge_error() -> str | None: def ensure_trace_bridge(repo_root: Path) -> str | None: native_path: Final = _native_module_path() native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None - if needs_rebuild(native_mtime, _newest_source_mtime(repo_root)): + rebuild_required: Final = needs_rebuild( + native_mtime, _newest_source_mtime(repo_root) + ) or not _installed_bridge_has_trace(repo_root) + if rebuild_required: print(f"Rebuilding native Rust bridge ({BRIDGE_FEATURE} feature)...", flush=True) succeeded: Final output: Final diff --git a/tests/rust-python-harness/shared/parity/fixtures/recording.py b/tests/rust-python-harness/shared/parity/fixtures/recording.py index ee0d4b6de7b..cb6ed968b90 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/recording.py +++ b/tests/rust-python-harness/shared/parity/fixtures/recording.py @@ -191,6 +191,7 @@ class _RecordingHandler(LocalHttpHandler): self.end_headers() self.wfile.write(body) + def _recording_provider(spec: UpstreamEndpoint) -> AbstractContextManager[_RecordingProvider]: return serve_in_thread(_RecordingProvider(spec)) diff --git a/tests/rust-python-harness/shared/parity/fixtures/store.py b/tests/rust-python-harness/shared/parity/fixtures/store.py index 270af2a7625..1145b5d7d27 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/store.py +++ b/tests/rust-python-harness/shared/parity/fixtures/store.py @@ -15,6 +15,8 @@ from .cassette import deserialize_cassette, serialize_cassette from .recording import RecordedInteraction FIXTURE_SCHEMA_VERSION: Final = 1 + + class FixtureInput(Protocol): def canonical_input(self) -> dict[str, object]: ... diff --git a/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py b/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py index 4535ba05bf6..b162e4949d2 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py +++ b/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py @@ -45,8 +45,8 @@ class _Upstream(LocalHttpServer): super().__init__(("127.0.0.1", 0), _UpstreamHandler) self.response_status: Final = status -class _UpstreamHandler(LocalHttpHandler): +class _UpstreamHandler(LocalHttpHandler): def do_POST(self) -> None: length: Final = int(self.headers.get("content-length") or "0") self.rfile.read(length) @@ -59,6 +59,7 @@ class _UpstreamHandler(LocalHttpHandler): self.end_headers() self.wfile.write(body) + def _upstream(status: int = 200) -> AbstractContextManager[_Upstream]: return serve_in_thread(_Upstream(status)) diff --git a/tests/rust-python-harness/shared/parity/fixtures/test_recording.py b/tests/rust-python-harness/shared/parity/fixtures/test_recording.py index 6181f18e89a..c1e7c2af8d8 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/test_recording.py +++ b/tests/rust-python-harness/shared/parity/fixtures/test_recording.py @@ -238,6 +238,7 @@ class _ControlledUpstreamHandler(LocalHttpHandler): self.end_headers() self.wfile.write(body) + def _controlled_upstream( stream_chunks: tuple[bytes, ...] = _SSE_CHUNKS, ) -> AbstractContextManager[_ControlledUpstream]: diff --git a/tests/rust-python-harness/shared/reporting/models.py b/tests/rust-python-harness/shared/reporting/models.py index 1ebba6c9793..4a46a6d4cfe 100644 --- a/tests/rust-python-harness/shared/reporting/models.py +++ b/tests/rust-python-harness/shared/reporting/models.py @@ -180,15 +180,11 @@ class HarnessRun: @property def unique_checks(self) -> int: - return len( - {nodeid for result in self.results.values() for nodeid in result.collected} - ) + return len({nodeid for result in self.results.values() for nodeid in result.collected}) @property def completed_checks(self) -> int: - return len( - {nodeid for result in self.results.values() for nodeid in result.completed} - ) + return len({nodeid for result in self.results.values() for nodeid in result.completed}) @classmethod def from_cases(cls, cases: Iterable[HarnessCase]) -> HarnessRun: diff --git a/tests/rust-python-harness/shared/reporting/strategy.py b/tests/rust-python-harness/shared/reporting/strategy.py index 7e76f035e20..d8e9d9e5ba9 100644 --- a/tests/rust-python-harness/shared/reporting/strategy.py +++ b/tests/rust-python-harness/shared/reporting/strategy.py @@ -67,6 +67,13 @@ class RunnerArgumentDefinition: metavar: str = "ARG" +@dataclass(frozen=True, slots=True) +class RunnerOptionDefinition: + option: str + help: str + choices: tuple[str, ...] + + class StrategyRunner(Protocol): def __call__( self, @@ -90,3 +97,4 @@ class StrategyDefinition: render: StrategyRenderer surfaces: tuple[Surface, ...] = () runner_argument: RunnerArgumentDefinition | None = None + runner_options: tuple[RunnerOptionDefinition, ...] = () diff --git a/tests/rust-python-harness/shared/test_native_build.py b/tests/rust-python-harness/shared/test_native_build.py index f3e5aead846..dc08bc1a2b6 100644 --- a/tests/rust-python-harness/shared/test_native_build.py +++ b/tests/rust-python-harness/shared/test_native_build.py @@ -89,8 +89,8 @@ def test_ensure_trace_bridge_reports_failed_rebuild(tmp_path: Final, monkeypatch assert "boom" in message -def test_ensure_trace_bridge_flags_missing_trace_feature_without_rebuild( - tmp_path: Final, monkeypatch: pytest.MonkeyPatch +def test_ensure_trace_bridge_rebuilds_when_trace_feature_is_missing( + tmp_path: Final, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] ) -> None: native: Final = tmp_path / "_native.abi3.so" native.write_bytes(b"") @@ -105,12 +105,17 @@ def test_ensure_trace_bridge_flags_missing_trace_feature_without_rebuild( state.rebuilt = True return True, "" + def fake_get_native_bridge() -> SimpleNamespace: + assert state.rebuilt + return SimpleNamespace(_trace=object()) + monkeypatch.setattr(native_build, "_native_module_path", lambda: native) monkeypatch.setattr(native_build, "_rebuild", fake_rebuild) - monkeypatch.setattr(native_build, "get_native_bridge", lambda: SimpleNamespace(_trace=None)) + monkeypatch.setattr(native_build, "_installed_bridge_has_trace", lambda repo_root: False) + monkeypatch.setattr(native_build, "get_native_bridge", fake_get_native_bridge) message: Final = native_build.ensure_trace_bridge(tmp_path) - assert message is not None - assert "_trace" in message - assert state.rebuilt is False + assert message is None + assert state.rebuilt is True + assert "Rebuilding native Rust bridge" in capsys.readouterr().out diff --git a/tests/rust-python-harness/shared/tracing/profiler.py b/tests/rust-python-harness/shared/tracing/profiler.py index abfb6a2425d..e72a36a5289 100644 --- a/tests/rust-python-harness/shared/tracing/profiler.py +++ b/tests/rust-python-harness/shared/tracing/profiler.py @@ -2,7 +2,8 @@ from __future__ import annotations import sys import threading -from collections.abc import Generator, Iterator, Mapping +import warnings +from collections.abc import Callable, Generator, Iterator, Mapping from contextlib import contextmanager from dataclasses import dataclass from functools import lru_cache @@ -32,6 +33,7 @@ class PythonProfiler: self._source_root: Final = str(source_root.resolve()) + "/" self._seen_frames: Final[set[FrameType]] = set() self._event_ids: Final[dict[FrameType, int]] = {} + self._lock: Final = threading.Lock() self.events: Final[list[FunctionTraceEvent]] = [] def __call__(self, frame: FrameType, event: str, _arg: object) -> None: @@ -40,14 +42,15 @@ class PythonProfiler: function_name: Final = self.function_name(frame) if function_name is None: return - event_id: Final = len(self.events) - parent_id: Final = next( - (self._event_ids[ancestor] for ancestor in _frame_ancestors(frame) if ancestor in self._event_ids), - None, - ) - self._seen_frames.add(frame) - self._event_ids[frame] = event_id - self.events.append(FunctionTraceEvent(id=event_id, parent_id=parent_id, function=function_name)) + with self._lock: + event_id: Final = len(self.events) + parent_id: Final = next( + (self._event_ids[ancestor] for ancestor in _frame_ancestors(frame) if ancestor in self._event_ids), + None, + ) + self._seen_frames.add(frame) + self._event_ids[frame] = event_id + self.events.append(FunctionTraceEvent(id=event_id, parent_id=parent_id, function=function_name)) def function_name(self, frame: FrameType) -> str | None: code: Final = frame.f_code @@ -137,21 +140,51 @@ def _frame_ancestors(frame: FrameType) -> Generator[FrameType]: @contextmanager -def profile_python(source_root: Path, *, threads: bool = False) -> Generator[PythonProfiler]: - profiler: Final = PythonProfiler(source_root) +def _installed_profiler(profiler: Callable[[FrameType, str, object], None], *, threads: bool) -> Generator[None]: + if threads and sys.version_info >= (3, 12): + tool_id: Final = next((slot for slot in (2, 3, 4, 0, 1, 5) if sys.monitoring.get_tool(slot) is None), None) + if tool_id is None: + raise RuntimeError("no sys.monitoring tool ID is available for Python trace collection") + + def started(_code: CodeType, _offset: int) -> None: + profiler(sys._getframe(1), "call", None) + + sys.monitoring.use_tool_id(tool_id, "litellm-python-trace") + try: + sys.monitoring.register_callback(tool_id, sys.monitoring.events.PY_START, started) + sys.monitoring.set_events(tool_id, sys.monitoring.events.PY_START) + yield + finally: + sys.monitoring.set_events(tool_id, 0) + sys.monitoring.register_callback(tool_id, sys.monitoring.events.PY_START, None) + sys.monitoring.free_tool_id(tool_id) + return + if threads: + warnings.warn( + "Python <3.12 cannot trace existing worker threads; use Python 3.12+ for complete threaded traces", + RuntimeWarning, + stacklevel=3, + ) previous_thread: Final = threading.getprofile() if threads: threading.setprofile(profiler) previous: Final = sys.getprofile() sys.setprofile(profiler) try: - yield profiler + yield finally: sys.setprofile(previous) if threads: threading.setprofile(previous_thread) +@contextmanager +def profile_python(source_root: Path, *, threads: bool = False) -> Generator[PythonProfiler]: + profiler: Final = PythonProfiler(source_root) + with _installed_profiler(profiler, threads=threads): + yield profiler + + @contextmanager def profile_python_function_usage( source_root: Path, @@ -160,14 +193,5 @@ def profile_python_function_usage( threads: bool = False, ) -> Generator[PythonFunctionUsageProfiler]: profiler: Final = PythonFunctionUsageProfiler(source_root, functions) - previous_thread: Final = threading.getprofile() - if threads: - threading.setprofile(profiler) - previous: Final = sys.getprofile() - sys.setprofile(profiler) - try: + with _installed_profiler(profiler, threads=threads): yield profiler - finally: - sys.setprofile(previous) - if threads: - threading.setprofile(previous_thread) diff --git a/tests/rust-python-harness/shared/tracing/steps.py b/tests/rust-python-harness/shared/tracing/steps.py index 2475f0fdcc5..492ffab64e5 100644 --- a/tests/rust-python-harness/shared/tracing/steps.py +++ b/tests/rust-python-harness/shared/tracing/steps.py @@ -76,7 +76,7 @@ def _span_for(engine: Engine, function: str, mappings: Sequence[TraceMapping]) - def pipeline_projection( - engine: Engine, events: Sequence[FunctionTraceEvent], mappings: Sequence[TraceMapping] + engine: Engine, events: Sequence[FunctionTraceEvent], mappings: Sequence[TraceMapping] | None = None ) -> PipelineProjection: raw_parents: dict[int, int | None] = {} projected_ids: set[int] = set() @@ -88,7 +88,7 @@ def pipeline_projection( if event.parent_id is not None and event.parent_id not in raw_parents: raise ValueError(f"trace event {event.id} references unknown or later parent {event.parent_id}") raw_parents[event.id] = event.parent_id - span = _span_for(engine, event.function, mappings) + span = event.function if mappings is None else _span_for(engine, event.function, mappings) if span is None: unmatched += 1 continue @@ -181,12 +181,7 @@ class TraceDiff: @property def matches(self) -> bool: - return ( - not self.python_only - and not self.rust_only - and not self.missing_mappings - and self.shared_order_matches - ) + return not self.python_only and not self.rust_only and not self.missing_mappings and self.shared_order_matches def _missing_mappings( @@ -257,9 +252,7 @@ def trace_diff( rust_counts: Final = Counter(rust_spans) python_only_counts: Final = python_counts - rust_counts rust_only_counts: Final = rust_counts - python_counts - python_only: Final = tuple( - span for span, count in python_only_counts.items() for _ in range(count) - ) + python_only: Final = tuple(span for span, count in python_only_counts.items() for _ in range(count)) rust_only: Final = tuple(span for span, count in rust_only_counts.items() for _ in range(count)) first_difference: Final = _first_difference(python, rust, mappings, contract) return TraceDiff( diff --git a/tests/rust-python-harness/shared/tracing/test_profiler.py b/tests/rust-python-harness/shared/tracing/test_profiler.py index 616e9c23e75..b9a07f249a3 100644 --- a/tests/rust-python-harness/shared/tracing/test_profiler.py +++ b/tests/rust-python-harness/shared/tracing/test_profiler.py @@ -4,9 +4,10 @@ import asyncio import sys import threading from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from functools import wraps from pathlib import Path -from types import FunctionType +from types import FrameType, FunctionType from typing import Final, ParamSpec, TypeVar, cast import pytest @@ -41,11 +42,12 @@ def _events_named(profiler: PythonProfiler, name: str) -> tuple[FunctionTraceEve return tuple(event for event in profiler.events if event.function.endswith(name)) -def test_profiler_keeps_repeated_calls() -> None: +@pytest.mark.parametrize("threads", (False, True)) +def test_profiler_keeps_repeated_calls(threads: bool) -> None: def called() -> None: return None - with profile_python(Path(__file__).parent) as profiler: + with profile_python(Path(__file__).parent, threads=threads) as profiler: called() called() @@ -60,14 +62,15 @@ def test_profiler_qualifies_decorated_methods_by_class() -> None: assert _module_qualnames(__name__)[cast(FunctionType, Decorated.call.__wrapped__).__code__] == "Decorated.call" -def test_profiler_records_real_frame_ancestry() -> None: +@pytest.mark.parametrize("threads", (False, True)) +def test_profiler_records_real_frame_ancestry(threads: bool) -> None: def called() -> None: return None def outer() -> None: called() - with profile_python(Path(__file__).parent) as profiler: + with profile_python(Path(__file__).parent, threads=threads) as profiler: outer() outer_event, called_event = (event for event in profiler.events if event.function.endswith(("outer", "called"))) @@ -84,18 +87,20 @@ def test_profiler_restores_previous_profiler_after_failure() -> None: assert sys.getprofile() is previous -def test_profiler_does_not_count_coroutine_resumption_as_another_call() -> None: +@pytest.mark.parametrize("threads", (False, True)) +def test_profiler_does_not_count_coroutine_resumption_as_another_call(threads: bool) -> None: async def suspended() -> None: await asyncio.sleep(0) await asyncio.sleep(0) - with profile_python(Path(__file__).parent) as profiler: + with profile_python(Path(__file__).parent, threads=threads) as profiler: asyncio.run(suspended()) assert len(_events_named(profiler, "suspended")) == 1 -def test_profiler_preserves_parent_across_coroutine_suspension() -> None: +@pytest.mark.parametrize("threads", (False, True)) +def test_profiler_preserves_parent_across_coroutine_suspension(threads: bool) -> None: def called() -> None: return None @@ -103,7 +108,7 @@ def test_profiler_preserves_parent_across_coroutine_suspension() -> None: await asyncio.sleep(0) called() - with profile_python(Path(__file__).parent) as profiler: + with profile_python(Path(__file__).parent, threads=threads) as profiler: asyncio.run(suspended()) suspended_event: Final = _events_named(profiler, "suspended")[0] @@ -124,6 +129,96 @@ def test_profiler_captures_worker_threads_when_enabled() -> None: assert called_event.parent_id is None +@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring") +@pytest.mark.parametrize("prewarm", (False, True)) +def test_profiler_captures_reused_workers_without_leaking_between_sessions(prewarm: bool) -> None: + def called() -> None: + return None + + with ThreadPoolExecutor(max_workers=1) as executor: + if prewarm: + executor.submit(called).result(timeout=5) + with profile_python(Path(__file__).parent, threads=True) as first: + executor.submit(called).result(timeout=5) + executor.submit(called).result(timeout=5) + with profile_python(Path(__file__).parent, threads=True) as second: + executor.submit(called).result(timeout=5) + executor.submit(called).result(timeout=5) + + assert len(_events_named(first, "called")) == 1 + assert len(_events_named(second, "called")) == 1 + + +def test_profiler_restores_main_and_worker_hooks_after_failure() -> None: + previous: Final = sys.getprofile() + previous_thread: Final = threading.getprofile() + + with ThreadPoolExecutor(max_workers=1) as executor: + worker_previous: Final = executor.submit(sys.getprofile).result(timeout=5) + with pytest.raises(RuntimeError, match="stop"): + with profile_python(Path(__file__).parent, threads=True): + raise RuntimeError("stop") + assert executor.submit(sys.getprofile).result(timeout=5) is worker_previous + + assert sys.getprofile() is previous + assert threading.getprofile() is previous_thread + + +@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring") +def test_function_usage_profiler_captures_reused_workers() -> None: + def selected() -> None: + return None + + function: Final = f"{Path(__file__).name}:{selected.__code__.co_firstlineno} {selected.__qualname__}" + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(selected).result(timeout=5) + with profile_python_function_usage(Path(__file__).parent, frozenset((function,)), threads=True) as profiler: + executor.submit(selected).result(timeout=5) + + assert profiler.called == {function} + + +@pytest.mark.skipif(sys.version_info < (3, 12), reason="independent thread hooks require sys.monitoring") +def test_threaded_profiler_preserves_custom_worker_hook_and_releases_monitoring_slot() -> None: + def worker_hook(_frame: FrameType, _event: str, _arg: object) -> None: + return None + + def fail_with_profile(executor: ThreadPoolExecutor) -> None: + with profile_python(Path(__file__).parent, threads=True): + assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook + raise RuntimeError("stop") + + tools_before: Final = tuple(sys.monitoring.get_tool(slot) for slot in range(6)) + with ThreadPoolExecutor(max_workers=1, initializer=lambda: sys.setprofile(worker_hook)) as executor: + assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook + with pytest.raises(RuntimeError, match="stop"): + fail_with_profile(executor) + assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook + + assert tuple(sys.monitoring.get_tool(slot) for slot in range(6)) == tools_before + + +@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring") +def test_threaded_profiler_keeps_concurrent_event_ids_and_parent_links() -> None: + def child() -> None: + return None + + def parent() -> None: + child() + + with ThreadPoolExecutor(max_workers=4) as executor: + with profile_python(Path(__file__).parent, threads=True) as profiler: + futures: Final = tuple(executor.submit(parent) for _ in range(200)) + for future in futures: + future.result(timeout=5) + + parent_ids: Final = frozenset(event.id for event in _events_named(profiler, "parent")) + children: Final = _events_named(profiler, "child") + assert len(parent_ids) == len(children) == 200 + assert frozenset(event.parent_id for event in children) == parent_ids + assert tuple(event.id for event in profiler.events) == tuple(range(len(profiler.events))) + + def test_function_usage_profiler_records_only_selected_functions() -> None: def selected() -> None: return None diff --git a/tests/rust-python-harness/shared/tracing/test_steps.py b/tests/rust-python-harness/shared/tracing/test_steps.py index 2efc6a3c579..ee5e0bafd28 100644 --- a/tests/rust-python-harness/shared/tracing/test_steps.py +++ b/tests/rust-python-harness/shared/tracing/test_steps.py @@ -40,6 +40,23 @@ def test_python_projection_collapses_unmapped_parents_and_counts_noise() -> None ] +@pytest.mark.parametrize("engine", ("python", "rust")) +def test_projection_without_mappings_keeps_every_call_and_parent(engine: Engine) -> None: + events: Final = ( + event(0, "module.py:1 entry"), + event(1, "module.py:2 internal_helper", 0), + event(2, "module.py:3 nested", 1), + event(3, "module.py:2 internal_helper", 0), + ) + + projection: Final = pipeline_projection(engine, events) + + assert projection.unmatched == 0 + assert tuple((step.id, step.parent_id, step.span, step.raw) for step in projection.steps) == tuple( + (item.id, item.parent_id, item.function, item.raw) for item in events + ) + + def test_rust_projection_keeps_unknown_spans() -> None: projection: Final = pipeline_projection("rust", (event(0, "route"), event(1, "new_span", 0)), MAPPINGS) assert [(step.span, step.parent_id) for step in projection.steps] == [("route", None), ("new_span", 0)] @@ -146,9 +163,7 @@ def test_trace_diff_allows_reordered_concurrent_children() -> None: def test_trace_diff_prunes_declared_engine_only_nodes_but_requires_them() -> None: mappings: Final = (MAPPINGS[0], mapping(rust_span="rust_prepare")) python: Final = pipeline_projection("python", (event(0, "module.py:1 entry"),), mappings).steps - rust: Final = pipeline_projection( - "rust", (event(0, "route"), event(1, "rust_prepare", 0)), mappings - ).steps + rust: Final = pipeline_projection("rust", (event(0, "route"), event(1, "rust_prepare", 0)), mappings).steps assert trace_diff(python, rust, mappings).matches assert trace_diff(python, rust[:1], mappings).missing_mappings == ("rust_prepare",) diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py index bcca8ac6d42..1b283d75ca5 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py @@ -264,9 +264,7 @@ def _formatting_strategy() -> SearchStrategy[ReductoFormatting]: ), st.sampled_from((False, True)).map(lambda value: {"add_page_markers": value}), st.sampled_from((False, True)).map(lambda value: {"merge_tables": value}), - st.sampled_from(REDUCTO_FORMATTING_INCLUDE_GROUPS) - .map(list) - .map(lambda value: {"include": value}), + st.sampled_from(REDUCTO_FORMATTING_INCLUDE_GROUPS).map(list).map(lambda value: {"include": value}), ) return values.map(ReductoFormatting.model_validate) diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py index 5c4abc78081..f14296fbf06 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py @@ -124,10 +124,7 @@ class RecordingCallback(CustomLogger): if isinstance(value, Mapping): if any(not isinstance(map_key, str) for map_key in value): raise TypeError("callback kwarg mappings must use string keys") - return { - map_key: self._normalized_kwargs(map_value, map_key) - for map_key, map_value in value.items() - } + return {map_key: self._normalized_kwargs(map_value, map_key) for map_key, map_value in value.items()} if isinstance(value, (list, tuple)): return [self._normalized_kwargs(item) for item in value] raise TypeError(f"unsupported callback kwarg type: {type(value)}") diff --git a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md index bb7cb8c91d8..19861aea2b4 100644 --- a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md +++ b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md @@ -1 +1 @@ -Maps Python profiler frames onto feature-gated Rust span names via an explicit per-case mapping (Rust span name is the identity) and compares steps, order, and nesting of both live traces against a replayed provider response. +Prints every collected Python call under litellm/ and every feature-gated Rust span from live traces against replayed HTTP responses. The two traces are independent and are not compared. API-key and Vertex credentials scenarios exercise separate authentication paths; credentials scenarios replay the token exchange locally. diff --git a/tests/rust-python-harness/strategies/trace_parity/__init__.py b/tests/rust-python-harness/strategies/trace_parity/__init__.py index ec88b0169fa..710bdaa3d39 100644 --- a/tests/rust-python-harness/strategies/trace_parity/__init__.py +++ b/tests/rust-python-harness/strategies/trace_parity/__init__.py @@ -7,6 +7,7 @@ from ...shared.reporting.strategy import ( ModuleCaseSpec, NotImplementedCaseSpec, RunnerArgumentDefinition, + RunnerOptionDefinition, StrategyDefinition, ) from .reporting import render_trace_results @@ -26,13 +27,17 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( ModuleCaseSpec( coverage=Coverage.PARTIAL, module="tests.rust-python-harness.strategies.trace_parity.sdk.messages.case", - note="Async only until anthropic_messages_handler supports sync calls.", + note="Success paths are async; sync tracing captures the currently unsupported behavior.", ), surface="sdk", ), CaseDefinition( "responses", - NotImplementedCaseSpec(reason="No Responses trace-parity case is registered."), + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.sdk.responses.case", + note="Core create paths: native, streaming, provider error, Azure override, and chat bridge.", + ), surface="sdk", ), CaseDefinition( @@ -70,13 +75,17 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( ModuleCaseSpec( coverage=Coverage.PARTIAL, module="tests.rust-python-harness.strategies.trace_parity.gateway.messages.case", - note="Non-streaming success paths only.", + note="Anthropic/Azure provider routes plus a fully consumed downstream streaming path.", ), surface="gateway", ), CaseDefinition( "responses", - NotImplementedCaseSpec(reason="No gateway Responses trace-parity case is registered."), + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.gateway.responses.case", + note="Native OpenAI non-streaming and fully consumed downstream streaming paths.", + ), surface="gateway", ), CaseDefinition( @@ -86,7 +95,11 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( ), CaseDefinition( "chat_completions", - NotImplementedCaseSpec(reason="No gateway chat trace-parity case is registered."), + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.gateway.chat_completions.case", + note="Anthropic non-streaming and fully consumed downstream streaming paths.", + ), surface="gateway", ), CaseDefinition( @@ -99,8 +112,8 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( STRATEGY: Final = StrategyDefinition( id="trace_parity", order=20, - label="Trace parity", - description="Compare pipeline steps, order, and nesting between Python profiler frames and Rust spans via an explicit mapping.", + label="Traces", + description="Print Python profiler frames and Rust spans for representative pipeline scenarios.", directory=Path(__file__).parent, runnable_spec=ModuleCaseSpec, cases=CASES, @@ -112,4 +125,11 @@ STRATEGY: Final = StrategyDefinition( metavar="NAME", help="run only this named trace scenario; repeat to select more than one", ), + runner_options=( + RunnerOptionDefinition( + option="--engine", + choices=("python", "rust"), + help="show only this engine's trace; omit to print both engines", + ), + ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/fixtures.py b/tests/rust-python-harness/strategies/trace_parity/fixtures.py new file mode 100644 index 00000000000..a3d084c97de --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/fixtures.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +import base64 +import binascii +import json +import struct +from collections.abc import Iterable, Mapping +from typing import Final + +from ...shared.parity.recorded_http import ( + HttpHeader, + RecordedHttpResponse, + RecordedHttpStreamResponse, + RecordedStreamChunk, +) + +JSON_HEADERS: Final = (HttpHeader(name="content-type", value="application/json"),) +SSE_HEADERS: Final = (HttpHeader(name="content-type", value="text/event-stream"),) +AWS_EVENT_STREAM_HEADERS: Final = (HttpHeader(name="content-type", value="application/vnd.amazon.eventstream"),) + + +def json_response(body: Mapping[str, object] | bytes, *, status: int = 200) -> RecordedHttpResponse: + encoded: Final = body if isinstance(body, bytes) else json.dumps(body).encode() + return RecordedHttpResponse.from_bytes(status, JSON_HEADERS, encoded) + + +def sse_event(event: str, payload: Mapping[str, object]) -> bytes: + return f"event: {event}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n".encode() + + +def sse_response(events: Iterable[tuple[str, Mapping[str, object]]]) -> RecordedHttpStreamResponse: + return RecordedHttpStreamResponse( + kind="http_stream", + status_code=200, + headers=SSE_HEADERS, + chunks=tuple(RecordedStreamChunk.from_bytes(sse_event(event, payload)) for event, payload in events), + ) + + +def _aws_string_header(name: str, value: str) -> bytes: + name_bytes: Final = name.encode() + value_bytes: Final = value.encode() + return ( + struct.pack("!B", len(name_bytes)) + + name_bytes + + struct.pack("!B", 7) + + struct.pack("!H", len(value_bytes)) + + value_bytes + ) + + +def aws_event_stream_frame(payload: Mapping[str, object]) -> bytes: + event_payload: Final = json.dumps( + {"bytes": base64.b64encode(json.dumps(payload, separators=(",", ":")).encode()).decode()}, + separators=(",", ":"), + ).encode() + headers: Final = ( + _aws_string_header(":event-type", "chunk") + + _aws_string_header(":content-type", "application/json") + + _aws_string_header(":message-type", "event") + ) + total_length: Final = 12 + len(headers) + len(event_payload) + 4 + prelude: Final = struct.pack("!II", total_length, len(headers)) + prelude_crc: Final = binascii.crc32(prelude) & 0xFFFFFFFF + prelude_crc_bytes: Final = struct.pack("!I", prelude_crc) + message_crc: Final = binascii.crc32(prelude_crc_bytes + headers + event_payload, prelude_crc) & 0xFFFFFFFF + return prelude + prelude_crc_bytes + headers + event_payload + struct.pack("!I", message_crc) + + +def aws_event_stream_response( + events: Iterable[Mapping[str, object]], *, corrupt_last_frame: bool = False +) -> RecordedHttpStreamResponse: + frames: Final = tuple(aws_event_stream_frame(event) for event in events) + body: Final = ( + b"".join((*frames[:-1], frames[-1][:-1] + bytes((frames[-1][-1] ^ 0xFF,)))) + if corrupt_last_frame + else b"".join(frames) + ) + return RecordedHttpStreamResponse( + kind="http_stream", + status_code=200, + headers=AWS_EVENT_STREAM_HEADERS, + chunks=(RecordedStreamChunk.from_bytes(body),), + ) + + +def anthropic_response_body(*, model: str = "claude-sonnet-5") -> dict[str, object]: + return { + "id": "msg_trace", + "type": "message", + "role": "assistant", + "model": model, + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 3}, + } + + +def anthropic_stream_events(*, model: str = "claude-sonnet-5") -> tuple[tuple[str, Mapping[str, object]], ...]: + return ( + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_trace", + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 0}, + }, + }, + ), + ( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello"}}, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 1}, + }, + ), + ("message_stop", {"type": "message_stop"}), + ) + + +def responses_body(*, model: str = "gpt-5", status: str = "completed") -> dict[str, object]: + return { + "id": "resp_trace", + "object": "response", + "created_at": 1_750_000_000, + "status": status, + "model": model, + "output": [ + { + "type": "message", + "id": "msg_trace", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hello", "annotations": []}], + } + ], + "usage": {"input_tokens": 2, "output_tokens": 3, "total_tokens": 5}, + } + + +def responses_stream_events(*, model: str = "gpt-5") -> tuple[tuple[str, Mapping[str, object]], ...]: + response: Final = responses_body(model=model) + return ( + ( + "response.created", + {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, + ), + ( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": "msg_trace", + "output_index": 0, + "content_index": 0, + "delta": "hello", + }, + ), + ("response.completed", {"type": "response.completed", "response": response}), + ) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py new file mode 100644 index 00000000000..3dc6d731b4b --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from typing import Final + +from .....shared.tracing.steps import Engine, mapping +from ...fixtures import anthropic_response_body, anthropic_stream_events, json_response, sse_response +from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite + +MAPPINGS: Final = ( + mapping(span="python_chat_gateway_route", python_frame=r"proxy_server\.py:\d+ chat_completion$"), + mapping(span="python_gateway_service", python_frame=r"ProxyBaseLLMRequestProcessing\.base_process_llm_request$"), + mapping(span="python_chat_entrypoint", python_frame=r"main\.py:\d+ a?completion$"), + mapping(span="python_provider_config", python_frame=r"ProviderConfigManager\.get_provider_chat_config$"), + mapping(rust_span="validate_environment", python_frame=r"(? RouteFixture: + return RouteFixture( + kwargs={ + "model_alias": "trace-model", + "provider_model": "anthropic/claude-sonnet-5", + "body": { + "model": "trace-model", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 16, + }, + }, + provider_responses=(json_response(anthropic_response_body()),), + ) + + +def _stream_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _fixture(engine, base_url) + return fixture.with_body(stream=True).derive( + provider_responses=(sse_response(anthropic_stream_events()),), + ) + + +TRACE_SUITE: Final = TraceSuite( + route=GatewayRouteSpec("chat_completions", rust_supported=False), + scenarios=( + TraceScenario(name="async-anthropic", fixture=_fixture, mappings=MAPPINGS, asynchronous=True), + TraceScenario( + name="async-anthropic-downstream-stream", + fixture=_stream_fixture, + mappings=(*MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + ), +) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py index 860e872dd44..94d6be7cebf 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py @@ -13,8 +13,8 @@ from ....shared.parity.replay import replay_server from ....shared.tracing.native import TraceResponsePayload, native_trace_events from ....shared.tracing.profiler import FunctionTraceEvent, profile_python from ....shared.tracing.steps import Engine, PipelineProjection, pipeline_projection -from ..models import GatewayRouteSpec, RouteFixture, TraceExecutionFailure, TraceMode, TraceScenario -from ..reporting import TraceComparisonArtifact +from ..models import GatewayRouteSpec, RouteFixture, TraceEngine, TraceExecutionFailure, TraceScenario +from ..reporting import TraceArtifact class _GatewayResponsePayload(BaseModel): @@ -28,7 +28,14 @@ class _GatewayClient(Protocol): def post(self, url: str, *, json: object, headers: dict[str, str]) -> httpx.Response: ... -def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: +_ROUTE_PATHS: Final = { + "messages": "/v1/messages", + "chat_completions": "/v1/chat/completions", + "responses": "/v1/responses", +} + + +def _collect_python(fixture: RouteFixture, route: GatewayRouteSpec) -> tuple[FunctionTraceEvent, ...]: from fastapi.testclient import TestClient import litellm @@ -61,7 +68,7 @@ def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: client: Final = cast(_GatewayClient, TestClient(proxy_server.app)) response: Final = client.post( - "/v1/messages", + _ROUTE_PATHS[route.route], json=fixture.kwargs["body"], headers={"authorization": "Bearer trace-key"}, ) @@ -76,9 +83,10 @@ def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: proxy_server.app.dependency_overrides[user_api_key_auth] = old_override -def _collect_rust(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: +def _collect_rust(fixture: RouteFixture, route: GatewayRouteSpec) -> tuple[FunctionTraceEvent, ...]: payload: Final = json.dumps( { + "path": _ROUTE_PATHS[route.route], "model_alias": fixture.kwargs["model_alias"], "provider_model": fixture.kwargs["provider_model"], "api_base": fixture.kwargs["api_base"], @@ -130,7 +138,9 @@ def _gateway_trace_binary() -> Path: return rust_root / "target" / "debug" / "trace-parity-gateway" -def _collect(scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: +def _collect( + route: GatewayRouteSpec, scenario: TraceScenario, engine: Engine +) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: try: with replay_server() as provider: base_fixture: Final = scenario.fixture(engine, provider.url) @@ -140,7 +150,7 @@ def _collect(scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEven ) for response in fixture.provider_responses: provider.enqueue_response(response) - events: Final = _collect_python(fixture) if engine == "python" else _collect_rust(fixture) + events: Final = _collect_python(fixture, route) if engine == "python" else _collect_rust(fixture, route) provider.take_requests(len(fixture.provider_responses)) return events except Exception as error: @@ -150,40 +160,38 @@ def _collect(scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEven def _projections( python_events: tuple[FunctionTraceEvent, ...], rust_events: tuple[FunctionTraceEvent, ...], - scenario: TraceScenario, - mode: TraceMode, ) -> tuple[PipelineProjection, PipelineProjection, str | None]: - mappings: Final = scenario.mappings_for(mode) try: return ( - pipeline_projection("python", python_events, mappings), - pipeline_projection("rust", rust_events, mappings), + pipeline_projection("python", python_events), + pipeline_projection("rust", rust_events), None, ) except ValueError as error: return PipelineProjection(), PipelineProjection(), f"harness: {error}" -def execute_gateway_trace(route: GatewayRouteSpec, scenario: TraceScenario, mode: TraceMode) -> TraceComparisonArtifact: - mappings: Final = scenario.mappings_for(mode) - python_trace: Final = _collect(scenario, "python") - rust_trace: Final = _collect(scenario, "rust") +def execute_gateway_trace( + route: GatewayRouteSpec, + scenario: TraceScenario, + engine: TraceEngine = "both", +) -> TraceArtifact: + effective_engine: Final[TraceEngine] = "python" if engine == "both" and not route.rust_supported else engine + python_trace: Final = _collect(route, scenario, "python") if effective_engine != "rust" else () + rust_trace: Final = _collect(route, scenario, "rust") if effective_engine != "python" else () collection_python_error: Final = None if isinstance(python_trace, tuple) else f"python: {python_trace.message}" rust_error: Final = None if isinstance(rust_trace, tuple) else f"rust: {rust_trace.message}" python_events: Final = python_trace if isinstance(python_trace, tuple) else () rust_events: Final = rust_trace if isinstance(rust_trace, tuple) else () - python, rust, projection_error = _projections(python_events, rust_events, scenario, mode) + python, rust, projection_error = _projections(python_events, rust_events) python_error: Final = projection_error or collection_python_error - return TraceComparisonArtifact.from_traces( + return TraceArtifact.from_traces( + engine=effective_engine, surface="gateway", sdk_function=route.route, scenario=scenario.name, - mode=mode, - mappings=mappings, - contract=scenario.contract, python=python.steps, rust=rust.steps, - python_unmatched=python.unmatched, python_error=python_error, rust_error=rust_error, ) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py index 30f51cee353..ca9c858f6b7 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py @@ -1,10 +1,9 @@ from __future__ import annotations -import json from typing import Final -from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse from .....shared.tracing.steps import Engine, mapping +from ...fixtures import anthropic_response_body, anthropic_stream_events, json_response, sse_response from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite @@ -49,24 +48,7 @@ def _fixture(_engine: Engine, provider: str) -> RouteFixture: "max_tokens": 16, }, }, - provider_responses=( - RecordedHttpResponse.from_bytes( - 200, - (HttpHeader(name="content-type", value="application/json"),), - json.dumps( - { - "id": "msg_trace", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-5", - "content": [{"type": "text", "text": "hello"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 3}, - } - ).encode(), - ), - ), + provider_responses=(json_response(anthropic_response_body()),), ) @@ -78,6 +60,13 @@ def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: return _fixture(engine, "azure_ai") +def _stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _anthropic_fixture(engine, _base_url) + return fixture.with_body(stream=True).derive( + provider_responses=(sse_response(anthropic_stream_events()),), + ) + + ANTHROPIC_MAPPINGS: Final = ( *GATEWAY_MAPPINGS, mapping( @@ -100,7 +89,22 @@ AZURE_MAPPINGS: Final = ( TRACE_SUITE: Final = TraceSuite( route=GatewayRouteSpec("messages"), scenarios=( - TraceScenario(name="anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, modes=("async",)), - TraceScenario(name="azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, modes=("async",)), + TraceScenario( + name="async-anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, asynchronous=True + ), + TraceScenario(name="async-azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, asynchronous=True), + TraceScenario( + name="async-anthropic-downstream-stream", + fixture=_stream_fixture, + mappings=( + *ANTHROPIC_MAPPINGS, + mapping(span="python_upstream_stream", python_frame=r"AnthropicMessagesStreamingResponse\.__anext__$"), + mapping( + span="python_downstream_stream", python_frame=r"DataGenerator\.__anext__$|async_data_generator$" + ), + mapping(span="python_stream_callback", python_frame=r"Logging\.async_success_handler$"), + ), + asynchronous=True, + ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py new file mode 100644 index 00000000000..7d805a0603c --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from typing import Final + +from .....shared.tracing.steps import Engine, mapping +from ...fixtures import json_response, responses_body, responses_stream_events, sse_response +from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite + +MAPPINGS: Final = ( + mapping( + span="python_responses_gateway_route", python_frame=r"response_api_endpoints/endpoints\.py:\d+ responses_api$" + ), + mapping(span="python_gateway_service", python_frame=r"ProxyBaseLLMRequestProcessing\.base_process_llm_request$"), + mapping(span="python_responses", python_frame=r"responses/main\.py:\d+ a?responses$"), + mapping(span="python_provider_config", python_frame=r"ProviderConfigManager\.get_provider_responses_api_config$"), + mapping(rust_span="validate_environment", python_frame=r"OpenAIResponsesAPIConfig\.validate_environment$"), + mapping(rust_span="complete_url", python_frame=r"OpenAIResponsesAPIConfig\.get_complete_url$"), + mapping(rust_span="transform_request", python_frame=r"OpenAIResponsesAPIConfig\.transform_responses_api_request$"), + mapping(span="python_logging_pre_call", python_frame=r"Logging\.pre_call$"), + mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), + mapping(rust_span="transform_response", python_frame=r"OpenAIResponsesAPIConfig\.transform_response_api_response$"), + mapping(span="python_success_callback", python_frame=r"Logging\.async_success_handler$|Logging\.success_handler$"), +) + +STREAM_MAPPINGS: Final = ( + mapping(span="python_stream_iterator", python_frame=r"ResponsesAPIStreamingIterator\.__init__$"), + mapping(span="python_stream_next", python_frame=r"ResponsesAPIStreamingIterator\.__anext__$"), + mapping(span="python_stream_transform", python_frame=r"OpenAIResponsesAPIConfig\.transform_streaming_response$"), + mapping(span="python_downstream_stream", python_frame=r"DataGenerator\.__anext__$|async_data_generator$"), +) + + +def _fixture(_engine: Engine, _base_url: str) -> RouteFixture: + return RouteFixture( + kwargs={ + "model_alias": "trace-model", + "provider_model": "openai/gpt-5", + "body": {"model": "trace-model", "input": "hello"}, + }, + provider_responses=(json_response(responses_body()),), + ) + + +def _stream_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _fixture(engine, base_url) + return fixture.with_body(stream=True).derive( + provider_responses=(sse_response(responses_stream_events()),), + ) + + +TRACE_SUITE: Final = TraceSuite( + route=GatewayRouteSpec("responses", rust_supported=False), + scenarios=( + TraceScenario(name="async-openai", fixture=_fixture, mappings=MAPPINGS, asynchronous=True), + TraceScenario( + name="async-openai-downstream-stream", + fixture=_stream_fixture, + mappings=(*MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + ), +) diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py index 04659b25382..d6ed42250c4 100644 --- a/tests/rust-python-harness/strategies/trace_parity/models.py +++ b/tests/rust-python-harness/strategies/trace_parity/models.py @@ -1,35 +1,61 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass -from typing import Final, Literal, TypeAlias +from typing import Final, Literal, TypeAlias, cast -from ...shared.parity.recorded_http import RecordedHttpResponse +from ...shared.parity.recorded_http import RecordedResponse from ...shared.reporting.models import SdkFunction -from ...shared.tracing.steps import Engine, TraceContract, TraceMapping +from ...shared.tracing.steps import Engine, TraceMapping -TraceMode = Literal["sync", "async"] +TraceEngine = Literal["python", "rust", "both"] TraceFailureSource = Literal["python", "rust", "harness"] @dataclass(frozen=True, slots=True) class RouteFixture: kwargs: dict[str, object] - provider_responses: tuple[RecordedHttpResponse, ...] + provider_responses: tuple[RecordedResponse, ...] expected_failure: bool = False + consume_stream: bool = False + environment: tuple[tuple[str, str], ...] = () + + def derive( + self, + *, + kwargs: Mapping[str, object] | None = None, + provider_responses: tuple[RecordedResponse, ...] | None = None, + expected_failure: bool | None = None, + consume_stream: bool | None = None, + ) -> RouteFixture: + return RouteFixture( + kwargs={**self.kwargs, **(kwargs or {})}, + provider_responses=self.provider_responses if provider_responses is None else provider_responses, + expected_failure=self.expected_failure if expected_failure is None else expected_failure, + consume_stream=self.consume_stream if consume_stream is None else consume_stream, + environment=self.environment, + ) + + def with_body(self, **updates: object) -> RouteFixture: + raw_body: Final = self.kwargs.get("body") + if not isinstance(raw_body, dict): + raise ValueError("route fixture does not contain an object body") + body: Final = cast(dict[str, object], raw_body) + return self.derive(kwargs={"body": {**body, **updates}}) @dataclass(frozen=True, slots=True) class RouteSpec: route: SdkFunction python_entrypoints: tuple[str, str] - rust_entrypoints: tuple[str, str] + rust_entrypoints: tuple[str, str] | None fixture: Callable[[Engine, str], RouteFixture] @dataclass(frozen=True, slots=True) class GatewayRouteSpec: route: SdkFunction + rust_supported: bool = True TraceRouteSpec: TypeAlias = RouteSpec | GatewayRouteSpec @@ -40,14 +66,7 @@ class TraceScenario: name: str fixture: Callable[[Engine, str], RouteFixture] mappings: tuple[TraceMapping, ...] - modes: tuple[TraceMode, ...] = ("sync", "async") - contract: TraceContract = TraceContract() - sync_mappings: tuple[TraceMapping, ...] | None = None - async_mappings: tuple[TraceMapping, ...] | None = None - - def mappings_for(self, mode: TraceMode) -> tuple[TraceMapping, ...]: - selected: Final = self.async_mappings if mode == "async" else self.sync_mappings - return self.mappings if selected is None else selected + asynchronous: bool @dataclass(frozen=True, slots=True) diff --git a/tests/rust-python-harness/strategies/trace_parity/reporting.py b/tests/rust-python-harness/strategies/trace_parity/reporting.py index 9c5bf9e88cd..e7c07ef9c0f 100644 --- a/tests/rust-python-harness/strategies/trace_parity/reporting.py +++ b/tests/rust-python-harness/strategies/trace_parity/reporting.py @@ -1,31 +1,24 @@ from __future__ import annotations import os -import re import sys from collections.abc import Sequence -from typing import Final, Literal +from typing import Final from pydantic import BaseModel, ConfigDict, ValidationError from ...shared.reporting.models import SURFACES, CaseResult, RunStatus, SdkFunction, Surface from ...shared.reporting.rendering import ReportSection from ...shared.reporting.strategy import NotImplementedCaseSpec, SkippedCaseSpec -from ...shared.tracing.steps import ( - PipelineStep, - TraceContract, - TraceDiff, - TraceMapping, - trace_depths, - trace_diff, -) +from ...shared.tracing.steps import PipelineStep, trace_depths +from .models import TraceEngine -TRACE_COMPARISON_ARTIFACT: Final = "trace_comparison" +TRACE_ARTIFACT: Final = "trace" TRACE_PARITY_HINT: Final = ( "rebuild the native bridge with the trace-parity feature, e.g. `uvx maturin develop --features trace-parity`" ) -_COLORS: Final[dict[str, str]] = {"green": "32", "yellow": "33", "red": "31", "cyan": "36"} +_COLORS: Final[dict[str, str]] = {"yellow": "33", "red": "31", "cyan": "36"} _RESET: Final = "\033[0m" @@ -47,26 +40,15 @@ class TraceEventArtifact(BaseModel): return PipelineStep(self.id, self.parent_id, self.span, self.raw) -class TraceMappingArtifact(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - - span: str - python: str | None - rust: str | None - - -class TraceComparisonArtifact(BaseModel): +class TraceArtifact(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") + engine: TraceEngine = "both" surface: Surface sdk_function: SdkFunction scenario: str - mode: Literal["sync", "async"] - mappings: tuple[TraceMappingArtifact, ...] python: tuple[TraceEventArtifact, ...] rust: tuple[TraceEventArtifact, ...] - python_unmatched: int - unordered_children_of: frozenset[str] python_error: str | None = None rust_error: str | None = None @@ -74,41 +56,27 @@ class TraceComparisonArtifact(BaseModel): def from_traces( cls, *, + engine: TraceEngine = "both", surface: Surface, sdk_function: SdkFunction, scenario: str, - mode: Literal["sync", "async"], - mappings: Sequence[TraceMapping], - contract: TraceContract, python: Sequence[PipelineStep], rust: Sequence[PipelineStep], - python_unmatched: int, python_error: str | None = None, rust_error: str | None = None, - ) -> TraceComparisonArtifact: + ) -> TraceArtifact: return cls( + engine=engine, surface=surface, sdk_function=sdk_function, scenario=scenario, - mode=mode, - mappings=tuple( - TraceMappingArtifact( - span=item.span, - python=item.python.pattern if item.python else None, - rust=item.rust, - ) - for item in mappings - ), python=tuple( TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw) for step in python ), rust=tuple( - TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw) - for step in rust + TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw) for step in rust ), - python_unmatched=python_unmatched, - unordered_children_of=contract.unordered_children_of, python_error=python_error, rust_error=rust_error, ) @@ -119,32 +87,9 @@ class TraceComparisonArtifact(BaseModel): def rust_steps(self) -> tuple[PipelineStep, ...]: return tuple(event.step() for event in self.rust) - def diff(self) -> TraceDiff: - return trace_diff( - self.python_steps(), - self.rust_steps(), - tuple( - TraceMapping( - item.span, - re.compile(item.python) if item.python is not None else None, - item.rust, - ) - for item in self.mappings - ), - TraceContract(self.unordered_children_of), - ) - - def exact_match(self) -> bool: - return self.diff().matches - def has_errors(self) -> bool: return self.python_error is not None or self.rust_error is not None - def contract_matches(self) -> bool: - if self.has_errors(): - return False - return self.diff().matches - def _split_raw(raw: str) -> tuple[str, str]: location, separator, name = raw.partition(" ") @@ -153,72 +98,28 @@ def _split_raw(raw: str) -> tuple[str, str]: return raw, "" -def _python_line(index: int, step: PipelineStep, depth: int, exclusive: frozenset[str]) -> str: +def _python_line(index: int, step: PipelineStep, depth: int) -> str: name: Final = _split_raw(step.raw)[0] location: Final = _split_raw(step.raw)[1] suffix: Final = f" ({location})" if location else "" - marker: Final = " [python only]" if step.span in exclusive else "" - return _paint(f"{index} {' ' * depth}{name}{suffix}{marker}", "cyan") + return _paint(f"{index} {' ' * depth}{name}{suffix}", "cyan") -def _python_lines(steps: tuple[PipelineStep, ...], exclusive: frozenset[str]) -> str: +def _python_lines(steps: tuple[PipelineStep, ...]) -> str: depths: Final = trace_depths(steps) - lines: Final = tuple( - _python_line(index, step, depths[step.id], exclusive) for index, step in enumerate(steps, start=1) - ) + lines: Final = tuple(_python_line(index, step, depths[step.id]) for index, step in enumerate(steps, start=1)) return f"{_paint('PYTHON', 'cyan')} ({len(steps)} steps)\n" + ("\n".join(lines) if lines else "(empty)") -def _python_references(steps: tuple[PipelineStep, ...]) -> dict[tuple[str, int], str]: - references: dict[tuple[str, int], str] = {} - occurrences: dict[str, int] = {} - for index, step in enumerate(steps, start=1): - name = _split_raw(step.raw)[0] - occurrence = occurrences.get(step.span, 0) + 1 - occurrences[step.span] = occurrence - references[(step.span, occurrence)] = f"{index} {name}" - return references - - -def _rust_line( - step: PipelineStep, - depth: int, - occurrence: int, - references: dict[tuple[str, int], str], -) -> str: - span: Final = _paint(step.span, "yellow") - key: Final = (step.span, occurrence) - reference: Final = ( - _paint(references[key], "cyan") if key in references else _paint("[rust only]", "yellow") - ) - suffix: Final = f"#{occurrence}" if occurrence > 1 else "" - return f"{' ' * depth}{span}{suffix} -> {reference}" - - -def _rust_lines(steps: tuple[PipelineStep, ...], references: dict[tuple[str, int], str]) -> str: +def _rust_lines(steps: tuple[PipelineStep, ...]) -> str: depths: Final = trace_depths(steps) - occurrences: dict[str, int] = {} - lines: list[str] = [] - for step in steps: - occurrence = occurrences.get(step.span, 0) + 1 - occurrences[step.span] = occurrence - lines.append(_rust_line(step, depths[step.id], occurrence, references)) + lines: Final = tuple( + _paint(f"{index} {' ' * depths[step.id]}{step.span}", "yellow") for index, step in enumerate(steps, 1) + ) return f"{_paint('RUST', 'yellow')} ({len(steps)} steps)\n" + ("\n".join(lines) if lines else "(empty)") -def _state_text(state: str, *, good: bool) -> str: - return _paint(state, "green" if good else "red") - - -def _contract_line(artifact: TraceComparisonArtifact) -> str: - matches: Final = artifact.contract_matches() - status: Final = _state_text("PASS" if matches else "FAIL", good=matches) - if artifact.python_error or artifact.rust_error: - return f"Contract: {status}" - return f"Contract: {status}" - - -def _error_lines(artifact: TraceComparisonArtifact) -> tuple[str, ...]: +def _error_lines(artifact: TraceArtifact) -> tuple[str, ...]: lines: list[str] = [] for engine, error in (("Python", artifact.python_error), ("Rust", artifact.rust_error)): if error is None: @@ -229,68 +130,20 @@ def _error_lines(artifact: TraceComparisonArtifact) -> tuple[str, ...]: return tuple(lines) -def _unseen_mappings( - artifact: TraceComparisonArtifact, - python: tuple[PipelineStep, ...], - rust: tuple[PipelineStep, ...], -) -> tuple[str, ...]: - return artifact.diff().missing_mappings - - -def _comparison_status_lines( - artifact: TraceComparisonArtifact, - python: tuple[PipelineStep, ...], - rust: tuple[PipelineStep, ...], -) -> tuple[str, ...]: - diff: Final = artifact.diff() - exact_match: Final = artifact.exact_match() - if artifact.has_errors(): - return (*_error_lines(artifact), _contract_line(artifact)) - unseen: Final = _unseen_mappings(artifact, python, rust) - unseen_line: Final[tuple[str, ...]] = (f"Unseen mappings: {', '.join(unseen)}",) if unseen else () - drift_lines: Final[tuple[str, ...]] = ( - (_state_text("Same steps, order, and nesting", good=True),) - if exact_match - else ( - _paint(f"Python only: {', '.join(diff.python_only) or 'none'}", "cyan"), - _paint(f"Rust only: {', '.join(diff.rust_only) or 'none'}", "yellow"), - f"First difference: {diff.first_difference or 'none'}", - f"Python frames outside mapping: {artifact.python_unmatched}", - ) - ) - return ( - f"Trace: {_state_text('MATCH' if exact_match else 'DRIFT', good=exact_match)}", - *drift_lines, - *unseen_line, - _contract_line(artifact), - ) - - -def _render_comparison(artifact: TraceComparisonArtifact) -> str: - python: Final = artifact.python_steps() - rust: Final = artifact.rust_steps() - diff: Final = artifact.diff() - python_exclusive: Final = frozenset(item.span for item in artifact.mappings if item.rust is None) - status_lines: Final = _comparison_status_lines(artifact, python, rust) - return "\n\n".join( - ( - _python_lines(python, python_exclusive | frozenset(diff.python_only)), - _rust_lines(rust, _python_references(python)), - "\n".join(status_lines), - ) - ) - - -def _mode(nodeid: str) -> str: - if "[" in nodeid: - return nodeid.rsplit("[", 1)[-1].removesuffix("]") - head, _, tail = nodeid.rpartition(":") - return tail if head else "unknown mode" +def _render_trace(artifact: TraceArtifact) -> str: + traces: tuple[str, ...] + if artifact.engine == "python": + traces = (_python_lines(artifact.python_steps()),) + elif artifact.engine == "rust": + traces = (_rust_lines(artifact.rust_steps()),) + else: + traces = (_python_lines(artifact.python_steps()), _rust_lines(artifact.rust_steps())) + return "\n\n".join((*traces, *_error_lines(artifact))) def _scenario(nodeid: str) -> str: parts: Final = nodeid.split(":") - return parts[-2] if len(parts) >= 5 else "default" + return parts[-1] if len(parts) >= 4 else "default" def _unavailable(status: RunStatus) -> str: @@ -299,20 +152,20 @@ def _unavailable(status: RunStatus) -> str: def _render_artifact(body: str) -> str: try: - artifact: Final = TraceComparisonArtifact.model_validate_json(body) + artifact: Final = TraceArtifact.model_validate_json(body) except ValidationError as error: - return f"Trace comparison artifact is invalid: {error}" - return _render_comparison(artifact) + return f"Trace artifact is invalid: {error}" + return _render_trace(artifact) -def _mode_section(result: CaseResult, nodeid: str, status: RunStatus) -> str: +def _scenario_section(result: CaseResult, nodeid: str, status: RunStatus) -> str: artifacts: Final = tuple( - artifact for artifact in result.artifacts.get(nodeid, ()) if artifact.kind == TRACE_COMPARISON_ARTIFACT + artifact for artifact in result.artifacts.get(nodeid, ()) if artifact.kind == TRACE_ARTIFACT ) body: Final = ( "\n\n".join(_render_artifact(artifact.body) for artifact in artifacts) if artifacts else _unavailable(status) ) - label: Final = f"Scenario: {_scenario(nodeid)} / Mode: {_mode(nodeid)}" + label: Final = f"Scenario: {_scenario(nodeid)}" return f"{label}\n{'-' * len(label)}\n\n{body}" @@ -321,7 +174,7 @@ def _case_block(result: CaseResult) -> str: outcomes: Final = tuple(result.outcomes.items()) or ( (nodeid, RunStatus.NOT_RUN) for nodeid in sorted(result.collected) ) - sections: Final = tuple(_mode_section(result, nodeid, status) for nodeid, status in outcomes) + sections: Final = tuple(_scenario_section(result, nodeid, status) for nodeid, status in outcomes) return "\n\n".join((f"{header}\n{'=' * len(header)}", *sections)) @@ -357,11 +210,11 @@ def _surface_section(surface: Surface, results: Sequence[CaseResult]) -> ReportS *((not_implemented,) if not_implemented else ()), *((skipped,) if skipped else ()), ) - return ReportSection(f"{surface.upper()} trace comparisons", blocks or ("No runnable trace comparisons",)) + return ReportSection(f"{surface.upper()} traces", blocks or ("No runnable traces",)) def render_trace_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: sections: Final = tuple( section for surface in SURFACES if (section := _surface_section(surface, results)) is not None ) - return sections or (ReportSection("Trace comparisons", ("No trace comparisons selected",)),) + return sections or (ReportSection("Traces", ("No traces selected",)),) diff --git a/tests/rust-python-harness/strategies/trace_parity/runner.py b/tests/rust-python-harness/strategies/trace_parity/runner.py index b78a3c7da3f..706e054bf52 100644 --- a/tests/rust-python-harness/strategies/trace_parity/runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/runner.py @@ -4,13 +4,20 @@ import importlib from collections.abc import Sequence from pathlib import Path from time import monotonic -from typing import Final +from typing import Final, cast +from ...shared.native_build import ensure_trace_bridge from ...shared.reporting.models import CaseResult, HarnessCase, HarnessRun, ResultArtifact, RunStatus, Surface from ...shared.reporting.strategy import ModuleCaseSpec, UpdateCallback -from ...shared.native_build import ensure_trace_bridge -from .models import GatewayRouteSpec, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario, TraceSuite -from .reporting import TRACE_COMPARISON_ARTIFACT, TraceComparisonArtifact +from .models import ( + GatewayRouteSpec, + RouteSpec, + TraceEngine, + TraceExecutionFailure, + TraceScenario, + TraceSuite, +) +from .reporting import TRACE_ARTIFACT, TraceArtifact from .sdk.execution import execute_trace @@ -32,15 +39,13 @@ def validate_trace_suite(suite: TraceSuite, harness_case: HarnessCase) -> str | names: Final = tuple(scenario.name for scenario in suite.scenarios) if not names or len(names) != len(set(names)) or any(not name or ":" in name for name in names): return "scenario names must be non-empty, unique, and colon-free" - invalid_modes: Final = tuple( + invalid_names: Final = tuple( scenario.name for scenario in suite.scenarios - if not scenario.modes - or len(scenario.modes) != len(set(scenario.modes)) - or any(mode not in {"sync", "async"} for mode in scenario.modes) + if not scenario.name.startswith("async-" if scenario.asynchronous else "sync-") ) - if invalid_modes: - return f"scenarios must use non-empty, unique sync/async modes: {', '.join(invalid_modes)}" + if invalid_names: + return f"scenario names must start with sync- or async-: {', '.join(invalid_names)}" surface: Final = harness_case.surface if surface == "sdk" and not isinstance(suite.route, RouteSpec): return "must use RouteSpec for the sdk surface" @@ -57,15 +62,14 @@ def scenario_nodeids( trace_suite: TraceSuite, harness_case: HarnessCase, selected_scenarios: frozenset[str] = frozenset(), -) -> tuple[tuple[TraceScenario, TraceMode, str], ...]: +) -> tuple[tuple[TraceScenario, str], ...]: surface: Final = harness_case.surface if surface is None: return () return tuple( - (scenario, mode, f"trace:{surface}:{harness_case.sdk_function}:{scenario.name}:{mode}") + (scenario, f"trace:{surface}:{harness_case.sdk_function}:{scenario.name}") for scenario in trace_suite.scenarios if not selected_scenarios or scenario.name in selected_scenarios - for mode in scenario.modes ) @@ -77,54 +81,49 @@ def _record_setup_failure(run: HarnessRun, case: HarnessCase, message: str, stag run.failures.append((nodeid, message)) -def run_trace_mode( +def run_trace_scenario( run: HarnessRun, result: CaseResult, trace_suite: TraceSuite, scenario: TraceScenario, - mode: TraceMode, surface: Surface, nodeid: str, on_update: UpdateCallback, + engine: TraceEngine = "both", ) -> None: started_at: Final = monotonic() - comparison: Final = _execute_mode(trace_suite, scenario, mode, surface) + trace: Final = _execute_scenario(trace_suite, scenario, surface, engine) duration: Final = monotonic() - started_at - if isinstance(comparison, TraceExecutionFailure): + if isinstance(trace, TraceExecutionFailure): result.record(nodeid, RunStatus.ERROR, duration) - run.failures.append((nodeid, comparison.message)) + run.failures.append((nodeid, trace.message)) on_update(run) return - artifact: Final = ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()) - if comparison.has_errors(): + artifact: Final = ResultArtifact(TRACE_ARTIFACT, trace.model_dump_json()) + if trace.has_errors(): result.record(nodeid, RunStatus.ERROR, duration, (artifact,)) - run.failures.append( - (nodeid, "\n".join(error for error in (comparison.python_error, comparison.rust_error) if error)) - ) + run.failures.append((nodeid, "\n".join(error for error in (trace.python_error, trace.rust_error) if error))) else: - status: Final = RunStatus.PASSED if comparison.contract_matches() else RunStatus.FAILED - result.record(nodeid, status, duration, (artifact,)) - if status is RunStatus.FAILED: - run.failures.append((nodeid, "trace contract mismatch; see the rendered comparison")) + result.record(nodeid, RunStatus.PASSED, duration, (artifact,)) on_update(run) -def _execute_mode( +def _execute_scenario( trace_suite: TraceSuite, scenario: TraceScenario, - mode: TraceMode, surface: Surface, -) -> TraceComparisonArtifact | TraceExecutionFailure: + engine: TraceEngine, +) -> TraceArtifact | TraceExecutionFailure: route: Final = trace_suite.route if isinstance(route, GatewayRouteSpec): if surface != "gateway": return TraceExecutionFailure("harness", "gateway route cannot run on the sdk surface") from .gateway.execution import execute_gateway_trace - return execute_gateway_trace(route, scenario, mode) + return execute_gateway_trace(route, scenario, engine) if surface != "sdk": return TraceExecutionFailure("harness", "sdk route cannot run on the gateway surface") - return execute_trace(route, scenario, mode, surface) + return execute_trace(route, scenario, surface, engine) def _run_case( @@ -132,6 +131,7 @@ def _run_case( harness_case: HarnessCase, selected_scenarios: frozenset[str], on_update: UpdateCallback, + engine: TraceEngine, ) -> None: result: Final = run.results[harness_case.key] spec: Final = harness_case.spec @@ -146,15 +146,29 @@ def _run_case( on_update(run) return nodeids: Final = scenario_nodeids(trace_suite, harness_case, selected_scenarios) - result.collected.update(nodeid for _, _, nodeid in nodeids) + result.collected.update(nodeid for _, nodeid in nodeids) if not nodeids: result.status = RunStatus.SKIPPED on_update(run) return result.status = RunStatus.RUNNING on_update(run) - for scenario, mode, nodeid in nodeids: - run_trace_mode(run, result, trace_suite, scenario, mode, surface, nodeid, on_update) + for scenario, nodeid in nodeids: + run_trace_scenario(run, result, trace_suite, scenario, surface, nodeid, on_update, engine) + + +def runner_selection(runner_args: Sequence[str]) -> tuple[frozenset[str], TraceEngine]: + engine: TraceEngine = "both" + scenarios: list[str] = [] + for argument in runner_args: + if argument.startswith("--engine="): + value = argument.removeprefix("--engine=") + if value not in {"python", "rust"}: + raise ValueError(f"invalid trace engine: {value}") + engine = cast(TraceEngine, value) + else: + scenarios.append(argument) + return frozenset(scenarios), engine def run_trace_cases( @@ -163,10 +177,10 @@ def run_trace_cases( on_update: UpdateCallback, runner_args: Sequence[str] = (), ) -> tuple[int, HarnessRun]: - selected_scenarios: Final = frozenset(runner_args) + selected_scenarios, engine = runner_selection(runner_args) run: Final = HarnessRun.from_cases(cases) runnable_cases: Final = tuple(case for case in cases if isinstance(case.spec, ModuleCaseSpec)) - bridge_error: Final = ensure_trace_bridge(repo_root) if runnable_cases else None + bridge_error: Final = ensure_trace_bridge(repo_root) if runnable_cases and engine != "python" else None if bridge_error is not None: for harness_case in runnable_cases: _record_setup_failure(run, harness_case, bridge_error, "bridge") @@ -174,7 +188,7 @@ def run_trace_cases( on_update(run) return 1, run for harness_case in cases: - _run_case(run, harness_case, selected_scenarios, on_update) + _run_case(run, harness_case, selected_scenarios, on_update, engine) run.finished_at = monotonic() on_update(run) failed: Final = any( diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py index 6be5afd60d6..1221f237570 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py @@ -1,10 +1,15 @@ from __future__ import annotations -import json from typing import Final -from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse from .....shared.tracing.steps import Engine, mapping +from ...fixtures import ( + anthropic_response_body, + anthropic_stream_events, + aws_event_stream_response, + json_response, + sse_response, +) from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite COMMON_MAPPINGS: Final = ( @@ -24,6 +29,22 @@ COMMON_MAPPINGS: Final = ( mapping(rust_span="execute_chat_completions_provider_call"), mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), mapping(rust_span="transform_response", python_frame=r"(? RouteFixture: - response: Final = json.dumps( - { - "id": "msg_trace", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-5", - "content": [{"type": "text", "text": "hello"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 3}, - } - ).encode() return RouteFixture( kwargs={ "model": "anthropic/claude-sonnet-5", "messages": [{"role": "user", "content": "hello"}], **({"optional_params": {"max_tokens": 16}} if engine == "rust" else {"max_tokens": 16}), }, - provider_responses=( - RecordedHttpResponse.from_bytes( - 200, (HttpHeader(name="content-type", value="application/json"),), response - ), - ), + provider_responses=(json_response(anthropic_response_body()),), ) def _bedrock_fixture(engine: Engine, _base_url: str) -> RouteFixture: - response: Final = json.dumps( - { - "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, - "stopReason": "end_turn", - "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5}, - "metrics": {"latencyMs": 1}, - } - ).encode() + response: Final[dict[str, object]] = { + "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5}, + "metrics": {"latencyMs": 1}, + } credentials: Final = { "aws_access_key_id": "test-access", "aws_secret_access_key": "test-secret", @@ -94,11 +97,60 @@ def _bedrock_fixture(engine: Engine, _base_url: str) -> RouteFixture: else {**credentials, "max_tokens": 16} ), }, + provider_responses=(json_response(response),), + ) + + +def _anthropic_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _anthropic_fixture(engine, _base_url) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(anthropic_stream_events()),), + consume_stream=True, + ) + + +def _bedrock_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _bedrock_fixture(engine, _base_url) + events: Final[tuple[dict[str, object], ...]] = ( + {"messageStart": {"role": "assistant"}}, + {"contentBlockStart": {"contentBlockIndex": 0, "start": {}}}, + {"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "hello"}}}, + {"contentBlockStop": {"contentBlockIndex": 0}}, + {"messageStop": {"stopReason": "end_turn"}}, + {"metadata": {"usage": {"inputTokens": 2, "outputTokens": 1, "totalTokens": 3}}}, + ) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(aws_event_stream_response(events),), + consume_stream=True, + ) + + +def _provider_error_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _anthropic_fixture(engine, _base_url) + return fixture.derive( provider_responses=( - RecordedHttpResponse.from_bytes( - 200, (HttpHeader(name="content-type", value="application/json"),), response + json_response( + {"type": "error", "error": {"type": "invalid_request_error", "message": "bad request"}}, + status=400, ), ), + expected_failure=True, + ) + + +def _stream_error_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _anthropic_fixture(engine, base_url) + events: Final = ( + anthropic_stream_events()[0], + ("error", {"type": "error", "error": {"type": "overloaded_error", "message": "overloaded"}}), + ) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(events),), + expected_failure=True, + consume_stream=True, ) @@ -132,18 +184,64 @@ TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( TraceScenario( - name="anthropic", + name="sync-anthropic", fixture=_anthropic_fixture, - mappings=COMMON_MAPPINGS, - sync_mappings=SYNC_MAPPINGS, - async_mappings=ASYNC_MAPPINGS, + mappings=SYNC_MAPPINGS, + asynchronous=False, ), TraceScenario( - name="bedrock", + name="async-anthropic", + fixture=_anthropic_fixture, + mappings=ASYNC_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="sync-anthropic-stream", + fixture=_anthropic_stream_fixture, + mappings=(*SYNC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=False, + ), + TraceScenario( + name="async-anthropic-stream", + fixture=_anthropic_stream_fixture, + mappings=(*ASYNC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-anthropic-provider-error", + fixture=_provider_error_fixture, + mappings=(*ASYNC_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-anthropic-stream-error", + fixture=_stream_error_fixture, + mappings=(*ASYNC_MAPPINGS, *STREAM_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="sync-bedrock", fixture=_bedrock_fixture, - mappings=BEDROCK_COMMON_MAPPINGS, - sync_mappings=BEDROCK_SYNC_MAPPINGS, - async_mappings=BEDROCK_ASYNC_MAPPINGS, + mappings=BEDROCK_SYNC_MAPPINGS, + asynchronous=False, + ), + TraceScenario( + name="async-bedrock", + fixture=_bedrock_fixture, + mappings=BEDROCK_ASYNC_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="sync-bedrock-event-stream", + fixture=_bedrock_stream_fixture, + mappings=(*BEDROCK_SYNC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=False, + ), + TraceScenario( + name="async-bedrock-event-stream", + fixture=_bedrock_stream_fixture, + mappings=(*BEDROCK_ASYNC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py index eb6c9233565..783c22a0dc0 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -1,18 +1,20 @@ from __future__ import annotations import asyncio -from collections.abc import Awaitable +import os +from collections.abc import AsyncIterable, Awaitable, Iterable from dataclasses import dataclass from pathlib import Path from typing import Final, Protocol, cast +from unittest.mock import patch from ....shared.parity.replay import replay_server from ....shared.reporting.models import Surface from ....shared.tracing.native import TraceResponsePayload, native_trace_events from ....shared.tracing.profiler import FunctionTraceEvent, profile_python from ....shared.tracing.steps import Engine, pipeline_projection -from ..models import RouteFixture, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario -from ..reporting import TraceComparisonArtifact +from ..models import RouteFixture, RouteSpec, TraceEngine, TraceExecutionFailure, TraceScenario +from ..reporting import TraceArtifact class SdkCall(Protocol): @@ -25,10 +27,20 @@ class _CollectedTrace: error: str | None = None -def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> object: +def _invoke( + function: SdkCall, + kwargs: dict[str, object], + *, + asynchronous: bool, + consume_stream: bool = False, +) -> object: async def invoke_async() -> object: try: - return await cast(Awaitable[object], function(**kwargs)) + response: Final = await cast(Awaitable[object], function(**kwargs)) + if consume_stream and isinstance(response, AsyncIterable): + stream = cast(AsyncIterable[object], response) + return tuple([item async for item in stream]) + return response finally: await asyncio.sleep(0) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -38,7 +50,10 @@ def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) if asynchronous: return asyncio.run(invoke_async()) - return function(**kwargs) + response: Final = function(**kwargs) + if consume_stream and isinstance(response, Iterable): + return tuple(cast(Iterable[object], response)) + return response def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCall | TraceExecutionFailure: @@ -47,6 +62,8 @@ def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCa from litellm.rust_bridge import get_native_bridge if engine == "rust": + if spec.rust_entrypoints is None: + return TraceExecutionFailure("rust", f"{spec.route} has no native Rust trace entrypoint") bridge: Final = cast(object | None, get_native_bridge()) if bridge is None: return TraceExecutionFailure("rust", "native Rust bridge is required for trace parity") @@ -62,14 +79,6 @@ def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCa return cast(SdkCall, getattr(owner, spec.python_entrypoints[int(asynchronous)])) -def _python_invocation_error(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> str | None: - try: - _invoke(function, kwargs, asynchronous=asynchronous) - except Exception as error: - return f"{type(error).__name__}: {error}" - return None - - def _collect( function: SdkCall, fixture: RouteFixture, @@ -83,14 +92,23 @@ def _collect( return _CollectedTrace(native_trace_events(payload), payload.error) import litellm - with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: - error: Final = _python_invocation_error(function, kwargs, asynchronous=asynchronous) + previous_suppress_debug_info: Final = litellm.suppress_debug_info + try: + if fixture.expected_failure: + litellm.suppress_debug_info = True + with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: + error: str | None + try: + _invoke(function, kwargs, asynchronous=asynchronous, consume_stream=fixture.consume_stream) + error = None + except Exception as caught: + error = f"{type(caught).__name__}: {caught}" + finally: + litellm.suppress_debug_info = previous_suppress_debug_info return _CollectedTrace(tuple(profiler.events), error) -def collect_trace( - spec: RouteSpec, engine: Engine, *, asynchronous: bool -) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: +def collect_trace(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: function: Final = _entrypoint(spec, engine, asynchronous=asynchronous) if isinstance(function, TraceExecutionFailure): return function @@ -101,15 +119,18 @@ def collect_trace( provider.enqueue_response(response) fixture: Final = RouteFixture( kwargs={ - **base_fixture.kwargs, "api_key": "test-key", + **base_fixture.kwargs, "api_base": provider.url, **({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}), }, provider_responses=base_fixture.provider_responses, expected_failure=base_fixture.expected_failure, + consume_stream=base_fixture.consume_stream, + environment=base_fixture.environment, ) - collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) + with patch.dict(os.environ, fixture.environment): + collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) provider.take_requests(len(fixture.provider_responses)) except Exception as error: return TraceExecutionFailure(engine, f"{type(error).__name__}: {error}") @@ -129,48 +150,56 @@ def _failure_message(result: tuple[FunctionTraceEvent, ...] | TraceExecutionFail def execute_trace( - route: RouteSpec, scenario: TraceScenario, mode: TraceMode, surface: Surface -) -> TraceComparisonArtifact: - asynchronous: Final = mode == "async" - mappings: Final = scenario.mappings_for(mode) + route: RouteSpec, + scenario: TraceScenario, + surface: Surface, + engine: TraceEngine = "both", +) -> TraceArtifact: + effective_engine: Final[TraceEngine] = "python" if engine == "both" and route.rust_entrypoints is None else engine scenario_route: Final = RouteSpec( route=route.route, python_entrypoints=route.python_entrypoints, rust_entrypoints=route.rust_entrypoints, fixture=scenario.fixture, ) - python_trace: Final = collect_trace(scenario_route, "python", asynchronous=asynchronous) - rust_trace: Final = collect_trace(scenario_route, "rust", asynchronous=asynchronous) + python_trace: Final = ( + collect_trace( + scenario_route, + "python", + asynchronous=scenario.asynchronous, + ) + if effective_engine != "rust" + else () + ) + rust_trace: Final = ( + collect_trace(scenario_route, "rust", asynchronous=scenario.asynchronous) + if effective_engine != "python" + else () + ) python_error: Final = _failure_message(python_trace) rust_error: Final = _failure_message(rust_trace) python_events: Final = python_trace if isinstance(python_trace, tuple) else () rust_events: Final = rust_trace if isinstance(rust_trace, tuple) else () try: - python: Final = pipeline_projection("python", python_events, mappings) - rust: Final = pipeline_projection("rust", rust_events, mappings) + python: Final = pipeline_projection("python", python_events) + rust: Final = pipeline_projection("rust", rust_events) except ValueError as error: - return TraceComparisonArtifact.from_traces( + return TraceArtifact.from_traces( + engine=effective_engine, surface=surface, sdk_function=route.route, scenario=scenario.name, - mode=mode, - mappings=mappings, - contract=scenario.contract, python=(), rust=(), - python_unmatched=0, python_error=f"harness: {error}", ) - return TraceComparisonArtifact.from_traces( + return TraceArtifact.from_traces( + engine=effective_engine, surface=surface, sdk_function=route.route, scenario=scenario.name, - mode=mode, - mappings=mappings, - contract=scenario.contract, python=python.steps, rust=rust.steps, - python_unmatched=python.unmatched, python_error=python_error, rust_error=rust_error, ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py index 27079c28cd8..211c454eadf 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py @@ -1,14 +1,26 @@ from __future__ import annotations -import json from typing import Final -from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse from .....shared.tracing.steps import Engine, mapping +from ...fixtures import ( + anthropic_response_body, + anthropic_stream_events, + aws_event_stream_response, + json_response, + sse_response, +) from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite COMMON_MAPPINGS: Final = ( mapping(rust_span="messages", python_frame=r"anthropic_interface/messages/__init__\.py:\d+ a?create$"), + mapping(span="python_sanitize_empty_content", python_frame=r"strip_empty_content_blocks_from_anthropic_messages$"), + mapping(span="python_sanitize_tool_ids", python_frame=r"sanitize_tool_use_ids_in_anthropic_messages$"), + mapping( + span="python_flatten_web_search", python_frame=r"flatten_unencrypted_web_search_results_in_anthropic_messages$" + ), + mapping(span="python_cache_control", python_frame=r"AnthropicCacheControlHook\.maybe_inject_cache_control$"), + mapping(span="python_pre_request_hooks", python_frame=r"_execute_pre_request_hooks$"), mapping( span="python_messages_provider_config", python_frame=r"ProviderConfigManager\.get_provider_anthropic_messages_config$", @@ -30,10 +42,42 @@ COMMON_MAPPINGS: Final = ( ), mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), mapping(rust_span="transform_response", python_frame=r"(? RouteFixture: conversation: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16} - response: Final = json.dumps( - { - "id": "msg_trace", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-5", - "content": [{"type": "text", "text": "hello"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 3}, - } - ).encode() return RouteFixture( kwargs={ "model": f"{provider}/claude-sonnet-5", **({"body": {**conversation, "model": "claude-sonnet-5"}} if engine == "rust" else conversation), }, - provider_responses=( - RecordedHttpResponse.from_bytes( - 200, (HttpHeader(name="content-type", value="application/json"),), response - ), - ), + provider_responses=(json_response(anthropic_response_body()),), ) @@ -88,11 +156,170 @@ def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: return _fixture(engine, "azure_ai") +def _bedrock_kwargs(engine: Engine) -> dict[str, object]: + conversation: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16} + return { + "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + **( + {"body": {**conversation, "model": "anthropic.claude-3-sonnet-20240229-v1:0"}} + if engine == "rust" + else conversation + ), + "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + "aws_region_name": "us-east-1", + } + + +def _bedrock_fixture(engine: Engine, _base_url: str) -> RouteFixture: + response_fixture: Final = _fixture(engine, "anthropic") + return RouteFixture(kwargs=_bedrock_kwargs(engine), provider_responses=response_fixture.provider_responses) + + +def _bedrock_retry_fixture(engine: Engine, _base_url: str) -> RouteFixture: + success_fixture: Final = _bedrock_fixture(engine, _base_url) + messages: Final = [ + {"role": "user", "content": "hello"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "old reasoning", "signature": ""}, + {"type": "text", "text": "partial answer"}, + ], + }, + {"role": "user", "content": "continue"}, + ] + kwargs: Final = { + **_bedrock_kwargs(engine), + **( + {"body": {"messages": messages, "max_tokens": 16, "model": "anthropic.claude-3-sonnet-20240229-v1:0"}} + if engine == "rust" + else {"messages": messages} + ), + } + return success_fixture.derive( + kwargs=kwargs, + provider_responses=( + json_response({"message": "messages.1.content.0: Invalid `signature` in `thinking` block"}, status=400), + *success_fixture.provider_responses, + ), + ) + + +def _mock_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _fixture(engine, "anthropic") + return fixture.derive(kwargs={"mock_response": "hello from mock"}, provider_responses=()) + + +def _provider_error_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _fixture(engine, "anthropic") + return fixture.derive( + provider_responses=( + json_response( + {"type": "error", "error": {"type": "invalid_request_error", "message": "bad request"}}, + status=400, + ), + ), + expected_failure=True, + ) + + +def _sync_unsupported_fixture(engine: Engine, base_url: str) -> RouteFixture: + if engine == "rust": + return _anthropic_fixture(engine, base_url) + fixture: Final = _fixture(engine, "anthropic") + return fixture.derive(provider_responses=(), expected_failure=True) + + +def _stream_fixture_for(engine: Engine, provider: str) -> RouteFixture: + fixture: Final = _fixture(engine, provider) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(anthropic_stream_events()),), + consume_stream=True, + ) + + +def _stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _stream_fixture_for(engine, "anthropic") + + +def _azure_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _stream_fixture_for(engine, "azure_ai") + + +def _bedrock_stream_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _bedrock_fixture(engine, base_url) + events: Final = tuple(payload for _, payload in anthropic_stream_events()) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(aws_event_stream_response(events),), + consume_stream=True, + ) + + +def _bedrock_stream_error_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _bedrock_fixture(engine, base_url) + start: Final = anthropic_stream_events(model="anthropic.claude-3-sonnet-20240229-v1:0")[0][1] + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(aws_event_stream_response((start, {"type": "message_stop"}), corrupt_last_frame=True),), + expected_failure=True, + consume_stream=True, + ) + + SPEC: Final = RouteSpec("messages", ("create", "acreate"), ("messages", "amessages"), _anthropic_fixture) TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( - TraceScenario(name="anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, modes=("async",)), - TraceScenario(name="azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, modes=("async",)), + TraceScenario( + name="async-anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, asynchronous=True + ), + TraceScenario(name="async-azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, asynchronous=True), + TraceScenario(name="async-bedrock", fixture=_bedrock_fixture, mappings=BEDROCK_MAPPINGS, asynchronous=True), + TraceScenario( + name="async-bedrock-invalid-thinking-retry", + fixture=_bedrock_retry_fixture, + mappings=RETRY_MAPPINGS, + asynchronous=True, + ), + TraceScenario(name="async-mock-response", fixture=_mock_fixture, mappings=MOCK_MAPPINGS, asynchronous=True), + TraceScenario( + name="async-anthropic-provider-error", + fixture=_provider_error_fixture, + mappings=ANTHROPIC_FAILURE_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="async-anthropic-stream", + fixture=_stream_fixture, + mappings=(*ANTHROPIC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-azure-ai-stream", + fixture=_azure_stream_fixture, + mappings=(*AZURE_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-bedrock-event-stream", + fixture=_bedrock_stream_fixture, + mappings=(*BEDROCK_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-bedrock-event-stream-error", + fixture=_bedrock_stream_error_fixture, + mappings=(*BEDROCK_MAPPINGS, *STREAM_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="sync-unsupported", + fixture=_sync_unsupported_fixture, + mappings=ANTHROPIC_MAPPINGS, + asynchronous=False, + ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py index 2a4a1b3a152..bb21e8ab0c5 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from typing import Final, cast +from typing import Final from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse from .....shared.tracing.steps import Engine, mapping @@ -53,6 +53,23 @@ ASYNC_MAPPINGS: Final = ( ), ) +PUBLIC_RUST_DISPATCH_MAPPINGS: Final = ( + mapping(span="public_sdk_entrypoint", python_frame=r"ocr/main\.py:\d+ a?ocr$"), + mapping(span="public_request", python_frame=r"ocr/main\.py:\d+ _public_request$"), + mapping(span="bind_request", python_frame=r"ocr/main\.py:\d+ _bind_request$"), + mapping(span="rust_ocr_enabled", python_frame=r"rust_bridge/configuration\.py:\d+ rust_ocr_enabled$"), + mapping(span="select_native_ocr", python_frame=r"rust_bridge/ocr_lifecycle\.py:\d+ select$"), + mapping(span="load_native_bridge", python_frame=r"rust_bridge/bindings\.py:\d+ NativeBinding\.load$"), + mapping(span="native_call_setup", python_frame=r"rust_bridge/lifecycle\.py:\d+ setup$"), + mapping(span="native_response", python_frame=r"rust_bridge/ocr\.py:\d+ _response$"), + mapping(span="native_call_finalize", python_frame=r"rust_bridge/lifecycle\.py:\d+ finalize$"), + mapping( + span="native_success_bookkeeping", + python_frame=r"rust_bridge/lifecycle\.py:\d+ success_bookkeeping$", + ), + *(mapping(rust_span=item.rust) for item in SYNC_MAPPINGS if item.rust is not None), +) + CALLBACK_SUCCESS_SYNC_MAPPINGS: Final = (*SYNC_MAPPINGS, SUCCESS_CALLBACK_SYNC_MAPPING) CALLBACK_SUCCESS_ASYNC_MAPPINGS: Final = (*ASYNC_MAPPINGS, SUCCESS_CALLBACK_ASYNC_MAPPING) CALLBACK_FAILURE_SYNC_MAPPINGS: Final = ( @@ -164,23 +181,6 @@ def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: ) -def _vertex_fixture(engine: Engine, _base_url: str) -> RouteFixture: - fixture: Final = _fixture( - engine, - "vertex_ai/mistral-ocr-maas", - {"type": "image_url", "image_url": "data:image/png;base64,aGVsbG8="}, - ) - vertex: Final = {"vertex_project": "trace-project", "vertex_location": "us-central1"} - optional_params: Final = cast(dict[str, object], fixture.kwargs.get("optional_params", {})) - return RouteFixture( - kwargs={ - **fixture.kwargs, - **({"optional_params": {**optional_params, **vertex}} if engine == "rust" else vertex), - }, - provider_responses=fixture.provider_responses, - ) - - def _vertex_deepseek_fixture(engine: Engine, _base_url: str) -> RouteFixture: vertex: Final = {"vertex_project": "trace-project", "vertex_location": "us-central1"} return RouteFixture( @@ -204,6 +204,62 @@ def _vertex_deepseek_fixture(engine: Engine, _base_url: str) -> RouteFixture: ) +def _vertex_deepseek_credentials_fixture(engine: Engine, base_url: str) -> RouteFixture: + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric import rsa + + fixture: Final = _vertex_deepseek_fixture(engine, base_url) + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + credentials: Final = json.dumps( + { + "type": "service_account", + "project_id": "trace-project", + "private_key_id": "trace-key", + "private_key": private_key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ).decode(), + "client_email": "trace@trace-project.iam.gserviceaccount.com", + "token_uri": f"{base_url}/token", + } + ) + return RouteFixture( + kwargs={**fixture.kwargs, "api_key": None}, + environment=(("VERTEXAI_CREDENTIALS", credentials), ("VERTEX_AI_API_KEY", "")), + provider_responses=( + RecordedHttpResponse.from_bytes( + 200, + (HttpHeader(name="content-type", value="application/json"),), + b'{"access_token":"trace-token","token_type":"Bearer","expires_in":3600}', + ), + *fixture.provider_responses, + ), + ) + + +def _cohere_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return RouteFixture( + kwargs={ + "model": "cohere/parse-v5.0", + "document": {"type": "image_url", "image_url": "data:image/png;base64,aGVsbG8="}, + **({"optional_params": {"output_format": "blocks"}} if engine == "rust" else {"output_format": "blocks"}), + }, + provider_responses=( + RecordedHttpResponse.from_bytes( + 200, + (HttpHeader(name="content-type", value="application/json"),), + json.dumps( + { + "pages": [{"index": 0, "blocks": [{"type": "text", "text": {"content": "hello"}}]}], + "meta": {"billed_units": {"pages": 1}}, + } + ).encode(), + ), + ), + ) + + def _azure_document_intelligence_fixture(engine: Engine, base_url: str) -> RouteFixture: completed: Final = json.dumps( { @@ -249,30 +305,6 @@ def _azure_document_intelligence_fixture(engine: Engine, base_url: str) -> Route ) -VERTEX_COMMON_MAPPINGS: Final = ( - *COMMON_MAPPINGS[:7], - mapping( - rust_span="transform_ocr_request", - python_frame=( - r"VertexAIOCRConfig\.(?:async_)?transform_ocr_request$" - r"|MistralOCRConfig\.transform_ocr_request$" - ), - ), - COMMON_MAPPINGS[-1], -) -VERTEX_SYNC_MAPPINGS: Final = ( - *VERTEX_COMMON_MAPPINGS, - mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"), - mapping(span="python_transform_ocr_response_wrapper", python_frame=r"BaseLLMHTTPHandler\._transform_ocr_response$"), - mapping(rust_span="transform_ocr_response", python_frame=r"MistralOCRConfig\.transform_ocr_response$"), -) -VERTEX_ASYNC_MAPPINGS: Final = ( - *VERTEX_COMMON_MAPPINGS, - mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"), - mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"), - mapping(rust_span="transform_ocr_response", python_frame=r"MistralOCRConfig\.transform_ocr_response$"), -) - DEEPSEEK_COMMON_MAPPINGS: Final = ( mapping(rust_span="ocr", python_frame=r"ocr/main\.py:\d+ a?ocr$"), mapping(rust_span="prepare_ocr_call", python_frame=r"ocr/main\.py:\d+ _prepare_ocr_request$"), @@ -352,59 +384,126 @@ DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS: Final = ( mapping(span="python_poll_http_request", python_frame=r"AsyncHTTPHandler\.get$"), ) +COHERE_COMMON_MAPPINGS: Final = ( + *COMMON_MAPPINGS[:7], + mapping( + rust_span="transform_ocr_request", + python_frame=r"CohereParseConfig\.(?:async_)?transform_ocr_request$", + ), + COMMON_MAPPINGS[-1], + mapping(rust_span="transform_ocr_response", python_frame=r"CohereParseConfig\.transform_ocr_response$"), +) +COHERE_ASYNC_MAPPINGS: Final = ( + *COHERE_COMMON_MAPPINGS, + mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"), + mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"), +) SPEC: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _mistral_fixture) TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( TraceScenario( - name="mistral", + name="sync-mistral", fixture=_mistral_fixture, - mappings=COMMON_MAPPINGS, - sync_mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, ), TraceScenario( - name="mistral-callback-success", + name="async-mistral", + fixture=_mistral_fixture, + mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, + ), + TraceScenario( + name="sync-mistral-callback-success", fixture=_mistral_callback_success_fixture, - mappings=COMMON_MAPPINGS, - sync_mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS, - async_mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS, + mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS, + asynchronous=False, ), TraceScenario( - name="mistral-callback-failure", + name="async-mistral-callback-success", + fixture=_mistral_callback_success_fixture, + mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="sync-mistral-callback-failure", fixture=_mistral_callback_failure_fixture, - mappings=(*COMMON_MAPPINGS, FAILURE_CALLBACK_MAPPING), - sync_mappings=CALLBACK_FAILURE_SYNC_MAPPINGS, - async_mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS, + mappings=CALLBACK_FAILURE_SYNC_MAPPINGS, + asynchronous=False, ), TraceScenario( - name="azure-ai", + name="async-mistral-callback-failure", + fixture=_mistral_callback_failure_fixture, + mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="sync-azure-ai", fixture=_azure_fixture, - mappings=AZURE_COMMON_MAPPINGS, - sync_mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, ), TraceScenario( - name="azure-document-intelligence", + name="async-azure-ai", + fixture=_azure_fixture, + mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, + ), + TraceScenario( + name="sync-azure-document-intelligence", fixture=_azure_document_intelligence_fixture, - mappings=DOCUMENT_INTELLIGENCE_COMMON_MAPPINGS, - sync_mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, ), TraceScenario( - name="vertex-ai", - fixture=_vertex_fixture, - mappings=VERTEX_COMMON_MAPPINGS, - sync_mappings=(*VERTEX_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*VERTEX_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + name="async-azure-document-intelligence", + fixture=_azure_document_intelligence_fixture, + mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, ), TraceScenario( - name="vertex-deepseek", + name="sync-vertex-deepseek", fixture=_vertex_deepseek_fixture, - mappings=DEEPSEEK_COMMON_MAPPINGS, - sync_mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, + ), + TraceScenario( + name="async-vertex-deepseek", + fixture=_vertex_deepseek_fixture, + mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, + ), + TraceScenario( + name="sync-vertex-deepseek-credentials", + fixture=_vertex_deepseek_credentials_fixture, + mappings=DEEPSEEK_SYNC_MAPPINGS, + asynchronous=False, + ), + TraceScenario( + name="async-vertex-deepseek-credentials", + fixture=_vertex_deepseek_credentials_fixture, + mappings=DEEPSEEK_ASYNC_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="async-cohere", + fixture=_cohere_fixture, + mappings=(*COHERE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, + ), + TraceScenario( + name="sync-public-rust-dispatch", + fixture=_mistral_fixture, + mappings=PUBLIC_RUST_DISPATCH_MAPPINGS, + asynchronous=False, + ), + TraceScenario( + name="async-public-rust-dispatch", + fixture=_mistral_fixture, + mappings=PUBLIC_RUST_DISPATCH_MAPPINGS, + asynchronous=True, ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py new file mode 100644 index 00000000000..3f6e540efc5 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +from typing import Final + +from .....shared.tracing.steps import Engine, mapping +from ...fixtures import ( + anthropic_response_body, + anthropic_stream_events, + json_response, + responses_body, + responses_stream_events, + sse_response, +) +from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite + +COMMON_MAPPINGS: Final = ( + mapping(span="python_responses", python_frame=r"responses/main\.py:\d+ a?responses$"), + mapping( + span="python_responses_provider_config", + python_frame=r"ProviderConfigManager\.get_provider_responses_api_config$", + ), + mapping(rust_span="responses_provider_config"), + mapping(rust_span="validate_environment", python_frame=r"validate_environment$"), + mapping(rust_span="complete_url", python_frame=r"get_complete_url$"), + mapping( + rust_span="transform_request", + python_frame=r"(? RouteFixture: + model: Final = "gpt-5" + return RouteFixture( + kwargs={ + "model": f"{provider}/{model}", + "input": "hello", + **({"body": {"model": model, "input": "hello"}} if engine == "rust" else {}), + }, + provider_responses=(json_response(responses_body(model=model)),), + ) + + +def _openai_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _native_fixture(engine, "openai") + + +def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _native_fixture(engine, "azure") + return fixture.derive(kwargs={"api_version": "2025-04-01-preview"}) + + +def _openai_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _openai_fixture(engine, _base_url) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(responses_stream_events()),), + consume_stream=True, + ) + + +def _provider_error_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _openai_fixture(engine, _base_url) + return fixture.derive( + provider_responses=( + json_response({"error": {"message": "bad request", "type": "invalid_request_error"}}, status=400), + ), + expected_failure=True, + ) + + +def _stream_failed_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _openai_fixture(engine, base_url) + failed_response: Final[dict[str, object]] = { + **responses_body(), + "status": "failed", + "output": [], + "error": {"message": "stream failed", "type": "server_error", "code": "server_error"}, + } + events: Final = ( + ( + "response.created", + {"type": "response.created", "response": {**failed_response, "status": "in_progress", "error": None}}, + ), + ("response.failed", {"type": "response.failed", "response": failed_response}), + ) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(events),), + expected_failure=True, + consume_stream=True, + ) + + +def _anthropic_bridge_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return RouteFixture( + kwargs={ + "model": "anthropic/claude-sonnet-5", + "input": "hello", + "max_output_tokens": 16, + **({"body": {"model": "claude-sonnet-5", "input": "hello"}} if engine == "rust" else {}), + }, + provider_responses=(json_response(anthropic_response_body()),), + ) + + +def _anthropic_bridge_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _anthropic_bridge_fixture(engine, _base_url) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(anthropic_stream_events()),), + consume_stream=True, + ) + + +SPEC: Final = RouteSpec("responses", ("responses", "aresponses"), None, _openai_fixture) +TRACE_SUITE: Final = TraceSuite( + route=SPEC, + scenarios=( + TraceScenario(name="sync-openai", fixture=_openai_fixture, mappings=COMMON_MAPPINGS, asynchronous=False), + TraceScenario(name="async-openai", fixture=_openai_fixture, mappings=COMMON_MAPPINGS, asynchronous=True), + TraceScenario( + name="sync-openai-stream", + fixture=_openai_stream_fixture, + mappings=(*COMMON_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=False, + ), + TraceScenario( + name="async-openai-stream", + fixture=_openai_stream_fixture, + mappings=(*COMMON_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-openai-provider-error", + fixture=_provider_error_fixture, + mappings=(*COMMON_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-openai-stream-failed", + fixture=_stream_failed_fixture, + mappings=(*COMMON_MAPPINGS, *STREAM_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario(name="async-azure", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, asynchronous=True), + TraceScenario( + name="async-anthropic-chat-bridge", + fixture=_anthropic_bridge_fixture, + mappings=BRIDGE_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="async-anthropic-chat-bridge-stream", + fixture=_anthropic_bridge_stream_fixture, + mappings=( + *BRIDGE_MAPPINGS, + mapping(span="python_chat_stream_wrapper", python_frame=r"CustomStreamWrapper\.__init__$"), + mapping(span="python_chat_stream_next", python_frame=r"CustomStreamWrapper\.__anext__$"), + mapping( + span="python_responses_bridge_stream_iterator", + python_frame=r"LiteLLMCompletionStreamingIterator\.__init__$|LiteLLMCompletionStreamingIterator\.__anext__$", + ), + ), + asynchronous=True, + ), + ), +) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py b/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py new file mode 100644 index 00000000000..d0dbd281a97 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +from importlib import import_module +from typing import Final, cast + +from ..models import TraceSuite + + +def _suite(module: str) -> TraceSuite: + loaded: Final = import_module(module) + candidate: Final = cast(object, getattr(loaded, "TRACE_SUITE")) + assert isinstance(candidate, TraceSuite) + return candidate + + +def test_core_sdk_scenario_matrix_keeps_distinct_migration_paths() -> None: + chat: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.chat_completions.case") + messages: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.messages.case") + ocr: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case") + responses: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.responses.case") + + assert {(scenario.name, scenario.asynchronous) for scenario in chat.scenarios} >= { + ("sync-anthropic", False), + ("async-anthropic", True), + ("sync-anthropic-stream", False), + ("async-anthropic-stream", True), + ("async-anthropic-provider-error", True), + ("async-anthropic-stream-error", True), + ("sync-bedrock", False), + ("async-bedrock", True), + ("sync-bedrock-event-stream", False), + ("async-bedrock-event-stream", True), + } + assert {(scenario.name, scenario.asynchronous) for scenario in messages.scenarios} >= { + ("async-anthropic-stream", True), + ("async-azure-ai-stream", True), + ("async-bedrock-event-stream", True), + ("async-bedrock-event-stream-error", True), + ("async-bedrock-invalid-thinking-retry", True), + ("sync-unsupported", False), + } + assert {(scenario.name, scenario.asynchronous) for scenario in ocr.scenarios} >= { + ("async-cohere", True), + ("sync-public-rust-dispatch", False), + ("async-public-rust-dispatch", True), + } + assert {(scenario.name, scenario.asynchronous) for scenario in responses.scenarios} >= { + ("sync-openai", False), + ("async-openai", True), + ("sync-openai-stream", False), + ("async-openai-stream", True), + ("async-openai-provider-error", True), + ("async-openai-stream-failed", True), + ("async-azure", True), + ("async-anthropic-chat-bridge", True), + ("async-anthropic-chat-bridge-stream", True), + } + + +def test_core_gateway_matrix_keeps_downstream_streams_separate() -> None: + modules: Final = ( + "tests.rust-python-harness.strategies.trace_parity.gateway.chat_completions.case", + "tests.rust-python-harness.strategies.trace_parity.gateway.messages.case", + "tests.rust-python-harness.strategies.trace_parity.gateway.responses.case", + ) + + for module in modules: + suite = _suite(module) + assert any("downstream-stream" in scenario.name for scenario in suite.scenarios) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/transcription/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/transcription/case.py index 2fb1054d10a..3b4d2e1447d 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/transcription/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/transcription/case.py @@ -92,11 +92,16 @@ TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( TraceScenario( - name="bedrock", + name="sync-bedrock", fixture=_fixture, - mappings=MAPPINGS, - sync_mappings=SYNC_MAPPINGS, - async_mappings=ASYNC_MAPPINGS, + mappings=SYNC_MAPPINGS, + asynchronous=False, + ), + TraceScenario( + name="async-bedrock", + fixture=_fixture, + mappings=ASYNC_MAPPINGS, + asynchronous=True, ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/test_reporting.py b/tests/rust-python-harness/strategies/trace_parity/test_reporting.py index 68264064c92..22cc87592b8 100644 --- a/tests/rust-python-harness/strategies/trace_parity/test_reporting.py +++ b/tests/rust-python-harness/strategies/trace_parity/test_reporting.py @@ -1,58 +1,46 @@ from __future__ import annotations -from collections.abc import Sequence from typing import Final, Literal import pytest from ...shared.reporting.models import CaseResult, Coverage, HarnessCase, ResultArtifact, RunStatus from ...shared.reporting.strategy import ModuleCaseSpec, NotImplementedCaseSpec -from ...shared.tracing.steps import PipelineStep, TraceContract, TraceMapping, mapping +from ...shared.tracing.steps import PipelineStep from . import reporting -from .reporting import TRACE_COMPARISON_ARTIFACT, TraceComparisonArtifact, render_trace_results - -MAPPINGS: Final = ( - mapping(rust_span="ocr", python_frame=r"ocr/main\.py:\d+ a?ocr$"), - mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$"), -) +from .reporting import TRACE_ARTIFACT, TraceArtifact, render_trace_results -def _result(comparison: TraceComparisonArtifact) -> CaseResult: +def _result(trace: TraceArtifact) -> CaseResult: case: Final = HarnessCase( strategy_id="trace_parity", strategy_label="Trace parity", - sdk_function=comparison.sdk_function, + sdk_function=trace.sdk_function, spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"), - surface=comparison.surface, + surface=trace.surface, ) result: Final = CaseResult(case=case) - nodeid: Final = f"trace:sdk:{comparison.sdk_function}:{comparison.scenario}:{comparison.mode}" + nodeid: Final = f"trace:{trace.surface}:{trace.sdk_function}:{trace.scenario}" result.collected.add(nodeid) - result.record( - nodeid, - RunStatus.PASSED, - artifacts=(ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()),), - ) + result.record(nodeid, RunStatus.PASSED, artifacts=(ResultArtifact(TRACE_ARTIFACT, trace.model_dump_json()),)) return result -def _comparison( +def _trace( python: tuple[PipelineStep, ...], rust: tuple[PipelineStep, ...], *, - mappings: Sequence[TraceMapping] = MAPPINGS, rust_error: str | None = None, -) -> TraceComparisonArtifact: - return TraceComparisonArtifact.from_traces( + engine: Literal["python", "rust", "both"] = "both", + scenario: str = "sync-default", +) -> TraceArtifact: + return TraceArtifact.from_traces( + engine=engine, surface="sdk", sdk_function="ocr", - scenario="default", - mode="sync", - mappings=mappings, - contract=TraceContract(), + scenario=scenario, python=python, rust=rust, - python_unmatched=796, rust_error=rust_error, ) @@ -67,107 +55,55 @@ def _events(*items: tuple[str, int, str | None]) -> tuple[PipelineStep, ...]: return tuple(steps) -def test_renderer_shows_matching_python_and_rust_paths() -> None: - rust: Final = _events(("ocr", 0, None), ("http_request", 1, None)) +def test_renderer_prints_python_and_rust_traces_independently() -> None: python: Final = _events( ("ocr", 0, "ocr/main.py:88 aocr"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), + ("python_prepare", 1, "prep.py:1 python_prepare"), ) - - section: Final = render_trace_results((_result(_comparison(python, rust)),))[0] - report: Final = "\n\n".join(section.blocks) - - assert section.title == "SDK trace comparisons" - assert "Case: ocr" in report - assert "PYTHON (2 steps)\n1 aocr (ocr/main.py:88)\n2 AsyncHTTPHandler.post (http_handler.py:673)" in report - assert "RUST (2 steps)\nocr -> 1 aocr\n http_request -> 2 AsyncHTTPHandler.post" in report - assert "Mapping (identifier -> span)" not in report - assert "Trace: MATCH" in report - assert "Same steps, order, and nesting" in report - assert "Unseen mappings:" not in report - - -def test_renderer_reports_mappings_that_matched_nothing() -> None: - events: Final = _events(("ocr", 0, None)) - - section: Final = render_trace_results((_result(_comparison(events, events)),))[0] - report: Final = "\n\n".join(section.blocks) - - assert "Unseen mappings: http_request" in report - assert "Contract: FAIL" in report - - -def test_renderer_numbers_repeated_span_occurrences() -> None: - mappings: Final = (MAPPINGS[0], MAPPINGS[1]) - rust: Final = _events(("ocr", 0, None), ("http_request", 1, None), ("http_request", 1, None)) - python: Final = _events( - ("ocr", 0, "ocr/main.py:88 aocr"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), - ) - - report: Final = "\n\n".join(render_trace_results((_result(_comparison(python, rust, mappings=mappings)),))[0].blocks) - - assert "http_request#2" in report - - -def test_renderer_accepts_declared_engine_specific_steps() -> None: - mappings: Final = ( - *MAPPINGS[:1], - mapping(span="python_prepare", python_frame=r"python_prepare$"), - mapping(rust_span="rust_prepare"), - ) - python: Final = _events(("ocr", 0, None), ("python_prepare", 1, "prep.py:1 python_prepare")) rust: Final = _events(("ocr", 0, None), ("rust_prepare", 1, None)) - section: Final = render_trace_results((_result(_comparison(python, rust, mappings=mappings)),))[0] + section: Final = render_trace_results((_result(_trace(python, rust)),))[0] report: Final = "\n\n".join(section.blocks) - assert "2 python_prepare (prep.py:1) [python only]" in report - assert "rust_prepare -> [rust only]" in report - assert "Trace: MATCH" in report - assert "Contract: PASS" in report + assert section.title == "SDK traces" + assert "PYTHON (2 steps)\n1 aocr (ocr/main.py:88)\n2 python_prepare (prep.py:1)" in report + assert "RUST (2 steps)\n1 ocr\n2 rust_prepare" in report + assert "python only" not in report + assert "rust only" not in report + assert " -> " not in report + assert "Trace: MATCH" not in report + assert "Trace: DRIFT" not in report + assert "Contract:" not in report -def test_unavailable_check_reports_mode_from_nodeid() -> None: - case: Final = HarnessCase( - strategy_id="trace_parity", - strategy_label="Trace parity", - sdk_function="ocr", - spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"), - surface="sdk", - ) - result: Final = CaseResult(case=case) - result.collected.add("trace:sdk:ocr:default:sync") - result.record("trace:sdk:ocr:default:sync", RunStatus.ERROR) +@pytest.mark.parametrize( + ("engine", "present", "absent"), + (("python", "PYTHON (1 steps)", "RUST"), ("rust", "RUST (1 steps)", "PYTHON")), +) +def test_renderer_prints_only_selected_engine(engine: Literal["python", "rust"], present: str, absent: str) -> None: + events: Final = _events(("ocr", 0, None)) - section: Final = render_trace_results((result,))[0] - report: Final = "\n\n".join(section.blocks) + report: Final = "\n\n".join(render_trace_results((_result(_trace(events, events, engine=engine)),))[0].blocks) - assert "Case: ocr" in report - assert "Scenario: default / Mode: sync" in report - assert "Trace: NOT AVAILABLE\nTest outcome: error" in report - assert "unknown mode" not in report + assert present in report + assert absent not in report def test_renderer_keeps_collected_trace_when_one_engine_errors() -> None: - python: Final = _events( - ("ocr", 0, "ocr/main.py:88 aocr"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), + python: Final = _events(("ocr", 0, "ocr/main.py:88 aocr")) + + report: Final = "\n\n".join( + render_trace_results( + (_result(_trace(python, (), rust_error="rust: native Rust bridge must include the trace-parity feature")),) + )[0].blocks ) - section: Final = render_trace_results( - (_result(_comparison(python, (), rust_error="rust: native Rust bridge must include the trace-parity feature")),) - )[0] - report: Final = "\n\n".join(section.blocks) - - assert "PYTHON (2 steps)\n1 aocr (ocr/main.py:88) [python only]" in report + assert "PYTHON (1 steps)\n1 aocr (ocr/main.py:88)" in report assert "Rust error: rust: native Rust bridge must include the trace-parity feature" in report assert "hint: rebuild the native bridge with the trace-parity feature" in report - assert "Contract: FAIL" in report -def test_renderer_groups_all_modes_under_one_case_header() -> None: +def test_unavailable_trace_reports_scenario_from_nodeid() -> None: case: Final = HarnessCase( strategy_id="trace_parity", strategy_label="Trace parity", @@ -176,76 +112,55 @@ def test_renderer_groups_all_modes_under_one_case_header() -> None: surface="sdk", ) result: Final = CaseResult(case=case) - events: Final = _events(("ocr", 0, None)) - modes: Final[tuple[Literal["sync", "async"], ...]] = ("sync", "async") - for mode in modes: - nodeid = f"trace:sdk:ocr:default:{mode}" - result.collected.add(nodeid) - comparison = TraceComparisonArtifact.from_traces( - surface="sdk", - sdk_function="ocr", - scenario="default", - mode=mode, - mappings=MAPPINGS, - contract=TraceContract(), - python=events, - rust=events, - python_unmatched=0, - ) - result.record( - nodeid, - RunStatus.PASSED, - artifacts=(ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()),), - ) + result.collected.add("trace:sdk:ocr:async-error") + result.record("trace:sdk:ocr:async-error", RunStatus.ERROR) - section: Final = render_trace_results((result,))[0] + report: Final = "\n\n".join(render_trace_results((result,))[0].blocks) + + assert "Scenario: async-error" in report + assert "Trace: NOT AVAILABLE\nTest outcome: error" in report + + +def test_renderer_groups_scenarios_under_one_case_header() -> None: + result: Final = _result(_trace(_events(("ocr", 0, None)), (), scenario="sync-default")) + async_trace: Final = _trace((), _events(("ocr", 0, None)), scenario="async-default") + nodeid: Final = "trace:sdk:ocr:async-default" + result.collected.add(nodeid) + result.record(nodeid, RunStatus.PASSED, artifacts=(ResultArtifact(TRACE_ARTIFACT, async_trace.model_dump_json()),)) + + report: Final = render_trace_results((result,))[0].blocks[0] - assert len(section.blocks) == 1 - report: Final = section.blocks[0] assert report.count("Case: ocr") == 1 - assert "Scenario: default / Mode: sync" in report - assert "Scenario: default / Mode: async" in report + assert "Scenario: sync-default" in report + assert "Scenario: async-default" in report def test_renderer_colors_every_trace_line_in_a_terminal(monkeypatch: pytest.MonkeyPatch) -> None: - rust: Final = _events(("ocr", 0, None), ("http_request", 1, None)) - python: Final = _events( - ("ocr", 0, "ocr/main.py:88 aocr"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), - ) + events: Final = _events(("ocr", 0, "ocr/main.py:88 aocr")) monkeypatch.setattr(reporting.sys.stdout, "isatty", lambda: True) monkeypatch.delenv("NO_COLOR", raising=False) - section: Final = render_trace_results((_result(_comparison(python, rust)),))[0] - report: Final = "\n\n".join(section.blocks) + report: Final = "\n\n".join(render_trace_results((_result(_trace(events, events)),))[0].blocks) - assert "\033[36mPYTHON\033[0m (2 steps)" in report + assert "\033[36mPYTHON\033[0m (1 steps)" in report assert "\033[36m1 aocr (ocr/main.py:88)\033[0m" in report - assert "\033[33mRUST\033[0m (2 steps)" in report - assert "\033[33mocr\033[0m -> \033[36m1 aocr\033[0m" in report - assert "\033[33mhttp_request\033[0m -> \033[36m2 AsyncHTTPHandler.post\033[0m" in report + assert "\033[33mRUST\033[0m (1 steps)" in report + assert "\033[33m1 ocr\033[0m" in report -def test_renderer_groups_cases_and_unavailable_entries_by_surface() -> None: - events: Final = _events(("ocr", 0, None)) - gateway_results: Final = tuple( - CaseResult( - case=HarnessCase( - strategy_id="trace_parity", - strategy_label="Trace parity", - sdk_function=sdk_function, - spec=NotImplementedCaseSpec(reason=f"No {sdk_function} case is registered."), - surface="gateway", - ), - status=RunStatus.NOT_IMPLEMENTED, - ) - for sdk_function in ("ocr", "messages") +def test_renderer_groups_unavailable_entries_by_surface() -> None: + gateway_result: Final = CaseResult( + case=HarnessCase( + strategy_id="trace_parity", + strategy_label="Trace parity", + sdk_function="messages", + spec=NotImplementedCaseSpec(reason="No messages case is registered."), + surface="gateway", + ), + status=RunStatus.NOT_IMPLEMENTED, ) - sections: Final = render_trace_results((_result(_comparison(events, events)), *gateway_results)) + sections: Final = render_trace_results((_result(_trace((), ())), gateway_result)) - assert tuple(section.title for section in sections) == ("SDK trace comparisons", "GATEWAY trace comparisons") - gateway_report: Final = "\n\n".join(sections[1].blocks) - assert gateway_report.count("Not implemented") == 1 - assert "- ocr: No ocr case is registered." in gateway_report - assert "- messages: No messages case is registered." in gateway_report + assert tuple(section.title for section in sections) == ("SDK traces", "GATEWAY traces") + assert "- messages: No messages case is registered." in "\n\n".join(sections[1].blocks) diff --git a/tests/rust-python-harness/strategies/trace_parity/test_runner.py b/tests/rust-python-harness/strategies/trace_parity/test_runner.py index 0fcc5860ff2..be25dd53b02 100644 --- a/tests/rust-python-harness/strategies/trace_parity/test_runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/test_runner.py @@ -1,12 +1,23 @@ from __future__ import annotations -from typing import Final +import importlib +import os +from pathlib import Path +from types import SimpleNamespace +from typing import Final, cast + +import pytest + +import litellm from ...shared.reporting.models import Coverage, HarnessCase, HarnessRun, RunStatus, SdkFunction, Surface from ...shared.reporting.strategy import ModuleCaseSpec -from ...shared.tracing.steps import Engine +from ...shared.tracing.profiler import FunctionTraceEvent +from ...shared.tracing.steps import Engine, PipelineStep, mapping from .models import GatewayRouteSpec, RouteFixture, RouteSpec, TraceScenario, TraceSuite -from .runner import run_trace_mode, scenario_nodeids, validate_trace_suite +from .reporting import TraceArtifact +from .runner import run_trace_cases, run_trace_scenario, runner_selection, scenario_nodeids, validate_trace_suite +from .sdk.execution import SdkCall, collect_trace, execute_trace def _fixture(_engine: Engine, _base_url: str) -> RouteFixture: @@ -27,46 +38,243 @@ def test_scenario_filtering_and_occurrence_node_ids() -> None: suite: Final = TraceSuite( route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture), scenarios=( - TraceScenario("one", _fixture, (), modes=("sync", "async")), - TraceScenario("two", _fixture, (), modes=("async",)), + TraceScenario("sync-one", _fixture, (), asynchronous=False), + TraceScenario("async-one", _fixture, (), asynchronous=True), + TraceScenario("async-two", _fixture, (), asynchronous=True), ), ) case: Final = _case() - nodes: Final = scenario_nodeids(suite, case, frozenset({"two"})) + nodes: Final = scenario_nodeids(suite, case, frozenset({"async-two"})) - assert tuple(nodeid for _, _, nodeid in nodes) == ("trace:sdk:ocr:two:async",) + assert tuple(nodeid for _, nodeid in nodes) == ("trace:sdk:ocr:async-two",) + + +def test_python_engine_is_separate_from_scenario_selection() -> None: + assert runner_selection(("mistral", "--engine=python")) == (frozenset({"mistral"}), "python") + + +def test_python_engine_skips_native_bridge(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner") + case: Final = _case() + selected: list[tuple[frozenset[str], str]] = [] + + def reject_bridge(_repo_root: Path) -> str | None: + raise AssertionError("Python-only tracing must not inspect or build the native bridge") + + def capture_case( + _run: HarnessRun, + _case: HarnessCase, + scenarios: frozenset[str], + _on_update: object, + engine: str, + ) -> None: + selected.append((scenarios, engine)) + + monkeypatch.setattr(runner, "ensure_trace_bridge", reject_bridge) + monkeypatch.setattr(runner, "_run_case", capture_case) + + exit_code, _ = run_trace_cases((case,), tmp_path, lambda _: None, ("mistral", "--engine=python")) + + assert exit_code == 0 + assert selected == [(frozenset({"mistral"}), "python")] + + +def test_python_trace_preserves_native_ocr_dispatch_setting(monkeypatch: pytest.MonkeyPatch) -> None: + execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.execution") + route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture) + observed: list[str | None] = [] + + def collect( + _function: SdkCall, + _fixture: RouteFixture, + _engine: Engine, + *, + asynchronous: bool, + ) -> SimpleNamespace: + observed.append(os.environ.get("LITELLM_RUST")) + return SimpleNamespace( + events=(FunctionTraceEvent(0, None, "aocr" if asynchronous else "ocr"),), + error=None, + ) + + monkeypatch.setattr(execution, "_collect", collect) + + monkeypatch.setenv("LITELLM_RUST", "0") + collect_trace(route, "python", asynchronous=False) + monkeypatch.setenv("LITELLM_RUST", "1") + collect_trace(route, "python", asynchronous=True) + + assert observed == ["0", "1"] + assert os.environ["LITELLM_RUST"] == "1" + + +def test_expected_provider_failure_omits_feedback_banner( + capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch +) -> None: + loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.responses.case") + suite: Final = cast(TraceSuite, loaded.TRACE_SUITE) + scenario: Final = next(item for item in suite.scenarios if item.name == "async-openai-provider-error") + monkeypatch.setattr(litellm, "suppress_debug_info", False) + assert isinstance(suite.route, RouteSpec) + + result: Final = execute_trace(suite.route, scenario, "sdk", engine="python") + + assert result.python_error is None + assert "Give Feedback / Get Help" not in capsys.readouterr().out + assert litellm.suppress_debug_info is False + + +@pytest.mark.parametrize("asynchronous", (False, True)) +def test_vertex_trace_keeps_unmapped_helpers_and_parents(asynchronous: bool) -> None: + loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case") + suite: Final = cast(TraceSuite, loaded.TRACE_SUITE) + name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek" + scenario: Final = next(item for item in suite.scenarios if item.name == name) + assert isinstance(suite.route, RouteSpec) + + trace: Final = execute_trace(suite.route, scenario, "sdk", engine="python") + + assert trace.python_error is None + url: Final = next( + event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.get_complete_url") + ) + project: Final = next( + event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_project") + ) + location: Final = next( + event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_location") + ) + assert project.parent_id == location.parent_id == url.id + assert not any(event.raw.endswith(" VertexBase.get_access_token") for event in trace.python) + + +@pytest.mark.parametrize("asynchronous", (False, True)) +def test_vertex_credentials_trace_runs_real_auth_helpers(asynchronous: bool, monkeypatch: pytest.MonkeyPatch) -> None: + loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case") + suite: Final = cast(TraceSuite, loaded.TRACE_SUITE) + name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek-credentials" + scenario: Final = next(item for item in suite.scenarios if item.name == name) + monkeypatch.setenv("VERTEXAI_CREDENTIALS", "original-credentials") + monkeypatch.setenv("VERTEX_AI_API_KEY", "original-api-key") + assert isinstance(suite.route, RouteSpec) + + trace: Final = execute_trace(suite.route, scenario, "sdk", engine="python") + + assert trace.python_error is None + validate: Final = next( + event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.validate_environment") + ) + helpers: Final = ( + "VertexBase.safe_get_vertex_ai_project", + "VertexBase.safe_get_vertex_ai_credentials", + "VertexBase.get_access_token", + ) + assert tuple(event.raw.split(" ", 1)[1] for event in trace.python if event.parent_id == validate.id) == helpers + token: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.get_access_token")) + load: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.load_auth")) + refresh: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.refresh_auth")) + assert load.parent_id == token.id + assert refresh.parent_id == load.id + assert os.environ["VERTEXAI_CREDENTIALS"] == "original-credentials" + assert os.environ["VERTEX_AI_API_KEY"] == "original-api-key" + + +def test_gateway_trace_keeps_calls_outside_scenario_mappings(monkeypatch: pytest.MonkeyPatch) -> None: + execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.gateway.execution") + events: Final = ( + FunctionTraceEvent(0, None, "route.py:1 entry"), + FunctionTraceEvent(1, 0, "auth.py:2 authenticate"), + FunctionTraceEvent(2, 1, "auth.py:3 credentials"), + ) + scenario: Final = TraceScenario( + "async-gateway", + _fixture, + (mapping(rust_span="entry", python_frame=r" entry$"),), + asynchronous=True, + ) + monkeypatch.setattr(execution, "_collect", lambda *_args: events) + + trace: Final = execution.execute_gateway_trace(GatewayRouteSpec("messages"), scenario, engine="python") + + assert trace.python_error is None + assert tuple((event.id, event.parent_id, event.raw) for event in trace.python) == tuple( + (event.id, event.parent_id, event.raw) for event in events + ) + + +def test_default_trace_skips_unavailable_rust_sdk_entrypoint(monkeypatch: pytest.MonkeyPatch) -> None: + execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.execution") + route: Final = RouteSpec("responses", ("responses", "aresponses"), None, _fixture) + scenario: Final = TraceScenario("sync-openai", _fixture, (), asynchronous=False) + engines: list[Engine] = [] + + def collect(_route: RouteSpec, engine: Engine, *, asynchronous: bool) -> tuple[FunctionTraceEvent, ...]: + engines.append(engine) + return (FunctionTraceEvent(0, None, "responses"),) + + monkeypatch.setattr(execution, "collect_trace", collect) + + trace: Final = execution.execute_trace(route, scenario, "sdk") + + assert engines == ["python"] + assert trace.engine == "python" + assert trace.rust_error is None + + +def test_default_trace_skips_unavailable_rust_gateway_route(monkeypatch: pytest.MonkeyPatch) -> None: + execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.gateway.execution") + route: Final = GatewayRouteSpec("responses", rust_supported=False) + scenario: Final = TraceScenario("async-openai", _fixture, (), asynchronous=True) + engines: list[Engine] = [] + + def collect(_route: GatewayRouteSpec, _scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEvent, ...]: + engines.append(engine) + return (FunctionTraceEvent(0, None, "responses"),) + + monkeypatch.setattr(execution, "_collect", collect) + + trace: Final = execution.execute_gateway_trace(route, scenario) + + assert engines == ["python"] + assert trace.engine == "python" + assert trace.rust_error is None def test_scenario_validation_rejects_duplicate_and_unsafe_names() -> None: route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture) duplicate: Final = TraceSuite( route=route, - scenarios=(TraceScenario("same", _fixture, ()), TraceScenario("same", _fixture, ())), + scenarios=( + TraceScenario("sync-same", _fixture, (), asynchronous=False), + TraceScenario("sync-same", _fixture, (), asynchronous=False), + ), + ) + unsafe: Final = TraceSuite( + route=route, scenarios=(TraceScenario("sync-bad:name", _fixture, (), asynchronous=False),) ) - unsafe: Final = TraceSuite(route=route, scenarios=(TraceScenario("bad:name", _fixture, ()),)) case: Final = _case() assert validate_trace_suite(duplicate, case) is not None assert validate_trace_suite(unsafe, case) is not None -def test_scenario_validation_rejects_invalid_modes_and_route_registration() -> None: - invalid_modes: Final = TraceSuite( +def test_scenario_validation_rejects_invalid_names_and_route_registration() -> None: + invalid_name: Final = TraceSuite( route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture), - scenarios=(TraceScenario("invalid", _fixture, (), modes=("sync", "sync")),), + scenarios=(TraceScenario("bedrock", _fixture, (), asynchronous=True),), ) wrong_function: Final = TraceSuite( route=RouteSpec("messages", ("create", "acreate"), ("messages", "amessages"), _fixture), - scenarios=(TraceScenario("one", _fixture, ()),), + scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),), ) wrong_surface: Final = TraceSuite( route=GatewayRouteSpec("ocr"), - scenarios=(TraceScenario("one", _fixture, ()),), + scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),), ) case: Final = _case() - assert "unique sync/async modes" in (validate_trace_suite(invalid_modes, case) or "") + assert "start with sync- or async-" in (validate_trace_suite(invalid_name, case) or "") assert "does not match case function" in (validate_trace_suite(wrong_function, case) or "") assert "must use RouteSpec" in (validate_trace_suite(wrong_surface, case) or "") @@ -77,11 +285,35 @@ def test_invalid_route_dispatch_records_harness_error() -> None: result: Final = run.results[case.key] suite: Final = TraceSuite( route=GatewayRouteSpec("ocr"), - scenarios=(TraceScenario("one", _fixture, (), modes=("sync",)),), + scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),), ) - nodeid: Final = "trace:sdk:ocr:one:sync" + nodeid: Final = "trace:sdk:ocr:sync-one" - run_trace_mode(run, result, suite, suite.scenarios[0], "sync", "sdk", nodeid, lambda _: None) + run_trace_scenario(run, result, suite, suite.scenarios[0], "sdk", nodeid, lambda _: None) assert result.outcomes[nodeid] is RunStatus.ERROR assert run.failures == [(nodeid, "gateway route cannot run on the sdk surface")] + + +def test_different_python_and_rust_traces_pass(monkeypatch: pytest.MonkeyPatch) -> None: + runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner") + case: Final = _case() + run: Final = HarnessRun.from_cases((case,)) + result: Final = run.results[case.key] + suite: Final = TraceSuite( + route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture), + scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),), + ) + trace: Final = TraceArtifact.from_traces( + surface="sdk", + sdk_function="ocr", + scenario="sync-one", + python=(PipelineStep(0, None, "python_step", "python.py:1 python_step"),), + rust=(PipelineStep(0, None, "rust_step", "rust_step"),), + ) + monkeypatch.setattr(runner, "_execute_scenario", lambda *_args: trace) + + run_trace_scenario(run, result, suite, suite.scenarios[0], "sdk", "trace:sdk:ocr:sync-one", lambda _: None) + + assert result.outcomes["trace:sdk:ocr:sync-one"] is RunStatus.PASSED + assert run.failures == [] diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py b/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py index 98ea0b02e68..9dd79e860e6 100644 --- a/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py @@ -38,30 +38,28 @@ def _trace_functions( python_functions: Final[dict[str, PythonFunctionIdentity]] = {} rust_functions: Final[dict[str, RustFunctionIdentity]] = {} for scenario in suite.scenarios: - for mode in scenario.modes: - route: Final = RouteSpec( - route=suite.route.route, - python_entrypoints=suite.route.python_entrypoints, - rust_entrypoints=suite.route.rust_entrypoints, - fixture=scenario.fixture, - ) - python_trace: Final = collect_trace(route, "python", asynchronous=mode == "async") - rust_trace: Final = collect_trace(route, "rust", asynchronous=mode == "async") - if isinstance(python_trace, TraceExecutionFailure): - raise ValueError(f"Python trace discovery failed for {scenario.name}/{mode}: {python_trace.message}") - if isinstance(rust_trace, TraceExecutionFailure): - raise ValueError(f"Rust trace discovery failed for {scenario.name}/{mode}: {rust_trace.message}") - mappings: Final = scenario.mappings_for(mode) - python_projection: Final = pipeline_projection("python", python_trace, mappings) - rust_projection: Final = pipeline_projection("rust", rust_trace, mappings) - for step in python_projection.steps: - if step.span in spec.trace_spans: - function: Final = PythonFunctionIdentity.from_trace(step.raw) - python_functions[function.raw] = function - for step in rust_projection.steps: - if step.span in spec.trace_spans: - function: Final = RustFunctionIdentity.from_trace(step.raw) - rust_functions[step.raw] = function + route: Final = RouteSpec( + route=suite.route.route, + python_entrypoints=suite.route.python_entrypoints, + rust_entrypoints=suite.route.rust_entrypoints, + fixture=scenario.fixture, + ) + python_trace: Final = collect_trace(route, "python", asynchronous=scenario.asynchronous) + rust_trace: Final = collect_trace(route, "rust", asynchronous=scenario.asynchronous) + if isinstance(python_trace, TraceExecutionFailure): + raise ValueError(f"Python trace discovery failed for {scenario.name}: {python_trace.message}") + if isinstance(rust_trace, TraceExecutionFailure): + raise ValueError(f"Rust trace discovery failed for {scenario.name}: {rust_trace.message}") + python_projection: Final = pipeline_projection("python", python_trace, scenario.mappings) + rust_projection: Final = pipeline_projection("rust", rust_trace, scenario.mappings) + for step in python_projection.steps: + if step.span in spec.trace_spans: + function: Final = PythonFunctionIdentity.from_trace(step.raw) + python_functions[function.raw] = function + for step in rust_projection.steps: + if step.span in spec.trace_spans: + function: Final = RustFunctionIdentity.from_trace(step.raw) + rust_functions[step.raw] = function if not python_functions or not rust_functions: raise ValueError(f"Python trace discovery found no functions for spans: {', '.join(spec.trace_spans)}") return ( diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 2a33d84ec78..7e6d4d24905 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -10,9 +10,9 @@ from websockets.exceptions import ConnectionClosed from websockets.frames import Close import litellm +from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import ( - REALTIME_SESSION_SUCCESS_LOGGED_KEY, RealTimeStreaming, client_sent_openai_beta_realtime_header, ) @@ -3399,6 +3399,26 @@ async def test_refused_session_does_not_stamp_the_reservation_ownership_marker() assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details +@pytest.mark.asyncio +async def test_refused_session_stamps_the_failure_ownership_marker(): + """LIT-6463: the enqueued failure callback releases the key's max_parallel_requests + slot from the logging worker, so a refusal stamps REALTIME_SESSION_FAILURE_LOGGED_KEY. + The proxy endpoint reads it to leave the slot to that callback instead of racing it. + A session that relayed frames logs a success and must not carry the failure stamp.""" + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + refused: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close)) + session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode() + relayed: Final = _relay_session( + _client_ws_that_never_sends(), _backend_ws_closing_with(session_created, upstream_close) + ) + + await refused.run() + await relayed.run() + + assert refused.logging.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY) is True + assert REALTIME_SESSION_FAILURE_LOGGED_KEY not in relayed.logging.model_call_details + + @pytest.mark.asyncio async def test_transformed_transcription_completion_never_sends_response_create(): from typing import Final diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 2eadaee9e1a..ac3a43b742f 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -8,6 +8,7 @@ from unittest.mock import MagicMock import pytest import litellm +from litellm.constants import REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig @@ -77,6 +78,7 @@ class UnavailableBedrockStream: class FakeLogging: def __init__(self, trace_id="trace-nova-sonic"): self.litellm_trace_id = trace_id + self.model_call_details = {} class DisconnectingClientWS: @@ -673,6 +675,31 @@ class TestBedrockRealtimeProviderFailurePropagation: await spend_dispatch["coro"] assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] + @pytest.mark.asyncio + async def test_success_dispatch_stamps_the_ownership_marker_only_when_spend_was_logged( + self, stub_aws_sdk_client, spend_dispatch + ): + stub_aws_sdk_client["streams"] = [ScriptedBedrockStream(self.TEXT_TURN)] + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=spend_dispatch["logging_obj"], + **self.AWS_PARAMS, + ) + await spend_dispatch["coro"] + assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] + assert spend_dispatch["logging_obj"].model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True + + idle_logging = FakeLogging() + stub_aws_sdk_client["streams"] = [ScriptedBedrockStream([])] + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=idle_logging, + **self.AWS_PARAMS, + ) + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in idle_logging.model_call_details + @pytest.mark.asyncio async def test_stream_failure_after_client_disconnect_is_not_a_provider_failure(self, stub_aws_sdk_client): stream = ScriptedBedrockStream([], receiver_type=BreakingBedrockReceiver) @@ -684,6 +711,25 @@ class TestBedrockRealtimeProviderFailurePropagation: assert stream.input_stream.closed + @pytest.mark.asyncio + async def test_client_disconnect_ends_the_session_while_bedrock_output_stays_open(self, stub_aws_sdk_client): + receiver = DrainedThenOpenBedrockReceiver([]) + stream = ScriptedBedrockStream([], receiver_type=lambda _payloads: receiver) + stub_aws_sdk_client["streams"] = [stream] + + await asyncio.wait_for( + BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=RealtimeClientWS(), + logging_obj=FakeLogging(), + **self.AWS_PARAMS, + ), + timeout=1, + ) + + assert receiver.drained.is_set(), "the handler must have been waiting on the open provider stream" + assert stream.input_stream.closed + @pytest.mark.asyncio async def test_session_updated_is_not_sent_before_bedrock_is_ready(self, stub_aws_models): handler = BedrockRealtime() diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index ea2c925212d..5f87f2def93 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5527,7 +5527,9 @@ async def _run_internal_user_budget_alert( with ( patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: common_checks has no database seam - patch("litellm.proxy.proxy_server.get_current_spend", _get_spend), # test-quality-ok: common_checks imports it locally + patch( # test-quality-ok: common_checks imports get_current_spend locally + "litellm.proxy.proxy_server.get_current_spend", _get_spend + ), patch.object(slack_alerting, "send_alert", send_alert), ): error: Final = await _check_for_error() @@ -6419,9 +6421,7 @@ async def test_get_team_membership_negative_caches_a_missing_row(): assert first is None assert second is None mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() - cached = await cache.async_get_cache( - key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1") - ) + cached = await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")) assert cached == NO_TEAM_MEMBERSHIP_SENTINEL @@ -6454,6 +6454,314 @@ async def test_get_team_membership_reads_sentinel_as_no_membership_not_a_model() mock_prisma_client.db.litellm_teammembership.find_unique.assert_not_awaited() +@pytest.mark.asyncio +async def test_get_team_membership_coalesces_parallel_db_fetches(): + from litellm.proxy.auth.auth_checks import get_team_membership + + started = asyncio.Event() + release = asyncio.Event() + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-parallel", "team_id": "t-parallel", "spend": 1.0} + + async def _slow_find_unique(*args, **kwargs): + started.set() + await release.wait() + return membership_row + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_slow_find_unique) + cache = UserApiKeyCache() + + async def _load(): + return await get_team_membership( + user_id="u-parallel", + team_id="t-parallel", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + + first = asyncio.create_task(_load()) + second = asyncio.create_task(_load()) + await started.wait() + await asyncio.sleep(0) + release.set() + results = await asyncio.gather(first, second) + + assert results[0] is not None and results[1] is not None + assert results[0].user_id == "u-parallel" + assert results[1].user_id == "u-parallel" + mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_get_team_membership_invalidation_waits_for_in_flight_load_then_evicts_it(): + from litellm.proxy._types import LiteLLM_TeamMembership + from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state + from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + started = asyncio.Event() + release_stale = asyncio.Event() + rows = iter(("budget-old", "budget-new")) + + async def _find_unique(*args, **kwargs): + budget_id = next(rows) + row = MagicMock() + row.dict = lambda: {"user_id": "u-inv", "team_id": "t-inv", "spend": 1.0, "budget_id": budget_id} + if budget_id == "budget-old": + started.set() + await release_stale.wait() + return row + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_find_unique) + cache = UserApiKeyCache() + _key = team_membership_reservation_cache_key(user_id="u-inv", team_id="t-inv") + + async def _load(): + return await get_team_membership( + user_id="u-inv", team_id="t-inv", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + + stale = asyncio.create_task(_load()) + await asyncio.wait_for(started.wait(), timeout=2) + invalidation = asyncio.create_task( + invalidate_team_member_spend_state(user_id="u-inv", team_id="t-inv", user_api_key_cache=cache) + ) + for _ in range(5): + await asyncio.sleep(0) + assert not invalidation.done() + + release_stale.set() + await asyncio.wait_for(invalidation, timeout=2) + stale_result = await stale + assert stale_result is not None and stale_result.budget_id == "budget-old" + assert await cache.async_get_cache(key=_key) is None + + fresh_result = await _load() + assert fresh_result is not None and fresh_result.budget_id == "budget-new" + assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 + cached = CacheCodec.deserialize(await cache.async_get_cache(key=_key), model_type=LiteLLM_TeamMembership) + assert cached is not None and cached.budget_id == "budget-new" + again = await _load() + assert again is not None and again.budget_id == "budget-new" + assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 + + +@pytest.mark.asyncio +async def test_get_team_membership_invalidation_during_cache_write_evicts_stale_entry(): + from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + write_started = asyncio.Event() + release_write = asyncio.Event() + + class _SlowWriteCache(UserApiKeyCache): + async def async_set_cache(self, key, value, local_only=False, **kwargs): + write_started.set() + await release_write.wait() + return await super().async_set_cache(key, value, local_only=local_only, **kwargs) + + row = MagicMock() + row.dict = lambda: {"user_id": "u-w", "team_id": "t-w", "spend": 1.0, "budget_id": "budget-old"} + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=row) + cache = _SlowWriteCache() + + stale = asyncio.create_task( + get_team_membership(user_id="u-w", team_id="t-w", prisma_client=mock_prisma_client, user_api_key_cache=cache) + ) + await asyncio.wait_for(write_started.wait(), timeout=2) + invalidation = asyncio.create_task( + invalidate_team_member_spend_state(user_id="u-w", team_id="t-w", user_api_key_cache=cache) + ) + for _ in range(5): + await asyncio.sleep(0) + assert not invalidation.done() + + release_write.set() + await asyncio.wait_for(invalidation, timeout=2) + stale_result = await stale + + assert stale_result is not None and stale_result.budget_id == "budget-old" + assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-w", team_id="t-w")) is None + + +@pytest.mark.asyncio +async def test_common_checks_calls_get_team_membership_once_per_request(): + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + team = LiteLLM_TeamTable(team_id="t-once") + token = UserAPIKeyAuth(token="k-once", user_id="u-once", team_id="t-once", models=["gpt-4o-mini"]) + membership = MagicMock() + membership.litellm_budget_table = None + membership.spend = 0.0 + + with ( + patch( # test-quality-ok: common_checks imports prisma_client from proxy_server + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: common_checks imports user_api_key_cache from proxy_server + "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache() + ), + patch( # test-quality-ok: counts membership loads; common_checks has no membership seam + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=membership, + ) as load_membership, + patch( # test-quality-ok: common_checks imports get_current_spend locally + "litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0 + ), + ): + result = await common_checks( + request_body={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, + team_object=team, + user_object=LiteLLM_UserTable(user_id="u-once"), + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + + assert result is True + assert load_membership.await_count == 1 + + +@pytest.mark.asyncio +async def test_common_checks_skips_membership_load_when_no_check_reads_it(): + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + team = LiteLLM_TeamTable(team_id="t-lazy") + token = UserAPIKeyAuth(token="k-lazy", user_id="u-lazy", team_id="t-lazy") + + with ( + patch( # test-quality-ok: common_checks imports prisma_client from proxy_server + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: common_checks imports user_api_key_cache from proxy_server + "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache() + ), + patch( # test-quality-ok: counts membership loads; common_checks has no membership seam + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + ) as load_membership, + ): + result = await common_checks( + request_body={}, + team_object=team, + user_object=LiteLLM_UserTable(user_id="u-lazy"), + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/key/info", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + + assert result is True + load_membership.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_team_membership_db_error_returns_none_and_retries_next_call(): + from litellm.proxy.auth.auth_checks import get_team_membership + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-fail", "team_id": "t-fail", "spend": 1.0} + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock( + side_effect=[RuntimeError("db down"), membership_row] + ) + cache = UserApiKeyCache() + + failed = await get_team_membership( + user_id="u-fail", + team_id="t-fail", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + cached_after_failure = await cache.async_get_cache( + key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail") + ) + recovered = await get_team_membership( + user_id="u-fail", + team_id="t-fail", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + + assert failed is None + assert cached_after_failure is None + assert recovered is not None + assert recovered.user_id == "u-fail" + assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 + + +@pytest.mark.asyncio +async def test_get_team_membership_string_prisma_client_returns_none(): + from litellm.proxy.auth.auth_checks import get_team_membership + + result = await get_team_membership( + user_id="u-str", + team_id="t-str", + prisma_client="hello-world", + user_api_key_cache=UserApiKeyCache(), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load(): + from litellm.proxy.auth.auth_checks import get_team_membership + + started = asyncio.Event() + release = asyncio.Event() + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-shield", "team_id": "t-shield", "spend": 1.0} + + async def _slow_find_unique(*args, **kwargs): + started.set() + await release.wait() + return membership_row + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_slow_find_unique) + cache = UserApiKeyCache() + + async def _load(): + return await get_team_membership( + user_id="u-shield", + team_id="t-shield", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + + owner = asyncio.create_task(_load()) + await started.wait() + waiter = asyncio.create_task(_load()) + await asyncio.sleep(0) + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + release.set() + result = await owner + + assert result is not None + assert result.user_id == "u-shield" + mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + + @pytest.mark.asyncio async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sentinel(): """ @@ -6477,10 +6785,7 @@ async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sent assert before is None await invalidate_team_member_spend_state(user_id="u-1", team_id="t-1", user_api_key_cache=cache) - assert ( - await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")) - is None - ) + assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")) is None after = await get_team_membership( user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache diff --git a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py index fb82d8708fd..7506bd031d9 100644 --- a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py +++ b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py @@ -31,6 +31,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key from litellm.proxy.common_utils.user_api_key_cache import ( + NO_TEAM_MEMBERSHIP_SENTINEL, UserApiKeyCache, model_access_group_registry_cache_key, model_access_group_spend_counter_key, @@ -95,6 +96,11 @@ async def _cache( ), model_type=LiteLLM_TeamMembership, ) + else: + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=USER_ID, team_id=TEAM_ID), + value=NO_TEAM_MEMBERSHIP_SENTINEL, + ) if org_models: await cache.async_set_cache( key=f"org_id:{ORG_ID}", @@ -314,9 +320,7 @@ class _RecordingPrismaClient: def __init__(self, *rows: _MagBudgetRow) -> None: self.rows = {row.access_group_name: row for row in rows} self.batches: list[list[str]] = [] - self.db = SimpleNamespace( - litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many) - ) + self.db = SimpleNamespace(litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many)) async def _find_many(self, **kwargs): requested = list(kwargs["where"]["access_group_name"]["in"]) @@ -345,7 +349,9 @@ async def _enforce( read, seen = _spend_reader(spend_by_counter_key or {}) # The check takes its client and cache as arguments, injected just below. get_current_spend is the # one collaborator it reaches by a lazy `from litellm.proxy.proxy_server import`, with no parameter. - with patch("litellm.proxy.proxy_server.get_current_spend", read): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point + with patch( # test-quality-ok: get_current_spend is lazily imported inside the budget check + "litellm.proxy.proxy_server.get_current_spend", read + ): await _model_access_group_max_budget_check( matched_model_access_groups=matched, prisma_client=prisma_client if prisma_client is not None else _RecordingPrismaClient(*rows), @@ -491,9 +497,7 @@ async def test_a_second_request_serves_the_budget_row_from_cache(): async def test_a_database_error_does_not_block_the_request(): class _FailingPrismaClient: def __init__(self) -> None: - self.db = SimpleNamespace( - litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom) - ) + self.db = SimpleNamespace(litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom)) async def _boom(self, **kwargs): raise RuntimeError("database unavailable") @@ -507,11 +511,15 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> read, _ = _spend_reader({MODEL_ACCESS_GROUP_COUNTER_KEY: 99.0}) with ( - # common_checks resolves all three off the proxy_server module at call time; its signature - # has no client, cache or spend-reader parameter to pass them through instead. - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter - patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter - patch("litellm.proxy.proxy_server.get_current_spend", read), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point + patch( # test-quality-ok: common_checks lazily imports prisma_client from proxy_server + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), + patch( # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server + "litellm.proxy.proxy_server.user_api_key_cache", cache + ), + patch( # test-quality-ok: get_current_spend is lazily imported inside the budget check + "litellm.proxy.proxy_server.get_current_spend", read + ), ): return await common_checks( request_body={"model": "gpt-4o", "messages": []}, @@ -524,7 +532,9 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> llm_router=Router(model_list=MODEL_LIST), proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), valid_token=UserAPIKeyAuth(api_key="hashed", models=["tier-a"], user_id=USER_ID), - request=SimpleNamespace(method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions")), + request=SimpleNamespace( + method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions") + ), skip_budget_checks=skip_budget_checks, ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index ccb66fd9534..9a4badab8a9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -14366,3 +14366,123 @@ async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids assert exc_info.value.status_code == 400 assert expected_error in str(exc_info.value.detail) mock_db_client.db.query_raw.assert_not_called() + + +class _TeamRowWithOrganization(LiteLLM_TeamTable): + litellm_organization_table: LiteLLM_OrganizationTable | None = None + + +@pytest.mark.parametrize( + "organization, expected_models", + [ + ( + LiteLLM_OrganizationTable( + organization_id="org-1", + budget_id="budget-1", + models=["all-proxy-models"], + created_by="admin", + updated_by="admin", + ), + ["all-proxy-models"], + ), + ( + LiteLLM_OrganizationTable( + organization_id="org-1", + budget_id="budget-1", + models=["gpt-4o"], + created_by="admin", + updated_by="admin", + ), + ["gpt-4o"], + ), + (None, None), + ], +) +@pytest.mark.asyncio +async def test_team_info_returns_parent_organization_models(organization, expected_models): + """/team/info must report the parent org's model ceiling. + + A team admin who is not an org admin gets a 403 from /organization/info, so this + is the only read that can tell the Admin UI whether the org allows all proxy + models. Without it the team edit form hides the "All Proxy Models" option and a + team admin cannot grant their team everything on the proxy. + """ + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + team_row = _TeamRowWithOrganization( + team_id="team-1", + organization_id="org-1" if organization is not None else None, + litellm_organization_table=organization, + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.get_data = AsyncMock(return_value=[]) + + memberships = AsyncMock(return_value=[]) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: no seam on team_info + patch.object(team_endpoints, "get_all_team_memberships", memberships), # test-quality-ok: no seam on team_info + ): + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert response["team_info"].organization_models == expected_models + + +@pytest.mark.parametrize( + "caller, expected_models", + [ + (UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.INTERNAL_USER), ["gpt-4o"]), + (UserAPIKeyAuth(user_id="member-1", user_role=LitellmUserRoles.INTERNAL_USER), None), + (UserAPIKeyAuth(team_id="team-1"), None), + ], +) +@pytest.mark.asyncio +async def test_team_info_reports_parent_organization_models_only_to_team_managers(caller, expected_models): + """Plain members and team keys can read their team, but not the org's wider allow-list.""" + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + team_row = _TeamRowWithOrganization( + team_id="team-1", + organization_id="org-1", + members_with_roles=[ + Member(user_id="admin-1", role="admin"), + Member(user_id="member-1", role="user"), + ], + litellm_organization_table=LiteLLM_OrganizationTable( + organization_id="org-1", + budget_id="budget-1", + models=["gpt-4o"], + created_by="admin", + updated_by="admin", + ), + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma.get_data = AsyncMock(return_value=[]) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: no seam on team_info + patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[])), # test-quality-ok: no seam on team_info + patch.object( # test-quality-ok: no seam on team_info + team_endpoints, "_is_user_org_admin_for_team", AsyncMock(return_value=False) + ), + ): + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=caller, + ) + + assert response["team_info"].organization_models == expected_models diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index a5e6cf285ad..e09dddfec5b 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -33,6 +33,7 @@ from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash from litellm.proxy.proxy_server import app, initialize from litellm.utils import _invalidate_model_cost_lowercase_map @@ -9999,6 +10000,7 @@ async def _lit6973_drive_realtime_session( reservation: dict, *, backend_logged_success: bool, + backend_logged_failure: bool = False, phase_one_exit: str | None = None, websocket: MagicMock | None = None, ) -> MagicMock: @@ -10006,8 +10008,9 @@ async def _lit6973_drive_realtime_session( phase_one_exit picks a rejection before the relay: "model_access" makes the key/model check raise ProxyException, "pre_call" makes pre-call processing - (rate limits, guardrails) raise. Neither reaches route_request, so no success - log can own the reservation and the endpoint has to release it on that exit. + (rate limits, guardrails) raise, "pre_call_cancelled" cancels the task inside + pre-call processing. None reaches route_request, so no success log can own the + reservation and the endpoint has to release it on that exit. route_request resolves normally in both cases: the relay owns the session once route_request returns. A successful session enqueues its success cost @@ -10017,7 +10020,7 @@ async def _lit6973_drive_realtime_session( logging object carries a real model_call_details dict so the stamp is observable, and the reservation has empty entries so the real release touches no counter store.""" - from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY + from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.proxy import proxy_server as ps user_api_key_dict: Final = UserAPIKeyAuth(api_key="sk-test", token="hashed-token") @@ -10029,6 +10032,8 @@ async def _lit6973_drive_realtime_session( async def fake_llm_call() -> None: if backend_logged_success: logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True + if backend_logged_failure: + logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True from litellm.proxy._types import ProxyException @@ -10037,7 +10042,13 @@ async def _lit6973_drive_realtime_session( if phase_one_exit == "model_access" else None ) - pre_call_error: Final = Exception("Rate limit exceeded") if phase_one_exit == "pre_call" else None + pre_call_error: Final = ( + asyncio.CancelledError() + if phase_one_exit == "pre_call_cancelled" + else Exception("Rate limit exceeded") + if phase_one_exit == "pre_call" + else None + ) pre_call: Final = AsyncMock( side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj) ) @@ -10155,6 +10166,114 @@ async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_c assert reservation["finalized"] is False +_LIT6463_COUNTER_KEY: Final = "{api_key:hashed-token}:max_parallel_requests" + + +async def _lit6463_drive_realtime_session_holding_a_max_parallel_slot( + *, + backend_logged_success: bool, + backend_logged_failure: bool = False, + phase_one_exit: str | None = None, +) -> tuple[DualCache, RequestRateLimiterStash]: + """Run the realtime endpoint with a real v3 limiter registered and the request's + stash already holding slot-1 of a two-slot counter, the state pre-call leaves + behind. Returns the limiter's cache and the stash so the test can read what the + endpoint did to the slot.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + _request_stash, + ) + from litellm.proxy.utils import InternalUsageCache + + dual_cache: Final = DualCache() + await dual_cache.async_set_cache( + key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True + ) + limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache)) + stash: Final = RequestRateLimiterStash( + parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]} + ) + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + stash_token: Final = _request_stash.set(stash) + try: + hooks: Final = patch.dict( # test-quality-ok: registers the real limiter the route's release reads + ps.proxy_logging_obj.proxy_hook_mapping, {"parallel_request_limiter": limiter} + ) + expected_exit: Final = ( + pytest.raises(asyncio.CancelledError) + if phase_one_exit == "pre_call_cancelled" + else contextlib.nullcontext() + ) + with hooks, expected_exit: + await _lit6973_drive_realtime_session( + reservation, + backend_logged_success=backend_logged_success, + backend_logged_failure=backend_logged_failure, + phase_one_exit=phase_one_exit, + ) + finally: + _request_stash.reset(stash_token) + return dual_cache, stash + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase_one_exit", [None, "pre_call", "pre_call_cancelled"]) +async def test_realtime_session_ending_without_llm_callbacks_releases_the_max_parallel_slot( + phase_one_exit: str | None, +): + """The rate limiter acquires the key's max_parallel_requests slot in pre-call and + only frees it from the LLM success/failure callbacks. A realtime session that ends + without either callback (Bedrock closes without usage events, a later pre-call hook + rejects the session, or the task is cancelled while still in pre-call) has to be + released by the route itself, or the slot stays occupied until its TTL and the key's + next session is refused with a 429.""" + dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot( + backend_logged_success=False, phase_one_exit=phase_one_exit + ) + + assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == {"slot-2": 2.0} + assert stash.parallel_slot is None + + +@pytest.mark.asyncio +async def test_successful_realtime_session_leaves_the_max_parallel_slot_for_the_limiter_callback(): + """A session that enqueued its success callback hands the slot to the limiter's + own success handler, which runs on the logging worker. If the route also released + it, the two releases would race on the same stashed acquisition and, under the + limiter's integer in-memory fallback, double-decrement the counter so the key + admits more sessions than max_parallel_requests allows. With the success stamp + present the route leaves the slot and the stash alone.""" + dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot( + backend_logged_success=True + ) + + assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == { + "slot-1": 1.0, + "slot-2": 2.0, + } + assert stash.parallel_slot == {"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]} + + +@pytest.mark.asyncio +async def test_refused_realtime_session_leaves_the_max_parallel_slot_for_the_limiter_failure_callback(): + """An upstream refusal before any frame enqueues the failure callback instead, and + the limiter's failure handler releases the slot from the logging worker just like + the success handler does. The route sees no success stamp, so it still settles the + budget reservation, but it must leave the slot to that callback or the two releases + race on the same acquisition.""" + dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot( + backend_logged_success=False, backend_logged_failure=True + ) + + assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == { + "slot-1": 1.0, + "slot-2": 2.0, + } + assert stash.parallel_slot == {"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]} + + @pytest.mark.asyncio async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): """If releasing the reservation itself fails (e.g. the counter store is down), diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/test_litellm/router_utils/test_cooldown_handlers.py index 7ee0ed3701b..6fed4be5909 100644 --- a/tests/test_litellm/router_utils/test_cooldown_handlers.py +++ b/tests/test_litellm/router_utils/test_cooldown_handlers.py @@ -437,3 +437,67 @@ class TestRoutingGroupCooldownAlternatives: ) is False ) + + +class TestTeamModelCooldownAlternatives: + def _router(self, team_deployments: int, blocked_ids: frozenset[str] = frozenset()) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": f"model_name_team-1_{i}", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": { + "id": f"team-deploy-{i}", + "team_id": "team-1", + "team_public_model_name": "team-gpt-4o-mini", + "blocked": f"team-deploy-{i}" in blocked_ids, + }, + } + for i in range(team_deployments) + ] + ) + + def test_429_on_team_deployment_with_sibling_cools_down(self): + from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment + + router = self._router(team_deployments=2) + assert ( + _should_cooldown_deployment( + litellm_router_instance=router, + deployment="team-deploy-0", + exception_status=429, + original_exception=Exception("rate limited"), + requested_model_group="team-gpt-4o-mini", + ) + is True + ) + + def test_429_on_only_team_deployment_keeps_single_deployment_exemption(self): + from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment + + router = self._router(team_deployments=1) + assert ( + _should_cooldown_deployment( + litellm_router_instance=router, + deployment="team-deploy-0", + exception_status=429, + original_exception=Exception("rate limited"), + requested_model_group="team-gpt-4o-mini", + ) + is False + ) + + def test_429_with_only_a_blocked_sibling_keeps_single_deployment_exemption(self): + from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment + + router = self._router(team_deployments=2, blocked_ids=frozenset({"team-deploy-1"})) + assert ( + _should_cooldown_deployment( + litellm_router_instance=router, + deployment="team-deploy-0", + exception_status=429, + original_exception=Exception("rate limited"), + requested_model_group="team-gpt-4o-mini", + ) + is False + ) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 3ffdeefcae3..f1b445fb1bd 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -896,6 +896,46 @@ def test_arouter_test_team_model(): assert result is not None +def test_team_model_has_alternatives(): + def team_deployment( + deployment_id: str, team_id: str, public_model_name: str, blocked: bool = False + ) -> DeploymentTypedDict: + return { + "model_name": f"model_name_{team_id}_{deployment_id}", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": { + "id": deployment_id, + "team_id": team_id, + "team_public_model_name": public_model_name, + "blocked": blocked, + }, + } + + router = litellm.Router( + model_list=[ + team_deployment("team-a-1", "team-a", "shared-model"), + team_deployment("team-a-2", "team-a", "shared-model"), + team_deployment("team-a-solo", "team-a", "solo-model"), + team_deployment("team-b-1", "team-b", "shared-model"), + team_deployment("team-c-1", "team-c", "paused-sibling-model"), + team_deployment("team-c-paused", "team-c", "paused-sibling-model", blocked=True), + { + "model_name": "plain-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": {"id": "plain-1"}, + }, + ], + ) + + assert router.team_model_has_alternatives("team-a-1") is True + assert router.team_model_has_alternatives("team-a-2") is True + assert router.team_model_has_alternatives("team-a-solo") is False + assert router.team_model_has_alternatives("team-b-1") is False + assert router.team_model_has_alternatives("team-c-1") is False + assert router.team_model_has_alternatives("plain-1") is False + assert router.team_model_has_alternatives("missing-deployment") is False + + def test_arouter_ignore_invalid_deployments(): """ Test that router.ignore_invalid_deployments is set to True diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts index eccd8a80748..56e45216ec5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts @@ -328,7 +328,8 @@ describe("useTeam", () => { }); it("should return team data when query is successful", async () => { - (teamInfoCall as any).mockResolvedValue(mockTeams[0]); + // /team/info answers with an envelope; the hook is typed as the team itself. + (teamInfoCall as any).mockResolvedValue({ team_id: "team-1", team_info: mockTeams[0], keys: [] }); const { result } = renderHook(() => useTeam("team-1"), { wrapper }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts index 14e95bcd543..05025adc5e6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts @@ -163,7 +163,8 @@ export const useTeam = (teamId?: string) => { throw new Error("Missing auth or teamId"); } - return teamInfoCall(accessToken, teamId); + const { team_info } = (await teamInfoCall(accessToken, teamId)) as { team_info: Team }; + return team_info; }, initialData: () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectDetailsPage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectDetailsPage.test.tsx index cf6576ce6e9..d22ac0d8742 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectDetailsPage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectDetailsPage.test.tsx @@ -242,13 +242,11 @@ describe("ProjectDetail", () => { it("should show team information when team data is available", () => { mockUseTeam.mockReturnValue({ data: { - team_info: { - team_id: "team-1", - team_alias: "Engineering", - models: ["gpt-4"], - spend: 50, - members_with_roles: [], - }, + team_id: "team-1", + team_alias: "Engineering", + models: ["gpt-4"], + spend: 50, + members_with_roles: [], }, isLoading: false, }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectDetailsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectDetailsPage.tsx index f94240f3c9e..2585b67c81b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectDetailsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectDetailsPage.tsx @@ -14,16 +14,6 @@ import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { EditProjectModal } from "./ProjectModals/EditProjectModal"; import { ProjectKeysSection } from "./ProjectKeysSection"; -interface TeamInfoShape { - team_id: string; - team_alias?: string; - models?: string[]; - max_budget?: number | null; - budget_duration?: string | null; - spend?: number; - members_with_roles?: { user_id: string; role: string }[]; -} - interface ProjectDetailProps { projectId: string; onBack: () => void; @@ -33,10 +23,7 @@ const utilisationTone = (percent: number) => (percent >= 90 ? "over" : percent > export function ProjectDetail({ projectId, onBack }: ProjectDetailProps) { const { data: project, isLoading } = useProjectDetails(projectId); - const { data: teamData } = useTeam(project?.team_id ?? undefined); - // teamInfoCall returns { team_id, team_info: {...}, keys, team_memberships } - const teamInfo: TeamInfoShape | undefined = ((teamData as unknown as { team_info?: TeamInfoShape })?.team_info ?? - teamData) as TeamInfoShape | undefined; + const { data: teamInfo } = useTeam(project?.team_id ?? undefined); const [isEditModalVisible, setIsEditModalVisible] = useState(false); const spend = project?.spend ?? 0; diff --git a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx index 34a21122027..93a43d5e533 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx @@ -415,6 +415,139 @@ describe("ModelSelect", () => { } }); + it("should take the org model ceiling from the team when /organization/info is not readable", async () => { + const testCases = [ + { + name: "org allows all proxy models", + organizationModels: ["all-proxy-models"], + shouldShowSentinel: true, + offered: ["gpt-4", "claude-3"], + notOffered: [] as string[], + }, + { + name: "org places no ceiling at all", + organizationModels: [], + shouldShowSentinel: true, + offered: ["gpt-4", "claude-3"], + notOffered: [] as string[], + }, + { + name: "org restricts the team to one model", + organizationModels: ["gpt-4"], + shouldShowSentinel: false, + offered: ["gpt-4"], + notOffered: ["claude-3"], + }, + ]; + + for (const testCase of testCases) { + const user = userEvent.setup(); + // A team admin gets a 403 from /organization/info, so the org query never resolves. + mockUseOrganization.mockReturnValue({ data: undefined, isLoading: false } as any); + mockUseTeam.mockReturnValue({ + data: { team_id: "team-1", organization_models: testCase.organizationModels }, + isLoading: false, + } as any); + + const { unmount } = renderWithProviders( + , + ); + + await openModelList(user); + if (testCase.shouldShowSentinel) { + expectOffered("All Proxy Models"); + } else { + expectNotOffered("All Proxy Models"); + } + expectOffered("No Default Models"); + testCase.offered.forEach(expectOffered); + testCase.notOffered.forEach(expectNotOffered); + + unmount(); + } + }); + + it("should stay in the loading state while a list-seeded team is still fetching its org ceiling", () => { + mockUseOrganization.mockReturnValue({ data: undefined, isLoading: false } as any); + mockUseTeam.mockReturnValue({ + data: { team_id: "team-1", models: [] }, + isLoading: false, + isFetching: true, + } as any); + + renderWithProviders( + , + ); + + expect(screen.queryAllByRole("combobox")).toHaveLength(0); + }); + + it("should not hold the loading state on a background refetch once the org ceiling is known", async () => { + const user = userEvent.setup(); + mockUseOrganization.mockReturnValue({ data: undefined, isLoading: false } as any); + mockUseTeam.mockReturnValue({ + data: { team_id: "team-1", organization_models: ["all-proxy-models"] }, + isLoading: false, + isFetching: true, + } as any); + + renderWithProviders( + , + ); + + await openModelList(user); + expectOffered("All Proxy Models"); + }); + + it("should offer no models for an org team when neither the team nor the org reports a ceiling", async () => { + const testCases = [ + { name: "/team/info withheld the ceiling", team: { team_id: "team-1", organization_models: null } }, + { name: "/team/info failed after the list seeded the team", team: { team_id: "team-1", models: [] } }, + ]; + + for (const testCase of testCases) { + const user = userEvent.setup(); + mockUseOrganization.mockReturnValue({ data: undefined, isLoading: false } as any); + mockUseTeam.mockReturnValue({ data: testCase.team, isLoading: false, isFetching: false } as any); + + const { unmount } = renderWithProviders( + , + ); + + await openModelList(user); + expectNotOffered("All Proxy Models"); + expectOffered("No Default Models"); + expectNotOffered("gpt-4"); + expectNotOffered("claude-3"); + + unmount(); + } + }); + it("should use custom dataTestId when provided", async () => { renderWithProviders( + organizationModels.length === 0 || organizationModels.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value); + +// useTeam seeds from the team list, which omits organization_models; /team/info is the only source of the org ceiling. +const isAwaitingOrganizationModels = (team: Team | undefined, isFetchingTeam: boolean) => + isFetchingTeam && team !== undefined && team.organization_models === undefined; + const contextFilters: Record string[]> = { user: ({ allProxyModels, userModels, options }) => { if (!userModels) return []; @@ -82,18 +89,10 @@ const contextFilters: Record { - if (selectedOrganization) { - if ( - selectedOrganization.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value) || - selectedOrganization.models.length === 0 - ) { - return allProxyModels; - } - return allProxyModels.filter((model) => selectedOrganization.models.includes(model)); - } - - return allProxyModels ?? []; + team: ({ allProxyModels, organizationID, organizationModels }) => { + if (organizationModels === undefined) return organizationID ? [] : allProxyModels; + if (isUncappedModelCeiling(organizationModels)) return allProxyModels; + return allProxyModels.filter((model) => organizationModels.includes(model)); }, organization: ({ allProxyModels }) => { @@ -108,7 +107,7 @@ const contextFilters: Record { const deduplicatedProxyModels = Array.from(new Map(allProxyModels.map((m) => [m.id, m])).values()).map( (model) => model.id, @@ -118,7 +117,13 @@ const filterModels = ( const filterFn = contextFilters[ctx.context]; if (!filterFn) return []; - return filterFn({ allProxyModels: deduplicatedProxyModels, ...extra, options: ctx.options }); + const filterArgs: FilterContextArgs = { + allProxyModels: deduplicatedProxyModels, + organizationID: ctx.organizationID, + ...extra, + options: ctx.options, + }; + return filterFn(filterArgs); }; export const ModelSelect = (props: ModelSelectProps) => { @@ -126,16 +131,17 @@ export const ModelSelect = (props: ModelSelectProps) => { const { id, teamID, organizationID, options, context, dataTestId, value = [], onChange, style } = props; const { showAllProxyModelsOverride, includeSpecialOptions } = options || {}; const { data: allProxyModels, isLoading: isLoadingAllProxyModels } = useAllProxyModels(); - const { data: team, isLoading: isLoadingTeam } = useTeam(teamID); + const { data: team, isLoading: isLoadingTeam, isFetching: isFetchingTeam } = useTeam(teamID); const { data: organization, isLoading: isLoadingOrganization } = useOrganization(organizationID); const { data: currentUser, isLoading: isCurrentUserLoading } = useCurrentUser(); const isSpecialOption = (value: string) => MODEL_SENTINEL_OPTIONS.some((sv) => sv.value === value); const hasSpecialOptionSelected = value.some(isSpecialOption); - const isLoading = isLoadingAllProxyModels || isLoadingTeam || isLoadingOrganization || isCurrentUserLoading; - const organizationHasAllProxyModels = - organization?.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value) || - organization?.models.length === 0; + const isTeamPending = isLoadingTeam || isAwaitingOrganizationModels(team, isFetchingTeam); + const isLoading = isLoadingAllProxyModels || isTeamPending || isLoadingOrganization || isCurrentUserLoading; + // The org's ceiling rides on /team/info, which a team admin may read; /organization/info 403s for them. + const organizationModels = team?.organization_models ?? organization?.models; + const organizationHasAllProxyModels = organizationModels !== undefined && isUncappedModelCeiling(organizationModels); const shouldShowAllProxyModels = showAllProxyModelsOverride || (organizationHasAllProxyModels && includeSpecialOptions) || context === "global"; @@ -159,8 +165,7 @@ export const ModelSelect = (props: ModelSelectProps) => { }; const filteredModels = filterModels(allProxyModels?.data ?? [], props, { - selectedTeam: team, - selectedOrganization: organization, + organizationModels, userModels: currentUser?.models, }); diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index a6cc940c7fd..d92e23af078 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -27,6 +27,8 @@ export interface Team { access_group_models?: string[]; access_group_mcp_server_ids?: string[]; access_group_agent_ids?: string[]; + // Parent org's model ceiling. undefined = no org / not loaded; [] or ["all-proxy-models"] = no ceiling. + organization_models?: string[] | null; } export interface KeyResponse { diff --git a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx index f5d5562ccfe..39542882798 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx @@ -1,12 +1,15 @@ import React from "react"; -import { Control } from "react-hook-form"; +import { Control, UseFormReturn } from "react-hook-form"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Textarea } from "@/components/ui/textarea"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { CircleHelp } from "lucide-react"; import { FormField } from "@/components/shared/form/FormField"; +import { toast } from "@/lib/toast"; import AgentSelector from "../agent_management/AgentSelector"; import NumericalInput from "../shared/numerical_input"; import SkillSelector from "../skills/SkillSelector"; +import { moveTagsOutOfMetadataJson } from "./keyEditFieldNormalizers"; import { AgentsAndGroups, KeyEditFormValues } from "./keyEditFormValues"; export const labelWithHint = (label: React.ReactNode, hint: string): React.ReactNode => ( @@ -85,6 +88,42 @@ export const KeyAgentAndSkillFields = ({ ); +type KeyEditForm = Pick< + UseFormReturn, + "control" | "getValues" | "setValue" +>; + +export const moveMetadataTagsToTagsField = (form: KeyEditForm): void => { + const moved = moveTagsOutOfMetadataJson(form.getValues("metadata"), form.getValues("tags")); + if (moved === null) return; + form.setValue("metadata", moved.metadata, { shouldDirty: true }); + form.setValue("tags", moved.tags, { shouldDirty: true }); + if (moved.movedTags.length > 0) { + toast.info(`Moved ${moved.movedTags.join(", ")} from metadata to the Tags field`); + } +}; + +export const KeyMetadataField = ({ form }: { form: KeyEditForm }) => ( + + {(field) => ( +