main.py

import sys
import argparse
from rag_pipeline import RAGPipeline
from config import Config

def main():
    # Validate configuration
    try:
        Config.validate()
    except ValueError as e:
        print(f"Configuration Error: {e}")
        print("\nPlease create a .env file with your OPENAI_API_KEY")
        print("Example: OPENAI_API_KEY=sk-your-key-here")
        return
    
    parser = argparse.ArgumentParser(description='RAG System - Query PDFs with OpenAI')
    parser.add_argument('--ingest', type=str, help='Path to PDF file to ingest')
    parser.add_argument('--query', type=str, help='Query to ask')
    parser.add_argument('--interactive', action='store_true', help='Start interactive mode')
    parser.add_argument('--no-vision', action='store_true', help='Skip vision analysis (faster)')
    parser.add_argument('--top-k', type=int, default=5, help='Number of chunks to retrieve')
    
    args = parser.parse_args()
    
    # Initialize pipeline
    rag = RAGPipeline()
    
    # Ingest PDF if provided
    if args.ingest:
        print(f"Ingesting PDF: {args.ingest}")
        use_vision = not args.no_vision
        rag.ingest_pdf(args.ingest, use_vision=use_vision)
    else:
        # Try to load existing index
        if not rag.load_existing_index():
            print("\nNo existing index found and no PDF provided.")
            print("Usage: python main.py --ingest <pdf_path>")
            return
    
    # Handle query modes
    if args.interactive:
        rag.interactive_query()
    elif args.query:
        result = rag.query(args.query, top_k=args.top_k)
        
        if 'error' in result:
            print(f"Error: {result['error']}")
        else:
            print("\n" + "=" * 50)
            print("QUESTION:")
            print(result['question'])
            print("\n" + "-" * 50)
            print("ANSWER:")
            print(result['answer'])
            print("\n" + "-" * 50)
            print("SOURCES:")
            for i, source in enumerate(result['sources'], 1):
                print(f"{i}. Page {source['page']} ({source['type']}) - Score: {source['similarity_score']:.4f}")
            print("=" * 50)
    else:
        print("\nIndex loaded successfully!")
        print("\nUsage examples:")
        print("  Interactive mode: python main.py --interactive")
        print("  Single query: python main.py --query 'What is the main topic?'")
        print("  Ingest new PDF: python main.py --ingest document.pdf")

if __name__ == "__main__":
    main()