Update app.py

app.py

@@ -4,9 +4,7 @@ import pandas as pd
 from functools import partial
 from ai_classroom_suite.UIBaseComponents import *
 
-# history is a list of list  
-# [[user_input_str, bot_response_str], ...]
-
+# TODO: 
 class EchoingTutor(SlightlyDelusionalTutor):
     def add_user_message(self, user_message):
         self.conversation_memory.append([user_message, None])
@@ -36,6 +34,33 @@ def get_tutor_reply(user_message, chat_tutor):
 def get_conversation_history(chat_tutor):
     return chat_tutor.conversation_memory, chat_tutor
 
+def create_prompt_store(chat_tutor, vs_button, upload_files, openai_auth):
+    text_segs = []
+    upload_segs = []
+    
+    if upload_files:
+        print(upload_files)
+        upload_fnames = [f.name for f in upload_files]
+        upload_segs = get_document_segments(upload_fnames, 'file', chunk_size=700, chunk_overlap=100)
+    
+    # get the full list of everything
+    all_segs = text_segs + upload_segs
+    print(all_segs)
+
+    # create the vector store and update tutor
+    vs_db, vs_retriever = create_local_vector_store(all_segs, search_kwargs={"k": 2})
+    chat_tutor.vector_store = vs_db
+    chat_tutor.vs_retriever = vs_retriever
+
+    # create the tutor chain
+    if not chat_tutor.api_key_valid or not chat_tutor.openai_auth:
+        chat_tutor = embed_key(openai_auth, chat_tutor)
+    qa_chain = create_tutor_mdl_chain(kind="retrieval_qa", mdl=chat_tutor.chat_llm, retriever = chat_tutor.vs_retriever, return_source_documents=True)
+    chat_tutor.tutor_chain = qa_chain
+
+    # return the store
+    return chat_tutor, gr.update(interactive=True, value='Tutor Initialized!')
+
 
 ### Instructor Interface Helper Functions ###
 def get_instructor_prompt(fileobj):
@@ -57,6 +82,7 @@ with gr.Blocks() as demo:
 
     # Student interface
     with gr.Tab("For Students"): 
+        
         # Chatbot interface
         gr.Markdown("""
         ## Chat with the Model
@@ -80,6 +106,7 @@ with gr.Blocks() as demo:
         chat_history = gr.JSON(label = "conversation history")
         test_btn.click(get_conversation_history, inputs=[study_tutor], outputs=[chat_history, study_tutor])
 
+        # Download conversation history file
         with gr.Blocks():
             gr.Markdown("""
             ## Export Your Chat History
@@ -101,7 +128,6 @@ with gr.Blocks() as demo:
         # API Authentication functionality
         # Instead of ask students to provide key, the key is now provided by the instructor
         with gr.Box():
-            
             gr.Markdown("### OpenAI API Key ")
             gr.HTML("""<span>Embed your OpenAI API key below; if you haven't created one already, visit
             <a href="https://platform.openai.com/account/api-keys">platform.openai.com/account/api-keys</a>
@@ -121,6 +147,7 @@ with gr.Blocks() as demo:
             # embed_key(api_input, study_tutor)
 
         # Upload secret prompt functionality
+        # The instructor will provide a secret prompt/persona to the tutor
         with gr.Blocks():
             # testing purpose, change visible to False at deployment
             test_secret = gr.Textbox(label="Current secret prompt", value=os.environ.get("SECRET_PROMPT"), visible=True)
@@ -128,13 +155,20 @@ with gr.Blocks() as demo:
             file_input = gr.File(label="Load a .txt or .py file",
                             file_types=['.py', '.txt'], type="file",
                             elem_classes="short-height")
+
+            # Verify prompt content
             instructor_prompt = gr.Textbox(label="Verify your prompt content", visible=True)
             file_input.upload(fn=get_instructor_prompt, inputs=file_input, outputs=instructor_prompt)
 
-            # Set the secret prompt in this session
+            # Set the secret prompt in this session and embed it to the study tutor
             prompt_submit_btn = gr.Button("Submit")
-            prompt_submit_btn.click(fn=embed_prompt, inputs=instructor_prompt, outputs=test_secret)
-
+            prompt_submit_btn.click(
+                fn=embed_prompt, inputs=instructor_prompt, outputs=test_secret
+            ).then(
+                fn=create_prompt_store, 
+                inputs=[study_tutor, prompt_submit_btn, file_input, api_input],
+                outputs=[study_tutor, prompt_submit_btn]
+            )
             # TODO: may need a way to set the secret prompt permanently in settings/secret
     
 demo.queue().launch(server_name='0.0.0.0', server_port=7860)