feat & bugfix: Added metadata features and fixed naming bug

src/banks/registries/directory.py

@@ -1,13 +1,19 @@
 # SPDX-FileCopyrightText: 2023-present Massimiliano Pippi <[email protected]>
 #
 # SPDX-License-Identifier: MIT
+import json
 from pathlib import Path
 
 from pydantic import BaseModel, Field
 
 from banks import Prompt
 from banks.registry import TemplateNotFoundError
 
+# Constants
+default_version = "0"
+default_index_name = "index.json"
+default_meta_path = "meta"
+
 
 class PromptFile(BaseModel):
     name: str
@@ -26,11 +32,12 @@ def __init__(self, directory_path: Path, *, force_reindex: bool = False):
             raise ValueError(msg)
 
         self._path = directory_path
-        self._index_path = self._path / "index.json"
+        self._index_path = self._path / default_index_name
         if not self._index_path.exists() or force_reindex:
             self._scan()
         else:
             self._load()
+        self._meta_path = self._path / default_meta_path
 
     def _load(self):
         self._index = PromptFileIndex.model_validate_json(self._index_path.read_text())
@@ -43,19 +50,42 @@ def _scan(self):
         self._index_path.write_text(self._index.model_dump_json())
 
     def get(self, *, name: str, version: str | None = None) -> "Prompt":
-        version = version or "0"
+        version = version or default_version
         for pf in self._index.files:
             if pf.name == name and pf.version == version and pf.path.exists():
                 return Prompt(pf.path.read_text())
         raise TemplateNotFoundError
 
     def set(self, *, name: str, prompt: Prompt, version: str | None = None, overwrite: bool = False):
-        version = version or "0"
+        version = version or default_version
         for pf in self._index.files:
             if pf.name == name and pf.version == version and overwrite:
                 pf.path.write_text(prompt.raw)
                 return
-        new_prompt_file = self._path / "{name}.{version}.jinja"
+        new_prompt_file = self._path / f"{name}.{version}.jinja"
         new_prompt_file.write_text(prompt.raw)
         pf = PromptFile(name=name, version=version, path=new_prompt_file)
         self._index.files.append(pf)
+
+    def get_meta(self, *, name: str, version: str | None = None) -> dict:
+        version = version or default_version
+        meta_path = self._meta_path / f"{name}.{version}.json"
+        if not meta_path.exists():
+            msg = f"Meta directory or file for prompt {name}:{version}.jinja not found."
+            raise FileNotFoundError(msg)
+        return json.loads(open(meta_path).read())
+
+    def set_meta(self, *, meta: dict, name: str, version: str | None = None, overwrite: bool = False):
+        version = version or default_version
+        if not self._meta_path.exists():
+            self._meta_path.mkdir()
+        if Path(self._path / f"{name}.{version}.jinja") not in [pf.path for pf in self._index.files]:
+            msg = f"Prompt {name}.{version}.jinja not found in the index. Cannot set meta for a non-existing prompt."
+            raise ValueError(msg)
+
+        if f"{name}:{version}.json" in self._meta_path.glob("*.json"):
+            if not overwrite:
+                msg = f"Meta file for prompt {name}:{version} already exists. Use overwrite=True to overwrite."
+                raise ValueError(msg)
+        meta_path = self._meta_path / f"{name}.{version}.json"
+        meta_path.write_text(json.dumps(meta))

tests/test_directory_registry.py

@@ -1,4 +1,5 @@
 import os
+import time
 from pathlib import Path
 
 import pytest
@@ -32,7 +33,7 @@ def test_init_from_existing_index(populated_dir):
 
 
 def test_init_from_existing_index_force(populated_dir):
-    r = DirectoryTemplateRegistry(populated_dir)  # creates the index
+    _ = DirectoryTemplateRegistry(populated_dir)  # creates the index
     # change the directory structure
     f = populated_dir / "blog.jinja"
     os.remove(f)
@@ -65,3 +66,29 @@ def test_set_existing_overwrite(populated_dir):
     new_prompt = Prompt("a new prompt!")
     r.set(name="blog", prompt=new_prompt, overwrite=True)
     assert r.get(name="blog").raw.startswith("a new prompt!")
+
+
+def test_set_multiple_templates(populated_dir):
+    r = DirectoryTemplateRegistry(Path(populated_dir))
+    new_prompt = Prompt("a very new prompt!")
+    old_prompt = Prompt("an old prompt!")
+    r.set(name="new", version="2", prompt=new_prompt)
+    r.set(name="old", version="1", prompt=old_prompt)
+    assert r.get(name="old", version="1").raw == "an old prompt!"
+    assert r.get(name="new", version="2").raw == "a very new prompt!"
+
+
+def test_empty_meta(populated_dir):
+    r = DirectoryTemplateRegistry(populated_dir)
+    with pytest.raises(FileNotFoundError):
+        r.get_meta(name="blog")
+
+
+def test_set_meta(populated_dir):
+    r = DirectoryTemplateRegistry(populated_dir)
+    new_prompt = Prompt("a very new prompt!")
+    r.set(name="new", version="3", prompt=new_prompt)
+    r.set_meta(name="new", version="3", meta={"accuracy": 91.2, "last_updated": time.ctime()})
+    assert r.get_meta(name="new", version="3") == {"accuracy": 91.2, "last_updated": time.ctime()}
+    with pytest.raises(ValueError):
+        r.set_meta(name="foo", version="bar", meta={"accuracy": 91.2, "last_updated": time.ctime()}, overwrite=False)