Onboarding with memory files and thinking tools

prompts/mcp_prompts.yml

@@ -1,12 +1,48 @@
 lang: en
 prompts:
   onboarding_prompt: |
-    You are to assist in generating overview documentation for onboarding developers of the project.
-    Specifically, you are to create the file `{{ onboarding_file }}` in the project root directory.
-    The file should contain 
-      * a brief description of the project, including the project's purpose, 
-      * an overview of the structure of the project,
-      * instructions on how to add and execute new automated tests.
-    Use functions for reading the project's directory structure and individual files to accomplish this.
-    Read only the necessary files and directories to generate the documentation.
-    You are to provide the content for the file in markdown format.
+    You are viewing the project for the first time.
+    Your task is to assemble relevant high-level information about the project which
+    will be saved to memory files in the following steps.
+    The information should be sufficient to understand what the project is about,
+    and the most important commands for developing code.
+    The project is being developed on the system: {{ system }}.
+
+    You need to identify at least the following information:
+    * the project's purpose
+    * the tech stack used
+    * the rough structure of the codebase
+    * the commands for testing, formatting, and linting
+    * the commands for running the entrypoints of the project
+    * the util commands for the system, like `git`, `ls`, `cd`, `grep`, `find`, etc. Keep in mind that the system is {{ system }},
+      so the commands might be different than on a regular unix system.
+    * whether there are particular guidelines, styles, design patterns, etc. that one should know about
+    
+    This list is not exhaustive, you can add more information if you think it is relevant.
+
+    For doing that, you will need to acquire information about the project with the corresponding tools.
+    Read only the necessary files and directories to avoid loading too much data into memory.
+    If you cannot find everything you need from the project itself, you should ask the user for more information.
+
+    After collecting all the information, use the `write_memory` tool (in multiple calls) to save it to various memory files.
+    A particularly important memory file will be the `suggested_commands.md` file, which should contain
+    a list of commands that the user should know about to develop code in this project.
+
+  think_about_collected_information: |
+    Have you collected all the information you need for solving the current task? If not, can the missing information be acquired by using the available tools,
+    in particular the tools related to symbol discovery? Or do you need to ask the user for more information?
+    Think about it step by step and give a summary of the missing information and how it could be acquired.
+
+  think_about_task_adherence: |
+    Are you deviating from the task at hand? Do you need any additional information to proceed?
+    Note that it is better to stop and ask the user for clarification
+    than to perform large changes which might not be aligned with the user's intentions.
+    If you feel like the conversation is deviating too much from the original task, apologize and suggest to the user
+    how to proceed. If the conversation became too long, create a summary of the current progress and suggest to the user
+    to start a new conversation based on that summary.
+
+  think_about_whether_you_are_done: |
+    Have you already performed all the steps required by the task? Is it appropriate to run tests and linting, and if so,
+    have you done that already? Should new tests be written to cover the changes?
+    Note that a task that is just about exploring the codebase does not require running tests or linting.
+  

scripts/gen_prompt_factory.py

@@ -35,8 +35,11 @@ def main():
         code = f.read()
     methods_str = "".join(methods)
     code = code.replace("# methods", methods_str)
-    with open(f"{package}/llm/prompt_factory.py", "w") as f:
+    
+    prompt_factory_module = f"{package}/llm/prompt_factory.py"
+    with open(prompt_factory_module, "w") as f:
         f.write(code)
+    log.info(f"Prompt factory generated successfully in {prompt_factory_module}")
 
 
 if __name__ == "__main__":

src/serena/llm/prompt_factory.py

