Refactor file handling prompts and improve error handling in view tool (#624)

Author: DavidKoleczek

assistants/document-assistant/assistant/chat.py

@@ -221,7 +221,9 @@ async def on_message_created(
         await archive_task_queues.enqueue_run(
             context=context,
             attachments=attachments,
-            archive_task_config=ArchiveTaskConfig(chunk_token_count_threshold=30_000),
+            archive_task_config=ArchiveTaskConfig(
+                chunk_token_count_threshold=config.orchestration.prompts.token_window
+            ),
             archive_summarizer=construct_archive_summarizer(
                 service_config=config.generative_ai_fast_client_config.service_config,
                 request_config=config.generative_ai_fast_client_config.request_config,

assistants/document-assistant/assistant/context_management/file_manager.py

@@ -20,19 +20,25 @@
 Files that are read-only are known as "attachments" and are initially appended to user's message at the time they uploaded them. \
 Eventually they might fall out of your context window and you will need to use the `view` tool to read them again if you need it. \
 A summary of the file content has been provided to you to better understand what the file is about.
+There are more files that you can access. First call the `ls` tool to list all files available in the filesystem.
 
 ### Recent & Relevant Files
 
 You can read the following files in again using the `view` tool if they are needed. \
-If they are editable you can also use the `edit_file` tool to edit them."""
+If they are editable you can also use the `edit_file` tool to edit them.
+Paths are mounted at different locations depending on the type of file and you must always use the absolute path to the file, starting with `/` for any path.
+- Editable files are mounted at `/editable_documents/editable_file.md`.
+- User uploaded files or "attachments" are mounted at `/attachments/attachment.pdf`."""
 
-FILESYSTEM_ADDON_PROMPT = """### Filesystem
 
-**Very important:** This current interaction with the user is long-running and due to context window limitations, the above section can only show a limited number of files. \
-There are more files that you can access. First call the `ls` tool to list all files available in the filesystem. \
-Then, you can use the `view` tool (use it multiple times if needed) to read any of the files that you find relevant to the user's request.\
-This is a similar concept to how you would explore a codebase in a code editor."""
+ARCHIVES_ADDON_PROMPT = """### Conversation Memories and Archives
 
+You have a limited context window, which means that some of the earlier parts of the conversation may fall out of your context.
+To help you with that, below you will find summaries of older parts of the conversation that have been "archived". \
+You should use these summaries as "memories" to help you understand the historical context and preferences of the user. \
+Note that some of these archived conversation may still be visible to you in the conversation history.
+If the current user's task requires you to access the full content of the conversation, you can use the `view` tool to read the archived conversations. \
+Historical conversations are mounted at `/archives/conversation_1234567890.json`"""
 
 VIEW_TOOL = {
     "type": "function",
@@ -45,7 +51,7 @@
             "properties": {
                 "path": {
                     "type": "string",
-                    "description": "The relative path to the file.",
+                    "description": "The absolute path to the file. Must start with `/` followed by the mount point, e.g. `/editable_documents/filename.md`.",
                 },
             },
             "required": ["path"],
@@ -90,8 +96,9 @@
 EDIT_TOOL_DESCRIPTION_HOSTED = """Edits the Markdown file at the provided path, focused on the given task.
 The user has Markdown editor available that is side by side with this chat.
 Remember that the editable files are the ones that have the `-rw-` permission bits. \
+They also must be mounted at `/editable_documents/` and have a `.md` extension. \
 If you provide a new file path, it will be created for you and then the editor will start to edit it (from scratch). \
-Name the file with capital letters and spacing like "Weekly AI Report.md" or "Email to Boss.md" since it will be directly shown to the user in that way.
+Name the file with capital letters and spacing like "/editable_documents/Weekly AI Report.md" or "/editable_documents/Email to Boss.md" since it will be directly shown to the user in that way.
 Provide a task that you want it to do in the document. For example, if you want to have it expand on one section, \
 you can say "expand on the section about <topic x>". The task should be at most a few sentences. \
 Do not provide it any additional context outside of the task parameter. It will automatically be fetched as needed by this tool.

assistants/document-assistant/assistant/filesystem/_prompts.py

@@ -10,7 +10,6 @@
 from assistant_extensions.mcp import (
     ExtendedCallToolRequestParams,
     MCPSession,
-    OpenAISamplingHandler,
     handle_mcp_tool_call,
 )
 from chat_context_toolkit.virtual_filesystem import VirtualFileSystem
@@ -29,7 +28,6 @@
     ConversationContext,
 )
 
-from assistant.filesystem import AttachmentsExtension
 from assistant.guidance.dynamic_ui_inspector import update_dynamic_ui_state
 from assistant.guidance.guidance_prompts import DYNAMIC_UI_TOOL_NAME, DYNAMIC_UI_TOOL_RESULT
 

assistants/document-assistant/assistant/response/completion_handler.py

@@ -69,7 +69,7 @@
 )
 from assistant.filesystem._file_sources import attachments_file_source_mount, editable_documents_file_source_mount
 from assistant.filesystem._filesystem import _files_drive_for_context
-from assistant.filesystem._prompts import FILES_PROMPT, LS_TOOL_OBJ
+from assistant.filesystem._prompts import ARCHIVES_ADDON_PROMPT, FILES_PROMPT, LS_TOOL_OBJ
 from assistant.guidance.dynamic_ui_inspector import get_dynamic_ui_state, update_dynamic_ui_state
 from assistant.guidance.guidance_prompts import DYNAMIC_UI_TOOL_NAME, DYNAMIC_UI_TOOL_OBJ
 from assistant.response.completion_handler import handle_completion
@@ -452,32 +452,43 @@ async def _construct_dynamic_ui_system_prompt(self) -> str:
         return system_prompt
 
     async def _construct_filesystem_system_prompt(self) -> str:
-        """Constructs the files available to the assistant that are out of context.
+        """Constructs the filesystem system prompt with available files.
+
+        Builds a system prompt that includes:
+        1. FILES_PROMPT with attachments and editable_documents (up to 25 files)
+        2. ARCHIVES_ADDON_PROMPT (if archives exist)
+        3. Archives files listing (up to 25 files)
+
+        Files are sorted by timestamp (newest first), limited to 25 per category,
+        then sorted alphabetically by path.
 
         This is an example of what gets added after the FILES_PROMPT:
         -r-- path2.pdf [File content summary: <summary>]
         -rw- path3.txt [File content summary: No summary available yet, use the context available to determine the use of this file]
         """
-        entries = (
-            list(await self.virtual_filesystem.list_directory(path="/attachments"))
-            + list(await self.virtual_filesystem.list_directory(path="/editable_documents"))
-            + list(await self.virtual_filesystem.list_directory(path="/archives"))
-        )
-        files = [entry for entry in entries if isinstance(entry, FileEntry)]
+        # Get all file entries
+        attachments_entries = list(await self.virtual_filesystem.list_directory(path="/attachments"))
+        editable_documents_entries = list(await self.virtual_filesystem.list_directory(path="/editable_documents"))
+        archives_entries = list(await self.virtual_filesystem.list_directory(path="/archives"))
+        
+        # Separate regular files from archives
+        regular_files = [entry for entry in (attachments_entries + editable_documents_entries) if isinstance(entry, FileEntry)]
+        archives_files = [entry for entry in archives_entries if isinstance(entry, FileEntry)]
 
         # TODO: Better ranking algorithm
-        # order the files by timestamp, newest first
-        files.sort(key=lambda f: f.timestamp, reverse=True)
-        # take the top 25 files
-        files = files[:25]
+        # order the regular files by timestamp, newest first
+        regular_files.sort(key=lambda f: f.timestamp, reverse=True)
+        # take the top 25 regular files
+        regular_files = regular_files[:25]
         # order them alphabetically by path
-        files.sort(key=lambda f: f.path.lower())
+        regular_files.sort(key=lambda f: f.path.lower())
 
+        # Start with FILES_PROMPT and add attachments/editable_documents
         system_prompt = FILES_PROMPT + "\n"
-        if not files:
+        if not regular_files:
             system_prompt += "\nNo files are currently available."
 
-        for file in files:
+        for file in regular_files:
             # Format permissions: -rw- for read_write, -r-- for read
             permissions = "-rw-" if file.permission == "read_write" else "-r--"
             # Use the file description as the summary, or provide a default message
@@ -487,6 +498,29 @@ async def _construct_filesystem_system_prompt(self) -> str:
                 else "No summary available yet, use the context available to determine the use of this file"
             )
             system_prompt += f"{permissions} {file.path} [File content summary: {summary}]\n"
+
+        # Add ARCHIVES_ADDON_PROMPT if there are archives
+        if archives_files:
+            system_prompt += "\n" + ARCHIVES_ADDON_PROMPT + "\n"
+            
+            # order the archives files by timestamp, newest first
+            archives_files.sort(key=lambda f: f.timestamp, reverse=True)
+            # take the top 25 archives files
+            archives_files = archives_files[:25]
+            # order them alphabetically by path
+            archives_files.sort(key=lambda f: f.path.lower())
+            
+            for file in archives_files:
+                # Format permissions: -rw- for read_write, -r-- for read
+                permissions = "-rw-" if file.permission == "read_write" else "-r--"
+                # Use the file description as the summary, or provide a default message
+                summary = (
+                    file.description
+                    if file.description
+                    else "No summary available yet, use the context available to determine the use of this file"
+                )
+                system_prompt += f"{permissions} {file.path} [File content summary: {summary}]\n"
+
         return system_prompt
 
     def _override_edit_file_description(self, tools: list[ChatCompletionToolParam]) -> list[ChatCompletionToolParam]:

assistants/document-assistant/assistant/response/responder.py

@@ -47,7 +47,13 @@ def mounts(self) -> Iterable[MountPoint]:
         return self._mounts.values()
 
     def _split_path(self, path: str) -> tuple[str, str]:
-        _, mount_path_segment, *source_path_segments = path.split("/")
+        path_segments = path.split("/")
+        if len(path_segments) < 2:
+            raise ValueError(
+                f"Invalid path format: {path}. Path must start with '/' and contain at least one directory."
+            )
+
+        _, mount_path_segment, *source_path_segments = path_segments
         mount_path = "/" + mount_path_segment
         source_path = "/" + "/".join(source_path_segments)
         return mount_path, source_path

libraries/python/chat-context-toolkit/chat_context_toolkit/virtual_filesystem/_virtual_filesystem.py

@@ -54,6 +54,8 @@ async def execute(self, args: dict) -> str | Iterable[ChatCompletionContentPartT
             file_content = await self.virtual_filesystem.read_file(path)
         except FileNotFoundError:
             return f"Error: File at path {path} not found. Please pay attention to the available files and try again."
+        except ValueError as e:
+            return f"Error: {str(e)}"
 
         result = f'<file path="{path}">\n{file_content}\n</file>'
         return result

libraries/python/chat-context-toolkit/chat_context_toolkit/virtual_filesystem/tools/_view_tool.py

@@ -29,3 +29,18 @@ async def test_view_tool_error_handling():
         result
         == "Error: File at path /nonexistent.txt not found. Please pay attention to the available files and try again."
     )
+
+
+async def test_view_tool_invalid_path_format():
+    """Test ViewTool.execute with invalid path formats."""
+    mock_vfs = MagicMock(spec=VirtualFileSystem)
+    mock_vfs.read_file.side_effect = ValueError(
+        "Invalid path format: Bair Haiku.md. Path must start with '/' and contain at least one directory."
+    )
+
+    view_tool = ViewTool(mock_vfs)
+    result = await view_tool.execute({"path": "Bair Haiku.md"})
+    assert (
+        result
+        == "Error: Invalid path format: Bair Haiku.md. Path must start with '/' and contain at least one directory."
+    )