cua-agent 0.1.20__tar.gz → 0.1.22__tar.gz

{cua_agent-0.1.20 → cua_agent-0.1.22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cua-agent
-Version: 0.1.20
+Version: 0.1.22
 Summary: CUA (Computer Use) Agent for AI-driven computer interaction
 Author-Email: TryCua <gh@trycua.com>
 Requires-Python: <3.13,>=3.10
@@ -33,6 +33,19 @@ Requires-Dist: openai<2.0.0,>=1.14.0; extra == "som"
 Requires-Dist: groq<0.5.0,>=0.4.0; extra == "som"
 Requires-Dist: dashscope<2.0.0,>=1.13.0; extra == "som"
 Requires-Dist: requests<3.0.0,>=2.31.0; extra == "som"
+Provides-Extra: omni
+Requires-Dist: torch>=2.2.1; extra == "omni"
+Requires-Dist: torchvision>=0.17.1; extra == "omni"
+Requires-Dist: ultralytics>=8.0.0; extra == "omni"
+Requires-Dist: transformers>=4.38.2; extra == "omni"
+Requires-Dist: cua-som<0.2.0,>=0.1.0; extra == "omni"
+Requires-Dist: anthropic<0.47.0,>=0.46.0; extra == "omni"
+Requires-Dist: boto3<2.0.0,>=1.35.81; extra == "omni"
+Requires-Dist: openai<2.0.0,>=1.14.0; extra == "omni"
+Requires-Dist: groq<0.5.0,>=0.4.0; extra == "omni"
+Requires-Dist: dashscope<2.0.0,>=1.13.0; extra == "omni"
+Requires-Dist: requests<3.0.0,>=2.31.0; extra == "omni"
+Requires-Dist: ollama<0.5.0,>=0.4.7; extra == "omni"
 Provides-Extra: all
 Requires-Dist: torch>=2.2.1; extra == "all"
 Requires-Dist: torchvision>=0.17.1; extra == "all"
@@ -45,6 +58,7 @@ Requires-Dist: openai<2.0.0,>=1.14.0; extra == "all"
 Requires-Dist: groq<0.5.0,>=0.4.0; extra == "all"
 Requires-Dist: dashscope<2.0.0,>=1.13.0; extra == "all"
 Requires-Dist: requests<3.0.0,>=2.31.0; extra == "all"
+Requires-Dist: ollama<0.5.0,>=0.4.7; extra == "all"
 Description-Content-Type: text/markdown
 
 <div align="center">
@@ -80,7 +94,7 @@ pip install "cua-agent[all]"
 # or install specific loop providers
 pip install "cua-agent[openai]" # OpenAI Cua Loop
 pip install "cua-agent[anthropic]" # Anthropic Cua Loop
-pip install "cua-agent[omni]" # Cua Loop based on OmniParser
+pip install "cua-agent[omni]" # Cua Loop based on OmniParser (includes Ollama for local models)
 ```
 
 ## Run
@@ -123,7 +137,7 @@ The `cua-agent` package provides three agent loops variations, based on differen
 |:-----------|:-----------------|:------------|:-------------|
 | `AgentLoop.OPENAI` | • `computer_use_preview` | Use OpenAI Operator CUA model | Not Required |
 | `AgentLoop.ANTHROPIC` | • `claude-3-5-sonnet-20240620`<br>• `claude-3-7-sonnet-20250219` | Use Anthropic Computer-Use | Not Required |
-| `AgentLoop.OMNI` <br>(preview) | • `claude-3-5-sonnet-20240620`<br>• `claude-3-7-sonnet-20250219`<br>• `gpt-4.5-preview`<br>• `gpt-4o`<br>• `gpt-4`<br>• `gpt-3.5-turbo` | Use OmniParser for element pixel-detection (SoM) and any VLMs | OmniParser |
+| `AgentLoop.OMNI` <br>(experimental) | • `claude-3-5-sonnet-20240620`<br>• `claude-3-7-sonnet-20250219`<br>• `gpt-4.5-preview`<br>• `gpt-4o`<br>• `gpt-4` | Use OmniParser for element pixel-detection (SoM) and any VLMs for UI Grounding and Reasoning | OmniParser |
 
 ## AgentResponse
 The `AgentResponse` class represents the structured output returned after each agent turn. It contains the agent's response, reasoning, tool usage, and other metadata. The response format aligns with the new [OpenAI Agent SDK specification](https://platform.openai.com/docs/api-reference/responses) for better consistency across different agent loops.

{cua_agent-0.1.20 → cua_agent-0.1.22}/README.md

@@ -31,7 +31,7 @@ pip install "cua-agent[all]"
 # or install specific loop providers
 pip install "cua-agent[openai]" # OpenAI Cua Loop
 pip install "cua-agent[anthropic]" # Anthropic Cua Loop
-pip install "cua-agent[omni]" # Cua Loop based on OmniParser
+pip install "cua-agent[omni]" # Cua Loop based on OmniParser (includes Ollama for local models)
 ```
 
 ## Run
@@ -74,7 +74,7 @@ The `cua-agent` package provides three agent loops variations, based on differen
 |:-----------|:-----------------|:------------|:-------------|
 | `AgentLoop.OPENAI` | • `computer_use_preview` | Use OpenAI Operator CUA model | Not Required |
 | `AgentLoop.ANTHROPIC` | • `claude-3-5-sonnet-20240620`<br>• `claude-3-7-sonnet-20250219` | Use Anthropic Computer-Use | Not Required |
-| `AgentLoop.OMNI` <br>(preview) | • `claude-3-5-sonnet-20240620`<br>• `claude-3-7-sonnet-20250219`<br>• `gpt-4.5-preview`<br>• `gpt-4o`<br>• `gpt-4`<br>• `gpt-3.5-turbo` | Use OmniParser for element pixel-detection (SoM) and any VLMs | OmniParser |
+| `AgentLoop.OMNI` <br>(experimental) | • `claude-3-5-sonnet-20240620`<br>• `claude-3-7-sonnet-20250219`<br>• `gpt-4.5-preview`<br>• `gpt-4o`<br>• `gpt-4` | Use OmniParser for element pixel-detection (SoM) and any VLMs for UI Grounding and Reasoning | OmniParser |
 
 ## AgentResponse
 The `AgentResponse` class represents the structured output returned after each agent turn. It contains the agent's response, reasoning, tool usage, and other metadata. The response format aligns with the new [OpenAI Agent SDK specification](https://platform.openai.com/docs/api-reference/responses) for better consistency across different agent loops.

{cua_agent-0.1.20 → cua_agent-0.1.22}/agent/core/agent.py

@@ -86,7 +86,8 @@ class ComputerAgent:
 
         # Get API key from environment if not provided
         actual_api_key = api_key or os.environ.get(ENV_VARS[self.provider], "")
-        if not actual_api_key:
+        # Ollama is local and doesn't require an API key
+        if not actual_api_key and str(self.provider) != "ollama":
             raise ValueError(f"No API key provided for {self.provider}")
 
         # Create the appropriate loop using the factory

{cua_agent-0.1.20 → cua_agent-0.1.22}/agent/core/provider_config.py

@@ -6,10 +6,12 @@ from ..providers.omni.types import LLMProvider
 DEFAULT_MODELS = {
     LLMProvider.OPENAI: "gpt-4o",
     LLMProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
+    LLMProvider.OLLAMA: "gemma3:4b-it-q4_K_M",
 }
 
 # Map providers to their environment variable names
 ENV_VARS = {
     LLMProvider.OPENAI: "OPENAI_API_KEY",
     LLMProvider.ANTHROPIC: "ANTHROPIC_API_KEY",
+    LLMProvider.OLLAMA: "OLLAMA_API_KEY",
 }

{cua_agent-0.1.20 → cua_agent-0.1.22}/agent/core/types.py

@@ -10,6 +10,7 @@ class AgentLoop(Enum):
     ANTHROPIC = auto()  # Anthropic implementation
     OMNI = auto()  # OmniLoop implementation
     OPENAI = auto()  # OpenAI implementation
+    OLLAMA = auto()  # OLLAMA implementation
     # Add more loop types as needed
 
 

{cua_agent-0.1.20 → cua_agent-0.1.22}/agent/providers/omni/clients/base.py

@@ -1,43 +1,34 @@
 """Base client implementation for Omni providers."""
 
-import os
 import logging
 from typing import Dict, List, Optional, Any, Tuple
-import aiohttp
-import json
 
 logger = logging.getLogger(__name__)
 
+
 class BaseOmniClient:
     """Base class for provider-specific clients."""
-    
-    def __init__(
-        self,
-        api_key: Optional[str] = None,
-        model: Optional[str] = None
-    ):
+
+    def __init__(self, api_key: Optional[str] = None, model: Optional[str] = None):
         """Initialize base client.
-        
+
         Args:
             api_key: Optional API key
             model: Optional model name
         """
         self.api_key = api_key
         self.model = model
-        
+
     async def run_interleaved(
-        self,
-        messages: List[Dict[str, Any]],
-        system: str,
-        max_tokens: Optional[int] = None
+        self, messages: List[Dict[str, Any]], system: str, max_tokens: Optional[int] = None
     ) -> Dict[str, Any]:
         """Run interleaved chat completion.
-        
+
         Args:
             messages: List of message dicts
             system: System prompt
             max_tokens: Optional max tokens override
-            
+
         Returns:
             Response dict
         """

cua_agent-0.1.22/agent/providers/omni/clients/ollama.py

@@ -0,0 +1,122 @@
+"""Ollama API client implementation."""
+
+import logging
+from typing import Any, Dict, List, Optional, Tuple, cast
+import asyncio
+from httpx import ConnectError, ReadTimeout
+
+from ollama import AsyncClient, Options
+from ollama import Message
+from .base import BaseOmniClient
+
+logger = logging.getLogger(__name__)
+
+
+class OllamaClient(BaseOmniClient):
+    """Client for making calls to Ollama API."""
+
+    def __init__(self, api_key: str, model: str, max_retries: int = 3, retry_delay: float = 1.0):
+        """Initialize the Ollama client.
+
+        Args:
+            api_key: Not used
+            model: Ollama model name (e.g. "gemma3:4b-it-q4_K_M")
+            max_retries: Maximum number of retries for API calls
+            retry_delay: Base delay between retries in seconds
+        """
+        if not model:
+            raise ValueError("Model name must be provided")
+
+        self.client = AsyncClient(
+            host="http://localhost:11434",
+        )
+        self.model: str = model  # Add explicit type annotation
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
+
+    def _convert_message_format(self, system: str, messages: List[Dict[str, Any]]) -> List[Any]:
+        """Convert messages from standard format to Ollama format.
+
+        Args:
+            messages: Messages in standard format
+
+        Returns:
+            Messages in Ollama format
+        """
+        ollama_messages = []
+
+        # Add system message
+        ollama_messages.append(
+            {
+                "role": "system",
+                "content": system,
+            }
+        )
+
+        for message in messages:
+            # Skip messages with empty content
+            if not message.get("content"):
+                continue
+            content = message.get("content", [{}])[0]
+            isImage = content.get("type", "") == "image_url"
+            isText = content.get("type", "") == "text"
+            if isText:
+                data = content.get("text", "")
+                ollama_messages.append({"role": message["role"], "content": data})
+            if isImage:
+                data = content.get("image_url", {}).get("url", "")
+                # remove header
+                data = data.removeprefix("data:image/png;base64,")
+                ollama_messages.append(
+                    {"role": message["role"], "content": "Use this image", "images": [data]}
+                )
+
+        # Cast the list to the correct type expected by Ollama
+        return cast(List[Any], ollama_messages)
+
+    async def run_interleaved(
+        self, messages: List[Dict[str, Any]], system: str, max_tokens: int
+    ) -> Any:
+        """Run model with interleaved conversation format.
+
+        Args:
+            messages: List of messages to process
+            system: System prompt
+            max_tokens: Not used
+
+        Returns:
+            Model response
+        """
+        last_error = None
+
+        for attempt in range(self.max_retries):
+            try:
+                # Convert messages to Ollama format
+                ollama_messages = self._convert_message_format(system, messages)
+
+                response = await self.client.chat(
+                    model=self.model,
+                    options=Options(
+                        temperature=0,
+                    ),
+                    messages=ollama_messages,
+                    format="json",
+                )
+
+                return response
+
+            except (ConnectError, ReadTimeout) as e:
+                last_error = e
+                logger.warning(
+                    f"Connection error on attempt {attempt + 1}/{self.max_retries}: {str(e)}"
+                )
+                if attempt < self.max_retries - 1:
+                    await asyncio.sleep(self.retry_delay * (attempt + 1))  # Exponential backoff
+                continue
+
+            except Exception as e:
+                logger.error(f"Unexpected error in Ollama API call: {str(e)}")
+                raise RuntimeError(f"Ollama API call failed: {str(e)}")
+
+        # If we get here, all retries failed
+        raise RuntimeError(f"Connection error after {self.max_retries} retries: {str(last_error)}")

{cua_agent-0.1.20 → cua_agent-0.1.22}/agent/providers/omni/clients/openai.py

@@ -4,11 +4,7 @@ import os
 import logging
 from typing import Dict, List, Optional, Any
 import aiohttp
-import base64
 import re
-import json
-import ssl
-import certifi
 from datetime import datetime
 from .base import BaseOmniClient
 

{cua_agent-0.1.20 → cua_agent-0.1.22}/agent/providers/omni/loop.py

@@ -19,6 +19,7 @@ from computer import Computer
 from .types import LLMProvider
 from .clients.openai import OpenAIClient
 from .clients.anthropic import AnthropicClient
+from .clients.ollama import OllamaClient
 from .prompts import SYSTEM_PROMPT
 from .api_handler import OmniAPIHandler
 from .tools.manager import ToolManager
@@ -135,6 +136,11 @@ class OmniLoop(BaseLoop):
                 api_key=self.api_key,
                 model=self.model,
             )
