diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 30887e9f60d..e8130741732 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -11,7 +11,7 @@ import TabItem from '@theme/TabItem'; | Description | Vertex AI is a fully-managed AI development platform for building and using generative AI. | | Provider Route on LiteLLM | `vertex_ai/` | | Link to Provider Doc | [Vertex AI ↗](https://cloud.google.com/vertex-ai) | -| Base URL | [https://{vertex_location}-aiplatform.googleapis.com/](https://{vertex_location}-aiplatform.googleapis.com/) | +| Base URL | 1. Regional endpoints
[https://{vertex_location}-aiplatform.googleapis.com/](https://{vertex_location}-aiplatform.googleapis.com/)
2. Global endpoints (limited availability)
[https://aiplatform.googleapis.com/](https://{aiplatform.googleapis.com/)| | Supported Operations | [`/chat/completions`](#sample-usage), `/completions`, [`/embeddings`](#embedding-models), [`/audio/speech`](#text-to-speech-apis), [`/fine_tuning`](#fine-tuning-apis), [`/batches`](#batch-apis), [`/files`](#batch-apis), [`/images`](#image-generation-models) | @@ -832,7 +832,7 @@ OR You can set: - `vertex_credentials` (str) - can be a json string or filepath to your vertex ai service account.json -- `vertex_location` (str) - place where vertex model is deployed (us-central1, asia-southeast1, etc.) +- `vertex_location` (str) - place where vertex model is deployed (us-central1, asia-southeast1, etc.). Some models support the global location, please see [Vertex AI documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#supported_models) - `vertex_project` Optional[str] - use if vertex project different from the one in vertex_credentials as dynamic params for a `litellm.completion` call. diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 8bef3337587..91ce58162f2 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -196,6 +196,7 @@ class DualCache(BaseCache): key, parent_otel_span: Optional[Span] = None, local_only: bool = False, + redis_only: bool = False, **kwargs, ): # Try to fetch from in-memory cache first @@ -204,7 +205,7 @@ class DualCache(BaseCache): f"async get cache: cache key: {key}; local_only: {local_only}" ) result = None - if self.in_memory_cache is not None: + if self.in_memory_cache is not None and not redis_only: in_memory_result = await self.in_memory_cache.async_get_cache( key, **kwargs ) @@ -213,7 +214,7 @@ class DualCache(BaseCache): if in_memory_result is not None: result = in_memory_result - if result is None and self.redis_cache is not None and local_only is False: + if result is None and self.redis_cache is not None and not local_only: # If not found in in-memory cache, try fetching from Redis redis_result = await self.redis_cache.async_get_cache( key, parent_otel_span=parent_otel_span diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index aa543ee4891..a66b1e755f6 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -833,6 +833,7 @@ class PrometheusLogger(CustomLogger): exception_status=str(getattr(original_exception, "status_code", None)), exception_class=self._get_exception_class_name(original_exception), tags=_tags, + route=user_api_key_dict.request_route, ) _labels = prometheus_label_factory( supported_enum_labels=PrometheusMetricLabels.get_labels( @@ -873,6 +874,7 @@ class PrometheusLogger(CustomLogger): user=user_api_key_dict.user_id, user_email=user_api_key_dict.user_email, status_code="200", + route=user_api_key_dict.request_route, ) _labels = prometheus_label_factory( supported_enum_labels=PrometheusMetricLabels.get_labels( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 785186b8ab8..a0767c09ee2 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -695,10 +695,13 @@ class CustomStreamWrapper: Ensure model id is always the same across all chunks. - If first chunk sent + id set, use that id for all chunks. + If a valid ID is received in any chunk, use it for the response. """ - if self.response_id is None: + if id and isinstance(id, str) and id.strip(): self.response_id = id + elif self.response_id is None: + self.response_id = id + if self.response_id is not None and isinstance(self.response_id, str): model_response.id = self.response_id return model_response diff --git a/litellm/llms/cohere/embed/transformation.py b/litellm/llms/cohere/embed/transformation.py index 03d7edd1262..b5b350a952c 100644 --- a/litellm/llms/cohere/embed/transformation.py +++ b/litellm/llms/cohere/embed/transformation.py @@ -50,7 +50,10 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): ) -> dict: for k, v in non_default_params.items(): if k == "encoding_format": - optional_params["embedding_types"] = v + if isinstance(v, list): + optional_params["embedding_types"] = v + else: + optional_params["embedding_types"] = [v] elif k == "dimensions": optional_params["output_dimension"] = v return optional_params diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index 5cf213950c1..b2e781e5665 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -1,10 +1,14 @@ from typing import Union +from litellm import verbose_logger -import httpx +# dynamic import to allow usage even if httpx is not installed in dev env +try: + import httpx +except ImportError: + httpx = None # type: ignore from litellm.llms.base_llm.chat.transformation import BaseLLMException - class OllamaError(BaseLLMException): def __init__( self, status_code: int, message: str, headers: Union[dict, httpx.Headers] @@ -43,3 +47,81 @@ def _convert_image(image): image_data.convert("RGB").save(jpeg_image, "JPEG") jpeg_image.seek(0) return base64.b64encode(jpeg_image.getvalue()).decode("utf-8") + + +from litellm.llms.base_llm.base_utils import BaseLLMModelInfo + +class OllamaModelInfo(BaseLLMModelInfo): + """ + Dynamic model listing for Ollama server. + Fetches /api/models and /api/tags, then for each tag also /api/models?tag=... + Returns the union of all model names. + """ + @staticmethod + def get_api_key(api_key=None) -> None: + return None # Ollama does not use an API key by default + + @staticmethod + def get_api_base(api_base: str | None = None) -> str: + from litellm.secret_managers.main import get_secret_str + # env var OLLAMA_API_BASE or default + return api_base or get_secret_str("OLLAMA_API_BASE") or "http://localhost:11434" + + def get_models(self, api_key=None, api_base: str | None = None) -> list[str]: + """ + List all models available on the Ollama server via /api/tags endpoint. + """ + import httpx + base = self.get_api_base(api_base) + names: set[str] = set() + try: + resp = httpx.get(f"{base}/api/tags") + resp.raise_for_status() + data = resp.json() + # Expecting a dict with a 'models' list + models_list = [] + if isinstance(data, dict) and 'models' in data and isinstance(data['models'], list): + models_list = data['models'] + elif isinstance(data, list): + models_list = data + # Extract model names + for entry in models_list: + if not isinstance(entry, dict): + continue + nm = entry.get('name') or entry.get('model') + if isinstance(nm, str): + names.add(nm) + except Exception as e: + verbose_logger.warning(f"Error retrieving ollama tag endpoint: {e}") + # If tags endpoint fails, fall back to static list + try: + from litellm import models_by_provider + static = models_by_provider.get("ollama", []) or [] + return [f"ollama/{m}" for m in static] + except Exception as e1: + verbose_logger.warning(f"Error retrieving static ollama models as fallback: {e1}") + return [] + # assemble full model names + result = sorted(names) + return result + def validate_environment( + self, + headers: dict, + model: str, + messages: list, + optional_params: dict, + litellm_params: dict, + api_key=None, + api_base=None, + ) -> dict: + """ + No-op environment validation for Ollama. + """ + return {} + + @staticmethod + def get_base_model(model: str) -> str: + """ + Return the base model name for Ollama (no-op). + """ + return model diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 477995a1578..f96848c6d56 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -84,9 +84,15 @@ def _get_vertex_url( endpoint = "generateContent" if stream is True: endpoint = "streamGenerateContent" - url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}?alt=sse" + if vertex_location== "global": + url = f"https://aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/global/publishers/google/models/{model}:{endpoint}?alt=sse" + else: + url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}?alt=sse" else: - url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}" + if vertex_location == "global": + url = f"https://aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/global/publishers/google/models/{model}:{endpoint}" + else: + url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}" # if model is only numeric chars then it's a fine tuned gemini model # model = 4965075652664360960 diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a1fd7f7366b..a6f1a5883bd 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5870,6 +5870,66 @@ "supported_output_modalities": ["text"], "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "gemini/gemini-2.5-flash-preview-tts": { + "max_tokens": 65535, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 1e-6, + "input_cost_per_token": 0.15e-6, + "output_cost_per_token": 0.6e-6, + "output_cost_per_reasoning_token": 3.5e-6, + "litellm_provider": "gemini", + "mode": "chat", + "rpm": 10, + "tpm": 250000, + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_audio_output": false, + "supports_tool_choice": true, + "supported_endpoints": ["/v1/chat/completions", "/v1/completions"], + "supported_modalities": ["text"], + "supported_output_modalities": ["audio"], + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview" + }, + "gemini/gemini-2.5-flash-preview-05-20": { + "max_tokens": 65535, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 1e-6, + "input_cost_per_token": 0.15e-6, + "output_cost_per_token": 0.6e-6, + "output_cost_per_reasoning_token": 3.5e-6, + "litellm_provider": "gemini", + "mode": "chat", + "rpm": 10, + "tpm": 250000, + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_audio_output": false, + "supports_tool_choice": true, + "supported_endpoints": ["/v1/chat/completions", "/v1/completions"], + "supported_modalities": ["text", "image", "audio", "video"], + "supported_output_modalities": ["text"], + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview" + }, "gemini/gemini-2.5-flash-preview-04-17": { "max_tokens": 65535, "max_input_tokens": 1048576, @@ -5900,6 +5960,34 @@ "supported_output_modalities": ["text"], "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview" }, + "gemini-2.5-flash-preview-05-20": { + "max_tokens": 65535, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 1e-6, + "input_cost_per_token": 0.15e-6, + "output_cost_per_token": 0.6e-6, + "output_cost_per_reasoning_token": 3.5e-6, + "litellm_provider": "vertex_ai-language-models", + "mode": "chat", + "supports_reasoning": true, + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_response_schema": true, + "supports_audio_output": false, + "supports_tool_choice": true, + "supported_endpoints": ["/v1/chat/completions", "/v1/completions", "/v1/batch"], + "supported_modalities": ["text", "image", "audio", "video"], + "supported_output_modalities": ["text"], + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview" + }, "gemini-2.5-flash-preview-04-17": { "max_tokens": 65535, "max_input_tokens": 1048576, @@ -6061,6 +6149,59 @@ "supported_output_modalities": ["text"], "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview" }, + "gemini-2.0-flash-preview-image-generation": { + "max_tokens": 8192, + "max_input_tokens": 1048576, + "max_output_tokens": 8192, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 0.0000007, + "input_cost_per_token": 0.0000001, + "output_cost_per_token": 0.0000004, + "litellm_provider": "vertex_ai-language-models", + "mode": "chat", + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_response_schema": true, + "supports_audio_output": true, + "supports_audio_input": true, + "supported_modalities": ["text", "image", "audio", "video"], + "supported_output_modalities": ["text", "image"], + "supports_tool_choice": true, + "source": "https://ai.google.dev/pricing#2_0flash" + }, + "gemini-2.5-pro-preview-tts": { + "max_tokens": 65535, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 0.0000007, + "input_cost_per_token": 0.00000125, + "input_cost_per_token_above_200k_tokens": 0.0000025, + "output_cost_per_token": 0.00001, + "output_cost_per_token_above_200k_tokens": 0.000015, + "litellm_provider": "vertex_ai-language-models", + "mode": "chat", + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_response_schema": true, + "supports_audio_output": false, + "supports_tool_choice": true, + "supported_modalities": ["text"], + "supported_output_modalities": ["audio"], + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview" + }, "gemini/gemini-2.0-pro-exp-02-05": { "max_tokens": 8192, "max_input_tokens": 2097152, @@ -6099,6 +6240,34 @@ "supports_tool_choice": true, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "gemini/gemini-2.0-flash-preview-image-generation": { + "max_tokens": 8192, + "max_input_tokens": 1048576, + "max_output_tokens": 8192, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 0.0000007, + "input_cost_per_token": 0.0000001, + "output_cost_per_token": 0.0000004, + "litellm_provider": "gemini", + "mode": "chat", + "rpm": 10000, + "tpm": 10000000, + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_response_schema": true, + "supports_audio_output": true, + "supports_audio_input": true, + "supported_modalities": ["text", "image", "audio", "video"], + "supported_output_modalities": ["text", "image"], + "supports_tool_choice": true, + "source": "https://ai.google.dev/pricing#2_0flash" + }, "gemini/gemini-2.0-flash": { "max_tokens": 8192, "max_input_tokens": 1048576, @@ -6180,6 +6349,35 @@ "supported_output_modalities": ["text", "image"], "source": "https://ai.google.dev/pricing#2_0flash" }, + "gemini/gemini-2.5-pro-preview-tts": { + "max_tokens": 65535, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 0.0000007, + "input_cost_per_token": 0.00000125, + "input_cost_per_token_above_200k_tokens": 0.0000025, + "output_cost_per_token": 0.00001, + "output_cost_per_token_above_200k_tokens": 0.000015, + "litellm_provider": "gemini", + "mode": "chat", + "rpm": 10000, + "tpm": 10000000, + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_response_schema": true, + "supports_audio_output": false, + "supports_tool_choice": true, + "supported_modalities": ["text"], + "supported_output_modalities": ["audio"], + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview" + }, "gemini/gemini-2.5-pro-preview-05-06": { "max_tokens": 65535, "max_input_tokens": 1048576, diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html deleted file mode 100644 index 6bd2ec7f07f..00000000000 --- 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 a6fb25f3a79..74f81459b88 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -71,3 +71,6 @@ model_list: model: mistral/* api_key: os.environ/MISTRAL_API_KEY access_groups: ["beta-models"] + +litellm_settings: + cache: true \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0c118898f0f..088d1dd7d4c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -702,7 +702,7 @@ class GenerateKeyRequest(KeyRequestBase): class GenerateKeyResponse(KeyRequestBase): key: str # type: ignore key_name: Optional[str] = None - expires: Optional[datetime] + expires: Optional[datetime] = None user_id: Optional[str] = None token_id: Optional[str] = None litellm_budget_table: Optional[Any] = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e18b358fa63..3c759e839ec 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -11,7 +11,8 @@ Run checks for: import asyncio import re import time -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast +from datetime import datetime +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast from fastapi import Request, status from pydantic import BaseModel @@ -33,6 +34,7 @@ from litellm.proxy._types import ( LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, + LiteLLM_VerificationToken, LiteLLMRoutes, LitellmUserRoles, ProxyErrorTypes, @@ -43,7 +45,12 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.route_llm_request import route_request -from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics +from litellm.proxy.utils import ( + InternalUsageCache, + PrismaClient, + ProxyLogging, + log_db_metrics, +) from litellm.router import Router from litellm.utils import get_utc_datetime @@ -640,6 +647,69 @@ async def _get_fuzzy_user_object( return response +class UserObjectCache: + def __init__( + self, + user_api_key_cache: DualCache, + internal_usage_cache: Optional[InternalUsageCache] = None, + ): + """ + - user_api_key_cache: cache for user api keys + - internal_usage_cache: cache for internal usage (connected to Redis) + """ + self.user_api_key_cache = user_api_key_cache + self.internal_usage_cache = internal_usage_cache + + async def update_user_object( + self, + user_id: str, + user_object: Union[dict, LiteLLM_UserTable], + litellm_parent_otel_span: Optional[Span] = None, + ): + """ + - update user object in cache + """ + if isinstance(user_object, LiteLLM_UserTable): + user_object = user_object.model_dump() + for k, v in user_object.items(): + if isinstance(v, datetime): + user_object[k] = v.isoformat() + await self.user_api_key_cache.async_set_cache(key=user_id, value=user_object) + if self.internal_usage_cache is not None: + await self.internal_usage_cache.async_set_cache( + key=user_id, + value=user_object, + litellm_parent_otel_span=litellm_parent_otel_span, + ) + + async def get_user_object( + self, user_id: str, litellm_parent_otel_span: Optional[Span] = None + ) -> Optional[LiteLLM_UserTable]: + """ + - get user object from cache + """ + cached_obj: Optional[Union[dict, LiteLLM_UserTable]] = None + + ## CHECK REDIS CACHE ## + if self.internal_usage_cache is not None: + cached_obj = await self.internal_usage_cache.async_get_cache( + key=user_id, + litellm_parent_otel_span=litellm_parent_otel_span, + redis_only=True, + ) + + if cached_obj is None: + cached_obj = await self.user_api_key_cache.async_get_cache(key=user_id) + + if cached_obj is not None: + if isinstance(cached_obj, dict): + return LiteLLM_UserTable(**cached_obj) + elif isinstance(cached_obj, LiteLLM_UserTable): + return cached_obj + + return None + + @log_db_metrics async def get_user_object( user_id: Optional[str], @@ -657,18 +727,23 @@ async def get_user_object( - if valid, return LiteLLM_UserTable object with defined limits - if not, then raise an error """ + user_object_cache = UserObjectCache( + user_api_key_cache=user_api_key_cache, + internal_usage_cache=proxy_logging_obj.internal_usage_cache + if proxy_logging_obj is not None + else None, + ) if user_id is None: return None # check if in cache if not check_db_only: - cached_user_obj = await user_api_key_cache.async_get_cache(key=user_id) + cached_user_obj = await user_object_cache.get_user_object( + user_id=user_id, litellm_parent_otel_span=parent_otel_span + ) if cached_user_obj is not None: - if isinstance(cached_user_obj, dict): - return LiteLLM_UserTable(**cached_user_obj) - elif isinstance(cached_user_obj, LiteLLM_UserTable): - return cached_user_obj + return cached_user_obj # else, check db if prisma_client is None: raise Exception("No db connected") @@ -697,8 +772,14 @@ async def get_user_object( if response is None: if user_id_upsert: + new_user_params: Dict[str, Any] = { + "user_id": user_id, + } + if litellm.default_internal_user_params is not None: + new_user_params.update(litellm.default_internal_user_params) + response = await prisma_client.db.litellm_usertable.create( - data={"user_id": user_id}, + data=new_user_params, include={"organization_memberships": True}, ) else: @@ -720,7 +801,9 @@ async def get_user_object( response_dict = _response.model_dump() # save the user object to cache - await user_api_key_cache.async_set_cache(key=user_id, value=response_dict) + await user_object_cache.update_user_object( + user_id=user_id, user_object=response_dict + ) # save to db access time _update_last_db_access_time( @@ -1014,6 +1097,38 @@ class ExperimentalUIJWTToken: ) +async def _get_object_from_cache( + key: str, + proxy_logging_obj: Optional[ProxyLogging], + user_api_key_cache: DualCache, + parent_otel_span: Optional[Span], + base_model: Type[BaseModel], +) -> Optional[BaseModel]: + cached_obj: Optional[Union[dict, BaseModel]] = None + + ## CHECK REDIS CACHE ## + if ( + proxy_logging_obj is not None + and proxy_logging_obj.internal_usage_cache.dual_cache + ): + cached_obj = ( + await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache( + key=key, parent_otel_span=parent_otel_span + ) + ) + + if cached_obj is None: + cached_obj = await user_api_key_cache.async_get_cache(key=key) + + if cached_obj is not None: + if isinstance(cached_obj, dict): + return base_model(**cached_obj) + elif isinstance(cached_obj, base_model): + return cached_obj + + return None + + @log_db_metrics async def get_key_object( hashed_token: str, @@ -1036,15 +1151,16 @@ async def get_key_object( # check if in cache key = hashed_token - cached_key_obj: Optional[UserAPIKeyAuth] = await user_api_key_cache.async_get_cache( - key=key + cached_key_obj = await _get_object_from_cache( + key=key, + proxy_logging_obj=proxy_logging_obj, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + base_model=LiteLLM_VerificationToken, ) if cached_key_obj is not None: - if isinstance(cached_key_obj, dict): - return UserAPIKeyAuth(**cached_key_obj) - elif isinstance(cached_key_obj, UserAPIKeyAuth): - return cached_key_obj + return UserAPIKeyAuth(**cached_key_obj.model_dump(exclude_none=True)) if check_cache_only: raise Exception( diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim.py index 0409068afc4..66a0aff6cb0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim.py @@ -14,8 +14,8 @@ from pydantic import BaseModel from websockets.asyncio.client import ClientConnection, connect from litellm import DualCache -from litellm._version import version as litellm_version from litellm._logging import verbose_proxy_logger +from litellm._version import version as litellm_version from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -125,17 +125,16 @@ class AimGuardrail(CustomGuardrail): if action_type is None: verbose_proxy_logger.debug("Aim: No required action specified") return data - match action_type: - case "monitor_action": - verbose_proxy_logger.info("Aim: monitor action") - case "block_action": - self._handle_block_action(res["analysis_result"], required_action) - case "anonymize_action": - return self._anonymize_request( - res["analysis_result"], required_action, data - ) - case _: - verbose_proxy_logger.error(f"Aim: {action_type} action") + if action_type == "monitor_action": + verbose_proxy_logger.info("Aim: monitor action") + elif action_type == "block_action": + self._handle_block_action(res["analysis_result"], required_action) + elif action_type == "anonymize_action": + return self._anonymize_request( + res["analysis_result"], required_action, data + ) + else: + verbose_proxy_logger.error(f"Aim: {action_type} action") return data def _handle_block_action(self, analysis_result: Any, required_action: Any) -> None: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index c00d282a810..bbdd8c81335 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -23,6 +23,7 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Request, status import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import * +from litellm.proxy.auth.auth_checks import UserObjectCache from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity @@ -54,9 +55,9 @@ def _update_internal_new_user_params(data_json: dict, data: NewUserRequest) -> d data_json["user_id"] = str(uuid.uuid4()) auto_create_key = data_json.pop("auto_create_key", True) if auto_create_key is False: - data_json["table_name"] = ( - "user" # only create a user, don't create key if 'auto_create_key' set to False - ) + data_json[ + "table_name" + ] = "user" # only create a user, don't create key if 'auto_create_key' set to False is_internal_user = False if data.user_role and data.user_role.is_internal_user_role: @@ -238,23 +239,15 @@ async def new_user( else: raise e - new_user_response = NewUserResponse( - key=response.get("token", ""), - expires=response.get("expires", None), - max_budget=response["max_budget"], - user_id=response["user_id"], - user_role=response.get("user_role", None), - user_email=response.get("user_email", None), - user_alias=response.get("user_alias", None), - teams=response.get("teams", None), - team_id=response.get("team_id", None), - metadata=response.get("metadata", None), - models=response.get("models", None), - tpm_limit=response.get("tpm_limit", None), - rpm_limit=response.get("rpm_limit", None), - budget_duration=response.get("budget_duration", None), - model_max_budget=response.get("model_max_budget", None), - ) + special_keys = ["token", "token_id"] + response_dict = {} + for key, value in response.items(): + if key in NewUserResponse.model_fields.keys() and key not in special_keys: + response_dict[key] = value + + response_dict["key"] = response.get("token", "") + + new_user_response = NewUserResponse(**response_dict) ######################################################### ########## USER CREATED HOOK ################ @@ -555,7 +548,6 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest) -> di not in ( [], {}, - 0, ) and k not in LiteLLM_ManagementEndpoint_MetadataFields ): # models default to [], spend defaults to 0, we should not reset these values @@ -582,9 +574,9 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest) -> di "budget_duration" not in non_default_values ): # applies internal user limits, if user role updated if is_internal_user and litellm.internal_user_budget_duration is not None: - non_default_values["budget_duration"] = ( - litellm.internal_user_budget_duration - ) + non_default_values[ + "budget_duration" + ] = litellm.internal_user_budget_duration from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time non_default_values["budget_reset_at"] = get_budget_reset_time( @@ -650,10 +642,15 @@ async def user_update( """ - from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + from litellm.proxy.proxy_server import ( + litellm_proxy_admin_name, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) try: - data_json: dict = data.json() + data_json: dict = data.model_dump(exclude_unset=True) # get the row from db if prisma_client is None: raise Exception("Not connected to DB!") @@ -734,6 +731,16 @@ async def user_update( user_row_litellm_typed = LiteLLM_UserTable( **user_row.model_dump(exclude_none=True) ) + + ## UPDATE CACHE ## + user_object_cache = UserObjectCache( + user_api_key_cache=user_api_key_cache, + internal_usage_cache=proxy_logging_obj.internal_usage_cache, + ) + await user_object_cache.update_user_object( + user_id=response["user_id"], user_object=user_row_litellm_typed + ) + asyncio.create_task( UserManagementEventHooks.create_internal_user_audit_log( user_id=user_row_litellm_typed.user_id, @@ -1227,13 +1234,13 @@ async def ui_view_users( } # Query users with pagination and filters - users: Optional[List[BaseModel]] = ( - await prisma_client.db.litellm_usertable.find_many( - where=where_conditions, - skip=skip, - take=page_size, - order={"created_at": "desc"}, - ) + users: Optional[ + List[BaseModel] + ] = await prisma_client.db.litellm_usertable.find_many( + where=where_conditions, + skip=skip, + take=page_size, + order={"created_at": "desc"}, ) if not users: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d1e725def55..64eb0a5cb46 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -146,12 +146,14 @@ class InternalUsageCache: key, litellm_parent_otel_span: Union[Span, None], local_only: bool = False, + redis_only: bool = False, **kwargs, ) -> Any: return await self.dual_cache.async_get_cache( key=key, local_only=local_only, parent_otel_span=litellm_parent_otel_span, + redis_only=redis_only, **kwargs, ) diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 8fdcce4cbb8..0155a0630c0 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -70,6 +70,7 @@ class UserAPIKeyLabelNames(Enum): EXCEPTION_CLASS = EXCEPTION_CLASS STATUS_CODE = "status_code" FALLBACK_MODEL = "fallback_model" + ROUTE = "route" DEFINED_PROMETHEUS_METRICS = Literal[ @@ -125,6 +126,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.STATUS_CODE.value, UserAPIKeyLabelNames.USER_EMAIL.value, + UserAPIKeyLabelNames.ROUTE.value, ] litellm_proxy_failed_requests_metric = [ @@ -137,6 +139,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.EXCEPTION_STATUS.value, UserAPIKeyLabelNames.EXCEPTION_CLASS.value, + UserAPIKeyLabelNames.ROUTE.value, ] litellm_deployment_latency_per_output_token = [ @@ -292,3 +295,6 @@ class UserAPIKeyLabelValues(BaseModel): fallback_model: Annotated[ Optional[str], Field(..., alias=UserAPIKeyLabelNames.FALLBACK_MODEL.value) ] = None + route: Annotated[ + Optional[str], Field(..., alias=UserAPIKeyLabelNames.ROUTE.value) + ] = None diff --git a/litellm/utils.py b/litellm/utils.py index 5d6a64cd345..0dd4c8cd1e3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5851,7 +5851,7 @@ def _get_valid_models_from_provider_api( _model_cache.set_cached_model_info(custom_llm_provider, litellm_params, models) return models except Exception as e: - verbose_logger.debug(f"Error getting valid models: {e}") + verbose_logger.warning(f"Error getting valid models: {e}") return [] @@ -5916,7 +5916,7 @@ def get_valid_models( return valid_models except Exception as e: - verbose_logger.debug(f"Error getting valid models: {e}") + verbose_logger.warning(f"Error getting valid models: {e}") return [] # NON-Blocking @@ -6599,6 +6599,10 @@ class ProviderConfigManager: return litellm.AnthropicModelInfo() elif LlmProviders.XAI == provider: return litellm.XAIModelInfo() + elif LlmProviders.OLLAMA == provider or LlmProviders.OLLAMA_CHAT == provider: + # Dynamic model listing for Ollama server + from litellm.llms.ollama.common_utils import OllamaModelInfo + return OllamaModelInfo() elif LlmProviders.VLLM == provider: from litellm.llms.vllm.common_utils import ( VLLMModelInfo, # experimental approach, to reduce bloat on __init__.py diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a1fd7f7366b..a6f1a5883bd 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5870,6 +5870,66 @@ "supported_output_modalities": ["text"], "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "gemini/gemini-2.5-flash-preview-tts": { + "max_tokens": 65535, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 1e-6, + "input_cost_per_token": 0.15e-6, + "output_cost_per_token": 0.6e-6, + "output_cost_per_reasoning_token": 3.5e-6, + "litellm_provider": "gemini", + "mode": "chat", + "rpm": 10, + "tpm": 250000, + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_audio_output": false, + "supports_tool_choice": true, + "supported_endpoints": ["/v1/chat/completions", "/v1/completions"], + "supported_modalities": ["text"], + "supported_output_modalities": ["audio"], + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview" + }, + "gemini/gemini-2.5-flash-preview-05-20": { + "max_tokens": 65535, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 1e-6, + "input_cost_per_token": 0.15e-6, + "output_cost_per_token": 0.6e-6, + "output_cost_per_reasoning_token": 3.5e-6, + "litellm_provider": "gemini", + "mode": "chat", + "rpm": 10, + "tpm": 250000, + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_audio_output": false, + "supports_tool_choice": true, + "supported_endpoints": ["/v1/chat/completions", "/v1/completions"], + "supported_modalities": ["text", "image", "audio", "video"], + "supported_output_modalities": ["text"], + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview" + }, "gemini/gemini-2.5-flash-preview-04-17": { "max_tokens": 65535, "max_input_tokens": 1048576, @@ -5900,6 +5960,34 @@ "supported_output_modalities": ["text"], "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview" }, + "gemini-2.5-flash-preview-05-20": { + "max_tokens": 65535, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 1e-6, + "input_cost_per_token": 0.15e-6, + "output_cost_per_token": 0.6e-6, + "output_cost_per_reasoning_token": 3.5e-6, + "litellm_provider": "vertex_ai-language-models", + "mode": "chat", + "supports_reasoning": true, + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_response_schema": true, + "supports_audio_output": false, + "supports_tool_choice": true, + "supported_endpoints": ["/v1/chat/completions", "/v1/completions", "/v1/batch"], + "supported_modalities": ["text", "image", "audio", "video"], + "supported_output_modalities": ["text"], + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview" + }, "gemini-2.5-flash-preview-04-17": { "max_tokens": 65535, "max_input_tokens": 1048576, @@ -6061,6 +6149,59 @@ "supported_output_modalities": ["text"], "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview" }, + "gemini-2.0-flash-preview-image-generation": { + "max_tokens": 8192, + "max_input_tokens": 1048576, + "max_output_tokens": 8192, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 0.0000007, + "input_cost_per_token": 0.0000001, + "output_cost_per_token": 0.0000004, + "litellm_provider": "vertex_ai-language-models", + "mode": "chat", + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_response_schema": true, + "supports_audio_output": true, + "supports_audio_input": true, + "supported_modalities": ["text", "image", "audio", "video"], + "supported_output_modalities": ["text", "image"], + "supports_tool_choice": true, + "source": "https://ai.google.dev/pricing#2_0flash" + }, + "gemini-2.5-pro-preview-tts": { + "max_tokens": 65535, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 0.0000007, + "input_cost_per_token": 0.00000125, + "input_cost_per_token_above_200k_tokens": 0.0000025, + "output_cost_per_token": 0.00001, + "output_cost_per_token_above_200k_tokens": 0.000015, + "litellm_provider": "vertex_ai-language-models", + "mode": "chat", + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_response_schema": true, + "supports_audio_output": false, + "supports_tool_choice": true, + "supported_modalities": ["text"], + "supported_output_modalities": ["audio"], + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview" + }, "gemini/gemini-2.0-pro-exp-02-05": { "max_tokens": 8192, "max_input_tokens": 2097152, @@ -6099,6 +6240,34 @@ "supports_tool_choice": true, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "gemini/gemini-2.0-flash-preview-image-generation": { + "max_tokens": 8192, + "max_input_tokens": 1048576, + "max_output_tokens": 8192, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 0.0000007, + "input_cost_per_token": 0.0000001, + "output_cost_per_token": 0.0000004, + "litellm_provider": "gemini", + "mode": "chat", + "rpm": 10000, + "tpm": 10000000, + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_response_schema": true, + "supports_audio_output": true, + "supports_audio_input": true, + "supported_modalities": ["text", "image", "audio", "video"], + "supported_output_modalities": ["text", "image"], + "supports_tool_choice": true, + "source": "https://ai.google.dev/pricing#2_0flash" + }, "gemini/gemini-2.0-flash": { "max_tokens": 8192, "max_input_tokens": 1048576, @@ -6180,6 +6349,35 @@ "supported_output_modalities": ["text", "image"], "source": "https://ai.google.dev/pricing#2_0flash" }, + "gemini/gemini-2.5-pro-preview-tts": { + "max_tokens": 65535, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_audio_token": 0.0000007, + "input_cost_per_token": 0.00000125, + "input_cost_per_token_above_200k_tokens": 0.0000025, + "output_cost_per_token": 0.00001, + "output_cost_per_token_above_200k_tokens": 0.000015, + "litellm_provider": "gemini", + "mode": "chat", + "rpm": 10000, + "tpm": 10000000, + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_response_schema": true, + "supports_audio_output": false, + "supports_tool_choice": true, + "supported_modalities": ["text"], + "supported_output_modalities": ["audio"], + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview" + }, "gemini/gemini-2.5-pro-preview-05-06": { "max_tokens": 65535, "max_input_tokens": 1048576, diff --git a/pyproject.toml b/pyproject.toml index 30d9182a941..2607937e7b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.70.2" +version = "1.70.3" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT" @@ -136,7 +136,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.70.2" +version = "1.70.3" version_files = [ "pyproject.toml:^version" ] diff --git a/tests/litellm/litellm_core_utils/test_streaming_handler.py b/tests/litellm/litellm_core_utils/test_streaming_handler.py index 81bde88f392..c38709735f0 100644 --- a/tests/litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/litellm/litellm_core_utils/test_streaming_handler.py @@ -612,6 +612,33 @@ def test_streaming_handler_with_stop_chunk( assert returned_chunk is None +def test_set_response_id_propagation_empty_to_valid(initialized_custom_stream_wrapper: CustomStreamWrapper): + """Test that response_id is properly set when first chunk has empty ID and second chunk has valid ID""" + + model_response1 = ModelResponseStream(id="", created=1742056047, model=None) + model_response1 = initialized_custom_stream_wrapper.set_model_id(model_response1.id, model_response1) + assert model_response1.id == "" + + model_response2 = ModelResponseStream(id="valid-id-123", created=1742056048, model=None) + model_response2 = initialized_custom_stream_wrapper.set_model_id("valid-id-123", model_response2) + assert model_response2.id == "valid-id-123" + assert initialized_custom_stream_wrapper.response_id == "valid-id-123" + + +def test_set_response_id_propagation_valid_to_invalid(initialized_custom_stream_wrapper: CustomStreamWrapper): + """Test that response_id is maintained when first chunk has valid ID and second chunk has invalid ID""" + + model_response1 = ModelResponseStream(id="first-valid-id", created=1742056049, model=None) + model_response1 = initialized_custom_stream_wrapper.set_model_id("first-valid-id", model_response1) + assert model_response1.id == "first-valid-id" + assert initialized_custom_stream_wrapper.response_id == "first-valid-id" + + model_response2 = ModelResponseStream(id="", created=1742056050, model=None) + model_response2 = initialized_custom_stream_wrapper.set_model_id("", model_response2) + assert model_response2.id == "first-valid-id" + assert initialized_custom_stream_wrapper.response_id == "first-valid-id" + + @pytest.mark.asyncio async def test_streaming_completion_start_time(logging_obj: Logging): """Test that the start time is set correctly""" diff --git a/tests/litellm/llms/ollama/test_ollama_model_info.py b/tests/litellm/llms/ollama/test_ollama_model_info.py new file mode 100644 index 00000000000..f5fea572ac9 --- /dev/null +++ b/tests/litellm/llms/ollama/test_ollama_model_info.py @@ -0,0 +1,106 @@ +import os +import sys +import json +import uuid +import pytest +from unittest.mock import MagicMock, patch + + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +""" +Unit tests for OllamaModelInfo.get_models functionality. +""" +# Ensure a dummy httpx module is available for import in tests +import sys, types +# Provide a dummy httpx module for import in get_models +if 'httpx' not in sys.modules: + # Create a minimal module with HTTPStatusError + httpx_mod = types.ModuleType('httpx') + httpx_mod.HTTPStatusError = Exception + sys.modules['httpx'] = httpx_mod + +import httpx + +from litellm.llms.ollama.common_utils import OllamaModelInfo + + +class DummyResponse: + """ + A dummy response object to simulate httpx responses. + """ + def __init__(self, json_data, status_code=200): + self._json = json_data + self.status_code = status_code + + def raise_for_status(self): + if self.status_code >= 400: + # Simulate an HTTP status error + raise httpx.HTTPStatusError("Error status code", request=None, response=None) + + def json(self): + return self._json + + +class TestOllamaModelInfo: + def test_get_models_from_dict_response(self, monkeypatch): + """ + When the /api/tags endpoint returns a dict with a 'models' list, + get_models should extract and return sorted unique model names. + """ + calls = [] + sample = {'models': [ + {'name': 'zeta'}, + {'model': 'alpha'}, + {'name': 123}, # non-str should be ignored + 'invalid', # non-dict should be ignored + ]} + + def mock_get(url): + calls.append(url) + return DummyResponse(sample, status_code=200) + + monkeypatch.setattr(httpx, 'get', mock_get) + info = OllamaModelInfo() + models = info.get_models() + # Only 'alpha' and 'zeta' should be returned, sorted alphabetically + assert models == ['alpha', 'zeta'] + # Ensure correct endpoint was called + assert calls and calls[0].endswith('/api/tags') + + + def test_get_models_from_list_response(self, monkeypatch): + """ + When the /api/tags endpoint returns a list of dicts, + get_models should extract and return sorted unique model names. + """ + sample = [ + {'name': 'm1'}, + {'model': 'm2'}, + {}, # no name/model key should be ignored + ] + + def mock_get(url): + return DummyResponse(sample, status_code=200) + + monkeypatch.setattr(httpx, 'get', mock_get) + info = OllamaModelInfo() + models = info.get_models() + assert models == ['m1', 'm2'] + + + def test_get_models_fallback_on_error(self, monkeypatch): + """ + If the httpx.get call raises an exception, get_models should + fall back to the static models_by_provider list prefixed by 'ollama/'. + """ + def mock_get(url): + raise Exception("connection failure") + + monkeypatch.setattr(httpx, 'get', mock_get) + info = OllamaModelInfo() + models = info.get_models() + # Default static ollama_models is ['llama2'], so expect ['ollama/llama2'] + assert models == ['ollama/llama2'] \ No newline at end of file diff --git a/tests/litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index 90a7fb30e19..fca89650a53 100644 --- a/tests/litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -17,6 +17,7 @@ from litellm.llms.vertex_ai.common_utils import ( get_vertex_location_from_url, get_vertex_project_id_from_url, set_schema_property_ordering, + _get_vertex_url ) @@ -516,3 +517,45 @@ def test_vertex_ai_complex_response_schema(): assert "additionalProperties" not in type2 assert "additionalProperties" not in type3 assert "additionalProperties" not in type3_prop3_items + +@pytest.mark.parametrize( + "stream, expected_endpoint_suffix", + [ + (True, "streamGenerateContent?alt=sse"), + (False, "generateContent"), + ], +) +def test_get_vertex_url_global_region(stream, expected_endpoint_suffix): + """ + Test _get_vertex_url when vertex_location is 'global' for chat mode. + """ + mode = "chat" + model = "gemini-1.5-pro-preview-0409" + vertex_project = "test-g-project" + vertex_location = "global" + vertex_api_version = "v1" + + # Mock litellm.VertexGeminiConfig.get_model_for_vertex_ai_url to return model as is + # as we are not testing that part here, just the URL construction + with patch("litellm.VertexGeminiConfig.get_model_for_vertex_ai_url", side_effect=lambda model: model): + url, endpoint = _get_vertex_url( + mode=mode, + model=model, + stream=stream, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_api_version=vertex_api_version, + ) + + expected_url_base = f"https://aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/global/publishers/google/models/{model}" + + if stream: + expected_endpoint = "streamGenerateContent" + expected_url = f"{expected_url_base}:{expected_endpoint}?alt=sse" + else: + expected_endpoint = "generateContent" + expected_url = f"{expected_url_base}:{expected_endpoint}" + + + assert endpoint == expected_endpoint + assert url == expected_url diff --git a/tests/litellm/proxy/auth/test_auth_checks.py b/tests/litellm/proxy/auth/test_auth_checks.py index 24e8506e971..aa60952add4 100644 --- a/tests/litellm/proxy/auth/test_auth_checks.py +++ b/tests/litellm/proxy/auth/test_auth_checks.py @@ -13,12 +13,14 @@ from datetime import datetime, timedelta import pytest import litellm +from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ( LiteLLM_UserTable, + LiteLLM_VerificationToken, LitellmUserRoles, SSOUserDefinedValues, ) -from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.utils import get_utc_datetime @@ -59,6 +61,8 @@ def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_v # Decrypt and verify token contents decrypted_token = decrypt_value_helper(token, exception_type="debug") + # Check that decrypted_token is not None before using json.loads + assert decrypted_token is not None token_data = json.loads(decrypted_token) assert token_data["user_id"] == "test_user" @@ -110,3 +114,117 @@ def test_get_key_object_from_ui_hash_key_invalid(): # Test with invalid token key_object = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key("invalid_token") assert key_object is None + + +@pytest.mark.asyncio +async def test_get_object_from_cache_redis(): + """Test that _get_object_from_cache retrieves from Redis cache when available""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_VerificationToken + from litellm.proxy.auth.auth_checks import _get_object_from_cache + + # Create mock objects + mock_proxy_logging = MagicMock() + mock_proxy_logging.internal_usage_cache.dual_cache = AsyncMock() + mock_user_api_key_cache = DualCache() + + # Create test data + test_key = "test_key" + test_data = { + "token": "test_token", + "key_name": "test_key_name", + "spend": 0.0, + "models": ["gpt-3.5-turbo"], + } + + # Mock Redis cache response + mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache.return_value = ( + test_data + ) + + # Call the function + result = await _get_object_from_cache( + key=test_key, + proxy_logging_obj=mock_proxy_logging, + user_api_key_cache=mock_user_api_key_cache, + parent_otel_span=None, + base_model=LiteLLM_VerificationToken, + ) + + # Verify Redis cache was checked + mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache.assert_called_once_with( + key=test_key, parent_otel_span=None + ) + + # Verify result is correct + assert isinstance(result, LiteLLM_VerificationToken) + assert result.token == "test_token" + assert result.key_name == "test_key_name" + assert result.spend == 0.0 + assert result.models == ["gpt-3.5-turbo"] + +async def test_default_internal_user_params_with_get_user_object(monkeypatch): + """Test that default_internal_user_params is used when creating a new user via get_user_object""" + # Set up default_internal_user_params + default_params = { + "models": ["gpt-4", "claude-3-opus"], + "max_budget": 200.0, + "user_role": "internal_user", + } + monkeypatch.setattr(litellm, "default_internal_user_params", default_params) + + # Mock the necessary dependencies + mock_prisma_client = MagicMock() + mock_db = AsyncMock() + mock_prisma_client.db = mock_db + + # Set up the user creation mock - create a complete user model that can be converted to a dict + mock_user = MagicMock() + mock_user.user_id = "new_test_user" + mock_user.models = ["gpt-4", "claude-3-opus"] + mock_user.max_budget = 200.0 + mock_user.user_role = "internal_user" + mock_user.organization_memberships = [] + + # Make the mock model_dump or dict method return appropriate data + mock_user.dict = lambda: { + "user_id": "new_test_user", + "models": ["gpt-4", "claude-3-opus"], + "max_budget": 200.0, + "user_role": "internal_user", + "organization_memberships": [], + } + + # Setup the mock returns + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.create = AsyncMock(return_value=mock_user) + + # Create a mock cache - use AsyncMock for async methods + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + # Call get_user_object with user_id_upsert=True to trigger user creation + try: + user_obj = await get_user_object( + user_id="new_test_user", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + user_id_upsert=True, + proxy_logging_obj=None, + ) + except Exception as e: + # this fails since the mock object is a MagicMock and not a LiteLLM_UserTable + print(e) + + # Verify the user was created with the default params + mock_prisma_client.db.litellm_usertable.create.assert_called_once() + creation_args = mock_prisma_client.db.litellm_usertable.create.call_args[1]["data"] + + # Verify defaults were applied to the creation args + assert "models" in creation_args + assert creation_args["models"] == ["gpt-4", "claude-3-opus"] + assert creation_args["max_budget"] == 200.0 + assert creation_args["user_role"] == "internal_user" diff --git a/tests/litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 360f21f1717..53653a2933c 100644 --- a/tests/litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -10,9 +10,14 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path -from litellm.proxy._types import LiteLLM_UserTableFiltered, UserAPIKeyAuth +from litellm.proxy._types import ( + LiteLLM_UserTableFiltered, + UpdateUserRequest, + UserAPIKeyAuth, +) from litellm.proxy.management_endpoints.internal_user_endpoints import ( LiteLLM_UserTableWithKeyCount, + _update_internal_user_params, get_user_key_counts, get_users, ui_view_users, @@ -169,3 +174,66 @@ def test_validate_sort_params(): assert _validate_sort_params("user_id", "desc") == {"user_id": "desc"} with pytest.raises(Exception): _validate_sort_params("user_id", "invalid") + + +def test_update_user_request_pydantic_object(): + """ + Test that _update_internal_user_params correctly processes an email-only update + """ + data = UpdateUserRequest(user_email="test@example.com") + + data_json = data.model_dump(exclude_unset=True) + + assert data_json == {"user_email": "test@example.com"} + + +def test_update_internal_user_params_email(): + """ + Test that _update_internal_user_params correctly processes an email-only update + """ + from litellm.proxy._types import UpdateUserRequest + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_internal_user_params, + ) + + # Create test data with only email update + data_json = {"user_email": "test@example.com"} + data = UpdateUserRequest(user_email="test@example.com") + + # Call the function + non_default_values = _update_internal_user_params(data_json=data_json, data=data) + + # Assertions + assert len(non_default_values) == 1 # Should only contain email + assert "user_email" in non_default_values + assert non_default_values["user_email"] == "test@example.com" + assert "user_id" not in non_default_values # Should not add user_id if not provided + assert "max_budget" not in non_default_values # Should not add default values + assert "budget_duration" not in non_default_values # Should not add default values + + +def test_update_internal_user_params_reset_spend_and_max_budget(): + """ + Relevant Issue: https://github.com/BerriAI/litellm/issues/10495 + """ + from litellm.proxy._types import UpdateUserRequest + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_internal_user_params, + ) + + # Create test data with only email update + data = UpdateUserRequest(spend=0, max_budget=0, user_id="test_user_id") + data_json = data.model_dump(exclude_unset=True) + + # Call the function + non_default_values = _update_internal_user_params(data_json=data_json, data=data) + + # Assertions + assert len(non_default_values) == 3 # Should only contain email + assert "spend" in non_default_values + assert non_default_values["spend"] == 0 + assert "max_budget" in non_default_values + assert non_default_values["max_budget"] == 0 + assert "user_id" in non_default_values # Should not add user_id if not provided + assert non_default_values["user_id"] == "test_user_id" + assert "budget_duration" not in non_default_values # Should not add default values diff --git a/tests/llm_translation/test_optional_params.py b/tests/llm_translation/test_optional_params.py index d1d10d50643..b29d20ba408 100644 --- a/tests/llm_translation/test_optional_params.py +++ b/tests/llm_translation/test_optional_params.py @@ -1459,3 +1459,11 @@ def test_azure_response_format(monkeypatch): response_format={"type": "json_object"}, ) assert optional_params["response_format"] == {"type": "json_object"} + +def test_cohere_embed_dimensions_param(): + optional_params = get_optional_params_embeddings( + model="embed-multilingual-v3.0", + custom_llm_provider="cohere", + encoding_format="float", + ) + assert optional_params["embedding_types"] == ["float"] \ No newline at end of file diff --git a/tests/logging_callback_tests/test_prometheus_unit_tests.py b/tests/logging_callback_tests/test_prometheus_unit_tests.py index 0b58bc7aafb..007de7e337e 100644 --- a/tests/logging_callback_tests/test_prometheus_unit_tests.py +++ b/tests/logging_callback_tests/test_prometheus_unit_tests.py @@ -694,6 +694,7 @@ async def test_async_post_call_failure_hook(prometheus_logger): team_alias="test_team_alias", user_id="test_user", end_user_id="test_end_user", + request_route="/chat/completions", ) # Call the function @@ -714,6 +715,7 @@ async def test_async_post_call_failure_hook(prometheus_logger): user="test_user", exception_status="429", exception_class="Openai.RateLimitError", + route=user_api_key_dict.request_route, ) prometheus_logger.litellm_proxy_failed_requests_metric.labels().inc.assert_called_once() @@ -728,6 +730,7 @@ async def test_async_post_call_failure_hook(prometheus_logger): user="test_user", status_code="429", user_email=None, + route=user_api_key_dict.request_route, ) prometheus_logger.litellm_proxy_total_requests_metric.labels().inc.assert_called_once() @@ -752,6 +755,7 @@ async def test_async_post_call_success_hook(prometheus_logger): team_alias="test_team_alias", user_id="test_user", end_user_id="test_end_user", + request_route="/chat/completions", ) response = {"choices": [{"message": {"content": "test response"}}]} @@ -772,6 +776,7 @@ async def test_async_post_call_success_hook(prometheus_logger): user="test_user", status_code="200", user_email=None, + route=user_api_key_dict.request_route, ) prometheus_logger.litellm_proxy_total_requests_metric.labels().inc.assert_called_once() diff --git a/tests/otel_tests/test_prometheus.py b/tests/otel_tests/test_prometheus.py index 7b80e23fb9b..95ecddc15d3 100644 --- a/tests/otel_tests/test_prometheus.py +++ b/tests/otel_tests/test_prometheus.py @@ -106,7 +106,7 @@ async def test_proxy_failure_metrics(): print("/metrics", metrics) # Check if the failure metric is present and correct - expected_metric = 'litellm_proxy_failed_requests_metric_total{api_key_alias="None",end_user="None",exception_class="Openai.RateLimitError",exception_status="429",hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",requested_model="fake-azure-endpoint",team="None",team_alias="None",user="default_user_id"} 1.0' + expected_metric = 'litellm_proxy_failed_requests_metric_total{api_key_alias="None",end_user="None",exception_class="Openai.RateLimitError",exception_status="429",hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",requested_model="fake-azure-endpoint",route="/chat/completions",team="None",team_alias="None",user="default_user_id"} 1.0' assert ( expected_metric in metrics @@ -115,7 +115,7 @@ async def test_proxy_failure_metrics(): assert expected_llm_deployment_failure assert ( - 'litellm_proxy_total_requests_metric_total{api_key_alias="None",end_user="None",hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",requested_model="fake-azure-endpoint",status_code="429",team="None",team_alias="None",user="default_user_id",user_email="None"} 1.0' + 'litellm_proxy_total_requests_metric_total{api_key_alias="None",end_user="None",hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",requested_model="fake-azure-endpoint",route="/chat/completions",status_code="429",team="None",team_alias="None",user="default_user_id",user_email="None"} 1.0' in metrics ) diff --git a/ui/litellm-dashboard/src/components/teams.tsx b/ui/litellm-dashboard/src/components/teams.tsx index 6c6b4dd5b58..5e768f73f26 100644 --- a/ui/litellm-dashboard/src/components/teams.tsx +++ b/ui/litellm-dashboard/src/components/teams.tsx @@ -665,7 +665,7 @@ const Teams: React.FC = ({ overflow: "hidden", }} > - {team["spend"]} + {Number(team["spend"]).toFixed(4)}