feat(agent): let the LLM pick the folder for semantic context searches [PC-5304]

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.18.55"
+version = "0.18.56"
 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

@@ -2,14 +2,14 @@
 
 import logging
 import uuid
-from typing import Any, Optional
+from typing import Annotated, Any, Optional
 
 from jsonpath_ng import parse  # type: ignore[import-untyped]
 from langchain_core.documents import Document
 from langchain_core.language_models import BaseChatModel
 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, WithJsonSchema, create_model
 from uipath.agent.models.agent import (
     AgentContextResourceConfig,
     AgentContextRetrievalMode,
@@ -69,6 +69,12 @@
 # downgrading to the old blocking behaviour, so the misconfiguration is visible.
 SEARCH_DURING_INGESTION_FF = "EnableContextGroundingSearchDuringIngestion"
 
+DYNAMIC_FOLDER_PATH_PREFIX_DESCRIPTION = (
+    "Optional folder path prefix that limits the search to documents under that "
+    "folder in the knowledge base, for example 'Reports/2024'. Leave empty to "
+    "search the whole knowledge base."
+)
+
 logger = logging.getLogger(__name__)
 
 _CONTEXT_GROUNDING_ERRORS: dict[
@@ -284,6 +290,17 @@ class ContextOutputSchemaModel(BaseModel):
         and resource.settings.folder_path_prefix.value
     )
 
+    folder_setting = resource.settings.folder_path_prefix
+    if folder_setting and folder_setting.variant == "dynamic":
+        schema_fields["folder_path_prefix"] = (
+            Annotated[Optional[str], WithJsonSchema({"type": "string"})],
+            Field(
+                default=None,
+                description=folder_setting.description
+                or DYNAMIC_FOLDER_PATH_PREFIX_DESCRIPTION,
+            ),
+        )
+
     _resolved_arg_folder_prefix: str | None = None
 
     input_model = create_model("SemanticSearchInput", **schema_fields)
@@ -297,9 +314,13 @@ class ContextOutputSchemaModel(BaseModel):
     )
     async def context_tool_fn(
         query: Optional[str] = None,
+        folder_path_prefix: Optional[str] = None,
     ) -> dict[str, Any]:
         resolved_folder_path_prefix = (
-            static_folder_path_prefix or _resolved_arg_folder_prefix
+            static_folder_path_prefix
+            or _resolved_arg_folder_prefix
+            or (folder_path_prefix or "").strip()
+            or None
         )
 
         debug_run = UiPathConfig.is_studio_project

tests/agent/tools/test_context_tool.py

