Update rag.py

rag.py

@@ -78,36 +78,39 @@ def generate_answer(query: str) -> str:
         # Load FAISS indices
         logger.info("Loading FAISS indices...")
         data_vec = load_faiss_index(
-            str(vectors_dir / "msd_data_vec"),
-            "sentence-transformers/all-MiniLM-L12-v2"
-        )
-        med_vec = load_faiss_index(
-            str(vectors_dir / "med_data_vec"),
+            str(vectors_dir / "faiss_v4"),
             "sentence-transformers/all-MiniLM-L12-v2"
         )
 
         # Create the LLM instance
         llm = ChatOpenAI(
-            model="gpt-3.5-turbo",
+            model="gpt-4o-mini",
             temperature=0,
             openai_api_key=OPENAI_API_KEY
         )
 
         # Define the prompt template
         template = """You are a helpful medical information assistant. Use the following pieces of context to answer the medical question at the end.
-        
-        Important notes:
-        - Base your answer strictly on the provided context and understandable for all readers
-        - If you don't know the answer, just say that you don't know
-        - Include relevant disclaimers about consulting healthcare professionals
-        - If suggesting medications (upon question), mention potential side effects if provided in the context
-        - Highlight if the information is general knowledge or requires professional medical advice
-        
-        Context: {context}
-        
-        Question: {question}
-        
-        Medical Information Assistant:"""
+
+                    Important notes:
+                    - Base your answer strictly on the provided context in a clear and reader-friendly way, using paragraph or bullet points as needed.
+                    - If you don't know the answer, say so honestly.
+                    - Include relevant disclaimers recommending consulting a healthcare professional for personalized advice.
+                    - When suggesting treatment options or medications (if applicable to the question), include:
+                        - Drug name
+                        - Drug class
+                        - Dosage
+                        - Frequency and duration
+                        - Potential adverse effects
+                        - Associated risks and additional recommendations
+                    - Highlight if the information is general knowledge or requires professional medical advice.
+                    - Prompt for additional details about the disease or patient characteristics if necessary to provide a thorough answer, and encourage users to ask clarifying questions.
+
+                    Context: {context}
+
+                    Question: {question}
+
+                    Medical Information Assistant:"""
 
         QA_CHAIN_PROMPT = PromptTemplate(
             input_variables=["context", "question"], 
@@ -117,15 +120,11 @@ def generate_answer(query: str) -> str:
         # Initialize and combine retrievers
         logger.info("Setting up retrieval chain...")
         data_retriever = data_vec.as_retriever()
-        med_retriever = med_vec.as_retriever()
-        combined_retriever = MergerRetriever(
-            retrievers=[data_retriever, med_retriever]
-        )
 
         # Initialize the RetrievalQA chain
         qa_chain = RetrievalQA.from_chain_type(
             llm=llm,
-            retriever=combined_retriever,
+            retriever=data_retriever,
             return_source_documents=True,
             chain_type_kwargs={"prompt": QA_CHAIN_PROMPT}
         )
@@ -138,12 +137,7 @@ def generate_answer(query: str) -> str:
         # Extracting the relevant documents from the result
         extracted_docs = result.get("source_documents", [])
         logger.info(f"Extracted documents: {extracted_docs}")  # Log the extracted documents
-        
-        # New organized printing of extracted documents
-        print("\nExtracted Documents:")
-        for doc in extracted_docs:
-            print(f"Source: {doc.metadata['source']}, Row: {doc.metadata['row']}")
-            print(f"Content: {doc.page_content}\n")
+
         
         return result["result"]