Update app.py

app.py

@@ -1,130 +1,97 @@
 import os
 import json
+import pandas as pd
 import gradio as gr
+import openai
+import time
 from typing import List, Dict
 from langchain.document_loaders import AirtableLoader
-from langchain.vectorstores import FAISS
-from langchain.embeddings import OpenAIEmbeddings
-from langchain.chains import RetrievalQA
-from langchain.chat_models import ChatOpenAI
-from langchain.schema import SystemMessage, HumanMessage
-from langchain.text_splitter import CharacterTextSplitter
-from langchain.docstore.document import Document
 
 # Set up API keys
 AIRTABLE_API_KEY = os.getenv("AIRTABLE_API_KEY")
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-
-base_id = os.getenv("base_id")
-table_id = os.getenv("table_id")
-view = os.getenv("view")
-
-def load_airtable_data() -> List[Dict]:
-    """Load data from Airtable and return as a list of dictionaries."""
-    loader = AirtableLoader(os.environ["AIRTABLE_API_KEY"], table_id, base_id, view=view)
+BASE_ID = os.getenv("base_id")
+TABLE_ID = os.getenv("table_id")
+VIEW = os.getenv("view")
+
+# Set up OpenAI client
+openai.api_key = OPENAI_API_KEY
+client = openai.Client(api_key=OPENAI_API_KEY)
+
+# Set up assistant
+ASSISTANT_ID = os.getenv('assistant_id')
+assistant = client.beta.assistants.retrieve(ASSISTANT_ID)
+thread = client.beta.threads.create()
+
+# Function to load data from Airtable and return as a DataFrame
+def load_airtable_data() -> pd.DataFrame:
+    loader = AirtableLoader(AIRTABLE_API_KEY, TABLE_ID, BASE_ID, view=VIEW)
     documents = loader.load()
     data = []
     for doc in documents:
         try:
-            # Try to parse the JSON content
             record = json.loads(doc.page_content)
             data.append(record)
         except json.JSONDecodeError:
-            # If JSON parsing fails, use the raw content
             print(f"Warning: Could not parse JSON for document: {doc.page_content[:100]}...")
             data.append({"raw_content": doc.page_content})
-    return data
+    return pd.DataFrame(data)
 
-# Load Airtable data
+# Load Airtable data into DataFrame
 try:
-    airtable_data = load_airtable_data()
-    print(f"Successfully loaded {len(airtable_data)} records from Airtable.")
+    airtable_data_df = load_airtable_data()
+    print(f"Successfully loaded data from Airtable.")
 except Exception as e:
     print(f"Error loading Airtable data: {str(e)}")
-    airtable_data = []
-
-# Prepare documents for embedding
-text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
-documents = [Document(page_content=json.dumps(record)) for record in airtable_data]
-split_documents = text_splitter.split_documents(documents)
-
-# Initialize the embedding model and FAISS index
-embedding_model = OpenAIEmbeddings()
-vectorstore = FAISS.from_documents(split_documents, embedding_model)
-
-# Define the retrieval model
-retriever = vectorstore.as_retriever()
-
-# Define the chat model
-chat_model = ChatOpenAI(model="gpt-4o")
-
-# Define a custom prompt for context
-system_message_content = """
-#Context
-You are a school assistant with strong database Q&A capabilities. 
-Your role is to help educators keep track of students' assignments in different classes.
-This is a complex problem, because each student has their own menu of classes (they choose their classes), so that it can be hard for a teacher to know what assignments their students might have
-in other classes. Solving this requires carefully analyzing a database.
-You have acces to a list of records with the following format:
--Class
--List of students enrolled in the class (student codes)
--List of DUE dates, when students turn in work done at home
--List of DO dates, when students take assessments in class
--List of DUE assignments
--List of DO assessments
-The policy is that students cannot have 0, 1, or 2 DO assessments the same day, but not 3 or more.
-HOWEVER, DUE assignments do not count towards this total.
-
-#Instructions
-When asked a question about a class and a date, follow this thought process internally without sharing it with the user. Only share the conclusion.
-[Thought process (hidden internal state):
--Store a dictionary of all students enrolled in this class
--Look up all other classes one by one, and check if it has a DO date the same day, and if some of the same students are enrolled. Make sure to double-check the student codes and rosters
--Deduce whether any student in the first class already has 2 or more "DO" assessments on that day in the other classes
--Double check by running the same steps again and testing your conclusion. Pay careful attention to student codes and dates.
-
-Keep this procedure in mind when answering other questions as well.
-]
+    airtable_data_df = pd.DataFrame()
 
-This is a comple time. Take your time and be thorough and this is critically important to the school.
-Be concise and factual in your answers unless asked for more details. 
-Base all of your answers on the data provided.
-Start your response with a thumb up 👍 or a thumb down 👎 depending on your answer.
-Double-check your answers, and if you don't know the answer, say that you don't know.
-"""
-
-# Create the QA chain
-qa_chain = RetrievalQA.from_chain_type(
-    llm=chat_model,
-    chain_type="stuff",
-    retriever=retriever,
-    return_source_documents=True
-)
-
-def ask_question(question: str) -> str:
-    """Ask a question about the Airtable data."""
-    # Combine the system message and user question
-    full_query = f"{system_message_content}\n\nHuman: {question}\n\nAssistant:"
+# Function to chat with the assistant
+def chat_with_assistant(message: str, dataframe: pd.DataFrame) -> str:
+    dataframe_csv = dataframe.to_csv(index=False)
+    
+    full_message = f"""
+    You are an assistant with code interpreter capabilities.
+    I have a DataFrame with the following content:
+    {dataframe_csv}
+    
+    Here is my question: {message}
+    
+    Please use the DataFrame and code to provide an answer.
+    """
+    
+    client.beta.threads.messages.create(
+        thread_id=thread.id,
+        role="user",
+        content=full_message
+    )
+    
+    run = client.beta.threads.runs.create(
+        thread_id=thread.id,
+        assistant_id=ASSISTANT_ID
+    )
     
-    # Get the response from the QA chain
-    response = qa_chain({"query": full_query})
+    while True:
+        run_status = client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
+        if run_status.status == 'completed':
+            messages = client.beta.threads.messages.list(thread_id=thread.id)
+            assistant_response = messages.data[0].content[0].text.value
+            break
+        time.sleep(1)
     
-    # Return the response content
-    return response['result']
+    return assistant_response
 
 # Define the Gradio interface
 def gradio_interface(question: str) -> str:
-    return ask_question(question)
+    return chat_with_assistant(question, airtable_data_df)
 
 # Set up Gradio interface
 iface = gr.Interface(
     fn=gradio_interface,
     inputs="text",
-    #outputs="text",
     outputs=gr.Markdown(),
     title="📅 Summative Assessment Tracker",
     description="I am here to help you schedule summative assessments for your students"
 )
 
 # Launch the Gradio app
-iface.launch(debug=True)
+iface.launch(debug=True)