@@ -28,6 +28,7 @@
     AgentStartupErrorCode,
 )
 from uipath_langchain.agent.tools.context_tool import (
+    DYNAMIC_FOLDER_PATH_PREFIX_DESCRIPTION,
     SEARCH_DURING_INGESTION_FF,
     _normalize_folder_prefix,
     build_glob_pattern,
@@ -1426,3 +1427,175 @@ async def test_opt_in_is_forwarded_to_the_retriever(self):
         kwargs = await self._retriever_kwargs_for(resource, flag_enabled=True)
 
         assert kwargs["search_during_ingestion"] is True
+
+
+class TestDynamicFolderPathPrefix:
+    @staticmethod
+    def _semantic_resource(
+        folder_path_prefix, query_variant="dynamic", query_value=None
+    ):
+        return _make_context_resource(
+            name="semantic_tool",
+            retrieval_mode=AgentContextRetrievalMode.SEMANTIC,
+            query_variant=query_variant,
+            query_value=query_value,
+            folder_path_prefix=folder_path_prefix,
+        )
+
+    @staticmethod
+    async def _search_scopes(resource, *calls_args):
+        sdk = MagicMock()
+        search_result = MagicMock()
+        search_result.semantic_results.values = []
+        search_result.semantic_results.metadata.operation_id = "op-1"
+        sdk.context_grounding.unified_search_async = AsyncMock(
+            return_value=search_result
+        )
+        with patch(
+            "uipath_langchain.retrievers.context_grounding_retriever.UiPath",
+            return_value=sdk,
+        ):
+            tool = handle_semantic_search("semantic_tool", resource)
+            for index, args in enumerate(calls_args):
+                call = {
+                    "type": "tool_call",
+                    "id": f"call-{index}",
+                    "name": "semantic_tool",
+                    "args": args,
+                }
+                wrapper = getattr(tool, "awrapper", None)
+                if wrapper is not None:
+                    await wrapper(tool, call, {})
+                else:
+                    await tool.ainvoke(call)
+
+        search = sdk.context_grounding.unified_search_async
+        assert search.await_count == len(calls_args)
+        return [call.kwargs["scope"] for call in search.call_args_list]
+
+    @classmethod
+    async def _search_scope(cls, resource, **tool_args):
+        [scope] = await cls._search_scopes(resource, tool_args)
+        return scope
+
+    def test_dynamic_folder_path_prefix_is_an_optional_tool_argument(self):
+        resource = self._semantic_resource(AgentContextQuerySetting(variant="dynamic"))
+
+        tool = handle_semantic_search("semantic_tool", resource)
+
+        assert tool.args_schema is not None
+        assert hasattr(tool.args_schema, "model_json_schema")
+        schema = tool.args_schema.model_json_schema()
+        assert schema["properties"]["folder_path_prefix"]["type"] == "string"
+        assert "folder_path_prefix" not in schema.get("required", [])
+
+    def test_dynamic_folder_path_prefix_uses_the_author_description(self):
+        resource = self._semantic_resource(
+            AgentContextQuerySetting(
+                variant="dynamic",
+                description="Docket folder, for example Patents/123",
+            )
+        )
+
+        tool = handle_semantic_search("semantic_tool", resource)
+
+        assert hasattr(tool.args_schema, "model_json_schema")
+        schema = tool.args_schema.model_json_schema()
+        assert (
+            schema["properties"]["folder_path_prefix"]["description"]
+            == "Docket folder, for example Patents/123"
+        )
+
+    @pytest.mark.parametrize("description", [None, ""])
+    def test_dynamic_folder_path_prefix_without_author_description_gets_the_default(
+        self, description
+    ):
+        resource = self._semantic_resource(
+            AgentContextQuerySetting(variant="dynamic", description=description)
+        )
+
+        tool = handle_semantic_search("semantic_tool", resource)
+
+        assert hasattr(tool.args_schema, "model_json_schema")
+        schema = tool.args_schema.model_json_schema()
+        assert (
+            schema["properties"]["folder_path_prefix"]["description"]
+            == DYNAMIC_FOLDER_PATH_PREFIX_DESCRIPTION
+        )
+
+    def test_static_folder_path_prefix_is_not_a_tool_argument(self):
+        resource = self._semantic_resource(
+            AgentContextQuerySetting(variant="static", value="Patents/123")
+        )
+
+        tool = handle_semantic_search("semantic_tool", resource)
+
+        assert hasattr(tool.args_schema, "model_json_schema")
+        schema = tool.args_schema.model_json_schema()
+        assert "folder_path_prefix" not in schema["properties"]
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("query_variant", "query_value", "query_args"),
+        [
+            ("dynamic", None, {"query": "invention summary"}),
+            ("static", "invention summary", {}),
+            ("argument", "{{topic}}", {"query": "invention summary"}),
+        ],
+    )
+    async def test_folder_chosen_by_the_llm_scopes_the_search(
+        self, query_variant, query_value, query_args
+    ):
+        resource = self._semantic_resource(
+            AgentContextQuerySetting(variant="dynamic"),
+            query_variant=query_variant,
+            query_value=query_value,
+        )
+
+        scope = await self._search_scope(
+            resource, **query_args, folder_path_prefix="Patents/123"
+        )
+
+        assert scope is not None
+        assert scope.folder == "Patents/123"
+
+    @pytest.mark.asyncio
+    async def test_each_call_uses_its_own_folder(self):
+        resource = self._semantic_resource(AgentContextQuerySetting(variant="dynamic"))
+
+        first, second = await self._search_scopes(
+            resource,
+            {"query": "invention summary", "folder_path_prefix": "Patents/123"},
+            {"query": "invention summary", "folder_path_prefix": "Patents/456"},
+        )
+
+        assert first.folder == "Patents/123"
+        assert second.folder == "Patents/456"
+
+    @pytest.mark.asyncio
+    async def test_no_folder_from_the_llm_searches_the_whole_index(self):
+        resource = self._semantic_resource(AgentContextQuerySetting(variant="dynamic"))
+
+        scope = await self._search_scope(resource, query="memory segregation")
+
+        assert scope is None
+
+    @pytest.mark.asyncio
+    async def test_blank_folder_from_the_llm_searches_the_whole_index(self):
+        resource = self._semantic_resource(AgentContextQuerySetting(variant="dynamic"))
+
+        scope = await self._search_scope(
+            resource, query="memory segregation", folder_path_prefix="  "
+        )
+
+        assert scope is None
+
+    @pytest.mark.asyncio
+    async def test_null_folder_from_the_llm_searches_the_whole_index(self):
+        resource = self._semantic_resource(AgentContextQuerySetting(variant="dynamic"))
+
+        scope = await self._search_scope(
+            resource, query="memory segregation", folder_path_prefix=None
+        )
+
+        assert scope is None