Merge pull request #115 from UiPath/feature/llm_agent_hub

agenthub: use llms from agenthub

Author: ionmincu

pyproject.toml

@@ -1,11 +1,11 @@
 [project]
 name = "uipath-langchain"
-version = "0.0.116"
+version = "0.0.117"
 description = "UiPath Langchain"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.10"
 dependencies = [
-    "uipath>=2.0.77, <2.1.0",
+    "uipath>=2.0.79, <2.1.0",
     "langgraph>=0.2.70",
     "langchain-core>=0.3.34",
     "langgraph-checkpoint-sqlite>=2.0.3",

src/uipath_langchain/_utils/_request_mixin.py

@@ -40,6 +40,8 @@ class Config:
 
     default_headers: Optional[Mapping[str, str]] = {
         "X-UiPath-Streaming-Enabled": "false",
+        "X-UiPath-JobKey": os.getenv("UIPATH_JOB_KEY", ""),
+        "X-UiPath-ProcessKey": os.getenv("UIPATH_PROCESS_KEY", ""),
     }
     model_name: Optional[str] = Field(
         default_factory=lambda: os.getenv("UIPATH_MODEL_NAME", "gpt-4o-2024-08-06"),

src/uipath_langchain/_utils/_settings.py

@@ -1,6 +1,5 @@
 # mypy: disable-error-code="syntax"
 import os
-from enum import Enum
 from typing import Any, Optional
 
 import httpx
@@ -47,14 +46,6 @@ class UiPathClientSettings(BaseSettings):
     action_id: str = Field(default="DefaultActionId", alias="UIPATH_ACTION_ID")
 
 
-class UiPathEndpoints(Enum):
-    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}"
-    )
-
-
 def get_uipath_token_header(
     settings: Any = None,
 ) -> str:

src/uipath_langchain/chat/models.py

@@ -1,4 +1,6 @@
 import json
+import logging
+from os import environ as env
 from typing import Any, Dict, List, Literal, Optional, Union
 
 from langchain_core.callbacks import (
@@ -12,9 +14,11 @@
 from langchain_core.runnables import Runnable
 from langchain_openai.chat_models import AzureChatOpenAI
 from pydantic import BaseModel
+from uipath.utils import EndpointManager
 
 from uipath_langchain._utils._request_mixin import UiPathRequestMixin
-from uipath_langchain._utils._settings import UiPathEndpoints
+
+logger = logging.getLogger(__name__)
 
 
 class UiPathAzureChatOpenAI(UiPathRequestMixin, AzureChatOpenAI):
@@ -71,14 +75,25 @@ def with_structured_output(
 
     @property
     def endpoint(self) -> str:
-        return UiPathEndpoints.PASSTHROUGH_COMPLETION_ENDPOINT.value.format(
+        endpoint = EndpointManager.get_passthrough_endpoint()
+        logger.debug("Using endpoint: %s", endpoint)
+        return endpoint.format(
             model=self.model_name, api_version=self.openai_api_version
         )
 
 
 class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
     """Custom LLM connector for LangChain integration with UiPath Normalized."""
 
+    def __init__(self, *args: Any, **kwargs: Any):
+        """Initialize the UiPath Azure Chat OpenAI model."""
+
+        super().__init__(*args, **kwargs)
+        self.default_headers = {
+            "X-UiPath-JobKey": env.get("UIPATH_JOB_KEY", ""),
+            "X-UiPath-ProcessKey": env.get("UIPATH_PROCESS_KEY", ""),
+        }
+
     def _create_chat_result(
         self,
         response: Union[Dict[str, Any], BaseModel],
@@ -252,7 +267,11 @@ def with_structured_output(
 
     @property
     def endpoint(self) -> str:
-        return UiPathEndpoints.NORMALIZED_COMPLETION_ENDPOINT.value
+        endpoint = EndpointManager.get_passthrough_endpoint()
+        logger.debug("Using endpoint: %s", endpoint)
+        return endpoint.format(
+            model=self.model_name, api_version=self.openai_api_version
+        )
 
     @property
     def is_normalized(self) -> bool:

src/uipath_langchain/embeddings/embeddings.py

@@ -5,9 +5,9 @@
 from langchain_community.callbacks.manager import openai_callback_var
 from langchain_openai.embeddings import AzureOpenAIEmbeddings, OpenAIEmbeddings
 from pydantic import Field
+from uipath.utils import EndpointManager
 
 from uipath_langchain._utils._request_mixin import UiPathRequestMixin
-from uipath_langchain._utils._settings import UiPathEndpoints
 
 
 class UiPathAzureOpenAIEmbeddings(UiPathRequestMixin, AzureOpenAIEmbeddings):
@@ -43,7 +43,8 @@ def __init__(self, **kwargs):
 
     @property
     def endpoint(self) -> str:
-        return UiPathEndpoints.EMBEDDING_ENDPOINT.value.format(
+        endpoint = EndpointManager.get_embeddings_endpoint()
+        return endpoint.format(
             model=self.model_name, api_version=self.openai_api_version
         )
 
@@ -105,6 +106,7 @@ async def aembed_documents(
 
     @property
     def endpoint(self) -> str:
-        return UiPathEndpoints.EMBEDDING_ENDPOINT.value.format(
+        endpoint = EndpointManager.get_embeddings_endpoint()
+        return endpoint.format(
             model=self.model_name, api_version=self.openai_api_version
         )