rag_query.py

#!/usr/bin/env python3
"""
RAG Query Engine for YouTube Transcripts
Uses LLM-based extraction and synthesis (works with LM Studio and OpenAI)
"""

import json
import sys
import argparse
import logging
from typing import List, Dict, Optional, Union
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

import config
from chunker import chunk_video_transcripts
from llm_client import LLMClient


def setup_logging(verbose: bool = False, log_file: Optional[str] = None):
    """Setup logging configuration."""
    level = logging.DEBUG if verbose else logging.INFO
    log_format = '%(asctime)s - %(levelname)s - %(message)s'
    
    # Configure root logger
    root_logger = logging.getLogger()
    root_logger.setLevel(level)
    
    # Clear existing handlers
    root_logger.handlers = []
    
    # Console handler (always) - unbuffered
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(level)
    console_handler.setFormatter(logging.Formatter(log_format, datefmt='%Y-%m-%d %H:%M:%S'))
    # Force immediate flushing
    console_handler.stream = sys.stdout
    root_logger.addHandler(console_handler)
    
    # Force Python to use unbuffered output
    sys.stdout.reconfigure(line_buffering=True) if hasattr(sys.stdout, 'reconfigure') else None
    
    # File handler (if specified)
    if log_file:
        file_handler = logging.FileHandler(log_file, mode='a', encoding='utf-8')
        file_handler.setLevel(logging.DEBUG)  # Always DEBUG for file logs
        file_handler.setFormatter(logging.Formatter(log_format, datefmt='%Y-%m-%d %H:%M:%S'))
        root_logger.addHandler(file_handler)
        print(f"Detailed logs will be saved to: {log_file}")
    
    logger = logging.getLogger(__name__)
    for handler in logger.handlers:
        handler.flush()
    
    return logger


