sql-chat / app.py
cloud-sean's picture
Update app.py
c260a8f
from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.sql_database import SQLDatabase
from langchain.llms.openai import OpenAI
from langchain.agents import AgentExecutor
import urllib, os
from io import StringIO
import sys
class Capturing(list):
def __enter__(self):
self._stdout = sys.stdout
sys.stdout = self._stringio = StringIO()
return self
def __exit__(self, *args):
self.extend(self._stringio.getvalue().splitlines())
del self._stringio # free up some memory
sys.stdout = self._stdout
def answer_question(question):
with Capturing() as printed_text:
answer = agent_executor.run("what are the top 3 most expensive items and how many customers bought them?")
import re
text = '\n'.join(printed_text) + '\n' + str(answer)
# Remove all escape characters
text = re.sub(r"\x1b\[\d+(;\d+)?m", "", text)
# Remove all characters inside angle brackets
text = re.sub(r"<.*?>", "", text)
# Remove all leading/trailing whitespaces
text = text.strip()
return text
db = SQLDatabase.from_uri("mssql+pyodbc:///?odbc_connect=Driver={ODBC Driver 18 for SQL Server};Server=tcp:tesserversean.database.windows.net,1433;Database=testdb-sean;Uid=sean;Pwd=abc123456!;Encrypt=yes;TrustServerCertificate=no;Connection Timeout=30;")
toolkit = SQLDatabaseToolkit(db=db)
agent_executor = create_sql_agent(
llm = OpenAI(model_name="gpt-4", temperature=0.0),
toolkit=toolkit,
verbose=True
)
import gradio as gr
with gr.Blocks(css="footer {visibility: hidden}", title="SQL Chat") as demo:
csv_file = gr.State([])
question = gr.Textbox(label="Question")
ask_question = gr.Button(label="Ask Question")
text_box = gr.TextArea(label="Output", lines=10)
ask_question.click(answer_question, inputs=[question], outputs=text_box)
demo.launch()