fix: files messages stack

Author: cotovanu-cristian

src/uipath_langchain/agent/react/llm_with_files.py

@@ -1,10 +1,11 @@
 """LLM invocation with file attachments support."""
 
-from dataclasses import dataclass
+import asyncio
+from dataclasses import asdict, dataclass
 from typing import Any
 
 from langchain_core.language_models import BaseChatModel
-from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
+from langchain_core.messages import AIMessage, HumanMessage
 
 from .file_type_handler import build_message_content_part_from_data
 
@@ -18,6 +19,18 @@ class FileInfo:
     mime_type: str
 
 
+@dataclass
+class TextContentPart:
+    """Text content part in LangChain messages."""
+
+    type: str
+    text: str
+
+
+ContentPart = str | dict[str, Any]
+MessageContent = str | list[ContentPart]
+
+
 def _get_model_name(model: BaseChatModel) -> str:
     """Extract model name from a BaseChatModel instance."""
     for attr in ["model_name", "_model_name", "model", "model_id"]:
@@ -44,33 +57,76 @@ async def create_part_for_file(
     )
 
 
-async def llm_call_with_files(
-    messages: list[AnyMessage],
+async def add_files_to_message(
+    message: HumanMessage,
     files: list[FileInfo],
     model: BaseChatModel,
-) -> AIMessage:
-    """Invoke an LLM with file attachments.
+) -> HumanMessage:
+    """Add file attachments to a HumanMessage.
 
-    Downloads files, creates provider-specific content parts, and appends them
-    as a HumanMessage. If no files are provided, equivalent to model.ainvoke().
+    Downloads files in parallel, creates provider-specific content parts, and merges them
+    with the existing message content.
+
+    Args:
+        message: The HumanMessage to add files to (must be HumanMessage as files are user inputs)
+        files: List of files to add as attachments
+        model: The LLM model (used to determine provider-specific file format)
+
+    Returns:
+        A new HumanMessage with file content parts merged with original content.
+        The returned message will always have content as list[ContentPart] when files are added.
+
+    Notes:
+        - If message.content is a string, it will be converted to a text content part
+        - If message.content is already a list, files are appended to existing parts
+        - Empty files list returns the original message unchanged
     """
     if not files:
-        response = await model.ainvoke(messages)
-        if not isinstance(response, AIMessage):
-            raise TypeError(
-                f"LLM returned {type(response).__name__} instead of AIMessage"
-            )
-        return response
-
-    content_parts: list[str | dict[Any, Any]] = []
-    for file_info in files:
-        content_part = await create_part_for_file(file_info, model)
-        content_parts.append(content_part)
-
-    file_message = HumanMessage(content=content_parts)
-    all_messages = list(messages) + [file_message]
-
-    response = await model.ainvoke(all_messages)
-    if not isinstance(response, AIMessage):
-        raise TypeError(f"LLM returned {type(response).__name__} instead of AIMessage")
-    return response
+        return message
+
+    file_content_parts: list[dict[str, Any]] = await asyncio.gather(
+        *[create_part_for_file(file_info, model) for file_info in files]
+    )
+
+    existing_content: MessageContent = message.content
+    content_parts: list[ContentPart]
+
+    if isinstance(existing_content, str):
+        # convert string content to text content part
+        text_part = TextContentPart(type="text", text=existing_content)
+        content_parts = [asdict(text_part)]
+    elif isinstance(existing_content, list):
+        # preserve existing content parts
+        content_parts = list(existing_content)
+    else:
+        # handle edge case of empty or unexpected content
+        content_parts = []
+
+    # append file content parts
+    content_parts.extend(file_content_parts)
+
+    return HumanMessage(content=content_parts)
+
+
+def extract_text_content(message: AIMessage) -> str:
+    """Extract text content from an AI message.
+    Handles both simple string content and structured content with multiple parts extracting only the text portions.
+    """
+    content = message.content
+
+    if isinstance(content, str):
+        return content
+
+    if isinstance(content, list):
+        text_parts: list[str] = []
+        for part in content:
+            if isinstance(part, dict) and part.get("type") == "text":
+                text = part.get("text", "")
+                if text:
+                    text_parts.append(text)
+            elif isinstance(part, str):
+                text_parts.append(part)
+
+        return "\n".join(text_parts)
+
+    return str(content)

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

@@ -11,7 +11,11 @@
 from uipath.platform import UiPath
 
 from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model
-from uipath_langchain.agent.react.llm_with_files import FileInfo, llm_call_with_files
+from uipath_langchain.agent.react.llm_with_files import (
+    FileInfo,
+    add_files_to_message,
+    extract_text_content,
+)
 from uipath_langchain.agent.tools.structured_tool_with_output_type import (
     StructuredToolWithOutputType,
 )
@@ -62,12 +66,17 @@ async def tool_fn(**kwargs: Any):
         if not files:
             return {"analysisResult": "No attachments provided to analyze."}
 
+        human_message = HumanMessage(content=analysisTask)
+        human_message_with_files = await add_files_to_message(human_message, files, llm)
+
         messages: list[AnyMessage] = [
             SystemMessage(content=ANALYZE_FILES_SYSTEM_MESSAGE),
-            HumanMessage(content=analysisTask),
+            human_message_with_files,
         ]
-        result = await llm_call_with_files(messages, files, llm)
-        return result
+        result = await llm.ainvoke(messages)
+
+        analysis_result = extract_text_content(result)
+        return analysis_result
 
     wrapper = get_job_attachment_wrapper(output_type=output_model)
     tool = AnalyzeFileTool(

tests/agent/tools/internal_tools/test_analyze_files_tool.py

@@ -4,7 +4,7 @@
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
-from langchain_core.messages import AIMessage
+from langchain_core.messages import AIMessage, HumanMessage
 from pydantic import BaseModel, ConfigDict, Field
 from uipath.agent.models.agent import (
     AgentInternalToolProperties,
@@ -76,15 +76,15 @@ def resource_config(self):
         "uipath_langchain.agent.wrappers.job_attachment_wrapper.get_job_attachment_wrapper"
     )
     @patch(
-        "uipath_langchain.agent.tools.internal_tools.analyze_files_tool.llm_call_with_files"
+        "uipath_langchain.agent.tools.internal_tools.analyze_files_tool.add_files_to_message"
     )
     @patch(
         "uipath_langchain.agent.tools.internal_tools.analyze_files_tool._resolve_job_attachment_arguments"
     )
     async def test_create_analyze_file_tool_success(
         self,
         mock_resolve_attachments,
-        mock_llm_call,
+        mock_add_files,
         mock_get_wrapper,
         resource_config,
         mock_llm,
@@ -98,7 +98,16 @@ async def test_create_analyze_file_tool_success(
                 mime_type="application/pdf",
             )
         ]
-        mock_llm_call.return_value = "Analysis complete"
+
+        # mock add_files_to_message to return a message with files added
+        mock_message_with_files = HumanMessage(
+            content=[
+                {"type": "text", "text": "Summarize the document"},
+                {"type": "file", "url": "https://example.com/file.pdf"},
+            ]
+        )
+        mock_add_files.return_value = mock_message_with_files
+
         mock_wrapper = Mock()
         mock_get_wrapper.return_value = mock_wrapper
 
@@ -121,18 +130,29 @@ async def test_create_analyze_file_tool_success(
         )
 
         # Verify calls
-        assert result == "Analysis complete"
+        assert result == "Analyzed result"
         mock_resolve_attachments.assert_called_once()
-        mock_llm_call.assert_called_once()
-
-        # Verify LLM call arguments
-        call_args = mock_llm_call.call_args
-        messages, files, llm = call_args[0]
-        assert len(messages) == 2
-        assert messages[0].content == ANALYZE_FILES_SYSTEM_MESSAGE
-        assert messages[1].content == "Summarize the document"
-        assert len(files) == 1
-        assert files[0].url == "https://example.com/file.pdf"
+        mock_add_files.assert_called_once()
+        mock_llm.ainvoke.assert_called_once()
+
+        # Verify add_files_to_message was called correctly
+        add_files_call_args = mock_add_files.call_args
+        message_arg = add_files_call_args[0][0]
+        files_arg = add_files_call_args[0][1]
+        llm_arg = add_files_call_args[0][2]
+
+        assert isinstance(message_arg, HumanMessage)
+        assert message_arg.content == "Summarize the document"
+        assert len(files_arg) == 1
+        assert files_arg[0].url == "https://example.com/file.pdf"
+        assert llm_arg == mock_llm
+
+        # Verify llm.ainvoke was called with correct messages
+        ainvoke_call_args = mock_llm.ainvoke.call_args
+        messages_arg = ainvoke_call_args[0][0]
+        assert len(messages_arg) == 2
+        assert messages_arg[0].content == ANALYZE_FILES_SYSTEM_MESSAGE
+        assert messages_arg[1] == mock_message_with_files
 
     @patch(
         "uipath_langchain.agent.wrappers.job_attachment_wrapper.get_job_attachment_wrapper"
@@ -176,15 +196,15 @@ async def test_create_analyze_file_tool_missing_attachments(
         "uipath_langchain.agent.wrappers.job_attachment_wrapper.get_job_attachment_wrapper"
     )
     @patch(
-        "uipath_langchain.agent.tools.internal_tools.analyze_files_tool.llm_call_with_files"
+        "uipath_langchain.agent.tools.internal_tools.analyze_files_tool.add_files_to_message"
     )
     @patch(
         "uipath_langchain.agent.tools.internal_tools.analyze_files_tool._resolve_job_attachment_arguments"
     )
     async def test_create_analyze_file_tool_with_multiple_attachments(
         self,
         mock_resolve_attachments,
-        mock_llm_call,
+        mock_add_files,
         mock_get_wrapper,
         resource_config,
         mock_llm,
@@ -202,10 +222,25 @@ async def test_create_analyze_file_tool_with_multiple_attachments(
                 mime_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
             ),
         ]
-        mock_llm_call.return_value = "Multiple files analyzed"
+
+        # mock add_files_to_message to return a message with multiple files
+        mock_message_with_files = HumanMessage(
+            content=[
+                {"type": "text", "text": "Compare these documents"},
+                {"type": "file", "url": "https://example.com/file1.pdf"},
+                {"type": "file", "url": "https://example.com/file2.docx"},
+            ]
+        )
+        mock_add_files.return_value = mock_message_with_files
+
         mock_wrapper = Mock()
         mock_get_wrapper.return_value = mock_wrapper
 
+        # setup llm to return analyzed result
+        mock_llm.ainvoke = AsyncMock(
+            return_value=AIMessage(content="Multiple files analyzed")
+        )
+
         tool = create_analyze_file_tool(resource_config, mock_llm)
 
         mock_attachments = [
@@ -227,8 +262,8 @@ async def test_create_analyze_file_tool_with_multiple_attachments(
         assert result == "Multiple files analyzed"
         mock_resolve_attachments.assert_called_once()
 
-        # Verify LLM received both files
-        call_args = mock_llm_call.call_args
+        # Verify add_files_to_message received both files
+        call_args = mock_add_files.call_args
         files = call_args[0][1]
         assert len(files) == 2