Update app.py

app.py

@@ -12,6 +12,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from sentence_transformers import SentenceTransformer
 import faiss
 import numpy as np
+import google.generativeai as genai
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -31,14 +32,26 @@ os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 
 # Global model variables
 embeddings_model = None
+gemini_model = None
 
 def initialize_models():
-    """Initialize lightweight embedding model only"""
-    global embeddings_model
+    """Initialize embedding model and Gemini API"""
+    global embeddings_model, gemini_model
     try:
-        logger.info("Initializing embedding model...")
+        logger.info("Initializing models...")
         
-        # Load small embeddings model (only 22MB)
+        # Get Gemini API key from environment
+        gemini_api_key = os.environ.get("GEMINI_API_KEY")
+        if not gemini_api_key:
+            logger.error("GEMINI_API_KEY not found in environment variables!")
+            return False
+        
+        # Configure Gemini
+        genai.configure(api_key=gemini_api_key)
+        gemini_model = genai.GenerativeModel('gemini-2.0-flash-exp')
+        logger.info("Gemini API configured successfully")
+        
+        # Load embeddings model (only 22MB!)
         logger.info("Loading all-MiniLM-L6-v2...")
         embeddings_model = SentenceTransformer(
             "all-MiniLM-L6-v2",
@@ -46,11 +59,11 @@ def initialize_models():
             cache_folder="/tmp"
         )
         
-        logger.info("Model initialized successfully")
+        logger.info("Models initialized successfully")
         return True
         
     except Exception as e:
-        logger.error(f"Error initializing model: {str(e)}")
+        logger.error(f"Error initializing models: {str(e)}")
         import traceback
         traceback.print_exc()
         return False
@@ -66,7 +79,6 @@ def load_pdf(filepath: str) -> List[str]:
             logger.warning("No pages extracted from PDF")
             return []
         
-        # Combine page content
         docs = [page.page_content for page in pages if page.page_content.strip()]
         logger.info(f"Loaded {len(pages)} pages")
         return docs
@@ -82,7 +94,6 @@ def create_faiss_index(chunks: List[str]):
     try:
         logger.info(f"Creating FAISS index for {len(chunks)} chunks")
         
-        # Encode chunks in batches to save memory
         batch_size = 32
         embeddings_list = []
         
@@ -93,14 +104,12 @@ def create_faiss_index(chunks: List[str]):
         
         embeddings = np.vstack(embeddings_list).astype('float32')
         
-        # Create FAISS index
         dim = embeddings.shape[1]
         index = faiss.IndexFlatL2(dim)
         index.add(embeddings)
         
         logger.info(f"FAISS index created with dimension {dim}")
         
-        # Clean up
         del embeddings_list
         gc.collect()
         
@@ -112,19 +121,16 @@ def create_faiss_index(chunks: List[str]):
         traceback.print_exc()
         raise
 
-def retrieve_context(question: str, chunks: List[str], index, k: int = 3) -> str:
+def retrieve_context(question: str, chunks: List[str], index, k: int = 5) -> str:
     """Retrieve relevant context for question"""
     try:
-        # Encode question
         q_embedding = embeddings_model.encode([question])
         q_embedding = np.array(q_embedding).astype('float32')
         
-        # Search index
         distances, indices = index.search(q_embedding, k)
         
-        # Get relevant chunks with distances
         relevant_chunks = []
-        for i, dist in zip(indices[0], distances[0]):
+        for i in indices[0]:
             if i < len(chunks):
                 relevant_chunks.append(chunks[i])
         
@@ -137,37 +143,37 @@ def retrieve_context(question: str, chunks: List[str], index, k: int = 3) -> str
         logger.error(f"Error retrieving context: {str(e)}")
         return ""
 
-def generate_extractive_answer(question: str, context: str) -> str:
-    """Generate extractive answer from context (no LLM needed)"""
+def generate_answer_with_gemini(question: str, context: str) -> str:
+    """Generate answer using Gemini API"""
     try:
