Update app.py

app.py

@@ -1,14 +1,20 @@
 import os
 import gradio as gr
-from langchain_community.vectorstores import FAISS
-from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain.chains import RetrievalQA
 from langchain_core.prompts import PromptTemplate
-from sentence_transformers import SentenceTransformer
-from collections import OrderedDict
-import google.generativeai as genai
 from langchain.llms.base import LLM
+from collections import OrderedDict
 from typing import Optional, List
+import google.generativeai as genai
+
+# Custom utility functions
+from embeddings import (
+    load_pdf_files,
+    create_chunks,
+    get_embedding_model,
+    store_embeddings,
+    load_faiss_db
+)
 
 # Constants
 DATA_PATH = "dataFolder/"
@@ -16,71 +22,57 @@ DB_FAISS_PATH = "/tmp/vectorstore/db_faiss"
 CACHE_DIR = "/tmp/models_cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
 
-# Google AI API setup with better error handling
+# Google AI API setup
 GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
 if not GOOGLE_API_KEY:
-    print("Warning: GOOGLE_API_KEY not found in environment variables!")
+    print("⚠️ GOOGLE_API_KEY not found in environment variables!")
     print("Please set your Google API key in Hugging Face Spaces secrets.")
 else:
     genai.configure(api_key=GOOGLE_API_KEY)
 
-# Load the embedding model
-embedding_model = HuggingFaceEmbeddings(
-    model_name="rishi002/all-MiniLM-L6-v2",
-    cache_folder=CACHE_DIR
-)
-
-# Load or create FAISS database
+# Load or create FAISS vector store
 def load_or_create_faiss():
+    embedding_model = get_embedding_model()
     if not os.path.exists(DB_FAISS_PATH):
-        print("🔄 Creating FAISS Database...")
-        from embeddings import load_pdf_files, create_chunks  # Your custom chunking logic
-
+        print("🔄 FAISS index not found. Creating new index...")
         documents = load_pdf_files(DATA_PATH)
         text_chunks = create_chunks(documents)
-
-        db = FAISS.from_documents(text_chunks, embedding_model)
-        db.save_local(DB_FAISS_PATH)
+        db = store_embeddings(text_chunks, embedding_model, DB_FAISS_PATH)
     else:
-        print("✅ FAISS Database Exists. Loading...")
-
-    return FAISS.load_local(DB_FAISS_PATH, embedding_model, allow_dangerous_deserialization=True)
+        print("✅ Existing FAISS index found. Loading it...")
+        db = load_faiss_db(DB_FAISS_PATH, embedding_model)
+    return db
 
 db = load_or_create_faiss()
 
-# Custom Gemini LLM wrapper for LangChain - Fixed for Hugging Face
-# Custom Gemini LLM wrapper for LangChain - Fixed for Hugging Face
+# ✅ Custom Gemini LLM wrapper for LangChain
 class GeminiLLM(LLM):
     model_name: str = "gemini-2.0-flash"
-    
+
     class Config:
-        """Configuration for this pydantic object."""
         extra = 'forbid'
         arbitrary_types_allowed = True
-    
+
     def __init__(self, model_name: str = "gemini-2.0-flash", **kwargs):
-        # Initialize only with pydantic-defined fields
         super().__init__(model_name=model_name, **kwargs)
-    
+
     def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
         try:
-            # Use local variable, not self.model
             model = genai.GenerativeModel(self.model_name)
             response = model.generate_content(prompt)
             return response.text
         except Exception as e:
             return f"Error generating response: {str(e)}"
-    
+
     @property
     def _identifying_params(self):
         return {"model_name": self.model_name}
-    
+
     @property
     def _llm_type(self):
         return "gemini"
 
-
-# Updated prompt template with health profile
+# Prompt template with user health profile
 CUSTOM_PROMPT_TEMPLATE = """
 Use the pieces of information provided in the context to answer the user's question.
 If you don't know the answer, just say that you don't know. Don't make up an answer.
@@ -96,17 +88,13 @@ Question: {question}
 Start the answer directly.
 """
 
-# No need for API function anymore - health info will be passed as parameter
-
-# Create qa_chain using Gemini
+# QA Chain constructor
 def create_qa_chain():
     prompt = PromptTemplate(
-        template=CUSTOM_PROMPT_TEMPLATE, 
+        template=CUSTOM_PROMPT_TEMPLATE,
         input_variables=["context", "question", "health_info"]
     )
-    
     gemini_llm = GeminiLLM()
-    
     return RetrievalQA.from_chain_type(
         llm=gemini_llm,
         chain_type="stuff",
@@ -115,41 +103,38 @@ def create_qa_chain():
         chain_type_kwargs={'prompt': prompt}
     )
 
-# Main QA Chain
 qa_chain = create_qa_chain()
 
-# Modified ask_question function with health_info as parameter
+# Function to handle question asking
 def ask_question(query: str, health_info: str = "No health profile provided"):
     try:
-        # Prepare inputs for the QA chain
         qa_inputs = {
             'query': query,
             'health_info': health_info
         }
-        
-        # Get response from QA chain
         response = qa_chain.invoke(qa_inputs)
         result = response["result"]
-        
-        # Clean up response to remove duplicates
+
+        # Deduplicate output
         sentences = [s.strip() for s in result.split('.') if s.strip()]
         unique_sentences = list(OrderedDict.fromkeys(sentences))
         cleaned_result = '. '.join(unique_sentences) + '.'
-        
+
         return cleaned_result, []
-        
+
     except Exception as e:
         return f"Error: {str(e)}", []
 
-# Gradio Interface with two inputs
+# Gradio Interface
 iface = gr.Interface(
-    fn=ask_question, 
+    fn=ask_question,
     inputs=[
         gr.Textbox(label="Question", placeholder="Enter your question here..."),
         gr.Textbox(label="Health Profile", placeholder="Enter your health information (optional)...", value="No health profile provided")
-    ], 
+    ],
     outputs=["text", "json"],
     title="Medical RAG Chatbot",
     description="Ask medical questions and optionally provide your health profile for personalized responses."
 )
-iface.launch(share=True)
+
+iface.launch(share=True)