From 67f1bceda676ef8b5c6acea7c547ce1e5c67b96a Mon Sep 17 00:00:00 2001
From: mahendrasinghbora <mail.mahendrasinghbora@gmail.com>
Date: Fri, 21 Jun 2024 16:43:54 +0530
Subject: [PATCH] Fix regex in extract_sql to handle CTEs without matching
 WITHIN

---
 src/vanna/base/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/vanna/base/base.py b/src/vanna/base/base.py
index 56360adc..492516ea 100644
--- a/src/vanna/base/base.py
+++ b/src/vanna/base/base.py
@@ -182,7 +182,7 @@ def extract_sql(self, llm_response: str) -> str:
         """
 
         # If the llm_response contains a CTE (with clause), extract the last sql between WITH and ;
-        sqls = re.findall(r"WITH.*?;", llm_response, re.DOTALL)
+        sqls = re.findall(r"\bWITH\b .*?;", llm_response, re.DOTALL)
         if sqls:
             sql = sqls[-1]
             self.log(title="Extracted SQL", message=f"{sql}")