From f8fb9f55eada1a222aa8ba03f1c62552fbde9ae0 Mon Sep 17 00:00:00 2001
From: Vignesh Aigal <vignesh.aigal@gmail.com>
Date: Sat, 12 Oct 2024 10:46:15 -0700
Subject: [PATCH] Add liquid variables extractor using ast

---
 llmstack/play/utils.py | 78 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 78 insertions(+)

diff --git a/llmstack/play/utils.py b/llmstack/play/utils.py
index 510aa649a96..43d22054721 100644
--- a/llmstack/play/utils.py
+++ b/llmstack/play/utils.py
@@ -2,9 +2,72 @@
 import re
 import threading
 
+from liquid.ast import ChildNode, ParseTree
+from liquid.expression import (
+    AssignmentExpression,
+    Blank,
+    BooleanExpression,
+    Continue,
+    Empty,
+    FilteredExpression,
+    Identifier,
+    InfixExpression,
+    LoopExpression,
+    Nil,
+)
 from pydantic import BaseModel
 
 
+def extract_nodes(node):
+    nodes = []
+    if isinstance(node, ParseTree):
+        for stmt in node.statements:
+            nodes.extend(extract_nodes(stmt) or [])
+    elif isinstance(node, ChildNode):
+        nodes.append(node)
+        if node.node:
+            nodes.extend(extract_nodes(node.node) or [])
+    elif hasattr(node, "children"):
+        for child in node.children():
+            nodes.extend(extract_nodes(child) or [])
+    return nodes
+
+
+def extract_variables(expression):
+    variables = []
+    if expression is None:
+        return []
+    if isinstance(expression, Identifier):
+        result = []
+        for element in expression.path:
+            if isinstance(element, Identifier):
+                result.extend(extract_variables(element))
+            else:
+                result.append(element.value)
+        return result
+    if (
+        isinstance(expression, Nil)
+        or isinstance(expression, Empty)
+        or isinstance(expression, Blank)
+        or isinstance(expression, Continue)
+    ):
+        return []
+    elif isinstance(expression, FilteredExpression):
+        variables.extend(extract_variables(expression.expression))
+    elif isinstance(expression, AssignmentExpression):
+        variables.extend(extract_variables(expression.expression))
+    elif isinstance(expression, LoopExpression):
+        variables.extend(extract_variables(expression.iterable))
+    elif isinstance(expression, BooleanExpression):
+        variables.extend(extract_variables(expression.expression))
+    elif isinstance(expression, InfixExpression):
+        variables.extend([extract_variables(expression.left)])
+        variables.extend([extract_variables(expression.right)])
+    else:
+        raise NotImplementedError(f"Unsupported expression: {expression} {type(expression)}")
+    return variables
+
+
 def run_coro_in_new_loop(coro):
     def start_loop(loop):
         asyncio.set_event_loop(loop)
@@ -104,6 +167,21 @@ def extract_from_string(s):
     return variables
 
 
+def extract_variables_from_liquid_template(liquid_template):
+    variables = []
+
+    nodes = extract_nodes(liquid_template.tree)
+    for node in nodes:
+        extracted_variables = extract_variables(node.expression)
+        if extracted_variables:
+            if isinstance(extracted_variables[0], list):
+                variables.extend(extracted_variables)
+            else:
+                variables.append(extracted_variables)
+
+    return variables
+
+
 # A utility function to recursively convert template vars of type
 # _inputs[0].xyz to _inputs0.xyz for backward compatibility in a
 # dictionary with nested string values