-        # Simple extractive approach: return most relevant sentences
-        sentences = context.split('.')
-        
-        # Score sentences by keyword overlap with question
-        question_words = set(question.lower().split())
-        scored_sentences = []
-        
-        for sent in sentences:
-            if len(sent.strip()) > 20:  # Only consider substantial sentences
-                sent_words = set(sent.lower().split())
-                score = len(question_words & sent_words)
-                scored_sentences.append((score, sent.strip()))
-        
-        # Sort by score and get top sentences
-        scored_sentences.sort(reverse=True)
-        top_sentences = [sent for score, sent in scored_sentences[:3] if score > 0]
-        
-        if top_sentences:
-            answer = ". ".join(top_sentences) + "."
-            logger.info(f"Generated extractive answer")
-            return answer
-        else:
-            # Return first part of context if no good match
-            return context[:500] + "..."
+        logger.info(f"Generating answer with Gemini for: {question}")
+        
+        prompt = f"""You are a helpful AI assistant that answers questions based on the provided context from a PDF document.
+
+Context from PDF:
+{context}
+
+Question: {question}
+
+Instructions:
+- Answer the question clearly and concisely based ONLY on the context provided
+- If the context doesn't contain enough information to answer, say so
+- Provide a well-structured, informative answer
+- If asked to summarize, provide a comprehensive summary
+
+Answer:"""
+        
+        response = gemini_model.generate_content(prompt)
+        answer = response.text.strip()
+        
+        logger.info(f"Generated answer: {answer[:100]}...")
+        return answer
         
     except Exception as e:
-        logger.error(f"Error generating answer: {str(e)}")
-        return context[:500] + "..."
+        logger.error(f"Error generating answer with Gemini: {str(e)}")
+        import traceback
+        traceback.print_exc()
+        return "Sorry, I couldn't generate an answer. Please try again."
 
 def cleanup_temp_files(filepath):
     """Clean up temporary files"""
@@ -181,9 +187,10 @@ def cleanup_temp_files(filepath):
 @app.route('/')
 def home():
     return jsonify({
-        "message": "PDF QA API is running!",
+        "message": "PDF QA API with Gemini 2.0 Flash is running!",
         "status": "healthy",
-        "model": "all-MiniLM-L6-v2 (extractive)"
+        "model": "Google Gemini 2.0 Flash",
+        "embeddings": "all-MiniLM-L6-v2"
     })
 
 @app.route('/health')
@@ -213,8 +220,8 @@ def ask():
         
         # Split into chunks
         splitter = RecursiveCharacterTextSplitter(
-            chunk_size=500,
-            chunk_overlap=50,
+            chunk_size=800,
+            chunk_overlap=100,
             separators=["\n\n", "\n", ". ", " ", ""]
         )
         
@@ -222,11 +229,6 @@ def ask():
         for doc in docs:
             chunks.extend(splitter.split_text(doc))
         
-        # Limit chunks to avoid memory issues
-        if len(chunks) > 200:
-            logger.warning(f"Too many chunks ({len(chunks)}), limiting to 200")
-            chunks = chunks[:200]
-        
         logger.info(f"Created {len(chunks)} chunks")
         
         if not chunks:
@@ -235,14 +237,14 @@ def ask():
         # Create FAISS index
         index, embeddings = create_faiss_index(chunks)
         
-        # Retrieve context
-        context = retrieve_context(question, chunks, index, k=5)
+        # Retrieve context (get more chunks for Gemini since it can handle it)
+        context = retrieve_context(question, chunks, index, k=7)
         
         if not context:
             return jsonify({"error": "Failed to retrieve context from PDF"}), 500
         
-        # Generate extractive answer
-        answer = generate_extractive_answer(question, context)
+        # Generate answer with Gemini
+        answer = generate_answer_with_gemini(question, context)
         
         if not answer or len(answer.strip()) < 10:
             return jsonify({"error": "Failed to generate answer from PDF content"}), 500
@@ -254,7 +256,7 @@ def ask():
         
         return jsonify({
             "answer": answer,
-            "method": "extractive"
+            "model": "gemini-2.0-flash-exp"
         })
 
     except Exception as e:
@@ -265,7 +267,6 @@ def ask():
     finally:
         if filepath:
             cleanup_temp_files(filepath)
-        # Force garbage collection
         gc.collect()
 
 if __name__ == "__main__":