cua-agent 0.1.17__py3-none-any.whl → 0.1.19__py3-none-any.whl

agent/__init__.py

@@ -49,7 +49,7 @@ except Exception as e:
     logger.warning(f"Error initializing telemetry: {e}")
 
 from .providers.omni.types import LLMProvider, LLM
-from .core.loop import AgentLoop
-from .core.computer_agent import ComputerAgent
+from .core.factory import AgentLoop
+from .core.agent import ComputerAgent
 
 __all__ = ["AgentLoop", "LLMProvider", "LLM", "ComputerAgent"]

agent/core/__init__.py

@@ -1,6 +1,6 @@
 """Core agent components."""
 
-from .loop import BaseLoop
+from .factory import BaseLoop
 from .messages import (
     BaseMessageManager,
     ImageRetentionConfig,

agent/core/{computer_agent.py → agent.py}

@@ -3,32 +3,18 @@
 import asyncio
 import logging
 import os
-from typing import Any, AsyncGenerator, Dict, Optional, cast, List
+from typing import AsyncGenerator, Optional
 
 from computer import Computer
-from ..providers.anthropic.loop import AnthropicLoop
-from ..providers.omni.loop import OmniLoop
-from ..providers.omni.parser import OmniParser
-from ..providers.omni.types import LLMProvider, LLM
+from ..providers.omni.types import LLM
 from .. import AgentLoop
-from .messages import StandardMessageManager, ImageRetentionConfig
 from .types import AgentResponse
+from .factory import LoopFactory
+from .provider_config import DEFAULT_MODELS, ENV_VARS
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Default models for different providers
-DEFAULT_MODELS = {
-    LLMProvider.OPENAI: "gpt-4o",
-    LLMProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
-}
-
-# Map providers to their environment variable names
-ENV_VARS = {
-    LLMProvider.OPENAI: "OPENAI_API_KEY",
-    LLMProvider.ANTHROPIC: "ANTHROPIC_API_KEY",
-}
-
 
 class ComputerAgent:
     """A computer agent that can perform automated tasks using natural language instructions."""
@@ -98,35 +84,27 @@ class ComputerAgent:
                     f"No model specified for provider {self.provider} and no default found"
                 )
 
-        # Ensure computer is properly cast for typing purposes
-        computer_instance = self.computer
-
         # 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:
             raise ValueError(f"No API key provided for {self.provider}")
 
