RAGSample.py

import os
from dotenv import load_dotenv
from langchain_community.document_loaders import WebBaseLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_ollama import ChatOllama
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.retrievers import BaseRetriever
from langchain_core.runnables import Runnable
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
import chromadb
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import pandas as pd
from typing import Optional, List
import re

# Disable ChromaDB telemetry to avoid the error
os.environ["ANONYMIZED_TELEMETRY"] = "False"
os.environ["CHROMA_SERVER_HOST"] = "localhost"
os.environ["CHROMA_SERVER_HTTP_PORT"] = "8000"


class ImprovedTFIDFEmbeddings(Embeddings):
    """Improved TF-IDF based embedding function with better preprocessing."""
    
    def __init__(self):
        self.vectorizer = TfidfVectorizer(
            max_features=5000,
            stop_words='english',
            ngram_range=(1, 3),
            min_df=1,
            max_df=0.85,
            lowercase=True,
            strip_accents='unicode',
            analyzer='word'
        )
        self.fitted = False
        self.documents = []
    
    def embed_documents(self, texts):
        """Create embeddings for a list of texts."""
        if not self.fitted:
            self.documents = texts
            self.vectorizer.fit(texts)
            self.fitted = True
        
        # Transform texts to TF-IDF vectors
        tfidf_matrix = self.vectorizer.transform(texts)
        
        # Convert to dense arrays and normalize
        embeddings = []
        for i in range(tfidf_matrix.shape[0]):
            embedding = tfidf_matrix[i].toarray().flatten()
            # Normalize the embedding
            norm = np.linalg.norm(embedding)
            if norm > 0:
                embedding = embedding / norm
            # Pad or truncate to 512 dimensions
            if len(embedding) < 512:
                embedding = np.pad(embedding, (0, 512 - len(embedding)))
            else:
                embedding = embedding[:512]
            embeddings.append(embedding.tolist())
        
        return embeddings
    
    def embed_query(self, text):
        """Create embedding for a single query text."""
        if not self.fitted:
            # If not fitted, fit with just this text
            self.vectorizer.fit([text])
            self.fitted = True
        
        # Transform query to TF-IDF vector
        tfidf_matrix = self.vectorizer.transform([text])
        embedding = tfidf_matrix[0].toarray().flatten()
        
        # Normalize the embedding
        norm = np.linalg.norm(embedding)
        if norm > 0:
            embedding = embedding / norm
        
        # Pad or truncate to 512 dimensions
        if len(embedding) < 512:
            embedding = np.pad(embedding, (0, 512 - len(embedding)))
        else:
            embedding = embedding[:512]
        
        return embedding.tolist()


class SmartFAQRetriever(BaseRetriever):
    """Smart retriever optimized for FAQ datasets with semantic similarity."""
    
    def __init__(self, documents: List[Document], k: int = 4):
        super().__init__()
        self._documents = documents
        self._k = k
        self._vectorizer = None  # Use private attribute
    
    @property
    def documents(self):
        return self._documents
    
    @property
    def k(self):
        return self._k
    
    def _get_relevant_documents(self, query: str) -> List[Document]:
        """Retrieve documents based on semantic similarity."""
        # Ensure vectorizer is fitted
        if not hasattr(self, '_vectorizer') or self._vectorizer is None or not hasattr(self._vectorizer, 'vocabulary_') or not self._vectorizer.vocabulary_:
            print("[SmartFAQRetriever] Fitting vectorizer...")
            self._vectorizer = TfidfVectorizer(
                max_features=3000,
                stop_words='english',
                ngram_range=(1, 2),
                min_df=1,
                max_df=0.9
            )
            questions = []
            for doc in self._documents:
                if "QUESTION:" in doc.page_content:
                    question_part = doc.page_content.split("ANSWER:")[0]
                    question = question_part.replace("QUESTION:", "").strip()
                    questions.append(question)
                else:
                    questions.append(doc.page_content)
            self._vectorizer.fit(questions)
        query_lower = query.lower().strip()
        
        # Extract questions from documents
        questions = []
        for doc in self._documents:
            if "QUESTION:" in doc.page_content:
                question_part = doc.page_content.split("ANSWER:")[0]
                question = question_part.replace("QUESTION:", "").strip()
                questions.append(question)
            else:
                questions.append(doc.page_content)
        
        # Transform query and questions to TF-IDF vectors
        query_vector = self._vectorizer.transform([query_lower])
        question_vectors = self._vectorizer.transform(questions)
        
        # Calculate cosine similarities
        similarities = cosine_similarity(query_vector, question_vectors).flatten()
        
        # Get top k documents
        top_indices = similarities.argsort()[-self._k:][::-1]
        
        # Return documents with highest similarity scores
        relevant_docs = [self._documents[i] for i in top_indices if similarities[i] > 0.1]
        
        if not relevant_docs:
            # Fallback to first k documents if no good matches
            relevant_docs = self._documents[:self._k]
        
        return relevant_docs
    
    async def _aget_relevant_documents(self, query: str) -> List[Document]:
        """Async version of get_relevant_documents."""
        return self._get_relevant_documents(query)