+        elif self.provider == LLMProvider.OLLAMA:
+            self.client = OllamaClient(
+                api_key=self.api_key,
+                model=self.model,
+            )
         else:
             raise ValueError(f"Unsupported provider: {self.provider}")
 
@@ -160,6 +166,11 @@ class OmniLoop(BaseLoop):
                     max_retries=self.max_retries,
                     retry_delay=self.retry_delay,
                 )
+            elif self.provider == LLMProvider.OLLAMA:
+                self.client = OllamaClient(
+                    api_key=self.api_key,
+                    model=self.model,
+                )
             else:
                 raise ValueError(f"Unsupported provider: {self.provider}")
 
@@ -370,6 +381,13 @@ class OmniLoop(BaseLoop):
                 else:
                     logger.warning("Invalid Anthropic response format")
                     return True, action_screenshot_saved
+            elif self.provider == LLMProvider.OLLAMA:
+                try:
+                    raw_text = response["message"]["content"]
+                    standard_content = [{"type": "text", "text": raw_text}]
+                except (KeyError, TypeError, IndexError) as e:
+                    logger.error(f"Invalid response format: {str(e)}")
+                    return True, action_screenshot_saved
             else:
                 # Assume OpenAI or compatible format
                 try:

{cua_agent-0.1.20 → cua_agent-0.1.22}/agent/providers/omni/types.py