@@ -1,24 +1,33 @@
-from .multilang_prompt import MultiLangContainer, MultiLangPromptTemplateCollection, PromptList
-
-
-class PromptFactory:
-    # NOTE: This class is auto-generated by gen_prompt_factory.py
-
-    def __init__(
-        self, lang_shortcode: str = "en", fallback_mode: MultiLangContainer.FallbackMode = MultiLangContainer.FallbackMode.EXCEPTION
-    ):
-        self.lang_shortcode = lang_shortcode
-        self.collection = MultiLangPromptTemplateCollection()
-        self.fallback_mode = fallback_mode
-
-    def _format_prompt(self, prompt_name: str, kwargs: dict) -> str:
-        del kwargs["self"]
-        mpt = self.collection.get_multilang_prompt_template(prompt_name)
-        return mpt.get_item(self.lang_shortcode, self.fallback_mode).instantiate(**kwargs)
-
-    def _get_list(self, prompt_name: str) -> PromptList:
-        mpl = self.collection.get_multilang_prompt_list(prompt_name)
-        return mpl.get_item(self.lang_shortcode, self.fallback_mode)
-
-    def create_onboarding_prompt(self, *, onboarding_file: str) -> str:
-        return self._format_prompt("onboarding_prompt", locals())
+from .multilang_prompt import MultiLangContainer, MultiLangPromptTemplateCollection, PromptList
+
+
+class PromptFactory:
+    # NOTE: This class is auto-generated by gen_prompt_factory.py
+
+    def __init__(self, lang_shortcode: str = "en", fallback_mode=MultiLangContainer.FallbackMode.EXCEPTION):
+        self.lang_shortcode = lang_shortcode
+        self.collection = MultiLangPromptTemplateCollection()
+        self.fallback_mode = fallback_mode
+
+    def _format_prompt(self, prompt_name: str, kwargs) -> str:
+        del kwargs["self"]
+        mpt = self.collection.get_multilang_prompt_template(prompt_name)
+        return mpt.get_item(self.lang_shortcode, self.fallback_mode).instantiate(**kwargs)
+
+    def _get_list(self, prompt_name: str) -> PromptList:
+        mpl = self.collection.get_multilang_prompt_list(prompt_name)
+        return mpl.get_item(self.lang_shortcode, self.fallback_mode)
+
+    def create_onboarding_prompt(self, *, system) -> str:
+        return self._format_prompt('onboarding_prompt', locals())
+
+    def create_think_about_collected_information(self) -> str:
+        return self._format_prompt('think_about_collected_information', locals())
+
+    def create_think_about_task_adherence(self) -> str:
+        return self._format_prompt('think_about_task_adherence', locals())
+
+    def create_think_about_whether_you_are_done(self) -> str:
+        return self._format_prompt('think_about_whether_you_are_done', locals())
+
+    

src/serena/mcp.py

@@ -4,6 +4,8 @@ The Serena Model Context Protocol (MCP) Server
 
 import json
 import os
+from pathlib import Path
+import platform
 import sys
 import traceback
 from abc import ABC, abstractmethod
@@ -47,7 +49,10 @@ class SerenaMCPRequestContext:
     project_root: str
     project_config: dict[str, Any]
     prompt_factory: PromptFactory
-
+    
+    def get_serena_managed_dir(self) -> str:
+        return os.path.join(self.project_root, ".serena")
+        
 
 @asynccontextmanager
 async def server_lifespan(mcp_server: FastMCP) -> AsyncIterator[SerenaMCPRequestContext]:
@@ -68,7 +73,7 @@ async def server_lifespan(mcp_server: FastMCP) -> AsyncIterator[SerenaMCPRequest
     with open(project_file, encoding="utf-8") as f:
         project_config = yaml.safe_load(f)
     language = Language(project_config["language"])
-    project_root = project_config["project_root"]
+    project_root = str(Path(project_config["project_root"]).resolve())
 
     # create and start the language server instance
     config = MultilspyConfig(code_language=language)
@@ -81,7 +86,35 @@ async def server_lifespan(mcp_server: FastMCP) -> AsyncIterator[SerenaMCPRequest
         )
     finally:
         language_server.stop()
-
+        
+        
+class MemoriesManager:
+    def __init__(self, memory_dir: str):
+        self._memory_dir = Path(memory_dir)
+        
+    def _get_memory_file_path(self, memory_file_name: str) -> Path:
+        return self._memory_dir / memory_file_name
+        
+    def load_memory(self, memory_file_name: str) -> str:
+        memory_file_path = self._get_memory_file_path(memory_file_name)
+        if not memory_file_path.exists():
+            return f"Memory file {memory_file_name} not found, consider creating it with the `write_memory` tool if you need it."
+        with open(memory_file_path, "r", encoding="utf-8") as f:
+            return f.read()
+        
+    def save_memory(self, memory_file_name: str, content: str) -> str:
+        memory_file_path = self._get_memory_file_path(memory_file_name)
+        with open(memory_file_path, "w", encoding="utf-8") as f:
+            f.write(content)
+        return f"Memory file {memory_file_name} written."
+    
+    def list_memories(self) -> list[str]:   
+        return [f.name for f in self._memory_dir.iterdir() if f.is_file()]
+    
+    def delete_memory(self, memory_file_name: str) -> str:
+        memory_file_path = self._get_memory_file_path(memory_file_name)
+        memory_file_path.unlink()
+        return f"Memory file {memory_file_name} deleted."
 
 mcp_settings = Settings(lifespan=server_lifespan)
 mcp = FastMCP(**mcp_settings.model_dump())
@@ -94,6 +127,9 @@ class Component(ABC):
         self.project_root = lifespan_context.project_root
         self.project_config = lifespan_context.project_config
         self.prompt_factory = lifespan_context.prompt_factory
+        
+        memories_dir = os.path.join(lifespan_context.get_serena_managed_dir(), "memories")
+        self.memories_manager = MemoriesManager(memories_dir)
 
 
 _DEFAULT_MAX_ANSWER_LENGTH = int(2e5)
@@ -462,21 +498,137 @@ def insert_before_symbol(
     return AppendAfterSymbolTool(ctx).execute()
 
 
+@mcp.tool()
+def check_onboarding_performed(ctx: Context) -> str:
+    """
+    Check if onboarding was performed yet.
+    You should always call this tool in the beginning of the conversation, 
+    before any question about code or the project is asked.
+    You will call this tool only once per conversation.
+    """
+    if len(list_memories(ctx)) == 0:
+        return "Onboarding not performed yet (no memories available). " + \
+            "You should perform onboarding by calling the `onboarding` tool before proceeding with the task."
+    else:
+        return "Onboarding already performed, no need to perform it again."
+
+
 @mcp.tool()
 def onboarding(ctx: Context) -> str:
     """
+    Call this tool if onboarding was not performed yet.
+    You will call this tool at most once per conversation.
+    
     :param ctx: the context object, which will be created and provided automatically
     :return: instructions on how to create the onboarding information
     """
-    onboarding_file = "serena_onboarding.md"
+    system = platform.system()
 
     class OnboardingPrompt(SimplePrompt):
         def _create_prompt(self) -> str:
-            return self.prompt_factory.create_onboarding_prompt(onboarding_file=onboarding_file)
+            return self.prompt_factory.create_onboarding_prompt(system=system)
 
     return OnboardingPrompt(ctx).create()
 
 
+@mcp.tool()
+def write_memory(ctx: Context, memory_file_name: str, content: str) -> str:
+    """
+    Write some general information about this project that can be useful for future tasks to a memory file.
+    The information should be short and to the point.
+    The memory file name should be meaningful, such that from the name you can infer what the information is about.
+    It is better to have multiple small memory files than to have a single large one because
+    memories will be read one by one and we only ever want to read relevant memories.
+    
+    This tool is either called during the onboarding process or when you have identified
+    something worth remembering about the project from the past conversation.
+    """
+    class WriteMemoryTool(Tool):
+        def _execute(self) -> str:
+            return self.memories_manager.save_memory(memory_file_name, content)
+
+    return WriteMemoryTool(ctx).execute()
+
+
+@mcp.tool()
+def read_memory(ctx: Context, memory_file_name: str) -> str:
+    """
+    Read the content of a memory file. This tool should only be used if the information 
+    is relevant to the current task. You should be able to infer whether the information
+    is relevant from the memory file name.
+    You should not read the same memory file multiple times in the same conversation.
+    """
+    class ReadMemoryTool(Tool):
+        def _execute(self) -> str:
+            return self.memories_manager.load_memory(memory_file_name)
+        
+    return ReadMemoryTool(ctx).execute()
+
+
+@mcp.tool()
+def list_memories(ctx: Context) -> str:
+    """
+    List available memories. Any memory can be read using the `read_memory` tool.
+    """ 
+    class ListMemoriesTool(Tool):
+        def _execute(self) -> str:
+            return json.dumps(self.memories_manager.list_memories())
+
+    return ListMemoriesTool(ctx).execute()
+
+
+@mcp.tool()
+def delete_memory(ctx: Context, memory_file_name: str) -> str:
+    """
+    Delete a memory file. Should only happen if a user asks for it explicitly,
+    for example by saying that the information retrieved from a memory file is no longer correct
+    or no longer relevant for the project.
+    """
+    class DeleteMemoryTool(Tool):
+        def _execute(self) -> str:
+            return self.memories_manager.delete_memory(memory_file_name)
+
+    return DeleteMemoryTool(ctx).execute()
+
+
+@mcp.tool()
+def think_about_collected_information(ctx: Context) -> str:
+    """
+    Think about the collected information and whether it is sufficient and relevant.
+    """
+    class ThinkAboutCollectedInformationTool(Tool):
+        def _execute(self) -> str:
+            return self.prompt_factory.create_think_about_collected_information()
+        
+    return ThinkAboutCollectedInformationTool(ctx).execute()
+
+
+@mcp.tool()
+def think_about_task_adherence(ctx: Context) -> str:
+    """
+    Think about the task at hand and whether you are still on track.
+    Especially important if the conversation has been going on for a while and there
+    has been a lot of back and forth.
+    """
+    class ThinkAboutTaskAdherenceTool(Tool):
+        def _execute(self) -> str:
+            return self.prompt_factory.create_think_about_task_adherence()
+
+    return ThinkAboutTaskAdherenceTool(ctx).execute()
+
+
+@mcp.tool()
+def think_about_whether_you_are_done(ctx: Context) -> str:
+    """
+    Think about whether you are done with the task.
+    """
+    class ThinkAboutWhetherYouAreDoneTool(Tool):
+        def _execute(self) -> str:
+            return self.prompt_factory.create_think_about_whether_you_are_done()
+
+    return ThinkAboutWhetherYouAreDoneTool(ctx).execute()
+
+
 @mcp.tool()
 def search_files_for_pattern(
     ctx: Context,