def setup_retriever(use_kaggle_data: bool = False, kaggle_dataset: Optional[str] = None, 
                   kaggle_username: Optional[str] = None, kaggle_key: Optional[str] = None,
                   use_local_mental_health_data: bool = False) -> BaseRetriever:
    """
    Creates a vector store with documents from test data, Kaggle datasets, or local mental health data.
    
    Args:
        use_kaggle_data: Whether to load Kaggle data instead of test documents
        kaggle_dataset: Kaggle dataset name (e.g., 'username/dataset-name')
        kaggle_username: Your Kaggle username (optional if using kaggle.json)
        kaggle_key: Your Kaggle API key (optional if using kaggle.json)
        use_local_mental_health_data: Whether to load local mental health FAQ data
    """
    print("Setting up the retriever...")
    
    if use_local_mental_health_data:
        try:
            print("Loading mental health FAQ data from local file...")
            mental_health_file = "data/Mental_Health_FAQ.csv"
            
            if not os.path.exists(mental_health_file):
                print(f"Mental health FAQ file not found: {mental_health_file}")
                use_local_mental_health_data = False
            else:
                # Load mental health FAQ data
                df = pd.read_csv(mental_health_file)
                documents = []
                
                for _, row in df.iterrows():
                    question = row['Questions']
                    answer = row['Answers']
                    # Create document in FAQ format
                    content = f"QUESTION: {question}\nANSWER: {answer}"
                    documents.append(Document(page_content=content))
                
                print(f"Loaded {len(documents)} mental health FAQ documents")
                for i, doc in enumerate(documents[:3]):
                    print(f"Sample FAQ {i+1}: {doc.page_content[:200]}...")
                
        except Exception as e:
            print(f"Error loading mental health data: {e}")
            use_local_mental_health_data = False
    
    if use_kaggle_data and kaggle_dataset:
        try:
            from src.kaggle_loader import KaggleDataLoader
            
            print(f"Loading Kaggle dataset: {kaggle_dataset}")
            # Create loader without parameters - it will auto-load from kaggle.json
            loader = KaggleDataLoader()
            
            # Download the dataset
            dataset_path = loader.download_dataset(kaggle_dataset)
            
            # Load documents based on file type - only process files from this specific dataset
            documents = []
            
            # Get the dataset name to identify the correct files
            dataset_name = kaggle_dataset.split('/')[-1]
            print(f"Processing files in dataset directory: {dataset_path}")
            
            for file in os.listdir(dataset_path):
                file_path = os.path.join(dataset_path, file)
                
                if file.endswith('.csv'):
                    print(f"Loading CSV file: {file}")
                    # For FAQ datasets, use the improved loading method
                    if 'faq' in file.lower() or 'mental' in file.lower():
                        documents.extend(loader.load_csv_dataset(file_path, [], chunk_size=50))
                    else:
                        # For other CSV files, use first few columns as text
                        df = pd.read_csv(file_path)
                        text_columns = df.columns[:3].tolist()  # Use first 3 columns
                        documents.extend(loader.load_csv_dataset(file_path, text_columns, chunk_size=50))
                
                elif file.endswith('.json'):
                    print(f"Loading JSON file: {file}")
                    documents.extend(loader.load_json_dataset(file_path))
                
                elif file.endswith('.txt'):
                    print(f"Loading text file: {file}")
                    documents.extend(loader.load_text_dataset(file_path))
            
            print(f"Loaded {len(documents)} documents from Kaggle dataset")
            for i, doc in enumerate(documents[:3]):
                print(f"Sample doc {i+1}: {doc.page_content[:200]}")
            
        except Exception as e:
            print(f"Error loading Kaggle data: {e}")
            print("Falling back to test documents...")
            use_kaggle_data = False
    
    if not use_kaggle_data and not use_local_mental_health_data:
        # No test documents - use mental health data as default
        print("No specific data source specified, loading mental health FAQ data as default...")
        try:
            mental_health_file = "data/Mental_Health_FAQ.csv"
            
            if not os.path.exists(mental_health_file):
                raise FileNotFoundError(f"Mental health FAQ file not found: {mental_health_file}")
            
            # Load mental health FAQ data
            df = pd.read_csv(mental_health_file)
            documents = []
            
            for _, row in df.iterrows():
                question = row['Questions']
                answer = row['Answers']
                # Create document in FAQ format
                content = f"QUESTION: {question}\nANSWER: {answer}"
                documents.append(Document(page_content=content))
            
            print(f"Loaded {len(documents)} mental health FAQ documents")
            for i, doc in enumerate(documents[:3]):
                print(f"Sample FAQ {i+1}: {doc.page_content[:200]}...")
                
        except Exception as e:
            print(f"Error loading mental health data: {e}")
            raise Exception("No valid data source available. Please ensure mental health FAQ data is present or provide Kaggle credentials.")

    print("Creating TF-IDF embeddings...")
    embeddings = ImprovedTFIDFEmbeddings()

    print("Creating ChromaDB vector store...")
    client = chromadb.PersistentClient(path="./src/chroma_db")
    
    # Clear existing collections to prevent mixing old and new data
    try:
        collections = client.list_collections()
        for collection in collections:
            print(f"Deleting existing collection: {collection.name}")
            client.delete_collection(collection.name)
    except Exception as e:
        print(f"Warning: Could not clear existing collections: {e}")
    
    print(f"Processing {len(documents)} documents...")
    
    # Check if this is a FAQ dataset and use smart retriever
    if any("QUESTION:" in doc.page_content for doc in documents):
        print("Using SmartFAQRetriever for better semantic matching...")
        return SmartFAQRetriever(documents, k=4)
    else:
        # Use vector store for non-FAQ datasets
        vectorstore = Chroma.from_documents(
            documents=documents,
            embedding=embeddings,
            client=client
    )
    print("Retriever setup complete.")
    return vectorstore.as_retriever(k=4)

