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

uipath_langchain/chat/openai.py

@@ -0,0 +1,133 @@
+import logging
+import os
+from typing import Optional
+
+import httpx
+from langchain_openai import AzureChatOpenAI
+from uipath._utils._ssl_context import get_httpx_client_kwargs
+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),
+                **get_httpx_client_kwargs(),
+            ),
+            http_client=httpx.Client(
+                transport=UiPathSyncURLRewriteTransport(verify=True),
+                **get_httpx_client_kwargs(),
+            ),
+            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"

uipath_langchain/chat/vertex.py

@@ -0,0 +1,255 @@
+import os
+from typing import Any, Optional
+
+import httpx
+from uipath._utils._ssl_context import get_httpx_client_kwargs
+from uipath.utils import EndpointManager
+
+from .supported_models import GeminiModels
+
+
+def _check_genai_dependencies() -> None:
+    """Check if required dependencies for UiPathChatVertex are installed."""
+    import importlib.util
+
+    missing_packages = []
+
+    if importlib.util.find_spec("langchain_google_genai") is None:
+        missing_packages.append("langchain-google-genai")
+
+    if importlib.util.find_spec("google.genai") is None:
+        missing_packages.append("google-genai")
+
+    if missing_packages:
+        packages_str = ", ".join(missing_packages)
+        raise ImportError(
+            f"The following packages are required to use UiPathChatVertex: {packages_str}\n"
+            "Please install them using one of the following methods:\n\n"
+            "  # Using pip:\n"
+            f"  pip install uipath-langchain[vertex]\n\n"
+            "  # Using uv:\n"
+            f"  uv add 'uipath-langchain[vertex]'\n\n"
+        )
+
+
+_check_genai_dependencies()
+
+import google.genai
+from google.genai import types as genai_types
+from langchain_google_genai import ChatGoogleGenerativeAI
+from pydantic import PrivateAttr
+
+
+def _rewrite_request_for_gateway(
+    request: httpx.Request, gateway_url: str
+) -> httpx.Request:
+    """Rewrite a request to redirect to the UiPath gateway."""
+    url_str = str(request.url)
+    if "generateContent" in url_str or "streamGenerateContent" in url_str:
+        is_streaming = "alt=sse" in url_str
+
+        headers = dict(request.headers)
+
+        headers["X-UiPath-Streaming-Enabled"] = "true" if is_streaming else "false"
+
+        gateway_url_parsed = httpx.URL(gateway_url)
+        if gateway_url_parsed.host:
+            headers["host"] = gateway_url_parsed.host
+
+        return httpx.Request(
+            method=request.method,
+            url=gateway_url,
+            headers=headers,
+            content=request.content,
+            extensions=request.extensions,
+        )
+    return request
+
+
+class _UrlRewriteTransport(httpx.BaseTransport):
+    """Transport that rewrites URLs to redirect to UiPath gateway."""
+
+    def __init__(self, gateway_url: str):
+        self.gateway_url = gateway_url
+        self._transport = httpx.HTTPTransport()
+
+    def handle_request(self, request: httpx.Request) -> httpx.Response:
+        request = _rewrite_request_for_gateway(request, self.gateway_url)
+        return self._transport.handle_request(request)
+
+    def close(self) -> None:
+        self._transport.close()
+
+
+class _AsyncUrlRewriteTransport(httpx.AsyncBaseTransport):
+    """Async transport that rewrites URLs to redirect to UiPath gateway."""
+
+    def __init__(self, gateway_url: str):
+        self.gateway_url = gateway_url
+        self._transport = httpx.AsyncHTTPTransport()
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        request = _rewrite_request_for_gateway(request, self.gateway_url)
+        return await self._transport.handle_async_request(request)
+
+    async def aclose(self) -> None:
+        await self._transport.aclose()
+
+
+class UiPathChatVertex(ChatGoogleGenerativeAI):
+    """UiPath Vertex AI Chat model that routes requests through UiPath's LLM Gateway."""
+
+    _vendor: str = PrivateAttr(default="vertexai")
+    _model_name: str = PrivateAttr()
+    _uipath_token: str = PrivateAttr()
+    _uipath_llmgw_url: Optional[str] = PrivateAttr(default=None)
+
+    def __init__(
+        self,
+        org_id: Optional[str] = None,
+        tenant_id: Optional[str] = None,
+        token: Optional[str] = None,
+        model_name: str = GeminiModels.gemini_2_5_flash,
+        temperature: Optional[float] = None,
+        **kwargs: Any,
+    ):
+        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"
+            )
+
+        uipath_url = self._build_base_url(model_name)
+        headers = self._build_headers(token)
+
+        http_options = genai_types.HttpOptions(
+            httpx_client=httpx.Client(
+                transport=_UrlRewriteTransport(uipath_url),
+                headers=headers,
+                **get_httpx_client_kwargs(),
+            ),
+            httpx_async_client=httpx.AsyncClient(
+                transport=_AsyncUrlRewriteTransport(uipath_url),
+                headers=headers,
+                **get_httpx_client_kwargs(),
+            ),
+        )
+
+        if temperature is None and (
+            "gemini-3" in model_name or "gemini-2" in model_name
+        ):
+            temperature = 1.0
+
+        super().__init__(
+            model=model_name,
+            google_api_key="uipath-gateway",
+            temperature=temperature,
+            **kwargs,
+        )
+
+        custom_client = google.genai.Client(
+            api_key="uipath-gateway",
+            http_options=http_options,
+        )
+
+        object.__setattr__(self, "client", custom_client)
+
+        self._model_name = model_name
+        self._uipath_token = token
+        self._uipath_llmgw_url = uipath_url
+
+        if self.temperature is not None and not 0 <= self.temperature <= 2.0:
+            raise ValueError("temperature must be in the range [0.0, 2.0]")
+
+        if self.top_p is not None and not 0 <= self.top_p <= 1:
+            raise ValueError("top_p must be in the range [0.0, 1.0]")
+
+        if self.top_k is not None and self.top_k <= 0:
+            raise ValueError("top_k must be positive")
+
+        additional_headers = self.additional_headers or {}
+        self.default_metadata = tuple(additional_headers.items())
+
+    @staticmethod
+    def _build_headers(token: str) -> dict[str, str]:
+        """Build HTTP headers for UiPath Gateway requests."""
+        headers = {
+            "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
+
+    @staticmethod
+    def _build_base_url(model_name: str) -> str:
+        """Build the full URL for the UiPath LLM Gateway."""
+        env_uipath_url = os.getenv("UIPATH_URL")
+
+        if not env_uipath_url:
+            raise ValueError("UIPATH_URL environment variable is required")
+
+        vendor_endpoint = EndpointManager.get_vendor_endpoint()
+        formatted_endpoint = vendor_endpoint.format(
+            vendor="vertexai",
+            model=model_name,
+        )
+        return f"{env_uipath_url.rstrip('/')}/{formatted_endpoint}"
+
+    def _stream(self, messages, stop=None, run_manager=None, **kwargs):
+        """Streaming fallback - calls _generate and yields single response."""
+        from langchain_core.messages import AIMessageChunk
+        from langchain_core.outputs import ChatGenerationChunk
+
+        result = self._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
+
+        if result.generations:
+            message = result.generations[0].message
+            chunk = AIMessageChunk(
+                content=message.content,
+                additional_kwargs=message.additional_kwargs,
+                response_metadata=getattr(message, "response_metadata", {}),
+                id=message.id,
+                tool_calls=getattr(message, "tool_calls", []),
+                tool_call_chunks=getattr(message, "tool_call_chunks", []),
+            )
+            if hasattr(message, "usage_metadata") and message.usage_metadata:
+                chunk.usage_metadata = message.usage_metadata
+
+            yield ChatGenerationChunk(message=chunk)
+
+    async def _astream(self, messages, stop=None, run_manager=None, **kwargs):
+        """Async streaming fallback - calls _agenerate and yields single response."""
+        from langchain_core.messages import AIMessageChunk
+        from langchain_core.outputs import ChatGenerationChunk
+
+        result = await self._agenerate(
+            messages, stop=stop, run_manager=run_manager, **kwargs
+        )
+
+        if result.generations:
+            message = result.generations[0].message
+            chunk = AIMessageChunk(
+                content=message.content,
+                additional_kwargs=message.additional_kwargs,
+                response_metadata=getattr(message, "response_metadata", {}),
+                id=message.id,
+                tool_calls=getattr(message, "tool_calls", []),
+                tool_call_chunks=getattr(message, "tool_call_chunks", []),
+            )
+            if hasattr(message, "usage_metadata") and message.usage_metadata:
+                chunk.usage_metadata = message.usage_metadata
+
+            yield ChatGenerationChunk(message=chunk)

uipath_langchain/embeddings/embeddings.py

@@ -1,8 +1,7 @@
 import os
-from typing import List, Optional
+from typing import Any
 
 import httpx
-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
@@ -11,9 +10,15 @@ from uipath_langchain._utils._request_mixin import UiPathRequestMixin
 
 
 class UiPathAzureOpenAIEmbeddings(UiPathRequestMixin, AzureOpenAIEmbeddings):
-    """Custom Embeddings connector for LangChain integration with UiPath, with minimal changes compared to AzureOpenAIEmbeddings."""
+    """Custom Embeddings connector for LangChain integration with UiPath.
 
-    model_name: Optional[str] = Field(
+    This class modifies the OpenAI client to:
+    - Use UiPath endpoints
+    - Log request/response durations
+    - Apply custom URL preparation and header building
+    """
+
+    model_name: str | None = Field(
         default_factory=lambda: os.getenv(
             "UIPATH_MODEL_NAME", "text-embedding-3-large"
         ),
@@ -36,6 +41,7 @@ class UiPathAzureOpenAIEmbeddings(UiPathRequestMixin, AzureOpenAIEmbeddings):
             ),
             **kwargs,
         )
