Update app.py

app.py

@@ -1,12 +1,11 @@
-# Import Libraries
 import gradio as gr
 import spaces
 from dotenv import load_dotenv
-from implementation.answer import answer_question
+# Ensure your implementation file exposes separate retrieval and generation functions
+from implementation.answer import get_context, generate_response_stream
 
 load_dotenv(override=True)
 
-
 def format_context(context):
     result = "<h2 style='color: #ff7800;'>Relevant Context</h2>\n\n"
     for doc in context:
@@ -14,17 +13,29 @@ def format_context(context):
         result += doc.page_content + "\n\n"
     return result
 
-@spaces.GPU
-def chat(history):
+# STEP 1: CPU-only step to fetch context immediately without waiting for a GPU slot
+def retrieve_context_step(history):
     last_message = (
         "\n".join(map(str, history[-1]["content"]))
         if isinstance(history[-1]["content"], list)
         else history[-1]["content"]
     )
     prior = history[:-1]
-    answer, context = answer_question(last_message, prior, use_rewrite=True)
-    history.append({"role": "assistant", "content": answer})
-    return history, format_context(context)
+    
+    # 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
 
 
 def main():
@@ -32,20 +43,12 @@ 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,
-                )
-                message = gr.Textbox(
-                    label="Your Question",
-                    placeholder="Ask anything about Python",
-                    show_label=False,
-                )
+                chatbot = gr.Chatbot(label="💬 Conversation", height=600, type="messages")
+                message = gr.Textbox(label="Your Question", placeholder="Ask anything about Python", show_label=False)
 
             with gr.Column(scale=1):
                 context_markdown = gr.Markdown(
@@ -55,13 +58,22 @@ def main():
                     height=600,
                 )
 
-        message.submit(
-            put_message_in_chatbot, inputs=[
-                message, chatbot], outputs=[message, chatbot]
-        ).then(chat, inputs=chatbot, outputs=[chatbot, context_markdown])
+        # 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]
+        )
 
     ui.launch()
 
-
 if __name__ == "__main__":
     main()