Skip to content

Commit

Permalink
Update base.py
Browse files Browse the repository at this point in the history
move .strip() as part of the _extract_python_code()
  • Loading branch information
brochure committed Nov 6, 2024
1 parent d49f09e commit 8293217
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion src/vanna/base/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]*?)```"

Expand Down Expand Up @@ -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 ----------------- #

Expand Down

0 comments on commit 8293217

Please sign in to comment.