Harrison/add template format

langchain/chains/llm.py

@@ -33,7 +33,7 @@ def output_keys(self) -> List[str]:
 
     def _run(self, inputs: Dict[str, Any]) -> Dict[str, str]:
         selected_inputs = {k: inputs[k] for k in self.prompt.input_variables}
-        prompt = self.prompt.template.format(**selected_inputs)
+        prompt = self.prompt.format(**selected_inputs)
 
         kwargs = {}
         if "stop" in inputs:

langchain/prompt.py

@@ -1,30 +1,47 @@
-"""Base schema types."""
-from typing import Dict, List
+"""Prompt schema definition."""
+from typing import Any, Dict, List
 
 from pydantic import BaseModel, Extra, root_validator
 
 from langchain.formatting import formatter
 
+_FORMATTER_MAPPING = {
+    "f-string": formatter.format,
+}
+
 
 class Prompt(BaseModel):
     """Schema to represent a prompt for an LLM."""
 
     input_variables: List[str]
     template: str
+    template_format: str = "f-string"
 
     class Config:
         """Configuration for this pydantic object."""
 
         extra = Extra.forbid
 
+    def format(self, **kwargs: Any) -> str:
+        """Format the prompt with the inputs."""
+        return _FORMATTER_MAPPING[self.template_format](self.template, **kwargs)
+
     @root_validator()
     def template_is_valid(cls, values: Dict) -> Dict:
         """Check that template and input variables are consistent."""
         input_variables = values["input_variables"]
         template = values["template"]
+        template_format = values["template_format"]
+        if template_format not in _FORMATTER_MAPPING:
+            valid_formats = list(_FORMATTER_MAPPING)
+            raise ValueError(
+                f"Invalid template format. Got `{template_format}`;"
+                f" should be one of {valid_formats}"
+            )
         dummy_inputs = {input_variable: "foo" for input_variable in input_variables}
         try:
-            formatter.format(template, **dummy_inputs)
+            formatter_func = _FORMATTER_MAPPING[template_format]
+            formatter_func(template, **dummy_inputs)
         except KeyError:
             raise ValueError("Invalid prompt schema.")
         return values

tests/unit_tests/test_prompt.py

@@ -35,3 +35,13 @@ def test_prompt_wrong_input_variables() -> None:
     input_variables = ["bar"]
     with pytest.raises(ValueError):
         Prompt(input_variables=input_variables, template=template)
+
+
+def test_prompt_invalid_template_format() -> None:
+    """Test initializing a prompt with invalid template format."""
+    template = "This is a {foo} test."
+    input_variables = ["foo"]
+    with pytest.raises(ValueError):
+        Prompt(
+            input_variables=input_variables, template=template, template_format="bar"
+        )