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

agent/__init__.py

@@ -5,19 +5,13 @@ agent - Decorator-based Computer Use Agent with liteLLM integration
 import logging
 import sys
 
-from .decorators import register_agent
-from .agent import ComputerAgent
-from .types import Messages, AgentResponse
-
 # Import loops to register them
 from . import loops
+from .agent import ComputerAgent
+from .decorators import register_agent
+from .types import AgentResponse, Messages
 
-__all__ = [
-    "register_agent",
-    "ComputerAgent",
-    "Messages",
-    "AgentResponse"
-]
+__all__ = ["register_agent", "ComputerAgent", "Messages", "AgentResponse"]
 
 __version__ = "0.4.0"
 
@@ -28,13 +22,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 +39,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/__main__.py

@@ -5,8 +5,9 @@ Usage:
     python -m agent.cli <model_string>
 """
 
-import sys
 import asyncio
+import sys
+
 from .cli import main
 
 if __name__ == "__main__":

agent/adapters/__init__.py

@@ -2,10 +2,16 @@
 Adapters package for agent - Custom LLM adapters for LiteLLM
 """
 
+from .azure_ml_adapter import AzureMLAdapter
+from .cua_adapter import CUAAdapter
 from .huggingfacelocal_adapter import HuggingFaceLocalAdapter
 from .human_adapter import HumanAdapter
+from .mlxvlm_adapter import MLXVLMAdapter
 
 __all__ = [
+    "AzureMLAdapter",
     "HuggingFaceLocalAdapter",
     "HumanAdapter",
+    "MLXVLMAdapter",
+    "CUAAdapter",
 ]

agent/adapters/azure_ml_adapter.py

