Upgrade models and update prompts (#534)

DavidKoleczek · web-flow · commit 786a089cdc2a · 2025-05-07T18:31:58.000-04:00
diff --git a/assistants/document-assistant/assistant/config.py b/assistants/document-assistant/assistant/config.py
@@ -194,6 +194,14 @@ class PromptsConfigModel(BaseModel):
                 - 🧪 Ask me to conduct research for example, *Find me the latest competitors in the wearables market*
                """).strip()
 
+    knowledge_cutoff: Annotated[
+        str,
+        Field(
+            title="Knowledge Cutoff",
+            description="Knowledge cutoff date for the models being used.",
+        ),
+    ] = "2024-05"
+
 
 class OrchestrationConfigModel(BaseModel):
     hosted_mcp_servers: Annotated[
@@ -278,7 +286,7 @@ class AssistantConfigModel(BaseModel):
     ] = AzureOpenAIClientConfigModel(
         service_config=azure_openai_service_config_construct(default_deployment="gpt-4.1"),
         request_config=OpenAIRequestConfig(
-            max_tokens=128_000,
+            max_tokens=180000,
             response_tokens=16_384,
             model="gpt-4.1",
             is_reasoning_model=False,
@@ -294,11 +302,11 @@ class AssistantConfigModel(BaseModel):
         ),
         UISchema(widget="radio", hide_title=True),
     ] = AzureOpenAIClientConfigModel(
-        service_config=azure_openai_service_config_reasoning_construct(default_deployment="o3-mini"),
+        service_config=azure_openai_service_config_reasoning_construct(default_deployment="o4-mini"),
         request_config=OpenAIRequestConfig(
             max_tokens=200_000,
             response_tokens=65_536,
-            model="o3-mini",
+            model="o4-mini",
             is_reasoning_model=True,
             reasoning_effort="high",
         ),
diff --git a/assistants/document-assistant/assistant/filesystem/__init__.py b/assistants/document-assistant/assistant/filesystem/__init__.py
@@ -1,6 +1,6 @@
 from ._filesystem import AttachmentProcessingErrorHandler, AttachmentsExtension
 from ._model import Attachment, AttachmentsConfigModel, DocumentEditorConfigModel
-from ._prompts import FILES_PROMPT, VIEW_TOOL, VIEW_TOOL_OBJ
+from ._prompts import EDIT_TOOL_DESCRIPTION_HOSTED, EDIT_TOOL_DESCRIPTION_LOCAL, FILES_PROMPT, VIEW_TOOL, VIEW_TOOL_OBJ
 
 __all__ = [
     "AttachmentsExtension",
@@ -11,4 +11,6 @@
     "FILES_PROMPT",
     "VIEW_TOOL",
     "VIEW_TOOL_OBJ",
+    "EDIT_TOOL_DESCRIPTION_HOSTED",
+    "EDIT_TOOL_DESCRIPTION_LOCAL",
 ]
diff --git a/assistants/document-assistant/assistant/filesystem/_prompts.py b/assistants/document-assistant/assistant/filesystem/_prompts.py
@@ -42,3 +42,28 @@
     ),
     type="function",
 )
+
+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. \
+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.
+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.
+
+Args:
+    path: The relative path to the file.
+    task: The specific task that you want the document editor to do."""
+
+EDIT_TOOL_DESCRIPTION_LOCAL = """The user has a file editor corresponding to the file type, open like VSCode, Word, PowerPoint, TeXworks (+ MiKTeX), open side by side with this chat.
+Use this tool to create new files or edit existing ones.
+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.
+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.
+
+Args:
+    path: The relative path to the file.
+    task: The specific task that you want the document editor to do."""
diff --git a/assistants/document-assistant/assistant/guidance/guidance_prompts.py b/assistants/document-assistant/assistant/guidance/guidance_prompts.py
@@ -34,6 +34,7 @@
 since there might be a lot of UI elements generated over time and we want to make sure the user can understand what it is referring to.
 
 ### Current Dynamic UI State
+
 - The current UI schema is after the "ui_elements" key in the JSON object.
 - The current selections (if any) are in the "form_data" key in the JSON object. \
 Note that textboxes are indexed by their order in the list of UI elements, so the first textbox if it has data will be "textbox_0", the second "textbox_1", etc.
diff --git a/assistants/document-assistant/assistant/response/prompts.py b/assistants/document-assistant/assistant/response/prompts.py
@@ -5,6 +5,8 @@
 freeing them up to be more focused on higher level objectives once you understand their needs and goals. \
 One of the core features is a Markdown document editor that will be open side by side whenever a document is opened or edited. \
 They are counting on you, so be creative, guiding, work hard, and find ways to be successful.
+Knowledge cutoff: {{knowledge_cutoff}}
+Current date: {{current_date}}
 
 # On Responding in Chat (Formatting)
 - **Text & Markdown:**
diff --git a/assistants/document-assistant/assistant/response/responder.py b/assistants/document-assistant/assistant/response/responder.py
@@ -6,6 +6,7 @@
 from typing import Any, Callable
 
 import deepmerge
+import pendulum
 from assistant_extensions.mcp import (
     ExtendedCallToolRequestParams,
     MCPClientSettings,
@@ -18,6 +19,7 @@
     refresh_mcp_sessions,
     sampling_message_to_chat_completion_message,
 )
+from liquid import render
 from mcp import SamplingMessage, ServerNotification
 from mcp.types import (
     TextContent,
@@ -47,12 +49,17 @@
 )
 
 from assistant.config import AssistantConfigModel
-from assistant.filesystem import VIEW_TOOL_OBJ, AttachmentsExtension
+from assistant.filesystem import (
+    EDIT_TOOL_DESCRIPTION_HOSTED,
+    EDIT_TOOL_DESCRIPTION_LOCAL,
+    VIEW_TOOL_OBJ,
+    AttachmentsExtension,
+)
+from assistant.filesystem._prompts import FILES_PROMPT
 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
 from assistant.response.models import StepResult
-from assistant.response.prompts import GUARDRAILS_POSTFIX, ORCHESTRATION_SYSTEM_PROMPT
 from assistant.response.utils import get_ai_client_configs, get_completion, get_openai_tools_from_mcp_sessions
 from assistant.response.utils.formatting_utils import format_message
 from assistant.response.utils.message_utils import (
@@ -299,9 +306,19 @@ async def _construct_prompt(self) -> tuple[list, list[ChatCompletionMessageParam
         # Remove any view tool that was added by an MCP server and replace it with ours
         tools = [tool for tool in tools if tool["function"]["name"] != "view"]
         tools.append(VIEW_TOOL_OBJ)
+        # Override the description of the edit_file depending on the environment
+        tools = self._override_edit_file_description(tools)
 
         # Start constructing main system prompt
-        main_system_prompt = ORCHESTRATION_SYSTEM_PROMPT
+        main_system_prompt = self.config.orchestration.prompts.orchestration_prompt
+        # Inject the {{knowledge_cutoff}} and {{current_date}} placeholders
+        main_system_prompt = render(
+            main_system_prompt,
+            **{
+                "knowledge_cutoff": self.config.orchestration.prompts.knowledge_cutoff,
+                "current_date": pendulum.now(tz="America/Los_Angeles").format("YYYY-MM-DD"),
+            },
+        )
 
         # Construct key parts of the system messages which are core capabilities.
         # Best practice is to have these start with a ## <heading content>
@@ -310,23 +327,23 @@ async def _construct_prompt(self) -> tuple[list, list[ChatCompletionMessageParam
             dynamic_ui_system_prompt = self.tokenizer.truncate_str(
                 await self._construct_dynamic_ui_system_prompt(), self.max_system_prompt_component_tokens
             )
-            main_system_prompt += "\n\n" + dynamic_ui_system_prompt
+            main_system_prompt += "\n\n" + dynamic_ui_system_prompt.strip()
 
         # Filesystem System Prompt
         filesystem_system_prompt = self.tokenizer.truncate_str(
             await self._construct_filesystem_system_prompt(), self.max_system_prompt_component_tokens
         )
-        main_system_prompt += "\n\n" + filesystem_system_prompt
+        main_system_prompt += "\n\n" + filesystem_system_prompt.strip()
 
         # Add specific guidance from MCP servers
         mcp_prompts = await get_mcp_server_prompts(self.mcp_sessions)
         mcp_prompt_string = self.tokenizer.truncate_str(
             "## MCP Servers" + "\n\n" + "\n\n".join(mcp_prompts), self.max_system_prompt_component_tokens
         )
-        main_system_prompt += "\n\n" + mcp_prompt_string
+        main_system_prompt += "\n\n" + mcp_prompt_string.strip()
 
         # Always append the guardrails postfix at the end.
-        main_system_prompt += "\n\n" + GUARDRAILS_POSTFIX
+        main_system_prompt += "\n\n" + self.config.orchestration.prompts.guardrails_prompt.strip()
 
         logging.info("The system prompt has been constructed.")
 
@@ -435,7 +452,7 @@ async def _construct_dynamic_ui_system_prompt(self) -> str:
         if not current_dynamic_ui_elements:
             current_dynamic_ui_elements = "No dynamic UI elements have been generated yet. Consider generating some."
 
-        system_prompt = "## On Dynamic UI Elements"
+        system_prompt = "## On Dynamic UI Elements\n"
         system_prompt += "\n" + self.config.orchestration.guidance.prompt
         system_prompt += "\n" + str(current_dynamic_ui_elements)
         return system_prompt
@@ -454,7 +471,7 @@ async def _construct_filesystem_system_prompt(self) -> str:
         all_files.extend([(filename, "-rw-") for filename in doc_editor_filenames])
         all_files.sort(key=lambda x: x[0])
 
-        system_prompt = "## Files\n"
+        system_prompt = f"{FILES_PROMPT}" + "\n\n### Files\n"
         if not all_files:
             system_prompt += "\nNo files have been added or created yet."
         else:
@@ -537,6 +554,39 @@ async def _context_management(
         preserved_messages = initial_messages + recent_messages
         return preserved_messages
 
+    def _override_edit_file_description(self, tools: list[ChatCompletionToolParam]) -> list[ChatCompletionToolParam]:
+        """
+        Override the edit_file description based on the root (the one that indicates the hosted env, otherwise assume the local env).
+        """
+        try:
+            # Get the root of the filesystem-edit tool
+            # Find the filesystem MCP by name
+            filesystem_mcp = next(
+                (mcp for mcp in self.mcp_sessions if mcp.config.server_config.key == "filesystem-edit"),
+                None,
+            )
+            filesystem_root = None
+            if filesystem_mcp:
+                # Get the root of the filesystem-edit tool
+                filesystem_root = next(
+                    (root for root in filesystem_mcp.config.server_config.roots if root.name == "root"),
+                    None,
+                )
+
+            edit_tool = next(
+                (tool for tool in tools if tool["function"]["name"] == "edit_file"),
+                None,
+            )
+            if filesystem_root and filesystem_root.uri == "file://workspace" and edit_tool:
+                edit_tool["function"]["description"] = EDIT_TOOL_DESCRIPTION_HOSTED
+            elif filesystem_root and edit_tool:
+                edit_tool["function"]["description"] = EDIT_TOOL_DESCRIPTION_LOCAL
+        except Exception as e:
+            logger.error(f"Failed to override edit_file description: {e}")
+            return tools
+
+        return tools
+
     # endregion
 
     # region MCP Sessions
diff --git a/assistants/document-assistant/assistant/response/utils/formatting_utils.py b/assistants/document-assistant/assistant/response/utils/formatting_utils.py
@@ -1,6 +1,7 @@
 import logging
 from textwrap import dedent
 
+import pendulum
 from semantic_workbench_api_model.workbench_model import (
     ConversationMessage,
     ConversationParticipant,
@@ -18,7 +19,8 @@ def format_message(message: ConversationMessage, participants: list[Conversation
         None,
     )
     participant_name = conversation_participant.name if conversation_participant else "unknown"
-    message_datetime = message.timestamp.strftime("%Y-%m-%d %H:%M:%S")
+    message_datetime = pendulum.instance(message.timestamp, tz="America/Los_Angeles")
+    message_datetime = message_datetime.format("Y-MM-DD HH:mm:ss")
     return f"[{participant_name} - {message_datetime}]: {message.content}"
 
 
diff --git a/assistants/document-assistant/pyproject.toml b/assistants/document-assistant/pyproject.toml
@@ -16,6 +16,7 @@ dependencies = [
     "openai-client>=0.1.0",
     "pdfplumber>=0.11.2",
     "pendulum>=3.1,<4.0",
+    "python-liquid>=2.0,<3.0",
     "tiktoken>=0.9.0",
 ]
 
diff --git a/assistants/document-assistant/uv.lock b/assistants/document-assistant/uv.lock
diff --git a/libraries/python/mcp-extensions/mcp_extensions/llm/mcp_chat_completion.py b/libraries/python/mcp-extensions/mcp_extensions/llm/mcp_chat_completion.py
@@ -73,11 +73,30 @@ async def mcp_chat_completion(request: ChatCompletionRequest, client: Context) -
     extra_args["response_format"] = response_format
 
     model = request.model
-    # Default to gpt-4o
-    model_preferences = ModelPreferences(intelligencePriority=0, speedPriority=1)
-    if model.startswith("o3"):
+    if model in [
+        "gpt-4o",
+        "gpt-4o-2024-11-20",
+        "gpt-4o-2024-08-06",
+        "gpt-4o-2024-05-13",
+        "gpt-4.1",
+        "gpt-4.1-2025-04-14",
+    ]:
+        model_preferences = ModelPreferences(intelligencePriority=0, speedPriority=1)
+    elif model in [
+        "o3",
+        "o3-2025-04-16",
+        "o3-mini",
+        "o3-mini-2025-01-31",
+        "o4-mini",
+        "o4-mini-2025-04-16",
+        "o1-mini",
+        "o1-mini-2024-09-12",
+    ]:
         model_preferences = ModelPreferences(intelligencePriority=1)
+    else:
+        model_preferences = ModelPreferences(intelligencePriority=0, speedPriority=1)
 
+    # Remove the keys that are not needed for the request
     extra_args.pop("messages", None)
     extra_args.pop("max_completion_tokens", None)
     extra_args.pop("model", None)
diff --git a/mcp-servers/mcp-server-filesystem-edit/mcp_server_filesystem_edit/config.py b/mcp-servers/mcp-server-filesystem-edit/mcp_server_filesystem_edit/config.py
@@ -4,6 +4,7 @@
 import sys
 from typing import Annotated
 
+import pendulum
 from pydantic import Field
 from pydantic_settings import BaseSettings
 
@@ -30,3 +31,16 @@ class Settings(BaseSettings):
     feedback_tool_prefix: str = "[Feedback Tool]: "
     file_tool_prefix: str = "[File Tool]: "
     rewrite_threshold: int = 250  # Tokens above this will trigger document edits instead of a rewrite
+
+    # Model choices. NOTE: When using MCP sampling, the value here != what might actually be used.
+    draft_path_model: str = "gpt-4.1"
+    edit_model: str = "o4-mini"
+    convert_tool_calls_model: str = "gpt-4.1"
+    summarization_model: str = "gpt-4.1"
+    comment_model: str = "o4-mini"
+    comment_analysis_model: str = "gpt-4.1"
+    knowledge_cutoff: str = "2024-05"
+
+    def current_date_func(self) -> str:
+        """Return the current date in Eastern US timezone formatted as YYYY-MM-DD."""
+        return pendulum.now(tz="America/Los_Angeles").format("YYYY-MM-DD")
diff --git a/mcp-servers/mcp-server-filesystem-edit/mcp_server_filesystem_edit/prompts/markdown_draft.py b/mcp-servers/mcp-server-filesystem-edit/mcp_server_filesystem_edit/prompts/markdown_draft.py
@@ -3,8 +3,8 @@
 from mcp_extensions.llm.llm_types import DeveloperMessage, UserMessage
 
 MD_DRAFT_REASONING_DEV_PROMPT = DeveloperMessage(
-    content="""You're a document writer working inside Microsoft Word with PhD-level expertise across multiple domains and are using the provided conversation history as the primary context. \
-You will also be provided context such as attached documents, which may be relevant to the task. \
+    content="""You're a document writer working side by side with a user and an AI assistant with PhD-level expertise across multiple domains and are using the provided conversation history as the primary context. \
+The user may also have provided context such as attached documents, other written documents, or the assistant may have called tools which may have context relevant to the task. \
 If the context is relevant, you must use it accurately and factually.
 Knowledge cutoff: {{knowledge_cutoff}}
 Current date: {{current_date}}
@@ -33,7 +33,7 @@
 
 ### Contextual Awareness
 - Pay attention to and incorporate the preferences the user has subtly provided throughout the conversation.
-- Ensure the writing style of the document aligns with the style of any uploaded documents from the user.
+- Ensure the writing style of the document aligns with the style of any documents that were written by the user.
 - If there are sources, attachments, or results of tool calls, you should use those factually and accurately. \
 Keep in mind that not all sources are quality, so you should use your judgment and own expertise to determine if you should include it.
 
diff --git a/mcp-servers/mcp-server-filesystem-edit/mcp_server_filesystem_edit/prompts/markdown_edit.py b/mcp-servers/mcp-server-filesystem-edit/mcp_server_filesystem_edit/prompts/markdown_edit.py
@@ -3,14 +3,14 @@
 from mcp_extensions.llm.llm_types import DeveloperMessage, SystemMessage, UserMessage
 
 MARKDOWN_EDIT_FORMAT_INSTRUCTIONS = """- Your focus is on writing documents, but you can use standard Markdown syntax for the content. You should use it sparingly.
-- Use headings (levels 1-4) to structure the content.
-- Include relevant dates in the content.
+- Use headings (levels 1-4) to structure the content. Do not add additional numbers or letters to the headings; leave it at something like ## Technical Implementation.
 - Do not insert extra newlines, such as after list items, unless it is already the existing style of the document."""
 
 WORD_EDIT_FORMAT_INSTRUCTIONS = """- Your focus is on writing documents, but you can use a subset of Markdown syntax for the content. You should use it sparingly.
 - The following are rules specific to environment where the content will be rendered.
     - You can ONLY use Markdown syntax for paragraphs with bold and italics, headings, lists (numbered and bulleted, but NON NESTED). All other Markdown syntax is unsupported and forbidden.
-    - You can only use heading levels 1-6.
+    - You can only use heading levels 1-6. \
+Do not add additional numbers or letters to the headings; leave it at something like ## Technical Implementation.
     - When creating numbered lists, you must use the syntax of "1." to start each item in the list.
     - Do NOT nest lists. For example, do not create a bulleted list "inside" of a numbered list.
     - You must use headings to create a hierarchy of information instead of using any form of nested lists.
diff --git a/mcp-servers/mcp-server-filesystem-edit/mcp_server_filesystem_edit/server.py b/mcp-servers/mcp-server-filesystem-edit/mcp_server_filesystem_edit/server.py
diff --git a/mcp-servers/mcp-server-filesystem-edit/mcp_server_filesystem_edit/tools/add_comments.py b/mcp-servers/mcp-server-filesystem-edit/mcp_server_filesystem_edit/tools/add_comments.py
diff --git a/mcp-servers/mcp-server-filesystem-edit/mcp_server_filesystem_edit/tools/edit.py b/mcp-servers/mcp-server-filesystem-edit/mcp_server_filesystem_edit/tools/edit.py

Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@ dependencies = [`
`16`	`16`	`"openai-client>=0.1.0",`
`17`	`17`	`"pdfplumber>=0.11.2",`
`18`	`18`	`"pendulum>=3.1,<4.0",`
	`19`	`+ "python-liquid>=2.0,<3.0",`
`19`	`20`	`"tiktoken>=0.9.0",`
`20`	`21`	`]`
`21`	`22`