diff --git a/.circleci/config.yml b/.circleci/config.yml index 133a7184f9b..e03d1086282 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1153,7 +1153,7 @@ jobs: pip install "pytest-asyncio==0.21.1" pip install "respx==0.22.0" pip install "pydantic==2.10.2" - pip install "mcp==1.10.1" + pip install "mcp==1.21.2" # Run pytest and generate JUnit XML report - run: name: Run tests diff --git a/README.md b/README.md index 75a23faa5c1..58ffa12c5e1 100644 --- a/README.md +++ b/README.md @@ -374,7 +374,9 @@ Support for more providers. Missing a provider or LLM Platform, raise a [feature 1. (In root) create virtual environment `python -m venv .venv` 2. Activate virtual environment `source .venv/bin/activate` 3. Install dependencies `pip install -e ".[all]"` -4. Start proxy backend `python litellm/proxy_cli.py` +4. `pip install prisma` +5. `prisma generate` +6. Start proxy backend `python litellm/proxy/proxy_cli.py` ### Frontend 1. Navigate to `ui/litellm-dashboard` diff --git a/cookbook/ai_coding_tool_guides/index.json b/cookbook/ai_coding_tool_guides/index.json index 7d022d6de3b..f879292aeff 100644 --- a/cookbook/ai_coding_tool_guides/index.json +++ b/cookbook/ai_coding_tool_guides/index.json @@ -95,4 +95,40 @@ "LiteLLM", "Quickstart" ] +}, +{ + "title": "AI Coding Tool Usage Tracking", + "description": "This is a guide to tracking usage for AI coding tools monitor the use of Claude Code , Google Antigravity, OpenAI Codex, Roo Code etc. through LiteLLM.", + "url": "https://docs.litellm.ai/docs/tutorials/cost_tracking_coding", + "date": "2026-01-17", + "version": "1.0.0", + "tags": [ + "Claude Code", + "Gemini CLI", + "OpenAI Codex", + "LiteLLM" + ] +}, +{ + "title": "Use Web Search with Claude Code (across OpenAI/Anthropic/Gemini/etc.)", + "description": "This is a guide for using Web Search with Claude Code via LiteLLM.", + "url": "https://docs.litellm.ai/docs/tutorials/claude_code_websearch", + "date": "2026-01-17", + "version": "1.0.0", + "tags": [ + "Claude Code", + "LiteLLM", + "Web Search" + ] +}, +{ + "title": "Track Claude Code Usage per user via Custom Headers", + "description": "This is a guide for tracking claude code user usage by passing a customer ID header.", + "url": "https://docs.litellm.ai/docs/tutorials/claude_code_customer_tracking", + "date": "2026-01-17", + "version": "1.0.0", + "tags": [ + "Claude Code", + "LiteLLM" + ] }] \ No newline at end of file diff --git a/docker/Dockerfile.health_check b/docker/Dockerfile.health_check new file mode 100644 index 00000000000..de62e4bd729 --- /dev/null +++ b/docker/Dockerfile.health_check @@ -0,0 +1,16 @@ +FROM python:3.11-slim + +WORKDIR /app + +# Copy health check script and requirements +COPY scripts/health_check/health_check_client.py /app/health_check_client.py +COPY scripts/health_check/health_check_requirements.txt /app/requirements.txt + +# Install dependencies +RUN pip install --no-cache-dir -r requirements.txt + +# Make script executable +RUN chmod +x /app/health_check_client.py + +# Set entrypoint +ENTRYPOINT ["python", "/app/health_check_client.py"] diff --git a/docker/supervisord.conf b/docker/supervisord.conf index c6855fe652b..877335804fe 100644 --- a/docker/supervisord.conf +++ b/docker/supervisord.conf @@ -1,6 +1,8 @@ [supervisord] nodaemon=true loglevel=info +logfile=/tmp/supervisord.log +pidfile=/tmp/supervisord.pid [group:litellm] programs=main,health diff --git a/docs/my-website/docs/contributing.md b/docs/my-website/docs/contributing.md index a88013ff1b3..be7222f6cb8 100644 --- a/docs/my-website/docs/contributing.md +++ b/docs/my-website/docs/contributing.md @@ -1,45 +1,100 @@ # Contributing - UI -Here's how to run the LiteLLM UI locally for making changes: +Thanks for contributing to the LiteLLM UI! This guide will help you set up your local development environment. + + +## 1. Clone the repo -## 1. Clone the repo ```bash git clone https://github.com/BerriAI/litellm.git +cd litellm ``` -## 2. Start the UI + Proxy +## 2. Start the Proxy -**2.1 Start the proxy on port 4000** +Create a config file (e.g., `config.yaml`): -Tell the proxy where the UI is located -```bash -DATABASE_URL = "postgresql://:@:/" -LITELLM_MASTER_KEY = "sk-1234" -STORE_MODEL_IN_DB = "True" +```yaml +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + +general_settings: + master_key: sk-1234 + database_url: postgresql://:@:/ + store_model_in_db: true ``` +Start the proxy on port 4000: + ```bash -cd litellm/litellm/proxy -python3 proxy_cli.py --config /path/to/config.yaml --port 4000 +poetry run litellm --config config.yaml --port 4000 ``` -**2.2 Start the UI** +The UI comes pre-built in the repo. Access it at `http://localhost:4000/ui` -Set the mode as development (this will assume the proxy is running on localhost:4000) -```bash -npm install # install dependencies -``` +## 3. UI Development + +There are two options for UI development: + +### Option A: Development Mode (Hot Reload) + +This runs the UI on port 3000 with hot reload. The proxy runs on port 4000. ```bash -cd litellm/ui/litellm-dashboard - +cd ui/litellm-dashboard +npm install npm run dev - -# starts on http://0.0.0.0:3000 ``` -## 3. Go to local UI +**Login flow:** +1. Go to `http://localhost:3000` +2. You'll be redirected to `http://localhost:4000/ui` for login +3. After logging in, manually navigate back to `http://localhost:3000/` +4. You're now authenticated and can develop with hot reload + +:::note +If you experience redirect loops or authentication issues, clear your browser cookies for localhost or use Build Mode instead. +::: + +### Option B: Build Mode + +This builds the UI and copies it to the proxy. Changes require rebuilding. + +1. Make your code changes in `ui/litellm-dashboard/src/` + +2. Build the UI +```bash +cd ui/litellm-dashboard +npm install +npm run build +``` + +After building, copy the output to the proxy: ```bash -http://0.0.0.0:3000 -``` \ No newline at end of file +cp -r out/* ../../litellm/proxy/_experimental/out/ +``` + +Then restart the proxy and access the UI at `http://localhost:4000/ui` + +## 4. Submitting a PR + +1. Create a new branch for your changes: +```bash +git checkout -b feat/your-feature-name +``` + +2. Stage and commit your changes: +```bash +git add . +git commit -m "feat: description of your changes" +``` + +3. Push to your fork: +```bash +git push origin feat/your-feature-name +``` + +4. Create a Pull Request on GitHub following the [PR template](https://github.com/BerriAI/litellm/blob/main/.github/pull_request_template.md) diff --git a/docs/my-website/docs/providers/stability.md b/docs/my-website/docs/providers/stability.md index 62a8ab43cd8..c4bc5376d1f 100644 --- a/docs/my-website/docs/providers/stability.md +++ b/docs/my-website/docs/providers/stability.md @@ -173,6 +173,14 @@ Stability AI returns images in base64 format. The response is OpenAI-compatible: Stability AI supports various image editing operations including inpainting, upscaling, outpainting, background removal, and more. +:::info Optional Parameters +**Important:** Different Stability models have different parameter requirements: +- Some models don't require a `prompt` (e.g., upscaling, background removal) +- The `style-transfer` model uses `init_image` and `style_image` instead of `image` +- The `outpaint` model requires numeric parameters (`left`, `right`, `up`, `down`) +LiteLLM automatically handles these differences for you. +::: + ### Usage - LiteLLM Python SDK #### Inpainting (Edit with Mask) @@ -217,11 +225,11 @@ response = image_edit( creativity=0.3, # 0-0.35, higher = more creative ) -# Fast upscaling - quick upscaling +# Fast upscaling - quick upscaling (no prompt needed) response = image_edit( model="stability/stable-fast-upscale-v1:0", image=open("low_res_image.png", "rb"), - prompt="Quickly upscale this image", + # No prompt required for fast upscale ) print(response) ``` @@ -259,7 +267,7 @@ os.environ['STABILITY_API_KEY'] = "your-api-key" response = image_edit( model="stability/stable-image-remove-background-v1:0", image=open("portrait.png", "rb"), - prompt="Remove the background", + # No prompt required for fast upscale ) print(response) ``` @@ -329,10 +337,29 @@ response = image_edit( model="stability/stable-image-erase-object-v1:0", image=open("scene.png", "rb"), mask=open("object_mask.png", "rb"), # Mask the object to erase - prompt="Remove the object", + # No prompt needed ) print(response) ``` +#### Style Transfer + +```python showLineNumbers +from litellm import image_edit +import os + +os.environ['STABILITY_API_KEY'] = "your-api-key" + +# Transfer style from one image to another +# Note: Uses init_image (via image param) and style_image +response = image_edit( + model="stability/stable-style-transfer-v1:0", + image=open("content_image.png", "rb"), # Maps to init_image + style_image=open("style_reference.png", "rb"), # Style to apply + fidelity=0.5, # 0-1, balance between content and style + # No prompt needed +) + +print(response) ### Supported Image Edit Models @@ -419,6 +446,23 @@ response = image_edit( ) print(response) ``` +# Fast upscale without prompt +response = image_edit( + model="bedrock/stability.stable-fast-upscale-v1:0", + image=open("low_res_image.png", "rb"), +) + +# Outpaint with numeric parameters +response = image_edit( + model="bedrock/stability.stable-outpaint-v1:0", + image=open("original_image.png", "rb"), + left=100, # Automatically converted to int + right=100, + up=50, + down=50, +) + +print(response) ### Supported Bedrock Stability Models diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 33ebf535d29..be2bf86ab10 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -1390,6 +1390,77 @@ model_list: +### **Workload Identity Federation** + +LiteLLM supports [Google Cloud Workload Identity Federation (WIF)](https://cloud.google.com/iam/docs/workload-identity-federation), which allows you to grant on-premises or multi-cloud workloads access to Google Cloud resources without using a service account key. This is the recommended approach for workloads running in other cloud environments (AWS, Azure, etc.) or on-premises. + +To use Workload Identity Federation, pass the path to your WIF credentials configuration file via `vertex_credentials`: + + + + +```python +from litellm import completion + +response = completion( + model="vertex_ai/gemini-1.5-pro", + messages=[{"role": "user", "content": "Hello!"}], + vertex_credentials="/path/to/wif-credentials.json", # 👈 WIF credentials file + vertex_project="your-gcp-project-id", + vertex_location="us-central1" +) +``` + + + + +```yaml +model_list: + - model_name: gemini-model + litellm_params: + model: vertex_ai/gemini-1.5-pro + vertex_project: your-gcp-project-id + vertex_location: us-central1 + vertex_credentials: /path/to/wif-credentials.json # 👈 WIF credentials file +``` + +Alternatively, you can create credentials in **LLM Credentials** in the LiteLLM UI and use those to authenticate your models: + +```yaml +model_list: + - model_name: gemini-model + litellm_params: + model: vertex_ai/gemini-1.5-pro + vertex_project: your-gcp-project-id + vertex_location: us-central1 + litellm_credential_name: my-vertex-wif-credential # 👈 Reference credential stored in UI +``` + + + + +**WIF Credentials File Format** + +Your WIF credentials JSON file typically looks like this (for AWS federation): + +```json +{ + "type": "external_account", + "audience": "//iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/POOL_ID/providers/PROVIDER_ID", + "subject_token_type": "urn:ietf:params:aws:token-type:aws4_request", + "service_account_impersonation_url": "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/SERVICE_ACCOUNT_EMAIL:generateAccessToken", + "token_url": "https://sts.googleapis.com/v1/token", + "credential_source": { + "environment_id": "aws1", + "region_url": "http://169.254.169.254/latest/meta-data/placement/availability-zone", + "url": "http://169.254.169.254/latest/meta-data/iam/security-credentials", + "regional_cred_verification_url": "https://sts.{region}.amazonaws.com?Action=GetCallerIdentity&Version=2011-06-15" + } +} +``` + +For more details on setting up Workload Identity Federation, see [Google Cloud WIF documentation](https://cloud.google.com/iam/docs/workload-identity-federation). + ### **Environment Variables** You can set: diff --git a/docs/my-website/docs/proxy/deploy.md b/docs/my-website/docs/proxy/deploy.md index 5686e9fd835..7393e73ba87 100644 --- a/docs/my-website/docs/proxy/deploy.md +++ b/docs/my-website/docs/proxy/deploy.md @@ -4,6 +4,10 @@ import Image from '@theme/IdealImage'; # Docker, Helm, Terraform +:::info No Limits on LiteLLM OSS +There are **no limits** on the number of users, keys, or teams you can create on LiteLLM OSS. +::: + You can find the Dockerfile to build litellm proxy [here](https://github.com/BerriAI/litellm/blob/main/Dockerfile) > Note: Production requires at least 4 CPU cores and 8 GB RAM. diff --git a/docs/my-website/docs/tutorials/claude_code_websearch.md b/docs/my-website/docs/tutorials/claude_code_websearch.md index cc2f79666da..478fc960348 100644 --- a/docs/my-website/docs/tutorials/claude_code_websearch.md +++ b/docs/my-website/docs/tutorials/claude_code_websearch.md @@ -1,12 +1,16 @@ +import Image from '@theme/IdealImage'; + # Claude Code - WebSearch Across All Providers Enable Claude Code's web search tool to work with any provider (Bedrock, Azure, Vertex, etc.). LiteLLM automatically intercepts web search requests and executes them server-side. + + ## Proxy Configuration Add WebSearch interception to your `litellm_config.yaml`: -```yaml +```yaml showLineNumbers title="litellm_config.yaml" model_list: - model_name: bedrock-sonnet litellm_params: @@ -37,7 +41,7 @@ search_tools: Create `config.yaml`: -```yaml +```yaml showLineNumbers title="config.yaml" model_list: - model_name: bedrock-sonnet litellm_params: @@ -58,14 +62,14 @@ search_tools: ### 2. Start Proxy -```bash +```bash showLineNumbers title="Start LiteLLM Proxy" export PERPLEXITY_API_KEY=your-key litellm --config config.yaml ``` ### 3. Use with Claude Code -```bash +```bash showLineNumbers title="Configure Claude Code" export ANTHROPIC_BASE_URL=http://localhost:4000 export ANTHROPIC_API_KEY=sk-1234 claude @@ -116,12 +120,19 @@ sequenceDiagram Configure which search provider to use. LiteLLM supports multiple search providers: -| Provider | Configuration | -|----------|---------------| -| **Perplexity** | `search_provider: perplexity` | -| **Tavily** | `search_provider: tavily` | +| Provider | `search_provider` Value | Environment Variable | +|----------|------------------------|----------------------| +| **Perplexity AI** | `perplexity` | `PERPLEXITYAI_API_KEY` | +| **Tavily** | `tavily` | `TAVILY_API_KEY` | +| **Exa AI** | `exa_ai` | `EXA_API_KEY` | +| **Parallel AI** | `parallel_ai` | `PARALLEL_AI_API_KEY` | +| **Google PSE** | `google_pse` | `GOOGLE_PSE_API_KEY`, `GOOGLE_PSE_ENGINE_ID` | +| **DataForSEO** | `dataforseo` | `DATAFORSEO_LOGIN`, `DATAFORSEO_PASSWORD` | +| **Firecrawl** | `firecrawl` | `FIRECRAWL_API_KEY` | +| **SearXNG** | `searxng` | `SEARXNG_API_BASE` (required) | +| **Linkup** | `linkup` | `LINKUP_API_KEY` | -See [all supported search providers](../search/index.md) for the complete list. +See [all supported search providers](../search/index.md) for detailed setup instructions and provider-specific parameters. ## Configuration Options @@ -145,7 +156,7 @@ Use these values in `enabled_providers`: ### Complete Configuration Example -```yaml +```yaml showLineNumbers title="Complete config.yaml" model_list: - model_name: bedrock-sonnet litellm_params: diff --git a/docs/my-website/docs/tutorials/cursor_integration.md b/docs/my-website/docs/tutorials/cursor_integration.md index 3f462e1ee5d..49f88bd0487 100644 --- a/docs/my-website/docs/tutorials/cursor_integration.md +++ b/docs/my-website/docs/tutorials/cursor_integration.md @@ -1,3 +1,5 @@ +import Image from '@theme/IdealImage'; + # Cursor Integration Route Cursor IDE requests through LiteLLM for unified logging, budget controls, and access to any model. @@ -76,6 +78,34 @@ Send a message. All requests now route through LiteLLM. --- +## Connecting MCP Servers + +You can also connect MCP servers to Cursor via LiteLLM Proxy. + +For official instructions on configuring MCP integration with Cursor, please refer to the Cursor documentation here: [https://cursor.com/en-US/docs/context/mcp](https://cursor.com/en-US/docs/context/mcp). + +1. In Cursor Settings, go to the "Tools & MCP" tab and click "New MCP Server". + +2. In your `mcp.json`, add the following configuration: + +``` +{ + "mcpServers": { + "litellm": { + "url": "http://localhost:4000/everything/mcp", + "type": "http", + "headers": { + "Authorization": "Bearer sk-LITELLM_VIRTUAL_KEY" + } + } + } +} +``` + +3. LiteLLM's MCP will now appear under "Installed MCP Servers" in Cursor. + + + ## Troubleshooting | Issue | Solution | diff --git a/docs/my-website/img/claude_code_websearch.png b/docs/my-website/img/claude_code_websearch.png new file mode 100644 index 00000000000..a0d8a3ba85a Binary files /dev/null and b/docs/my-website/img/claude_code_websearch.png differ diff --git a/docs/my-website/img/cursor_mcp_installed.png b/docs/my-website/img/cursor_mcp_installed.png new file mode 100644 index 00000000000..f2339bcec3d Binary files /dev/null and b/docs/my-website/img/cursor_mcp_installed.png differ diff --git a/docs/my-website/img/release_notes/claude_code_websearch.png b/docs/my-website/img/release_notes/claude_code_websearch.png new file mode 100644 index 00000000000..eec4b6d70e8 Binary files /dev/null and b/docs/my-website/img/release_notes/claude_code_websearch.png differ diff --git a/docs/my-website/release_notes/v1.81.0/index.md b/docs/my-website/release_notes/v1.81.0/index.md index 071422f96be..edd720f6bfd 100644 --- a/docs/my-website/release_notes/v1.81.0/index.md +++ b/docs/my-website/release_notes/v1.81.0/index.md @@ -1,5 +1,5 @@ --- -title: "v1.81.0 - Claude Code - Web Search with all LiteLLM Providers" +title: "v1.81.0 - Claude Code - Web Search Across All Providers" slug: "v1-81-0" date: 2026-01-18T10:00:00 authors: @@ -47,6 +47,22 @@ pip install litellm==1.81.0 - **Claude Code** - Support for using web search across Bedrock, Vertex AI, and all LiteLLM providers - **Major Change** - [50MB limit on image URL downloads](#major-change---chatcompletions-image-url-download-size-limit) to improve reliability +- **Performance** - [25% CPU Usage Reduction](#performance---25-cpu-usage-reduction) by removing premature model.dump() calls from the hot path +- **Deleted Keys Audit Table on UI** - [View deleted keys and teams for audit purposes](../../docs/proxy/deleted_keys_teams.md) with spend and budget information at the time of deletion + +--- + +## Claude Code - Web Search Across All Providers + + + +This release brings web search support to Claude Code across all LiteLLM providers (Bedrock, Azure, Vertex AI, and more), enabling AI coding assistants to search the web for real-time information. + +This means you can now use Claude Code's web search tool with any provider, not just Anthropic's native API. LiteLLM automatically intercepts web search requests and executes them server-side using your configured search provider (Perplexity, Tavily, Exa AI, and more). + +Proxy Admins can configure web search interception in their LiteLLM proxy config to enable this capability for their teams using Claude Code with Bedrock, Azure, or any other supported provider. + +[**Learn more →**](../../docs/tutorials/claude_code_websearch.md) --- @@ -140,6 +156,20 @@ This feature improves reliability by: --- +## Performance - 25% CPU Usage Reduction + +LiteLLM now reduces CPU usage by removing premature `model.dump()` calls from the hot path in request processing. Previously, Pydantic model serialization was performed earlier and more frequently than necessary, causing unnecessary CPU overhead on every request. By deferring serialization until it is actually needed, LiteLLM reduces CPU usage and improves request throughput under high load. + +--- + +## Deleted Keys Audit Table on UI + + + +LiteLLM now provides a comprehensive audit table for deleted API keys and teams directly in the UI. This feature allows you to easily track the spend of deleted keys, view their associated team information, and maintain accurate financial records for auditing and compliance purposes. The table displays key details including key aliases, team associations, and spend information captured at the time of deletion. For more information on how to use this feature, see the [Deleted Keys & Teams documentation](../../docs/proxy/deleted_keys_teams.md). + +--- + ## New Models / Updated Models #### New Model Support diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251115120021_baseline_diff/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251115120021_baseline_diff/migration.sql deleted file mode 100644 index 2f725d83806..00000000000 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251115120021_baseline_diff/migration.sql +++ /dev/null @@ -1,2 +0,0 @@ --- This is an empty migration. - diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251115120539_baseline_diff/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251115120539_baseline_diff/migration.sql deleted file mode 100644 index 2f725d83806..00000000000 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251115120539_baseline_diff/migration.sql +++ /dev/null @@ -1,2 +0,0 @@ --- This is an empty migration. - diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 167aad7959a..2d36dbeacda 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -113,7 +113,9 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str: litellm_logging_obj.model = model litellm_logging_obj.custom_llm_provider = custom_llm_provider litellm_logging_obj.model_call_details["model"] = model - litellm_logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider + litellm_logging_obj.model_call_details[ + "custom_llm_provider" + ] = custom_llm_provider return agent_name @@ -197,7 +199,11 @@ async def asend_message( ) # Extract params from request - params = request.params.model_dump(mode="json") if hasattr(request.params, "model_dump") else dict(request.params) + params = ( + request.params.model_dump(mode="json") + if hasattr(request.params, "model_dump") + else dict(request.params) + ) response_dict = await A2ACompletionBridgeHandler.handle_non_streaming( request_id=str(request.id), @@ -216,7 +222,9 @@ async def asend_message( # Create A2A client if not provided but api_base is available if a2a_client is None: if api_base is None: - raise ValueError("Either a2a_client or api_base is required for standard A2A flow") + raise ValueError( + "Either a2a_client or api_base is required for standard A2A flow" + ) a2a_client = await create_a2a_client(base_url=api_base) # Type assertion: a2a_client is guaranteed to be non-None here @@ -235,7 +243,11 @@ async def asend_message( # Calculate token usage from request and response response_dict = a2a_response.model_dump(mode="json", exclude_none=True) - prompt_tokens, completion_tokens, _ = A2ARequestUtils.calculate_usage_from_request_response( + ( + prompt_tokens, + completion_tokens, + _, + ) = A2ARequestUtils.calculate_usage_from_request_response( request=request, response_dict=response_dict, ) @@ -280,7 +292,9 @@ def send_message( if loop is not None: return asend_message(a2a_client=a2a_client, request=request, **kwargs) else: - return asyncio.run(asend_message(a2a_client=a2a_client, request=request, **kwargs)) + return asyncio.run( + asend_message(a2a_client=a2a_client, request=request, **kwargs) + ) async def asend_message_streaming( @@ -347,7 +361,11 @@ async def asend_message_streaming( ) # Extract params from request - params = request.params.model_dump(mode="json") if hasattr(request.params, "model_dump") else dict(request.params) + params = ( + request.params.model_dump(mode="json") + if hasattr(request.params, "model_dump") + else dict(request.params) + ) async for chunk in A2ACompletionBridgeHandler.handle_streaming( request_id=str(request.id), @@ -365,7 +383,9 @@ async def asend_message_streaming( # Create A2A client if not provided but api_base is available if a2a_client is None: if api_base is None: - raise ValueError("Either a2a_client or api_base is required for standard A2A flow") + raise ValueError( + "Either a2a_client or api_base is required for standard A2A flow" + ) a2a_client = await create_a2a_client(base_url=api_base) # Type assertion: a2a_client is guaranteed to be non-None here @@ -378,7 +398,9 @@ async def asend_message_streaming( stream = a2a_client.send_message_streaming(request) # Build logging object for streaming completion callbacks - agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr(a2a_client, "agent_card", None) + agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr( + a2a_client, "agent_card", None + ) agent_name = getattr(agent_card, "name", "unknown") if agent_card else "unknown" model = f"a2a_agent/{agent_name}" @@ -456,7 +478,7 @@ async def create_a2a_client( if not A2A_SDK_AVAILABLE: raise ImportError( "The 'a2a' package is required for A2A agent invocation. " - "Install it with: pip install a2a" + "Install it with: pip install a2a-sdk" ) verbose_logger.info(f"Creating A2A client for {base_url}") @@ -512,7 +534,7 @@ async def aget_agent_card( if not A2A_SDK_AVAILABLE: raise ImportError( "The 'a2a' package is required for A2A agent invocation. " - "Install it with: pip install a2a" + "Install it with: pip install a2a-sdk" ) verbose_logger.info(f"Fetching agent card from {base_url}") @@ -534,5 +556,3 @@ async def aget_agent_card( f"Fetched agent card: {agent_card.name if hasattr(agent_card, 'name') else 'unknown'}" ) return agent_card - - diff --git a/litellm/images/main.py b/litellm/images/main.py index 1b09c20d350..6c4c502a7b0 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -714,8 +714,8 @@ def image_variation( @client def image_edit( # noqa: PLR0915 - image: Union[FileTypes, List[FileTypes]], - prompt: str, + image: Optional[Union[FileTypes, List[FileTypes]]] = None, + prompt: Optional[str]= None, model: Optional[str] = None, mask: Optional[str] = None, n: Optional[int] = None, @@ -766,7 +766,7 @@ def image_edit( # noqa: PLR0915 _is_async = kwargs.pop("async_call", False) is True # add images / or return a single image - images = image if isinstance(image, list) else [image] + images = image if isinstance(image, list) else ([image] if image is not None else []) headers_from_kwargs = kwargs.get("headers") merged_extra_headers: Dict[str, Any] = {} diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index a223925d59a..9fbccac68dd 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -987,7 +987,10 @@ class OpenTelemetry(CustomLogger): # TODO: Refactor to use the proper OTEL Logs API instead of directly creating SDK LogRecords from opentelemetry._logs import SeverityNumber, get_logger, get_logger_provider - from opentelemetry.sdk._logs import LogRecord as SdkLogRecord + try: + from opentelemetry.sdk._logs import LogRecord as SdkLogRecord # OTEL < 1.39.0 + except ImportError: + from opentelemetry.sdk._logs._internal import LogRecord as SdkLogRecord # OTEL >= 1.39.0 otel_logger = get_logger(LITELLM_LOGGER_NAME) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 43ed23587d8..21a79af2bd4 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4410,9 +4410,10 @@ def _bedrock_tools_pt(tools: List) -> List[BedrockToolBlock]: defs = parameters.pop("$defs", {}) defs_copy = copy.deepcopy(defs) - # flatten the defs - for _, value in defs_copy.items(): - unpack_defs(value, defs_copy) + # Expand $ref references in parameters using the definitions + # Note: We don't pre-flatten defs as that causes exponential memory growth + # with circular references (see issue #19098). unpack_defs handles nested + # refs recursively and correctly detects/skips circular references. unpack_defs(parameters, defs_copy) tool_input_schema = BedrockToolInputSchemaBlock( json=BedrockToolJsonSchemaBlock( diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 5b1b663e855..86378b97d2e 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -934,8 +934,15 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) return tools - def _ensure_context_management_beta_header(self, headers: dict) -> None: - beta_value = ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value + def _ensure_beta_header(self, headers: dict, beta_value: str) -> None: + """ + Ensure a beta header value is present in the anthropic-beta header. + Merges with existing values instead of overriding them. + + Args: + headers: Dictionary of headers to update + beta_value: The beta header value to add + """ existing_beta = headers.get("anthropic-beta") if existing_beta is None: headers["anthropic-beta"] = beta_value @@ -944,6 +951,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if beta_value not in existing_values: headers["anthropic-beta"] = f"{existing_beta}, {beta_value}" + def _ensure_context_management_beta_header(self, headers: dict) -> None: + beta_value = ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value + self._ensure_beta_header(headers, beta_value) + def update_headers_with_optional_anthropic_beta( self, headers: dict, optional_params: dict ) -> dict: @@ -960,20 +971,20 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if tool.get("type", None) and tool.get("type").startswith( ANTHROPIC_HOSTED_TOOLS.WEB_FETCH.value ): - headers["anthropic-beta"] = ( - ANTHROPIC_BETA_HEADER_VALUES.WEB_FETCH_2025_09_10.value + self._ensure_beta_header( + headers, ANTHROPIC_BETA_HEADER_VALUES.WEB_FETCH_2025_09_10.value ) elif tool.get("type", None) and tool.get("type").startswith( ANTHROPIC_HOSTED_TOOLS.MEMORY.value ): - headers["anthropic-beta"] = ( - ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value + self._ensure_beta_header( + headers, ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value ) if optional_params.get("context_management") is not None: self._ensure_context_management_beta_header(headers) if optional_params.get("output_format") is not None: - headers["anthropic-beta"] = ( - ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value + self._ensure_beta_header( + headers, ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value ) return headers diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index 87bae59ba0f..77d46ff9179 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -88,7 +88,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): self, model: str, prompt: Optional[str], - image: FileTypes, + image: Optional[FileTypes], image_edit_optional_request_params: Dict, litellm_params: GenericLiteLLMParams, headers: dict, @@ -102,6 +102,9 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): if prompt is None: raise ValueError("FLUX 2 image edit requires a prompt.") + if image is None: + raise ValueError("FLUX 2 image edit requires an image.") + image_b64 = self._convert_image_to_base64(image) # Build request body with required params diff --git a/litellm/llms/base_llm/image_edit/transformation.py b/litellm/llms/base_llm/image_edit/transformation.py index cc723480371..b088cdf37f6 100644 --- a/litellm/llms/base_llm/image_edit/transformation.py +++ b/litellm/llms/base_llm/image_edit/transformation.py @@ -93,7 +93,7 @@ class BaseImageEditConfig(ABC): self, model: str, prompt: Optional[str], - image: FileTypes, + image: Optional[FileTypes], image_edit_optional_request_params: Dict, litellm_params: GenericLiteLLMParams, headers: dict, diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 0f1dcff6294..ef441fa5039 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -62,7 +62,7 @@ class BedrockImageEdit(BaseAWSLLM): self, model: str, image: list, - prompt: str, + prompt: Optional[str], model_response: ImageResponse, optional_params: dict, logging_obj: LitellmLogging, @@ -127,7 +127,7 @@ class BedrockImageEdit(BaseAWSLLM): timeout: Optional[Union[float, httpx.Timeout]], model: str, logging_obj: LitellmLogging, - prompt: str, + prompt: Optional[str], model_response: ImageResponse, client: Optional[AsyncHTTPHandler] = None, ) -> ImageResponse: @@ -163,7 +163,7 @@ class BedrockImageEdit(BaseAWSLLM): self, model: str, image: list, - prompt: str, + prompt: Optional[str], optional_params: dict, api_base: Optional[str], extra_headers: Optional[dict], @@ -176,7 +176,7 @@ class BedrockImageEdit(BaseAWSLLM): Args: model (str): The model to use for the image edit image (list): The images to edit - prompt (str): The prompt for the edit + prompt (Optional[str]): The prompt for the edit optional_params (dict): The optional parameters for the image edit api_base (Optional[str]): The base URL for the Bedrock API extra_headers (Optional[dict]): The extra headers to include in the request @@ -248,7 +248,7 @@ class BedrockImageEdit(BaseAWSLLM): self, model: str, image: list, - prompt: str, + prompt: Optional[str], optional_params: dict, ) -> dict: """ @@ -276,7 +276,7 @@ class BedrockImageEdit(BaseAWSLLM): model_response: ImageResponse, model: str, logging_obj: LitellmLogging, - prompt: str, + prompt: Optional[str], response: httpx.Response, data: dict, ) -> ImageResponse: diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index e8b77812988..fc14b571a8c 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -150,11 +150,11 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): return mapped_params - def transform_image_edit_request( + def transform_image_edit_request( #noqa: PLR0915 self, model: str, prompt: Optional[str], - image: FileTypes, + image: Optional[FileTypes], image_edit_optional_request_params: Dict, litellm_params: GenericLiteLLMParams, headers: dict, @@ -164,32 +164,38 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): Returns the request body dict that will be JSON-encoded by the handler. """ - if prompt is None: - raise ValueError("Bedrock Stability image edit requires a prompt.") - # Build Bedrock Stability request data: Dict[str, Any] = { - "prompt": prompt, "output_format": "png", # Default to PNG } - # Convert image to base64 - image_b64: str - if hasattr(image, 'read') and callable(getattr(image, 'read', None)): - # File-like object (e.g., BufferedReader from open()) - image_bytes = image.read() # type: ignore - image_b64 = base64.b64encode(image_bytes).decode('utf-8') # type: ignore - elif isinstance(image, bytes): - # Raw bytes - image_b64 = base64.b64encode(image).decode('utf-8') - elif isinstance(image, str): - # Already a base64 string - image_b64 = image - else: - # Try to handle as bytes - image_b64 = base64.b64encode(bytes(image)).decode('utf-8') # type: ignore + # Add prompt only if provided (some models don't require it) + if prompt is not None and prompt != "": + data["prompt"] = prompt + + # Convert image to base64 if provided + if image is not None: + image_b64: str + if hasattr(image, 'read') and callable(getattr(image, 'read', None)): + # File-like object (e.g., BufferedReader from open()) + image_bytes = image.read() # type: ignore + image_b64 = base64.b64encode(image_bytes).decode('utf-8') # type: ignore + elif isinstance(image, bytes): + # Raw bytes + image_b64 = base64.b64encode(image).decode('utf-8') + elif isinstance(image, str): + # Already a base64 string + image_b64 = image + else: + # Try to handle as bytes + image_b64 = base64.b64encode(bytes(image)).decode('utf-8') # type: ignore - data["image"] = image_b64 + # For style-transfer models, map image to init_image + model_lower = model.lower() + if "style-transfer" in model_lower: + data["init_image"] = image_b64 + else: + data["image"] = image_b64 # Add optional params (already mapped in map_openai_params) for key, value in image_edit_optional_request_params.items(): # type: ignore @@ -221,30 +227,43 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): file_b64 = str(file_bytes) data[key] = file_b64 continue - - # Supported text fields - if key in [ - "negative_prompt", - "aspect_ratio", - "seed", - "output_format", - "model", - "mode", + + # Numeric fields that need to be converted to int/float + numeric_int_fields = ["left", "right", "up", "down", "seed"] + numeric_float_fields = [ "strength", - "style_preset", "creativity", "control_strength", "grow_mask", - "left", - "right", - "up", - "down", - "select_prompt", - "search_prompt", "fidelity", "composition_fidelity", "style_strength", "change_strength", + ] + + if key in numeric_int_fields: + # Convert to int (these are pixel values for outpaint) + try: + data[key] = int(value) # type: ignore + except (ValueError, TypeError): + data[key] = value # type: ignore + elif key in numeric_float_fields: + # Convert to float + try: + data[key] = float(value) # type: ignore + except (ValueError, TypeError): + data[key] = value # type: ignore + + # Supported text fields + elif key in [ + "negative_prompt", + "aspect_ratio", + "output_format", + "model", + "mode", + "style_preset", + "select_prompt", + "search_prompt", ]: data[key] = value # type: ignore diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ab1e735fca7..6a87967c3aa 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -3080,10 +3080,8 @@ class BaseLLMHTTPHandler: transformed_request, bytes ): # Handle traditional file uploads - # Ensure transformed_request is a string for httpx compatibility - if isinstance(transformed_request, bytes): - transformed_request = transformed_request.decode("utf-8") - + # Note: transformed_request can be bytes (for binary files like PDFs) + # or str (for text files like JSONL). httpx handles both correctly. # Use the HTTP method specified by the provider config http_method = provider_config.file_upload_http_method.upper() if http_method == "PUT": diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index 0015155b47f..16541138217 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -81,21 +81,23 @@ class GeminiImageEditConfig(BaseImageEditConfig): self, model: str, prompt: Optional[str], - image: FileTypes, + image: Optional[FileTypes], image_edit_optional_request_params: Dict[str, Any], litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[Dict[str, Any], Optional[RequestFiles]]: - inline_parts = self._prepare_inline_image_parts(image) + inline_parts = self._prepare_inline_image_parts(image) if image else [] if not inline_parts: raise ValueError("Gemini image edit requires at least one image.") - if prompt is None: - raise ValueError("Gemini image edit requires a prompt.") + # Build parts list with image and prompt (if provided) + parts = inline_parts.copy() + if prompt is not None and prompt != "": + parts.append({"text": prompt}) contents = [ { - "parts": inline_parts + [{"text": prompt}], + "parts": parts, } ] diff --git a/litellm/llms/openai/image_edit/dalle2_transformation.py b/litellm/llms/openai/image_edit/dalle2_transformation.py index 13531546d2e..fd697b210ee 100644 --- a/litellm/llms/openai/image_edit/dalle2_transformation.py +++ b/litellm/llms/openai/image_edit/dalle2_transformation.py @@ -31,7 +31,7 @@ class DallE2ImageEditConfig(OpenAIImageEditConfig): self, model: str, prompt: Optional[str], - image: FileTypes, + image: Optional[FileTypes], image_edit_optional_request_params: Dict, litellm_params: GenericLiteLLMParams, headers: dict, @@ -40,18 +40,20 @@ class DallE2ImageEditConfig(OpenAIImageEditConfig): Transform image edit request for DALL-E-2. DALL-E-2 only accepts a single image with field name "image" (not "image[]"). - """ - if prompt is None: - raise ValueError("DALL-E-2 image edit requires a prompt.") - - request = ImageEditRequestParams( - model=model, - image=image, - prompt=prompt, + """ + request_params = { + "model": model, **image_edit_optional_request_params, - ) + } + if image is not None: + request_params["image"] = image + if prompt is not None: + request_params["prompt"] = prompt + + request = ImageEditRequestParams(**request_params) request_dict = cast(Dict, request) + ######################################################### # Separate images and masks as `files` and send other parameters as `data` ######################################################### diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index 9edad9ee2c9..a1e5375d098 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -80,7 +80,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig): self, model: str, prompt: Optional[str], - image: FileTypes, + image: Optional[FileTypes], image_edit_optional_request_params: Dict, litellm_params: GenericLiteLLMParams, headers: dict, @@ -91,15 +91,17 @@ class OpenAIImageEditConfig(BaseImageEditConfig): Handles multipart/form-data for images. Uses "image[]" field name to support multiple images (e.g., for gpt-image-1). """ - if prompt is None: - raise ValueError("OpenAI image edit requires a prompt.") - - request = ImageEditRequestParams( - model=model, - image=image, - prompt=prompt, + # Build request params, only including non-None values + request_params = { + "model": model, **image_edit_optional_request_params, - ) + } + if image is not None: + request_params["image"] = image + if prompt is not None: + request_params["prompt"] = prompt + + request = ImageEditRequestParams(**request_params) request_dict = cast(Dict, request) ######################################################### diff --git a/litellm/llms/recraft/image_edit/transformation.py b/litellm/llms/recraft/image_edit/transformation.py index 9bf46704ed1..d2a56236819 100644 --- a/litellm/llms/recraft/image_edit/transformation.py +++ b/litellm/llms/recraft/image_edit/transformation.py @@ -102,7 +102,7 @@ class RecraftImageEditConfig(BaseImageEditConfig): self, model: str, prompt: Optional[str], - image: FileTypes, + image: Optional[FileTypes], image_edit_optional_request_params: Dict, litellm_params: GenericLiteLLMParams, headers: dict, @@ -114,15 +114,15 @@ class RecraftImageEditConfig(BaseImageEditConfig): https://www.recraft.ai/docs#image-to-image """ - if prompt is None: - raise ValueError("Recraft image edit requires a prompt.") - - request_body: RecraftImageEditRequestParams = RecraftImageEditRequestParams( - model=model, - prompt=prompt, - strength=image_edit_optional_request_params.pop("strength", self.DEFAULT_STRENGTH), + request_params = { + "model": model, + "strength": image_edit_optional_request_params.pop("strength", self.DEFAULT_STRENGTH), **image_edit_optional_request_params, - ) + } + if prompt is not None: + request_params["prompt"] = prompt + + request_body = RecraftImageEditRequestParams(**request_params) request_dict = cast(Dict, request_body) ######################################################### # Reuse OpenAI logic: Separate images as `files` and send other parameters as `data` diff --git a/litellm/llms/replicate/chat/handler.py b/litellm/llms/replicate/chat/handler.py index 4c75db5abc6..c37473b3183 100644 --- a/litellm/llms/replicate/chat/handler.py +++ b/litellm/llms/replicate/chat/handler.py @@ -83,19 +83,27 @@ async def async_handle_prediction_response_streaming( await asyncio.sleep( REPLICATE_POLLING_DELAY_SECONDS ) # prevent being rate limited by replicate - print_verbose(f"replicate: polling endpoint: {prediction_url}") response = await http_client.get(prediction_url, headers=headers) if response.status_code == 200: response_data = response.json() - status = response_data["status"] - if "output" in response_data: + status = response_data.get("status", "") + # Check that "output" exists and is not None or empty + output_present = "output" in response_data and response_data["output"] is not None + if output_present: try: - output_string = "".join(response_data["output"]) + # If output is None or not a list, treat as empty string + if isinstance(response_data["output"], list): + output_string = "".join(response_data["output"]) + elif response_data["output"] is None: + output_string = "" + else: + # fallback for other types; convert to string safely + output_string = str(response_data["output"]) except Exception: raise ReplicateError( status_code=422, message="Unable to parse response. Got={}".format( - response_data["output"] + response_data.get("output", None) ), headers=response.headers, ) @@ -103,7 +111,7 @@ async def async_handle_prediction_response_streaming( print_verbose(f"New chunk: {new_output}") yield {"output": new_output, "status": status} previous_output = output_string - status = response_data["status"] + status = response_data.get("status", "") if status == "failed": replicate_error = response_data.get("error", "") raise ReplicateError( diff --git a/litellm/llms/stability/image_edit/transformations.py b/litellm/llms/stability/image_edit/transformations.py index 013e3f27a02..53bdc825dd4 100644 --- a/litellm/llms/stability/image_edit/transformations.py +++ b/litellm/llms/stability/image_edit/transformations.py @@ -171,7 +171,7 @@ class StabilityImageEditConfig(BaseImageEditConfig): self, model: str, prompt: Optional[str], - image: FileTypes, + image: Optional[FileTypes], image_edit_optional_request_params: Dict, litellm_params: GenericLiteLLMParams, headers: dict, @@ -190,11 +190,14 @@ class StabilityImageEditConfig(BaseImageEditConfig): } # Add prompt only if provided (some Stability endpoints don't require it) - if prompt is not None: + if prompt is not None and prompt != "": data["prompt"] = prompt # Handle image parameter - could be a single file or list image_file = image[0] if isinstance(image, list) else image # type: ignore - files: Dict[str, Any] = {"image": image_file} + files: Dict[str, Any] = {} + if image is not None: + image_file = image[0] if isinstance(image, list) else image # type: ignore + files["image"] = image_file # Add optional params (already mapped in map_openai_params) for key, value in image_edit_optional_request_params.items(): # type: ignore diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 5aa7662f175..5b09bcc0289 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -453,9 +453,10 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False): valid_schema_fields = set(get_type_hints(Schema).keys()) defs = parameters.pop("$defs", {}) - # flatten the defs - for name, value in defs.items(): - unpack_defs(value, defs) + # Expand $ref references in parameters using the definitions + # Note: We don't pre-flatten defs as that causes exponential memory growth + # with circular references (see issue #19098). unpack_defs handles nested + # refs recursively and correctly detects/skips circular references. unpack_defs(parameters, defs) # 5. Nullable fields: diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 154d5669eb8..8fcd285824d 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -152,22 +152,24 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): self, model: str, prompt: Optional[str], - image: FileTypes, + image: Optional[FileTypes], image_edit_optional_request_params: Dict[str, Any], litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[Dict[str, Any], Optional[RequestFiles]]: - inline_parts = self._prepare_inline_image_parts(image) + inline_parts = self._prepare_inline_image_parts(image) if image else [] if not inline_parts: raise ValueError("Vertex AI Gemini image edit requires at least one image.") - if prompt is None: - raise ValueError("Vertex AI Gemini image edit requires a prompt.") + # Build parts list with image and prompt (if provided) + parts = inline_parts.copy() + if prompt is not None and prompt != "": + parts.append({"text": prompt}) # Correct format for Vertex AI Gemini image editing contents = { "role": "USER", - "parts": inline_parts + [{"text": prompt}] + "parts": parts } request_body: Dict[str, Any] = {"contents": contents} diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index 337a4bd4dd6..b58825e1faa 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -144,7 +144,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): self, model: str, prompt: Optional[str], - image: FileTypes, + image: Optional[FileTypes], image_edit_optional_request_params: Dict[str, Any], litellm_params: GenericLiteLLMParams, headers: dict, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 470d598a25f..43f7bde5da3 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -7857,6 +7857,24 @@ "supports_tool_choice": true, "supports_vision": true }, + "dall-e-2": { + "input_cost_per_image": 0.02, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits", + "/v1/images/variations" + ] + }, + "dall-e-3": { + "input_cost_per_image": 0.04, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "deepseek-chat": { "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.8e-07, diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404.html deleted file mode 100644 index c6035eb40ca..00000000000 --- a/litellm/proxy/_experimental/out/404.html +++ /dev/null @@ -1 +0,0 @@ -404: This page could not be found.LiteLLM Dashboard

