rag_cli.py

#!/usr/bin/env python3
import sys
import json
import argparse
import logging
import os
from pathlib import Path

os.environ['GRPC_VERBOSITY'] = 'ERROR'
os.environ['GLOG_minloglevel'] = '2'
logging.getLogger("chromadb").setLevel(logging.ERROR)
logging.getLogger("google").setLevel(logging.ERROR)


sys.path.insert(0, str(Path(__file__).parent / 'src'))

from crawler.web_crawler import WebCrawler
from indexer.indexer import index_documents
from qa.qa_engine import ask
import config

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

_vector_store_instance = None

def get_vector_store():
    """Get or create a singleton vector store instance."""
    global _vector_store_instance
    if _vector_store_instance is None:
        from indexer.vector_store import VectorStore
        _vector_store_instance = VectorStore(
            config.VECTOR_STORE_PATH,
            collection_name=config.COLLECTION_NAME
        )
    return _vector_store_instance

def crawl_site(args):
    """Crawl a website and save results."""
    crawler = WebCrawler(
        max_pages=args.max_pages,
        max_depth=args.max_depth,
        crawl_delay_ms=args.crawl_delay_ms,
        request_timeout=config.REQUEST_TIMEOUT
    )
    
    logger.info(f"Crawling {args.url} (max {args.max_pages} pages, depth {args.max_depth})")
    result = crawler.crawl(args.url)
    
    output = config.CRAWLED_DIR / "last_crawl.json"
    with open(output, 'w', encoding='utf-8') as f:
        json.dump(result, f, indent=2)
    
    logger.info(f"Crawled {result['page_count']} pages")
    logger.info(f"Results saved to {output}")
    return result

def index_content(args, pages=None):
    """Index previously crawled content."""
    if not pages:
        try:
            with open(config.CRAWLED_DIR / "last_crawl.json") as f:
                data = json.load(f)
                pages = data.get('pages', {})
        except FileNotFoundError:
            logger.error("No crawl results found. Run crawl first.")
            sys.exit(1)
    
    logger.info(f"Indexing {len(pages)} pages")
    
    vs = get_vector_store()
    
    result = index_documents(
        pages=pages,
        chunk_size=args.chunk_size,
        chunk_overlap=args.chunk_overlap,
        embedding_model=args.embedding_model,
        vector_store=vs  # Pass the vector store instance
    )
    
    logger.info(f"Indexed {result['vector_count']} chunks")
    if result['errors']:
        logger.warning(f"Encountered {len(result['errors'])} errors while indexing")
    
    return result

def ask_question(args):
    vs = get_vector_store()
    
    debug = getattr(args, 'verbose', False)
    
    result = ask(
        question=args.question,
        top_k=args.top_k,
        vector_store=vs,  
        debug=debug
    )
    
    # Pretty print results
    print("\nQuestion:", args.question)
    print("\nAnswer:", result['answer'])
    print("\nSources:")
    max_sources = getattr(args, "max_sources", 2)
    for i, src in enumerate(result['sources'][:max_sources], 1):
        print(f"\n{i}. {src['url']}")
        snippet = src['snippet'][:200]
        print(f"   Snippet: {snippet}...")
    
    print("\nTimings:")
    for k, v in result['timings'].items():
        print(f"  {k}: {v}ms")

def main():
    parser = argparse.ArgumentParser(description="RAG Service CLI")
    subparsers = parser.add_subparsers(dest='command', help='Command to run')
    
    crawl_parser = subparsers.add_parser('crawl', help='Crawl a website')
    crawl_parser.add_argument('url', help='Starting URL to crawl')
    crawl_parser.add_argument('--max-pages', type=int, default=config.MAX_PAGES)
    crawl_parser.add_argument('--max-depth', type=int, default=config.MAX_DEPTH)
    crawl_parser.add_argument('--crawl-delay-ms', type=int, default=config.CRAWL_DELAY_MS)
    
    index_parser = subparsers.add_parser('index', help='Index crawled content')
    index_parser.add_argument('--chunk-size', type=int, default=config.CHUNK_SIZE)
    index_parser.add_argument('--chunk-overlap', type=int, default=config.CHUNK_OVERLAP)
    index_parser.add_argument('--embedding-model', default=config.EMBEDDING_MODEL)
    
    ask_parser = subparsers.add_parser('ask', help='Ask a question')
    ask_parser.add_argument('question', help='Question to ask')
    ask_parser.add_argument('--top-k', type=int, default=config.TOP_K)
    ask_parser.add_argument('--max-sources', type=int, default=2, help='Number of source URLs to show')

    
    pipeline_parser = subparsers.add_parser('pipeline', help='Run full pipeline')
    pipeline_parser.add_argument('url', help='Starting URL to crawl')
    pipeline_parser.add_argument('question', help='Question to ask')
    pipeline_parser.add_argument('--max-pages', type=int, default=config.MAX_PAGES)
    pipeline_parser.add_argument('--max-depth', type=int, default=config.MAX_DEPTH)
    pipeline_parser.add_argument('--crawl-delay-ms', type=int, default=config.CRAWL_DELAY_MS)
    pipeline_parser.add_argument('--chunk-size', type=int, default=config.CHUNK_SIZE)
    pipeline_parser.add_argument('--chunk-overlap', type=int, default=config.CHUNK_OVERLAP)
    pipeline_parser.add_argument('--embedding-model', default=config.EMBEDDING_MODEL)
    pipeline_parser.add_argument('--top-k', type=int, default=config.TOP_K)
    
    args = parser.parse_args()
    
    if args.command == 'crawl':
        crawl_site(args)
    elif args.command == 'index':
        index_content(args)
    elif args.command == 'ask':
        ask_question(args)
    elif args.command == 'pipeline':
        logger.info("Starting pipeline...")
        pages = crawl_site(args)
        index_content(args, pages=pages.get('pages'))
        ask_question(args)
    else:
        parser.print_help()
        sys.exit(1)

if __name__ == '__main__':
    main()