feat: add bedrock chat models

Author: radugheo

.gitignore

@@ -177,3 +177,8 @@ cython_debug/
 **/.uipath
 **/**.nupkg
 **/__uipath/
+.claude/settings.local.json
+
+/.vscode/launch.json
+
+playground.py

pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
     "llama-index>=0.14.8",
     "llama-index-embeddings-azure-openai>=0.4.1",
     "llama-index-llms-azure-openai>=0.4.2",
+    "llama-index-llms-google-genai>=0.8.0",
     "openinference-instrumentation-llama-index>=4.3.9",
     "uipath>=2.2.26, <2.3.0",
 ]
@@ -24,6 +25,19 @@ maintainers = [
     { name = "Cristian Pufu", email = "cristian.pufu@uipath.com" }
 ]
 
+[project.optional-dependencies]
+bedrock = [
+    "llama-index-llms-bedrock>=0.3.0",
+    "llama-index-llms-bedrock-converse>=0.3.0",
+    "boto3>=1.28.0",
+    "aiobotocore>=2.5.0",
+]
+vertex = [
+    "llama-index-llms-vertex>=0.4.0",
+    "google-cloud-aiplatform>=1.38.0",
+    "httpx>=0.27.0",
+]
+
 [project.entry-points."uipath.middlewares"]
 register = "uipath_llamaindex.middlewares:register_middleware"
 

src/uipath_llamaindex/llms/__init__.py

@@ -2,8 +2,18 @@
     OpenAIModel,
     UiPathOpenAI,
 )
+from .supported_models import (
+    BedrockModels,
+    GeminiModels,
+    OpenAIModels,
+)
+from .vertex import UiPathVertex
 
 __all__ = [
     "UiPathOpenAI",
+    "UiPathVertex",
     "OpenAIModel",
+    "OpenAIModels",
+    "GeminiModels",
+    "BedrockModels",
 ]

src/uipath_llamaindex/llms/bedrock.py

@@ -0,0 +1,218 @@
+import logging
+import os
+from typing import Optional
+
+from uipath.utils import EndpointManager
+
+from .supported_models import BedrockModels
+
+logger = logging.getLogger(__name__)
+
+
+def _check_bedrock_dependencies() -> None:
+    """Check if required dependencies for UiPath Bedrock LLMs are installed."""
+    import importlib.util
+
+    missing_packages = []
+
+    if importlib.util.find_spec("llama_index.llms.bedrock") is None:
+        missing_packages.append("llama-index-llms-bedrock")
+
+    if importlib.util.find_spec("llama_index.llms.bedrock_converse") is None:
+        missing_packages.append("llama-index-llms-bedrock-converse")
+
+    if importlib.util.find_spec("boto3") is None:
+        missing_packages.append("boto3")
+
+    if importlib.util.find_spec("aiobotocore") is None:
+        missing_packages.append("aiobotocore")
+
+    if missing_packages:
+        packages_str = ", ".join(missing_packages)
+        raise ImportError(
+            f"The following packages are required to use UiPath Bedrock LLMs: {packages_str}\n"
+            "Please install them using one of the following methods:\n\n"
+            "  # Using pip:\n"
+            f"  pip install uipath-llamaindex[bedrock]\n\n"
+            "  # Using uv:\n"
+            f"  uv add 'uipath-llamaindex[bedrock]'\n\n"
+        )
+
+
+_check_bedrock_dependencies()
+
+import boto3
+from llama_index.llms.bedrock import Bedrock
+from llama_index.llms.bedrock_converse import BedrockConverse
+
+
+class AwsBedrockCompletionsPassthroughClient:
+    def __init__(
+        self,
+        model: str,
+        token: str,
+        api_flavor: str,
+    ):
+        self.model = model
+        self.token = token
+        self.api_flavor = api_flavor
+        self._vendor = "awsbedrock"
+        self._url: Optional[str] = None
+
+    @property
+    def endpoint(self) -> str:
+        vendor_endpoint = EndpointManager.get_vendor_endpoint()
+        formatted_endpoint = vendor_endpoint.format(
+            vendor=self._vendor,
+            model=self.model,
+        )
+        return formatted_endpoint
+
+    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
+
+    def get_client(self):
+        client = boto3.client(
+            "bedrock-runtime",
+            region_name="us-east-1",
+            aws_access_key_id="none",
+            aws_secret_access_key="none",
+            verify=False,
+        )
+        client.meta.events.register(
+            "before-send.bedrock-runtime.*", self._modify_request
+        )
+        return client
+
+    def get_session(self):
+        """Get aiobotocore session for async operations with custom event handlers."""
+        from aiobotocore.session import get_session
+
+        session = get_session()
+        session.get_component("event_emitter").register(
+            "before-send.bedrock-runtime.*", self._modify_request
+        )
+        return session
+
+    def _modify_request(self, request, **kwargs):
+        """Intercept boto3 request and redirect to LLM Gateway"""
+        # Detect streaming based on URL suffix:
+        # - converse-stream / invoke-with-response-stream -> streaming
+        # - converse / invoke -> non-streaming
+        streaming = "true" if request.url.endswith("-stream") else "false"
+        request.url = self._build_base_url()
+
+        headers = {
+            "Authorization": f"Bearer {self.token}",
+            "X-UiPath-LlmGateway-ApiFlavor": self.api_flavor,
+            "X-UiPath-Streaming-Enabled": streaming,
+        }
+
+        job_key = os.getenv("UIPATH_JOB_KEY")
+        process_key = os.getenv("UIPATH_PROCESS_KEY")
+        if job_key:
+            headers["X-UiPath-JobKey"] = job_key
+        if process_key:
+            headers["X-UiPath-ProcessKey"] = process_key
+
+        request.headers.update(headers)
+
+
+class UiPathChatBedrockConverse(BedrockConverse):
+    def __init__(
+        self,
+        org_id: Optional[str] = None,
+        tenant_id: Optional[str] = None,
+        token: Optional[str] = None,
+        model: str = BedrockModels.anthropic_claude_haiku_4_5,
+        **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"
+            )
+
+        passthrough_client = AwsBedrockCompletionsPassthroughClient(
+            model=model,
+            token=token,
+            api_flavor="converse",
+        )
+
+        client = passthrough_client.get_client()
+        botocore_session = passthrough_client.get_session()
+
+        super().__init__(
+            model=model,
+            client=client,
+            botocore_session=botocore_session,
+            region_name="us-east-1",
+            aws_access_key_id="none",
+            aws_secret_access_key="none",
+            **kwargs,
+        )
+
+
+class UiPathChatBedrock(Bedrock):
+    def __init__(
+        self,
+        org_id: Optional[str] = None,
+        tenant_id: Optional[str] = None,
+        token: Optional[str] = None,
+        model: str = BedrockModels.anthropic_claude_haiku_4_5,
+        context_size: int = 200000,
+        **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"
+            )
+
+        passthrough_client = AwsBedrockCompletionsPassthroughClient(
+            model=model,
+            token=token,
+            api_flavor="invoke",
+        )
+
+        client = passthrough_client.get_client()
+
+        super().__init__(
+            model=model,
+            client=client,
+            context_size=context_size,
+            aws_access_key_id="none",
+            aws_secret_access_key="none",
+            region_name="us-east-1",
+            **kwargs,
+        )

src/uipath_llamaindex/llms/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"