mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-09 18:25:53 +00:00
chore: resolve merge conflicts with upstream/main
Accept upstream's new models and re-apply removal of 133 deprecated models.
This commit is contained in:
Chesars
@@ -1,12 +1,19 @@
|
||||
name: "LiteLLM CodeQL config"
|
||||
|
||||
# Exclude queries that produce result sets > 2 GiB on this codebase,
|
||||
# causing 49+ minute runs that fail and block CI resources.
|
||||
# Use security-extended suite instead of security-and-quality to avoid
|
||||
# result sets > 2 GiB on this codebase that cause fatal OOM failures.
|
||||
queries:
|
||||
- uses: security-extended
|
||||
|
||||
# These two queries are security queries included in security-extended that
|
||||
# individually produce result sets > 2 GiB on this codebase, causing fatal
|
||||
# OOM failures. Exclude them as a safety net until CI confirms they no longer
|
||||
# OOM; drop these exclusions in a follow-up once verified.
|
||||
query-filters:
|
||||
- exclude:
|
||||
id: py/clear-text-logging-sensitive-data # CWE-312/CleartextLogging.ql — result set > 2 GiB
|
||||
id: py/clear-text-logging-sensitive-data # CWE-312 — > 2 GiB result set
|
||||
- exclude:
|
||||
id: py/polynomial-redos # CWE-730/PolynomialReDoS.ql — result set > 2 GiB
|
||||
id: py/polynomial-redos # CWE-730 — > 2 GiB result set
|
||||
|
||||
paths-ignore:
|
||||
- tests
|
||||
|
||||
@@ -0,0 +1,169 @@
|
||||
---
|
||||
slug: gemini_embedding_2_multimodal
|
||||
title: "Gemini Embedding 2 Preview: Multimodal Embeddings on LiteLLM"
|
||||
date: 2025-03-11T10:00:00
|
||||
authors:
|
||||
- name: Sameer Kankute
|
||||
title: SWE @ LiteLLM (LLM Translation)
|
||||
url: https://www.linkedin.com/in/sameer-kankute/
|
||||
image_url: https://pbs.twimg.com/profile_images/2001352686994907136/ONgNuSk5_400x400.jpg
|
||||
description: "Generate embeddings from text, images, audio, video, and PDFs with gemini-embedding-2-preview on LiteLLM via Gemini API and Vertex AI."
|
||||
tags: [gemini, embeddings, multimodal, vertex ai]
|
||||
hide_table_of_contents: false
|
||||
---
|
||||
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
|
||||
# Gemini Embedding 2 Preview: Multimodal Embeddings
|
||||
|
||||
LiteLLM now supports **multimodal embeddings** with `gemini-embedding-2-preview`—generating a single embedding from a mix of text, images, audio, video, and PDF content. Available via both the **Gemini API** (API key) and **Vertex AI** (GCP credentials).
|
||||
|
||||
## Supported Input Types
|
||||
|
||||
| Modality | Supported Formats |
|
||||
|----------|-------------------|
|
||||
| **Text** | Plain text |
|
||||
| **Image** | PNG, JPEG |
|
||||
| **Audio** | MP3, WAV |
|
||||
| **Video** | MP4, MOV |
|
||||
| **Documents** | PDF |
|
||||
|
||||
## Input Formats
|
||||
|
||||
LiteLLM accepts three input formats for multimodal content:
|
||||
|
||||
1. **Data URIs** – Base64-encoded inline: `data:image/png;base64,<encoded_data>`
|
||||
2. **GCS URLs** – Cloud Storage paths (Vertex AI): `gs://bucket/path/to/file.png`
|
||||
3. **Gemini File References** – Pre-uploaded files (Gemini API): `files/abc123`
|
||||
|
||||
## Quick Start
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="gemini" label="Gemini API">
|
||||
|
||||
```python
|
||||
from litellm import embedding
|
||||
import os
|
||||
|
||||
os.environ["GEMINI_API_KEY"] = "your-api-key"
|
||||
|
||||
# Text + Image (base64)
|
||||
response = embedding(
|
||||
model="gemini/gemini-embedding-2-preview",
|
||||
input=[
|
||||
"The food was delicious and the waiter...",
|
||||
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII"
|
||||
],
|
||||
)
|
||||
print(response)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="vertex" label="Vertex AI">
|
||||
|
||||
```python
|
||||
import litellm
|
||||
from litellm import embedding
|
||||
|
||||
litellm.vertex_project = "your-project-id"
|
||||
litellm.vertex_location = "us-central1"
|
||||
|
||||
# Text + Image (GCS URL)
|
||||
response = embedding(
|
||||
model="vertex_ai/gemini-embedding-2-preview",
|
||||
input=[
|
||||
"Describe this image",
|
||||
"gs://my-bucket/images/photo.png"
|
||||
],
|
||||
)
|
||||
print(response)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="proxy" label="LiteLLM Proxy">
|
||||
|
||||
**1. Config (config.yaml)**
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gemini-embedding-2-preview
|
||||
litellm_params:
|
||||
model: gemini/gemini-embedding-2-preview
|
||||
api_key: os.environ/GEMINI_API_KEY
|
||||
- model_name: vertex-gemini-embedding-2-preview
|
||||
litellm_params:
|
||||
model: vertex_ai/gemini-embedding-2-preview
|
||||
vertex_project: os.environ/VERTEXAI_PROJECT
|
||||
vertex_location: os.environ/VERTEXAI_LOCATION
|
||||
|
||||
general_settings:
|
||||
master_key: sk-1234
|
||||
```
|
||||
|
||||
**2. Start proxy**
|
||||
|
||||
```bash
|
||||
litellm --config config.yaml
|
||||
```
|
||||
|
||||
**3. Call embeddings**
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:4000/embeddings \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "gemini-embedding-2-preview",
|
||||
"input": [
|
||||
"The food was delicious and the waiter...",
|
||||
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII"
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Input Format Examples
|
||||
|
||||
| Format | Example | Provider |
|
||||
|--------|---------|----------|
|
||||
| **Data URI** | `data:image/png;base64,...` | Gemini, Vertex AI |
|
||||
| **GCS URL** | `gs://bucket/path/image.png` | Vertex AI |
|
||||
| **File reference** | `files/abc123` | Gemini API only |
|
||||
|
||||
### Supported MIME Types for Data URIs
|
||||
|
||||
- **Images:** `image/png`, `image/jpeg`
|
||||
- **Audio:** `audio/mpeg`, `audio/wav`
|
||||
- **Video:** `video/mp4`, `video/quicktime`
|
||||
- **Documents:** `application/pdf`
|
||||
|
||||
### GCS URL MIME Inference
|
||||
|
||||
For Vertex AI, MIME types are inferred from file extensions:
|
||||
|
||||
- `.png` → `image/png`
|
||||
- `.jpg` / `.jpeg` → `image/jpeg`
|
||||
- `.mp3` → `audio/mpeg`
|
||||
- `.wav` → `audio/wav`
|
||||
- `.mp4` → `video/mp4`
|
||||
- `.mov` → `video/quicktime`
|
||||
- `.pdf` → `application/pdf`
|
||||
|
||||
## Optional Parameters
|
||||
|
||||
| Parameter | Description | Maps to |
|
||||
|-----------|-------------|---------|
|
||||
| `dimensions` | Output embedding size | `outputDimensionality` |
|
||||
|
||||
```python
|
||||
response = embedding(
|
||||
model="gemini/gemini-embedding-2-preview",
|
||||
input=["text to embed"],
|
||||
dimensions=768, # Optional: control output vector size
|
||||
)
|
||||
```
|
||||
@@ -514,6 +514,57 @@ All models listed [here](https://ai.google.dev/gemini-api/docs/models/gemini) ar
|
||||
| Model Name | Function Call |
|
||||
| :--- | :--- |
|
||||
| text-embedding-004 | `embedding(model="gemini/text-embedding-004", input)` |
|
||||
| gemini-embedding-2-preview | `embedding(model="gemini/gemini-embedding-2-preview", input)` | [Multimodal docs](#gemini-embedding-2-preview-multimodal) |
|
||||
|
||||
### Gemini Embedding 2 Preview (Multimodal)
|
||||
|
||||
`gemini-embedding-2-preview` supports **multimodal embeddings**—text, images, audio, video, and PDF in a single request. See [blog post](/blog/gemini_embedding_2_multimodal) for details.
|
||||
|
||||
**Input formats:**
|
||||
- **Data URIs:** `data:image/png;base64,<encoded_data>`
|
||||
- **Gemini file references:** `files/abc123` (pre-uploaded via Gemini Files API)
|
||||
|
||||
**Supported MIME types:** `image/png`, `image/jpeg`, `audio/mpeg`, `audio/wav`, `video/mp4`, `video/quicktime`, `application/pdf`
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="sdk" label="SDK">
|
||||
|
||||
```python
|
||||
from litellm import embedding
|
||||
import os
|
||||
os.environ["GEMINI_API_KEY"] = ""
|
||||
|
||||
# Text + Image (base64)
|
||||
response = embedding(
|
||||
model="gemini/gemini-embedding-2-preview",
|
||||
input=[
|
||||
"The food was delicious and the waiter...",
|
||||
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII"
|
||||
],
|
||||
)
|
||||
print(response)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="proxy" label="PROXY">
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:4000/embeddings \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "gemini-embedding-2-preview",
|
||||
"input": [
|
||||
"The food was delicious and the waiter...",
|
||||
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII"
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
**Optional:** `dimensions` maps to Gemini's `outputDimensionality`.
|
||||
|
||||
|
||||
## Vertex AI Embedding Models
|
||||
|
||||
@@ -133,6 +133,21 @@ LiteLLM attempts [OAuth 2.0 Authorization Server Discovery](https://datatracker.
|
||||
|
||||
<br/>
|
||||
|
||||
### AWS SigV4 Authentication
|
||||
|
||||
For MCP servers hosted on [AWS Bedrock AgentCore](https://docs.aws.amazon.com/bedrock/latest/userguide/agentcore.html), select **AWS SigV4** as the authentication type. LiteLLM will sign every outgoing MCP request with your AWS credentials using [Signature Version 4](https://docs.aws.amazon.com/general/latest/gr/signature-version-4.html).
|
||||
|
||||
<Image
|
||||
img={require('../img/mcp_aws_sigv4_ui.png')}
|
||||
style={{width: '80%', display: 'block', margin: '0'}}
|
||||
/>
|
||||
|
||||
Fill in your AWS region, service name (defaults to `bedrock-agentcore`), and optionally your AWS access key and secret. If credentials are omitted, LiteLLM falls back to the boto3 credential chain (IAM roles, environment variables, etc.).
|
||||
|
||||
[**See full SigV4 setup guide**](./mcp_aws_sigv4.md)
|
||||
|
||||
<br/>
|
||||
|
||||
### Static Headers
|
||||
|
||||
Sometimes your MCP server needs specific headers on every request. Maybe it's an API key, maybe it's a custom header the server expects. Instead of configuring auth, you can just set them directly.
|
||||
|
||||
@@ -1,3 +1,7 @@
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
import Image from '@theme/IdealImage';
|
||||
|
||||
# MCP - AWS SigV4 Auth
|
||||
|
||||
Use AWS SigV4 authentication to connect LiteLLM to MCP servers hosted on [AWS Bedrock AgentCore](https://docs.aws.amazon.com/bedrock/latest/userguide/agentcore.html).
|
||||
@@ -10,6 +14,36 @@ LiteLLM's `aws_sigv4` auth type handles this automatically: every outgoing MCP r
|
||||
|
||||
## Quick Start
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="ui" label="LiteLLM UI">
|
||||
|
||||
1. Navigate to **MCP Servers** and click **Add New MCP Server**
|
||||
2. Set the transport to **Streamable HTTP**
|
||||
3. Select **AWS SigV4** as the authentication type
|
||||
4. Fill in your AWS credentials:
|
||||
|
||||
<Image
|
||||
img={require('../img/mcp_aws_sigv4_ui.png')}
|
||||
style={{width: '80%', display: 'block', margin: '0'}}
|
||||
/>
|
||||
|
||||
<br/>
|
||||
|
||||
| Field | Required | Description |
|
||||
|-------|----------|-------------|
|
||||
| **AWS Region** | Yes | AWS region for SigV4 signing (e.g., `us-east-1`) |
|
||||
| **AWS Service Name** | No | Defaults to `bedrock-agentcore` |
|
||||
| **AWS Access Key ID** | No | Falls back to boto3 credential chain if blank |
|
||||
| **AWS Secret Access Key** | No | Required if Access Key ID is provided |
|
||||
| **AWS Session Token** | No | Only needed for temporary STS credentials |
|
||||
|
||||
Once created, LiteLLM will sign every outgoing MCP request with SigV4. The server's tools appear automatically in the MCP Tools list.
|
||||
|
||||
**Editing credentials:** When editing an existing SigV4 server, leave credential fields blank to keep the current values. Only fields you fill in will be updated.
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="config" label="config.yaml">
|
||||
|
||||
### 1. Set AWS credentials
|
||||
|
||||
```bash
|
||||
@@ -60,9 +94,12 @@ arn%3Aaws%3Abedrock-agentcore%3Aus-east-1%3A123456789012%3Aruntime%2Fmy-mcp-serv
|
||||
litellm --config config.yaml
|
||||
```
|
||||
|
||||
### 4. Use the MCP tools
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
Once started, your AgentCore MCP tools are available through LiteLLM like any other MCP server:
|
||||
## Use the MCP tools
|
||||
|
||||
Once configured, your AgentCore MCP tools are available through LiteLLM like any other MCP server:
|
||||
|
||||
```bash title="List available tools"
|
||||
curl http://localhost:4000/mcp-rest/tools/list \
|
||||
|
||||
@@ -79,6 +79,7 @@ All models listed [here](https://github.com/BerriAI/litellm/blob/57f37f743886a02
|
||||
| textembedding-gecko@003 | `embedding(model="vertex_ai/textembedding-gecko@003", input)` |
|
||||
| text-embedding-preview-0409 | `embedding(model="vertex_ai/text-embedding-preview-0409", input)` |
|
||||
| text-multilingual-embedding-preview-0409 | `embedding(model="vertex_ai/text-multilingual-embedding-preview-0409", input)` |
|
||||
| gemini-embedding-2-preview | `embedding(model="vertex_ai/gemini-embedding-2-preview", input)` | [Multimodal docs](#gemini-embedding-2-preview-multimodal) |
|
||||
| Fine-tuned OR Custom Embedding models | `embedding(model="vertex_ai/<your-model-id>", input)` |
|
||||
|
||||
### Supported OpenAI (Unified) Params
|
||||
@@ -257,6 +258,71 @@ model_list:
|
||||
|
||||
## **Multi-Modal Embeddings**
|
||||
|
||||
### Gemini Embedding 2 Preview (Multimodal)
|
||||
|
||||
`gemini-embedding-2-preview` supports **unified multimodal embeddings**—text, images, audio, video, and PDF in a single request. See [blog post](/blog/gemini_embedding_2_multimodal) for details.
|
||||
|
||||
**Input formats:**
|
||||
- **Data URIs:** `data:image/png;base64,<encoded_data>`
|
||||
- **GCS URLs:** `gs://bucket/path/to/file.png` (MIME type inferred from extension)
|
||||
|
||||
**Supported MIME types:** `image/png`, `image/jpeg`, `audio/mpeg`, `audio/wav`, `video/mp4`, `video/quicktime`, `application/pdf`
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="sdk" label="SDK">
|
||||
|
||||
```python
|
||||
import litellm
|
||||
from litellm import embedding
|
||||
|
||||
litellm.vertex_project = "your-project-id"
|
||||
litellm.vertex_location = "us-central1"
|
||||
|
||||
# Text + Image (GCS URL)
|
||||
response = embedding(
|
||||
model="vertex_ai/gemini-embedding-2-preview",
|
||||
input=[
|
||||
"Describe this image",
|
||||
"gs://my-bucket/images/photo.png"
|
||||
],
|
||||
)
|
||||
|
||||
# Text + Image (base64)
|
||||
response = embedding(
|
||||
model="vertex_ai/gemini-embedding-2-preview",
|
||||
input=[
|
||||
"The food was delicious",
|
||||
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII"
|
||||
],
|
||||
)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="proxy" label="LiteLLM PROXY">
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: vertex-gemini-embedding-2-preview
|
||||
litellm_params:
|
||||
model: vertex_ai/gemini-embedding-2-preview
|
||||
vertex_project: "your-project-id"
|
||||
vertex_location: "us-central1"
|
||||
```
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:4000/embeddings \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "vertex-gemini-embedding-2-preview",
|
||||
"input": ["Describe this", "gs://bucket/image.png"]
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
### multimodalembedding@001 (Legacy)
|
||||
|
||||
Known Limitations:
|
||||
- Only supports 1 image / video / image per request
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 72 KiB |
@@ -247,23 +247,27 @@ def _get_embedding_url(
|
||||
- bge/endpoint_id -> strips to endpoint_id for endpoints/ routing
|
||||
- numeric model -> routes to endpoints/
|
||||
- regular model -> routes to publishers/google/models/
|
||||
"""
|
||||
endpoint = "predict"
|
||||
|
||||
# Strip routing prefixes (bge/, gemma/, etc.) for endpoint URL construction
|
||||
- models with uses_embed_content flag -> use embedContent endpoint instead of predict
|
||||
"""
|
||||
original_model = model
|
||||
model = get_vertex_base_model_name(model=model)
|
||||
|
||||
# Get base URL (handles global vs regional)
|
||||
try:
|
||||
model_info = litellm.get_model_info(
|
||||
model=original_model,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
uses_embed_content = model_info.get("uses_embed_content", False)
|
||||
except Exception:
|
||||
uses_embed_content = False
|
||||
|
||||
endpoint = "embedContent" if uses_embed_content else "predict"
|
||||
|
||||
base_url = get_vertex_base_url(vertex_location)
|
||||
|
||||
if model.isdigit():
|
||||
# https://us-central1-aiplatform.googleapis.com/v1/projects/$PROJECT_ID/locations/us-central1/endpoints/$ENDPOINT_ID:predict
|
||||
# https://aiplatform.googleapis.com/v1/projects/$PROJECT_ID/locations/global/endpoints/$ENDPOINT_ID:predict
|
||||
url = f"{base_url}/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}:{endpoint}"
|
||||
else:
|
||||
# Regular model -> publisher model
|
||||
# https://us-central1-aiplatform.googleapis.com/v1/projects/$PROJECT_ID/locations/us-central1/publishers/google/models/{model}:predict
|
||||
# https://aiplatform.googleapis.com/v1/projects/$PROJECT_ID/locations/global/publishers/google/models/{model}:predict
|
||||
url = f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}"
|
||||
|
||||
return url, endpoint
|
||||
|
||||
@@ -3,12 +3,11 @@ Google AI Studio /batchEmbedContents Embeddings Endpoint
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Literal, Optional, Union
|
||||
from typing import Any, Dict, Literal, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
@@ -19,15 +18,98 @@ from litellm.types.llms.vertex_ai import (
|
||||
VertexAIBatchEmbeddingsRequestBody,
|
||||
VertexAIBatchEmbeddingsResponseObject,
|
||||
)
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
|
||||
from ..gemini.vertex_and_google_ai_studio_gemini import VertexLLM
|
||||
from .batch_embed_content_transformation import (
|
||||
_is_file_reference,
|
||||
_is_multimodal_input,
|
||||
process_embed_content_response,
|
||||
process_response,
|
||||
transform_openai_input_gemini_content,
|
||||
transform_openai_input_gemini_embed_content,
|
||||
)
|
||||
|
||||
|
||||
class GoogleBatchEmbeddings(VertexLLM):
|
||||
def _resolve_file_references(
|
||||
self,
|
||||
input: EmbeddingInput,
|
||||
api_key: str,
|
||||
sync_handler: HTTPHandler,
|
||||
) -> Dict[str, Dict[str, str]]:
|
||||
"""
|
||||
Resolve Gemini file references (files/...) to get mime_type and uri.
|
||||
|
||||
Args:
|
||||
input: EmbeddingInput that may contain file references
|
||||
api_key: Gemini API key
|
||||
sync_handler: HTTP client
|
||||
|
||||
Returns:
|
||||
Dict mapping file name to {mime_type, uri}
|
||||
"""
|
||||
input_list = [input] if isinstance(input, str) else input
|
||||
resolved_files: Dict[str, Dict[str, str]] = {}
|
||||
|
||||
for element in input_list:
|
||||
if isinstance(element, str) and _is_file_reference(element):
|
||||
url = f"https://generativelanguage.googleapis.com/v1beta/{element}"
|
||||
headers = {"x-goog-api-key": api_key}
|
||||
response = sync_handler.get(url=url, headers=headers)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise Exception(
|
||||
f"Error fetching file {element}: {response.status_code} {response.text}"
|
||||
)
|
||||
|
||||
file_data = response.json()
|
||||
resolved_files[element] = {
|
||||
"mime_type": file_data.get("mimeType", ""),
|
||||
"uri": file_data.get("uri", element),
|
||||
}
|
||||
|
||||
return resolved_files
|
||||
|
||||
async def _async_resolve_file_references(
|
||||
self,
|
||||
input: EmbeddingInput,
|
||||
api_key: str,
|
||||
async_handler: AsyncHTTPHandler,
|
||||
) -> Dict[str, Dict[str, str]]:
|
||||
"""
|
||||
Async version of _resolve_file_references.
|
||||
|
||||
Args:
|
||||
input: EmbeddingInput that may contain file references
|
||||
api_key: Gemini API key
|
||||
async_handler: Async HTTP client
|
||||
|
||||
Returns:
|
||||
Dict mapping file name to {mime_type, uri}
|
||||
"""
|
||||
input_list = [input] if isinstance(input, str) else input
|
||||
resolved_files: Dict[str, Dict[str, str]] = {}
|
||||
|
||||
for element in input_list:
|
||||
if isinstance(element, str) and _is_file_reference(element):
|
||||
url = f"https://generativelanguage.googleapis.com/v1beta/{element}"
|
||||
headers = {"x-goog-api-key": api_key}
|
||||
response = await async_handler.get(url=url, headers=headers)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise Exception(
|
||||
f"Error fetching file {element}: {response.status_code} {response.text}"
|
||||
)
|
||||
|
||||
file_data = response.json()
|
||||
resolved_files[element] = {
|
||||
"mime_type": file_data.get("mimeType", ""),
|
||||
"uri": file_data.get("uri", element),
|
||||
}
|
||||
|
||||
return resolved_files
|
||||
|
||||
def batch_embeddings(
|
||||
self,
|
||||
model: str,
|
||||
@@ -54,20 +136,6 @@ class GoogleBatchEmbeddings(VertexLLM):
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
auth_header, url = self._get_token_and_url(
|
||||
model=model,
|
||||
auth_header=_auth_header,
|
||||
gemini_api_key=api_key,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
vertex_credentials=vertex_credentials,
|
||||
stream=None,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
should_use_v1beta1_features=False,
|
||||
mode="batch_embedding",
|
||||
)
|
||||
|
||||
if client is None:
|
||||
_params = {}
|
||||
if timeout is not None:
|
||||
@@ -83,9 +151,25 @@ class GoogleBatchEmbeddings(VertexLLM):
|
||||
|
||||
optional_params = optional_params or {}
|
||||
|
||||
### TRANSFORMATION ###
|
||||
request_data = transform_openai_input_gemini_content(
|
||||
input=input, model=model, optional_params=optional_params
|
||||
is_multimodal = _is_multimodal_input(input)
|
||||
use_embed_content = is_multimodal or (custom_llm_provider == "vertex_ai")
|
||||
if use_embed_content:
|
||||
mode = "embedding"
|
||||
else:
|
||||
mode = "batch_embedding"
|
||||
|
||||
auth_header, url = self._get_token_and_url(
|
||||
model=model,
|
||||
auth_header=_auth_header,
|
||||
gemini_api_key=api_key,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
vertex_credentials=vertex_credentials,
|
||||
stream=None,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
should_use_v1beta1_features=False,
|
||||
mode=mode,
|
||||
)
|
||||
|
||||
headers = {
|
||||
@@ -93,14 +177,46 @@ class GoogleBatchEmbeddings(VertexLLM):
|
||||
}
|
||||
if auth_header is not None:
|
||||
if isinstance(auth_header, dict):
|
||||
# For Gemini with custom api_base: auth_header is {"x-goog-api-key": "..."}
|
||||
headers.update(auth_header)
|
||||
else:
|
||||
# For Vertex AI: auth_header is a Bearer token string
|
||||
headers["Authorization"] = f"Bearer {auth_header}"
|
||||
if extra_headers is not None:
|
||||
headers.update(extra_headers)
|
||||
|
||||
if aembedding is True:
|
||||
return self.async_batch_embeddings( # type: ignore
|
||||
model=model,
|
||||
api_base=api_base,
|
||||
url=url,
|
||||
data=None,
|
||||
model_response=model_response,
|
||||
timeout=timeout,
|
||||
headers=headers,
|
||||
input=input,
|
||||
use_embed_content=use_embed_content,
|
||||
api_key=api_key,
|
||||
optional_params=optional_params,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
### TRANSFORMATION (sync path) ###
|
||||
if use_embed_content:
|
||||
resolved_files = {}
|
||||
if api_key:
|
||||
resolved_files = self._resolve_file_references(
|
||||
input=input, api_key=api_key, sync_handler=sync_handler
|
||||
)
|
||||
request_data = transform_openai_input_gemini_embed_content(
|
||||
input=input,
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
resolved_files=resolved_files,
|
||||
)
|
||||
else:
|
||||
request_data = transform_openai_input_gemini_content(
|
||||
input=input, model=model, optional_params=optional_params
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=input,
|
||||
@@ -112,18 +228,6 @@ class GoogleBatchEmbeddings(VertexLLM):
|
||||
},
|
||||
)
|
||||
|
||||
if aembedding is True:
|
||||
return self.async_batch_embeddings( # type: ignore
|
||||
model=model,
|
||||
api_base=api_base,
|
||||
url=url,
|
||||
data=request_data,
|
||||
model_response=model_response,
|
||||
timeout=timeout,
|
||||
headers=headers,
|
||||
input=input,
|
||||
)
|
||||
|
||||
response = sync_handler.post(
|
||||
url=url,
|
||||
headers=headers,
|
||||
@@ -134,26 +238,38 @@ class GoogleBatchEmbeddings(VertexLLM):
|
||||
raise Exception(f"Error: {response.status_code} {response.text}")
|
||||
|
||||
_json_response = response.json()
|
||||
_predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore
|
||||
|
||||
return process_response(
|
||||
model=model,
|
||||
model_response=model_response,
|
||||
_predictions=_predictions,
|
||||
input=input,
|
||||
)
|
||||
|
||||
if use_embed_content:
|
||||
return process_embed_content_response(
|
||||
input=input,
|
||||
model_response=model_response,
|
||||
model=model,
|
||||
response_json=_json_response,
|
||||
)
|
||||
else:
|
||||
_predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore
|
||||
return process_response(
|
||||
model=model,
|
||||
model_response=model_response,
|
||||
_predictions=_predictions,
|
||||
input=input,
|
||||
)
|
||||
|
||||
async def async_batch_embeddings(
|
||||
self,
|
||||
model: str,
|
||||
api_base: Optional[str],
|
||||
url: str,
|
||||
data: VertexAIBatchEmbeddingsRequestBody,
|
||||
data: Optional[Union[VertexAIBatchEmbeddingsRequestBody, dict]],
|
||||
model_response: EmbeddingResponse,
|
||||
input: EmbeddingInput,
|
||||
timeout: Optional[Union[float, httpx.Timeout]],
|
||||
headers={},
|
||||
client: Optional[AsyncHTTPHandler] = None,
|
||||
use_embed_content: bool = False,
|
||||
api_key: Optional[str] = None,
|
||||
optional_params: Optional[dict] = None,
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> EmbeddingResponse:
|
||||
if client is None:
|
||||
_params = {}
|
||||
@@ -171,6 +287,36 @@ class GoogleBatchEmbeddings(VertexLLM):
|
||||
else:
|
||||
async_handler = client # type: ignore
|
||||
|
||||
### TRANSFORMATION (async path) ###
|
||||
if use_embed_content:
|
||||
resolved_files = {}
|
||||
if api_key:
|
||||
resolved_files = await self._async_resolve_file_references(
|
||||
input=input, api_key=api_key, async_handler=async_handler
|
||||
)
|
||||
data = transform_openai_input_gemini_embed_content(
|
||||
input=input,
|
||||
model=model,
|
||||
optional_params=optional_params or {},
|
||||
resolved_files=resolved_files,
|
||||
)
|
||||
else:
|
||||
data = transform_openai_input_gemini_content(
|
||||
input=input, model=model, optional_params=optional_params or {}
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
if logging_obj is not None:
|
||||
logging_obj.pre_call(
|
||||
input=input,
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": data,
|
||||
"api_base": url,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
|
||||
response = await async_handler.post(
|
||||
url=url,
|
||||
headers=headers,
|
||||
@@ -181,11 +327,19 @@ class GoogleBatchEmbeddings(VertexLLM):
|
||||
raise Exception(f"Error: {response.status_code} {response.text}")
|
||||
|
||||
_json_response = response.json()
|
||||
_predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore
|
||||
|
||||
return process_response(
|
||||
model=model,
|
||||
model_response=model_response,
|
||||
_predictions=_predictions,
|
||||
input=input,
|
||||
)
|
||||
|
||||
if use_embed_content:
|
||||
return process_embed_content_response(
|
||||
input=input,
|
||||
model_response=model_response,
|
||||
model=model,
|
||||
response_json=_json_response,
|
||||
)
|
||||
else:
|
||||
_predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore
|
||||
return process_response(
|
||||
model=model,
|
||||
model_response=model_response,
|
||||
_predictions=_predictions,
|
||||
input=input,
|
||||
)
|
||||
|
||||
@@ -4,20 +4,142 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc
|
||||
Why separate file? Make it easy to see how transformation works
|
||||
"""
|
||||
|
||||
from typing import List
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
from litellm.types.llms.openai import EmbeddingInput
|
||||
from litellm.types.llms.vertex_ai import (
|
||||
BlobType,
|
||||
ContentType,
|
||||
EmbedContentRequest,
|
||||
FileDataType,
|
||||
PartType,
|
||||
VertexAIBatchEmbeddingsRequestBody,
|
||||
VertexAIBatchEmbeddingsResponseObject,
|
||||
)
|
||||
from litellm.types.utils import Embedding, Usage
|
||||
from litellm.types.utils import Embedding, EmbeddingResponse, Usage
|
||||
from litellm.utils import get_formatted_prompt, token_counter
|
||||
|
||||
SUPPORTED_EMBEDDING_MIME_TYPES = {
|
||||
"image/png",
|
||||
"image/jpeg",
|
||||
"audio/mpeg",
|
||||
"audio/wav",
|
||||
"video/mp4",
|
||||
"video/quicktime",
|
||||
"application/pdf",
|
||||
}
|
||||
|
||||
|
||||
def _is_file_reference(s: str) -> bool:
|
||||
"""Check if string is a Gemini file reference (files/...)."""
|
||||
return isinstance(s, str) and s.startswith("files/")
|
||||
|
||||
|
||||
def _is_gcs_url(s: str) -> bool:
|
||||
"""Check if string is a GCS URL (gs://...)."""
|
||||
return isinstance(s, str) and s.startswith("gs://")
|
||||
|
||||
|
||||
def _infer_mime_type_from_gcs_url(gcs_url: str) -> str:
|
||||
"""
|
||||
Infer MIME type from GCS URL file extension.
|
||||
|
||||
Args:
|
||||
gcs_url: GCS URL like gs://bucket/path/to/file.png
|
||||
|
||||
Returns:
|
||||
str: Inferred MIME type
|
||||
|
||||
Raises:
|
||||
ValueError: If file extension is not supported
|
||||
"""
|
||||
extension_to_mime = {
|
||||
".png": "image/png",
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".mp3": "audio/mpeg",
|
||||
".wav": "audio/wav",
|
||||
".mp4": "video/mp4",
|
||||
".mov": "video/quicktime",
|
||||
".pdf": "application/pdf",
|
||||
}
|
||||
|
||||
gcs_url_lower = gcs_url.lower()
|
||||
for ext, mime_type in extension_to_mime.items():
|
||||
if gcs_url_lower.endswith(ext):
|
||||
return mime_type
|
||||
|
||||
raise ValueError(
|
||||
f"Unable to infer MIME type from GCS URL: {gcs_url}. "
|
||||
f"Supported extensions: {', '.join(extension_to_mime.keys())}"
|
||||
)
|
||||
|
||||
|
||||
def _parse_data_url(data_url: str) -> Tuple[str, str]:
|
||||
"""
|
||||
Parse a data URL to extract the media type and base64 data.
|
||||
|
||||
Args:
|
||||
data_url: Data URL in format: data:image/jpeg;base64,/9j/4AAQ...
|
||||
|
||||
Returns:
|
||||
tuple: (media_type, base64_data)
|
||||
media_type: e.g., "image/jpeg", "video/mp4", "audio/mpeg"
|
||||
base64_data: The base64-encoded data without the prefix
|
||||
|
||||
Raises:
|
||||
ValueError: If data URL format is invalid or MIME type is unsupported
|
||||
"""
|
||||
if not data_url.startswith("data:"):
|
||||
raise ValueError(f"Invalid data URL format: {data_url[:50]}...")
|
||||
|
||||
if "," not in data_url:
|
||||
raise ValueError(f"Invalid data URL format (missing comma): {data_url[:50]}...")
|
||||
|
||||
metadata, base64_data = data_url.split(",", 1)
|
||||
|
||||
metadata = metadata[5:]
|
||||
|
||||
if ";" in metadata:
|
||||
media_type = metadata.split(";")[0]
|
||||
else:
|
||||
media_type = metadata
|
||||
|
||||
if media_type not in SUPPORTED_EMBEDDING_MIME_TYPES:
|
||||
raise ValueError(
|
||||
f"Unsupported MIME type for embedding: {media_type}. "
|
||||
f"Supported types: {', '.join(sorted(SUPPORTED_EMBEDDING_MIME_TYPES))}"
|
||||
)
|
||||
|
||||
return media_type, base64_data
|
||||
|
||||
|
||||
def _is_multimodal_input(input: EmbeddingInput) -> bool:
|
||||
"""
|
||||
Check if the input contains multimodal data (data URIs, file references, or GCS URLs).
|
||||
|
||||
Args:
|
||||
input: EmbeddingInput (str or List[str])
|
||||
|
||||
Returns:
|
||||
bool: True if any element is a data URI, file reference, or GCS URL
|
||||
"""
|
||||
if isinstance(input, str):
|
||||
input_list = [input]
|
||||
else:
|
||||
input_list = input
|
||||
|
||||
for element in input_list:
|
||||
if isinstance(element, str):
|
||||
if element.startswith("data:") and ";base64," in element:
|
||||
return True
|
||||
if _is_file_reference(element):
|
||||
return True
|
||||
if _is_gcs_url(element):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def transform_openai_input_gemini_content(
|
||||
input: EmbeddingInput, model: str, optional_params: dict
|
||||
@@ -26,12 +148,17 @@ def transform_openai_input_gemini_content(
|
||||
The content to embed. Only the parts.text fields will be counted.
|
||||
"""
|
||||
gemini_model_name = "models/{}".format(model)
|
||||
|
||||
gemini_params = optional_params.copy()
|
||||
if "dimensions" in gemini_params:
|
||||
gemini_params["outputDimensionality"] = gemini_params.pop("dimensions")
|
||||
|
||||
requests: List[EmbedContentRequest] = []
|
||||
if isinstance(input, str):
|
||||
request = EmbedContentRequest(
|
||||
model=gemini_model_name,
|
||||
content=ContentType(parts=[PartType(text=input)]),
|
||||
**optional_params
|
||||
**gemini_params
|
||||
)
|
||||
requests.append(request)
|
||||
else:
|
||||
@@ -39,13 +166,119 @@ def transform_openai_input_gemini_content(
|
||||
request = EmbedContentRequest(
|
||||
model=gemini_model_name,
|
||||
content=ContentType(parts=[PartType(text=i)]),
|
||||
**optional_params
|
||||
**gemini_params
|
||||
)
|
||||
requests.append(request)
|
||||
|
||||
return VertexAIBatchEmbeddingsRequestBody(requests=requests)
|
||||
|
||||
|
||||
def transform_openai_input_gemini_embed_content(
|
||||
input: EmbeddingInput,
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
resolved_files: Optional[Dict[str, Dict[str, str]]] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Transform OpenAI embedding input to Gemini embedContent format (multimodal).
|
||||
|
||||
Args:
|
||||
input: EmbeddingInput (str or List[str]) with text, data URIs, or file references
|
||||
model: Model name
|
||||
optional_params: Additional parameters (taskType, outputDimensionality, etc.)
|
||||
resolved_files: Dict mapping file names (files/abc) to {mime_type, uri}
|
||||
|
||||
Returns:
|
||||
dict: Gemini embedContent request body with content.parts
|
||||
"""
|
||||
resolved_files = resolved_files or {}
|
||||
|
||||
gemini_params = optional_params.copy()
|
||||
if "dimensions" in gemini_params:
|
||||
gemini_params["outputDimensionality"] = gemini_params.pop("dimensions")
|
||||
|
||||
input_list = [input] if isinstance(input, str) else input
|
||||
parts: List[PartType] = []
|
||||
|
||||
for element in input_list:
|
||||
if not isinstance(element, str):
|
||||
raise ValueError(f"Unsupported input type: {type(element)}")
|
||||
|
||||
if element.startswith("data:") and ";base64," in element:
|
||||
mime_type, base64_data = _parse_data_url(element)
|
||||
blob: BlobType = {"mime_type": mime_type, "data": base64_data}
|
||||
parts.append(PartType(inline_data=blob))
|
||||
elif _is_gcs_url(element):
|
||||
mime_type = _infer_mime_type_from_gcs_url(element)
|
||||
file_data: FileDataType = {
|
||||
"mime_type": mime_type,
|
||||
"file_uri": element,
|
||||
}
|
||||
parts.append(PartType(file_data=file_data))
|
||||
elif _is_file_reference(element):
|
||||
if element not in resolved_files:
|
||||
raise ValueError(f"File reference {element} not resolved")
|
||||
file_info = resolved_files[element]
|
||||
file_data_ref: FileDataType = {
|
||||
"mime_type": file_info["mime_type"],
|
||||
"file_uri": file_info["uri"],
|
||||
}
|
||||
parts.append(PartType(file_data=file_data_ref))
|
||||
else:
|
||||
parts.append(PartType(text=element))
|
||||
|
||||
request_body: dict = {
|
||||
"content": ContentType(parts=parts),
|
||||
**gemini_params,
|
||||
}
|
||||
|
||||
return request_body
|
||||
|
||||
|
||||
def process_embed_content_response(
|
||||
input: EmbeddingInput,
|
||||
model_response: EmbeddingResponse,
|
||||
model: str,
|
||||
response_json: dict,
|
||||
) -> EmbeddingResponse:
|
||||
"""
|
||||
Process Gemini embedContent response (single embedding for multimodal input).
|
||||
|
||||
Args:
|
||||
input: Original input
|
||||
model_response: EmbeddingResponse to populate
|
||||
model: Model name
|
||||
response_json: Raw JSON response from embedContent endpoint
|
||||
|
||||
Returns:
|
||||
EmbeddingResponse with single embedding
|
||||
"""
|
||||
if "embedding" not in response_json:
|
||||
raise ValueError(f"embedContent response missing 'embedding' field: {response_json}")
|
||||
|
||||
embedding_data = response_json["embedding"]
|
||||
|
||||
openai_embedding = Embedding(
|
||||
embedding=embedding_data["values"],
|
||||
index=0,
|
||||
object="embedding",
|
||||
)
|
||||
|
||||
model_response.data = [openai_embedding]
|
||||
model_response.model = model
|
||||
|
||||
if _is_multimodal_input(input):
|
||||
prompt_tokens = 0
|
||||
else:
|
||||
input_text = get_formatted_prompt(data={"input": input}, call_type="embedding")
|
||||
prompt_tokens = token_counter(model=model, text=input_text)
|
||||
model_response.usage = Usage(
|
||||
prompt_tokens=prompt_tokens, total_tokens=prompt_tokens
|
||||
)
|
||||
|
||||
return model_response
|
||||
|
||||
|
||||
def process_response(
|
||||
input: EmbeddingInput,
|
||||
model_response: EmbeddingResponse,
|
||||
|
||||
+27
-2
@@ -132,6 +132,7 @@ from litellm.utils import (
|
||||
create_tokenizer,
|
||||
get_api_key,
|
||||
get_llm_provider,
|
||||
get_model_info,
|
||||
get_non_default_completion_params,
|
||||
get_non_default_transcription_params,
|
||||
get_optional_params_embeddings,
|
||||
@@ -5190,13 +5191,37 @@ def embedding( # noqa: PLR0915
|
||||
or get_secret_str("VERTEX_API_BASE")
|
||||
)
|
||||
|
||||
if (
|
||||
try:
|
||||
model_info = get_model_info(model=model, custom_llm_provider="vertex_ai")
|
||||
uses_embed_content = model_info.get("uses_embed_content", False)
|
||||
except Exception:
|
||||
uses_embed_content = False
|
||||
|
||||
if uses_embed_content:
|
||||
response = google_batch_embeddings.batch_embeddings( # type: ignore
|
||||
model=model,
|
||||
input=input,
|
||||
encoding=_get_encoding(),
|
||||
logging_obj=logging,
|
||||
optional_params=optional_params,
|
||||
model_response=EmbeddingResponse(),
|
||||
vertex_project=vertex_ai_project,
|
||||
vertex_location=vertex_ai_location,
|
||||
vertex_credentials=vertex_credentials,
|
||||
aembedding=aembedding,
|
||||
print_verbose=print_verbose,
|
||||
custom_llm_provider="vertex_ai",
|
||||
api_key=None,
|
||||
api_base=api_base,
|
||||
client=client,
|
||||
extra_headers=headers,
|
||||
)
|
||||
elif (
|
||||
"image" in optional_params
|
||||
or "video" in optional_params
|
||||
or model
|
||||
in vertex_multimodal_embedding.SUPPORTED_MULTIMODAL_EMBEDDING_MODELS
|
||||
):
|
||||
# multimodal embedding is supported on vertex httpx
|
||||
response = vertex_multimodal_embedding.multimodal_embedding(
|
||||
model=model,
|
||||
input=input,
|
||||
|
||||
@@ -2565,6 +2565,19 @@
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"azure/gpt-35-turbo-1106": {
|
||||
"deprecation_date": "2025-03-31",
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "azure",
|
||||
"max_input_tokens": 16384,
|
||||
"max_output_tokens": 4096,
|
||||
"max_tokens": 4096,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2e-06,
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"azure/gpt-35-turbo-16k": {
|
||||
"input_cost_per_token": 3e-06,
|
||||
"litellm_provider": "azure",
|
||||
@@ -8045,6 +8058,20 @@
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"cerebras/zai-glm-4.6": {
|
||||
"deprecation_date": "2026-01-20",
|
||||
"input_cost_per_token": 2.25e-06,
|
||||
"litellm_provider": "cerebras",
|
||||
"max_input_tokens": 128000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.75e-06,
|
||||
"source": "https://www.cerebras.ai/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"cerebras/zai-glm-4.7": {
|
||||
"input_cost_per_token": 2.25e-06,
|
||||
"litellm_provider": "cerebras",
|
||||
@@ -13697,6 +13724,96 @@
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"gemini-live-2.5-flash-preview-native-audio-09-2025": {
|
||||
"cache_read_input_token_cost": 7.5e-08,
|
||||
"input_cost_per_audio_token": 3e-06,
|
||||
"input_cost_per_token": 3e-07,
|
||||
"litellm_provider": "vertex_ai-language-models",
|
||||
"max_audio_length_hours": 8.4,
|
||||
"max_audio_per_prompt": 1,
|
||||
"max_images_per_prompt": 3000,
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 65535,
|
||||
"max_pdf_size_mb": 30,
|
||||
"max_tokens": 65535,
|
||||
"max_video_length": 1,
|
||||
"max_videos_per_prompt": 10,
|
||||
"mode": "realtime",
|
||||
"output_cost_per_audio_token": 1.2e-05,
|
||||
"output_cost_per_token": 2e-06,
|
||||
"source": "https://ai.google.dev/gemini-api/docs/pricing",
|
||||
"supported_endpoints": [
|
||||
"/vertex_ai/live"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image",
|
||||
"audio",
|
||||
"video"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text",
|
||||
"audio"
|
||||
],
|
||||
"supports_audio_input": true,
|
||||
"supports_audio_output": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": 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/gemini-live-2.5-flash-preview-native-audio-09-2025": {
|
||||
"cache_read_input_token_cost": 7.5e-08,
|
||||
"input_cost_per_audio_token": 3e-06,
|
||||
"input_cost_per_token": 3e-07,
|
||||
"litellm_provider": "gemini",
|
||||
"max_audio_length_hours": 8.4,
|
||||
"max_audio_per_prompt": 1,
|
||||
"max_images_per_prompt": 3000,
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 65535,
|
||||
"max_pdf_size_mb": 30,
|
||||
"max_tokens": 65535,
|
||||
"max_video_length": 1,
|
||||
"max_videos_per_prompt": 10,
|
||||
"mode": "realtime",
|
||||
"output_cost_per_audio_token": 1.2e-05,
|
||||
"output_cost_per_token": 2e-06,
|
||||
"rpm": 100000,
|
||||
"source": "https://ai.google.dev/gemini-api/docs/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/realtime"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image",
|
||||
"audio",
|
||||
"video"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text",
|
||||
"audio"
|
||||
],
|
||||
"supports_audio_input": true,
|
||||
"supports_audio_output": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"tpm": 8000000
|
||||
},
|
||||
"gemini-2.5-flash-lite-preview-06-17": {
|
||||
"deprecation_date": "2025-11-18",
|
||||
"cache_read_input_token_cost": 2.5e-08,
|
||||
@@ -14321,6 +14438,32 @@
|
||||
"output_vector_size": 3072,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models"
|
||||
},
|
||||
"gemini-embedding-2-preview": {
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"input_cost_per_video_per_second": 0.0237,
|
||||
"litellm_provider": "vertex_ai-embedding-models",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0,
|
||||
"output_vector_size": 3072,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
|
||||
"uses_embed_content": true
|
||||
},
|
||||
"vertex_ai/gemini-embedding-2-preview": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0,
|
||||
"output_vector_size": 3072,
|
||||
"source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal",
|
||||
"supports_multimodal": true,
|
||||
"uses_embed_content": true
|
||||
},
|
||||
"gemini/gemini-embedding-001": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "gemini",
|
||||
@@ -14333,6 +14476,19 @@
|
||||
"source": "https://ai.google.dev/gemini-api/docs/embeddings#model-versions",
|
||||
"tpm": 10000000
|
||||
},
|
||||
"gemini/gemini-embedding-2-preview": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0,
|
||||
"output_vector_size": 3072,
|
||||
"rpm": 10000,
|
||||
"source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal",
|
||||
"supports_multimodal": true,
|
||||
"tpm": 10000000
|
||||
},
|
||||
"gemini/gemini-2.0-flash": {
|
||||
"cache_read_input_token_cost": 2.5e-08,
|
||||
"deprecation_date": "2026-06-01",
|
||||
@@ -16588,6 +16744,20 @@
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"gpt-4-0613": {
|
||||
"deprecation_date": "2025-06-06",
|
||||
"input_cost_per_token": 3e-05,
|
||||
"litellm_provider": "openai",
|
||||
"max_input_tokens": 8192,
|
||||
"max_output_tokens": 4096,
|
||||
"max_tokens": 4096,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-05,
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"gpt-4-1106-preview": {
|
||||
"deprecation_date": "2026-03-26",
|
||||
"input_cost_per_token": 1e-05,
|
||||
@@ -27107,6 +27277,16 @@
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 1536
|
||||
},
|
||||
"text-embedding-ada-002-v2": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"input_cost_per_token_batches": 5e-08,
|
||||
"litellm_provider": "openai",
|
||||
"max_input_tokens": 8191,
|
||||
"max_tokens": 8191,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_cost_per_token_batches": 0.0
|
||||
},
|
||||
"text-embedding-large-exp-03-07": {
|
||||
"input_cost_per_character": 2.5e-08,
|
||||
"input_cost_per_token": 1e-07,
|
||||
@@ -36686,32 +36866,5 @@
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"azure/gpt-35-turbo-1106": {
|
||||
"deprecation_date": "2025-03-31",
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "azure",
|
||||
"max_input_tokens": 16384,
|
||||
"max_output_tokens": 4096,
|
||||
"max_tokens": 4096,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2e-06,
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"cerebras/zai-glm-4.6": {
|
||||
"deprecation_date": "2026-01-20",
|
||||
"input_cost_per_token": 2.25e-06,
|
||||
"litellm_provider": "cerebras",
|
||||
"max_input_tokens": 128000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.75e-06,
|
||||
"source": "https://www.cerebras.ai/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
}
|
||||
}
|
||||
|
||||
@@ -100,6 +100,50 @@ def encrypt_credentials(
|
||||
value=client_secret,
|
||||
new_encryption_key=encryption_key,
|
||||
)
|
||||
# AWS SigV4 credential fields
|
||||
aws_access_key_id = credentials.get("aws_access_key_id")
|
||||
if aws_access_key_id is not None:
|
||||
credentials["aws_access_key_id"] = encrypt_value_helper(
|
||||
value=aws_access_key_id,
|
||||
new_encryption_key=encryption_key,
|
||||
)
|
||||
aws_secret_access_key = credentials.get("aws_secret_access_key")
|
||||
if aws_secret_access_key is not None:
|
||||
credentials["aws_secret_access_key"] = encrypt_value_helper(
|
||||
value=aws_secret_access_key,
|
||||
new_encryption_key=encryption_key,
|
||||
)
|
||||
aws_session_token = credentials.get("aws_session_token")
|
||||
if aws_session_token is not None:
|
||||
credentials["aws_session_token"] = encrypt_value_helper(
|
||||
value=aws_session_token,
|
||||
new_encryption_key=encryption_key,
|
||||
)
|
||||
# aws_region_name and aws_service_name are NOT secrets — stored as-is
|
||||
return credentials
|
||||
|
||||
|
||||
def decrypt_credentials(
|
||||
credentials: MCPCredentials,
|
||||
) -> MCPCredentials:
|
||||
"""Decrypt all secret fields in an MCPCredentials dict using the global salt key."""
|
||||
secret_fields = [
|
||||
"auth_value",
|
||||
"client_id",
|
||||
"client_secret",
|
||||
"aws_access_key_id",
|
||||
"aws_secret_access_key",
|
||||
"aws_session_token",
|
||||
]
|
||||
for field in secret_fields:
|
||||
value = credentials.get(field)
|
||||
if value is not None:
|
||||
credentials[field] = decrypt_value_helper(
|
||||
value=value,
|
||||
key=field,
|
||||
exception_type="debug",
|
||||
return_original_value=True,
|
||||
)
|
||||
return credentials
|
||||
|
||||
|
||||
@@ -350,9 +394,57 @@ async def update_mcp_server(
|
||||
"""
|
||||
Update a new mcp server record in the db
|
||||
"""
|
||||
import json
|
||||
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
|
||||
# Use helper to prepare data with proper JSON serialization
|
||||
data_dict = _prepare_mcp_server_data(data)
|
||||
|
||||
# Pre-fetch existing record once if we need it for auth_type or credential logic
|
||||
existing = None
|
||||
has_credentials = "credentials" in data_dict and data_dict["credentials"] is not None
|
||||
if data.auth_type or has_credentials:
|
||||
existing = await prisma_client.db.litellm_mcpservertable.find_unique(
|
||||
where={"server_id": data.server_id}
|
||||
)
|
||||
|
||||
# Clear stale credentials when auth_type changes but no new credentials provided
|
||||
if (
|
||||
data.auth_type
|
||||
and "credentials" not in data_dict
|
||||
and existing
|
||||
and existing.auth_type is not None
|
||||
and existing.auth_type != data.auth_type
|
||||
):
|
||||
data_dict["credentials"] = None
|
||||
|
||||
# Merge credentials: preserve existing fields not present in the update.
|
||||
# Without this, a partial credential update (e.g. changing only region)
|
||||
# would wipe encrypted secrets that the UI cannot display back.
|
||||
if "credentials" in data_dict and data_dict["credentials"] is not None:
|
||||
if existing and existing.credentials:
|
||||
# Only merge when auth_type is unchanged. Switching auth types
|
||||
# (e.g. oauth2 → api_key) should replace credentials entirely
|
||||
# to avoid stale secrets from the previous auth type lingering.
|
||||
auth_type_unchanged = (
|
||||
data.auth_type is None or data.auth_type == existing.auth_type
|
||||
)
|
||||
if auth_type_unchanged:
|
||||
existing_creds = (
|
||||
json.loads(existing.credentials)
|
||||
if isinstance(existing.credentials, str)
|
||||
else dict(existing.credentials)
|
||||
)
|
||||
new_creds = (
|
||||
json.loads(data_dict["credentials"])
|
||||
if isinstance(data_dict["credentials"], str)
|
||||
else dict(data_dict["credentials"])
|
||||
)
|
||||
# New values override existing; existing keys not in update are preserved
|
||||
merged = {**existing_creds, **new_creds}
|
||||
data_dict["credentials"] = safe_dumps(merged)
|
||||
|
||||
# Add audit fields
|
||||
data_dict["updated_by"] = touched_by
|
||||
|
||||
@@ -374,8 +466,12 @@ async def rotate_mcp_server_credentials_master_key(
|
||||
continue
|
||||
|
||||
credentials_copy = dict(credentials)
|
||||
encrypted_credentials = encrypt_credentials(
|
||||
# Decrypt with current key first, then re-encrypt with new key
|
||||
decrypted_credentials = decrypt_credentials(
|
||||
credentials=cast(MCPCredentials, credentials_copy),
|
||||
)
|
||||
encrypted_credentials = encrypt_credentials(
|
||||
credentials=decrypted_credentials,
|
||||
encryption_key=new_master_key,
|
||||
)
|
||||
|
||||
|
||||
@@ -597,9 +597,10 @@ class MCPServerManager:
|
||||
else:
|
||||
client_secret_value = encrypted_client_secret
|
||||
|
||||
# TODO: Add AWS SigV4 credential decryption here when DB-stored
|
||||
# SigV4 MCP servers are supported. Requires corresponding changes
|
||||
# to encrypt_credentials() in db.py and MCPCredentials TypedDict.
|
||||
# AWS SigV4 credential fields
|
||||
aws_creds = self._extract_aws_credentials(
|
||||
credentials_dict, credentials_are_encrypted
|
||||
)
|
||||
|
||||
scopes: Optional[List[str]] = None
|
||||
if credentials_dict:
|
||||
@@ -679,6 +680,12 @@ class MCPServerManager:
|
||||
is_byok=bool(getattr(mcp_server, "is_byok", False)),
|
||||
byok_description=getattr(mcp_server, "byok_description", None) or [],
|
||||
byok_api_key_help_url=getattr(mcp_server, "byok_api_key_help_url", None),
|
||||
# AWS SigV4 fields
|
||||
aws_access_key_id=aws_creds.get("aws_access_key_id"),
|
||||
aws_secret_access_key=aws_creds.get("aws_secret_access_key"),
|
||||
aws_session_token=aws_creds.get("aws_session_token"),
|
||||
aws_region_name=aws_creds.get("aws_region_name"),
|
||||
aws_service_name=aws_creds.get("aws_service_name"),
|
||||
)
|
||||
return new_server
|
||||
|
||||
@@ -1518,6 +1525,52 @@ class MCPServerManager:
|
||||
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _decrypt_credential_field(
|
||||
encrypted_value: Optional[str],
|
||||
key: str,
|
||||
credentials_are_encrypted: bool,
|
||||
) -> Optional[str]:
|
||||
"""Decrypt a single credential field, or return as-is if not encrypted."""
|
||||
if not encrypted_value:
|
||||
return None
|
||||
if credentials_are_encrypted:
|
||||
return decrypt_value_helper(
|
||||
value=encrypted_value,
|
||||
key=key,
|
||||
exception_type="debug",
|
||||
return_original_value=True,
|
||||
)
|
||||
return encrypted_value
|
||||
|
||||
def _extract_aws_credentials(
|
||||
self,
|
||||
credentials_dict: Optional[Dict[str, str]],
|
||||
credentials_are_encrypted: bool,
|
||||
) -> Dict[str, Optional[str]]:
|
||||
"""Extract and decrypt AWS SigV4 credential fields from credentials dict."""
|
||||
if not credentials_dict:
|
||||
return {}
|
||||
return {
|
||||
"aws_access_key_id": self._decrypt_credential_field(
|
||||
credentials_dict.get("aws_access_key_id"),
|
||||
"aws_access_key_id",
|
||||
credentials_are_encrypted,
|
||||
),
|
||||
"aws_secret_access_key": self._decrypt_credential_field(
|
||||
credentials_dict.get("aws_secret_access_key"),
|
||||
"aws_secret_access_key",
|
||||
credentials_are_encrypted,
|
||||
),
|
||||
"aws_session_token": self._decrypt_credential_field(
|
||||
credentials_dict.get("aws_session_token"),
|
||||
"aws_session_token",
|
||||
credentials_are_encrypted,
|
||||
),
|
||||
"aws_region_name": credentials_dict.get("aws_region_name"),
|
||||
"aws_service_name": credentials_dict.get("aws_service_name"),
|
||||
}
|
||||
|
||||
def _extract_scopes(self, scopes_value: Any) -> Optional[List[str]]:
|
||||
if isinstance(scopes_value, str):
|
||||
scopes = [s.strip() for s in scopes_value.split() if s.strip()]
|
||||
|
||||
@@ -412,6 +412,17 @@ if MCP_AVAILABLE:
|
||||
inherited_credentials["client_secret"] = existing_server.client_secret
|
||||
if existing_server.scopes:
|
||||
inherited_credentials["scopes"] = existing_server.scopes
|
||||
# AWS SigV4 fields
|
||||
if existing_server.aws_access_key_id:
|
||||
inherited_credentials["aws_access_key_id"] = existing_server.aws_access_key_id
|
||||
if existing_server.aws_secret_access_key:
|
||||
inherited_credentials["aws_secret_access_key"] = existing_server.aws_secret_access_key
|
||||
if existing_server.aws_session_token:
|
||||
inherited_credentials["aws_session_token"] = existing_server.aws_session_token
|
||||
if existing_server.aws_region_name:
|
||||
inherited_credentials["aws_region_name"] = existing_server.aws_region_name
|
||||
if existing_server.aws_service_name:
|
||||
inherited_credentials["aws_service_name"] = existing_server.aws_service_name
|
||||
|
||||
if not inherited_credentials:
|
||||
return payload
|
||||
|
||||
@@ -404,7 +404,9 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
|
||||
headers=headers,
|
||||
params=requested_query_params,
|
||||
)
|
||||
elif HttpPassThroughEndpointHelpers.is_multipart(request) is True:
|
||||
elif HttpPassThroughEndpointHelpers.is_multipart(request) is True and not _parsed_body:
|
||||
# Only use multipart handler if we don't have a parsed body
|
||||
# (parsed body means it was JSON despite multipart content-type header)
|
||||
return await HttpPassThroughEndpointHelpers.make_multipart_http_request(
|
||||
request=request,
|
||||
async_client=async_client,
|
||||
@@ -677,8 +679,15 @@ async def pass_through_request( # noqa: PLR0915
|
||||
str(url)
|
||||
)
|
||||
|
||||
# Skip body parsing for multipart requests - make_multipart_http_request will handle it
|
||||
# But if custom_body is provided (e.g., JSON parsed despite multipart content-type), use it
|
||||
is_multipart = HttpPassThroughEndpointHelpers.is_multipart(request) and not custom_body
|
||||
|
||||
if custom_body:
|
||||
_parsed_body = custom_body
|
||||
elif is_multipart:
|
||||
# Don't parse multipart body here - it will be handled by make_multipart_http_request
|
||||
_parsed_body = {}
|
||||
else:
|
||||
_parsed_body = await _read_request_body(request)
|
||||
verbose_proxy_logger.debug(
|
||||
@@ -1043,30 +1052,22 @@ async def _parse_request_data_by_content_type(
|
||||
# Handle requests with no body (e.g., DELETE requests)
|
||||
pass
|
||||
elif "multipart/form-data" in content_type:
|
||||
# ✅ Handle multipart form-data
|
||||
form = await request.form()
|
||||
if "query_params" in form:
|
||||
form_value = form["query_params"]
|
||||
if isinstance(form_value, str):
|
||||
try:
|
||||
query_params_data = json.loads(form_value)
|
||||
except Exception:
|
||||
query_params_data = form_value
|
||||
else:
|
||||
query_params_data = form_value
|
||||
|
||||
if "custom_body" in form:
|
||||
form_value = form["custom_body"]
|
||||
if isinstance(form_value, str):
|
||||
try:
|
||||
custom_body_data = json.loads(form_value)
|
||||
except Exception:
|
||||
custom_body_data = form_value
|
||||
else:
|
||||
custom_body_data = form_value
|
||||
|
||||
if "file" in form:
|
||||
file_data = form["file"] # this is a Starlette UploadFile object
|
||||
# ✅ Try to parse as JSON first (handles misconfigured clients sending JSON with multipart content-type)
|
||||
# If that fails, skip parsing - pass_through_request will handle actual multipart
|
||||
try:
|
||||
body = await request.json()
|
||||
# Successfully parsed as JSON - treat as JSON body
|
||||
query_params_data = body.get("query_params")
|
||||
custom_body_data = body.get("custom_body")
|
||||
stream = body.get("stream")
|
||||
# If custom_body is not set, use the entire body
|
||||
if custom_body_data is None and body:
|
||||
custom_body_data = body
|
||||
except (json.JSONDecodeError, Exception):
|
||||
# Not JSON - this is actual multipart data
|
||||
# Skip parsing here to avoid consuming the request body stream
|
||||
# make_multipart_http_request will handle it
|
||||
pass
|
||||
|
||||
elif "application/x-www-form-urlencoded" in content_type:
|
||||
# ✅ Handle URL-encoded form data
|
||||
@@ -1132,7 +1133,6 @@ def create_pass_through_route(
|
||||
fastapi_response: Response,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
subpath: str = "", # captures sub-paths when include_subpath=True
|
||||
custom_body: Optional[dict] = None,
|
||||
):
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
InitPassThroughEndpointHelpers,
|
||||
@@ -1208,12 +1208,9 @@ def create_pass_through_route(
|
||||
)
|
||||
if query_params:
|
||||
final_query_params.update(query_params)
|
||||
# When a caller (e.g. bedrock_proxy_route) supplies a pre-built
|
||||
# body, use it instead of the body parsed from the raw request.
|
||||
# Use the body parsed from the raw request
|
||||
final_custom_body: Optional[dict] = None
|
||||
if custom_body is not None:
|
||||
final_custom_body = custom_body
|
||||
elif isinstance(custom_body_data, dict):
|
||||
if isinstance(custom_body_data, dict):
|
||||
final_custom_body = custom_body_data
|
||||
|
||||
return await pass_through_request( # type: ignore
|
||||
|
||||
@@ -853,7 +853,12 @@ def run_server( # noqa: PLR0915
|
||||
):
|
||||
check_prisma_schema_diff(db_url=None)
|
||||
else:
|
||||
PrismaManager.setup_database(use_migrate=not use_prisma_db_push)
|
||||
if not PrismaManager.setup_database(use_migrate=not use_prisma_db_push):
|
||||
print( # noqa
|
||||
"\033[1;31mLiteLLM Proxy: Database setup failed after multiple retries. "
|
||||
"The proxy cannot start safely. Please check your database connection and migration status.\033[0m"
|
||||
)
|
||||
sys.exit(1)
|
||||
else:
|
||||
print( # noqa
|
||||
f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa
|
||||
|
||||
@@ -557,6 +557,17 @@ class VertexAIBatchEmbeddingsResponseObject(TypedDict):
|
||||
embeddings: List[ContentEmbeddings]
|
||||
|
||||
|
||||
class GeminiEmbedContentRequestBody(TypedDict, total=False):
|
||||
content: Required[ContentType]
|
||||
taskType: TaskTypeEnum
|
||||
title: str
|
||||
outputDimensionality: int
|
||||
|
||||
|
||||
class GeminiEmbedContentResponseObject(TypedDict):
|
||||
embedding: ContentEmbeddings
|
||||
|
||||
|
||||
# Vertex AI Batch Prediction
|
||||
|
||||
|
||||
|
||||
@@ -95,6 +95,22 @@ class MCPCredentials(TypedDict, total=False):
|
||||
OAuth 2.0 scopes to request when exchanging the client credentials
|
||||
"""
|
||||
|
||||
# AWS SigV4 fields
|
||||
aws_access_key_id: Optional[str]
|
||||
"""AWS access key ID for SigV4 signing. Optional — falls back to boto3 credential chain."""
|
||||
|
||||
aws_secret_access_key: Optional[str]
|
||||
"""AWS secret access key for SigV4 signing. Optional — falls back to boto3 credential chain."""
|
||||
|
||||
aws_session_token: Optional[str]
|
||||
"""AWS session token for temporary STS credentials. Optional."""
|
||||
|
||||
aws_region_name: Optional[str]
|
||||
"""AWS region for SigV4 signing (e.g., 'us-east-1'). Not a secret — stored unencrypted."""
|
||||
|
||||
aws_service_name: Optional[str]
|
||||
"""AWS service name for SigV4 signing (e.g., 'bedrock-agentcore'). Not a secret — stored unencrypted."""
|
||||
|
||||
|
||||
class MCPServerCostInfo(TypedDict, total=False):
|
||||
default_cost_per_query: Optional[float]
|
||||
|
||||
@@ -253,6 +253,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
||||
tpm: Optional[int]
|
||||
rpm: Optional[int]
|
||||
provider_specific_entry: Optional[Dict[str, float]]
|
||||
uses_embed_content: Optional[bool]
|
||||
|
||||
|
||||
class ModelInfo(ModelInfoBase, total=False):
|
||||
|
||||
@@ -5779,6 +5779,7 @@ def _get_model_info_helper( # noqa: PLR0915
|
||||
provider_specific_entry=_model_info.get(
|
||||
"provider_specific_entry", None
|
||||
),
|
||||
uses_embed_content=_model_info.get("uses_embed_content", None),
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug(f"Error getting model info: {e}")
|
||||
|
||||
+12314
-12277
File diff suppressed because it is too large
Load Diff
@@ -15,8 +15,16 @@ from unittest.mock import MagicMock, patch
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import (
|
||||
_is_multimodal_input,
|
||||
_parse_data_url,
|
||||
process_embed_content_response,
|
||||
transform_openai_input_gemini_embed_content,
|
||||
)
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
|
||||
|
||||
def test_gemini_batch_embeddings_with_custom_api_base_and_auth_header():
|
||||
@@ -47,11 +55,9 @@ def test_gemini_batch_embeddings_with_custom_api_base_and_auth_header():
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"predictions": [
|
||||
"embeddings": [
|
||||
{
|
||||
"embeddings": {
|
||||
"values": [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
}
|
||||
"values": [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -109,11 +115,9 @@ def test_gemini_batch_embeddings_with_extra_headers():
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"predictions": [
|
||||
"embeddings": [
|
||||
{
|
||||
"embeddings": {
|
||||
"values": [0.1, 0.2, 0.3]
|
||||
}
|
||||
"values": [0.1, 0.2, 0.3]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -143,3 +147,380 @@ def test_gemini_batch_embeddings_with_extra_headers():
|
||||
assert "X-Custom" in headers
|
||||
assert headers["X-Custom"] == "custom-value"
|
||||
|
||||
|
||||
def test_is_multimodal_input_detection():
|
||||
"""Test that _is_multimodal_input correctly detects multimodal inputs."""
|
||||
assert _is_multimodal_input("plain text") is False
|
||||
assert _is_multimodal_input(["text1", "text2"]) is False
|
||||
|
||||
assert _is_multimodal_input("data:image/png;base64,iVBORw0KGgo=") is True
|
||||
assert _is_multimodal_input(["text", "data:image/png;base64,abc"]) is True
|
||||
|
||||
assert _is_multimodal_input("files/abc123") is True
|
||||
assert _is_multimodal_input(["text", "files/myfile"]) is True
|
||||
|
||||
|
||||
def test_parse_data_url():
|
||||
"""Test that _parse_data_url correctly extracts MIME type and base64 data."""
|
||||
mime_type, base64_data = _parse_data_url("data:image/png;base64,iVBORw0KGgo=")
|
||||
assert mime_type == "image/png"
|
||||
assert base64_data == "iVBORw0KGgo="
|
||||
|
||||
mime_type, base64_data = _parse_data_url("data:audio/mpeg;base64,SUQzBAA=")
|
||||
assert mime_type == "audio/mpeg"
|
||||
assert base64_data == "SUQzBAA="
|
||||
|
||||
mime_type, base64_data = _parse_data_url("data:video/mp4;base64,AAAAIGZ0eXA=")
|
||||
assert mime_type == "video/mp4"
|
||||
assert base64_data == "AAAAIGZ0eXA="
|
||||
|
||||
mime_type, base64_data = _parse_data_url("data:application/pdf;base64,JVBERi0=")
|
||||
assert mime_type == "application/pdf"
|
||||
assert base64_data == "JVBERi0="
|
||||
|
||||
|
||||
def test_mime_type_validation():
|
||||
"""Test that unsupported MIME types raise ValueError."""
|
||||
with pytest.raises(ValueError, match="Unsupported MIME type"):
|
||||
_parse_data_url("data:text/plain;base64,SGVsbG8=")
|
||||
|
||||
with pytest.raises(ValueError, match="Unsupported MIME type"):
|
||||
_parse_data_url("data:application/json;base64,e30=")
|
||||
|
||||
|
||||
def test_parse_data_url_invalid_format():
|
||||
"""Test that invalid data URL formats raise ValueError."""
|
||||
with pytest.raises(ValueError, match="Invalid data URL format"):
|
||||
_parse_data_url("not-a-data-url")
|
||||
|
||||
with pytest.raises(ValueError, match="missing comma"):
|
||||
_parse_data_url("data:image/png;base64")
|
||||
|
||||
|
||||
def test_transform_multimodal_text_and_image():
|
||||
"""Test transformation of mixed text and image input."""
|
||||
input_data = [
|
||||
"The food was delicious",
|
||||
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
|
||||
]
|
||||
|
||||
result = transform_openai_input_gemini_embed_content(
|
||||
input=input_data,
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
resolved_files=None,
|
||||
)
|
||||
|
||||
assert "content" in result
|
||||
assert "parts" in result["content"]
|
||||
parts = result["content"]["parts"]
|
||||
|
||||
assert len(parts) == 2
|
||||
assert parts[0]["text"] == "The food was delicious"
|
||||
assert "inline_data" in parts[1]
|
||||
assert parts[1]["inline_data"]["mime_type"] == "image/png"
|
||||
assert "data" in parts[1]["inline_data"]
|
||||
|
||||
|
||||
def test_transform_multimodal_with_file_reference():
|
||||
"""Test transformation with Gemini file reference."""
|
||||
input_data = ["Some text", "files/abc123"]
|
||||
|
||||
resolved_files = {
|
||||
"files/abc123": {
|
||||
"mime_type": "image/jpeg",
|
||||
"uri": "https://generativelanguage.googleapis.com/v1beta/files/abc123"
|
||||
}
|
||||
}
|
||||
|
||||
result = transform_openai_input_gemini_embed_content(
|
||||
input=input_data,
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
resolved_files=resolved_files,
|
||||
)
|
||||
|
||||
assert "content" in result
|
||||
parts = result["content"]["parts"]
|
||||
|
||||
assert len(parts) == 2
|
||||
assert parts[0]["text"] == "Some text"
|
||||
assert "file_data" in parts[1]
|
||||
assert parts[1]["file_data"]["mime_type"] == "image/jpeg"
|
||||
assert parts[1]["file_data"]["file_uri"] == "https://generativelanguage.googleapis.com/v1beta/files/abc123"
|
||||
|
||||
|
||||
def test_embed_content_response_processing():
|
||||
"""Test processing of embedContent response (single embedding)."""
|
||||
response_json = {
|
||||
"embedding": {
|
||||
"values": [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
}
|
||||
}
|
||||
|
||||
model_response = EmbeddingResponse()
|
||||
result = process_embed_content_response(
|
||||
input=["test input"],
|
||||
model_response=model_response,
|
||||
model="gemini-embedding-2-preview",
|
||||
response_json=response_json,
|
||||
)
|
||||
|
||||
assert len(result.data) == 1
|
||||
assert result.data[0].embedding == [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
assert result.data[0].index == 0
|
||||
assert result.data[0].object == "embedding"
|
||||
assert result.model == "gemini-embedding-2-preview"
|
||||
assert result.usage.prompt_tokens > 0
|
||||
|
||||
|
||||
def test_embed_content_response_multimodal_sets_prompt_tokens_zero():
|
||||
"""Test that multimodal input sets prompt_tokens=0 (cannot accurately count)."""
|
||||
response_json = {
|
||||
"embedding": {
|
||||
"values": [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
}
|
||||
}
|
||||
|
||||
model_response = EmbeddingResponse()
|
||||
result = process_embed_content_response(
|
||||
input=["text", "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="],
|
||||
model_response=model_response,
|
||||
model="gemini-embedding-2-preview",
|
||||
response_json=response_json,
|
||||
)
|
||||
|
||||
assert result.usage.prompt_tokens == 0
|
||||
|
||||
|
||||
def test_gemini_multimodal_embedding_e2e():
|
||||
"""Test end-to-end multimodal embedding call through litellm.embedding()."""
|
||||
client = HTTPHandler()
|
||||
|
||||
def mock_auth_token(*args, **kwargs):
|
||||
return None, "test-project"
|
||||
|
||||
with patch.object(client, "post") as mock_post, patch(
|
||||
"litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._ensure_access_token",
|
||||
side_effect=mock_auth_token
|
||||
), patch(
|
||||
"litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._get_token_and_url"
|
||||
) as mock_get_token:
|
||||
mock_get_token.return_value = (
|
||||
{"x-goog-api-key": "test-key"},
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent?key=test-key"
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"embedding": {
|
||||
"values": [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
}
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
response = litellm.embedding(
|
||||
model="gemini/gemini-embedding-2-preview",
|
||||
input=["The food was delicious", "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="],
|
||||
api_key="test-key",
|
||||
client=client
|
||||
)
|
||||
|
||||
mock_post.assert_called_once()
|
||||
|
||||
call_args = mock_post.call_args
|
||||
kwargs = call_args.kwargs if hasattr(call_args, 'kwargs') else call_args[1]
|
||||
|
||||
request_body = json.loads(kwargs.get("data", "{}"))
|
||||
|
||||
assert "content" in request_body
|
||||
assert "parts" in request_body["content"]
|
||||
parts = request_body["content"]["parts"]
|
||||
|
||||
assert len(parts) == 2
|
||||
assert parts[0]["text"] == "The food was delicious"
|
||||
assert "inline_data" in parts[1]
|
||||
assert parts[1]["inline_data"]["mime_type"] == "image/png"
|
||||
|
||||
assert len(response.data) == 1
|
||||
assert response.data[0].embedding == [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
|
||||
|
||||
def test_gemini_multimodal_embedding_with_audio():
|
||||
"""Test multimodal embedding with audio input."""
|
||||
input_data = ["Audio description", "data:audio/mpeg;base64,SUQzBAAAAAA="]
|
||||
|
||||
result = transform_openai_input_gemini_embed_content(
|
||||
input=input_data,
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
resolved_files=None,
|
||||
)
|
||||
|
||||
parts = result["content"]["parts"]
|
||||
assert len(parts) == 2
|
||||
assert parts[0]["text"] == "Audio description"
|
||||
assert parts[1]["inline_data"]["mime_type"] == "audio/mpeg"
|
||||
|
||||
|
||||
def test_gemini_multimodal_embedding_with_video():
|
||||
"""Test multimodal embedding with video input."""
|
||||
input_data = ["data:video/mp4;base64,AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAAIZnJlZQAA"]
|
||||
|
||||
result = transform_openai_input_gemini_embed_content(
|
||||
input=input_data,
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
resolved_files=None,
|
||||
)
|
||||
|
||||
parts = result["content"]["parts"]
|
||||
assert len(parts) == 1
|
||||
assert parts[0]["inline_data"]["mime_type"] == "video/mp4"
|
||||
|
||||
|
||||
|
||||
def test_transform_with_optional_params():
|
||||
"""Test that optional params like outputDimensionality are passed through."""
|
||||
input_data = ["test text"]
|
||||
|
||||
result = transform_openai_input_gemini_embed_content(
|
||||
input=input_data,
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={"outputDimensionality": 768, "taskType": "SEMANTIC_SIMILARITY"},
|
||||
resolved_files=None,
|
||||
)
|
||||
|
||||
assert result["outputDimensionality"] == 768
|
||||
assert result["taskType"] == "SEMANTIC_SIMILARITY"
|
||||
|
||||
|
||||
def test_dimensions_mapped_to_output_dimensionality():
|
||||
"""Test that OpenAI 'dimensions' param is mapped to Gemini 'outputDimensionality'."""
|
||||
input_data = ["test text"]
|
||||
|
||||
result = transform_openai_input_gemini_embed_content(
|
||||
input=input_data,
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={"dimensions": 768},
|
||||
resolved_files=None,
|
||||
)
|
||||
|
||||
assert "outputDimensionality" in result
|
||||
assert result["outputDimensionality"] == 768
|
||||
assert "dimensions" not in result
|
||||
|
||||
|
||||
def test_is_gcs_url():
|
||||
"""Test GCS URL detection."""
|
||||
from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import (
|
||||
_is_gcs_url,
|
||||
)
|
||||
|
||||
assert _is_gcs_url("gs://my-bucket/path/to/file.png") is True
|
||||
assert _is_gcs_url("gs://bucket/image.jpg") is True
|
||||
assert _is_gcs_url("https://storage.googleapis.com/bucket/file.png") is False
|
||||
assert _is_gcs_url("files/abc123") is False
|
||||
assert _is_gcs_url("data:image/png;base64,abc") is False
|
||||
assert _is_gcs_url("regular text") is False
|
||||
|
||||
|
||||
def test_infer_mime_type_from_gcs_url():
|
||||
"""Test MIME type inference from GCS URL."""
|
||||
from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import (
|
||||
_infer_mime_type_from_gcs_url,
|
||||
)
|
||||
|
||||
assert _infer_mime_type_from_gcs_url("gs://bucket/image.png") == "image/png"
|
||||
assert _infer_mime_type_from_gcs_url("gs://bucket/photo.jpg") == "image/jpeg"
|
||||
assert _infer_mime_type_from_gcs_url("gs://bucket/photo.JPEG") == "image/jpeg"
|
||||
assert _infer_mime_type_from_gcs_url("gs://bucket/audio.mp3") == "audio/mpeg"
|
||||
assert _infer_mime_type_from_gcs_url("gs://bucket/audio.wav") == "audio/wav"
|
||||
assert _infer_mime_type_from_gcs_url("gs://bucket/video.mp4") == "video/mp4"
|
||||
assert _infer_mime_type_from_gcs_url("gs://bucket/video.mov") == "video/quicktime"
|
||||
assert _infer_mime_type_from_gcs_url("gs://bucket/doc.pdf") == "application/pdf"
|
||||
|
||||
with pytest.raises(ValueError, match="Unable to infer MIME type"):
|
||||
_infer_mime_type_from_gcs_url("gs://bucket/file.txt")
|
||||
|
||||
|
||||
def test_transform_multimodal_with_gcs_url():
|
||||
"""Test transformation with GCS URL."""
|
||||
input_data = [
|
||||
"Describe this image",
|
||||
"gs://my-bucket/images/photo.png"
|
||||
]
|
||||
|
||||
result = transform_openai_input_gemini_embed_content(
|
||||
input=input_data,
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
resolved_files=None,
|
||||
)
|
||||
|
||||
parts = result["content"]["parts"]
|
||||
assert len(parts) == 2
|
||||
assert parts[0]["text"] == "Describe this image"
|
||||
assert parts[1]["file_data"]["mime_type"] == "image/png"
|
||||
assert parts[1]["file_data"]["file_uri"] == "gs://my-bucket/images/photo.png"
|
||||
|
||||
|
||||
def test_multimodal_input_detection_with_gcs():
|
||||
"""Test that GCS URLs are detected as multimodal."""
|
||||
from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import (
|
||||
_is_multimodal_input,
|
||||
)
|
||||
|
||||
assert _is_multimodal_input(["text", "gs://bucket/file.png"]) is True
|
||||
assert _is_multimodal_input("gs://bucket/video.mp4") is True
|
||||
assert _is_multimodal_input(["just text", "more text"]) is False
|
||||
|
||||
|
||||
def test_vertex_ai_text_only_embedding_uses_embed_content():
|
||||
"""
|
||||
Test that vertex_ai/gemini-embedding-2-preview with text-only input uses
|
||||
embedContent endpoint (not batchEmbedContents) and returns a single embedding.
|
||||
"""
|
||||
client = HTTPHandler()
|
||||
embed_content_url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test/locations/us-central1/publishers/google/models/gemini-embedding-2-preview:embedContent"
|
||||
|
||||
def mock_auth_token(*args, **kwargs):
|
||||
return "Bearer test-token", "test-project"
|
||||
|
||||
with patch.object(client, "post") as mock_post, patch(
|
||||
"litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._ensure_access_token",
|
||||
side_effect=mock_auth_token,
|
||||
), patch(
|
||||
"litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._get_token_and_url"
|
||||
) as mock_get_token:
|
||||
mock_get_token.return_value = (
|
||||
{"Authorization": "Bearer test-token"},
|
||||
embed_content_url,
|
||||
)
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"embedding": {"values": [0.1, 0.2, 0.3, 0.4, 0.5]}
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
response = litellm.embedding(
|
||||
model="vertex_ai/gemini-embedding-2-preview",
|
||||
input=["Hello, world!"],
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
client=client,
|
||||
)
|
||||
|
||||
mock_post.assert_called_once()
|
||||
call_args = mock_post.call_args
|
||||
post_url = call_args.kwargs.get("url", call_args.args[0] if call_args.args else "")
|
||||
assert "embedContent" in str(post_url)
|
||||
data = json.loads(call_args.kwargs["data"])
|
||||
assert "content" in data
|
||||
assert "parts" in data["content"]
|
||||
assert len(data["content"]["parts"]) == 1
|
||||
assert data["content"]["parts"][0]["text"] == "Hello, world!"
|
||||
assert len(response.data) == 1
|
||||
|
||||
|
||||
@@ -2,11 +2,14 @@
|
||||
Tests for AWS SigV4 authentication in MCP client.
|
||||
|
||||
Tests the MCPSigV4Auth httpx.Auth subclass that enables per-request
|
||||
SigV4 signing for Bedrock AgentCore MCP servers.
|
||||
SigV4 signing for Bedrock AgentCore MCP servers, plus DB/UI path
|
||||
tests for credential encryption, merge-on-update, and build_from_table.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock
|
||||
from unittest.mock import patch, MagicMock, AsyncMock
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -315,3 +318,568 @@ class TestMCPServerManagerSigV4:
|
||||
client = await manager._create_mcp_client(server=server)
|
||||
|
||||
assert client._aws_auth is None
|
||||
|
||||
|
||||
class TestSigV4CredentialEncryption:
|
||||
"""Test encrypt/decrypt round-trip for AWS SigV4 credentials."""
|
||||
|
||||
def test_encrypt_credentials_handles_aws_fields(self):
|
||||
"""AWS credential fields are encrypted in the credentials dict."""
|
||||
from litellm.proxy._experimental.mcp_server.db import encrypt_credentials
|
||||
|
||||
creds = {
|
||||
"aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
|
||||
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
||||
"aws_session_token": "FwoGZX...",
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_service_name": "bedrock-agentcore",
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.encrypt_value_helper",
|
||||
side_effect=lambda value, new_encryption_key: f"enc:{value}",
|
||||
):
|
||||
result = encrypt_credentials(credentials=creds, encryption_key="test-key")
|
||||
|
||||
# Secrets should be encrypted
|
||||
assert result["aws_access_key_id"] == "enc:AKIAIOSFODNN7EXAMPLE"
|
||||
assert (
|
||||
result["aws_secret_access_key"]
|
||||
== "enc:wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
|
||||
)
|
||||
assert result["aws_session_token"] == "enc:FwoGZX..."
|
||||
# Non-secrets should be unchanged
|
||||
assert result["aws_region_name"] == "us-east-1"
|
||||
assert result["aws_service_name"] == "bedrock-agentcore"
|
||||
|
||||
def test_encrypt_credentials_skips_absent_aws_fields(self):
|
||||
"""encrypt_credentials does not fail when AWS fields are absent."""
|
||||
from litellm.proxy._experimental.mcp_server.db import encrypt_credentials
|
||||
|
||||
creds = {"auth_value": "some-token"}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.encrypt_value_helper",
|
||||
side_effect=lambda value, new_encryption_key: f"enc:{value}",
|
||||
):
|
||||
result = encrypt_credentials(credentials=creds, encryption_key="test-key")
|
||||
|
||||
assert result["auth_value"] == "enc:some-token"
|
||||
assert "aws_access_key_id" not in result
|
||||
|
||||
|
||||
class TestCredentialMergeOnUpdate:
|
||||
"""Test that partial credential updates preserve existing fields."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_partial_update_preserves_existing_credentials(self):
|
||||
"""Updating only aws_region_name should not wipe aws_secret_access_key."""
|
||||
from litellm.proxy._experimental.mcp_server.db import update_mcp_server
|
||||
from litellm.proxy._types import UpdateMCPServerRequest
|
||||
|
||||
existing_record = MagicMock()
|
||||
existing_record.auth_type = "aws_sigv4"
|
||||
existing_record.credentials = json.dumps(
|
||||
{
|
||||
"aws_access_key_id": "enc:AKI",
|
||||
"aws_secret_access_key": "enc:SAK",
|
||||
"aws_region_name": "us-east-1",
|
||||
}
|
||||
)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(
|
||||
return_value=existing_record
|
||||
)
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(
|
||||
return_value=MagicMock()
|
||||
)
|
||||
|
||||
data = UpdateMCPServerRequest(
|
||||
server_id="test-server",
|
||||
auth_type="aws_sigv4",
|
||||
credentials={"aws_region_name": "eu-west-1"},
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._get_salt_key",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.encrypt_value_helper",
|
||||
side_effect=lambda value, new_encryption_key: value,
|
||||
):
|
||||
await update_mcp_server(mock_prisma, data, "test-user")
|
||||
|
||||
# Grab the data dict passed to prisma update
|
||||
update_call = mock_prisma.db.litellm_mcpservertable.update
|
||||
assert update_call.called
|
||||
data_dict = update_call.call_args[1]["data"]
|
||||
merged_creds = json.loads(data_dict["credentials"])
|
||||
|
||||
# Existing encrypted secrets should be preserved
|
||||
assert merged_creds["aws_access_key_id"] == "enc:AKI"
|
||||
assert merged_creds["aws_secret_access_key"] == "enc:SAK"
|
||||
# New region value should be updated
|
||||
assert merged_creds["aws_region_name"] == "eu-west-1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_without_credentials_preserves_all(self):
|
||||
"""Update with no credentials field should not touch existing credentials."""
|
||||
from litellm.proxy._experimental.mcp_server.db import update_mcp_server
|
||||
from litellm.proxy._types import UpdateMCPServerRequest
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(
|
||||
return_value=MagicMock()
|
||||
)
|
||||
|
||||
data = UpdateMCPServerRequest(
|
||||
server_id="test-server",
|
||||
description="Updated description",
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._get_salt_key",
|
||||
return_value=None,
|
||||
):
|
||||
await update_mcp_server(mock_prisma, data, "test-user")
|
||||
|
||||
data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"]
|
||||
assert "credentials" not in data_dict
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_new_server_no_merge(self):
|
||||
"""Update with credentials on a server that has no existing credentials."""
|
||||
from litellm.proxy._experimental.mcp_server.db import update_mcp_server
|
||||
from litellm.proxy._types import UpdateMCPServerRequest
|
||||
|
||||
existing_record = MagicMock()
|
||||
existing_record.auth_type = "aws_sigv4"
|
||||
existing_record.credentials = None
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(
|
||||
return_value=existing_record
|
||||
)
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(
|
||||
return_value=MagicMock()
|
||||
)
|
||||
|
||||
data = UpdateMCPServerRequest(
|
||||
server_id="test-server",
|
||||
auth_type="aws_sigv4",
|
||||
credentials={"aws_region_name": "us-east-1"},
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._get_salt_key",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.encrypt_value_helper",
|
||||
side_effect=lambda value, new_encryption_key: value,
|
||||
):
|
||||
await update_mcp_server(mock_prisma, data, "test-user")
|
||||
|
||||
data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"]
|
||||
stored_creds = json.loads(data_dict["credentials"])
|
||||
assert stored_creds == {"aws_region_name": "us-east-1"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_type_change_replaces_credentials_entirely(self):
|
||||
"""Switching auth_type should replace credentials, not merge."""
|
||||
from litellm.proxy._experimental.mcp_server.db import update_mcp_server
|
||||
from litellm.proxy._types import UpdateMCPServerRequest
|
||||
|
||||
existing_record = MagicMock()
|
||||
existing_record.auth_type = "aws_sigv4"
|
||||
existing_record.credentials = json.dumps(
|
||||
{
|
||||
"aws_access_key_id": "enc:AKI",
|
||||
"aws_secret_access_key": "enc:SAK",
|
||||
"aws_region_name": "us-east-1",
|
||||
}
|
||||
)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(
|
||||
return_value=existing_record
|
||||
)
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(
|
||||
return_value=MagicMock()
|
||||
)
|
||||
|
||||
data = UpdateMCPServerRequest(
|
||||
server_id="test-server",
|
||||
auth_type="api_key",
|
||||
credentials={"auth_value": "my-key"},
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._get_salt_key",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.encrypt_value_helper",
|
||||
side_effect=lambda value, new_encryption_key: f"enc:{value}",
|
||||
):
|
||||
await update_mcp_server(mock_prisma, data, "test-user")
|
||||
|
||||
data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"]
|
||||
stored_creds = json.loads(data_dict["credentials"])
|
||||
# Should only have the new api_key credential, no stale aws_* fields
|
||||
assert stored_creds == {"auth_value": "enc:my-key"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_auth_type_merges_credentials(self):
|
||||
"""Same auth_type should merge credentials (preserve untouched fields)."""
|
||||
from litellm.proxy._experimental.mcp_server.db import update_mcp_server
|
||||
from litellm.proxy._types import UpdateMCPServerRequest
|
||||
|
||||
existing_record = MagicMock()
|
||||
existing_record.auth_type = "oauth2"
|
||||
existing_record.credentials = json.dumps(
|
||||
{
|
||||
"client_id": "enc:id",
|
||||
"client_secret": "enc:secret",
|
||||
"scopes": ["read"],
|
||||
}
|
||||
)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(
|
||||
return_value=existing_record
|
||||
)
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(
|
||||
return_value=MagicMock()
|
||||
)
|
||||
|
||||
data = UpdateMCPServerRequest(
|
||||
server_id="test-server",
|
||||
auth_type="oauth2",
|
||||
credentials={"scopes": ["read", "write"]},
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._get_salt_key",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.encrypt_value_helper",
|
||||
side_effect=lambda value, new_encryption_key: value,
|
||||
):
|
||||
await update_mcp_server(mock_prisma, data, "test-user")
|
||||
|
||||
data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"]
|
||||
merged_creds = json.loads(data_dict["credentials"])
|
||||
assert merged_creds["client_id"] == "enc:id"
|
||||
assert merged_creds["client_secret"] == "enc:secret"
|
||||
assert merged_creds["scopes"] == ["read", "write"]
|
||||
|
||||
|
||||
class TestSigV4BuildFromTable:
|
||||
"""Test build_mcp_server_from_table correctly loads AWS SigV4 credentials."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_mcp_server_from_table_with_sigv4_credentials(self):
|
||||
"""SigV4 credentials from DB are decrypted and mapped to MCPServer fields."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
|
||||
table_record = MagicMock()
|
||||
table_record.server_id = "test-sigv4-server"
|
||||
table_record.server_name = "sigv4_server"
|
||||
table_record.alias = None
|
||||
table_record.description = None
|
||||
table_record.url = "https://bedrock-agentcore.us-east-1.amazonaws.com/invocations"
|
||||
table_record.spec_path = None
|
||||
table_record.transport = "http"
|
||||
table_record.auth_type = "aws_sigv4"
|
||||
table_record.mcp_info = {"server_name": "sigv4_server"}
|
||||
table_record.credentials = json.dumps(
|
||||
{
|
||||
"aws_access_key_id": "enc:AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "enc:SECRET",
|
||||
"aws_session_token": "enc:TOKEN",
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_service_name": "bedrock-agentcore",
|
||||
}
|
||||
)
|
||||
table_record.extra_headers = None
|
||||
table_record.static_headers = None
|
||||
table_record.command = None
|
||||
table_record.args = []
|
||||
table_record.env = None
|
||||
table_record.mcp_access_groups = []
|
||||
table_record.allowed_tools = []
|
||||
table_record.disallowed_tools = None
|
||||
table_record.allow_all_keys = False
|
||||
table_record.available_on_public_internet = True
|
||||
table_record.authorization_url = None
|
||||
table_record.token_url = None
|
||||
table_record.registration_url = None
|
||||
table_record.created_at = None
|
||||
table_record.updated_at = None
|
||||
table_record.client_id = None
|
||||
table_record.client_secret = None
|
||||
table_record.tool_name_to_display_name = None
|
||||
table_record.tool_name_to_description = None
|
||||
table_record.byok_api_key_help_url = None
|
||||
table_record.oauth2_flow = None
|
||||
|
||||
manager = MCPServerManager()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.decrypt_value_helper",
|
||||
side_effect=lambda value, key, exception_type, return_original_value: value.replace(
|
||||
"enc:", ""
|
||||
),
|
||||
):
|
||||
server = await manager.build_mcp_server_from_table(table_record)
|
||||
|
||||
assert server.auth_type == "aws_sigv4"
|
||||
assert server.aws_access_key_id == "AKIAEXAMPLE"
|
||||
assert server.aws_secret_access_key == "SECRET"
|
||||
assert server.aws_session_token == "TOKEN"
|
||||
assert server.aws_region_name == "us-east-1"
|
||||
assert server.aws_service_name == "bedrock-agentcore"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_mcp_server_from_table_without_sigv4_credentials(self):
|
||||
"""Non-SigV4 servers still work — AWS fields default to None."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
|
||||
table_record = MagicMock()
|
||||
table_record.server_id = "test-bearer-server"
|
||||
table_record.server_name = "bearer_server"
|
||||
table_record.alias = None
|
||||
table_record.description = None
|
||||
table_record.url = "https://example.com/mcp"
|
||||
table_record.spec_path = None
|
||||
table_record.transport = "http"
|
||||
table_record.auth_type = "bearer_token"
|
||||
table_record.mcp_info = {"server_name": "bearer_server"}
|
||||
table_record.credentials = json.dumps({"auth_value": "enc:tok"})
|
||||
table_record.extra_headers = None
|
||||
table_record.static_headers = None
|
||||
table_record.command = None
|
||||
table_record.args = []
|
||||
table_record.env = None
|
||||
table_record.mcp_access_groups = []
|
||||
table_record.allowed_tools = []
|
||||
table_record.disallowed_tools = None
|
||||
table_record.allow_all_keys = False
|
||||
table_record.available_on_public_internet = True
|
||||
table_record.authorization_url = None
|
||||
table_record.token_url = None
|
||||
table_record.registration_url = None
|
||||
table_record.created_at = None
|
||||
table_record.updated_at = None
|
||||
table_record.client_id = None
|
||||
table_record.client_secret = None
|
||||
table_record.tool_name_to_display_name = None
|
||||
table_record.tool_name_to_description = None
|
||||
table_record.byok_api_key_help_url = None
|
||||
table_record.oauth2_flow = None
|
||||
|
||||
manager = MCPServerManager()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.decrypt_value_helper",
|
||||
side_effect=lambda value, key, exception_type, return_original_value: value.replace(
|
||||
"enc:", ""
|
||||
),
|
||||
):
|
||||
server = await manager.build_mcp_server_from_table(table_record)
|
||||
|
||||
assert server.auth_type == "bearer_token"
|
||||
assert server.aws_access_key_id is None
|
||||
assert server.aws_secret_access_key is None
|
||||
assert server.aws_session_token is None
|
||||
assert server.aws_region_name is None
|
||||
assert server.aws_service_name is None
|
||||
|
||||
|
||||
class TestDecryptCredentials:
|
||||
"""Test decrypt_credentials helper."""
|
||||
|
||||
def test_decrypt_credentials_handles_all_secret_fields(self):
|
||||
"""All secret fields are decrypted; non-secret fields are left as-is."""
|
||||
from litellm.proxy._experimental.mcp_server.db import decrypt_credentials
|
||||
|
||||
creds = {
|
||||
"auth_value": "enc:tok",
|
||||
"client_id": "enc:cid",
|
||||
"client_secret": "enc:csec",
|
||||
"aws_access_key_id": "enc:AKI",
|
||||
"aws_secret_access_key": "enc:SAK",
|
||||
"aws_session_token": "enc:TOK",
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_service_name": "bedrock-agentcore",
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.decrypt_value_helper",
|
||||
side_effect=lambda value, key, exception_type="error", return_original_value=False: value.replace("enc:", ""),
|
||||
):
|
||||
result = decrypt_credentials(credentials=creds)
|
||||
|
||||
assert result["auth_value"] == "tok"
|
||||
assert result["client_id"] == "cid"
|
||||
assert result["client_secret"] == "csec"
|
||||
assert result["aws_access_key_id"] == "AKI"
|
||||
assert result["aws_secret_access_key"] == "SAK"
|
||||
assert result["aws_session_token"] == "TOK"
|
||||
# Non-secrets untouched
|
||||
assert result["aws_region_name"] == "us-east-1"
|
||||
assert result["aws_service_name"] == "bedrock-agentcore"
|
||||
|
||||
def test_decrypt_credentials_skips_absent_fields(self):
|
||||
"""Absent fields are not touched."""
|
||||
from litellm.proxy._experimental.mcp_server.db import decrypt_credentials
|
||||
|
||||
creds = {"aws_access_key_id": "enc:AKI"}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.decrypt_value_helper",
|
||||
side_effect=lambda value, key, exception_type="error", return_original_value=False: value.replace("enc:", ""),
|
||||
):
|
||||
result = decrypt_credentials(credentials=creds)
|
||||
|
||||
assert result["aws_access_key_id"] == "AKI"
|
||||
assert "aws_secret_access_key" not in result
|
||||
|
||||
|
||||
class TestRotateCredentials:
|
||||
"""Test rotate_mcp_server_credentials_master_key decrypts before re-encrypting."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rotation_decrypts_then_reencrypts(self):
|
||||
"""Key rotation should decrypt with old key then encrypt with new key."""
|
||||
from litellm.proxy._experimental.mcp_server.db import (
|
||||
rotate_mcp_server_credentials_master_key,
|
||||
)
|
||||
|
||||
server = MagicMock()
|
||||
server.server_id = "srv-1"
|
||||
server.credentials = {
|
||||
"aws_access_key_id": "enc_old:AKI",
|
||||
"aws_secret_access_key": "enc_old:SAK",
|
||||
"aws_region_name": "us-east-1",
|
||||
}
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(
|
||||
return_value=[server]
|
||||
)
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._get_salt_key",
|
||||
return_value="old-key",
|
||||
), patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.decrypt_value_helper",
|
||||
side_effect=lambda value, key, exception_type="error", return_original_value=False: value.replace("enc_old:", ""),
|
||||
), patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.encrypt_value_helper",
|
||||
side_effect=lambda value, new_encryption_key: f"enc_new:{value}",
|
||||
):
|
||||
await rotate_mcp_server_credentials_master_key(
|
||||
mock_prisma, "admin", "new-key"
|
||||
)
|
||||
|
||||
update_call = mock_prisma.db.litellm_mcpservertable.update
|
||||
assert update_call.called
|
||||
stored_creds = json.loads(update_call.call_args[1]["data"]["credentials"])
|
||||
# Should be decrypted from old, then encrypted with new
|
||||
assert stored_creds["aws_access_key_id"] == "enc_new:AKI"
|
||||
assert stored_creds["aws_secret_access_key"] == "enc_new:SAK"
|
||||
# Non-secret fields should pass through unchanged
|
||||
assert stored_creds["aws_region_name"] == "us-east-1"
|
||||
|
||||
|
||||
class TestAuthTypeSwitchClearsCredentials:
|
||||
"""Test that switching auth_type without credentials clears stale secrets."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_type_change_without_credentials_clears_stale(self):
|
||||
"""Changing auth_type without providing credentials should clear old ones."""
|
||||
from litellm.proxy._experimental.mcp_server.db import update_mcp_server
|
||||
from litellm.proxy._types import UpdateMCPServerRequest
|
||||
|
||||
existing_record = MagicMock()
|
||||
existing_record.auth_type = "oauth2"
|
||||
existing_record.credentials = json.dumps(
|
||||
{"client_id": "enc:cid", "client_secret": "enc:csec"}
|
||||
)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(
|
||||
return_value=existing_record
|
||||
)
|
||||
mock_prisma.db.litellm_mcpservertable.update = AsyncMock(
|
||||
return_value=MagicMock()
|
||||
)
|
||||
|
||||
data = UpdateMCPServerRequest(
|
||||
server_id="test-server",
|
||||
auth_type="aws_sigv4",
|
||||
# No credentials provided
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._get_salt_key",
|
||||
return_value=None,
|
||||
):
|
||||
await update_mcp_server(mock_prisma, data, "test-user")
|
||||
|
||||
data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"]
|
||||
# Credentials should be cleared (set to None)
|
||||
assert data_dict.get("credentials") is None
|
||||
|
||||
|
||||
class TestInheritCredentials:
|
||||
"""Test _inherit_credentials_from_existing_server copies AWS fields."""
|
||||
|
||||
def test_inherits_sigv4_credentials(self):
|
||||
"""SigV4 fields are copied from existing server to inherited credentials."""
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
|
||||
_inherit_credentials_from_existing_server,
|
||||
)
|
||||
from litellm.proxy._types import NewMCPServerRequest
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
existing = MCPServer(
|
||||
server_id="existing-sigv4",
|
||||
name="sigv4_server",
|
||||
server_name="sigv4_server",
|
||||
url="https://bedrock-agentcore.us-east-1.amazonaws.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.aws_sigv4,
|
||||
aws_access_key_id="AKIAEXAMPLE",
|
||||
aws_secret_access_key="SECRET",
|
||||
aws_session_token="TOKEN",
|
||||
aws_region_name="us-east-1",
|
||||
aws_service_name="bedrock-agentcore",
|
||||
)
|
||||
|
||||
payload = NewMCPServerRequest(
|
||||
server_id="existing-sigv4",
|
||||
server_name="sigv4_server",
|
||||
url="https://bedrock-agentcore.us-east-1.amazonaws.com/mcp",
|
||||
transport="http",
|
||||
auth_type="aws_sigv4",
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager"
|
||||
) as mock_manager:
|
||||
mock_manager.get_mcp_server_by_id.return_value = existing
|
||||
result = _inherit_credentials_from_existing_server(payload)
|
||||
|
||||
assert result.credentials is not None
|
||||
assert result.credentials["aws_access_key_id"] == "AKIAEXAMPLE"
|
||||
assert result.credentials["aws_secret_access_key"] == "SECRET"
|
||||
assert result.credentials["aws_session_token"] == "TOKEN"
|
||||
assert result.credentials["aws_region_name"] == "us-east-1"
|
||||
assert result.credentials["aws_service_name"] == "bedrock-agentcore"
|
||||
|
||||
@@ -2369,3 +2369,107 @@ def test_get_registered_pass_through_route_with_custom_root():
|
||||
|
||||
# Clean up
|
||||
_registered_pass_through_routes.clear()
|
||||
|
||||
|
||||
def test_mapped_pass_through_routes_with_server_root_path():
|
||||
"""
|
||||
Mapped passthrough routes (vertex_ai, bedrock, etc) should match
|
||||
even when SERVER_ROOT_PATH is set and the incoming route is prefixed.
|
||||
|
||||
Regression test for https://github.com/BerriAI/litellm/issues/22272
|
||||
"""
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
InitPassThroughEndpointHelpers,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path"
|
||||
) as mock_get_root:
|
||||
mock_get_root.return_value = "/litellm"
|
||||
|
||||
# prefixed route should match mapped routes like /vertex_ai
|
||||
assert (
|
||||
InitPassThroughEndpointHelpers.is_registered_pass_through_route(
|
||||
"/litellm/vertex_ai/v1/projects/foo"
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
InitPassThroughEndpointHelpers.is_registered_pass_through_route(
|
||||
"/litellm/bedrock/model/invoke"
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
# bare route without prefix should not match when root is set
|
||||
assert (
|
||||
InitPassThroughEndpointHelpers.is_registered_pass_through_route(
|
||||
"/vertex_ai/v1/projects/foo"
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multipart_passthrough_preserves_boundary():
|
||||
"""
|
||||
Test that multipart/form-data requests through passthrough preserve the boundary
|
||||
and can be correctly parsed by the upstream server.
|
||||
|
||||
Regression test for multipart boundary stripping issue.
|
||||
"""
|
||||
from io import BytesIO
|
||||
|
||||
# Mock the httpx request to verify files are passed correctly
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = httpx.Headers({"content-type": "application/json"})
|
||||
mock_response.aread = AsyncMock(return_value=b'{"filename": "test.txt", "size": 17}')
|
||||
mock_response.text = '{"filename": "test.txt", "size": 17}'
|
||||
|
||||
async def mock_httpx_request(method, url, **kwargs):
|
||||
# Verify that files parameter is passed (not json)
|
||||
assert "files" in kwargs, "Files should be passed for multipart requests"
|
||||
assert "file" in kwargs["files"], "File field should be in files dict"
|
||||
|
||||
# Verify content-type is NOT in headers (httpx will set it with correct boundary)
|
||||
headers = kwargs.get("headers", {})
|
||||
assert "content-type" not in headers, "content-type should be removed for multipart"
|
||||
|
||||
filename, content, content_type = kwargs["files"]["file"]
|
||||
assert filename == "test.txt"
|
||||
assert content == b"test file content"
|
||||
assert content_type == "text/plain"
|
||||
|
||||
return mock_response
|
||||
|
||||
async_client = MagicMock()
|
||||
async_client.request = AsyncMock(side_effect=mock_httpx_request)
|
||||
|
||||
# Create mock request
|
||||
request = MagicMock(spec=Request)
|
||||
request.method = "POST"
|
||||
request.headers = Headers({"content-type": "multipart/form-data; boundary=test123"})
|
||||
|
||||
# Mock form data
|
||||
file_content = b"test file content"
|
||||
file = BytesIO(file_content)
|
||||
headers = Headers({"content-type": "text/plain"})
|
||||
upload_file = UploadFile(file=file, filename="test.txt", headers=headers)
|
||||
upload_file.read = AsyncMock(return_value=file_content)
|
||||
|
||||
form_data = {"file": upload_file}
|
||||
request.form = AsyncMock(return_value=form_data)
|
||||
|
||||
# Test the multipart handler directly
|
||||
response = await HttpPassThroughEndpointHelpers.make_multipart_http_request(
|
||||
request=request,
|
||||
async_client=async_client,
|
||||
url=httpx.URL("http://test.com/upload"),
|
||||
headers={},
|
||||
requested_query_params=None,
|
||||
)
|
||||
|
||||
# Verify the response
|
||||
assert response.status_code == 200
|
||||
async_client.request.assert_called_once()
|
||||
|
||||
@@ -664,6 +664,64 @@ class TestHealthAppFactory:
|
||||
)
|
||||
mock_setup_database.assert_called_with(use_migrate=False)
|
||||
|
||||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
|
||||
def test_startup_fails_when_db_setup_fails(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
):
|
||||
"""Test that proxy exits with code 1 when PrismaManager.setup_database returns False"""
|
||||
from litellm.proxy.proxy_cli import run_server
|
||||
|
||||
mock_subprocess_run.return_value = MagicMock(returncode=0)
|
||||
mock_should_update_schema.return_value = True
|
||||
mock_setup_database.return_value = False
|
||||
|
||||
mock_proxy_module = MagicMock(
|
||||
app=MagicMock(),
|
||||
ProxyConfig=MagicMock(),
|
||||
KeyManagementSettings=MagicMock(),
|
||||
save_worker_config=MagicMock(),
|
||||
)
|
||||
|
||||
clean_env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("DATABASE_URL", "DIRECT_URL")
|
||||
}
|
||||
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
|
||||
|
||||
with patch.dict(
|
||||
os.environ, clean_env, clear=True
|
||||
), patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"proxy_server": mock_proxy_module,
|
||||
"litellm.proxy.proxy_server": mock_proxy_module,
|
||||
},
|
||||
), patch(
|
||||
"litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"
|
||||
) as mock_get_args:
|
||||
mock_get_args.return_value = {
|
||||
"app": "litellm.proxy.proxy_server:app",
|
||||
"host": "localhost",
|
||||
"port": 8000,
|
||||
}
|
||||
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
run_server.main(
|
||||
["--local", "--skip_server_startup"], standalone_mode=False
|
||||
)
|
||||
assert exc_info.value.code == 1
|
||||
mock_setup_database.assert_called_once_with(use_migrate=True)
|
||||
|
||||
|
||||
# --- Module-level helpers for worker startup hook tests ---
|
||||
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
export type { MCPEvent } from "../mcp_tools/types";
|
||||
import type { MCPEvent } from "../mcp_tools/types";
|
||||
export type { MCPEvent };
|
||||
|
||||
export interface ChatMessage {
|
||||
id: string;
|
||||
|
||||
@@ -33,7 +33,7 @@ interface CreateMCPServerProps {
|
||||
}
|
||||
|
||||
const AUTH_TYPES_REQUIRING_AUTH_VALUE = [AUTH_TYPE.API_KEY, AUTH_TYPE.BEARER_TOKEN, AUTH_TYPE.TOKEN, AUTH_TYPE.BASIC];
|
||||
const AUTH_TYPES_REQUIRING_CREDENTIALS = [...AUTH_TYPES_REQUIRING_AUTH_VALUE, AUTH_TYPE.OAUTH2];
|
||||
const AUTH_TYPES_REQUIRING_CREDENTIALS = [...AUTH_TYPES_REQUIRING_AUTH_VALUE, AUTH_TYPE.OAUTH2, AUTH_TYPE.AWS_SIGV4];
|
||||
const CREATE_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-create-state";
|
||||
|
||||
const reduceStaticHeaders = (list: unknown): Record<string, string> => {
|
||||
@@ -85,6 +85,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
|
||||
const authType = formValues.auth_type as string | undefined;
|
||||
const shouldShowAuthValueField = authType ? AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(authType) : false;
|
||||
const isOAuthAuthType = authType === AUTH_TYPE.OAUTH2;
|
||||
const isAwsSigV4AuthType = authType === AUTH_TYPE.AWS_SIGV4;
|
||||
const isM2MFlow = isOAuthAuthType && formValues.oauth_flow_type === OAUTH_FLOW.M2M;
|
||||
|
||||
const persistCreateUiState = () => {
|
||||
@@ -767,6 +768,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
|
||||
<Select.Option value="token">Token</Select.Option>
|
||||
<Select.Option value="basic">Basic Auth</Select.Option>
|
||||
<Select.Option value="oauth2">OAuth</Select.Option>
|
||||
<Select.Option value="aws_sigv4">AWS SigV4 (Bedrock AgentCore MCPs)</Select.Option>
|
||||
</Select>
|
||||
</Form.Item>
|
||||
|
||||
@@ -818,6 +820,122 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
|
||||
/>
|
||||
)}
|
||||
|
||||
{transportType !== "stdio" && transportType !== "" && isAwsSigV4AuthType && (
|
||||
<>
|
||||
<p className="text-sm text-gray-500 mb-2">
|
||||
For MCP servers hosted on AWS Bedrock AgentCore.{" "}
|
||||
<a href="https://docs.litellm.ai/docs/mcp_aws_sigv4" target="_blank" rel="noopener noreferrer" className="text-blue-500 hover:text-blue-700">
|
||||
View docs →
|
||||
</a>
|
||||
</p>
|
||||
<Form.Item
|
||||
label={
|
||||
<span className="text-sm font-medium text-gray-700 flex items-center">
|
||||
AWS Region
|
||||
<Tooltip title="AWS region for SigV4 signing (e.g., us-east-1)">
|
||||
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name={["credentials", "aws_region_name"]}
|
||||
rules={[{ required: true, message: "AWS region is required for SigV4 auth" }]}
|
||||
>
|
||||
<Input
|
||||
placeholder="us-east-1"
|
||||
className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
|
||||
/>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<span className="text-sm font-medium text-gray-700 flex items-center">
|
||||
AWS Service Name
|
||||
<Tooltip title="AWS service name for SigV4 signing. Defaults to 'bedrock-agentcore'.">
|
||||
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name={["credentials", "aws_service_name"]}
|
||||
>
|
||||
<Input
|
||||
placeholder="bedrock-agentcore"
|
||||
className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
|
||||
/>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<span className="text-sm font-medium text-gray-700 flex items-center">
|
||||
AWS Access Key ID
|
||||
<Tooltip title="Optional. If not provided, falls back to the boto3 credential chain (IAM role, env vars, etc.).">
|
||||
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name={["credentials", "aws_access_key_id"]}
|
||||
dependencies={[["credentials", "aws_secret_access_key"]]}
|
||||
rules={[
|
||||
({ getFieldValue }) => ({
|
||||
validator(_, value) {
|
||||
const secretKey = getFieldValue(["credentials", "aws_secret_access_key"]);
|
||||
if (secretKey && !value) {
|
||||
return Promise.reject(new Error("Access Key ID is required when Secret Access Key is provided"));
|
||||
}
|
||||
return Promise.resolve();
|
||||
},
|
||||
}),
|
||||
]}
|
||||
>
|
||||
<Input.Password
|
||||
placeholder="AKIA... (optional — uses IAM role if blank)"
|
||||
className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
|
||||
/>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<span className="text-sm font-medium text-gray-700 flex items-center">
|
||||
AWS Secret Access Key
|
||||
<Tooltip title="Optional. Required if AWS Access Key ID is provided.">
|
||||
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name={["credentials", "aws_secret_access_key"]}
|
||||
dependencies={[["credentials", "aws_access_key_id"]]}
|
||||
rules={[
|
||||
({ getFieldValue }) => ({
|
||||
validator(_, value) {
|
||||
const accessKeyId = getFieldValue(["credentials", "aws_access_key_id"]);
|
||||
if (accessKeyId && !value) {
|
||||
return Promise.reject(new Error("Secret Access Key is required when Access Key ID is provided"));
|
||||
}
|
||||
return Promise.resolve();
|
||||
},
|
||||
}),
|
||||
]}
|
||||
>
|
||||
<Input.Password
|
||||
placeholder="Enter secret key (optional — uses IAM role if blank)"
|
||||
className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
|
||||
/>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<span className="text-sm font-medium text-gray-700 flex items-center">
|
||||
AWS Session Token
|
||||
<Tooltip title="Optional. Only needed for temporary STS credentials.">
|
||||
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name={["credentials", "aws_session_token"]}
|
||||
>
|
||||
<Input.Password
|
||||
placeholder="Enter session token (optional)"
|
||||
className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
|
||||
/>
|
||||
</Form.Item>
|
||||
</>
|
||||
)}
|
||||
|
||||
{/* Stdio Configuration - only show for stdio transport */}
|
||||
<StdioConfiguration isVisible={transportType === "stdio"} />
|
||||
</div>
|
||||
|
||||
@@ -22,7 +22,7 @@ interface MCPServerEditProps {
|
||||
}
|
||||
|
||||
const AUTH_TYPES_REQUIRING_AUTH_VALUE = [AUTH_TYPE.API_KEY, AUTH_TYPE.BEARER_TOKEN, AUTH_TYPE.TOKEN, AUTH_TYPE.BASIC];
|
||||
const AUTH_TYPES_REQUIRING_CREDENTIALS = [...AUTH_TYPES_REQUIRING_AUTH_VALUE, AUTH_TYPE.OAUTH2];
|
||||
const AUTH_TYPES_REQUIRING_CREDENTIALS = [...AUTH_TYPES_REQUIRING_AUTH_VALUE, AUTH_TYPE.OAUTH2, AUTH_TYPE.AWS_SIGV4];
|
||||
const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state";
|
||||
|
||||
const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
||||
@@ -50,6 +50,7 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
||||
const isMCPTransport = !isStdioTransport && !isOpenAPITransport;
|
||||
const shouldShowAuthValueField = authType ? AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(authType) : false;
|
||||
const isOAuthAuthType = authType === AUTH_TYPE.OAUTH2;
|
||||
const isAwsSigV4AuthType = authType === AUTH_TYPE.AWS_SIGV4;
|
||||
const oauthFlowTypeValue = Form.useWatch("oauth_flow_type", form) as string | undefined;
|
||||
const isM2MFlow = isOAuthAuthType && oauthFlowTypeValue === OAUTH_FLOW.M2M;
|
||||
|
||||
@@ -665,6 +666,7 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
||||
<Select.Option value="token">Token</Select.Option>
|
||||
<Select.Option value="basic">Basic Auth</Select.Option>
|
||||
<Select.Option value="oauth2">OAuth</Select.Option>
|
||||
<Select.Option value="aws_sigv4">AWS SigV4 (Bedrock AgentCore MCPs)</Select.Option>
|
||||
</Select>
|
||||
</Form.Item>
|
||||
)}
|
||||
@@ -883,6 +885,100 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
||||
</>
|
||||
)}
|
||||
|
||||
{!isStdioTransport && isAwsSigV4AuthType && (
|
||||
<>
|
||||
<p className="text-sm text-gray-500 mb-2">
|
||||
For MCP servers hosted on AWS Bedrock AgentCore.{" "}
|
||||
<a href="https://docs.litellm.ai/docs/mcp_aws_sigv4" target="_blank" rel="noopener noreferrer" className="text-blue-500 hover:text-blue-700">
|
||||
View docs →
|
||||
</a>
|
||||
</p>
|
||||
<Form.Item
|
||||
label={
|
||||
<span className="text-sm font-medium text-gray-700 flex items-center">
|
||||
AWS Region
|
||||
<Tooltip title="AWS region for SigV4 signing (e.g., us-east-1)">
|
||||
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name={["credentials", "aws_region_name"]}
|
||||
rules={[]}
|
||||
>
|
||||
<Input
|
||||
placeholder="us-east-1 (leave blank to keep existing)"
|
||||
className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
|
||||
/>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<span className="text-sm font-medium text-gray-700 flex items-center">
|
||||
AWS Service Name
|
||||
<Tooltip title="AWS service name for SigV4 signing. Defaults to 'bedrock-agentcore'.">
|
||||
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name={["credentials", "aws_service_name"]}
|
||||
>
|
||||
<Input
|
||||
placeholder="bedrock-agentcore (leave blank to keep existing)"
|
||||
className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
|
||||
/>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<span className="text-sm font-medium text-gray-700 flex items-center">
|
||||
AWS Access Key ID
|
||||
<Tooltip title="Optional. If not provided, falls back to the boto3 credential chain (IAM role, env vars, etc.).">
|
||||
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name={["credentials", "aws_access_key_id"]}
|
||||
rules={[]}
|
||||
>
|
||||
<Input.Password
|
||||
placeholder="Leave blank to keep existing"
|
||||
className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
|
||||
/>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<span className="text-sm font-medium text-gray-700 flex items-center">
|
||||
AWS Secret Access Key
|
||||
<Tooltip title="Optional. Required if AWS Access Key ID is provided.">
|
||||
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name={["credentials", "aws_secret_access_key"]}
|
||||
rules={[]}
|
||||
>
|
||||
<Input.Password
|
||||
placeholder="Leave blank to keep existing"
|
||||
className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
|
||||
/>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<span className="text-sm font-medium text-gray-700 flex items-center">
|
||||
AWS Session Token
|
||||
<Tooltip title="Optional. Only needed for temporary STS credentials.">
|
||||
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name={["credentials", "aws_session_token"]}
|
||||
>
|
||||
<Input.Password
|
||||
placeholder="Leave blank to keep existing"
|
||||
className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
|
||||
/>
|
||||
</Form.Item>
|
||||
</>
|
||||
)}
|
||||
|
||||
{/* Permission Management / Access Control Section */}
|
||||
<div className="mt-6">
|
||||
<MCPPermissionManagement
|
||||
|
||||
@@ -39,6 +39,7 @@ export const AUTH_TYPE = {
|
||||
TOKEN: "token",
|
||||
BASIC: "basic",
|
||||
OAUTH2: "oauth2",
|
||||
AWS_SIGV4: "aws_sigv4",
|
||||
};
|
||||
|
||||
export const OAUTH_FLOW = {
|
||||
|
||||
Reference in New Issue
Block a user