Add RAG capabilities and fix requirements.txt location

.gitignore

@@ -0,0 +1,22 @@
+# Ignore PDF files in guideline-docs (keep them local)
+Oncolife/guideline-docs/*.pdf
+
+# Ignore other binary files
+*.pdf
+*.docx
+*.doc
+
+# Keep specific files we want in git
+!Oncolife/guideline-docs/*.json
+!Oncolife/guideline-docs/*.txt
+!Oncolife/guideline-docs/oncolife_alerts_configuration*.docx
+!Oncolife/guideline-docs/WrittenChatbotDocument*.docx
+
+# System files
+.DS_Store
+__pycache__/
+*.pyc
+
+# Environment files
+oncolife_env311/
+.env 

Oncolife/app.py

@@ -4,15 +4,26 @@ OncoLife Symptom & Triage Assistant
 A medical chatbot that performs both symptom assessment and clinical triage for chemotherapy patients.
 Updated: Using BioMistral-7B base model for medical conversations.
 REBUILD: Simplified to use only base model, no adapters.
+RAG: Added document retrieval capabilities for PDFs and other reference materials.
 """
 
 import gradio as gr
 import os
 import json
+from pathlib import Path
 from transformers import AutoTokenizer, MistralForCausalLM
 import torch
 from spaces import GPU
 
+# RAG imports
+import chromadb
+from sentence_transformers import SentenceTransformer
+import PyPDF2
+import pdfplumber
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+import fitz  # PyMuPDF for better PDF handling
+
 # Force GPU detection for HF Spaces
 @GPU
 def force_gpu_detection():
@@ -37,14 +48,210 @@ class OncoLifeAssistant:
         
         self._load_model(BASE, gpu_available)
         
-        # Initialize conversation state
-        self.conversation_state = {
-            "symptoms": [],
-            "asked_ids": [],
-            "answers": {},
-            "current_symptom": None,
-            "conversation_phase": "initial"  # initial, symptom_assessment, triage, summary
-        }
+        # Load the OncoLife instructions
+        self._load_instructions()
+        
+        # Initialize RAG system
+        self._initialize_rag()
+    
+    def _load_instructions(self):
+        """Load the OncoLife instructions from the text file"""
+        try:
+            instructions_file = Path(__file__).parent / "oncolifebot_instructions.txt"
+            if instructions_file.exists():
+                with open(instructions_file, 'r') as f:
+                    self.instructions = f.read()
+                print("✅ Loaded oncolifebot_instructions.txt")
+            else:
+                print("⚠️ oncolifebot_instructions.txt not found")
+                self.instructions = ""
+        except Exception as e:
+            print(f"❌ Error loading instructions: {e}")
+            self.instructions = ""
+    
+    def _initialize_rag(self):
+        """Initialize the RAG system with document embeddings"""
+        try:
+            print("🔍 Initializing RAG system...")
+            
+            # Initialize embedding model
+            self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+            print("✅ Loaded embedding model")
+            
+            # Initialize ChromaDB
+            self.chroma_client = chromadb.Client()
+            self.collection = self.chroma_client.create_collection(
+                name="oncolife_documents",
+                metadata={"description": "OncoLife reference documents"}
+            )
+            print("✅ Initialized ChromaDB collection")
+            
+            # Load and process documents
+            self._load_documents()
+            
+        except Exception as e:
+            print(f"❌ Error initializing RAG: {e}")
+            self.embedding_model = None
+            self.collection = None
+    
+    def _load_documents(self):
+        """Load and process all reference documents"""
+        try:
+            docs_path = Path(__file__).parent / "guideline-docs"
+            print(f"📚 Loading documents from: {docs_path}")
+            
+            if not docs_path.exists():
+                print("⚠️ guideline-docs directory not found")
+                return
+            
+            # Text splitter for chunking documents
+            text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=1000,
+                chunk_overlap=200,
+                separators=["\n\n", "\n", ". ", " ", ""]
+            )
+            
+            documents_loaded = 0
+            
+            # Process PDF files
+            for pdf_file in docs_path.glob("*.pdf"):
+                try:
+                    print(f"📄 Processing PDF: {pdf_file.name}")
+                    text = self._extract_pdf_text(pdf_file)
+                    if text:
+                        chunks = text_splitter.split_text(text)
+                        self._add_chunks_to_db(chunks, pdf_file.name)
+                        documents_loaded += 1
+                        print(f"✅ Added {len(chunks)} chunks from {pdf_file.name}")
+                except Exception as e:
+                    print(f"❌ Error processing {pdf_file.name}: {e}")
+            
+            # Process JSON files
+            for json_file in docs_path.glob("*.json"):
+                try:
+                    print(f"📄 Processing JSON: {json_file.name}")
+                    with open(json_file, 'r') as f:
+                        data = json.load(f)
+                    # Convert JSON to text representation
+                    text = json.dumps(data, indent=2)
+                    chunks = text_splitter.split_text(text)
+                    self._add_chunks_to_db(chunks, json_file.name)
+                    documents_loaded += 1
+                    print(f"✅ Added {len(chunks)} chunks from {json_file.name}")
+                except Exception as e:
+                    print(f"❌ Error processing {json_file.name}: {e}")
+            
+            # Process text files
+            for txt_file in docs_path.glob("*.txt"):
+                try:
+                    print(f"📄 Processing TXT: {txt_file.name}")
+                    with open(txt_file, 'r', encoding='utf-8') as f:
+                        text = f.read()
+                    chunks = text_splitter.split_text(text)
+                    self._add_chunks_to_db(chunks, txt_file.name)
+                    documents_loaded += 1
+                    print(f"✅ Added {len(chunks)} chunks from {txt_file.name}")
+                except Exception as e:
+                    print(f"❌ Error processing {txt_file.name}: {e}")
+            
+            print(f"✅ RAG system initialized with {documents_loaded} documents")
+            
+        except Exception as e:
+            print(f"❌ Error loading documents: {e}")
+    
+    def _extract_pdf_text(self, pdf_path):
+        """Extract text from PDF using multiple methods"""
+        try:
+            # Try PyMuPDF first (better for complex PDFs)
+            try:
+                doc = fitz.open(pdf_path)
+                text = ""
+                for page in doc:
+                    text += page.get_text()
+                doc.close()
+                if text.strip():
+                    return text
+            except Exception as e:
+                print(f"PyMuPDF failed for {pdf_path.name}: {e}")
+            
+            # Fallback to pdfplumber
+            try:
+                with pdfplumber.open(pdf_path) as pdf:
+                    text = ""
+                    for page in pdf.pages:
+                        if page.extract_text():
+                            text += page.extract_text() + "\n"
+                    return text
+            except Exception as e:
+                print(f"pdfplumber failed for {pdf_path.name}: {e}")
+            
+            # Final fallback to PyPDF2
+            try:
+                with open(pdf_path, 'rb') as file:
+                    reader = PyPDF2.PdfReader(file)
+                    text = ""
+                    for page in reader.pages:
+                        text += page.extract_text() + "\n"
+                    return text
+            except Exception as e:
+                print(f"PyPDF2 failed for {pdf_path.name}: {e}")
+            
+            return None
+            
+        except Exception as e:
+            print(f"❌ Error extracting text from {pdf_path.name}: {e}")
+            return None
+    
+    def _add_chunks_to_db(self, chunks, source_name):
+        """Add document chunks to the vector database"""
+        try:
+            if not chunks or not self.collection:
+                return
+            
+            # Generate embeddings
+            embeddings = self.embedding_model.encode(chunks)
+            
+            # Add to ChromaDB
+            self.collection.add(
+                embeddings=embeddings.tolist(),
+                documents=chunks,
+                metadatas=[{"source": source_name, "chunk_id": i} for i in range(len(chunks))],
+                ids=[f"{source_name}_chunk_{i}" for i in range(len(chunks))]
+            )
+            
+        except Exception as e:
+            print(f"❌ Error adding chunks to database: {e}")
+    
+    def _retrieve_relevant_documents(self, query, top_k=5):
+        """Retrieve relevant document chunks for a query"""
+        try:
+            if not self.collection or not self.embedding_model:
+                return []
+            
+            # Generate query embedding
+            query_embedding = self.embedding_model.encode([query])
+            
+            # Search for similar documents
+            results = self.collection.query(
+                query_embeddings=query_embedding.tolist(),
+                n_results=top_k
+            )
+            
+            # Format results
+            relevant_docs = []
+            if results['documents']:
+                for i, doc in enumerate(results['documents'][0]):
+                    relevant_docs.append({
+                        'content': doc,
+                        'source': results['metadatas'][0][i]['source'],
+                        'similarity': results['distances'][0][i] if 'distances' in results else None
+                    })
+            
+            return relevant_docs
+            
+        except Exception as e:
+            print(f"❌ Error retrieving documents: {e}")
+            return []
     
     def _load_model(self, model_id, gpu_available):
         """Load the BioMistral base model"""
@@ -90,7 +297,7 @@ class OncoLifeAssistant:
             self.tokenizer = None
     
     def generate_oncolife_response(self, user_input, conversation_history):
-        """Generate response using OncoLife Symptom & Triage Assistant protocol"""
+        """Generate response using OncoLife instructions and RAG"""
         try:
             if self.model is None or self.tokenizer is None:
                 return """❌ **Model Loading Error**
@@ -104,24 +311,24 @@ Please check the Space logs for details."""
             
             print(f"🔄 Generating OncoLife response for: {user_input}")
             
-            # Create OncoLife-specific prompt
-            system_prompt = """You are the OncoLife Symptom & Triage Assistant, a medical chatbot that performs both symptom assessment and clinical triage for chemotherapy patients. Your task is to guide users through structured symptom reporting and decide whether any responses require escalation to their care team.
+            # Retrieve relevant documents using RAG
+            relevant_docs = self._retrieve_relevant_documents(user_input, top_k=3)
+            
+            # Format retrieved documents
+            context_text = ""
+            if relevant_docs:
+                context_text = "\n\n**Relevant Reference Information:**\n"
+                for i, doc in enumerate(relevant_docs):
+                    context_text += f"\n--- Source: {doc['source']} ---\n{doc['content'][:500]}...\n"
+            
+            # Create prompt using the loaded instructions and retrieved context
+            system_prompt = f"""You are the OncoLife Symptom & Triage Assistant. Follow these instructions exactly:
 
-Follow this workflow:
-1. Ask for symptoms if none provided
-2. For each symptom, ask severity rating (mild/moderate/severe)
-3. Check for red flags and immediate escalation needs
-4. Grade severity using CTCAE or UKONS criteria
-5. Ask targeted questions based on utility scoring
-6. Provide structured summary with triage recommendations
+{self.instructions}
 
-Safety protocols:
-- Never provide medical advice or treatment recommendations
-- Always redirect to oncology team for medical decisions
-- Escalate immediately for dangerous symptoms
-- Add legal disclaimer at session end
+{context_text}
 
-Current conversation state: {conversation_state}"""
+Current user input: {user_input}"""
             
             # Format conversation history
             history_text = ""
@@ -190,10 +397,6 @@ Current conversation state: {conversation_state}"""
             else:
                 answer = response.strip()
             
-            # Add legal disclaimer if this appears to be end of session
-            if any(keyword in user_input.lower() for keyword in ['done', 'finished', 'complete', 'summary']):
-                answer += "\n\n" + self._get_legal_disclaimer()
-            
             print("✅ OncoLife response generated successfully")
             return answer
             
@@ -210,14 +413,6 @@ This could be due to:
 
 Please try a simpler question or check the logs for more details."""
     
-    def _get_legal_disclaimer(self):
-        """Return the legal disclaimer as specified in the instructions"""
-        return """**Legal Disclaimer:**
-
-Patient verbalizes agreement with plan of care and understanding of the information we have gone over today and has no further comments, questions or concerns at this time. Will follow up with Doctor or ONN if symptoms worsen, do not improve, or any other symptoms develop. Agrees to seek emergency care if pt believes is needed, including for increased dizziness, depression, or any thoughts of SI.
-
-**Important:** I cannot provide medical advice or treatment recommendations. Please call your oncology team to confirm what's appropriate for your specific situation."""
-    
     def chat(self, message, history):
         """Main chat interface for OncoLife Assistant"""
         if not message.strip():
@@ -231,7 +426,7 @@ Patient verbalizes agreement with plan of care and understanding of the informat
                 "assistant": assistant_msg
             })
         
-        # Generate response using OncoLife protocol
+        # Generate response using OncoLife instructions and RAG
         response = self.generate_oncolife_response(message, conversation_history)
         
         return response
@@ -241,7 +436,7 @@ assistant = OncoLifeAssistant()
 interface = gr.ChatInterface(
     fn=assistant.chat,
     title="🏥 OncoLife Symptom & Triage Assistant",
-    description="I'm here to help assess your symptoms and determine if you need to contact your care team. Please describe your symptoms or concerns.",
+    description="I'm here to help assess your symptoms and determine if you need to contact your care team. I can access your medical guidelines and reference documents to provide accurate information.",
     examples=[
         ["I'm feeling nauseous and tired"],
         ["I have a fever of 101"],

requirements.txt

@@ -1,12 +1,12 @@
 # Medical Chatbot HF Space Requirements
 
 # Web framework
-gradio>=4.44.0
+gradio==4.44.0
 
 # Machine learning libraries - specific versions for compatibility
-torch>=2.1.0,<3.0.0
-transformers>=4.35.0,<5.0.0
-accelerate>=0.24.0
+torch==2.1.2
+transformers==4.36.2
+accelerate==0.25.0
 
 # HF Spaces GPU support
 spaces>=0.1.0
@@ -16,5 +16,14 @@ numpy>=1.21.0,<2.0.0
 requests>=2.28.0
 
 # Additional dependencies for better device handling
-safetensors>=0.4.0
+safetensors==0.4.1
 tokenizers>=0.15.0
+
+# RAG implementation
+bitsandbytes==0.41.3
+sentence-transformers==2.2.2
+chromadb==0.4.22
+pypdf2==3.0.1
+pdfplumber==0.10.3
+langchain==0.1.0
+langchain-community==0.0.10