cua-agent 0.4.14__py3-none-any.whl → 0.4.16__py3-none-any.whl

agent/__init__.py

@@ -28,13 +28,9 @@ try:
     # Import from core telemetry for basic functions
     from core.telemetry import (
         is_telemetry_enabled,
-        flush,
         record_event,
     )
 
-    # Import set_dimension from our own telemetry module
-    from .telemetry import set_dimension
-
     # Check if telemetry is enabled
     if is_telemetry_enabled():
         logger.info("Telemetry is enabled")
@@ -49,11 +45,6 @@ try:
             },
         )
 
-        # Set the package version as a dimension
-        set_dimension("agent_version", __version__)
-
-        # Flush events to ensure they're sent
-        flush()
     else:
         logger.info("Telemetry is disabled")
 except ImportError as e:

agent/agent.py

@@ -7,7 +7,13 @@ from typing import Dict, List, Any, Optional, AsyncGenerator, Union, cast, Calla
 
 from litellm.responses.utils import Usage
 
-from .types import Messages, AgentCapability
+from .types import (
+    Messages,
+    AgentCapability,
+    ToolError,
+    IllegalArgumentError
+)
+from .responses import make_tool_error_item, replace_failed_computer_calls_with_function_calls
 from .decorators import find_agent_config
 import json
 import litellm
@@ -30,6 +36,15 @@ from .computers import (
     make_computer_handler
 )
 
+def assert_callable_with(f, *args, **kwargs):
+   """Check if function can be called with given arguments."""
+   try:
+       inspect.signature(f).bind(*args, **kwargs)
+       return True
+   except TypeError as e:
+       sig = inspect.signature(f)
+       raise IllegalArgumentError(f"Expected {sig}, got args={args} kwargs={kwargs}") from e
+
 def get_json(obj: Any, max_depth: int = 10) -> Any:
     def custom_serializer(o: Any, depth: int = 0, seen: Optional[Set[int]] = None) -> Any:
         if seen is None:
@@ -405,7 +420,8 @@ class ComputerAgent:
     
     async def _handle_item(self, item: Any, computer: Optional[AsyncComputerHandler] = None, ignore_call_ids: Optional[List[str]] = None) -> List[Dict[str, Any]]:
         """Handle each item; may cause a computer action + screenshot."""
-        if ignore_call_ids and item.get("call_id") and item.get("call_id") in ignore_call_ids:
+        call_id = item.get("call_id")
+        if ignore_call_ids and call_id and call_id in ignore_call_ids:
             return []
         
         item_type = item.get("type", None)
@@ -419,96 +435,103 @@ class ComputerAgent:
             #             print(content_item.get("text"))
             return []
         
-        if item_type == "computer_call":
-            await self._on_computer_call_start(item)
-            if not computer:
-                raise ValueError("Computer handler is required for computer calls")
-
-            # Perform computer actions
-            action = item.get("action")
-            action_type = action.get("type")
-            if action_type is None:
-                print(f"Action type cannot be `None`: action={action}, action_type={action_type}")
-                return []
-            
-            # Extract action arguments (all fields except 'type')
-            action_args = {k: v for k, v in action.items() if k != "type"}
-            
-            # print(f"{action_type}({action_args})")
-            
-            # Execute the computer action
-            computer_method = getattr(computer, action_type, None)
-            if computer_method:
-                await computer_method(**action_args)
-            else:
-                print(f"Unknown computer action: {action_type}")
-                return []
-            
-            # Take screenshot after action
-            if self.screenshot_delay and self.screenshot_delay > 0:
-                await asyncio.sleep(self.screenshot_delay)
-            screenshot_base64 = await computer.screenshot()
-            await self._on_screenshot(screenshot_base64, "screenshot_after")
+        try:
+            if item_type == "computer_call":
+                await self._on_computer_call_start(item)
+                if not computer:
+                    raise ValueError("Computer handler is required for computer calls")
+
+                # Perform computer actions
+                action = item.get("action")
+                action_type = action.get("type")
+                if action_type is None:
+                    print(f"Action type cannot be `None`: action={action}, action_type={action_type}")
+                    return []
+                
+                # Extract action arguments (all fields except 'type')
+                action_args = {k: v for k, v in action.items() if k != "type"}
+                
+                # print(f"{action_type}({action_args})")
+                
+                # Execute the computer action
+                computer_method = getattr(computer, action_type, None)
+                if computer_method:
+                    assert_callable_with(computer_method, **action_args)
+                    await computer_method(**action_args)
+                else:
+                    print(f"Unknown computer action: {action_type}")
+                    return []
+                
+                # Take screenshot after action
+                if self.screenshot_delay and self.screenshot_delay > 0:
+                    await asyncio.sleep(self.screenshot_delay)
+                screenshot_base64 = await computer.screenshot()
+                await self._on_screenshot(screenshot_base64, "screenshot_after")
+                
+                # Handle safety checks
+                pending_checks = item.get("pending_safety_checks", [])
+                acknowledged_checks = []
+                for check in pending_checks:
+                    check_message = check.get("message", str(check))
+                    acknowledged_checks.append(check)
+                    # TODO: implement a callback for safety checks
+                    # if acknowledge_safety_check_callback(check_message, allow_always=True):
+                    #     acknowledged_checks.append(check)
+                    # else:
+                    #     raise ValueError(f"Safety check failed: {check_message}")
+                
+                # Create call output
+                call_output = {
+                    "type": "computer_call_output",
+                    "call_id": item.get("call_id"),
+                    "acknowledged_safety_checks": acknowledged_checks,
+                    "output": {
+                        "type": "input_image",
+                        "image_url": f"data:image/png;base64,{screenshot_base64}",
+                    },
+                }
+                
+                # # Additional URL safety checks for browser environments
+                # if await computer.get_environment() == "browser":
+                #     current_url = await computer.get_current_url()
+                #     call_output["output"]["current_url"] = current_url
+                #     # TODO: implement a callback for URL safety checks
+                #     # check_blocklisted_url(current_url)
+                
+                result = [call_output]
+                await self._on_computer_call_end(item, result)
+                return result
             
