cua-agent 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

agent/__init__.py

@@ -5,6 +5,6 @@ __version__ = "0.1.0"
 from .core.factory import AgentFactory
 from .core.agent import ComputerAgent
 from .types.base import Provider, AgenticLoop
-from .providers.omni.types import APIProvider
+from .providers.omni.types import LLMProvider, LLM, Model, LLMModel, APIProvider
 
-__all__ = ["AgentFactory", "Provider", "ComputerAgent", "AgenticLoop", "APIProvider"]
+__all__ = ["AgentFactory", "Provider", "ComputerAgent", "AgenticLoop", "LLMProvider", "LLM", "Model", "LLMModel", "APIProvider"]

agent/core/agent.py

@@ -3,7 +3,7 @@
 import os
 import logging
 import asyncio
-from typing import Any, AsyncGenerator, Dict, List, Optional, TYPE_CHECKING
+from typing import Any, AsyncGenerator, Dict, List, Optional, TYPE_CHECKING, Union, cast
 from datetime import datetime
 
 from computer import Computer
@@ -17,23 +17,23 @@ if TYPE_CHECKING:
     from ..providers.omni.loop import OmniLoop
     from ..providers.omni.parser import OmniParser
 
-# Import the APIProvider enum without importing the whole module
-from ..providers.omni.types import APIProvider
+# Import the provider types
+from ..providers.omni.types import LLMProvider, LLM, Model, LLMModel, APIProvider
 
 logger = logging.getLogger(__name__)
 
 # Default models for different providers
 DEFAULT_MODELS = {
-    APIProvider.OPENAI: "gpt-4o",
-    APIProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
-    APIProvider.GROQ: "llama3-70b-8192",
+    LLMProvider.OPENAI: "gpt-4o",
+    LLMProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
+    LLMProvider.GROQ: "llama3-70b-8192",
 }
 
 # Map providers to their environment variable names
 ENV_VARS = {
-    APIProvider.OPENAI: "OPENAI_API_KEY",
-    APIProvider.GROQ: "GROQ_API_KEY",
-    APIProvider.ANTHROPIC: "ANTHROPIC_API_KEY",
+    LLMProvider.OPENAI: "OPENAI_API_KEY",
+    LLMProvider.GROQ: "GROQ_API_KEY",
+    LLMProvider.ANTHROPIC: "ANTHROPIC_API_KEY",
 }
 
 
@@ -48,9 +48,8 @@ class ComputerAgent(BaseComputerAgent):
         self,
         computer: Computer,
         loop_type: AgenticLoop = AgenticLoop.OMNI,
-        ai_provider: APIProvider = APIProvider.OPENAI,
+        model: Optional[Union[LLM, Dict[str, str], str]] = None,
         api_key: Optional[str] = None,
-        model: Optional[str] = None,
         save_trajectory: bool = True,
         trajectory_dir: Optional[str] = "trajectories",
         only_n_most_recent_images: Optional[int] = None,
@@ -63,9 +62,12 @@ class ComputerAgent(BaseComputerAgent):
         Args:
             computer: Computer instance to control
             loop_type: The type of loop to use (Anthropic or Omni)
-            ai_provider: AI provider to use (required for Cua loop)
+            model: LLM configuration. Can be:
+                  - LLM object with provider and name
+                  - Dict with 'provider' and 'name' keys
+                  - String with model name (defaults to OpenAI provider)
+                  - None (defaults based on loop_type)
             api_key: Optional API key (will use environment variable if not provided)
-            model: Optional model name (will use provider default if not specified)
             save_trajectory: Whether to save screenshots and logs
             trajectory_dir: Directory to save trajectories (defaults to "trajectories")
             only_n_most_recent_images: Limit history to N most recent images
@@ -88,7 +90,6 @@ class ComputerAgent(BaseComputerAgent):
         )
 
         self.loop_type = loop_type
-        self.provider = ai_provider
         self.save_trajectory = save_trajectory
         self.trajectory_dir = trajectory_dir
         self.only_n_most_recent_images = only_n_most_recent_images
@@ -98,14 +99,19 @@ class ComputerAgent(BaseComputerAgent):
         # Configure logging based on verbosity
         self._configure_logging(verbosity)
 
