From a1d581e728ed7395a5e113ee40c9ccc6acc527e3 Mon Sep 17 00:00:00 2001
From: Massimiliano Pippi <mpippi@gmail.com>
Date: Sat, 28 Sep 2024 17:04:56 +0200
Subject: [PATCH] feat: add name field to prompt

---
 src/banks/prompt.py  | 6 ++++++
 tests/test_prompt.py | 2 ++
 2 files changed, 8 insertions(+)

diff --git a/src/banks/prompt.py b/src/banks/prompt.py
index ba933b5..caf11b1 100644
--- a/src/banks/prompt.py
+++ b/src/banks/prompt.py
@@ -15,6 +15,7 @@ def __init__(
         self,
         text: str,
         *,
+        name: str | None = None,
         version: str | None = None,
         metadata: dict[str, Any] | None = None,
         canary_word: str | None = None,
@@ -31,6 +32,7 @@ def __init__(
                 be used.
         """
         self._metadata = metadata or {}
+        self._name = name
         self._raw: str = text
         self._render_cache = render_cache or DefaultCache()
         self._template = env.from_string(text)
@@ -47,6 +49,10 @@ def _get_context(self, data: dict | None) -> dict:
     def metadata(self) -> dict[str, Any]:
         return self._metadata
 
+    @property
+    def name(self) -> str | None:
+        return self._name
+
     @property
     def raw(self) -> str:
         """Returns the raw text of the prompt."""
diff --git a/tests/test_prompt.py b/tests/test_prompt.py
index f3f3b2d..b74594c 100644
--- a/tests/test_prompt.py
+++ b/tests/test_prompt.py
@@ -35,12 +35,14 @@ def test_ctor():
         text="This is raw text",
         version="1.0",
         metadata={"LLM": "GPT-3.5"},
+        name="test_prompt",
         canary_word="FOO",
         render_cache=DefaultCache(),
     )
     assert p.raw == "This is raw text"
     assert p.version == "1.0"
     assert p.metadata == {"LLM": "GPT-3.5"}
+    assert p.name == "test_prompt"
     assert p.canary_leaked("The message is FOO")
     assert p.text() == "This is raw text"
     assert p.text() == "This is raw text"