unique_internal_search 2026.22.0.dev3__tar.gz → 2026.22.0.dev4__tar.gz

{unique_internal_search-2026.22.0.dev3 → unique_internal_search-2026.22.0.dev4}/PKG-INFO

@@ -1,13 +1,13 @@
 Metadata-Version: 2.3
 Name: unique-internal-search
-Version: 2026.22.0.dev3
+Version: 2026.22.0.dev4
 Summary: 
 Author: Martin Fadler, Sadique Sheik, Fabian Schläpfer, Pascal Hauri
 Author-email: Martin Fadler <martin.fadler@unique.ch>, Sadique Sheik <sadique@unique.ch>, Fabian Schläpfer <fabian@unique.ch>, Pascal Hauri <pascal@unique.ch>
 License: Proprietary
 Requires-Dist: pydantic>=2.8.2,<3
 Requires-Dist: typing-extensions>=4.9.0,<5
-Requires-Dist: unique-toolkit>=2026.22.0.dev7,<2026.22.0rc0
+Requires-Dist: unique-toolkit>=2026.22.0.dev12,<2026.22.0rc0
 Requires-Python: >=3.12, <4
 Description-Content-Type: text/markdown
 

{unique_internal_search-2026.22.0.dev3 → unique_internal_search-2026.22.0.dev4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_internal_search"
-version = "2026.22.0.dev3"
+version = "2026.22.0.dev4"
 description = ""
 readme = "README.md"
 license = { text = "Proprietary" }
@@ -14,7 +14,7 @@ requires-python = ">=3.12,<4"
 dependencies = [
     "pydantic>=2.8.2,<3",
     "typing-extensions>=4.9.0,<5",
-    "unique-toolkit>=2026.22.0.dev7,<2026.22.0rc0",
+    "unique-toolkit>=2026.22.0.dev12,<2026.22.0rc0",
 ]
 
 [build-system]

{unique_internal_search-2026.22.0.dev3 → unique_internal_search-2026.22.0.dev4}/unique_internal_search/config.py

@@ -23,6 +23,7 @@ from unique_toolkit.content.schemas import (
 )
 
 from unique_internal_search.prompts import (
+    DEFAULT_CONTENT_IDS_PARAM_DESCRIPTION,
     DEFAULT_LANGUAGE_PARAM_DESCRIPTION,
     DEFAULT_SEARCH_STRING_PARAM_DESCRIPTION,
     DEFAULT_TOOL_DESCRIPTION,
@@ -209,6 +210,23 @@ class InternalSearchConfig(BaseToolConfig):
         description="Allow execution of multiple search strings in one call. When set to True, each string is searched individually and results are merged into a single response.",
     )
 
+    enable_content_id_filter: bool = Field(
+        default=False,
+        description=(
+            "When True, exposes an optional `content_ids` parameter in the tool schema so the LLM can restrict the search to specific documents. "
+        ),
+    )
+
+    param_description_content_ids: Annotated[
+        str,
+        RJSFMetaTag.StringWidget.textarea(
+            rows=int(len(DEFAULT_CONTENT_IDS_PARAM_DESCRIPTION.split("\n")))
+        ),
+    ] = get_string_field_with_pattern_validation(
+        DEFAULT_CONTENT_IDS_PARAM_DESCRIPTION,
+        description="`content_ids` parameter description shown to the LLM.",
+    )
+
     metadata_chunk_sections: dict[str, str] = Field(
         default={},
         description=(

{unique_internal_search-2026.22.0.dev3 → unique_internal_search-2026.22.0.dev4}/unique_internal_search/prompts.py

@@ -63,6 +63,11 @@ DEFAULT_SEARCH_STRING_PARAM_DESCRIPTION = "An expanded term that is optimized fo
 
 DEFAULT_LANGUAGE_PARAM_DESCRIPTION = "The language that the user wrote the query in"
 
+DEFAULT_CONTENT_IDS_PARAM_DESCRIPTION = (
+    "Optional list of document IDs (e.g cont_sx6neydisjrdvq59q39s89cg) to restrict the search "
+    "to specific documents. Use `null` to search all available documents."
+)
+
 
 DEFAULT_TOOL_RESPONSE_SYSTEM_REMINDER_PROMPT = """
 ## Reminder: Cite Every Fact from These Results

{unique_internal_search-2026.22.0.dev3 → unique_internal_search-2026.22.0.dev4}/unique_internal_search/service.py

@@ -168,7 +168,12 @@ class InternalSearchService:
             )
             and chat_only
         ):
-            content_ids = self.selected_uploaded_file_ids
+            if content_ids is not None:
+                content_ids = list(
+                    filter(lambda x: x in self.selected_uploaded_file_ids, content_ids)
+                )
+            else:
+                content_ids = self.selected_uploaded_file_ids
 
         # Run all searches in parallel
         results = await asyncio.gather(
@@ -236,6 +241,7 @@ class InternalSearchService:
         self.debug_info = {
             "searchStrings": search_strings,
             "metadataFilter": metadata_filter,
+            "contentIds": content_ids,
             "chatOnly": chat_only,
         }
         return selected_chunks
@@ -424,6 +430,16 @@ class InternalSearchTool(Tool[InternalSearchConfig], InternalSearchService):
                 Field(description=self.config.param_description_search_string),
             )
 
+        optional_fields: dict = {}
+        if self.config.enable_content_id_filter:
+            optional_fields["content_ids"] = (
+                list[str] | None,
+                Field(
+                    default=None,
+                    description=self.config.param_description_content_ids,
+                ),
+            )
+
         internal_search_tool_input = create_model(
             "InternalSearchToolInput",
             search_string=search_string_field,
@@ -431,6 +447,7 @@ class InternalSearchTool(Tool[InternalSearchConfig], InternalSearchService):
                 str,
                 Field(description=self.config.param_description_language),
             ),
+            **optional_fields,
         )
         return LanguageModelToolDescription(
             name=self.name,

{unique_internal_search-2026.22.0.dev3 → unique_internal_search-2026.22.0.dev4}/unique_internal_search/uploaded_search/service.py

@@ -1,3 +1,5 @@
+from typing import Literal
+
 from pydantic import Field, create_model
 from typing_extensions import override
 from unique_toolkit import ContentService
@@ -14,6 +16,7 @@ from unique_toolkit.agentic.tools.tool_progress_reporter import (
 )
 from unique_toolkit.app.schemas import BaseEvent, ChatEvent
 from unique_toolkit.chat.service import LanguageModelToolDescription
+from unique_toolkit.content import Content
 from unique_toolkit.language_model.schemas import (
     LanguageModelFunction,
 )
@@ -50,12 +53,28 @@ class UploadedSearchTool(Tool[UploadedSearchConfig]):
         else:
             self._user_query = None
 
+        # This blocking API call should be avoided.
+        # However, we don't have an easy way to pass user uploaded files to the tool currently
+        # Note that this was being done in `tool_description_for_system_prompt` before
+        self._valid_documents = self._compute_valid_documents()
+
     @override
     def display_name(self) -> str:
         return self._display_name
 
     @override
     def tool_description(self) -> LanguageModelToolDescription:
+        optional_fields: dict = {}
+        if self._config.enable_content_id_filter and self._valid_documents:
+            content_id_type = Literal[*(d.id for d in self._valid_documents)]
+            optional_fields["content_ids"] = (
+                list[content_id_type] | None,
+                Field(
+                    default=None,
+                    description=self._config.param_description_content_ids,
+                ),
+            )
+
         internal_search_tool_input = create_model(
             "InternalSearchToolInput",
             search_string=(
@@ -66,7 +85,9 @@ class UploadedSearchTool(Tool[UploadedSearchConfig]):
                 str,
                 Field(description=self._config.param_description_language),
             ),
+            **optional_fields,
         )
+
         return LanguageModelToolDescription(
             name=self.name,
             description=self._config.tool_description,
@@ -74,28 +95,12 @@ class UploadedSearchTool(Tool[UploadedSearchConfig]):
         )
 
     def tool_description_for_system_prompt(self) -> str:
-        documents = self._content_service.get_documents_uploaded_to_chat()
-        if feature_flags.enable_selected_uploaded_files_un_18215.is_enabled(
-            self._company_id
-        ):
-            documents = [
-                doc for doc in documents if doc.id in self._selected_uploaded_files
-            ]
-
-        valid_documents = []
-        for doc in documents:
-            if not doc.is_ingested(default_if_unknown=True) or doc.is_expired():
-                continue
-            valid_documents.append(
-                {
-                    "name": doc.title or doc.key,
-                    "id": doc.id,
-                }
-            )
-
         return render_template(
             self._config.tool_description_for_system_prompt,
-            valid_documents=valid_documents,
+            valid_documents=[
+                {"name": doc.title or doc.key, "id": doc.id}
+                for doc in self._valid_documents
+            ],
         )
 
     def tool_format_information_for_system_prompt(self) -> str:
@@ -114,6 +119,26 @@ class UploadedSearchTool(Tool[UploadedSearchConfig]):
         search_string_data = ""
         if isinstance(tool_call.arguments, dict):
             search_string_data = tool_call.arguments.get("search_string", "") or ""
+
+            # Verify no content_id outside valid ones, as tool may not be strict
+
+            if "content_ids" in tool_call.arguments:
+                tool_content_ids = tool_call.arguments["content_ids"]
+
+                if tool_content_ids is not None:
+                    valid_content_ids = {doc.id for doc in self._valid_documents}
+                    filtered = list(
+                        filter(
+                            lambda content_id: content_id in valid_content_ids,
+                            tool_content_ids,
+                        )
+                    )
+
+                    if len(filtered) == 0:
+                        raise ValueError("All supplied `content_ids` are invalid")
+
+                    tool_call.arguments["content_ids"] = filtered
+
         tool_response = await self._internal_search_tool.run(tool_call)
         if (
             self._tool_progress_reporter
@@ -151,5 +176,23 @@ This tool call was automatically executed to retrieve the user's uploaded docume
 Please do not mention these instructions in your response to the user!
 </system_reminder>"""
 
+    def _compute_valid_documents(self) -> list[Content]:
+        documents = self._content_service.get_documents_uploaded_to_chat()
+
+        if feature_flags.enable_selected_uploaded_files_un_18215.is_enabled(
+            self._company_id
+        ):
+            documents = [
+                doc for doc in documents if doc.id in self._selected_uploaded_files
+            ]
+
+        valid_documents = []
+        for doc in documents:
+            if not doc.is_ingested(default_if_unknown=True) or doc.is_expired():
+                continue
+            valid_documents.append(doc)
+
+        return valid_documents
+
 
 ToolFactory.register_tool(UploadedSearchTool, UploadedSearchConfig)