diff --git a/.circleci/requirements.txt b/.circleci/requirements.txt index 8c44dc1830..a5ec74424f 100644 --- a/.circleci/requirements.txt +++ b/.circleci/requirements.txt @@ -16,4 +16,5 @@ uvloop==0.21.0 mcp==1.25.0 # for MCP server semantic_router==0.1.10 # for auto-routing with litellm fastuuid==0.12.0 -responses==0.25.7 # for proxy client tests \ No newline at end of file +responses==0.25.7 # for proxy client tests +pytest-retry==1.6.3 # for automatic test retries \ No newline at end of file diff --git a/AGENTS.md b/AGENTS.md index 61afbd035f..5a48049ef4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -51,12 +51,14 @@ LiteLLM is a unified interface for 100+ LLMs that: ### MAKING CODE CHANGES FOR THE UI (IGNORE FOR BACKEND) -1. **Use Common Components as much as possible**: +1. **Tremor is DEPRECATED, do not use Tremor components in new features/changes** + - The only exception is the Tremor Table component and its required Tremor Table sub components. + +2. **Use Common Components as much as possible**: - These are usually defined in the `common_components` directory - Use these components as much as possible and avoid building new components unless needed - - Tremor components are deprecated; prefer using Ant Design (AntD) as much as possible -2. **Testing**: +3. **Testing**: - The codebase uses **Vitest** and **React Testing Library** - **Query Priority Order**: Use query methods in this order: `getByRole`, `getByLabelText`, `getByPlaceholderText`, `getByText`, `getByTestId` - **Always use `screen`** instead of destructuring from `render()` (e.g., use `screen.getByText()` not `getByText`) diff --git a/Dockerfile b/Dockerfile index 0e7a8412bb..2c54e2dec2 100644 --- a/Dockerfile +++ b/Dockerfile @@ -69,8 +69,8 @@ RUN find /usr/lib -type f -path "*/tornado/test/*" -delete && \ # Convert Windows line endings to Unix and make executable RUN sed -i 's/\r$//' docker/install_auto_router.sh && chmod +x docker/install_auto_router.sh && ./docker/install_auto_router.sh -# Generate prisma client -RUN prisma generate +# Generate prisma client using the correct schema +RUN prisma generate --schema=./litellm/proxy/schema.prisma # Convert Windows line endings to Unix for entrypoint scripts RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh RUN sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh diff --git a/README.md b/README.md index 914fda384b..77adddf897 100644 --- a/README.md +++ b/README.md @@ -267,6 +267,7 @@ Support for more providers. Missing a provider or LLM Platform, raise a [feature Greptile OpenHands

Netflix

+ OpenAI Agents SDK diff --git a/docs/my-website/docs/a2a.md b/docs/my-website/docs/a2a.md index d7145e4b83..a7e8b52d99 100644 --- a/docs/my-website/docs/a2a.md +++ b/docs/my-website/docs/a2a.md @@ -68,7 +68,7 @@ Follow [this guide, to add your pydantic ai agent to LiteLLM Agent Gateway](./pr ## Invoking your Agents -Use the [A2A Python SDK](https://pypi.org/project/a2a/) to invoke agents through LiteLLM. +Use the [A2A Python SDK](https://pypi.org/project/a2a-sdk) to invoke agents through LiteLLM. This example shows how to: 1. **List available agents** - Query `/v1/agents` to see which agents your key can access @@ -193,6 +193,120 @@ The logs show: style={{width: '100%', display: 'block', margin: '2rem auto'}} /> + +## Forwarding LiteLLM Context Headers + +When LiteLLM invokes your A2A agent, it sends special headers that enable: +- **Trace Grouping**: All LLM calls from the same agent execution appear under one trace +- **Agent Spend Tracking**: Costs are attributed to the specific agent + +| Header | Purpose | +|--------|---------| +| `X-LiteLLM-Trace-Id` | Links all LLM calls to the same execution flow | +| `X-LiteLLM-Agent-Id` | Attributes spend to the correct agent | + + +To enable these features, your A2A server must **forward these headers** to any LLM calls it makes back to LiteLLM. + +### Implementation Steps + +**Step 1: Extract headers from incoming A2A request** +```python def get_litellm_headers(request) -> dict: + """Extract X-LiteLLM-* headers from incoming A2A request.""" + all_headers = request.call_context.state.get('headers', {}) + return { + k: v for k, v in all_headers.items() + if k.lower().startswith('x-litellm-') + } +``` + +**Step 2: Forward headers to your LLM calls** +Pass the extracted headers when making calls back to LiteLLM: + + + +```python from openai import OpenAI + +headers = get_litellm_headers(request) + +client = OpenAI( + api_key="sk-your-litellm-key", + base_url="http://localhost:4000", + default_headers=headers, # Forward headers +) + +response = client.chat.completions.create( + model="gpt-4o", + messages=[{"role": "user", "content": "Hello"}] +) +``` + + + + +```python +from langchain_openai import ChatOpenAI + +headers = get_litellm_headers(request) + +llm = ChatOpenAI( + model="gpt-4o", + openai_api_key="sk-your-litellm-key", + base_url="http://localhost:4000", + default_headers=headers, # Forward headers +) +``` + + + +```python +import litellm + +headers = get_litellm_headers(request) + +response = litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + api_base="http://localhost:4000", + extra_headers=headers, # Forward headers +) +``` + + + +```python +import httpx + +headers = get_litellm_headers(request) +headers["Authorization"] = "Bearer sk-your-litellm-key" + +response = httpx.post( + "http://localhost:4000/v1/chat/completions", + headers=headers, + json={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]} +) +``` + + + +### Result + +With header forwarding enabled, you'll see: + +**Trace Grouping in Langfuse:** + + + +**Agent Spend Attribution:** + + + ## API Reference ### Endpoint diff --git a/docs/my-website/docs/observability/datadog.md b/docs/my-website/docs/observability/datadog.md index 7cf91ced34..6f785be101 100644 --- a/docs/my-website/docs/observability/datadog.md +++ b/docs/my-website/docs/observability/datadog.md @@ -7,6 +7,7 @@ import TabItem from '@theme/TabItem'; LiteLLM Supports logging to the following Datdog Integrations: - `datadog` [Datadog Logs](https://docs.datadoghq.com/logs/) - `datadog_llm_observability` [Datadog LLM Observability](https://www.datadoghq.com/product/llm-observability/) +- `datadog_cost_management` [Datadog Cloud Cost Management](#datadog-cloud-cost-management) - `ddtrace-run` [Datadog Tracing](#datadog-tracing) ## Datadog Logs @@ -73,7 +74,7 @@ Send logs through a local DataDog agent (useful for containerized environments): ```shell LITELLM_DD_AGENT_HOST="localhost" # hostname or IP of DataDog agent LITELLM_DD_AGENT_PORT="10518" # [OPTIONAL] port of DataDog agent (default: 10518) -DD_API_KEY="5f2d0f310***********" # [OPTIONAL] your datadog API Key (agent handles auth) +DD_API_KEY="5f2d0f310***********" # [OPTIONAL] your datadog API Key (Agent handles auth for Logs. REQUIRED for LLM Observability) DD_SOURCE="litellm_dev" # [OPTIONAL] your datadog source ``` @@ -84,6 +85,9 @@ When `LITELLM_DD_AGENT_HOST` is set, logs are sent to the agent instead of direc **Note:** We use `LITELLM_DD_AGENT_HOST` instead of `DD_AGENT_HOST` to avoid conflicts with `ddtrace` which automatically sets `DD_AGENT_HOST` for APM tracing. +> [!IMPORTANT] +> **Datadog LLM Observability**: `DD_API_KEY` is **REQUIRED** even when using the Datadog Agent (`LITELLM_DD_AGENT_HOST`). The agent acts as a proxy but the API key header is mandatory for the LLM Observability endpoint. + **Step 3**: Start the proxy, make a test request Start proxy @@ -161,6 +165,50 @@ On the Datadog LLM Observability page, you should see that both input messages a + + + +## Datadog Cloud Cost Management + +| Feature | Details | +|---------|---------| +| **What is logged** | Aggregated LLM Costs (FOCUS format) | +| **Events** | Periodic Uploads of Aggregated Cost Data | +| **Product Link** | [Datadog Cloud Cost Management](https://docs.datadoghq.com/cost_management/) | + +We will use the `--config` to set `litellm.callbacks = ["datadog_cost_management"]`. This will periodically upload aggregated LLM cost data to Datadog. + +**Step 1**: Create a `config.yaml` file and set `litellm_settings`: `success_callback` + +```yaml +model_list: + - model_name: gpt-3.5-turbo + litellm_params: + model: gpt-3.5-turbo +litellm_settings: + callbacks: ["datadog_cost_management"] +``` + +**Step 2**: Set Required env variables + +```shell +DD_API_KEY="your-api-key" +DD_APP_KEY="your-app-key" # REQUIRED for Cost Management +DD_SITE="us5.datadoghq.com" +``` + +**Step 3**: Start the proxy + +```shell +litellm --config config.yaml +``` + +**How it works** +* LiteLLM aggregates costs in-memory by Provider, Model, Date, and Tags. +* Requires `DD_APP_KEY` for the Custom Costs API. +* Costs are uploaded periodically (flushed). + + ### Datadog Tracing Use `ddtrace-run` to enable [Datadog Tracing](https://ddtrace.readthedocs.io/en/stable/installation_quickstart.html) on litellm proxy @@ -203,5 +251,5 @@ LiteLLM supports customizing the following Datadog environment variables | `POD_NAME` | Pod name tag (useful for Kubernetes deployments) | "unknown" | ❌ No | \* **Required when using Direct API** (default): `DD_API_KEY` and `DD_SITE` are required -\* **Optional when using DataDog Agent**: Set `LITELLM_DD_AGENT_HOST` to use agent mode; `DD_API_KEY` and `DD_SITE` are not required +\* **Optional when using DataDog Agent**: Set `LITELLM_DD_AGENT_HOST` to use agent mode; `DD_API_KEY` and `DD_SITE` are not required for **Datadog Logs**. (**Note: `DD_API_KEY` IS REQUIRED for Datadog LLM Observability**) diff --git a/docs/my-website/docs/providers/vercel_ai_gateway.md b/docs/my-website/docs/providers/vercel_ai_gateway.md index 91f0a18ea1..3ff007171e 100644 --- a/docs/my-website/docs/providers/vercel_ai_gateway.md +++ b/docs/my-website/docs/providers/vercel_ai_gateway.md @@ -11,7 +11,7 @@ import TabItem from '@theme/TabItem'; | Provider Route on LiteLLM | `vercel_ai_gateway/` | | Link to Provider Doc | [Vercel AI Gateway Documentation ↗](https://vercel.com/docs/ai-gateway) | | Base URL | `https://ai-gateway.vercel.sh/v1` | -| Supported Operations | `/chat/completions`, `/models` | +| Supported Operations | `/chat/completions`, `/embeddings`, `/models` |

@@ -73,7 +73,7 @@ messages = [{"content": "Hello, how are you?", "role": "user"}] # Vercel AI Gateway call with streaming response = completion( - model="vercel_ai_gateway/openai/gpt-4o", + model="vercel_ai_gateway/openai/gpt-4o", messages=messages, stream=True ) @@ -82,6 +82,33 @@ for chunk in response: print(chunk) ``` +### Embeddings + +```python showLineNumbers title="Vercel AI Gateway Embeddings" +import os +from litellm import embedding + +os.environ["VERCEL_AI_GATEWAY_API_KEY"] = "your-api-key" + +# Vercel AI Gateway embedding call +response = embedding( + model="vercel_ai_gateway/openai/text-embedding-3-small", + input="Hello world" +) + +print(response.data[0]["embedding"][:5]) # Print first 5 dimensions +``` + +You can also specify the `dimensions` parameter: + +```python showLineNumbers title="Vercel AI Gateway Embeddings with Dimensions" +response = embedding( + model="vercel_ai_gateway/openai/text-embedding-3-small", + input=["Hello world", "Goodbye world"], + dimensions=768 +) +``` + ## Usage - LiteLLM Proxy Add the following to your LiteLLM Proxy configuration file: @@ -97,6 +124,11 @@ model_list: litellm_params: model: vercel_ai_gateway/anthropic/claude-4-sonnet api_key: os.environ/VERCEL_AI_GATEWAY_API_KEY + + - model_name: text-embedding-3-small-gateway + litellm_params: + model: vercel_ai_gateway/openai/text-embedding-3-small + api_key: os.environ/VERCEL_AI_GATEWAY_API_KEY ``` Start your LiteLLM Proxy server: diff --git a/docs/my-website/docs/proxy/guardrails/onyx_security.md b/docs/my-website/docs/proxy/guardrails/onyx_security.md index 85b0ba9f83..d240902eb5 100644 --- a/docs/my-website/docs/proxy/guardrails/onyx_security.md +++ b/docs/my-website/docs/proxy/guardrails/onyx_security.md @@ -128,6 +128,7 @@ guardrails: mode: ["pre_call", "post_call", "during_call"] # Run at multiple stages api_key: os.environ/ONYX_API_KEY api_base: os.environ/ONYX_API_BASE + timeout: 10.0 # Optional, defaults to 10 seconds ``` ### Required Parameters @@ -137,6 +138,7 @@ guardrails: ### Optional Parameters - **`api_base`**: Onyx API base URL (defaults to `https://ai-guard.onyx.security`) +- **`timeout`**: Request timeout in seconds (defaults to `10.0`) ## Environment Variables @@ -145,4 +147,5 @@ You can set these environment variables instead of hardcoding values in your con ```shell export ONYX_API_KEY="your-api-key-here" export ONYX_API_BASE="https://ai-guard.onyx.security" # Optional +export ONYX_TIMEOUT=10 # Optional, timeout in seconds ``` diff --git a/docs/my-website/img/a2a_agent_spend.png b/docs/my-website/img/a2a_agent_spend.png new file mode 100644 index 0000000000..15ec769392 Binary files /dev/null and b/docs/my-website/img/a2a_agent_spend.png differ diff --git a/docs/my-website/img/a2a_trace_grouping.png b/docs/my-website/img/a2a_trace_grouping.png new file mode 100644 index 0000000000..05130420aa Binary files /dev/null and b/docs/my-website/img/a2a_trace_grouping.png differ diff --git a/docs/my-website/release_notes/v1.81.3-stable/index.md b/docs/my-website/release_notes/v1.81.3-stable/index.md new file mode 100644 index 0000000000..22b6f43dee --- /dev/null +++ b/docs/my-website/release_notes/v1.81.3-stable/index.md @@ -0,0 +1,423 @@ +--- +title: "v1.81.3-stable - Performance - 25% CPU Usage Reduction" +slug: "v1-81-3" +date: 2026-01-26T10:00:00 +authors: + - name: Krrish Dholakia + title: CEO, LiteLLM + url: https://www.linkedin.com/in/krish-d/ + image_url: https://pbs.twimg.com/profile_images/1298587542745358340/DZv3Oj-h_400x400.jpg + - name: Ishaan Jaff + title: CTO, LiteLLM + url: https://www.linkedin.com/in/reffajnaahsi/ + image_url: https://pbs.twimg.com/profile_images/1613813310264340481/lz54oEiB_400x400.jpg +hide_table_of_contents: false +--- + +import Image from '@theme/IdealImage'; +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +## Deploy this version + + + + +``` showLineNumbers title="docker run litellm" +docker run \ +-e STORE_MODEL_IN_DB=True \ +-p 4000:4000 \ +docker.litellm.ai/berriai/litellm:v1.81.3.rc.2 +``` + + + + + +``` showLineNumbers title="pip install litellm" +pip install litellm==1.81.3.rc.2 +``` + + + + +--- + +## New Models / Updated Models + +### New Model Support + +| Provider | Model | Context Window | Input ($/1M tokens) | Output ($/1M tokens) | Deprecation Date | +| -------- | ----- | -------------- | ------------------- | -------------------- | ---------------- | +| OpenAI | `gpt-audio`, `gpt-audio-2025-08-28` | 128K | $32/1M audio tokens, $2.5/1M text tokens | $64/1M audio tokens, $10/1M text tokens | - | +| OpenAI | `gpt-audio-mini`, `gpt-audio-mini-2025-08-28` | 128K | $10/1M audio tokens, $0.6/1M text tokens | $20/1M audio tokens, $2.4/1M text tokens | - | +| Deepinfra, Vertex AI, Google AI Studio, OpenRouter, Vercel AI Gateway | `gemini-2.0-flash-001`, `gemini-2.0-flash` | - | - | - | 2026-03-31 | +| Groq | `openai/gpt-oss-120b` | 131K | 0.075/1M cache read | 0.6/1M output tokens | - | +| Groq | `groq/openai/gpt-oss-20b` | 131K | 0.0375/1M cache read, $0.075/1M text tokens | 0.3/1M output tokens | - | +| Vertex AI | `gemini-2.5-computer-use-preview-10-2025` | 128K | $1.25 | $10 | - | +| Azure AI | `claude-haiku-4-5` | $1.25/1M cache read, $2/1M cache read above 1 hr, $0.1/1M text tokens | $5/1M output tokens | - | +| Azure AI | `claude-sonnet-4-5` | $3.75/1M cache read, $6/1M cache read above 1 hr, $3/1M text tokens | $15/1M output tokens | - | +| Azure AI | `claude-opus-4-5` | $6.25/1M cache read, $10/1M cache read above 1 hr, $0.5/1M text tokens | $25/1M output tokens | - | +| Azure AI | `claude-opus-4-1` | $18.75/1M cache read, $30/1M cache read above 1 hr, $1.5/1M text tokens | $75/1M output tokens | - | + +### Features + +- **[OpenAI](../../docs/providers/openai)** + - Add gpt-audio and gpt-audio-mini models to pricing - [PR #19509](https://github.com/BerriAI/litellm/pull/19509) + - correct audio token costs for gpt-4o-audio-preview models - [PR #19500](https://github.com/BerriAI/litellm/pull/19500) + - Limit stop sequence as per openai spec (ensures JetBrains IDE compatibility) - [PR #19562](https://github.com/BerriAI/litellm/pull/19562) + +- **[VertexAI](../../docs/providers/vertex)** + - Docs - Google Workload Identity Federation (WIF) support - [PR #19320](https://github.com/BerriAI/litellm/pull/19320) + +- **[Agentcore](../../docs/providers/bedrock_agentcore)** + - Fixes streaming issues with AWS Bedrock AgentCore where responses would stop after the first chunk, particularly affecting OAuth-enabled agents - [PR #17141](https://github.com/BerriAI/litellm/pull/17141) + +- **[Chatgpt](../../docs/providers/chatgpt)** + - Adds support for calling chatgpt subscription via LiteLLM - [PR #19030](https://github.com/BerriAI/litellm/pull/19030) + - Adds responses API bridge support for chatgpt subscription provider - [PR #19030](https://github.com/BerriAI/litellm/pull/19030) + +- **[Bedrock](../../docs/providers/bedrock)** + - support for output format for bedrock invoke via v1/messages - [PR #19560](https://github.com/BerriAI/litellm/pull/19560) + +- **[Azure](../../docs/providers/azure/azure)** + - Add support for Azure OpenAI v1 API - [PR #19313](https://github.com/BerriAI/litellm/pull/19313) + - preserve content_policy_violation details for images (#19328) - [PR #19372](https://github.com/BerriAI/litellm/pull/19372) + - Support OpenAI-format nested tool definitions for Responses API - [PR #19526](https://github.com/BerriAI/litellm/pull/19526) + +- **Gemini([Vertex AI](../../docs/providers/vertex), [Google AI Studio](../../docs/providers/gemini))** + - use responseJsonSchema for Gemini 2.0+ models - [PR #19314](https://github.com/BerriAI/litellm/pull/19314) + +- **[Volcengine](../../docs/providers/volcano)** + - Support Volcengine responses api - [PR #18508](https://github.com/BerriAI/litellm/pull/18508) + +- **[Anthropic](../../docs/providers/anthropic)** + - Add Support for calling Claude Code Max subscriptions via LiteLLM - [PR #19453](https://github.com/BerriAI/litellm/pull/19453) + - Add Structured output for /v1/messages with Anthropic API, Azure Anthropic API, Bedrock Converse - [PR #19545](https://github.com/BerriAI/litellm/pull/19545) + +- **[Brave Search](../../docs/search/brave)** + - New Search provider - [PR #19433](https://github.com/BerriAI/litellm/pull/19433) + +- **Sarvam ai** + - Add support for new sarvam models - [PR #19479](https://github.com/BerriAI/litellm/pull/19479) + +- **[GMI](../../docs/providers/gmi)** + - add GMI Cloud provider support - [PR #19376](https://github.com/BerriAI/litellm/pull/19376) + + +### Bug Fixes + +- **[Anthropic](../../docs/providers/anthropic)** + - Fix anthropic-beta sent client side being overridden instead of appended to - [PR #19343](https://github.com/BerriAI/litellm/pull/19343) + - Filter out unsupported fields from JSON schema for Anthropic's output_format API - [PR #19482](https://github.com/BerriAI/litellm/pull/19482) + +- **[Bedrock](../../docs/providers/bedrock)** + - Expose stability models via /image_edits endpoint and ensure proper request transformation - [PR #19323](https://github.com/BerriAI/litellm/pull/19323) + - Claude Code x Bedrock Invoke fails with advanced-tool-use-2025-11-20 - [PR #19373](https://github.com/BerriAI/litellm/pull/19373) + - deduplicate tool calls in assistant history - [PR #19324](https://github.com/BerriAI/litellm/pull/19324) + - fix: correct us.anthropic.claude-opus-4-5 In-region pricing - [PR #19310](https://github.com/BerriAI/litellm/pull/19310) + - Fix request validation errors when using Claude 4 via bedrock invoke - [PR #19381](https://github.com/BerriAI/litellm/pull/19381) + - Handle thinking with tool calls for Claude 4 models - [PR #19506](https://github.com/BerriAI/litellm/pull/19506) + - correct streaming choice index for tool calls - [PR #19506](https://github.com/BerriAI/litellm/pull/19506) + +- **[Ollama](../../docs/providers/ollama)** + - Fix tool call errors due with improved message extraction - [PR #19369](https://github.com/BerriAI/litellm/pull/19369) + +- **[VertexAI](../../docs/providers/vertex)** + - Removed optional vertex_count_tokens_location param before request is sent to vertex - [PR #19359](https://github.com/BerriAI/litellm/pull/19359) + +- **Gemini([Vertex AI](../../docs/providers/vertex), [Google AI Studio](../../docs/providers/gemini))** + - Supports setting media_resolution and fps parameters on each video file, when using Gemini video understanding - [PR #19273](https://github.com/BerriAI/litellm/pull/19273) + - handle reasoning_effort as dict from OpenAI Agents SDK - [PR #19419](https://github.com/BerriAI/litellm/pull/19419) + - add file content support in tool results - [PR #19416](https://github.com/BerriAI/litellm/pull/19416) + +- **[Azure](../../docs/providers/azure_ai)** + - Fix Azure AI costs for Anthropic models - [PR #19530](https://github.com/BerriAI/litellm/pull/19530) + +- **[Giga Chat](../../docs/providers/gigachat)** + - Add tool choice mapping - [PR #19645](https://github.com/BerriAI/litellm/pull/19645) +--- + +## AI API Endpoints (LLMs, MCP, Agents) + +### Features + +- **[Files API](../../docs/files_endpoints)** + - Add managed files support when load_balancing is True - [PR #19338](https://github.com/BerriAI/litellm/pull/19338) + +- **[Claude Plugin Marketplace](../../docs/tutorials/claude_code_plugin_marketplace)** + - Add self hosted Claude Code Plugin Marketplace - [PR #19378](https://github.com/BerriAI/litellm/pull/19378) + +- **[MCP](../../docs/mcp)** + - Add MCP Protocol version 2025-11-25 support - [PR #19379](https://github.com/BerriAI/litellm/pull/19379) + - Log MCP tool calls and list tools in the LiteLLM Spend Logs table for easier debugging - [PR #19469](https://github.com/BerriAI/litellm/pull/19469) + +- **[Vertex AI](../../docs/providers/vertex)** + - Ensure only anthropic betas are forwarded down to LLM API (by default) - [PR #19542](https://github.com/BerriAI/litellm/pull/19542) + - Allow overriding to support forwarding incoming headers are forwarded down to target - [PR #19524](https://github.com/BerriAI/litellm/pull/19524) + +- **[Chat/Completions](../../docs/completion/input)** + - Add MCP tools response to chat completions - [PR #19552](https://github.com/BerriAI/litellm/pull/19552) + - Add custom vertex ai finish reasons to the output - [PR #19558](https://github.com/BerriAI/litellm/pull/19558) + - Return MCP execution in /chat/completions before model output during streaming - [PR #19623](https://github.com/BerriAI/litellm/pull/19623) + +### Bugs + +- **[Responses API](../../docs/response_api)** + - Fix duplicate messages during MCP streaming tool execution - [PR #19317](https://github.com/BerriAI/litellm/pull/19317) + - Fix pickle error when using OpenAI's Responses API with stream=True and tool_choice of type allowed_tools (an OpenAI-native parameter) - [PR #17205](https://github.com/BerriAI/litellm/pull/17205) + - stream tool call events for non-openai models - [PR #19368](https://github.com/BerriAI/litellm/pull/19368) + - preserve tool output ordering for gemini in responses bridge - [PR #19360](https://github.com/BerriAI/litellm/pull/19360) + - Add ID caching to prevent ID mismatch text-start and text-delta - [PR #19390](https://github.com/BerriAI/litellm/pull/19390) + - Include output_item, reasoning_summary_Text_done and reasoning_summary_part_done events for non-openai models - [PR #19472](https://github.com/BerriAI/litellm/pull/19472) + +- **[Chat/Completions](../../docs/completion/input)** + - fix: drop_params not dropping prompt_cache_key for non-OpenAI providers - [PR #19346](https://github.com/BerriAI/litellm/pull/19346) + +- **[Realtime API](../../docs/realtime)** + - disable SSL for ws:// WebSocket connections - [PR #19345](https://github.com/BerriAI/litellm/pull/19345) + +- **[Generate Content](../../docs/generateContent)** + - Log actual user input when google genai/vertex endpoints are called client-side - [PR #19156](https://github.com/BerriAI/litellm/pull/19156) + +- **[/messages/count_tokens Anthropic Token Counting](../../docs/anthropic_count_tokens)** + - ensure it works for Anthropic, Azure AI Anthropic on AI Gateway - [PR #19432](https://github.com/BerriAI/litellm/pull/19432) + +- **[MCP](../../docs/mcp)** + - forward static_headers to MCP servers - [PR #19366](https://github.com/BerriAI/litellm/pull/19366) + +- **[Batch API](../../docs/batches)** + - Fix: generation config empty for batch - [PR #19556](https://github.com/BerriAI/litellm/pull/19556) + +- **[Pass Through Endpoints](../../docs/proxy/pass_through)** + - Always reupdate registry - [PR #19420](https://github.com/BerriAI/litellm/pull/19420) +--- + +## Management Endpoints / UI + +### Features + +- **Cost Estimator** + - Fix model dropdown - [PR #19529](https://github.com/BerriAI/litellm/pull/19529) + +- **Claude Code Plugins** + - Allow Adding Claude Code Plugins via UI - [PR #19387](https://github.com/BerriAI/litellm/pull/19387) + +- **Guardrails** + - New Policy management UI - [PR #19668](https://github.com/BerriAI/litellm/pull/19668) + - Allow adding policies on Keys/Teams + Viewing on Info panels - [PR #19688](https://github.com/BerriAI/litellm/pull/19688) + +- **General** + - respects custom authentication header override - [PR #19276](https://github.com/BerriAI/litellm/pull/19276) + +- **Playground** + - Button to Fill Custom API Base - [PR #19440](https://github.com/BerriAI/litellm/pull/19440) + - display mcp output on the play ground - [PR #19553](https://github.com/BerriAI/litellm/pull/19553) + +- **Models** + - Paginate /v2/models/info - [PR #19521](https://github.com/BerriAI/litellm/pull/19521) + - All Model Tab Pagination - [PR #19525](https://github.com/BerriAI/litellm/pull/19525) + - Adding Optional scope Param to /models - [PR #19539](https://github.com/BerriAI/litellm/pull/19539) + - Model Search - [PR #19622](https://github.com/BerriAI/litellm/pull/19622) + - Filter by Model ID and Team ID - [PR #19713](https://github.com/BerriAI/litellm/pull/19713) + +- **MCP Servers** + - MCP Tools Tab Resetting to Overview - [PR #19468](https://github.com/BerriAI/litellm/pull/19468) + +- **Organizations** + - Prevent org admin from creating a new user with proxy_admin permissions - [PR #19296](https://github.com/BerriAI/litellm/pull/19296) + - Edit Page: Reusable Model Select - [PR #19601](https://github.com/BerriAI/litellm/pull/19601) + +- **Teams** + - Reusable Model Select - [PR #19543](https://github.com/BerriAI/litellm/pull/19543) + - [Fix] Team Update with Organization having All Proxy Models - [PR #19604](https://github.com/BerriAI/litellm/pull/19604) + +- **Logs** + - Include tool arguments in spend logs table - [PR #19640](https://github.com/BerriAI/litellm/pull/19640) + +- **Fallbacks / Loadbalancing** + - New fallbacks modal - [PR #19673](https://github.com/BerriAI/litellm/pull/19673) + - Set fallbacks/loadbalancing by team/key - [PR #19686](https://github.com/BerriAI/litellm/pull/19686) + +### Bugs + +- **Playground** + - increase model selector width in playground Compare view - [PR #19423](https://github.com/BerriAI/litellm/pull/19423) + +- **Virtual Keys** + - Sorting Shows Incorrect Entries - [PR #19534](https://github.com/BerriAI/litellm/pull/19534) + +- **General** + - UI 404 error when SERVER_ROOT_PATH is set - [PR #19467](https://github.com/BerriAI/litellm/pull/19467) + - Redirect to ui/login on expired JWT - [PR #19687](https://github.com/BerriAI/litellm/pull/19687) + +- **SSO** + - Fix SSO user roles not updating for existing users - [PR #19621](https://github.com/BerriAI/litellm/pull/19621) + +- **Guardrails** + - ensure guardrail patterns persist on edit and mode toggle - [PR #19265](https://github.com/BerriAI/litellm/pull/19265) +--- + +## AI Integrations + +### Logging + +- **General Logging** + - prevent printing duplicate StandardLoggingPayload logs - [PR #19325](https://github.com/BerriAI/litellm/pull/19325) + - Fix: log duplication when json_logs is enabled - [PR #19705](https://github.com/BerriAI/litellm/pull/19705) +- **Langfuse OTEL** + - ignore service logs and fix callback shadowing - [PR #19298](https://github.com/BerriAI/litellm/pull/19298) +- **Langfuse** + - Send litellm_trace_id - [PR #19528](https://github.com/BerriAI/litellm/pull/19528) + - Add Langfuse mock mode for testing without API calls - [PR #19676](https://github.com/BerriAI/litellm/pull/19676) +- **GCS Bucket** + - prevent unbounded queue growth due to slow API calls - [PR #19297](https://github.com/BerriAI/litellm/pull/19297) + - Add GCS mock mode for testing without API calls - [PR #19683](https://github.com/BerriAI/litellm/pull/19683) +- **Responses API Logging** + - Fix pydantic serialization error - [PR #19486](https://github.com/BerriAI/litellm/pull/19486) +- **Arize Phoenix** + - add openinference span kinds to arize phoenix - [PR #19267](https://github.com/BerriAI/litellm/pull/19267) +- **Prometheus** + - Added new prometheus metrics for user count and team count - [PR #19520](https://github.com/BerriAI/litellm/pull/19520) + +### Guardrails + +- **Bedrock Guardrails** + - Ensure post_call guardrail checks input+output - [PR #19151](https://github.com/BerriAI/litellm/pull/19151) +- **Prompt Security** + - fixing prompt-security's guardrail implementation - [PR #19374](https://github.com/BerriAI/litellm/pull/19374) +- **Presidio** + - Fixes crash in Presidio Guardrail when running in background threads (logging_hook) - [PR #19714](https://github.com/BerriAI/litellm/pull/19714) +- **Pillar Security** + - Migrate Pillar Security to Generic Guardrail API - [PR #19364](https://github.com/BerriAI/litellm/pull/19364) +- **Policy Engine** + - New LiteLLM Policy engine - create policies to manage guardrails, conditions - permissions per Key, Team - [PR #19612](https://github.com/BerriAI/litellm/pull/19612) +- **General** + - add case-insensitive support for guardrail mode and actions - [PR #19480](https://github.com/BerriAI/litellm/pull/19480) + +### Prompt Management + +- **General** + - fix prompt info lookup and delete using correct IDs - [PR #19358](https://github.com/BerriAI/litellm/pull/19358) + +### Secret Manager + +- **AWS Secret Manager** + - ensure auto-rotation updates existing AWS secret instead of creating new one - [PR #19455](https://github.com/BerriAI/litellm/pull/19455) +- **Hashicorp Vault** + - Ensure key rotations work with Vault - [PR #19634](https://github.com/BerriAI/litellm/pull/19634) + +--- + +## Spend Tracking, Budgets and Rate Limiting + +- **Pricing Updates** + - Add openai/dall-e base pricing entries - [PR #19133](https://github.com/BerriAI/litellm/pull/19133) + - Add `input_cost_per_video_per_second` in ModelInfoBase - [PR #19398](https://github.com/BerriAI/litellm/pull/19398) + +--- + +## Performance / Loadbalancing / Reliability improvements + + +- **General** + - Fix date overflow/division by zero in proxy utils - [PR #19527](https://github.com/BerriAI/litellm/pull/19527) + - Fix in-flight request termination on SIGTERM when health-check runs in a separate process - [PR #19427](https://github.com/BerriAI/litellm/pull/19427) + - Fix Pass through routes to work with server root path - [PR #19383](https://github.com/BerriAI/litellm/pull/19383) + - Fix logging error for stop iteration - [PR #19649](https://github.com/BerriAI/litellm/pull/19649) + - prevent retrying 4xx client errors - [PR #19275](https://github.com/BerriAI/litellm/pull/19275) + - add better error handling for misconfig on health check - [PR #19441](https://github.com/BerriAI/litellm/pull/19441) + +- **Router** + - Fix Azure RPM calculation formula - [PR #19513](https://github.com/BerriAI/litellm/pull/19513) + - Persist scheduler request queue to redis - [PR #19304](https://github.com/BerriAI/litellm/pull/19304) + - pass search_tools to Router during DB-triggered initialization - [PR #19388](https://github.com/BerriAI/litellm/pull/19388) + - Fixed PromptCachingCache to correctly handle messages where cache_control is a sibling key of string content - [PR #19266](https://github.com/BerriAI/litellm/pull/19266) + +- **Memory Leaks/OOM** + - prevent OOM with nested $defs in tool schemas - [PR #19112](https://github.com/BerriAI/litellm/pull/19112) + - fix: HTTP client memory leaks in Presidio, OpenAI, and Gemini - [PR #19190](https://github.com/BerriAI/litellm/pull/19190) + +- **Non root** + - fix logfile and pidfile of supervisor for non root environment - [PR #17267](https://github.com/BerriAI/litellm/pull/17267) + - resolve Read-only file system error in non-root images - [PR #19449](https://github.com/BerriAI/litellm/pull/19449) + +- **Dockerfile** + - Redis Semantic Caching - add missing redisvl dependency to requirements.txt - [PR #19417](https://github.com/BerriAI/litellm/pull/19417) + - Bump OTEL versions to support a2a dependency - resolves modulenotfounderror for Microsoft Agents by @Harshit28j in #18991 + +- **DB** + - Handle PostgreSQL cached plan errors during rolling deployments - [PR #19424](https://github.com/BerriAI/litellm/pull/19424) + +- **Timeouts** + - Fix: total timeout is not respected - [PR #19389](https://github.com/BerriAI/litellm/pull/19389) + +- **SDK** + - Field-Existence Checks to Type Classes to Prevent Attribute Errors - [PR #18321](https://github.com/BerriAI/litellm/pull/18321) + - add google-cloud-aiplatform as optional dependency with clear error message - [PR #19437](https://github.com/BerriAI/litellm/pull/19437) + - Make grpc dependency optional - [PR #19447](https://github.com/BerriAI/litellm/pull/19447) + - Add support for retry policies - [PR #19645](https://github.com/BerriAI/litellm/pull/19645) + +- **Performance** + - Cut chat_completion latency by ~21% by reducing pre-call processing time - [PR #19535](https://github.com/BerriAI/litellm/pull/19535) + - Optimize strip_trailing_slash with O(1) index check - [PR #19679](https://github.com/BerriAI/litellm/pull/19679) + - Optimize use_custom_pricing_for_model with set intersection - [PR #19677](https://github.com/BerriAI/litellm/pull/19677) + - perf: skip pattern_router.route() for non-wildcard models - [PR #19664](https://github.com/BerriAI/litellm/pull/19664) + - perf: Add LRU caching to get_model_info for faster cost lookups - [PR #19606](https://github.com/BerriAI/litellm/pull/19606) + +--- + +## General Proxy Improvements + +### Doc Improvements + - new tutorial for adding MCPs to Cursor via LiteLLM - [PR #19317](https://github.com/BerriAI/litellm/pull/19317) + - fix vertex_region to vertex_location in Vertex AI pass-through docs - [PR #19380](https://github.com/BerriAI/litellm/pull/19380) + - clarify Gemini and Vertex AI model prefix in json file - [PR #19443](https://github.com/BerriAI/litellm/pull/19443) + - update Claude Code integration guides - [PR #19415](https://github.com/BerriAI/litellm/pull/19415) + - adjust opencode tutorial - [PR #19605](https://github.com/BerriAI/litellm/pull/19605) + - add spend-queue-troubleshooting docs - [PR #19659](https://github.com/BerriAI/litellm/pull/19659) + - docs: add litellm-enterprise requirement for managed files - [PR #19689](https://github.com/BerriAI/litellm/pull/19689) + +### Helm + - Add support for keda in helm chart - [PR #19337](https://github.com/BerriAI/litellm/pull/19337) + - sync Helm chart version with LiteLLM release version - [PR #19438](https://github.com/BerriAI/litellm/pull/19438) + - Enable PreStop hook configuration in values.yaml - [PR #19613](https://github.com/BerriAI/litellm/pull/19613) + +### General + - Add health check scripts and parallel execution support - [PR #19295](https://github.com/BerriAI/litellm/pull/19295) + + +--- + +## New Contributors + + +* @dushyantzz made their first contribution in [PR #19158](https://github.com/BerriAI/litellm/pull/19158) +* @obod-mpw made their first contribution in [PR #19133](https://github.com/BerriAI/litellm/pull/19133) +* @msexxeta made their first contribution in [PR #19030](https://github.com/BerriAI/litellm/pull/19030) +* @rsicart made their first contribution in [PR #19337](https://github.com/BerriAI/litellm/pull/19337) +* @cluebbehusen made their first contribution in [PR #19311](https://github.com/BerriAI/litellm/pull/19311) +* @Lucky-Lodhi2004 made their first contribution in [PR #19315](https://github.com/BerriAI/litellm/pull/19315) +* @binbandit made their first contribution in [PR #19324](https://github.com/BerriAI/litellm/pull/19324) +* @flex-myeonghyeon made their first contribution in [PR #19381](https://github.com/BerriAI/litellm/pull/19381) +* @Lrakotoson made their first contribution in [PR #18321](https://github.com/BerriAI/litellm/pull/18321) +* @bensi94 made their first contribution in [PR #18787](https://github.com/BerriAI/litellm/pull/18787) +* @victorigualada made their first contribution in [PR #19368](https://github.com/BerriAI/litellm/pull/19368) +* @VedantMadane made their first contribution in #19266 +* @stiyyagura0901 made their first contribution in #19276 +* @kamilio made their first contribution in [PR #19447](https://github.com/BerriAI/litellm/pull/19447) +* @jonathansampson made their first contribution in [PR #19433](https://github.com/BerriAI/litellm/pull/19433) +* @rynecarbone made their first contribution in [PR #19416](https://github.com/BerriAI/litellm/pull/19416) +* @jayy-77 made their first contribution in #19366 +* @davida-ps made their first contribution in [PR #19374](https://github.com/BerriAI/litellm/pull/19374) +* @joaodinissf made their first contribution in [PR #19506](https://github.com/BerriAI/litellm/pull/19506) +* @ecao310 made their first contribution in [PR #19520](https://github.com/BerriAI/litellm/pull/19520) +* @mpcusack-altos made their first contribution in [PR #19577](https://github.com/BerriAI/litellm/pull/19577) +* @milan-berri made their first contribution in [PR #19602](https://github.com/BerriAI/litellm/pull/19602) +* @xqe2011 made their first contribution in #19621 + +--- + +## Full Changelog + +**[View complete changelog on GitHub](https://github.com/BerriAI/litellm/releases/tag/v1.81.3.rc)** diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 2d36dbeacd..ae95faede2 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -9,7 +9,7 @@ import datetime from typing import TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, Optional, Union import litellm -from litellm._logging import verbose_logger +from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator from litellm.a2a_protocol.utils import A2ARequestUtils from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT @@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.agents import LiteLLMSendMessageResponse from litellm.utils import client +import uuid if TYPE_CHECKING: from a2a.client import A2AClient as A2AClientType @@ -225,7 +226,11 @@ async def asend_message( raise ValueError( "Either a2a_client or api_base is required for standard A2A flow" ) - a2a_client = await create_a2a_client(base_url=api_base) + trace_id = str(uuid.uuid4()) + extra_headers = {"X-LiteLLM-Trace-Id": trace_id} + if agent_id: + extra_headers["X-LiteLLM-Agent-Id"] = agent_id + a2a_client = await create_a2a_client(base_url=api_base, extra_headers=extra_headers) # Type assertion: a2a_client is guaranteed to be non-None here assert a2a_client is not None @@ -490,6 +495,10 @@ async def create_a2a_client( ) httpx_client = http_handler.client + if extra_headers: + httpx_client.headers.update(extra_headers) + verbose_proxy_logger.debug(f"A2A client created with extra_headers={extra_headers}") + # Resolve agent card resolver = A2ACardResolver( httpx_client=httpx_client, diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 3cbb02259b..e2de3cd502 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -4,13 +4,17 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers. import asyncio import base64 -from typing import Any, Awaitable, Callable, Dict, List, Optional, TypeVar, Union +from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple, TypeVar, Union import httpx from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client -from mcp.client.streamable_http import streamable_http_client + +try: + from mcp.client.streamable_http import streamable_http_client # type: ignore +except ImportError: + streamable_http_client = None from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import ( @@ -76,104 +80,100 @@ class MCPClient: def _create_transport_context( self, - ) -> tuple[Any, Optional[httpx.AsyncClient]]: - """Create the appropriate transport context based on transport type.""" + ) -> Tuple[Any, Optional[httpx.AsyncClient]]: + """ + Create the appropriate transport context based on transport type. + + Returns: + Tuple of (transport_context, http_client). + http_client is only set for HTTP transport and needs cleanup. + """ http_client: Optional[httpx.AsyncClient] = None if self.transport_type == MCPTransport.stdio: if not self.stdio_config: raise ValueError("stdio_config is required for stdio transport") - server_params = StdioServerParameters( command=self.stdio_config.get("command", ""), args=self.stdio_config.get("args", []), env=self.stdio_config.get("env", {}), ) - transport_ctx = stdio_client(server_params) - elif self.transport_type == MCPTransport.sse: + return stdio_client(server_params), None + + if self.transport_type == MCPTransport.sse: headers = self._get_auth_headers() httpx_client_factory = self._create_httpx_client_factory() - transport_ctx = sse_client( + return sse_client( url=self.server_url, timeout=self.timeout, headers=headers, httpx_client_factory=httpx_client_factory, - ) - else: - headers = self._get_auth_headers() - httpx_client_factory = self._create_httpx_client_factory() - verbose_logger.debug( - "litellm headers for streamable_http_client: %s", headers - ) - http_client = httpx_client_factory( - headers=headers, - timeout=httpx.Timeout(self.timeout), - ) - transport_ctx = streamable_http_client( - url=self.server_url, - http_client=http_client, - ) - - if transport_ctx is None: - raise RuntimeError("Failed to create transport context") + ), None + # HTTP transport (default) + headers = self._get_auth_headers() + httpx_client_factory = self._create_httpx_client_factory() + verbose_logger.debug( + "litellm headers for streamable_http_client: %s", headers + ) + http_client = httpx_client_factory( + headers=headers, + timeout=httpx.Timeout(self.timeout), + ) + transport_ctx = streamable_http_client( + url=self.server_url, + http_client=http_client, + ) return transport_ctx, http_client + async def _execute_session_operation( + self, + transport_ctx: Any, + operation: Callable[[ClientSession], Awaitable[TSessionResult]], + ) -> TSessionResult: + """ + Execute an operation within a transport and session context. + + Handles entering/exiting contexts and running the operation. + """ + transport = await transport_ctx.__aenter__() + try: + read_stream, write_stream = transport[0], transport[1] + session_ctx = ClientSession(read_stream, write_stream) + session = await session_ctx.__aenter__() + try: + await session.initialize() + return await operation(session) + finally: + try: + await session_ctx.__aexit__(None, None, None) + except BaseException as e: + verbose_logger.debug(f"Error during session context exit: {e}") + finally: + try: + await transport_ctx.__aexit__(None, None, None) + except BaseException as e: + verbose_logger.debug(f"Error during transport context exit: {e}") + async def run_with_session( self, operation: Callable[[ClientSession], Awaitable[TSessionResult]] ) -> TSessionResult: """Open a session, run the provided coroutine, and clean up.""" - transport_ctx = None http_client: Optional[httpx.AsyncClient] = None - session_ctx = None - try: transport_ctx, http_client = self._create_transport_context() - - # Enter transport context - transport = await transport_ctx.__aenter__() - try: - read_stream, write_stream = transport[0], transport[1] - session_ctx = ClientSession(read_stream, write_stream) - - # Enter session context - session = await session_ctx.__aenter__() - try: - await session.initialize() - result = await operation(session) - return result - finally: - # Ensure session context is properly exited - if session_ctx is not None: - try: - await session_ctx.__aexit__(None, None, None) - except Exception as e: - verbose_logger.debug( - f"Error during session context exit: {e}" - ) - finally: - # Ensure transport context is properly exited - if transport_ctx is not None: - try: - await transport_ctx.__aexit__(None, None, None) - except Exception as e: - verbose_logger.debug( - f"Error during transport context exit: {e}" - ) + return await self._execute_session_operation(transport_ctx, operation) except Exception: verbose_logger.warning( "MCP client run_with_session failed for %s", self.server_url or "stdio" ) raise finally: - # Always clean up http_client if it was created if http_client is not None: try: await http_client.aclose() - except Exception as e: - verbose_logger.debug( - f"Error during http_client cleanup: {e}" - ) + except BaseException as e: + verbose_logger.debug(f"Error during http_client cleanup: {e}") def update_auth_value(self, mcp_auth_value: Union[str, Dict[str, str]]): """ diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index 6b30b6b736..6a003b8c49 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -83,6 +83,33 @@ }, "description": "Datadog Logging Integration" }, + { + "id": "datadog_cost_management", + "displayName": "Datadog Cost Management", + "logo": "datadog.png", + "supports_key_team_logging": false, + "dynamic_params": { + "dd_api_key": { + "type": "password", + "ui_name": "API Key", + "description": "Datadog API key for authentication", + "required": true + }, + "dd_app_key": { + "type": "password", + "ui_name": "App Key", + "description": "Datadog Application Key for Cloud Cost Management", + "required": true + }, + "dd_site": { + "type": "text", + "ui_name": "Site", + "description": "Datadog site URL (e.g., us5.datadoghq.com)", + "required": true + } + }, + "description": "Datadog Cloud Cost Management Integration" + }, { "id": "lago", "displayName": "Lago", @@ -407,4 +434,4 @@ }, "description": "SQS Queue (AWS) Logging Integration" } -] +] \ No newline at end of file diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 6a76b57e7f..a5bb530fc5 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -516,7 +516,9 @@ class CustomGuardrail(CustomLogger): from litellm.types.utils import GuardrailMode # Use event_type if provided, otherwise fall back to self.event_hook - guardrail_mode: Union[GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks]] + guardrail_mode: Union[ + GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks] + ] if event_type is not None: guardrail_mode = event_type elif isinstance(self.event_hook, Mode): @@ -524,11 +526,21 @@ class CustomGuardrail(CustomLogger): else: guardrail_mode = self.event_hook # type: ignore[assignment] + from litellm.litellm_core_utils.core_helpers import ( + filter_exceptions_from_params, + ) + + # Sanitize the response to ensure it's JSON serializable and free of circular refs + # This prevents RecursionErrors in downstream loggers (Langfuse, Datadog, etc.) + clean_guardrail_response = filter_exceptions_from_params( + guardrail_json_response + ) + slg = StandardLoggingGuardrailInformation( guardrail_name=self.guardrail_name, guardrail_provider=guardrail_provider, guardrail_mode=guardrail_mode, - guardrail_response=guardrail_json_response, + guardrail_response=clean_guardrail_response, guardrail_status=guardrail_status, start_time=start_time, end_time=end_time, diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 503e8d8c87..735d1005d2 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -32,6 +32,7 @@ from litellm.integrations.datadog.datadog_handler import ( get_datadog_service, get_datadog_source, get_datadog_tags, + get_datadog_base_url_from_env, ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.llms.custom_httpx.http_handler import ( @@ -100,7 +101,9 @@ class DataDogLogger( self._configure_dd_direct_api() # Optional override for testing - self._apply_dd_base_url_override() + dd_base_url = get_datadog_base_url_from_env() + if dd_base_url: + self.intake_url = f"{dd_base_url}/api/v2/logs" self.sync_client = _get_httpx_client() asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() @@ -159,18 +162,6 @@ class DataDogLogger( self.DD_API_KEY = os.getenv("DD_API_KEY") self.intake_url = f"https://http-intake.logs.{os.getenv('DD_SITE')}/api/v2/logs" - def _apply_dd_base_url_override(self) -> None: - """ - Apply base URL override for testing purposes - """ - dd_base_url: Optional[str] = ( - os.getenv("_DATADOG_BASE_URL") - or os.getenv("DATADOG_BASE_URL") - or os.getenv("DD_BASE_URL") - ) - if dd_base_url is not None: - self.intake_url = f"{dd_base_url}/api/v2/logs" - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ Async Log success events to Datadog diff --git a/litellm/integrations/datadog/datadog_cost_management.py b/litellm/integrations/datadog/datadog_cost_management.py new file mode 100644 index 0000000000..2eb94b59dd --- /dev/null +++ b/litellm/integrations/datadog/datadog_cost_management.py @@ -0,0 +1,204 @@ +import asyncio +import os +import time +from datetime import datetime +from typing import Dict, List, Optional, Tuple + +from litellm._logging import verbose_logger +from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.integrations.datadog_cost_management import ( + DatadogFOCUSCostEntry, +) +from litellm.types.utils import StandardLoggingPayload + + +class DatadogCostManagementLogger(CustomBatchLogger): + def __init__(self, **kwargs): + self.dd_api_key = os.getenv("DD_API_KEY") + self.dd_app_key = os.getenv("DD_APP_KEY") + self.dd_site = os.getenv("DD_SITE", "datadoghq.com") + + if not self.dd_api_key or not self.dd_app_key: + verbose_logger.warning( + "Datadog Cost Management: DD_API_KEY and DD_APP_KEY are required. Integration will not work." + ) + + self.upload_url = f"https://api.{self.dd_site}/api/v2/cost/custom_costs" + + self.async_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) + + # Initialize lock and start periodic flush task + self.flush_lock = asyncio.Lock() + asyncio.create_task(self.periodic_flush()) + + # Check if flush_lock is already in kwargs to avoid double passing (unlikely but safe) + if "flush_lock" not in kwargs: + kwargs["flush_lock"] = self.flush_lock + + super().__init__(**kwargs) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + try: + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object", None + ) + + if standard_logging_object is None: + return + + # Only log if there is a cost associated + if standard_logging_object.get("response_cost", 0) > 0: + self.log_queue.append(standard_logging_object) + + if len(self.log_queue) >= self.batch_size: + await self.async_send_batch() + + except Exception as e: + verbose_logger.exception( + f"Datadog Cost Management: Error in async_log_success_event: {str(e)}" + ) + + async def async_send_batch(self): + if not self.log_queue: + return + + try: + # Aggregate costs from the batch + aggregated_entries = self._aggregate_costs(self.log_queue) + + if not aggregated_entries: + return + + # Send to Datadog + await self._upload_to_datadog(aggregated_entries) + + # Clear queue only on success (or if we decide to drop on failure) + # CustomBatchLogger clears queue in flush_queue, so we just process here + + except Exception as e: + verbose_logger.exception( + f"Datadog Cost Management: Error in async_send_batch: {str(e)}" + ) + + def _aggregate_costs( + self, logs: List[StandardLoggingPayload] + ) -> List[DatadogFOCUSCostEntry]: + """ + Aggregates costs by Provider, Model, and Date. + Returns a list of DatadogFOCUSCostEntry. + """ + aggregator: Dict[Tuple[str, str, str, Tuple[Tuple[str, str], ...]], DatadogFOCUSCostEntry] = {} + + for log in logs: + try: + # Extract keys for aggregation + provider = log.get("custom_llm_provider") or "unknown" + model = log.get("model") or "unknown" + cost = log.get("response_cost", 0) + + if cost == 0: + continue + + # Get date strings (FOCUS format requires specific keys, but for aggregation we group by Day) + # UTC date + # We interpret "ChargePeriod" as the day of the request. + ts = log.get("startTime") or time.time() + dt = datetime.fromtimestamp(ts) + date_str = dt.strftime("%Y-%m-%d") + + # ChargePeriodStart and End + # If we want daily granularity, end date is usually same day or next day? + # Datadog Custom Costs usually expects periods. + # "ChargePeriodStart": "2023-01-01", "ChargePeriodEnd": "2023-12-31" in example. + # If we send daily, we can say Start=Date, End=Date. + + # Grouping Key: Provider + Model + Date + Tags? + # For simplicity, let's aggregate by Provider + Model + Date first. + # If we handle tags, we need to include them in the key. + + tags = self._extract_tags(log) + tags_key = tuple(sorted(tags.items())) if tags else () + + key = (provider, model, date_str, tags_key) + + if key not in aggregator: + aggregator[key] = { + "ProviderName": provider, + "ChargeDescription": f"LLM Usage for {model}", + "ChargePeriodStart": date_str, + "ChargePeriodEnd": date_str, + "BilledCost": 0.0, + "BillingCurrency": "USD", + "Tags": tags if tags else None, + } + + aggregator[key]["BilledCost"] += cost + + except Exception as e: + verbose_logger.warning( + f"Error processing log for cost aggregation: {e}" + ) + continue + + return list(aggregator.values()) + + def _extract_tags(self, log: StandardLoggingPayload) -> Dict[str, str]: + from litellm.integrations.datadog.datadog_handler import ( + get_datadog_env, + get_datadog_hostname, + get_datadog_pod_name, + get_datadog_service, + ) + + tags = { + "env": get_datadog_env(), + "service": get_datadog_service(), + "host": get_datadog_hostname(), + "pod_name": get_datadog_pod_name(), + } + + # Add metadata as tags + metadata = log.get("metadata", {}) + if metadata: + # Add user info + if "user_api_key_alias" in metadata: + tags["user"] = str(metadata["user_api_key_alias"]) + if "user_api_key_team_alias" in metadata: + tags["team"] = str(metadata["user_api_key_team_alias"]) + # model_group is not in StandardLoggingMetadata TypedDict, so we need to access it via dict.get() + model_group = metadata.get("model_group") # type: ignore[misc] + if model_group: + tags["model_group"] = str(model_group) + + return tags + + async def _upload_to_datadog(self, payload: List[Dict]): + if not self.dd_api_key or not self.dd_app_key: + return + + headers = { + "Content-Type": "application/json", + "DD-API-KEY": self.dd_api_key, + "DD-APPLICATION-KEY": self.dd_app_key, + } + + # The API endpoint expects a list of objects directly in the body (file content behavior) + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + + data_json = safe_dumps(payload) + + response = await self.async_client.put( + self.upload_url, content=data_json, headers=headers + ) + + response.raise_for_status() + + verbose_logger.debug( + f"Datadog Cost Management: Uploaded {len(payload)} cost entries. Status: {response.status_code}" + ) diff --git a/litellm/integrations/datadog/datadog_handler.py b/litellm/integrations/datadog/datadog_handler.py index 26fab77759..e2f30f2f61 100644 --- a/litellm/integrations/datadog/datadog_handler.py +++ b/litellm/integrations/datadog/datadog_handler.py @@ -20,6 +20,14 @@ def get_datadog_hostname() -> str: return os.getenv("HOSTNAME", "") +def get_datadog_base_url_from_env() -> Optional[str]: + """ + Get base URL override from common DD_BASE_URL env var. + This is useful for testing or custom endpoints. + """ + return os.getenv("DD_BASE_URL") + + def get_datadog_env() -> str: return os.getenv("DD_ENV", "unknown") diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 6ffdbc0a00..4f6a5b339a 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -21,6 +21,7 @@ from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.integrations.datadog.datadog_handler import ( get_datadog_service, get_datadog_tags, + get_datadog_base_url_from_env, ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -43,24 +44,22 @@ class DataDogLLMObsLogger(CustomBatchLogger): def __init__(self, **kwargs): try: verbose_logger.debug("DataDogLLMObs: Initializing logger") - if os.getenv("DD_API_KEY", None) is None: - raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>'") - if os.getenv("DD_SITE", None) is None: - raise Exception( - "DD_SITE is not set, set 'DD_SITE=<>', example sit = `us5.datadoghq.com`" - ) + # Configure DataDog endpoint (Agent or Direct API) + # Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST + dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST") self.async_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback ) self.DD_API_KEY = os.getenv("DD_API_KEY") - self.DD_SITE = os.getenv("DD_SITE") - self.intake_url = ( - f"https://api.{self.DD_SITE}/api/intake/llm-obs/v1/trace/spans" - ) - # testing base url - dd_base_url = os.getenv("DD_BASE_URL") + if dd_agent_host: + self._configure_dd_agent(dd_agent_host=dd_agent_host) + else: + self._configure_dd_direct_api() + + # Optional override for testing + dd_base_url = get_datadog_base_url_from_env() if dd_base_url: self.intake_url = f"{dd_base_url}/api/intake/llm-obs/v1/trace/spans" @@ -78,6 +77,38 @@ class DataDogLLMObsLogger(CustomBatchLogger): verbose_logger.exception(f"DataDogLLMObs: Error initializing - {str(e)}") raise e + def _configure_dd_agent(self, dd_agent_host: str): + """ + Configure the Datadog logger to send traces to the Agent. + """ + # When using the Agent, LLM Observability Intake does NOT require the API Key + # Reference: https://docs.datadoghq.com/llm_observability/setup/sdk/#agent-setup + + # Use specific port for LLM Obs (Trace Agent) to avoid conflict with Logs Agent (10518) + agent_port = os.getenv("LITELLM_DD_LLM_OBS_PORT", "8126") + self.DD_SITE = "localhost" # Not used for URL construction in agent mode + self.intake_url = ( + f"http://{dd_agent_host}:{agent_port}/api/intake/llm-obs/v1/trace/spans" + ) + verbose_logger.debug(f"DataDogLLMObs: Using DD Agent at {self.intake_url}") + + def _configure_dd_direct_api(self): + """ + Configure the Datadog logger to send traces directly to the Datadog API. + """ + if not self.DD_API_KEY: + raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>'") + + self.DD_SITE = os.getenv("DD_SITE") + if not self.DD_SITE: + raise Exception( + "DD_SITE is not set, set 'DD_SITE=<>', example site = `us5.datadoghq.com`" + ) + + self.intake_url = ( + f"https://api.{self.DD_SITE}/api/intake/llm-obs/v1/trace/spans" + ) + def _get_datadog_llm_obs_params(self) -> Dict: """ Get the datadog_llm_observability_params from litellm.datadog_llm_observability_params @@ -164,13 +195,14 @@ class DataDogLLMObsLogger(CustomBatchLogger): json_payload = safe_dumps(payload) + headers = {"Content-Type": "application/json"} + if self.DD_API_KEY: + headers["DD-API-KEY"] = self.DD_API_KEY + response = await self.async_client.post( url=self.intake_url, content=json_payload, - headers={ - "DD-API-KEY": self.DD_API_KEY, - "Content-Type": "application/json", - }, + headers=headers, ) if response.status_code != 202: diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 46ada3c393..7bf97665fd 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -23,6 +23,7 @@ from litellm.constants import MAX_LANGFUSE_INITIALIZED_CLIENTS from litellm.litellm_core_utils.core_helpers import ( safe_deep_copy, reconstruct_model_name, + filter_exceptions_from_params, ) from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info from litellm.integrations.langfuse.langfuse_mock_client import ( @@ -75,9 +76,8 @@ def _extract_cache_read_input_tokens(usage_obj) -> int: # Check prompt_tokens_details.cached_tokens (used by Gemini and other providers) if hasattr(usage_obj, "prompt_tokens_details"): prompt_tokens_details = getattr(usage_obj, "prompt_tokens_details", None) - if ( - prompt_tokens_details is not None - and hasattr(prompt_tokens_details, "cached_tokens") + if prompt_tokens_details is not None and hasattr( + prompt_tokens_details, "cached_tokens" ): cached_tokens = getattr(prompt_tokens_details, "cached_tokens", None) if ( @@ -540,7 +540,6 @@ class LangFuseLogger: verbose_logger.debug("Langfuse Layer Logging - logging to langfuse v2") try: - metadata = metadata or {} standard_logging_object: Optional[StandardLoggingPayload] = cast( Optional[StandardLoggingPayload], kwargs.get("standard_logging_object", None), @@ -706,9 +705,10 @@ class LangFuseLogger: clean_metadata["litellm_response_cost"] = cost if standard_logging_object is not None: - clean_metadata["hidden_params"] = standard_logging_object[ - "hidden_params" - ] + hidden_params = standard_logging_object.get("hidden_params", {}) + clean_metadata["hidden_params"] = filter_exceptions_from_params( + hidden_params + ) if ( litellm.langfuse_default_tags is not None diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index af569b9ed6..2c897cb069 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -229,14 +229,18 @@ class PrometheusLogger(CustomLogger): self.litellm_remaining_api_key_requests_for_model = self._gauge_factory( "litellm_remaining_api_key_requests_for_model", "Remaining Requests API Key can make for model (model based rpm limit on key)", - labelnames=["hashed_api_key", "api_key_alias", "model"], + labelnames=self.get_labels_for_metric( + "litellm_remaining_api_key_requests_for_model" + ), ) # Remaining MODEL TPM limit for API Key self.litellm_remaining_api_key_tokens_for_model = self._gauge_factory( "litellm_remaining_api_key_tokens_for_model", "Remaining Tokens API Key can make for model (model based tpm limit on key)", - labelnames=["hashed_api_key", "api_key_alias", "model"], + labelnames=self.get_labels_for_metric( + "litellm_remaining_api_key_tokens_for_model" + ), ) ######################################## @@ -312,6 +316,18 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric("litellm_deployment_state"), ) + self.litellm_deployment_tpm_limit = self._gauge_factory( + "litellm_deployment_tpm_limit", + "Deployment TPM limit found in config", + labelnames=self.get_labels_for_metric("litellm_deployment_tpm_limit"), + ) + + self.litellm_deployment_rpm_limit = self._gauge_factory( + "litellm_deployment_rpm_limit", + "Deployment RPM limit found in config", + labelnames=self.get_labels_for_metric("litellm_deployment_rpm_limit"), + ) + self.litellm_deployment_cooled_down = self._counter_factory( "litellm_deployment_cooled_down", "LLM Deployment Analytics - Number of times a deployment has been cooled down by LiteLLM load balancing logic. exception_status is the status of the exception that caused the deployment to be cooled down", @@ -373,15 +389,9 @@ class PrometheusLogger(CustomLogger): self.litellm_llm_api_failed_requests_metric = self._counter_factory( name="litellm_llm_api_failed_requests_metric", documentation="deprecated - use litellm_proxy_failed_requests_metric", - labelnames=[ - "end_user", - "hashed_api_key", - "api_key_alias", - "model", - "team", - "team_alias", - "user", - ], + labelnames=self.get_labels_for_metric( + "litellm_llm_api_failed_requests_metric" + ), ) self.litellm_requests_metric = self._counter_factory( @@ -954,6 +964,8 @@ class PrometheusLogger(CustomLogger): route=standard_logging_payload["metadata"].get( "user_api_key_request_route" ), + client_ip=standard_logging_payload["metadata"].get("requester_ip_address"), + user_agent=standard_logging_payload["metadata"].get("user_agent"), ) if ( @@ -1011,6 +1023,7 @@ class PrometheusLogger(CustomLogger): user_api_key_alias=user_api_key_alias, kwargs=kwargs, metadata=_metadata, + model_id=enum_values.model_id, ) # set latency metrics @@ -1234,6 +1247,7 @@ class PrometheusLogger(CustomLogger): user_api_key_alias: Optional[str], kwargs: dict, metadata: dict, + model_id: Optional[str] = None, ): from litellm.proxy.common_utils.callback_utils import ( get_model_group_from_litellm_kwargs, @@ -1255,11 +1269,11 @@ class PrometheusLogger(CustomLogger): ) self.litellm_remaining_api_key_requests_for_model.labels( - user_api_key, user_api_key_alias, model_group + user_api_key, user_api_key_alias, model_group, model_id ).set(remaining_requests) self.litellm_remaining_api_key_tokens_for_model.labels( - user_api_key, user_api_key_alias, model_group + user_api_key, user_api_key_alias, model_group, model_id ).set(remaining_tokens) def _set_latency_metrics( @@ -1387,6 +1401,7 @@ class PrometheusLogger(CustomLogger): user_api_team, user_api_team_alias, user_id, + standard_logging_payload.get("model_id", ""), ).inc() self.set_llm_deployment_failure_metrics(kwargs) except Exception as e: @@ -1575,6 +1590,10 @@ class PrometheusLogger(CustomLogger): litellm_params=request_data, proxy_server_request=request_data.get("proxy_server_request", {}), ) + _metadata = request_data.get("metadata", {}) or {} + model_id = _metadata.get("model_info", {}).get("id") or request_data.get( + "model_info", {} + ).get("id") enum_values = UserAPIKeyLabelValues( end_user=user_api_key_dict.end_user_id, user=user_api_key_dict.user_id, @@ -1589,6 +1608,9 @@ class PrometheusLogger(CustomLogger): exception_class=self._get_exception_class_name(original_exception), tags=_tags, route=user_api_key_dict.request_route, + client_ip=_metadata.get("requester_ip_address"), + user_agent=_metadata.get("user_agent"), + model_id=model_id, ) _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric( @@ -1628,6 +1650,7 @@ class PrometheusLogger(CustomLogger): ): return + _metadata = data.get("metadata", {}) or {} enum_values = UserAPIKeyLabelValues( end_user=user_api_key_dict.end_user_id, hashed_api_key=user_api_key_dict.api_key, @@ -1643,6 +1666,8 @@ class PrometheusLogger(CustomLogger): litellm_params=data, proxy_server_request=data.get("proxy_server_request", {}), ), + client_ip=_metadata.get("requester_ip_address"), + user_agent=_metadata.get("user_agent"), ) _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric( @@ -1715,6 +1740,10 @@ class PrometheusLogger(CustomLogger): "user_api_key_team_alias" ], tags=standard_logging_payload.get("request_tags", []), + client_ip=standard_logging_payload["metadata"].get( + "requester_ip_address" + ), + user_agent=standard_logging_payload["metadata"].get("user_agent"), ) """ @@ -1752,6 +1781,49 @@ class PrometheusLogger(CustomLogger): ) ) + def _set_deployment_tpm_rpm_limit_metrics( + self, + model_info: dict, + litellm_params: dict, + litellm_model_name: Optional[str], + model_id: Optional[str], + api_base: Optional[str], + llm_provider: Optional[str], + ): + """ + Set the deployment TPM and RPM limits metrics + """ + tpm = model_info.get("tpm") or litellm_params.get("tpm") + rpm = model_info.get("rpm") or litellm_params.get("rpm") + + if tpm is not None: + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_deployment_tpm_limit" + ), + enum_values=UserAPIKeyLabelValues( + litellm_model_name=litellm_model_name, + model_id=model_id, + api_base=api_base, + api_provider=llm_provider, + ), + ) + self.litellm_deployment_tpm_limit.labels(**_labels).set(tpm) + + if rpm is not None: + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_deployment_rpm_limit" + ), + enum_values=UserAPIKeyLabelValues( + litellm_model_name=litellm_model_name, + model_id=model_id, + api_base=api_base, + api_provider=llm_provider, + ), + ) + self.litellm_deployment_rpm_limit.labels(**_labels).set(rpm) + def set_llm_deployment_success_metrics( self, request_kwargs: dict, @@ -1785,6 +1857,16 @@ class PrometheusLogger(CustomLogger): _model_info = _metadata.get("model_info") or {} model_id = _model_info.get("id", None) + if _model_info or _litellm_params: + self._set_deployment_tpm_rpm_limit_metrics( + model_info=_model_info, + litellm_params=_litellm_params, + litellm_model_name=litellm_model_name, + model_id=model_id, + api_base=api_base, + llm_provider=llm_provider, + ) + remaining_requests: Optional[int] = None remaining_tokens: Optional[int] = None if additional_headers := standard_logging_payload["hidden_params"][ @@ -2262,7 +2344,10 @@ class PrometheusLogger(CustomLogger): async def fetch_keys( page_size: int, page: int - ) -> Tuple[List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]], Optional[int]]: + ) -> Tuple[ + List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]], + Optional[int], + ]: key_list_response = await _list_key_helper( prisma_client=prisma_client, page=page, @@ -2378,12 +2463,16 @@ class PrometheusLogger(CustomLogger): # Get total user count total_users = await prisma_client.db.litellm_usertable.count() self.litellm_total_users_metric.set(total_users) - verbose_logger.debug(f"Prometheus: set litellm_total_users to {total_users}") + verbose_logger.debug( + f"Prometheus: set litellm_total_users to {total_users}" + ) # Get total team count total_teams = await prisma_client.db.litellm_teamtable.count() self.litellm_teams_count_metric.set(total_teams) - verbose_logger.debug(f"Prometheus: set litellm_teams_count to {total_teams}") + verbose_logger.debug( + f"Prometheus: set litellm_teams_count to {total_teams}" + ) except Exception as e: verbose_logger.exception( f"Error initializing user/team count metrics: {str(e)}" diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 9cb0a00d9f..00695cbfb5 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -351,9 +351,9 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: # Skip callable objects (functions, methods, lambdas) but not classes (type objects) if callable(data) and not isinstance(data, type): return None - # Skip known non-serializable object types (Logging, etc.) + # Skip known non-serializable object types (Logging, Router, etc.) obj_type_name = type(data).__name__ - if obj_type_name in ["Logging", "LiteLLMLoggingObj"]: + if obj_type_name in ["Logging", "LiteLLMLoggingObj", "Router"]: return None if isinstance(data, dict): diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index e290101f8b..060e98fd49 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -93,8 +93,11 @@ def get_litellm_params( "text_completion": text_completion, "azure_ad_token_provider": azure_ad_token_provider, "user_continue_message": user_continue_message, - "base_model": base_model or ( - _get_base_model_from_litellm_call_metadata(metadata=metadata) if metadata else None + "base_model": base_model + or ( + _get_base_model_from_litellm_call_metadata(metadata=metadata) + if metadata + else None ), "litellm_trace_id": litellm_trace_id, "litellm_session_id": litellm_session_id, @@ -139,5 +142,7 @@ def get_litellm_params( "aws_sts_endpoint": kwargs.get("aws_sts_endpoint"), "aws_external_id": kwargs.get("aws_external_id"), "aws_bedrock_runtime_endpoint": kwargs.get("aws_bedrock_runtime_endpoint"), + "tpm": kwargs.get("tpm"), + "rpm": kwargs.get("rpm"), } return litellm_params diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index fadeeffa9c..66b14946c3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -335,7 +335,9 @@ class Logging(LiteLLMLoggingBaseClass): self.start_time = start_time # log the call start time self.call_type = call_type self.litellm_call_id = litellm_call_id - self.litellm_trace_id: str = litellm_trace_id if litellm_trace_id else str(uuid.uuid4()) + self.litellm_trace_id: str = ( + litellm_trace_id if litellm_trace_id else str(uuid.uuid4()) + ) self.function_id = function_id self.streaming_chunks: List[Any] = [] # for generating complete stream response self.sync_streaming_chunks: List[ @@ -544,7 +546,10 @@ class Logging(LiteLLMLoggingBaseClass): if "stream_options" in additional_params: self.stream_options = additional_params["stream_options"] ## check if custom pricing set ## - if any(litellm_params.get(key) is not None for key in _CUSTOM_PRICING_KEYS & litellm_params.keys()): + if any( + litellm_params.get(key) is not None + for key in _CUSTOM_PRICING_KEYS & litellm_params.keys() + ): self.custom_pricing = True if "custom_llm_provider" in self.model_call_details: @@ -3299,6 +3304,7 @@ def _get_masked_values( "token", "key", "secret", + "vertex_credentials", ] return { k: ( @@ -4453,6 +4459,7 @@ class StandardLoggingPayloadSetup: user_api_key_request_route=None, spend_logs_metadata=None, requester_ip_address=None, + user_agent=None, requester_metadata=None, prompt_management_metadata=prompt_management_metadata, applied_guardrails=applied_guardrails, @@ -5138,6 +5145,7 @@ def get_standard_logging_object_payload( model_group=_model_group, model_id=_model_id, requester_ip_address=clean_metadata.get("requester_ip_address", None), + user_agent=clean_metadata.get("user_agent", None), messages=StandardLoggingPayloadSetup.append_system_prompt_messages( kwargs=kwargs, messages=kwargs.get("messages") ), @@ -5203,6 +5211,7 @@ def get_standard_logging_metadata( user_api_key_team_alias=None, spend_logs_metadata=None, requester_ip_address=None, + user_agent=None, requester_metadata=None, user_api_key_end_user_id=None, prompt_management_metadata=None, diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index bbe28e3ec2..25ad0a570c 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -21,11 +21,13 @@ from litellm.types.utils import ( ChatCompletionMessageToolCall, ChatCompletionRedactedThinkingBlock, Choices, + CompletionTokensDetailsWrapper, Delta, EmbeddingResponse, Function, HiddenParams, ImageResponse, + PromptTokensDetailsWrapper, ) from litellm.types.utils import Logprobs as TextCompletionLogprobs from litellm.types.utils import ( @@ -304,6 +306,22 @@ class LiteLLMResponseObjectHandler: "text_tokens": 0, } + # Map Responses API naming to Chat Completions API naming for cost calculator + if usage.get("prompt_tokens") is None: + usage["prompt_tokens"] = usage.get("input_tokens", 0) + if usage.get("completion_tokens") is None: + usage["completion_tokens"] = usage.get("output_tokens", 0) + + # Convert dicts to wrapper objects so getattr() works in cost calculation + if isinstance(usage.get("input_tokens_details"), dict): + usage["prompt_tokens_details"] = PromptTokensDetailsWrapper( + **usage["input_tokens_details"] + ) + if isinstance(usage.get("output_tokens_details"), dict): + usage["completion_tokens_details"] = CompletionTokensDetailsWrapper( + **usage["output_tokens_details"] + ) + if model_response_object is None: model_response_object = ImageResponse(**response_object) return model_response_object diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 30263543fc..1d1e38c09d 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4408,7 +4408,7 @@ def _bedrock_tools_pt(tools: List) -> List[BedrockToolBlock]: ] """ """ - Bedrock toolConfig looks like: + Bedrock toolConfig looks like: "tools": [ { "toolSpec": { @@ -4436,6 +4436,7 @@ def _bedrock_tools_pt(tools: List) -> List[BedrockToolBlock]: tool_block_list: List[BedrockToolBlock] = [] for tool in tools: + # Handle regular function tools parameters = tool.get("function", {}).get( "parameters", {"type": "object", "properties": {}} ) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9d50cc4d92..71d74121a3 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -110,6 +110,10 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["tools"] = tools_to_check if structured_messages: inputs["structured_messages"] = structured_messages + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=data, @@ -309,6 +313,14 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["images"] = images_to_check if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check + # Include model information from the response if available + response_model = None + if isinstance(response, dict): + response_model = response.get("model") + elif hasattr(response, "model"): + response_model = getattr(response, "model", None) + if response_model: + inputs["model"] = response_model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -552,7 +564,7 @@ class AnthropicMessagesHandler(BaseTranslation): response_content = response.get("content", []) else: response_content = getattr(response, "content", None) or [] - + if not response_content: return False for content_block in response_content: diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 58120bab1f..5ba0754b74 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -170,21 +170,33 @@ class LiteLLMAnthropicMessagesAdapter: def _add_cache_control_if_applicable( self, - source: Dict[str, Any], - target: Dict[str, Any], + source: Any, + target: Any, model: Optional[str], ) -> None: """ Extract cache_control from source and add to target if it should be preserved. + This method accepts Any type to support both regular dicts and TypedDict objects. + TypedDict objects (like ChatCompletionTextObject, ChatCompletionImageObject, etc.) + are dicts at runtime but have specific types at type-check time. Using Any allows + this method to work with both while maintaining runtime correctness. + Args: - source: Dict containing potential cache_control field - target: Dict to add cache_control to + source: Dict or TypedDict containing potential cache_control field + target: Dict or TypedDict to add cache_control to model: Model name to check if cache_control should be preserved """ - cache_control = source.get("cache_control") + # TypedDict objects are dicts at runtime, so .get() works + cache_control = source.get("cache_control") if isinstance(source, dict) else getattr(source, "cache_control", None) if cache_control and model and self.is_anthropic_claude_model(model): - target["cache_control"] = cache_control + # TypedDict objects support dict operations at runtime + # Use type ignore consistent with codebase pattern (see anthropic/chat/transformation.py:432) + if isinstance(target, dict): + target["cache_control"] = cache_control # type: ignore[typeddict-item] + else: + # Fallback for non-dict objects (shouldn't happen in practice) + cast(Dict[str, Any], target)["cache_control"] = cache_control def translatable_anthropic_params(self) -> List: """ @@ -220,7 +232,7 @@ class LiteLLMAnthropicMessagesAdapter: elif message_content and isinstance(message_content, list): for content in message_content: if content.get("type") == "text": - text_obj: Dict[str, Any] = ChatCompletionTextObject( + text_obj = ChatCompletionTextObject( type="text", text=content.get("text", "") ) self._add_cache_control_if_applicable(content, text_obj, model) @@ -236,7 +248,7 @@ class LiteLLMAnthropicMessagesAdapter: image_url_obj = ChatCompletionImageUrlObject( url=openai_image_url ) - image_obj: Dict[str, Any] = ChatCompletionImageObject( + image_obj = ChatCompletionImageObject( type="image_url", image_url=image_url_obj ) self._add_cache_control_if_applicable(content, image_obj, model) @@ -245,21 +257,21 @@ class LiteLLMAnthropicMessagesAdapter: # Convert Anthropic document format (PDF, etc.) to OpenAI format source = content.get("source", {}) openai_image_url = ( - self._translate_anthropic_image_to_openai(source) + self._translate_anthropic_image_to_openai(cast(dict, source)) ) if openai_image_url: image_url_obj = ChatCompletionImageUrlObject( url=openai_image_url ) - doc_obj: Dict[str, Any] = ChatCompletionImageObject( + doc_obj = ChatCompletionImageObject( type="image_url", image_url=image_url_obj ) self._add_cache_control_if_applicable(content, doc_obj, model) new_user_content_list.append(doc_obj) # type: ignore elif content.get("type") == "tool_result": if "content" not in content: - tool_result: Dict[str, Any] = ChatCompletionToolMessage( + tool_result = ChatCompletionToolMessage( role="tool", tool_call_id=content.get("tool_use_id", ""), content="", @@ -383,7 +395,7 @@ class LiteLLMAnthropicMessagesAdapter: assistant_message_str: Optional[str] = None assistant_content_list: List[Dict[str, Any]] = [] # For content blocks with cache_control has_cache_control_in_text = False - tool_calls: List[Dict[str, Any]] = [] + tool_calls: List[ChatCompletionAssistantToolCall] = [] thinking_blocks: List[ Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] ] = [] @@ -427,7 +439,7 @@ class LiteLLMAnthropicMessagesAdapter: provider_specific_fields ) - tool_call: Dict[str, Any] = ChatCompletionAssistantToolCall( + tool_call = ChatCompletionAssistantToolCall( id=content.get("id", ""), type="function", function=function_chunk, @@ -603,7 +615,7 @@ class LiteLLMAnthropicMessagesAdapter: for k, v in tool.items(): if k not in mapped_tool_params: # pass additional computer kwargs function_chunk.setdefault("parameters", {}).update({k: v}) - tool_param: Dict[str, Any] = ChatCompletionToolParam(type="function", function=function_chunk) + tool_param = ChatCompletionToolParam(type="function", function=function_chunk) self._add_cache_control_if_applicable(tool, tool_param, model) new_tools.append(tool_param) # type: ignore[arg-type] @@ -645,6 +657,41 @@ class LiteLLMAnthropicMessagesAdapter: }, } + def _add_system_message_to_messages( + self, + new_messages: List[AllMessageValues], + anthropic_message_request: AnthropicMessagesRequest, + ) -> None: + """Add system message to messages list if present in request.""" + if "system" not in anthropic_message_request: + return + system_content = anthropic_message_request["system"] + if not system_content: + return + # Handle system as string or array of content blocks + if isinstance(system_content, str): + new_messages.insert( + 0, + ChatCompletionSystemMessage(role="system", content=system_content), + ) + elif isinstance(system_content, list): + # Convert Anthropic system content blocks to OpenAI format + openai_system_content: List[Dict[str, Any]] = [] + model_name = anthropic_message_request.get("model", "") + for block in system_content: + if isinstance(block, dict) and block.get("type") == "text": + text_block: Dict[str, Any] = { + "type": "text", + "text": block.get("text", ""), + } + self._add_cache_control_if_applicable(block, text_block, model_name) + openai_system_content.append(text_block) + if openai_system_content: + new_messages.insert( + 0, + ChatCompletionSystemMessage(role="system", content=openai_system_content), # type: ignore + ) + def translate_anthropic_to_openai( self, anthropic_message_request: AnthropicMessagesRequest ) -> ChatCompletionRequest: @@ -673,32 +720,7 @@ class LiteLLMAnthropicMessagesAdapter: model=anthropic_message_request.get("model"), ) ## ADD SYSTEM MESSAGE TO MESSAGES - if "system" in anthropic_message_request: - system_content = anthropic_message_request["system"] - if system_content: - # Handle system as string or array of content blocks - if isinstance(system_content, str): - new_messages.insert( - 0, - ChatCompletionSystemMessage(role="system", content=system_content), - ) - elif isinstance(system_content, list): - # Convert Anthropic system content blocks to OpenAI format - openai_system_content: List[Dict[str, Any]] = [] - model_name = anthropic_message_request.get("model", "") - for block in system_content: - if isinstance(block, dict) and block.get("type") == "text": - text_block: Dict[str, Any] = { - "type": "text", - "text": block.get("text", ""), - } - self._add_cache_control_if_applicable(block, text_block, model_name) - openai_system_content.append(text_block) - if openai_system_content: - new_messages.insert( - 0, - ChatCompletionSystemMessage(role="system", content=openai_system_content), # type: ignore - ) + self._add_system_message_to_messages(new_messages, anthropic_message_request) new_kwargs: ChatCompletionRequest = { "model": anthropic_message_request["model"], @@ -902,7 +924,7 @@ class LiteLLMAnthropicMessagesAdapter: ) # extract usage usage: Usage = getattr(response, "usage") - anthropic_usage: Dict[str, Any] = AnthropicUsage( + anthropic_usage = AnthropicUsage( input_tokens=usage.prompt_tokens or 0, output_tokens=usage.completion_tokens or 0, ) @@ -1055,7 +1077,7 @@ class LiteLLMAnthropicMessagesAdapter: else: litellm_usage_chunk = None if litellm_usage_chunk is not None: - usage_delta: Dict[str, Any] = UsageDelta( + usage_delta = UsageDelta( input_tokens=litellm_usage_chunk.prompt_tokens or 0, output_tokens=litellm_usage_chunk.completion_tokens or 0, ) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index cb26a22edf..ec66514207 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -298,6 +298,39 @@ class AmazonConverseConfig(BaseConfig): # Check if the model is specifically Nova Lite 2 return "nova-2-lite" in model_without_region + def _map_web_search_options( + self, + web_search_options: dict, + model: str + ) -> Optional[BedrockToolBlock]: + """ + Map web_search_options to Nova grounding systemTool. + + Nova grounding (web search) is only supported on Amazon Nova models. + Returns None for non-Nova models. + + Args: + web_search_options: The web_search_options dict from the request + model: The model identifier string + + Returns: + BedrockToolBlock with systemTool for Nova models, None otherwise + + Reference: https://docs.aws.amazon.com/nova/latest/userguide/grounding.html + """ + # Only Nova models support nova_grounding + # Model strings can be like: "amazon.nova-pro-v1:0", "us.amazon.nova-pro-v1:0", etc. + if "nova" not in model.lower(): + verbose_logger.debug( + f"web_search_options passed but model {model} is not a Nova model. " + "Nova grounding is only supported on Amazon Nova models." + ) + return None + + # Nova doesn't support search_context_size or user_location params + # (unlike Anthropic), so we just enable grounding with no options + return BedrockToolBlock(systemTool={"name": "nova_grounding"}) + def _transform_reasoning_effort_to_reasoning_config( self, reasoning_effort: str ) -> dict: @@ -438,6 +471,10 @@ class AmazonConverseConfig(BaseConfig): ): supported_params.append("tools") + # Nova models support web_search_options (mapped to nova_grounding systemTool) + if base_model.startswith("amazon.nova"): + supported_params.append("web_search_options") + if litellm.utils.supports_tool_choice( model=model, custom_llm_provider=self.custom_llm_provider ) or litellm.utils.supports_tool_choice( @@ -730,6 +767,13 @@ class AmazonConverseConfig(BaseConfig): if bedrock_tier in ("default", "flex", "priority"): optional_params["serviceTier"] = {"type": bedrock_tier} + if param == "web_search_options" and value and isinstance(value, dict): + grounding_tool = self._map_web_search_options(value, model) + if grounding_tool is not None: + optional_params = self._add_tools_to_optional_params( + optional_params=optional_params, tools=[grounding_tool] + ) + # Only update thinking tokens for non-GPT-OSS models and non-Nova-Lite-2 models # Nova Lite 2 handles token budgeting differently through reasoningConfig if "gpt-oss" not in model and not self._is_nova_lite_2_model(model): @@ -1388,20 +1432,23 @@ class AmazonConverseConfig(BaseConfig): str, List[ChatCompletionToolCallChunk], Optional[List[BedrockConverseReasoningContentBlock]], + Optional[List[CitationsContentBlock]], ]: """ - Translate the message content to a string and a list of tool calls and reasoning content blocks + Translate the message content to a string and a list of tool calls, reasoning content blocks, and citations. Returns: content_str: str tools: List[ChatCompletionToolCallChunk] reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] + citationsContentBlocks: Optional[List[CitationsContentBlock]] - Citations from Nova grounding """ content_str = "" tools: List[ChatCompletionToolCallChunk] = [] reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = ( None ) + citationsContentBlocks: Optional[List[CitationsContentBlock]] = None for idx, content in enumerate(content_blocks): """ - Content is either a tool response or text @@ -1446,10 +1493,15 @@ class AmazonConverseConfig(BaseConfig): if reasoningContentBlocks is None: reasoningContentBlocks = [] reasoningContentBlocks.append(content["reasoningContent"]) + # Handle Nova grounding citations content + if "citationsContent" in content: + if citationsContentBlocks is None: + citationsContentBlocks = [] + citationsContentBlocks.append(content["citationsContent"]) - return content_str, tools, reasoningContentBlocks + return content_str, tools, reasoningContentBlocks, citationsContentBlocks - def _transform_response( + def _transform_response( # noqa: PLR0915 self, model: str, response: httpx.Response, @@ -1525,18 +1577,27 @@ class AmazonConverseConfig(BaseConfig): reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = ( None ) + citationsContentBlocks: Optional[List[CitationsContentBlock]] = None if message is not None: ( content_str, tools, reasoningContentBlocks, + citationsContentBlocks, ) = self._translate_message_content(message["content"]) + # Initialize provider_specific_fields if we have any special content blocks + provider_specific_fields: dict = {} + if reasoningContentBlocks is not None: + provider_specific_fields["reasoningContentBlocks"] = reasoningContentBlocks + if citationsContentBlocks is not None: + provider_specific_fields["citationsContent"] = citationsContentBlocks + + if provider_specific_fields: + chat_completion_message["provider_specific_fields"] = provider_specific_fields + if reasoningContentBlocks is not None: - chat_completion_message["provider_specific_fields"] = { - "reasoningContentBlocks": reasoningContentBlocks, - } chat_completion_message["reasoning_content"] = ( self._transform_reasoning_content(reasoningContentBlocks) ) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 17474fa022..1c58a11eeb 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -1476,6 +1476,11 @@ class AWSEventStreamDecoder: reasoning_content = ( "" # set to non-empty string to ensure consistency with Anthropic ) + elif "citationsContent" in delta_obj: + # Handle Nova grounding citations in streaming responses + provider_specific_fields = { + "citationsContent": delta_obj["citationsContent"], + } return ( text, tool_use, diff --git a/litellm/llms/cohere/rerank/guardrail_translation/handler.py b/litellm/llms/cohere/rerank/guardrail_translation/handler.py index 6893a5991c..b8133c59f7 100644 --- a/litellm/llms/cohere/rerank/guardrail_translation/handler.py +++ b/litellm/llms/cohere/rerank/guardrail_translation/handler.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail @@ -49,8 +50,13 @@ class CohereRerankHandler(BaseTranslation): # Process query only query = data.get("query") if query is not None and isinstance(query, str): + inputs = GenericGuardrailAPIInputs(texts=[query]) + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [query]}, + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index d0ed3f165c..fb00aa28f4 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -87,6 +87,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tools = data.get("tools") if tools: inputs["tools"] = tools + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -297,6 +301,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs["images"] = images_to_check if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check # type: ignore + # Include model information from the response if available + if hasattr(response, "model") and response.model: + inputs["model"] = response.model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -417,6 +424,13 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: inputs["images"] = images_to_check + # Include model information from the first response if available + if ( + responses_so_far + and hasattr(responses_so_far[0], "model") + and responses_so_far[0].model + ): + inputs["model"] = responses_so_far[0].model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py index 73d08cfead..1f8c6159da 100644 --- a/litellm/llms/openai/completion/guardrail_translation/handler.py +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail @@ -53,8 +54,13 @@ class OpenAITextCompletionHandler(BaseTranslation): if isinstance(prompt, str): # Single string prompt + inputs = GenericGuardrailAPIInputs(texts=[prompt]) + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [prompt]}, + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, @@ -80,8 +86,13 @@ class OpenAITextCompletionHandler(BaseTranslation): text_indices.append(idx) if texts_to_check: + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": texts_to_check}, + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, @@ -154,8 +165,12 @@ class OpenAITextCompletionHandler(BaseTranslation): if user_metadata: request_data["litellm_metadata"] = user_metadata + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + # Include model information from the response if available + if hasattr(response, "model") and response.model: + inputs["model"] = response.model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": texts_to_check}, + inputs=inputs, request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/openai/image_generation/cost_calculator.py b/litellm/llms/openai/image_generation/cost_calculator.py index 35caaf6e9b..988d562613 100644 --- a/litellm/llms/openai/image_generation/cost_calculator.py +++ b/litellm/llms/openai/image_generation/cost_calculator.py @@ -8,8 +8,7 @@ from typing import Optional from litellm import verbose_logger from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token -from litellm.responses.utils import ResponseAPILoggingUtils -from litellm.types.utils import ImageResponse +from litellm.types.utils import ImageResponse, Usage def cost_calculator( @@ -39,11 +38,18 @@ def cost_calculator( ) return 0.0 - # Transform ImageUsage to Usage using the existing helper - # ImageUsage has the same format as ResponseAPIUsage - chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + # If usage is already a Usage object with completion_tokens_details set, + # use it directly (it was already transformed in convert_to_image_response) + if isinstance(usage, Usage) and usage.completion_tokens_details is not None: + chat_usage = usage + else: + # Transform ImageUsage to Usage using the existing helper + # ImageUsage has the same format as ResponseAPIUsage + from litellm.responses.utils import ResponseAPILoggingUtils + + chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + usage + ) # Use generic_cost_per_token for cost calculation prompt_cost, completion_cost = generic_cost_per_token( diff --git a/litellm/llms/openai/image_generation/guardrail_translation/handler.py b/litellm/llms/openai/image_generation/guardrail_translation/handler.py index 842a64b187..e6340ba470 100644 --- a/litellm/llms/openai/image_generation/guardrail_translation/handler.py +++ b/litellm/llms/openai/image_generation/guardrail_translation/handler.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail @@ -52,8 +53,13 @@ class OpenAIImageGenerationHandler(BaseTranslation): # Apply guardrail to the prompt if isinstance(prompt, str): + inputs = GenericGuardrailAPIInputs(texts=[prompt]) + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [prompt]}, + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 9b8f15c762..d943662f9e 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -105,6 +105,10 @@ class OpenAIResponsesHandler(BaseTranslation): inputs["tools"] = tools_to_check if structured_messages: inputs["structured_messages"] = structured_messages # type: ignore + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -150,6 +154,10 @@ class OpenAIResponsesHandler(BaseTranslation): inputs["tools"] = tools_to_check if structured_messages: inputs["structured_messages"] = structured_messages # type: ignore + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=data, @@ -344,6 +352,14 @@ class OpenAIResponsesHandler(BaseTranslation): inputs["images"] = images_to_check if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check + # Include model information from the response if available + response_model = None + if isinstance(response, dict): + response_model = response.get("model") + elif hasattr(response, "model"): + response_model = getattr(response, "model", None) + if response_model: + inputs["model"] = response_model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -388,12 +404,15 @@ class OpenAIResponsesHandler(BaseTranslation): tool_calls = model_response_stream.choices[0].delta.tool_calls if tool_calls: + inputs = GenericGuardrailAPIInputs() + inputs["tool_calls"] = cast( + List[ChatCompletionToolCallChunk], tool_calls + ) + # Include model information if available + if hasattr(model_response_stream, "model") and model_response_stream.model: + inputs["model"] = model_response_stream.model _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={ - "tool_calls": cast( - List[ChatCompletionToolCallChunk], tool_calls - ) - }, + inputs=inputs, request_data={}, input_type="response", logging_obj=litellm_logging_obj, @@ -417,7 +436,11 @@ class OpenAIResponsesHandler(BaseTranslation): guardrail_inputs["tool_calls"] = cast( List[ChatCompletionToolCallChunk], tool_calls ) - if tool_calls: + # Include model information from the response if available + response_model = final_chunk.get("response", {}).get("model") + if response_model: + guardrail_inputs["model"] = response_model + if tool_calls or text: _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=guardrail_inputs, request_data={}, @@ -429,8 +452,14 @@ class OpenAIResponsesHandler(BaseTranslation): # tool_calls = model_response_stream.choices[0].tool_calls # convert openai response to model response string_so_far = self.get_streaming_string_so_far(responses_so_far) + inputs = GenericGuardrailAPIInputs(texts=[string_so_far]) + # Try to get model from the final chunk if available + if isinstance(final_chunk, dict): + response_model = final_chunk.get("response", {}).get("model") if isinstance(final_chunk.get("response"), dict) else None + if response_model: + inputs["model"] = response_model _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [string_so_far]}, + inputs=inputs, request_data={}, input_type="response", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/openai/speech/guardrail_translation/handler.py b/litellm/llms/openai/speech/guardrail_translation/handler.py index 4c2f71477b..e6796fbac2 100644 --- a/litellm/llms/openai/speech/guardrail_translation/handler.py +++ b/litellm/llms/openai/speech/guardrail_translation/handler.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail @@ -50,8 +51,13 @@ class OpenAITextToSpeechHandler(BaseTranslation): return data if isinstance(input_text, str): + inputs = GenericGuardrailAPIInputs(texts=[input_text]) + # Include model information if available (voice model) + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [input_text]}, + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py index ac416f42c8..3d76a21c38 100644 --- a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py +++ b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail @@ -88,8 +89,12 @@ class OpenAIAudioTranscriptionHandler(BaseTranslation): if user_metadata: request_data["litellm_metadata"] = user_metadata + inputs = GenericGuardrailAPIInputs(texts=[original_text]) + # Include model information from the response if available + if hasattr(response, "model") and response.model: + inputs["model"] = response.model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [original_text]}, + inputs=inputs, request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index c0979e37e6..40433d5341 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, List, Optional from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.proxy._types import PassThroughGuardrailSettings +from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail @@ -118,8 +119,13 @@ class PassThroughEndpointHandler(BaseTranslation): return data # Apply guardrail (pass-through doesn't modify the text, just checks it) + inputs = GenericGuardrailAPIInputs(texts=[text_to_check]) + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [text_to_check]}, + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, @@ -178,8 +184,13 @@ class PassThroughEndpointHandler(BaseTranslation): request_data["litellm_metadata"] = user_metadata # Apply guardrail (pass-through doesn't modify the text, just checks it) + inputs = GenericGuardrailAPIInputs(texts=[text_to_check]) + # Include model information from the response if available + response_model = response.get("model") if isinstance(response, dict) else None + if response_model: + inputs["model"] = response_model _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [text_to_check]}, + inputs=inputs, request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/vercel_ai_gateway/embedding/__init__.py b/litellm/llms/vercel_ai_gateway/embedding/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/litellm/llms/vercel_ai_gateway/embedding/transformation.py b/litellm/llms/vercel_ai_gateway/embedding/transformation.py new file mode 100644 index 0000000000..7238b05f10 --- /dev/null +++ b/litellm/llms/vercel_ai_gateway/embedding/transformation.py @@ -0,0 +1,176 @@ +""" +Vercel AI Gateway Embedding API Configuration. + +This module provides the configuration for Vercel AI Gateway's Embedding API. +Vercel AI Gateway is OpenAI-compatible and supports embeddings via the /v1/embeddings endpoint. + +Docs: https://vercel.com/docs/ai-gateway/openai-compat/embeddings +""" + +from typing import TYPE_CHECKING, Any, Optional + +import httpx + +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllEmbeddingInputValues +from litellm.types.utils import EmbeddingResponse +from litellm.utils import convert_to_model_response_object + +from ..common_utils import VercelAIGatewayException + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class VercelAIGatewayEmbeddingConfig(BaseEmbeddingConfig): + """ + Configuration for Vercel AI Gateway's Embedding API. + + Reference: https://vercel.com/docs/ai-gateway/openai-compat/embeddings + """ + + def validate_environment( + self, + headers: dict, + model: str, + messages: list, + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Validate environment and set up headers for Vercel AI Gateway API. + + Vercel AI Gateway requires: + - Authorization header with Bearer token (API key or OIDC token) + """ + vercel_headers = { + "Content-Type": "application/json", + } + + # Add Authorization header if api_key is provided + if api_key: + vercel_headers["Authorization"] = f"Bearer {api_key}" + + # Merge with existing headers (user's extra_headers take priority) + merged_headers = {**vercel_headers, **headers} + + return merged_headers + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + """ + Get the complete URL for Vercel AI Gateway Embedding API endpoint. + """ + if api_base: + api_base = api_base.rstrip("/") + else: + api_base = ( + get_secret_str("VERCEL_AI_GATEWAY_API_BASE") + or "https://ai-gateway.vercel.sh/v1" + ) + + return f"{api_base}/embeddings" + + def transform_embedding_request( + self, + model: str, + input: AllEmbeddingInputValues, + optional_params: dict, + headers: dict, + ) -> dict: + """ + Transform embedding request to Vercel AI Gateway format (OpenAI-compatible). + """ + # Ensure input is a list + if isinstance(input, str): + input = [input] + + # Strip 'vercel_ai_gateway/' prefix if present + if model.startswith("vercel_ai_gateway/"): + model = model.replace("vercel_ai_gateway/", "", 1) + + return { + "model": model, + "input": input, + **optional_params, + } + + def transform_embedding_response( + self, + model: str, + raw_response: httpx.Response, + model_response: EmbeddingResponse, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str], + request_data: dict, + optional_params: dict, + litellm_params: dict, + ) -> EmbeddingResponse: + """ + Transform embedding response from Vercel AI Gateway format (OpenAI-compatible). + """ + logging_obj.post_call(original_response=raw_response.text) + + # Vercel AI Gateway returns standard OpenAI-compatible embedding response + response_json = raw_response.json() + + return convert_to_model_response_object( + response_object=response_json, + model_response_object=model_response, + response_type="embedding", + ) + + def get_supported_openai_params(self, model: str) -> list: + """ + Get list of supported OpenAI parameters for Vercel AI Gateway embeddings. + + Vercel AI Gateway supports the standard OpenAI embeddings parameters + and auto-maps 'dimensions' to each provider's expected field. + """ + return [ + "timeout", + "dimensions", + "encoding_format", + "user", + ] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + """ + Map OpenAI parameters to Vercel AI Gateway format. + """ + for param, value in non_default_params.items(): + if param in self.get_supported_openai_params(model): + optional_params[param] = value + return optional_params + + def get_error_class( + self, error_message: str, status_code: int, headers: Any + ) -> Any: + """ + Get the error class for Vercel AI Gateway errors. + """ + return VercelAIGatewayException( + message=error_message, + status_code=status_code, + headers=headers, + ) diff --git a/litellm/main.py b/litellm/main.py index ce84c8988e..5b8c569a39 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -148,7 +148,7 @@ from litellm.utils import ( validate_and_fix_openai_messages, validate_and_fix_openai_tools, validate_chat_completion_tool_choice, - validate_openai_optional_params + validate_openai_optional_params, ) from ._logging import verbose_logger @@ -368,7 +368,7 @@ class AsyncCompletions: @tracer.wrap() @client -async def acompletion( # noqa: PLR0915 +async def acompletion( # noqa: PLR0915 model: str, # Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create messages: List = [], @@ -603,12 +603,11 @@ async def acompletion( # noqa: PLR0915 if timeout is not None and isinstance(timeout, (int, float)): timeout_value = float(timeout) init_response = await asyncio.wait_for( - loop.run_in_executor(None, func_with_context), - timeout=timeout_value + loop.run_in_executor(None, func_with_context), timeout=timeout_value ) else: init_response = await loop.run_in_executor(None, func_with_context) - + if isinstance(init_response, dict) or isinstance( init_response, ModelResponse ): ## CACHING SCENARIO @@ -640,6 +639,7 @@ async def acompletion( # noqa: PLR0915 except asyncio.TimeoutError: custom_llm_provider = custom_llm_provider or "openai" from litellm.exceptions import Timeout + raise Timeout( message=f"Request timed out after {timeout} seconds", model=model, @@ -1118,7 +1118,6 @@ def completion( # type: ignore # noqa: PLR0915 # validate optional params stop = validate_openai_optional_params(stop=stop) - ######### unpacking kwargs ##################### args = locals() @@ -1135,7 +1134,9 @@ def completion( # type: ignore # noqa: PLR0915 # Check if MCP tools are present (following responses pattern) # Cast tools to Optional[Iterable[ToolParam]] for type checking tools_for_mcp = cast(Optional[Iterable[ToolParam]], tools) - if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools_for_mcp): + if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway( + tools=tools_for_mcp + ): # Return coroutine - acompletion will await it # completion() can return a coroutine when MCP tools are present, which acompletion() awaits return acompletion_with_mcp( # type: ignore[return-value] @@ -1536,6 +1537,8 @@ def completion( # type: ignore # noqa: PLR0915 max_retries=max_retries, timeout=timeout, litellm_request_debug=kwargs.get("litellm_request_debug", False), + tpm=kwargs.get("tpm"), + rpm=kwargs.get("rpm"), ) cast(LiteLLMLoggingObj, logging).update_environment_variables( model=model, @@ -2361,11 +2364,7 @@ def completion( # type: ignore # noqa: PLR0915 input=messages, api_key=api_key, original_response=response ) elif custom_llm_provider == "minimax": - api_key = ( - api_key - or get_secret_str("MINIMAX_API_KEY") - or litellm.api_key - ) + api_key = api_key or get_secret_str("MINIMAX_API_KEY") or litellm.api_key api_base = ( api_base @@ -2413,7 +2412,9 @@ def completion( # type: ignore # noqa: PLR0915 or custom_llm_provider == "wandb" or custom_llm_provider == "clarifai" or custom_llm_provider in litellm.openai_compatible_providers - or JSONProviderRegistry.exists(custom_llm_provider) # JSON-configured providers + or JSONProviderRegistry.exists( + custom_llm_provider + ) # JSON-configured providers or "ft:gpt-3.5-turbo" in model # finetune gpt-3.5-turbo ): # allow user to make an openai call with a custom base # note: if a user sets a custom base - we should ensure this works @@ -4724,7 +4725,7 @@ def embedding( # noqa: PLR0915 if headers is not None and headers != {}: optional_params["extra_headers"] = headers - + if encoding_format is not None: optional_params["encoding_format"] = encoding_format else: @@ -4866,6 +4867,36 @@ def embedding( # noqa: PLR0915 headers = openrouter_headers + response = base_llm_http_handler.embedding( + model=model, + input=input, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + logging_obj=logging, + timeout=timeout, + model_response=EmbeddingResponse(), + optional_params=optional_params, + client=client, + aembedding=aembedding, + litellm_params=litellm_params_dict, + headers=headers, + ) + elif custom_llm_provider == "vercel_ai_gateway": + api_base = ( + api_base + or litellm.api_base + or get_secret_str("VERCEL_AI_GATEWAY_API_BASE") + or "https://ai-gateway.vercel.sh/v1" + ) + + api_key = ( + api_key + or litellm.api_key + or get_secret_str("VERCEL_AI_GATEWAY_API_KEY") + or get_secret_str("VERCEL_OIDC_TOKEN") + ) + response = base_llm_http_handler.embedding( model=model, input=input, @@ -6759,9 +6790,7 @@ def speech( # noqa: PLR0915 if text_to_speech_provider_config is None: text_to_speech_provider_config = MinimaxTextToSpeechConfig() - minimax_config = cast( - MinimaxTextToSpeechConfig, text_to_speech_provider_config - ) + minimax_config = cast(MinimaxTextToSpeechConfig, text_to_speech_provider_config) if api_base is not None: litellm_params_dict["api_base"] = api_base @@ -6901,7 +6930,7 @@ async def ahealth_check( custom_llm_provider_from_params = model_params.get("custom_llm_provider", None) api_base_from_params = model_params.get("api_base", None) api_key_from_params = model_params.get("api_key", None) - + model, custom_llm_provider, _, _ = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider_from_params, @@ -7275,8 +7304,9 @@ def __getattr__(name: str) -> Any: _encoding = tiktoken.get_encoding("cl100k_base") # Cache it in the module's __dict__ for subsequent accesses import sys + sys.modules[__name__].__dict__["encoding"] = _encoding global _encoding_cache _encoding_cache = _encoding return _encoding - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") \ No newline at end of file + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ded591a8f5..56feff548a 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -387,25 +387,57 @@ async def callback(code: str, state: str): 1. Try resource_metadata from WWW-Authenticate header (if present) 2. Fall back to path-based well-known URI: /.well-known/oauth-protected-resource/{path} ( - If the resource identifier value contains a path or query component, any terminating slash (/) - following the host component MUST be removed before inserting /.well-known/ and the well-known - URI path suffix between the host component and the path(include root path) and/or query components. + If the resource identifier value contains a path or query component, any terminating slash (/) + following the host component MUST be removed before inserting /.well-known/ and the well-known + URI path suffix between the host component and the path(include root path) and/or query components. https://datatracker.ietf.org/doc/html/rfc9728#section-3.1) 3. Fall back to root-based well-known URI: /.well-known/oauth-protected-resource + + Dual Pattern Support: + - Standard MCP pattern: /mcp/{server_name} (recommended, used by mcp-inspector, VSCode Copilot) + - LiteLLM legacy pattern: /{server_name}/mcp (backward compatibility) + + The resource URL returned matches the pattern used in the discovery request. """ -@router.get(f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}/mcp") -@router.get("/.well-known/oauth-protected-resource") -async def oauth_protected_resource_mcp( - request: Request, mcp_server_name: Optional[str] = None -): + + +def _build_oauth_protected_resource_response( + request: Request, + mcp_server_name: Optional[str], + use_standard_pattern: bool, +) -> dict: + """ + Build OAuth protected resource response with the appropriate URL pattern. + + Args: + request: FastAPI Request object + mcp_server_name: Name of the MCP server + use_standard_pattern: If True, use /mcp/{server_name} pattern; + if False, use /{server_name}/mcp pattern + + Returns: + OAuth protected resource metadata dict + """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - # Get the correct base URL considering X-Forwarded-* headers + request_base_url = get_request_base_url(request) mcp_server: Optional[MCPServer] = None if mcp_server_name: mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name) + + # Build resource URL based on the pattern + if mcp_server_name: + if use_standard_pattern: + # Standard MCP pattern: /mcp/{server_name} + resource_url = f"{request_base_url}/mcp/{mcp_server_name}" + else: + # LiteLLM legacy pattern: /{server_name}/mcp + resource_url = f"{request_base_url}/{mcp_server_name}/mcp" + else: + resource_url = f"{request_base_url}/mcp" + return { "authorization_servers": [ ( @@ -414,14 +446,55 @@ async def oauth_protected_resource_mcp( else f"{request_base_url}" ) ], - "resource": ( - f"{request_base_url}/{mcp_server_name}/mcp" - if mcp_server_name - else f"{request_base_url}/mcp" - ), # this is what Claude will call + "resource": resource_url, "scopes_supported": mcp_server.scopes if mcp_server else [], } + +# Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} +# This is the pattern expected by standard MCP clients (mcp-inspector, VSCode Copilot) +@router.get(f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}") +async def oauth_protected_resource_mcp_standard( + request: Request, mcp_server_name: str +): + """ + OAuth protected resource discovery endpoint using standard MCP URL pattern. + + Standard pattern: /mcp/{server_name} + Discovery path: /.well-known/oauth-protected-resource/mcp/{server_name} + + This endpoint is compliant with MCP specification and works with standard + MCP clients like mcp-inspector and VSCode Copilot. + """ + return _build_oauth_protected_resource_response( + request=request, + mcp_server_name=mcp_server_name, + use_standard_pattern=True, + ) + + +# LiteLLM legacy pattern: /.well-known/oauth-protected-resource/{server_name}/mcp +# Kept for backward compatibility with existing deployments +@router.get(f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}/mcp") +@router.get("/.well-known/oauth-protected-resource") +async def oauth_protected_resource_mcp( + request: Request, mcp_server_name: Optional[str] = None +): + """ + OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern. + + Legacy pattern: /{server_name}/mcp + Discovery path: /.well-known/oauth-protected-resource/{server_name}/mcp + + This endpoint is kept for backward compatibility. New integrations should + use the standard MCP pattern (/mcp/{server_name}) instead. + """ + return _build_oauth_protected_resource_response( + request=request, + mcp_server_name=mcp_server_name, + use_standard_pattern=False, + ) + """ https://datatracker.ietf.org/doc/html/rfc8414#section-3.1 RFC 8414: Path-aware OAuth discovery @@ -430,15 +503,26 @@ async def oauth_protected_resource_mcp( the well-known URI suffix between the host component and the path(include root path) component. """ -@router.get(f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}") -@router.get("/.well-known/oauth-authorization-server") -async def oauth_authorization_server_mcp( - request: Request, mcp_server_name: Optional[str] = None -): + + +def _build_oauth_authorization_server_response( + request: Request, + mcp_server_name: Optional[str], +) -> dict: + """ + Build OAuth authorization server metadata response. + + Args: + request: FastAPI Request object + mcp_server_name: Name of the MCP server + + Returns: + OAuth authorization server metadata dict + """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - # Get the correct base URL considering X-Forwarded-* headers + request_base_url = get_request_base_url(request) authorization_endpoint = ( @@ -470,18 +554,58 @@ async def oauth_authorization_server_mcp( } +# Standard MCP pattern: /.well-known/oauth-authorization-server/mcp/{server_name} +@router.get(f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}") +async def oauth_authorization_server_mcp_standard( + request: Request, mcp_server_name: str +): + """ + OAuth authorization server discovery endpoint using standard MCP URL pattern. + + Standard pattern: /mcp/{server_name} + Discovery path: /.well-known/oauth-authorization-server/mcp/{server_name} + """ + return _build_oauth_authorization_server_response( + request=request, + mcp_server_name=mcp_server_name, + ) + + +# LiteLLM legacy pattern and root endpoint +@router.get(f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}") +@router.get("/.well-known/oauth-authorization-server") +async def oauth_authorization_server_mcp( + request: Request, mcp_server_name: Optional[str] = None +): + """ + OAuth authorization server discovery endpoint. + + Supports both legacy pattern (/{server_name}) and root endpoint. + """ + return _build_oauth_authorization_server_response( + request=request, + mcp_server_name=mcp_server_name, + ) + + # Alias for standard OpenID discovery @router.get("/.well-known/openid-configuration") async def openid_configuration(request: Request): return await oauth_authorization_server_mcp(request) +# Additional legacy pattern support @router.get("/.well-known/oauth-authorization-server/{mcp_server_name}/mcp") -@router.get("/.well-known/oauth-authorization-server") -async def oauth_authorization_server_root( - request: Request, mcp_server_name: Optional[str] = None +async def oauth_authorization_server_legacy( + request: Request, mcp_server_name: str ): - return await oauth_authorization_server_mcp(request, mcp_server_name) + """ + OAuth authorization server discovery for legacy /{server_name}/mcp pattern. + """ + return _build_oauth_authorization_server_response( + request=request, + mcp_server_name=mcp_server_name, + ) @router.post("/{mcp_server_name}/register") diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py index 8d6d236b88..4d53ae7059 100644 --- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py +++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py @@ -46,8 +46,13 @@ class MCPGuardrailTranslationHandler(BaseTranslation): ) return data + inputs = GenericGuardrailAPIInputs(texts=[content]) + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs=GenericGuardrailAPIInputs(texts=[content]), + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 81806d6c06..92fd54e877 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -63,7 +63,20 @@ from litellm.types.mcp_server.mcp_server_manager import ( MCPOAuthMetadata, MCPServer, ) -from mcp.shared.tool_name_validation import SEP_986_URL, validate_tool_name + +try: + from mcp.shared.tool_name_validation import SEP_986_URL, validate_tool_name # type: ignore +except ImportError: + SEP_986_URL = "https://github.com/modelcontextprotocol/protocol/blob/main/proposals/0001-tool-name-validation.md" + + def validate_tool_name(name: str): + from pydantic import BaseModel + + class MockResult(BaseModel): + is_valid: bool = True + warnings: list = [] + + return MockResult() # Probe includes characters on both sides of the separator to mimic real prefixed tool names. @@ -90,7 +103,9 @@ def _warn_on_server_name_fields( if result.is_valid: return - warning_text = "; ".join(result.warnings) if result.warnings else "Validation failed" + warning_text = ( + "; ".join(result.warnings) if result.warnings else "Validation failed" + ) verbose_logger.warning( "MCP server '%s' has invalid %s '%s': %s", server_id, @@ -103,7 +118,6 @@ def _warn_on_server_name_fields( _warn("server_name", server_name) - def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]: """ Deserialize optional JSON mappings stored in the database. @@ -391,10 +405,13 @@ class MCPServerManager: # Note: `extra_headers` on MCPServer is a List[str] of header names to forward # from the client request (not available in this OpenAPI tool generation step). # `static_headers` is a dict of concrete headers to always send. - headers = merge_mcp_headers( - extra_headers=headers, - static_headers=server.static_headers, - ) or {} + headers = ( + merge_mcp_headers( + extra_headers=headers, + static_headers=server.static_headers, + ) + or {} + ) verbose_logger.debug( f"Using headers for OpenAPI tools (excluding sensitive values): " diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 51f476fad0..6d54c3871e 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -73,7 +73,11 @@ if MCP_AVAILABLE: AuthContextMiddleware, auth_context_var, ) - from mcp.server.streamable_http_manager import StreamableHTTPSessionManager + + try: + from mcp.server.streamable_http_manager import StreamableHTTPSessionManager + except ImportError: + StreamableHTTPSessionManager = None # type: ignore from mcp.types import ( CallToolResult, EmbeddedResource, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 185de9ec0a..c854d81ec7 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -425,12 +425,12 @@ class LiteLLMRoutes(enum.Enum): ] google_routes = [ - "/v1beta/models/{model_name}:countTokens", - "/v1beta/models/{model_name}:generateContent", - "/v1beta/models/{model_name}:streamGenerateContent", - "/models/{model_name}:countTokens", - "/models/{model_name}:generateContent", - "/models/{model_name}:streamGenerateContent", + "/v1beta/models/{model_name:path}:countTokens", + "/v1beta/models/{model_name:path}:generateContent", + "/v1beta/models/{model_name:path}:streamGenerateContent", + "/models/{model_name:path}:countTokens", + "/models/{model_name:path}:generateContent", + "/models/{model_name:path}:streamGenerateContent", # Google Interactions API "/interactions", "/v1beta/interactions", @@ -3428,6 +3428,8 @@ class LitellmMetadataFromRequestHeaders(TypedDict, total=False): """ spend_logs_metadata: Optional[dict] + agent_id: Optional[str] + trace_id: Optional[str] class JWTKeyItem(TypedDict, total=False): diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 9b9a988c07..2bd84a1d98 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -758,11 +758,27 @@ def get_model_from_request( if match: model = match.group(1) + # If still not found, extract model from Google generateContent-style routes. + # These routes put the model in the path and allow "/" inside the model id. + # Examples: + # - /v1beta/models/gemini-2.0-flash:generateContent + # - /v1beta/models/bedrock/claude-sonnet-3.7:generateContent + # - /models/custom/ns/model:streamGenerateContent + if model is None and not route.lower().startswith("/vertex"): + google_match = re.search(r"/(?:v1beta|beta)/models/([^:]+):", route) + if google_match: + model = google_match.group(1) + + if model is None and not route.lower().startswith("/vertex"): + google_match = re.search(r"^/models/([^:]+):", route) + if google_match: + model = google_match.group(1) + # If still not found, extract from Vertex AI passthrough route # Pattern: /vertex_ai/.../models/{model_id}:* # Example: /vertex_ai/v1/.../models/gemini-1.5-pro:generateContent - if model is None and "/vertex" in route.lower(): - vertex_match = re.search(r"/models/([^/:]+)", route) + if model is None and route.lower().startswith("/vertex"): + vertex_match = re.search(r"/models/([^:]+)", route) if vertex_match: model = vertex_match.group(1) diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 5be44f479b..939cfefadc 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -197,9 +197,9 @@ async def authenticate_user( # noqa: PLR0915 - Login with UI_USERNAME and UI_PASSWORD - Login with Invite Link `user_email` and `password` combination """ - if secrets.compare_digest(username, ui_username) and secrets.compare_digest( - password, ui_password - ): + if secrets.compare_digest( + username.encode("utf-8"), ui_username.encode("utf-8") + ) and secrets.compare_digest(password.encode("utf-8"), ui_password.encode("utf-8")): # Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin user_role = LitellmUserRoles.PROXY_ADMIN user_id = LITELLM_PROXY_ADMIN_NAME @@ -313,9 +313,9 @@ async def authenticate_user( # noqa: PLR0915 # check if password == _user_row.password hash_password = hash_token(token=password) - if secrets.compare_digest(password, _password) or secrets.compare_digest( - hash_password, _password - ): + if secrets.compare_digest( + password.encode("utf-8"), _password.encode("utf-8") + ) or secrets.compare_digest(hash_password.encode("utf-8"), _password.encode("utf-8")): if os.getenv("DATABASE_URL") is not None: # Expire any previous UI session tokens for this user await expire_previous_ui_session_tokens( diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 96a70b8016..fa63ff01b3 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -392,7 +392,15 @@ class RouteChecks: # Ensure route is a string before attempting regex matching if not isinstance(route, str): return False - pattern = re.sub(r"\{[^}]+\}", r"[^/]+", pattern) + + def _placeholder_to_regex(match: re.Match) -> str: + placeholder = match.group(0).strip("{}") + if placeholder.endswith(":path"): + # allow "/" in the placeholder value, but don't eat the route suffix after ":" + return r"[^:]+" + return r"[^/]+" + + pattern = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern) # Anchor the pattern to match the entire string pattern = f"^{pattern}$" if re.match(pattern, route): diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index c0cff84bac..faeca9b2ae 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -274,11 +274,20 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915 WebSearchInterceptionLogger, ) - websearch_interception_obj = WebSearchInterceptionLogger.initialize_from_proxy_config( - litellm_settings=litellm_settings, - callback_specific_params=callback_specific_params, + websearch_interception_obj = ( + WebSearchInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params=callback_specific_params, + ) ) imported_list.append(websearch_interception_obj) + elif isinstance(callback, str) and callback == "datadog_cost_management": + from litellm.integrations.datadog.datadog_cost_management import ( + DatadogCostManagementLogger, + ) + + datadog_cost_management_obj = DatadogCostManagementLogger() + imported_list.append(datadog_cost_management_obj) elif isinstance(callback, CustomLogger): imported_list.append(callback) else: @@ -353,17 +362,17 @@ def get_remaining_tokens_and_requests_from_request_data(data: Dict) -> Dict[str, remaining_requests_variable_name = f"litellm-key-remaining-requests-{model_group}" remaining_requests = _metadata.get(remaining_requests_variable_name, None) if remaining_requests: - headers[f"x-litellm-key-remaining-requests-{h11_model_group_name}"] = ( - remaining_requests - ) + headers[ + f"x-litellm-key-remaining-requests-{h11_model_group_name}" + ] = remaining_requests # Remaining Tokens remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}" remaining_tokens = _metadata.get(remaining_tokens_variable_name, None) if remaining_tokens: - headers[f"x-litellm-key-remaining-tokens-{h11_model_group_name}"] = ( - remaining_tokens - ) + headers[ + f"x-litellm-key-remaining-tokens-{h11_model_group_name}" + ] = remaining_tokens return headers @@ -438,9 +447,9 @@ def add_guardrail_response_to_standard_logging_object( ): if litellm_logging_obj is None: return - standard_logging_object: Optional[StandardLoggingPayload] = ( - litellm_logging_obj.model_call_details.get("standard_logging_object") - ) + standard_logging_object: Optional[ + StandardLoggingPayload + ] = litellm_logging_obj.model_call_details.get("standard_logging_object") if standard_logging_object is None: return guardrail_information = standard_logging_object.get("guardrail_information", []) @@ -469,7 +478,9 @@ def get_metadata_variable_name_from_kwargs( return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" -def process_callback(_callback: str, callback_type: str, environment_variables: dict) -> dict: +def process_callback( + _callback: str, callback_type: str, environment_variables: dict +) -> dict: """Process a single callback and return its data with environment variables""" env_vars = CustomLogger.get_callback_env_vars(_callback) @@ -481,11 +492,9 @@ def process_callback(_callback: str, callback_type: str, environment_variables: else: env_vars_dict[_var] = env_variable - return { - "name": _callback, - "variables": env_vars_dict, - "type": callback_type - } + return {"name": _callback, "variables": env_vars_dict, "type": callback_type} + + def normalize_callback_names(callbacks: Iterable[Any]) -> List[Any]: if callbacks is None: return [] diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index d5a6266146..b37074e25e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -185,6 +185,7 @@ class GenericGuardrailAPI(CustomGuardrail): tools = inputs.get("tools") structured_messages = inputs.get("structured_messages") tool_calls = inputs.get("tool_calls") + model = inputs.get("model") # Use provided request_data or create an empty dict if request_data is None: @@ -215,6 +216,7 @@ class GenericGuardrailAPI(CustomGuardrail): tool_calls=tool_calls, additional_provider_specific_params=additional_params, input_type=input_type, + model=model, ) # Prepare headers diff --git a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py index 5f57cab1db..3598dbe741 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py +++ b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py @@ -8,6 +8,7 @@ import os import uuid from typing import TYPE_CHECKING, Any, Literal, Optional, Type +import httpx from fastapi import HTTPException from litellm._logging import verbose_proxy_logger @@ -25,10 +26,12 @@ if TYPE_CHECKING: class OnyxGuardrail(CustomGuardrail): def __init__( - self, api_base: Optional[str] = None, api_key: Optional[str] = None, **kwargs + self, api_base: Optional[str] = None, api_key: Optional[str] = None, timeout: Optional[float] = 10.0, **kwargs ): + timeout = timeout or int(os.getenv("ONYX_TIMEOUT", 10.0)) self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback + llm_provider=httpxSpecialProvider.GuardrailCallback, + params={"timeout": httpx.Timeout(timeout=timeout, connect=5.0)}, ) self.api_base = api_base or os.getenv( "ONYX_API_BASE", diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 064538ef3b..9be78264e8 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -558,6 +558,16 @@ class LiteLLMProxyRequestSetup: ######################################################################################### # Finally update the requests metadata with the `metadata_from_headers` ######################################################################################### + agent_id_from_header = headers.get("x-litellm-agent-id") + trace_id_from_header = headers.get("x-litellm-trace-id") + if agent_id_from_header: + metadata_from_headers["agent_id"] = agent_id_from_header + verbose_proxy_logger.debug(f"Extracted agent_id from header: {agent_id_from_header}") + + if trace_id_from_header: + metadata_from_headers["trace_id"] = trace_id_from_header + verbose_proxy_logger.debug(f"Extracted trace_id from header: {trace_id_from_header}") + if isinstance(data[_metadata_variable_name], dict): data[_metadata_variable_name].update(metadata_from_headers) return data @@ -846,7 +856,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915 # Add headers to metadata for guardrails to access (fixes #17477) # Guardrails use metadata["headers"] to access request headers (e.g., User-Agent) - if _metadata_variable_name in data and isinstance(data[_metadata_variable_name], dict): + if _metadata_variable_name in data and isinstance( + data[_metadata_variable_name], dict + ): data[_metadata_variable_name]["headers"] = _headers # check for forwardable headers @@ -1002,7 +1014,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915 # User spend, budget - used by prometheus.py # Follow same pattern as team and API key budgets - data[_metadata_variable_name]["user_api_key_user_spend"] = user_api_key_dict.user_spend + data[_metadata_variable_name][ + "user_api_key_user_spend" + ] = user_api_key_dict.user_spend data[_metadata_variable_name][ "user_api_key_user_max_budget" ] = user_api_key_dict.user_max_budget @@ -1029,8 +1043,8 @@ async def add_litellm_data_to_request( # noqa: PLR0915 ## [Enterprise Only] # Add User-IP Address requester_ip_address = "" - if premium_user is True: - # Only set the IP Address for Enterprise Users + if True: # Always set the IP Address if available + # logic for tracking IP Address # logic for tracking IP Address if ( @@ -1050,6 +1064,16 @@ async def add_litellm_data_to_request( # noqa: PLR0915 requester_ip_address = request.client.host data[_metadata_variable_name]["requester_ip_address"] = requester_ip_address + # Add User-Agent + user_agent = "" + if ( + request is not None + and hasattr(request, "headers") + and "user-agent" in request.headers + ): + user_agent = request.headers["user-agent"] + data[_metadata_variable_name]["user_agent"] = user_agent + # Check if using tag based routing tags = LiteLLMProxyRequestSetup.add_request_tag_to_metadata( llm_router=llm_router, @@ -1532,7 +1556,9 @@ def add_guardrails_from_policy_engine( f"policy_count={len(registry.get_all_policies())}" ) if not registry.is_initialized(): - verbose_proxy_logger.debug("Policy engine not initialized, skipping policy matching") + verbose_proxy_logger.debug( + "Policy engine not initialized, skipping policy matching" + ) return # Build context from request @@ -1550,13 +1576,17 @@ def add_guardrails_from_policy_engine( # Get matching policies via attachments matching_policy_names = PolicyMatcher.get_matching_policies(context=context) - verbose_proxy_logger.debug(f"Policy engine: matched policies via attachments: {matching_policy_names}") + verbose_proxy_logger.debug( + f"Policy engine: matched policies via attachments: {matching_policy_names}" + ) # Combine attachment-based policies with dynamic request body policies all_policy_names = set(matching_policy_names) if request_body_policies and isinstance(request_body_policies, list): all_policy_names.update(request_body_policies) - verbose_proxy_logger.debug(f"Policy engine: added dynamic policies from request body: {request_body_policies}") + verbose_proxy_logger.debug( + f"Policy engine: added dynamic policies from request body: {request_body_policies}" + ) if not all_policy_names: return @@ -1567,7 +1597,9 @@ def add_guardrails_from_policy_engine( context=context, ) - verbose_proxy_logger.debug(f"Policy engine: applied policies (conditions matched): {applied_policy_names}") + verbose_proxy_logger.debug( + f"Policy engine: applied policies (conditions matched): {applied_policy_names}" + ) # Track applied policies in metadata for response headers for policy_name in applied_policy_names: @@ -1578,7 +1610,9 @@ def add_guardrails_from_policy_engine( # Resolve guardrails from matching policies resolved_guardrails = PolicyResolver.resolve_guardrails_for_context(context=context) - verbose_proxy_logger.debug(f"Policy engine: resolved guardrails: {resolved_guardrails}") + verbose_proxy_logger.debug( + f"Policy engine: resolved guardrails: {resolved_guardrails}" + ) if not resolved_guardrails: return diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index a15f47d13b..83d7f3fde4 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -56,7 +56,19 @@ except ImportError as e: MCP_AVAILABLE = False if MCP_AVAILABLE: - from mcp.shared.tool_name_validation import validate_tool_name + try: + from mcp.shared.tool_name_validation import validate_tool_name # type: ignore + except ImportError: + + def validate_tool_name(name: str): + from pydantic import BaseModel + + class MockResult(BaseModel): + is_valid: bool = True + warnings: list = [] + + return MockResult() + from litellm.proxy._experimental.mcp_server.db import ( create_mcp_server, delete_mcp_server, @@ -122,9 +134,7 @@ if MCP_AVAILABLE: ) if validation_result.warnings: error_messages_text = ( - error_messages_text - + "\n" - + "\n".join(validation_result.warnings) + error_messages_text + "\n" + "\n".join(validation_result.warnings) ) raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 613390c0ec..1b827411e5 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -99,24 +99,24 @@ def determine_role_from_groups( ) -> Optional[LitellmUserRoles]: """ Determine the highest privilege role for a user based on their groups. - + Role hierarchy (highest to lowest): - proxy_admin - proxy_admin_viewer - internal_user - internal_user_viewer - + Args: user_groups: List of group names from the SSO token role_mappings: RoleMappings configuration object - + Returns: The highest privilege role found, or default_role if no matches, or None """ if not role_mappings.roles: # No role mappings configured, return default_role return role_mappings.default_role - + # Role hierarchy (highest to lowest) role_hierarchy = [ LitellmUserRoles.PROXY_ADMIN, @@ -124,20 +124,22 @@ def determine_role_from_groups( LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, ] - + # Convert user_groups to a set for efficient lookup user_groups_set = set(user_groups) if isinstance(user_groups, list) else set() - + # Find the highest privilege role the user belongs to for role in role_hierarchy: if role in role_mappings.roles: role_groups = role_mappings.roles[role] - if isinstance(role_groups, list) and user_groups_set.intersection(set(role_groups)): + if isinstance(role_groups, list) and user_groups_set.intersection( + set(role_groups) + ): verbose_proxy_logger.debug( f"User groups {user_groups} matched role '{role.value}' via groups: {role_groups}" ) return role - + # No matching groups found, return default_role verbose_proxy_logger.debug( f"User groups {user_groups} did not match any role mappings, using default_role: {role_mappings.default_role}" @@ -326,9 +328,7 @@ def generic_response_convertor( "GENERIC_USER_PROVIDER_ATTRIBUTE", "provider" ) - generic_user_role_attribute_name = os.getenv( - "GENERIC_USER_ROLE_ATTRIBUTE", "role" - ) + generic_user_role_attribute_name = os.getenv("GENERIC_USER_ROLE_ATTRIBUTE", "role") verbose_proxy_logger.debug( f" generic_user_id_attribute_name: {generic_user_id_attribute_name}\n generic_user_email_attribute_name: {generic_user_email_attribute_name}" @@ -345,12 +345,15 @@ def generic_response_convertor( # Determine user role based on role_mappings if available # Only apply role_mappings for GENERIC SSO provider user_role: Optional[LitellmUserRoles] = None - - if role_mappings is not None and role_mappings.provider.lower() in ["generic", "okta"]: + + if role_mappings is not None and role_mappings.provider.lower() in [ + "generic", + "okta", + ]: # Use role_mappings to determine role from groups group_claim = role_mappings.group_claim user_groups_raw: Any = get_nested_value(response, group_claim) - + # Handle different formats: could be a list, string (comma-separated), or single value user_groups: List[str] = [] if isinstance(user_groups_raw, list): @@ -361,7 +364,7 @@ def generic_response_convertor( elif user_groups_raw is not None: # Single value user_groups = [str(user_groups_raw)] - + if user_groups: user_role = determine_role_from_groups(user_groups, role_mappings) verbose_proxy_logger.debug( @@ -373,10 +376,12 @@ def generic_response_convertor( verbose_proxy_logger.debug( f"No groups found in '{group_claim}', using default_role: {role_mappings.default_role}" ) - + # Fallback to existing logic if role_mappings not used if user_role is None: - user_role_from_sso = get_nested_value(response, generic_user_role_attribute_name) + user_role_from_sso = get_nested_value( + response, generic_user_role_attribute_name + ) if user_role_from_sso is not None: role = get_litellm_user_role(user_role_from_sso) if role is not None: @@ -399,7 +404,9 @@ def generic_response_convertor( ) -def _setup_generic_sso_env_vars(generic_client_id: str, redirect_url: str) -> Tuple[str, List[str], str, str, str, bool]: +def _setup_generic_sso_env_vars( + generic_client_id: str, redirect_url: str +) -> Tuple[str, List[str], str, str, str, bool]: """Setup and validate Generic SSO environment variables.""" generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None) generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ") @@ -492,7 +499,43 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: verbose_proxy_logger.debug( f"Could not load role_mappings from database: {e}. Continuing with existing role logic." ) + + generic_role_mappings = os.getenv("GENERIC_ROLE_MAPPINGS_ROLES", None) + generic_role_mappings_group_claim = os.getenv( + "GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", None + ) + generic_role_mappoings_default_role = os.getenv( + "GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", None + ) + if generic_role_mappings is not None: + verbose_proxy_logger.debug( + "Found role_mappings for generic provider in environment variables" + ) + import ast + try: + generic_user_role_mappings_data: Dict[ + LitellmUserRoles, List[str] + ] = ast.literal_eval(generic_role_mappings) + if isinstance(generic_user_role_mappings_data, dict): + from litellm.types.proxy.management_endpoints.ui_sso import ( + RoleMappings, + ) + + role_mappings_data = { + "provider": "generic", + "group_claim": generic_role_mappings_group_claim, + "default_role": generic_role_mappoings_default_role, + "roles": generic_user_role_mappings_data, + } + + role_mappings = RoleMappings(**role_mappings_data) + verbose_proxy_logger.debug( + f"Loaded role_mappings from environments for provider '{role_mappings.provider}'." + ) + return role_mappings + except TypeError as e: + verbose_proxy_logger.warning(f"Error decoding role mappings from environment variables: {e}. Continuing with existing role logic.") return role_mappings @@ -529,7 +572,7 @@ async def get_generic_sso_response( # Get role_mappings from SSO settings if available role_mappings = await _setup_role_mappings() - + def response_convertor(response, client): nonlocal received_response # return for user debugging received_response = response @@ -1217,20 +1260,24 @@ async def insert_sso_user( role_mappings_configured = False try: from litellm.proxy.utils import get_prisma_client_or_throw - + prisma_client = get_prisma_client_or_throw( "Prisma client is None, connect a database to your proxy" ) - + # Get SSO config from dedicated table sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( where={"id": "sso_config"} ) - + if sso_db_record and sso_db_record.sso_settings: sso_settings_dict = dict(sso_db_record.sso_settings) role_mappings_data = sso_settings_dict.get("role_mappings") role_mappings_configured = role_mappings_data is not None + generic_user_role_mappings = os.getenv("GENERIC_USER_ROLE_MAPPINGS", None) + if generic_user_role_mappings is not None: + role_mappings_configured = True + except Exception as e: # If we can't check role_mappings, continue with existing logic verbose_proxy_logger.debug( @@ -1240,7 +1287,10 @@ async def insert_sso_user( # Apply default_internal_user_params if litellm.default_internal_user_params: # If role_mappings is configured and user_role is already set from SSO, preserve it - if role_mappings_configured and user_defined_values.get("user_role") is not None: + if ( + role_mappings_configured + and user_defined_values.get("user_role") is not None + ): # Preserve the SSO-extracted role, but apply other defaults preserved_role = user_defined_values.get("user_role") user_defined_values.update(litellm.default_internal_user_params) # type: ignore diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 16cdd9da64..183c25ed46 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5420,6 +5420,24 @@ async def chat_completion( # noqa: PLR0915 global general_settings, user_debug, proxy_logging_obj, llm_model_list global user_temperature, user_request_timeout, user_max_tokens, user_api_base data = await _read_request_body(request=request) + if user_api_key_dict is not None: + if data.get("metadata") is None: + data["metadata"] = {} + if ( + hasattr(user_api_key_dict, "user_id") + and user_api_key_dict.user_id is not None + ): + data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + if ( + hasattr(user_api_key_dict, "team_id") + and user_api_key_dict.team_id is not None + ): + data["metadata"]["user_api_key_team_id"] = user_api_key_dict.team_id + if ( + hasattr(user_api_key_dict, "org_id") + and user_api_key_dict.org_id is not None + ): + data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) try: result = await base_llm_response_processor.base_process_llm_request( @@ -5571,6 +5589,24 @@ async def completion( # noqa: PLR0915 data = {} try: data = await _read_request_body(request=request) + if user_api_key_dict is not None: + if data.get("metadata") is None: + data["metadata"] = {} + if ( + hasattr(user_api_key_dict, "user_id") + and user_api_key_dict.user_id is not None + ): + data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + if ( + hasattr(user_api_key_dict, "team_id") + and user_api_key_dict.team_id is not None + ): + data["metadata"]["user_api_key_team_id"] = user_api_key_dict.team_id + if ( + hasattr(user_api_key_dict, "org_id") + and user_api_key_dict.org_id is not None + ): + data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) return await base_llm_response_processor.base_process_llm_request( request=request, @@ -5790,6 +5826,25 @@ async def embeddings( # noqa: PLR0915 ) data["input"] = input_list + if user_api_key_dict is not None: + if data.get("metadata") is None: + data["metadata"] = {} + if ( + hasattr(user_api_key_dict, "user_id") + and user_api_key_dict.user_id is not None + ): + data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + if ( + hasattr(user_api_key_dict, "team_id") + and user_api_key_dict.team_id is not None + ): + data["metadata"]["user_api_key_team_id"] = user_api_key_dict.team_id + if ( + hasattr(user_api_key_dict, "org_id") + and user_api_key_dict.org_id is not None + ): + data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id + # Use unified request processor (same as chat/completions and responses) base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) @@ -9645,7 +9700,7 @@ def get_logo_url(): @app.get("/get_image", include_in_schema=False) -def get_image(): +async def get_image(): """Get logo to show on admin UI""" # get current_dir @@ -9664,25 +9719,37 @@ def get_image(): if is_non_root and not os.path.exists(default_logo): default_logo = default_site_logo + cache_dir = assets_dir if is_non_root else current_dir + cache_path = os.path.join(cache_dir, "cached_logo.jpg") + + # [OPTIMIZATION] Check if the cached image exists first + if os.path.exists(cache_path): + return FileResponse(cache_path, media_type="image/jpeg") + logo_path = os.getenv("UI_LOGO_PATH", default_logo) verbose_proxy_logger.debug("Reading logo from path: %s", logo_path) # Check if the logo path is an HTTP/HTTPS URL if logo_path.startswith(("http://", "https://")): - # Download the image and cache it - client = HTTPHandler() - response = client.get(logo_path) - if response.status_code == 200: - # Save the image to a local file - cache_dir = assets_dir if is_non_root else current_dir - cache_path = os.path.join(cache_dir, "cached_logo.jpg") - with open(cache_path, "wb") as f: - f.write(response.content) + try: + # Download the image and cache it + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - # Return the cached image as a FileResponse - return FileResponse(cache_path, media_type="image/jpeg") - else: - # Handle the case when the image cannot be downloaded + async_client = AsyncHTTPHandler(timeout=5.0) + response = await async_client.get(logo_path) + if response.status_code == 200: + # Save the image to a local file + with open(cache_path, "wb") as f: + f.write(response.content) + + # Return the cached image as a FileResponse + return FileResponse(cache_path, media_type="image/jpeg") + else: + # Handle the case when the image cannot be downloaded + return FileResponse(default_logo, media_type="image/jpeg") + except Exception as e: + # Handle any exceptions during the download (e.g., timeout, connection error) + verbose_proxy_logger.debug(f"Error downloading logo from {logo_path}: {e}") return FileResponse(default_logo, media_type="image/jpeg") else: # Return the local image file if the logo path is not an HTTP/HTTPS URL diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 79b4fd6873..79f182817f 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -26,6 +26,184 @@ from litellm.proxy.common_utils.http_parsing_utils import ( router = APIRouter() +def _build_file_metadata_entry( + response: Any, + file_data: Optional[Tuple[str, bytes, str]] = None, + file_url: Optional[str] = None, +) -> Dict[str, Any]: + """ + Build a file metadata entry for storing in vector_store_metadata. + + Args: + response: The response from litellm.aingest containing file_id + file_data: Optional tuple of (filename, content, content_type) + file_url: Optional URL if file was ingested from URL + + Returns: + Dictionary with file metadata (file_id, filename, file_url, ingested_at, etc.) + """ + from datetime import datetime, timezone + + # Extract file_id from response + file_id = None + if hasattr(response, "get"): + file_id = response.get("file_id") + elif hasattr(response, "file_id"): + file_id = response.file_id + + # Extract file information from file_data tuple + filename = None + file_size = None + content_type = None + + if file_data: + filename = file_data[0] + file_size = len(file_data[1]) if len(file_data) > 1 else None + content_type = file_data[2] if len(file_data) > 2 else None + + # Build file metadata entry + file_entry = { + "file_id": file_id, + "filename": filename, + "file_url": file_url, + "ingested_at": datetime.now(timezone.utc).isoformat(), + } + + # Add optional fields if available + if file_size is not None: + file_entry["file_size"] = file_size + if content_type is not None: + file_entry["content_type"] = content_type + + return file_entry + + +async def _save_vector_store_to_db_from_rag_ingest( + response: Any, + ingest_options: Dict[str, Any], + prisma_client, + user_api_key_dict: UserAPIKeyAuth, + file_data: Optional[Tuple[str, bytes, str]] = None, + file_url: Optional[str] = None, +) -> None: + """ + Helper function to save a newly created vector store from RAG ingest to the database. + + This function: + - Extracts vector store ID and config from the ingest response + - Checks if the vector store already exists in the database + - Creates a new database entry if it doesn't exist + - Adds the vector store to the registry + + Args: + response: The response from litellm.aingest() + ingest_options: The ingest options containing vector store config + prisma_client: The Prisma database client + user_api_key_dict: User API key authentication info + """ + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + create_vector_store_in_db, + ) + + # Handle both dict and object responses + if hasattr(response, "get"): + vector_store_id = response.get("vector_store_id") + elif hasattr(response, "vector_store_id"): + vector_store_id = response.vector_store_id + else: + verbose_proxy_logger.warning( + f"Unable to extract vector_store_id from response type: {type(response)}" + ) + return + + if vector_store_id is None or not isinstance(vector_store_id, str): + verbose_proxy_logger.warning( + "Vector store ID is None or not a string, skipping database save" + ) + return + + vector_store_config = ingest_options.get("vector_store", {}) + custom_llm_provider = vector_store_config.get("custom_llm_provider") + + # Extract litellm_vector_store_params for custom name and description + litellm_vector_store_params = ingest_options.get("litellm_vector_store_params", {}) + custom_vector_store_name = litellm_vector_store_params.get("vector_store_name") + custom_vector_store_description = litellm_vector_store_params.get("vector_store_description") + + # Build file metadata entry using helper + file_entry = _build_file_metadata_entry( + response=response, + file_data=file_data, + file_url=file_url, + ) + + try: + # Check if vector store already exists in database + existing_vector_store = ( + await prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": vector_store_id} + ) + ) + + # Only create if it doesn't exist + if existing_vector_store is None: + verbose_proxy_logger.info( + f"Saving newly created vector store {vector_store_id} to database" + ) + + # Initialize metadata with first file + initial_metadata = { + "ingested_files": [file_entry] + } + + # Use custom name if provided, otherwise default + vector_store_name = custom_vector_store_name or f"RAG Vector Store - {vector_store_id[:8]}" + vector_store_description = custom_vector_store_description or "Created via RAG ingest endpoint" + + await create_vector_store_in_db( + vector_store_id=vector_store_id, + custom_llm_provider=custom_llm_provider or "openai", + prisma_client=prisma_client, + vector_store_name=vector_store_name, + vector_store_description=vector_store_description, + vector_store_metadata=initial_metadata, + ) + + verbose_proxy_logger.info( + f"Vector store {vector_store_id} saved to database successfully" + ) + else: + verbose_proxy_logger.info( + f"Vector store {vector_store_id} already exists, appending file to metadata" + ) + + # Update existing vector store with new file + existing_metadata = existing_vector_store.vector_store_metadata or {} + if isinstance(existing_metadata, str): + import json + existing_metadata = json.loads(existing_metadata) + + ingested_files = existing_metadata.get("ingested_files", []) + ingested_files.append(file_entry) + existing_metadata["ingested_files"] = ingested_files + + # Update the vector store + from litellm.proxy.utils import safe_dumps + await prisma_client.db.litellm_managedvectorstorestable.update( + where={"vector_store_id": vector_store_id}, + data={"vector_store_metadata": safe_dumps(existing_metadata)} + ) + + verbose_proxy_logger.info( + f"Added file {file_entry.get('filename') or file_entry.get('file_url', 'Unknown')} to vector store {vector_store_id} metadata" + ) + except Exception as db_error: + # Log the error but don't fail the request since ingestion succeeded + verbose_proxy_logger.exception( + f"Failed to save vector store {vector_store_id} to database: {db_error}" + ) + + async def parse_rag_ingest_request( request: Request, ) -> Tuple[Dict[str, Any], Optional[Tuple[str, bytes, str]], Optional[str], Optional[str]]: @@ -158,6 +336,7 @@ async def rag_ingest( add_litellm_data_to_request, general_settings, llm_router, + prisma_client, proxy_config, version, ) @@ -189,6 +368,25 @@ async def rag_ingest( **request_data, ) + # Save vector store to database if it was newly created and prisma_client is available + verbose_proxy_logger.debug( + f"RAG Ingest - Checking database save conditions: prisma_client={prisma_client is not None}, response={response is not None}, response_type={type(response)}" + ) + + if prisma_client is not None and response is not None: + await _save_vector_store_to_db_from_rag_ingest( + response=response, + ingest_options=ingest_options, + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + file_data=file_data, + file_url=file_url, + ) + else: + verbose_proxy_logger.warning( + f"Skipping database save: prisma_client={prisma_client is not None}, response={response is not None}" + ) + return response except HTTPException: diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 5d348d35cb..db4e4beec2 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -396,7 +396,7 @@ def get_logging_payload( # noqa: PLR0915 ) # Extract agent_id for A2A requests (set directly on model_call_details) - agent_id: Optional[str] = kwargs.get("agent_id") + agent_id: Optional[str] = kwargs.get("agent_id") or metadata.get("agent_id") custom_llm_provider = kwargs.get("custom_llm_provider") raw_model = cast(str, kwargs.get("model") or "") model_name = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index bc61a60fe5..f3787e62f4 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -133,6 +133,112 @@ async def _resolve_embedding_config_from_db( return None +######################################################## +# Helper Functions +######################################################## +async def create_vector_store_in_db( + vector_store_id: str, + custom_llm_provider: str, + prisma_client, + vector_store_name: Optional[str] = None, + vector_store_description: Optional[str] = None, + vector_store_metadata: Optional[Dict] = None, + litellm_params: Optional[Dict] = None, + litellm_credential_name: Optional[str] = None, +) -> LiteLLM_ManagedVectorStore: + """ + Helper function to create a vector store in the database. + + This function handles: + - Checking if vector store already exists + - Creating the vector store in the database + - Adding it to the vector store registry + + Returns: + LiteLLM_ManagedVectorStore: The created vector store object + + Raises: + HTTPException: If vector store already exists or database error occurs + """ + from litellm.types.router import GenericLiteLLMParams + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + # Check if vector store already exists + existing_vector_store = ( + await prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": vector_store_id} + ) + ) + if existing_vector_store is not None: + raise HTTPException( + status_code=400, + detail=f"Vector store with ID {vector_store_id} already exists", + ) + + # Prepare data for database + data_to_create: Dict[str, Any] = { + "vector_store_id": vector_store_id, + "custom_llm_provider": custom_llm_provider, + } + + if vector_store_name is not None: + data_to_create["vector_store_name"] = vector_store_name + if vector_store_description is not None: + data_to_create["vector_store_description"] = vector_store_description + if vector_store_metadata is not None: + data_to_create["vector_store_metadata"] = safe_dumps(vector_store_metadata) + if litellm_credential_name is not None: + data_to_create["litellm_credential_name"] = litellm_credential_name + + # Handle litellm_params - always provide at least an empty dict + if litellm_params: + # Auto-resolve embedding config if embedding model is provided but config is not + embedding_model = litellm_params.get("litellm_embedding_model") + if embedding_model and not 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: + litellm_params["litellm_embedding_config"] = resolved_config + verbose_proxy_logger.info( + f"Auto-resolved embedding config for model {embedding_model}" + ) + + litellm_params_dict = GenericLiteLLMParams( + **litellm_params + ).model_dump(exclude_none=True) + data_to_create["litellm_params"] = safe_dumps(litellm_params_dict) + else: + # Provide empty dict if no litellm_params provided + data_to_create["litellm_params"] = safe_dumps({}) + + # Create in database + _new_vector_store = ( + await prisma_client.db.litellm_managedvectorstorestable.create( + data=data_to_create + ) + ) + + new_vector_store: LiteLLM_ManagedVectorStore = LiteLLM_ManagedVectorStore( + **_new_vector_store.model_dump() + ) + + # Add vector store to registry + if litellm.vector_store_registry is not None: + litellm.vector_store_registry.add_vector_store_to_registry( + vector_store=new_vector_store + ) + + verbose_proxy_logger.info( + f"Vector store {vector_store_id} created in database successfully" + ) + + return new_vector_store + + ######################################################## # Management Endpoints ######################################################## @@ -156,71 +262,34 @@ async def new_vector_store( - vector_store_metadata: Optional[Dict] - Additional metadata for the vector store """ 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") try: - # Check if vector store already exists - existing_vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.find_unique( - where={"vector_store_id": vector_store.get("vector_store_id")} - ) - ) - if existing_vector_store is not None: + vector_store_id = vector_store.get("vector_store_id") + custom_llm_provider = vector_store.get("custom_llm_provider") + + if not vector_store_id or not custom_llm_provider: raise HTTPException( status_code=400, - detail=f"Vector store with ID {vector_store.get('vector_store_id')} already exists", - ) - - if vector_store.get("vector_store_metadata") is not None: - vector_store["vector_store_metadata"] = safe_dumps( - vector_store.get("vector_store_metadata") - ) - - # Safely handle JSON serialization of litellm_params - litellm_params_json: Optional[str] = None - _input_litellm_params: dict = vector_store.get("litellm_params", {}) or {} - if _input_litellm_params is not None: - # 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) - litellm_params_json = safe_dumps(litellm_params_dict) - del vector_store["litellm_params"] - - _new_vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.create( - data={ - **vector_store, - "litellm_params": litellm_params_json, - } + detail="vector_store_id and custom_llm_provider are required" ) + + # Extract and validate metadata + metadata = vector_store.get("vector_store_metadata") + validated_metadata: Optional[Dict] = None + if metadata is not None and isinstance(metadata, dict): + validated_metadata = metadata + + new_vector_store = await create_vector_store_in_db( + vector_store_id=vector_store_id, + custom_llm_provider=custom_llm_provider, + prisma_client=prisma_client, + vector_store_name=vector_store.get("vector_store_name"), + vector_store_description=vector_store.get("vector_store_description"), + vector_store_metadata=validated_metadata, + litellm_params=vector_store.get("litellm_params"), + litellm_credential_name=vector_store.get("litellm_credential_name"), ) - new_vector_store: LiteLLM_ManagedVectorStore = LiteLLM_ManagedVectorStore( - **_new_vector_store.model_dump() - ) - - # Add vector store to registry - if litellm.vector_store_registry is not None: - litellm.vector_store_registry.add_vector_store_to_registry( - vector_store=new_vector_store - ) - return { "status": "success", "message": f"Vector store {vector_store.get('vector_store_id')} created successfully", diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 0ccbb435e5..620097f83d 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -198,6 +198,10 @@ async def _execute_query_pipeline( """ Execute the RAG query pipeline. """ + # Extract router from kwargs - use it for completion if available + # to properly resolve virtual model names + router: Optional["Router"] = kwargs.pop("router", None) + # 1. Extract query from last user message query_text = RAGQuery.extract_query_from_messages(messages) if not query_text: @@ -233,12 +237,21 @@ async def _execute_query_pipeline( context_message = RAGQuery.build_context_message(context_chunks) modified_messages = messages[:-1] + [context_message] + [messages[-1]] - response = await litellm.acompletion( - model=model, - messages=modified_messages, - stream=stream, - **kwargs, - ) + # Use router if available to properly resolve virtual model names + if router is not None: + response = await router.acompletion( + model=model, + messages=modified_messages, + stream=stream, + **kwargs, + ) + else: + response = await litellm.acompletion( + model=model, + messages=modified_messages, + stream=stream, + **kwargs, + ) # 5. Attach search results to response if not stream and isinstance(response, ModelResponse): diff --git a/litellm/router.py b/litellm/router.py index 54650b120a..9fe37efa3a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1707,8 +1707,11 @@ class Router: litellm_params = deployment.get("litellm_params", {}) dep_num_retries = litellm_params.get("num_retries") - if dep_num_retries is not None and isinstance(dep_num_retries, int): - exception.num_retries = dep_num_retries # type: ignore + if dep_num_retries is not None: + try: + exception.num_retries = int(dep_num_retries) # type: ignore # Handle both int and str + except (ValueError, TypeError): + pass # Skip if value can't be converted to int def _update_kwargs_with_default_litellm_params( self, kwargs: dict, metadata_variable_name: Optional[str] = "metadata" diff --git a/litellm/types/integrations/datadog_cost_management.py b/litellm/types/integrations/datadog_cost_management.py new file mode 100644 index 0000000000..fe04f43ea0 --- /dev/null +++ b/litellm/types/integrations/datadog_cost_management.py @@ -0,0 +1,27 @@ +from typing import Dict, Optional, TypedDict + + +from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams + + +class DatadogCostManagementInitParams(StandardCustomLoggerInitParams): + """ + Init params for Datadog Cost Management + """ + + datadog_cost_management_params: Optional[Dict] = None + + +class DatadogFOCUSCostEntry(TypedDict): + """ + Represents a single cost line item in the FOCUS format. + Ref: https://focus.finops.org/#specification + """ + + ProviderName: str + ChargeDescription: str + ChargePeriodStart: str + ChargePeriodEnd: str + BilledCost: float + BillingCurrency: str + Tags: Optional[Dict[str, str]] diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 670f8215d8..146c390363 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -150,6 +150,9 @@ class UserAPIKeyLabelNames(Enum): FALLBACK_MODEL = "fallback_model" ROUTE = "route" MODEL_GROUP = "model_group" + CLIENT_IP = "client_ip" + USER_AGENT = "user_agent" + CALLBACK_NAME = "callback_name" DEFINED_PROMETHEUS_METRICS = Literal[ @@ -196,6 +199,12 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_cache_hits_metric", "litellm_cache_misses_metric", "litellm_cached_tokens_metric", + "litellm_deployment_tpm_limit", + "litellm_deployment_rpm_limit", + "litellm_remaining_api_key_requests_for_model", + "litellm_remaining_api_key_tokens_for_model", + "litellm_llm_api_failed_requests_metric", + "litellm_callback_logging_failures_metric", ] @@ -209,6 +218,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.REQUESTED_MODEL.value, UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.USER.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_llm_api_time_to_first_token_metric = [ @@ -244,6 +254,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.TEAM_ALIAS.value, UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] # Guardrail metrics - these use custom labels (guardrail_name, status, error_type, hook_type) @@ -263,6 +274,8 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.STATUS_CODE.value, UserAPIKeyLabelNames.USER_EMAIL.value, UserAPIKeyLabelNames.ROUTE.value, + UserAPIKeyLabelNames.CLIENT_IP.value, + UserAPIKeyLabelNames.USER_AGENT.value, UserAPIKeyLabelNames.MODEL_ID.value, ] @@ -278,6 +291,8 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.EXCEPTION_STATUS.value, UserAPIKeyLabelNames.EXCEPTION_CLASS.value, UserAPIKeyLabelNames.ROUTE.value, + UserAPIKeyLabelNames.CLIENT_IP.value, + UserAPIKeyLabelNames.USER_AGENT.value, UserAPIKeyLabelNames.MODEL_ID.value, ] @@ -299,6 +314,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, UserAPIKeyLabelNames.API_KEY_HASH.value, UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_remaining_requests_metric = [ @@ -308,6 +324,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, UserAPIKeyLabelNames.API_KEY_HASH.value, UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_remaining_tokens_metric = [ @@ -317,6 +334,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, UserAPIKeyLabelNames.API_KEY_HASH.value, UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_requests_metric = [ @@ -328,6 +346,9 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.TEAM_ALIAS.value, UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.USER_EMAIL.value, + UserAPIKeyLabelNames.CLIENT_IP.value, + UserAPIKeyLabelNames.USER_AGENT.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_spend_metric = [ @@ -339,6 +360,9 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.TEAM_ALIAS.value, UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.USER_EMAIL.value, + UserAPIKeyLabelNames.CLIENT_IP.value, + UserAPIKeyLabelNames.USER_AGENT.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_input_tokens_metric = [ @@ -351,6 +375,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.USER_EMAIL.value, UserAPIKeyLabelNames.REQUESTED_MODEL.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_total_tokens_metric = [ @@ -363,6 +388,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.USER_EMAIL.value, UserAPIKeyLabelNames.REQUESTED_MODEL.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_output_tokens_metric = [ @@ -375,6 +401,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.USER_EMAIL.value, UserAPIKeyLabelNames.REQUESTED_MODEL.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_deployment_state = [ @@ -384,6 +411,15 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.API_PROVIDER.value, ] + litellm_deployment_tpm_limit = [ + UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.MODEL_ID.value, + UserAPIKeyLabelNames.API_BASE.value, + UserAPIKeyLabelNames.API_PROVIDER.value, + ] + + litellm_deployment_rpm_limit = litellm_deployment_tpm_limit + litellm_deployment_cooled_down = [ UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, UserAPIKeyLabelNames.MODEL_ID.value, @@ -401,6 +437,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.TEAM_ALIAS.value, UserAPIKeyLabelNames.EXCEPTION_STATUS.value, UserAPIKeyLabelNames.EXCEPTION_CLASS.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_deployment_failed_fallbacks = litellm_deployment_successful_fallbacks @@ -443,6 +480,26 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, ] + litellm_user_budget_remaining_hours_metric = [ + UserAPIKeyLabelNames.USER.value, + ] + + litellm_remaining_api_key_requests_for_model = [ + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + ] + + litellm_remaining_api_key_tokens_for_model = [ + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + ] + + litellm_callback_logging_failures_metric = [ + UserAPIKeyLabelNames.CALLBACK_NAME.value, + ] + # Add deployment metrics litellm_deployment_failure_responses = [ UserAPIKeyLabelNames.REQUESTED_MODEL.value, @@ -456,6 +513,8 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.API_KEY_ALIAS.value, UserAPIKeyLabelNames.TEAM.value, UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.CLIENT_IP.value, + UserAPIKeyLabelNames.USER_AGENT.value, ] litellm_deployment_total_requests = [ @@ -468,10 +527,37 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.API_KEY_ALIAS.value, UserAPIKeyLabelNames.TEAM.value, UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.CLIENT_IP.value, + UserAPIKeyLabelNames.USER_AGENT.value, ] litellm_deployment_success_responses = litellm_deployment_total_requests + litellm_remaining_api_key_requests_for_model = [ + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.MODEL_ID.value, + ] + + litellm_remaining_api_key_tokens_for_model = [ + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.MODEL_ID.value, + ] + + litellm_llm_api_failed_requests_metric = [ + UserAPIKeyLabelNames.END_USER.value, + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.USER.value, + UserAPIKeyLabelNames.MODEL_ID.value, + ] + # Buffer monitoring metrics - these typically don't need additional labels litellm_pod_lock_manager_size: List[str] = [] @@ -492,6 +578,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.TEAM_ALIAS.value, UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.USER.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_cache_hits_metric = _cache_metric_labels @@ -584,6 +671,12 @@ class UserAPIKeyLabelValues(BaseModel): route: Annotated[ Optional[str], Field(..., alias=UserAPIKeyLabelNames.ROUTE.value) ] = None + client_ip: Annotated[ + Optional[str], Field(..., alias=UserAPIKeyLabelNames.CLIENT_IP.value) + ] = None + user_agent: Annotated[ + Optional[str], Field(..., alias=UserAPIKeyLabelNames.USER_AGENT.value) + ] = None class PrometheusMetricsConfig(BaseModel): diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index ef2f1ba4d5..a85aaafe23 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -93,6 +93,67 @@ class GuardrailConverseContentBlock(TypedDict, total=False): text: GuardrailConverseTextBlock +class CitationWebLocationBlock(TypedDict, total=False): + """ + Web location block for Nova grounding citations. + Contains the URL and domain from web search results. + + Reference: https://docs.aws.amazon.com/nova/latest/userguide/grounding.html + """ + + url: str + domain: str + + +class CitationLocationBlock(TypedDict, total=False): + """ + Location block containing the web location for a citation. + """ + + web: CitationWebLocationBlock + + +class CitationReferenceBlock(TypedDict, total=False): + """ + Citation reference block containing a single citation with its location. + + Each citation contains: + - location.web.url: The URL of the source + - location.web.domain: The domain of the source + """ + + location: CitationLocationBlock + + +class CitationsContentBlock(TypedDict, total=False): + """ + Citations content block returned by Nova grounding (web search) tool. + + When Nova grounding is enabled via systemTool, the model may return + citationsContent blocks containing web search citation references. + + Reference: https://docs.aws.amazon.com/nova/latest/userguide/grounding.html + + Example response structure: + { + "citationsContent": { + "citations": [ + { + "location": { + "web": { + "url": "https://example.com/article", + "domain": "example.com" + } + } + } + ] + } + } + """ + + citations: List[CitationReferenceBlock] + + class ContentBlock(TypedDict, total=False): text: str image: ImageBlock @@ -103,6 +164,7 @@ class ContentBlock(TypedDict, total=False): cachePoint: CachePointBlock reasoningContent: BedrockConverseReasoningContentBlock guardContent: GuardrailConverseContentBlock + citationsContent: CitationsContentBlock class MessageBlock(TypedDict): @@ -159,8 +221,24 @@ class ToolSpecBlock(TypedDict, total=False): description: str +class SystemToolBlock(TypedDict, total=False): + """ + System tool block for Nova grounding and other built-in tools. + + Example: + { + "systemTool": { + "name": "nova_grounding" + } + } + """ + + name: Required[str] + + class ToolBlock(TypedDict, total=False): toolSpec: Optional[ToolSpecBlock] + systemTool: Optional[SystemToolBlock] cachePoint: Optional[CachePointBlock] @@ -210,11 +288,13 @@ class ContentBlockStartEvent(TypedDict, total=False): class ContentBlockDeltaEvent(TypedDict, total=False): """ Either 'text' or 'toolUse' will be specified for Converse API streaming response. + May also include 'citationsContent' when Nova grounding is enabled. """ text: str toolUse: ToolBlockDeltaEvent reasoningContent: BedrockConverseReasoningContentBlockDelta + citationsContent: CitationsContentBlock class PerformanceConfigBlock(TypedDict): @@ -879,3 +959,8 @@ class BedrockGetBatchResponse(TypedDict, total=False): outputDataConfig: BedrockOutputDataConfig timeoutDurationInHours: Optional[int] clientRequestToken: Optional[str] + +class BedrockToolBlock(TypedDict, total=False): + toolSpec: Optional[ToolSpecBlock] + systemTool: Optional[SystemToolBlock] # For Nova grounding + cachePoint: Optional[CachePointBlock] diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index cbca58e651..96d78cf882 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -51,19 +51,20 @@ class GenericGuardrailAPIRequest(BaseModel): """Request model for the Generic Guardrail API""" input_type: Literal["request", "response"] - litellm_call_id: Optional[str] # the call id of the individual LLM call + litellm_call_id: Optional[str] = None # the call id of the individual LLM call litellm_trace_id: Optional[ str - ] # the trace id of the LLM call - useful if there are multiple LLM calls for the same conversation - structured_messages: Optional[List[AllMessageValues]] - images: Optional[List[str]] - tools: Optional[List[ChatCompletionToolParam]] - texts: Optional[List[str]] + ] = None # the trace id of the LLM call - useful if there are multiple LLM calls for the same conversation + structured_messages: Optional[List[AllMessageValues]] = None + images: Optional[List[str]] = None + tools: Optional[List[ChatCompletionToolParam]] = None + texts: Optional[List[str]] = None request_data: GenericGuardrailAPIMetadata - additional_provider_specific_params: Optional[Dict[str, Any]] + additional_provider_specific_params: Optional[Dict[str, Any]] = None tool_calls: Optional[ Union[List[ChatCompletionToolCallChunk], List[ChatCompletionMessageToolCall]] - ] + ] = None + model: Optional[str] = None # the model being used for the LLM call class GenericGuardrailAPIResponse: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/onyx.py b/litellm/types/proxy/guardrails/guardrail_hooks/onyx.py index aa5b9d7a3f..42d7e94829 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/onyx.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/onyx.py @@ -16,6 +16,11 @@ class OnyxGuardrailConfigModel(GuardrailConfigModel): description="The API key for the Onyx Guard server. If not provided, the `ONYX_API_KEY` environment variable is checked.", ) + timeout: Optional[float] = Field( + default=None, + description="The timeout for the Onyx Guard server in seconds. If not provided, the `ONYX_TIMEOUT` environment variable is checked.", + ) + @staticmethod def ui_friendly_name() -> str: return "Onyx Guardrail" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cd797dd1e5..2d67e13e92 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3,25 +3,26 @@ import time from enum import Enum from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Union -from aiohttp import FormData from openai._models import BaseModel as OpenAIObject -from openai.types.audio.transcription_create_params import FileTypes # type: ignore -from openai.types.chat.chat_completion import ChatCompletion +from openai.types.audio.transcription_create_params import FileTypes as FileTypes # type: ignore +from openai.types.chat.chat_completion import ChatCompletion as ChatCompletion from openai.types.completion_usage import ( CompletionTokensDetails, CompletionUsage, PromptTokensDetails, ) from openai.types.moderation import ( - Categories, - CategoryAppliedInputTypes, - CategoryScores, + Categories as Categories, + CategoryAppliedInputTypes as CategoryAppliedInputTypes, + CategoryScores as CategoryScores, +) +from openai.types.moderation_create_response import ( + Moderation as Moderation, + ModerationCreateResponse as ModerationCreateResponse, ) -from openai.types.moderation_create_response import Moderation, ModerationCreateResponse from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator -from typing_extensions import Callable, Dict, Required, TypedDict, override +from typing_extensions import Required, TypedDict -import litellm from litellm._uuid import uuid from litellm.types.llms.base import ( BaseLiteLLMOpenAIResponseObject, @@ -52,7 +53,7 @@ from .llms.openai import ( ResponsesAPIResponse, WebSearchOptions, ) -from .rerank import RerankResponse +from .rerank import RerankResponse as RerankResponse if TYPE_CHECKING: from .vector_stores import VectorStoreSearchResponse @@ -1411,7 +1412,7 @@ class Usage(SafeAttributeModel, CompletionUsage): prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None """Breakdown of tokens used in the prompt.""" - def __init__( + def __init__( # noqa: PLR0915 self, prompt_tokens: Optional[int] = None, completion_tokens: Optional[int] = None, @@ -2501,6 +2502,7 @@ class StandardLoggingMetadata(StandardLoggingUserAPIKeyMetadata): dict ] # special param to log k,v pairs to spendlogs for a call requester_ip_address: Optional[str] + user_agent: Optional[str] requester_metadata: Optional[dict] requester_custom_headers: Optional[ Dict[str, str] @@ -2686,6 +2688,7 @@ class StandardLoggingPayload(TypedDict): request_tags: list end_user: Optional[str] requester_ip_address: Optional[str] + user_agent: Optional[str] messages: Optional[Union[str, list, dict]] response: Optional[Union[str, list, dict]] error_str: Optional[str] @@ -3428,3 +3431,4 @@ class GenericGuardrailAPIInputs(TypedDict, total=False): structured_messages: List[ AllMessageValues ] # structured messages sent to the LLM - indicates if text is from system or user + model: Optional[str] # the model being used for the LLM call diff --git a/litellm/utils.py b/litellm/utils.py index 584ab8805a..bb95be05b5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8053,6 +8053,12 @@ class ProviderConfigManager: ) return OpenrouterEmbeddingConfig() + elif litellm.LlmProviders.VERCEL_AI_GATEWAY == provider: + from litellm.llms.vercel_ai_gateway.embedding.transformation import ( + VercelAIGatewayEmbeddingConfig, + ) + + return VercelAIGatewayEmbeddingConfig() elif litellm.LlmProviders.GIGACHAT == provider: return litellm.GigaChatEmbeddingConfig() elif litellm.LlmProviders.SAGEMAKER == provider: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d0a3bec900..421f2b8b4d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -10231,6 +10231,48 @@ "mode": "completion", "output_cost_per_token": 5e-07 }, + "deepseek-v3-2-251201": { + "input_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "max_input_tokens": 98304, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "glm-4-7-251222": { + "input_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "max_input_tokens": 204800, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 0.0, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "kimi-k2-thinking-251104": { + "input_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "max_input_tokens": 229376, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "doubao-embedding": { "input_cost_per_token": 0.0, "litellm_provider": "volcengine", diff --git a/pyproject.toml b/pyproject.toml index cae519b480..1e4836ed19 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -145,6 +145,7 @@ mypy = "^1.0" pytest = "^7.4.3" pytest-mock = "^3.12.0" pytest-asyncio = "^0.21.1" +pytest-retry = "^1.6.3" requests-mock = "^1.12.1" responses = "^0.25.7" respx = "^0.22.0" @@ -183,6 +184,8 @@ plugins = "pydantic.mypy" [tool.pytest.ini_options] asyncio_mode = "auto" +retries = 20 +retry_delay = 5 markers = [ "asyncio: mark test as an asyncio test", "limit_leaks: mark test with memory limit for leak detection (e.g., '40 MB')", diff --git a/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py new file mode 100644 index 0000000000..ac2c945baa --- /dev/null +++ b/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -0,0 +1,1108 @@ +"""Tests for MCP OAuth discoverable endpoints""" +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + + +@pytest.mark.asyncio +async def test_authorize_endpoint_includes_response_type(): + """Test that authorize endpoint includes response_type=code parameter (fixes #15684)""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server + oauth2_server = MCPServer( + server_id="test_oauth_server", + name="test_oauth", + server_name="test_oauth", + alias="test_oauth", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="test_client_id", + client_secret="test_client_secret", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + scopes=["read", "write"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # Mock request + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + # Mock the encryption functions to avoid needing a signing key + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" + ) as mock_encrypt: + mock_encrypt.return_value = "mocked_encrypted_state" + + # Call authorize endpoint + response = await authorize( + request=mock_request, + client_id="test_client_id", + mcp_server_name="test_oauth", + redirect_uri="https://client.example.com/callback", + state="test_state", + ) + + # Verify response is a redirect + assert response.status_code == 307 # FastAPI RedirectResponse default + + # Verify response_type is in the redirect URL + assert "response_type=code" in response.headers["location"] + assert "https://provider.com/oauth/authorize" in response.headers["location"] + assert "client_id=test_client_id" in response.headers["location"] + assert "scope=read+write" in response.headers["location"] + + +@pytest.mark.asyncio +async def test_authorize_endpoint_forwards_pkce_parameters(): + """Test that authorize endpoint forwards PKCE parameters (code_challenge and code_challenge_method)""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server (simulating Google OAuth) + oauth2_server = MCPServer( + server_id="google_mcp", + name="google_mcp", + server_name="google_mcp", + alias="google_mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="669428968603-test.apps.googleusercontent.com", + client_secret="GOCSPX-test_secret", + authorization_url="https://accounts.google.com/o/oauth2/v2/auth", + token_url="https://oauth2.googleapis.com/token", + scopes=["https://www.googleapis.com/auth/drive", "openid", "email"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # Mock request + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm-proxy.example.com/" + mock_request.headers = {} + + # Mock the encryption function + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" + ) as mock_encrypt: + mock_encrypt.return_value = "mocked_encrypted_state_with_pkce" + + # Call authorize endpoint with PKCE parameters + response = await authorize( + request=mock_request, + client_id="669428968603-test.apps.googleusercontent.com", + mcp_server_name="google_mcp", + redirect_uri="http://localhost:60108/callback", + state="test_client_state", + code_challenge="x6YH_qgwbvOzbsHDuL1sW9gYkR9-gObUiIB5RkPwxDk", + code_challenge_method="S256", + ) + + # Verify response is a redirect + assert response.status_code == 307 + + # Verify PKCE parameters are included in the redirect URL + location = response.headers["location"] + assert "https://accounts.google.com/o/oauth2/v2/auth" in location + assert "code_challenge=x6YH_qgwbvOzbsHDuL1sW9gYkR9-gObUiIB5RkPwxDk" in location + assert "code_challenge_method=S256" in location + assert "client_id=669428968603-test.apps.googleusercontent.com" in location + assert "response_type=code" in location + + +@pytest.mark.asyncio +async def test_token_endpoint_forwards_code_verifier(): + """Test that token endpoint forwards code_verifier for PKCE flow""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + token_endpoint, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + import httpx + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server + oauth2_server = MCPServer( + server_id="google_mcp", + name="google_mcp", + server_name="google_mcp", + alias="google_mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="669428968603-test.apps.googleusercontent.com", + client_secret="GOCSPX-test_secret", + authorization_url="https://accounts.google.com/o/oauth2/v2/auth", + token_url="https://oauth2.googleapis.com/token", + scopes=["https://www.googleapis.com/auth/drive", "openid", "email"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # Mock request + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm-proxy.example.com/" + mock_request.headers = {} + + # Mock httpx client response + mock_response = MagicMock() + mock_response.json.return_value = { + "access_token": "ya29.test_access_token", + "token_type": "Bearer", + "expires_in": 3599, + "scope": "openid email https://www.googleapis.com/auth/drive", + } + mock_response.raise_for_status = MagicMock() + + # Mock the async httpx client with AsyncMock for async methods + from unittest.mock import AsyncMock + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" + ) as mock_get_client: + mock_async_client = MagicMock() + # Use AsyncMock for the async post method + mock_async_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_async_client + + # Call token endpoint with code_verifier + response = await token_endpoint( + request=mock_request, + grant_type="authorization_code", + code="4/test_authorization_code", + redirect_uri="http://localhost:60108/callback", + client_id="669428968603-test.apps.googleusercontent.com", + mcp_server_name="google_mcp", + client_secret="GOCSPX-test_secret", + code_verifier="test_code_verifier_from_client", + ) + + # Verify that the token endpoint was called with code_verifier + mock_async_client.post.assert_called_once() + call_args = mock_async_client.post.call_args + + # Check the data parameter includes code_verifier + assert call_args[1]["data"]["code_verifier"] == "test_code_verifier_from_client" + assert call_args[1]["data"]["code"] == "4/test_authorization_code" + assert ( + call_args[1]["data"]["client_id"] + == "669428968603-test.apps.googleusercontent.com" + ) + assert call_args[1]["data"]["client_secret"] == "GOCSPX-test_secret" + assert call_args[1]["data"]["grant_type"] == "authorization_code" + + # Verify response + response_data = response.body + import json + + token_data = json.loads(response_data) + assert token_data["access_token"] == "ya29.test_access_token" + assert token_data["token_type"] == "Bearer" + + +@pytest.mark.asyncio +async def test_register_client_without_mcp_server_name_returns_dummy(): + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client, + ) + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", + new=AsyncMock(return_value={}), + ): + result = await register_client(request=mock_request) + + assert result == { + "client_id": "dummy_client", + "client_secret": "dummy", + "redirect_uris": ["https://proxy.litellm.example/callback"], + } + + +@pytest.mark.asyncio +async def test_register_client_returns_existing_server_credentials(): + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + oauth2_server = MCPServer( + server_id="stored_server", + name="stored_server", + server_name="stored_server", + alias="stored_server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="existing-client", + client_secret="existing-secret", + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + + try: + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", + new=AsyncMock(return_value={}), + ): + result = await register_client( + request=mock_request, mcp_server_name=oauth2_server.server_name + ) + finally: + global_mcp_server_manager.registry.clear() + + assert result == { + "client_id": "stored_server", + "client_secret": "dummy", + "redirect_uris": ["https://proxy.litellm.example/callback"], + } + + +@pytest.mark.asyncio +async def test_register_client_remote_registration_success(): + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + oauth2_server = MCPServer( + server_id="remote_server", + name="remote_server", + server_name="remote_server", + alias="remote_server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + + request_payload = { + "client_name": "Litellm Proxy", + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + "token_endpoint_auth_method": "client_secret_post", + } + + mock_response = MagicMock() + mock_response.json.return_value = { + "client_id": "generated-client", + "client_secret": "generated-secret", + } + mock_response.raise_for_status = MagicMock() + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + try: + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", + new=AsyncMock(return_value=request_payload), + ), patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ): + response = await register_client( + request=mock_request, mcp_server_name=oauth2_server.server_name + ) + finally: + global_mcp_server_manager.registry.clear() + + import json + + assert response.status_code == 200 + payload = json.loads(response.body.decode("utf-8")) + assert payload == mock_response.json.return_value + + mock_async_client.post.assert_called_once() + call_args = mock_async_client.post.call_args + assert call_args.args[0] == oauth2_server.registration_url + assert call_args.kwargs["headers"] == { + "Content-Type": "application/json", + "Accept": "application/json", + } + assert call_args.kwargs["json"]["redirect_uris"] == [ + "https://proxy.litellm.example/callback" + ] + assert call_args.kwargs["json"]["grant_types"] == request_payload["grant_types"] + assert ( + call_args.kwargs["json"]["token_endpoint_auth_method"] + == request_payload["token_endpoint_auth_method"] + ) + + +@pytest.mark.asyncio +async def test_authorize_endpoint_respects_x_forwarded_proto(): + """Test that authorize endpoint uses X-Forwarded-Proto header to construct correct redirect_uri""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server + oauth2_server = MCPServer( + server_id="test_oauth_server", + name="test_oauth", + server_name="test_oauth", + alias="test_oauth", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="test_client_id", + client_secret="test_client_secret", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + scopes=["read", "write"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # Mock request with http base_url but X-Forwarded-Proto: https + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://litellm.example.com/" # HTTP + mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + + # Mock the encryption functions + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" + ) as mock_encrypt: + mock_encrypt.return_value = "mocked_encrypted_state" + + # Call authorize endpoint + response = await authorize( + request=mock_request, + client_id="test_client_id", + mcp_server_name="test_oauth", + redirect_uri="https://client.example.com/callback", + state="test_state", + ) + + # Verify redirect URL uses HTTPS in the redirect_uri parameter + location = response.headers["location"] + + # The redirect_uri parameter sent to the OAuth provider should use HTTPS + assert ( + "redirect_uri=https%3A%2F%2Flitellm.example.com%2Fcallback" in location + or "redirect_uri=https://litellm.example.com/callback" in location + ) + + +@pytest.mark.asyncio +async def test_token_endpoint_respects_x_forwarded_proto(): + """Test that token endpoint uses X-Forwarded-Proto header for redirect_uri""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + token_endpoint, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server + oauth2_server = MCPServer( + server_id="google_mcp", + name="google_mcp", + server_name="google_mcp", + alias="google_mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="test_client_id", + client_secret="test_secret", + authorization_url="https://accounts.google.com/o/oauth2/v2/auth", + token_url="https://oauth2.googleapis.com/token", + scopes=["openid", "email"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # Mock request with http base_url but X-Forwarded-Proto: https + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://litellm-proxy.example.com/" # HTTP + mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + + # Mock httpx client response + mock_response = MagicMock() + mock_response.json.return_value = { + "access_token": "test_token", + "token_type": "Bearer", + "expires_in": 3599, + } + mock_response.raise_for_status = MagicMock() + + # Mock the async httpx client + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = mock_async_client + + # Call token endpoint + response = await token_endpoint( + request=mock_request, + grant_type="authorization_code", + code="test_code", + redirect_uri="http://localhost:60108/callback", + client_id="test_client_id", + mcp_server_name="google_mcp", + client_secret="test_secret", + ) + + # Verify that the redirect_uri sent to the provider uses HTTPS + call_args = mock_async_client.post.call_args + assert ( + call_args[1]["data"]["redirect_uri"] + == "https://litellm-proxy.example.com/callback" + ) + + +@pytest.mark.asyncio +async def test_oauth_protected_resource_standard_pattern(): + """Test that oauth_protected_resource_mcp_standard returns standard MCP URL pattern (/mcp/{server_name})""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + oauth_protected_resource_mcp_standard, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server + oauth2_server = MCPServer( + server_id="test_server", + name="test_server", + server_name="test_server", + alias="test_server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="test_client_id", + client_secret="test_client_secret", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + scopes=["read", "write"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # Mock request + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + # Call the standard pattern endpoint + response = await oauth_protected_resource_mcp_standard( + request=mock_request, + mcp_server_name="test_server", + ) + + # Verify response uses standard MCP pattern: /mcp/{server_name} + assert response["resource"] == "https://litellm.example.com/mcp/test_server" + assert response["authorization_servers"][0] == "https://litellm.example.com/test_server" + assert response["scopes_supported"] == oauth2_server.scopes + + +@pytest.mark.asyncio +async def test_oauth_protected_resource_legacy_pattern(): + """Test that oauth_protected_resource_mcp returns legacy URL pattern (/{server_name}/mcp)""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + oauth_protected_resource_mcp, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server + oauth2_server = MCPServer( + server_id="test_server", + name="test_server", + server_name="test_server", + alias="test_server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="test_client_id", + client_secret="test_client_secret", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + scopes=["read", "write"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # Mock request + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + # Call the legacy pattern endpoint + response = await oauth_protected_resource_mcp( + request=mock_request, + mcp_server_name="test_server", + ) + + # Verify response uses legacy pattern: /{server_name}/mcp + assert response["resource"] == "https://litellm.example.com/test_server/mcp" + assert response["authorization_servers"][0] == "https://litellm.example.com/test_server" + assert response["scopes_supported"] == oauth2_server.scopes + + +@pytest.mark.asyncio +async def test_oauth_protected_resource_respects_x_forwarded_proto(): + """Test that oauth_protected_resource_mcp uses X-Forwarded-Proto for URLs""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + oauth_protected_resource_mcp, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server + oauth2_server = MCPServer( + server_id="test_oauth_server", + name="test_oauth", + server_name="test_oauth", + alias="test_oauth", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="test_client_id", + client_secret="test_client_secret", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + scopes=["read", "write"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # Mock request with http base_url but X-Forwarded-Proto: https + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://litellm.example.com/" # HTTP + mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + + # Call the endpoint + response = await oauth_protected_resource_mcp( + request=mock_request, + mcp_server_name="test_oauth", + ) + + # Verify response uses HTTPS URLs + assert response["authorization_servers"][0].startswith( + "https://litellm.example.com/" + ) + assert response["scopes_supported"] == oauth2_server.scopes + + +@pytest.mark.asyncio +async def test_oauth_authorization_server_respects_x_forwarded_proto(): + """Test that oauth_authorization_server_mcp uses X-Forwarded-Proto for URLs""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + oauth_authorization_server_mcp, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server + oauth2_server = MCPServer( + server_id="test_oauth_server", + name="test_oauth", + server_name="test_oauth", + alias="test_oauth", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="test_client_id", + client_secret="test_client_secret", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + scopes=["read", "write"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # Mock request with http base_url but X-Forwarded-Proto: https + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://litellm.example.com/" # HTTP + mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + + # Call the endpoint + response = await oauth_authorization_server_mcp( + request=mock_request, + mcp_server_name="test_oauth", + ) + + # Verify response uses HTTPS URLs + assert response["authorization_endpoint"].startswith("https://litellm.example.com/") + assert response["token_endpoint"].startswith("https://litellm.example.com/") + assert response["registration_endpoint"].startswith("https://litellm.example.com/") + assert response["grant_types_supported"] == ["authorization_code", "refresh_token"] + assert response["scopes_supported"] == oauth2_server.scopes + + +@pytest.mark.asyncio +async def test_register_client_respects_x_forwarded_proto(): + """Test that register_client uses X-Forwarded-Proto for redirect_uris""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client, + ) + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Mock request with http base_url but X-Forwarded-Proto: https + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://proxy.litellm.example/" # HTTP + mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", + new=AsyncMock(return_value={}), + ): + result = await register_client(request=mock_request) + + # Verify the redirect_uris use HTTPS + assert result == { + "client_id": "dummy_client", + "client_secret": "dummy", + "redirect_uris": ["https://proxy.litellm.example/callback"], + } + + +@pytest.mark.asyncio +async def test_authorize_endpoint_respects_x_forwarded_host(): + """Test that authorize endpoint uses X-Forwarded-Host and X-Forwarded-Proto to construct correct redirect_uri""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server + oauth2_server = MCPServer( + server_id="test_oauth_server", + name="test_oauth", + server_name="test_oauth", + alias="test_oauth", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="test_client_id", + client_secret="test_client_secret", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + scopes=["read", "write"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # Mock request simulating nginx proxy: + # Internal: http://localhost:8888/github/mcp + # External: https://proxy.example.com/github/mcp + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:8888/github/mcp" + mock_request.headers = { + "X-Forwarded-Proto": "https", + "X-Forwarded-Host": "proxy.example.com", + } + + # Mock the encryption functions + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" + ) as mock_encrypt: + mock_encrypt.return_value = "mocked_encrypted_state" + + # Call authorize endpoint + response = await authorize( + request=mock_request, + client_id="test_client_id", + mcp_server_name="test_oauth", + redirect_uri="https://client.example.com/callback", + state="test_state", + ) + + # Verify redirect URL uses the forwarded host and scheme + location = response.headers["location"] + + # The redirect_uri parameter should use the external URL + assert ( + "redirect_uri=https%3A%2F%2Fproxy.example.com%2Fgithub%2Fmcp%2Fcallback" + in location + or "redirect_uri=https://proxy.example.com/github/mcp/callback" in location + ) + + +@pytest.mark.asyncio +async def test_token_endpoint_respects_x_forwarded_host(): + """Test that token endpoint uses X-Forwarded-Host and X-Forwarded-Proto for redirect_uri""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + token_endpoint, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server + oauth2_server = MCPServer( + server_id="google_mcp", + name="google_mcp", + server_name="google_mcp", + alias="google_mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="test_client_id", + client_secret="test_secret", + authorization_url="https://accounts.google.com/o/oauth2/v2/auth", + token_url="https://oauth2.googleapis.com/token", + scopes=["openid", "email"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # Mock request simulating nginx proxy without port in host + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:8888/github/mcp" + mock_request.headers = { + "X-Forwarded-Proto": "https", + "X-Forwarded-Host": "proxy.example.com", + } + + # Mock httpx client response + mock_response = MagicMock() + mock_response.json.return_value = { + "access_token": "test_token", + "token_type": "Bearer", + "expires_in": 3599, + } + mock_response.raise_for_status = MagicMock() + + # Mock the async httpx client + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = mock_async_client + + # Call token endpoint + response = await token_endpoint( + request=mock_request, + grant_type="authorization_code", + code="test_code", + redirect_uri="http://localhost:60108/callback", + client_id="test_client_id", + mcp_server_name="google_mcp", + client_secret="test_secret", + ) + + # Verify that the redirect_uri sent to the provider uses the external URL + call_args = mock_async_client.post.call_args + assert ( + call_args[1]["data"]["redirect_uri"] + == "https://proxy.example.com/github/mcp/callback" + ) + + +@pytest.mark.parametrize( + "base_url,x_forwarded_proto,x_forwarded_host,x_forwarded_port,expected_url", + [ + # Case 1: No forwarded headers - use original URL as-is (no trailing slash) + ( + "http://localhost:4000/", + None, + None, + None, + "http://localhost:4000", + ), + # Case 2: Only X-Forwarded-Proto - change scheme only + ( + "http://localhost:4000/", + "https", + None, + None, + "https://localhost:4000", + ), + # Case 3: X-Forwarded-Proto + X-Forwarded-Host - change scheme and host + ( + "http://localhost:4000/", + "https", + "proxy.example.com", + None, + "https://proxy.example.com", + ), + # Case 4: X-Forwarded-Host with port included in host header + ( + "http://localhost:4000/", + "https", + "proxy.example.com:8080", + None, + "https://proxy.example.com:8080", + ), + # Case 5: X-Forwarded-Host + X-Forwarded-Port as separate headers + ( + "http://localhost:4000/", + "https", + "proxy.example.com", + "8443", + "https://proxy.example.com:8443", + ), + # Case 6: Only X-Forwarded-Host without proto - use original scheme + ( + "http://localhost:4000/", + None, + "proxy.example.com", + None, + "http://proxy.example.com", + ), + # Case 7: Only X-Forwarded-Port without host - preserves original port if present + # (This is safer behavior - X-Forwarded-Port alone is unusual) + ( + "http://localhost:4000/", + None, + None, + "8443", + "http://localhost:4000", # Original port preserved when already present + ), + # Case 8: Complex internal URL with path (path is preserved) + ( + "http://localhost:8888/github/mcp", + "https", + "proxy.example.com", + None, + "https://proxy.example.com/github/mcp", + ), + # Case 9: IPv6 address in X-Forwarded-Host (should not treat :: as port separator) + ( + "http://localhost:4000/", + "https", + "[2001:db8::1]", + None, + "https://[2001:db8::1]", + ), + # Case 10: IPv6 address with port + ( + "http://localhost:4000/", + "https", + "[2001:db8::1]:8080", + None, + "https://[2001:db8::1]:8080", + ), + # Case 11: X-Forwarded-Host already has port, X-Forwarded-Port also provided (host wins) + ( + "http://localhost:4000/", + "https", + "proxy.example.com:9000", + "8443", + "https://proxy.example.com:9000", + ), + # Case 12: Standard proxy setup (most common case) + ( + "http://127.0.0.1:8888/", + "https", + "chatproxy.company.com", + None, + "https://chatproxy.company.com", + ), + # Case 13: Internal URL already has port, X-Forwarded-Port does NOT override + # (safer behavior - preserves original port when X-Forwarded-Host not provided) + ( + "http://localhost:4000/", + None, + None, + "443", + "http://localhost:4000", # Original port preserved + ), + # Case 14: Original URL with existing port in netloc, X-Forwarded-Host replaces it + ( + "http://internal.local:8888/", + "https", + "external.com", + None, + "https://external.com", + ), + ], +) +def test_get_request_base_url_comprehensive( + base_url, x_forwarded_proto, x_forwarded_host, x_forwarded_port, expected_url +): + """Comprehensive test for get_request_base_url with various header combinations""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + get_request_base_url, + ) + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Create mock request + mock_request = MagicMock(spec=Request) + mock_request.base_url = base_url + + # Build headers dict + headers = {} + if x_forwarded_proto: + headers["X-Forwarded-Proto"] = x_forwarded_proto + if x_forwarded_host: + headers["X-Forwarded-Host"] = x_forwarded_host + if x_forwarded_port: + headers["X-Forwarded-Port"] = x_forwarded_port + + # Mock headers.get() to return our test values + def mock_get(header_name, default=None): + return headers.get(header_name, default) + + mock_request.headers.get = mock_get + + # Test the function + result = get_request_base_url(mock_request) + + # Verify result + assert result == expected_url, ( + f"Expected '{expected_url}' but got '{result}'\n" + f"Input: base_url={base_url}, " + f"X-Forwarded-Proto={x_forwarded_proto}, " + f"X-Forwarded-Host={x_forwarded_host}, " + f"X-Forwarded-Port={x_forwarded_port}" + ) diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index 7c0db41d13..d48dd1bfd9 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -3954,3 +3954,288 @@ def test_bedrock_openai_error_handling(): assert exc_info.value.status_code == 422 print("✓ Error handling works correctly") + +# ============================================================================ +# Nova Grounding (web_search_options) Unit Tests (Mocked) +# ============================================================================ + +def test_bedrock_nova_grounding_web_search_options_non_streaming(): + """ + Unit test for Nova grounding using web_search_options parameter (non-streaming). + + This test mocks the HTTP call to verify: + 1. web_search_options is correctly mapped to systemTool for Nova models + 2. The request structure is correct + + Related: https://docs.aws.amazon.com/nova/latest/userguide/grounding.html + """ + from unittest.mock import patch, MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler() + + messages = [ + { + "role": "user", + "content": "What is the current population of Tokyo, Japan?", + } + ] + + with patch.object(client, "post") as mock_post: + try: + completion( + model="us.amazon.nova-pro-v1:0", # No bedrock/ prefix when using api_base + messages=messages, + web_search_options={}, # Enables Nova grounding + max_tokens=500, + client=client, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com", + ) + except Exception: + pass # Expected - we're just checking the request structure + + # Verify the request was made correctly + if mock_post.called: + request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}")) + print(f"Request body: {json.dumps(request_body, indent=2)}") + + # Verify toolConfig is present with systemTool + assert "toolConfig" in request_body, "toolConfig should be in request" + tool_config = request_body["toolConfig"] + assert "tools" in tool_config, "tools should be in toolConfig" + + # Find the systemTool for nova_grounding + system_tool_found = False + for tool in tool_config["tools"]: + if "systemTool" in tool: + assert tool["systemTool"]["name"] == "nova_grounding" + system_tool_found = True + break + + assert system_tool_found, "systemTool with nova_grounding should be present" + print(f"✓ web_search_options correctly transformed to systemTool (non-streaming)") + + +def test_bedrock_nova_grounding_with_function_tools(): + """ + Unit test for Nova grounding combined with regular function tools. + + This tests the scenario where users want both web grounding AND + custom function calling capabilities. + """ + from unittest.mock import patch + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler() + + # Regular function tool + tools = [ + { + "type": "function", + "function": { + "name": "get_stock_price", + "description": "Get the current stock price for a given ticker symbol", + "parameters": { + "type": "object", + "properties": { + "ticker": { + "type": "string", + "description": "The stock ticker symbol, e.g. AAPL, GOOGL", + } + }, + "required": ["ticker"], + }, + }, + } + ] + + messages = [ + { + "role": "user", + "content": "What is the current market cap of Apple Inc?", + } + ] + + with patch.object(client, "post") as mock_post: + try: + completion( + model="us.amazon.nova-pro-v1:0", # No bedrock/ prefix when using api_base + messages=messages, + tools=tools, + web_search_options={}, # Also enable web grounding + max_tokens=500, + client=client, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com", + ) + except Exception: + pass # Expected - we're just checking the request structure + + # Verify the request was made correctly + if mock_post.called: + request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}")) + print(f"Request body: {json.dumps(request_body, indent=2)}") + + # Verify toolConfig has both function tool and systemTool + assert "toolConfig" in request_body, "toolConfig should be in request" + tool_config = request_body["toolConfig"] + assert "tools" in tool_config, "tools should be in toolConfig" + + tools_in_request = tool_config["tools"] + + # Should have both the function tool and the systemTool + function_tool_found = False + system_tool_found = False + + for tool in tools_in_request: + if "toolSpec" in tool: + assert tool["toolSpec"]["name"] == "get_stock_price" + function_tool_found = True + if "systemTool" in tool: + assert tool["systemTool"]["name"] == "nova_grounding" + system_tool_found = True + + assert function_tool_found, "Function tool (get_stock_price) should be present" + assert system_tool_found, "systemTool (nova_grounding) should be present" + print(f"✓ Both function tools and web_search_options correctly combined") + + +@pytest.mark.asyncio +async def test_bedrock_nova_grounding_async(): + """ + Async unit test for Nova grounding via web_search_options. + + This test verifies the request transformation for async calls. + """ + from unittest.mock import patch, AsyncMock + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + client = AsyncHTTPHandler() + + messages = [ + { + "role": "user", + "content": "What is the weather forecast for New York City today?", + } + ] + + with patch.object(client, "post", new=AsyncMock()) as mock_post: + try: + await litellm.acompletion( + model="us.amazon.nova-pro-v1:0", # No bedrock/ prefix when using api_base + messages=messages, + web_search_options={}, + max_tokens=500, + client=client, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com", + ) + except Exception: + pass # Expected - we're just checking the request structure + + # Verify the request was made correctly + if mock_post.called: + request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}")) + print(f"Request body: {json.dumps(request_body, indent=2)}") + + # Verify toolConfig is present with systemTool + assert "toolConfig" in request_body, "toolConfig should be in request" + tool_config = request_body["toolConfig"] + assert "tools" in tool_config, "tools should be in toolConfig" + + # Find the systemTool for nova_grounding + system_tool_found = False + for tool in tool_config["tools"]: + if "systemTool" in tool: + assert tool["systemTool"]["name"] == "nova_grounding" + system_tool_found = True + break + + assert system_tool_found, "systemTool with nova_grounding should be present" + print(f"✓ Async web_search_options correctly transformed to systemTool") + + +def test_bedrock_nova_web_search_options_ignored_for_non_nova(): + """ + Test that web_search_options is ignored for non-Nova Bedrock models. + + Nova grounding is only supported on Nova models. For other models, + the parameter should be silently ignored. + """ + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + config = AmazonConverseConfig() + + # Should return None for non-Nova models + result = config._map_web_search_options({}, "anthropic.claude-3-sonnet-v1") + assert result is None + + result = config._map_web_search_options({}, "amazon.titan-text-express-v1") + assert result is None + + # Should return systemTool for Nova models + result = config._map_web_search_options({}, "amazon.nova-pro-v1:0") + assert result is not None + system_tool = result.get("systemTool") + assert system_tool is not None + assert system_tool["name"] == "nova_grounding" + + result2 = config._map_web_search_options({}, "us.amazon.nova-premier-v1:0") + assert result2 is not None + system_tool2 = result2.get("systemTool") + assert system_tool2 is not None + assert system_tool2["name"] == "nova_grounding" + + +def test_bedrock_nova_grounding_request_transformation(): + """ + Unit test to verify that web_search_options transforms to systemTool in the request. + """ + from unittest.mock import patch, MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler() + + messages = [{"role": "user", "content": "What is the population of Tokyo?"}] + + with patch.object(client, "post") as mock_post: + mock_post.return_value = MagicMock( + status_code=200, + json=lambda: { + "output": {"message": {"role": "assistant", "content": [{"text": "Test"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 10, "outputTokens": 5} + } + ) + + try: + response = completion( + model="bedrock/us.amazon.nova-pro-v1:0", + messages=messages, + web_search_options={}, + max_tokens=100, + client=client, + ) + except Exception: + pass # Expected - we're just checking the request + + if mock_post.called: + request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}")) + print(f"Request body: {json.dumps(request_body, indent=2)}") + + # Verify toolConfig is present with systemTool + assert "toolConfig" in request_body, "toolConfig should be in request" + + tool_config = request_body["toolConfig"] + assert "tools" in tool_config, "tools should be in toolConfig" + + tools_in_request = tool_config["tools"] + + # Find the systemTool + system_tool_found = False + for tool in tools_in_request: + if "systemTool" in tool: + assert tool["systemTool"]["name"] == "nova_grounding" + system_tool_found = True + break + + assert system_tool_found, "systemTool with nova_grounding should be present" + print("✓ web_search_options correctly transformed to systemTool") diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/mcp_tests/test_aresponses_api_with_mcp.py index 6c8f51201d..a7bbfef14a 100644 --- a/tests/mcp_tests/test_aresponses_api_with_mcp.py +++ b/tests/mcp_tests/test_aresponses_api_with_mcp.py @@ -683,9 +683,11 @@ async def test_streaming_responses_api_with_mcp_tools( Return the user the result of request 2 """ - # Skip test if ANTHROPIC_API_KEY is not set for anthropic models - if "anthropic" in model.lower() and not os.getenv("ANTHROPIC_API_KEY"): + # Skip test if required API keys are not set + if ("anthropic" in model.lower() or "claude" in model.lower()) and not os.getenv("ANTHROPIC_API_KEY"): pytest.skip("ANTHROPIC_API_KEY not set, skipping anthropic model test") + if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv("OPENAI_API_KEY"): + pytest.skip("OPENAI_API_KEY not set, skipping openai model test") from unittest.mock import AsyncMock, patch diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index f56e8321b5..5785ff750f 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1105,14 +1105,24 @@ async def test_mcp_server_manager_config_integration_with_database(): test_manager.get_allowed_mcp_servers = mock_get_allowed_servers - # Mock _create_mcp_client to return a client that completes immediately - # This avoids network calls while preserving the actual conversion logic - def mock_create_mcp_client(*args, **kwargs): - mock_client = MagicMock() - mock_client.run_with_session = AsyncMock(return_value="ok") - return mock_client - - test_manager._create_mcp_client = mock_create_mcp_client + # Mock health_check_server to avoid real network calls that timeout + async def mock_health_check(server_id: str, mcp_auth_header=None): + server = test_manager.get_mcp_server_by_id(server_id) + if not server: + return None + return LiteLLM_MCPServerTable( + server_id=server_id, + server_name=server.name, + url=server.url, + transport=server.transport, + description=server.mcp_info.get("description") if server.mcp_info else None, + mcp_access_groups=server.access_groups, + status="healthy", + last_health_check=datetime.datetime.now(), + mcp_info=server.mcp_info, + ) + + test_manager.health_check_server = mock_health_check # Test the method (this tests our second fix) servers_list = await test_manager.get_all_mcp_servers_with_health_and_teams( diff --git a/tests/proxy_unit_tests/test_get_image.py b/tests/proxy_unit_tests/test_get_image.py new file mode 100644 index 0000000000..ad8c267275 --- /dev/null +++ b/tests/proxy_unit_tests/test_get_image.py @@ -0,0 +1,89 @@ +import os +import sys +from unittest import mock + +# Standard path insertion +sys.path.insert(0, os.path.abspath("../..")) + +import pytest +import httpx +from litellm.proxy.proxy_server import app + + +@pytest.mark.asyncio +async def test_get_image_error_handling(): + """ + Test that get_image handles network errors gracefully and doesn't hang. + """ + # Set an unreachable URL + os.environ["UI_LOGO_PATH"] = "http://invalid-url-12345.com/logo.jpg" + + # Clear cache + parent_dir = os.path.dirname( + os.path.dirname( + app.__file__ + if hasattr(app, "__file__") + else "litellm/proxy/proxy_server.py" + ) + ) + cache_path = os.path.join(parent_dir, "proxy", "cached_logo.jpg") + if os.path.exists(cache_path): + os.remove(cache_path) + + # Mock AsyncHTTPHandler to simulate a timeout or connection error + with mock.patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" + ) as mock_get: + mock_get.side_effect = httpx.ConnectError("Network is unreachable") + + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://testserver" + ) as ac: + response = await ac.get("/get_image") + + assert response.status_code == 200 + assert response.headers["content-type"] == "image/jpeg" + + +@pytest.mark.asyncio +async def test_get_image_cache_logic(): + """ + Test that once cached, get_image doesn't hit the network. + """ + os.environ["UI_LOGO_PATH"] = "http://example.com/logo.jpg" + + # Clear cache + parent_dir = os.path.dirname( + os.path.dirname( + app.__file__ + if hasattr(app, "__file__") + else "litellm/proxy/proxy_server.py" + ) + ) + cache_path = os.path.join(parent_dir, "proxy", "cached_logo.jpg") + if os.path.exists(cache_path): + os.remove(cache_path) + + # Mock response + mock_response = mock.Mock() + mock_response.status_code = 200 + mock_response.content = b"fake image data" + + with mock.patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" + ) as mock_get: + mock_get.return_value = mock_response + + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://testserver" + ) as ac: + # First call - should hit download logic + response1 = await ac.get("/get_image") + assert response1.status_code == 200 + assert mock_get.call_count == 1 + + # Second call - should hit cache + response2 = await ac.get("/get_image") + assert response2.status_code == 200 + # If cache works, mock_get shouldn't be called again + assert mock_get.call_count == 1 diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py index 31bcaff532..8db6d98f13 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py @@ -2,7 +2,9 @@ import os import sys import unittest.mock as mock +import httpx import pytest +import respx from httpx import Response sys.path.insert(0, os.path.abspath("../../..")) @@ -11,6 +13,8 @@ from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ( ResendEmailLogger, ) +# Test file for Resend email integration + @pytest.fixture def mock_env_vars(): @@ -37,7 +41,13 @@ def mock_httpx_client(): @pytest.mark.asyncio +@respx.mock async def test_send_email_success(mock_env_vars, mock_httpx_client): + # Block all HTTP requests at network level to prevent real API calls + respx.post("https://api.resend.com/emails").mock( + return_value=httpx.Response(200, json={"id": "test_email_id"}) + ) + # Initialize the logger logger = ResendEmailLogger() @@ -71,7 +81,13 @@ async def test_send_email_success(mock_env_vars, mock_httpx_client): @pytest.mark.asyncio +@respx.mock async def test_send_email_missing_api_key(mock_httpx_client): + # Block all HTTP requests at network level to prevent real API calls + respx.post("https://api.resend.com/emails").mock( + return_value=httpx.Response(200, json={"id": "test_email_id"}) + ) + # Remove the API key from environment before initializing logger original_key = os.environ.pop("RESEND_API_KEY", None) @@ -109,7 +125,13 @@ async def test_send_email_missing_api_key(mock_httpx_client): @pytest.mark.asyncio +@respx.mock async def test_send_email_multiple_recipients(mock_env_vars, mock_httpx_client): + # Block all HTTP requests at network level to prevent real API calls + respx.post("https://api.resend.com/emails").mock( + return_value=httpx.Response(200, json={"id": "test_email_id"}) + ) + # Initialize the logger logger = ResendEmailLogger() diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py index f8b5e40bf3..836b717bd6 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py @@ -2,7 +2,9 @@ import os import sys import unittest.mock as mock +import httpx import pytest +import respx from httpx import Response sys.path.insert(0, os.path.abspath("../../..")) @@ -101,7 +103,13 @@ async def test_send_email_missing_api_key(): @pytest.mark.asyncio +@respx.mock async def test_send_email_multiple_recipients(mock_env_vars, mock_httpx_client): + # Block all HTTP requests at network level to prevent real API calls + respx.post("https://api.sendgrid.com/v3/mail/send").mock( + return_value=httpx.Response(202, text="accepted") + ) + logger = SendGridEmailLogger() from_email = "test@example.com" diff --git a/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py new file mode 100644 index 0000000000..be2084969a --- /dev/null +++ b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py @@ -0,0 +1,169 @@ +import os +import time +from unittest.mock import AsyncMock + +import pytest +from httpx import Response + +from litellm.integrations.datadog.datadog_cost_management import ( + DatadogCostManagementLogger, +) +from litellm.types.utils import StandardLoggingPayload + + +@pytest.fixture +def clean_env(): + # Save original env + original_api_key = os.environ.get("DD_API_KEY") + original_app_key = os.environ.get("DD_APP_KEY") + original_site = os.environ.get("DD_SITE") + + # Set test env + os.environ["DD_API_KEY"] = "test_api_key" + os.environ["DD_APP_KEY"] = "test_app_key" + os.environ["DD_SITE"] = "test.datadoghq.com" + + yield + + # Restore original env + if original_api_key: + os.environ["DD_API_KEY"] = original_api_key + else: + del os.environ["DD_API_KEY"] + + if original_app_key: + os.environ["DD_APP_KEY"] = original_app_key + else: + del os.environ["DD_APP_KEY"] + + if original_site: + os.environ["DD_SITE"] = original_site + else: + del os.environ["DD_SITE"] + + +@pytest.mark.asyncio +async def test_init(clean_env): + """ + Test initialization sets up clients and url correctly + """ + logger = DatadogCostManagementLogger() + assert logger.dd_api_key == "test_api_key" + assert logger.dd_app_key == "test_app_key" + assert ( + logger.upload_url == "https://api.test.datadoghq.com/api/v2/cost/custom_costs" + ) + + +@pytest.mark.asyncio +async def test_aggregate_costs(clean_env): + """ + Test that costs are correctly aggregated by provider, model, and date + """ + logger = DatadogCostManagementLogger() + + # Mock some log payloads + now = time.time() + day_str = time.strftime("%Y-%m-%d", time.localtime(now)) + + logs = [ + StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.01, + startTime=now, + metadata={"user_api_key_team_alias": "team-a"}, + ), + StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.02, + startTime=now, + metadata={"user_api_key_team_alias": "team-a"}, + ), + StandardLoggingPayload( + custom_llm_provider="anthropic", + model="claude-3", + response_cost=0.05, + startTime=now, + ), + ] + + aggregated = logger._aggregate_costs(logs) + + assert len(aggregated) == 2 + + # Check OpenAI entry + openai_entry = next(e for e in aggregated if e["ProviderName"] == "openai") + assert openai_entry["BilledCost"] == 0.03 + assert openai_entry["ChargeDescription"] == "LLM Usage for gpt-4" + assert openai_entry["ChargePeriodStart"] == day_str + assert openai_entry["Tags"]["team"] == "team-a" + assert "env" in openai_entry["Tags"] + assert "service" in openai_entry["Tags"] + + # Check Anthropic entry + anthropic_entry = next(e for e in aggregated if e["ProviderName"] == "anthropic") + assert anthropic_entry["BilledCost"] == 0.05 + + +@pytest.mark.asyncio +async def test_async_log_success_event(clean_env): + """ + Test that logs are added to queue + """ + logger = DatadogCostManagementLogger(batch_size=10) + + await logger.async_log_success_event( + kwargs={"standard_logging_object": {"response_cost": 0.01}}, + response_obj={}, + start_time=time.time(), + end_time=time.time(), + ) + + assert len(logger.log_queue) == 1 + assert logger.log_queue[0]["response_cost"] == 0.01 + + # Test zero cost ignored + await logger.async_log_success_event( + kwargs={"standard_logging_object": {"response_cost": 0.0}}, + response_obj={}, + start_time=time.time(), + end_time=time.time(), + ) + + assert len(logger.log_queue) == 1 + + +@pytest.mark.asyncio +async def test_async_send_batch(clean_env): + """ + Test that batch is aggregated and uploaded + """ + logger = DatadogCostManagementLogger() + logger.async_client = AsyncMock() + logger.async_client.put.return_value = Response(202, json={"status": "ok"}) + + # Add logs directly to queue + logger.log_queue = [ + StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.01, + startTime=time.time(), + ) + ] + + await logger.async_send_batch() + + # Verify API called + assert logger.async_client.put.called + call_args = logger.async_client.put.call_args + assert call_args[0][0] == "https://api.test.datadoghq.com/api/v2/cost/custom_costs" + + import json + + # Use call_args.kwargs['content'] + content = json.loads(call_args[1]["content"]) + assert content[0]["ProviderName"] == "openai" + assert content[0]["BilledCost"] == 0.01 diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py new file mode 100644 index 0000000000..2bb51e1e1b --- /dev/null +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py @@ -0,0 +1,62 @@ +import os +from unittest.mock import patch +from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger + + +def test_datadog_llm_obs_agent_configuration(): + """ + Test that DataDog LLM Obs logger correctly configures agent endpoint. + """ + test_env = { + "LITELLM_DD_AGENT_HOST": "localhost", + "LITELLM_DD_LLM_OBS_PORT": "10518", + "DD_API_KEY": "test-api-key", # Optional, but checking if it's preserved + } + + # Ensure DD_SITE is NOT set to verify we don't need it in agent mode + + with patch.dict(os.environ, test_env, clear=True): + with patch("asyncio.create_task"): # Prevent periodic flush task from running + dd_logger = DataDogLLMObsLogger() + + expected_url = "http://localhost:10518/api/intake/llm-obs/v1/trace/spans" + assert dd_logger.intake_url == expected_url + assert dd_logger.DD_API_KEY == "test-api-key" + + +def test_datadog_llm_obs_agent_no_api_key_ok(): + """ + Test that agent mode works WITHOUT DD_API_KEY (agent handles auth). + """ + test_env = { + "LITELLM_DD_AGENT_HOST": "localhost", + # No DD_API_KEY + } + + with patch.dict(os.environ, test_env, clear=True): + with patch("asyncio.create_task"): + # Should NOT raise exception anymore + dd_logger = DataDogLLMObsLogger() + + assert dd_logger.DD_API_KEY is None + # Default port is 8126 if not set + expected_url = "http://localhost:8126/api/intake/llm-obs/v1/trace/spans" + assert dd_logger.intake_url == expected_url + + +def test_datadog_llm_obs_direct_api_configuration(): + """ + Test that direct API configuration still works as expected. + """ + test_env = { + "DD_API_KEY": "direct-api-key", + "DD_SITE": "us5.datadoghq.com", + } + + with patch.dict(os.environ, test_env, clear=True): + with patch("asyncio.create_task"): + dd_logger = DataDogLLMObsLogger() + + expected_url = "https://api.us5.datadoghq.com/api/intake/llm-obs/v1/trace/spans" + assert dd_logger.intake_url == expected_url + assert dd_logger.DD_API_KEY == "direct-api-key" diff --git a/tests/test_litellm/integrations/test_custom_guardrail_recursion.py b/tests/test_litellm/integrations/test_custom_guardrail_recursion.py new file mode 100644 index 0000000000..f05b5848bd --- /dev/null +++ b/tests/test_litellm/integrations/test_custom_guardrail_recursion.py @@ -0,0 +1,73 @@ +import pytest +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.types.guardrails import GuardrailEventHooks +import json + + +class TestCustomGuardrailRecursion: + """ + Specific tests for the circular reference / RecursionError fix in logging. + """ + + def test_log_guardrail_information_handles_circular_references(self): + """ + Test that add_standard_logging method sanitizes input data containing circular references + instead of crashing. + + This reproduces the Langfuse crash scenario: + Request -> Metadata -> GuardrailResponse -> DebugContext -> Request + """ + guardrail = CustomGuardrail( + guardrail_name="recursion_test_guardrail", + event_hook=GuardrailEventHooks.pre_call, + ) + + # 1. Setup Circular Data + request_data = {"user_id": "test_recursive_user"} + metadata = {"session_id": "123"} + request_data["metadata"] = metadata + + # Create the danger: Guardrail Response holding a reference back to request_data + dirty_response = { + "flagged": False, + "debug_context": request_data, # <--- ACCESS TO ROOT (Circular Ref) + } + + # 2. Invoke the logging method + # If the fix is working, this will NOT raise RecursionError + try: + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=dirty_response, + request_data=request_data, + guardrail_status="success", + start_time=1.0, + end_time=2.0, + duration=1.0, + masked_entity_count={}, + event_type=GuardrailEventHooks.pre_call, + ) + except RecursionError: + pytest.fail( + "RecursionError raised! The cyclic reference sanitization failed." + ) + + # 3. Verify the data stored is safe + stored_info = request_data["metadata"][ + "standard_logging_guardrail_information" + ][0] + stored_response = stored_info["guardrail_response"] + + # Check that we can dump it to JSON without crashing (Ultimate proof) + try: + json.dumps(stored_response) + except Exception as e: + pytest.fail(f"Stored data is not JSON serializable: {e}") + + # Check content - keys should be preserved but recursion broken + assert "debug_context" in stored_response + debug_context = stored_response["debug_context"] + + # In a sanitized copy, the nested metadata should be a copy, not the original live dict + assert debug_context["user_id"] == "test_recursive_user" + # The 'metadata' inside 'debug_context' would be where recursion stops or is filtered + assert "metadata" in debug_context diff --git a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py new file mode 100644 index 0000000000..4a9fa3de5f --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py @@ -0,0 +1,203 @@ +import pytest +from unittest.mock import MagicMock, patch +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import ( + UserAPIKeyLabelValues, +) +from litellm.proxy._types import UserAPIKeyAuth + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_includes_client_ip_user_agent(): + """ + Test that async_post_call_failure_hook includes client_ip and user_agent in UserAPIKeyLabelValues + """ + # Mocking + # Mocking + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + # Initialize attributes manually as __init__ is mocked + logger.litellm_proxy_failed_requests_metric = MagicMock() + logger.litellm_proxy_total_requests_metric = MagicMock() + logger.get_labels_for_metric = MagicMock( + return_value=["client_ip", "user_agent"] + ) + + request_data = { + "model": "gpt-4", + "metadata": { + "requester_ip_address": "127.0.0.1", + "user_agent": "test-agent", + }, + } + user_api_key_dict = UserAPIKeyAuth(token="test_token") + original_exception = Exception("Test exception") + + # Mock prometheus_label_factory to inspect arguments + with patch( + "litellm.integrations.prometheus.prometheus_label_factory" + ) as mock_label_factory: + mock_label_factory.return_value = {} + + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=original_exception, + user_api_key_dict=user_api_key_dict, + ) + + # Verification + assert mock_label_factory.call_count >= 1 + + # Check calls + calls = mock_label_factory.call_args_list + found = False + for call in calls: + kwargs = call.kwargs + enum_values = kwargs.get("enum_values") + if isinstance(enum_values, UserAPIKeyLabelValues): + if ( + enum_values.client_ip == "127.0.0.1" + and enum_values.user_agent == "test-agent" + ): + found = True + break + + assert ( + found + ), "UserAPIKeyLabelValues should contain client_ip='127.0.0.1' and user_agent='test-agent'" + + +@pytest.mark.asyncio +async def test_async_post_call_success_hook_includes_client_ip_user_agent(): + """ + Test that async_post_call_success_hook includes client_ip and user_agent in UserAPIKeyLabelValues + """ + # Mocking + # Mocking + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + logger.litellm_proxy_total_requests_metric = MagicMock() + logger.get_labels_for_metric = MagicMock( + return_value=["client_ip", "user_agent"] + ) + + data = { + "model": "gpt-4", + "metadata": { + "requester_ip_address": "192.168.1.1", + "user_agent": "success-agent", + }, + } + user_api_key_dict = UserAPIKeyAuth(token="test_token") + response = MagicMock() + + # Mock prometheus_label_factory to inspect arguments + with patch( + "litellm.integrations.prometheus.prometheus_label_factory" + ) as mock_label_factory: + mock_label_factory.return_value = {} + + await logger.async_post_call_success_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + ) + + # Verification + assert mock_label_factory.call_count >= 1 + + # Check calls + calls = mock_label_factory.call_args_list + found = False + for call in calls: + kwargs = call.kwargs + enum_values = kwargs.get("enum_values") + if isinstance(enum_values, UserAPIKeyLabelValues): + if ( + enum_values.client_ip == "192.168.1.1" + and enum_values.user_agent == "success-agent" + ): + found = True + break + + assert ( + found + ), "UserAPIKeyLabelValues should contain client_ip='192.168.1.1' and user_agent='success-agent'" + + +def test_set_llm_deployment_failure_metrics_includes_client_ip_user_agent(): + """ + Test that set_llm_deployment_failure_metrics includes client_ip and user_agent in UserAPIKeyLabelValues + """ + # Mocking + # Mocking + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + logger.litellm_deployment_failure_responses = MagicMock() + logger.litellm_deployment_total_requests = MagicMock() + logger.get_labels_for_metric = MagicMock( + return_value=["client_ip", "user_agent"] + ) + logger.set_deployment_partial_outage = MagicMock() + + request_kwargs = { + "model": "gpt-4", + "standard_logging_object": { + "metadata": { + "requester_ip_address": "10.0.0.1", + "user_agent": "failure-deployment", + "user_api_key_team_id": "team_1", + "user_api_key_team_alias": "team_alias_1", + "user_api_key_alias": "key_alias_1", + }, + "model_group": "group_1", + "api_base": "http://api.base", + "model_id": "model_1", + }, + "litellm_params": {}, + "exception": Exception("Deployment failure"), + } + + # Mock prometheus_label_factory to inspect arguments + with patch( + "litellm.integrations.prometheus.prometheus_label_factory" + ) as mock_label_factory: + mock_label_factory.return_value = {} + + logger.set_llm_deployment_failure_metrics(request_kwargs=request_kwargs) + + # Verification + assert mock_label_factory.call_count >= 1 + + # Check calls + calls = mock_label_factory.call_args_list + found = False + for call in calls: + kwargs = call.kwargs + enum_values = kwargs.get("enum_values") + if isinstance(enum_values, UserAPIKeyLabelValues): + if ( + enum_values.client_ip == "10.0.0.1" + and enum_values.user_agent == "failure-deployment" + ): + found = True + break + + assert ( + found + ), "UserAPIKeyLabelValues should contain client_ip='10.0.0.1' and user_agent='failure-deployment'" + + +if __name__ == "__main__": + import asyncio + + asyncio.run(test_async_post_call_failure_hook_includes_client_ip_user_agent()) + asyncio.run(test_async_post_call_success_hook_includes_client_ip_user_agent()) + test_set_llm_deployment_failure_metrics_includes_client_ip_user_agent() + print("✅ All client_ip and user_agent tests passed!") diff --git a/tests/test_litellm/integrations/test_prometheus_labels.py b/tests/test_litellm/integrations/test_prometheus_labels.py index d9a57df161..9d4626b88a 100644 --- a/tests/test_litellm/integrations/test_prometheus_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_labels.py @@ -26,15 +26,49 @@ def test_user_email_in_required_metrics(): "litellm_input_tokens_metric", "litellm_output_tokens_metric", "litellm_requests_metric", - "litellm_spend_metric" + "litellm_spend_metric", ] for metric_name in metrics_with_user_email: labels = PrometheusMetricLabels.get_labels(metric_name) - assert user_email_label in labels, f"Metric {metric_name} should contain user_email label" + assert ( + user_email_label in labels + ), f"Metric {metric_name} should contain user_email label" print(f"✅ {metric_name} contains user_email label") +def test_model_id_in_required_metrics(): + """ + Test that model_id label is present in all the metrics that should have it + """ + model_id_label = UserAPIKeyLabelNames.MODEL_ID.value + + # Metrics that should have model_id + metrics_with_model_id = [ + "litellm_proxy_total_requests_metric", + "litellm_proxy_failed_requests_metric", + "litellm_input_tokens_metric", + "litellm_output_tokens_metric", + "litellm_requests_metric", + "litellm_spend_metric", + "litellm_llm_api_latency_metric", + "litellm_remaining_requests_metric", + "litellm_deployment_successful_fallbacks", + "litellm_cache_hits_metric", + "litellm_cache_misses_metric", + "litellm_remaining_api_key_requests_for_model", + "litellm_remaining_api_key_tokens_for_model", + "litellm_llm_api_failed_requests_metric", + ] + + for metric_name in metrics_with_model_id: + labels = PrometheusMetricLabels.get_labels(metric_name) + assert ( + model_id_label in labels + ), f"Metric {metric_name} should contain model_id label" + print(f"✅ {metric_name} contains model_id label") + + def test_user_email_label_exists(): """Test that the USER_EMAIL label is properly defined""" assert UserAPIKeyLabelNames.USER_EMAIL.value == "user_email" @@ -52,12 +86,14 @@ def test_prometheus_metric_labels_structure(): "litellm_proxy_failed_requests_metric", "litellm_input_tokens_metric", "litellm_output_tokens_metric", - "litellm_spend_metric" + "litellm_spend_metric", ] for metric_name in test_metrics: # Check metric is in DEFINED_PROMETHEUS_METRICS - assert metric_name in get_args(DEFINED_PROMETHEUS_METRICS), f"{metric_name} should be in DEFINED_PROMETHEUS_METRICS" + assert metric_name in get_args( + DEFINED_PROMETHEUS_METRICS + ), f"{metric_name} should be in DEFINED_PROMETHEUS_METRICS" # Check labels can be retrieved labels = PrometheusMetricLabels.get_labels(metric_name) @@ -98,11 +134,11 @@ def test_route_normalization_for_responses_api(): """ Test that route normalization prevents high cardinality in Prometheus metrics for the /v1/responses/{response_id} endpoint. - + Issue: https://github.com/BerriAI/litellm/issues/XXXX Each unique response ID was creating a separate metric line, causing the /metrics endpoint to grow to ~30MB and take ~40 seconds to respond. - + Fix: Routes are normalized to collapse dynamic IDs into placeholders. """ from litellm.proxy.auth.auth_utils import normalize_request_route @@ -115,43 +151,53 @@ def test_route_normalization_for_responses_api(): ("/v1/responses/resp_abc123", "/v1/responses/{response_id}"), ("/v1/responses/litellm_poll_xyz", "/v1/responses/{response_id}"), ] - + for original, expected in responses_routes: normalized = normalize_request_route(original) - assert normalized == expected, \ - f"Failed: {original} -> {normalized} (expected {expected})" - + assert ( + normalized == expected + ), f"Failed: {original} -> {normalized} (expected {expected})" + # Verify cardinality reduction - unique_normalized = set(normalize_request_route(route) for route, _ in responses_routes) - assert len(unique_normalized) == 1, \ - f"Expected 1 unique normalized route, got {len(unique_normalized)}: {unique_normalized}" - - print(f"✅ Responses API routes: {len(responses_routes)} different IDs normalized to 1 metric label") - + unique_normalized = set( + normalize_request_route(route) for route, _ in responses_routes + ) + assert ( + len(unique_normalized) == 1 + ), f"Expected 1 unique normalized route, got {len(unique_normalized)}: {unique_normalized}" + + print( + f"✅ Responses API routes: {len(responses_routes)} different IDs normalized to 1 metric label" + ) + def test_route_normalization_for_sub_routes(): """Test that sub-routes like /cancel and /input_items are normalized correctly""" from litellm.proxy.auth.auth_utils import normalize_request_route - + sub_routes = [ ("/v1/responses/id1/cancel", "/v1/responses/{response_id}/cancel"), ("/v1/responses/id2/cancel", "/v1/responses/{response_id}/cancel"), ("/v1/responses/id3/input_items", "/v1/responses/{response_id}/input_items"), - ("/openai/v1/responses/id4/input_items", "/openai/v1/responses/{response_id}/input_items"), + ( + "/openai/v1/responses/id4/input_items", + "/openai/v1/responses/{response_id}/input_items", + ), ] - + for original, expected in sub_routes: normalized = normalize_request_route(original) - assert normalized == expected, \ - f"Failed: {original} -> {normalized} (expected {expected})" - + assert ( + normalized == expected + ), f"Failed: {original} -> {normalized} (expected {expected})" + print("✅ Sub-routes normalized correctly") def test_route_normalization_preserves_static_routes(): """Test that static routes are not affected by normalization""" from litellm.proxy.auth.auth_utils import normalize_request_route - + static_routes = [ "/chat/completions", "/v1/chat/completions", @@ -161,46 +207,47 @@ def test_route_normalization_preserves_static_routes(): "/v1/models", "/v1/responses", # List endpoint without ID ] - + for route in static_routes: normalized = normalize_request_route(route) - assert normalized == route, \ - f"Static route should not be modified: {route} -> {normalized}" - + assert ( + normalized == route + ), f"Static route should not be modified: {route} -> {normalized}" + print(f"✅ {len(static_routes)} static routes preserved") def test_route_normalization_other_dynamic_apis(): """Test normalization for other OpenAI-compatible APIs with dynamic IDs""" from litellm.proxy.auth.auth_utils import normalize_request_route - + test_cases = [ # Threads API ("/v1/threads/thread_123", "/v1/threads/{thread_id}"), ("/v1/threads/thread_abc/messages", "/v1/threads/{thread_id}/messages"), - ("/v1/threads/thread_abc/runs/run_123", "/v1/threads/{thread_id}/runs/{run_id}"), - + ( + "/v1/threads/thread_abc/runs/run_123", + "/v1/threads/{thread_id}/runs/{run_id}", + ), # Vector Stores API ("/v1/vector_stores/vs_123", "/v1/vector_stores/{vector_store_id}"), ("/v1/vector_stores/vs_123/files", "/v1/vector_stores/{vector_store_id}/files"), - # Assistants API ("/v1/assistants/asst_123", "/v1/assistants/{assistant_id}"), - # Files API ("/v1/files/file_123", "/v1/files/{file_id}"), ("/v1/files/file_123/content", "/v1/files/{file_id}/content"), - # Batches API ("/v1/batches/batch_123", "/v1/batches/{batch_id}"), ("/v1/batches/batch_123/cancel", "/v1/batches/{batch_id}/cancel"), ] - + for original, expected in test_cases: normalized = normalize_request_route(original) - assert normalized == expected, \ - f"Failed: {original} -> {normalized} (expected {expected})" - + assert ( + normalized == expected + ), f"Failed: {original} -> {normalized} (expected {expected})" + print(f"✅ {len(test_cases)} other API routes normalized correctly") @@ -219,25 +266,39 @@ def test_prometheus_metrics_use_normalized_routes(): # Create a mock PrometheusLogger prometheus_logger = MagicMock() - prometheus_logger.get_labels_for_metric = PrometheusLogger.get_labels_for_metric.__get__(prometheus_logger) - + prometheus_logger.get_labels_for_metric = ( + PrometheusLogger.get_labels_for_metric.__get__(prometheus_logger) + ) + # Test with a normalized route enum_values = UserAPIKeyLabelValues( route="/v1/responses/{response_id}", # Normalized route status_code="200", requested_model="gpt-4", ) - + labels = prometheus_label_factory( supported_enum_labels=prometheus_logger.get_labels_for_metric( metric_name="litellm_proxy_total_requests_metric" ), enum_values=enum_values, ) - + # Verify the route is normalized in labels - assert labels["route"] == "/v1/responses/{response_id}", \ - f"Expected normalized route in labels, got: {labels.get('route')}" - + assert ( + labels["route"] == "/v1/responses/{response_id}" + ), f"Expected normalized route in labels, got: {labels.get('route')}" + print("✅ Prometheus metrics use normalized routes in labels") + +if __name__ == "__main__": + test_user_email_in_required_metrics() + test_user_email_label_exists() + test_prometheus_metric_labels_structure() + test_route_normalization_for_responses_api() + test_route_normalization_for_sub_routes() + test_route_normalization_preserves_static_routes() + test_route_normalization_other_dynamic_apis() + test_prometheus_metrics_use_normalized_routes() + print("\n✅ All prometheus label tests passed!") diff --git a/tests/test_litellm/integrations/test_prometheus_missing_metrics.py b/tests/test_litellm/integrations/test_prometheus_missing_metrics.py new file mode 100644 index 0000000000..7fcfb21ed4 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_missing_metrics.py @@ -0,0 +1,77 @@ +""" +Unit tests for the new Prometheus metrics that were previously missing from validation. + +Tests for: +- litellm_remaining_api_key_requests_for_model +- litellm_remaining_api_key_tokens_for_model +- litellm_callback_logging_failures_metric +""" +from typing import get_args +from litellm.types.integrations.prometheus import ( + DEFINED_PROMETHEUS_METRICS, + PrometheusMetricLabels, + UserAPIKeyLabelNames, +) + + +def test_new_metrics_in_defined_metrics(): + """ + Test that the new metrics are present in DEFINED_PROMETHEUS_METRICS. + """ + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + + new_metrics = [ + "litellm_remaining_api_key_requests_for_model", + "litellm_remaining_api_key_tokens_for_model", + "litellm_callback_logging_failures_metric", + ] + + for metric in new_metrics: + assert ( + metric in defined_metrics + ), f"{metric} should be in DEFINED_PROMETHEUS_METRICS" + + +def test_new_metrics_have_correct_labels(): + """ + Test that the new metrics have the correct labels defined. + """ + # Test API Key limits metrics labels + api_key_metrics = [ + "litellm_remaining_api_key_requests_for_model", + "litellm_remaining_api_key_tokens_for_model", + ] + + expected_api_key_labels = [ + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + ] + + for metric in api_key_metrics: + labels = PrometheusMetricLabels.get_labels(metric) + for expected_label in expected_api_key_labels: + assert ( + expected_label in labels + ), f"{metric} should have label {expected_label}" + + # Test Callback failure metric labels + callback_metric = "litellm_callback_logging_failures_metric" + callback_labels = PrometheusMetricLabels.get_labels(callback_metric) + + assert ( + UserAPIKeyLabelNames.CALLBACK_NAME.value in callback_labels + ), f"{callback_metric} should have label {UserAPIKeyLabelNames.CALLBACK_NAME.value}" + + +def test_callback_name_label_definition(): + """ + Test that CALLBACK_NAME is defined correctly in UserAPIKeyLabelNames. + """ + assert UserAPIKeyLabelNames.CALLBACK_NAME.value == "callback_name" + + +if __name__ == "__main__": + test_new_metrics_in_defined_metrics() + test_new_metrics_have_correct_labels() + test_callback_name_label_definition() 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 a22fe13798..e87233a52a 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 @@ -1138,6 +1138,73 @@ def test_bedrock_create_bedrock_block_different_document_formats(): assert block["document"]["name"].endswith(f"_{format_type}") assert block["document"]["format"] == format_type +def test_bedrock_nova_web_search_options_mapping(): + """ + Test that web_search_options is correctly mapped to Nova grounding. + + This follows the LiteLLM pattern for web search where: + - Vertex AI maps web_search_options to {"googleSearch": {}} + - Anthropic maps web_search_options to {"type": "web_search_20250305", ...} + - Nova should map web_search_options to {"systemTool": {"name": "nova_grounding"}} + """ + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + config = AmazonConverseConfig() + + # Test basic mapping for Nova model + result = config._map_web_search_options({}, "amazon.nova-pro-v1:0") + + assert result is not None + system_tool = result.get("systemTool") + assert system_tool is not None + assert system_tool["name"] == "nova_grounding" + + # Test with search_context_size (should be ignored for Nova) + result2 = config._map_web_search_options( + {"search_context_size": "high"}, + "us.amazon.nova-premier-v1:0" + ) + + assert result2 is not None + system_tool2 = result2.get("systemTool") + assert system_tool2 is not None + assert system_tool2["name"] == "nova_grounding" + # Nova doesn't support search_context_size, so it's just ignored + +def test_bedrock_tools_pt_does_not_handle_system_tool(): + """ + Verify that _bedrock_tools_pt does NOT handle system_tool format. + + System tools (nova_grounding) should be added via web_search_options, + not via the tools parameter directly. + """ + + from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_tools_pt + + # Regular function tools should still work + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"} + }, + "required": ["location"] + } + } + } + ] + + result = _bedrock_tools_pt(tools=tools) + + assert len(result) == 1 + tool_spec = result[0].get("toolSpec") + assert tool_spec is not None + assert tool_spec["name"] == "get_weather" def test_convert_to_anthropic_tool_result_image_with_cache_control(): """ @@ -1305,12 +1372,12 @@ def test_convert_to_anthropic_tool_result_image_url_as_http(): assert result["content"][0]["cache_control"]["type"] == "ephemeral" def test_anthropic_messages_pt_server_tool_use_passthrough(): """ - Test that anthropic_messages_pt passes through server_tool_use and + Test that anthropic_messages_pt passes through server_tool_use and tool_search_tool_result blocks in assistant message content. - + These are Anthropic-native content types used for tool search functionality that need to be preserved when reconstructing multi-turn conversations. - + Fixes: https://github.com/BerriAI/litellm/issues/XXXXX """ from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt @@ -1359,15 +1426,15 @@ def test_anthropic_messages_pt_server_tool_use_passthrough(): # Verify we have 3 messages (user, assistant, user) assert len(result) == 3 - + # Verify the assistant message content assistant_msg = result[1] assert assistant_msg["role"] == "assistant" assert isinstance(assistant_msg["content"], list) - + # Find the different content block types content_types = [block.get("type") for block in assistant_msg["content"]] - + # Verify server_tool_use block is preserved assert "server_tool_use" in content_types server_tool_use_block = next( @@ -1376,7 +1443,7 @@ def test_anthropic_messages_pt_server_tool_use_passthrough(): assert server_tool_use_block["id"] == "srvtoolu_01ABC123" assert server_tool_use_block["name"] == "tool_search_tool_regex" assert server_tool_use_block["input"] == {"query": ".*time.*"} - + # Verify tool_search_tool_result block is preserved assert "tool_search_tool_result" in content_types tool_result_block = next( @@ -1385,7 +1452,7 @@ def test_anthropic_messages_pt_server_tool_use_passthrough(): assert tool_result_block["tool_use_id"] == "srvtoolu_01ABC123" assert tool_result_block["content"]["type"] == "tool_search_tool_search_result" assert tool_result_block["content"]["tool_references"][0]["tool_name"] == "get_time" - + # Verify text block is also preserved assert "text" in content_types text_block = next( diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index e035e193fe..1f3f558a49 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -787,11 +787,13 @@ def test_get_masked_values(): "presidio_ad_hoc_recognizers": None, "aws_bedrock_runtime_endpoint": None, "presidio_anonymizer_api_base": None, + "vertex_credentials": "{sensitive_api_key}", } masked_values = _get_masked_values( sensitive_object, unmasked_length=4, number_of_asterisks=4 ) assert masked_values["presidio_anonymizer_api_base"] is None + assert masked_values["vertex_credentials"] == "{s****y}" @pytest.mark.asyncio diff --git a/tests/test_litellm/llms/vercel_ai_gateway/embedding/__init__.py b/tests/test_litellm/llms/vercel_ai_gateway/embedding/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py b/tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py new file mode 100644 index 0000000000..af1e1df92f --- /dev/null +++ b/tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py @@ -0,0 +1,218 @@ +import os +import sys +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.vercel_ai_gateway.embedding.transformation import ( + VercelAIGatewayEmbeddingConfig, +) +from litellm.llms.vercel_ai_gateway.common_utils import VercelAIGatewayException +from litellm.types.utils import EmbeddingResponse + + +def test_vercel_ai_gateway_embedding_get_complete_url(): + """Test URL generation for embeddings endpoint""" + config = VercelAIGatewayEmbeddingConfig() + + # Test with default API base + url = config.get_complete_url( + api_base=None, + api_key=None, + model="openai/text-embedding-3-small", + optional_params={}, + litellm_params={}, + ) + assert url == "https://ai-gateway.vercel.sh/v1/embeddings" + + # Test with custom API base + url = config.get_complete_url( + api_base="https://custom.vercel.sh/v1", + api_key=None, + model="openai/text-embedding-3-small", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.vercel.sh/v1/embeddings" + + # Test with trailing slash + url = config.get_complete_url( + api_base="https://custom.vercel.sh/v1/", + api_key=None, + model="openai/text-embedding-3-small", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.vercel.sh/v1/embeddings" + + +def test_vercel_ai_gateway_embedding_transform_request(): + """Test request transformation for embeddings""" + config = VercelAIGatewayEmbeddingConfig() + + # Test with string input + request = config.transform_embedding_request( + model="openai/text-embedding-3-small", + input="Hello world", + optional_params={}, + headers={}, + ) + assert request["model"] == "openai/text-embedding-3-small" + assert request["input"] == ["Hello world"] + + # Test with list input + request = config.transform_embedding_request( + model="openai/text-embedding-3-small", + input=["Hello", "World"], + optional_params={}, + headers={}, + ) + assert request["model"] == "openai/text-embedding-3-small" + assert request["input"] == ["Hello", "World"] + + # Test stripping vercel_ai_gateway/ prefix + request = config.transform_embedding_request( + model="vercel_ai_gateway/openai/text-embedding-3-small", + input="Hello", + optional_params={}, + headers={}, + ) + assert request["model"] == "openai/text-embedding-3-small" + + +def test_vercel_ai_gateway_embedding_transform_request_with_dimensions(): + """Test request transformation with dimensions parameter""" + config = VercelAIGatewayEmbeddingConfig() + + request = config.transform_embedding_request( + model="openai/text-embedding-3-small", + input="Hello world", + optional_params={"dimensions": 768}, + headers={}, + ) + assert request["model"] == "openai/text-embedding-3-small" + assert request["input"] == ["Hello world"] + assert request["dimensions"] == 768 + + +def test_vercel_ai_gateway_embedding_validate_environment(): + """Test header validation and setup""" + config = VercelAIGatewayEmbeddingConfig() + + headers = config.validate_environment( + headers={}, + model="openai/text-embedding-3-small", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test_key", + ) + assert headers["Content-Type"] == "application/json" + assert headers["Authorization"] == "Bearer test_key" + + # Test with existing headers (should merge) + headers = config.validate_environment( + headers={"X-Custom": "value"}, + model="openai/text-embedding-3-small", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test_key", + ) + assert headers["X-Custom"] == "value" + assert headers["Authorization"] == "Bearer test_key" + + +def test_vercel_ai_gateway_embedding_get_supported_params(): + """Test supported OpenAI parameters""" + config = VercelAIGatewayEmbeddingConfig() + supported = config.get_supported_openai_params("openai/text-embedding-3-small") + + assert "dimensions" in supported + assert "encoding_format" in supported + assert "timeout" in supported + assert "user" in supported + + +def test_vercel_ai_gateway_embedding_map_openai_params(): + """Test OpenAI parameter mapping""" + config = VercelAIGatewayEmbeddingConfig() + + optional_params = config.map_openai_params( + non_default_params={"dimensions": 768, "encoding_format": "float"}, + optional_params={}, + model="openai/text-embedding-3-small", + drop_params=False, + ) + assert optional_params["dimensions"] == 768 + assert optional_params["encoding_format"] == "float" + + +def test_vercel_ai_gateway_embedding_error_class(): + """Test error class creation""" + config = VercelAIGatewayEmbeddingConfig() + + error = config.get_error_class( + error_message="Test error", + status_code=400, + headers={"Content-Type": "application/json"}, + ) + + assert isinstance(error, VercelAIGatewayException) + assert error.message == "Test error" + assert error.status_code == 400 + + +def test_vercel_ai_gateway_embedding_transform_response(): + """Test response transformation""" + config = VercelAIGatewayEmbeddingConfig() + + mock_response = MagicMock(spec=httpx.Response) + mock_response.text = '{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2,0.3]}],"model":"openai/text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}' + mock_response.json.return_value = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "openai/text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + } + + mock_logging = MagicMock() + + response = config.transform_embedding_response( + model="openai/text-embedding-3-small", + raw_response=mock_response, + model_response=EmbeddingResponse(), + logging_obj=mock_logging, + api_key="test_key", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + assert response is not None + mock_logging.post_call.assert_called_once() + + +def test_vercel_ai_gateway_embedding_env_vars(): + """Test environment variable handling""" + config = VercelAIGatewayEmbeddingConfig() + + with patch.dict( + os.environ, + { + "VERCEL_AI_GATEWAY_API_BASE": "https://env.vercel.sh/v1", + }, + ): + url = config.get_complete_url( + api_base=None, + api_key=None, + model="openai/text-embedding-3-small", + optional_params={}, + litellm_params={}, + ) + assert url == "https://env.vercel.sh/v1/embeddings" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index ecdc75ede5..f241b2aa0d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1885,7 +1885,7 @@ class TestMCPServerManager: # Create mock client that tracks call_tool usage mock_client = AsyncMock() - async def mock_call_tool(params): + async def mock_call_tool(params, host_progress_callback=None): # Return a mock CallToolResult result = MagicMock(spec=CallToolResult) result.content = [{"type": "text", "text": "Tool executed successfully"}] diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 62f9cc33b6..82920ce1d8 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -8,6 +8,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( _get_customer_id_from_standard_headers, get_end_user_id_from_request_body, + get_model_from_request, get_key_model_rpm_limit, get_key_model_tpm_limit, ) @@ -186,3 +187,25 @@ class TestGetEndUserIdFromRequestBodyWithStandardHeaders: request_body=request_body, request_headers=headers ) assert result == "body-user" + + +def test_get_model_from_request_supports_google_model_names_with_slashes(): + assert ( + get_model_from_request( + request_data={}, + route="/v1beta/models/bedrock/claude-sonnet-3.7:generateContent", + ) + == "bedrock/claude-sonnet-3.7" + ) + assert ( + get_model_from_request( + request_data={}, + route="/models/hosted_vllm/gpt-oss-20b:generateContent", + ) + == "hosted_vllm/gpt-oss-20b" + ) + + +def test_get_model_from_request_vertex_passthrough_still_works(): + route = "/vertex_ai/v1/projects/p/locations/l/publishers/google/models/gemini-1.5-pro:generateContent" + assert get_model_from_request(request_data={}, route=route) == "gemini-1.5-pro" diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index e7b27908c1..a0e29e0610 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -626,3 +626,133 @@ async def test_expire_previous_ui_session_tokens_exception_handling(): # Should not raise exception despite database error await expire_previous_ui_session_tokens(user_id, mock_prisma_client) + + +@pytest.mark.asyncio +async def test_authenticate_user_admin_login_with_non_ascii_characters(): + """Test admin login with non-ASCII characters in password (issue #19559)""" + master_key = "sk-1234" + ui_username = "admin£test" + ui_password = "sk-1234£pass" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict( + os.environ, + { + "UI_USERNAME": ui_username, + "UI_PASSWORD": ui_password, + "DATABASE_URL": "postgresql://test:test@localhost/test", + }, + ): + with patch( + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_generate_key: + mock_generate_key.return_value = { + "token": "test-token-123", + "user_id": LITELLM_PROXY_ADMIN_NAME, + } + + with patch( + "litellm.proxy.auth.login_utils.user_update", + new_callable=AsyncMock, + return_value=None, + ) as mock_user_update: + with patch( + "litellm.proxy.auth.login_utils.get_secret_bool", + return_value=False, + ): + result = await authenticate_user( + username=ui_username, + password=ui_password, + master_key=master_key, + prisma_client=mock_prisma_client, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == LITELLM_PROXY_ADMIN_NAME + assert result.key == "test-token-123" + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + + +def test_authenticate_user_non_ascii_direct_comparison(): + """Test that non-ASCII characters can be compared directly (unit test for fix)""" + import secrets + + # This test verifies the fix handles non-ASCII by encoding to bytes + username = "admin£test" + password = "pass£word" + + # This would fail without encoding: + # secrets.compare_digest(username, username) # TypeError! + + # But works with the fix: + result = secrets.compare_digest( + username.encode("utf-8"), username.encode("utf-8") + ) + assert result is True + + # And correctly returns False for different passwords + result = secrets.compare_digest( + password.encode("utf-8"), "different£pass".encode("utf-8") + ) + assert result is False + + +@pytest.mark.asyncio +async def test_authenticate_user_database_login_with_non_ascii_password(): + """Test database user login with non-ASCII characters in password (issue #19559)""" + master_key = "sk-1234" + user_email = "test@example.com" + password_with_special_char = "correct£password" + hashed_password = hash_token(token=password_with_special_char) + + mock_user = MagicMock() + mock_user.user_id = "test-user-123" + mock_user.user_email = user_email + mock_user.password = hashed_password + mock_user.user_role = LitellmUserRoles.INTERNAL_USER + + def mock_find_first(**kwargs): + where = kwargs.get("where", {}) + user_email_filter = where.get("user_email", {}) + if str(user_email_filter.get("equals", "")).lower() == user_email.lower(): + return mock_user + return None + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( + side_effect=mock_find_first + ) + + with patch.dict( + os.environ, + { + "DATABASE_URL": "postgresql://test:test@localhost/test", + "UI_USERNAME": "admin", + "UI_PASSWORD": "admin-password", + }, + ): + with patch( + "litellm.proxy.auth.login_utils.expire_previous_ui_session_tokens", + new_callable=AsyncMock, + return_value=None, + ): + with patch( + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_generate_key: + mock_generate_key.return_value = {"token": "token-123"} + + result = await authenticate_user( + username=user_email, + password=password_with_special_char, + master_key=master_key, + prisma_client=mock_prisma_client, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == "test-user-123" + assert result.user_email == user_email diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index b8084906fa..a745ac3de1 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -161,9 +161,11 @@ def test_virtual_key_llm_api_route_includes_passthrough_prefix(route): [ "/v1beta/models/gemini-2.5-flash:countTokens", "/v1beta/models/gemini-2.0-flash:generateContent", + "/v1beta/models/bedrock/claude-sonnet-3.7:generateContent", "/v1beta/models/gemini-1.5-pro:streamGenerateContent", "/models/gemini-2.5-flash:countTokens", "/models/gemini-2.0-flash:generateContent", + "/models/bedrock/claude-sonnet-3.7:generateContent", "/models/gemini-1.5-pro:streamGenerateContent", ], ) @@ -187,9 +189,11 @@ def test_virtual_key_llm_api_routes_allows_google_routes(route): "/v1beta/models/google-gemini-2-5-pro-code-reviewer-k8s:generateContent", "/v1beta/models/gemini-2.5-flash-exp:countTokens", "/v1beta/models/custom-model-name-123:streamGenerateContent", + "/v1beta/models/bedrock/claude-sonnet-3.7:generateContent", "/models/google-gemini-2-5-pro-code-reviewer-k8s:generateContent", "/models/gemini-2.5-flash-exp:countTokens", "/models/custom-model-name-123:streamGenerateContent", + "/models/bedrock/claude-sonnet-3.7:generateContent", ], ) def test_google_routes_with_dynamic_model_names_recognized_as_llm_api_route(route): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 61d44e46da..5c03914192 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -549,6 +549,62 @@ class TestAdditionalParams: ) +class TestModelParameter: + """Test model parameter handling in guardrail requests""" + + @pytest.mark.asyncio + async def test_model_passed_from_inputs( + self, generic_guardrail, mock_request_data_input + ): + """Test that model is passed to the API when provided in inputs""" + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "NONE", + "texts": ["test"], + } + mock_response.raise_for_status = MagicMock() + + with patch.object( + generic_guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + await generic_guardrail.apply_guardrail( + inputs={"texts": ["test"], "model": "gpt-4"}, + request_data=mock_request_data_input, + input_type="request", + ) + + # Verify API was called with model + call_args = mock_post.call_args + json_payload = call_args.kwargs["json"] + assert json_payload["model"] == "gpt-4" + + @pytest.mark.asyncio + async def test_model_none_when_not_provided( + self, generic_guardrail, mock_request_data_input + ): + """Test that model is None when not provided in inputs""" + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "NONE", + "texts": ["test"], + } + mock_response.raise_for_status = MagicMock() + + with patch.object( + generic_guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + await generic_guardrail.apply_guardrail( + inputs={"texts": ["test"]}, # No model in inputs + request_data=mock_request_data_input, + input_type="request", + ) + + # Verify API was called with model=None + call_args = mock_post.call_args + json_payload = call_args.kwargs["json"] + assert json_payload["model"] is None + + class TestErrorHandling: """Test error handling scenarios""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py index 9ede649f39..fb7480d263 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py @@ -3,6 +3,7 @@ import sys import uuid from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException from httpx import Request, Response @@ -47,20 +48,129 @@ def test_onyx_guard_config(): del os.environ["ONYX_API_KEY"] +def test_onyx_guard_with_custom_timeout_from_kwargs(): + """Test Onyx guard instantiation with custom timeout passed via kwargs.""" + # Set environment variables for testing + os.environ["ONYX_API_BASE"] = "https://test.onyx.security" + os.environ["ONYX_API_KEY"] = "test-api-key" + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + + # Simulate how guardrail is instantiated from config with timeout + guardrail = OnyxGuardrail( + guardrail_name="onyx-guard-custom-timeout", + event_hook="pre_call", + default_on=True, + timeout=45.0, + ) + + # Verify the client was initialized with custom timeout + mock_get_client.assert_called() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 45.0 + assert timeout_param.connect == 5.0 + + # Clean up + if "ONYX_API_BASE" in os.environ: + del os.environ["ONYX_API_BASE"] + if "ONYX_API_KEY" in os.environ: + del os.environ["ONYX_API_KEY"] + + +def test_onyx_guard_with_timeout_none_uses_env_var(): + """Test Onyx guard with timeout=None uses ONYX_TIMEOUT env var. + + When timeout=None is passed (as it would be from config model with default None), + the ONYX_TIMEOUT environment variable should be used. + """ + # Set environment variables for testing + os.environ["ONYX_API_BASE"] = "https://test.onyx.security" + os.environ["ONYX_API_KEY"] = "test-api-key" + os.environ["ONYX_TIMEOUT"] = "60" + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + + # Pass timeout=None to simulate config model behavior + guardrail = OnyxGuardrail( + guardrail_name="onyx-guard-env-timeout", + event_hook="pre_call", + default_on=True, + timeout=None, # This triggers env var lookup + ) + + # Verify the client was initialized with timeout from env var + mock_get_client.assert_called() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 60.0 + assert timeout_param.connect == 5.0 + + # Clean up + if "ONYX_API_BASE" in os.environ: + del os.environ["ONYX_API_BASE"] + if "ONYX_API_KEY" in os.environ: + del os.environ["ONYX_API_KEY"] + if "ONYX_TIMEOUT" in os.environ: + del os.environ["ONYX_TIMEOUT"] + + +def test_onyx_guard_with_timeout_none_defaults_to_10(): + """Test Onyx guard with timeout=None and no env var defaults to 10 seconds.""" + # Set environment variables for testing + os.environ["ONYX_API_BASE"] = "https://test.onyx.security" + os.environ["ONYX_API_KEY"] = "test-api-key" + # Ensure ONYX_TIMEOUT is not set + if "ONYX_TIMEOUT" in os.environ: + del os.environ["ONYX_TIMEOUT"] + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + + # Pass timeout=None with no env var - should default to 10.0 + guardrail = OnyxGuardrail( + guardrail_name="onyx-guard-default-timeout", + event_hook="pre_call", + default_on=True, + timeout=None, + ) + + # Verify the client was initialized with default timeout of 10.0 + mock_get_client.assert_called() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 10.0 + assert timeout_param.connect == 5.0 + + # Clean up + if "ONYX_API_BASE" in os.environ: + del os.environ["ONYX_API_BASE"] + if "ONYX_API_KEY" in os.environ: + del os.environ["ONYX_API_KEY"] + + class TestOnyxGuardrail: """Test suite for Onyx Security Guardrail integration.""" def setup_method(self): """Setup test environment.""" # Clean up any existing environment variables - for key in ["ONYX_API_BASE", "ONYX_API_KEY"]: + for key in ["ONYX_API_BASE", "ONYX_API_KEY", "ONYX_TIMEOUT"]: if key in os.environ: del os.environ[key] def teardown_method(self): """Clean up test environment.""" # Clean up any environment variables set during tests - for key in ["ONYX_API_BASE", "ONYX_API_KEY"]: + for key in ["ONYX_API_BASE", "ONYX_API_KEY", "ONYX_TIMEOUT"]: if key in os.environ: del os.environ[key] @@ -103,6 +213,95 @@ class TestOnyxGuardrail: ): OnyxGuardrail(guardrail_name="test-guard", event_hook="pre_call") + def test_initialization_with_default_timeout(self): + """Test that default timeout is 10.0 seconds.""" + os.environ["ONYX_API_KEY"] = "test-api-key" + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + guardrail = OnyxGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True + ) + + # Verify the client was initialized with correct timeout + mock_get_client.assert_called_once() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 10.0 + assert timeout_param.connect == 5.0 + + def test_initialization_with_custom_timeout_parameter(self): + """Test initialization with custom timeout parameter.""" + os.environ["ONYX_API_KEY"] = "test-api-key" + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + guardrail = OnyxGuardrail( + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + timeout=30.0, + ) + + # Verify the client was initialized with custom timeout + mock_get_client.assert_called_once() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 30.0 + assert timeout_param.connect == 5.0 + + def test_initialization_with_timeout_from_env_var(self): + """Test initialization with timeout from ONYX_TIMEOUT environment variable. + + Note: The env var is only used when timeout=None is explicitly passed, + since the default parameter value is 10.0 (not None). + """ + os.environ["ONYX_API_KEY"] = "test-api-key" + os.environ["ONYX_TIMEOUT"] = "25" + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + # Must pass timeout=None explicitly to trigger env var lookup + guardrail = OnyxGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True, timeout=None + ) + + # Verify the client was initialized with timeout from env var + mock_get_client.assert_called_once() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 25.0 + assert timeout_param.connect == 5.0 + + def test_initialization_timeout_parameter_overrides_env_var(self): + """Test that timeout parameter overrides ONYX_TIMEOUT environment variable.""" + os.environ["ONYX_API_KEY"] = "test-api-key" + os.environ["ONYX_TIMEOUT"] = "25" + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + guardrail = OnyxGuardrail( + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + timeout=15.0, + ) + + # Verify the client was initialized with parameter timeout (not env var) + mock_get_client.assert_called_once() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 15.0 + assert timeout_param.connect == 5.0 + @pytest.mark.asyncio async def test_apply_guardrail_request_no_violations(self): """Test apply_guardrail for request with no violations detected.""" @@ -388,6 +587,105 @@ class TestOnyxGuardrail: assert result == inputs + @pytest.mark.asyncio + async def test_apply_guardrail_timeout_error_handling(self): + """Test handling of timeout errors in apply_guardrail (graceful degradation).""" + # Set required API key + os.environ["ONYX_API_KEY"] = "test-api-key" + + guardrail = OnyxGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True, timeout=1.0 + ) + + inputs = GenericGuardrailAPIInputs() + + request_data = { + "proxy_server_request": { + "messages": [{"role": "user", "content": "Test message"}], + "model": "gpt-3.5-turbo", + } + } + + # Test httpx timeout error + with patch.object( + guardrail.async_handler, "post", side_effect=httpx.TimeoutException("Request timed out") + ): + # Should return original inputs on timeout (graceful degradation) + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + logging_obj=None, + ) + + assert result == inputs + + @pytest.mark.asyncio + async def test_apply_guardrail_read_timeout_error_handling(self): + """Test handling of read timeout errors in apply_guardrail.""" + # Set required API key + os.environ["ONYX_API_KEY"] = "test-api-key" + + guardrail = OnyxGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True, timeout=5.0 + ) + + inputs = GenericGuardrailAPIInputs() + + request_data = { + "proxy_server_request": { + "messages": [{"role": "user", "content": "Test message"}], + "model": "gpt-3.5-turbo", + } + } + + # Test httpx ReadTimeout error + with patch.object( + guardrail.async_handler, "post", side_effect=httpx.ReadTimeout("Read timed out") + ): + # Should return original inputs on timeout (graceful degradation) + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + logging_obj=None, + ) + + assert result == inputs + + @pytest.mark.asyncio + async def test_apply_guardrail_connect_timeout_error_handling(self): + """Test handling of connect timeout errors in apply_guardrail.""" + # Set required API key + os.environ["ONYX_API_KEY"] = "test-api-key" + + guardrail = OnyxGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True, timeout=5.0 + ) + + inputs = GenericGuardrailAPIInputs() + + request_data = { + "proxy_server_request": { + "messages": [{"role": "user", "content": "Test message"}], + "model": "gpt-3.5-turbo", + } + } + + # Test httpx ConnectTimeout error + with patch.object( + guardrail.async_handler, "post", side_effect=httpx.ConnectTimeout("Connect timed out") + ): + # Should return original inputs on timeout (graceful degradation) + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + logging_obj=None, + ) + + assert result == inputs + @pytest.mark.asyncio async def test_apply_guardrail_no_logging_obj(self): """Test apply_guardrail without logging object (uses UUID).""" diff --git a/tests/test_litellm/proxy/test_chat_completion_metadata.py b/tests/test_litellm/proxy/test_chat_completion_metadata.py new file mode 100644 index 0000000000..38dcdc13c5 --- /dev/null +++ b/tests/test_litellm/proxy/test_chat_completion_metadata.py @@ -0,0 +1,154 @@ +import pytest +from unittest.mock import MagicMock, AsyncMock, patch +from litellm.proxy.proxy_server import chat_completion, completion, embeddings +from litellm.proxy._types import UserAPIKeyAuth +from fastapi import Request, Response + + +@pytest.mark.asyncio +async def test_chat_completion_metadata_population(): + # Setup + request = MagicMock(spec=Request) + # Mock _read_request_body to return a dict + with patch( + "litellm.proxy.proxy_server._read_request_body", new_callable=AsyncMock + ) as mock_read_body: + mock_read_body.return_value = {"model": "gpt-3.5-turbo", "messages": []} + + user_api_key_dict = UserAPIKeyAuth( + user_id="test_user_id", team_id="test_team_id", org_id="test_org_id" + ) + + fastapi_response = MagicMock(spec=Response) + + # Mock ProxyBaseLLMRequestProcessing + with patch( + "litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + mock_instance = MockProcessor.return_value + mock_instance.base_process_llm_request = AsyncMock( + return_value={"choices": []} + ) + + # Execute + await chat_completion( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + ) + + # Verify + # Check if ProxyBaseLLMRequestProcessing was initialized with data containing metadata + call_args = MockProcessor.call_args + assert call_args is not None + data_arg = call_args.kwargs.get("data") + assert data_arg is not None + + assert "metadata" in data_arg + assert data_arg["metadata"]["user_api_key_user_id"] == "test_user_id" + assert data_arg["metadata"]["user_api_key_team_id"] == "test_team_id" + assert data_arg["metadata"]["user_api_key_org_id"] == "test_org_id" + + +@pytest.mark.asyncio +async def test_embedding_metadata_population(): + """ + Test that the embedding endpoint correctly populates metadata + from UserAPIKeyAuth. + """ + # Setup + with patch( + "litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing.base_process_llm_request" + ): + with patch( + "litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing.__init__", + return_value=None, + ) as mock_base_process_init: + # Create a mock UserAPIKeyAuth object + mock_user_auth = MagicMock(spec=UserAPIKeyAuth) + mock_user_auth.user_id = "test_user_id_emb" + mock_user_auth.team_id = "test_team_id_emb" + mock_user_auth.org_id = "test_org_id_emb" + + # Create a mock Request object + mock_request = MagicMock(spec=Request) + mock_request.json = AsyncMock( + return_value={"model": "gpt-3.5-turbo", "input": "hello"} + ) + # Mock _read_request_body to return our data + with patch( + "litellm.proxy.proxy_server._read_request_body", + new=AsyncMock( + return_value={"model": "gpt-3.5-turbo", "input": "hello"} + ), + ): + # Call the endpoint function directly + await embeddings( + request=mock_request, + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=mock_user_auth, + ) + + # Check if ProxyBaseLLMRequestProcessing was initialized with the correct metadata + mock_base_process_init.assert_called_once() + call_args = mock_base_process_init.call_args + # handle both positional and keyword args for data + if "data" in call_args.kwargs: + data_arg = call_args.kwargs["data"] + else: + data_arg = call_args.args[0] + + assert ( + data_arg["metadata"]["user_api_key_user_id"] == "test_user_id_emb" + ) + assert ( + data_arg["metadata"]["user_api_key_team_id"] == "test_team_id_emb" + ) + assert data_arg["metadata"]["user_api_key_org_id"] == "test_org_id_emb" + + +@pytest.mark.asyncio +async def test_completion_metadata_population(): + # Setup + request = MagicMock(spec=Request) + # Mock _read_request_body to return a dict + with patch( + "litellm.proxy.proxy_server._read_request_body", new_callable=AsyncMock + ) as mock_read_body: + mock_read_body.return_value = { + "model": "gpt-3.5-turbo-instruct", + "prompt": "test", + } + + user_api_key_dict = UserAPIKeyAuth( + user_id="test_user_id_2", team_id="test_team_id_2", org_id="test_org_id_2" + ) + + fastapi_response = MagicMock(spec=Response) + + # Mock ProxyBaseLLMRequestProcessing + with patch( + "litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + mock_instance = MockProcessor.return_value + mock_instance.base_process_llm_request = AsyncMock( + return_value={"choices": []} + ) + + # Execute + await completion( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + ) + + # Verify + call_args = MockProcessor.call_args + assert call_args is not None + data_arg = call_args.kwargs.get("data") + assert data_arg is not None + + assert "metadata" in data_arg + assert data_arg["metadata"]["user_api_key_user_id"] == "test_user_id_2" + assert data_arg["metadata"]["user_api_key_team_id"] == "test_team_id_2" + assert data_arg["metadata"]["user_api_key_org_id"] == "test_org_id_2" diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index ad4f53dac4..60bdb7d12c 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -1227,3 +1227,113 @@ class TestProxySettingEndpoints: assert retrieved_role_mappings["provider"] == "google" assert retrieved_role_mappings["group_claim"] == "groups" assert retrieved_role_mappings["default_role"] == LitellmUserRoles.INTERNAL_USER + + def test_setup_role_mappings_custom_logic_with_env_vars(self, monkeypatch): + """Test the _setup_role_mappings function directly with custom role mapping logic from environment variables""" + import asyncio + import os + from litellm.proxy.management_endpoints.ui_sso import _setup_role_mappings + from litellm.proxy._types import LitellmUserRoles + + # Set up environment variables for custom role mappings using valid Python dict format + monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_ROLES", "{'proxy_admin': ['custom-admin-group'], 'internal_user': ['custom-user-group'], 'proxy_admin_viewer': ['custom-viewer-group']}") + monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", "custom-groups") + monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", "internal_user_viewer") + + # Debug: Print environment variables + print("GENERIC_ROLE_MAPPINGS_ROLES:", os.getenv("GENERIC_ROLE_MAPPINGS_ROLES")) + print("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM:", os.getenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM")) + print("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE:", os.getenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE")) + + # Run the async function + role_mappings = asyncio.run(_setup_role_mappings()) + + # Debug: Print result + print("role_mappings result:", role_mappings) + + # Verify role_mappings is returned correctly from environment variables + assert role_mappings is not None + assert role_mappings.provider == "generic" + assert role_mappings.group_claim == "custom-groups" + assert role_mappings.default_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + assert role_mappings.roles[LitellmUserRoles.PROXY_ADMIN] == ["custom-admin-group"] + assert role_mappings.roles[LitellmUserRoles.INTERNAL_USER] == ["custom-user-group"] + assert role_mappings.roles[LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY] == ["custom-viewer-group"] + + def test_setup_role_mappings_custom_logic_with_no_config(self, monkeypatch): + """Test the _setup_role_mappings function returns None when no configuration is available""" + import asyncio + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy.management_endpoints.ui_sso import _setup_role_mappings + + # Ensure environment variables are not set + monkeypatch.delenv("GENERIC_ROLE_MAPPINGS_ROLES", raising=False) + monkeypatch.delenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", raising=False) + monkeypatch.delenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", raising=False) + + # Mock the prisma client to return None (no database record) + mock_prisma = MagicMock() + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) + # Run the async function + role_mappings = asyncio.run(_setup_role_mappings()) + + # Should return None when no configuration is available + assert role_mappings is None + + def test_get_sso_settings_with_env_role_mappings(self, mock_proxy_config, mock_auth, monkeypatch): + import json + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_ROLES", '{"proxy_admin": ["custom-admin-group"], "internal_user": ["custom-user-group"], "proxy_admin_viewer": ["custom-viewer-group"]}') + monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", "custom-groups") + monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", "internal_user_viewer") + + mock_prisma = MagicMock() + mock_db_record = MagicMock() + mock_db_record.sso_settings = { + "google_client_id": "test_google_client_id", + "role_mappings": { + "provider": "google", + "group_claim": "db-groups", + "default_role": "proxy_admin", + "roles": { + "proxy_admin": ["db-admin-group"], + }, + }, + } + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_db_record) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + from litellm.proxy.proxy_server import proxy_config + monkeypatch.setattr( + proxy_config, "_decrypt_and_set_db_env_variables", lambda environment_variables: environment_variables + ) + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + data = response.json() + + values = data["values"] + assert "role_mappings" in values + assert values["role_mappings"] is not None + + # The database values shoeld override the environment variables + assert values["role_mappings"]["provider"] == "google" + assert values["role_mappings"]["group_claim"] == "db-groups" + assert values["role_mappings"]["default_role"] == LitellmUserRoles.PROXY_ADMIN + assert values["role_mappings"]["roles"][LitellmUserRoles.PROXY_ADMIN] == ["db-admin-group"] + + # Verify that the database was checked but environment variables took priority + mock_prisma.db.litellm_ssoconfig.find_unique.assert_called_once_with( + where={"id": "sso_config"} + ) + + # Verify other SSO settings are still correctly returned + assert values["google_client_id"] == "test_google_client_id" + + # Verify field_schema is still present + assert "field_schema" in data + assert "properties" in data["field_schema"] + assert "role_mappings" in data["field_schema"]["properties"] diff --git a/tests/test_litellm/secret_managers/test_secret_managers_main.py b/tests/test_litellm/secret_managers/test_secret_managers_main.py index cadbd883cc..eaef6956cd 100644 --- a/tests/test_litellm/secret_managers/test_secret_managers_main.py +++ b/tests/test_litellm/secret_managers/test_secret_managers_main.py @@ -48,7 +48,8 @@ def mock_env(): @patch("litellm.secret_managers.main.oidc_cache") @patch("litellm.secret_managers.main._get_oidc_http_handler") -def test_oidc_google_success(mock_get_http_handler, mock_oidc_cache): +@patch("httpx.Client") # Prevent any real HTTP connections +def test_oidc_google_success(mock_httpx_client, mock_get_http_handler, mock_oidc_cache): mock_oidc_cache.get_cache.return_value = None mock_handler = MockHTTPHandler(timeout=600.0) mock_get_http_handler.return_value = mock_handler @@ -144,7 +145,7 @@ def test_oidc_azure_file_success(mock_env, tmp_path): mock_env["AZURE_FEDERATED_TOKEN_FILE"] = str(token_file) secret_name = "oidc/azure/azure-audience" - result = get_secret(secret_name) + result = get_secret(secret_name) assert result == "azure_token" @@ -156,16 +157,22 @@ def test_oidc_azure_ad_token_success(mock_get_azure_ad_token_provider): if "AZURE_FEDERATED_TOKEN_FILE" in os.environ: del os.environ["AZURE_FEDERATED_TOKEN_FILE"] + # Mock the token provider function that gets returned and called mock_token_provider = Mock(return_value="azure_ad_token") mock_get_azure_ad_token_provider.return_value = mock_token_provider - secret_name = "oidc/azure/api://azure-audience" - result = get_secret(secret_name) + + # Also mock the Azure Identity SDK to prevent any real Azure calls + with patch("azure.identity.get_bearer_token_provider") as mock_bearer: + mock_bearer.return_value = mock_token_provider + + secret_name = "oidc/azure/api://azure-audience" + result = get_secret(secret_name) - assert result == "azure_ad_token" - mock_get_azure_ad_token_provider.assert_called_once_with( - azure_scope="api://azure-audience" - ) - mock_token_provider.assert_called_once_with() + assert result == "azure_ad_token" + mock_get_azure_ad_token_provider.assert_called_once_with( + azure_scope="api://azure-audience" + ) + mock_token_provider.assert_called_once_with() def test_oidc_file_success(tmp_path): diff --git a/tests/test_litellm/test_gpt_image_cost_calculator.py b/tests/test_litellm/test_gpt_image_cost_calculator.py index 0a2a62b6c9..620c073498 100644 --- a/tests/test_litellm/test_gpt_image_cost_calculator.py +++ b/tests/test_litellm/test_gpt_image_cost_calculator.py @@ -19,10 +19,13 @@ import pytest import litellm from litellm.types.utils import ( + CompletionTokensDetailsWrapper, ImageResponse, ImageObject, ImageUsage, ImageUsageInputTokensDetails, + PromptTokensDetailsWrapper, + Usage, ) @@ -202,6 +205,71 @@ class TestGPTImageCostRouting: assert cost >= 0 +class TestGPTImage15OutputImageTokens: + """ + Test for GitHub issue #19508: + Image usage calculation does not include image tokens in gpt-image-1.5 + + gpt-image-1.5 returns output_tokens_details with separate image_tokens and text_tokens, + and these must be correctly included in cost calculation. + """ + + def test_gpt_image_15_output_image_tokens_cost(self): + """ + Test that output image tokens are correctly included in cost calculation. + + This tests the fix for issue #19508 where output_tokens_details.image_tokens + were not being included in the cost calculation, causing costs to be + underreported (e.g., $0.046 instead of $0.14). + """ + # Simulate gpt-image-1.5 response with output_tokens_details + # This is what the API returns and what convert_to_image_response transforms + usage = Usage( + prompt_tokens=169, + completion_tokens=4599, + total_tokens=4768, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=169, + image_tokens=0, + ), + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=439, + image_tokens=4160, + ), + ) + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(b64_json="test")], + ) + image_response.usage = usage + image_response._hidden_params = {"custom_llm_provider": "openai"} + + cost = litellm.completion_cost( + completion_response=image_response, + model="gpt-image-1.5", + call_type="image_generation", + custom_llm_provider="openai", + ) + + # gpt-image-1.5 pricing: + # - input_cost_per_token: 5e-06 ($5/1M for text input) + # - output_cost_per_token: 1e-05 ($10/1M for text output) + # - output_cost_per_image_token: 3.2e-05 ($32/1M for image output) + # + # Expected cost: + # Input text: 169 * $5/1M = $0.000845 + # Output text: 439 * $10/1M = $0.00439 + # Output image: 4160 * $32/1M = $0.13312 + # Total: $0.138355 + expected_cost = 169 * 5e-06 + 439 * 1e-05 + 4160 * 3.2e-05 + + assert abs(cost - expected_cost) < 1e-6, ( + f"Expected {expected_cost}, got {cost}. " + f"Image tokens may not be included in cost calculation." + ) + + class TestCompletionCostIntegration: """Test the full completion_cost integration for gpt-image-1""" diff --git a/tests/test_litellm/test_router_per_deployment_num_retries.py b/tests/test_litellm/test_router_per_deployment_num_retries.py index 4021ca2807..154ba579e4 100644 --- a/tests/test_litellm/test_router_per_deployment_num_retries.py +++ b/tests/test_litellm/test_router_per_deployment_num_retries.py @@ -32,17 +32,17 @@ class TestPerDeploymentNumRetries: ) deployment = router.model_list[0] - + # Create a mock exception without num_retries class MockException(Exception): pass - + exc = MockException("test error") assert not hasattr(exc, "num_retries") or exc.num_retries is None - + # Call the helper router._set_deployment_num_retries_on_exception(exc, deployment) - + # Verify num_retries was set from deployment assert exc.num_retries == 5 @@ -66,16 +66,16 @@ class TestPerDeploymentNumRetries: ) deployment = router.model_list[0] - + # Create an exception that already has num_retries class MockException(Exception): num_retries = 10 # Already set - + exc = MockException("test error") - + # Call the helper router._set_deployment_num_retries_on_exception(exc, deployment) - + # Verify num_retries was NOT overridden assert exc.num_retries == 10 @@ -99,15 +99,15 @@ class TestPerDeploymentNumRetries: ) deployment = router.model_list[0] - + class MockException(Exception): pass - + exc = MockException("test error") - + # Call the helper router._set_deployment_num_retries_on_exception(exc, deployment) - + # Verify num_retries was not set (deployment has no num_retries) assert not hasattr(exc, "num_retries") or exc.num_retries is None @@ -155,3 +155,36 @@ class TestPerDeploymentNumRetries: kwargs = {} router._update_kwargs_before_fallbacks(model="test-model", kwargs=kwargs) assert kwargs["num_retries"] == 7 # Uses global + + def test_set_deployment_num_retries_with_string_value(self): + """ + Test that _set_deployment_num_retries_on_exception handles string values + from environment variables correctly. + GitHub Issue: #19481 + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "test-key", + "num_retries": "6", # String value (as from env var) + }, + }, + ], + num_retries=0, # Global setting + ) + + deployment = router.model_list[0] + + class MockException(Exception): + pass + + exc = MockException("test error") + + # Call the helper + router._set_deployment_num_retries_on_exception(exc, deployment) + + # Verify num_retries was converted from string to int + assert exc.num_retries == 6 diff --git a/tests/test_litellm/vector_stores/test_vector_store_registry.py b/tests/test_litellm/vector_stores/test_vector_store_registry.py index 3c4308dd3a..a3af476bc7 100644 --- a/tests/test_litellm/vector_stores/test_vector_store_registry.py +++ b/tests/test_litellm/vector_stores/test_vector_store_registry.py @@ -119,8 +119,12 @@ def test_add_vector_store_to_registry(): +@respx.mock def test_search_uses_registry_credentials(): """search() should pull credentials from vector_store_registry when available""" + # Block all HTTP requests at the network level to prevent real API calls + respx.route().mock(return_value=httpx.Response(200, json={"object": "list", "data": []})) + vector_store = LiteLLM_ManagedVectorStore( vector_store_id="vs1", custom_llm_provider="bedrock", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowPrompts.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowPrompts.ts new file mode 100644 index 0000000000..801fbdbb99 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowPrompts.ts @@ -0,0 +1,35 @@ +// hooks/useDisableShowPrompts.ts +import { useSyncExternalStore } from "react"; +import { getLocalStorageItem } from "@/utils/localStorageUtils"; +import { LOCAL_STORAGE_EVENT } from "@/utils/localStorageUtils"; + +function subscribe(callback: () => void) { + const onStorage = (e: StorageEvent) => { + if (e.key === "disableShowPrompts") { + callback(); + } + }; + + const onCustom = (e: Event) => { + const { key } = (e as CustomEvent).detail; + if (key === "disableShowPrompts") { + callback(); + } + }; + + window.addEventListener("storage", onStorage); + window.addEventListener(LOCAL_STORAGE_EVENT, onCustom); + + return () => { + window.removeEventListener("storage", onStorage); + window.removeEventListener(LOCAL_STORAGE_EVENT, onCustom); + }; +} + +function getSnapshot() { + return getLocalStorageItem("disableShowPrompts") === "true"; +} + +export function useDisableShowPrompts() { + return useSyncExternalStore(subscribe, getSnapshot); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 97837ff8e0..97e4c799e7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -56,8 +56,10 @@ export default function Layout({ children }: { children: React.ReactNode }) { userRole={userRole} premiumUser={premiumUser} proxySettings={undefined} - setProxySettings={() => {}} + setProxySettings={() => { }} accessToken={accessToken} + isDarkMode={false} + toggleDarkMode={() => { }} />
diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 8ac1f756f9..23c80acf97 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -45,6 +45,7 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { jwtDecode } from "jwt-decode"; import { useSearchParams } from "next/navigation"; import { Suspense, useEffect, useState } from "react"; +import { ConfigProvider, theme } from "antd"; function getCookie(name: string) { // Safer cookie read + decoding; handles '=' inside values @@ -131,6 +132,12 @@ export default function CreateKeyPage() { const [showClaudeCodePrompt, setShowClaudeCodePrompt] = useState(false); const [showClaudeCodeModal, setShowClaudeCodeModal] = useState(false); + // Dark mode state + const [isDarkMode, setIsDarkMode] = useState(false); + const toggleDarkMode = () => { + setIsDarkMode(!isDarkMode); + }; + const invitation_id = searchParams.get("invitation_id"); // Get page from URL, default to 'api-keys' if not present @@ -282,7 +289,7 @@ export default function CreateKeyPage() { 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); @@ -362,225 +369,231 @@ export default function CreateKeyPage() { return ( }> - - {invitation_id ? ( - - ) : ( -
- + + {invitation_id ? ( + -
-
- -
+ ) : ( +
+ +
+
+ +
- {page == "api-keys" ? ( - - ) : page == "models" ? ( - - ) : page == "llm-playground" ? ( - - ) : page == "users" ? ( - - ) : page == "teams" ? ( - - ) : page == "organizations" ? ( - - ) : page == "admin-panel" ? ( - - ) : page == "api_ref" ? ( - - ) : page == "logging-and-alerts" ? ( - - ) : page == "budgets" ? ( - - ) : page == "guardrails" ? ( - - ) : page == "policies" ? ( - - ) : page == "agents" ? ( - - ) : page == "prompts" ? ( - - ) : page == "transform-request" ? ( - - ) : page == "router-settings" ? ( - - ) : page == "ui-theme" ? ( - - ) : page == "cost-tracking" ? ( - - ) : page == "model-hub-table" ? ( - isAdminRole(userRole) ? ( - + ) : page == "models" ? ( + + ) : page == "llm-playground" ? ( + + ) : page == "users" ? ( + + ) : page == "teams" ? ( + + ) : page == "organizations" ? ( + + ) : page == "admin-panel" ? ( + + ) : page == "api_ref" ? ( + + ) : page == "logging-and-alerts" ? ( + + ) : page == "budgets" ? ( + + ) : page == "guardrails" ? ( + + ) : page == "policies" ? ( + + ) : page == "agents" ? ( + + ) : page == "prompts" ? ( + + ) : page == "transform-request" ? ( + + ) : page == "router-settings" ? ( + + ) : page == "ui-theme" ? ( + + ) : page == "cost-tracking" ? ( + + ) : page == "model-hub-table" ? ( + isAdminRole(userRole) ? ( + + ) : ( + + ) + ) : page == "caching" ? ( + + ) : page == "pass-through-settings" ? ( + + ) : page == "logs" ? ( + + ) : page == "mcp-servers" ? ( + + ) : page == "search-tools" ? ( + + ) : page == "tag-management" ? ( + + ) : page == "claude-code-plugins" ? ( + + ) : page == "vector-stores" ? ( + + ) : page == "new_usage" ? ( + ) : ( - - ) - ) : page == "caching" ? ( - - ) : page == "pass-through-settings" ? ( - - ) : page == "logs" ? ( - - ) : page == "mcp-servers" ? ( - - ) : page == "search-tools" ? ( - - ) : page == "tag-management" ? ( - - ) : page == "claude-code-plugins" ? ( - - ) : page == "vector-stores" ? ( - - ) : page == "new_usage" ? ( - - ) : ( - - )} + + )} +
+ + {/* Survey Components */} + + + + {/* Claude Code Components */} + +
- - {/* Survey Components */} - - - - {/* Claude Code Components */} - - -
- )} - + )} + + ); diff --git a/ui/litellm-dashboard/src/components/BulkEditUsers.test.tsx b/ui/litellm-dashboard/src/components/BulkEditUsers.test.tsx new file mode 100644 index 0000000000..4185625e74 --- /dev/null +++ b/ui/litellm-dashboard/src/components/BulkEditUsers.test.tsx @@ -0,0 +1,343 @@ +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi, beforeEach } from "vitest"; +import { renderWithProviders, screen, waitFor } from "../../tests/test-utils"; +import BulkEditUserModal from "./BulkEditUsers"; +import { userBulkUpdateUserCall, teamBulkMemberAddCall } from "./networking"; +import NotificationsManager from "./molecules/notifications_manager"; + +vi.mock("./networking", () => ({ + userBulkUpdateUserCall: vi.fn(), + teamBulkMemberAddCall: vi.fn(), +})); + +vi.mock("./user_edit_view", () => ({ + UserEditView: ({ onSubmit, onCancel }: { onSubmit: (values: any) => void; onCancel: () => void }) => ( +
+ + +
+ ), +})); + +const mockUserBulkUpdateUserCall = vi.mocked(userBulkUpdateUserCall); +const mockTeamBulkMemberAddCall = vi.mocked(teamBulkMemberAddCall); + +const defaultProps = { + open: true, + onCancel: vi.fn(), + selectedUsers: [ + { user_id: "user1", user_email: "user1@example.com", user_role: "user", max_budget: 50 }, + { user_id: "user2", user_email: "user2@example.com", user_role: "admin", max_budget: null }, + ], + possibleUIRoles: { + admin: { ui_label: "Admin", description: "Administrator role" }, + user: { ui_label: "User", description: "Regular user role" }, + }, + accessToken: "test-token", + onSuccess: vi.fn(), + teams: [ + { team_id: "team1", team_alias: "Team 1" }, + { team_id: "team2", team_alias: "Team 2" }, + ], + userRole: "Admin", + userModels: ["gpt-4", "gpt-3.5-turbo"], + allowAllUsers: false, +}; + +describe("BulkEditUserModal", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUserBulkUpdateUserCall.mockResolvedValue({ + results: [], + total_requested: 2, + successful_updates: 2, + failed_updates: 0, + }); + mockTeamBulkMemberAddCall.mockResolvedValue({ + successful_additions: 2, + failed_additions: 0, + }); + }); + + it("should render without crashing", () => { + renderWithProviders(); + + expect(screen.getByText(`Bulk Edit ${defaultProps.selectedUsers.length} User(s)`)).toBeInTheDocument(); + }); + + it("should display modal title with correct user count", () => { + renderWithProviders(); + + expect(screen.getByText("Bulk Edit 2 User(s)")).toBeInTheDocument(); + }); + + it("should display selected users table when modal is open", () => { + renderWithProviders(); + + expect(screen.getByText("Selected Users (2):")).toBeInTheDocument(); + expect(screen.getByText("user1")).toBeInTheDocument(); + expect(screen.getByText("user2")).toBeInTheDocument(); + expect(screen.getByText("user1@example.com")).toBeInTheDocument(); + expect(screen.getByText("user2@example.com")).toBeInTheDocument(); + }); + + it("should display user roles in table", () => { + renderWithProviders(); + + expect(screen.getByText("User")).toBeInTheDocument(); + expect(screen.getByText("Admin")).toBeInTheDocument(); + }); + + it("should display budget information in table", () => { + renderWithProviders(); + + expect(screen.getByText("$50")).toBeInTheDocument(); + expect(screen.getByText("Unlimited")).toBeInTheDocument(); + }); + + it("should call onCancel when cancel button is clicked", async () => { + const user = userEvent.setup(); + const onCancel = vi.fn(); + renderWithProviders(); + + const cancelButton = screen.getByRole("button", { name: "Cancel" }); + await user.click(cancelButton); + + expect(onCancel).toHaveBeenCalledTimes(1); + }); + + + it("should show update all users checkbox when allowAllUsers is true", () => { + renderWithProviders(); + + expect(screen.getByRole("checkbox", { name: /update all users/i })).toBeInTheDocument(); + }); + + it("should not show update all users checkbox when allowAllUsers is false", () => { + renderWithProviders(); + + expect(screen.queryByRole("checkbox", { name: /update all users/i })).not.toBeInTheDocument(); + }); + + it("should toggle update all users mode", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const checkbox = screen.getByRole("checkbox", { name: /update all users/i }); + expect(checkbox).not.toBeChecked(); + + await user.click(checkbox); + + expect(checkbox).toBeChecked(); + expect(screen.getByText("Bulk Edit All Users")).toBeInTheDocument(); + }); + + it("should show warning message when update all users is enabled", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const checkbox = screen.getByRole("checkbox", { name: /update all users/i }); + await user.click(checkbox); + + expect(screen.getByText(/this will apply changes to all users/i)).toBeInTheDocument(); + }); + + it("should hide selected users table when update all users is enabled", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + expect(screen.getByText("Selected Users (2):")).toBeInTheDocument(); + + const checkbox = screen.getByRole("checkbox", { name: /update all users/i }); + await user.click(checkbox); + + expect(screen.queryByText("Selected Users (2):")).not.toBeInTheDocument(); + }); + + it("should display team management section", () => { + renderWithProviders(); + + expect(screen.getByText("Team Management")).toBeInTheDocument(); + expect(screen.getByRole("checkbox", { name: /add selected users to teams/i })).toBeInTheDocument(); + }); + + it("should show team budget input when add to teams is checked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const addToTeamsCheckbox = screen.getByRole("checkbox", { name: /add selected users to teams/i }); + await user.click(addToTeamsCheckbox); + + expect(screen.getByText("Team Budget (Optional):")).toBeInTheDocument(); + expect(screen.getByPlaceholderText("Max budget per user in team")).toBeInTheDocument(); + }); + + it("should render UserEditView component", () => { + renderWithProviders(); + + expect(screen.getByTestId("user-edit-view")).toBeInTheDocument(); + }); + + it("should show error when access token is missing", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const submitButton = screen.getByRole("button", { name: "Submit" }); + await user.click(submitButton); + + await waitFor(() => { + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith("Access token not found"); + }); + }); + + it("should call userBulkUpdateUserCall with correct payload for selected users", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const submitButton = screen.getByRole("button", { name: "Submit" }); + await user.click(submitButton); + + await waitFor(() => { + expect(mockUserBulkUpdateUserCall).toHaveBeenCalledWith( + "test-token", + { user_role: "admin", max_budget: 100 }, + ["user1", "user2"], + ); + }); + }); + + it("should call userBulkUpdateUserCall with allUsers flag when update all users is enabled", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const updateAllCheckbox = screen.getByRole("checkbox", { name: /update all users/i }); + await user.click(updateAllCheckbox); + + const submitButton = screen.getByRole("button", { name: "Submit" }); + await user.click(submitButton); + + await waitFor(() => { + expect(mockUserBulkUpdateUserCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ user_role: "admin", max_budget: 100 }), + undefined, + true, + ); + }); + }); + + + it("should show success message after successful user update", async () => { + const user = userEvent.setup(); + mockUserBulkUpdateUserCall.mockResolvedValue({ + results: [], + total_requested: 2, + successful_updates: 2, + failed_updates: 0, + }); + + renderWithProviders(); + + const submitButton = screen.getByRole("button", { name: "Submit" }); + await user.click(submitButton); + + await waitFor(() => { + expect(NotificationsManager.success).toHaveBeenCalledWith("Updated 2 user(s)"); + }); + }); + + it("should show success message for all users update", async () => { + const user = userEvent.setup(); + mockUserBulkUpdateUserCall.mockResolvedValue({ + results: [], + total_requested: 100, + successful_updates: 100, + failed_updates: 0, + }); + + renderWithProviders(); + + const updateAllCheckbox = screen.getByRole("checkbox", { name: /update all users/i }); + await user.click(updateAllCheckbox); + + const submitButton = screen.getByRole("button", { name: "Submit" }); + await user.click(submitButton); + + await waitFor(() => { + expect(NotificationsManager.success).toHaveBeenCalledWith("Updated all users (100 total)"); + }); + }); + + + it("should show error message when bulk update fails", async () => { + const user = userEvent.setup(); + mockUserBulkUpdateUserCall.mockRejectedValueOnce(new Error("Update failed")); + + renderWithProviders(); + + const submitButton = screen.getByRole("button", { name: "Submit" }); + await user.click(submitButton); + + await waitFor(() => { + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith("Failed to perform bulk operations"); + }); + }); + + it("should call onSuccess and onCancel after successful update", async () => { + const user = userEvent.setup(); + const onSuccess = vi.fn(); + const onCancel = vi.fn(); + + renderWithProviders(); + + const submitButton = screen.getByRole("button", { name: "Submit" }); + await user.click(submitButton); + + await waitFor(() => { + expect(onSuccess).toHaveBeenCalledTimes(1); + expect(onCancel).toHaveBeenCalledTimes(1); + }); + }); + + it("should truncate long user IDs in table", () => { + const longUserId = "a".repeat(30); + const propsWithLongId = { + ...defaultProps, + selectedUsers: [{ user_id: longUserId, user_email: "test@example.com", user_role: "user", max_budget: null }], + }; + + renderWithProviders(); + + expect(screen.getByText(new RegExp(`${longUserId.slice(0, 20)}...`))).toBeInTheDocument(); + }); + + it("should display no email text when user email is missing", () => { + const propsWithoutEmail = { + ...defaultProps, + selectedUsers: [{ user_id: "user1", user_email: null, user_role: "user", max_budget: null }], + }; + + renderWithProviders(); + + expect(screen.getByText("No email")).toBeInTheDocument(); + }); + + it("should display role label from possibleUIRoles when available", () => { + renderWithProviders(); + + expect(screen.getByText("Admin")).toBeInTheDocument(); + expect(screen.getByText("User")).toBeInTheDocument(); + }); + + it("should display role key when ui_label is not available", () => { + const propsWithoutUIRoles = { + ...defaultProps, + possibleUIRoles: null, + }; + + renderWithProviders(); + + expect(screen.getByText("user")).toBeInTheDocument(); + expect(screen.getByText("admin")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/bulk_edit_user.tsx b/ui/litellm-dashboard/src/components/BulkEditUsers.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/bulk_edit_user.tsx rename to ui/litellm-dashboard/src/components/BulkEditUsers.tsx index b3847911cc..2f3e57ff2a 100644 --- a/ui/litellm-dashboard/src/components/bulk_edit_user.tsx +++ b/ui/litellm-dashboard/src/components/BulkEditUsers.tsx @@ -18,7 +18,7 @@ import NotificationsManager from "./molecules/notifications_manager"; const { Text, Title } = Typography; interface BulkEditUserModalProps { - visible: boolean; + open: boolean; onCancel: () => void; selectedUsers: any[]; possibleUIRoles: Record> | null; @@ -31,7 +31,7 @@ interface BulkEditUserModalProps { } const BulkEditUserModal: React.FC = ({ - visible, + open, onCancel, selectedUsers, possibleUIRoles, @@ -75,7 +75,7 @@ const BulkEditUserModal: React.FC = ({ keys: [], teams: teams || [], }), - [teams, visible], + [teams, open], ); const handleSubmit = async (formValues: any) => { @@ -145,7 +145,7 @@ const BulkEditUserModal: React.FC = ({ if (updateAllUsers) { members = null; } else { - const members = selectedUsers.map((user) => ({ + members = selectedUsers.map((user) => ({ user_id: user.user_id, role: "user" as const, // Default role for bulk add user_email: user.user_email || null, @@ -214,7 +214,7 @@ const BulkEditUserModal: React.FC = ({ return ( { + const Option = ({ children, value }: any) => ( + + ); + const Select = ({ children, value, onChange, placeholder }: any) => ( + + ); + Select.Option = Option; + return { + Select, + Tooltip: ({ children, title }: any) => ( +
+ {children} +
+ ), + Switch: ({ checked, onChange }: any) => ( + onChange(e.target.checked)} + /> + ), + Divider: () =>
, + }; +}); + +vi.mock("@ant-design/icons", () => ({ + InfoCircleOutlined: () => ℹ, +})); + +vi.mock("@tremor/react", () => ({ + TextInput: ({ value, onValueChange, onChange, placeholder, name, className }: any) => { + const handleChange = (e: React.ChangeEvent) => { + if (onChange) { + onChange(e); + } + if (onValueChange) { + onValueChange(e.target.value); + } + }; + return ( + + ); + }, +})); + +describe("KeyLifecycleSettings", () => { + const mockForm = { + getFieldValue: vi.fn(), + setFieldValue: vi.fn(), + setFieldsValue: vi.fn(), + }; + + const defaultProps = { + form: mockForm, + autoRotationEnabled: false, + onAutoRotationChange: vi.fn(), + rotationInterval: "", + onRotationIntervalChange: vi.fn(), + isCreateMode: false, + }; + + beforeEach(() => { + vi.clearAllMocks(); + mockForm.getFieldValue.mockReturnValue(""); + }); + + it("should render without crashing", () => { + renderWithProviders(); + + expect(screen.getByText("Key Expiry Settings")).toBeInTheDocument(); + expect(screen.getByText("Auto-Rotation Settings")).toBeInTheDocument(); + }); + + describe("Key Expiry Settings", () => { + it("should render expiry input field", () => { + renderWithProviders(); + + expect(screen.getByText("Expire Key")).toBeInTheDocument(); + expect(screen.getByTestId("duration-input")).toBeInTheDocument(); + }); + + it("should show correct placeholder in create mode", () => { + renderWithProviders(); + + const input = screen.getByTestId("duration-input"); + expect(input).toHaveAttribute( + "placeholder", + "e.g., 30d or leave empty to never expire" + ); + }); + + it("should show correct placeholder in edit mode", () => { + renderWithProviders(); + + const input = screen.getByTestId("duration-input"); + expect(input).toHaveAttribute("placeholder", "e.g., 30d or -1 to never expire"); + }); + + it("should show correct tooltip in create mode", () => { + renderWithProviders(); + + const tooltips = screen.getAllByTestId("tooltip"); + const expiryTooltip = tooltips.find((tooltip) => + tooltip.getAttribute("title")?.includes("Leave empty to never expire") + ); + expect(expiryTooltip).toBeInTheDocument(); + expect(expiryTooltip).toHaveAttribute( + "title", + "Set when this key should expire. Format: 30s (seconds), 30m (minutes), 30h (hours), 30d (days). Leave empty to never expire." + ); + }); + + it("should show correct tooltip in edit mode", () => { + renderWithProviders(); + + const tooltips = screen.getAllByTestId("tooltip"); + const expiryTooltip = tooltips.find((tooltip) => + tooltip.getAttribute("title")?.includes("Use -1 to never expire") + ); + expect(expiryTooltip).toBeInTheDocument(); + expect(expiryTooltip).toHaveAttribute( + "title", + "Set when this key should expire. Format: 30s (seconds), 30m (minutes), 30h (hours), 30d (days). Use -1 to never expire." + ); + }); + + it("should initialize with form value if present", () => { + mockForm.getFieldValue.mockReturnValue("30d"); + renderWithProviders(); + + const input = screen.getByTestId("duration-input") as HTMLInputElement; + expect(input.value).toBe("30d"); + }); + + it("should update form using setFieldValue when duration changes", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const input = screen.getByTestId("duration-input"); + await user.type(input, "60d"); + + expect(mockForm.setFieldValue).toHaveBeenCalledWith("duration", "60d"); + }); + + it("should update form using setFieldsValue when setFieldValue is not available", async () => { + const user = userEvent.setup(); + const formWithoutSetFieldValue = { + getFieldValue: vi.fn().mockReturnValue(""), + setFieldsValue: vi.fn(), + }; + renderWithProviders( + + ); + + const input = screen.getByTestId("duration-input"); + await user.type(input, "90d"); + + expect(formWithoutSetFieldValue.setFieldsValue).toHaveBeenCalledWith({ duration: "90d" }); + }); + }); + + describe("Auto-Rotation Settings", () => { + it("should render auto-rotation switch", () => { + renderWithProviders(); + + expect(screen.getByText("Enable Auto-Rotation")).toBeInTheDocument(); + expect(screen.getByTestId("switch")).toBeInTheDocument(); + }); + + it("should show switch as unchecked when autoRotationEnabled is false", () => { + renderWithProviders(); + + const switchElement = screen.getByTestId("switch") as HTMLInputElement; + expect(switchElement.checked).toBe(false); + }); + + it("should show switch as checked when autoRotationEnabled is true", () => { + renderWithProviders(); + + const switchElement = screen.getByTestId("switch") as HTMLInputElement; + expect(switchElement.checked).toBe(true); + }); + + it("should call onAutoRotationChange when switch is toggled", async () => { + const user = userEvent.setup(); + const onAutoRotationChange = vi.fn(); + renderWithProviders( + + ); + + const switchElement = screen.getByTestId("switch"); + await user.click(switchElement); + + expect(onAutoRotationChange).toHaveBeenCalledWith(true); + }); + + it("should not show rotation interval section when auto-rotation is disabled", () => { + renderWithProviders(); + + expect(screen.queryByText("Rotation Interval")).not.toBeInTheDocument(); + expect(screen.queryByTestId("select")).not.toBeInTheDocument(); + }); + + it("should show rotation interval section when auto-rotation is enabled", () => { + renderWithProviders( + + ); + + expect(screen.getByText("Rotation Interval")).toBeInTheDocument(); + expect(screen.getByTestId("select")).toBeInTheDocument(); + }); + + it("should show all predefined interval options", () => { + renderWithProviders( + + ); + + expect(screen.getByText("7 days")).toBeInTheDocument(); + expect(screen.getByText("30 days")).toBeInTheDocument(); + expect(screen.getByText("90 days")).toBeInTheDocument(); + expect(screen.getByText("180 days")).toBeInTheDocument(); + expect(screen.getByText("365 days")).toBeInTheDocument(); + expect(screen.getByText("Custom interval")).toBeInTheDocument(); + }); + + it("should display current rotation interval in select", () => { + renderWithProviders( + + ); + + const select = screen.getByTestId("select") as HTMLSelectElement; + expect(select.value).toBe("90d"); + }); + + it("should call onRotationIntervalChange when predefined interval is selected", async () => { + const user = userEvent.setup(); + const onRotationIntervalChange = vi.fn(); + renderWithProviders( + + ); + + const select = screen.getByTestId("select"); + await user.selectOptions(select, "30d"); + + expect(onRotationIntervalChange).toHaveBeenCalledWith("30d"); + }); + + it("should show custom input when custom option is selected", async () => { + const user = userEvent.setup(); + renderWithProviders( + + ); + + const select = screen.getByTestId("select"); + await user.selectOptions(select, "custom"); + + expect(screen.getByTestId("custom-interval-input")).toBeInTheDocument(); + expect(screen.getByText("Supported formats: seconds (s), minutes (m), hours (h), days (d)")).toBeInTheDocument(); + }); + + it("should hide custom input when predefined interval is selected after custom", async () => { + const user = userEvent.setup(); + const onRotationIntervalChange = vi.fn(); + renderWithProviders( + + ); + + const select = screen.getByTestId("select"); + await user.selectOptions(select, "7d"); + + expect(screen.queryByTestId("custom-interval-input")).not.toBeInTheDocument(); + expect(onRotationIntervalChange).toHaveBeenCalledWith("7d"); + }); + + it("should call onRotationIntervalChange when custom interval is entered", async () => { + const user = userEvent.setup(); + const onRotationIntervalChange = vi.fn(); + renderWithProviders( + + ); + + const select = screen.getByTestId("select"); + await user.selectOptions(select, "custom"); + + const customInput = screen.getByTestId("custom-interval-input"); + await user.type(customInput, "14d"); + + expect(onRotationIntervalChange).toHaveBeenCalledWith("14d"); + }); + + it("should show info message when auto-rotation is enabled", () => { + renderWithProviders(); + + expect( + screen.getByText( + "When rotation occurs, you'll receive a notification with the new key. The old key will be deactivated after a brief grace period." + ) + ).toBeInTheDocument(); + }); + + it("should not show info message when auto-rotation is disabled", () => { + renderWithProviders(); + + expect( + screen.queryByText( + "When rotation occurs, you'll receive a notification with the new key. The old key will be deactivated after a brief grace period." + ) + ).not.toBeInTheDocument(); + }); + + it("should initialize with custom interval input visible when custom interval is provided", () => { + renderWithProviders( + + ); + + expect(screen.getByTestId("custom-interval-input")).toBeInTheDocument(); + const customInput = screen.getByTestId("custom-interval-input") as HTMLInputElement; + expect(customInput.value).toBe("14d"); + }); + + it("should show custom option selected when custom interval is provided", () => { + renderWithProviders( + + ); + + const select = screen.getByTestId("select") as HTMLSelectElement; + expect(select.value).toBe("custom"); + }); + + it("should not call onRotationIntervalChange when selecting custom option", async () => { + const user = userEvent.setup(); + const onRotationIntervalChange = vi.fn(); + renderWithProviders( + + ); + + const select = screen.getByTestId("select"); + await user.selectOptions(select, "custom"); + + expect(onRotationIntervalChange).not.toHaveBeenCalled(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx b/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx index 81d22b5634..0f29a47d1d 100644 --- a/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx +++ b/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx @@ -11,6 +11,7 @@ interface KeyLifecycleSettingsProps { onAutoRotationChange: (enabled: boolean) => void; rotationInterval: string; onRotationIntervalChange: (interval: string) => void; + isCreateMode?: boolean; // If true, shows "leave empty to never expire" instead of "-1 to never expire" } const KeyLifecycleSettings: React.FC = ({ @@ -19,6 +20,7 @@ const KeyLifecycleSettings: React.FC = ({ onAutoRotationChange, rotationInterval, onRotationIntervalChange, + isCreateMode = false, }) => { // Predefined intervals const predefinedIntervals = ["7d", "30d", "90d", "180d", "365d"]; @@ -64,13 +66,19 @@ const KeyLifecycleSettings: React.FC = ({
({ let mockUseThemeImpl = () => ({ logoUrl: null as string | null }); let mockUseHealthReadinessImpl = () => ({ data: null as any }); let mockGetLocalStorageItemImpl = () => null as string | null; +let mockUseDisableShowPromptsImpl = () => false; vi.mock("@/contexts/ThemeContext", () => ({ useTheme: () => mockUseThemeImpl(), @@ -25,7 +26,12 @@ vi.mock("@/app/(dashboard)/hooks/healthReadiness/useHealthReadiness", () => ({ useHealthReadiness: () => mockUseHealthReadinessImpl(), })); +vi.mock("@/app/(dashboard)/hooks/useDisableShowPrompts", () => ({ + useDisableShowPrompts: () => mockUseDisableShowPromptsImpl(), +})); + vi.mock("@/utils/localStorageUtils", () => ({ + LOCAL_STORAGE_EVENT: "local-storage-change", getLocalStorageItem: () => mockGetLocalStorageItemImpl(), setLocalStorageItem: vi.fn(), removeLocalStorageItem: vi.fn(), @@ -52,6 +58,8 @@ describe("Navbar", () => { setProxySettings: vi.fn(), accessToken: "test-token", isPublicPage: false, + isDarkMode: false, + toggleDarkMode: vi.fn(), }; it("should render without crashing", () => { @@ -198,4 +206,11 @@ describe("Navbar", () => { expect(cookieUtils.clearTokenCookies).toHaveBeenCalled(); expect(window.location.href).toBe(""); }); + + it("should not render dark mode toggle slider", () => { + renderWithProviders(); + + // DO NOT RENDER THIS UNTIL ALL COMPONENTS ARE CONFIRMED TO SUPPORT DARK MODE STYLES. IT IS AN ISSUE IF THIS TEST FAILS. + expect(screen.queryByTestId("dark-mode-toggle")).not.toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index 6dac073b3a..c78f355ff1 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -1,4 +1,5 @@ import { useHealthReadiness } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadiness"; +import { useDisableShowPrompts } from "@/app/(dashboard)/hooks/useDisableShowPrompts"; import { getProxyBaseUrl } from "@/components/networking"; import { useTheme } from "@/contexts/ThemeContext"; import { clearTokenCookies } from "@/utils/cookieUtils"; @@ -16,8 +17,10 @@ import { MailOutlined, MenuFoldOutlined, MenuUnfoldOutlined, + MoonOutlined, SafetyOutlined, SlackOutlined, + SunOutlined, UserOutlined, } from "@ant-design/icons"; import type { MenuProps } from "antd"; @@ -36,6 +39,8 @@ interface NavbarProps { isPublicPage: boolean; sidebarCollapsed?: boolean; onToggleSidebar?: () => void; + isDarkMode: boolean; + toggleDarkMode: () => void; } const Navbar: React.FC = ({ @@ -49,11 +54,13 @@ const Navbar: React.FC = ({ isPublicPage = false, sidebarCollapsed = false, onToggleSidebar, + isDarkMode, + toggleDarkMode }) => { const baseUrl = getProxyBaseUrl(); - console.log("baseUrl", baseUrl); const [logoutUrl, setLogoutUrl] = useState(""); const [disableShowNewBadge, setDisableShowNewBadge] = useState(false); + const disableShowPrompts = useDisableShowPrompts(); const { logoUrl } = useTheme(); const { data: healthData } = useHealthReadiness(); const version = healthData?.litellm_version; @@ -152,6 +159,27 @@ const Navbar: React.FC = ({ aria-label="Toggle hide new feature indicators" />
+
e.stopPropagation()} + > + Hide All Prompts + { + if (checked) { + setLocalStorageItem("disableShowPrompts", "true"); + emitLocalStorageChange("disableShowPrompts"); + } else { + removeLocalStorageItem("disableShowPrompts"); + emitLocalStorageChange("disableShowPrompts"); + } + }} + aria-label="Toggle hide all prompts" + /> +
), @@ -227,6 +255,15 @@ const Navbar: React.FC = ({ > Star us on GitHub + {/* Dark mode is currently a work in progress. To test, you can change 'false' to 'true' below. + Do not set this to true by default until all components are confirmed to support dark mode styles. */} + {false && } + unCheckedChildren={} + />} => { + try { + let url = proxyBaseUrl ? `${proxyBaseUrl}/rag/ingest` : `/rag/ingest`; + + const formData = new FormData(); + formData.append("file", file); + + const ingestOptions: any = { + ingest_options: { + vector_store: { + custom_llm_provider: customLlmProvider, + ...(vectorStoreId && { vector_store_id: vectorStoreId }), + }, + }, + }; + + // Add litellm_vector_store_params if name or description provided + if (vectorStoreName || vectorStoreDescription) { + ingestOptions.ingest_options.litellm_vector_store_params = {}; + if (vectorStoreName) { + ingestOptions.ingest_options.litellm_vector_store_params.vector_store_name = vectorStoreName; + } + if (vectorStoreDescription) { + ingestOptions.ingest_options.litellm_vector_store_params.vector_store_description = vectorStoreDescription; + } + } + + formData.append("request", JSON.stringify(ingestOptions)); + + const response = await fetch(url, { + method: "POST", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + }, + body: formData, + }); + + if (!response.ok) { + const error = await response.json(); + throw new Error(error.error?.message || error.detail || "Failed to ingest document"); + } + + return await response.json(); + } catch (error) { + console.error("Error ingesting document:", error); + throw error; + } +}; + export const getEmailEventSettings = async (accessToken: string): Promise => { try { const url = proxyBaseUrl ? `${proxyBaseUrl}/email/event_settings` : `/email/event_settings`; diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 1edbba28af..8e89b77bc1 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -162,7 +162,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { const [userSearchLoading, setUserSearchLoading] = useState(false); const [mcpAccessGroups, setMcpAccessGroups] = useState([]); const [disabledCallbacks, setDisabledCallbacks] = useState([]); - const [keyType, setKeyType] = useState("default"); + const [keyType, setKeyType] = useState("llm_api"); const [modelAliases, setModelAliases] = useState<{ [key: string]: string }>({}); const [autoRotationEnabled, setAutoRotationEnabled] = useState(false); const [rotationInterval, setRotationInterval] = useState("30d"); @@ -173,7 +173,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { form.resetFields(); setLoggingSettings([]); setDisabledCallbacks([]); - setKeyType("default"); + setKeyType("llm_api"); setModelAliases({}); setAutoRotationEnabled(false); setRotationInterval("30d"); @@ -188,7 +188,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { form.resetFields(); setLoggingSettings([]); setDisabledCallbacks([]); - setKeyType("default"); + setKeyType("llm_api"); setModelAliases({}); setAutoRotationEnabled(false); setRotationInterval("30d"); @@ -321,9 +321,9 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { formValues.rotation_interval = rotationInterval; } - // Handle duration field for key expiry - if (formValues.duration) { - formValues.duration = formValues.duration; + // Handle duration field for key expiry - convert empty string to null + if (!formValues.duration || formValues.duration.trim() === "") { + formValues.duration = null; } // Update the formValues with the final metadata @@ -707,11 +707,11 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { } name="key_type" - initialValue="default" + initialValue="llm_api" className="mt-4" > setVectorStoreName(e.target.value)} + placeholder="e.g., Product Documentation, Customer Support KB" + size="large" + className="rounded-md" + /> + + + + Description{" "} + + + + + } + > + setVectorStoreDescription(e.target.value)} + placeholder="e.g., Contains all product documentation and user guides" + rows={2} + size="large" + className="rounded-md" + /> + + + + Provider{" "} + + + + + } + required + > + + + + +
+ +
+
+ + + {/* Success Message */} + {ingestResults.length > 0 && ( + +

+ Vector Store ID: {ingestResults[0]?.vector_store_id} +

+

+ Documents Ingested: {ingestResults.length} +

+ + } + type="success" + showIcon + closable + /> + )} + + ); +}; + +export default CreateVectorStore; diff --git a/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.test.tsx new file mode 100644 index 0000000000..761bc2b7df --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.test.tsx @@ -0,0 +1,102 @@ +import { render, screen, fireEvent, act } from "@testing-library/react"; +import { describe, it, expect, vi } from "vitest"; +import DocumentsTable from "./DocumentsTable"; +import { DocumentUpload } from "./types"; + +// Mock antd message +vi.mock("antd", async () => { + const actual = await vi.importActual("antd"); + return { + ...actual, + message: { + success: vi.fn(), + }, + }; +}); + +describe("DocumentsTable", () => { + const mockDocuments: DocumentUpload[] = [ + { + uid: "1", + name: "test1.pdf", + status: "done", + size: 1024000, + type: "application/pdf", + }, + { + uid: "2", + name: "test2.txt", + status: "uploading", + size: 2048000, + type: "text/plain", + }, + { + uid: "3", + name: "test3.docx", + status: "error", + size: 512000, + type: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + }, + ]; + + it("should render the table successfully", () => { + const onRemove = vi.fn(); + render(); + + expect(screen.getByText("test1.pdf")).toBeInTheDocument(); + expect(screen.getByText("test2.txt")).toBeInTheDocument(); + expect(screen.getByText("test3.docx")).toBeInTheDocument(); + }); + + it("should display correct status badges", () => { + const onRemove = vi.fn(); + render(); + + expect(screen.getByText("Ready")).toBeInTheDocument(); + expect(screen.getByText("Uploading")).toBeInTheDocument(); + expect(screen.getByText("Error")).toBeInTheDocument(); + }); + + it("should display file sizes", () => { + const onRemove = vi.fn(); + render(); + + expect(screen.getByText(/1000.00 KB/)).toBeInTheDocument(); + expect(screen.getByText(/2.00 MB/)).toBeInTheDocument(); + expect(screen.getByText(/500.00 KB/)).toBeInTheDocument(); + }); + + it("should call onRemove when delete button is clicked", () => { + const onRemove = vi.fn(); + render(); + + const deleteButtons = screen.getAllByLabelText(/delete/i); + + act(() => { + fireEvent.click(deleteButtons[0]); + }); + + expect(onRemove).toHaveBeenCalledWith("1"); + }); + + it("should show empty state when no documents", () => { + const onRemove = vi.fn(); + render(); + + expect(screen.getByText(/No documents uploaded yet/)).toBeInTheDocument(); + }); + + it("should have action buttons for each document", () => { + const onRemove = vi.fn(); + render(); + + // Each document should have 3 action buttons (view, copy, delete) + const viewButtons = screen.getAllByLabelText(/eye/i); + const copyButtons = screen.getAllByLabelText(/copy/i); + const deleteButtons = screen.getAllByLabelText(/delete/i); + + expect(viewButtons).toHaveLength(3); + expect(copyButtons).toHaveLength(3); + expect(deleteButtons).toHaveLength(3); + }); +}); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx new file mode 100644 index 0000000000..aeb4240d36 --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx @@ -0,0 +1,98 @@ +import React from "react"; +import { Table, Badge, Tooltip, message } from "antd"; +import { EyeOutlined, CopyOutlined, DeleteOutlined } from "@ant-design/icons"; +import { DocumentUpload } from "./types"; + +interface DocumentsTableProps { + documents: DocumentUpload[]; + onRemove: (uid: string) => void; +} + +const DocumentsTable: React.FC = ({ documents, onRemove }) => { + const handleCopyId = (uid: string) => { + navigator.clipboard.writeText(uid); + message.success("Document ID copied to clipboard"); + }; + + const getStatusBadge = (status: DocumentUpload["status"]) => { + const statusConfig = { + uploading: { color: "blue", text: "Uploading" }, + done: { color: "green", text: "Ready" }, + error: { color: "red", text: "Error" }, + removed: { color: "default", text: "Removed" }, + }; + + const config = statusConfig[status]; + return ; + }; + + const formatFileSize = (bytes?: number) => { + if (!bytes) return "-"; + const kb = bytes / 1024; + if (kb < 1024) return `${kb.toFixed(2)} KB`; + return `${(kb / 1024).toFixed(2)} MB`; + }; + + const columns = [ + { + title: "Name", + dataIndex: "name", + key: "name", + render: (name: string, record: DocumentUpload) => ( +
+ {name} + {record.size && ({formatFileSize(record.size)})} +
+ ), + }, + { + title: "Status", + dataIndex: "status", + key: "status", + width: 150, + render: (status: DocumentUpload["status"]) => getStatusBadge(status), + }, + { + title: "Actions", + key: "actions", + width: 120, + render: (_: any, record: DocumentUpload) => ( +
+ + console.log("View", record)} + /> + + + handleCopyId(record.uid)} + /> + + + onRemove(record.uid)} + /> + +
+ ), + }, + ]; + + return ( + + ); +}; + +export default DocumentsTable; diff --git a/ui/litellm-dashboard/src/components/vector_store_management/TestVectorStoreTab.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/TestVectorStoreTab.test.tsx new file mode 100644 index 0000000000..ad6ef02e87 --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/TestVectorStoreTab.test.tsx @@ -0,0 +1,90 @@ +import { render, screen, fireEvent } from "@testing-library/react"; +import { describe, it, expect, vi } from "vitest"; +import TestVectorStoreTab from "./TestVectorStoreTab"; +import { VectorStore } from "./types"; + +// Mock VectorStoreTester component +vi.mock("./VectorStoreTester", () => ({ + VectorStoreTester: ({ vectorStoreId, accessToken }: { vectorStoreId: string; accessToken: string }) => ( +
+
{vectorStoreId}
+
{accessToken}
+
+ ), +})); + +const mockVectorStores: VectorStore[] = [ + { + vector_store_id: "vs_123", + custom_llm_provider: "openai", + vector_store_name: "Test Store 1", + vector_store_description: "Description 1", + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + { + vector_store_id: "vs_456", + custom_llm_provider: "bedrock", + vector_store_name: "Test Store 2", + vector_store_description: "Description 2", + created_at: "2024-01-02T00:00:00Z", + updated_at: "2024-01-02T00:00:00Z", + }, +]; + +describe("TestVectorStoreTab", () => { + it("should render the component successfully", () => { + render(); + + expect(screen.getByText("Select Vector Store")).toBeInTheDocument(); + expect(screen.getByText("Choose a vector store to test search queries against")).toBeInTheDocument(); + }); + + it("should show message when no access token", () => { + render(); + + expect(screen.getByText("Access token is required to test vector stores.")).toBeInTheDocument(); + }); + + it("should show message when no vector stores available", () => { + render(); + + expect(screen.getByText("No vector stores available. Create one first to test it.")).toBeInTheDocument(); + }); + + it("should render VectorStoreTester with first vector store by default", () => { + render(); + + expect(screen.getByTestId("vector-store-tester")).toBeInTheDocument(); + expect(screen.getByTestId("tester-vector-store-id")).toHaveTextContent("vs_123"); + expect(screen.getByTestId("tester-access-token")).toHaveTextContent("test-token"); + }); + + it("should update VectorStoreTester when selecting different vector store", () => { + render(); + + // Find the select component + const selectElement = screen.getByRole("combobox"); + + // Change selection + fireEvent.mouseDown(selectElement); + + // Wait for options to appear and click the second one + const option2 = screen.getByText("Test Store 2"); + fireEvent.click(option2); + + // Verify the tester component updated + expect(screen.getByTestId("tester-vector-store-id")).toHaveTextContent("vs_456"); + }); + + it("should display vector store names in select options", () => { + render(); + + const selectElement = screen.getByRole("combobox"); + fireEvent.mouseDown(selectElement); + + // Use getAllByText since the selected value also shows the name + expect(screen.getAllByText("Test Store 1").length).toBeGreaterThan(0); + expect(screen.getByText("Test Store 2")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/TestVectorStoreTab.tsx b/ui/litellm-dashboard/src/components/vector_store_management/TestVectorStoreTab.tsx new file mode 100644 index 0000000000..da005491ca --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/TestVectorStoreTab.tsx @@ -0,0 +1,75 @@ +import React, { useState } from "react"; +import { Card, Select, Typography } from "antd"; +import { VectorStoreTester } from "./VectorStoreTester"; +import { VectorStore } from "./types"; + +const { Text, Title } = Typography; + +interface TestVectorStoreTabProps { + accessToken: string | null; + vectorStores: VectorStore[]; +} + +const TestVectorStoreTab: React.FC = ({ accessToken, vectorStores }) => { + const [selectedVectorStoreId, setSelectedVectorStoreId] = useState( + vectorStores.length > 0 ? vectorStores[0].vector_store_id : undefined + ); + + if (!accessToken) { + return ( + + Access token is required to test vector stores. + + ); + } + + if (vectorStores.length === 0) { + return ( + +
+ No vector stores available. Create one first to test it. +
+ + ); + } + + return ( +
+ +
+
+ Select Vector Store + Choose a vector store to test search queries against +
+ + +
+ + + {selectedVectorStoreId && ( + + )} +
+ ); +}; + +export default TestVectorStoreTab; diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx index 65d15260c4..0e2be7f62d 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx @@ -130,9 +130,9 @@ describe("VectorStoreTable", () => { expect(screen.getByText("Provider")).toBeInTheDocument(); expect(screen.getByText("Created At")).toBeInTheDocument(); expect(screen.getByText("Updated At")).toBeInTheDocument(); - // Check that we have the expected number of header cells (6 data + 1 actions) + // Check that we have the expected number of header cells (7 data + 1 actions) const headers = screen.getAllByRole("columnheader"); - expect(headers).toHaveLength(7); + expect(headers).toHaveLength(8); }); it("should render all vector store rows", () => { @@ -183,7 +183,7 @@ describe("VectorStoreTable", () => { it("should render fallback for missing name", () => { renderComponent(); const fallbackElements = screen.getAllByText("-"); - expect(fallbackElements.length).toBe(2); // One for missing name, one for missing description + expect(fallbackElements.length).toBe(3); // One for missing name, one for missing description, one for missing files }); it("should wrap name in tooltip", () => { @@ -203,7 +203,7 @@ describe("VectorStoreTable", () => { it("should render fallback for missing description", () => { renderComponent(); const fallbackElements = screen.getAllByText("-"); - expect(fallbackElements.length).toBe(2); // One for missing name, one for missing description + expect(fallbackElements.length).toBe(3); // One for missing name, one for missing description, one for missing files }); it("should wrap description in tooltip", () => { @@ -386,7 +386,7 @@ describe("VectorStoreTable", () => { it("should span all columns in empty state", () => { renderComponent({ data: [] }); const emptyCell = screen.getByText("No vector stores found").closest("td"); - expect(emptyCell).toHaveAttribute("colSpan", "7"); // 6 data columns + 1 actions column + expect(emptyCell).toHaveAttribute("colSpan", "8"); // 7 data columns + 1 actions column }); }); @@ -403,7 +403,7 @@ describe("VectorStoreTable", () => { renderComponent({ data: minimalData }); expect(screen.getByText("minimal")).toBeInTheDocument(); - expect(screen.getAllByText("-")).toHaveLength(2); // Name and description fallbacks + expect(screen.getAllByText("-")).toHaveLength(3); // Name, description, and files fallbacks }); it("should handle single vector store", () => { diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.tsx index 52462c02e9..41e2b63112 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.tsx @@ -66,6 +66,32 @@ const VectorStoreTable: React.FC = ({ data, onView, onEdi ); }, }, + { + header: "Files", + accessorKey: "vector_store_metadata", + cell: ({ row }) => { + const vectorStore = row.original; + const ingestedFiles = vectorStore.vector_store_metadata?.ingested_files || []; + + if (ingestedFiles.length === 0) { + return -; + } + + const filenames = ingestedFiles + .map((file) => file.filename || file.file_url || "Unknown") + .join(", "); + + const displayText = ingestedFiles.length === 1 + ? ingestedFiles[0].filename || ingestedFiles[0].file_url || "1 file" + : `${ingestedFiles.length} files`; + + return ( + + {displayText} + + ); + }, + }, { header: "Provider", accessorKey: "custom_llm_provider", diff --git a/ui/litellm-dashboard/src/components/vector_store_management/index.tsx b/ui/litellm-dashboard/src/components/vector_store_management/index.tsx index 6d21e861d4..9cb57b8b9f 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/index.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/index.tsx @@ -1,5 +1,5 @@ import React, { useState, useEffect } from "react"; -import { Icon, Button as TremorButton, Col, Text, Grid } from "@tremor/react"; +import { Icon, Button as TremorButton, Col, Text, Grid, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react"; import { RefreshIcon } from "@heroicons/react/outline"; import { vectorStoreListCall, vectorStoreDeleteCall, credentialListCall, CredentialItem } from "../networking"; import { VectorStore } from "./types"; @@ -7,6 +7,8 @@ import VectorStoreTable from "./VectorStoreTable"; import VectorStoreForm from "./VectorStoreForm"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; import VectorStoreInfoView from "./vector_store_info"; +import CreateVectorStore from "./CreateVectorStore"; +import TestVectorStoreTab from "./TestVectorStoreTab"; import { isAdminRole } from "@/utils/roles"; import NotificationsManager from "../molecules/notifications_manager"; @@ -101,6 +103,12 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID fetchVectorStores(); }; + const handleVectorStoreCreated = (vectorStoreId: string) => { + console.log("Vector store created:", vectorStoreId); + fetchVectorStores(); + // Optionally switch to the manage tab + }; + useEffect(() => { fetchVectorStores(); fetchCredentials(); @@ -134,18 +142,46 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID -

You can use vector stores to store and retrieve LLM embeddings..

+

You can use vector stores to store and retrieve LLM embeddings.

 - setIsCreateModalVisible(true)}> - + Add Vector Store - + + + Create Vector Store + Manage Vector Stores + Test Vector Store + - - - - - + + {/* Tab 1: Create Vector Store */} + + + + + {/* Tab 2: Manage Vector Stores */} + + setIsCreateModalVisible(true)}> + + Add Vector Store + + + + + + + + + + {/* Tab 3: Test Vector Store */} + + + + + {/* Create Vector Store Modal */} ; + vector_store_metadata?: VectorStoreMetadata; created_at: string; updated_at: string; created_by?: string; @@ -41,3 +55,32 @@ export interface VectorStoreListResponse { current_page: number; total_pages: number; } + +// Document ingestion types +export interface DocumentUpload { + uid: string; + name: string; + status: "uploading" | "done" | "error" | "removed"; + size?: number; + type?: string; + originFileObj?: File; +} + +export interface RAGIngestRequest { + file_url?: string; + file_id?: string; + ingest_options: { + vector_store: { + custom_llm_provider: string; + vector_store_id?: string; + }; + }; +} + +export interface RAGIngestResponse { + id: string; + status: "completed" | "processing" | "failed"; + vector_store_id: string; + file_id: string; + error?: string; +} diff --git a/ui/litellm-dashboard/src/components/view_users.tsx b/ui/litellm-dashboard/src/components/view_users.tsx index f8cf8302a5..8a09c6d1f9 100644 --- a/ui/litellm-dashboard/src/components/view_users.tsx +++ b/ui/litellm-dashboard/src/components/view_users.tsx @@ -2,7 +2,7 @@ import { Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react"; import React, { useEffect, useState } from "react"; import { Button } from "@tremor/react"; -import BulkEditUserModal from "./bulk_edit_user"; +import BulkEditUserModal from "./BulkEditUsers"; import CreateUser from "./create_user_button"; import EditUserModal from "./edit_user"; import { @@ -286,7 +286,7 @@ const ViewUserDashboard: React.FC = ({ accessToken, toke }, handleDelete, handleResetPassword, - () => {}, // placeholder function, will be overridden in UserDataTable + () => { }, // placeholder function, will be overridden in UserDataTable ); return ( @@ -415,7 +415,7 @@ const ViewUserDashboard: React.FC = ({ accessToken, toke /> setIsBulkEditModalVisible(false)} selectedUsers={selectedUsers} possibleUIRoles={possibleUIRoles}