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/deploy/charts/litellm-helm/templates/deployment.yaml b/deploy/charts/litellm-helm/templates/deployment.yaml index c3e0055e38..4ac5582d06 100644 --- a/deploy/charts/litellm-helm/templates/deployment.yaml +++ b/deploy/charts/litellm-helm/templates/deployment.yaml @@ -38,6 +38,10 @@ spec: serviceAccountName: {{ include "litellm.serviceAccountName" . }} securityContext: {{- toYaml .Values.podSecurityContext | nindent 8 }} + {{- with .Values.extraInitContainers }} + initContainers: + {{- toYaml . | nindent 8 }} + {{- end }} containers: - name: {{ include "litellm.name" . }} securityContext: diff --git a/deploy/charts/litellm-helm/templates/migrations-job.yaml b/deploy/charts/litellm-helm/templates/migrations-job.yaml index f8893a47af..3459fa12d1 100644 --- a/deploy/charts/litellm-helm/templates/migrations-job.yaml +++ b/deploy/charts/litellm-helm/templates/migrations-job.yaml @@ -35,6 +35,10 @@ spec: {{- toYaml . | nindent 8 }} {{- end }} serviceAccountName: {{ include "litellm.serviceAccountName" . }} + {{- with .Values.migrationJob.extraInitContainers }} + initContainers: + {{- toYaml . | nindent 8 }} + {{- end }} containers: - name: prisma-migrations image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default (printf "main-%s" .Chart.AppVersion) }}" diff --git a/deploy/charts/litellm-helm/values.yaml b/deploy/charts/litellm-helm/values.yaml index b75ce64037..cea25974bb 100644 --- a/deploy/charts/litellm-helm/values.yaml +++ b/deploy/charts/litellm-helm/values.yaml @@ -281,6 +281,7 @@ migrationJob: # cpu: 100m # memory: 100Mi extraContainers: [] + extraInitContainers: [] # Hook configuration hooks: 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/gemini.md b/docs/my-website/docs/providers/gemini.md index 23a02f7365..b9ad7820dd 100644 --- a/docs/my-website/docs/providers/gemini.md +++ b/docs/my-website/docs/providers/gemini.md @@ -1840,6 +1840,57 @@ content = response.get('choices', [{}])[0].get('message', {}).get('content') print(content) ``` +## gemini-robotics-er-1.5-preview Usage + +```python +from litellm import api_base +from openai import OpenAI +import os +import base64 + +client = OpenAI(base_url="http://0.0.0.0:4000", api_key="sk-12345") +base64_image = base64.b64encode(open("closeup-object-on-table-many-260nw-1216144471.webp", "rb").read()).decode() + +import json +import re +tools = [{"codeExecution": {}}] +response = client.chat.completions.create( + model="gemini/gemini-robotics-er-1.5-preview", + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Point to no more than 10 items in the image. The label returned should be an identifying name for the object detected. The answer should follow the json format: [{\"point\": [y, x], \"label\": }, ...]. The points are in [y, x] format normalized to 0-1000." + }, + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"} + } + ] + } + ], + tools=tools +) + +# Extract JSON from markdown code block if present +content = response.choices[0].message.content +# Look for triple-backtick JSON block +match = re.search(r'```json\s*(.*?)\s*```', content, re.DOTALL) +if match: + json_str = match.group(1) +else: + json_str = content + +try: + data = json.loads(json_str) + print(json.dumps(data, indent=2)) +except Exception as e: + print("Error parsing response as JSON:", e) + print("Response content:", content) +``` + ## Usage - PDF / Videos / etc. Files ### Inline Data (e.g. audio stream) diff --git a/docs/my-website/docs/providers/sarvam.md b/docs/my-website/docs/providers/sarvam.md new file mode 100644 index 0000000000..d77e9c0c75 --- /dev/null +++ b/docs/my-website/docs/providers/sarvam.md @@ -0,0 +1,89 @@ +# Sarvam.ai + +LiteLLM supports all the text models from [Sarvam ai](https://docs.sarvam.ai/api-reference-docs/chat/chat-completions) + +## Usage + +```python +import os +from litellm import completion + +# Set your Sarvam API key +os.environ["SARVAM_API_KEY"] = "" + +messages = [{"role": "user", "content": "Hello"}] + +response = completion( + model="sarvam/sarvam-m", + messages=messages, +) +print(response) +``` + +## Usage with LiteLLM Proxy Server + +Here's how to call a Sarvam.ai model with the LiteLLM Proxy Server + +1. **Modify the `config.yaml`:** + + ```yaml + model_list: + - model_name: my-model + litellm_params: + model: sarvam/ # add sarvam/ prefix to route as Sarvam provider + api_key: api-key # api key to send your model + ``` + +2. **Start the proxy:** + + ```bash + $ litellm --config /path/to/config.yaml + ``` + +3. **Send a request to LiteLLM Proxy Server:** + + + + + + ```python + import openai + + client = openai.OpenAI( + api_key="sk-1234", # pass litellm proxy key, if you're using virtual keys + base_url="http://0.0.0.0:4000" # litellm-proxy-base url + ) + + response = client.chat.completions.create( + model="my-model", + messages=[ + { + "role": "user", + "content": "what llm are you" + } + ], + ) + + print(response) + ``` + + + + + ```shell + curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "model": "my-model", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ] + }' + ``` + + + 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/docs/proxy/guardrails/quick_start.md b/docs/my-website/docs/proxy/guardrails/quick_start.md index cb6379d49f..ddb215fcb6 100644 --- a/docs/my-website/docs/proxy/guardrails/quick_start.md +++ b/docs/my-website/docs/proxy/guardrails/quick_start.md @@ -405,14 +405,10 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ ## **Proxy Admin Controls** -### ✨ Monitoring Guardrails +### Monitoring Guardrails Monitor which guardrails were executed and whether they passed or failed. e.g. guardrail going rogue and failing requests we don't intend to fail -:::info - -✨ This is an Enterprise only feature [Get a free trial](https://www.litellm.ai/enterprise#trial) - ::: #### Setup diff --git a/docs/my-website/docs/rag_ingest.md b/docs/my-website/docs/rag_ingest.md index 1133b85f20..7adc2d70b5 100644 --- a/docs/my-website/docs/rag_ingest.md +++ b/docs/my-website/docs/rag_ingest.md @@ -5,7 +5,7 @@ All-in-one document ingestion pipeline: **Upload → Chunk → Embed → Vector | Feature | Supported | |---------|-----------| | Logging | Yes | -| Supported Providers | `openai`, `bedrock`, `vertex_ai`, `gemini` | +| Supported Providers | `openai`, `bedrock`, `vertex_ai`, `gemini`, `s3_vectors` | :::tip After ingesting documents, use [/rag/query](./rag_query.md) to search and generate responses with your ingested content. @@ -75,6 +75,31 @@ curl -X POST "http://localhost:4000/v1/rag/ingest" \ }" ``` +### AWS S3 Vectors + +```bash showLineNumbers title="Ingest to S3 Vectors" +curl -X POST "http://localhost:4000/v1/rag/ingest" \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d "{ + \"file\": { + \"filename\": \"document.txt\", + \"content\": \"$(base64 -i document.txt)\", + \"content_type\": \"text/plain\" + }, + \"ingest_options\": { + \"embedding\": { + \"model\": \"text-embedding-3-small\" + }, + \"vector_store\": { + \"custom_llm_provider\": \"s3_vectors\", + \"vector_bucket_name\": \"my-embeddings\", + \"aws_region_name\": \"us-west-2\" + } + } + }" +``` + ## Response ```json @@ -265,6 +290,57 @@ When `vector_store_id` is omitted, LiteLLM automatically creates: 4. Install: `pip install 'google-cloud-aiplatform>=1.60.0'` ::: +### vector_store (AWS S3 Vectors) + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `custom_llm_provider` | string | - | `"s3_vectors"` | +| `vector_bucket_name` | string | **required** | S3 vector bucket name | +| `index_name` | string | auto-create | Vector index name | +| `dimension` | integer | auto-detect | Vector dimension (auto-detected from embedding model) | +| `distance_metric` | string | `cosine` | Distance metric: `cosine` or `euclidean` | +| `non_filterable_metadata_keys` | array | `["source_text"]` | Metadata keys excluded from filtering | +| `aws_region_name` | string | `us-west-2` | AWS region | +| `aws_access_key_id` | string | env | AWS access key | +| `aws_secret_access_key` | string | env | AWS secret key | + +:::info S3 Vectors Auto-Creation +When `index_name` is omitted, LiteLLM automatically creates: +- S3 vector bucket (if it doesn't exist) +- Vector index with auto-detected dimensions from your embedding model + +**Dimension Auto-Detection**: The vector dimension is automatically detected by making a test embedding request to your specified model. No need to manually specify dimensions! + +**Supported Embedding Models**: Works with any LiteLLM-supported embedding model (OpenAI, Cohere, Bedrock, Azure, etc.) +::: + +**Example with auto-detection:** +```json +{ + "embedding": { + "model": "text-embedding-3-small" // Dimension auto-detected as 1536 + }, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": "my-embeddings" + } +} +``` + +**Example with custom embedding provider:** +```json +{ + "embedding": { + "model": "cohere/embed-english-v3.0" // Dimension auto-detected as 1024 + }, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": "my-embeddings", + "distance_metric": "cosine" + } +} +``` + ## Input Examples ### File (Base64) diff --git a/docs/my-website/docs/routing.md b/docs/my-website/docs/routing.md index 47967775e1..b5ece7237a 100644 --- a/docs/my-website/docs/routing.md +++ b/docs/my-website/docs/routing.md @@ -828,7 +828,12 @@ asyncio.run(router_acompletion()) ``` - + +## Traffic Mirroring / Silent Experiments + +Traffic mirroring allows you to "mimic" production traffic to a secondary (silent) model for evaluation purposes. The silent model's response is gathered in the background and does not affect the latency or result of the primary request. + +[**See detailed guide on A/B Testing - Traffic Mirroring here**](./traffic_mirroring.md) ## Basic Reliability diff --git a/docs/my-website/docs/traffic_mirroring.md b/docs/my-website/docs/traffic_mirroring.md new file mode 100644 index 0000000000..3bdcb0f161 --- /dev/null +++ b/docs/my-website/docs/traffic_mirroring.md @@ -0,0 +1,83 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# A/B Testing - Traffic Mirroring + +Traffic mirroring allows you to "mimic" production traffic to a secondary (silent) model for evaluation purposes. The silent model's response is gathered in the background and does not affect the latency or result of the primary request. + +This is useful for: +- Testing a new model's performance on production prompts before switching. +- Comparing costs and latency between different providers. +- Debugging issues by mirroring traffic to a more verbose model. + +## Quick Start + +To enable traffic mirroring, add `silent_model` to the `litellm_params` of a deployment. + + + + +```python +from litellm import Router + +model_list = [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "azure/chatgpt-v-2", + "api_key": "...", + "silent_model": "gpt-4" # 👈 Mirror traffic to gpt-4 + }, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "..." + }, + } +] + +router = Router(model_list=model_list) + +# The request to "gpt-3.5-turbo" will trigger a background call to "gpt-4" +response = await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "How does traffic mirroring work?"}] +) +``` + + + + +Add `silent_model` to your `config.yaml`: + +```yaml +model_list: + - model_name: primary-model + litellm_params: + model: azure/gpt-35-turbo + api_key: os.environ/AZURE_API_KEY + silent_model: evaluation-model # 👈 Mirror traffic here + - model_name: evaluation-model + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY +``` + + + + +## How it works +1. **Request Received**: A request is made to a model group (e.g. `primary-model`). +2. **Deployment Picked**: LiteLLM picks a deployment from the group. +3. **Primary Call**: LiteLLM makes the call to the primary deployment. +4. **Mirroring**: If `silent_model` is present, LiteLLM triggers a background call to that model. + - For **Sync** calls: Uses a shared thread pool. + - For **Async** calls: Uses `asyncio.create_task`. +5. **Isolation**: The background call uses a `deepcopy` of the original request parameters and sets `metadata["is_silent_experiment"] = True`. It also strips out logging IDs to prevent collisions in usage tracking. + +## Key Features +- **Latency Isolation**: The primary request returns as soon as it's ready. The background (silent) call does not block. +- **Unified Logging**: Background calls are processed via the Router, meaning they are automatically logged to your configured observability tools (Langfuse, S3, etc.). +- **Evaluation**: Use the `is_silent_experiment: True` flag in your logs to filter and compare results between the primary and mirrored calls. 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/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 85573599e4..ef7573e9ce 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -364,6 +364,7 @@ const sidebars = { label: "Load Balancing, Routing, Fallbacks", href: "https://docs.litellm.ai/docs/routing-load-balancing", }, + "traffic_mirroring", { type: "category", label: "Logging, Alerting, Metrics", @@ -775,6 +776,7 @@ const sidebars = { "providers/oci", "providers/ollama", "providers/openrouter", + "providers/sarvam", "providers/ovhcloud", "providers/perplexity", "providers/petals", diff --git a/enterprise/litellm_enterprise/proxy/auth/route_checks.py b/enterprise/litellm_enterprise/proxy/auth/route_checks.py index 6cce781faf..6f7cf9143f 100644 --- a/enterprise/litellm_enterprise/proxy/auth/route_checks.py +++ b/enterprise/litellm_enterprise/proxy/auth/route_checks.py @@ -36,7 +36,7 @@ class EnterpriseRouteChecks: if not premium_user: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"🚨🚨🚨 DISABLING LLM API ENDPOINTS is an Enterprise feature\n🚨 {CommonProxyErrors.not_premium_user.value}", + detail=f"🚨🚨🚨 DISABLING ADMIN ENDPOINTS is an Enterprise feature\n🚨 {CommonProxyErrors.not_premium_user.value}", ) return get_secret_bool("DISABLE_ADMIN_ENDPOINTS") is True 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/constants.py b/litellm/constants.py index 0525bf3843..0c9d1d941b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1327,6 +1327,13 @@ COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY = int( DEFAULT_CHUNK_SIZE = int(os.getenv("DEFAULT_CHUNK_SIZE", 1000)) DEFAULT_CHUNK_OVERLAP = int(os.getenv("DEFAULT_CHUNK_OVERLAP", 200)) +########################### S3 Vectors RAG Constants ########################### +S3_VECTORS_DEFAULT_DIMENSION = int(os.getenv("S3_VECTORS_DEFAULT_DIMENSION", 1024)) +S3_VECTORS_DEFAULT_DISTANCE_METRIC = str( + os.getenv("S3_VECTORS_DEFAULT_DISTANCE_METRIC", "cosine") +) +S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS = ["source_text"] + ########################### Microsoft SSO Constants ########################### MICROSOFT_USER_EMAIL_ATTRIBUTE = str( os.getenv("MICROSOFT_USER_EMAIL_ATTRIBUTE", "userPrincipalName") 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 bafb0d88c8..6772520bd5 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 @@ -1245,6 +1258,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, @@ -1266,11 +1280,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( @@ -1365,14 +1379,14 @@ class PrometheusLogger(CustomLogger): standard_logging_payload: StandardLoggingPayload = kwargs.get( "standard_logging_object", {} ) - + if self._should_skip_metrics_for_invalid_key( kwargs=kwargs, standard_logging_payload=standard_logging_payload ): return - + model = kwargs.get("model", "") - + litellm_params = kwargs.get("litellm_params", {}) or {} get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking() @@ -1396,6 +1410,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: @@ -1413,49 +1428,57 @@ class PrometheusLogger(CustomLogger): ) -> Optional[int]: """ Extract HTTP status code from various input formats for validation. - + This is a centralized helper to extract status code from different callback function signatures. Handles both ProxyException (uses 'code') and standard exceptions (uses 'status_code'). - + Args: kwargs: Dictionary potentially containing 'exception' key enum_values: Object with 'status_code' attribute exception: Exception object to extract status code from directly - + Returns: Status code as integer if found, None otherwise """ status_code = None - + # Try from enum_values first (most common in our callbacks) - if enum_values and hasattr(enum_values, "status_code") and enum_values.status_code: + if ( + enum_values + and hasattr(enum_values, "status_code") + and enum_values.status_code + ): try: status_code = int(enum_values.status_code) except (ValueError, TypeError): pass - + if not status_code and exception: # ProxyException uses 'code' attribute, other exceptions may use 'status_code' - status_code = getattr(exception, "status_code", None) or getattr(exception, "code", None) + status_code = getattr(exception, "status_code", None) or getattr( + exception, "code", None + ) if status_code is not None: try: status_code = int(status_code) except (ValueError, TypeError): status_code = None - + if not status_code and kwargs: exception_in_kwargs = kwargs.get("exception") if exception_in_kwargs: - status_code = getattr(exception_in_kwargs, "status_code", None) or getattr(exception_in_kwargs, "code", None) + status_code = getattr( + exception_in_kwargs, "status_code", None + ) or getattr(exception_in_kwargs, "code", None) if status_code is not None: try: status_code = int(status_code) except (ValueError, TypeError): status_code = None - + return status_code - + def _is_invalid_api_key_request( self, status_code: Optional[int], @@ -1463,23 +1486,23 @@ class PrometheusLogger(CustomLogger): ) -> bool: """ Determine if a request has an invalid API key based on status code and exception. - + This method prevents invalid authentication attempts from being recorded in Prometheus metrics. A 401 status code is the definitive indicator of authentication failure. Additionally, we check exception messages for authentication error patterns to catch cases where the exception hasn't been converted to a ProxyException yet. - + Args: status_code: HTTP status code (401 indicates authentication error) exception: Exception object to check for auth-related error messages - + Returns: True if the request has an invalid API key and metrics should be skipped, False otherwise """ if status_code == 401: return True - + # Handle cases where AssertionError is raised before conversion to ProxyException if exception is not None: exception_str = str(exception).lower() @@ -1492,9 +1515,9 @@ class PrometheusLogger(CustomLogger): ] if any(pattern in exception_str for pattern in auth_error_patterns): return True - + return False - + def _should_skip_metrics_for_invalid_key( self, kwargs: Optional[dict] = None, @@ -1505,18 +1528,18 @@ class PrometheusLogger(CustomLogger): ) -> bool: """ Determine if Prometheus metrics should be skipped for invalid API key requests. - + This is a centralized validation method that extracts status code and exception information from various callback function signatures and determines if the request represents an invalid API key attempt that should be filtered from metrics. - + Args: kwargs: Dictionary potentially containing exception and other data user_api_key_dict: User API key authentication object (currently unused) enum_values: Object with status_code attribute standard_logging_payload: Standard logging payload dictionary exception: Exception object to check directly - + Returns: True if metrics should be skipped (invalid key detected), False otherwise """ @@ -1525,17 +1548,17 @@ class PrometheusLogger(CustomLogger): enum_values=enum_values, exception=exception, ) - + if exception is None and kwargs: exception = kwargs.get("exception") - + if self._is_invalid_api_key_request(status_code, exception=exception): verbose_logger.debug( "Skipping Prometheus metrics for invalid API key request: " f"status_code={status_code}, exception={type(exception).__name__ if exception else None}" ) return True - + return False async def async_post_call_failure_hook( @@ -1576,6 +1599,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, @@ -1590,6 +1617,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( @@ -1629,6 +1659,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, @@ -1644,6 +1675,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( @@ -1684,7 +1717,7 @@ class PrometheusLogger(CustomLogger): exception = request_kwargs.get("exception", None) llm_provider = _litellm_params.get("custom_llm_provider", None) - + if self._should_skip_metrics_for_invalid_key( kwargs=request_kwargs, standard_logging_payload=standard_logging_payload, @@ -1716,6 +1749,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"), ) """ @@ -1753,6 +1790,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, @@ -1786,6 +1866,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"][ @@ -2263,7 +2353,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, @@ -2379,12 +2472,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/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 5d0bedb776..7137a4e422 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -31,15 +31,19 @@ def _process_image_response(response: Response, url: str) -> str: f"Error: Image size ({size_mb:.2f}MB) exceeds maximum allowed size ({MAX_IMAGE_URL_DOWNLOAD_SIZE_MB}MB). url={url}" ) - image_bytes = response.content + # Stream download with size checking to prevent downloading huge files + max_bytes = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024) + image_bytes = bytearray() + bytes_downloaded = 0 - # Check actual size after download if Content-Length was not available - if content_length is None: - size_mb = len(image_bytes) / (1024 * 1024) - if size_mb > MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: + for chunk in response.iter_bytes(chunk_size=8192): + bytes_downloaded += len(chunk) + if bytes_downloaded > max_bytes: + size_mb = bytes_downloaded / (1024 * 1024) raise litellm.ImageFetchError( f"Error: Image size ({size_mb:.2f}MB) exceeds maximum allowed size ({MAX_IMAGE_URL_DOWNLOAD_SIZE_MB}MB). url={url}" ) + image_bytes.extend(chunk) base64_image = base64.b64encode(image_bytes).decode("utf-8") 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/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 89f2094d5d..935fd53c19 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -5,8 +5,8 @@ import httpx from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( - BaseVectorStoreAuthCredentials, VECTOR_STORE_OPENAI_PARAMS, + BaseVectorStoreAuthCredentials, VectorStoreCreateOptionalRequestParams, VectorStoreCreateResponse, VectorStoreIndexEndpoints, @@ -64,6 +64,30 @@ class BaseVectorStoreConfig: pass + async def atransform_search_vector_store_request( + self, + vector_store_id: str, + query: Union[str, List[str]], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> Tuple[str, Dict]: + """ + Optional async version of transform_search_vector_store_request. + If not implemented, the handler will fall back to the sync version. + Providers that need to make async calls (e.g., generating embeddings) should override this. + """ + # Default implementation: call the sync version + return self.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=litellm_logging_obj, + litellm_params=litellm_params, + ) + @abstractmethod def transform_search_vector_store_response( self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 642d15fe3e..1de1c40c43 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1163,7 +1163,7 @@ class BaseAWSLLM: def _sign_request( self, - service_name: Literal["bedrock", "sagemaker", "bedrock-agentcore"], + service_name: Literal["bedrock", "sagemaker", "bedrock-agentcore", "s3vectors"], headers: dict, optional_params: dict, request_data: dict, 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/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 89b42f5e94..65d237bdbd 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -797,7 +797,7 @@ class BedrockEventStreamDecoderBase: def get_anthropic_beta_from_headers(headers: dict) -> List[str]: """ Extract anthropic-beta header values and convert them to a list. - Supports comma-separated values from user headers. + Supports both JSON array format and comma-separated values from user headers. Used by both converse and invoke transformations for consistent handling of anthropic-beta headers that should be passed to AWS Bedrock. @@ -812,8 +812,25 @@ def get_anthropic_beta_from_headers(headers: dict) -> List[str]: if not anthropic_beta_header: return [] - # Split comma-separated values and strip whitespace - return [beta.strip() for beta in anthropic_beta_header.split(",")] + # If it's already a list, return it + if isinstance(anthropic_beta_header, list): + return anthropic_beta_header + + # Try to parse as JSON array first (e.g., '["interleaved-thinking-2025-05-14", "claude-code-20250219"]') + if isinstance(anthropic_beta_header, str): + anthropic_beta_header = anthropic_beta_header.strip() + if anthropic_beta_header.startswith("[") and anthropic_beta_header.endswith("]"): + try: + parsed = json.loads(anthropic_beta_header) + if isinstance(parsed, list): + return [str(beta).strip() for beta in parsed] + except json.JSONDecodeError: + pass # Fall through to comma-separated parsing + + # Fall back to comma-separated values + return [beta.strip() for beta in anthropic_beta_header.split(",")] + + return [] class CommonBatchFilesUtils: diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index a7065caece..fc4220c054 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -162,6 +162,49 @@ class AmazonAnthropicClaudeMessagesConfig( return any(pattern in model_lower for pattern in supported_patterns) + def _is_claude_opus_4_5(self, model: str) -> bool: + """ + Check if the model is Claude Opus 4.5. + + Args: + model: The model name + + Returns: + True if the model is Claude Opus 4.5 + """ + model_lower = model.lower() + opus_4_5_patterns = [ + "opus-4.5", "opus_4.5", "opus-4-5", "opus_4_5", + ] + return any(pattern in model_lower for pattern in opus_4_5_patterns) + + def _supports_tool_search_on_bedrock(self, model: str) -> bool: + """ + Check if the model supports tool search on Bedrock. + + On Amazon Bedrock, server-side tool search is supported on Claude Opus 4.5 + and Claude Sonnet 4.5 with the tool-search-tool-2025-10-19 beta header. + + Ref: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool + + Args: + model: The model name + + Returns: + True if the model supports tool search on Bedrock + """ + model_lower = model.lower() + + # Supported models for tool search on Bedrock + supported_patterns = [ + # Opus 4.5 + "opus-4.5", "opus_4.5", "opus-4-5", "opus_4_5", + # Sonnet 4.5 + "sonnet-4.5", "sonnet_4.5", "sonnet-4-5", "sonnet_4_5", + ] + + return any(pattern in model_lower for pattern in supported_patterns) + def _filter_unsupported_beta_headers_for_bedrock( self, model: str, beta_set: set ) -> None: @@ -169,25 +212,33 @@ class AmazonAnthropicClaudeMessagesConfig( Remove beta headers that are not supported on Bedrock for the given model. Extended thinking beta headers are only supported on specific Claude 4+ models. - Advanced tool use headers are not supported on Bedrock Invoke API. + Advanced tool use headers are not supported on Bedrock Invoke API, but need to be + translated to Bedrock-specific headers for models that support tool search + (Claude Opus 4.5, Sonnet 4.5). This prevents 400 "invalid beta flag" errors on Bedrock. Note: Bedrock Invoke API fails with a 400 error when unsupported beta headers are sent, returning: {"message":"invalid beta flag"} + Translation for models supporting tool search (Opus 4.5, Sonnet 4.5): + - advanced-tool-use-2025-11-20 -> tool-search-tool-2025-10-19 + tool-examples-2025-10-29 + Args: model: The model name beta_set: The set of beta headers to filter in-place """ beta_headers_to_remove = set() + has_advanced_tool_use = False - # 1. Filter out beta headers that are universally unsupported on Bedrock Invoke + # 1. Filter out beta headers that are universally unsupported on Bedrock Invoke and track if advanced-tool-use header is present for beta in beta_set: for unsupported_pattern in self.UNSUPPORTED_BEDROCK_INVOKE_BETA_PATTERNS: if unsupported_pattern in beta.lower(): beta_headers_to_remove.add(beta) + has_advanced_tool_use = True break - + + # 2. Filter out extended thinking headers for models that don't support them extended_thinking_patterns = [ "extended-thinking", @@ -204,6 +255,14 @@ class AmazonAnthropicClaudeMessagesConfig( for beta in beta_headers_to_remove: beta_set.discard(beta) + # 3. Translate advanced-tool-use to Bedrock-specific headers for models that support tool search + # Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html + # Ref: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool + if has_advanced_tool_use and self._supports_tool_search_on_bedrock(model): + beta_set.add("tool-search-tool-2025-10-19") + beta_set.add("tool-examples-2025-10-29") + + def _get_tool_search_beta_header_for_bedrock( self, model: str, @@ -256,7 +315,7 @@ class AmazonAnthropicClaudeMessagesConfig( Ref: https://aws.amazon.com/blogs/machine-learning/structured-data-response-with-amazon-bedrock-prompt-engineering-and-tool-use/ """ import json - + # Extract schema from output_format schema = output_format.get("schema") if not schema: 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/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 6a87967c3a..d2ea7e872a 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -7033,17 +7033,31 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - ( - url, - request_body, - ) = vector_store_provider_config.transform_search_vector_store_request( - vector_store_id=vector_store_id, - query=query, - vector_store_search_optional_params=vector_store_search_optional_params, - api_base=api_base, - litellm_logging_obj=logging_obj, - litellm_params=dict(litellm_params), - ) + # Check if provider has async transform method + if hasattr(vector_store_provider_config, "atransform_search_vector_store_request"): + ( + url, + request_body, + ) = await vector_store_provider_config.atransform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), + ) + else: + ( + url, + request_body, + ) = vector_store_provider_config.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), + ) all_optional_params: Dict[str, Any] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) headers, signed_json_body = vector_store_provider_config.sign_request( diff --git a/litellm/llms/groq/chat/transformation.py b/litellm/llms/groq/chat/transformation.py index a75ecd8cc7..34ea7b03dd 100644 --- a/litellm/llms/groq/chat/transformation.py +++ b/litellm/llms/groq/chat/transformation.py @@ -323,4 +323,12 @@ class GroqChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): status_code=error.get("code"), message=error.get("message"), body=error ) + # Map Groq's 'reasoning' field to LiteLLM's 'reasoning_content' field + # Groq returns delta.reasoning, but LiteLLM expects delta.reasoning_content + choices = chunk.get("choices", []) + for choice in choices: + delta = choice.get("delta", {}) + if "reasoning" in delta: + delta["reasoning_content"] = delta.pop("reasoning") + return super().chunk_parser(chunk) diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index 1d21490ea3..e955800b94 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -23,7 +23,7 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig class HostedVLLMChatConfig(OpenAIGPTConfig): def get_supported_openai_params(self, model: str) -> List[str]: params = super().get_supported_openai_params(model) - params.append("reasoning_effort") + params.extend(["reasoning_effort", "thinking"]) return params def map_openai_params( @@ -41,6 +41,27 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): _tools = _remove_strict_from_schema(_tools) if _tools is not None: non_default_params["tools"] = _tools + + # Handle thinking parameter - convert Anthropic-style to OpenAI-style reasoning_effort + # vLLM is OpenAI-compatible, so it understands reasoning_effort, not thinking + # Reference: https://github.com/BerriAI/litellm/issues/19761 + thinking = non_default_params.pop("thinking", None) + if thinking is not None and isinstance(thinking, dict): + if thinking.get("type") == "enabled": + # Only convert if reasoning_effort not already set + if "reasoning_effort" not in non_default_params: + budget_tokens = thinking.get("budget_tokens", 0) + # Map budget_tokens to reasoning_effort level + # Same logic as Anthropic adapter (translate_anthropic_thinking_to_reasoning_effort) + if budget_tokens >= 10000: + non_default_params["reasoning_effort"] = "high" + elif budget_tokens >= 5000: + non_default_params["reasoning_effort"] = "medium" + elif budget_tokens >= 2000: + non_default_params["reasoning_effort"] = "low" + else: + non_default_params["reasoning_effort"] = "minimal" + return super().map_openai_params( non_default_params, optional_params, model, drop_params ) 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/s3_vectors/__init__.py b/litellm/llms/s3_vectors/__init__.py new file mode 100644 index 0000000000..e8367949c3 --- /dev/null +++ b/litellm/llms/s3_vectors/__init__.py @@ -0,0 +1 @@ +# S3 Vectors LLM integration diff --git a/litellm/llms/s3_vectors/vector_stores/__init__.py b/litellm/llms/s3_vectors/vector_stores/__init__.py new file mode 100644 index 0000000000..ac24b4a38d --- /dev/null +++ b/litellm/llms/s3_vectors/vector_stores/__init__.py @@ -0,0 +1 @@ +# S3 Vectors vector store integration diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py new file mode 100644 index 0000000000..df81a78289 --- /dev/null +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -0,0 +1,254 @@ +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + +import httpx + +from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.types.router import GenericLiteLLMParams +from litellm.types.vector_stores import ( + VECTOR_STORE_OPENAI_PARAMS, + BaseVectorStoreAuthCredentials, + VectorStoreIndexEndpoints, + VectorStoreResultContent, + VectorStoreSearchOptionalRequestParams, + VectorStoreSearchResponse, + VectorStoreSearchResult, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): + """Vector store configuration for AWS S3 Vectors.""" + + def __init__(self) -> None: + BaseVectorStoreConfig.__init__(self) + BaseAWSLLM.__init__(self) + + def get_auth_credentials( + self, litellm_params: dict + ) -> BaseVectorStoreAuthCredentials: + return {} + + def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints: + return { + "read": [("POST", "/QueryVectors")], + "write": [], + } + + def get_supported_openai_params( + self, model: str + ) -> List[VECTOR_STORE_OPENAI_PARAMS]: + return ["max_num_results"] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + drop_params: bool, + ) -> dict: + for param, value in non_default_params.items(): + if param == "max_num_results": + optional_params["maxResults"] = value + return optional_params + + def validate_environment( + self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] + ) -> dict: + headers = headers or {} + headers.setdefault("Content-Type", "application/json") + return headers + + def get_complete_url(self, api_base: Optional[str], litellm_params: dict) -> str: + aws_region_name = litellm_params.get("aws_region_name") + if not aws_region_name: + raise ValueError("aws_region_name is required for S3 Vectors") + return f"https://s3vectors.{aws_region_name}.api.aws" + + def transform_search_vector_store_request( + self, + vector_store_id: str, + query: Union[str, List[str]], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> Tuple[str, Dict]: + """Sync version - generates embedding synchronously.""" + # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name + # If not in that format, try to construct it from litellm_params + bucket_name: str + index_name: str + + if ":" in vector_store_id: + bucket_name, index_name = vector_store_id.split(":", 1) + else: + # Try to get bucket_name from litellm_params + bucket_name_from_params = litellm_params.get("vector_bucket_name") + if not bucket_name_from_params or not isinstance(bucket_name_from_params, str): + raise ValueError( + "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " + "or vector_bucket_name must be provided in litellm_params" + ) + bucket_name = bucket_name_from_params + index_name = vector_store_id + + if isinstance(query, list): + query = " ".join(query) + + # Generate embedding for the query + embedding_model = litellm_params.get("embedding_model", "text-embedding-3-small") + + import litellm as litellm_module + embedding_response = litellm_module.embedding(model=embedding_model, input=[query]) + query_embedding = embedding_response.data[0]["embedding"] + + url = f"{api_base}/QueryVectors" + + request_body: Dict[str, Any] = { + "vectorBucketName": bucket_name, + "indexName": index_name, + "queryVector": {"float32": query_embedding}, + "topK": vector_store_search_optional_params.get("max_num_results", 5), # Default to 5 + "returnDistance": True, + "returnMetadata": True, + } + + litellm_logging_obj.model_call_details["query"] = query + return url, request_body + + async def atransform_search_vector_store_request( + self, + vector_store_id: str, + query: Union[str, List[str]], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> Tuple[str, Dict]: + """Async version - generates embedding asynchronously.""" + # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name + # If not in that format, try to construct it from litellm_params + bucket_name: str + index_name: str + + if ":" in vector_store_id: + bucket_name, index_name = vector_store_id.split(":", 1) + else: + # Try to get bucket_name from litellm_params + bucket_name_from_params = litellm_params.get("vector_bucket_name") + if not bucket_name_from_params or not isinstance(bucket_name_from_params, str): + raise ValueError( + "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " + "or vector_bucket_name must be provided in litellm_params" + ) + bucket_name = bucket_name_from_params + index_name = vector_store_id + + if isinstance(query, list): + query = " ".join(query) + + # Generate embedding for the query asynchronously + embedding_model = litellm_params.get("embedding_model", "text-embedding-3-small") + + import litellm as litellm_module + embedding_response = await litellm_module.aembedding(model=embedding_model, input=[query]) + query_embedding = embedding_response.data[0]["embedding"] + + url = f"{api_base}/QueryVectors" + + request_body: Dict[str, Any] = { + "vectorBucketName": bucket_name, + "indexName": index_name, + "queryVector": {"float32": query_embedding}, + "topK": vector_store_search_optional_params.get("max_num_results", 5), # Default to 5 + "returnDistance": True, + "returnMetadata": True, + } + + litellm_logging_obj.model_call_details["query"] = query + return url, request_body + + def sign_request( + self, + headers: dict, + optional_params: Dict, + request_data: Dict, + api_base: str, + api_key: Optional[str] = None, + ) -> Tuple[dict, Optional[bytes]]: + return self._sign_request( + service_name="s3vectors", + headers=headers, + optional_params=optional_params, + request_data=request_data, + api_base=api_base, + api_key=api_key, + ) + + def transform_search_vector_store_response( + self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj + ) -> VectorStoreSearchResponse: + try: + response_data = response.json() + results: List[VectorStoreSearchResult] = [] + + for item in response_data.get("vectors", []) or []: + metadata = item.get("metadata", {}) or {} + source_text = metadata.get("source_text", "") + + if not source_text: + continue + + # Extract file information from metadata + chunk_index = metadata.get("chunk_index", "0") + file_id = f"s3-vectors-chunk-{chunk_index}" + filename = metadata.get("filename", f"document-{chunk_index}") + + # S3 Vectors returns distance, convert to similarity score (0-1) + # Lower distance = higher similarity + # We'll normalize using 1 / (1 + distance) to get a 0-1 score + distance = item.get("distance") + score = None + if distance is not None: + # Convert distance to similarity score between 0 and 1 + # For cosine distance: similarity = 1 - distance + # For euclidean: use 1 / (1 + distance) + # Assuming cosine distance here + score = max(0.0, min(1.0, 1.0 - float(distance))) + + results.append( + VectorStoreSearchResult( + score=score, + content=[VectorStoreResultContent(text=source_text, type="text")], + file_id=file_id, + filename=filename, + attributes=metadata, + ) + ) + + return VectorStoreSearchResponse( + object="vector_store.search_results.page", + search_query=litellm_logging_obj.model_call_details.get("query", ""), + data=results, + ) + except Exception as e: + raise self.get_error_class( + error_message=str(e), + status_code=response.status_code, + headers=response.headers, + ) + + # Vector store creation is not yet implemented + def transform_create_vector_store_request( + self, + vector_store_create_optional_params, + api_base: str, + ) -> Tuple[str, Dict]: + raise NotImplementedError + + def transform_create_vector_store_response(self, response: httpx.Response): + raise NotImplementedError 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 1fc4521b29..23922e3c8b 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 = [], @@ -1098,7 +1098,6 @@ def completion( # type: ignore # noqa: PLR0915 # validate optional params stop = validate_openai_optional_params(stop=stop) - ######### unpacking kwargs ##################### args = locals() @@ -1115,7 +1114,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] @@ -1516,6 +1517,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, @@ -2341,11 +2344,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 @@ -2393,7 +2392,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 @@ -4704,7 +4705,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: @@ -4846,6 +4847,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, @@ -6739,9 +6770,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 @@ -6881,7 +6910,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, @@ -7255,8 +7284,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/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d958ea4503..4ac4159558 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -10232,6 +10232,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", @@ -13480,6 +13522,43 @@ "supports_vision": true, "supports_web_search": true }, + "gemini-robotics-er-1.5-preview": { + "cache_read_input_token_cost": 0, + "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 1e-06, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "output_cost_per_reasoning_token": 2.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-robotics-er-1-5-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions" + ], + "supported_modalities": [ + "text", + "image", + "video", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true + }, "gemini-2.5-computer-use-preview-10-2025": { "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, 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/_experimental/out/assets/logos/s3_vector.png b/litellm/proxy/_experimental/out/assets/logos/s3_vector.png new file mode 100644 index 0000000000..15a1a456e1 Binary files /dev/null and b/litellm/proxy/_experimental/out/assets/logos/s3_vector.png differ 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_checks.py b/litellm/proxy/auth/auth_checks.py index 61d9044f92..e0b056d450 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1016,8 +1016,10 @@ async def _get_fuzzy_user_object( ) if response is None and user_email is not None: + # Use case-insensitive query to handle emails with different casing + # This matches the pattern used in _check_duplicate_user_email response = await prisma_client.db.litellm_usertable.find_first( - where={"user_email": user_email}, + where={"user_email": {"equals": user_email, "mode": "insensitive"}}, include={"organization_memberships": True}, ) 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/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index bba5b87024..5c41b371ca 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -264,7 +264,7 @@ async def create_interaction( general_settings=general_settings, proxy_config=proxy_config, select_data_generator=select_data_generator, - model=data.get("model"), + model=data.get("model") or data.get("agent"), user_model=user_model, user_temperature=user_temperature, user_request_timeout=user_request_timeout, 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/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index d0afb62e8b..63db2d72fe 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1771,6 +1771,113 @@ async def _add_team_members_to_team( return updated_team, updated_users, updated_team_memberships +async def _validate_and_populate_member_user_info( + member: Member, + prisma_client: PrismaClient, +) -> Member: + """ + Validate and populate user_email/user_id for a member. + + Logic: + 1. If both user_email and user_id are provided, verify they belong to the same user (use user_email as source of truth) + 2. If only user_email is provided, populate user_id from DB + 3. If only user_id is provided, populate user_email from DB (if user exists) + 4. If only user_id is provided and doesn't exist, allow it to pass with user_email as None (will be upserted later) + 5. If user_email and user_id mismatch, throw error + + Returns a Member with user_email and user_id populated (user_email may be None if only user_id provided and user doesn't exist). + """ + if member.user_email is None and member.user_id is None: + raise HTTPException( + status_code=400, + detail={"error": "Either user_id or user_email must be provided"}, + ) + + # Case 1: Both user_email and user_id provided - verify they match + if member.user_email is not None and member.user_id is not None: + # Use user_email as source of truth + # Check for multiple users with same email first + users_by_email = await prisma_client.get_data( + key_val={"user_email": member.user_email}, + table_name="user", + query_type="find_all", + ) + + if users_by_email is None or ( + isinstance(users_by_email, list) and len(users_by_email) == 0 + ): + # User doesn't exist yet - this is fine, will be created later + return member + + if isinstance(users_by_email, list) and len(users_by_email) > 1: + raise HTTPException( + status_code=400, + detail={ + "error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead." + }, + ) + + # Get the single user + user_by_email = users_by_email[0] + + # Verify the user_id matches + if user_by_email.user_id != member.user_id: + raise HTTPException( + status_code=400, + detail={ + "error": f"user_email '{member.user_email}' and user_id '{member.user_id}' do not belong to the same user." + }, + ) + + # Both match, return as is + return member + + # Case 2: Only user_email provided - populate user_id from DB + if member.user_email is not None and member.user_id is None: + user_by_email = await prisma_client.db.litellm_usertable.find_first( + where={"user_email": {"equals": member.user_email, "mode": "insensitive"}} + ) + + if user_by_email is None: + # User doesn't exist yet - this is fine, will be created later + return member + + # Check for multiple users with same email + users_by_email = await prisma_client.get_data( + key_val={"user_email": member.user_email}, + table_name="user", + query_type="find_all", + ) + + if users_by_email and isinstance(users_by_email, list) and len(users_by_email) > 1: + raise HTTPException( + status_code=400, + detail={ + "error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead." + }, + ) + + # Populate user_id + member.user_id = user_by_email.user_id + return member + + # Case 3: Only user_id provided - populate user_email from DB if user exists + if member.user_id is not None and member.user_email is None: + user_by_id = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": member.user_id} + ) + + if user_by_id is None: + # User doesn't exist yet - allow it to pass with user_email as None + # Will be upserted later with just user_id and null email + return member + + # Populate user_email + member.user_email = user_by_id.user_email + return member + + return member + @router.post( "/team/member_add", tags=["team management"], @@ -1846,6 +1953,19 @@ async def team_member_add( complete_team_data=complete_team_data, ) + # Validate and populate user_email/user_id for members before processing + if isinstance(data.member, Member): + await _validate_and_populate_member_user_info( + member=data.member, + prisma_client=prisma_client, + ) + elif isinstance(data.member, List): + for m in data.member: + await _validate_and_populate_member_user_info( + member=m, + prisma_client=prisma_client, + ) + updated_team, updated_users, updated_team_memberships = ( await _add_team_members_to_team( data=data, diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 613390c0ec..4048b3731c 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -93,30 +93,50 @@ else: router = APIRouter() +def normalize_email(email: Optional[str]) -> Optional[str]: + """ + Normalize email address to lowercase for consistent storage and comparison. + + Email addresses should be treated as case-insensitive for SSO purposes, + even though RFC 5321 technically allows case-sensitive local parts. + This prevents issues where SSO providers return emails with different casing + than what's stored in the database. + + Args: + email: Email address to normalize, can be None + + Returns: + Lowercased email address, or None if input is None + """ + if email is None: + return None + return email.lower() if isinstance(email, str) else email + + def determine_role_from_groups( user_groups: List[str], role_mappings: "RoleMappings", ) -> 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 +144,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 +348,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 +365,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 +384,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 +396,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: @@ -390,7 +415,7 @@ def generic_response_convertor( display_name=get_nested_value( response, generic_user_display_name_attribute_name ), - email=get_nested_value(response, generic_user_email_attribute_name), + email=normalize_email(get_nested_value(response, generic_user_email_attribute_name)), first_name=get_nested_value(response, generic_user_first_name_attribute_name), last_name=get_nested_value(response, generic_user_last_name_attribute_name), provider=get_nested_value(response, generic_provider_attribute_name), @@ -399,7 +424,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 +519,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 +592,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 @@ -688,7 +751,7 @@ async def get_user_info_from_db( if _id is not None and isinstance(_id, str): potential_user_ids.append(_id) - user_email = ( + user_email = normalize_email( getattr(result, "email", None) if not isinstance(result, dict) else result.get("email", None) @@ -763,8 +826,8 @@ def _build_sso_user_update_data( Returns: dict: Update data containing user_email and optionally user_role if valid - """ - update_data: dict = {"user_email": user_email} + """ + update_data: dict = {"user_email": normalize_email(user_email)} # Get SSO role from result and include if valid sso_role = getattr(result, "user_role", None) @@ -1217,20 +1280,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 +1307,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 @@ -1266,7 +1336,7 @@ async def insert_sso_user( new_user_request = NewUserRequest( user_id=user_defined_values["user_id"], - user_email=user_defined_values["user_email"], + user_email=normalize_email(user_defined_values["user_email"]), user_role=user_defined_values["user_role"], # type: ignore max_budget=user_defined_values["max_budget"], budget_duration=user_defined_values["budget_duration"], @@ -1931,7 +2001,7 @@ class SSOAuthenticationHandler: """ Gets the user email and id from the OpenID result after validating the email domain """ - user_email: Optional[str] = getattr(result, "email", None) + user_email: Optional[str] = normalize_email(getattr(result, "email", None)) user_id: Optional[str] = ( getattr(result, "id", None) if result is not None else None ) @@ -1970,7 +2040,7 @@ class SSOAuthenticationHandler: "GENERIC_USER_ROLE_ATTRIBUTE", "role" ) user_id = getattr(result, "id", None) - user_email = getattr(result, "email", None) + user_email = normalize_email(getattr(result, "email", None)) if user_role is None: _role_from_attr = getattr(result, generic_user_role_attribute_name, None) # type: ignore if _role_from_attr is not None: @@ -2363,7 +2433,7 @@ class MicrosoftSSOHandler: response = response or {} verbose_proxy_logger.debug(f"Microsoft SSO Callback Response: {response}") openid_response = CustomOpenID( - email=response.get(MICROSOFT_USER_EMAIL_ATTRIBUTE) or response.get("mail"), + email=normalize_email(response.get(MICROSOFT_USER_EMAIL_ATTRIBUTE) or response.get("mail")), display_name=response.get(MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE), provider="microsoft", id=response.get(MICROSOFT_USER_ID_ATTRIBUTE), diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 16cdd9da64..12e34e221d 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) @@ -7752,6 +7807,7 @@ async def _apply_search_filter_to_models( size: int, prisma_client: Optional[Any], proxy_config: Any, + sort_by: Optional[str] = None, ) -> Tuple[List[Dict[str, Any]], Optional[int]]: """ Apply search filter to models, querying database for additional matching models. @@ -7763,6 +7819,7 @@ async def _apply_search_filter_to_models( size: Page size prisma_client: Prisma client for database queries proxy_config: Proxy config for decrypting models + sort_by: Optional sort field - if provided, fetch all matching models instead of paginating at DB level Returns: Tuple of (filtered_models, total_count). total_count is None if not searching. @@ -7826,17 +7883,15 @@ async def _apply_search_filter_to_models( # Calculate total count for search results search_total_count = router_models_count + db_models_total_count - # Fetch database models if we need more for the current page - if router_models_count < models_needed_for_page: - models_to_fetch = min( - models_needed_for_page - router_models_count, db_models_total_count - ) - - if models_to_fetch > 0: + # If sorting is requested, we need to fetch ALL matching models to sort correctly + # Otherwise, we can optimize by only fetching what's needed for the current page + if sort_by: + # Fetch all matching database models for sorting + if db_models_total_count > 0: db_models_raw = ( await prisma_client.db.litellm_proxymodeltable.find_many( where=db_where_condition, - take=models_to_fetch, + take=db_models_total_count, # Fetch all matching models ) ) @@ -7847,6 +7902,28 @@ async def _apply_search_filter_to_models( ) if decrypted_models: db_models.extend(decrypted_models) + else: + # Fetch database models if we need more for the current page + if router_models_count < models_needed_for_page: + models_to_fetch = min( + models_needed_for_page - router_models_count, db_models_total_count + ) + + if models_to_fetch > 0: + db_models_raw = ( + await prisma_client.db.litellm_proxymodeltable.find_many( + where=db_where_condition, + take=models_to_fetch, + ) + ) + + # Convert database models to router format + for db_model in db_models_raw: + decrypted_models = proxy_config.decrypt_model_list_from_db( + [db_model] + ) + if decrypted_models: + db_models.extend(decrypted_models) except Exception as e: verbose_proxy_logger.exception( f"Error querying database models with search: {str(e)}" @@ -7862,6 +7939,80 @@ async def _apply_search_filter_to_models( return filtered_models, search_total_count +def _sort_models( + all_models: List[Dict[str, Any]], + sort_by: Optional[str], + sort_order: str = "asc", +) -> List[Dict[str, Any]]: + """ + Sort models by the specified field and order. + + Args: + all_models: List of models to sort + sort_by: Field to sort by (model_name, created_at, updated_at, costs, status) + sort_order: Sort order (asc or desc) + + Returns: + Sorted list of models + """ + if not sort_by or sort_by not in ["model_name", "created_at", "updated_at", "costs", "status"]: + return all_models + + reverse = sort_order.lower() == "desc" + + def get_sort_key(model: Dict[str, Any]) -> Any: + model_info = model.get("model_info", {}) + + if sort_by == "model_name": + return model.get("model_name", "").lower() + + elif sort_by == "created_at": + created_at = model_info.get("created_at") + if created_at is None: + # Put None values at the end for asc, at the start for desc + return (datetime.max if not reverse else datetime.min) + if isinstance(created_at, str): + try: + return datetime.fromisoformat(created_at.replace("Z", "+00:00")) + except (ValueError, AttributeError): + return datetime.min if not reverse else datetime.max + return created_at + + elif sort_by == "updated_at": + updated_at = model_info.get("updated_at") + if updated_at is None: + return (datetime.max if not reverse else datetime.min) + if isinstance(updated_at, str): + try: + return datetime.fromisoformat(updated_at.replace("Z", "+00:00")) + except (ValueError, AttributeError): + return datetime.min if not reverse else datetime.max + return updated_at + + elif sort_by == "costs": + input_cost = model_info.get("input_cost_per_token", 0) or 0 + output_cost = model_info.get("output_cost_per_token", 0) or 0 + total_cost = input_cost + output_cost + # Put 0 or None costs at the end for asc, at the start for desc + if total_cost == 0: + return (float("inf") if not reverse else float("-inf")) + return total_cost + + elif sort_by == "status": + # False (config) comes before True (db) for asc + db_model = model_info.get("db_model", False) + return db_model + + return None + + try: + sorted_models = sorted(all_models, key=get_sort_key, reverse=reverse) + return sorted_models + except Exception as e: + verbose_proxy_logger.exception(f"Error sorting models by {sort_by}: {str(e)}") + return all_models + + def _paginate_models_response( all_models: List[Dict[str, Any]], page: int, @@ -8023,6 +8174,55 @@ async def _filter_models_by_team_id( return filtered_models +async def _find_model_by_id( + model_id: str, + search: Optional[str], + llm_router, + prisma_client, + proxy_config, +) -> tuple[list, Optional[int]]: + """Find a model by its ID and optionally filter by search term.""" + found_model = None + + # First, search in config + if llm_router is not None: + found_model = llm_router.get_model_info(id=model_id) + if found_model: + found_model = copy.deepcopy(found_model) + + # If not found in config, search in database + if found_model is None: + try: + db_model = await prisma_client.db.litellm_proxymodeltable.find_unique( + where={"model_id": model_id} + ) + if db_model: + # Convert database model to router format + decrypted_models = proxy_config.decrypt_model_list_from_db( + [db_model] + ) + if decrypted_models: + found_model = decrypted_models[0] + except Exception as e: + verbose_proxy_logger.exception( + f"Error querying database for modelId {model_id}: {str(e)}" + ) + + # If model found, verify search filter if provided + if found_model is not None: + if search is not None and search.strip(): + search_lower = search.lower().strip() + model_name = found_model.get("model_name", "") + if search_lower not in model_name.lower(): + # Model found but doesn't match search filter + found_model = None + + # Set all_models to the found model or empty list + all_models = [found_model] if found_model is not None else [] + search_total_count: Optional[int] = len(all_models) + return all_models, search_total_count + + @router.get( "/v2/model/info", description="v2 - returns models available to the user based on their API key permissions. Shows model info from config.yaml (except api key and api base). Filter to just user-added models with ?user_models_only=true", @@ -8054,6 +8254,14 @@ async def model_info_v2( None, description="Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids", ), + sortBy: Optional[str] = fastapi.Query( + None, + description="Field to sort by. Options: model_name, created_at, updated_at, costs, status", + ), + sortOrder: Optional[str] = fastapi.Query( + "asc", + description="Sort order. Options: asc, desc", + ), ): """ BETA ENDPOINT. Might change unexpectedly. Use `/v1/model/info` for now. @@ -8081,44 +8289,13 @@ async def model_info_v2( # If modelId is provided, search for the specific model if modelId is not None: - found_model = None - - # First, search in config - if llm_router is not None: - found_model = llm_router.get_model_info(id=modelId) - if found_model: - found_model = copy.deepcopy(found_model) - - # If not found in config, search in database - if found_model is None: - try: - db_model = await prisma_client.db.litellm_proxymodeltable.find_unique( - where={"model_id": modelId} - ) - if db_model: - # Convert database model to router format - decrypted_models = proxy_config.decrypt_model_list_from_db( - [db_model] - ) - if decrypted_models: - found_model = decrypted_models[0] - except Exception as e: - verbose_proxy_logger.exception( - f"Error querying database for modelId {modelId}: {str(e)}" - ) - - # If model found, verify search filter if provided - if found_model is not None: - if search is not None and search.strip(): - search_lower = search.lower().strip() - model_name = found_model.get("model_name", "") - if search_lower not in model_name.lower(): - # Model found but doesn't match search filter - found_model = None - - # Set all_models to the found model or empty list - all_models = [found_model] if found_model is not None else [] - search_total_count: Optional[int] = len(all_models) + all_models, search_total_count = await _find_model_by_id( + model_id=modelId, + search=search, + llm_router=llm_router, + prisma_client=prisma_client, + proxy_config=proxy_config, + ) else: # Normal flow when modelId is not provided all_models = copy.deepcopy(llm_router.model_list) @@ -8138,6 +8315,7 @@ async def model_info_v2( size=size, prisma_client=prisma_client, proxy_config=proxy_config, + sort_by=sortBy, ) if user_models_only: @@ -8181,6 +8359,20 @@ async def model_info_v2( if modelId is not None: search_total_count = len(all_models) + # Apply sorting before pagination + if sortBy: + # Validate sortOrder + if sortOrder and sortOrder.lower() not in ["asc", "desc"]: + raise HTTPException( + status_code=400, + detail=f"Invalid sortOrder: {sortOrder}. Must be 'asc' or 'desc'", + ) + all_models = _sort_models( + all_models=all_models, + sort_by=sortBy, + sort_order=sortOrder or "asc", + ) + verbose_proxy_logger.debug("all_models: %s", all_models) return _paginate_models_response( @@ -9645,7 +9837,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 +9856,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..2c34457c02 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -26,6 +26,193 @@ 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") + + # Extract provider-specific params from vector_store_config to save as litellm_params + # This ensures params like aws_region_name, embedding_model, etc. are available for search + provider_specific_params = {} + excluded_keys = {"custom_llm_provider", "vector_store_id"} + for key, value in vector_store_config.items(): + if key not in excluded_keys and value is not None: + provider_specific_params[key] = value + + # 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, + litellm_params=provider_specific_params if provider_specific_params else None, + ) + + 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 +345,7 @@ async def rag_ingest( add_litellm_data_to_request, general_settings, llm_router, + prisma_client, proxy_config, version, ) @@ -189,6 +377,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/route_llm_request.py b/litellm/proxy/route_llm_request.py index 236f58bb82..af441ee43e 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -221,8 +221,9 @@ async def route_request( "aretrieve_container_file_content", ]: return getattr(llm_router, f"{route_type}")(**data) - # Interactions API: get/delete/cancel don't need model routing + # Interactions API: create with agent, get/delete/cancel don't need model routing if route_type in [ + "acreate_interaction", "aget_interaction", "adelete_interaction", "acancel_interaction", 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/ingestion/base_ingestion.py b/litellm/rag/ingestion/base_ingestion.py index 20059487b4..3daa767188 100644 --- a/litellm/rag/ingestion/base_ingestion.py +++ b/litellm/rag/ingestion/base_ingestion.py @@ -24,6 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.rag.ingestion.file_parsers import extract_text_from_pdf from litellm.rag.text_splitters import RecursiveCharacterTextSplitter from litellm.types.rag import RAGIngestOptions, RAGIngestResponse @@ -193,11 +194,23 @@ class BaseRAGIngestion(ABC): if text: text_to_chunk = text elif file_content and not ocr_was_used: + # Try UTF-8 decode first try: text_to_chunk = file_content.decode("utf-8") except UnicodeDecodeError: - verbose_logger.debug("Binary file detected, skipping text chunking") - return [] + # Check if it's a PDF and try to extract text + if file_content.startswith(b"%PDF"): + verbose_logger.debug("PDF detected, attempting text extraction") + text_to_chunk = extract_text_from_pdf(file_content) + if not text_to_chunk: + verbose_logger.debug( + "PDF text extraction failed. Install 'pypdf' or 'PyPDF2' for PDF support, " + "or enable OCR with a vision model." + ) + return [] + else: + verbose_logger.debug("Binary file detected, skipping text chunking") + return [] if not text_to_chunk: return [] diff --git a/litellm/rag/ingestion/file_parsers/__init__.py b/litellm/rag/ingestion/file_parsers/__init__.py new file mode 100644 index 0000000000..5be68cdbb7 --- /dev/null +++ b/litellm/rag/ingestion/file_parsers/__init__.py @@ -0,0 +1,9 @@ +""" +File parsers for RAG ingestion. + +Provides text extraction utilities for various file formats. +""" + +from .pdf_parser import extract_text_from_pdf + +__all__ = ["extract_text_from_pdf"] diff --git a/litellm/rag/ingestion/file_parsers/pdf_parser.py b/litellm/rag/ingestion/file_parsers/pdf_parser.py new file mode 100644 index 0000000000..9a533ccf13 --- /dev/null +++ b/litellm/rag/ingestion/file_parsers/pdf_parser.py @@ -0,0 +1,70 @@ +""" +PDF text extraction utilities. + +Provides text extraction from PDF files using pypdf or PyPDF2. +""" + +from typing import Optional + +from litellm._logging import verbose_logger + + +def extract_text_from_pdf(file_content: bytes) -> Optional[str]: + """ + Extract text from PDF using pypdf if available. + + Args: + file_content: Raw PDF bytes + + Returns: + Extracted text or None if extraction fails + """ + try: + from io import BytesIO + + # Try pypdf first (most common) + try: + from pypdf import PdfReader as PypdfReader + + pdf_file = BytesIO(file_content) + reader = PypdfReader(pdf_file) + + text_parts = [] + for page in reader.pages: + text = page.extract_text() + if text: + text_parts.append(text) + + if text_parts: + extracted_text = "\n\n".join(text_parts) + verbose_logger.debug(f"Extracted {len(extracted_text)} characters from PDF using pypdf") + return extracted_text + + except ImportError: + verbose_logger.debug("pypdf not available, trying PyPDF2") + + # Fallback to PyPDF2 + try: + from PyPDF2 import PdfReader as PyPDF2Reader + + pdf_file = BytesIO(file_content) + reader = PyPDF2Reader(pdf_file) + + text_parts = [] + for page in reader.pages: + text = page.extract_text() + if text: + text_parts.append(text) + + if text_parts: + extracted_text = "\n\n".join(text_parts) + verbose_logger.debug(f"Extracted {len(extracted_text)} characters from PDF using PyPDF2") + return extracted_text + + except ImportError: + verbose_logger.debug("PyPDF2 not available, PDF extraction requires OCR or pypdf/PyPDF2 library") + + except Exception as e: + verbose_logger.debug(f"PDF text extraction failed: {e}") + + return None diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py new file mode 100644 index 0000000000..e6c166a101 --- /dev/null +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -0,0 +1,573 @@ +""" +S3 Vectors-specific RAG Ingestion implementation. + +S3 Vectors is AWS's native vector storage service that provides: +- Purpose-built vector buckets for storing and querying vectors +- Vector indexes with configurable dimensions and distance metrics +- Metadata filtering for semantic search + +This implementation: +1. Auto-creates vector buckets and indexes if not provided +2. Uses LiteLLM's embedding API (supports any provider) +3. Uses httpx + AWS SigV4 signing (no boto3 dependency for S3 Vectors APIs) +4. Stores vectors with metadata using PutVectors API +""" + +from __future__ import annotations + +import hashlib +import uuid +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +import litellm +from litellm._logging import verbose_logger +from litellm.constants import ( + S3_VECTORS_DEFAULT_DIMENSION, + S3_VECTORS_DEFAULT_DISTANCE_METRIC, + S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS, +) +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion + +if TYPE_CHECKING: + from litellm import Router + from litellm.types.rag import RAGIngestOptions + + +class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): + """ + S3 Vectors RAG ingestion using httpx + AWS SigV4 signing. + + Workflow: + 1. Auto-create vector bucket if needed (CreateVectorBucket API) + 2. Auto-create vector index if needed (CreateVectorIndex API) + 3. Generate embeddings using LiteLLM (supports any provider) + 4. Store vectors with PutVectors API + + Configuration: + - vector_bucket_name: S3 vector bucket name (required) + - index_name: Vector index name (auto-creates if not provided) + - dimension: Vector dimension (default: S3_VECTORS_DEFAULT_DIMENSION) + - distance_metric: "cosine" or "euclidean" (default: S3_VECTORS_DEFAULT_DISTANCE_METRIC) + - non_filterable_metadata_keys: List of metadata keys to exclude from filtering + """ + + def __init__( + self, + ingest_options: "RAGIngestOptions", + router: Optional["Router"] = None, + ): + BaseRAGIngestion.__init__(self, ingest_options=ingest_options, router=router) + BaseAWSLLM.__init__(self) + + # Extract config + self.vector_bucket_name = self.vector_store_config["vector_bucket_name"] + self.index_name = self.vector_store_config.get("index_name") + self.distance_metric = self.vector_store_config.get( + "distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC + ) + self.non_filterable_metadata_keys = self.vector_store_config.get( + "non_filterable_metadata_keys", + S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS, + ) + + # Get dimension from config (will be auto-detected on first use if not provided) + self.dimension = self._get_dimension_from_config() + + # Get AWS region using BaseAWSLLM method + _aws_region = self.vector_store_config.get("aws_region_name") + self.aws_region_name = self.get_aws_region_name_for_non_llm_api_calls( + aws_region_name=str(_aws_region) if _aws_region else None + ) + + # Create httpx client (similar to s3_v2.py) + ssl_verify = self._get_ssl_verify( + ssl_verify=self.vector_store_config.get("ssl_verify") + ) + self.async_httpx_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.RAG, + params={"ssl_verify": ssl_verify} if ssl_verify is not None else None, + ) + + # Track if infrastructure is initialized + self._config_initialized = False + + async def _get_dimension_from_embedding_request(self) -> int: + """ + Auto-detect dimension by making a test embedding request. + + Makes a single embedding request with a test string to determine + the output dimension of the embedding model. + """ + if not self.embedding_config or "model" not in self.embedding_config: + return S3_VECTORS_DEFAULT_DIMENSION + + try: + model_name = self.embedding_config["model"] + verbose_logger.debug( + f"Auto-detecting dimension by making test embedding request to {model_name}" + ) + + # Make a test embedding request + test_input = "test" + if self.router: + response = await self.router.aembedding(model=model_name, input=[test_input]) + else: + response = await litellm.aembedding(model=model_name, input=[test_input]) + + # Get dimension from the response + if response.data and len(response.data) > 0: + dimension = len(response.data[0]["embedding"]) + verbose_logger.debug( + f"Auto-detected dimension {dimension} for embedding model {model_name}" + ) + return dimension + except Exception as e: + verbose_logger.warning( + f"Could not auto-detect dimension from embedding model: {e}. " + f"Using default dimension of {S3_VECTORS_DEFAULT_DIMENSION}." + ) + + return S3_VECTORS_DEFAULT_DIMENSION + + def _get_dimension_from_config(self) -> Optional[int]: + """ + Get vector dimension from config if explicitly provided. + + Returns None if dimension should be auto-detected. + """ + if "dimension" in self.vector_store_config: + return int(self.vector_store_config["dimension"]) + return None + + async def _ensure_config_initialized(self): + """Lazily initialize S3 Vectors infrastructure.""" + if self._config_initialized: + return + + # Auto-detect dimension if not provided + if self.dimension is None: + self.dimension = await self._get_dimension_from_embedding_request() + + # Ensure vector bucket exists + await self._ensure_vector_bucket_exists() + + # Ensure vector index exists + if not self.index_name: + # Auto-generate index name + unique_id = uuid.uuid4().hex[:8] + self.index_name = f"litellm-index-{unique_id}" + + await self._ensure_vector_index_exists() + + self._config_initialized = True + + async def _sign_and_execute_request( + self, + method: str, + url: str, + data: Optional[str] = None, + headers: Optional[Dict[str, str]] = None, + ) -> Any: + """ + Helper to sign and execute AWS API requests using httpx + SigV4. + + Pattern from litellm/integrations/s3_v2.py + """ + try: + import requests + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError( + "Missing botocore to call S3 Vectors. Run 'pip install boto3'." + ) + + # Get AWS credentials using BaseAWSLLM's get_credentials method + credentials = self.get_credentials( + aws_access_key_id=self.vector_store_config.get("aws_access_key_id"), + aws_secret_access_key=self.vector_store_config.get("aws_secret_access_key"), + aws_session_token=self.vector_store_config.get("aws_session_token"), + aws_region_name=self.aws_region_name, + aws_session_name=self.vector_store_config.get("aws_session_name"), + aws_profile_name=self.vector_store_config.get("aws_profile_name"), + aws_role_name=self.vector_store_config.get("aws_role_name"), + aws_web_identity_token=self.vector_store_config.get("aws_web_identity_token"), + aws_sts_endpoint=self.vector_store_config.get("aws_sts_endpoint"), + aws_external_id=self.vector_store_config.get("aws_external_id"), + ) + + # Prepare headers + if headers is None: + headers = {} + + if data: + headers["Content-Type"] = "application/json" + # Calculate SHA256 hash of the content + content_hash = hashlib.sha256(data.encode("utf-8")).hexdigest() + headers["x-amz-content-sha256"] = content_hash + else: + # For requests without body, use hash of empty string + headers["x-amz-content-sha256"] = hashlib.sha256(b"").hexdigest() + + # Prepare the request + req = requests.Request(method, url, data=data, headers=headers) + prepped = req.prepare() + + # Sign the request + aws_request = AWSRequest( + method=prepped.method, + url=prepped.url, + data=prepped.body, + headers=prepped.headers, + ) + SigV4Auth(credentials, "s3vectors", self.aws_region_name).add_auth(aws_request) + + # Prepare the signed headers + signed_headers = dict(aws_request.headers.items()) + + # Make the request using specific method (pattern from s3_v2.py) + method_upper = method.upper() + if method_upper == "PUT": + response = await self.async_httpx_client.put( + url, data=data, headers=signed_headers + ) + elif method_upper == "POST": + response = await self.async_httpx_client.post( + url, data=data, headers=signed_headers + ) + elif method_upper == "GET": + response = await self.async_httpx_client.get(url, headers=signed_headers) + else: + raise ValueError(f"Unsupported HTTP method: {method}") + + return response + + async def _ensure_vector_bucket_exists(self): + """Create vector bucket if it doesn't exist using GetVectorBucket and CreateVectorBucket APIs.""" + verbose_logger.debug( + f"Ensuring S3 vector bucket exists: {self.vector_bucket_name}" + ) + + # Validate bucket name (AWS S3 naming rules) + if len(self.vector_bucket_name) < 3: + raise ValueError( + f"Invalid vector_bucket_name '{self.vector_bucket_name}': " + f"AWS S3 bucket names must be at least 3 characters long. " + f"Please provide a valid bucket name (e.g., 'my-vector-bucket')." + ) + if not self.vector_bucket_name.replace("-", "").replace(".", "").isalnum(): + raise ValueError( + f"Invalid vector_bucket_name '{self.vector_bucket_name}': " + f"AWS S3 bucket names can only contain lowercase letters, numbers, hyphens, and periods. " + f"Please provide a valid bucket name (e.g., 'my-vector-bucket')." + ) + + # Try to get bucket info using GetVectorBucket API + get_url = f"https://s3vectors.{self.aws_region_name}.api.aws/GetVectorBucket" + get_body = safe_dumps({"vectorBucketName": self.vector_bucket_name}) + + try: + response = await self._sign_and_execute_request("POST", get_url, data=get_body) + if response.status_code == 200: + verbose_logger.debug(f"Vector bucket {self.vector_bucket_name} exists") + return + except Exception as e: + verbose_logger.debug( + f"Bucket check failed (may not exist): {e}, attempting to create" + ) + + # Create vector bucket using CreateVectorBucket API + try: + verbose_logger.debug(f"Creating vector bucket: {self.vector_bucket_name}") + create_url = f"https://s3vectors.{self.aws_region_name}.api.aws/CreateVectorBucket" + create_body = safe_dumps({ + "vectorBucketName": self.vector_bucket_name + }) + + response = await self._sign_and_execute_request("POST", create_url, data=create_body) + + if response.status_code in (200, 201): + verbose_logger.info(f"Created vector bucket: {self.vector_bucket_name}") + elif response.status_code == 409: + # Bucket already exists (ConflictException) + verbose_logger.debug( + f"Vector bucket {self.vector_bucket_name} already exists" + ) + else: + verbose_logger.error(f"CreateVectorBucket failed: {response.status_code} - {response.text}") + response.raise_for_status() + except Exception as e: + verbose_logger.exception(f"Error creating vector bucket: {e}") + raise + + async def _ensure_vector_index_exists(self): + """Create vector index if it doesn't exist using GetIndex and CreateIndex APIs.""" + verbose_logger.debug( + f"Ensuring vector index exists: {self.vector_bucket_name}/{self.index_name}" + ) + + # Try to get index info using GetIndex API + get_url = f"https://s3vectors.{self.aws_region_name}.api.aws/GetIndex" + get_body = safe_dumps({ + "vectorBucketName": self.vector_bucket_name, + "indexName": self.index_name + }) + + try: + response = await self._sign_and_execute_request("POST", get_url, data=get_body) + if response.status_code == 200: + verbose_logger.debug(f"Vector index {self.index_name} exists") + return + except Exception as e: + verbose_logger.debug( + f"Index check failed (may not exist): {e}, attempting to create" + ) + + # Create vector index using CreateIndex API + try: + verbose_logger.debug( + f"Creating vector index: {self.index_name} with dimension={self.dimension}, metric={self.distance_metric}" + ) + + # Prepare index configuration per AWS API docs + index_config = { + "vectorBucketName": self.vector_bucket_name, + "indexName": self.index_name, + "dataType": "float32", + "dimension": self.dimension, + "distanceMetric": self.distance_metric, + } + + if self.non_filterable_metadata_keys: + index_config["metadataConfiguration"] = { + "nonFilterableMetadataKeys": self.non_filterable_metadata_keys + } + + create_url = f"https://s3vectors.{self.aws_region_name}.api.aws/CreateIndex" + response = await self._sign_and_execute_request( + "POST", create_url, data=safe_dumps(index_config) + ) + + if response.status_code in (200, 201): + verbose_logger.info(f"Created vector index: {self.index_name}") + elif response.status_code == 409: + verbose_logger.debug(f"Vector index {self.index_name} already exists") + else: + verbose_logger.error(f"CreateIndex failed: {response.status_code} - {response.text}") + response.raise_for_status() + except Exception as e: + verbose_logger.exception(f"Error creating vector index: {e}") + raise + + async def _put_vectors(self, vectors: List[Dict[str, Any]]): + """ + Call PutVectors API to store vectors in S3 Vectors. + + Args: + vectors: List of vector objects with keys: "key", "data", "metadata" + """ + verbose_logger.debug( + f"Storing {len(vectors)} vectors in {self.vector_bucket_name}/{self.index_name}" + ) + + url = f"https://s3vectors.{self.aws_region_name}.api.aws/PutVectors" + + # Prepare request body per AWS API docs + request_body = { + "vectorBucketName": self.vector_bucket_name, + "indexName": self.index_name, + "vectors": vectors + } + + try: + response = await self._sign_and_execute_request( + "POST", url, data=safe_dumps(request_body) + ) + + if response.status_code in (200, 201): + verbose_logger.info( + f"Successfully stored {len(vectors)} vectors in index {self.index_name}" + ) + else: + verbose_logger.error( + f"PutVectors failed with status {response.status_code}: {response.text}" + ) + response.raise_for_status() + except Exception as e: + verbose_logger.exception(f"Error storing vectors: {e}") + raise + + async def embed( + self, + chunks: List[str], + ) -> Optional[List[List[float]]]: + """ + Generate embeddings using LiteLLM's embedding API. + + Supports any embedding provider (OpenAI, Bedrock, Cohere, etc.) + """ + if not chunks: + return None + + # Use embedding config from ingest_options or default + if not self.embedding_config: + verbose_logger.warning( + "No embedding config provided, using default text-embedding-3-small" + ) + self.embedding_config = {"model": "text-embedding-3-small"} + + embedding_model = self.embedding_config.get("model", "text-embedding-3-small") + + verbose_logger.debug( + f"Generating embeddings for {len(chunks)} chunks using {embedding_model}" + ) + + # Convert to list to ensure type compatibility + input_chunks: List[str] = list(chunks) + + if self.router: + response = await self.router.aembedding(model=embedding_model, input=input_chunks) + else: + response = await litellm.aembedding(model=embedding_model, input=input_chunks) + + return [item["embedding"] for item in response.data] + + async def store( + self, + file_content: Optional[bytes], + filename: Optional[str], + content_type: Optional[str], + chunks: List[str], + embeddings: Optional[List[List[float]]], + ) -> Tuple[Optional[str], Optional[str]]: + """ + Store vectors in S3 Vectors using PutVectors API. + + Steps: + 1. Ensure vector bucket exists (auto-create if needed) + 2. Ensure vector index exists (auto-create if needed) + 3. Prepare vector data with metadata + 4. Call PutVectors API with httpx + SigV4 signing + + Args: + file_content: Raw file bytes (not used for S3 Vectors) + filename: Name of the file + content_type: MIME type (not used for S3 Vectors) + chunks: Text chunks + embeddings: Vector embeddings + + Returns: + Tuple of (index_name, filename) + """ + # Ensure infrastructure exists + await self._ensure_config_initialized() + + if not embeddings or not chunks: + error_msg = ( + "No text content could be extracted from the file for embedding. " + "Possible causes:\n" + " 1. PDF files require OCR - add 'ocr' config with a vision model (e.g., 'anthropic/claude-3-5-sonnet-20241022')\n" + " 2. Binary files cannot be processed - convert to text first\n" + " 3. File is empty or contains no extractable text\n" + "For PDFs, either enable OCR or use a PDF extraction library to convert to text before ingestion." + ) + verbose_logger.error(error_msg) + raise ValueError(error_msg) + + # Prepare vectors for PutVectors API + vectors = [] + for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)): + # Build metadata dict + metadata: Dict[str, str] = { + "source_text": chunk, # Non-filterable (for reference) + "chunk_index": str(i), # Filterable + } + + if filename: + metadata["filename"] = filename # Filterable + + vector_obj = { + "key": f"{filename}_{i}" if filename else f"chunk_{i}", + "data": {"float32": embedding}, + "metadata": metadata, + } + + vectors.append(vector_obj) + + # Call PutVectors API + await self._put_vectors(vectors) + + # Return vector_store_id in format bucket_name:index_name for S3 Vectors search compatibility + vector_store_id = f"{self.vector_bucket_name}:{self.index_name}" + return vector_store_id, filename + + async def query_vector_store( + self, vector_store_id: str, query: str, top_k: int = 5 + ) -> Optional[Dict[str, Any]]: + """ + Query S3 Vectors using QueryVectors API. + + Args: + vector_store_id: Index name + query: Query text + top_k: Number of results to return + + Returns: + Query results with vectors and metadata + """ + verbose_logger.debug(f"Querying index {vector_store_id} with query: {query}") + + # Generate query embedding + if not self.embedding_config: + self.embedding_config = {"model": "text-embedding-3-small"} + + embedding_model = self.embedding_config.get("model", "text-embedding-3-small") + + response = await litellm.aembedding(model=embedding_model, input=[query]) + query_embedding = response.data[0]["embedding"] + + # Call QueryVectors API + url = f"https://s3vectors.{self.aws_region_name}.api.aws/QueryVectors" + + request_body = { + "vectorBucketName": self.vector_bucket_name, + "indexName": vector_store_id, + "queryVector": {"float32": query_embedding}, + "topK": top_k, + "returnDistance": True, + "returnMetadata": True, + } + + try: + response = await self._sign_and_execute_request( + "POST", url, data=safe_dumps(request_body) + ) + + if response.status_code == 200: + results = response.json() + verbose_logger.debug(f"Query returned {len(results.get('vectors', []))} results") + + # Check if query terms appear in results + if results.get("vectors"): + for result in results["vectors"]: + metadata = result.get("metadata", {}) + source_text = metadata.get("source_text", "") + if query.lower() in source_text.lower(): + return results + + # Return results even if exact match not found + return results + else: + verbose_logger.error( + f"QueryVectors failed with status {response.status_code}: {response.text}" + ) + return None + except Exception as e: + verbose_logger.exception(f"Error querying vectors: {e}") + return None diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 0ccbb435e5..571f78f2f7 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -31,6 +31,7 @@ from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion from litellm.rag.ingestion.bedrock_ingestion import BedrockRAGIngestion from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion from litellm.rag.ingestion.openai_ingestion import OpenAIRAGIngestion +from litellm.rag.ingestion.s3_vectors_ingestion import S3VectorsRAGIngestion from litellm.rag.rag_query import RAGQuery from litellm.types.rag import ( RAGIngestOptions, @@ -48,6 +49,7 @@ INGESTION_REGISTRY: Dict[str, Type[BaseRAGIngestion]] = { "openai": OpenAIRAGIngestion, "bedrock": BedrockRAGIngestion, "gemini": GeminiRAGIngestion, + "s3_vectors": S3VectorsRAGIngestion, } @@ -198,6 +200,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 +239,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..09d71b6b49 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -58,6 +58,7 @@ from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_metadata_variable_name_from_kwargs, ) +from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dd_tracing import tracer @@ -1250,10 +1251,25 @@ class Router: specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + # Check for silent model experiment + # Make a local copy of litellm_params to avoid mutating the Router's state + litellm_params = deployment["litellm_params"].copy() + silent_model = litellm_params.pop("silent_model", None) + if silent_model is not None: + # Mirroring traffic to a secondary model + # Use shared thread pool for background calls + executor.submit( + self._silent_experiment_completion, + silent_model, + messages, + **kwargs, + ) + + self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + kwargs.pop("silent_model", None) # Ensure it's not in kwargs either # No copy needed - data is only read and spread into new dict below - data = deployment["litellm_params"] + data = litellm_params.copy() # Use the local copy of litellm_params model_name = data["model"] potential_model_client = self._get_client( deployment=deployment, kwargs=kwargs @@ -1274,15 +1290,14 @@ class Router: if not self.has_model_id(model): self.routing_strategy_pre_call_checks(deployment=deployment) - response = litellm.completion( - **{ - **data, - "messages": messages, - "caching": self.cache_responses, - "client": model_client, - **kwargs, - } - ) + input_kwargs = { + **data, + "messages": messages, + "caching": self.cache_responses, + "client": model_client, + **kwargs, + } + response = litellm.completion(**input_kwargs) verbose_router_logger.info( f"litellm.completion(model={model_name})\033[32m 200 OK\033[0m" ) @@ -1309,6 +1324,56 @@ class Router: self._set_deployment_num_retries_on_exception(e, deployment) raise e + def _get_silent_experiment_kwargs(self, **kwargs) -> dict: + """ + Prepare kwargs for a silent experiment by ensuring isolation from the primary call. + """ + # Copy kwargs to ensure isolation + silent_kwargs = copy.deepcopy(kwargs) + if "metadata" not in silent_kwargs: + silent_kwargs["metadata"] = {} + + silent_kwargs["metadata"]["is_silent_experiment"] = True + + # Pop logging objects and call IDs to ensure a fresh logging context + # This prevents collisions in the Proxy's database (spend_logs) + silent_kwargs.pop("litellm_call_id", None) + silent_kwargs.pop("litellm_logging_obj", None) + silent_kwargs.pop("standard_logging_object", None) + silent_kwargs.pop("proxy_server_request", None) + + return silent_kwargs + + def _silent_experiment_completion( + self, silent_model: str, messages: List[Any], **kwargs + ): + """ + Run a silent experiment in the background (thread). + """ + try: + # Prevent infinite recursion if silent model also has a silent model + if kwargs.get("metadata", {}).get("is_silent_experiment", False): + return + + messages = copy.deepcopy(messages) + + verbose_router_logger.info( + f"Starting silent experiment for model {silent_model}" + ) + + silent_kwargs = self._get_silent_experiment_kwargs(**kwargs) + + # Trigger the silent request + self.completion( + model=silent_model, + messages=cast(List[Dict[str, str]], messages), + **silent_kwargs, + ) + except Exception as e: + verbose_router_logger.error( + f"Silent experiment failed for model {silent_model}: {str(e)}" + ) + # fmt: off @overload @@ -1517,6 +1582,36 @@ class Router: return FallbackStreamWrapper(stream_with_fallbacks()) + async def _silent_experiment_acompletion( + self, silent_model: str, messages: List[Any], **kwargs + ): + """ + Run a silent experiment in the background. + """ + try: + # Prevent infinite recursion if silent model also has a silent model + if kwargs.get("metadata", {}).get("is_silent_experiment", False): + return + + messages = copy.deepcopy(messages) + + verbose_router_logger.info( + f"Starting silent experiment for model {silent_model}" + ) + + silent_kwargs = self._get_silent_experiment_kwargs(**kwargs) + + # Trigger the silent request + await self.acompletion( + model=silent_model, + messages=cast(List[AllMessageValues], messages), + **silent_kwargs, + ) + except Exception as e: + verbose_router_logger.error( + f"Silent experiment failed for model {silent_model}: {str(e)}" + ) + async def _acompletion( # noqa: PLR0915 self, model: str, messages: List[Dict[str, str]], **kwargs ) -> Union[ModelResponse, CustomStreamWrapper,]: @@ -1563,9 +1658,27 @@ class Router: self._track_deployment_metrics( deployment=deployment, parent_otel_span=parent_otel_span ) + + # Check for silent model experiment + # Make a local copy of litellm_params to avoid mutating the Router's state + litellm_params = deployment["litellm_params"].copy() + silent_model = litellm_params.pop("silent_model", None) + + if silent_model is not None: + # Mirroring traffic to a secondary model + # This is a silent experiment, so we don't want to block the primary request + asyncio.create_task( + self._silent_experiment_acompletion( + silent_model=silent_model, + messages=messages, # Use messages instead of *args + **kwargs, + ) + ) + self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + kwargs.pop("silent_model", None) # Ensure it's not in kwargs either # No copy needed - data is only read and spread into new dict below - data = deployment["litellm_params"] + data = litellm_params.copy() # Use the local copy of litellm_params model_name = data["model"] @@ -1582,6 +1695,7 @@ class Router: "client": model_client, **kwargs, } + input_kwargs.pop("silent_model", None) _response = litellm.acompletion(**input_kwargs) @@ -1707,8 +1821,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 fd9b722287..ee49ba1a19 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 = [ @@ -217,6 +227,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.API_KEY_ALIAS.value, UserAPIKeyLabelNames.TEAM.value, UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_request_total_latency_metric = [ @@ -228,6 +239,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.TEAM_ALIAS.value, UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_request_queue_time_seconds = [ @@ -239,6 +251,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) @@ -258,6 +271,9 @@ 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, ] litellm_proxy_failed_requests_metric = [ @@ -272,6 +288,9 @@ 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, ] litellm_deployment_latency_per_output_token = [ @@ -292,6 +311,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 = [ @@ -301,6 +321,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 = [ @@ -310,6 +331,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 = [ @@ -321,6 +343,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 = [ @@ -332,6 +357,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 = [ @@ -344,6 +372,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.USER_EMAIL.value, UserAPIKeyLabelNames.REQUESTED_MODEL.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_total_tokens_metric = [ @@ -356,6 +385,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.USER_EMAIL.value, UserAPIKeyLabelNames.REQUESTED_MODEL.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_output_tokens_metric = [ @@ -368,6 +398,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.USER_EMAIL.value, UserAPIKeyLabelNames.REQUESTED_MODEL.value, + UserAPIKeyLabelNames.MODEL_ID.value, ] litellm_deployment_state = [ @@ -377,6 +408,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, @@ -394,6 +434,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 @@ -436,6 +477,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, @@ -449,6 +510,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 = [ @@ -461,10 +524,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] = [] @@ -485,6 +575,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 @@ -577,6 +668,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/rag.py b/litellm/types/rag.py index fe237a1343..cae0770868 100644 --- a/litellm/types/rag.py +++ b/litellm/types/rag.py @@ -129,9 +129,53 @@ class VertexAIVectorStoreOptions(TypedDict, total=False): import_timeout: Optional[int] # Timeout in seconds (default: 600) +class S3VectorsVectorStoreOptions(TypedDict, total=False): + """ + AWS S3 Vectors configuration. + + Example (auto-create): + {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "my-embeddings"} + + Example (use existing): + {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "my-embeddings", + "index_name": "my-index"} + + Example (with credentials): + {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "my-embeddings", + "litellm_credential_name": "my-aws-creds"} + + Auto-creation creates: S3 vector bucket and vector index (if not provided). + Embeddings are generated using LiteLLM's embedding API (supports any provider). + """ + + custom_llm_provider: Literal["s3_vectors"] + vector_bucket_name: str # Required - S3 vector bucket name + index_name: Optional[str] # Vector index name (auto-creates if not provided) + + # Index configuration (for auto-creation) + dimension: Optional[int] # Vector dimension (auto-detected from embedding model, or default: 1024) + distance_metric: Optional[Literal["cosine", "euclidean"]] # Default: cosine + non_filterable_metadata_keys: Optional[List[str]] # Keys excluded from filtering (e.g., ["source_text"]) + + # Credentials (loaded from litellm.credential_list if litellm_credential_name is provided) + litellm_credential_name: Optional[str] # Credential name to load from litellm.credential_list + + # AWS auth (uses BaseAWSLLM) + aws_access_key_id: Optional[str] + aws_secret_access_key: Optional[str] + aws_session_token: Optional[str] + aws_region_name: Optional[str] # default: us-west-2 + aws_role_name: Optional[str] + aws_session_name: Optional[str] + aws_profile_name: Optional[str] + aws_web_identity_token: Optional[str] + aws_sts_endpoint: Optional[str] + aws_external_id: Optional[str] + + # Union type for vector store options RAGIngestVectorStoreOptions = Union[ - OpenAIVectorStoreOptions, BedrockVectorStoreOptions, VertexAIVectorStoreOptions + OpenAIVectorStoreOptions, BedrockVectorStoreOptions, VertexAIVectorStoreOptions, S3VectorsVectorStoreOptions ] diff --git a/litellm/types/router.py b/litellm/types/router.py index 8ea7a20753..43943d9e07 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -404,6 +404,10 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): aws_access_key_id: Optional[str] aws_secret_access_key: Optional[str] aws_region_name: Optional[str] + ## AWS S3 VECTORS ## + vector_bucket_name: Optional[str] + index_name: Optional[str] + embedding_model: Optional[str] ## IBM WATSONX ## watsonx_region_name: Optional[str] ## CUSTOM PRICING ## diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cd797dd1e5..49c903502d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3,25 +3,28 @@ 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 as Categories from openai.types.moderation import ( - Categories, - CategoryAppliedInputTypes, - CategoryScores, + CategoryAppliedInputTypes as CategoryAppliedInputTypes, +) +from openai.types.moderation import CategoryScores as CategoryScores +from openai.types.moderation_create_response import Moderation as Moderation +from openai.types.moderation_create_response import ( + 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 +55,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 +1414,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 +2504,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 +2690,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] @@ -3072,6 +3077,7 @@ class LlmProviders(str, Enum): LLAMA = "meta_llama" NSCALE = "nscale" PG_VECTOR = "pg_vector" + S3_VECTORS = "s3_vectors" HELICONE = "helicone" HYPERBOLIC = "hyperbolic" RECRAFT = "recraft" @@ -3428,3 +3434,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..d7fb4855a4 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: @@ -8446,6 +8452,12 @@ class ProviderConfigManager: ) return RAGFlowVectorStoreConfig() + elif litellm.LlmProviders.S3_VECTORS == provider: + from litellm.llms.s3_vectors.vector_stores.transformation import ( + S3VectorsVectorStoreConfig, + ) + + return S3VectorsVectorStoreConfig() return None @staticmethod @@ -8693,9 +8705,9 @@ class ProviderConfigManager: """ Get Search configuration for a given provider. """ + from litellm.llms.brave.search.transformation import BraveSearchConfig from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig from litellm.llms.exa_ai.search.transformation import ExaAISearchConfig - from litellm.llms.brave.search.transformation import BraveSearchConfig from litellm.llms.firecrawl.search.transformation import FirecrawlSearchConfig from litellm.llms.google_pse.search.transformation import GooglePSESearchConfig from litellm.llms.linkup.search.transformation import LinkupSearchConfig diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d958ea4503..4ac4159558 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -10232,6 +10232,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", @@ -13480,6 +13522,43 @@ "supports_vision": true, "supports_web_search": true }, + "gemini-robotics-er-1.5-preview": { + "cache_read_input_token_cost": 0, + "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 1e-06, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "output_cost_per_reasoning_token": 2.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-robotics-er-1-5-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions" + ], + "supported_modalities": [ + "text", + "image", + "video", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true + }, "gemini-2.5-computer-use-preview-10-2025": { "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, diff --git a/requirements.txt b/requirements.txt index 36328a081a..f250fb213e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -68,6 +68,7 @@ jsonschema>=4.23.0,<5.0.0 # validating json schema - aligned with openapi-core + websockets==15.0.1 # for realtime API soundfile==0.12.1 # for audio file processing openapi-core==0.21.0 # for OpenAPI compliance tests +pypdf>=6.6.2 # for PDF text extraction in RAG ingestion ######################## # LITELLM ENTERPRISE DEPENDENCIES diff --git a/tests/enterprise/litellm_enterprise/proxy/auth/test_route_checks.py b/tests/enterprise/litellm_enterprise/proxy/auth/test_route_checks.py index ce77ff69e0..706e3b7187 100644 --- a/tests/enterprise/litellm_enterprise/proxy/auth/test_route_checks.py +++ b/tests/enterprise/litellm_enterprise/proxy/auth/test_route_checks.py @@ -180,3 +180,68 @@ class TestEnterpriseRouteChecks: # Should not raise exception since management routes are enabled EnterpriseRouteChecks.should_call_route("/config/update") + + +class TestEnterpriseRouteChecksErrorMessages: + """Test that error messages correctly identify which feature requires Enterprise license""" + + @patch("litellm.secret_managers.main.get_secret_bool") + @patch("litellm.proxy.proxy_server.premium_user", False) + def test_disable_llm_api_endpoints_error_message(self, mock_get_secret_bool): + """ + Test that when DISABLE_LLM_API_ENDPOINTS is set without Enterprise license, + the error message correctly mentions 'LLM API ENDPOINTS' + """ + with patch.dict(os.environ, {"DISABLE_LLM_API_ENDPOINTS": "true"}): + with pytest.raises(HTTPException) as exc_info: + EnterpriseRouteChecks.is_llm_api_route_disabled() + + assert exc_info.value.status_code == 500 + assert "DISABLING LLM API ENDPOINTS is an Enterprise feature" in str( + exc_info.value.detail + ) + + @patch("litellm.secret_managers.main.get_secret_bool") + @patch("litellm.proxy.proxy_server.premium_user", False) + def test_disable_admin_endpoints_error_message(self, mock_get_secret_bool): + """ + Test that when DISABLE_ADMIN_ENDPOINTS is set without Enterprise license, + the error message correctly mentions 'ADMIN ENDPOINTS' (not 'LLM API ENDPOINTS') + + This is a regression test for a bug where the error message incorrectly said + 'DISABLING LLM API ENDPOINTS' when the actual issue was DISABLE_ADMIN_ENDPOINTS. + """ + with patch.dict(os.environ, {"DISABLE_ADMIN_ENDPOINTS": "true"}): + with pytest.raises(HTTPException) as exc_info: + EnterpriseRouteChecks.is_management_routes_disabled() + + assert exc_info.value.status_code == 500 + assert "DISABLING ADMIN ENDPOINTS is an Enterprise feature" in str( + exc_info.value.detail + ) + # Ensure it does NOT mention LLM API ENDPOINTS (the old buggy message) + assert "LLM API ENDPOINTS" not in str(exc_info.value.detail) + + @patch("litellm.secret_managers.main.get_secret_bool") + @patch("litellm.proxy.proxy_server.premium_user", True) + def test_disable_llm_api_endpoints_with_premium_user(self, mock_get_secret_bool): + """ + Test that premium users can use DISABLE_LLM_API_ENDPOINTS without error + """ + mock_get_secret_bool.return_value = True + with patch.dict(os.environ, {"DISABLE_LLM_API_ENDPOINTS": "true"}): + # Should not raise exception for premium users + result = EnterpriseRouteChecks.is_llm_api_route_disabled() + assert result is True + + @patch("litellm.secret_managers.main.get_secret_bool") + @patch("litellm.proxy.proxy_server.premium_user", True) + def test_disable_admin_endpoints_with_premium_user(self, mock_get_secret_bool): + """ + Test that premium users can use DISABLE_ADMIN_ENDPOINTS without error + """ + mock_get_secret_bool.return_value = True + with patch.dict(os.environ, {"DISABLE_ADMIN_ENDPOINTS": "true"}): + # Should not raise exception for premium users + result = EnterpriseRouteChecks.is_management_routes_disabled() + assert result is True 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/llm_translation/test_groq.py b/tests/llm_translation/test_groq.py index 33c1a42587..b82cffd418 100644 --- a/tests/llm_translation/test_groq.py +++ b/tests/llm_translation/test_groq.py @@ -11,7 +11,10 @@ import pytest import litellm from base_llm_unit_tests import BaseLLMChatTest -from litellm.llms.groq.chat.transformation import GroqChatConfig +from litellm.llms.groq.chat.transformation import ( + GroqChatConfig, + GroqChatCompletionStreamingHandler, +) class TestGroq(BaseLLMChatTest): def get_base_completion_call_args(self) -> dict: @@ -164,3 +167,124 @@ class TestGroqStructuredOutputs: if "tools" in result: tool_names = [t.get("function", {}).get("name") for t in result["tools"]] assert "json_tool_call" not in tool_names + + +class TestGroqReasoning: + """ + Tests for Groq reasoning field mapping. + + Groq returns 'reasoning' field in delta, but LiteLLM expects 'reasoning_content'. + """ + + def test_reasoning_field_mapping_in_streaming_chunks(self): + """ + Test that Groq's 'reasoning' field in streaming chunks is properly mapped + to LiteLLM's 'reasoning_content' field. + """ + handler = GroqChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + # Simulate a chunk with reasoning field as returned by Groq + groq_chunk = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1769511767, + "model": "qwen/qwen3-32b", + "choices": [ + { + "delta": { + "reasoning": "This is reasoning content", + "role": None, + }, + "finish_reason": None, + "index": 0, + } + ], + } + + # Parse the chunk + parsed_chunk = handler.chunk_parser(groq_chunk) + + # Verify that reasoning was mapped to reasoning_content + assert parsed_chunk.choices[0].delta.reasoning_content == "This is reasoning content" + # Verify that the original 'reasoning' field was removed + assert not hasattr(parsed_chunk.choices[0].delta, "reasoning") + + def test_reasoning_field_not_present(self): + """ + Test that chunks without reasoning field still work correctly. + """ + handler = GroqChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + # Simulate a chunk without reasoning field + groq_chunk = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1769511767, + "model": "qwen/qwen3-32b", + "choices": [ + { + "delta": { + "content": "Regular content", + "role": "assistant", + }, + "finish_reason": None, + "index": 0, + } + ], + } + + # Parse the chunk + parsed_chunk = handler.chunk_parser(groq_chunk) + + # Verify that content is present + assert parsed_chunk.choices[0].delta.content == "Regular content" + assert parsed_chunk.choices[0].delta.role == "assistant" + # Verify that reasoning_content is not set (it should be deleted by Delta.__init__) + assert not hasattr(parsed_chunk.choices[0].delta, "reasoning_content") + + def test_reasoning_with_tool_calls(self): + """ + Test that reasoning field is properly mapped even when tool_calls are present. + """ + handler = GroqChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + # Simulate a chunk with both reasoning and tool_calls + groq_chunk = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1769511767, + "model": "qwen/qwen3-32b", + "choices": [ + { + "delta": { + "reasoning": "Reasoning before tool call", + "tool_calls": [ + { + "index": 0, + "id": "call_123", + "function": {"name": "test_function", "arguments": "{}"}, + "type": "function", + } + ], + }, + "finish_reason": None, + "index": 0, + } + ], + } + + # Parse the chunk + parsed_chunk = handler.chunk_parser(groq_chunk) + + # Verify that reasoning was mapped to reasoning_content + assert parsed_chunk.choices[0].delta.reasoning_content == "Reasoning before tool call" + # Verify tool_calls are still present + assert parsed_chunk.choices[0].delta.tool_calls is not None + assert len(parsed_chunk.choices[0].delta.tool_calls) == 1 + assert parsed_chunk.choices[0].delta.tool_calls[0]["function"]["name"] == "test_function" 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_auth_checks.py b/tests/proxy_unit_tests/test_auth_checks.py index 24c005443a..05c6e4984a 100644 --- a/tests/proxy_unit_tests/test_auth_checks.py +++ b/tests/proxy_unit_tests/test_auth_checks.py @@ -584,7 +584,7 @@ async def test_get_fuzzy_user_object(): ) assert result == test_user mock_prisma.db.litellm_usertable.find_first.assert_called_with( - where={"user_email": "test@example.com"}, + where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}}, include={"organization_memberships": True}, ) @@ -612,7 +612,7 @@ async def test_get_fuzzy_user_object(): ) assert result == test_user mock_prisma.db.litellm_usertable.find_first.assert_called_with( - where={"user_email": "test@example.com"}, + where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}}, include={"organization_memberships": True}, ) 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/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 839d09b52d..d864a0d986 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -1117,7 +1117,10 @@ async def test_create_team_member_add(prisma_client, new_member_method): ) as mock_litellm_usertable, patch( "litellm.proxy.auth.auth_checks._get_team_object_from_user_api_key_cache", new=AsyncMock(return_value=team_obj), - ) as mock_team_obj: + ) as mock_team_obj, patch( + "litellm.proxy.proxy_server.prisma_client.get_data", + new=AsyncMock(return_value=[]), + ) as mock_get_data: mock_client = AsyncMock( return_value=LiteLLM_UserTable( @@ -1126,6 +1129,10 @@ async def test_create_team_member_add(prisma_client, new_member_method): ) mock_litellm_usertable.upsert = mock_client mock_litellm_usertable.find_many = AsyncMock(return_value=None) + # Mock find_first for user_email validation (returns None for new users) + mock_litellm_usertable.find_first = AsyncMock(return_value=None) + # Mock find_unique for user_id validation (returns None for new users) + mock_litellm_usertable.find_unique = AsyncMock(return_value=None) team_mock_client = AsyncMock() original_val = getattr( litellm.proxy.proxy_server.prisma_client.db, "litellm_teamtable" @@ -1299,7 +1306,10 @@ async def test_create_team_member_add_team_admin( ) as mock_litellm_usertable, patch( "litellm.proxy.auth.auth_checks._get_team_object_from_user_api_key_cache", new=AsyncMock(return_value=team_obj), - ) as mock_team_obj: + ) as mock_team_obj, patch( + "litellm.proxy.proxy_server.prisma_client.get_data", + new=AsyncMock(return_value=[]), + ) as mock_get_data: mock_client = AsyncMock( return_value=LiteLLM_UserTable( user_id="1234", max_budget=100, user_email="1234" @@ -1307,6 +1317,10 @@ async def test_create_team_member_add_team_admin( ) mock_litellm_usertable.upsert = mock_client mock_litellm_usertable.find_many = AsyncMock(return_value=None) + # Mock find_first for user_email validation (returns None for new users) + mock_litellm_usertable.find_first = AsyncMock(return_value=None) + # Mock find_unique for user_id validation (returns None for new users) + mock_litellm_usertable.find_unique = AsyncMock(return_value=None) team_mock_client = AsyncMock() original_val = getattr( 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 c0b863ef6e..a83bc1df1e 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) @@ -74,11 +110,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 @@ -91,43 +127,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", @@ -137,46 +183,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") @@ -195,26 +242,29 @@ 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") @@ -227,4 +277,4 @@ if __name__ == "__main__": 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!") \ No newline at end of file + 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_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index b15d75a414..9c2939b2da 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -66,6 +66,41 @@ class LargeImageClient: ) +class StreamingLargeImageClient: + """ + Client that streams a large image to test streaming download protection. + This simulates a huge file without actually creating it all in memory. + """ + + def __init__(self, size_mb=100, include_content_length=False): + self.size_mb = size_mb + self.include_content_length = include_content_length + + def get(self, url, follow_redirects=True): + size_bytes = int(self.size_mb * 1024 * 1024) + headers = {"Content-Type": "image/jpeg"} + if self.include_content_length: + headers["Content-Length"] = str(size_bytes) + + # Create a generator that yields chunks without creating the whole file in memory + def generate_chunks(total_size, chunk_size=8192): + bytes_sent = 0 + while bytes_sent < total_size: + chunk = b"x" * min(chunk_size, total_size - bytes_sent) + bytes_sent += len(chunk) + yield chunk + + # Create response with streaming content + response = Response( + status_code=200, + headers=headers, + request=Request("GET", url), + ) + # Mock the iter_bytes method to return our generator + response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size) + return response + + def test_image_exceeds_size_limit_with_content_length(monkeypatch): """ Test that images exceeding MAX_IMAGE_URL_DOWNLOAD_SIZE_MB are rejected when Content-Length header is present. @@ -83,6 +118,7 @@ def test_image_exceeds_size_limit_with_content_length(monkeypatch): def test_image_exceeds_size_limit_without_content_length(monkeypatch): """ Test that images exceeding MAX_IMAGE_URL_DOWNLOAD_SIZE_MB are rejected even without Content-Length header. + This uses the old non-streaming mock for backward compatibility. """ monkeypatch.setattr( litellm, "module_level_client", LargeImageClient(size_mb=100, include_content_length=False) @@ -94,6 +130,29 @@ def test_image_exceeds_size_limit_without_content_length(monkeypatch): assert "exceeds maximum allowed size" in str(excinfo.value) +def test_streaming_download_protects_against_huge_files(monkeypatch): + """ + Test that streaming download aborts early when file exceeds size limit, + preventing memory exhaustion from huge files (e.g., petabyte-sized files). + + This test verifies that the streaming implementation doesn't download the entire + file into memory before checking size. Instead, it should abort as soon as the + limit is exceeded during streaming. + """ + # Simulate a 1GB file - far larger than the 50MB default limit + client = StreamingLargeImageClient(size_mb=1024, include_content_length=False) + monkeypatch.setattr(litellm, "module_level_client", client) + + with pytest.raises(litellm.ImageFetchError) as excinfo: + convert_url_to_base64("https://example.com/huge-image.jpg") + + # Verify the error message shows it was caught during streaming + assert "exceeds maximum allowed size" in str(excinfo.value) + + # The error should be raised after downloading just slightly more than the limit + # not after downloading the full 1GB + + class SmallImageClient: """ Client that returns a small valid image. @@ -124,6 +183,26 @@ def test_image_within_size_limit(monkeypatch): assert result.startswith("data:image/jpeg;base64,") +def test_streaming_download_handles_petabyte_file(monkeypatch): + """ + Test that streaming download can handle extremely large file URLs (e.g., petabyte-sized) + without attempting to download the entire file or causing memory exhaustion. + + This simulates what happens if a malicious actor or misconfiguration provides + a URL to an extremely large file. + """ + # Simulate a 1 petabyte file (1,000,000 GB) + # Without streaming protection, this would cause OOM or hang indefinitely + client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False) + monkeypatch.setattr(litellm, "module_level_client", client) + + with pytest.raises(litellm.ImageFetchError) as excinfo: + convert_url_to_base64("https://example.com/petabyte-file.jpg") + + # Should fail fast without downloading anywhere near 1 petabyte + assert "exceeds maximum allowed size" in str(excinfo.value) + + def test_image_size_limit_disabled(monkeypatch): """ Test that setting MAX_IMAGE_URL_DOWNLOAD_SIZE_MB to 0 disables all image URL downloads. 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/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 1cb84d32c1..98b392a353 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -279,3 +279,188 @@ def test_output_format_with_no_schema(): # Content should remain as string (not converted to list) assert isinstance(last_user_message["content"], str) assert last_user_message["content"] == "Hello" + + +def test_advanced_tool_use_header_translation_for_opus_4_5(): + """ + Test that advanced-tool-use-2025-11-20 header is translated to Bedrock-specific headers + for Claude Opus 4.5. + + Regression test for: Claude Code sends advanced-tool-use header which needs to be + translated to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 for Bedrock + Invoke API on Claude Opus 4.5. + + Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html + """ + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + config = AmazonAnthropicClaudeMessagesConfig() + + messages = [ + {"role": "user", "content": "What's the weather like?"} + ] + + anthropic_messages_optional_request_params = { + "max_tokens": 100, + } + + # Simulate advanced-tool-use header from Claude Code + headers = { + "anthropic-beta": "advanced-tool-use-2025-11-20" + } + + # Test with Claude Opus 4.5 + result = config.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-5-20250514-v1:0", + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params={}, + headers=headers, + ) + + # Verify advanced-tool-use header was removed + assert "anthropic_beta" in result + beta_headers = result["anthropic_beta"] + assert "advanced-tool-use-2025-11-20" not in beta_headers, \ + "advanced-tool-use header should be removed for Bedrock" + + # Verify Bedrock-specific headers were added + assert "tool-search-tool-2025-10-19" in beta_headers, \ + "tool-search-tool-2025-10-19 should be added for Opus 4.5" + assert "tool-examples-2025-10-29" in beta_headers, \ + "tool-examples-2025-10-29 should be added for Opus 4.5" + + +def test_advanced_tool_use_header_filtered_for_non_opus_4_5(): + """ + Test that advanced-tool-use-2025-11-20 header is filtered out for non-Opus 4.5 models + without adding Bedrock-specific headers. + + The translation to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 should + only happen for Claude Opus 4.5. + """ + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + config = AmazonAnthropicClaudeMessagesConfig() + + messages = [ + {"role": "user", "content": "What's the weather like?"} + ] + + anthropic_messages_optional_request_params = { + "max_tokens": 100, + } + + # Simulate advanced-tool-use header from Claude Code + headers = { + "anthropic-beta": "advanced-tool-use-2025-11-20" + } + + # Test with Claude Sonnet 4.5 (not Opus 4.5) + result = config.transform_anthropic_messages_request( + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params={}, + headers=headers, + ) + + # Verify advanced-tool-use header was removed + beta_headers = result.get("anthropic_beta", []) + assert "advanced-tool-use-2025-11-20" not in beta_headers, \ + "advanced-tool-use header should be removed for Bedrock" + + # Verify Bedrock-specific headers were NOT added (only for Opus 4.5) + assert "tool-search-tool-2025-10-19" not in beta_headers, \ + "tool-search-tool should not be added for non-Opus 4.5 models" + assert "tool-examples-2025-10-29" not in beta_headers, \ + "tool-examples should not be added for non-Opus 4.5 models" + + +def test_advanced_tool_use_header_translation_with_multiple_beta_headers(): + """ + Test that advanced-tool-use header translation works correctly when multiple + beta headers are present. + """ + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + config = AmazonAnthropicClaudeMessagesConfig() + + messages = [ + {"role": "user", "content": "What's the weather like?"} + ] + + anthropic_messages_optional_request_params = { + "max_tokens": 100, + } + + # Multiple beta headers including advanced-tool-use + headers = { + "anthropic-beta": "claude-code-20250219,advanced-tool-use-2025-11-20,interleaved-thinking-2025-05-14" + } + + # Test with Claude Opus 4.5 + result = config.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-5-20250514-v1:0", + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params={}, + headers=headers, + ) + + beta_headers = result.get("anthropic_beta", []) + + # Verify advanced-tool-use was removed + assert "advanced-tool-use-2025-11-20" not in beta_headers + + # Verify Bedrock-specific headers were added + assert "tool-search-tool-2025-10-19" in beta_headers + assert "tool-examples-2025-10-29" in beta_headers + + # Verify other beta headers are preserved + assert "claude-code-20250219" in beta_headers + assert "interleaved-thinking-2025-05-14" in beta_headers + + +def test_opus_4_5_model_detection(): + """ + Test that the _is_claude_opus_4_5 method correctly identifies Opus 4.5 models + with various naming conventions. + """ + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + config = AmazonAnthropicClaudeMessagesConfig() + + # Test various Opus 4.5 naming patterns + opus_4_5_models = [ + "anthropic.claude-opus-4-5-20250514-v1:0", + "anthropic.claude-opus-4.5-20250514-v1:0", + "anthropic.claude-opus_4_5-20250514-v1:0", + "anthropic.claude-opus_4.5-20250514-v1:0", + "us.anthropic.claude-opus-4-5-20250514-v1:0", + "ANTHROPIC.CLAUDE-OPUS-4-5-20250514-V1:0", # Case insensitive + ] + + for model in opus_4_5_models: + assert config._is_claude_opus_4_5(model), \ + f"Should detect {model} as Opus 4.5" + + # Test non-Opus 4.5 models + non_opus_4_5_models = [ + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-opus-4-20250514-v1:0", # Opus 4, not 4.5 + "anthropic.claude-opus-4-1-20250514-v1:0", # Opus 4.1, not 4.5 + "anthropic.claude-haiku-4-5-20251001-v1:0", + ] + + for model in non_opus_4_5_models: + assert not config._is_claude_opus_4_5(model), \ + f"Should not detect {model} as Opus 4.5" diff --git a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py index b9324e4966..e7b6de29b6 100644 --- a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py +++ b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py @@ -390,3 +390,103 @@ class TestAnthropicBetaHeaderSupport: "anthropic_beta SHOULD be added for Anthropic models with cross-region prefix." ) assert "context-1m-2025-08-07" in additional_fields["anthropic_beta"] + + def test_messages_advanced_tool_use_translation_opus_4_5(self): + """Test that advanced-tool-use header is translated to Bedrock-specific headers for Opus 4.5. + + Regression test for: Claude Code sends advanced-tool-use-2025-11-20 header which needs + to be translated to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 for + Bedrock Invoke API on Claude Opus 4.5. + + Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html + """ + config = AmazonAnthropicClaudeMessagesConfig() + headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} + + result = config.transform_anthropic_messages_request( + model="us.anthropic.claude-opus-4-5-20250514-v1:0", + messages=[{"role": "user", "content": "Test"}], + anthropic_messages_optional_request_params={"max_tokens": 100}, + litellm_params={}, + headers=headers + ) + + assert "anthropic_beta" in result + beta_headers = result["anthropic_beta"] + + # advanced-tool-use should be removed + assert "advanced-tool-use-2025-11-20" not in beta_headers, ( + "advanced-tool-use-2025-11-20 should be removed for Bedrock Invoke API" + ) + + # Bedrock-specific headers should be added for Opus 4.5 + assert "tool-search-tool-2025-10-19" in beta_headers, ( + "tool-search-tool-2025-10-19 should be added for Opus 4.5" + ) + assert "tool-examples-2025-10-29" in beta_headers, ( + "tool-examples-2025-10-29 should be added for Opus 4.5" + ) + + def test_messages_advanced_tool_use_translation_sonnet_4_5(self): + """Test that advanced-tool-use header is translated to Bedrock-specific headers for Sonnet 4.5. + + Regression test for: Claude Code sends advanced-tool-use-2025-11-20 header which needs + to be translated to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 for + Bedrock Invoke API on Claude Sonnet 4.5. + + Ref: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool + """ + config = AmazonAnthropicClaudeMessagesConfig() + headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} + + result = config.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "Test"}], + anthropic_messages_optional_request_params={"max_tokens": 100}, + litellm_params={}, + headers=headers + ) + + assert "anthropic_beta" in result + beta_headers = result["anthropic_beta"] + + # advanced-tool-use should be removed + assert "advanced-tool-use-2025-11-20" not in beta_headers, ( + "advanced-tool-use-2025-11-20 should be removed for Bedrock Invoke API" + ) + + # Bedrock-specific headers should be added for Sonnet 4.5 + assert "tool-search-tool-2025-10-19" in beta_headers, ( + "tool-search-tool-2025-10-19 should be added for Sonnet 4.5" + ) + assert "tool-examples-2025-10-29" in beta_headers, ( + "tool-examples-2025-10-29 should be added for Sonnet 4.5" + ) + + def test_messages_advanced_tool_use_filtered_unsupported_model(self): + """Test that advanced-tool-use header is filtered out for models that don't support tool search. + + The translation to Bedrock-specific headers should only happen for models that + support tool search on Bedrock (Opus 4.5, Sonnet 4.5). + For other models, the advanced-tool-use header should just be removed. + """ + config = AmazonAnthropicClaudeMessagesConfig() + headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} + + # Test with Claude 3.5 Sonnet (does NOT support tool search on Bedrock) + result = config.transform_anthropic_messages_request( + model="us.anthropic.claude-3-5-sonnet-20241022-v2:0", + messages=[{"role": "user", "content": "Test"}], + anthropic_messages_optional_request_params={"max_tokens": 100}, + litellm_params={}, + headers=headers + ) + + beta_headers = result.get("anthropic_beta", []) + + # advanced-tool-use should be removed + assert "advanced-tool-use-2025-11-20" not in beta_headers + + # Bedrock-specific headers should NOT be added for unsupported models + assert "tool-search-tool-2025-10-19" not in beta_headers + assert "tool-examples-2025-10-29" not in beta_headers diff --git a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py b/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py index 3749a5a8ca..9b3b6aeaea 100644 --- a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py @@ -101,3 +101,50 @@ def test_hosted_vllm_supports_reasoning_effort(): drop_params=False, ) assert optional_params["reasoning_effort"] == "high" + + +def test_hosted_vllm_supports_thinking(): + """ + Test that hosted_vllm supports the 'thinking' parameter. + + Anthropic-style thinking is converted to OpenAI-style reasoning_effort + since vLLM is OpenAI-compatible. + + Related issue: https://github.com/BerriAI/litellm/issues/19761 + """ + config = HostedVLLMChatConfig() + supported_params = config.get_supported_openai_params( + model="hosted_vllm/GLM-4.6-FP8" + ) + assert "thinking" in supported_params + + # Test thinking with low budget_tokens -> "minimal" (for < 2000) + optional_params = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 1024}}, + optional_params={}, + model="hosted_vllm/GLM-4.6-FP8", + drop_params=False, + ) + assert "thinking" not in optional_params # thinking should NOT be passed + assert optional_params["reasoning_effort"] == "minimal" + + # Test thinking with high budget_tokens -> "high" + optional_params = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 15000}}, + optional_params={}, + model="hosted_vllm/GLM-4.6-FP8", + drop_params=False, + ) + assert optional_params["reasoning_effort"] == "high" + + # Test that existing reasoning_effort is not overwritten + optional_params = config.map_openai_params( + non_default_params={ + "thinking": {"type": "enabled", "budget_tokens": 15000}, + "reasoning_effort": "low", + }, + optional_params={}, + model="hosted_vllm/GLM-4.6-FP8", + drop_params=False, + ) + assert optional_params["reasoning_effort"] == "low" diff --git a/tests/test_litellm/llms/s3_vectors/__init__.py b/tests/test_litellm/llms/s3_vectors/__init__.py new file mode 100644 index 0000000000..d4b0c4d855 --- /dev/null +++ b/tests/test_litellm/llms/s3_vectors/__init__.py @@ -0,0 +1 @@ +# S3 Vectors tests diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py b/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py new file mode 100644 index 0000000000..231735c1de --- /dev/null +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py @@ -0,0 +1 @@ +# S3 Vectors vector store tests diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py new file mode 100644 index 0000000000..3a84da3154 --- /dev/null +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -0,0 +1,115 @@ +from unittest.mock import MagicMock, Mock + +import httpx +import pytest + +from litellm.llms.s3_vectors.vector_stores.transformation import ( + S3VectorsVectorStoreConfig, +) +from litellm.types.vector_stores import VectorStoreSearchResponse + + +class TestS3VectorsVectorStoreConfig: + def test_init(self): + """Test that S3VectorsVectorStoreConfig initializes correctly""" + config = S3VectorsVectorStoreConfig() + assert config is not None + + def test_get_supported_openai_params(self): + """Test that supported OpenAI params are returned""" + config = S3VectorsVectorStoreConfig() + params = config.get_supported_openai_params("test-model") + assert "max_num_results" in params + + def test_get_complete_url(self): + """Test URL generation for S3 Vectors""" + config = S3VectorsVectorStoreConfig() + litellm_params = {"aws_region_name": "us-west-2"} + url = config.get_complete_url(None, litellm_params) + assert url == "https://s3vectors.us-west-2.api.aws" + + def test_get_complete_url_missing_region(self): + """Test that missing region raises error""" + config = S3VectorsVectorStoreConfig() + litellm_params = {} + with pytest.raises(ValueError, match="aws_region_name is required"): + config.get_complete_url(None, litellm_params) + + @pytest.mark.skip(reason="Requires embedding API call, tested in integration tests") + def test_transform_search_request(self): + """Test search request transformation""" + # This test requires making an actual embedding API call + # It's better tested in integration tests + pass + + def test_transform_search_request_invalid_vector_store_id(self): + """Test that invalid vector_store_id format raises error""" + config = S3VectorsVectorStoreConfig() + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {} + + with pytest.raises( + ValueError, match="vector_store_id must be in format 'bucket_name:index_name'" + ): + config.transform_search_vector_store_request( + vector_store_id="invalid-format", + query="test query", + vector_store_search_optional_params={}, + api_base="https://s3vectors.us-west-2.api.aws", + litellm_logging_obj=mock_logging_obj, + litellm_params={}, + ) + + def test_transform_search_response(self): + """Test search response transformation""" + config = S3VectorsVectorStoreConfig() + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {"query": "test query"} + + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = { + "vectors": [ + { + "distance": 0.05, # S3 Vectors returns distance, not score + "metadata": { + "source_text": "This is test content", + "chunk_index": "0", + "filename": "test.pdf", + }, + }, + { + "distance": 0.15, + "metadata": { + "source_text": "More test content", + "chunk_index": "1", + }, + }, + ] + } + mock_response.status_code = 200 + mock_response.headers = {} + + result = config.transform_search_vector_store_response( + mock_response, mock_logging_obj + ) + + # VectorStoreSearchResponse is a TypedDict, so check structure instead of isinstance + assert result["object"] == "vector_store.search_results.page" + assert result["search_query"] == "test query" + assert len(result["data"]) == 2 + # Score should be 1 - distance (cosine similarity) + assert result["data"][0]["score"] == 0.95 # 1 - 0.05 + assert result["data"][0]["content"][0]["text"] == "This is test content" + assert result["data"][0]["filename"] == "test.pdf" + assert result["data"][1]["score"] == 0.85 # 1 - 0.15 + assert result["data"][1]["content"][0]["text"] == "More test content" + + def test_map_openai_params(self): + """Test OpenAI parameter mapping""" + config = S3VectorsVectorStoreConfig() + non_default_params = {"max_num_results": 5} + optional_params = {} + + result = config.map_openai_params(non_default_params, optional_params, False) + + assert result["maxResults"] == 5 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/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 3df0dc881e..af0591e5f5 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -28,6 +28,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, _can_object_call_vector_stores, + _get_fuzzy_user_object, _get_team_db_check, _virtual_key_max_budget_alert_check, _virtual_key_soft_budget_check, @@ -1331,3 +1332,42 @@ async def test_virtual_key_max_budget_alert_check_scenarios( assert ( alert_triggered == expect_alert ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, max_budget={max_budget}" + + +@pytest.mark.asyncio +async def test_get_fuzzy_user_object_case_insensitive_email(): + """Test that _get_fuzzy_user_object uses case-insensitive email lookup""" + # Setup mock Prisma client + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_usertable = MagicMock() + + # Mock user data with mixed case email + test_user = LiteLLM_UserTable( + user_id="test_123", + sso_user_id=None, + user_email="Test@Example.com", # Mixed case in DB + organization_memberships=[], + max_budget=None, + ) + + # Test: SSO ID not found, find by email with different casing + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=test_user) + + # Search with lowercase email (different from DB) + result = await _get_fuzzy_user_object( + prisma_client=mock_prisma, + sso_user_id=None, + user_email="test@example.com", # Lowercase search + ) + + # Verify user was found despite case difference + assert result == test_user + + # Verify the query used case-insensitive mode + mock_prisma.db.litellm_usertable.find_first.assert_called_once() + call_args = mock_prisma.db.litellm_usertable.find_first.call_args + assert call_args.kwargs["where"]["user_email"]["equals"] == "test@example.com" + assert call_args.kwargs["where"]["user_email"]["mode"] == "insensitive" + assert call_args.kwargs["include"] == {"organization_memberships": True} 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/google_endpoints/test_interactions_agent_param.py b/tests/test_litellm/proxy/google_endpoints/test_interactions_agent_param.py new file mode 100644 index 0000000000..1063f59afb --- /dev/null +++ b/tests/test_litellm/proxy/google_endpoints/test_interactions_agent_param.py @@ -0,0 +1,75 @@ +""" +Test for interactions endpoint agent parameter handling. + +Tests that the /v1beta/interactions endpoint correctly extracts +the `agent` parameter as a fallback when `model` is not provided. +""" + +import pytest + + +class TestInteractionsAgentParameter: + """Test agent parameter handling in interactions endpoint.""" + + def test_agent_parameter_fallback_logic(self): + """ + Test the core logic: model or agent extraction. + + This tests the fix in endpoints.py line ~267: + model=data.get("model") or data.get("agent") + """ + # Case 1: Only agent provided (Deep Research use case) + data = { + "agent": "deep-research-pro-preview-12-2025", + "input": "Research quantum computing", + "background": True, + } + model = data.get("model") or data.get("agent") + assert model == "deep-research-pro-preview-12-2025" + + # Case 2: Only model provided (normal use case) + data = { + "model": "gemini-2.5-flash", + "input": "Hello world", + } + model = data.get("model") or data.get("agent") + assert model == "gemini-2.5-flash" + + # Case 3: Both provided (model takes precedence) + data = { + "model": "gemini-2.5-flash", + "agent": "deep-research-pro-preview-12-2025", + "input": "Test", + } + model = data.get("model") or data.get("agent") + assert model == "gemini-2.5-flash" + + # Case 4: Neither provided + data = { + "input": "Test", + } + model = data.get("model") or data.get("agent") + assert model is None + + def test_route_type_in_skip_model_routing_list(self): + """ + Test that acreate_interaction is in the list of routes + that skip model-based routing. + + This tests the fix in route_llm_request.py. + """ + # The list of routes that skip model routing for interactions + skip_model_routing_routes = [ + "acreate_interaction", + "aget_interaction", + "adelete_interaction", + "acancel_interaction", + ] + + # acreate_interaction should be in the list (this is the fix) + assert "acreate_interaction" in skip_model_routing_routes + + # All interaction routes should be covered + assert "aget_interaction" in skip_model_routing_routes + assert "adelete_interaction" in skip_model_routing_routes + assert "acancel_interaction" in skip_model_routing_routes 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/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a97ed93a85..467ee3661d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _persist_deleted_team_records, _save_deleted_team_records, _transform_teams_to_deleted_records, + _validate_and_populate_member_user_info, delete_team, router, team_member_add_duplication_check, @@ -5466,3 +5467,122 @@ async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client) ) and mock_db_client.db.litellm_verificationtoken.find_many.called: # If it was called, that's unexpected for admin users assert False, "API keys should not be fetched for team admin users" + + +@pytest.mark.asyncio +async def test_validate_and_populate_member_user_info_both_provided_match(): + """ + Test _validate_and_populate_member_user_info when both user_email and user_id + are provided and they match the same user in the database. + """ + # Create member with both user_email and user_id + member = Member(user_email="test@example.com", user_id="user-123", role="user") + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock user object that matches both email and user_id + mock_user = MagicMock() + mock_user.user_id = "user-123" + mock_user.user_email = "test@example.com" + + # Mock get_data to return single user matching email + mock_prisma_client.get_data = AsyncMock(return_value=[mock_user]) + + # Call the function + result = await _validate_and_populate_member_user_info( + member=member, + prisma_client=mock_prisma_client, + ) + + # Verify result matches input (both already provided and match) + assert result.user_email == "test@example.com" + assert result.user_id == "user-123" + + # Verify get_data was called with correct parameters + mock_prisma_client.get_data.assert_called_once_with( + key_val={"user_email": "test@example.com"}, + table_name="user", + query_type="find_all", + ) + + +@pytest.mark.asyncio +async def test_validate_and_populate_member_user_info_only_email_provided(): + """ + Test _validate_and_populate_member_user_info when only user_email is provided. + Should populate user_id from database. + """ + # Create member with only user_email + member = Member(user_email="test@example.com", user_id=None, role="user") + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock user object from find_first + mock_user_find_first = MagicMock() + mock_user_find_first.user_id = "user-456" + mock_user_find_first.user_email = "test@example.com" + + # Mock find_first to return the user + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( + return_value=mock_user_find_first + ) + + # Mock get_data to return single user (no duplicates) + mock_prisma_client.get_data = AsyncMock(return_value=[mock_user_find_first]) + + # Call the function + result = await _validate_and_populate_member_user_info( + member=member, + prisma_client=mock_prisma_client, + ) + + # Verify user_id was populated + assert result.user_email == "test@example.com" + assert result.user_id == "user-456" + + # Verify find_first was called with correct parameters + mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with( + where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}} + ) + + # Verify get_data was called to check for duplicates + mock_prisma_client.get_data.assert_called_once_with( + key_val={"user_email": "test@example.com"}, + table_name="user", + query_type="find_all", + ) + + +@pytest.mark.asyncio +async def test_validate_and_populate_member_user_info_only_user_id_not_found(): + """ + Test _validate_and_populate_member_user_info when only user_id is provided + but the user doesn't exist in the database. Should allow it to pass with + user_email as None (will be upserted later). + """ + # Create member with only user_id + member = Member(user_email=None, user_id="nonexistent-user", role="user") + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock find_unique to return None (user not found) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + # Call the function - should NOT raise an exception + result = await _validate_and_populate_member_user_info( + member=member, + prisma_client=mock_prisma_client, + ) + + # Verify the result - should return member with user_id set and user_email as None + assert result.user_id == "nonexistent-user" + assert result.user_email is None + assert result.role == "user" + + # Verify find_unique was called with correct parameters + mock_prisma_client.db.litellm_usertable.find_unique.assert_called_once_with( + where={"user_id": "nonexistent-user"} + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index f983af2d0b..5e9078ea87 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -24,6 +24,7 @@ from litellm.proxy.management_endpoints.ui_sso import ( GoogleSSOHandler, MicrosoftSSOHandler, SSOAuthenticationHandler, + normalize_email, ) from litellm.types.proxy.management_endpoints.ui_sso import ( DefaultTeamSSOParams, @@ -667,6 +668,85 @@ def test_build_sso_user_update_data_without_role(): assert "user_role" not in update_data +def test_normalize_email(): + """ + Test that normalize_email correctly lowercases email addresses and handles edge cases. + """ + # Test with lowercase email + assert normalize_email("test@example.com") == "test@example.com" + + # Test with uppercase email + assert normalize_email("TEST@EXAMPLE.COM") == "test@example.com" + + # Test with mixed case email + assert normalize_email("Test.User@Example.COM") == "test.user@example.com" + + # Test with None + assert normalize_email(None) is None + + # Test with empty string + assert normalize_email("") == "" + + +def test_build_sso_user_update_data_normalizes_email(): + """ + Test that _build_sso_user_update_data normalizes email addresses to lowercase. + """ + from litellm.proxy.management_endpoints.types import CustomOpenID + from litellm.proxy.management_endpoints.ui_sso import _build_sso_user_update_data + + sso_result = CustomOpenID( + id="test-user-789", + email="Test.User@Example.COM", + display_name="Test User", + provider="microsoft", + team_ids=[], + user_role=None, + ) + + update_data = _build_sso_user_update_data( + result=sso_result, + user_email="Test.User@Example.COM", + user_id="test-user-789", + ) + + # Email should be normalized to lowercase + assert update_data["user_email"] == "test.user@example.com" + assert "user_role" not in update_data + + +def test_generic_response_convertor_normalizes_email(): + """ + Test that generic_response_convertor normalizes email addresses. + """ + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + mock_response = { + "preferred_username": "user123", + "email": "Test.User@Example.COM", + "sub": "Test User", + "first_name": "Test", + "last_name": "User", + "provider": "generic", + } + + # Mock JWT handler + mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler.get_team_ids_from_jwt.return_value = [] + + result = generic_response_convertor( + response=mock_response, + jwt_handler=mock_jwt_handler, + sso_jwt_handler=None, + role_mappings=None, + ) + + # Email should be normalized to lowercase + assert result.email == "test.user@example.com" + assert result.id == "user123" + assert result.display_name == "Test User" + + @pytest.mark.asyncio async def test_upsert_sso_user_updates_role_for_existing_user(): """ 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/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 2f6eccaf04..22a60f220e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -4013,6 +4013,279 @@ async def test_model_info_v2_filter_by_team_id(monkeypatch): app.dependency_overrides = original_overrides +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sort_by,sort_order,expected_order", + [ + # Test model_name sorting + ("model_name", "asc", ["a-model", "b-model", "z-model"]), + ("model_name", "desc", ["z-model", "b-model", "a-model"]), + # Test created_at sorting + ("created_at", "asc", ["old-model", "mid-model", "new-model"]), + ("created_at", "desc", ["new-model", "mid-model", "old-model"]), + # Test updated_at sorting + ("updated_at", "asc", ["old-updated", "mid-updated", "new-updated"]), + ("updated_at", "desc", ["new-updated", "mid-updated", "old-updated"]), + # Test costs sorting + ("costs", "asc", ["low-cost", "mid-cost", "high-cost"]), + ("costs", "desc", ["high-cost", "mid-cost", "low-cost"]), + # Test status sorting (False/config models come before True/db models in asc) + ("status", "asc", ["config-model-1", "config-model-2", "db-model"]), + ("status", "desc", ["db-model", "config-model-1", "config-model-2"]), + ], +) +async def test_model_info_v2_sorting(monkeypatch, sort_by, sort_order, expected_order): + """ + Test sorting functionality for /v2/model/info endpoint. + Tests all sortBy fields (model_name, created_at, updated_at, costs, status) + with both asc and desc sort orders. + """ + from datetime import datetime, timedelta + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + # Create base time for date comparisons + base_time = datetime(2024, 1, 1, 12, 0, 0) + + # Create mock models with different values for each sort field + mock_models = [] + + if sort_by == "model_name": + # Models with different names + mock_models = [ + { + "model_name": "z-model", + "litellm_params": {"model": "z-model"}, + "model_info": {"id": "z-model"}, + }, + { + "model_name": "a-model", + "litellm_params": {"model": "a-model"}, + "model_info": {"id": "a-model"}, + }, + { + "model_name": "b-model", + "litellm_params": {"model": "b-model"}, + "model_info": {"id": "b-model"}, + }, + ] + elif sort_by == "created_at": + # Models with different created_at timestamps + mock_models = [ + { + "model_name": "new-model", + "litellm_params": {"model": "new-model"}, + "model_info": { + "id": "new-model", + "created_at": (base_time + timedelta(days=3)).isoformat(), + }, + }, + { + "model_name": "old-model", + "litellm_params": {"model": "old-model"}, + "model_info": { + "id": "old-model", + "created_at": (base_time - timedelta(days=3)).isoformat(), + }, + }, + { + "model_name": "mid-model", + "litellm_params": {"model": "mid-model"}, + "model_info": { + "id": "mid-model", + "created_at": base_time.isoformat(), + }, + }, + ] + elif sort_by == "updated_at": + # Models with different updated_at timestamps + mock_models = [ + { + "model_name": "new-updated", + "litellm_params": {"model": "new-updated"}, + "model_info": { + "id": "new-updated", + "updated_at": (base_time + timedelta(days=3)).isoformat(), + }, + }, + { + "model_name": "old-updated", + "litellm_params": {"model": "old-updated"}, + "model_info": { + "id": "old-updated", + "updated_at": (base_time - timedelta(days=3)).isoformat(), + }, + }, + { + "model_name": "mid-updated", + "litellm_params": {"model": "mid-updated"}, + "model_info": { + "id": "mid-updated", + "updated_at": base_time.isoformat(), + }, + }, + ] + elif sort_by == "costs": + # Models with different costs (input_cost + output_cost) + mock_models = [ + { + "model_name": "high-cost", + "litellm_params": {"model": "high-cost"}, + "model_info": { + "id": "high-cost", + "input_cost_per_token": 0.00005, + "output_cost_per_token": 0.00015, + }, + }, + { + "model_name": "low-cost", + "litellm_params": {"model": "low-cost"}, + "model_info": { + "id": "low-cost", + "input_cost_per_token": 0.00001, + "output_cost_per_token": 0.00003, + }, + }, + { + "model_name": "mid-cost", + "litellm_params": {"model": "mid-cost"}, + "model_info": { + "id": "mid-cost", + "input_cost_per_token": 0.00003, + "output_cost_per_token": 0.00007, + }, + }, + ] + elif sort_by == "status": + # Models with different db_model status (False = config, True = db) + mock_models = [ + { + "model_name": "db-model", + "litellm_params": {"model": "db-model"}, + "model_info": {"id": "db-model", "db_model": True}, + }, + { + "model_name": "config-model-1", + "litellm_params": {"model": "config-model-1"}, + "model_info": {"id": "config-model-1", "db_model": False}, + }, + { + "model_name": "config-model-2", + "litellm_params": {"model": "config-model-2"}, + "model_info": {"id": "config-model-2", "db_model": False}, + }, + ] + + # Mock llm_router + mock_router = MagicMock() + mock_router.model_list = mock_models + + # Mock prisma_client + mock_prisma_client = MagicMock() + + # Mock proxy_config.get_config + mock_get_config = AsyncMock(return_value={}) + + # Mock user authentication + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.api_key = "test-key" + mock_user_api_key_dict.team_models = [] + mock_user_api_key_dict.models = [] + + # Apply monkeypatches + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + + # Override auth dependency + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_api_key_dict + + client = TestClient(app) + try: + # Test sorting with specified sortBy and sortOrder + response = client.get( + "/v2/model/info", params={"sortBy": sort_by, "sortOrder": sort_order} + ) + assert response.status_code == 200 + data = response.json() + assert len(data["data"]) == len(expected_order) + + # Verify models are in expected order + actual_order = [m["model_name"] for m in data["data"]] + assert actual_order == expected_order, ( + f"Sorting failed for sortBy={sort_by}, sortOrder={sort_order}. " + f"Expected: {expected_order}, Got: {actual_order}" + ) + + finally: + app.dependency_overrides = original_overrides + + +@pytest.mark.asyncio +async def test_model_info_v2_sorting_invalid_sort_order(monkeypatch): + """ + Test that invalid sortOrder values return a 400 error. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + # Create mock models + mock_models = [ + { + "model_name": "test-model", + "litellm_params": {"model": "test-model"}, + "model_info": {"id": "test-model"}, + } + ] + + # Mock llm_router + mock_router = MagicMock() + mock_router.model_list = mock_models + + # Mock prisma_client + mock_prisma_client = MagicMock() + + # Mock proxy_config.get_config + mock_get_config = AsyncMock(return_value={}) + + # Mock user authentication + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.api_key = "test-key" + mock_user_api_key_dict.team_models = [] + mock_user_api_key_dict.models = [] + + # Apply monkeypatches + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + + # Override auth dependency + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_api_key_dict + + client = TestClient(app) + try: + # Test invalid sortOrder + response = client.get( + "/v2/model/info", params={"sortBy": "model_name", "sortOrder": "invalid"} + ) + assert response.status_code == 400 + data = response.json() + assert "Invalid sortOrder" in data["detail"] + + finally: + app.dependency_overrides = original_overrides + + @pytest.mark.asyncio async def test_apply_search_filter_to_models(monkeypatch): """ 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/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/test_router_silent_experiment.py b/tests/test_litellm/test_router_silent_experiment.py new file mode 100644 index 0000000000..9b82cde13c --- /dev/null +++ b/tests/test_litellm/test_router_silent_experiment.py @@ -0,0 +1,157 @@ +import asyncio +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +from litellm.router import Router + + +@pytest.mark.asyncio +async def test_router_silent_experiment_acompletion(): + """ + Test that silent_model triggers a background acompletion call + and that the silent_model parameter is stripped from both calls. + """ + model_list = [ + { + "model_name": "primary-model", + "litellm_params": { + "model": "openai/gpt-3.5-turbo", + "api_key": "fake-key", + "silent_model": "silent-model", + }, + }, + { + "model_name": "silent-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "fake-key", + }, + }, + ] + + router = Router(model_list=model_list) + + # Mock litellm.acompletion + mock_acompletion = MagicMock() + # Create a future that resolves to a ModelResponse + mock_response = litellm.ModelResponse(choices=[{"message": {"content": "hello"}}]) + future = asyncio.Future() + future.set_result(mock_response) + mock_acompletion.return_value = future + + with patch("litellm.acompletion", mock_acompletion): + response = await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "hi"}], + ) + + assert response.choices[0].message.content == "hello" + + # Give the background task a moment to trigger (it's an asyncio task) + await asyncio.sleep(0.1) + + # Should have 2 calls: one for primary, one for silent + assert mock_acompletion.call_count == 2 + + # Check call arguments + call_args_list = mock_acompletion.call_args_list + + # Verify no silent_model in any call to litellm.acompletion + for call in call_args_list: + args, kwargs = call + assert "silent_model" not in kwargs + if "metadata" in kwargs: + # One call should have is_silent_experiment=True + pass + + # Find the silent call + silent_call = next( + ( + c + for c in call_args_list + if c[1].get("metadata", {}).get("is_silent_experiment") is True + ), + None, + ) + assert silent_call is not None + assert silent_call[1]["model"] == "openai/gpt-4" + + # Find the primary call + primary_call = next( + ( + c + for c in call_args_list + if not c[1].get("metadata", {}).get("is_silent_experiment") + ), + None, + ) + assert primary_call is not None + assert primary_call[1]["model"] == "openai/gpt-3.5-turbo" + + +def test_router_silent_experiment_completion(): + """ + Test that silent_model triggers a background completion call (sync) + and that the silent_model parameter is stripped. + """ + model_list = [ + { + "model_name": "primary-model", + "litellm_params": { + "model": "openai/gpt-3.5-turbo", + "api_key": "fake-key", + "silent_model": "silent-model", + }, + }, + { + "model_name": "silent-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "fake-key", + }, + }, + ] + + router = Router(model_list=model_list) + + # Mock litellm.completion + mock_completion = MagicMock() + mock_response = litellm.ModelResponse(choices=[{"message": {"content": "hello"}}]) + mock_completion.return_value = mock_response + + with patch("litellm.completion", mock_completion): + response = router.completion( + model="primary-model", + messages=[{"role": "user", "content": "hi"}], + ) + + assert response.choices[0].message.content == "hello" + + # The sync background call uses a thread pool. We might need to wait a bit. + import time + + time.sleep(0.5) + + # Should have 2 calls + assert mock_completion.call_count == 2 + + call_args_list = mock_completion.call_args_list + + # Verify no silent_model in any call + for call in call_args_list: + args, kwargs = call + assert "silent_model" not in kwargs + + # Find the silent call + silent_call = next( + ( + c + for c in call_args_list + if c[1].get("metadata", {}).get("is_silent_experiment") is True + ), + None, + ) + assert silent_call is not None + assert silent_call[1]["model"] == "openai/gpt-4" diff --git a/tests/vector_store_tests/rag/test_rag_s3_vectors.py b/tests/vector_store_tests/rag/test_rag_s3_vectors.py new file mode 100644 index 0000000000..cd8a362a7b --- /dev/null +++ b/tests/vector_store_tests/rag/test_rag_s3_vectors.py @@ -0,0 +1,107 @@ +""" +S3 Vectors RAG ingestion tests. + +Requires environment variables: +- AWS_ACCESS_KEY_ID +- AWS_SECRET_ACCESS_KEY +- AWS_REGION_NAME (optional, defaults to us-west-2) + +Optional: +- S3_VECTOR_BUCKET_NAME (optional, auto-generates if not set) +""" + +import os +import sys +from typing import Any, Dict, Optional + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +import litellm +from litellm.types.rag import RAGIngestOptions +from tests.vector_store_tests.rag.base_rag_tests import BaseRAGTest + + +class TestRAGS3Vectors(BaseRAGTest): + """Test RAG Ingest with AWS S3 Vectors.""" + + @pytest.fixture(autouse=True) + def check_env_vars(self): + """Check required environment variables before each test.""" + aws_key = os.environ.get("AWS_ACCESS_KEY_ID") + aws_secret = os.environ.get("AWS_SECRET_ACCESS_KEY") + + if not aws_key or not aws_secret: + pytest.skip("Skipping S3 Vectors test: AWS credentials required") + + def get_base_ingest_options(self) -> RAGIngestOptions: + """ + Return S3 Vectors-specific ingest options. + + Chunking is configured via chunking_strategy (unified interface). + Embeddings are generated using LiteLLM's embedding API. + """ + vector_bucket_name = os.environ.get( + "S3_VECTOR_BUCKET_NAME", "test-litellm-vectors" + ) + aws_region = os.environ.get("AWS_REGION_NAME", "us-west-2") + + return { + "chunking_strategy": { + "chunk_size": 512, + "chunk_overlap": 100, + }, + "embedding": { + "model": "text-embedding-3-small" # Can use any LiteLLM-supported model + }, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": vector_bucket_name, + "index_name": "test-index", + # dimension is auto-detected from embedding model (text-embedding-3-small = 1536) + "distance_metric": "cosine", + "non_filterable_metadata_keys": ["source_text"], + "aws_region_name": aws_region, + }, + } + + async def query_vector_store( + self, + vector_store_id: str, + query: str, + ) -> Optional[Dict[str, Any]]: + """Query S3 Vectors index.""" + try: + # Import the ingestion class to use its query method + from litellm.rag.ingestion.s3_vectors_ingestion import ( + S3VectorsRAGIngestion, + ) + except ImportError: + pytest.skip("S3 Vectors ingestion not available") + + vector_bucket_name = os.environ.get( + "S3_VECTOR_BUCKET_NAME", "test-litellm-vectors" + ) + aws_region = os.environ.get("AWS_REGION_NAME", "us-west-2") + + # Create ingestion instance to use query method + ingest_options = { + "embedding": {"model": "text-embedding-3-small"}, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": vector_bucket_name, + "aws_region_name": aws_region, + }, + } + + ingestion = S3VectorsRAGIngestion(ingest_options=ingest_options) + + # Query the index + results = await ingestion.query_vector_store( + vector_store_id=vector_store_id, + query=query, + top_k=5, + ) + + return results diff --git a/tests/vector_store_tests/test_s3_vectors_vector_store.py b/tests/vector_store_tests/test_s3_vectors_vector_store.py new file mode 100644 index 0000000000..a7a1568c1c --- /dev/null +++ b/tests/vector_store_tests/test_s3_vectors_vector_store.py @@ -0,0 +1,42 @@ +from base_vector_store_test import BaseVectorStoreTest +import os +import pytest + + +class TestS3VectorsVectorStore(BaseVectorStoreTest): + @pytest.fixture(autouse=True) + def check_env_vars(self): + """Check if required environment variables are set""" + required_vars = ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"] + missing_vars = [var for var in required_vars if not os.getenv(var)] + if missing_vars: + pytest.skip(f"Missing required environment variables: {', '.join(missing_vars)}") + + def get_base_request_args(self) -> dict: + """ + Must return the base request args for searching. + For S3 Vectors, vector_store_id should be in format: bucket_name:index_name + """ + return { + "custom_llm_provider": "s3_vectors", + "vector_store_id": os.getenv( + "S3_VECTORS_VECTOR_STORE_ID", "test-litellm-vectors:test-index" + ), + "query": "What is machine learning?", + "aws_region_name": os.getenv("AWS_REGION_NAME", "us-west-2"), + "aws_access_key_id": os.getenv("AWS_ACCESS_KEY_ID"), + "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"), + } + + def get_base_create_vector_store_args(self) -> dict: + """ + Vector store creation is not yet implemented for S3 Vectors. + This test will be skipped. + """ + return {} + + @pytest.mark.parametrize("sync_mode", [True, False]) + @pytest.mark.asyncio + async def test_basic_create_vector_store(self, sync_mode): + """S3 Vectors doesn't support vector store creation via this API yet""" + pytest.skip("Vector store creation not yet implemented for S3 Vectors") diff --git a/ui/litellm-dashboard/public/assets/logos/s3_vector.png b/ui/litellm-dashboard/public/assets/logos/s3_vector.png new file mode 100644 index 0000000000..15a1a456e1 Binary files /dev/null and b/ui/litellm-dashboard/public/assets/logos/s3_vector.png differ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts index 3db24cea9a..4985206092 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts @@ -104,6 +104,8 @@ describe("useModelsInfo", () => { 50, undefined, undefined, + undefined, + undefined, undefined ); expect(modelInfoCall).toHaveBeenCalledTimes(1); @@ -126,6 +128,8 @@ describe("useModelsInfo", () => { 25, undefined, undefined, + undefined, + undefined, undefined ); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index 13630f19bd..c57de675e0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -27,7 +27,7 @@ const modelHubKeys = createQueryKeys("modelHub"); const allProxyModelsKeys = createQueryKeys("allProxyModels"); const selectedTeamModelsKeys = createQueryKeys("selectedTeamModels"); -export const useModelsInfo = (page: number = 1, size: number = 50, search?: string, modelId?: string, teamId?: string) => { +export const useModelsInfo = (page: number = 1, size: number = 50, search?: string, modelId?: string, teamId?: string, sortBy?: string, sortOrder?: string) => { const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ queryKey: modelKeys.list({ @@ -39,9 +39,11 @@ export const useModelsInfo = (page: number = 1, size: number = 50, search?: stri ...(search && { search }), ...(modelId && { modelId }), ...(teamId && { teamId }), + ...(sortBy && { sortBy }), + ...(sortOrder && { sortOrder }), }, }), - queryFn: async () => await modelInfoCall(accessToken!, userId!, userRole!, page, size, search, modelId, teamId), + queryFn: async () => await modelInfoCall(accessToken!, userId!, userRole!, page, size, search, modelId, teamId, sortBy, sortOrder), enabled: Boolean(accessToken && userId && userRole), }); }; 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/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index 42189e391a..a5ec9d3578 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -2,11 +2,11 @@ import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap" import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { Team } from "@/components/key_team_helpers/key_list"; -import { ModelDataTable } from "@/components/model_dashboard/table"; +import { AllModelsDataTable } from "@/components/model_dashboard/all_models_table"; import { columns } from "@/components/molecules/models/columns"; import { getDisplayModelName } from "@/components/view_model/model_name_display"; import { InfoCircleOutlined } from "@ant-design/icons"; -import { PaginationState } from "@tanstack/react-table"; +import { PaginationState, SortingState } from "@tanstack/react-table"; import { Grid, Select, SelectItem, TabPanel, Text } from "@tremor/react"; import { Skeleton, Spin } from "antd"; import debounce from "lodash/debounce"; @@ -49,6 +49,7 @@ const AllModelsTab = ({ pageIndex: 0, pageSize: 50, }); + const [sorting, setSorting] = useState([]); // Debounce search input const debouncedUpdateSearch = useMemo( @@ -72,12 +73,33 @@ const AllModelsTab = ({ // Determine teamId to pass to the query - only pass if not "personal" const teamIdForQuery = currentTeam === "personal" ? undefined : currentTeam.team_id; + // Convert sorting state to sortBy and sortOrder for API + const sortBy = useMemo(() => { + if (sorting.length === 0) return undefined; + const sort = sorting[0]; + // Map column IDs to server-side field names + // The server expects field names like "model_name", "created_at", etc. + const columnIdToServerField: Record = { + input_cost: "costs", // Map input_cost column to "costs" for server-side sorting + model_info_db_model: "status", // Map model_info.db_model column to "status" for server-side sorting + }; + return columnIdToServerField[sort.id] || sort.id; + }, [sorting]); + + const sortOrder = useMemo(() => { + if (sorting.length === 0) return undefined; + const sort = sorting[0]; + return sort.desc ? "desc" : "asc"; + }, [sorting]); + const { data: rawModelData, isLoading: isLoadingModelsInfo } = useModelsInfo( currentPage, pageSize, debouncedSearch || undefined, undefined, - teamIdForQuery + teamIdForQuery, + sortBy, + sortOrder ); const isLoading = isLoadingModelsInfo || isLoadingModelCostMap; @@ -139,6 +161,7 @@ const AllModelsTab = ({ useEffect(() => { setPagination((prev: PaginationState) => ({ ...prev, pageIndex: 0 })); + setCurrentPage(1); }, [selectedModelGroup, selectedModelAccessGroupFilter]); // Reset pagination when team changes @@ -147,6 +170,12 @@ const AllModelsTab = ({ setPagination((prev: PaginationState) => ({ ...prev, pageIndex: 0 })); }, [teamIdForQuery]); + // Reset pagination when sorting changes + useEffect(() => { + setCurrentPage(1); + setPagination((prev: PaginationState) => ({ ...prev, pageIndex: 0 })); + }, [sorting]); + const resetFilters = () => { setModelNameSearch(""); setSelectedModelGroup("all"); @@ -155,6 +184,7 @@ const AllModelsTab = ({ setModelViewMode("current_team"); setCurrentPage(1); setPagination({ pageIndex: 0, pageSize: 50 }); + setSorting([]); }; return ( @@ -439,7 +469,7 @@ const AllModelsTab = ({
- }> - - {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 = ({
{ + className?: string; + } +} + +interface AllModelsDataTableProps { + data: TData[]; + columns: ColumnDef[]; + isLoading?: boolean; + sorting?: SortingState; + onSortingChange?: OnChangeFn; + pagination?: PaginationState; + onPaginationChange?: OnChangeFn; + enablePagination?: boolean; +} + +export function AllModelsDataTable({ + data = [], + columns, + isLoading = false, + sorting = [], + onSortingChange, + pagination, + onPaginationChange, + enablePagination = false, +}: AllModelsDataTableProps) { + const [columnResizeMode] = React.useState("onChange"); + const [columnSizing, setColumnSizing] = React.useState({}); + const [columnVisibility, setColumnVisibility] = React.useState({}); + + const tableInstance = useReactTable({ + data, + columns, + state: { + sorting, + columnSizing, + columnVisibility, + ...(enablePagination && pagination ? { pagination } : {}), + }, + columnResizeMode, + onSortingChange: onSortingChange, + onColumnSizingChange: setColumnSizing, + onColumnVisibilityChange: setColumnVisibility, + ...(enablePagination && onPaginationChange ? { onPaginationChange } : {}), + getCoreRowModel: getCoreRowModel(), + // NO getSortedRowModel - sorting is handled server-side + ...(enablePagination ? { getPaginationRowModel: getPaginationRowModel() } : {}), + enableSorting: true, + enableColumnResizing: true, + manualSorting: true, // Enable manual sorting for server-side sorting + defaultColumn: { + minSize: 40, + maxSize: 500, + }, + }); + + const getHeaderText = (header: any): string => { + if (typeof header === "string") { + return header; + } + if (typeof header === "function") { + const headerElement = header(); + if (headerElement && headerElement.props && headerElement.props.children) { + const children = headerElement.props.children; + if (typeof children === "string") { + return children; + } + if (children.props && children.props.children) { + return children.props.children; + } + } + } + return ""; + }; + + return ( +
+
+
+ + + {tableInstance.getHeaderGroups().map((headerGroup) => ( + + {headerGroup.headers.map((header) => ( + +
+
+ {header.isPlaceholder + ? null + : flexRender(header.column.columnDef.header, header.getContext())} +
+ {header.id !== "actions" && header.column.getCanSort() && ( +
+ {header.column.getIsSorted() ? ( + { + asc: , + desc: , + }[header.column.getIsSorted() as string] + ) : ( + + )} +
+ )} +
+ {header.column.getCanResize() && ( +
+ )} + + ))} + + ))} + + + {isLoading ? ( + + +
+

🚅 Loading models...

+
+ + + ) : tableInstance.getRowModel().rows.length > 0 ? ( + tableInstance.getRowModel().rows.map((row) => ( + + {row.getVisibleCells().map((cell) => ( + + {flexRender(cell.column.columnDef.cell, cell.getContext())} + + ))} + + )) + ) : ( + + +
+

No models found

+
+ + + )} + +
+
+
+
+ ); +} diff --git a/ui/litellm-dashboard/src/components/molecules/models/columns.tsx b/ui/litellm-dashboard/src/components/molecules/models/columns.tsx index 43813852f8..1fc08e502a 100644 --- a/ui/litellm-dashboard/src/components/molecules/models/columns.tsx +++ b/ui/litellm-dashboard/src/components/molecules/models/columns.tsx @@ -17,297 +17,300 @@ export const columns = ( expandedRows: Set, setExpandedRows: (expandedRows: Set) => void, ): ColumnDef[] => [ - { - header: () => Model ID, - accessorKey: "model_info.id", - cell: ({ row }) => { - const model = row.original; - return ( - -
setSelectedModelId(model.model_info.id)} - > - {model.model_info.id} -
- - ); - }, - }, - { - header: () => Model Information, - accessorKey: "model_name", - size: 250, // Fixed column width - cell: ({ row }) => { - const model = row.original; - const displayName = getDisplayModelName(row.original) || "-"; - const tooltipContent = ( -
-
- Provider: {model.provider || "-"} -
-
- Public Model Name: {displayName} -
-
- LiteLLM Model Name: {model.litellm_model_name || "-"} -
-
- ); - - return ( - -
- {/* Provider Icon */} -
- {model.provider ? ( - - ) : ( -
-
- )} -
- - {/* Model Names Container */} -
- {/* Public Model Name */} -
{displayName}
- {/* LiteLLM Model Name */} -
- {model.litellm_model_name || "-"} -
-
-
- - ); - }, - }, - { - header: () => Credentials, - accessorKey: "litellm_credential_name", - size: 180, // Fixed column width - cell: ({ row }) => { - const model = row.original; - const credentialName = model.litellm_params?.litellm_credential_name; - - return credentialName ? ( - -
- - - {credentialName} - -
- - ) : ( -
- - No credentials -
- ); - }, - }, - { - header: () => Created By, - accessorKey: "model_info.created_by", - sortingFn: "datetime", - size: 160, // Fixed column width - cell: ({ row }) => { - const model = row.original; - const isConfigModel = !model.model_info?.db_model; - const createdBy = model.model_info.created_by; - const createdAt = model.model_info.created_at ? new Date(model.model_info.created_at).toLocaleDateString() : null; - - return ( -
- {/* Created By - Primary */} -
- {isConfigModel ? "Defined in config" : createdBy || "Unknown"} -
- {/* Created At - Secondary */} -
- {isConfigModel ? "-" : createdAt || "Unknown date"} -
-
- ); - }, - }, - { - header: () => Updated At, - accessorKey: "model_info.updated_at", - sortingFn: "datetime", - cell: ({ row }) => { - const model = row.original; - return ( - - {model.model_info.updated_at ? new Date(model.model_info.updated_at).toLocaleDateString() : "-"} - - ); - }, - }, - { - header: () => Costs, - accessorKey: "input_cost", - size: 120, // Fixed column width - cell: ({ row }) => { - const model = row.original; - const inputCost = model.input_cost; - const outputCost = model.output_cost; - - // If both costs are missing or undefined, show "-" - if (!inputCost && !outputCost) { + { + header: () => Model ID, + accessorKey: "model_info.id", + enableSorting: false, + cell: ({ row }) => { + const model = row.original; return ( -
- - + +
setSelectedModelId(model.model_info.id)} + > + {model.model_info.id} +
+ + ); + }, + }, + { + header: () => Model Information, + accessorKey: "model_name", + size: 250, // Fixed column width + cell: ({ row }) => { + const model = row.original; + const displayName = getDisplayModelName(row.original) || "-"; + const tooltipContent = ( +
+
+ Provider: {model.provider || "-"} +
+
+ Public Model Name: {displayName} +
+
+ LiteLLM Model Name: {model.litellm_model_name || "-"} +
); - } - return ( - -
- {/* Input Cost - Primary */} - {inputCost &&
In: ${inputCost}
} - {/* Output Cost - Secondary */} - {outputCost &&
Out: ${outputCost}
} -
- - ); - }, - }, - { - header: () => Team ID, - accessorKey: "model_info.team_id", - cell: ({ row }) => { - const model = row.original; - return model.model_info.team_id ? ( -
- - + return ( + +
+ {/* Provider Icon */} +
+ {model.provider ? ( + + ) : ( +
-
+ )} +
+ + {/* Model Names Container */} +
+ {/* Public Model Name */} +
{displayName}
+ {/* LiteLLM Model Name */} +
+ {model.litellm_model_name || "-"} +
+
+
-
- ) : ( - "-" - ); + ); + }, }, - }, - { - header: () => Model Access Group, - accessorKey: "model_info.model_access_group", - enableSorting: false, - cell: ({ row }) => { - const model = row.original; - const accessGroups = model.model_info.access_groups; + { + header: () => Credentials, + accessorKey: "litellm_credential_name", + enableSorting: false, + size: 180, // Fixed column width + cell: ({ row }) => { + const model = row.original; + const credentialName = model.litellm_params?.litellm_credential_name; - if (!accessGroups || accessGroups.length === 0) { - return "-"; - } + return credentialName ? ( + +
+ + + {credentialName} + +
+ + ) : ( +
+ + No credentials +
+ ); + }, + }, + { + header: () => Created By, + accessorKey: "model_info.created_by", + sortingFn: "datetime", + size: 160, // Fixed column width + cell: ({ row }) => { + const model = row.original; + const isConfigModel = !model.model_info?.db_model; + const createdBy = model.model_info.created_by; + const createdAt = model.model_info.created_at ? new Date(model.model_info.created_at).toLocaleDateString() : null; - const modelId = model.model_info.id; - const isExpanded = expandedRows.has(modelId); - const shouldShowExpandButton = accessGroups.length > 1; - - const toggleExpanded = () => { - const newExpanded = new Set(expandedRows); - if (isExpanded) { - newExpanded.delete(modelId); - } else { - newExpanded.add(modelId); - } - setExpandedRows(newExpanded); - }; - - return ( -
- - {accessGroups[0]} - - - {(isExpanded || (!shouldShowExpandButton && accessGroups.length === 2)) && - accessGroups.slice(1).map((group: string, index: number) => ( - - {group} - - ))} - - {shouldShowExpandButton && ( - - )} -
- ); + {isConfigModel ? "Defined in config" : createdBy || "Unknown"} +
+ {/* Created At - Secondary */} +
+ {isConfigModel ? "-" : createdAt || "Unknown date"} +
+
+ ); + }, }, - }, - { - header: () => Status, - accessorKey: "model_info.db_model", - cell: ({ row }) => { - const model = row.original; - return ( -
Updated At, + accessorKey: "model_info.updated_at", + sortingFn: "datetime", + cell: ({ row }) => { + const model = row.original; + return ( + + {model.model_info.updated_at ? new Date(model.model_info.updated_at).toLocaleDateString() : "-"} + + ); + }, + }, + { + header: () => Costs, + accessorKey: "input_cost", + size: 120, // Fixed column width + cell: ({ row }) => { + const model = row.original; + const inputCost = model.input_cost; + const outputCost = model.output_cost; + + // If both costs are missing or undefined, show "-" + if (!inputCost && !outputCost) { + return ( +
+ - +
+ ); + } + + return ( + +
+ {/* Input Cost - Primary */} + {inputCost &&
In: ${inputCost}
} + {/* Output Cost - Secondary */} + {outputCost &&
Out: ${outputCost}
} +
+ + ); + }, + }, + { + header: () => Team ID, + accessorKey: "model_info.team_id", + enableSorting: false, + cell: ({ row }) => { + const model = row.original; + return model.model_info.team_id ? ( +
+ + + +
+ ) : ( + "-" + ); + }, + }, + { + header: () => Model Access Group, + accessorKey: "model_info.model_access_group", + enableSorting: false, + cell: ({ row }) => { + const model = row.original; + const accessGroups = model.model_info.access_groups; + + if (!accessGroups || accessGroups.length === 0) { + return "-"; + } + + const modelId = model.model_info.id; + const isExpanded = expandedRows.has(modelId); + const shouldShowExpandButton = accessGroups.length > 1; + + const toggleExpanded = () => { + const newExpanded = new Set(expandedRows); + if (isExpanded) { + newExpanded.delete(modelId); + } else { + newExpanded.add(modelId); + } + setExpandedRows(newExpanded); + }; + + return ( +
+ + {accessGroups[0]} + + + {(isExpanded || (!shouldShowExpandButton && accessGroups.length === 2)) && + accessGroups.slice(1).map((group: string, index: number) => ( + + {group} + + ))} + + {shouldShowExpandButton && ( + + )} +
+ ); + }, + }, + { + header: () => Status, + accessorKey: "model_info.db_model", + cell: ({ row }) => { + const model = row.original; + return ( +
- {model.model_info.db_model ? "DB Model" : "Config Model"} -
- ); + > + {model.model_info.db_model ? "DB Model" : "Config Model"} +
+ ); + }, }, - }, - { - id: "actions", - header: () => Actions, - cell: ({ row }) => { - const model = row.original; - const canEditModel = userRole === "Admin" || model.model_info?.created_by === userID; - const isConfigModel = !model.model_info?.db_model; - return ( -
- {isConfigModel ? ( - - - - ) : ( - - { - if (canEditModel) { - setSelectedModelId(model.model_info.id); - } - }} - className={!canEditModel ? "opacity-50 cursor-not-allowed" : "cursor-pointer hover:text-red-600"} - /> - - )} -
- ); + { + id: "actions", + header: () => Actions, + cell: ({ row }) => { + const model = row.original; + const canEditModel = userRole === "Admin" || model.model_info?.created_by === userID; + const isConfigModel = !model.model_info?.db_model; + return ( +
+ {isConfigModel ? ( + + + + ) : ( + + { + if (canEditModel) { + setSelectedModelId(model.model_info.id); + } + }} + className={!canEditModel ? "opacity-50 cursor-not-allowed" : "cursor-pointer hover:text-red-600"} + /> + + )} +
+ ); + }, }, - }, -]; + ]; diff --git a/ui/litellm-dashboard/src/components/navbar.test.tsx b/ui/litellm-dashboard/src/components/navbar.test.tsx index 5d3c7254ff..9fa32cf9cb 100644 --- a/ui/litellm-dashboard/src/components/navbar.test.tsx +++ b/ui/litellm-dashboard/src/components/navbar.test.tsx @@ -16,6 +16,7 @@ vi.mock("@/utils/proxyUtils", () => ({ 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..8e4fda0ba5 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" + /> +
), @@ -191,7 +219,7 @@ const Navbar: React.FC = ({ style={{ animationDuration: "2s" }} title="Happy Holidays!" > - 🎄 + ❄️ @@ -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={} + />} { +export const modelInfoCall = async (accessToken: string, userID: string, userRole: string, page: number = 1, size: number = 50, search?: string, modelId?: string, teamId?: string, sortBy?: string, sortOrder?: string) => { /** * Get all models on proxy */ try { - console.log("modelInfoCall:", accessToken, userID, userRole, page, size, search, modelId, teamId); + console.log("modelInfoCall:", accessToken, userID, userRole, page, size, search, modelId, teamId, sortBy, sortOrder); let url = proxyBaseUrl ? `${proxyBaseUrl}/v2/model/info` : `/v2/model/info`; const params = new URLSearchParams(); params.append("include_team_models", "true"); @@ -2027,6 +2027,12 @@ export const modelInfoCall = async (accessToken: string, userID: string, userRol if (teamId && teamId.trim()) { params.append("teamId", teamId.trim()); } + if (sortBy && sortBy.trim()) { + params.append("sortBy", sortBy.trim()); + } + if (sortOrder && sortOrder.trim()) { + params.append("sortOrder", sortOrder.trim()); + } if (params.toString()) { url += `?${params.toString()}`; } @@ -6948,6 +6954,64 @@ export const vectorStoreUpdateCall = async (accessToken: string, formValues: Rec } }; +export const ragIngestCall = async ( + accessToken: string, + file: File, + customLlmProvider: string, + vectorStoreId?: string, + vectorStoreName?: string, + vectorStoreDescription?: string, + providerSpecificParams?: Record +): Promise => { + 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 }), + ...(providerSpecificParams && providerSpecificParams), + }, + }, + }; + + // 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 + > + + + + {/* S3 Vectors Configuration */} + {selectedProvider === "s3_vectors" && ( + + )} + + {/* Other Provider-specific fields */} + {selectedProvider !== "s3_vectors" && + getProviderSpecificFields(selectedProvider).map((field: VectorStoreFieldConfig) => { + if (field.type === "select") { + // For embedding model selection, we'd need to fetch available models + // For now, provide a text input as fallback + return ( + + {field.label}{" "} + + + + + } + required={field.required} + > + + setProviderParams((prev) => ({ ...prev, [field.name]: e.target.value })) + } + placeholder={field.placeholder} + size="large" + className="rounded-md" + /> + + ); + } + + return ( + + {field.label}{" "} + + + + + } + required={field.required} + > + + setProviderParams((prev) => ({ ...prev, [field.name]: e.target.value })) + } + placeholder={field.placeholder} + size="large" + className="rounded-md" + /> + + ); + })} + + +
+ +
+ + + + {/* 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..507ad0537d --- /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(/1.95 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/S3VectorsConfig.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.test.tsx new file mode 100644 index 0000000000..28f885a9dd --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.test.tsx @@ -0,0 +1,203 @@ +import { render, screen, fireEvent, waitFor, act } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import S3VectorsConfig from "./S3VectorsConfig"; +import * as fetchModels from "../playground/llm_calls/fetch_models"; + +// Mock fetchAvailableModels +vi.mock("../playground/llm_calls/fetch_models", () => ({ + fetchAvailableModels: vi.fn(), +})); + +describe("S3VectorsConfig", () => { + const mockOnParamsChange = vi.fn(); + const defaultProps = { + accessToken: "test-token", + providerParams: {}, + onParamsChange: mockOnParamsChange, + }; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render the component successfully", () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + render(); + + expect(screen.getByText("AWS S3 Vectors Setup")).toBeInTheDocument(); + expect(screen.getByText("Vector Bucket Name")).toBeInTheDocument(); + expect(screen.getByText("Index Name")).toBeInTheDocument(); + expect(screen.getByText("AWS Region")).toBeInTheDocument(); + expect(screen.getByText("Embedding Model")).toBeInTheDocument(); + }); + + it("should display setup instructions", () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + render(); + + expect( + screen.getByText(/AWS S3 Vectors allows you to store and query vector embeddings directly in S3/) + ).toBeInTheDocument(); + expect(screen.getByText(/Vector buckets and indexes will be automatically created/)).toBeInTheDocument(); + expect(screen.getByText(/Vector dimensions are auto-detected/)).toBeInTheDocument(); + }); + + it("should fetch embedding models on mount", async () => { + const mockModels = [ + { model_group: "text-embedding-3-small", mode: "embedding" }, + { model_group: "text-embedding-3-large", mode: "embedding" }, + { model_group: "gpt-4", mode: "chat" }, + ]; + + const fetchSpy = vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue(mockModels); + + render(); + + await waitFor(() => { + expect(fetchSpy).toHaveBeenCalledWith("test-token"); + }); + }); + + it("should filter and display only embedding models", async () => { + const mockModels = [ + { model_group: "text-embedding-3-small", mode: "embedding" }, + { model_group: "text-embedding-3-large", mode: "embedding" }, + { model_group: "gpt-4", mode: "chat" }, + { model_group: "gpt-3.5-turbo", mode: "chat" }, + ]; + + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue(mockModels); + + render(); + + // Wait for models to load + await waitFor(() => { + expect(fetchModels.fetchAvailableModels).toHaveBeenCalled(); + }); + + // The component should filter to only embedding models internally + // We can verify this by checking the component loaded successfully + expect(screen.getByText("Embedding Model")).toBeInTheDocument(); + }); + + it("should call onParamsChange when vector bucket name changes", async () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + render(); + + const bucketInput = screen.getByPlaceholderText("my-vector-bucket (min 3 chars)"); + + await act(async () => { + fireEvent.change(bucketInput, { target: { value: "test-bucket" } }); + }); + + expect(mockOnParamsChange).toHaveBeenCalledWith({ + vector_bucket_name: "test-bucket", + }); + }); + + it("should call onParamsChange when AWS region changes", async () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + render(); + + const regionInput = screen.getByPlaceholderText("us-west-2"); + + await act(async () => { + fireEvent.change(regionInput, { target: { value: "us-east-1" } }); + }); + + expect(mockOnParamsChange).toHaveBeenCalledWith({ + aws_region_name: "us-east-1", + }); + }); + + it("should call onParamsChange when embedding model is selected", async () => { + const mockModels = [ + { model_group: "text-embedding-3-small", mode: "embedding" }, + { model_group: "text-embedding-3-large", mode: "embedding" }, + ]; + + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue(mockModels); + + render(); + + await waitFor(() => { + expect(fetchModels.fetchAvailableModels).toHaveBeenCalled(); + }); + + // Find the Select component and trigger change directly + const selectElement = screen.getByRole("combobox"); + + await act(async () => { + // Simulate selecting a value by firing the change event + fireEvent.change(selectElement, { target: { value: "text-embedding-3-small" } }); + }); + + // The component should handle the selection + expect(screen.getByText("Embedding Model")).toBeInTheDocument(); + }); + + it("should preserve existing params when updating a field", async () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + const existingParams = { + vector_bucket_name: "existing-bucket", + aws_region_name: "us-west-2", + }; + + render(); + + const indexInput = screen.getByPlaceholderText("my-vector-index (optional, min 3 chars)"); + + await act(async () => { + fireEvent.change(indexInput, { target: { value: "my-index" } }); + }); + + expect(mockOnParamsChange).toHaveBeenCalledWith({ + vector_bucket_name: "existing-bucket", + aws_region_name: "us-west-2", + index_name: "my-index", + }); + }); + + it("should display existing param values", () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + const existingParams = { + vector_bucket_name: "my-bucket", + index_name: "my-index", + aws_region_name: "eu-west-1", + embedding_model: "text-embedding-3-small", + }; + + render(); + + expect(screen.getByDisplayValue("my-bucket")).toBeInTheDocument(); + expect(screen.getByDisplayValue("my-index")).toBeInTheDocument(); + expect(screen.getByDisplayValue("eu-west-1")).toBeInTheDocument(); + }); + + it("should handle model fetch error gracefully", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + vi.spyOn(fetchModels, "fetchAvailableModels").mockRejectedValue(new Error("Failed to fetch models")); + + render(); + + await waitFor(() => { + expect(consoleErrorSpy).toHaveBeenCalledWith("Error fetching embedding models:", expect.any(Error)); + }); + + consoleErrorSpy.mockRestore(); + }); + + it("should not fetch models if accessToken is null", () => { + const fetchSpy = vi.spyOn(fetchModels, "fetchAvailableModels"); + + render(); + + expect(fetchSpy).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx new file mode 100644 index 0000000000..bf71492312 --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx @@ -0,0 +1,192 @@ +import React, { useState, useEffect } from "react"; +import { Alert, Form, Input, Select, Tooltip } from "antd"; +import { InfoCircleOutlined } from "@ant-design/icons"; +import { fetchAvailableModels, ModelGroup } from "../playground/llm_calls/fetch_models"; + +interface S3VectorsConfigProps { + accessToken: string | null; + providerParams: Record; + onParamsChange: (params: Record) => void; +} + +const S3VectorsConfig: React.FC = ({ + accessToken, + providerParams, + onParamsChange, +}) => { + const [embeddingModels, setEmbeddingModels] = useState([]); + const [isLoadingModels, setIsLoadingModels] = useState(false); + + useEffect(() => { + if (!accessToken) return; + + const loadModels = async () => { + setIsLoadingModels(true); + try { + const models = await fetchAvailableModels(accessToken); + // Filter for embedding models only + const embeddingOnly = models.filter((model) => model.mode === "embedding"); + setEmbeddingModels(embeddingOnly); + } catch (error) { + console.error("Error fetching embedding models:", error); + } finally { + setIsLoadingModels(false); + } + }; + + loadModels(); + }, [accessToken]); + + const handleFieldChange = (fieldName: string, value: string) => { + onParamsChange({ + ...providerParams, + [fieldName]: value, + }); + }; + + return ( + <> + {/* S3 Vectors Setup Instructions */} + +

AWS S3 Vectors allows you to store and query vector embeddings directly in S3:

+
    +
  • Vector buckets and indexes will be automatically created if they don't exist
    • +
  • Vector dimensions are auto-detected from your selected embedding model
    • +
  • Ensure your AWS credentials have permissions for S3 Vectors operations
    • +
  • + Learn more:{" "} + + AWS S3 Vectors Documentation + +
    • +
+ + } + type="info" + showIcon + style={{ marginBottom: "16px" }} + /> + + {/* Vector Bucket Name */} + + Vector Bucket Name{" "} + + + + + } + required + validateStatus={ + providerParams.vector_bucket_name && providerParams.vector_bucket_name.length < 3 + ? "error" + : undefined + } + help={ + providerParams.vector_bucket_name && providerParams.vector_bucket_name.length < 3 + ? "Bucket name must be at least 3 characters" + : undefined + } + > + handleFieldChange("vector_bucket_name", e.target.value)} + placeholder="my-vector-bucket (min 3 chars)" + size="large" + className="rounded-md" + /> + + + {/* Index Name (Optional) */} + + Index Name{" "} + + + + + } + validateStatus={ + providerParams.index_name && providerParams.index_name.length > 0 && providerParams.index_name.length < 3 + ? "error" + : undefined + } + help={ + providerParams.index_name && providerParams.index_name.length > 0 && providerParams.index_name.length < 3 + ? "Index name must be at least 3 characters if provided" + : undefined + } + > + handleFieldChange("index_name", e.target.value)} + placeholder="my-vector-index (optional, min 3 chars)" + size="large" + className="rounded-md" + /> + + + {/* AWS Region */} + + AWS Region{" "} + + + + + } + required + > + handleFieldChange("aws_region_name", e.target.value)} + placeholder="us-west-2" + size="large" + className="rounded-md" + /> + + + {/* Embedding Model */} + + Embedding Model{" "} + + + + + } + required + > + + {vectorStores.map((vs) => ( + +
+ {vs.vector_store_name || vs.vector_store_id} + {vs.vector_store_name && ( + {vs.vector_store_id} + )} +
+ + ))} + + + + + {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..ac79012250 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(5); // One for missing name, one for missing description, three for missing files (one per store) }); 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(5); // One for missing name, one for missing description, three for missing files (one per store) }); 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/vector_store_providers.tsx b/ui/litellm-dashboard/src/components/vector_store_providers.tsx index 3dcf8eed87..9efda3a5a7 100644 --- a/ui/litellm-dashboard/src/components/vector_store_providers.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_providers.tsx @@ -1,5 +1,6 @@ export enum VectorStoreProviders { Bedrock = "Amazon Bedrock", + S3Vectors = "Amazon S3 Vectors", PgVector = "PostgreSQL pgvector (LiteLLM Connector)", VertexRagEngine = "Vertex AI RAG Engine", OpenAI = "OpenAI", @@ -14,6 +15,7 @@ export const vectorStoreProviderMap: Record = { OpenAI: "openai", Azure: "azure", Milvus: "milvus", + S3Vectors: "s3_vectors", }; const asset_logos_folder = "../ui/assets/logos/"; @@ -25,6 +27,7 @@ export const vectorStoreProviderLogoMap: Record = { [VectorStoreProviders.OpenAI]: `${asset_logos_folder}openai_small.svg`, [VectorStoreProviders.Azure]: `${asset_logos_folder}microsoft_azure.svg`, [VectorStoreProviders.Milvus]: `${asset_logos_folder}milvus.svg`, + [VectorStoreProviders.S3Vectors]: `${asset_logos_folder}s3_vector.png`, }; // Define field types for provider-specific configurations @@ -114,6 +117,40 @@ export const vectorStoreProviderFields: Record type: "select", }, ], + s3_vectors: [ + { + name: "vector_bucket_name", + label: "Vector Bucket Name", + tooltip: "S3 bucket name for vector storage (will be auto-created if it doesn't exist)", + placeholder: "my-vector-bucket", + required: true, + type: "text", + }, + { + name: "index_name", + label: "Index Name", + tooltip: "Name for the vector index (optional, will be auto-generated if not provided)", + placeholder: "my-vector-index", + required: false, + type: "text", + }, + { + name: "aws_region_name", + label: "AWS Region", + tooltip: "AWS region where the S3 bucket is located (e.g., us-west-2)", + placeholder: "us-west-2", + required: true, + type: "text", + }, + { + name: "embedding_model", + label: "Embedding Model", + tooltip: "Select the embedding model to use for vector generation", + placeholder: "text-embedding-3-small", + required: true, + type: "select", + }, + ], }; export const getVectorStoreProviderLogoAndName = (providerValue: string): { logo: string; displayName: 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}