-            # Handle safety checks
-            pending_checks = item.get("pending_safety_checks", [])
-            acknowledged_checks = []
-            for check in pending_checks:
-                check_message = check.get("message", str(check))
-                acknowledged_checks.append(check)
-                # TODO: implement a callback for safety checks
-                # if acknowledge_safety_check_callback(check_message, allow_always=True):
-                #     acknowledged_checks.append(check)
-                # else:
-                #     raise ValueError(f"Safety check failed: {check_message}")
+            if item_type == "function_call":
+                await self._on_function_call_start(item)
+                # Perform function call
+                function = self._get_tool(item.get("name"))
+                if not function:
+                    raise ValueError(f"Function {item.get("name")} not found")
             
-            # Create call output
-            call_output = {
-                "type": "computer_call_output",
-                "call_id": item.get("call_id"),
-                "acknowledged_safety_checks": acknowledged_checks,
-                "output": {
-                    "type": "input_image",
-                    "image_url": f"data:image/png;base64,{screenshot_base64}",
-                },
-            }
+                args = json.loads(item.get("arguments"))
+
+                # Validate arguments before execution
+                assert_callable_with(function, **args)
+
+                # Execute function - use asyncio.to_thread for non-async functions
+                if inspect.iscoroutinefunction(function):
+                    result = await function(**args)
+                else:
+                    result = await asyncio.to_thread(function, **args)
             
-            # # Additional URL safety checks for browser environments
-            # if await computer.get_environment() == "browser":
-            #     current_url = await computer.get_current_url()
-            #     call_output["output"]["current_url"] = current_url
-            #     # TODO: implement a callback for URL safety checks
-            #     # check_blocklisted_url(current_url)
+                # Create function call output
+                call_output = {
+                    "type": "function_call_output",
+                    "call_id": item.get("call_id"),
+                    "output": str(result),
+                }
             
-            result = [call_output]
-            await self._on_computer_call_end(item, result)
-            return result
-        
-        if item_type == "function_call":
-            await self._on_function_call_start(item)
-            # Perform function call
-            function = self._get_tool(item.get("name"))
-            if not function:
-                raise ValueError(f"Function {item.get("name")} not found")
-        
-            args = json.loads(item.get("arguments"))
-
-            # Execute function - use asyncio.to_thread for non-async functions
-            if inspect.iscoroutinefunction(function):
-                result = await function(**args)
-            else:
-                result = await asyncio.to_thread(function, **args)
-        
-            # Create function call output
-            call_output = {
-                "type": "function_call_output",
-                "call_id": item.get("call_id"),
-                "output": str(result),
-            }
-        
-            result = [call_output]
-            await self._on_function_call_end(item, result)
-            return result
+                result = [call_output]
+                await self._on_function_call_end(item, result)
+                return result
+        except ToolError as e:
+            return [make_tool_error_item(repr(e), call_id)]
 
         return []
 
