feat(agenthub): add agenthub llm

ionmincu · ionmincu · commit 5f6f59261f6c · 2025-07-07T12:17:33.000+03:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "uipath"
-version = "2.0.73"
+version = "2.0.74"
 description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.10"
diff --git a/src/uipath/_services/llm_gateway_service.py b/src/uipath/_services/llm_gateway_service.py
@@ -10,9 +10,9 @@
     TextEmbedding,
     ToolChoice,
     ToolDefinition,
-    UsageInfo,
 )
 from ..tracing._traced import traced
+from ..utils import EndpointManager
 from ._base_service import BaseService
 
 # Common constants
@@ -54,36 +54,12 @@ class UiPathOpenAIService(BaseService):
     def __init__(self, config: Config, execution_context: ExecutionContext) -> None:
         super().__init__(config=config, execution_context=execution_context)
 
-    @traced(name="llm_embeddings_usage", run_type="uipath")
-    async def embeddings_usage(
-        self, input: str, embedding_model: str = EmbeddingModels.text_embedding_ada_002
-    ):
-        """Embedd the input text using llm gateway service.
-
-        Args:
-            input (str): The input text to embedd.
-            embedding_model (str, optional): The embedding model to use. Defaults to text-embedding-ada-002.
-
-        Returns:
-            EmbeddingUsageInfo: The embedding usage information.
-        """
-        endpoint = Endpoint(
-            f"/llmgateway_/openai/deployments/{embedding_model}/embeddings/usage"
-        )
-
-        response = await self.request_async(
-            "POST",
-            endpoint,
-            content=json.dumps({"input": input}),
-            params={"api-version": API_VERSION},
-            headers=DEFAULT_LLM_HEADERS,
-        )
-
-        return UsageInfo.model_validate(response.json())
-
     @traced(name="llm_embeddings", run_type="uipath")
     async def embeddings(
-        self, input: str, embedding_model: str = EmbeddingModels.text_embedding_ada_002
+        self,
+        input: str,
+        embedding_model: str = EmbeddingModels.text_embedding_ada_002,
+        openai_api_version: str = API_VERSION,
     ):
         """Embed the input text using llm gateway service.
 
@@ -93,9 +69,10 @@ async def embeddings(
         Returns:
             TextEmbedding: The embedding response.
         """
