rag.py

import os
import time
from typing import List, Tuple, Optional
# from markitdown import MarkItDown
from rank_bm25 import BM25Okapi
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import json
import re
# import nltk
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
from dataclasses import dataclass
from difflib import SequenceMatcher
# nltk.download('punkt')
# nltk.download('stopwords')

@dataclass
class Section:
    content: str
    context: str
    heading: str = ""
    source_file: str = ""

def similar(a: str, b: str, threshold: float = 0.8) -> bool:
    return SequenceMatcher(None, a, b).ratio() > threshold

class Retriever:
    def __init__(self, docs_dir: str, db_dir: str):
        self.docs_dir = docs_dir
        self.db_dir = db_dir
        self.documents = []
        self.tokenized_corpus = []
        self.bm25 = None
        self.vectorizer = TfidfVectorizer(
            min_df=1,
            stop_words=None,
            lowercase=True,
            token_pattern=r'(?u)\b\w+\b'
        )
        self.doc_vectors = None
        self.sections = []
        self.context_window = 3  # Number of sentences for context
        os.makedirs(db_dir, exist_ok=True)

    def preprocess_text(self, text: str) -> str:
        """Clean and tokenize text while preserving important terms"""
        # Remove all utf8 incompatible characters
        text = re.sub(r'[^\x00-\x7F]+', '', text)
        text = text.lower()
        # Remove special characters
        text = ' '.join(word for word in text.split() if any(c.isalnum() for c in word))
        return text

    def tokenize_text(self, text: str) -> List[str]:
        """Tokenize text into words"""
        tokens = word_tokenize(text)
        # Remove very short tokens and number tokens
        tokens = [token for token in tokens if len(token) > 1 and not token.isnumeric()]
        return tokens

    def extract_context(self, text: str, query: str, window_size: int = 100) -> str:
        """Extract context around query matches including subheading points"""
        lines = text.split('\n')

        # Check if any line matches the subheading format and query
        for i, line in enumerate(lines):
            line = line.strip()
            # Skip lines starting with '> '
            if line.startswith('> '):
                continue
            if line.startswith('- ') and line.endswith(':') and query.lower() in line.lower():
                # Found matching subheading, collect all bullet points
                result = [line]
                j = i + 1
                while j < len(lines):
                    current_line = lines[j].strip()

                    # Check for bullet points, numbers, roman numerals, or indented content
                    is_bullet = current_line.startswith('* ') or current_line.startswith('- ')
                    is_number = bool(re.match(r'^\d+\.', current_line))
                    is_roman = bool(re.match(r'^[ivxIVX]+\)', current_line))
                    is_paren = bool(re.match(r'^\([^)]+\)', current_line))
                    is_indented = lines[j].startswith('    ') or lines[j].startswith('  ') or lines[j].startswith('\t')

                    # If line is empty or doesn't match any continuation pattern, break
                    if not current_line or (not any([is_bullet, is_number, is_roman, is_paren, is_indented]) 
                                           and not current_line.startswith('(')):
                        break

                    # Only append if not a quote line
                    if not current_line.startswith('> '):
                        result.append(current_line)
                    j += 1
                return '\n'.join(result)

        # If no subheading match, use regular word-based context extraction
        words = text.split()
        query_words = query.lower().split()

        for i, word in enumerate(words):
            if any(qw in word.lower() for qw in query_words):
                start = max(0, i - window_size)
                end = min(len(words), i + window_size + 1)
                return ' '.join(words[start:end])

        return text

    def preprocess_documents(self, file_filter: Optional[List[str]] = None):
        """Process documents and store in database
        Args:
            file_filter: Optional list of filenames to process. If None, process all files.
        """
        t0 = time.perf_counter()
        current_heading = ""

        for filename in os.listdir(self.docs_dir):
            if filename.endswith(('.txt', '.md')):
                if file_filter and filename not in file_filter:
                    continue

                filepath = os.path.join(self.docs_dir, filename)
                with open(filepath, 'r', encoding='utf-8') as file:
                    content = file.read()
                    # # Handle markdown formatting
                    # content = re.sub(r'^#+ ', '> ', content, flags=re.MULTILINE) # Add > to headings
                    # content = re.sub(r'^\s*(\d+\. )', '- ', content, flags=re.MULTILINE) # Convert numbered lists
                    # content = re.sub(r'^\s*[-*] ', '- ', content, flags=re.MULTILINE) # Standardize bullet points
                    # content = re.sub(r'^\s*(\w+:)', '- \1', content, flags=re.MULTILINE) # Format subheadings

                    # Split into sections based on headers
                    sections = content.split('\n\n')

                    for section in sections:
                        if section.strip().startswith('**'):
                            current_heading = section.strip('*').strip()
                            continue

                        if section.strip():
                            # Filter out lines starting with '>'
                            filtered_lines = [line for line in section.split('\n') 
                                           if not line.strip().startswith('>')]
                            filtered_section = '\n'.join(filtered_lines)
                            
                            if filtered_section.strip():  # Only process if content remains
                                cleaned_text = self.preprocess_text(section)
                                self.sections.append(Section(
                                    content=cleaned_text,
                                    context=section.strip(),
                                    heading=current_heading,
                                    source_file=filename
                                ))
                                self.documents.append(cleaned_text)

        self.tokenized_corpus = [self.tokenize_text(doc) for doc in self.documents]
        try:
            self.bm25 = BM25Okapi(self.tokenized_corpus)
        except ZeroDivisionError:
            print("Error: No documents found for BM25. Please check your documents directory.")
            return

        self.doc_vectors = self.vectorizer.fit_transform(self.documents)

        # Save processed documents
        with open(os.path.join(self.db_dir, 'processed_docs.json'), 'w') as f:
            json.dump({
                'documents': self.documents,
                'sections': [(s.content, s.context, s.heading, s.source_file) for s in self.sections]
            }, f)
        print(f"Preprocessed {len([file for file in os.listdir(self.docs_dir) if file.endswith(('.txt', '.md'))])} documents into {len(self.documents)} vectors in {time.perf_counter() - t0:.4f} seconds.")

    def load_database(self):
        """Load preprocessed documents"""
        t0 = time.perf_counter()
        with open(os.path.join(self.db_dir, 'processed_docs.json'), 'r') as f:
            data = json.load(f)
            self.documents = data['documents']
            self.sections = [Section(content=s[0], context=s[1], heading=s[2], source_file=s[3]) 
                           for s in data['sections']]

        self.tokenized_corpus = [self.tokenize_text(doc) for doc in self.documents]
        self.bm25 = BM25Okapi(self.tokenized_corpus)
        self.doc_vectors = self.vectorizer.fit_transform(self.documents)
        print(f"Loaded {len(self.documents)} vectors in {time.perf_counter() - t0:.4f} seconds.")

    def retrieve(self, query: str, top_k: int = 6, file_filter: Optional[List[str]] = None) -> List[Tuple[str, float, str, str]]:
        """
        Retrieve documents using BM25 and Vector similarity

        Args:
            query: Search query
            top_k: Number of results to return
            file_filter: List of filenames to search through (optional)
        Returns:
            List of tuples (context, score, heading, source_file)
        """
        t0 = time.perf_counter()
        # Get BM25 scores
        tokenized_query = self.tokenize_text(query)
        bm25_scores = self.bm25.get_scores(tokenized_query)

        # Get vector similarity scores
        query_vec = self.vectorizer.transform([query])
        vector_scores = cosine_similarity(query_vec, self.doc_vectors).flatten()

        # Combine scores (simple average)
        combined_scores = (bm25_scores + vector_scores) / 2

        # Get top-k documents with deduplication
        results = []
        seen_content = set()

        for idx in combined_scores.argsort()[::-1]:
            score = combined_scores[idx]
            if score == 0:
                continue

            section = self.sections[idx]

            # Skip if not in file_filter
            if file_filter and section.source_file not in file_filter:
                continue

            context = self.extract_context(section.context, query)

            # Skip if too similar to existing results
            if any(similar(context, seen) for seen in seen_content):
                continue

            seen_content.add(context)
            results.append((context, score, section.heading, section.source_file))

            if len(results) >= top_k:
                break

        print(f"Retrieved {len(results)} results in {time.perf_counter() - t0:.4f} seconds.")
        return results

def main():
    retriever = Retriever('./rag/documents', './rag/database')

    # Check if database exists, create if not
    if not os.path.exists('./rag/database/processed_docs.json'):
        print("Processing documents...")
        retriever.preprocess_documents()
    else:
        print("Loading existing database...")
        retriever.load_database()

    while True:
        query = input("Enter your question (or '/quit' to exit): ")
        if query.lower() == '/quit':
            print("Exiting...")
            break

        # file filtering
        file_filter = input("Enter specific files to search (comma separated) or press Enter for all: ")
        file_filter = [f.strip() for f in file_filter.split(',')] if file_filter else None

        results = retriever.retrieve(query, file_filter=file_filter)

        if not results:
            print("No relevant documents found.")
            continue

        print("\nRelevant information:")
        for context, score, heading, source_file in results:
            if heading:
                print(f"\nSection: {heading}")
            print(f"Source: {source_file}")
            print(f"Score: {score:.3f}")
            print(f"Content: {context}\n")
            print("-" * 80)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nExiting...")