Adds update_prompt_template function

jupyterlab · Aug 2, 2023 · dd6935f · dd6935f
1 parent f6ef7b2
commit dd6935f
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
@@ -176,6 +176,12 @@ async def _call_in_executor(self, *args, **kwargs) -> Coroutine[Any, Any, str]:
         _call_with_args = functools.partial(self._call, *args, **kwargs)
         return await loop.run_in_executor(executor, _call_with_args)
 
+    def update_prompt_template(format: str, template: str):
+        """
+        Changes the class-level prompt template for a given format.
+        """
+        prompt_templates[format] = PromptTemplate.from_template(template)
+
     def prompt_template(self, format) -> PromptTemplate:
         """
         Produce a prompt template suitable for use with a particular model, to