class RAGQueryEngine:
    """Query engine for YouTube transcripts using LLM-based extraction"""
    
    def __init__(self, transcript_files: Union[str, List[str]], api_type: Optional[str] = None, logger: Optional[logging.Logger] = None):
        """
        Initialize query engine.
        
        Args:
            transcript_files: Path to JSON file(s) with transcripts (string or list of strings)
            api_type: 'lmstudio' or 'openai' (legacy, uses EXTRACTION_API_TYPE and SYNTHESIS_API_TYPE if None)
            logger: Logger instance (creates one if None)
        """
        # Support both single file and list of files
        if isinstance(transcript_files, str):
            self.transcript_files = [transcript_files]
        else:
            self.transcript_files = transcript_files
        
        # Keep backward compatibility
        self.transcript_file = self.transcript_files[0] if self.transcript_files else None
        self.logger = logger or logging.getLogger(__name__)
        
        # Use separate API types for extraction and synthesis
        extraction_api_type = config.EXTRACTION_API_TYPE if api_type is None else api_type
        synthesis_api_type = config.SYNTHESIS_API_TYPE if api_type is None else api_type
        
        self.logger.info(f"Initializing LLM clients... (Extraction: {extraction_api_type}, Synthesis: {synthesis_api_type})")
        self.extraction_client = LLMClient(api_type=extraction_api_type, logger=self.logger)
        self.synthesis_client = LLMClient(api_type=synthesis_api_type, logger=self.logger)
        # Keep llm_client for backward compatibility (uses extraction client)
        self.llm_client = self.extraction_client
        self.logger.info("LLM clients initialized")
        
        self.transcripts_data = None
        self.chunks = None
        
        # Load transcripts
        self.logger.info("Loading transcripts...")
        self.load_transcripts()
        
        # Pre-chunk transcripts
        self.logger.info("Chunking transcripts...")
        self.chunk_transcripts()
        self.logger.info("Initialization complete")
    
    def load_transcripts(self):
        """Load transcripts from JSON file(s)"""
        self.transcripts_data = []
        
        for transcript_file in self.transcript_files:
            try:
                with open(transcript_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                
                # Get transcripts array
                file_transcripts = data.get('transcripts', [])
                
                # Filter only successful transcripts
                file_transcripts = [
                    t for t in file_transcripts 
                    if t.get('success') and t.get('transcript_full_text')
                ]
                
                # Add source file info to each transcript
                for t in file_transcripts:
                    t['_source_file'] = transcript_file
                
                self.transcripts_data.extend(file_transcripts)
                self.logger.info(f"Loaded {len(file_transcripts)} videos from {transcript_file}")
                
            except FileNotFoundError:
                self.logger.error(f"Transcript file not found: {transcript_file}")
            except json.JSONDecodeError as e:
                self.logger.error(f"Invalid JSON file {transcript_file}: {e}")
        
        if not self.transcripts_data:
            self.logger.error("No transcripts loaded from any file!")
            sys.exit(1)
        
        self.logger.info(f"Total: {len(self.transcripts_data)} videos loaded from {len(self.transcript_files)} file(s)")
    
    def chunk_transcripts(self):
        """Chunk all transcripts by combining videos (videos are never split)"""
        self.logger.info("Chunking transcripts (combining videos, not splitting)...")
        self.chunks = chunk_video_transcripts(
            self.transcripts_data,
            max_tokens=config.MAX_CHUNK_TOKENS,
            model=config.EXTRACTION_MODEL
        )
        total_videos_in_chunks = sum(chunk.get('num_videos', 1) for chunk in self.chunks)
        avg_videos = total_videos_in_chunks / len(self.chunks) if self.chunks else 0
        self.logger.info(f"Created {len(self.chunks)} chunks containing {total_videos_in_chunks} videos (avg {avg_videos:.1f} videos per chunk)")
    
    def extract_relevant_content_parallel(
        self,
        question: str,
        max_parallel: Optional[int] = None,
        return_logs: bool = False
    ) -> tuple:
        """
        Extract relevant content from all chunks in parallel.
        
        Args:
            question: User's question
            max_parallel: Maximum parallel LLM calls (uses config if None)
            return_logs: If True, returns logs along with extracted content
        
        Returns:
            List of extracted content with metadata, and optionally extraction logs
        """
        max_parallel = max_parallel or config.MAX_PARALLEL_EXTRACTIONS
        
        self.logger.info(f"Extracting relevant content from {len(self.chunks)} chunks...")
        self.logger.info(f"Using {max_parallel} parallel workers")
        
        extracted_content = []
        extraction_logs = []  # Store logs for each extraction
        
        # Process chunks in parallel
        with ThreadPoolExecutor(max_workers=max_parallel) as executor:
            # Submit all tasks
            futures = {
                executor.submit(
                    self.extraction_client.extract_relevant_content,
                    chunk['text'],
                    question,
                    chunk['videos'],  # Pass list of videos in chunk
                    chunk['chunk_id'],  # Pass chunk_id for logging
                    None,  # model (uses default)
                    None,  # temperature (uses default)
                    return_logs  # return_logs flag
                ): chunk
                for chunk in self.chunks
            }
            
            # Process results with progress bar (show time estimates)
            with tqdm(total=len(futures), desc="Extracting content", unit="chunk", 
                     bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]') as pbar:
                for future in as_completed(futures):
                    chunk = futures[future]
                    try:
                        result = future.result(timeout=300)  # 5 minute timeout per chunk
                        
                        if return_logs:
                            extracted_text, log_data = result
                            extraction_logs.append(log_data)
                            # Check log_data to see what the extraction actually returned
                            if log_data.get('result') == 'NO_RELEVANT_CONTENT':
                                self.logger.debug(f"Chunk {chunk['chunk_id']} - Log indicates NO_RELEVANT_CONTENT")
                                extracted_text = None  # Treat as no content
                            elif log_data.get('result') == 'ERROR':
                                self.logger.warning(f"Chunk {chunk['chunk_id']} - Extraction had error: {log_data.get('error', 'Unknown error')}")
                                extracted_text = None  # Treat as no content
                        else:
                            extracted_text = result
                        
                        # Check if we got valid content (not empty, not "NO_RELEVANT_CONTENT")
                        if extracted_text and extracted_text.strip():
                            # Check if it's explicitly NO_RELEVANT_CONTENT (case-insensitive)
                            text_upper = extracted_text.strip().upper()
                            if text_upper == "NO_RELEVANT_CONTENT" or text_upper.startswith("NO_RELEVANT_CONTENT"):
                                self.logger.debug(f"Chunk {chunk['chunk_id']} - Explicitly marked as NO_RELEVANT_CONTENT")
                            else:
                                # For chunks with multiple videos, create entries for each video
                                # The extracted text should already have video attribution
                                for video in chunk['videos']:
                                    extracted_content.append({
                                        'content': extracted_text.strip(),  # Contains video attribution
                                        'video_id': video['video_id'],
                                        'video_title': video['video_title'],
                                        'video_url': video.get('video_url', ''),
                                        'chunk_id': chunk['chunk_id']
                                    })
                                self.logger.info(f"Added extracted content from chunk {chunk['chunk_id']} ({len(chunk['videos'])} videos, {len(extracted_text)} chars)")
                        else:
                            self.logger.debug(f"Chunk {chunk['chunk_id']} - No valid content (extracted_text is empty or None)")
                    except Exception as e:
                        self.logger.error(f"Error processing chunk {chunk['chunk_id']}: {e}", exc_info=True)
                        if return_logs:
                            extraction_logs.append({
                                "chunk_id": chunk.get('chunk_id'),
                                "error": str(e),
                                "result": "ERROR"
                            })
                    
                    pbar.update(1)
        
        self.logger.info(f"Extracted relevant content from {len(extracted_content)} chunks")
        
        if return_logs:
            return extracted_content, extraction_logs
        return extracted_content
    
    def query(
        self,
        question: str,
        max_parallel: Optional[int] = None,
        return_logs: bool = False
    ) -> Dict[str, any]:
        """
        Query the transcript database.
        
        Args:
            question: User's question
            max_parallel: Maximum parallel extractions (uses config if None)
            return_logs: If True, includes detailed logs in the response
        
        Returns:
            Dict with 'answer', 'sources', metadata, and optionally 'logs'
        """
        # Stage 1: Extract relevant content from chunks
        if return_logs:
            extracted_content, extraction_logs = self.extract_relevant_content_parallel(question, max_parallel, return_logs=True)
        else:
            extracted_content = self.extract_relevant_content_parallel(question, max_parallel, return_logs=False)
            extraction_logs = []
        
        if not extracted_content:
            result = {
                'answer': "No relevant content found in the transcripts for your question.",
                'sources': [],
                'extracted_chunks': 0,
                'total_chunks': len(self.chunks)
            }
            if return_logs:
                result['logs'] = {
                    'extraction_logs': extraction_logs,
                    'synthesis_log': None
                }
            return result
        
        # Stage 2: Deduplicate similar content (simple approach: keep all for now)
        # TODO: Add deduplication if needed
        
        # Stage 3: Synthesize final answer
        self.logger.info("Synthesizing final answer...")
        result = self.synthesis_client.synthesize_answer(question, extracted_content, return_logs=return_logs)
        
        result['extracted_chunks'] = len(extracted_content)
        result['total_chunks'] = len(self.chunks)
        result['total_videos'] = len(self.transcripts_data)
        result['extracted_content'] = extracted_content  # Store for context reuse
        
        if return_logs:
            result['logs'] = {
                'extraction_logs': extraction_logs,
                'synthesis_log': result.pop('synthesis_log', None)
            }
        
        return result


def main():
    """Main CLI interface"""
    parser = argparse.ArgumentParser(
        description='Query YouTube transcripts using RAG with LLM extraction'
    )
    parser.add_argument(
        'transcript_file',
        help='Path to JSON file with transcripts (from transcript_fetcher.py)'
    )
    parser.add_argument(
        'question',
        help='Question to ask about the transcripts'
    )
    parser.add_argument(
        '--api-type',
        choices=['lmstudio', 'openai'],
        default=None,
        help='API type to use (default: from config or environment)'
    )
    parser.add_argument(
        '--max-parallel',
        type=int,
        default=None,
        help=f'Maximum parallel extractions (default: {config.MAX_PARALLEL_EXTRACTIONS})'
    )
    parser.add_argument(
        '--config',
        action='store_true',
        help='Show configuration and exit'
    )
    parser.add_argument(
        '--verbose', '-v',
        action='store_true',
        help='Enable verbose logging (show all LLM responses)'
    )
    parser.add_argument(
        '--log-file',
        type=str,
        default=None,
        help='Save detailed logs to file (includes full AI inputs/outputs)'
    )
    
    args = parser.parse_args()
    
    if args.config:
        config.print_config()
        return
    
    # Setup logging
    logger = setup_logging(verbose=args.verbose, log_file=args.log_file)
    
    # Initialize query engine
    logger.info("Starting query engine initialization...")
    sys.stdout.flush()
    try:
        engine = RAGQueryEngine(args.transcript_file, api_type=args.api_type, logger=logger)
        logger.info("Query engine initialized successfully")
    except KeyboardInterrupt:
        logger.error("Interrupted by user")
        sys.exit(1)
    except Exception as e:
        logger.error(f"Error initializing query engine: {e}", exc_info=True)
        sys.exit(1)
    
    # Query
    print(f"\nQuestion: {args.question}\n")
    result = engine.query(args.question, max_parallel=args.max_parallel)
    
    # Print results
    print("\n" + "="*80)
    print("ANSWER:")
    print("="*80)
    print(result['answer'])
    print("\n" + "="*80)
    print("SOURCES:")
    print("="*80)
    for source in result['sources']:
        print(f"- {source['video_title']}")
        if source.get('video_url'):
            print(f"  URL: {source['video_url']}")
    
    print("\n" + "="*80)
    print("STATISTICS:")
    print("="*80)
    print(f"Total videos: {result.get('total_videos', 0)}")
    print(f"Total chunks: {result.get('total_chunks', 0)}")
    print(f"Relevant chunks: {result.get('extracted_chunks', 0)}")
    print(f"Sources: {len(result['sources'])}")


if __name__ == '__main__':
    main()