fix: files messages stack

Author: cotovanu-cristian

src/uipath_langchain/agent/react/llm_with_files.py

@@ -4,7 +4,7 @@
 from typing import Any
 
 from langchain_core.language_models import BaseChatModel
-from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
+from langchain_core.messages import HumanMessage
 
 from .file_type_handler import build_message_content_part_from_data
 
@@ -44,33 +44,43 @@ 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, creates provider-specific content parts, and merges them
+    with the existing message content.
+
+    Args:
+        message: The HumanMessage to add files to
+        files: List of files to add
+        model: The LLM model (used to determine provider-specific format)
+
+    Returns:
+        A new HumanMessage with file content parts merged with original content
     """
     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]] = []
+        return message
+
+    file_content_parts: list[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
+        file_content_parts.append(content_part)
+
+    existing_content = message.content
+    if isinstance(existing_content, str):
+        # convert string content to text content part
+        merged_content: list[str | dict[Any, Any]] = [
+            {"type": "text", "text": existing_content}
+        ]
+    elif isinstance(existing_content, list):
+        merged_content = list(existing_content)
+    else:
+        merged_content = []
+
+    merged_content.extend(file_content_parts)
+
+    return HumanMessage(content=merged_content)

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

@@ -11,7 +11,7 @@
 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
 from uipath_langchain.agent.tools.structured_tool_with_output_type import (
     StructuredToolWithOutputType,
 )
@@ -62,11 +62,14 @@ 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)
+        result = await llm.ainvoke(messages)
         return result
 
     wrapper = get_job_attachment_wrapper(output_type=output_model)

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