+        # Monkey-patch the OpenAI client to use your custom methods
         self.client._client._prepare_url = self._prepare_url
         self.client._client._build_headers = self._build_headers
         self.async_client._client._prepare_url = self._prepare_url
@@ -50,63 +56,154 @@ class UiPathAzureOpenAIEmbeddings(UiPathRequestMixin, AzureOpenAIEmbeddings):
 
 
 class UiPathOpenAIEmbeddings(UiPathRequestMixin, OpenAIEmbeddings):
-    """Custom Embeddings connector for LangChain integration with UiPath, with full control over the embedding call."""
+    """Custom Embeddings connector for LangChain integration with UiPath.
+
+    This implementation uses custom _call and _acall methods for full control
+    over the API request/response cycle.
+    """
 
-    model_name: Optional[str] = Field(
+    model_name: str | None = Field(
         default_factory=lambda: os.getenv(
             "UIPATH_MODEL_NAME", "text-embedding-3-large"
         ),
         alias="model",
     )
 
+    # Add instance variables for tracking if needed
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._total_tokens = 0
+        self._total_requests = 0
+
     def embed_documents(
-        self, texts: List[str], chunk_size: Optional[int] = None, **kwargs
-    ) -> List[List[float]]:
-        """Embed a list of documents using the UiPath."""
-        embeddings = []
-        total_tokens = 0
-        # Process in chunks if specified
-        chunk_size_ = chunk_size or self.chunk_size or len(texts)
+        self, texts: list[str], chunk_size: int | None = None, **kwargs: Any
+    ) -> list[list[float]]:
+        """Embed a list of documents using UiPath endpoint.
+
+        Args:
+            texts: List of texts to embed
+            chunk_size: Number of texts to process in each batch
+            **kwargs: Additional arguments passed to the API
+
+        Returns:
+            List of embeddings for each text
+        """
+        chunk_size_ = chunk_size or self.chunk_size
+        embeddings: list[list[float]] = []
+
         for i in range(0, len(texts), chunk_size_):
             chunk = texts[i : i + chunk_size_]
