From 342e80d48d030c8d20748b7f33fd71990e5364c8 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Sep 2025 15:25:42 -0700 Subject: [PATCH 01/76] feat: initial commit for v2 oauth flow --- .../mcp_server/auth/user_api_key_auth_mcp.py | 27 +-- .../mcp_server/discoverable_endpoints.py | 204 ++++++++++++++++++ .../mcp_server/mcp_server_manager.py | 58 +++-- .../mcp_server/rest_endpoints.py | 119 +++++++++- .../proxy/_experimental/mcp_server/server.py | 50 +++-- .../index.html} | 0 .../proxy/_experimental/out/onboarding.html | 1 - litellm/proxy/_new_secret_config.yaml | 13 ++ litellm/proxy/auth/user_api_key_auth.py | 47 ++-- litellm/proxy/proxy_server.py | 83 ++++++- litellm/types/mcp.py | 20 +- .../types/mcp_server/mcp_server_manager.py | 7 +- 12 files changed, 544 insertions(+), 85 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py rename litellm/proxy/_experimental/out/{model_hub_table.html => model_hub_table/index.html} (100%) delete mode 100644 litellm/proxy/_experimental/out/onboarding.html diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 6afed97fe9..4e7745392b 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Tuple, Dict, Set +from typing import Dict, List, Optional, Set, Tuple from starlette.datastructures import Headers from starlette.requests import Request @@ -96,9 +96,12 @@ class MCPRequestHandler: return b"{}" request.body = mock_body # type: ignore - validated_user_api_key_auth = await user_api_key_auth( - api_key=litellm_api_key, request=request - ) + if ".well-known" in str(request.url): # public routes + validated_user_api_key_auth = UserAPIKeyAuth() + else: + validated_user_api_key_auth = await user_api_key_auth( + api_key=litellm_api_key, request=request + ) return ( validated_user_api_key_auth, mcp_auth_header, @@ -359,10 +362,10 @@ class MCPRequestHandler: return [] try: - team_obj: Optional[ - LiteLLM_TeamTable - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": user_api_key_auth.team_id}, + team_obj: Optional[LiteLLM_TeamTable] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": user_api_key_auth.team_id}, + ) ) if team_obj is None: verbose_logger.debug("team_obj is None") @@ -535,10 +538,10 @@ class MCPRequestHandler: verbose_logger.debug("prisma_client is None") return [] - team_obj: Optional[ - LiteLLM_TeamTable - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": user_api_key_auth.team_id}, + team_obj: Optional[LiteLLM_TeamTable] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": user_api_key_auth.team_id}, + ) ) if team_obj is None: verbose_logger.debug("team_obj is None") diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py new file mode 100644 index 0000000000..d38edf5dae --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -0,0 +1,204 @@ +from typing import Optional +from urllib.parse import urlencode, urlparse, urlunparse + +import httpx +from fastapi import APIRouter, Form, HTTPException, Request +from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse + +STATE_MAP = {} + +router = APIRouter( + tags=["mcp"], +) + + +@router.get("/authorize") +async def authorize( + request: Request, client_id: str, redirect_uri: str, state: str = "" +): + # Redirect to real GitHub OAuth + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + mcp_server = global_mcp_server_manager.get_mcp_server_by_name(client_id) + if mcp_server is None: + raise HTTPException(status_code=404, detail="MCP server not found") + if mcp_server.auth_type != "oauth2": + raise HTTPException(status_code=400, detail="MCP server is not OAuth2") + if mcp_server.client_id is None: + raise HTTPException(status_code=400, detail="MCP server client id is not set") + if mcp_server.authorization_url is None: + raise HTTPException( + status_code=400, detail="MCP server authorization url is not set" + ) + if mcp_server.scopes is None: + raise HTTPException(status_code=400, detail="MCP server scopes is not set") + + # Parse it to remove any existing query + parsed = urlparse(redirect_uri) + base_url = urlunparse(parsed._replace(query="")) + request_base_url = str(request.base_url).rstrip("/") + STATE_MAP[state] = base_url + params = { + "client_id": mcp_server.client_id, + "redirect_uri": f"{request_base_url}/callback", + "scope": " ".join(mcp_server.scopes), + "state": state, + } + return RedirectResponse(f"{mcp_server.authorization_url}?{urlencode(params)}") + + +@router.post("/token") +async def token_endpoint( + request: Request, + grant_type: str = Form(...), + code: str = Form(None), + redirect_uri: str = Form(None), + client_id: str = Form(...), + client_secret: str = Form(...), +): + """ + Accept the authorization code from Claude and exchange it for GitHub token. + Forward the GitHub token back to Claude in standard OAuth format. + + 1. Call the token endpoint + 2. Store the user's PAT in the db - and generate a LiteLLM virtual key + 2. Return the token + 3. Return a virtual key in this response + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + mcp_server = global_mcp_server_manager.get_mcp_server_by_name(client_id) + if mcp_server is None: + raise HTTPException(status_code=404, detail="MCP server not found") + + if grant_type != "authorization_code": + raise HTTPException(status_code=400, detail="Unsupported grant_type") + + if mcp_server.token_url is None: + raise HTTPException(status_code=400, detail="MCP server token url is not set") + + proxy_base_url = str(request.base_url).rstrip("/") + + # Exchange code for real GitHub token + async with httpx.AsyncClient() as client: + resp = await client.post( + mcp_server.token_url, + headers={"Accept": "application/json"}, + data={ + "client_id": mcp_server.client_id, + "client_secret": mcp_server.client_secret, + "code": code, + "redirect_uri": f"{proxy_base_url}/callback", + }, + ) + resp.raise_for_status() + github_token = resp.json()["access_token"] + + # Return to Claude in expected OAuth 2 format + + ### return a virtual key in this response + + return JSONResponse( + {"access_token": github_token, "token_type": "Bearer", "expires_in": 3600} + ) + + +@router.get("/callback") +async def callback(code: str, state: str): + # Exchange code for token with GitHub + params = {"code": code, "state": state} + + # Forward token to Claude ephemeral endpoint + redirect_uri = STATE_MAP.pop(state, None) + + complete_returned_url = f"{redirect_uri}?{urlencode(params)}" + if redirect_uri: + async with httpx.AsyncClient() as client: + await client.post( + complete_returned_url, + json={}, + ) + # Return simple page so the browser doesn't hang + return HTMLResponse( + "Authentication complete. You can close this window." + ) + + # fallback if redirect_uri not found + return HTMLResponse( + "Authentication incomplete. You can close this window." + ) + + +# ------------------------------ +# Optional .well-known endpoints for MCP + OAuth discovery +# ------------------------------ +@router.get("/.well-known/oauth-protected-resource/{mcp_server_name}/mcp") +@router.get("/.well-known/oauth-protected-resource") +async def oauth_protected_resource_mcp( + request: Request, mcp_server_name: Optional[str] = None +): + request_base_url = str(request.base_url).rstrip("/") + return { + "authorization_servers": [ + ( + f"{request_base_url}/{mcp_server_name}" + if mcp_server_name + else f"{request_base_url}" + ) + ], + "resource": ( + f"{request_base_url}/{mcp_server_name}/mcp" + if mcp_server_name + else f"{request_base_url}/mcp" + ), # this is what Claude will call + } + + +@router.get("/.well-known/oauth-authorization-server/{mcp_server_name}") +@router.get("/.well-known/oauth-authorization-server") +async def oauth_authorization_server_mcp( + request: Request, mcp_server_name: Optional[str] = None +): + request_base_url = str(request.base_url).rstrip("/") + return { + "issuer": request_base_url, # point to your proxy + "authorization_endpoint": f"{request_base_url}/authorize", + "token_endpoint": f"{request_base_url}/token", + "response_types_supported": ["code"], + "grant_types_supported": ["authorization_code"], + "code_challenge_methods_supported": ["S256"], + "token_endpoint_auth_methods_supported": ["client_secret_post"], + # Claude expects a registration endpoint, even if we just fake it + "registration_endpoint": f"{request_base_url}/{mcp_server_name}/register", + } + + +# Alias for standard OpenID discovery +@router.get("/.well-known/openid-configuration") +async def openid_configuration(request: Request): + return await oauth_authorization_server_mcp(request) + + +@router.get("/.well-known/oauth-authorization-server/{mcp_server_name}/mcp") +@router.get("/.well-known/oauth-authorization-server") +async def oauth_authorization_server_root( + request: Request, mcp_server_name: Optional[str] = None +): + return await oauth_authorization_server_mcp(request, mcp_server_name) + + +@router.post("/{mcp_server_name}/register") +@router.post("/register") +async def register_client(request: Request, mcp_server_name: Optional[str] = None): + request_base_url = str(request.base_url).rstrip("/") + + # return fixed GitHub client credentials + return { + "client_id": mcp_server_name or "dummy_client", + "client_secret": "dummy", + "redirect_uris": [f"{request_base_url}/mcp/callback"], + } diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index ab5d1b10bf..254f036df0 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -199,6 +199,12 @@ class MCPServerManager: command=server_config.get("command", None) or "", args=server_config.get("args", None) or [], env=server_config.get("env", None) or {}, + # oauth specific fields + client_id=server_config.get("client_id", None), + client_secret=server_config.get("client_secret", None), + scopes=server_config.get("scopes", None), + authorization_url=server_config.get("authorization_url", None), + token_url=server_config.get("token_url", None), # TODO: utility fn the default values transport=server_config.get("transport", MCPTransport.http), auth_type=server_config.get("auth_type", None), @@ -607,20 +613,26 @@ class MCPServerManager: "arguments": arguments, "server_name": server_name_from_prefix, "user_api_key_auth": user_api_key_auth, - "user_api_key_user_id": getattr(user_api_key_auth, "user_id", None) - if user_api_key_auth - else None, - "user_api_key_team_id": getattr(user_api_key_auth, "team_id", None) - if user_api_key_auth - else None, - "user_api_key_end_user_id": getattr( - user_api_key_auth, "end_user_id", None - ) - if user_api_key_auth - else None, - "user_api_key_hash": getattr(user_api_key_auth, "api_key_hash", None) - if user_api_key_auth - else None, + "user_api_key_user_id": ( + getattr(user_api_key_auth, "user_id", None) + if user_api_key_auth + else None + ), + "user_api_key_team_id": ( + getattr(user_api_key_auth, "team_id", None) + if user_api_key_auth + else None + ), + "user_api_key_end_user_id": ( + getattr(user_api_key_auth, "end_user_id", None) + if user_api_key_auth + else None + ), + "user_api_key_hash": ( + getattr(user_api_key_auth, "api_key_hash", None) + if user_api_key_auth + else None + ), } # Create MCP request object for processing @@ -834,6 +846,16 @@ class MCPServerManager: return server return None + def get_mcp_server_by_name(self, server_name: str) -> Optional[MCPServer]: + """ + Get the MCP Server from the server name + """ + registry = self.get_registry() + for server in registry.values(): + if server.server_name == server_name: + return server + return None + def _generate_stable_server_id( self, server_name: str, @@ -1023,9 +1045,11 @@ class MCPServerManager: auth_type=_server_config.auth_type, created_at=datetime.datetime.now(), updated_at=datetime.datetime.now(), - description=_server_config.mcp_info.get("description") - if _server_config.mcp_info - else None, + description=( + _server_config.mcp_info.get("description") + if _server_config.mcp_info + else None + ), mcp_info=_server_config.mcp_info, mcp_access_groups=_server_config.access_groups or [], # Stdio-specific fields diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 2a9174717d..332d1d4e05 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,12 +1,17 @@ import importlib from typing import Dict, List, Optional +from urllib.parse import urlencode, urlparse, urlunparse -from fastapi import APIRouter, Depends, Query, Request +import httpx +from fastapi import APIRouter, Depends, HTTPException, Query, Request +from fastapi.params import Form +from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse from litellm._logging import verbose_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +STATE_MAP = {} MCP_AVAILABLE: bool = True try: importlib.import_module("mcp") @@ -177,9 +182,9 @@ if MCP_AVAILABLE: return { "tools": list_tools_result, "error": "partial_failure" if error_message else None, - "message": error_message - if error_message - else "Successfully retrieved tools", + "message": ( + error_message if error_message else "Successfully retrieved tools" + ), } except Exception as e: @@ -333,3 +338,109 @@ if MCP_AVAILABLE: } return await _execute_with_mcp_client(request, _list_tools_operation) + + @router.get("/authorize") + async def authorize( + request: Request, client_id: str, redirect_uri: str, state: str = "" + ): + # Redirect to real GitHub OAuth + from mcp_server_manager import global_mcp_server_manager + + mcp_server = global_mcp_server_manager.get_mcp_server_by_id(client_id) + if mcp_server is None: + raise HTTPException(status_code=404, detail="MCP server not found") + + if mcp_server.auth_type != "oauth2": + raise HTTPException(status_code=400, detail="MCP server is not OAuth2") + + if mcp_server.client_id is None: + raise HTTPException( + status_code=400, detail="MCP server client id is not set" + ) + + if mcp_server.authorization_url is None: + raise HTTPException( + status_code=400, detail="MCP server authorization url is not set" + ) + + if mcp_server.scopes is None: + raise HTTPException(status_code=400, detail="MCP server scopes is not set") + + request_base_url = str(request.base_url).rstrip("/") + # Parse it to remove any existing query + parsed = urlparse(redirect_uri) + base_url = urlunparse(parsed._replace(query="")) + STATE_MAP[state] = base_url + params = { + "client_id": mcp_server.client_id, + "redirect_uri": f"{request_base_url}/callback", + "scope": " ".join(mcp_server.scopes), + "state": state, + } + return RedirectResponse(f"{mcp_server.authorization_url}?{urlencode(params)}") + + @router.post("/token") + async def token_endpoint( + request: Request, + grant_type=Form(...), + code=Form(None), + client_id=Form(...), + ): + """ + Accept the authorization code from Claude and exchange it for GitHub token. + Forward the GitHub token back to Claude in standard OAuth format. + """ + from mcp_server_manager import global_mcp_server_manager + + mcp_server = global_mcp_server_manager.get_mcp_server_by_id(str(client_id)) + if mcp_server is None: + raise HTTPException(status_code=404, detail="MCP server not found") + + if mcp_server.token_url is None: + raise HTTPException( + status_code=400, detail="MCP server token url is not set" + ) + + if grant_type != "authorization_code": + raise HTTPException(status_code=400, detail="Unsupported grant_type") + + request_base_url = str(request.base_url).rstrip("/") + # Exchange code for real GitHub token + async with httpx.AsyncClient() as client: + resp = await client.post( + mcp_server.token_url, + headers={"Accept": "application/json"}, + data={ + "client_id": mcp_server.client_id, + "client_secret": mcp_server.client_secret, + "code": code, + "redirect_uri": f"{request_base_url}/callback", + }, + ) + resp.raise_for_status() + return JSONResponse(resp.json()) + + @router.get("/callback") + async def callback(code: str, state: str): + # Exchange code for token with GitHub + params = {"code": code, "state": state} + + # Forward token to Claude ephemeral endpoint + redirect_uri = STATE_MAP.pop(state, None) + + complete_returned_url = f"{redirect_uri}?{urlencode(params)}" + if redirect_uri: + async with httpx.AsyncClient() as client: + await client.post( + complete_returned_url, + json={}, + ) + # Return simple page so the browser doesn't hang + return HTMLResponse( + "Authentication complete. You can close this window." + ) + + # fallback if redirect_uri not found + return HTMLResponse( + "Authentication incomplete. You can close this window." + ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 2cf91c84dc..f02b079196 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -519,16 +519,16 @@ if MCP_AVAILABLE: "litellm_logging_obj", None ) if litellm_logging_obj: - litellm_logging_obj.model_call_details[ - "mcp_tool_call_metadata" - ] = standard_logging_mcp_tool_call + litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( + standard_logging_mcp_tool_call + ) litellm_logging_obj.model = f"MCP: {name}" # Try managed server tool first (pass the full prefixed name) # Primary and recommended way to use MCP servers ######################################################### - mcp_server: Optional[ - MCPServer - ] = global_mcp_server_manager._get_mcp_server_from_tool_name(name) + mcp_server: Optional[MCPServer] = ( + global_mcp_server_manager._get_mcp_server_from_tool_name(name) + ) if mcp_server: standard_logging_mcp_tool_call["mcp_server_cost_info"] = ( mcp_server.mcp_info or {} @@ -638,26 +638,32 @@ if MCP_AVAILABLE: mcp_path_match = re.match(r"^/mcp/([^?#]+)(?:\?.*)?(?:#.*)?$", path) if mcp_path_match: servers_and_path = mcp_path_match.group(1) - + if servers_and_path: # Check if it contains commas (comma-separated servers) - if ',' in servers_and_path: + if "," in servers_and_path: # For comma-separated, look for a path at the end # Common patterns: /tools, /chat/completions, etc. - path_match = re.search(r'/([^/,]+(?:/[^/,]+)*)$', servers_and_path) + path_match = re.search(r"/([^/,]+(?:/[^/,]+)*)$", servers_and_path) if path_match: # Path found at the end, remove it from servers - path_part = '/' + path_match.group(1) - servers_part = servers_and_path[:-len(path_part)] - mcp_servers_from_path = [s.strip() for s in servers_part.split(',') if s.strip()] + path_part = "/" + path_match.group(1) + servers_part = servers_and_path[: -len(path_part)] + mcp_servers_from_path = [ + s.strip() for s in servers_part.split(",") if s.strip() + ] else: # No path, just comma-separated servers - mcp_servers_from_path = [s.strip() for s in servers_and_path.split(',') if s.strip()] + mcp_servers_from_path = [ + s.strip() for s in servers_and_path.split(",") if s.strip() + ] else: # Single server case - use regex approach for server/path separation # This handles cases like "custom_solutions/user_123/chat/completions" # where we want to extract "custom_solutions/user_123" as the server name - single_server_match = re.match(r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path) + single_server_match = re.match( + r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path + ) if single_server_match: server_name = single_server_match.group(1) mcp_servers_from_path = [server_name] @@ -809,6 +815,8 @@ if MCP_AVAILABLE: # Mount the MCP handlers app.mount("/", handle_streamable_http_mcp) + app.mount("/mcp", handle_streamable_http_mcp) + app.mount("/{mcp_server_name}/mcp", handle_streamable_http_mcp) app.mount("/sse", handle_sse_mcp) app.add_middleware(AuthContextMiddleware) @@ -839,14 +847,12 @@ if MCP_AVAILABLE: ) auth_context_var.set(auth_user) - def get_auth_context() -> ( - Tuple[ - Optional[UserAPIKeyAuth], - Optional[str], - Optional[List[str]], - Optional[Dict[str, str]], - ] - ): + def get_auth_context() -> Tuple[ + Optional[UserAPIKeyAuth], + Optional[str], + Optional[List[str]], + Optional[Dict[str, str]], + ]: """ Get the UserAPIKeyAuth from the auth context variable. diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html deleted file mode 100644 index 1dc5fdb1f8..0000000000 --- a/litellm/proxy/_experimental/out/onboarding.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 8ffe5e8aef..c50d391e05 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -15,3 +15,16 @@ model_list: model: hosted_vllm/whisper-v3 api_base: "https://webhook.site/2f385e05-00aa-402b-86d1-efc9261471a5" api_key: dummy + +mcp_servers: + github_mcp: + url: "https://api.githubcopilot.com/mcp" + auth_type: oauth2 + authorization_url: https://github.com/login/oauth/authorize + 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"] + # allowed_tools: ["list_tools"] + # disallowed_tools: ["repo_delete"] + diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 6dea634a80..9edfe677ca 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -390,19 +390,18 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 pass_through_endpoints: Optional[List[dict]] = general_settings.get( "pass_through_endpoints", None ) - passed_in_key: Optional[str] = None ## CHECK IF X-LITELM-API-KEY IS PASSED IN - supercedes Authorization header - api_key, passed_in_key = get_api_key( - custom_litellm_key_header=custom_litellm_key_header, - api_key=api_key, - azure_api_key_header=azure_api_key_header, - anthropic_api_key_header=anthropic_api_key_header, - google_ai_studio_api_key_header=google_ai_studio_api_key_header, - azure_apim_header=azure_apim_header, - pass_through_endpoints=pass_through_endpoints, - route=route, - request=request, - ) + # api_key, passed_in_key = get_api_key( + # custom_litellm_key_header=custom_litellm_key_header, + # api_key=api_key, + # azure_api_key_header=azure_api_key_header, + # anthropic_api_key_header=anthropic_api_key_header, + # google_ai_studio_api_key_header=google_ai_studio_api_key_header, + # azure_apim_header=azure_apim_header, + # pass_through_endpoints=pass_through_endpoints, + # route=route, + # request=request, + # ) # if user wants to pass LiteLLM_Master_Key as a custom header, example pass litellm keys as X-LiteLLM-Key: Bearer sk-1234 custom_litellm_key_header_name = general_settings.get("litellm_key_header_name") if custom_litellm_key_header_name is not None: @@ -502,7 +501,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 end_user_object = result["end_user_object"] org_id = result["org_id"] token = result["token"] - team_membership: Optional[LiteLLM_TeamMembership] = result.get("team_membership", None) + team_membership: Optional[LiteLLM_TeamMembership] = result.get( + "team_membership", None + ) global_proxy_spend = await get_global_proxy_spend( litellm_proxy_admin_name=litellm_proxy_admin_name, @@ -537,10 +538,22 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 org_id=org_id, parent_otel_span=parent_otel_span, end_user_id=end_user_id, - user_tpm_limit=user_object.tpm_limit if user_object is not None else None, - user_rpm_limit=user_object.rpm_limit if user_object is not None else None, - team_member_rpm_limit=team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None, - team_member_tpm_limit=team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None, + user_tpm_limit=( + user_object.tpm_limit if user_object is not None else None + ), + user_rpm_limit=( + user_object.rpm_limit if user_object is not None else None + ), + team_member_rpm_limit=( + team_membership.safe_get_team_member_rpm_limit() + if team_membership is not None + else None + ), + team_member_tpm_limit=( + team_membership.safe_get_team_member_tpm_limit() + if team_membership is not None + else None + ), ) # run through common checks _ = await common_checks( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 472aeb140c..628af159ed 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -35,13 +35,13 @@ from litellm.constants import ( LITELLM_SETTINGS_SAFE_DB_OVERRIDES, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.utils import load_credentials_from_list from litellm.types.utils import ( ModelResponse, ModelResponseStream, TextCompletionResponse, TokenCountResponse, ) +from litellm.utils import load_credentials_from_list if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -151,6 +151,9 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + router as mcp_discoverable_endpoints_router, +) from litellm.proxy._experimental.mcp_server.rest_endpoints import ( router as mcp_rest_endpoints_router, ) @@ -249,9 +252,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, @@ -298,9 +299,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, ) @@ -9470,5 +9469,75 @@ app.include_router(ui_discovery_endpoints_router) ######################################################## # MCP Server ######################################################## + + +# Dynamic MCP server routes - handle /{mcp_server_name}/mcp +@app.api_route( + "/{mcp_server_name}/mcp", + methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], +) +async def dynamic_mcp_route(mcp_server_name: str, request: Request): + """Handle dynamic MCP server routes like /github_mcp/mcp""" + try: + # Validate that the MCP server exists + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name) + if mcp_server is None: + raise HTTPException( + status_code=404, detail=f"MCP server '{mcp_server_name}' not found" + ) + + # Create a new scope with the correct path format that the MCP handler expects + # Transform /{mcp_server_name}/mcp to /mcp/{mcp_server_name} + scope = dict(request.scope) + scope["path"] = f"/mcp/{mcp_server_name}" + + # Import the MCP handler + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + ) + + # Create a custom send function to capture the response + response_started = False + response_body = b"" + response_status = 200 + response_headers = [] + + async def custom_send(message): + nonlocal response_started, response_body, response_status, response_headers + if message["type"] == "http.response.start": + response_started = True + response_status = message["status"] + response_headers = message.get("headers", []) + elif message["type"] == "http.response.body": + response_body += message.get("body", b"") + + # Call the existing MCP handler + await handle_streamable_http_mcp( + scope, receive=request.receive, send=custom_send + ) + + # Return the response + from starlette.responses import Response + + headers_dict = {k.decode(): v.decode() for k, v in response_headers} + return Response( + content=response_body, + status_code=response_status, + headers=headers_dict, + media_type=headers_dict.get("content-type", "application/json"), + ) + + except Exception as e: + verbose_proxy_logger.error( + f"Error handling dynamic MCP route for {mcp_server_name}: {str(e)}" + ) + raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") + + app.mount(path=BASE_MCP_ROUTE, app=mcp_app) app.include_router(mcp_rest_endpoints_router) +app.include_router(mcp_discoverable_endpoints_router) diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 4cd32915ce..bbb9e7d7b4 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -15,6 +15,7 @@ else: MCPImageContent = Any MCPTextContent = Any + class MCPTransport(str, enum.Enum): sse = "sse" http = "http" @@ -26,17 +27,21 @@ class MCPSpecVersion(str, enum.Enum): mar_2025 = "2025-03-26" jun_2025 = "2025-06-18" + class MCPAuth(str, enum.Enum): none = "none" api_key = "api_key" bearer_token = "bearer_token" basic = "basic" authorization = "authorization" + oauth2 = "oauth2" # MCP Literals MCPTransportType = Literal[MCPTransport.sse, MCPTransport.http, MCPTransport.stdio] -MCPSpecVersionType = Literal[MCPSpecVersion.nov_2024, MCPSpecVersion.mar_2025, MCPSpecVersion.jun_2025] +MCPSpecVersionType = Literal[ + MCPSpecVersion.nov_2024, MCPSpecVersion.mar_2025, MCPSpecVersion.jun_2025 +] MCPAuthType = Optional[ Literal[ MCPAuth.none, @@ -44,11 +49,11 @@ MCPAuthType = Optional[ MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization, + MCPAuth.oauth2, ] ] - class MCPServerCostInfo(TypedDict, total=False): default_cost_per_query: Optional[float] """ @@ -82,6 +87,7 @@ class MCPPreCallRequestObject(BaseModel): """ Pydantic object used for MCP pre_call_hook request validation and modification """ + tool_name: str arguments: Dict[str, Any] server_name: Optional[str] = None @@ -93,6 +99,7 @@ class MCPPreCallResponseObject(BaseModel): """ Pydantic object used for MCP pre_call_hook response """ + should_proceed: bool = True modified_arguments: Optional[Dict[str, Any]] = None error_message: Optional[str] = None @@ -103,6 +110,7 @@ class MCPDuringCallRequestObject(BaseModel): """ Pydantic object used for MCP during_call_hook request """ + tool_name: str arguments: Dict[str, Any] server_name: Optional[str] = None @@ -114,6 +122,7 @@ class MCPDuringCallResponseObject(BaseModel): """ Pydantic object used for MCP during_call_hook response """ + should_continue: bool = True error_message: Optional[str] = None hidden_params: HiddenParams = HiddenParams() @@ -123,5 +132,8 @@ class MCPPostCallResponseObject(BaseModel): """ Pydantic object used for MCP post_call_hook response """ - mcp_tool_call_response: List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]] - hidden_params: HiddenParams \ No newline at end of file + + mcp_tool_call_response: List[ + Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource] + ] + hidden_params: HiddenParams diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index bd09ef9c19..0dec1b23c6 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -6,7 +6,6 @@ from typing_extensions import TypedDict from litellm.proxy._types import MCPAuthType, MCPTransportType from litellm.types.mcp import MCPServerCostInfo - # MCPInfo now allows arbitrary additional fields for custom metadata MCPInfo = Dict[str, Any] @@ -21,6 +20,12 @@ class MCPServer(BaseModel): auth_type: Optional[MCPAuthType] = None authentication_token: Optional[str] = None mcp_info: Optional[MCPInfo] = None + # OAuth-specific fields + client_id: Optional[str] = None + client_secret: Optional[str] = None + scopes: Optional[List[str]] = None + authorization_url: Optional[str] = None + token_url: Optional[str] = None # Stdio-specific fields command: Optional[str] = None args: Optional[List[str]] = None From 667df3c3db1423218492bd4800ad48366049c64b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Sep 2025 17:21:42 -0700 Subject: [PATCH 02/76] fix: revert commit --- litellm/proxy/auth/user_api_key_auth.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9edfe677ca..63047f37b1 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -391,17 +391,17 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 "pass_through_endpoints", None ) ## CHECK IF X-LITELM-API-KEY IS PASSED IN - supercedes Authorization header - # api_key, passed_in_key = get_api_key( - # custom_litellm_key_header=custom_litellm_key_header, - # api_key=api_key, - # azure_api_key_header=azure_api_key_header, - # anthropic_api_key_header=anthropic_api_key_header, - # google_ai_studio_api_key_header=google_ai_studio_api_key_header, - # azure_apim_header=azure_apim_header, - # pass_through_endpoints=pass_through_endpoints, - # route=route, - # request=request, - # ) + api_key, passed_in_key = get_api_key( + custom_litellm_key_header=custom_litellm_key_header, + api_key=api_key, + azure_api_key_header=azure_api_key_header, + anthropic_api_key_header=anthropic_api_key_header, + google_ai_studio_api_key_header=google_ai_studio_api_key_header, + azure_apim_header=azure_apim_header, + pass_through_endpoints=pass_through_endpoints, + route=route, + request=request, + ) # if user wants to pass LiteLLM_Master_Key as a custom header, example pass litellm keys as X-LiteLLM-Key: Bearer sk-1234 custom_litellm_key_header_name = general_settings.get("litellm_key_header_name") if custom_litellm_key_header_name is not None: From 0d36e6cbe9e5def82e1a7be5d9f293568b6d5728 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Sep 2025 18:13:12 -0700 Subject: [PATCH 03/76] fix(discoverable_endpoints.py): redirect to redirect uri --- .../mcp_server/discoverable_endpoints.py | 20 +++++++++---------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index d38edf5dae..513b590b8d 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -12,9 +12,14 @@ router = APIRouter( ) +@router.get("/{mcp_server_name}/authorize") @router.get("/authorize") async def authorize( - request: Request, client_id: str, redirect_uri: str, state: str = "" + request: Request, + client_id: str, + redirect_uri: str, + state: str = "", + mcp_server_name: Optional[str] = None, ): # Redirect to real GitHub OAuth from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -115,17 +120,10 @@ async def callback(code: str, state: str): # Forward token to Claude ephemeral endpoint redirect_uri = STATE_MAP.pop(state, None) - complete_returned_url = f"{redirect_uri}?{urlencode(params)}" if redirect_uri: - async with httpx.AsyncClient() as client: - await client.post( - complete_returned_url, - json={}, - ) - # Return simple page so the browser doesn't hang - return HTMLResponse( - "Authentication complete. You can close this window." - ) + complete_returned_url = f"{redirect_uri}?{urlencode(params)}" + + return RedirectResponse(url=complete_returned_url, status_code=302) # fallback if redirect_uri not found return HTMLResponse( From 1c939c70e58c049203f9214a2ca419023274f114 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Sep 2025 18:39:10 -0700 Subject: [PATCH 04/76] feat(discoverable_endpoints.py): use encryption + encoding to securely handle state + redirect uri without storing in db would needlessly flood the db --- .../mcp_server/discoverable_endpoints.py | 74 +++++++++++++++---- 1 file changed, 60 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 513b590b8d..05dae916af 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,3 +1,4 @@ +import json from typing import Optional from urllib.parse import urlencode, urlparse, urlunparse @@ -5,13 +6,54 @@ import httpx from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse -STATE_MAP = {} +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + encrypt_value_helper, +) router = APIRouter( tags=["mcp"], ) +def encode_state_with_base_url(base_url: str, original_state: str) -> str: + """ + Encode the base_url and original state using encryption. + + Args: + base_url: The base URL to encode + original_state: The original state parameter + + Returns: + An encrypted string that encodes both values + """ + state_data = {"base_url": base_url, "original_state": original_state} + state_json = json.dumps(state_data, sort_keys=True) + encrypted_state = encrypt_value_helper(state_json) + return encrypted_state + + +def decode_state_hash(encrypted_state: str) -> tuple[str, str]: + """ + Decode an encrypted state to retrieve the base_url and original state. + + Args: + encrypted_state: The encrypted string to decode + + Returns: + A tuple of (base_url, original_state) + + Raises: + Exception: If decryption fails or data is malformed + """ + decrypted_json = decrypt_value_helper(encrypted_state, "oauth_state") + if decrypted_json is None: + raise ValueError("Failed to decrypt state parameter") + + state_data = json.loads(decrypted_json) + return state_data["base_url"], state_data["original_state"] + + @router.get("/{mcp_server_name}/authorize") @router.get("/authorize") async def authorize( @@ -44,12 +86,15 @@ async def authorize( parsed = urlparse(redirect_uri) base_url = urlunparse(parsed._replace(query="")) request_base_url = str(request.base_url).rstrip("/") - STATE_MAP[state] = base_url + + # Encode the base_url and original state in a unique hash + encoded_state = encode_state_with_base_url(base_url, state) + params = { "client_id": mcp_server.client_id, "redirect_uri": f"{request_base_url}/callback", "scope": " ".join(mcp_server.scopes), - "state": state, + "state": encoded_state, } return RedirectResponse(f"{mcp_server.authorization_url}?{urlencode(params)}") @@ -114,21 +159,22 @@ async def token_endpoint( @router.get("/callback") async def callback(code: str, state: str): - # Exchange code for token with GitHub - params = {"code": code, "state": state} + try: + # Decode the state hash to get base_url and original state + base_url, original_state = decode_state_hash(state) - # Forward token to Claude ephemeral endpoint - redirect_uri = STATE_MAP.pop(state, None) - - if redirect_uri: - complete_returned_url = f"{redirect_uri}?{urlencode(params)}" + # Exchange code for token with GitHub + params = {"code": code, "state": original_state} + # Forward token to Claude ephemeral endpoint + complete_returned_url = f"{base_url}?{urlencode(params)}" return RedirectResponse(url=complete_returned_url, status_code=302) - # fallback if redirect_uri not found - return HTMLResponse( - "Authentication incomplete. You can close this window." - ) + except Exception: + # fallback if state hash not found + return HTMLResponse( + "Authentication incomplete. You can close this window." + ) # ------------------------------ From c73f4914de49d46256923305836f1512d05a4c87 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Sep 2025 19:35:05 -0700 Subject: [PATCH 05/76] feat(mcp/): forward oauth2 headers, when applicable if server is of auth type oauth, forward along the oauth headers Allows github mcp on litellm to work via claude code --- litellm/experimental_mcp_client/client.py | 14 ++++--- .../mcp_server/auth/litellm_auth_handler.py | 16 ++++++-- .../mcp_server/auth/user_api_key_auth_mcp.py | 22 ++++++++++- .../mcp_server/mcp_server_manager.py | 14 ++++++- .../proxy/_experimental/mcp_server/server.py | 37 +++++++++++++++++-- litellm/proxy/proxy_server.py | 1 + 6 files changed, 90 insertions(+), 14 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index ecb58e1822..d4db9251fc 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -1,10 +1,11 @@ """ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers. """ + import asyncio import base64 from datetime import timedelta -from typing import List, Optional +from typing import Dict, List, Optional from mcp import ClientSession, StdioServerParameters from mcp.client.sse import sse_client @@ -46,6 +47,7 @@ class MCPClient: auth_value: Optional[str] = None, timeout: float = 60.0, stdio_config: Optional[MCPStdioConfig] = None, + extra_headers: Optional[Dict[str, str]] = None, ): self.server_url: str = server_url self.transport_type: MCPTransport = transport_type @@ -59,7 +61,7 @@ class MCPClient: self._session_ctx = None self._task: Optional[asyncio.Task] = None self.stdio_config: Optional[MCPStdioConfig] = stdio_config - + self.extra_headers: Optional[Dict[str, str]] = extra_headers # handle the basic auth value if provided if auth_value: self.update_auth_value(auth_value) @@ -186,9 +188,7 @@ class MCPClient: def _get_auth_headers(self) -> dict: """Generate authentication headers based on auth type.""" - headers = { - "MCP-Protocol-Version": "2025-06-18" - } + headers = {"MCP-Protocol-Version": "2025-06-18"} if self._mcp_auth_value: if self.auth_type == MCPAuth.bearer_token: @@ -200,6 +200,10 @@ class MCPClient: elif self.auth_type == MCPAuth.authorization: headers["Authorization"] = self._mcp_auth_value + # update the headers with the extra headers + if self.extra_headers: + headers.update(self.extra_headers) + return headers async def list_tools(self) -> List[MCPTool]: diff --git a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py index 058f45d712..56a22040f0 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py +++ b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Dict +from typing import Dict, List, Optional from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser @@ -8,17 +8,27 @@ from litellm.proxy._types import UserAPIKeyAuth class MCPAuthenticatedUser(AuthenticatedUser): """ Wrapper class to make LiteLLM's authentication and configuration compatible with MCP's AuthenticatedUser. - + This class handles: 1. User API key authentication information 2. MCP authentication header (deprecated) 3. MCP server configuration (can include access groups) 4. Server-specific authentication headers + 5. OAuth2 headers """ - def __init__(self, user_api_key_auth: UserAPIKeyAuth, mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None, mcp_server_auth_headers: Optional[Dict[str, str]] = None, mcp_protocol_version: Optional[str] = None): + def __init__( + self, + user_api_key_auth: UserAPIKeyAuth, + mcp_auth_header: Optional[str] = None, + mcp_servers: Optional[List[str]] = None, + mcp_server_auth_headers: Optional[Dict[str, str]] = None, + mcp_protocol_version: Optional[str] = None, + oauth2_headers: Optional[Dict[str, str]] = None, + ): self.user_api_key_auth = user_api_key_auth self.mcp_auth_header = mcp_auth_header self.mcp_servers = mcp_servers self.mcp_server_auth_headers = mcp_server_auth_headers or {} self.mcp_protocol_version = mcp_protocol_version + self.oauth2_headers = oauth2_headers diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 4e7745392b..281edf38ad 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -36,7 +36,11 @@ class MCPRequestHandler: async def process_mcp_request( scope: Scope, ) -> Tuple[ - UserAPIKeyAuth, Optional[str], Optional[List[str]], Optional[Dict[str, str]] + UserAPIKeyAuth, + Optional[str], + Optional[List[str]], + Optional[Dict[str, str]], + Optional[Dict[str, str]], ]: """ Process and validate MCP request headers from the ASGI scope. @@ -44,6 +48,7 @@ class MCPRequestHandler: 1. Extracting and validating authentication headers 2. Processing MCP server configuration 3. Handling MCP-specific headers + 4. Handling oauth2 headers Args: scope: ASGI scope containing request information @@ -70,6 +75,9 @@ class MCPRequestHandler: MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) ) + # Get the oauth2 headers + oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers) + # Parse MCP servers from header mcp_servers_header = headers.get( MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME @@ -107,6 +115,7 @@ class MCPRequestHandler: mcp_auth_header, mcp_servers, mcp_server_auth_headers, + oauth2_headers, ) @staticmethod @@ -177,6 +186,17 @@ class MCPRequestHandler: return server_auth_headers + @staticmethod + def _get_oauth2_headers_from_headers(headers: Headers) -> Dict[str, str]: + """ + Get the oauth2 headers from the request headers. + """ + oauth2_headers = {} + for header_name, header_value in headers.items(): + if header_name.lower().startswith("authorization"): + oauth2_headers["Authorization"] = header_value + return oauth2_headers + @staticmethod def _get_mcp_client_side_auth_header_name() -> str: """ diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 254f036df0..bc69314bcc 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -39,7 +39,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.utils import ProxyLogging -from litellm.types.mcp import MCPStdioConfig +from litellm.types.mcp import MCPAuth, MCPStdioConfig from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer @@ -382,6 +382,7 @@ class MCPServerManager: self, server: MCPServer, mcp_auth_header: Optional[str] = None, + extra_headers: Optional[Dict[str, str]] = None, ) -> MCPClient: """ Create an MCPClient instance for the given server. @@ -411,6 +412,7 @@ class MCPServerManager: auth_value=mcp_auth_header or server.authentication_token, timeout=60.0, stdio_config=stdio_config, + extra_headers=extra_headers, ) else: # For HTTP/SSE transports @@ -421,12 +423,14 @@ class MCPServerManager: auth_type=server.auth_type, auth_value=mcp_auth_header or server.authentication_token, timeout=60.0, + extra_headers=extra_headers, ) async def _get_tools_from_server( self, server: MCPServer, mcp_auth_header: Optional[str] = None, + extra_headers: Optional[Dict[str, str]] = None, ) -> List[MCPTool]: """ Helper method to get tools from a single MCP server with prefixed names. @@ -447,6 +451,7 @@ class MCPServerManager: client = self._create_mcp_client( server=server, mcp_auth_header=mcp_auth_header, + extra_headers=extra_headers, ) tools = await self._fetch_tools_with_timeout(client, server.name) @@ -564,6 +569,7 @@ class MCPServerManager: mcp_auth_header: Optional[str] = None, mcp_server_auth_headers: Optional[Dict[str, str]] = None, proxy_logging_obj: Optional[ProxyLogging] = None, + oauth2_headers: Optional[Dict[str, str]] = None, ) -> CallToolResult: """ Call a tool with the given name and arguments (handles prefixed tool names) @@ -684,9 +690,15 @@ class MCPServerManager: 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 + client = self._create_mcp_client( server=mcp_server, mcp_auth_header=server_auth_header, + extra_headers=extra_headers, ) async with client: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index f02b079196..f79d40ca61 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -22,6 +22,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_VERSION, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer from litellm.types.utils import StandardLoggingMCPToolCall from litellm.utils import client @@ -178,6 +179,7 @@ if MCP_AVAILABLE: mcp_auth_header, mcp_servers, mcp_server_auth_headers, + oauth2_headers, ) = get_auth_context() verbose_logger.debug( f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}" @@ -195,6 +197,7 @@ if MCP_AVAILABLE: mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, ) verbose_logger.info( f"MCP list_tools - Successfully returned {len(tools)} tools" @@ -235,6 +238,7 @@ if MCP_AVAILABLE: mcp_auth_header, _, mcp_server_auth_headers, + oauth2_headers, ) = get_auth_context() verbose_logger.debug( @@ -357,15 +361,17 @@ if MCP_AVAILABLE: mcp_auth_header: Optional[str], mcp_servers: Optional[List[str]], mcp_server_auth_headers: Optional[Dict[str, str]] = None, + oauth2_headers: Optional[Dict[str, str]] = None, ) -> List[MCPTool]: - """ + f""" Helper method to fetch tools from MCP servers based on server filtering criteria. Args: user_api_key_auth: User authentication info for access control mcp_auth_header: Optional auth header for MCP server (deprecated) mcp_servers: Optional list of server names/aliases to filter by - mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} + mcp_server_auth_headers: Optional dict of server-specific auth headers + oauth2_headers: Optional dict of oauth2 headers Returns: List[MCPTool]: Combined list of tools from filtered servers @@ -398,6 +404,10 @@ if MCP_AVAILABLE: elif mcp_server_auth_headers and server.server_name is not None: server_auth_header = mcp_server_auth_headers.get(server.server_name) + extra_headers: Optional[Dict[str, str]] = None + if server.auth_type == MCPAuth.oauth2: + extra_headers = oauth2_headers + # 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 @@ -406,6 +416,7 @@ if MCP_AVAILABLE: tools = await global_mcp_server_manager._get_tools_from_server( server=server, mcp_auth_header=server_auth_header, + extra_headers=extra_headers, ) all_tools.extend(tools) verbose_logger.debug( @@ -427,6 +438,7 @@ if MCP_AVAILABLE: mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None, mcp_server_auth_headers: Optional[Dict[str, str]] = None, + oauth2_headers: Optional[Dict[str, str]] = None, ) -> List[MCPTool]: """ List all available MCP tools. @@ -450,6 +462,7 @@ if MCP_AVAILABLE: mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, ) verbose_logger.debug( f"Successfully fetched {len(managed_tools)} tools from managed MCP servers" @@ -683,6 +696,7 @@ if MCP_AVAILABLE: mcp_auth_header, _, mcp_server_auth_headers, + oauth2_headers, ) = await MCPRequestHandler.process_mcp_request(scope) mcp_servers = mcp_servers_from_path else: @@ -691,8 +705,15 @@ if MCP_AVAILABLE: mcp_auth_header, mcp_servers, mcp_server_auth_headers, + oauth2_headers, ) = await MCPRequestHandler.process_mcp_request(scope) - return user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers + return ( + user_api_key_auth, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + ) async def handle_streamable_http_mcp( scope: Scope, receive: Receive, send: Send @@ -705,6 +726,7 @@ if MCP_AVAILABLE: mcp_auth_header, mcp_servers, mcp_server_auth_headers, + oauth2_headers, ) = await extract_mcp_auth_context(scope, path) verbose_logger.debug( f"MCP request mcp_servers (header/path): {mcp_servers}" @@ -718,6 +740,7 @@ if MCP_AVAILABLE: mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, ) # Ensure session managers are initialized @@ -756,6 +779,7 @@ if MCP_AVAILABLE: mcp_auth_header, mcp_servers, mcp_server_auth_headers, + oauth2_headers, ) = await extract_mcp_auth_context(scope, path) verbose_logger.debug( f"MCP request mcp_servers (header/path): {mcp_servers}" @@ -768,6 +792,7 @@ if MCP_AVAILABLE: mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, ) if not _SESSION_MANAGERS_INITIALIZED: @@ -829,6 +854,7 @@ if MCP_AVAILABLE: mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None, mcp_server_auth_headers: Optional[Dict[str, str]] = None, + oauth2_headers: Optional[Dict[str, str]] = None, ) -> None: """ Set the UserAPIKeyAuth in the auth context variable. @@ -844,6 +870,7 @@ if MCP_AVAILABLE: mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, ) auth_context_var.set(auth_user) @@ -852,6 +879,7 @@ if MCP_AVAILABLE: Optional[str], Optional[List[str]], Optional[Dict[str, str]], + Optional[Dict[str, str]], ]: """ Get the UserAPIKeyAuth from the auth context variable. @@ -867,8 +895,9 @@ if MCP_AVAILABLE: auth_user.mcp_auth_header, auth_user.mcp_servers, auth_user.mcp_server_auth_headers, + auth_user.oauth2_headers, ) - return None, None, None, None + return None, None, None, None, None ######################################################## ############ End of Auth Context Functions ############# diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 628af159ed..42bbae77f4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9483,6 +9483,7 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + from litellm.types.mcp import MCPAuth mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name) if mcp_server is None: From 47ef6a7ee676064a02f21148fd2719ea1e7ae1bf Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 09:02:08 -0700 Subject: [PATCH 06/76] fix(mcp/): pass oauth2 headers on call tools --- litellm/experimental_mcp_client/client.py | 3 +++ litellm/proxy/_experimental/mcp_server/server.py | 5 +++++ 2 files changed, 8 insertions(+) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index d4db9251fc..2e02f460b6 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -117,6 +117,9 @@ class MCPClient: await self._session.initialize() else: # http headers = self._get_auth_headers() + verbose_logger.debug( + "litellm headers for streamablehttp_client: ", headers + ) self._transport_ctx = streamablehttp_client( url=self.server_url, timeout=timedelta(seconds=self.timeout), diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index f79d40ca61..7f16179390 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -270,6 +270,7 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, **data, # for logging ) except BlockedPiiEntityError as e: @@ -505,6 +506,7 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth] = None, mcp_auth_header: Optional[str] = None, mcp_server_auth_headers: Optional[Dict[str, str]] = None, + oauth2_headers: Optional[Dict[str, str]] = None, **kwargs: Any, ) -> List[Union[TextContent, ImageContent, EmbeddedResource]]: """ @@ -552,6 +554,7 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, litellm_logging_obj=litellm_logging_obj, ) @@ -604,6 +607,7 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth] = None, mcp_auth_header: Optional[str] = None, mcp_server_auth_headers: Optional[Dict[str, str]] = None, + oauth2_headers: Optional[Dict[str, str]] = None, litellm_logging_obj: Optional[Any] = None, ) -> List[Union[TextContent, ImageContent, EmbeddedResource]]: """Handle tool execution for managed server tools""" @@ -616,6 +620,7 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, proxy_logging_obj=proxy_logging_obj, ) verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) From e04df1416d3138a61214b38b79b7a07c8c727eda Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 09:07:44 -0700 Subject: [PATCH 07/76] test_async_chat_azure --- .../test_custom_callback_router.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tests/logging_callback_tests/test_custom_callback_router.py b/tests/logging_callback_tests/test_custom_callback_router.py index e8fc1ac867..87ae211c02 100644 --- a/tests/logging_callback_tests/test_custom_callback_router.py +++ b/tests/logging_callback_tests/test_custom_callback_router.py @@ -393,14 +393,14 @@ async def test_async_chat_azure(): litellm.set_verbose = True model_list = [ { - "model_name": "gpt-3.5-turbo", # openai model name + "model_name": "gpt-4.1-nano", # openai model name "litellm_params": { # params for litellm completion/embedding call - "model": "azure/gpt-4o-new-test", + "model": "azure/gpt-4.1-nano", "api_key": os.getenv("AZURE_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), "api_base": os.getenv("AZURE_API_BASE"), }, - "model_info": {"base_model": "azure/gpt-4-1106-preview"}, + "model_info": {"base_model": "azure/gpt-4.1-nano"}, "tpm": 240000, "rpm": 1800, }, @@ -471,7 +471,6 @@ async def test_async_chat_azure(): pytest.fail(f"An exception occurred - {str(e)}") -# asyncio.run(test_async_chat_azure()) ## EMBEDDING @pytest.mark.asyncio async def test_async_embedding_azure(): @@ -483,7 +482,7 @@ async def test_async_embedding_azure(): { "model_name": "azure-embedding-model", # openai model name "litellm_params": { # params for litellm completion/embedding call - "model": "azure/azure-embedding-model", + "model": "azure/text-embedding-ada-002", "api_key": os.getenv("AZURE_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), "api_base": os.getenv("AZURE_API_BASE"), From 8510c7041631681429cbd13c85b8132c972eab18 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 09:11:43 -0700 Subject: [PATCH 08/76] test fixes --- .../test_custom_callback_input.py | 119 ------------------ .../test_custom_callback_router.py | 79 ++++++++++++ 2 files changed, 79 insertions(+), 119 deletions(-) diff --git a/tests/local_testing/test_custom_callback_input.py b/tests/local_testing/test_custom_callback_input.py index fb7ca51040..90e9adb789 100644 --- a/tests/local_testing/test_custom_callback_input.py +++ b/tests/local_testing/test_custom_callback_input.py @@ -914,125 +914,6 @@ async def test_async_embedding_bedrock(): pytest.fail(f"An exception occurred: {str(e)}") -# asyncio.run(test_async_embedding_bedrock()) - - -# CACHING -## Test Azure - completion, embedding -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=1) -async def test_async_completion_azure_caching(): - litellm.set_verbose = True - customHandler_caching = CompletionCustomHandler() - litellm.cache = Cache( - type="redis", - host=os.environ["REDIS_HOST"], - port=os.environ["REDIS_PORT"], - password=os.environ["REDIS_PASSWORD"], - ) - litellm.callbacks = [customHandler_caching] - unique_time = time.time() - response1 = await litellm.acompletion( - model="azure/chatgpt-v-3", - messages=[ - {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} - ], - caching=True, - ) - await asyncio.sleep(1) - print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}") - response2 = await litellm.acompletion( - model="azure/chatgpt-v-3", - messages=[ - {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} - ], - caching=True, - ) - await asyncio.sleep(1) # success callbacks are done in parallel - print( - f"customHandler_caching.states post-cache hit: {customHandler_caching.states}" - ) - assert len(customHandler_caching.errors) == 0 - assert len(customHandler_caching.states) == 4 # pre, post, success, success - - -@pytest.mark.asyncio -async def test_async_completion_azure_caching_streaming(): - import copy - - litellm.set_verbose = True - customHandler_caching = CompletionCustomHandler() - litellm.cache = Cache( - type="redis", - host=os.environ["REDIS_HOST"], - port=os.environ["REDIS_PORT"], - password=os.environ["REDIS_PASSWORD"], - ) - litellm.callbacks = [customHandler_caching] - unique_time = uuid.uuid4() - response1 = await litellm.acompletion( - model="azure/chatgpt-v-3", - messages=[ - {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} - ], - caching=True, - stream=True, - ) - async for chunk in response1: - print(f"chunk in response1: {chunk}") - await asyncio.sleep(1) - initial_customhandler_caching_states = len(customHandler_caching.states) - print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}") - response2 = await litellm.acompletion( - model="azure/chatgpt-v-3", - messages=[ - {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} - ], - caching=True, - stream=True, - ) - async for chunk in response2: - print(f"chunk in response2: {chunk}") - await asyncio.sleep(1) # success callbacks are done in parallel - print( - f"customHandler_caching.states post-cache hit: {customHandler_caching.states}" - ) - assert len(customHandler_caching.errors) == 0 - assert ( - len(customHandler_caching.states) > initial_customhandler_caching_states - ) # pre, post, streaming .., success, success - - -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=2) -async def test_async_embedding_azure_caching(): - print("Testing custom callback input - Azure Caching") - customHandler_caching = CompletionCustomHandler() - litellm.cache = Cache( - type="redis", - host=os.environ["REDIS_HOST"], - port=os.environ["REDIS_PORT"], - password=os.environ["REDIS_PASSWORD"], - ) - litellm.callbacks = [customHandler_caching] - unique_time = time.time() - response1 = await litellm.aembedding( - model="azure/azure-embedding-model", - input=[f"good morning from litellm1 {unique_time}"], - caching=True, - ) - await asyncio.sleep(1) # set cache is async for aembedding() - response2 = await litellm.aembedding( - model="azure/azure-embedding-model", - input=[f"good morning from litellm1 {unique_time}"], - caching=True, - ) - await asyncio.sleep(1) # success callbacks are done in parallel - print(customHandler_caching.states) - print(customHandler_caching.errors) - assert len(customHandler_caching.errors) == 0 - assert len(customHandler_caching.states) == 4 # pre, post, success, success - # Image Generation diff --git a/tests/logging_callback_tests/test_custom_callback_router.py b/tests/logging_callback_tests/test_custom_callback_router.py index 87ae211c02..f831df2e2d 100644 --- a/tests/logging_callback_tests/test_custom_callback_router.py +++ b/tests/logging_callback_tests/test_custom_callback_router.py @@ -649,6 +649,84 @@ async def test_async_completion_azure_caching(): assert len(customHandler_caching.states) == 4 # pre, post, success, success +@pytest.mark.asyncio +async def test_async_completion_azure_caching_streaming(): + import copy + + litellm.set_verbose = True + customHandler_caching = CompletionCustomHandler() + litellm.cache = Cache( + type="redis", + host=os.environ["REDIS_HOST"], + port=os.environ["REDIS_PORT"], + password=os.environ["REDIS_PASSWORD"], + ) + litellm.callbacks = [customHandler_caching] + unique_time = uuid.uuid4() + response1 = await litellm.acompletion( + model="azure/chatgpt-v-3", + messages=[ + {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} + ], + caching=True, + stream=True, + ) + async for chunk in response1: + print(f"chunk in response1: {chunk}") + await asyncio.sleep(1) + initial_customhandler_caching_states = len(customHandler_caching.states) + print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}") + response2 = await litellm.acompletion( + model="azure/chatgpt-v-3", + messages=[ + {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} + ], + caching=True, + stream=True, + ) + async for chunk in response2: + print(f"chunk in response2: {chunk}") + await asyncio.sleep(1) # success callbacks are done in parallel + print( + f"customHandler_caching.states post-cache hit: {customHandler_caching.states}" + ) + assert len(customHandler_caching.errors) == 0 + assert ( + len(customHandler_caching.states) > initial_customhandler_caching_states + ) # pre, post, streaming .., success, success + + +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=2) +async def test_async_embedding_azure_caching(): + print("Testing custom callback input - Azure Caching") + customHandler_caching = CompletionCustomHandler() + litellm.cache = Cache( + type="redis", + host=os.environ["REDIS_HOST"], + port=os.environ["REDIS_PORT"], + password=os.environ["REDIS_PASSWORD"], + ) + litellm.callbacks = [customHandler_caching] + unique_time = time.time() + response1 = await litellm.aembedding( + model="azure/azure-embedding-model", + input=[f"good morning from litellm1 {unique_time}"], + caching=True, + ) + await asyncio.sleep(1) # set cache is async for aembedding() + response2 = await litellm.aembedding( + model="azure/azure-embedding-model", + input=[f"good morning from litellm1 {unique_time}"], + caching=True, + ) + await asyncio.sleep(1) # success callbacks are done in parallel + print(customHandler_caching.states) + print(customHandler_caching.errors) + assert len(customHandler_caching.errors) == 0 + assert len(customHandler_caching.states) == 4 # pre, post, success, success + + @pytest.mark.asyncio async def test_rate_limit_error_callback(): """ @@ -716,3 +794,4 @@ async def test_rate_limit_error_callback(): assert "original_model_group" in mock_client.call_args.kwargs assert mock_client.call_args.kwargs["original_model_group"] == "my-test-gpt" + From afbebcd626e430243491792798a193f61c20b980 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 09:12:37 -0700 Subject: [PATCH 09/76] test_async_chat_azure --- tests/logging_callback_tests/test_custom_callback_router.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/logging_callback_tests/test_custom_callback_router.py b/tests/logging_callback_tests/test_custom_callback_router.py index f831df2e2d..91d7961919 100644 --- a/tests/logging_callback_tests/test_custom_callback_router.py +++ b/tests/logging_callback_tests/test_custom_callback_router.py @@ -407,7 +407,7 @@ async def test_async_chat_azure(): ] router = Router(model_list=model_list, num_retries=0) # type: ignore response = await router.acompletion( - model="gpt-3.5-turbo", + model="gpt-4.1-nano", messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}], ) print("got response, sleeping 5 seconds....") @@ -422,7 +422,7 @@ async def test_async_chat_azure(): litellm.callbacks = [customHandler_streaming_azure_router] router2 = Router(model_list=model_list, num_retries=0) # type: ignore response = await router2.acompletion( - model="gpt-3.5-turbo", + model="gpt-4.1-nano", messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}], stream=True, ) From 0202f1d034ea85f2342cef4d92087d65bac47e2b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 09:22:30 -0700 Subject: [PATCH 10/76] refactor: cleanup dup code --- .../mcp_server/rest_endpoints.py | 107 ------------------ 1 file changed, 107 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 332d1d4e05..d2f16154f3 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -11,7 +11,6 @@ from litellm._logging import verbose_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -STATE_MAP = {} MCP_AVAILABLE: bool = True try: importlib.import_module("mcp") @@ -338,109 +337,3 @@ if MCP_AVAILABLE: } return await _execute_with_mcp_client(request, _list_tools_operation) - - @router.get("/authorize") - async def authorize( - request: Request, client_id: str, redirect_uri: str, state: str = "" - ): - # Redirect to real GitHub OAuth - from mcp_server_manager import global_mcp_server_manager - - mcp_server = global_mcp_server_manager.get_mcp_server_by_id(client_id) - if mcp_server is None: - raise HTTPException(status_code=404, detail="MCP server not found") - - if mcp_server.auth_type != "oauth2": - raise HTTPException(status_code=400, detail="MCP server is not OAuth2") - - if mcp_server.client_id is None: - raise HTTPException( - status_code=400, detail="MCP server client id is not set" - ) - - if mcp_server.authorization_url is None: - raise HTTPException( - status_code=400, detail="MCP server authorization url is not set" - ) - - if mcp_server.scopes is None: - raise HTTPException(status_code=400, detail="MCP server scopes is not set") - - request_base_url = str(request.base_url).rstrip("/") - # Parse it to remove any existing query - parsed = urlparse(redirect_uri) - base_url = urlunparse(parsed._replace(query="")) - STATE_MAP[state] = base_url - params = { - "client_id": mcp_server.client_id, - "redirect_uri": f"{request_base_url}/callback", - "scope": " ".join(mcp_server.scopes), - "state": state, - } - return RedirectResponse(f"{mcp_server.authorization_url}?{urlencode(params)}") - - @router.post("/token") - async def token_endpoint( - request: Request, - grant_type=Form(...), - code=Form(None), - client_id=Form(...), - ): - """ - Accept the authorization code from Claude and exchange it for GitHub token. - Forward the GitHub token back to Claude in standard OAuth format. - """ - from mcp_server_manager import global_mcp_server_manager - - mcp_server = global_mcp_server_manager.get_mcp_server_by_id(str(client_id)) - if mcp_server is None: - raise HTTPException(status_code=404, detail="MCP server not found") - - if mcp_server.token_url is None: - raise HTTPException( - status_code=400, detail="MCP server token url is not set" - ) - - if grant_type != "authorization_code": - raise HTTPException(status_code=400, detail="Unsupported grant_type") - - request_base_url = str(request.base_url).rstrip("/") - # Exchange code for real GitHub token - async with httpx.AsyncClient() as client: - resp = await client.post( - mcp_server.token_url, - headers={"Accept": "application/json"}, - data={ - "client_id": mcp_server.client_id, - "client_secret": mcp_server.client_secret, - "code": code, - "redirect_uri": f"{request_base_url}/callback", - }, - ) - resp.raise_for_status() - return JSONResponse(resp.json()) - - @router.get("/callback") - async def callback(code: str, state: str): - # Exchange code for token with GitHub - params = {"code": code, "state": state} - - # Forward token to Claude ephemeral endpoint - redirect_uri = STATE_MAP.pop(state, None) - - complete_returned_url = f"{redirect_uri}?{urlencode(params)}" - if redirect_uri: - async with httpx.AsyncClient() as client: - await client.post( - complete_returned_url, - json={}, - ) - # Return simple page so the browser doesn't hang - return HTMLResponse( - "Authentication complete. You can close this window." - ) - - # fallback if redirect_uri not found - return HTMLResponse( - "Authentication incomplete. You can close this window." - ) From 9790451ccecf2c231a245278781d36634b40f9b1 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 09:23:49 -0700 Subject: [PATCH 11/76] test_custom_callback_router.py --- .../test_custom_callback_router.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/logging_callback_tests/test_custom_callback_router.py b/tests/logging_callback_tests/test_custom_callback_router.py index 91d7961919..27d9e17edf 100644 --- a/tests/logging_callback_tests/test_custom_callback_router.py +++ b/tests/logging_callback_tests/test_custom_callback_router.py @@ -605,9 +605,9 @@ async def test_async_completion_azure_caching(): unique_time = time.time() model_list = [ { - "model_name": "gpt-3.5-turbo", # openai model name + "model_name": "gpt-4.1-nano", # openai model name "litellm_params": { # params for litellm completion/embedding call - "model": "azure/chatgpt-v-3", + "model": "azure/gpt-4.1-nano", "api_key": os.getenv("AZURE_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), "api_base": os.getenv("AZURE_API_BASE"), @@ -626,7 +626,7 @@ async def test_async_completion_azure_caching(): ] router = Router(model_list=model_list) # type: ignore response1 = await router.acompletion( - model="gpt-3.5-turbo", + model="gpt-4.1-nano", messages=[ {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} ], @@ -635,7 +635,7 @@ async def test_async_completion_azure_caching(): await asyncio.sleep(1) print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}") response2 = await router.acompletion( - model="gpt-3.5-turbo", + model="gpt-4.1-nano", messages=[ {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} ], @@ -652,6 +652,7 @@ async def test_async_completion_azure_caching(): @pytest.mark.asyncio async def test_async_completion_azure_caching_streaming(): import copy + import uuid litellm.set_verbose = True customHandler_caching = CompletionCustomHandler() @@ -664,7 +665,7 @@ async def test_async_completion_azure_caching_streaming(): litellm.callbacks = [customHandler_caching] unique_time = uuid.uuid4() response1 = await litellm.acompletion( - model="azure/chatgpt-v-3", + model="azure/gpt-4.1-nano", messages=[ {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} ], @@ -677,7 +678,7 @@ async def test_async_completion_azure_caching_streaming(): initial_customhandler_caching_states = len(customHandler_caching.states) print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}") response2 = await litellm.acompletion( - model="azure/chatgpt-v-3", + model="azure/gpt-4.1-nano", messages=[ {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} ], @@ -710,13 +711,13 @@ async def test_async_embedding_azure_caching(): litellm.callbacks = [customHandler_caching] unique_time = time.time() response1 = await litellm.aembedding( - model="azure/azure-embedding-model", + model="azure/text-embedding-ada-002", input=[f"good morning from litellm1 {unique_time}"], caching=True, ) await asyncio.sleep(1) # set cache is async for aembedding() response2 = await litellm.aembedding( - model="azure/azure-embedding-model", + model="azure/text-embedding-ada-002", input=[f"good morning from litellm1 {unique_time}"], caching=True, ) From 4054eeea20bbb46ab8c4fefad8e5d30ca6939b86 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 09:26:38 -0700 Subject: [PATCH 12/76] test build and test --- proxy_server_config.yaml | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/proxy_server_config.yaml b/proxy_server_config.yaml index 16efe9ffd0..f786568118 100644 --- a/proxy_server_config.yaml +++ b/proxy_server_config.yaml @@ -7,14 +7,14 @@ model_list: id: "1" - model_name: gpt-3.5-turbo-end-user-test litellm_params: - model: azure/gpt-4o-new-test - api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ + model: azure/gpt-4.1-nano + api_base: https://krris-m2f9a9i7-eastus2.openai.azure.com/ api_version: "2023-05-15" api_key: os.environ/AZURE_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault - model_name: gpt-3.5-turbo litellm_params: - model: azure/gpt-4o-new-test - api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ + model: azure/gpt-4.1-nano + api_base: https://krris-m2f9a9i7-eastus2.openai.azure.com/ api_version: "2023-05-15" api_key: os.environ/AZURE_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault - model_name: gpt-3.5-turbo-large @@ -26,8 +26,8 @@ model_list: stream_timeout: 60 - model_name: gpt-4 litellm_params: - model: azure/gpt-4o-new-test - api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ + model: azure/gpt-4.1-nano + api_base: https://krris-m2f9a9i7-eastus2.openai.azure.com/ api_version: "2023-05-15" api_key: os.environ/AZURE_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault rpm: 480 @@ -39,9 +39,9 @@ model_list: input_cost_per_second: 0.000420 - model_name: text-embedding-ada-002 litellm_params: - model: azure/azure-embedding-model + model: azure/text-embedding-ada-002 api_key: os.environ/AZURE_API_KEY - api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ + api_base: https://krris-m2f9a9i7-eastus2.openai.azure.com/ api_version: "2023-05-15" model_info: mode: embedding From 22c033d3d77344fae4f6987c07efade338921fb7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 09:28:24 -0700 Subject: [PATCH 13/76] test_acompletion_caching_on_router_caching_groups --- tests/local_testing/test_router_caching.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/tests/local_testing/test_router_caching.py b/tests/local_testing/test_router_caching.py index 574f133ace..d9ee4fdeb5 100644 --- a/tests/local_testing/test_router_caching.py +++ b/tests/local_testing/test_router_caching.py @@ -213,10 +213,8 @@ async def test_acompletion_caching_with_ttl_on_router(): { "model_name": "gpt-3.5-turbo", "litellm_params": { - "model": "azure/chatgpt-v-3", - "api_key": os.getenv("AZURE_API_KEY"), - "api_base": os.getenv("AZURE_API_BASE"), - "api_version": os.getenv("AZURE_API_VERSION"), + "model": "gpt-4.1-nano", + "api_key": os.getenv("OPENAI_API_KEY"), }, "tpm": 100000, "rpm": 10000, @@ -277,9 +275,9 @@ async def test_acompletion_caching_on_router_caching_groups(): "rpm": 10000, }, { - "model_name": "azure-gpt-3.5-turbo", + "model_name": "azure-gpt-4.1-nano", "litellm_params": { - "model": "azure/chatgpt-v-3", + "model": "azure/gpt-4.1-nano", "api_key": os.getenv("AZURE_API_KEY"), "api_base": os.getenv("AZURE_API_BASE"), "api_version": os.getenv("AZURE_API_VERSION"), From ac09d29c0ed5748a0ee3cee94c078ac8b35cc5fa Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 09:30:07 -0700 Subject: [PATCH 14/76] test(tests/): add unit test to confirm oauth2 headers forwarded on receipt --- .../mcp_server/test_mcp_server.py | 120 +++++++++++++++++- 1 file changed, 118 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 8fa9964b1f..10a6ab8cb0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -102,7 +102,9 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): working_server if server_id == "working_server" else failing_server ) - async def mock_get_tools_from_server(server, mcp_auth_header=None): + async def mock_get_tools_from_server( + server, mcp_auth_header=None, extra_headers=None + ): if server.name == "working_server": # Working server returns tools tool1 = MagicMock() @@ -184,7 +186,9 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): failing_server1 if server_id == "failing_server1" else failing_server2 ) - async def mock_get_tools_from_server(server, mcp_auth_header=None): + async def mock_get_tools_from_server( + server, mcp_auth_header=None, extra_headers=None + ): # All servers fail raise Exception(f"Server {server.name} connection failed") @@ -448,3 +452,115 @@ async def test_mcp_routing_with_conflicting_alias_and_group_name(): assert ( called_servers[0].server_id == specific_server.server_id ), "Should have contacted the specific server alias, not the group." + + +@pytest.mark.asyncio +async def test_oauth2_headers_passed_to_mcp_client(): + """Test that OAuth2 headers are properly passed through to the MCP client for OAuth2 servers like github_mcp""" + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.server import ( + _get_tools_from_mcp_servers, + set_auth_context, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP server not available") + + # Clear the registry to avoid conflicts with other tests + global_mcp_server_manager.registry.clear() + + # Create an OAuth2 MCP server similar to github_mcp configuration + oauth2_server = MCPServer( + server_id="github_mcp_server_id", + name="github_mcp", + alias="github_mcp", + transport=MCPTransport.http, + url="https://api.githubcopilot.com/mcp", + auth_type=MCPAuth.oauth2, + client_id="test_github_client_id", + client_secret="test_github_client_secret", + scopes=["public_repo", "user:email"], + authorization_url="https://github.com/login/oauth/authorize", + token_url="https://github.com/login/oauth/access_token", + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # Mock user auth + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user") + + # Set up OAuth2 headers that would come from the client + oauth2_headers = {"Authorization": "Bearer github_oauth_token_12345"} + + # Set auth context with OAuth2 headers + set_auth_context(user_api_key_auth=user_api_key_auth, oauth2_headers=oauth2_headers) + + # This will capture the arguments passed to _create_mcp_client + captured_client_args = {} + + def mock_create_mcp_client(server, mcp_auth_header=None, extra_headers=None): + # Capture the arguments for verification + captured_client_args.update( + { + "server": server, + "mcp_auth_header": mcp_auth_header, + "extra_headers": extra_headers, + } + ) + # Return a mock client that doesn't actually connect + mock_client = MagicMock() + mock_client.disconnect = AsyncMock() + return mock_client + + # Mock _fetch_tools_with_timeout to avoid actual network calls + async def mock_fetch_tools_with_timeout(client, server_name): + return [] # Return empty list of tools + + with patch.object( + global_mcp_server_manager, + "_create_mcp_client", + side_effect=mock_create_mcp_client, + ) as mock_create_client, patch.object( + global_mcp_server_manager, + "_fetch_tools_with_timeout", + side_effect=mock_fetch_tools_with_timeout, + ), patch.object( + global_mcp_server_manager, + "get_allowed_mcp_servers", + AsyncMock(return_value=[oauth2_server.server_id]), + ): + # Call _get_tools_from_mcp_servers which should eventually call _create_mcp_client + await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=None, + mcp_servers=None, # Will use all allowed servers + oauth2_headers=oauth2_headers, + ) + + # Verify that _create_mcp_client was called + assert ( + mock_create_client.call_count == 1 + ), "Expected _create_mcp_client to be called once" + + # Verify the server passed to _create_mcp_client is the OAuth2 server + assert captured_client_args["server"].server_id == oauth2_server.server_id + assert captured_client_args["server"].auth_type == MCPAuth.oauth2 + + # Most importantly: verify that OAuth2 headers were passed as extra_headers + assert ( + captured_client_args["extra_headers"] is not None + ), "Expected extra_headers to be passed for OAuth2 server" + assert ( + captured_client_args["extra_headers"] == oauth2_headers + ), f"Expected OAuth2 headers to be passed as extra_headers, got {captured_client_args['extra_headers']}" + + # Verify the Authorization header specifically + assert "Authorization" in captured_client_args["extra_headers"] + assert ( + captured_client_args["extra_headers"]["Authorization"] + == "Bearer github_oauth_token_12345" + ) From 6451f10f2c9f685fde5cd2b463bb2bd2e79bb3ab Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 09:34:01 -0700 Subject: [PATCH 15/76] fix: fix ruff check --- .../mcp_server/mcp_server_manager.py | 140 ++++++++++-------- .../mcp_server/rest_endpoints.py | 6 +- .../proxy/_experimental/mcp_server/server.py | 2 +- 3 files changed, 79 insertions(+), 69 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index bc69314bcc..9fb39e7498 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -561,6 +561,77 @@ class MCPServerManager: ) return prefixed_tools + async def pre_call_tool_check( + self, + name: str, + arguments: Dict[str, Any], + server_name_from_prefix: str, + user_api_key_auth: Optional[UserAPIKeyAuth], + proxy_logging_obj: ProxyLogging, + ): + pre_hook_kwargs = { + "name": name, + "arguments": arguments, + "server_name": server_name_from_prefix, + "user_api_key_auth": user_api_key_auth, + "user_api_key_user_id": ( + getattr(user_api_key_auth, "user_id", None) + if user_api_key_auth + else None + ), + "user_api_key_team_id": ( + getattr(user_api_key_auth, "team_id", None) + if user_api_key_auth + else None + ), + "user_api_key_end_user_id": ( + getattr(user_api_key_auth, "end_user_id", None) + if user_api_key_auth + else None + ), + "user_api_key_hash": ( + getattr(user_api_key_auth, "api_key_hash", None) + if user_api_key_auth + else None + ), + } + + # Create MCP request object for processing + mcp_request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs( + pre_hook_kwargs + ) + + # Convert to LLM format for existing guardrail compatibility + synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format( + mcp_request_obj, pre_hook_kwargs + ) + + try: + # Use standard pre_call_hook with call_type="mcp_call" + modified_data = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_auth, # type: ignore + data=synthetic_llm_data, + call_type="mcp_call", # type: ignore + ) + if modified_data: + # Convert response back to MCP format and apply modifications + modified_kwargs = ( + proxy_logging_obj._convert_mcp_hook_response_to_kwargs( + modified_data, pre_hook_kwargs + ) + ) + if modified_kwargs.get("arguments") != arguments: + arguments = modified_kwargs["arguments"] + + 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 pre call: {str(e)}") + raise e + async def call_tool( self, name: str, @@ -614,71 +685,14 @@ class MCPServerManager: # Using standard pre_call_hook with call_type="mcp_call" ######################################################### if proxy_logging_obj: - pre_hook_kwargs = { - "name": name, - "arguments": arguments, - "server_name": server_name_from_prefix, - "user_api_key_auth": user_api_key_auth, - "user_api_key_user_id": ( - getattr(user_api_key_auth, "user_id", None) - if user_api_key_auth - else None - ), - "user_api_key_team_id": ( - getattr(user_api_key_auth, "team_id", None) - if user_api_key_auth - else None - ), - "user_api_key_end_user_id": ( - getattr(user_api_key_auth, "end_user_id", None) - if user_api_key_auth - else None - ), - "user_api_key_hash": ( - getattr(user_api_key_auth, "api_key_hash", None) - if user_api_key_auth - else None - ), - } - - # Create MCP request object for processing - mcp_request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs( - pre_hook_kwargs + await self.pre_call_tool_check( + name=original_tool_name, + arguments=arguments, + server_name_from_prefix=server_name_from_prefix, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, ) - # Convert to LLM format for existing guardrail compatibility - synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format( - mcp_request_obj, pre_hook_kwargs - ) - - try: - # Use standard pre_call_hook with call_type="mcp_call" - modified_data = await proxy_logging_obj.pre_call_hook( - user_api_key_dict=user_api_key_auth, # type: ignore - data=synthetic_llm_data, - call_type="mcp_call", # type: ignore - ) - if modified_data: - # Convert response back to MCP format and apply modifications - modified_kwargs = ( - proxy_logging_obj._convert_mcp_hook_response_to_kwargs( - modified_data, pre_hook_kwargs - ) - ) - if modified_kwargs.get("arguments") != arguments: - arguments = modified_kwargs["arguments"] - - 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 pre call: {str(e)}" - ) - raise e - # Get server-specific auth header if available server_auth_header = None if mcp_server_auth_headers and mcp_server.alias: diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index d2f16154f3..ef770a9d43 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,11 +1,7 @@ import importlib from typing import Dict, List, Optional -from urllib.parse import urlencode, urlparse, urlunparse -import httpx -from fastapi import APIRouter, Depends, HTTPException, Query, Request -from fastapi.params import Form -from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse +from fastapi import APIRouter, Depends, Query, Request from litellm._logging import verbose_logger from litellm.proxy._types import UserAPIKeyAuth diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 7f16179390..38b0cfd99a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -364,7 +364,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers: Optional[Dict[str, str]] = None, oauth2_headers: Optional[Dict[str, str]] = None, ) -> List[MCPTool]: - f""" + """ Helper method to fetch tools from MCP servers based on server filtering criteria. Args: From 3a31e9cfa0766a5347db5aa4d2308f23249f17d3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 09:39:13 -0700 Subject: [PATCH 16/76] fix: fix perf issue --- .../mcp_server/discoverable_endpoints.py | 31 +++++++++++-------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 05dae916af..af586098be 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -6,6 +6,10 @@ import httpx from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -134,19 +138,20 @@ async def token_endpoint( proxy_base_url = str(request.base_url).rstrip("/") # Exchange code for real GitHub token - async with httpx.AsyncClient() as client: - resp = await client.post( - mcp_server.token_url, - headers={"Accept": "application/json"}, - data={ - "client_id": mcp_server.client_id, - "client_secret": mcp_server.client_secret, - "code": code, - "redirect_uri": f"{proxy_base_url}/callback", - }, - ) - resp.raise_for_status() - github_token = resp.json()["access_token"] + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) + response = await async_client.post( + mcp_server.token_url, + headers={"Accept": "application/json"}, + data={ + "client_id": mcp_server.client_id, + "client_secret": mcp_server.client_secret, + "code": code, + "redirect_uri": f"{proxy_base_url}/callback", + }, + ) + + response.raise_for_status() + github_token = response.json()["access_token"] # Return to Claude in expected OAuth 2 format From 2ba52a1f5bc2f50a26abd643e24edc11c1a8e585 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 27 Sep 2025 09:47:16 -0700 Subject: [PATCH 17/76] Revert "Azure Managed Batches - api key error (#14932)" This reverts commit b039374c718ecf7089ba74eca5a4aa95b6df6a09. --- litellm/llms/azure/batches/handler.py | 8 ++++---- litellm/llms/azure/common_utils.py | 13 +------------ litellm/llms/azure/files/handler.py | 10 +++++----- 3 files changed, 10 insertions(+), 21 deletions(-) diff --git a/litellm/llms/azure/batches/handler.py b/litellm/llms/azure/batches/handler.py index c038f62561..7fc6388ba8 100644 --- a/litellm/llms/azure/batches/handler.py +++ b/litellm/llms/azure/batches/handler.py @@ -62,7 +62,7 @@ class AzureBatchesAPI(BaseAzureLLM): ) if azure_client is None: raise ValueError( - "Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment." + "OpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment." ) if _is_async is True: @@ -108,7 +108,7 @@ class AzureBatchesAPI(BaseAzureLLM): ) if azure_client is None: raise ValueError( - "Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment." + "OpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment." ) if _is_async is True: @@ -156,7 +156,7 @@ class AzureBatchesAPI(BaseAzureLLM): ) if azure_client is None: raise ValueError( - "Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment." + "OpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment." ) response = azure_client.batches.cancel(**cancel_batch_data) return response @@ -195,7 +195,7 @@ class AzureBatchesAPI(BaseAzureLLM): ) if azure_client is None: raise ValueError( - "Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment." + "OpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment." ) if _is_async is True: diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 15569581ee..09b1888e04 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -576,19 +576,8 @@ class BaseAzureLLM(BaseOpenAILLM): verbose_logger.debug( f"Initializing Azure OpenAI Client for {model_name}, Api Base: {str(api_base)}, Api Key:{_api_key}" ) - - # Extract API key from multiple sources with proper precedence - resolved_api_key = ( - api_key - or litellm_params.get("api_key") - or litellm.api_key - or litellm.azure_key - or get_secret_str("AZURE_OPENAI_API_KEY") - or get_secret_str("AZURE_API_KEY") - ) - azure_client_params = { - "api_key": resolved_api_key, + "api_key": api_key, "azure_endpoint": api_base, "api_version": api_version, "azure_ad_token": azure_ad_token, diff --git a/litellm/llms/azure/files/handler.py b/litellm/llms/azure/files/handler.py index d7b16d304d..50c122ccf2 100644 --- a/litellm/llms/azure/files/handler.py +++ b/litellm/llms/azure/files/handler.py @@ -58,7 +58,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): ) if openai_client is None: raise ValueError( - "Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment." + "AzureOpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment." ) if _is_async is True: @@ -106,7 +106,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): ) if openai_client is None: raise ValueError( - "Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment." + "AzureOpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment." ) if _is_async is True: @@ -156,7 +156,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): ) if openai_client is None: raise ValueError( - "Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment." + "AzureOpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment." ) if _is_async is True: @@ -208,7 +208,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): ) if openai_client is None: raise ValueError( - "Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment." + "AzureOpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment." ) if _is_async is True: @@ -262,7 +262,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): ) if openai_client is None: raise ValueError( - "Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment." + "AzureOpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment." ) if _is_async is True: From 1f42ac1aece5d1e996fdebf05f70cbc2c181b043 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 09:48:16 -0700 Subject: [PATCH 18/76] test_async_completion_azure_caching_streaming --- .../test_custom_callback_router.py | 25 ++++++++++++++++--- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/tests/logging_callback_tests/test_custom_callback_router.py b/tests/logging_callback_tests/test_custom_callback_router.py index 27d9e17edf..5948fb6a2e 100644 --- a/tests/logging_callback_tests/test_custom_callback_router.py +++ b/tests/logging_callback_tests/test_custom_callback_router.py @@ -664,8 +664,25 @@ async def test_async_completion_azure_caching_streaming(): ) litellm.callbacks = [customHandler_caching] unique_time = uuid.uuid4() - response1 = await litellm.acompletion( - model="azure/gpt-4.1-nano", + + # Use Router instead of direct litellm.acompletion to get router-specific metadata + model_list = [ + { + "model_name": "gpt-4.1-nano", + "litellm_params": { + "model": "azure/gpt-4.1-nano", + "api_key": os.getenv("AZURE_API_KEY"), + "api_version": os.getenv("AZURE_API_VERSION"), + "api_base": os.getenv("AZURE_API_BASE"), + }, + "tpm": 240000, + "rpm": 1800, + }, + ] + router = Router(model_list=model_list) + + response1 = await router.acompletion( + model="gpt-4.1-nano", messages=[ {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} ], @@ -677,8 +694,8 @@ async def test_async_completion_azure_caching_streaming(): await asyncio.sleep(1) initial_customhandler_caching_states = len(customHandler_caching.states) print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}") - response2 = await litellm.acompletion( - model="azure/gpt-4.1-nano", + response2 = await router.acompletion( + model="gpt-4.1-nano", messages=[ {"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"} ], From d6e882163afecaa719338aa06051a8fd7d06cb3f Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 09:53:37 -0700 Subject: [PATCH 19/76] test_acompletion_caching_on_router_caching_groups --- tests/local_testing/test_router_caching.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/local_testing/test_router_caching.py b/tests/local_testing/test_router_caching.py index d9ee4fdeb5..6f92d5bea4 100644 --- a/tests/local_testing/test_router_caching.py +++ b/tests/local_testing/test_router_caching.py @@ -275,7 +275,7 @@ async def test_acompletion_caching_on_router_caching_groups(): "rpm": 10000, }, { - "model_name": "azure-gpt-4.1-nano", + "model_name": "azure-gpt-3.5-turbo", "litellm_params": { "model": "azure/gpt-4.1-nano", "api_key": os.getenv("AZURE_API_KEY"), From 3a8e4f31600f8dce7a420e13adfbf100e406847c Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 09:54:54 -0700 Subject: [PATCH 20/76] test_async_embedding_azure_caching --- .../test_custom_callback_router.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/logging_callback_tests/test_custom_callback_router.py b/tests/logging_callback_tests/test_custom_callback_router.py index 5948fb6a2e..1e3f5f4cfb 100644 --- a/tests/logging_callback_tests/test_custom_callback_router.py +++ b/tests/logging_callback_tests/test_custom_callback_router.py @@ -725,16 +725,22 @@ async def test_async_embedding_azure_caching(): port=os.environ["REDIS_PORT"], password=os.environ["REDIS_PASSWORD"], ) + router = Router(model_list=[{ + "model_name": "text-embedding-ada-002", + "litellm_params": { + "model": "openai/text-embedding-ada-002", + }, + }]) litellm.callbacks = [customHandler_caching] unique_time = time.time() - response1 = await litellm.aembedding( - model="azure/text-embedding-ada-002", + response1 = await router.aembedding( + model="text-embedding-ada-002", input=[f"good morning from litellm1 {unique_time}"], caching=True, ) await asyncio.sleep(1) # set cache is async for aembedding() - response2 = await litellm.aembedding( - model="azure/text-embedding-ada-002", + response2 = await router.aembedding( + model="text-embedding-ada-002", input=[f"good morning from litellm1 {unique_time}"], caching=True, ) From 6cca81291b924def25c6865f5aee3b6a6d99193c Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 09:57:07 -0700 Subject: [PATCH 21/76] test_acompletion_caching_on_router --- tests/local_testing/test_router_caching.py | 26 +++++++--------------- 1 file changed, 8 insertions(+), 18 deletions(-) diff --git a/tests/local_testing/test_router_caching.py b/tests/local_testing/test_router_caching.py index 6f92d5bea4..52fc5b8858 100644 --- a/tests/local_testing/test_router_caching.py +++ b/tests/local_testing/test_router_caching.py @@ -6,7 +6,7 @@ import sys import time import traceback from unittest.mock import patch - +from typing import Union import pytest sys.path.insert( @@ -85,21 +85,11 @@ async def test_acompletion_caching_on_router(): litellm.set_verbose = True model_list = [ { - "model_name": "gpt-3.5-turbo", + "model_name": "gpt-4.1-nano", "litellm_params": { - "model": "gpt-3.5-turbo", + "model": "gpt-4.1-nano", "api_key": os.getenv("OPENAI_API_KEY"), - }, - "tpm": 100000, - "rpm": 10000, - }, - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "azure/chatgpt-v-3", - "api_key": os.getenv("AZURE_API_KEY"), - "api_base": os.getenv("AZURE_API_BASE"), - "api_version": os.getenv("AZURE_API_VERSION"), + "mock_response": "Hello world", }, "tpm": 100000, "rpm": 10000, @@ -120,13 +110,13 @@ async def test_acompletion_caching_on_router(): routing_strategy="simple-shuffle", ) response1 = await router.acompletion( - model="gpt-3.5-turbo", messages=messages, temperature=1 + model="gpt-4.1-nano", messages=messages, temperature=1 ) print(f"response1: {response1}") await asyncio.sleep(5) # add cache is async, async sleep for cache to get set response2 = await router.acompletion( - model="gpt-3.5-turbo", messages=messages, temperature=1 + model="gpt-4.1-nano", messages=messages, temperature=1 ) print(f"response2: {response2}") assert response1.id == response2.id @@ -341,8 +331,8 @@ async def test_acompletion_caching_on_router_caching_groups(): ], ) def test_create_correct_redis_cache_instance( - startup_nodes: list[dict] | None, - expected_cache_type: type[RedisClusterCache | RedisCache], + startup_nodes: Union[list[dict], None], + expected_cache_type: Union[type[RedisClusterCache], type[RedisCache]], ): cache_config = dict( host="mockhost", From a86e78f13ad6069e2dafaabd2fb03d5659ae8858 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 09:58:00 -0700 Subject: [PATCH 22/76] test_azure_ai_request_format --- tests/llm_translation/test_azure_ai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/llm_translation/test_azure_ai.py b/tests/llm_translation/test_azure_ai.py index 9f648aee05..b38a58a7ca 100644 --- a/tests/llm_translation/test_azure_ai.py +++ b/tests/llm_translation/test_azure_ai.py @@ -283,7 +283,7 @@ async def test_azure_ai_request_format(): # Set up the test parameters api_key = os.getenv("AZURE_API_KEY") api_base = f"{os.getenv('AZURE_API_BASE')}/openai/deployments/gpt-4o-new-test/chat/completions?api-version=2024-08-01-preview" - model = "azure_ai/gpt-4o" + model = "azure_ai/gpt-4.1-nano" messages = [ {"role": "user", "content": "hi"}, {"role": "assistant", "content": "Hello! How can I assist you today?"}, From 5787f868612ceb259b8aaf6be59646b6983f85c0 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 09:59:47 -0700 Subject: [PATCH 23/76] fix: fix ruff check --- litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index af586098be..9c4ab70f0d 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -2,7 +2,6 @@ import json from typing import Optional from urllib.parse import urlencode, urlparse, urlunparse -import httpx from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse From 9e7ee0614a3c802a8852c12a31b2d33c42f7d105 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:01:14 -0700 Subject: [PATCH 24/76] test fix --- tests/test_models.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index a0fa052527..b1a5b84b3b 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -83,10 +83,8 @@ async def add_models( data = { "model_name": model_name, "litellm_params": { - "model": "azure/chatgpt-v-3", - "api_key": "os.environ/AZURE_API_KEY", - "api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/", - "api_version": "2023-05-15", + "model": "openai/gpt-4.1-nano", + "api_key": "os.environ/OPENAI_API_KEY", }, "model_info": {"id": model_id}, } @@ -119,10 +117,8 @@ async def update_model( data = { "model_name": model_name, "litellm_params": { - "model": "azure/chatgpt-v-3", - "api_key": "os.environ/AZURE_API_KEY", - "api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/", - "api_version": "2023-05-15", + "model": "openai/gpt-4.1-nano", + "api_key": "os.environ/OPENAI_API_KEY", }, "model_info": {"id": model_id}, } From cf1ed225d7f4d02116486269cd071ecb0741fca6 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 10:01:48 -0700 Subject: [PATCH 25/76] fix: fix vllm test --- .../test_hosted_vllm_rerank_transformation.py | 37 +++++++++++++------ 1 file changed, 26 insertions(+), 11 deletions(-) diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index 3d4d0e7b70..9e6fa608c5 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -1,8 +1,19 @@ -import sys import os +import sys + import pytest + from litellm.llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig -from litellm.types.rerank import OptionalRerankParams, RerankResponse, RerankResponseResult, RerankResponseMeta, RerankBilledUnits, RerankTokens, RerankResponseDocument +from litellm.types.rerank import ( + OptionalRerankParams, + RerankBilledUnits, + RerankResponse, + RerankResponseDocument, + RerankResponseMeta, + RerankResponseResult, + RerankTokens, +) + class TestHostedVLLMRerankTransform: def setup_method(self): @@ -27,23 +38,27 @@ class TestHostedVLLMRerankTransform: assert params["return_documents"] is True def test_map_cohere_rerank_params_raises_on_max_chunks_per_doc(self): - with pytest.raises(ValueError, match="Hosted VLLM does not support max_chunks_per_doc"): + with pytest.raises( + ValueError, match="Hosted VLLM does not support max_chunks_per_doc" + ): self.config.map_cohere_rerank_params( non_default_params=None, model=self.model, drop_params=False, query="test query", documents=["doc1"], - max_chunks_per_doc=5 + max_chunks_per_doc=5, ) def test_get_complete_url(self): base = "https://api.example.com" url = self.config.get_complete_url(base, self.model) - assert url == "https://api.example.com/v1/rerank" - # Already ends with /v1/rerank - url2 = self.config.get_complete_url("https://api.example.com/v1/rerank", self.model) - assert url2 == "https://api.example.com/v1/rerank" + assert url == "https://api.example.com/rerank" + # Already ends with /rerank + url2 = self.config.get_complete_url( + "https://api.example.com/rerank", self.model + ) + assert url2 == "https://api.example.com/rerank" # Raises if api_base is None with pytest.raises(ValueError): self.config.get_complete_url(None, self.model) @@ -55,7 +70,7 @@ class TestHostedVLLMRerankTransform: {"index": 0, "relevance_score": 0.9, "document": {"text": "doc1 text"}}, {"index": 1, "relevance_score": 0.7, "document": {"text": "doc2 text"}}, ], - "usage": {"total_tokens": 42} + "usage": {"total_tokens": 42}, } result = self.config._transform_response(response_dict) assert result.id == "abc123" @@ -75,7 +90,7 @@ class TestHostedVLLMRerankTransform: response_dict = { "id": "abc123", "results": [{"relevance_score": 0.5}], - "usage": {"total_tokens": 10} + "usage": {"total_tokens": 10}, } with pytest.raises(ValueError, match="Missing required fields in the result="): - self.config._transform_response(response_dict) \ No newline at end of file + self.config._transform_response(response_dict) From 7278ae5305f407b6fcafc814fc7f49dcac85c142 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 10:03:54 -0700 Subject: [PATCH 26/76] test: fix tests --- .../_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 2a5cc7bcb7..fb15dbc14e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -543,6 +543,7 @@ class TestMCPRequestHandler: mcp_auth_header, mcp_servers_result, mcp_server_auth_headers, + oauth2_headers, ) = await MCPRequestHandler.process_mcp_request(scope) assert auth_result == mock_auth_result assert mcp_auth_header == expected_result["mcp_auth"] From 9220ddc598a7903364b95ea5697c6995b74b9a7f Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:04:02 -0700 Subject: [PATCH 27/76] test router policy violation --- .../test_router_policy_violation.py | 137 ------------------ 1 file changed, 137 deletions(-) delete mode 100644 tests/local_testing/test_router_policy_violation.py diff --git a/tests/local_testing/test_router_policy_violation.py b/tests/local_testing/test_router_policy_violation.py deleted file mode 100644 index 1e72868db6..0000000000 --- a/tests/local_testing/test_router_policy_violation.py +++ /dev/null @@ -1,137 +0,0 @@ -#### What this tests #### -# This tests if the router sends back a policy violation, without retries - -import sys, os, time -import traceback, asyncio -import pytest - -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path - -import litellm -from litellm import Router -from litellm.integrations.custom_logger import CustomLogger - - -class MyCustomHandler(CustomLogger): - success: bool = False - failure: bool = False - previous_models: int = 0 - - def log_pre_api_call(self, model, messages, kwargs): - print(f"Pre-API Call") - print( - f"previous_models: {kwargs['litellm_params']['metadata']['previous_models']}" - ) - self.previous_models += len( - kwargs["litellm_params"]["metadata"]["previous_models"] - ) # {"previous_models": [{"model": litellm_model_name, "exception_type": AuthenticationError, "exception_string": }]} - print(f"self.previous_models: {self.previous_models}") - - def log_post_api_call(self, kwargs, response_obj, start_time, end_time): - print( - f"Post-API Call - response object: {response_obj}; model: {kwargs['model']}" - ) - - def log_stream_event(self, kwargs, response_obj, start_time, end_time): - print(f"On Stream") - - def async_log_stream_event(self, kwargs, response_obj, start_time, end_time): - print(f"On Stream") - - def log_success_event(self, kwargs, response_obj, start_time, end_time): - print(f"On Success") - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - print(f"On Success") - - def log_failure_event(self, kwargs, response_obj, start_time, end_time): - print(f"On Failure") - - -kwargs = { - "model": "azure/gpt-3.5-turbo", - "messages": [ - {"role": "system", "content": "You are a helpful assistant."}, - { - "role": "user", - "content": "vorrei vedere la cosa più bella ad Ercolano. Qual’è?", - }, - ], -} - - -@pytest.mark.asyncio -async def test_async_fallbacks(): - litellm.set_verbose = False - model_list = [ - { # list of model deployments - "model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name - "litellm_params": { # params for litellm completion/embedding call - "model": "azure/chatgpt-v-3", - "api_key": os.getenv("AZURE_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), - }, - "tpm": 240000, - "rpm": 1800, - }, - { - "model_name": "azure/gpt-3.5-turbo", # openai model name - "litellm_params": { # params for litellm completion/embedding call - "model": "azure/chatgpt-functioncalling", - "api_key": os.getenv("AZURE_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), - }, - "tpm": 240000, - "rpm": 1800, - }, - { - "model_name": "gpt-3.5-turbo", # openai model name - "litellm_params": { # params for litellm completion/embedding call - "model": "gpt-3.5-turbo", - "api_key": os.getenv("OPENAI_API_KEY"), - }, - "tpm": 1000000, - "rpm": 9000, - }, - { - "model_name": "gpt-3.5-turbo-16k", # openai model name - "litellm_params": { # params for litellm completion/embedding call - "model": "gpt-3.5-turbo-16k", - "api_key": os.getenv("OPENAI_API_KEY"), - }, - "tpm": 1000000, - "rpm": 9000, - }, - ] - - router = Router( - model_list=model_list, - num_retries=3, - fallbacks=[{"azure/gpt-3.5-turbo": ["gpt-3.5-turbo"]}], - # context_window_fallbacks=[ - # {"azure/gpt-3.5-turbo-context-fallback": ["gpt-3.5-turbo-16k"]}, - # {"gpt-3.5-turbo": ["gpt-3.5-turbo-16k"]}, - # ], - set_verbose=False, - ) - customHandler = MyCustomHandler() - litellm.callbacks = [customHandler] - try: - response = await router.acompletion(**kwargs) - pytest.fail( - f"An exception occurred: {e}" - ) # should've raised azure policy error - except litellm.Timeout as e: - pass - except Exception as e: - await asyncio.sleep( - 0.05 - ) # allow a delay as success_callbacks are on a separate thread - assert customHandler.previous_models == 0 # 0 retries, 0 fallback - router.reset() - finally: - router.reset() From c9c52b3ccc3a0186331ec5972006ad87a6f0b0e1 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:05:57 -0700 Subject: [PATCH 28/76] test_azure_ai_request_format --- tests/llm_translation/test_azure_ai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/llm_translation/test_azure_ai.py b/tests/llm_translation/test_azure_ai.py index b38a58a7ca..943020c1bb 100644 --- a/tests/llm_translation/test_azure_ai.py +++ b/tests/llm_translation/test_azure_ai.py @@ -282,7 +282,7 @@ async def test_azure_ai_request_format(): # Set up the test parameters api_key = os.getenv("AZURE_API_KEY") - api_base = f"{os.getenv('AZURE_API_BASE')}/openai/deployments/gpt-4o-new-test/chat/completions?api-version=2024-08-01-preview" + api_base = os.getenv("AZURE_API_BASE") model = "azure_ai/gpt-4.1-nano" messages = [ {"role": "user", "content": "hi"}, From a2e72fbcc91bed370e68942f3faf0b857d312943 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:08:32 -0700 Subject: [PATCH 29/76] test_caching_with_models_v2 --- tests/local_testing/test_caching.py | 32 ++++++++++++++--------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/tests/local_testing/test_caching.py b/tests/local_testing/test_caching.py index 39019ef35a..8bd73af609 100644 --- a/tests/local_testing/test_caching.py +++ b/tests/local_testing/test_caching.py @@ -325,7 +325,7 @@ def test_caching_with_models_v2(): litellm.set_verbose = True response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True) response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True) - response3 = completion(model="azure/chatgpt-v-3", messages=messages, caching=True) + response3 = completion(model="gpt-4.1-nano", messages=messages, caching=True) print(f"response1: {response1}") print(f"response2: {response2}") print(f"response3: {response3}") @@ -527,7 +527,7 @@ def test_embedding_caching_azure(): print(api_key) print(api_base) embedding1 = embedding( - model="azure/azure-embedding-model", + model="text-embedding-ada-002", input=["good morning from litellm", "this is another item"], api_key=api_key, api_base=api_base, @@ -540,7 +540,7 @@ def test_embedding_caching_azure(): time.sleep(1) start_time = time.time() embedding2 = embedding( - model="azure/azure-embedding-model", + model="text-embedding-ada-002", input=["good morning from litellm", "this is another item"], api_key=api_key, api_base=api_base, @@ -595,10 +595,10 @@ async def test_embedding_caching_azure_individual_items(): ] embedding_val_1 = await aembedding( - model="azure/azure-embedding-model", input=embedding_1, caching=True + model="text-embedding-ada-002", input=embedding_1, caching=True ) embedding_val_2 = await aembedding( - model="azure/azure-embedding-model", input=embedding_2, caching=True + model="text-embedding-ada-002", input=embedding_2, caching=True ) print(f"embedding_val_2._hidden_params: {embedding_val_2._hidden_params}") assert embedding_val_2._hidden_params["cache_hit"] == True @@ -633,11 +633,11 @@ async def test_embedding_caching_azure_individual_items_reordered(): ] embedding_val_1 = await aembedding( - model="azure/azure-embedding-model", input=embedding_1, caching=True + model="text-embedding-ada-002", input=embedding_1, caching=True ) print("embedding val 1", embedding_val_1) embedding_val_2 = await aembedding( - model="azure/azure-embedding-model", input=embedding_2, caching=True + model="text-embedding-ada-002", input=embedding_2, caching=True ) print("embedding val 2", embedding_val_2) print(f"embedding_val_2._hidden_params: {embedding_val_2._hidden_params}") @@ -667,7 +667,7 @@ async def test_embedding_caching_base_64(): ] embedding_val_1 = await aembedding( - model="azure/azure-embedding-model", + model="text-embedding-ada-002", input=inputs, caching=True, encoding_format="base64", @@ -675,7 +675,7 @@ async def test_embedding_caching_base_64(): await asyncio.sleep(5) print("\n\nCALL2\n\n") embedding_val_2 = await aembedding( - model="azure/azure-embedding-model", + model="text-embedding-ada-002", input=inputs, caching=True, encoding_format="base64", @@ -718,7 +718,7 @@ async def test_embedding_caching_redis_ttl(): # Call the embedding method embedding_val_1 = await litellm.aembedding( - model="azure/azure-embedding-model", + model="text-embedding-ada-002", input=inputs, encoding_format="base64", ) @@ -1226,7 +1226,7 @@ async def test_s3_cache_stream_azure(sync_mode): if sync_mode: response1 = litellm.completion( - model="azure/chatgpt-v-3", + model="azure/gpt-4.1-nano", messages=messages, max_tokens=40, temperature=1, @@ -1239,7 +1239,7 @@ async def test_s3_cache_stream_azure(sync_mode): print(response_1_content) else: response1 = await litellm.acompletion( - model="azure/chatgpt-v-3", + model="azure/gpt-4.1-nano", messages=messages, max_tokens=40, temperature=1, @@ -1259,7 +1259,7 @@ async def test_s3_cache_stream_azure(sync_mode): if sync_mode: response2 = litellm.completion( - model="azure/chatgpt-v-3", + model="azure/gpt-4.1-nano", messages=messages, max_tokens=40, temperature=1, @@ -1272,7 +1272,7 @@ async def test_s3_cache_stream_azure(sync_mode): print(response_2_content) else: response2 = await litellm.acompletion( - model="azure/chatgpt-v-3", + model="azure/gpt-4.1-nano", messages=messages, max_tokens=40, temperature=1, @@ -1335,7 +1335,7 @@ async def test_s3_cache_acompletion_azure(): print("s3 Cache: test for caching, streaming + completion") response1 = await litellm.acompletion( - model="azure/chatgpt-v-3", + model="azure/gpt-4.1-nano", messages=messages, max_tokens=40, temperature=1, @@ -1345,7 +1345,7 @@ async def test_s3_cache_acompletion_azure(): time.sleep(2) response2 = await litellm.acompletion( - model="azure/chatgpt-v-3", + model="azure/gpt-4.1-nano", messages=messages, max_tokens=40, temperature=1, From fdfff392226d876740f64763256220179aa0ff53 Mon Sep 17 00:00:00 2001 From: Teddy Amkie Date: Sat, 27 Sep 2025 10:13:16 -0700 Subject: [PATCH 30/76] Add Azure passthrough documentation and sidebar entry (#14958) Co-authored-by: Cursor Agent --- .../docs/pass_through/azure_passthrough.md | 89 +++++++++++++++++++ docs/my-website/sidebars.js | 1 + 2 files changed, 90 insertions(+) create mode 100644 docs/my-website/docs/pass_through/azure_passthrough.md diff --git a/docs/my-website/docs/pass_through/azure_passthrough.md b/docs/my-website/docs/pass_through/azure_passthrough.md new file mode 100644 index 0000000000..cac0633358 --- /dev/null +++ b/docs/my-website/docs/pass_through/azure_passthrough.md @@ -0,0 +1,89 @@ +# Azure Passthrough + +Pass-through endpoints for `/azure` + +## Overview + +| Feature | Supported | Notes | +|-------|-------|-------| +| Cost Tracking | ❌ | Not supported | +| Logging | ✅ | Works across all integrations | +| Streaming | ✅ | Fully supported | + +### When to use this? + +- For most use cases, you should use the [native LiteLLM Azure OpenAI Integration](../providers/azure/azure) (`/chat/completions`, `/embeddings`, `/completions`, `/images`, etc.) +- Use this passthrough to call newer or less common Azure OpenAI endpoints that LiteLLM doesn't fully support yet, such as `/assistants`, `/threads`, `/vector_stores` + +Simply replace your Azure endpoint (e.g. `https://.openai.azure.com`) with `LITELLM_PROXY_BASE_URL/azure` + +## Usage Examples + +### Assistants API + +#### Create Azure OpenAI Client + +Make sure you do the following: +- Point `azure_endpoint` to your `LITELLM_PROXY_BASE_URL/azure` +- Use your `LITELLM_API_KEY` as the `api_key` + +```python +import openai + +client = openai.AzureOpenAI( + azure_endpoint="http://0.0.0.0:4000/azure", # /azure + api_key="sk-anything", # + api_version="2024-05-01-preview" # required Azure API version +) +``` + +#### Create an Assistant + +```python +assistant = client.beta.assistants.create( + name="Math Tutor", + instructions="You are a math tutor. Help solve equations.", + model="gpt-4o", +) +``` + +#### Create a Thread +```python +thread = client.beta.threads.create() +``` + +#### Add a Message to the Thread +```python +message = client.beta.threads.messages.create( + thread_id=thread.id, + role="user", + content="Solve 3x + 11 = 14", +) +``` + +#### Run the Assistant +```python +run = client.beta.threads.runs.create( + thread_id=thread.id, + assistant_id=assistant.id, +) + +# Check run status +run_status = client.beta.threads.runs.retrieve( + thread_id=thread.id, + run_id=run.id +) +``` + +#### Retrieve Messages +```python +messages = client.beta.threads.messages.list( + thread_id=thread.id +) +``` + +#### Delete the Assistant + +```python +client.beta.assistants.delete(assistant.id) +``` \ No newline at end of file diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index cb68e9024d..07110b938d 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -342,6 +342,7 @@ const sidebars = { "pass_through/anthropic_completion", "pass_through/assembly_ai", "pass_through/bedrock", + "pass_through/azure_passthrough", "pass_through/cohere", "pass_through/google_ai_studio", "pass_through/langfuse", From d5c716fab996e8e59d396d4f73f2c353f763acfc Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:15:19 -0700 Subject: [PATCH 31/76] fix CustomLoggerRegistry --- litellm/litellm_core_utils/custom_logger_registry.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index fbf685c370..75dc0918ab 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -15,6 +15,7 @@ from litellm.integrations.agentops import AgentOps from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.integrations.argilla import ArgillaLogger from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger +from litellm.integrations.bitbucket import BitBucketPromptManager from litellm.integrations.braintrust_logging import BraintrustLogger from litellm.integrations.datadog.datadog import DataDogLogger from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger @@ -94,6 +95,7 @@ class CustomLoggerRegistry: "bitbucket": BitBucketPromptManager, "cloudzero": CloudZeroLogger, "posthog": PostHogLogger, + "bitbucket": BitBucketPromptManager, } try: From f2f839138445a9c66eccf5517c5ce6da761cbd96 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:17:16 -0700 Subject: [PATCH 32/76] test azure openai fix --- tests/llm_translation/test_azure_openai.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/llm_translation/test_azure_openai.py b/tests/llm_translation/test_azure_openai.py index 5dab7e6d68..33947e25a6 100644 --- a/tests/llm_translation/test_azure_openai.py +++ b/tests/llm_translation/test_azure_openai.py @@ -339,7 +339,7 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version): with patch.object(client.chat.completions.with_raw_response, "create") as mock_post: response = litellm.completion( - model="azure/gpt-4o-new-test", + model="azure/gpt-4.1-nano", messages=[ { "role": "system", @@ -474,7 +474,7 @@ def test_azure_max_retries_0( try: completion( - model="azure/gpt-4o-new-test", + model="azure/gpt-4.1-nano", messages=[{"role": "user", "content": "Hello world"}], max_retries=max_retries, stream=stream, @@ -502,7 +502,7 @@ async def test_async_azure_max_retries_0( try: await acompletion( - model="azure/gpt-4o-new-test", + model="azure/gpt-4.1-nano", messages=[{"role": "user", "content": "Hello world"}], max_retries=max_retries, stream=stream, @@ -565,7 +565,7 @@ async def test_azure_embedding_max_retries_0( from litellm import aembedding, embedding args = { - "model": "azure/azure-embedding-model", + "model": "text-embedding-ada-002", "input": "Hello world", "max_retries": max_retries, } @@ -598,7 +598,7 @@ def test_azure_safety_result(): litellm._turn_on_debug() response = completion( - model="azure/gpt-4o-new-test", + model="azure/gpt-4.1-nano", messages=[{"role": "user", "content": "Hello world"}], ) print(f"response: {response}") From 2b3966329101d74e9e97c2c6a71f33e3cde3738d Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:18:18 -0700 Subject: [PATCH 33/76] test_async_fallbacks_embeddings --- tests/local_testing/test_router_fallbacks.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index 5463374268..2159c77f0d 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -250,10 +250,8 @@ def test_sync_fallbacks_embeddings(): { # list of model deployments "model_name": "good-azure-embedding-model", # openai model name "litellm_params": { # params for litellm completion/embedding call - "model": "azure/azure-embedding-model", - "api_key": os.getenv("AZURE_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "model": "text-embedding-ada-002", + "api_key": os.getenv("OPENAI_API_KEY"), }, "tpm": 240000, "rpm": 1800, @@ -302,10 +300,8 @@ async def test_async_fallbacks_embeddings(): { # list of model deployments "model_name": "good-azure-embedding-model", # openai model name "litellm_params": { # params for litellm completion/embedding call - "model": "azure/azure-embedding-model", - "api_key": os.getenv("AZURE_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "model": "text-embedding-ada-002", + "api_key": os.getenv("OPENAI_API_KEY"), }, "tpm": 240000, "rpm": 1800, From 516e169f9e504f6fa99e7607e158c4e75e5645af Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 10:22:59 -0700 Subject: [PATCH 34/76] ci: use mypy ini in github linting tests --- .../llm_cost_calc/tool_call_cost_tracking.py | 171 +++++++++++------- 1 file changed, 110 insertions(+), 61 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 21ff44ab08..52175989cb 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -47,7 +47,7 @@ class StandardBuiltInToolCostTracking: - Code Interpreter (Azure) """ standard_built_in_tools_params = standard_built_in_tools_params or {} - + # Handle web search if StandardBuiltInToolCostTracking.response_object_includes_web_search_call( response_object=response_object, usage=usage @@ -58,7 +58,7 @@ class StandardBuiltInToolCostTracking: usage=usage, standard_built_in_tools_params=standard_built_in_tools_params, ) - + # Handle file search if StandardBuiltInToolCostTracking.response_object_includes_file_search_call( response_object=response_object @@ -68,7 +68,7 @@ class StandardBuiltInToolCostTracking: custom_llm_provider=custom_llm_provider, standard_built_in_tools_params=standard_built_in_tools_params, ) - + # Handle Azure assistant features return StandardBuiltInToolCostTracking._handle_azure_assistant_costs( model=model, @@ -85,14 +85,14 @@ class StandardBuiltInToolCostTracking: ) -> float: """Handle web search cost calculation.""" from litellm.llms import get_cost_for_web_search_request - + model_info = StandardBuiltInToolCostTracking._safe_get_model_info( model=model, custom_llm_provider=custom_llm_provider ) - + if custom_llm_provider is None and model_info is not None: custom_llm_provider = model_info["litellm_provider"] - + if ( model_info is not None and usage is not None @@ -105,9 +105,11 @@ class StandardBuiltInToolCostTracking: ) if result is not None: return result - + return StandardBuiltInToolCostTracking.get_cost_for_web_search( - web_search_options=standard_built_in_tools_params.get("web_search_options", None), + web_search_options=standard_built_in_tools_params.get( + "web_search_options", None + ), model_info=model_info, ) @@ -121,12 +123,17 @@ class StandardBuiltInToolCostTracking: model_info = StandardBuiltInToolCostTracking._safe_get_model_info( model=model, custom_llm_provider=custom_llm_provider ) - file_search_usage = standard_built_in_tools_params.get("file_search", {}) - + file_search_raw = standard_built_in_tools_params.get("file_search", {}) + file_search_usage: Optional[FileSearchTool] = ( + FileSearchTool(**file_search_raw) if file_search_raw else None + ) + # Convert model_info to dict and extract usage parameters model_info_dict = dict(model_info) if model_info is not None else None - storage_gb, days = StandardBuiltInToolCostTracking._extract_file_search_params(file_search_usage) - + storage_gb, days = StandardBuiltInToolCostTracking._extract_file_search_params( + file_search_usage + ) + return StandardBuiltInToolCostTracking.get_cost_for_file_search( file_search=file_search_usage, provider=custom_llm_provider, @@ -144,11 +151,11 @@ class StandardBuiltInToolCostTracking: """Handle Azure assistant features cost calculation.""" if custom_llm_provider != "azure": return 0.0 - + model_info = StandardBuiltInToolCostTracking._safe_get_model_info( model=model, custom_llm_provider=custom_llm_provider ) - + total_cost = 0.0 total_cost += StandardBuiltInToolCostTracking._get_vector_store_cost( model_info, custom_llm_provider, standard_built_in_tools_params @@ -159,31 +166,33 @@ class StandardBuiltInToolCostTracking: total_cost += StandardBuiltInToolCostTracking._get_code_interpreter_cost( model_info, custom_llm_provider, standard_built_in_tools_params ) - + return total_cost @staticmethod - def _extract_file_search_params(file_search_usage: Any) -> Tuple[Optional[float], Optional[float]]: + def _extract_file_search_params( + file_search_usage: Any, + ) -> Tuple[Optional[float], Optional[float]]: """Extract and convert file search parameters safely.""" storage_gb = None days = None - + if isinstance(file_search_usage, dict): storage_gb_val = file_search_usage.get("storage_gb") days_val = file_search_usage.get("days") - + if storage_gb_val is not None: try: storage_gb = float(storage_gb_val) # type: ignore except (TypeError, ValueError): storage_gb = None - + if days_val is not None: try: days = float(days_val) # type: ignore except (TypeError, ValueError): days = None - + return storage_gb, days @staticmethod @@ -193,13 +202,17 @@ class StandardBuiltInToolCostTracking: standard_built_in_tools_params: StandardBuiltInToolsParams, ) -> float: """Calculate vector store cost.""" - vector_store_usage = standard_built_in_tools_params.get("vector_store_usage", None) + vector_store_usage = standard_built_in_tools_params.get( + "vector_store_usage", None + ) if not vector_store_usage: return 0.0 - + model_info_dict = dict(model_info) if model_info is not None else None - vector_store_dict = vector_store_usage if isinstance(vector_store_usage, dict) else {} - + vector_store_dict = ( + vector_store_usage if isinstance(vector_store_usage, dict) else {} + ) + return StandardBuiltInToolCostTracking.get_cost_for_vector_store( vector_store_usage=vector_store_dict, provider=custom_llm_provider, @@ -213,13 +226,17 @@ class StandardBuiltInToolCostTracking: standard_built_in_tools_params: StandardBuiltInToolsParams, ) -> float: """Calculate computer use cost.""" - computer_use_usage = standard_built_in_tools_params.get("computer_use_usage", {}) + computer_use_usage = standard_built_in_tools_params.get( + "computer_use_usage", {} + ) if not computer_use_usage: return 0.0 - + model_info_dict = dict(model_info) if model_info is not None else None - input_tokens, output_tokens = StandardBuiltInToolCostTracking._extract_token_counts(computer_use_usage) - + input_tokens, output_tokens = ( + StandardBuiltInToolCostTracking._extract_token_counts(computer_use_usage) + ) + return StandardBuiltInToolCostTracking.get_cost_for_computer_use( input_tokens=input_tokens, output_tokens=output_tokens, @@ -234,13 +251,17 @@ class StandardBuiltInToolCostTracking: standard_built_in_tools_params: StandardBuiltInToolsParams, ) -> float: """Calculate code interpreter cost.""" - code_interpreter_sessions = standard_built_in_tools_params.get("code_interpreter_sessions", None) + code_interpreter_sessions = standard_built_in_tools_params.get( + "code_interpreter_sessions", None + ) if not code_interpreter_sessions: return 0.0 - + model_info_dict = dict(model_info) if model_info is not None else None - sessions = StandardBuiltInToolCostTracking._safe_convert_to_int(code_interpreter_sessions) - + sessions = StandardBuiltInToolCostTracking._safe_convert_to_int( + code_interpreter_sessions + ) + return StandardBuiltInToolCostTracking.get_cost_for_code_interpreter( sessions=sessions, provider=custom_llm_provider, @@ -248,18 +269,24 @@ class StandardBuiltInToolCostTracking: ) @staticmethod - def _extract_token_counts(computer_use_usage: Any) -> Tuple[Optional[int], Optional[int]]: + def _extract_token_counts( + computer_use_usage: Any, + ) -> Tuple[Optional[int], Optional[int]]: """Extract and convert token counts safely.""" input_tokens = None output_tokens = None - + if isinstance(computer_use_usage, dict): input_tokens_val = computer_use_usage.get("input_tokens") output_tokens_val = computer_use_usage.get("output_tokens") - - input_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int(input_tokens_val) - output_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int(output_tokens_val) - + + input_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int( + input_tokens_val + ) + output_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int( + output_tokens_val + ) + return input_tokens, output_tokens @staticmethod @@ -400,8 +427,11 @@ class StandardBuiltInToolCostTracking: if model_info is None: return 0.0 + search_context_raw = model_info.get("search_context_cost_per_query", {}) search_context_pricing: SearchContextCostPerQuery = ( - model_info.get("search_context_cost_per_query", {}) or {} + SearchContextCostPerQuery(**search_context_raw) + if search_context_raw + else SearchContextCostPerQuery() ) if web_search_options.get("search_context_size", None) == "low": return search_context_pricing.get("search_context_size_low", 0.0) @@ -424,9 +454,12 @@ class StandardBuiltInToolCostTracking: """ if model_info is None: return 0.0 + search_context_raw = model_info.get("search_context_cost_per_query", {}) or {} search_context_pricing: SearchContextCostPerQuery = ( - model_info.get("search_context_cost_per_query", {}) or {} - ) or {} + SearchContextCostPerQuery(**search_context_raw) + if search_context_raw + else SearchContextCostPerQuery() + ) return search_context_pricing.get("search_context_size_medium", 0.0) @staticmethod @@ -445,22 +478,27 @@ class StandardBuiltInToolCostTracking: """ if file_search is None: return 0.0 - + # Check if model-specific pricing is available - if model_info and "file_search_cost_per_gb_per_day" in model_info and provider == "azure": + if ( + model_info + and "file_search_cost_per_gb_per_day" in model_info + and provider == "azure" + ): if storage_gb and days: return storage_gb * days * model_info["file_search_cost_per_gb_per_day"] elif model_info and "file_search_cost_per_1k_calls" in model_info: return model_info["file_search_cost_per_1k_calls"] - + # Azure has storage-based pricing for file search if provider == "azure": from litellm.constants import AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY + if storage_gb and days: return storage_gb * days * AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY # Default to 0 if no storage info provided return 0.0 - + # Default to OpenAI pricing (per-call based) return OPENAI_FILE_SEARCH_COST_PER_1K_CALLS @@ -472,24 +510,25 @@ class StandardBuiltInToolCostTracking: ) -> float: """ Calculate cost for vector store usage. - + Azure charges based on storage size and duration. """ if vector_store_usage is None: return 0.0 - + storage_gb = vector_store_usage.get("storage_gb", 0.0) days = vector_store_usage.get("days", 0.0) - + # Check if model-specific pricing is available if model_info and "vector_store_cost_per_gb_per_day" in model_info: return storage_gb * days * model_info["vector_store_cost_per_gb_per_day"] - + # Azure has different pricing structure for vector store if provider == "azure": from litellm.constants import AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY + return storage_gb * days * AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY - + # OpenAI doesn't charge separately for vector store (included in embeddings) return 0.0 @@ -502,14 +541,18 @@ class StandardBuiltInToolCostTracking: ) -> float: """ Calculate cost for computer use feature. - + Azure: $0.003 USD per 1K input tokens, $0.012 USD per 1K output tokens """ if provider == "azure" and (input_tokens or output_tokens): # Check if model-specific pricing is available if model_info: - input_cost = model_info.get("computer_use_input_cost_per_1k_tokens", 0.0) - output_cost = model_info.get("computer_use_output_cost_per_1k_tokens", 0.0) + input_cost = model_info.get( + "computer_use_input_cost_per_1k_tokens", 0.0 + ) + output_cost = model_info.get( + "computer_use_output_cost_per_1k_tokens", 0.0 + ) if input_cost or output_cost: total_cost = 0.0 if input_tokens: @@ -517,19 +560,24 @@ class StandardBuiltInToolCostTracking: if output_tokens: total_cost += (output_tokens / 1000.0) * output_cost return total_cost - + # Azure default pricing from litellm.constants import ( AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS, AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS, ) + total_cost = 0.0 if input_tokens: - total_cost += (input_tokens / 1000.0) * AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS + total_cost += ( + input_tokens / 1000.0 + ) * AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS if output_tokens: - total_cost += (output_tokens / 1000.0) * AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS + total_cost += ( + output_tokens / 1000.0 + ) * AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS return total_cost - + # OpenAI doesn't charge separately for computer use yet return 0.0 @@ -541,21 +589,22 @@ class StandardBuiltInToolCostTracking: ) -> float: """ Calculate cost for code interpreter feature. - + Azure: $0.03 USD per session """ if sessions is None or sessions == 0: return 0.0 - + # Check if model-specific pricing is available if model_info and "code_interpreter_cost_per_session" in model_info: return sessions * model_info["code_interpreter_cost_per_session"] - + # Azure pricing for code interpreter if provider == "azure": from litellm.constants import AZURE_CODE_INTERPRETER_COST_PER_SESSION + return sessions * AZURE_CODE_INTERPRETER_COST_PER_SESSION - + # OpenAI doesn't charge separately for code interpreter yet return 0.0 From a569e4d7114e1f563e8a8f9dd30848677b4e8647 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 10:23:46 -0700 Subject: [PATCH 35/76] test: update test to not stop on first error (allows us to see all issues) --- .github/workflows/test-linting.yml | 2 +- .github/workflows/test-litellm.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index ffca305a0d..6162c53f95 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -44,7 +44,7 @@ jobs: - name: Run MyPy type checking run: | cd litellm - poetry run mypy . --ignore-missing-imports + poetry run mypy . --config-file=mypy.ini cd .. - name: Check for circular imports diff --git a/.github/workflows/test-litellm.yml b/.github/workflows/test-litellm.yml index 0d3a9f2b5d..b6b77ad326 100644 --- a/.github/workflows/test-litellm.yml +++ b/.github/workflows/test-litellm.yml @@ -40,4 +40,4 @@ jobs: cd .. - name: Run tests run: | - poetry run pytest tests/test_litellm -x -vv -n 4 + poetry run pytest tests/test_litellm -vv -n 4 From abe250dcd8184482302f5f97af039787b4d9e243 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:32:47 -0700 Subject: [PATCH 36/76] test_azure_embedding_max_retries_0 --- tests/llm_translation/test_azure_openai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/llm_translation/test_azure_openai.py b/tests/llm_translation/test_azure_openai.py index 33947e25a6..77b8acfcbe 100644 --- a/tests/llm_translation/test_azure_openai.py +++ b/tests/llm_translation/test_azure_openai.py @@ -565,7 +565,7 @@ async def test_azure_embedding_max_retries_0( from litellm import aembedding, embedding args = { - "model": "text-embedding-ada-002", + "model": "azure/text-embedding-ada-002", "input": "Hello world", "max_retries": max_retries, } From ba3e20060ae2f0b6c91d991a2b48d6b94bc7fb17 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:34:36 -0700 Subject: [PATCH 37/76] llm_translation_testing fix --- .circleci/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 9761167ec1..a41191aaeb 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -810,7 +810,7 @@ jobs: command: | pwd ls - python -m pytest -vv tests/llm_translation --cov=litellm --cov-report=xml -x -v --junitxml=test-results/junit.xml --durations=5 -n 4 + python -m pytest -vv tests/llm_translation --cov=litellm --cov-report=xml -v --junitxml=test-results/junit.xml --durations=5 -n 4 no_output_timeout: 120m - run: name: Rename the coverage files From bcf26b6fce2213b919022f3a50e8dd3cac12782f Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:36:43 -0700 Subject: [PATCH 38/76] CALLBACK_CLASS_STR_TO_CLASS_TYPE --- litellm/litellm_core_utils/custom_logger_registry.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 75dc0918ab..87fc2173b1 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -95,7 +95,6 @@ class CustomLoggerRegistry: "bitbucket": BitBucketPromptManager, "cloudzero": CloudZeroLogger, "posthog": PostHogLogger, - "bitbucket": BitBucketPromptManager, } try: From 3baa3aff1b86f9f7e4142936e0ba46d0fbcaa0ea Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:38:35 -0700 Subject: [PATCH 39/76] test fix --- .../test_configs/test_custom_logger.yaml | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/proxy_unit_tests/test_configs/test_custom_logger.yaml b/tests/proxy_unit_tests/test_configs/test_custom_logger.yaml index 2ad500b36f..6ff1b08635 100644 --- a/tests/proxy_unit_tests/test_configs/test_custom_logger.yaml +++ b/tests/proxy_unit_tests/test_configs/test_custom_logger.yaml @@ -1,20 +1,16 @@ model_list: - model_name: Azure OpenAI GPT-4 Canada litellm_params: - model: azure/chatgpt-v-3 - api_base: os.environ/AZURE_API_BASE - api_key: os.environ/AZURE_API_KEY - api_version: "2023-07-01-preview" + model: gpt-4.1-nano + api_key: os.environ/OPENAI_API_KEY model_info: mode: chat input_cost_per_token: 0.0002 id: gm - model_name: azure-embedding-model litellm_params: - model: azure/azure-embedding-model - api_base: os.environ/AZURE_API_BASE - api_key: os.environ/AZURE_API_KEY - api_version: "2023-07-01-preview" + model: text-embedding-ada-002 + api_key: os.environ/OPENAI_API_KEY model_info: mode: embedding input_cost_per_token: 0.002 From 5d6889dc928f7038fc27dcaac038b85a2b53b26e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 10:40:06 -0700 Subject: [PATCH 40/76] test: update tests --- tests/mcp_tests/test_mcp_server.py | 42 +++++++++++++++--------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index e55fdc8516..d0489bac58 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -98,9 +98,9 @@ async def test_mcp_server_manager_https_server(): assert tools[0].name == f"{expected_prefix}-gmail_send_email" # Manually set up the tool mapping for the call_tool test - mcp_server_manager.tool_name_to_mcp_server_name_mapping[ - "gmail_send_email" - ] = expected_prefix + mcp_server_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = ( + expected_prefix + ) mcp_server_manager.tool_name_to_mcp_server_name_mapping[ f"{expected_prefix}-gmail_send_email" ] = expected_prefix @@ -260,9 +260,9 @@ async def test_mcp_http_transport_call_tool_mock(): ) # Manually set up tool mapping (normally done by list_tools) - test_manager.tool_name_to_mcp_server_name_mapping[ - "gmail_send_email" - ] = "test_http_server" + test_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = ( + "test_http_server" + ) # Call the tool result = await test_manager.call_tool( @@ -326,9 +326,9 @@ async def test_mcp_http_transport_call_tool_error_mock(): ) # Manually set up tool mapping - test_manager.tool_name_to_mcp_server_name_mapping[ - "gmail_send_email" - ] = "test_http_server" + test_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = ( + "test_http_server" + ) # Call the tool with invalid data result = await test_manager.call_tool( @@ -792,18 +792,18 @@ async def test_get_tools_from_mcp_servers(): assert len(result) == 1, "Should only return tools from server1" assert result[0].name == "tool1", "Should return tool from server1" - # Test Case 2: Without specific MCP servers - # Create a different mock manager for the second test case - mock_manager_2 = AsyncMock() - mock_manager_2.get_allowed_mcp_servers = AsyncMock( - return_value=["server1_id", "server2_id"] - ) - mock_manager_2.get_mcp_server_by_id = mock_get_server_by_id - mock_manager_2._get_tools_from_server = AsyncMock( - side_effect=lambda server, mcp_auth_header=None: [mock_tool_1] - if server.server_id == "server1_id" - else [mock_tool_2] - ) + # Test Case 2: Without specific MCP servers + # Create a different mock manager for the second test case + mock_manager_2 = AsyncMock() + mock_manager_2.get_allowed_mcp_servers = AsyncMock( + return_value=["server1_id", "server2_id"] + ) + mock_manager_2.get_mcp_server_by_id = mock_get_server_by_id + mock_manager_2._get_tools_from_server = AsyncMock( + side_effect=lambda server, mcp_auth_header=None, extra_headers=None: ( + [mock_tool_1] if server.server_id == "server1_id" else [mock_tool_2] + ) + ) with patch( "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", From 936e3466acbbdf55ea8c8321fc3375e0bee4e3d9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:42:51 -0700 Subject: [PATCH 41/76] litellm_router_testing --- .circleci/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index a41191aaeb..69057d1a3f 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -464,7 +464,7 @@ jobs: command: | pwd ls - python -m pytest tests/local_testing --cov=litellm --cov-report=xml -vv -k "router" -x -v --junitxml=test-results/junit.xml --durations=5 + python -m pytest tests/local_testing --cov=litellm --cov-report=xml -vv -k "router" -v --junitxml=test-results/junit.xml --durations=5 no_output_timeout: 120m - run: name: Rename the coverage files From e31972f5db82bbe09aaf657722d5b7d993e8bb14 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:45:36 -0700 Subject: [PATCH 42/76] litellm_mapped_tests --- .circleci/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 69057d1a3f..41c3ba1133 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1042,7 +1042,7 @@ jobs: command: | pwd ls - python -m pytest -vv tests/test_litellm --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit-litellm.xml --durations=10 -n 8 + python -m pytest -vv tests/test_litellm --cov=litellm --cov-report=xml -s -v --junitxml=test-results/junit-litellm.xml --durations=10 -n 8 no_output_timeout: 120m - run: name: Rename the coverage files From 82eeca78708543efef008a1dd4989a36a89e981e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 10:49:21 -0700 Subject: [PATCH 43/76] test: update test --- .github/workflows/test-linting.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 6162c53f95..7e169d6a27 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -44,7 +44,7 @@ jobs: - name: Run MyPy type checking run: | cd litellm - poetry run mypy . --config-file=mypy.ini + poetry run mypy . --ignore-missing-imports --disable-error-code=var-annotated cd .. - name: Check for circular imports From 8d943172600b7f090a4cd2ad89f25fc7948c3933 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:53:44 -0700 Subject: [PATCH 44/76] test fix BitBucketPromptManager redef --- litellm/litellm_core_utils/custom_logger_registry.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 87fc2173b1..4957c97e5b 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -40,7 +40,6 @@ try: from litellm_enterprise.integrations.prometheus import PrometheusLogger except Exception: PrometheusLogger = None -from litellm.integrations.bitbucket import BitBucketPromptManager from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger from litellm.integrations.dotprompt import DotpromptManager from litellm.integrations.s3_v2 import S3Logger From 68da131728e1f947fb99127b34ba75c2b7c14d9c Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:55:46 -0700 Subject: [PATCH 45/76] TestAzureEmbedding --- tests/llm_translation/test_azure_openai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/llm_translation/test_azure_openai.py b/tests/llm_translation/test_azure_openai.py index 77b8acfcbe..131e23b1c5 100644 --- a/tests/llm_translation/test_azure_openai.py +++ b/tests/llm_translation/test_azure_openai.py @@ -204,7 +204,7 @@ def test_process_azure_endpoint_url(api_base, model, expected_endpoint): class TestAzureEmbedding(BaseLLMEmbeddingTest): def get_base_embedding_call_args(self) -> dict: return { - "model": "azure/azure-embedding-model", + "model": "azure/text-embedding-ada-002", "api_key": os.getenv("AZURE_API_KEY"), "api_base": os.getenv("AZURE_API_BASE"), } From 97f0f3fe32eeea390ccc757989e4dad039e79a11 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:57:22 -0700 Subject: [PATCH 46/76] test_gemini_finish_reason --- tests/llm_translation/test_gemini.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index aa349253e5..c34e59a73e 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -452,7 +452,7 @@ def test_gemini_finish_reason(): litellm._turn_on_debug() response = completion( - model="gemini/gemini-1.5-pro", + model="gemini/gemini-2.5-flash-lite", messages=[{"role": "user", "content": "give me 3 random words"}], max_tokens=2, ) From 09556d0a44ac1c6464d465cdbc5a48e4e73ec7ac Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 10:58:14 -0700 Subject: [PATCH 47/76] test: add debugging step to linting --- .github/workflows/test-linting.yml | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 7e169d6a27..e2c6500b98 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -11,6 +11,9 @@ jobs: steps: - uses: actions/checkout@v4 + with: + fetch-depth: 0 + clean: true - name: Set up Python uses: actions/setup-python@v4 @@ -20,6 +23,11 @@ jobs: - name: Install Poetry uses: snok/install-poetry@v1 + - name: Clean Python cache + run: | + find . -type d -name "__pycache__" -exec rm -rf {} + || true + find . -name "*.pyc" -delete || true + - name: Install dependencies run: | poetry install --with dev @@ -31,6 +39,15 @@ jobs: poetry run black . cd .. + - name: Debug - Check file state + run: | + echo "Current branch:" + git branch --show-current + echo "Last 3 commits:" + git log --oneline -3 + echo "File content around line 43:" + head -50 litellm/litellm_core_utils/custom_logger_registry.py | tail -10 + - name: Run Ruff linting run: | cd litellm From 55572ce3cd511fbbeb824d3952e05c6ed7e3bb12 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 10:59:25 -0700 Subject: [PATCH 48/76] test_add_model_run_health --- tests/test_models.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index b1a5b84b3b..67e77dcaaf 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -307,10 +307,8 @@ async def add_model_for_health_checking(session, model_id="123"): data = { "model_name": f"azure-model-health-check-{model_id}", "litellm_params": { - "model": "azure/chatgpt-v-3", - "api_key": os.getenv("AZURE_API_KEY"), - "api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/", - "api_version": "2023-05-15", + "model": "gpt-4.1-nano", + "api_key": os.getenv("OPENAI_API_KEY"), }, "model_info": {"id": model_id}, } @@ -432,7 +430,7 @@ async def test_add_model_run_health(): assert _health_info["healthy_count"] == 1 assert ( - _healthy_endpooint["model"] == "azure/chatgpt-v-3" + _healthy_endpooint["model"] == "gpt-4.1-nano" ) # this is the model that got added # assert httpx client is is unchanges From 5cdeb63cdd6fe39ef56d0d643ed85d3d4a7fc10c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 11:09:25 -0700 Subject: [PATCH 49/76] test: fix tests --- .github/workflows/test-litellm.yml | 2 +- .../mcp_server/auth/test_user_api_key_auth_mcp.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-litellm.yml b/.github/workflows/test-litellm.yml index b6b77ad326..b7f4a25d59 100644 --- a/.github/workflows/test-litellm.yml +++ b/.github/workflows/test-litellm.yml @@ -40,4 +40,4 @@ jobs: cd .. - name: Run tests run: | - poetry run pytest tests/test_litellm -vv -n 4 + poetry run pytest tests/test_litellm --tb=short -vv --maxfail=10 -n 4 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index fb15dbc14e..afafa510a4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -364,6 +364,7 @@ class TestMCPRequestHandler: mcp_auth_header, mcp_servers, mcp_server_auth_headers, + oauth2_headers, ) = await MCPRequestHandler.process_mcp_request(scope) # Assert the results @@ -717,6 +718,7 @@ class TestMCPCustomHeaderName: mcp_auth_header, mcp_servers, mcp_server_auth_headers, + oauth2_headers, ) = await MCPRequestHandler.process_mcp_request(scope) # Assert the results @@ -887,6 +889,7 @@ class TestMCPAccessGroupsE2E: mcp_auth_header, mcp_servers, mcp_server_auth_headers, + oauth2_headers, ) = await MCPRequestHandler.process_mcp_request(scope) # Assert the results @@ -936,6 +939,7 @@ class TestMCPAccessGroupsE2E: mcp_auth_header, mcp_servers, mcp_server_auth_headers, + oauth2_headers, ) = await MCPRequestHandler.process_mcp_request(scope) # Assert the results @@ -964,6 +968,7 @@ def test_mcp_path_based_server_segregation(monkeypatch): mcp_auth_header, mcp_servers, mcp_server_auth_headers, + oauth2_headers, ) = get_auth_context() # Capture the MCP servers for testing From 5c6a863e40d6ebd8d1920b138bfd0ced4ac2ca50 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 11:10:44 -0700 Subject: [PATCH 50/76] docs: cleanup --- docs/my-website/sidebars.js | 1 - 1 file changed, 1 deletion(-) diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index a131e5c34e..47289ca1d4 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -562,7 +562,6 @@ const sidebars = { items: [ "set_keys", "completion/token_usage", - "sdk/headers", "sdk_custom_pricing", "embedding/async_embedding", "embedding/moderation", From bd77bdabe634fd9def69101c39bcb71b155d15a5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 11:27:34 -0700 Subject: [PATCH 51/76] test_call_one_endpoint --- tests/local_testing/test_router.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index e7e61d378c..fbe0a2fb8a 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -544,7 +544,7 @@ def test_call_one_endpoint(): { "model_name": "gpt-3.5-turbo", # openai model name "litellm_params": { # params for litellm completion/embedding call - "model": "azure/gpt-4o-new-test", + "model": "azure/gpt-4.1-nano", "api_key": old_api_key, "api_version": os.getenv("AZURE_API_VERSION"), "api_base": os.getenv("AZURE_API_BASE"), @@ -555,7 +555,7 @@ def test_call_one_endpoint(): { "model_name": "text-embedding-ada-002", "litellm_params": { - "model": "azure/azure-embedding-model", + "model": "azure/text-embedding-ada-002", "api_key": os.environ["AZURE_API_KEY"], "api_base": os.environ["AZURE_API_BASE"], }, @@ -574,7 +574,7 @@ def test_call_one_endpoint(): async def call_azure_completion(): response = await router.acompletion( - model="azure/gpt-4o-new-test", + model="azure/gpt-4.1-nano", messages=[{"role": "user", "content": "hello this request will pass"}], specific_deployment=True, ) @@ -582,7 +582,7 @@ def test_call_one_endpoint(): async def call_azure_embedding(): response = await router.aembedding( - model="azure/azure-embedding-model", + model="azure/text-embedding-ada-002", input=["good morning from litellm"], specific_deployment=True, ) From 66b86e46e3b768813f5c4ca7214cc3aa3ff92bc8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 11:27:56 -0700 Subject: [PATCH 52/76] test_azure_embedding_on_router --- tests/local_testing/test_router.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index fbe0a2fb8a..7ef426778c 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -1274,7 +1274,7 @@ def test_azure_embedding_on_router(): { "model_name": "text-embedding-ada-002", "litellm_params": { - "model": "azure/azure-embedding-model", + "model": "azure/text-embedding-ada-002", "api_key": os.environ["AZURE_API_KEY"], "api_base": os.environ["AZURE_API_BASE"], }, From 401ec6f2cf1c8cffaa9ac1a1285536a3fbe2171c Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 11:28:20 -0700 Subject: [PATCH 53/76] test_cooldown_badrequest_error --- tests/local_testing/test_router_cooldown_handlers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/local_testing/test_router_cooldown_handlers.py b/tests/local_testing/test_router_cooldown_handlers.py index d3fd78063b..7d3fdda22c 100644 --- a/tests/local_testing/test_router_cooldown_handlers.py +++ b/tests/local_testing/test_router_cooldown_handlers.py @@ -44,7 +44,7 @@ async def test_cooldown_badrequest_error(): { "model_name": "gpt-3.5-turbo", "litellm_params": { - "model": "azure/chatgpt-v-3", + "model": "azure/gpt-4.1-nano", "api_key": os.getenv("AZURE_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), "api_base": os.getenv("AZURE_API_BASE"), From 172dc07e52d695559fd2d89791f8eff6167eeba3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 11:29:05 -0700 Subject: [PATCH 54/76] test_cooldown_same_model_name --- tests/local_testing/test_acooldowns_router.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/tests/local_testing/test_acooldowns_router.py b/tests/local_testing/test_acooldowns_router.py index 8427fd2be8..0c385133e3 100644 --- a/tests/local_testing/test_acooldowns_router.py +++ b/tests/local_testing/test_acooldowns_router.py @@ -143,20 +143,16 @@ async def test_cooldown_same_model_name(sync_mode): { "model_name": "gpt-3.5-turbo", "litellm_params": { - "model": "azure/chatgpt-v-3", + "model": "gpt-4.1-nano", "api_key": "bad-key", - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), "tpm": 90, }, }, { "model_name": "gpt-3.5-turbo", "litellm_params": { - "model": "azure/chatgpt-v-3", - "api_key": os.getenv("AZURE_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "model": "gpt-4.1-nano", + "api_key": os.getenv("OPENAI_API_KEY"), "tpm": 1, }, }, From d6165454833eda3c066c08e64308d055e5cdc3d0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 11:30:09 -0700 Subject: [PATCH 55/76] test_reading_key_from_model_list --- tests/local_testing/test_router.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index 7ef426778c..4be7e78f97 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -473,16 +473,12 @@ def test_reading_key_from_model_list(): try: print("testing if router raises an exception") - old_api_key = os.environ["AZURE_API_KEY"] - os.environ.pop("AZURE_API_KEY", None) model_list = [ { "model_name": "gpt-3.5-turbo", # openai model name "litellm_params": { # params for litellm completion/embedding call - "model": "azure/chatgpt-v-3", - "api_key": old_api_key, - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "model": "gpt-4.1-nano", + "api_key": os.getenv("OPENAI_API_KEY"), }, "tpm": 240000, "rpm": 1800, @@ -521,10 +517,8 @@ def test_reading_key_from_model_list(): print("\n completed_response", completed_response) assert len(completed_response) > 0 print("\n Passed Streaming") - os.environ["AZURE_API_KEY"] = old_api_key router.reset() except Exception as e: - os.environ["AZURE_API_KEY"] = old_api_key print(f"FAILED TEST") pytest.fail(f"Got unexpected exception on router! - {e}") From 95590ba81cbe133edb9088ccda018795330097fe Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 11:35:33 -0700 Subject: [PATCH 56/76] test_router_retries --- tests/local_testing/test_router.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index 4be7e78f97..3a07238eb3 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -334,10 +334,8 @@ async def test_router_retries(sync_mode): { "model_name": "gpt-3.5-turbo", "litellm_params": { - "model": "azure/gpt-4o-new-test", - "api_key": os.getenv("AZURE_API_KEY"), - "api_base": os.getenv("AZURE_API_BASE"), - "api_version": os.getenv("AZURE_API_VERSION"), + "model": "gpt-4.1-nano", + "api_key": os.getenv("OPENAI_API_KEY"), }, }, ] From 5e2479d7a3aaf84d555a36832111978698d16a8b Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 11:38:08 -0700 Subject: [PATCH 57/76] test_router_with_caching --- tests/local_testing/test_prometheus_service.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/local_testing/test_prometheus_service.py b/tests/local_testing/test_prometheus_service.py index be620a0c7c..3bcae6336d 100644 --- a/tests/local_testing/test_prometheus_service.py +++ b/tests/local_testing/test_prometheus_service.py @@ -92,24 +92,22 @@ async def test_router_with_caching(): """ try: - def get_azure_params(deployment_name: str): + def get_openai_params(): params = { - "model": f"azure/{deployment_name}", - "api_key": os.environ["AZURE_API_KEY"], - "api_version": os.environ["AZURE_API_VERSION"], - "api_base": os.environ["AZURE_API_BASE"], + "model": "gpt-4.1-nano", + "api_key": os.environ["OPENAI_API_KEY"], } return params model_list = [ { "model_name": "azure/gpt-4", - "litellm_params": get_azure_params("gpt-4o-new-test"), + "litellm_params": get_openai_params(), "tpm": 100, }, { "model_name": "azure/gpt-4", - "litellm_params": get_azure_params("gpt-4o-new-test"), + "litellm_params": get_openai_params(), "tpm": 1000, }, ] From c99c9ad87df8ad9dd31cdfb81ac6077804bda657 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 11:41:08 -0700 Subject: [PATCH 58/76] test_service_unavailable_fallbacks --- tests/local_testing/test_router_fallbacks.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index 2159c77f0d..5b4e6ab580 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -989,10 +989,8 @@ async def test_service_unavailable_fallbacks(sync_mode): { "model_name": "gpt-3.5-turbo-0125-preview", "litellm_params": { - "model": "azure/chatgpt-v-3", - "api_key": os.getenv("AZURE_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "model": "gpt-4.1-nano", + "api_key": os.getenv("OPENAI_API_KEY"), }, }, ], @@ -1010,7 +1008,7 @@ async def test_service_unavailable_fallbacks(sync_mode): messages=[{"role": "user", "content": "Hey, how's it going?"}], ) - assert response.model == "gpt-3.5-turbo-0125" + assert "gpt-4.1-nano" in response.model @pytest.mark.parametrize("sync_mode", [True, False]) From 9097aaa65c1f3e887d95d1260ce3c64ebf043ad5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 11:42:34 -0700 Subject: [PATCH 59/76] test_usage_based_routing --- tests/local_testing/test_router_get_deployments.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/local_testing/test_router_get_deployments.py b/tests/local_testing/test_router_get_deployments.py index 4d98097341..8d064be9a8 100644 --- a/tests/local_testing/test_router_get_deployments.py +++ b/tests/local_testing/test_router_get_deployments.py @@ -397,7 +397,7 @@ def test_usage_based_routing(): "model": f"azure/{deployment_name}", "api_key": os.environ["AZURE_API_KEY"], "api_version": os.environ["AZURE_API_VERSION"], - "api_base": os.environ["AZURE_API_BASE"], + "api_base": "https://fake-api.openai.com/v1", } return params From 284a8549a1bed0a14519695752542646b7eb5b5c Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 11:43:20 -0700 Subject: [PATCH 60/76] test_chat_completion --- tests/proxy_unit_tests/test_proxy_custom_logger.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/proxy_unit_tests/test_proxy_custom_logger.py b/tests/proxy_unit_tests/test_proxy_custom_logger.py index e013bbc216..0fbe5a31c3 100644 --- a/tests/proxy_unit_tests/test_proxy_custom_logger.py +++ b/tests/proxy_unit_tests/test_proxy_custom_logger.py @@ -164,7 +164,7 @@ def test_chat_completion(client): my_custom_logger.async_success == True ) # checks if the status of async_success is True, only the async_log_success_event can set this to true assert ( - my_custom_logger.async_completion_kwargs["model"] == "chatgpt-v-3" + my_custom_logger.async_completion_kwargs["model"] == "gpt-4.1-nano" ) # checks if kwargs passed to async_log_success_event are correct print( "\n\n Custom Logger Async Completion args", From ddf918cb42d00992d5ce02a5523ae6bb52aa4e09 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 11:55:41 -0700 Subject: [PATCH 61/76] use_callback_in_llm_call --- .../test_unit_tests_init_callbacks.py | 25 ++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py index a0ba4f34d9..0226fd66fd 100644 --- a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py +++ b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py @@ -102,6 +102,26 @@ async def use_callback_in_llm_call( elif callback == "openmeter": # it's currently handled in jank way, TODO: fix openmete and then actually run it's test return + elif callback == "bitbucket": + # Set up mock bitbucket configuration required for initialization + litellm.global_bitbucket_config = { + "workspace": "test-workspace", + "repository": "test-repo", + "access_token": "test-token", + "branch": "main" + } + # Mock BitBucket HTTP calls to prevent actual API requests + import httpx + from unittest.mock import MagicMock + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"values": []} + mock_response.text = "" + + patch.object( + litellm.module_level_client, "get", return_value=mock_response + ).start() elif callback == "prometheus": # pytest teardown - clear existing prometheus collectors collectors = list(REGISTRY._collector_to_names.keys()) @@ -175,7 +195,10 @@ async def use_callback_in_llm_call( if callback == "argilla": patch.stopall() - if callback == "argilla": + if callback == "bitbucket": + # Clean up bitbucket configuration and patches + if hasattr(litellm, 'global_bitbucket_config'): + delattr(litellm, 'global_bitbucket_config') patch.stopall() From 3e793d71222cffd37b07ed742cc0d4cb63b65617 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 27 Sep 2025 11:56:45 -0700 Subject: [PATCH 62/76] test - fix mypy linting on ci/cd (#14981) * mypy_linting make 1 job for it * mypy_linting * fix fastuuid * fix mypy * mypy fix --- .circleci/config.yml | 41 ++++++++++++++++++++++++++++++++++------- 1 file changed, 34 insertions(+), 7 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 41c3ba1133..14e59a63e1 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -51,6 +51,33 @@ jobs: command: | python -m pytest tests/windows_tests/test_litellm_on_windows.py -v + mypy_linting: + docker: + - image: cimg/python:3.12 + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + working_directory: ~/project + resource_class: medium + + steps: + - checkout + - setup_google_dns + - run: + name: Install Dependencies + command: | + python -m pip install --upgrade pip + python -m pip install -r requirements.txt + pip uninstall fastuuid -y + pip install "mypy==1.18.2" + - run: + name: MyPy Type Checking + command: | + cd litellm + # Use the same approach as GitHub Actions, explicitly exclude fastuuid to avoid segfaults + python -m mypy . --ignore-missing-imports + cd .. + no_output_timeout: 10m local_testing: docker: - image: cimg/python:3.12 @@ -140,13 +167,6 @@ jobs: python -m pip install black python -m black . cd .. - - run: - name: Linting Testing - command: | - cd litellm - # Use the same simple approach that works in GitHub Actions - python -m mypy . --ignore-missing-imports --show-traceback - cd .. # Run pytest and generate JUnit XML report - run: @@ -3071,6 +3091,12 @@ workflows: only: - main - /litellm_.*/ + - mypy_linting: + filters: + branches: + only: + - main + - /litellm_.*/ - local_testing: filters: branches: @@ -3317,6 +3343,7 @@ workflows: - main - publish_to_pypi: requires: + - mypy_linting - local_testing - build_and_test - e2e_openai_endpoints From 200dd0ef397faff7ac787fb76436d570a038504a Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:04:30 -0700 Subject: [PATCH 63/76] local_testing fix --- .circleci/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 14e59a63e1..d41c5aea20 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -174,7 +174,7 @@ jobs: command: | pwd ls - python -m pytest -vv tests/local_testing --cov=litellm --cov-report=xml -x --junitxml=test-results/junit.xml --durations=5 -k "not test_python_38.py and not test_basic_python_version.py and not router and not assistants and not langfuse and not caching and not cache" -n 4 + python -m pytest -vv tests/local_testing --cov=litellm --cov-report=xml --junitxml=test-results/junit.xml --durations=5 -k "not test_python_38.py and not test_basic_python_version.py and not router and not assistants and not langfuse and not caching and not cache" -n 4 no_output_timeout: 120m - run: name: Rename the coverage files From bbf5761b4987afe5d193c91b77bc291799c00660 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:06:26 -0700 Subject: [PATCH 64/76] tets health check --- .../test_health_check.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename tests/{local_testing => litellm_utils_tests}/test_health_check.py (99%) diff --git a/tests/local_testing/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py similarity index 99% rename from tests/local_testing/test_health_check.py rename to tests/litellm_utils_tests/test_health_check.py index f14bb617da..6cff1f08d9 100644 --- a/tests/local_testing/test_health_check.py +++ b/tests/litellm_utils_tests/test_health_check.py @@ -20,7 +20,7 @@ import litellm async def test_azure_health_check(): response = await litellm.ahealth_check( model_params={ - "model": "azure/chatgpt-v-3", + "model": "azure/gpt-4.1-nano", "messages": [{"role": "user", "content": "Hey, how's it going?"}], "api_key": os.getenv("AZURE_API_KEY"), "api_base": os.getenv("AZURE_API_BASE"), @@ -51,7 +51,7 @@ async def test_text_completion_health_check(): async def test_azure_embedding_health_check(): response = await litellm.ahealth_check( model_params={ - "model": "azure/azure-embedding-model", + "model": "azure/text-embedding-ada-002", "api_key": os.getenv("AZURE_API_KEY"), "api_base": os.getenv("AZURE_API_BASE"), "api_version": os.getenv("AZURE_API_VERSION"), From 499c988d847d74f406f9396271e13f69fbadd713 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:09:46 -0700 Subject: [PATCH 65/76] mypy lint fixes --- litellm/proxy/auth/oauth2_proxy_hook.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/oauth2_proxy_hook.py b/litellm/proxy/auth/oauth2_proxy_hook.py index a1db5d842c..a900372463 100644 --- a/litellm/proxy/auth/oauth2_proxy_hook.py +++ b/litellm/proxy/auth/oauth2_proxy_hook.py @@ -15,8 +15,8 @@ async def handle_oauth2_proxy_request(request: Request) -> UserAPIKeyAuth: verbose_proxy_logger.debug("Handling oauth2 proxy request") # Define the OAuth2 config mappings oauth2_config_mappings: Dict[str, str] = general_settings.get( - "oauth2_config_mappings", None - ) + "oauth2_config_mappings", {} + ) or {} verbose_proxy_logger.debug(f"Oauth2 config mappings: {oauth2_config_mappings}") if not oauth2_config_mappings: From 4baa48962e73df499f10126642638f463a28f380 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:10:28 -0700 Subject: [PATCH 66/76] _group_keys_by_hash_tag fix mypy --- litellm/proxy/hooks/parallel_request_limiter_v3.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 03ddfa1f7a..3e4cbec29f 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -297,7 +297,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Group keys by their Redis hash tag to ensure cluster compatibility. Keys with the same hash tag will be processed together. """ - groups = {} + groups: Dict[str, List[str]] = {} for key in keys: # Extract hash tag from key like "{api_key:sk-123}:requests" if "{" in key and "}" in key: @@ -378,7 +378,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for descriptor in descriptors: descriptor_key = descriptor["key"] descriptor_value = descriptor["value"] - rate_limit = descriptor.get("rate_limit", {}) or {} + rate_limit: Optional[RateLimitDescriptorRateLimitObject] = descriptor.get("rate_limit", {}) or {} requests_limit = rate_limit.get("requests_per_unit") tokens_limit = rate_limit.get("tokens_per_unit") max_parallel_requests_limit = rate_limit.get("max_parallel_requests") From 609f27b4ad347f880c9df60dfcedc434c0a7aef6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:11:00 -0700 Subject: [PATCH 67/76] calculate_mcp_tool_call_cost --- litellm/proxy/_experimental/mcp_server/cost_calculator.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/cost_calculator.py b/litellm/proxy/_experimental/mcp_server/cost_calculator.py index eea10924a1..3fe33b6dd5 100644 --- a/litellm/proxy/_experimental/mcp_server/cost_calculator.py +++ b/litellm/proxy/_experimental/mcp_server/cost_calculator.py @@ -38,7 +38,8 @@ class MCPCostCalculator: # Unpack the mcp_tool_call_metadata ######################################################### mcp_tool_call_metadata: StandardLoggingMCPToolCall = cast(StandardLoggingMCPToolCall, litellm_logging_obj.model_call_details.get("mcp_tool_call_metadata", {})) or {} - mcp_server_cost_info: MCPServerCostInfo = mcp_tool_call_metadata.get("mcp_server_cost_info", {}) or {} + mcp_server_cost_info_raw = mcp_tool_call_metadata.get("mcp_server_cost_info", {}) or {} + mcp_server_cost_info: MCPServerCostInfo = cast(MCPServerCostInfo, mcp_server_cost_info_raw) ######################################################### # User defined cost per query ######################################################### From a1b7ce442615d591a3f0bf6cea74c9fe973ccd9f Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:12:30 -0700 Subject: [PATCH 68/76] make_bedrock_api_request --- .../guardrail_hooks/bedrock_guardrails.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 462582b1f3..eec3c1a215 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -363,7 +363,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prepared_request.headers, ) - response = await self.async_handler.post( + httpx_response = await self.async_handler.post( url=prepared_request.url, data=prepared_request.body, # type: ignore headers=prepared_request.headers, # type: ignore @@ -373,19 +373,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, - guardrail_json_response=response.json(), + guardrail_json_response=httpx_response.json(), request_data=request_data or {}, guardrail_status=self._get_bedrock_guardrail_response_status( - response=response + response=httpx_response ), start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), ) ######################################################### - if response.status_code == 200: + if httpx_response.status_code == 200: # check if the response was flagged - _json_response = response.json() + _json_response = httpx_response.json() redacted_response = _redact_pii_matches(_json_response) verbose_proxy_logger.debug("Bedrock AI response : %s", redacted_response) bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response) @@ -398,8 +398,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): else: verbose_proxy_logger.error( "Bedrock AI: error in response. Status code: %s, response: %s", - response.status_code, - response.text, + httpx_response.status_code, + httpx_response.text, ) return bedrock_guardrail_response From 1bb7ad4f69d3d8dba5b8d310e5e9649b6a2c0b29 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:13:51 -0700 Subject: [PATCH 69/76] litellm_utils_testing --- .circleci/config.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.circleci/config.yml b/.circleci/config.yml index d41c5aea20..459dd70e8d 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1200,6 +1200,7 @@ jobs: pip install "pytest-cov==5.0.0" pip install "google-generativeai==0.3.2" pip install "google-cloud-aiplatform==1.43.0" + pip install pytest-mock # Run pytest and generate JUnit XML report - run: name: Run tests From f7213bacfe242637aaab9119f4857bfbf4d1c248 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:18:17 -0700 Subject: [PATCH 70/76] test_async_custom_handler_embedding_optional_param --- tests/local_testing/test_custom_logger.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/local_testing/test_custom_logger.py b/tests/local_testing/test_custom_logger.py index ba9973e11d..ee212658c8 100644 --- a/tests/local_testing/test_custom_logger.py +++ b/tests/local_testing/test_custom_logger.py @@ -392,7 +392,7 @@ async def test_async_custom_handler_embedding_optional_param(): customHandler_optional_params = MyCustomHandler() litellm.callbacks = [customHandler_optional_params] response = await litellm.aembedding( - model="azure/azure-embedding-model", input=["hello world"], user="John" + model="text-embedding-ada-002", input=["hello world"], user="John" ) await asyncio.sleep(1) # success callback is async assert customHandler_optional_params.user == "John" From 82978f4c5c8b2f1b14ab245c5689c22ca4902db2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:21:32 -0700 Subject: [PATCH 71/76] fix mypy --- .circleci/config.yml | 2 +- litellm/mypy.ini | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 459dd70e8d..0cedeb7168 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -75,7 +75,7 @@ jobs: command: | cd litellm # Use the same approach as GitHub Actions, explicitly exclude fastuuid to avoid segfaults - python -m mypy . --ignore-missing-imports + python -m mypy . cd .. no_output_timeout: 10m local_testing: diff --git a/litellm/mypy.ini b/litellm/mypy.ini index d4e5e7f6bd..4702b59112 100644 --- a/litellm/mypy.ini +++ b/litellm/mypy.ini @@ -5,7 +5,8 @@ mypy_path = litellm/stubs namespace_packages = True disable_error_code = valid-type, - annotation-unchecked + annotation-unchecked, + import-untyped [mypy-google.*] ignore_missing_imports = True From 9ea8c72ebccc8e0ab071d329ea24a6759a86ba31 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:26:06 -0700 Subject: [PATCH 72/76] mypy lint fix --- .../llm_cost_calc/tool_call_cost_tracking.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 52175989cb..b611366177 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -123,7 +123,7 @@ class StandardBuiltInToolCostTracking: model_info = StandardBuiltInToolCostTracking._safe_get_model_info( model=model, custom_llm_provider=custom_llm_provider ) - file_search_raw = standard_built_in_tools_params.get("file_search", {}) + file_search_raw: Any = standard_built_in_tools_params.get("file_search", {}) file_search_usage: Optional[FileSearchTool] = ( FileSearchTool(**file_search_raw) if file_search_raw else None ) @@ -427,7 +427,7 @@ class StandardBuiltInToolCostTracking: if model_info is None: return 0.0 - search_context_raw = model_info.get("search_context_cost_per_query", {}) + search_context_raw: Any = model_info.get("search_context_cost_per_query", {}) search_context_pricing: SearchContextCostPerQuery = ( SearchContextCostPerQuery(**search_context_raw) if search_context_raw @@ -454,7 +454,7 @@ class StandardBuiltInToolCostTracking: """ if model_info is None: return 0.0 - search_context_raw = model_info.get("search_context_cost_per_query", {}) or {} + search_context_raw: Any = model_info.get("search_context_cost_per_query", {}) or {} search_context_pricing: SearchContextCostPerQuery = ( SearchContextCostPerQuery(**search_context_raw) if search_context_raw From 4b06d6d6d6650a34f3b09737372baef04e1cbabd Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:26:40 -0700 Subject: [PATCH 73/76] test_provider_budgets_e2e_test --- tests/local_testing/test_router_budget_limiter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/local_testing/test_router_budget_limiter.py b/tests/local_testing/test_router_budget_limiter.py index b799da97ff..3242e997ae 100644 --- a/tests/local_testing/test_router_budget_limiter.py +++ b/tests/local_testing/test_router_budget_limiter.py @@ -74,7 +74,7 @@ async def test_provider_budgets_e2e_test(): { "model_name": "gpt-3.5-turbo", # openai model name "litellm_params": { # params for litellm completion/embedding call - "model": "azure/chatgpt-v-3", + "model": "azure/gpt-4.1-nano", "api_key": os.getenv("AZURE_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), "api_base": os.getenv("AZURE_API_BASE"), From 99314221e1c7195dd9bdaa7b62f8e0e8a3977a2e Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:27:34 -0700 Subject: [PATCH 74/76] test_router_azure_acompletion --- tests/local_testing/test_router.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index 3a07238eb3..4f9cf7aa95 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -612,7 +612,7 @@ def test_router_azure_acompletion(): { "model_name": "gpt-3.5-turbo", # openai model name "litellm_params": { # params for litellm completion/embedding call - "model": "azure/gpt-4o-new-test", + "model": "azure/gpt-4.1-nano", "api_key": old_api_key, "api_version": os.getenv("AZURE_API_VERSION"), "api_base": os.getenv("AZURE_API_BASE"), From 2ba34d22bd0c0022dce43d18c298493646b47aa5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:27:54 -0700 Subject: [PATCH 75/76] test_router_mock_request_with_mock_timeout_with_fallbacks --- tests/local_testing/test_mock_request.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/local_testing/test_mock_request.py b/tests/local_testing/test_mock_request.py index 6a9c5239f4..de8764edba 100644 --- a/tests/local_testing/test_mock_request.py +++ b/tests/local_testing/test_mock_request.py @@ -157,7 +157,7 @@ def test_router_mock_request_with_mock_timeout_with_fallbacks(): { "model_name": "azure-gpt", "litellm_params": { - "model": "azure/chatgpt-v-3", + "model": "azure/gpt-4.1-nano", "api_key": os.getenv("AZURE_API_KEY"), "api_base": os.getenv("AZURE_API_BASE"), }, From 973d4b71e7c72116d9ca00bf88bfce01ca68f8a3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 12:34:49 -0700 Subject: [PATCH 76/76] test_embedding_caching_azure --- tests/local_testing/test_caching.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/local_testing/test_caching.py b/tests/local_testing/test_caching.py index 8bd73af609..3fb3ae3db4 100644 --- a/tests/local_testing/test_caching.py +++ b/tests/local_testing/test_caching.py @@ -527,7 +527,7 @@ def test_embedding_caching_azure(): print(api_key) print(api_base) embedding1 = embedding( - model="text-embedding-ada-002", + model="azure/text-embedding-ada-002", input=["good morning from litellm", "this is another item"], api_key=api_key, api_base=api_base, @@ -540,7 +540,7 @@ def test_embedding_caching_azure(): time.sleep(1) start_time = time.time() embedding2 = embedding( - model="text-embedding-ada-002", + model="azure/text-embedding-ada-002", input=["good morning from litellm", "this is another item"], api_key=api_key, api_base=api_base,