def setup_rag_chain() -> Runnable:
    """Sets up the RAG chain with a prompt template and an LLM."""
    # Define the prompt template for the LLM
    prompt = PromptTemplate(
        template="""You are an assistant for question-answering tasks.
Use the following documents to answer the question.
If you don't know the answer, just say that you don't know.
Use three sentences maximum and keep the answer concise:
Question: {question}
Documents: {documents}
Answer:
""",
        input_variables=["question", "documents"],
    )

    # Initialize the LLM with dolphin-llama3:8b model
    # Note: This requires the Ollama server to be running with the specified model
    llm = ChatOllama(
        model="dolphin-llama3:8b",
        temperature=0,
    )

    # Create a chain combining the prompt template and LLM
    return prompt | llm | StrOutputParser()


# Define the RAG application class
class RAGApplication:
    def __init__(self, retriever: BaseRetriever, rag_chain: Runnable):
        self.retriever = retriever
        self.rag_chain = rag_chain

    def run(self, question: str) -> str:
        """Runs the RAG pipeline for a given question."""
        # Retrieve relevant documents
        documents = self.retriever.invoke(question)
        
        # Debug: Print retrieved documents
        print(f"\nDEBUG: Retrieved {len(documents)} documents for question: '{question}'")
        for i, doc in enumerate(documents):
            print(f"DEBUG: Document {i+1}: {doc.page_content[:200]}...")
        
        # Extract content from retrieved documents
        doc_texts = "\n\n".join([doc.page_content for doc in documents])
        
        # Debug: Print the combined document text
        print(f"DEBUG: Combined document text: {doc_texts[:300]}...")
        
        # Get the answer from the language model
        answer = self.rag_chain.invoke({"question": question, "documents": doc_texts})
        return answer

# Main execution block
if __name__ == "__main__":
    load_dotenv()

    # 1. Setup the components
    retriever = setup_retriever()
    rag_chain = setup_rag_chain()

    # 2. Initialize the RAG application
    rag_application = RAGApplication(retriever, rag_chain)

    # 3. Run an example query
    question = "What is prompt engineering"
    print("\n--- Running RAG Application ---")
    print(f"Question: {question}")
    answer = rag_application.run(question)
    print(f"Answer: {answer}")