uipath-langchain 0.0.112__py3-none-any.whl → 0.1.24__py3-none-any.whl

uipath_langchain/chat/models.py

@@ -1,20 +1,27 @@
 import json
-from typing import Any, Dict, List, Literal, Optional, Union
+import logging
+from typing import Any, AsyncIterator, Iterator, Literal, Union
 
 from langchain_core.callbacks import (
     AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
 )
 from langchain_core.language_models import LanguageModelInput
-from langchain_core.messages import AIMessage, BaseMessage
+from langchain_core.language_models.chat_models import (
+    agenerate_from_stream,
+    generate_from_stream,
+)
+from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage
 from langchain_core.messages.ai import UsageMetadata
-from langchain_core.outputs import ChatGeneration, ChatResult
+from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
 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):
@@ -22,37 +29,118 @@ class UiPathAzureChatOpenAI(UiPathRequestMixin, AzureChatOpenAI):
 
     def _generate(
         self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
         if "tools" in kwargs and not kwargs["tools"]:
             del kwargs["tools"]
+
+        if self.streaming:
+            stream_iter = self._stream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            )
+            return generate_from_stream(stream_iter)
+
         payload = self._get_request_payload(messages, stop=stop, **kwargs)
         response = self._call(self.url, payload, self.auth_headers)
         return self._create_chat_result(response)
 
     async def _agenerate(
         self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
         if "tools" in kwargs and not kwargs["tools"]:
             del kwargs["tools"]
+
+        if self.streaming:
+            stream_iter = self._astream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            )
+            return await agenerate_from_stream(stream_iter)
+
         payload = self._get_request_payload(messages, stop=stop, **kwargs)
         response = await self._acall(self.url, payload, self.auth_headers)
         return self._create_chat_result(response)
 
+    def _stream(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> Iterator[ChatGenerationChunk]:
+        if "tools" in kwargs and not kwargs["tools"]:
+            del kwargs["tools"]
+        kwargs["stream"] = True
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+
+        default_chunk_class = AIMessageChunk
+
+        for chunk in self._stream_request(self.url, payload, self.auth_headers):
+            if self.logger:
+                self.logger.debug(f"[Stream] Got chunk from _stream_request: {chunk}")
+            generation_chunk = self._convert_chunk(
+                chunk, default_chunk_class, include_tool_calls=True
+            )
+            if generation_chunk is None:
+                if self.logger:
+                    self.logger.debug("[Stream] Skipping None generation_chunk")
+                continue
+
+            if self.logger:
+                self.logger.debug(
+                    f"[Stream] Yielding generation_chunk: {generation_chunk}"
+                )
+
+            if run_manager:
+                run_manager.on_llm_new_token(
+                    generation_chunk.text,
+                    chunk=generation_chunk,
+                )
+
+            yield generation_chunk
+
+    async def _astream(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[ChatGenerationChunk]:
+        if "tools" in kwargs and not kwargs["tools"]:
+            del kwargs["tools"]
+        kwargs["stream"] = True
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+
+        default_chunk_class = AIMessageChunk
+
+        async for chunk in self._astream_request(self.url, payload, self.auth_headers):
+            generation_chunk = self._convert_chunk(
+                chunk, default_chunk_class, include_tool_calls=True
+            )
+            if generation_chunk is None:
+                continue
+
+            if run_manager:
+                await run_manager.on_llm_new_token(
+                    generation_chunk.text,
+                    chunk=generation_chunk,
+                )
+
+            yield generation_chunk
+
     def with_structured_output(
         self,
-        schema: Optional[Any] = None,
+        schema: Any = None,
         *,
         method: Literal["function_calling", "json_mode", "json_schema"] = "json_schema",
         include_raw: bool = False,
-        strict: Optional[bool] = None,
+        strict: bool | None = None,
         **kwargs: Any,
     ) -> Runnable[LanguageModelInput, Any]:
         """Model wrapper that returns outputs formatted to match the given schema."""
@@ -71,7 +159,9 @@ class UiPathAzureChatOpenAI(UiPathRequestMixin, AzureChatOpenAI):
 
     @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
         )
 
@@ -81,8 +171,8 @@ class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
 
     def _create_chat_result(
         self,
-        response: Union[Dict[str, Any], BaseModel],
-        generation_info: Optional[Dict[Any, Any]] = None,
+        response: Union[dict[str, Any], BaseModel],
+        generation_info: dict[Any, Any] | None = None,
     ) -> ChatResult:
         if not isinstance(response, dict):
             response = response.model_dump()
@@ -112,6 +202,7 @@ class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
                     "id": tool["id"],
                     "name": tool["name"],
                     "args": tool["arguments"],
+                    "type": "tool_call",
                 }
                 for tool in message["tool_calls"]
             ]
