diff --git a/.gitignore b/.gitignore index c2ac5137cbe..e1045032d46 100644 --- a/.gitignore +++ b/.gitignore @@ -97,3 +97,5 @@ litellm_config.yaml .vscode/launch.json litellm/proxy/to_delete_loadtest_work/* update_model_cost_map.py +tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +litellm/proxy/_experimental/out/guardrails/index.html diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index ece2b496bf6..73065b050b7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -131,7 +131,9 @@ class ProxyExtrasDBManager: ) @staticmethod - def _resolve_all_migrations(migrations_dir: str, schema_path: str): + def _resolve_all_migrations( + migrations_dir: str, schema_path: str, mark_all_applied: bool = True + ): """ 1. Compare the current database state to schema.prisma and generate a migration for the diff. 2. Run prisma migrate deploy to apply any pending migrations. @@ -210,6 +212,8 @@ class ProxyExtrasDBManager: logger.warning("Migration diff application timed out.") # 3. Mark all migrations as applied + if not mark_all_applied: + return migration_names = ProxyExtrasDBManager._get_migration_names(migrations_dir) logger.info(f"Resolving {len(migration_names)} migrations") for migration_name in migration_names: @@ -263,6 +267,13 @@ class ProxyExtrasDBManager: logger.info(f"prisma migrate deploy stdout: {result.stdout}") logger.info("prisma migrate deploy completed") + + # Run sanity check to ensure DB matches schema + logger.info("Running post-migration sanity check...") + ProxyExtrasDBManager._resolve_all_migrations( + migrations_dir, schema_path, mark_all_applied=False + ) + logger.info("✅ Post-migration sanity check completed") return True except subprocess.CalledProcessError as e: logger.info(f"prisma db error: {e.stderr}, e: {e.stdout}") diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 683f06ee0c7..b060f22d355 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -57,22 +57,22 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) -> Tuple[Optional[Any], int]: """ Handle raw dict response items from Responses API (e.g., GPT-5 Codex format). - + Args: item: Raw dict response item with 'type' field index: Current choice index - + Returns: Tuple of (Choice object or None, updated index) """ from litellm.types.utils import Choices, Message item_type = item.get("type") - + # Ignore reasoning items for now if item_type == "reasoning": return None, index - + # Handle message items with output_text content if item_type == "message": content_list = item.get("content", []) @@ -83,13 +83,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): response_text = content_item.get("text", "") msg = Message( role=item.get("role", "assistant"), - content=response_text if response_text else "" - ) - choice = Choices( - message=msg, finish_reason="stop", index=index + content=response_text if response_text else "", ) + choice = Choices(message=msg, finish_reason="stop", index=index) return choice, index + 1 - + # Unknown or unsupported type return None, index @@ -294,8 +292,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if isinstance(item, ResponseReasoningItem): - for content in item.summary: - response_text = getattr(content, "text", "") + for summary_item in item.summary: + response_text = getattr(summary_item, "text", "") reasoning_content = response_text if response_text else "" elif isinstance(item, ResponseOutputMessage): @@ -340,7 +338,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): index += 1 elif isinstance(item, dict): # Handle raw dict responses (e.g., from GPT-5 Codex) - choice, index = self._handle_raw_dict_response_item(item=item, index=index) + choice, index = self._handle_raw_dict_response_item( + item=item, index=index + ) if choice is not None: choices.append(choice) else: diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index b10ddc9e812..6aa671a5011 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -86,8 +86,15 @@ class MCPClient: async def connect(self): """Initialize the transport and session.""" if self._session: + verbose_logger.debug( + f"MCP client already connected to {self.server_url or 'stdio'}" + ) return # Already connected + verbose_logger.info( + f"MCP client connecting to {self.server_url or 'stdio'} via {self.transport_type}" + ) + try: if self.transport_type == MCPTransport.stdio: # For stdio transport, use stdio_client with command-line parameters @@ -107,6 +114,9 @@ class MCPClient: ) self._session = await self._session_ctx.__aenter__() await self._session.initialize() + verbose_logger.info( + f"MCP client successfully connected via stdio: {self.stdio_config.get('command', '')}" + ) elif self.transport_type == MCPTransport.sse: headers = self._get_auth_headers() httpx_client_factory = self._create_httpx_client_factory() @@ -122,6 +132,9 @@ class MCPClient: ) self._session = await self._session_ctx.__aenter__() await self._session.initialize() + verbose_logger.info( + f"MCP client successfully connected via SSE to {self.server_url}" + ) else: # http headers = self._get_auth_headers() httpx_client_factory = self._create_httpx_client_factory() @@ -140,6 +153,9 @@ class MCPClient: ) self._session = await self._session_ctx.__aenter__() await self._session.initialize() + verbose_logger.info( + f"MCP client successfully connected via HTTP to {self.server_url}" + ) except ValueError as e: # Re-raise ValueError exceptions (like missing stdio_config) verbose_logger.warning(f"MCP client connection failed: {str(e)}") @@ -159,7 +175,12 @@ class MCPClient: async def disconnect(self): """Clean up session and connections.""" + verbose_logger.info( + f"MCP client disconnecting from {self.server_url or 'stdio'}" + ) + if self._task and not self._task.done(): + verbose_logger.debug("MCP client cancelling background task") self._task.cancel() try: await self._task @@ -168,16 +189,24 @@ class MCPClient: if self._session: try: + verbose_logger.debug("MCP client closing session") await self._session_ctx.__aexit__(None, None, None) # type: ignore - except Exception: + except Exception as e: + verbose_logger.debug( + f"Error closing MCP session: {type(e).__name__}: {str(e)}" + ) pass self._session = None self._session_ctx = None if self._transport_ctx: try: + verbose_logger.debug("MCP client closing transport") await self._transport_ctx.__aexit__(None, None, None) - except Exception: + except Exception as e: + verbose_logger.debug( + f"Error closing MCP transport: {type(e).__name__}: {str(e)}" + ) pass self._transport_ctx = None self._transport = None @@ -261,25 +290,55 @@ class MCPClient: async def list_tools(self) -> List[MCPTool]: """List available tools from the server.""" + verbose_logger.debug( + f"MCP client listing tools from {self.server_url or 'stdio'}" + ) + if not self._session: + verbose_logger.debug("MCP client session not found, attempting to connect") try: await self.connect() except Exception as e: - verbose_logger.warning(f"MCP client connection failed: {str(e)}") + verbose_logger.error( + f"MCP client connection failed during list_tools: {type(e).__name__}: {str(e)}" + ) return [] if self._session is None: - verbose_logger.warning("MCP client session is not initialized") + verbose_logger.error( + "MCP client session is not initialized after connection attempt" + ) return [] try: result = await self._session.list_tools() + tool_count = len(result.tools) + tool_names = [tool.name for tool in result.tools] + verbose_logger.info( + f"MCP client listed {tool_count} tools from {self.server_url or 'stdio'}: {tool_names}" + ) return result.tools except asyncio.CancelledError: + verbose_logger.warning("MCP client list_tools was cancelled") await self.disconnect() raise except Exception as e: - verbose_logger.warning(f"MCP client list_tools failed: {str(e)}") + error_type = type(e).__name__ + verbose_logger.error( + f"MCP client list_tools failed - " + f"Error Type: {error_type}, " + f"Error: {str(e)}, " + f"Server: {self.server_url or 'stdio'}, " + f"Transport: {self.transport_type}" + ) + + # Check if it's a stream/connection error + if "BrokenResourceError" in error_type or "Broken" in error_type: + verbose_logger.error( + "MCP client detected broken connection/stream during list_tools - " + "the MCP server may have crashed, disconnected, or timed out" + ) + await self.disconnect() # Return empty list instead of raising to allow graceful degradation return [] @@ -290,17 +349,28 @@ class MCPClient: """ Call an MCP Tool. """ + verbose_logger.info( + f"MCP client calling tool '{call_tool_request_params.name}' with arguments: {call_tool_request_params.arguments}" + ) + if not self._session: + verbose_logger.warning( + "MCP client session not found, attempting to connect" + ) try: await self.connect() except Exception as e: - verbose_logger.warning(f"MCP client connection failed: {str(e)}") + verbose_logger.error( + f"MCP client connection failed before tool call: {type(e).__name__}: {str(e)}" + ) return MCPCallToolResult( content=[TextContent(type="text", text=f"{str(e)}")], isError=True ) if self._session is None: - verbose_logger.warning("MCP client session is not initialized") + verbose_logger.error( + "MCP client session is not initialized after connection attempt" + ) return MCPCallToolResult( content=[ TextContent( @@ -310,22 +380,59 @@ class MCPClient: isError=True, ) + # Check session and transport state before calling tool + verbose_logger.debug( + f"MCP client state before tool call - " + f"session: {'active' if self._session else 'none'}, " + f"transport: {'active' if self._transport else 'none'}, " + f"session_ctx: {'active' if self._session_ctx else 'none'}, " + f"transport_ctx: {'active' if self._transport_ctx else 'none'}" + ) + try: + verbose_logger.debug("MCP client sending tool call to session") tool_result = await self._session.call_tool( name=call_tool_request_params.name, arguments=call_tool_request_params.arguments, ) + verbose_logger.info( + f"MCP client tool call '{call_tool_request_params.name}' completed successfully" + ) return tool_result except asyncio.CancelledError: + verbose_logger.warning("MCP client tool call was cancelled") await self.disconnect() raise except Exception as e: - verbose_logger.warning(f"MCP client call_tool failed: {str(e)}") + import traceback + + error_trace = traceback.format_exc() + verbose_logger.debug(f"MCP client tool call traceback:\n{error_trace}") + + # Log detailed error information + error_type = type(e).__name__ + verbose_logger.error( + f"MCP client call_tool failed - " + f"Error Type: {error_type}, " + f"Error: {str(e)}, " + f"Tool: {call_tool_request_params.name}, " + f"Server: {self.server_url or 'stdio'}, " + f"Transport: {self.transport_type}" + ) + + # Check if it's a stream/connection error + if "BrokenResourceError" in error_type or "Broken" in error_type: + verbose_logger.error( + "MCP client detected broken connection/stream - " + "the MCP server may have crashed, disconnected, or timed out. " + "Session and transport will be disconnected." + ) + await self.disconnect() # Return a default error result instead of raising return MCPCallToolResult( content=[ - TextContent(type="text", text=f"{str(e)}") + TextContent(type="text", text=f"{error_type}: {str(e)}") ], # Empty content for error case isError=True, ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index da1f35a608a..f313a673827 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1058,6 +1058,103 @@ class MCPServerManager: ) ) + async def _call_regular_mcp_tool( + self, + mcp_server: MCPServer, + original_tool_name: str, + arguments: Dict[str, Any], + tasks: List, + mcp_auth_header: Optional[str], + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], + oauth2_headers: Optional[Dict[str, str]], + raw_headers: Optional[Dict[str, str]], + proxy_logging_obj: Optional[ProxyLogging], + ) -> CallToolResult: + """ + Call a regular MCP tool using the MCP client. + + Args: + mcp_server: The MCP server configuration + original_tool_name: The original tool name (without prefix) + arguments: Tool arguments + tasks: List of async tasks to append to (for during hooks) + mcp_auth_header: MCP auth header (deprecated) + mcp_server_auth_headers: Optional dict of server-specific auth headers + oauth2_headers: Optional OAuth2 headers + raw_headers: Optional raw headers from the request + proxy_logging_obj: Optional ProxyLogging object for hook integration + + Returns: + CallToolResult from the MCP server + + Raises: + BlockedPiiEntityError: If PII is blocked by guardrails + GuardrailRaisedException: If guardrails block the call + HTTPException: If an HTTP error occurs + """ + # Get server-specific auth header if available + server_auth_header: Optional[Union[Dict[str, str], str]] = None + if mcp_server_auth_headers and mcp_server.alias: + server_auth_header = mcp_server_auth_headers.get(mcp_server.alias) + elif mcp_server_auth_headers and mcp_server.server_name: + server_auth_header = mcp_server_auth_headers.get(mcp_server.server_name) + + # Fall back to deprecated mcp_auth_header if no server-specific header found + if server_auth_header is None: + server_auth_header = mcp_auth_header + + # oauth2 headers + extra_headers: Optional[Dict[str, str]] = None + if mcp_server.auth_type == MCPAuth.oauth2: + extra_headers = oauth2_headers + + if mcp_server.extra_headers and raw_headers: + if extra_headers is None: + extra_headers = {} + for header in mcp_server.extra_headers: + if header in raw_headers: + extra_headers[header] = raw_headers[header] + + client = self._create_mcp_client( + server=mcp_server, + mcp_auth_header=server_auth_header, + extra_headers=extra_headers, + ) + + call_tool_params = MCPCallToolRequestParams( + name=original_tool_name, + arguments=arguments, + ) + + async def _call_tool_via_client(client, params): + async with client: + return await client.call_tool(params) + + tasks.append( + asyncio.create_task(_call_tool_via_client(client, call_tool_params)) + ) + + # IMPORTANT: Must await tasks INSIDE the context manager to keep connection alive + try: + mcp_responses = await asyncio.gather(*tasks) + except ( + BlockedPiiEntityError, + GuardrailRaisedException, + HTTPException, + ) as e: + # Re-raise guardrail exceptions to properly fail the MCP call + verbose_logger.error( + f"Guardrail blocked MCP tool call during result check: {str(e)}" + ) + raise e + + # If proxy_logging_obj is None, the tool call result is at index 0 + # If proxy_logging_obj is not None, the tool call result is at index 1 (after the during hook task) + result_index = 1 if proxy_logging_obj else 0 + result = mcp_responses[result_index] + + return cast(CallToolResult, result) + async def call_tool( self, name: str, @@ -1146,48 +1243,19 @@ class MCPServerManager: ) else: # For regular MCP servers, use the MCP client - # Get server-specific auth header if available - server_auth_header: Optional[Union[Dict[str, str], str]] = None - if mcp_server_auth_headers and mcp_server.alias: - server_auth_header = mcp_server_auth_headers.get(mcp_server.alias) - elif mcp_server_auth_headers and mcp_server.server_name: - server_auth_header = mcp_server_auth_headers.get(mcp_server.server_name) - - # Fall back to deprecated mcp_auth_header if no server-specific header found - if server_auth_header is None: - server_auth_header = mcp_auth_header - - # oauth2 headers - extra_headers: Optional[Dict[str, str]] = None - if mcp_server.auth_type == MCPAuth.oauth2: - extra_headers = oauth2_headers - - if mcp_server.extra_headers and raw_headers: - if extra_headers is None: - extra_headers = {} - for header in mcp_server.extra_headers: - if header in raw_headers: - extra_headers[header] = raw_headers[header] - - client = self._create_mcp_client( - server=mcp_server, - mcp_auth_header=server_auth_header, - extra_headers=extra_headers, - ) - - call_tool_params = MCPCallToolRequestParams( - name=original_tool_name, + return await self._call_regular_mcp_tool( + mcp_server=mcp_server, + original_tool_name=original_tool_name, arguments=arguments, + tasks=tasks, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + proxy_logging_obj=proxy_logging_obj, ) - async def _call_tool_via_client(client, params): - async with client: - return await client.call_tool(params) - - tasks.append( - asyncio.create_task(_call_tool_via_client(client, call_tool_params)) - ) - + # For OpenAPI tools, await outside the client context try: mcp_responses = await asyncio.gather(*tasks) diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index d2e68067acd..f92c2eeb3af 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -11,4 +11,4 @@ mcp_servers: token_url: https://github.com/login/oauth/access_token client_id: os.environ/GITHUB_OAUTH_CLIENT_ID client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET - scopes: ["public_repo", "user:email"] \ No newline at end of file + scopes: ["public_repo", "user:email"] diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e2d68df29d3..41b36d46c75 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -776,6 +776,7 @@ class KeyRequestBase(GenerateRequestBase): tags: Optional[List[str]] = None enforced_params: Optional[List[str]] = None allowed_routes: Optional[list] = [] + allowed_passthrough_routes: Optional[list] = None rpm_limit_type: Optional[ Literal["guaranteed_throughput", "best_effort_throughput"] ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm @@ -1281,6 +1282,7 @@ class NewTeamRequest(TeamBase): guardrails: Optional[List[str]] = None prompts: Optional[List[str]] = None object_permission: Optional[LiteLLM_ObjectPermissionBase] = None + allowed_passthrough_routes: Optional[list] = None team_member_budget: Optional[float] = ( None # allow user to set a budget for all team members ) @@ -1336,6 +1338,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): team_member_rpm_limit: Optional[int] = None team_member_tpm_limit: Optional[int] = None team_member_key_duration: Optional[str] = None + allowed_passthrough_routes: Optional[list] = None class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase): @@ -1443,7 +1446,7 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): "1234567890": ["tool_name_1", "tool_name_2"] } """ - + vector_stores: Optional[List[str]] = [] @@ -3138,6 +3141,7 @@ LiteLLM_ManagementEndpoint_MetadataFields_Premium = [ "team_member_key_duration", "prompts", "logging", + "allowed_passthrough_routes", ] diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 39f11e64bb7..12b25fd95ee 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -52,9 +52,12 @@ class RouteChecks: if len(valid_token.allowed_routes) == 0: return True - # explicit check for allowed routes - if route in valid_token.allowed_routes: - return True + # explicit check for allowed routes (exact match or prefix match) + for allowed_route in valid_token.allowed_routes: + if RouteChecks._route_matches_allowed_route( + route=route, allowed_route=allowed_route + ): + return True ## check if 'allowed_route' is a field name in LiteLLMRoutes if any( @@ -62,13 +65,13 @@ class RouteChecks: for allowed_route in valid_token.allowed_routes ): for allowed_route in valid_token.allowed_routes: - if allowed_route in LiteLLMRoutes._member_names_: + if allowed_route in LiteLLMRoutes._member_names_: if RouteChecks.check_route_access( route=route, allowed_routes=LiteLLMRoutes._member_map_[allowed_route].value, ): return True - + ################################################ # For llm_api_routes, also check registered pass-through endpoints ################################################ @@ -76,7 +79,10 @@ class RouteChecks: from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( InitPassThroughEndpointHelpers, ) - if InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route): + + if InitPassThroughEndpointHelpers.is_registered_pass_through_route( + route=route + ): return True # check if wildcard pattern is allowed @@ -111,6 +117,32 @@ class RouteChecks: return masker._mask_value(user_id) + @staticmethod + def _raise_admin_only_route_exception( + user_obj: Optional[LiteLLM_UserTable], + route: str, + ) -> None: + """ + Raise exception for routes that require proxy admin access + + Args: + user_obj (Optional[LiteLLM_UserTable]): The user object + route (str): The route being accessed + + Raises: + Exception: With user role and masked user_id information + """ + user_role = "unknown" + user_id = "unknown" + if user_obj is not None: + user_role = user_obj.user_role or "unknown" + user_id = user_obj.user_id or "unknown" + + masked_user_id = RouteChecks._mask_user_id(user_id) + raise Exception( + f"Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route={route}. Your role={user_role}. Your user_id={masked_user_id}" + ) + @staticmethod def non_proxy_admin_allowed_routes_check( user_obj: Optional[LiteLLM_UserTable], @@ -195,16 +227,27 @@ class RouteChecks: pass elif route.startswith("/v1/mcp/") or route.startswith("/mcp-rest/"): pass # authN/authZ handled by api itself - else: - user_role = "unknown" - user_id = "unknown" - if user_obj is not None: - user_role = user_obj.user_role or "unknown" - user_id = user_obj.user_id or "unknown" + elif RouteChecks.check_passthrough_route_access( + route=route, user_api_key_dict=valid_token + ): + pass + elif valid_token.allowed_routes is not None: + # check if route is in allowed_routes (exact match or prefix match) + route_allowed = False + for allowed_route in valid_token.allowed_routes: + if RouteChecks._route_matches_allowed_route( + route=route, allowed_route=allowed_route + ): + route_allowed = True + break - masked_user_id = RouteChecks._mask_user_id(user_id) - raise Exception( - f"Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route={route}. Your role={user_role}. Your user_id={masked_user_id}" + if not route_allowed: + RouteChecks._raise_admin_only_route_exception( + user_obj=user_obj, route=route + ) + else: + RouteChecks._raise_admin_only_route_exception( + user_obj=user_obj, route=route ) @staticmethod @@ -347,6 +390,32 @@ class RouteChecks: # If there's no wildcard, the pattern and route should match exactly return route == pattern + @staticmethod + def _route_matches_allowed_route(route: str, allowed_route: str) -> bool: + """ + Check if route matches the allowed_route pattern. + Supports both exact match and prefix match. + + Examples: + - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-6" -> True (exact match) + - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-6/v1/chat/completions" -> True (prefix match) + - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-600" -> False (not a valid prefix) + + Args: + route: The actual route being accessed + allowed_route: The allowed route pattern + + Returns: + bool: True if route matches (exact or prefix), False otherwise + """ + # Exact match + if route == allowed_route: + return True + # Prefix match - ensure we add "/" to prevent false matches like /fake-openai-proxy-600 + if route.startswith(allowed_route + "/"): + return True + return False + @staticmethod def check_route_access(route: str, allowed_routes: List[str]) -> bool: """ @@ -394,6 +463,44 @@ class RouteChecks: return False + @staticmethod + def check_passthrough_route_access( + route: str, user_api_key_dict: UserAPIKeyAuth + ) -> bool: + """ + Check if route is a passthrough route. + Supports both exact match and prefix match. + """ + metadata = user_api_key_dict.metadata + team_metadata = user_api_key_dict.team_metadata or {} + if metadata is None and team_metadata is None: + return False + if ( + "allowed_passthrough_routes" not in metadata + and "allowed_passthrough_routes" not in team_metadata + ): + return False + if ( + metadata.get("allowed_passthrough_routes") is None + and team_metadata.get("allowed_passthrough_routes") is None + ): + return False + + allowed_passthrough_routes = ( + metadata.get("allowed_passthrough_routes") + or team_metadata.get("allowed_passthrough_routes") + or [] + ) + + # Check if route matches any allowed passthrough route (exact or prefix match) + for allowed_route in allowed_passthrough_routes: + if RouteChecks._route_matches_allowed_route( + route=route, allowed_route=allowed_route + ): + return True + + return False + @staticmethod def _is_assistants_api_request(request: Request) -> bool: """ diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 3226ea01798..aa486433396 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -47,9 +47,9 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, ) from litellm.proxy.management_helpers.object_permission_utils import ( + _set_object_permission, attach_object_permission_to_dict, handle_update_object_permission_common, - _set_object_permission, ) from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, @@ -563,6 +563,7 @@ async def _common_key_generation_helper( # noqa: PLR0915 data_json = data.model_dump(exclude_unset=True, exclude_none=True) # type: ignore data_json = handle_key_type(data, data_json) + # if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users if "max_budget" in data_json: data_json["key_max_budget"] = data_json.pop("max_budget", None) @@ -838,6 +839,7 @@ async def generate_key_fn( - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests) - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"] + - allowed_passthrough_routes: Optional[list] - List of allowed pass through endpoints for the key. Store the actual endpoint or store a wildcard pattern for a set of endpoints. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through endpoints the key can access, without specifying the routes. If allowed_routes is specified, allowed_pass_through_endpoints is ignored. - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "mcp_tool_permissions": {"server_id_1": ["tool1", "tool2"]}}. IF null or {} then no object permission. - key_type: Optional[str] - Type of key that determines default allowed routes. Options: "llm_api" (can call LLM API routes), "management" (can call management routes), "read_only" (can only call info/read routes), "default" (uses default allowed routes). Defaults to "default". - prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts. @@ -1114,8 +1116,6 @@ def prepare_metadata_fields( return non_default_values - - async def prepare_key_update_data( data: Union[UpdateKeyRequest, RegenerateKeyRequest], existing_key_row: LiteLLM_VerificationToken, @@ -1277,6 +1277,7 @@ async def update_key_fn( - temp_budget_increase: Optional[float] - Temporary budget increase for the key (Enterprise only). - temp_budget_expiry: Optional[str] - Expiry time for the temporary budget increase (Enterprise only). - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"] + - allowed_passthrough_routes: Optional[list] - List of allowed pass through routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through routes the key can access, without specifying the routes. If allowed_routes is specified, allowed_passthrough_routes is ignored. - prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts. - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "mcp_tool_permissions": {"server_id_1": ["tool1", "tool2"]}}. IF null or {} then no object permission. - auto_rotate: Optional[bool] - Whether this key should be automatically rotated @@ -2819,6 +2820,7 @@ async def list_keys( code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) + @router.get( "/key/aliases", tags=["key management"], diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index fd7d150a7f4..b8f6b4a4460 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -4,7 +4,6 @@ This is an enterprise feature and requires a premium license. """ -from litellm._uuid import uuid from typing import Any, Dict, List, Optional, Set, Tuple from fastapi import ( @@ -17,11 +16,12 @@ from fastapi import ( Request, Response, ) -from typing_extensions import TypedDict from pydantic import BaseModel +from typing_extensions import TypedDict import litellm from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( LiteLLM_TeamTable, @@ -51,32 +51,31 @@ from litellm.types.proxy.management_endpoints.scim_v2 import * class UserProvisionerHelpers: """Helper methods for user provisioning operations.""" - + @staticmethod async def handle_existing_user_by_email( - prisma_client, - new_user_request: NewUserRequest + prisma_client, new_user_request: NewUserRequest ) -> Optional[SCIMUser]: """ Check if a user with the given email already exists and update them if found. - + Args: prisma_client: Database client new_user_request: New user request data - + Returns: SCIMUser if user was updated, None if no existing user found """ if not new_user_request.user_email: return None - + existing_user = await prisma_client.db.litellm_usertable.find_first( where={"user_email": new_user_request.user_email} ) - + if not existing_user: return None - + # Update the user updated_user = await prisma_client.db.litellm_usertable.update( where={"user_id": existing_user.user_id}, @@ -88,12 +87,15 @@ class UserProvisionerHelpers: "metadata": safe_dumps(new_user_request.metadata), }, ) - - return await ScimTransformations.transform_litellm_user_to_scim_user(updated_user) + + return await ScimTransformations.transform_litellm_user_to_scim_user( + updated_user + ) class ScimUserData(TypedDict): """Typed structure for extracted SCIM user data.""" + user_email: Optional[str] user_alias: Optional[str] sso_user_id: Optional[str] @@ -105,6 +107,7 @@ class ScimUserData(TypedDict): class GroupMemberExtractionResult(BaseModel): """Result of extracting and processing group members.""" + existing_member_ids: List[str] created_users: List[NewUserResponse] all_member_ids: List[str] # existing + newly created @@ -121,7 +124,7 @@ scim_router = APIRouter( async def _get_prisma_client_or_raise_exception(): """Check if database is connected and raise HTTPException if not.""" from litellm.proxy.proxy_server import prisma_client - + if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No database connected"}) return prisma_client @@ -130,32 +133,32 @@ async def _get_prisma_client_or_raise_exception(): async def _check_user_exists(user_id: str): """Check if user exists and return user, raise 404 if not found.""" prisma_client = await _get_prisma_client_or_raise_exception() - + user = await prisma_client.db.litellm_usertable.find_unique( where={"user_id": user_id} ) - + if not user: raise HTTPException( status_code=404, detail={"error": f"User not found with ID: {user_id}"} ) - + return user async def _check_team_exists(team_id: str): """Check if team exists and return team, raise 404 if not found.""" prisma_client = await _get_prisma_client_or_raise_exception() - + team = await prisma_client.db.litellm_teamtable.find_unique( where={"team_id": team_id} ) - + if not team: raise HTTPException( status_code=404, detail={"error": f"Group not found with ID: {team_id}"} ) - + return team @@ -184,7 +187,9 @@ def _extract_scim_user_data(user: SCIMUser) -> ScimUserData: } -def _build_scim_metadata(given_name: Optional[str], family_name: Optional[str], active: Optional[bool] = None) -> Dict[str, Any]: +def _build_scim_metadata( + given_name: Optional[str], family_name: Optional[str], active: Optional[bool] = None +) -> Dict[str, Any]: """Build metadata dictionary with SCIM data.""" metadata: Dict[str, Any] = { "scim_metadata": LiteLLM_UserScimMetadata( @@ -192,17 +197,17 @@ def _build_scim_metadata(given_name: Optional[str], family_name: Optional[str], familyName=family_name, ).model_dump() } - + if active is not None: metadata["scim_active"] = active - + return metadata async def _extract_group_member_ids(group: SCIMGroup) -> GroupMemberExtractionResult: """ Extract member IDs from SCIMGroup, creating users that don't exist. - + Returns: GroupMemberExtractionResult with existing members, created users, and all member IDs """ @@ -210,35 +215,34 @@ async def _extract_group_member_ids(group: SCIMGroup) -> GroupMemberExtractionRe existing_member_ids = [] created_users = [] all_member_ids = [] - + if group.members: for member in group.members: user_id = member.value - + # Check if user exists user = await prisma_client.db.litellm_usertable.find_unique( where={"user_id": user_id} ) - + if user: existing_member_ids.append(user_id) all_member_ids.append(user_id) else: # Create the user if they don't exist using our helper created_user = await _create_user_if_not_exists( - user_id=user_id, - created_via="scim_group_membership" + user_id=user_id, created_via="scim_group_membership" ) - + if created_user: created_users.append(created_user) all_member_ids.append(user_id) # If creation failed, user is skipped (logged in helper) - + return GroupMemberExtractionResult( existing_member_ids=existing_member_ids, created_users=created_users, - all_member_ids=all_member_ids + all_member_ids=all_member_ids, ) @@ -246,7 +250,7 @@ async def _get_team_members_display(member_ids: List[str]) -> List[SCIMMember]: """Get SCIMMember objects with display names for a list of member IDs.""" prisma_client = await _get_prisma_client_or_raise_exception() members: List[SCIMMember] = [] - + for member_id in member_ids: user = await prisma_client.db.litellm_usertable.find_unique( where={"user_id": member_id} @@ -254,18 +258,20 @@ async def _get_team_members_display(member_ids: List[str]) -> List[SCIMMember]: if user: display_name = user.user_email or user.user_id members.append(SCIMMember(value=user.user_id, display=display_name)) - + return members -async def _handle_team_membership_changes(user_id: str, existing_teams: List[str], new_teams: List[str]) -> None: +async def _handle_team_membership_changes( + user_id: str, existing_teams: List[str], new_teams: List[str] +) -> None: """Handle adding/removing user from teams based on changes.""" existing_teams_set = set(existing_teams) new_teams_set = set(new_teams) - + teams_to_add = new_teams_set - existing_teams_set teams_to_remove = existing_teams_set - new_teams_set - + if teams_to_add or teams_to_remove: await patch_team_membership( user_id=user_id, @@ -274,19 +280,21 @@ async def _handle_team_membership_changes(user_id: str, existing_teams: List[str ) -async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_group") -> Optional[NewUserResponse]: +async def _create_user_if_not_exists( + user_id: str, created_via: str = "scim_group" +) -> Optional[NewUserResponse]: """ Helper function to create a user if they don't exist. - + Args: user_id: The user ID to create created_via: Context for where the user was created from - + Returns: LiteLLM_UserTable if user was created, None if creation failed """ from litellm.proxy.management_endpoints.internal_user_endpoints import new_user - + try: # Get default role for new internal users default_role: Optional[ @@ -313,7 +321,7 @@ async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_grou created_user = await new_user(data=new_user_request) verbose_proxy_logger.info(f"Created user {user_id} via {created_via}") return created_user - + except Exception as e: verbose_proxy_logger.exception(f"Failed to create user {user_id}: {e}") return None @@ -324,7 +332,7 @@ async def _get_team_member_user_ids_from_team(team: LiteLLM_TeamTable) -> List[s Get the IDs of the members from a team. Use one source of truth for the member IDs: team.members_with_roles - + """ member_user_ids: List[str] = [] for member in team.members_with_roles or []: @@ -337,7 +345,6 @@ async def _get_team_member_user_ids_from_team(team: LiteLLM_TeamTable) -> List[s return member_user_ids - # Dependency to set the correct SCIM Content-Type async def set_scim_content_type(response: Response): """Sets the Content-Type header to application/scim+json""" @@ -450,7 +457,7 @@ async def get_user( verbose_proxy_logger.debug("SCIM GET USER request for user_id=%s", user_id) try: user = await _check_user_exists(user_id) - + # Convert to SCIM format scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user) return scim_user @@ -458,6 +465,7 @@ async def get_user( except Exception as e: raise handle_exception_on_proxy(e) + @scim_router.post( "/Users", response_model=SCIMUser, @@ -471,11 +479,9 @@ async def create_user( Create a user according to SCIM v2 protocol """ try: - verbose_proxy_logger.debug( - "SCIM CREATE USER request: %s", user.model_dump() - ) + verbose_proxy_logger.debug("SCIM CREATE USER request: %s", user.model_dump()) prisma_client = await _get_prisma_client_or_raise_exception() - + # Extract data from SCIM user user_data = _extract_scim_user_data(user) @@ -487,20 +493,24 @@ async def create_user( if existing_user: raise HTTPException( status_code=409, - detail={"error": f"User already exists with username: {user.userName}"}, + detail={ + "error": f"User already exists with username: {user.userName}" + }, ) # Create user in database user_id = user.userName or str(uuid.uuid4()) - metadata = _build_scim_metadata(user_data["given_name"], user_data["family_name"]) + metadata = _build_scim_metadata( + user_data["given_name"], user_data["family_name"] + ) default_role: Optional[ Literal[ - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - ] + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + ] ] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY if litellm.default_internal_user_params: default_role = litellm.default_internal_user_params.get("user_role") @@ -517,22 +527,23 @@ async def create_user( # Check if user with email already exists and update if found existing_user_scim = await UserProvisionerHelpers.handle_existing_user_by_email( - prisma_client=prisma_client, - new_user_request=new_user_request + prisma_client=prisma_client, new_user_request=new_user_request ) - + if existing_user_scim: return existing_user_scim created_user = await new_user( data=new_user_request, ) - + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( user=created_user ) return scim_user - except HTTPException as e: # allow exceptions like SCIMUserAlreadyExists to be raised + except ( + HTTPException + ) as e: # allow exceptions like SCIMUserAlreadyExists to be raised raise e except Exception as e: raise handle_exception_on_proxy(e) @@ -564,18 +575,16 @@ async def update_user( # Extract data from SCIM user user_data = _extract_scim_user_data(user) - # Build metadata with SCIM data + # Build metadata with SCIM data metadata = _build_scim_metadata( - user_data["given_name"], - user_data["family_name"], - user_data["active"] + user_data["given_name"], user_data["family_name"], user_data["active"] ) # Handle team membership changes await _handle_team_membership_changes( user_id=user_id, existing_teams=existing_user.teams or [], - new_teams=user_data["teams"] + new_teams=user_data["teams"], ) # Update user with all new data (full replacement) @@ -590,6 +599,7 @@ async def update_user( # Serialize metadata to JSON string for Prisma to avoid GraphQL parsing issues if "metadata" in update_data and isinstance(update_data["metadata"], dict): from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + update_data["metadata"] = safe_dumps(update_data["metadata"]) updated_user = await prisma_client.db.litellm_usertable.update( @@ -598,8 +608,10 @@ async def update_user( ) # Convert back to SCIM format - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(updated_user) - + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( + updated_user + ) + return scim_user except Exception as e: @@ -617,9 +629,7 @@ async def delete_user( """ Delete a user according to SCIM v2 protocol """ - verbose_proxy_logger.debug( - "SCIM DELETE USER request for user_id=%s", user_id - ) + verbose_proxy_logger.debug("SCIM DELETE USER request for user_id=%s", user_id) try: prisma_client = await _get_prisma_client_or_raise_exception() existing_user = await _check_user_exists(user_id) @@ -668,7 +678,9 @@ def _extract_group_values(value: Any) -> List[str]: return group_values -def _handle_displayname_update(op_type: str, value: Any, update_data: Dict[str, Any]) -> None: +def _handle_displayname_update( + op_type: str, value: Any, update_data: Dict[str, Any] +) -> None: """Handle displayname updates.""" if op_type == "remove": update_data["user_alias"] = None @@ -676,7 +688,9 @@ def _handle_displayname_update(op_type: str, value: Any, update_data: Dict[str, update_data["user_alias"] = str(value) -def _handle_externalid_update(op_type: str, value: Any, update_data: Dict[str, Any]) -> None: +def _handle_externalid_update( + op_type: str, value: Any, update_data: Dict[str, Any] +) -> None: """Handle externalid updates.""" if op_type == "remove": update_data["sso_user_id"] = None @@ -697,7 +711,9 @@ def _handle_active_update(op_type: str, value: Any, metadata: Dict[str, Any]) -> metadata["scim_active"] = bool_val -def _handle_name_update(path: str, op_type: str, value: Any, scim_metadata: Dict[str, Any]) -> None: +def _handle_name_update( + path: str, op_type: str, value: Any, scim_metadata: Dict[str, Any] +) -> None: """Handle name field updates (givenName, familyName).""" if path == "name.givenname": if op_type == "remove": @@ -711,7 +727,9 @@ def _handle_name_update(path: str, op_type: str, value: Any, scim_metadata: Dict scim_metadata["familyName"] = str(value) -def _handle_group_operations(op_type: str, value: Any, teams_set: Set[str]) -> Optional[Set[str]]: +def _handle_group_operations( + op_type: str, value: Any, teams_set: Set[str] +) -> Optional[Set[str]]: """Handle group/team membership operations.""" group_values = _extract_group_values(value) if op_type == "replace": @@ -724,7 +742,9 @@ def _handle_group_operations(op_type: str, value: Any, teams_set: Set[str]) -> O return None -def _handle_generic_metadata(path: str, op_type: str, value: Any, metadata: Dict[str, Any]) -> None: +def _handle_generic_metadata( + path: str, op_type: str, value: Any, metadata: Dict[str, Any] +) -> None: """Handle generic metadata operations for unknown paths.""" if op_type == "remove": metadata.pop(path, None) @@ -769,6 +789,7 @@ def _apply_patch_ops( update_data["metadata"] = metadata return update_data, final_team_set + async def patch_team_membership( user_id: str, teams_ids_to_add_user_to: List[str], @@ -778,29 +799,35 @@ async def patch_team_membership( Add or remove user from teams """ for _team_id in teams_ids_to_add_user_to: - try: - await team_member_add( - data=TeamMemberAddRequest( - team_id=_team_id, - member=Member(user_id=user_id, role="user"), - ), - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), - ) - except Exception as e: - verbose_proxy_logger.exception(f"Error adding user to team {_team_id}: {e}") + try: + await team_member_add( + data=TeamMemberAddRequest( + team_id=_team_id, + member=Member(user_id=user_id, role="user"), + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ), + ) + except Exception as e: + verbose_proxy_logger.exception(f"Error adding user to team {_team_id}: {e}") for _team_id in teams_ids_to_remove_user_from: try: await team_member_delete( data=TeamMemberDeleteRequest(team_id=_team_id, user_id=user_id), - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ), ) except Exception as e: - verbose_proxy_logger.exception(f"Error removing user from team {_team_id}: {e}") - + verbose_proxy_logger.exception( + f"Error removing user from team {_team_id}: {e}" + ) return True + @scim_router.patch( "/Users/{user_id}", response_model=SCIMUser, @@ -833,7 +860,7 @@ async def patch_user( await _handle_team_membership_changes( user_id=user_id, existing_teams=existing_user.teams or [], - new_teams=list(final_team_set) + new_teams=list(final_team_set), ) update_data["teams"] = list(final_team_set) @@ -841,6 +868,7 @@ async def patch_user( # Serialize metadata to JSON string for Prisma to avoid GraphQL parsing issues if "metadata" in update_data and isinstance(update_data["metadata"], dict): from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + update_data["metadata"] = safe_dumps(update_data["metadata"]) updated_user = await prisma_client.db.litellm_usertable.update( @@ -848,7 +876,9 @@ async def patch_user( data=update_data, ) - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(updated_user) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( + updated_user + ) return scim_user @@ -947,9 +977,7 @@ async def get_group( """ Get a single group by ID according to SCIM v2 protocol """ - verbose_proxy_logger.debug( - "SCIM GET GROUP request for group_id=%s", group_id - ) + verbose_proxy_logger.debug("SCIM GET GROUP request for group_id=%s", group_id) try: team = await _check_team_exists(group_id) @@ -981,9 +1009,9 @@ async def create_group( ) try: prisma_client = await _get_prisma_client_or_raise_exception() - + # Generate ID if not provided - team_id = group.id or str(uuid.uuid4()) + team_id = group.id or group.externalId or str(uuid.uuid4()) # Check if team already exists existing_team = await prisma_client.db.litellm_teamtable.find_unique( @@ -998,7 +1026,10 @@ async def create_group( # Extract and process group members (creating users that don't exist) member_result = await _extract_group_member_ids(group) - members_with_roles = [Member(user_id=member_id, role="user") for member_id in member_result.all_member_ids] + members_with_roles = [ + Member(user_id=member_id, role="user") + for member_id in member_result.all_member_ids + ] # Create team in database created_team = await new_team( @@ -1043,13 +1074,17 @@ async def update_group( # Extract and process group members (creating users that don't exist) member_result = await _extract_group_member_ids(group) - verbose_proxy_logger.debug(f"SCIM PUT GROUP all_member_ids: {member_result.all_member_ids}") - verbose_proxy_logger.debug(f"SCIM PUT GROUP created_users: {len(member_result.created_users)}") + verbose_proxy_logger.debug( + f"SCIM PUT GROUP all_member_ids: {member_result.all_member_ids}" + ) + verbose_proxy_logger.debug( + f"SCIM PUT GROUP created_users: {len(member_result.created_users)}" + ) # Prepare update data existing_metadata = existing_team.metadata if existing_team.metadata else {} updated_metadata = {**existing_metadata, "scim_data": group.model_dump()} - + update_data = { "team_alias": group.displayName, "metadata": safe_dumps(updated_metadata), @@ -1066,7 +1101,7 @@ async def update_group( verbose_proxy_logger.debug(f"SCIM PUT GROUP current_members: {current_members}") final_members = set(member_result.all_member_ids) verbose_proxy_logger.debug(f"SCIM PUT GROUP final_members: {final_members}") - + await _handle_group_membership_changes( group_id=group_id, current_members=current_members, @@ -1094,9 +1129,7 @@ async def delete_group( """ Delete a group according to SCIM v2 protocol """ - verbose_proxy_logger.debug( - "SCIM DELETE GROUP request for group_id=%s", group_id - ) + verbose_proxy_logger.debug("SCIM DELETE GROUP request for group_id=%s", group_id) try: prisma_client = await _get_prisma_client_or_raise_exception() existing_team = await _check_team_exists(group_id) @@ -1124,21 +1157,19 @@ async def delete_group( async def _process_group_patch_operations( - patch_ops: SCIMPatchOp, - existing_team, - prisma_client + patch_ops: SCIMPatchOp, existing_team, prisma_client ) -> Tuple[Dict[str, Any], Set[str]]: """Process patch operations for a group and return update data and final members.""" update_data: Dict[str, Any] = {} - + # Create a fresh copy of existing metadata to avoid Prisma issues existing_metadata = existing_team.metadata or {} metadata = dict(existing_metadata) if existing_metadata else {} - + # Track member changes current_members = set(existing_team.members or []) final_members = current_members.copy() - + # Process each patch operation for op in patch_ops.Operations: path = (op.path or "").lower() @@ -1169,14 +1200,13 @@ async def _process_group_patch_operations( else: # Create the user if they don't exist using our helper created_user = await _create_user_if_not_exists( - user_id=member_id, - created_via="scim_group_patch" + user_id=member_id, created_via="scim_group_patch" ) - + if created_user: valid_members.append(member_id) # If creation failed, user is skipped (logged in helper) - + if op_type == "replace": final_members = set(valid_members) elif op_type == "add": @@ -1194,21 +1224,18 @@ async def _process_group_patch_operations( # Include metadata in update data if it exists if metadata: update_data["metadata"] = metadata - + return update_data, final_members async def _apply_group_patch_updates( - group_id: str, - update_data: Dict[str, Any], - final_members: Set[str], - prisma_client + group_id: str, update_data: Dict[str, Any], final_members: Set[str], prisma_client ): """Apply patch updates to the group in the database.""" # Serialize metadata if present if "metadata" in update_data and isinstance(update_data["metadata"], dict): update_data["metadata"] = safe_dumps(update_data["metadata"]) - + # Update members list update_data["members"] = list(final_members) @@ -1217,22 +1244,20 @@ async def _apply_group_patch_updates( where={"team_id": group_id}, data=update_data, ) - + return updated_team async def _handle_group_membership_changes( - group_id: str, - current_members: Set[str], - final_members: Set[str] + group_id: str, current_members: Set[str], final_members: Set[str] ): """Handle adding/removing members from the group.""" members_to_add = final_members - current_members members_to_remove = current_members - final_members - + verbose_proxy_logger.debug(f"members_to_add: {members_to_add}") verbose_proxy_logger.debug(f"members_to_remove: {members_to_remove}") - + # Use existing helper functions for team membership changes for member_id in members_to_add: await patch_team_membership( @@ -1276,7 +1301,7 @@ async def patch_group( update_data, final_members = await _process_group_patch_operations( patch_ops, existing_team, prisma_client ) - + # Track current members for comparison current_members = set(await _get_team_member_user_ids_from_team(existing_team)) @@ -1286,9 +1311,7 @@ async def patch_group( ) # Handle user-team relationship changes - await _handle_group_membership_changes( - group_id, current_members, final_members - ) + await _handle_group_membership_changes(group_id, current_members, final_members) # Convert to SCIM format and return scim_group = await ScimTransformations.transform_litellm_team_to_scim_group( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index d05a0259f46..f2d2cef3d6f 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -27,6 +27,7 @@ from litellm.proxy._types import ( CommonProxyErrors, DeleteTeamRequest, LiteLLM_AuditLogs, + LiteLLM_ManagementEndpoint_MetadataFields, LiteLLM_ManagementEndpoint_MetadataFields_Premium, LiteLLM_ModelTable, LiteLLM_OrganizationTable, @@ -56,9 +57,6 @@ from litellm.proxy._types import ( UpdateTeamRequest, UserAPIKeyAuth, ) -from litellm.proxy.management_helpers.object_permission_utils import ( - _set_object_permission, -) from litellm.proxy.auth.auth_checks import ( allowed_route_check_inside_route, can_org_access_model, @@ -76,6 +74,7 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import ( get_daily_activity, ) from litellm.proxy.management_helpers.object_permission_utils import ( + _set_object_permission, handle_update_object_permission_common, ) from litellm.proxy.management_helpers.team_member_permission_checks import ( @@ -320,7 +319,9 @@ async def new_team( # noqa: PLR0915 - team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members. - team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo" - prompts: Optional[List[str]] - List of allowed prompts for the team. If specified, the team will only be able to use these specific prompts. + - allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team. + Returns: - team_id: (str) Unique team id - used for tracking spend across multiple keys for same team id. @@ -478,7 +479,7 @@ async def new_team( # noqa: PLR0915 ## Create Team Member Budget Table data_json = data.json() - + ## Handle Object Permission - MCP, Vector Stores etc. data_json = await _set_object_permission( data_json=data_json, @@ -514,6 +515,14 @@ async def new_team( # noqa: PLR0915 value=getattr(data, field), ) + for field in LiteLLM_ManagementEndpoint_MetadataFields: + if getattr(data, field, None) is not None: + _set_object_metadata_field( + object_data=complete_team_data, + field_name=field, + value=getattr(data, field), + ) + # If budget_duration is set, set `budget_reset_at` if complete_team_data.budget_duration is not None: from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time @@ -619,6 +628,53 @@ async def _update_model_table( return _model_id +async def fetch_and_validate_organization( + organization_id: str, + existing_team_row: Any, + llm_router: Optional[Router], + prisma_client: Any, +) -> Any: + """ + Fetch and validate an organization for team update operations. + + Args: + organization_id: The organization ID to fetch + existing_team_row: The existing team row being updated + llm_router: The LLM router instance + prisma_client: The Prisma database client + + Returns: + The organization row from the database + + Raises: + HTTPException: If llm_router is None, organization not found, or validation fails + """ + if llm_router is None: + raise HTTPException( + status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} + ) + + organization_row = await prisma_client.db.litellm_organizationtable.find_unique( + where={"organization_id": organization_id}, + include={"litellm_budget_table": True, "users": True}, + ) + + if organization_row is None: + raise HTTPException( + status_code=404, + detail={ + "error": f"Organization not found, passed organization_id={organization_id}" + }, + ) + + validate_team_org_change( + team=LiteLLM_TeamTable(**existing_team_row.model_dump()), + organization=LiteLLM_OrganizationTable(**organization_row.model_dump()), + llm_router=llm_router, + ) + + return organization_row + def validate_team_org_change( team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTable, llm_router: Router @@ -754,6 +810,7 @@ async def update_team( - team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members. - team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members. - team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo" + - allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team. Example - update team TPM Limit ``` @@ -810,25 +867,11 @@ async def update_team( if ( data.organization_id is not None and len(data.organization_id) > 0 ): # allow unsetting the organization_id - if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} - ) - organization_row = await prisma_client.db.litellm_organizationtable.find_unique( - where={"organization_id": data.organization_id}, - include={"litellm_budget_table": True, "users": True}, - ) - if organization_row is None: - raise HTTPException( - status_code=404, - detail={ - "error": f"Organization not found, passed organization_id={data.organization_id}" - }, - ) - validate_team_org_change( - team=LiteLLM_TeamTable(**existing_team_row.model_dump()), - organization=LiteLLM_OrganizationTable(**organization_row.model_dump()), + await fetch_and_validate_organization( + organization_id=data.organization_id, + existing_team_row=existing_team_row, llm_router=llm_router, + prisma_client=prisma_client, ) elif data.organization_id is not None and len(data.organization_id) == 0: # unsetting the organization_id @@ -877,6 +920,13 @@ async def update_team( field_name=field, ) + for field in LiteLLM_ManagementEndpoint_MetadataFields: + if field in updated_kv and updated_kv[field] is not None: + _update_team_metadata_field( + updated_kv=updated_kv, + field_name=field, + ) + if "model_aliases" in updated_kv: updated_kv.pop("model_aliases") _model_id = await _update_model_table( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 227ea4db9ff..478966d61b3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -5,7 +5,7 @@ import json import traceback from base64 import b64encode from datetime import datetime -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union, cast from urllib.parse import urlencode, urlparse import httpx @@ -42,6 +42,7 @@ from litellm.passthrough import BasePassthroughUtils from litellm.proxy._types import ( ConfigFieldInfo, ConfigFieldUpdate, + LiteLLMRoutes, PassThroughEndpointResponse, PassThroughGenericEndpoint, ProxyException, @@ -66,7 +67,7 @@ router = APIRouter() pass_through_endpoint_logging = PassThroughEndpointLogging() # Global registry to track registered pass-through routes and prevent memory leaks -_registered_pass_through_routes: Dict[str, Dict[str, str]] = {} +_registered_pass_through_routes: Dict[str, Dict[str, Union[str, Dict[str, Any]]]] = {} def get_response_body(response: httpx.Response) -> Optional[dict]: @@ -974,24 +975,74 @@ def create_pass_through_route( ] = None, # if pass-through endpoint is a streaming request subpath: str = "", # captures sub-paths when include_subpath=True ): + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + + path = request.url.path + + if not InitPassThroughEndpointHelpers.is_registered_pass_through_route( + route=path + ): + raise HTTPException( + status_code=404, + detail=f"Pass-through endpoint {endpoint} not found. This could have been deleted or not yet added to the proxy.", + ) + + passthrough_params = ( + InitPassThroughEndpointHelpers.get_registered_pass_through_route( + route=path + ) + ) + target_params = { + "target": target, + "custom_headers": custom_headers, + "forward_headers": _forward_headers, + "merge_query_params": _merge_query_params, + "cost_per_request": cost_per_request, + } + + if passthrough_params is not None: + target_params.update(passthrough_params.get("passthrough_params", {})) + + # Extract and cast parameters with proper types + param_target = target_params.get("target") or target + param_custom_headers = target_params.get("custom_headers", custom_headers) + param_forward_headers = target_params.get( + "forward_headers", _forward_headers + ) + param_merge_query_params = target_params.get( + "merge_query_params", _merge_query_params + ) + param_cost_per_request = target_params.get( + "cost_per_request", cost_per_request + ) + # Construct the full target URL with subpath if needed full_target = ( HttpPassThroughEndpointHelpers.construct_target_url_with_subpath( - base_target=target, subpath=subpath, include_subpath=include_subpath + base_target=cast(str, param_target), + subpath=subpath, + include_subpath=include_subpath, ) ) + # Ensure custom_headers is a dict + headers_dict = ( + param_custom_headers if isinstance(param_custom_headers, dict) else {} + ) + return await pass_through_request( # type: ignore request=request, target=full_target, - custom_headers=custom_headers or {}, + custom_headers=headers_dict, user_api_key_dict=user_api_key_dict, - forward_headers=_forward_headers, - merge_query_params=_merge_query_params, + forward_headers=cast(Optional[bool], param_forward_headers), + merge_query_params=cast(Optional[bool], param_merge_query_params), query_params=query_params, stream=stream, custom_body=custom_body, - cost_per_request=cost_per_request, + cost_per_request=cast(Optional[float], param_cost_per_request), custom_llm_provider=custom_llm_provider, ) @@ -1592,6 +1643,14 @@ class InitPassThroughEndpointHelpers: "endpoint_id": endpoint_id, "path": path, "type": "exact", + "passthrough_params": { + "target": target, + "custom_headers": custom_headers, + "forward_headers": forward_headers, + "merge_query_params": merge_query_params, + "dependencies": dependencies, + "cost_per_request": cost_per_request, + }, } @staticmethod @@ -1645,6 +1704,14 @@ class InitPassThroughEndpointHelpers: "endpoint_id": endpoint_id, "path": path, "type": "subpath", + "passthrough_params": { + "target": target, + "custom_headers": custom_headers, + "forward_headers": forward_headers, + "merge_query_params": merge_query_params, + "dependencies": dependencies, + "cost_per_request": cost_per_request, + }, } @staticmethod @@ -1661,6 +1728,11 @@ class InitPassThroughEndpointHelpers: "Removed pass-through route from registry: %s", key ) + @staticmethod + def clear_all_pass_through_routes(): + """Clear all pass-through routes from the registry""" + _registered_pass_through_routes.clear() + @staticmethod def is_registered_pass_through_route(route: str) -> bool: """ @@ -1675,6 +1747,12 @@ class InitPassThroughEndpointHelpers: Returns: bool: True if route is a registered pass-through endpoint, False otherwise """ + + ## CHECK IF MAPPED PASS THROUGH ENDPOINT + for mapped_route in LiteLLMRoutes.mapped_pass_through_routes.value: + if route.startswith(mapped_route): + return True + # Fast path: check if any registered route key contains this path # Keys are in format: "{endpoint_id}:exact:{path}" or "{endpoint_id}:subpath:{path}" # Extract unique paths from keys for quick checking @@ -1683,7 +1761,6 @@ class InitPassThroughEndpointHelpers: if len(parts) == 3: route_type = parts[1] registered_path = parts[2] - if route_type == "exact" and route == registered_path: return True elif route_type == "subpath": @@ -1694,11 +1771,42 @@ class InitPassThroughEndpointHelpers: return False + @staticmethod + def get_registered_pass_through_route(route: str) -> Optional[Dict[str, Any]]: + """Get passthrough params for a given route""" + for key in _registered_pass_through_routes.keys(): + parts = key.split(":", 2) # Split into [endpoint_id, type, path] + if len(parts) == 3: + route_type = parts[1] + registered_path = parts[2] + + if route_type == "exact" and route == registered_path: + return _registered_pass_through_routes[key] + elif route_type == "subpath": + if route == registered_path or route.startswith( + registered_path + "/" + ): + return _registered_pass_through_routes[key] + + return None + + +def _get_combined_pass_through_endpoints( + pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]], + config_pass_through_endpoints: List[Dict], +): + """Get combined pass-through endpoints from db + config""" + return pass_through_endpoints + config_pass_through_endpoints + async def initialize_pass_through_endpoints( pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]], ): """ + 1. Create a global list of pass-through endpoints (db + config) + 2. Clear all existing pass-through endpoints from the FastAPI app routes + 3. Add new endpoints to the in-memory registry + Initialize a list of pass-through endpoints by adding them to the FastAPI app routes Args: @@ -1711,9 +1819,26 @@ async def initialize_pass_through_endpoints( verbose_proxy_logger.debug("initializing pass through endpoints") from litellm.proxy._types import CommonProxyErrors, LiteLLMRoutes - from litellm.proxy.proxy_server import app, premium_user + from litellm.proxy.proxy_server import ( + app, + config_passthrough_endpoints, + premium_user, + ) - for endpoint in pass_through_endpoints: + ## get combined pass-through endpoints from db + config + combined_pass_through_endpoints: List[Union[Dict, PassThroughGenericEndpoint]] + + if config_passthrough_endpoints is not None: + combined_pass_through_endpoints = _get_combined_pass_through_endpoints( # type: ignore + pass_through_endpoints, config_passthrough_endpoints + ) + else: + combined_pass_through_endpoints = pass_through_endpoints # type: ignore + + ## clear all existing pass-through endpoints from the FastAPI app routes + InitPassThroughEndpointHelpers.clear_all_pass_through_routes() + + for endpoint in combined_pass_through_endpoints: if isinstance(endpoint, PassThroughGenericEndpoint): endpoint = endpoint.model_dump() @@ -1818,23 +1943,91 @@ async def _get_pass_through_endpoints_from_db( return returned_endpoints +async def _filter_endpoints_by_team_allowed_routes( + team_id: str, + pass_through_endpoints: List[PassThroughGenericEndpoint], + prisma_client, +) -> List[PassThroughGenericEndpoint]: + """ + Filter pass-through endpoints based on team's allowed_passthrough_routes metadata. + + Args: + team_id: The team ID to check permissions for + pass_through_endpoints: List of endpoints to filter + prisma_client: Database client + + Returns: + Filtered list of endpoints based on team permissions + + Raises: + HTTPException: If team is not found + """ + # retrieve team from db + team = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + ) + if team is None: + raise HTTPException( + status_code=404, + detail={"error": "Team not found"}, + ) + + # retrieve team metadata + team_metadata = team.metadata + if ( + team_metadata is not None + and team_metadata.get("allowed_passthrough_routes") is not None + ): + ## FILTER pass_through_endpoints by allowed_passthrough_routes + pass_through_endpoints = [ + endpoint + for endpoint in pass_through_endpoints + if endpoint.path in team_metadata.get("allowed_passthrough_routes") + ] + + return pass_through_endpoints + + @router.get( "/config/pass_through_endpoint", dependencies=[Depends(user_api_key_auth)], response_model=PassThroughEndpointResponse, ) +@router.get( + "/config/pass_through_endpoint/team/{team_id}", + dependencies=[Depends(user_api_key_auth)], + response_model=PassThroughEndpointResponse, +) async def get_pass_through_endpoints( endpoint_id: Optional[str] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + team_id: Optional[str] = None, ): """ GET configured pass through endpoint. If no endpoint_id given, return all configured endpoints. """ ## Get existing pass-through endpoint field value + from litellm.proxy._types import CommonProxyErrors + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + pass_through_endpoints = await _get_pass_through_endpoints_from_db( endpoint_id=endpoint_id, user_api_key_dict=user_api_key_dict ) + + if team_id is not None: + pass_through_endpoints = await _filter_endpoints_by_team_allowed_routes( + team_id=team_id, + pass_through_endpoints=pass_through_endpoints, + prisma_client=prisma_client, + ) + return PassThroughEndpointResponse(endpoints=pass_through_endpoints) @@ -1930,6 +2123,7 @@ async def update_pass_through_endpoints( field_value=pass_through_endpoint_data, config_type="general_settings", ) + await update_config_general_settings( data=updated_data, user_api_key_dict=user_api_key_dict ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index be0f7d6be08..8f5a046be03 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -259,9 +259,7 @@ from litellm.proxy.management_endpoints.customer_endpoints import ( from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) -from litellm.proxy.management_endpoints.internal_user_endpoints import ( - user_update, -) +from litellm.proxy.management_endpoints.internal_user_endpoints import user_update from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, duration_in_seconds, @@ -308,9 +306,7 @@ from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMi from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) -from litellm.proxy.openai_files_endpoints.files_endpoints import ( - set_files_config, -) +from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( passthrough_endpoint_router, ) @@ -580,7 +576,7 @@ async def _initialize_shared_aiohttp_session(): ttl_dns_cache=AIOHTTP_TTL_DNS_CACHE, enable_cleanup_closed=True, ) - + session = ClientSession(connector=connector) verbose_proxy_logger.info( f"SESSION REUSE: Created shared aiohttp session for connection pooling (ID: {id(session)})" @@ -723,7 +719,7 @@ async def proxy_startup_event(app: FastAPI): verbose_proxy_logger.info("SESSION REUSE: Closed shared aiohttp session") except Exception as e: verbose_proxy_logger.error(f"Error closing shared aiohttp session: {e}") - + await proxy_shutdown_event() @@ -995,13 +991,16 @@ experimental = False llm_router: Optional[Router] = None llm_model_list: Optional[list] = None general_settings: dict = {} +config_passthrough_endpoints: Optional[List[Dict[str, Any]]] = None callback_settings: dict = {} log_file = "api_log.json" worker_config = None master_key: Optional[str] = None otel_logging = False prisma_client: Optional[PrismaClient] = None -shared_aiohttp_session: Optional["ClientSession"] = None # Global shared session for connection reuse +shared_aiohttp_session: Optional["ClientSession"] = ( + None # Global shared session for connection reuse +) user_api_key_cache = DualCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) @@ -1385,31 +1384,33 @@ async def update_cache( # noqa: PLR0915 """ if tags is None or response_cost is None: return - + try: for tag_name in tags: if not tag_name or not isinstance(tag_name, str): continue - + cache_key = f"tag:{tag_name}" # Fetch the existing tag object from cache - existing_tag_obj = await user_api_key_cache.async_get_cache(key=cache_key) + existing_tag_obj = await user_api_key_cache.async_get_cache( + key=cache_key + ) if existing_tag_obj is None: # do nothing if tag not in api key cache continue - + verbose_proxy_logger.debug( f"_update_tag_cache: existing spend for tag={tag_name}: {existing_tag_obj}; response_cost: {response_cost}" ) - + if isinstance(existing_tag_obj, dict): existing_spend = existing_tag_obj.get("spend", 0) or 0 else: existing_spend = getattr(existing_tag_obj, "spend", 0) or 0 - + # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - + # Update the spend column for the given tag if isinstance(existing_tag_obj, dict): existing_tag_obj["spend"] = new_spend @@ -1482,6 +1483,7 @@ async def _run_background_health_check(): from litellm.proxy.health_check_utils.shared_health_check_manager import ( SharedHealthCheckManager, ) + shared_health_manager = SharedHealthCheckManager( redis_cache=redis_usage_cache, health_check_ttl=DEFAULT_SHARED_HEALTH_CHECK_TTL, @@ -1502,16 +1504,21 @@ async def _run_background_health_check(): # Use shared health check if available, otherwise fall back to direct health check # Convert health_check_details to bool for perform_shared_health_check (defaults to True if None) - details_bool = health_check_details if health_check_details is not None else True - + details_bool = ( + health_check_details if health_check_details is not None else True + ) + if shared_health_manager is not None: try: - healthy_endpoints, unhealthy_endpoints = await shared_health_manager.perform_shared_health_check( - model_list=_llm_model_list, details=details_bool + healthy_endpoints, unhealthy_endpoints = ( + await shared_health_manager.perform_shared_health_check( + model_list=_llm_model_list, details=details_bool + ) ) except Exception as e: verbose_proxy_logger.error( - "Error in shared health check, falling back to direct health check: %s", str(e) + "Error in shared health check, falling back to direct health check: %s", + str(e), ) healthy_endpoints, unhealthy_endpoints = await perform_health_check( model_list=_llm_model_list, details=health_check_details @@ -1879,7 +1886,7 @@ class ProxyConfig: """ Load config values into proxy global state """ - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, health_check_interval, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, callback_settings, proxy_batch_polling_interval + global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, health_check_interval, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, callback_settings, proxy_batch_polling_interval, config_passthrough_endpoints config: dict = await self.get_config(config_file_path=config_file_path) @@ -2234,9 +2241,13 @@ class ProxyConfig: ## pass through endpoints if general_settings.get("pass_through_endpoints", None) is not None: + config_passthrough_endpoints = general_settings[ + "pass_through_endpoints" + ] await initialize_pass_through_endpoints( pass_through_endpoints=general_settings["pass_through_endpoints"] ) + ## ADMIN UI ACCESS ## ui_access_mode = general_settings.get( "ui_access_mode", "all" @@ -3055,7 +3066,9 @@ class ProxyConfig: return current_config # For dictionary values, update only non-none values - if isinstance(current_config[param_name], dict) and isinstance(db_param_value, dict): + if isinstance(current_config[param_name], dict) and isinstance( + db_param_value, dict + ): _deep_merge_dicts(current_config[param_name], db_param_value) else: # Non-dict or mismatched types: DB value replaces config (unchanged behavior) diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 9317bf26178..7e0b4cfa243 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -58,7 +58,7 @@ class LiteLLMCompletionTransformationHandler: responses_api_request=responses_api_request, **kwargs, ) - + completion_args = {} completion_args.update(kwargs) completion_args.update(litellm_completion_request) @@ -83,6 +83,7 @@ class LiteLLMCompletionTransformationHandler: elif isinstance(litellm_completion_response, litellm.CustomStreamWrapper): return LiteLLMCompletionStreamingIterator( + model=model, litellm_custom_stream_wrapper=litellm_completion_response, request_input=input, responses_api_request=responses_api_request, @@ -106,7 +107,7 @@ class LiteLLMCompletionTransformationHandler: previous_response_id=previous_response_id, litellm_completion_request=litellm_completion_request, ) - + acompletion_args = {} acompletion_args.update(kwargs) acompletion_args.update(litellm_completion_request) @@ -130,9 +131,12 @@ class LiteLLMCompletionTransformationHandler: elif isinstance(litellm_completion_response, litellm.CustomStreamWrapper): return LiteLLMCompletionStreamingIterator( + model=litellm_completion_request.get("model") or "", litellm_custom_stream_wrapper=litellm_completion_response, request_input=request_input, responses_api_request=responses_api_request, - custom_llm_provider=litellm_completion_request.get("custom_llm_provider"), + custom_llm_provider=litellm_completion_request.get( + "custom_llm_provider" + ), litellm_metadata=kwargs.get("litellm_metadata", {}), ) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 20822b3628e..a4fb2d96032 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -1,4 +1,6 @@ -from typing import List, Optional, Union +import time +import uuid +from typing import List, Optional, Union, cast import litellm from litellm.main import stream_chunk_builder @@ -8,11 +10,23 @@ from litellm.responses.litellm_completion_transformation.transformation import ( from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( + PART_UNION_TYPES, + BaseLiteLLMOpenAIResponseObject, + ContentPartAddedEvent, + ContentPartDoneEvent, + ContentPartDonePartOutputText, + ContentPartDonePartReasoningText, + OutputItemAddedEvent, + OutputItemDoneEvent, OutputTextDeltaEvent, + OutputTextDoneEvent, ReasoningSummaryTextDeltaEvent, ResponseCompletedEvent, + ResponseCreatedEvent, + ResponseInProgressEvent, ResponseInputParam, ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, ) @@ -32,12 +46,14 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def __init__( self, + model: str, litellm_custom_stream_wrapper: litellm.CustomStreamWrapper, request_input: Union[str, ResponseInputParam], responses_api_request: ResponsesAPIOptionalRequestParams, custom_llm_provider: Optional[str] = None, litellm_metadata: Optional[dict] = None, ): + self.model: str = model self.litellm_custom_stream_wrapper: litellm.CustomStreamWrapper = ( litellm_custom_stream_wrapper ) @@ -50,14 +66,274 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.collected_chat_completion_chunks: List[ModelResponseStream] = [] self.finished: bool = False self.litellm_logging_obj = litellm_custom_stream_wrapper.logging_obj + self.sent_response_created_event: bool = False + self.sent_response_in_progress_event: bool = False + self.sent_output_item_added_event: bool = False + self.sent_content_part_added_event: bool = False + self.sent_output_text_done_event: bool = False + self.sent_output_content_part_done_event: bool = False + self.sent_output_item_done_event: bool = False + self.litellm_model_response: Optional[ + Union[ModelResponse, TextCompletionResponse] + ] = None + self.final_text: str = "" + + def _default_response_created_event_data(self) -> dict: + response_created_event_data = { + "id": f"resp_{str(uuid.uuid4())}", + "object": "response", + "created_at": int(time.time()), + "status": "in_progress", + "error": None, + "incomplete_details": None, + "instructions": self.request_input, + "max_output_tokens": None, + "model": self.model, + "output": [], + "parallel_tool_calls": True, + "previous_response_id": None, + "reasoning": {"effort": None, "summary": None}, + "store": True, + } + if "temperature" in self.responses_api_request: + response_created_event_data["temperature"] = self.responses_api_request[ + "temperature" + ] + if "text" in self.responses_api_request: + response_created_event_data["text"] = self.responses_api_request["text"] + if "tool_choice" in self.responses_api_request: + response_created_event_data["tool_choice"] = self.responses_api_request[ + "tool_choice" + ] + else: + response_created_event_data["tool_choice"] = "auto" + if "tools" in self.responses_api_request: + response_created_event_data["tools"] = self.responses_api_request["tools"] + else: + response_created_event_data["tools"] = [] + if "top_p" in self.responses_api_request: + response_created_event_data["top_p"] = self.responses_api_request["top_p"] + else: + response_created_event_data["top_p"] = 1.0 + if "truncation" in self.responses_api_request: + response_created_event_data["truncation"] = self.responses_api_request[ + "truncation" + ] + if "user" in self.responses_api_request: + response_created_event_data["user"] = self.responses_api_request["user"] + if "metadata" in self.responses_api_request: + response_created_event_data["metadata"] = self.responses_api_request[ + "metadata" + ] + return response_created_event_data + + def create_response_created_event(self) -> ResponseCreatedEvent: + """ + data: {"type":"response.created","response":{"id":"resp_67c9fdcecf488190bdd9a0409de3a1ec07b8b0ad4e5eb654","object":"response","created_at":1741290958,"status":"in_progress","error":null,"incomplete_details":null,"instructions":"You are a helpful assistant.","max_output_tokens":null,"model":"gpt-4.1-2025-04-14","output":[],"parallel_tool_calls":true,"previous_response_id":null,"reasoning":{"effort":null,"summary":null},"store":true,"temperature":1.0,"text":{"format":{"type":"text"}},"tool_choice":"auto","tools":[],"top_p":1.0,"truncation":"disabled","usage":null,"user":null,"metadata":{}}} + + """ + response_created_event_data = self._default_response_created_event_data() + return ResponseCreatedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_CREATED, + response=ResponsesAPIResponse(**response_created_event_data), + ) + + def create_response_in_progress_event(self) -> ResponseInProgressEvent: + response_in_progress_event_data = self._default_response_created_event_data() + response_in_progress_event_data["status"] = "in_progress" + return ResponseInProgressEvent( + type=ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, + response=ResponsesAPIResponse(**response_in_progress_event_data), + ) + + def create_output_item_added_event(self) -> OutputItemAddedEvent: + return OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=0, + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": f"msg_{str(uuid.uuid4())}", + "type": "message", + "status": "in_progress", + "role": "assistant", + "content": [], + } + ), + ) + + def create_content_part_added_event(self) -> ContentPartAddedEvent: + return ContentPartAddedEvent( + type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED, + item_id=f"msg_{str(uuid.uuid4())}", + output_index=0, + content_index=0, + part=BaseLiteLLMOpenAIResponseObject( + **{"type": "output_text", "text": "", "annotations": []} + ), + ) + + def create_litellm_model_response( + self, + ) -> Optional[ModelResponse]: + return cast( + Optional[ModelResponse], + stream_chunk_builder( + chunks=self.collected_chat_completion_chunks, + logging_obj=self.litellm_logging_obj, + ), + ) + + def create_output_text_done_event( + self, litellm_complete_object: ModelResponse + ) -> OutputTextDoneEvent: + return OutputTextDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + item_id=f"msg_{str(uuid.uuid4())}", + output_index=0, + content_index=0, + text=getattr(litellm_complete_object.choices[0].message, "content", "") # type: ignore + or "", + ) + + def create_output_content_part_done_event( + self, litellm_complete_object: ModelResponse + ) -> ContentPartDoneEvent: + + text = getattr(litellm_complete_object.choices[0].message, "content", "") or "" # type: ignore + reasoning_content = getattr(litellm_complete_object.choices[0].message, "reasoning_content", "") or "" # type: ignore + + part: Optional[PART_UNION_TYPES] = None + if reasoning_content: + part = ContentPartDonePartReasoningText( + type="reasoning_text", + reasoning=reasoning_content, + ) + + else: + part = ContentPartDonePartOutputText( + type="output_text", + text=text, + annotations=[], + logprobs=None, + ) + + return ContentPartDoneEvent( + type=ResponsesAPIStreamEvents.CONTENT_PART_DONE, + item_id=f"msg_{str(uuid.uuid4())}", + output_index=0, + content_index=0, + part=part, + ) + + def create_output_item_done_event( + self, litellm_complete_object: ModelResponse + ) -> OutputItemDoneEvent: + text = self.litellm_model_response.choices[0].message.content or "" # type: ignore + return OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=0, + sequence_number=1, + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": f"msg_{str(uuid.uuid4())}", + "status": "completed", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": text, + "annotations": [], + } + ], + } + ), + ) + + def return_default_done_events( + self, litellm_complete_object: ModelResponse + ) -> Optional[BaseLiteLLMOpenAIResponseObject]: + if self.sent_output_text_done_event is False: + self.sent_output_text_done_event = True + return self.create_output_text_done_event(litellm_complete_object) + if self.sent_output_content_part_done_event is False: + self.sent_output_content_part_done_event = True + return self.create_output_content_part_done_event(litellm_complete_object) + if self.sent_output_item_done_event is False: + self.sent_output_item_done_event = True + return self.create_output_item_done_event(litellm_complete_object) + return None + + def return_default_initial_events( + self, + ) -> Optional[BaseLiteLLMOpenAIResponseObject]: + if self.sent_response_created_event is False: + self.sent_response_created_event = True + return self.create_response_created_event() + elif self.sent_response_in_progress_event is False: + self.sent_response_in_progress_event = True + return self.create_response_in_progress_event() + elif self.sent_output_item_added_event is False: + self.sent_output_item_added_event = True + return self.create_output_item_added_event() + elif self.sent_content_part_added_event is False: + self.sent_content_part_added_event = True + return self.create_content_part_added_event() + return None + + def is_stream_finished(self) -> bool: + if ( + self.sent_output_text_done_event is True + and self.sent_output_content_part_done_event is True + and self.sent_output_item_done_event is True + ): + return True + return False + + def common_done_event_logic( + self, sync_mode: bool = True + ) -> BaseLiteLLMOpenAIResponseObject: + if not self.litellm_model_response or isinstance( + self.litellm_model_response, TextCompletionResponse + ): + self.litellm_model_response = self.create_litellm_model_response() + if self.litellm_model_response: + done_event = self.return_default_done_events(self.litellm_model_response) + if done_event: + return done_event + else: + if sync_mode: + raise StopIteration + else: + raise StopAsyncIteration + + self.finished = self.is_stream_finished() + response_completed_event = self._emit_response_completed_event( + self.litellm_model_response + ) + if response_completed_event: + return response_completed_event + else: + if sync_mode: + raise StopIteration + else: + raise StopAsyncIteration async def __anext__( self, - ) -> Union[ResponsesAPIStreamingResponse, ResponseCompletedEvent]: + ) -> Union[ + ResponsesAPIStreamingResponse, + ResponseCompletedEvent, + BaseLiteLLMOpenAIResponseObject, + ]: try: while True: if self.finished is True: raise StopAsyncIteration + + result = self.return_default_initial_events() + if result: + return result # Get the next chunk from the stream try: chunk = await self.litellm_custom_stream_wrapper.__anext__() @@ -70,12 +346,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if response_api_chunk: return response_api_chunk except StopAsyncIteration: - self.finished = True - response_completed_event = self._emit_response_completed_event() - if response_completed_event: - return response_completed_event - else: - raise StopAsyncIteration + return self.common_done_event_logic(sync_mode=False) except Exception as e: # Handle HTTP errors @@ -87,12 +358,20 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def __next__( self, - ) -> Union[ResponsesAPIStreamingResponse, ResponseCompletedEvent]: + ) -> Union[ + ResponsesAPIStreamingResponse, + ResponseCompletedEvent, + BaseLiteLLMOpenAIResponseObject, + ]: try: while True: if self.finished is True: raise StopIteration # Get the next chunk from the stream + + result = self.return_default_initial_events() + if result: + return result try: chunk = self.litellm_custom_stream_wrapper.__next__() self.collected_chat_completion_chunks.append(chunk) @@ -104,13 +383,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if response_api_chunk: return response_api_chunk except StopIteration: - self.finished = True - response_completed_event = self._emit_response_completed_event() - if response_completed_event: - return response_completed_event - else: - raise StopIteration - + return self.common_done_event_logic(sync_mode=True) except Exception as e: # Handle HTTP errors self.finished = True @@ -165,26 +438,33 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): chat_completion_delta: ChatCompletionDelta = choice.delta return chat_completion_delta.content or "" - def _emit_response_completed_event(self) -> Optional[ResponseCompletedEvent]: - litellm_model_response: Optional[ - Union[ModelResponse, TextCompletionResponse] - ] = stream_chunk_builder(chunks=self.collected_chat_completion_chunks, logging_obj=self.litellm_logging_obj) - if litellm_model_response and isinstance(litellm_model_response, ModelResponse): + def _emit_response_completed_event( + self, litellm_model_response: ModelResponse + ) -> Optional[ResponseCompletedEvent]: + + if litellm_model_response: # Add cost to usage object if include_cost_in_streaming_usage is True - if litellm.include_cost_in_streaming_usage and self.litellm_logging_obj is not None: + if ( + litellm.include_cost_in_streaming_usage + and self.litellm_logging_obj is not None + ): usage = getattr(litellm_model_response, "usage", None) if usage is not None: setattr( - usage, "cost", self.litellm_logging_obj._response_cost_calculator(result=litellm_model_response) + usage, + "cost", + self.litellm_logging_obj._response_cost_calculator( + result=litellm_model_response + ), ) - + # Transform the response responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( request_input=self.request_input, chat_completion_response=litellm_model_response, responses_api_request=self.responses_api_request, ) - + # Encode the response ID to match non-streaming behavior encoded_response = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( responses_api_response=responses_api_response, diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 7d53452c1c0..8801e561915 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -1,19 +1,10 @@ -from litellm._uuid import uuid -from typing import ( - TYPE_CHECKING, - Any, - Dict, - List, - Optional, - Union, - cast, -) +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast from litellm._logging import verbose_logger -from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, -) +from litellm._uuid import uuid +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( + BaseLiteLLMOpenAIResponseObject, MCPCallArgumentsDeltaEvent, MCPCallArgumentsDoneEvent, MCPCallCompletedEvent, @@ -38,22 +29,24 @@ async def create_mcp_list_tools_events( mcp_tools_with_litellm_proxy: List[ToolParam], user_api_key_auth: Any, base_item_id: str, - pre_processed_mcp_tools: List[Any] + pre_processed_mcp_tools: List[Any], ) -> List[ResponsesAPIStreamingResponse]: """Create MCP discovery events using pre-processed tools from the parent""" - + events: List[ResponsesAPIStreamingResponse] = [] - + try: # Extract MCP server names mcp_servers = [] for tool in mcp_tools_with_litellm_proxy: if isinstance(tool, dict) and "server_url" in tool: server_url = tool.get("server_url") - if isinstance(server_url, str) and server_url.startswith("litellm_proxy/mcp/"): + if isinstance(server_url, str) and server_url.startswith( + "litellm_proxy/mcp/" + ): server_name = server_url.split("/")[-1] mcp_servers.append(server_name) - + # Emit list tools in progress event in_progress_event = MCPListToolsInProgressEvent( type=ResponsesAPIStreamEvents.MCP_LIST_TOOLS_IN_PROGRESS, @@ -62,21 +55,21 @@ async def create_mcp_list_tools_events( item_id=base_item_id, ) events.append(in_progress_event) - + # Use the pre-processed MCP tools that were already fetched, filtered, and deduplicated by the parent filtered_mcp_tools = pre_processed_mcp_tools - + # Convert tools to dict format for the event mcp_tools_dict = [] for tool in filtered_mcp_tools: - if hasattr(tool, 'model_dump') and callable(getattr(tool, 'model_dump')): + if hasattr(tool, "model_dump") and callable(getattr(tool, "model_dump")): # Type cast to help mypy understand this is safe after hasattr check mcp_tools_dict.append(cast(Any, tool).model_dump()) - elif hasattr(tool, '__dict__'): + elif hasattr(tool, "__dict__"): mcp_tools_dict.append(tool.__dict__) else: - mcp_tools_dict.append({"name": getattr(tool, 'name', str(tool))}) - + mcp_tools_dict.append({"name": getattr(tool, "name", str(tool))}) + # Emit list tools completed event completed_event = MCPListToolsCompletedEvent( type=ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED, @@ -85,7 +78,7 @@ async def create_mcp_list_tools_events( item_id=base_item_id, ) events.append(completed_event) - + # Add output_item.done event with the actual tools list (matching OpenAI format) from litellm.types.llms.openai import OutputItemDoneEvent @@ -95,45 +88,50 @@ async def create_mcp_list_tools_events( first_tool = mcp_tools_with_litellm_proxy[0] if isinstance(first_tool, dict): server_label_value = first_tool.get("server_label", "") - server_label = str(server_label_value) if server_label_value is not None else "" - + server_label = ( + str(server_label_value) if server_label_value is not None else "" + ) + # Format tools for OpenAI output_item.done format formatted_tools = [] for tool in filtered_mcp_tools: tool_dict = { - "name": getattr(tool, 'name', 'unknown'), - "description": getattr(tool, 'description', ''), + "name": getattr(tool, "name", "unknown"), + "description": getattr(tool, "description", ""), "annotations": {"read_only": False}, } - + # Add input_schema if available - if hasattr(tool, 'inputSchema'): - tool_dict["input_schema"] = getattr(tool, 'inputSchema') - elif hasattr(tool, 'input_schema'): - tool_dict["input_schema"] = getattr(tool, 'input_schema') - + if hasattr(tool, "inputSchema"): + tool_dict["input_schema"] = getattr(tool, "inputSchema") + elif hasattr(tool, "input_schema"): + tool_dict["input_schema"] = getattr(tool, "input_schema") + formatted_tools.append(tool_dict) - + # Create the output_item.done event with MCP tools list output_item_done_event = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=0, - item={ - "id": base_item_id, - "type": "mcp_list_tools", - "server_label": server_label, - "tools": formatted_tools - } + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": base_item_id, + "type": "mcp_list_tools", + "server_label": server_label, + "tools": formatted_tools, + } + ), ) events.append(output_item_done_event) - + verbose_logger.debug(f"Created {len(events)} MCP discovery events") - + except Exception as e: verbose_logger.error(f"Error creating MCP list tools events: {e}") import traceback + traceback.print_exc() - + # Emit failed event on error failed_event = MCPListToolsFailedEvent( type=ResponsesAPIStreamEvents.MCP_LIST_TOOLS_FAILED, @@ -142,37 +140,39 @@ async def create_mcp_list_tools_events( item_id=base_item_id, ) events.append(failed_event) - + # Still emit output_item.done event even on failure (with empty tools list) from litellm.types.llms.openai import OutputItemDoneEvent - + output_item_done_event = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=0, - item={ - "id": base_item_id, - "type": "mcp_list_tools", - "server_label": "", - "tools": [] - } + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": base_item_id, + "type": "mcp_list_tools", + "server_label": "", + "tools": [], + } + ), ) events.append(output_item_done_event) - + return events def create_mcp_call_events( - tool_name: str, - tool_call_id: str, + tool_name: str, + tool_call_id: str, arguments: str, result: Optional[str] = None, base_item_id: Optional[str] = None, - sequence_start: int = 1 + sequence_start: int = 1, ) -> List[ResponsesAPIStreamingResponse]: """Create MCP call events following OpenAI's specification""" events: List[ResponsesAPIStreamingResponse] = [] item_id = base_item_id or f"mcp_{uuid.uuid4().hex[:8]}" - + # MCP call in progress event in_progress_event = MCPCallInProgressEvent( type=ResponsesAPIStreamEvents.MCP_CALL_IN_PROGRESS, @@ -181,7 +181,7 @@ def create_mcp_call_events( item_id=item_id, ) events.append(in_progress_event) - + # MCP call arguments delta event (streaming the arguments) arguments_delta_event = MCPCallArgumentsDeltaEvent( type=ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA, @@ -191,7 +191,7 @@ def create_mcp_call_events( sequence_number=sequence_start + 1, ) events.append(arguments_delta_event) - + # MCP call arguments done event arguments_done_event = MCPCallArgumentsDoneEvent( type=ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DONE, @@ -201,7 +201,7 @@ def create_mcp_call_events( sequence_number=sequence_start + 2, ) events.append(arguments_done_event) - + # MCP call completed event (or failed if result indicates failure) if result is not None: completed_event = MCPCallCompletedEvent( @@ -211,23 +211,25 @@ def create_mcp_call_events( output_index=0, ) events.append(completed_event) - + # Add output_item.done event with the tool call result from litellm.types.llms.openai import OutputItemDoneEvent - + output_item_done_event = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=0, - item={ - "id": item_id, - "type": "mcp_call", - "approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}", - "arguments": arguments, - "error": None, - "name": tool_name, - "output": result, - "server_label": "litellm" - }, + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": item_id, + "type": "mcp_call", + "approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}", + "arguments": arguments, + "error": None, + "name": tool_name, + "output": result, + "server_label": "litellm", + } + ), ) events.append(output_item_done_event) else: @@ -238,7 +240,7 @@ def create_mcp_call_events( output_index=0, ) events.append(failed_event) - + return events @@ -250,51 +252,60 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): 3. Handles tool execution and follow-up calls for auto-execute tools 4. Emits tool execution events in the stream """ - + def __init__( self, base_iterator: Any, # Can be None - will be created internally mcp_events: List[ResponsesAPIStreamingResponse], mcp_tools_with_litellm_proxy: Optional[List[Any]] = None, user_api_key_auth: Any = None, - original_request_params: Optional[Dict[str, Any]] = None + original_request_params: Optional[Dict[str, Any]] = None, ): # MCP setup self.mcp_tools_with_litellm_proxy = mcp_tools_with_litellm_proxy or [] self.user_api_key_auth = user_api_key_auth self.original_request_params = original_request_params or {} self.should_auto_execute = self._should_auto_execute_tools() - + # Streaming state management self.phase = "mcp_discovery" # mcp_discovery -> initial_response -> tool_execution -> follow_up_response -> finished self.finished = False - + # Event queues and generation flags - self.mcp_discovery_events: List[ResponsesAPIStreamingResponse] = mcp_events # Pre-generated MCP discovery events + self.mcp_discovery_events: List[ResponsesAPIStreamingResponse] = ( + mcp_events # Pre-generated MCP discovery events + ) self.tool_execution_events: List[ResponsesAPIStreamingResponse] = [] self.mcp_discovery_generated = True # Events are already generated - self.mcp_events = mcp_events # Store the initial MCP events for backward compatibility - + self.mcp_events = ( + mcp_events # Store the initial MCP events for backward compatibility + ) + # Iterator references - self.base_iterator: Optional[Union[Any, ResponsesAPIResponse]] = base_iterator # Will be created when needed + self.base_iterator: Optional[Union[Any, ResponsesAPIResponse]] = ( + base_iterator # Will be created when needed + ) self.follow_up_iterator: Optional[Any] = None - + # Response collection for tool execution self.collected_response: Optional[ResponsesAPIResponse] = None - + # Set up model metadata (will be updated when we get the real iterator) - self.model = self.original_request_params.get('model', 'unknown') + self.model = self.original_request_params.get("model", "unknown") self.litellm_metadata = {} - self.custom_llm_provider = self.original_request_params.get('custom_llm_provider', None) - + self.custom_llm_provider = self.original_request_params.get( + "custom_llm_provider", None + ) + # Mark as async iterator self.is_async = True - + def _should_auto_execute_tools(self) -> bool: """Check if tools should be auto-executed""" from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) + return LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools( self.mcp_tools_with_litellm_proxy ) @@ -306,45 +317,49 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): """ Phase-based streaming: 1. mcp_discovery - Emit MCP discovery events - 2. initial_response - Stream the first LLM response + 2. initial_response - Stream the first LLM response 3. tool_execution - Emit tool execution events 4. follow_up_response - Stream the follow-up response 5. finished - End iteration """ - + # Phase 1: MCP Discovery Events if self.phase == "mcp_discovery": # Generate MCP discovery events if not already done # MCP discovery events are already generated and available - + # Emit MCP discovery events if self.mcp_discovery_events: return self.mcp_discovery_events.pop(0) - + # All MCP discovery events emitted, move to next phase - verbose_logger.debug("MCP discovery phase complete, transitioning to initial_response") + verbose_logger.debug( + "MCP discovery phase complete, transitioning to initial_response" + ) self.phase = "initial_response" await self._create_initial_response_iterator() # Fall through to process the initial response immediately - + # Phase 2: Initial Response Stream if self.phase == "initial_response": if self.base_iterator: # Check if base_iterator is actually iterable - if hasattr(self.base_iterator, '__anext__'): + if hasattr(self.base_iterator, "__anext__"): try: chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined] - + # If auto-execution is enabled, check for completed responses - if self.should_auto_execute and self._is_response_completed(chunk): + if self.should_auto_execute and self._is_response_completed( + chunk + ): # Collect the response for tool execution - response_obj = getattr(chunk, 'response', None) + response_obj = getattr(chunk, "response", None) if isinstance(response_obj, ResponsesAPIResponse): self.collected_response = response_obj # Move to tool execution phase after emitting this chunk self.phase = "tool_execution" await self._generate_tool_execution_events() - + return chunk except StopAsyncIteration: # Initial response ended, move to next phase @@ -357,24 +372,26 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): else: # base_iterator is not async iterable (likely a ResponsesAPIResponse) # Collect it for tool execution if needed - if self.should_auto_execute and isinstance(self.base_iterator, ResponsesAPIResponse): + if self.should_auto_execute and isinstance( + self.base_iterator, ResponsesAPIResponse + ): self.collected_response = self.base_iterator self.phase = "tool_execution" await self._generate_tool_execution_events() else: self.phase = "finished" raise StopAsyncIteration - + # Phase 3: Tool Execution Events if self.phase == "tool_execution": # Emit any queued tool execution events if self.tool_execution_events: return self.tool_execution_events.pop(0) - + # Move to follow-up response phase self.phase = "follow_up_response" await self._create_follow_up_iterator() - + # Phase 4: Follow-up Response Stream if self.phase == "follow_up_response": if self.follow_up_iterator: @@ -386,20 +403,22 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): else: self.phase = "finished" raise StopAsyncIteration - + # Phase 5: Finished if self.phase == "finished": raise StopAsyncIteration - + # Should not reach here raise StopAsyncIteration - + def _is_response_completed(self, chunk: ResponsesAPIStreamingResponse) -> bool: """Check if this chunk indicates the response is completed""" from litellm.types.llms.openai import ResponsesAPIStreamEvents - return getattr(chunk, 'type', None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED - - + + return ( + getattr(chunk, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + async def _create_initial_response_iterator(self) -> None: """Create the initial response iterator by making the first LLM call""" try: @@ -408,38 +427,45 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Make the initial response API call - but avoid the MCP wrapper params = self.original_request_params.copy() - params['stream'] = True # Ensure streaming - + params["stream"] = True # Ensure streaming + # Use the pre-fetched all_tools from original_request_params (no re-processing needed) params_for_llm = {} for key, value in params.items(): - params_for_llm[key] = value # Copy all params as-is since tools are already processed - - tools_count = len(params_for_llm.get('tools', [])) + params_for_llm[key] = ( + value # Copy all params as-is since tools are already processed + ) + + tools_count = len(params_for_llm.get("tools", [])) verbose_logger.debug(f"Making LLM call with {tools_count} tools") response = await aresponses(**params_for_llm) - + # Set the base iterator - if hasattr(response, '__aiter__') or hasattr(response, '__iter__'): + if hasattr(response, "__aiter__") or hasattr(response, "__iter__"): self.base_iterator = response # Copy metadata from the real iterator - self.model = getattr(response, 'model', self.model) - self.litellm_metadata = getattr(response, 'litellm_metadata', {}) - self.custom_llm_provider = getattr(response, 'custom_llm_provider', self.custom_llm_provider) - verbose_logger.debug(f"Created base iterator: {type(self.base_iterator)}") + self.model = getattr(response, "model", self.model) + self.litellm_metadata = getattr(response, "litellm_metadata", {}) + self.custom_llm_provider = getattr( + response, "custom_llm_provider", self.custom_llm_provider + ) + verbose_logger.debug( + f"Created base iterator: {type(self.base_iterator)}" + ) else: # Non-streaming response - this shouldn't happen but handle it verbose_logger.warning(f"Got non-streaming response: {type(response)}") self.base_iterator = None self.phase = "finished" - + except Exception as e: verbose_logger.error(f"Error creating initial response iterator: {e}") import traceback + traceback.print_exc() self.base_iterator = None self.phase = "finished" - + async def _generate_tool_execution_events(self) -> None: """Generate tool execution events and execute tools""" if not self.collected_response: @@ -447,7 +473,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) - + try: # Extract tool calls from the response if self.collected_response is not None: @@ -456,9 +482,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): tool_calls = [] if not tool_calls: return - + for tool_call in tool_calls: - tool_name, tool_arguments, tool_call_id = LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call) + tool_name, tool_arguments, tool_call_id = ( + LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call) + ) if tool_name and tool_call_id: # Create MCP call events for this tool execution call_events = create_mcp_call_events( @@ -467,34 +495,35 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): arguments=tool_arguments or "{}", # JSON string with arguments result=None, # Will be set after execution base_item_id=f"mcp_{uuid.uuid4().hex[:8]}", - sequence_start=len(self.tool_execution_events) + 1 + sequence_start=len(self.tool_execution_events) + 1, ) # Add the in_progress and arguments events (not the completed event yet) self.tool_execution_events.extend(call_events[:-1]) - + # Execute the tools tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( - tool_calls=tool_calls, - user_api_key_auth=self.user_api_key_auth + tool_calls=tool_calls, user_api_key_auth=self.user_api_key_auth ) - + # Create completion events and output_item.done events for tool execution for tool_result in tool_results: tool_call_id = tool_result.get("tool_call_id", "unknown") result_text = tool_result.get("result", "") - + # Find matching tool name and arguments tool_name = "unknown" tool_arguments = "{}" for tool_call in tool_calls: - name, args, call_id = LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call) + name, args, call_id = ( + LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call) + ) if call_id == tool_call_id: tool_name = name or "unknown" tool_arguments = args or "{}" break - + item_id = f"mcp_{uuid.uuid4().hex[:8]}" - + # Create the completion event completed_event = MCPCallCompletedEvent( type=ResponsesAPIStreamEvents.MCP_CALL_COMPLETED, @@ -503,79 +532,84 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): output_index=0, ) self.tool_execution_events.append(completed_event) - + # Create output_item.done event with the tool call result from litellm.types.llms.openai import OutputItemDoneEvent - + output_item_done_event = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=0, - item={ - "id": item_id, - "type": "mcp_call", - "approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}", - "arguments": tool_arguments, - "error": None, - "name": tool_name, - "output": result_text, - "server_label": "litellm" # or extract from tool config - }, + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": item_id, + "type": "mcp_call", + "approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}", + "arguments": tool_arguments, + "error": None, + "name": tool_name, + "output": result_text, + "server_label": "litellm", # or extract from tool config + } + ), ) self.tool_execution_events.append(output_item_done_event) - + # Store tool results for follow-up call self.tool_results = tool_results - + except Exception as e: verbose_logger.error(f"Error in tool execution: {e}") import traceback + traceback.print_exc() self.tool_results = [] - + async def _create_follow_up_iterator(self) -> None: """Create the follow-up response iterator with tool results""" - if not self.collected_response or not hasattr(self, 'tool_results'): + if not self.collected_response or not hasattr(self, "tool_results"): return - + from litellm.responses.main import aresponses from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) - + try: # Create follow-up input if self.collected_response is not None: follow_up_input = LiteLLM_Proxy_MCP_Handler._create_follow_up_input( response=self.collected_response, # type: ignore[arg-type] tool_results=self.tool_results, - original_input=self.original_request_params.get('input') + original_input=self.original_request_params.get("input"), ) - + # Make follow-up call with streaming follow_up_params = self.original_request_params.copy() - follow_up_params.update({ - 'input': follow_up_input, - 'previous_response_id': self.collected_response.id, # type: ignore[attr-defined] - 'stream': True - }) + follow_up_params.update( + { + "input": follow_up_input, + "previous_response_id": self.collected_response.id, # type: ignore[attr-defined] + "stream": True, + } + ) else: return # Remove tool_choice to avoid forcing more tool calls - follow_up_params.pop('tool_choice', None) - + follow_up_params.pop("tool_choice", None) + follow_up_response = await aresponses(**follow_up_params) - + # Set up the follow-up iterator - if hasattr(follow_up_response, '__aiter__'): + if hasattr(follow_up_response, "__aiter__"): self.follow_up_iterator = follow_up_response - + except Exception as e: verbose_logger.error(f"Error creating follow-up iterator: {e}") import traceback + traceback.print_exc() self.follow_up_iterator = None - def __iter__(self): return self @@ -583,11 +617,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # First, emit any queued MCP events if self.mcp_events: # type: ignore[attr-defined] return self.mcp_events.pop(0) # type: ignore[attr-defined] - + # Then delegate to the base iterator if not self.is_async: try: - if self.base_iterator and hasattr(self.base_iterator, '__next__'): + if self.base_iterator and hasattr(self.base_iterator, "__next__"): return next(cast(Any, self.base_iterator)) # type: ignore[arg-type] else: raise StopIteration diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index eda3e6921da..b78913402ff 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -93,24 +93,35 @@ class BaseResponsesAPIStreamingIterator: # Store the completed response if ( openai_responses_api_chunk - and openai_responses_api_chunk.type + and getattr(openai_responses_api_chunk, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED ): self.completed_response = openai_responses_api_chunk # Add cost to usage object if include_cost_in_streaming_usage is True - if litellm.include_cost_in_streaming_usage and self.logging_obj is not None: - response_obj: Optional[ResponsesAPIResponse] = getattr(openai_responses_api_chunk, "response", None) + if ( + litellm.include_cost_in_streaming_usage + and self.logging_obj is not None + ): + response_obj: Optional[ResponsesAPIResponse] = getattr( + openai_responses_api_chunk, "response", None + ) if response_obj: - usage_obj: Optional[ResponseAPIUsage] = getattr(response_obj, "usage", None) + usage_obj: Optional[ResponseAPIUsage] = getattr( + response_obj, "usage", None + ) if usage_obj is not None: try: - cost: Optional[float] = self.logging_obj._response_cost_calculator(result=response_obj) + cost: Optional[float] = ( + self.logging_obj._response_cost_calculator( + result=response_obj + ) + ) if cost is not None: setattr(usage_obj, "cost", cost) except Exception: # If cost calculation fails, continue without cost pass - + self._handle_logging_completed_response() return openai_responses_api_chunk diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 005aec31e2e..56971b91e5d 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -43,12 +43,7 @@ from openai.types.responses.response import ( # Handle OpenAI SDK version compatibility for Text type try: - # fmt: off - from openai.types.responses.response_create_params import ( # type: ignore[attr-defined] - Text as ResponseText, # type: ignore[attr-defined] - ) - - # fmt: on + from openai.types.responses.response_create_params import ( Text as ResponseText ) # type: ignore[attr-defined] # fmt: skip # isort: skip except (ImportError, AttributeError): # Fall back to the concrete config type available in all SDK versions from openai.types.responses.response_text_config_param import ( @@ -992,7 +987,9 @@ class ResponsesAPIOptionalRequestParams(TypedDict, total=False): prompt_cache_key: Optional[str] stream_options: Optional[dict] top_logprobs: Optional[int] - partial_images: Optional[int] # Number of partial images to generate (1-3) for streaming image generation + partial_images: Optional[ + int + ] # Number of partial images to generate (1-3) for streaming image generation class ResponsesAPIRequestParams(ResponsesAPIOptionalRequestParams, total=False): @@ -1056,7 +1053,9 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): parallel_tool_calls: Optional[bool] = None temperature: Optional[float] = None tool_choice: Optional[ToolChoice] = None - tools: Optional[Union[List[Tool], List[ResponseFunctionToolCall], List[Dict[str, Any]]]] = None + tools: Optional[ + Union[List[Tool], List[ResponseFunctionToolCall], List[Dict[str, Any]]] + ] = None top_p: Optional[float] = None max_output_tokens: Optional[int] = None previous_response_id: Optional[str] = None @@ -1180,13 +1179,27 @@ class ReasoningSummaryTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): class OutputItemAddedEvent(BaseLiteLLMOpenAIResponseObject): type: Literal[ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED] output_index: int - item: Optional[dict] + item: Optional[BaseLiteLLMOpenAIResponseObject] class OutputItemDoneEvent(BaseLiteLLMOpenAIResponseObject): type: Literal[ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE] output_index: int - item: dict + sequence_number: int = 1 + item: BaseLiteLLMOpenAIResponseObject + + +class OpenAIChatCompletionLogprobsContentTopLogprobs(TypedDict, total=False): + bytes: List + logprob: Required[float] + token: Required[str] + + +class OpenAIChatCompletionLogprobsContent(TypedDict, total=False): + bytes: List + logprob: Required[float] + token: Required[str] + top_logprobs: List[OpenAIChatCompletionLogprobsContentTopLogprobs] class ContentPartAddedEvent(BaseLiteLLMOpenAIResponseObject): @@ -1194,7 +1207,31 @@ class ContentPartAddedEvent(BaseLiteLLMOpenAIResponseObject): item_id: str output_index: int content_index: int - part: dict + part: BaseLiteLLMOpenAIResponseObject + + +class ContentPartDonePartOutputText(BaseLiteLLMOpenAIResponseObject): + type: Literal["output_text"] + text: str + annotations: List[BaseLiteLLMOpenAIResponseObject] + logprobs: Optional[List[OpenAIChatCompletionLogprobsContent]] + + +class ContentPartDonePartRefusal(BaseLiteLLMOpenAIResponseObject): + type: Literal["refusal"] + refusal: str + + +class ContentPartDonePartReasoningText(BaseLiteLLMOpenAIResponseObject): + type: Literal["reasoning_text"] + reasoning: str + + +PART_UNION_TYPES = Union[ + ContentPartDonePartOutputText, + ContentPartDonePartRefusal, + ContentPartDonePartReasoningText, +] class ContentPartDoneEvent(BaseLiteLLMOpenAIResponseObject): @@ -1202,7 +1239,7 @@ class ContentPartDoneEvent(BaseLiteLLMOpenAIResponseObject): item_id: str output_index: int content_index: int - part: dict + part: PART_UNION_TYPES class OutputTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): @@ -1413,6 +1450,7 @@ ResponsesAPIStreamingResponse = Annotated[ ImageGenerationPartialImageEvent, ErrorEvent, GenericEvent, + BaseLiteLLMOpenAIResponseObject, ], Discriminator("type"), ] @@ -1731,19 +1769,6 @@ class OpenAIModerationResponse(BaseLiteLLMOpenAIResponseObject): _hidden_params: dict = PrivateAttr(default_factory=dict) -class OpenAIChatCompletionLogprobsContentTopLogprobs(TypedDict, total=False): - bytes: List - logprob: Required[float] - token: Required[str] - - -class OpenAIChatCompletionLogprobsContent(TypedDict, total=False): - bytes: List - logprob: Required[float] - token: Required[str] - top_logprobs: List[OpenAIChatCompletionLogprobsContentTopLogprobs] - - class OpenAIChatCompletionLogprobs(TypedDict, total=False): content: List[OpenAIChatCompletionLogprobsContent] refusal: List[OpenAIChatCompletionLogprobsContent] diff --git a/log.txt b/log.txt new file mode 100644 index 00000000000..a0f2afcdb18 --- /dev/null +++ b/log.txt @@ -0,0 +1,78 @@ +============================= test session starts ============================== +platform darwin -- Python 3.11.4, pytest-7.4.1, pluggy-1.2.0 +rootdir: /Users/krrishdholakia/Documents/litellm +plugins: snapshot-0.9.0, cov-5.0.0, timeout-2.2.0, postgresql-7.0.1, respx-0.21.1, asyncio-0.21.1, langsmith-0.3.4, anyio-4.8.0, mock-3.11.1, Faker-25.9.2 +asyncio: mode=Mode.STRICT +collected 1 item + +tests/llm_translation/test_gemini.py . [100%] + +=============================== warnings summary =============================== +tests/llm_translation/base_llm_unit_tests.py:481 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:481: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=6, delay=1) + +tests/llm_translation/base_llm_unit_tests.py:523 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:523: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=6, delay=1) + +tests/llm_translation/base_llm_unit_tests.py:601 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:601: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=6, delay=1) + +tests/llm_translation/base_llm_unit_tests.py:641 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:641: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=6, delay=1) + +tests/llm_translation/base_llm_unit_tests.py:650 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:650: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=6, delay=1) + +tests/llm_translation/base_llm_unit_tests.py:694 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:694: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=6, delay=1) + +tests/llm_translation/base_llm_unit_tests.py:745 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:745: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=6, delay=1) + +tests/llm_translation/base_llm_unit_tests.py:783 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:783: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=6, delay=1) + +tests/llm_translation/base_llm_unit_tests.py:859 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:859: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=4, delay=2) + +tests/llm_translation/base_llm_unit_tests.py:955 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:955: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=4, delay=1) + +tests/llm_translation/base_llm_unit_tests.py:1073 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:1073: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=3, delay=1) + +tests/llm_translation/base_llm_unit_tests.py:1109 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:1109: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=3, delay=1) + +tests/llm_translation/base_llm_unit_tests.py:1232 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:1232: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=3, delay=1) + +tests/llm_translation/test_gemini.py:36 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/test_gemini.py:36: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=3, delay=2) + +tests/llm_translation/test_gemini.py:510 + /Users/krrishdholakia/Documents/litellm/tests/llm_translation/test_gemini.py:510: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html + @pytest.mark.flaky(retries=3, delay=2) + +tests/llm_translation/test_gemini.py::test_gemini_image_generation_async + /Users/krrishdholakia/Library/Python/3.11/lib/python/site-packages/pydantic/main.py:463: UserWarning: Pydantic serializer warnings: + PydanticSerializationUnexpectedValue(Expected 10 fields but got 7: Expected `Message` - serialized value may not be as expected [input_value=Message(content="Here's t...er_specific_fields=None), input_type=Message]) + PydanticSerializationUnexpectedValue(Expected `StreamingChoices` - serialized value may not be as expected [input_value=Choices(finish_reason='st...r_specific_fields=None)), input_type=Choices]) + return self.__pydantic_serializer__.to_python( + +-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html +======================== 1 passed, 16 warnings in 5.86s ======================== diff --git a/tests/llm_translation/test_anthropic_completion.py b/tests/llm_translation/test_anthropic_completion.py index 4eeb80c4191..7d9ef56f164 100644 --- a/tests/llm_translation/test_anthropic_completion.py +++ b/tests/llm_translation/test_anthropic_completion.py @@ -361,6 +361,7 @@ def test_process_anthropic_headers_with_no_matching_headers(): def test_anthropic_tool_use(tool_type, tool_config, message_content): """Test Anthropic tool use with computer use and web fetch tools.""" from litellm import completion + litellm._turn_on_debug() tools = [tool_config] @@ -1518,3 +1519,126 @@ def test_anthropic_streaming(): role_set_count += 1 assert role_set_count == 1 + + +def test_anthropic_via_responses_api(): + from litellm.types.llms.openai import ResponsesAPIStreamEvents + + response = litellm.responses( + model="anthropic/claude-sonnet-4-5", + input="Who won the World Cup in 2022?", + max_output_tokens=100, + stream=True, + ) + + assert response is not None + + # Expected event sequence + expected_events = [ + ResponsesAPIStreamEvents.RESPONSE_CREATED, + ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, + ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + ResponsesAPIStreamEvents.CONTENT_PART_ADDED, + ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, # Can occur multiple times + ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + ResponsesAPIStreamEvents.CONTENT_PART_DONE, + ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + ] + + events_seen = [] + text_delta_count = 0 + + for chunk in response: + print(f"chunk: {chunk}") + + # Each chunk should have a type attribute + assert hasattr(chunk, "type"), f"Chunk missing 'type' attribute: {chunk}" + + event_type = chunk.type + + # Track events seen + if event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA: + text_delta_count += 1 + if ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA not in events_seen: + events_seen.append(event_type) + else: + events_seen.append(event_type) + + # Assert specific structures for each event type + if event_type == ResponsesAPIStreamEvents.RESPONSE_CREATED: + assert chunk.type == ResponsesAPIStreamEvents.RESPONSE_CREATED + assert hasattr(chunk, "response") + assert chunk.response.status == "in_progress" + assert hasattr(chunk.response, "id") + assert hasattr(chunk.response, "model") + + elif event_type == ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS: + assert chunk.type == ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS + assert hasattr(chunk, "response") + assert chunk.response.status == "in_progress" + + elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED: + assert chunk.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED + assert hasattr(chunk, "output_index") + assert hasattr(chunk, "item") + assert chunk.item.type == "message" + assert chunk.item.role == "assistant" + + elif event_type == ResponsesAPIStreamEvents.CONTENT_PART_ADDED: + assert chunk.type == ResponsesAPIStreamEvents.CONTENT_PART_ADDED + assert hasattr(chunk, "item_id") + assert hasattr(chunk, "output_index") + assert hasattr(chunk, "content_index") + assert hasattr(chunk, "part") + assert chunk.part.type == "output_text" + + elif event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA: + assert chunk.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA + assert hasattr(chunk, "item_id") + assert hasattr(chunk, "output_index") + assert hasattr(chunk, "content_index") + assert hasattr(chunk, "delta") + assert isinstance(chunk.delta, str) + + elif event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE: + assert chunk.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE + assert hasattr(chunk, "item_id") + assert hasattr(chunk, "output_index") + assert hasattr(chunk, "content_index") + assert hasattr(chunk, "text") + + elif event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE: + assert chunk.type == ResponsesAPIStreamEvents.CONTENT_PART_DONE + assert hasattr(chunk, "item_id") + assert hasattr(chunk, "output_index") + assert hasattr(chunk, "content_index") + assert hasattr(chunk, "part") + assert chunk.part.type == "output_text" + + elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE: + assert chunk.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE + assert hasattr(chunk, "output_index") + assert hasattr(chunk, "item") + assert chunk.item.status == "completed" + + elif event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: + assert chunk.type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert hasattr(chunk, "response") + assert chunk.response.status == "completed" + assert hasattr(chunk.response, "usage") + assert hasattr(chunk.response, "output") + + # Assert we saw all expected events + print(f"Events seen: {events_seen}") + assert ( + events_seen == expected_events + ), f"Event sequence mismatch. Expected: {expected_events}, Got: {events_seen}" + + # Assert we saw at least one text delta + assert ( + text_delta_count > 0 + ), f"Expected at least one response.output_text.delta event, got {text_delta_count}" + + print(f"✓ All {len(events_seen)} events matched expected structure") + print(f"✓ Received {text_delta_count} text delta chunks") diff --git a/tests/llm_translation/test_openai.py b/tests/llm_translation/test_openai.py index 4b63311f1a7..15bb5fec66d 100644 --- a/tests/llm_translation/test_openai.py +++ b/tests/llm_translation/test_openai.py @@ -728,22 +728,6 @@ def test_openai_safety_identifier_parameter_sync(): assert request_body["safety_identifier"] == "user_code_123456" -def test_gpt_5_reasoning(): - litellm._turn_on_debug() - response = litellm.completion( - model="openai/responses/gpt-5-mini", - messages=[ - { - "role": "user", - "content": "Think of the capital of France, and then write it.", - } - ], - reasoning_effort="low", - ) - print("response: ", response) - assert response.choices[0].message.reasoning_content is not None - - def test_gpt_5_reasoning_streaming(): litellm._turn_on_debug() response = litellm.completion( diff --git a/tests/pass_through_tests/test_openai_assistants_passthrough.py b/tests/pass_through_tests/test_openai_assistants_passthrough.py index e5783877ec0..28568005fd6 100644 --- a/tests/pass_through_tests/test_openai_assistants_passthrough.py +++ b/tests/pass_through_tests/test_openai_assistants_passthrough.py @@ -9,9 +9,12 @@ from openai import AssistantEventHandler client = openai.OpenAI(base_url="http://0.0.0.0:4000/openai", api_key="sk-1234") + def test_pass_through_file_operations(): # Create a temporary file - with tempfile.NamedTemporaryFile(mode='w+', suffix='.txt', delete=False) as temp_file: + with tempfile.NamedTemporaryFile( + mode="w+", suffix=".txt", delete=False + ) as temp_file: temp_file.write("This is a test file for the OpenAI Assistants API.") temp_file.flush() @@ -26,6 +29,7 @@ def test_pass_through_file_operations(): delete_file = client.files.delete(file.id) print("file deleted", delete_file) + def test_openai_assistants_e2e_operations(): assistant = client.beta.assistants.create( name="Math Tutor", @@ -98,13 +102,13 @@ def test_openai_assistants_e2e_operations_stream(): stream.until_done() - def test_azure_openai_assistants_e2e_operations_stream(): from openai import AzureOpenAI + client = AzureOpenAI( - base_url="http://0.0.0.0:4000/azure-config-passthrough/openai", + base_url="http://0.0.0.0:4000/azure-config-passthrough/openai", api_key="sk-1234", - api_version="2025-01-01-preview" + api_version="2025-01-01-preview", ) assistant = client.beta.assistants.create( name="Math Tutor", @@ -134,4 +138,4 @@ def test_azure_openai_assistants_e2e_operations_stream(): instructions="Please address the user as Jane Doe. The user has a premium account.", event_handler=EventHandler(), ) as stream: - stream.until_done() \ No newline at end of file + stream.until_done() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 0476f290a09..bfb2f4bf195 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1222,6 +1222,104 @@ class TestMCPServerManager: "Contact proxy admin to allow this tool" in exc_info.value.detail["error"] ) + @pytest.mark.asyncio + async def test_call_tool_without_broken_pipe_error(self): + """ + Test that call_tool properly uses async context manager to avoid broken pipe errors. + This test ensures that tasks are awaited INSIDE the context manager, keeping the connection alive. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from mcp.types import CallToolResult + + manager = MCPServerManager() + + # Create a test server + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + url="http://test-server.com", + ) + + # Register the server and map a tool to it + manager.registry = {"test-server": server} + manager.tool_name_to_mcp_server_name_mapping["test_tool"] = "test-server" + + # Create mock client that tracks context manager usage + mock_client = MagicMock() + context_entered = False + context_exited = False + call_tool_called_inside_context = False + + async def mock_aenter(self): + nonlocal context_entered + context_entered = True + return self + + async def mock_aexit(self, exc_type, exc_val, exc_tb): + nonlocal context_exited + context_exited = True + # Verify that call_tool was called before context exit + assert ( + call_tool_called_inside_context + ), "call_tool must be awaited inside context manager" + return False + + async def mock_call_tool(params): + nonlocal call_tool_called_inside_context + # Verify we're inside the context when this is called + assert context_entered, "call_tool called outside context manager" + assert not context_exited, "call_tool called after context exit" + call_tool_called_inside_context = True + + # Return a mock CallToolResult + result = MagicMock(spec=CallToolResult) + result.content = [{"type": "text", "text": "Tool executed successfully"}] + result.isError = False + return result + + mock_client.__aenter__ = mock_aenter + mock_client.__aexit__ = mock_aexit + mock_client.call_tool = mock_call_tool + + # Mock _create_mcp_client to return our mock client + manager._create_mcp_client = MagicMock(return_value=mock_client) + + # Mock user auth with no restrictions + user_api_key_auth = MagicMock() + user_api_key_auth.object_permission = None + user_api_key_auth.object_permission_id = None + + # Mock proxy logging + proxy_logging_obj = MagicMock() + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock( + return_value={} + ) + proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) + proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + + # Call the tool + result = await manager.call_tool( + name="test_tool", + arguments={"param": "value"}, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + ) + + # Verify the result + assert result is not None + assert result.isError is False + assert len(result.content) > 0 + + # Verify context manager was used properly + assert context_entered, "Context manager __aenter__ was not called" + assert context_exited, "Context manager __aexit__ was not called" + assert ( + call_tool_called_inside_context + ), "call_tool was not awaited inside context" + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 37cff2bd94a..0c72c8bb163 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -130,22 +130,24 @@ def test_virtual_key_allowed_routes_with_litellm_routes_member_name_denied(): assert "Only allowed to call routes: ['info_routes']" in str(exc_info.value) assert "Tried to call route: /chat/completions" in str(exc_info.value) -@pytest.mark.parametrize("route", [ - "/anthropic/v1/messages", - "/anthropic/v1/count_tokens", - "/gemini/v1/models", - "/gemini/countTokens", -]) + +@pytest.mark.parametrize( + "route", + [ + "/anthropic/v1/messages", + "/anthropic/v1/count_tokens", + "/gemini/v1/models", + "/gemini/countTokens", + ], +) def test_virtual_key_llm_api_route_includes_passthrough_prefix(route): """ Virtual key with llm_api_routes should allow passthrough routes like /anthropic/v1/messages - + Relevant issue: https://github.com/BerriAI/litellm/issues/14017 """ - valid_token = UserAPIKeyAuth( - user_id="test_user", allowed_routes=["llm_api_routes"] - ) + valid_token = UserAPIKeyAuth(user_id="test_user", allowed_routes=["llm_api_routes"]) result = RouteChecks.is_virtual_key_allowed_to_call_route( route=route, valid_token=valid_token @@ -232,7 +234,7 @@ def test_virtual_key_allowed_routes_with_no_member_names_only_explicit(): def test_anthropic_count_tokens_route_is_llm_api_route(): """Test that /v1/messages/count_tokens is recognized as an LLM API route for Anthropic""" - + # Test the core anthropic routes assert RouteChecks.is_llm_api_route("/v1/messages") is True assert RouteChecks.is_llm_api_route("/v1/messages/count_tokens") is True @@ -240,11 +242,11 @@ def test_anthropic_count_tokens_route_is_llm_api_route(): def test_anthropic_count_tokens_route_accessible_to_internal_users(): """Test that internal users can access the Anthropic count_tokens route""" - + # Test that the route is recognized as an LLM API route (which means it's accessible to internal users) # This is the core check that was failing in the original issue assert RouteChecks.is_llm_api_route("/v1/messages/count_tokens") is True - + # Also test that the regular messages route still works assert RouteChecks.is_llm_api_route("/v1/messages") is True @@ -252,7 +254,7 @@ def test_anthropic_count_tokens_route_accessible_to_internal_users(): def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints(): """ Test that virtual keys with llm_api_routes permission can access registered pass-through endpoints. - + This tests the scenario where a pass-through endpoint is registered from the DB (e.g., /azure-assistant) and a virtual key with llm_api_routes permission should be able to access both the exact path and subpaths (e.g., /azure-assistant/openai/assistants). @@ -272,7 +274,7 @@ def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints(): "type": "subpath", }, } - + with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", mock_registered_routes, @@ -282,21 +284,21 @@ def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints(): user_id="test_user", allowed_routes=["llm_api_routes"], ) - + # Test exact match for registered pass-through endpoint result1 = RouteChecks.is_virtual_key_allowed_to_call_route( route="/azure-assistant", valid_token=valid_token, ) assert result1 is True - + # Test subpath for registered pass-through endpoint with subpath type result2 = RouteChecks.is_virtual_key_allowed_to_call_route( route="/custom-endpoint/openai/assistants", valid_token=valid_token, ) assert result2 is True - + # Test exact match for subpath type result3 = RouteChecks.is_virtual_key_allowed_to_call_route( route="/custom-endpoint", @@ -319,7 +321,7 @@ def test_virtual_key_without_llm_api_routes_cannot_access_pass_through(): "type": "exact", }, } - + with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", mock_registered_routes, @@ -329,12 +331,289 @@ def test_virtual_key_without_llm_api_routes_cannot_access_pass_through(): user_id="test_user", allowed_routes=["info_routes"], ) - + # Test that access is denied with pytest.raises(Exception) as exc_info: RouteChecks.is_virtual_key_allowed_to_call_route( route="/azure-assistant", valid_token=valid_token, ) - + assert "Virtual key is not allowed to call this route" in str(exc_info.value) + + +def test_check_passthrough_route_access_key_metadata_exact_match(): + """Test that key metadata allowed_passthrough_routes allows exact match""" + + # Create a UserAPIKeyAuth with allowed_passthrough_routes in metadata + valid_token = UserAPIKeyAuth( + user_id="test_user", + metadata={"allowed_passthrough_routes": ["/custom-endpoint"]}, + ) + + # Test exact match + result = RouteChecks.check_passthrough_route_access( + route="/custom-endpoint", + user_api_key_dict=valid_token, + ) + + assert result is True + + +def test_check_passthrough_route_access_key_metadata_prefix_match(): + """Test that key metadata allowed_passthrough_routes allows prefix match""" + + # Create a UserAPIKeyAuth with allowed_passthrough_routes in metadata + valid_token = UserAPIKeyAuth( + user_id="test_user", + metadata={"allowed_passthrough_routes": ["/custom-endpoint"]}, + ) + + # Test prefix match + result = RouteChecks.check_passthrough_route_access( + route="/custom-endpoint/v1/chat/completions", + user_api_key_dict=valid_token, + ) + + assert result is True + + +def test_check_passthrough_route_access_key_metadata_no_match(): + """Test that key metadata allowed_passthrough_routes denies non-matching routes""" + + # Create a UserAPIKeyAuth with allowed_passthrough_routes in metadata + valid_token = UserAPIKeyAuth( + user_id="test_user", + metadata={"allowed_passthrough_routes": ["/custom-endpoint"]}, + ) + + # Test non-matching route + result = RouteChecks.check_passthrough_route_access( + route="/other-endpoint", + user_api_key_dict=valid_token, + ) + + assert result is False + + +def test_check_passthrough_route_access_team_metadata_exact_match(): + """Test that team metadata allowed_passthrough_routes allows exact match""" + + # Create a UserAPIKeyAuth with allowed_passthrough_routes in team_metadata + valid_token = UserAPIKeyAuth( + user_id="test_user", + metadata={}, + team_metadata={"allowed_passthrough_routes": ["/team-endpoint"]}, + ) + + # Test exact match + result = RouteChecks.check_passthrough_route_access( + route="/team-endpoint", + user_api_key_dict=valid_token, + ) + + assert result is True + + +def test_check_passthrough_route_access_team_metadata_prefix_match(): + """Test that team metadata allowed_passthrough_routes allows prefix match""" + + # Create a UserAPIKeyAuth with allowed_passthrough_routes in team_metadata + valid_token = UserAPIKeyAuth( + user_id="test_user", + metadata={}, + team_metadata={"allowed_passthrough_routes": ["/team-endpoint"]}, + ) + + # Test prefix match + result = RouteChecks.check_passthrough_route_access( + route="/team-endpoint/v1/messages", + user_api_key_dict=valid_token, + ) + + assert result is True + + +def test_check_passthrough_route_access_team_metadata_no_match(): + """Test that team metadata allowed_passthrough_routes denies non-matching routes""" + + # Create a UserAPIKeyAuth with allowed_passthrough_routes in team_metadata + valid_token = UserAPIKeyAuth( + user_id="test_user", + metadata={}, + team_metadata={"allowed_passthrough_routes": ["/team-endpoint"]}, + ) + + # Test non-matching route + result = RouteChecks.check_passthrough_route_access( + route="/other-endpoint", + user_api_key_dict=valid_token, + ) + + assert result is False + + +def test_check_passthrough_route_access_key_metadata_takes_precedence(): + """Test that key metadata takes precedence over team metadata""" + + # Create a UserAPIKeyAuth with different allowed_passthrough_routes in both metadata + valid_token = UserAPIKeyAuth( + user_id="test_user", + metadata={"allowed_passthrough_routes": ["/key-endpoint"]}, + team_metadata={"allowed_passthrough_routes": ["/team-endpoint"]}, + ) + + # Test that key endpoint is allowed + result1 = RouteChecks.check_passthrough_route_access( + route="/key-endpoint", + user_api_key_dict=valid_token, + ) + + # Test that team endpoint is NOT allowed (key metadata takes precedence) + result2 = RouteChecks.check_passthrough_route_access( + route="/team-endpoint", + user_api_key_dict=valid_token, + ) + + assert result1 is True + assert result2 is False + + +def test_check_passthrough_route_access_no_metadata(): + """Test that route is denied when metadata and team_metadata don't have allowed_passthrough_routes""" + + # Create a UserAPIKeyAuth without allowed_passthrough_routes + valid_token = UserAPIKeyAuth( + user_id="test_user", + ) + + # Test that route is denied + result = RouteChecks.check_passthrough_route_access( + route="/any-endpoint", + user_api_key_dict=valid_token, + ) + + assert result is False + + +def test_check_passthrough_route_access_no_allowed_passthrough_routes_key(): + """Test that route is denied when allowed_passthrough_routes is not in metadata""" + + # Create a UserAPIKeyAuth with metadata but no allowed_passthrough_routes + valid_token = UserAPIKeyAuth( + user_id="test_user", + metadata={"other_field": "value"}, + team_metadata={}, + ) + + # Test that route is denied + result = RouteChecks.check_passthrough_route_access( + route="/any-endpoint", + user_api_key_dict=valid_token, + ) + + assert result is False + + +def test_check_passthrough_route_access_allowed_passthrough_routes_is_none(): + """Test that route is denied when allowed_passthrough_routes is None""" + + # Create a UserAPIKeyAuth with allowed_passthrough_routes set to None + valid_token = UserAPIKeyAuth( + user_id="test_user", + metadata={"allowed_passthrough_routes": None}, + team_metadata={"allowed_passthrough_routes": None}, + ) + + # Test that route is denied + result = RouteChecks.check_passthrough_route_access( + route="/any-endpoint", + user_api_key_dict=valid_token, + ) + + assert result is False + + +def test_check_passthrough_route_access_multiple_routes(): + """Test that multiple allowed_passthrough_routes work correctly""" + + # Create a UserAPIKeyAuth with multiple allowed_passthrough_routes + valid_token = UserAPIKeyAuth( + user_id="test_user", + metadata={ + "allowed_passthrough_routes": [ + "/endpoint-1", + "/endpoint-2", + "/endpoint-3", + ] + }, + ) + + # Test that all allowed routes work + result1 = RouteChecks.check_passthrough_route_access( + route="/endpoint-1/v1/chat", + user_api_key_dict=valid_token, + ) + result2 = RouteChecks.check_passthrough_route_access( + route="/endpoint-2", + user_api_key_dict=valid_token, + ) + result3 = RouteChecks.check_passthrough_route_access( + route="/endpoint-3/completions", + user_api_key_dict=valid_token, + ) + + # Test that non-allowed route fails + result4 = RouteChecks.check_passthrough_route_access( + route="/endpoint-4", + user_api_key_dict=valid_token, + ) + + assert result1 is True + assert result2 is True + assert result3 is True + assert result4 is False + + +def test_check_passthrough_route_access_prevents_false_prefix_match(): + """Test that prefix matching doesn't allow false matches like /endpoint vs /endpoint-2""" + + # Create a UserAPIKeyAuth with allowed_passthrough_routes + valid_token = UserAPIKeyAuth( + user_id="test_user", + metadata={"allowed_passthrough_routes": ["/endpoint"]}, + ) + + # Test that /endpoint-2 is NOT allowed (not a valid prefix match) + result = RouteChecks.check_passthrough_route_access( + route="/endpoint-2", + user_api_key_dict=valid_token, + ) + + assert result is False + + # Test that /endpoint/something IS allowed (valid prefix match) + result2 = RouteChecks.check_passthrough_route_access( + route="/endpoint/something", + user_api_key_dict=valid_token, + ) + + assert result2 is True + + +def test_check_passthrough_route_access_empty_list(): + """Test that empty allowed_passthrough_routes list denies all routes""" + + # Create a UserAPIKeyAuth with empty allowed_passthrough_routes + valid_token = UserAPIKeyAuth( + user_id="test_user", + metadata={"allowed_passthrough_routes": []}, + ) + + # Test that route is denied + result = RouteChecks.check_passthrough_route_access( + route="/any-endpoint", + user_api_key_dict=valid_token, + ) + + assert result is False diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 1a26f8a39aa..ce90c1ed8ad 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -670,7 +670,7 @@ async def test_create_pass_through_route_with_cost_per_request(): # Create a proper UserAPIKeyAuth mock mock_user_api_key_dict = MagicMock() mock_user_api_key_dict.api_key = "test-key" - + await endpoint_func( request=mock_request, user_api_key_dict=mock_user_api_key_dict, @@ -747,12 +747,12 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): """ Test that pass_through_request (parent method) correctly includes proxy_server_request in kwargs passed to the success handler. - - Critical Test: Ensures that when pass_through_request is called, the kwargs passed to + + Critical Test: Ensures that when pass_through_request is called, the kwargs passed to downstream methods contain the proxy server request details (url, method, body). """ print("running test_pass_through_request_contains_proxy_server_request_in_kwargs") - + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.HttpPassThroughEndpointHelpers.non_streaming_http_request_handler" @@ -766,38 +766,44 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_response_body" ) as mock_get_response_body: - # Setup mock for pre_call_hook and post_call_failure_hook - mock_proxy_logging.pre_call_hook = AsyncMock(return_value={"test": "data"}) + # Setup mock for pre_call_hook and post_call_failure_hook + mock_proxy_logging.pre_call_hook = AsyncMock( + return_value={"test": "data"} + ) mock_proxy_logging.post_call_failure_hook = AsyncMock() - + # Setup mock for http response mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {} - mock_response.aread = AsyncMock(return_value=b'{"success": true}') + mock_response.aread = AsyncMock( + return_value=b'{"success": true}' + ) mock_response.text = '{"success": true}' mock_response.raise_for_status = MagicMock() - + # Mock the HTTP request handler directly mock_http_handler.return_value = mock_response - + # Mock response body parser mock_get_response_body.return_value = {"success": True} - + # Mock headers for custom headers mock_processing.get_custom_headers.return_value = {} - + # Mock success handler to capture kwargs mock_success_handler.return_value = None - + # Create mock request mock_request = MagicMock(spec=Request) mock_request.method = "POST" mock_request.url = "http://test-proxy.com/api/endpoint" - mock_request.body = AsyncMock(return_value=b'{"message": "test request"}') + mock_request.body = AsyncMock( + return_value=b'{"message": "test request"}' + ) mock_request.headers = Headers({}) mock_request.query_params = QueryParams({}) - + # Create mock user API key dict mock_user_api_key_dict = MagicMock() mock_user_api_key_dict.api_key = "test-api-key" @@ -809,7 +815,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): mock_user_api_key_dict.team_alias = "test-team-alias" mock_user_api_key_dict.end_user_id = "test-end-user-id" mock_user_api_key_dict.request_route = "/api/endpoint" - + # Call pass_through_request (the parent method) result = await pass_through_request( request=mock_request, @@ -817,38 +823,38 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): custom_headers={"X-Custom": "header"}, user_api_key_dict=mock_user_api_key_dict, ) - + # Verify the success handler was called mock_success_handler.assert_called_once() - + # Extract the kwargs passed to the success handler call_kwargs = mock_success_handler.call_args[1] - + # Verify that litellm_params exists in kwargs assert "litellm_params" in call_kwargs litellm_params = call_kwargs["litellm_params"] - + # Verify that proxy_server_request exists in litellm_params assert "proxy_server_request" in litellm_params proxy_server_request = litellm_params["proxy_server_request"] - + # Verify the proxy_server_request contains expected fields assert "url" in proxy_server_request assert "method" in proxy_server_request assert "body" in proxy_server_request - + # Verify the values match the original request assert proxy_server_request["url"] == str(mock_request.url) assert proxy_server_request["method"] == mock_request.method # The body should be the value returned by pre_call_hook, not the original request body assert proxy_server_request["body"] == {"test": "data"} - + # Verify other required kwargs are present assert "call_type" in call_kwargs assert call_kwargs["call_type"] == "pass_through_endpoint" assert "litellm_call_id" in call_kwargs assert "passthrough_logging_payload" in call_kwargs - + # Verify metadata contains user information assert "metadata" in litellm_params metadata = litellm_params["metadata"] @@ -862,7 +868,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): async def test_create_pass_through_endpoint(): """ Test creating a new pass-through endpoint - + This test verifies that the create_pass_through_endpoints function: 1. Accepts a PassThroughGenericEndpoint object 2. Auto-generates an ID if not provided @@ -881,36 +887,38 @@ async def test_create_pass_through_endpoint(): ) # Mock the database functions - with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: - with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config: + with patch( + "litellm.proxy.proxy_server.get_config_general_settings" + ) as mock_get_config: + with patch( + "litellm.proxy.proxy_server.update_config_general_settings" + ) as mock_update_config: # Mock existing config (empty list) mock_get_config.return_value = ConfigFieldInfo( - field_name="pass_through_endpoints", - field_value=[] + field_name="pass_through_endpoints", field_value=[] ) - + # Create test endpoint data test_endpoint = PassThroughGenericEndpoint( path="/test/endpoint", target="http://example.com/api", headers={"Authorization": "Bearer test-token"}, include_subpath=True, - cost_per_request=0.50 + cost_per_request=0.50, ) - + # Mock user API key dict mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) - + # Call the create function result = await create_pass_through_endpoints( - data=test_endpoint, - user_api_key_dict=mock_user_api_key_dict + data=test_endpoint, user_api_key_dict=mock_user_api_key_dict ) - + # Verify the result assert isinstance(result, PassThroughEndpointResponse) assert len(result.endpoints) == 1 - + created_endpoint = result.endpoints[0] assert created_endpoint.path == "/test/endpoint" assert created_endpoint.target == "http://example.com/api" @@ -918,13 +926,13 @@ async def test_create_pass_through_endpoint(): assert created_endpoint.include_subpath is True assert created_endpoint.cost_per_request == 0.50 assert created_endpoint.id is not None # Should be auto-generated - + # Verify database calls mock_get_config.assert_called_once_with( field_name="pass_through_endpoints", - user_api_key_dict=mock_user_api_key_dict + user_api_key_dict=mock_user_api_key_dict, ) - + mock_update_config.assert_called_once() update_call_args = mock_update_config.call_args[1] assert update_call_args["data"].field_name == "pass_through_endpoints" @@ -937,7 +945,7 @@ async def test_create_pass_through_endpoint(): async def test_update_pass_through_endpoint(): """ Test updating an existing pass-through endpoint - + This test verifies that the update_pass_through_endpoints function: 1. Finds the existing endpoint by ID 2. Updates only the provided fields (partial updates) @@ -957,8 +965,12 @@ async def test_update_pass_through_endpoint(): ) # Mock the database functions - with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: - with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config: + with patch( + "litellm.proxy.proxy_server.get_config_general_settings" + ) as mock_get_config: + with patch( + "litellm.proxy.proxy_server.update_config_general_settings" + ) as mock_update_config: # Create existing endpoint data existing_endpoint_id = "test-endpoint-123" existing_endpoints = [ @@ -968,53 +980,58 @@ async def test_update_pass_through_endpoint(): "target": "http://example.com/api", "headers": {"Authorization": "Bearer old-token"}, "include_subpath": False, - "cost_per_request": 0.25 + "cost_per_request": 0.25, } ] - + # Mock existing config mock_get_config.return_value = ConfigFieldInfo( - field_name="pass_through_endpoints", - field_value=existing_endpoints + field_name="pass_through_endpoints", field_value=existing_endpoints ) - + # Create update data (partial update) update_data = PassThroughGenericEndpoint( path="/test/endpoint", # Keep same path target="http://newapi.com/v2", # Update target - headers={"Authorization": "Bearer new-token", "X-Custom": "header"}, # Update headers - cost_per_request=0.75 # Update cost + headers={ + "Authorization": "Bearer new-token", + "X-Custom": "header", + }, # Update headers + cost_per_request=0.75, # Update cost # include_subpath not provided - should preserve existing value ) - + # Mock user API key dict mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) - + # Call the update function result = await update_pass_through_endpoints( endpoint_id=existing_endpoint_id, data=update_data, - user_api_key_dict=mock_user_api_key_dict + user_api_key_dict=mock_user_api_key_dict, ) - + # Verify the result assert isinstance(result, PassThroughEndpointResponse) assert len(result.endpoints) == 1 - + updated_endpoint = result.endpoints[0] assert updated_endpoint.id == existing_endpoint_id # ID preserved assert updated_endpoint.path == "/test/endpoint" assert updated_endpoint.target == "http://newapi.com/v2" # Updated - assert updated_endpoint.headers == {"Authorization": "Bearer new-token", "X-Custom": "header"} # Updated + assert updated_endpoint.headers == { + "Authorization": "Bearer new-token", + "X-Custom": "header", + } # Updated assert updated_endpoint.include_subpath is False # Preserved existing value assert updated_endpoint.cost_per_request == 0.75 # Updated - + # Verify database calls mock_get_config.assert_called_once_with( field_name="pass_through_endpoints", - user_api_key_dict=mock_user_api_key_dict + user_api_key_dict=mock_user_api_key_dict, ) - + mock_update_config.assert_called_once() update_call_args = mock_update_config.call_args[1] assert update_call_args["data"].field_name == "pass_through_endpoints" @@ -1042,7 +1059,9 @@ async def test_update_pass_through_endpoint_not_found(): ) # Mock the database functions - with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: + with patch( + "litellm.proxy.proxy_server.get_config_general_settings" + ) as mock_get_config: # Mock existing config with different endpoint existing_endpoints = [ { @@ -1051,32 +1070,30 @@ async def test_update_pass_through_endpoint_not_found(): "target": "http://different.com/api", "headers": {}, "include_subpath": False, - "cost_per_request": 0.0 + "cost_per_request": 0.0, } ] - + mock_get_config.return_value = ConfigFieldInfo( - field_name="pass_through_endpoints", - field_value=existing_endpoints + field_name="pass_through_endpoints", field_value=existing_endpoints ) - + # Create update data update_data = PassThroughGenericEndpoint( - path="/test/endpoint", - target="http://newapi.com/v2" + path="/test/endpoint", target="http://newapi.com/v2" ) - + # Mock user API key dict mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) - + # Call the update function with non-existent ID with pytest.raises(HTTPException) as exc_info: await update_pass_through_endpoints( endpoint_id="non-existent-endpoint-123", data=update_data, - user_api_key_dict=mock_user_api_key_dict + user_api_key_dict=mock_user_api_key_dict, ) - + # Verify the exception assert exc_info.value.status_code == 404 assert "not found" in str(exc_info.value.detail).lower() @@ -1086,7 +1103,7 @@ async def test_update_pass_through_endpoint_not_found(): async def test_delete_pass_through_endpoint(): """ Test deleting an existing pass-through endpoint - + This test verifies that the delete_pass_through_endpoints function: 1. Finds the existing endpoint by ID 2. Removes it from the database @@ -1103,8 +1120,12 @@ async def test_delete_pass_through_endpoint(): ) # Mock the database functions - with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: - with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config: + with patch( + "litellm.proxy.proxy_server.get_config_general_settings" + ) as mock_get_config: + with patch( + "litellm.proxy.proxy_server.update_config_general_settings" + ) as mock_update_config: # Create existing endpoint data endpoint_to_delete_id = "test-endpoint-123" other_endpoint_id = "other-endpoint-456" @@ -1115,7 +1136,7 @@ async def test_delete_pass_through_endpoint(): "target": "http://example.com/api", "headers": {"Authorization": "Bearer test-token"}, "include_subpath": True, - "cost_per_request": 0.50 + "cost_per_request": 0.50, }, { "id": other_endpoint_id, @@ -1123,29 +1144,28 @@ async def test_delete_pass_through_endpoint(): "target": "http://other.com/api", "headers": {}, "include_subpath": False, - "cost_per_request": 0.25 - } + "cost_per_request": 0.25, + }, ] - + # Mock existing config mock_get_config.return_value = ConfigFieldInfo( - field_name="pass_through_endpoints", - field_value=existing_endpoints + field_name="pass_through_endpoints", field_value=existing_endpoints ) - + # Mock user API key dict mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) - + # Call the delete function result = await delete_pass_through_endpoints( endpoint_id=endpoint_to_delete_id, - user_api_key_dict=mock_user_api_key_dict + user_api_key_dict=mock_user_api_key_dict, ) - + # Verify the result assert isinstance(result, PassThroughEndpointResponse) assert len(result.endpoints) == 1 - + deleted_endpoint = result.endpoints[0] assert deleted_endpoint.id == endpoint_to_delete_id assert deleted_endpoint.path == "/test/endpoint" @@ -1153,13 +1173,13 @@ async def test_delete_pass_through_endpoint(): assert deleted_endpoint.headers == {"Authorization": "Bearer test-token"} assert deleted_endpoint.include_subpath is True assert deleted_endpoint.cost_per_request == 0.50 - + # Verify database calls mock_get_config.assert_called_once_with( field_name="pass_through_endpoints", - user_api_key_dict=mock_user_api_key_dict + user_api_key_dict=mock_user_api_key_dict, ) - + mock_update_config.assert_called_once() update_call_args = mock_update_config.call_args[1] assert update_call_args["data"].field_name == "pass_through_endpoints" @@ -1183,7 +1203,9 @@ async def test_delete_pass_through_endpoint_not_found(): ) # Mock the database functions - with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: + with patch( + "litellm.proxy.proxy_server.get_config_general_settings" + ) as mock_get_config: # Mock existing config with different endpoint existing_endpoints = [ { @@ -1192,25 +1214,24 @@ async def test_delete_pass_through_endpoint_not_found(): "target": "http://different.com/api", "headers": {}, "include_subpath": False, - "cost_per_request": 0.0 + "cost_per_request": 0.0, } ] - + mock_get_config.return_value = ConfigFieldInfo( - field_name="pass_through_endpoints", - field_value=existing_endpoints + field_name="pass_through_endpoints", field_value=existing_endpoints ) - + # Mock user API key dict mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) - + # Call the delete function with non-existent ID with pytest.raises(HTTPException) as exc_info: await delete_pass_through_endpoints( endpoint_id="non-existent-endpoint-123", - user_api_key_dict=mock_user_api_key_dict + user_api_key_dict=mock_user_api_key_dict, ) - + # Verify the exception assert exc_info.value.status_code == 400 assert "not found" in str(exc_info.value.detail).lower() @@ -1229,34 +1250,33 @@ async def test_delete_pass_through_endpoint_empty_list(): ) # Mock the database functions - with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: + with patch( + "litellm.proxy.proxy_server.get_config_general_settings" + ) as mock_get_config: # Mock empty config mock_get_config.return_value = ConfigFieldInfo( - field_name="pass_through_endpoints", - field_value=None + field_name="pass_through_endpoints", field_value=None ) - + # Mock user API key dict mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) - + # Call the delete function with pytest.raises(HTTPException) as exc_info: await delete_pass_through_endpoints( - endpoint_id="any-endpoint-123", - user_api_key_dict=mock_user_api_key_dict + endpoint_id="any-endpoint-123", user_api_key_dict=mock_user_api_key_dict ) - + # Verify the exception assert exc_info.value.status_code == 400 assert "no pass-through endpoints setup" in str(exc_info.value.detail).lower() - @pytest.mark.asyncio async def test_pass_through_request_query_params_forwarding(): """ Test that query parameters from the original request are properly forwarded to the target URL. - + This test verifies the fix for the bug where query parameters like api-version were being lost when forwarding requests to Azure OpenAI and other pass-through endpoints. """ @@ -1275,42 +1295,57 @@ async def test_pass_through_request_query_params_forwarding(): ) as mock_get_response_body: # Setup mock for pre_call_hook test_body = {"name": "Azure Assistant", "model": "gpt-4o"} - mock_proxy_logging.pre_call_hook = AsyncMock(return_value=test_body) - + mock_proxy_logging.pre_call_hook = AsyncMock( + return_value=test_body + ) + # Setup mock for http response mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} - mock_response.aread = AsyncMock(return_value=b'{"id": "asst_123", "object": "assistant"}') + mock_response.aread = AsyncMock( + return_value=b'{"id": "asst_123", "object": "assistant"}' + ) mock_response.text = '{"id": "asst_123", "object": "assistant"}' mock_response.raise_for_status = MagicMock() - + # Mock the HTTP request handler to capture the call mock_http_handler.return_value = mock_response - + # Mock response body parser - mock_get_response_body.return_value = {"id": "asst_123", "object": "assistant"} - + mock_get_response_body.return_value = { + "id": "asst_123", + "object": "assistant", + } + # Mock headers for custom headers mock_processing.get_custom_headers.return_value = {} - + # Mock success handler mock_success_handler.return_value = None - + # Create mock request with query parameters (Azure API version) mock_request = MagicMock(spec=Request) mock_request.method = "POST" - mock_request.url = "http://localhost:4000/azure-assistant/openai/assistants" - mock_request.body = AsyncMock(return_value=json.dumps(test_body).encode()) - mock_request.headers = Headers({"Content-Type": "application/json"}) - + mock_request.url = ( + "http://localhost:4000/azure-assistant/openai/assistants" + ) + mock_request.body = AsyncMock( + return_value=json.dumps(test_body).encode() + ) + mock_request.headers = Headers( + {"Content-Type": "application/json"} + ) + # Create QueryParams with api-version parameter - mock_request.query_params = QueryParams([("api-version", "2025-01-01-preview")]) - + mock_request.query_params = QueryParams( + [("api-version", "2025-01-01-preview")] + ) + # Create mock user API key dict mock_user_api_key_dict = MagicMock() mock_user_api_key_dict.api_key = "sk-1234" - + # Call pass_through_request result = await pass_through_request( request=mock_request, @@ -1318,20 +1353,25 @@ async def test_pass_through_request_query_params_forwarding(): custom_headers={"Authorization": "Bearer azure_token"}, user_api_key_dict=mock_user_api_key_dict, ) - + # Verify the HTTP handler was called mock_http_handler.assert_called_once() - + # Extract the call arguments to verify query parameters were passed call_kwargs = mock_http_handler.call_args[1] - + # The key assertion: query parameters should be preserved and passed to the HTTP handler assert "requested_query_params" in call_kwargs - assert call_kwargs["requested_query_params"] == {"api-version": "2025-01-01-preview"} - + assert call_kwargs["requested_query_params"] == { + "api-version": "2025-01-01-preview" + } + # Verify the target URL is correct - assert str(call_kwargs["url"]) == "https://krris-m2f9a9i7-eastus2.openai.azure.com/openai/assistants" - + assert ( + str(call_kwargs["url"]) + == "https://krris-m2f9a9i7-eastus2.openai.azure.com/openai/assistants" + ) + # Verify the request body is preserved assert call_kwargs["_parsed_body"] == test_body @@ -1356,15 +1396,16 @@ async def test_pass_through_with_httpbin_redirect(): mock_request.method = "GET" mock_request.headers = Headers({}) mock_request.query_params = QueryParams("") - + # Mock the body method to return empty bytes for GET request async def mock_body(): return b"" + mock_request.body = mock_body - + # Mock user API key dict mock_user_api_key_dict = MagicMock() - + try: # Test with httpbin.org redirect endpoint # This will redirect to httpbin.org/get @@ -1372,19 +1413,283 @@ async def test_pass_through_with_httpbin_redirect(): request=mock_request, target="https://httpbin.org/redirect/1", custom_headers={}, - user_api_key_dict=mock_user_api_key_dict + user_api_key_dict=mock_user_api_key_dict, ) - + # Should get the final response (200) from /get endpoint, not the redirect (302) assert response.status_code == 200 - + # The response should be from the /get endpoint - response_content = response.body.decode('utf-8') - + response_content = response.body.decode("utf-8") + # httpbin.org/get returns JSON with info about the request assert '"url": "https://httpbin.org/get"' in response_content print("GOT A Response from HTTPBIN=", response_content) except Exception as e: # If httpbin.org is not accessible, skip the test import pytest + pytest.skip(f"Could not reach httpbin.org for integration test: {e}") + + +@pytest.mark.asyncio +async def test_filter_endpoints_by_team_allowed_routes_with_filter(): + """ + Test that _filter_endpoints_by_team_allowed_routes correctly filters endpoints + when team has allowed_passthrough_routes in metadata + """ + from litellm.proxy._types import PassThroughGenericEndpoint + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _filter_endpoints_by_team_allowed_routes, + ) + + # Create test endpoints + endpoints = [ + PassThroughGenericEndpoint( + id="endpoint-1", path="/api/allowed1", target="http://example.com/api1" + ), + PassThroughGenericEndpoint( + id="endpoint-2", path="/api/allowed2", target="http://example.com/api2" + ), + PassThroughGenericEndpoint( + id="endpoint-3", path="/api/notallowed", target="http://example.com/api3" + ), + ] + + # Mock prisma client + mock_prisma_client = MagicMock() + mock_team = MagicMock() + mock_team.metadata = { + "allowed_passthrough_routes": ["/api/allowed1", "/api/allowed2"] + } + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_team + ) + + # Call the function + result = await _filter_endpoints_by_team_allowed_routes( + team_id="test-team-123", + pass_through_endpoints=endpoints, + prisma_client=mock_prisma_client, + ) + + # Should only return allowed endpoints + assert len(result) == 2 + assert result[0].path == "/api/allowed1" + assert result[1].path == "/api/allowed2" + + # Verify database call + mock_prisma_client.db.litellm_teamtable.find_unique.assert_called_once_with( + where={"team_id": "test-team-123"} + ) + + +@pytest.mark.asyncio +async def test_filter_endpoints_by_team_allowed_routes_team_not_found(): + """ + Test that _filter_endpoints_by_team_allowed_routes raises HTTPException + when team is not found + """ + from fastapi import HTTPException + + from litellm.proxy._types import PassThroughGenericEndpoint + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _filter_endpoints_by_team_allowed_routes, + ) + + # Create test endpoints + endpoints = [ + PassThroughGenericEndpoint( + id="endpoint-1", path="/api/test", target="http://example.com/api" + ), + ] + + # Mock prisma client to return None (team not found) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + + # Call the function and expect HTTPException + with pytest.raises(HTTPException) as exc_info: + await _filter_endpoints_by_team_allowed_routes( + team_id="non-existent-team", + pass_through_endpoints=endpoints, + prisma_client=mock_prisma_client, + ) + + # Verify the exception + assert exc_info.value.status_code == 404 + assert "Team not found" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_filter_endpoints_by_team_allowed_routes_no_metadata(): + """ + Test that _filter_endpoints_by_team_allowed_routes returns all endpoints + when team has no metadata + """ + from litellm.proxy._types import PassThroughGenericEndpoint + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _filter_endpoints_by_team_allowed_routes, + ) + + # Create test endpoints + endpoints = [ + PassThroughGenericEndpoint( + id="endpoint-1", path="/api/test1", target="http://example.com/api1" + ), + PassThroughGenericEndpoint( + id="endpoint-2", path="/api/test2", target="http://example.com/api2" + ), + ] + + # Mock prisma client with team that has None metadata + mock_prisma_client = MagicMock() + mock_team = MagicMock() + mock_team.metadata = None + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_team + ) + + # Call the function + result = await _filter_endpoints_by_team_allowed_routes( + team_id="test-team-123", + pass_through_endpoints=endpoints, + prisma_client=mock_prisma_client, + ) + + # Should return all endpoints when no metadata + assert len(result) == 2 + assert result[0].path == "/api/test1" + assert result[1].path == "/api/test2" + + +@pytest.mark.asyncio +async def test_filter_endpoints_by_team_allowed_routes_no_allowed_routes_key(): + """ + Test that _filter_endpoints_by_team_allowed_routes returns all endpoints + when team metadata doesn't have allowed_passthrough_routes key + """ + from litellm.proxy._types import PassThroughGenericEndpoint + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _filter_endpoints_by_team_allowed_routes, + ) + + # Create test endpoints + endpoints = [ + PassThroughGenericEndpoint( + id="endpoint-1", path="/api/test1", target="http://example.com/api1" + ), + PassThroughGenericEndpoint( + id="endpoint-2", path="/api/test2", target="http://example.com/api2" + ), + ] + + # Mock prisma client with team that has metadata but no allowed_passthrough_routes + mock_prisma_client = MagicMock() + mock_team = MagicMock() + mock_team.metadata = {"some_other_key": "some_value"} + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_team + ) + + # Call the function + result = await _filter_endpoints_by_team_allowed_routes( + team_id="test-team-123", + pass_through_endpoints=endpoints, + prisma_client=mock_prisma_client, + ) + + # Should return all endpoints when allowed_passthrough_routes key doesn't exist + assert len(result) == 2 + assert result[0].path == "/api/test1" + assert result[1].path == "/api/test2" + + +@pytest.mark.asyncio +async def test_filter_endpoints_by_team_allowed_routes_empty_allowed_list(): + """ + Test that _filter_endpoints_by_team_allowed_routes returns empty list + when team has empty allowed_passthrough_routes list + """ + from litellm.proxy._types import PassThroughGenericEndpoint + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _filter_endpoints_by_team_allowed_routes, + ) + + # Create test endpoints + endpoints = [ + PassThroughGenericEndpoint( + id="endpoint-1", path="/api/test1", target="http://example.com/api1" + ), + PassThroughGenericEndpoint( + id="endpoint-2", path="/api/test2", target="http://example.com/api2" + ), + ] + + # Mock prisma client with team that has empty allowed_passthrough_routes + mock_prisma_client = MagicMock() + mock_team = MagicMock() + mock_team.metadata = {"allowed_passthrough_routes": []} + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_team + ) + + # Call the function + result = await _filter_endpoints_by_team_allowed_routes( + team_id="test-team-123", + pass_through_endpoints=endpoints, + prisma_client=mock_prisma_client, + ) + + # Should return empty list when allowed_passthrough_routes is empty + assert len(result) == 0 + + +@pytest.mark.asyncio +async def test_filter_endpoints_by_team_allowed_routes_partial_match(): + """ + Test that _filter_endpoints_by_team_allowed_routes correctly filters + when only some endpoints match allowed routes + """ + from litellm.proxy._types import PassThroughGenericEndpoint + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _filter_endpoints_by_team_allowed_routes, + ) + + # Create test endpoints + endpoints = [ + PassThroughGenericEndpoint( + id="endpoint-1", path="/api/openai", target="http://example.com/openai" + ), + PassThroughGenericEndpoint( + id="endpoint-2", + path="/api/anthropic", + target="http://example.com/anthropic", + ), + PassThroughGenericEndpoint( + id="endpoint-3", path="/api/azure", target="http://example.com/azure" + ), + PassThroughGenericEndpoint( + id="endpoint-4", path="/api/cohere", target="http://example.com/cohere" + ), + ] + + # Mock prisma client with team that allows only 2 routes + mock_prisma_client = MagicMock() + mock_team = MagicMock() + mock_team.metadata = {"allowed_passthrough_routes": ["/api/openai", "/api/azure"]} + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_team + ) + + # Call the function + result = await _filter_endpoints_by_team_allowed_routes( + team_id="test-team-123", + pass_through_endpoints=endpoints, + prisma_client=mock_prisma_client, + ) + + # Should return only the 2 allowed endpoints + assert len(result) == 2 + assert result[0].path == "/api/openai" + assert result[1].path == "/api/azure" diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py index 5323589818b..020b5de0a2a 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py @@ -4,14 +4,20 @@ Test reasoning content preservation in Responses API transformation from unittest.mock import AsyncMock -from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) -from litellm.types.utils import ModelResponse, Choices, Message +from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) class TestReasoningContentStreaming: @@ -41,6 +47,7 @@ class TestReasoningContentStreaming: mock_stream = AsyncMock() iterator = LiteLLMCompletionStreamingIterator( + model="test-model", litellm_custom_stream_wrapper=mock_stream, request_input="Test input", responses_api_request={}, @@ -78,6 +85,7 @@ class TestReasoningContentStreaming: mock_stream = AsyncMock() iterator = LiteLLMCompletionStreamingIterator( + model="test-model", litellm_custom_stream_wrapper=mock_stream, request_input="Test input", responses_api_request={}, @@ -114,6 +122,7 @@ class TestReasoningContentStreaming: mock_stream = AsyncMock() iterator = LiteLLMCompletionStreamingIterator( + model="test-model", litellm_custom_stream_wrapper=mock_stream, request_input="Test input", responses_api_request={}, @@ -272,6 +281,7 @@ def test_streaming_chunk_id_raw(): ) iterator = LiteLLMCompletionStreamingIterator( + model="test-model", litellm_custom_stream_wrapper=AsyncMock(), request_input="Test input", responses_api_request={}, diff --git a/ui/litellm-dashboard/src/components/common_components/PassThroughRoutesSelector.tsx b/ui/litellm-dashboard/src/components/common_components/PassThroughRoutesSelector.tsx new file mode 100644 index 00000000000..116e3f6aa39 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/PassThroughRoutesSelector.tsx @@ -0,0 +1,69 @@ +import React, { useEffect, useState } from "react"; +import { Select } from "antd"; +import { getPassThroughEndpointsCall } from "../networking"; + +interface PassThroughRoutesSelectorProps { + onChange: (selectedRoutes: string[]) => void; + value?: string[]; + className?: string; + accessToken: string; + placeholder?: string; + disabled?: boolean; + teamId?: string | null; +} + +const PassThroughRoutesSelector: React.FC = ({ + onChange, + value, + className, + accessToken, + placeholder = "Select pass through routes", + disabled = false, + teamId, +}) => { + const [passThroughRoutes, setPassThroughRoutes] = useState([]); + const [loading, setLoading] = useState(false); + + useEffect(() => { + const fetchPassThroughRoutes = async () => { + if (!accessToken) return; + + setLoading(true); + try { + const response = await getPassThroughEndpointsCall(accessToken, teamId); + if (response.endpoints) { + const routes = response.endpoints.map((route: { path: string }) => route.path); + setPassThroughRoutes(routes); + } + } catch (error) { + console.error("Error fetching pass through routes:", error); + } finally { + setLoading(false); + } + }; + + fetchPassThroughRoutes(); + }, [accessToken, teamId]); + + return ( +