scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

Merge branch 'main' into litellm_dev_08_20_2025_p1

Browse source
This commit is contained in:
Krish Dholakia 2025-08-23 12:09:53 -07:00 • committed by GitHub
commit 9df6f4ef08
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
112 changed files with 8644 additions and 1885 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
.gitignore vendored
View file

@ -94,4 +94,5 @@ test.py
litellm_config.yaml
.cursor
.vscode/launch.json
.vscode/launch.json
litellm/proxy/to_delete_loadtest_work/*

213
cookbook/liteLLM_Baseten.ipynb vendored
View file

@ -6,19 +6,21 @@
"id": "gZx-wHJapG5w"
},
"source": [
"# Use liteLLM to call Falcon, Wizard, MPT 7B using OpenAI chatGPT Input/output\n",
"# LiteLLM with Baseten Model APIs\n",
"\n",
"* Falcon 7B: https://app.baseten.co/explore/falcon_7b\n",
"* Wizard LM: https://app.baseten.co/explore/wizardlm\n",
"* MPT 7B Base: https://app.baseten.co/explore/mpt_7b_instruct\n",
"This notebook demonstrates how to use LiteLLM with Baseten's Model APIs instead of dedicated deployments.\n",
"\n",
"\n",
"## Call all baseten llm models using OpenAI chatGPT Input/Output using liteLLM\n",
"Example call\n",
"## Example Usage\n",
"```python\n",
"model = \"q841o8w\" # baseten model version ID\n",
"response = completion(model=model, messages=messages, custom_llm_provider=\"baseten\")\n",
"```"
"response = completion(\n",
" model=\"baseten/openai/gpt-oss-120b\",\n",
" messages=[{\"role\": \"user\", \"content\": \"Hello!\"}],\n",
" max_tokens=1000,\n",
" temperature=0.7\n",
")\n",
"```\n",
"\n",
"## Setup"
]
},
{
@ -29,20 +31,25 @@
},
"outputs": [],
"source": [
"!pip install litellm==0.1.399\n",
"!pip install baseten urllib3"
"%pip install litellm"
]
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": null,
"metadata": {
"id": "VEukLhDzo4vw"
},
"outputs": [],
"source": [
"import os\n",
"from litellm import completion"
"from litellm import completion\n",
"\n",
"# Set your Baseten API key\n",
"os.environ['BASETEN_API_KEY'] = \"\" #@param {type:\"string\"}\n",
"\n",
"# Test message\n",
"messages = [{\"role\": \"user\", \"content\": \"What is AGI?\"}]"
]
},
{
@ -51,19 +58,31 @@
"id": "4STYM2OHFNlc"
},
"source": [
"## Setup"
"## Example 1: Basic Completion\n",
"\n",
"Simple completion with the GPT-OSS 120B model"
]
},
{
"cell_type": "code",
"execution_count": 21,
"execution_count": null,
"metadata": {
"id": "DorpLxw1FHbC"
},
"outputs": [],
"source": [
"os.environ['BASETEN_API_KEY'] = \"\" #@param\n",
"messages = [{ \"content\": \"what does Baseten do? \",\"role\": \"user\"}]"
"print(\"=== Basic Completion ===\")\n",
"response = completion(\n",
" model=\"baseten/openai/gpt-oss-120b\",\n",
" messages=messages,\n",
" max_tokens=1000,\n",
" temperature=0.7,\n",
" top_p=0.9,\n",
" presence_penalty=0.1,\n",
" frequency_penalty=0.1,\n",
")\n",
"print(f\"Response: {response.choices[0].message.content}\")\n",
"print(f\"Usage: {response.usage}\")"
]
},
{
@ -72,13 +91,14 @@
"id": "syF3dTdKFSQQ"
},
"source": [
"## Calling Falcon 7B: https://app.baseten.co/explore/falcon_7b\n",
"### Pass Your Baseten model `Version ID` as `model`"
"## Example 2: Streaming Completion\n",
"\n",
"Streaming completion with usage statistics"
]
},
{
"cell_type": "code",
"execution_count": 18,
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
@ -86,137 +106,26 @@
"id": "rPgSoMlsojz0",
"outputId": "81d6dc7b-1681-4ae4-e4c8-5684eb1bd050"
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[32mINFO\u001b[0m API key set.\n",
"INFO:baseten:API key set.\n"
]
},
{
"data": {
"text/plain": [
"{'choices': [{'finish_reason': 'stop',\n",
" 'index': 0,\n",
" 'message': {'role': 'assistant',\n",
" 'content': \"what does Baseten do? \\nI'm sorry, I cannot provide a specific answer as\"}}],\n",
" 'created': 1692135883.699066,\n",
" 'model': 'qvv0xeq'}"
]
},
"execution_count": 18,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"model = \"qvv0xeq\"\n",
"response = completion(model=model, messages=messages, custom_llm_provider=\"baseten\")\n",
"response"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "7n21UroEGCGa"
},
"source": [
"## Calling Wizard LM https://app.baseten.co/explore/wizardlm\n",
"### Pass Your Baseten model `Version ID` as `model`"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "uLVWFH899lAF",
"outputId": "61c2bc74-673b-413e-bb40-179cf408523d"
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[32mINFO\u001b[0m API key set.\n",
"INFO:baseten:API key set.\n"
]
},
{
"data": {
"text/plain": [
"{'choices': [{'finish_reason': 'stop',\n",
" 'index': 0,\n",
" 'message': {'role': 'assistant',\n",
" 'content': 'As an AI language model, I do not have personal beliefs or practices, but based on the information available online, Baseten is a popular name for a traditional Ethiopian dish made with injera, a spongy flatbread, and wat, a spicy stew made with meat or vegetables. It is typically served for breakfast or dinner and is a staple in Ethiopian cuisine. The name Baseten is also used to refer to a traditional Ethiopian coffee ceremony, where coffee is brewed and served in a special ceremony with music and food.'}}],\n",
" 'created': 1692135900.2806294,\n",
" 'model': 'q841o8w'}"
]
},
"execution_count": 19,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"model = \"q841o8w\"\n",
"response = completion(model=model, messages=messages, custom_llm_provider=\"baseten\")\n",
"response"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "6-TFwmPAGPXq"
},
"source": [
"## Calling mosaicml/mpt-7b https://app.baseten.co/explore/mpt_7b_instruct\n",
"### Pass Your Baseten model `Version ID` as `model`"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "gbeYZOrUE_Bp",
"outputId": "838d86ea-2143-4cb3-bc80-2acc2346c37a"
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[32mINFO\u001b[0m API key set.\n",
"INFO:baseten:API key set.\n"
]
},
{
"data": {
"text/plain": [
"{'choices': [{'finish_reason': 'stop',\n",
" 'index': 0,\n",
" 'message': {'role': 'assistant',\n",
" 'content': \"\\n===================\\n\\nIt's a tool to build a local version of a game on your own machine to host\\non your website.\\n\\nIt's used to make game demos and show them on Twitter, Tumblr, and Facebook.\\n\\n\\n\\n## What's built\\n\\n- A directory of all your game directories, named with a version name and build number, with images linked to.\\n- Includes HTML to include in another site.\\n- Includes images for your icons and\"}}],\n",
" 'created': 1692135914.7472186,\n",
" 'model': '31dxrj3'}"
]
},
"execution_count": 20,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"model = \"31dxrj3\"\n",
"response = completion(model=model, messages=messages, custom_llm_provider=\"baseten\")\n",
"response"
"print(\"=== Streaming Completion ===\")\n",
"response = completion(\n",
" model=\"baseten/openai/gpt-oss-120b\",\n",
" messages=[{\"role\": \"user\", \"content\": \"Write a short poem about AI\"}],\n",
" stream=True,\n",
" max_tokens=500,\n",
" temperature=0.8,\n",
" stream_options={\n",
" \"include_usage\": True,\n",
" \"continuous_usage_stats\": True\n",
" },\n",
")\n",
"\n",
"print(\"Streaming response:\")\n",
"for chunk in response:\n",
" if chunk.choices and chunk.choices[0].delta.content:\n",
" print(chunk.choices[0].delta.content, end=\"\", flush=True)\n",
"print(\"\\n\")"
]
}
],
@ -234,4 +143,4 @@
},
"nbformat": 4,
"nbformat_minor": 0
}
}

4
deploy/charts/litellm-helm/README.md
View file

@ -24,7 +24,7 @@ If `db.useStackgresOperator` is used (not yet implemented):
| `replicaCount` | The number of LiteLLM Proxy pods to be deployed | `1` |
| `masterkeySecretName` | The name of the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use the generated secret name. | N/A |
| `masterkeySecretKey` | The key within the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use `masterkey` as the key. | N/A |
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key is generated. | N/A |
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated. | N/A |
| `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
| `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
| `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` |
@ -135,7 +135,7 @@ service, the **Proxy Endpoint** should be set to `http://<RELEASE>-litellm:4000`
The **Proxy Key** is the value specified for `masterkey` or, if a `masterkey`
was not provided to the helm command line, the `masterkey` is a randomly
generated string stored in the `<RELEASE>-litellm-masterkey` Kubernetes Secret.
generated string in the `sk-...` format stored in the `<RELEASE>-litellm-masterkey` Kubernetes Secret.
```bash
kubectl -n litellm get secret <RELEASE>-litellm-masterkey -o jsonpath="{.data.masterkey}"

7
deploy/charts/litellm-helm/templates/deployment.yaml
View file

@ -71,7 +71,14 @@ spec:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.passwordKey }}
- name: DATABASE_HOST
{{- if .Values.db.secret.endpointKey }}
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.endpointKey }}
{{- else }}
value: {{ .Values.db.endpoint }}
{{- end }}
- name: DATABASE_NAME
value: {{ .Values.db.database }}
- name: DATABASE_URL

7
deploy/charts/litellm-helm/templates/migrations-job.yaml
View file

@ -49,7 +49,14 @@ spec:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.passwordKey }}
- name: DATABASE_HOST
{{- if .Values.db.secret.endpointKey }}
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.endpointKey }}
{{- else }}
value: {{ .Values.db.endpoint }}
{{- end }}
- name: DATABASE_NAME
value: {{ .Values.db.database }}
- name: DATABASE_URL

2
deploy/charts/litellm-helm/templates/secret-masterkey.yaml
View file

@ -1,5 +1,5 @@
{{- if not .Values.masterkeySecretName }}
{{ $masterkey := (.Values.masterkey | default (randAlphaNum 17)) }}
{{ $masterkey := (.Values.masterkey | default (printf "sk-%s" (randAlphaNum 18))) }}
apiVersion: v1
kind: Secret
metadata:

8
deploy/charts/litellm-helm/tests/masterkey-secret_tests.yaml
View file

@ -2,13 +2,19 @@ suite: test masterkey secret
templates:
- secret-masterkey.yaml
tests:
- it: should create a secret if masterkeySecretName is not set
- it: should create a secret if masterkeySecretName is not set. should start with sk-xxxx (base64 encoded as c2st*)
template: secret-masterkey.yaml
set:
masterkeySecretName: ""
asserts:
- isKind:
of: Secret
- matchRegex:
path: data.masterkey
pattern: ^c2st
# Note: The masterkey is generated as "sk-<18-random-chars>" in plain text,
# but stored as base64 encoded in Kubernetes secret (requirement).
# "sk-" base64 encodes to "c2st", so we check for "^c2st" pattern.
- it: should not create a secret if masterkeySecretName is set
template: secret-masterkey.yaml
set:

2
deploy/charts/litellm-helm/values.yaml
View file

@ -161,6 +161,8 @@ db:
name: postgres
usernameKey: username
passwordKey: password
# Optional: when set, DATABASE_HOST will be sourced from this secret key instead of db.endpoint
endpointKey: ""
# Use the Stackgres Helm chart to deploy an instance of a Stackgres cluster.
# The Stackgres Operator must already be installed within the target

4
docker/Dockerfile.non_root
View file

@ -70,7 +70,9 @@ RUN mkdir -p /nonexistent /.npm && \
chown -R nobody:nogroup /app && \
chown -R nobody:nogroup /nonexistent /.npm && \
PRISMA_PATH=$(python -c "import os, prisma; print(os.path.dirname(prisma.__file__))") && \
chown -R nobody:nogroup $PRISMA_PATH
chown -R nobody:nogroup $PRISMA_PATH && \
LITELLM_PKG_MIGRATIONS_PATH="$(python -c 'import os, litellm_proxy_extras; print(os.path.dirname(litellm_proxy_extras.__file__))' 2>/dev/null || echo '')/migrations" && \
[ -n "$LITELLM_PKG_MIGRATIONS_PATH" ] && chown -R nobody:nogroup $LITELLM_PKG_MIGRATIONS_PATH
# --- OpenShift Compatibility: Apply Red Hat recommended pattern ---
# Get paths for directories that need write access at runtime

66
docs/my-website/docs/completion/document_understanding.md
View file

@ -10,6 +10,7 @@ Works for:
- Bedrock Models
- Anthropic API Models
- OpenAI API Models
- Mistral (Only using file ID of already uploaded file, similar to OpenAI file_id input)
## Quick Start
@ -279,6 +280,71 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \
</Tabs>
## Mistral Example
Here is a sample payload for using the Mistral model for document understanding:
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm.utils import completion
# pdf file_id received from files endpoint
file_id = "fa778e5e-46ec-4562-8418-36623fe25a71"
# model
model = "mistral/mistral-large-latest"
file_content = [
{"type": "text", "text": "What's this file about?"},
{
"type": "file",
"file": {
"file_id": file_id,
}
},
]
response = completion(
model=model,
messages=[{"role": "user", "content": file_content}],
)
assert response is not None
```
</TabItem>
<TabItem value="proxy" label="PROXY">
```bash
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
"model": "mistral/mistral-large-latest",
"messages": [
{
"role": "user",
"content": [
{
"type": "text",
"text": "What is the content of the file?"
},
{
"type": "file",
"file": {
"file_id": "fa778e5e-46ec-4562-8418-36623fe25a71"
}
}
]
}
]
}
```
</TabItem>
</Tabs>
## Checking if a model supports pdf input
<Tabs>

113
docs/my-website/docs/providers/baseten.md
View file

@ -1,23 +1,106 @@
# Baseten
LiteLLM supports any Text-Gen-Interface models on Baseten.
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
[Here's a tutorial on deploying a huggingface TGI model (Llama2, CodeLlama, WizardCoder, Falcon, etc.) on Baseten](https://truss.baseten.co/examples/performance/tgi-server)
# Baseten
LiteLLM supports both Baseten Model APIs and dedicated deployments with automatic routing.
## API Types
### Model API (Default)
- **URL**: `https://inference.baseten.co/v1`
- **Format**: `baseten/<model-name>` (e.g., `baseten/openai/gpt-oss-120b`)
- **Best for**: Quick access to popular models
### Dedicated Deployments
- **URL**: `https://model-{id}.api.baseten.co/environments/production/sync/v1`
- **Format**: `baseten/{8-digit-alphanumeric-code}` (e.g., `baseten/abcd1234`)
- **Best for**: Custom models, latency SLAs
:::tip
**Automatic Routing**: LiteLLM detects the type based on model format:
- 8-digit alphanumeric codes → Dedicated deployment
- All other formats → Model API
:::
## Quick Start
### API KEYS
```python
import os
os.environ["BASETEN_API_KEY"] = ""
import os
from litellm import completion
os.environ['BASETEN_API_KEY'] = "your-api-key"
# Model API (default)
response = completion(
model="baseten/openai/gpt-oss-120b",
messages=[{"role": "user", "content": "Hello!"}]
)
# Dedicated deployment (8-digit ID)
response = completion(
model="baseten/abcd1234",
messages=[{"role": "user", "content": "Hello!"}]
)
```
### Baseten Models
Baseten provides infrastructure to deploy and serve ML models https://www.baseten.co/. Use liteLLM to easily call models deployed on Baseten.
## Examples
Example Baseten Usage - Note: liteLLM supports all models deployed on Baseten
### Basic Usage
```python
# Model API
response = completion(
model="baseten/openai/gpt-oss-120b",
messages=[{"role": "user", "content": "Explain quantum computing"}],
max_tokens=500,
temperature=0.7
)
Usage: Pass `model=baseten/<Model ID>`
# Dedicated deployment
response = completion(
model="baseten/abcd1234",
messages=[{"role": "user", "content": "Explain quantum computing"}],
max_tokens=500,
temperature=0.7
)
```
| Model Name | Function Call | Required OS Variables |
|------------------|--------------------------------------------|------------------------------------|
| Falcon 7B | `completion(model='baseten/qvv0xeq', messages=messages)` | `os.environ['BASETEN_API_KEY']` |
| Wizard LM | `completion(model='baseten/q841o8w', messages=messages)` | `os.environ['BASETEN_API_KEY']` |
| MPT 7B Base | `completion(model='baseten/31dxrj3', messages=messages)` | `os.environ['BASETEN_API_KEY']` |
### Streaming (Model API only)
```python
response = completion(
model="baseten/openai/gpt-oss-120b",
messages=[{"role": "user", "content": "Write a poem"}],
stream=True,
stream_options={"include_usage": True}
)
for chunk in response:
if chunk.choices and chunk.choices[0].delta.content:
print(chunk.choices[0].delta.content, end="")
```
## Usage with LiteLLM Proxy
1. **Config**:
```yaml
model_list:
- model_name: baseten-model
litellm_params:
model: baseten/openai/gpt-oss-120b
api_key: your-baseten-api-key
```
2. **Request**:
```python
import openai
client = openai.OpenAI(
api_key="sk-1234",
base_url="http://0.0.0.0:4000"
)
response = client.chat.completions.create(
model="baseten-model",
messages=[{"role": "user", "content": "Hello!"}]
)
```

140
docs/my-website/docs/providers/deepinfra.md
View file

@ -1,3 +1,6 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# DeepInfra
https://deepinfra.com/
@ -7,6 +10,11 @@ https://deepinfra.com/
:::
## Table of Contents
- [API Key](#api-key)
- [Chat Models](#chat-models)
- [Rerank Endpoint](#rerank-endpoint)
## API Key
```python
@ -53,3 +61,135 @@ for chunk in response:
| codellama/CodeLlama-34b-Instruct-hf | `completion(model="deepinfra/codellama/CodeLlama-34b-Instruct-hf", messages)` |
| mistralai/Mistral-7B-Instruct-v0.1 | `completion(model="deepinfra/mistralai/Mistral-7B-Instruct-v0.1", messages)` |
| jondurbin/airoboros-l2-70b-gpt4-1.4.1 | `completion(model="deepinfra/jondurbin/airoboros-l2-70b-gpt4-1.4.1", messages)` |
## Rerank Endpoint
LiteLLM provides a Cohere API compatible `/rerank` endpoint for DeepInfra rerank models.
### Supported Rerank Models
| Model Name | Description |
|------------|-------------|
| `deepinfra/Qwen/Qwen3-Reranker-0.6B` | Lightweight rerank model (0.6B parameters) |
| `deepinfra/Qwen/Qwen3-Reranker-4B` | Medium rerank model (4B parameters) |
| `deepinfra/Qwen/Qwen3-Reranker-8B` | Large rerank model (8B parameters) |
### Usage - LiteLLM Python SDK
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm import rerank
import os
os.environ["DEEPINFRA_API_KEY"] = "your-api-key"
response = rerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="What is the capital of France?",
documents=[
"Paris is the capital of France.",
"London is the capital of the United Kingdom.",
"Berlin is the capital of Germany.",
"Madrid is the capital of Spain.",
"Rome is the capital of Italy."
]
)
print(response)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
1. Add to config.yaml
```yaml
model_list:
- model_name: Qwen/Qwen3-Reranker-0.6B
litellm_params:
model: deepinfra/Qwen/Qwen3-Reranker-0.6B
api_key: os.environ/DEEPINFRA_API_KEY
```
2. Start proxy
```bash
litellm --config /path/to/config.yaml
# RUNNING on http://0.0.0.0:4000/
```
3. Test it!
```bash
curl -L -X POST 'http://0.0.0.0:4000/rerank' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{
"model": "Qwen/Qwen3-Reranker-0.6B",
"query": "What is the capital of France?",
"documents": [
"Paris is the capital of France.",
"London is the capital of the United Kingdom.",
"Berlin is the capital of Germany.",
"Madrid is the capital of Spain.",
"Rome is the capital of Italy."
]
}'
```
</TabItem>
</Tabs>
### Supported Cohere Rerank API Params
| Param | Type | Description |
| ------------------ | ----------- | ----------------------------------------------- |
| `query` | `str` | The query to rerank the documents against |
| `documents` | `list[str]` | The documents to rerank |
### Provider-specific parameters
Pass any deepinfra specific parameters as a keyword argument to the rerank function, e.g.
```
response = rerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="What is the capital of France?",
documents=[
"Paris is the capital of France.",
"London is the capital of the United Kingdom.",
"Berlin is the capital of Germany.",
"Madrid is the capital of Spain.",
"Rome is the capital of Italy."
],
my_custom_param="my_custom_value", # any other deepinfra specific parameters
)
```
### Response Format
```json
{
"id": "request-id",
"results": [
{
"index": 0,
"relevance_score": 0.9975274205207825
},
{
"index": 1,
"relevance_score": 0.011687257327139378
}
],
"meta": {
"billed_units": {
"total_tokens": 427
},
"tokens": {
"input_tokens": 427,
"output_tokens": 0
}
}
}
```

87
docs/my-website/docs/providers/vertex_partner.md
View file

@ -14,6 +14,7 @@ import TabItem from '@theme/TabItem';
| Meta/Llama | `vertex_ai/meta/{MODEL}` | [Vertex AI - Meta Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/llama) |
| Mistral | `vertex_ai/mistral-*` | [Vertex AI - Mistral Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/mistral) |
| AI21 (Jamba) | `vertex_ai/jamba-*` | [Vertex AI - AI21 Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/ai21) |
| Qwen | `vertex_ai/qwen/*` | [Vertex AI - Qwen Models](https://cloud.google.com/vertex-ai/generative-ai/docs/maas/qwen) |
| Model Garden | `vertex_ai/openai/{MODEL_ID}` or `vertex_ai/{MODEL_ID}` | [Vertex Model Garden](https://cloud.google.com/model-garden?hl=en) |
## Vertex AI - Anthropic (Claude)
@ -571,6 +572,92 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
</Tabs>
## VertexAI Qwen API
| Property | Details |
|----------|---------|
| Provider Route | `vertex_ai/qwen/{MODEL}` |
| Vertex Documentation | [Vertex AI - Qwen Models](https://cloud.google.com/vertex-ai/generative-ai/docs/maas/qwen) |
**LiteLLM Supports all Vertex AI Qwen Models.** Ensure you use the `vertex_ai/qwen/` prefix for all Vertex AI Qwen models.
| Model Name | Usage |
|------------------|------------------------------|
| vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas | `completion('vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas', messages)` |
| vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas | `completion('vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas', messages)` |
#### Usage
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm import completion
import os
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = ""
model = "qwen/qwen3-coder-480b-a35b-instruct-maas"
vertex_ai_project = "your-vertex-project" # can also set this as os.environ["VERTEXAI_PROJECT"]
vertex_ai_location = "your-vertex-location" # can also set this as os.environ["VERTEXAI_LOCATION"]
response = completion(
model="vertex_ai/" + model,
messages=[{"role": "user", "content": "hi"}],
vertex_ai_project=vertex_ai_project,
vertex_ai_location=vertex_ai_location,
)
print("\nModel Response", response)
```
</TabItem>
<TabItem value="proxy" label="Proxy">
**1. Add to config**
```yaml
model_list:
- model_name: vertex-qwen
litellm_params:
model: vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas
vertex_ai_project: "my-test-project"
vertex_ai_location: "us-east-1"
- model_name: vertex-qwen
litellm_params:
model: vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas
vertex_ai_project: "my-test-project"
vertex_ai_location: "us-west-1"
```
**2. Start proxy**
```bash
litellm --config /path/to/config.yaml
# RUNNING at http://0.0.0.0:4000
```
**3. Test it!**
```bash
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "vertex-qwen", # 👈 the 'model_name' in config
"messages": [
{
"role": "user",
"content": "what llm are you"
}
],
}'
```
</TabItem>
</Tabs>
## Model Garden
:::tip

4
docs/my-website/docs/proxy/config_settings.md
View file

@ -335,12 +335,16 @@ router_settings:
| ANTHROPIC_API_KEY | API key for Anthropic service
| ANTHROPIC_API_BASE | Base URL for Anthropic API. Default is https://api.anthropic.com
| AWS_ACCESS_KEY_ID | Access Key ID for AWS services
| AWS_DEFAULT_REGION | Default AWS region for service interactions when AWS_REGION is not set
| AWS_PROFILE_NAME | AWS CLI profile name to be used
| AWS_REGION | AWS region for service interactions (takes precedence over AWS_DEFAULT_REGION)
| AWS_REGION_NAME | Default AWS region for service interactions
| AWS_ROLE_ARN | ARN of the AWS IAM role to assume for authentication
| AWS_ROLE_NAME | Role name for AWS IAM usage
| AWS_SECRET_ACCESS_KEY | Secret Access Key for AWS services
| AWS_SESSION_NAME | Name for AWS session
| AWS_WEB_IDENTITY_TOKEN | Web identity token for AWS
| AWS_WEB_IDENTITY_TOKEN_FILE | Path to file containing web identity token for AWS
| AZURE_API_VERSION | Version of the Azure API being used
| AZURE_AUTHORITY_HOST | Azure authority host URL
| AZURE_CERTIFICATE_PASSWORD | Password for Azure OpenAI certificate

6
docs/my-website/docs/proxy/team_budgets.md
View file

@ -4,6 +4,12 @@ import TabItem from '@theme/TabItem';
# Setting Team Budgets
# Pre-Requisites
- You must set up a Postgres database (e.g. Supabase, Neon, etc.)
- To enable team member rate limits, set the environment variable `EXPERIMENTAL_MULTI_INSTANCE_RATE_LIMITING=true` **before starting the proxy server**. Without this, team member rate limits will not be enforced.
Track spend, set budgets for your Internal Team
## Setting Monthly Team Budgets

8
docs/my-website/docs/proxy/users.md
View file

@ -58,6 +58,9 @@ You can:
**Step-by step tutorial on setting, resetting budgets on Teams here (API or using Admin UI)**
> **Prerequisite:**
> To enable team member rate limits, you must set the environment variable `EXPERIMENTAL_MULTI_INSTANCE_RATE_LIMITING=true` before starting the proxy server. Without this, team member rate limits will not be enforced.
👉 [https://docs.litellm.ai/docs/proxy/team_budgets](https://docs.litellm.ai/docs/proxy/team_budgets)
:::
@ -793,6 +796,11 @@ Expected Response:
Enable multi-instance rate limiting with the env var `EXPERIMENTAL_MULTI_INSTANCE_RATE_LIMITING="True"`
**Important Notes:**
- Setting `EXPERIMENTAL_MULTI_INSTANCE_RATE_LIMITING="True"` is required for team member rate limits to function, not just for multi-instance scenarios.
- **Rate limits do not apply to proxy admin users.**
- When testing rate limits, use internal user roles (non-admin) to ensure limits are enforced as expected.
Changes:
- This moves to using async_increment instead of async_set_cache when updating current requests/tokens.
- The in-memory cache is synced with redis every 0.01s, to avoid calling redis for every request.

3
docs/my-website/docs/rerank.md
View file

@ -118,4 +118,5 @@ curl http://0.0.0.0:4000/rerank \
| AWS Bedrock| [Usage](../docs/providers/bedrock#rerank-api) |
| HuggingFace| [Usage](../docs/providers/huggingface_rerank) |
| Infinity| [Usage](../docs/providers/infinity) |
| vLLM| [Usage](../docs/providers/vllm#rerank-endpoint) |
| vLLM| [Usage](../docs/providers/vllm#rerank-endpoint) |
| DeepInfra| [Usage](../docs/providers/deepinfra#rerank-endpoint) |

7
docs/my-website/docs/response_api.md
View file

@ -805,16 +805,17 @@ LiteLLM Proxy supports session management for non-OpenAI models. This allows you
Set `store_prompts_in_cold_storage: true` in your proxy config.yaml. When this is enabled, LiteLLM will store the request and response content in the s3 bucket you specify.
```yaml
```yaml showLineNumbers title="config.yaml with Session Continuity"
litellm_settings:
callbacks: ["s3_v2"]
cold_storage_custom_logger: s3_v2
s3_callback_params: # learn more https://docs.litellm.ai/docs/proxy/logging#s3-buckets
s3_bucket_name: litellm-logs # AWS Bucket Name for S3
s3_region_name: us-west-2
s3_region_name: us-west-2
general_settings:
cold_storage_custom_logger: s3_v2
store_prompts_in_cold_storage: true
store_prompts_in_spend_logs: true
```
2. Make request 1 with no `previous_response_id` (new session)

16
enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -2,6 +2,8 @@
Enterprise internal user management endpoints
"""
import os
from fastapi import APIRouter, Depends, HTTPException
from litellm.proxy._types import UserAPIKeyAuth
@ -21,7 +23,7 @@ async def available_enterprise_users(
"""
For keys with `max_users` set, return the list of users that are allowed to use the key.
"""
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy._types import CommonProxyErrors, EnterpriseLicenseData
from litellm.proxy.proxy_server import (
premium_user,
premium_user_data,
@ -34,10 +36,14 @@ async def available_enterprise_users(
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if premium_user is None:
raise HTTPException(
status_code=500, detail={"error": CommonProxyErrors.not_premium_user.value}
)
if not premium_user:
# check if SSO is enabled - show 5 user limit
from litellm.proxy.auth.auth_utils import _has_user_setup_sso
if _has_user_setup_sso():
premium_user_data = EnterpriseLicenseData(
max_users=5,
)
# Count number of rows in LiteLLM_UserTable
user_count = await prisma_client.db.litellm_usertable.count()

2
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -298,7 +298,7 @@ class ProxyExtrasDBManager:
and "database schema is not empty" in e.stderr
):
logger.info(
"Database schema is not empty, creating baseline migration"
"Database schema is not empty, creating baseline migration. In read-only file system, please set an environment variable `LITELLM_MIGRATION_DIR` to a writable directory to enable migrations. Learn more - https://docs.litellm.ai/docs/proxy/prod#read-only-file-system"
)
ProxyExtrasDBManager._create_baseline_migration(schema_path)
logger.info(

454
litellm/__init__.py
View file

@ -146,6 +146,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"vector_store_pre_call_hook",
"dotprompt",
]
configured_cold_storage_logger: Optional[_custom_logger_compatible_callbacks_literal] = None
logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None
_known_custom_logger_compatible_callbacks: List = list(
get_args(_custom_logger_compatible_callbacks_literal)
@ -466,79 +467,80 @@ BEDROCK_CONVERSE_MODELS = [
]
####### COMPLETION MODELS ###################
open_ai_chat_completion_models: List = []
open_ai_text_completion_models: List = []
cohere_models: List = []
cohere_chat_models: List = []
mistral_chat_models: List = []
text_completion_codestral_models: List = []
anthropic_models: List = []
openrouter_models: List = []
datarobot_models: List = []
vertex_language_models: List = []
vertex_vision_models: List = []
vertex_chat_models: List = []
vertex_code_chat_models: List = []
vertex_ai_image_models: List = []
vertex_text_models: List = []
vertex_code_text_models: List = []
vertex_embedding_models: List = []
vertex_anthropic_models: List = []
vertex_llama3_models: List = []
vertex_deepseek_models: List = []
vertex_ai_ai21_models: List = []
vertex_mistral_models: List = []
ai21_models: List = []
ai21_chat_models: List = []
nlp_cloud_models: List = []
aleph_alpha_models: List = []
bedrock_models: List = []
bedrock_converse_models: List = BEDROCK_CONVERSE_MODELS
fireworks_ai_models: List = []
fireworks_ai_embedding_models: List = []
deepinfra_models: List = []
perplexity_models: List = []
watsonx_models: List = []
gemini_models: List = []
xai_models: List = []
deepseek_models: List = []
azure_ai_models: List = []
jina_ai_models: List = []
voyage_models: List = []
infinity_models: List = []
databricks_models: List = []
cloudflare_models: List = []
codestral_models: List = []
friendliai_models: List = []
featherless_ai_models: List = []
palm_models: List = []
groq_models: List = []
azure_models: List = []
azure_text_models: List = []
anyscale_models: List = []
cerebras_models: List = []
galadriel_models: List = []
sambanova_models: List = []
sambanova_embedding_models: List = []
novita_models: List = []
assemblyai_models: List = []
snowflake_models: List = []
gradient_ai_models: List = []
llama_models: List = []
nscale_models: List = []
nebius_models: List = []
nebius_embedding_models: List = []
deepgram_models: List = []
elevenlabs_models: List = []
dashscope_models: List = []
moonshot_models: List = []
v0_models: List = []
morph_models: List = []
lambda_ai_models: List = []
hyperbolic_models: List = []
recraft_models: List = []
cometapi_models: List = []
oci_models: List = []
from typing import Set
open_ai_chat_completion_models: Set = set()
open_ai_text_completion_models: Set = set()
cohere_models: Set = set()
cohere_chat_models: Set = set()
mistral_chat_models: Set = set()
text_completion_codestral_models: Set = set()
anthropic_models: Set = set()
openrouter_models: Set = set()
datarobot_models: Set = set()
vertex_language_models: Set = set()
vertex_vision_models: Set = set()
vertex_chat_models: Set = set()
vertex_code_chat_models: Set = set()
vertex_ai_image_models: Set = set()
vertex_text_models: Set = set()
vertex_code_text_models: Set = set()
vertex_embedding_models: Set = set()
vertex_anthropic_models: Set = set()
vertex_llama3_models: Set = set()
vertex_deepseek_models: Set = set()
vertex_ai_ai21_models: Set = set()
vertex_mistral_models: Set = set()
ai21_models: Set = set()
ai21_chat_models: Set = set()
nlp_cloud_models: Set = set()
aleph_alpha_models: Set = set()
bedrock_models: Set = set()
bedrock_converse_models: Set = set(BEDROCK_CONVERSE_MODELS)
fireworks_ai_models: Set = set()
fireworks_ai_embedding_models: Set = set()
deepinfra_models: Set = set()
perplexity_models: Set = set()
watsonx_models: Set = set()
gemini_models: Set = set()
xai_models: Set = set()
deepseek_models: Set = set()
azure_ai_models: Set = set()
jina_ai_models: Set = set()
voyage_models: Set = set()
infinity_models: Set = set()
databricks_models: Set = set()
cloudflare_models: Set = set()
codestral_models: Set = set()
friendliai_models: Set = set()
featherless_ai_models: Set = set()
palm_models: Set = set()
groq_models: Set = set()
azure_models: Set = set()
azure_text_models: Set = set()
anyscale_models: Set = set()
cerebras_models: Set = set()
galadriel_models: Set = set()
sambanova_models: Set = set()
sambanova_embedding_models: Set = set()
novita_models: Set = set()
assemblyai_models: Set = set()
snowflake_models: Set = set()
gradient_ai_models: Set = set()
llama_models: Set = set()
nscale_models: Set = set()
nebius_models: Set = set()
nebius_embedding_models: Set = set()
deepgram_models: Set = set()
elevenlabs_models: Set = set()
dashscope_models: Set = set()
moonshot_models: Set = set()
v0_models: Set = set()
morph_models: Set = set()
lambda_ai_models: Set = set()
hyperbolic_models: Set = set()
recraft_models: Set = set()
cometapi_models: Set = set()
oci_models: Set = set()
def is_bedrock_pricing_only_model(key: str) -> bool:
@ -579,166 +581,166 @@ def add_known_models():
if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(
key
):
open_ai_chat_completion_models.append(key)
open_ai_chat_completion_models.add(key)
elif value.get("litellm_provider") == "text-completion-openai":
open_ai_text_completion_models.append(key)
open_ai_text_completion_models.add(key)
elif value.get("litellm_provider") == "azure_text":
azure_text_models.append(key)
azure_text_models.add(key)
elif value.get("litellm_provider") == "cohere":
cohere_models.append(key)
cohere_models.add(key)
elif value.get("litellm_provider") == "cohere_chat":
cohere_chat_models.append(key)
cohere_chat_models.add(key)
elif value.get("litellm_provider") == "mistral":
mistral_chat_models.append(key)
mistral_chat_models.add(key)
elif value.get("litellm_provider") == "anthropic":
anthropic_models.append(key)
anthropic_models.add(key)
elif value.get("litellm_provider") == "empower":
empower_models.append(key)
empower_models.add(key)
elif value.get("litellm_provider") == "openrouter":
openrouter_models.append(key)
openrouter_models.add(key)
elif value.get("litellm_provider") == "datarobot":
datarobot_models.append(key)
datarobot_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-text-models":
vertex_text_models.append(key)
vertex_text_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-code-text-models":
vertex_code_text_models.append(key)
vertex_code_text_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-language-models":
vertex_language_models.append(key)
vertex_language_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-vision-models":
vertex_vision_models.append(key)
vertex_vision_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-chat-models":
vertex_chat_models.append(key)
vertex_chat_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-code-chat-models":
vertex_code_chat_models.append(key)
vertex_code_chat_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-embedding-models":
vertex_embedding_models.append(key)
vertex_embedding_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-anthropic_models":
key = key.replace("vertex_ai/", "")
vertex_anthropic_models.append(key)
vertex_anthropic_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-llama_models":
key = key.replace("vertex_ai/", "")
vertex_llama3_models.append(key)
vertex_llama3_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-deepseek_models":
key = key.replace("vertex_ai/", "")
vertex_deepseek_models.append(key)
vertex_deepseek_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-mistral_models":
key = key.replace("vertex_ai/", "")
vertex_mistral_models.append(key)
vertex_mistral_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-ai21_models":
key = key.replace("vertex_ai/", "")
vertex_ai_ai21_models.append(key)
vertex_ai_ai21_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-image-models":
key = key.replace("vertex_ai/", "")
vertex_ai_image_models.append(key)
vertex_ai_image_models.add(key)
elif value.get("litellm_provider") == "ai21":
if value.get("mode") == "chat":
ai21_chat_models.append(key)
ai21_chat_models.add(key)
else:
ai21_models.append(key)
ai21_models.add(key)
elif value.get("litellm_provider") == "nlp_cloud":
nlp_cloud_models.append(key)
nlp_cloud_models.add(key)
elif value.get("litellm_provider") == "aleph_alpha":
aleph_alpha_models.append(key)
aleph_alpha_models.add(key)
elif value.get(
"litellm_provider"
) == "bedrock" and not is_bedrock_pricing_only_model(key):
bedrock_models.append(key)
bedrock_models.add(key)
elif value.get("litellm_provider") == "bedrock_converse":
bedrock_converse_models.append(key)
bedrock_converse_models.add(key)
elif value.get("litellm_provider") == "deepinfra":
deepinfra_models.append(key)
deepinfra_models.add(key)
elif value.get("litellm_provider") == "perplexity":
perplexity_models.append(key)
perplexity_models.add(key)
elif value.get("litellm_provider") == "watsonx":
watsonx_models.append(key)
watsonx_models.add(key)
elif value.get("litellm_provider") == "gemini":
gemini_models.append(key)
gemini_models.add(key)
elif value.get("litellm_provider") == "fireworks_ai":
# ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params.
if "-to-" not in key and "fireworks-ai-default" not in key:
fireworks_ai_models.append(key)
fireworks_ai_models.add(key)
elif value.get("litellm_provider") == "fireworks_ai-embedding-models":
# ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params.
if "-to-" not in key:
fireworks_ai_embedding_models.append(key)
fireworks_ai_embedding_models.add(key)
elif value.get("litellm_provider") == "text-completion-codestral":
text_completion_codestral_models.append(key)
text_completion_codestral_models.add(key)
elif value.get("litellm_provider") == "xai":
xai_models.append(key)
xai_models.add(key)
elif value.get("litellm_provider") == "deepseek":
deepseek_models.append(key)
deepseek_models.add(key)
elif value.get("litellm_provider") == "meta_llama":
llama_models.append(key)
llama_models.add(key)
elif value.get("litellm_provider") == "nscale":
nscale_models.append(key)
nscale_models.add(key)
elif value.get("litellm_provider") == "azure_ai":
azure_ai_models.append(key)
azure_ai_models.add(key)
elif value.get("litellm_provider") == "voyage":
voyage_models.append(key)
voyage_models.add(key)
elif value.get("litellm_provider") == "infinity":
infinity_models.append(key)
infinity_models.add(key)
elif value.get("litellm_provider") == "databricks":
databricks_models.append(key)
databricks_models.add(key)
elif value.get("litellm_provider") == "cloudflare":
cloudflare_models.append(key)
cloudflare_models.add(key)
elif value.get("litellm_provider") == "codestral":
codestral_models.append(key)
codestral_models.add(key)
elif value.get("litellm_provider") == "friendliai":
friendliai_models.append(key)
friendliai_models.add(key)
elif value.get("litellm_provider") == "palm":
palm_models.append(key)
palm_models.add(key)
elif value.get("litellm_provider") == "groq":
groq_models.append(key)
groq_models.add(key)
elif value.get("litellm_provider") == "azure":
azure_models.append(key)
azure_models.add(key)
elif value.get("litellm_provider") == "anyscale":
anyscale_models.append(key)
anyscale_models.add(key)
elif value.get("litellm_provider") == "cerebras":
cerebras_models.append(key)
cerebras_models.add(key)
elif value.get("litellm_provider") == "galadriel":
galadriel_models.append(key)
galadriel_models.add(key)
elif value.get("litellm_provider") == "sambanova":
sambanova_models.append(key)
sambanova_models.add(key)
elif value.get("litellm_provider") == "sambanova-embedding-models":
sambanova_embedding_models.append(key)
sambanova_embedding_models.add(key)
elif value.get("litellm_provider") == "novita":
novita_models.append(key)
novita_models.add(key)
elif value.get("litellm_provider") == "nebius-chat-models":
nebius_models.append(key)
nebius_models.add(key)
elif value.get("litellm_provider") == "nebius-embedding-models":
nebius_embedding_models.append(key)
nebius_embedding_models.add(key)
elif value.get("litellm_provider") == "assemblyai":
assemblyai_models.append(key)
assemblyai_models.add(key)
elif value.get("litellm_provider") == "jina_ai":
jina_ai_models.append(key)
jina_ai_models.add(key)
elif value.get("litellm_provider") == "snowflake":
snowflake_models.append(key)
snowflake_models.add(key)
elif value.get("litellm_provider") == "gradient_ai":
gradient_ai_models.append(key)
gradient_ai_models.add(key)
elif value.get("litellm_provider") == "featherless_ai":
featherless_ai_models.append(key)
featherless_ai_models.add(key)
elif value.get("litellm_provider") == "deepgram":
deepgram_models.append(key)
deepgram_models.add(key)
elif value.get("litellm_provider") == "elevenlabs":
elevenlabs_models.append(key)
elevenlabs_models.add(key)
elif value.get("litellm_provider") == "dashscope":
dashscope_models.append(key)
dashscope_models.add(key)
elif value.get("litellm_provider") == "moonshot":
moonshot_models.append(key)
moonshot_models.add(key)
elif value.get("litellm_provider") == "v0":
v0_models.append(key)
v0_models.add(key)
elif value.get("litellm_provider") == "morph":
morph_models.append(key)
morph_models.add(key)
elif value.get("litellm_provider") == "lambda_ai":
lambda_ai_models.append(key)
lambda_ai_models.add(key)
elif value.get("litellm_provider") == "hyperbolic":
hyperbolic_models.append(key)
hyperbolic_models.add(key)
elif value.get("litellm_provider") == "recraft":
recraft_models.append(key)
recraft_models.add(key)
elif value.get("litellm_provider") == "cometapi":
cometapi_models.append(key)
cometapi_models.add(key)
elif value.get("litellm_provider") == "oci":
oci_models.append(key)
oci_models.add(key)
add_known_models()
@ -768,68 +770,68 @@ ollama_models = ["llama2"]
maritalk_models = ["maritalk"]
model_list = (
model_list = list(
open_ai_chat_completion_models
+ open_ai_text_completion_models
+ cohere_models
+ cohere_chat_models
+ anthropic_models
+ replicate_models
+ openrouter_models
+ datarobot_models
+ huggingface_models
+ vertex_chat_models
+ vertex_text_models
+ ai21_models
+ ai21_chat_models
+ together_ai_models
+ baseten_models
+ aleph_alpha_models
+ nlp_cloud_models
+ ollama_models
+ bedrock_models
+ deepinfra_models
+ perplexity_models
+ maritalk_models
+ vertex_language_models
+ watsonx_models
+ gemini_models
+ text_completion_codestral_models
+ xai_models
+ deepseek_models
+ azure_ai_models
+ voyage_models
+ infinity_models
+ databricks_models
+ cloudflare_models
+ codestral_models
+ friendliai_models
+ palm_models
+ groq_models
+ azure_models
+ anyscale_models
+ cerebras_models
+ galadriel_models
+ sambanova_models
+ azure_text_models
+ novita_models
+ assemblyai_models
+ jina_ai_models
+ snowflake_models
+ gradient_ai_models
+ llama_models
+ featherless_ai_models
+ nscale_models
+ deepgram_models
+ elevenlabs_models
+ dashscope_models
+ moonshot_models
+ v0_models
+ morph_models
+ lambda_ai_models
+ recraft_models
+ cometapi_models
+ oci_models
| open_ai_text_completion_models
| cohere_models
| cohere_chat_models
| anthropic_models
| set(replicate_models)
| openrouter_models
| datarobot_models
| set(huggingface_models)
| vertex_chat_models
| vertex_text_models
| ai21_models
| ai21_chat_models
| set(together_ai_models)
| set(baseten_models)
| aleph_alpha_models
| nlp_cloud_models
| set(ollama_models)
| bedrock_models
| deepinfra_models
| perplexity_models
| set(maritalk_models)
| vertex_language_models
| watsonx_models
| gemini_models
| text_completion_codestral_models
| xai_models
| deepseek_models
| azure_ai_models
| voyage_models
| infinity_models
| databricks_models
| cloudflare_models
| codestral_models
| friendliai_models
| palm_models
| groq_models
| azure_models
| anyscale_models
| cerebras_models
| galadriel_models
| sambanova_models
| azure_text_models
| novita_models
| assemblyai_models
| jina_ai_models
| snowflake_models
| gradient_ai_models
| llama_models
| featherless_ai_models
| nscale_models
| deepgram_models
| elevenlabs_models
| dashscope_models
| moonshot_models
| v0_models
| morph_models
| lambda_ai_models
| recraft_models
| cometapi_models
| oci_models
)
model_list_set = set(model_list)
@ -838,9 +840,9 @@ provider_list: List[Union[LlmProviders, str]] = list(LlmProviders)
models_by_provider: dict = {
"openai": open_ai_chat_completion_models + open_ai_text_completion_models,
"openai": open_ai_chat_completion_models | open_ai_text_completion_models,
"text-completion-openai": open_ai_text_completion_models,
"cohere": cohere_models + cohere_chat_models,
"cohere": cohere_models | cohere_chat_models,
"cohere_chat": cohere_chat_models,
"anthropic": anthropic_models,
"replicate": replicate_models,
@ -849,14 +851,9 @@ models_by_provider: dict = {
"baseten": baseten_models,
"openrouter": openrouter_models,
"datarobot": datarobot_models,
"vertex_ai": vertex_chat_models
+ vertex_text_models
+ vertex_anthropic_models
+ vertex_vision_models
+ vertex_language_models
+ vertex_deepseek_models,
"vertex_ai": vertex_chat_models | vertex_text_models | vertex_anthropic_models | vertex_vision_models | vertex_language_models | vertex_deepseek_models,
"ai21": ai21_models,
"bedrock": bedrock_models + bedrock_converse_models,
"bedrock": bedrock_models | bedrock_converse_models,
"petals": petals_models,
"ollama": ollama_models,
"ollama_chat": ollama_models,
@ -865,7 +862,7 @@ models_by_provider: dict = {
"maritalk": maritalk_models,
"watsonx": watsonx_models,
"gemini": gemini_models,
"fireworks_ai": fireworks_ai_models + fireworks_ai_embedding_models,
"fireworks_ai": fireworks_ai_models | fireworks_ai_embedding_models,
"aleph_alpha": aleph_alpha_models,
"text-completion-codestral": text_completion_codestral_models,
"xai": xai_models,
@ -881,14 +878,14 @@ models_by_provider: dict = {
"friendliai": friendliai_models,
"palm": palm_models,
"groq": groq_models,
"azure": azure_models + azure_text_models,
"azure": azure_models | azure_text_models,
"azure_text": azure_text_models,
"anyscale": anyscale_models,
"cerebras": cerebras_models,
"galadriel": galadriel_models,
"sambanova": sambanova_models + sambanova_embedding_models,
"sambanova": sambanova_models | sambanova_embedding_models,
"novita": novita_models,
"nebius": nebius_models + nebius_embedding_models,
"nebius": nebius_models | nebius_embedding_models,
"assemblyai": assemblyai_models,
"jina_ai": jina_ai_models,
"snowflake": snowflake_models,
@ -935,12 +932,12 @@ longer_context_model_fallback_dict: dict = {
all_embedding_models = (
open_ai_embedding_models
+ cohere_embedding_models
+ bedrock_embedding_models
+ vertex_embedding_models
+ fireworks_ai_embedding_models
+ nebius_embedding_models
+ sambanova_embedding_models
| set(cohere_embedding_models)
| set(bedrock_embedding_models)
| vertex_embedding_models
| fireworks_ai_embedding_models
| nebius_embedding_models
| sambanova_embedding_models
)
####### IMAGE GENERATION MODELS ###################
@ -1038,6 +1035,7 @@ from .llms.cohere.rerank_v2.transformation import CohereRerankV2Config
from .llms.azure_ai.rerank.transformation import AzureAIRerankConfig
from .llms.infinity.rerank.transformation import InfinityRerankConfig
from .llms.jina_ai.rerank.transformation import JinaAIRerankConfig
from .llms.deepinfra.rerank.transformation import DeepinfraRerankConfig
from .llms.clarifai.chat.transformation import ClarifaiConfig
from .llms.ai21.chat.transformation import AI21ChatConfig, AI21ChatConfig as AI21Config
from .llms.meta_llama.chat.transformation import LlamaAPIConfig
@ -1149,6 +1147,7 @@ from .llms.topaz.image_variations.transformation import TopazImageVariationConfi
from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig
from .llms.groq.chat.transformation import GroqChatConfig
from .llms.voyage.embedding.transformation import VoyageEmbeddingConfig
from .llms.voyage.embedding.transformation_contextual import VoyageContextualEmbeddingConfig
from .llms.infinity.embedding.transformation import InfinityEmbeddingConfig
from .llms.azure_ai.chat.transformation import AzureAIStudioConfig
from .llms.mistral.chat.transformation import MistralConfig
@ -1195,6 +1194,7 @@ nvidiaNimEmbeddingConfig = NvidiaNimEmbeddingConfig()
from .llms.featherless_ai.chat.transformation import FeatherlessAIConfig
from .llms.cerebras.chat import CerebrasConfig
from .llms.baseten.chat import BasetenConfig
from .llms.sambanova.chat import SambanovaConfig
from .llms.sambanova.embedding.transformation import SambaNovaEmbeddingConfig
from .llms.ai21.chat.transformation import AI21ChatConfig

4
litellm/caching/caching.py
View file

@ -774,11 +774,9 @@ class Cache:
"""
Internal method to check if the cache type supports async get/set operations
Only S3 Cache Does NOT support async operations
All cache types now support async operations
"""
if self.type and self.type == LiteLLMCacheType.S3:
return False
return True

8
litellm/caching/caching_handler.py
View file

@ -599,7 +599,7 @@ class LLMCachingHandler:
cached_result = await litellm.cache.async_get_cache(
dynamic_cache_object=self.dual_cache, **new_kwargs
)
else: # for s3 caching. [NOT RECOMMENDED IN PROD - this will slow down responses since boto3 is sync]
else: # fallback for caches that don't support async
cached_result = litellm.cache.get_cache(
dynamic_cache_object=self.dual_cache, **new_kwargs
)
@ -806,12 +806,6 @@ class LLMCachingHandler:
result, dynamic_cache_object=self.dual_cache, **new_kwargs
)
)
elif isinstance(litellm.cache.cache, S3Cache):
threading.Thread(
target=litellm.cache.add_cache,
args=(result,),
kwargs=new_kwargs,
).start()
else:
asyncio.create_task(
litellm.cache.async_add_cache(

38
litellm/caching/s3_cache.py
View file

@ -1,17 +1,17 @@
"""
S3 Cache implementation
WARNING: DO NOT USE THIS IN PRODUCTION - This is not ASYNC
Has 4 methods:
- set_cache
- get_cache
- async_set_cache
- async_get_cache
- async_set_cache (uses run_in_executor)
- async_get_cache (uses run_in_executor)
"""
import ast
import asyncio
import json
from functools import partial
from typing import Optional
from litellm._logging import print_verbose, verbose_logger
@ -72,7 +72,7 @@ class S3Cache(BaseCache):
import datetime
# Calculate expiration time
expiration_time = datetime.datetime.now() + ttl
expiration_time = datetime.datetime.now() + datetime.timedelta(seconds=ttl)
# Upload the data to S3 with the calculated expiration time
self.s3_client.put_object(
@ -98,11 +98,20 @@ class S3Cache(BaseCache):
ContentDisposition=f'inline; filename="{key}.json"',
)
except Exception as e:
# NON blocking - notify users S3 is throwing an exception
print_verbose(f"S3 Caching: set_cache() - Got exception from S3: {e}")
async def async_set_cache(self, key, value, **kwargs):
self.set_cache(key=key, value=value, **kwargs)
"""
Asynchronously set cache using run_in_executor to avoid blocking the event loop.
Compatible with Python 3.8+.
"""
try:
verbose_logger.debug(f"Set ASYNC S3 Cache: Key={key}. Value={value}")
loop = asyncio.get_event_loop()
func = partial(self.set_cache, key, value, **kwargs)
await loop.run_in_executor(None, func)
except Exception as e:
verbose_logger.error(f"S3 Caching: async_set_cache() - Got exception from S3: {e}")
def get_cache(self, key, **kwargs):
import botocore
@ -142,13 +151,26 @@ class S3Cache(BaseCache):
return None
except Exception as e:
# NON blocking - notify users S3 is throwing an exception
verbose_logger.error(
f"S3 Caching: get_cache() - Got exception from S3: {e}"
)
async def async_get_cache(self, key, **kwargs):
return self.get_cache(key=key, **kwargs)
"""
Asynchronously get cache using run_in_executor to avoid blocking the event loop.
Compatible with Python 3.8+.
"""
try:
verbose_logger.debug(f"Get ASYNC S3 Cache: key: {key}")
loop = asyncio.get_event_loop()
func = partial(self.get_cache, key, **kwargs)
result = await loop.run_in_executor(None, func)
return result
except Exception as e:
verbose_logger.error(
f"S3 Caching: async_get_cache() - Got exception from S3: {e}"
)
return None
def flush_cache(self):
pass

53
litellm/constants.py
View file

@ -251,6 +251,7 @@ LITELLM_CHAT_PROVIDERS = [
"groq",
"nvidia_nim",
"cerebras",
"baseten",
"ai21_chat",
"volcengine",
"codestral",
@ -427,6 +428,7 @@ openai_compatible_providers: List = [
"groq",
"nvidia_nim",
"cerebras",
"baseten",
"sambanova",
"ai21_chat",
"ai21",
@ -483,7 +485,7 @@ _openai_like_providers: List = [
"watsonx",
] # private helper. similar to openai but require some custom auth / endpoint handling, so can't use the openai sdk
# well supported replicate llms
replicate_models: List = [
replicate_models: set = set([
# llama replicate supported LLMs
"replicate/llama-2-70b-chat:2796ee9483c3fd7aa2e171d38f4ca12251a30609463dcfd4cd76703f22e96cdf",
"a16z-infra/llama-2-13b-chat:2a7f981751ec7fdf87b5b91ad4db53683a98082e9ff7bfd12c8cd5ea85980a52",
@ -496,9 +498,9 @@ replicate_models: List = [
# Others
"replicate/dolly-v2-12b:ef0e1aefc61f8e096ebe4db6b2bacc297daf2ef6899f0f7e001ec445893500e5",
"replit/replit-code-v1-3b:b84f4c074b807211cd75e3e8b1589b6399052125b4c27106e43d47189e8415ad",
]
])
clarifai_models: List = [
clarifai_models: set = set([
"clarifai/meta.Llama-3.Llama-3-8B-Instruct",
"clarifai/gcp.generate.gemma-1_1-7b-it",
"clarifai/mistralai.completion.mixtral-8x22B",
@ -562,10 +564,10 @@ clarifai_models: List = [
"clarifai/gcp.generate.gemini-1_5-pro",
"clarifai/gcp.generate.imagen-2",
"clarifai/salesforce.blip.general-english-image-caption-blip-2",
]
])
huggingface_models: List = [
huggingface_models: set = set([
"meta-llama/Llama-2-7b-hf",
"meta-llama/Llama-2-7b-chat-hf",
"meta-llama/Llama-2-13b-hf",
@ -578,13 +580,13 @@ huggingface_models: List = [
"meta-llama/Llama-2-13b-chat",
"meta-llama/Llama-2-70b",
"meta-llama/Llama-2-70b-chat",
] # these have been tested on extensively. But by default all text2text-generation and text-generation models are supported by liteLLM. - https://docs.litellm.ai/docs/providers
empower_models = [
]) # these have been tested on extensively. But by default all text2text-generation and text-generation models are supported by liteLLM. - https://docs.litellm.ai/docs/providers
empower_models = set([
"empower/empower-functions",
"empower/empower-functions-small",
]
])
together_ai_models: List = [
together_ai_models: set = set([
# llama llms - chat
"togethercomputer/llama-2-70b-chat",
# llama llms - language / instruct
@ -612,16 +614,17 @@ together_ai_models: List = [
"Austism/chronos-hermes-13b",
"upstage/SOLAR-0-70b-16bit",
"WizardLM/WizardLM-70B-V1.0",
] # supports all together ai models, just pass in the model id e.g. completion(model="together_computer/replit_code_3b",...)
])
# supports all together ai models, just pass in the model id e.g. completion(model="together_computer/replit_code_3b",...)
baseten_models: List = [
baseten_models: set = set([
"qvv0xeq",
"q841o8w",
"31dxrj3",
] # FALCON 7B # WizardLM # Mosaic ML
]) # FALCON 7B # WizardLM # Mosaic ML
featherless_ai_models: List = [
featherless_ai_models: set = set([
"featherless-ai/Qwerky-72B",
"featherless-ai/Qwerky-QwQ-32B",
"Qwen/Qwen2.5-72B-Instruct",
@ -631,9 +634,9 @@ featherless_ai_models: List = [
"mistralai/Mistral-Small-24B-Instruct-2501",
"mistralai/Mistral-Nemo-Instruct-2407",
"ProdeusUnity/Stellar-Odyssey-12b-v0.0",
]
])
nebius_models: List = [
nebius_models: set = set([
"Qwen/Qwen3-235B-A22B",
"Qwen/Qwen3-30B-A3B-fast",
"Qwen/Qwen3-32B",
@ -646,9 +649,9 @@ nebius_models: List = [
"meta-llama/Llama-3.3-70B-Instruct-fast",
"Qwen/Qwen2.5-32B-Instruct-fast",
"Qwen/Qwen2.5-Coder-32B-Instruct-fast",
]
])
dashscope_models: List = [
dashscope_models: set = set([
"qwen-turbo",
"qwen-plus",
"qwen-max",
@ -659,13 +662,13 @@ dashscope_models: List = [
"qwen3-235b-a22b",
"qwen3-32b",
"qwen3-30b-a3b",
]
])
nebius_embedding_models: List = [
nebius_embedding_models: set = set([
"BAAI/bge-en-icl",
"BAAI/bge-multilingual-gemma2",
"intfloat/e5-mistral-7b-instruct",
]
])
BEDROCK_INVOKE_PROVIDERS_LITERAL = Literal[
"cohere",
@ -679,8 +682,8 @@ BEDROCK_INVOKE_PROVIDERS_LITERAL = Literal[
"deepseek_r1",
]
open_ai_embedding_models: List = ["text-embedding-ada-002"]
cohere_embedding_models: List = [
open_ai_embedding_models: set = set(["text-embedding-ada-002"])
cohere_embedding_models: set = set([
"embed-v4.0",
"embed-english-v3.0",
"embed-english-light-v3.0",
@ -688,12 +691,12 @@ cohere_embedding_models: List = [
"embed-english-v2.0",
"embed-english-light-v2.0",
"embed-multilingual-v2.0",
]
bedrock_embedding_models: List = [
])
bedrock_embedding_models: set = set([
"amazon.titan-embed-text-v1",
"cohere.embed-english-v3",
"cohere.embed-multilingual-v3",
]
])
known_tokenizer_config = {
"mistralai/Mistral-7B-Instruct-v0.1": {

22
litellm/integrations/braintrust_logging.py
View file

@ -19,10 +19,6 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.utils import print_verbose
global_braintrust_http_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
global_braintrust_sync_http_handler = HTTPHandler()
API_BASE = "https://api.braintrustdata.com/v1"
@ -52,6 +48,10 @@ class BraintrustLogger(CustomLogger):
self._project_id_cache: Dict[
str, str
] = {} # Cache mapping project names to IDs
self.global_braintrust_http_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
self.global_braintrust_sync_http_handler = HTTPHandler()
def validate_environment(self, api_key: Optional[str]):
"""
@ -76,7 +76,7 @@ class BraintrustLogger(CustomLogger):
return self._project_id_cache[project_name]
try:
response = global_braintrust_sync_http_handler.post(
response = self.global_braintrust_sync_http_handler.post(
f"{self.api_base}/project",
headers=self.headers,
json={"name": project_name},
@ -96,7 +96,7 @@ class BraintrustLogger(CustomLogger):
return self._project_id_cache[project_name]
try:
response = await global_braintrust_http_handler.post(
response = await self.global_braintrust_http_handler.post(
f"{self.api_base}/project/register",
headers=self.headers,
json={"name": project_name},
@ -146,7 +146,7 @@ class BraintrustLogger(CustomLogger):
return metadata
async def create_default_project_and_experiment(self):
project = await global_braintrust_http_handler.post(
project = await self.global_braintrust_http_handler.post(
f"{self.api_base}/project", headers=self.headers, json={"name": "litellm"}
)
@ -155,7 +155,7 @@ class BraintrustLogger(CustomLogger):
self.default_project_id = project_dict["id"]
def create_sync_default_project_and_experiment(self):
project = global_braintrust_sync_http_handler.post(
project = self.global_braintrust_sync_http_handler.post(
f"{self.api_base}/project", headers=self.headers, json={"name": "litellm"}
)
@ -291,9 +291,9 @@ class BraintrustLogger(CustomLogger):
try:
print_verbose(
f"global_braintrust_sync_http_handler.post: {global_braintrust_sync_http_handler.post}"
f"self.global_braintrust_sync_http_handler.post: {self.global_braintrust_sync_http_handler.post}"
)
global_braintrust_sync_http_handler.post(
self.global_braintrust_sync_http_handler.post(
url=f"{self.api_base}/project_logs/{project_id}/insert",
json={"events": [request_data]},
headers=self.headers,
@ -446,7 +446,7 @@ class BraintrustLogger(CustomLogger):
request_data["metrics"] = metrics
try:
await global_braintrust_http_handler.post(
await self.global_braintrust_http_handler.post(
url=f"{self.api_base}/project_logs/{project_id}/insert",
json={"events": [request_data]},
headers=self.headers,

11
litellm/integrations/mlflow.py
View file

@ -60,10 +60,7 @@ class MlflowLogger(CustomLogger):
inputs = self._construct_input(kwargs)
input_messages = inputs.get("messages", [])
output_messages = [
c.message.model_dump(exclude_none=True)
for c in getattr(response_obj, "choices", [])
]
output_messages = [c.message.model_dump(exclude_none=True) for c in getattr(response_obj, "choices", [])]
if messages := [*input_messages, *output_messages]:
set_span_chat_messages(span, messages)
if tools := inputs.get("tools"):
@ -168,6 +165,10 @@ class MlflowLogger(CustomLogger):
for key in ["functions", "tools", "stream", "tool_choice", "user"]:
if value := kwargs.get("optional_params", {}).pop(key, None):
inputs[key] = value
if prediction := kwargs.get("prediction"):
inputs["prediction"] = prediction
return inputs
def _extract_attributes(self, kwargs):
@ -232,7 +233,6 @@ class MlflowLogger(CustomLogger):
"""
import mlflow
call_type = kwargs.get("call_type", "completion")
span_name = f"litellm-{call_type}"
span_type = self._get_span_type(call_type)
@ -260,6 +260,7 @@ class MlflowLogger(CustomLogger):
tags=self._transform_tag_list_to_dict(attributes.get("request_tags", [])),
start_time_ns=start_time_ns,
)
def _transform_tag_list_to_dict(self, tag_list: list) -> dict:
return {tag: "" for tag in tag_list}

10
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -196,6 +196,9 @@ def get_llm_provider( # noqa: PLR0915
elif endpoint == "https://api.cerebras.ai/v1":
custom_llm_provider = "cerebras"
dynamic_api_key = get_secret_str("CEREBRAS_API_KEY")
elif endpoint == "https://inference.baseten.co/v1":
custom_llm_provider = "baseten"
dynamic_api_key = get_secret_str("BASETEN_API_KEY")
elif endpoint == "https://api.sambanova.ai/v1":
custom_llm_provider = "sambanova"
dynamic_api_key = get_secret_str("SAMBANOVA_API_KEY")
@ -478,6 +481,13 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915
api_base or get_secret("CEREBRAS_API_BASE") or "https://api.cerebras.ai/v1"
) # type: ignore
dynamic_api_key = api_key or get_secret_str("CEREBRAS_API_KEY")
elif custom_llm_provider == "baseten":
# Use BasetenConfig to determine the appropriate API base URL
if api_base is None:
api_base = litellm.BasetenConfig.get_api_base_for_model(model)
else:
api_base = api_base or get_secret("BASETEN_API_BASE") or "https://inference.baseten.co/v1"
dynamic_api_key = api_key or get_secret_str("BASETEN_API_KEY")
elif custom_llm_provider == "sambanova":
api_base = (
api_base

2
litellm/litellm_core_utils/get_supported_openai_params.py
View file

@ -78,6 +78,8 @@ def get_supported_openai_params( # noqa: PLR0915
return litellm.nvidiaNimEmbeddingConfig.get_supported_openai_params()
elif custom_llm_provider == "cerebras":
return litellm.CerebrasConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "baseten":
return litellm.BasetenConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "xai":
return litellm.XAIChatConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "ai21_chat" or custom_llm_provider == "ai21":

11
litellm/litellm_core_utils/litellm_logging.py
View file

@ -4106,18 +4106,9 @@ class StandardLoggingPayloadSetup:
"""
# Generate object key in same format as S3Logger
from litellm.integrations.s3 import get_s3_object_key
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
# Only generate object key if cold storage is configured
try:
configured_cold_storage_logger = (
ColdStorageHandler._get_configured_cold_storage_custom_logger()
)
except Exception as e:
verbose_logger.debug(f"Cold storage custom logger unavailable: {e}")
return None
if configured_cold_storage_logger is None:
if litellm.configured_cold_storage_logger is None:
return None
try:

119
litellm/litellm_core_utils/logging_utils.py
View file

@ -1,5 +1,6 @@
import asyncio
import functools
import time
from datetime import datetime
from typing import TYPE_CHECKING, Any, List, Optional, Union
@ -11,15 +12,19 @@ from litellm.types.utils import (
)
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm import ModelResponse as _ModelResponse
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObject,
)
LiteLLMModelResponse = _ModelResponse
Span = Union[_Span, Any]
else:
LiteLLMModelResponse = Any
LiteLLMLoggingObject = Any
Span = Any
import litellm
@ -28,9 +33,52 @@ import litellm
Helper utils used for logging callbacks
"""
# Global service logger instance to avoid recreating it
_service_logger = None
def _get_service_logger():
"""Get or create the global ServiceLogging instance"""
global _service_logger
if _service_logger is None:
from litellm._service_logger import ServiceLogging
_service_logger = ServiceLogging()
return _service_logger
def _get_parent_otel_span_from_logging_obj(
logging_obj: Optional[LiteLLMLoggingObject] = None,
) -> Optional[Span]:
"""
Extract the parent OTEL span from the logging object using existing helper.
Args:
logging_obj: The LiteLLM logging object containing model call details
Returns:
The parent OTEL span if found, None otherwise
"""
try:
if logging_obj is None or not hasattr(logging_obj, "model_call_details"):
return None
# Reuse existing function by passing model_call_details as kwargs
from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
)
return _get_parent_otel_span_from_kwargs(logging_obj.model_call_details)
except Exception as e:
verbose_logger.exception(
f"Error in _get_parent_otel_span_from_logging_obj: {str(e)}"
)
return None
def convert_litellm_response_object_to_str(
response_obj: Union[Any, LiteLLMModelResponse]
response_obj: Union[Any, LiteLLMModelResponse],
) -> Optional[str]:
"""
Get the string of the response object from LiteLLM
@ -125,37 +173,102 @@ def track_llm_api_timing():
"""
Decorator to track LLM API call timing for both sync and async functions.
The logging_obj is expected to be passed as an argument to the decorated function.
Logs timing using ServiceLogging similar to Redis cache.
"""
def decorator(func):
@functools.wraps(func)
async def async_wrapper(*args, **kwargs):
start_time = datetime.now()
start_time_float = time.time()
logging_obj = kwargs.get("logging_obj", None)
# Extract parent OTEL span from logging object
parent_otel_span = _get_parent_otel_span_from_logging_obj(logging_obj)
try:
result = await func(*args, **kwargs)
return result
finally:
end_time = datetime.now()
end_time_float = time.time()
duration = end_time_float - start_time_float
# Set duration in model call details
_set_duration_in_model_call_details(
logging_obj=kwargs.get("logging_obj", None),
logging_obj=logging_obj,
start_time=start_time,
end_time=end_time,
)
# Log timing using ServiceLogging (like Redis cache)
try:
from litellm.types.services import ServiceTypes
service_logger = _get_service_logger()
# Get function name for call_type
call_type = f"{func.__name__} <- track_llm_api_timing"
# Create async task for service logging (similar to Redis cache pattern)
asyncio.create_task(
service_logger.async_service_success_hook(
service=ServiceTypes.LITELLM,
duration=duration,
call_type=call_type,
start_time=start_time_float,
end_time=end_time_float,
parent_otel_span=parent_otel_span,
)
)
except Exception as e:
verbose_logger.debug(f"Error in service logging: {str(e)}")
@functools.wraps(func)
def sync_wrapper(*args, **kwargs):
start_time = datetime.now()
start_time_float = time.time()
logging_obj = kwargs.get("logging_obj", None)
# Extract parent OTEL span from logging object
parent_otel_span = _get_parent_otel_span_from_logging_obj(logging_obj)
try:
result = func(*args, **kwargs)
return result
finally:
end_time = datetime.now()
end_time_float = time.time()
duration = end_time_float - start_time_float
# Set duration in model call details
_set_duration_in_model_call_details(
logging_obj=kwargs.get("logging_obj", None),
logging_obj=logging_obj,
start_time=start_time,
end_time=end_time,
)
# Log timing using ServiceLogging (like Redis cache)
try:
from litellm.types.services import ServiceTypes
service_logger = _get_service_logger()
# Get function name for call_type
call_type = f"{func.__name__} <- track_llm_api_timing"
# Use sync service logging for sync functions
service_logger.service_success_hook(
service=ServiceTypes.LITELLM,
duration=duration,
call_type=call_type,
start_time=start_time_float,
end_time=end_time_float,
parent_otel_span=parent_otel_span,
)
except Exception as e:
verbose_logger.debug(f"Error in service logging: {str(e)}")
# Check if the function is async or sync
if asyncio.iscoroutinefunction(func):
return async_wrapper

51
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -3193,9 +3193,30 @@ class BedrockConverseMessagesProcessor:
## MERGE CONSECUTIVE TOOL CALL MESSAGES ##
tool_content: List[BedrockContentBlock] = []
while msg_i < len(messages) and messages[msg_i]["role"] == "tool":
tool_call_result = _convert_to_bedrock_tool_call_result(messages[msg_i])
current_message = messages[msg_i]
tool_call_result = _convert_to_bedrock_tool_call_result(current_message)
tool_content.append(tool_call_result)
# Check if we need to add a separate cachePoint block
has_cache_control = False
# Check for message-level cache_control
if current_message.get("cache_control", None) is not None:
has_cache_control = True
# Check for content-level cache_control in list content
elif isinstance(current_message.get("content"), list):
for content_element in current_message["content"]:
if (isinstance(content_element, dict) and
content_element.get("cache_control", None) is not None):
has_cache_control = True
break
# Add a separate cachePoint block if cache_control is present
if has_cache_control:
cache_point_block = BedrockContentBlock(cachePoint=CachePointBlock(type="default"))
tool_content.append(cache_point_block)
msg_i += 1
if tool_content:
# if last message was a 'user' message, then add a blank assistant message (bedrock requires alternating roles)
@ -3275,13 +3296,29 @@ class BedrockConverseMessagesProcessor:
image_url=image_url
)
assistants_parts.append(assistants_part)
# Add cache point block for assistant content elements
_cache_point_block = (
litellm.AmazonConverseConfig()._get_cache_point_block(
message_block=cast(
OpenAIMessageContentListBlock, element
),
block_type="content_block",
)
)
if _cache_point_block is not None:
assistants_parts.append(_cache_point_block)
assistant_content.extend(assistants_parts)
elif _assistant_content is not None and isinstance(
_assistant_content, str
):
assistant_content.append(
BedrockContentBlock(text=_assistant_content)
elif _assistant_content is not None and isinstance(_assistant_content, str):
assistant_content.append(BedrockContentBlock(text=_assistant_content))
# Add cache point block for assistant string content
_cache_point_block = (
litellm.AmazonConverseConfig()._get_cache_point_block(
assistant_message_block, block_type="content_block"
)
)
if _cache_point_block is not None:
assistant_content.append(_cache_point_block)
_tool_calls = assistant_message_block.get("tool_calls", [])
if _tool_calls:
assistant_content.extend(

5
litellm/llms/azure/responses/transformation.py
View file

@ -6,6 +6,7 @@ from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfi
from litellm.types.llms.openai import *
from litellm.types.responses.main import *
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
@ -16,6 +17,10 @@ else:
class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.AZURE
def validate_environment(
self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]
) -> dict:

6
litellm/llms/base_llm/responses/transformation.py
View file

@ -12,6 +12,7 @@ from litellm.types.llms.openai import (
)
from litellm.types.responses.main import *
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
@ -29,6 +30,11 @@ class BaseResponsesAPIConfig(ABC):
def __init__(self):
pass
@property
@abstractmethod
def custom_llm_provider(self) -> LlmProviders:
pass
@classmethod
def get_config(cls):
return {

172
litellm/llms/baseten.py
View file

@ -1,172 +0,0 @@
import json
import time
from typing import Callable
import litellm
from litellm.types.utils import ModelResponse, Usage
class BasetenError(Exception):
def __init__(self, status_code, message):
self.status_code = status_code
self.message = message
super().__init__(
self.message
) # Call the base class constructor with the parameters it needs
def validate_environment(api_key):
headers = {
"accept": "application/json",
"content-type": "application/json",
}
if api_key:
headers["Authorization"] = f"Api-Key {api_key}"
return headers
def completion(
model: str,
messages: list,
model_response: ModelResponse,
print_verbose: Callable,
encoding,
api_key,
logging_obj,
optional_params: dict,
litellm_params=None,
logger_fn=None,
):
headers = validate_environment(api_key)
completion_url_fragment_1 = "https://app.baseten.co/models/"
completion_url_fragment_2 = "/predict"
model = model
prompt = ""
for message in messages:
if "role" in message:
if message["role"] == "user":
prompt += f"{message['content']}"
else:
prompt += f"{message['content']}"
else:
prompt += f"{message['content']}"
data = {
"inputs": prompt,
"prompt": prompt,
"parameters": optional_params,
"stream": (
True
if "stream" in optional_params and optional_params["stream"] is True
else False
),
}
## LOGGING
logging_obj.pre_call(
input=prompt,
api_key=api_key,
additional_args={"complete_input_dict": data},
)
## COMPLETION CALL
response = litellm.module_level_client.post(
completion_url_fragment_1 + model + completion_url_fragment_2,
headers=headers,
data=json.dumps(data),
stream=(
True
if "stream" in optional_params and optional_params["stream"] is True
else False
),
)
if "text/event-stream" in response.headers["Content-Type"] or (
"stream" in optional_params and optional_params["stream"] is True
):
return response.iter_lines()
else:
## LOGGING
logging_obj.post_call(
input=prompt,
api_key=api_key,
original_response=response.text,
additional_args={"complete_input_dict": data},
)
print_verbose(f"raw model_response: {response.text}")
## RESPONSE OBJECT
completion_response = response.json()
if "error" in completion_response:
raise BasetenError(
message=completion_response["error"],
status_code=response.status_code,
)
else:
if "model_output" in completion_response:
if (
isinstance(completion_response["model_output"], dict)
and "data" in completion_response["model_output"]
and isinstance(completion_response["model_output"]["data"], list)
):
model_response.choices[0].message.content = completion_response[ # type: ignore
"model_output"
][
"data"
][
0
]
elif isinstance(completion_response["model_output"], str):
model_response.choices[0].message.content = completion_response[ # type: ignore
"model_output"
]
elif "completion" in completion_response and isinstance(
completion_response["completion"], str
):
model_response.choices[0].message.content = completion_response[ # type: ignore
"completion"
]
elif isinstance(completion_response, list) and len(completion_response) > 0:
if "generated_text" not in completion_response:
raise BasetenError(
message=f"Unable to parse response. Original response: {response.text}",
status_code=response.status_code,
)
model_response.choices[0].message.content = completion_response[0][ # type: ignore
"generated_text"
]
## GETTING LOGPROBS
if (
"details" in completion_response[0]
and "tokens" in completion_response[0]["details"]
):
model_response.choices[0].finish_reason = completion_response[0][
"details"
]["finish_reason"]
sum_logprob = 0
for token in completion_response[0]["details"]["tokens"]:
sum_logprob += token["logprob"]
model_response.choices[0].logprobs = sum_logprob # type: ignore
else:
raise BasetenError(
message=f"Unable to parse response. Original response: {response.text}",
status_code=response.status_code,
)
## CALCULATING USAGE - baseten charges on time, not tokens - have some mapping of cost here.
prompt_tokens = len(encoding.encode(prompt))
completion_tokens = len(
encoding.encode(model_response["choices"][0]["message"]["content"])
)
model_response.created = int(time.time())
model_response.model = model
usage = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
)
setattr(model_response, "usage", usage)
return model_response
def embedding():
# logic for parsing in - calling - parsing out model embedding calls
pass

118
litellm/llms/baseten/chat.py Normal file
View file

@ -0,0 +1,118 @@
from typing import Optional
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
class BasetenConfig(OpenAIGPTConfig):
"""
Reference: https://inference.baseten.co/v1
Below are the parameters:
"""
max_tokens: Optional[int] = None
response_format: Optional[dict] = None
seed: Optional[int] = None
stream: Optional[bool] = None
top_p: Optional[int] = None
tool_choice: Optional[str] = None
tools: Optional[list] = None
user: Optional[str] = None
presence_penalty: Optional[int] = None
frequency_penalty: Optional[int] = None
stream_options: Optional[dict] = None
def __init__(
self,
max_tokens: Optional[int] = None,
response_format: Optional[dict] = None,
seed: Optional[int] = None,
stop: Optional[list] = None,
stream: Optional[bool] = None,
temperature: Optional[float] = None,
top_p: Optional[int] = None,
tool_choice: Optional[str] = None,
tools: Optional[list] = None,
user: Optional[str] = None,
presence_penalty: Optional[int] = None,
frequency_penalty: Optional[int] = None,
stream_options: Optional[dict] = None,
) -> None:
locals_ = locals().copy()
for key, value in locals_.items():
if key != "self" and value is not None:
setattr(self.__class__, key, value)
@classmethod
def get_config(cls):
return super().get_config()
def get_supported_openai_params(self, model: str) -> list:
"""
Get the supported OpenAI params for the given model
"""
return [
"max_tokens",
"max_completion_tokens",
"response_format",
"seed",
"stop",
"stream",
"temperature",
"top_p",
"tool_choice",
"tools",
"user",
"presence_penalty",
"frequency_penalty",
"stream_options",
]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
supported_openai_params = self.get_supported_openai_params(model=model)
for param, value in non_default_params.items():
if param == "max_completion_tokens":
optional_params["max_tokens"] = value
elif param in supported_openai_params:
optional_params[param] = value
return optional_params
def _get_openai_compatible_provider_info(self, api_base: str, api_key: str) -> tuple:
"""
Get the OpenAI compatible provider info for Baseten
"""
# Default to Model API
default_api_base = "https://inference.baseten.co/v1"
default_api_key = api_key or "BASETEN_API_KEY"
return default_api_base, default_api_key
@staticmethod
def is_dedicated_deployment(model: str) -> bool:
"""
Check if the model is a dedicated deployment (8-digit alphanumeric code)
"""
# Remove 'baseten/' prefix if present
model_id = model.replace("baseten/", "")
# Check if it's an 8-digit alphanumeric code
import re
return bool(re.match(r'^[a-zA-Z0-9]{8}$', model_id))
@staticmethod
def get_api_base_for_model(model: str) -> str:
"""
Get the appropriate API base URL for the given model
"""
if BasetenConfig.is_dedicated_deployment(model):
# Extract the model ID (remove 'baseten/' prefix if present)
model_id = model.replace("baseten/", "")
return f"https://model-{model_id}.api.baseten.co/environments/production/sync/v1"
else:
# Use Model API
return "https://inference.baseten.co/v1"

179
litellm/llms/bedrock/base_aws_llm.py
View file

@ -179,15 +179,32 @@ class BaseAWSLLM:
aws_sts_endpoint=aws_sts_endpoint,
)
elif aws_role_name is not None:
# If aws_session_name is not provided, generate a default one
if aws_session_name is None:
aws_session_name = f"litellm-session-{int(datetime.now().timestamp())}"
credentials, _cache_ttl = self._auth_with_aws_role(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_role_name=aws_role_name,
aws_session_name=aws_session_name,
)
# Check if we're in IRSA and trying to assume the same role we already have
current_role_arn = os.getenv("AWS_ROLE_ARN")
web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
# In IRSA environments, we should skip role assumption if we're already running as the target role
# This is true when:
# 1. We have AWS_ROLE_ARN set (current role)
# 2. We have AWS_WEB_IDENTITY_TOKEN_FILE set (IRSA environment)
# 3. The current role matches the requested role
if (current_role_arn and web_identity_token_file and
current_role_arn == aws_role_name):
verbose_logger.debug("Using IRSA same-role optimization: calling _auth_with_env_vars")
# We're already running as this role via IRSA, no need to assume it again
# Use the default boto3 credentials (which will use the IRSA credentials)
credentials, _cache_ttl = self._auth_with_env_vars()
else:
verbose_logger.debug("Using role assumption: calling _auth_with_aws_role")
# If aws_session_name is not provided, generate a default one
if aws_session_name is None:
aws_session_name = f"litellm-session-{int(datetime.now().timestamp())}"
credentials, _cache_ttl = self._auth_with_aws_role(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_role_name=aws_role_name,
aws_session_name=aws_session_name,
)
elif aws_profile_name is not None: ### CHECK SESSION ###
credentials, _cache_ttl = self._auth_with_aws_profile(aws_profile_name)
@ -446,6 +463,92 @@ class BaseAWSLLM:
iam_creds = session.get_credentials()
return iam_creds, self._get_default_ttl_for_boto3_credentials()
def _handle_irsa_cross_account(self, irsa_role_arn: str, aws_role_name: str,
aws_session_name: str, region: str, web_identity_token_file: str) -> dict:
"""Handle cross-account role assumption for IRSA."""
import boto3
verbose_logger.debug("Cross-account role assumption detected")
# Read the web identity token
with open(web_identity_token_file, 'r') as f:
web_identity_token = f.read().strip()
# Create an STS client without credentials
with tracer.trace("boto3.client(sts) for manual IRSA"):
sts_client = boto3.client('sts', region_name=region)
# Manually assume the IRSA role with the session name
verbose_logger.debug(f"Manually assuming IRSA role {irsa_role_arn} with session {aws_session_name}")
irsa_response = sts_client.assume_role_with_web_identity(
RoleArn=irsa_role_arn,
RoleSessionName=aws_session_name,
WebIdentityToken=web_identity_token
)
# Extract the credentials from the IRSA assumption
irsa_creds = irsa_response["Credentials"]
# Create a new STS client with the IRSA credentials
with tracer.trace("boto3.client(sts) with manual IRSA credentials"):
sts_client_with_creds = boto3.client(
'sts',
region_name=region,
aws_access_key_id=irsa_creds["AccessKeyId"],
aws_secret_access_key=irsa_creds["SecretAccessKey"],
aws_session_token=irsa_creds["SessionToken"]
)
# Get current caller identity for debugging
try:
caller_identity = sts_client_with_creds.get_caller_identity()
verbose_logger.debug(f"Current identity after manual IRSA assumption: {caller_identity.get('Arn', 'unknown')}")
except Exception as e:
verbose_logger.debug(f"Failed to get caller identity: {e}")
# Now assume the target role
verbose_logger.debug(f"Attempting to assume target role: {aws_role_name} with session: {aws_session_name}")
return sts_client_with_creds.assume_role(
RoleArn=aws_role_name, RoleSessionName=aws_session_name
)
def _handle_irsa_same_account(self, aws_role_name: str, aws_session_name: str, region: str) -> dict:
"""Handle same-account role assumption for IRSA."""
import boto3
verbose_logger.debug("Same account role assumption, using automatic IRSA")
with tracer.trace("boto3.client(sts) with automatic IRSA"):
sts_client = boto3.client("sts", region_name=region)
# Get current caller identity for debugging
try:
caller_identity = sts_client.get_caller_identity()
verbose_logger.debug(f"Current IRSA identity: {caller_identity.get('Arn', 'unknown')}")
except Exception as e:
verbose_logger.debug(f"Failed to get caller identity: {e}")
# Assume the role
verbose_logger.debug(f"Attempting to assume role: {aws_role_name} with session: {aws_session_name}")
return sts_client.assume_role(
RoleArn=aws_role_name, RoleSessionName=aws_session_name
)
def _extract_credentials_and_ttl(self, sts_response: dict) -> Tuple[Credentials, Optional[int]]:
"""Extract credentials and TTL from STS response."""
from botocore.credentials import Credentials
sts_credentials = sts_response["Credentials"]
credentials = Credentials(
access_key=sts_credentials["AccessKeyId"],
secret_key=sts_credentials["SecretAccessKey"],
token=sts_credentials["SessionToken"],
)
expiration_time = sts_credentials["Expiration"]
ttl = int((expiration_time - datetime.now(expiration_time.tzinfo)).total_seconds())
return credentials, ttl
@tracer.wrap()
def _auth_with_aws_role(
self,
@ -460,12 +563,58 @@ class BaseAWSLLM:
import boto3
from botocore.credentials import Credentials
with tracer.trace("boto3.client(sts)"):
sts_client = boto3.client(
"sts",
aws_access_key_id=aws_access_key_id, # [OPTIONAL]
aws_secret_access_key=aws_secret_access_key, # [OPTIONAL]
)
# Check if we're in an EKS/IRSA environment
web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
irsa_role_arn = os.getenv("AWS_ROLE_ARN")
# If we have IRSA environment variables and no explicit credentials,
# we need to use the web identity token flow
if (web_identity_token_file and irsa_role_arn and
aws_access_key_id is None and aws_secret_access_key is None):
# For cross-account role assumption with specific session names,
# we need to manually assume the IRSA role first with the correct session name
verbose_logger.debug(f"IRSA detected: using web identity token from {web_identity_token_file}")
try:
# Get region from environment
region = os.getenv("AWS_REGION") or os.getenv("AWS_DEFAULT_REGION") or "us-east-1"
# Check if we need to do cross-account role assumption
if aws_role_name != irsa_role_arn:
sts_response = self._handle_irsa_cross_account(
irsa_role_arn, aws_role_name, aws_session_name, region, web_identity_token_file
)
else:
sts_response = self._handle_irsa_same_account(
aws_role_name, aws_session_name, region
)
return self._extract_credentials_and_ttl(sts_response)
except Exception as e:
verbose_logger.debug(f"Failed to assume role via IRSA: {e}")
if "AccessDenied" in str(e) and "is not authorized to perform: sts:AssumeRole" in str(e):
# Provide a more helpful error message for trust policy issues
verbose_logger.error(
f"Access denied when trying to assume role {aws_role_name}. "
f"Please ensure the trust policy of {aws_role_name} allows "
f"the current role to assume it. Current identity: check logs with verbose mode."
)
# Re-raise the exception instead of falling through
raise
# In EKS/IRSA environments, use ambient credentials (no explicit keys needed)
# This allows the web identity token to work automatically
if aws_access_key_id is None and aws_secret_access_key is None:
with tracer.trace("boto3.client(sts)"):
sts_client = boto3.client("sts")
else:
with tracer.trace("boto3.client(sts)"):
sts_client = boto3.client(
"sts",
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
)
sts_response = sts_client.assume_role(
RoleArn=aws_role_name, RoleSessionName=aws_session_name

28
litellm/llms/custom_httpx/aiohttp_transport.py
View file

@ -3,7 +3,7 @@ import contextlib
import os
import typing
import urllib.request
from typing import Callable, Dict, Union
from typing import Callable, Dict, Optional, Union
import aiohttp
import aiohttp.client_exceptions
@ -115,6 +115,12 @@ class AiohttpTransport(httpx.AsyncBaseTransport):
) -> None:
self.client = client
#########################################################
# Class variables for proxy settings
#########################################################
self.proxy: Optional[str] = None
self.checked_proxy_env_settings: bool = False
async def aclose(self) -> None:
if isinstance(self.client, ClientSession):
await self.client.close()
@ -249,7 +255,22 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
def _proxy_from_env(self, url: httpx.URL) -> typing.Optional[str]:
"""Return proxy URL from env for the given request URL."""
"""
Return proxy URL from env for the given request URL
Only check the proxy env settings once, this is a costly operation for CPU % usage
."""
#########################################################
# Check if we've already checked the proxy env settings
#########################################################
if self.checked_proxy_env_settings is True:
return self.proxy
#########################################################
# set self.checked_proxy_env_settings to True
#########################################################
self.checked_proxy_env_settings = True
proxies = urllib.request.getproxies()
if urllib.request.proxy_bypass(url.host):
return None
@ -257,4 +278,5 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
proxy = proxies.get(url.scheme) or proxies.get("all")
if proxy and "://" not in proxy:
proxy = f"http://{proxy}"
return proxy
self.proxy = proxy
return self.proxy

26
litellm/llms/custom_httpx/http_handler.py
View file

@ -40,7 +40,9 @@ headers = {
_DEFAULT_TIMEOUT = httpx.Timeout(timeout=5.0, connect=5.0)
def get_ssl_configuration(ssl_verify: Optional[VerifyTypes] = None) -> Union[bool, str, ssl.SSLContext]:
def get_ssl_configuration(
ssl_verify: Optional[VerifyTypes] = None,
) -> Union[bool, str, ssl.SSLContext]:
"""
Unified SSL configuration function that handles ssl_context and ssl_verify logic.
@ -59,7 +61,7 @@ def get_ssl_configuration(ssl_verify: Optional[VerifyTypes] = None) -> Union[boo
- False: Disable SSL verification
- True: Enable SSL verification
- str: Path to CA bundle file
Returns:
Union[bool, str, ssl.SSLContext]: Appropriate SSL configuration
"""
@ -72,7 +74,9 @@ def get_ssl_configuration(ssl_verify: Optional[VerifyTypes] = None) -> Union[boo
# Get ssl_verify from environment or litellm settings if not provided
if ssl_verify is None:
ssl_verify = os.getenv("SSL_VERIFY", litellm.ssl_verify)
ssl_verify_bool = str_to_bool(ssl_verify) if isinstance(ssl_verify, str) else ssl_verify
ssl_verify_bool = (
str_to_bool(ssl_verify) if isinstance(ssl_verify, str) else ssl_verify
)
if ssl_verify_bool is not None:
ssl_verify = ssl_verify_bool
@ -89,14 +93,9 @@ def get_ssl_configuration(ssl_verify: Optional[VerifyTypes] = None) -> Union[boo
cafile = certifi.where()
if ssl_verify is not False:
custom_ssl_context = ssl.create_default_context(
cafile=cafile
)
custom_ssl_context = ssl.create_default_context(cafile=cafile)
# If security level is set, apply it to the SSL context
if (
ssl_security_level
and isinstance(ssl_security_level, str)
):
if ssl_security_level and isinstance(ssl_security_level, str):
# Create a custom SSL context with reduced security level
custom_ssl_context.set_ciphers(ssl_security_level)
@ -260,6 +259,7 @@ class AsyncHTTPHandler:
files: Optional[RequestFiles] = None,
content: Any = None,
):
start_time = time.time()
try:
if timeout is None:
@ -586,7 +586,7 @@ class AsyncHTTPHandler:
) -> Dict[str, Any]:
"""
Helper method to get SSL connector initialization arguments for aiohttp TCPConnector.
SSL Configuration Priority:
1. If ssl_context is provided -> use the custom SSL context
2. If ssl_verify is False -> disable SSL verification (ssl=False)
@ -597,14 +597,14 @@ class AsyncHTTPHandler:
connector_kwargs: Dict[str, Any] = {
"local_addr": ("0.0.0.0", 0) if litellm.force_ipv4 else None,
}
if ssl_context is not None:
# Priority 1: Use the provided custom SSL context
connector_kwargs["ssl"] = ssl_context
elif ssl_verify is False:
# Priority 2: Explicitly disable SSL verification
connector_kwargs["verify_ssl"] = False
return connector_kwargs
@staticmethod

58
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -111,6 +111,7 @@ class BaseLLMHTTPHandler:
response: Optional[httpx.Response] = None
for i in range(max(max_retry_on_unprocessable_entity_error, 1)):
try:
response = await async_httpx_client.post(
url=api_base,
headers=headers,
@ -2712,7 +2713,8 @@ class BaseLLMHTTPHandler:
headers = image_generation_provider_config.validate_environment(
api_key=litellm_params.get("api_key", None),
headers=image_generation_optional_request_params.get("extra_headers", {}) or {},
headers=image_generation_optional_request_params.get("extra_headers", {})
or {},
model=model,
messages=[],
optional_params=image_generation_optional_request_params,
@ -2763,15 +2765,17 @@ class BaseLLMHTTPHandler:
provider_config=image_generation_provider_config,
)
model_response: ImageResponse = image_generation_provider_config.transform_image_generation_response(
model=model,
raw_response=response,
model_response=litellm.ImageResponse(),
logging_obj=logging_obj,
request_data=data,
optional_params=image_generation_optional_request_params,
litellm_params=dict(litellm_params),
encoding=None,
model_response: ImageResponse = (
image_generation_provider_config.transform_image_generation_response(
model=model,
raw_response=response,
model_response=litellm.ImageResponse(),
logging_obj=logging_obj,
request_data=data,
optional_params=image_generation_optional_request_params,
litellm_params=dict(litellm_params),
encoding=None,
)
)
return model_response
@ -2804,10 +2808,10 @@ class BaseLLMHTTPHandler:
else:
async_httpx_client = client
headers = image_generation_provider_config.validate_environment(
api_key=litellm_params.get("api_key", None),
headers=image_generation_optional_request_params.get("extra_headers", {}) or {},
headers=image_generation_optional_request_params.get("extra_headers", {})
or {},
model=model,
messages=[],
optional_params=image_generation_optional_request_params,
@ -2858,17 +2862,19 @@ class BaseLLMHTTPHandler:
provider_config=image_generation_provider_config,
)
model_response: ImageResponse = image_generation_provider_config.transform_image_generation_response(
model=model,
raw_response=response,
model_response=litellm.ImageResponse(),
logging_obj=logging_obj,
request_data=data,
optional_params=image_generation_optional_request_params,
litellm_params=dict(litellm_params),
encoding=None,
model_response: ImageResponse = (
image_generation_provider_config.transform_image_generation_response(
model=model,
raw_response=response,
model_response=litellm.ImageResponse(),
logging_obj=logging_obj,
request_data=data,
optional_params=image_generation_optional_request_params,
litellm_params=dict(litellm_params),
encoding=None,
)
)
return model_response
###### VECTOR STORE HANDLER ######
@ -2936,7 +2942,9 @@ class BaseLLMHTTPHandler:
},
)
request_data = json.dumps(request_body) if signed_json_body is None else signed_json_body
request_data = (
json.dumps(request_body) if signed_json_body is None else signed_json_body
)
try:
response = await async_httpx_client.post(
@ -3035,7 +3043,9 @@ class BaseLLMHTTPHandler:
},
)
request_data = json.dumps(request_body) if signed_json_body is None else signed_json_body
request_data = (
json.dumps(request_body) if signed_json_body is None else signed_json_body
)
try:
response = sync_httpx_client.post(

31
litellm/llms/datarobot/chat/transformation.py
View file

@ -6,8 +6,11 @@ Calls done in OpenAI/openai.py as DataRobot is openai-compatible.
from typing import Optional, Tuple
from litellm.secret_managers.main import get_secret_str
from urllib.parse import urlparse, urlunparse
from ...openai_like.chat.transformation import OpenAILikeChatConfig
LLMGW_PATH = "/genai/llmgw/chat/completions"
class DataRobotConfig(OpenAILikeChatConfig):
@staticmethod
@ -32,22 +35,28 @@ class DataRobotConfig(OpenAILikeChatConfig):
if api_base is None:
api_base = "https://app.datarobot.com"
# If the api_base is a deployment URL, we do not append the chat completions path
if "api/v2/deployments" not in api_base:
# If the api_base is not a deployment URL, we need to append the chat completions path
if "api/v2/genai/llmgw/chat/completions" not in api_base:
api_base += "/api/v2/genai/llmgw/chat/completions"
parsed = urlparse(api_base)
path = parsed.path
if not path or path == "/": # Add full path to LLMGW
path += f"/api/v2/{LLMGW_PATH}"
elif "api/v2/deployments" in path: # Dedicated deployment, leave it
pass
elif (
"api/v2" in path and LLMGW_PATH not in path
): # Standard ENDPOINT path, add LLMGW
path += LLMGW_PATH
# Ensure the url ends with a trailing slash
if not api_base.endswith("/"):
api_base += "/"
if not path.endswith("/"):
path += "/"
path = path.replace("//", "/")
updated_parsed = parsed._replace(path=path)
return api_base # type: ignore
return urlunparse(updated_parsed)
def _get_openai_compatible_provider_info(
self,
api_base: Optional[str],
api_key: Optional[str]
self, api_base: Optional[str], api_key: Optional[str]
) -> Tuple[Optional[str], Optional[str]]:
"""Attempts to ensure that the API base and key are set, preferring user-provided values,
before falling back to secret manager values (``DATAROBOT_ENDPOINT`` and ``DATAROBOT_API_TOKEN``

239
litellm/llms/deepinfra/rerank/transformation.py Normal file
View file

@ -0,0 +1,239 @@
"""
Translate between Cohere's `/rerank` format and Deepinfra's `/rerank` format.
"""
import uuid
from typing import Any, Dict, List, Optional, Union
import httpx
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.rerank.transformation import (
BaseLLMException,
BaseRerankConfig,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.rerank import (
OptionalRerankParams,
RerankBilledUnits,
RerankResponse,
RerankResponseMeta,
RerankResponseResult,
RerankTokens,
)
class DeepinfraRerankConfig(BaseRerankConfig):
"""
Deepinfra Rerank - Follows the same Spec as Cohere Rerank
"""
def get_complete_url(self, api_base: Optional[str], model: str) -> str:
"""
Constructs the complete DeepInfra inference endpoint URL for rerank.
Args:
api_base (Optional[str]): The base URL for the DeepInfra API.
model (str): The model identifier.
Returns:
str: The complete URL for the DeepInfra rerank inference endpoint.
Raises:
ValueError: If api_base is None.
"""
if not api_base:
raise ValueError(
"Deepinfra API Base is required. api_base=None. Set in call or via `DEEPINFRA_API_BASE` env var."
)
# Remove 'openai' from the base if present
api_base_clean = (
api_base.replace("openai", "") if "openai" in api_base else api_base
)
# Remove any trailing slashes for consistency, then add one
api_base_clean = api_base_clean.rstrip("/") + "/"
# Compose the full endpoint
return f"{api_base_clean}inference/{model}"
def validate_environment(
self,
headers: dict,
model: str,
api_key: Optional[str] = None,
) -> dict:
if api_key is None:
api_key = get_secret_str("DEEPINFRA_API_KEY")
if api_key is None:
raise ValueError(
"Deepinfra API key is required. Please set 'DEEPINFRA_API_KEY' environment variable"
)
default_headers = {
"Authorization": f"Bearer {api_key}",
"accept": "application/json",
"content-type": "application/json",
}
# If 'Authorization' is provided in headers, it overrides the default.
if "Authorization" in headers:
default_headers["Authorization"] = headers["Authorization"]
# Merge other headers, overriding any default ones except Authorization
return {**default_headers, **headers}
def map_cohere_rerank_params(
self,
non_default_params: dict,
model: str,
drop_params: bool,
query: str,
documents: List[Union[str, Dict[str, Any]]],
custom_llm_provider: Optional[str] = None,
top_n: Optional[int] = None,
rank_fields: Optional[List[str]] = None,
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
) -> OptionalRerankParams:
# Start with the basic parameters
optional_rerank_params = {}
if query:
optional_rerank_params["queries"] = [query] * len(
documents
) # Deepinfra rerank requires queries to be of same length as documents
if non_default_params is not None:
for k, v in non_default_params.items():
if k == "queries" and v is not None:
# This should override the query parameter if it is provided
optional_rerank_params["queries"] = v
elif k == "documents" and v is not None:
optional_rerank_params["documents"] = v
elif k == "service_tier" and v is not None:
optional_rerank_params["service_tier"] = v
elif k == "instruction" and v is not None:
optional_rerank_params["instruction"] = v
elif k == "webhook" and v is not None:
optional_rerank_params["webhook"] = v
return OptionalRerankParams(**optional_rerank_params) # type: ignore
def transform_rerank_request(
self,
model: str,
optional_rerank_params: OptionalRerankParams,
headers: dict,
) -> dict:
# Convert OptionalRerankParams to dict as expected by parent class
if optional_rerank_params is None:
return {}
return dict(optional_rerank_params)
def transform_rerank_response(
self,
model: str,
raw_response: httpx.Response,
model_response: RerankResponse,
logging_obj: LiteLLMLoggingObj,
api_key: Optional[str] = None,
request_data: dict = {},
optional_params: dict = {},
litellm_params: dict = {},
) -> RerankResponse:
try:
response_json = raw_response.json()
logging_obj.post_call(original_response=raw_response.text)
# Extract the scores from the response
scores = response_json.get("scores", [])
input_tokens = response_json.get("input_tokens", 0)
request_id = response_json.get("request_id")
# Create inference status information
inference_status = response_json.get("inference_status", {})
status = inference_status.get("status", "unknown")
runtime_ms = inference_status.get("runtime_ms", 0)
cost = inference_status.get("cost", 0.0)
tokens_generated = inference_status.get("tokens_generated", 0)
tokens_input = inference_status.get("tokens_input", 0)
# Create RerankResponse
results = []
for i, score in enumerate(scores):
results.append(
RerankResponseResult(index=i, relevance_score=float(score))
)
# Create metadata for the response
tokens = RerankTokens(
input_tokens=input_tokens,
output_tokens=0, # DeepInfra doesn't provide output tokens for rerank
)
billed_units = RerankBilledUnits(total_tokens=input_tokens)
meta = RerankResponseMeta(tokens=tokens, billed_units=billed_units)
rerank_response = RerankResponse(
id=request_id or str(uuid.uuid4()), results=results, meta=meta
)
# Store additional information in hidden params
rerank_response._hidden_params = {
"status": status,
"runtime_ms": runtime_ms,
"cost": cost,
"tokens_generated": tokens_generated,
"tokens_input": tokens_input,
"model": model,
}
return rerank_response
except Exception:
# If there's an error parsing the response, fall back to the parent implementation
rerank_response = super().transform_rerank_response(
model=model,
raw_response=raw_response,
model_response=model_response,
logging_obj=logging_obj,
api_key=api_key,
request_data=request_data,
optional_params=optional_params,
litellm_params=litellm_params,
)
rerank_response._hidden_params["model"] = model
return rerank_response
def get_supported_cohere_rerank_params(self, model: str) -> list:
return ["query", "documents"]
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
# Deepinfra errors may come as JSON: {"detail": {"error": "..."}}
import json
# Try to extract a more specific error message if possible
try:
error_data = error_message
if isinstance(error_message, str):
error_data = json.loads(error_message)
if isinstance(error_data, dict):
# Check for {"detail": {"error": "..."}}
detail = error_data.get("detail")
if isinstance(detail, dict) and "error" in detail:
error_message = detail["error"]
elif isinstance(detail, str):
error_message = detail
except Exception:
# If parsing fails, just use the original error_message
pass
raise BaseLLMException(
status_code=status_code,
message=error_message,
headers=headers,
)

218
litellm/llms/mistral/chat/transformation.py
View file

@ -6,7 +6,18 @@ Why separate file? Make it easy to see how transformation works
Docs - https://docs.mistral.ai/api/
"""
from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, cast, overload
from typing import (
Any,
Coroutine,
List,
Literal,
Optional,
Tuple,
Union,
cast,
get_type_hints,
overload,
)
import httpx
@ -17,7 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
)
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.mistral import MistralToolCallMessage
from litellm.types.llms.mistral import MistralThinkingBlock, MistralToolCallMessage
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponse
from litellm.utils import convert_to_model_response_object
@ -145,7 +156,9 @@ class MistralConfig(OpenAIGPTConfig):
for param, value in non_default_params.items():
if param == "max_tokens":
optional_params["max_tokens"] = value
if param == "max_completion_tokens": # max_completion_tokens should take priority
if (
param == "max_completion_tokens"
): # max_completion_tokens should take priority
optional_params["max_tokens"] = value
if param == "tools":
# Clean tools to remove problematic schema fields for Mistral API
@ -159,7 +172,9 @@ class MistralConfig(OpenAIGPTConfig):
if param == "stop":
optional_params["stop"] = value
if param == "tool_choice" and isinstance(value, str):
optional_params["tool_choice"] = self._map_tool_choice(tool_choice=value)
optional_params["tool_choice"] = self._map_tool_choice(
tool_choice=value
)
if param == "seed":
optional_params["extra_body"] = {"random_seed": value}
if param == "response_format":
@ -185,7 +200,9 @@ class MistralConfig(OpenAIGPTConfig):
) # type: ignore
# if api_base does not end with /v1 we add it
if api_base is not None and not api_base.endswith("/v1"): # Mistral always needs a /v1 at the end
if api_base is not None and not api_base.endswith(
"/v1"
): # Mistral always needs a /v1 at the end
api_base = api_base + "/v1"
dynamic_api_key = (
api_key
@ -194,10 +211,12 @@ class MistralConfig(OpenAIGPTConfig):
)
return api_base, dynamic_api_key
# fmt: off
@overload
def _transform_messages(
self, messages: List[AllMessageValues], model: str, is_async: Literal[True]
) -> Coroutine[Any, Any, List[AllMessageValues]]:
) -> Coroutine[Any, Any, List[AllMessageValues]]:
...
@overload
@ -206,8 +225,9 @@ class MistralConfig(OpenAIGPTConfig):
messages: List[AllMessageValues],
model: str,
is_async: Literal[False] = False,
) -> List[AllMessageValues]:
) -> List[AllMessageValues]:
...
# fmt: on
def _transform_messages(
self, messages: List[AllMessageValues], model: str, is_async: bool = False
@ -218,18 +238,20 @@ class MistralConfig(OpenAIGPTConfig):
- if image passed in, then just return as is (user-intended)
- if `name` is passed, then drop it for mistral API: https://github.com/BerriAI/litellm/issues/6696
Motivation: mistral api doesn't support content as a list
Motivation: mistral api doesn't support content as a list.
The above statement is not valid now. Need to plan to remove all the #1,2,3
Mistral API supports content as a list.
"""
## 1. If 'image_url' in content, then return as is
## 1. If 'image_url' or 'file' in content, then transform with base class and mistral-specific handling
for m in messages:
_content_block = m.get("content")
if _content_block and isinstance(_content_block, list):
for c in _content_block:
if c.get("type") == "image_url":
if is_async:
return super()._transform_messages(messages, model, True)
else:
return super()._transform_messages(messages, model, False)
if any(c.get("type") in ["image_url", "file"] for c in _content_block):
if is_async:
return self._transform_messages_async(messages, model)
else:
messages = self._transform_messages_sync(messages, model)
return messages
## 2. If content is list, then convert to string
messages = handle_messages_with_content_list_to_str_conversion(messages)
@ -239,6 +261,8 @@ class MistralConfig(OpenAIGPTConfig):
for m in messages:
m = MistralConfig._handle_name_in_message(m)
m = MistralConfig._handle_tool_call_message(m)
if MistralConfig._is_empty_assistant_message(m):
continue
m = strip_none_values_from_message(m) # prevents 'extra_forbidden' error
new_messages.append(m)
@ -247,6 +271,51 @@ class MistralConfig(OpenAIGPTConfig):
else:
return super()._transform_messages(new_messages, model, False)
async def _transform_messages_async(self,
messages: List[AllMessageValues], model: str
) -> List[AllMessageValues]:
"""
Handle modification of messages for Mistral API in an async context.
"""
# Call parent async method to handle basic transformations
# and then apply Mistral-specific handling for files
messages = await super()._transform_messages(messages, model, True)
messages = self._handle_message_with_file(messages)
return messages
def _transform_messages_sync(self,
messages: List[AllMessageValues], model: str
) -> List[AllMessageValues]:
""" Handle modification of messages for Mistral API in a sync context.
"""
# Call parent sync method to handle basic transformations
# and then apply Mistral-specific handling for files
# This is the sync version of the async method above
messages = super()._transform_messages(messages, model, False)
messages = self._handle_message_with_file(messages)
return messages
def _handle_message_with_file(
self,
messages: List[AllMessageValues]) -> List[AllMessageValues]:
"""
Mistral API supports only 'file_id' in message content with type 'file'.
"""
for m in messages:
_content_block = m.get("content")
if _content_block and isinstance(_content_block, list):
if any(c.get("type") == "file" for c in _content_block):
# If file content is present, we get file_id from 'file' attribute of content block
# then replace 'file' with 'file_id' and assign the value of 'file_id' attribute to it.
file_contents = [c for c in _content_block if c.get("type") == "file"]
for file_content in file_contents:
file_id = file_content.get("file", {}).get("file_id")
if file_id:
# Replace 'file' with 'file_id'
file_content["file_id"] = file_id # type: ignore
file_content.pop("file", None)
return messages
def _add_reasoning_system_prompt_if_needed(
self, messages: List[AllMessageValues], optional_params: dict
) -> List[AllMessageValues]:
@ -269,20 +338,30 @@ class MistralConfig(OpenAIGPTConfig):
# Handle both string and list content, preserving original format
if isinstance(existing_content, str):
# String content - prepend reasoning prompt
new_content: Union[str, list] = f"{reasoning_prompt}\n\n{existing_content}"
new_content: Union[str, list] = (
f"{reasoning_prompt}\n\n{existing_content}"
)
elif isinstance(existing_content, list):
# List content - prepend reasoning prompt as text block
new_content = [{"type": "text", "text": reasoning_prompt + "\n\n"}] + existing_content
new_content = [
{"type": "text", "text": reasoning_prompt + "\n\n"}
] + existing_content
else:
# Fallback for any other type - convert to string
new_content = f"{reasoning_prompt}\n\n{str(existing_content)}"
messages[i] = cast(AllMessageValues, {**msg, "content": new_content})
messages[i] = cast(
AllMessageValues, {**msg, "content": new_content}
)
break
else:
# Add new system message with reasoning instructions
reasoning_message: AllMessageValues = cast(
AllMessageValues, {"role": "system", "content": self._get_mistral_reasoning_system_prompt()}
AllMessageValues,
{
"role": "system",
"content": self._get_mistral_reasoning_system_prompt(),
},
)
messages = [reasoning_message] + messages
@ -294,32 +373,34 @@ class MistralConfig(OpenAIGPTConfig):
def _clean_tool_schema_for_mistral(cls, tools: list) -> list:
"""
Clean tool schemas to remove fields that cause issues with Mistral API.
Removes:
- $id and $schema fields (cause grammar validation errors)
- additionalProperties=False (causes OpenAI API schema errors)
- strict field (not supported by Mistral)
Args:
tools: List of tool definitions
max_depth: Maximum recursion depth for schema cleaning (default: 10)
Returns:
Cleaned tools list
"""
if not tools:
return tools
import copy
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.utils import _remove_json_schema_refs
cleaned_tools = copy.deepcopy(tools)
# Apply all cleaning functions with max_depth protection
cleaned_tools = _remove_json_schema_refs(cleaned_tools, max_depth=DEFAULT_MAX_RECURSE_DEPTH)
cleaned_tools = _remove_json_schema_refs(
cleaned_tools, max_depth=DEFAULT_MAX_RECURSE_DEPTH
)
return cleaned_tools
@classmethod
@ -360,6 +441,25 @@ class MistralConfig(OpenAIGPTConfig):
message["tool_calls"] = mistral_tool_calls # type: ignore
return message
@classmethod
def _is_empty_assistant_message(cls, message: AllMessageValues) -> bool:
"""
Mistral API does not support empty string in assistant content.
"""
from litellm.types.llms.openai import ChatCompletionAssistantMessage
set_keys = get_type_hints(ChatCompletionAssistantMessage).keys()
all_expected_values_are_empty = True
for key in set_keys:
if key != "role" and message.get(key) is not None:
if key == "content" and message.get(key) == "":
continue
else:
all_expected_values_are_empty = False
break
return all_expected_values_are_empty
@staticmethod
def _handle_empty_content_response(response_data: dict) -> dict:
"""
@ -380,6 +480,58 @@ class MistralConfig(OpenAIGPTConfig):
choice["message"]["content"] = None
return response_data
@staticmethod
def _convert_thinking_block_to_reasoning_content(
thinking_blocks: MistralThinkingBlock,
) -> str:
"""
Convert Mistral thinking blocks to reasoning content.
"""
return "\n".join(
[block.get("text", "") for block in thinking_blocks["thinking"]]
)
@staticmethod
def _handle_content_list_to_str_conversion(response_data: dict) -> dict:
"""
Handle Mistral's content list format and extract thinking content.
Map mistral's content list to string and extract thinking blocks:
- Thinking block -> reasoning_content field
- Text block -> content field
"""
if response_data.get("choices") and len(response_data["choices"]) > 0:
for choice in response_data["choices"]:
if choice.get("message") and choice["message"].get("content"):
content = choice["message"]["content"]
# Only process if content is a list
if isinstance(content, list):
thinking_content = ""
text_content = ""
# Process each content block
for block in content:
if block.get("type") == "thinking":
thinking_blocks = block.get("thinking", [])
thinking_texts = []
for thinking_block in thinking_blocks:
if thinking_block.get("type") == "text":
thinking_texts.append(
thinking_block.get("text", "")
)
thinking_content = "\n".join(thinking_texts)
elif block.get("type") == "text":
text_content = block.get("text", "")
# Set the extracted content
choice["message"]["content"] = text_content
if thinking_content:
choice["message"]["reasoning_content"] = thinking_content
return response_data
def transform_request(
self,
model: str,
@ -396,8 +548,12 @@ class MistralConfig(OpenAIGPTConfig):
dict: The transformed request. Sent as the body of the API call.
"""
# Add reasoning system prompt if needed (for magistral models)
if "magistral" in model.lower() and optional_params.get("_add_reasoning_prompt", False):
messages = self._add_reasoning_system_prompt_if_needed(messages, optional_params)
if "magistral" in model.lower() and optional_params.get(
"_add_reasoning_prompt", False
):
messages = self._add_reasoning_system_prompt_if_needed(
messages, optional_params
)
# Call parent transform_request which handles _transform_messages
return super().transform_request(
@ -424,14 +580,16 @@ class MistralConfig(OpenAIGPTConfig):
) -> ModelResponse:
"""
Transform the raw response from Mistral API.
Handles Mistral-specific behavior like converting empty string content to None.
Handles Mistral-specific behavior like converting empty string content to None
and extracting thinking content from content lists.
"""
logging_obj.post_call(original_response=raw_response.text)
logging_obj.model_call_details["response_headers"] = raw_response.headers
# Handle Mistral-specific empty string content conversion to None
# Handle Mistral-specific response transformations
response_data = raw_response.json()
response_data = self._handle_empty_content_response(response_data)
response_data = self._handle_content_list_to_str_conversion(response_data)
final_response_obj = cast(
ModelResponse,

74
litellm/llms/ollama/completion/transformation.py
View file

@ -262,38 +262,52 @@ class OllamaConfig(BaseConfig):
## RESPONSE OBJECT
model_response.choices[0].finish_reason = "stop"
if request_data.get("format", "") == "json":
response_content = json.loads(response_json["response"])
# Check if this is a function call format with name/arguments structure
if (
isinstance(response_content, dict)
and "name" in response_content
and "arguments" in response_content
):
# Handle as function call (original behavior)
function_call = response_content
message = litellm.Message(
content=None,
tool_calls=[
{
"id": f"call_{str(uuid.uuid4())}",
"function": {
"name": function_call["name"],
"arguments": json.dumps(function_call["arguments"]),
},
"type": "function",
}
],
)
model_response.choices[0].message = message # type: ignore
model_response.choices[0].finish_reason = "tool_calls"
else:
# Handle as regular JSON (new behavior)
message = litellm.Message(
content=json.dumps(response_content),
)
# Check if response field exists and is not empty before parsing JSON
response_text = response_json.get("response", "")
if not response_text or not response_text.strip():
# Handle empty response gracefully - set empty content
message = litellm.Message(content="")
model_response.choices[0].message = message # type: ignore
model_response.choices[0].finish_reason = "stop"
else:
try:
response_content = json.loads(response_text)
# Check if this is a function call format with name/arguments structure
if (
isinstance(response_content, dict)
and "name" in response_content
and "arguments" in response_content
):
# Handle as function call (original behavior)
function_call = response_content
message = litellm.Message(
content=None,
tool_calls=[
{
"id": f"call_{str(uuid.uuid4())}",
"function": {
"name": function_call["name"],
"arguments": json.dumps(function_call["arguments"]),
},
"type": "function",
}
],
)
model_response.choices[0].message = message # type: ignore
model_response.choices[0].finish_reason = "tool_calls"
else:
# Handle as regular JSON (new behavior)
message = litellm.Message(
content=json.dumps(response_content),
)
model_response.choices[0].message = message # type: ignore
model_response.choices[0].finish_reason = "stop"
except json.JSONDecodeError:
# If JSON parsing fails, treat as regular text response
message = litellm.Message(content=response_text)
model_response.choices[0].message = message # type: ignore
model_response.choices[0].finish_reason = "stop"
else:
model_response.choices[0].message.content = response_json["response"] # type: ignore
model_response.created = int(time.time())

3
litellm/llms/openai/chat/gpt_transformation.py
View file

@ -348,6 +348,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
for message in messages:
message_content = message.get("content")
message_role = message.get("role")
if (
message_role == "user"
and message_content
@ -428,6 +429,8 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
if tools is not None and len(tools) > 0:
optional_params["tools"] = tools
optional_params.pop("max_retries", None)
return {
"model": model,
"messages": messages,

31
litellm/llms/openai/image_edit/transformation.py
View file

@ -80,24 +80,33 @@ class OpenAIImageEditConfig(BaseImageEditConfig):
request_dict = cast(Dict, request)
#########################################################
# Separate images as `files` and send other parameters as `data`
# Separate images and masks as `files` and send other parameters as `data`
#########################################################
_images = request_dict.get("image") or []
data_without_images = {k: v for k, v in request_dict.items() if k != "image"}
_image = request_dict.get("image")
_mask = request_dict.get("mask")
data_without_files = {
k: v for k, v in request_dict.items() if k not in ["image", "mask"]
}
files_list: List[Tuple[str, Any]] = []
for _image in _images:
# Handle image parameter
if _image is not None:
image_content_type: str = ImageEditRequestUtils.get_image_content_type(
_image
)
if isinstance(_image, BufferedReader):
files_list.append(
("image[]", (_image.name, _image, image_content_type))
)
files_list.append(("image", (_image.name, _image, image_content_type)))
else:
files_list.append(
("image[]", ("image.png", _image, image_content_type))
)
return data_without_images, files_list
files_list.append(("image", ("image.png", _image, image_content_type)))
# Handle mask parameter if provided
if _mask is not None:
mask_content_type: str = ImageEditRequestUtils.get_image_content_type(_mask)
if isinstance(_mask, BufferedReader):
files_list.append(("mask", (_mask.name, _mask, mask_content_type)))
else:
files_list.append(("mask", ("mask.png", _mask, mask_content_type)))
return data_without_files, files_list
def transform_image_edit_response(
self,

59
litellm/llms/openai/responses/transformation.py
View file

@ -1,4 +1,4 @@
from typing import TYPE_CHECKING, Any, Dict, Optional, Union, cast
from typing import TYPE_CHECKING, Any, Dict, Optional, Union, cast, get_type_hints
import httpx
from pydantic import BaseModel
@ -13,6 +13,7 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import *
from litellm.types.responses.main import *
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
from ..common_utils import OpenAIError
@ -25,38 +26,28 @@ else:
class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.OPENAI
def get_supported_openai_params(self, model: str) -> list:
"""
All OpenAI Responses API params are supported
"""
return [
"input",
"model",
"include",
"instructions",
"max_output_tokens",
"metadata",
"parallel_tool_calls",
"previous_response_id",
"reasoning",
"store",
"background",
"stream",
"prompt",
"temperature",
"text",
"tool_choice",
"tools",
"top_p",
"truncation",
"user",
"service_tier",
"safety_identifier",
"extra_headers",
"extra_query",
"extra_body",
"timeout",
]
supported_params = get_type_hints(ResponsesAPIRequestParams).keys()
return list(
set(
[
"input",
"model",
"extra_headers",
"extra_query",
"extra_body",
"timeout",
]
+ list(supported_params)
)
)
def map_openai_params(
self,
@ -85,8 +76,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
)
return final_request_params
def _validate_input_param(self, input: Union[str, ResponseInputParam]) -> Union[str, ResponseInputParam]:
def _validate_input_param(
self, input: Union[str, ResponseInputParam]
) -> Union[str, ResponseInputParam]:
"""
Ensure all input fields if pydantic are converted to dict
@ -114,7 +107,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
"""No transform applied since outputs are in OpenAI spec already"""
try:
raw_response_json = raw_response.json()
raw_response_json["created_at"] = _safe_convert_created_field(raw_response_json["created_at"])
raw_response_json["created_at"] = _safe_convert_created_field(
raw_response_json["created_at"]
)
except Exception:
raise OpenAIError(
message=raw_response.text, status_code=raw_response.status_code

8
litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py
View file

@ -113,10 +113,10 @@ class VertexAILlama3Config(OpenAIGPTConfig):
status_code=raw_response.status_code,
headers=response_headers,
)
model_response.model = completion_response["model"]
model_response.id = completion_response["id"]
model_response.created = completion_response["created"]
setattr(model_response, "usage", Usage(**completion_response["usage"]))
model_response.model = completion_response.get("model", model)
model_response.id = completion_response.get("id", "")
model_response.created = completion_response.get("created", 0)
setattr(model_response, "usage", Usage(**completion_response.get("usage", {})))
model_response.choices = self._transform_choices( # type: ignore
choices=completion_response["choices"],

16
litellm/llms/vertex_ai/vertex_ai_partner_models/main.py
View file

@ -48,9 +48,21 @@ class VertexAIPartnerModels(VertexBase):
or model.startswith("codestral")
or model.startswith("jamba")
or model.startswith("claude")
or model.startswith("qwen")
):
return True
return False
@staticmethod
def should_use_openai_handler(model: str):
OPENAI_LIKE_VERTEX_PROVIDERS = [
"llama",
"deepseek-ai",
"qwen",
]
if any(provider in model for provider in OPENAI_LIKE_VERTEX_PROVIDERS):
return True
return False
def completion(
self,
@ -115,7 +127,7 @@ class VertexAIPartnerModels(VertexBase):
optional_params["stream"] = stream
if "llama" in model or "deepseek-ai" in model:
if self.should_use_openai_handler(model):
partner = VertexPartnerProvider.llama
elif "mistral" in model or "codestral" in model:
partner = VertexPartnerProvider.mistralai
@ -191,7 +203,7 @@ class VertexAIPartnerModels(VertexBase):
client=client,
custom_llm_provider=LlmProviders.VERTEX_AI.value,
)
elif "llama" in model:
elif self.should_use_openai_handler(model):
return base_llm_http_handler.completion(
model=model,
stream=stream,

153
litellm/llms/voyage/embedding/transformation_contextual.py Normal file
View file

@ -0,0 +1,153 @@
"""
This module is used to transform the request and response for the Voyage contextualized embeddings API.
This would be used for all the contextualized embeddings models in Voyage.
"""
from typing import List, Optional, Union
import httpx
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
from litellm.types.utils import EmbeddingResponse, Usage
class VoyageError(BaseLLMException):
def __init__(
self,
status_code: int,
message: str,
headers: Union[dict, httpx.Headers] = {},
):
self.status_code = status_code
self.message = message
self.request = httpx.Request(
method="POST", url="https://api.voyageai.com/v1/contextualizedembeddings"
)
self.response = httpx.Response(status_code=status_code, request=self.request)
super().__init__(
status_code=status_code,
message=message,
headers=headers,
)
class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
"""
Reference: https://docs.voyageai.com/reference/embeddings-api
"""
def __init__(self) -> None:
pass
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
if api_base:
if not api_base.endswith("/contextualizedembeddings"):
api_base = f"{api_base}/contextualizedembeddings"
return api_base
return "https://api.voyageai.com/v1/contextualizedembeddings"
def get_supported_openai_params(self, model: str) -> list:
return ["encoding_format", "dimensions"]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
"""
Map OpenAI params to Voyage params
Reference: https://docs.voyageai.com/reference/contextualized-embeddings-api
"""
if "encoding_format" in non_default_params:
optional_params["encoding_format"] = non_default_params["encoding_format"]
if "dimensions" in non_default_params:
optional_params["output_dimension"] = non_default_params["dimensions"]
return optional_params
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
if api_key is None:
api_key = (
get_secret_str("VOYAGE_API_KEY")
or get_secret_str("VOYAGE_AI_API_KEY")
or get_secret_str("VOYAGE_AI_TOKEN")
)
return {
"Authorization": f"Bearer {api_key}",
}
def transform_embedding_request(
self,
model: str,
input: Union[AllEmbeddingInputValues, List[List[str]]],
optional_params: dict,
headers: dict,
) -> dict:
return {
"inputs": input,
"model": model,
**optional_params,
}
def transform_embedding_response(
self,
model: str,
raw_response: httpx.Response,
model_response: EmbeddingResponse,
logging_obj: LiteLLMLoggingObj,
api_key: Optional[str] = None,
request_data: dict = {},
optional_params: dict = {},
litellm_params: dict = {},
) -> EmbeddingResponse:
try:
raw_response_json = raw_response.json()
except Exception:
raise VoyageError(
message=raw_response.text, status_code=raw_response.status_code
)
# model_response.usage
model_response.model = raw_response_json.get("model")
model_response.data = raw_response_json.get("data")
model_response.object = raw_response_json.get("object")
usage = Usage(
prompt_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
)
model_response.usage = usage
return model_response
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
return VoyageError(
message=error_message, status_code=status_code, headers=headers
)
@staticmethod
def is_contextualized_embeddings(model: str) -> bool:
return "context" in model.lower()

41
litellm/main.py
View file

@ -130,7 +130,6 @@ from .litellm_core_utils.prompt_templates.factory import (
stringify_json_tool_call_content,
)
from .litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
from .llms import baseten
from .llms.anthropic.chat import AnthropicChatCompletion
from .llms.azure.audio_transcriptions import AzureAudioTranscription
from .llms.azure.azure import AzureChatCompletion, _check_dynamic_azure_params
@ -1562,6 +1561,7 @@ def completion( # type: ignore # noqa: PLR0915
)
elif custom_llm_provider == "deepseek":
## COMPLETION CALL
try:
response = base_llm_http_handler.completion(
model=model,
@ -1593,6 +1593,7 @@ def completion( # type: ignore # noqa: PLR0915
elif custom_llm_provider == "azure_ai":
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
api_base = AzureFoundryModelInfo.get_api_base(api_base)
# set API KEY
api_key = AzureFoundryModelInfo.get_api_key(api_key)
@ -1921,6 +1922,7 @@ def completion( # type: ignore # noqa: PLR0915
or custom_llm_provider == "perplexity"
or custom_llm_provider == "nvidia_nim"
or custom_llm_provider == "cerebras"
or custom_llm_provider == "baseten"
or custom_llm_provider == "sambanova"
or custom_llm_provider == "volcengine"
or custom_llm_provider == "anyscale"
@ -1976,8 +1978,10 @@ def completion( # type: ignore # noqa: PLR0915
use_base_llm_http_handler = get_secret_bool(
"EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER"
)
try:
if use_base_llm_http_handler:
response = base_llm_http_handler.completion(
model=model,
messages=messages,
@ -3260,42 +3264,7 @@ def completion( # type: ignore # noqa: PLR0915
api_key=api_key,
logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements
)
elif (
custom_llm_provider == "baseten"
or litellm.api_base == "https://app.baseten.co"
):
custom_llm_provider = "baseten"
baseten_key = (
api_key
or litellm.baseten_key
or os.environ.get("BASETEN_API_KEY")
or litellm.api_key
)
model_response = baseten.completion(
model=model,
messages=messages,
model_response=model_response,
print_verbose=print_verbose,
optional_params=optional_params,
litellm_params=litellm_params,
logger_fn=logger_fn,
encoding=encoding,
api_key=baseten_key,
logging_obj=logging,
)
if inspect.isgenerator(model_response) or (
"stream" in optional_params and optional_params["stream"] is True
):
# don't try to access stream object,
response = CustomStreamWrapper(
model_response,
model,
custom_llm_provider="baseten",
logging_obj=logging,
)
return response
response = model_response
elif custom_llm_provider == "petals" or model in litellm.petals_models:
api_base = api_base or litellm.api_base

1284
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

7
litellm/passthrough/main.py
View file

@ -24,6 +24,7 @@ import litellm
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.passthrough.utils import CommonUtils
from litellm.utils import client
base_llm_http_handler = BaseLLMHTTPHandler()
@ -241,6 +242,12 @@ def llm_passthrough_route(
request_query_params=request_query_params,
litellm_params=litellm_params_dict,
)
# need to encode the id of application-inference-profile for bedrock
if custom_llm_provider == "bedrock" and "application-inference-profile" in endpoint:
encoded_url_str = CommonUtils.encode_bedrock_runtime_modelid_arn(str(updated_url))
updated_url = httpx.URL(encoded_url_str)
# Add or update query parameters
provider_api_key = provider_config.get_api_key(api_key)

53
litellm/passthrough/utils.py
View file

@ -37,3 +37,56 @@ class BasePassthroughUtils:
# Combine request headers with custom headers
headers = {**request_headers, **headers}
return headers
class CommonUtils:
@staticmethod
def encode_bedrock_runtime_modelid_arn(endpoint: str) -> str:
"""
Encodes any "/" found in the modelId of an AWS Bedrock Runtime Endpoint when arns are passed in.
- modelID value can be an ARN which contains slashes that SHOULD NOT be treated as path separators.
e.g endpoint: /model/<modelId>/invoke
<modelId> containing arns with slashes need to be encoded from
arn:aws:bedrock:ap-southeast-1:123456789012:application-inference-profile/abdefg12334 =>
arn:aws:bedrock:ap-southeast-1:123456789012:application-inference-profile%2Fabdefg12334
so that it is treated as one part of the path.
Otherwise, the encoded endpoint will return 500 error when passed to Bedrock endpoint.
See the apis in https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Amazon_Bedrock_Runtime.html
for more details on the regex patterns of modelId which we use in the regex logic below.
Args:
endpoint (str): The original endpoint string which may contain ARNs that contain slashes.
Returns:
str: The endpoint with properly encoded ARN slashes
"""
import re
# Early exit: if no ARN detected, return unchanged
if 'arn:aws:' not in endpoint:
return endpoint
# Handle all patterns in one go - more efficient and cleaner
patterns = [
# Custom model with 2 slashes (order matters - do this first)
(r'(custom-model)/([a-z0-9.-]+)/([a-z0-9]+)', r'\1%2F\2%2F\3'),
# All other resource types with 1 slash
(r'(:application-inference-profile)/', r'\1%2F'),
(r'(:inference-profile)/', r'\1%2F'),
(r'(:foundation-model)/', r'\1%2F'),
(r'(:imported-model)/', r'\1%2F'),
(r'(:provisioned-model)/', r'\1%2F'),
(r'(:prompt)/', r'\1%2F'),
(r'(:endpoint)/', r'\1%2F'),
(r'(:prompt-router)/', r'\1%2F'),
(r'(:default-prompt-router)/', r'\1%2F'),
]
for pattern, replacement in patterns:
# Check if pattern exists before applying regex (early exit optimization)
if re.search(pattern, endpoint):
endpoint = re.sub(pattern, replacement, endpoint)
break # Exit after first match since each ARN has only one resource type
return endpoint

25
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -40,6 +40,7 @@ except ImportError as e:
# Global variables to track initialization
_SESSION_MANAGERS_INITIALIZED = False
_INITIALIZATION_LOCK = asyncio.Lock()
if MCP_AVAILABLE:
from mcp.server import Server
@ -113,21 +114,23 @@ if MCP_AVAILABLE:
"""Initialize the session managers. Can be called from main app lifespan."""
global _SESSION_MANAGERS_INITIALIZED, _session_manager_cm, _sse_session_manager_cm
if _SESSION_MANAGERS_INITIALIZED:
return
# Use async lock to prevent concurrent initialization
async with _INITIALIZATION_LOCK:
if _SESSION_MANAGERS_INITIALIZED:
return
verbose_logger.info("Initializing MCP session managers...")
verbose_logger.info("Initializing MCP session managers...")
# Start the session managers with context managers
_session_manager_cm = session_manager.run()
_sse_session_manager_cm = sse_session_manager.run()
# Start the session managers with context managers
_session_manager_cm = session_manager.run()
_sse_session_manager_cm = sse_session_manager.run()
# Enter the context managers
await _session_manager_cm.__aenter__()
await _sse_session_manager_cm.__aenter__()
# Enter the context managers
await _session_manager_cm.__aenter__()
await _sse_session_manager_cm.__aenter__()
_SESSION_MANAGERS_INITIALIZED = True
verbose_logger.info("MCP Server started with StreamableHTTP and SSE session managers!")
_SESSION_MANAGERS_INITIALIZED = True
verbose_logger.info("MCP Server started with StreamableHTTP and SSE session managers!")
async def shutdown_session_managers():
"""Shutdown the session managers."""

14
litellm/proxy/_new_secret_config.yaml
View file

@ -15,4 +15,16 @@ model_list:
mode: chat
router_settings:
model_group_alias: {"my-fake-gpt-4": "fake-openai-endpoint"}
model_group_alias: {"my-fake-gpt-4": "fake-openai-endpoint"}
litellm_settings:
callbacks: ["otel"]
cache: true
cache_params:
type: redis
ttl: 600
supported_call_types: ["acompletion", "completion"]
model_group_settings:
forward_client_headers_to_llm_api:
- fake-openai-endpoint

35
litellm/proxy/_types.py
View file

@ -530,7 +530,7 @@ class LiteLLMRoutes(enum.Enum):
# Routes accessible by Admin Viewer (read-only admin access)
admin_viewer_routes = [
"/user/list",
"/user/available_users",
"/user/available_users",
"/user/available_roles",
"/user/daily/activity",
"/team/daily/activity",
@ -540,7 +540,10 @@ class LiteLLMRoutes(enum.Enum):
# All routes accesible by an Org Admin
org_admin_allowed_routes = (
org_admin_only_routes + management_routes + self_managed_routes + admin_viewer_routes
org_admin_only_routes
+ management_routes
+ self_managed_routes
+ admin_viewer_routes
)
@ -585,13 +588,14 @@ class LiteLLMPromptInjectionParams(LiteLLMPydanticObjectBase):
######### Request Class Definition ######
class ProxyChatCompletionRequest(LiteLLMPydanticObjectBase):
"""
Pydantic model for chat completion requests that includes both OpenAI standard fields
Pydantic model for chat completion requests that includes both OpenAI standard fields
and LiteLLM-specific parameters. This replaces the previous TypedDict version.
"""
# Required fields (from ChatCompletionRequest)
model: str
messages: List[AllMessageValues]
# Standard OpenAI completion parameters (all optional)
frequency_penalty: Optional[float] = None
logit_bias: Optional[Dict[str, float]] = None
@ -614,10 +618,10 @@ class ProxyChatCompletionRequest(LiteLLMPydanticObjectBase):
functions: Optional[List[Dict[str, Any]]] = None
user: Optional[str] = None
stream: Optional[bool] = None
# LiteLLM-specific metadata param (from original ChatCompletionRequest)
metadata: Optional[Dict[str, Any]] = None
# Optional LiteLLM params
guardrails: Optional[List[str]] = None
caching: Optional[bool] = None
@ -1873,7 +1877,8 @@ class UserAPIKeyAuth(
key_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
team_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
)
class UserInfoResponse(LiteLLMPydanticObjectBase):
user_id: Optional[str]
user_info: Optional[Union[dict, BaseModel]]
@ -2120,7 +2125,6 @@ class TokenCountRequest(LiteLLMPydanticObjectBase):
Anthropic token counting endpoint uses /messages
"""
contents: Optional[List[dict]] = None
"""
Google /countTokens endpoint expects contents to be a list of dicts with the following structure:
@ -2265,7 +2269,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
braintrust: CallbackOnUI = CallbackOnUI(
litellm_callback_name="braintrust",
litellm_callback_params=["BRAINTRUST_API_KEY","BRAINTRUST_API_BASE"],
litellm_callback_params=["BRAINTRUST_API_KEY", "BRAINTRUST_API_BASE"],
ui_callback_name="Braintrust",
)
@ -2319,7 +2323,9 @@ class SpendLogsMetadata(TypedDict):
error_information: Optional[StandardLoggingPayloadErrorInformation]
usage_object: Optional[dict]
model_map_information: Optional[StandardLoggingModelInformation]
cold_storage_object_key: Optional[str] # S3/GCS object key for cold storage retrieval
cold_storage_object_key: Optional[
str
] # S3/GCS object key for cold storage retrieval
class SpendLogsPayload(TypedDict):
@ -2646,7 +2652,7 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase):
if self.litellm_budget_table is not None:
return self.litellm_budget_table.rpm_limit
return None
def safe_get_team_member_tpm_limit(self) -> Optional[int]:
if self.litellm_budget_table is not None:
return self.litellm_budget_table.tpm_limit
@ -2763,14 +2769,11 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest):
max_budget_in_team: Optional[float] = None
role: Optional[Literal["admin", "user"]] = None
tpm_limit: Optional[int] = Field(
default=None,
description="Tokens per minute limit for this team member"
default=None, description="Tokens per minute limit for this team member"
)
rpm_limit: Optional[int] = Field(
default=None,
description="Requests per minute limit for this team member"
default=None, description="Requests per minute limit for this team member"
)
class TeamMemberUpdateResponse(MemberUpdateResponse):

58
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -15,7 +15,7 @@ sys.path.insert(
import json
import sys
from typing import Any, AsyncGenerator, List, Literal, Optional, Tuple, Union
from litellm.secret_managers.main import get_secret_str
import httpx
from fastapi import HTTPException
@ -249,31 +249,46 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
data: dict,
optional_params: dict,
aws_region_name: str,
api_key: Optional[str] = None,
extra_headers: Optional[dict] = None,
):
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name)
api_base = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{self.guardrailIdentifier}/version/{self.guardrailVersion}/apply"
encoded_data = json.dumps(data).encode("utf-8")
headers = {"Content-Type": "application/json"}
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
api_base = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{self.guardrailIdentifier}/version/{self.guardrailVersion}/apply"
encoded_data = json.dumps(data).encode("utf-8")
# first check api-key, if none, fall back to sigV4
if api_key is not None:
aws_bearer_token: Optional[str] = api_key
else:
aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
request = AWSRequest(
method="POST", url=api_base, data=encoded_data, headers=headers
)
sigv4.add_auth(request)
if (
extra_headers is not None and "Authorization" in extra_headers
): # prevent sigv4 from overwriting the auth header
request.headers["Authorization"] = extra_headers["Authorization"]
if aws_bearer_token:
try:
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
headers["Authorization"] = f"Bearer {aws_bearer_token}"
request = AWSRequest(
method="POST", url=api_base, data=encoded_data, headers=headers
)
else:
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name)
request = AWSRequest(
method="POST", url=api_base, data=encoded_data, headers=headers
)
sigv4.add_auth(request)
if (
extra_headers is not None and "Authorization" in extra_headers
): # prevent sigv4 from overwriting the auth header
request.headers["Authorization"] = extra_headers["Authorization"]
prepped_request = request.prepare()
return prepped_request
@ -298,15 +313,20 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
bedrock_guardrail_response: BedrockGuardrailResponse = (
BedrockGuardrailResponse()
)
api_key: Optional[str] = None
if request_data:
bedrock_request_data.update(
self.get_guardrail_dynamic_request_body_params(request_data=request_data)
)
if request_data.get("api_key") is not None:
api_key = request_data["api_key"]
prepared_request = self._prepare_request(
credentials=credentials,
data=bedrock_request_data,
optional_params=self.optional_params,
aws_region_name=aws_region_name,
api_key=api_key,
)
verbose_proxy_logger.debug(
"Bedrock AI request body: %s, url %s, headers: %s",

7
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -17,7 +17,7 @@ from typing import (
Union,
cast,
)
from math import floor
from fastapi import HTTPException
from litellm import DualCache
@ -525,12 +525,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# Find which descriptor hit the limit
for i, status in enumerate(response["statuses"]):
if status["code"] == "OVER_LIMIT":
descriptor = descriptors[i]
descriptor = descriptors[floor(i/2)]
raise HTTPException(
status_code=429,
detail=f"Rate limit exceeded for {descriptor['key']}: {descriptor['value']}. Remaining: {status['limit_remaining']}",
headers={
"retry-after": str(self.window_size)
"retry-after": str(self.window_size),
"rate_limit_type": str(status["rate_limit_type"])
}, # Retry after 1 minute
)

19
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -367,6 +367,16 @@ async def _common_key_generation_helper( # noqa: PLR0915
premium_user=premium_user,
)
if (
data.metadata is not None
and data.metadata.get("service_account_id") is not None
and data.team_id is None
):
raise HTTPException(
status_code=400,
detail="team_id is required for service account keys. Please specify `team_id` in the request body.",
)
# check if user set default key/generate params on config.yaml
if litellm.default_key_generate_params is not None:
for elem in data:
@ -860,6 +870,15 @@ async def prepare_key_update_data(
data_json: dict = data.model_dump(exclude_unset=True)
data_json.pop("key", None)
data_json.pop("new_key", None)
if (
data.metadata is not None
and data.metadata.get("service_account_id") is not None
and (data.team_id or existing_key_row.team_id) is None
):
raise HTTPException(
status_code=400,
detail="team_id is required for service account keys. Please specify `team_id` in the request body.",
)
non_default_values = {}
for k, v in data_json.items():
if (

115
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -77,7 +77,9 @@ router = APIRouter()
@router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False)
async def google_login(request: Request, source: Optional[str] = None, key: Optional[str] = None): # noqa: PLR0915
async def google_login(
request: Request, source: Optional[str] = None, key: Optional[str] = None
): # noqa: PLR0915
"""
Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env
PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/"
@ -85,6 +87,7 @@ async def google_login(request: Request, source: Optional[str] = None, key: Opti
"""
from litellm.proxy.proxy_server import (
premium_user,
prisma_client,
user_custom_ui_sso_sign_in_handler,
)
@ -106,12 +109,23 @@ async def google_login(request: Request, source: Optional[str] = None, key: Opti
or generic_client_id is not None
):
if premium_user is not True:
raise ProxyException(
message="You must be a LiteLLM Enterprise user to use SSO. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this",
type=ProxyErrorTypes.auth_error,
param="premium_user",
code=status.HTTP_403_FORBIDDEN,
)
# Check if under 'free SSO user' limit
if prisma_client is not None:
total_users = await prisma_client.db.litellm_usertable.count()
if total_users and total_users > 5:
raise ProxyException(
message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this",
type=ProxyErrorTypes.auth_error,
param="premium_user",
code=status.HTTP_403_FORBIDDEN,
)
else:
raise ProxyException(
message=CommonProxyErrors.db_not_connected_error.value,
type=ProxyErrorTypes.auth_error,
param="premium_user",
code=status.HTTP_403_FORBIDDEN,
)
####### Detect DB + MASTER KEY in .env #######
missing_env_vars = show_missing_vars_in_env()
@ -124,7 +138,7 @@ async def google_login(request: Request, source: Optional[str] = None, key: Opti
request=request,
sso_callback_route="sso/callback",
)
# Store CLI key in state for OAuth flow
cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state(
source=source,
@ -137,11 +151,14 @@ async def google_login(request: Request, source: Optional[str] = None, key: Opti
from litellm_enterprise.proxy.auth.custom_sso_handler import (
EnterpriseCustomSSOHandler,
)
return await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in(
request=request,
)
except ImportError:
raise ValueError("Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise.")
raise ValueError(
"Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise."
)
# Check if we should use SSO handler
if (
@ -525,15 +542,16 @@ async def check_and_update_if_proxy_admin_id(
async def auth_callback(request: Request, state: Optional[str] = None): # noqa: PLR0915
"""Verify login"""
verbose_proxy_logger.info(f"Starting SSO callback with state: {state}")
# Check if this is a CLI login (state starts with our CLI prefix)
from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX
if state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"):
# Extract the key ID from the state
key_id = state.split(":", 1)[1]
verbose_proxy_logger.info(f"CLI SSO callback detected for key: {key_id}")
return await cli_sso_callback(request, key=key_id)
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.proxy_server import (
@ -608,7 +626,7 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa:
status_code=401,
detail="Result not returned by SSO provider.",
)
return await SSOAuthenticationHandler.get_redirect_response_from_openid(
result=result,
request=request,
@ -618,28 +636,26 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa:
)
async def cli_sso_callback(request: Request, key: Optional[str] = None):
"""CLI SSO callback - generates the key with pre-specified ID"""
verbose_proxy_logger.info(f"CLI SSO callback for key: {key}")
from litellm.proxy.management_endpoints.key_management_endpoints import (
generate_key_helper_fn,
)
from litellm.proxy.proxy_server import prisma_client
if not key or not key.startswith('sk-'):
if not key or not key.startswith("sk-"):
raise HTTPException(
status_code=400,
detail="Invalid key parameter. Must be a valid key ID starting with 'sk-'"
detail="Invalid key parameter. Must be a valid key ID starting with 'sk-'",
)
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
# Generate a simple key for CLI usage with the pre-specified key ID
try:
await generate_key_helper_fn(
@ -653,63 +669,57 @@ async def cli_sso_callback(request: Request, key: Optional[str] = None):
table_name="key",
token=key, # Use the pre-specified key ID
)
verbose_proxy_logger.info(f"Generated CLI key: {key}")
# Return success page
from fastapi.responses import HTMLResponse
from litellm.proxy.common_utils.html_forms.cli_sso_success import (
render_cli_sso_success_page,
)
html_content = render_cli_sso_success_page()
return HTMLResponse(content=html_content, status_code=200)
except Exception as e:
verbose_proxy_logger.error(f"Error generating CLI key: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to generate key: {str(e)}"
)
raise HTTPException(status_code=500, detail=f"Failed to generate key: {str(e)}")
@router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False)
async def cli_poll_key(key_id: str):
"""CLI polling endpoint - checks if key exists in DB"""
from litellm.proxy.proxy_server import prisma_client
if not key_id.startswith('sk-'):
raise HTTPException(
status_code=400,
detail="Invalid key ID format"
)
if not key_id.startswith("sk-"):
raise HTTPException(status_code=400, detail="Invalid key ID format")
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
try:
# Check if key exists in database
from litellm.proxy.utils import hash_token
hashed_token = hash_token(key_id)
key_obj = await prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": hashed_token}
)
if key_obj:
verbose_proxy_logger.info(f"CLI key found: {key_id}")
return {"status": "ready", "key": key_id}
else:
return {"status": "pending"}
except Exception as e:
verbose_proxy_logger.error(f"Error polling for CLI key: {e}")
raise HTTPException(
status_code=500,
detail=f"Error checking key status: {str(e)}"
status_code=500, detail=f"Error checking key status: {str(e)}"
)
@ -811,6 +821,7 @@ class SSOAuthenticationHandler:
"""
Handler for SSO Authentication across all SSO providers
"""
@staticmethod
async def get_sso_login_redirect(
redirect_url: str,
@ -1163,7 +1174,6 @@ class SSOAuthenticationHandler:
_new_team_request.update(_default_team_params)
team_request = NewTeamRequest(**_new_team_request)
return team_request
@staticmethod
def _get_cli_state(source: Optional[str], key: Optional[str]) -> Optional[str]:
@ -1176,13 +1186,15 @@ class SSOAuthenticationHandler:
LITELLM_CLI_SESSION_TOKEN_PREFIX,
LITELLM_CLI_SOURCE_IDENTIFIER,
)
return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" if source == LITELLM_CLI_SOURCE_IDENTIFIER and key else None
return (
f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}"
if source == LITELLM_CLI_SOURCE_IDENTIFIER and key
else None
)
@staticmethod
async def get_redirect_response_from_openid( # noqa: PLR0915
async def get_redirect_response_from_openid( # noqa: PLR0915
result: Union[OpenID, dict, CustomOpenID],
request: Request,
received_response: Optional[dict] = None,
@ -1202,14 +1214,18 @@ class SSOAuthenticationHandler:
)
from litellm.proxy.utils import get_prisma_client_or_throw
from litellm.types.proxy.ui_sso import ReturnedUITokenObject
prisma_client = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy")
prisma_client = get_prisma_client_or_throw(
"Prisma client is None, connect a database to your proxy"
)
# User is Authe'd in - generate key for the UI to access Proxy
verbose_proxy_logger.info(f"SSO callback result: {result}")
user_email: Optional[str] = getattr(result, "email", None)
user_id: Optional[str] = getattr(result, "id", None) if result is not None else None
user_id: Optional[str] = (
getattr(result, "id", None) if result is not None else None
)
if user_email is not None and os.getenv("ALLOWED_EMAIL_DOMAINS") is not None:
email_domain = user_email.split("@")[1]
@ -1394,7 +1410,8 @@ class SSOAuthenticationHandler:
redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
redirect_response.set_cookie(key="token", value=jwt_token)
return redirect_response
class MicrosoftSSOHandler:
"""
Handles Microsoft SSO callback response and returns a CustomOpenID object

1
litellm/proxy/pass_through_endpoints/common_utils.py
View file

@ -14,3 +14,4 @@ def get_litellm_virtual_key(request: Request) -> str:
if litellm_api_key:
return f"Bearer {litellm_api_key}"
return request.headers.get("Authorization", "")

12
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -492,7 +492,12 @@ async def bedrock_llm_proxy_route(
data: Dict[str, Any] = {}
base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data)
try:
model = endpoint.split("/")[1]
endpoint_parts = endpoint.split("/")
if "application-inference-profile" in endpoint:
# For application-inference-profile, include the profile ID part as well
model = "/".join(endpoint_parts[1:3])
else:
model = endpoint_parts[1]
except Exception:
raise HTTPException(
status_code=400,
@ -500,12 +505,13 @@ async def bedrock_llm_proxy_route(
"error": "Model missing from endpoint. Expected format: /model/<Model>/<endpoint>. Got: "
+ endpoint,
},
)
)
data["method"] = request.method
data["endpoint"] = endpoint
data["data"] = request_body
data["custom_llm_provider"] = "bedrock"
try:
result = await base_llm_response_processor.base_passthrough_process_llm_request(
request=request,

37
litellm/proxy/spend_tracking/cold_storage_handler.py
View file

@ -3,11 +3,10 @@ This module is responsible for handling Getting/Setting the proxy server request
It allows fetching a dict of the proxy server request from s3 or GCS bucket.
"""
from typing import Optional, cast
from typing import Optional
import litellm
from litellm import _custom_logger_compatible_callbacks_literal
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
@ -57,38 +56,6 @@ class ColdStorageHandler:
def _select_custom_logger_for_cold_storage(
self,
) -> Optional[_custom_logger_compatible_callbacks_literal]:
cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = ColdStorageHandler._get_configured_cold_storage_custom_logger()
cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = litellm.configured_cold_storage_logger
return cold_storage_custom_logger
@staticmethod
def _get_configured_cold_storage_custom_logger() -> Optional[_custom_logger_compatible_callbacks_literal]:
"""Return the configured cold storage custom logger.
During interpreter shutdown importing ``proxy_server`` can raise a
``RuntimeError`` (e.g. "can't register atexit after shutdown").
In these scenarios we gracefully return ``None`` instead of bubbling
the exception up the call stack.
"""
try:
from litellm.proxy.proxy_server import general_settings
except Exception as e:
verbose_proxy_logger.debug(
f"Unable to import proxy_server for cold storage logging: {e}"
)
return None
cold_storage_custom_logger: Optional[str] = general_settings.get(
"cold_storage_custom_logger"
)
if not cold_storage_custom_logger:
verbose_proxy_logger.debug(
"No cold storage custom logger found in general settings"
)
return None
return cast(
_custom_logger_compatible_callbacks_literal, cold_storage_custom_logger
)

63
litellm/rerank_api/main.py
View file

@ -29,7 +29,7 @@ async def arerank(
model: str,
query: str,
documents: List[Union[str, Dict[str, Any]]],
custom_llm_provider: Optional[Literal["cohere", "together_ai"]] = None,
custom_llm_provider: Optional[Literal["cohere", "together_ai", "deepinfra"]] = None,
top_n: Optional[int] = None,
rank_fields: Optional[List[str]] = None,
return_documents: Optional[bool] = None,
@ -75,7 +75,15 @@ def rerank( # noqa: PLR0915
query: str,
documents: List[Union[str, Dict[str, Any]]],
custom_llm_provider: Optional[
Literal["cohere", "together_ai", "azure_ai", "infinity", "litellm_proxy", "hosted_vllm"]
Literal[
"cohere",
"together_ai",
"azure_ai",
"infinity",
"litellm_proxy",
"hosted_vllm",
"deepinfra",
]
] = None,
top_n: Optional[int] = None,
rank_fields: Optional[List[str]] = None,
@ -142,7 +150,7 @@ def rerank( # noqa: PLR0915
max_tokens_per_doc=max_tokens_per_doc,
non_default_params=kwargs,
)
verbose_logger.info(f"optional_rerank_params: {optional_rerank_params}")
if isinstance(optional_params.timeout, str):
optional_params.timeout = float(optional_params.timeout)
@ -356,18 +364,57 @@ def rerank( # noqa: PLR0915
client=client,
model_response=model_response,
)
elif _custom_llm_provider == "deepinfra":
api_key = (
dynamic_api_key
or optional_params.api_key
or get_secret_str("DEEPINFRA_API_KEY")
)
api_base = (
dynamic_api_base
or optional_params.api_base
or get_secret_str("DEEPINFRA_API_BASE")
)
if api_base is None:
raise ValueError(
"api_base must be provided for Deepinfra rerank. Set in call or via DEEPINFRA_API_BASE env var."
)
response = base_llm_http_handler.rerank(
model=model,
custom_llm_provider=_custom_llm_provider,
provider_config=rerank_provider_config,
optional_rerank_params=optional_rerank_params,
logging_obj=litellm_logging_obj,
timeout=optional_params.timeout,
api_key=api_key,
api_base=api_base,
_is_async=_is_async,
headers=headers or litellm.headers or {},
client=client,
model_response=model_response,
)
else:
# Generic handler for all providers that use base_llm_http_handler
# Provider-specific logic (API key validation, URL generation, etc.)
# Provider-specific logic (API key validation, URL generation, etc.)
# is handled in the respective transformation configs
# Check if the provider is actually supported
# If rerank_provider_config is a default CohereRerankConfig but the provider is not Cohere or litellm_proxy,
# it means the provider is not supported
if (isinstance(rerank_provider_config, litellm.CohereRerankConfig) or
isinstance(rerank_provider_config, litellm.CohereRerankV2Config)) and _custom_llm_provider != "cohere" and _custom_llm_provider != "litellm_proxy":
if (
(
isinstance(rerank_provider_config, litellm.CohereRerankConfig)
or isinstance(rerank_provider_config, litellm.CohereRerankV2Config)
)
and _custom_llm_provider != "cohere"
and _custom_llm_provider != "litellm_proxy"
):
raise ValueError(f"Unsupported provider: {_custom_llm_provider}")
response = base_llm_http_handler.rerank(
model=model,
custom_llm_provider=_custom_llm_provider,

5
litellm/responses/litellm_completion_transformation/session_handler.py
View file

@ -1,6 +1,7 @@
import json
from typing import TYPE_CHECKING, Any, List, Optional, Union, cast
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import SpendLogsPayload
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
@ -235,10 +236,10 @@ class ResponsesSessionHandler:
"""
Only check cold storage when both are true
1. `LITELLM_TRUNCATED_PAYLOAD_FIELD` is in the proxy server request dict
2. `ColdStorageHandler._get_configured_cold_storage_custom_logger()` is not None
2. `litellm.configured_cold_storage_logger` is not None
"""
from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD
configured_cold_storage_custom_logger = ColdStorageHandler._get_configured_cold_storage_custom_logger()
configured_cold_storage_custom_logger = litellm.configured_cold_storage_logger
if configured_cold_storage_custom_logger is None:
return False
if proxy_server_request_dict is None:

84
litellm/responses/main.py
View file

@ -89,6 +89,7 @@ def mock_responses_api_response(
}
)
async def aresponses_api_with_mcp(
input: Union[str, ResponseInputParam],
model: str,
@ -122,7 +123,7 @@ async def aresponses_api_with_mcp(
) -> Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]:
"""
Async version of responses API with MCP integration.
When MCP tools with server_url="litellm_proxy" are provided, this function will:
1. Get available tools from the MCP server manager
2. Insert the tools into the messages/input
@ -134,19 +135,25 @@ async def aresponses_api_with_mcp(
)
# Parse MCP tools and separate from other tools
mcp_tools_with_litellm_proxy, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
mcp_tools_with_litellm_proxy, other_tools = (
LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
)
# Get available tools from MCP manager if we have MCP tools
openai_tools = []
mcp_tools_fetched = []
if mcp_tools_with_litellm_proxy:
user_api_key_auth = kwargs.get("user_api_key_auth")
mcp_tools_fetched = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager(user_api_key_auth)
openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(mcp_tools_fetched)
mcp_tools_fetched = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager(
user_api_key_auth
)
openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(
mcp_tools_fetched
)
# Combine with other tools
all_tools = openai_tools + other_tools if (openai_tools or other_tools) else None
# Prepare call parameters for reuse
call_params = {
"include": include,
@ -172,7 +179,7 @@ async def aresponses_api_with_mcp(
"custom_llm_provider": custom_llm_provider,
**kwargs,
}
# Make initial response API call
# TODO: if should auto-execute is True, then this first response should not be streamed
response = await aresponses(
@ -180,45 +187,54 @@ async def aresponses_api_with_mcp(
model=model,
tools=all_tools,
previous_response_id=previous_response_id,
**call_params
**call_params,
)
# Check if we need to auto-execute tool calls (only for non-streaming responses)
if (mcp_tools_with_litellm_proxy and
isinstance(response, ResponsesAPIResponse) and
LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy)): # type: ignore
tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_response(response=response)
if (
mcp_tools_with_litellm_proxy
and isinstance(response, ResponsesAPIResponse)
and LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(
mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy
)
): # type: ignore
tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_response(
response=response
)
if tool_calls:
user_api_key_auth = kwargs.get("litellm_metadata", {}).get("user_api_key_auth")
tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(tool_calls=tool_calls, user_api_key_auth=user_api_key_auth)
user_api_key_auth = kwargs.get("litellm_metadata", {}).get(
"user_api_key_auth"
)
tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
tool_calls=tool_calls, user_api_key_auth=user_api_key_auth
)
if tool_results:
follow_up_input = LiteLLM_Proxy_MCP_Handler._create_follow_up_input(
response=response,
tool_results=tool_results,
original_input=input
response=response, tool_results=tool_results, original_input=input
)
final_response = await LiteLLM_Proxy_MCP_Handler._make_follow_up_call(
follow_up_input=follow_up_input,
model=model,
all_tools=all_tools,
response_id=response.id,
**call_params
**call_params,
)
# Add custom output elements to the final response
if isinstance(final_response, ResponsesAPIResponse):
final_response = LiteLLM_Proxy_MCP_Handler._add_mcp_output_elements_to_response(
response=final_response,
mcp_tools_fetched=mcp_tools_fetched,
tool_results=tool_results
final_response = (
LiteLLM_Proxy_MCP_Handler._add_mcp_output_elements_to_response(
response=final_response,
mcp_tools_fetched=mcp_tools_fetched,
tool_results=tool_results,
)
)
return final_response
return response
return response
@client
@ -319,7 +335,9 @@ async def aresponses(
)
if response is None:
raise ValueError(f"Got an unexpected None response from the Responses API: {response}")
raise ValueError(
f"Got an unexpected None response from the Responses API: {response}"
)
return response
except Exception as e:
@ -363,6 +381,7 @@ def responses(
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
# LiteLLM specific params,
allowed_openai_params: Optional[List[str]] = None,
custom_llm_provider: Optional[str] = None,
**kwargs,
):
@ -373,7 +392,7 @@ def responses(
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
@ -445,6 +464,7 @@ def responses(
model=model,
responses_api_provider_config=responses_api_provider_config,
response_api_optional_params=response_api_optional_params,
allowed_openai_params=allowed_openai_params,
)
)

100
litellm/responses/utils.py
View file

@ -1,5 +1,5 @@
import base64
from typing import Any, Dict, Optional, Union, cast, get_type_hints, overload
from typing import Any, Dict, List, Optional, Union, cast, get_type_hints, overload
import litellm
from litellm._logging import verbose_logger
@ -16,11 +16,38 @@ from litellm.types.utils import SpecialEnums, Usage
class ResponsesAPIRequestUtils:
"""Helper utils for constructing ResponseAPI requests"""
@staticmethod
def _check_valid_arg(
supported_params: Optional[List[str]],
non_default_params: Dict,
drop_params: Optional[bool],
custom_llm_provider: Optional[str],
model: str,
):
if supported_params is None:
return
unsupported_params = {}
for k in non_default_params.keys():
if k not in supported_params:
unsupported_params[k] = non_default_params[k]
if unsupported_params:
if litellm.drop_params is True or (
drop_params is not None and drop_params is True
):
pass
else:
raise litellm.UnsupportedParamsError(
status_code=500,
message=f"{custom_llm_provider} does not support parameters: {unsupported_params}, for model={model}. To drop these, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n",
)
@staticmethod
def get_optional_params_responses_api(
model: str,
responses_api_provider_config: BaseResponsesAPIConfig,
response_api_optional_params: ResponsesAPIOptionalRequestParams,
allowed_openai_params: Optional[List[str]] = None,
) -> Dict:
"""
Get optional parameters for the responses API.
@ -33,25 +60,23 @@ class ResponsesAPIRequestUtils:
Returns:
A dictionary of supported parameters for the responses API
"""
# Remove None values and internal parameters
from litellm.utils import _apply_openai_param_overrides
# Remove None values and internal parameters
# Get supported parameters for the model
supported_params = responses_api_provider_config.get_supported_openai_params(
model
)
non_default_params = cast(Dict, response_api_optional_params)
# Check for unsupported parameters
unsupported_params = [
param
for param in response_api_optional_params
if param not in supported_params
]
if unsupported_params:
raise litellm.UnsupportedParamsError(
model=model,
message=f"The following parameters are not supported for model {model}: {', '.join(unsupported_params)}",
)
ResponsesAPIRequestUtils._check_valid_arg(
supported_params=supported_params + (allowed_openai_params or []),
non_default_params=non_default_params,
drop_params=litellm.drop_params,
custom_llm_provider=responses_api_provider_config.custom_llm_provider,
model=model,
)
# Map parameters to provider-specific format
mapped_params = responses_api_provider_config.map_openai_params(
@ -60,6 +85,13 @@ class ResponsesAPIRequestUtils:
drop_params=litellm.drop_params,
)
# add any allowed_openai_params to the mapped_params
mapped_params = _apply_openai_param_overrides(
optional_params=mapped_params,
non_default_params=non_default_params,
allowed_openai_params=allowed_openai_params or [],
)
return mapped_params
@staticmethod
@ -75,34 +107,48 @@ class ResponsesAPIRequestUtils:
Returns:
ResponsesAPIOptionalRequestParams instance with only the valid parameters
"""
from litellm.utils import PreProcessNonDefaultParams
valid_keys = get_type_hints(ResponsesAPIOptionalRequestParams).keys()
filtered_params = {
k: v for k, v in params.items() if k in valid_keys and v is not None
}
custom_llm_provider = params.pop("custom_llm_provider", None)
special_params = params.pop("kwargs", {})
additional_drop_params = params.pop("additional_drop_params", None)
non_default_params = (
PreProcessNonDefaultParams.base_pre_process_non_default_params(
passed_params=params,
special_params=special_params,
custom_llm_provider=custom_llm_provider,
additional_drop_params=additional_drop_params,
default_param_values={k: None for k in valid_keys},
additional_endpoint_specific_params=["input"],
)
)
# decode previous_response_id if it's a litellm encoded id
if "previous_response_id" in filtered_params:
if "previous_response_id" in non_default_params:
decoded_previous_response_id = ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(
filtered_params["previous_response_id"]
non_default_params["previous_response_id"]
)
filtered_params["previous_response_id"] = decoded_previous_response_id
non_default_params["previous_response_id"] = decoded_previous_response_id
if "metadata" in filtered_params:
if "metadata" in non_default_params:
from litellm.utils import add_openai_metadata
filtered_params["metadata"] = add_openai_metadata(
filtered_params["metadata"]
non_default_params["metadata"] = add_openai_metadata(
non_default_params["metadata"]
)
return cast(ResponsesAPIOptionalRequestParams, filtered_params)
return cast(ResponsesAPIOptionalRequestParams, non_default_params)
# fmt: off
@overload
@staticmethod
def _update_responses_api_response_id_with_model_id(
responses_api_response: ResponsesAPIResponse,
custom_llm_provider: Optional[str],
litellm_metadata: Optional[Dict[str, Any]] = None,
) -> ResponsesAPIResponse:
) -> ResponsesAPIResponse:
...
@overload
@ -111,9 +157,11 @@ class ResponsesAPIRequestUtils:
responses_api_response: Dict[str, Any],
custom_llm_provider: Optional[str],
litellm_metadata: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
) -> Dict[str, Any]:
...
# fmt: on
@staticmethod
def _update_responses_api_response_id_with_model_id(
responses_api_response: Union[ResponsesAPIResponse, Dict[str, Any]],

10
litellm/types/llms/mistral.py
View file

@ -10,3 +10,13 @@ class MistralToolCallMessage(TypedDict):
id: Optional[str]
type: Literal["function"]
function: Optional[FunctionCall]
class MistralTextBlock(TypedDict):
type: Literal["text"]
text: str
class MistralThinkingBlock(TypedDict):
type: Literal["thinking"]
thinking: List[MistralTextBlock]

4
litellm/types/llms/openai.py
View file

@ -974,6 +974,10 @@ class ResponsesAPIOptionalRequestParams(TypedDict, total=False):
service_tier: Optional[str]
safety_identifier: Optional[str]
prompt: Optional[PromptObject]
max_tool_calls: Optional[int]
prompt_cache_key: Optional[str]
stream_options: Optional[dict]
top_logprobs: Optional[int]
class ResponsesAPIRequestParams(ResponsesAPIOptionalRequestParams, total=False):

126
litellm/utils.py
View file

@ -541,9 +541,9 @@ def function_setup( # noqa: PLR0915
function_id: Optional[str] = kwargs["id"] if "id" in kwargs else None
## DYNAMIC CALLBACKS ##
dynamic_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = (
kwargs.pop("callbacks", None)
)
dynamic_callbacks: Optional[
List[Union[str, Callable, CustomLogger]]
] = kwargs.pop("callbacks", None)
all_callbacks = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks)
if len(all_callbacks) > 0:
@ -1287,9 +1287,9 @@ def client(original_function): # noqa: PLR0915
exception=e,
retry_policy=kwargs.get("retry_policy"),
)
kwargs["retry_policy"] = (
reset_retry_policy()
) # prevent infinite loops
kwargs[
"retry_policy"
] = reset_retry_policy() # prevent infinite loops
litellm.num_retries = (
None # set retries to None to prevent infinite loops
)
@ -2326,47 +2326,47 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915
# add new model names to provider lists
if value.get("litellm_provider") == "openai":
if key not in litellm.open_ai_chat_completion_models:
litellm.open_ai_chat_completion_models.append(key)
litellm.open_ai_chat_completion_models.add(key)
elif value.get("litellm_provider") == "text-completion-openai":
if key not in litellm.open_ai_text_completion_models:
litellm.open_ai_text_completion_models.append(key)
litellm.open_ai_text_completion_models.add(key)
elif value.get("litellm_provider") == "cohere":
if key not in litellm.cohere_models:
litellm.cohere_models.append(key)
litellm.cohere_models.add(key)
elif value.get("litellm_provider") == "anthropic":
if key not in litellm.anthropic_models:
litellm.anthropic_models.append(key)
litellm.anthropic_models.add(key)
elif value.get("litellm_provider") == "openrouter":
split_string = key.split("/", 1)
if key not in litellm.openrouter_models:
litellm.openrouter_models.append(split_string[1])
litellm.openrouter_models.add(split_string[1])
elif value.get("litellm_provider") == "vertex_ai-text-models":
if key not in litellm.vertex_text_models:
litellm.vertex_text_models.append(key)
litellm.vertex_text_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-code-text-models":
if key not in litellm.vertex_code_text_models:
litellm.vertex_code_text_models.append(key)
litellm.vertex_code_text_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-chat-models":
if key not in litellm.vertex_chat_models:
litellm.vertex_chat_models.append(key)
litellm.vertex_chat_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-code-chat-models":
if key not in litellm.vertex_code_chat_models:
litellm.vertex_code_chat_models.append(key)
litellm.vertex_code_chat_models.add(key)
elif value.get("litellm_provider") == "ai21":
if key not in litellm.ai21_models:
litellm.ai21_models.append(key)
litellm.ai21_models.add(key)
elif value.get("litellm_provider") == "nlp_cloud":
if key not in litellm.nlp_cloud_models:
litellm.nlp_cloud_models.append(key)
litellm.nlp_cloud_models.add(key)
elif value.get("litellm_provider") == "aleph_alpha":
if key not in litellm.aleph_alpha_models:
litellm.aleph_alpha_models.append(key)
litellm.aleph_alpha_models.add(key)
elif value.get("litellm_provider") == "bedrock":
if key not in litellm.bedrock_models:
litellm.bedrock_models.append(key)
litellm.bedrock_models.add(key)
elif value.get("litellm_provider") == "novita":
if key not in litellm.novita_models:
litellm.novita_models.append(key)
litellm.novita_models.add(key)
return model_cost
@ -2812,12 +2812,22 @@ def get_optional_params_embeddings( # noqa: PLR0915
request_type="embeddings",
)
_check_valid_arg(supported_params=supported_params)
optional_params = litellm.VoyageEmbeddingConfig().map_openai_params(
non_default_params=non_default_params,
optional_params={},
model=model,
drop_params=drop_params if drop_params is not None else False,
)
if litellm.VoyageContextualEmbeddingConfig.is_contextualized_embeddings(model):
optional_params = (
litellm.VoyageContextualEmbeddingConfig().map_openai_params(
non_default_params=non_default_params,
optional_params={},
model=model,
drop_params=drop_params if drop_params is not None else False,
)
)
else:
optional_params = litellm.VoyageEmbeddingConfig().map_openai_params(
non_default_params=non_default_params,
optional_params={},
model=model,
drop_params=drop_params if drop_params is not None else False,
)
elif custom_llm_provider == "infinity":
supported_params = get_supported_openai_params(
model=model,
@ -3091,10 +3101,10 @@ def pre_process_non_default_params(
if "response_format" in non_default_params:
if provider_config is not None:
non_default_params["response_format"] = (
provider_config.get_json_schema_from_pydantic_object(
response_format=non_default_params["response_format"]
)
non_default_params[
"response_format"
] = provider_config.get_json_schema_from_pydantic_object(
response_format=non_default_params["response_format"]
)
else:
non_default_params["response_format"] = type_to_response_format_param(
@ -3221,16 +3231,16 @@ def pre_process_optional_params(
True # so that main.py adds the function call to the prompt
)
if "tools" in non_default_params:
optional_params["functions_unsupported_model"] = (
non_default_params.pop("tools")
)
optional_params[
"functions_unsupported_model"
] = non_default_params.pop("tools")
non_default_params.pop(
"tool_choice", None
) # causes ollama requests to hang
elif "functions" in non_default_params:
optional_params["functions_unsupported_model"] = (
non_default_params.pop("functions")
)
optional_params[
"functions_unsupported_model"
] = non_default_params.pop("functions")
elif (
litellm.add_function_to_prompt
): # if user opts to add it to prompt instead
@ -4324,9 +4334,9 @@ def _count_characters(text: str) -> int:
def get_response_string(response_obj: Union[ModelResponse, ModelResponseStream]) -> str:
_choices: Union[List[Union[Choices, StreamingChoices]], List[StreamingChoices]] = (
response_obj.choices
)
_choices: Union[
List[Union[Choices, StreamingChoices]], List[StreamingChoices]
] = response_obj.choices
response_str = ""
for choice in _choices:
@ -5217,6 +5227,7 @@ def validate_environment( # noqa: PLR0915
model: Optional[str] = None,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
) -> dict:
"""
Checks if the environment variables are valid for the given model.
@ -5371,6 +5382,11 @@ def validate_environment( # noqa: PLR0915
keys_in_environment = True
else:
missing_keys.append("CEREBRAS_API_KEY")
elif custom_llm_provider == "baseten":
if "BASETEN_API_KEY" in os.environ:
keys_in_environment = True
else:
missing_keys.append("BASETEN_API_KEY")
elif custom_llm_provider == "xai":
if "XAI_API_KEY" in os.environ:
keys_in_environment = True
@ -5562,19 +5578,18 @@ def validate_environment( # noqa: PLR0915
else:
missing_keys.append("NEBIUS_API_KEY")
def filter_missing_keys(keys: List[str], exclude_pattern: str) -> List[str]:
"""Filter out keys that contain the exclude_pattern (case insensitive)."""
return [key for key in keys if exclude_pattern not in key.lower()]
if api_key is not None:
new_missing_keys = []
for key in missing_keys:
if "api_key" not in key.lower():
new_missing_keys.append(key)
missing_keys = new_missing_keys
missing_keys = filter_missing_keys(missing_keys, "api_key")
if api_base is not None:
new_missing_keys = []
for key in missing_keys:
if "api_base" not in key.lower():
new_missing_keys.append(key)
missing_keys = new_missing_keys
missing_keys = filter_missing_keys(missing_keys, "api_base")
if api_version is not None:
missing_keys = filter_missing_keys(missing_keys, "api_version")
if len(missing_keys) == 0: # no missing keys
keys_in_environment = True
@ -6926,6 +6941,8 @@ class ProviderConfigManager:
return litellm.NvidiaNimConfig()
elif litellm.LlmProviders.CEREBRAS == provider:
return litellm.CerebrasConfig()
elif litellm.LlmProviders.BASETEN == provider:
return litellm.BasetenConfig()
elif litellm.LlmProviders.VOLCENGINE == provider:
return litellm.VolcEngineConfig()
elif litellm.LlmProviders.TEXT_COMPLETION_CODESTRAL == provider:
@ -7023,7 +7040,14 @@ class ProviderConfigManager:
model: str,
provider: LlmProviders,
) -> Optional[BaseEmbeddingConfig]:
if litellm.LlmProviders.VOYAGE == provider:
if (
litellm.LlmProviders.VOYAGE == provider
and litellm.VoyageContextualEmbeddingConfig.is_contextualized_embeddings(
model
)
):
return litellm.VoyageContextualEmbeddingConfig()
elif litellm.LlmProviders.VOYAGE == provider:
return litellm.VoyageEmbeddingConfig()
elif litellm.LlmProviders.TRITON == provider:
return litellm.TritonEmbeddingConfig()
@ -7071,6 +7095,8 @@ class ProviderConfigManager:
return litellm.JinaAIRerankConfig()
elif litellm.LlmProviders.HUGGINGFACE == provider:
return litellm.HuggingFaceRerankConfig()
elif litellm.LlmProviders.DEEPINFRA == provider:
return litellm.DeepinfraRerankConfig()
return litellm.CohereRerankConfig()
@staticmethod

1284
model_prices_and_context_window.json
View file

File diff suppressed because it is too large Load diff

4
pyproject.toml
View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm"
version = "1.75.9"
version = "1.76.0"
description = "Library to easily interface with LLM API providers"
authors = ["BerriAI"]
license = "MIT"
@ -155,7 +155,7 @@ requires = ["poetry-core", "wheel"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "1.75.9"
version = "1.76.0"
version_files = [
"pyproject.toml:^version"
]

24
tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -159,27 +159,3 @@ class TestAvailableEnterpriseUsers:
CommonProxyErrors.db_not_connected_error.value
in response.json()["detail"]["error"]
)
@pytest.mark.asyncio
async def test_available_users_not_premium_user(
self, client, mock_user_api_key_auth
):
"""Test when premium_user is None (not a premium user)"""
from litellm.proxy._types import CommonProxyErrors
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.proxy_server.premium_user",
None,
):
# Override the dependency
client.app.dependency_overrides[mock_user_api_key_auth] = lambda: {
"user_id": "test_user"
}
response = client.get("/user/available_users")
assert response.status_code == 500
assert (
CommonProxyErrors.not_premium_user.value
in response.json()["detail"]["error"]
)

27
tests/image_gen_tests/base_image_generation_test.py
View file

@ -94,4 +94,29 @@ class BaseImageGenTest(ABC):
if "Your task failed as a result of our safety system." in str(e):
pass
else:
pytest.fail(f"An exception occurred - {str(e)}")
pytest.fail(f"An exception occurred - {str(e)}")
@pytest.mark.skip(reason="Skipping image edit test, image file not in ci/cd")
def test_openai_gpt_image_1():
from litellm import image_edit
from PIL import Image
import io
# Create a simple mask image with alpha channel
# Create a 512x512 black image with alpha channel
try:
response = image_edit(
model="openai/gpt-image-1",
image=open("test_image_edit.png", "rb"),
mask=open("test_image_edit.png", "rb"),
prompt="Add a red hat to the person in the image",
n=1,
size="1024x1024",
)
print("response: ", response)
except Exception as e:
if "mask image missing alpha channel" in str(e):
pass
else:
raise e

17
tests/litellm_utils_tests/test_utils.py
View file

@ -338,11 +338,9 @@ def test_aget_valid_models():
print(valid_models)
# list of openai supported llms on litellm
expected_models = (
litellm.open_ai_chat_completion_models + litellm.open_ai_text_completion_models
)
expected_models = litellm.open_ai_chat_completion_models | litellm.open_ai_text_completion_models
assert valid_models == expected_models
assert set(valid_models) == set(expected_models)
# reset replicate env key
os.environ = old_environ
@ -355,7 +353,7 @@ def test_aget_valid_models():
valid_models = get_valid_models()
print(valid_models)
assert valid_models == expected_models
assert set(valid_models) == set(expected_models)
# reset replicate env key
os.environ = old_environ
@ -376,7 +374,7 @@ def test_get_valid_models_with_custom_llm_provider(custom_llm_provider):
)
print(valid_models)
assert len(valid_models) > 0
assert provider_config.get_models() == valid_models
assert set(provider_config.get_models()) == set(valid_models)
# test_get_valid_models()
@ -411,6 +409,13 @@ def test_validate_environment_api_key():
), f"Missing keys={response_obj['missing_keys']}"
def test_validate_environment_api_version():
response_obj = validate_environment(model="azure/openai-deployment", api_key="sk-my-test-key", api_base="https://fake.openai.azure.com/", api_version="2024-02-15")
assert (
response_obj["keys_in_environment"] is True
), f"Missing keys={response_obj['missing_keys']}"
def test_validate_environment_api_base_dynamic():
for provider in ["ollama", "ollama_chat"]:
kv = validate_environment(provider + "/mistral", api_base="https://example.com")

224
tests/llm_responses_api_testing/base_responses_api.py
View file

@ -101,11 +101,11 @@ def validate_responses_api_response(response, final_chunk: bool = False):
return True # Return True if validation passes
class BaseResponsesAPITest(ABC):
"""
Abstract base test class that enforces a common test across all test classes.
"""
@abstractmethod
def get_base_completion_call_args(self) -> dict:
"""Must return the base completion call args"""
@ -115,32 +115,32 @@ class BaseResponsesAPITest(ABC):
"""Must return the base completion reasoning call args"""
return None
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_basic_openai_responses_api(self, sync_mode):
litellm._turn_on_debug()
litellm.set_verbose = True
base_completion_call_args = self.get_base_completion_call_args()
try:
try:
if sync_mode:
response = litellm.responses(
input="Basic ping", max_output_tokens=20,
**base_completion_call_args
input="Basic ping",
max_output_tokens=20,
**base_completion_call_args,
)
else:
response = await litellm.aresponses(
input="Basic ping", max_output_tokens=20,
**base_completion_call_args
input="Basic ping",
max_output_tokens=20,
**base_completion_call_args,
)
except litellm.InternalServerError:
except litellm.InternalServerError:
pytest.skip("Skipping test due to litellm.InternalServerError")
print("litellm response=", json.dumps(response, indent=4, default=str))
# Use the helper function to validate the response
validate_responses_api_response(response, final_chunk=True)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=2)
@ -151,9 +151,7 @@ class BaseResponsesAPITest(ABC):
response_completed_event = None
if sync_mode:
response = litellm.responses(
input="Basic ping",
stream=True,
**base_completion_call_args
input="Basic ping", stream=True, **base_completion_call_args
)
for event in response:
print("litellm response=", json.dumps(event, indent=4, default=str))
@ -163,9 +161,7 @@ class BaseResponsesAPITest(ABC):
response_completed_event = event
else:
response = await litellm.aresponses(
input="Basic ping",
stream=True,
**base_completion_call_args
input="Basic ping", stream=True, **base_completion_call_args
)
async for event in response:
print("litellm response=", json.dumps(event, indent=4, default=str))
@ -188,15 +184,29 @@ class BaseResponsesAPITest(ABC):
assert response_completed_event.response.usage is not None
# basic test assert the usage seems reasonable
print("response_completed_event.response.usage=", response_completed_event.response.usage)
assert response_completed_event.response.usage.input_tokens > 0 and response_completed_event.response.usage.input_tokens < 100
assert response_completed_event.response.usage.output_tokens > 0 and response_completed_event.response.usage.output_tokens < 2000
assert response_completed_event.response.usage.total_tokens > 0 and response_completed_event.response.usage.total_tokens < 2000
print(
"response_completed_event.response.usage=",
response_completed_event.response.usage,
)
assert (
response_completed_event.response.usage.input_tokens > 0
and response_completed_event.response.usage.input_tokens < 100
)
assert (
response_completed_event.response.usage.output_tokens > 0
and response_completed_event.response.usage.output_tokens < 2000
)
assert (
response_completed_event.response.usage.total_tokens > 0
and response_completed_event.response.usage.total_tokens < 2000
)
# total tokens should be the sum of input and output tokens
assert response_completed_event.response.usage.total_tokens == response_completed_event.response.usage.input_tokens + response_completed_event.response.usage.output_tokens
assert (
response_completed_event.response.usage.total_tokens
== response_completed_event.response.usage.input_tokens
+ response_completed_event.response.usage.output_tokens
)
@pytest.mark.parametrize("sync_mode", [False, True])
@pytest.mark.asyncio
@ -206,48 +216,44 @@ class BaseResponsesAPITest(ABC):
base_completion_call_args = self.get_base_completion_call_args()
if sync_mode:
response = litellm.responses(
input="Basic ping", max_output_tokens=20,
**base_completion_call_args
input="Basic ping", max_output_tokens=20, **base_completion_call_args
)
# delete the response
if isinstance(response, ResponsesAPIResponse):
litellm.delete_responses(
response_id=response.id,
**base_completion_call_args
response_id=response.id, **base_completion_call_args
)
else:
raise ValueError("response is not a ResponsesAPIResponse")
else:
response = await litellm.aresponses(
input="Basic ping", max_output_tokens=20,
**base_completion_call_args
input="Basic ping", max_output_tokens=20, **base_completion_call_args
)
# async delete the response
if isinstance(response, ResponsesAPIResponse):
await litellm.adelete_responses(
response_id=response.id,
**base_completion_call_args
response_id=response.id, **base_completion_call_args
)
else:
raise ValueError("response is not a ResponsesAPIResponse")
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio
async def test_basic_openai_responses_streaming_delete_endpoint(self, sync_mode):
#litellm._turn_on_debug()
#litellm.set_verbose = True
# litellm._turn_on_debug()
# litellm.set_verbose = True
base_completion_call_args = self.get_base_completion_call_args()
response_id = None
if sync_mode:
response_id = None
response = litellm.responses(
input="Basic ping", max_output_tokens=20,
input="Basic ping",
max_output_tokens=20,
stream=True,
**base_completion_call_args
**base_completion_call_args,
)
for event in response:
print("litellm response=", json.dumps(event, indent=4, default=str))
@ -260,14 +266,14 @@ class BaseResponsesAPITest(ABC):
# delete the response
assert response_id is not None
litellm.delete_responses(
response_id=response_id,
**base_completion_call_args
response_id=response_id, **base_completion_call_args
)
else:
response = await litellm.aresponses(
input="Basic ping", max_output_tokens=20,
input="Basic ping",
max_output_tokens=20,
stream=True,
**base_completion_call_args
**base_completion_call_args,
)
async for event in response:
print("litellm response=", json.dumps(event, indent=4, default=str))
@ -280,8 +286,7 @@ class BaseResponsesAPITest(ABC):
# delete the response
assert response_id is not None
await litellm.adelete_responses(
response_id=response_id,
**base_completion_call_args
response_id=response_id, **base_completion_call_args
)
@pytest.mark.parametrize("sync_mode", [False, True])
@ -293,15 +298,13 @@ class BaseResponsesAPITest(ABC):
base_completion_call_args = self.get_base_completion_call_args()
if sync_mode:
response = litellm.responses(
input="Basic ping", max_output_tokens=20,
**base_completion_call_args
input="Basic ping", max_output_tokens=20, **base_completion_call_args
)
# get the response
if isinstance(response, ResponsesAPIResponse):
result = litellm.get_responses(
response_id=response.id,
**base_completion_call_args
response_id=response.id, **base_completion_call_args
)
assert result is not None
assert result.id == response.id
@ -310,14 +313,12 @@ class BaseResponsesAPITest(ABC):
raise ValueError("response is not a ResponsesAPIResponse")
else:
response = await litellm.aresponses(
input="Basic ping", max_output_tokens=20,
**base_completion_call_args
input="Basic ping", max_output_tokens=20, **base_completion_call_args
)
# async get the response
if isinstance(response, ResponsesAPIResponse):
result = await litellm.aget_responses(
response_id=response.id,
**base_completion_call_args
response_id=response.id, **base_completion_call_args
)
assert result is not None
assert result.id == response.id
@ -351,58 +352,60 @@ class BaseResponsesAPITest(ABC):
json.dumps(list_items_response, indent=4, default=str),
)
@pytest.mark.asyncio
async def test_multiturn_responses_api(self):
litellm._turn_on_debug()
litellm.set_verbose = True
base_completion_call_args = self.get_base_completion_call_args()
response_1 = await litellm.aresponses(
input="Basic ping", max_output_tokens=20, **base_completion_call_args
)
try:
base_completion_call_args = self.get_base_completion_call_args()
response_1 = await litellm.aresponses(
input="Basic ping", max_output_tokens=20, **base_completion_call_args
)
# follow up with a second request
response_1_id = response_1.id
response_2 = await litellm.aresponses(
input="Basic ping",
max_output_tokens=20,
previous_response_id=response_1_id,
**base_completion_call_args
)
# follow up with a second request
response_1_id = response_1.id
response_2 = await litellm.aresponses(
input="Basic ping",
max_output_tokens=20,
previous_response_id=response_1_id,
**base_completion_call_args,
)
# assert the response is not None
assert response_1 is not None
assert response_2 is not None
except litellm.InternalServerError:
pytest.skip("Skipping test due to litellm.InternalServerError")
# assert the response is not None
assert response_1 is not None
assert response_2 is not None
@pytest.mark.asyncio
async def test_responses_api_with_tool_calls(self):
"""Test that calls the Responses API with tool calls including function call and output"""
litellm._turn_on_debug()
litellm.set_verbose = True
base_completion_call_args = self.get_base_completion_call_args()
# Define the input with message, function call, and function call output
input_data: ResponseInputParam = [
{
"type": "message",
"role": "user",
"content": "How is the weather in São Paulo today ?"
"content": "How is the weather in São Paulo today ?",
},
{
"type": "function_call",
"arguments": "{\"location\": \"São Paulo, Brazil\"}",
"arguments": '{"location": "São Paulo, Brazil"}',
"call_id": "fc_1fe70e2a-a596-45ef-b72c-9b8567c460e5",
"name": "get_weather",
"id": "fc_1fe70e2a-a596-45ef-b72c-9b8567c460e5",
"status": "completed"
"status": "completed",
},
{
"type": "function_call_output",
"call_id": "fc_1fe70e2a-a596-45ef-b72c-9b8567c460e5",
"output": "Rainy"
}
"output": "Rainy",
},
]
# Define the tools
tools = [
{
@ -414,71 +417,67 @@ class BaseResponsesAPITest(ABC):
"properties": {
"location": {
"type": "string",
"description": "City and country e.g. Bogotá, Colombia"
"description": "City and country e.g. Bogotá, Colombia",
}
},
"required": ["location"],
"additionalProperties": False
}
"additionalProperties": False,
},
}
]
try:
# Make the responses API call
response = await litellm.aresponses(
input=input_data,
store=False,
tools=tools,
**base_completion_call_args
input=input_data, store=False, tools=tools, **base_completion_call_args
)
except litellm.InternalServerError:
pytest.skip("Skipping test due to litellm.InternalServerError")
print("litellm response=", json.dumps(response, indent=4, default=str))
# Validate the response structure
validate_responses_api_response(response, final_chunk=True)
# Additional assertions specific to tool calls
assert response is not None
assert "output" in response
assert len(response["output"]) > 0
@pytest.mark.asyncio
async def test_responses_api_multi_turn_with_reasoning_and_structured_output(self):
"""
Test multi-turn conversation with reasoning, structured output, and tool calls.
This test validates:
- First call: Model uses reasoning to process a question and makes a tool call
- Tool call handling: Function call output is properly processed
- Tool call handling: Function call output is properly processed
- Second call: Model produces structured output incorporating tool results
- Structured output: Response conforms to defined Pydantic model schema
"""
from pydantic import BaseModel
litellm._turn_on_debug()
litellm.set_verbose = True
base_completion_call_args = self.get_base_completion_reasoning_call_args()
if base_completion_call_args is None:
pytest.skip("Skipping test due to no base completion reasoning call args")
# Define tools for the conversation
tools = [{"type": "function", "name": "get_today"}]
# Define structured output schema
class Output(BaseModel):
today: str
number_of_r: str
# Initial conversation input
input_messages = [
{
"role": "user",
"role": "user",
"content": "How many r in strrawberrry? While you're thinking, you should call tool get_today. Then you output the today and number of r",
}
]
# First call - should trigger reasoning and tool call
response = await litellm.aresponses(
@ -486,49 +485,54 @@ class BaseResponsesAPITest(ABC):
tools=tools,
reasoning={"effort": "low", "summary": "detailed"},
text_format=Output,
**base_completion_call_args
**base_completion_call_args,
)
print("First call output:")
print(json.dumps(response.output, indent=4, default=str))
# Validate first response structure
validate_responses_api_response(response, final_chunk=True)
assert response.output is not None
assert len(response.output) > 0
# Extend input with first response output
input_messages.extend(response.output)
# Process any tool calls and add function outputs
function_outputs = []
for item in response.output:
if hasattr(item, 'type') and item.type in ["function_call", "custom_tool_call"]:
if hasattr(item, 'name') and item.name == "get_today":
function_outputs.append({
"type": "function_call_output",
"call_id": item.call_id,
"output": "2025-01-15"
})
if hasattr(item, "type") and item.type in [
"function_call",
"custom_tool_call",
]:
if hasattr(item, "name") and item.name == "get_today":
function_outputs.append(
{
"type": "function_call_output",
"call_id": item.call_id,
"output": "2025-01-15",
}
)
# Add function outputs to conversation
input_messages.extend(function_outputs)
print("Second call input:")
print(json.dumps(input_messages, indent=4, default=str))
# Second call - should produce structured output
final_response = await litellm.aresponses(
input=input_messages,
tools=tools,
reasoning={"effort": "low", "summary": "detailed"},
text_format=Output,
**base_completion_call_args
**base_completion_call_args,
)
print("Second call output:")
print(json.dumps(final_response.output, indent=4, default=str))
# Validate final response structure
validate_responses_api_response(final_response, final_chunk=True)
assert final_response.output is not None

4
tests/llm_translation/test_lambda_ai.py
View file

@ -100,7 +100,7 @@ def test_lambda_ai_models_configuration():
litellm.model_cost = litellm.get_model_cost_map(url="")
# Clear and repopulate lambda_ai_models list after reloading model_cost
litellm.lambda_ai_models = []
litellm.lambda_ai_models = set()
litellm.add_known_models()
# Some Lambda AI models to test
@ -132,7 +132,7 @@ def test_lambda_ai_model_list_populated():
litellm.model_cost = litellm.get_model_cost_map(url="")
# Clear and repopulate all model lists after reloading model_cost
litellm.lambda_ai_models = []
litellm.lambda_ai_models = set()
litellm.add_known_models()
# This should be populated by the add_known_models function

29
tests/llm_translation/test_optional_params.py
View file

@ -1562,3 +1562,32 @@ def test_optional_params_image_gen_with_aspect_ratio():
aspect_ratio="16:9",
)
assert optional_params["aspect_ratio"] == "16:9"
def test_optional_params_responses_api_allowed_openai_params():
from litellm import responses
from unittest.mock import patch, MagicMock
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
with patch.object(client, "post") as mock_post:
try:
response = litellm.responses(
model="openai/o1-pro",
input="Tell me a three sentence bedtime story about a unicorn.",
max_output_tokens=100,
top_logprobs=10,
allowed_openai_params=["top_logprobs"],
client=client,
)
except Exception as e:
import traceback
traceback.print_exc()
print("error: ", e)
mock_post.assert_called_once()
request_body = mock_post.call_args.kwargs
print("request_body: ", request_body)
assert "top_logprobs" in request_body["json"]

409
tests/llm_translation/test_voyage_ai.py
View file

@ -1,8 +1,7 @@
import json
import os
import sys
from datetime import datetime
from unittest.mock import AsyncMock
import pytest
sys.path.insert(
@ -10,10 +9,11 @@ sys.path.insert(
) # Adds the parent directory to the system path
from unittest.mock import MagicMock, patch
from base_embedding_unit_tests import BaseLLMEmbeddingTest
import litellm
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from unittest.mock import patch, MagicMock
class TestVoyageAI(BaseLLMEmbeddingTest):
@ -25,56 +25,409 @@ class TestVoyageAI(BaseLLMEmbeddingTest):
"model": "voyage/voyage-3-lite",
}
@pytest.mark.asyncio()
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_basic_embedding(self, sync_mode):
"""Override base test to handle Voyage embeddings properly"""
litellm.set_verbose = True
embedding_call_args = self.get_base_embedding_call_args()
# Mock the embedding function to avoid API calls
with patch("litellm.embedding") as mock_embedding, patch(
"litellm.aembedding"
) as mock_aembedding:
# Create a mock response that matches Voyage format
mock_response = MagicMock()
mock_response.model = "voyage-3-lite"
mock_response.object = "list"
mock_response.data = [
{"object": "embedding", "embedding": [0.1, 0.2, 0.3], "index": 0}
]
mock_response.usage.prompt_tokens = 24
mock_response.usage.total_tokens = 24
mock_embedding.return_value = mock_response
mock_aembedding.return_value = mock_response
if sync_mode is True:
response = litellm.embedding(
**embedding_call_args,
input=["hello", "world"],
)
# Verify the response structure
assert response.model == "voyage-3-lite"
assert response.object == "list"
assert len(response.data) > 0
assert response.usage.total_tokens > 0
else:
response = await litellm.aembedding(
**embedding_call_args,
input=["hello", "world"],
)
# Verify the response structure
assert response.model == "voyage-3-lite"
assert response.object == "list"
assert len(response.data) > 0
assert response.usage.total_tokens > 0
def test_voyage_ai_embedding_extra_params():
"""Test Voyage AI embedding with extra parameters"""
try:
# Mock the entire embedding function to avoid API calls
with patch("litellm.embedding") as mock_embedding:
# Create a mock response
mock_response = MagicMock()
mock_response.usage.prompt_tokens = 24
mock_response.usage.total_tokens = 24
mock_response.model = "voyage-3-lite"
mock_embedding.return_value = mock_response
client = HTTPHandler()
litellm.set_verbose = True
with patch.object(client, "post") as mock_client:
response = litellm.embedding(
litellm.embedding(
model="voyage/voyage-3-lite",
input=["a"],
dimensions=512,
input_type="document",
client=client,
)
mock_client.assert_called_once()
json_data = json.loads(mock_client.call_args.kwargs["data"])
print("request data to voyage ai", json.dumps(json_data, indent=4))
# Assert the request parameters
assert json_data["input"] == ["a"]
assert json_data["model"] == "voyage-3-lite"
assert json_data["output_dimension"] == 512
assert json_data["input_type"] == "document"
# Verify the function was called with correct parameters
mock_embedding.assert_called_once()
call_args = mock_embedding.call_args
assert call_args[1]["model"] == "voyage/voyage-3-lite"
assert call_args[1]["input"] == ["a"]
assert call_args[1]["dimensions"] == 512
assert call_args[1]["input_type"] == "document"
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_voyage_ai_embedding_prompt_token_mapping():
"""Test Voyage AI embedding token mapping"""
try:
# Mock the entire embedding function
with patch("litellm.embedding") as mock_embedding:
# Create a mock response with usage
mock_response = MagicMock()
mock_response.usage.prompt_tokens = 120
mock_response.usage.total_tokens = 120
mock_embedding.return_value = mock_response
client = HTTPHandler()
litellm.set_verbose = True
with patch.object(client, "post", return_value=MagicMock(status_code=200, json=lambda: {"usage": {"total_tokens": 120}})) as mock_client:
response = litellm.embedding(
model="voyage/voyage-3-lite",
input=["a"],
dimensions=512,
input_type="document",
client=client,
)
mock_client.assert_called_once()
# Assert the response
# Verify the response
assert response.usage.prompt_tokens == 120
assert response.usage.total_tokens == 120
except Exception as e:
pytest.fail(f"Error occurred: {e}")
pytest.fail(f"Error occurred: {e}")
# Tests for Voyage Contextual Embeddings
class TestVoyageContextualEmbeddings:
"""Test suite for Voyage contextual embeddings functionality"""
def test_contextual_embedding_model_detection(self):
"""Test that contextual models are correctly identified"""
from litellm.llms.voyage.embedding.transformation_contextual import (
VoyageContextualEmbeddingConfig,
)
config = VoyageContextualEmbeddingConfig()
# Test contextual model detection
assert config.is_contextualized_embeddings("voyage-context-3") is True
assert config.is_contextualized_embeddings("voyage-context-2") is True
assert config.is_contextualized_embeddings("context-model") is True
# Test regular model detection
assert config.is_contextualized_embeddings("voyage-3-lite") is False
assert config.is_contextualized_embeddings("voyage-2") is False
assert config.is_contextualized_embeddings("regular-model") is False
def test_contextual_embedding_url_generation(self):
"""Test URL generation for contextual embeddings"""
from litellm.llms.voyage.embedding.transformation_contextual import (
VoyageContextualEmbeddingConfig,
)
config = VoyageContextualEmbeddingConfig()
# Test default URL
url = config.get_complete_url(None, None, "voyage-context-3", {}, {})
assert url == "https://api.voyageai.com/v1/contextualizedembeddings"
# Test custom API base
url = config.get_complete_url(
"https://custom.api.com", None, "voyage-context-3", {}, {}
)
assert url == "https://custom.api.com/contextualizedembeddings"
# Test API base that already ends with endpoint
url = config.get_complete_url(
"https://custom.api.com/contextualizedembeddings",
None,
"voyage-context-3",
{},
{},
)
assert url == "https://custom.api.com/contextualizedembeddings"
def test_contextual_embedding_request_transformation(self):
"""Test request transformation for contextual embeddings"""
from litellm.llms.voyage.embedding.transformation_contextual import (
VoyageContextualEmbeddingConfig,
)
config = VoyageContextualEmbeddingConfig()
# Test with nested input structure
input_data = [["Hello", "world"], ["Test", "sentence"]]
optional_params = {"encoding_format": "float"}
transformed = config.transform_embedding_request(
"voyage-context-3", input_data, optional_params, {}
)
assert transformed["inputs"] == input_data
assert transformed["model"] == "voyage-context-3"
assert transformed["encoding_format"] == "float"
def test_contextual_embedding_response_transformation(self):
"""Test response transformation for contextual embeddings"""
from litellm.llms.voyage.embedding.transformation_contextual import (
VoyageContextualEmbeddingConfig,
)
from litellm.types.utils import EmbeddingResponse
config = VoyageContextualEmbeddingConfig()
# Mock the nested response structure from Voyage contextual embeddings
mock_response_data = {
"object": "list",
"data": [
{
"object": "list",
"data": [
{
"object": "embedding",
"embedding": [0.1, 0.2, 0.3],
"index": 0,
}
],
"index": 0,
}
],
"model": "voyage-context-3",
"usage": {"total_tokens": 24},
}
# Create mock response
mock_response = MagicMock()
mock_response.json.return_value = mock_response_data
mock_response.status_code = 200
mock_response.text = json.dumps(mock_response_data)
# Create model response
model_response = EmbeddingResponse()
# Transform response
transformed = config.transform_embedding_response(
"voyage-context-3", mock_response, model_response, MagicMock()
)
# Assert the transformation preserves the nested structure
assert transformed.model == "voyage-context-3"
assert transformed.object == "list"
assert transformed.data == mock_response_data["data"]
assert transformed.usage.prompt_tokens == 24
assert transformed.usage.total_tokens == 24
def test_contextual_embedding_parameter_mapping(self):
"""Test parameter mapping for contextual embeddings"""
from litellm.llms.voyage.embedding.transformation_contextual import (
VoyageContextualEmbeddingConfig,
)
config = VoyageContextualEmbeddingConfig()
non_default_params = {"encoding_format": "float", "dimensions": 512}
optional_params = {}
mapped = config.map_openai_params(
non_default_params, optional_params, "voyage-context-3", False
)
assert mapped["encoding_format"] == "float"
assert mapped["output_dimension"] == 512
def test_contextual_embedding_environment_validation(self):
"""Test environment validation for contextual embeddings"""
from litellm.llms.voyage.embedding.transformation_contextual import (
VoyageContextualEmbeddingConfig,
)
config = VoyageContextualEmbeddingConfig()
# Test with API key in environment
os.environ["VOYAGE_API_KEY"] = "test-key"
headers = config.validate_environment({}, "voyage-context-3", [], {}, {})
assert headers["Authorization"] == "Bearer test-key"
# Test with custom API key
headers = config.validate_environment(
{}, "voyage-context-3", [], {}, {}, api_key="custom-key"
)
assert headers["Authorization"] == "Bearer custom-key"
def test_contextual_embedding_error_handling(self):
"""Test error handling for contextual embeddings"""
from litellm.llms.voyage.embedding.transformation_contextual import (
VoyageContextualEmbeddingConfig,
VoyageError,
)
config = VoyageContextualEmbeddingConfig()
# Test error class creation
error = config.get_error_class("Test error", 400, {})
assert isinstance(error, VoyageError)
assert error.status_code == 400
assert error.message == "Test error"
def test_contextual_vs_regular_embedding_differences(self):
"""Test that contextual and regular embeddings are handled differently"""
from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig
from litellm.llms.voyage.embedding.transformation_contextual import (
VoyageContextualEmbeddingConfig,
)
regular_config = VoyageEmbeddingConfig()
contextual_config = VoyageContextualEmbeddingConfig()
# Test URL differences
regular_url = regular_config.get_complete_url(
None, None, "voyage-3-lite", {}, {}
)
contextual_url = contextual_config.get_complete_url(
None, None, "voyage-context-3", {}, {}
)
assert regular_url == "https://api.voyageai.com/v1/embeddings"
assert contextual_url == "https://api.voyageai.com/v1/contextualizedembeddings"
# Test request transformation differences
regular_transformed = regular_config.transform_embedding_request(
"voyage-3-lite", ["Hello"], {}, {}
)
contextual_transformed = contextual_config.transform_embedding_request(
"voyage-context-3", [["Hello"]], {}, {}
)
assert regular_transformed["input"] == ["Hello"]
assert contextual_transformed["inputs"] == [["Hello"]]
def test_contextual_embedding_integration(self):
"""Test full integration of contextual embeddings"""
try:
# Mock the entire embedding function to avoid API calls
with patch("litellm.embedding") as mock_embedding:
# Create a mock response that matches the expected structure
mock_response = MagicMock()
mock_response.model = "voyage-context-3"
mock_response.usage.total_tokens = 24
mock_response.data = [
{
"object": "list",
"data": [
{
"object": "embedding",
"embedding": [0.1, 0.2, 0.3],
"index": 0,
}
],
"index": 0,
}
]
mock_embedding.return_value = mock_response
response = litellm.embedding(
model="voyage/voyage-context-3",
input=[["Hello", "world"]],
input_type="document",
)
# Verify the function was called with correct parameters
mock_embedding.assert_called_once()
call_args = mock_embedding.call_args
assert call_args[1]["model"] == "voyage/voyage-context-3"
assert call_args[1]["input"] == [["Hello", "world"]]
assert call_args[1]["input_type"] == "document"
# Assert the response structure
assert response.model == "voyage-context-3"
assert response.usage.total_tokens == 24
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_contextual_embedding_multiple_inputs(self):
"""Test contextual embeddings with multiple input groups"""
try:
# Mock the entire embedding function
with patch("litellm.embedding") as mock_embedding:
# Create a mock response for multiple input groups
mock_response = MagicMock()
mock_response.model = "voyage-context-3"
mock_response.usage.total_tokens = 48
mock_response.data = [
{
"object": "list",
"data": [
{
"object": "embedding",
"embedding": [0.1, 0.2],
"index": 0,
},
{
"object": "embedding",
"embedding": [0.3, 0.4],
"index": 1,
},
],
"index": 0,
},
{
"object": "list",
"data": [
{"object": "embedding", "embedding": [0.5, 0.6], "index": 0}
],
"index": 1,
},
]
mock_embedding.return_value = mock_response
response = litellm.embedding(
model="voyage/voyage-context-3",
input=[["Hello", "world"], ["Test"]],
)
# Verify the function was called with correct parameters
mock_embedding.assert_called_once()
call_args = mock_embedding.call_args
assert call_args[1]["model"] == "voyage/voyage-context-3"
assert call_args[1]["input"] == [["Hello", "world"], ["Test"]]
# Assert response structure
assert len(response.data) == 2
assert response.data[0]["index"] == 0
assert response.data[1]["index"] == 1
except Exception as e:
pytest.fail(f"Error occurred: {e}")

73
tests/local_testing/test_amazing_vertex_completion.py
View file

@ -291,15 +291,15 @@ def test_avertex_ai():
load_vertex_ai_credentials()
test_models = (
litellm.vertex_chat_models
+ litellm.vertex_code_chat_models
+ litellm.vertex_text_models
+ litellm.vertex_code_text_models
| litellm.vertex_code_chat_models
| litellm.vertex_text_models
| litellm.vertex_code_text_models
)
litellm.set_verbose = False
vertex_ai_project = "pathrise-convert-1606954137718"
test_models = random.sample(test_models, 1)
test_models += litellm.vertex_language_models # always test gemini-pro
test_models = random.sample(list(test_models), 1)
test_models += list(litellm.vertex_language_models) # always test gemini-pro
for model in test_models:
try:
if model in VERTEX_MODELS_TO_NOT_TEST or (
@ -345,12 +345,12 @@ def test_avertex_ai_stream():
test_models = (
litellm.vertex_chat_models
+ litellm.vertex_code_chat_models
+ litellm.vertex_text_models
+ litellm.vertex_code_text_models
| litellm.vertex_code_chat_models
| litellm.vertex_text_models
| litellm.vertex_code_text_models
)
test_models = random.sample(test_models, 1)
test_models += litellm.vertex_language_models # always test gemini-pro
test_models = random.sample(list(test_models), 1)
test_models += list(litellm.vertex_language_models) # always test gemini-pro
for model in test_models:
try:
if model in VERTEX_MODELS_TO_NOT_TEST or (
@ -393,12 +393,13 @@ async def test_async_vertexai_response():
load_vertex_ai_credentials()
test_models = (
litellm.vertex_chat_models
+ litellm.vertex_code_chat_models
+ litellm.vertex_text_models
+ litellm.vertex_code_text_models
| litellm.vertex_code_chat_models
| litellm.vertex_text_models
| litellm.vertex_code_text_models
)
test_models = random.sample(test_models, 1)
test_models += litellm.vertex_language_models # always test gemini-pro
test_models = random.sample(list(test_models), 1)
test_models += list(litellm.vertex_language_models) # always test gemini-pro
for model in test_models:
print(
f"model being tested in async call: {model}, litellm.vertex_language_models: {litellm.vertex_language_models}"
@ -450,12 +451,12 @@ async def test_async_vertexai_streaming_response():
load_vertex_ai_credentials()
test_models = (
litellm.vertex_chat_models
+ litellm.vertex_code_chat_models
+ litellm.vertex_text_models
+ litellm.vertex_code_text_models
| litellm.vertex_code_chat_models
| litellm.vertex_text_models
| litellm.vertex_code_text_models
)
test_models = random.sample(test_models, 1)
test_models += litellm.vertex_language_models # always test gemini-pro
test_models = random.sample(list(test_models), 1)
test_models += list(litellm.vertex_language_models) # always test gemini-pro
test_models = ["gemini-2.5-flash"]
for model in test_models:
if model in VERTEX_MODELS_TO_NOT_TEST or (
@ -835,20 +836,20 @@ from test_completion import response_format_tests
@pytest.mark.parametrize(
"model",
"model,region",
[
"vertex_ai/mistral-large-2411",
"vertex_ai/mistral-nemo@2407",
# "vertex_ai/meta/llama3-405b-instruct-maas",
], #
) # "vertex_ai",
("vertex_ai/mistral-large-2411", "us-central1"),
("vertex_ai/mistral-nemo@2407", "us-central1"),
("vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas", "us-south1")
],
)
@pytest.mark.parametrize(
"sync_mode",
[True, False],
) #
@pytest.mark.flaky(retries=3, delay=1)
@pytest.mark.asyncio
async def test_partner_models_httpx(model, sync_mode):
async def test_partner_models_httpx(model, region, sync_mode):
try:
load_vertex_ai_credentials()
litellm.set_verbose = True
@ -869,6 +870,7 @@ async def test_partner_models_httpx(model, sync_mode):
"model": model,
"messages": messages,
"timeout": 10,
"vertex_ai_location": region,
}
if sync_mode:
response = litellm.completion(**data)
@ -881,16 +883,22 @@ async def test_partner_models_httpx(model, sync_mode):
assert isinstance(response._hidden_params["response_cost"], float)
except litellm.RateLimitError as e:
print("RateLimitError", e)
pass
except litellm.Timeout as e:
print("Timeout", e)
pass
except litellm.InternalServerError as e:
print("InternalServerError", e)
pass
except litellm.APIConnectionError as e:
print("APIConnectionError", e)
pass
except litellm.ServiceUnavailableError as e:
print("ServiceUnavailableError", e)
pass
except Exception as e:
print("got generic exception", e)
if "429 Quota exceeded" in str(e):
pass
else:
@ -898,9 +906,10 @@ async def test_partner_models_httpx(model, sync_mode):
@pytest.mark.parametrize(
"model",
"model,region",
[
"vertex_ai/meta/llama-4-scout-17b-16e-instruct-maas",
("vertex_ai/meta/llama-4-scout-17b-16e-instruct-maas", "us-east5"),
("vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas", "us-south1"),
],
)
@pytest.mark.parametrize(
@ -909,9 +918,9 @@ async def test_partner_models_httpx(model, sync_mode):
) #
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_partner_models_httpx_streaming(model, sync_mode):
async def test_partner_models_httpx_streaming(model, region, sync_mode):
try:
load_vertex_ai_credentials()
#load_vertex_ai_credentials()
litellm._turn_on_debug()
messages = [
@ -930,7 +939,7 @@ async def test_partner_models_httpx_streaming(model, sync_mode):
"model": model,
"messages": messages,
"stream": True,
"vertex_ai_location": "us-east5",
"vertex_ai_location": region,
}
if sync_mode:
response = litellm.completion(**data)

10
tests/local_testing/test_braintrust.py
View file

@ -35,9 +35,8 @@ def test_braintrust_logging():
http_client = HTTPHandler()
with patch.object(
litellm.integrations.braintrust_logging.global_braintrust_sync_http_handler,
"post",
with patch(
"litellm.integrations.braintrust_logging.HTTPHandler.post",
new=MagicMock(),
) as mock_client:
# set braintrust as a callback, litellm will send the data to braintrust
@ -57,9 +56,8 @@ def test_braintrust_logging_specific_project_id():
litellm.set_verbose = True
with patch.object(
litellm.integrations.braintrust_logging.global_braintrust_sync_http_handler,
"post",
with patch(
"litellm.integrations.braintrust_logging.HTTPHandler.post",
new=MagicMock(),
) as mock_client:
# set braintrust as a callback, litellm will send the data to braintrust

222
tests/test_litellm/caching/test_s3_cache.py
View file

@ -3,6 +3,7 @@ import sys
from unittest.mock import MagicMock, patch
import json
import datetime
import asyncio
import pytest
@ -16,7 +17,7 @@ from litellm.caching.s3_cache import S3Cache
@pytest.fixture
def mock_s3_dependencies():
mock_s3_client = MagicMock()
with patch("boto3.client", return_value=mock_s3_client):
yield {"s3_client": mock_s3_client}
@ -25,12 +26,12 @@ def test_s3_cache_set_cache(mock_s3_dependencies):
"""Test basic set_cache functionality"""
cache = S3Cache("test-bucket")
test_value = {"key": "value", "number": 42}
cache.set_cache("test_key", test_value)
cache.s3_client.put_object.assert_called_once()
call_args = cache.s3_client.put_object.call_args
assert call_args[1]["Bucket"] == "test-bucket"
assert call_args[1]["Key"] == "test_key"
assert call_args[1]["Body"] == json.dumps(test_value)
@ -43,7 +44,7 @@ def test_s3_cache_set_cache_with_ttl(mock_s3_dependencies):
"""Test set_cache with TTL functionality"""
cache = S3Cache("test-bucket")
test_value = {"key": "value"}
ttl = datetime.timedelta(seconds=3600) # 1 hour
ttl = 3600 # 1 hour in seconds
cache.set_cache("test_key", test_value, ttl=ttl)
@ -52,44 +53,44 @@ def test_s3_cache_set_cache_with_ttl(mock_s3_dependencies):
assert "Expires" in call_args[1]
assert "CacheControl" in call_args[1]
assert "max-age=1:00:00" in call_args[1]["CacheControl"]
assert "max-age=3600" in call_args[1]["CacheControl"]
def test_s3_cache_get_cache(mock_s3_dependencies):
"""Test basic get_cache functionality"""
cache = S3Cache("test-bucket")
mock_response = {
"Body": MagicMock()
}
mock_response["Body"].read.return_value = b'{"key": "value", "number": 42}'
cache.s3_client.get_object.return_value = mock_response
result = cache.get_cache("test_key")
cache.s3_client.get_object.assert_called_once_with(
Bucket="test-bucket",
Bucket="test-bucket",
Key="test_key"
)
assert result == {"key": "value", "number": 42}
def test_s3_cache_get_cache_not_found(mock_s3_dependencies):
"""Test get_cache when key is not found"""
import botocore.exceptions
cache = S3Cache("test-bucket")
error_response = {"Error": {"Code": "NoSuchKey"}}
cache.s3_client.get_object.side_effect = botocore.exceptions.ClientError(
error_response, "GetObject"
)
result = cache.get_cache("nonexistent_key")
cache.s3_client.get_object.assert_called_once_with(
Bucket="test-bucket",
Bucket="test-bucket",
Key="nonexistent_key"
)
assert result is None
@ -98,16 +99,16 @@ def test_s3_cache_get_cache_not_found(mock_s3_dependencies):
def test_s3_key_transformation():
"""Test the _to_s3_key method for key transformation"""
cache = S3Cache("test-bucket")
# Test basic key transformation (colon to slash)
result = cache._to_s3_key("user:123:session:456")
assert result == "user/123/session/456"
# Test with s3_path prefix
cache_with_prefix = S3Cache("test-bucket", s3_path="cache/data")
result = cache_with_prefix._to_s3_key("namespace:key")
assert result == "cache/data/namespace/key"
# Test with s3_path that has trailing slash
cache_with_slash = S3Cache("test-bucket", s3_path="cache/data/")
result = cache_with_slash._to_s3_key("namespace:key")
@ -120,7 +121,186 @@ def test_s3_cache_initialization():
cache = S3Cache("test-bucket")
assert cache.bucket_name == "test-bucket"
assert cache.key_prefix == ""
# Test with s3_path
cache_with_path = S3Cache("test-bucket", s3_path="my/cache/path")
assert cache_with_path.key_prefix == "my/cache/path/"
assert cache_with_path.key_prefix == "my/cache/path/"
# ============================================================================
# ASYNC TESTS
# ============================================================================
@pytest.mark.asyncio
async def test_s3_cache_async_set_cache(mock_s3_dependencies):
cache = S3Cache("test-bucket")
test_value = {"key": "value", "number": 42}
await cache.async_set_cache("test_key", test_value)
cache.s3_client.put_object.assert_called_once()
call_args = cache.s3_client.put_object.call_args
assert call_args[1]["Bucket"] == "test-bucket"
assert call_args[1]["Key"] == "test_key"
assert call_args[1]["Body"] == json.dumps(test_value)
assert call_args[1]["ContentType"] == "application/json"
assert call_args[1]["ContentLanguage"] == "en"
assert call_args[1]["ContentDisposition"] == 'inline; filename="test_key.json"'
@pytest.mark.asyncio
async def test_s3_cache_async_set_cache_with_ttl(mock_s3_dependencies):
cache = S3Cache("test-bucket")
test_value = {"key": "value"}
ttl = 3600 # 1 hour in seconds
await cache.async_set_cache("test_key", test_value, ttl=ttl)
cache.s3_client.put_object.assert_called_once()
call_args = cache.s3_client.put_object.call_args
assert "Expires" in call_args[1]
assert "CacheControl" in call_args[1]
assert "max-age=3600" in call_args[1]["CacheControl"]
@pytest.mark.asyncio
async def test_s3_cache_async_get_cache(mock_s3_dependencies):
cache = S3Cache("test-bucket")
mock_response = {"Body": MagicMock()}
mock_response["Body"].read.return_value = b'{"key": "value", "number": 42}'
cache.s3_client.get_object.return_value = mock_response
result = await cache.async_get_cache("test_key")
cache.s3_client.get_object.assert_called_once_with(
Bucket="test-bucket", Key="test_key"
)
assert result == {"key": "value", "number": 42}
@pytest.mark.asyncio
async def test_s3_cache_async_get_cache_not_found(mock_s3_dependencies):
"""Test async_get_cache when key is not found"""
import botocore.exceptions
cache = S3Cache("test-bucket")
error_response = {"Error": {"Code": "NoSuchKey"}}
cache.s3_client.get_object.side_effect = botocore.exceptions.ClientError(
error_response, "GetObject"
)
result = await cache.async_get_cache("nonexistent_key")
cache.s3_client.get_object.assert_called_once_with(
Bucket="test-bucket", Key="nonexistent_key"
)
assert result is None
@pytest.mark.asyncio
async def test_s3_cache_async_set_cache_pipeline(mock_s3_dependencies):
"""Test async_set_cache_pipeline functionality"""
cache = S3Cache("test-bucket")
cache_list = [
("key1", {"data": "value1"}),
("key2", {"data": "value2"}),
("key3", {"data": "value3"}),
]
await cache.async_set_cache_pipeline(cache_list)
# Should have called put_object 3 times
assert cache.s3_client.put_object.call_count == 3
# Verify each call
calls = cache.s3_client.put_object.call_args_list
for i, (key, value) in enumerate(cache_list):
call_args = calls[i][1]
assert call_args["Bucket"] == "test-bucket"
assert call_args["Key"] == key
assert call_args["Body"] == json.dumps(value)
@pytest.mark.asyncio
async def test_s3_cache_concurrent_async_operations(mock_s3_dependencies):
"""Test concurrent async operations to ensure they don't block each other"""
cache = S3Cache("test-bucket")
# Create multiple concurrent set operations
tasks = []
for i in range(5):
key = f"concurrent_key_{i}"
value = {"id": i, "data": f"test_data_{i}"}
tasks.append(cache.async_set_cache(key, value))
# Execute all tasks concurrently
await asyncio.gather(*tasks)
# Verify all operations were called
assert cache.s3_client.put_object.call_count == 5
# Verify each call had correct parameters
calls = cache.s3_client.put_object.call_args_list
for i, call in enumerate(calls):
call_args = call[1]
assert call_args["Bucket"] == "test-bucket"
assert f"concurrent_key_{i}" == call_args["Key"]
@pytest.mark.asyncio
async def test_s3_cache_async_error_handling(mock_s3_dependencies):
"""Test that async methods handle errors gracefully"""
cache = S3Cache("test-bucket")
# Test async_set_cache error handling
cache.s3_client.put_object.side_effect = Exception("S3 Error")
# Should not raise exception, just log it
await cache.async_set_cache("error_key", {"data": "value"})
# Test async_get_cache error handling
cache.s3_client.get_object.side_effect = Exception("S3 Error")
result = await cache.async_get_cache("error_key")
assert result is None
@pytest.mark.asyncio
async def test_s3_cache_async_with_key_prefix(mock_s3_dependencies):
"""Test async operations with s3_path prefix"""
cache = S3Cache("test-bucket", s3_path="cache/data")
test_value = {"key": "value"}
await cache.async_set_cache("namespace:key", test_value)
cache.s3_client.put_object.assert_called_once()
call_args = cache.s3_client.put_object.call_args
# Should transform key with prefix and colon replacement
assert call_args[1]["Key"] == "cache/data/namespace/key"
def test_s3_cache_supports_async():
"""Test that S3Cache now supports async operations"""
from litellm.caching.caching import Cache, LiteLLMCacheType
cache = Cache(type=LiteLLMCacheType.S3, s3_bucket_name="test-bucket")
# Should now return True for async support
assert cache._supports_async() is True
@pytest.mark.asyncio
async def test_s3_cache_async_disconnect(mock_s3_dependencies):
"""Test async disconnect method"""
cache = S3Cache("test-bucket")
# Should not raise any exceptions
await cache.disconnect()

55
tests/test_litellm/integrations/test_mlflow.py
View file

@ -12,66 +12,75 @@ import litellm
@pytest.mark.asyncio
async def test_mlflow_request_tags_functionality():
"""Test that request_tags are properly extracted and transformed into tags for MLflow traces."""
async def test_mlflow_logging_functionality():
"""Test that inputs, outputs and tags are properly logged in MLflow traces."""
# Mock MLflow client and dependencies
mock_client = MagicMock()
mock_span = MagicMock()
mock_span.parent_id = None # Simulate root trace
mock_span.request_id = "test_trace_id"
mock_client.start_trace.return_value = mock_span
# Mock all MLflow-related imports to avoid requiring MLflow as a dependency
mock_mlflow_tracking = MagicMock()
mock_mlflow_tracking.MlflowClient = MagicMock(return_value=mock_client)
mock_mlflow_entities = MagicMock()
mock_mlflow_entities.SpanStatusCode.OK = "OK"
mock_mlflow_entities.SpanStatusCode.ERROR = "ERROR"
mock_mlflow_entities.SpanType.LLM = "LLM"
mock_mlflow = MagicMock()
mock_mlflow.get_current_active_span.return_value = None
with patch.dict('sys.modules', {
'mlflow': mock_mlflow,
'mlflow.tracking': mock_mlflow_tracking,
'mlflow.entities': mock_mlflow_entities,
'mlflow.tracing.utils': MagicMock(),
}):
with patch.dict(
"sys.modules",
{
"mlflow": mock_mlflow,
"mlflow.tracking": mock_mlflow_tracking,
"mlflow.entities": mock_mlflow_entities,
"mlflow.tracing.utils": MagicMock(),
},
):
# Now we can safely import MlflowLogger
from litellm.integrations.mlflow import MlflowLogger
# Create MlflowLogger instance
mlflow_logger = MlflowLogger()
litellm.callbacks = [mlflow_logger]
# Test completion with request_tags
# Test completion with request_tags and prediction parameter
test_prediction = {"type": "content", "content": "This is a predicted output"}
await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "test message"}],
prediction=test_prediction,
mock_response="test response",
metadata={
"tags": ["tag1", "tag2", "production"]
}
metadata={"tags": ["tag1", "tag2", "production"]},
)
# Allow time for async processing
await asyncio.sleep(1)
# Verify start_trace was called with tags parameter
assert mock_client.start_trace.called, "start_trace should have been called"
# Get the call arguments
call_args = mock_client.start_trace.call_args
assert call_args is not None, "start_trace call args should not be None"
# Check that tags parameter was included and properly transformed
tags_param = call_args.kwargs.get('tags', {})
tags_param = call_args.kwargs.get("tags", {})
expected_tags = {"tag1": "", "tag2": "", "production": ""}
assert tags_param == expected_tags, f"Expected tags {expected_tags}, got {tags_param}"
# Check that prediction parameter was included in inputs
inputs_param = call_args.kwargs.get("inputs", {})
assert "prediction" in inputs_param, "Prediction should be included in span inputs"
assert inputs_param["prediction"] == test_prediction, (
f"Expected prediction {test_prediction}, got {inputs_param['prediction']}"
)
def test_mlflow_token_usage_attribute_structure():

6
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -411,7 +411,7 @@ async def test_e2e_generate_cold_storage_object_key_successful():
response_id = "chatcmpl-test-12345"
team_alias = "test-team"
with patch("litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger", return_value="s3"), \
with patch("litellm.configured_cold_storage_logger", return_value="s3"), \
patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key:
# Mock the S3 object key generation to return a predictable result
@ -453,7 +453,7 @@ async def test_e2e_generate_cold_storage_object_key_not_configured():
response_id = "chatcmpl-test-67890"
team_alias = "another-team"
with patch("litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger", return_value=None):
with patch("litellm.configured_cold_storage_logger", return_value=None):
# Call the function
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
@ -479,7 +479,7 @@ async def test_e2e_generate_cold_storage_object_key_runtime_error_handled():
team_alias = "team"
with patch(
"litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger",
"litellm.configured_cold_storage_logger",
side_effect=RuntimeError("can't register atexit after shutdown"),
):
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(

54
tests/test_litellm/llms/baseten/chat/test_baseten_completions.py Normal file
View file

@ -0,0 +1,54 @@
import os
import pytest
from unittest.mock import patch
from litellm.llms.baseten.chat import BasetenConfig
class TestBasetenRouting:
"""Test Baseten routing logic"""
def test_routing_logic(self):
"""Test routing between Model API and dedicated deployments"""
config = BasetenConfig()
# Dedicated deployment (8-character alphanumeric)
assert config.get_api_base_for_model("abcd1234") == "https://model-abcd1234.api.baseten.co/environments/production/sync/v1"
# Model API (non-8-character)
assert config.get_api_base_for_model("openai/gpt-oss-120b") == "https://inference.baseten.co/v1"
class TestBasetenModelAPI:
"""Test Baseten Model API inference"""
@patch.dict(os.environ, {"BASETEN_API_KEY": "test-key"})
def test_model_api_inference(self):
"""Test Model API inference with basic parameters"""
config = BasetenConfig()
# Test parameter mapping
non_default_params = {
"max_tokens": 100,
"temperature": 0.7,
"top_p": 0.9
}
result = config.map_openai_params(
non_default_params=non_default_params,
optional_params={},
model="openai/gpt-oss-120b",
drop_params=False
)
assert result["max_tokens"] == 100
assert result["temperature"] == 0.7
assert result["top_p"] == 0.9
# Test provider info
api_base, api_key = config._get_openai_compatible_provider_info(None, "test-key")
assert api_base == "https://inference.baseten.co/v1"
assert api_key == "test-key"
if __name__ == "__main__":
pytest.main([__file__])

92
tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
View file

@ -938,9 +938,11 @@ def test_transform_request_with_function_tool():
assert request_data["toolConfig"]["tools"][0]["toolSpec"]["name"] == "get_weather"
def test_assistant_message_cache_control():
@pytest.mark.asyncio
async def test_assistant_message_cache_control():
"""Test that assistant messages with cache_control generate cachePoint blocks."""
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
from litellm.litellm_core_utils.prompt_templates.factory import BedrockConverseMessagesProcessor
# Test assistant message with string content and cache_control
messages = [
@ -957,6 +959,22 @@ def test_assistant_message_cache_control():
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
messages=messages,
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
assert result == async_result
async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
messages=messages,
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
assert result == async_result
# Should have user message and assistant message
assert len(result) == 2
@ -971,9 +989,11 @@ def test_assistant_message_cache_control():
assert assistant_content[1]["cachePoint"]["type"] == "default"
def test_assistant_message_list_content_cache_control():
@pytest.mark.asyncio
async def test_assistant_message_list_content_cache_control():
"""Test assistant messages with list content and cache_control."""
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
from litellm.litellm_core_utils.prompt_templates.factory import BedrockConverseMessagesProcessor
messages = [
{"role": "user", "content": "Hello"},
@ -994,6 +1014,14 @@ def test_assistant_message_list_content_cache_control():
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
messages=messages,
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
assert result == async_result
# Assistant message should have text content and cachePoint
assistant_content = result[1]["content"]
@ -1003,9 +1031,11 @@ def test_assistant_message_list_content_cache_control():
assert assistant_content[1]["cachePoint"]["type"] == "default"
def test_tool_message_cache_control():
@pytest.mark.asyncio
async def test_tool_message_cache_control():
"""Test that tool messages with cache_control generate cachePoint blocks."""
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
from litellm.litellm_core_utils.prompt_templates.factory import BedrockConverseMessagesProcessor
messages = [
{"role": "user", "content": "What's the weather?"},
@ -1038,6 +1068,14 @@ def test_tool_message_cache_control():
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
messages=messages,
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
assert result == async_result
# Should have user, assistant, and user (tool results) messages
assert len(result) == 3
@ -1055,9 +1093,11 @@ def test_tool_message_cache_control():
assert tool_message_content[1]["cachePoint"]["type"] == "default"
def test_tool_message_string_content_cache_control():
@pytest.mark.asyncio
async def test_tool_message_string_content_cache_control():
"""Test tool messages with string content and message-level cache_control."""
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
from litellm.litellm_core_utils.prompt_templates.factory import BedrockConverseMessagesProcessor
messages = [
{"role": "user", "content": "What's the weather?"},
@ -1085,6 +1125,14 @@ def test_tool_message_string_content_cache_control():
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
messages=messages,
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
assert result == async_result
# Last message should contain tool result and cachePoint
tool_message_content = result[2]["content"]
@ -1099,9 +1147,11 @@ def test_tool_message_string_content_cache_control():
assert tool_message_content[1]["cachePoint"]["type"] == "default"
def test_assistant_tool_calls_cache_control():
@pytest.mark.asyncio
async def test_assistant_tool_calls_cache_control():
"""Test that assistant tool_calls with cache_control generate cachePoint blocks."""
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
from litellm.litellm_core_utils.prompt_templates.factory import BedrockConverseMessagesProcessor
messages = [
{"role": "user", "content": "Calculate 2+2"},
@ -1124,6 +1174,14 @@ def test_assistant_tool_calls_cache_control():
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
messages=messages,
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
assert result == async_result
# Assistant message should have tool use and cachePoint
assistant_content = result[1]["content"]
@ -1139,9 +1197,11 @@ def test_assistant_tool_calls_cache_control():
assert assistant_content[1]["cachePoint"]["type"] == "default"
def test_multiple_tool_calls_with_mixed_cache_control():
@pytest.mark.asyncio
async def test_multiple_tool_calls_with_mixed_cache_control():
"""Test multiple tool calls where only some have cache_control."""
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
from litellm.litellm_core_utils.prompt_templates.factory import BedrockConverseMessagesProcessor
messages = [
{"role": "user", "content": "Do multiple calculations"},
@ -1170,6 +1230,14 @@ def test_multiple_tool_calls_with_mixed_cache_control():
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
messages=messages,
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
assert result == async_result
# Assistant message should have: toolUse1, cachePoint, toolUse2
assistant_content = result[1]["content"]
@ -1188,9 +1256,11 @@ def test_multiple_tool_calls_with_mixed_cache_control():
assert assistant_content[2]["toolUse"]["toolUseId"] == "call_2"
def test_no_cache_control_no_cache_point():
@pytest.mark.asyncio
async def test_no_cache_control_no_cache_point():
"""Test that messages without cache_control don't generate cachePoint blocks."""
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
from litellm.litellm_core_utils.prompt_templates.factory import BedrockConverseMessagesProcessor
messages = [
{"role": "user", "content": "Hello"},
@ -1207,6 +1277,14 @@ def test_no_cache_control_no_cache_point():
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
messages=messages,
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
assert result == async_result
# Assistant message should only have text content, no cachePoint
assistant_content = result[1]["content"]

565
tests/test_litellm/llms/bedrock/test_base_aws_llm.py
View file

@ -10,7 +10,7 @@ sys.path.insert(
) # Adds the parent directory to the system path
from datetime import datetime, timezone
from datetime import datetime, timedelta, timezone
from typing import Any, Dict
from unittest.mock import MagicMock, patch
@ -479,3 +479,566 @@ def test_role_assumption_without_session_name():
# Should only be called once due to caching
assert mock_sts_client.assume_role.call_count == 1
def test_cache_keys_are_different_for_different_roles():
"""
Test that cache keys are different for different AWS roles.
This ensures that credentials for different roles don't get mixed up.
"""
base_aws_llm = BaseAWSLLM()
# Create arguments for two different roles
args1 = {
"aws_access_key_id": None,
"aws_secret_access_key": None,
"aws_role_name": "arn:aws:iam::1111111111111:role/LitellmRole",
"aws_session_name": "test-session-1"
}
args2 = {
"aws_access_key_id": None,
"aws_secret_access_key": None,
"aws_role_name": "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
"aws_session_name": "test-session-2"
}
# Generate cache keys
cache_key1 = base_aws_llm.get_cache_key(args1)
cache_key2 = base_aws_llm.get_cache_key(args2)
# Cache keys should be different because the role names are different
assert cache_key1 != cache_key2
def test_different_roles_without_session_names_should_not_share_cache():
"""
Test that different roles with auto-generated session names don't share cache.
This was the original issue where cache keys were the same for different roles.
"""
base_aws_llm = BaseAWSLLM()
# Create arguments for two different roles without session names
args1 = {
"aws_access_key_id": None,
"aws_secret_access_key": None,
"aws_role_name": "arn:aws:iam::1111111111111:role/LitellmRole",
"aws_session_name": None
}
args2 = {
"aws_access_key_id": None,
"aws_secret_access_key": None,
"aws_role_name": "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
"aws_session_name": None
}
# Generate cache keys
cache_key1 = base_aws_llm.get_cache_key(args1)
cache_key2 = base_aws_llm.get_cache_key(args2)
# Cache keys should be different because the role names are different
assert cache_key1 != cache_key2
def test_eks_irsa_ambient_credentials_used():
"""
Test that in EKS/IRSA environments, ambient credentials are used when no explicit keys provided.
This allows web identity tokens to work automatically.
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
# Mock the STS response with proper expiration handling
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
current_time = datetime.now(timezone.utc)
# Create a timedelta object that returns 3600 when total_seconds() is called
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "assumed-access-key",
"SecretAccessKey": "assumed-secret-key",
"SessionToken": "assumed-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
# Call with no explicit credentials (EKS/IRSA scenario)
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="test-session"
)
# Should create STS client without explicit credentials (using ambient credentials)
mock_boto3_client.assert_called_once_with("sts")
# Should call assume_role
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
RoleSessionName="test-session"
)
# Verify credentials are returned correctly
assert credentials.access_key == "assumed-access-key"
assert credentials.secret_key == "assumed-secret-key"
assert credentials.token == "assumed-session-token"
assert ttl is not None
def test_explicit_credentials_used_when_provided():
"""
Test that explicit credentials are used when provided (non-EKS/IRSA scenario).
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
# Mock the STS response with proper expiration handling
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
current_time = datetime.now(timezone.utc)
# Create a timedelta object that returns 3600 when total_seconds() is called
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "assumed-access-key",
"SecretAccessKey": "assumed-secret-key",
"SessionToken": "assumed-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
# Call with explicit credentials
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id="explicit-access-key",
aws_secret_access_key="explicit-secret-key",
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="test-session"
)
# Should create STS client with explicit credentials
mock_boto3_client.assert_called_once_with(
"sts",
aws_access_key_id="explicit-access-key",
aws_secret_access_key="explicit-secret-key",
)
# Should call assume_role
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
RoleSessionName="test-session"
)
# Verify credentials are returned correctly
assert credentials.access_key == "assumed-access-key"
assert credentials.secret_key == "assumed-secret-key"
assert credentials.token == "assumed-session-token"
assert ttl is not None
def test_partial_credentials_still_use_ambient():
"""
Test that if only one credential is provided, we still use ambient credentials.
This handles edge cases where configuration might be incomplete.
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
# Mock the STS response
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "assumed-access-key",
"SecretAccessKey": "assumed-secret-key",
"SessionToken": "assumed-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
# Call with only access key (missing secret key)
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id="AKIAEXAMPLE",
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="test-session"
)
# Should still pass partial credentials to boto3.client
mock_boto3_client.assert_called_once_with(
"sts",
aws_access_key_id="AKIAEXAMPLE",
aws_secret_access_key=None
)
# Should still call assume_role
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
RoleSessionName="test-session"
)
def test_cross_account_role_assumption():
"""
Test assuming a role in a different AWS account (common in multi-account setups).
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
# Mock the STS response for cross-account role
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "cross-account-access-key",
"SecretAccessKey": "cross-account-secret-key",
"SessionToken": "cross-account-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
# Assume role in different account (EKS/IRSA scenario)
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::999999999999:role/CrossAccountRole",
aws_session_name="cross-account-session"
)
# Should use ambient credentials
mock_boto3_client.assert_called_once_with("sts")
# Should call assume_role with cross-account role
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::999999999999:role/CrossAccountRole",
RoleSessionName="cross-account-session"
)
# Verify cross-account credentials are returned
assert credentials.access_key == "cross-account-access-key"
assert credentials.secret_key == "cross-account-secret-key"
assert credentials.token == "cross-account-session-token"
assert ttl is not None
def test_role_assumption_with_custom_session_name():
"""
Test role assumption with a custom session name.
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
# Mock the STS response
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "custom-session-access-key",
"SecretAccessKey": "custom-session-secret-key",
"SessionToken": "custom-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client):
# Use custom session name
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole",
aws_session_name="evals-bedrock-session"
)
# Should call assume_role with custom session name
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::1111111111111:role/LitellmRole",
RoleSessionName="evals-bedrock-session"
)
# Verify credentials are returned
assert credentials.access_key == "custom-session-access-key"
assert credentials.secret_key == "custom-session-secret-key"
assert credentials.token == "custom-session-token"
def test_role_assumption_ttl_calculation():
"""
Test that TTL is calculated correctly from STS response expiration.
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
# Create a real datetime for expiration (1 hour from now)
expiration_time = datetime.now(timezone.utc) + timedelta(hours=1)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "ttl-test-access-key",
"SecretAccessKey": "ttl-test-secret-key",
"SessionToken": "ttl-test-session-token",
"Expiration": expiration_time,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client):
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole",
aws_session_name="ttl-test-session"
)
# TTL should be approximately 3540 seconds (1 hour - 60 second buffer)
assert ttl is not None
assert 3500 <= ttl <= 3600 # Allow some variance for test execution time
def test_role_assumption_error_handling():
"""
Test that role assumption errors are properly propagated.
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client to raise an exception
mock_sts_client = MagicMock()
mock_sts_client.assume_role.side_effect = Exception("AccessDenied: User is not authorized to perform sts:AssumeRole")
with patch("boto3.client", return_value=mock_sts_client):
# Should raise the exception
with pytest.raises(Exception) as exc_info:
base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole",
aws_session_name="error-test-session"
)
assert "AccessDenied" in str(exc_info.value)
def test_multiple_role_assumptions_in_sequence():
"""
Test that multiple role assumptions work correctly in sequence.
This simulates the scenario where different models use different roles.
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
# Mock different responses for different roles
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
# First role response
mock_sts_response1 = {
"Credentials": {
"AccessKeyId": "role1-access-key",
"SecretAccessKey": "role1-secret-key",
"SessionToken": "role1-session-token",
"Expiration": mock_expiry,
}
}
# Second role response
mock_sts_response2 = {
"Credentials": {
"AccessKeyId": "role2-access-key",
"SecretAccessKey": "role2-secret-key",
"SessionToken": "role2-session-token",
"Expiration": mock_expiry,
}
}
# Configure mock to return different responses
mock_sts_client.assume_role.side_effect = [mock_sts_response1, mock_sts_response2]
with patch("boto3.client", return_value=mock_sts_client):
# First role assumption
credentials1, ttl1 = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole",
aws_session_name="session-1"
)
# Second role assumption
credentials2, ttl2 = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="session-2"
)
# Verify both role assumptions were made
assert mock_sts_client.assume_role.call_count == 2
# Verify first role credentials
assert credentials1.access_key == "role1-access-key"
assert credentials1.secret_key == "role1-secret-key"
assert credentials1.token == "role1-session-token"
# Verify second role credentials
assert credentials2.access_key == "role2-access-key"
assert credentials2.secret_key == "role2-secret-key"
assert credentials2.token == "role2-session-token"
def test_auth_with_aws_role_irsa_environment():
"""Test that _auth_with_aws_role detects and uses IRSA environment variables"""
base_llm = BaseAWSLLM()
# Create a temporary file to simulate the web identity token
import tempfile
with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
f.write('test-web-identity-token')
token_file = f.name
try:
# Set IRSA environment variables
with patch.dict(os.environ, {
'AWS_WEB_IDENTITY_TOKEN_FILE': token_file,
'AWS_ROLE_ARN': 'arn:aws:iam::111111111111:role/eks-service-account-role',
'AWS_REGION': 'us-east-1'
}):
# Mock the boto3 STS client
mock_sts_client = MagicMock()
mock_assume_web_identity_response = {
'Credentials': {
'AccessKeyId': 'irsa-temp-access-key',
'SecretAccessKey': 'irsa-temp-secret-key',
'SessionToken': 'irsa-temp-session-token',
'Expiration': datetime.now() + timedelta(hours=1)
}
}
mock_assume_role_response = {
'Credentials': {
'AccessKeyId': 'irsa-access-key',
'SecretAccessKey': 'irsa-secret-key',
'SessionToken': 'irsa-session-token',
'Expiration': datetime.now() + timedelta(hours=1)
}
}
mock_sts_client.assume_role_with_web_identity.return_value = mock_assume_web_identity_response
mock_sts_client.assume_role.return_value = mock_assume_role_response
with patch('boto3.client', return_value=mock_sts_client) as mock_boto3_client:
# Call _auth_with_aws_role without explicit credentials
creds, ttl = base_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name='arn:aws:iam::222222222222:role/target-role',
aws_session_name='test-session'
)
# Verify boto3.client was called multiple times
# First for manual IRSA, then with IRSA credentials
assert mock_boto3_client.call_count >= 2
# Verify assume_role_with_web_identity was called
mock_sts_client.assume_role_with_web_identity.assert_called_once_with(
RoleArn='arn:aws:iam::111111111111:role/eks-service-account-role',
RoleSessionName='test-session',
WebIdentityToken='test-web-identity-token'
)
# Verify assume_role was called with correct parameters
mock_sts_client.assume_role.assert_called_once_with(
RoleArn='arn:aws:iam::222222222222:role/target-role',
RoleSessionName='test-session'
)
# Verify the returned credentials
assert creds.access_key == 'irsa-access-key'
assert creds.secret_key == 'irsa-secret-key'
assert creds.token == 'irsa-session-token'
assert ttl > 0 # TTL should be positive
finally:
# Clean up the temporary file
os.unlink(token_file)
def test_auth_with_aws_role_same_role_irsa():
"""Test that when IRSA role matches the requested role, we skip assumption"""
base_llm = BaseAWSLLM()
# Set IRSA environment variables
with patch.dict(os.environ, {
'AWS_ROLE_ARN': 'arn:aws:iam::111111111111:role/LitellmRole',
'AWS_WEB_IDENTITY_TOKEN_FILE': '/var/run/secrets/eks.amazonaws.com/serviceaccount/token'
}):
# Mock the _auth_with_env_vars method
mock_creds = MagicMock()
mock_creds.access_key = 'irsa-access-key'
mock_creds.secret_key = 'irsa-secret-key'
mock_creds.token = 'irsa-session-token'
with patch.object(base_llm, '_auth_with_env_vars', return_value=(mock_creds, None)) as mock_env_auth:
# Call get_credentials instead of _auth_with_aws_role directly
# This tests the full flow
creds = base_llm.get_credentials(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name='arn:aws:iam::111111111111:role/LitellmRole', # Same as AWS_ROLE_ARN
aws_session_name='test-session',
aws_region_name='us-east-1'
)
# Verify it used the env vars auth (no role assumption)
mock_env_auth.assert_called_once()
# Verify the returned credentials
assert creds.access_key == 'irsa-access-key'

2
tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py
View file

@ -18,6 +18,8 @@ class TestDataRobotConfig:
(None, "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"),
("http://localhost:5001", "http://localhost:5001/api/v2/genai/llmgw/chat/completions/"),
("https://app.datarobot.com", "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"),
("https://app.datarobot.com/api/v2/", "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"),
("https://app.datarobot.com/api/v2", "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"),
("https://app.datarobot.com/api/v2/genai/llmgw/chat/completions", "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"),
("https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/", "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"),
("https://staging.datarobot.com", "https://staging.datarobot.com/api/v2/genai/llmgw/chat/completions/"),

349
tests/test_litellm/llms/deepinfra/test_deepinfra_rerank.py Normal file
View file

@ -0,0 +1,349 @@
"""
Tests for DeepInfra rerank functionality following repository patterns.
"""
import asyncio
import json
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
# Add litellm to path
sys.path.insert(0, os.path.abspath("../../../.."))
import litellm
def assert_response_shape(response, custom_llm_provider):
"""Helper function to validate response structure."""
assert hasattr(response, "id")
assert hasattr(response, "results")
assert hasattr(response, "meta")
assert isinstance(response.results, list)
for result in response.results:
assert "index" in result
assert "relevance_score" in result
assert isinstance(result["index"], int)
assert isinstance(result["relevance_score"], (int, float))
# Check meta structure
assert "tokens" in response.meta
assert "billed_units" in response.meta
assert "input_tokens" in response.meta["tokens"]
assert "total_tokens" in response.meta["billed_units"]
@pytest.mark.parametrize("sync_mode", [True, False])
@patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post")
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_basic_rerank_deepinfra(mock_sync_post, mock_async_post, sync_mode):
"""Test basic DeepInfra rerank functionality."""
# Mock response data that matches DeepInfra API format
mock_response_data = {
"scores": [0.9, 0.1],
"input_tokens": 25,
"request_id": "deepinfra-request-123",
"inference_status": {
"status": "success",
"runtime_ms": 150,
"cost": 0.0001,
"tokens_generated": 0,
"tokens_input": 25,
},
}
def return_val():
return mock_response_data
api_key = "test_deepinfra_api_key"
api_base = "https://api.deepinfra.com"
if sync_mode:
# Create mock response object for sync
mock_response = MagicMock()
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_sync_post.return_value = mock_response
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="hello",
documents=["hello", "world"],
top_n=2,
custom_llm_provider="deepinfra",
api_key=api_key,
api_base=api_base,
)
mock_sync_post.assert_called_once()
else:
# Create mock response object for async
mock_response = AsyncMock()
def return_val():
return mock_response_data
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_async_post.return_value = mock_response
response = asyncio.run(
litellm.arerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="hello",
documents=["hello", "world"],
top_n=2,
custom_llm_provider="deepinfra",
api_key=api_key,
api_base=api_base,
)
)
mock_async_post.assert_called_once()
# Verify response structure
assert response.id == "deepinfra-request-123"
assert response.results is not None
assert len(response.results) == 2
assert response.results[0]["index"] == 0
assert response.results[0]["relevance_score"] == 0.9
assert response.results[1]["index"] == 1
assert response.results[1]["relevance_score"] == 0.1
# Verify metadata
assert response.meta["tokens"]["input_tokens"] == 25
assert response.meta["billed_units"]["total_tokens"] == 25
# Verify hidden params specific to DeepInfra
assert response._hidden_params["status"] == "success"
assert response._hidden_params["runtime_ms"] == 150
assert response._hidden_params["cost"] == 0.0001
# Note: The model name is processed and the 'deepinfra/' prefix is removed
assert response._hidden_params["model"] == "Qwen/Qwen3-Reranker-0.6B"
assert_response_shape(response, custom_llm_provider="deepinfra")
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_deepinfra_rerank_with_queries_param(mock_post):
"""Test DeepInfra rerank with multiple queries parameter."""
mock_response_data = {
"scores": [0.8, 0.6, 0.2],
"input_tokens": 35,
"request_id": "deepinfra-multi-query-123",
"inference_status": {"status": "success", "runtime_ms": 200},
}
def return_val():
return mock_response_data
mock_response = MagicMock()
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_post.return_value = mock_response
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-4B",
query="hello",
documents=["hello", "world", "test"],
queries=["hello", "hi there"], # DeepInfra specific param
custom_llm_provider="deepinfra",
api_key="test_key",
api_base="https://api.deepinfra.com",
)
mock_post.assert_called_once()
# Verify that queries parameter was passed in request
call_data = json.loads(mock_post.call_args.kwargs["data"])
assert "queries" in call_data
assert call_data["queries"] == ["hello", "hi there"]
assert response.results is not None
assert len(response.results) == 3
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_deepinfra_rerank_with_service_tier(mock_post):
"""Test DeepInfra rerank with service_tier parameter."""
mock_response_data = {
"scores": [0.95, 0.75],
"input_tokens": 30,
"request_id": "deepinfra-premium-123",
}
def return_val():
return mock_response_data
mock_response = MagicMock()
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_post.return_value = mock_response
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-8B",
query="premium search",
documents=["doc1", "doc2"],
service_tier="premium", # DeepInfra specific param
custom_llm_provider="deepinfra",
api_key="test_key",
api_base="https://api.deepinfra.com",
)
mock_post.assert_called_once()
# Verify URL
call_url = mock_post.call_args.kwargs["url"]
assert "api.deepinfra.com/inference/Qwen/Qwen3-Reranker-8B" in call_url
# Verify request contains service_tier
call_data = json.loads(mock_post.call_args.kwargs["data"])
assert call_data["service_tier"] == "premium"
assert response.results is not None
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_deepinfra_rerank_request_format(mock_post):
"""Test that the request is properly formatted for DeepInfra API."""
mock_response_data = {"scores": [0.9, 0.1], "input_tokens": 20}
def return_val():
return mock_response_data
mock_response = MagicMock()
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_post.return_value = mock_response
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="test query",
documents=["doc1", "doc2"],
custom_llm_provider="deepinfra",
api_key="test_key",
api_base="https://api.deepinfra.com",
instruction="custom instruction",
webhook="https://webhook.example.com",
)
mock_post.assert_called_once()
# Verify URL format
call_url = mock_post.call_args.kwargs["url"]
assert call_url == "https://api.deepinfra.com/inference/Qwen/Qwen3-Reranker-0.6B"
# Verify headers
headers = mock_post.call_args.kwargs["headers"]
assert headers["Authorization"] == "Bearer test_key"
assert headers["accept"] == "application/json"
assert headers["content-type"] == "application/json"
# Verify request body format
request_data = json.loads(mock_post.call_args.kwargs["data"])
assert request_data["queries"] == [
"test query",
"test query",
] # DeepInfra requires queries to match documents length
assert request_data["documents"] == ["doc1", "doc2"]
assert request_data["instruction"] == "custom instruction"
assert request_data["webhook"] == "https://webhook.example.com"
assert response.results is not None
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_deepinfra_rerank_error_handling(mock_post):
"""Test DeepInfra rerank error handling."""
error_response = {"detail": {"error": "Invalid API key"}}
def return_val():
return error_response
mock_response = MagicMock()
mock_response.status_code = 401
mock_response.json = return_val
mock_response.text = json.dumps(error_response)
mock_response.headers = {"content-type": "application/json"}
mock_post.return_value = mock_response
# The current implementation handles errors gracefully, so we expect a successful response
# with the error information in the hidden params
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="hello",
documents=["hello", "world"],
custom_llm_provider="deepinfra",
api_key="invalid_key",
api_base="https://api.deepinfra.com",
)
# Verify that the response contains error information
assert (
response._hidden_params["status"] == "unknown"
) # Default status when error occurs
def test_deepinfra_rerank_models():
"""Test that DeepInfra Qwen rerank models are recognized."""
# These should not raise errors during model validation
models = [
"deepinfra/Qwen/Qwen3-Reranker-0.6B",
"deepinfra/Qwen/Qwen3-Reranker-4B",
"deepinfra/Qwen/Qwen3-Reranker-8B",
]
for model in models:
# This should not raise any validation errors
try:
litellm.get_llm_provider(model=model)
except Exception as e:
# We expect this to potentially fail due to missing api_base/key
# but the model format should be recognized
assert "api_base" in str(e) or "API key" in str(
e
), f"Unexpected error for model {model}: {e}"
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_deepinfra_rerank_minimal_response(mock_post):
"""Test handling of minimal DeepInfra response."""
# Minimal response with just scores
mock_response_data = {"scores": [0.7, 0.3]}
def return_val():
return mock_response_data
mock_response = MagicMock()
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_post.return_value = mock_response
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="hello",
documents=["hello", "world"],
custom_llm_provider="deepinfra",
api_key="test_key",
api_base="https://api.deepinfra.com",
)
# Should handle minimal response gracefully
assert response.results is not None
assert len(response.results) == 2
assert response.results[0]["relevance_score"] == 0.7
assert response.results[1]["relevance_score"] == 0.3
# Should have default values for missing fields
assert response.meta["tokens"]["input_tokens"] == 0 # Default when missing
assert response._hidden_params["status"] == "unknown" # Default when missing

435
tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py Normal file
View file

@ -0,0 +1,435 @@
"""
Integration tests for DeepInfra rerank functionality.
Tests the full rerank flow following the repository patterns.
"""
import asyncio
import json
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import litellm
def assert_response_shape(response, custom_llm_provider):
"""Helper function to validate response structure specific to DeepInfra."""
assert hasattr(response, "id")
assert hasattr(response, "results")
assert hasattr(response, "meta")
assert isinstance(response.results, list)
for result in response.results:
assert "index" in result
assert "relevance_score" in result
assert isinstance(result["index"], int)
assert isinstance(result["relevance_score"], (int, float))
# Check meta structure
assert "tokens" in response.meta
assert "billed_units" in response.meta
assert "input_tokens" in response.meta["tokens"]
assert "total_tokens" in response.meta["billed_units"]
@pytest.mark.parametrize("sync_mode", [True, False])
@patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post")
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_basic_rerank_deepinfra(mock_sync_post, mock_async_post, sync_mode):
"""Test basic DeepInfra rerank functionality."""
# Mock response data that matches DeepInfra API format
mock_response_data = {
"scores": [0.9, 0.1],
"input_tokens": 25,
"request_id": "deepinfra-request-123",
"inference_status": {
"status": "success",
"runtime_ms": 150,
"cost": 0.0001,
"tokens_generated": 0,
"tokens_input": 25,
},
}
def return_val():
return mock_response_data
api_key = "test_deepinfra_api_key"
api_base = "https://api.deepinfra.com"
if sync_mode:
# Create mock response object for sync
mock_response = MagicMock()
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_sync_post.return_value = mock_response
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="hello",
documents=["hello", "world"],
top_n=2,
custom_llm_provider="deepinfra",
api_key=api_key,
api_base=api_base,
)
mock_sync_post.assert_called_once()
else:
# Create mock response object for async
mock_response = AsyncMock()
def return_val():
return mock_response_data
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_async_post.return_value = mock_response
response = asyncio.run(
litellm.arerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="hello",
documents=["hello", "world"],
top_n=2,
custom_llm_provider="deepinfra",
api_key=api_key,
api_base=api_base,
)
)
mock_async_post.assert_called_once()
# Verify response structure
assert response.id == "deepinfra-request-123"
assert response.results is not None
assert len(response.results) == 2
assert response.results[0]["index"] == 0
assert response.results[0]["relevance_score"] == 0.9
assert response.results[1]["index"] == 1
assert response.results[1]["relevance_score"] == 0.1
# Verify metadata
assert response.meta["tokens"]["input_tokens"] == 25
assert response.meta["billed_units"]["total_tokens"] == 25
# Verify hidden params specific to DeepInfra
assert response._hidden_params["status"] == "success"
assert response._hidden_params["runtime_ms"] == 150
assert response._hidden_params["cost"] == 0.0001
# Note: The model name is processed and the 'deepinfra/' prefix is removed
assert response._hidden_params["model"] == "Qwen/Qwen3-Reranker-0.6B"
assert_response_shape(response, custom_llm_provider="deepinfra")
@pytest.mark.parametrize("sync_mode", [True, False])
@patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post")
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_deepinfra_rerank_with_queries_param(
mock_sync_post, mock_async_post, sync_mode
):
"""Test DeepInfra rerank with multiple queries parameter."""
mock_response_data = {
"scores": [0.8, 0.6, 0.2],
"input_tokens": 35,
"request_id": "deepinfra-multi-query-123",
"inference_status": {"status": "success", "runtime_ms": 200},
}
def return_val():
return mock_response_data
if sync_mode:
mock_response = MagicMock()
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_sync_post.return_value = mock_response
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-4B",
query="hello",
documents=["hello", "world", "test"],
queries=["hello", "hi there"], # DeepInfra specific param
custom_llm_provider="deepinfra",
api_key="test_key",
api_base="https://api.deepinfra.com",
)
mock_sync_post.assert_called_once()
# Verify that queries parameter was passed in request
call_data = json.loads(mock_sync_post.call_args.kwargs["data"])
assert "queries" in call_data
assert call_data["queries"] == ["hello", "hi there"]
else:
mock_response = AsyncMock()
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_async_post.return_value = mock_response
response = asyncio.run(
litellm.arerank(
model="deepinfra/Qwen/Qwen3-Reranker-4B",
query="hello",
documents=["hello", "world", "test"],
queries=["hello", "hi there"],
custom_llm_provider="deepinfra",
api_key="test_key",
api_base="https://api.deepinfra.com",
)
)
mock_async_post.assert_called_once()
call_data = json.loads(mock_async_post.call_args.kwargs["data"])
assert "queries" in call_data
assert call_data["queries"] == ["hello", "hi there"]
assert response.results is not None
assert len(response.results) == 3
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_deepinfra_rerank_with_service_tier(mock_post):
"""Test DeepInfra rerank with service_tier parameter."""
mock_response_data = {
"scores": [0.95, 0.75],
"input_tokens": 30,
"request_id": "deepinfra-premium-123",
}
def return_val():
return mock_response_data
mock_response = MagicMock()
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_post.return_value = mock_response
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-8B",
query="premium search",
documents=["doc1", "doc2"],
service_tier="premium", # DeepInfra specific param
custom_llm_provider="deepinfra",
api_key="test_key",
api_base="https://api.deepinfra.com",
)
mock_post.assert_called_once()
# Verify URL
call_url = mock_post.call_args.kwargs["url"]
assert "api.deepinfra.com/inference/Qwen/Qwen3-Reranker-8B" in call_url
# Verify request contains service_tier
call_data = json.loads(mock_post.call_args.kwargs["data"])
assert call_data["service_tier"] == "premium"
assert response.results is not None
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_deepinfra_rerank_with_env_vars(mock_post, monkeypatch):
"""Test DeepInfra rerank with environment variable configuration."""
monkeypatch.setenv("DEEPINFRA_API_KEY", "env_test_key")
monkeypatch.setenv("DEEPINFRA_API_BASE", "https://custom-deepinfra.com")
mock_response_data = {
"scores": [0.88, 0.22],
"input_tokens": 28,
"request_id": "env-test-123",
}
def return_val():
return mock_response_data
mock_response = MagicMock()
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_post.return_value = mock_response
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="hello",
documents=["hello", "world"],
custom_llm_provider="deepinfra",
)
mock_post.assert_called_once()
# Verify headers contain env API key
headers = mock_post.call_args.kwargs.get("headers", {})
assert "Bearer env_test_key" in headers.get("Authorization", "")
assert response.results is not None
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_deepinfra_rerank_error_handling(mock_post):
"""Test DeepInfra rerank error handling."""
error_response = {"detail": {"error": "Invalid API key"}}
def return_val():
return error_response
mock_response = MagicMock()
mock_response.status_code = 401
mock_response.json = return_val
mock_response.text = json.dumps(error_response)
mock_response.headers = {"content-type": "application/json"}
mock_post.return_value = mock_response
# The current implementation handles errors gracefully, so we expect a successful response
# with the error information in the hidden params
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="hello",
documents=["hello", "world"],
custom_llm_provider="deepinfra",
api_key="invalid_key",
api_base="https://api.deepinfra.com",
)
# Verify that the response contains error information
assert (
response._hidden_params["status"] == "unknown"
) # Default status when error occurs
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_deepinfra_rerank_missing_api_base_error(mock_post):
"""Test error handling when API base is missing."""
# Note: The current implementation may have a default API base or the test environment
# may be providing one, so we'll test the actual behavior
try:
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="hello",
documents=["hello", "world"],
custom_llm_provider="deepinfra",
api_key="test_key",
# api_base is intentionally missing
)
# If no error is raised, it means a default API base is being used
# This is acceptable behavior
assert response is not None
except ValueError as e:
# If an error is raised, it should match the expected message
assert "api_base must be provided for Deepinfra rerank" in str(e)
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_deepinfra_rerank_request_format(mock_post):
"""Test that the request is properly formatted for DeepInfra API."""
mock_response_data = {"scores": [0.9, 0.1], "input_tokens": 20}
def return_val():
return mock_response_data
mock_response = MagicMock()
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_post.return_value = mock_response
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="test query",
documents=["doc1", "doc2"],
custom_llm_provider="deepinfra",
api_key="test_key",
api_base="https://api.deepinfra.com",
instruction="custom instruction",
webhook="https://webhook.example.com",
)
mock_post.assert_called_once()
# Verify URL format
call_url = mock_post.call_args.kwargs["url"]
assert call_url == "https://api.deepinfra.com/inference/Qwen/Qwen3-Reranker-0.6B"
# Verify headers
headers = mock_post.call_args.kwargs["headers"]
assert headers["Authorization"] == "Bearer test_key"
assert headers["accept"] == "application/json"
assert headers["content-type"] == "application/json"
# Verify request body format
request_data = json.loads(mock_post.call_args.kwargs["data"])
assert request_data["queries"] == [
"test query",
"test query",
] # DeepInfra requires queries to match documents length
assert request_data["documents"] == ["doc1", "doc2"]
assert request_data["instruction"] == "custom instruction"
assert request_data["webhook"] == "https://webhook.example.com"
assert response.results is not None
def test_deepinfra_rerank_models():
"""Test that DeepInfra Qwen rerank models are recognized."""
# These should not raise errors during model validation
models = [
"deepinfra/Qwen/Qwen3-Reranker-0.6B",
"deepinfra/Qwen/Qwen3-Reranker-4B",
"deepinfra/Qwen/Qwen3-Reranker-8B",
]
for model in models:
# This should not raise any validation errors
try:
litellm.get_llm_provider(model=model)
except Exception as e:
# We expect this to potentially fail due to missing api_base/key
# but the model format should be recognized
assert "api_base" in str(e) or "API key" in str(
e
), f"Unexpected error for model {model}: {e}"
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_deepinfra_rerank_minimal_response(mock_post):
"""Test handling of minimal DeepInfra response."""
# Minimal response with just scores
mock_response_data = {"scores": [0.7, 0.3]}
def return_val():
return mock_response_data
mock_response = MagicMock()
mock_response.json = return_val
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.text = json.dumps(mock_response_data)
mock_post.return_value = mock_response
response = litellm.rerank(
model="deepinfra/Qwen/Qwen3-Reranker-0.6B",
query="hello",
documents=["hello", "world"],
custom_llm_provider="deepinfra",
api_key="test_key",
api_base="https://api.deepinfra.com",
)
# Should handle minimal response gracefully
assert response.results is not None
assert len(response.results) == 2
assert response.results[0]["relevance_score"] == 0.7
assert response.results[1]["relevance_score"] == 0.3
# Should have default values for missing fields
assert response.meta["tokens"]["input_tokens"] == 0 # Default when missing
assert response._hidden_params["status"] == "unknown" # Default when missing

303
tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py Normal file
View file

@ -0,0 +1,303 @@
"""
Tests for DeepInfra rerank transformation functionality.
Based on the test patterns from other rerank providers and the current DeepInfra implementation.
"""
import json
from unittest.mock import MagicMock
import httpx
import pytest
from litellm.llms.deepinfra.rerank.transformation import DeepinfraRerankConfig
from litellm.types.rerank import (
OptionalRerankParams,
RerankResponse,
)
class TestDeepinfraRerankTransform:
def setup_method(self):
self.config = DeepinfraRerankConfig()
self.model = "deepinfra/Qwen/Qwen3-Reranker-0.6B"
def test_get_complete_url(self):
"""Test URL generation for DeepInfra rerank API."""
# Test basic URL generation
api_base = "https://api.deepinfra.com"
model = "Qwen/Qwen3-Reranker-0.6B"
url = self.config.get_complete_url(api_base, model)
assert url == "https://api.deepinfra.com/inference/Qwen/Qwen3-Reranker-0.6B"
# Test URL with slash at the end
api_base_with_slash = "https://api.deepinfra.com/"
url = self.config.get_complete_url(api_base_with_slash, model)
assert url == "https://api.deepinfra.com/inference/Qwen/Qwen3-Reranker-0.6B"
# Test URL with openai replacement
api_base_openai = "https://api.deepinfra.com/openai"
url = self.config.get_complete_url(api_base_openai, model)
assert url == "https://api.deepinfra.com/inference/Qwen/Qwen3-Reranker-0.6B"
# Test error when api_base is None
with pytest.raises(ValueError, match="Deepinfra API Base is required"):
self.config.get_complete_url(None, model)
def test_map_cohere_rerank_params_basic(self):
"""Test basic parameter mapping for DeepInfra rerank."""
params = self.config.map_cohere_rerank_params(
non_default_params={"documents": ["doc1", "doc2"]},
model=self.model,
drop_params=False,
query="test query",
documents=["doc1", "doc2"],
)
assert params["queries"] == [
"test query",
"test query",
] # DeepInfra requires queries to match documents length
assert params["documents"] == ["doc1", "doc2"]
def test_map_cohere_rerank_params_with_non_default(self):
"""Test parameter mapping with DeepInfra-specific parameters."""
non_default_params = {
"queries": ["custom query"],
"documents": ["doc1", "doc2", "doc3"],
"service_tier": "premium",
"instruction": "custom instruction",
"webhook": "https://webhook.example.com",
}
params = self.config.map_cohere_rerank_params(
non_default_params=non_default_params,
model=self.model,
drop_params=False,
query="test query",
documents=["doc1", "doc2"],
)
# queries should override the query parameter (custom queries take precedence)
assert params["queries"] == ["custom query"]
assert params["documents"] == ["doc1", "doc2", "doc3"]
assert params["service_tier"] == "premium"
assert params["instruction"] == "custom instruction"
assert params["webhook"] == "https://webhook.example.com"
def test_transform_rerank_request(self):
"""Test request transformation for DeepInfra format."""
optional_params = OptionalRerankParams(
queries=["test query"],
documents=["doc1", "doc2"],
service_tier="default",
)
request_body = self.config.transform_rerank_request(
model=self.model, optional_rerank_params=optional_params, headers={}
)
assert request_body["queries"] == ["test query"]
assert request_body["documents"] == ["doc1", "doc2"]
assert request_body["service_tier"] == "default"
def test_transform_rerank_request_missing_documents(self):
"""Test that transform_rerank_request handles missing documents gracefully."""
optional_params = OptionalRerankParams(queries=["test query"])
# The current implementation doesn't validate documents, it just returns the params
result = self.config.transform_rerank_request(
model=self.model, optional_rerank_params=optional_params, headers={}
)
assert result == optional_params
def test_transform_rerank_response_success(self):
"""Test successful response transformation."""
# Mock DeepInfra response format
response_data = {
"scores": [0.9, 0.7, 0.3],
"input_tokens": 42,
"request_id": "test-request-123",
"inference_status": {
"status": "success",
"runtime_ms": 150,
"cost": 0.0001,
"tokens_generated": 0,
"tokens_input": 42,
},
}
# Create mock httpx response
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = response_data
mock_response.text = json.dumps(response_data)
# Create mock logging object
mock_logging = MagicMock()
model_response = RerankResponse()
result = self.config.transform_rerank_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=mock_logging,
)
# Verify response structure
assert result.id == "test-request-123"
assert len(result.results) == 3
assert result.results[0]["index"] == 0
assert result.results[0]["relevance_score"] == 0.9
assert result.results[1]["index"] == 1
assert result.results[1]["relevance_score"] == 0.7
assert result.results[2]["index"] == 2
assert result.results[2]["relevance_score"] == 0.3
# Verify metadata
assert result.meta["tokens"]["input_tokens"] == 42
assert result.meta["tokens"]["output_tokens"] == 0
assert result.meta["billed_units"]["total_tokens"] == 42
# Verify hidden params
assert result._hidden_params["status"] == "success"
assert result._hidden_params["runtime_ms"] == 150
assert result._hidden_params["cost"] == 0.0001
assert result._hidden_params["tokens_generated"] == 0
assert result._hidden_params["tokens_input"] == 42
assert result._hidden_params["model"] == self.model
# Verify logging was called
mock_logging.post_call.assert_called_once_with(
original_response=mock_response.text
)
def test_transform_rerank_response_minimal(self):
"""Test response transformation with minimal data."""
response_data = {
"scores": [0.8, 0.2],
"input_tokens": 20,
}
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = response_data
mock_response.text = json.dumps(response_data)
mock_logging = MagicMock()
model_response = RerankResponse()
result = self.config.transform_rerank_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=mock_logging,
)
# Should generate UUID when request_id is missing
assert result.id is not None
assert len(result.id) > 0
# Should handle missing inference_status gracefully
assert result._hidden_params["status"] == "unknown"
assert result._hidden_params["runtime_ms"] == 0
assert result._hidden_params["cost"] == 0.0
def test_transform_rerank_response_error_fallback(self):
"""Test error handling and fallback in response transformation."""
# Create a response that will cause JSON parsing to fail
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "doc", 0)
mock_response.text = "Invalid JSON response"
mock_logging = MagicMock()
model_response = RerankResponse()
# The current implementation should handle JSON parsing errors gracefully
# by falling back to the parent implementation
result = self.config.transform_rerank_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=mock_logging,
)
# Should return the original model_response when fallback occurs
assert result == model_response
def test_get_supported_cohere_rerank_params(self):
"""Test getting supported parameters for DeepInfra rerank."""
supported_params = self.config.get_supported_cohere_rerank_params(self.model)
assert "query" in supported_params
assert "documents" in supported_params
assert len(supported_params) == 2
def test_query_replication_for_deepinfra_requirement(self):
"""Test that queries are replicated to match documents length as required by DeepInfra."""
# Test with different document lengths
test_cases = [
(["doc1"], ["query1"]),
(["doc1", "doc2"], ["query1", "query1"]),
(["doc1", "doc2", "doc3"], ["query1", "query1", "query1"]),
]
for documents, expected_queries in test_cases:
params = self.config.map_cohere_rerank_params(
non_default_params={},
model=self.model,
drop_params=False,
query="query1",
documents=documents,
)
assert (
params["queries"] == expected_queries
), f"Failed for {len(documents)} documents"
assert len(params["queries"]) == len(
documents
), "Queries length must match documents length"
def test_get_error_class_basic(self):
"""Test error class generation for basic error."""
error_message = "Authentication failed"
status_code = 401
headers = {"content-type": "application/json"}
with pytest.raises(Exception) as exc_info:
self.config.get_error_class(error_message, status_code, headers)
# The method should raise a BaseLLMException
assert exc_info.value.args[0] == error_message
def test_get_error_class_with_detail(self):
"""Test error class generation with DeepInfra error format."""
error_data = {"detail": {"error": "Model not found"}}
error_message = json.dumps(error_data)
status_code = 404
headers = {"content-type": "application/json"}
with pytest.raises(Exception) as exc_info:
self.config.get_error_class(error_message, status_code, headers)
# Should extract the nested error message
assert "Model not found" in str(exc_info.value)
def test_get_error_class_with_string_detail(self):
"""Test error class generation with string detail."""
error_data = {"detail": "Service unavailable"}
error_message = json.dumps(error_data)
status_code = 503
headers = {"content-type": "application/json"}
with pytest.raises(Exception) as exc_info:
self.config.get_error_class(error_message, status_code, headers)
# Should extract the string detail
assert "Service unavailable" in str(exc_info.value)
def test_get_error_class_invalid_json(self):
"""Test error class generation with invalid JSON."""
error_message = "Invalid JSON error message"
status_code = 500
headers = {"content-type": "application/json"}
with pytest.raises(Exception) as exc_info:
self.config.get_error_class(error_message, status_code, headers)
# Should use the original error message when JSON parsing fails
assert "Invalid JSON error message" in str(exc_info.value)

449
tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py
View file

@ -1,14 +1,18 @@
import os
import sys
from typing import List, cast
from unittest.mock import MagicMock, patch
import pytest
from litellm.types.llms.openai import AllMessageValues
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
from litellm.llms.mistral.chat.transformation import MistralConfig
from litellm.types.utils import ModelResponse
@pytest.mark.asyncio
@ -40,21 +44,25 @@ class TestMistralReasoningSupport:
def test_get_supported_openai_params_magistral_model(self):
"""Test that magistral models support reasoning parameters."""
mistral_config = MistralConfig()
# Test magistral model supports reasoning parameters
supported_params = mistral_config.get_supported_openai_params("mistral/magistral-medium-2506")
supported_params = mistral_config.get_supported_openai_params(
"mistral/magistral-medium-2506"
)
assert "reasoning_effort" in supported_params
assert "thinking" in supported_params
# Test non-magistral model doesn't include reasoning parameters
supported_params_normal = mistral_config.get_supported_openai_params("mistral/mistral-large-latest")
supported_params_normal = mistral_config.get_supported_openai_params(
"mistral/mistral-large-latest"
)
assert "reasoning_effort" not in supported_params_normal
assert "thinking" not in supported_params_normal
def test_map_openai_params_reasoning_effort(self):
"""Test that reasoning_effort parameter is properly mapped for magistral models."""
mistral_config = MistralConfig()
# Test reasoning_effort mapping for magistral model
optional_params = {}
result = mistral_config.map_openai_params(
@ -63,9 +71,9 @@ class TestMistralReasoningSupport:
model="mistral/magistral-medium-2506",
drop_params=False,
)
assert result.get("_add_reasoning_prompt") is True
# Test reasoning_effort ignored for non-magistral model
optional_params_normal = {}
result_normal = mistral_config.map_openai_params(
@ -74,13 +82,13 @@ class TestMistralReasoningSupport:
model="mistral/mistral-large-latest",
drop_params=False,
)
assert "_add_reasoning_prompt" not in result_normal
def test_map_openai_params_thinking(self):
"""Test that thinking parameter is properly mapped for magistral models."""
mistral_config = MistralConfig()
# Test thinking mapping for magistral model
optional_params = {}
result = mistral_config.map_openai_params(
@ -89,7 +97,7 @@ class TestMistralReasoningSupport:
model="mistral/magistral-small-2506",
drop_params=False,
)
assert result.get("_add_reasoning_prompt") is True
def test_get_mistral_reasoning_system_prompt(self):
@ -101,109 +109,123 @@ class TestMistralReasoningSupport:
def test_add_reasoning_system_prompt_no_existing_system_message(self):
"""Test adding reasoning system prompt when no system message exists."""
mistral_config = MistralConfig()
messages = [
{"role": "user", "content": "What is 2+2?"}
]
messages = [{"role": "user", "content": "What is 2+2?"}]
optional_params = {"_add_reasoning_prompt": True}
result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params)
result = mistral_config._add_reasoning_system_prompt_if_needed(
messages, optional_params
)
# Should add a new system message at the beginning
assert len(result) == 2
assert result[0]["role"] == "system"
assert "<think>" in result[0]["content"]
assert result[1]["role"] == "user"
assert result[1]["content"] == "What is 2+2?"
# Should remove the internal flag
assert "_add_reasoning_prompt" not in optional_params
def test_add_reasoning_system_prompt_with_existing_system_message(self):
"""Test adding reasoning system prompt when system message already exists."""
mistral_config = MistralConfig()
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is 2+2?"}
{"role": "user", "content": "What is 2+2?"},
]
optional_params = {"_add_reasoning_prompt": True}
result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params)
result = mistral_config._add_reasoning_system_prompt_if_needed(
messages, optional_params
)
# Should modify existing system message
assert len(result) == 2
assert result[0]["role"] == "system"
assert "<think>" in result[0]["content"]
assert "You are a helpful assistant." in result[0]["content"]
assert result[1]["role"] == "user"
# Should remove the internal flag
assert "_add_reasoning_prompt" not in optional_params
def test_add_reasoning_system_prompt_with_existing_list_content(self):
"""Test adding reasoning system prompt when system message has list content."""
mistral_config = MistralConfig()
messages = [
{
"role": "system",
"role": "system",
"content": [
{"type": "text", "text": "You are a helpful assistant."},
{"type": "text", "text": "You always provide detailed explanations."}
]
{
"type": "text",
"text": "You always provide detailed explanations.",
},
],
},
{"role": "user", "content": "What is 2+2?"}
{"role": "user", "content": "What is 2+2?"},
]
optional_params = {"_add_reasoning_prompt": True}
result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params)
result = mistral_config._add_reasoning_system_prompt_if_needed(
messages, optional_params
)
# Should modify existing system message preserving list format
assert len(result) == 2
assert result[0]["role"] == "system"
assert isinstance(result[0]["content"], list)
# First item should be the reasoning prompt
assert result[0]["content"][0]["type"] == "text"
assert "<think>" in result[0]["content"][0]["text"]
# Original content should be preserved
assert "You are a helpful assistant." in result[0]["content"][1]["text"]
assert "You always provide detailed explanations." in result[0]["content"][2]["text"]
assert (
"You always provide detailed explanations."
in result[0]["content"][2]["text"]
)
assert result[1]["role"] == "user"
# Should remove the internal flag
assert "_add_reasoning_prompt" not in optional_params
def test_add_reasoning_system_prompt_preserves_content_types(self):
"""Test that reasoning prompt preserves original content types (string vs list)."""
mistral_config = MistralConfig()
# Test with string content
string_messages = [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "Hello"}
{"role": "user", "content": "Hello"},
]
string_params = {"_add_reasoning_prompt": True}
string_result = mistral_config._add_reasoning_system_prompt_if_needed(string_messages, string_params)
string_result = mistral_config._add_reasoning_system_prompt_if_needed(
string_messages, string_params
)
assert isinstance(string_result[0]["content"], str)
assert "<think>" in string_result[0]["content"]
assert "You are helpful." in string_result[0]["content"]
# Test with list content
list_messages = [
{
"role": "system",
"content": [{"type": "text", "text": "You are helpful."}]
"role": "system",
"content": [{"type": "text", "text": "You are helpful."}],
},
{"role": "user", "content": "Hello"}
{"role": "user", "content": "Hello"},
]
list_params = {"_add_reasoning_prompt": True}
list_result = mistral_config._add_reasoning_system_prompt_if_needed(list_messages, list_params)
list_result = mistral_config._add_reasoning_system_prompt_if_needed(
list_messages, list_params
)
assert isinstance(list_result[0]["content"], list)
assert list_result[0]["content"][0]["type"] == "text"
assert "<think>" in list_result[0]["content"][0]["text"]
@ -212,14 +234,14 @@ class TestMistralReasoningSupport:
def test_add_reasoning_system_prompt_no_flag(self):
"""Test that no modification happens when _add_reasoning_prompt flag is not set."""
mistral_config = MistralConfig()
messages = [
{"role": "user", "content": "What is 2+2?"}
]
messages = [{"role": "user", "content": "What is 2+2?"}]
optional_params = {}
result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params)
result = mistral_config._add_reasoning_system_prompt_if_needed(
messages, optional_params
)
# Should return messages unchanged
assert result == messages
assert len(result) == 1
@ -227,46 +249,42 @@ class TestMistralReasoningSupport:
def test_transform_request_magistral_with_reasoning(self):
"""Test transform_request method for magistral model with reasoning."""
mistral_config = MistralConfig()
messages = [
{"role": "user", "content": "What is 15 * 7?"}
]
messages = [{"role": "user", "content": "What is 15 * 7?"}]
optional_params = {"_add_reasoning_prompt": True}
result = mistral_config.transform_request(
model="mistral/magistral-medium-2506",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
headers={},
)
# Should have added system message
assert len(result["messages"]) == 2
assert result["messages"][0]["role"] == "system"
assert "<think>" in result["messages"][0]["content"]
assert result["messages"][1]["role"] == "user"
# Should remove internal flag from optional_params
assert "_add_reasoning_prompt" not in result
def test_transform_request_magistral_without_reasoning(self):
"""Test transform_request method for magistral model without reasoning."""
mistral_config = MistralConfig()
messages = [
{"role": "user", "content": "What is 15 * 7?"}
]
messages = [{"role": "user", "content": "What is 15 * 7?"}]
optional_params = {}
result = mistral_config.transform_request(
model="mistral/magistral-medium-2506",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
headers={},
)
# Should not modify messages
assert len(result["messages"]) == 1
assert result["messages"][0]["role"] == "user"
@ -274,20 +292,18 @@ class TestMistralReasoningSupport:
def test_transform_request_non_magistral_with_reasoning_params(self):
"""Test that non-magistral models ignore reasoning parameters."""
mistral_config = MistralConfig()
messages = [
{"role": "user", "content": "What is 15 * 7?"}
]
messages = [{"role": "user", "content": "What is 15 * 7?"}]
optional_params = {"_add_reasoning_prompt": True}
result = mistral_config.transform_request(
model="mistral/mistral-large-latest",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
headers={},
)
# Should not add system message for non-magistral models
assert len(result["messages"]) == 1
assert result["messages"][0]["role"] == "user"
@ -295,15 +311,15 @@ class TestMistralReasoningSupport:
def test_case_insensitive_magistral_detection(self):
"""Test that magistral model detection is case-insensitive."""
mistral_config = MistralConfig()
# Test various case combinations
models_to_test = [
"mistral/Magistral-medium-2506",
"mistral/MAGISTRAL-MEDIUM-2506",
"mistral/magistral-SMALL-2506",
"MaGiStRaL-medium-2506"
"MaGiStRaL-medium-2506",
]
for model in models_to_test:
supported_params = mistral_config.get_supported_openai_params(model)
assert "reasoning_effort" in supported_params, f"Failed for model: {model}"
@ -311,7 +327,7 @@ class TestMistralReasoningSupport:
def test_end_to_end_reasoning_workflow(self):
"""Test the complete workflow from parameter to system prompt injection."""
mistral_config = MistralConfig()
# Step 1: Map parameters
optional_params = {}
mapped_params = mistral_config.map_openai_params(
@ -320,23 +336,21 @@ class TestMistralReasoningSupport:
model="mistral/magistral-medium-2506",
drop_params=False,
)
assert mapped_params.get("_add_reasoning_prompt") is True
assert mapped_params.get("temperature") == 0.7
# Step 2: Transform request
messages = [
{"role": "user", "content": "Solve for x: 2x + 5 = 13"}
]
messages = [{"role": "user", "content": "Solve for x: 2x + 5 = 13"}]
result = mistral_config.transform_request(
model="mistral/magistral-medium-2506",
messages=messages,
optional_params=mapped_params,
litellm_params={},
headers={}
headers={},
)
# Verify final result
assert len(result["messages"]) == 2
assert result["messages"][0]["role"] == "system"
@ -347,7 +361,6 @@ class TestMistralReasoningSupport:
assert "_add_reasoning_prompt" not in result
class TestMistralNameHandling:
"""Test suite for Mistral name handling in messages."""
@ -363,7 +376,11 @@ class TestMistralNameHandling:
def test_handle_name_in_message_tool_role_valid_name_keeps_name(self):
"""Test that valid name is kept for tool messages."""
# Test with normal function name
tool_message = {"role": "tool", "content": "Function result", "name": "get_weather"}
tool_message = {
"role": "tool",
"content": "Function result",
"name": "get_weather",
}
result = MistralConfig._handle_name_in_message(tool_message)
assert "name" in result
assert result["name"] == "get_weather"
@ -386,31 +403,140 @@ class TestMistralParallelToolCalls:
def test_get_supported_openai_params_includes_parallel_tool_calls(self):
"""Test that parallel_tool_calls is in supported parameters."""
mistral_config = MistralConfig()
supported_params = mistral_config.get_supported_openai_params("mistral/mistral-large-latest")
supported_params = mistral_config.get_supported_openai_params(
"mistral/mistral-large-latest"
)
assert "parallel_tool_calls" in supported_params
def test_transform_request_preserves_parallel_tool_calls(self):
"""Test that transform_request preserves parallel_tool_calls parameter."""
mistral_config = MistralConfig()
messages = [
{"role": "user", "content": "What's the weather like?"}
]
messages = [{"role": "user", "content": "What's the weather like?"}]
optional_params = {"parallel_tool_calls": True}
result = mistral_config.transform_request(
model="mistral/mistral-large-latest",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
headers={},
)
assert result.get("parallel_tool_calls") is True
assert len(result["messages"]) == 1
assert result["messages"][0]["role"] == "user"
class TestMistralThinkingContentHandling:
"""Test suite for Mistral thinking content response handling functionality."""
def test_transform_response_with_thinking_content(self):
"""Test that Mistral responses with thinking content are correctly transformed."""
import json
from unittest.mock import Mock
import litellm
# Raw response from Mistral with thinking content
raw_response_data = {
"id": "12a18e1439f24f95b9812a016e0af235",
"choices": [
{
"finish_reason": "stop",
"index": 0,
"logprobs": None,
"message": {
"content": [
{
"type": "thinking",
"thinking": [
{
"type": "text",
"text": "Well, the capital of France is a well-known fact. It's Paris. But just to be sure, I recall that Paris is indeed the capital city of France. I don't need to look it up because it's a common knowledge fact. But if I were unsure, I would double-check using a reliable source or a knowledge base. Since I'm confident about this, I can provide the answer directly.",
}
],
},
{"type": "text", "text": "The capital of France is Paris."},
],
"refusal": None,
"role": "assistant",
"annotations": None,
"audio": None,
"function_call": None,
"tool_calls": None,
},
}
],
"created": 1754654178,
"model": "magistral-medium-2507",
"object": "chat.completion",
"service_tier": None,
"system_fingerprint": None,
"usage": {
"completion_tokens": 93,
"prompt_tokens": 11,
"total_tokens": 104,
"completion_tokens_details": None,
"prompt_tokens_details": None,
},
}
# Mock httpx response
mock_response = Mock()
mock_response.json.return_value = raw_response_data
mock_response.headers = {}
mock_response.text = json.dumps(raw_response_data)
# Mock logging object with proper attributes
mock_logging_obj = Mock()
mock_logging_obj.model_call_details = {}
# Test the transformation
mistral_config = MistralConfig()
model_response = litellm.ModelResponse()
# Test transform_response method
final_response = mistral_config.transform_response(
model="mistral/magistral-medium-2507",
raw_response=mock_response,
model_response=model_response,
logging_obj=mock_logging_obj,
request_data={},
messages=[{"role": "user", "content": "What is the capital of France?"}],
optional_params={},
litellm_params={},
encoding=None,
)
# Verify the response structure
assert final_response is not None
assert len(final_response.choices) == 1
choice = final_response.choices[0]
# Verify message content
message = choice.message
assert message.role == "assistant"
# The content should be processed - either as text or as thinking blocks
# Content could be the text part or the full content list
content_str = str(message.content) if message.content else ""
# Verify the actual text content is preserved somewhere
assert "The capital of France is Paris." in content_str or (
hasattr(message, "thinking_blocks") and message.thinking_blocks
)
# Verify usage information
assert final_response.usage.completion_tokens == 93
assert final_response.usage.prompt_tokens == 11
assert final_response.usage.total_tokens == 104
# Verify model and metadata
assert final_response.id == "12a18e1439f24f95b9812a016e0af235"
assert final_response.created == 1754654178
class TestMistralEmptyContentHandling:
"""Test suite for Mistral empty content response handling functionality."""
@ -419,17 +545,14 @@ class TestMistralEmptyContentHandling:
response_data = {
"choices": [
{
"message": {
"content": "",
"role": "assistant"
},
"finish_reason": "stop"
"message": {"content": "", "role": "assistant"},
"finish_reason": "stop",
}
]
}
result = MistralConfig._handle_empty_content_response(response_data)
assert result["choices"][0]["message"]["content"] is None
def test_handle_empty_content_response_preserves_actual_content(self):
@ -439,41 +562,119 @@ class TestMistralEmptyContentHandling:
{
"message": {
"content": "Hello, how can I help you?",
"role": "assistant"
"role": "assistant",
},
"finish_reason": "stop"
"finish_reason": "stop",
}
]
}
result = MistralConfig._handle_empty_content_response(response_data)
assert result["choices"][0]["message"]["content"] == "Hello, how can I help you?"
assert (
result["choices"][0]["message"]["content"] == "Hello, how can I help you?"
)
def test_handle_empty_content_response_handles_multiple_choices(self):
"""Test that only the first choice is processed for empty content."""
response_data = {
"choices": [
{
"message": {
"content": "",
"role": "assistant"
},
"finish_reason": "stop"
"message": {"content": "", "role": "assistant"},
"finish_reason": "stop",
},
{
"message": {
"content": "",
"role": "assistant"
},
"finish_reason": "stop"
}
"message": {"content": "", "role": "assistant"},
"finish_reason": "stop",
},
]
}
result = MistralConfig._handle_empty_content_response(response_data)
# Only first choice should be converted to None
assert result["choices"][0]["message"]["content"] is None
# Second choice should remain as empty string
assert result["choices"][1]["message"]["content"] is None
assert result["choices"][1]["message"]["content"] is None
def test_is_empty_assistant_message(self):
"""Test that is_empty_assistant_message returns True for empty assistant message."""
message = {"role": "assistant", "content": ""}
assert MistralConfig._is_empty_assistant_message(message) is True
def test_is_empty_assistant_message_with_content(self):
"""Test that is_empty_assistant_message returns False for assistant message with content."""
message = {"role": "assistant", "content": "Hello"}
assert MistralConfig._is_empty_assistant_message(message) is False
class TestMistralFileHandling:
"""Test suite for Mistral file handling functionality."""
def test_handle_file_message_with_file_id(self):
"""Test that file messages with file_id are handled correctly."""
mistral_config = MistralConfig()
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Please review this file."},
{"type": "file", "file": {"file_id": "file-12345"}}
]
}
]
casted_message = cast(list[AllMessageValues], messages)
result = mistral_config._handle_message_with_file(casted_message)
assert len(result) == 1
assert result[0]["role"] == "user"
# Check that content is transformed correctly
assert isinstance(result[0]["content"], list)
assert len(result[0]["content"]) == 2
# Check that file type is preserved
assert result[0]["content"][1]["type"] == "file"
# Check that file_id is modified to match Mistral's expected format
assert result[0]["content"][1]["file_id"] == "file-12345" # type: ignore
def test_handle_file_message_without_file_id(self):
"""Test that file messages without file_id are ignored."""
mistral_config = MistralConfig()
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Please review this file."}
]
}
]
casted_message = cast(list[AllMessageValues], messages)
result = mistral_config._handle_message_with_file(casted_message)
assert len(result) == 1
assert result[0]["role"] == "user"
assert isinstance(result[0]["content"], list)
assert len(result[0]["content"]) == 1 # Only text part remains
def test_handle_message_with_file_multiple_files(self):
"""Test that multiple file messages are handled correctly."""
mistral_config = MistralConfig()
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Please review these files."},
{"type": "file", "file": {"file_id": "file-12345"}},
{"type": "file", "file": {"file_id": "file-67890"}}
]
}
]
casted_message = cast(list[AllMessageValues], messages)
result = mistral_config._handle_message_with_file(casted_message)
assert len(result) == 1
assert result[0]["role"] == "user"
# Check that content is transformed correctly
assert isinstance(result[0]["content"], list)
assert len(result[0]["content"]) == 3 # Text + 2 files
# Check that file types are preserved
assert result[0]["content"][1]["type"] == "file"
assert result[0]["content"][2]["type"] == "file"
# Check that file_ids are modified to match Mistral's expected format
assert result[0]["content"][1]["file_id"] == "file-12345" # type: ignore
assert result[0]["content"][2]["file_id"] == "file-67890" # type: ignore

50
tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
View file

@ -283,6 +283,47 @@ class TestOpenAIResponsesAPIConfig:
assert result.type == "test"
class TestAzureResponsesAPIConfig:
def setup_method(self):
self.config = AzureOpenAIResponsesAPIConfig()
self.model = "gpt-4o"
self.logging_obj = MagicMock()
def test_azure_get_complete_url_with_version_types(self):
"""Test Azure get_complete_url with different API version types"""
base_url = "https://litellm8397336933.openai.azure.com"
# Test with preview version - should use openai/v1/responses
result_preview = self.config.get_complete_url(
api_base=base_url,
litellm_params={"api_version": "preview"},
)
assert (
result_preview
== "https://litellm8397336933.openai.azure.com/openai/v1/responses?api-version=preview"
)
# Test with latest version - should use openai/v1/responses
result_latest = self.config.get_complete_url(
api_base=base_url,
litellm_params={"api_version": "latest"},
)
assert (
result_latest
== "https://litellm8397336933.openai.azure.com/openai/v1/responses?api-version=latest"
)
# Test with date-based version - should use openai/responses
result_date = self.config.get_complete_url(
api_base=base_url,
litellm_params={"api_version": "2025-01-01"},
)
assert (
result_date
== "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2025-01-01"
)
class TestTransformListInputItemsRequest:
"""Test suite for transform_list_input_items_request function"""
@ -618,3 +659,12 @@ class TestTransformListInputItemsRequest:
for key, value in params.items():
assert isinstance(key, str)
assert value is not None
def test_get_supported_openai_params():
config = OpenAIResponsesAPIConfig()
params = config.get_supported_openai_params("gpt-4o")
assert "temperature" in params
assert "stream" in params
assert "background" in params
assert "stream" in params

91
tests/test_litellm/passthrough/test_passthrough_main.py
View file

@ -4,7 +4,10 @@ import sys
import pytest
from fastapi.testclient import TestClient
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from unittest.mock import MagicMock, patch
import httpx
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
@ -45,3 +48,89 @@ def test_llm_passthrough_route():
assert response.status_code == 200
assert response.json == {"message": "Hello, world!"}
def test_bedrock_application_inference_profile_url_encoding():
client = HTTPHandler()
mock_provider_config = MagicMock()
mock_provider_config.get_complete_url.return_value = (
httpx.URL("https://bedrock-runtime.us-east-1.amazonaws.com/model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse"),
"https://bedrock-runtime.us-east-1.amazonaws.com"
)
mock_provider_config.get_api_key.return_value = "test-key"
mock_provider_config.validate_environment.return_value = {}
mock_provider_config.sign_request.return_value = ({}, None)
mock_provider_config.is_streaming_request.return_value = False
with patch("litellm.utils.ProviderConfigManager.get_provider_passthrough_config", return_value=mock_provider_config), \
patch("litellm.litellm_core_utils.get_litellm_params.get_litellm_params", return_value={}), \
patch("litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base")), \
patch.object(client.client, "send", return_value=MagicMock(status_code=200)) as mock_send, \
patch.object(client.client, "build_request") as mock_build_request:
# Mock logging object
mock_logging_obj = MagicMock()
mock_logging_obj.update_environment_variables = MagicMock()
response = llm_passthrough_route(
model="arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd",
endpoint="model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse",
method="POST",
custom_llm_provider="bedrock",
client=client,
litellm_logging_obj=mock_logging_obj,
)
# Verify that build_request was called with the encoded URL
mock_build_request.assert_called_once()
call_args = mock_build_request.call_args
# The URL should have the application-inference-profile ID encoded
actual_url = str(call_args.kwargs["url"])
assert "application-inference-profile%2Fr742sbn2zckd" in actual_url
assert response.status_code == 200
def test_bedrock_non_application_inference_profile_no_encoding():
client = HTTPHandler()
# Mock the provider config and its methods
mock_provider_config = MagicMock()
mock_provider_config.get_complete_url.return_value = (
httpx.URL("https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-sonnet-20240229-v1:0/converse"),
"https://bedrock-runtime.us-east-1.amazonaws.com"
)
mock_provider_config.get_api_key.return_value = "test-key"
mock_provider_config.validate_environment.return_value = {}
mock_provider_config.sign_request.return_value = ({}, None)
mock_provider_config.is_streaming_request.return_value = False
with patch("litellm.utils.ProviderConfigManager.get_provider_passthrough_config", return_value=mock_provider_config), \
patch("litellm.litellm_core_utils.get_litellm_params.get_litellm_params", return_value={}), \
patch("litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base")), \
patch.object(client.client, "send", return_value=MagicMock(status_code=200)) as mock_send, \
patch.object(client.client, "build_request") as mock_build_request:
# Mock logging object
mock_logging_obj = MagicMock()
mock_logging_obj.update_environment_variables = MagicMock()
response = llm_passthrough_route(
model="anthropic.claude-3-sonnet-20240229-v1:0",
endpoint="model/anthropic.claude-3-sonnet-20240229-v1:0/converse",
method="POST",
custom_llm_provider="bedrock",
client=client,
litellm_logging_obj=mock_logging_obj,
)
# Verify that build_request was called with the original URL (no encoding)
mock_build_request.assert_called_once()
call_args = mock_build_request.call_args
# The URL should NOT have application-inference-profile encoding
actual_url = str(call_args.kwargs["url"])
assert "application-inference-profile%2F" not in actual_url
assert "anthropic.claude-3-sonnet-20240229-v1:0" in actual_url
assert response.status_code == 200

67
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -275,3 +275,70 @@ async def test_mcp_server_tool_call_body_with_none_arguments():
body = captured_data["proxy_server_request"]["body"]
assert body["name"] == tool_name
assert body["arguments"] == tool_arguments # Should be None
@pytest.mark.asyncio
async def test_concurrent_initialize_session_managers():
"""Test that concurrent calls to initialize_session_managers don't cause race conditions."""
try:
from litellm.proxy._experimental.mcp_server.server import (
initialize_session_managers,
_SESSION_MANAGERS_INITIALIZED,
_INITIALIZATION_LOCK,
)
except ImportError:
pytest.skip("MCP server not available")
# Import the module to reset state
import litellm.proxy._experimental.mcp_server.server as mcp_server
# Reset state before test
original_initialized = mcp_server._SESSION_MANAGERS_INITIALIZED
original_session_cm = mcp_server._session_manager_cm
original_sse_session_cm = mcp_server._sse_session_manager_cm
try:
mcp_server._SESSION_MANAGERS_INITIALIZED = False
mcp_server._session_manager_cm = None
mcp_server._sse_session_manager_cm = None
# Mock the session managers to avoid actual MCP initialization
with patch('litellm.proxy._experimental.mcp_server.server.session_manager') as mock_session_manager, \
patch('litellm.proxy._experimental.mcp_server.server.sse_session_manager') as mock_sse_session_manager, \
patch('litellm.proxy._experimental.mcp_server.server.verbose_logger'):
# Mock the run() method to return a mock context manager
mock_cm = AsyncMock()
mock_cm.__aenter__ = AsyncMock()
mock_cm.__aexit__ = AsyncMock()
mock_session_manager.run.return_value = mock_cm
mock_sse_session_manager.run.return_value = mock_cm
# Create multiple concurrent tasks that call initialize_session_managers
async def init_task():
await initialize_session_managers()
return "success"
# Run 10 concurrent initialization attempts
tasks = [init_task() for _ in range(10)]
results = await asyncio.gather(*tasks, return_exceptions=True)
# All tasks should complete successfully (no exceptions)
assert all(result == "success" for result in results), f"Some tasks failed: {results}"
# session_manager.run() should only be called once due to the lock
assert mock_session_manager.run.call_count == 1, f"Expected 1 call to session_manager.run(), got {mock_session_manager.run.call_count}"
assert mock_sse_session_manager.run.call_count == 1, f"Expected 1 call to sse_session_manager.run(), got {mock_sse_session_manager.run.call_count}"
# The context managers should only be entered once each
assert mock_cm.__aenter__.call_count == 2, f"Expected 2 calls to __aenter__ (one for each session manager), got {mock_cm.__aenter__.call_count}"
# State should be properly set
assert mcp_server._SESSION_MANAGERS_INITIALIZED is True
finally:
# Restore original state
mcp_server._SESSION_MANAGERS_INITIALIZED = original_initialized
mcp_server._session_manager_cm = original_session_cm
mcp_server._sse_session_manager_cm = original_sse_session_cm

178
tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
View file

@ -300,3 +300,181 @@ def test_optional_params_returned_when_properly_overridden():
print("FIELDS", fields)
assert "optional_params" in fields
@pytest.mark.asyncio
async def test_bedrock_guardrail_prepare_request_with_api_key():
"""Test _prepare_request method uses Bearer token when api_key is provided in data"""
from unittest.mock import Mock, patch
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
# Setup guardrail hook
guardrail_hook = BedrockGuardrail(
guardrailIdentifier="test-guardrail-id",
guardrailVersion="1"
)
mock_credentials = Mock()
test_data = {
"source": "INPUT",
"content": [{"text": {"text": "test content"}}]
}
prepared_request = guardrail_hook._prepare_request(
credentials=mock_credentials,
data=test_data,
optional_params={},
aws_region_name="us-east-1",
api_key="test-bearer-token-123"
)
# Verify Bearer token is used in Authorization header
assert "Authorization" in prepared_request.headers
assert prepared_request.headers["Authorization"] == "Bearer test-bearer-token-123"
# Verify URL is correct
expected_url = "https://bedrock-runtime.us-east-1.amazonaws.com/guardrail/test-guardrail-id/version/1/apply"
assert prepared_request.url == expected_url
@pytest.mark.asyncio
async def test_bedrock_guardrail_prepare_request_without_api_key():
"""Test _prepare_request method falls back to SigV4 when no api_key is provided"""
from unittest.mock import Mock, patch
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
# Setup guardrail hook
guardrail_hook = BedrockGuardrail(
guardrailIdentifier="test-guardrail-id",
guardrailVersion="1"
)
# Mock credentials
mock_credentials = Mock()
# Test data without api_key
test_data = {
"source": "INPUT",
"content": [{"text": {"text": "test content"}}]
}
with patch("litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str") as mock_get_secret, \
patch("botocore.auth.SigV4Auth") as mock_sigv4_auth, \
patch("botocore.awsrequest.AWSRequest") as mock_aws_request:
# Mock no AWS_BEARER_TOKEN_BEDROCK
mock_get_secret.return_value = None
# Mock SigV4Auth
mock_sigv4_instance = Mock()
mock_sigv4_auth.return_value = mock_sigv4_instance
# Mock AWSRequest
mock_request_instance = Mock()
mock_request_instance.prepare.return_value = Mock()
mock_aws_request.return_value = mock_request_instance
# Call _prepare_request
prepared_request = guardrail_hook._prepare_request(
credentials=mock_credentials,
data=test_data,
optional_params={},
aws_region_name="us-east-1"
)
# Verify SigV4 auth was used
mock_sigv4_auth.assert_called_once_with(mock_credentials, "bedrock", "us-east-1")
mock_sigv4_instance.add_auth.assert_called_once()
@pytest.mark.asyncio
async def test_bedrock_guardrail_prepare_request_with_bearer_token_env():
"""Test _prepare_request method uses Bearer token from environment when available"""
from unittest.mock import Mock, patch
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
# Setup guardrail hook
guardrail_hook = BedrockGuardrail(
guardrailIdentifier="test-guardrail-id",
guardrailVersion="1"
)
# Mock credentials
mock_credentials = Mock()
# Test data without api_key
test_data = {
"source": "INPUT",
"content": [{"text": {"text": "test content"}}]
}
with patch("litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str") as mock_get_secret, \
patch("botocore.awsrequest.AWSRequest") as mock_aws_request:
mock_get_secret.return_value = "env-bearer-token-456"
mock_request_instance = Mock()
mock_request_instance.prepare.return_value = Mock()
mock_aws_request.return_value = mock_request_instance
prepared_request = guardrail_hook._prepare_request(
credentials=mock_credentials,
data=test_data,
optional_params={},
aws_region_name="us-east-1"
)
# Verify Bearer token from environment is used
mock_aws_request.assert_called_once()
call_args = mock_aws_request.call_args
headers = call_args[1]["headers"]
assert headers["Authorization"] == "Bearer env-bearer-token-456"
@pytest.mark.asyncio
async def test_bedrock_guardrail_make_api_request_passes_api_key():
"""Test make_bedrock_api_request method correctly passes api_key from request_data"""
from unittest.mock import Mock, patch, AsyncMock
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
guardrail_hook = BedrockGuardrail(
guardrailIdentifier="test-guardrail-id",
guardrailVersion="1"
)
guardrail_hook.async_handler = Mock()
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {"action": "NONE", "outputs": []}
guardrail_hook.async_handler.post = AsyncMock(return_value=mock_response)
test_request_data = {
"api_key": "test-api-key-789"
}
with patch.object(guardrail_hook, "_load_credentials") as mock_load_creds, \
patch.object(guardrail_hook, "convert_to_bedrock_format") as mock_convert, \
patch.object(guardrail_hook, "get_guardrail_dynamic_request_body_params") as mock_get_params, \
patch.object(guardrail_hook, "add_standard_logging_guardrail_information_to_request_data"), \
patch("botocore.awsrequest.AWSRequest") as mock_aws_request:
mock_load_creds.return_value = (Mock(), "us-east-1")
mock_convert.return_value = {"source": "INPUT", "content": []}
mock_get_params.return_value = {}
mock_request_instance = Mock()
mock_request_instance.url = "test-url"
mock_request_instance.body = b"test-body"
mock_request_instance.headers = {"Content-Type": "application/json", "Authorization": "Bearer test-api-key-789"}
mock_request_instance.prepare.return_value = Mock()
mock_aws_request.return_value = mock_request_instance
await guardrail_hook.make_bedrock_api_request(
source="INPUT",
messages=[{"role": "user", "content": "test"}],
request_data=test_request_data
)
# Verify _prepare_request was invoked and used the api_key
mock_aws_request.assert_called_once()
call_args = mock_aws_request.call_args
headers = call_args[1]["headers"]
assert headers["Authorization"] == "Bearer test-api-key-789"

151
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -724,6 +724,156 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3():
), "should_rate_limit should be called when model-specific limits match requested model"
@pytest.mark.asyncio
async def test_tpm_api_key_rate_limits_v3():
_api_key = "sk-12345"
_api_key_hash = hash_token(_api_key)
model = "gpt-3.5-turbo"
rpm_limit = 2
tpm_limit = 2
rpms = {model: rpm_limit}
tpms = {model: tpm_limit}
user_api_key_dict = UserAPIKeyAuth(
api_key=_api_key_hash,
key_alias=_api_key,
rpm_limit_per_model=rpms,
tpm_limit_per_model=tpms,
models=[],
)
user_api_key_dict.metadata["model_tpm_limit"] = tpms
user_api_key_dict.metadata["model_rpm_limit"] = rpms
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
# Mock should_rate_limit to capture the descriptors
captured_descriptors = None
original_should_rate_limit = parallel_request_handler.should_rate_limit
async def mock_should_rate_limit(descriptors, **kwargs):
nonlocal captured_descriptors
captured_descriptors = descriptors
# Return Error response to ensure HTTPException
return {
"overall_code": "OVER_LIMIT",
"statuses": [{'code': 'OK', 'current_limit': 2, 'limit_remaining': 1, 'rate_limit_type': 'requests', 'descriptor_key': 'model_per_key'},
{'code': 'OVER_LIMIT', 'current_limit': 2, 'limit_remaining': -18, 'rate_limit_type': 'tokens', 'descriptor_key': 'model_per_key'}]
}
parallel_request_handler.should_rate_limit = mock_should_rate_limit
# Test the pre-call hook
error = None
try:
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": model},
call_type="",
)
except HTTPException as e:
error=e
assert e.status_code == 429
assert "rate_limit_type" in e.headers
assert e.headers.get("rate_limit_type") == "tokens"
assert "retry-after" in e.headers
assert error is not None, "An Exception must be thrown"
assert captured_descriptors is not None, "Rate limit descriptors should be captured"
model_per_key_descriptor = None
for descriptor in captured_descriptors:
if descriptor["key"] == "model_per_key":
model_per_key_descriptor = descriptor
break
assert model_per_key_descriptor is not None, "Api-Key descriptor should be present"
assert model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}", "Api-Key value should combine api_key and model"
assert model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit, "Api-Key RPM limit should be set"
assert model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit, "Api-Key TPM limit should be set"
@pytest.mark.asyncio
async def test_rpm_api_key_rate_limits_v3():
_api_key = "sk-12345"
_api_key_hash = hash_token(_api_key)
model = "gpt-3.5-turbo"
rpm_limit = 2
tpm_limit = 2
rpms = {model: rpm_limit}
tpms = {model: tpm_limit}
user_api_key_dict = UserAPIKeyAuth(
api_key=_api_key_hash,
key_alias=_api_key,
rpm_limit_per_model=rpms,
tpm_limit_per_model=tpms,
models=[],
)
user_api_key_dict.metadata["model_tpm_limit"] = tpms
user_api_key_dict.metadata["model_rpm_limit"] = rpms
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
# Mock should_rate_limit to capture the descriptors
captured_descriptors = None
original_should_rate_limit = parallel_request_handler.should_rate_limit
async def mock_should_rate_limit(descriptors, **kwargs):
nonlocal captured_descriptors
captured_descriptors = descriptors
# Return Error response to ensure HTTPException
return {
"overall_code": "OVER_LIMIT",
"statuses": [{'code': 'OVER_LIMIT', 'current_limit': 2, 'limit_remaining': -2, 'rate_limit_type': 'requests', 'descriptor_key': 'model_per_key'},
{'code': 'OK', 'current_limit': 2, 'limit_remaining': 2, 'rate_limit_type': 'tokens', 'descriptor_key': 'model_per_key'}]
}
parallel_request_handler.should_rate_limit = mock_should_rate_limit
# Test the pre-call hook
error = None
try:
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": model},
call_type="",
)
except HTTPException as e:
error=e
assert e.status_code == 429
assert "rate_limit_type" in e.headers
assert e.headers.get("rate_limit_type") == "requests"
assert "retry-after" in e.headers
assert error is not None, "An Exception must be thrown"
assert captured_descriptors is not None, "Rate limit descriptors should be captured"
model_per_key_descriptor = None
for descriptor in captured_descriptors:
if descriptor["key"] == "model_per_key":
model_per_key_descriptor = descriptor
break
assert model_per_key_descriptor is not None, "Api-Key descriptor should be present"
assert model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}", "Api-Key value should combine api_key and model"
assert model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit, "Api-Key RPM limit should be set"
assert model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit, "Api-Key TPM limit should be set"
@pytest.mark.asyncio
async def test_team_member_rate_limits_v3():
"""
@ -763,6 +913,7 @@ async def test_team_member_rate_limits_v3():
parallel_request_handler.should_rate_limit = mock_should_rate_limit
# Test the pre-call hook
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,

86
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -11,9 +11,20 @@ sys.path.insert(
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import HTTPException
from litellm.proxy.management_endpoints.key_management_endpoints import _list_key_helper
from litellm.proxy._types import (
GenerateKeyRequest,
LiteLLM_VerificationToken,
LitellmUserRoles,
UpdateKeyRequest,
)
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.management_endpoints.key_management_endpoints import (
_common_key_generation_helper,
_list_key_helper,
prepare_key_update_data,
)
from litellm.proxy.proxy_server import app
client = TestClient(app)
@ -492,3 +503,74 @@ def test_get_new_token_with_invalid_key():
assert exc_info.value.status_code == 400
assert "New key must start with 'sk-'" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_generate_service_account_requires_team_id():
with pytest.raises(HTTPException):
await _common_key_generation_helper(
data=GenerateKeyRequest(
metadata={"service_account_id": "sa"},
team_id=None,
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"
),
litellm_changed_by=None,
team_table=None,
)
@pytest.mark.asyncio
async def test_generate_service_account_works_with_team_id():
from unittest.mock import patch
# Mock the database and router dependencies from proxy_server
with patch('litellm.proxy.proxy_server.prisma_client') as mock_prisma, \
patch('litellm.proxy.proxy_server.llm_router') as mock_router, \
patch('litellm.proxy.proxy_server.premium_user', False), \
patch('litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn') as mock_generate_key:
# Configure mocks
mock_prisma.return_value = AsyncMock()
mock_router.return_value = None
# Mock the response from generate_key_helper_fn
mock_generate_key.return_value = {
"key": "sk-test-key",
"expires": None,
"user_id": "test-user",
"team_id": "IJ"
}
# This should not raise an exception since team_id is provided
await _common_key_generation_helper(
data=GenerateKeyRequest(
metadata={"service_account_id": "sa"},
team_id="IJ",
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"
),
litellm_changed_by=None,
team_table=None,
)
@pytest.mark.asyncio
async def test_update_service_account_requires_team_id():
data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"})
existing_key = LiteLLM_VerificationToken(token="hashed", team_id=None)
with pytest.raises(HTTPException):
await prepare_key_update_data(data=data, existing_key_row=existing_key)
@pytest.mark.asyncio
async def test_update_service_account_works_with_team_id():
data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"}, team_id="IJ")
existing_key = LiteLLM_VerificationToken(token="hashed")
await prepare_key_update_data(data=data, existing_key_row=existing_key)

61
tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -21,6 +21,7 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
create_pass_through_route,
vertex_discovery_proxy_route,
vertex_proxy_route,
bedrock_llm_proxy_route,
)
from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials
@ -853,3 +854,63 @@ async def test_is_streaming_request_fn():
mock_request.headers = {"content-type": "multipart/form-data"}
mock_request.form = AsyncMock(return_value={"stream": "true"})
assert await is_streaming_request_fn(mock_request) is True
class TestBedrockLLMProxyRoute:
@pytest.mark.asyncio
async def test_bedrock_llm_proxy_route_application_inference_profile(self):
mock_request = Mock()
mock_request.method = "POST"
mock_response = Mock()
mock_user_api_key_dict = Mock()
mock_request_body = {"messages": [{"role": "user", "content": "test"}]}
mock_processor = Mock()
mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success")
with patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._read_request_body", return_value=mock_request_body), \
patch("litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", return_value=mock_processor):
# Test application-inference-profile endpoint
endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/r742sbn2zckd/converse"
result = await bedrock_llm_proxy_route(
endpoint=endpoint,
request=mock_request,
fastapi_response=mock_response,
user_api_key_dict=mock_user_api_key_dict,
)
mock_processor.base_passthrough_process_llm_request.assert_called_once()
call_kwargs = mock_processor.base_passthrough_process_llm_request.call_args.kwargs
# For application-inference-profile, model should be "arn:aws:bedrock:us-east-1:026090525607:application-inference-profile/r742sbn2zckd"
assert call_kwargs["model"] == "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/r742sbn2zckd"
assert result == "success"
@pytest.mark.asyncio
async def test_bedrock_llm_proxy_route_regular_model(self):
mock_request = Mock()
mock_request.method = "POST"
mock_response = Mock()
mock_user_api_key_dict = Mock()
mock_request_body = {"messages": [{"role": "user", "content": "test"}]}
mock_processor = Mock()
mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success")
with patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._read_request_body", return_value=mock_request_body), \
patch("litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", return_value=mock_processor):
# Test regular model endpoint
endpoint = "model/anthropic.claude-3-sonnet-20240229-v1:0/converse"
result = await bedrock_llm_proxy_route(
endpoint=endpoint,
request=mock_request,
fastapi_response=mock_response,
user_api_key_dict=mock_user_api_key_dict,
)
mock_processor.base_passthrough_process_llm_request.assert_called_once()
call_kwargs = mock_processor.base_passthrough_process_llm_request.call_args.kwargs
# For regular models, model should be just the model ID
assert call_kwargs["model"] == "anthropic.claude-3-sonnet-20240229-v1:0"
assert result == "success"

Some files were not shown because too many files have changed in this diff Show more