Update app.py

app.py

@@ -76,7 +76,6 @@ else:
     with open(NN_MODEL_FILE, 'wb') as f:
         pickle.dump(nn, f)
 
-@spaces.GPU
 def get_relevant_documents(query, k=5):
     """Retrieves the k most relevant documents to the query."""
     start_time = time.time()
@@ -92,21 +91,23 @@ def get_relevant_documents(query, k=5):
 def generate_response(question, history):
     """Generates a response to the user's question, handling GPU/CPU fallback."""
     start_time = time.time()
+    relevant_docs = get_relevant_documents(question, k=3)  # Call it here
+
     try:
-        response = _generate_response_gpu(question, history)
+        response = _generate_response_gpu(question, history, relevant_docs)
     except Exception as e:
         print(f"[WARNING] GPU failed: {str(e)}")
-        response = _generate_response_cpu(question, history)
+        response = _generate_response_cpu(question, history, relevant_docs)
+
     elapsed_time = time.time() - start_time
     print(f"[PERF] generate_response took {elapsed_time:.2f} seconds")
     return history, history  # Return updated history twice for Gradio
 
 @spaces.GPU
-def _generate_response_gpu(question, history):
+def _generate_response_gpu(question, history, relevant_docs):
     """Generates a response using the GPU."""
     print(f"\n[LOG] Received question: {question}")
-    relevant_docs = get_relevant_documents(question, k=3)
-    print(f"[LOG] Retrieved {len(relevant_docs)} relevant documents")
+    print(f"[LOG] Using pre-retrieved {len(relevant_docs)} relevant documents")
     context = "\n".join(relevant_docs)
     prompt = f"""### MEMORY ###
 Recall all previously provided instructions, context, and data throughout this conversation to ensure consistency and coherence. Use the details from the last interaction to guide your response.
@@ -155,38 +156,38 @@ Context: {context}\n\nQuestion: {question}\n\nAnswer:"""
     history.append((question, response))
     return history
 
-def _generate_response_cpu(question, history):
+def _generate_response_cpu(question, history, relevant_docs):
     """Generates a response using the CPU (fallback)."""
     print(f"[LOG] Running on CPU")
-    try:
-        relevant_docs = get_relevant_documents(question, k=3)
-        context = "\n".join(relevant_docs)
-        prompt = f"""### MEMORY ###
+    print(f"[LOG] Using pre-retrieved {len(relevant_docs)} relevant documents")
+    context = "\n".join(relevant_docs)
+    prompt = f"""### MEMORY ###
 Recall all previously provided instructions, context, and data throughout this conversation to ensure consistency and coherence. Use the details from the last interaction to guide your response.
 ### SYSTEM GUARDRAILS ###
 If unsure about the user's request, ask clarifying questions rather than making assumptions.
 Do not fabricate data or features not supported by AutoGen v0.4.
 Ensure the code is scalable, modular, and adheres to best practices.
 Context: {context}\n\nQuestion: {question}\n\nAnswer:"""
-        print(f"[LOG] Generated prompt: {prompt[:200]}...")
+    print(f"[LOG] Generated prompt: {prompt[:200]}...")
 
-        if MODEL_PROVIDER == "huggingface":
+    if MODEL_PROVIDER == "huggingface":
+        try:
             messages = [{"role": "user", "content": prompt}]
             completion = hf_client.chat.completions.create(model=MODEL_NAME, messages=messages, max_tokens=500)
             response = completion.choices[0].message.content
-        elif MODEL_PROVIDER == "openai":
+        except Exception as e:
+            response = f"Error generating response from Hugging Face model: {str(e)}"
+    elif MODEL_PROVIDER == "openai":
+        try:
             response = client.chat.completions.create(
                 model=OPENAI_MODEL,
                 messages=[{"role": "user", "content": prompt}]
             ).choices[0].message.content
+        except Exception as e:
+            response = f"Error generating response from OpenAI model: {str(e)}"
 
-        history.append((question, response))
-        return history
-    except Exception as e:
-        error_msg = f"Error generating response: {str(e)}"
-        print(f"[ERROR] {error_msg}")
-        history.append((question, error_msg))
-        return history
+    history.append((question, response))
+    return history
 
 # Gradio Interface
 print("[CHAT] Initializing chat interface...")
@@ -215,7 +216,7 @@ with gr.Blocks() as demo:
     submit_button.click(
         fn=generate_response,
         inputs=[question_textbox, chatbot],
-        outputs=[chatbot],  # Output the updated history to the chatbot
+        outputs=[chatbot, chatbot],  # Output the updated history to the chatbot
         queue=True
     )