fix: tests

CalebMartinUiPath · CalebMartinUiPath · commit bea930eba901 · 2026-01-29T17:10:20.000-08:00
diff --git a/src/uipath_langchain/agent/tools/internal_tools/batch_transform_tool.py b/src/uipath_langchain/agent/tools/internal_tools/batch_transform_tool.py
@@ -25,6 +25,9 @@
 
 from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model
 from uipath_langchain.agent.react.types import AgentGraphState
+from uipath_langchain.agent.tools.internal_tools.schema_utils import (
+    add_query_field_to_schema,
+)
 from uipath_langchain.agent.tools.static_args import handle_static_args
 from uipath_langchain.agent.tools.structured_tool_with_argument_properties import (
     StructuredToolWithArgumentProperties,
@@ -53,11 +56,9 @@ def create_batch_transform_tool(
     output_columns_setting = settings.output_columns
     web_search_grounding_setting = settings.web_search_grounding
 
-    # Check if query is dynamic or static
     is_query_static = query_setting and query_setting.variant == "static"
     static_query = query_setting.value if is_query_static else None
 
-    # Get static values for other settings
     static_folder_path_prefix = None
     if folder_path_prefix_setting:
         static_folder_path_prefix = getattr(folder_path_prefix_setting, "value", None)
@@ -67,7 +68,6 @@ def create_batch_transform_tool(
         value = getattr(web_search_grounding_setting, "value", None)
         static_web_search = value == "Enabled" if value else False
 
-    # Convert output columns to BatchTransformOutputColumn
     batch_transform_output_columns = [
         BatchTransformOutputColumn(name=col.name, description=col.description or "")
         for col in output_columns_setting
@@ -76,18 +76,11 @@ def create_batch_transform_tool(
     # Use resource input schema and add query field if dynamic
     input_schema = dict(resource.input_schema)
     if not is_query_static:
-        if "properties" not in input_schema:
-            input_schema["properties"] = {}
-        input_schema["properties"]["query"] = {
-            "type": "string",
-            "description": query_setting.description
-            if query_setting and query_setting.description
-            else "The query to create a batch transform off of",
-        }
-        if "required" not in input_schema:
-            input_schema["required"] = []
-        if "query" not in input_schema["required"]:
-            input_schema["required"].append("query")
+        add_query_field_to_schema(
+            input_schema,
+            query_description=query_setting.description if query_setting else None,
+            default_description="Describe the task: what to research, what to synthesize.",
+        )
 
     # Create input model from modified schema
     input_model = create_model(input_schema)
@@ -101,7 +94,6 @@ def create_batch_transform_tool(
         example_calls=[],  # Examples cannot be provided for internal tools
     )
     async def batch_transform_tool_fn(**kwargs: Any) -> dict[str, Any]:
-        # Get query - dynamic from kwargs or static from settings
         query = kwargs.get("query") if not is_query_static else static_query
         if not query:
             raise ValueError("Query is required for Batch Transform tool")
@@ -113,27 +105,22 @@ async def batch_transform_tool_fn(**kwargs: Any) -> dict[str, Any]:
         if not attachment:
             raise ValueError("Attachment is required for Batch Transform tool")
 
-        # Extract attachment ID using getattr (works for Pydantic models)
         attachment_id = getattr(attachment, "ID", None)
         if not attachment_id:
             raise ValueError("Attachment ID is required")
 
-        # Get destination path, default to output.csv if not provided
         destination_path = kwargs.get("destination_path", "output.csv")
 
-        # Create ephemeral index directly via SDK
         uipath = UiPath()
         ephemeral_index = await uipath.context_grounding.create_ephemeral_index_async(
             usage=EphemeralIndexUsage.BATCH_RAG,
             attachments=[attachment_id],
         )
 
-        # Wait for index ingestion only if in progress
         if ephemeral_index.in_progress_ingestion():
             ephemeral_index_dict = interrupt(WaitEphemeralIndex(index=ephemeral_index))
             ephemeral_index = ContextGroundingIndex(**ephemeral_index_dict)
 
-        # Create Batch Transform request using interrupt
         return interrupt(
             CreateBatchTransform(
                 name=f"task-{uuid.uuid4()}",
@@ -144,7 +131,7 @@ async def batch_transform_tool_fn(**kwargs: Any) -> dict[str, Any]:
                 storage_bucket_folder_path_prefix=static_folder_path_prefix,
                 enable_web_search_grounding=static_web_search,
                 destination_path=destination_path,
-                is_ephemeral=True,
+                is_ephemeral_index=True,
             )
         )
 
diff --git a/src/uipath_langchain/agent/tools/internal_tools/deeprag_tool.py b/src/uipath_langchain/agent/tools/internal_tools/deeprag_tool.py
@@ -25,6 +25,9 @@
 
 from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model
 from uipath_langchain.agent.react.types import AgentGraphState
+from uipath_langchain.agent.tools.internal_tools.schema_utils import (
+    add_query_field_to_schema,
+)
 from uipath_langchain.agent.tools.static_args import handle_static_args
 from uipath_langchain.agent.tools.structured_tool_with_argument_properties import (
     StructuredToolWithArgumentProperties,
@@ -51,36 +54,23 @@ def create_deeprag_tool(
     query_setting = settings.query
     citation_mode_setting = settings.citation_mode
 
-    # Determine citation mode
     citation_mode = (
         CitationMode(citation_mode_setting.value)
         if citation_mode_setting
         else CitationMode.INLINE
     )
 
-    # Check if query is dynamic or static
     is_query_static = query_setting and query_setting.variant == "static"
     static_query = query_setting.value if is_query_static else None
 
-    # Use resource input schema and add query field if dynamic
     input_schema = dict(resource.input_schema)
     if not is_query_static:
-        # Add query field to the schema
-        if "properties" not in input_schema:
-            input_schema["properties"] = {}
-        input_schema["properties"]["query"] = {
-            "type": "string",
-            "description": query_setting.description
-            if query_setting and query_setting.description
-            else "The query to create a deeprag off of",
-        }
-        # Add query to required fields
-        if "required" not in input_schema:
-            input_schema["required"] = []
-        if "query" not in input_schema["required"]:
-            input_schema["required"].append("query")
-
-    # Create input model from modified schema
+        add_query_field_to_schema(
+            input_schema,
+            query_description=query_setting.description if query_setting else None,
+            default_description="Describe the task: what to research across documents, what to synthesize and how to cite sources.",
+        )
+
     input_model = create_model(input_schema)
     output_model = create_model(resource.output_schema)
 
@@ -92,7 +82,6 @@ def create_deeprag_tool(
         example_calls=[],  # Examples cannot be provided for internal tools
     )
     async def deeprag_tool_fn(**kwargs: Any) -> dict[str, Any]:
-        # Get query - dynamic from kwargs or static from settings
         query = kwargs.get("query") if not is_query_static else static_query
         if not query:
             raise ValueError("Query is required for DeepRAG tool")
@@ -104,32 +93,28 @@ async def deeprag_tool_fn(**kwargs: Any) -> dict[str, Any]:
         if not attachment:
             raise ValueError("Attachment is required for DeepRAG tool")
 
-        # Extract attachment ID using getattr (works for Pydantic models)
         attachment_id = getattr(attachment, "ID", None)
         if not attachment_id:
             raise ValueError("Attachment ID is required")
 
-        # Create ephemeral index directly via SDK
         uipath = UiPath()
         ephemeral_index = await uipath.context_grounding.create_ephemeral_index_async(
             usage=EphemeralIndexUsage.DEEP_RAG,
             attachments=[attachment_id],
         )
 
-        # Wait for index ingestion only if in progress
         if ephemeral_index.in_progress_ingestion():
             ephemeral_index_dict = interrupt(WaitEphemeralIndex(index=ephemeral_index))
             ephemeral_index = ContextGroundingIndex(**ephemeral_index_dict)
 
-        # Create DeepRAG request using interrupt
         return interrupt(
             CreateDeepRag(
                 name=f"task-{uuid.uuid4()}",
                 index_name=ephemeral_index.name,
                 index_id=ephemeral_index.id,
                 prompt=query,
                 citation_mode=citation_mode,
-                is_ephemeral=True,
+                is_ephemeral_index=True,
             )
         )
 
diff --git a/src/uipath_langchain/agent/tools/internal_tools/schema_utils.py b/src/uipath_langchain/agent/tools/internal_tools/schema_utils.py
@@ -0,0 +1,31 @@
+"""Utility functions for internal tool schema manipulation."""
+
+
+def add_query_field_to_schema(
+    input_schema: dict,
+    query_description: str | None = None,
+    default_description: str = "Query or prompt for the operation.",
+) -> None:
+    """Add a dynamic query field to an input schema.
+
+    This modifies the input schema in-place by adding a 'query' property
+    and marking it as required.
+
+    Args:
+        input_schema: The JSON schema dict to modify
+        query_description: Custom description for the query field
+        default_description: Default description if query_description is not provided
+    """
+    if "properties" not in input_schema:
+        input_schema["properties"] = {}
+
+    input_schema["properties"]["query"] = {
+        "type": "string",
+        "description": query_description if query_description else default_description,
+    }
+
+    if "required" not in input_schema:
+        input_schema["required"] = []
+
+    if "query" not in input_schema["required"]:
+        input_schema["required"].append("query")
diff --git a/tests/agent/tools/internal_tools/test_analyze_files_tool.py b/tests/agent/tools/internal_tools/test_analyze_files_tool.py
@@ -7,7 +7,7 @@
 from langchain_core.messages import AIMessage, HumanMessage
 from pydantic import BaseModel, ConfigDict, Field
 from uipath.agent.models.agent import (
-    AgentInternalToolProperties,
+    AgentInternalAnalyzeFilesToolProperties,
     AgentInternalToolResourceConfig,
     AgentInternalToolType,
 )
@@ -60,7 +60,7 @@ def resource_config(self):
         }
         output_schema = {"type": "object", "properties": {"result": {"type": "string"}}}
 
-        properties = AgentInternalToolProperties(
+        properties = AgentInternalAnalyzeFilesToolProperties(
             tool_type=AgentInternalToolType.ANALYZE_FILES
         )
 
diff --git a/tests/agent/tools/internal_tools/test_batch_transform_tool.py b/tests/agent/tools/internal_tools/test_batch_transform_tool.py
diff --git a/tests/agent/tools/internal_tools/test_deeprag_tool.py b/tests/agent/tools/internal_tools/test_deeprag_tool.py

Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@`
`7`	`7`	`from langchain_core.messages import AIMessage, HumanMessage`
`8`	`8`	`from pydantic import BaseModel, ConfigDict, Field`
`9`	`9`	`from uipath.agent.models.agent import (`
`10`		`- AgentInternalToolProperties,`
	`10`	`+ AgentInternalAnalyzeFilesToolProperties,`
`11`	`11`	`AgentInternalToolResourceConfig,`
`12`	`12`	`AgentInternalToolType,`
`13`	`13`	`)`
`@@ -60,7 +60,7 @@ def resource_config(self):`
`60`	`60`	`}`
`61`	`61`	`output_schema = {"type": "object", "properties": {"result": {"type": "string"}}}`
`62`	`62`
`63`		`- properties = AgentInternalToolProperties(`
	`63`	`+ properties = AgentInternalAnalyzeFilesToolProperties(`
`64`	`64`	`tool_type=AgentInternalToolType.ANALYZE_FILES`
`65`	`65`	`)`
`66`	`66`