fix(ECS): honor query argument bindings in DeepRAG and Batch Transform

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.18.18"
+version = "0.18.19"
 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/tools/internal_tools/batch_transform_tool.py

@@ -9,6 +9,7 @@
 from uipath.agent.models.agent import (
     AgentInternalBatchTransformToolProperties,
     AgentInternalToolResourceConfig,
+    AgentToolArgumentArgumentProperties,
 )
 from uipath.core.feature_flags import FeatureFlags
 from uipath.eval.mocks import mockable
@@ -82,6 +83,13 @@ def create_batch_transform_tool(
     is_query_static = query_setting and query_setting.variant == "static"
     static_query = query_setting.value if is_query_static else None
 
+    argument_properties = dict(resource.argument_properties)
+    if query_setting and query_setting.variant == "argument":
+        argument_properties["query"] = AgentToolArgumentArgumentProperties(
+            is_sensitive=False,
+            argument_path=(query_setting.value or "").strip("{}"),
+        )
+
     static_folder_path_prefix = None
     if folder_path_prefix_setting:
         static_folder_path_prefix = getattr(folder_path_prefix_setting, "value", None)
@@ -219,7 +227,7 @@ async def upload_result_attachment():
         args_schema=input_model,
         coroutine=batch_transform_tool_fn,
         output_type=output_model,
-        argument_properties=resource.argument_properties,
+        argument_properties=argument_properties,
         metadata={
             "tool_type": "context_grounding",
             "display_name": tool_name,

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

@@ -9,6 +9,7 @@
 from uipath.agent.models.agent import (
     AgentInternalDeepRagToolProperties,
     AgentInternalToolResourceConfig,
+    AgentToolArgumentArgumentProperties,
 )
 from uipath.core.feature_flags import FeatureFlags
 from uipath.eval.mocks import mockable
@@ -84,6 +85,13 @@ def create_deeprag_tool(
     is_query_static = query_setting and query_setting.variant == "static"
     static_query = query_setting.value if is_query_static else None
 
+    argument_properties = dict(resource.argument_properties)
+    if query_setting and query_setting.variant == "argument":
+        argument_properties["query"] = AgentToolArgumentArgumentProperties(
+            is_sensitive=False,
+            argument_path=(query_setting.value or "").strip("{}"),
+        )
+
     input_schema = dict(resource.input_schema)
     if not is_query_static:
         add_query_field_to_schema(
@@ -180,7 +188,7 @@ async def create_deeprag():
         args_schema=input_model,
         coroutine=deeprag_tool_fn,
         output_type=output_model,
-        argument_properties=resource.argument_properties,
+        argument_properties=argument_properties,
         metadata={
             "tool_type": resource.type.lower(),
             "display_name": tool_name,

tests/agent/tools/internal_tools/test_batch_transform_tool.py

@@ -5,7 +5,8 @@
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
-from pydantic import BaseModel, ConfigDict, Field
+from langchain_core.messages import ToolCall
+from pydantic import BaseModel, ConfigDict, Field, create_model
 from uipath.agent.models.agent import (
     AgentContextOutputColumn,
     AgentContextQuerySetting,
@@ -26,6 +27,7 @@
 from uipath_langchain.agent.tools.internal_tools.batch_transform_tool import (
     create_batch_transform_tool,
 )
+from uipath_langchain.agent.tools.static_args import StaticArgsHandler
 
 
 class MockAttachment(BaseModel):
@@ -333,6 +335,7 @@ async def test_create_batch_transform_tool_static_query_wait_for_ingestion(
         lambda **kwargs: lambda f: f,
     )
     @patch.dict(os.environ, {"UIPATH_FEATURE_DisableBatchTransformFromAttachment": "1"})
+    @pytest.mark.parametrize("variant", ["dynamic", "argument"])
     async def test_create_batch_transform_tool_dynamic_query(
         self,
         mock_interrupt,
@@ -341,8 +344,9 @@ async def test_create_batch_transform_tool_dynamic_query(
         mock_get_wrapper,
         resource_config_dynamic,
         mock_llm,
+        variant,
     ):
-        """Test Batch Transform tool with dynamic query."""
+        """Test model-supplied and argument-bound Batch Transform queries."""
         # Setup mocks
         mock_uipath = AsyncMock()
         mock_uipath_class.return_value = mock_uipath
@@ -371,18 +375,28 @@ async def test_create_batch_transform_tool_dynamic_query(
         mock_wrapper = Mock()
         mock_get_wrapper.return_value = mock_wrapper
 
-        # Create tool
+        if variant == "argument":
+            resource_config_dynamic.properties.settings.query = (
+                AgentContextQuerySetting(variant="argument", value="{{task}}")
+            )
         tool = create_batch_transform_tool(resource_config_dynamic, mock_llm)
+        expected = "Extract all names"
+        call = ToolCall(name=tool.name, args={"query": expected}, id="call-1")
+        if variant == "argument":
+            expected = "  exact {{literal}}\n task "
+            input_model = create_model("Input", task=(str, ...))
+            handler = StaticArgsHandler()
+            handler.initialize([tool], input_model(task=expected), input_model)
+            handler.apply_to_response([call])
 
         # Test tool execution with dynamic query
         mock_attachment = MockAttachment(
             ID=str(uuid.uuid4()), FullName="input.csv", MimeType="text/csv"
         )
 
         assert tool.coroutine is not None
-        result = await tool.coroutine(
-            attachment=mock_attachment, query="Extract all names"
-        )
+        result = await tool.coroutine(attachment=mock_attachment, **call["args"])
+        assert mock_interrupt.call_args.args[0].prompt == expected
 
         # Verify result contains attachment info
         assert result == {

tests/agent/tools/internal_tools/test_deeprag_tool.py

@@ -5,7 +5,8 @@
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
-from pydantic import BaseModel, ConfigDict, Field
+from langchain_core.messages import ToolCall
+from pydantic import BaseModel, ConfigDict, Field, create_model
 from uipath.agent.models.agent import (
     AgentContextQuerySetting,
     AgentInternalDeepRagSettings,
@@ -25,6 +26,7 @@
 from uipath_langchain.agent.tools.internal_tools.deeprag_tool import (
     create_deeprag_tool,
 )
+from uipath_langchain.agent.tools.static_args import StaticArgsHandler
 
 
 class MockAttachment(BaseModel):
@@ -464,3 +466,29 @@
         assert isinstance(create_payload, CreateDeepRag)
         assert create_payload.attachments == [attachment_id]
         assert create_payload.prompt == "What are the main points?"
+
+    @patch("uipath_langchain.agent.tools.internal_tools.deeprag_tool.interrupt")
+    @patch(
+        "uipath_langchain.agent.tools.internal_tools.deeprag_tool.mockable",
+        lambda **kwargs: lambda f: f,
+    )
+    async def test_argument_query_uses_selected_input(
+        self, mock_interrupt, resource_config_static, mock_llm
+    ):
+        resource_config_static.properties.settings.query = AgentContextQuerySetting(
+            variant="argument", value="{{task}}"
+        )
+        tool = create_deeprag_tool(resource_config_static, mock_llm)
+        input_model = create_model("Input", task=(str, ...))
+        expected = "  exact {{literal}}\n task "
+        handler = StaticArgsHandler()
+        handler.initialize([tool], input_model(task=expected), input_model)
+        call = ToolCall(name=tool.name, args={"query": "model query"}, id="call-1")
+        handler.apply_to_response([call])
+        attachment = MockAttachment(
+            ID=str(uuid.uuid4()), FullName="test.pdf", MimeType="application/pdf"
+        )
+        mock_interrupt.return_value = {"text": "done"}
+        assert tool.coroutine is not None
+        await tool.coroutine(**call["args"], attachment=attachment)
+        assert mock_interrupt.call_args.args[0].prompt == expected