-            payload = {"input": chunk}
+
+            # Build payload matching OpenAI API format
+            payload: dict[str, Any] = {
+                "input": chunk,
+                "model": self.model,
+            }
+
+            # Add optional parameters
+            if self.dimensions is not None:
+                payload["dimensions"] = self.dimensions
+
+            # Add model_kwargs and any additional kwargs
+            payload.update(self.model_kwargs)
+            payload.update(kwargs)
+
+            # Make the API call using custom _call method
             response = self._call(self.url, payload, self.auth_headers)
+
+            # Extract embeddings
             chunk_embeddings = [r["embedding"] for r in response["data"]]
-            total_tokens += response["usage"]["prompt_tokens"]
             embeddings.extend(chunk_embeddings)
-        if contextvar := openai_callback_var.get():
-            contextvar.prompt_tokens += total_tokens
-            contextvar.total_tokens += total_tokens
-            contextvar.successful_requests += 1
+
+            # Track usage internally (optional)
+            if "usage" in response:
+                self._total_tokens += response["usage"].get("total_tokens", 0)
+                self._total_requests += 1
+
         return embeddings
 
     async def aembed_documents(
         self,
-        texts: List[str],
-        chunk_size: Optional[int] = None,
-        **kwargs,
-    ) -> List[List[float]]:
-        """Async version of embed_documents."""
-        embeddings = []
-        total_tokens = 0
-        # Process in chunks if specified
-        chunk_size_ = chunk_size or self.chunk_size or len(texts)
+        texts: list[str],
+        chunk_size: int | None = None,
+        **kwargs: Any,
+    ) -> list[list[float]]:
+        """Async version of embed_documents.
+
+        Args:
+            texts: List of texts to embed
+            chunk_size: Number of texts to process in each batch
+            **kwargs: Additional arguments passed to the API
+
+        Returns:
+            List of embeddings for each text
+        """
+        chunk_size_ = chunk_size or self.chunk_size
+        embeddings: list[list[float]] = []
+
         for i in range(0, len(texts), chunk_size_):
             chunk = texts[i : i + chunk_size_]
