unique_internal_search 2026.22.0.dev2__tar.gz → 2026.22.0.dev4__tar.gz

{unique_internal_search-2026.22.0.dev2 → 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.dev2
+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.dev5,<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.dev2 → unique_internal_search-2026.22.0.dev4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_internal_search"
-version = "2026.22.0.dev2"
+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.dev5,<2026.22.0rc0",
+    "unique-toolkit>=2026.22.0.dev12,<2026.22.0rc0",
 ]
 
 [build-system]

{unique_internal_search-2026.22.0.dev2 → 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.dev2 → 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.dev2 → 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.dev2 → unique_internal_search-2026.22.0.dev4}/unique_internal_search/uploaded_search/config.py

@@ -1,5 +1,7 @@
 from typing import Annotated
 
+from pydantic import Field
+from pydantic.json_schema import SkipJsonSchema
 from unique_toolkit._common.pydantic.rjsf_tags import RJSFMetaTag
 
 from unique_internal_search.config import InternalSearchConfig
@@ -51,3 +53,9 @@ class UploadedSearchConfig(InternalSearchConfig):
         DEFAULT_TOOL_FORMAT_INFORMATION_FOR_SYSTEM_PROMPT,
         description="Tool format information for the system prompt.",
     )
+    # This is a hack to ensure that the system reminder is enabled when the tool is forced,
+    # as we cannot at the moment easily pass extra values to the tool constructor
+    enable_tool_call_system_reminder: SkipJsonSchema[bool] = Field(
+        default=True,
+        description="Whether to attach the uploaded-search system reminder to tool responses.",
+    )

unique_internal_search-2026.22.0.dev2/unique_internal_search/uploaded_search/prompts.py → unique_internal_search-2026.22.0.dev4/unique_internal_search/uploaded_search/prompts/__init__.py