@@ -0,0 +1,283 @@
+"""
+Azure ML Custom Provider Adapter for LiteLLM.
+
+This adapter provides direct OpenAI-compatible API access to Azure ML endpoints
+without message transformation, specifically for models like Fara-7B that require
+exact OpenAI message formatting.
+"""
+
+import json
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional
+
+import httpx
+from litellm import acompletion, completion
+from litellm.llms.custom_llm import CustomLLM
+from litellm.types.utils import GenericStreamingChunk, ModelResponse
+
+
+class AzureMLAdapter(CustomLLM):
+    """
+    Azure ML Adapter for OpenAI-compatible endpoints.
+
+    Makes direct HTTP calls to Azure ML foundry inference endpoints
+    using the OpenAI-compatible API format without transforming messages.
+
+    Usage:
+        model = "azure_ml/Fara-7B"
+        api_base = "https://foundry-inference-xxx.centralus.inference.ml.azure.com"
+        api_key = "your-api-key"
+
+        response = litellm.completion(
+            model=model,
+            messages=[...],
+            api_base=api_base,
+            api_key=api_key
+        )
+    """
+
+    def __init__(self, **kwargs):
+        """Initialize the adapter."""
+        super().__init__()
+        self._client: Optional[httpx.Client] = None
+        self._async_client: Optional[httpx.AsyncClient] = None
+
+    def _get_client(self) -> httpx.Client:
+        """Get or create sync HTTP client."""
+        if self._client is None:
+            self._client = httpx.Client(timeout=600.0)
+        return self._client
+
+    def _get_async_client(self) -> httpx.AsyncClient:
+        """Get or create async HTTP client."""
+        if self._async_client is None:
+            self._async_client = httpx.AsyncClient(timeout=600.0)
+        return self._async_client
+
+    def _prepare_request(self, **kwargs) -> tuple[str, dict, dict]:
+        """
+        Prepare the HTTP request without transforming messages.
+
+        Applies Azure ML workaround: double-encodes function arguments to work around
+        Azure ML's bug where it parses arguments before validation.
+
+        Returns:
+            Tuple of (url, headers, json_data)
+        """
+        # Extract required params
+        api_base = kwargs.get("api_base")
+        api_key = kwargs.get("api_key")
+        model = kwargs.get("model", "").replace("azure_ml/", "")
+        messages = kwargs.get("messages", [])
+
+        if not api_base:
+            raise ValueError("api_base is required for azure_ml provider")
+        if not api_key:
+            raise ValueError("api_key is required for azure_ml provider")
+
+        # Build OpenAI-compatible endpoint URL
+        base_url = api_base.rstrip("/")
+        url = f"{base_url}/chat/completions"
+
+        # Prepare headers
+        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+
+        # WORKAROUND for Azure ML bug:
+        # Azure ML incorrectly parses the arguments field before validation,
+        # causing it to reject valid JSON strings. We double-encode arguments
+        # so that after Azure ML's parse, they remain as strings.
+        messages_copy = []
+        for message in messages:
+            msg_copy = message.copy()
+
+            # Check if message has tool_calls that need double-encoding
+            if "tool_calls" in msg_copy:
+                tool_calls_copy = []
+                for tool_call in msg_copy["tool_calls"]:
+                    tc_copy = tool_call.copy()
+
+                    if "function" in tc_copy and "arguments" in tc_copy["function"]:
+                        func_copy = tc_copy["function"].copy()
+                        arguments = func_copy["arguments"]
+
+                        # If arguments is already a string, double-encode it
+                        if isinstance(arguments, str):
+                            func_copy["arguments"] = json.dumps(arguments)
+
+                        tc_copy["function"] = func_copy
+
+                    tool_calls_copy.append(tc_copy)
+
+                msg_copy["tool_calls"] = tool_calls_copy
+
+            messages_copy.append(msg_copy)
+
+        # Prepare request body with double-encoded messages
+        json_data = {"model": model, "messages": messages_copy}
+
+        # Add optional parameters if provided
+        optional_params = [
+            "temperature",
+            "top_p",
+            "n",
+            "stream",
+            "stop",
+            "max_tokens",
+            "presence_penalty",
+            "frequency_penalty",
+            "logit_bias",
+            "user",
+            "response_format",
+            "seed",
+            "tools",
+            "tool_choice",
+        ]
+
+        for param in optional_params:
+            if param in kwargs and kwargs[param] is not None:
+                json_data[param] = kwargs[param]
+
+        return url, headers, json_data
+
+    def completion(self, *args, **kwargs) -> ModelResponse:
+        """
+        Synchronous completion method.
+
+        Makes a direct HTTP POST to Azure ML's OpenAI-compatible endpoint.
+        """
+        url, headers, json_data = self._prepare_request(**kwargs)
+
+        client = self._get_client()
+        response = client.post(url, headers=headers, json=json_data)
+        response.raise_for_status()
+
+        # Parse response
+        response_json = response.json()
+
+        # Return using litellm's completion with the actual response
+        return completion(
+            model=f"azure_ml/{kwargs.get('model', '')}",
+            mock_response=response_json["choices"][0]["message"]["content"],
+            messages=kwargs.get("messages", []),
+        )
+
+    async def acompletion(self, *args, **kwargs) -> ModelResponse:
+        """
+        Asynchronous completion method.
+
+        Makes a direct async HTTP POST to Azure ML's OpenAI-compatible endpoint.
+        """
+        url, headers, json_data = self._prepare_request(**kwargs)
+
+        client = self._get_async_client()
+        response = await client.post(url, headers=headers, json=json_data)
+        response.raise_for_status()
+
+        # Parse response
+        response_json = response.json()
+
+        # Return using litellm's acompletion with the actual response
+        return await acompletion(
+            model=f"azure_ml/{kwargs.get('model', '')}",
+            mock_response=response_json["choices"][0]["message"]["content"],
+            messages=kwargs.get("messages", []),
+        )
+
+    def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]:
+        """
+        Synchronous streaming method.
+
+        Makes a streaming HTTP POST to Azure ML's OpenAI-compatible endpoint.
+        """
+        url, headers, json_data = self._prepare_request(**kwargs)
+        json_data["stream"] = True
+
+        client = self._get_client()
+
+        with client.stream("POST", url, headers=headers, json=json_data) as response:
+            response.raise_for_status()
+
+            for line in response.iter_lines():
+                if line.startswith("data: "):
+                    data = line[6:]  # Remove "data: " prefix
+                    if data == "[DONE]":
+                        break
+
+                    try:
+                        chunk_json = json.loads(data)
+                        delta = chunk_json["choices"][0].get("delta", {})
+                        content = delta.get("content", "")
+                        finish_reason = chunk_json["choices"][0].get("finish_reason")
+
+                        generic_streaming_chunk: GenericStreamingChunk = {
+                            "finish_reason": finish_reason,
+                            "index": 0,
+                            "is_finished": finish_reason is not None,
+                            "text": content,
+                            "tool_use": None,
+                            "usage": chunk_json.get(
+                                "usage",
+                                {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
+                            ),
+                        }
+
+                        yield generic_streaming_chunk
+                    except json.JSONDecodeError:
+                        continue
+
+    async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]:
+        """
+        Asynchronous streaming method.
+
+        Makes an async streaming HTTP POST to Azure ML's OpenAI-compatible endpoint.
+        """
+        url, headers, json_data = self._prepare_request(**kwargs)
+        json_data["stream"] = True
+
+        client = self._get_async_client()
+
+        async with client.stream("POST", url, headers=headers, json=json_data) as response:
+            response.raise_for_status()
+
+            async for line in response.aiter_lines():
+                if line.startswith("data: "):
+                    data = line[6:]  # Remove "data: " prefix
+                    if data == "[DONE]":
+                        break
+
+                    try:
+                        chunk_json = json.loads(data)
+                        delta = chunk_json["choices"][0].get("delta", {})
+                        content = delta.get("content", "")
+                        finish_reason = chunk_json["choices"][0].get("finish_reason")
+
+                        generic_streaming_chunk: GenericStreamingChunk = {
+                            "finish_reason": finish_reason,
+                            "index": 0,
+                            "is_finished": finish_reason is not None,
+                            "text": content,
+                            "tool_use": None,
+                            "usage": chunk_json.get(
+                                "usage",
+                                {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
+                            ),
+                        }
+
+                        yield generic_streaming_chunk
+                    except json.JSONDecodeError:
+                        continue
+
+    def __del__(self):
+        """Cleanup HTTP clients."""
+        if self._client is not None:
+            self._client.close()
+        if self._async_client is not None:
+            import asyncio
+
+            try:
+                loop = asyncio.get_event_loop()
+                if loop.is_running():
+                    loop.create_task(self._async_client.aclose())
+                else:
+                    loop.run_until_complete(self._async_client.aclose())
+            except Exception:
+                pass

agent/adapters/cua_adapter.py

@@ -0,0 +1,161 @@
+import os
+from typing import Any, AsyncIterator, Iterator
+
+from litellm import acompletion, completion
+from litellm.llms.custom_llm import CustomLLM
+from litellm.types.utils import GenericStreamingChunk, ModelResponse
+
+
+class CUAAdapter(CustomLLM):
+    def __init__(self, base_url: str | None = None, api_key: str | None = None, **_: Any):
+        super().__init__()
+        self.base_url = base_url or os.environ.get("CUA_BASE_URL") or "https://inference.cua.ai/v1"
+        self.api_key = (
+            api_key or os.environ.get("CUA_INFERENCE_API_KEY") or os.environ.get("CUA_API_KEY")
+        )
+
+    def _normalize_model(self, model: str) -> str:
+        # Accept either "cua/<model>" or raw "<model>"
+        return model.split("/", 1)[1] if model and model.startswith("cua/") else model
+
+    def completion(self, *args, **kwargs) -> ModelResponse:
+        model = kwargs.get("model", "")
+        api_base = kwargs.get("api_base") or self.base_url
+        if "anthropic/" in model:
+            model = f"anthropic/{self._normalize_model(model)}"
+            api_base = api_base.removesuffix("/v1")
+        elif "gemini/" in model or "google/" in model:
+            # Route to Gemini pass-through endpoint
+            model = f"gemini/{self._normalize_model(model)}"
+            api_base = api_base + "/gemini"
+        else:
+            model = f"openai/{self._normalize_model(model)}"
+
+        params = {
+            "model": model,
+            "messages": kwargs.get("messages", []),
+            "api_base": api_base,
+            "api_key": kwargs.get("api_key") or self.api_key,
+            "stream": False,
+        }
+
+        # Forward tools if provided
+        if "tools" in kwargs:
+            params["tools"] = kwargs["tools"]
+
+        if "optional_params" in kwargs:
+            params.update(kwargs["optional_params"])
+            del kwargs["optional_params"]
+
+        if "headers" in kwargs:
+            params["headers"] = kwargs["headers"]
+            del kwargs["headers"]
+
+        # Print dropped parameters
+        original_keys = set(kwargs.keys())
+        used_keys = set(params.keys())  # Only these are extracted from kwargs
+        ignored_keys = {
+            "litellm_params",
+            "client",
+            "print_verbose",
+            "acompletion",
+            "timeout",
+            "logging_obj",
+            "encoding",
+            "custom_prompt_dict",
+            "model_response",
+            "logger_fn",
+        }
+        dropped_keys = original_keys - used_keys - ignored_keys
+        if dropped_keys:
+            dropped_keyvals = {k: kwargs[k] for k in dropped_keys}
+            # print(f"CUAAdapter.completion: Dropped parameters: {dropped_keyvals}")
+
+        return completion(**params)  # type: ignore
+
+    async def acompletion(self, *args, **kwargs) -> ModelResponse:
+        model = kwargs.get("model", "")
+        api_base = kwargs.get("api_base") or self.base_url
+        if "anthropic/" in model:
+            model = f"anthropic/{self._normalize_model(model)}"
+            api_base = api_base.removesuffix("/v1")
+        elif "gemini/" in model or "google/" in model:
+            # Route to Gemini pass-through endpoint
+            model = f"gemini/{self._normalize_model(model)}"
+            api_base = api_base + "/gemini"
+        else:
+            model = f"openai/{self._normalize_model(model)}"
+
+        params = {
+            "model": model,
+            "messages": kwargs.get("messages", []),
+            "api_base": api_base,
+            "api_key": kwargs.get("api_key") or self.api_key,
+            "stream": False,
+        }
+
+        # Forward tools if provided
+        if "tools" in kwargs:
+            params["tools"] = kwargs["tools"]
+
+        if "optional_params" in kwargs:
+            params.update(kwargs["optional_params"])
+            del kwargs["optional_params"]
+
+        if "headers" in kwargs:
+            params["headers"] = kwargs["headers"]
+            del kwargs["headers"]
+
+        # Print dropped parameters
+        original_keys = set(kwargs.keys())
+        used_keys = set(params.keys())  # Only these are extracted from kwargs
+        ignored_keys = {
+            "litellm_params",
+            "client",
+            "print_verbose",
+            "acompletion",
+            "timeout",
+            "logging_obj",
+            "encoding",
+            "custom_prompt_dict",
+            "model_response",
+            "logger_fn",
+        }
+        dropped_keys = original_keys - used_keys - ignored_keys
+        if dropped_keys:
+            dropped_keyvals = {k: kwargs[k] for k in dropped_keys}
+            # print(f"CUAAdapter.acompletion: Dropped parameters: {dropped_keyvals}")
+
+        response = await acompletion(**params)  # type: ignore
+
+        return response
+
+    def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]:
+        params = dict(kwargs)
+        inner_model = self._normalize_model(params.get("model", ""))
+        params.update(
+            {
+                "model": f"openai/{inner_model}",
+                "api_base": self.base_url,
+                "api_key": self.api_key,
+                "stream": True,
+            }
+        )
+        # Yield chunks directly from LiteLLM's streaming generator
+        for chunk in completion(**params):  # type: ignore
+            yield chunk  # type: ignore
+
+    async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]:
+        params = dict(kwargs)
+        inner_model = self._normalize_model(params.get("model", ""))
+        params.update(
+            {
+                "model": f"openai/{inner_model}",
+                "api_base": self.base_url,
+                "api_key": self.api_key,
+                "stream": True,
+            }
+        )
+        stream = await acompletion(**params)  # type: ignore
+        async for chunk in stream:  # type: ignore
+            yield chunk  # type: ignore