tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 22:25:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add support for computer use for gemini

Browse Source
This commit is contained in:
Sameer Kankute
2025-12-10 10:34:08 +05:30
parent 254c1155a2
commit bcac9e41f6
8 changed files with 555 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/providers/gemini.md
+159
View File
@@ -1019,7 +1019,166 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \
</Tabs>
### Computer Use Tool
<Tabs>
<TabItem value="sdk" label="LiteLLM Python SDK">
```python
from litellm import completion
import os
os.environ["GEMINI_API_KEY"] = "your-api-key"
# Computer Use tool with browser environment
tools = [
{
"type": "computer_use",
"environment": "browser", # optional: "browser" or "unspecified"
"excluded_predefined_functions": ["drag_and_drop"] # optional
}
]
messages = [
{
"role": "user",
"content": [
{
"type": "text",
"text": "Navigate to google.com and search for 'LiteLLM'"
},
{
"type": "image_url",
"image_url": {
"url": "data:image/png;base64,..." # screenshot of current browser state
}
}
]
}
]
response = completion(
model="gemini/gemini-2.5-computer-use-preview-10-2025",
messages=messages,
tools=tools,
)
print(response)
# Handling tool responses with screenshots
# When the model makes a tool call, send the response back with a screenshot:
if response.choices[0].message.tool_calls:
tool_call = response.choices[0].message.tool_calls[0]
# Add assistant message with tool call
messages.append(response.choices[0].message.model_dump())
# Add tool response with screenshot
messages.append({
"role": "tool",
"tool_call_id": tool_call.id,
"content": [
{
"type": "text",
"text": '{"url": "https://example.com", "status": "completed"}'
},
{
"type": "input_image",
"image_url": "data:image/png;base64,..." # New screenshot after action (Can send an image url as well, litellm handles the conversion)
}
]
})
# Continue conversation with updated screenshot
response = completion(
model="gemini/gemini-2.5-computer-use-preview-10-2025",
messages=messages,
tools=tools,
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM Proxy Server">
1. Add model to config.yaml
```yaml
model_list:
- model_name: gemini-computer-use
litellm_params:
model: gemini/gemini-2.5-computer-use-preview-10-2025
api_key: os.environ/GEMINI_API_KEY
```
2. Start proxy
```bash
litellm --config /path/to/config.yaml
```
3. Make request
```bash
curl http://0.0.0.0:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-1234" \
-d '{
"model": "gemini-computer-use",
"messages": [
{
"role": "user",
"content": [
{
"type": "text",
"text": "Click on the search button"
},
{
"type": "image_url",
"image_url": {
"url": "data:image/png;base64,..."
}
}
]
}
],
"tools": [
{
"type": "computer_use",
"environment": "browser"
}
]
}'
```
**Tool Response Format:**
When responding to Computer Use tool calls, include the URL and screenshot:
```json
{
"role": "tool",
"tool_call_id": "call_abc123",
"content": [
{
"type": "text",
"text": "{\"url\": \"https://example.com\", \"status\": \"completed\"}"
},
{
"type": "input_image",
"image_url": "data:image/png;base64,..."
}
]
}
```
### Environment Mapping
| LiteLLM Input | Gemini API Value |
|--------------|------------------|
| `"browser"` | `ENVIRONMENT_BROWSER` |
| `"unspecified"` | `ENVIRONMENT_UNSPECIFIED` |
| `ENVIRONMENT_BROWSER` | `ENVIRONMENT_BROWSER` (passed through) |
| `ENVIRONMENT_UNSPECIFIED` | `ENVIRONMENT_UNSPECIFIED` (passed through) |
litellm/litellm_core_utils/prompt_templates/factory.py
+75 -12
View File
@@ -6,7 +6,7 @@ import mimetypes
import re
import xml.etree.ElementTree as ET
from enum import Enum
from typing import Any, List, Optional, Tuple, Union, cast, overload
from typing import Any, Dict, List, Optional, Tuple, Union, cast, overload
from jinja2.sandbox import ImmutableSandboxedEnvironment
@@ -1455,7 +1455,7 @@ def convert_to_gemini_tool_call_invoke(
def convert_to_gemini_tool_call_result(
message: Union[ChatCompletionToolMessage, ChatCompletionFunctionMessage],
last_message_with_tool_calls: Optional[dict],
) -> VertexPartType:
) -> Union[VertexPartType, List[VertexPartType]]:
"""
OpenAI message with a tool result looks like:
{
@@ -1471,16 +1471,47 @@ def convert_to_gemini_tool_call_result(
"name": "get_current_weather",
"content": "function result goes here",
}
Supports content with images for Computer Use:
{
"role": "tool",
"tool_call_id": "call_abc123",
"content": [
{"type": "text", "text": "I found the requested image:"},
{"type": "input_image", "image_url": "https://example.com/image.jpg" }
]
}
"""
from litellm.types.llms.vertex_ai import BlobType
content_str: str = ""
inline_data: Optional[BlobType] = None
if "content" in message:
if isinstance(message["content"], str):
content_str = message["content"]
elif isinstance(message["content"], List):
content_list = message["content"]
for content in content_list:
if content["type"] == "text":
content_str += content["text"]
content_type = content.get("type", "")
if content_type == "text":
content_str += content.get("text", "")
elif content_type == "input_image":
# Extract image for inline_data (for Computer Use screenshots)
image_url = content.get("image_url", "")
if image_url:
# Convert image to base64 blob format for Gemini
try:
image_obj = convert_to_anthropic_image_obj(image_url, format=None)
inline_data = BlobType(
data=image_obj["data"],
mime_type=image_obj["media_type"]
)
except Exception as e:
verbose_logger.warning(
f"Failed to process image in tool response: {e}"
)
name: Optional[str] = message.get("name", "") # type: ignore
# Recover name from last message with tool calls
@@ -1503,14 +1534,41 @@ def convert_to_gemini_tool_call_result(
)
)
# Parse response data - support both JSON string and plain string
# For Computer Use, the response should contain structured data like {"url": "..."}
response_data: dict
try:
import json
if content_str.strip().startswith("{") or content_str.strip().startswith("["):
# Try to parse as JSON (for Computer Use structured responses)
parsed = json.loads(content_str)
if isinstance(parsed, dict):
response_data = parsed # Use the parsed JSON directly
else:
response_data = {"content": content_str}
else:
response_data = {"content": content_str}
except (json.JSONDecodeError, ValueError):
# Not valid JSON, wrap in content field
response_data = {"content": content_str}
# We can't determine from openai message format whether it's a successful or
# error call result so default to the successful result template
_function_response = VertexFunctionResponse(
name=name, response={"content": content_str} # type: ignore
name=name, response=response_data # type: ignore
)
_part = VertexPartType(function_response=_function_response)
# Create part with function_response, and optionally inline_data for images (Computer Use)
_part: VertexPartType = {"function_response": _function_response}
# For Computer Use, if we have an image, we need separate parts:
# - One part with function_response
# - One part with inline_data
# Gemini's PartType is a oneof, so we can't have both in the same part
if inline_data:
image_part: VertexPartType = {"inline_data": inline_data}
return [_part, image_part]
return _part
@@ -2085,13 +2143,18 @@ def anthropic_messages_pt( # noqa: PLR0915
): # support assistant tool invoke conversion
# Get web_search_results from provider_specific_fields for server_tool_use reconstruction
# Fixes: https://github.com/BerriAI/litellm/issues/17737
_provider_specific_fields = assistant_content_block.get("provider_specific_fields") or {}
_provider_specific_fields_raw = assistant_content_block.get("provider_specific_fields")
_provider_specific_fields: Dict[str, Any] = {}
if isinstance(_provider_specific_fields_raw, dict):
_provider_specific_fields = cast(Dict[str, Any], _provider_specific_fields_raw)
_web_search_results = _provider_specific_fields.get("web_search_results")
tool_invoke_results = convert_to_anthropic_tool_invoke(
assistant_tool_calls,
web_search_results=_web_search_results,
)
# AnthropicMessagesAssistantMessageValues includes AnthropicMessagesToolUseParam
assistant_content.extend(
convert_to_anthropic_tool_invoke(
assistant_tool_calls,
web_search_results=_web_search_results,
)
cast(List[AnthropicMessagesAssistantMessageValues], tool_invoke_results)
)
assistant_function_call = assistant_content_block.get("function_call")
litellm/llms/vertex_ai/gemini/transformation.py
+6 -2
View File
@@ -3,7 +3,7 @@ Transformation logic from OpenAI format to Gemini format.
Why separate file? Make it easy to see how transformation works
"""
import json
import os
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union, cast
@@ -418,7 +418,11 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
messages[msg_i], last_message_with_tool_calls # type: ignore
)
msg_i += 1
tool_call_responses.append(_part)
# Handle both single part and list of parts (for Computer Use with images)
if isinstance(_part, list):
tool_call_responses.extend(_part)
else:
tool_call_responses.append(_part)
if msg_i < len(messages) and (
messages[msg_i]["role"] not in tool_call_message_roles
):
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+61 -2
View File
@@ -309,6 +309,44 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"""
return Tools(googleSearch={})
def _transform_computer_use_config(
self, computer_use_config: dict
) -> dict:
"""
Transform Computer Use configuration to Gemini API format.
Args:
computer_use_config: The computer use configuration from LiteLLM
Returns:
Transformed computer use configuration for Gemini API
"""
transformed_config = {}
# Transform environment values if needed
if "environment" in computer_use_config:
env_value = computer_use_config["environment"]
if env_value == "browser":
transformed_config["environment"] = "ENVIRONMENT_BROWSER"
elif env_value == "unspecified":
transformed_config["environment"] = "ENVIRONMENT_UNSPECIFIED"
elif env_value in ["ENVIRONMENT_BROWSER", "ENVIRONMENT_UNSPECIFIED"]:
# Already in correct format
transformed_config["environment"] = env_value
else:
verbose_logger.info(
f"Invalid environment value for computer_use: {env_value}. "
f"Supported: 'browser', 'unspecified', 'ENVIRONMENT_BROWSER', 'ENVIRONMENT_UNSPECIFIED'"
)
# Transform excluded_predefined_functions to camelCase
if "excluded_predefined_functions" in computer_use_config:
transformed_config["excludedPredefinedFunctions"] = computer_use_config["excluded_predefined_functions"]
elif "excludedPredefinedFunctions" in computer_use_config:
transformed_config["excludedPredefinedFunctions"] = computer_use_config["excludedPredefinedFunctions"]
return transformed_config
def _extract_google_maps_retrieval_config(
self, google_maps_config: dict
) -> Tuple[dict, Optional[dict]]:
@@ -400,6 +438,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
code_execution: Optional[dict] = None
googleMaps: Optional[dict] = None
google_maps_retrieval_config: Optional[dict] = None
computerUse: Optional[dict] = None
# remove 'additionalProperties' from tools
value = _remove_additional_properties(value)
# remove 'strict' from tools
@@ -428,10 +467,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
elif "name" in tool: # functions list
openai_function_object = ChatCompletionToolParamFunctionChunk(**tool) # type: ignore
if "type" in tool and tool["type"] == "computer_use":
computer_use_config = {k: v for k, v in tool.items() if k != "type"}
tool = {VertexToolName.COMPUTER_USE.value: computer_use_config}
# Handle tools with 'type' field (OpenAI spec compliance) Ignore this field -> https://github.com/BerriAI/litellm/issues/14644#issuecomment-3342061838
if "type" in tool:
elif "type" in tool:
tool = {k: tool[k] for k in tool if k != "type"}
tool_name = list(tool.keys())[0] if len(tool.keys()) == 1 else None
if tool_name and (
tool_name == "codeExecution"
@@ -473,6 +514,22 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
) = self._extract_google_maps_retrieval_config(
google_maps_config=google_maps_value
)
elif tool_name and (
tool_name == VertexToolName.COMPUTER_USE.value
or tool_name == "computer_use"
):
computer_use_value = self.get_tool_value(
tool, VertexToolName.COMPUTER_USE.value
)
# Transform Computer Use configuration to Gemini API format
if computer_use_value is not None:
computerUse = self._transform_computer_use_config(
computer_use_config=computer_use_value
)
else:
# Empty config - Gemini will use defaults
computerUse = {}
elif openai_function_object is not None:
gtool_func_declaration = FunctionDeclaration(
name=openai_function_object["name"],
@@ -510,6 +567,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
_tools[VertexToolName.URL_CONTEXT.value] = urlContext
if googleMaps is not None:
_tools[VertexToolName.GOOGLE_MAPS.value] = googleMaps
if computerUse is not None:
_tools[VertexToolName.COMPUTER_USE.value] = computerUse
# Add retrieval config to toolConfig if googleMaps has location data
if google_maps_retrieval_config is not None:
litellm/model_prices_and_context_window_backup.json
+31
View File
@@ -14428,6 +14428,37 @@
"supports_web_search": true,
"tpm": 800000
},
"gemini/gemini-2.5-computer-use-preview-10-2025": {
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_above_200k_tokens": 2.5e-06,
"litellm_provider": "gemini",
"max_images_per_prompt": 3000,
"max_input_tokens": 128000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_above_200k_tokens": 1.5e-05,
"rpm": 2000,
"source": "https://ai.google.dev/gemini-api/docs/computer-use",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/completions"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true,
"tpm": 800000
},
"gemini/gemini-3-pro-preview": {
"cache_read_input_token_cost": 2e-07,
"cache_read_input_token_cost_above_200k_tokens": 4e-07,
litellm/types/llms/vertex_ai.py
+3 -8
View File
@@ -1,16 +1,9 @@
import json
from enum import Enum
from typing import Any, Dict, List, Literal, Optional, Tuple, Union
from typing import Any, Dict, List, Literal, Optional, Union
from typing_extensions import (
Protocol,
Required,
Self,
TypedDict,
TypeGuard,
get_origin,
override,
runtime_checkable,
)
@@ -230,6 +223,7 @@ class VertexToolName(str, Enum):
URL_CONTEXT = "url_context"
CODE_EXECUTION = "code_execution"
GOOGLE_MAPS = "googleMaps"
COMPUTER_USE = "computerUse"
class Tools(TypedDict, total=False):
@@ -240,6 +234,7 @@ class Tools(TypedDict, total=False):
url_context: dict
code_execution: dict
googleMaps: dict
computerUse: dict
retrieval: Retrieval
model_prices_and_context_window.json
+31
View File
@@ -14428,6 +14428,37 @@
"supports_web_search": true,
"tpm": 800000
},
"gemini/gemini-2.5-computer-use-preview-10-2025": {
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_above_200k_tokens": 2.5e-06,
"litellm_provider": "gemini",
"max_images_per_prompt": 3000,
"max_input_tokens": 128000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_above_200k_tokens": 1.5e-05,
"rpm": 2000,
"source": "https://ai.google.dev/gemini-api/docs/computer-use",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/completions"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true,
"tpm": 800000
},
"gemini/gemini-3-pro-preview": {
"cache_read_input_token_cost": 2e-07,
"cache_read_input_token_cost_above_200k_tokens": 4e-07,
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py
+189 -6
View File
@@ -1,8 +1,12 @@
from litellm.llms.vertex_ai.gemini.transformation import (
check_if_part_exists_in_parts,
_transform_request_body,
_gemini_convert_messages_with_history,
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_gemini_tool_call_result,
)
from litellm.llms.vertex_ai.gemini.transformation import (
_gemini_convert_messages_with_history,
_transform_request_body,
check_if_part_exists_in_parts,
)
from litellm.types.llms.vertex_ai import BlobType
def test_check_if_part_exists_in_parts():
@@ -394,10 +398,11 @@ def test_thought_signature_with_function_call_mode():
def test_dummy_signature_added_for_gemini_3_conversation_history():
"""Test that dummy signatures are added when transferring conversation history from older models (like gemini-2.5-flash) to gemini-3."""
import base64
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_gemini_tool_call_invoke,
)
import base64
# Simulate conversation history from gemini-2.5-flash (no thought signature)
assistant_message_from_older_model = {
@@ -509,10 +514,11 @@ def test_dummy_signature_not_added_when_signature_exists():
def test_dummy_signature_with_function_call_mode():
"""Test that dummy signatures are added for function_call mode when converting to gemini-3."""
import base64
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_gemini_tool_call_invoke,
)
import base64
# Assistant message with function_call (not tool_calls) and no signature
assistant_message_function_call = {
@@ -538,3 +544,180 @@ def test_dummy_signature_with_function_call_mode():
# Verify it's the expected dummy signature
expected_dummy = base64.b64encode(b"skip_thought_signature_validator").decode("utf-8")
assert gemini_parts[0]["thoughtSignature"] == expected_dummy
def test_convert_tool_response_with_base64_image():
"""Test tool response with base64 data URI image."""
# Create a small test image (1x1 red pixel PNG)
test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
image_data_uri = f"data:image/png;base64,{test_image_base64}"
# Create tool message with image
tool_message = {
"role": "tool",
"tool_call_id": "call_test123",
"content": [
{
"type": "text",
"text": '{"url": "https://example.com", "status": "success"}'
},
{
"type": "input_image",
"image_url": image_data_uri
}
]
}
# Mock last message with tool calls
last_message_with_tool_calls = {
"tool_calls": [
{
"id": "call_test123",
"function": {
"name": "click_at",
"arguments": '{"x": 100, "y": 200}'
}
}
]
}
# Convert tool response (returns list when image is present)
result = convert_to_gemini_tool_call_result(
tool_message, last_message_with_tool_calls
)
# Verify results - should be a list with 2 parts (function_response + inline_data)
assert isinstance(result, list), f"Expected list when image present, got {type(result)}"
assert len(result) == 2, f"Expected 2 parts, got {len(result)}"
# Find function_response part and inline_data part
function_response_part = None
inline_data_part = None
for part in result:
if "function_response" in part:
function_response_part = part
elif "inline_data" in part:
inline_data_part = part
# Check function_response exists
assert function_response_part is not None, "Missing function_response part"
function_response = function_response_part["function_response"]
assert function_response["name"] == "click_at"
assert "response" in function_response
# Verify JSON response is parsed correctly
assert "url" in function_response["response"]
assert function_response["response"]["url"] == "https://example.com"
# Check inline_data exists
assert inline_data_part is not None, "Missing inline_data part"
inline_data: BlobType = inline_data_part["inline_data"]
assert "data" in inline_data
assert "mime_type" in inline_data
assert inline_data["mime_type"] == "image/png"
assert inline_data["data"] == test_image_base64
def test_convert_tool_response_with_url_image():
"""Test tool response with HTTP URL image (will download and convert)."""
import pytest
# Use a publicly accessible test image URL
test_image_url = "https://via.placeholder.com/1x1.png"
tool_message = {
"role": "tool",
"tool_call_id": "call_test456",
"content": [
{
"type": "text",
"text": '{"url": "https://example.com"}'
},
{
"type": "input_image",
"image_url": test_image_url
}
]
}
last_message_with_tool_calls = {
"tool_calls": [
{
"id": "call_test456",
"function": {
"name": "type_text_at",
"arguments": '{"x": 300, "y": 400, "text": "hello"}'
}
}
]
}
try:
result = convert_to_gemini_tool_call_result(
tool_message, last_message_with_tool_calls
)
# Should be a list with 2 parts when image is present
assert isinstance(result, list), f"Expected list when image present, got {type(result)}"
assert len(result) == 2, f"Expected 2 parts, got {len(result)}"
# Find parts
function_response_part = next(p for p in result if "function_response" in p)
inline_data_part = next(p for p in result if "inline_data" in p)
# Check function_response exists
assert function_response_part is not None, "Missing function_response part"
function_response = function_response_part["function_response"]
assert function_response["name"] == "type_text_at"
# Check inline_data exists (URL should be downloaded and converted)
assert inline_data_part is not None, "Missing inline_data part"
inline_data: BlobType = inline_data_part["inline_data"]
assert "data" in inline_data
assert "mime_type" in inline_data
except Exception as e:
# Skip test if URL download fails (no internet connection, etc.)
pytest.skip(f"Failed to download image from URL: {e}")
def test_convert_tool_response_text_only():
"""Test tool response with only text (no image)."""
tool_message = {
"role": "tool",
"tool_call_id": "call_test789",
"content": [
{
"type": "text",
"text": '{"status": "completed", "result": "success"}'
}
]
}
last_message_with_tool_calls = {
"tool_calls": [
{
"id": "call_test789",
"function": {
"name": "wait_5_seconds",
"arguments": "{}"
}
}
]
}
result = convert_to_gemini_tool_call_result(
tool_message, last_message_with_tool_calls
)
# Should be a single part (no list) when no image
assert not isinstance(result, list), "Should return single part when no image"
# Check function_response exists
assert "function_response" in result
function_response = result["function_response"]
assert function_response["name"] == "wait_5_seconds"
# Verify JSON response is parsed correctly
assert "status" in function_response["response"]
assert function_response["response"]["status"] == "completed"
# Check inline_data does NOT exist (no image provided)
assert "inline_data" not in result