fix: prevent OOM in analyze files tool for large attachments

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.8.12"
+version = "0.8.13"
 description = "Python SDK that enables developers to build and deploy LangGraph agents to the UiPath Cloud Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath_langchain/agent/exceptions/exceptions.py

@@ -45,6 +45,8 @@ class AgentRuntimeErrorCode(str, Enum):
     # State
     STATE_ERROR = "STATE_ERROR"
 
+    FILE_ERROR = "FILE_ERROR"
+
     LLM_INVALID_RESPONSE = "LLM_INVALID_RESPONSE"
     TOOL_INVALID_WRAPPER_STATE = "TOOL_INVALID_WRAPPER_STATE"
 

src/uipath_langchain/agent/multimodal/invoke.py

@@ -1,6 +1,6 @@
 """LLM invocation with multimodal file attachments."""
 
-import asyncio
+import logging
 from typing import Any
 
 from langchain_core.language_models import BaseChatModel
@@ -12,25 +12,39 @@
 )
 from langchain_core.messages.content import create_file_block, create_image_block
 
-from .types import FileInfo
+from .types import MAX_FILE_SIZE_BYTES, FileInfo
 from .utils import download_file_base64, is_image, is_pdf, sanitize_filename
 
+logger = logging.getLogger("uipath")
+
 
 async def build_file_content_block(
     file_info: FileInfo,
+    *,
+    max_size: int = MAX_FILE_SIZE_BYTES,
 ) -> DataContentBlock:
     """Build a LangChain content block for a file attachment.
 
+    Downloads the file with size enforcement and creates the content block.
+    Size validation happens during download (via Content-Length check and
+    streaming guard) to avoid loading oversized files into memory.
+
     Args:
         file_info: File URL, name, and MIME type.
+        max_size: Maximum allowed raw file size in bytes. LLM providers
+            enforce payload limits; base64 encoding adds ~30% overhead.
 
     Returns:
         A DataContentBlock for the file (image or PDF).
 
     Raises:
-        ValueError: If the MIME type is not supported.
+        ValueError: If the MIME type is not supported or the file exceeds
+            the size limit for LLM payloads.
     """
-    base64_file = await download_file_base64(file_info.url)
+    try:
+        base64_file = await download_file_base64(file_info.url, max_size=max_size)
+    except ValueError as exc:
+        raise ValueError(f"File '{file_info.name}': {exc}") from exc
 
     if is_image(file_info.mime_type):
         return create_image_block(base64=base64_file, mime_type=file_info.mime_type)
