fix: add async session to BedrockConverse client

Author: radugheo

playground.py

@@ -1,5 +1,6 @@
 from llama_index.core.llms import ChatMessage
-from uipath_llamaindex.llms import UiPathVertex, GeminiModels
+
+from uipath_llamaindex.llms import GeminiModels, UiPathVertex
 
 
 def test_all_methods():

pyproject.toml

@@ -29,6 +29,7 @@ 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",

src/uipath_llamaindex/llms/__init__.py

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

src/uipath_llamaindex/llms/bedrock.py

@@ -24,6 +24,9 @@ def _check_bedrock_dependencies() -> None:
     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(
@@ -89,6 +92,16 @@ def get_client(self):
         )
         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:
@@ -146,14 +159,16 @@ def __init__(
         )
 
         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
+            **kwargs,
         )
 
 
@@ -199,5 +214,5 @@ def __init__(
             aws_access_key_id="none",
             aws_secret_access_key="none",
             region_name="us-east-1",
-            **kwargs
+            **kwargs,
         )

src/uipath_llamaindex/llms/vertex.py

@@ -67,7 +67,9 @@ def redirected_request(method, url, **kwargs_inner):
             is_streaming = kwargs_inner.get("stream", False)
             headers["X-UiPath-Streaming-Enabled"] = "true" if is_streaming else "false"
 
-            return original_request(method, self.llmgw_url, headers=headers, **kwargs_inner)
+            return original_request(
+                method, self.llmgw_url, headers=headers, **kwargs_inner
+            )
 
         self._session.request = redirected_request  # type: ignore[method-assign]
 
@@ -199,7 +201,9 @@ def _build_base_url(self) -> str:
             env_uipath_url = os.getenv("UIPATH_URL")
 
             if env_uipath_url:
-                self._uipath_url = f"{env_uipath_url.rstrip('/')}/{self._uipath_endpoint}"
+                self._uipath_url = (
+                    f"{env_uipath_url.rstrip('/')}/{self._uipath_endpoint}"
+                )
             else:
                 raise ValueError("UIPATH_URL environment variable is required")
 

testcases/chat-models/src/assert.py

@@ -35,7 +35,7 @@
 result_summary = output_content["result_summary"]
 assert result_summary and result_summary.strip() != "", "Result summary is empty"
 
-print(f"\nTest Results:")
+print("\nTest Results:")
 print(f"  Success: {success}")
 print(f"  Summary:\n{result_summary}")
 
@@ -49,7 +49,7 @@
             print(line)
     print("=" * 80)
 
-assert success is True, f"Test did not succeed. See detailed results above."
+assert success is True, "Test did not succeed. See detailed results above."
 
 assert "chunks_received" in output_content, "Missing 'chunks_received' field"
 chunks_received = output_content["chunks_received"]

testcases/chat-models/src/main.py

@@ -228,7 +228,7 @@ async def test_single_model_all(
             model_results[method_name] = "⊘ Not available"
 
     # Test tool calling with predict_and_call
-    logger.info(f"  Testing tool_calling...")
+    logger.info("  Testing tool_calling...")
     try:
         if hasattr(model, "apredict_and_call"):
             response = await model.apredict_and_call(
@@ -241,17 +241,17 @@ async def test_single_model_all(
                 logger.info(f"     Tool calls detected: {tool_calls_count}")
                 model_results["tool_calling"] = f"✓ ({tool_calls_count} calls)"
             else:
-                logger.warning(f"     No tool calls detected")
+                logger.warning("     No tool calls detected")
                 model_results["tool_calling"] = "✗ No tool calls detected"
         else:
-            logger.warning(f"     Tool calling not available")
+            logger.warning("     Tool calling not available")
             model_results["tool_calling"] = "⊘ Not available"
     except Exception as e:
         logger.error(f"     Tool calling failed: {e}")
         model_results["tool_calling"] = f"✗ {format_error_message(str(e))}"
 
     # Test structured output
-    logger.info(f"  Testing structured_output...")
+    logger.info("  Testing structured_output...")
     try:
         if hasattr(model, "as_structured_llm"):
             sllm = model.as_structured_llm(PersonInfo)
@@ -269,7 +269,7 @@ async def test_single_model_all(
                 logger.warning(f"     Response is not PersonInfo: {type(response)}")
                 model_results["structured_output"] = f"✗ Wrong type: {type(response)}"
         else:
-            logger.warning(f"     Structured output not available")
+            logger.warning("     Structured output not available")
             model_results["structured_output"] = "⊘ Not available"
     except Exception as e:
         logger.error(f"     Structured output failed: {e}")
@@ -357,7 +357,7 @@ async def run_tests(self, ev: TestInput) -> TestOutput:
         logger.info("TEST RESULTS")
         logger.info("=" * 80)
         logger.info(f"Success: {not has_failures}")
-        logger.info(f"Summary:\n" + "\n".join(summary_lines))
+        logger.info("Summary:\n" + "\n".join(summary_lines))
         logger.info(f"Chunks Received: {total_result.chunks}")
         logger.info(f"Content Length: {total_result.content_length}")
         logger.info(f"Tool Calls: {total_result.tool_calls}")