Update app.py

app.py

@@ -1,11 +1,12 @@
+# Import Libraries
 import gradio as gr
 import spaces
 from dotenv import load_dotenv
-# Ensure your implementation file exposes separate retrieval and generation functions
-from implementation.answer import get_context, generate_response_stream
+from implementation.answer import answer_question
 
 load_dotenv(override=True)
 
+
 def format_context(context):
     result = "<h2 style='color: #ff7800;'>Relevant Context</h2>\n\n"
     for doc in context:
@@ -13,29 +14,17 @@ def format_context(context):
         result += doc.page_content + "\n\n"
     return result
 
-# STEP 1: CPU-only step to fetch context immediately without waiting for a GPU slot
-def retrieve_context_step(history):
+@spaces.GPU
+def chat(history):
     last_message = (
         "\n".join(map(str, history[-1]["content"]))
         if isinstance(history[-1]["content"], list)
         else history[-1]["content"]
     )
     prior = history[:-1]
-    
-    # Run vector DB lookup or API calls on the CPU
-    context = get_context(last_message, prior, use_rewrite=True)
-    return format_context(context)
-
-# STEP 2: Dedicated GPU step that only kicks in for actual model inference
-@spaces.GPU
-def generate_answer_step(history, context_html):
-    # Append placeholder for assistant response
-    history.append({"role": "assistant", "content": ""})
-    
-    # Ensure your LLM function yields chunks of text (Streaming)
-    for text_chunk in generate_response_stream(history[:-1], context_html):
-        history[-1]["content"] = text_chunk
-        yield history
+    answer, context = answer_question(last_message, prior, use_rewrite=True)
+    history.append({"role": "assistant", "content": answer})
+    return history, format_context(context)
 
 
 def main():
@@ -43,12 +32,20 @@ def main():
         return "", history + [{"role": "user", "content": message}]
 
     with gr.Blocks(title="PyComp: Simple Python Companion") as ui:
-        gr.Markdown("# 🏢 Meet PyComp: Simple Python Companion\nAsk me anything about Python!")
+        gr.Markdown(
+            "# 🏢 Meet PyComp: Simple Python Companion\nAsk me anything about Python!")
 
         with gr.Row():
             with gr.Column(scale=1):
-                chatbot = gr.Chatbot(label="💬 Conversation", height=600, type="messages")
-                message = gr.Textbox(label="Your Question", placeholder="Ask anything about Python", show_label=False)
+                chatbot = gr.Chatbot(
+                    label="💬 Conversation",
+                    height=600,
+                )
+                message = gr.Textbox(
+                    label="Your Question",
+                    placeholder="Ask anything about Python",
+                    show_label=False,
+                )
 
             with gr.Column(scale=1):
                 context_markdown = gr.Markdown(
@@ -58,22 +55,13 @@ def main():
                     height=600,
                 )
 
-        # Chain the events: First retrieve context on CPU, then stream generation on GPU
-        submit_event = message.submit(
-            put_message_in_chatbot, 
-            inputs=[message, chatbot], 
-            outputs=[message, chatbot]
-        ).then(
-            retrieve_context_step, 
-            inputs=[chatbot], 
-            outputs=[context_markdown]
-        ).then(
-            generate_answer_step, 
-            inputs=[chatbot, context_markdown], 
-            outputs=[chatbot]
-        )
+        message.submit(
+            put_message_in_chatbot, inputs=[
+                message, chatbot], outputs=[message, chatbot]
+        ).then(chat, inputs=chatbot, outputs=[chatbot, context_markdown])
 
     ui.launch()
 
+
 if __name__ == "__main__":
-    main()
+    main()