diff --git a/docs/my-website/docs/providers/perplexity.md b/docs/my-website/docs/providers/perplexity.md index 68adf9939c6d..e3991c63bffb 100644 --- a/docs/my-website/docs/providers/perplexity.md +++ b/docs/my-website/docs/providers/perplexity.md @@ -120,7 +120,7 @@ All models listed here https://docs.perplexity.ai/docs/model-cards are supported -## Agentic Research API (Responses API) +## Agent API (Responses API) Requires v1.72.6+ @@ -196,7 +196,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/openai/gpt-4o", + model="perplexity/openai/gpt-5.2", input="Explain quantum computing in simple terms", custom_llm_provider="perplexity", max_output_tokens=500, @@ -215,7 +215,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/anthropic/claude-3-5-sonnet-20241022", + model="perplexity/anthropic/claude-sonnet-4-5", input="Write a short story about a robot learning to paint", custom_llm_provider="perplexity", max_output_tokens=500, @@ -234,7 +234,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/google/gemini-2.0-flash-exp", + model="perplexity/google/gemini-2.5-flash", input="Explain the concept of neural networks", custom_llm_provider="perplexity", max_output_tokens=500, @@ -253,7 +253,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/xai/grok-2-1212", + model="perplexity/xai/grok-4-1-fast-non-reasoning", input="What makes a good AI assistant?", custom_llm_provider="perplexity", max_output_tokens=500, @@ -276,7 +276,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/openai/gpt-4o", + model="perplexity/openai/gpt-5.2", input="What's the weather in San Francisco today?", custom_llm_provider="perplexity", tools=[{"type": "web_search"}], @@ -286,6 +286,78 @@ response = responses( print(response.output) ``` +### Function Calling + +The Agent API supports custom function tools. Pass function tools through unchanged: + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +response = responses( + model="perplexity/openai/gpt-5.2", + input="What's the weather in San Francisco?", + custom_llm_provider="perplexity", + tools=[ + {"type": "web_search"}, + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"}, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + }, + }, + }, + ], + instructions="Use tools when appropriate.", +) + +print(response.output) +``` + +### Structured Outputs + +Request JSON schema structured outputs via the `text` parameter: + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +response = responses( + model="perplexity/preset/pro-search", + input="Extract key facts about the Eiffel Tower", + custom_llm_provider="perplexity", + text={ + "format": { + "type": "json_schema", + "name": "facts", + "schema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "height_meters": {"type": "number"}, + "year_built": {"type": "integer"}, + }, + "required": ["name", "height_meters", "year_built"], + }, + "strict": True, + } + }, +) + +print(response.output) +``` + ### Reasoning Effort (Responses API) @@ -319,7 +391,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/anthropic/claude-3-5-sonnet-20241022", + model="perplexity/anthropic/claude-sonnet-4-5", input=[ {"type": "message", "role": "system", "content": "You are a helpful assistant."}, {"type": "message", "role": "user", "content": "What are the latest AI developments?"}, @@ -343,7 +415,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/openai/gpt-4o", + model="perplexity/openai/gpt-5.2", input="Tell me a story about space exploration", custom_llm_provider="perplexity", stream=True, @@ -360,23 +432,28 @@ for chunk in response: | Provider | Model Name | Function Call | |----------|------------|---------------| -| OpenAI | gpt-4o | `responses(model="perplexity/openai/gpt-4o", ...)` | -| OpenAI | gpt-4o-mini | `responses(model="perplexity/openai/gpt-4o-mini", ...)` | | OpenAI | gpt-5.2 | `responses(model="perplexity/openai/gpt-5.2", ...)` | -| Anthropic | claude-3-5-sonnet-20241022 | `responses(model="perplexity/anthropic/claude-3-5-sonnet-20241022", ...)` | -| Anthropic | claude-3-5-haiku-20241022 | `responses(model="perplexity/anthropic/claude-3-5-haiku-20241022", ...)` | -| Google | gemini-2.0-flash-exp | `responses(model="perplexity/google/gemini-2.0-flash-exp", ...)` | -| Google | gemini-2.0-flash-thinking-exp | `responses(model="perplexity/google/gemini-2.0-flash-thinking-exp", ...)` | -| xAI | grok-2-1212 | `responses(model="perplexity/xai/grok-2-1212", ...)` | -| xAI | grok-2-vision-1212 | `responses(model="perplexity/xai/grok-2-vision-1212", ...)` | +| OpenAI | gpt-5.1 | `responses(model="perplexity/openai/gpt-5.1", ...)` | +| OpenAI | gpt-5-mini | `responses(model="perplexity/openai/gpt-5-mini", ...)` | +| Anthropic | claude-opus-4-6 | `responses(model="perplexity/anthropic/claude-opus-4-6", ...)` | +| Anthropic | claude-opus-4-5 | `responses(model="perplexity/anthropic/claude-opus-4-5", ...)` | +| Anthropic | claude-sonnet-4-5 | `responses(model="perplexity/anthropic/claude-sonnet-4-5", ...)` | +| Anthropic | claude-haiku-4-5 | `responses(model="perplexity/anthropic/claude-haiku-4-5", ...)` | +| Google | gemini-3-pro-preview | `responses(model="perplexity/google/gemini-3-pro-preview", ...)` | +| Google | gemini-3-flash-preview | `responses(model="perplexity/google/gemini-3-flash-preview", ...)` | +| Google | gemini-2.5-pro | `responses(model="perplexity/google/gemini-2.5-pro", ...)` | +| Google | gemini-2.5-flash | `responses(model="perplexity/google/gemini-2.5-flash", ...)` | +| xAI | grok-4-1-fast-non-reasoning | `responses(model="perplexity/xai/grok-4-1-fast-non-reasoning", ...)` | +| Perplexity | sonar | `responses(model="perplexity/perplexity/sonar", ...)` | ### Available Presets -| Preset Name | Function Call | -|----------------|--------------------------------------------------------| -| fast-search | `responses(model="perplexity/preset/fast-search", ...)`| -| pro-search | `responses(model="perplexity/preset/pro-search", ...)` | -| deep-research | `responses(model="perplexity/preset/deep-research", ...)`| +| Preset Name | Function Call | +|-------------|---------------| +| fast-search | `responses(model="perplexity/preset/fast-search", ...)` | +| pro-search | `responses(model="perplexity/preset/pro-search", ...)` | +| deep-research | `responses(model="perplexity/preset/deep-research", ...)` | +| advanced-deep-research | `responses(model="perplexity/preset/advanced-deep-research", ...)` | ### Complete Example @@ -388,7 +465,7 @@ os.environ['PERPLEXITY_API_KEY'] = "" # Comprehensive example with multiple features response = responses( - model="perplexity/openai/gpt-4o", + model="perplexity/openai/gpt-5.2", input="Research the latest developments in quantum computing and provide sources", custom_llm_provider="perplexity", tools=[ diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index ad6eb6b4f324..cc0f818b0a01 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1242,6 +1242,16 @@ def completion_cost( # noqa: PLR0915 ) elif call_type in _VIDEO_CALL_TYPES: ### VIDEO GENERATION COST CALCULATION ### + # Extract custom model_info for deployment-specific pricing + _video_model_info: Optional[ModelInfo] = None + if custom_pricing and litellm_logging_obj is not None: + _litellm_params = getattr( + litellm_logging_obj, "litellm_params", None + ) + if _litellm_params is not None: + _metadata = _litellm_params.get("metadata", {}) or {} + _video_model_info = _metadata.get("model_info", None) + usage_obj = getattr(completion_response, "usage", None) if completion_response is not None and usage_obj: # Handle both dict and Pydantic Usage object @@ -1262,12 +1272,14 @@ def completion_cost( # noqa: PLR0915 model=model, duration_seconds=duration_seconds, custom_llm_provider=custom_llm_provider, + model_info=_video_model_info, ) # Fallback to default video cost calculation if no duration available return default_video_cost_calculator( model=model, duration_seconds=0.0, # Default to 0 if no duration available custom_llm_provider=custom_llm_provider, + model_info=_video_model_info, ) elif call_type in _SPEECH_CALL_TYPES: prompt_characters = litellm.utils._count_characters(text=prompt) @@ -1892,6 +1904,7 @@ def default_video_cost_calculator( model: str, duration_seconds: float, custom_llm_provider: Optional[str] = None, + model_info: Optional[ModelInfo] = None, ) -> float: """ Default video cost calculator for video generation @@ -1900,6 +1913,9 @@ def default_video_cost_calculator( model (str): Model name duration_seconds (float): Duration of the generated video in seconds custom_llm_provider (Optional[str]): Custom LLM provider + model_info (Optional[ModelInfo]): Deployment-level model info containing + custom video pricing. When provided, used before falling back to + the global litellm.model_cost lookup. Returns: float: Cost in USD for the video generation @@ -1907,42 +1923,47 @@ def default_video_cost_calculator( Raises: Exception: If model pricing not found in cost map """ - # Build model names for cost lookup - base_model_name = model - model_name_without_custom_llm_provider: Optional[str] = None - if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"): - model_name_without_custom_llm_provider = model.replace( - f"{custom_llm_provider}/", "" - ) - base_model_name = ( - f"{custom_llm_provider}/{model_name_without_custom_llm_provider}" - ) + # Use custom model_info pricing if provided (deployment-specific pricing) + cost_info: Optional[dict] = None + if model_info is not None: + cost_info = dict(model_info) + else: + # Build model names for cost lookup + base_model_name = model + model_name_without_custom_llm_provider: Optional[str] = None + if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"): + model_name_without_custom_llm_provider = model.replace( + f"{custom_llm_provider}/", "" + ) + base_model_name = ( + f"{custom_llm_provider}/{model_name_without_custom_llm_provider}" + ) - verbose_logger.debug(f"Looking up cost for video model: {base_model_name}") + verbose_logger.debug(f"Looking up cost for video model: {base_model_name}") - model_without_provider = model.split("/")[-1] + model_without_provider = model.split("/")[-1] - # Try model with provider first, fall back to base model name - cost_info: Optional[dict] = None - models_to_check: List[Optional[str]] = [ - base_model_name, - model, - model_without_provider, - model_name_without_custom_llm_provider, - ] - for _model in models_to_check: - if _model is not None and _model in litellm.model_cost: - cost_info = litellm.model_cost[_model] - break + # Try model with provider first, fall back to base model name + models_to_check: List[Optional[str]] = [ + base_model_name, + model, + model_without_provider, + model_name_without_custom_llm_provider, + ] + for _model in models_to_check: + if _model is not None and _model in litellm.model_cost: + cost_info = litellm.model_cost[_model] + break + + # If still not found, try with custom_llm_provider prefix + if cost_info is None and custom_llm_provider: + prefixed_model = f"{custom_llm_provider}/{model}" + if prefixed_model in litellm.model_cost: + cost_info = litellm.model_cost[prefixed_model] - # If still not found, try with custom_llm_provider prefix - if cost_info is None and custom_llm_provider: - prefixed_model = f"{custom_llm_provider}/{model}" - if prefixed_model in litellm.model_cost: - cost_info = litellm.model_cost[prefixed_model] if cost_info is None: raise Exception( - f"Model not found in cost map. Tried checking {models_to_check}" + f"Model not found in cost map for model={model}" ) # Check for video-specific cost per second first diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3648dc691ecd..0601e7e84551 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4709,6 +4709,7 @@ def get_model_cost_information( custom_pricing: Optional[bool], custom_llm_provider: Optional[str], init_response_obj: Union[Any, BaseModel, dict], + api_base: Optional[str] = None, ) -> StandardLoggingModelInformation: model_cost_name = _select_model_name_for_cost_calc( model=None, @@ -4723,7 +4724,9 @@ def get_model_cost_information( else: try: _model_cost_information = litellm.get_model_info( - model=model_cost_name, custom_llm_provider=custom_llm_provider + model=model_cost_name, + custom_llm_provider=custom_llm_provider, + api_base=api_base, ) model_cost_information = StandardLoggingModelInformation( model_map_key=model_cost_name, @@ -5236,6 +5239,7 @@ def get_standard_logging_object_payload( custom_pricing=custom_pricing, custom_llm_provider=kwargs.get("custom_llm_provider"), init_response_obj=init_response_obj, + api_base=litellm_params.get("api_base"), ) response_cost: float = kwargs.get("response_cost", 0) or 0.0 diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 7b485501f614..ba415af9a5a7 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1848,9 +1848,10 @@ def convert_to_anthropic_tool_invoke( break else: # Regular tool_use + sanitized_tool_id = _sanitize_anthropic_tool_use_id(tool_id) _anthropic_tool_use_param = AnthropicMessagesToolUseParam( type="tool_use", - id=tool_id, + id=sanitized_tool_id, name=tool_name, input=tool_input, ) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 3789f546d73b..3dfef07d4269 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -28,6 +28,7 @@ AIOHTTP_CONNECTOR_LIMIT, AIOHTTP_CONNECTOR_LIMIT_PER_HOST, AIOHTTP_KEEPALIVE_TIMEOUT, + AIOHTTP_NEEDS_CLEANUP_CLOSED, AIOHTTP_TTL_DNS_CACHE, DEFAULT_SSL_CIPHERS, ) @@ -876,9 +877,10 @@ def _create_aiohttp_transport( transport_connector_kwargs = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, - "enable_cleanup_closed": True, **connector_kwargs, } + if AIOHTTP_NEEDS_CLEANUP_CLOSED: + transport_connector_kwargs["enable_cleanup_closed"] = True if AIOHTTP_CONNECTOR_LIMIT > 0: transport_connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index c4d08c83a2a1..ed14b6a3318f 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -6,7 +6,7 @@ from httpx._models import Headers, Response import litellm -from litellm._logging import verbose_proxy_logger +from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -223,7 +223,9 @@ def get_api_key() -> Optional[str]: or get_secret_str("OLLAMA_API_KEY") ) - def get_model_info(self, model: str) -> ModelInfoBase: + def get_model_info( + self, model: str, api_base: Optional[str] = None + ) -> ModelInfoBase: """ curl http://localhost:11434/api/show -d '{ "name": "mistral" @@ -231,7 +233,11 @@ def get_model_info(self, model: str) -> ModelInfoBase: """ if model.startswith("ollama/") or model.startswith("ollama_chat/"): model = model.split("/", 1)[1] - api_base = get_secret_str("OLLAMA_API_BASE") or "http://localhost:11434" + api_base = ( + api_base + or get_secret_str("OLLAMA_API_BASE") + or "http://localhost:11434" + ) api_key = self.get_api_key() headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} @@ -242,8 +248,21 @@ def get_model_info(self, model: str) -> ModelInfoBase: headers=headers, ) except Exception as e: - raise Exception( - f"OllamaError: Error getting model info for {model}. Set Ollama API Base via `OLLAMA_API_BASE` environment variable. Error: {e}" + verbose_logger.debug( + "OllamaError: Could not get model info for %s from %s. Error: %s", + model, + api_base, + e, + ) + return ModelInfoBase( + key=model, + litellm_provider="ollama", + mode="chat", + input_cost_per_token=0.0, + output_cost_per_token=0.0, + max_tokens=None, + max_input_tokens=None, + max_output_tokens=None, ) model_info = response.json() diff --git a/litellm/llms/openai/cost_calculation.py b/litellm/llms/openai/cost_calculation.py index e5349db3af75..ac1e4a6b08f8 100644 --- a/litellm/llms/openai/cost_calculation.py +++ b/litellm/llms/openai/cost_calculation.py @@ -7,7 +7,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token -from litellm.types.utils import CallTypes, Usage +from litellm.types.utils import CallTypes, ModelInfo, Usage from litellm.utils import get_model_info @@ -129,7 +129,10 @@ def cost_per_second( def video_generation_cost( - model: str, duration_seconds: float, custom_llm_provider: Optional[str] = None + model: str, + duration_seconds: float, + custom_llm_provider: Optional[str] = None, + model_info: Optional[ModelInfo] = None, ) -> float: """ Calculates the cost for video generation based on duration in seconds. @@ -138,14 +141,18 @@ def video_generation_cost( - model: str, the model name without provider prefix - duration_seconds: float, the duration of the generated video in seconds - custom_llm_provider: str, the custom llm provider + - model_info: Optional[dict], deployment-level model info containing + custom video pricing. When provided, skips the global + get_model_info() lookup so that deployment-specific pricing is used. Returns: float - total_cost_in_usd """ ## GET MODEL INFO - model_info = get_model_info( - model=model, custom_llm_provider=custom_llm_provider or "openai" - ) + if model_info is None: + model_info = get_model_info( + model=model, custom_llm_provider=custom_llm_provider or "openai" + ) # Check for video-specific cost per second video_cost_per_second = model_info.get("output_cost_per_video_per_second") diff --git a/litellm/llms/perplexity/responses/__init__.py b/litellm/llms/perplexity/responses/__init__.py index 9bdf810e8390..3285a472113f 100644 --- a/litellm/llms/perplexity/responses/__init__.py +++ b/litellm/llms/perplexity/responses/__init__.py @@ -1,5 +1,5 @@ """ -Perplexity Agentic Research API (Responses API) module +Perplexity Agent API (Responses API) module """ from .transformation import PerplexityResponsesConfig diff --git a/litellm/llms/perplexity/responses/transformation.py b/litellm/llms/perplexity/responses/transformation.py index 178e76ea9704..6d2ed51600c1 100644 --- a/litellm/llms/perplexity/responses/transformation.py +++ b/litellm/llms/perplexity/responses/transformation.py @@ -1,5 +1,5 @@ """ -Transformation logic for Perplexity Agentic Research API (Responses API) +Transformation logic for Perplexity Agent API (Responses API) This module handles the translation between OpenAI's Responses API format and Perplexity's Responses API format, which supports: @@ -32,10 +32,10 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): """ - Configuration for Perplexity Agentic Research API (Responses API) + Configuration for Perplexity Agent API (Responses API) - - Reference: https://docs.perplexity.ai/agentic-research/quickstart + + Reference: https://docs.perplexity.ai/docs/agent-api/overview """ @property @@ -45,8 +45,9 @@ def custom_llm_provider(self) -> LlmProviders: def get_supported_openai_params(self, model: str) -> list: """ Perplexity Responses API supports a different set of parameters - + Ref: https://docs.perplexity.ai/api-reference/responses-post + Params aligned with response-echo fields and Open Responses spec. """ return [ "max_output_tokens", @@ -58,6 +59,23 @@ def get_supported_openai_params(self, model: str) -> list: "preset", "instructions", "models", # Model fallback support + "tool_choice", + "parallel_tool_calls", + "max_tool_calls", + "text", + "previous_response_id", + "store", + "background", + "truncation", + "metadata", + "safety_identifier", + "user", + "stream_options", + "top_logprobs", + "prompt_cache_key", + "frequency_penalty", + "presence_penalty", + "service_tier", ] def validate_environment( @@ -65,16 +83,15 @@ def validate_environment( ) -> dict: """Validate environment and set up headers""" # Get API key from environment - api_key = ( - get_secret_str("PERPLEXITYAI_API_KEY") - or get_secret_str("PERPLEXITY_API_KEY") + api_key = get_secret_str("PERPLEXITYAI_API_KEY") or get_secret_str( + "PERPLEXITY_API_KEY" ) - + if api_key: headers["Authorization"] = f"Bearer {api_key}" - + headers["Content-Type"] = "application/json" - + return headers def get_complete_url( @@ -84,15 +101,17 @@ def get_complete_url( ) -> str: """Get the complete URL for the Perplexity Responses API""" if api_base is None: - api_base = get_secret_str("PERPLEXITY_API_BASE") or "https://api.perplexity.ai" - + api_base = ( + get_secret_str("PERPLEXITY_API_BASE") or "https://api.perplexity.ai" + ) + # Ensure api_base doesn't end with a slash api_base = api_base.rstrip("/") - + # Add the responses endpoint return f"{api_base}/v1/responses" - def map_openai_params( + def map_openai_params( # noqa: PLR0915 self, response_api_optional_params: ResponsesAPIOptionalRequestParams, model: str, @@ -100,78 +119,136 @@ def map_openai_params( ) -> Dict: """ Map OpenAI Responses API parameters to Perplexity format - + Key differences: - Supports 'preset' parameter for predefined configurations - Supports 'instructions' parameter for system-level guidance - Tools are specified differently (web_search, fetch_url) """ mapped_params: Dict[str, Any] = {} - + # Map standard parameters if response_api_optional_params.get("max_output_tokens"): - mapped_params["max_output_tokens"] = response_api_optional_params["max_output_tokens"] - + mapped_params["max_output_tokens"] = response_api_optional_params[ + "max_output_tokens" + ] + if response_api_optional_params.get("temperature"): mapped_params["temperature"] = response_api_optional_params["temperature"] - + if response_api_optional_params.get("top_p"): mapped_params["top_p"] = response_api_optional_params["top_p"] - + if response_api_optional_params.get("stream"): mapped_params["stream"] = response_api_optional_params["stream"] - + if response_api_optional_params.get("stream_options"): - mapped_params["stream_options"] = response_api_optional_params["stream_options"] - + mapped_params["stream_options"] = response_api_optional_params[ + "stream_options" + ] + # Map Perplexity-specific parameters (using .get() with Any dict access) preset = response_api_optional_params.get("preset") # type: ignore if preset: mapped_params["preset"] = preset - + instructions = response_api_optional_params.get("instructions") # type: ignore if instructions: mapped_params["instructions"] = instructions - + if response_api_optional_params.get("reasoning"): mapped_params["reasoning"] = response_api_optional_params["reasoning"] - + tools = response_api_optional_params.get("tools") if tools: # Convert tools to list of dicts for transformation - tools_list = [dict(tool) if hasattr(tool, '__dict__') else tool for tool in tools] # type: ignore + tools_list = [dict(tool) if hasattr(tool, "__dict__") else tool for tool in tools] # type: ignore mapped_params["tools"] = self._transform_tools(tools_list) # type: ignore - + + # Tool control + if response_api_optional_params.get("tool_choice"): + mapped_params["tool_choice"] = response_api_optional_params["tool_choice"] + if response_api_optional_params.get("parallel_tool_calls") is not None: + mapped_params["parallel_tool_calls"] = response_api_optional_params[ + "parallel_tool_calls" + ] + if response_api_optional_params.get("max_tool_calls"): + mapped_params["max_tool_calls"] = response_api_optional_params[ + "max_tool_calls" + ] + + # Structured outputs + text_param = response_api_optional_params.get("text") + if text_param: + mapped_params["text"] = text_param + + # Conversation continuity + if response_api_optional_params.get("previous_response_id"): + mapped_params["previous_response_id"] = response_api_optional_params[ + "previous_response_id" + ] + + # Storage and lifecycle + if response_api_optional_params.get("store") is not None: + mapped_params["store"] = response_api_optional_params["store"] + if response_api_optional_params.get("background") is not None: + mapped_params["background"] = response_api_optional_params["background"] + if response_api_optional_params.get("truncation"): + mapped_params["truncation"] = response_api_optional_params["truncation"] + + # Metadata + if response_api_optional_params.get("metadata"): + mapped_params["metadata"] = response_api_optional_params["metadata"] + if response_api_optional_params.get("safety_identifier"): + mapped_params["safety_identifier"] = response_api_optional_params[ + "safety_identifier" + ] + if response_api_optional_params.get("user"): + mapped_params["user"] = response_api_optional_params["user"] + + # Additional + if response_api_optional_params.get("top_logprobs") is not None: + mapped_params["top_logprobs"] = response_api_optional_params["top_logprobs"] + if response_api_optional_params.get("prompt_cache_key"): + mapped_params["prompt_cache_key"] = response_api_optional_params[ + "prompt_cache_key" + ] + if response_api_optional_params.get("frequency_penalty") is not None: + mapped_params["frequency_penalty"] = response_api_optional_params[ + "frequency_penalty" # type: ignore[typeddict-item] + ] + if response_api_optional_params.get("presence_penalty") is not None: + mapped_params["presence_penalty"] = response_api_optional_params[ + "presence_penalty" # type: ignore[typeddict-item] + ] + if response_api_optional_params.get("service_tier"): + mapped_params["service_tier"] = response_api_optional_params["service_tier"] + return mapped_params def _transform_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """ - Transform tools to Perplexity format - - Perplexity supports: + Transform tools to Perplexity format. + + Perplexity supports (per public OpenAPI spec): - web_search: Performs web searches - fetch_url: Fetches content from URLs + - function: Function Calling """ perplexity_tools = [] - + for tool in tools: if isinstance(tool, dict): - tool_type = tool.get("type") - + tool_type = tool.get("type", "") + # Direct Perplexity tool format if tool_type in ["web_search", "fetch_url"]: perplexity_tools.append(tool) - - # OpenAI function format - try to map to Perplexity tools + + # Function tools: Perplexity supports them natively elif tool_type == "function": - function = tool.get("function", {}) - function_name = function.get("name", "") - - if function_name == "web_search" or "search" in function_name.lower(): - perplexity_tools.append({"type": "web_search"}) - elif function_name == "fetch_url" or "fetch" in function_name.lower(): - perplexity_tools.append({"type": "fetch_url"}) - + perplexity_tools.append(tool) + return perplexity_tools def transform_responses_api_request( @@ -204,24 +281,26 @@ def transform_responses_api_request( "model": model, "input": self._format_input(input), } - + # Add all optional parameters for key, value in response_api_optional_request_params.items(): data[key] = value - + return data - def _format_input(self, input: Union[str, ResponseInputParam]) -> Union[str, List[Dict[str, Any]]]: + def _format_input( + self, input: Union[str, ResponseInputParam] + ) -> Union[str, List[Dict[str, Any]]]: """ Format input for Perplexity Responses API - + The API accepts either: - A simple string for single-turn queries - An array of message objects for multi-turn conversations """ if isinstance(input, str): return input - + # Handle ResponseInputParam format if isinstance(input, list): formatted_messages = [] @@ -234,7 +313,7 @@ def _format_input(self, input: Union[str, ResponseInputParam]) -> Union[str, Lis } formatted_messages.append(formatted_message) return formatted_messages - + return str(input) def transform_response_api_response( @@ -267,10 +346,14 @@ def transform_response_api_response( # Transform usage to handle Perplexity's cost structure usage_data = raw_response_json.get("usage", {}) transformed_usage_dict = self._transform_usage(usage_data) - + # Convert usage dict to ResponseAPIUsage object - usage_obj = ResponseAPIUsage(**transformed_usage_dict) if transformed_usage_dict else None - + usage_obj = ( + ResponseAPIUsage(**transformed_usage_dict) + if transformed_usage_dict + else None + ) + # Map Perplexity response to OpenAI Responses API format response = ResponsesAPIResponse( id=raw_response_json.get("id", ""), @@ -283,11 +366,11 @@ def transform_response_api_response( ) return response - + def _transform_usage(self, usage_data: Dict[str, Any]) -> Dict[str, Any]: """ Transform Perplexity usage data to OpenAI format - + Perplexity returns: { "input_tokens": 100, @@ -300,7 +383,7 @@ def _transform_usage(self, usage_data: Dict[str, Any]) -> Dict[str, Any]: "total_cost": 0.0003 } } - + OpenAI expects: { "input_tokens": 100, @@ -314,7 +397,7 @@ def _transform_usage(self, usage_data: Dict[str, Any]) -> Dict[str, Any]: "output_tokens": usage_data.get("output_tokens", 0), "total_tokens": usage_data.get("total_tokens", 0), } - + # Transform cost from Perplexity format (dict) to OpenAI format (float) cost_obj = usage_data.get("cost") if isinstance(cost_obj, dict) and "total_cost" in cost_obj: @@ -322,20 +405,20 @@ def _transform_usage(self, usage_data: Dict[str, Any]) -> Dict[str, Any]: verbose_logger.debug( "Transformed Perplexity cost object to float: %s -> %s", cost_obj, - cost_obj["total_cost"] + cost_obj["total_cost"], ) elif cost_obj is not None: # If cost is already a float/number, use it as-is transformed["cost"] = cost_obj - + # Add input_tokens_details if present if "input_tokens_details" in usage_data: transformed["input_tokens_details"] = usage_data["input_tokens_details"] - + # Add output_tokens_details if present if "output_tokens_details" in usage_data: transformed["output_tokens_details"] = usage_data["output_tokens_details"] - + return transformed def transform_streaming_response( @@ -353,10 +436,10 @@ def transform_streaming_response( event_pydantic_model = PerplexityResponsesConfig.get_event_model_class( event_type=event_type ) - + # Transform Perplexity-specific fields to OpenAI format parsed_chunk = self._transform_perplexity_chunk(parsed_chunk) - + # Defensive: Handle error.code being null (similar to OpenAI implementation) try: error_obj = parsed_chunk.get("error") @@ -375,13 +458,13 @@ def transform_streaming_response( def _transform_perplexity_chunk(self, chunk: dict) -> dict: """ Transform Perplexity-specific fields in a streaming chunk to OpenAI format. - + This handles: - Converting Perplexity's cost object to a simple float """ # Make a copy to avoid modifying the original chunk = dict(chunk) - + # Transform usage.cost from Perplexity format to OpenAI format # Perplexity: {"currency": "USD", "input_cost": 0.0001, "output_cost": 0.0002, "total_cost": 0.0003} # OpenAI: 0.0003 (just the total_cost as a float) @@ -400,10 +483,10 @@ def _transform_perplexity_chunk(self, chunk: dict) -> dict: verbose_logger.debug( "Transformed Perplexity cost object to float: %s -> %s", cost_obj, - cost_obj["total_cost"] + cost_obj["total_cost"], ) except Exception as e: # If transformation fails, log and continue with original chunk verbose_logger.debug("Failed to transform Perplexity cost object: %s", e) - + return chunk diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9f2765652fd8..35a75d4caf9f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -714,8 +714,9 @@ async def _initialize_shared_aiohttp_session(): connector_kwargs = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, - "enable_cleanup_closed": True, } + if AIOHTTP_NEEDS_CLEANUP_CLOSED: + connector_kwargs["enable_cleanup_closed"] = True if AIOHTTP_CONNECTOR_LIMIT > 0: connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 536cb73de9e2..c4ff325db1f7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -13,8 +13,6 @@ from typing import ( TYPE_CHECKING, Any, - Callable, - Coroutine, Dict, List, Literal, @@ -2277,6 +2275,11 @@ def __init__( 0.0, float(os.getenv("PRISMA_AUTH_RECONNECT_LOCK_TIMEOUT_SECONDS", "0.1")), ) + self._engine_pidfd: int = -1 + self._engine_pid: int = 0 + self._watching_engine: bool = False + self._engine_confirmed_dead: bool = False + self._engine_wait_thread: Optional[threading.Thread] = None verbose_proxy_logger.debug("Success - Created Prisma Client") def get_request_status( @@ -3544,31 +3547,368 @@ async def disconnect(self): ) raise e - async def _run_reconnect_cycle( - self, timeout_seconds: Optional[float] = None - ) -> None: + def _get_engine_pid(self) -> int: + try: + engine = self.db._original_prisma._engine # type: ignore[attr-defined] + if engine is not None and engine.process is not None: + return engine.process.pid + except (AttributeError, TypeError): + pass + return 0 + + def _is_engine_alive(self) -> bool: + if self._engine_pid <= 0: + return True + try: + os.kill(self._engine_pid, 0) + return True + except ProcessLookupError: + return False + except (PermissionError, OSError): + return True + + @staticmethod + def _reap_all_zombies() -> set: + """Reap ALL zombie child processes via waitpid(-1, WNOHANG). + + Returns a set of reaped PIDs. As PID 1 in Docker (or any + process that spawns children), we must reap ALL terminated + children to prevent zombie accumulation. + """ + reaped: set = set() + while True: + try: + pid, _ = os.waitpid(-1, os.WNOHANG) + if pid == 0: + break + reaped.add(pid) + except ChildProcessError: + break + return reaped + + def _try_waitpid_watch(self, pid: int) -> bool: + """Watch engine PID via os.waitpid() in a dedicated thread. + + The thread blocks on os.waitpid(pid, 0) which is a kernel-level + wait and with zero CPU overhead, instant detection when the process exits. + When the process dies, the thread notifies the asyncio event loop + via call_soon_threadsafe. + + Returns True if the thread was started, False on failure. """ - Run a reconnect cycle with direct db operations and a single overall timeout - budget to avoid long retries on hot paths (e.g. auth). + try: + probe_pid, _ = os.waitpid(pid, os.WNOHANG) + except ChildProcessError: + verbose_proxy_logger.debug( + "PID %s is not a child process; skipping waitpid watch.", pid, + ) + return False + + if probe_pid == pid: + verbose_proxy_logger.warning( + "prisma-query-engine PID %s already dead at watch start.", pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + asyncio.create_task( + self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + ) + return True + + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return False + + thread = threading.Thread( + target=self._waitpid_thread_func, + args=(pid, loop), + daemon=True, + name=f"prisma-engine-waitpid-{pid}", + ) + thread.start() + self._engine_wait_thread = thread + return True + + def _waitpid_thread_func(self, pid: int, loop: asyncio.AbstractEventLoop) -> None: + """Thread function: block until engine PID exits, then notify event loop. + + Note: uvloop/libuv may reap the child first via waitpid(-1, WNOHANG) + in its SIGCHLD handler. In that case our waitpid raises ChildProcessError. + we still notify the event loop because the engine is dead either way. """ - async def _do_direct_reconnect() -> None: + try: + os.waitpid(pid, 0) + except ChildProcessError: + pass + except OSError: + pass + try: + loop.call_soon_threadsafe(self._on_engine_death_from_thread, pid) + except RuntimeError: + pass + + def _on_engine_death_from_thread(self, dead_pid: int) -> None: + """Called on the event loop thread when the waitpid thread detects engine death.""" + if self._engine_confirmed_dead: + return + if dead_pid != self._engine_pid: + return + verbose_proxy_logger.error( + "prisma-query-engine PID %s exited (waitpid thread); triggering reconnect.", + dead_pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + asyncio.create_task( + self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + ) + + def _try_pidfd_watch(self, pid: int) -> bool: + """ + Watch engine PID via pidfd_open + asyncio event loop reader. + + Returns True if pidfd watch was set up, False if unavailable or failed. + Broad OSError catch handles both ENOSYS and SECCOMP-blocked syscalls. + """ + if not hasattr(os, "pidfd_open"): + return False + fd = -1 + try: + fd = os.pidfd_open(pid, 0) # type: ignore[attr-defined] + asyncio.get_running_loop().add_reader(fd, self._on_pidfd_readable) + self._engine_pidfd = fd + return True + except OSError: + if fd >= 0: + os.close(fd) + return False + + def _on_pidfd_readable(self) -> None: + """pidfd became readable: engine process exited or became zombie. + + Sets _engine_confirmed_dead BEFORE cleanup so _run_reconnect_cycle + takes the heavy path (recreate Prisma client + re-arm watcher). + """ + if self._engine_confirmed_dead: + # Already handled -- just clean up pidfd resources. + if self._engine_pidfd >= 0: + try: + asyncio.get_running_loop().remove_reader(self._engine_pidfd) + except Exception: + pass + try: + os.close(self._engine_pidfd) + except OSError: + pass + self._engine_pidfd = -1 + return + dead_pid = self._engine_pid + verbose_proxy_logger.error( + "prisma-query-engine PID %s exited (pidfd event); triggering reconnect.", + dead_pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + asyncio.create_task( + self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + ) + + async def _poll_engine_proc(self) -> None: + """poll via os.kill(pid, 0) every 1s. + Only used when BOTH waitpid thread and pidfd are unavailable + (e.g., PID is not our child process and pidfd_open fails) + """ + while self._watching_engine and self._engine_pid > 0: try: - await self.db.disconnect() - except Exception as disconnect_err: + os.kill(self._engine_pid, 0) + except ProcessLookupError: + verbose_proxy_logger.error( + "prisma-query-engine PID %s gone; triggering reconnect.", + self._engine_pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + await self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + return + except (PermissionError, OSError): verbose_proxy_logger.debug( - "Prisma DB disconnect before reconnect failed (ignored): %s", - disconnect_err, + "Cannot signal PID %s; stopping engine poll.", + self._engine_pid, ) + self._cleanup_engine_watcher() + return + await asyncio.sleep(1) - await self.db.connect() - await self.db.query_raw("SELECT 1") + def _cleanup_engine_watcher(self) -> None: + """Clean up pidfd reader, waitpid thread ref, or stop polling and reset state.""" + self._watching_engine = False + if self._engine_pidfd >= 0: + try: + asyncio.get_running_loop().remove_reader(self._engine_pidfd) + except Exception: + pass + try: + os.close(self._engine_pidfd) + except OSError: + pass + self._engine_pidfd = -1 + self._engine_wait_thread = None + self._engine_pid = 0 + + async def _start_engine_watcher(self) -> None: + """ + Start watching the Prisma query engine process for death. + + Detection priority: + 1. os.waitpid() in a dedicated thread, works with all event loops. + 2. pidfd_open kernel fd registered with asyncio. + 3. os.kill(pid, 0) polling (1s), last-resort fallback when neither + waitpid thread nor pidfd are available. + """ + if self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None: + return + pid = self._get_engine_pid() + if pid == 0: + verbose_proxy_logger.debug("Could not find prisma-query-engine PID; engine death detection unavailable.") + return + self._engine_pid = pid + self._engine_confirmed_dead = False + verbose_proxy_logger.info("Found prisma-query-engine at PID %s.", pid) + waitpid_ok = self._try_waitpid_watch(pid) + pidfd_ok = False if waitpid_ok else self._try_pidfd_watch(pid) + if waitpid_ok: + verbose_proxy_logger.info( + "Watching engine PID %s via waitpid thread.", pid, + ) + elif pidfd_ok: + verbose_proxy_logger.info( + "Watching engine PID %s via pidfd.", pid, + ) + else: + verbose_proxy_logger.info( + "Watching engine PID %s via os.kill polling.", pid, + ) + self._watching_engine = True + asyncio.create_task(self._poll_engine_proc()) + + def _stop_engine_watcher(self) -> None: + """Stop watching the engine process and clean up all resources.""" + self._cleanup_engine_watcher() + self._engine_confirmed_dead = False + verbose_proxy_logger.debug("Stopped engine process watcher.") + + async def _run_reconnect_cycle( + self, timeout_seconds: Optional[float] = None + ) -> None: + """ + Run a reconnect cycle with a single overall timeout budget. + + Uses the _engine_confirmed_dead flag (set by waitpid thread / pidfd / poll + handlers) to choose between heavy reconnect (engine dead -- recreate + Prisma client, re-arm watcher) and lightweight reconnect (network + blip -- disconnect, connect, SELECT 1). + """ effective_timeout = ( - timeout_seconds - if timeout_seconds is not None - else self._db_watchdog_reconnect_timeout_seconds + timeout_seconds if timeout_seconds is not None else self._db_watchdog_reconnect_timeout_seconds + ) + + engine_is_dead = self._engine_confirmed_dead or ( + self._engine_pid > 0 and not self._is_engine_alive() + ) + + if engine_is_dead: + dead_pid = self._engine_pid + verbose_proxy_logger.warning( + "prisma-query-engine PID %s is dead; reconnecting.", + dead_pid, + ) + self._reap_all_zombies() + self._cleanup_engine_watcher() + self._engine_confirmed_dead = False + + async def _do_heavy_reconnect() -> None: + db_url = os.getenv("DATABASE_URL", "") + if not db_url: + verbose_proxy_logger.error("DATABASE_URL not set; cannot recreate Prisma client.") + raise RuntimeError("DATABASE_URL not set") + await self.db.recreate_prisma_client(db_url) + await self._start_engine_watcher() + + await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout) + else: + verbose_proxy_logger.debug("Performing Prisma DB reconnect (engine alive or unknown).") + + async def _do_direct_reconnect() -> None: + try: + await self.db.disconnect() + except Exception as disconnect_err: + verbose_proxy_logger.debug( + "Prisma DB disconnect before reconnect failed (ignored): %s", + disconnect_err, + ) + + await self.db.connect() + await self.db.query_raw("SELECT 1") + + await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout) + + async def _attempt_reconnect_inside_lock( + self, + force: bool, + reason: str, + timeout_seconds: Optional[float], + ) -> bool: + now = time.time() + if ( + force is False + and now - self._db_last_reconnect_attempt_ts + < self._db_reconnect_cooldown_seconds + ): + verbose_proxy_logger.debug( + "Skipping DB reconnect attempt inside lock due to cooldown. reason=%s", + reason, + ) + return False + + verbose_proxy_logger.warning( + "Attempting Prisma DB reconnect. reason=%s", reason ) - await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout) + + reconnect_succeeded = False + try: + await self._run_reconnect_cycle(timeout_seconds=timeout_seconds) + reconnect_succeeded = True + verbose_proxy_logger.info( + "Prisma DB reconnect succeeded. reason=%s", reason + ) + except Exception as reconnect_err: + verbose_proxy_logger.error( + "Prisma DB reconnect failed. reason=%s error=%s", + reason, + reconnect_err, + ) + finally: + self._db_last_reconnect_attempt_ts = time.time() + + return reconnect_succeeded async def attempt_db_reconnect( self, @@ -3595,59 +3935,10 @@ async def attempt_db_reconnect( ) return False - async def _attempt_reconnect_inside_lock() -> bool: - now = time.time() - if ( - force is False - and now - self._db_last_reconnect_attempt_ts - < self._db_reconnect_cooldown_seconds - ): - verbose_proxy_logger.debug( - "Skipping DB reconnect attempt inside lock due to cooldown. reason=%s", - reason, - ) - return False - - verbose_proxy_logger.warning( - "Attempting Prisma DB reconnect. reason=%s", reason - ) - - reconnect_succeeded = False - try: - await self._run_reconnect_cycle(timeout_seconds=timeout_seconds) - reconnect_succeeded = True - verbose_proxy_logger.info( - "Prisma DB reconnect succeeded. reason=%s", reason - ) - except Exception as reconnect_err: - verbose_proxy_logger.error( - "Prisma DB reconnect failed. reason=%s error=%s", - reason, - reconnect_err, - ) - finally: - # Start cooldown after reconnect attempt has completed. - self._db_last_reconnect_attempt_ts = time.time() - - return reconnect_succeeded - if lock_timeout_seconds is None: async with self._db_reconnect_lock: - return await _attempt_reconnect_inside_lock() - - return await self._attempt_reconnect_with_lock_timeout( - _attempt_reconnect_inside_lock, - reason=reason, - lock_timeout_seconds=lock_timeout_seconds, - ) + return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) - async def _attempt_reconnect_with_lock_timeout( - self, - reconnect_fn: Callable[[], Coroutine[Any, Any, bool]], - reason: str, - lock_timeout_seconds: float, - ) -> bool: - """Acquire the reconnect lock with a timeout, then run reconnect_fn.""" lock_acquired_by_timeout_task = False async def _acquire_reconnect_lock() -> bool: @@ -3695,14 +3986,14 @@ async def _acquire_reconnect_lock() -> bool: return False try: - return await reconnect_fn() + return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) finally: self._db_reconnect_lock.release() async def start_db_health_watchdog_task(self) -> None: - """ - Start a background task that probes DB health and attempts reconnect on failure. - """ + """Start background tasks that monitor DB health: + - A periodic SELECT 1 probe that triggers reconnect on network/connection failure. + - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling.""" if self._db_health_watchdog_enabled is not True: verbose_proxy_logger.debug( "Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED" @@ -3720,11 +4011,11 @@ async def start_db_health_watchdog_task(self) -> None: self._db_health_watchdog_probe_timeout_seconds, self._db_watchdog_reconnect_timeout_seconds, ) + await self._start_engine_watcher() async def stop_db_health_watchdog_task(self) -> None: - """ - Stop DB health watchdog task gracefully. - """ + """Stop DB health watchdog task and engine watcher gracefully.""" + self._stop_engine_watcher() if self._db_health_watchdog_task is None: return self._db_health_watchdog_task.cancel() diff --git a/litellm/utils.py b/litellm/utils.py index fef99c8b2010..b4de1e00efc6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5385,14 +5385,18 @@ def _get_max_position_embeddings(model_name: str) -> Optional[int]: @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) def _cached_get_model_info_helper( - model: str, custom_llm_provider: Optional[str] + model: str, + custom_llm_provider: Optional[str], + api_base: Optional[str] = None, ) -> ModelInfoBase: """ _get_model_info_helper wrapped with lru_cache Speed Optimization to hit high RPS """ - return _get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + return _get_model_info_helper( + model=model, custom_llm_provider=custom_llm_provider, api_base=api_base + ) def get_provider_info( @@ -5428,7 +5432,9 @@ def _is_potential_model_name_in_model_cost( def _get_model_info_helper( # noqa: PLR0915 - model: str, custom_llm_provider: Optional[str] = None + model: str, + custom_llm_provider: Optional[str] = None, + api_base: Optional[str] = None, ) -> ModelInfoBase: """ Helper for 'get_model_info'. Separated out to avoid infinite loop caused by returning 'supported_openai_param's @@ -5486,7 +5492,7 @@ def _get_model_info_helper( # noqa: PLR0915 elif ( custom_llm_provider == "ollama" or custom_llm_provider == "ollama_chat" ) and not _is_potential_model_name_in_model_cost(potential_model_names): - return litellm.OllamaConfig().get_model_info(model) + return litellm.OllamaConfig().get_model_info(model, api_base=api_base) else: """ Check if: (in order of specificity) @@ -5725,8 +5731,6 @@ def _get_model_info_helper( # noqa: PLR0915 ) except Exception as e: verbose_logger.debug(f"Error getting model info: {e}") - if "OllamaError" in str(e): - raise e raise Exception( "This model isn't mapped yet. model={}, custom_llm_provider={}. Add it here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json.".format( model, custom_llm_provider @@ -5735,7 +5739,11 @@ def _get_model_info_helper( # noqa: PLR0915 @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) -def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> ModelInfo: +def get_model_info( + model: str, + custom_llm_provider: Optional[str] = None, + api_base: Optional[str] = None, +) -> ModelInfo: """ Get a dict for the maximum tokens (context window), input_cost_per_token, output_cost_per_token for a given model. @@ -5813,6 +5821,7 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod _model_info = _get_model_info_helper( model=model, custom_llm_provider=custom_llm_provider, + api_base=api_base, ) provider_info = get_provider_info( diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3909ce4c8b08..624714feb872 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -26555,65 +26555,124 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "perplexity/preset/fast-search": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_preset": true, + "supports_function_calling": true + }, "perplexity/preset/pro-search": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_preset": true + "supports_preset": true, + "supports_function_calling": true }, - "perplexity/openai/gpt-4o": { + "perplexity/preset/deep-research": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_preset": true, + "supports_function_calling": true }, - "perplexity/openai/gpt-4o-mini": { + "perplexity/preset/advanced-deep-research": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_preset": true, + "supports_function_calling": true }, "perplexity/openai/gpt-5.2": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": true + "supports_reasoning": true, + "supports_function_calling": true + }, + "perplexity/openai/gpt-5.1": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/anthropic/claude-3-5-sonnet-20241022": { + "perplexity/openai/gpt-5-mini": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/anthropic/claude-3-5-haiku-20241022": { + "perplexity/anthropic/claude-opus-4-6": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/google/gemini-2.0-flash-exp": { + "perplexity/anthropic/claude-opus-4-5": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/google/gemini-2.0-flash-thinking-exp": { + "perplexity/anthropic/claude-sonnet-4-5": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": true + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/anthropic/claude-haiku-4-5": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-3-pro-preview": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-3-flash-preview": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-2.5-pro": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/xai/grok-2-1212": { + "perplexity/google/gemini-2.5-flash": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/xai/grok-2-vision-1212": { + "perplexity/xai/grok-4-1-fast-non-reasoning": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/perplexity/sonar": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true }, "publicai/aisingapore/Qwen-SEA-LION-v4-32B-IT": { "input_cost_per_token": 0.0, diff --git a/tests/litellm/proxy/test_prisma_engine_watchdog.py b/tests/litellm/proxy/test_prisma_engine_watchdog.py new file mode 100644 index 000000000000..011b8002db29 --- /dev/null +++ b/tests/litellm/proxy/test_prisma_engine_watchdog.py @@ -0,0 +1,446 @@ +""" +Tests for PrismaClient engine watchdog: death detection and automatic reconnect. + +Covers: +- Engine PID discovery and liveness check +- Engine process gone (os.kill raises ProcessLookupError) → reconnect triggered +- PermissionError from os.kill → treated as alive (process exists but not ours) +- pidfd handler → schedules attempt_db_reconnect even when lock is held +- waitpid thread → instant cross-platform detection, triggers reconnect +- _run_reconnect_cycle branches: heavy path (engine dead) vs lightweight path (engine alive) +- _engine_confirmed_dead flag ensures heavy reconnect even after _engine_pid reset +- Successful heavy reconnect → watcher re-armed for new process +- Missing DATABASE_URL → graceful RuntimeError in reconnect cycle +- Shutdown → polling loop exits cleanly +""" + +import asyncio +import os +import threading +import time +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +@pytest.fixture(autouse=True) +def mock_prisma_binary(): + """Mock prisma.Prisma to avoid requiring generated Prisma binaries for unit tests.""" + import sys + + mock_module = MagicMock() + with patch.dict(sys.modules, {"prisma": mock_module}): + yield + + +@pytest.fixture +def mock_proxy_logging(): + proxy_logging = AsyncMock(spec=ProxyLogging) + proxy_logging.failure_handler = AsyncMock() + return proxy_logging + + +@pytest.fixture +def engine_client(mock_proxy_logging) -> PrismaClient: + """ + Minimal PrismaClient fixture for engine watchdog tests. + Uses the real constructor pattern from PR #21706 (database_url). + """ + client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) + client.db = MagicMock() + client.db.recreate_prisma_client = AsyncMock() + client.db.disconnect = AsyncMock(return_value=None) + client.db.connect = AsyncMock(return_value=None) + client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + return client + + +# --------------------------------------------------------------------------- +# _is_engine_alive +# --------------------------------------------------------------------------- + + +def test_is_engine_alive_returns_true_when_pid_unknown(engine_client): + """_is_engine_alive returns True when no engine PID is tracked.""" + engine_client._engine_pid = 0 + assert engine_client._is_engine_alive() is True + + +def test_is_engine_alive_returns_false_when_process_gone(engine_client): + """_is_engine_alive returns False when os.kill raises ProcessLookupError.""" + engine_client._engine_pid = 9999 + with patch("os.kill", side_effect=ProcessLookupError): + assert engine_client._is_engine_alive() is False + + +def test_is_engine_alive_returns_true_on_permission_error(engine_client): + """_is_engine_alive returns True when os.kill raises PermissionError (process exists but not ours).""" + engine_client._engine_pid = 1234 + with patch("os.kill", side_effect=PermissionError): + assert engine_client._is_engine_alive() is True + + +def test_is_engine_alive_returns_true_for_running_process(engine_client): + """_is_engine_alive returns True when os.kill succeeds (process running).""" + engine_client._engine_pid = 1234 + with patch("os.kill"): + assert engine_client._is_engine_alive() is True + + +# --------------------------------------------------------------------------- +# _poll_engine_proc — calls attempt_db_reconnect on death +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_poll_missing_process_triggers_reconnect(engine_client) -> None: + """Polling loop triggers attempt_db_reconnect when os.kill raises ProcessLookupError.""" + engine_client._engine_pid = 1234 + engine_client._watching_engine = True + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + with patch("os.kill", side_effect=ProcessLookupError): + await engine_client._poll_engine_proc() + + engine_client.attempt_db_reconnect.assert_awaited_once_with( + reason="engine_process_death", + force=True, + ) + + +@pytest.mark.asyncio +async def test_poll_permission_error_stops_polling(engine_client) -> None: + """Polling loop stops cleanly when os.kill raises PermissionError (process not ours).""" + engine_client._engine_pid = 1234 + engine_client._watching_engine = True + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + with patch("os.kill", side_effect=PermissionError): + await engine_client._poll_engine_proc() + + # PermissionError means process exists but isn't ours — no reconnect, just stop polling + engine_client.attempt_db_reconnect.assert_not_awaited() + assert engine_client._watching_engine is False + assert engine_client._engine_pid == 0 + + +@pytest.mark.asyncio +async def test_stop_loop_halts_polling(engine_client) -> None: + """Polling loop exits cleanly when _stop_engine_watcher is called.""" + engine_client._engine_pid = 1234 + engine_client._watching_engine = True + + async def stop_during_sleep(_duration: float) -> None: + engine_client._stop_engine_watcher() + + with ( + patch("os.kill"), + patch("asyncio.sleep", side_effect=stop_during_sleep), + ): + await engine_client._poll_engine_proc() + + assert engine_client._watching_engine is False + assert engine_client._engine_pid == 0 + + +# --------------------------------------------------------------------------- +# _on_pidfd_readable — calls attempt_db_reconnect +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_pidfd_readable_schedules_reconnect(engine_client) -> None: + """pidfd handler schedules attempt_db_reconnect via asyncio.create_task.""" + engine_client._engine_pid = 1234 + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + with patch("asyncio.create_task", side_effect=capture_task): + engine_client._on_pidfd_readable() + + # Run the captured coroutine to completion + assert len(created_coros) == 1 + await created_coros[0] + + engine_client.attempt_db_reconnect.assert_awaited_once_with( + reason="engine_process_death", + force=True, + ) + + +@pytest.mark.asyncio +async def test_pidfd_schedules_reconnect_task_when_lock_held(engine_client) -> None: + """pidfd handler schedules reconnect task even when _db_reconnect_lock is held.""" + engine_client._engine_pid = 1234 + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + async with engine_client._db_reconnect_lock: + with patch("asyncio.create_task", side_effect=capture_task): + engine_client._on_pidfd_readable() + + for coro in created_coros: + coro.close() + + assert len(created_coros) == 1 + + +# --------------------------------------------------------------------------- +# _run_reconnect_cycle — engine liveness branching +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_heavy_path_when_engine_dead( + engine_client, +) -> None: + """_run_reconnect_cycle calls recreate_prisma_client when engine is dead.""" + engine_client._engine_pid = 1234 + engine_client._start_engine_watcher = AsyncMock() + + with ( + patch.object(engine_client, "_is_engine_alive", return_value=False), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + patch("os.waitpid", side_effect=ChildProcessError), + ): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") + engine_client._start_engine_watcher.assert_awaited_once() + engine_client.db.connect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_heavy_path_when_confirmed_dead( + engine_client, +) -> None: + """_run_reconnect_cycle takes heavy path when _engine_confirmed_dead is set. + + This is the critical race-condition fix: SIGCHLD/pidfd handlers set + _engine_confirmed_dead BEFORE _cleanup_engine_watcher resets _engine_pid + to 0, so the heavy path executes even after cleanup. + """ + engine_client._engine_pid = 0 # Already reset by cleanup! + engine_client._engine_confirmed_dead = True # But flag survives cleanup + engine_client._start_engine_watcher = AsyncMock() + + with ( + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + patch("os.waitpid", side_effect=ChildProcessError), + ): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") + engine_client._start_engine_watcher.assert_awaited_once() + engine_client.db.connect.assert_not_awaited() + assert engine_client._engine_confirmed_dead is False # Reset after use + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_lightweight_path_when_engine_alive( + engine_client, +) -> None: + """_run_reconnect_cycle uses disconnect/connect when engine is alive.""" + engine_client._engine_pid = 1234 + + with patch.object(engine_client, "_is_engine_alive", return_value=True): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.connect.assert_awaited_once() + engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") + engine_client.db.recreate_prisma_client.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_lightweight_path_when_pid_unknown( + engine_client, +) -> None: + """_run_reconnect_cycle uses lightweight path when engine PID is not tracked.""" + engine_client._engine_pid = 0 + + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.connect.assert_awaited_once() + engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") + engine_client.db.recreate_prisma_client.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_heavy_path_raises_without_database_url( + engine_client, +) -> None: + """Heavy reconnect raises RuntimeError when DATABASE_URL is not set.""" + engine_client._engine_pid = 1234 + + with ( + patch.object(engine_client, "_is_engine_alive", return_value=False), + patch.dict(os.environ, {}, clear=True), + patch("os.waitpid", side_effect=ChildProcessError), + ): + with pytest.raises(RuntimeError, match="DATABASE_URL not set"): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.recreate_prisma_client.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# start/stop lifecycle integration +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_start_watchdog_task_also_starts_engine_watcher( + engine_client, +) -> None: + """start_db_health_watchdog_task() also starts engine watcher.""" + engine_client._start_engine_watcher = AsyncMock() + + loop = asyncio.get_running_loop() + dummy_task = loop.create_task(asyncio.sleep(3600)) + + def fake_create_task(coro): + coro.close() + return dummy_task + + with patch("asyncio.create_task", side_effect=fake_create_task): + await engine_client.start_db_health_watchdog_task() + + engine_client._start_engine_watcher.assert_awaited_once() + dummy_task.cancel() + try: + await dummy_task + except asyncio.CancelledError: + pass + + +@pytest.mark.asyncio +async def test_stop_watchdog_task_also_stops_engine_watcher( + engine_client, +) -> None: + """stop_db_health_watchdog_task() also stops engine watcher.""" + engine_client._stop_engine_watcher = MagicMock() + + loop = asyncio.get_running_loop() + dummy_task = loop.create_task(asyncio.sleep(3600)) + engine_client._db_health_watchdog_task = dummy_task + + await engine_client.stop_db_health_watchdog_task() + + engine_client._stop_engine_watcher.assert_called_once() + assert engine_client._db_health_watchdog_task is None + + +# --------------------------------------------------------------------------- +# waitpid thread (cross-platform) +# --------------------------------------------------------------------------- + + +def test_try_waitpid_watch_returns_false_when_not_child(engine_client): + """_try_waitpid_watch returns False when PID is not our child process.""" + engine_client._engine_pid = 9999 + with patch("os.waitpid", side_effect=ChildProcessError): + assert engine_client._try_waitpid_watch(9999) is False + assert engine_client._engine_wait_thread is None + + +def test_try_waitpid_watch_starts_thread_for_child(engine_client): + """_try_waitpid_watch starts a daemon thread when PID is our child.""" + engine_client._engine_pid = 1234 + mock_thread = MagicMock() + mock_loop = MagicMock() + with ( + patch("os.waitpid", return_value=(0, 0)), + patch("asyncio.get_running_loop", return_value=mock_loop), + patch("threading.Thread", return_value=mock_thread) as mock_thread_cls, + ): + result = engine_client._try_waitpid_watch(1234) + assert result is True + mock_thread.start.assert_called_once() + assert engine_client._engine_wait_thread is mock_thread + + +@pytest.mark.asyncio +async def test_try_waitpid_watch_handles_already_dead_engine(engine_client) -> None: + """_try_waitpid_watch detects engine already dead at watch start.""" + engine_client._engine_pid = 1234 + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + waitpid_calls = iter([(1234, 0)]) + + def mock_waitpid(pid, flags): + if pid == -1: + raise ChildProcessError + return next(waitpid_calls) + + with ( + patch("os.waitpid", side_effect=mock_waitpid), + patch("asyncio.create_task", side_effect=capture_task), + ): + result = engine_client._try_waitpid_watch(1234) + + assert result is True + assert engine_client._engine_confirmed_dead is True + assert len(created_coros) == 1 + created_coros[0].close() + + +@pytest.mark.asyncio +async def test_on_engine_death_from_thread_triggers_reconnect(engine_client) -> None: + """waitpid thread callback schedules attempt_db_reconnect.""" + engine_client._engine_pid = 1234 + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + with patch("asyncio.create_task", side_effect=capture_task): + engine_client._on_engine_death_from_thread(1234) + + assert len(created_coros) == 1 + await created_coros[0] + + engine_client.attempt_db_reconnect.assert_awaited_once_with( + reason="engine_process_death", + force=True, + ) + + +def test_on_engine_death_from_thread_no_double_trigger(engine_client): + """waitpid thread callback does not trigger reconnect if already confirmed dead.""" + engine_client._engine_pid = 1234 + engine_client._engine_confirmed_dead = True + + with patch("asyncio.create_task") as mock_create_task: + engine_client._on_engine_death_from_thread(1234) + + mock_create_task.assert_not_called() + + +def test_on_engine_death_from_thread_ignores_stale_pid(engine_client): + """waitpid thread callback ignores death notification for a stale PID.""" + engine_client._engine_pid = 5678 + + with patch("asyncio.create_task") as mock_create_task: + engine_client._on_engine_death_from_thread(1234) + + mock_create_task.assert_not_called() diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index 8974632631dd..88bca007740d 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -973,6 +973,54 @@ def test_convert_to_anthropic_tool_invoke_regular_tool(): assert result[0]["input"] == {"location": "San Francisco"} +def test_convert_to_anthropic_tool_invoke_sanitizes_invalid_ids(): + """Test that tool_use IDs with invalid characters are sanitized. + + Anthropic requires tool_use_id to match ^[a-zA-Z0-9_-]+$. + IDs from external frameworks (e.g. MiniMax) may contain characters + like colons that violate this pattern. + """ + tool_calls = [ + { + "id": "sessions_history:183", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "Boston"}', + }, + }, + { + "id": "composio.NOTION_SEARCH", + "type": "function", + "function": { + "name": "search_notes", + "arguments": '{"query": "test"}', + }, + }, + ] + + result = convert_to_anthropic_tool_invoke(tool_calls) + + assert len(result) == 2 + # Colons replaced with underscores + assert result[0]["id"] == "sessions_history_183" + # Dots replaced with underscores + assert result[1]["id"] == "composio_NOTION_SEARCH" + # Valid IDs should pass through unchanged + valid_tool_calls = [ + { + "id": "toolu_01ABC-xyz_123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "NYC"}', + }, + } + ] + valid_result = convert_to_anthropic_tool_invoke(valid_tool_calls) + assert valid_result[0]["id"] == "toolu_01ABC-xyz_123" + + def test_convert_to_anthropic_tool_invoke_server_tool(): """ Test that server_tool_use (srvtoolu_) is reconstructed as server_tool_use. diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py new file mode 100644 index 000000000000..c8c0e09c08c0 --- /dev/null +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py @@ -0,0 +1,31 @@ +from unittest.mock import MagicMock, patch + + +def test_create_aiohttp_transport_sets_enable_cleanup_closed_when_needed(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + monkeypatch.setattr(http_handler_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", True) + + with patch.object(http_handler_module, "TCPConnector", return_value=connector_mock) as mock_tcp_connector: + with patch.object(http_handler_module, "ClientSession", return_value=session_mock): + transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(shared_session=None) + transport._get_valid_client_session() + + assert mock_tcp_connector.call_args.kwargs["enable_cleanup_closed"] is True + + +def test_create_aiohttp_transport_omits_enable_cleanup_closed_when_not_needed(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + monkeypatch.setattr(http_handler_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", False) + + with patch.object(http_handler_module, "TCPConnector", return_value=connector_mock) as mock_tcp_connector: + with patch.object(http_handler_module, "ClientSession", return_value=session_mock): + transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(shared_session=None) + transport._get_valid_client_session() + + assert "enable_cleanup_closed" not in mock_tcp_connector.call_args.kwargs diff --git a/tests/test_litellm/llms/ollama/test_ollama_model_info.py b/tests/test_litellm/llms/ollama/test_ollama_model_info.py index 7eef15cd4d29..5585e9d1e0ee 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_model_info.py +++ b/tests/test_litellm/llms/ollama/test_ollama_model_info.py @@ -138,6 +138,88 @@ def mock_get(url, headers): assert models == ["ollama/llama2"] +class TestOllamaGetModelInfo: + """Tests for OllamaConfig.get_model_info() api_base threading and graceful fallback.""" + + def test_get_model_info_uses_provided_api_base(self, monkeypatch): + """When api_base is passed, get_model_info should use it instead of env var or default.""" + from litellm.llms.ollama.completion.transformation import OllamaConfig + + captured_urls = [] + + def mock_post(url, json, headers=None): + captured_urls.append(url) + resp = DummyResponse( + {"template": "{{ .System }} tools {{ .Prompt }}", "model_info": {"context_length": 4096}}, + status_code=200, + ) + return resp + + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + + config = OllamaConfig() + result = config.get_model_info("llama3", api_base="http://my-remote-server:11434") + + assert captured_urls[0] == "http://my-remote-server:11434/api/show" + assert result["max_tokens"] == 4096 + + def test_get_model_info_falls_back_to_env_var(self, monkeypatch): + """When no api_base is passed, should fall back to OLLAMA_API_BASE env var.""" + from litellm.llms.ollama.completion.transformation import OllamaConfig + + captured_urls = [] + + def mock_post(url, json, headers=None): + captured_urls.append(url) + return DummyResponse({"template": "", "model_info": {}}, status_code=200) + + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + monkeypatch.setenv("OLLAMA_API_BASE", "http://env-server:11434") + + config = OllamaConfig() + config.get_model_info("llama3") + + assert captured_urls[0] == "http://env-server:11434/api/show" + + def test_get_model_info_graceful_fallback_on_connection_error(self, monkeypatch): + """When the Ollama server is unreachable, should return defaults instead of raising.""" + from litellm.llms.ollama.completion.transformation import OllamaConfig + + def mock_post(url, json, headers=None): + raise ConnectionError("Connection refused") + + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + + config = OllamaConfig() + result = config.get_model_info("llama3", api_base="http://unreachable:11434") + + assert result["key"] == "llama3" + assert result["litellm_provider"] == "ollama" + assert result["input_cost_per_token"] == 0.0 + assert result["output_cost_per_token"] == 0.0 + assert result["max_tokens"] is None + + def test_get_model_info_strips_ollama_prefix(self, monkeypatch): + """Should strip 'ollama/' or 'ollama_chat/' prefix from model name.""" + from litellm.llms.ollama.completion.transformation import OllamaConfig + + captured_json = [] + + def mock_post(url, json, headers=None): + captured_json.append(json) + return DummyResponse({"template": "", "model_info": {}}, status_code=200) + + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + + config = OllamaConfig() + config.get_model_info("ollama/llama3", api_base="http://localhost:11434") + assert captured_json[0]["name"] == "llama3" + + config.get_model_info("ollama_chat/llama3", api_base="http://localhost:11434") + assert captured_json[1]["name"] == "llama3" + + class TestOllamaAuthHeaders: """Tests for Ollama authentication header handling in completion calls.""" diff --git a/tests/test_litellm/llms/perplexity/responses/test_perplexity_responses_transformation.py b/tests/test_litellm/llms/perplexity/responses/test_perplexity_responses_transformation.py new file mode 100644 index 000000000000..cdd4ef913f97 --- /dev/null +++ b/tests/test_litellm/llms/perplexity/responses/test_perplexity_responses_transformation.py @@ -0,0 +1,381 @@ +""" +Tests for Perplexity Responses API transformation + +Tests the PerplexityResponsesConfig class that handles Perplexity-specific +transformations for the Agent API (Responses API). + +Source: litellm/llms/perplexity/responses/transformation.py +""" + +import os +import sys + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.perplexity.responses.transformation import PerplexityResponsesConfig +from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + + +class TestPerplexityResponsesTransformation: + """Test Perplexity Responses API configuration and transformations""" + + def test_function_tool_passthrough(self): + """Function tools with name/description/parameters are preserved""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"}, + "unit": { + "type": "string", + "enum": ["celsius", "fahrenheit"], + }, + }, + }, + }, + } + ] + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert "tools" in result + assert len(result["tools"]) == 1 + assert result["tools"][0]["type"] == "function" + assert result["tools"][0]["function"]["name"] == "get_weather" + assert ( + result["tools"][0]["function"]["description"] == "Get the current weather" + ) + assert "parameters" in result["tools"][0]["function"] + + def test_web_search_tool_passthrough(self): + """web_search tools are passed through unchanged""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams(tools=[{"type": "web_search"}]) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert "tools" in result + assert len(result["tools"]) == 1 + assert result["tools"][0]["type"] == "web_search" + + def test_fetch_url_tool_passthrough(self): + """fetch_url tools are passed through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams(tools=[{"type": "fetch_url"}]) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert "tools" in result + assert len(result["tools"]) == 1 + assert result["tools"][0]["type"] == "fetch_url" + + def test_mixed_tools_function_and_web_search(self): + """Mixed function and web_search tools are transformed correctly""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + tools=[ + {"type": "web_search"}, + { + "type": "function", + "function": { + "name": "custom_tool", + "description": "A custom tool", + "parameters": {"type": "object"}, + }, + }, + ] + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert len(result["tools"]) == 2 + assert result["tools"][0]["type"] == "web_search" + assert result["tools"][1]["type"] == "function" + assert result["tools"][1]["function"]["name"] == "custom_tool" + + def test_tool_choice_mapping(self): + """tool_choice passes through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + tool_choice="required", temperature=0.7 + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert result.get("tool_choice") == "required" + + def test_parallel_tool_calls(self): + """parallel_tool_calls passes through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + parallel_tool_calls=True, temperature=0.7 + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert result.get("parallel_tool_calls") is True + + def test_max_tool_calls_mapping(self): + """max_tool_calls passes through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams(max_tool_calls=5, temperature=0.7) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert result.get("max_tool_calls") == 5 + + def test_text_passthrough(self): + """text param passes through as-is (Perplexity accepts Open Responses format directly)""" + config = PerplexityResponsesConfig() + + text_value = { + "format": { + "type": "json_schema", + "name": "weather_response", + "schema": { + "type": "object", + "properties": {"temp": {"type": "number"}}, + }, + "strict": True, + } + } + + params = ResponsesAPIOptionalRequestParams( + text=text_value, + temperature=0.7, + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert "text" in result + assert result["text"] == text_value + assert "response_format" not in result + + def test_previous_response_id(self): + """previous_response_id passes through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + previous_response_id="resp_abc123", + temperature=0.7, + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert result.get("previous_response_id") == "resp_abc123" + + def test_store_background_truncation(self): + """Lifecycle params pass through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + store=True, + background=False, + truncation="auto", + temperature=0.7, + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert result.get("store") is True + assert result.get("background") is False + assert result.get("truncation") == "auto" + + def test_metadata_safety_identifier_user(self): + """Metadata params pass through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + metadata={"request_id": "req_123"}, + safety_identifier="safety_123", + user="user_456", + temperature=0.7, + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert result.get("metadata") == {"request_id": "req_123"} + assert result.get("safety_identifier") == "safety_123" + assert result.get("user") == "user_456" + + def test_all_supported_params_declared(self): + """get_supported_openai_params returns complete list""" + config = PerplexityResponsesConfig() + supported = config.get_supported_openai_params("perplexity/openai/gpt-5.2") + + expected = [ + "max_output_tokens", + "stream", + "temperature", + "top_p", + "tools", + "reasoning", + "preset", + "instructions", + "models", + "tool_choice", + "parallel_tool_calls", + "max_tool_calls", + "text", + "previous_response_id", + "store", + "background", + "truncation", + "metadata", + "safety_identifier", + "user", + "stream_options", + "top_logprobs", + "prompt_cache_key", + "frequency_penalty", + "presence_penalty", + "service_tier", + ] + + for param in expected: + assert param in supported, f"Missing supported param: {param}" + + def test_cost_transformation(self): + """Perplexity cost dict to OpenAI float""" + config = PerplexityResponsesConfig() + + usage_data = { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost": { + "currency": "USD", + "input_cost": 0.0001, + "output_cost": 0.0002, + "total_cost": 0.0003, + }, + } + + result = config._transform_usage(usage_data) + + assert result["input_tokens"] == 100 + assert result["output_tokens"] == 200 + assert result["total_tokens"] == 300 + assert result["cost"] == 0.0003 + + def test_cost_transformation_float_passthrough(self): + """Cost already float passes through""" + config = PerplexityResponsesConfig() + + usage_data = { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost": 0.0005, + } + + result = config._transform_usage(usage_data) + + assert result["cost"] == 0.0005 + + def test_preset_handling(self): + """Preset model names work""" + config = PerplexityResponsesConfig() + + data = config.transform_responses_api_request( + model="preset/pro-search", + input="What is AI?", + response_api_optional_request_params={"temperature": 0.7}, + litellm_params={}, + headers={}, + ) + + assert data["preset"] == "pro-search" + assert data["input"] == "What is AI?" + assert "temperature" in data + + def test_get_complete_url(self): + """Correct endpoint URL""" + config = PerplexityResponsesConfig() + + url = config.get_complete_url(api_base=None, litellm_params={}) + assert url == "https://api.perplexity.ai/v1/responses" + + custom_url = config.get_complete_url( + api_base="https://custom.perplexity.ai", + litellm_params={}, + ) + assert custom_url == "https://custom.perplexity.ai/v1/responses" + + url_with_slash = config.get_complete_url( + api_base="https://api.perplexity.ai/", + litellm_params={}, + ) + assert url_with_slash == "https://api.perplexity.ai/v1/responses" + + def test_perplexity_provider_config_registration(self): + """Test that Perplexity provider returns PerplexityResponsesConfig""" + config = ProviderConfigManager.get_provider_responses_api_config( + model="perplexity/openai/gpt-5.2", + provider=LlmProviders.PERPLEXITY, + ) + + assert config is not None + assert isinstance(config, PerplexityResponsesConfig) + assert config.custom_llm_provider == LlmProviders.PERPLEXITY diff --git a/tests/test_litellm/proxy/test_aiohttp_cleanup_closed.py b/tests/test_litellm/proxy/test_aiohttp_cleanup_closed.py new file mode 100644 index 000000000000..f16b687a2486 --- /dev/null +++ b/tests/test_litellm/proxy/test_aiohttp_cleanup_closed.py @@ -0,0 +1,34 @@ +import asyncio +from unittest.mock import MagicMock, patch + + +def test_initialize_shared_aiohttp_session_sets_enable_cleanup_closed_when_needed( + monkeypatch, +): + from litellm.proxy import proxy_server as proxy_server_module + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + monkeypatch.setattr(proxy_server_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", True) + + with patch("aiohttp.TCPConnector", return_value=connector_mock) as mock_tcp_connector: + with patch("aiohttp.ClientSession", return_value=session_mock): + asyncio.run(proxy_server_module._initialize_shared_aiohttp_session()) + + assert mock_tcp_connector.call_args.kwargs["enable_cleanup_closed"] is True + + +def test_initialize_shared_aiohttp_session_omits_enable_cleanup_closed_when_not_needed( + monkeypatch, +): + from litellm.proxy import proxy_server as proxy_server_module + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + monkeypatch.setattr(proxy_server_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", False) + + with patch("aiohttp.TCPConnector", return_value=connector_mock) as mock_tcp_connector: + with patch("aiohttp.ClientSession", return_value=session_mock): + asyncio.run(proxy_server_module._initialize_shared_aiohttp_session()) + + assert "enable_cleanup_closed" not in mock_tcp_connector.call_args.kwargs diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index 47a60b09af7f..661cdd870994 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -243,6 +243,77 @@ def test_video_generation_cost_calculation_unknown_model(self): custom_llm_provider="openai" ) + def test_video_generation_cost_with_custom_model_info(self): + """Test that custom model_info pricing is applied for video generation. + + When a deployment has custom pricing via model_info, it should be used + instead of looking up the global litellm.model_cost map. + + Related: https://github.com/BerriAI/litellm/issues/21907 + """ + model_info = { + "output_cost_per_video_per_second": 0.05, + } + cost = default_video_cost_calculator( + model="my-custom-video-model", + duration_seconds=10.0, + model_info=model_info, + ) + assert cost == 0.5 + + def test_video_generation_cost_custom_model_info_fallback_to_per_second(self): + """Test that output_cost_per_second is used as fallback when + output_cost_per_video_per_second is not set in custom model_info. + + Related: https://github.com/BerriAI/litellm/issues/21907 + """ + model_info = { + "output_cost_per_second": 0.10, + } + cost = default_video_cost_calculator( + model="my-custom-video-model", + duration_seconds=5.0, + model_info=model_info, + ) + assert cost == 0.5 + + def test_video_generation_cost_custom_pricing_through_completion_cost(self): + """Test that custom video pricing flows through completion_cost via litellm_logging_obj. + + This tests the full cost calculation path: completion_cost extracts model_info + from litellm_logging_obj.litellm_params.metadata.model_info and passes it to + the video cost calculator. + + Related: https://github.com/BerriAI/litellm/issues/21907 + """ + from litellm.cost_calculator import completion_cost + + # Create mock response with usage containing duration_seconds + mock_response = MagicMock() + mock_response.usage = MagicMock() + mock_response.usage.duration_seconds = 10.0 + type(mock_response)._hidden_params = {} + + # Create mock litellm_logging_obj with custom pricing + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_params = { + "metadata": { + "model_info": { + "output_cost_per_video_per_second": 0.05, + } + } + } + + cost = completion_cost( + completion_response=mock_response, + model="openai/hunyuanvideo", + call_type="create_video", + custom_llm_provider="openai", + custom_pricing=True, + litellm_logging_obj=mock_logging_obj, + ) + assert cost == 0.5 + def test_video_generation_with_files(self): """Test video generation with file uploads.""" config = OpenAIVideoConfig()