feat(prompts): add util for variable name extraction (#1046)

Co-authored-by: Roberto Montalti <37136851+rhighs@users.noreply.github.com>

Author: hassiebp

langfuse/model.py

@@ -1,7 +1,7 @@
 """@private"""
 
 from abc import ABC, abstractmethod
-from typing import Optional, TypedDict, Any, Dict, Union, List
+from typing import Optional, TypedDict, Any, Dict, Union, List, Tuple
 import re
 
 from langfuse.api.resources.commons.types.dataset import (
@@ -54,6 +54,72 @@ class ChatMessageDict(TypedDict):
     content: str
 
 
+class TemplateParser:
+    OPENING = "{{"
+    CLOSING = "}}"
+
+    @staticmethod
+    def _parse_next_variable(
+        content: str, start_idx: int
+    ) -> Optional[Tuple[str, int, int]]:
+        """Returns (variable_name, start_pos, end_pos) or None if no variable found"""
+        var_start = content.find(TemplateParser.OPENING, start_idx)
+        if var_start == -1:
+            return None
+
+        var_end = content.find(TemplateParser.CLOSING, var_start)
+        if var_end == -1:
+            return None
+
+        variable_name = content[
+            var_start + len(TemplateParser.OPENING) : var_end
+        ].strip()
+        return (variable_name, var_start, var_end + len(TemplateParser.CLOSING))
+
+    @staticmethod
+    def find_variable_names(content: str) -> List[str]:
+        names = []
+        curr_idx = 0
+
+        while curr_idx < len(content):
+            result = TemplateParser._parse_next_variable(content, curr_idx)
+            if not result:
+                break
+            names.append(result[0])
+            curr_idx = result[2]
+
+        return names
+
+    @staticmethod
+    def compile_template(content: str, data: Optional[Dict[str, Any]] = None) -> str:
+        if data is None:
+            return content
+
+        result_list = []
+        curr_idx = 0
+
+        while curr_idx < len(content):
+            result = TemplateParser._parse_next_variable(content, curr_idx)
+
+            if not result:
+                result_list.append(content[curr_idx:])
+                break
+
+            variable_name, var_start, var_end = result
+            result_list.append(content[curr_idx:var_start])
+
+            if variable_name in data:
+                result_list.append(
+                    str(data[variable_name]) if data[variable_name] is not None else ""
+                )
+            else:
+                result_list.append(content[var_start:var_end])
+
+            curr_idx = var_end
+
+        return "".join(result_list)
+
+
 class BasePromptClient(ABC):
     name: str
     version: int
@@ -73,6 +139,11 @@ def __init__(self, prompt: Prompt, is_fallback: bool = False):
     def compile(self, **kwargs) -> Union[str, List[ChatMessage]]:
         pass
 
+    @property
+    @abstractmethod
+    def variables(self) -> List[str]:
+        pass
+
     @abstractmethod
     def __eq__(self, other):
         pass
@@ -85,55 +156,19 @@ def get_langchain_prompt(self):
     def _get_langchain_prompt_string(content: str):
         return re.sub(r"{{\s*(\w+)\s*}}", r"{\g<1>}", content)
 
-    @staticmethod
-    def _compile_template_string(content: str, data: Dict[str, Any] = {}) -> str:
-        opening = "{{"
-        closing = "}}"
-
-        result_list = []
-        curr_idx = 0
-
-        while curr_idx < len(content):
-            # Find the next opening tag
-            var_start = content.find(opening, curr_idx)
-
-            if var_start == -1:
-                result_list.append(content[curr_idx:])
-                break
-
-            # Find the next closing tag
-            var_end = content.find(closing, var_start)
-
-            if var_end == -1:
-                result_list.append(content[curr_idx:])
-                break
-
-            # Append the content before the variable
-            result_list.append(content[curr_idx:var_start])
-
-            # Extract the variable name
-            variable_name = content[var_start + len(opening) : var_end].strip()
-
-            # Append the variable value
-            if variable_name in data:
-                result_list.append(
-                    str(data[variable_name]) if data[variable_name] is not None else ""
-                )
-            else:
-                result_list.append(content[var_start : var_end + len(closing)])
-
-            curr_idx = var_end + len(closing)
-
-        return "".join(result_list)
-
 
 class TextPromptClient(BasePromptClient):
     def __init__(self, prompt: Prompt_Text, is_fallback: bool = False):
         super().__init__(prompt, is_fallback)
         self.prompt = prompt.prompt
 
     def compile(self, **kwargs) -> str:
-        return self._compile_template_string(self.prompt, kwargs)
+        return TemplateParser.compile_template(self.prompt, kwargs)
+
+    @property
+    def variables(self) -> List[str]:
+        """Return all the variable names in the prompt template."""
+        return TemplateParser.find_variable_names(self.prompt)
 
     def __eq__(self, other):
         if isinstance(self, other.__class__):
@@ -160,7 +195,7 @@ def get_langchain_prompt(self, **kwargs) -> str:
             str: The string that can be plugged into Langchain's PromptTemplate.
         """
         prompt = (
-            self._compile_template_string(self.prompt, kwargs)
+            TemplateParser.compile_template(self.prompt, kwargs)
             if kwargs
             else self.prompt
         )
@@ -178,12 +213,23 @@ def __init__(self, prompt: Prompt_Chat, is_fallback: bool = False):
     def compile(self, **kwargs) -> List[ChatMessageDict]:
         return [
             ChatMessageDict(
-                content=self._compile_template_string(chat_message["content"], kwargs),
+                content=TemplateParser.compile_template(
+                    chat_message["content"], kwargs
+                ),
                 role=chat_message["role"],
             )
             for chat_message in self.prompt
         ]
 
+    @property
+    def variables(self) -> List[str]:
+        """Return all the variable names in the chat prompt template."""
+        return [
+            variable
+            for chat_message in self.prompt
+            for variable in TemplateParser.find_variable_names(chat_message["content"])
+        ]
+
     def __eq__(self, other):
         if isinstance(self, other.__class__):
             return (
@@ -215,7 +261,7 @@ def get_langchain_prompt(self, **kwargs):
             (
                 msg["role"],
                 self._get_langchain_prompt_string(
-                    self._compile_template_string(msg["content"], kwargs)
+                    TemplateParser.compile_template(msg["content"], kwargs)
                     if kwargs
                     else msg["content"]
                 ),

tests/test_prompt.py

@@ -994,3 +994,100 @@ def test_do_not_link_observation_if_fallback():
 
     assert len(trace.observations) == 1
     assert trace.observations[0].prompt_id is None
+
+
+def test_variable_names_on_content_with_variable_names():
+    langfuse = Langfuse()
+
+    prompt_client = langfuse.create_prompt(
+        name="test_variable_names_1",
+        prompt="test prompt with var names {{ var1 }} {{ var2 }}",
+        is_active=True,
+        type="text",
+    )
+
+    second_prompt_client = langfuse.get_prompt("test_variable_names_1")
+
+    assert prompt_client.name == second_prompt_client.name
+    assert prompt_client.version == second_prompt_client.version
+    assert prompt_client.prompt == second_prompt_client.prompt
+    assert prompt_client.labels == ["production", "latest"]
+
+    var_names = second_prompt_client.variables
+
+    assert var_names == ["var1", "var2"]
+
+
+def test_variable_names_on_content_with_no_variable_names():
+    langfuse = Langfuse()
+
+    prompt_client = langfuse.create_prompt(
+        name="test_variable_names_2",
+        prompt="test prompt with no var names",
+        is_active=True,
+        type="text",
+    )
+
+    second_prompt_client = langfuse.get_prompt("test_variable_names_2")
+
+    assert prompt_client.name == second_prompt_client.name
+    assert prompt_client.version == second_prompt_client.version
+    assert prompt_client.prompt == second_prompt_client.prompt
+    assert prompt_client.labels == ["production", "latest"]
+
+    var_names = second_prompt_client.variables
+
+    assert var_names == []
+
+
+def test_variable_names_on_content_with_variable_names_chat_messages():
+    langfuse = Langfuse()
+
+    prompt_client = langfuse.create_prompt(
+        name="test_variable_names_3",
+        prompt=[
+            {
+                "role": "system",
+                "content": "test prompt with template vars {{ var1 }} {{ var2 }}",
+            },
+            {"role": "user", "content": "test prompt 2 with template vars {{ var3 }}"},
+        ],
+        is_active=True,
+        type="chat",
+    )
+
+    second_prompt_client = langfuse.get_prompt("test_variable_names_3")
+
+    assert prompt_client.name == second_prompt_client.name
+    assert prompt_client.version == second_prompt_client.version
+    assert prompt_client.prompt == second_prompt_client.prompt
+    assert prompt_client.labels == ["production", "latest"]
+
+    var_names = second_prompt_client.variables
+
+    assert var_names == ["var1", "var2", "var3"]
+
+
+def test_variable_names_on_content_with_no_variable_names_chat_messages():
+    langfuse = Langfuse()
+
+    prompt_client = langfuse.create_prompt(
+        name="test_variable_names_4",
+        prompt=[
+            {"role": "system", "content": "test prompt with no template vars"},
+            {"role": "user", "content": "test prompt 2 with no template vars"},
+        ],
+        is_active=True,
+        type="chat",
+    )
+
+    second_prompt_client = langfuse.get_prompt("test_variable_names_4")
+
+    assert prompt_client.name == second_prompt_client.name
+    assert prompt_client.version == second_prompt_client.version
+    assert prompt_client.prompt == second_prompt_client.prompt
+    assert prompt_client.labels == ["production", "latest"]
+
+    var_names = second_prompt_client.variables
+
+    assert var_names == []