feat: add bedrock chat models

Author: radugheo

pyproject.toml

@@ -24,6 +24,13 @@ 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",
+]
+
 [project.entry-points."uipath.middlewares"]
 register = "uipath_llamaindex.middlewares:register_middleware"
 

src/uipath_llamaindex/llms/bedrock.py

@@ -0,0 +1,203 @@
+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 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 _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()
+
+        super().__init__(
+            model=model,
+            client=client,
+            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"

testcases/chat-models/input.json

@@ -0,0 +1,3 @@
+{
+  "prompt": "Count from 1 to 5."
+}

testcases/chat-models/llama_index.json

@@ -0,0 +1,5 @@
+{
+  "workflows": {
+    "agent": "src/main.py:agent"
+  }
+}

testcases/chat-models/pyproject.toml

@@ -0,0 +1,15 @@
+[project]
+name = "chat-models-agent"
+version = "0.0.1"
+description = "Chat models testing agent for UiPath LlamaIndex"
+authors = [{ name = "John Doe", email = "john.doe@myemail.com" }]
+dependencies = [
+    "python-dotenv>=1.0.1",
+    "uipath-llamaindex",
+    "pydantic>=2.10.6",
+    "typing-extensions>=4.12.2",
+]
+requires-python = ">=3.11"
+
+[tool.uv.sources]
+uipath-llamaindex = { path = "../../", editable = true }

testcases/chat-models/run.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+set -e
+
+echo "Syncing dependencies..."
+uv sync
+
+echo "Adding UiPath LlamaIndex Chat Models Extras..."
+uv add uipath-llamaindex[bedrock]
+
+# uv add uipath-llamaindex[vertex]
+
+echo "Authenticating with UiPath..."
+uv run uipath auth --client-id="$CLIENT_ID" --client-secret="$CLIENT_SECRET" --base-url="$BASE_URL"
+
+echo "Initializing the project..."
+uv run uipath init
+
+echo "Packing agent..."
+uv run uipath pack
+
+echo "Running agent..."
+uv run uipath run agent --file input.json
+
+echo "Running agent again with empty UIPATH_JOB_KEY..."
+export UIPATH_JOB_KEY=""
+uv run uipath run agent --trace-file .uipath/traces.jsonl --file input.json >> local_run_output.log

testcases/chat-models/src/assert.py

@@ -0,0 +1,82 @@
+import json
+import os
+
+print("Checking chat models test agent output...")
+
+uipath_dir = ".uipath"
+assert os.path.exists(uipath_dir), "NuGet package directory (.uipath) not found"
+
+nupkg_files = [f for f in os.listdir(uipath_dir) if f.endswith(".nupkg")]
+assert nupkg_files, "NuGet package file (.nupkg) not found in .uipath directory"
+
+print(f"NuGet package found: {nupkg_files[0]}")
+
+output_file = "__uipath/output.json"
+assert os.path.isfile(output_file), "Agent output file not found"
+
+print("Agent output file found")
+
+with open(output_file, "r", encoding="utf-8") as f:
+    output_data = json.load(f)
+
+status = output_data.get("status")
+assert status == "successful", f"Agent execution failed with status: {status}"
+
+print("Agent execution status: successful")
+
+assert "output" in output_data, "Missing 'output' field in agent response"
+
+output_content = output_data["output"]
+
+assert "success" in output_content, "Missing 'success' field in output"
+success = output_content["success"]
+
+assert "result_summary" in output_content, "Missing 'result_summary' field in output"
+result_summary = output_content["result_summary"]
+assert result_summary and result_summary.strip() != "", "Result summary is empty"
+
+print(f"\nTest Results:")
+print(f"  Success: {success}")
+print(f"  Summary:\n{result_summary}")
+
+if not success:
+    print("\n" + "=" * 80)
+    print("TEST FAILURES")
+    print("=" * 80)
+    lines = result_summary.split("\n")
+    for line in lines:
+        if line.strip():
+            print(line)
+    print("=" * 80)
+
+assert success is True, f"Test did not succeed. See detailed results above."
+
+assert "chunks_received" in output_content, "Missing 'chunks_received' field"
+chunks_received = output_content["chunks_received"]
+assert chunks_received is not None and chunks_received > 0, (
+    f"Expected positive chunks_received, got: {chunks_received}"
+)
+print(f"Chunks received: {chunks_received}")
+
+assert "content_length" in output_content, "Missing 'content_length' field"
+content_length = output_content["content_length"]
+assert content_length is not None and content_length > 0, (
+    f"Expected positive content_length, got: {content_length}"
+)
+print(f"Content length: {content_length}")
+
+assert "tool_calls_count" in output_content, "Missing 'tool_calls_count' field"
+tool_calls_count = output_content["tool_calls_count"]
+assert tool_calls_count is not None and tool_calls_count > 0, (
+    f"Expected positive tool_calls_count, got: {tool_calls_count}"
+)
+print(f"Tool calls count: {tool_calls_count}")
+
+with open("local_run_output.log", "r", encoding="utf-8") as f:
+    local_run_output = f.read()
+
+assert "Successful execution." in local_run_output, (
+    f"Response does not contain 'Successful execution.'. Actual response: {local_run_output}"
+)
+
+print("All validations passed successfully!")