@@ -11,6 +11,7 @@ class LLMProvider(StrEnum):
     ANTHROPIC = "anthropic"
     OMNI = "omni"
     OPENAI = "openai"
+    OLLAMA = "ollama"
 
 
 @dataclass
@@ -35,10 +36,12 @@ Model = LLM
 PROVIDER_TO_DEFAULT_MODEL: Dict[LLMProvider, str] = {
     LLMProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
     LLMProvider.OPENAI: "gpt-4o",
+    LLMProvider.OLLAMA: "gemma3:4b-it-q4_K_M",
 }
 
 # Environment variable names for each provider
 PROVIDER_TO_ENV_VAR: Dict[LLMProvider, str] = {
     LLMProvider.ANTHROPIC: "ANTHROPIC_API_KEY",
     LLMProvider.OPENAI: "OPENAI_API_KEY",
+    LLMProvider.OLLAMA: "none",
 }

{cua_agent-0.1.20 → cua_agent-0.1.22}/agent/providers/openai/api_handler.py

@@ -132,6 +132,9 @@ class OpenAIAPIHandler:
                 }
             ],
             "input": input_array,
+            "reasoning": {
+                "generate_summary": "concise",
+            },
             "truncation": "auto",
         }
 

{cua_agent-0.1.20 → cua_agent-0.1.22}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "cua-agent"
-version = "0.1.20"
+version = "0.1.22"
 description = "CUA (Computer Use) Agent for AI-driven computer interaction"
 readme = "README.md"
 authors = [
@@ -49,6 +49,20 @@ som = [
     "dashscope>=1.13.0,<2.0.0",
     "requests>=2.31.0,<3.0.0",
 ]
+omni = [
+    "torch>=2.2.1",
+    "torchvision>=0.17.1",
+    "ultralytics>=8.0.0",
+    "transformers>=4.38.2",
+    "cua-som>=0.1.0,<0.2.0",
+    "anthropic>=0.46.0,<0.47.0",
+    "boto3>=1.35.81,<2.0.0",
+    "openai>=1.14.0,<2.0.0",
+    "groq>=0.4.0,<0.5.0",
+    "dashscope>=1.13.0,<2.0.0",
+    "requests>=2.31.0,<3.0.0",
+    "ollama>=0.4.7,<0.5.0",
+]
 all = [
     "torch>=2.2.1",
     "torchvision>=0.17.1",
@@ -61,6 +75,7 @@ all = [
     "groq>=0.4.0,<0.5.0",
     "dashscope>=1.13.0,<2.0.0",
     "requests>=2.31.0,<3.0.0",
+    "ollama>=0.4.7,<0.5.0",
 ]
 
 [tool.pdm]
@@ -84,7 +99,7 @@ target-version = [
 
 [tool.ruff]
 line-length = 100
-target-version = "0.1.20"
+target-version = "0.1.22"
 select = [
     "E",
     "F",
@@ -98,7 +113,7 @@ docstring-code-format = true
 
 [tool.mypy]
 strict = true
-python_version = "0.1.20"
+python_version = "0.1.22"
 ignore_missing_imports = true
 disallow_untyped_defs = true
 check_untyped_defs = true