404

This page could not be found.

\ No newline at end of file diff --git a/litellm/proxy/_experimental/out/api-reference.html b/litellm/proxy/_experimental/out/api-reference/index.html similarity index 100% rename from litellm/proxy/_experimental/out/api-reference.html rename to litellm/proxy/_experimental/out/api-reference/index.html diff --git a/litellm/proxy/_experimental/out/experimental/api-playground.html b/litellm/proxy/_experimental/out/experimental/api-playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/api-playground.html rename to litellm/proxy/_experimental/out/experimental/api-playground/index.html diff --git a/litellm/proxy/_experimental/out/experimental/budgets.html b/litellm/proxy/_experimental/out/experimental/budgets/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/budgets.html rename to litellm/proxy/_experimental/out/experimental/budgets/index.html diff --git a/litellm/proxy/_experimental/out/experimental/caching.html b/litellm/proxy/_experimental/out/experimental/caching/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/caching.html rename to litellm/proxy/_experimental/out/experimental/caching/index.html diff --git a/litellm/proxy/_experimental/out/experimental/old-usage.html b/litellm/proxy/_experimental/out/experimental/old-usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/old-usage.html rename to litellm/proxy/_experimental/out/experimental/old-usage/index.html diff --git a/litellm/proxy/_experimental/out/experimental/prompts.html b/litellm/proxy/_experimental/out/experimental/prompts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/prompts.html rename to litellm/proxy/_experimental/out/experimental/prompts/index.html diff --git a/litellm/proxy/_experimental/out/experimental/tag-management.html b/litellm/proxy/_experimental/out/experimental/tag-management/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/tag-management.html rename to litellm/proxy/_experimental/out/experimental/tag-management/index.html diff --git a/litellm/proxy/_experimental/out/guardrails.html b/litellm/proxy/_experimental/out/guardrails.html deleted file mode 100644 index 6dae8818641..00000000000 --- a/litellm/proxy/_experimental/out/guardrails.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_experimental/out/login.html b/litellm/proxy/_experimental/out/login/index.html similarity index 100% rename from litellm/proxy/_experimental/out/login.html rename to litellm/proxy/_experimental/out/login/index.html diff --git a/litellm/proxy/_experimental/out/logs.html b/litellm/proxy/_experimental/out/logs/index.html similarity index 100% rename from litellm/proxy/_experimental/out/logs.html rename to litellm/proxy/_experimental/out/logs/index.html diff --git a/litellm/proxy/_experimental/out/mcp/oauth/callback.html b/litellm/proxy/_experimental/out/mcp/oauth/callback/index.html similarity index 100% rename from litellm/proxy/_experimental/out/mcp/oauth/callback.html rename to litellm/proxy/_experimental/out/mcp/oauth/callback/index.html diff --git a/litellm/proxy/_experimental/out/model-hub.html b/litellm/proxy/_experimental/out/model-hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model-hub.html rename to litellm/proxy/_experimental/out/model-hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub.html deleted file mode 100644 index fb4d18d07a2..00000000000 --- a/litellm/proxy/_experimental/out/model_hub.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/models-and-endpoints.html b/litellm/proxy/_experimental/out/models-and-endpoints/index.html similarity index 100% rename from litellm/proxy/_experimental/out/models-and-endpoints.html rename to litellm/proxy/_experimental/out/models-and-endpoints/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html deleted file mode 100644 index 77ec4ff6f48..00000000000 --- a/litellm/proxy/_experimental/out/onboarding.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_experimental/out/organizations.html b/litellm/proxy/_experimental/out/organizations/index.html similarity index 100% rename from litellm/proxy/_experimental/out/organizations.html rename to litellm/proxy/_experimental/out/organizations/index.html diff --git a/litellm/proxy/_experimental/out/playground.html b/litellm/proxy/_experimental/out/playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/playground.html rename to litellm/proxy/_experimental/out/playground/index.html diff --git a/litellm/proxy/_experimental/out/settings/admin-settings.html b/litellm/proxy/_experimental/out/settings/admin-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/admin-settings.html rename to litellm/proxy/_experimental/out/settings/admin-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/logging-and-alerts.html b/litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/logging-and-alerts.html rename to litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html diff --git a/litellm/proxy/_experimental/out/settings/router-settings.html b/litellm/proxy/_experimental/out/settings/router-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/router-settings.html rename to litellm/proxy/_experimental/out/settings/router-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/ui-theme.html b/litellm/proxy/_experimental/out/settings/ui-theme/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/ui-theme.html rename to litellm/proxy/_experimental/out/settings/ui-theme/index.html diff --git a/litellm/proxy/_experimental/out/teams.html b/litellm/proxy/_experimental/out/teams/index.html similarity index 100% rename from litellm/proxy/_experimental/out/teams.html rename to litellm/proxy/_experimental/out/teams/index.html diff --git a/litellm/proxy/_experimental/out/test-key.html b/litellm/proxy/_experimental/out/test-key/index.html similarity index 100% rename from litellm/proxy/_experimental/out/test-key.html rename to litellm/proxy/_experimental/out/test-key/index.html diff --git a/litellm/proxy/_experimental/out/tools/mcp-servers.html b/litellm/proxy/_experimental/out/tools/mcp-servers/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/mcp-servers.html rename to litellm/proxy/_experimental/out/tools/mcp-servers/index.html diff --git a/litellm/proxy/_experimental/out/tools/vector-stores.html b/litellm/proxy/_experimental/out/tools/vector-stores/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/vector-stores.html rename to litellm/proxy/_experimental/out/tools/vector-stores/index.html diff --git a/litellm/proxy/_experimental/out/usage.html b/litellm/proxy/_experimental/out/usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/usage.html rename to litellm/proxy/_experimental/out/usage/index.html diff --git a/litellm/proxy/_experimental/out/users.html b/litellm/proxy/_experimental/out/users/index.html similarity index 100% rename from litellm/proxy/_experimental/out/users.html rename to litellm/proxy/_experimental/out/users/index.html diff --git a/litellm/proxy/_experimental/out/virtual-keys.html b/litellm/proxy/_experimental/out/virtual-keys/index.html similarity index 100% rename from litellm/proxy/_experimental/out/virtual-keys.html rename to litellm/proxy/_experimental/out/virtual-keys/index.html diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index c2d53b40b7b..a21f3291d31 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -55,9 +55,30 @@ async def _handle_stream_message( proxy_server_request: Optional[dict] = None, ) -> StreamingResponse: """Handle message/stream method via SDK functions.""" - from a2a.types import MessageSendParams, SendStreamingMessageRequest - from litellm.a2a_protocol import asend_message_streaming + from litellm.a2a_protocol.main import A2A_SDK_AVAILABLE + + # Check is handled in invoke_agent_a2a, but if called directly: + if not A2A_SDK_AVAILABLE: + # Return a streaming response that yields an error + async def _error_stream(): + yield json.dumps( + { + "jsonrpc": "2.0", + "id": request_id, + "error": { + "code": -32603, + "message": "Server error: 'a2a' package not installed", + }, + } + ) + "\n" + + return StreamingResponse(_error_stream(), media_type="application/x-ndjson") + + from a2a.types import ( + MessageSendParams, + SendStreamingMessageRequest, + ) async def stream_response(): try: @@ -75,16 +96,20 @@ async def _handle_stream_message( ): # Chunk may be dict or object depending on bridge vs standard path if hasattr(chunk, "model_dump"): - yield json.dumps(chunk.model_dump(mode="json", exclude_none=True)) + "\n" + yield json.dumps( + chunk.model_dump(mode="json", exclude_none=True) + ) + "\n" else: yield json.dumps(chunk) + "\n" except Exception as e: verbose_proxy_logger.exception(f"Error streaming A2A response: {e}") - yield json.dumps({ - "jsonrpc": "2.0", - "id": request_id, - "error": {"code": -32603, "message": f"Streaming error: {str(e)}"}, - }) + "\n" + yield json.dumps( + { + "jsonrpc": "2.0", + "id": request_id, + "error": {"code": -32603, "message": f"Streaming error: {str(e)}"}, + } + ) + "\n" return StreamingResponse(stream_response(), media_type="application/x-ndjson") @@ -169,9 +194,8 @@ async def invoke_agent_a2a( - message/send: Send a message and get a response - message/stream: Send a message and stream the response """ - from a2a.types import MessageSendParams, SendMessageRequest - from litellm.a2a_protocol import asend_message + from litellm.a2a_protocol.main import A2A_SDK_AVAILABLE from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, ) @@ -189,16 +213,28 @@ async def invoke_agent_a2a( # Validate JSON-RPC format if body.get("jsonrpc") != "2.0": - return _jsonrpc_error(body.get("id"), -32600, "Invalid Request: jsonrpc must be '2.0'") + return _jsonrpc_error( + body.get("id"), -32600, "Invalid Request: jsonrpc must be '2.0'" + ) request_id = body.get("id") method = body.get("method") params = body.get("params", {}) + if not A2A_SDK_AVAILABLE: + return _jsonrpc_error( + request_id, + -32603, + "Server error: 'a2a' package not installed. Please install 'a2a-sdk'.", + 500, + ) + # Find the agent agent = _get_agent(agent_id) if agent is None: - return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' not found", 404) + return _jsonrpc_error( + request_id, -32000, f"Agent '{agent_id}' not found", 404 + ) is_allowed = await AgentRequestHandler.is_agent_allowed( agent_id=agent.agent_id, @@ -213,23 +249,29 @@ async def invoke_agent_a2a( # Get backend URL and agent name agent_url = agent.agent_card_params.get("url") agent_name = agent.agent_card_params.get("name", agent_id) - + # Get litellm_params (may include custom_llm_provider for completion bridge) litellm_params = agent.litellm_params or {} custom_llm_provider = litellm_params.get("custom_llm_provider") - + # URL is required unless using completion bridge with a provider that derives endpoint from model # (e.g., bedrock/agentcore derives endpoint from ARN in model string) if not agent_url and not custom_llm_provider: - return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' has no URL configured", 500) + return _jsonrpc_error( + request_id, -32000, f"Agent '{agent_id}' has no URL configured", 500 + ) - verbose_proxy_logger.info(f"Proxying A2A request to agent '{agent_id}' at {agent_url or 'completion-bridge'}") + verbose_proxy_logger.info( + f"Proxying A2A request to agent '{agent_id}' at {agent_url or 'completion-bridge'}" + ) # Set up data dict for litellm processing - body.update({ - "model": f"a2a_agent/{agent_name}", - "custom_llm_provider": "a2a_agent", - }) + body.update( + { + "model": f"a2a_agent/{agent_name}", + "custom_llm_provider": "a2a_agent", + } + ) # Add litellm data (user_api_key, user_id, team_id, etc.) data = await add_litellm_data_to_request( @@ -243,6 +285,8 @@ async def invoke_agent_a2a( # Route through SDK functions if method == "message/send": + from a2a.types import MessageSendParams, SendMessageRequest + a2a_request = SendMessageRequest( id=request_id, params=MessageSendParams(**params), @@ -255,7 +299,9 @@ async def invoke_agent_a2a( metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), ) - return JSONResponse(content=response.model_dump(mode="json", exclude_none=True)) + return JSONResponse( + content=response.model_dump(mode="json", exclude_none=True) + ) elif method == "message/stream": return await _handle_stream_message( diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index cbe28849b1e..3cbca27ce0c 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -1,5 +1,6 @@ #### Analytics Endpoints ##### import os + from fastapi import APIRouter from litellm.types.proxy.discovery_endpoints.ui_discovery_endpoints import ( @@ -14,10 +15,12 @@ router = APIRouter() "/litellm/.well-known/litellm-ui-config", response_model=UiDiscoveryEndpoints ) # if mounted at root path async def get_ui_config(): - from litellm.proxy.utils import get_proxy_base_url, get_server_root_path from litellm.proxy.auth.auth_utils import _has_user_setup_sso + from litellm.proxy.utils import get_proxy_base_url, get_server_root_path - auto_redirect_ui_login_to_sso = os.getenv("AUTO_REDIRECT_UI_LOGIN_TO_SSO", "true").lower() == "true" + auto_redirect_ui_login_to_sso = ( + os.getenv("AUTO_REDIRECT_UI_LOGIN_TO_SSO", "true").lower() == "true" + ) admin_ui_disabled = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true" return UiDiscoveryEndpoints( diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index a1453e10dbf..4a2c05f8590 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -244,8 +244,10 @@ async def image_edit_api( if mask is None and mask_array is not None: mask = mask_array - if image is None: - raise HTTPException(status_code=422, detail="Field required: image") + # if image is None: + # raise HTTPException(status_code=422, detail="Field required: image") + # Note: Image is optional for some models (e.g., Bedrock Stability style-transfer) + # The validation will be done at the model level if image is truly required from litellm.proxy.proxy_server import ( _read_request_body, @@ -272,6 +274,10 @@ async def image_edit_api( data["image"] = image_files if mask_files: data["mask"] = mask_files + + # Ensure prompt exists in data (default to None for models that don't require it) + if "prompt" not in data: + data["prompt"] = None data["model"] = ( model diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 7e3f5820814..2c6b378ae38 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -146,8 +146,8 @@ async def route_create_file( Priority: 1. If target_storage is specified and not "default" -> use storage backend 2. If model parameter provided -> use model credentials and encode ID - 3. If enable_loadbalancing_on_batch_endpoints -> deprecated loadbalancing - 4. If target_model_names_list -> managed files (requires DB) + 3. If target_model_names_list -> managed files (requires DB, supports loadbalancing) + 4. If enable_loadbalancing_on_batch_endpoints -> deprecated loadbalancing 5. Else -> use custom_llm_provider with files_settings """ @@ -202,18 +202,9 @@ async def route_create_file( return response - # EXISTING: Deprecated loadbalancing approach - if ( - litellm.enable_loadbalancing_on_batch_endpoints is True - and is_router_model - and router_model is not None - ): - response = await _deprecated_loadbalanced_create_file( - llm_router=llm_router, - router_model=router_model, - _create_file_request=_create_file_request, - ) - elif target_model_names_list: + # Handle managed files (supports loadbalancing via llm_router.acreate_file) + # Priority: Check for managed files BEFORE deprecated loadbalancing + if target_model_names_list: managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") if managed_files_obj is None: raise ProxyException( @@ -236,6 +227,7 @@ async def route_create_file( param="None", code=500, ) + # Managed files internally calls llm_router.acreate_file() which includes loadbalancing response = await managed_files_obj.acreate_file( llm_router=llm_router, create_file_request=_create_file_request, @@ -243,6 +235,17 @@ async def route_create_file( litellm_parent_otel_span=user_api_key_dict.parent_otel_span, user_api_key_dict=user_api_key_dict, ) + # EXISTING: Deprecated loadbalancing approach (for backwards compatibility when not using managed files) + elif ( + litellm.enable_loadbalancing_on_batch_endpoints is True + and is_router_model + and router_model is not None + ): + response = await _deprecated_loadbalanced_create_file( + llm_router=llm_router, + router_model=router_model, + _create_file_request=_create_file_request, + ) else: # get configs for custom_llm_provider llm_provider_config = get_files_provider_config( diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 7db76fd31dd..09c14bc42c1 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -1,8 +1,8 @@ #### CRUD ENDPOINTS for UI Settings ##### import json -from typing import Any, Dict, List, Union, Optional +from typing import Any, Dict, List, Optional, Union -from fastapi import APIRouter, Depends, HTTPException, UploadFile, File +from fastapi import APIRouter, Depends, File, HTTPException, UploadFile import litellm from litellm._logging import verbose_proxy_logger @@ -10,6 +10,7 @@ from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.proxy.management_endpoints.ui_sso import ( DefaultTeamSSOParams, + InProductNudgeResponse, SSOConfig, ) @@ -22,11 +23,11 @@ class IPAddress(BaseModel): class UIThemeConfig(BaseModel): """Configuration for UI theme customization""" - + # Logo configuration logo_url: Optional[str] = Field( default=None, - description="URL or path to custom logo image. Can be a local file path or HTTP/HTTPS URL" + description="URL or path to custom logo image. Can be a local file path or HTTP/HTTPS URL", ) @@ -85,7 +86,10 @@ class UISettingsResponse(SettingsResponse): # Allowlist of UI settings that can be stored -ALLOWED_UI_SETTINGS_FIELDS = {"disable_model_add_for_internal_users", "disable_team_admin_delete_team_user"} +ALLOWED_UI_SETTINGS_FIELDS = { + "disable_model_add_for_internal_users", + "disable_team_admin_delete_team_user", +} @router.get( @@ -434,7 +438,7 @@ async def get_sso_settings(): # Initialize with defaults sso_settings_dict = {} - + if sso_db_record and sso_db_record.sso_settings: # Load settings from database sso_settings_dict = dict(sso_db_record.sso_settings) @@ -444,26 +448,43 @@ async def get_sso_settings(): role_mappings = None if role_mappings_data: from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings + if isinstance(role_mappings_data, dict): role_mappings = RoleMappings(**role_mappings_data) elif isinstance(role_mappings_data, RoleMappings): role_mappings = role_mappings_data - - decrypted_sso_settings_dict = proxy_config._decrypt_and_set_db_env_variables(environment_variables=sso_settings_dict) + + decrypted_sso_settings_dict = proxy_config._decrypt_and_set_db_env_variables( + environment_variables=sso_settings_dict + ) # Build SSO config with database values or environment fallback - + sso_config = SSOConfig( google_client_id=decrypted_sso_settings_dict.get("google_client_id", None), - google_client_secret=decrypted_sso_settings_dict.get("google_client_secret", None), - microsoft_client_id=decrypted_sso_settings_dict.get("microsoft_client_id", None), - microsoft_client_secret=decrypted_sso_settings_dict.get("microsoft_client_secret", None), + google_client_secret=decrypted_sso_settings_dict.get( + "google_client_secret", None + ), + microsoft_client_id=decrypted_sso_settings_dict.get( + "microsoft_client_id", None + ), + microsoft_client_secret=decrypted_sso_settings_dict.get( + "microsoft_client_secret", None + ), microsoft_tenant=decrypted_sso_settings_dict.get("microsoft_tenant", None), generic_client_id=decrypted_sso_settings_dict.get("generic_client_id", None), - generic_client_secret=decrypted_sso_settings_dict.get("generic_client_secret", None), - generic_authorization_endpoint=decrypted_sso_settings_dict.get("generic_authorization_endpoint", None), - generic_token_endpoint=decrypted_sso_settings_dict.get("generic_token_endpoint", None), - generic_userinfo_endpoint=decrypted_sso_settings_dict.get("generic_userinfo_endpoint", None), + generic_client_secret=decrypted_sso_settings_dict.get( + "generic_client_secret", None + ), + generic_authorization_endpoint=decrypted_sso_settings_dict.get( + "generic_authorization_endpoint", None + ), + generic_token_endpoint=decrypted_sso_settings_dict.get( + "generic_token_endpoint", None + ), + generic_userinfo_endpoint=decrypted_sso_settings_dict.get( + "generic_userinfo_endpoint", None + ), proxy_base_url=decrypted_sso_settings_dict.get("proxy_base_url", None), user_email=decrypted_sso_settings_dict.get("user_email"), ui_access_mode=decrypted_sso_settings_dict.get("ui_access_mode"), @@ -506,10 +527,14 @@ async def update_sso_settings(sso_config: SSOConfig): """ Update SSO configuration by saving to the dedicated SSO table. """ - import os import json + import os - from litellm.proxy.proxy_server import prisma_client, store_model_in_db, proxy_config + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_config, + store_model_in_db, + ) if prisma_client is None: raise HTTPException( @@ -562,7 +587,9 @@ async def update_sso_settings(sso_config: SSOConfig): # Clear environment variable if value is null/empty os.environ.pop(env_var_name, None) - encrypted_sso_data = proxy_config._encrypt_env_variables(environment_variables=sso_data) + encrypted_sso_data = proxy_config._encrypt_env_variables( + environment_variables=sso_data + ) # Save to dedicated SSO table await prisma_client.db.litellm_ssoconfig.upsert( @@ -655,9 +682,10 @@ async def update_ui_theme_settings(theme_config: UIThemeConfig): Update UI theme configuration. Updates logo settings for the admin UI. """ - from litellm.proxy.proxy_server import proxy_config, store_model_in_db import os + from litellm.proxy.proxy_server import proxy_config, store_model_in_db + if store_model_in_db is not True: raise HTTPException( status_code=500, @@ -668,28 +696,30 @@ async def update_ui_theme_settings(theme_config: UIThemeConfig): # Load existing config config = await proxy_config.get_config() - + # Update config with UI theme settings if "general_settings" not in config: config["general_settings"] = {} - + if "environment_variables" not in config: config["environment_variables"] = {} # Convert theme config to dict theme_data = theme_config.model_dump(exclude_none=True) - + # Store UI theme config in litellm_settings (where it's retrieved from) if "litellm_settings" not in config: config["litellm_settings"] = {} config["litellm_settings"]["ui_theme_config"] = theme_data - + # Update UI_LOGO_PATH environment variable if logo_url is provided # If logo_url is empty string, None, or null, remove the environment variable to use default logo_url = theme_data.get("logo_url") verbose_proxy_logger.debug(f"Updating logo_url: {logo_url}") - - if logo_url and isinstance(logo_url, str) and logo_url.strip(): # Check if logo_url exists and is not empty/whitespace + + if ( + logo_url and isinstance(logo_url, str) and logo_url.strip() + ): # Check if logo_url exists and is not empty/whitespace config["environment_variables"]["UI_LOGO_PATH"] = logo_url os.environ["UI_LOGO_PATH"] = logo_url verbose_proxy_logger.debug(f"Set UI_LOGO_PATH to: {logo_url}") @@ -704,12 +734,15 @@ async def update_ui_theme_settings(theme_config: UIThemeConfig): # Handle environment variable encryption if needed stored_config = config.copy() - if "environment_variables" in stored_config and len(stored_config["environment_variables"]) > 0: + if ( + "environment_variables" in stored_config + and len(stored_config["environment_variables"]) > 0 + ): # Only encrypt if there are environment variables to encrypt stored_config["environment_variables"] = proxy_config._encrypt_env_variables( environment_variables=stored_config["environment_variables"] ) - + # Save the updated config await proxy_config.save_config(new_config=stored_config) @@ -720,6 +753,34 @@ async def update_ui_theme_settings(theme_config: UIThemeConfig): } +@router.get( + "/in_product_nudges", + tags=["UI Settings"], + dependencies=[Depends(user_api_key_auth)], + response_model=InProductNudgeResponse, +) +async def get_in_product_nudges(): + """ + Get in-product nudges configuration. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected. Please connect a database."}, + ) + + db_record = await prisma_client.db.litellm_dailytagspend.find_first( + where={"tag": "User-Agent: claude-cli"} + ) + + if db_record: + return InProductNudgeResponse(is_claude_code_enabled=True) + + return InProductNudgeResponse(is_claude_code_enabled=False) + + @router.get( "/get/ui_settings", tags=["UI Settings"], @@ -752,7 +813,9 @@ async def get_ui_settings(): ui_settings = dict(ui_settings_json) # Sanitize any unexpected keys from persisted config before returning - ui_settings = {k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS} + ui_settings = { + k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS + } # Build config-like object for schema helper config: Dict[str, Any] = {"litellm_settings": {"ui_settings": ui_settings}} @@ -823,6 +886,7 @@ async def update_ui_settings( "settings": ui_settings, } + @router.post( "/upload/logo", tags=["UI Theme Settings"], @@ -839,35 +903,35 @@ async def upload_logo(file: UploadFile = File(...)): # Validate file type allowed_extensions = {".png", ".jpg", ".jpeg", ".svg"} file_extension = Path(file.filename or "").suffix.lower() - + if file_extension not in allowed_extensions: raise HTTPException( status_code=400, - detail=f"Invalid file type. Allowed types: {', '.join(allowed_extensions)}" + detail=f"Invalid file type. Allowed types: {', '.join(allowed_extensions)}", ) - + # Validate file size (max 5MB) file_content = await file.read() if len(file_content) > 5 * 1024 * 1024: # 5MB raise HTTPException( - status_code=400, - detail="File size too large. Maximum size is 5MB." + status_code=400, detail="File size too large. Maximum size is 5MB." ) - + # Create uploads directory if it doesn't exist current_dir = os.path.dirname(os.path.abspath(__file__)) upload_dir = os.path.join(current_dir, "..", "uploads") os.makedirs(upload_dir, exist_ok=True) - + # Generate unique filename from litellm._uuid import uuid + unique_filename = f"logo_{uuid.uuid4().hex}{file_extension}" file_path = os.path.join(upload_dir, unique_filename) - + # Save the file with open(file_path, "wb") as buffer: buffer.write(file_content) - + return { "message": "Logo uploaded successfully", "status": "success", diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 661f94e5f04..bc61a60fe5a 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -245,6 +245,7 @@ async def list_vector_stores( """ List all available vector stores with optional filtering and pagination. Combines both in-memory vector stores and those stored in the database. + Database is the source of truth - deleted stores are removed from memory, updated stores sync to memory. Parameters: - page: int - Page number for pagination (default: 1) @@ -252,29 +253,65 @@ async def list_vector_stores( """ from litellm.proxy.proxy_server import prisma_client - seen_vector_store_ids = set() + vector_store_map: Dict[str, LiteLLM_ManagedVectorStore] = {} + db_vector_store_ids: set = set() try: - # Get in-memory vector stores - in_memory_vector_stores: List[LiteLLM_ManagedVectorStore] = [] + # Get vector stores from database first (source of truth) + vector_stores_from_db = await VectorStoreRegistry._get_vector_stores_from_db( + prisma_client=prisma_client + ) + + # Build map from database vector stores + for vector_store in vector_stores_from_db: + vector_store_id = vector_store.get("vector_store_id", None) + if vector_store_id: + vector_store_map[vector_store_id] = vector_store + db_vector_store_ids.add(vector_store_id) + + # Process in-memory vector stores if litellm.vector_store_registry is not None: in_memory_vector_stores = copy.deepcopy( litellm.vector_store_registry.vector_stores ) + + vector_stores_to_delete_from_memory: List[str] = [] + + for vector_store in in_memory_vector_stores: + vector_store_id = vector_store.get("vector_store_id", None) + if not vector_store_id: + continue + + # If vector store is in memory but NOT in database, it was deleted + if vector_store_id not in db_vector_store_ids: + verbose_proxy_logger.info( + f"Vector store {vector_store_id} exists in memory but not in database - marking for deletion from cache" + ) + vector_stores_to_delete_from_memory.append(vector_store_id) + # If not in our map yet, add it (only in-memory, not in DB) + elif vector_store_id not in vector_store_map: + vector_store_map[vector_store_id] = vector_store + + # Synchronize in-memory registry with database + # 1. Remove deleted vector stores from memory + for vs_id in vector_stores_to_delete_from_memory: + litellm.vector_store_registry.delete_vector_store_from_registry( + vector_store_id=vs_id + ) + verbose_proxy_logger.debug( + f"Removed deleted vector store {vs_id} from in-memory registry" + ) + + # 2. Update in-memory registry with database versions (for updates) + for vector_store in vector_stores_from_db: + vector_store_id = vector_store.get("vector_store_id", None) + if vector_store_id: + litellm.vector_store_registry.update_vector_store_in_registry( + vector_store_id=vector_store_id, + updated_data=vector_store + ) - # Get vector stores from database - vector_stores_from_db = await VectorStoreRegistry._get_vector_stores_from_db( - prisma_client=prisma_client - ) - - # Combine in-memory and database vector stores - combined_vector_stores: List[LiteLLM_ManagedVectorStore] = [] - for vector_store in in_memory_vector_stores + vector_stores_from_db: - vector_store_id = vector_store.get("vector_store_id", None) - if vector_store_id not in seen_vector_store_ids: - combined_vector_stores.append(vector_store) - seen_vector_store_ids.add(vector_store_id) - + combined_vector_stores = list(vector_store_map.values()) total_count = len(combined_vector_stores) total_pages = (total_count + page_size - 1) // page_size @@ -303,7 +340,7 @@ async def delete_vector_store( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - Delete a vector store. + Delete a vector store from both database and in-memory registry. Parameters: - vector_store_id: str - ID of the vector store to delete @@ -314,31 +351,53 @@ async def delete_vector_store( raise HTTPException(status_code=500, detail="Database not connected") try: - # Check if vector store exists + # Check if vector store exists in database or in-memory registry + db_vector_store_exists = False + memory_vector_store_exists = False + existing_vector_store = ( await prisma_client.db.litellm_managedvectorstorestable.find_unique( where={"vector_store_id": data.vector_store_id} ) ) - if existing_vector_store is None: + if existing_vector_store is not None: + db_vector_store_exists = True + + # Check in-memory registry + if litellm.vector_store_registry is not None: + memory_vector_store = litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry( + vector_store_id=data.vector_store_id + ) + if memory_vector_store is not None: + memory_vector_store_exists = True + + # If not found in either location, raise 404 + if not db_vector_store_exists and not memory_vector_store_exists: raise HTTPException( status_code=404, detail=f"Vector store with ID {data.vector_store_id} not found", ) - # Delete vector store - await prisma_client.db.litellm_managedvectorstorestable.delete( - where={"vector_store_id": data.vector_store_id} - ) + # Delete from database if exists + if db_vector_store_exists: + await prisma_client.db.litellm_managedvectorstorestable.delete( + where={"vector_store_id": data.vector_store_id} + ) - # Delete vector store from registry - if litellm.vector_store_registry is not None: + # Delete from in-memory registry if exists + if memory_vector_store_exists and litellm.vector_store_registry is not None: litellm.vector_store_registry.delete_vector_store_from_registry( vector_store_id=data.vector_store_id ) - return {"message": f"Vector store {data.vector_store_id} deleted successfully"} + return { + "status": "success", + "message": f"Vector store {data.vector_store_id} deleted successfully" + } + except HTTPException: + raise except Exception as e: + verbose_proxy_logger.exception(f"Error deleting vector store: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) @@ -415,8 +474,12 @@ async def update_vector_store( data: VectorStoreUpdateRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - """Update vector store details""" + """ + Update vector store details in both database and in-memory registry. + The updated data is immediately synchronized to the in-memory registry. + """ from litellm.proxy.proxy_server import prisma_client + from litellm.types.router import GenericLiteLLMParams if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") @@ -424,11 +487,36 @@ async def update_vector_store( try: update_data = data.model_dump(exclude_unset=True) vector_store_id = update_data.pop("vector_store_id") + + # Handle metadata serialization if update_data.get("vector_store_metadata") is not None: update_data["vector_store_metadata"] = safe_dumps( update_data["vector_store_metadata"] ) + + # Handle litellm_params if provided + if "litellm_params" in update_data: + _input_litellm_params: dict = update_data.get("litellm_params", {}) or {} + + # Auto-resolve embedding config if embedding model is provided but config is not + embedding_model = _input_litellm_params.get("litellm_embedding_model") + if embedding_model and not _input_litellm_params.get("litellm_embedding_config"): + resolved_config = await _resolve_embedding_config_from_db( + embedding_model=embedding_model, + prisma_client=prisma_client + ) + if resolved_config: + _input_litellm_params["litellm_embedding_config"] = resolved_config + verbose_proxy_logger.info( + f"Auto-resolved embedding config for model {embedding_model}" + ) + + litellm_params_dict = GenericLiteLLMParams( + **_input_litellm_params + ).model_dump(exclude_none=True) + update_data["litellm_params"] = safe_dumps(litellm_params_dict) + # Update in database updated = await prisma_client.db.litellm_managedvectorstorestable.update( where={"vector_store_id": vector_store_id}, data=update_data, @@ -436,13 +524,21 @@ async def update_vector_store( updated_vs = LiteLLM_ManagedVectorStore(**updated.model_dump()) + # Immediately update in-memory registry to keep it in sync if litellm.vector_store_registry is not None: litellm.vector_store_registry.update_vector_store_in_registry( vector_store_id=vector_store_id, updated_data=updated_vs, ) + verbose_proxy_logger.debug( + f"Updated vector store {vector_store_id} in both database and in-memory registry" + ) - return {"vector_store": updated_vs} + return { + "status": "success", + "message": f"Vector store {vector_store_id} updated successfully", + "vector_store": updated_vs + } except Exception as e: verbose_proxy_logger.exception(f"Error updating vector store: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index ad910c9cd97..eaa80c6cfe4 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2,7 +2,8 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion API) """ -from typing import Any, Dict, List, Literal, Optional, Tuple, Union, cast +from collections.abc import Sequence +from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, cast from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.tool_param import FunctionToolParam @@ -378,6 +379,7 @@ class LiteLLMCompletionResponsesConfig: if isinstance(input, str): messages.append(ChatCompletionUserMessage(role="user", content=input)) elif isinstance(input, list): + existing_tool_call_ids: Set[str] = set() for _input in input: chat_completion_messages = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( input_item=_input @@ -390,12 +392,98 @@ class LiteLLMCompletionResponsesConfig: input_item=_input ): tool_call_output_messages.extend(chat_completion_messages) - else: - messages.extend(chat_completion_messages) + continue - messages.extend(tool_call_output_messages) + if LiteLLMCompletionResponsesConfig._is_input_item_function_call( + input_item=_input + ): + call_id_raw = _input.get("call_id") or _input.get("id") or "" + if call_id_raw: + existing_tool_call_ids.add(str(call_id_raw)) + + messages.extend(chat_completion_messages) + + deduped_tool_call_messages = ( + LiteLLMCompletionResponsesConfig._deduplicate_tool_call_output_messages( + tool_call_output_messages=tool_call_output_messages, + existing_tool_call_ids=existing_tool_call_ids, + ) + ) + messages.extend(deduped_tool_call_messages) return messages + @staticmethod + def _deduplicate_tool_call_output_messages( + tool_call_output_messages: List[ + Union[ + AllMessageValues, + GenericChatCompletionMessage, + ChatCompletionMessageToolCall, + ChatCompletionResponseMessage, + ] + ], + existing_tool_call_ids: Set[str], + ) -> List[ + Union[ + AllMessageValues, + GenericChatCompletionMessage, + ChatCompletionMessageToolCall, + ChatCompletionResponseMessage, + ] + ]: + """Return tool call outputs after dropping assistant entries with duplicate call_ids.""" + if not tool_call_output_messages: + return [] + + filtered_messages: List[ + Union[ + AllMessageValues, + GenericChatCompletionMessage, + ChatCompletionMessageToolCall, + ChatCompletionResponseMessage, + ] + ] = [] + seen_tool_call_ids: Set[str] = set(existing_tool_call_ids) + + for tool_call_message in tool_call_output_messages: + if isinstance(tool_call_message, dict): + role = tool_call_message.get("role", "") + else: + role = getattr(tool_call_message, "role", "") + call_id = "" + + if role == "assistant": + tool_calls: Any = None + if isinstance(tool_call_message, dict): + tool_calls = tool_call_message.get("tool_calls") + else: + tool_calls = getattr(tool_call_message, "tool_calls", None) + + if ( + isinstance(tool_calls, Sequence) + and not isinstance(tool_calls, (str, bytes)) + and len(tool_calls) > 0 + ): + first_call = tool_calls[0] + call_id_raw = None + if isinstance(first_call, dict): + call_id_raw = first_call.get("id") + else: + call_id_raw = getattr(first_call, "id", None) + + if call_id_raw: + call_id = str(call_id_raw) + + if call_id and call_id in seen_tool_call_ids and role == "assistant": + continue + + if call_id and role == "assistant": + seen_tool_call_ids.add(call_id) + + filtered_messages.append(tool_call_message) + + return filtered_messages + @staticmethod def _ensure_tool_call_output_has_corresponding_tool_call( messages: List[Union[AllMessageValues, GenericChatCompletionMessage]], diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index c00c2a2f3b2..ac040d3d6ec 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -655,7 +655,6 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): follow_up_params.update( { "input": follow_up_input, - "previous_response_id": self.collected_response.id, # type: ignore[attr-defined] "stream": True, } ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 0b838f916e2..7abd4f90f2f 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -55,6 +55,7 @@ class BaseResponsesAPIStreamingIterator: self.responses_api_provider_config = responses_api_provider_config self.completed_response: Optional[ResponsesAPIStreamingResponse] = None self.start_time = getattr(logging_obj, "start_time", datetime.now()) + self._failure_handled = False # Track if failure handler has been called # track request context for hooks self.litellm_metadata = litellm_metadata @@ -169,7 +170,8 @@ class BaseResponsesAPIStreamingIterator: # If we can't parse the chunk, continue return None except Exception as e: - # Ensure failures trigger failure hooks + # Trigger failure hooks before re-raising + # This ensures failures are logged even when _process_chunk is called directly self._handle_failure(e) raise @@ -287,7 +289,13 @@ class BaseResponsesAPIStreamingIterator: def _handle_failure(self, exception: Exception): """ Trigger failure handlers before bubbling the exception. + Only calls handlers once even if called multiple times. """ + # Prevent double-calling failure handlers + if self._failure_handled: + return + self._failure_handled = True + traceback_exception = traceback.format_exc() try: run_async_function( @@ -383,11 +391,20 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def _handle_logging_completed_response(self): """Handle logging for completed responses in async context""" - # Create a deep copy for logging to avoid modifying the response object that will be returned to the user + # Create a copy for logging to avoid modifying the response object that will be returned to the user # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) # to chat completion format (prompt_tokens/completion_tokens) for internal logging - import copy - logging_response = copy.deepcopy(self.completed_response) + # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with + # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) + logging_response = self.completed_response + if self.completed_response is not None and hasattr(self.completed_response, 'model_dump'): + try: + logging_response = type(self.completed_response).model_validate( + self.completed_response.model_dump() + ) + except Exception: + # Fallback to original if serialization fails + pass asyncio.create_task( self.logging_obj.async_success_handler( @@ -469,11 +486,20 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def _handle_logging_completed_response(self): """Handle logging for completed responses in sync context""" - # Create a deep copy for logging to avoid modifying the response object that will be returned to the user + # Create a copy for logging to avoid modifying the response object that will be returned to the user # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) # to chat completion format (prompt_tokens/completion_tokens) for internal logging - import copy - logging_response = copy.deepcopy(self.completed_response) + # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with + # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) + logging_response = self.completed_response + if self.completed_response is not None and hasattr(self.completed_response, 'model_dump'): + try: + logging_response = type(self.completed_response).model_validate( + self.completed_response.model_dump() + ) + except Exception: + # Fallback to original if serialization fails + pass run_async_function( async_function=self.logging_obj.async_success_handler, diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index 187d8c97c05..c9d998f6a92 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -1,11 +1,10 @@ from typing import Dict, List, Literal, Optional, Union -from pydantic import Field +from pydantic import BaseModel, Field from typing_extensions import TypedDict -from litellm.types.utils import LiteLLMPydanticObjectBase - from litellm.proxy._types import LitellmUserRoles +from litellm.types.utils import LiteLLMPydanticObjectBase class LiteLLM_UpperboundKeyGenerateParams(LiteLLMPydanticObjectBase): @@ -28,6 +27,7 @@ class LiteLLM_UpperboundKeyGenerateParams(LiteLLMPydanticObjectBase): tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None + class MicrosoftGraphAPIUserGroupDirectoryObject(TypedDict, total=False): """Model for Microsoft Graph API directory object""" @@ -65,7 +65,7 @@ class AccessControl_UI_AccessMode(LiteLLMPydanticObjectBase): class RoleMappings(LiteLLMPydanticObjectBase): """ Configuration for mapping SSO groups to LiteLLM roles. - + The system will look at the group_claim field in the SSO token to determine which role to assign the user based on the roles mapping. """ @@ -78,11 +78,11 @@ class RoleMappings(LiteLLMPydanticObjectBase): ) default_role: Optional[LitellmUserRoles] = Field( default=None, - description="Default role to assign if user's groups don't match any role mappings. Must be a valid LitellmUserRoles value (e.g., 'proxy_admin', 'internal_user', 'proxy_admin_viewer')" + description="Default role to assign if user's groups don't match any role mappings. Must be a valid LitellmUserRoles value (e.g., 'proxy_admin', 'internal_user', 'proxy_admin_viewer')", ) roles: Dict[LitellmUserRoles, List[str]] = Field( default_factory=dict, - description="Mapping of LiteLLM role names to arrays of SSO group names. Example: {'proxy_admin': ['group-1', 'group-2'], 'proxy_admin_viewer': ['group-3']}" + description="Mapping of LiteLLM role names to arrays of SSO group names. Example: {'proxy_admin': ['group-1', 'group-2'], 'proxy_admin_viewer': ['group-3']}", ) @@ -185,3 +185,10 @@ class DefaultTeamSSOParams(LiteLLMPydanticObjectBase): default=None, description="Default rpm limit for new automatically created teams", ) + + +class InProductNudgeResponse(BaseModel): + is_claude_code_enabled: bool = Field( + default=False, + description="Whether the Claude Code nudge should be shown.", + ) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 470d598a25f..43f7bde5da3 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -7857,6 +7857,24 @@ "supports_tool_choice": true, "supports_vision": true }, + "dall-e-2": { + "input_cost_per_image": 0.02, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits", + "/v1/images/variations" + ] + }, + "dall-e-3": { + "input_cost_per_image": 0.04, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "deepseek-chat": { "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.8e-07, diff --git a/poetry.lock b/poetry.lock index b76f9c30f00..6a66d5fdf6a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,36 @@ -# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. + +[[package]] +name = "a2a-sdk" +version = "0.3.22" +description = "A2A Python SDK" +optional = true +python-versions = ">=3.10" +groups = ["main"] +markers = "python_version >= \"3.10\" and extra == \"extra-proxy\"" +files = [ + {file = "a2a_sdk-0.3.22-py3-none-any.whl", hash = "sha256:b98701135bb90b0ff85d35f31533b6b7a299bf810658c1c65f3814a6c15ea385"}, + {file = "a2a_sdk-0.3.22.tar.gz", hash = "sha256:77a5694bfc4f26679c11b70c7f1062522206d430b34bc1215cfbb1eba67b7e7d"}, +] + +[package.dependencies] +google-api-core = ">=1.26.0" +httpx = ">=0.28.1" +httpx-sse = ">=0.4.0" +protobuf = ">=5.29.5" +pydantic = ">=2.11.3" + +[package.extras] +all = ["cryptography (>=43.0.0)", "fastapi (>=0.115.2)", "grpcio (>=1.60)", "grpcio-reflection (>=1.7.0)", "grpcio-tools (>=1.60)", "opentelemetry-api (>=1.33.0)", "opentelemetry-sdk (>=1.33.0)", "pyjwt (>=2.0.0)", "sqlalchemy[aiomysql,asyncio] (>=2.0.0)", "sqlalchemy[aiosqlite,asyncio] (>=2.0.0)", "sqlalchemy[asyncio,postgresql-asyncpg] (>=2.0.0)", "sse-starlette", "starlette"] +encryption = ["cryptography (>=43.0.0)"] +grpc = ["grpcio (>=1.60)", "grpcio-reflection (>=1.7.0)", "grpcio-tools (>=1.60)"] +http-server = ["fastapi (>=0.115.2)", "sse-starlette", "starlette"] +mysql = ["sqlalchemy[aiomysql,asyncio] (>=2.0.0)"] +postgresql = ["sqlalchemy[asyncio,postgresql-asyncpg] (>=2.0.0)"] +signing = ["pyjwt (>=2.0.0)"] +sql = ["sqlalchemy[aiomysql,asyncio] (>=2.0.0)", "sqlalchemy[aiosqlite,asyncio] (>=2.0.0)", "sqlalchemy[asyncio,postgresql-asyncpg] (>=2.0.0)"] +sqlite = ["sqlalchemy[aiosqlite,asyncio] (>=2.0.0)"] +telemetry = ["opentelemetry-api (>=1.33.0)", "opentelemetry-sdk (>=1.33.0)"] [[package]] name = "aiofiles" @@ -1268,25 +1300,6 @@ dev = ["autoflake", "black", "build", "databricks-connect", "httpx", "ipython", notebook = ["ipython (>=8,<10)", "ipywidgets (>=8,<9)"] openai = ["httpx", "langchain-openai ; python_version > \"3.7\"", "openai"] -[[package]] -name = "deprecated" -version = "1.3.1" -description = "Python @deprecated decorator to deprecate old python classes, functions or methods." -optional = false -python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,>=2.7" -groups = ["main", "dev", "proxy-dev"] -files = [ - {file = "deprecated-1.3.1-py2.py3-none-any.whl", hash = "sha256:597bfef186b6f60181535a29fbe44865ce137a5079f295b479886c82729d5f3f"}, - {file = "deprecated-1.3.1.tar.gz", hash = "sha256:b1b50e0ff0c1fddaa5708a2c6b0a6588bb09b892825ab2b214ac9ea9d92a5223"}, -] -markers = {main = "python_version >= \"3.10\""} - -[package.dependencies] -wrapt = ">=1.10,<3" - -[package.extras] -dev = ["PyTest", "PyTest-Cov", "bump2version (<1)", "setuptools ; python_version >= \"3.12\"", "tox"] - [[package]] name = "diskcache" version = "5.6.3" @@ -2521,7 +2534,7 @@ description = "Consume Server-Sent Event (SSE) messages with HTTPX." optional = true python-versions = ">=3.9" groups = ["main"] -markers = "python_version >= \"3.10\" and extra == \"proxy\"" +markers = "python_version >= \"3.10\" and (extra == \"proxy\" or extra == \"extra-proxy\")" files = [ {file = "httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc"}, {file = "httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d"}, @@ -4036,143 +4049,153 @@ voice-helpers = ["numpy (>=2.0.2)", "sounddevice (>=0.5.1)"] [[package]] name = "opentelemetry-api" -version = "1.25.0" +version = "1.39.1" description = "OpenTelemetry Python API" optional = false -python-versions = ">=3.8" +python-versions = ">=3.9" groups = ["main", "dev", "proxy-dev"] files = [ - {file = "opentelemetry_api-1.25.0-py3-none-any.whl", hash = "sha256:757fa1aa020a0f8fa139f8959e53dec2051cc26b832e76fa839a6d76ecefd737"}, - {file = "opentelemetry_api-1.25.0.tar.gz", hash = "sha256:77c4985f62f2614e42ce77ee4c9da5fa5f0bc1e1821085e9a47533a9323ae869"}, + {file = "opentelemetry_api-1.39.1-py3-none-any.whl", hash = "sha256:2edd8463432a7f8443edce90972169b195e7d6a05500cd29e6d13898187c9950"}, + {file = "opentelemetry_api-1.39.1.tar.gz", hash = "sha256:fbde8c80e1b937a2c61f20347e91c0c18a1940cecf012d62e65a7caf08967c9c"}, ] markers = {main = "python_version >= \"3.10\""} [package.dependencies] -deprecated = ">=1.2.6" -importlib-metadata = ">=6.0,<=7.1" +importlib-metadata = ">=6.0,<8.8.0" +typing-extensions = ">=4.5.0" [[package]] name = "opentelemetry-exporter-otlp" -version = "1.25.0" +version = "1.39.1" description = "OpenTelemetry Collector Exporters" optional = false -python-versions = ">=3.8" +python-versions = ">=3.9" groups = ["dev", "proxy-dev"] files = [ - {file = "opentelemetry_exporter_otlp-1.25.0-py3-none-any.whl", hash = "sha256:d67a831757014a3bc3174e4cd629ae1493b7ba8d189e8a007003cacb9f1a6b60"}, - {file = "opentelemetry_exporter_otlp-1.25.0.tar.gz", hash = "sha256:ce03199c1680a845f82e12c0a6a8f61036048c07ec7a0bd943142aca8fa6ced0"}, + {file = "opentelemetry_exporter_otlp-1.39.1-py3-none-any.whl", hash = "sha256:68ae69775291f04f000eb4b698ff16ff685fdebe5cb52871bc4e87938a7b00fe"}, + {file = "opentelemetry_exporter_otlp-1.39.1.tar.gz", hash = "sha256:7cf7470e9fd0060c8a38a23e4f695ac686c06a48ad97f8d4867bc9b420180b9c"}, ] [package.dependencies] -opentelemetry-exporter-otlp-proto-grpc = "1.25.0" -opentelemetry-exporter-otlp-proto-http = "1.25.0" +opentelemetry-exporter-otlp-proto-grpc = "1.39.1" +opentelemetry-exporter-otlp-proto-http = "1.39.1" [[package]] name = "opentelemetry-exporter-otlp-proto-common" -version = "1.25.0" +version = "1.39.1" description = "OpenTelemetry Protobuf encoding" optional = false -python-versions = ">=3.8" +python-versions = ">=3.9" groups = ["dev", "proxy-dev"] files = [ - {file = "opentelemetry_exporter_otlp_proto_common-1.25.0-py3-none-any.whl", hash = "sha256:15637b7d580c2675f70246563363775b4e6de947871e01d0f4e3881d1848d693"}, - {file = "opentelemetry_exporter_otlp_proto_common-1.25.0.tar.gz", hash = "sha256:c93f4e30da4eee02bacd1e004eb82ce4da143a2f8e15b987a9f603e0a85407d3"}, + {file = "opentelemetry_exporter_otlp_proto_common-1.39.1-py3-none-any.whl", hash = "sha256:08f8a5862d64cc3435105686d0216c1365dc5701f86844a8cd56597d0c764fde"}, + {file = "opentelemetry_exporter_otlp_proto_common-1.39.1.tar.gz", hash = "sha256:763370d4737a59741c89a67b50f9e39271639ee4afc999dadfe768541c027464"}, ] [package.dependencies] -opentelemetry-proto = "1.25.0" +opentelemetry-proto = "1.39.1" [[package]] name = "opentelemetry-exporter-otlp-proto-grpc" -version = "1.25.0" +version = "1.39.1" description = "OpenTelemetry Collector Protobuf over gRPC Exporter" optional = false -python-versions = ">=3.8" +python-versions = ">=3.9" groups = ["dev", "proxy-dev"] files = [ - {file = "opentelemetry_exporter_otlp_proto_grpc-1.25.0-py3-none-any.whl", hash = "sha256:3131028f0c0a155a64c430ca600fd658e8e37043cb13209f0109db5c1a3e4eb4"}, - {file = "opentelemetry_exporter_otlp_proto_grpc-1.25.0.tar.gz", hash = "sha256:c0b1661415acec5af87625587efa1ccab68b873745ca0ee96b69bb1042087eac"}, + {file = "opentelemetry_exporter_otlp_proto_grpc-1.39.1-py3-none-any.whl", hash = "sha256:fa1c136a05c7e9b4c09f739469cbdb927ea20b34088ab1d959a849b5cc589c18"}, + {file = "opentelemetry_exporter_otlp_proto_grpc-1.39.1.tar.gz", hash = "sha256:772eb1c9287485d625e4dbe9c879898e5253fea111d9181140f51291b5fec3ad"}, ] [package.dependencies] -deprecated = ">=1.2.6" -googleapis-common-protos = ">=1.52,<2.0" -grpcio = ">=1.0.0,<2.0.0" +googleapis-common-protos = ">=1.57,<2.0" +grpcio = [ + {version = ">=1.63.2,<2.0.0", markers = "python_version < \"3.13\""}, + {version = ">=1.66.2,<2.0.0", markers = "python_version >= \"3.13\""}, +] opentelemetry-api = ">=1.15,<2.0" -opentelemetry-exporter-otlp-proto-common = "1.25.0" -opentelemetry-proto = "1.25.0" -opentelemetry-sdk = ">=1.25.0,<1.26.0" +opentelemetry-exporter-otlp-proto-common = "1.39.1" +opentelemetry-proto = "1.39.1" +opentelemetry-sdk = ">=1.39.1,<1.40.0" +typing-extensions = ">=4.6.0" + +[package.extras] +gcp-auth = ["opentelemetry-exporter-credential-provider-gcp (>=0.59b0)"] [[package]] name = "opentelemetry-exporter-otlp-proto-http" -version = "1.25.0" +version = "1.39.1" description = "OpenTelemetry Collector Protobuf over HTTP Exporter" optional = false -python-versions = ">=3.8" +python-versions = ">=3.9" groups = ["dev", "proxy-dev"] files = [ - {file = "opentelemetry_exporter_otlp_proto_http-1.25.0-py3-none-any.whl", hash = "sha256:2eca686ee11b27acd28198b3ea5e5863a53d1266b91cda47c839d95d5e0541a6"}, - {file = "opentelemetry_exporter_otlp_proto_http-1.25.0.tar.gz", hash = "sha256:9f8723859e37c75183ea7afa73a3542f01d0fd274a5b97487ea24cb683d7d684"}, + {file = "opentelemetry_exporter_otlp_proto_http-1.39.1-py3-none-any.whl", hash = "sha256:d9f5207183dd752a412c4cd564ca8875ececba13be6e9c6c370ffb752fd59985"}, + {file = "opentelemetry_exporter_otlp_proto_http-1.39.1.tar.gz", hash = "sha256:31bdab9745c709ce90a49a0624c2bd445d31a28ba34275951a6a362d16a0b9cb"}, ] [package.dependencies] -deprecated = ">=1.2.6" googleapis-common-protos = ">=1.52,<2.0" opentelemetry-api = ">=1.15,<2.0" -opentelemetry-exporter-otlp-proto-common = "1.25.0" -opentelemetry-proto = "1.25.0" -opentelemetry-sdk = ">=1.25.0,<1.26.0" +opentelemetry-exporter-otlp-proto-common = "1.39.1" +opentelemetry-proto = "1.39.1" +opentelemetry-sdk = ">=1.39.1,<1.40.0" requests = ">=2.7,<3.0" +typing-extensions = ">=4.5.0" + +[package.extras] +gcp-auth = ["opentelemetry-exporter-credential-provider-gcp (>=0.59b0)"] [[package]] name = "opentelemetry-proto" -version = "1.25.0" +version = "1.39.1" description = "OpenTelemetry Python Proto" optional = false -python-versions = ">=3.8" +python-versions = ">=3.9" groups = ["main", "dev", "proxy-dev"] files = [ - {file = "opentelemetry_proto-1.25.0-py3-none-any.whl", hash = "sha256:f07e3341c78d835d9b86665903b199893befa5e98866f63d22b00d0b7ca4972f"}, - {file = "opentelemetry_proto-1.25.0.tar.gz", hash = "sha256:35b6ef9dc4a9f7853ecc5006738ad40443701e52c26099e197895cbda8b815a3"}, + {file = "opentelemetry_proto-1.39.1-py3-none-any.whl", hash = "sha256:22cdc78efd3b3765d09e68bfbd010d4fc254c9818afd0b6b423387d9dee46007"}, + {file = "opentelemetry_proto-1.39.1.tar.gz", hash = "sha256:6c8e05144fc0d3ed4d22c2289c6b126e03bcd0e6a7da0f16cedd2e1c2772e2c8"}, ] markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""} [package.dependencies] -protobuf = ">=3.19,<5.0" +protobuf = ">=5.0,<7.0" [[package]] name = "opentelemetry-sdk" -version = "1.25.0" +version = "1.39.1" description = "OpenTelemetry Python SDK" optional = false -python-versions = ">=3.8" +python-versions = ">=3.9" groups = ["main", "dev", "proxy-dev"] files = [ - {file = "opentelemetry_sdk-1.25.0-py3-none-any.whl", hash = "sha256:d97ff7ec4b351692e9d5a15af570c693b8715ad78b8aafbec5c7100fe966b4c9"}, - {file = "opentelemetry_sdk-1.25.0.tar.gz", hash = "sha256:ce7fc319c57707ef5bf8b74fb9f8ebdb8bfafbe11898410e0d2a761d08a98ec7"}, + {file = "opentelemetry_sdk-1.39.1-py3-none-any.whl", hash = "sha256:4d5482c478513ecb0a5d938dcc61394e647066e0cc2676bee9f3af3f3f45f01c"}, + {file = "opentelemetry_sdk-1.39.1.tar.gz", hash = "sha256:cf4d4563caf7bff906c9f7967e2be22d0d6b349b908be0d90fb21c8e9c995cc6"}, ] markers = {main = "python_version >= \"3.10\""} [package.dependencies] -opentelemetry-api = "1.25.0" -opentelemetry-semantic-conventions = "0.46b0" -typing-extensions = ">=3.7.4" +opentelemetry-api = "1.39.1" +opentelemetry-semantic-conventions = "0.60b1" +typing-extensions = ">=4.5.0" [[package]] name = "opentelemetry-semantic-conventions" -version = "0.46b0" +version = "0.60b1" description = "OpenTelemetry Semantic Conventions" optional = false -python-versions = ">=3.8" +python-versions = ">=3.9" groups = ["main", "dev", "proxy-dev"] files = [ - {file = "opentelemetry_semantic_conventions-0.46b0-py3-none-any.whl", hash = "sha256:6daef4ef9fa51d51855d9f8e0ccd3a1bd59e0e545abe99ac6203804e36ab3e07"}, - {file = "opentelemetry_semantic_conventions-0.46b0.tar.gz", hash = "sha256:fbc982ecbb6a6e90869b15c1673be90bd18c8a56ff1cffc0864e38e2edffaefa"}, + {file = "opentelemetry_semantic_conventions-0.60b1-py3-none-any.whl", hash = "sha256:9fa8c8b0c110da289809292b0591220d3a7b53c1526a23021e977d68597893fb"}, + {file = "opentelemetry_semantic_conventions-0.60b1.tar.gz", hash = "sha256:87c228b5a0669b748c76d76df6c364c369c28f1c465e50f661e39737e84bc953"}, ] markers = {main = "python_version >= \"3.10\""} [package.dependencies] -opentelemetry-api = "1.25.0" +opentelemetry-api = "1.39.1" +typing-extensions = ">=4.5.0" [[package]] name = "orjson" @@ -4828,23 +4851,23 @@ testing = ["google-api-core (>=1.31.5)"] [[package]] name = "protobuf" -version = "4.25.8" +version = "5.29.5" description = "" optional = false python-versions = ">=3.8" groups = ["main", "dev", "proxy-dev"] files = [ - {file = "protobuf-4.25.8-cp310-abi3-win32.whl", hash = "sha256:504435d831565f7cfac9f0714440028907f1975e4bed228e58e72ecfff58a1e0"}, - {file = "protobuf-4.25.8-cp310-abi3-win_amd64.whl", hash = "sha256:bd551eb1fe1d7e92c1af1d75bdfa572eff1ab0e5bf1736716814cdccdb2360f9"}, - {file = "protobuf-4.25.8-cp37-abi3-macosx_10_9_universal2.whl", hash = "sha256:ca809b42f4444f144f2115c4c1a747b9a404d590f18f37e9402422033e464e0f"}, - {file = "protobuf-4.25.8-cp37-abi3-manylinux2014_aarch64.whl", hash = "sha256:9ad7ef62d92baf5a8654fbb88dac7fa5594cfa70fd3440488a5ca3bfc6d795a7"}, - {file = "protobuf-4.25.8-cp37-abi3-manylinux2014_x86_64.whl", hash = "sha256:83e6e54e93d2b696a92cad6e6efc924f3850f82b52e1563778dfab8b355101b0"}, - {file = "protobuf-4.25.8-cp38-cp38-win32.whl", hash = "sha256:27d498ffd1f21fb81d987a041c32d07857d1d107909f5134ba3350e1ce80a4af"}, - {file = "protobuf-4.25.8-cp38-cp38-win_amd64.whl", hash = "sha256:d552c53d0415449c8d17ced5c341caba0d89dbf433698e1436c8fa0aae7808a3"}, - {file = "protobuf-4.25.8-cp39-cp39-win32.whl", hash = "sha256:077ff8badf2acf8bc474406706ad890466274191a48d0abd3bd6987107c9cde5"}, - {file = "protobuf-4.25.8-cp39-cp39-win_amd64.whl", hash = "sha256:f4510b93a3bec6eba8fd8f1093e9d7fb0d4a24d1a81377c10c0e5bbfe9e4ed24"}, - {file = "protobuf-4.25.8-py3-none-any.whl", hash = "sha256:15a0af558aa3b13efef102ae6e4f3efac06f1eea11afb3a57db2901447d9fb59"}, - {file = "protobuf-4.25.8.tar.gz", hash = "sha256:6135cf8affe1fc6f76cced2641e4ea8d3e59518d1f24ae41ba97bcad82d397cd"}, + {file = "protobuf-5.29.5-cp310-abi3-win32.whl", hash = "sha256:3f1c6468a2cfd102ff4703976138844f78ebd1fb45f49011afc5139e9e283079"}, + {file = "protobuf-5.29.5-cp310-abi3-win_amd64.whl", hash = "sha256:3f76e3a3675b4a4d867b52e4a5f5b78a2ef9565549d4037e06cf7b0942b1d3fc"}, + {file = "protobuf-5.29.5-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:e38c5add5a311f2a6eb0340716ef9b039c1dfa428b28f25a7838ac329204a671"}, + {file = "protobuf-5.29.5-cp38-abi3-manylinux2014_aarch64.whl", hash = "sha256:fa18533a299d7ab6c55a238bf8629311439995f2e7eca5caaff08663606e9015"}, + {file = "protobuf-5.29.5-cp38-abi3-manylinux2014_x86_64.whl", hash = "sha256:63848923da3325e1bf7e9003d680ce6e14b07e55d0473253a690c3a8b8fd6e61"}, + {file = "protobuf-5.29.5-cp38-cp38-win32.whl", hash = "sha256:ef91363ad4faba7b25d844ef1ada59ff1604184c0bcd8b39b8a6bef15e1af238"}, + {file = "protobuf-5.29.5-cp38-cp38-win_amd64.whl", hash = "sha256:7318608d56b6402d2ea7704ff1e1e4597bee46d760e7e4dd42a3d45e24b87f2e"}, + {file = "protobuf-5.29.5-cp39-cp39-win32.whl", hash = "sha256:6f642dc9a61782fa72b90878af134c5afe1917c89a568cd3476d758d3c3a0736"}, + {file = "protobuf-5.29.5-cp39-cp39-win_amd64.whl", hash = "sha256:470f3af547ef17847a28e1f47200a1cbf0ba3ff57b7de50d22776607cd2ea353"}, + {file = "protobuf-5.29.5-py3-none-any.whl", hash = "sha256:6cf42630262c59b2d8de33954443d94b746c952b01434fc58a417fdbd2e84bd5"}, + {file = "protobuf-5.29.5.tar.gz", hash = "sha256:bc1463bafd4b0929216c35f437a8e28731a2b7fe3d98bb77a600efced5a15c84"}, ] markers = {main = "python_version >= \"3.10\" and (extra == \"mlflow\" or extra == \"extra-proxy\") or extra == \"extra-proxy\""} @@ -7687,7 +7710,7 @@ version = "1.17.3" description = "Module for decorators, wrappers and monkey patching." optional = false python-versions = ">=3.8" -groups = ["main", "dev", "proxy-dev"] +groups = ["dev"] files = [ {file = "wrapt-1.17.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:88bbae4d40d5a46142e70d58bf664a89b6b4befaea7b2ecc14e03cedb8e06c04"}, {file = "wrapt-1.17.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e6b13af258d6a9ad602d57d889f83b9d5543acd471eee12eb51f5b01f8eb1bc2"}, @@ -7771,7 +7794,6 @@ files = [ {file = "wrapt-1.17.3-py3-none-any.whl", hash = "sha256:7171ae35d2c33d326ac19dd8facb1e82e5fd04ef8c6c0e394d7af55a55051c22"}, {file = "wrapt-1.17.3.tar.gz", hash = "sha256:f66eb08feaa410fe4eebd17f2a2c8e2e46d3476e9f8c783daa8e09e0faa666d0"}, ] -markers = {main = "python_version >= \"3.10\""} [[package]] name = "wsproto" @@ -7972,7 +7994,7 @@ type = ["pytest-mypy"] [extras] caching = ["diskcache"] -extra-proxy = ["azure-identity", "azure-keyvault-secrets", "google-cloud-iam", "google-cloud-kms", "prisma", "redisvl", "resend"] +extra-proxy = ["a2a-sdk", "azure-identity", "azure-keyvault-secrets", "google-cloud-iam", "google-cloud-kms", "prisma", "redisvl", "resend"] mlflow = ["mlflow"] proxy = ["PyJWT", "apscheduler", "azure-identity", "azure-storage-blob", "backoff", "boto3", "cryptography", "fastapi", "fastapi-sso", "gunicorn", "litellm-enterprise", "litellm-proxy-extras", "mcp", "orjson", "polars", "pynacl", "python-multipart", "pyyaml", "rich", "rq", "soundfile", "uvicorn", "uvloop", "websockets"] semantic-router = ["semantic-router"] diff --git a/pyproject.toml b/pyproject.toml index d6242b27786..4b22b2fffed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,6 +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"} +a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"} litellm-proxy-extras = {version = "0.4.23", optional = true} rich = {version = "13.7.1", optional = true} litellm-enterprise = {version = "0.1.27", optional = true} @@ -111,7 +112,8 @@ extra_proxy = [ "google-cloud-kms", "google-cloud-iam", "resend", - "redisvl" + "redisvl", + "a2a-sdk" ] utils = [ @@ -147,9 +149,9 @@ types-requests = "*" types-setuptools = "*" types-redis = "*" types-PyYAML = "*" -opentelemetry-api = "1.25.0" -opentelemetry-sdk = "1.25.0" -opentelemetry-exporter-otlp = "1.25.0" +opentelemetry-api = "^1.28.0" +opentelemetry-sdk = "^1.28.0" +opentelemetry-exporter-otlp = "^1.28.0" langfuse = "^2.45.0" fastapi-offline = "^1.7.3" @@ -157,9 +159,9 @@ fastapi-offline = "^1.7.3" prisma = "0.11.0" hypercorn = "^0.15.0" prometheus-client = "0.20.0" -opentelemetry-api = "1.25.0" -opentelemetry-sdk = "1.25.0" -opentelemetry-exporter-otlp = "1.25.0" +opentelemetry-api = "^1.28.0" +opentelemetry-sdk = "^1.28.0" +opentelemetry-exporter-otlp = "^1.28.0" azure-identity = {version = "^1.15.0", python = ">=3.9"} [build-system] diff --git a/requirements.txt b/requirements.txt index a49a94ca274..68b85c3935d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,12 +16,12 @@ prisma==0.11.0 # for db nodejs-wheel-binaries==24.12.0 ## required by prisma for migrations, prevents runtime download (updated from nodejs-bin for security fixes) mangum==0.17.0 # for aws lambda functions pynacl==1.6.2 # for encrypting keys -google-cloud-aiplatform==1.47.0 # for vertex ai calls +google-cloud-aiplatform==1.133.0 # for vertex ai calls google-cloud-iam==2.19.1 # for GCP IAM Redis authentication -google-genai==1.22.0 +google-genai==1.37.0 anthropic[vertex]==0.54.0 mcp==1.25.0 ; python_version >= "3.10" # for MCP server -google-generativeai==0.5.0 # for vertex ai calls +# google-generativeai removed - deprecated, replaced by google-genai (line 21) async_generator==1.10.0 # for async ollama calls langfuse==2.59.7 # for langfuse self-hosted logging prometheus_client==0.20.0 # for /metrics endpoint on proxy @@ -38,9 +38,10 @@ azure-ai-contentsafety==1.0.0 # for azure content safety azure-identity==1.16.1 ; python_version >= "3.9" # for azure content safety azure-keyvault==4.2.0 # for azure KMS integration azure-storage-file-datalake==12.20.0 # for azure buck storage logging -opentelemetry-api==1.25.0 -opentelemetry-sdk==1.25.0 -opentelemetry-exporter-otlp==1.25.0 +opentelemetry-api==1.28.0 +opentelemetry-sdk==1.28.0 +opentelemetry-exporter-otlp==1.28.0 +a2a-sdk>=0.3.22 ; python_version >= "3.10" # grpcio: 1.68.0-1.68.1 has reconnect bug (#38290), 1.75+ has Python 3.14 wheels + fix grpcio>=1.62.3,!=1.68.*,!=1.69.*,!=1.70.*,!=1.71.0,!=1.71.1,!=1.72.0,!=1.72.1,!=1.73.0; python_version < "3.14" grpcio>=1.75.0; python_version >= "3.14" diff --git a/scripts/health_check/health_check_client.py b/scripts/health_check/health_check_client.py new file mode 100644 index 00000000000..337c754a75c --- /dev/null +++ b/scripts/health_check/health_check_client.py @@ -0,0 +1,406 @@ +#!/usr/bin/env python3 +""" +LiteLLM Health Check Client + +A sentinel health check tool that tests all configured models on a LiteLLM proxy. +Similar to HRT's health check system, this script: +- Can read models from YAML config file (like HRT) or fetch from proxy API +- Sends a simple test request to each model concurrently +- Reports health status for each model +- Supports both chat/completion and embedding models +""" + +import asyncio +import json +import os +import sys +import time +from typing import Dict, List, Optional, Tuple + +import httpx +import yaml + + +class LiteLLMHealthCheckClient: + """Client for health checking LiteLLM proxy models.""" + + def __init__( + self, + base_url: str, + api_key: str, + timeout: int = 120, # Match Go implementation's 120s timeout + completion_prompt: str = "Say this is a test", # Match Go implementation + embedding_text: str = "This is a test for vectorization.", # Match Go implementation + ): + """ + Initialize the health check client. + + Args: + base_url: Base URL of the LiteLLM proxy (e.g., https://litellm.example.com) + api_key: API key for authentication + timeout: Request timeout in seconds (default: 120, matching Go implementation) + completion_prompt: Test prompt for chat/completion models + embedding_text: Test text for embedding models + """ + self.base_url = base_url.rstrip("/") + self.api_key = api_key + self.timeout = timeout + self.completion_prompt = completion_prompt + self.embedding_text = embedding_text + self.headers = { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + } + + def load_models_from_yaml(self, yaml_path: str) -> List[Dict]: + """ + Load models from a YAML config file (similar to Go implementation). + + Args: + yaml_path: Path to the YAML config file + + Returns: + List of model dictionaries with 'id' and 'mode' keys + """ + try: + with open(yaml_path, "r") as f: + config = yaml.safe_load(f) + + model_list = config.get("model_list", []) + models = [] + + for entry in model_list: + model_name = entry.get("model_name", "") + litellm_params = entry.get("litellm_params", {}) + model_info = litellm_params.get("model_info", {}) + mode = model_info.get("mode", "") + + # Use model_name as the ID (this is what gets sent to the API) + models.append( + { + "id": model_name, + "mode": mode.lower() if mode else "", + "provider": model_info.get("provider", ""), + } + ) + + return models + except Exception as e: + print(f"Error loading models from YAML file {yaml_path}: {e}", file=sys.stderr) + return [] + + async def fetch_models(self, client: httpx.AsyncClient) -> List[Dict]: + """ + Fetch all available models from the proxy API. + + Returns: + List of model dictionaries with 'id' and 'mode' keys + """ + try: + # Try /v1/models first (OpenAI-compatible endpoint) + response = await client.get( + f"{self.base_url}/v1/models", + headers=self.headers, + timeout=self.timeout, + ) + response.raise_for_status() + data = response.json() + models_data = data.get("data", []) + models = [] + for m in models_data: + models.append({"id": m["id"], "mode": "", "provider": ""}) + return models + except Exception as e: + print(f"Error fetching models from /v1/models: {e}", file=sys.stderr) + # Fallback to /model/info endpoint which has more details + try: + response = await client.get( + f"{self.base_url}/model/info", + headers=self.headers, + timeout=self.timeout, + ) + response.raise_for_status() + data = response.json() + if isinstance(data, dict) and "data" in data: + models_data = data["data"] + elif isinstance(data, list): + models_data = data + else: + models_data = [] + + models = [] + for m in models_data: + model_info = m.get("model_info", {}) + mode = model_info.get("mode", "") + models.append( + { + "id": m.get("model_name", m.get("id", "unknown")), + "mode": mode.lower() if mode else "", + "provider": model_info.get("provider", ""), + } + ) + return models + except Exception as e2: + print(f"Error fetching models from /model/info: {e2}", file=sys.stderr) + return [] + + async def check_model_health( + self, client: httpx.AsyncClient, model: Dict + ) -> Tuple[str, Dict]: + """ + Check health of a single model by sending a test request. + + Args: + client: HTTP client + model: Model dictionary with 'id' and 'mode' keys + + Returns: + Tuple of (model_id, result_dict) + """ + model_id = model["id"] + mode = model.get("mode", "") + + start_time = time.time() + result = { + "model": model_id, + "healthy": False, + "error": None, + "response_time_ms": None, + "mode": mode, + } + + try: + # Determine if this is an embedding model + # Check mode first (from config), then fall back to name-based detection + is_embedding = ( + mode == "embedding" + or any( + keyword in model_id.lower() + for keyword in ["embedding", "embed", "text-embedding"] + ) + ) + + if is_embedding: + # Test embedding endpoint (matching Go implementation) + embedding_response = await client.post( + f"{self.base_url}/v1/embeddings", + headers=self.headers, + json={ + "model": model_id, + "input": self.embedding_text, + }, + timeout=self.timeout, + ) + embedding_response.raise_for_status() + embedding_data = embedding_response.json() + dimensions = 0 + if "data" in embedding_data and len(embedding_data["data"]) > 0: + dimensions = len(embedding_data["data"][0].get("embedding", [])) + + result["healthy"] = True + result["mode"] = "embedding" + result["dimensions"] = dimensions + else: + # Test chat completion endpoint (matching Go implementation) + completion_response = await client.post( + f"{self.base_url}/v1/chat/completions", + headers=self.headers, + json={ + "model": model_id, + "messages": [ + {"role": "user", "content": self.completion_prompt} + ], + "max_tokens": 10, # Minimal tokens for health check + }, + timeout=self.timeout, + ) + completion_response.raise_for_status() + completion_data = completion_response.json() + response_text = "" + if "choices" in completion_data and len(completion_data["choices"]) > 0: + response_text = ( + completion_data["choices"][0] + .get("message", {}) + .get("content", "") + ) + + result["healthy"] = True + result["mode"] = "chat" + result["response_text"] = response_text[:100] # Truncate for display + + elapsed_ms = (time.time() - start_time) * 1000 + result["response_time_ms"] = round(elapsed_ms, 2) + + except httpx.HTTPStatusError as e: + result["error"] = f"HTTP {e.response.status_code}: {e.response.text[:200]}" + except httpx.TimeoutException: + result["error"] = f"Request timeout after {self.timeout}s" + except Exception as e: + result["error"] = str(e)[:200] + + return model_id, result + + async def run_health_checks( + self, + models: Optional[List[Dict]] = None, + models_only: Optional[List[str]] = None, + ) -> Dict[str, Dict]: + """ + Run health checks on all models concurrently. + + Args: + models: Optional list of models to check. If None, fetches from proxy. + models_only: Optional list of model IDs to check. If set, only these + models are health-checked (must exist in the models list). + + Returns: + Dictionary mapping model_id to health check result + """ + async with httpx.AsyncClient() as client: + if models is None: + models = await self.fetch_models(client) + + if not models: + print("No models found to health check", file=sys.stderr) + return {} + + if models_only: + allowlist = {m.strip() for m in models_only if m and m.strip()} + models = [m for m in models if m.get("id") in allowlist] + print( + f"Filtering to only check {len(models)} models: {', '.join(sorted(allowlist))}", + file=sys.stderr, + ) + if not models: + print( + "No models matched LITELLM_MODELS_ONLY filter", + file=sys.stderr, + ) + return {} + + print(f"Running health checks on {len(models)} models...", file=sys.stderr) + + # Run all health checks concurrently + tasks = [self.check_model_health(client, model) for model in models] + results_list = await asyncio.gather(*tasks, return_exceptions=True) + + # Convert to dictionary format + results = {} + for result in results_list: + if isinstance(result, Exception): + print( + f"Exception in health check task: {result}", file=sys.stderr + ) + continue + # Type narrowing: after checking it's not an Exception, it's a Tuple + if isinstance(result, tuple) and len(result) == 2: + model_id, result_dict = result + results[model_id] = result_dict + + return results + + def print_results(self, results: Dict[str, Dict], json_output: bool = False): + """ + Print health check results. + + Args: + results: Dictionary of health check results + json_output: If True, output as JSON + """ + if json_output: + print(json.dumps(results, indent=2)) + return + + healthy_count = sum(1 for r in results.values() if r.get("healthy")) + unhealthy_count = len(results) - healthy_count + + # Print detailed results for each model (matching Go output format) + print(f"\n{'='*60}", file=sys.stderr) + print(f"Starting health check queries\n", file=sys.stderr) + + for model_id, result in results.items(): + if result.get("healthy"): + if result.get("mode") == "embedding": + dimensions = result.get("dimensions", 0) + print( + f"---- {model_id} ----\n✅ Success. " + f"Generated embedding vector with {dimensions} dimensions.\n\n", + file=sys.stderr, + ) + else: + response_text = result.get("response_text", "") + print( + f"---- {model_id} ----\n✅ Success. " + f"Response:\n{response_text}\n\n", + file=sys.stderr, + ) + else: + error = result.get("error", "Unknown error") + print(f"---- {model_id} ----\n❌ ERROR: {error}\n\n", file=sys.stderr) + + print(f"{'='*60}", file=sys.stderr) + print(f"Health Check Summary", file=sys.stderr) + print(f"{'='*60}", file=sys.stderr) + print(f"Total models: {len(results)}", file=sys.stderr) + print(f"Healthy: {healthy_count}", file=sys.stderr) + print(f"Unhealthy: {unhealthy_count}", file=sys.stderr) + print(f"{'='*60}\n", file=sys.stderr) + + # Exit with non-zero code if any models are unhealthy + if unhealthy_count > 0: + sys.exit(1) + else: + sys.exit(0) + + +async def main(): + """Main entry point.""" + base_url = os.environ.get("LITELLM_BASE_URL", "http://localhost:4000") + api_key = os.environ.get("LITELLM_API_KEY", "sk-1234") + yaml_path = os.environ.get("LITELLM_MODELS_YAML") + + if not base_url: + print("Error: LITELLM_BASE_URL environment variable not set", file=sys.stderr) + sys.exit(1) + + if not api_key: + print("Error: LITELLM_API_KEY environment variable not set", file=sys.stderr) + sys.exit(1) + + timeout = int(os.environ.get("LITELLM_TIMEOUT", "120")) # Match Go's 120s default + completion_prompt = os.environ.get( + "LITELLM_COMPLETION_PROMPT", "Say this is a test" + ) + embedding_text = os.environ.get( + "LITELLM_EMBEDDING_TEXT", "This is a test for vectorization." + ) + json_output = os.environ.get("LITELLM_JSON_OUTPUT", "").lower() == "true" + # Optional: only health-check these model IDs (comma-separated). E.g.: + # LITELLM_MODELS_ONLY=claude-3.7-sonnet,claude-3.5-sonnet,claude-4.5-haiku + models_only_raw = os.environ.get("LITELLM_MODELS_ONLY", "") + models_only = [m.strip() for m in models_only_raw.split(",") if m.strip()] or None + + client = LiteLLMHealthCheckClient( + base_url=base_url, + api_key=api_key, + timeout=timeout, + completion_prompt=completion_prompt, + embedding_text=embedding_text, + ) + + # Load models from YAML if provided, otherwise fetch from API + models = None + if yaml_path: + models = client.load_models_from_yaml(yaml_path) + if models: + print( + f"Successfully loaded {len(models)} models from {yaml_path}", + file=sys.stderr, + ) + + results = await client.run_health_checks(models=models, models_only=models_only) + client.print_results(results, json_output=json_output) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/scripts/health_check/health_check_client_README.md b/scripts/health_check/health_check_client_README.md new file mode 100644 index 00000000000..e3132499e04 --- /dev/null +++ b/scripts/health_check/health_check_client_README.md @@ -0,0 +1,246 @@ +# LiteLLM Health Check Client + +A health check tool for testing all configured models on a LiteLLM proxy. Tests each model with completion/embedding requests and reports health status, errors, and response times. + +## Features + +- **YAML Config Support**: Reads models from YAML config file OR fetches from proxy API +- **Smart Mode Detection**: Detects embedding vs chat models from config or model name +- **Concurrent Testing**: Tests all models concurrently using asyncio +- **Containerized**: Docker image for easy deployment +- **Parallel Execution**: Supports parallel execution for stress testing +- **Configurable**: Customizable timeouts (default 120s) and test prompts + +## Quick Start + +### As a Python Script + +**Option 1: Fetch models from proxy API** +```bash +export LITELLM_BASE_URL="https://litellm.example.com" +export LITELLM_API_KEY="your-api-key" +python scripts/health_check/health_check_client.py +``` + +**Option 2: Use YAML config file** +```bash +export LITELLM_BASE_URL="https://litellm.example.com" +export LITELLM_API_KEY="your-api-key" +export LITELLM_MODELS_YAML="/path/to/config.yaml" +python scripts/health_check/health_check_client.py +``` + +### As a Docker Container + +1. Build the Docker image: + +```bash +docker build -f docker/Dockerfile.health_check -t litellm/litellm-health-check:latest . +``` + +2. Run a single health check: + +```bash +docker run --rm \ + -e LITELLM_BASE_URL="https://litellm.example.com" \ + -e LITELLM_API_KEY="your-api-key" \ + litellm/litellm-health-check:latest +``` + +### Parallel Execution (Stress Testing) + +Run multiple health check containers in parallel: + +**PowerShell:** +```powershell +$env:LITELLM_BASE_URL="https://litellm.example.com" +$env:LITELLM_API_KEY="your-api-key" +.\scripts\health_check\run_parallel_health_checks.ps1 16 +``` + +**Bash/Shell:** +```bash +export LITELLM_BASE_URL="https://litellm.example.com" +export LITELLM_API_KEY="your-api-key" +./scripts/health_check/run_parallel_health_checks.sh 16 +``` + + +## Configuration + +### Environment Variables + +- `LITELLM_BASE_URL` (required): Base URL of the LiteLLM proxy + - Example: `https://litellm.example.com` +- `LITELLM_API_KEY` (required): API key for authentication +- `LITELLM_MODELS_YAML` (optional): Path to YAML config file with model_list + - If provided, reads models from YAML instead of fetching from API + - Example: `/path/to/config.yaml` +- `LITELLM_TIMEOUT` (optional): Request timeout in seconds (default: 120) +- `LITELLM_COMPLETION_PROMPT` (optional): Test prompt for chat/completion models (default: "Say this is a test") +- `LITELLM_EMBEDDING_TEXT` (optional): Test text for embedding models (default: "This is a test for vectorization.") +- `LITELLM_JSON_OUTPUT` (optional): Output results as JSON (default: false) + +## Output + +### Standard Output (Human-Readable) + +Example output format: + +``` +============================================================ +Starting health check queries + +---- gpt-4o ---- +✅ Success. Response: +This is a test + +---- text-embedding-3-small ---- +✅ Success. Generated embedding vector with 1536 dimensions. + +---- gpt-5-codex ---- +❌ ERROR: HTTP 503: Service unavailable + +============================================================ +Health Check Summary +============================================================ +Total models: 47 +Healthy: 45 +Unhealthy: 2 +============================================================ +``` + +Exit code: `0` if all models are healthy, `1` if any models are unhealthy. + +### JSON Output + +When `LITELLM_JSON_OUTPUT=true`, outputs JSON: + +```json +{ + "gpt-4o": { + "model": "gpt-4o", + "healthy": true, + "error": null, + "response_time_ms": 245.67, + "mode": "chat", + "response_text": "This is a test" + }, + "text-embedding-3-small": { + "model": "text-embedding-3-small", + "healthy": true, + "error": null, + "response_time_ms": 123.45, + "mode": "embedding", + "dimensions": 1536 + } +} +``` + +## How It Works + +1. **Model Discovery**: + - If `LITELLM_MODELS_YAML` is set: Reads models from YAML config file + - Otherwise: Queries `/v1/models` (OpenAI-compatible) or `/model/info` to get all configured models +2. **Mode Detection**: + - Checks `mode` field from YAML config, or falls back to model name patterns (embedding, embed, text-embedding) +3. **Concurrent Testing**: + - Chat models: `POST /v1/chat/completions` with configurable prompt (default: "Say this is a test") + - Embedding models: `POST /v1/embeddings` with configurable text (default: "This is a test for vectorization.") +4. **Reporting**: Health status, errors, response times, and response details are reported + +## Use Cases + +### 1. Regular Health Monitoring + +Run as a cron job or scheduled task: + +```bash +# Cron job: Run every 5 minutes +*/5 * * * * /path/to/health_check.sh +``` + +### 2. Load/Stress Testing + +Run multiple health checks in parallel: + +**PowerShell:** +```powershell +.\scripts\health_check\run_parallel_health_checks.ps1 16 +``` + +### 3. CI/CD Integration + +Add to your deployment pipeline: + +```yaml +# GitHub Actions example +- name: Health Check + run: | + docker run --rm \ + -e LITELLM_BASE_URL="${{ secrets.LITELLM_BASE_URL }}" \ + -e LITELLM_API_KEY="${{ secrets.LITELLM_API_KEY }}" \ + litellm/litellm-health-check:latest +``` + +### 4. Kubernetes Deployment + +Deploy as a CronJob: + +```yaml +apiVersion: batch/v1 +kind: CronJob +metadata: + name: litellm-health-check +spec: + schedule: "*/5 * * * *" # Every 5 minutes + jobTemplate: + spec: + template: + spec: + containers: + - name: health-check + image: litellm/litellm-health-check:latest + env: + - name: LITELLM_BASE_URL + value: "https://litellm.example.com" + - name: LITELLM_API_KEY + valueFrom: + secretKeyRef: + name: litellm-secrets + key: api-key + restartPolicy: OnFailure +``` + +## Troubleshooting + +### No Models Found + +- Verify `LITELLM_BASE_URL` is correct +- Check that the API key has permissions to list models +- Ensure the proxy is running and accessible +- If using YAML, verify `LITELLM_MODELS_YAML` path is correct + +### Timeout Errors + +- Increase `LITELLM_TIMEOUT` for slower models (default is 120s) +- Check network connectivity to the proxy +- Verify proxy isn't overloaded + +### Authentication Errors + +- Verify `LITELLM_API_KEY` is correct +- Check API key has not expired +- Ensure the key has necessary permissions + +## Dependencies + +- Python 3.11+ +- httpx (for async HTTP requests) +- pyyaml (for YAML config file support) +- Docker or Podman (for containerized execution) +- PowerShell (for parallel execution script on Windows) + +## License + +Same as LiteLLM project. diff --git a/scripts/health_check/health_check_requirements.txt b/scripts/health_check/health_check_requirements.txt new file mode 100644 index 00000000000..c9d2650c884 --- /dev/null +++ b/scripts/health_check/health_check_requirements.txt @@ -0,0 +1,2 @@ +httpx>=0.24.0 +pyyaml>=6.0 diff --git a/scripts/health_check/run_parallel_health_checks.ps1 b/scripts/health_check/run_parallel_health_checks.ps1 new file mode 100644 index 00000000000..856e7f20ec9 --- /dev/null +++ b/scripts/health_check/run_parallel_health_checks.ps1 @@ -0,0 +1,69 @@ +# Parallel LiteLLM Health Check Runner (PowerShell version) +# +# This script runs multiple health check containers in parallel. +# +# Usage: +# $env:LITELLM_BASE_URL="https://litellm.example.com" +# $env:LITELLM_API_KEY="your-api-key" +# .\run_parallel_health_checks.ps1 [num_parallel_jobs] [image_name] +# +# Defaults: +# - num_parallel_jobs: 16 +# - image_name: litellm/litellm-health-check:latest + +param( + [int]$NumParallelJobs = 16, + [string]$ImageName = "litellm/litellm-health-check:latest", + [string]$ContainerRuntime = "docker" +) + +# Set defaults for environment variables if not provided +if (-not $env:LITELLM_BASE_URL) { + $env:LITELLM_BASE_URL = "https://litellm-perf-cache-and-router.onrender.com" + Write-Warning "LITELLM_BASE_URL not set, using default: $env:LITELLM_BASE_URL" +} + +if (-not $env:LITELLM_API_KEY) { + $env:LITELLM_API_KEY = "sk-1234" + Write-Warning "LITELLM_API_KEY not set, using default: $env:LITELLM_API_KEY" +} + +# Check if container runtime is available +$runtimeExists = Get-Command $ContainerRuntime -ErrorAction SilentlyContinue +if (-not $runtimeExists) { + Write-Error "Error: $ContainerRuntime is not installed" + exit 1 +} + +Write-Host "Running $NumParallelJobs parallel health check containers..." -ForegroundColor Yellow +Write-Host "Using image: $ImageName" -ForegroundColor Yellow +Write-Host "Container runtime: $ContainerRuntime" -ForegroundColor Yellow +Write-Host "LiteLLM Base URL: $env:LITELLM_BASE_URL" -ForegroundColor Cyan +Write-Host "" +Write-Host "NOTE: This will run continuously. Press Ctrl+C to stop." -ForegroundColor Red +Write-Host "" +Write-Host "Troubleshooting:" -ForegroundColor Yellow +Write-Host " - If you see 'All connection attempts failed', check:" -ForegroundColor Yellow +Write-Host " 1. Is the LiteLLM proxy running on the expected port?" -ForegroundColor Yellow +Write-Host " 2. Set LITELLM_BASE_URL to the correct URL (e.g., http://host.docker.internal:PORT)" -ForegroundColor Yellow +Write-Host " 3. On Linux, you may need to use the host IP instead of host.docker.internal" -ForegroundColor Yellow +Write-Host "" + +# Run parallel health checks +# This creates an infinite loop that keeps spawning containers +# Each container tests all models, then exits, and a new one starts +while ($true) { + # Start up to NumParallelJobs containers in parallel + 1..$NumParallelJobs | ForEach-Object -Parallel { + $runtime = $using:ContainerRuntime + $imageName = $using:ImageName + $baseUrl = $env:LITELLM_BASE_URL + $apiKey = $env:LITELLM_API_KEY + + & $runtime run --rm ` + -e LITELLM_BASE_URL="$baseUrl" ` + -e LITELLM_API_KEY="$apiKey" ` + -e LITELLM_JSON_OUTPUT="true" ` + $imageName + } -ThrottleLimit $NumParallelJobs +} diff --git a/scripts/health_check/run_parallel_health_checks.sh b/scripts/health_check/run_parallel_health_checks.sh new file mode 100644 index 00000000000..9b6c5d9f393 --- /dev/null +++ b/scripts/health_check/run_parallel_health_checks.sh @@ -0,0 +1,79 @@ +#!/bin/bash +# Parallel LiteLLM Health Check Runner (Bash version) +# +# This script runs multiple health check containers in parallel. +# +# Usage: +# export LITELLM_BASE_URL="https://litellm.example.com" +# export LITELLM_API_KEY="your-api-key" +# ./run_parallel_health_checks.sh [num_parallel_jobs] [image_name] [container_runtime] +# +# Defaults: +# - num_parallel_jobs: 16 +# - image_name: litellm/litellm-health-check:latest +# - container_runtime: docker + +set -e + +# Default values +NUM_PARALLEL_JOBS="${1:-16}" +IMAGE_NAME="${2:-litellm/litellm-health-check:latest}" +CONTAINER_RUNTIME="${3:-docker}" + +# Set defaults for environment variables if not provided +if [ -z "$LITELLM_BASE_URL" ]; then + export LITELLM_BASE_URL="https://litellm-perf-cache-and-router.onrender.com" + echo "Warning: LITELLM_BASE_URL not set, using default: $LITELLM_BASE_URL" >&2 +fi + +if [ -z "$LITELLM_API_KEY" ]; then + export LITELLM_API_KEY="sk-1234" + echo "Warning: LITELLM_API_KEY not set, using default: $LITELLM_API_KEY" >&2 +fi + +# Check if container runtime is available +if ! command -v "$CONTAINER_RUNTIME" &> /dev/null; then + echo "Error: $CONTAINER_RUNTIME is not installed" >&2 + exit 1 +fi + +# Print configuration +echo "Running $NUM_PARALLEL_JOBS parallel health check containers..." +echo "Using image: $IMAGE_NAME" +echo "Container runtime: $CONTAINER_RUNTIME" +echo "LiteLLM Base URL: $LITELLM_BASE_URL" +echo "" +echo "NOTE: This will run continuously. Press Ctrl+C to stop." +echo "" +echo "Troubleshooting:" +echo " - If you see 'All connection attempts failed', check:" +echo " 1. Is the LiteLLM proxy running on the expected port?" +echo " 2. Set LITELLM_BASE_URL to the correct URL (e.g., http://host.docker.internal:PORT)" +echo " 3. On Linux, you may need to use the host IP instead of host.docker.internal" +echo "" + +# Function to run a single health check container +run_health_check() { + "$CONTAINER_RUNTIME" run --rm \ + -e LITELLM_BASE_URL="$LITELLM_BASE_URL" \ + -e LITELLM_API_KEY="$LITELLM_API_KEY" \ + -e LITELLM_JSON_OUTPUT="true" \ + "$IMAGE_NAME" +} + +# Run parallel health checks +# This creates an infinite loop that keeps spawning containers +# Each container tests all models, then exits, and a new one starts +while true; do + # Start containers in parallel using background jobs + pids=() + for ((i=1; i<=NUM_PARALLEL_JOBS; i++)); do + run_health_check & + pids+=($!) + done + + # Wait for all background jobs to complete + for pid in "${pids[@]}"; do + wait "$pid" 2>/dev/null || true + done +done diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index feb182921db..ea87b56ff30 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/tests/code_coverage_tests/liccheck.ini @@ -89,6 +89,7 @@ tokenizers: >=0.20.2 # Apache 2.0 License jinja2: >=3.1.4 # BSD 3-Clause License litellm-proxy-extras: >=0.1.1 # MIT License litellm-enterprise: >=0.1.1 # LiteLLM Enterprise License +a2a-sdk: >=0.3.22 # Apache 2.0 license anyio: >=4.5.0 # Unknown license httpx-aiohttp: >=0.1.4 # Unknown license backoff: >=2.2.1 # Unknown license diff --git a/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py b/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py index 1103eaf92be..b81161881ea 100644 --- a/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py +++ b/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py @@ -231,7 +231,7 @@ class TestBaseResponsesAPIStreamingIterator: mock_logging_obj = Mock(spec=LiteLLMLoggingObj) mock_logging_obj.model_call_details = {"litellm_params": {}} mock_config = Mock(spec=BaseResponsesAPIConfig) - + # Create the iterator instance iterator = BaseResponsesAPIStreamingIterator( response=mock_response, @@ -239,11 +239,73 @@ class TestBaseResponsesAPIStreamingIterator: responses_api_provider_config=mock_config, logging_obj=mock_logging_obj ) - + # Test with empty chunk result = iterator._process_chunk("") assert result is None - + # Test with None chunk result = iterator._process_chunk(None) - assert result is None \ No newline at end of file + assert result is None + + def test_handle_logging_completed_response_with_unpickleable_objects(self): + """ + Test that _handle_logging_completed_response handles responses containing + objects that cannot be pickled (like Pydantic ValidatorIterator). + + This test verifies the fix for issue #17192 where streaming with tool_choice + containing allowed_tools would fail with: + "cannot pickle 'pydantic_core._pydantic_core.ValidatorIterator' object" + + The fix uses model_dump + model_validate instead of copy.deepcopy. + """ + import asyncio + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + + # Mock dependencies + mock_response = Mock() + mock_response.headers = {} + mock_response.aiter_lines = Mock() + mock_logging_obj = Mock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {"litellm_params": {}} + mock_logging_obj.async_success_handler = Mock() + mock_logging_obj.success_handler = Mock() + mock_config = Mock(spec=BaseResponsesAPIConfig) + + # Create the iterator instance + iterator = ResponsesAPIStreamingIterator( + response=mock_response, + model="gpt-4", + responses_api_provider_config=mock_config, + logging_obj=mock_logging_obj, + litellm_metadata={"model_info": {"id": "model_123"}}, + custom_llm_provider="openai" + ) + + # Create a ResponseCompletedEvent with tool_choice that has model_dump + mock_completed_response = Mock() + mock_completed_response.model_dump.return_value = { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [{"type": "function_call", "name": "search_web"}], + "tool_choice": {"type": "function", "name": "search_web"} + } + } + # model_validate should return a new mock (the copy) + type(mock_completed_response).model_validate = Mock(return_value=Mock()) + + iterator.completed_response = mock_completed_response + + # This should NOT raise an exception + # Previously it would fail with: TypeError: cannot pickle 'ValidatorIterator' + # Mock asyncio.create_task and executor.submit since we're not in async context + with patch('asyncio.create_task') as mock_create_task, \ + patch('litellm.responses.streaming_iterator.executor') as mock_executor: + try: + iterator._handle_logging_completed_response() + except TypeError as e: + if "pickle" in str(e): + pytest.fail(f"_handle_logging_completed_response failed with pickle error: {e}") + raise + diff --git a/tests/mcp_tests/mcp_server.py b/tests/mcp_tests/mcp_server.py index 99a67edd021..bc6accbb721 100644 --- a/tests/mcp_tests/mcp_server.py +++ b/tests/mcp_tests/mcp_server.py @@ -1,9 +1,33 @@ # math_server.py +import argparse +import os + from mcp.server.fastmcp import FastMCP mcp = FastMCP("Math") +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="MCP math test server") + parser.add_argument( + "--transport", + default=os.getenv("MCP_TRANSPORT", "stdio"), + help="Transport to use (stdio or http)", + ) + parser.add_argument( + "--host", + default=os.getenv("MCP_HOST", "127.0.0.1"), + help="Host to bind when serving over HTTP", + ) + parser.add_argument( + "--port", + type=int, + default=int(os.getenv("MCP_PORT", "0")), + help="Port to bind when serving over HTTP", + ) + return parser.parse_args() + + @mcp.tool() def add(a: int, b: int) -> int: """Add two numbers""" @@ -16,5 +40,24 @@ def multiply(a: int, b: int) -> int: return a * b +def main() -> None: + args = _parse_args() + transport = (args.transport or "stdio").lower() + + if transport == "stdio": + mcp.run(transport="stdio") + return + + if transport in {"http", "streamable_http", "streamable-http"}: + if args.port <= 0: + raise ValueError("HTTP transport requires a valid --port value") + mcp.settings.host = args.host + mcp.settings.port = args.port + mcp.run(transport="streamable-http") + return + + raise ValueError(f"Unsupported transport: {transport}") + + if __name__ == "__main__": - mcp.run(transport="stdio") + main() diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/mcp_tests/test_aresponses_api_with_mcp.py index 865a580f0ca..57c79039ee0 100644 --- a/tests/mcp_tests/test_aresponses_api_with_mcp.py +++ b/tests/mcp_tests/test_aresponses_api_with_mcp.py @@ -1,3 +1,4 @@ +import logging import os import sys import pytest @@ -660,16 +661,25 @@ async def test_streaming_mcp_events_validation(): @pytest.mark.asyncio -async def test_streaming_responses_api_with_mcp_tools(): +@pytest.mark.parametrize( + "model", + [ + pytest.param("gpt-4o-mini", id="openai"), + pytest.param("claude-haiku-4-5", id="anthropic"), + ], +) +async def test_streaming_responses_api_with_mcp_tools( + model: str, caplog: pytest.LogCaptureFixture +): """ Test the streaming responses API with MCP tools when using server_url="litellm_proxy" Under the hood the follow occurs - MCP: responses called litellm MCP manager.list_tools (MOCKED) - - Request 1: Made to gpt-4o with fetched tools (REAL LLM CALL) + - Request 1: Made to model under test with fetched tools (REAL LLM CALL) - MCP: Execute tool call from request 1 and returns result (MOCKED) - - Request 2: Made to gpt-4o with fetched tools and tool results (REAL LLM CALL) + - Request 2: Made to model under test with fetched tools and tool results (REAL LLM CALL) Return the user the result of request 2 """ @@ -693,75 +703,101 @@ async def test_streaming_responses_api_with_mcp_tools(): ] # Only mock the MCP-specific operations, let LLM responses be real - with patch.object(LiteLLM_Proxy_MCP_Handler, '_get_mcp_tools_from_manager', new_callable=AsyncMock) as mock_get_tools, \ - patch.object(LiteLLM_Proxy_MCP_Handler, '_execute_tool_calls', new_callable=AsyncMock) as mock_execute_tools: - - # Setup MCP mocks only - mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) - - # Create a dynamic mock that will match the actual tool call ID from the LLM response - def mock_execute_tool_calls_side_effect(tool_calls, user_api_key_auth): - """Mock function that returns results matching the actual tool call IDs from the LLM""" - results = [] - for tool_call in tool_calls: - # Extract call_id from the tool call - call_id = None - if isinstance(tool_call, dict): - call_id = tool_call.get("call_id") or tool_call.get("id") - elif hasattr(tool_call, 'call_id'): - call_id = tool_call.call_id - elif hasattr(tool_call, 'id'): - call_id = tool_call.id - - if call_id: - results.append({ - "tool_call_id": call_id, - "result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output." - }) - return results - - mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect - - # Make the actual call - LLM responses will be real - mcp_tool_config = cast(Any, { - "type": "mcp", - "server_url": "litellm_proxy", - "require_approval": "never" - }) - response = await litellm.aresponses( - model="gpt-4o-mini", - tools=[mcp_tool_config], - tool_choice="required", - input=[ + with caplog.at_level(logging.ERROR): + with patch.object( + LiteLLM_Proxy_MCP_Handler, + '_get_mcp_tools_from_manager', + new_callable=AsyncMock, + ) as mock_get_tools, patch.object( + LiteLLM_Proxy_MCP_Handler, + '_execute_tool_calls', + new_callable=AsyncMock, + ) as mock_execute_tools: + # Setup MCP mocks only + mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) + + # Create a dynamic mock that will match the actual tool call ID from the LLM response + def mock_execute_tool_calls_side_effect( + tool_calls, user_api_key_auth, **kwargs + ): + """Mock function that returns results matching the actual tool call IDs from the LLM""" + results = [] + for tool_call in tool_calls: + # Extract call_id from the tool call + call_id = None + if isinstance(tool_call, dict): + call_id = tool_call.get("call_id") or tool_call.get("id") + elif hasattr(tool_call, 'call_id'): + call_id = tool_call.call_id + elif hasattr(tool_call, 'id'): + call_id = tool_call.id + + if call_id: + results.append( + { + "tool_call_id": call_id, + "result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output.", + } + ) + return results + + mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect + + # Make the actual call - LLM responses will be real + mcp_tool_config = cast( + Any, { - "role": "user", - "type": "message", - "content": "give me a TLDR of what BerriAI/litellm is about" - } - ], - stream=True - ) - - print(f"📋 Response type: {type(response)}") - assert hasattr(response, '__aiter__'), "Response should be an async streaming response" - - # Collect streaming chunks - chunks = [] - async for chunk in response: - chunks.append(chunk) - print(f"📦 Chunk type: {getattr(chunk, 'type', 'unknown')}") - - print(f"📊 Total chunks received: {len(chunks)}") - - # Verify MCP mocks were called (may be called multiple times in streaming) - assert mock_get_tools.call_count >= 1, f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}" - print(f"MCP tools fetched: {len(mock_mcp_tools)}") - - # Verify we got a response - assert response is not None - assert len(chunks) > 0, "Should have received streaming chunks" - - print("Basic streaming responses API with MCP tools test passed!") + "type": "mcp", + "server_url": "litellm_proxy", + "require_approval": "never", + }, + ) + response = await litellm.aresponses( + model=model, + tools=[mcp_tool_config], + tool_choice="required", + input=[ + { + "role": "user", + "type": "message", + "content": "give me a TLDR of what BerriAI/litellm is about", + } + ], + stream=True, + ) + + print(f"📋 Response type: {type(response)}") + assert hasattr(response, '__aiter__'), "Response should be an async streaming response" + + # Collect streaming chunks + chunks = [] + async for chunk in response: + chunks.append(chunk) + print(f"📦 Chunk type: {getattr(chunk, 'type', 'unknown')}") + + print(f"📊 Total chunks received: {len(chunks)}") + + # Verify MCP mocks were called (may be called multiple times in streaming) + assert ( + mock_get_tools.call_count >= 1 + ), f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}" + print(f"MCP tools fetched: {len(mock_mcp_tools)}") + + # Verify we got a response + assert response is not None + assert len(chunks) > 0, "Should have received streaming chunks" + + print("Basic streaming responses API with MCP tools test passed!") + + lite_errors = [ + record + for record in caplog.records + if record.levelno >= logging.ERROR + and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) + ] + assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( + record.getMessage() for record in lite_errors + ) @pytest.mark.asyncio diff --git a/tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml b/tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml new file mode 100644 index 00000000000..37d5359e3ce --- /dev/null +++ b/tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml @@ -0,0 +1,23 @@ +general_settings: + master_key: sk-1234 + +litellm_settings: + drop_params: true + +model_list: + - model_name: openai-gpt-4o-mini + litellm_params: + model: gpt-4o-mini + - model_name: anthropic-claude-haiku-4-5 + litellm_params: + model: anthropic/claude-haiku-4-5 + +mcp_servers: + math_stdio: + transport: stdio + command: python3 + args: + - tests/mcp_tests/mcp_server.py + math_streamable_http: + transport: http + url: http://127.0.0.1:0/mcp diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 66dd2bdc6b3..e8a1231c6fb 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1001,34 +1001,39 @@ async def test_mcp_server_manager_access_groups_from_config(): MCPRequestHandler, ) - # Patch global_mcp_server_manager for this test + # Patch global_mcp_server_manager for this test and restore afterwards to + # avoid leaking state into other tests (e.g. the proxy MCP e2e suite). import litellm.proxy._experimental.mcp_server.mcp_server_manager as mcp_server_manager_mod + original_manager = mcp_server_manager_mod.global_mcp_server_manager mcp_server_manager_mod.global_mcp_server_manager = test_manager - # Should find config_server for group-a, both for group-b, other_server for group-c - import asyncio + try: + # Should find config_server for group-a, both for group-b, other_server for group-c + import asyncio - server_ids_a = await MCPRequestHandler._get_mcp_servers_from_access_groups([ - "group-a" - ]) - server_ids_b = await MCPRequestHandler._get_mcp_servers_from_access_groups([ - "group-b" - ]) - server_ids_c = await MCPRequestHandler._get_mcp_servers_from_access_groups([ - "group-c" - ]) - assert any(config_server.server_id == sid for sid in server_ids_a) - assert set(server_ids_b) == set( - [ - s.server_id + server_ids_a = await MCPRequestHandler._get_mcp_servers_from_access_groups([ + "group-a" + ]) + server_ids_b = await MCPRequestHandler._get_mcp_servers_from_access_groups([ + "group-b" + ]) + server_ids_c = await MCPRequestHandler._get_mcp_servers_from_access_groups([ + "group-c" + ]) + assert any(config_server.server_id == sid for sid in server_ids_a) + assert set(server_ids_b) == set( + [ + s.server_id + for s in test_manager.config_mcp_servers.values() + if "group-b" in s.access_groups + ] + ) + assert any( + s.name == "other_server" and s.server_id in server_ids_c for s in test_manager.config_mcp_servers.values() - if "group-b" in s.access_groups - ] - ) - assert any( - s.name == "other_server" and s.server_id in server_ids_c - for s in test_manager.config_mcp_servers.values() - ) + ) + finally: + mcp_server_manager_mod.global_mcp_server_manager = original_manager async def test_mcp_server_manager_config_integration_with_database(): diff --git a/tests/mcp_tests/test_proxy_mcp_e2e.py b/tests/mcp_tests/test_proxy_mcp_e2e.py new file mode 100644 index 00000000000..2b8cde54710 --- /dev/null +++ b/tests/mcp_tests/test_proxy_mcp_e2e.py @@ -0,0 +1,236 @@ +import asyncio +import os +import socket +import subprocess +import sys +import threading +import time +import typing +from pathlib import Path + +import pytest +import uvicorn +import yaml +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client + +from litellm.proxy.proxy_server import ( + app as proxy_app, + cleanup_router_config_variables, + initialize, +) + + +CONFIG_TEMPLATE_PATH = Path("tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml") +MCP_SERVER_SCRIPT = Path("tests/mcp_tests/mcp_server.py") +PROJECT_ROOT = Path(__file__).resolve().parents[2] +PROXY_START_TIMEOUT = 30 + + +PROXY_AUTHORIZATION_HEADER = "Bearer sk-1234" + + +@pytest.fixture(scope="session", autouse=True) +def _clear_proxy_database_env() -> typing.Iterator[None]: + """Ensure local proxy DB settings don't leak into tests.""" + mp = pytest.MonkeyPatch() + mp.delenv("DATABASE_URL", raising=False) + try: + yield + finally: + mp.undo() + + +def _initialize_proxy(config_path: str) -> None: + cleanup_router_config_variables() + asyncio.run(initialize(config=config_path, debug=True)) + + +def _start_proxy_server(config_path: str) -> tuple[str, uvicorn.Server, threading.Thread, socket.socket]: + _initialize_proxy(config_path) + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("127.0.0.1", 0)) + host, port = sock.getsockname() + + config = uvicorn.Config(proxy_app, host=host, port=port, log_level="warning") + server = uvicorn.Server(config) + + def _run() -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(server.serve(sockets=[sock])) + + thread = threading.Thread(target=_run, daemon=True) + thread.start() + + start_time = time.time() + while not server.started: + if not thread.is_alive(): + raise RuntimeError("Proxy server failed to start") + if time.time() - start_time > PROXY_START_TIMEOUT: + raise TimeoutError("Proxy server did not start in time") + time.sleep(0.05) + + return f"http://{host}:{port}", server, thread, sock + + +@pytest.fixture(scope="session") +def math_streamable_http_server() -> str: + host = "127.0.0.1" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind((host, 0)) + _, port = sock.getsockname() + + cmd = [ + sys.executable, + str(MCP_SERVER_SCRIPT), + "--transport", + "http", + "--host", + host, + "--port", + str(port), + ] + + env = os.environ.copy() + server_process = subprocess.Popen( + cmd, + cwd=str(PROJECT_ROOT), + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + start_time = time.time() + while True: + if server_process.poll() is not None: + stdout, stderr = server_process.communicate() + raise RuntimeError( + f"Streamable HTTP MCP server exited early.\nSTDOUT: {stdout.decode()}\nSTDERR: {stderr.decode()}" + ) + try: + with socket.create_connection((host, port), timeout=0.1): + break + except OSError: + if time.time() - start_time > PROXY_START_TIMEOUT: + server_process.terminate() + raise TimeoutError("Streamable HTTP MCP server did not start in time") + time.sleep(0.05) + + yield f"http://{host}:{port}" + + server_process.terminate() + try: + server_process.wait(timeout=5) + except subprocess.TimeoutExpired: + server_process.kill() + + +@pytest.fixture(scope="session") +def proxy_server_url( + tmp_path_factory: pytest.TempPathFactory, math_streamable_http_server: str +): + config_dir = tmp_path_factory.mktemp("mcp_e2e") + config_path = config_dir / "config.yaml" + config = yaml.safe_load(CONFIG_TEMPLATE_PATH.read_text()) + config["mcp_servers"]["math_streamable_http"][ + "url" + ] = f"{math_streamable_http_server}/mcp" + config_path.write_text(yaml.safe_dump(config)) + + server_url, server, thread, sock = _start_proxy_server(str(config_path)) + + yield server_url + + server.should_exit = True + thread.join(timeout=10) + sock.close() + + +class TestProxyMcpSimpleConnections: + @pytest.mark.asyncio + async def test_proxy_mcp_stdio_roundtrip(self, proxy_server_url: str) -> None: + async with asyncio.timeout(20): + async with streamablehttp_client( + url=f"{proxy_server_url}/mcp", + headers={ + "Authorization": PROXY_AUTHORIZATION_HEADER, + "x-mcp-servers": "math_stdio", + }, + ) as (read, write, _get_session_id): + async with ClientSession(read, write) as session: + await session.initialize() + tools_result = await session.list_tools() + assert any(tool.name.endswith("add") for tool in tools_result.tools) + + result = await session.call_tool( + "add", arguments={"a": 3, "b": 4} + ) + assert result.content + first_content = result.content[0] + text = getattr(first_content, "text", None) + assert text == "7" + + @pytest.mark.asyncio + async def test_proxy_mcp_streamable_http_roundtrip( + self, proxy_server_url: str + ) -> None: + async with asyncio.timeout(20): + async with streamablehttp_client( + url=f"{proxy_server_url}/mcp", + headers={ + "Authorization": PROXY_AUTHORIZATION_HEADER, + "x-mcp-servers": "math_streamable_http", + }, + ) as (read, write, _get_session_id): + async with ClientSession(read, write) as session: + await session.initialize() + tools_result = await session.list_tools() + assert any(tool.name.endswith("add") for tool in tools_result.tools) + + result = await session.call_tool( + "add", arguments={"a": 5, "b": 6} + ) + assert result.content + first_content = result.content[0] + text = getattr(first_content, "text", None) + assert text == "11" + + @pytest.mark.asyncio + async def test_proxy_mcp_lists_all_servers_without_header( + self, proxy_server_url: str + ) -> None: + async with asyncio.timeout(20): + async with streamablehttp_client( + url=f"{proxy_server_url}/mcp", + headers={"Authorization": PROXY_AUTHORIZATION_HEADER}, + ) as (read, write, _get_session_id): + async with ClientSession(read, write) as session: + await session.initialize() + tools_result = await session.list_tools() + tool_names = {tool.name for tool in tools_result.tools} + expected_tool_names = { + "math_stdio-add", + "math_stdio-multiply", + "math_streamable_http-add", + "math_streamable_http-multiply", + } + assert expected_tool_names <= tool_names + + async def _call_and_get_text( + tool_name: str, *, a: int, b: int + ) -> str | None: + result = await session.call_tool(tool_name, arguments={"a": a, "b": b}) + assert result.content + first_content = result.content[0] + return getattr(first_content, "text", None) + + stdio_result = await _call_and_get_text( + "math_stdio-add", a=2, b=3 + ) + streamable_result = await _call_and_get_text( + "math_streamable_http-add", a=4, b=5 + ) + assert stdio_result == "5" + assert streamable_result == "9" diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 42a2b5d0971..a22fe13798f 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1392,3 +1392,134 @@ def test_anthropic_messages_pt_server_tool_use_passthrough(): b for b in assistant_msg["content"] if b.get("type") == "text" ) assert text_block["text"] == "I found the time tool. How can I help you?" + + +def test_bedrock_tools_unpack_defs_no_oom_with_nested_refs(): + """ + Regression test for issue #19098: unpack_defs() causes OOM with nested tool schemas. + + The old implementation had a "flatten defs" loop that would pre-expand each def + using unpack_defs(), but since defs often reference each other, each subsequent + call would copy already-expanded content, causing exponential memory growth. + + This test creates a schema with multiple nested $defs that reference each other + to verify the fix prevents memory explosion while still correctly resolving refs. + """ + import sys + import copy + + from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_tools_pt + + # Schema with multiple nested $defs that reference each other + # This pattern would cause OOM with the old "flatten defs" loop + complex_nested_schema = { + "type": "object", + "properties": { + "query": {"$ref": "#/$defs/Expression"}, + }, + "$defs": { + "Expression": { + "type": "object", + "properties": { + "type": {"type": "string", "enum": ["and", "or", "not", "comparison"]}, + "left": {"$ref": "#/$defs/Operand"}, + "right": {"$ref": "#/$defs/Operand"}, + "operator": {"$ref": "#/$defs/Operator"}, + }, + }, + "Operand": { + "type": "object", + "anyOf": [ + {"$ref": "#/$defs/Literal"}, + {"$ref": "#/$defs/FieldRef"}, + {"$ref": "#/$defs/Expression"}, # Circular: Operand -> Expression -> Operand + ], + }, + "Literal": { + "type": "object", + "properties": { + "type": {"type": "string", "const": "literal"}, + "value": {"$ref": "#/$defs/LiteralValue"}, + }, + }, + "LiteralValue": { + "oneOf": [ + {"type": "string"}, + {"type": "number"}, + {"type": "boolean"}, + {"type": "null"}, + ], + }, + "FieldRef": { + "type": "object", + "properties": { + "type": {"type": "string", "const": "field"}, + "name": {"type": "string"}, + "table": {"$ref": "#/$defs/TableRef"}, + }, + }, + "TableRef": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "alias": {"type": "string"}, + }, + }, + "Operator": { + "type": "string", + "enum": ["=", "!=", "<", ">", "<=", ">=", "LIKE", "IN"], + }, + }, + } + + tools = [ + { + "type": "function", + "function": { + "name": "execute_query", + "description": "Execute a query with complex expressions", + "parameters": complex_nested_schema, + }, + } + ] + + # Measure initial size + def get_size(obj, seen=None): + size = sys.getsizeof(obj) + if seen is None: + seen = set() + obj_id = id(obj) + if obj_id in seen: + return 0 + seen.add(obj_id) + if isinstance(obj, dict): + size += sum([get_size(v, seen) for v in obj.values()]) + size += sum([get_size(k, seen) for k in obj.keys()]) + elif hasattr(obj, "__iter__") and not isinstance(obj, (str, bytes, bytearray)): + size += sum([get_size(i, seen) for i in obj]) + return size + + initial_size = get_size(tools) + + # Process through _bedrock_tools_pt - this should complete without OOM + tools_copy = copy.deepcopy(tools) + result = _bedrock_tools_pt(tools=tools_copy) + + final_size = get_size(result) + + # The expansion factor should be reasonable (< 100x), not exponential (35000x as in #19098) + expansion_factor = final_size / initial_size + assert expansion_factor < 100, ( + f"Memory expansion factor {expansion_factor:.1f}x is too high. " + f"Initial: {initial_size} bytes, Final: {final_size} bytes" + ) + + # Verify the result is valid Bedrock tools format + assert isinstance(result, list) + assert len(result) == 1 + assert "toolSpec" in result[0] + assert result[0]["toolSpec"]["name"] == "execute_query" + + # Verify $defs have been removed (Bedrock doesn't support them) + tool_schema = result[0]["toolSpec"].get("inputSchema", {}).get("json", {}) + assert "$defs" not in tool_schema, "$defs should be removed after expansion" diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index e7a123aa8c3..7e96c4634fd 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -681,6 +681,73 @@ def test_anthropic_chat_headers_add_context_management_beta(): assert headers["anthropic-beta"] == "context-management-2025-06-27" +def test_anthropic_beta_header_merging_with_output_format(): + """ + Test that anthropic-beta headers from extra_headers are merged with + output_format beta headers instead of being overridden. + + This is a regression test for: https://github.com/BerriAI/litellm/issues/... + When using response_format with a Pydantic model AND extra_headers with + anthropic-beta (e.g., for context-1m extension), both beta headers should + be present in the final request. + """ + config = AnthropicConfig() + + # Simulate headers that already have the context-1m beta header from extra_headers + headers = {"anthropic-beta": "context-1m-2025-08-07"} + + # Simulate output_format being set (happens when using response_format with Sonnet 4.5) + optional_params = { + "output_format": { + "type": "json_schema", + "schema": {"type": "object", "properties": {}} + } + } + + result_headers = config.update_headers_with_optional_anthropic_beta( + headers, optional_params + ) + + # Both beta headers should be present + beta_value = result_headers["anthropic-beta"] + assert "context-1m-2025-08-07" in beta_value, \ + f"User's context-1m beta header missing from: {beta_value}" + assert "structured-outputs-2025-11-13" in beta_value, \ + f"Structured output beta header missing from: {beta_value}" + + +def test_anthropic_beta_header_merging_with_multiple_features(): + """ + Test that multiple beta headers can be merged when using multiple features. + """ + config = AnthropicConfig() + + # Start with a user-provided beta header + headers = {"anthropic-beta": "context-1m-2025-08-07"} + + # Use multiple features that require beta headers + optional_params = { + "output_format": { + "type": "json_schema", + "schema": {"type": "object", "properties": {}} + }, + "context_management": _sample_context_management_payload(), + "tools": [{"type": "web_fetch_20250910", "name": "web_fetch"}] + } + + result_headers = config.update_headers_with_optional_anthropic_beta( + headers, optional_params + ) + + beta_value = result_headers["anthropic-beta"] + + # All beta headers should be present + assert "context-1m-2025-08-07" in beta_value + assert "structured-outputs-2025-11-13" in beta_value + assert "context-management-2025-06-27" in beta_value + assert "web-fetch-2025-09-10" in beta_value + + def test_anthropic_chat_transform_request_includes_context_management(): config = AnthropicConfig() headers = {} diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py new file mode 100644 index 00000000000..ceea3d0b16c --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py @@ -0,0 +1,260 @@ +""" +Test Vertex AI binary file upload functionality + +This test ensures that binary files (like PDFs, images) are correctly handled +during upload without attempting UTF-8 decoding, which would cause errors. + +Regression test for: UTF-8 codec error when uploading binary files +""" + +import io +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx + +from litellm.llms.custom_httpx.llm_http_handler import AsyncHTTPHandler +from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig +from litellm.types.llms.openai import CreateFileRequest + + +class TestVertexAIBinaryFileUpload: + """Test binary file upload handling for Vertex AI""" + + def setup_method(self): + """Setup test method""" + self.http_handler = AsyncHTTPHandler() + self.vertex_config = VertexAIFilesConfig() + + @pytest.mark.asyncio + async def test_pdf_file_upload_bytes_handling(self): + """ + Test that PDF binary data is correctly handled without UTF-8 decoding. + + This is a regression test for the error: + 'utf-8' codec can't decode byte 0xc4 in position 10: invalid continuation byte + """ + # Create mock PDF binary data (with non-UTF-8 bytes) + # PDF files start with %PDF- and contain binary data + mock_pdf_content = b"%PDF-1.4\n%\xc4\xe5\xf2\xe5\xeb\xa7\xf3\xa0\xd0\xc4\xc6\n" + mock_pdf_content += b"\x00\x01\x02\x03\xff\xfe\xfd" * 100 # Add more binary data + + # Create file object + file_obj = io.BytesIO(mock_pdf_content) + file_obj.name = "test_document.pdf" + + # Create file request + create_file_data: CreateFileRequest = { + "file": file_obj, + "purpose": "user_data", + } + + # Transform the request + transformed_request = self.vertex_config.transform_create_file_request( + model="vertex_ai/gemini-flash", + create_file_data=create_file_data, + optional_params={}, + litellm_params={}, + ) + + # Verify the transformation returns bytes (not string) + assert isinstance(transformed_request, bytes), ( + f"Expected bytes for binary file, got {type(transformed_request)}" + ) + + # Verify the bytes match the original content + assert transformed_request == mock_pdf_content, ( + "Transformed request should preserve binary content exactly" + ) + + # Verify that the bytes contain non-UTF-8 characters + # This should raise UnicodeDecodeError if we try to decode + with pytest.raises(UnicodeDecodeError): + transformed_request.decode("utf-8") + + @pytest.mark.asyncio + async def test_image_file_upload_bytes_handling(self): + """Test that image binary data (PNG) is correctly handled""" + # Create mock PNG binary data (PNG signature + some binary data) + mock_png_content = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR" + mock_png_content += b"\x00\x01\x02\x03\xff\xfe\xfd" * 50 + + file_obj = io.BytesIO(mock_png_content) + file_obj.name = "test_image.png" + + create_file_data: CreateFileRequest = { + "file": file_obj, + "purpose": "user_data", + } + + transformed_request = self.vertex_config.transform_create_file_request( + model="vertex_ai/gemini-flash", + create_file_data=create_file_data, + optional_params={}, + litellm_params={}, + ) + + # Verify bytes are preserved + assert isinstance(transformed_request, bytes) + assert transformed_request == mock_png_content + + @pytest.mark.asyncio + async def test_http_handler_accepts_bytes_without_decoding(self): + """ + Test that httpx correctly accepts binary data without decoding. + + This test verifies that bytes can be passed to httpx's post/put methods + without needing UTF-8 decoding, which is the core of our fix. + """ + # Create mock binary data with non-UTF-8 bytes + mock_binary_data = b"\x00\x01\x02\x03\xff\xfe\xfd\xc4\xe5\xf2" + + # Test that httpx accepts bytes in the data parameter + # We're testing the behavior, not making an actual request + + # Verify that attempting to decode would fail (proving it's binary) + with pytest.raises(UnicodeDecodeError): + mock_binary_data.decode("utf-8") + + # Verify that httpx Request accepts bytes + try: + request = httpx.Request( + method="POST", + url="https://example.com/upload", + data=mock_binary_data, + headers={"Content-Type": "application/octet-stream"}, + ) + # If we get here, httpx accepts bytes - which is what we need + assert request.content == mock_binary_data + except Exception as e: + pytest.fail(f"httpx should accept bytes in data parameter: {e}") + + # Document the expected behavior + assert isinstance(mock_binary_data, bytes), ( + "Binary file data should remain as bytes" + ) + + @pytest.mark.asyncio + async def test_jsonl_file_upload_returns_string(self): + """ + Test that JSONL files (text) are correctly transformed to strings. + + This ensures we handle both binary and text files correctly. + """ + # Create mock JSONL content + mock_jsonl_content = ( + '{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", ' + '"body": {"model": "gemini-flash", "messages": [{"role": "user", "content": "Hello"}]}}\n' + ) + + file_obj = io.BytesIO(mock_jsonl_content.encode("utf-8")) + file_obj.name = "batch_requests.jsonl" + + create_file_data: CreateFileRequest = { + "file": file_obj, + "purpose": "batch", + } + + transformed_request = self.vertex_config.transform_create_file_request( + model="vertex_ai/gemini-flash", + create_file_data=create_file_data, + optional_params={}, + litellm_params={}, + ) + + # JSONL files should be transformed to string + assert isinstance(transformed_request, str), ( + f"Expected string for JSONL file, got {type(transformed_request)}" + ) + + @pytest.mark.asyncio + async def test_mixed_file_types_in_sequence(self): + """ + Test uploading different file types in sequence to ensure no state pollution. + """ + # Test 1: Upload binary file + binary_content = b"\x00\x01\x02\x03\xff\xfe\xfd" + binary_file = io.BytesIO(binary_content) + binary_file.name = "binary.dat" + + binary_request: CreateFileRequest = { + "file": binary_file, + "purpose": "user_data", + } + + result1 = self.vertex_config.transform_create_file_request( + model="vertex_ai/gemini-flash", + create_file_data=binary_request, + optional_params={}, + litellm_params={}, + ) + assert isinstance(result1, bytes) + + # Test 2: Upload JSONL file + jsonl_content = '{"test": "data"}\n' + jsonl_file = io.BytesIO(jsonl_content.encode("utf-8")) + jsonl_file.name = "batch.jsonl" + + jsonl_request: CreateFileRequest = { + "file": jsonl_file, + "purpose": "batch", + } + + result2 = self.vertex_config.transform_create_file_request( + model="vertex_ai/gemini-flash", + create_file_data=jsonl_request, + optional_params={}, + litellm_params={}, + ) + assert isinstance(result2, str) + + # Test 3: Upload another binary file + binary_content2 = b"\xc4\xe5\xf2\xe5\xeb" + binary_file2 = io.BytesIO(binary_content2) + binary_file2.name = "binary2.dat" + + binary_request2: CreateFileRequest = { + "file": binary_file2, + "purpose": "user_data", + } + + result3 = self.vertex_config.transform_create_file_request( + model="vertex_ai/gemini-flash", + create_file_data=binary_request2, + optional_params={}, + litellm_params={}, + ) + assert isinstance(result3, bytes) + + def test_bytes_type_preservation_documentation(self): + """ + Documentation test: Verify that bytes are the correct type for binary uploads. + + This test documents the expected behavior: + - Binary files (PDF, images, etc.) should remain as bytes + - Text files (JSONL) should be strings + - httpx accepts both bytes and strings in the 'data' parameter + - bytes should NEVER be decoded to UTF-8 for binary files + """ + # This is a documentation test - it always passes + # but serves as a reference for the expected behavior + + expected_behavior = { + "binary_files": { + "input_type": "bytes", + "output_type": "bytes", + "examples": ["PDF", "PNG", "JPEG", "binary data"], + "http_method": "POST or PUT", + "encoding": "none - preserve raw bytes", + }, + "text_files": { + "input_type": "str or bytes", + "output_type": "str", + "examples": ["JSONL", "CSV", "TXT"], + "http_method": "POST", + "encoding": "UTF-8", + }, + } + + assert expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes" + assert expected_behavior["text_files"]["encoding"] == "UTF-8" diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index 061e27da919..9588c3b55c3 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -49,18 +49,20 @@ async def test_invoke_agent_a2a_adds_litellm_data(): # Mock request mock_request = MagicMock() - mock_request.json = AsyncMock(return_value={ - "jsonrpc": "2.0", - "id": "test-id", - "method": "message/send", - "params": { - "message": { - "role": "user", - "parts": [{"kind": "text", "text": "Hello"}], - "messageId": "msg-123", - } - }, - }) + mock_request.json = AsyncMock( + return_value={ + "jsonrpc": "2.0", + "id": "test-id", + "method": "message/send", + "params": { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Hello"}], + "messageId": "msg-123", + } + }, + } + ) mock_user_api_key_dict = UserAPIKeyAuth( api_key="sk-test-key", @@ -77,40 +79,44 @@ async def test_invoke_agent_a2a_adds_litellm_data(): SendMessageRequest, SendStreamingMessageRequest, ) + # Real types available - use them - use_real_types = True + pass except ImportError: # Real types not available - create realistic mocks - use_real_types = False - + pass + def make_mock_pydantic_class(name): """Create a mock class that behaves like a Pydantic model.""" + class MockPydanticClass: def __init__(self, **kwargs): self.__dict__.update(kwargs) # Store kwargs for model_dump() if needed self._kwargs = kwargs - + def model_dump(self, mode="json", exclude_none=False): """Mock model_dump method.""" result = dict(self._kwargs) if exclude_none: result = {k: v for k, v in result.items() if v is not None} return result - + MockPydanticClass.__name__ = name return MockPydanticClass - + MessageSendParams = make_mock_pydantic_class("MessageSendParams") SendMessageRequest = make_mock_pydantic_class("SendMessageRequest") - SendStreamingMessageRequest = make_mock_pydantic_class("SendStreamingMessageRequest") - + SendStreamingMessageRequest = make_mock_pydantic_class( + "SendStreamingMessageRequest" + ) + # Create a mock module for a2a.types mock_a2a_types = MagicMock() mock_a2a_types.MessageSendParams = MessageSendParams mock_a2a_types.SendMessageRequest = SendMessageRequest mock_a2a_types.SendStreamingMessageRequest = SendStreamingMessageRequest - + # Patch at the source modules with patch( "litellm.proxy.agent_endpoints.a2a_endpoints._get_agent", @@ -137,12 +143,15 @@ async def test_invoke_agent_a2a_adds_litellm_data(): ), patch.dict( sys.modules, {"a2a": MagicMock(), "a2a.types": mock_a2a_types}, + ), patch( + "litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", + True, ): from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a mock_fastapi_response = MagicMock() - result = await invoke_agent_a2a( + await invoke_agent_a2a( agent_id="test-agent", request=mock_request, fastapi_response=mock_fastapi_response, diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 4651bf59b40..b86f927ea00 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -856,3 +856,97 @@ def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, l result = response.json() assert result["id"] == "file-abc123" assert result["purpose"] == "fine-tune" + + +def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, llm_router: Router): + """ + Test that managed files work with loadbalancing when both target_model_names + and enable_loadbalancing_on_batch_endpoints are enabled. + + This ensures that the priority order is correct: + - managed files should take precedence over deprecated loadbalancing + - managed files internally use llm_router.acreate_file() which provides loadbalancing + """ + from litellm.llms.base_llm.files.transformation import BaseFileEndpoints + from litellm.types.llms.openai import OpenAIFileObject + + # Enable loadbalancing on batch endpoints + monkeypatch.setattr("litellm.enable_loadbalancing_on_batch_endpoints", True) + + proxy_logging_obj = ProxyLogging( + user_api_key_cache=DualCache(default_in_memory_ttl=1) + ) + proxy_logging_obj._add_proxy_hooks(llm_router) + + # Track calls to verify loadbalancing through router + router_acreate_file_calls = [] + + class ManagedFilesWithLoadbalancing(BaseFileEndpoints): + async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + # Verify we receive the target model names + assert len(target_model_names_list) > 0, "Should have target_model_names_list" + + # Simulate what managed files does - call llm_router.acreate_file for each model + # This is where loadbalancing happens internally + for model in target_model_names_list: + router_acreate_file_calls.append({ + "model": model, + "via_router": True + }) + + # Return a managed file ID (base64 encoded) + return OpenAIFileObject( + id="litellm_managed_file_abc123", + object="file", + bytes=100, + created_at=1234567890, + filename="batch_data.jsonl", + purpose="batch", + status="uploaded", + ) + + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): + raise NotImplementedError("Not implemented for test") + + async def afile_list(self, purpose, litellm_parent_otel_span): + raise NotImplementedError("Not implemented for test") + + async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + raise NotImplementedError("Not implemented for test") + + async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + raise NotImplementedError("Not implemented for test") + + proxy_logging_obj.proxy_hook_mapping["managed_files"] = ManagedFilesWithLoadbalancing() + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj + ) + + # Create batch file content + test_file_content = b'{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}]}}' + test_file = ("batch_data.jsonl", test_file_content, "application/jsonl") + + # Make request with both target_model_names AND enable_loadbalancing_on_batch_endpoints + response = client.post( + "/v1/files", + files={"file": test_file}, + data={ + "purpose": "batch", + "target_model_names": "azure-gpt-3-5-turbo,gpt-3.5-turbo", # Multiple models + }, + headers={"Authorization": "Bearer test-key"}, + ) + + # Verify success + assert response.status_code == 200 + result = response.json() + assert result["id"] == "litellm_managed_file_abc123" + assert result["purpose"] == "batch" + + # Verify that managed files was called (via router for loadbalancing) + # This proves that managed files took precedence over deprecated loadbalancing + assert len(router_acreate_file_calls) == 2, "Should have called router for both models" + assert router_acreate_file_calls[0]["model"] == "azure-gpt-3-5-turbo" + assert router_acreate_file_calls[1]["model"] == "gpt-3.5-turbo" + assert all(call["via_router"] for call in router_acreate_file_calls), "All calls should go through router" diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 352e84719f1..558fe18ae38 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -1051,6 +1051,189 @@ async def test_vector_store_synchronization_across_instances(): ) +@pytest.mark.asyncio +async def test_vector_store_update_and_list_synchronization(): + """ + Test that vector store updates are properly synchronized across multiple instances. + + This test simulates the scenario where: + 1. Instance 1 creates a vector store + 2. Instance 2 caches it in memory + 3. Instance 1 updates the vector store in the database + 4. Instance 2 should see the updated data when listing (database is source of truth) + + This is a regression test to prevent the bug where Instance 2 would show + stale cached data instead of the updated database version. + """ + from datetime import datetime, timezone + from unittest.mock import AsyncMock, MagicMock + + from litellm.types.vector_stores import LiteLLM_ManagedVectorStore + from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + + # Simulate two instances with separate in-memory registries + instance_1_registry = VectorStoreRegistry(vector_stores=[]) + instance_2_registry = VectorStoreRegistry(vector_stores=[]) + + # Mock database that both instances share + mock_db_vector_stores = [] + + async def mock_find_many(order=None): + """Mock find_many for listing vector stores""" + result = [] + for vs in mock_db_vector_stores: + class MockVectorStore: + def __init__(self, data): + for key, value in data.items(): + setattr(self, key, value) + self._data = data + + def __iter__(self): + return iter(self._data.items()) + result.append(MockVectorStore(vs)) + return result + + async def mock_create(data): + """Mock create for adding vector store to DB""" + vector_store = data.copy() + mock_db_vector_stores.append(vector_store) + mock_obj = MagicMock() + mock_obj.model_dump.return_value = vector_store + return mock_obj + + async def mock_update(where, data): + """Mock update for modifying vector store in DB""" + vector_store_id = where.get("vector_store_id") + for i, vs in enumerate(mock_db_vector_stores): + if vs.get("vector_store_id") == vector_store_id: + # Update the vector store + mock_db_vector_stores[i].update(data) + mock_obj = MagicMock() + mock_obj.model_dump.return_value = mock_db_vector_stores[i] + return mock_obj + raise Exception(f"Vector store {vector_store_id} not found") + + # Create mock prisma client + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_managedvectorstorestable.find_many = AsyncMock( + side_effect=mock_find_many + ) + mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock( + side_effect=mock_create + ) + mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock( + side_effect=mock_update + ) + + # Test vector store data + test_vector_store_id = "test-update-store-001" + original_name = "Original Name" + updated_name = "Updated Name" + + test_vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": test_vector_store_id, + "custom_llm_provider": "bedrock", + "vector_store_name": original_name, + "vector_store_description": "Testing update synchronization", + "litellm_params": { + "vector_store_id": test_vector_store_id, + "custom_llm_provider": "bedrock", + "region_name": "us-east-1" + }, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + } + + # Step 1: Create vector store on Instance 1 + await mock_prisma_client.db.litellm_managedvectorstorestable.create( + data=test_vector_store + ) + instance_1_registry.add_vector_store_to_registry(vector_store=test_vector_store) + + # Step 2: Instance 2 fetches and caches the vector store + vector_stores_from_db = await VectorStoreRegistry._get_vector_stores_from_db( + prisma_client=mock_prisma_client + ) + for vs in vector_stores_from_db: + if vs.get("vector_store_id") == test_vector_store_id: + instance_2_registry.add_vector_store_to_registry(vector_store=vs) + + # Verify both instances have the original data + instance_1_vs = instance_1_registry.get_litellm_managed_vector_store_from_registry( + test_vector_store_id + ) + instance_2_vs = instance_2_registry.get_litellm_managed_vector_store_from_registry( + test_vector_store_id + ) + assert instance_1_vs.get("vector_store_name") == original_name + assert instance_2_vs.get("vector_store_name") == original_name + + # Step 3: Instance 1 updates the vector store in the database + # (Simulating what happens in update_vector_store endpoint) + update_data = {"vector_store_name": updated_name} + await mock_prisma_client.db.litellm_managedvectorstorestable.update( + where={"vector_store_id": test_vector_store_id}, + data=update_data + ) + + # Instance 1 updates its own cache + updated_vs_instance_1 = test_vector_store.copy() + updated_vs_instance_1["vector_store_name"] = updated_name + instance_1_registry.update_vector_store_in_registry( + vector_store_id=test_vector_store_id, + updated_data=updated_vs_instance_1 + ) + + # Verify Instance 1 has the updated data + instance_1_vs_after_update = instance_1_registry.get_litellm_managed_vector_store_from_registry( + test_vector_store_id + ) + assert instance_1_vs_after_update.get("vector_store_name") == updated_name + + # Verify Instance 2 still has stale data in cache + instance_2_vs_before_list = instance_2_registry.get_litellm_managed_vector_store_from_registry( + test_vector_store_id + ) + assert instance_2_vs_before_list.get("vector_store_name") == original_name, ( + "Instance 2 should still have stale cached data before list operation" + ) + + # Step 4: Instance 2 calls list endpoint (which should sync with database) + # This simulates what list_vector_stores endpoint does + vector_stores_from_db_after_update = await VectorStoreRegistry._get_vector_stores_from_db( + prisma_client=mock_prisma_client + ) + + # Build map from database vector stores (database is source of truth) + vector_store_map = {} + for vector_store in vector_stores_from_db_after_update: + vector_store_id = vector_store.get("vector_store_id") + if vector_store_id: + vector_store_map[vector_store_id] = vector_store + + # Update in-memory registry with database versions (this is the key fix) + instance_2_registry.update_vector_store_in_registry( + vector_store_id=vector_store_id, + updated_data=vector_store + ) + + # Step 5: Verify Instance 2 now has the updated data + instance_2_vs_after_list = instance_2_registry.get_litellm_managed_vector_store_from_registry( + test_vector_store_id + ) + assert instance_2_vs_after_list.get("vector_store_name") == updated_name, ( + "Instance 2 should have updated data after list operation syncs with database" + ) + + # Verify the list returned the correct data + combined_vector_stores = list(vector_store_map.values()) + assert len(combined_vector_stores) == 1 + assert combined_vector_stores[0].get("vector_store_id") == test_vector_store_id + assert combined_vector_stores[0].get("vector_store_name") == updated_name, ( + "List should return updated data from database" + ) + + @pytest.mark.asyncio async def test_resolve_embedding_config_from_db(): """Test that _resolve_embedding_config_from_db correctly resolves embedding config from database.""" diff --git a/ui/litellm-dashboard/build_release_ui.sh b/ui/litellm-dashboard/build_release_ui.sh new file mode 100755 index 00000000000..4f1168502c6 --- /dev/null +++ b/ui/litellm-dashboard/build_release_ui.sh @@ -0,0 +1,22 @@ +#!/bin/bash +set -e + +destination_dir="../../litellm/proxy/_experimental/out" + +chmod +x ./build_ui.sh +./build_ui.sh + +commit_message="chore: update Next.js build artifacts ($(date -u +"%Y-%m-%d %H:%M UTC"), node $(node -v))" + +if git rev-parse --is-inside-work-tree > /dev/null 2>&1; then + git add -f "$destination_dir"/ + + if ! git diff --cached --quiet; then + git commit -m "$commit_message" + echo "Git commit created." + else + echo "No changes to commit." + fi +else + echo "Not a git repository. Skipping commit." +fi diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index ac019a7a8c5..8a56287e156 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -17,7 +17,7 @@ import { Team } from "@/components/key_team_helpers/key_list"; import { MCPServers } from "@/components/mcp_tools"; import ModelHubTable from "@/components/AIHub/ModelHubTable"; import Navbar from "@/components/navbar"; -import { getUiConfig, Organization, proxyBaseUrl, setGlobalLitellmHeaderName } from "@/components/networking"; +import { getUiConfig, Organization, proxyBaseUrl, setGlobalLitellmHeaderName, getInProductNudgesCall } from "@/components/networking"; import NewUsagePage from "@/components/UsagePage/components/UsagePageView"; import OldTeams from "@/components/OldTeams"; import { fetchUserModels } from "@/components/organisms/create_key_button"; @@ -27,7 +27,7 @@ import PromptsPanel from "@/components/prompts"; import PublicModelHub from "@/components/public_model_hub"; import { SearchTools } from "@/components/search_tools"; import Settings from "@/components/settings"; -import { SurveyPrompt, SurveyModal } from "@/components/survey"; +import { SurveyPrompt, SurveyModal, ClaudeCodePrompt, ClaudeCodeModal } from "@/components/survey"; import TagManagement from "@/components/tag_management"; import TransformRequestPanel from "@/components/transform_request"; import UIThemeSettings from "@/components/ui_theme_settings"; @@ -124,6 +124,11 @@ export default function CreateKeyPage() { const [showSurveyPrompt, setShowSurveyPrompt] = useState(true); const [showSurveyModal, setShowSurveyModal] = useState(false); + // Claude Code feedback state + const [isClaudeCode, setIsClaudeCode] = useState(false); + const [showClaudeCodePrompt, setShowClaudeCodePrompt] = useState(false); + const [showClaudeCodeModal, setShowClaudeCodeModal] = useState(false); + const invitation_id = searchParams.get("invitation_id"); // Get page from URL, default to 'api-keys' if not present @@ -267,6 +272,29 @@ export default function CreateKeyPage() { } }, [accessToken, userID, userRole]); + // Fetch in-product nudges configuration from backend + useEffect(() => { + if (accessToken && token) { + (async () => { + try { + const nudgesConfig = await getInProductNudgesCall(accessToken); + const isUsingClaudeCode = nudgesConfig?.is_claude_code_enabled || false; + setIsClaudeCode(isUsingClaudeCode); + + // Show Claude Code prompt on login if enabled + if (isUsingClaudeCode) { + setShowClaudeCodePrompt(true); + // Don't show the regular survey prompt if showing Claude Code prompt + setShowSurveyPrompt(false); + } + } catch (error) { + console.error("Failed to fetch in-product nudges:", error); + // Silently fail and don't show Claude Code nudge + } + })(); + } + }, [accessToken, token]); + // Auto-dismiss survey prompt after 15 seconds useEffect(() => { if (showSurveyPrompt && !showSurveyModal) { @@ -277,6 +305,16 @@ export default function CreateKeyPage() { } }, [showSurveyPrompt, showSurveyModal]); + // Auto-dismiss Claude Code prompt after 15 seconds + useEffect(() => { + if (showClaudeCodePrompt && !showClaudeCodeModal) { + const timer = setTimeout(() => { + setShowClaudeCodePrompt(false); + }, 15000); + return () => clearTimeout(timer); + } + }, [showClaudeCodePrompt, showClaudeCodeModal]); + const handleOpenSurvey = () => { setShowSurveyPrompt(false); setShowSurveyModal(true); @@ -296,6 +334,25 @@ export default function CreateKeyPage() { setShowSurveyPrompt(true); }; + const handleOpenClaudeCode = () => { + setShowClaudeCodePrompt(false); + setShowClaudeCodeModal(true); + }; + + const handleDismissClaudeCodePrompt = () => { + setShowClaudeCodePrompt(false); + }; + + const handleClaudeCodeComplete = () => { + setShowClaudeCodeModal(false); + }; + + const handleClaudeCodeModalClose = () => { + // If they close the modal without completing, show the prompt again + setShowClaudeCodeModal(false); + setShowClaudeCodePrompt(true); + }; + if (authLoading || redirectToLogin) { return ; } @@ -503,6 +560,18 @@ export default function CreateKeyPage() { onClose={handleSurveyModalClose} onComplete={handleSurveyComplete} /> + + {/* Claude Code Components */} + + )} diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 9079fd6a7ed..5bd756f1d1d 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -7,6 +7,7 @@ import { BarChartOutlined, BgColorsOutlined, BlockOutlined, + BookOutlined, CreditCardOutlined, DatabaseOutlined, ExperimentOutlined, @@ -47,6 +48,7 @@ interface MenuItem { roles?: string[]; children?: MenuItem[]; icon?: React.ReactNode; + external_url?: string; } // Group configuration @@ -213,6 +215,13 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse label: "AI Hub", icon: , }, + { + key: "learning-resources", + page: "learning-resources", + label: "Learning Resources", + icon: , + external_url: "https://models.litellm.ai/cookbook", + }, { key: "experimental", page: "experimental", @@ -252,7 +261,7 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse page: "usage", label: "Old Usage", icon: , - }, + } ], }, ], @@ -364,9 +373,23 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse key: child.key, icon: child.icon, label: child.label, - onClick: () => navigateToPage(child.page), + onClick: () => { + if (child.external_url) { + window.open(child.external_url, "_blank"); + } else { + navigateToPage(child.page); + } + }, })), - onClick: !item.children ? () => navigateToPage(item.page) : undefined, + onClick: !item.children + ? () => { + if (item.external_url) { + window.open(item.external_url, "_blank"); + } else { + navigateToPage(item.page); + } + } + : undefined, })), }); }); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 82894dfb0e2..c975877d06c 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -35,6 +35,36 @@ export const getCallbackConfigsCall = async (accessToken: string) => { throw error; } }; + +export const getInProductNudgesCall = async (accessToken: string) => { + /** + * Get in-product nudges configuration. + */ + try { + let url = proxyBaseUrl ? `${proxyBaseUrl}/in_product_nudges` : `/in_product_nudges`; + + const response = await fetch(url, { + method: "GET", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + const data = await response.json(); + return data; + } catch (error) { + console.error("Failed to get in-product nudges:", error); + throw error; + } +}; /** * Helper file for calls being made to proxy */ diff --git a/ui/litellm-dashboard/src/components/survey/ClaudeCodeModal.tsx b/ui/litellm-dashboard/src/components/survey/ClaudeCodeModal.tsx new file mode 100644 index 00000000000..eac5e8b7a41 --- /dev/null +++ b/ui/litellm-dashboard/src/components/survey/ClaudeCodeModal.tsx @@ -0,0 +1,69 @@ +import React from "react"; +import { X, Code, ExternalLink } from "lucide-react"; +import { Button } from "antd"; + +interface ClaudeCodeModalProps { + isOpen: boolean; + onClose: () => void; + onComplete: () => void; +} + +const GOOGLE_FORM_URL = "https://forms.gle/LZeJQ3XytBakckYa9"; + +export function ClaudeCodeModal({ isOpen, onClose, onComplete }: ClaudeCodeModalProps) { + if (!isOpen) return null; + + const handleOpenForm = () => { + window.open(GOOGLE_FORM_URL, "_blank", "noopener,noreferrer"); + onComplete(); + }; + + return ( +
+ {/* Backdrop */} +
+ + {/* Modal */} +
+ {/* Header */} +
+
+ + Claude Code Feedback +
+ +
+ + {/* Content */} +
+

