visual-rag-toolkit 0.1.1__py3-none-any.whl

visual_rag/cli/main.py

@@ -0,0 +1,629 @@
+#!/usr/bin/env python3
+"""
+Visual RAG Toolkit CLI
+
+Provides command-line interface for:
+- Processing PDFs (embedding, Cloudinary upload, Qdrant indexing)
+- Searching documents
+- Managing collections
+
+Usage:
+    # Process PDFs (like process_pdfs_saliency_v2.py)
+    visual-rag process --reports-dir ./pdfs --metadata-file metadata.json
+
+    # Search
+    visual-rag search --query "budget allocation" --collection my_docs
+
+    # Show collection info
+    visual-rag info --collection my_docs
+"""
+
+import argparse
+import logging
+import os
+import sys
+from pathlib import Path
+from urllib.parse import urlparse
+
+from dotenv import load_dotenv
+
+logger = logging.getLogger(__name__)
+
+
+def setup_logging(debug: bool = False):
+    """Configure logging."""
+    level = logging.DEBUG if debug else logging.INFO
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s - %(levelname)s - %(message)s",
+        force=True,
+    )
+
+
+def cmd_process(args):
+    """
+    Process PDFs: convert → embed → upload to Cloudinary → index in Qdrant.
+
+    Equivalent to process_pdfs_saliency_v2.py
+    """
+    from visual_rag import CloudinaryUploader, QdrantIndexer, VisualEmbedder, load_config
+    from visual_rag.indexing.pipeline import ProcessingPipeline
+
+    # Load environment
+    load_dotenv()
+
+    # Load config
+    config = {}
+    if args.config and Path(args.config).exists():
+        config = load_config(args.config)
+
+    # Get PDFs
+    reports_dir = Path(args.reports_dir)
+    if not reports_dir.exists():
+        logger.error(f"❌ Reports directory not found: {reports_dir}")
+        sys.exit(1)
+
+    pdf_paths = sorted(reports_dir.glob("*.pdf")) + sorted(reports_dir.glob("*.PDF"))
+    if not pdf_paths:
+        logger.error(f"❌ No PDF files found in: {reports_dir}")
+        sys.exit(1)
+
+    logger.info(f"📁 Found {len(pdf_paths)} PDF files")
+
+    # Load metadata mapping
+    metadata_mapping = {}
+    if args.metadata_file:
+        metadata_mapping = ProcessingPipeline.load_metadata_mapping(Path(args.metadata_file))
+
+    # Dry run - just show summary
+    if args.dry_run:
+        logger.info("🏃 DRY RUN MODE")
+        logger.info(f"   PDFs: {len(pdf_paths)}")
+        logger.info(f"   Metadata entries: {len(metadata_mapping)}")
+        logger.info(f"   Collection: {args.collection}")
+        logger.info(f"   Cloudinary: {'ENABLED' if not args.no_cloudinary else 'DISABLED'}")
+
+        for pdf in pdf_paths[:10]:
+            has_meta = "✓" if pdf.stem.lower() in metadata_mapping else "✗"
+            logger.info(f"   {has_meta} {pdf.name}")
+        if len(pdf_paths) > 10:
+            logger.info(f"   ... and {len(pdf_paths) - 10} more")
+        return
+
+    # Get settings
+    model_name = args.model or config.get("model", {}).get("name", "vidore/colSmol-500M")
+    collection_name = args.collection or config.get("qdrant", {}).get(
+        "collection_name", "visual_documents"
+    )
+
+    torch_dtype = None
+    if args.torch_dtype != "auto":
+        import torch
+
+        torch_dtype = {
+            "float32": torch.float32,
+            "float16": torch.float16,
+            "bfloat16": torch.bfloat16,
+        }[args.torch_dtype]
+
+    logger.info(f"🤖 Initializing embedder: {model_name}")
+    embedder = VisualEmbedder(
+        model_name=model_name,
+        batch_size=args.batch_size,
+        torch_dtype=torch_dtype,
+        processor_speed=str(getattr(args, "processor_speed", "fast")),
+    )
+
+    # Initialize Qdrant indexer
+    qdrant_url = (
+        os.getenv("SIGIR_QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("QDRANT_URL")
+    )
+    qdrant_api_key = (
+        os.getenv("SIGIR_QDRANT_KEY")
+        or os.getenv("SIGIR_QDRANT_API_KEY")
+        or os.getenv("DEST_QDRANT_API_KEY")
+        or os.getenv("QDRANT_API_KEY")
+    )
+
+    if not qdrant_url:
+        logger.error("❌ QDRANT_URL environment variable not set")
+        sys.exit(1)
+
+    logger.info(f"🔌 Connecting to Qdrant: {qdrant_url}")
+    indexer = QdrantIndexer(
+        url=qdrant_url,
+        api_key=qdrant_api_key,
+        collection_name=collection_name,
+        prefer_grpc=args.prefer_grpc,
+        vector_datatype=args.qdrant_vector_dtype,
+    )
+
+    # Create collection if needed
+    indexer.create_collection(force_recreate=args.force_recreate)
+    inferred_fields = []
+    inferred_fields.append({"field": "filename", "type": "keyword"})
+    inferred_fields.append({"field": "page_number", "type": "integer"})
+    inferred_fields.append({"field": "has_text", "type": "bool"})
+
+    if metadata_mapping:
+        keys = set()
+        for _, meta in metadata_mapping.items():
+            if isinstance(meta, dict):
+                keys.update(meta.keys())
+        for k in sorted(keys):
+            if k in ("filename", "page_number", "has_text"):
+                continue
+            inferred_type = "keyword"
+            for _, meta in metadata_mapping.items():
+                if not isinstance(meta, dict):
+                    continue
+                v = meta.get(k)
+                if isinstance(v, bool):
+                    inferred_type = "bool"
+                    break
+                if isinstance(v, int):
+                    inferred_type = "integer"
+                    break
+                if isinstance(v, float):
+                    inferred_type = "float"
+                    break
+            inferred_fields.append({"field": k, "type": inferred_type})
+
+    indexer.create_payload_indexes(fields=inferred_fields)
+
+    # Initialize Cloudinary uploader (optional)
+    cloudinary_uploader = None
+    if not args.no_cloudinary:
+        try:
+            project_name = config.get("project_name", "visual_docs")
+            cloudinary_uploader = CloudinaryUploader(folder=project_name)
+        except ValueError as e:
+            logger.warning(f"⚠️ Cloudinary not configured: {e}")
+            logger.warning("   Continuing without Cloudinary uploads")
+
+    # Create pipeline
+    pipeline = ProcessingPipeline(
+        embedder=embedder,
+        indexer=indexer,
+        cloudinary_uploader=cloudinary_uploader,
+        metadata_mapping=metadata_mapping,
+        config=config,
+        embedding_strategy=args.strategy,
+        crop_empty=bool(getattr(args, "crop_empty", False)),
+        crop_empty_percentage_to_remove=float(
+            getattr(args, "crop_empty_percentage_to_remove", 0.9)
+        ),
+        crop_empty_remove_page_number=bool(getattr(args, "crop_empty_remove_page_number", False)),
+    )
+
+    # Process PDFs
+    total_uploaded = 0
+    total_skipped = 0
+    total_failed = 0
+
+    skip_existing = not args.no_skip_existing
+
+    for pdf_idx, pdf_path in enumerate(pdf_paths, 1):
+        logger.info(f"\n{'='*60}")
+        logger.info(f"📄 [{pdf_idx}/{len(pdf_paths)}] {pdf_path.name}")
+        logger.info(f"{'='*60}")
+
+        result = pipeline.process_pdf(
+            pdf_path,
+            skip_existing=skip_existing,
+            upload_to_cloudinary=(not args.no_cloudinary),
+            upload_to_qdrant=True,
+        )
+
+        total_uploaded += result["uploaded"]
+        total_skipped += result["skipped"]
+        total_failed += result["failed"]
+
+    # Summary
+    logger.info(f"\n{'='*60}")
+    logger.info("📊 SUMMARY")
+    logger.info(f"{'='*60}")
+    logger.info(f"   Total PDFs: {len(pdf_paths)}")
+    logger.info(f"   Uploaded: {total_uploaded}")
+    logger.info(f"   Skipped: {total_skipped}")
+    logger.info(f"   Failed: {total_failed}")
+
+    info = indexer.get_collection_info()
+    if info:
+        logger.info(f"   Collection points: {info.get('points_count', 'N/A')}")
+
+
+def cmd_search(args):
+    """Search documents."""
+    from qdrant_client import QdrantClient
+
+    from visual_rag import VisualEmbedder
+    from visual_rag.retrieval import SingleStageRetriever, TwoStageRetriever
+
+    load_dotenv()
+
+    qdrant_url = (
+        os.getenv("SIGIR_QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("QDRANT_URL")
+    )
+    qdrant_api_key = (
+        os.getenv("SIGIR_QDRANT_KEY")
+        or os.getenv("SIGIR_QDRANT_API_KEY")
+        or os.getenv("DEST_QDRANT_API_KEY")
+        or os.getenv("QDRANT_API_KEY")
+    )
+
+    if not qdrant_url:
+        logger.error("❌ QDRANT_URL not set")
+        sys.exit(1)
+
+    # Initialize
+    logger.info(f"🤖 Loading model: {args.model}")
+    embedder = VisualEmbedder(
+        model_name=args.model, processor_speed=str(getattr(args, "processor_speed", "fast"))
+    )
+
+    logger.info("🔌 Connecting to Qdrant")
+    grpc_port = 6334 if args.prefer_grpc and urlparse(qdrant_url).port == 6333 else None
+    client = QdrantClient(
+        url=qdrant_url,
+        api_key=qdrant_api_key,
+        prefer_grpc=args.prefer_grpc,
+        grpc_port=grpc_port,
+        check_compatibility=False,
+    )
+    two_stage = TwoStageRetriever(client, args.collection)
+    single_stage = SingleStageRetriever(client, args.collection)
+
+    # Embed query
+    logger.info(f"🔍 Query: {args.query}")
+    query_embedding = embedder.embed_query(args.query)
+
+    # Build filter
+    filter_obj = None
+    if args.year or args.source or args.district:
+        filter_obj = two_stage.build_filter(
+            year=args.year,
+            source=args.source,
+            district=args.district,
+        )
+
+    # Search
+    query_np = query_embedding.detach().cpu().numpy()
+    if args.strategy == "single_full":
+        results = single_stage.search(
+            query_embedding=query_np,
+            top_k=args.top_k,
+            strategy="multi_vector",
+            filter_obj=filter_obj,
+        )
+    elif args.strategy == "single_tiles":
+        results = single_stage.search(
+            query_embedding=query_np,
+            top_k=args.top_k,
+            strategy="tiles_maxsim",
+            filter_obj=filter_obj,
+        )
+    elif args.strategy == "single_global":
+        results = single_stage.search(
+            query_embedding=query_np,
+            top_k=args.top_k,
+            strategy="pooled_global",
+            filter_obj=filter_obj,
+        )
+    else:
+        results = two_stage.search(
+            query_embedding=query_np,
+            top_k=args.top_k,
+            prefetch_k=args.prefetch_k,
+            filter_obj=filter_obj,
+            stage1_mode=args.stage1_mode,
+        )
+
+    # Display results
+    logger.info(f"\n📊 Results ({len(results)}):")
+    for i, result in enumerate(results, 1):
+        payload = result.get("payload", {})
+        score = result.get("score_final", result.get("score_stage1", 0))
+
+        filename = payload.get("filename", "N/A")
+        page_num = payload.get("page_number", "N/A")
+        year = payload.get("year", "N/A")
+        source = payload.get("source", "N/A")
+
+        logger.info(f"  {i}. {filename} p.{page_num}")
+        logger.info(f"     Score: {score:.4f} | Year: {year} | Source: {source}")
+
+        # Text snippet
+        text = payload.get("text", "")
+        if text and args.show_text:
+            snippet = text[:200].replace("\n", " ")
+            logger.info(f"     Text: {snippet}...")
+
+
+def cmd_info(args):
+    """Show collection info."""
+    from qdrant_client import QdrantClient
+
+    load_dotenv()
+
+    qdrant_url = (
+        os.getenv("SIGIR_QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("QDRANT_URL")
+    )
+    qdrant_api_key = (
+        os.getenv("SIGIR_QDRANT_KEY")
+        or os.getenv("SIGIR_QDRANT_API_KEY")
+        or os.getenv("DEST_QDRANT_API_KEY")
+        or os.getenv("QDRANT_API_KEY")
+    )
+
+    if not qdrant_url:
+        logger.error("❌ QDRANT_URL not set")
+        sys.exit(1)
+
+    grpc_port = 6334 if args.prefer_grpc and urlparse(qdrant_url).port == 6333 else None
+    client = QdrantClient(
+        url=qdrant_url,
+        api_key=qdrant_api_key,
+        prefer_grpc=args.prefer_grpc,
+        grpc_port=grpc_port,
+        check_compatibility=False,
+    )
+
+    try:
+        info = client.get_collection(args.collection)
+
+        status = info.status
+        if hasattr(status, "value"):
+            status = status.value
+
+        indexed_count = getattr(info, "indexed_vectors_count", 0) or 0
+        if isinstance(indexed_count, dict):
+            indexed_count = sum(indexed_count.values())
+
+        logger.info(f"📊 Collection: {args.collection}")
+        logger.info(f"   Status: {status}")
+        logger.info(f"   Points: {info.points_count}")
+        logger.info(f"   Indexed vectors: {indexed_count}")
+
+        # Show vector config
+        if hasattr(info, "config") and hasattr(info.config, "params"):
+            vectors = getattr(info.config.params, "vectors", {})
+            if vectors:
+                logger.info(f"   Vectors: {list(vectors.keys())}")
+
+    except Exception as e:
+        logger.error(f"❌ Could not get collection info: {e}")
+        sys.exit(1)
+
+
+def main():
+    """Main CLI entry point."""
+    parser = argparse.ArgumentParser(
+        prog="visual-rag",
+        description="Visual RAG Toolkit - Visual document retrieval with ColPali",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Process PDFs (like process_pdfs_saliency_v2.py)
+  visual-rag process --reports-dir ./pdfs --metadata-file metadata.json
+
+  # Process without Cloudinary
+  visual-rag process --reports-dir ./pdfs --no-cloudinary
+
+  # Search
+  visual-rag search --query "budget allocation" --collection my_docs
+
+  # Search with filters
+  visual-rag search --query "budget" --year 2023 --source "Local Government"
+
+  # Show collection info
+  visual-rag info --collection my_docs
+        """,
+    )
+    parser.add_argument("--debug", action="store_true", help="Enable debug logging")
+
+    subparsers = parser.add_subparsers(dest="command", help="Command")
+
+    # =========================================================================
+    # PROCESS command
+    # =========================================================================
+    process_parser = subparsers.add_parser(
+        "process",
+        help="Process PDFs: embed, upload to Cloudinary, index in Qdrant",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    process_parser.add_argument(
+        "--reports-dir", type=str, required=True, help="Directory containing PDF files"
+    )
+    process_parser.add_argument(
+        "--metadata-file",
+        type=str,
+        help="JSON file with filename → metadata mapping (like filename_metadata.json)",
+    )
+    process_parser.add_argument(
+        "--collection", type=str, default="visual_documents", help="Qdrant collection name"
+    )
+    process_parser.add_argument(
+        "--model",
+        type=str,
+        default="vidore/colSmol-500M",
+        help="Model name (vidore/colSmol-500M, vidore/colpali-v1.3, etc.)",
+    )
+    process_parser.add_argument("--batch-size", type=int, default=8, help="Embedding batch size")
+    process_parser.add_argument("--config", type=str, help="Path to config.yaml file")
+    process_parser.add_argument(
+        "--no-cloudinary", action="store_true", help="Skip Cloudinary uploads"
+    )
+    process_parser.add_argument(
+        "--crop-empty",
+        action="store_true",
+        help="Crop empty whitespace from page images before embedding (default: off).",
+    )
+    process_parser.add_argument(
+        "--crop-empty-percentage-to-remove",
+        type=float,
+        default=0.9,
+        help="Kept for traceability; currently does not affect cropping behavior (default: 0.9).",
+    )
+    process_parser.add_argument(
+        "--crop-empty-remove-page-number",
+        action="store_true",
+        help="If set, attempts to crop away the bottom region that contains sparse page numbers (default: off).",
+    )
+    process_parser.add_argument(
+        "--no-skip-existing",
+        action="store_true",
+        help="Process all pages even if they exist in Qdrant",
+    )
+    process_parser.add_argument(
+        "--force-recreate", action="store_true", help="Delete and recreate collection"
+    )
+    process_parser.add_argument(
+        "--dry-run", action="store_true", help="Show what would be processed without doing it"
+    )
+    process_parser.add_argument(
+        "--strategy",
+        type=str,
+        default="pooling",
+        choices=["pooling", "standard", "all"],
+        help="Embedding strategy: 'pooling' (NOVEL), 'standard' (BASELINE), "
+        "'all' (embed once, store BOTH for comparison)",
+    )
+    process_parser.add_argument(
+        "--torch-dtype",
+        type=str,
+        default="auto",
+        choices=["auto", "float32", "float16", "bfloat16"],
+        help="Torch dtype for model weights (default: auto; CUDA->bfloat16, else float32).",
+    )
+    process_parser.add_argument(
+        "--qdrant-vector-dtype",
+        type=str,
+        default="float16",
+        choices=["float16", "float32"],
+        help="Datatype for vectors stored in Qdrant (default: float16).",
+    )
+    process_parser.add_argument(
+        "--processor-speed",
+        type=str,
+        default="fast",
+        choices=["fast", "slow", "auto"],
+        help="Processor implementation: fast (default, with fallback to slow), slow, or auto.",
+    )
+    process_grpc_group = process_parser.add_mutually_exclusive_group()
+    process_grpc_group.add_argument(
+        "--prefer-grpc",
+        dest="prefer_grpc",
+        action="store_true",
+        default=True,
+        help="Use gRPC for Qdrant client (recommended).",
+    )
+    process_grpc_group.add_argument(
+        "--no-prefer-grpc",
+        dest="prefer_grpc",
+        action="store_false",
+        help="Disable gRPC for Qdrant client.",
+    )
+    process_parser.set_defaults(func=cmd_process)
+
+    # =========================================================================
+    # SEARCH command
+    # =========================================================================
+    search_parser = subparsers.add_parser(
+        "search",
+        help="Search documents",
+    )
+    search_parser.add_argument("--query", type=str, required=True, help="Search query")
+    search_parser.add_argument(
+        "--collection", type=str, default="visual_documents", help="Qdrant collection name"
+    )
+    search_parser.add_argument(
+        "--model", type=str, default="vidore/colSmol-500M", help="Model name"
+    )
+    search_parser.add_argument(
+        "--processor-speed",
+        type=str,
+        default="fast",
+        choices=["fast", "slow", "auto"],
+        help="Processor implementation: fast (default, with fallback to slow), slow, or auto.",
+    )
+    search_parser.add_argument("--top-k", type=int, default=10, help="Number of results")
+    search_parser.add_argument(
+        "--strategy",
+        type=str,
+        default="single_full",
+        choices=["single_full", "single_tiles", "single_global", "two_stage"],
+        help="Search strategy",
+    )
+    search_parser.add_argument(
+        "--prefetch-k", type=int, default=200, help="Prefetch candidates for two-stage retrieval"
+    )
+    search_parser.add_argument(
+        "--stage1-mode",
+        type=str,
+        default="pooled_query_vs_tiles",
+        choices=["pooled_query_vs_tiles", "tokens_vs_tiles", "pooled_query_vs_global"],
+        help="Stage 1 mode for two-stage retrieval",
+    )
+    search_parser.add_argument("--year", type=int, help="Filter by year")
+    search_parser.add_argument("--source", type=str, help="Filter by source")
+    search_parser.add_argument("--district", type=str, help="Filter by district")
+    search_parser.add_argument(
+        "--show-text", action="store_true", help="Show text snippets in results"
+    )
+    search_grpc_group = search_parser.add_mutually_exclusive_group()
+    search_grpc_group.add_argument(
+        "--prefer-grpc",
+        dest="prefer_grpc",
+        action="store_true",
+        default=True,
+        help="Use gRPC for Qdrant client (recommended).",
+    )
+    search_grpc_group.add_argument(
+        "--no-prefer-grpc",
+        dest="prefer_grpc",
+        action="store_false",
+        help="Disable gRPC for Qdrant client.",
+    )
+    search_parser.set_defaults(func=cmd_search)
+
+    # =========================================================================
+    # INFO command
+    # =========================================================================
+    info_parser = subparsers.add_parser(
+        "info",
+        help="Show collection info",
+    )
+    info_parser.add_argument(
+        "--collection", type=str, default="visual_documents", help="Qdrant collection name"
+    )
+    info_grpc_group = info_parser.add_mutually_exclusive_group()
+    info_grpc_group.add_argument(
+        "--prefer-grpc",
+        dest="prefer_grpc",
+        action="store_true",
+        default=True,
+        help="Use gRPC for Qdrant client (recommended).",
+    )
+    info_grpc_group.add_argument(
+        "--no-prefer-grpc",
+        dest="prefer_grpc",
+        action="store_false",
+        help="Disable gRPC for Qdrant client.",
+    )
+    info_parser.set_defaults(func=cmd_info)
+
+    # Parse and execute
+    args = parser.parse_args()
+
+    setup_logging(args.debug)
+
+    if not args.command:
+        parser.print_help()
+        sys.exit(0)
+
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()