From d49f09eacde311409c288e203be4e58c48f16d9d Mon Sep 17 00:00:00 2001 From: Robert Date: Tue, 5 Nov 2024 12:51:28 +0800 Subject: [PATCH 1/2] Update base.py remove leading and trailing whitespace of plotly code generated from llm --- 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 696ad681..24f4d6c4 100644 --- a/src/vanna/base/base.py +++ b/src/vanna/base/base.py @@ -734,7 +734,7 @@ def generate_plotly_code( plotly_code = self.submit_prompt(message_log, kwargs=kwargs) - return self._sanitize_plotly_code(self._extract_python_code(plotly_code)) + return self._sanitize_plotly_code(self._extract_python_code(plotly_code.strip())) # ----------------- Connect to Any Database to run the Generated SQL ----------------- # From 829321712765cccbaefdad12efe3a0a0de52269d Mon Sep 17 00:00:00 2001 From: Robert Date: Wed, 6 Nov 2024 11:09:01 +0800 Subject: [PATCH 2/2] Update base.py move .strip() as part of the _extract_python_code() --- src/vanna/base/base.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/vanna/base/base.py b/src/vanna/base/base.py index 24f4d6c4..cf9f05ef 100644 --- a/src/vanna/base/base.py +++ b/src/vanna/base/base.py @@ -689,6 +689,9 @@ def generate_question(self, sql: str, **kwargs) -> str: return response def _extract_python_code(self, markdown_string: str) -> str: + # Strip whitespace to avoid indentation errors in LLM-generated code + markdown_string = markdown_string.strip() + # Regex pattern to match Python code blocks pattern = r"```[\w\s]*python\n([\s\S]*?)```|```([\s\S]*?)```" @@ -734,7 +737,7 @@ def generate_plotly_code( plotly_code = self.submit_prompt(message_log, kwargs=kwargs) - return self._sanitize_plotly_code(self._extract_python_code(plotly_code.strip())) + return self._sanitize_plotly_code(self._extract_python_code(plotly_code)) # ----------------- Connect to Any Database to run the Generated SQL ----------------- #