merge conlict

app.py

@@ -10,6 +10,11 @@ from langchain.agents import create_sql_agent
 from langchain import OpenAI
 from langchain import PromptTemplate, OpenAI, LLMChain
 
+JSON_DATA_LABEL = 'json_data'
+
+if JSON_DATA_LABEL not in st.session_state:
+    st.session_state[JSON_DATA_LABEL] = {}
+
 
 def tables_from_db():
     db = sqlite3.connect('switrs.sqlite')
@@ -32,6 +37,10 @@ def from_db(table: str):
     print(columns)
     return alt.Chart(df).mark_circle().encode(x='case_id', y=y_column_name)
 
+def parse_into_json(ans: str) -> (str, str):
+    # json parsing goes here
+    return ans
+
 def from_gpt(query: str):
     # Create the agent executor
     db = SQLDatabase.from_uri("sqlite:///switrs.sqlite")
@@ -49,7 +58,7 @@ def from_gpt(query: str):
     )
 
     # Run the query using the agent executor
-    return agent_executor.run(query)
+    ans = agent_executor.run(query)
 
 def parse_into_json(result: str):
     prompt_template = "Reformat this {result} in JSON format"
@@ -62,13 +71,19 @@ def parse_into_json(result: str):
 
     Final_Results['text']
 
-def parse_into_chart(response: str):
-    if ":" not in response:
-        return response
+    # Extract json data from response
+    (response, json_data) = parse_into_json(ans)
 
-    data = response.split(":")[1]
-    return alt.Chart(data).mark_circle()
+    st.session_state[JSON_DATA_LABEL] = json_data
 
+    return response
+
+def plot():
+    """
+    Tries to plot the last generated answer
+    """
+    json = st.session_state[JSON_DATA_LABEL]
+    return {} if not json else alt.Chart(json).mark_circle()
 
 
 def circle_chart():
@@ -96,10 +111,12 @@ def get_response(prompt: str, *kargs):
         table = prompt_lower.split(":")[1]
         on_render = st.write
         response = from_db(table)
+    elif prompt_lower == 'plot':
+        on_render = st.write
+        response = plot()
     else:
         on_render = st.write
         response = from_gpt(prompt)
-        # response = parse_into_chart(response)
 
     return (response, on_render)