+        # Process model configuration
+        self.model_config = self._process_model_config(model, loop_type)
+        
         # Get API key from environment if not provided
         if api_key is None:
             env_var = (
-                ENV_VARS.get(ai_provider) if loop_type == AgenticLoop.OMNI else "ANTHROPIC_API_KEY"
+                ENV_VARS.get(self.model_config.provider) 
+                if loop_type == AgenticLoop.OMNI 
+                else "ANTHROPIC_API_KEY"
             )
             if not env_var:
                 raise ValueError(
-                    f"Unsupported provider: {ai_provider}. Please use one of: {list(ENV_VARS.keys())}"
+                    f"Unsupported provider: {self.model_config.provider}. Please use one of: {list(ENV_VARS.keys())}"
                 )
 
             api_key = os.environ.get(env_var)
@@ -119,17 +125,51 @@ class ComputerAgent(BaseComputerAgent):
                 )
         self.api_key = api_key
 
-        # Set model based on provider if not specified
-        if model is None:
-            if loop_type == AgenticLoop.OMNI:
-                self.model = DEFAULT_MODELS[ai_provider]
-            else:  # Anthropic loop
-                self.model = DEFAULT_MODELS[APIProvider.ANTHROPIC]
-        else:
-            self.model = model
-
         # Initialize the appropriate loop based on loop_type
         self.loop = self._init_loop()
+        
+    def _process_model_config(
+        self, model_input: Optional[Union[LLM, Dict[str, str], str]], loop_type: AgenticLoop
+    ) -> LLM:
+        """Process and normalize model configuration.
+        
+        Args:
+            model_input: Input model configuration (LLM, dict, string, or None)
+            loop_type: The loop type being used
+            
+        Returns:
+            Normalized LLM instance
+        """
+        # Handle case where model_input is None
+        if model_input is None:
+            # Use Anthropic for Anthropic loop, OpenAI for Omni loop
+            default_provider = (
+                LLMProvider.ANTHROPIC if loop_type == AgenticLoop.ANTHROPIC else LLMProvider.OPENAI
+            )
+            return LLM(provider=default_provider)
+            
+        # Handle case where model_input is already a LLM or one of its aliases
+        if isinstance(model_input, (LLM, Model, LLMModel)):
+            return model_input
+            
+        # Handle case where model_input is a dict
+        if isinstance(model_input, dict):
+            provider = model_input.get("provider", LLMProvider.OPENAI)
+            if isinstance(provider, str):
+                provider = LLMProvider(provider)
+            return LLM(
+                provider=provider,
+                name=model_input.get("name")
+            )
+            
+        # Handle case where model_input is a string (model name)
+        if isinstance(model_input, str):
+            default_provider = (
+                LLMProvider.ANTHROPIC if loop_type == AgenticLoop.ANTHROPIC else LLMProvider.OPENAI
+            )
+            return LLM(provider=default_provider, name=model_input)
+            
+        raise ValueError(f"Unsupported model configuration: {model_input}")
 
     def _configure_logging(self, verbosity: int):
         """Configure logging based on verbosity level."""
@@ -162,9 +202,12 @@ class ComputerAgent(BaseComputerAgent):
         if self.loop_type == AgenticLoop.ANTHROPIC:
             from ..providers.anthropic.loop import AnthropicLoop
 
