feat(agenthub): add agenthub llm

Author: ionmincu

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"

src/uipath/_services/llm_gateway_service.py

@@ -3,7 +3,7 @@
 
 from .._config import Config
 from .._execution_context import ExecutionContext
-from .._utils import Endpoint
+from .._utils import Endpoint, EndpointManager
 from ..models.llm_gateway import (
     ChatCompletion,
     SpecificToolChoice,
@@ -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,

src/uipath/_utils/__init__.py

@@ -1,4 +1,5 @@
 from ._endpoint import Endpoint
+from ._endpoints_manager import EndpointManager
 from ._infer_bindings import get_inferred_bindings_names, infer_bindings
 from ._logs import setup_logging
 from ._request_override import header_folder
@@ -8,6 +9,7 @@
 
 __all__ = [
     "Endpoint",
+    "EndpointManager",
     "setup_logging",
     "RequestSpec",
     "header_folder",

src/uipath/_utils/_endpoints_manager.py

@@ -0,0 +1,69 @@
+import os
+from enum import Enum
+from typing import Any, Dict, Optional
+
+import httpx
+
+
+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_capabilities_endpoint = os.getenv(
+        "UIPATH_AGENTHUB_CAPABILITIES_ENDPOINT", "/agenthub_/llm/api/capabilities"
+    )
+    _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:
+                capabilities_url = f"{cls.base_url.rstrip('/')}/{UiPathEndpoints.AH_CAPABILITIES_ENDPOINT.value}"
+                response = http_client.get(capabilities_url)
+                return response.status_code == 200
+        except Exception:
+            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

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."""