+ Help us improve your experience +

+

+ We'd love to hear about your experience using LiteLLM with Claude Code. Your feedback helps us improve the product for everyone. +

+

+ This brief survey takes about 2-3 minutes to complete. +

+ + +
+
+
+ ); +} + diff --git a/ui/litellm-dashboard/src/components/survey/ClaudeCodePrompt.tsx b/ui/litellm-dashboard/src/components/survey/ClaudeCodePrompt.tsx new file mode 100644 index 00000000000..9006575ce0a --- /dev/null +++ b/ui/litellm-dashboard/src/components/survey/ClaudeCodePrompt.tsx @@ -0,0 +1,26 @@ +import React from "react"; +import { Code } from "lucide-react"; +import { NudgePrompt } from "./NudgePrompt"; + +interface ClaudeCodePromptProps { + onOpen: () => void; + onDismiss: () => void; + isVisible: boolean; +} + +export function ClaudeCodePrompt({ onOpen, onDismiss, isVisible }: ClaudeCodePromptProps) { + return ( + + ); +} + diff --git a/ui/litellm-dashboard/src/components/survey/NudgePrompt.tsx b/ui/litellm-dashboard/src/components/survey/NudgePrompt.tsx new file mode 100644 index 00000000000..9095c6c21c5 --- /dev/null +++ b/ui/litellm-dashboard/src/components/survey/NudgePrompt.tsx @@ -0,0 +1,91 @@ +import React, { useEffect, useState } from "react"; +import { X, LucideIcon } from "lucide-react"; +import { Button } from "antd"; + +interface NudgePromptProps { + onOpen: () => void; + onDismiss: () => void; + isVisible: boolean; + title: string; + description: string; + buttonText: string; + icon: LucideIcon; + accentColor: string; + buttonStyle?: React.CSSProperties; +} + +const DISMISS_DURATION = 15000; // 15 seconds + +export function NudgePrompt({ + onOpen, + onDismiss, + isVisible, + title, + description, + buttonText, + icon: Icon, + accentColor, + buttonStyle, +}: NudgePromptProps) { + const [progress, setProgress] = useState(100); + + useEffect(() => { + if (!isVisible) { + setProgress(100); + return; + } + + const startTime = Date.now(); + const interval = setInterval(() => { + const elapsed = Date.now() - startTime; + const remaining = Math.max(0, 100 - (elapsed / DISMISS_DURATION) * 100); + setProgress(remaining); + + if (remaining <= 0) { + clearInterval(interval); + } + }, 50); + + return () => clearInterval(interval); + }, [isVisible]); + + if (!isVisible) return null; + + return ( +
+ {/* Progress bar at top showing time remaining */} +
+
+
+ +
+
+
+ + {title} +
+ +
+ +