@@ -122,9 +213,9 @@ class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
         self,
         input_: LanguageModelInput,
         *,
-        stop: Optional[List[str]] = None,
+        stop: list[str] | None = None,
         **kwargs: Any,
-    ) -> Dict[Any, Any]:
+    ) -> dict[Any, Any]:
         payload = super()._get_request_payload(input_, stop=stop, **kwargs)
         # hacks to make the request work with uipath normalized
         for message in payload["messages"]:
@@ -143,11 +234,40 @@ class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
                 }
         return payload
 
+    def _normalize_tool_choice(self, kwargs: dict[str, Any]) -> None:
+        """Normalize tool_choice for UiPath Gateway compatibility.
+
+        Converts LangChain tool_choice formats to UiPath Gateway format:
+        - String "required" -> {"type": "required"}
+        - String "auto" -> {"type": "auto"}
+        - Dict with function -> {"type": "tool", "name": "function_name"}
+        """
+        if "tool_choice" in kwargs:
+            tool_choice = kwargs["tool_choice"]
+
+            if isinstance(tool_choice, str):
+                if tool_choice in ("required", "auto", "none"):
+                    logger.debug(
+                        f"Converting tool_choice from '{tool_choice}' to {{'type': '{tool_choice}'}}"
+                    )
+                    kwargs["tool_choice"] = {"type": tool_choice}
+            elif (
+                isinstance(tool_choice, dict) and tool_choice.get("type") == "function"
+            ):
+                function_name = tool_choice["function"]["name"]
+                logger.debug(
+                    f"Converting tool_choice from function '{function_name}' to tool format"
+                )
+                kwargs["tool_choice"] = {
+                    "type": "tool",
+                    "name": function_name,
+                }
+
     def _generate(
         self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
         """Override the _generate method to implement the chat model logic.
@@ -167,21 +287,23 @@ class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
         """
         if kwargs.get("tools"):
             kwargs["tools"] = [tool["function"] for tool in kwargs["tools"]]
-        if "tool_choice" in kwargs and kwargs["tool_choice"]["type"] == "function":
-            kwargs["tool_choice"] = {
-                "type": "tool",
-                "name": kwargs["tool_choice"]["function"]["name"],
-            }
-        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+        self._normalize_tool_choice(kwargs)
+
+        if self.streaming:
+            stream_iter = self._stream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            )
+            return generate_from_stream(stream_iter)
 
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
         response = self._call(self.url, payload, self.auth_headers)
         return self._create_chat_result(response)
 
     async def _agenerate(
         self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
         """Override the _generate method to implement the chat model logic.
@@ -201,25 +323,122 @@ class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
         """
         if kwargs.get("tools"):
             kwargs["tools"] = [tool["function"] for tool in kwargs["tools"]]
-        if "tool_choice" in kwargs and kwargs["tool_choice"]["type"] == "function":
-            kwargs["tool_choice"] = {
-                "type": "tool",
-                "name": kwargs["tool_choice"]["function"]["name"],
-            }
-        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+        self._normalize_tool_choice(kwargs)
+
+        if self.streaming:
+            stream_iter = self._astream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            )
+            return await agenerate_from_stream(stream_iter)
 
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
         response = await self._acall(self.url, payload, self.auth_headers)
         return self._create_chat_result(response)
 
