diff --git a/src/utils/prompting.py b/src/utils/prompting.py index 55e927c..cd90ae4 100644 --- a/src/utils/prompting.py +++ b/src/utils/prompting.py @@ -10,7 +10,7 @@ class Prompt(TypedDict): def get_prompt( - name, *, location: str = "src/chore/prompts", by_file_location: str = None + name, *, location: str = "src/chore/prompts", by_file_location: str = None ) -> tuple[str, str]: if by_file_location: path = os.path.join( @@ -25,3 +25,21 @@ def get_prompt( with open(path, "r") as file: prompt: Prompt = yaml.safe_load(file) return prompt["system"], prompt["chat"] + + +def get_prompts( + name, *, location: str = "src/chore/prompts", by_file_location: str = None +) -> dict[str, Prompt]: + if by_file_location: + path = os.path.join( + os.path.dirname(os.path.abspath(by_file_location)), + "prompts", + f"{name}.yaml", + ) + else: + path = os.path.join(os.getcwd(), location, f"{name}.yaml") + if not os.path.exists(path): + raise FileNotFoundError(f"Prompt file {path} does not exist.") + with open(path, "r") as file: + prompts: dict[str, Prompt] = yaml.safe_load(file) + return prompts