""" Vector RAG Manager for Too Many Cables Implements true RAG with embeddings and vector database using ChromaDB """ import os import json import uuid import hashlib import logging from pathlib import Path from datetime import datetime from typing import Dict, List, Optional, Tuple, Any # Disable all external telemetry BEFORE importing chromadb to avoid PostHog issues os.environ.setdefault('ANONYMIZED_TELEMETRY', 'False') os.environ.setdefault('CHROMA_TELEMETRY_ENABLED', 'false') os.environ.setdefault('CHROMADB_TELEMETRY_DISABLED', 'true') os.environ.setdefault('POSTHOG_DISABLED', 'true') # Set HuggingFace cache directories BEFORE importing transformers/sentence-transformers os.environ['TRANSFORMERS_CACHE'] = '/app/data/transformers_cache' os.environ['HF_HOME'] = '/app/data/huggingface_cache' os.environ['HF_DATASETS_CACHE'] = '/app/data/huggingface_cache' os.environ['SENTENCE_TRANSFORMERS_HOME'] = '/app/data/sentence_transformers_cache' import chromadb from chromadb.config import Settings from sentence_transformers import SentenceTransformer import numpy as np from .knowledge_base_manager import KnowledgeBaseManager class VectorRAGManager: def __init__(self, knowledge_base_path: str = "knowledge_base", vector_db_path: str = "vector_db", embedding_model: str = "all-MiniLM-L6-v2"): """Initialize Vector RAG Manager with ChromaDB and embeddings""" self.kb_path = Path(knowledge_base_path) self.vector_db_path = Path(vector_db_path) self.embedding_model_name = embedding_model # Create directories self.kb_path.mkdir(exist_ok=True) self.vector_db_path.mkdir(exist_ok=True) # Initialize knowledge base manager self.kb_manager = KnowledgeBaseManager(knowledge_base_path) # Initialize ChromaDB client (explicitly disable telemetry) self.chroma_client = chromadb.PersistentClient( path=str(self.vector_db_path), settings=Settings( anonymized_telemetry=False, allow_reset=True ) ) # Initialize embedding model logging.info(f"Loading embedding model: {embedding_model}") self.embedding_model = SentenceTransformer(embedding_model) logging.info(f"Embedding model loaded. Dimension: {self.embedding_model.get_sentence_embedding_dimension()}") # Collection name for document chunks self.collection_name = "tmc_documents" # Document chunking parameters (token-based for better quality) # Reduced from 300 to 200 tokens for more focused, relevant chunks self.chunk_size_tokens = 200 # target tokens per chunk (smaller = more focused) self.chunk_overlap_tokens = 30 # ~15% overlap (maintains context) # Approximate token conversion (rough estimate: 1 token ≈ 4 characters) self.chunk_size = self.chunk_size_tokens * 4 # ~800 chars self.chunk_overlap = self.chunk_overlap_tokens * 4 # ~120 chars # Initialize or get collection self._init_collection() def _init_collection(self): """Initialize or get the ChromaDB collection""" try: self.collection = self.chroma_client.get_collection( name=self.collection_name ) logging.info(f"Loaded existing collection '{self.collection_name}' with {self.collection.count()} documents") except Exception: # Collection doesn't exist, create it (handle older/newer API differences) try: self.collection = self.chroma_client.create_collection( name=self.collection_name, metadata={"description": "Too Many Cables knowledge base embeddings"} ) except Exception: # Fallback without metadata self.collection = self.chroma_client.create_collection( name=self.collection_name ) logging.info(f"Created new collection '{self.collection_name}'") def estimate_tokens(self, text: str) -> int: """Rough token estimation: ~1 token per 4 characters for English text""" return max(1, len(text) // 4) def normalize_content(self, content: str) -> str: """ Clean and normalize content before chunking Removes boilerplate, headers/footers, nav menus, TOCs, contact blocks """ import re # Remove common boilerplate patterns boilerplate_patterns = [ r'Copyright \d{4}.*?All rights reserved\.?', r'© \d{4}.*?(?:\n|$)', r'Contact us:.*?(?:\n\n|$)', r'For more information.*?(?:\n\n|$)', r'Visit our website.*?(?:\n\n|$)', r'Call us at.*?(?:\n\n|$)', r'Email:.*?(?:\n\n|$)', r'Phone:.*?(?:\n\n|$)', r'Address:.*?(?:\n\n|$)', r'Table of Contents.*?(?:\n\n|\n(?=[A-Z]))', r'TOC:.*?(?:\n\n|\n(?=[A-Z]))', r'Navigation:.*?(?:\n\n|$)', r'Home \| Products \| Support.*?(?:\n|$)', r'Back to top.*?(?:\n|$)', r'Print this page.*?(?:\n|$)', r'Share:.*?(?:\n|$)', r'Related articles:.*?(?:\n\n|$)', r'See also:.*?(?:\n\n|$)', r'Last updated:.*?(?:\n|$)', r'Page \d+ of \d+.*?(?:\n|$)', ] cleaned_content = content for pattern in boilerplate_patterns: cleaned_content = re.sub(pattern, '', cleaned_content, flags=re.IGNORECASE | re.DOTALL) # Remove excessive whitespace and normalize line breaks cleaned_content = re.sub(r'\n\s*\n\s*\n+', '\n\n', cleaned_content) # Max 2 consecutive newlines cleaned_content = re.sub(r'[ \t]+', ' ', cleaned_content) # Normalize spaces cleaned_content = re.sub(r'\n ', '\n', cleaned_content) # Remove leading spaces on lines # Remove repetitive elements (like repeated contact info) lines = cleaned_content.split('\n') deduped_lines = [] seen_lines = set() for line in lines: line_clean = line.strip().lower() # Skip very short lines or highly repetitive content if len(line_clean) < 10: deduped_lines.append(line) continue # Check for repetitive patterns if line_clean not in seen_lines: deduped_lines.append(line) seen_lines.add(line_clean) elif len(seen_lines) < 50: # Allow some repetition in small docs deduped_lines.append(line) return '\n'.join(deduped_lines).strip() def chunk_text(self, text: str, chunk_size_tokens: Optional[int] = None, chunk_overlap_tokens: Optional[int] = None) -> List[str]: """ Split text into overlapping chunks optimized for RAG retrieval Uses token-based sizing (200-400 tokens) with 10-15% overlap Enhanced with product manual structure awareness """ target_tokens = chunk_size_tokens or self.chunk_size_tokens overlap_tokens = chunk_overlap_tokens or self.chunk_overlap_tokens # Convert to character estimates target_chars = target_tokens * 4 overlap_chars = overlap_tokens * 4 if self.estimate_tokens(text) <= target_tokens: return [text] chunks = [] start = 0 while start < len(text): end = start + target_chars # Try to end at optimal boundaries (product sections > sentences > paragraphs > words) if end < len(text): # Priority 1: Product section boundaries (for product manuals) product_boundaries = [ '\n### TMC-', # Product headers like ### TMC-HDMI-8K-10FT '\n## Product ', # Section headers like ## Product Overview '\n## Compatibility', # Other major sections '\n## Installation', '\n## Troubleshooting' ] best_end = end # Look for product boundaries within last 30% of chunk search_range = int(target_chars * 0.3) for i in range(max(end - search_range, start), end): for boundary in product_boundaries: if text[i:].startswith(boundary): best_end = i break if best_end != end: break # Priority 2: Sentence boundaries within reasonable range if best_end == end: sentence_ends = ['. ', '! ', '? ', '.\n', '!\n', '?\n'] # Look for sentence endings within last 20% of chunk search_range = int(target_chars * 0.2) for i in range(max(end - search_range, start), end): for sentence_end in sentence_ends: if text[i:i+len(sentence_end)] == sentence_end: best_end = i + len(sentence_end) break if best_end != end: break # Priority 3: Paragraph boundaries if no sentence found if best_end == end: search_range = int(target_chars * 0.2) for i in range(max(end - search_range, start), end): if text[i:i+2] == '\n\n': best_end = i + 2 break # Priority 4: Word boundaries as fallback if best_end == end: for i in range(end - 1, max(end - 50, start), -1): if text[i] == ' ': best_end = i + 1 break end = best_end chunk = text[start:end].strip() if chunk and len(chunk) > 50: # Only include meaningful chunks # Verify token count is reasonable chunk_tokens = self.estimate_tokens(chunk) if chunk_tokens >= 50: # Minimum chunk size chunks.append(chunk) # Move start position with overlap start = end - overlap_chars if start >= len(text): break return chunks def generate_embedding(self, text: str) -> List[float]: """Generate embedding for a text using sentence transformer""" embedding = self.embedding_model.encode(text, convert_to_tensor=False) return embedding.tolist() def index_documents(self, force_reindex: bool = False) -> Dict[str, int]: """ Index all documents in the knowledge base into ChromaDB """ logging.info("Starting document indexing...") # Scan knowledge base self.kb_manager.scan_documents() # Check if we need to reindex if not force_reindex: current_count = self.collection.count() total_docs = sum(len(cat['documents']) for cat in self.kb_manager.documents.values()) if current_count > 0: logging.info(f"Collection already has {current_count} chunks. Use force_reindex=True to rebuild.") return {"existing_chunks": current_count, "documents": total_docs} # Clear existing collection if force reindex if force_reindex and self.collection.count() > 0: logging.info("Clearing existing collection for reindexing...") self.chroma_client.delete_collection(self.collection_name) self.collection = self.chroma_client.create_collection( name=self.collection_name, metadata={"description": "Too Many Cables knowledge base embeddings"} ) # Process all documents stats = {"documents": 0, "chunks": 0, "categories": {}} for category, cat_info in self.kb_manager.documents.items(): stats["categories"][category] = 0 for doc_info in cat_info["documents"]: logging.info(f"Processing: {doc_info['title']}") # Load document content content = self.kb_manager.load_document_content(doc_info['path']) if not content: continue # Normalize and clean content before chunking normalized_content = self.normalize_content(content) if not normalized_content or len(normalized_content) < 100: logging.warning(f"Skipping {doc_info['title']} - content too short after normalization") continue # Chunk the normalized document chunks = self.chunk_text(normalized_content) # Process each chunk chunk_ids = [] chunk_texts = [] chunk_embeddings = [] chunk_metadatas = [] for i, chunk in enumerate(chunks): chunk_id = f"{doc_info['filename']}_{i}" chunk_ids.append(chunk_id) chunk_texts.append(chunk) # Generate embedding embedding = self.generate_embedding(chunk) chunk_embeddings.append(embedding) # Create metadata metadata = { "document_title": doc_info['title'], "document_path": doc_info['path'], "category": category, "chunk_index": i, "total_chunks": len(chunks), "document_filename": doc_info['filename'] } chunk_metadatas.append(metadata) # Add chunks to collection if chunk_ids: self.collection.add( ids=chunk_ids, embeddings=chunk_embeddings, documents=chunk_texts, metadatas=chunk_metadatas ) stats["chunks"] += len(chunk_ids) stats["categories"][category] += len(chunk_ids) stats["documents"] += 1 logging.info(f"Indexing complete! Processed {stats['documents']} documents into {stats['chunks']} chunks") return stats def calculate_relevance_score(self, query: str, chunk: str, metadata: Dict) -> float: """ Calculate enhanced relevance score using multiple signals """ query_lower = query.lower() chunk_lower = chunk.lower() score = 0.0 # Exact phrase matches (high value) for word in query_lower.split(): if len(word) > 3: # Skip short words if word in chunk_lower: score += 0.3 # Bonus for multiple occurrences score += 0.1 * (chunk_lower.count(word) - 1) # Category relevance bonus category = metadata.get('category', '').lower() if 'policy' in query_lower and 'policies' in category: score += 0.2 elif 'manual' in query_lower and 'manual' in category: score += 0.2 elif 'faq' in query_lower and 'faq' in category: score += 0.2 # Policy-specific query matching (warranty, return, shipping) policy_types = { 'warranty': ['warranty', 'guarantee', 'defect', 'lifetime', 'coverage', 'claim'], 'return': ['return', 'refund', 'money-back', 'exchange'], 'shipping': ['shipping', 'delivery', 'freight', 'ship'] } # Detect policy type in query detected_policy = None for policy_type, keywords in policy_types.items(): if any(kw in query_lower for kw in keywords): detected_policy = policy_type break # Boost chunks that match the detected policy type if detected_policy: policy_keywords = policy_types[detected_policy] matches = sum(1 for kw in policy_keywords if kw in chunk_lower) score += 0.20 * matches # Increased from 0.15 to boost correct policy type more # Penalize chunks about OTHER policy types (stronger penalty) for other_type, other_keywords in policy_types.items(): if other_type != detected_policy: # Count how many keywords from the WRONG policy appear wrong_matches = sum(1 for kw in other_keywords[:3] if kw in chunk_lower) if wrong_matches > 0: score -= 0.25 * wrong_matches # Much stronger penalty # Penalize boilerplate/footer content more aggressively boilerplate_indicators = [ 'contact our customer service', 'email:', 'phone:', 'mailing address', 'policy effective', 'subject to change', '@toomanycables.com', 'live chat:', 'social media:', '[corporate' ] boilerplate_count = sum(1 for indicator in boilerplate_indicators if indicator in chunk_lower) if boilerplate_count > 0: score -= 0.25 * boilerplate_count # Stronger penalty, scales with amount of boilerplate # Product type matching product_terms = { 'usb-c': ['usb-c', 'usbc', 'type-c'], 'hdmi': ['hdmi', '4k', '8k', 'display'], 'lightning': ['lightning', 'iphone', 'apple'], 'usb-a': ['usb-a', 'usba', 'standard usb'] } for product, variants in product_terms.items(): if any(variant in query_lower for variant in variants): if any(variant in chunk_lower for variant in variants): score += 0.25 # Length penalty for very long chunks (prefer concise answers) if len(chunk) > 800: score -= 0.1 # Chunk position bonus (earlier chunks often have key info) chunk_index = metadata.get('chunk_index', 0) if chunk_index == 0: score += 0.1 elif chunk_index == 1: score += 0.05 return score def semantic_search(self, query: str, n_results: int = 5) -> List[Dict[str, Any]]: """ Perform semantic search using ChromaDB """ if self.collection.count() == 0: logging.warning("Collection is empty. Run index_documents() first.") return [] # Generate query embedding query_embedding = self.generate_embedding(query) # Search in ChromaDB results = self.collection.query( query_embeddings=[query_embedding], n_results=min(n_results, self.collection.count()), include=["documents", "metadatas", "distances"] ) # Format results formatted_results = [] if results['documents'] and results['documents'][0]: documents = results['documents'][0] metadatas = results['metadatas'][0] distances = results['distances'][0] for i, (doc, metadata, distance) in enumerate(zip(documents, metadatas, distances)): # ChromaDB uses L2 (squared Euclidean) distance by default # For normalized embeddings, we can convert L2 to cosine similarity: # cosine_similarity = 1 - (L2_distance^2 / 2) # But distances returned might already be squared, so let's use inverse distance as similarity # Lower distance = higher similarity # For L2 distance on normalized vectors: distance typically ranges 0-2 # Convert to similarity score where 0 distance = 1.0 similarity if distance < 0.0001: # Perfect match similarity = 1.0 else: # For L2 distance on normalized vectors: similarity ≈ 1 - (distance²/2) # But ChromaDB already returns squared distance, so: similarity = 1 - (distance/2) similarity = max(0.0, 1.0 - (distance / 2.0)) formatted_results.append({ "document": doc, "metadata": metadata, "similarity": similarity, "distance": distance, # Keep original distance for debugging "rank": i + 1 }) return formatted_results def retrieve_and_rerank(self, query: str, initial_k: int = 20, final_k: int = 5, similarity_threshold: float = 0.30) -> List[Dict[str, Any]]: """ Retrieve top-k candidates and re-rank to final top results This is the core optimization for better RAG quality Args: query: Search query initial_k: Number of candidates to retrieve initially final_k: Number of final results to return similarity_threshold: Minimum similarity score (0.20 for L2 distance conversion) """ # Step 1: Retrieve top 20 candidates using vector similarity candidates = self.semantic_search(query, n_results=initial_k) if not candidates: return [] # Filter by similarity threshold before reranking candidates = [c for c in candidates if c['similarity'] >= similarity_threshold] if not candidates: logging.info(f"No candidates meet similarity threshold {similarity_threshold}") return [] logging.debug(f"Filtered to {len(candidates)} candidates meeting threshold {similarity_threshold}") # Step 2: Re-rank using enhanced relevance scoring for candidate in candidates: # Combine vector similarity with relevance scoring relevance_score = self.calculate_relevance_score( query, candidate['document'], candidate['metadata'] ) # Weighted combination: 70% vector similarity + 30% relevance features candidate['final_score'] = (0.7 * candidate['similarity']) + (0.3 * relevance_score) candidate['relevance_score'] = relevance_score # Debug logging for top candidates if candidate['similarity'] > 0.40: doc_title = candidate['metadata'].get('document_title', 'Unknown')[:30] content_preview = candidate['document'][:50].replace('\n', ' ') logging.debug(f"Rerank: '{doc_title}' sim={candidate['similarity']:.3f}, rel={relevance_score:.3f}, final={candidate['final_score']:.3f}, preview='{content_preview}'") # Step 3: Re-sort by final score and take top results reranked = sorted(candidates, key=lambda x: x['final_score'], reverse=True) # Step 4: Diversity filtering - avoid too many chunks from same document final_results = [] seen_documents = set() for result in reranked: doc_title = result['metadata']['document_title'] # Allow max 2 chunks per document in final results doc_count = sum(1 for r in final_results if r['metadata']['document_title'] == doc_title) if doc_count < 2 or len(final_results) < final_k // 2: final_results.append(result) seen_documents.add(doc_title) if len(final_results) >= final_k: break return final_results def retrieve_and_rerank_filtered(self, query: str, target_categories: List[str], initial_k: int = 20, final_k: int = 5, similarity_threshold: float = 0.30) -> List[Dict[str, Any]]: """ Category-filtered retrieve and rerank for section routing Args: query: Search query target_categories: Categories to filter by initial_k: Number of candidates to retrieve initially final_k: Number of final results to return similarity_threshold: Minimum similarity score (0.20 for L2 distance conversion) """ if self.collection.count() == 0: logging.warning("Collection is empty. Run index_documents() first.") return [] # Generate query embedding query_embedding = self.generate_embedding(query) # Search with larger initial pool for filtering results = self.collection.query( query_embeddings=[query_embedding], n_results=min(initial_k * 2, self.collection.count()), # Get more for filtering include=["documents", "metadatas", "distances"] ) # Filter by target categories first category_filtered = [] if results['documents'] and results['documents'][0]: documents = results['documents'][0] metadatas = results['metadatas'][0] distances = results['distances'][0] for doc, metadata, distance in zip(documents, metadatas, distances): if metadata.get('category', '').lower() in [cat.lower() for cat in target_categories]: # Use same similarity calculation as semantic_search if distance < 0.0001: similarity = 1.0 else: similarity = max(0.0, 1.0 - (distance / 2.0)) # Apply similarity threshold to filter out low-similarity documents if similarity >= similarity_threshold: category_filtered.append({ "document": doc, "metadata": metadata, "similarity": similarity }) else: # Log filtered documents at debug level doc_title = metadata.get('document_title', 'Unknown')[:40] logging.debug(f"Filtered out '{doc_title}': similarity={similarity:.3f} < threshold {similarity_threshold}") # If we don't have enough results from target categories, fall back to general search if len(category_filtered) < final_k: logging.info(f"Only {len(category_filtered)} results in target categories meet threshold, falling back to general search") return self.retrieve_and_rerank(query, initial_k, final_k, similarity_threshold) # Apply reranking to filtered results for candidate in category_filtered[:initial_k]: # Limit to initial_k for reranking relevance_score = self.calculate_relevance_score( query, candidate['document'], candidate['metadata'] ) candidate['final_score'] = (0.7 * candidate['similarity']) + (0.3 * relevance_score) candidate['relevance_score'] = relevance_score # Sort by final score and apply diversity filtering reranked = sorted(category_filtered[:initial_k], key=lambda x: x['final_score'], reverse=True) # Log top results for debugging if reranked: logging.info(f"Top result: similarity={reranked[0]['similarity']:.3f}, final_score={reranked[0]['final_score']:.3f}") # Diversity filtering final_results = [] seen_documents = set() for result in reranked: doc_title = result['metadata']['document_title'] doc_count = sum(1 for r in final_results if r['metadata']['document_title'] == doc_title) if doc_count < 2 or len(final_results) < final_k // 2: final_results.append(result) seen_documents.add(doc_title) if len(final_results) >= final_k: break return final_results def get_relevant_context(self, query: str, n_results: int = 5, similarity_threshold: float = 0.20) -> str: """ Get relevant context for a query using semantic search Args: query: Search query n_results: Maximum number of results to return similarity_threshold: Minimum similarity score (0.20 for L2 distance conversion) """ # Perform semantic search results = self.semantic_search(query, n_results) # Filter by similarity threshold relevant_results = [ r for r in results if r["similarity"] >= similarity_threshold ] if not relevant_results: return "" # Build context string context_parts = [] total_length = 0 max_context_length = 2000 # Group chunks by document to avoid repetition doc_chunks = {} for result in relevant_results: doc_title = result["metadata"]["document_title"] if doc_title not in doc_chunks: doc_chunks[doc_title] = [] doc_chunks[doc_title].append(result) # Build context from best chunks per document for doc_title, chunks in doc_chunks.items(): # Sort chunks by similarity chunks.sort(key=lambda x: x["similarity"], reverse=True) # Take best chunk from this document best_chunk = chunks[0] doc_header = f"\n--- {doc_title} (Category: {best_chunk['metadata']['category']}) ---\n" chunk_content = best_chunk["document"] combined_length = len(doc_header) + len(chunk_content) if total_length + combined_length > max_context_length: # Truncate to fit remaining_space = max_context_length - total_length - len(doc_header) if remaining_space > 100: chunk_content = chunk_content[:remaining_space] + "..." context_parts.append(doc_header + chunk_content) break context_parts.append(doc_header + chunk_content) total_length += combined_length return "".join(context_parts) def get_collection_stats(self) -> Dict[str, Any]: """Get statistics about the vector database""" try: count = self.collection.count() # Get sample of metadata to analyze categories if count > 0: sample_results = self.collection.get(limit=min(count, 100)) categories = {} if sample_results["metadatas"]: for metadata in sample_results["metadatas"]: cat = metadata.get("category", "unknown") categories[cat] = categories.get(cat, 0) + 1 return { "total_chunks": count, "categories": categories, "embedding_dimension": self.embedding_model.get_sentence_embedding_dimension(), "collection_name": self.collection_name } else: return { "total_chunks": 0, "categories": {}, "embedding_dimension": self.embedding_model.get_sentence_embedding_dimension(), "collection_name": self.collection_name } except Exception as e: return {"error": str(e)} # Production ready - no test code included