uipath-langchain 0.0.133__py3-none-any.whl → 0.1.28__py3-none-any.whl

uipath_langchain/_utils/_request_mixin.py

@@ -3,13 +3,16 @@ import json
 import logging
 import os
 import time
-from typing import Any, Dict, List, Mapping, Optional
+from typing import Any, AsyncIterator, Dict, Iterator, Mapping
 
 import httpx
 import openai
 from langchain_core.embeddings import Embeddings
 from langchain_core.language_models.chat_models import _cleanup_llm_representation
-from pydantic import BaseModel, ConfigDict, Field, SecretStr
+from langchain_core.messages import AIMessageChunk
+from langchain_core.messages.ai import UsageMetadata
+from langchain_core.outputs import ChatGenerationChunk
+from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
 from tenacity import (
     AsyncRetrying,
     Retrying,
@@ -17,16 +20,23 @@ from tenacity import (
     stop_after_attempt,
     wait_exponential_jitter,
 )
-from uipath._cli._runtime._contracts import UiPathErrorCategory
 from uipath._utils._ssl_context import get_httpx_client_kwargs
+from uipath.runtime.errors import (
+    UiPathErrorCategory,
+    UiPathErrorCode,
+    UiPathRuntimeError,
+)
 
-from uipath_langchain._cli._runtime._exception import LangGraphRuntimeError
 from uipath_langchain._utils._settings import (
     UiPathClientFactorySettings,
     UiPathClientSettings,
     get_uipath_token_header,
 )
 from uipath_langchain._utils._sleep_policy import before_sleep_log
+from uipath_langchain.runtime.errors import (
+    LangGraphErrorCode,
+    LangGraphRuntimeError,
+)
 
 
 def get_from_uipath_url():
@@ -36,45 +46,60 @@ def get_from_uipath_url():
     return None
 
 
+def _get_access_token(data):
+    """Get access token from settings, environment variables, or UiPath client factory."""
+    token = (
+        getattr(data["settings"], "access_token", None)
+        or os.getenv("UIPATH_ACCESS_TOKEN")
+        or os.getenv("UIPATH_SERVICE_TOKEN")
+    )
+
+    if token:
+        return token
+
+    try:
+        settings = UiPathClientFactorySettings(
+            UIPATH_BASE_URL=data["base_url"],
+            UIPATH_CLIENT_ID=data["client_id"],
+            UIPATH_CLIENT_SECRET=data["client_secret"],
+        )
+        return get_uipath_token_header(settings)
+    except ValidationError:
+        raise UiPathRuntimeError(
+            UiPathErrorCode.EXECUTION_ERROR,
+            title="Authorization required",
+            detail="Authorization required. Please run uipath auth",
+            category=UiPathErrorCategory.USER,
+        ) from None
+
+
 class UiPathRequestMixin(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    default_headers: Optional[Mapping[str, str]] = {
+    default_headers: Mapping[str, str] | None = {
         "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(
+    model_name: str | None = Field(
         default_factory=lambda: os.getenv("UIPATH_MODEL_NAME", "gpt-4o-2024-08-06"),
         alias="model",
     )
-    settings: Optional[UiPathClientSettings] = None
-    client_id: Optional[str] = Field(
-        default_factory=lambda: os.getenv("UIPATH_CLIENT_ID")
-    )
-    client_secret: Optional[str] = Field(
+    settings: UiPathClientSettings | None = None
+    client_id: str | None = Field(default_factory=lambda: os.getenv("UIPATH_CLIENT_ID"))
+    client_secret: str | None = Field(
         default_factory=lambda: os.getenv("UIPATH_CLIENT_SECRET")
     )
-    base_url: Optional[str] = Field(
+    base_url: str | None = Field(
         default_factory=lambda data: getattr(data["settings"], "base_url", None)
         or os.getenv("UIPATH_BASE_URL")
         or get_from_uipath_url(),
         alias="azure_endpoint",
     )
-    access_token: Optional[str] = Field(
-        default_factory=lambda data: (
-            getattr(data["settings"], "access_token", None)
-            or os.getenv("UIPATH_ACCESS_TOKEN")  # Environment variable
-            or os.getenv("UIPATH_SERVICE_TOKEN")  # Environment variable
-            or get_uipath_token_header(
-                UiPathClientFactorySettings(
-                    UIPATH_BASE_URL=data["base_url"],
-                    UIPATH_CLIENT_ID=data["client_id"],
-                    UIPATH_CLIENT_SECRET=data["client_secret"],
-                )
-            )  # Get service token from UiPath
-        )
+    access_token: str | None = Field(
+        default_factory=lambda data: _get_access_token(data)
     )
+
     org_id: Any = Field(
         default_factory=lambda data: getattr(data["settings"], "org_id", None)
         or os.getenv("UIPATH_ORGANIZATION_ID", "")
@@ -103,39 +128,39 @@ class UiPathRequestMixin(BaseModel):
         alias="timeout",
     )
 
-    openai_api_version: Optional[str] = Field(
+    openai_api_version: str | None = Field(
         default_factory=lambda: os.getenv("OPENAI_API_VERSION", "2024-08-01-preview"),
         alias="api_version",
     )
     include_account_id: bool = False
-    temperature: Optional[float] = 0.0
-    max_tokens: Optional[int] = 1000
-    frequency_penalty: Optional[float] = None
-    presence_penalty: Optional[float] = None
+    temperature: float | None = 0.0
+    max_tokens: int | None = 1000
+    frequency_penalty: float | None = None
+    presence_penalty: float | None = None
 
-    logger: Optional[logging.Logger] = None
-    max_retries: Optional[int] = 5
+    logger: logging.Logger | None = None
+    max_retries: int | None = 5
     base_delay: float = 5.0
     max_delay: float = 60.0
 
-    _url: Optional[str] = None
-    _auth_headers: Optional[Dict[str, str]] = None
+    _url: str | None = None
+    _auth_headers: dict[str, str] | None = None
 
     # required to instantiate AzureChatOpenAI subclasses
-    azure_endpoint: Optional[str] = Field(
+    azure_endpoint: str | None = Field(
         default="placeholder", description="Bypassed Azure endpoint"
     )
-    openai_api_key: Optional[SecretStr] = Field(
+    openai_api_key: SecretStr | None = Field(
         default=SecretStr("placeholder"), description="Bypassed API key"
     )
     # required to instatiate ChatAnthropic subclasses (will be needed when passthrough is implemented for Anthropic models)
-    stop_sequences: Optional[List[str]] = Field(
+    stop_sequences: list[str] | None = Field(
         default=None, description="Bypassed stop sequence"
     )
 
     def _request(
-        self, url: str, request_body: Dict[str, Any], headers: Dict[str, str]
-    ) -> Dict[str, Any]:
+        self, url: str, request_body: dict[str, Any], headers: dict[str, str]
+    ) -> dict[str, Any]:
         """Run an asynchronous call to the LLM."""
         # if self.logger:
         #     self.logger.info(f"Completion request: {request_body['messages'][:2]}")
@@ -156,6 +181,8 @@ class UiPathRequestMixin(BaseModel):
 
             # Handle HTTP errors and map them to OpenAI exceptions
             try:
+                content = response.content  # Read content to avoid closed stream issues
+                print(f"Response content: {content.decode('utf-8')}")
                 response.raise_for_status()
             except httpx.HTTPStatusError as err:
                 if self.logger:
@@ -175,8 +202,8 @@ class UiPathRequestMixin(BaseModel):
             return response.json()
 
     def _call(
-        self, url: str, request_body: Dict[str, Any], headers: Dict[str, str]
-    ) -> Dict[str, Any]:
+        self, url: str, request_body: dict[str, Any], headers: dict[str, str]
+    ) -> dict[str, Any]:
         """Run a synchronous call with retries to LLM"""
         if self.max_retries is None:
             return self._request(url, request_body, headers)
@@ -199,6 +226,7 @@ class UiPathRequestMixin(BaseModel):
 
         try:
             return retryer(self._request, url, request_body, headers)
+            # return self._request(url, request_body, headers)
         except openai.APIStatusError as err:
             if self.logger:
                 self.logger.error(
@@ -214,8 +242,8 @@ class UiPathRequestMixin(BaseModel):
             raise err
 
     async def _arequest(
-        self, url: str, request_body: Dict[str, Any], headers: Dict[str, str]
-    ) -> Dict[str, Any]:
+        self, url: str, request_body: dict[str, Any], headers: dict[str, str]
+    ) -> dict[str, Any]:
         # if self.logger:
         #     self.logger.info(f"Completion request: {request_body['messages'][:2]}")
         client_kwargs = get_httpx_client_kwargs()
@@ -253,8 +281,8 @@ class UiPathRequestMixin(BaseModel):
             return response.json()
 
     async def _acall(
-        self, url: str, request_body: Dict[str, Any], headers: Dict[str, str]
-    ) -> Dict[str, Any]:
+        self, url: str, request_body: dict[str, Any], headers: dict[str, str]
+    ) -> dict[str, Any]:
         """Run an asynchronous call with retries to the LLM."""
         if self.max_retries is None:
             return await self._arequest(url, request_body, headers)
@@ -310,6 +338,256 @@ class UiPathRequestMixin(BaseModel):
                 )
             raise err
 
+    def _convert_chunk(
+        self,
+        chunk: Dict[str, Any],
+        default_chunk_class: type,
+        include_tool_calls: bool = False,
+    ) -> ChatGenerationChunk | None:
+        """Convert a streaming chunk to a ChatGenerationChunk.
+
+        Args:
+            chunk: The raw SSE chunk dictionary
+            default_chunk_class: The default message chunk class to use
+            include_tool_calls: Whether to parse and include tool call chunks
+
+        Returns:
+            A ChatGenerationChunk or None if the chunk should be skipped
+        """
+
+        token_usage = chunk.get("usage")
+        choices = chunk.get("choices", [])
+
+        usage_metadata: UsageMetadata | None = None
+        if token_usage:
+            usage_metadata = UsageMetadata(
+                input_tokens=token_usage.get("prompt_tokens", 0),
+                output_tokens=token_usage.get("completion_tokens", 0),
+                total_tokens=token_usage.get("total_tokens", 0),
+            )
+
+        if len(choices) == 0:
+            return ChatGenerationChunk(
+                message=default_chunk_class(content="", usage_metadata=usage_metadata),
+                generation_info={},
+            )
+
+        choice = choices[0]
+        delta = choice.get("delta")
+        if delta is None:
+            return None
+
+        # Extract content from delta
+        content = delta.get("content", "")
+
+        # Build the message chunk
+        message_kwargs = {
+            "content": content or "",
+            "usage_metadata": usage_metadata,
+        }
+
+        # Handle tool calls if requested (for normalized API)
+        if include_tool_calls:
+            tool_calls = delta.get("tool_calls", [])
+            tool_call_chunks = []
+            if tool_calls:
+                for tc in tool_calls:
+                    # Tool call structure: {'function': {'name': '...', 'arguments': '...'}, 'id': '...', 'index': 0}
+                    function = tc.get("function", {})
+                    tool_call_chunks.append(
+                        {
+                            "id": tc.get("id"),
+                            "name": function.get("name"),
+                            "args": function.get("arguments", ""),
+                            "index": tc.get("index", 0),
+                        }
+                    )
+            if tool_call_chunks:
+                message_kwargs["tool_call_chunks"] = tool_call_chunks
+
+        message_chunk = AIMessageChunk(**message_kwargs)
+
+        generation_info = {}
+        if finish_reason := choice.get("finish_reason"):
+            generation_info["finish_reason"] = finish_reason
+            if model_name := chunk.get("model"):
+                generation_info["model_name"] = model_name
+
+        return ChatGenerationChunk(
+            message=message_chunk,
+            generation_info=generation_info or None,
+        )
+
+    def _stream_request(
+        self, url: str, request_body: Dict[str, Any], headers: Dict[str, str]
+    ) -> Iterator[Dict[str, Any]]:
+        """Stream SSE responses from the LLM."""
+        client_kwargs = get_httpx_client_kwargs()
+        with httpx.Client(
+            **client_kwargs,
+            event_hooks={
+                "request": [self._log_request_duration],
+                "response": [self._log_response_duration],
+            },
+        ) as client:
+            with client.stream(
+                "POST",
+                url,
+                headers=headers,
+                json=request_body,
+                timeout=self.default_request_timeout,
+            ) as response:
+                try:
+                    response.raise_for_status()
+                except httpx.HTTPStatusError as err:
+                    if self.logger:
+                        self.logger.error(
+                            "Error querying UiPath: %s (%s)",
+                            err.response.reason_phrase,
+                            err.response.status_code,
+                            extra={
+                                "ActionName": self.settings.action_name,
+                                "ActionId": self.settings.action_id,
+                            }
+                            if self.settings
+                            else None,
+                        )
+                    # Read the response body for streaming responses
+                    err.response.read()
+                    raise self._make_status_error_from_response(err.response) from err
+
+                for line in response.iter_lines():
+                    line = line.strip()
+                    if not line:
+                        continue
+
+                    if self.logger:
+                        self.logger.debug(f"[SSE] Raw line: {line}")
+
+                    if line.startswith("data:"):
+                        data = line[
+                            5:
+                        ].strip()  # Remove "data:" prefix and strip whitespace
+                        if data == "[DONE]":
+                            break
+                        if not data:  # Skip empty data lines
+                            continue
+                        try:
+                            parsed = json.loads(data)
+                            # Skip empty chunks (some APIs send them as keepalive)
+                            # Check for truly empty: empty id AND (no choices or empty choices list)
+                            if (not parsed.get("id") or parsed.get("id") == "") and (
+                                not parsed.get("choices")
+                                or len(parsed.get("choices", [])) == 0
+                            ):
+                                if self.logger:
+                                    self.logger.debug(
+                                        "[SSE] Skipping empty keepalive chunk"
+                                    )
+                                continue
+                            yield parsed
+                        except json.JSONDecodeError as e:
+                            if self.logger:
+                                self.logger.warning(
+                                    f"Failed to parse SSE chunk: {data}, error: {e}"
+                                )
+                            continue
+                    else:
+                        # Handle lines without "data: " prefix (some APIs send raw JSON)
+                        try:
+                            parsed = json.loads(line)
+                            if self.logger:
+                                self.logger.debug(f"[SSE] Parsed raw JSON: {parsed}")
+                            yield parsed
+                        except json.JSONDecodeError:
+                            # Not JSON, skip
+                            pass
+
+    async def _astream_request(
+        self, url: str, request_body: Dict[str, Any], headers: Dict[str, str]
+    ) -> AsyncIterator[Dict[str, Any]]:
+        """Async stream SSE responses from the LLM."""
+        client_kwargs = get_httpx_client_kwargs()
+        async with httpx.AsyncClient(
+            **client_kwargs,
+            event_hooks={
+                "request": [self._alog_request_duration],
+                "response": [self._alog_response_duration],
+            },
+        ) as client:
+            async with client.stream(
+                "POST",
+                url,
+                headers=headers,
+                json=request_body,
+                timeout=self.default_request_timeout,
+            ) as response:
+                try:
+                    response.raise_for_status()
+                except httpx.HTTPStatusError as err:
+                    if self.logger:
+                        self.logger.error(
+                            "Error querying LLM: %s (%s)",
+                            err.response.reason_phrase,
+                            err.response.status_code,
+                            extra={
+                                "ActionName": self.settings.action_name,
+                                "ActionId": self.settings.action_id,
+                            }
+                            if self.settings
+                            else None,
+                        )
+                    # Read the response body for streaming responses
+                    await err.response.aread()
+                    raise self._make_status_error_from_response(err.response) from err
+
+                async for line in response.aiter_lines():
+                    line = line.strip()
+                    if not line:
+                        continue
+
+                    if self.logger:
+                        self.logger.debug(f"[SSE] Raw line: {line}")
+
+                    if line.startswith("data:"):
+                        data = line[
+                            5:
+                        ].strip()  # Remove "data:" prefix and strip whitespace
+                        if data == "[DONE]":
+                            break
+                        if not data:  # Skip empty data lines
+                            continue
+                        try:
+                            parsed = json.loads(data)
+                            # Skip empty chunks (some APIs send them as keepalive)
+                            # Check for truly empty: empty id AND (no choices or empty choices list)
+                            if (not parsed.get("id") or parsed.get("id") == "") and (
+                                not parsed.get("choices")
+                                or len(parsed.get("choices", [])) == 0
+                            ):
+                                if self.logger:
+                                    self.logger.debug(
+                                        "[SSE] Skipping empty keepalive chunk"
+                                    )
+                                continue
+                            yield parsed
+                        except json.JSONDecodeError as e:
+                            if self.logger:
+                                self.logger.warning(
+                                    f"Failed to parse SSE chunk: {data}, error: {e}"
+                                )
+                            continue
+                    else:
+                        # Handle lines without "data: " prefix (some APIs send raw JSON)
+                        try:
+                            parsed = json.loads(line)
+                            if self.logger:
+                                self.logger.debug(f"[SSE] Parsed raw JSON: {parsed}")
+                            yield parsed
+                        except json.JSONDecodeError:
+                            # Not JSON, skip
+                            pass
+
     def _make_status_error_from_response(
         self,
         response: httpx.Response,
@@ -354,7 +632,7 @@ class UiPathRequestMixin(BaseModel):
                 title = body.get("title", "").lower()
                 if title == "license not available":
                     raise LangGraphRuntimeError(
-                        code="LICENSE_NOT_AVAILABLE",
+                        code=LangGraphErrorCode.LICENSE_NOT_AVAILABLE,
                         title=body.get("title", "License Not Available"),
                         detail=body.get(
                             "detail", "License not available for this service"
@@ -427,7 +705,7 @@ class UiPathRequestMixin(BaseModel):
         return "uipath"
 
     @property
-    def _identifying_params(self) -> Dict[str, Any]:
+    def _identifying_params(self) -> dict[str, Any]:
         return {
             "url": self.url,
             "model": self.model_name,
@@ -463,13 +741,11 @@ class UiPathRequestMixin(BaseModel):
         )
 
     @property
-    def auth_headers(self) -> Dict[str, str]:
+    def auth_headers(self) -> dict[str, str]:
         if not self._auth_headers:
             self._auth_headers = {
                 **self.default_headers,  # type: ignore
                 "Authorization": f"Bearer {self.access_token}",
-                "X-UiPath-LlmGateway-RequestingProduct": self.requesting_product,
-                "X-UiPath-LlmGateway-RequestingFeature": self.requesting_feature,
                 "X-UiPath-LlmGateway-TimeoutSeconds": str(self.default_request_timeout),
             }
             if self.is_normalized and self.model_name:
@@ -481,7 +757,7 @@ class UiPathRequestMixin(BaseModel):
                 self._auth_headers["x-uipath-internal-tenantid"] = self.tenant_id
         return self._auth_headers
 
-    def _get_llm_string(self, stop: Optional[list[str]] = None, **kwargs: Any) -> str:
+    def _get_llm_string(self, stop: list[str] | None = None, **kwargs: Any) -> str:
         serialized_repr = getattr(self, "_serialized", self.model_dump())
         _cleanup_llm_representation(serialized_repr, 1)
         kwargs = serialized_repr.get("kwargs", serialized_repr)

uipath_langchain/_utils/_settings.py

@@ -1,6 +1,6 @@
 # mypy: disable-error-code="syntax"
 import os
-from typing import Any, Optional
+from typing import Any
 
 import httpx
 from pydantic import Field
@@ -21,7 +21,7 @@ class UiPathCachedPathsSettings(BaseSettings):
 
 
 uipath_cached_paths_settings = UiPathCachedPathsSettings()
-uipath_token_header: Optional[str] = None
+uipath_token_header: str | None = None
 
 
 class UiPathClientFactorySettings(BaseSettings):

uipath_langchain/agent/exceptions/__init__.py

@@ -0,0 +1,6 @@
+from .exceptions import AgentNodeRoutingException, AgentTerminationException
+
+__all__ = [
+    "AgentNodeRoutingException",
+    "AgentTerminationException",
+]

uipath_langchain/agent/exceptions/exceptions.py

@@ -0,0 +1,11 @@
+"""Exceptions for the basic agent loop."""
+
+from uipath.runtime.errors import UiPathRuntimeError
+
+
+class AgentNodeRoutingException(Exception):
+    pass
+
+
+class AgentTerminationException(UiPathRuntimeError):
+    pass

uipath_langchain/agent/guardrails/__init__.py

@@ -0,0 +1,21 @@
+from .guardrail_nodes import (
+    create_agent_guardrail_node,
+    create_llm_guardrail_node,
+    create_tool_guardrail_node,
+)
+from .guardrails_factory import build_guardrails_with_actions
+from .guardrails_subgraph import (
+    create_agent_guardrails_subgraph,
+    create_llm_guardrails_subgraph,
+    create_tool_guardrails_subgraph,
+)
+
+__all__ = [
+    "create_llm_guardrails_subgraph",
+    "create_agent_guardrails_subgraph",
+    "create_tool_guardrails_subgraph",
+    "create_llm_guardrail_node",
+    "create_agent_guardrail_node",
+    "create_tool_guardrail_node",
+    "build_guardrails_with_actions",
+]

uipath_langchain/agent/guardrails/actions/__init__.py

@@ -0,0 +1,11 @@
+from .base_action import GuardrailAction
+from .block_action import BlockAction
+from .escalate_action import EscalateAction
+from .log_action import LogAction
+
+__all__ = [
+    "GuardrailAction",
+    "BlockAction",
+    "LogAction",
+    "EscalateAction",
+]

uipath_langchain/agent/guardrails/actions/base_action.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from uipath.platform.guardrails import BaseGuardrail, GuardrailScope
+
+from uipath_langchain.agent.guardrails.types import ExecutionStage
+
+GuardrailActionNode = tuple[str, Any]
+
+
+class GuardrailAction(ABC):
+    """Extensible action interface producing a node to enforce the action on guardrail validation failure."""
+
+    @abstractmethod
+    def action_node(
+        self,
+        *,
+        guardrail: BaseGuardrail,
+        scope: GuardrailScope,
+        execution_stage: ExecutionStage,
+        guarded_component_name: str,
+    ) -> GuardrailActionNode:
+        """Create and return the Action node to execute on validation failure."""
+        ...

uipath_langchain/agent/guardrails/actions/block_action.py

@@ -0,0 +1,42 @@
+import re
+
+from uipath.platform.guardrails import BaseGuardrail, GuardrailScope
+from uipath.runtime.errors import UiPathErrorCategory, UiPathErrorCode
+
+from uipath_langchain.agent.guardrails.types import ExecutionStage
+
+from ...exceptions import AgentTerminationException
+from ..types import AgentGuardrailsGraphState
+from .base_action import GuardrailAction, GuardrailActionNode
+
+
+class BlockAction(GuardrailAction):
+    """Action that terminates execution when a guardrail fails.
+
+    Args:
+        reason: Reason string to include in the raised exception title.
+    """
+
+    def __init__(self, reason: str) -> None:
+        self.reason = reason
+
+    def action_node(
+        self,
+        *,
+        guardrail: BaseGuardrail,
+        scope: GuardrailScope,
+        execution_stage: ExecutionStage,
+        guarded_component_name: str,
+    ) -> GuardrailActionNode:
+        raw_node_name = f"{scope.name}_{execution_stage.name}_{guardrail.name}_block"
+        node_name = re.sub(r"\W+", "_", raw_node_name.lower()).strip("_")
+
+        async def _node(_state: AgentGuardrailsGraphState):
+            raise AgentTerminationException(
+                code=UiPathErrorCode.EXECUTION_ERROR,
+                title="Guardrail violation",
+                detail=self.reason,
+                category=UiPathErrorCategory.USER,
+            )
+
+        return node_name, _node