fixed sql chain

app.py

@@ -1,13 +1,14 @@
 from os import write
+from langchain.llms.base import LLM
 import streamlit as st
 import numpy as np
 import pandas as pd
 import altair as alt
 import chat as idf_chat
-import sqlite3
+# import sqlite3
 from langchain.sql_database import SQLDatabase
 from langchain.agents.agent_toolkits import SQLDatabaseToolkit
-from langchain.agents import create_sql_agent
+from langchain.agents import LLMSingleActionAgent, agent_types, create_sql_agent, initialize_agent, AgentType,load_tools
 from langchain import OpenAI
 from langchain import PromptTemplate, OpenAI, LLMChain
 from langchain.chains import SimpleSequentialChain
@@ -16,6 +17,7 @@ from langchain import SQLDatabaseChain
 JSON_DATA_LABEL = 'json_data'
 DB_CHAIN = 'db_chain'
 llm=OpenAI(temperature=0)
+db = SQLDatabase.from_uri("sqlite:///FXTrades.db")
 
 
 if JSON_DATA_LABEL not in st.session_state:
@@ -24,38 +26,12 @@ if JSON_DATA_LABEL not in st.session_state:
 if DB_CHAIN not in st.session_state:
     st.session_state[DB_CHAIN] = {}
 
-def tables_from_db():
-#    db = sqlite3.connect('switrs.sqlite')
-    db = sqlite3.connect('FXTrades.db')
-    cursor = db.cursor()
-    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
-    tables = cursor.fetchall()
-    cursor.close()
-    db.close()
-    return tables
-
-def from_db(table: str):
-#   db = sqlite3.connect('switrs.sqlite')
-    db = sqlite3.connect('FXTrades.db')
-    # Read into a panda DataFrame
-    df = pd.read_sql_query(f"SELECT * FROM {table} LIMIT 50", db)
-    db.close()
-    columns = df.columns
-    # Pick a random column for y axis
-    column_index = np.random.randint(0, columns.size, 1)
-    y_column_name = columns[column_index][0]
-    print(columns)
-    return alt.Chart(df).mark_circle().encode(x='case_id', y=y_column_name)
-
-
-
 def get_sql_agent():
 #    db = SQLDatabase.from_uri("sqlite:///switrs.sqlite")
     db = SQLDatabase.from_uri("sqlite:///FXTrades.db")
 
     toolkit = SQLDatabaseToolkit(llm=llm,db=db)
 
-
     return create_sql_agent(
         llm=llm,
         toolkit=toolkit,
@@ -65,32 +41,31 @@ def get_sql_agent():
     )
 
 def get_db_chain():
-    
-    db = SQLDatabase.from_uri("sqlite:///FXTrades.db")
 
-    _DEFAULT_TEMPLATE = """Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
+    template = """Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
     Use the following format:
     ccyBoughtccyBought
     Question: "Question here"
     SQLQuery: "SQL Query to run"
     SQLResult: "Result of the SQLQuery"
     Answer: "Final answer here"
-    
+
     Only use the following tables:
-    
+
     {table_info}
-    
+
     If someone asks any question involving client name, you need to join with Client table
     volume: you need to count records
     Amounts: you need to use USD amount
     Trades: you need to get volume of trades
     Currency Bought: you need to use ccyBought
     Currency Sold: you need to use ccySold
-    
-    
+
+
     Question: {input}"""
     PROMPT = PromptTemplate(
-        input_variables=["input", "table_info", "dialect"], template=_DEFAULT_TEMPLATE
+        input_variables=["input", "table_info", "dialect"],
+        template=template
         )
 
     return SQLDatabaseChain.from_llm(llm, db, prompt=PROMPT, verbose=True)
@@ -114,23 +89,15 @@ def plot_chart():
 
     return alt.Chart(dataframe).mark_line()
 
+# sql_agent = get_sql_agent()
+db_chain = get_db_chain()
+json_chain = get_json_chain()
 
 
 def from_gpt(query: str, plot: bool):
 
-#    sql_agent = get_sql_agent()
-    sql_agent = st.session_state[DB_CHAIN]
-    if not sql_agent:
-        sql_agent = get_db_chain()
-        st.session_state[DB_CHAIN] = sql_agent
-    
-    json_chain = get_json_chain()
-
-    chains = [sql_agent, json_chain] if plot else [sql_agent]
-
-    # Run the query using the agent executor
+    chains = [db_chain, json_chain] if plot else [db_chain]
     main_chain = SimpleSequentialChain(chains=chains, verbose=True)
-
     ans = main_chain.run(query)
 
     # Save data as json if plot
@@ -156,13 +123,6 @@ def get_response(prompt: str, *kargs):
     elif prompt_lower == 'circle chart':
         on_render = st.write
         response = circle_chart()
-    elif prompt_lower == 'db':
-        on_render = st.write
-        response = tables_from_db()
-    elif prompt_lower.startswith('db:'):
-        table = prompt_lower.split(":")[1]
-        on_render = st.write
-        response = from_db(table)
     elif prompt_lower.startswith('json'):
         p = prompt_lower.split('json ')[1]
         on_render = st.write
@@ -206,6 +166,6 @@ What trades did client {client} do in May 2022
 with st.sidebar:
     st.markdown(sidebar_text)
 
-prompt = chat.get_promt("Ask IDF Anything")
+# prompt = chat.get_promt("Ask IDF Anything")
 
-chat.process(prompt, get_response, llm)
+chat.process(get_response, llm)

chat.py

@@ -7,7 +7,6 @@ from typing import Callable
 
 RESPONSE_LABEL = 'chat_response'
 PROMPT_LABEL = 'chat_prompt'
-TEXT_INPUT_LABEL = "chat_input"
 
 class Chat:
 
@@ -18,18 +17,7 @@ class Chat:
         if PROMPT_LABEL not in st.session_state:
             st.session_state[PROMPT_LABEL] = []
 
-        if TEXT_INPUT_LABEL not in st.session_state:
-            st.session_state[TEXT_INPUT_LABEL] = ''
-
-    def get_promt(self, placeholder: str):
-        # return st.text_input(
-        #     label="ChatIDF",
-        #     placeholder=placeholder,
-        #     key="chat_widget",
-        # )
-        return st.chat_input(placeholder=placeholder)
-
-    def process(self, prompt: str, process_prompt: Callable, *args):
+    def process(self, process_prompt: Callable, *args):
         """
         process_prompt(promt: str, *args) -> tuple(Any, Callable)
             callback to process the chat promt, it takes the promt for input
@@ -45,7 +33,7 @@ class Chat:
                 on_render(response)
 
         # Compute prompt
-        if prompt:
+        if prompt:= st.chat_input("Ask IDF Anything"):
             st.session_state[PROMPT_LABEL].append(prompt)
             (response, on_render) = process_prompt(prompt, *args)
             st.session_state[RESPONSE_LABEL].append((response, on_render))