reddit_rag_system.py

import pandas as pd
import numpy as np
from sentence_transformers import SentenceTransformer
from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility
import os
import json
from datetime import datetime
from tqdm import tqdm
import logging
import pyarrow.parquet as pq

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class RedditRAGSystem:
    def __init__(self, model_name="all-MiniLM-L6-v2", collection_name="reddit_comments"):
        """
        Initialize the Reddit RAG System
        
        Args:
            model_name: Sentence transformer model for embeddings
            collection_name: Name of the Milvus collection
        """
        self.model_name = model_name
        self.collection_name = collection_name
        self.model = None
        self.collection = None
        self.embedding_dim = 384  # Default for all-MiniLM-L6-v2
        
        # Connect to Milvus
        self._connect_to_milvus()
        
    def _connect_to_milvus(self):
        """Connect to Milvus database"""
        try:
            connections.connect(
                alias="default",
                host="localhost",
                port="19530"
            )
            logger.info("✅ Connected to Milvus successfully")
        except Exception as e:
            logger.error(f"❌ Failed to connect to Milvus: {e}")
            logger.info("💡 Make sure Milvus is running. You can start it with: docker run -d --name milvus_standalone -p 19530:19530 -p 9091:9091 milvusdb/milvus:latest standalone")
            raise
    
    def load_embedding_model(self):
        """Load the sentence transformer model"""
        logger.info(f"🔄 Loading embedding model: {self.model_name}")
        self.model = SentenceTransformer(self.model_name)
        self.embedding_dim = self.model.get_sentence_embedding_dimension()
        logger.info(f"✅ Model loaded. Embedding dimension: {self.embedding_dim}")
    
    def create_collection(self):
        """Create the Milvus collection for storing comment embeddings"""
        if utility.has_collection(self.collection_name):
            logger.info(f"📋 Collection '{self.collection_name}' already exists")
            self.collection = Collection(self.collection_name)
            return
        
        # Define schema
        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="comment_id", dtype=DataType.VARCHAR, max_length=50),
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.embedding_dim),
            FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
            FieldSchema(name="subreddit", dtype=DataType.VARCHAR, max_length=100),
            FieldSchema(name="username", dtype=DataType.VARCHAR, max_length=100),
            FieldSchema(name="score", dtype=DataType.INT64),
            FieldSchema(name="created_utc", dtype=DataType.INT64),
            FieldSchema(name="sentiment", dtype=DataType.VARCHAR, max_length=20),
        ]
        
        schema = CollectionSchema(fields, description="Reddit comments for RAG system")
        
        # Create collection
        self.collection = Collection(self.collection_name, schema)
        
        # Create index
        index_params = {
            "metric_type": "COSINE",
            "index_type": "IVF_FLAT",
            "params": {"nlist": 1024}
        }
        self.collection.create_index("embedding", index_params)
        logger.info(f"✅ Collection '{self.collection_name}' created with index")
    
    def process_comments_chunk(self, df_chunk, batch_id):
        """
        Process a chunk of comments and add to Milvus
        
        Args:
            df_chunk: DataFrame chunk of comments
            batch_id: Batch identifier
        """
        # Filter out empty or null comments
        df_chunk = df_chunk.dropna(subset=['body_clean'])
        df_chunk = df_chunk[df_chunk['body_clean'].str.strip() != '']
        
        if df_chunk.empty:
            return 0
        
        # Prepare texts for embedding
        texts = df_chunk['body_clean'].tolist()
        print(texts)
        
        # Generate embeddings
        embeddings = self.model.encode(texts, show_progress_bar=False, batch_size=32)
        print("Successfully generated embeddings.")
        
        # Prepare data for insertion
        data = []
        for idx, row in df_chunk.iterrows():
            # Determine dominant sentiment
            negative = row.get('negative', 0)
            neutral = row.get('neutral', 0)
            positive = row.get('positive', 0)
            
            if positive > negative and positive > neutral:
                sentiment = "positive"
            elif negative > positive and negative > neutral:
                sentiment = "negative"
            else:
                sentiment = "neutral"
            
            data.append([
                row.get('id', f"comment_{idx}_{batch_id}"),  # comment_id
                embeddings[idx].tolist(),  # embedding
                row.get('body_clean', '')[:65534],  # text (truncate if too long)
                row.get('subreddit', 'unknown')[:99],  # subreddit
                row.get('username', 'unknown')[:99],  # username
                row.get('score', 0),  # score
                row.get('created_utc', 0),  # created_utc
                sentiment,  # sentiment
            ])
        
        # Insert into collection
        #self.collection.insert(data)
        self.collection.insert(list(zip(*data)))
        print("Successfully inserted data into collection.")
        return len(data)
    
    def build_knowledge_base(self, parquet_file, batch_size=1000, max_samples=None):
        """
        Build the vector knowledge base from parquet file
        
        Args:
            parquet_file: Path to the parquet file
            batch_size: Number of records to process in each batch
            max_samples: Maximum number of samples to process (for testing)
        """
        logger.info(f"🏗️ Building knowledge base from {parquet_file}")
        
        # Load model if not loaded
        if self.model is None:
            self.load_embedding_model()
        
        # Create collection
        self.create_collection()
        
        # Load collection
        self.collection.load()
        
        # Process file in chunks
        total_processed = 0
        batch_id = 0
        
        try:
            # Read the entire file first (for small samples) or use pyarrow for chunking
            if max_samples and max_samples <= 10000:
                # For small samples, read all at once and then chunk
                logger.info(f"Reading {max_samples} samples from parquet file...")
                df = pd.read_parquet(parquet_file, engine="pyarrow")
                df = df.head(max_samples)
                
                # Process in batches
                for i in range(0, len(df), batch_size):
                    chunk = df.iloc[i:i+batch_size]
                    processed = self.process_comments_chunk(chunk, batch_id)
                    total_processed += processed
                    batch_id += 1
                    
                    # Flush every 10 batches
                    if batch_id % 10 == 0:
                        self.collection.flush()
                        logger.info(f"📊 Processed {total_processed:,} comments so far")
                        
            else:
                # For large files, use pyarrow to read in chunks
                parquet_file_obj = pq.ParquetFile(parquet_file)
                
                # Get total row groups
                num_row_groups = parquet_file_obj.num_row_groups
                logger.info(f"Processing {num_row_groups} row groups...")
                
                for row_group_idx in range(num_row_groups):
                    if max_samples and total_processed >= max_samples:
                        break
                    
                    # Read row group
                    chunk = parquet_file_obj.read_row_group(row_group_idx).to_pandas()
                    
                    # Process in smaller batches within the row group
                    for i in range(0, len(chunk), batch_size):
                        if max_samples and total_processed >= max_samples:
                            break
                            
                        sub_chunk = chunk.iloc[i:i+batch_size]
                        processed = self.process_comments_chunk(sub_chunk, batch_id)
                        total_processed += processed
                        batch_id += 1
                        
                        # Flush every 10 batches
                        if batch_id % 10 == 0:
                            self.collection.flush()
                            logger.info(f"📊 Processed {total_processed:,} comments so far")
            
            # Final flush
            self.collection.flush()
            
        except Exception as e:
            logger.error(f"❌ Error processing file: {e}")
            raise
        
        logger.info(f"✅ Knowledge base built! Total comments: {total_processed:,}")
        return total_processed
    
    def search_similar_comments(self, query, n_results=5, subreddit_filter=None):
        """
        Search for similar comments
        
        Args:
            query: Search query
            n_results: Number of results to return
            subreddit_filter: Optional subreddit filter
        """
        if not self.collection:
            logger.error("❌ No collection loaded")
            return []
        
        # Generate query embedding
        query_embedding = self.model.encode([query])
        
        # Prepare search parameters
        search_params = {"metric_type": "COSINE", "params": {"nprobe": 10}}
        
        # Build expression for filtering
        expr = None
        if subreddit_filter:
            expr = f'subreddit == "{subreddit_filter}"'
        
        # Search
        results = self.collection.search(
            data=query_embedding.tolist(),
            anns_field="embedding",
            param=search_params,
            limit=n_results,
            expr=expr,
            output_fields=["comment_id", "text", "subreddit", "username", "score", "created_utc", "sentiment"]
        )
        
        return results
    
    def get_collection_stats(self):
        """Get statistics about the collection"""
        if not self.collection:
            logger.error("❌ No collection loaded")
            return None
        
        stats = {
            "total_entities": self.collection.num_entities,
            "collection_name": self.collection_name
        }
        
        logger.info(f"📊 Collection stats: {stats}")
        return stats
    
    def format_search_results(self, results):
        """Format search results for display"""
        if not results or not results[0]:
            return []
        
        formatted_results = []
        for hit in results[0]:
            result = {
                "comment_id": hit.entity.get("comment_id"),
                "text": hit.entity.get("text"),
                "subreddit": hit.entity.get("subreddit"),
                "username": hit.entity.get("username"),
                "score": hit.entity.get("score"),
                "created_utc": hit.entity.get("created_utc"),
                "sentiment": hit.entity.get("sentiment"),
                "similarity_score": hit.score
            }
            formatted_results.append(result)
        
        return formatted_results

