Fix pegasus response and add doc

This commit is contained in:
Sameer Kankute
2025-11-28 21:41:25 +05:30
parent 247160277e
commit 9d058398df
4 changed files with 301 additions and 3 deletions
+127
View File
@@ -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
@@ -5,16 +5,32 @@ Reference:
https://docs.twelvelabs.io/docs/models/pegasus
"""
from typing import Any, Dict, List, Optional
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):
"""
@@ -53,7 +69,35 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
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
@@ -72,9 +116,18 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
if media_source is not None:
request_data["mediaSource"] = media_source
for key in ("temperature", "maxOutputTokens", "responseFormat"):
# 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]:
@@ -131,3 +184,97 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
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
@@ -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")
@@ -38,7 +38,9 @@ def test_map_openai_params_translates_fields():
assert optional_params["maxOutputTokens"] == 20
assert optional_params["temperature"] == 0.6
assert "responseFormat" in optional_params
assert optional_params["responseFormat"]["json_schema"]["name"] == "video_schema"
# 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():