diff --git a/code_templates/prompt_factory_template.py b/code_templates/prompt_factory_template.py new file mode 100644 index 00000000..81fe7701 --- /dev/null +++ b/code_templates/prompt_factory_template.py @@ -0,0 +1,21 @@ +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) + + # methods diff --git a/myproject.demo.yml b/myproject.demo.yml index 38198798..1ef3bfac 100644 --- a/myproject.demo.yml +++ b/myproject.demo.yml @@ -2,3 +2,5 @@ project_root: /path/to/project # language of the project (csharp, python, rust, java, typescript, javascript, go, or ruby) language: python +# list of directories to ignore: either names (e.g. "temp") or relative paths (e.g. "build/foo") +ignored_dirs: [".git", "temp"] diff --git a/prompts/mcp_prompts.yml b/prompts/mcp_prompts.yml new file mode 100644 index 00000000..d564052a --- /dev/null +++ b/prompts/mcp_prompts.yml @@ -0,0 +1,12 @@ +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. diff --git a/pyproject.toml b/pyproject.toml index f978614e..f1d20608 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ "python-dotenv>=1.0.0, <2", "mcp>=1.5.0", "fastmcp>=0.4.1", + "sensai-utils>=1.4.0", ] [project.license] @@ -230,6 +231,10 @@ ignore = [ "C401", "C901", "ASYNC230", + "ISC003", + "B024", + "B007", + "SIM102", ] unfixable = [ "F841", diff --git a/scripts/gen_prompt_factory.py b/scripts/gen_prompt_factory.py new file mode 100644 index 00000000..ca67d11a --- /dev/null +++ b/scripts/gen_prompt_factory.py @@ -0,0 +1,43 @@ +from sensai.util import logging + +from serena.llm.multilang_prompt import MultiLangPromptTemplateCollection + +log = logging.getLogger(__name__) + + +def main(): + coll = MultiLangPromptTemplateCollection() + + package = "src/serena" + + # collect methods to generate + indent = " " + methods = [] + for mpt in coll.prompt_templates.values(): + prompt_name = mpt.name + params = mpt.get_parameters() + if len(params) == 0: + params_str = "" + else: + params_str = ", *, " + ", ".join(params) + methods.append( + f"def create_{prompt_name}(self{params_str}) -> str:" + + f"\n{indent}{indent}return self._format_prompt('{prompt_name}', locals())\n\n{indent}" + ) + for mpl in coll.prompt_lists.values(): + prompt_name = mpl.name + methods.append( + f"def get_list_{prompt_name}(self) -> PromptList:" + f"\n{indent}{indent}return self._get_list('{prompt_name}')\n\n{indent}" + ) + + # write prompt factory with added methods + with open("code_templates/prompt_factory_template.py") as f: + code = f.read() + methods_str = "".join(methods) + code = code.replace("# methods", methods_str) + with open(f"{package}/llm/prompt_factory.py", "w") as f: + f.write(code) + + +if __name__ == "__main__": + logging.run_main(main) diff --git a/scripts/mcp_server.py b/scripts/mcp_server.py index 62917076..1d0b3323 100644 --- a/scripts/mcp_server.py +++ b/scripts/mcp_server.py @@ -2,21 +2,25 @@ The Serena Model Context Protocol (MCP) Server """ +import json import os import sys -from abc import abstractmethod +from abc import ABC, abstractmethod from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import cast +from typing import Any, cast import yaml +from mcp.server.fastmcp.prompts.base import Message, UserMessage from mcp.server.fastmcp.server import Context, FastMCP, Settings from sensai.util import logging from multilspy import SyncLanguageServer from multilspy.multilspy_config import Language, MultilspyConfig from multilspy.multilspy_logger import MultilspyLogger +from serena.file_utils import scan_directory +from serena.llm.prompt_factory import PromptFactory log = logging.getLogger(__name__) @@ -24,6 +28,9 @@ log = logging.getLogger(__name__) @dataclass class SerenaMCPRequestContext: language_server: SyncLanguageServer + project_root: str + project_config: dict[str, Any] + prompt_factory: PromptFactory @asynccontextmanager @@ -41,9 +48,9 @@ async def server_lifespan(mcp_server: FastMCP) -> AsyncIterator[SerenaMCPRequest # read project configuration with open(project_file, encoding="utf-8") as f: - project = yaml.safe_load(f) - language = Language(project["language"]) - project_root = project["project_root"] + project_config = yaml.safe_load(f) + language = Language(project_config["language"]) + project_root = project_config["project_root"] # create and start the language server instance config = MultilspyConfig(code_language=language) @@ -51,7 +58,9 @@ async def server_lifespan(mcp_server: FastMCP) -> AsyncIterator[SerenaMCPRequest language_server = SyncLanguageServer.create(config, logger, project_root) language_server.start() try: - yield SerenaMCPRequestContext(language_server=language_server) + yield SerenaMCPRequestContext( + language_server=language_server, project_root=project_root, project_config=project_config, prompt_factory=PromptFactory() + ) finally: language_server.stop() @@ -60,11 +69,16 @@ mcp_settings = Settings(lifespan=server_lifespan) mcp = FastMCP(**mcp_settings.model_dump()) -class Tool: +class Component(ABC): def __init__(self, ctx: Context): lifespan_context = cast(SerenaMCPRequestContext, ctx.request_context.lifespan_context) self.langsrv = lifespan_context.language_server + self.project_root = lifespan_context.project_root + self.project_config = lifespan_context.project_config + self.prompt_factory = lifespan_context.prompt_factory + +class Tool(Component): def execute(self) -> str: try: return self._execute() @@ -73,12 +87,47 @@ class Tool: return f"Error executing tool: {e}" @abstractmethod - def _execute(self): + def _execute(self) -> str: pass +class SimplePrompt(Component): + def __init__(self, ctx: Context): + super().__init__(ctx) + self.messages = [] + + def create(self) -> str: + return self._create_prompt() + + @abstractmethod + def _create_prompt(self) -> str: + pass + + +class SequentialPrompt(Component): + def __init__(self, ctx: Context): + super().__init__(ctx) + self.messages = [] + + def create(self) -> list[Message]: + self._add_messages() + return self.messages + + @abstractmethod + def _add_messages(self) -> None: + pass + + def _add_user_message(self, msg: str) -> None: + self.messages.append(UserMessage(content=msg)) + + @mcp.tool() def read_file(ctx: Context, relative_path: str) -> str: + """ + :param ctx: the context object, which will be created and provided automatically + :param relative_path: the relative path to the file to read + :return: the full text of the file at the given relative path + """ log.info(f"read_file: {relative_path=}") class ReadFileTool(Tool): @@ -89,6 +138,41 @@ def read_file(ctx: Context, relative_path: str) -> str: return ReadFileTool(ctx).execute() +@mcp.tool() +def list_dir(ctx: Context, relative_path: str, recursive: bool) -> str: + """ + :param ctx: the context object, which will be created and provided automatically + :param relative_path: the relative path to the directory to list; pass "." to scan the project root + :param recursive: whether to scan subdirectories recursively + :return: a JSON object with the names of directories and files within the given directory + """ + log.info(f"list_dir: {relative_path=}") + + class ListDirTool(Tool): + def _execute(self): + dirs, files = scan_directory( + os.path.join(self.project_root, relative_path), recursive=recursive, ignored_dirs=self.project_config["ignored_dirs"] + ) + return json.dumps({"dirs": dirs, "files": files}) + + return ListDirTool(ctx).execute() + + +@mcp.tool() +def onboarding(ctx: Context) -> str: + """ + :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" + + class OnboardingPrompt(SimplePrompt): + def _create_prompt(self): + return self.prompt_factory.create_onboarding_prompt(onboarding_file=onboarding_file) + + return OnboardingPrompt(ctx).create() + + if __name__ == "__main__": logging.basicConfig( level=logging.DEBUG, stream=sys.stderr, format="%(levelname)-5s %(asctime)-15s %(name)s:%(funcName)s:%(lineno)d - %(message)s" diff --git a/src/serena/llm/jinja_template.py b/src/serena/llm/jinja_template.py new file mode 100644 index 00000000..cc70b1b3 --- /dev/null +++ b/src/serena/llm/jinja_template.py @@ -0,0 +1,30 @@ +import jinja2 +import jinja2.meta +import jinja2.nodes +import jinja2.visitor + +from serena.util.class_decorators import singleton + + +@singleton +class JinjaEnvProvider: + def __init__(self): + self._env = None + + def get_env(self): + if self._env is None: + self._env = jinja2.Environment() + return self._env + + +class JinjaTemplate: + def __init__(self, template_string: str): + self._template_string = template_string + self._template = JinjaEnvProvider().get_env().from_string(self._template_string) + + def render(self, **kwargs) -> str: + return self._template.render(**kwargs) + + def get_parameters(self) -> set[str]: + parsed_content = self._template.environment.parse(self._template_string) + return jinja2.meta.find_undeclared_variables(parsed_content) diff --git a/src/serena/llm/multilang_prompt.py b/src/serena/llm/multilang_prompt.py new file mode 100644 index 00000000..f4395316 --- /dev/null +++ b/src/serena/llm/multilang_prompt.py @@ -0,0 +1,203 @@ +import os +from collections.abc import Callable +from enum import Enum +from typing import Any, Generic, TypeVar + +import yaml +from sensai.util.string import ToStringMixin + +from .jinja_template import JinjaTemplate + +LANG_CODES = ["en", "de"] + + +class PromptTemplate(ToStringMixin): + def __init__(self, name: str, jinja_template_string: str): + self.name = name + self.jinja_template = JinjaTemplate(jinja_template_string.strip()) + self.parameters = self.jinja_template.get_parameters() + + def _tostring_excludes(self) -> list[str]: + return ["jinja_template"] + + def instantiate(self, **kwargs) -> str: + return self.jinja_template.render(**kwargs) + + +class PromptList: + def __init__(self, items: list[str]): + self.items = [x.strip() for x in items] + + def to_string(self): + bullet = " * " + indent = " " * len(bullet) + items = [x.replace("\n", "\n" + indent) for x in self.items] + return "\n * ".join(items) + + +T = TypeVar("T") + + +class MultiLangContainer(Generic[T], ToStringMixin): + """ + Represents a container of items which are associated with different languages + """ + + def __init__(self, name: str): + self.name = name + self.lang2item: dict[str, T] = {} + + def _tostring_excludes(self) -> list[str]: + return ["lang2item"] + + def _tostring_additional_entries(self) -> dict[str, Any]: + return dict(languages=list(self.lang2item.keys())) + + class FallbackMode(Enum): + ANY = "any" + """ + If the requested language is not found, return the item for any language + """ + EXCEPTION = "exception" + """ + If the requested language is not found, raise an exception + """ + + def add_item(self, item: T, lang: str = ""): + self.lang2item[lang] = item + + def get_item(self, lang: str, fallback_mode: FallbackMode = FallbackMode.EXCEPTION) -> T: + """ + Gets the item for the given language. + + :param lang: the language shortcode for which to obtain the prompt template + :param fallback_mode: defines what to do if there is no item for the given language + :return: the item + """ + item = self.lang2item.get(lang) + if item is not None: + return item + else: + if fallback_mode == self.FallbackMode.EXCEPTION: + raise KeyError(f"Item for language '{lang}' not found for name '{self.name}'") + if fallback_mode == self.FallbackMode.ANY: + return next(iter(self.lang2item.values())) + else: + raise ValueError(f"Invalid fallback mode '{fallback_mode}'") + + +class MultiLangPromptTemplate(MultiLangContainer[PromptTemplate]): + """ + Represents a prompt template which potentially supports multiple languages. + """ + + def get_parameters(self) -> list[str]: + """ + Gets the parameters of the prompt template (which must be the same for all languages). + + :return: the list of parameters required to instantiate the template + """ + prev_params = None + for template in self.lang2item.values(): + params = set(template.parameters) + if prev_params is not None: + assert ( + params == prev_params + ), f"Parameters of MLPT '{self.name}' are inconsistent: {sorted(params)} vs {sorted(prev_params)}" + prev_params = params + return sorted(prev_params) + + +class MultiLangPromptList(MultiLangContainer[PromptList]): + pass + + +class MultiLangPromptTemplateCollection: + """ + Represents a collection of multi-language prompts, which are read from the prompts/ directory. + The language of a prompt is determined by suffix in the prompt name, e.g. "_de" for German. + For yaml files, the language can also be set globally for all prompts within it. + + The prompts/ folder may contain: + - .txt files with a single prompt template (filename defines prompt name, content is the prompt template) + - .yml files with multiple prompt templates under the top-level key 'prompts'. + Each key under 'prompts' defines a prompt name and the corresponding value is the prompt template. + The value can also be a list of prompts instead of a single template. + The language of all can be set by specifying the key 'lang' in addition to 'prompts'. + """ + + def __init__(self): + self.prompt_templates: dict[str, MultiLangPromptTemplate] = {} + self.prompt_lists: dict[str, MultiLangPromptList] = {} + prompts_dir = self._prompt_template_folder() + self._read_prompt_templates(prompts_dir) + + @classmethod + def _prompt_template_folder(cls) -> str: + dir_path = os.path.dirname(os.path.realpath(__file__)) + prompts_dir = None + for i in range(5): + dir_path = os.path.join(dir_path, "..") + prompts_dir = os.path.join(dir_path, "prompts") + if os.path.isdir(prompts_dir): + break + if not os.path.isdir(prompts_dir): + raise FileNotFoundError("Could not find the 'prompts' directory") + return prompts_dir + + @staticmethod + def _container_lang(prompt_name: str, collection: dict[str, T], container_factory: Callable[[str], T]) -> tuple[T, str]: + lang = "" + if prompt_name[-2:] in LANG_CODES and prompt_name[-3] == "_": + lang = prompt_name[-2:] + prompt_name = prompt_name[:-3] + + container = collection.get(prompt_name) + if container is None: + container = container_factory(prompt_name) + collection[prompt_name] = container + + return container, lang + + def _add_prompt_template(self, prompt_name: str, jinja_prompt_template: str): + """ + :param prompt_name: a prompt name, which may have a language shortcode suffix (e.g. "_de") + :param jinja_prompt_template: the actual prompt string which may contain placeholders/parameters (e.g. "{name}") + """ + multilang_prompt_template, lang = self._container_lang(prompt_name, self.prompt_templates, MultiLangPromptTemplate) + multilang_prompt_template.add_item(PromptTemplate(prompt_name, jinja_prompt_template), lang=lang) + + def _add_prompt_list(self, prompt_name: str, prompt_list: list[str]): + """ + :param prompt_name: a prompt name, which may have a language shortcode suffix (e.g. "_de") + :param prompt_list: a list of prompts + """ + multilang_prompt_list, lang = self._container_lang(prompt_name, self.prompt_lists, MultiLangPromptList) + multilang_prompt_list.add_item(PromptList(prompt_list), lang=lang) + + def _read_prompt_templates(self, prompts_dir: str): + for fn in os.listdir(prompts_dir): + path = os.path.join(prompts_dir, fn) + if fn.endswith(".txt"): + with open(path) as file: + prompt_template_string = file.read() + self._add_prompt_template(fn[:-4], prompt_template_string) + elif fn.endswith(".yml"): + with open(path, encoding="utf-8") as f: + data = yaml.safe_load(f) + assert "prompts" in data, f"YAML file {fn} must contain key 'prompts'" + lang_suffix = ("_" + data["lang"]) if "lang" in data else "" + for prompt_name, prompt in data["prompts"].items(): + prompt_name += lang_suffix + if isinstance(prompt, list): + self._add_prompt_list(prompt_name, prompt) + elif isinstance(prompt, str): + self._add_prompt_template(prompt_name, prompt) + else: + raise ValueError(f"Invalid prompt type: {prompt}") + + def get_multilang_prompt_template(self, prompt_name: str) -> MultiLangPromptTemplate: + return self.prompt_templates[prompt_name] + + def get_multilang_prompt_list(self, prompt_name: str) -> MultiLangPromptList: + return self.prompt_lists[prompt_name] diff --git a/src/serena/llm/prompt_factory.py b/src/serena/llm/prompt_factory.py new file mode 100644 index 00000000..56c61de2 --- /dev/null +++ b/src/serena/llm/prompt_factory.py @@ -0,0 +1,22 @@ +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, *, onboarding_file) -> str: + return self._format_prompt("onboarding_prompt", locals()) diff --git a/src/serena/util/class_decorators.py b/src/serena/util/class_decorators.py new file mode 100644 index 00000000..be36ffd6 --- /dev/null +++ b/src/serena/util/class_decorators.py @@ -0,0 +1,10 @@ +def singleton(cls): + instance = None + + def get_instance(*args, **kwargs): + nonlocal instance + if instance is None: + instance = cls(*args, **kwargs) + return instance + + return get_instance diff --git a/src/serena/util/file_system.py b/src/serena/util/file_system.py new file mode 100644 index 00000000..f3a06df7 --- /dev/null +++ b/src/serena/util/file_system.py @@ -0,0 +1,66 @@ +import os +from collections.abc import Sequence + + +def scan_directory( + path: str, + recursive=False, + relative_to: str | None = None, + ignored_dirs: Sequence[str] = (), + ignored_files: Sequence[str] = (), +) -> tuple[list[str], list[str]]: + """ + :param path: the path to scan + :param recursive: whether to recursively scan subdirectories + :param relative_to: the path to which the results should be relative to; if None, provide absolute paths + :param ignored_dirs: a list of directory names or relative paths to ignore + :param ignored_files: a list of file names or relative paths to ignore + :return: the list of directories and files + """ + files = [] + directories = [] + + abs_path = os.path.abspath(path) + rel_base = os.path.abspath(relative_to) if relative_to else None + + # Helper function to check if an item should be ignored + def is_ignored(entry_path, ignored_items): + entry_name = os.path.basename(entry_path) + + # Check if name is directly in ignored list + if entry_name in ignored_items: + return True + + # Check if relative path matches any ignored path + if rel_base: + rel_path = os.path.relpath(entry_path, rel_base) + for item in ignored_items: + if rel_path == item or rel_path.startswith(f"{item}/"): + return True + + return False + + with os.scandir(abs_path) as entries: + for entry in entries: + entry_path = entry.path + + if rel_base: + result_path = os.path.relpath(entry_path, rel_base) + else: + result_path = entry_path + + if entry.is_file(): + if not is_ignored(entry_path, ignored_files): + files.append(result_path) + elif entry.is_dir(): + if not is_ignored(entry_path, ignored_dirs): + directories.append(result_path) + if recursive: + sub_dirs, sub_files = scan_directory( + entry_path, recursive=True, relative_to=relative_to, ignored_dirs=ignored_dirs, ignored_files=ignored_files + ) + files.extend(sub_files) + directories.extend(sub_dirs) + + # Note: I've swapped the return order to match your function signature + return directories, files diff --git a/uv.lock b/uv.lock index 3f8daa35..f0623ecd 100644 --- a/uv.lock +++ b/uv.lock @@ -2060,6 +2060,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/40/b0/4562db6223154aa4e22f939003cb92514c79f3d4dccca3444253fd17f902/Send2Trash-1.8.3-py3-none-any.whl", hash = "sha256:0c31227e0bd08961c7665474a3d1ef7193929fedda4233843689baa056be46c9", size = 18072 }, ] +[[package]] +name = "sensai-utils" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e9/b1/9dba08e2d8ea739b58d1ac393c066169f1ca6506c86f0322d6e56e600ae2/sensai_utils-1.4.0.tar.gz", hash = "sha256:2d32bdcc91fd1428c5cae0181e98623142d2d5f7e115e23d585a842dd9dc59ba", size = 56317 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c6/df/49cc942b653427df0bbd3b9a2fb1e85652df3d4c1d9f6dcbfce1b6155716/sensai_utils-1.4.0-py3-none-any.whl", hash = "sha256:ed6fc57552620e43b33cf364ea0bc0fd7df39391069dd7b621b113ef55547507", size = 63214 }, +] + [[package]] name = "serena" version = "0.1.0" @@ -2073,6 +2085,7 @@ dependencies = [ { name = "pyright" }, { name = "python-dotenv" }, { name = "requests" }, + { name = "sensai-utils" }, ] [package.optional-dependencies] @@ -2125,6 +2138,7 @@ requires-dist = [ { name = "python-dotenv", specifier = ">=1.0.0,<2" }, { name = "requests", specifier = ">=2.32.3,<3" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.0.285" }, + { name = "sensai-utils", specifier = ">=1.4.0" }, { name = "sphinx", marker = "extra == 'dev'", specifier = ">=7" }, { name = "sphinx-book-theme", marker = "extra == 'dev'", specifier = ">=1.0.1" }, { name = "sphinx-comments", marker = "extra == 'dev'", specifier = ">=0.0.3" },