diff --git a/README.md b/README.md index 4f1a86a648f..9fed1c6dbc7 100644 --- a/README.md +++ b/README.md @@ -274,8 +274,6 @@ echo 'LITELLM_MASTER_KEY="sk-1234"' > .env # password generator to get a random hash for litellm salt key echo 'LITELLM_SALT_KEY="sk-1234"' >> .env -source .env - # Start docker compose up ``` diff --git a/docs/my-website/docs/providers/gemini.md b/docs/my-website/docs/providers/gemini.md index 1b21ed8d03c..f33c492f182 100644 --- a/docs/my-website/docs/providers/gemini.md +++ b/docs/my-website/docs/providers/gemini.md @@ -2006,3 +2006,34 @@ curl -L -X POST 'http://localhost:4000/v1/chat/completions' \ +### Image Generation Pricing + +Gemini image generation models (like `gemini-3-pro-image-preview`) return `image_tokens` in the response usage. These tokens are priced differently from text tokens: + +| Token Type | Price per 1M tokens | Price per token | +|------------|---------------------|-----------------| +| Text output | $12 | $0.000012 | +| Image output | $120 | $0.00012 | + +The number of image tokens depends on the output resolution: + +| Resolution | Tokens per image | Cost per image | +|------------|------------------|----------------| +| 1K-2K (1024x1024 to 2048x2048) | 1,120 | $0.134 | +| 4K (4096x4096) | 2,000 | $0.24 | + +LiteLLM automatically calculates costs using `output_cost_per_image_token` from the model pricing configuration. + +**Example response usage:** +```json +{ + "completion_tokens_details": { + "reasoning_tokens": 225, + "text_tokens": 0, + "image_tokens": 1120 + } +} +``` + +For more details, see [Google's Gemini pricing documentation](https://ai.google.dev/gemini-api/docs/pricing). + diff --git a/docs/my-website/docs/proxy/deploy.md b/docs/my-website/docs/proxy/deploy.md index e40d7acc7c8..0f0e5f678d3 100644 --- a/docs/my-website/docs/proxy/deploy.md +++ b/docs/my-website/docs/proxy/deploy.md @@ -26,8 +26,6 @@ echo 'LITELLM_MASTER_KEY="sk-1234"' > .env # password generator to get a random hash for litellm salt key echo 'LITELLM_SALT_KEY="sk-1234"' >> .env -source .env - # Start docker compose up ``` @@ -1072,4 +1070,4 @@ A: We explored MySQL but that was hard to maintain and led to bugs for customers **Q: If there is Postgres downtime, how does LiteLLM react? Does it fail-open or is there API downtime?** -A: You can gracefully handle DB unavailability if it's on your VPC. See our production guide for more details: [Gracefully Handle DB Unavailability](https://docs.litellm.ai/docs/proxy/prod#6-if-running-litellm-on-vpc-gracefully-handle-db-unavailability) \ No newline at end of file +A: You can gracefully handle DB unavailability if it's on your VPC. See our production guide for more details: [Gracefully Handle DB Unavailability](https://docs.litellm.ai/docs/proxy/prod#6-if-running-litellm-on-vpc-gracefully-handle-db-unavailability) diff --git a/docs/my-website/docs/proxy/docker_quick_start.md b/docs/my-website/docs/proxy/docker_quick_start.md index d82a0b01d1d..35d9923e92c 100644 --- a/docs/my-website/docs/proxy/docker_quick_start.md +++ b/docs/my-website/docs/proxy/docker_quick_start.md @@ -52,8 +52,6 @@ echo 'LITELLM_MASTER_KEY="sk-1234"' > .env # password generator to get a random hash for litellm salt key echo 'LITELLM_SALT_KEY="sk-1234"' >> .env -source .env - # Start docker compose up ``` diff --git a/docs/my-website/docs/proxy/dynamic_rate_limit.md b/docs/my-website/docs/proxy/dynamic_rate_limit.md index 9c875a51eba..f5438b5a6f5 100644 --- a/docs/my-website/docs/proxy/dynamic_rate_limit.md +++ b/docs/my-website/docs/proxy/dynamic_rate_limit.md @@ -175,7 +175,37 @@ general_settings: litellm --config /path/to/config.yaml ``` -#### 2. Create Keys with Priority Levels +### Set priority on either a team or a key + +Priority can be set at either the **team level** or **key level**. Team-level priority takes precedence over key-level priority. + +**Option A: Set Priority on Team (Recommended)** + +All keys within a team will inherit the team's priority. This is useful when you want all keys for a specific environment or project to have the same priority. + +```bash +curl -X POST 'http://0.0.0.0:4000/team/new' \ +-H 'Authorization: Bearer sk-1234' \ +-H 'Content-Type: application/json' \ +-d '{ + "team_alias": "production-team", + "metadata": {"priority": "prod"} +}' +``` + +Create a key for this team: +```bash +curl -X POST 'http://0.0.0.0:4000/key/generate' \ +-H 'Authorization: Bearer sk-1234' \ +-H 'Content-Type: application/json' \ +-d '{ + "team_id": "team-id-from-previous-response" +}' +``` + +**Option B: Set Priority on Individual Keys** + +Set priority directly on the key. This is useful when you need fine-grained control per key. **Production Key:** ```bash @@ -205,7 +235,7 @@ curl -X POST 'http://0.0.0.0:4000/key/generate' \ -d '{}' ``` -**Expected Response for both:** +**Expected Response:** ```json { "key": "sk-...", @@ -214,6 +244,11 @@ curl -X POST 'http://0.0.0.0:4000/key/generate' \ } ``` +**Priority Resolution Order:** +1. If key belongs to a team with `metadata.priority` set → use team priority +2. Else if key has `metadata.priority` set → use key priority +3. Else → use `default_priority` from config + #### 3. Test Priority Allocation **Test Production Key (should get 9 RPM):** diff --git a/docs/my-website/docs/response_api.md b/docs/my-website/docs/response_api.md index 96bfc196d0e..52e4c1e26f5 100644 --- a/docs/my-website/docs/response_api.md +++ b/docs/my-website/docs/response_api.md @@ -81,6 +81,85 @@ for event in stream: f.write(image_bytes) ``` +#### Image Generation (Non-streaming) + +Image generation is supported for models that generate images. Generated images are returned in the `output` array with `type: "image_generation_call"`. + +**Gemini (Google AI Studio):** +```python showLineNumbers title="Gemini Image Generation" +import litellm +import base64 + +# Gemini image generation models don't require tools parameter +response = litellm.responses( + model="gemini/gemini-2.5-flash-image", + input="Generate a cute cat playing with yarn" +) + +# Access generated images from output +for item in response.output: + if item.type == "image_generation_call": + # item.result contains pure base64 (no data: prefix) + image_bytes = base64.b64decode(item.result) + + # Save the image + with open(f"generated_{item.id}.png", "wb") as f: + f.write(image_bytes) + +print(f"Image saved: generated_{response.output[0].id}.png") +``` + +**OpenAI:** +```python showLineNumbers title="OpenAI Image Generation" +import litellm +import base64 + +# OpenAI models require tools parameter for image generation +response = litellm.responses( + model="openai/gpt-4o", + input="Generate a futuristic city at sunset", + tools=[{"type": "image_generation"}] +) + +# Access generated images from output +for item in response.output: + if item.type == "image_generation_call": + image_bytes = base64.b64decode(item.result) + with open(f"generated_{item.id}.png", "wb") as f: + f.write(image_bytes) +``` + +**Response Format:** + +When image generation is successful, the response contains: + +```json +{ + "id": "resp_abc123", + "status": "completed", + "output": [ + { + "type": "image_generation_call", + "id": "resp_abc123_img_0", + "status": "completed", + "result": "iVBORw0KGgo..." // Pure base64 string (no data: prefix) + } + ] +} +``` + +**Supported Models:** + +| Provider | Models | Requires `tools` Parameter | +|----------|--------|---------------------------| +| Google AI Studio | `gemini/gemini-2.5-flash-image` | ❌ No | +| Vertex AI | `vertex_ai/gemini-2.5-flash-image-preview` | ❌ No | +| OpenAI | `gpt-4o`, `gpt-4o-mini`, `gpt-4.1`, `gpt-4.1-mini`, `gpt-4.1-nano`, `o3` | ✅ Yes | +| AWS Bedrock | Stability AI, Amazon Nova Canvas models | Model-specific | +| Fal AI | Various image generation models | Check model docs | + +**Note:** The `result` field contains pure base64-encoded image data without the `data:image/png;base64,` prefix. You must decode it with `base64.b64decode()` before saving. + #### GET a Response ```python showLineNumbers title="Get Response by ID" import litellm diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.10-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.10-py3-none-any.whl new file mode 100644 index 00000000000..ce4e805663a Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.10-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.10.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.10.tar.gz new file mode 100644 index 00000000000..a4e218ee2fa Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.10.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251204124859_add_end_user_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251204124859_add_end_user_spend_table/migration.sql new file mode 100644 index 00000000000..c4234785c54 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251204124859_add_end_user_spend_table/migration.sql @@ -0,0 +1,42 @@ +-- CreateTable +CREATE TABLE "LiteLLM_DailyEndUserSpend" ( + "id" TEXT NOT NULL, + "end_user_id" TEXT, + "date" TEXT NOT NULL, + "api_key" TEXT NOT NULL, + "model" TEXT, + "model_group" TEXT, + "custom_llm_provider" TEXT, + "mcp_namespaced_tool_name" TEXT, + "prompt_tokens" BIGINT NOT NULL DEFAULT 0, + "completion_tokens" BIGINT NOT NULL DEFAULT 0, + "cache_read_input_tokens" BIGINT NOT NULL DEFAULT 0, + "cache_creation_input_tokens" BIGINT NOT NULL DEFAULT 0, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "api_requests" BIGINT NOT NULL DEFAULT 0, + "successful_requests" BIGINT NOT NULL DEFAULT 0, + "failed_requests" BIGINT NOT NULL DEFAULT 0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_DailyEndUserSpend_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyEndUserSpend_date_idx" ON "LiteLLM_DailyEndUserSpend"("date"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyEndUserSpend_end_user_id_idx" ON "LiteLLM_DailyEndUserSpend"("end_user_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyEndUserSpend_api_key_idx" ON "LiteLLM_DailyEndUserSpend"("api_key"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyEndUserSpend_model_idx" ON "LiteLLM_DailyEndUserSpend"("model"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyEndUserSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyEndUserSpend"("mcp_namespaced_tool_name"); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 583c493adc6..40f00437e58 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -465,6 +465,34 @@ model LiteLLM_DailyOrganizationSpend { @@index([mcp_namespaced_tool_name]) } +// Track daily end user (customer) spend metrics per model and key +model LiteLLM_DailyEndUserSpend { + id String @id @default(uuid()) + end_user_id String? + date String + api_key String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + @@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name]) + @@index([date]) + @@index([end_user_id]) + @@index([api_key]) + @@index([model]) + @@index([mcp_namespaced_tool_name]) +} + // Track daily team spend metrics per model and key model LiteLLM_DailyTeamSpend { id String @id @default(uuid()) diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 6bc576e4f3d..cc8a92b9c67 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm-proxy-extras" -version = "0.4.9" +version = "0.4.10" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." authors = ["BerriAI"] readme = "README.md" @@ -22,7 +22,7 @@ requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "0.4.9" +version = "0.4.10" version_files = [ "pyproject.toml:version", "../requirements.txt:litellm-proxy-extras==", diff --git a/litellm/__init__.py b/litellm/__init__.py index e312169ffc7..0e88873d1f5 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -520,6 +520,7 @@ perplexity_models: Set = set() watsonx_models: Set = set() gemini_models: Set = set() xai_models: Set = set() +zai_models: Set = set() deepseek_models: Set = set() runwayml_models: Set = set() azure_ai_models: Set = set() @@ -711,6 +712,8 @@ def add_known_models(): text_completion_codestral_models.add(key) elif value.get("litellm_provider") == "xai": xai_models.add(key) + elif value.get("litellm_provider") == "zai": + zai_models.add(key) elif value.get("litellm_provider") == "fal_ai": fal_ai_models.add(key) elif value.get("litellm_provider") == "deepseek": @@ -872,6 +875,7 @@ model_list = list( | gemini_models | text_completion_codestral_models | xai_models + | zai_models | fal_ai_models | deepseek_models | azure_ai_models @@ -960,6 +964,7 @@ models_by_provider: dict = { "aleph_alpha": aleph_alpha_models, "text-completion-codestral": text_completion_codestral_models, "xai": xai_models, + "zai": zai_models, "fal_ai": fal_ai_models, "deepseek": deepseek_models, "runwayml": runwayml_models, @@ -1300,6 +1305,7 @@ from .llms.friendliai.chat.transformation import FriendliaiChatConfig from .llms.jina_ai.embedding.transformation import JinaAIEmbeddingConfig from .llms.xai.chat.transformation import XAIChatConfig from .llms.xai.common_utils import XAIModelInfo +from .llms.zai.chat.transformation import ZAIChatConfig from .llms.aiml.chat.transformation import AIMLChatConfig from .llms.volcengine.chat.transformation import ( VolcEngineChatConfig as VolcEngineConfig, @@ -1497,10 +1503,46 @@ def set_global_gitlab_config(config: Dict[str, Any]) -> None: # Lazy loading system for heavy modules to reduce initial import time and memory usage if TYPE_CHECKING: + from litellm.types.utils import ModelInfo as _ModelInfoType + + # Cost calculator functions cost_per_token: Callable[..., Tuple[float, float]] completion_cost: Callable[..., float] response_cost_calculator: Any modify_integration: Any + + # Utils functions - type stubs for truly lazy loaded functions only + # (functions NOT imported via "from .main import *") + get_response_string: Callable[..., str] + supports_function_calling: Callable[..., bool] + supports_web_search: Callable[..., bool] + supports_url_context: Callable[..., bool] + supports_response_schema: Callable[..., bool] + supports_parallel_function_calling: Callable[..., bool] + supports_vision: Callable[..., bool] + supports_audio_input: Callable[..., bool] + supports_audio_output: Callable[..., bool] + supports_system_messages: Callable[..., bool] + supports_reasoning: Callable[..., bool] + acreate: Callable[..., Any] + get_max_tokens: Callable[..., int] + get_model_info: Callable[..., _ModelInfoType] + register_prompt_template: Callable[..., None] + validate_environment: Callable[..., dict] + check_valid_key: Callable[..., bool] + register_model: Callable[..., None] + encode: Callable[..., list] + decode: Callable[..., str] + _calculate_retry_after: Callable[..., float] + _should_retry: Callable[..., bool] + get_supported_openai_params: Callable[..., Optional[list]] + get_api_base: Callable[..., Optional[str]] + get_first_chars_messages: Callable[..., str] + get_provider_fields: Callable[..., List] + get_valid_models: Callable[..., list] + + # Response types - truly lazy loaded only (not in main.py or elsewhere) + ModelResponseListIterator: Type[Any] def __getattr__(name: str) -> Any: diff --git a/litellm/constants.py b/litellm/constants.py index fd3eb762395..0a87f336729 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -149,6 +149,7 @@ REDIS_UPDATE_BUFFER_KEY = "litellm_spend_update_buffer" REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_spend_update_buffer" REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_team_spend_update_buffer" REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_org_spend_update_buffer" +REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_end_user_spend_update_buffer" REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", 10000)) @@ -538,6 +539,7 @@ openai_compatible_endpoints: List = [ "https://api.friendli.ai/serverless/v1", "api.sambanova.ai/v1", "api.x.ai/v1", + "ollama.com", "api.galadriel.ai/v1", "api.llama.com/compat/v1/", "api.featherless.ai/v1", diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 57eab6d29a3..4407134fa92 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -860,9 +860,9 @@ def completion_cost( # noqa: PLR0915 or isinstance(completion_response, dict) ): # tts returns a custom class if isinstance(completion_response, dict): - usage_obj: Optional[Union[dict, Usage]] = ( - completion_response.get("usage", {}) - ) + usage_obj: Optional[ + Union[dict, Usage] + ] = completion_response.get("usage", {}) else: usage_obj = getattr(completion_response, "usage", {}) if isinstance(usage_obj, BaseModel) and not _is_known_usage_objects( @@ -1066,13 +1066,14 @@ def completion_cost( # noqa: PLR0915 # If model is like "tavily-search", construct "tavily/search" for cost lookup search_model = f"{custom_llm_provider}/search" - prompt_cost, completion_cost_result = ( - search_provider_cost_per_query( - model=search_model, - custom_llm_provider=custom_llm_provider, - number_of_queries=number_of_queries, - optional_params=optional_params, - ) + ( + prompt_cost, + completion_cost_result, + ) = search_provider_cost_per_query( + model=search_model, + custom_llm_provider=custom_llm_provider, + number_of_queries=number_of_queries, + optional_params=optional_params, ) # Return the total cost (prompt_cost + completion_cost, but for search it's just prompt_cost) @@ -1080,11 +1081,13 @@ def completion_cost( # noqa: PLR0915 # Apply discount original_cost = _final_cost - _final_cost, discount_percent, discount_amount = ( - _apply_cost_discount( - base_cost=_final_cost, - custom_llm_provider=custom_llm_provider, - ) + ( + _final_cost, + discount_percent, + discount_amount, + ) = _apply_cost_discount( + base_cost=_final_cost, + custom_llm_provider=custom_llm_provider, ) # Store cost breakdown in logging object if available @@ -1329,9 +1332,8 @@ def response_cost_calculator( response_cost = 0.0 else: if isinstance(response_object, BaseModel): - response_object._hidden_params["optional_params"] = optional_params - if hasattr(response_object, "_hidden_params"): + response_object._hidden_params["optional_params"] = optional_params provider_response_cost = get_response_cost_from_hidden_params( response_object._hidden_params ) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 288c122e0e7..9d0eff15d0c 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -229,6 +229,9 @@ def get_llm_provider( # noqa: PLR0915 elif endpoint == "api.deepseek.com/v1": custom_llm_provider = "deepseek" dynamic_api_key = get_secret_str("DEEPSEEK_API_KEY") + elif endpoint == "ollama.com": + custom_llm_provider = "ollama" + dynamic_api_key = get_secret_str("OLLAMA_API_KEY") elif endpoint == "https://api.friendli.ai/serverless/v1": custom_llm_provider = "friendliai" dynamic_api_key = get_secret_str( @@ -469,11 +472,13 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 model = model.split("/", 1)[1] # Check JSON providers FIRST (before hardcoded ones) - from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry if JSONProviderRegistry.exists(custom_llm_provider): provider_config = JSONProviderRegistry.get(custom_llm_provider) + if provider_config is None: + raise ValueError(f"Provider {custom_llm_provider} not found") config_class = create_config_class(provider_config) api_base, dynamic_api_key = config_class()._get_openai_compatible_provider_info( api_base, api_key @@ -556,6 +561,13 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 or "https://api.studio.nebius.ai/v1" ) # type: ignore dynamic_api_key = api_key or get_secret_str("NEBIUS_API_KEY") + elif custom_llm_provider == "ollama": + api_base = ( + api_base + or get_secret("OLLAMA_API_BASE") + or "http://localhost:11434" + ) # type: ignore + dynamic_api_key = api_key or get_secret_str("OLLAMA_API_KEY") elif (custom_llm_provider == "ai21_chat") or ( custom_llm_provider == "ai21" and model in litellm.ai21_chat_models ): @@ -675,12 +687,12 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 api_base, api_key ) elif custom_llm_provider == "zai": - api_base = ( - api_base - or get_secret_str("ZAI_API_BASE") - or "https://api.z.ai/api/paas/v4" + ( + api_base, + dynamic_api_key, + ) = litellm.ZAIChatConfig()._get_openai_compatible_provider_info( + api_base, api_key ) - dynamic_api_key = api_key or get_secret_str("ZAI_API_KEY") elif custom_llm_provider == "together_ai": api_base = ( api_base diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 9717f442b82..ef2183a4556 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -583,9 +583,11 @@ def generic_cost_per_token( reasoning_tokens = completion_tokens_details["reasoning_tokens"] image_tokens = completion_tokens_details["image_tokens"] - if text_tokens == 0: + # Only assume all tokens are text if there's NO breakdown at all + # If image_tokens, audio_tokens, or reasoning_tokens exist, respect text_tokens=0 + has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0 + if text_tokens == 0 and not has_token_breakdown: text_tokens = usage.completion_tokens - if text_tokens == usage.completion_tokens: is_text_tokens_total = True ## TEXT COST completion_cost = float(text_tokens) * completion_base_cost diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index a7f460fab59..4ffb7ace5b0 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -737,6 +737,7 @@ class CustomStreamWrapper: or ( "tool_calls" in model_response.choices[0].delta and model_response.choices[0].delta["tool_calls"] is not None + and len(model_response.choices[0].delta["tool_calls"]) > 0 ) or ( "function_call" in model_response.choices[0].delta diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 705f3c9e630..2a1d7f2e3a3 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -960,7 +960,10 @@ class AmazonConverseConfig(BaseConfig): bedrock_tools = _bedrock_tools_pt(filtered_tools) # Set anthropic_beta in additional_request_params if we have any beta features - if anthropic_beta_list: + # ONLY apply to Anthropic/Claude models - other models (e.g., Qwen, Llama) don't support this field + # and will error with "unknown variant anthropic_beta" if included + base_model = BedrockModelInfo.get_base_model(model) + if anthropic_beta_list and base_model.startswith("anthropic"): # Remove duplicates while preserving order unique_betas = [] seen = set() diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 5107f76fc84..038895a39e5 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -416,15 +416,34 @@ class OCIChatConfig(BaseConfig): "Please install it with: pip install cryptography" ) from e + # Handle oci_key - it should be a string (PEM content) + oci_key_content = None + if oci_key: + if isinstance(oci_key, str): + oci_key_content = oci_key + # Fix common issues with PEM content + # Replace escaped newlines with actual newlines + oci_key_content = oci_key_content.replace("\\n", "\n") + # Ensure proper line endings + if "\r\n" in oci_key_content: + oci_key_content = oci_key_content.replace("\r\n", "\n") + else: + raise OCIError( + status_code=400, + message=f"oci_key must be a string containing the PEM private key content. " + f"Got type: {type(oci_key).__name__}", + ) + private_key = ( - load_private_key_from_str(oci_key) - if oci_key + load_private_key_from_str(oci_key_content) + if oci_key_content else load_private_key_from_file(oci_key_file) if oci_key_file else None ) if private_key is None: - raise Exception( - "Private key is required for OCI authentication. Please provide either oci_key or oci_key_file." + raise OCIError( + status_code=400, + message="Private key is required for OCI authentication. Please provide either oci_key or oci_key_file.", ) signature = private_key.sign( diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index aa2580453a8..809c3e4d3e0 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -164,7 +164,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): for tool_call_idx, tool_call in enumerate(tool_calls): if isinstance(tool_call, dict): # Add the full tool call object to the list - tool_calls_to_check.append(ChatCompletionToolParam(**tool_call)) + tool_calls_to_check.append(cast(ChatCompletionToolParam, tool_call)) tool_call_task_mappings.append((msg_idx, int(tool_call_idx))) async def _apply_guardrail_responses_to_input_texts( @@ -380,20 +380,20 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if isinstance(content, str): # String content - accumulate for this choice - key = (choice_idx, None) - if key not in combined_texts: - combined_texts[key] = "" - combined_texts[key] += content + str_key: Tuple[int, Optional[int]] = (choice_idx, None) + if str_key not in combined_texts: + combined_texts[str_key] = "" + combined_texts[str_key] += content elif isinstance(content, list): # List content - accumulate for each content item for content_idx, content_item in enumerate(content): text_str = content_item.get("text") if text_str: - key = (choice_idx, content_idx) - if key not in combined_texts: - combined_texts[key] = "" - combined_texts[key] += text_str + list_key: Tuple[int, Optional[int]] = (choice_idx, content_idx) + if list_key not in combined_texts: + combined_texts[list_key] = "" + combined_texts[list_key] += text_str # Step 2: Create lists for guardrail processing texts_to_check: List[str] = [] @@ -401,9 +401,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): task_mappings: List[Tuple[int, Optional[int]]] = [] # Track (choice_index, content_index) for each combined text - for (choice_idx, content_idx), combined_text in combined_texts.items(): + for (map_choice_idx, map_content_idx), combined_text in combined_texts.items(): texts_to_check.append(combined_text) - task_mappings.append((choice_idx, content_idx)) + task_mappings.append((map_choice_idx, map_content_idx)) # Step 3: Apply guardrail to all combined texts in batch if texts_to_check: @@ -503,7 +503,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Determine content source and tool calls based on choice type content = None - tool_calls = None + tool_calls: Optional[List[Any]] = None if isinstance(choice, litellm.Choices): content = choice.message.content tool_calls = choice.message.tool_calls @@ -686,15 +686,15 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if isinstance(content, str): # String content - key = (choice_idx_in_response, None) - if key in guardrail_map: - if key not in already_set: + str_key: Tuple[int, Optional[int]] = (choice_idx_in_response, None) + if str_key in guardrail_map: + if str_key not in already_set: # First chunk - set the complete guardrailed text if isinstance(choice, litellm.StreamingChoices): - choice.delta.content = guardrail_map[key] + choice.delta.content = guardrail_map[str_key] elif isinstance(choice, litellm.Choices): - choice.message.content = guardrail_map[key] - already_set[key] = True + choice.message.content = guardrail_map[str_key] + already_set[str_key] = True else: # Subsequent chunks - clear the content if isinstance(choice, litellm.StreamingChoices): @@ -706,12 +706,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # List content - handle each content item for content_idx, content_item in enumerate(content): if "text" in content_item: - key = (choice_idx_in_response, content_idx) - if key in guardrail_map: - if key not in already_set: + list_key: Tuple[int, Optional[int]] = (choice_idx_in_response, content_idx) + if list_key in guardrail_map: + if list_key not in already_set: # First chunk - set the complete guardrailed text - content_item["text"] = guardrail_map[key] - already_set[key] = True + content_item["text"] = guardrail_map[list_key] + already_set[list_key] = True else: # Subsequent chunks - clear the text content_item["text"] = "" diff --git a/litellm/llms/openai_like/dynamic_config.py b/litellm/llms/openai_like/dynamic_config.py index ca2489799c2..1e7866bebbe 100644 --- a/litellm/llms/openai_like/dynamic_config.py +++ b/litellm/llms/openai_like/dynamic_config.py @@ -19,11 +19,11 @@ def create_config_class(provider: SimpleProviderConfig): """Generate config class dynamically from JSON configuration""" # Choose base class - base_class = ( + base_class: type = ( OpenAIGPTConfig if provider.base_class == "openai_gpt" else OpenAILikeChatConfig ) - class JSONProviderConfig(base_class): + class JSONProviderConfig(base_class): # type: ignore[valid-type,misc] @overload def _transform_messages( self, messages: List[AllMessageValues], model: str, is_async: Literal[True] @@ -87,6 +87,9 @@ def create_config_class(provider: SimpleProviderConfig): if not api_base: api_base = provider.base_url + if api_base is None: + raise ValueError(f"api_base is required for provider {provider.slug}") + if not api_base.endswith("/chat/completions"): api_base = f"{api_base}/chat/completions" diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 3151a6d667e..a95d5447e97 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -116,7 +116,7 @@ def _process_gemini_image( is not None ): file_data = FileDataType(file_uri=image_url, mime_type=image_type) - part: PartType = {"file_data": file_data} + part = {"file_data": file_data} if media_resolution_enum is not None and model is not None: from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig @@ -129,7 +129,7 @@ def _process_gemini_image( image = convert_to_anthropic_image_obj(image_url, format=format) _blob: BlobType = {"data": image["data"], "mime_type": image["media_type"]} - part: PartType = {"inline_data": cast(BlobType, _blob)} + part = {"inline_data": cast(BlobType, _blob)} if media_resolution_enum is not None and model is not None: from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index e604bd392a6..106074811f6 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1085,24 +1085,26 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): def _extract_thinking_blocks_from_parts( self, parts: List[HttpxPartType] ) -> List[ChatCompletionThinkingBlock]: - """Extract thinking blocks from parts if present""" + """Extract thinking blocks from parts if present. + + Per Google's docs (https://ai.google.dev/gemini-api/docs/thinking): + - Parts with `thought: true` contain thinking/reasoning content + - `thoughtSignature` is a separate token for multi-turn context preservation, + it does NOT indicate that the content is thinking (a part can have + thoughtSignature without thought: true, e.g., function calls) + """ thinking_blocks: List[ChatCompletionThinkingBlock] = [] for part in parts: - if "thoughtSignature" in part: - part_copy = part.copy() - part_copy.pop("thoughtSignature") - - text_content = part_copy.get("text") - if isinstance(text_content, str) and text_content.strip() == "": - continue - - thinking_blocks.append( - ChatCompletionThinkingBlock( - type="thinking", - thinking=json.dumps(part_copy), - signature=part["thoughtSignature"], - ) - ) + if part.get("thought") is True: + thinking_text = part.get("text", "") + block: ChatCompletionThinkingBlock = { + "type": "thinking", + "thinking": thinking_text, + } + signature = part.get("thoughtSignature") + if signature is not None: + block["signature"] = signature + thinking_blocks.append(block) return thinking_blocks def _extract_image_response_from_parts( diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 416c611d863..b9747652362 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -140,7 +140,11 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): if not vertex_project or not vertex_location: raise ValueError("vertex_project and vertex_location are required for Vertex AI") - base_url = f"https://{vertex_location}-aiplatform.googleapis.com" + # Handle global location differently (no region prefix in URL) + if vertex_location == "global": + base_url = "https://aiplatform.googleapis.com" + else: + base_url = f"https://{vertex_location}-aiplatform.googleapis.com" return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}:generateContent" diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index aff14b1004f..18ca077c4da 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -220,7 +220,7 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): Returns: Tuple of (mapped_voice_str, mapped_params) """ - mapped_params = {} + mapped_params: Dict[str, Any] = {} ########################################################## # Map voice using helper diff --git a/litellm/llms/zai/__init__.py b/litellm/llms/zai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/zai/chat/__init__.py b/litellm/llms/zai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/zai/chat/transformation.py b/litellm/llms/zai/chat/transformation.py new file mode 100644 index 00000000000..47b314d4e0d --- /dev/null +++ b/litellm/llms/zai/chat/transformation.py @@ -0,0 +1,33 @@ +from typing import Optional, Tuple + +from litellm.secret_managers.main import get_secret_str + +from ...openai.chat.gpt_transformation import OpenAIGPTConfig + +ZAI_API_BASE = "https://api.z.ai/api/paas/v4" + + +class ZAIChatConfig(OpenAIGPTConfig): + @property + def custom_llm_provider(self) -> Optional[str]: + return "zai" + + def _get_openai_compatible_provider_info( + self, api_base: Optional[str], api_key: Optional[str] + ) -> Tuple[Optional[str], Optional[str]]: + api_base = api_base or get_secret_str("ZAI_API_BASE") or ZAI_API_BASE + dynamic_api_key = api_key or get_secret_str("ZAI_API_KEY") + return api_base, dynamic_api_key + + def get_supported_openai_params(self, model: str) -> list: + return [ + "max_tokens", + "stream", + "stream_options", + "temperature", + "top_p", + "stop", + "tools", + "tool_choice", + ] + diff --git a/litellm/main.py b/litellm/main.py index 0e19699309d..3eaa4527f8b 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3503,6 +3503,9 @@ def completion( # type: ignore # noqa: PLR0915 or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" ) + if api_key is not None and "Authorization" not in headers: + headers["Authorization"] = f"Bearer {api_key}" + response = base_llm_http_handler.completion( model=model, stream=stream, @@ -3536,6 +3539,9 @@ def completion( # type: ignore # noqa: PLR0915 or os.environ.get("OLLAMA_API_KEY") or litellm.api_key ) + if api_key is not None and "Authorization" not in headers: + headers["Authorization"] = f"Bearer {api_key}" + response = base_llm_http_handler.completion( model=model, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 634ea6dc48a..c26aada0a50 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5164,6 +5164,19 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "azure_ai/mistral-large-3": { + "input_cost_per_token": 5e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 256000, + "max_output_tokens": 8191, + "max_tokens": 8191, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://azure.microsoft.com/en-us/blog/introducing-mistral-large-3-in-microsoft-foundry-open-capable-and-ready-for-production-workloads/", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure_ai/mistral-medium-2505": { "input_cost_per_token": 4e-07, "litellm_provider": "azure_ai", @@ -12134,6 +12147,7 @@ "max_tokens": 65536, "mode": "image_generation", "output_cost_per_image": 0.134, + "output_cost_per_image_token": 1.2e-04, "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", @@ -13871,6 +13885,7 @@ "max_tokens": 65536, "mode": "image_generation", "output_cost_per_image": 0.134, + "output_cost_per_image_token": 1.2e-04, "output_cost_per_token": 1.2e-05, "rpm": 1000, "tpm": 4000000, @@ -18745,6 +18760,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/mistral-large-3": { + "input_cost_per_token": 5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 256000, + "max_output_tokens": 8191, + "max_tokens": 8191, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.mistral.ai/models/mistral-large-3-25-12", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/mistral-medium": { "input_cost_per_token": 2.7e-06, "litellm_provider": "mistral", @@ -25774,6 +25804,7 @@ "max_tokens": 65536, "mode": "image_generation", "output_cost_per_image": 0.134, + "output_cost_per_image_token": 1.2e-04, "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d42fd80cbd7..8c9d8630457 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1272,19 +1272,19 @@ class MCPServerManager: prefix = get_server_prefix(server) for tool in tools: - prefixed_name = add_server_prefix_to_name(tool.name, prefix) + tool_copy = tool.model_copy(deep=True) - name_to_use = prefixed_name if add_prefix else tool.name + original_name = tool_copy.name + prefixed_name = add_server_prefix_to_name(original_name, prefix) - tool_obj = MCPTool( - name=name_to_use, - description=tool.description, - inputSchema=tool.inputSchema, - ) - prefixed_tools.append(tool_obj) + name_to_use = prefixed_name if add_prefix else original_name + + # Preserve all tool fields including metadata/_meta by avoiding mutation + tool_copy.name = name_to_use + prefixed_tools.append(tool_copy) # Update tool to server mapping for resolution (support both forms) - self.tool_name_to_mcp_server_name_mapping[tool.name] = prefix + self.tool_name_to_mcp_server_name_mapping[original_name] = prefix self.tool_name_to_mcp_server_name_mapping[prefixed_name] = prefix verbose_logger.info( @@ -1953,9 +1953,9 @@ class MCPServerManager: server_name_from_prefix ): return server - elif normalize_server_name(server.server_name) == normalize_server_name( - server_name_from_prefix - ): + elif normalize_server_name( + server.server_name + ) == normalize_server_name(server_name_from_prefix): return server return None diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 412a6de0059..edf53e99573 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -6,7 +6,7 @@ LiteLLM MCP Server Routes import asyncio import contextlib from datetime import datetime -from typing import Any, AsyncIterator, Dict, List, Optional, Tuple, Union +from typing import Any, AsyncIterator, Dict, List, Optional, Tuple, Union, cast from fastapi import FastAPI, HTTPException from pydantic import AnyUrl, ConfigDict @@ -72,7 +72,13 @@ if MCP_AVAILABLE: auth_context_var, ) from mcp.server.streamable_http_manager import StreamableHTTPSessionManager - from mcp.types import EmbeddedResource, ImageContent, Prompt, TextContent + from mcp.types import ( + CallToolResult, + EmbeddedResource, + ImageContent, + Prompt, + TextContent, + ) from mcp.types import Tool as MCPTool from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( @@ -234,7 +240,7 @@ if MCP_AVAILABLE: @server.call_tool() async def mcp_server_tool_call( name: str, arguments: Dict[str, Any] | None - ) -> List[Union[TextContent, ImageContent, EmbeddedResource]]: + ) -> CallToolResult: """ Call a specific tool with the provided arguments @@ -300,26 +306,37 @@ if MCP_AVAILABLE: ) except BlockedPiiEntityError as e: verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {str(e)}") - # Return error as text content for MCP protocol - return [ - TextContent( - text=f"Error: Blocked PII entity detected - {str(e)}", type="text" - ) - ] + return CallToolResult( + content=[ + TextContent( + text=f"Error: Blocked PII entity detected - {str(e)}", + type="text", + ) + ], + isError=True, + ) except GuardrailRaisedException as e: verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {str(e)}") - # Return error as text content for MCP protocol - return [ - TextContent(text=f"Error: Guardrail violation - {str(e)}", type="text") - ] + return CallToolResult( + content=[ + TextContent( + text=f"Error: Guardrail violation - {str(e)}", type="text" + ) + ], + isError=True, + ) except HTTPException as e: verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}") - # Return error as text content for MCP protocol - return [TextContent(text=f"Error: {str(e.detail)}", type="text")] + return CallToolResult( + content=[TextContent(text=f"Error: {str(e.detail)}", type="text")], + isError=True, + ) except Exception as e: verbose_logger.exception(f"MCP mcp_server_tool_call - error: {e}") - # Return error as text content for MCP protocol - return [TextContent(text=f"Error: {str(e)}", type="text")] + return CallToolResult( + content=[TextContent(text=f"Error: {str(e)}", type="text")], + isError=True, + ) return response @@ -1173,7 +1190,7 @@ if MCP_AVAILABLE: oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, **kwargs: Any, - ) -> List[Union[TextContent, ImageContent, EmbeddedResource]]: + ) -> CallToolResult: """ Call a specific tool with the provided arguments (handles prefixed tool names) """ @@ -1237,9 +1254,9 @@ if MCP_AVAILABLE: "litellm_logging_obj", None ) if litellm_logging_obj: - litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( - standard_logging_mcp_tool_call - ) + litellm_logging_obj.model_call_details[ + "mcp_tool_call_metadata" + ] = standard_logging_mcp_tool_call litellm_logging_obj.model = f"MCP: {name}" # Check if tool exists in local registry first (for OpenAPI-based tools) # These tools are registered with their prefixed names @@ -1247,7 +1264,8 @@ if MCP_AVAILABLE: local_tool = global_mcp_tool_registry.get_tool(name) if local_tool: verbose_logger.debug(f"Executing local registry tool: {name}") - response = await _handle_local_mcp_tool(name, arguments) + local_content = await _handle_local_mcp_tool(name, arguments) + response = CallToolResult(content=cast(Any, local_content), isError=False) # Try managed MCP server tool (pass the full prefixed name) # Primary and recommended way to use external MCP servers @@ -1279,7 +1297,12 @@ if MCP_AVAILABLE: # Deprecated: Local MCP Server Tool ######################################################### else: - response = await _handle_local_mcp_tool(original_tool_name, arguments) + local_content = await _handle_local_mcp_tool( + original_tool_name, arguments + ) + response = CallToolResult( + content=cast(Any, local_content), isError=False + ) ######################################################### # Post MCP Tool Call Hook @@ -1432,7 +1455,7 @@ if MCP_AVAILABLE: oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, litellm_logging_obj: Optional[Any] = None, - ) -> List[Union[TextContent, ImageContent, EmbeddedResource]]: + ) -> CallToolResult: """Handle tool execution for managed server tools""" # Import here to avoid circular import from litellm.proxy.proxy_server import proxy_logging_obj @@ -1449,7 +1472,7 @@ if MCP_AVAILABLE: proxy_logging_obj=proxy_logging_obj, ) verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) - return call_tool_result.content # type: ignore[return-value] + return call_tool_result async def _handle_local_mcp_tool( name: str, arguments: Dict[str, Any] @@ -1741,14 +1764,16 @@ if MCP_AVAILABLE: ) auth_context_var.set(auth_user) - def get_auth_context() -> Tuple[ - Optional[UserAPIKeyAuth], - Optional[str], - Optional[List[str]], - Optional[Dict[str, Dict[str, str]]], - Optional[Dict[str, str]], - Optional[Dict[str, str]], - ]: + def get_auth_context() -> ( + Tuple[ + Optional[UserAPIKeyAuth], + Optional[str], + Optional[List[str]], + Optional[Dict[str, Dict[str, str]]], + Optional[Dict[str, str]], + Optional[Dict[str, str]], + ] + ): """ Get the UserAPIKeyAuth from the auth context variable. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9edc93bfffa..5eb05be3bf9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3639,6 +3639,8 @@ class DailyOrganizationSpendTransaction(BaseDailySpendTransaction): class DailyUserSpendTransaction(BaseDailySpendTransaction): user_id: str +class DailyEndUserSpendTransaction(BaseDailySpendTransaction): + end_user_id: str class DailyTagSpendTransaction(BaseDailySpendTransaction): request_id: Optional[str] diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 6ef983b221c..fb9757ca647 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -9,9 +9,9 @@ import os import secrets from typing import Literal, Optional, cast -import litellm from fastapi import HTTPException +import litellm from litellm.constants import LITELLM_PROXY_ADMIN_NAME from litellm.proxy._types import ( LiteLLM_UserTable, @@ -64,13 +64,19 @@ def get_ui_credentials(master_key: Optional[str]) -> tuple[str, str]: class LoginResult: """Result object containing authentication data from login.""" + user_id: str + key: str + user_email: Optional[str] + user_role: str + login_method: Literal["sso", "username_password"] + def __init__( self, user_id: str, key: str, user_email: Optional[str], user_role: str, - login_method: str = "username_password", + login_method: Literal["sso", "username_password"] = "username_password", ): self.user_id = user_id self.key = key @@ -193,14 +199,14 @@ async def authenticate_user( key = response["token"] # type: ignore if get_secret_bool("EXPERIMENTAL_UI_LOGIN"): + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + user_info: Optional[LiteLLM_UserTable] = None if _user_row is not None: user_info = _user_row elif ( user_id is not None ): # if user_id is not None, we are using the UI_USERNAME and UI_PASSWORD - from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken - user_info = LiteLLM_UserTable( user_id=user_id, user_role=user_role, diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 6c9289e3ff6..715a6ebd25d 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -25,6 +25,7 @@ from litellm.proxy._types import ( DailyTagSpendTransaction, DailyOrganizationSpendTransaction, DailyTeamSpendTransaction, + DailyEndUserSpendTransaction, DailyUserSpendTransaction, DBSpendUpdateTransactions, Litellm_EntityType, @@ -65,6 +66,7 @@ class DBSpendUpdateWriter: self.spend_update_queue = SpendUpdateQueue() self.daily_spend_update_queue = DailySpendUpdateQueue() self.daily_team_spend_update_queue = DailySpendUpdateQueue() + self.daily_end_user_spend_update_queue = DailySpendUpdateQueue() self.daily_org_spend_update_queue = DailySpendUpdateQueue() self.daily_tag_spend_update_queue = DailySpendUpdateQueue() @@ -182,6 +184,13 @@ class DBSpendUpdateWriter: ) ) + asyncio.create_task( + self.add_spend_log_transaction_to_daily_end_user_transaction( + payload=payload, + prisma_client=prisma_client, + ) + ) + asyncio.create_task( self.add_spend_log_transaction_to_daily_team_transaction( payload=payload, @@ -475,6 +484,7 @@ class DBSpendUpdateWriter: daily_spend_update_queue=self.daily_spend_update_queue, daily_team_spend_update_queue=self.daily_team_spend_update_queue, daily_org_spend_update_queue=self.daily_org_spend_update_queue, + daily_end_user_spend_update_queue=self.daily_end_user_spend_update_queue, daily_tag_spend_update_queue=self.daily_tag_spend_update_queue, ) @@ -538,6 +548,16 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_tag_spend_update_transactions, ) + daily_end_user_spend_update_transactions = ( + await self.redis_update_buffer.get_all_daily_end_user_spend_update_transactions_from_redis_buffer() + ) + if daily_end_user_spend_update_transactions is not None: + await DBSpendUpdateWriter.update_daily_end_user_spend( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_end_user_spend_update_transactions, + ) except Exception as e: verbose_proxy_logger.error(f"Error committing spend updates: {e}") finally: @@ -627,6 +647,20 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_tag_spend_update_transactions, ) + ################## Daily End-User Spend Update Transactions ################## + # Aggregate all in memory daily end-user spend transactions and commit to db + daily_end_user_spend_update_transactions = cast( + Dict[str, DailyEndUserSpendTransaction], + await self.daily_end_user_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(), + ) + + await DBSpendUpdateWriter.update_daily_end_user_spend( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_end_user_spend_update_transactions, + ) + async def _commit_spend_updates_to_db( # noqa: PLR0915 self, prisma_client: PrismaClient, @@ -990,6 +1024,20 @@ class DBSpendUpdateWriter: ) -> None: ... + @overload + @staticmethod + async def _update_daily_spend( + n_retry_times: int, + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + daily_spend_transactions: Dict[str, DailyEndUserSpendTransaction], + entity_type: Literal["end_user"], + entity_id_field: str, + table_name: str, + unique_constraint_name: str, + ) -> None: + ... + @overload @staticmethod async def _update_daily_spend( @@ -1015,14 +1063,15 @@ class DBSpendUpdateWriter: Dict[str, DailyTeamSpendTransaction], Dict[str, DailyTagSpendTransaction], Dict[str, DailyOrganizationSpendTransaction], + Dict[str, DailyEndUserSpendTransaction], ], - entity_type: Literal["user", "team", "org", "tag"], + entity_type: Literal["user", "team", "org", "tag", "end_user"], entity_id_field: str, table_name: str, unique_constraint_name: str, ) -> None: """ - Generic function to update daily spend for any entity type (user, team, org, tag) + Generic function to update daily spend for any entity type (user, team, org, tag, end_user) """ from litellm.proxy.utils import _raise_failed_update_spend_exception @@ -1267,6 +1316,27 @@ class DBSpendUpdateWriter: unique_constraint_name="organization_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name", ) + @staticmethod + async def update_daily_end_user_spend( + n_retry_times: int, + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + daily_spend_transactions: Dict[str, DailyEndUserSpendTransaction], + ): + """ + Batch job to update LiteLLM_DailyEndUserSpend table using in-memory daily_spend_transactions + """ + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_spend_transactions, + entity_type="end_user", + entity_id_field="end_user_id", + table_name="litellm_dailyenduserspend", + unique_constraint_name="end_user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name", + ) + @staticmethod async def update_daily_tag_spend( n_retry_times: int, @@ -1292,7 +1362,7 @@ class DBSpendUpdateWriter: self, payload: Union[dict, SpendLogsPayload], prisma_client: PrismaClient, - type: Literal["user", "team", "org", "request_tags"] = "user", + type: Literal["user", "team", "org", "request_tags", "end_user"] = "user", ) -> Optional[BaseDailySpendTransaction]: common_expected_keys = ["startTime", "api_key"] if type == "user": @@ -1303,6 +1373,8 @@ class DBSpendUpdateWriter: expected_keys = ["organization_id", *common_expected_keys] elif type == "request_tags": expected_keys = ["request_tags", *common_expected_keys] + elif type == "end_user": + expected_keys = ["end_user_id", *common_expected_keys] else: raise ValueError(f"Invalid type: {type}") if not all(key in payload for key in expected_keys): @@ -1474,6 +1546,48 @@ class DBSpendUpdateWriter: update={daily_transaction_key: daily_transaction} ) + async def add_spend_log_transaction_to_daily_end_user_transaction( + self, + payload: SpendLogsPayload, + prisma_client: Optional[PrismaClient] = None, + ) -> None: + if prisma_client is None: + verbose_proxy_logger.debug( + "prisma_client is None. Skipping writing spend logs to db." + ) + return + + end_user_id = payload.get("end_user") + if end_user_id is None or end_user_id == "": + verbose_proxy_logger.debug( + "end_user is None or empty for request. Skipping incrementing end user spend." + ) + return + + payload_with_end_user_id = cast( + SpendLogsPayload, + { + **payload, + "end_user_id": end_user_id, + }, + ) + + base_daily_transaction = ( + await self._common_add_spend_log_transaction_to_daily_transaction( + payload_with_end_user_id, prisma_client, "end_user" + ) + ) + if base_daily_transaction is None: + return + + daily_transaction_key = f"{end_user_id}_{base_daily_transaction['date']}_{payload_with_end_user_id['api_key']}_{payload_with_end_user_id['model']}_{payload_with_end_user_id['custom_llm_provider']}" + daily_transaction = DailyEndUserSpendTransaction( + end_user_id=end_user_id, **base_daily_transaction + ) + await self.daily_end_user_spend_update_queue.add_update( + update={daily_transaction_key: daily_transaction} + ) + async def add_spend_log_transaction_to_daily_tag_transaction( self, payload: SpendLogsPayload, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 921fd9701bd..e3b20d7266d 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -16,6 +16,7 @@ from litellm.constants import ( REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -24,6 +25,7 @@ from litellm.proxy._types import ( DailyTeamSpendTransaction, DailyUserSpendTransaction, DailyOrganizationSpendTransaction, + DailyEndUserSpendTransaction, DBSpendUpdateTransactions, ) from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj @@ -107,6 +109,7 @@ class RedisUpdateBuffer: daily_spend_update_queue: DailySpendUpdateQueue, daily_team_spend_update_queue: DailySpendUpdateQueue, daily_org_spend_update_queue: DailySpendUpdateQueue, + daily_end_user_spend_update_queue: DailySpendUpdateQueue, daily_tag_spend_update_queue: DailySpendUpdateQueue, ): """ @@ -172,6 +175,9 @@ class RedisUpdateBuffer: daily_org_spend_update_transactions = ( await daily_org_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) + daily_end_user_spend_update_transactions = ( + await daily_end_user_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) daily_tag_spend_update_transactions = ( await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) @@ -207,6 +213,12 @@ class RedisUpdateBuffer: service_type=ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE, ) + await self._store_transactions_in_redis( + transactions=daily_end_user_spend_update_transactions, + redis_key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + service_type=ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE, + ) + await self._store_transactions_in_redis( transactions=daily_tag_spend_update_transactions, redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, @@ -365,6 +377,30 @@ class RedisUpdateBuffer: ), ) + async def get_all_daily_end_user_spend_update_transactions_from_redis_buffer( + self, + ) -> Optional[Dict[str, DailyEndUserSpendTransaction]]: + """ + Gets all the daily end-user spend update transactions from Redis + """ + if self.redis_cache is None: + return None + list_of_transactions = await self.redis_cache.async_lpop( + key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + ) + if list_of_transactions is None: + return None + list_of_daily_spend_update_transactions = [ + json.loads(transaction) for transaction in list_of_transactions + ] + return cast( + Dict[str, DailyEndUserSpendTransaction], + DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( + list_of_daily_spend_update_transactions + ), + ) + async def get_all_daily_tag_spend_update_transactions_from_redis_buffer( self, ) -> Optional[Dict[str, DailyTagSpendTransaction]]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index c7b4f19a089..6ad21a4758a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -127,7 +127,7 @@ class GenericGuardrailAPI(CustomGuardrail): for field_name in GenericGuardrailAPIMetadata.__annotations__.keys(): value = metadata_dict.get(field_name) if value is not None: - result_metadata[field_name] = value + result_metadata[field_name] = value # type: ignore[literal-required] # handle user_api_key_token = user_api_key_hash if metadata_dict.get("user_api_key_token") is not None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index d183b688edd..8666f6add53 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -207,6 +207,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): Send text to the Presidio analyzer endpoint and get analysis results """ try: + # Skip empty or whitespace-only text to avoid Presidio errors + # Common in tool/function calling where assistant content is empty + if not text or len(text.strip()) == 0: + verbose_proxy_logger.debug( + "Skipping Presidio analysis for empty/whitespace-only text" + ) + return [] + async with aiohttp.ClientSession() as session: if self.mock_redacted_text is not None: return self.mock_redacted_text @@ -231,9 +239,42 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async with session.post(analyze_url, json=analyze_payload) as response: analyze_results = await response.json() verbose_proxy_logger.debug("analyze_results: %s", analyze_results) + + # Handle error responses from Presidio (e.g., {'error': 'No text provided'}) + # Presidio may return a dict instead of a list when errors occur + if isinstance(analyze_results, dict): + if "error" in analyze_results: + verbose_proxy_logger.warning( + "Presidio analyzer returned error: %s, returning empty list", + analyze_results.get("error") + ) + return [] + # If it's a dict but not an error, try to process it as a single item + verbose_proxy_logger.debug( + "Presidio returned dict (not list), attempting to process as single item" + ) + try: + return [PresidioAnalyzeResponseItem(**analyze_results)] + except Exception as e: + verbose_proxy_logger.warning( + "Failed to parse Presidio dict response: %s, returning empty list", + e + ) + return [] + + # Normal case: list of results final_results = [] for item in analyze_results: - final_results.append(PresidioAnalyzeResponseItem(**item)) + try: + final_results.append(PresidioAnalyzeResponseItem(**item)) + except TypeError as te: + # Handle case where item is not a dict (shouldn't happen, but be defensive) + verbose_proxy_logger.warning( + "Skipping invalid Presidio result item: %s (error: %s)", + item, + te + ) + continue return final_results except Exception as e: raise e diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 5e4784d709e..79e9838d115 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -409,7 +409,7 @@ def _build_model_param_to_info_mapping(model_list: list) -> dict: Returns: Dictionary mapping model parameter to list of model info dicts """ - model_param_to_info = {} + model_param_to_info: dict = {} for model in model_list: model_info = model.get("model_info", {}) model_name = model.get("model_name") @@ -1048,28 +1048,6 @@ async def health_readiness(): index_info = "index does not exist - error: " + str(e) cache_type = {"type": cache_type, "index_info": index_info} - # build license metadata - try: - from litellm.proxy.proxy_server import _license_check # type: ignore - - license_available: bool = _license_check.is_premium() if _license_check else False - license_expiration: Optional[str] = None - - if getattr(_license_check, "airgapped_license_data", None): - license_expiration = _license_check.airgapped_license_data.get( # type: ignore[arg-type] - "expiration_date" - ) - - license_metadata = { - "license": { - "has_license": license_available, - "expiration_date": license_expiration, - } - } - except Exception: - # fail closed: don't let license check break readiness - license_metadata = {"license": {"has_license": False, "expiration_date": None}} - # check DB if prisma_client is not None: # if db passed in, check if it's connected db_health_status = await _db_health_readiness_check() @@ -1080,7 +1058,6 @@ async def health_readiness(): "litellm_version": version, "success_callbacks": success_callback_names, "use_aiohttp_transport": AsyncHTTPHandler._should_use_aiohttp_transport(), - **license_metadata, **db_health_status, } else: @@ -1091,7 +1068,6 @@ async def health_readiness(): "litellm_version": version, "success_callbacks": success_callback_names, "use_aiohttp_transport": AsyncHTTPHandler._should_use_aiohttp_transport(), - **license_metadata, } except Exception as e: raise HTTPException(status_code=503, detail=f"Service Unhealthy ({str(e)})") diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 7e6ec1dc151..53419ef6ad7 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -80,6 +80,32 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): weight = convert_priority_to_percent(value, model_info) return weight + def _get_priority_from_user_api_key_dict( + self, user_api_key_dict: UserAPIKeyAuth + ) -> Optional[str]: + """ + Get priority from user_api_key_dict. + + Checks team metadata first (takes precedence), then falls back to key metadata. + + Args: + user_api_key_dict: User authentication info + + Returns: + Priority string if found, None otherwise + """ + priority: Optional[str] = None + + # Check team metadata first (takes precedence) + if user_api_key_dict.team_metadata is not None: + priority = user_api_key_dict.team_metadata.get("priority", None) + + # Fall back to key metadata + if priority is None: + priority = user_api_key_dict.metadata.get("priority", None) + + return priority + def _normalize_priority_weights( self, model_info: ModelGroupInfo ) -> Dict[str, float]: @@ -328,7 +354,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): model: str, model_group_info: ModelGroupInfo, user_api_key_dict: UserAPIKeyAuth, - key_priority: Optional[str], + priority: Optional[str], saturation: float, data: dict, ) -> None: @@ -355,7 +381,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): model: Model name model_group_info: Model configuration user_api_key_dict: User authentication info - key_priority: User's priority level + priority: User's priority level saturation: Current saturation level data: Request data dictionary @@ -384,7 +410,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): priority_descriptors = self._create_priority_based_descriptors( model=model, user_api_key_dict=user_api_key_dict, - priority=key_priority, + priority=priority, ) if priority_descriptors: descriptors_to_check.extend(priority_descriptors) @@ -412,14 +438,14 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): status_code=429, detail={ "error": f"Model capacity reached for {model}. " - f"Priority: {key_priority}, " + f"Priority: {priority}, " f"Rate limit type: {status['rate_limit_type']}, " f"Remaining: {status['limit_remaining']}" }, headers={ "retry-after": str(self.v3_limiter.window_size), "rate_limit_type": str(status["rate_limit_type"]), - "x-litellm-priority": key_priority or "default", + "x-litellm-priority": priority or "default", }, ) @@ -427,13 +453,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): elif descriptor_key == "priority_model" and should_enforce_priority: verbose_proxy_logger.debug( f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, " - f"priority: {key_priority}" + f"priority: {priority}" ) raise HTTPException( status_code=429, detail={ "error": f"Priority-based rate limit exceeded. " - f"Priority: {key_priority}, " + f"Priority: {priority}, " f"Rate limit type: {status['rate_limit_type']}, " f"Remaining: {status['limit_remaining']}, " f"Model saturation: {saturation:.1%}" @@ -441,7 +467,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): headers={ "retry-after": str(self.v3_limiter.window_size), "rate_limit_type": str(status["rate_limit_type"]), - "x-litellm-priority": key_priority or "default", + "x-litellm-priority": priority or "default", "x-litellm-saturation": f"{saturation:.2%}", }, ) @@ -521,7 +547,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return None model = data["model"] - key_priority: Optional[str] = user_api_key_dict.metadata.get("priority", None) + priority = self._get_priority_from_user_api_key_dict( + user_api_key_dict=user_api_key_dict + ) # Get model configuration model_group_info: Optional[ModelGroupInfo] = ( @@ -543,7 +571,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): verbose_proxy_logger.debug( f"[Dynamic Rate Limiter] Model={model}, Saturation={saturation:.1%}, " - f"Threshold={saturation_threshold:.1%}, Priority={key_priority}" + f"Threshold={saturation_threshold:.1%}, Priority={priority}" ) # STEP 2: Check rate limits in THREE phases @@ -555,7 +583,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): model=model, model_group_info=model_group_info, user_api_key_dict=user_api_key_dict, - key_priority=key_priority, + priority=priority, saturation=saturation, data=data, ) @@ -586,8 +614,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # Add additional priority-specific headers if isinstance(response, ModelResponse): - key_priority: Optional[str] = user_api_key_dict.metadata.get( - "priority", None + priority = self._get_priority_from_user_api_key_dict( + user_api_key_dict=user_api_key_dict ) # Get existing additional headers @@ -599,7 +627,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ) # Add priority information - additional_headers["x-litellm-priority"] = key_priority or "default" + additional_headers["x-litellm-priority"] = priority or "default" additional_headers["x-litellm-rate-limiter-version"] = "v3" # Update response @@ -614,3 +642,121 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): f"Error in dynamic rate limiter v3 post-call hook: {str(e)}" ) return response + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + """ + Update token usage for priority-based rate limiting after successful API calls. + + Increments token counters for: + - model_saturation_check: Model-wide token tracking + - priority_model: Priority-specific token tracking + """ + from litellm.litellm_core_utils.core_helpers import ( + _get_parent_otel_span_from_kwargs, + ) + from litellm.proxy.common_utils.callback_utils import ( + get_model_group_from_litellm_kwargs, + ) + from litellm.types.caching import RedisPipelineIncrementOperation + from litellm.types.utils import Usage + + try: + verbose_proxy_logger.debug( + "INSIDE dynamic rate limiter ASYNC SUCCESS LOGGING" + ) + + litellm_parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) + + # Get metadata from standard_logging_object + standard_logging_object = kwargs.get("standard_logging_object") or {} + standard_logging_metadata = standard_logging_object.get("metadata") or {} + + # Get model and priority + model_group = get_model_group_from_litellm_kwargs(kwargs) + if not model_group: + return + + # Get priority from user_api_key_auth_metadata in standard_logging_metadata + # This is where user_api_key_dict.metadata is stored during pre-call + user_api_key_auth_metadata = standard_logging_metadata.get("user_api_key_auth_metadata") or {} + key_priority: Optional[str] = user_api_key_auth_metadata.get("priority") + + # Get total tokens from response + total_tokens = 0 + rate_limit_type = self.v3_limiter.get_rate_limit_type() + + if isinstance(response_obj, ModelResponse): + _usage = getattr(response_obj, "usage", None) + if _usage and isinstance(_usage, Usage): + if rate_limit_type == "output": + total_tokens = _usage.completion_tokens + elif rate_limit_type == "input": + total_tokens = _usage.prompt_tokens + elif rate_limit_type == "total": + total_tokens = _usage.total_tokens + + if total_tokens == 0: + return + + # Create pipeline operations for token increments + pipeline_operations: List[RedisPipelineIncrementOperation] = [] + + # Model-wide token tracking (model_saturation_check) + model_token_key = self.v3_limiter.create_rate_limit_keys( + key="model_saturation_check", + value=model_group, + rate_limit_type="tokens", + ) + pipeline_operations.append( + RedisPipelineIncrementOperation( + key=model_token_key, + increment_value=total_tokens, + ttl=self.v3_limiter.window_size, + ) + ) + + # Priority-specific token tracking (priority_model) + # Determine priority key (same logic as _get_priority_allocation) + has_explicit_priority = ( + key_priority is not None + and litellm.priority_reservation is not None + and key_priority in litellm.priority_reservation + ) + + if has_explicit_priority and key_priority is not None: + priority_key = f"{model_group}:{key_priority}" + else: + priority_key = f"{model_group}:default_pool" + + priority_token_key = self.v3_limiter.create_rate_limit_keys( + key="priority_model", + value=priority_key, + rate_limit_type="tokens", + ) + pipeline_operations.append( + RedisPipelineIncrementOperation( + key=priority_token_key, + increment_value=total_tokens, + ttl=self.v3_limiter.window_size, + ) + ) + + # Execute token increments with TTL preservation + if pipeline_operations: + await self.v3_limiter.async_increment_tokens_with_ttl_preservation( + pipeline_operations=pipeline_operations, + parent_otel_span=litellm_parent_otel_span, + ) + + # Only log 'priority' if it's known safe; otherwise, redact. + SAFE_PRIORITIES = {"low", "medium", "high", "default"} + logged_priority = key_priority if key_priority in SAFE_PRIORITIES else "REDACTED" + verbose_proxy_logger.debug( + f"[Dynamic Rate Limiter] Incremented tokens by {total_tokens} for " + f"model={model_group}, priority={logged_priority}" + ) + + except Exception as e: + verbose_proxy_logger.exception( + f"Error in dynamic rate limiter success event: {str(e)}" + ) diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 44be6bbe656..3aa62eeeede 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -7,7 +7,6 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy._types import ( - CommonProxyErrors, GenerateKeyRequest, GenerateKeyResponse, KeyRequest, @@ -45,9 +44,13 @@ class KeyManagementEventHooks: ) from litellm.proxy.proxy_server import litellm_proxy_admin_name - await KeyManagementEventHooks._send_key_created_email( - response.model_dump(exclude_none=True) - ) + # Send email notification - non-blocking, independent operation + try: + await KeyManagementEventHooks._send_key_created_email( + response.model_dump(exclude_none=True) + ) + except Exception as e: + verbose_proxy_logger.warning(f"Failed to send key created email: {e}") # Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True if litellm.store_audit_logs is True: @@ -69,11 +72,17 @@ class KeyManagementEventHooks: ) ) ) - # store the generated key in the secret manager - await KeyManagementEventHooks._store_virtual_key_in_secret_manager( - secret_name=data.key_alias or f"virtual-key-{response.token_id}", - secret_token=response.key, - ) + + # Store the generated key in the secret manager - non-blocking, independent operation + try: + await KeyManagementEventHooks._store_virtual_key_in_secret_manager( + secret_name=data.key_alias or f"virtual-key-{response.token_id}", + secret_token=response.key, + ) + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to store virtual key in secret manager: {e}" + ) @staticmethod async def async_key_updated_hook( @@ -132,22 +141,31 @@ class KeyManagementEventHooks: ) from litellm.proxy.proxy_server import litellm_proxy_admin_name - # store the generated key in the secret manager + # Store the generated key in the secret manager - non-blocking, independent operation if data is not None and response.token_id is not None: - initial_secret_name = ( - existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}" - ) - await KeyManagementEventHooks._rotate_virtual_key_in_secret_manager( - current_secret_name=initial_secret_name, - new_secret_name=data.key_alias or f"virtual-key-{response.token_id}", - new_secret_value=response.key, - ) + try: + initial_secret_name = ( + existing_key_row.key_alias + or f"virtual-key-{existing_key_row.token}" + ) + await KeyManagementEventHooks._rotate_virtual_key_in_secret_manager( + current_secret_name=initial_secret_name, + new_secret_name=data.key_alias or f"virtual-key-{response.token_id}", + new_secret_value=response.key, + ) + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to rotate virtual key in secret manager: {e}" + ) - # send key rotated email if configured - await KeyManagementEventHooks._send_key_rotated_email( - response=response.model_dump(exclude_none=True), - existing_key_alias=existing_key_row.key_alias, - ) + # Send key rotated email if configured - non-blocking, independent operation + try: + await KeyManagementEventHooks._send_key_rotated_email( + response=response.model_dump(exclude_none=True), + existing_key_alias=existing_key_row.key_alias, + ) + except Exception as e: + verbose_proxy_logger.warning(f"Failed to send key rotated email: {e}") # store the audit log if litellm.store_audit_logs is True and existing_key_row.token is not None: @@ -324,66 +342,109 @@ class KeyManagementEventHooks: ) @staticmethod - async def _send_key_created_email(response: dict): + def _is_email_sending_enabled() -> bool: + """ + Check if email sending is enabled via v2 enterprise loggers or v0 alerting config. + + Returns True only if email is actually configured, preventing any email + processing when the user has not opted in. + """ + # Check v2 enterprise email loggers try: from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( BaseEmailLogger, ) - except ImportError: - raise Exception( - "Trying to use Email Hooks" - + CommonProxyErrors.missing_enterprise_package.value + + initialized_email_loggers = ( + litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=BaseEmailLogger + ) ) + if len(initialized_email_loggers) > 0: + return True + except ImportError: + pass + + # Check v0 alerting config + from litellm.proxy.proxy_server import general_settings + + if "email" in general_settings.get("alerting", []): + return True + + return False + + @staticmethod + async def _send_key_created_email(response: dict): + """ + Send key created email if email sending is enabled. + + This method is non-blocking - it will return silently if email is not + configured, and will log warnings instead of raising exceptions on failure. + """ + # Early exit if email is not enabled + if not KeyManagementEventHooks._is_email_sending_enabled(): + verbose_proxy_logger.debug( + "Email sending not enabled, skipping key created email" + ) + return from litellm.proxy.proxy_server import general_settings, proxy_logging_obj + ########################## + # v2 integration for emails (enterprise) + ########################## try: + from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( + BaseEmailLogger, + ) from litellm_enterprise.types.enterprise_callbacks.send_emails import ( SendKeyCreatedEmailEvent, ) + + initialized_email_loggers = ( + litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=BaseEmailLogger + ) + ) + if len(initialized_email_loggers) > 0: + event = SendKeyCreatedEmailEvent( + virtual_key=response.get("key", ""), + event="key_created", + event_group=Litellm_EntityType.KEY, + event_message="API Key Created", + token=response.get("token", ""), + spend=response.get("spend", 0.0), + max_budget=response.get("max_budget", 0.0), + user_id=response.get("user_id", None), + team_id=response.get("team_id", "Default Team"), + key_alias=response.get("key_alias", None), + ) + for email_logger in initialized_email_loggers: + if isinstance(email_logger, BaseEmailLogger): + await email_logger.send_key_created_email( + send_key_created_email_event=event, + ) + return except ImportError: - raise Exception( - "Trying to use Email Hooks" - + CommonProxyErrors.missing_enterprise_package.value - ) - - event = SendKeyCreatedEmailEvent( - virtual_key=response.get("key", ""), - event="key_created", - event_group=Litellm_EntityType.KEY, - event_message="API Key Created", - token=response.get("token", ""), - spend=response.get("spend", 0.0), - max_budget=response.get("max_budget", 0.0), - user_id=response.get("user_id", None), - team_id=response.get("team_id", "Default Team"), - key_alias=response.get("key_alias", None), - ) - - ########################## - # v2 integration for emails - ########################## - initialized_email_loggers = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=BaseEmailLogger - ) - ) - if len(initialized_email_loggers) > 0: - for email_logger in initialized_email_loggers: - if isinstance(email_logger, BaseEmailLogger): - await email_logger.send_key_created_email( - send_key_created_email_event=event, - ) + pass ########################## # v0 integration for emails ########################## - else: - if "email" not in general_settings.get("alerting", []): - raise ValueError( - "Email alerting not setup on config.yaml. Please set `alerting=['email']. \nDocs: https://docs.litellm.ai/docs/proxy/email`" - ) + if "email" in general_settings.get("alerting", []): + from litellm.proxy._types import WebhookEvent + event = WebhookEvent( + event="key_created", + event_group=Litellm_EntityType.KEY, + event_message="API Key Created", + token=response.get("token", ""), + spend=response.get("spend", 0.0), + max_budget=response.get("max_budget", 0.0), + user_id=response.get("user_id", None), + team_id=response.get("team_id", "Default Team"), + key_alias=response.get("key_alias", None), + ) # If user configured email alerting - send an Email letting their end-user know the key was created asyncio.create_task( proxy_logging_obj.slack_alerting_instance.send_key_created_or_user_invited_email( @@ -393,25 +454,39 @@ class KeyManagementEventHooks: @staticmethod async def _send_key_rotated_email(response: dict, existing_key_alias: Optional[str]): + """ + Send key rotated email if email sending is enabled. + + This method is non-blocking - it will return silently if email is not + configured, and will log warnings instead of raising exceptions on failure. + """ + # Early exit if email is not enabled + if not KeyManagementEventHooks._is_email_sending_enabled(): + verbose_proxy_logger.debug( + "Email sending not enabled, skipping key rotated email" + ) + return + try: from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( BaseEmailLogger, ) except ImportError: - raise Exception( - "Trying to use Email Hooks" - + CommonProxyErrors.missing_enterprise_package.value + # Enterprise package not installed - v0 doesn't support key rotated email + verbose_proxy_logger.debug( + "Enterprise package not installed, skipping key rotated email" ) + return try: from litellm_enterprise.types.enterprise_callbacks.send_emails import ( SendKeyRotatedEmailEvent, ) except ImportError: - raise Exception( - "Trying to use Email Hooks" - + CommonProxyErrors.missing_enterprise_package.value + verbose_proxy_logger.debug( + "Enterprise types not available, skipping key rotated email" ) + return event = SendKeyRotatedEmailEvent( virtual_key=response.get("key", ""), diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c462493de6c..6ef281e5dcd 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -65,6 +65,12 @@ for i = 1, #KEYS, 2 do table.insert(results, increment_value) -- counter else local counter = redis.call('INCR', counter_key) + -- This happens when window_key exists but counter_key doesn't (e.g., tokens key + -- created after requests key when both share the same window_key) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 then + redis.call('EXPIRE', counter_key, window_size) + end table.insert(results, window_start) -- window_start table.insert(results, counter) -- counter end diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 3afbbdd5a4b..9ff0fe6e590 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -20,6 +20,10 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.utils import handle_exception_on_proxy +from litellm.types.proxy.management_endpoints.common_daily_activity import ( + SpendAnalyticsPaginatedResponse, +) +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity router = APIRouter() @@ -673,4 +677,78 @@ async def list_end_user( str(e) ) ) - raise handle_exception_on_proxy(e) \ No newline at end of file + raise handle_exception_on_proxy(e) + +@router.get( + "/customer/daily/activity", + tags=["Customer Management"], + dependencies=[Depends(user_api_key_auth)], + response_model=SpendAnalyticsPaginatedResponse, +) +@router.get( + "/end_user/daily/activity", + tags=["Customer Management"], + include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], +) +async def get_customer_daily_activity( + end_user_ids: Optional[str] = None, + start_date: Optional[str] = None, + end_date: Optional[str] = None, + model: Optional[str] = None, + api_key: Optional[str] = None, + page: int = 1, + page_size: int = 10, + exclude_end_user_ids: Optional[str] = None, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + + """ + Get daily activity for specific organizations or all accessible organizations. + """ + from litellm.proxy.proxy_server import ( + prisma_client, + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + # Parse comma-separated ids + end_user_ids_list = end_user_ids.split(",") if end_user_ids else None + exclude_end_user_ids_list: Optional[List[str]] = None + if exclude_end_user_ids: + exclude_end_user_ids_list = ( + exclude_end_user_ids.split(",") if exclude_end_user_ids else None + ) + + + # Fetch organization aliases for metadata + where_condition = {} + if end_user_ids_list: + where_condition["user_id"] = {"in": list(end_user_ids_list)} + end_user_aliases = await prisma_client.db.litellm_endusertable.find_many( + where=where_condition + ) + end_user_alias_metadata = { + e.user_id: {"alias": e.alias} + for e in end_user_aliases + } + + # Query daily activity for organizations + return await get_daily_activity( + prisma_client=prisma_client, + table_name="litellm_dailyenduserspend", + entity_id_field="end_user_id", + entity_id=end_user_ids_list, + entity_metadata_field=end_user_alias_metadata, + exclude_entity_ids=exclude_end_user_ids_list, + start_date=start_date, + end_date=end_date, + model=model, + api_key=api_key, + page=page, + page_size=page_size, + ) \ No newline at end of file diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 4b62e490a82..9009ce8995b 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -31,6 +31,7 @@ from litellm.proxy._types import ( LiteLLM_ManagementEndpoint_MetadataFields_Premium, LiteLLM_ModelTable, LiteLLM_OrganizationTable, + LiteLLM_OrganizationTableWithMembers, LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj, @@ -1051,7 +1052,7 @@ async def fetch_and_validate_organization( organization_row = await prisma_client.db.litellm_organizationtable.find_unique( where={"organization_id": organization_id}, - include={"litellm_budget_table": True, "users": True}, + include={"litellm_budget_table": True, "members": True}, ) if organization_row is None: @@ -1064,7 +1065,7 @@ async def fetch_and_validate_organization( validate_team_org_change( team=LiteLLM_TeamTable(**existing_team_row.model_dump()), - organization=LiteLLM_OrganizationTable(**organization_row.model_dump()), + organization=LiteLLM_OrganizationTableWithMembers(**organization_row.model_dump()), llm_router=llm_router, ) @@ -1072,7 +1073,7 @@ async def fetch_and_validate_organization( def validate_team_org_change( - team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTable, llm_router: Router + team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTableWithMembers, llm_router: Router ) -> bool: """ Validate that a team can be moved to an organization. @@ -1123,7 +1124,7 @@ def validate_team_org_change( # Check if the team's user_id is a member of the org team_members = [m.user_id for m in team.members_with_roles] - org_members = [m.user_id for m in organization.users] if organization.users else [] + org_members = [m.user_id for m in organization.members] if organization.members else [] not_in_org = [ m for m in team_members diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index a447dad0b19..b2eda5a8d38 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -971,18 +971,23 @@ def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> di Used for google and vertex JS SDKs, and Azure passthrough Checks both 'tags' and 'x-litellm-tags' headers """ - # Initialize tags list if it doesn't exist - if "tags" not in metadata: - metadata["tags"] = [] - + tags_to_add = [] + # Check for 'tags' header first _tags = request.headers.get("tags") if _tags: - metadata["tags"].extend([tag.strip() for tag in _tags.split(",")]) + tags_to_add.extend([tag.strip() for tag in _tags.split(",")]) _tags = request.headers.get("x-litellm-tags") if _tags: - metadata["tags"].extend([tag.strip() for tag in _tags.split(",")]) + tags_to_add.extend([tag.strip() for tag in _tags.split(",")]) + + # Only add tags key if there are tags to add + if tags_to_add: + if "tags" not in metadata: + metadata["tags"] = [] + metadata["tags"].extend(tags_to_add) + return metadata diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 098cdb80e04..a33f56b0327 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,67 +1,16 @@ model_list: - - model_name: qwen-25vl-72b + - model_name: openai/gpt-4o-mini litellm_params: - model: bedrock/openai/arn:aws:bedrock:us-east-1:046319184608:imported-model/0m2lasirsp6z + model: openai/gpt-4o-mini + tpm: 1000 -guardrails: - - guardrail_name: "bedrock-pre-guard" - litellm_params: - guardrail: bedrock - mode: "pre_call" - guardrailIdentifier: ff6ujrregl1q - guardrailVersion: "DRAFT" - - -# like MCPs/vector stores -search_tools: - - search_tool_name: litellm-search - litellm_params: - search_provider: perplexity - api_key: os.environ/PERPLEXITYAI_API_KEY - - search_tool_name: firecrawl-search - litellm_params: - search_provider: firecrawl - api_key: os.environ/FIRECRAWL_API_KEY - litellm_settings: - max_end_user_budget_id: "2f6634cd-c631-4d3b-96c7-ad510ea06eaf" - # Comprehensive logging settings - store_audit_logs: true - verbose: true - log_level: "DEBUG" # Options: DEBUG, INFO, WARNING, ERROR - callbacks: ["s3_v2", "smtp_email"] - s3_callback_params: - s3_endpoint_url: "https://localhost:443" # Replace with your Minio server URL and port - s3_aws_access_key_id: "minioadmin" - s3_aws_secret_access_key: "minioadmin" - s3_region_name: "minio" # This can be any value for Minio - s3_bucket_name: "litellm-test" # Replace with your bucket name - s3_use_ssl: False - s3_verify: False - cache: True - cache_params: - type: local - drop_params: True + callbacks: ["dynamic_rate_limiter_v3"] + priority_reservation: + "prod": 0.9 # 90% reserved for production + "dev": 0.1 # 10% reserved for development -general_settings: - store_prompts_in_spend_logs: True - pass_through_endpoints: - - path: "/special/rerank" - target: "https://api.cohere.com/v1/rerank" - headers: - Authorization: "Bearer os.environ/COHERE_API_KEY" - guardrails: - bedrock-pre-guard: - request_fields: ["documents[*].text"] -vector_store_registry: - - vector_store_name: "bedrock-litellm-website-knowledgebase" - litellm_params: - vector_store_id: "T37J8R4WTM" - custom_llm_provider: "bedrock" - vector_store_description: "Bedrock vector store for the Litellm website knowledgebase" - vector_store_metadata: - source: "https://www.litellm.com/docs" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c1a5fd6c924..2caaec47243 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -45,7 +45,10 @@ from litellm.constants import ( LITELLM_SETTINGS_SAFE_DB_OVERRIDES, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy.common_utils.callback_utils import normalize_callback_names +from litellm.proxy.common_utils.callback_utils import ( + normalize_callback_names, + process_callback, +) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.types.utils import ( ModelResponse, @@ -54,7 +57,6 @@ from litellm.types.utils import ( TokenCountResponse, ) from litellm.utils import load_credentials_from_list -from litellm.proxy.common_utils.callback_utils import process_callback if TYPE_CHECKING: from aiohttp import ClientSession @@ -168,8 +170,8 @@ from litellm.constants import ( PROXY_BUDGET_RESCHEDULER_MIN_TIME, ) from litellm.exceptions import RejectedRequestError -from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_litellm_metadata_from_kwargs, @@ -613,7 +615,7 @@ async def proxy_shutdown_event(): await jwt_handler.close() if db_writer_client is not None: - await db_writer_client.close() + await db_writer_client.close() # type: ignore[reportGeneralTypeIssues] # flush remaining langfuse logs if "langfuse" in litellm.success_callback: @@ -792,7 +794,7 @@ async def proxy_startup_event(app: FastAPI): except Exception as e: verbose_proxy_logger.error(f"Error closing shared aiohttp session: {e}") - await proxy_shutdown_event() + await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues] app = FastAPI( @@ -802,7 +804,7 @@ app = FastAPI( description=_description, version=version, root_path=server_root_path, # check if user passed root path, FastAPI defaults this value to "" - lifespan=proxy_startup_event, + lifespan=proxy_startup_event, # type: ignore[reportGeneralTypeIssues] ) vertex_live_passthrough_vertex_base = VertexBase() @@ -8330,9 +8332,9 @@ async def login(request: Request): # noqa: PLR0915 # Generate JWT token import jwt - jwt_token = jwt.encode( # type: ignore + jwt_token = jwt.encode( cast(dict, returned_ui_token_object), - master_key, + cast(str, master_key), algorithm="HS256", ) @@ -8377,9 +8379,9 @@ async def login_v2(request: Request): # noqa: PLR0915 import jwt - jwt_token = jwt.encode( # type: ignore + jwt_token = jwt.encode( cast(dict, returned_ui_token_object), - master_key, + cast(str, master_key), algorithm="HS256", ) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index a23775f122d..ddd41ca0b1d 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2124,17 +2124,67 @@ "litellm_provider": "oci", "credential_fields": [ { - "key": "api_key", - "label": "API Key", - "placeholder": null, - "tooltip": null, + "key": "oci_user", + "label": "OCI User OCID", + "placeholder": "ocid1.user.oc1..aaaaaaaaexample", + "tooltip": "The OCID of the user making the API call", "required": true, - "field_type": "password", + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "oci_fingerprint", + "label": "OCI Key Fingerprint", + "placeholder": "aa:bb:cc:dd:ee:ff:00:11:22:33:44:55:66:77:88:99", + "tooltip": "The fingerprint of the API signing key", + "required": true, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "oci_tenancy", + "label": "OCI Tenancy OCID", + "placeholder": "ocid1.tenancy.oc1..aaaaaaaaexample", + "tooltip": "The OCID of your tenancy", + "required": true, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "oci_region", + "label": "OCI Region", + "placeholder": "us-ashburn-1", + "tooltip": "The OCI region identifier (e.g., us-ashburn-1, eu-frankfurt-1)", + "required": true, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "oci_compartment_id", + "label": "OCI Compartment OCID", + "placeholder": "ocid1.compartment.oc1..aaaaaaaaexample", + "tooltip": "The OCID of the compartment containing the GenAI resources", + "required": true, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "oci_key", + "label": "OCI Private Key (PEM)", + "placeholder": "-----BEGIN RSA PRIVATE KEY-----\n...\n-----END RSA PRIVATE KEY-----", + "tooltip": "The full PEM-encoded private key content for API signing. Paste the entire key including BEGIN/END markers.", + "required": true, + "field_type": "textarea", "options": null, "default_value": null } ], - "default_model_placeholder": "oci/xai.grok-4" + "default_model_placeholder": "oci/xai.grok-3" }, { "provider": "OVHCLOUD", diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 583c493adc6..40f00437e58 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -465,6 +465,34 @@ model LiteLLM_DailyOrganizationSpend { @@index([mcp_namespaced_tool_name]) } +// Track daily end user (customer) spend metrics per model and key +model LiteLLM_DailyEndUserSpend { + id String @id @default(uuid()) + end_user_id String? + date String + api_key String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + @@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name]) + @@index([date]) + @@index([end_user_id]) + @@index([api_key]) + @@index([model]) + @@index([mcp_namespaced_tool_name]) +} + // Track daily team spend metrics per model and key model LiteLLM_DailyTeamSpend { id String @id @default(uuid()) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 57af32339f2..9359c20c67e 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -39,6 +39,7 @@ from litellm.types.responses.main import ( GenericResponseOutputItem, GenericResponseOutputItemContentAnnotation, OutputFunctionToolCall, + OutputImageGenerationCall, OutputText, ) from litellm.types.utils import ( @@ -830,9 +831,9 @@ class LiteLLMCompletionResponsesConfig: def _transform_chat_completion_choices_to_responses_output( chat_completion_response: ModelResponse, choices: List[Choices], - ) -> List[Union[GenericResponseOutputItem, OutputFunctionToolCall]]: + ) -> List[Union[GenericResponseOutputItem, OutputFunctionToolCall, OutputImageGenerationCall]]: responses_output: List[ - Union[GenericResponseOutputItem, OutputFunctionToolCall] + Union[GenericResponseOutputItem, OutputFunctionToolCall, OutputImageGenerationCall] ] = [] responses_output.extend( @@ -881,28 +882,130 @@ class LiteLLMCompletionResponsesConfig: ] return [] + @staticmethod + def _extract_image_generation_output_items( + chat_completion_response: ModelResponse, + choice: Choices, + ) -> List[OutputImageGenerationCall]: + """ + Extract image generation outputs from a choice that contains images. + + Transforms message.images from chat completion format: + { + 'image_url': {'url': 'data:image/png;base64,iVBORw0...'}, + 'type': 'image_url', + 'index': 0 + } + + To Responses API format: + { + 'type': 'image_generation_call', + 'id': 'img_...', + 'status': 'completed', + 'result': 'iVBORw0...' # Pure base64 without data: prefix + } + """ + image_generation_items: List[OutputImageGenerationCall] = [] + + images = getattr(choice.message, 'images', []) + if not images: + return image_generation_items + + for idx, image_item in enumerate(images): + # Extract base64 from data URL + image_url = image_item.get('image_url', {}).get('url', '') + base64_data = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(image_url) + + if base64_data: + image_generation_items.append( + OutputImageGenerationCall( + type="image_generation_call", + id=f"{chat_completion_response.id}_img_{idx}", + status=LiteLLMCompletionResponsesConfig._map_finish_reason_to_image_generation_status( + choice.finish_reason + ), + result=base64_data, + ) + ) + + return image_generation_items + + @staticmethod + def _map_finish_reason_to_image_generation_status( + finish_reason: Optional[str], + ) -> Literal["in_progress", "completed", "incomplete", "failed"]: + """ + Map finish_reason to image generation status. + + Image generation status only supports: in_progress, completed, incomplete, failed + (does not support: cancelled, queued like general ResponsesAPIStatus) + """ + if finish_reason == "stop": + return "completed" + elif finish_reason == "length": + return "incomplete" + elif finish_reason in ["content_filter", "error"]: + return "failed" + else: + # Default to completed for other cases + return "completed" + + @staticmethod + def _extract_base64_from_data_url(data_url: str) -> Optional[str]: + """ + Extract pure base64 string from a data URL. + + Input: 'data:image/png;base64,iVBORw0KGgoAAAANS...' + Output: 'iVBORw0KGgoAAAANS...' + + If input is already pure base64 (no prefix), return as-is. + """ + if not data_url: + return None + + # Check if it's a data URL with prefix + if data_url.startswith('data:'): + # Split by comma to separate prefix from base64 data + parts = data_url.split(',', 1) + if len(parts) == 2: + return parts[1] # Return the base64 part + return None + else: + # Already pure base64 + return data_url + @staticmethod def _extract_message_output_items( chat_completion_response: ModelResponse, choices: List[Choices], - ) -> List[GenericResponseOutputItem]: - message_output_items = [] + ) -> List[Union[GenericResponseOutputItem, OutputImageGenerationCall]]: + message_output_items: List[Union[GenericResponseOutputItem, OutputImageGenerationCall]] = [] for choice in choices: - message_output_items.append( - GenericResponseOutputItem( - type="message", - id=chat_completion_response.id, - status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( - choice.finish_reason - ), - role=choice.message.role, - content=[ - LiteLLMCompletionResponsesConfig._transform_chat_message_to_response_output_text( - choice.message - ) - ], + # Check if message has images (image generation) + if hasattr(choice.message, 'images') and choice.message.images: + # Extract image generation output + image_generation_items = LiteLLMCompletionResponsesConfig._extract_image_generation_output_items( + chat_completion_response=chat_completion_response, + choice=choice, + ) + message_output_items.extend(image_generation_items) + else: + # Regular message output + message_output_items.append( + GenericResponseOutputItem( + type="message", + id=chat_completion_response.id, + status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( + choice.finish_reason + ), + role=choice.message.role, + content=[ + LiteLLMCompletionResponsesConfig._transform_chat_message_to_response_output_text( + choice.message + ) + ], + ) ) - ) return message_output_items @staticmethod diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 9fb97d47d17..377ba546988 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -76,6 +76,7 @@ from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject from litellm.types.responses.main import ( GenericResponseOutputItem, OutputFunctionToolCall, + OutputImageGenerationCall, ) FileContent = Union[IO[bytes], bytes, PathLike] @@ -1071,7 +1072,7 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): object: Optional[str] = None output: Union[ List[Union[ResponseOutputItem, Dict]], - List[Union[GenericResponseOutputItem, OutputFunctionToolCall]], + List[Union[GenericResponseOutputItem, OutputFunctionToolCall, OutputImageGenerationCall]], ] parallel_tool_calls: Optional[bool] = None temperature: Optional[float] = None diff --git a/litellm/types/proxy/public_endpoints/public_endpoints.py b/litellm/types/proxy/public_endpoints/public_endpoints.py index 7edf05dc947..aca58e36921 100644 --- a/litellm/types/proxy/public_endpoints/public_endpoints.py +++ b/litellm/types/proxy/public_endpoints/public_endpoints.py @@ -16,7 +16,7 @@ class ProviderCredentialField(BaseModel): placeholder: Optional[str] = None tooltip: Optional[str] = None required: bool = False - field_type: Literal["text", "password", "select", "upload"] = "text" + field_type: Literal["text", "password", "select", "upload", "textarea"] = "text" options: Optional[List[str]] = None default_value: Optional[str] = None diff --git a/litellm/types/responses/main.py b/litellm/types/responses/main.py index 25ba4d0a075..7d0620af23a 100644 --- a/litellm/types/responses/main.py +++ b/litellm/types/responses/main.py @@ -36,6 +36,15 @@ class OutputFunctionToolCall(BaseLiteLLMOpenAIResponseObject): status: Literal["in_progress", "completed", "incomplete"] +class OutputImageGenerationCall(BaseLiteLLMOpenAIResponseObject): + """An image generation call output""" + + type: Literal["image_generation_call"] + id: str + status: Literal["in_progress", "completed", "incomplete", "failed"] + result: Optional[str] # Base64 encoded image data (without data:image prefix) + + class GenericResponseOutputItem(BaseLiteLLMOpenAIResponseObject): """ Generic response API output item diff --git a/litellm/types/services.py b/litellm/types/services.py index 78cb0cb3ceb..bb34f817f29 100644 --- a/litellm/types/services.py +++ b/litellm/types/services.py @@ -1,10 +1,11 @@ import enum -from litellm._uuid import uuid from typing import List, Optional from pydantic import BaseModel, Field from typing_extensions import TypedDict +from litellm._uuid import uuid + class ServiceMetrics(enum.Enum): COUNTER = "counter" @@ -33,6 +34,8 @@ class ServiceTypes(str, enum.Enum): # daily spend update queue - actual transaction events IN_MEMORY_DAILY_SPEND_UPDATE_QUEUE = "in_memory_daily_spend_update_queue" REDIS_DAILY_SPEND_UPDATE_QUEUE = "redis_daily_spend_update_queue" + REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE = "redis_daily_end_user_spend_update_queue" + REDIS_DAILY_ORG_SPEND_UPDATE_QUEUE = "redis_daily_org_spend_update_queue" REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE = "redis_daily_team_spend_update_queue" REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE = "redis_daily_tag_spend_update_queue" # spend update queue - current spend of key, user, team @@ -87,6 +90,9 @@ DEFAULT_SERVICE_CONFIGS = { ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE.value: { "metrics": [ServiceMetrics.GAUGE] }, + ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE.value: { + "metrics": [ServiceMetrics.GAUGE] + }, ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE.value: { "metrics": [ServiceMetrics.GAUGE] }, diff --git a/litellm/utils.py b/litellm/utils.py index 0db84d3f5b9..244eedc9cb4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7023,11 +7023,13 @@ class ProviderConfigManager: """ # Check JSON providers FIRST - from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry if JSONProviderRegistry.exists(provider.value): provider_config = JSONProviderRegistry.get(provider.value) + if provider_config is None: + raise ValueError(f"Provider {provider.value} not found") return create_config_class(provider_config)() if ( @@ -7050,6 +7052,8 @@ class ProviderConfigManager: return litellm.DatabricksConfig() elif litellm.LlmProviders.XAI == provider: return litellm.XAIChatConfig() + elif litellm.LlmProviders.ZAI == provider: + return litellm.ZAIChatConfig() elif litellm.LlmProviders.LAMBDA_AI == provider: return litellm.LambdaAIChatConfig() elif litellm.LlmProviders.LLAMA == provider: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 634ea6dc48a..c26aada0a50 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5164,6 +5164,19 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "azure_ai/mistral-large-3": { + "input_cost_per_token": 5e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 256000, + "max_output_tokens": 8191, + "max_tokens": 8191, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://azure.microsoft.com/en-us/blog/introducing-mistral-large-3-in-microsoft-foundry-open-capable-and-ready-for-production-workloads/", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure_ai/mistral-medium-2505": { "input_cost_per_token": 4e-07, "litellm_provider": "azure_ai", @@ -12134,6 +12147,7 @@ "max_tokens": 65536, "mode": "image_generation", "output_cost_per_image": 0.134, + "output_cost_per_image_token": 1.2e-04, "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", @@ -13871,6 +13885,7 @@ "max_tokens": 65536, "mode": "image_generation", "output_cost_per_image": 0.134, + "output_cost_per_image_token": 1.2e-04, "output_cost_per_token": 1.2e-05, "rpm": 1000, "tpm": 4000000, @@ -18745,6 +18760,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/mistral-large-3": { + "input_cost_per_token": 5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 256000, + "max_output_tokens": 8191, + "max_tokens": 8191, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.mistral.ai/models/mistral-large-3-25-12", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/mistral-medium": { "input_cost_per_token": 2.7e-06, "litellm_provider": "mistral", @@ -25774,6 +25804,7 @@ "max_tokens": 65536, "mode": "image_generation", "output_cost_per_image": 0.134, + "output_cost_per_image_token": 1.2e-04, "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" diff --git a/pyproject.toml b/pyproject.toml index cb2003dab2d..1b3588bc95f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,7 +59,7 @@ websockets = {version = "^15.0.1", optional = true} boto3 = {version = "1.36.0", optional = true} redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} mcp = {version = "^1.21.2", optional = true, python = ">=3.10"} -litellm-proxy-extras = {version = "0.4.9", optional = true} +litellm-proxy-extras = {version = "0.4.10", optional = true} rich = {version = "13.7.1", optional = true} litellm-enterprise = {version = "0.1.23", optional = true} diskcache = {version = "^5.6.1", optional = true} diff --git a/requirements.txt b/requirements.txt index 7c9d10481a0..ffdcad5c008 100644 --- a/requirements.txt +++ b/requirements.txt @@ -44,7 +44,7 @@ sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests cryptography==44.0.1 tzdata==2025.1 # IANA time zone database -litellm-proxy-extras==0.4.9 # for proxy extras - e.g. prisma migrations +litellm-proxy-extras==0.4.10 # for proxy extras - e.g. prisma migrations ### LITELLM PACKAGE DEPENDENCIES python-dotenv==1.0.1 # for env tiktoken==0.8.0 # for calculating usage diff --git a/schema.prisma b/schema.prisma index 583c493adc6..40f00437e58 100644 --- a/schema.prisma +++ b/schema.prisma @@ -465,6 +465,34 @@ model LiteLLM_DailyOrganizationSpend { @@index([mcp_namespaced_tool_name]) } +// Track daily end user (customer) spend metrics per model and key +model LiteLLM_DailyEndUserSpend { + id String @id @default(uuid()) + end_user_id String? + date String + api_key String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + @@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name]) + @@index([date]) + @@index([end_user_id]) + @@index([api_key]) + @@index([model]) + @@index([mcp_namespaced_tool_name]) +} + // Track daily team spend metrics per model and key model LiteLLM_DailyTeamSpend { id String @id @default(uuid()) diff --git a/tests/basic_proxy_startup_tests/test_basic_proxy_startup.py b/tests/basic_proxy_startup_tests/test_basic_proxy_startup.py index 43250af44e0..62dd3297da7 100644 --- a/tests/basic_proxy_startup_tests/test_basic_proxy_startup.py +++ b/tests/basic_proxy_startup_tests/test_basic_proxy_startup.py @@ -25,7 +25,10 @@ async def test_health_and_chat_completion(): readiness_response = await response.json() assert readiness_response["status"] == "connected" +<<<<<<< HEAD +======= +>>>>>>> parent of d89990e0c5 (Add license metadata to health/readiness endpoint. (#15997)) # Test liveness endpoint async with session.get("http://0.0.0.0:4000/health/liveness") as response: assert response.status == 200 diff --git a/tests/litellm/llms/oci/chat/test_oci_chat_transformation.py b/tests/litellm/llms/oci/chat/test_oci_chat_transformation.py new file mode 100644 index 00000000000..c23f4501a2c --- /dev/null +++ b/tests/litellm/llms/oci/chat/test_oci_chat_transformation.py @@ -0,0 +1,199 @@ +""" +Tests for OCI Chat Transformation module. + +These tests verify the OCI credential handling, particularly the PEM key +normalization logic for handling different newline formats. +""" + +import os +import sys +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.oci.chat.transformation import OCIChatConfig +from litellm.llms.oci.common_utils import OCIError + + +@pytest.fixture +def config(): + return OCIChatConfig() + + +class TestOCIKeyNormalization: + """Tests for OCI private key content normalization.""" + + def test_oci_key_with_escaped_newlines(self, config): + """Test that escaped newlines (\\n) are converted to actual newlines.""" + # Simulate PEM content with escaped newlines (as would come from JSON/UI input) + escaped_pem = "-----BEGIN RSA PRIVATE KEY-----\\nMIIEowIBAAKCAQEA...\\n-----END RSA PRIVATE KEY-----" + + optional_params = { + "oci_user": "ocid1.user.oc1..test", + "oci_fingerprint": "aa:bb:cc:dd", + "oci_tenancy": "ocid1.tenancy.oc1..test", + "oci_region": "us-ashburn-1", + "oci_key": escaped_pem, + } + + # We can't fully test signing without a real key, but we can verify + # the error message indicates the key was processed (not a type error) + with pytest.raises(Exception) as exc_info: + config._sign_with_manual_credentials( + headers={}, + optional_params=optional_params, + request_data={"test": "data"}, + api_base="https://test.oci.oraclecloud.com/api", + ) + + # The error should be about key format/loading, not about type + # This confirms the string was processed and newlines were normalized + error_message = str(exc_info.value) + assert "must be a string" not in error_message.lower() + + def test_oci_key_with_crlf_newlines(self, config): + """Test that Windows-style CRLF newlines are normalized to LF.""" + # Simulate PEM content with CRLF newlines + crlf_pem = "-----BEGIN RSA PRIVATE KEY-----\r\nMIIEowIBAAKCAQEA...\r\n-----END RSA PRIVATE KEY-----" + + optional_params = { + "oci_user": "ocid1.user.oc1..test", + "oci_fingerprint": "aa:bb:cc:dd", + "oci_tenancy": "ocid1.tenancy.oc1..test", + "oci_region": "us-ashburn-1", + "oci_key": crlf_pem, + } + + with pytest.raises(Exception) as exc_info: + config._sign_with_manual_credentials( + headers={}, + optional_params=optional_params, + request_data={"test": "data"}, + api_base="https://test.oci.oraclecloud.com/api", + ) + + error_message = str(exc_info.value) + assert "must be a string" not in error_message.lower() + + def test_oci_key_rejects_non_string_type(self, config): + """Test that non-string oci_key values raise OCIError.""" + optional_params = { + "oci_user": "ocid1.user.oc1..test", + "oci_fingerprint": "aa:bb:cc:dd", + "oci_tenancy": "ocid1.tenancy.oc1..test", + "oci_region": "us-ashburn-1", + "oci_key": {"invalid": "dict"}, # Wrong type + } + + with pytest.raises(OCIError) as exc_info: + config._sign_with_manual_credentials( + headers={}, + optional_params=optional_params, + request_data={"test": "data"}, + api_base="https://test.oci.oraclecloud.com/api", + ) + + assert exc_info.value.status_code == 400 + assert "must be a string" in str(exc_info.value.message) + assert "dict" in str(exc_info.value.message) + + def test_oci_key_rejects_list_type(self, config): + """Test that list oci_key values raise OCIError.""" + optional_params = { + "oci_user": "ocid1.user.oc1..test", + "oci_fingerprint": "aa:bb:cc:dd", + "oci_tenancy": "ocid1.tenancy.oc1..test", + "oci_region": "us-ashburn-1", + "oci_key": ["invalid", "list"], # Wrong type + } + + with pytest.raises(OCIError) as exc_info: + config._sign_with_manual_credentials( + headers={}, + optional_params=optional_params, + request_data={"test": "data"}, + api_base="https://test.oci.oraclecloud.com/api", + ) + + assert exc_info.value.status_code == 400 + assert "must be a string" in str(exc_info.value.message) + assert "list" in str(exc_info.value.message) + + +class TestOCIValidateEnvironment: + """Tests for OCI environment validation.""" + + def test_missing_required_credentials_raises_error(self, config): + """Test that missing required credentials raise an error.""" + with pytest.raises(Exception) as exc_info: + config.validate_environment( + headers={}, + model="oci/xai.grok-3", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, # No credentials provided + litellm_params={}, + api_key=None, + api_base=None, + ) + + error_message = str(exc_info.value) + assert "oci_user" in error_message + assert "oci_fingerprint" in error_message + assert "oci_tenancy" in error_message + + def test_validate_environment_with_all_credentials(self, config): + """Test that validation passes with all required credentials.""" + headers = config.validate_environment( + headers={}, + model="oci/xai.grok-3", + messages=[{"role": "user", "content": "Hello"}], + optional_params={ + "oci_user": "ocid1.user.oc1..test", + "oci_fingerprint": "aa:bb:cc:dd", + "oci_tenancy": "ocid1.tenancy.oc1..test", + "oci_region": "us-ashburn-1", + "oci_compartment_id": "ocid1.compartment.oc1..test", + "oci_key": "-----BEGIN RSA PRIVATE KEY-----\ntest\n-----END RSA PRIVATE KEY-----", + }, + litellm_params={}, + api_key=None, + api_base=None, + ) + + assert headers["content-type"] == "application/json" + assert "user-agent" in headers + + +class TestOCIGetCompleteUrl: + """Tests for OCI URL generation.""" + + def test_get_complete_url_default_region(self, config): + """Test URL generation with default region.""" + url = config.get_complete_url( + api_base=None, + api_key=None, + model="oci/xai.grok-3", + optional_params={}, + litellm_params={}, + stream=False, + ) + + assert "us-ashburn-1" in url + assert "inference.generativeai" in url + assert "/20231130/actions/chat" in url + + def test_get_complete_url_custom_region(self, config): + """Test URL generation with custom region.""" + url = config.get_complete_url( + api_base=None, + api_key=None, + model="oci/xai.grok-3", + optional_params={"oci_region": "eu-frankfurt-1"}, + litellm_params={}, + stream=False, + ) + + assert "eu-frankfurt-1" in url + assert "inference.generativeai" in url diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index c75edba8500..a679fa4e4cb 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -1065,7 +1065,7 @@ def test_parse_content_for_reasoning(content, expected_reasoning, expected_conte ("gemini/gemini-1.5-pro", True), ("predibase/llama3-8b-instruct", True), ("gpt-3.5-turbo", False), - ("groq/llama-3.3-70b-versatile", True), + ("groq/llama-3.3-70b-versatile", False), ], ) def test_supports_response_schema(model, expected_bool): diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py index 44acd60ad49..d27626dd447 100644 --- a/tests/mcp_tests/test_mcp_logging.py +++ b/tests/mcp_tests/test_mcp_logging.py @@ -109,7 +109,11 @@ async def test_mcp_cost_tracking(): # Add assertions assert response is not None - response_list = list(response) # Convert iterable to list + # Handle CallToolResult - access .content for the list of content items + if isinstance(response, CallToolResult): + response_list = response.content + else: + response_list = list(response) # Convert iterable to list for backward compatibility assert len(response_list) == 1 assert isinstance(response_list[0], TextContent) assert response_list[0].text == "Test response" @@ -238,8 +242,8 @@ async def test_mcp_cost_tracking_per_tool(): assert response1 is not None assert response2 is not None - response_list_1 = list(response1) - response_list_2 = list(response2) + response_list_1 = list(response1.content) + response_list_2 = list(response2.content) assert len(response_list_1) == 1 assert len(response_list_2) == 1 diff --git a/tests/proxy_admin_ui_tests/e2e_ui_tests/view_internal_user.spec.ts b/tests/proxy_admin_ui_tests/e2e_ui_tests/view_internal_user.spec.ts index 8be5ff0c540..832832d8ae8 100644 --- a/tests/proxy_admin_ui_tests/e2e_ui_tests/view_internal_user.spec.ts +++ b/tests/proxy_admin_ui_tests/e2e_ui_tests/view_internal_user.spec.ts @@ -39,9 +39,8 @@ test("view internal user page", async ({ page }) => { const rowCount = await page.locator("tbody tr").count(); expect(rowCount).toBeGreaterThan(0); - const userIdHeader = page.locator("th", { hasText: "User ID" }); - page.screenshot({ path: "test-results/user_id_header.png" }); - await expect(userIdHeader).toBeVisible(); + const userIdHeader = await page.locator("th", { hasText: "User ID" }); + await expect(userIdHeader).toBeVisible({ timeout: 10000 }); // test pagination // Wait for pagination controls to be visible diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/proxy_unit_tests/test_jwt.py index ed75a56361a..b8df16d135b 100644 --- a/tests/proxy_unit_tests/test_jwt.py +++ b/tests/proxy_unit_tests/test_jwt.py @@ -1389,7 +1389,9 @@ async def test_custom_validate_called(): jwt_handler = MagicMock() jwt_handler.litellm_jwtauth = MagicMock( - custom_validate=mock_custom_validate, allowed_routes=["/chat/completions"] + custom_validate=mock_custom_validate, + allowed_routes=["/chat/completions"], + oidc_userinfo_enabled=False, ) jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "test_user"}) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3d3dd98f490..65e3dbec8bd 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -438,9 +438,9 @@ def test_string_cost_values_edge_cases(): # Expected costs: # Prompt: 1000 * 1e-6 + 100 * 0 (invalid string becomes 0) - # Completion: 500 * 2e-6 (text_tokens == completion_tokens, so is_text_tokens_total=True, no separate audio cost) + # Completion: 500 * 2e-6 + 50 * 2e-6 (audio tokens fall back to base cost when output_cost_per_audio_token is None) expected_prompt_cost = 1000 * 1e-6 - expected_completion_cost = 500 * 2e-6 + expected_completion_cost = 500 * 2e-6 + 50 * 2e-6 assert round(prompt_cost, 12) == round(expected_prompt_cost, 12) assert round(completion_cost, 12) == round(expected_completion_cost, 12) @@ -720,6 +720,72 @@ def test_service_tier_fallback_pricing(): assert abs(std_cost[1] - expected_standard_completion) < 1e-10, f"Standard completion cost mismatch: {std_cost[1]} vs {expected_standard_completion}" +def test_gemini_image_generation_cost_with_zero_text_tokens(): + """ + Test that image_tokens are correctly costed when text_tokens=0. + + Reproduces issue #17410: completion_cost calculates incorrectly for + Gemini-3-pro-image model - image_tokens were treated as text tokens + when text_tokens=0. + + https://github.com/BerriAI/litellm/issues/17410 + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model = "gemini-3-pro-image-preview" + custom_llm_provider = "vertex_ai" + + # Usage from the issue: text_tokens=0, image_tokens=1120, reasoning_tokens=225 + usage = Usage( + completion_tokens=1345, + prompt_tokens=10, + total_tokens=1355, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=None, + audio_tokens=None, + reasoning_tokens=225, + rejected_prediction_tokens=None, + text_tokens=0, # This is the key: text_tokens=0 + image_tokens=1120, + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=None, cached_tokens=None, text_tokens=10, image_tokens=None + ), + ) + + model_cost_map = litellm.model_cost[model] + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + + # Expected costs: + # - text_tokens: 0 * output_cost_per_token = 0 + # - image_tokens: 1120 * output_cost_per_image_token = 1120 * 1.2e-04 = 0.1344 + # - reasoning_tokens: 225 * output_cost_per_token = 225 * 1.2e-05 = 0.0027 + # Total completion: ~0.1371 + + output_cost_per_image_token = model_cost_map.get("output_cost_per_image_token", 0) + output_cost_per_token = model_cost_map.get("output_cost_per_token", 0) + + expected_image_cost = 1120 * output_cost_per_image_token + expected_reasoning_cost = 225 * output_cost_per_token # reasoning uses base token cost + expected_completion_cost = expected_image_cost + expected_reasoning_cost + + # The bug was: all 1345 tokens were treated as text = 1345 * 1.2e-05 = 0.01614 + # Fixed: image_tokens use image pricing = ~0.137 + + assert completion_cost > 0.10, ( + f"Completion cost should be > $0.10 (image tokens are expensive), got ${completion_cost:.6f}. " + f"Bug: tokens may be incorrectly treated as text tokens." + ) + assert round(completion_cost, 4) == round(expected_completion_cost, 4), ( + f"Expected completion cost ${expected_completion_cost:.6f}, got ${completion_cost:.6f}" + ) + + def test_bedrock_anthropic_prompt_caching(): """Test Bedrock Anthropic models with prompt caching return correct costs.""" model = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 3456f69eebd..6a528fef8f0 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1079,3 +1079,86 @@ def test_has_special_delta_attribute( assert not initialized_custom_stream_wrapper._has_special_delta_attribute( delta_with_none, "audio" ) + + +def test_is_chunk_non_empty_with_empty_tool_calls( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """ + Test that is_chunk_non_empty returns False when tool_calls is an empty list. + + Regression test for https://github.com/BerriAI/litellm/issues/17425 + Empty tool_calls in delta should not be considered non-empty chunks. + """ + chunk = { + "id": "test-chunk-id", + "object": "chat.completion.chunk", + "created": 1741037890, + "model": "claude-sonnet-4-20250514", + "choices": [ + { + "index": 0, + "delta": { + "content": None, + "tool_calls": [], # Empty tool_calls list + }, + "logprobs": None, + "finish_reason": None, + } + ], + } + # Empty tool_calls should return False + assert ( + initialized_custom_stream_wrapper.is_chunk_non_empty( + completion_obj={}, # completion_obj has no tool_calls + model_response=ModelResponseStream(**chunk), + response_obj={}, + ) + is False + ) + + +def test_is_chunk_non_empty_with_valid_tool_calls( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """ + Test that is_chunk_non_empty returns True when tool_calls has valid entries. + + Companion test for https://github.com/BerriAI/litellm/issues/17425 + Non-empty tool_calls in delta should be considered non-empty chunks. + """ + chunk = { + "id": "test-chunk-id", + "object": "chat.completion.chunk", + "created": 1741037890, + "model": "claude-sonnet-4-20250514", + "choices": [ + { + "index": 0, + "delta": { + "content": None, + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "NYC"}', + }, + } + ], + }, + "logprobs": None, + "finish_reason": None, + } + ], + } + # Non-empty tool_calls should return True + assert ( + initialized_custom_stream_wrapper.is_chunk_non_empty( + completion_obj={}, + model_response=ModelResponseStream(**chunk), + response_obj={}, + ) + is True + ) diff --git a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py index 7de2294954c..b9324e4966f 100644 --- a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py +++ b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py @@ -290,3 +290,103 @@ class TestAnthropicBetaHeaderSupport: else: # If no beta headers, that's also fine assert True + + def test_converse_non_anthropic_model_no_anthropic_beta(self): + """Test that non-Anthropic models (e.g., Qwen) do NOT get anthropic_beta in additionalModelRequestFields. + + This is critical because non-Anthropic models on Bedrock will error with + "unknown variant anthropic_beta" if this field is included. + """ + config = AmazonConverseConfig() + # Even if headers contain anthropic-beta, non-Anthropic models should NOT get it + headers = {"anthropic-beta": "context-1m-2025-08-07,interleaved-thinking-2025-05-14"} + + # Test with Qwen model (using ARN format like the user's config) + result = config._transform_request_helper( + model="qwen.qwen3-coder-480b-a35b-v1:0", + system_content_blocks=[], + optional_params={}, + messages=[{"role": "user", "content": "Test"}], + headers=headers + ) + + additional_fields = result.get("additionalModelRequestFields", {}) + assert "anthropic_beta" not in additional_fields, ( + "anthropic_beta should NOT be added for non-Anthropic models like Qwen. " + "This field is only supported by Anthropic/Claude models on Bedrock." + ) + + def test_converse_llama_model_no_anthropic_beta(self): + """Test that Llama models do NOT get anthropic_beta in additionalModelRequestFields.""" + config = AmazonConverseConfig() + headers = {"anthropic-beta": "context-1m-2025-08-07"} + + result = config._transform_request_helper( + model="meta.llama3-2-11b-instruct-v1:0", + system_content_blocks=[], + optional_params={}, + messages=[{"role": "user", "content": "Test"}], + headers=headers + ) + + additional_fields = result.get("additionalModelRequestFields", {}) + assert "anthropic_beta" not in additional_fields, ( + "anthropic_beta should NOT be added for Llama models." + ) + + def test_converse_nova_model_no_anthropic_beta(self): + """Test that Amazon Nova models do NOT get anthropic_beta in additionalModelRequestFields.""" + config = AmazonConverseConfig() + headers = {"anthropic-beta": "computer-use-2024-10-22"} + + result = config._transform_request_helper( + model="amazon.nova-pro-v1:0", + system_content_blocks=[], + optional_params={}, + messages=[{"role": "user", "content": "Test"}], + headers=headers + ) + + additional_fields = result.get("additionalModelRequestFields", {}) + assert "anthropic_beta" not in additional_fields, ( + "anthropic_beta should NOT be added for Amazon Nova models." + ) + + def test_converse_anthropic_model_gets_anthropic_beta(self): + """Test that Anthropic models DO get anthropic_beta in additionalModelRequestFields.""" + config = AmazonConverseConfig() + headers = {"anthropic-beta": "context-1m-2025-08-07"} + + result = config._transform_request_helper( + model="anthropic.claude-3-5-sonnet-20241022-v2:0", + system_content_blocks=[], + optional_params={}, + messages=[{"role": "user", "content": "Test"}], + headers=headers + ) + + additional_fields = result.get("additionalModelRequestFields", {}) + assert "anthropic_beta" in additional_fields, ( + "anthropic_beta SHOULD be added for Anthropic models." + ) + assert "context-1m-2025-08-07" in additional_fields["anthropic_beta"] + + def test_converse_anthropic_model_with_cross_region_prefix(self): + """Test that Anthropic models with cross-region prefix still get anthropic_beta.""" + config = AmazonConverseConfig() + headers = {"anthropic-beta": "context-1m-2025-08-07"} + + # Model with 'us.' cross-region prefix + result = config._transform_request_helper( + model="us.anthropic.claude-3-5-sonnet-20241022-v2:0", + system_content_blocks=[], + optional_params={}, + messages=[{"role": "user", "content": "Test"}], + headers=headers + ) + + additional_fields = result.get("additionalModelRequestFields", {}) + assert "anthropic_beta" in additional_fields, ( + "anthropic_beta SHOULD be added for Anthropic models with cross-region prefix." + ) + assert "context-1m-2025-08-07" in additional_fields["anthropic_beta"] 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 b0aa464bf9d..7eef15cd4d2 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_model_info.py +++ b/tests/test_litellm/llms/ollama/test_ollama_model_info.py @@ -136,3 +136,327 @@ class TestOllamaModelInfo: models = info.get_models() # Default static ollama_models is ['llama2'], so expect ['ollama/llama2'] assert models == ["ollama/llama2"] + + +class TestOllamaAuthHeaders: + """Tests for Ollama authentication header handling in completion calls.""" + + def test_ollama_completion_with_api_key_adds_auth_header(self, monkeypatch): + """ + Test that when an api_key is provided to ollama completion, + the Authorization header is added with Bearer token format. + + This tests the bug fix where Ollama requests with API keys + were not including the Authorization header. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers that were passed to the completion call + captured_headers = {} + + def mock_completion(*args, **kwargs): + # Capture the headers that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion with ollama provider and api_key + litellm.completion( + model="ollama/llama2", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-api-key-12345", + api_base="http://localhost:11434" + ) + + # Verify that Authorization header was added + assert "Authorization" in captured_headers, \ + "Authorization header should be present when api_key is provided" + assert captured_headers["Authorization"] == "Bearer test-api-key-12345", \ + f"Authorization header should be 'Bearer test-api-key-12345', got {captured_headers.get('Authorization')}" + + except Exception as e: + pytest.fail(f"Ollama completion with api_key failed: {e}") + + def test_ollama_chat_completion_with_api_key_adds_auth_header(self, monkeypatch): + """ + Test that when an api_key is provided to ollama_chat completion, + the Authorization header is added with Bearer token format. + + This tests the bug fix for the ollama_chat provider variant. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers that were passed to the completion call + captured_headers = {} + + def mock_completion(*args, **kwargs): + # Capture the headers that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion with ollama_chat provider and api_key + litellm.completion( + model="ollama_chat/llama2", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-api-key-67890", + api_base="http://localhost:11434" + ) + + # Verify that Authorization header was added + assert "Authorization" in captured_headers, \ + "Authorization header should be present when api_key is provided" + assert captured_headers["Authorization"] == "Bearer test-api-key-67890", \ + f"Authorization header should be 'Bearer test-api-key-67890', got {captured_headers.get('Authorization')}" + + except Exception as e: + pytest.fail(f"Ollama_chat completion with api_key failed: {e}") + + def test_ollama_completion_without_api_key_no_auth_header(self, monkeypatch): + """ + Test that when no api_key is provided to ollama completion, + no Authorization header is added. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers that were passed to the completion call + captured_headers = {} + + def mock_completion(*args, **kwargs): + # Capture the headers that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion without api_key + litellm.completion( + model="ollama/llama2", + messages=[{"role": "user", "content": "Hello"}], + api_base="http://localhost:11434" + ) + + # Verify that Authorization header was NOT added + assert "Authorization" not in captured_headers, \ + "Authorization header should not be present when api_key is not provided" + + except Exception as e: + pytest.fail(f"Ollama completion without api_key failed: {e}") + + def test_ollama_completion_preserves_existing_auth_header(self, monkeypatch): + """ + Test that when an Authorization header is already present in headers, + it is not overwritten even if api_key is provided. + + This ensures the fix respects existing Authorization headers. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers that were passed to the completion call + captured_headers = {} + + def mock_completion(*args, **kwargs): + # Capture the headers that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion with both api_key and existing Authorization header + existing_auth = "Bearer existing-token" + litellm.completion( + model="ollama/llama2", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-api-key-should-not-be-used", + api_base="http://localhost:11434", + headers={"Authorization": existing_auth} + ) + + # Verify that existing Authorization header was preserved + assert "Authorization" in captured_headers, \ + "Authorization header should be present" + assert captured_headers["Authorization"] == existing_auth, \ + f"Existing Authorization header should be preserved, got {captured_headers.get('Authorization')}" + + except Exception as e: + pytest.fail(f"Ollama completion with existing auth header failed: {e}") + + def test_ollama_completion_with_ollama_com_api_base(self, monkeypatch): + """ + Test that when using https://ollama.com as api_base with an api_key, + the Authorization header is correctly added. + + This tests the real-world use case of using Ollama's hosted service. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers and api_base that were passed to the completion call + captured_headers = {} + captured_api_base = None + + def mock_completion(*args, **kwargs): + nonlocal captured_api_base + # Capture the headers and api_base that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + if 'api_base' in kwargs: + captured_api_base = kwargs['api_base'] + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion with ollama.com as api_base and api_key + litellm.completion( + model="ollama/qwen3-vl:235b-cloud", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-ollama-com-api-key", + api_base="https://ollama.com" + ) + + # Verify that Authorization header was added + assert "Authorization" in captured_headers, \ + "Authorization header should be present when using ollama.com with api_key" + assert captured_headers["Authorization"] == "Bearer test-ollama-com-api-key", \ + f"Authorization header should be 'Bearer test-ollama-com-api-key', got {captured_headers.get('Authorization')}" + + # Verify the api_base was passed correctly + assert captured_api_base == "https://ollama.com", \ + f"API base should be 'https://ollama.com', got {captured_api_base}" + + except Exception as e: + pytest.fail(f"Ollama completion with ollama.com api_base failed: {e}") + + def test_ollama_chat_completion_with_ollama_com_api_base(self, monkeypatch): + """ + Test that when using https://ollama.com as api_base with an api_key + for ollama_chat provider, the Authorization header is correctly added. + + This tests the real-world use case for the ollama_chat variant. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers and api_base that were passed to the completion call + captured_headers = {} + captured_api_base = None + + def mock_completion(*args, **kwargs): + nonlocal captured_api_base + # Capture the headers and api_base that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + if 'api_base' in kwargs: + captured_api_base = kwargs['api_base'] + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion with ollama.com as api_base and api_key + litellm.completion( + model="ollama_chat/qwen3-vl:235b-cloud", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-ollama-com-chat-key", + api_base="https://ollama.com" + ) + + # Verify that Authorization header was added + assert "Authorization" in captured_headers, \ + "Authorization header should be present when using ollama.com with api_key" + assert captured_headers["Authorization"] == "Bearer test-ollama-com-chat-key", \ + f"Authorization header should be 'Bearer test-ollama-com-chat-key', got {captured_headers.get('Authorization')}" + + # Verify the api_base was passed correctly + assert captured_api_base == "https://ollama.com", \ + f"API base should be 'https://ollama.com', got {captured_api_base}" + + except Exception as e: + pytest.fail(f"Ollama_chat completion with ollama.com api_base failed: {e}") + + def test_ollama_completion_with_ollama_com_without_api_key_fails_gracefully(self, monkeypatch): + """ + Test that when using https://ollama.com as api_base without an api_key, + no Authorization header is added (which would likely fail on the server side, + but we're testing the client behavior). + + This ensures we don't add empty or None Authorization headers. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers that were passed to the completion call + captured_headers = {} + + def mock_completion(*args, **kwargs): + # Capture the headers that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion with ollama.com but no api_key + litellm.completion( + model="ollama/llama2", + messages=[{"role": "user", "content": "Hello"}], + api_base="https://ollama.com" + ) + + # Verify that Authorization header was NOT added + assert "Authorization" not in captured_headers, \ + "Authorization header should not be present when api_key is not provided, even with ollama.com" + + except Exception as e: + pytest.fail(f"Ollama completion with ollama.com without api_key failed: {e}") diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 41afec9cd14..7d45ce4091a 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -390,13 +390,13 @@ def test_streaming_chunk_includes_reasoning_content(): ) -def test_streaming_chunk_with_tool_calls_includes_reasoning_content(): +def test_streaming_chunk_with_tool_calls_and_thought_includes_reasoning_content(): """ - Test for issue #16805: Ensure that when Gemini returns a streaming chunk with - tool calls AND thoughtSignature, the reasoning_content is included in the delta. + Test that when Gemini returns a streaming chunk with both thought: true parts + AND tool calls, the reasoning_content is correctly extracted from the thought parts. - Previously, thinking_blocks were only added to non-streaming responses, causing - reasoning_content to be missing in streaming mode when tools were enabled. + Per Google's docs: thought: true indicates reasoning content, NOT thoughtSignature. + thoughtSignature is just a token for multi-turn context preservation. """ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator, @@ -409,12 +409,16 @@ def test_streaming_chunk_with_tool_calls_includes_reasoning_content(): { "content": { "parts": [ + { + "text": "Let me think about how to get the time...", + "thought": True, # This indicates reasoning content + }, { "functionCall": { "name": "get_current_time", "args": {"timezone": "America/New_York"}, }, - "thoughtSignature": "EsEDCr4DAdHtim...", # Base64 signature + "thoughtSignature": "EsEDCr4DAdHtim...", # Just a token, not reasoning } ] }, @@ -433,8 +437,8 @@ def test_streaming_chunk_with_tool_calls_includes_reasoning_content(): ) streaming_chunk = iterator.chunk_parser(chunk) - # Verify that reasoning_content is present in the streaming delta - assert streaming_chunk.choices[0].delta.reasoning_content is not None + # Verify reasoning_content comes from the thought: true part + assert streaming_chunk.choices[0].delta.reasoning_content == "Let me think about how to get the time..." # Verify tool calls are also present assert streaming_chunk.choices[0].delta.tool_calls is not None @@ -442,6 +446,59 @@ def test_streaming_chunk_with_tool_calls_includes_reasoning_content(): assert streaming_chunk.choices[0].delta.tool_calls[0].function.name == "get_current_time" +def test_streaming_chunk_with_tool_calls_no_thought_no_reasoning_content(): + """ + Test that when Gemini returns tool calls with thoughtSignature but WITHOUT + thought: true, there is NO reasoning_content. + + This is a regression test for the bug where functionCall data was incorrectly + being placed into reasoning_content when thoughtSignature was present. + Per Google's docs: thoughtSignature is just a token for multi-turn, not reasoning. + """ + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + litellm_logging = MagicMock() + + chunk = { + "candidates": [ + { + "content": { + "parts": [ + { + "functionCall": { + "name": "get_current_time", + "args": {"timezone": "America/New_York"}, + }, + "thoughtSignature": "EsEDCr4DAdHtim...", # Just a token, NOT thought: true + } + ] + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 68, + "candidatesTokenCount": 120, + "totalTokenCount": 188, + }, + } + + iterator = ModelResponseIterator( + streaming_response=[], sync_stream=True, logging_obj=litellm_logging + ) + streaming_chunk = iterator.chunk_parser(chunk) + + # reasoning_content should be None - thoughtSignature alone does NOT mean reasoning + assert getattr(streaming_chunk.choices[0].delta, 'reasoning_content', None) is None + + # Tool calls should still work + assert streaming_chunk.choices[0].delta.tool_calls is not None + assert len(streaming_chunk.choices[0].delta.tool_calls) == 1 + assert streaming_chunk.choices[0].delta.tool_calls[0].function.name == "get_current_time" + + def test_check_finish_reason(): finish_reason_mappings = VertexGeminiConfig.get_finish_reason_mapping() for k, v in finish_reason_mappings.items(): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py new file mode 100644 index 00000000000..6311b6d74b2 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py @@ -0,0 +1,87 @@ +""" +Tests for MCP metadata preservation. + +This module tests that tool metadata is preserved when creating prefixed tools, +which is critical for ChatGPT UI widget rendering. +""" + +import sys + +import pytest + +# Add the parent directory to the path so we can import litellm +sys.path.insert(0, "../../../../../") + +from mcp.types import Tool as MCPTool + +from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager +from litellm.proxy._types import MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +class TestMCPMetadataPreservation: + """Test that metadata is preserved when creating prefixed tools""" + + def test_create_prefixed_tools_preserves_metadata(self): + """Test that _create_prefixed_tools preserves metadata and _meta fields""" + manager = MCPServerManager() + + # Create a mock server + mock_server = MCPServer( + server_id="test-server-1", + name="test_server", + alias="test", + server_name="Test Server", + url="https://test-server.com/mcp", + transport=MCPTransport.http, + ) + + # Create a tool with metadata + tool_with_metadata = MCPTool( + name="hello_widget", + description="Display a greeting widget", + inputSchema={"type": "object", "properties": {}}, + ) + # Add metadata using setattr since MCPTool might not have it in the constructor + tool_with_metadata.metadata = { + "openai/outputTemplate": "ui://widget/hello.html", + "openai/widgetDescription": "A greeting widget", + } + tool_with_metadata._meta = { + "openai/toolInvocation/invoking": "Preparing greeting...", + } + + # Create prefixed tools + prefixed_tools = manager._create_prefixed_tools( + [tool_with_metadata], mock_server, add_prefix=True + ) + + # Verify + assert len(prefixed_tools) == 1 + prefixed_tool = prefixed_tools[0] + + # Check that name is prefixed + assert prefixed_tool.name == "test-hello_widget" + + # Check that metadata is preserved + assert hasattr(prefixed_tool, "metadata") + assert prefixed_tool.metadata == { + "openai/outputTemplate": "ui://widget/hello.html", + "openai/widgetDescription": "A greeting widget", + } + + # Check that _meta is preserved + assert hasattr(prefixed_tool, "_meta") + assert prefixed_tool._meta == { + "openai/toolInvocation/invoking": "Preparing greeting...", + } + + # Check that other fields are preserved + assert prefixed_tool.description == "Display a greeting widget" + assert prefixed_tool.inputSchema == {"type": "object", "properties": {}} + + + +if __name__ == "__main__": + pytest.main([__file__]) + diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 86e25c46c8e..7a6e5ad17f6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -10,7 +10,13 @@ sys.path.insert(0, "../../../../../") import httpx from mcp import ReadResourceResult, Resource -from mcp.types import GetPromptResult, Prompt, ResourceTemplate, TextResourceContents +from mcp.types import ( + GetPromptResult, + Prompt, + ResourceTemplate, + TextResourceContents, + Tool as MCPTool, +) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, @@ -995,10 +1001,11 @@ class TestMCPServerManager: manager._create_mcp_client = MagicMock(return_value=object()) # Tools returned upstream (unprefixed from provider) - upstream_tool = MagicMock() - upstream_tool.name = "send_email" - upstream_tool.description = "Send an email" - upstream_tool.inputSchema = {} + upstream_tool = MCPTool( + name="send_email", + description="Send an email", + inputSchema={}, + ) manager._fetch_tools_with_timeout = AsyncMock(return_value=[upstream_tool]) @@ -1025,14 +1032,16 @@ class TestMCPServerManager: ) # Input tools as would come from upstream - t1 = MagicMock() - t1.name = "create_issue" - t1.description = "" - t1.inputSchema = {} - t2 = MagicMock() - t2.name = "close_issue" - t2.description = "" - t2.inputSchema = {} + t1 = MCPTool( + name="create_issue", + description="", + inputSchema={}, + ) + t2 = MCPTool( + name="close_issue", + description="", + inputSchema={}, + ) # Do not add prefix in returned objects out_tools = manager._create_prefixed_tools([t1, t2], server, add_prefix=False) @@ -1066,10 +1075,11 @@ class TestMCPServerManager: manager.registry = {server.server_id: server} # Populate mapping (add_prefix value doesn't matter for mapping population) - base_tool = MagicMock() - base_tool.name = "create_zap" - base_tool.description = "" - base_tool.inputSchema = {} + base_tool = MCPTool( + name="create_zap", + description="", + inputSchema={}, + ) _ = manager._create_prefixed_tools([base_tool], server, add_prefix=False) # Unprefixed resolution diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 181d21b44f6..db6c318357c 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -572,4 +572,77 @@ async def test_add_spend_log_transaction_to_daily_org_transaction_skips_when_org org_id=None, ) - writer.daily_org_spend_update_queue.add_update.assert_not_called() \ No newline at end of file + writer.daily_org_spend_update_queue.add_update.assert_not_called() + + +@pytest.mark.asyncio +async def test_add_spend_log_transaction_to_daily_end_user_transaction_injects_end_user_id_and_queues_update(): + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + mock_prisma.get_request_status = MagicMock(return_value="success") + + end_user_id = "end-user-xyz" + payload = { + "request_id": "req-1", + "user": "test-user", + "end_user": end_user_id, + "startTime": "2024-01-01T12:00:00", + "api_key": "test-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "model_group": "gpt-4-group", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": 0.2, + "metadata": '{"usage_object": {}}', + } + + writer.daily_end_user_spend_update_queue.add_update = AsyncMock() + + await writer.add_spend_log_transaction_to_daily_end_user_transaction( + payload=payload, + prisma_client=mock_prisma, + ) + + writer.daily_end_user_spend_update_queue.add_update.assert_called_once() + + call_args = writer.daily_end_user_spend_update_queue.add_update.call_args[1] + update_dict = call_args["update"] + assert len(update_dict) == 1 + for key, transaction in update_dict.items(): + assert key == f"{end_user_id}_2024-01-01_test-key_gpt-4_openai" + assert transaction["end_user_id"] == end_user_id + assert transaction["date"] == "2024-01-01" + assert transaction["api_key"] == "test-key" + assert transaction["model"] == "gpt-4" + assert transaction["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_add_spend_log_transaction_to_daily_end_user_transaction_skips_when_end_user_id_missing(): + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + mock_prisma.get_request_status = MagicMock(return_value="success") + + payload = { + "request_id": "req-2", + "user": "test-user", + "startTime": "2024-01-01T12:00:00", + "api_key": "test-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "model_group": "gpt-4-group", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": 0.2, + "metadata": '{"usage_object": {}}', + } + + writer.daily_end_user_spend_update_queue.add_update = AsyncMock() + + await writer.add_spend_log_transaction_to_daily_end_user_transaction( + payload=payload, + prisma_client=mock_prisma, + ) + + writer.daily_end_user_spend_update_queue.add_update.assert_not_called() \ No newline at end of file diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index a6b2ae5b3a0..6450b9a63b0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -634,6 +634,228 @@ async def test_request_data_flows_to_apply_guardrail(): print("✓ request_data correctly passed to apply_guardrail") +@pytest.mark.asyncio +async def test_empty_content_handling(presidio_guardrail, mock_user_api_key, mock_cache): + """ + Test that Presidio handles empty content gracefully. + + This is common in tool/function calling where assistant messages have + empty content but include tool_calls. + + Bug fix: Previously crashed with: + TypeError: argument after ** must be a mapping, not str + """ + test_data = { + "messages": [ + {"role": "user", "content": "What is 2+2?"}, + { + "role": "assistant", + "content": "", # Empty content - common in tool calls + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": {"name": "calculator", "arguments": '{"a":2,"b":2}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_123", "content": "4"}, + ], + "model": "gpt-4", + } + + # Mock check_pii to simulate PII processing without needing Presidio API + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + # Empty text returns as-is (this is what our fix ensures) + return text + + presidio_guardrail.check_pii = mock_check_pii + + # This should not raise an exception + result = await presidio_guardrail.async_pre_call_hook( + user_api_key_dict=mock_user_api_key, + cache=mock_cache, + data=test_data, + call_type="completion", + ) + + assert result is not None + assert "messages" in result + # Verify messages are preserved + assert len(result["messages"]) == 3 + + print("✓ Empty content handling test passed") + + +@pytest.mark.asyncio +async def test_whitespace_only_content(presidio_guardrail, mock_user_api_key, mock_cache): + """ + Test that Presidio handles whitespace-only content gracefully. + + Whitespace-only content should be treated the same as empty content. + """ + test_data = { + "messages": [ + {"role": "user", "content": " "}, # Whitespace only + {"role": "assistant", "content": "\n\t "}, # Tabs and newlines + {"role": "user", "content": "Real question here"}, + ], + "model": "gpt-4", + } + + # Mock check_pii to simulate PII processing + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + return text + + presidio_guardrail.check_pii = mock_check_pii + + result = await presidio_guardrail.async_pre_call_hook( + user_api_key_dict=mock_user_api_key, + cache=mock_cache, + data=test_data, + call_type="completion", + ) + + assert result is not None + assert len(result["messages"]) == 3 + + print("✓ Whitespace-only content test passed") + + +@pytest.mark.asyncio +async def test_analyze_text_with_empty_string(): + """ + Test analyze_text method directly with empty string. + + Should return empty list without making API call to Presidio. + """ + presidio = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test:5002/", + presidio_anonymizer_api_base="http://test:5001/", + output_parse_pii=False, + ) + + # Test with empty string - should return immediately without API call + result = await presidio.analyze_text( + text="", + presidio_config=None, + request_data={}, + ) + assert result == [], "Empty text should return empty list" + + # Test with whitespace only - should return immediately + result = await presidio.analyze_text( + text=" \n\t ", + presidio_config=None, + request_data={}, + ) + assert result == [], "Whitespace-only text should return empty list" + + print("✓ analyze_text empty string test passed") + + +@pytest.mark.asyncio +async def test_analyze_text_error_dict_handling(): + """ + Test that analyze_text handles error dict responses from Presidio API. + + When Presidio returns {'error': 'No text provided'}, should handle gracefully + instead of crashing with TypeError. + """ + presidio = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://mock-presidio:5002/", + presidio_anonymizer_api_base="http://mock-presidio:5001/", + output_parse_pii=False, + ) + + # Mock the HTTP response to return error dict + class MockResponse: + async def json(self): + return {"error": "No text provided"} + async def __aenter__(self): + return self + async def __aexit__(self, *args): + pass + + class MockSession: + def post(self, *args, **kwargs): + return MockResponse() + async def __aenter__(self): + return self + async def __aexit__(self, *args): + pass + + with patch("aiohttp.ClientSession", return_value=MockSession()): + result = await presidio.analyze_text( + text="some text", + presidio_config=None, + request_data={}, + ) + # Should return empty list when error dict is received + assert result == [], "Error dict should be handled gracefully" + + print("✓ analyze_text error dict handling test passed") + + +@pytest.mark.asyncio +async def test_tool_calling_complete_scenario(presidio_guardrail, mock_user_api_key, mock_cache): + """ + Test complete tool calling scenario with PII in user message. + + This tests the real-world scenario where: + 1. User provides a query with PII + 2. Assistant responds with empty content + tool_calls + 3. Tool provides response + 4. Assistant provides final answer + """ + test_data = { + "messages": [ + { + "role": "user", + "content": "My email is john.doe@example.com. Can you look up my account?", + }, + { + "role": "assistant", + "content": "", # Empty - tool call + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": {"name": "lookup_account", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_abc", "content": "Account found"}, + {"role": "assistant", "content": "I found your account information."}, + ], + "model": "gpt-4", + } + + # Mock check_pii to simulate PII masking + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + if "john.doe@example.com" in text: + return text.replace("john.doe@example.com", "[EMAIL]") + return text + + presidio_guardrail.check_pii = mock_check_pii + + result = await presidio_guardrail.async_pre_call_hook( + user_api_key_dict=mock_user_api_key, + cache=mock_cache, + data=test_data, + call_type="completion", + ) + + assert result is not None + # Verify PII was masked in user message + assert "[EMAIL]" in result["messages"][0]["content"] + assert "john.doe@example.com" not in result["messages"][0]["content"] + # Verify other messages preserved + assert len(result["messages"]) == 4 + + print("✓ Tool calling complete scenario test passed") + + if __name__ == "__main__": # Run tests asyncio.run( diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 1f76013e237..d9e10e6f4b8 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -1323,3 +1323,187 @@ async def test_default_priority_shared_pool(): print(f" - 3 keys without priority share ONE pool: {desc_a[0]['value']}") print(f" - Shared pool limit: {desc_a[0]['rate_limit']['requests_per_unit']} RPM") print(f" - Explicit priority 'prod' uses separate pool: {desc_prod[0]['value']}") + + +@pytest.mark.asyncio +async def test_async_log_success_event_increments_by_actual_tokens(): + """ + Test that async_log_success_event increments token counters by actual token usage. + + This validates the fix for Bug 1: Token count was incrementing by 1 instead of actual usage. + The async_log_success_event should increment both model_saturation_check and priority_model + counters by the actual completion_tokens (when rate_limit_type=output). + """ + from unittest.mock import MagicMock + + from litellm.types.utils import ModelResponse, Usage + + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"dev": 0.1, "prod": 0.9} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "test-token-increment" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": 1000, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Track what gets incremented + increment_calls = [] + + async def mock_increment(pipeline_operations, parent_otel_span=None): + for op in pipeline_operations: + increment_calls.append({ + "key": op["key"], + "increment_value": op["increment_value"], + }) + + handler.v3_limiter.async_increment_tokens_with_ttl_preservation = mock_increment + + # Create mock response with 50 completion tokens + mock_response = MagicMock(spec=ModelResponse) + mock_response.usage = MagicMock(spec=Usage) + mock_response.usage.prompt_tokens = 10 + mock_response.usage.completion_tokens = 50 + mock_response.usage.total_tokens = 60 + + # Create kwargs with priority in user_api_key_auth_metadata + kwargs = { + "standard_logging_object": { + "metadata": { + "user_api_key_auth_metadata": {"priority": "dev"}, + }, + "model_group": model, + }, + "litellm_params": { + "metadata": {"model_group": model}, + }, + } + + with patch( + "litellm.proxy.common_utils.callback_utils.get_model_group_from_litellm_kwargs", + return_value=model, + ): + await handler.async_log_success_event( + kwargs=kwargs, + response_obj=mock_response, + start_time=None, + end_time=None, + ) + + # Verify increments happened with actual token count (50 completion tokens) + assert len(increment_calls) == 2, f"Expected 2 increment calls, got {len(increment_calls)}" + + # Both should increment by 50 (completion_tokens, since rate_limit_type defaults to 'output') + for call in increment_calls: + assert call["increment_value"] == 50, ( + f"Expected increment of 50 tokens, got {call['increment_value']} for key {call['key']}" + ) + + # Verify correct keys were used + keys = [call["key"] for call in increment_calls] + assert any("model_saturation_check" in k for k in keys), "Should increment model_saturation_check" + assert any("priority_model" in k and "dev" in k for k in keys), "Should increment priority_model with 'dev' priority" + + +@pytest.mark.asyncio +async def test_async_log_success_event_uses_team_priority_from_auth_metadata(): + """ + Test that async_log_success_event correctly retrieves priority from user_api_key_auth_metadata. + + This validates the fix where priority is retrieved from standard_logging_metadata.user_api_key_auth_metadata + instead of just standard_logging_metadata.priority. This is important for team-based priority inheritance. + """ + from unittest.mock import MagicMock + + from litellm.types.utils import ModelResponse, Usage + + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"team_priority": 0.8, "default": 0.2} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "test-team-priority" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": 1000, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Track incremented keys to verify priority is used correctly + incremented_keys = [] + + async def mock_increment(pipeline_operations, parent_otel_span=None): + for op in pipeline_operations: + incremented_keys.append(op["key"]) + + handler.v3_limiter.async_increment_tokens_with_ttl_preservation = mock_increment + + # Create mock response + mock_response = MagicMock(spec=ModelResponse) + mock_response.usage = MagicMock(spec=Usage) + mock_response.usage.prompt_tokens = 10 + mock_response.usage.completion_tokens = 20 + mock_response.usage.total_tokens = 30 + + # Simulate team metadata inheritance: priority is in user_api_key_auth_metadata + # This is how the proxy passes team metadata to the callback + kwargs = { + "standard_logging_object": { + "metadata": { + # Priority NOT at top level (this would fail before the fix) + # Priority IS in user_api_key_auth_metadata (team inheritance) + "user_api_key_auth_metadata": {"priority": "team_priority"}, + }, + "model_group": model, + }, + "litellm_params": { + "metadata": {"model_group": model}, + }, + } + + with patch( + "litellm.proxy.common_utils.callback_utils.get_model_group_from_litellm_kwargs", + return_value=model, + ): + await handler.async_log_success_event( + kwargs=kwargs, + response_obj=mock_response, + start_time=None, + end_time=None, + ) + + # Verify the priority_model key uses 'team_priority' (not 'default_pool') + priority_keys = [k for k in incremented_keys if "priority_model" in k] + assert len(priority_keys) == 1, f"Expected 1 priority_model key, got {len(priority_keys)}" + + # The key should contain 'team_priority', not 'default_pool' + assert "team_priority" in priority_keys[0], ( + f"Expected priority key to use 'team_priority' from user_api_key_auth_metadata, " + f"got key: {priority_keys[0]}" + ) + assert "default_pool" not in priority_keys[0], ( + f"Priority key should NOT use 'default_pool', should use team's priority. Got: {priority_keys[0]}" + ) diff --git a/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py b/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py new file mode 100644 index 00000000000..f731d9e298a --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py @@ -0,0 +1,130 @@ +""" +Tests for KeyManagementEventHooks. + +Validates that email and secret manager operations are independent and non-blocking. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks + + +class TestKeyManagementEventHooksIndependentOperations: + """Tests that email and secret manager operations are independent.""" + + @pytest.mark.asyncio + async def test_email_failure_does_not_block_secret_manager(self): + """ + Test that if email sending fails, secret manager operation still runs. + + This validates the independent operation design where one failure + does not block the other operation. + """ + secret_manager_called = {"called": False} + + # Mock the email method to raise an exception + async def mock_send_email_raises(*args, **kwargs): + raise Exception("Email service unavailable") + + # Mock the secret manager method to track if it was called + async def mock_store_secret(*args, **kwargs): + secret_manager_called["called"] = True + + # Create mock objects for the hook parameters + mock_data = MagicMock() + mock_data.key_alias = "test-key-alias" + + mock_response = MagicMock() + mock_response.model_dump.return_value = {"key": "sk-test", "token": "test-token"} + mock_response.model_dump_json.return_value = '{"key": "sk-test"}' + mock_response.token_id = "token-123" + mock_response.key = "sk-test-key" + + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.user_id = "user-123" + mock_user_api_key_dict.api_key = "api-key-123" + + with patch.object( + KeyManagementEventHooks, + "_send_key_created_email", + side_effect=mock_send_email_raises, + ), patch.object( + KeyManagementEventHooks, + "_store_virtual_key_in_secret_manager", + side_effect=mock_store_secret, + ), patch( + "litellm.store_audit_logs", False + ), patch( + "litellm.proxy.hooks.key_management_event_hooks.verbose_proxy_logger" + ): + # Should not raise even though email fails + await KeyManagementEventHooks.async_key_generated_hook( + data=mock_data, + response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + # Secret manager should have been called despite email failure + assert secret_manager_called["called"] is True + + @pytest.mark.asyncio + async def test_secret_manager_failure_does_not_block_email(self): + """ + Test that if secret manager fails, email operation still runs. + + This validates the independent operation design where one failure + does not block the other operation. + """ + email_called = {"called": False} + + # Mock the email method to track if it was called + async def mock_send_email(*args, **kwargs): + email_called["called"] = True + + # Mock the secret manager method to raise an exception + async def mock_store_secret_raises(*args, **kwargs): + raise Exception("Secret manager unavailable") + + # Create mock objects for the hook parameters + mock_data = MagicMock() + mock_data.key_alias = "test-key-alias" + + mock_response = MagicMock() + mock_response.model_dump.return_value = {"key": "sk-test", "token": "test-token"} + mock_response.model_dump_json.return_value = '{"key": "sk-test"}' + mock_response.token_id = "token-123" + mock_response.key = "sk-test-key" + + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.user_id = "user-123" + mock_user_api_key_dict.api_key = "api-key-123" + + with patch.object( + KeyManagementEventHooks, + "_send_key_created_email", + side_effect=mock_send_email, + ), patch.object( + KeyManagementEventHooks, + "_store_virtual_key_in_secret_manager", + side_effect=mock_store_secret_raises, + ), patch( + "litellm.store_audit_logs", False + ), patch( + "litellm.proxy.hooks.key_management_event_hooks.verbose_proxy_logger" + ): + # Should not raise even though secret manager fails + await KeyManagementEventHooks.async_key_generated_hook( + data=mock_data, + response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + # Email should have been called despite secret manager failure + assert email_called["called"] is True + diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 86a6ceec25e..25ff6f89427 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -1,4 +1,4 @@ -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import FastAPI, HTTPException, Request, status @@ -301,3 +301,99 @@ def test_customer_endpoints_error_schema_consistency(mock_prisma_client, mock_us for key in ["message", "type", "code"]: assert isinstance(error1[key], str), f"error1[{key}] should be a string" assert isinstance(error2[key], str), f"error2[{key}] should be a string" + + +@pytest.mark.asyncio +async def test_get_customer_daily_activity_admin_param_passing(monkeypatch): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints import customer_endpoints + from litellm.proxy.management_endpoints.customer_endpoints import ( + get_customer_daily_activity, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock( + return_value=[] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse") + get_daily_activity_mock = AsyncMock(return_value=mocked_response) + monkeypatch.setattr( + customer_endpoints, "get_daily_activity", get_daily_activity_mock + ) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1") + result = await get_customer_daily_activity( + end_user_ids="end-user-1,end-user-2", + start_date="2024-01-01", + end_date="2024-01-31", + model="gpt-4", + api_key="test-key", + page=2, + page_size=5, + exclude_end_user_ids="end-user-3", + user_api_key_dict=auth, + ) + + get_daily_activity_mock.assert_awaited_once() + kwargs = get_daily_activity_mock.call_args.kwargs + assert kwargs["table_name"] == "litellm_dailyenduserspend" + assert kwargs["entity_id_field"] == "end_user_id" + assert kwargs["entity_id"] == ["end-user-1", "end-user-2"] + assert kwargs["exclude_entity_ids"] == ["end-user-3"] + assert kwargs["start_date"] == "2024-01-01" + assert kwargs["end_date"] == "2024-01-31" + assert kwargs["model"] == "gpt-4" + assert kwargs["api_key"] == "test-key" + assert kwargs["page"] == 2 + assert kwargs["page_size"] == 5 + + assert result is mocked_response + + +@pytest.mark.asyncio +async def test_get_customer_daily_activity_with_end_user_aliases(monkeypatch): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints import customer_endpoints + from litellm.proxy.management_endpoints.customer_endpoints import ( + get_customer_daily_activity, + ) + + mock_prisma_client = AsyncMock() + mock_end_user1 = MagicMock() + mock_end_user1.user_id = "end-user-1" + mock_end_user1.alias = "Customer One" + mock_end_user2 = MagicMock() + mock_end_user2.user_id = "end-user-2" + mock_end_user2.alias = "Customer Two" + + mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock( + return_value=[mock_end_user1, mock_end_user2] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse") + get_daily_activity_mock = AsyncMock(return_value=mocked_response) + monkeypatch.setattr( + customer_endpoints, "get_daily_activity", get_daily_activity_mock + ) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1") + await get_customer_daily_activity( + end_user_ids="end-user-1,end-user-2", + start_date="2024-01-01", + end_date="2024-01-31", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_end_user_ids=None, + user_api_key_dict=auth, + ) + + kwargs = get_daily_activity_mock.call_args.kwargs + assert kwargs["entity_metadata_field"] == { + "end-user-1": {"alias": "Customer One"}, + "end-user-2": {"alias": "Customer Two"}, + } diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 06ec71a84f8..d096b5515a0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -16,7 +16,9 @@ sys.path.insert( ) # Adds the parent directory to the system path from litellm.proxy._types import UserAPIKeyAuth # Import UserAPIKeyAuth from litellm.proxy._types import ( + LiteLLM_OrganizationMembershipTable, LiteLLM_OrganizationTable, + LiteLLM_OrganizationTableWithMembers, LiteLLM_TeamTable, LitellmUserRoles, Member, @@ -95,7 +97,7 @@ async def test_validate_team_org_change_same_org_id(): team.members_with_roles = [] # Mock organization - organization = MagicMock(spec=LiteLLM_OrganizationTable) + organization = MagicMock(spec=LiteLLM_OrganizationTableWithMembers) organization.organization_id = org_id organization.models = [] organization.litellm_budget_table = MagicMock() @@ -108,7 +110,7 @@ async def test_validate_team_org_change_same_org_id(): organization.litellm_budget_table.rpm_limit = ( 50 # This would normally fail validation ) - organization.users = [] + organization.members = [] # Mock Router mock_router = MagicMock(spec=Router) @@ -126,6 +128,114 @@ async def test_validate_team_org_change_same_org_id(): mock_access_check.assert_not_called() # Ensure access check wasn't called +@pytest.mark.asyncio +async def test_validate_team_org_change_members_in_org(): + """ + Test that validate_team_org_change passes when team members are in organization.members. + + This tests the fix for issue #17552 where membership was incorrectly checked against + organization.users (deprecated) instead of organization.members (correct). + """ + team_org_id = "team-org-123" + new_org_id = "new-org-456" + user_id_1 = "user-123" + user_id_2 = "user-456" + + # Mock team with members + team = MagicMock(spec=LiteLLM_TeamTable) + team.organization_id = team_org_id + team.models = [] + team.max_budget = None + team.tpm_limit = None + team.rpm_limit = None + + # Create mock team members + team_member_1 = MagicMock() + team_member_1.user_id = user_id_1 + team_member_2 = MagicMock() + team_member_2.user_id = user_id_2 + team.members_with_roles = [team_member_1, team_member_2] + + # Mock organization with members (using LiteLLM_OrganizationMembershipTable structure) + organization = MagicMock(spec=LiteLLM_OrganizationTableWithMembers) + organization.organization_id = new_org_id + organization.models = [] + organization.litellm_budget_table = None + + # Create mock organization members - these should match team members + org_member_1 = MagicMock(spec=LiteLLM_OrganizationMembershipTable) + org_member_1.user_id = user_id_1 + org_member_2 = MagicMock(spec=LiteLLM_OrganizationMembershipTable) + org_member_2.user_id = user_id_2 + organization.members = [org_member_1, org_member_2] + + # Mock Router + mock_router = MagicMock(spec=Router) + + # Test should pass - all team members are in org members + result = validate_team_org_change( + team=team, organization=organization, llm_router=mock_router + ) + assert result is True + + +@pytest.mark.asyncio +async def test_validate_team_org_change_member_not_in_org(): + """ + Test that validate_team_org_change raises HTTPException when team members + are NOT in organization.members. + + This tests the fix for issue #17552 where membership was incorrectly checked against + organization.users (deprecated) instead of organization.members (correct). + """ + team_org_id = "team-org-123" + new_org_id = "new-org-456" + user_id_1 = "user-123" + user_id_2 = "user-456" + user_id_not_in_org = "user-not-in-org-789" + + # Mock team with members (including one not in org) + team = MagicMock(spec=LiteLLM_TeamTable) + team.organization_id = team_org_id + team.models = [] + team.max_budget = None + team.tpm_limit = None + team.rpm_limit = None + + # Create mock team members - user_id_not_in_org is not in the org + team_member_1 = MagicMock() + team_member_1.user_id = user_id_1 + team_member_2 = MagicMock() + team_member_2.user_id = user_id_not_in_org + team.members_with_roles = [team_member_1, team_member_2] + + # Mock organization with members (missing user_id_not_in_org) + organization = MagicMock(spec=LiteLLM_OrganizationTableWithMembers) + organization.organization_id = new_org_id + organization.models = [] + organization.litellm_budget_table = None + + # Create mock organization members - only user_id_1 and user_id_2 are members + org_member_1 = MagicMock(spec=LiteLLM_OrganizationMembershipTable) + org_member_1.user_id = user_id_1 + org_member_2 = MagicMock(spec=LiteLLM_OrganizationMembershipTable) + org_member_2.user_id = user_id_2 + organization.members = [org_member_1, org_member_2] + + # Mock Router + mock_router = MagicMock(spec=Router) + + # Test should fail - user_id_not_in_org is not in org members + with pytest.raises(HTTPException) as exc_info: + validate_team_org_change( + team=team, organization=organization, llm_router=mock_router + ) + + assert exc_info.value.status_code == 403 + assert "not a member of the organization" in str(exc_info.value.detail) + assert user_id_not_in_org in str(exc_info.value.detail) + + # Test for /team/permissions_list endpoint (GET) @pytest.mark.asyncio async def test_get_team_permissions_list_success(mock_db_client, mock_admin_auth): diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7b3157f8ee1..1c77a914960 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -91,6 +91,7 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") client = TestClient(app) response = client.post( diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_image_generation_output.py b/tests/test_litellm/responses/litellm_completion_transformation/test_image_generation_output.py new file mode 100644 index 00000000000..29f1063a070 --- /dev/null +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_image_generation_output.py @@ -0,0 +1,172 @@ +""" +Unit tests for Responses API image generation support + +Tests the fix for Issue #16227: +https://github.com/BerriAI/litellm/issues/16227 + +Verifies that image generation outputs are correctly transformed +from /chat/completions format to /responses API format. +""" +import pytest +from unittest.mock import Mock +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) +from litellm.types.responses.main import OutputImageGenerationCall +from litellm.types.utils import ModelResponse, Choices, Message + + +class TestExtractBase64FromDataUrl: + """Tests for _extract_base64_from_data_url helper function""" + + def test_extracts_base64_from_data_url(self): + """Should extract pure base64 from data URL with prefix""" + data_url = "data:image/png;base64,iVBORw0KGgoAAAANS" + result = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url( + data_url + ) + assert result == "iVBORw0KGgoAAAANS" + + def test_returns_base64_as_is_if_no_prefix(self): + """Should return base64 as-is if no data: prefix""" + pure_base64 = "iVBORw0KGgoAAAANS" + result = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url( + pure_base64 + ) + assert result == pure_base64 + + def test_handles_invalid_inputs(self): + """Should return None for empty/None/malformed inputs""" + assert LiteLLMCompletionResponsesConfig._extract_base64_from_data_url("") is None + assert LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(None) is None + assert LiteLLMCompletionResponsesConfig._extract_base64_from_data_url("data:image/png;base64") is None + + +class TestExtractImageGenerationOutputItems: + """Tests for _extract_image_generation_output_items function""" + + def test_extracts_images_correctly(self): + """Should extract OutputImageGenerationCall objects from images""" + mock_response = Mock(spec=ModelResponse) + mock_response.id = "test_123" + + mock_message = Mock(spec=Message) + mock_message.images = [ + {"image_url": {"url": "data:image/png;base64,IMG1"}, "type": "image_url", "index": 0}, + {"image_url": {"url": "data:image/jpeg;base64,IMG2"}, "type": "image_url", "index": 1}, + ] + + mock_choice = Mock(spec=Choices) + mock_choice.message = mock_message + mock_choice.finish_reason = "stop" + + result = LiteLLMCompletionResponsesConfig._extract_image_generation_output_items( + chat_completion_response=mock_response, + choice=mock_choice, + ) + + assert len(result) == 2 + assert result[0].type == "image_generation_call" + assert result[0].result == "IMG1" + assert result[1].result == "IMG2" + assert result[0].id == "test_123_img_0" + assert result[1].id == "test_123_img_1" + assert result[0].status == "completed" + + def test_returns_empty_for_no_images(self): + """Should return empty list if no images""" + mock_response = Mock(spec=ModelResponse) + mock_message = Mock(spec=Message) + mock_message.images = [] + + mock_choice = Mock(spec=Choices) + mock_choice.message = mock_message + mock_choice.finish_reason = "stop" + + result = LiteLLMCompletionResponsesConfig._extract_image_generation_output_items( + chat_completion_response=mock_response, + choice=mock_choice, + ) + + assert result == [] + + def test_maps_finish_reason_to_status(self): + """Should correctly map finish_reason to status""" + mock_response = Mock(spec=ModelResponse) + mock_response.id = "test_finish" + + mock_message = Mock(spec=Message) + mock_message.images = [ + {"image_url": {"url": "data:image/png;base64,TEST"}, "type": "image_url", "index": 0} + ] + + mock_choice = Mock(spec=Choices) + mock_choice.message = mock_message + mock_choice.finish_reason = "length" + + result = LiteLLMCompletionResponsesConfig._extract_image_generation_output_items( + chat_completion_response=mock_response, + choice=mock_choice, + ) + + assert result[0].status == "incomplete" + + +class TestExtractMessageOutputItemsIntegration: + """Integration tests for _extract_message_output_items with images""" + + def test_detects_images_and_creates_image_generation_call(self): + """Should detect images in message and create image_generation_call output""" + mock_response = Mock(spec=ModelResponse) + mock_response.id = "integration_test_123" + + mock_message = Mock(spec=Message) + mock_message.images = [ + { + "image_url": {"url": "data:image/png;base64,INTEGRATION_TEST"}, + "type": "image_url", + "index": 0, + } + ] + mock_message.role = "assistant" + mock_message.content = "Here's your image!" + + mock_choice = Mock(spec=Choices) + mock_choice.message = mock_message + mock_choice.finish_reason = "stop" + + result = LiteLLMCompletionResponsesConfig._extract_message_output_items( + chat_completion_response=mock_response, + choices=[mock_choice], + ) + + # Should return image_generation_call, NOT regular message + assert len(result) == 1 + assert isinstance(result[0], OutputImageGenerationCall) + assert result[0].type == "image_generation_call" + assert result[0].result == "INTEGRATION_TEST" + + def test_creates_regular_message_when_no_images(self): + """Should create regular GenericResponseOutputItem when no images""" + from litellm.types.responses.main import GenericResponseOutputItem + + mock_response = Mock(spec=ModelResponse) + mock_response.id = "no_images_123" + + mock_message = Mock(spec=Message) + # No images attribute or empty + mock_message.role = "assistant" + mock_message.content = "Just text, no images" + + mock_choice = Mock(spec=Choices) + mock_choice.message = mock_message + mock_choice.finish_reason = "stop" + + result = LiteLLMCompletionResponsesConfig._extract_message_output_items( + chat_completion_response=mock_response, + choices=[mock_choice], + ) + + assert len(result) == 1 + assert isinstance(result[0], GenericResponseOutputItem) + assert result[0].type == "message" diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 9907113c7eb..641db8b107d 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -324,6 +324,7 @@ "resolved": "https://registry.npmjs.org/@babel/core/-/core-7.28.5.tgz", "integrity": "sha512-e7jT4DxYvIDLk1ZHmU/m/mB19rex9sv0c2ftBtjSBv+kVM/902eh0fINUzD7UwLLNR+jU585GxUJ8/EBfAM5fw==", "license": "MIT", + "peer": true, "dependencies": { "@babel/code-frame": "^7.27.1", "@babel/generator": "^7.28.5", @@ -2185,6 +2186,7 @@ } ], "license": "MIT", + "peer": true, "engines": { "node": ">=18" }, @@ -2227,6 +2229,7 @@ } ], "license": "MIT", + "peer": true, "engines": { "node": ">=18" } @@ -2336,6 +2339,7 @@ "resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.0.tgz", "integrity": "sha512-8sLjZwK0R+JlxlYcTuVnyT2v+htpdrjDOKuMcOVdYjt52Lh8hWRYpxBPoKx/Zg+bcjc3wx6fmQevMmUztS/ccA==", "license": "MIT", + "peer": true, "dependencies": { "cssesc": "^3.0.0", "util-deprecate": "^1.0.2" @@ -2757,6 +2761,7 @@ "resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.0.tgz", "integrity": "sha512-8sLjZwK0R+JlxlYcTuVnyT2v+htpdrjDOKuMcOVdYjt52Lh8hWRYpxBPoKx/Zg+bcjc3wx6fmQevMmUztS/ccA==", "license": "MIT", + "peer": true, "dependencies": { "cssesc": "^3.0.0", "util-deprecate": "^1.0.2" @@ -5824,6 +5829,7 @@ "integrity": "sha512-o4PXJQidqJl82ckFaXUeoAW+XysPLauYI43Abki5hABd853iMhitooc6znOnczgbTYmEP6U6/y1ZyKAIsvMKGg==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@babel/code-frame": "^7.10.4", "@babel/runtime": "^7.12.5", @@ -6617,6 +6623,7 @@ "resolved": "https://registry.npmjs.org/@types/react/-/react-18.2.48.tgz", "integrity": "sha512-qboRCl6Ie70DQQG9hhNREz81jqC1cs9EVNcjQ1AU+jH6NFfSAhVVbrrY/+nSF+Bsk4AOwm9Qa61InvMCyV+H3w==", "license": "MIT", + "peer": true, "dependencies": { "@types/prop-types": "*", "@types/scheduler": "*", @@ -6639,6 +6646,7 @@ "integrity": "sha512-MEe3UeoENYVFXzoXEWsvcpg6ZvlrFNlOQ7EOsvhI3CfAXwzPfO8Qwuxd40nepsYKqyyVQnTdEfv68q91yLcKrQ==", "dev": true, "license": "MIT", + "peer": true, "peerDependencies": { "@types/react": "^18.0.0" } @@ -6835,6 +6843,7 @@ "integrity": "sha512-lJi3PfxVmo0AkEY93ecfN+r8SofEqZNGByvHAI3GBLrvt1Cw6H5k1IM02nSzu0RfUafr2EvFSw0wAsZgubNplQ==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@typescript-eslint/scope-manager": "8.47.0", "@typescript-eslint/types": "8.47.0", @@ -7496,6 +7505,7 @@ "integrity": "sha512-hGISOaP18plkzbWEcP/QvtRW1xDXF2+96HbEX6byqQhAUbiS5oH6/9JwW+QsQCIYON2bI6QZBF+2PvOmrRZ9wA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@vitest/utils": "3.2.4", "fflate": "^0.8.2", @@ -7724,6 +7734,7 @@ "resolved": "https://registry.npmjs.org/acorn/-/acorn-8.15.0.tgz", "integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==", "license": "MIT", + "peer": true, "bin": { "acorn": "bin/acorn" }, @@ -7813,6 +7824,7 @@ "resolved": "https://registry.npmjs.org/ajv/-/ajv-6.12.6.tgz", "integrity": "sha512-j3fVLgvTo527anyYyJOGTYJbG+vnnQYvE0m5mmkc1TK+nxAppkCLMIL0aZ4dblVCNoGShhm+kzE4ZUykBoMg4g==", "license": "MIT", + "peer": true, "dependencies": { "fast-deep-equal": "^3.1.1", "fast-json-stable-stringify": "^2.0.0", @@ -8666,6 +8678,7 @@ } ], "license": "MIT", + "peer": true, "dependencies": { "baseline-browser-mapping": "^2.8.25", "caniuse-lite": "^1.0.30001754", @@ -8990,6 +9003,7 @@ "resolved": "https://registry.npmjs.org/chevrotain/-/chevrotain-11.0.3.tgz", "integrity": "sha512-ci2iJH6LeIkvP9eJW6gpueU8cnZhv85ELY8w8WiFtNjMHA5ad6pQLaJo9mEly/9qUyCpvqX8/POVUTf18/HFdw==", "license": "Apache-2.0", + "peer": true, "dependencies": { "@chevrotain/cst-dts-gen": "11.0.3", "@chevrotain/gast": "11.0.3", @@ -9718,6 +9732,7 @@ "resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.0.tgz", "integrity": "sha512-8sLjZwK0R+JlxlYcTuVnyT2v+htpdrjDOKuMcOVdYjt52Lh8hWRYpxBPoKx/Zg+bcjc3wx6fmQevMmUztS/ccA==", "license": "MIT", + "peer": true, "dependencies": { "cssesc": "^3.0.0", "util-deprecate": "^1.0.2" @@ -10080,6 +10095,7 @@ "resolved": "https://registry.npmjs.org/cytoscape/-/cytoscape-3.33.1.tgz", "integrity": "sha512-iJc4TwyANnOGR1OmWhsS9ayRS3s+XQ185FmuHObThD+5AeJCakAAbWv8KimMTt08xCCLNgneQwFp+JRJOr9qGQ==", "license": "MIT", + "peer": true, "engines": { "node": ">=0.10" } @@ -10489,6 +10505,7 @@ "resolved": "https://registry.npmjs.org/d3-selection/-/d3-selection-3.0.0.tgz", "integrity": "sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==", "license": "ISC", + "peer": true, "engines": { "node": ">=12" } @@ -10663,6 +10680,7 @@ "resolved": "https://registry.npmjs.org/date-fns/-/date-fns-3.6.0.tgz", "integrity": "sha512-fRHTG8g/Gif+kSh50gaGEdToemgfj74aRX3swtiouboip5JDLAyDE9F11nHMIcvOaXeOC6D7SpNhi7uFyB7Uww==", "license": "MIT", + "peer": true, "funding": { "type": "github", "url": "https://github.com/sponsors/kossnocorp" @@ -11540,6 +11558,7 @@ "deprecated": "This version is no longer supported. Please see https://eslint.org/version-support for other options.", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@eslint-community/eslint-utils": "^4.2.0", "@eslint-community/regexpp": "^4.6.1", @@ -11725,6 +11744,7 @@ "integrity": "sha512-whOE1HFo/qJDyX4SnXzP4N6zOWn79WhnCUY/iDR0mPfQZO8wcYE4JClzI2oZrhBnnMUCBCHZhO6VQyoBU95mZA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@rtsao/scc": "^1.1.0", "array-includes": "^3.1.9", @@ -14988,6 +15008,7 @@ "integrity": "sha512-454TI39PeRDW1LgpyLPyURtB4Zx1tklSr6+OFOipsxGUH1WMTvk6C65JQdrj455+DP2uJ1+veBEHTGFKWVLFoA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@acemir/cssom": "^0.9.23", "@asamuzakjp/dom-selector": "^6.7.4", @@ -15142,12 +15163,12 @@ } }, "node_modules/jws": { - "version": "3.2.2", - "resolved": "https://registry.npmjs.org/jws/-/jws-3.2.2.tgz", - "integrity": "sha512-YHlZCB6lMTllWDtSPHz/ZXTsi8S00usEV6v1tjq8tOUZzw7DpSDWVXjXDre6ed1w/pd495ODpHZYSdkRTsa0HA==", + "version": "3.2.3", + "resolved": "https://registry.npmjs.org/jws/-/jws-3.2.3.tgz", + "integrity": "sha512-byiJ0FLRdLdSVSReO/U4E7RoEyOCKnEnEPMjq3HxWtvzLsV08/i5RQKsFVNkCldrCaPr2vDNAOMsfs8T/Hze7g==", "license": "MIT", "dependencies": { - "jwa": "^1.4.1", + "jwa": "^1.4.2", "safe-buffer": "^5.0.1" } }, @@ -15991,9 +16012,9 @@ } }, "node_modules/mdast-util-to-hast": { - "version": "13.2.0", - "resolved": "https://registry.npmjs.org/mdast-util-to-hast/-/mdast-util-to-hast-13.2.0.tgz", - "integrity": "sha512-QGYKEuUsYT9ykKBCMOEDLsU5JRObWQusAolFMeko/tYPufNkRffBAQjIE+99jbA87xv6FgmjLtwjh9wBWajwAA==", + "version": "13.2.1", + "resolved": "https://registry.npmjs.org/mdast-util-to-hast/-/mdast-util-to-hast-13.2.1.tgz", + "integrity": "sha512-cctsq2wp5vTsLIcaymblUriiTcZd0CwWtCbLvrOzYCDZoWyMNV8sZ7krj09FSnsiJi3WVsHLM4k6Dq/yaPyCXA==", "license": "MIT", "dependencies": { "@types/hast": "^3.0.0", @@ -18115,6 +18136,7 @@ "resolved": "https://registry.npmjs.org/moment/-/moment-2.30.1.tgz", "integrity": "sha512-uEmtNhbDOrWPFS+hdjFCBfy9f2YoyzRpwcl+DqpC6taX21FzsTLQVbMV/W7PzNSX6x/bhC1zA3c2UQ5NzH6how==", "license": "MIT", + "peer": true, "engines": { "node": "*" } @@ -19306,6 +19328,7 @@ } ], "license": "MIT", + "peer": true, "dependencies": { "nanoid": "^3.3.11", "picocolors": "^1.1.1", @@ -20318,6 +20341,7 @@ "resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.0.tgz", "integrity": "sha512-8sLjZwK0R+JlxlYcTuVnyT2v+htpdrjDOKuMcOVdYjt52Lh8hWRYpxBPoKx/Zg+bcjc3wx6fmQevMmUztS/ccA==", "license": "MIT", + "peer": true, "dependencies": { "cssesc": "^3.0.0", "util-deprecate": "^1.0.2" @@ -21812,6 +21836,7 @@ "resolved": "https://registry.npmjs.org/react/-/react-18.3.1.tgz", "integrity": "sha512-wS+hAgJShR0KhEvPJArfuPVN1+Hz1t0Y6n5jLrGQbkb4urgPE/0Rve+1kMB1v/oWgHgm4WIcV+i7F2pTVj+2iQ==", "license": "MIT", + "peer": true, "dependencies": { "loose-envify": "^1.1.0" }, @@ -21851,6 +21876,7 @@ "resolved": "https://registry.npmjs.org/react-dom/-/react-dom-18.3.1.tgz", "integrity": "sha512-5m4nQKp+rZRb09LNH59GM4BxTh9251/ylbKIbpe7TpGxfJ+9kv6BLkLBXIjjspbgbnIBNqlI23tRnTWT0snUIw==", "license": "MIT", + "peer": true, "dependencies": { "loose-envify": "^1.1.0", "scheduler": "^0.23.2" @@ -21908,6 +21934,7 @@ "resolved": "https://registry.npmjs.org/@docusaurus/react-loadable/-/react-loadable-6.0.0.tgz", "integrity": "sha512-YMMxTUQV/QFSnbgrP3tjDzLHRg7vsbMn8e9HAa8o/1iXoiomo48b7sk/kkmWEuWNDPJVlKSJRB6Y2fHqdJk+SQ==", "license": "MIT", + "peer": true, "dependencies": { "@types/react": "*" }, @@ -21973,6 +22000,7 @@ "resolved": "https://registry.npmjs.org/react-router/-/react-router-5.3.4.tgz", "integrity": "sha512-Ys9K+ppnJah3QuaRiLxk+jDWOR1MekYQrlytiXxC1RyfbdsZkS5pvKAzCCr031xHixZwpnsYNT5xysdFHQaYsA==", "license": "MIT", + "peer": true, "dependencies": { "@babel/runtime": "^7.12.13", "history": "^4.9.0", @@ -23010,8 +23038,7 @@ "version": "1.1.5", "resolved": "https://registry.npmjs.org/schema-dts/-/schema-dts-1.1.5.tgz", "integrity": "sha512-RJr9EaCmsLzBX2NDiO5Z3ux2BVosNZN5jo0gWgsyKvxKIUL5R3swNvoorulAeL9kLB0iTSX7V6aokhla2m7xbg==", - "license": "Apache-2.0", - "peer": true + "license": "Apache-2.0" }, "node_modules/schema-utils": { "version": "4.3.3", @@ -23037,6 +23064,7 @@ "resolved": "https://registry.npmjs.org/ajv/-/ajv-8.17.1.tgz", "integrity": "sha512-B/gBuNg5SiMTrPkC+A2+cW0RszwxYmn6VYxB/inlBStS5nx6xHIt/ehKRhIMhqusl7a8LjQoZnjCs5vhwxOQ1g==", "license": "MIT", + "peer": true, "dependencies": { "fast-deep-equal": "^3.1.3", "fast-uri": "^3.0.1", @@ -24268,6 +24296,7 @@ "resolved": "https://registry.npmjs.org/tailwindcss/-/tailwindcss-3.4.18.tgz", "integrity": "sha512-6A2rnmW5xZMdw11LYjhcI5846rt9pbLSabY5XPxo+XWdxwZaFEn47Go4NzFiHu9sNNmr/kXivP1vStfvMaK1GQ==", "license": "MIT", + "peer": true, "dependencies": { "@alloc/quick-lru": "^5.2.0", "arg": "^5.0.2", @@ -24578,6 +24607,7 @@ "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz", "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "license": "MIT", + "peer": true, "engines": { "node": ">=12" }, @@ -24790,7 +24820,8 @@ "version": "2.8.1", "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", - "license": "0BSD" + "license": "0BSD", + "peer": true }, "node_modules/type-check": { "version": "0.4.0", @@ -24923,6 +24954,7 @@ "integrity": "sha512-pXWcraxM0uxAS+tN0AG/BF2TyqmHO014Z070UsJ+pFvYuRSq8KH8DmWpnbXe0pEPDHXZV3FcAbJkijJ5oNEnWw==", "devOptional": true, "license": "Apache-2.0", + "peer": true, "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" @@ -25465,6 +25497,7 @@ "integrity": "sha512-NL8jTlbo0Tn4dUEXEsUg8KeyG/Lkmc4Fnzb8JXN/Ykm9G4HNImjtABMJgkQoVjOBN/j2WAwDTRytdqJbZsah7w==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "esbuild": "^0.25.0", "fdir": "^6.5.0", @@ -25581,6 +25614,7 @@ "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "dev": true, "license": "MIT", + "peer": true, "engines": { "node": ">=12" }, @@ -25594,6 +25628,7 @@ "integrity": "sha512-LUCP5ev3GURDysTWiP47wRRUpLKMOfPh+yKTx3kVIEiu5KOMeqzpnYNsKyOoVrULivR8tLcks4+lga33Whn90A==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@types/chai": "^5.2.2", "@vitest/expect": "3.2.4", @@ -25799,6 +25834,7 @@ "resolved": "https://registry.npmjs.org/webpack/-/webpack-5.103.0.tgz", "integrity": "sha512-HU1JOuV1OavsZ+mfigY0j8d1TgQgbZ6M+J75zDkpEAwYeXjWSqrGJtgnPblJjd/mAyTNQ7ygw0MiKOn6etz8yw==", "license": "MIT", + "peer": true, "dependencies": { "@types/eslint-scope": "^3.7.7", "@types/estree": "^1.0.8", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.ts new file mode 100644 index 00000000000..10cbedc04d3 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.ts @@ -0,0 +1,41 @@ +import { allEndUsersCall } from "@/components/networking"; +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { all_admin_roles } from "@/utils/roles"; + +const customersKeys = createQueryKeys("customers"); + +export interface Customer { + user_id: string; + alias?: string | null; + spend: number; + blocked: boolean; + allowed_model_region?: string | null; + default_model?: string | null; + budget_id?: string | null; + litellm_budget_table?: { + budget_id: string; + max_budget?: number | null; + soft_budget?: number | null; + max_parallel_requests?: number | null; + tpm_limit?: number | null; + rpm_limit?: number | null; + model_max_budget?: Record | null; + budget_duration?: string | null; + budget_reset_at?: string | null; + created_at: string; + created_by: string; + updated_at: string; + updated_by: string; + } | null; +} + +export type CustomersResponse = Customer[]; + +export const useCustomers = (accessToken: string | null, userRole: string | null) => { + return useQuery({ + queryKey: customersKeys.list({}), + queryFn: async () => await allEndUsersCall(accessToken!), + enabled: Boolean(accessToken) && all_admin_roles.includes(userRole || ""), + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts new file mode 100644 index 00000000000..9198450a63d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts @@ -0,0 +1,82 @@ +/* @vitest-environment jsdom */ +import { renderHook } from "@testing-library/react"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import useAuthorized from "./useAuthorized"; + +const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock } = vi.hoisted(() => ({ + replaceMock: vi.fn(), + clearTokenCookiesMock: vi.fn(), + getProxyBaseUrlMock: vi.fn(() => "http://proxy.example"), +})); + +vi.mock("next/navigation", () => ({ + useRouter: () => ({ + replace: replaceMock, + }), +})); + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: getProxyBaseUrlMock, +})); + +vi.mock("@/utils/cookieUtils", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + clearTokenCookies: clearTokenCookiesMock, + }; +}); + +const createJwt = (payload: Record) => { + const base64Url = btoa(JSON.stringify(payload)).replace(/=+$/, "").replace(/\+/g, "-").replace(/\//g, "_"); + return `eyJhbGciOiJub25lIn0.${base64Url}.signature`; +}; + +const clearCookie = () => { + document.cookie = "token=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=/;"; +}; + +describe("useAuthorized", () => { + afterEach(() => { + replaceMock.mockReset(); + clearTokenCookiesMock.mockReset(); + getProxyBaseUrlMock.mockClear(); + clearCookie(); + }); + + it("should decode the token and expose user details", () => { + const token = createJwt({ + key: "api-key-123", + user_id: "user-1", + user_email: "user@example.com", + user_role: "app_admin", + premium_user: true, + disabled_non_admin_personal_key_creation: false, + login_method: "username_password", + }); + document.cookie = `token=${token}; path=/;`; + + const { result } = renderHook(() => useAuthorized()); + + expect(result.current.token).toBe(token); + expect(result.current.accessToken).toBe("api-key-123"); + expect(result.current.userId).toBe("user-1"); + expect(result.current.userEmail).toBe("user@example.com"); + expect(result.current.userRole).toBe("Admin"); + expect(result.current.premiumUser).toBe(true); + expect(result.current.disabledPersonalKeyCreation).toBe(false); + expect(result.current.showSSOBanner).toBe(true); + expect(replaceMock).not.toHaveBeenCalled(); + }); + + it("should clear cookies and redirect on an invalid token", () => { + document.cookie = "token=invalid-token; path=/;"; + + const { result } = renderHook(() => useAuthorized()); + + expect(clearTokenCookiesMock).toHaveBeenCalled(); + expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login"); + expect(result.current.accessToken).toBeNull(); + expect(result.current.userRole).toBe("Undefined Role"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts index cba7c1a3dc8..7610c6346be 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts @@ -4,6 +4,7 @@ import { useEffect, useMemo } from "react"; import { useRouter } from "next/navigation"; import { jwtDecode } from "jwt-decode"; import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; +import { getProxyBaseUrl } from "@/components/networking"; function formatUserRole(userRole: string) { if (!userRole) { @@ -42,7 +43,7 @@ const useAuthorized = () => { // Redirect after mount if missing/invalid token useEffect(() => { if (!token) { - router.replace("/sso/key/generate"); + router.replace(`${getProxyBaseUrl()}/ui/login`); } }, [token, router]); @@ -54,7 +55,7 @@ const useAuthorized = () => { } catch { // Bad token in cookie — clear and bounce clearTokenCookies(); - router.replace("/sso/key/generate"); + router.replace(`${getProxyBaseUrl()}/ui/login`); return null; } }, [token, router]); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx index 43e6f986dfb..17a833deacb 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx @@ -1,11 +1,11 @@ import React from "react"; import { Radio } from "antd"; -import type { ExportScope } from "./types"; +import type { ExportScope, EntityType } from "./types"; interface ExportTypeSelectorProps { value: ExportScope; onChange: (value: ExportScope) => void; - entityType: "tag" | "team" | "organization"; + entityType: EntityType; } const ExportTypeSelector: React.FC = ({ value, onChange, entityType }) => { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx index 3547d65379e..e326183d880 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx @@ -7,7 +7,7 @@ import type { EntitySpendData } from "./types"; interface UsageExportHeaderProps { dateValue: DateRangePickerValue; - entityType: "tag" | "team" | "organization"; + entityType: "tag" | "team" | "organization" | "customer"; spendData: EntitySpendData; // Optional filter props showFilters?: boolean; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts index ea11701f7ee..ded2731c945 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts @@ -2,6 +2,7 @@ import type { DateRangePickerValue } from "@tremor/react"; export type ExportFormat = "csv" | "json"; export type ExportScope = "daily" | "daily_with_models"; +export type EntityType = "tag" | "team" | "organization" | "customer"; export interface EntitySpendData { results: any[]; @@ -17,7 +18,7 @@ export interface EntitySpendData { export interface EntityUsageExportModalProps { isOpen: boolean; onClose: () => void; - entityType: "tag" | "team" | "organization"; + entityType: EntityType; spendData: EntitySpendData; dateRange: DateRangePickerValue; selectedFilters: string[]; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index 1327e158a6e..c93155feb23 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -1,6 +1,6 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; import Papa from "papaparse"; -import type { EntitySpendData, EntityBreakdown, ExportMetadata, ExportScope } from "./types"; +import type { EntitySpendData, EntityBreakdown, ExportMetadata, ExportScope, EntityType } from "./types"; import type { DateRangePickerValue } from "@tremor/react"; export const getEntityBreakdown = (spendData: EntitySpendData): EntityBreakdown[] => { @@ -139,7 +139,7 @@ export const generateExportData = ( }; export const generateMetadata = ( - entityType: "tag" | "team" | "organization", + entityType: EntityType, dateRange: DateRangePickerValue, selectedFilters: string[], exportScope: ExportScope, @@ -166,7 +166,7 @@ export const handleExportCSV = ( spendData: EntitySpendData, exportScope: ExportScope, entityLabel: string, - entityType: "tag" | "team" | "organization", + entityType: EntityType, ): void => { const data = generateExportData(spendData, exportScope, entityLabel); const csv = Papa.unparse(data); @@ -186,7 +186,7 @@ export const handleExportJSON = ( spendData: EntitySpendData, exportScope: ExportScope, entityLabel: string, - entityType: "tag" | "team" | "organization", + entityType: EntityType, dateRange: DateRangePickerValue, selectedFilters: string[], ): void => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx index 201bbb3c180..5f165bb7f85 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx @@ -42,7 +42,11 @@ const createQueryClient = () => defaultOptions: { queries: { retry: false, - gcTime: 0, + staleTime: Infinity, + gcTime: Infinity, + refetchOnWindowFocus: false, + refetchOnReconnect: false, + refetchOnMount: false, }, }, }); @@ -140,7 +144,7 @@ describe("Add Model Tab", () => { ); expect(await screen.findByRole("tab", { name: "Add Model" })).toBeInTheDocument(); - }); + }, 10000); // This test is flaky, adding a timeout until we find a better solution it("should display both Add Model and Add Auto Router tabs", async () => { const props = createTestProps(); diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx index 892e9e0197e..24df0ac21ef 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx @@ -1,7 +1,7 @@ import { useProviderFields } from "@/app/(dashboard)/hooks/providers/useProviderFields"; import { UploadOutlined } from "@ant-design/icons"; import { Text, TextInput } from "@tremor/react"; -import { Button as Button2, Col, Form, Row, Select, Typography, Upload, UploadProps } from "antd"; +import { Button as Button2, Col, Form, Input, Row, Select, Typography, Upload, UploadProps } from "antd"; import React from "react"; import { CredentialItem, ProviderCredentialFieldMetadata } from "../networking"; import { provider_map, Providers } from "../provider_info_helpers"; @@ -18,7 +18,7 @@ interface ProviderCredentialField { placeholder?: string; tooltip?: string; required?: boolean; - type?: "text" | "password" | "select" | "upload"; + type?: "text" | "password" | "select" | "upload" | "textarea"; options?: string[]; defaultValue?: string; } @@ -36,7 +36,9 @@ const mapFieldMetadataToUiField = (field: ProviderCredentialFieldMetadata): Prov ? "select" : field.field_type === "upload" ? "upload" - : "text"; + : field.field_type === "textarea" + ? "textarea" + : "text"; return { key: field.key, @@ -247,6 +249,13 @@ const ProviderSpecificFields: React.FC = ({ selecte > }>Click to Upload + ) : field.type === "textarea" ? ( + ) : ( ({ tagDailyActivityCall: vi.fn(), teamDailyActivityCall: vi.fn(), organizationDailyActivityCall: vi.fn(), + customerDailyActivityCall: vi.fn(), })); // Mock the child components to simplify testing @@ -42,6 +43,7 @@ describe("EntityUsage", () => { const mockTagDailyActivityCall = vi.mocked(networking.tagDailyActivityCall); const mockTeamDailyActivityCall = vi.mocked(networking.teamDailyActivityCall); const mockOrganizationDailyActivityCall = vi.mocked(networking.organizationDailyActivityCall); + const mockCustomerDailyActivityCall = vi.mocked(networking.customerDailyActivityCall); const mockSpendData = { results: [ @@ -128,9 +130,11 @@ describe("EntityUsage", () => { mockTagDailyActivityCall.mockClear(); mockTeamDailyActivityCall.mockClear(); mockOrganizationDailyActivityCall.mockClear(); + mockCustomerDailyActivityCall.mockClear(); mockTagDailyActivityCall.mockResolvedValue(mockSpendData); mockTeamDailyActivityCall.mockResolvedValue(mockSpendData); mockOrganizationDailyActivityCall.mockResolvedValue(mockSpendData); + mockCustomerDailyActivityCall.mockResolvedValue(mockSpendData); }); it("should render with tag entity type and display spend metrics", async () => { @@ -182,6 +186,21 @@ describe("EntityUsage", () => { }); }); + it("should render with customer entity type and call customer API", async () => { + render(); + + await waitFor(() => { + expect(mockCustomerDailyActivityCall).toHaveBeenCalled(); + }); + + expect(screen.getByText("Customer Spend Overview")).toBeInTheDocument(); + + await waitFor(() => { + const spendElements = screen.getAllByText("$100.50"); + expect(spendElements.length).toBeGreaterThan(0); + }); + }); + it("should switch between tabs", async () => { render(); diff --git a/ui/litellm-dashboard/src/components/entity_usage.tsx b/ui/litellm-dashboard/src/components/entity_usage.tsx index 501eac7124b..ca30ded9494 100644 --- a/ui/litellm-dashboard/src/components/entity_usage.tsx +++ b/ui/litellm-dashboard/src/components/entity_usage.tsx @@ -23,12 +23,18 @@ import { } from "@tremor/react"; import { ActivityMetrics, processActivityData } from "./activity_metrics"; import { DailyData, BreakdownMetrics, KeyMetricWithMetadata, EntityMetricWithMetadata, TagUsage } from "./usage/types"; -import { organizationDailyActivityCall, tagDailyActivityCall, teamDailyActivityCall } from "./networking"; +import { + organizationDailyActivityCall, + tagDailyActivityCall, + teamDailyActivityCall, + customerDailyActivityCall, +} from "./networking"; import TopKeyView from "./top_key_view"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { valueFormatterSpend } from "./usage/utils/value_formatters"; import { getProviderLogoAndName } from "./provider_info_helpers"; import { UsageExportHeader } from "./EntityUsageExport"; +import type { EntityType } from "./EntityUsageExport/types"; import TopModelView from "./top_model_view"; interface EntityMetrics { @@ -68,7 +74,7 @@ export interface EntityList { interface EntityUsageProps { accessToken: string | null; - entityType: "tag" | "team" | "organization"; + entityType: EntityType; entityId?: string | null; userID: string | null; userRole: string | null; @@ -135,6 +141,15 @@ const EntityUsage: React.FC = ({ selectedTags.length > 0 ? selectedTags : null, ); setSpendData(data); + } else if (entityType === "customer") { + const data = await customerDailyActivityCall( + accessToken, + startTime, + endTime, + 1, + selectedTags.length > 0 ? selectedTags : null, + ); + setSpendData(data); } else { throw new Error("Invalid entity type"); } diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index c5dc3114ff6..5d58add81d7 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -182,7 +182,7 @@ export interface ProviderCredentialFieldMetadata { placeholder?: string | null; tooltip?: string | null; required?: boolean; - field_type?: "text" | "password" | "select" | "upload"; + field_type?: "text" | "password" | "select" | "upload" | "textarea"; options?: string[] | null; default_value?: string | null; } @@ -1736,6 +1736,25 @@ export const organizationDailyActivityCall = async ( }); }; +export const customerDailyActivityCall = async ( + accessToken: string, + startTime: Date, + endTime: Date, + page: number = 1, + customerIds: string[] | null = null, +) => { + return fetchDailyActivity({ + accessToken, + endpoint: "/customer/daily/activity", + startTime, + endTime, + page, + extraQueryParams: { + end_user_ids: customerIds, + }, + }); +}; + export const getTotalSpendCall = async (accessToken: string) => { /** * Get all models on proxy @@ -2511,7 +2530,7 @@ export const allEndUsersCall = async (accessToken: string) => { console.log(data); return data; } catch (error) { - console.error("Failed to create key:", error); + console.error("Failed to fetch end users:", error); throw error; } }; diff --git a/ui/litellm-dashboard/src/components/new_usage.test.tsx b/ui/litellm-dashboard/src/components/new_usage.test.tsx index a06045137d7..aec07765e7e 100644 --- a/ui/litellm-dashboard/src/components/new_usage.test.tsx +++ b/ui/litellm-dashboard/src/components/new_usage.test.tsx @@ -3,6 +3,7 @@ import { describe, it, expect, vi, beforeEach, beforeAll } from "vitest"; import NewUsagePage from "./new_usage"; import type { Organization } from "./networking"; import * as networking from "./networking"; +import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers"; // Polyfill ResizeObserver for test environment beforeAll(() => { @@ -53,9 +54,14 @@ vi.mock("./EntityUsageExport", () => ({ default: () =>
Entity Usage Export Modal
, })); +vi.mock("@/app/(dashboard)/hooks/customers/useCustomers", () => ({ + useCustomers: vi.fn(), +})); + describe("NewUsage", () => { const mockUserDailyActivityAggregatedCall = vi.mocked(networking.userDailyActivityAggregatedCall); const mockTagListCall = vi.mocked(networking.tagListCall); + const mockUseCustomers = vi.mocked(useCustomers); const mockSpendData = { results: [ @@ -174,6 +180,19 @@ describe("NewUsage", () => { }, ]; + const mockCustomers = [ + { + user_id: "customer-123", + alias: "Test Customer", + spend: 0, + blocked: false, + allowed_model_region: null, + default_model: null, + budget_id: null, + litellm_budget_table: null, + }, + ]; + const defaultProps = { accessToken: "test-token", userRole: "Admin", @@ -205,6 +224,11 @@ describe("NewUsage", () => { mockTagListCall.mockClear(); mockUserDailyActivityAggregatedCall.mockResolvedValue(mockSpendData); mockTagListCall.mockResolvedValue({}); + mockUseCustomers.mockReturnValue({ + data: [], + isLoading: false, + error: null, + } as any); }); it("should render and fetch usage data on mount", async () => { @@ -289,4 +313,26 @@ describe("NewUsage", () => { expect(entityUsageElements.length).toBeGreaterThan(0); }); }); + + it("should show customer usage tab for admins", async () => { + mockUseCustomers.mockReturnValue({ + data: mockCustomers, + isLoading: false, + error: null, + } as any); + + const { getByText, getAllByText } = render(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + const customerTab = getByText("Customer Usage"); + fireEvent.click(customerTab); + + await waitFor(() => { + const entityUsageElements = getAllByText("Entity Usage"); + expect(entityUsageElements.length).toBeGreaterThan(0); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/new_usage.tsx b/ui/litellm-dashboard/src/components/new_usage.tsx index a8d30885493..f9280ae2bf2 100644 --- a/ui/litellm-dashboard/src/components/new_usage.tsx +++ b/ui/litellm-dashboard/src/components/new_usage.tsx @@ -27,9 +27,10 @@ import { Text, Title, } from "@tremor/react"; -import React, { useCallback, useEffect, useMemo, useState } from "react"; import { Alert } from "antd"; +import React, { useCallback, useEffect, useMemo, useState } from "react"; +import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { Button } from "@tremor/react"; import { all_admin_roles } from "../utils/roles"; @@ -86,6 +87,7 @@ const NewUsagePage: React.FC = ({ }); const [allTags, setAllTags] = useState([]); + const { data: customers = [] } = useCustomers(accessToken, userRole); const [modelViewType, setModelViewType] = useState<"groups" | "individual">("groups"); const [isCloudZeroModalOpen, setIsCloudZeroModalOpen] = useState(false); const [isGlobalExportModalOpen, setIsGlobalExportModalOpen] = useState(false); @@ -430,6 +432,7 @@ const NewUsagePage: React.FC = ({ Your Organization Usage )} Team Usage + {all_admin_roles.includes(userRole || "") ? Customer Usage : <>} {all_admin_roles.includes(userRole || "") ? Tag Usage : <>} {all_admin_roles.includes(userRole || "") ? User Agent Activity : <>} @@ -798,6 +801,23 @@ const NewUsagePage: React.FC = ({ /> + {/* Customer Usage Panel */} + + ({ + label: customer.alias || customer.user_id, + value: customer.user_id, + })) || null + } + premiumUser={premiumUser} + dateValue={dateValue} + /> + {/* Tag Usage Panel */} ({ fetchAvailableModels: vi.fn().mockResolvedValue([{ model_group: "gpt-4" }, { model_group: "gpt-3.5-turbo" }]), @@ -11,6 +12,34 @@ vi.mock("../llm_calls/chat_completion", () => ({ makeOpenAIChatCompletionRequest: vi.fn().mockResolvedValue(undefined), })); +let capturedOnImageUpload: ((file: File) => false) | null = null; + +vi.mock("../chat_ui/ChatImageUpload", () => ({ + default: ({ onImageUpload }: { onImageUpload: (file: File) => false }) => { + capturedOnImageUpload = onImageUpload; + return ( +
+ +
+ ); + }, +})); + +vi.mock("../chat_ui/ChatImageUtils", () => ({ + createChatMultimodalMessage: vi.fn().mockResolvedValue({ + role: "user", + content: [ + { type: "text", text: "test message" }, + { type: "image_url", image_url: { url: "data:image/png;base64,test" } }, + ], + }), + createChatDisplayMessage: vi.fn().mockReturnValue({ + role: "user", + content: "test message [Image attached]", + imagePreviewUrl: "blob:test-url", + }), +})); + vi.mock("./components/ComparisonPanel", () => ({ ComparisonPanel: ({ comparison, onRemove }: { comparison: any; onRemove: () => void }) => (
@@ -22,8 +51,9 @@ vi.mock("./components/ComparisonPanel", () => ({ })); vi.mock("./components/MessageInput", () => ({ - MessageInput: ({ value, onChange, onSend, disabled }: any) => ( + MessageInput: ({ value, onChange, onSend, disabled, hasAttachment, uploadComponent }: any) => (
+ {uploadComponent &&
{uploadComponent}
}