@@ -47,6 +61,9 @@ async def build_file_content_block(
 async def build_file_content_blocks(files: list[FileInfo]) -> list[DataContentBlock]:
     """Build content blocks from file attachments.
 
+    Files are processed sequentially to avoid loading multiple large files
+    into memory simultaneously.
+
     Args:
         files: List of file information to convert to content blocks
 
@@ -56,9 +73,10 @@ async def build_file_content_blocks(files: list[FileInfo]) -> list[DataContentBl
     if not files:
         return []
 
-    file_content_blocks: list[DataContentBlock] = await asyncio.gather(
-        *[build_file_content_block(file) for file in files]
-    )
+    file_content_blocks: list[DataContentBlock] = []
+    for file in files:
+        block = await build_file_content_block(file)
+        file_content_blocks.append(block)
     return file_content_blocks
 
 
@@ -100,6 +118,9 @@ async def llm_call_with_files(
     all_messages = list(messages) + [file_message]
 
     response = await model.ainvoke(all_messages)
+
+    del all_messages, file_message, content_blocks
+
     if not isinstance(response, AIMessage):
         raise TypeError(f"LLM returned {type(response).__name__} instead of AIMessage")
     return response

src/uipath_langchain/agent/multimodal/types.py

@@ -2,6 +2,15 @@
 
 from dataclasses import dataclass
 
+MAX_FILE_SIZE_BYTES: int = 30 * 1024 * 1024  # 30MB
+
+IMAGE_MIME_TYPES: set[str] = {
+    "image/png",
+    "image/jpeg",
+    "image/gif",
+    "image/webp",
+}
+
 
 @dataclass
 class FileInfo:
@@ -10,11 +19,3 @@ class FileInfo:
     url: str
     name: str
     mime_type: str
-
-
-IMAGE_MIME_TYPES: set[str] = {
-    "image/png",
-    "image/jpeg",
-    "image/gif",
-    "image/webp",
-}

src/uipath_langchain/agent/multimodal/utils.py

@@ -2,6 +2,7 @@
 
 import base64
 import re
+from collections.abc import AsyncIterator
 
 import httpx
 from uipath._utils._ssl_context import get_httpx_client_kwargs
@@ -36,10 +37,99 @@ def is_image(mime_type: str) -> bool:
     return mime_type.lower() in IMAGE_MIME_TYPES
 
 
-async def download_file_base64(url: str) -> str:
-    """Download a file from a URL and return its content as a base64 string."""
+def _format_mb(size_bytes: int, decimals: int = 1) -> str:
+    """Format a byte count as MB.
+
+    Args:
+        size_bytes: Size in bytes.
+        decimals: Number of decimal places (0 for rounded integer).
+    """
+    return f"{size_bytes / (1024 * 1024):.{decimals}f} MB"
+
+
+async def encode_streamed_base64(
+    chunks: AsyncIterator[bytes],
+    *,
+    max_size: int = 0,
+) -> str:
+    """Incrementally base64-encode an async stream of byte chunks.
+
+    Encodes chunks as they arrive so the raw file bytes are never assembled
+    into a single contiguous buffer. base64 processes 3-byte groups, so a
+    remainder of 0-2 bytes is buffered between chunks.
+
+    Args:
+        chunks: Async iterator yielding raw byte chunks.
+        max_size: Maximum allowed total size in bytes. 0 means unlimited.
+
+    Returns:
+        The full base64-encoded string.
+
+    Raises:
+        ValueError: If the total size exceeds max_size.
+    """
+    encoded_buf = bytearray()
+    remainder = b""
+    total = 0
+
+    async for chunk in chunks:
+        total += len(chunk)
+        if max_size > 0 and total > max_size:
+            raise ValueError(
+                f"File exceeds the {_format_mb(max_size, decimals=0)}"
+                f" limit for LLM payloads"
+                f" (downloaded {_format_mb(total)} so far)"
+            )
+
+        data = remainder + chunk
+        usable = len(data) - (len(data) % 3)
+        if usable > 0:
+            encoded_buf += base64.b64encode(data[:usable])
+            remainder = data[usable:]
+        else:
+            remainder = data
+
+    if remainder:
+        encoded_buf += base64.b64encode(remainder)
+
+    result = encoded_buf.decode("ascii")
+    del encoded_buf
+    return result
+
+
+async def download_file_base64(url: str, *, max_size: int = 0) -> str:
+    """Download a file from a URL and return its content as a base64 string.
+
+    Args:
+        url: The URL to download from.
+        max_size: Maximum allowed file size in bytes. 0 means unlimited.
+
+    Raises:
+        ValueError: If the file exceeds max_size.
+        httpx.HTTPStatusError: If the HTTP request fails.
+    """
     async with httpx.AsyncClient(**get_httpx_client_kwargs()) as client:
-        response = await client.get(url)
-        response.raise_for_status()
-        file_content = response.content
-    return base64.b64encode(file_content).decode("utf-8")
+        async with client.stream("GET", url) as response:
+            response.raise_for_status()
+
+            # Fast reject via Content-Length before reading the body
+            if max_size > 0:
+                content_length = response.headers.get("content-length")
+                if content_length:
+                    try:
+                        content_length_value = int(content_length)
+                    except ValueError:
+                        content_length_value = None
+                    if (
+                        content_length_value is not None
+                        and content_length_value > max_size
+                    ):
+                        raise ValueError(
+                            f"File is {_format_mb(content_length_value)}"
+                            f" which exceeds the {_format_mb(max_size, decimals=0)}"
+                            f" limit for Agent LLM payloads"
+                        )
+
+            return await encode_streamed_base64(
+                response.aiter_bytes(), max_size=max_size
+            )

src/uipath_langchain/agent/tools/internal_tools/analyze_files_tool.py

@@ -1,4 +1,3 @@
-import asyncio
 import uuid
 from typing import Any, cast
 
@@ -19,7 +18,12 @@
 )
 from uipath.eval.mocks import mockable
 from uipath.platform import UiPath
+from uipath.runtime.errors import UiPathErrorCategory
 
+from uipath_langchain.agent.exceptions import (
+    AgentRuntimeError,
+    AgentRuntimeErrorCode,
+)
 from uipath_langchain.agent.multimodal import FileInfo, build_file_content_block
 from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model
 from uipath_langchain.agent.react.types import AgentGraphState
@@ -77,8 +81,16 @@ async def tool_fn(**kwargs: Any):
         if not files:
             return {"analysisResult": "No attachments provided to analyze."}
 
-        human_message = HumanMessage(content=analysis_task)
-        human_message_with_files = await add_files_to_message(human_message, files)
+        try:
+            human_message = HumanMessage(content=analysis_task)
+            human_message_with_files = await add_files_to_message(human_message, files)
+        except ValueError as exc:
+            raise AgentRuntimeError(
+                code=AgentRuntimeErrorCode.FILE_ERROR,
+                title="File attachment too large",
+                detail=str(exc),
+                category=UiPathErrorCategory.USER,
+            ) from exc
 
         messages: list[AnyMessage] = [
             SystemMessage(content=ANALYZE_FILES_SYSTEM_MESSAGE),
@@ -87,6 +99,8 @@ async def tool_fn(**kwargs: Any):
         config = var_child_runnable_config.get(None)
         result = await non_streaming_llm.ainvoke(messages, config=config)
 
+        del messages, human_message_with_files, files
+
         analysis_result = extract_text_content(result)
         return {"analysisResult": analysis_result}
 
@@ -172,9 +186,10 @@ async def add_files_to_message(
     if not files:
         return message
 
-    file_content_blocks: list[DataContentBlock] = await asyncio.gather(
-        *[build_file_content_block(file) for file in files]
-    )
+    file_content_blocks: list[DataContentBlock] = []
+    for file in files:
+        block = await build_file_content_block(file)
+        file_content_blocks.append(block)
     return append_content_blocks_to_message(
         message, cast(list[ContentBlock], file_content_blocks)
     )