-        # Initialize the appropriate loop based on the loop parameter
-        if loop == AgentLoop.ANTHROPIC:
-            self._loop = AnthropicLoop(
-                api_key=actual_api_key,
-                model=actual_model_name,
-                computer=computer_instance,
-                save_trajectory=save_trajectory,
-                base_dir=trajectory_dir,
-                only_n_most_recent_images=only_n_most_recent_images,
-            )
-        else:
-            self._loop = OmniLoop(
+        # Create the appropriate loop using the factory
+        try:
+            # Let the factory create the appropriate loop with needed components
+            self._loop = LoopFactory.create_loop(
+                loop_type=loop,
                 provider=self.provider,
+                computer=self.computer,
+                model_name=actual_model_name,
                 api_key=actual_api_key,
-                model=actual_model_name,
-                computer=computer_instance,
                 save_trajectory=save_trajectory,
-                base_dir=trajectory_dir,
+                trajectory_dir=trajectory_dir,
                 only_n_most_recent_images=only_n_most_recent_images,
-                parser=OmniParser(),
             )
+        except ValueError as e:
+            logger.error(f"Failed to create loop: {str(e)}")
+            raise
 
         # Initialize the message manager from the loop
         self.message_manager = self._loop.message_manager
@@ -152,21 +130,6 @@ class ComputerAgent:
             else:
                 logger.info("Computer already initialized, skipping initialization")
 
-            # Take a test screenshot to verify the computer is working
-            logger.info("Testing computer with a screenshot...")
-            try:
-                test_screenshot = await self.computer.interface.screenshot()
-                # Determine the screenshot size based on its type
-                if isinstance(test_screenshot, (bytes, bytearray, memoryview)):
-                    size = len(test_screenshot)
-                elif hasattr(test_screenshot, "base64_image"):
-                    size = len(test_screenshot.base64_image)
-                else:
-                    size = "unknown"
-                logger.info(f"Screenshot test successful, size: {size}")
-            except Exception as e:
-                logger.error(f"Screenshot test failed: {str(e)}")
-                # Even though screenshot failed, we continue since some tests might not need it
         except Exception as e:
             logger.error(f"Error initializing computer in __aenter__: {str(e)}")
             raise
@@ -232,7 +195,6 @@ class ComputerAgent:
 
             # Execute the task and yield results
             async for result in self._loop.run(self.message_manager.messages):
-                # Yield the result to the caller
                 yield result
 
         except Exception as e:

agent/core/{loop.py → base.py}

@@ -1,35 +1,21 @@
-"""Base agent loop implementation."""
+"""Base loop definitions."""
 
 import logging
 import asyncio
 from abc import ABC, abstractmethod
-from enum import Enum, auto
-from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple
-from datetime import datetime
+from typing import Any, AsyncGenerator, Dict, List, Optional
 
 from computer import Computer
-from .experiment import ExperimentManager
 from .messages import StandardMessageManager, ImageRetentionConfig
 from .types import AgentResponse
+from .experiment import ExperimentManager
 
 logger = logging.getLogger(__name__)
 
 
-class AgentLoop(Enum):
-    """Enumeration of available loop types."""
-
-    ANTHROPIC = auto()  # Anthropic implementation
-    OMNI = auto()  # OmniLoop implementation
-    # Add more loop types as needed
-
-
 class BaseLoop(ABC):
     """Base class for agent loops that handle message processing and tool execution."""
 
-    ###########################################
-    # INITIALIZATION AND CONFIGURATION
-    ###########################################
-
     def __init__(
         self,
         computer: Computer,
@@ -68,6 +54,11 @@ class BaseLoop(ABC):
         self.only_n_most_recent_images = only_n_most_recent_images
         self._kwargs = kwargs
 
+        # Initialize message manager
+        self.message_manager = StandardMessageManager(
+            config=ImageRetentionConfig(num_images_to_keep=only_n_most_recent_images)
+        )
+
         # Initialize experiment manager
         if self.save_trajectory and self.base_dir:
             self.experiment_manager = ExperimentManager(
@@ -110,8 +101,7 @@ class BaseLoop(ABC):
                     )
                     raise RuntimeError(f"Failed to initialize: {str(e)}")
 
-        ###########################################
-
+    ###########################################
     # ABSTRACT METHODS TO BE IMPLEMENTED BY SUBCLASSES
     ###########################################
 
@@ -125,17 +115,14 @@ class BaseLoop(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    async def run(self, messages: List[Dict[str, Any]]) -> AsyncGenerator[AgentResponse, None]:
+    def run(self, messages: List[Dict[str, Any]]) -> AsyncGenerator[AgentResponse, None]:
         """Run the agent loop with provided messages.
 
-        This method handles the main agent loop including message processing,
-        API calls, response handling, and action execution.
-
         Args:
             messages: List of message objects
 
-        Yields:
-            Agent response format
+        Returns:
+            An async generator that yields agent responses
         """
         raise NotImplementedError
 

agent/core/factory.py

@@ -0,0 +1,104 @@
+"""Base agent loop implementation."""
+
+import logging
+import importlib.util
+from typing import Dict, Optional, Type, TYPE_CHECKING, Any, cast, Callable, Awaitable
+
+from computer import Computer
+from .types import AgentLoop
+from .base import BaseLoop
+
+# For type checking only
+if TYPE_CHECKING:
+    from ..providers.omni.types import LLMProvider
+
+logger = logging.getLogger(__name__)
+
+
+class LoopFactory:
+    """Factory class for creating agent loops."""
+
+    # Registry to store loop implementations
+    _loop_registry: Dict[AgentLoop, Type[BaseLoop]] = {}
+
+    @classmethod
+    def create_loop(
+        cls,
+        loop_type: AgentLoop,
+        api_key: str,
+        model_name: str,
+        computer: Computer,
+        provider: Any = None,
+        save_trajectory: bool = True,
+        trajectory_dir: str = "trajectories",
+        only_n_most_recent_images: Optional[int] = None,
+        acknowledge_safety_check_callback: Optional[Callable[[str], Awaitable[bool]]] = None,
+    ) -> BaseLoop:
+        """Create and return an appropriate loop instance based on type."""
+        if loop_type == AgentLoop.ANTHROPIC:
+            # Lazy import AnthropicLoop only when needed
+            try:
+                from ..providers.anthropic.loop import AnthropicLoop
+            except ImportError:
+                raise ImportError(
+                    "The 'anthropic' provider is not installed. "
+                    "Install it with 'pip install cua-agent[anthropic]'"
+                )
+
+            return AnthropicLoop(
+                api_key=api_key,
+                model=model_name,
+                computer=computer,
+                save_trajectory=save_trajectory,
+                base_dir=trajectory_dir,
+                only_n_most_recent_images=only_n_most_recent_images,
+            )
+        elif loop_type == AgentLoop.OPENAI:
+            # Lazy import OpenAILoop only when needed
+            try:
+                from ..providers.openai.loop import OpenAILoop
+            except ImportError:
+                raise ImportError(
+                    "The 'openai' provider is not installed. "
+                    "Install it with 'pip install cua-agent[openai]'"
+                )
+
+            return OpenAILoop(
+                api_key=api_key,
+                model=model_name,
+                computer=computer,
+                save_trajectory=save_trajectory,
+                base_dir=trajectory_dir,
+                only_n_most_recent_images=only_n_most_recent_images,
+                acknowledge_safety_check_callback=acknowledge_safety_check_callback,
+            )
+        elif loop_type == AgentLoop.OMNI:
+            # Lazy import OmniLoop and related classes only when needed
+            try:
+                from ..providers.omni.loop import OmniLoop
+                from ..providers.omni.parser import OmniParser
+                from ..providers.omni.types import LLMProvider
+            except ImportError:
+                raise ImportError(
+                    "The 'omni' provider is not installed. "
+                    "Install it with 'pip install cua-agent[all]'"
+                )
+
+            if provider is None:
+                raise ValueError("Provider is required for OMNI loop type")
+
+            # We know provider is the correct type at this point, so cast it
+            provider_instance = cast(LLMProvider, provider)
+
+            return OmniLoop(
+                provider=provider_instance,
+                api_key=api_key,
+                model=model_name,
+                computer=computer,
+                save_trajectory=save_trajectory,
+                base_dir=trajectory_dir,
+                only_n_most_recent_images=only_n_most_recent_images,
+                parser=OmniParser(),
+            )
+        else:
+            raise ValueError(f"Unsupported loop type: {loop_type}")

agent/core/provider_config.py

@@ -0,0 +1,15 @@
+"""Provider-specific configurations and constants."""
+
+from ..providers.omni.types import LLMProvider
+
+# Default models for different providers
+DEFAULT_MODELS = {
+    LLMProvider.OPENAI: "gpt-4o",
+    LLMProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
+}
+
+# Map providers to their environment variable names
+ENV_VARS = {
+    LLMProvider.OPENAI: "OPENAI_API_KEY",
+    LLMProvider.ANTHROPIC: "ANTHROPIC_API_KEY",
+}

agent/core/types.py

@@ -1,6 +1,16 @@
 """Core type definitions."""
 
 from typing import Any, Dict, List, Optional, TypedDict, Union
+from enum import Enum, auto
+
+
+class AgentLoop(Enum):
+    """Enumeration of available loop types."""
+
+    ANTHROPIC = auto()  # Anthropic implementation
+    OMNI = auto()  # OmniLoop implementation
+    OPENAI = auto()  # OpenAI implementation
+    # Add more loop types as needed
 
 
 class AgentResponse(TypedDict, total=False):

agent/providers/anthropic/loop.py

@@ -16,7 +16,7 @@ from datetime import datetime
 from computer import Computer
 
 # Base imports
-from ...core.loop import BaseLoop
+from ...core.base import BaseLoop
 from ...core.messages import StandardMessageManager, ImageRetentionConfig
 from ...core.types import AgentResponse
 

agent/providers/anthropic/response_handler.py

@@ -1,14 +1,11 @@
 """Response and tool handling for Anthropic provider."""
 
 import logging
-from typing import Any, Dict, List, Optional, Tuple, cast
+from typing import Any, Dict, List, Tuple, cast
 
 from anthropic.types.beta import (
     BetaMessage,
-    BetaMessageParam,
     BetaTextBlock,
-    BetaTextBlockParam,
-    BetaToolUseBlockParam,
     BetaContentBlockParam,
 )
 

agent/providers/anthropic/utils.py

@@ -1,14 +1,12 @@
 """Utility functions for Anthropic message handling."""
 
-import time
 import logging
 import re
 from typing import Any, Dict, List, Optional, Tuple, cast
-from anthropic.types.beta import BetaMessage, BetaMessageParam, BetaTextBlock
+from anthropic.types.beta import BetaMessage
 from ..omni.parser import ParseResult
 from ...core.types import AgentResponse
 from datetime import datetime
-import json
 
 # Configure module logger
 logger = logging.getLogger(__name__)

agent/providers/omni/loop.py

@@ -10,7 +10,7 @@ from httpx import ConnectError, ReadTimeout
 from typing import cast
 
 from .parser import OmniParser, ParseResult
-from ...core.loop import BaseLoop
+from ...core.base import BaseLoop
 from ...core.visualization import VisualizationHelper
 from ...core.messages import StandardMessageManager, ImageRetentionConfig
 from .utils import to_openai_agent_response_format

agent/providers/omni/types.py

@@ -9,8 +9,10 @@ class LLMProvider(StrEnum):
     """Supported LLM providers."""
 
     ANTHROPIC = "anthropic"
+    OMNI = "omni"
     OPENAI = "openai"
 
+
 @dataclass
 class LLM:
     """Configuration for LLM model and provider."""

agent/providers/openai/__init__.py

@@ -0,0 +1,6 @@
+"""OpenAI Agent Response API provider for computer control."""
+
+from .types import LLMProvider
+from .loop import OpenAILoop
+
+__all__ = ["OpenAILoop", "LLMProvider"]