+            # Ensure we always have a valid model name
+            model_name = self.model_config.name or DEFAULT_MODELS[LLMProvider.ANTHROPIC]
+            
             return AnthropicLoop(
                 api_key=self.api_key,
-                model=self.model,
+                model=model_name,
                 computer=self.computer,
                 save_trajectory=self.save_trajectory,
                 base_dir=self.trajectory_dir,
@@ -176,10 +219,13 @@ class ComputerAgent(BaseComputerAgent):
         if "parser" not in self._kwargs:
             self._kwargs["parser"] = OmniParser()
 
+        # Ensure we always have a valid model name
+        model_name = self.model_config.name or DEFAULT_MODELS[self.model_config.provider]
+        
         return OmniLoop(
-            provider=self.provider,
+            provider=self.model_config.provider,
             api_key=self.api_key,
-            model=self.model,
+            model=model_name,
             computer=self.computer,
             save_trajectory=self.save_trajectory,
             base_dir=self.trajectory_dir,

agent/core/messages.py

@@ -37,6 +37,17 @@ class BaseMessageManager:
         if self.image_retention_config.min_removal_threshold < 1:
             raise ValueError("min_removal_threshold must be at least 1")
 
+        # Track provider for message formatting
+        self.provider = "openai"  # Default provider
+
+    def set_provider(self, provider: str) -> None:
+        """Set the current provider to format messages for.
+
+        Args:
+            provider: Provider name (e.g., 'openai', 'anthropic')
+        """
+        self.provider = provider.lower()
+
     def prepare_messages(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
         """Prepare messages by applying image retention and caching as configured.
 
@@ -96,6 +107,10 @@ class BaseMessageManager:
         Args:
             messages: Messages to inject caching into
         """
+        # Only apply cache_control for Anthropic API, not OpenAI
+        if self.provider != "anthropic":
+            return
+
         # Default to caching last 3 turns
         turns_to_cache = 3
         for message in reversed(messages):

agent/providers/omni/loop.py

@@ -219,9 +219,13 @@ class OmniLoop(BaseLoop):
                     if self.client is None:
                         raise RuntimeError("Failed to initialize client")
 
+                # Set the provider in message manager based on current provider
+                provider_name = str(self.provider).split(".")[-1].lower()  # Extract name from enum
+                self.message_manager.set_provider(provider_name)
+
                 # Apply image retention and prepare messages
                 # This will limit the number of images based on only_n_most_recent_images
-                prepared_messages = self.message_manager.prepare_messages(messages.copy())
+                prepared_messages = self.message_manager.get_formatted_messages(provider_name)
 
                 # Filter out system messages for Anthropic
                 if self.provider == APIProvider.ANTHROPIC:

agent/providers/omni/messages.py

@@ -103,6 +103,9 @@ class OmniMessageManager(BaseMessageManager):
         Returns:
             List of formatted messages
         """
+        # Set the provider for message formatting
+        self.set_provider(provider)
+
         if provider == "anthropic":
             return self._format_for_anthropic()
         elif provider == "openai":

agent/providers/omni/types.py

@@ -1,11 +1,12 @@
 """Type definitions for the Omni provider."""
 
 from enum import StrEnum
-from typing import Dict
+from typing import Dict, Optional
+from dataclasses import dataclass
 
 
-class APIProvider(StrEnum):
-    """Supported API providers."""
+class LLMProvider(StrEnum):
+    """Supported LLM providers."""
 
     ANTHROPIC = "anthropic"
     OPENAI = "openai"
@@ -13,18 +14,40 @@ class APIProvider(StrEnum):
     QWEN = "qwen"
 
 
+# For backward compatibility
+APIProvider = LLMProvider
+
+
+@dataclass
+class LLM:
+    """Configuration for LLM model and provider."""
+    
+    provider: LLMProvider
+    name: Optional[str] = None
+    
+    def __post_init__(self):
+        """Set default model name if not provided."""
+        if self.name is None:
+            self.name = PROVIDER_TO_DEFAULT_MODEL.get(self.provider)
+
+
+# For backward compatibility
+LLMModel = LLM
+Model = LLM
+
+
 # Default models for each provider
-PROVIDER_TO_DEFAULT_MODEL: Dict[APIProvider, str] = {
-    APIProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
-    APIProvider.OPENAI: "gpt-4o",
-    APIProvider.GROQ: "deepseek-r1-distill-llama-70b",
-    APIProvider.QWEN: "qwen2.5-vl-72b-instruct",
+PROVIDER_TO_DEFAULT_MODEL: Dict[LLMProvider, str] = {
+    LLMProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
+    LLMProvider.OPENAI: "gpt-4o",
+    LLMProvider.GROQ: "deepseek-r1-distill-llama-70b",
+    LLMProvider.QWEN: "qwen2.5-vl-72b-instruct",
 }
 
 # Environment variable names for each provider
-PROVIDER_TO_ENV_VAR: Dict[APIProvider, str] = {
-    APIProvider.ANTHROPIC: "ANTHROPIC_API_KEY",
-    APIProvider.OPENAI: "OPENAI_API_KEY",
-    APIProvider.GROQ: "GROQ_API_KEY",
-    APIProvider.QWEN: "QWEN_API_KEY",
+PROVIDER_TO_ENV_VAR: Dict[LLMProvider, str] = {
+    LLMProvider.ANTHROPIC: "ANTHROPIC_API_KEY",
+    LLMProvider.OPENAI: "OPENAI_API_KEY",
+    LLMProvider.GROQ: "GROQ_API_KEY",
+    LLMProvider.QWEN: "QWEN_API_KEY",
 }

{cua_agent-0.1.0.dist-info → cua_agent-0.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cua-agent
-Version: 0.1.0
+Version: 0.1.1
 Summary: CUA (Computer Use) Agent for AI-driven computer interaction
 Author-Email: TryCua <gh@trycua.com>
 Requires-Python: <3.13,>=3.10

{cua_agent-0.1.0.dist-info → cua_agent-0.1.1.dist-info}/RECORD

@@ -1,15 +1,15 @@
 agent/README.md,sha256=8EFnLrKejthEcL9bZflQSbvA-KwpiPanBz8TEEwRub8,2153
-agent/__init__.py,sha256=16Q828puFb7Ucq_-de49moVCzl1-iDO8Uo5dzFwX0Ag,347
+agent/__init__.py,sha256=5IxjivBoXkpBQyPP3uwCrCoMx7gNZbM1rdVaIn2jxZ4,425
 agent/core/README.md,sha256=RY4kKEjm_-_Ul2xgY7ntzsXdPe0Tg1wvtOSZ4xp4DN0,3559
 agent/core/__init__.py,sha256=0htZ-VfsH9ixHB8j_SXu_uv6r3XXsq5TrghFNd-yRNE,709
-agent/core/agent.py,sha256=q2x0vFykIavX_FBi4Eq222QCSFmuuekAin4FPrtSGbY,11711
+agent/core/agent.py,sha256=HSVNTiEhlDHMGJiA4CKi33TtwEdsMIvkgZ9nHfk2M8E,13730
 agent/core/base_agent.py,sha256=MgaMKTwgqNJ1-TgS_mxALoC9COzc7Acg9y7Q8HAFX2c,6266
 agent/core/callbacks.py,sha256=VbGIf5QkHh3Q0KsLM6wv7hRdIA5WExTVYLm64bckyUA,4306
 agent/core/computer_agent.py,sha256=JGLMl_PwImUttmQh2amdLlXHS9CUyZ9MW20J1Xid7dM,2417
 agent/core/experiment.py,sha256=AST1t83eqaGzjoW6KvrhfVIs3ELAR_I70VHq2NsMmNk,7446
 agent/core/factory.py,sha256=WraOEHWPXBSN4R3DO7M2ctyadodeA8tzHM3dUjdQ_3A,3441
 agent/core/loop.py,sha256=E-0pz7MaguZQrHs5GP98Oc8C_Iz8ier0vXrD9Ny2HL8,8999
-agent/core/messages.py,sha256=Ou0lLEwa2EQCartcTszsvNjCP6sHUxmr2_C9PGzbASg,7163
+agent/core/messages.py,sha256=N8pV8Eh-AJpMuDPRI5OGWUIOU6DRr-pQjK9XU0go9Hk,7637
 agent/core/tools/__init__.py,sha256=xZen-PqUp2dUaMEHJowXCQm33_5Sxhsx9PSoD0rq6tI,489
 agent/core/tools/base.py,sha256=CdzRFNuOjNfzgyTUN4ZoCGkUDR5HI0ECQVpvrUdEij8,2295
 agent/core/tools/bash.py,sha256=jnJKVlHn8np8e0gWd8EO0_qqjMkfQzutSugA_Iol4jE,1585
@@ -43,8 +43,8 @@ agent/providers/omni/clients/openai.py,sha256=E4TAXMUFoYTunJETCWCNx5XAc6xutiN4rB
 agent/providers/omni/clients/utils.py,sha256=Ani9CVVBm_J2Dl51WG6p1GVuoI6cq8scISrG0pmQ37o,688
 agent/providers/omni/experiment.py,sha256=JGAdHi7Nf73I48c9k3TY1Xpr_i6D2VG1wurOzw5cNGk,9888
 agent/providers/omni/image_utils.py,sha256=qIFuNi5cIMVwrqYBXG1T6PxUlbxz7gIngFFP39bZIlU,2782
-agent/providers/omni/loop.py,sha256=Xr2QeedAVJ_jHn3KMopRuH3mrm2Qn4ncxKjqj9hWxAw,43577
-agent/providers/omni/messages.py,sha256=6LkQfzYDWq2FvIHpqhs5pc0l6AmFx_xKCjj1R5czMPo,6047
+agent/providers/omni/loop.py,sha256=U1R_ayfN4T25hvbLMp97qeqSrqVtSL-U03G8Sqf4AaM,43827
+agent/providers/omni/messages.py,sha256=zdjQCAMH-hOyrQQesHhTiIsQbw43KqVSmVIzS8JOIFA,6134
 agent/providers/omni/parser.py,sha256=Iv-cXWG2qzdYjyZJH5pGUzfv6nOaiHQ2OXdQSe00Ydw,9151
 agent/providers/omni/prompts.py,sha256=29qy8ppbLOjLil3aiqryjaiBf8CQx-xXHN44O-85Q00,4503
 agent/providers/omni/tool_manager.py,sha256=O6DxyEI-Vg6jt99phh011o4q4me_vNhH2YffIxkO4GM,2585
@@ -52,14 +52,14 @@ agent/providers/omni/tools/__init__.py,sha256=l636hx9Q5z9eaFdPanPwPENUE-w-Xm8kAZ
 agent/providers/omni/tools/bash.py,sha256=y_ibfP9iRcbiU_E0faAoa4DCP_BlkMlKOOURdBBIGZE,2030
 agent/providers/omni/tools/computer.py,sha256=xkMmAR0e_kbf0Zs2mggCDyWrQOJZyXOKPFjkutaQb94,9108
 agent/providers/omni/tools/manager.py,sha256=V_tav2yU92PyQnFlxNXG1wvNEaJoEYudtKx5sRjj06Q,2619
-agent/providers/omni/types.py,sha256=cEH6M5fcRN8ZIv_jfcYkTYboGBM4EzglLZo1_Xk7Ip8,800
+agent/providers/omni/types.py,sha256=6x-n3MLvvKOFAdvzYDf6Zzw-i118kvWHXE37qxa_L4o,1284
 agent/providers/omni/utils.py,sha256=JqSye1bEp4wxhUgmaMyZi172fTlgXtygJ7XlnvKdUtE,6337
 agent/providers/omni/visualization.py,sha256=N3qVQLxYmia3iSVC5oCt5YRlMPuVfylCOyB99R33u8U,3924
 agent/types/__init__.py,sha256=61UFJT-w0CT4YRn0LiTx4A7fsMdVQjlXO9vnmbI1A7Y,604
 agent/types/base.py,sha256=rVb4mPWp1SOHfrzOCDqx0pfCV5bgIsdrIzgM_kX_xVs,1090
 agent/types/messages.py,sha256=4-hwtxeAhto90_EZpHFducddtsHUsHauvXzYrpKG4RE,953
 agent/types/tools.py,sha256=Jes2CFCFqC727WWHbO-sG7V03rBHnQe5X7Oi9ZkuScI,877
-cua_agent-0.1.0.dist-info/METADATA,sha256=Q4nPzYL_UQwx82vuaRLBUFmA_Sgd37TVoGA9FNYDRmU,1890
-cua_agent-0.1.0.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
-cua_agent-0.1.0.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-cua_agent-0.1.0.dist-info/RECORD,,
+cua_agent-0.1.1.dist-info/METADATA,sha256=JdbHHQ7uBAlcnLZpZ1eWCiPCDODEOaB50j6XwtIs0Ss,1890
+cua_agent-0.1.1.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
+cua_agent-0.1.1.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+cua_agent-0.1.1.dist-info/RECORD,,