cleanup

README.md

@@ -21,3 +21,11 @@ $ python -m venv .venv
 $ pip install streamlit
 $ python3 app.py
 ```
+
+There are 4 special promts:
+- `db`: returns a list of the table names
+- `db:<table>`: loads the table `<table>` and returns a chart where:
+  - x: column `case_id`
+  - y: random column from table
+- `line chart`: a randomly generated line chart
+- `circle chart`: a randomly generated circle chart

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import numpy as np
 import pandas as pd
 import altair as alt
-from chat import chat as idf_chat
+import chat as idf_chat
 import sqlite3
 
 
@@ -17,9 +17,15 @@ def tables_from_db():
 
 def from_db(table: str):
     db = sqlite3.connect('switrs.sqlite')
-    df = pd.read_sql_query(f"SELECT * FROM {table} LIMIT 100", db)
+    # Read into a panda DataFrame
+    df = pd.read_sql_query(f"SELECT * FROM {table} LIMIT 50", db)
     db.close()
-    return alt.Chart(df).mark_circle()
+    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)
 
 
 
@@ -29,6 +35,7 @@ def circle_chart():
         x='a', y='b', size='c', color='c', tooltip=['a', 'b', 'c']
     )
 
+# Parse the prompt to pick an example and a render function
 def get_response(prompt: str, *kargs):
     on_render = st.write
     response = f"Here's what you asked: '{prompt}'"
@@ -40,12 +47,13 @@ def get_response(prompt: str, *kargs):
     elif prompt_lower == 'circle chart':
         on_render = st.write
         response = circle_chart()
-    elif prompt_lower == 'sql':
+    elif prompt_lower == 'db':
         on_render = st.write
         response = tables_from_db()
-    elif prompt_lower == 'parties':
+    elif prompt_lower.startswith('db:'):
+        table = prompt_lower.split(":")[1]
         on_render = st.write
-        response = from_db(prompt_lower)
+        response = from_db(table)
 
     return (response, on_render)
 
@@ -53,4 +61,3 @@ def get_response(prompt: str, *kargs):
 
 chat = idf_chat.Chat()
 chat.get_input(get_response)
-chat.display()

chat/chat.py

@@ -1,39 +0,0 @@
-import streamlit as st
-from typing import Callable
-
-
-
-RESPONSE_LABEL = 'chat_response'
-PROMPT_LABEL = 'chat_user_input'
-
-class Chat:
-    def __init__(self): 
-        if RESPONSE_LABEL not in st.session_state:
-            st.session_state[RESPONSE_LABEL] = []
-
-        if PROMPT_LABEL not in st.session_state:
-            st.session_state[PROMPT_LABEL] = []
-
-    def get_input(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
-            and returns a tuple with the response and a render callback
-        """
-        promt = st.chat_input(placeholder="Ask Me Anything", key='chat_text_input')
-        # promt = st.text_input(label="Ask Me Anything", key='chat_text_input')
-
-        if promt:
-            st.session_state[PROMPT_LABEL].append(promt)
-            (response, on_render) = process_prompt(promt, *args)
-            st.session_state[RESPONSE_LABEL].append((response, on_render))
-
-
-    def display(self):
-        if st.session_state[RESPONSE_LABEL]:
-            messages = zip(st.session_state[PROMPT_LABEL], st.session_state[RESPONSE_LABEL])
-            for prompt, (response, on_render) in messages:
-                st.chat_message("user").write(prompt)
-                with st.chat_message("assistant"):
-                    on_render(response)
-