fix: deeprag prefix fix (#674)

Co-authored-by: Gabriel Martin <91462031+gcuip@users.noreply.github.com>

Author: CalebMartinUiPath

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.8.8"
+version = "0.8.10"
 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/context_tool.py

@@ -1,15 +1,16 @@
 """Context tool creation for semantic index retrieval."""
 
 import uuid
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
 from langchain_core.documents import Document
 from langchain_core.messages import ToolCall
 from langchain_core.tools import BaseTool, StructuredTool
-from pydantic import BaseModel, Field, create_model
+from pydantic import BaseModel, Field, TypeAdapter, create_model
 from uipath.agent.models.agent import (
     AgentContextResourceConfig,
     AgentContextRetrievalMode,
+    AgentToolArgumentProperties,
 )
 from uipath.eval.mocks import mockable
 from uipath.platform import UiPath
@@ -41,6 +42,44 @@
 from .tool_node import ToolWrapperReturnType
 from .utils import sanitize_tool_name
 
+_ARG_PROPS_ADAPTER = TypeAdapter(Dict[str, AgentToolArgumentProperties])
+
+
+def _get_argument_properties(
+    resource: AgentContextResourceConfig,
+) -> dict[str, AgentToolArgumentProperties]:
+    """Extract argumentProperties from the resource's extra fields.
+
+    AgentContextResourceConfig doesn't declare argument_properties yet,
+    but BaseCfg(extra="allow") preserves the raw JSON value.
+    """
+    raw = (
+        resource.model_extra.get("argumentProperties") if resource.model_extra else None
+    )
+    if not raw:
+        return {}
+    return _ARG_PROPS_ADAPTER.validate_python(raw)
+
+
+def _build_folder_path_prefix_arg_props(
+    resource: AgentContextResourceConfig,
+) -> dict[str, Any]:
+    """Build argument_properties for folder_path_prefix from settings.
+
+    Fallback for when settings bag doesn't include argumentProperties
+    at the resource level but does set settings.folder_path_prefix
+    with variant="argument".
+    """
+    assert resource.settings.folder_path_prefix is not None
+    argument_path = (resource.settings.folder_path_prefix.value or "").strip("{}")
+    return {
+        "folder_path_prefix": {
+            "variant": "argument",
+            "argumentPath": argument_path,
+            "isSensitive": False,
+        }
+    }
+
 
 def is_static_query(resource: AgentContextResourceConfig) -> bool:
     """Check if the resource configuration uses a static query variant."""
@@ -152,27 +191,31 @@ def handle_deep_rag(
     if static:
         assert prompt is not None
 
-    folder_path_prefix = None
+    static_folder_path_prefix = None
     if (
         resource.settings.folder_path_prefix
         and resource.settings.folder_path_prefix.value
+        and resource.settings.folder_path_prefix.variant == "static"
     ):
-        folder_path_prefix = resource.settings.folder_path_prefix.value
+        static_folder_path_prefix = resource.settings.folder_path_prefix.value
 
     file_extension = None
     if resource.settings.file_extension and resource.settings.file_extension.value:
         file_extension = resource.settings.file_extension.value
 
-    glob_pattern = build_glob_pattern(
-        folder_path_prefix=folder_path_prefix, file_extension=file_extension
-    )
-
     output_model = create_model(
         "DeepRagOutputModel",
         __base__=DeepRagContent,
         deep_rag_id=(str, Field(alias="deepRagId")),
     )
 
+    arg_props = _get_argument_properties(resource)
+
+    has_folder_path_prefix_arg = "folder_path_prefix" in arg_props or (
+        resource.settings.folder_path_prefix
+        and resource.settings.folder_path_prefix.variant == "argument"
+    )
+
     schema_fields: dict[str, Any] = (
         {}
         if static
@@ -186,6 +229,18 @@ def handle_deep_rag(
             ),
         }
     )
+
+    if has_folder_path_prefix_arg:
+        schema_fields["folder_path_prefix"] = (
+            str,
+            Field(
+                default=None,
+                description="The folder path prefix within the index to filter on",
+            ),
+        )
+        if "folder_path_prefix" not in arg_props:
+            arg_props = _build_folder_path_prefix_arg_props(resource)
+
     input_model = create_model("DeepRagInput", **schema_fields)
 
     @mockable(
@@ -195,8 +250,14 @@ def handle_deep_rag(
         output_schema=output_model.model_json_schema(),
         example_calls=[],  # Examples cannot be provided for context.
     )
-    async def context_tool_fn(query: Optional[str] = None) -> dict[str, Any]:
+    async def context_tool_fn(
+        query: Optional[str] = None, folder_path_prefix: Optional[str] = None
+    ) -> dict[str, Any]:
         actual_prompt = prompt or query
+        glob_pattern = build_glob_pattern(
+            folder_path_prefix=static_folder_path_prefix or folder_path_prefix,
+            file_extension=file_extension,
+        )
 
         @durable_interrupt
         async def create_deep_rag():
@@ -211,12 +272,13 @@ async def create_deep_rag():
 
         return await create_deep_rag()
 
-    return StructuredToolWithOutputType(
+    return StructuredToolWithArgumentProperties(
         name=tool_name,
         description=resource.description,
         args_schema=input_model,
         coroutine=context_tool_fn,
         output_type=output_model,
+        argument_properties=arg_props,
         metadata={
             "tool_type": "context",
             "display_name": resource.name,
@@ -271,15 +333,19 @@ def handle_batch_transform(
     if static:
         assert prompt is not None
 
-    folder_path_prefix = None
+    static_folder_path_prefix = None
     if (
         resource.settings.folder_path_prefix
         and resource.settings.folder_path_prefix.value
+        and resource.settings.folder_path_prefix.variant == "static"
     ):
-        folder_path_prefix = resource.settings.folder_path_prefix.value
+        static_folder_path_prefix = resource.settings.folder_path_prefix.value
 
-    glob_pattern = build_glob_pattern(
-        folder_path_prefix=folder_path_prefix, file_extension=None
+    arg_props = _get_argument_properties(resource)
+
+    has_folder_path_prefix_arg = "folder_path_prefix" in arg_props or (
+        resource.settings.folder_path_prefix
+        and resource.settings.folder_path_prefix.variant == "argument"
     )
 
     output_model = create_model_from_schema(BATCH_TRANSFORM_OUTPUT_SCHEMA)
@@ -300,6 +366,16 @@ def handle_batch_transform(
             description="The relative file path destination for the modified csv file",
         ),
     )
+    if has_folder_path_prefix_arg:
+        schema_fields["folder_path_prefix"] = (
+            str,
+            Field(
+                default=None,
+                description="The folder path prefix within the index to filter on",
+            ),
+        )
+        if "folder_path_prefix" not in arg_props:
+            arg_props = _build_folder_path_prefix_arg_props(resource)
     input_model = create_model("BatchTransformInput", **schema_fields)
 
     @mockable(
@@ -310,9 +386,15 @@ def handle_batch_transform(
         example_calls=[],  # Examples cannot be provided for context.
     )
     async def context_tool_fn(
-        query: Optional[str] = None, destination_path: str = "output.csv"
+        query: Optional[str] = None,
+        destination_path: str = "output.csv",
+        folder_path_prefix: Optional[str] = None,
     ) -> dict[str, Any]:
         actual_prompt = prompt or query
+        glob_pattern = build_glob_pattern(
+            folder_path_prefix=static_folder_path_prefix or folder_path_prefix,
+            file_extension=None,
+        )
 
         @durable_interrupt
         async def create_batch_transform():
@@ -362,7 +444,7 @@ async def context_batch_transform_wrapper(
         args_schema=input_model,
         coroutine=context_tool_fn,
         output_type=output_model,
-        argument_properties={},
+        argument_properties=arg_props,
         metadata={
             "tool_type": "context",
             "display_name": resource.name,