""" RAG Data Preparation Service for OpenTriage. Uses Spark for high-speed chunking and cleaning of documents for Vector DB. Falls back to pure Python when Spark is not available (e.g., HuggingFace Spaces). """ import logging import re from typing import List, Dict, Any, Optional from datetime import datetime, timezone from pydantic import BaseModel # Try to import PySpark - may not be available on HuggingFace Spaces SPARK_AVAILABLE = False try: from pyspark.sql import SparkSession, DataFrame from pyspark.sql.functions import ( col, udf, explode, posexplode, lit, concat, concat_ws, length, size, array, struct, row_number, monotonically_increasing_id, regexp_replace, trim, lower, split ) from pyspark.sql.types import ( StructType, StructField, StringType, IntegerType, ArrayType, MapType ) from spark_manager import get_or_create_spark_session SPARK_AVAILABLE = True except ImportError: pass # PySpark not installed - use pure Python fallback logger = logging.getLogger(__name__) class DocumentChunk(BaseModel): """Model for a document chunk ready for embedding.""" chunk_id: str document_id: str document_type: str # issue, pr, doc, comment source_repo: str chunk_index: int total_chunks: int content: str metadata: Dict[str, Any] = {} token_count: int = 0 class RAGDataPrep: """ Service for preparing documents for Vector Database ingestion. Uses Spark for parallel processing of: - Text cleaning (remove markdown artifacts, code blocks) - Chunking with overlap - Metadata enrichment """ def __init__( self, chunk_size: int = 512, chunk_overlap: int = 64, min_chunk_size: int = 50 ): """ Initialize RAG data prep service. Args: chunk_size: Target tokens per chunk chunk_overlap: Token overlap between chunks min_chunk_size: Minimum tokens for a valid chunk """ self.chunk_size = chunk_size self.chunk_overlap = chunk_overlap self.min_chunk_size = min_chunk_size def _clean_text(self, text: str) -> str: """ Clean text but PRESERVE code blocks and structure. """ if not text: return "" # Remove markdown images but keep alt text if useful? # For now, just remove images to reduce noise text = re.sub(r'!\[([^\]]*)\]\([^\)]+\)', r'[IMAGE: \1]', text) # Remove HTML tags (basic) text = re.sub(r'<[^>]+>', ' ', text) # We KEEP code blocks now because they are crucial for setup instructions # Just ensure newlines around them are clean return text.strip() def _recursive_split(self, text: str, separators: List[str], chunk_size: int, chunk_overlap: int) -> List[str]: """ Recursively split text by separators. """ final_chunks = [] if len(text) <= chunk_size: return [text] if not separators: # If no separators left, hard split by size return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size-chunk_overlap)] separator = separators[0] next_separators = separators[1:] # Split by current separator splits = text.split(separator) current_chunk = [] current_length = 0 for split_part in splits: part_len = len(split_part) if current_length + part_len + len(separator) > chunk_size: # Flush current chunk if current_chunk: joined_chunk = separator.join(current_chunk) if len(joined_chunk) > chunk_size: # Recursively split this too-large chunk final_chunks.extend(self._recursive_split(joined_chunk, next_separators, chunk_size, chunk_overlap)) else: final_chunks.append(joined_chunk) # Start new chunk with overlaps if needed (simplified here for Python) # For strict LangChain parity we'd need more complex overlap logic. # We will just start a new chunk. current_chunk = [] current_length = 0 current_chunk.append(split_part) current_length += part_len + (len(separator) if current_length > 0 else 0) if current_chunk: final_chunks.append(separator.join(current_chunk)) return final_chunks def _chunk_text(self, text: str) -> List[Dict[str, Any]]: """ Chunk text using proper recursive character splitting to preserve context. """ if not text: return [] # 1 char approx 1 token? No, usually 4 chars ~ 1 token. # But here we work with characters as proxy. # Adjusted char_limit based on chunk_size (tokens) char_limit = self.chunk_size * 4 char_overlap = self.chunk_overlap * 4 separators = ["\n\n", "\n", ". ", " ", ""] chunks_text = self._recursive_split(text, separators, char_limit, char_overlap) chunks = [] for i, chunk_content in enumerate(chunks_text): if len(chunk_content.strip()) < self.min_chunk_size: continue chunks.append({ "content": chunk_content.strip(), "token_count": len(chunk_content) // 4, # Approx "chunk_index": i }) return chunks def create_chunking_udf(self): """ Create a UDF for text chunking. """ chunk_size = self.chunk_size chunk_overlap = self.chunk_overlap min_chunk_size = self.min_chunk_size def chunk_text_udf(text: str) -> List[Dict[str, Any]]: if not text: return [] # Simple tokenization tokens = re.findall(r'\b\w+\b', text.lower()) if len(tokens) < min_chunk_size: return [{"content": text, "token_count": len(tokens), "chunk_index": 0}] chunks = [] step = chunk_size - chunk_overlap for i in range(0, len(tokens), step): chunk_tokens = tokens[i:i + chunk_size] if len(chunk_tokens) < min_chunk_size and chunks: continue chunks.append({ "content": " ".join(chunk_tokens), "token_count": len(chunk_tokens), "chunk_index": len(chunks) }) return chunks return udf(chunk_text_udf, ArrayType( StructType([ StructField("content", StringType()), StructField("token_count", IntegerType()), StructField("chunk_index", IntegerType()) ]) )) def create_cleaning_udf(self): """ Create a UDF for text cleaning. """ def clean_text_udf(text: str) -> str: if not text: return "" # Remove code blocks text = re.sub(r'```[\s\S]*?```', ' [CODE_BLOCK] ', text) text = re.sub(r'`[^`]+`', ' [CODE] ', text) # Remove markdown links text = re.sub(r'\[([^\]]+)\]\([^\)]+\)', r'\1', text) # Remove images text = re.sub(r'!\[[^\]]*\]\([^\)]+\)', ' [IMAGE] ', text) # Remove HTML text = re.sub(r'<[^>]+>', ' ', text) # Clean whitespace text = re.sub(r'\s+', ' ', text) return text.strip() return udf(clean_text_udf, StringType()) async def fetch_documents( self, doc_types: List[str] = None, repo_names: List[str] = None, github_access_token: Optional[str] = None ) -> List[Dict[str, Any]]: """ Fetch documents from MongoDB for RAG preparation. Args: doc_types: Types to fetch (issue, pr, comment, readme, contributing) repo_names: Optional filter by repository Returns: List of document records """ from config.database import db doc_types = doc_types or ["issue", "pr", "comment", "readme", "contributing"] documents = [] # Fetch README and CONTRIBUTING files with high priority if repo_names: from services.github_service import github_service for repo in repo_names: # Fetch README if "readme" in doc_types: try: content = await github_service.fetch_repository_readme(repo, github_access_token) if content: # Extract key sections for high priority tagging priority = self._detect_priority_sections(content, "readme") documents.append({ "document_id": f"{repo}_readme", "document_type": "readme", "source_repo": repo, "title": "Project README", "body": content, "author": "System", "number": 0, "state": "active", "priority": priority, "created_at": datetime.now(timezone.utc).isoformat() }) except Exception as e: logger.error(f"Failed to fetch README for {repo}: {e}") # Fetch CONTRIBUTING.md (high priority for contributor context) if "contributing" in doc_types: try: content = await github_service.fetch_contributing_file(repo, github_access_token) if content: documents.append({ "document_id": f"{repo}_contributing", "document_type": "contributing", "source_repo": repo, "title": "Contributor Guidelines", "body": content, "author": "System", "number": 0, "state": "active", "priority": "high", # Always high priority "created_at": datetime.now(timezone.utc).isoformat() }) except Exception as e: logger.error(f"Failed to fetch CONTRIBUTING for {repo}: {e}") if "issue" in doc_types or "pr" in doc_types: query = {} if repo_names: query["repoName"] = {"$in": repo_names} cursor = db.issues.find(query, {"_id": 0}) items = await cursor.to_list(length=None) for item in items: doc_type = "pr" if item.get("isPR") else "issue" if doc_type not in doc_types: continue documents.append({ "document_id": item.get("id", ""), "document_type": doc_type, "source_repo": item.get("repoName", ""), "title": item.get("title", ""), "body": item.get("body", ""), "author": item.get("authorName", ""), "number": item.get("number", 0), "state": item.get("state", ""), "created_at": item.get("createdAt", "") }) if "comment" in doc_types: query = {} if repo_names: query["repoName"] = {"$in": repo_names} cursor = db.comments.find(query, {"_id": 0}) comments = await cursor.to_list(length=None) for comment in comments: documents.append({ "document_id": comment.get("id", ""), "document_type": "comment", "source_repo": comment.get("repoName", ""), "title": "", "body": comment.get("body", ""), "author": comment.get("author", ""), "number": 0, "state": "", "created_at": comment.get("createdAt", "") }) return documents def prepare_documents(self, documents: List[Dict[str, Any]]) -> Any: """ Prepare documents for Vector DB. Uses Spark if available, otherwise falls back to pure Python. Args: documents: List of document records Returns: DataFrame (Spark) or List[Dict] (Python fallback) with chunked documents """ if not documents: return None # Use pure Python if Spark is not available if not SPARK_AVAILABLE: return self._prepare_documents_python(documents) # Spark path spark = get_or_create_spark_session() # Create DataFrame df = spark.createDataFrame(documents) # Combine title and body df = df.withColumn( "full_text", concat_ws(" ", col("title"), col("body")) ) # Clean text using UDF cleaning_udf = self.create_cleaning_udf() df = df.withColumn("cleaned_text", cleaning_udf(col("full_text"))) # Chunk text using UDF chunking_udf = self.create_chunking_udf() df = df.withColumn("chunks", chunking_udf(col("cleaned_text"))) # Explode chunks into separate rows df = df.withColumn("chunk_data", explode(col("chunks"))) # Extract chunk fields df = df.withColumn("chunk_content", col("chunk_data.content")) df = df.withColumn("token_count", col("chunk_data.token_count")) df = df.withColumn("chunk_index", col("chunk_data.chunk_index")) # Calculate total chunks per document from pyspark.sql import Window window_spec = Window.partitionBy("document_id") df = df.withColumn("total_chunks", size(col("chunks"))) # Generate chunk IDs df = df.withColumn( "chunk_id", concat(col("document_id"), lit("_chunk_"), col("chunk_index")) ) # Select final columns result_df = df.select( "chunk_id", "document_id", "document_type", "source_repo", "chunk_index", "total_chunks", "chunk_content", "token_count", "author", "number", "state", "created_at" ) return result_df def _prepare_documents_python(self, documents: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """ Pure Python fallback for document preparation. Used when Spark is not available (e.g., HuggingFace Spaces). """ all_chunks = [] for doc in documents: # Combine title and body full_text = f"{doc.get('title', '')} {doc.get('body', '')}".strip() # Clean text cleaned_text = self._clean_text(full_text) # Chunk text chunks = self._chunk_text(cleaned_text) total_chunks = len(chunks) for chunk in chunks: all_chunks.append({ "chunk_id": f"{doc.get('document_id', '')}_chunk_{chunk['chunk_index']}", "document_id": doc.get("document_id", ""), "document_type": doc.get("document_type", ""), "source_repo": doc.get("source_repo", ""), "chunk_index": chunk["chunk_index"], "total_chunks": total_chunks, "chunk_content": chunk["content"], "token_count": chunk["token_count"], "author": doc.get("author", ""), "number": doc.get("number", 0), "state": doc.get("state", ""), "created_at": doc.get("created_at", "") }) return all_chunks async def prepare_and_store( self, doc_types: List[str] = None, repo_names: List[str] = None, collection_name: str = "rag_chunks", github_access_token: Optional[str] = None ) -> Dict[str, Any]: """ Prepare documents and store chunks in MongoDB. Args: doc_types: Document types to process repo_names: Optional repository filter collection_name: Target collection for chunks Returns: Preparation summary """ from config.database import db start_time = datetime.now(timezone.utc) # Fetch documents documents = await self.fetch_documents(doc_types, repo_names, github_access_token) if not documents: return { "status": "no_documents", "documents_processed": 0, "chunks_created": 0 } logger.info(f"Preparing {len(documents)} documents for RAG") # Prepare documents (returns DataFrame if Spark, List if Python fallback) result = self.prepare_documents(documents) if result is None: return { "status": "preparation_failed", "documents_processed": len(documents), "chunks_created": 0 } # Handle both Spark DataFrame and Python list if SPARK_AVAILABLE and hasattr(result, 'cache'): # Spark path result.cache() chunk_count = result.count() chunks = result.collect() # Store chunks (Spark Row objects) chunk_docs = [] for row in chunks: chunk_docs.append({ "chunkId": row.chunk_id, "documentId": row.document_id, "documentType": row.document_type, "sourceRepo": row.source_repo, "chunkIndex": row.chunk_index, "totalChunks": row.total_chunks, "content": row.chunk_content, "tokenCount": row.token_count, "metadata": { "author": row.author, "number": row.number, "state": row.state, "createdAt": row.created_at }, "preparedAt": datetime.now(timezone.utc).isoformat() }) result.unpersist() else: # Python fallback path chunk_count = len(result) # Store chunks (Python dicts) chunk_docs = [] for chunk in result: chunk_docs.append({ "chunkId": chunk["chunk_id"], "documentId": chunk["document_id"], "documentType": chunk["document_type"], "sourceRepo": chunk["source_repo"], "chunkIndex": chunk["chunk_index"], "totalChunks": chunk["total_chunks"], "content": chunk["chunk_content"], "tokenCount": chunk["token_count"], "metadata": { "author": chunk["author"], "number": chunk["number"], "state": chunk["state"], "createdAt": chunk["created_at"] }, "preparedAt": datetime.now(timezone.utc).isoformat() }) # Clear existing chunks for these repos if repo_names: await db[collection_name].delete_many({"sourceRepo": {"$in": repo_names}}) else: # Fallback if no repo names provided (shouldn't happen in single-repo index) await db[collection_name].delete_many({}) if chunk_docs: await db[collection_name].insert_many(chunk_docs) end_time = datetime.now(timezone.utc) duration = (end_time - start_time).total_seconds() logger.info(f"RAG preparation complete: {chunk_count} chunks from {len(documents)} documents") return { "status": "completed", "documents_processed": len(documents), "chunks_created": chunk_count, "collection": collection_name, "duration_seconds": duration, "avg_chunks_per_doc": chunk_count / len(documents) if documents else 0 } async def get_chunks_for_embedding( self, batch_size: int = 100, skip_embedded: bool = True ) -> List[Dict[str, Any]]: """ Get chunks ready for embedding. Args: batch_size: Number of chunks to return skip_embedded: Skip chunks that already have embeddings Returns: List of chunks for embedding """ from config.database import db query = {} if skip_embedded: query["embedding"] = {"$exists": False} cursor = db.rag_chunks.find(query, {"_id": 0}).limit(batch_size) chunks = await cursor.to_list(length=batch_size) return chunks async def store_embeddings( self, chunk_id: str, embedding: List[float] ) -> bool: """ Store embedding for a chunk. Args: chunk_id: Chunk identifier embedding: Vector embedding Returns: Success status """ from config.database import db result = await db.rag_chunks.update_one( {"chunkId": chunk_id}, { "$set": { "embedding": embedding, "embeddedAt": datetime.now(timezone.utc).isoformat() } } ) return result.modified_count > 0 def get_stats(self, df: Any) -> Dict[str, Any]: """ Get statistics about prepared chunks. Only works when SPARK_AVAILABLE is True. """ if df is None or not SPARK_AVAILABLE: return {} return { "total_chunks": df.count(), "by_type": { row.document_type: row["count"] for row in df.groupBy("document_type").count().collect() }, "avg_token_count": df.agg({"token_count": "avg"}).first()[0], "total_documents": df.select("document_id").distinct().count() } # Singleton instance rag_data_prep = RAGDataPrep( chunk_size=512, chunk_overlap=64, min_chunk_size=50 )