+    def _stream(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> Iterator[ChatGenerationChunk]:
+        """Stream the LLM on a given prompt.
+
+        Args:
+            messages: the prompt composed of a list of messages.
+            stop: a list of strings on which the model should stop generating.
+            run_manager: A run manager with callbacks for the LLM.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            An iterator of ChatGenerationChunk objects.
+        """
+        if kwargs.get("tools"):
+            kwargs["tools"] = [tool["function"] for tool in kwargs["tools"]]
+        self._normalize_tool_choice(kwargs)
+        kwargs["stream"] = True
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+
+        default_chunk_class = AIMessageChunk
+
+        for chunk in self._stream_request(self.url, payload, self.auth_headers):
+            if self.logger:
+                self.logger.debug(f"[Stream] Got chunk from _stream_request: {chunk}")
+            generation_chunk = self._convert_chunk(
+                chunk, default_chunk_class, include_tool_calls=True
+            )
+            if generation_chunk is None:
+                if self.logger:
+                    self.logger.debug("[Stream] Skipping None generation_chunk")
+                continue
+
+            if self.logger:
+                self.logger.debug(
+                    f"[Stream] Yielding generation_chunk: {generation_chunk}"
+                )
+
+            if run_manager:
+                run_manager.on_llm_new_token(
+                    generation_chunk.text,
+                    chunk=generation_chunk,
+                )
+
+            yield generation_chunk
+
+    async def _astream(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[ChatGenerationChunk]:
+        """Async stream the LLM on a given prompt.
+
+        Args:
+            messages: the prompt composed of a list of messages.
+            stop: a list of strings on which the model should stop generating.
+            run_manager: A run manager with callbacks for the LLM.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            An async iterator of ChatGenerationChunk objects.
+        """
+        if kwargs.get("tools"):
+            kwargs["tools"] = [tool["function"] for tool in kwargs["tools"]]
+        self._normalize_tool_choice(kwargs)
+        kwargs["stream"] = True
+        payload = self._get_request_payload(messages, stop=stop, **kwargs)
+
+        # Update headers to enable streaming
+        headers = {**self.auth_headers}
+        headers["X-UiPath-Streaming-Enabled"] = "true"
+
+        default_chunk_class = AIMessageChunk
+
+        async for chunk in self._astream_request(self.url, payload, headers):
+            generation_chunk = self._convert_chunk(
+                chunk, default_chunk_class, include_tool_calls=True
+            )
+            if generation_chunk is None:
+                continue
+
+            if run_manager:
+                await run_manager.on_llm_new_token(
+                    generation_chunk.text,
+                    chunk=generation_chunk,
+                )
+
+            yield generation_chunk
+
     def with_structured_output(
         self,
-        schema: Optional[Any] = None,
+        schema: Any = None,
         *,
         method: Literal[
             "function_calling", "json_mode", "json_schema"
         ] = "function_calling",
         include_raw: bool = False,
-        strict: Optional[bool] = None,
+        strict: bool | None = None,
         **kwargs: Any,
     ) -> Runnable[LanguageModelInput, Any]:
         """Model wrapper that returns outputs formatted to match the given schema."""
@@ -252,7 +471,11 @@ class UiPathChat(UiPathRequestMixin, AzureChatOpenAI):
 
     @property
     def endpoint(self) -> str:
-        return UiPathEndpoints.NORMALIZED_COMPLETION_ENDPOINT.value
+        endpoint = EndpointManager.get_normalized_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:

uipath_langchain/chat/openai.py

@@ -0,0 +1,132 @@
+import logging
+import os
+from typing import Optional
+
+import httpx
+from langchain_openai import AzureChatOpenAI
+from uipath.utils import EndpointManager
+
+from .supported_models import OpenAIModels
+
+logger = logging.getLogger(__name__)
+
+
+class UiPathURLRewriteTransport(httpx.AsyncHTTPTransport):
+    def __init__(self, verify: bool = True, **kwargs):
+        super().__init__(verify=verify, **kwargs)
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        original_url = str(request.url)
+
+        if "/openai/deployments/" in original_url:
+            base_url = original_url.split("/openai/deployments/")[0]
+            query_string = request.url.params
+            new_url_str = f"{base_url}/completions"
+            if query_string:
+                request.url = httpx.URL(new_url_str, params=query_string)
+            else:
+                request.url = httpx.URL(new_url_str)
+
+        return await super().handle_async_request(request)
+
+
+class UiPathSyncURLRewriteTransport(httpx.HTTPTransport):
+    def __init__(self, verify: bool = True, **kwargs):
+        super().__init__(verify=verify, **kwargs)
+
+    def handle_request(self, request: httpx.Request) -> httpx.Response:
+        original_url = str(request.url)
+
+        if "/openai/deployments/" in original_url:
+            base_url = original_url.split("/openai/deployments/")[0]
+            query_string = request.url.params
+            new_url_str = f"{base_url}/completions"
+            if query_string:
+                request.url = httpx.URL(new_url_str, params=query_string)
+            else:
+                request.url = httpx.URL(new_url_str)
+
+        return super().handle_request(request)
+
+
+class UiPathChatOpenAI(AzureChatOpenAI):
+    def __init__(
+        self,
+        token: Optional[str] = None,
+        model_name: str = OpenAIModels.gpt_5_mini_2025_08_07,
+        api_version: str = "2024-12-01-preview",
+        org_id: Optional[str] = None,
+        tenant_id: Optional[str] = None,
+        **kwargs,
+    ):
+        org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID")
+        tenant_id = tenant_id or os.getenv("UIPATH_TENANT_ID")
+        token = token or os.getenv("UIPATH_ACCESS_TOKEN")
+
+        if not org_id:
+            raise ValueError(
+                "UIPATH_ORGANIZATION_ID environment variable or org_id parameter is required"
+            )
+        if not tenant_id:
+            raise ValueError(
+                "UIPATH_TENANT_ID environment variable or tenant_id parameter is required"
+            )
+        if not token:
+            raise ValueError(
+                "UIPATH_ACCESS_TOKEN environment variable or token parameter is required"
+            )
+
+        self._openai_api_version = api_version
+        self._vendor = "openai"
+        self._model_name = model_name
+        self._url: Optional[str] = None
+
+        super().__init__(
+            azure_endpoint=self._build_base_url(),
+            model_name=model_name,
+            default_headers=self._build_headers(token),
+            http_async_client=httpx.AsyncClient(
+                transport=UiPathURLRewriteTransport(verify=True),
+                verify=True,
+            ),
+            http_client=httpx.Client(
+                transport=UiPathSyncURLRewriteTransport(verify=True),
+                verify=True,
+            ),
+            api_key=token,
+            api_version=api_version,
+            validate_base_url=False,
+            **kwargs,
+        )
+
+    def _build_headers(self, token: str) -> dict[str, str]:
+        headers = {
+            "X-UiPath-LlmGateway-ApiFlavor": "auto",
+            "Authorization": f"Bearer {token}",
+        }
+        if job_key := os.getenv("UIPATH_JOB_KEY"):
+            headers["X-UiPath-JobKey"] = job_key
+        if process_key := os.getenv("UIPATH_PROCESS_KEY"):
+            headers["X-UiPath-ProcessKey"] = process_key
+        return headers
+
+    @property
+    def endpoint(self) -> str:
+        vendor_endpoint = EndpointManager.get_vendor_endpoint()
+        formatted_endpoint = vendor_endpoint.format(
+            vendor=self._vendor,
+            model=self._model_name,
+            api_version=self._openai_api_version,
+        )
+        return formatted_endpoint.replace("/completions", "")
+
+    def _build_base_url(self) -> str:
+        if not self._url:
+            env_uipath_url = os.getenv("UIPATH_URL")
+
+            if env_uipath_url:
+                self._url = f"{env_uipath_url.rstrip('/')}/{self.endpoint}"
+            else:
+                raise ValueError("UIPATH_URL environment variable is required")
+
+        return self._url

uipath_langchain/chat/supported_models.py

@@ -0,0 +1,42 @@
+class OpenAIModels:
+    """Supported OpenAI model identifiers."""
+
+    # GPT-4o models
+    gpt_4o_2024_05_13 = "gpt-4o-2024-05-13"
+    gpt_4o_2024_08_06 = "gpt-4o-2024-08-06"
+    gpt_4o_2024_11_20 = "gpt-4o-2024-11-20"
+    gpt_4o_mini_2024_07_18 = "gpt-4o-mini-2024-07-18"
+
+    # GPT-4.1 models
+    gpt_4_1_2025_04_14 = "gpt-4.1-2025-04-14"
+    gpt_4_1_mini_2025_04_14 = "gpt-4.1-mini-2025-04-14"
+    gpt_4_1_nano_2025_04_14 = "gpt-4.1-nano-2025-04-14"
+
+    # GPT-5 models
+    gpt_5_2025_08_07 = "gpt-5-2025-08-07"
+    gpt_5_chat_2025_08_07 = "gpt-5-chat-2025-08-07"
+    gpt_5_mini_2025_08_07 = "gpt-5-mini-2025-08-07"
+    gpt_5_nano_2025_08_07 = "gpt-5-nano-2025-08-07"
+
+    # GPT-5.1 models
+    gpt_5_1_2025_11_13 = "gpt-5.1-2025-11-13"
+
+
+class GeminiModels:
+    """Supported Google Gemini model identifiers."""
+
+    gemini_2_5_pro = "gemini-2.5-pro"
+    gemini_2_5_flash = "gemini-2.5-flash"
+    gemini_2_0_flash_001 = "gemini-2.0-flash-001"
+
+
+class BedrockModels:
+    """Supported AWS Bedrock model identifiers."""
+
+    # Claude 3.7 models
+    anthropic_claude_3_7_sonnet = "anthropic.claude-3-7-sonnet-20250219-v1:0"
+
+    # Claude 4 models
+    anthropic_claude_sonnet_4 = "anthropic.claude-sonnet-4-20250514-v1:0"
+    anthropic_claude_sonnet_4_5 = "anthropic.claude-sonnet-4-5-20250929-v1:0"
+    anthropic_claude_haiku_4_5 = "anthropic.claude-haiku-4-5-20251001-v1:0"