{description}

+ + +
+
+ ); +} + diff --git a/ui/litellm-dashboard/src/components/survey/SurveyPrompt.tsx b/ui/litellm-dashboard/src/components/survey/SurveyPrompt.tsx index 69a886bddeb..a41acb265a7 100644 --- a/ui/litellm-dashboard/src/components/survey/SurveyPrompt.tsx +++ b/ui/litellm-dashboard/src/components/survey/SurveyPrompt.tsx @@ -1,6 +1,6 @@ -import React, { useEffect, useState } from "react"; -import { MessageSquare, X } from "lucide-react"; -import { Button } from "antd"; +import React from "react"; +import { MessageSquare } from "lucide-react"; +import { NudgePrompt } from "./NudgePrompt"; interface SurveyPromptProps { onOpen: () => void; @@ -8,70 +8,18 @@ interface SurveyPromptProps { isVisible: boolean; } -const DISMISS_DURATION = 15000; // 15 seconds - export function SurveyPrompt({ onOpen, onDismiss, isVisible }: SurveyPromptProps) { - const [progress, setProgress] = useState(100); - - useEffect(() => { - if (!isVisible) { - setProgress(100); - return; - } - - const startTime = Date.now(); - const interval = setInterval(() => { - const elapsed = Date.now() - startTime; - const remaining = Math.max(0, 100 - (elapsed / DISMISS_DURATION) * 100); - setProgress(remaining); - - if (remaining <= 0) { - clearInterval(interval); - } - }, 50); - - return () => clearInterval(interval); - }, [isVisible]); - - if (!isVisible) return null; - return ( -
- {/* Progress bar at top showing time remaining */} -
-
-
- -
-
-
- - Quick feedback -
- -
- -

- Help us improve LiteLLM! Share your experience in 5 quick questions. -

- - -
-
+ ); } diff --git a/ui/litellm-dashboard/src/components/survey/index.tsx b/ui/litellm-dashboard/src/components/survey/index.tsx index 7c36419980c..fde05084b7e 100644 --- a/ui/litellm-dashboard/src/components/survey/index.tsx +++ b/ui/litellm-dashboard/src/components/survey/index.tsx @@ -1,3 +1,6 @@ export { SurveyPrompt } from "./SurveyPrompt"; export { SurveyModal } from "./SurveyModal"; +export { ClaudeCodePrompt } from "./ClaudeCodePrompt"; +export { ClaudeCodeModal } from "./ClaudeCodeModal"; +export { NudgePrompt } from "./NudgePrompt";