def main():
    """Main function to demonstrate the RAG system"""
    print("🚀 REDDIT RAG SYSTEM")
    print("=" * 50)
    
    # Initialize the system
    rag_system = RedditRAGSystem()
    
    # Ask user what they want to do
    print("\nWhat would you like to do?")
    print("1. Build knowledge base (process comments)")
    print("2. Search similar comments")
    print("3. Get collection stats")
    
    choice = input("\nEnter your choice (1-3): ")
    
    if choice == "1":
        # Build knowledge base
        print("\n🏗️ BUILDING KNOWLEDGE BASE")
        print("-" * 30)
        
        # Ask for sample size for testing
        sample_size = input("Enter number of samples to process (or press Enter for all): ")
        max_samples = int(sample_size) if sample_size.strip() else None
        
        if max_samples:
            print(f"🧪 Processing {max_samples:,} samples for testing...")
        
        try:
            total_processed = rag_system.build_knowledge_base(
                "all_comments_merged.parquet", 
                batch_size=1000, 
                max_samples=max_samples
            )
            print(f"✅ Successfully processed {total_processed:,} comments")
        except Exception as e:
            print(f"❌ Error: {e}")
    
    elif choice == "2":
        # Search functionality
        print("\n🔍 SEARCH SIMILAR COMMENTS")
        print("-" * 30)
        
        # Load model if needed
        if rag_system.model is None:
            rag_system.load_embedding_model()
        
        # Load collection
        if utility.has_collection(rag_system.collection_name):
            rag_system.collection = Collection(rag_system.collection_name)
            rag_system.collection.load()
        else:
            print("❌ No collection found. Please build the knowledge base first.")
            return
        
        while True:
            query = input("\nEnter your search query (or 'quit' to exit): ")
            if query.lower() == 'quit':
                break
            
            subreddit_filter = input("Filter by subreddit (optional, press Enter to skip): ")
            subreddit_filter = subreddit_filter.strip() if subreddit_filter.strip() else None
            
            try:
                results = rag_system.search_similar_comments(query, n_results=5, subreddit_filter=subreddit_filter)
                formatted_results = rag_system.format_search_results(results)
                
                print(f"\n📋 Found {len(formatted_results)} similar comments:")
                for i, result in enumerate(formatted_results, 1):
                    print(f"\n{i}. Similarity: {result['similarity_score']:.3f}")
                    print(f"   Comment: {result['text'][:200]}...")
                    print(f"   Subreddit: r/{result['subreddit']}")
                    print(f"   User: u/{result['username']}")
                    print(f"   Score: {result['score']}")
                    print(f"   Sentiment: {result['sentiment']}")
                    print(f"   Comment ID: {result['comment_id']}")
                    
            except Exception as e:
                print(f"❌ Search error: {e}")
    
    elif choice == "3":
        # Get stats
        print("\n📊 COLLECTION STATISTICS")
        print("-" * 30)
        
        if utility.has_collection(rag_system.collection_name):
            rag_system.collection = Collection(rag_system.collection_name)
            rag_system.get_collection_stats()
        else:
            print("❌ No collection found")
    
    else:
        print("❌ Invalid choice")

if __name__ == "__main__":
    main()