Update Baseten LiteLLM integration

2025-08-19 12:21:05 -07:00 · 2025-08-19 12:21:05 -07:00 · 7c3d522435
commit 7c3d522435
parent 195ea6515e
10 changed files with 311 additions and 339 deletions
--- 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
-}
+}
--- 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/<model-name>` (e.g., `baseten/openai/gpt-oss-120b`)
+- **Best for**: Quick access to popular models
+
+### Dedicated Deployments
+- **URL**: `https://model-{id}.api.baseten.co/environments/production/sync/v1`
+- **Format**: `baseten/{8-digit-alphanumeric-code}` (e.g., `baseten/abcd1234`)
+- **Best for**: Custom models, latency SLAs
+
+:::tip
+**Automatic Routing**: LiteLLM detects the type based on model format:
+- 8-digit alphanumeric codes → Dedicated deployment
+- All other formats → Model API
+:::
+
+
+## Quick Start

-### API KEYS
 ```python
-import os 
-os.environ["BASETEN_API_KEY"] = ""
+import os
+from litellm import completion
+
+os.environ['BASETEN_API_KEY'] = "your-api-key"
+
+# Model API (default)
+response = completion(
+    model="baseten/openai/gpt-oss-120b",
+    messages=[{"role": "user", "content": "Hello!"}]
+)
+
+# Dedicated deployment (8-digit ID)
+response = completion(
+    model="baseten/abcd1234",
+    messages=[{"role": "user", "content": "Hello!"}]
+)
 ```

-### Baseten Models
-Baseten provides infrastructure to deploy and serve ML models https://www.baseten.co/. Use liteLLM to easily call models deployed on Baseten.
+## Examples

-Example Baseten Usage - Note: liteLLM supports all models deployed on Baseten
+### Basic Usage
+```python
+# Model API
+response = completion(
+    model="baseten/openai/gpt-oss-120b",
+    messages=[{"role": "user", "content": "Explain quantum computing"}],
+    max_tokens=500,
+    temperature=0.7
+)

-Usage: Pass `model=baseten/<Model ID>`
+# Dedicated deployment
+response = completion(
+    model="baseten/abcd1234",
+    messages=[{"role": "user", "content": "Explain quantum computing"}],
+    max_tokens=500,
+    temperature=0.7
+)
+```

-| Model Name       | Function Call                                  | Required OS Variables              |
-|------------------|--------------------------------------------|------------------------------------|
-| Falcon 7B        | `completion(model='baseten/qvv0xeq', messages=messages)`         | `os.environ['BASETEN_API_KEY']`     |
-| Wizard LM        | `completion(model='baseten/q841o8w', messages=messages)`         | `os.environ['BASETEN_API_KEY']`     |
-| MPT 7B Base      | `completion(model='baseten/31dxrj3', messages=messages)`         | `os.environ['BASETEN_API_KEY']`     |
+### Streaming (Model API only)
+```python
+response = completion(
+    model="baseten/openai/gpt-oss-120b",
+    messages=[{"role": "user", "content": "Write a poem"}],
+    stream=True,
+    stream_options={"include_usage": True}
+)
+
+for chunk in response:
+    if chunk.choices and chunk.choices[0].delta.content:
+        print(chunk.choices[0].delta.content, end="")
+```
+
+## Usage with LiteLLM Proxy
+
+1. **Config**:
+```yaml
+model_list:
+  - model_name: baseten-model
+    litellm_params:
+      model: baseten/openai/gpt-oss-120b
+      api_key: your-baseten-api-key
+```
+
+2. **Request**:
+```python
+import openai
+client = openai.OpenAI(
+    api_key="sk-1234",
+    base_url="http://0.0.0.0:4000"
+)
+
+response = client.chat.completions.create(
+    model="baseten-model",
+    messages=[{"role": "user", "content": "Hello!"}]
+)
+```
--- a/litellm/init.py
+++ b/litellm/init.py
@ -1193,6 +1193,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
--- 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",
--- 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
--- 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":
--- a/litellm/llms/baseten.py
+++ b/litellm/llms/baseten.py
@ -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
--- a/litellm/llms/baseten/chat.py
+++ 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[float] = None
+    tool_choice: Optional[str] = None
+    tools: Optional[list] = None
+    user: Optional[str] = None
+    presence_penalty: Optional[float] = None
+    frequency_penalty: Optional[float] = 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[float] = None,
+        tool_choice: Optional[str] = None,
+        tools: Optional[list] = None,
+        user: Optional[str] = None,
+        presence_penalty: Optional[float] = None,
+        frequency_penalty: Optional[float] = 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" 
--- a/litellm/main.py
+++ b/litellm/main.py
@ -1921,6 +1921,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"
--- a/litellm/utils.py
+++ b/litellm/utils.py
@ -3783,6 +3783,17 @@ def get_optional_params(  # noqa: PLR0915
                else False
            ),
        )
+    elif custom_llm_provider == "baseten":
+        optional_params = litellm.BasetenConfig().map_openai_params(
+            non_default_params=non_default_params,
+            optional_params=optional_params,
+            model=model,
+            drop_params=(
+                drop_params
+                if drop_params is not None and isinstance(drop_params, bool)
+                else False
+            ),
+        )
    elif custom_llm_provider == "xai":
        optional_params = litellm.XAIChatConfig().map_openai_params(
            model=model,
@ -5361,6 +5372,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
@ -6914,6 +6930,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: