smart-models/Placebo_AI
0
1from langchain_groq import ChatGroq2from langchain_huggingface import HuggingFaceEmbeddings3from langchain_chroma import Chroma4from langchain_classic.chains import ConversationalRetrievalChain5from langchain_classic.memory import ConversationBufferMemory6from langchain_core.prompts import PromptTemplate7from langchain_core.retrievers import BaseRetriever8from langchain_core.callbacks import CallbackManagerForRetrieverRun9from typing import List10from langchain_core.documents import Document11import os12import re13 14class MedicalChatbot:15 def __init__(self, vector_store_path=None):16 if vector_store_path is None:17 import os18 # Intelligently scan for the database (Local vs Hugging Face Dataset Mounts)19 possible_paths = [20 os.path.join(os.path.dirname(os.path.dirname(__file__)), "db", "vector_store"),21 "/data/TrinetraLabs/Placebo_AI_DB/db/vector_store",22 "/data/smart-models/Placebo_AI_DB/db/vector_store",23 "/data/db/vector_store"24 ]25 26 vector_store_path = possible_paths[0] # Default fallback27 for p in possible_paths:28 if os.path.exists(p):29 vector_store_path = p30 break31 32 # --- FUSE / NETWORK MOUNT FIX ---33 # SQLite struggles with file locks on HuggingFace Dataset mounts.34 # We copy the sqlite file to local fast storage and symlink the indices.35 if vector_store_path.startswith("/data/"):36 import shutil37 local_db_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "db", "vector_store")38 os.makedirs(local_db_path, exist_ok=True)39 40 print(f"Dataset mount detected. Copying SQLite DB to local storage to prevent locking timeout... (This may take a minute)")41 42 # Copy the sqlite file43 sqlite_src = os.path.join(vector_store_path, "chroma.sqlite3")44 sqlite_dst = os.path.join(local_db_path, "chroma.sqlite3")45 if os.path.exists(sqlite_src) and not os.path.exists(sqlite_dst):46 shutil.copy2(sqlite_src, sqlite_dst)47 48 # Symlink all other directories (the heavy vector indices)49 for item in os.listdir(vector_store_path):50 if item == "chroma.sqlite3":51 continue52 src_item = os.path.join(vector_store_path, item)53 dst_item = os.path.join(local_db_path, item)54 if os.path.isdir(src_item) and not os.path.exists(dst_item):55 os.symlink(src_item, dst_item)56 57 vector_store_path = local_db_path58 print("Database ready!")59 self.embeddings = HuggingFaceEmbeddings(60 model_name="nomic-ai/nomic-embed-text-v1.5",61 model_kwargs={'trust_remote_code': True}62 )63 64 # Connect to existing vector store65 if not os.path.exists(vector_store_path):66 raise FileNotFoundError(f"Vector store not found at {vector_store_path}. Please run indexer.py first.")67 68 self.vectorstore = Chroma(69 persist_directory=vector_store_path,70 embedding_function=self.embeddings71 )72 73 # Use Llama 3 for grounded inference via Groq74 self.llm = ChatGroq(75 model_name="llama-3.1-8b-instant",76 temperature=0.0, # Zero creativity for clinical safety77 streaming=True,78 groq_api_key=os.getenv("GROQ_API_KEY")79 )80 81 self.memory = ConversationBufferMemory(82 memory_key="chat_history",83 return_messages=True,84 output_key="answer"85 )86 87 self.prompt_template = """88 You are 'Placebo AI', a specialized Medical Knowledge Retrieval Assistant. 89 Your primary role is to extract and summarize medical information with 100% fidelity.90 91 STRICT CLINICAL PROTOCOLS:92 1. GROUNDING: Provide information ONLY from the provided CONTEXT. 93 2. DATA PRECISION: Preserve all numerical values, chemical formulas, and dosages EXACTLY as written.94 3. STRUCTURE & FORMATTING: 95 - **Pain Points & Symptoms**: Always bold key issues, pathological symptoms, patient pain points, or critical clinical risks (e.g., **severe nausea**, **respiratory depression**).96 - **Reactions & Mechanisms**: Stated clearly, using step-by-step reaction pathways, equations, or numbered sequences to detail chemical or physiological mechanisms.97 - Use Markdown Tables for data comparisons.98 - Use Bullet Points for lists.99 4. CITATION: You MUST end every factual statement with its source in [Book Name, Page #] format.100 5. FALLBACK: If the answer is truly not present in the context, say: "I'm sorry, but I couldn't find specific details on this topic in the current medical textbook library."101 102 CONTEXT: 103 {context}104 105 CHAT HISTORY: 106 {chat_history}107 108 QUESTION: 109 {question}110 111 PLACEBO AI RESPONSE (Start with 'Based on the clinical data in the textbooks:'):112 Based on the clinical data in the textbooks:113 """114 115 self.QA_PROMPT = PromptTemplate(116 template=self.prompt_template,117 input_variables=["context", "chat_history", "question"]118 )119 120 # Custom Retriever Logic for Exhaustive Reference Search121 class KeywordAugmentedRetriever(BaseRetriever):122 vectorstore: Chroma123 k: int = 10124 track_filter_state: str = "unified"125 126 def _get_relevant_documents(127 self, query: str, *, run_manager: CallbackManagerForRetrieverRun128 ) -> List[Document]:129 return self.get_relevant_documents_with_filter(query, track_filter=self.track_filter_state)130 131 def get_relevant_documents_with_filter(132 self, query: str, track_filter: str = "unified"133 ) -> List[Document]:134 # Build filter dictionary135 filter_dict = None136 if track_filter == "mbbs":137 filter_dict = {"track": "mbbs"}138 elif track_filter == "pharmacy":139 filter_dict = {"track": "pharmacy"}140 141 # 1. Vector Search (MMR for diversity)142 docs = self.vectorstore.max_marginal_relevance_search(143 query, k=self.k, fetch_k=50, lambda_mult=0.3, filter=filter_dict144 )145 146 return docs147 148 self.custom_retriever = KeywordAugmentedRetriever(vectorstore=self.vectorstore, k=10)149 150 self.chain = ConversationalRetrievalChain.from_llm(151 llm=self.llm,152 retriever=self.custom_retriever,153 memory=self.memory,154 combine_docs_chain_kwargs={"prompt": self.QA_PROMPT},155 return_source_documents=True156 )157 158 def ask(self, query, mode="unified"):159 """160 Processes a medical query using exhaustive retrieval and returns a verified response.161 Routes to the correct domain track.162 """163 self.custom_retriever.track_filter_state = mode164 result = self.chain.invoke({"question": query})165 answer = result["answer"]166 167 # Deduplicate and format clinical sources168 seen_sources = set()169 formatted_sources = []170 171 for doc in result["source_documents"]:172 book = doc.metadata.get("book_name", "Unknown Source")173 page = doc.metadata.get("page_number", "N/A")174 source_key = f"{book}_P{page}"175 176 if source_key not in seen_sources:177 formatted_sources.append({178 "book": book,179 "page": page,180 "subject": doc.metadata.get("subject", "General"),181 "image_path": doc.metadata.get("image_path"),182 "snippet": doc.page_content[:200] + "..."183 })184 seen_sources.add(source_key)185 186 return {187 "answer": answer,188 "sources": formatted_sources,189 "citation_footer": "\n\n---\n**Verified Clinical Sources:**\n" + \190 "\n".join([f"- {s['book']} (Page {s['page']})" for s in formatted_sources])191 }192 193if __name__ == "__main__":194 # Internal test check195 try:196 bot = MedicalChatbot()197 print("Medical Chatbot Initialized Successfully.")198 except Exception as e:199 print(f"Chatbot wait-state: {e}")