From 6be42dfe5d6689b5a01e743443aa3bf7b4dbd6c6 Mon Sep 17 00:00:00 2001
From: Paillat <me@paillat.dev>
Date: Sun, 21 Apr 2024 21:50:11 +0200
Subject: [PATCH] :sparkles: feat(prompting): Add new function get_prompts for
 loading multiple prompts from one yaml file at a time

---
 src/utils/prompting.py | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

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