clenaup

app.py

@@ -11,11 +11,7 @@ from langchain import OpenAI
 from langchain import PromptTemplate, OpenAI, LLMChain
 from langchain.chains import SimpleSequentialChain
 
-JSON_DATA_LABEL = 'json_data'
-
-if JSON_DATA_LABEL not in st.session_state:
-    st.session_state[JSON_DATA_LABEL] = {}
-
+llm=OpenAI(temperature=0)
 
 def tables_from_db():
     db = sqlite3.connect('switrs.sqlite')
@@ -40,7 +36,7 @@ def from_db(table: str):
 
 
 
-def get_sql_agent(llm):
+def get_sql_agent():
     db = SQLDatabase.from_uri("sqlite:///switrs.sqlite")
 
     toolkit = SQLDatabaseToolkit(llm=llm,db=db)
@@ -54,7 +50,7 @@ def get_sql_agent(llm):
         return_intermediate_steps=True
     )
 
-def get_json_chain(llm):
+def get_json_chain():
     prompt_template = "Reformat this {result} in JSON format"
     return LLMChain(
         llm=llm,
@@ -63,10 +59,10 @@ def get_json_chain(llm):
 
 
 
-def from_gpt(query: str, plot: bool, llm):
+def from_gpt(query: str, plot: bool):
 
-    sql_agent = get_sql_agent(llm)
-    json_chain = get_json_chain(llm)
+    sql_agent = get_sql_agent()
+    json_chain = get_json_chain()
 
     chains = [sql_agent, json_chain] if plot else [sql_agent]
 
@@ -111,7 +107,6 @@ def get_response(prompt: str, llm, *kargs):
     return (response, on_render)
 
 
-llm=OpenAI(temperature=0)
 
 
 chat = idf_chat.Chat()