diff --git a/.gitignore b/.gitignore index a58dae81ead..ed8c88c8990 100644 --- a/.gitignore +++ b/.gitignore @@ -94,4 +94,5 @@ test.py litellm_config.yaml .cursor -.vscode/launch.json \ No newline at end of file +.vscode/launch.json +litellm/proxy/to_delete_loadtest_work/* \ No newline at end of file diff --git a/cookbook/liteLLM_Baseten.ipynb b/cookbook/liteLLM_Baseten.ipynb index e03bb3254a5..0a5bc5f1df7 100644 --- a/cookbook/liteLLM_Baseten.ipynb +++ b/cookbook/liteLLM_Baseten.ipynb @@ -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 -} \ No newline at end of file +} diff --git a/deploy/charts/litellm-helm/README.md b/deploy/charts/litellm-helm/README.md index cef2b8d162d..73b722b64c6 100644 --- a/deploy/charts/litellm-helm/README.md +++ b/deploy/charts/litellm-helm/README.md @@ -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://-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 `-litellm-masterkey` Kubernetes Secret. +generated string in the `sk-...` format stored in the `-litellm-masterkey` Kubernetes Secret. ```bash kubectl -n litellm get secret -litellm-masterkey -o jsonpath="{.data.masterkey}" diff --git a/deploy/charts/litellm-helm/templates/deployment.yaml b/deploy/charts/litellm-helm/templates/deployment.yaml index 71412b8052b..b30b8829325 100644 --- a/deploy/charts/litellm-helm/templates/deployment.yaml +++ b/deploy/charts/litellm-helm/templates/deployment.yaml @@ -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 diff --git a/deploy/charts/litellm-helm/templates/migrations-job.yaml b/deploy/charts/litellm-helm/templates/migrations-job.yaml index ec80e86b21d..4c8925564af 100644 --- a/deploy/charts/litellm-helm/templates/migrations-job.yaml +++ b/deploy/charts/litellm-helm/templates/migrations-job.yaml @@ -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 diff --git a/deploy/charts/litellm-helm/templates/secret-masterkey.yaml b/deploy/charts/litellm-helm/templates/secret-masterkey.yaml index 5632957dc05..7c8560cc2cc 100644 --- a/deploy/charts/litellm-helm/templates/secret-masterkey.yaml +++ b/deploy/charts/litellm-helm/templates/secret-masterkey.yaml @@ -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: diff --git a/deploy/charts/litellm-helm/tests/masterkey-secret_tests.yaml b/deploy/charts/litellm-helm/tests/masterkey-secret_tests.yaml index eb1d3c3967f..bbbade9d802 100644 --- a/deploy/charts/litellm-helm/tests/masterkey-secret_tests.yaml +++ b/deploy/charts/litellm-helm/tests/masterkey-secret_tests.yaml @@ -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: diff --git a/deploy/charts/litellm-helm/values.yaml b/deploy/charts/litellm-helm/values.yaml index e9a96b23de7..0bd95003c10 100644 --- a/deploy/charts/litellm-helm/values.yaml +++ b/deploy/charts/litellm-helm/values.yaml @@ -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 diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index c24c82b2898..bc2bf169c28 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -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 diff --git a/docs/my-website/docs/completion/document_understanding.md b/docs/my-website/docs/completion/document_understanding.md index b831a7b9da2..172e0792801 100644 --- a/docs/my-website/docs/completion/document_understanding.md +++ b/docs/my-website/docs/completion/document_understanding.md @@ -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' \ +## Mistral Example + +Here is a sample payload for using the Mistral model for document understanding: + + + + + +```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 +``` + + + + +```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" + } + } + ] + } + ] +} +``` + + + ## Checking if a model supports pdf input diff --git a/docs/my-website/docs/providers/baseten.md b/docs/my-website/docs/providers/baseten.md index 902b1548faa..4e42cdf0447 100644 --- a/docs/my-website/docs/providers/baseten.md +++ b/docs/my-website/docs/providers/baseten.md @@ -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/` (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/` +# 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!"}] +) +``` diff --git a/docs/my-website/docs/providers/deepinfra.md b/docs/my-website/docs/providers/deepinfra.md index 1360117445f..ddf6122cac8 100644 --- a/docs/my-website/docs/providers/deepinfra.md +++ b/docs/my-website/docs/providers/deepinfra.md @@ -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 + + + + +```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) +``` + + + + +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." + ] +}' +``` + + + + +### 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 + } + } +} +``` diff --git a/docs/my-website/docs/providers/vertex_partner.md b/docs/my-website/docs/providers/vertex_partner.md index c6e324f2958..cf780e35dbd 100644 --- a/docs/my-website/docs/providers/vertex_partner.md +++ b/docs/my-website/docs/providers/vertex_partner.md @@ -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' \ +## 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 + + + + +```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) +``` + + + +**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" + } + ], + }' +``` + + + + + ## Model Garden :::tip diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index a1288a47a79..73b16af2f95 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -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 diff --git a/docs/my-website/docs/proxy/team_budgets.md b/docs/my-website/docs/proxy/team_budgets.md index 854d6edf304..66ba679c65e 100644 --- a/docs/my-website/docs/proxy/team_budgets.md +++ b/docs/my-website/docs/proxy/team_budgets.md @@ -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 diff --git a/docs/my-website/docs/proxy/users.md b/docs/my-website/docs/proxy/users.md index c812dccb199..d098e38de4a 100644 --- a/docs/my-website/docs/proxy/users.md +++ b/docs/my-website/docs/proxy/users.md @@ -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. diff --git a/docs/my-website/docs/rerank.md b/docs/my-website/docs/rerank.md index 11dcae777e4..c57eacbb224 100644 --- a/docs/my-website/docs/rerank.md +++ b/docs/my-website/docs/rerank.md @@ -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) | \ No newline at end of file +| vLLM| [Usage](../docs/providers/vllm#rerank-endpoint) | +| DeepInfra| [Usage](../docs/providers/deepinfra#rerank-endpoint) | \ No newline at end of file diff --git a/docs/my-website/docs/response_api.md b/docs/my-website/docs/response_api.md index c0bb003c096..94d7c73be05 100644 --- a/docs/my-website/docs/response_api.md +++ b/docs/my-website/docs/response_api.md @@ -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) diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py index d17946171bb..e60b4d69905 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py @@ -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() diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index c194c49765c..ece2b496bf6 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -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( diff --git a/litellm/__init__.py b/litellm/__init__.py index 8ec0c812100..9ac442cb3d2 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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 diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 0a8b6ef18d1..82fc37e0cb4 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -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 diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 7580752c307..1dcc0f1fdb2 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -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( diff --git a/litellm/caching/s3_cache.py b/litellm/caching/s3_cache.py index 6095d547aab..15f7a5c1e16 100644 --- a/litellm/caching/s3_cache.py +++ b/litellm/caching/s3_cache.py @@ -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 diff --git a/litellm/constants.py b/litellm/constants.py index afdb2073627..365dfa84510 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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": { diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 8149a6131e8..5cf403adee4 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -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, diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py index 634d0c1fdc9..86af800d732 100644 --- a/litellm/integrations/mlflow.py +++ b/litellm/integrations/mlflow.py @@ -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} diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index e9a1eba719e..02b542ed165 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -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 diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index 5fcd2ddb70a..6584defe5b7 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -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": diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 385df83c904..7bc7702684d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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: diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index c7512ea146b..bf43519afc6 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -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 diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0d8c3bacbf5..77cbe4c9a8e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -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( diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 063d1af9c33..488a711669d 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -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: diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index e2f89da5e86..4da4f7652e0 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -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 { diff --git a/litellm/llms/baseten.py b/litellm/llms/baseten.py deleted file mode 100644 index e1d513d6d11..00000000000 --- a/litellm/llms/baseten.py +++ /dev/null @@ -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 diff --git a/litellm/llms/baseten/chat.py b/litellm/llms/baseten/chat.py new file mode 100644 index 00000000000..05fc9961ac5 --- /dev/null +++ b/litellm/llms/baseten/chat.py @@ -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" \ No newline at end of file diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index cc205e62dc9..55aad673ced 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -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 diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 3ed7d04bde6..ab69ea1f8c3 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -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 diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index cf2187153a9..4d8781fff2a 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -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 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3e7dff91820..c9d70088d04 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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( diff --git a/litellm/llms/datarobot/chat/transformation.py b/litellm/llms/datarobot/chat/transformation.py index e334c94e517..23ce63c25b2 100644 --- a/litellm/llms/datarobot/chat/transformation.py +++ b/litellm/llms/datarobot/chat/transformation.py @@ -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`` diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py new file mode 100644 index 00000000000..8259c6075bb --- /dev/null +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -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, + ) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index b38a4982471..51fa65244a0 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -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, diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 3565f090b04..4f7be507cc2 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -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()) diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 396d59145ff..be0ca3a7086 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -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, diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index c8a1e8f0e1c..b638f185dc6 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -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, diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 501941fdc59..e70cadddaf7 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -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 diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py index 7e965313a0b..748a5f5fb40 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py @@ -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"], diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index 7303ab0786c..f281cab3b58 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -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, diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py new file mode 100644 index 00000000000..4df2fa4ba31 --- /dev/null +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -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() diff --git a/litellm/main.py b/litellm/main.py index 4a8f82e114d..c8442d483e9 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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 diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 55c42d928fa..202e7f0759f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -9723,6 +9723,30 @@ "supports_tool_choice": true, "supports_prompt_caching": true }, + "vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas": { + "max_tokens": 32768, + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 4e-06, + "litellm_provider": "vertex_ai-qwen_models", + "mode": "chat", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { + "max_tokens": 16384, + "max_input_tokens": 262144, + "max_output_tokens": 16384, + "input_cost_per_token": 0.25e-06, + "output_cost_per_token": 1e-06, + "litellm_provider": "vertex_ai-qwen_models", + "mode": "chat", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supports_function_calling": true, + "supports_tool_choice": true + }, "vertex_ai/meta/llama3-405b-instruct-maas": { "max_tokens": 32000, "max_input_tokens": 32000, @@ -10072,6 +10096,12 @@ "mode": "image_generation", "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "vertex_ai/imagen-4.0-generate-001": { + "output_cost_per_image": 0.04, + "litellm_provider": "vertex_ai-image-models", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, "vertex_ai/imagen-4.0-generate-preview-06-06": { "output_cost_per_image": 0.04, "litellm_provider": "vertex_ai-image-models", @@ -10084,6 +10114,18 @@ "mode": "image_generation", "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "vertex_ai/imagen-4.0-ultra-generate-001": { + "output_cost_per_image": 0.06, + "litellm_provider": "vertex_ai-image-models", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, + "vertex_ai/imagen-4.0-fast-generate-001": { + "output_cost_per_image": 0.02, + "litellm_provider": "vertex_ai-image-models", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, "vertex_ai/imagen-4.0-fast-generate-preview-06-06": { "output_cost_per_image": 0.02, "litellm_provider": "vertex_ai-image-models", @@ -10790,18 +10832,36 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", "supports_tool_choice": true }, + "gemini/imagen-4.0-generate-001": { + "output_cost_per_image": 0.04, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, "gemini/imagen-4.0-generate-preview-06-06": { "output_cost_per_image": 0.04, "litellm_provider": "gemini", "mode": "image_generation", "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "gemini/imagen-4.0-ultra-generate-001": { + "output_cost_per_image": 0.06, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, "gemini/imagen-4.0-ultra-generate-preview-06-06": { "output_cost_per_image": 0.06, "litellm_provider": "gemini", "mode": "image_generation", "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "gemini/imagen-4.0-fast-generate-001": { + "output_cost_per_image": 0.02, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, "gemini/imagen-4.0-fast-generate-preview-06-06": { "output_cost_per_image": 0.02, "litellm_provider": "gemini", @@ -11176,6 +11236,21 @@ "supports_tool_choice": true, "supports_prompt_caching": true }, + "openrouter/deepseek/deepseek-chat-v3.1": { + "max_tokens": 8192, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 2e-07, + "input_cost_per_token_cache_hit": 2e-08, + "output_cost_per_token": 8e-07, + "litellm_provider": "openrouter", + "mode": "chat", + "supports_function_calling": true, + "supports_assistant_prefill": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_prompt_caching": true + }, "openrouter/x-ai/grok-4": { "max_tokens": 256000, "max_input_tokens": 256000, @@ -15020,95 +15095,1036 @@ "litellm_provider": "ollama", "mode": "completion" }, - "deepinfra/lizpreciatior/lzlv_70b_fp16_hf": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 7e-07, - "output_cost_per_token": 9e-07, + "deepinfra/deepseek-ai/DeepSeek-V3": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 3.8e-07, + "output_cost_per_token": 8.9e-07, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, + "deepinfra/Phind/Phind-CodeLlama-34B-v2": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1.5e-08, + "output_cost_per_token": 2e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/gemma-2-9b-it": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3e-08, + "output_cost_per_token": 6e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/deepseek-ai/DeepSeek-R1-0528-Turbo": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen2-7B-Instruct": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/QVQ-72B-Preview": { + "max_tokens": 32000, + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-3.3-70B-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 2.3e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/microsoft/Phi-4-multimodal-instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 5e-08, + "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/mistralai/Devstral-Small-2507": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/microsoft/WizardLM-2-7B": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-3.2-90B-Vision-Instruct": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/mistralai/Mistral-Small-3.2-24B-Instruct-2506": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 5e-08, + "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-V3-0324": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 2.8e-07, + "output_cost_per_token": 8.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/mistralai/Mixtral-8x7B-Instruct-v0.1": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 8e-08, + "output_cost_per_token": 2.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/anthropic/claude-3-7-sonnet-latest": { + "max_tokens": 200000, + "max_input_tokens": 200000, + "max_output_tokens": 200000, + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/cognitivecomputations/dolphin-2.9.1-llama-3-70b": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen2.5-Coder-32B-Instruct": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 6e-08, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen3-235B-A22B": { + "max_tokens": 40960, + "max_input_tokens": 40960, + "max_output_tokens": 40960, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-V3-0324-Turbo": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/microsoft/WizardLM-2-8x22B": { + "max_tokens": 65536, + "max_input_tokens": 65536, + "max_output_tokens": 65536, + "input_cost_per_token": 4.8e-07, + "output_cost_per_token": 4.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Sao10K/L3-8B-Lunaris-v1-Turbo": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-Guard-4-12B": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 1.8e-07, + "output_cost_per_token": 1.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, "deepinfra/Gryphe/MythoMax-L2-13b": { "max_tokens": 4096, "max_input_tokens": 4096, "max_output_tokens": 4096, - "input_cost_per_token": 2.2e-07, - "output_cost_per_token": 2.2e-07, + "input_cost_per_token": 7.2e-08, + "output_cost_per_token": 7.2e-08, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, - "deepinfra/mistralai/Mistral-7B-Instruct-v0.1": { - "max_tokens": 8191, + "deepinfra/meta-llama/Llama-3.2-1B-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 5e-09, + "output_cost_per_token": 1e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/gemma-2-27b-it": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 2.7e-07, + "output_cost_per_token": 2.7e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen2.5-VL-32B-Instruct": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/nvidia/Llama-3.1-Nemotron-70B-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/mistralai/Mixtral-8x22B-Instruct-v0.1": { + "max_tokens": 65536, + "max_input_tokens": 65536, + "max_output_tokens": 65536, + "input_cost_per_token": 6.5e-07, + "output_cost_per_token": 6.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen2.5-7B-Instruct": { + "max_tokens": 32768, "max_input_tokens": 32768, - "max_output_tokens": 8191, + "max_output_tokens": 32768, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/google/gemini-1.5-flash-8b": { + "max_tokens": 1000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "input_cost_per_token": 3.75e-08, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/NousResearch/Hermes-3-Llama-3.1-70B": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 2.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-Guard-3-8B": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/mistralai/Mistral-Small-24B-Instruct-2501": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 5e-08, + "output_cost_per_token": 8e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/lizpreciatior/lzlv_70b_fp16_hf": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-2-13b-chat-hf": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, "input_cost_per_token": 1.3e-07, "output_cost_per_token": 1.3e-07, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, + "deepinfra/anthropic/claude-4-opus": { + "max_tokens": 200000, + "max_input_tokens": 200000, + "max_output_tokens": 200000, + "input_cost_per_token": 1.65e-05, + "output_cost_per_token": 8.25e-05, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/openchat/openchat-3.6-8b": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/google/gemma-3-27b-it": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 1.7e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Austism/chronos-hermes-13b-v2": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 1.3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Sao10K/L3.1-70B-Euryale-v2.2": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 6.5e-07, + "output_cost_per_token": 7.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/QwQ-32B-Preview": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 1.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/anthropic/claude-4-sonnet": { + "max_tokens": 200000, + "max_input_tokens": 200000, + "max_output_tokens": 200000, + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/microsoft/Phi-3-medium-4k-instruct": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 1.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/mattshumer/Reflection-Llama-3.1-70B": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/openchat/openchat_3.5": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Sao10K/L3.3-70B-Euryale-v2.3": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 6.5e-07, + "output_cost_per_token": 7.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 2.3e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-V3.1": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 3.2e-07, + "output_cost_per_token": 1.15e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen2.5-Coder-7B": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 2.5e-08, + "output_cost_per_token": 5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/cognitivecomputations/dolphin-2.6-mixtral-8x7b": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 2.4e-07, + "output_cost_per_token": 2.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Meta-Llama-3.1-405B-Instruct": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 8e-07, + "output_cost_per_token": 8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-Prover-V2-671B": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.18e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/zai-org/GLM-4.5": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 5.5e-07, + "output_cost_per_token": 2e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Llama-3.2-3B-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1.2e-08, + "output_cost_per_token": 2.4e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Meta-Llama-3-70B-Instruct": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/gemini-1.5-flash": { + "max_tokens": 1000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "input_cost_per_token": 7.5e-08, + "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/KoboldAI/LLaMA2-13B-Tiefighter": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/gemini-2.5-pro": { + "max_tokens": 1000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "input_cost_per_token": 8.75e-07, + "output_cost_per_token": 7e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-30B-A3B": { + "max_tokens": 40960, + "max_input_tokens": 40960, + "max_output_tokens": 40960, + "input_cost_per_token": 8e-08, + "output_cost_per_token": 2.9e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/QwQ-32B": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 7.5e-08, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/moonshotai/Kimi-K2-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Sao10K/L3-70B-Euryale-v2.1": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 7e-07, + "output_cost_per_token": 8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/microsoft/phi-4-reasoning-plus": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 3.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/google/gemma-3-12b-it": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 5e-08, + "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/gemini-2.5-flash": { + "max_tokens": 1000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "input_cost_per_token": 2.1e-07, + "output_cost_per_token": 1.75e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-R1": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 4.5e-07, + "output_cost_per_token": 2.15e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/mistralai/Mistral-7B-Instruct-v0.3": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 2.8e-08, + "output_cost_per_token": 5.4e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen2.5-72B-Instruct": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 3.9e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-14B": { + "max_tokens": 40960, + "max_input_tokens": 40960, + "max_output_tokens": 40960, + "input_cost_per_token": 6e-08, + "output_cost_per_token": 2.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/allenai/olmOCR-7B-0725-FP8": { + "max_tokens": 16384, + "max_input_tokens": 16384, + "max_output_tokens": 16384, + "input_cost_per_token": 2.7e-07, + "output_cost_per_token": 1.5e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/microsoft/phi-4": { + "max_tokens": 16384, + "max_input_tokens": 16384, + "max_output_tokens": 16384, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 1.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/NousResearch/Hermes-3-Llama-3.1-405B": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 7e-07, + "output_cost_per_token": 8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/zai-org/GLM-4.5-Air": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 1.1e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 7.5e-08, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/openai/gpt-oss-120b": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 4.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/codegemma-7b-it": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 7e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct-Turbo": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/mistralai/Mistral-Nemo-Instruct-2407": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 4e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/openbmb/MiniCPM-Llama3-V-2_5": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3.4e-07, + "output_cost_per_token": 3.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/bigcode/starcoder2-15b-instruct-v0.1": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "max_tokens": 1048576, + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Llama-4-Scout-17B-16E-Instruct": { + "max_tokens": 327680, + "max_input_tokens": 327680, + "max_output_tokens": 327680, + "input_cost_per_token": 8e-08, + "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/gemini-2.0-flash-001": { + "max_tokens": 1000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Gryphe/MythoMax-L2-13b-turbo": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 1.3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/google/gemma-1.1-7b-it": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 7e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 2.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 3e-08, + "output_cost_per_token": 5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-32B": { + "max_tokens": 40960, + "max_input_tokens": 40960, + "max_output_tokens": 40960, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-235B-A22B-Thinking-2507": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, "deepinfra/meta-llama/Llama-2-70b-chat-hf": { "max_tokens": 4096, "max_input_tokens": 4096, "max_output_tokens": 4096, - "input_cost_per_token": 7e-07, - "output_cost_per_token": 9e-07, + "input_cost_per_token": 6.4e-07, + "output_cost_per_token": 8e-07, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, - "deepinfra/cognitivecomputations/dolphin-2.6-mixtral-8x7b": { - "max_tokens": 8191, - "max_input_tokens": 32768, - "max_output_tokens": 8191, - "input_cost_per_token": 2.7e-07, - "output_cost_per_token": 2.7e-07, - "litellm_provider": "deepinfra", - "mode": "chat", - "supports_tool_choice": true - }, - "deepinfra/codellama/CodeLlama-34b-Instruct-hf": { + "deepinfra/nvidia/Nemotron-4-340B-Instruct": { "max_tokens": 4096, "max_input_tokens": 4096, "max_output_tokens": 4096, - "input_cost_per_token": 6e-07, - "output_cost_per_token": 6e-07, + "input_cost_per_token": 4.2e-06, + "output_cost_per_token": 4.2e-06, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, - "deepinfra/deepinfra/mixtral": { - "max_tokens": 4096, - "max_input_tokens": 32000, - "max_output_tokens": 4096, - "input_cost_per_token": 2.7e-07, - "output_cost_per_token": 2.7e-07, - "litellm_provider": "deepinfra", - "mode": "completion" - }, - "deepinfra/Phind/Phind-CodeLlama-34B-v2": { - "max_tokens": 4096, - "max_input_tokens": 16384, - "max_output_tokens": 4096, - "input_cost_per_token": 6e-07, - "output_cost_per_token": 6e-07, + "deepinfra/deepseek-ai/DeepSeek-R1-0528": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.15e-06, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, - "deepinfra/mistralai/Mixtral-8x7B-Instruct-v0.1": { - "max_tokens": 8191, + "deepinfra/deepseek-ai/DeepSeek-R1-Turbo": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/NovaSky-AI/Sky-T1-32B-Preview": { + "max_tokens": 32768, "max_input_tokens": 32768, - "max_output_tokens": 8191, - "input_cost_per_token": 2.7e-07, - "output_cost_per_token": 2.7e-07, + "max_output_tokens": 32768, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 1.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen3-235B-A22B-Instruct-2507": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 6e-07, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, + "deepinfra/mistralai/Mistral-Small-3.1-24B-Instruct-2503": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 5e-08, + "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/mistralai/Mistral-7B-Instruct-v0.1": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen2-72B-Instruct": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Llama-4-Maverick-17B-128E-Instruct-Turbo": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Sao10K/L3-8B-Lunaris-v1": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3e-08, + "output_cost_per_token": 6e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, "deepinfra/deepinfra/airoboros-70b": { "max_tokens": 4096, "max_input_tokens": 4096, @@ -15119,112 +16135,72 @@ "mode": "chat", "supports_tool_choice": true }, - "deepinfra/01-ai/Yi-34B-Chat": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 6e-07, - "output_cost_per_token": 6e-07, - "litellm_provider": "deepinfra", - "mode": "chat", - "supports_tool_choice": true - }, - "deepinfra/01-ai/Yi-6B-200K": { - "max_tokens": 4096, - "max_input_tokens": 200000, - "max_output_tokens": 4096, - "input_cost_per_token": 1.3e-07, - "output_cost_per_token": 1.3e-07, - "litellm_provider": "deepinfra", - "mode": "completion" - }, - "deepinfra/jondurbin/airoboros-l2-70b-gpt4-1.4.1": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 7e-07, - "output_cost_per_token": 9e-07, - "litellm_provider": "deepinfra", - "mode": "chat", - "supports_tool_choice": true - }, - "deepinfra/meta-llama/Llama-2-13b-chat-hf": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 2.2e-07, - "output_cost_per_token": 2.2e-07, - "litellm_provider": "deepinfra", - "mode": "chat", - "supports_tool_choice": true - }, - "deepinfra/amazon/MistralLite": { - "max_tokens": 8191, - "max_input_tokens": 32768, - "max_output_tokens": 8191, - "input_cost_per_token": 2e-07, - "output_cost_per_token": 2e-07, - "litellm_provider": "deepinfra", - "mode": "chat", - "supports_tool_choice": true - }, - "deepinfra/meta-llama/Llama-2-7b-chat-hf": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 1.3e-07, - "output_cost_per_token": 1.3e-07, + "deepinfra/google/gemma-3-4b-it": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 4e-08, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3-8B-Instruct": { - "max_tokens": 8191, - "max_input_tokens": 8191, - "max_output_tokens": 4096, - "input_cost_per_token": 8e-08, - "output_cost_per_token": 8e-08, + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3e-08, + "output_cost_per_token": 6e-08, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, - "deepinfra/meta-llama/Meta-Llama-3-70B-Instruct": { - "max_tokens": 8191, - "max_input_tokens": 8191, - "max_output_tokens": 4096, - "input_cost_per_token": 5.9e-07, - "output_cost_per_token": 7.9e-07, - "litellm_provider": "deepinfra", - "mode": "chat", - "supports_tool_choice": true - }, - "deepinfra/meta-llama/Meta-Llama-3.1-405B-Instruct": { + "deepinfra/mistralai/Mistral-7B-Instruct-v0.2": { "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "input_cost_per_token": 9e-07, - "output_cost_per_token": 9e-07, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-3.3-70B-Instruct-Turbo": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 3.8e-08, + "output_cost_per_token": 1.2e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, "supports_tool_choice": true }, - "deepinfra/01-ai/Yi-34B-200K": { - "max_tokens": 4096, - "max_input_tokens": 200000, - "max_output_tokens": 4096, - "input_cost_per_token": 6e-07, - "output_cost_per_token": 6e-07, + "deepinfra/mistralai/Devstral-Small-2505": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 6e-08, + "output_cost_per_token": 1.2e-07, "litellm_provider": "deepinfra", - "mode": "completion" + "mode": "chat", + "supports_tool_choice": true }, - "deepinfra/openchat/openchat_3.5": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 1.3e-07, - "output_cost_per_token": 1.3e-07, + "deepinfra/meta-llama/Llama-3.2-11B-Vision-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 4.9e-08, + "output_cost_per_token": 4.9e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/openai/gpt-oss-20b": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 1.6e-07, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true @@ -16650,6 +17626,14 @@ "litellm_provider": "voyage", "mode": "embedding" }, + "voyage/voyage-context-3": { + "max_tokens": 120000, + "max_input_tokens": 120000, + "input_cost_per_token": 1.8e-07, + "output_cost_per_token": 0.0, + "litellm_provider": "voyage", + "mode": "embedding" + }, "voyage/rerank-2": { "max_tokens": 16000, "max_input_tokens": 16000, @@ -16707,22 +17691,6 @@ }, "supports_tool_choice": true }, - "databricks/databricks-meta-llama-3-1-70b-instruct": { - "max_tokens": 128000, - "max_input_tokens": 128000, - "max_output_tokens": 128000, - "input_cost_per_token": 1.00002e-06, - "input_dbu_cost_per_token": 1.4286e-05, - "output_cost_per_token": 2.99999e-06, - "output_dbu_cost_per_token": 4.2857e-05, - "litellm_provider": "databricks", - "mode": "chat", - "source": "https://www.databricks.com/product/pricing/foundation-model-serving", - "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." - }, - "supports_tool_choice": true - }, "databricks/databricks-meta-llama-3-3-70b-instruct": { "max_tokens": 128000, "max_input_tokens": 128000, @@ -16755,22 +17723,6 @@ }, "supports_tool_choice": true }, - "databricks/databricks-dbrx-instruct": { - "max_tokens": 32768, - "max_input_tokens": 32768, - "max_output_tokens": 32768, - "input_cost_per_token": 7.4998e-07, - "input_dbu_cost_per_token": 1.0714e-05, - "output_cost_per_token": 2.24901e-06, - "output_dbu_cost_per_token": 3.2143e-05, - "litellm_provider": "databricks", - "mode": "chat", - "source": "https://www.databricks.com/product/pricing/foundation-model-serving", - "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." - }, - "supports_tool_choice": true - }, "databricks/databricks-meta-llama-3-70b-instruct": { "max_tokens": 128000, "max_input_tokens": 128000, diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 59fab1b3369..f4dc1ef6c84 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -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) diff --git a/litellm/passthrough/utils.py b/litellm/passthrough/utils.py index c52d0e3688d..4bf66d49881 100644 --- a/litellm/passthrough/utils.py +++ b/litellm/passthrough/utils.py @@ -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//invoke + 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 \ No newline at end of file diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index aa2a6ee1a75..38619112ccc 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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.""" diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 2675c2f687d..b3653c31435 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -15,4 +15,16 @@ model_list: mode: chat router_settings: - model_group_alias: {"my-fake-gpt-4": "fake-openai-endpoint"} \ No newline at end of file + 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 \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3e155135d79..2239526a316 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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): diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index d02671d69d8..6222233d502 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -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", diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index b73765781ec..b04d14bcc8b 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -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 ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a1650b4968b..ddce7481ce7 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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 ( diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 451f110a894..98165a5a0ce 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -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 diff --git a/litellm/proxy/pass_through_endpoints/common_utils.py b/litellm/proxy/pass_through_endpoints/common_utils.py index 3a3783dd57c..804960cdee9 100644 --- a/litellm/proxy/pass_through_endpoints/common_utils.py +++ b/litellm/proxy/pass_through_endpoints/common_utils.py @@ -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", "") + diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index d743d41ac61..82c5b3e343d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -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//. 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, diff --git a/litellm/proxy/spend_tracking/cold_storage_handler.py b/litellm/proxy/spend_tracking/cold_storage_handler.py index b3ab7b6224c..133403deae2 100644 --- a/litellm/proxy/spend_tracking/cold_storage_handler.py +++ b/litellm/proxy/spend_tracking/cold_storage_handler.py @@ -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 - ) \ No newline at end of file diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 678d4963ac8..5b9337e852b 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -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, diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 886c46e7373..bfb996a2385 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -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: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 2a26c54a56e..6757e1ad31f 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -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, ) ) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index b50c84ad463..ac59d28a50d 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -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]], diff --git a/litellm/types/llms/mistral.py b/litellm/types/llms/mistral.py index e9563a9ae3f..375120e5376 100644 --- a/litellm/types/llms/mistral.py +++ b/litellm/types/llms/mistral.py @@ -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] diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 30bf7f7c956..74ff831503e 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -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): diff --git a/litellm/utils.py b/litellm/utils.py index 8d7360b4fb2..bf8aefb95d7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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 diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 55c42d928fa..202e7f0759f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -9723,6 +9723,30 @@ "supports_tool_choice": true, "supports_prompt_caching": true }, + "vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas": { + "max_tokens": 32768, + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 4e-06, + "litellm_provider": "vertex_ai-qwen_models", + "mode": "chat", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { + "max_tokens": 16384, + "max_input_tokens": 262144, + "max_output_tokens": 16384, + "input_cost_per_token": 0.25e-06, + "output_cost_per_token": 1e-06, + "litellm_provider": "vertex_ai-qwen_models", + "mode": "chat", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supports_function_calling": true, + "supports_tool_choice": true + }, "vertex_ai/meta/llama3-405b-instruct-maas": { "max_tokens": 32000, "max_input_tokens": 32000, @@ -10072,6 +10096,12 @@ "mode": "image_generation", "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "vertex_ai/imagen-4.0-generate-001": { + "output_cost_per_image": 0.04, + "litellm_provider": "vertex_ai-image-models", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, "vertex_ai/imagen-4.0-generate-preview-06-06": { "output_cost_per_image": 0.04, "litellm_provider": "vertex_ai-image-models", @@ -10084,6 +10114,18 @@ "mode": "image_generation", "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "vertex_ai/imagen-4.0-ultra-generate-001": { + "output_cost_per_image": 0.06, + "litellm_provider": "vertex_ai-image-models", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, + "vertex_ai/imagen-4.0-fast-generate-001": { + "output_cost_per_image": 0.02, + "litellm_provider": "vertex_ai-image-models", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, "vertex_ai/imagen-4.0-fast-generate-preview-06-06": { "output_cost_per_image": 0.02, "litellm_provider": "vertex_ai-image-models", @@ -10790,18 +10832,36 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", "supports_tool_choice": true }, + "gemini/imagen-4.0-generate-001": { + "output_cost_per_image": 0.04, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, "gemini/imagen-4.0-generate-preview-06-06": { "output_cost_per_image": 0.04, "litellm_provider": "gemini", "mode": "image_generation", "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "gemini/imagen-4.0-ultra-generate-001": { + "output_cost_per_image": 0.06, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, "gemini/imagen-4.0-ultra-generate-preview-06-06": { "output_cost_per_image": 0.06, "litellm_provider": "gemini", "mode": "image_generation", "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "gemini/imagen-4.0-fast-generate-001": { + "output_cost_per_image": 0.02, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, "gemini/imagen-4.0-fast-generate-preview-06-06": { "output_cost_per_image": 0.02, "litellm_provider": "gemini", @@ -11176,6 +11236,21 @@ "supports_tool_choice": true, "supports_prompt_caching": true }, + "openrouter/deepseek/deepseek-chat-v3.1": { + "max_tokens": 8192, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 2e-07, + "input_cost_per_token_cache_hit": 2e-08, + "output_cost_per_token": 8e-07, + "litellm_provider": "openrouter", + "mode": "chat", + "supports_function_calling": true, + "supports_assistant_prefill": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_prompt_caching": true + }, "openrouter/x-ai/grok-4": { "max_tokens": 256000, "max_input_tokens": 256000, @@ -15020,95 +15095,1036 @@ "litellm_provider": "ollama", "mode": "completion" }, - "deepinfra/lizpreciatior/lzlv_70b_fp16_hf": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 7e-07, - "output_cost_per_token": 9e-07, + "deepinfra/deepseek-ai/DeepSeek-V3": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 3.8e-07, + "output_cost_per_token": 8.9e-07, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, + "deepinfra/Phind/Phind-CodeLlama-34B-v2": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1.5e-08, + "output_cost_per_token": 2e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/gemma-2-9b-it": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3e-08, + "output_cost_per_token": 6e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/deepseek-ai/DeepSeek-R1-0528-Turbo": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen2-7B-Instruct": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/QVQ-72B-Preview": { + "max_tokens": 32000, + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-3.3-70B-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 2.3e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/microsoft/Phi-4-multimodal-instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 5e-08, + "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/mistralai/Devstral-Small-2507": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/microsoft/WizardLM-2-7B": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-3.2-90B-Vision-Instruct": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/mistralai/Mistral-Small-3.2-24B-Instruct-2506": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 5e-08, + "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-V3-0324": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 2.8e-07, + "output_cost_per_token": 8.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/mistralai/Mixtral-8x7B-Instruct-v0.1": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 8e-08, + "output_cost_per_token": 2.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/anthropic/claude-3-7-sonnet-latest": { + "max_tokens": 200000, + "max_input_tokens": 200000, + "max_output_tokens": 200000, + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/cognitivecomputations/dolphin-2.9.1-llama-3-70b": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen2.5-Coder-32B-Instruct": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 6e-08, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen3-235B-A22B": { + "max_tokens": 40960, + "max_input_tokens": 40960, + "max_output_tokens": 40960, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-V3-0324-Turbo": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/microsoft/WizardLM-2-8x22B": { + "max_tokens": 65536, + "max_input_tokens": 65536, + "max_output_tokens": 65536, + "input_cost_per_token": 4.8e-07, + "output_cost_per_token": 4.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Sao10K/L3-8B-Lunaris-v1-Turbo": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-Guard-4-12B": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 1.8e-07, + "output_cost_per_token": 1.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, "deepinfra/Gryphe/MythoMax-L2-13b": { "max_tokens": 4096, "max_input_tokens": 4096, "max_output_tokens": 4096, - "input_cost_per_token": 2.2e-07, - "output_cost_per_token": 2.2e-07, + "input_cost_per_token": 7.2e-08, + "output_cost_per_token": 7.2e-08, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, - "deepinfra/mistralai/Mistral-7B-Instruct-v0.1": { - "max_tokens": 8191, + "deepinfra/meta-llama/Llama-3.2-1B-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 5e-09, + "output_cost_per_token": 1e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/gemma-2-27b-it": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 2.7e-07, + "output_cost_per_token": 2.7e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen2.5-VL-32B-Instruct": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/nvidia/Llama-3.1-Nemotron-70B-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/mistralai/Mixtral-8x22B-Instruct-v0.1": { + "max_tokens": 65536, + "max_input_tokens": 65536, + "max_output_tokens": 65536, + "input_cost_per_token": 6.5e-07, + "output_cost_per_token": 6.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen2.5-7B-Instruct": { + "max_tokens": 32768, "max_input_tokens": 32768, - "max_output_tokens": 8191, + "max_output_tokens": 32768, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/google/gemini-1.5-flash-8b": { + "max_tokens": 1000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "input_cost_per_token": 3.75e-08, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/NousResearch/Hermes-3-Llama-3.1-70B": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 2.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-Guard-3-8B": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/mistralai/Mistral-Small-24B-Instruct-2501": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 5e-08, + "output_cost_per_token": 8e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/lizpreciatior/lzlv_70b_fp16_hf": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-2-13b-chat-hf": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, "input_cost_per_token": 1.3e-07, "output_cost_per_token": 1.3e-07, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, + "deepinfra/anthropic/claude-4-opus": { + "max_tokens": 200000, + "max_input_tokens": 200000, + "max_output_tokens": 200000, + "input_cost_per_token": 1.65e-05, + "output_cost_per_token": 8.25e-05, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/openchat/openchat-3.6-8b": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/google/gemma-3-27b-it": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 1.7e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Austism/chronos-hermes-13b-v2": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 1.3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Sao10K/L3.1-70B-Euryale-v2.2": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 6.5e-07, + "output_cost_per_token": 7.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/QwQ-32B-Preview": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 1.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/anthropic/claude-4-sonnet": { + "max_tokens": 200000, + "max_input_tokens": 200000, + "max_output_tokens": 200000, + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/microsoft/Phi-3-medium-4k-instruct": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 1.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/mattshumer/Reflection-Llama-3.1-70B": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/openchat/openchat_3.5": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Sao10K/L3.3-70B-Euryale-v2.3": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 6.5e-07, + "output_cost_per_token": 7.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 2.3e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-V3.1": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 3.2e-07, + "output_cost_per_token": 1.15e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen2.5-Coder-7B": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 2.5e-08, + "output_cost_per_token": 5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/cognitivecomputations/dolphin-2.6-mixtral-8x7b": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 2.4e-07, + "output_cost_per_token": 2.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Meta-Llama-3.1-405B-Instruct": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 8e-07, + "output_cost_per_token": 8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-Prover-V2-671B": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.18e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/zai-org/GLM-4.5": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 5.5e-07, + "output_cost_per_token": 2e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Llama-3.2-3B-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1.2e-08, + "output_cost_per_token": 2.4e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Meta-Llama-3-70B-Instruct": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/gemini-1.5-flash": { + "max_tokens": 1000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "input_cost_per_token": 7.5e-08, + "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/KoboldAI/LLaMA2-13B-Tiefighter": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/gemini-2.5-pro": { + "max_tokens": 1000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "input_cost_per_token": 8.75e-07, + "output_cost_per_token": 7e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-30B-A3B": { + "max_tokens": 40960, + "max_input_tokens": 40960, + "max_output_tokens": 40960, + "input_cost_per_token": 8e-08, + "output_cost_per_token": 2.9e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/QwQ-32B": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 7.5e-08, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/moonshotai/Kimi-K2-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Sao10K/L3-70B-Euryale-v2.1": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 7e-07, + "output_cost_per_token": 8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/microsoft/phi-4-reasoning-plus": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 3.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/google/gemma-3-12b-it": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 5e-08, + "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/gemini-2.5-flash": { + "max_tokens": 1000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "input_cost_per_token": 2.1e-07, + "output_cost_per_token": 1.75e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-R1": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 4.5e-07, + "output_cost_per_token": 2.15e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/mistralai/Mistral-7B-Instruct-v0.3": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 2.8e-08, + "output_cost_per_token": 5.4e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen2.5-72B-Instruct": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 3.9e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-14B": { + "max_tokens": 40960, + "max_input_tokens": 40960, + "max_output_tokens": 40960, + "input_cost_per_token": 6e-08, + "output_cost_per_token": 2.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/allenai/olmOCR-7B-0725-FP8": { + "max_tokens": 16384, + "max_input_tokens": 16384, + "max_output_tokens": 16384, + "input_cost_per_token": 2.7e-07, + "output_cost_per_token": 1.5e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/microsoft/phi-4": { + "max_tokens": 16384, + "max_input_tokens": 16384, + "max_output_tokens": 16384, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 1.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/NousResearch/Hermes-3-Llama-3.1-405B": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 7e-07, + "output_cost_per_token": 8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/zai-org/GLM-4.5-Air": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 1.1e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 7.5e-08, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/openai/gpt-oss-120b": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 4.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/codegemma-7b-it": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 7e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct-Turbo": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/mistralai/Mistral-Nemo-Instruct-2407": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 4e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/openbmb/MiniCPM-Llama3-V-2_5": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3.4e-07, + "output_cost_per_token": 3.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/bigcode/starcoder2-15b-instruct-v0.1": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "max_tokens": 1048576, + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Llama-4-Scout-17B-16E-Instruct": { + "max_tokens": 327680, + "max_input_tokens": 327680, + "max_output_tokens": 327680, + "input_cost_per_token": 8e-08, + "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/google/gemini-2.0-flash-001": { + "max_tokens": 1000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Gryphe/MythoMax-L2-13b-turbo": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 1.3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/google/gemma-1.1-7b-it": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 7e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 2.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 3e-08, + "output_cost_per_token": 5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-32B": { + "max_tokens": 40960, + "max_input_tokens": 40960, + "max_output_tokens": 40960, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-235B-A22B-Thinking-2507": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, "deepinfra/meta-llama/Llama-2-70b-chat-hf": { "max_tokens": 4096, "max_input_tokens": 4096, "max_output_tokens": 4096, - "input_cost_per_token": 7e-07, - "output_cost_per_token": 9e-07, + "input_cost_per_token": 6.4e-07, + "output_cost_per_token": 8e-07, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, - "deepinfra/cognitivecomputations/dolphin-2.6-mixtral-8x7b": { - "max_tokens": 8191, - "max_input_tokens": 32768, - "max_output_tokens": 8191, - "input_cost_per_token": 2.7e-07, - "output_cost_per_token": 2.7e-07, - "litellm_provider": "deepinfra", - "mode": "chat", - "supports_tool_choice": true - }, - "deepinfra/codellama/CodeLlama-34b-Instruct-hf": { + "deepinfra/nvidia/Nemotron-4-340B-Instruct": { "max_tokens": 4096, "max_input_tokens": 4096, "max_output_tokens": 4096, - "input_cost_per_token": 6e-07, - "output_cost_per_token": 6e-07, + "input_cost_per_token": 4.2e-06, + "output_cost_per_token": 4.2e-06, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, - "deepinfra/deepinfra/mixtral": { - "max_tokens": 4096, - "max_input_tokens": 32000, - "max_output_tokens": 4096, - "input_cost_per_token": 2.7e-07, - "output_cost_per_token": 2.7e-07, - "litellm_provider": "deepinfra", - "mode": "completion" - }, - "deepinfra/Phind/Phind-CodeLlama-34B-v2": { - "max_tokens": 4096, - "max_input_tokens": 16384, - "max_output_tokens": 4096, - "input_cost_per_token": 6e-07, - "output_cost_per_token": 6e-07, + "deepinfra/deepseek-ai/DeepSeek-R1-0528": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.15e-06, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, - "deepinfra/mistralai/Mixtral-8x7B-Instruct-v0.1": { - "max_tokens": 8191, + "deepinfra/deepseek-ai/DeepSeek-R1-Turbo": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/NovaSky-AI/Sky-T1-32B-Preview": { + "max_tokens": 32768, "max_input_tokens": 32768, - "max_output_tokens": 8191, - "input_cost_per_token": 2.7e-07, - "output_cost_per_token": 2.7e-07, + "max_output_tokens": 32768, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 1.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Qwen/Qwen3-235B-A22B-Instruct-2507": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 6e-07, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, + "deepinfra/mistralai/Mistral-Small-3.1-24B-Instruct-2503": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 5e-08, + "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/mistralai/Mistral-7B-Instruct-v0.1": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen2-72B-Instruct": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/meta-llama/Llama-4-Maverick-17B-128E-Instruct-Turbo": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/Sao10K/L3-8B-Lunaris-v1": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3e-08, + "output_cost_per_token": 6e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, "deepinfra/deepinfra/airoboros-70b": { "max_tokens": 4096, "max_input_tokens": 4096, @@ -15119,112 +16135,72 @@ "mode": "chat", "supports_tool_choice": true }, - "deepinfra/01-ai/Yi-34B-Chat": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 6e-07, - "output_cost_per_token": 6e-07, - "litellm_provider": "deepinfra", - "mode": "chat", - "supports_tool_choice": true - }, - "deepinfra/01-ai/Yi-6B-200K": { - "max_tokens": 4096, - "max_input_tokens": 200000, - "max_output_tokens": 4096, - "input_cost_per_token": 1.3e-07, - "output_cost_per_token": 1.3e-07, - "litellm_provider": "deepinfra", - "mode": "completion" - }, - "deepinfra/jondurbin/airoboros-l2-70b-gpt4-1.4.1": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 7e-07, - "output_cost_per_token": 9e-07, - "litellm_provider": "deepinfra", - "mode": "chat", - "supports_tool_choice": true - }, - "deepinfra/meta-llama/Llama-2-13b-chat-hf": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 2.2e-07, - "output_cost_per_token": 2.2e-07, - "litellm_provider": "deepinfra", - "mode": "chat", - "supports_tool_choice": true - }, - "deepinfra/amazon/MistralLite": { - "max_tokens": 8191, - "max_input_tokens": 32768, - "max_output_tokens": 8191, - "input_cost_per_token": 2e-07, - "output_cost_per_token": 2e-07, - "litellm_provider": "deepinfra", - "mode": "chat", - "supports_tool_choice": true - }, - "deepinfra/meta-llama/Llama-2-7b-chat-hf": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 1.3e-07, - "output_cost_per_token": 1.3e-07, + "deepinfra/google/gemma-3-4b-it": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 4e-08, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3-8B-Instruct": { - "max_tokens": 8191, - "max_input_tokens": 8191, - "max_output_tokens": 4096, - "input_cost_per_token": 8e-08, - "output_cost_per_token": 8e-08, + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 3e-08, + "output_cost_per_token": 6e-08, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true }, - "deepinfra/meta-llama/Meta-Llama-3-70B-Instruct": { - "max_tokens": 8191, - "max_input_tokens": 8191, - "max_output_tokens": 4096, - "input_cost_per_token": 5.9e-07, - "output_cost_per_token": 7.9e-07, - "litellm_provider": "deepinfra", - "mode": "chat", - "supports_tool_choice": true - }, - "deepinfra/meta-llama/Meta-Llama-3.1-405B-Instruct": { + "deepinfra/mistralai/Mistral-7B-Instruct-v0.2": { "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "input_cost_per_token": 9e-07, - "output_cost_per_token": 9e-07, + "input_cost_per_token": 5.5e-08, + "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/meta-llama/Llama-3.3-70B-Instruct-Turbo": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 3.8e-08, + "output_cost_per_token": 1.2e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, "supports_tool_choice": true }, - "deepinfra/01-ai/Yi-34B-200K": { - "max_tokens": 4096, - "max_input_tokens": 200000, - "max_output_tokens": 4096, - "input_cost_per_token": 6e-07, - "output_cost_per_token": 6e-07, + "deepinfra/mistralai/Devstral-Small-2505": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 6e-08, + "output_cost_per_token": 1.2e-07, "litellm_provider": "deepinfra", - "mode": "completion" + "mode": "chat", + "supports_tool_choice": true }, - "deepinfra/openchat/openchat_3.5": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 1.3e-07, - "output_cost_per_token": 1.3e-07, + "deepinfra/meta-llama/Llama-3.2-11B-Vision-Instruct": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 4.9e-08, + "output_cost_per_token": 4.9e-08, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": false + }, + "deepinfra/openai/gpt-oss-20b": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 1.6e-07, "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true @@ -16650,6 +17626,14 @@ "litellm_provider": "voyage", "mode": "embedding" }, + "voyage/voyage-context-3": { + "max_tokens": 120000, + "max_input_tokens": 120000, + "input_cost_per_token": 1.8e-07, + "output_cost_per_token": 0.0, + "litellm_provider": "voyage", + "mode": "embedding" + }, "voyage/rerank-2": { "max_tokens": 16000, "max_input_tokens": 16000, @@ -16707,22 +17691,6 @@ }, "supports_tool_choice": true }, - "databricks/databricks-meta-llama-3-1-70b-instruct": { - "max_tokens": 128000, - "max_input_tokens": 128000, - "max_output_tokens": 128000, - "input_cost_per_token": 1.00002e-06, - "input_dbu_cost_per_token": 1.4286e-05, - "output_cost_per_token": 2.99999e-06, - "output_dbu_cost_per_token": 4.2857e-05, - "litellm_provider": "databricks", - "mode": "chat", - "source": "https://www.databricks.com/product/pricing/foundation-model-serving", - "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." - }, - "supports_tool_choice": true - }, "databricks/databricks-meta-llama-3-3-70b-instruct": { "max_tokens": 128000, "max_input_tokens": 128000, @@ -16755,22 +17723,6 @@ }, "supports_tool_choice": true }, - "databricks/databricks-dbrx-instruct": { - "max_tokens": 32768, - "max_input_tokens": 32768, - "max_output_tokens": 32768, - "input_cost_per_token": 7.4998e-07, - "input_dbu_cost_per_token": 1.0714e-05, - "output_cost_per_token": 2.24901e-06, - "output_dbu_cost_per_token": 3.2143e-05, - "litellm_provider": "databricks", - "mode": "chat", - "source": "https://www.databricks.com/product/pricing/foundation-model-serving", - "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." - }, - "supports_tool_choice": true - }, "databricks/databricks-meta-llama-3-70b-instruct": { "max_tokens": 128000, "max_input_tokens": 128000, diff --git a/pyproject.toml b/pyproject.toml index 574348979ba..15ba5af40b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" ] diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_internal_user_endpoints.py index d69c613a4f5..69c3b4cb59a 100644 --- a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_internal_user_endpoints.py @@ -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"] - ) diff --git a/tests/image_gen_tests/base_image_generation_test.py b/tests/image_gen_tests/base_image_generation_test.py index c3a5cfb2251..6e8470525d6 100644 --- a/tests/image_gen_tests/base_image_generation_test.py +++ b/tests/image_gen_tests/base_image_generation_test.py @@ -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)}") \ No newline at end of file + 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 diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index e9bfa643f3b..3228bd92189 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -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") diff --git a/tests/llm_responses_api_testing/base_responses_api.py b/tests/llm_responses_api_testing/base_responses_api.py index 72f6fb8c9cf..5cb8295b1af 100644 --- a/tests/llm_responses_api_testing/base_responses_api.py +++ b/tests/llm_responses_api_testing/base_responses_api.py @@ -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 diff --git a/tests/llm_translation/test_lambda_ai.py b/tests/llm_translation/test_lambda_ai.py index 3b73629b4a7..e50c6b09ef9 100644 --- a/tests/llm_translation/test_lambda_ai.py +++ b/tests/llm_translation/test_lambda_ai.py @@ -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 diff --git a/tests/llm_translation/test_optional_params.py b/tests/llm_translation/test_optional_params.py index 52f6d99fdf5..0e1e34362b0 100644 --- a/tests/llm_translation/test_optional_params.py +++ b/tests/llm_translation/test_optional_params.py @@ -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"] diff --git a/tests/llm_translation/test_voyage_ai.py b/tests/llm_translation/test_voyage_ai.py index d1c01024a05..a0b9ee0a44b 100644 --- a/tests/llm_translation/test_voyage_ai.py +++ b/tests/llm_translation/test_voyage_ai.py @@ -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}") \ No newline at end of file + 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}") diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 639d49d02a4..632e14a76eb 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -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) diff --git a/tests/local_testing/test_braintrust.py b/tests/local_testing/test_braintrust.py index c898594683e..13b23a97585 100644 --- a/tests/local_testing/test_braintrust.py +++ b/tests/local_testing/test_braintrust.py @@ -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 diff --git a/tests/test_litellm/caching/test_s3_cache.py b/tests/test_litellm/caching/test_s3_cache.py index a37f7f57064..dce3f7d585d 100644 --- a/tests/test_litellm/caching/test_s3_cache.py +++ b/tests/test_litellm/caching/test_s3_cache.py @@ -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/" \ No newline at end of file + 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() diff --git a/tests/test_litellm/integrations/test_mlflow.py b/tests/test_litellm/integrations/test_mlflow.py index f2ca8d992b9..3b75268d556 100644 --- a/tests/test_litellm/integrations/test_mlflow.py +++ b/tests/test_litellm/integrations/test_mlflow.py @@ -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(): diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 84edbb3ac5c..8fd73dbb79f 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -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( diff --git a/tests/test_litellm/llms/baseten/chat/test_baseten_completions.py b/tests/test_litellm/llms/baseten/chat/test_baseten_completions.py new file mode 100644 index 00000000000..9420149a8e4 --- /dev/null +++ b/tests/test_litellm/llms/baseten/chat/test_baseten_completions.py @@ -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__]) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index b60a30bee80..1c91cc0fe8b 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -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"] diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index a2bbaa620b2..01903737580 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -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' diff --git a/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py b/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py index 4999b078caa..1230a1fd2aa 100644 --- a/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py +++ b/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py @@ -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/"), diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank.py new file mode 100644 index 00000000000..0dda7d08da4 --- /dev/null +++ b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank.py @@ -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 diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py new file mode 100644 index 00000000000..3655f5c643b --- /dev/null +++ b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py @@ -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 diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py new file mode 100644 index 00000000000..252eb40532c --- /dev/null +++ b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py @@ -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) diff --git a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py index 3e4d3ddb304..e6d7ed78d6e 100644 --- a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py @@ -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 "" 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 "" 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 "" 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 "" 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 "" 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 "" 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 \ No newline at end of file + 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 diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index 9b8e56ab499..6d46a40f6c6 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -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 diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index afbaba60cf3..9b7e20159e4 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -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 \ No newline at end of file diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index d67e22b6d3b..42c64c15814 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -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 diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 8567e84950f..da052e71b58 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -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" \ No newline at end of file diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 3f7fdc55f36..694a49159c0 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -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, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 54909999038..e0102f8cd7a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -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) + diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index dc1803d5569..1702a317a29 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -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" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py index bd8c5f5a99a..97ef05100de 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py @@ -10,6 +10,8 @@ import pytest from fastapi import Request, Response from fastapi.testclient import TestClient +from litellm.passthrough.utils import CommonUtils + sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path @@ -42,3 +44,55 @@ async def test_get_litellm_virtual_key(): } result = get_litellm_virtual_key(mock_request) assert result == "Bearer test-key-123" + +def test_encode_bedrock_runtime_modelid_arn(): + # Test application-inference-profile ARN + endpoint = "model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse" + expected = "model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile%2Fr742sbn2zckd/converse" + result = CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) + assert result == expected + + # Test inference-profile ARN + endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:inference-profile/test-profile/invoke" + expected = "model/arn:aws:bedrock:us-east-1:123456789012:inference-profile%2Ftest-profile/invoke" + result = CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) + assert result == expected + + # Test foundation-model ARN + endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:foundation-model/anthropic.claude-3/converse" + expected = "model/arn:aws:bedrock:us-east-1:123456789012:foundation-model%2Fanthropic.claude-3/converse" + result = CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) + assert result == expected + + # Test custom-model ARN (2 slashes) + endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:custom-model/my-model.fine-tuned/abc123/invoke" + expected = "model/arn:aws:bedrock:us-east-1:123456789012:custom-model%2Fmy-model.fine-tuned%2Fabc123/invoke" + result = CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) + assert result == expected + + # Test provisioned-model ARN + endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:provisioned-model/test-model/converse" + expected = "model/arn:aws:bedrock:us-east-1:123456789012:provisioned-model%2Ftest-model/converse" + result = CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) + assert result == expected + + +def test_encode_bedrock_runtime_modelid_arn_no_arn(): + # Test regular model ID (no ARN) + endpoint = "model/anthropic.claude-3-sonnet-20240229-v1:0/converse" + result = CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) + assert result == endpoint + + +def test_encode_bedrock_runtime_modelid_arn_edge_cases(): + # Test multiple ARN types (should only encode first match) + endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/test1/converse" + expected = "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile%2Ftest1/converse" + result = CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) + assert result == expected + + # Test ARN with special characters in resource ID + endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/test-profile.v1/invoke" + expected = "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile%2Ftest-profile.v1/invoke" + result = CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) + assert result == expected \ No newline at end of file diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py index 1cc159bf963..c683b571345 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py @@ -223,7 +223,7 @@ async def test_e2e_cold_storage_successful_retrieval(): new_callable=AsyncMock, ) as mock_get_spend_logs, \ patch.object(session_handler, "COLD_STORAGE_HANDLER") as mock_cold_storage, \ - patch("litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger", return_value="s3"): + patch("litellm.configured_cold_storage_logger", return_value="s3"): # Setup mocks mock_get_spend_logs.return_value = mock_spend_logs @@ -343,7 +343,7 @@ async def test_should_check_cold_storage_for_full_payload(): # Test case 4: None request (should return True) proxy_request_none = None - 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"): # Test case 1: Should return True for truncated content result1 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf) assert result1 == True, "Should return True for proxy request with truncated PDF content" @@ -361,6 +361,6 @@ async def test_should_check_cold_storage_for_full_payload(): assert result4 == True, "Should return True for None proxy request" # Test case 5: Should return False when cold storage is not configured - 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): result5 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf) assert result5 == False, "Should return False when cold storage is not configured, even with truncated content" diff --git a/ui/litellm-dashboard/out/assets/logos/vllm.png b/ui/litellm-dashboard/out/assets/logos/vllm.png new file mode 100644 index 00000000000..6026ec741dc Binary files /dev/null and b/ui/litellm-dashboard/out/assets/logos/vllm.png differ diff --git a/ui/litellm-dashboard/public/assets/logos/vllm.png b/ui/litellm-dashboard/public/assets/logos/vllm.png new file mode 100644 index 00000000000..6026ec741dc Binary files /dev/null and b/ui/litellm-dashboard/public/assets/logos/vllm.png differ diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx index d74283ec972..23b903808ba 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx @@ -400,6 +400,20 @@ const PROVIDER_CREDENTIAL_FIELDS: Record = placeholder: "http://localhost:8000/generate", required: false }], + [Providers.Hosted_Vllm]: [ + { + key: "api_base", + label: "API Base", + placeholder: "https://...", + required: true + }, + { + key: "api_key", + label: "OpenAI API Key", + type: "password", + required: true + } + ], [Providers.Voyage]: [{ key: "api_key", label: "API Key", diff --git a/ui/litellm-dashboard/src/components/all_keys_table.tsx b/ui/litellm-dashboard/src/components/all_keys_table.tsx index 34116d3f4b1..b65f8911a44 100644 --- a/ui/litellm-dashboard/src/components/all_keys_table.tsx +++ b/ui/litellm-dashboard/src/components/all_keys_table.tsx @@ -7,8 +7,8 @@ import { Button } from "@tremor/react" import KeyInfoView from "./templates/key_info_view" import { Tooltip } from "antd" import { Team, KeyResponse } from "./key_team_helpers/key_list" -import FilterComponent from "./common_components/filter" -import { FilterOption } from "./common_components/filter" +import FilterComponent from "./molecules/filter" +import { FilterOption } from "./molecules/filter" import { keyListCall, Organization, userListCall } from "./networking" import { createTeamSearchFunction } from "./key_team_helpers/team_search_fn" import { createOrgSearchFunction } from "./key_team_helpers/organization_search_fn" diff --git a/ui/litellm-dashboard/src/components/chat_ui/SessionManagement.tsx b/ui/litellm-dashboard/src/components/chat_ui/SessionManagement.tsx index 67b01018e93..1bf34a1758a 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/SessionManagement.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/SessionManagement.tsx @@ -2,6 +2,7 @@ import React from "react"; import { Switch, Tooltip, message } from "antd"; import { InfoCircleOutlined, CopyOutlined } from "@ant-design/icons"; import { EndpointType } from "./mode_endpoint_mapping"; +import NotificationsManager from "../molecules/notifications_manager"; interface SessionManagementProps { endpointType: string; diff --git a/ui/litellm-dashboard/src/components/common_components/filter.tsx b/ui/litellm-dashboard/src/components/molecules/filter.tsx similarity index 75% rename from ui/litellm-dashboard/src/components/common_components/filter.tsx rename to ui/litellm-dashboard/src/components/molecules/filter.tsx index 20b11447b21..e52cbb463b9 100644 --- a/ui/litellm-dashboard/src/components/common_components/filter.tsx +++ b/ui/litellm-dashboard/src/components/molecules/filter.tsx @@ -1,4 +1,4 @@ -import React, { useState, useCallback } from "react"; +import React, { useState, useCallback, useEffect } from "react"; import { Button, Input, Select } from "antd"; import { FilterIcon } from "@heroicons/react/outline"; import debounce from "lodash/debounce"; @@ -43,6 +43,9 @@ const FilterComponent: React.FC = ({ const [searchInputValueMap, setSearchInputValueMap] = useState<{ [key: string]: string; }>({}); + const [initialOptionsLoaded, setInitialOptionsLoaded] = useState<{ + [key: string]: boolean; + }>({}); const debouncedSearch = useCallback( debounce(async (value: string, option: FilterOption) => { @@ -62,6 +65,36 @@ const FilterComponent: React.FC = ({ [] ); + // Load initial options for searchable filters + const loadInitialOptions = useCallback(async (option: FilterOption) => { + if (!option.isSearchable || !option.searchFn || initialOptionsLoaded[option.name]) return; + + setSearchLoadingMap((prev) => ({ ...prev, [option.name]: true })); + setInitialOptionsLoaded((prev) => ({ ...prev, [option.name]: true })); + + try { + // Load initial options with empty search to get some default results + const results = await option.searchFn(""); + setSearchOptionsMap((prev) => ({ ...prev, [option.name]: results })); + } catch (error) { + console.error("Error loading initial options:", error); + setSearchOptionsMap((prev) => ({ ...prev, [option.name]: [] })); + } finally { + setSearchLoadingMap((prev) => ({ ...prev, [option.name]: false })); + } + }, [initialOptionsLoaded]); + + // Load initial options when filters are shown + useEffect(() => { + if (showFilters) { + options.forEach(option => { + if (option.isSearchable && !initialOptionsLoaded[option.name]) { + loadInitialOptions(option); + } + }); + } + }, [showFilters, options, loadInitialOptions, initialOptionsLoaded]); + const handleFilterChange = (name: string, value: string) => { const newValues = { ...tempValues, @@ -80,6 +113,13 @@ const FilterComponent: React.FC = ({ onResetFilters(); }; + // Handle dropdown open to load initial options + const handleDropdownVisibleChange = (open: boolean, option: FilterOption) => { + if (open && option.isSearchable && !initialOptionsLoaded[option.name]) { + loadInitialOptions(option); + } + }; + // Define the order of filters const orderedFilters = [ "Team ID", @@ -125,6 +165,7 @@ const FilterComponent: React.FC = ({ placeholder={`Search ${option.label || option.name}...`} value={tempValues[option.name] || undefined} onChange={(value) => handleFilterChange(option.name, value)} + onDropdownVisibleChange={(open) => handleDropdownVisibleChange(open, option)} onSearch={(value) => { setSearchInputValueMap((prev) => ({ ...prev, @@ -138,6 +179,7 @@ const FilterComponent: React.FC = ({ loading={searchLoadingMap[option.name]} options={searchOptionsMap[option.name] || []} allowClear + notFoundContent={searchLoadingMap[option.name] ? "Loading..." : "No results found"} /> ) : option.options ? (