-        endpoint = Endpoint(
-            f"/llmgateway_/openai/deployments/{embedding_model}/embeddings"
+        endpoint = EndpointManager.get_embeddings_endpoint().format(
+            model=embedding_model, api_version=openai_api_version
         )
+        endpoint = Endpoint("/" + endpoint)
 
         response = await self.request_async(
             "POST",
@@ -114,6 +91,7 @@ async def chat_completions(
         model: str = ChatModels.gpt_4o_mini_2024_07_18,
         max_tokens: int = 50,
         temperature: float = 0,
+        api_version: str = API_VERSION,
     ):
         """Get chat completions using llm gateway service.
 
@@ -139,59 +117,10 @@ async def chat_completions(
         Returns:
             ChatCompletion: The chat completion response.
         """
-        endpoint = Endpoint(f"/llmgateway_/openai/deployments/{model}/chat/completions")
-
-        request_body = {
-            "messages": messages,
-            "max_tokens": max_tokens,
-            "temperature": temperature,
-        }
-
-        response = await self.request_async(
-            "POST",
-            endpoint,
-            content=json.dumps(request_body),
-            params={"api-version": API_VERSION},
-            headers=DEFAULT_LLM_HEADERS,
-        )
-
-        return ChatCompletion.model_validate(response.json())
-
-    @traced(name="llm_chat_completions_usage", run_type="uipath")
-    async def chat_completions_usage(
-        self,
-        messages: List[Dict[str, str]],
-        model: str = ChatModels.gpt_4o_mini_2024_07_18,
-        max_tokens: int = 50,
-        temperature: float = 0,
-    ):
-        """Get chat completions usage using llm gateway service.
-
-        Args:
-            messages (List[Dict[str, str]]): List of message dictionaries with 'role' and 'content' keys.
-                The supported roles are 'system', 'user', and 'assistant'.
-
-        Example:
-                ```
-                [
-                    {"role": "system", "content": "You are a helpful Python programming assistant."},
-                    {"role": "user", "content": "How do I read a file in Python?"},
-                    {"role": "assistant", "content": "You can use the built-in open() function."},
-                    {"role": "user", "content": "Can you show an example?"}
-                ]
-                ```
-                The conversation history can be included to provide context to the model.
-            model (str, optional): The model to use for chat completion. Defaults to ChatModels.gpt_4o_mini_2024_07_18.
-            max_tokens (int, optional): Maximum number of tokens to generate. Defaults to 50.
-            temperature (float, optional): Temperature for sampling, between 0 and 1.
-                Lower values make output more deterministic. Defaults to 0.
-
-        Returns:
-            ChatCompletion: The chat completion usage response.
-        """
-        endpoint = Endpoint(
-            f"/llmgateway_/openai/deployments/{model}/chat/completions/usage"
+        endpoint = EndpointManager.get_passthrough_endpoint().format(
+            model=model, api_version=api_version
         )
+        endpoint = Endpoint("/" + endpoint)
 
         request_body = {
             "messages": messages,
@@ -207,7 +136,7 @@ async def chat_completions_usage(
             headers=DEFAULT_LLM_HEADERS,
         )
 
-        return UsageInfo.model_validate(response.json())
+        return ChatCompletion.model_validate(response.json())
 
 
 class UiPathLlmChatService(BaseService):
@@ -229,6 +158,7 @@ async def chat_completions(
         top_p: float = 1,
         tools: Optional[List[ToolDefinition]] = None,
         tool_choice: Optional[ToolChoice] = None,
+        api_version: str = NORMALIZED_API_VERSION,
     ):
         """Get chat completions using UiPath's normalized LLM Gateway API.
 
@@ -250,7 +180,10 @@ async def chat_completions(
         Returns:
             ChatCompletion: The chat completion response.
         """
-        endpoint = Endpoint("/llmgateway_/api/chat/completions")
+        endpoint = EndpointManager.get_normalized_endpoint().format(
+            model=model, api_version=api_version
+        )
+        endpoint = Endpoint("/" + endpoint)
 
         request_body = {
             "messages": messages,
diff --git a/src/uipath/utils/__init__.py b/src/uipath/utils/__init__.py
@@ -0,0 +1,5 @@
+from ._endpoints_manager import EndpointManager
+
+__all__ = [
+    "EndpointManager",
+]
diff --git a/src/uipath/utils/_endpoints_manager.py b/src/uipath/utils/_endpoints_manager.py
@@ -0,0 +1,72 @@
+import logging
+import os
+from enum import Enum
+from typing import Optional
+
+import httpx
+
+loggger = logging.getLogger(__name__)
+
+
+class UiPathEndpoints(Enum):
+    AH_NORMALIZED_COMPLETION_ENDPOINT = "agenthub_/llm/api/chat/completions"
+    AH_PASSTHROUGH_COMPLETION_ENDPOINT = "agenthub_/llm/openai/deployments/{model}/chat/completions?api-version={api_version}"
+    AH_EMBEDDING_ENDPOINT = (
+        "agenthub_/llm/openai/deployments/{model}/embeddings?api-version={api_version}"
+    )
+    AH_CAPABILITIES_ENDPOINT = "agenthub_/llm/api/capabilities"
+
+    NORMALIZED_COMPLETION_ENDPOINT = "llmgateway_/api/chat/completions"
+    PASSTHROUGH_COMPLETION_ENDPOINT = "llmgateway_/openai/deployments/{model}/chat/completions?api-version={api_version}"
+    EMBEDDING_ENDPOINT = (
+        "llmgateway_/openai/deployments/{model}/embeddings?api-version={api_version}"
+    )
+
+
+class EndpointManager:
+    """Manages and caches the UiPath endpoints."""
+
+    _base_url = os.getenv("UIPATH_URL", "")
+    _agenthub_available: Optional[bool] = None
+
+    @classmethod
+    def is_agenthub_available(cls) -> bool:
+        """Check if AgentHub is available and cache the result."""
+        if cls._agenthub_available is None:
+            cls._agenthub_available = cls._check_agenthub()
+        return cls._agenthub_available
+
+    @classmethod
+    def _check_agenthub(cls) -> bool:
+        """Perform the actual check for AgentHub capabilities."""
+        try:
+            with httpx.Client() as http_client:
+                base_url = os.getenv("UIPATH_URL", "")
+                capabilities_url = f"{base_url.rstrip('/')}/{UiPathEndpoints.AH_CAPABILITIES_ENDPOINT.value}"
+                loggger.debug(f"Checking AgentHub capabilities at {capabilities_url}")
+                response = http_client.get(capabilities_url)
+                return response.status_code == 200
+        except Exception as e:
+            loggger.error(f"Error checking AgentHub capabilities: {e}", exc_info=True)
+            return False
+
+    @classmethod
+    def get_passthrough_endpoint(cls) -> str:
+        if cls.is_agenthub_available():
+            return UiPathEndpoints.AH_PASSTHROUGH_COMPLETION_ENDPOINT.value
+
+        return UiPathEndpoints.PASSTHROUGH_COMPLETION_ENDPOINT.value
+
+    @classmethod
+    def get_normalized_endpoint(cls) -> str:
+        if cls.is_agenthub_available():
+            return UiPathEndpoints.AH_NORMALIZED_COMPLETION_ENDPOINT.value
+
+        return UiPathEndpoints.NORMALIZED_COMPLETION_ENDPOINT.value
+
+    @classmethod
+    def get_embeddings_endpoint(cls) -> str:
+        if cls.is_agenthub_available():
+            return UiPathEndpoints.AH_EMBEDDING_ENDPOINT.value
+
+        return UiPathEndpoints.EMBEDDING_ENDPOINT.value
diff --git a/tests/sdk/services/test_llm_integration.py b/tests/sdk/services/test_llm_integration.py
@@ -74,20 +74,6 @@ async def test_embeddings_real(self, llm_service):
         assert hasattr(result, "usage")
         assert result.usage.prompt_tokens > 0
 
-    @pytest.mark.asyncio
-    async def test_embeddings_usage_real(self, llm_service):
-        """Test the embeddings_usage function with a real API call."""
-        input_text = "Testing the embedding usage endpoint."
-
-        # Make the actual API call
-        result = await llm_service.embeddings_usage(input=input_text)
-
-        # Validate the response
-        assert result is not None
-        assert hasattr(result, "encoding")
-        assert hasattr(result, "prompt_tokens")
-        assert result.prompt_tokens > 0
-
     @pytest.mark.asyncio
     async def test_chat_completions_real(self, llm_service):
         """Test the chat_completions function with a real API call."""
@@ -115,31 +101,6 @@ async def test_chat_completions_real(self, llm_service):
         assert hasattr(result, "usage")
         assert result.usage.prompt_tokens > 0
 
-    @pytest.mark.asyncio
-    async def test_chat_completions_usage_real(self, llm_service):
-        """Test the chat_completions_usage function with a real API call."""
-        messages = [
-            {"role": "system", "content": "You are a helpful assistant."},
-            {"role": "user", "content": "What is the capital of France?"},
-        ]
-
-        # Make the actual API call
-        result = await llm_service.chat_completions_usage(
-            messages=messages,
-            model=ChatModels.gpt_4o_mini_2024_07_18,
-            max_tokens=50,
-            temperature=0.7,
-        )
-
-        # Validate the response
-        assert result is not None
-        assert hasattr(result, "encoding")
-        assert hasattr(result, "prompt_tokens")
-        assert result.prompt_tokens > 0
-        assert isinstance(result.prompt_tokens, int)
-        assert isinstance(result.encoding, str)
-        assert len(result.encoding) > 0
-
     @pytest.mark.asyncio
     async def test_embeddings_with_custom_model_real(self, llm_service):
         """Test the embeddings function with a custom model."""
diff --git a/tests/sdk/services/test_llm_service.py b/tests/sdk/services/test_llm_service.py
@@ -35,26 +35,6 @@ def test_init(self, config, execution_context):
         assert service._config == config
         assert service._execution_context == execution_context
 
-    @patch.object(UiPathOpenAIService, "request_async")
-    @pytest.mark.asyncio
-    async def test_embeddings_usage(self, mock_request, openai_service):
-        # Mock response
-        mock_response = MagicMock()
-        mock_response.json.return_value = {
-            "encoding": "cl100k_base",
-            "prompt_tokens": 4,
-        }
-        mock_request.return_value = mock_response
-
-        # Call the method
-        result = await openai_service.embeddings_usage(input="Test input")
-
-        # Assertions
-        mock_request.assert_called_once()
-        assert isinstance(result, UsageInfo)
-        assert result.encoding == "cl100k_base"
-        assert result.prompt_tokens == 4
-
     @patch.object(UiPathOpenAIService, "request_async")
     @pytest.mark.asyncio
     async def test_embeddings(self, mock_request, openai_service):
@@ -78,40 +58,6 @@ async def test_embeddings(self, mock_request, openai_service):
         assert result.model == "text-embedding-ada-002"
         assert result.usage.prompt_tokens == 4
 
-    @patch.object(UiPathOpenAIService, "request_async")
-    @pytest.mark.asyncio
-    async def test_chat_completions_usage(self, mock_request, openai_service):
-        # Mock response
-        mock_response = MagicMock()
-        mock_response.json.return_value = {
-            "encoding": "cl100k_base",
-            "prompt_tokens": 10,
-        }
-        mock_request.return_value = mock_response
-
-        # Test messages
-        messages = [
-            {"role": "system", "content": "You are a helpful assistant"},
-            {"role": "user", "content": "Hello"},
-        ]
-
-        # Call the method
-        result = await openai_service.chat_completions_usage(
-            messages=messages, max_tokens=50, temperature=0.5
-        )
-
-        # Assertions
-        mock_request.assert_called_once()
-        assert isinstance(result, UsageInfo)
-        assert result.encoding == "cl100k_base"
-        assert result.prompt_tokens == 10
-
-        # Verify the correct endpoint was called
-        args, kwargs = mock_request.call_args
-        assert json.loads(kwargs["content"])["messages"] == messages
-        assert json.loads(kwargs["content"])["max_tokens"] == 50
-        assert json.loads(kwargs["content"])["temperature"] == 0.5
-
     @patch.object(UiPathOpenAIService, "request_async")
     @pytest.mark.asyncio
     async def test_embeddings_with_custom_model(self, mock_request, openai_service):
diff --git a/uv.lock b/uv.lock