@@ -569,6 +592,7 @@ class ComputerAgent:
             # - PII anonymization
             # - Image retention policy
             combined_messages = old_items + new_items
+            combined_messages = replace_failed_computer_calls_with_function_calls(combined_messages)
             preprocessed_messages = await self._on_llm_start(combined_messages)
             
             loop_kwargs = {

agent/callbacks/telemetry.py

@@ -7,13 +7,18 @@ import uuid
 from typing import List, Dict, Any, Optional, Union
 
 from .base import AsyncCallbackHandler
-from ..telemetry import (
+from core.telemetry import (
     record_event,
     is_telemetry_enabled,
-    set_dimension,
-    SYSTEM_INFO,
 )
 
+import platform
+
+SYSTEM_INFO = {
+    "os": platform.system().lower(),
+    "os_version": platform.release(),
+    "python_version": platform.python_version(),
+}
 
 class TelemetryCallback(AsyncCallbackHandler):
     """
@@ -65,11 +70,6 @@ class TelemetryCallback(AsyncCallbackHandler):
             **SYSTEM_INFO
         }
         
-        # Set session-level dimensions
-        set_dimension("session_id", self.session_id)
-        set_dimension("agent_type", agent_info["agent_type"])
-        set_dimension("model", agent_info["model"])
-        
         record_event("agent_session_start", agent_info)
     
     async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
@@ -98,7 +98,6 @@ class TelemetryCallback(AsyncCallbackHandler):
             if trajectory:
                 run_data["uploaded_trajectory"] = trajectory
         
-        set_dimension("run_id", self.run_id)
         record_event("agent_run_start", run_data)
     
     async def on_run_end(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]], new_items: List[Dict[str, Any]]) -> None:

agent/responses.py

@@ -252,6 +252,53 @@ def make_failed_tool_call_items(tool_name: str, tool_kwargs: Dict[str, Any], err
         }
     ]
 
+def make_tool_error_item(error_message: str, call_id: Optional[str] = None) -> Dict[str, Any]:
+    call_id = call_id if call_id else random_id()
+    return {
+        "type": "function_call_output",
+        "call_id": call_id,
+        "output": json.dumps({"error": error_message}),
+    }
+
+def replace_failed_computer_calls_with_function_calls(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """
+    Replace computer_call items with function_call items if they share a call_id with a function_call_output.
+    This indicates the computer call failed and should be treated as a function call instead.
+    We do this because the computer_call_output items do not support text output.
+    
+    Args:
+        messages: List of message items to process
+    """
+    messages = messages.copy()
+
+    # Find all call_ids that have function_call_output items
+    failed_call_ids = set()
+    for msg in messages:
+        if msg.get("type") == "function_call_output":
+            call_id = msg.get("call_id")
+            if call_id:
+                failed_call_ids.add(call_id)
+    
+    # Replace computer_call items that have matching call_ids
+    for i, msg in enumerate(messages):
+        if (msg.get("type") == "computer_call" and 
+            msg.get("call_id") in failed_call_ids):
+            
+            # Extract action from computer_call
+            action = msg.get("action", {})
+            call_id = msg.get("call_id")
+            
+            # Create function_call replacement
+            messages[i] = {
+                "type": "function_call",
+                "id": msg.get("id", random_id()),
+                "call_id": call_id,
+                "name": "computer",
+                "arguments": json.dumps(action),
+            }
+    
+    return messages
+
 # Conversion functions between element descriptions and coordinates
 def convert_computer_calls_desc2xy(responses_items: List[Dict[str, Any]], desc2xy: Dict[str, tuple]) -> List[Dict[str, Any]]:
     """

agent/types.py

@@ -16,6 +16,15 @@ Tools = Optional[Iterable[ToolParam]]
 AgentResponse = ResponsesAPIResponse 
 AgentCapability = Literal["step", "click"]
 
+# Exception types
+class ToolError(RuntimeError):
+    """Base exception for tool-related errors"""
+    pass
+
+class IllegalArgumentError(ToolError):
+    """Exception raised when function arguments are invalid"""
+    pass
+
 
 # Agent config registration
 class AgentConfigInfo(BaseModel):

agent/ui/gradio/app.py

@@ -114,14 +114,14 @@ MODEL_MAPPINGS = {
         "Anthropic: Claude 4 Opus (20250514)": "anthropic/claude-opus-4-20250514",
         "Anthropic: Claude 4 Sonnet (20250514)": "anthropic/claude-sonnet-4-20250514",
         "Anthropic: Claude 3.7 Sonnet (20250219)": "anthropic/claude-3-7-sonnet-20250219",
-        "Anthropic: Claude 3.5 Sonnet (20240620)": "anthropic/claude-3-5-sonnet-20240620",
+        "Anthropic: Claude 3.5 Sonnet (20241022)": "anthropic/claude-3-5-sonnet-20241022",
     },
     "omni": {
         "default": "omniparser+openai/gpt-4o",
         "OMNI: OpenAI GPT-4o": "omniparser+openai/gpt-4o",
         "OMNI: OpenAI GPT-4o mini": "omniparser+openai/gpt-4o-mini",
         "OMNI: Claude 3.7 Sonnet (20250219)": "omniparser+anthropic/claude-3-7-sonnet-20250219",
-        "OMNI: Claude 3.5 Sonnet (20240620)": "omniparser+anthropic/claude-3-5-sonnet-20240620",
+        "OMNI: Claude 3.5 Sonnet (20241022)": "omniparser+anthropic/claude-3-5-sonnet-20241022",
     },
     "uitars": {
         "default": "huggingface-local/ByteDance-Seed/UI-TARS-1.5-7B" if is_mac else "ui-tars",

agent/ui/gradio/ui_components.py

@@ -38,13 +38,13 @@ def create_gradio_ui() -> gr.Blocks:
         "Anthropic: Claude 4 Opus (20250514)",
         "Anthropic: Claude 4 Sonnet (20250514)",
         "Anthropic: Claude 3.7 Sonnet (20250219)",
-        "Anthropic: Claude 3.5 Sonnet (20240620)",
+        "Anthropic: Claude 3.5 Sonnet (20241022)",
     ]
     omni_models = [
         "OMNI: OpenAI GPT-4o",
         "OMNI: OpenAI GPT-4o mini",
         "OMNI: Claude 3.7 Sonnet (20250219)", 
-        "OMNI: Claude 3.5 Sonnet (20240620)"
+        "OMNI: Claude 3.5 Sonnet (20241022)"
     ]
     
     # Check if API keys are available

{cua_agent-0.4.14.dist-info → cua_agent-0.4.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cua-agent
-Version: 0.4.14
+Version: 0.4.16
 Summary: CUA (Computer Use) Agent for AI-driven computer interaction
 Author-Email: TryCua <gh@trycua.com>
 Requires-Python: >=3.12
@@ -56,8 +56,8 @@ Description-Content-Type: text/markdown
 <h1>
   <div class="image-wrapper" style="display: inline-block;">
     <picture>
-      <source media="(prefers-color-scheme: dark)" alt="logo" height="150" srcset="../../../img/logo_white.png" style="display: block; margin: auto;">
-      <source media="(prefers-color-scheme: light)" alt="logo" height="150" srcset="../../../img/logo_black.png" style="display: block; margin: auto;">
+      <source media="(prefers-color-scheme: dark)" alt="logo" height="150" srcset="https://raw.githubusercontent.com/trycua/cua/main/img/logo_white.png" style="display: block; margin: auto;">
+      <source media="(prefers-color-scheme: light)" alt="logo" height="150" srcset="https://raw.githubusercontent.com/trycua/cua/main/img/logo_black.png" style="display: block; margin: auto;">
       <img alt="Shows my svg">
     </picture>
   </div>
@@ -138,7 +138,7 @@ if __name__ == "__main__":
 ### Anthropic Claude (Computer Use API)
 ```python
 model="anthropic/claude-3-5-sonnet-20241022"
-model="anthropic/claude-3-5-sonnet-20240620"
+model="anthropic/claude-3-7-sonnet-20250219"
 model="anthropic/claude-opus-4-20250514"
 model="anthropic/claude-sonnet-4-20250514"
 ```

{cua_agent-0.4.14.dist-info → cua_agent-0.4.16.dist-info}/RECORD

@@ -1,16 +1,16 @@
-agent/__init__.py,sha256=vWbQYgjkzIso7zILSm4OAbNU_vrmN4HyYkfX8vC-Yi0,1547
+agent/__init__.py,sha256=MaW-BczJ-lCACPYH39DvFhE7ZWiSo7sBO6pBfyO7Nxc,1269
 agent/__main__.py,sha256=lBUe8Niqa5XoCjwFfXyX7GtnUwjjZXC1-j4V9mvUYSc,538
 agent/adapters/__init__.py,sha256=lNH6srgIMmZOI7dgicJs3LCk_1MeqLF0lou9n7b23Ts,238
 agent/adapters/huggingfacelocal_adapter.py,sha256=Uqjtcohhzd33VFh38Ra2y4Uv_lTghMswoqS1t-KKFkw,8480
 agent/adapters/human_adapter.py,sha256=xT4nnfNXb1z-vnGFlLmFEZN7TMcoMBGS40MtR1Zwv4o,13079
-agent/agent.py,sha256=mEbmN5G6y8jZ0FrlUnHfJQFE7r_GlXrHxqC93twv54k,27881
+agent/agent.py,sha256=Zo7K7FacdCwGUDW3cuMCpJ1h9tkULH3hrgx94tdjCsk,29123
 agent/callbacks/__init__.py,sha256=yxxBXUqpXQ-jRi_ixJMtmQPxoNRy5Vz1PUBzNNa1Dwg,538
 agent/callbacks/base.py,sha256=UnnnYlh6XCm6HKZZsAPaT_Eyo9LUYLyjyNwF-QRm6Ns,4691
 agent/callbacks/budget_manager.py,sha256=RyKM-7iXQcDotYvrw3eURzeEHEXvQjID-NobtvQWE7k,1832
 agent/callbacks/image_retention.py,sha256=tiuRT5ke9xXTb2eP8Gz-2ITyAMY29LURUH6AbjX3RP8,6165
 agent/callbacks/logging.py,sha256=OOxU97EzrxlnUAtiEnvy9FB7SwCUK90-rdpDFA2Ae4E,10921
 agent/callbacks/pii_anonymization.py,sha256=NEkUTUjQBi82nqus7kT-1E4RaeQ2hQrY7YCnKndLhP8,3272
-agent/callbacks/telemetry.py,sha256=PU7pkK7W1v1xjDN-9gA30lGvn4-WhqK3BPHGW3HpTOc,7497
+agent/callbacks/telemetry.py,sha256=RbUDhE41mTi8g9hNre0EpltK_NUZkLj8buJLWBzs0Ek,7363
 agent/callbacks/trajectory_saver.py,sha256=VHbiDQzI_XludkWhZIVqIMrsxgwKfFWwVtqaRot_D4U,12231
 agent/cli.py,sha256=AgaXwywHd3nGQWuqMRj6SbPyFaCPjfo5980Y1ApQOTQ,12413
 agent/computers/__init__.py,sha256=39ISJsaREaQIZckpzxSuLhuR763wUU3TxUux78EKjAg,1477
@@ -36,15 +36,14 @@ agent/loops/model_types.csv,sha256=GmFn4x80yoUpQZuQ-GXtJkPVlOLYWZ5u_5A73HRyeNE,1
 agent/loops/omniparser.py,sha256=-db8JUL2Orn47ERIaLbuNShAXn4LeIgYzRWphn_9Dg4,15071
 agent/loops/openai.py,sha256=8Ad_XufpENmLq1nEnhzF3oswPrPK1EPz-C5NU8UOEs0,8035
 agent/loops/uitars.py,sha256=PVNOdwcn2K6RgaxoU-9I4HjBTsEH073M11LTqTrN7C4,31849
-agent/responses.py,sha256=TTJ3wXN_eb0J26GKhO3cVQngOiZ1AgUPIUadozLUQyE,28991
-agent/telemetry.py,sha256=87ZTyBaT0wEPQn4v76II3g0V3GERuIVbypoX-Ug6FKQ,4786
-agent/types.py,sha256=ZoWY8a3GZtB8V0SnOzoI7DQy4nP_GRubxJKbuLPOc8c,840
+agent/responses.py,sha256=_SoN4BkaTxMHMB21EOtDc_aDBIJlfDwsCzszMBnIkH0,30764
+agent/types.py,sha256=h6SnmTAEAaryVCjwVZFAuCbio9UW13OqgQEV7HKmZVM,1060
 agent/ui/__init__.py,sha256=DTZpK85QXscXK2nM9HtpAhVBF13yAamUrtwrQSuV-kM,126
 agent/ui/__main__.py,sha256=vudWXYvGM0aNT5aZ94HPtGW8YXOZ4cLXepHyhUM_k1g,73
 agent/ui/gradio/__init__.py,sha256=yv4Mrfo-Sj2U5sVn_UJHAuwYCezo-5O4ItR2C9jzNko,145
-agent/ui/gradio/app.py,sha256=m2yDd6Tua_lMMZT1zCzOty2meYEy756d8OlFF7lpdeU,9117
-agent/ui/gradio/ui_components.py,sha256=vfsqVo_COsFfw11ouMHClib9fdBf3q52G-qbuo0RyOY,36068
-cua_agent-0.4.14.dist-info/METADATA,sha256=uLZvYOCvm5B-sX4z-WNZ1CYLmpZYQd31Rk-33MwzOMM,12616
-cua_agent-0.4.14.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-cua_agent-0.4.14.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-cua_agent-0.4.14.dist-info/RECORD,,
+agent/ui/gradio/app.py,sha256=Ol97YEbwREZZQ9_PMjVHlfOcu9BGsawxgAGAm79hT80,9117
+agent/ui/gradio/ui_components.py,sha256=dJUvKDmc1oSejtoR_gU_oWWYwxaOOQyPloSYRGMrUCQ,36068
+cua_agent-0.4.16.dist-info/METADATA,sha256=iKOsC1Me0Rq_jQrhfgLN-uY1KR5OtKIQXpRGXpnH_LE,12698
+cua_agent-0.4.16.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+cua_agent-0.4.16.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+cua_agent-0.4.16.dist-info/RECORD,,

agent/telemetry.py

@@ -1,142 +0,0 @@
-"""Agent telemetry for tracking anonymous usage and feature usage."""
-
-import logging
-import os
-import platform
-import sys
-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("agent.telemetry")
-
-try:
-    # Import from core telemetry
-    from core.telemetry import (
-        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():
-            core_increment(counter_name, value)
-
-    def set_dimension(name: str, value: Any) -> None:
-        """Set a dimension that will be attached to all events."""
-        logger.debug(f"Setting dimension {name}={value}")
-
-    TELEMETRY_AVAILABLE = True
-    logger.info("Successfully imported telemetry")
-except ImportError as e:
-    logger.warning(f"Could not import telemetry: {e}")
-    logger.debug("Telemetry not available, using no-op functions")
-
-# Get system info once to use in telemetry
-SYSTEM_INFO = {
-    "os": platform.system().lower(),
-    "os_version": platform.release(),
-    "python_version": platform.python_version(),
-}
-
-
-def enable_telemetry() -> bool:
-    """Enable telemetry if available.
-
-    Returns:
-        bool: True if telemetry was successfully enabled, False otherwise
-    """
-    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():
-        logger.info("Telemetry is globally disabled via environment variable - cannot enable")
-        return False
-
-    # Already enabled
-    if TELEMETRY_AVAILABLE:
-        return True
-
-    # Try to import and enable
-    try:
-        from core.telemetry import (
-            record_event,
-            increment,
-            get_telemetry_client,
-            flush,
-            is_telemetry_globally_disabled,
-        )
-
-        # Check again after import
-        if is_telemetry_globally_disabled():
-            logger.info("Telemetry is globally disabled via environment variable - cannot enable")
-            return False
-
-        TELEMETRY_AVAILABLE = True
-        logger.info("Telemetry successfully enabled")
-        return True
-    except ImportError as e:
-        logger.warning(f"Could not enable telemetry: {e}")
-        return False
-
-
-def is_telemetry_enabled() -> bool:
-    """Check if telemetry is enabled.
-
-    Returns:
-        bool: True if telemetry is enabled, False otherwise
-    """
-    # Use the core function if available, otherwise use our local flag
-    if TELEMETRY_AVAILABLE:
-        from core.telemetry import is_telemetry_enabled as core_is_enabled
-
-        return core_is_enabled()
-    return False
-
-
-def record_agent_initialization() -> None:
-    """Record when an agent instance is initialized."""
-    if TELEMETRY_AVAILABLE and is_telemetry_enabled():
-        record_event("agent_initialized", SYSTEM_INFO)
-
-        # Set dimensions that will be attached to all events
-        set_dimension("os", SYSTEM_INFO["os"])
-        set_dimension("os_version", SYSTEM_INFO["os_version"])
-        set_dimension("python_version", SYSTEM_INFO["python_version"])