Add model_name property to templates

When rendering templates with special variables representing, e.g., BOS
tokens we need to know which model the prompt is going to be passed to
so as to render the template with the appropriate token. This PR adds
a `model_name` attribute to the `Template` class.

Author: rlouf

prompts/templates.py

@@ -37,6 +37,7 @@ class Template:
 
     template: str
     signature: inspect.Signature
+    model: Optional[str] = None
     registry: Dict[str, Callable] = field(default_factory=dict)
 
     def __call__(self, *args, **kwargs) -> str:
@@ -83,6 +84,7 @@ def register(self, model_name: str):
 
         def wrapper(fn: Callable):
             tpl = template(fn)
+            tpl.model = model_name
             self.registry[model_name] = tpl
             return tpl
 

tests/test_templates.py

@@ -197,6 +197,10 @@ def simple_prompt_name(query: str):
     assert callable(simple_prompt)
     assert callable(simple_prompt["provider/name"])
 
+    assert simple_prompt.model is None
+    assert simple_prompt_name.model == "provider/name"
+    assert simple_prompt["provider/name"].model == "provider/name"
+
     assert simple_prompt("test") == "test"
     assert simple_prompt_name("test") == "name: test"