From 9931b4ad7adba90b52a113a2baf7e39744685120 Mon Sep 17 00:00:00 2001 From: Michael Panchenko Date: Sat, 29 Mar 2025 14:06:33 +0100 Subject: [PATCH] Onboarding with memory files and thinking tools --- prompts/mcp_prompts.yml | 54 +++++++++-- scripts/gen_prompt_factory.py | 5 +- src/serena/llm/prompt_factory.py | 57 ++++++----- src/serena/mcp.py | 162 ++++++++++++++++++++++++++++++- 4 files changed, 239 insertions(+), 39 deletions(-) diff --git a/prompts/mcp_prompts.yml b/prompts/mcp_prompts.yml index d564052a..812e6193 100644 --- a/prompts/mcp_prompts.yml +++ b/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. + diff --git a/scripts/gen_prompt_factory.py b/scripts/gen_prompt_factory.py index ca67d11a..916cf99c 100644 --- a/scripts/gen_prompt_factory.py +++ b/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__": diff --git a/src/serena/llm/prompt_factory.py b/src/serena/llm/prompt_factory.py index 3c70099e..d4ca4284 100644 --- a/src/serena/llm/prompt_factory.py +++ b/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()) + + diff --git a/src/serena/mcp.py b/src/serena/mcp.py index ce449ba1..664c9bcd 100644 --- a/src/serena/mcp.py +++ b/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,