Update app.py

app.py

@@ -8,6 +8,13 @@ from qdrant_client import QdrantClient
 from datetime import datetime
 import dspy
 import json
+import google.generativeai as genai
+
+# Configure Gemini API
+genai.configure(api_key="AIzaSyBO3-HG-WcITn58PdpK7mMyvFQitoH00qA")  # Replace with your actual Gemini API key
+
+# Load Gemini model
+gemini_model = genai.GenerativeModel('gemini-pro')
 
 # === Load Models ===
 print("Loading zero-shot classifier...")
@@ -20,12 +27,6 @@ print("Loading text generation model...")
 # Use a lighter model for testing
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
 
-print("Loading text generation model...")
-model_name = "google/flan-t5-large"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-qa_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
-
 
 # === Qdrant Setup ===
 print("Connecting to Qdrant...")
@@ -64,6 +65,8 @@ class MathAnswer(dspy.Signature):
     answer = dspy.OutputField()
 
 # === DSPy Programs ===
+
+       # return dspy.Output(answer=answer, retrieved_context=context)
 class MathRetrievalQA(dspy.Program):
     def forward(self, question):
         print("Inside MathRetrievalQA...")
@@ -71,16 +74,27 @@ class MathRetrievalQA(dspy.Program):
         context = "\n".join([item["solution"] for item in context_items if "solution" in item])
         print("Context for generation:", context)
         if not context:
-            return dspy.Output(answer="", retrieved_context="")
+            return {"answer": "", "retrieved_context": ""}
 
-        prompt = f"Question: {question}\nStep-by-step solution:\n{context}\nAnswer:"
-        print("Generating answer...")
-        answer = qa_pipeline(prompt, max_new_tokens=150)[0]["generated_text"]
+        # Step 1: Generate raw answer (e.g., using GPT2 or any pipeline)
+        prompt = f"Question: {question}\nContext: {context}\nAnswer:"
+        raw_answer = qa_pipeline(prompt, max_new_tokens=100)[0]["generated_text"]
 
-        print("Generated answer:", answer)
-        return {"answer": answer, "retrieved_context": context}
+        # Step 2: Send raw answer to Gemini for formatting
+        format_prompt = f"""You are a helpful math assistant. Please format the following answer into a clear, step-by-step solution for better readability.
 
-       # return dspy.Output(answer=answer, retrieved_context=context)
+Question: {question}
+
+Raw Answer:
+{raw_answer}
+
+Formatted Step-by-Step Answer:"""
+
+        response = gemini_model.generate_content(format_prompt)
+        formatted_answer = response.text
+
+        print("Formatted answer:", formatted_answer)
+        return {"answer": formatted_answer, "retrieved_context": context}
 
 class WebFallbackQA(dspy.Program):
     def forward(self, question):