cua-agent 0.1.5__tar.gz → 0.1.6__tar.gz

{cua_agent-0.1.5 → cua_agent-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cua-agent
-Version: 0.1.5
+Version: 0.1.6
 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.5 → cua_agent-0.1.6}/agent/__init__.py

@@ -48,9 +48,7 @@ except Exception as e:
     # Other issues with telemetry
     logger.warning(f"Error initializing telemetry: {e}")
 
-from .core.factory import AgentFactory
-from .core.agent import ComputerAgent
 from .providers.omni.types import LLMProvider, LLM
-from .types.base import Provider, AgentLoop
+from .types.base import AgentLoop
 
-__all__ = ["AgentFactory", "Provider", "ComputerAgent", "AgentLoop", "LLMProvider", "LLM"]
+__all__ = ["AgentLoop", "LLMProvider", "LLM"]

{cua_agent-0.1.5 → cua_agent-0.1.6}/agent/core/__init__.py

@@ -1,6 +1,5 @@
 """Core agent components."""
 
-from .base_agent import BaseComputerAgent
 from .loop import BaseLoop
 from .messages import (
     create_user_message,
@@ -12,7 +11,7 @@ from .messages import (
     ImageRetentionConfig,
 )
 from .callbacks import (
-    CallbackManager, 
+    CallbackManager,
     CallbackHandler,
     BaseCallbackManager,
     ContentCallback,
@@ -21,9 +20,8 @@ from .callbacks import (
 )
 
 __all__ = [
-    "BaseComputerAgent", 
-    "BaseLoop", 
-    "CallbackManager", 
+    "BaseLoop",
+    "CallbackManager",
     "CallbackHandler",
     "BaseMessageManager",
     "ImageRetentionConfig",

cua_agent-0.1.6/agent/core/computer_agent.py

@@ -0,0 +1,251 @@
+"""Main entry point for computer agents."""
+
+import asyncio
+import logging
+import os
+from typing import Any, AsyncGenerator, Dict, Optional, cast
+from dataclasses import dataclass
+
+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 .. import AgentLoop
+
+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."""
+
+    def __init__(
+        self,
+        computer: Computer,
+        model: LLM,
+        loop: AgentLoop,
+        max_retries: int = 3,
+        screenshot_dir: Optional[str] = None,
+        log_dir: Optional[str] = None,
+        api_key: Optional[str] = None,
+        save_trajectory: bool = True,
+        trajectory_dir: str = "trajectories",
+        only_n_most_recent_images: Optional[int] = None,
+        parser: Optional[OmniParser] = None,
+        verbosity: int = logging.INFO,
+    ):
+        """Initialize the ComputerAgent.
+
+        Args:
+            computer: Computer instance. If not provided, one will be created with default settings.
+            max_retries: Maximum number of retry attempts.
+            screenshot_dir: Directory to save screenshots.
+            log_dir: Directory to save logs (set to None to disable logging to files).
+            model: LLM object containing provider and model name. Takes precedence over provider/model_name.
+            provider: The AI provider to use (e.g., LLMProvider.ANTHROPIC). Only used if model is None.
+            api_key: The API key for the provider. If not provided, will look for environment variable.
+            model_name: The model name to use. Only used if model is None.
+            save_trajectory: Whether to save the trajectory.
+            trajectory_dir: Directory to save the trajectory.
+            only_n_most_recent_images: Maximum number of recent screenshots to include in API requests.
+            parser: Parser instance for the OmniLoop. Only used if provider is not ANTHROPIC.
+            verbosity: Logging level.
+        """
+        # Basic agent configuration
+        self.max_retries = max_retries
+        self.computer = computer or Computer()
+        self.queue = asyncio.Queue()
+        self.screenshot_dir = screenshot_dir
+        self.log_dir = log_dir
+        self._retry_count = 0
+        self._initialized = False
+        self._in_context = False
+
+        # Set logging level
+        logger.setLevel(verbosity)
+
+        # Setup logging
+        if self.log_dir:
+            os.makedirs(self.log_dir, exist_ok=True)
+            logger.info(f"Created logs directory: {self.log_dir}")
+
+        # Setup screenshots directory
+        if self.screenshot_dir:
+            os.makedirs(self.screenshot_dir, exist_ok=True)
+            logger.info(f"Created screenshots directory: {self.screenshot_dir}")
+
+        # Use the provided LLM object
+        self.provider = model.provider
+        actual_model_name = model.name or DEFAULT_MODELS.get(self.provider, "")
+
+        # Ensure we have a valid model name
+        if not actual_model_name:
+            actual_model_name = DEFAULT_MODELS.get(self.provider, "")
+            if not actual_model_name:
+                raise ValueError(
+                    f"No model specified for provider {self.provider} and no default found"
+                )
+
+        # Ensure computer is properly cast for typing purposes
+        computer_instance = cast(Computer, 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:
+            # Default to OmniLoop for other loop types
+            # Initialize parser if not provided
+            actual_parser = parser or OmniParser()
+
+            self._loop = OmniLoop(
+                provider=self.provider,
+                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,
+                parser=actual_parser,
+            )
+
+        logger.info(
+            f"ComputerAgent initialized with provider: {self.provider}, model: {actual_model_name}"
+        )
+
+    async def __aenter__(self):
+        """Initialize the agent when used as a context manager."""
+        logger.info("Entering ComputerAgent context")
+        self._in_context = True
+
+        # In case the computer wasn't initialized
+        try:
+            # Initialize the computer only if not already initialized
+            logger.info("Checking if computer is already initialized...")
+            if not self.computer._initialized:
+                logger.info("Initializing computer in __aenter__...")
+                # Use the computer's __aenter__ directly instead of calling run()
+                await self.computer.__aenter__()
+                logger.info("Computer initialized in __aenter__")
+            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
+
+        await self.initialize()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Cleanup agent resources if needed."""
+        logger.info("Cleaning up agent resources")
+        self._in_context = False
+
+        # Do any necessary cleanup
+        # We're not shutting down the computer here as it might be shared
+        # Just log that we're exiting
+        if exc_type:
+            logger.error(f"Exiting agent context with error: {exc_type.__name__}: {exc_val}")
+        else:
+            logger.info("Exiting agent context normally")
+
+        # If we have a queue, make sure to signal it's done
+        if hasattr(self, "queue") and self.queue:
+            await self.queue.put(None)  # Signal that we're done
+
+    async def initialize(self) -> None:
+        """Initialize the agent and its components."""
+        if not self._initialized:
+            # Always initialize the computer if available
+            if self.computer and not self.computer._initialized:
+                await self.computer.run()
+            self._initialized = True
+
+    async def _init_if_needed(self):
+        """Initialize the computer interface if it hasn't been initialized yet."""
+        if not self.computer._initialized:
+            logger.info("Computer not initialized, initializing now...")
+            try:
+                # Call run directly
+                await self.computer.run()
+                logger.info("Computer interface initialized successfully")
+            except Exception as e:
+                logger.error(f"Error initializing computer interface: {str(e)}")
+                raise
+
+    async def run(self, task: str) -> AsyncGenerator[Dict[str, Any], None]:
+        """Run a task using the computer agent.
+
+        Args:
+            task: Task description
+
+        Yields:
+            Task execution updates
+        """
+        try:
+            logger.info(f"Running task: {task}")
+
+            # Initialize the computer if needed
+            if not self._initialized:
+                await self.initialize()
+
+            # Format task as a message
+            messages = [{"role": "user", "content": task}]
+
+            # Pass properly formatted messages to the loop
+            if self._loop is None:
+                logger.error("Loop not initialized properly")
+                yield {"error": "Loop not initialized properly"}
+                return
+
+            # Execute the task and yield results
+            async for result in self._loop.run(messages):
+                yield result
+
+        except Exception as e:
+            logger.error(f"Error in agent run method: {str(e)}")
+            yield {
+                "role": "assistant",
+                "content": f"Error: {str(e)}",
+                "metadata": {"title": "❌ Error"},
+            }

{cua_agent-0.1.5 → cua_agent-0.1.6}/agent/core/experiment.py

@@ -84,7 +84,21 @@ class ExperimentManager:
         if isinstance(data, dict):
             result = {}
             for k, v in data.items():
-                result[k] = self.sanitize_log_data(v)
+                # Special handling for 'data' field in Anthropic message source
+                if k == "data" and isinstance(v, str) and len(v) > 1000:
+                    result[k] = f"[BASE64_DATA_LENGTH_{len(v)}]"
+                # Special handling for the 'media_type' key which indicates we're in an image block
+                elif k == "media_type" and "image" in str(v):
+                    result[k] = v
+                    # If we're in an image block, look for a sibling 'data' field with base64 content
+                    if (
+                        "data" in result
+                        and isinstance(result["data"], str)
+                        and len(result["data"]) > 1000
+                    ):
+                        result["data"] = f"[BASE64_DATA_LENGTH_{len(result['data'])}]"
+                else:
+                    result[k] = self.sanitize_log_data(v)
             return result
         elif isinstance(data, list):
             return [self.sanitize_log_data(item) for item in data]
@@ -93,15 +107,18 @@ class ExperimentManager:
         else:
             return data
 
-    def save_screenshot(self, img_base64: str, action_type: str = "") -> None:
+    def save_screenshot(self, img_base64: str, action_type: str = "") -> Optional[str]:
         """Save a screenshot to the experiment directory.
 
         Args:
             img_base64: Base64 encoded screenshot
             action_type: Type of action that triggered the screenshot
+
+        Returns:
+            Path to the saved screenshot or None if there was an error
         """
         if not self.current_turn_dir:
-            return
+            return None
 
         try:
             # Increment screenshot counter

{cua_agent-0.1.5 → cua_agent-0.1.6}/agent/core/loop.py

@@ -141,9 +141,6 @@ class BaseLoop(ABC):
                 # Initialize API client
                 await self.initialize_client()
 
-                # Initialize computer
-                await self.computer.initialize()
-
                 logger.info("Initialization complete.")
                 return
             except Exception as e:
@@ -173,15 +170,22 @@ class BaseLoop(ABC):
             base64_image = ""
 
             # Handle different types of screenshot returns
-            if isinstance(screenshot, bytes):
+            if isinstance(screenshot, (bytes, bytearray, memoryview)):
                 # Raw bytes screenshot
                 base64_image = base64.b64encode(screenshot).decode("utf-8")
             elif hasattr(screenshot, "base64_image"):
                 # Object-style screenshot with attributes
-                base64_image = screenshot.base64_image
-                if hasattr(screenshot, "width") and hasattr(screenshot, "height"):
-                    width = screenshot.width
-                    height = screenshot.height
+                # Type checking can't infer these attributes, but they exist at runtime
+                # on certain screenshot return types
+                base64_image = getattr(screenshot, "base64_image")
+                width = (
+                    getattr(screenshot, "width", width) if hasattr(screenshot, "width") else width
+                )
+                height = (
+                    getattr(screenshot, "height", height)
+                    if hasattr(screenshot, "height")
+                    else height
+                )
 
             # Create parsed screen data
             parsed_screen = {

{cua_agent-0.1.5 → cua_agent-0.1.6}/agent/core/telemetry.py

@@ -4,58 +4,70 @@ import logging
 import os
 import platform
 import sys
-from typing import Dict, Any
+from typing import Dict, Any, Callable
 
 # Import the core telemetry module
 TELEMETRY_AVAILABLE = False
 
+
+# Local fallbacks in case core telemetry isn't available
+def _noop(*args: Any, **kwargs: Any) -> None:
+    """No-op function for when telemetry is not available."""
+    pass
+
+
+# Define default functions with unique names to avoid shadowing
+_default_record_event = _noop
+_default_increment_counter = _noop
+_default_set_dimension = _noop
+_default_get_telemetry_client = lambda: None
+_default_flush = _noop
+_default_is_telemetry_enabled = lambda: False
+_default_is_telemetry_globally_disabled = lambda: True
+
+# Set the actual functions to the defaults initially
+record_event = _default_record_event
+increment_counter = _default_increment_counter
+set_dimension = _default_set_dimension
+get_telemetry_client = _default_get_telemetry_client
+flush = _default_flush
+is_telemetry_enabled = _default_is_telemetry_enabled
+is_telemetry_globally_disabled = _default_is_telemetry_globally_disabled
+
+logger = logging.getLogger("cua.agent.telemetry")
+
 try:
+    # Import from core telemetry
     from core.telemetry import (
-        record_event,
-        increment,
-        get_telemetry_client,
-        flush,
-        is_telemetry_enabled,
-        is_telemetry_globally_disabled,
+        record_event as core_record_event,
+        increment as core_increment,
+        get_telemetry_client as core_get_telemetry_client,
+        flush as core_flush,
+        is_telemetry_enabled as core_is_telemetry_enabled,
+        is_telemetry_globally_disabled as core_is_telemetry_globally_disabled,
     )
 
+    # Override the default functions with actual implementations
+    record_event = core_record_event
+    get_telemetry_client = core_get_telemetry_client
+    flush = core_flush
+    is_telemetry_enabled = core_is_telemetry_enabled
+    is_telemetry_globally_disabled = core_is_telemetry_globally_disabled
+
     def increment_counter(counter_name: str, value: int = 1) -> None:
         """Wrapper for increment to maintain backward compatibility."""
         if is_telemetry_enabled():
-            increment(counter_name, value)
+            core_increment(counter_name, value)
 
     def set_dimension(name: str, value: Any) -> None:
         """Set a dimension that will be attached to all events."""
-        logger = logging.getLogger("cua.agent.telemetry")
         logger.debug(f"Setting dimension {name}={value}")
 
     TELEMETRY_AVAILABLE = True
-    logger = logging.getLogger("cua.agent.telemetry")
     logger.info("Successfully imported telemetry")
 except ImportError as e:
-    logger = logging.getLogger("cua.agent.telemetry")
     logger.warning(f"Could not import telemetry: {e}")
-    TELEMETRY_AVAILABLE = False
-
-
-# Local fallbacks in case core telemetry isn't available
-def _noop(*args: Any, **kwargs: Any) -> None:
-    """No-op function for when telemetry is not available."""
-    pass
-
-
-logger = logging.getLogger("cua.agent.telemetry")
-
-# If telemetry isn't available, use no-op functions
-if not TELEMETRY_AVAILABLE:
     logger.debug("Telemetry not available, using no-op functions")
-    record_event = _noop  # type: ignore
-    increment_counter = _noop  # type: ignore
-    set_dimension = _noop  # type: ignore
-    get_telemetry_client = lambda: None  # type: ignore
-    flush = _noop  # type: ignore
-    is_telemetry_enabled = lambda: False  # type: ignore
-    is_telemetry_globally_disabled = lambda: True  # type: ignore
 
 # Get system info once to use in telemetry
 SYSTEM_INFO = {
@@ -71,7 +83,7 @@ def enable_telemetry() -> bool:
     Returns:
         bool: True if telemetry was successfully enabled, False otherwise
     """
-    global TELEMETRY_AVAILABLE
+    global TELEMETRY_AVAILABLE, record_event, increment_counter, get_telemetry_client, flush, is_telemetry_enabled, is_telemetry_globally_disabled
 
     # Check if globally disabled using core function
     if TELEMETRY_AVAILABLE and is_telemetry_globally_disabled():

{cua_agent-0.1.5 → cua_agent-0.1.6}/agent/providers/anthropic/loop.py

@@ -17,6 +17,7 @@ from anthropic.types.beta import (
     BetaTextBlock,
     BetaTextBlockParam,
     BetaToolUseBlockParam,
+    BetaContentBlockParam,
 )
 
 # Computer
@@ -24,12 +25,12 @@ from computer import Computer
 
 # Base imports
 from ...core.loop import BaseLoop
-from ...core.messages import ImageRetentionConfig
+from ...core.messages import ImageRetentionConfig as CoreImageRetentionConfig
 
 # Anthropic provider-specific imports
 from .api.client import AnthropicClientFactory, BaseAnthropicClient
 from .tools.manager import ToolManager
-from .messages.manager import MessageManager
+from .messages.manager import MessageManager, ImageRetentionConfig
 from .callbacks.manager import CallbackManager
 from .prompts import SYSTEM_PROMPT
 from .types import LLMProvider
@@ -48,8 +49,8 @@ class AnthropicLoop(BaseLoop):
     def __init__(
         self,
         api_key: str,
+        computer: Computer,
         model: str = "claude-3-7-sonnet-20250219",  # Fixed model
-        computer: Optional[Computer] = None,
         only_n_most_recent_images: Optional[int] = 2,
         base_dir: Optional[str] = "trajectories",
         max_retries: int = 3,
@@ -69,7 +70,7 @@ class AnthropicLoop(BaseLoop):
             retry_delay: Delay between retries in seconds
             save_trajectory: Whether to save trajectory data
         """
-        # Initialize base class
+        # Initialize base class with core config
         super().__init__(
             computer=computer,
             model=model,
@@ -93,8 +94,8 @@ class AnthropicLoop(BaseLoop):
         self.message_manager = None
         self.callback_manager = None
 
-        # Configure image retention
-        self.image_retention_config = ImageRetentionConfig(
+        # Configure image retention with core config
+        self.image_retention_config = CoreImageRetentionConfig(
             num_images_to_keep=only_n_most_recent_images
         )
 
@@ -113,7 +114,7 @@ class AnthropicLoop(BaseLoop):
 
             # Initialize message manager
             self.message_manager = MessageManager(
-                ImageRetentionConfig(
+                image_retention_config=ImageRetentionConfig(
                     num_images_to_keep=self.only_n_most_recent_images, enable_caching=True
                 )
             )
@@ -250,6 +251,10 @@ class AnthropicLoop(BaseLoop):
                 await self._process_screen(parsed_screen, self.message_history)
 
                 # Prepare messages and make API call
+                if self.message_manager is None:
+                    raise RuntimeError(
+                        "Message manager not initialized. Call initialize_client() first."
+                    )
                 prepared_messages = self.message_manager.prepare_messages(
                     cast(List[BetaMessageParam], self.message_history.copy())
                 )
@@ -257,7 +262,7 @@ class AnthropicLoop(BaseLoop):
                 # Create new turn directory for this API call
                 self._create_turn_dir()
 
-                # Make API call
+                # Use _make_api_call instead of direct client call to ensure logging
                 response = await self._make_api_call(prepared_messages)
 
                 # Handle the response
@@ -287,6 +292,11 @@ class AnthropicLoop(BaseLoop):
         Returns:
             API response
         """
+        if self.client is None:
+            raise RuntimeError("Client not initialized. Call initialize_client() first.")
+        if self.tool_manager is None:
+            raise RuntimeError("Tool manager not initialized. Call initialize_client() first.")
+
         last_error = None
 
         for attempt in range(self.max_retries):
@@ -297,6 +307,7 @@ class AnthropicLoop(BaseLoop):
                     "max_tokens": self.max_tokens,
                     "system": SYSTEM_PROMPT,
                 }
+                # Let ExperimentManager handle sanitization
                 self._log_api_call("request", request_data)
 
                 # Setup betas and system
@@ -320,7 +331,7 @@ class AnthropicLoop(BaseLoop):
                     betas=betas,
                 )
 
-                # Log success response
+                # Let ExperimentManager handle sanitization
                 self._log_api_call("response", request_data, response)
 
                 return response
@@ -365,25 +376,38 @@ class AnthropicLoop(BaseLoop):
                 }
             )
 
+            if self.callback_manager is None:
+                raise RuntimeError(
+                    "Callback manager not initialized. Call initialize_client() first."
+                )
+
             # Handle tool use blocks and collect results
             tool_result_content = []
             for content_block in response_params:
                 # Notify callback of content
-                self.callback_manager.on_content(content_block)
+                self.callback_manager.on_content(cast(BetaContentBlockParam, content_block))
 
                 # Handle tool use
                 if content_block.get("type") == "tool_use":
+                    if self.tool_manager is None:
+                        raise RuntimeError(
+                            "Tool manager not initialized. Call initialize_client() first."
+                        )
                     result = await self.tool_manager.execute_tool(
                         name=content_block["name"],
                         tool_input=cast(Dict[str, Any], content_block["input"]),
                     )
 
                     # Create tool result and add to content
-                    tool_result = self._make_tool_result(result, content_block["id"])
+                    tool_result = self._make_tool_result(
+                        cast(ToolResult, result), content_block["id"]
+                    )
                     tool_result_content.append(tool_result)
 
                     # Notify callback of tool result
-                    self.callback_manager.on_tool_result(result, content_block["id"])
+                    self.callback_manager.on_tool_result(
+                        cast(ToolResult, result), content_block["id"]
+                    )
 
             # If no tool results, we're done
             if not tool_result_content:
@@ -495,13 +519,13 @@ class AnthropicLoop(BaseLoop):
             result_text = f"<s>{result.system}</s>\n{result_text}"
         return result_text
 
-    def _handle_content(self, content: Dict[str, Any]) -> None:
+    def _handle_content(self, content: BetaContentBlockParam) -> None:
         """Handle content updates from the assistant."""
         if content.get("type") == "text":
-            text = content.get("text", "")
+            text_content = cast(BetaTextBlockParam, content)
+            text = text_content["text"]
             if text == "<DONE>":
                 return
-
             logger.info(f"Assistant: {text}")
 
     def _handle_tool_result(self, result: ToolResult, tool_id: str) -> None:
@@ -517,5 +541,10 @@ class AnthropicLoop(BaseLoop):
         """Handle API interactions."""
         if error:
             logger.error(f"API error: {error}")
+            self._log_api_call("error", request, error=error)
         else:
             logger.debug(f"API request: {request}")
+            if response:
+                self._log_api_call("response", request, response)
+            else:
+                self._log_api_call("request", request)

{cua_agent-0.1.5 → cua_agent-0.1.6}/agent/providers/anthropic/messages/manager.py

@@ -90,7 +90,9 @@ class MessageManager:
                     blocks_with_cache_control += 1
                     # Add cache control to the last content block only
                     if content and len(content) > 0:
-                        content[-1]["cache_control"] = {"type": "ephemeral"}
+                        content[-1]["cache_control"] = BetaCacheControlEphemeralParam(
+                            type="ephemeral"
+                        )
                 else:
                     # Remove any existing cache control
                     if content and len(content) > 0: