From 0c808cd21683783ca3f4cafb85c4a1675de3f495 Mon Sep 17 00:00:00 2001 From: Alan Konarski <129968242+akonarski-ds@users.noreply.github.com> Date: Wed, 2 Oct 2024 15:34:14 +0200 Subject: [PATCH] feat(prompts): integration with promptfoo (#54) --- .pre-commit-config.yaml | 2 +- .../src/ragbits/core/prompt/prompt.py | 26 +++++++-- .../tests/unit/prompts/test_prompt.py | 53 +++++++++++++++++++ packages/ragbits-dev-kit/README.md | 25 +++++++++ .../src/ragbits/dev_kit/cli.py | 2 + .../src/ragbits/dev_kit/promptfoo.py | 33 ++++++++++++ 6 files changed, 137 insertions(+), 4 deletions(-) create mode 100644 packages/ragbits-dev-kit/src/ragbits/dev_kit/promptfoo.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b8e6b3bd..7428c1b0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -44,7 +44,7 @@ repos: - id: mypy # You can add additional plugins for mypy below # such as types-python-dateutil - additional_dependencies: [pydantic>=2.8.2] + additional_dependencies: [pydantic>=2.8.2, types-pyyaml>=6.0.12] exclude: (/test_|setup.py|/tests/|docs/) # Sort imports alphabetically, and automatically separated into sections and by type. diff --git a/packages/ragbits-core/src/ragbits/core/prompt/prompt.py b/packages/ragbits-core/src/ragbits/core/prompt/prompt.py index 69e87efa..75b3c093 100644 --- a/packages/ragbits-core/src/ragbits/core/prompt/prompt.py +++ b/packages/ragbits-core/src/ragbits/core/prompt/prompt.py @@ -22,7 +22,7 @@ class Prompt(Generic[InputT, OutputT], BasePromptWithParser[OutputT], metaclass= system_prompt: Optional[str] = None user_prompt: str - additional_messages: ChatFormat = [] + additional_messages: Optional[ChatFormat] = None # function that parses the response from the LLM to specific output type # if not provided, the class tries to set it automatically based on the output type @@ -125,10 +125,13 @@ def chat(self) -> ChatFormat: Returns: ChatFormat: A list of dictionaries, each containing the role and content of a message. """ - return [ + chat = [ *([{"role": "system", "content": self.system_message}] if self.system_message is not None else []), {"role": "user", "content": self.user_message}, - ] + self.additional_messages + ] + if self.additional_messages: + chat.extend(self.additional_messages) + return chat def add_user_message(self, message: str) -> "Prompt[InputT, OutputT]": """ @@ -140,6 +143,8 @@ def add_user_message(self, message: str) -> "Prompt[InputT, OutputT]": Returns: Prompt[InputT, OutputT]: The current prompt instance in order to allow chaining. """ + if self.additional_messages is None: + self.additional_messages = [] self.additional_messages.append({"role": "user", "content": message}) return self @@ -153,6 +158,8 @@ def add_assistant_message(self, message: str) -> "Prompt[InputT, OutputT]": Returns: Prompt[InputT, OutputT]: The current prompt instance in order to allow chaining. """ + if self.additional_messages is None: + self.additional_messages = [] self.additional_messages.append({"role": "assistant", "content": message}) return self @@ -190,3 +197,16 @@ def parse_response(self, response: str) -> OutputT: ResponseParsingError: If the response cannot be parsed. """ return self.response_parser(response) + + @classmethod + def to_promptfoo(cls, config: dict[str, Any]) -> ChatFormat: + """ + Generate a prompt in the promptfoo format from a promptfoo test configuration. + + Args: + config: The promptfoo test configuration. + + Returns: + ChatFormat: The prompt in the format used by promptfoo. + """ + return cls(cls.input_type.model_validate(config["vars"])).chat # type: ignore diff --git a/packages/ragbits-core/tests/unit/prompts/test_prompt.py b/packages/ragbits-core/tests/unit/prompts/test_prompt.py index da0352cd..557849f0 100644 --- a/packages/ragbits-core/tests/unit/prompts/test_prompt.py +++ b/packages/ragbits-core/tests/unit/prompts/test_prompt.py @@ -198,3 +198,56 @@ class TestPrompt(Prompt[_PromptInput, str]): prompt = TestPrompt(_PromptInput(name="John", age=15, theme="pop")) assert prompt.output_schema() is None + + +def test_to_promptfoo(): + """Test that a prompt can be converted to a promptfoo prompt.""" + promptfoo_test_config = { + "vars": {"name": "John", "age": 25, "theme": "pop"}, + } + + class TestPrompt(Prompt[_PromptInput, str]): # pylint: disable=unused-variable + """A test prompt""" + + system_prompt = """ + You are a song generator for a {% if age > 18 %}adult{% else %}child{% endif %} named {{ name }}. + """ + user_prompt = "Theme for the song is {{ theme }}." + + assert TestPrompt.to_promptfoo(promptfoo_test_config) == [ + {"role": "system", "content": "You are a song generator for a adult named John."}, + {"role": "user", "content": "Theme for the song is pop."}, + ] + + +def test_two_instances_do_not_share_additional_messages(): + """ + Test that two instances of a prompt do not share additional messages. + """ + + class TestPrompt(Prompt[_PromptInput, str]): # pylint: disable=unused-variable + """A test prompt""" + + system_prompt = """ + You are a song generator for a {% if age > 18 %}adult{% else %}child{% endif %} named {{ name }}. + """ + user_prompt = "Theme for the song is {{ theme }}." + + prompt1 = TestPrompt(_PromptInput(name="John", age=15, theme="pop")) + prompt1.add_assistant_message("It's a really catchy tune.").add_user_message("I like it.") + + prompt2 = TestPrompt(_PromptInput(name="Alice", age=30, theme="rock")) + prompt2.add_assistant_message("It's a nice tune.") + + assert prompt1.chat == [ + {"role": "system", "content": "You are a song generator for a child named John."}, + {"role": "user", "content": "Theme for the song is pop."}, + {"role": "assistant", "content": "It's a really catchy tune."}, + {"role": "user", "content": "I like it."}, + ] + + assert prompt2.chat == [ + {"role": "system", "content": "You are a song generator for a adult named Alice."}, + {"role": "user", "content": "Theme for the song is rock."}, + {"role": "assistant", "content": "It's a nice tune."}, + ] diff --git a/packages/ragbits-dev-kit/README.md b/packages/ragbits-dev-kit/README.md index f27a2678..099abe68 100644 --- a/packages/ragbits-dev-kit/README.md +++ b/packages/ragbits-dev-kit/README.md @@ -1 +1,26 @@ # Ragbits Development Kit + +## Promptfoo Integration + +Ragbits' `Prompt` abstraction can be seamlessly integrated with the `promptfoo` tool. After installing `promptfoo` as +specified in the [promptfoo documentation](https://www.promptfoo.dev/docs/installation/), you can generate promptfoo +configuration files for all the prompts discovered by our autodiscover mechanism by running the following command: + +```bash +rbts prompts generate-promptfoo-configs +``` + +This command will generate a YAML files in the directory specified by `--target-path` (`promptfooconfigs` by +default). The generated file should look like this: + +```yaml +prompts: + - file:///path/to/your/prompt:PromptClass.to_promptfoo +``` + +You can then edit the generated file to add your custom `promptfoo` configurations. Once your `promptfoo` configuration +file is ready, you can run `promptfoo` with the following command: + +```bash +promptfoo -c /path/to/generated/promptfoo-config.yaml eval +``` \ No newline at end of file diff --git a/packages/ragbits-dev-kit/src/ragbits/dev_kit/cli.py b/packages/ragbits-dev-kit/src/ragbits/dev_kit/cli.py index 22475bb8..73f5267f 100644 --- a/packages/ragbits-dev-kit/src/ragbits/dev_kit/cli.py +++ b/packages/ragbits-dev-kit/src/ragbits/dev_kit/cli.py @@ -1,6 +1,7 @@ import typer from .prompt_lab.app import lab_app +from .promptfoo import generate_configs prompts_app = typer.Typer(no_args_is_help=True) @@ -13,4 +14,5 @@ def register(app: typer.Typer) -> None: app: The Typer object to register the commands with. """ prompts_app.command(name="lab")(lab_app) + prompts_app.command(name="generate-promptfoo-configs")(generate_configs) app.add_typer(prompts_app, name="prompts", help="Commands for managing prompts") diff --git a/packages/ragbits-dev-kit/src/ragbits/dev_kit/promptfoo.py b/packages/ragbits-dev-kit/src/ragbits/dev_kit/promptfoo.py new file mode 100644 index 00000000..d7357621 --- /dev/null +++ b/packages/ragbits-dev-kit/src/ragbits/dev_kit/promptfoo.py @@ -0,0 +1,33 @@ +import os +from pathlib import Path + +import yaml +from rich.console import Console + +from ragbits.dev_kit.prompt_lab.discovery import PromptDiscovery +from ragbits.dev_kit.prompt_lab.discovery.prompt_discovery import DEFAULT_FILE_PATTERN + + +def generate_configs( + file_pattern: str = DEFAULT_FILE_PATTERN, root_path: Path = Path.cwd(), target_path: Path = Path("promptfooconfigs") +) -> None: + """ + Generates promptfoo configuration files for all discovered prompts. + + Args: + file_pattern: The file pattern to search for Prompt objects. Defaults to "**/prompt_*.py" + root_path: The root path to search for Prompt objects. Defaults to the directory where the script is run. + target_path: The path to save the promptfoo configuration files. Defaults to "promptfooconfigs". + """ + prompts = PromptDiscovery(file_pattern=file_pattern, root_path=root_path).discover() + Console().print( + f"Discovered {len(prompts)} prompts." + f" Saving promptfoo configuration files to [bold green]{target_path}[/] folder ..." + ) + + if not target_path.exists(): + target_path.mkdir() + for prompt in prompts: + with open(target_path / f"{prompt.__qualname__}.yaml", "w", encoding="utf-8") as f: + prompt_path = f'file://{prompt.__module__.replace(".", os.sep)}.py:{prompt.__qualname__}.to_promptfoo' + yaml.dump({"prompts": [prompt_path]}, f)