From 4abae2f2628f404ec180fcd64299d4910ada4be3 Mon Sep 17 00:00:00 2001 From: Dillon DuPont Date: Thu, 23 Oct 2025 10:49:40 -0700 Subject: [PATCH 1/3] Add "cua/" LLM provider --- libs/python/agent/agent/adapters/__init__.py | 2 + .../agent/agent/adapters/cua_adapter.py | 73 +++++++++++++++++++ libs/python/agent/agent/agent.py | 8 +- 3 files changed, 78 insertions(+), 5 deletions(-) create mode 100644 libs/python/agent/agent/adapters/cua_adapter.py diff --git a/libs/python/agent/agent/adapters/__init__.py b/libs/python/agent/agent/adapters/__init__.py index 1f07a9fc..cded1d34 100644 --- a/libs/python/agent/agent/adapters/__init__.py +++ b/libs/python/agent/agent/adapters/__init__.py @@ -2,6 +2,7 @@ Adapters package for agent - Custom LLM adapters for LiteLLM """ +from .cua_adapter import CUAAdapter from .huggingfacelocal_adapter import HuggingFaceLocalAdapter from .human_adapter import HumanAdapter from .mlxvlm_adapter import MLXVLMAdapter @@ -10,4 +11,5 @@ __all__ = [ "HuggingFaceLocalAdapter", "HumanAdapter", "MLXVLMAdapter", + "CUAAdapter", ] diff --git a/libs/python/agent/agent/adapters/cua_adapter.py b/libs/python/agent/agent/adapters/cua_adapter.py new file mode 100644 index 00000000..e201ce54 --- /dev/null +++ b/libs/python/agent/agent/adapters/cua_adapter.py @@ -0,0 +1,73 @@ +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://gateway.cua.ai/v1" + self.api_key = api_key or os.environ.get("CUA_API_KEY") + + def _normalize_model(self, model: str) -> str: + # Accept either "cua/" or raw "" + return model.split("/", 1)[1] if model and model.startswith("cua/") else model + + def completion(self, *args, **kwargs) -> ModelResponse: + 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": False, + } + ) + return completion(**params) # type: ignore + + async def acompletion(self, *args, **kwargs) -> ModelResponse: + 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": False, + } + ) + return await acompletion(**params) # type: ignore + + 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 diff --git a/libs/python/agent/agent/agent.py b/libs/python/agent/agent/agent.py index 370af997..ca1aae37 100644 --- a/libs/python/agent/agent/agent.py +++ b/libs/python/agent/agent/agent.py @@ -23,11 +23,7 @@ import litellm import litellm.utils from litellm.responses.utils import Usage -from .adapters import ( - HuggingFaceLocalAdapter, - HumanAdapter, - MLXVLMAdapter, -) +from .adapters import CUAAdapter, HuggingFaceLocalAdapter, HumanAdapter, MLXVLMAdapter from .callbacks import ( BudgetManagerCallback, ImageRetentionCallback, @@ -272,10 +268,12 @@ class ComputerAgent: ) human_adapter = HumanAdapter() mlx_adapter = MLXVLMAdapter() + cua_adapter = CUAAdapter() litellm.custom_provider_map = [ {"provider": "huggingface-local", "custom_handler": hf_adapter}, {"provider": "human", "custom_handler": human_adapter}, {"provider": "mlx", "custom_handler": mlx_adapter}, + {"provider": "cua", "custom_handler": cua_adapter}, ] litellm.suppress_debug_info = True From cfedce280626191b935022f1c98d1ca893309fde Mon Sep 17 00:00:00 2001 From: Dillon DuPont Date: Thu, 23 Oct 2025 15:45:59 -0700 Subject: [PATCH 2/3] adjust url --- libs/python/agent/agent/adapters/cua_adapter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/python/agent/agent/adapters/cua_adapter.py b/libs/python/agent/agent/adapters/cua_adapter.py index e201ce54..b3f25bc3 100644 --- a/libs/python/agent/agent/adapters/cua_adapter.py +++ b/libs/python/agent/agent/adapters/cua_adapter.py @@ -9,7 +9,7 @@ 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://gateway.cua.ai/v1" + 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_API_KEY") def _normalize_model(self, model: str) -> str: From ac1511210b4b6b3a148338837a3720620a678935 Mon Sep 17 00:00:00 2001 From: Dillon DuPont Date: Tue, 11 Nov 2025 10:13:14 -0500 Subject: [PATCH 3/3] add a more specific env CUA_INFERENCE_API_KEY --- libs/python/agent/agent/adapters/cua_adapter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/python/agent/agent/adapters/cua_adapter.py b/libs/python/agent/agent/adapters/cua_adapter.py index b3f25bc3..76e13977 100644 --- a/libs/python/agent/agent/adapters/cua_adapter.py +++ b/libs/python/agent/agent/adapters/cua_adapter.py @@ -10,7 +10,7 @@ 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_API_KEY") + 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/" or raw ""