Merge pull request #17193 from BerriAI/litellm_twelvelabs_int

Added support for twelvelabs pegasus

docs/my-website/docs/providers/bedrock.md

@@ -1683,6 +1683,131 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
 </TabItem>
 </Tabs>
 
+## TwelveLabs Pegasus - Video Understanding
+
+TwelveLabs Pegasus 1.2 is a video understanding model that can analyze and describe video content. LiteLLM supports this model through Bedrock's `/invoke` endpoint.
+
+| Property | Details |
+|----------|---------|
+| Provider Route | `bedrock/us.twelvelabs.pegasus-1-2-v1:0`, `bedrock/eu.twelvelabs.pegasus-1-2-v1:0` |
+| Provider Documentation | [TwelveLabs Pegasus Docs ↗](https://docs.twelvelabs.io/docs/models/pegasus) |
+| Supported Parameters | `max_tokens`, `temperature`, `response_format` |
+| Media Input | S3 URI or base64-encoded video |
+
+### Supported Features
+
+- **Video Analysis**: Analyze video content from S3 or base64 input
+- **Structured Output**: Support for JSON schema response format
+- **S3 Integration**: Support for S3 video URLs with bucket owner specification
+
+### Usage with S3 Video
+
+<Tabs>
+<TabItem value="sdk" label="SDK">
+
+```python title="TwelveLabs Pegasus SDK Usage" showLineNumbers
+from litellm import completion
+import os
+
+# Set AWS credentials
+os.environ["AWS_ACCESS_KEY_ID"] = "your-aws-access-key"
+os.environ["AWS_SECRET_ACCESS_KEY"] = "your-aws-secret-key"
+os.environ["AWS_REGION_NAME"] = "us-east-1"
+
+response = completion(
+    model="bedrock/us.twelvelabs.pegasus-1-2-v1:0",
+    messages=[{"role": "user", "content": "Describe what happens in this video."}],
+    mediaSource={
+        "s3Location": {
+            "uri": "s3://your-bucket/video.mp4",
+            "bucketOwner": "123456789012",  # 12-digit AWS account ID
+        }
+    },
+    temperature=0.2
+)
+
+print(response.choices[0].message.content)
+```
+
+</TabItem>
+
+<TabItem value="proxy" label="Proxy">
+
+**1. Add to config**
+
+```yaml title="config.yaml" showLineNumbers
+model_list:
+  - model_name: pegasus-video
+    litellm_params:
+      model: bedrock/us.twelvelabs.pegasus-1-2-v1:0
+      aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
+      aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
+      aws_region_name: os.environ/AWS_REGION_NAME
+```
+
+**2. Start proxy**
+
+```bash title="Start LiteLLM Proxy" showLineNumbers
+litellm --config /path/to/config.yaml
+
+# RUNNING at http://0.0.0.0:4000
+```
+
+**3. Test it!**
+
+```bash title="Test Pegasus via Proxy" showLineNumbers
+curl --location 'http://0.0.0.0:4000/chat/completions' \
+  --header 'Authorization: Bearer sk-1234' \
+  --header 'Content-Type: application/json' \
+  --data '{
+    "model": "pegasus-video",
+    "messages": [
+      {
+        "role": "user",
+        "content": "Describe what happens in this video."
+      }
+    ],
+    "mediaSource": {
+      "s3Location": {
+        "uri": "s3://your-bucket/video.mp4",
+        "bucketOwner": "123456789012"
+      }
+    },
+    "temperature": 0.2
+  }'
+```
+
+</TabItem>
+</Tabs>
+
+### Usage with Base64 Video
+
+You can also pass video content directly as base64:
+
+```python title="Base64 Video Input" showLineNumbers
+from litellm import completion
+import base64
+
+# Read video file and encode to base64
+with open("video.mp4", "rb") as video_file:
+    video_base64 = base64.b64encode(video_file.read()).decode("utf-8")
+
+response = completion(
+    model="bedrock/us.twelvelabs.pegasus-1-2-v1:0",
+    messages=[{"role": "user", "content": "What is happening in this video?"}],
+    mediaSource={
+        "base64String": video_base64
+    },
+    temperature=0.2,
+)
+
+print(response.choices[0].message.content)
+```
+
+### Important Notes
+
+- **Response Format**: The model supports structured output via `response_format` with JSON schema
+
 ## Provisioned throughput models
 To use provisioned throughput Bedrock models pass 
 - `model=bedrock/<base-model>`, example `model=bedrock/anthropic.claude-v2`. Set `model` to any of the [Supported AWS models](#supported-aws-bedrock-models)
@@ -1743,6 +1868,8 @@ Here's an example of using a bedrock model with LiteLLM. For a complete list, re
 | Meta Llama 2 Chat 70b      | `completion(model='bedrock/meta.llama2-70b-chat-v1', messages=messages)`   | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
 | Mistral 7B Instruct        | `completion(model='bedrock/mistral.mistral-7b-instruct-v0:2', messages=messages)`   | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
 | Mixtral 8x7B Instruct      | `completion(model='bedrock/mistral.mixtral-8x7b-instruct-v0:1', messages=messages)`   | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
+| TwelveLabs Pegasus 1.2 (US) | `completion(model='bedrock/us.twelvelabs.pegasus-1-2-v1:0', messages=messages, mediaSource={...})`   | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
+| TwelveLabs Pegasus 1.2 (EU) | `completion(model='bedrock/eu.twelvelabs.pegasus-1-2-v1:0', messages=messages, mediaSource={...})`   | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
 
 
 ## Bedrock Embedding

litellm/__init__.py

@@ -1234,6 +1234,9 @@ from .llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation imp
 from .llms.bedrock.chat.invoke_transformations.amazon_titan_transformation import (
     AmazonTitanConfig,
 )
+from .llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import (
+    AmazonTwelveLabsPegasusConfig,
+)
 from .llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
     AmazonInvokeConfig,
 )

litellm/constants.py

@@ -856,6 +856,7 @@ BEDROCK_INVOKE_PROVIDERS_LITERAL = Literal[
     "nova",
     "deepseek_r1",
     "qwen3",
+    "twelvelabs",
 ]
 
 BEDROCK_EMBEDDING_PROVIDERS_LITERAL = Literal[

litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py

@@ -0,0 +1,280 @@
+"""
+Transforms OpenAI-style requests into TwelveLabs Pegasus 1.2 requests for Bedrock.
+
+Reference:
+https://docs.twelvelabs.io/docs/models/pegasus
+"""
+
+import json
+import time
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+import httpx
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.core_helpers import map_finish_reason
+from litellm.llms.base_llm.base_utils import type_to_response_format_param
+from litellm.llms.base_llm.chat.transformation import BaseConfig
+from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
+    AmazonInvokeConfig,
+)
+from litellm.llms.bedrock.common_utils import BedrockError
+from litellm.types.llms.openai import AllMessageValues
+from litellm.types.utils import ModelResponse, Usage
+from litellm.utils import get_base64_str
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
+
+    LiteLLMLoggingObj = _LiteLLMLoggingObj
+else:
+    LiteLLMLoggingObj = Any
+
+
+class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
+    """
+    Handles transforming OpenAI-style requests into Bedrock InvokeModel requests for
+    `twelvelabs.pegasus-1-2-v1:0`.
+
+    Pegasus 1.2 requires an `inputPrompt` and a `mediaSource` that either references
+    an S3 object or a base64-encoded clip. Optional OpenAI params (temperature,
+    response_format, max_tokens) are translated to the TwelveLabs schema.
+    """
+
+    def get_supported_openai_params(self, model: str) -> List[str]:
+        return [
+            "max_tokens",
+            "max_completion_tokens",
+            "temperature",
+            "response_format",
+        ]
+
+    def map_openai_params(
+        self,
+        non_default_params: dict,
+        optional_params: dict,
+        model: str,
+        drop_params: bool,
+    ) -> dict:
+        for param, value in non_default_params.items():
+            if param in {"max_tokens", "max_completion_tokens"}:
+                optional_params["maxOutputTokens"] = value
+            if param == "temperature":
+                optional_params["temperature"] = value
+            if param == "response_format":
+                optional_params["responseFormat"] = self._normalize_response_format(
+                    value
+                )
+        return optional_params
+
+    def _normalize_response_format(self, value: Any) -> Any:
+        """Normalize response_format to TwelveLabs format.
+        
+        TwelveLabs expects:
+        {
+            "jsonSchema": {...}
+        }
+        
+        But OpenAI format is:
+        {
+            "type": "json_schema",
+            "json_schema": {
+                "name": "...",
+                "schema": {...}
+            }
+        }
+        """
+        if isinstance(value, dict):
+            # If it has json_schema field, extract and transform it
+            if "json_schema" in value:
+                json_schema = value["json_schema"]
+                # Extract the schema if nested
+                if isinstance(json_schema, dict) and "schema" in json_schema:
+                    return {"jsonSchema": json_schema["schema"]}
+                # Otherwise use json_schema directly
+                return {"jsonSchema": json_schema}
+            # If it already has jsonSchema, return as is
+            if "jsonSchema" in value:
+                return value
+            # Otherwise return the dict as is
+            return value
+        return type_to_response_format_param(response_format=value) or value
+
+    def transform_request(
+        self,
+        model: str,
+        messages: List[AllMessageValues],
+        optional_params: dict,
+        litellm_params: dict,
+        headers: dict,
+    ) -> dict:
+        input_prompt = self._convert_messages_to_prompt(messages=messages)
+        request_data: Dict[str, Any] = {"inputPrompt": input_prompt}
+
+        media_source = self._build_media_source(optional_params)
+        if media_source is not None:
+            request_data["mediaSource"] = media_source
+
+        # Handle temperature and maxOutputTokens
+        for key in ("temperature", "maxOutputTokens"):
+            if key in optional_params:
+                request_data[key] = optional_params.get(key)
+        
+        # Handle responseFormat - transform to TwelveLabs format
+        if "responseFormat" in optional_params:
+            response_format = optional_params["responseFormat"]
+            transformed_format = self._normalize_response_format(response_format)
+            if transformed_format:
+                request_data["responseFormat"] = transformed_format
+        
+        return request_data
+
+    def _build_media_source(self, optional_params: dict) -> Optional[dict]:
+        direct_source = optional_params.get("mediaSource") or optional_params.get(
+            "media_source"
+        )
+        if isinstance(direct_source, dict):
+            return direct_source
+
+        base64_input = optional_params.get("video_base64") or optional_params.get(
+            "base64_string"
+        )
+        if base64_input:
+            return {"base64String": get_base64_str(base64_input)}
+
+        s3_uri = (
+            optional_params.get("video_s3_uri")
+            or optional_params.get("s3_uri")
+            or optional_params.get("media_source_s3_uri")
+        )
+        if s3_uri:
+            s3_location = {"uri": s3_uri}
+            bucket_owner = (
+                optional_params.get("video_s3_bucket_owner")
+                or optional_params.get("s3_bucket_owner")
+                or optional_params.get("media_source_bucket_owner")
+            )
+            if bucket_owner:
+                s3_location["bucketOwner"] = bucket_owner
+            return {"s3Location": s3_location}
+        return None
+
+    def _convert_messages_to_prompt(self, messages: List[AllMessageValues]) -> str:
+        prompt_parts: List[str] = []
+        for message in messages:
+            role = message.get("role", "user")
+            content = message.get("content", "")
+            if isinstance(content, list):
+                text_fragments = []
+                for item in content:
+                    if isinstance(item, dict):
+                        item_type = item.get("type")
+                        if item_type == "text":
+                            text_fragments.append(item.get("text", ""))
+                        elif item_type == "image_url":
+                            text_fragments.append("<image>")
+                        elif item_type == "video_url":
+                            text_fragments.append("<video>")
+                        elif item_type == "audio_url":
+                            text_fragments.append("<audio>")
+                    elif isinstance(item, str):
+                        text_fragments.append(item)
+                content = " ".join(text_fragments)
+            prompt_parts.append(f"{role}: {content}")
+        return "\n".join(part for part in prompt_parts if part).strip()
+
+    def transform_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        model_response: ModelResponse,
+        logging_obj: LiteLLMLoggingObj,
+        request_data: dict,
+        messages: List[AllMessageValues],
+        optional_params: dict,
+        litellm_params: dict,
+        encoding: Any,
+        api_key: Optional[str] = None,
+        json_mode: Optional[bool] = None,
+    ) -> ModelResponse:
+        """
+        Transform TwelveLabs Pegasus response to LiteLLM format.
+        
+        TwelveLabs response format:
+        {
+            "message": "...",
+            "finishReason": "stop" | "length"
+        }
+        
+        LiteLLM format:
+        ModelResponse with choices[0].message.content and finish_reason
+        """
+        try:
+            completion_response = raw_response.json()
+        except Exception as e:
+            raise BedrockError(
+                message=f"Error parsing response: {raw_response.text}, error: {str(e)}",
+                status_code=raw_response.status_code,
+            )
+        
+        verbose_logger.debug(
+            "twelvelabs pegasus response: %s",
+            json.dumps(completion_response, indent=4, default=str),
+        )
+        
+        # Extract message content
+        message_content = completion_response.get("message", "")
+        
+        # Extract finish reason and map to LiteLLM format
+        finish_reason_raw = completion_response.get("finishReason", "stop")
+        finish_reason = map_finish_reason(finish_reason_raw)
+        
+        # Set the response content
+        try:
+            if (
+                message_content
+                and hasattr(model_response.choices[0], "message")
+                and getattr(model_response.choices[0].message, "tool_calls", None) is None
+            ):
+                model_response.choices[0].message.content = message_content  # type: ignore
+                model_response.choices[0].finish_reason = finish_reason
+            else:
+                raise Exception("Unable to set message content")
+        except Exception as e:
+            raise BedrockError(
+                message=f"Error setting response content: {str(e)}. Response: {completion_response}",
+                status_code=raw_response.status_code,
+            )
+        
+        # Calculate usage from headers
+        bedrock_input_tokens = raw_response.headers.get(
+            "x-amzn-bedrock-input-token-count", None
+        )
+        bedrock_output_tokens = raw_response.headers.get(
+            "x-amzn-bedrock-output-token-count", None
+        )
+        
+        prompt_tokens = int(
+            bedrock_input_tokens or litellm.token_counter(messages=messages)
+        )
+        
+        completion_tokens = int(
+            bedrock_output_tokens
+            or litellm.token_counter(
+                text=model_response.choices[0].message.content,  # type: ignore
+                count_response_tokens=True,
+            )
+        )
+        
+        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
+

litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py

@@ -250,6 +250,14 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
                 ):  # completion(top_k=3) > anthropic_config(top_k=3) <- allows for dynamic variables to be passed in
                     inference_params[k] = v
             request_data = {"prompt": prompt, **inference_params}
+        elif provider == "twelvelabs":
+            return litellm.AmazonTwelveLabsPegasusConfig().transform_request(
+                model=model,
+                messages=messages,
+                optional_params=optional_params,
+                litellm_params=litellm_params,
+                headers=headers,
+            )
         else:
             raise BedrockError(
                 status_code=404,
@@ -321,6 +329,20 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
                     litellm_params=litellm_params,
                     encoding=encoding,
                 )
+            elif provider == "twelvelabs":
+                return litellm.AmazonTwelveLabsPegasusConfig().transform_response(
+                    model=model,
+                    raw_response=raw_response,
+                    model_response=model_response,
+                    logging_obj=logging_obj,
+                    request_data=request_data,
+                    messages=messages,
+                    optional_params=optional_params,
+                    litellm_params=litellm_params,
+                    encoding=encoding,
+                    api_key=api_key,
+                    json_mode=json_mode,
+                )
             elif provider == "ai21":
                 outputText = (
                     completion_response.get("completions")[0].get("data").get("text")

litellm/llms/bedrock/common_utils.py

@@ -616,6 +616,8 @@ def get_bedrock_chat_config(model: str):
         return litellm.AmazonInvokeNovaConfig()
     elif bedrock_invoke_provider == "qwen3":
         return litellm.AmazonQwen3Config()
+    elif bedrock_invoke_provider == "twelvelabs":
+        return litellm.AmazonTwelveLabsPegasusConfig()
     else:
         return litellm.AmazonInvokeConfig()
 

tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py

@@ -0,0 +1,87 @@
+from litellm.llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import (
+    AmazonTwelveLabsPegasusConfig,
+)
+
+
+def _make_messages() -> list[dict]:
+    return [
+        {"role": "system", "content": "You are an assistant"},
+        {"role": "user", "content": "Summarize the attached video."},
+    ]
+
+
+def test_supported_openai_params():
+    config = AmazonTwelveLabsPegasusConfig()
+    supported = config.get_supported_openai_params("twelvelabs.pegasus-1-2-v1:0")
+    assert "max_tokens" in supported
+    assert "temperature" in supported
+    assert "response_format" in supported
+
+
+def test_map_openai_params_translates_fields():
+    config = AmazonTwelveLabsPegasusConfig()
+    optional_params: dict = {}
+    config.map_openai_params(
+        non_default_params={
+            "max_tokens": 20,
+            "temperature": 0.6,
+            "response_format": {
+                "type": "json_schema",
+                "json_schema": {"name": "video_schema", "schema": {"type": "object"}},
+            },
+        },
+        optional_params=optional_params,
+        model="twelvelabs.pegasus-1-2-v1:0",
+        drop_params=False,
+    )
+
+    assert optional_params["maxOutputTokens"] == 20
+    assert optional_params["temperature"] == 0.6
+    assert "responseFormat" in optional_params
+    # TwelveLabs format: responseFormat contains jsonSchema directly (not json_schema)
+    assert "jsonSchema" in optional_params["responseFormat"]
+    assert optional_params["responseFormat"]["jsonSchema"]["type"] == "object"
+
+
+def test_transform_request_includes_base64_media():
+    config = AmazonTwelveLabsPegasusConfig()
+    optional_params = config.map_openai_params(
+        non_default_params={"max_tokens": 10},
+        optional_params={},
+        model="twelvelabs.pegasus-1-2-v1:0",
+        drop_params=False,
+    )
+    optional_params["video_base64"] = "data:video/mp4;base64,AAA"
+
+    request = config.transform_request(
+        model="twelvelabs.pegasus-1-2-v1:0",
+        messages=_make_messages(),
+        optional_params=optional_params,
+        litellm_params={},
+        headers={},
+    )
+
+    assert request["inputPrompt"].startswith("system:")
+    assert request["mediaSource"]["base64String"] == "AAA"
+    assert request["maxOutputTokens"] == 10
+
+
+def test_transform_request_includes_s3_media():
+    config = AmazonTwelveLabsPegasusConfig()
+    optional_params = {
+        "video_s3_uri": "s3://test-bucket/video.mp4",
+        "video_s3_bucket_owner": "123456789012",
+    }
+
+    request = config.transform_request(
+        model="twelvelabs.pegasus-1-2-v1:0",
+        messages=_make_messages(),
+        optional_params=optional_params,
+        litellm_params={},
+        headers={},
+    )
+
+    s3_location = request["mediaSource"]["s3Location"]
+    assert s3_location["uri"] == "s3://test-bucket/video.mp4"
+    assert s3_location["bucketOwner"] == "123456789012"
+