@@ -1,18 +1,11 @@
-DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT = (
-    "You can use the UploadedSearch tool to access and analyze documents uploaded by users during a chat. This tool is designed to handle a variety of document-related tasks, including summarization, explanation, and detailed information retrieval. "
-    "Use cases for the UploadedSearch tool include:\n"
-    "- Document Analysis: When a user uploads a document and asks for a summary, explanation, or specific details, this tool can extract and provide the requested information.\n"
-    "- Named Document Queries: If a user refers to a previously uploaded document by name (e.g., 'What does the Q2_Report.pdf say about revenue?'), this tool can locate and analyze the document to answer the query.\n"
-    "- Policy and Procedure Verification: Use the tool to find the most current company policies, procedures, or guidelines within uploaded documents.\n"
-    "- Project-Specific Information: Access project documents, reports, or meeting notes uploaded by users to provide precise details.\n"
-    "- Confidential and Proprietary Information: Ensure that sensitive topics requiring proprietary knowledge or confidential data are sourced securely from uploaded documents.\n\n"
-    "**Instruction Query Splitting**\n"
-    "You should split the user question into multiple search strings when the user's question needs to be decomposed / rewritten to find different facts. Perform an individual tool call for each search string. Avoid overly broad queries that may return unrelated results. Ensure the search string is specific and relevant to the uploaded document(s).\n\n"
-    "Examples:\n"
-    'User: "What does the Q2_Report.pdf say about revenue and expenses?" => search strings: ["Q2_Report.pdf revenue", "Q2_Report.pdf expenses"]\n'
-    'User: "Summarize the uploaded document." => search string: ["Summarize the uploaded document"]\n\n'
-    "{system_prompt_valid_documents}"
-    "{system_prompt_expired_documents}"
+from pathlib import Path
+
+from unique_toolkit._common.utils.jinja.helpers import load_template
+
+_PROMPTS_DIR = Path(__file__).parent
+
+DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT = load_template(
+    _PROMPTS_DIR, "tool_description_for_system_prompt.j2"
 )
 
 DEFAULT_TOOL_FORMAT_INFORMATION_FOR_SYSTEM_PROMPT = (

unique_internal_search-2026.22.0.dev4/unique_internal_search/uploaded_search/prompts/tool_description_for_system_prompt.j2

@@ -0,0 +1,21 @@
+You can use the UploadedSearch tool to access and analyze documents uploaded by users during a chat. This tool is designed to handle a variety of document-related tasks, including summarization, explanation, and detailed information retrieval. Use cases for the UploadedSearch tool include:
+- Document Analysis: When a user uploads a document and asks for a summary, explanation, or specific details, this tool can extract and provide the requested information.
+- Named Document Queries: If a user refers to a previously uploaded document by name (e.g., 'What does the Q2_Report.pdf say about revenue?'), this tool can locate and analyze the document to answer the query.
+- Policy and Procedure Verification: Use the tool to find the most current company policies, procedures, or guidelines within uploaded documents.
+- Project-Specific Information: Access project documents, reports, or meeting notes uploaded by users to provide precise details.
+- Confidential and Proprietary Information: Ensure that sensitive topics requiring proprietary knowledge or confidential data are sourced securely from uploaded documents.
+
+**Instruction Query Splitting**
+You should split the user question into multiple search strings when the user's question needs to be decomposed / rewritten to find different facts. Perform an individual tool call for each search string. Avoid overly broad queries that may return unrelated results. Ensure the search string is specific and relevant to the uploaded document(s).
+
+Examples:
+User: "What does the Q2_Report.pdf say about revenue and expenses?" => search strings: ["Q2_Report.pdf revenue", "Q2_Report.pdf expenses"]
+User: "Summarize the uploaded document." => search string: ["Summarize the uploaded document"]
+
+{% if valid_documents %}
+**The currently uploaded and searchable documents are the following**
+{% for document in valid_documents %}
+- {{ document.name }} (content_id: {{ document.id }})
+{% endfor %}
+
+{% endif %}

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

@@ -1,6 +1,9 @@
+from typing import Literal
+
 from pydantic import Field, create_model
 from typing_extensions import override
 from unique_toolkit import ContentService
+from unique_toolkit._common.utils.jinja.render import render_template
 from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
 from unique_toolkit.agentic.feature_flags import feature_flags
 from unique_toolkit.agentic.tools.factory import ToolFactory
@@ -13,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,
 )
@@ -49,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=(
@@ -65,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,
@@ -73,54 +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 = []
-        expired_documents = []
-        for doc in documents:
-            if not doc.is_ingested(default_if_unknown=True):
-                continue
-            if doc.is_expired():
-                expired_documents.append(doc)
-            else:
-                valid_documents.append(doc)
-
-        system_prompt_valid_documents = ""
-        system_prompt_expired_documents = ""
-        if valid_documents:
-            system_prompt_valid_documents = (
-                "**The currently uploaded and valid documents are the following**\n"
-            )
-            system_prompt_valid_documents = (
-                system_prompt_valid_documents
-                + "\n".join(
-                    f"- {doc.title or doc.key} (content_id: {doc.id})"
-                    for doc in valid_documents
-                )
-                + "\n"
-            )
-
-        if expired_documents:
-            system_prompt_expired_documents = (
-                "**The currently uploaded and expired documents are the following**\n"
-            )
-            system_prompt_expired_documents = (
-                system_prompt_expired_documents
-                + "\n".join(
-                    f"- {doc.title or doc.key} (content_id: {doc.id})"
-                    for doc in expired_documents
-                )
-            )
-
-        return self._config.tool_description_for_system_prompt.format(
-            system_prompt_valid_documents=system_prompt_valid_documents,
-            system_prompt_expired_documents=system_prompt_expired_documents,
+        return render_template(
+            self._config.tool_description_for_system_prompt,
+            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:
@@ -139,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
@@ -153,7 +153,10 @@ class UploadedSearchTool(Tool[UploadedSearchConfig]):
                 state=ProgressState.FINISHED,
             )
         tool_response.name = self.name
-        tool_response.system_reminder = self._get_tool_call_response_system_reminder()
+        if self._config.enable_tool_call_system_reminder:
+            tool_response.system_reminder = (
+                self._get_tool_call_response_system_reminder()
+            )
         return tool_response
 
     def _get_tool_call_response_system_reminder(self) -> str:
@@ -173,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)