-            payload = {"input": chunk}
+
+            # Build payload matching OpenAI API format
+            payload: dict[str, Any] = {
+                "input": chunk,
+                "model": self.model,
+            }
+
+            # Add optional parameters
+            if self.dimensions is not None:
+                payload["dimensions"] = self.dimensions
+
+            # Add model_kwargs and any additional kwargs
+            payload.update(self.model_kwargs)
+            payload.update(kwargs)
+
+            # Make the async API call using custom _acall method
             response = await self._acall(self.url, payload, self.auth_headers)
+
+            # Extract embeddings
             chunk_embeddings = [r["embedding"] for r in response["data"]]
-            total_tokens += response["usage"]["prompt_tokens"]
             embeddings.extend(chunk_embeddings)
-        if contextvar := openai_callback_var.get():
-            contextvar.prompt_tokens += total_tokens
-            contextvar.total_tokens += total_tokens
-            contextvar.successful_requests += 1
+
+            # Track usage internally (optional)
+            if "usage" in response:
+                self._total_tokens += response["usage"].get("total_tokens", 0)
+                self._total_requests += 1
+
         return embeddings
 
     @property
     def endpoint(self) -> str:
+        """Get the UiPath endpoint for embeddings."""
         endpoint = EndpointManager.get_embeddings_endpoint()
         return endpoint.format(
             model=self.model_name, api_version=self.openai_api_version
         )
+
+    @property
+    def url(self) -> str:
+        """Get the full URL for API requests."""
+        return self.endpoint
+
+    @property
+    def auth_headers(self) -> dict[str, str]:
+        """Get authentication headers for API requests."""
+        headers = {}
+        if self.openai_api_key:
+            headers["Authorization"] = (
+                f"Bearer {self.openai_api_key.get_secret_value()}"
+            )
+        if self.default_headers:
+            headers.update(self.default_headers)
+        return headers
+
+    def get_usage_stats(self) -> dict[str, int]:
+        """Get token usage statistics.
+
+        Returns:
+            Dictionary with total_tokens and total_requests
+        """
+        return {
+            "total_tokens": self._total_tokens,
+            "total_requests": self._total_requests,
+        }

uipath_langchain/middlewares.py

@@ -1,16 +1,10 @@
 from uipath._cli.middlewares import Middlewares
 
-from ._cli.cli_dev import langgraph_dev_middleware
-from ._cli.cli_eval import langgraph_eval_middleware
 from ._cli.cli_init import langgraph_init_middleware
 from ._cli.cli_new import langgraph_new_middleware
-from ._cli.cli_run import langgraph_run_middleware
 
 
 def register_middleware():
     """This function will be called by the entry point system when uipath_langchain is installed"""
     Middlewares.register("init", langgraph_init_middleware)
-    Middlewares.register("run", langgraph_run_middleware)
     Middlewares.register("new", langgraph_new_middleware)
-    Middlewares.register("dev", langgraph_dev_middleware)
-    Middlewares.register("eval", langgraph_eval_middleware)