cloud-sean/sql-chat

1from langchain.agents import create_sql_agent2from langchain.agents.agent_toolkits import SQLDatabaseToolkit3from langchain.sql_database import SQLDatabase4from langchain.llms.openai import OpenAI5from langchain.agents import AgentExecutor6import urllib, os7 8from io import StringIO 9import sys10 11class Capturing(list):12    def __enter__(self):13        self._stdout = sys.stdout14        sys.stdout = self._stringio = StringIO()15        return self16    def __exit__(self, *args):17        self.extend(self._stringio.getvalue().splitlines())18        del self._stringio    # free up some memory19        sys.stdout = self._stdout20 21 22 23def answer_question(question):24    with Capturing() as printed_text:25        answer = agent_executor.run("what are the top 3 most expensive items and how many customers bought them?")26    import re27    text = '\n'.join(printed_text) + '\n' + str(answer)28    # Remove all escape characters29    text = re.sub(r"\x1b\[\d+(;\d+)?m", "", text)30 31    # Remove all characters inside angle brackets32    text = re.sub(r"<.*?>", "", text)33 34    # Remove all leading/trailing whitespaces35    text = text.strip()36    return text37 38db = 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;")39toolkit = SQLDatabaseToolkit(db=db)40 41 42agent_executor = create_sql_agent(43    llm = OpenAI(model_name="gpt-4", temperature=0.0),44    toolkit=toolkit,45    verbose=True46)47 48 49 50import gradio as gr51 52with gr.Blocks(css="footer {visibility: hidden}", title="SQL Chat") as demo:53    csv_file = gr.State([])54    question = gr.Textbox(label="Question")55    ask_question = gr.Button(label="Ask Question")56    text_box = gr.TextArea(label="Output", lines=10)57 58    ask_question.click(answer_question, inputs=[question], outputs=text_box)59    60 61 62demo.launch()63 64 65