plot

app.py

@@ -1,3 +1,4 @@
+from os import write
 import streamlit as st
 import numpy as np
 import pandas as pd
@@ -11,8 +12,13 @@ from langchain import OpenAI
 from langchain import PromptTemplate, OpenAI, LLMChain
 from langchain.chains import SimpleSequentialChain
 
+JSON_DATA_LABEL = 'json_data'
 llm=OpenAI(temperature=0)
 
+
+if JSON_DATA_LABEL not in st.session_state:
+    st.session_state[JSON_DATA_LABEL] = {}
+
 def tables_from_db():
     db = sqlite3.connect('switrs.sqlite')
     cursor = db.cursor()
@@ -57,6 +63,13 @@ def get_json_chain():
         prompt=PromptTemplate.from_template(prompt_template)
     )
 
+def plot_chart():
+    json = st.session_state[JSON_DATA_LABEL]
+    if not json:
+        return "no data to plot"
+
+    return alt.Chart(json).mark_line()
+
 
 
 def from_gpt(query: str, plot: bool):
@@ -69,7 +82,12 @@ def from_gpt(query: str, plot: bool):
     # Run the query using the agent executor
     main_chain = SimpleSequentialChain(chains=chains, verbose=True)
 
-    return main_chain.run(query)
+    ans = main_chain.run(query)
+
+    # Save data as json if plot
+    st.session_state[JSON_DATA_LABEL] = ans if plot else {}
+
+    return ans
 
 def circle_chart():
     df = pd.DataFrame(np.random.randn(200, 3), columns=['a', 'b', 'c'])
@@ -96,10 +114,13 @@ def get_response(prompt: str, *kargs):
         table = prompt_lower.split(":")[1]
         on_render = st.write
         response = from_db(table)
-    elif prompt_lower.startswith('plot'):
-        p = prompt_lower.split('plot ')[1]
+    elif prompt_lower.startswith('json'):
+        p = prompt_lower.split('json ')[1]
         on_render = st.write
         response = from_gpt(p, plot=True)
+    elif prompt_lower == 'plot':
+        on_render = st.write
+        response = plot_chart()
     else:
         on_render = st.write
         response = from_gpt(prompt, plot=False)