diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 669b2fa..e12db89 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -41,12 +41,13 @@ jobs: - name: Install Python dependencies run: | python -m pip install --upgrade pip - pip install maturin numpy pytest pillow pyyaml requests fastapi uvicorn + pip install torch --index-url https://download.pytorch.org/whl/cpu + pip install maturin numpy pytest httpx pillow pyyaml requests fastapi uvicorn # 7. Build and install BBQ package via pip - name: Build and install BBQ (Release) run: | - pip install . + pip install -e . # 8. Check BBQ version and CLI - name: Verify BBQ version and CLI diff --git a/Cargo.toml b/Cargo.toml index c3e453b..099c4c9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ description = "This packages is the simd implementations of max sim function for license = "MIT" # readme = "README.md" keywords = ["maxsimd"] +build = "build.rs" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [lib] diff --git a/README.md b/README.md index 92fd0bb..9563491 100644 --- a/README.md +++ b/README.md @@ -5,12 +5,12 @@ [![License](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE) [![Python](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12%20%7C%203.13-0284c7)](https://www.python.org/) [![Rust](https://img.shields.io/badge/rust-1.75%2B-ea580c.svg)](https://www.rust-lang.org/) -[![CPU]()](#performance--benchmarks) -[![Model](https://img.shields.io/badge/Model-ColPali%20%2F%20SmolVLM-7c3aed.svg)](#key-features) -[![FFI](https://img.shields.io/badge/FFI-PyO3%200.29-e11d48.svg)](https://pyo3.rs/) -[![Repo Size](https://img.shields.io/badge/repo%20size-~3.2%20MB-0ea5e9.svg)](#) -[![Environment](https://img.shields.io/badge/environment-~2.5--3.0%20GB-8b5cf6.svg)](#) +[![SIMD](https://img.shields.io/badge/SIMD-AVX2%20%2B%20FMA-16a34a.svg)](#performance--benchmarks) +[![VLM](https://img.shields.io/badge/VLM-ColPali%20%2F%20colSmol-7c3aed.svg)](#quickstart--usage) +[![Scoring](https://img.shields.io/badge/Scoring-MaxSim%20Multi--Vector-0ea5e9.svg)](#performance--benchmarks) +[![FFI](https://img.shields.io/badge/FFI-PyO3%20Zero--Copy-e11d48.svg)](https://pyo3.rs/) [![Throughput](https://img.shields.io/badge/throughput->%20150k%20pages/s-f97316.svg)](#performance--benchmarks) +[![Multimodal RAG](https://img.shields.io/badge/LLM-Google%20Gemini-4285f4.svg)](#quickstart--usage) [![About](https://img.shields.io/badge/Lineage-ColPali%20%26%20maxsim--cpu-f59e0b.svg)](ABOUT.md) @@ -21,62 +21,32 @@ It provides a complete end-to-end local RAG pipeline: an automated PDF directory --- -## Key Features - -- **Fused AVX2/FMA SIMD MaxSim Engine**: 4-way query unrolled SIMD dot-product pipeline with in-register maximum tracking written in Rust (`maxsimd`), eliminating intermediate similarity matrix allocation and cutting memory load traffic by 27x. -- **Adaptive Multi-Threaded Parallelism**: Dynamic execution routing that runs small batches sequentially on the main thread to eliminate work-stealing overhead, and switches automatically to Rayon chunked work pools for large multi-page collections. -- **Unified Zero-Copy FFI**: A single unified Python entry point (`maxsimd.maxsim`) accepting contiguous NumPy arrays and PyTorch tensors directly with zero memory copying or overhead. -- **Vision-Language Indexing Server**: Automated folder monitoring (`data/watch/`), background PDF page rasterization at 150 DPI via PyMuPDF, and persistent SQLite metadata and multi-vector embedding storage. -- **Client-Side Gemini Multimodal RAG**: Seamless grounded answer synthesis with Google Gemini (`gemini-2.5-flash` / `gemini-2.0-flash`) over top retrieved visual pages, with automatic fallback to page viewing when offline or without an API key. -- **Production Benchmarking Suite**: Comprehensive micro- and macro-benchmark suites with CLI controls (`--suite maxsimd` for throughput scaling, `--suite vidore` evaluating retrieval quality, QPS, and numerical parity against PyTorch across 10 official ViDoRe datasets). - ---- - ## Performance & Benchmarks -> For in-depth benchmark comparisons, multi-baseline analysis, and full ViDoRe multimodal retrieval evaluations, see [**MORE INFO ON THE BENCHMARK README**](benchmarks/README.md). +BBQ-RAG's fused AVX2/FMA Rust SIMD kernel evaluates over 150,000 document pages per second on standard multi-core CPUs with zero intermediate heap allocations. For detailed methodology, multi-baseline comparisons, and ViDoRe multimodal retrieval evaluations, see the [benchmark report](benchmarks/README.md). ![MaxSim SIMD Throughput Scaling](assets/b1.png) ## Installation -### Prerequisites - -- Python 3.10 or higher -- Rust 1.75 or higher (with `cargo`) -- x86_64 CPU supporting AVX2 and FMA instructions +> [!IMPORTANT] +> **Hardware Support: x86_64 with FMA Only** +> BBQ-RAG's native SIMD compute engine (`maxsimd`) is strictly optimized for **x86_64 (64-bit AMD / Intel)** CPUs with **FMA (Fused Multiply-Add)** and **AVX2** vector instruction sets. +> - **Supported**: 64-bit x86 CPUs with FMA (Intel Haswell / Core 4th Gen+ and AMD Piledriver / Zen+). +> - **Unsupported**: ARM architectures (e.g., Apple Silicon M1/M2/M3/M4, AWS Graviton, Raspberry Pi) and legacy x86 CPUs lacking FMA support. +> +> Running `pip install .` automatically triggers an architecture and CPU instruction check. If your machine does not meet the hardware requirements, the installation will halt and inform you that BBQ-RAG cannot be run on your machine. -### Step 1: Clone Repository & Create Virtual Environment +Install BBQ-RAG and compile the native Rust SIMD extension in one step: ```bash git clone https://github.com/cmd-HMN/bbq-rag.git cd bbq-rag -python3 -m venv .venv -source .venv/bin/activate -``` - -### Step 2: Install Dependencies - -```bash -pip install --upgrade pip -pip install -r requirements.txt +pip install . ``` -### Step 3: Compile Rust Extension - -Build the high-performance release binary using `maturin`: - -```bash -maturin develop --release -``` - -Or build a redistributable wheel: - -```bash -maturin build --release -o dist/ -pip install dist/bbq_rag-*.whl -``` +> **Requirements**: Python 3.10+, Rust 1.75+ (with `cargo`), and an x86_64 CPU supporting AVX2 and FMA instructions. --- @@ -87,19 +57,43 @@ pip install dist/bbq_rag-*.whl Edit `config.yaml` to set your model IDs, watch directory, and Gemini preferences: ```yaml +# only inference supported base_model_id: "HuggingFaceTB/SmolVLM-256M-Instruct" +# only support this model lora_adapter_id: "vidore/colSmol-256M" + +# change the embedding_dim if you want embedding_dim: 128 + +# can be auto selected device: "auto" + +# data type torch_dtype: "bfloat16" +# masking for images in the model +mask_non_image_embeddings: false + +# prompt for vision model +visual_prompt_command: "Describe this image." + +# watching folder watch_folder_path: "data/watch" + +# pdf image dpi rendered for embeddings & viewing pdf_render_dpi: 150 -# Optional Google Gemini multimodal RAG settings +# Output directory for saved retrieved page images +images_output_dir: "data/rr" + +# Google Gemini multimodal RAG settings +# If gemini_api_key is omitted/empty, GEMINI_API_KEY / GOOGLE_API_KEY env vars are used. gemini_api_key: "" gemini_model: "gemini-3.6-flash" -rag_top_k: 3 +rag_top_k: 5 + +# Quantization precision mode: "f32" (default) or "qi8" (quantized INT8) +quantization: "f32" ``` ### 2. Start the Indexing Server @@ -119,26 +113,38 @@ python -m bbq.src.main server --config config.yaml Search indexed documents from the command line: ```bash -# Fetch top 10 matching pages (prints ASCII BBQ banner on run) +# Fetch top matching pages (prints ASCII BBQ banner on run) bbq client "What was the operating margin in Q3?" # Or: python -m bbq.src.main query "What was the operating margin in Q3?" # Fetch without ASCII logo banner bbq client --without-logo "What was the operating margin in Q3?" +# Save retrieved page images to disk (default: data/rr/) in background +bbq client "What is UX design?" --save-images +# Or using shorthand flags with custom output directory: +bbq client "What is UX design?" -i -o data/rr/ + # Query with Gemini Multimodal LLM (displays live cooking spinner and multimodal answer) export GEMINI_API_KEY="your-gemini-api-key" bbq client "Summarize the revenue growth" --use-llm # Continuous interactive query prompt loop (--infinite / -inf) bbq client --infinite -# Or with LLM cooking enabled: -bbq client --infinite --use-llm +# Or with image saving and LLM cooking enabled: +bbq client --infinite -i --use-llm # Run directly via the client module python -m bbq.src.client --infinite --use-llm ``` +#### Interactive Document Page Opener +When querying in an interactive terminal, BBQ displays an interactive page picker: +- **`↑` / `↓`** or **`j` / `k`**: Navigate between top retrieved PDF document pages. +- **`Enter`**: Open the selected document at that exact page in your system PDF viewer (`evince`, `okular`, or `xdg-open`). +- **`Esc` / `q`** (or selecting **`Done`**): Proceed to display results and return to prompt. +- Pass **`--without-opener`** (or **`--no-opener`**) to bypass the interactive picker and immediately print matching pages. + In `--infinite` mode, an interactive prompt (`bbq[query] >> `) appears for continuous querying: - Type any query to search. - Type `help` or `?` to show the commands menu. @@ -149,24 +155,41 @@ If no Gemini API key is provided or the API is unavailable, the client automatic ### 4. Python API Usage -#### Client Query & RAG +#### Client Query, Image Saving & RAG ```python from bbq.src.client import BBQClient -from bbq.src.config import load_configuration_from_yaml_file +from bbq.src.config import Config -config = load_configuration_from_yaml_file("config.yaml") +config = Config.from_yaml("config.yaml") client = BBQClient(server_url="http://localhost:8000", config=config) +# 1. Standard retrieval query +results = client.query(query_text="Explain the cash flow breakdown in the report", top_k=5) + +# 2. Save retrieved page images to disk (default configured dir: data/rr) +saved_paths = client.save_page_images(results) +print(f"Saved {len(saved_paths)} images to data/rr") + +# Or save asynchronously in a background daemon thread: +# thread = client.save_page_images_threaded( +# results, +# on_progress=lambda curr, total, path: print(f"Saved {curr}/{total}: {path}"), +# on_complete=lambda paths: print(f"All {len(paths)} images saved!") +# ) + +# 3. Multimodal answer generation with Gemini response = client.query_and_answer( query_text="Explain the cash flow breakdown in the report", top_k=3, + use_llm=True, + save_images=True, ) if response["answer"]: print("Gemini Multimodal Answer:\n", response["answer"]) else: - print("Retrieved Book Pages (Fallback):") + print("Retrieved Document Pages (Fallback):") for source in response["sources"]: print(f"File: {source['file_path']} | Page: {source['page_number']} | Score: {source['score']:.4f}") ``` @@ -198,10 +221,13 @@ print("Top 5 Document Scores:", scores[:5]) Run the full integration test suite and Rust unit tests: ```bash -# Run Python integration and regression tests (7 test suites) -python3 -m pytest +# Run Python unit and regression tests +pytest -# Run Rust unit tests (47 tests for BLAS and AVX2 kernels) +# Run tests in CI mode (headless, skips GUI viewer tests) +CI=true pytest + +# Run Rust unit tests (BLAS and AVX2 kernels) cargo test # Run scaling and ViDoRe benchmark suite @@ -210,13 +236,6 @@ python3 benchmarks --- -## Roadmap & TODO -- [ ] **Audit Boilerplate & Fix Errors for Scalability**: Review and harden all boilerplate code, fix edge-case runtime errors, eliminate redundant allocations, and optimize the codebase for production-grade throughput and scalability. -- [ ] **Eliminate Circular Dependencies**: Audit and refactor inter-module imports across `client`, `server`, `storage`, and `config`. -- [ ] **Origins & Attribution Reference**: See [ABOUT.md](ABOUT.md) for full project lineage, paper citations ([ColPali arXiv:2407.01449](https://arxiv.org/abs/2407.01449)), `maxsim-cpu` references, and reserved rights notices. - ---- - ## Acknowledgments & Lineage BBQ-RAG builds upon the foundational research and open-source contributions of the visual document retrieval community: diff --git a/bbq/maxsimd/__init__.py b/bbq/maxsimd/__init__.py index b51762c..83d080f 100644 --- a/bbq/maxsimd/__init__.py +++ b/bbq/maxsimd/__init__.py @@ -6,6 +6,13 @@ except ImportError: torch = None # type: ignore +try: + from bbq.src.common.hardware import verify_hardware_or_exit + + verify_hardware_or_exit() +except ImportError: + pass + from .maxsimd import ( maxsim as _raw_maxsim, __version__, diff --git a/bbq/src/client/client.py b/bbq/src/client/client.py index 741df2f..f0d27ea 100644 --- a/bbq/src/client/client.py +++ b/bbq/src/client/client.py @@ -7,7 +7,8 @@ import os import io import logging -from typing import TYPE_CHECKING, Any, Dict, List, Optional +import threading +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional import requests from PIL import Image @@ -70,13 +71,22 @@ def get_page_image( self, file_path: str, page_number: int = 1, - dpi: int = 150, + dpi: Optional[int] = None, save_path: Optional[str] = None, ) -> Image.Image: """Fetches rendered PNG page image from the server and optionally saves it.""" + render_dpi = ( + dpi + if dpi is not None + else (getattr(self.config, "pdf_render_dpi", 150) if self.config else 150) + ) resp = requests.get( f"{self.server_url}/page_image", - params={"file_path": file_path, "page_number": page_number, "dpi": dpi}, + params={ + "file_path": file_path, + "page_number": page_number, + "dpi": render_dpi, + }, timeout=15, ) resp.raise_for_status() @@ -86,41 +96,52 @@ def get_page_image( img.save(save_path, format="PNG") return img - def generate_answer_from_results( + def fetch_page_images( self, - query_text: str, results: List[Dict[str, Any]], - gemini_api_key: Optional[str] = None, - gemini_model: Optional[str] = None, save_images: bool = False, images_output_dir: Optional[str] = None, - ) -> Dict[str, Any]: + ) -> List[Image.Image]: """ - Fetches page images for retrieved results and generates a grounded multimodal answer with Gemini. + Fetches page images for retrieved results and optionally saves them to disk. + Default output directory is data/rr. + Updates each result dict in results with 'image_available' and 'saved_image_path'. """ + out_dir = images_output_dir + if not out_dir and self.config: + out_dir = getattr(self.config, "images_output_dir", None) + out_dir = out_dir or "data/rr" + + if save_images: + os.makedirs(out_dir, exist_ok=True) + page_images: List[Image.Image] = [] for i, res in enumerate(results, 1): + pnum = res.get("page_number", 1) save_path = ( - os.path.join( - images_output_dir or ".", - f"retrieved_rank_{i}_page_{res['page_number']}.png", - ) - if save_images - else None + os.path.join(out_dir, f"rr_{i}_{pnum}.png") if save_images else None ) img = None try: - img = self.get_page_image( - res["file_path"], res["page_number"], save_path=save_path + img = self.get_page_image(res["file_path"], pnum, save_path=save_path) + except Exception as e: + logger.debug( + f"Remote page image request failed for {res.get('file_path')}: {e}" ) - except Exception: if os.path.exists(res.get("file_path", "")): try: + dpi = ( + getattr(self.config, "pdf_render_dpi", 150) + if self.config + else 150 + ) img = get_local_pdf_page_image( - res["file_path"], res["page_number"], save_path=save_path + res["file_path"], pnum, dpi=dpi, save_path=save_path + ) + except Exception as le: + logger.debug( + f"Local page image extraction failed for {res.get('file_path')}: {le}" ) - except Exception: - pass if img: page_images.append(img) res["image_available"] = True @@ -128,6 +149,124 @@ def generate_answer_from_results( res["saved_image_path"] = save_path else: res["image_available"] = False + return page_images + + def save_page_images( + self, + results: List[Dict[str, Any]], + images_output_dir: Optional[str] = None, + ) -> List[str]: + """ + Fetches and saves rendered page images for retrieved results to disk (default: data/rr). + Returns a list of saved file paths. + """ + self.fetch_page_images( + results=results, + save_images=True, + images_output_dir=images_output_dir, + ) + return [ + res["saved_image_path"] for res in results if res.get("saved_image_path") + ] + + def save_page_images_threaded( + self, + results: List[Dict[str, Any]], + images_output_dir: Optional[str] = None, + on_progress: Optional[Callable[[int, int, str], None]] = None, + on_complete: Optional[Callable[[List[str]], None]] = None, + ) -> threading.Thread: + """ + Launches a background daemon thread to fetch and save page images asynchronously. + Calls on_progress(current_idx, total_count, saved_path) as each image is saved. + Calls on_complete(saved_paths) when all images have finished saving. + Returns the spawned background threading.Thread instance. + """ + + def _worker() -> None: + out_dir = images_output_dir + if not out_dir and self.config: + out_dir = getattr(self.config, "images_output_dir", None) + out_dir = out_dir or "data/rr" + os.makedirs(out_dir, exist_ok=True) + + saved_paths: List[str] = [] + total = len(results) + + for i, res in enumerate(results, 1): + pnum = res.get("page_number", 1) + save_path = os.path.join(out_dir, f"rr_{i}_{pnum}.png") + img = None + try: + img = self.get_page_image( + res["file_path"], pnum, save_path=save_path + ) + except Exception as e: + logger.debug( + f"Remote page image request failed for {res.get('file_path')}: {e}" + ) + if os.path.exists(res.get("file_path", "")): + try: + dpi = ( + getattr(self.config, "pdf_render_dpi", 150) + if self.config + else 150 + ) + img = get_local_pdf_page_image( + res["file_path"], pnum, dpi=dpi, save_path=save_path + ) + except Exception as le: + logger.debug( + f"Local page image extraction failed for {res.get('file_path')}: {le}" + ) + if img: + res["image_available"] = True + res["saved_image_path"] = save_path + saved_paths.append(save_path) + if on_progress: + try: + on_progress(i, total, save_path) + except Exception: + pass + else: + res["image_available"] = False + if on_progress: + try: + on_progress(i, total, "") + except Exception: + pass + + if on_complete: + try: + on_complete(saved_paths) + except Exception: + pass + + thread = threading.Thread( + target=_worker, + name="BBQImageSaverThread", + daemon=True, + ) + thread.start() + return thread + + def generate_answer_from_results( + self, + query_text: str, + results: List[Dict[str, Any]], + gemini_api_key: Optional[str] = None, + gemini_model: Optional[str] = None, + save_images: bool = False, + images_output_dir: Optional[str] = None, + ) -> Dict[str, Any]: + """ + Fetches page images for retrieved results and generates a grounded multimodal answer with Gemini. + """ + page_images = self.fetch_page_images( + results=results, + save_images=save_images, + images_output_dir=images_output_dir, + ) if self.config: gemini_api_key = gemini_api_key or getattr( @@ -191,6 +330,7 @@ def query_and_answer( """ Retrieves top_k pages. If use_llm is True, passes retrieved pages to Gemini LLM. Otherwise, returns retrieved matches directly without calling LLM. + If save_images is True, saves page images to disk (default: data/rr). """ suppress_logo = self.without_logo if without_logo is None else without_logo if not suppress_logo: @@ -208,6 +348,10 @@ def query_and_answer( } if not use_llm: + if save_images: + self.save_page_images( + results=results, images_output_dir=images_output_dir + ) return { "query": query_text, "answer": None, diff --git a/bbq/src/common/hardware.py b/bbq/src/common/hardware.py new file mode 100644 index 0000000..c44f67d --- /dev/null +++ b/bbq/src/common/hardware.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import platform +import sys +from typing import Tuple + + +def check_cpu_support() -> Tuple[bool, str]: + """ + Checks whether the host machine supports the necessary architecture + (x86_64) and CPU SIMD instructions (FMA). + + Returns: + (is_supported, error_message): Tuple where is_supported is True if compatible, + or False with a descriptive error message explaining why the machine is unsupported. + """ + machine = platform.machine().lower() + if machine not in ("x86_64", "amd64", "x64"): + return ( + False, + f"Unsupported CPU architecture: '{machine}'. " + f"BBQ-RAG is only supported on x86_64 systems with FMA instruction support.", + ) + + system = sys.platform.lower() + has_fma = False + + if system.startswith("linux"): + try: + with open("/proc/cpuinfo", "r", encoding="utf-8") as f: + for line in f: + if line.startswith("flags") or line.startswith("Features"): + flags = set(line.split(":", 1)[1].strip().lower().split()) + if "fma" in flags: + has_fma = True + break + except Exception: + # Fallback if /proc/cpuinfo cannot be read + has_fma = True + elif system == "darwin": + import subprocess + + try: + out = subprocess.check_output( + ["sysctl", "-n", "machdep.cpu.features", "machdep.cpu.leaf7_features"], + text=True, + stderr=subprocess.DEVNULL, + ) + if "FMA" in out.upper(): + has_fma = True + except Exception: + has_fma = True + elif system.startswith("win"): + has_fma = True + else: + has_fma = True + + if not has_fma: + return ( + False, + "Your x86_64 processor does not support FMA (Fused Multiply-Add) instructions. " + "BBQ-RAG requires FMA-capable processors (Intel Haswell+, AMD Piledriver+).", + ) + + return True, "" + + +def verify_hardware_or_exit(exit_code: int = 1) -> None: + """ + Verifies that the current machine meets CPU and instruction set requirements. + If incompatible, prints a clear error message to stderr and terminates. + """ + is_supported, reason = check_cpu_support() + if not is_supported: + sys.stderr.write( + f"ERROR: BBQ-RAG Hardware Compatibility Check Failed\n" + f"{reason}\n" + f"BBQ-RAG cannot be installed or run on this machine.\n" + ) + sys.stderr.flush() + sys.exit(exit_code) diff --git a/bbq/src/config/__init__.py b/bbq/src/config/__init__.py index 96269da..786c82d 100644 --- a/bbq/src/config/__init__.py +++ b/bbq/src/config/__init__.py @@ -1,4 +1,3 @@ from .config import Config __all__ = ["Config"] - diff --git a/bbq/src/config/config.py b/bbq/src/config/config.py index 0ac9e5c..c2608ab 100644 --- a/bbq/src/config/config.py +++ b/bbq/src/config/config.py @@ -27,6 +27,7 @@ class Config: watch_folder_path: str = "data/watch" embeddings_output_path: Optional[str] = None sqlite_db_path: Optional[str] = None + images_output_dir: str = "data/rr" pdf_render_dpi: int = 150 gemini_api_key: Optional[str] = None gemini_model: str = "gemini-3.6-flash" diff --git a/bbq/src/main.py b/bbq/src/main.py index c961af3..5e307a7 100644 --- a/bbq/src/main.py +++ b/bbq/src/main.py @@ -1,24 +1,26 @@ import sys import argparse import logging -from bbq.src.common.version import __version__, _read_project_version +from bbq.src.common.version import __version__ def run_server_command(args: argparse.Namespace) -> None: """Launches the persistent document-indexing server.""" - from bbq.src.config import Config - from bbq.src.server import start_document_indexing_server from bbq.src.terminal.tui import print_bbq if not getattr(args, "without_logo", False): print_bbq(tag="SERVER", server=getattr(args, "port", 8000)) + from bbq.src.config import Config + from bbq.src.server import start_document_indexing_server + config = Config.from_yaml(config_filepath=args.config) start_document_indexing_server( config=config, config_filepath=args.config, host=args.host, port=args.port, + without_logo=True, ) @@ -118,12 +120,67 @@ def run_client_query_command(args: argparse.Namespace) -> None: if not results: console.print("[yellow]No matching PDF pages found.[/yellow]") else: + save_images_requested = getattr(args, "save_images", False) or bool( + getattr(args, "images_output_dir", None) + ) + images_out_dir = getattr(args, "images_output_dir", None) or getattr( + config, "images_output_dir", "data/rr" + ) + + save_status = None + save_thread = None + if save_images_requested and not use_llm: + save_status = { + "active": True, + "completed": False, + "saved_count": 0, + "total": len(results), + "output_dir": images_out_dir, + "saved_paths": [], + } + + def _on_progress(idx: int, total: int, path: str) -> None: + if path and save_status is not None: + save_status["saved_count"] += 1 + save_status["saved_paths"].append(path) + + def _on_complete(paths: list) -> None: + if save_status is not None: + save_status["completed"] = True + save_status["active"] = False + + save_thread = client.save_page_images_threaded( + results=results, + images_output_dir=images_out_dir, + on_progress=_on_progress, + on_complete=_on_complete, + ) + # 2. Interactive arrow-key page picker (Option C) unless --without-opener is passed if not getattr(args, "without_opener", False) and sys.stdin.isatty(): interactive_page_picker(results, console=console) else: render_query_results_rich(results, console=console) + if save_images_requested and not use_llm: + if save_thread and save_thread.is_alive(): + if not is_infinite: + save_thread.join(timeout=5.0) + + saved_count = ( + save_status["saved_count"] + if save_status + else len([r for r in results if r.get("saved_image_path")]) + ) + if saved_count > 0: + console.print( + f"[bold green]✔ Saved {saved_count} page image(s) to {images_out_dir}[/bold green]\n" + ) + elif save_status and not save_status.get("active"): + console.print( + f"[bold yellow]⚠ Could not save page images to {images_out_dir}[/bold yellow]\n" + ) + # 3. If --use-llm is specified, pass top-k to LLM with cooking spinner if use_llm: with console.status( @@ -135,9 +192,21 @@ def run_client_query_command(args: argparse.Namespace) -> None: results=results, gemini_api_key=gemini_key, gemini_model=gemini_model, - save_images=args.save_images, + save_images=save_images_requested, + images_output_dir=images_out_dir, ) + if save_images_requested: + saved_paths = [ + r["saved_image_path"] + for r in results + if r.get("saved_image_path") + ] + if saved_paths: + console.print( + f"[bold green]✔ Saved {len(saved_paths)} page image(s) to {images_out_dir}[/bold green]\n" + ) + answer = rag_response.get("answer") if answer: render_llm_answer_rich( @@ -247,7 +316,9 @@ def build_argument_parser() -> argparse.ArgumentParser: # --- Server Command --- server_parser = subparsers.add_parser( - "server", help="Start persistent document indexing HTTP server" + "server", + parents=[client_parent], + help="Start persistent document indexing HTTP server", ) server_parser.add_argument( "--config", @@ -333,7 +404,17 @@ def build_argument_parser() -> argparse.ArgumentParser: help="Gemini model name (default from config: gemini-3.6-flash)", ) query_parser.add_argument( - "--save-images", "-i", action="store_true", help="Save page images to disk" + "--save-images", + "-i", + action="store_true", + help="Save page images to disk (default output: data/rr)", + ) + query_parser.add_argument( + "--images-output-dir", + "-o", + type=str, + default=None, + help="Directory to save page images (default from config: data/rr)", ) query_parser.add_argument( "--verbose", @@ -427,7 +508,17 @@ def build_client_main_parser() -> argparse.ArgumentParser: "--gemini-model", type=str, default=None, help="Gemini model name" ) parser.add_argument( - "--save-images", "-i", action="store_true", help="Save page images to disk" + "--save-images", + "-i", + action="store_true", + help="Save page images to disk (default output: data/rr)", + ) + parser.add_argument( + "--images-output-dir", + "-o", + type=str, + default=None, + help="Directory to save page images (default from config: data/rr)", ) parser.add_argument( "--without-logo", @@ -451,6 +542,10 @@ def build_client_main_parser() -> argparse.ArgumentParser: def main() -> None: + from bbq.src.common.hardware import verify_hardware_or_exit + + verify_hardware_or_exit() + # Backward compatibility fallback: if run as 'python -m bbq.src.main [config.yaml]' without subcommands if ( len(sys.argv) > 1 diff --git a/bbq/src/server/__init__.py b/bbq/src/server/__init__.py index b877ee6..df70766 100644 --- a/bbq/src/server/__init__.py +++ b/bbq/src/server/__init__.py @@ -1,12 +1,30 @@ -from typing import Any +from __future__ import annotations + +import importlib +from typing import Any, TYPE_CHECKING + +if TYPE_CHECKING: + from bbq.src.server.app import ( + create_bbq_fastapi_app, + run_http_server_in_thread, + ) + from bbq.src.server.ingestion import ( + process_single_pdf_file_ingestion, + process_single_pdf_file_deletion, + scan_and_ingest_existing_pdf_folder, + sync_and_cleanup_deleted_files, + ) + from bbq.src.server.retrieval import query_indexed_documents + from bbq.src.server.server import ( + BBQServer, + start_document_indexing_server, + ) __all__ = [ "BBQServer", "start_document_indexing_server", - "execute_background_server_pipeline", "query_indexed_documents", "create_bbq_fastapi_app", - "create_colpali_fastapi_app", "run_http_server_in_thread", "process_single_pdf_file_ingestion", "process_single_pdf_file_deletion", @@ -14,48 +32,22 @@ "sync_and_cleanup_deleted_files", ] +_EXPORTS = { + "BBQServer": "bbq.src.server.server", + "start_document_indexing_server": "bbq.src.server.server", + "query_indexed_documents": "bbq.src.server.retrieval", + "create_bbq_fastapi_app": "bbq.src.server.app", + "run_http_server_in_thread": "bbq.src.server.app", + "process_single_pdf_file_ingestion": "bbq.src.server.ingestion", + "process_single_pdf_file_deletion": "bbq.src.server.ingestion", + "scan_and_ingest_existing_pdf_folder": "bbq.src.server.ingestion", + "sync_and_cleanup_deleted_files": "bbq.src.server.ingestion", +} + def __getattr__(name: str) -> Any: - if name in ("BBQServer", "start_document_indexing_server", "execute_background_server_pipeline"): - from bbq.src.server.server import ( - BBQServer, - start_document_indexing_server, - execute_background_server_pipeline, - ) - mapping = { - "BBQServer": BBQServer, - "start_document_indexing_server": start_document_indexing_server, - "execute_background_server_pipeline": execute_background_server_pipeline, - } - return mapping[name] - elif name == "query_indexed_documents": - from bbq.src.server.retrieval import query_indexed_documents - return query_indexed_documents - elif name in ("create_bbq_fastapi_app", "create_colpali_fastapi_app", "run_http_server_in_thread"): - from bbq.src.server.app import create_bbq_fastapi_app, create_colpali_fastapi_app, run_http_server_in_thread - mapping = { - "create_bbq_fastapi_app": create_bbq_fastapi_app, - "create_colpali_fastapi_app": create_colpali_fastapi_app, - "run_http_server_in_thread": run_http_server_in_thread, - } - return mapping[name] - elif name in ( - "process_single_pdf_file_ingestion", - "process_single_pdf_file_deletion", - "scan_and_ingest_existing_pdf_folder", - "sync_and_cleanup_deleted_files", - ): - from bbq.src.server.ingestion import ( - process_single_pdf_file_ingestion, - process_single_pdf_file_deletion, - scan_and_ingest_existing_pdf_folder, - sync_and_cleanup_deleted_files, - ) - mapping = { - "process_single_pdf_file_ingestion": process_single_pdf_file_ingestion, - "process_single_pdf_file_deletion": process_single_pdf_file_deletion, - "scan_and_ingest_existing_pdf_folder": scan_and_ingest_existing_pdf_folder, - "sync_and_cleanup_deleted_files": sync_and_cleanup_deleted_files, - } - return mapping[name] + module_path = _EXPORTS.get(name) + if module_path: + module = importlib.import_module(module_path) + return getattr(module, name) raise AttributeError(f"module '{__name__}' has no attribute '{name}'") diff --git a/bbq/src/server/app.py b/bbq/src/server/app.py index 75a71e8..bfb1c89 100644 --- a/bbq/src/server/app.py +++ b/bbq/src/server/app.py @@ -1,13 +1,15 @@ +from __future__ import annotations + import io import logging -import time import threading +import time from typing import Any, Dict, Optional import uvicorn from fastapi import FastAPI, HTTPException, Request, Response -from .retrieval import query_indexed_documents +from bbq.src.server.retrieval import query_indexed_documents from bbq.src.storage.sql import SqlliteDB from bbq.src.utils.pdf_utils import extract_single_pdf_page_image @@ -15,9 +17,10 @@ def create_bbq_fastapi_app( - engine: Optional[Any], - tracker: SqlliteDB, + engine: Optional[Any] = None, + tracker: Optional[SqlliteDB] = None, get_engine_callback: Optional[Any] = None, + config: Optional[Any] = None, ) -> FastAPI: """ Creates the FastAPI web server for query retrieval, image rendering, and server status. @@ -42,17 +45,26 @@ def get_status(request: Request): client_host = request.client.host if request.client else "unknown" logger.info(f"Status check requested by client [{client_host}]") active_engine = _get_active_engine() - records = tracker.fetch_all_records() + records = tracker.fetch_all_records() if tracker else [] done_count = sum(1 for r in records if r.get("status") == "done") engine_ready = active_engine is not None - base_model_id = active_engine.config.base_model_id if active_engine else "Loading..." - device = ( - str(next(active_engine.model.parameters()).device) - if active_engine and hasattr(active_engine, "model") and active_engine.model is not None - else "Loading..." - ) - watch_folder = active_engine.config.watch_folder_path if active_engine else "data/watch" + cfg = getattr(active_engine, "config", None) or config + base_model_id = cfg.base_model_id if cfg else "Loading..." + watch_folder = cfg.watch_folder_path if cfg else "data/watch" + quantization_mode = ( + getattr(cfg, "quantization", None) or "f32" + ).strip().lower() + + device = "Loading..." + if ( + active_engine + and hasattr(active_engine, "model") + and active_engine.model is not None + ): + param = next(active_engine.model.parameters(), None) + if param is not None: + device = str(param.device) return { "status": "online" if engine_ready else "loading_engine", @@ -60,6 +72,7 @@ def get_status(request: Request): "base_model_id": base_model_id, "device": device, "watch_folder": watch_folder, + "quantization": quantization_mode, "total_documents": len(records), "indexed_documents": done_count, } @@ -68,22 +81,23 @@ def get_status(request: Request): def get_documents(request: Request): client_host = request.client.host if request.client else "unknown" logger.info(f"Document list requested by client [{client_host}]") - return {"documents": tracker.fetch_all_records()} + return {"documents": tracker.fetch_all_records() if tracker else []} - @app.post("/query") - def post_query(body: Dict[str, Any], request: Request): - client_host = request.client.host if request.client else "unknown" - query_text = body.get("query", "") - top_k = body.get("top_k", 5) - - logger.info(f"Received POST /query request from client [{client_host}]: '{query_text}' (top_k={top_k})") + def _handle_query( + query_text: str, top_k: int, client_host: str + ) -> Dict[str, Any]: + query_text = (query_text or "").strip() if not query_text: - logger.warning(f"Rejected empty POST query from client [{client_host}]") - raise HTTPException(status_code=400, detail="Field 'query' cannot be empty.") + logger.warning(f"Rejected empty query from client [{client_host}]") + raise HTTPException( + status_code=400, detail="Query cannot be empty." + ) active_engine = _get_active_engine() if not active_engine: - logger.warning(f"Rejected query from client [{client_host}]: Engine is still loading weights.") + logger.warning( + f"Rejected query from client [{client_host}]: Engine is still loading weights." + ) raise HTTPException( status_code=503, detail="AI Model weights are currently initializing. Please try again in a few seconds.", @@ -98,38 +112,29 @@ def post_query(body: Dict[str, Any], request: Request): ) elapsed = time.time() - start_time logger.info( - f"Successfully served POST /query to [{client_host}] in {elapsed:.3f}s: returned {len(results)} match(es)" + f"Successfully served query to [{client_host}] in {elapsed:.3f}s: returned {len(results)} match(es)" ) return {"query": query_text, "top_k": top_k, "results": results} + @app.post("/query") + def post_query(body: Dict[str, Any], request: Request): + client_host = request.client.host if request.client else "unknown" + query_text = body.get("query", "") + top_k = body.get("top_k", 5) + logger.info( + f"Received POST /query request from client [{client_host}]: '{query_text}' (top_k={top_k})" + ) + return _handle_query( + query_text=query_text, top_k=top_k, client_host=client_host + ) + @app.get("/query") def get_query(q: str, request: Request, top_k: int = 5): client_host = request.client.host if request.client else "unknown" - logger.info(f"Received GET /query request from client [{client_host}]: '{q}' (top_k={top_k})") - if not q: - logger.warning(f"Rejected empty GET query from client [{client_host}]") - raise HTTPException(status_code=400, detail="Query parameter 'q' cannot be empty.") - - active_engine = _get_active_engine() - if not active_engine: - logger.warning(f"Rejected query from client [{client_host}]: Engine is still loading weights.") - raise HTTPException( - status_code=503, - detail="AI Model weights are currently initializing. Please try again in a few seconds.", - ) - - start_time = time.time() - results = query_indexed_documents( - query_text=q, - engine=active_engine, - tracker=tracker, - top_k=top_k, - ) - elapsed = time.time() - start_time logger.info( - f"Successfully served GET /query to [{client_host}] in {elapsed:.3f}s: returned {len(results)} match(es)" + f"Received GET /query request from client [{client_host}]: '{q}' (top_k={top_k})" ) - return {"query": q, "top_k": top_k, "results": results} + return _handle_query(query_text=q, top_k=top_k, client_host=client_host) @app.get("/page_image") def get_page_image( @@ -139,24 +144,28 @@ def get_page_image( dpi: int = 150, ): client_host = request.client.host if request.client else "unknown" - logger.info(f"Client [{client_host}] requested page image for '{file_path}' (page={page_number}, dpi={dpi})") + logger.info( + f"Client [{client_host}] requested page image for '{file_path}' (page={page_number}, dpi={dpi})" + ) try: - pil_img = extract_single_pdf_page_image(pdf_filepath=file_path, page_number=page_number, dpi=dpi) + pil_img = extract_single_pdf_page_image( + pdf_filepath=file_path, page_number=page_number, dpi=dpi + ) buf = io.BytesIO() pil_img.save(buf, format="PNG") - logger.info(f"Successfully served page image to client [{client_host}]") + logger.info( + f"Successfully served page image to client [{client_host}]" + ) return Response(content=buf.getvalue(), media_type="image/png") except Exception as exc: - logger.error(f"Failed to serve page image to client [{client_host}]: {exc}") + logger.error( + f"Failed to serve page image to client [{client_host}]: {exc}" + ) raise HTTPException(status_code=400, detail=str(exc)) return app -# Backward compatibility alias -create_colpali_fastapi_app = create_bbq_fastapi_app - - def run_http_server_in_thread( app: FastAPI, host: str = "0.0.0.0", diff --git a/bbq/src/server/ingestion.py b/bbq/src/server/ingestion.py index 0d02d6c..995c663 100644 --- a/bbq/src/server/ingestion.py +++ b/bbq/src/server/ingestion.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import gc import logging import os @@ -21,6 +23,7 @@ def process_single_pdf_file_ingestion( tracker: SqlliteDB, console: Optional[Console] = None, page_chunk_size: int = 8, + batch_size: int = 4, ) -> None: if not os.path.exists(pdf_filepath): logger.warning(f"File not found for processing: {pdf_filepath}") @@ -29,22 +32,29 @@ def process_single_pdf_file_ingestion( import numpy as np import torch + file_hash: Optional[str] = None try: - file_hash: str = compute_file_sha256_hash(pdf_filepath) - existing_record: Optional[Dict[str, Any]] = tracker.fetch_file_record_by_hash(file_hash) + file_hash = compute_file_sha256_hash(pdf_filepath) + existing_record: Optional[Dict[str, Any]] = tracker.fetch_file_record_by_hash( + file_hash + ) if existing_record and existing_record.get("status") == "done": - logger.info(f"Skipping already processed PDF (hash={file_hash[:10]}): {pdf_filepath}") + logger.info( + f"Skipping already processed PDF (hash={file_hash[:10]}): {pdf_filepath}" + ) return logger.info(f"Starting PDF ingestion (hash={file_hash[:10]}): {pdf_filepath}") - tracker.update_file_status_to_processing(file_hash=file_hash, file_path=pdf_filepath) + tracker.update_file_status_to_processing( + file_hash=file_hash, file_path=pdf_filepath + ) filename = os.path.basename(pdf_filepath) num_pages: int = get_pdf_total_pages(pdf_filepath) if num_pages == 0: - error_msg: str = "PDF contains zero pages." + error_msg = "PDF contains zero pages." logger.error(f"Failed PDF ingestion for {pdf_filepath}: {error_msg}") tracker.update_file_status_to_failed( file_hash=file_hash, @@ -57,7 +67,10 @@ def process_single_pdf_file_ingestion( f"PDF {filename} has {num_pages} total pages. Streaming ingestion in chunks of {page_chunk_size} pages..." ) status_ctx = ( - console.status(f"[bold yellow]Processing {filename} (0/{num_pages} pages)...[/bold yellow]", spinner="dots") + console.status( + f"[bold yellow]Processing {filename} (0/{num_pages} pages)...[/bold yellow]", + spinner="dots", + ) if console else None ) @@ -82,7 +95,9 @@ def process_single_pdf_file_ingestion( dpi=engine.config.pdf_render_dpi, ) - chunk_tensor: torch.Tensor = engine.encode_multimodal_document_images(chunk_images, batch_size=4) + chunk_tensor: torch.Tensor = engine.encode_multimodal_document_images( + chunk_images, batch_size=batch_size + ) all_page_embeddings.append(chunk_tensor.cpu().float().numpy()) del chunk_images @@ -91,34 +106,72 @@ def process_single_pdf_file_ingestion( embeddings_numpy: np.ndarray = np.concatenate(all_page_embeddings, axis=0) os.makedirs(engine.config.embeddings_output_path, exist_ok=True) - embedding_filename: str = f"{file_hash}.npy" - output_npy_filepath: str = os.path.join(engine.config.embeddings_output_path, embedding_filename) - np.save(output_npy_filepath, embeddings_numpy) - logger.info(f"Saved embedding matrix to {output_npy_filepath} with shape {embeddings_numpy.shape}") + # Quantization checking: only apply qi8 if configured and dim is 128 + quant_mode = ( + getattr(engine.config, "quantization", None) or "f32" + ).strip().lower() + + is_quantized = False + if quant_mode in ("qi8", "int8"): + dim = embeddings_numpy.shape[-1] + if dim == 128: + is_quantized = True + else: + logger.warning( + f"Quantization '{quant_mode}' requested but vector dimension is {dim} " + f"(requires 128). Saving unquantized float32 embeddings." + ) + + if is_quantized: + from bbq.maxsimd.quantization import qi8 + + val, scale = qi8(embeddings_numpy, dim=128) + output_filepath = os.path.join( + engine.config.embeddings_output_path, f"{file_hash}.npz" + ) + np.savez(output_filepath, values=val, scales=scale) + logger.info( + f"Saved quantized ({quant_mode}) embedding matrix to {output_filepath} " + f"(values: {val.shape}, scales: {scale.shape})" + ) + else: + output_filepath = os.path.join( + engine.config.embeddings_output_path, f"{file_hash}.npy" + ) + np.save(output_filepath, embeddings_numpy) + logger.info( + f"Saved unquantized embedding matrix to {output_filepath} with shape {embeddings_numpy.shape}" + ) tracker.update_file_status_to_done( file_hash=file_hash, file_path=pdf_filepath, num_pages=num_pages, - embedding_path=output_npy_filepath, + embedding_path=output_filepath, + ) + logger.info( + f"Successfully finished PDF ingestion for {pdf_filepath} ({num_pages} pages)" ) - logger.info(f"Successfully finished PDF ingestion for {pdf_filepath} ({num_pages} pages)") finally: if status_ctx: status_ctx.stop() except Exception as exception_instance: - logger.exception(f"Unhandled exception during PDF ingestion for {pdf_filepath}: {exception_instance}") + logger.exception( + f"Unhandled exception during PDF ingestion for {pdf_filepath}" + ) try: - file_hash_fallback: str = compute_file_sha256_hash(pdf_filepath) + target_hash = file_hash or compute_file_sha256_hash(pdf_filepath) tracker.update_file_status_to_failed( - file_hash=file_hash_fallback, + file_hash=target_hash, file_path=pdf_filepath, error_message=str(exception_instance), ) - except Exception: - pass + except Exception as fallback_err: + logger.error( + f"Failed to record ingestion error for {pdf_filepath}: {fallback_err}" + ) def process_single_pdf_file_deletion( @@ -134,18 +187,26 @@ def process_single_pdf_file_deletion( ) return True else: - logger.warning(f"No record found in database for deleted PDF file: {pdf_filepath}") + logger.warning( + f"No record found in database for deleted PDF file: {pdf_filepath}" + ) return False def sync_and_cleanup_deleted_files( - watch_folder: str, + watch_folder: Optional[str], tracker: SqlliteDB, ) -> None: records = tracker.fetch_all_records() for record in records: file_path = record.get("file_path") - if file_path and not os.path.exists(file_path): + if not file_path: + continue + if watch_folder and not os.path.abspath(file_path).startswith( + os.path.abspath(watch_folder) + ): + continue + if not os.path.exists(file_path): logger.info(f"Detected removed file during startup sync: {file_path}") process_single_pdf_file_deletion(pdf_filepath=file_path, tracker=tracker) @@ -166,7 +227,7 @@ def scan_and_ingest_existing_pdf_folder( for root_dir, _, file_names in os.walk(watch_folder): for file_name in sorted(file_names): if file_name.lower().endswith(".pdf"): - full_pdf_path: str = os.path.join(root_dir, file_name) + full_pdf_path = os.path.join(root_dir, file_name) process_single_pdf_file_ingestion( pdf_filepath=full_pdf_path, engine=engine, diff --git a/bbq/src/server/retrieval.py b/bbq/src/server/retrieval.py index ed1c9dc..d89b89a 100644 --- a/bbq/src/server/retrieval.py +++ b/bbq/src/server/retrieval.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import logging import os -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from bbq.src.storage.sql import SqlliteDB @@ -11,18 +13,23 @@ def query_indexed_documents( query_text: str, engine: Any, tracker: SqlliteDB, - top_k: int = 5, + top_k: Optional[int] = None, ) -> List[Dict[str, Any]]: """ - Encodes the query text, loads indexed PDF page embeddings, calculates MaxSim scores, - and returns top-k matching PDF pages. + Encodes query text, loads indexed page embeddings, calculates MaxSim scores + (supporting float32 and quantized INT8 qi8 modes via maxsimd), and returns top-k matching PDF pages. """ - if not query_text.strip(): + query_text = (query_text or "").strip() + if not query_text: return [] import numpy as np import torch - from maxsimd import maxsim + + try: + from bbq.maxsimd import maxsim + except ImportError: + from maxsimd import maxsim records = tracker.fetch_all_records() completed_records = [ @@ -33,10 +40,32 @@ def query_indexed_documents( logger.info("No completed PDF embeddings found in tracker database.") return [] + effective_top_k = ( + int(top_k) if top_k is not None else getattr(engine.config, "rag_top_k", 5) + ) + # Encode query: shape [1, q_len, dim] -> [q_len, dim] q_tensor: torch.Tensor = engine.encode_query_text_inputs([query_text]) q_mat = q_tensor[0].cpu().float().numpy() + # Check configured quantization mode + quant_mode = (getattr(engine.config, "quantization", None) or "f32").strip().lower() + use_quant = quant_mode in ("qi8", "int8") + + q_val = None + q_scale = None + if use_quant: + if q_mat.shape[-1] == 128: + from bbq.maxsimd.quantization import qi8 + + q_val, q_scale = qi8(q_mat, dim=128) + else: + logger.warning( + f"Quantization '{quant_mode}' requested but query embedding dimension is " + f"{q_mat.shape[-1]} (requires 128). Falling back to unquantized retrieval." + ) + use_quant = False + all_page_results: List[Dict[str, Any]] = [] for record in completed_records: @@ -44,16 +73,55 @@ def query_indexed_documents( if not emb_path or not os.path.exists(emb_path): continue - doc_emb = np.load(emb_path) - if doc_emb.ndim == 2: - doc_emb = np.expand_dims(doc_emb, axis=0) + try: + loaded = np.load(emb_path) + except Exception as load_err: + logger.error(f"Failed to load embedding file '{emb_path}': {load_err}") + continue - num_pages = doc_emb.shape[0] + # Check if saved file is quantized (.npz with values and scales) + if isinstance(loaded, np.lib.npyio.NpzFile) or ( + hasattr(loaded, "files") and "values" in loaded.files + ): + doc_val = loaded["values"] + doc_scale = loaded["scales"] + num_pages = doc_val.shape[0] if doc_val.ndim == 3 else 1 - # Multi-core Rayon parallelized MaxSim scoring for pages of this document - scores = maxsim(q_mat, doc_emb) - scores_list = [float(s) for s in scores] + if not use_quant: + # Query was unquantized float32, quantize on-the-fly to score against quantized doc + if q_mat.shape[-1] == 128: + from bbq.maxsimd.quantization import qi8 + + curr_q_val, curr_q_scale = qi8(q_mat, dim=128) + scores = maxsim( + curr_q_val, + doc_val, + q_scale=curr_q_scale, + d_scale=doc_scale, + ) + else: + logger.error( + f"Cannot score quantized document '{emb_path}' with non-128 dim query" + ) + continue + else: + scores = maxsim(q_val, doc_val, q_scale=q_scale, d_scale=doc_scale) + else: + doc_emb = loaded + if doc_emb.ndim == 2: + doc_emb = np.expand_dims(doc_emb, axis=0) + num_pages = doc_emb.shape[0] + if use_quant and doc_emb.shape[-1] == 128: + # Document was unquantized float32, quantize on-the-fly to match query + from bbq.maxsimd.quantization import qi8 + + doc_val, doc_scale = qi8(doc_emb, dim=128) + scores = maxsim(q_val, doc_val, q_scale=q_scale, d_scale=doc_scale) + else: + scores = maxsim(q_mat, doc_emb) + + scores_list = [float(s) for s in scores] filename = os.path.basename(record.get("file_path", "")) for page_idx, score in enumerate(scores_list): all_page_results.append( @@ -68,4 +136,4 @@ def query_indexed_documents( ) all_page_results.sort(key=lambda x: x["score"], reverse=True) - return all_page_results[:top_k] + return all_page_results[: max(0, effective_top_k)] diff --git a/bbq/src/server/server.py b/bbq/src/server/server.py index 9bf0973..d84a0f9 100644 --- a/bbq/src/server/server.py +++ b/bbq/src/server/server.py @@ -1,16 +1,17 @@ +from __future__ import annotations + import logging import os import signal import sys import threading import time -from typing import Any, Optional, Tuple +import warnings +from typing import Any, Optional from rich.console import Console -from bbq.src.config import ( - Config, -) +from bbq.src.config import Config from bbq.src.server.app import create_bbq_fastapi_app, run_http_server_in_thread from bbq.src.server.ingestion import ( process_single_pdf_file_deletion, @@ -20,10 +21,22 @@ from bbq.src.storage.sql import SqlliteDB from bbq.src.terminal import ( configure_server_logging, - render_server_status_rich_panel, + print_bbq, ) + from bbq.src.utils.watcher import start_pdf_folder_watcher +# Disable progress bars and warnings during weight loading +os.environ["TQDM_DISABLE"] = "1" +os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" +os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1" + +# Silence bitsandbytes BNB_CUDA_VERSION override warning +logging.getLogger("bitsandbytes").setLevel(logging.ERROR) +warnings.filterwarnings("ignore", message=".*BNB_CUDA_VERSION.*") +warnings.filterwarnings("ignore", module=".*bitsandbytes.*") + + logger = logging.getLogger("bbq.server") @@ -32,10 +45,11 @@ class BBQServer: BBQServer: Object-oriented server orchestrator for document indexing & retrieval. Features: - - Threaded fast engine weight loading with live Rich status spinner. - - Clean KeyboardInterrupt / signal handling without traceback. + - Clean ASCII logo banner and Rich status spinners matching client interface. + - Threaded fast engine weight loading without progress bar clutter or library warnings. - Automatic folder watching and PDF sync. - Non-blocking daemon HTTP server thread. + - Clean KeyboardInterrupt / signal handling. """ def __init__( @@ -44,14 +58,18 @@ def __init__( config_filepath: str = "config.yaml", host: str = "0.0.0.0", port: int = 8000, + without_logo: bool = False, ) -> None: self.config_filepath = config_filepath self.host = host self.port = port + self.without_logo = without_logo self.console = Console() self.bbq_logger = configure_server_logging() - self.config: Config = config if config is not None else Config.from_yaml(config_filepath) + self.config: Config = ( + config if config is not None else Config.from_yaml(config_filepath) + ) self.tracker = SqlliteDB(db_filepath=self.config.sqlite_db_path) self.engine: Optional[Any] = None @@ -60,23 +78,51 @@ def __init__( self.loading_thread: Optional[threading.Thread] = None self._is_ready = False + def print_logo( + self, + name: Optional[str] = None, + tag: str = "SERVER", + server: Optional[Any] = None, + ) -> None: + """Prints the BBQ banner unless without_logo is enabled.""" + if not self.without_logo: + srv = server if server is not None else self.port + print_bbq(name=name, tag=tag, server=srv) + def load_engine_threaded(self) -> None: """ Loads AI model and processor weights in a background thread while displaying a responsive Rich status spinner interface. Handles Ctrl+C (KeyboardInterrupt) cleanly. """ - logger.info(f"Loading engine and model (base_model_id={self.config.base_model_id})...") + logger.info( + f"Loading engine and model (base_model_id={self.config.base_model_id})..." + ) def _loader(): try: + os.environ["TQDM_DISABLE"] = "1" + os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" + os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1" + logging.getLogger("bitsandbytes").setLevel(logging.ERROR) + warnings.filterwarnings("ignore", message=".*BNB_CUDA_VERSION.*") + warnings.filterwarnings("ignore", module=".*bitsandbytes.*") + + try: + import transformers.utils.logging as tfl + + tfl.disable_progress_bar() + tfl.set_verbosity_error() + except Exception: + pass + from bbq.src.utils.model_loader import initialize_engine self.engine = initialize_engine(self.config) self._is_ready = True logger.info("Engine model and processor loaded successfully.") - except Exception as exc: - logger.exception(f"Failed to load engine weights: {exc}") + except Exception: + logger.exception("Failed to load engine weights") self.loading_thread = threading.Thread(target=_loader, daemon=True) self.loading_thread.start() @@ -93,45 +139,60 @@ def _loader(): self.stop() sys.exit(0) + if self.engine: + self.console.print( + f"[bold green]✔ Model and weights loaded successfully ({self.config.base_model_id})[/bold green]" + ) + def start(self) -> None: """ Starts the persistent document indexing server pipeline. """ - self.console.print(render_server_status_rich_panel(self.config)) - logger.info(f"Starting persistent document-indexing server (Logs saved to: {self.bbq_logger.log_file_path})...") + self.print_logo() + logger.info( + f"Starting persistent document-indexing server (Logs saved to: {self.bbq_logger.log_file_path})..." + ) reset_count: int = self.tracker.reset_in_progress_processing_to_pending() if reset_count > 0: - logger.info(f"Recovered startup state: reset {reset_count} leftover processing record(s) to pending.") + logger.info( + f"Recovered startup state: reset {reset_count} leftover processing record(s) to pending." + ) - # Threaded Engine Weight Loading with Spinner - self.load_engine_threaded() + # Register Signal Handlers + self._register_signal_handlers() - if not self.engine: - logger.error("Server startup aborted: Engine failed to load.") - return + # Start FastAPI HTTP Server Thread immediately so port is open and status can be queried + fastapi_app = create_bbq_fastapi_app( + engine=self.engine, + tracker=self.tracker, + get_engine_callback=lambda: self.engine, + config=self.config, + ) + self.server_thread = run_http_server_in_thread( + fastapi_app, host=self.host, port=self.port + ) - # Scan watch folder and ingest existing PDFs - try: - with self.console.status( - f"[bold green]Scanning '{self.config.watch_folder_path}' for existing PDF documents...[/bold green]", - spinner="dots", - ): - scan_and_ingest_existing_pdf_folder(engine=self.engine, tracker=self.tracker, console=self.console) - except (KeyboardInterrupt, SystemExit): - logger.info("PDF scan interrupted by user (Ctrl+C). Shutting down...") - self.stop() - sys.exit(0) + self.console.print( + f"[bold green]✔ BBQ server running on http://{self.host}:{self.port} " + f"(quantization: {self.config.quantization})[/bold green]\n" + f"[dim]Watching '{self.config.watch_folder_path}' for document changes. Press Ctrl+C to stop.[/dim]\n" + ) # File system watcher callbacks def pdf_detected_event_callback(pdf_filepath: str) -> None: logger.info(f"Folder watcher detected new/modified PDF: {pdf_filepath}") - process_single_pdf_file_ingestion( - pdf_filepath=pdf_filepath, - engine=self.engine, - tracker=self.tracker, - console=self.console, - ) + if self.engine: + process_single_pdf_file_ingestion( + pdf_filepath=pdf_filepath, + engine=self.engine, + tracker=self.tracker, + console=self.console, + ) + else: + logger.info( + f"Engine loading in progress; '{pdf_filepath}' will be ingested once engine is ready." + ) def pdf_deleted_event_callback(pdf_filepath: str) -> None: logger.info(f"Folder watcher detected deleted PDF: {pdf_filepath}") @@ -140,32 +201,45 @@ def pdf_deleted_event_callback(pdf_filepath: str) -> None: tracker=self.tracker, ) - # Start Folder Watcher - logger.info(f"Starting file system watcher on directory: {self.config.watch_folder_path}") + # Start Folder Watcher immediately + logger.info( + f"Starting file system watcher on directory: {self.config.watch_folder_path}" + ) self.observer, _ = start_pdf_folder_watcher( watch_directory_path=self.config.watch_folder_path, callback_on_pdf_ready=pdf_detected_event_callback, callback_on_pdf_deleted=pdf_deleted_event_callback, ) - # Start FastAPI HTTP Server Thread - fastapi_app = create_bbq_fastapi_app( - engine=self.engine, - tracker=self.tracker, - get_engine_callback=lambda: self.engine, - ) - self.server_thread = run_http_server_in_thread(fastapi_app, host=self.host, port=self.port) + # Threaded Engine Weight Loading with Spinner + self.load_engine_threaded() - # Register Signal Handlers - self._register_signal_handlers() + if not self.engine: + logger.error("Server startup aborted: Engine failed to load.") + return - logger.info( - f"Server is actively watching folder '{self.config.watch_folder_path}' and listening on port {self.port}. Press Ctrl+C to stop." - ) + # Scan watch folder and ingest existing PDFs + try: + with self.console.status( + f"[bold green]Scanning '{self.config.watch_folder_path}' for existing PDF documents...[/bold green]", + spinner="dots", + ): + scan_and_ingest_existing_pdf_folder( + engine=self.engine, tracker=self.tracker, console=self.console + ) + self.console.print( + f"[bold green]✔ PDF scan complete for '{self.config.watch_folder_path}'[/bold green]" + ) + except (KeyboardInterrupt, SystemExit): + logger.info("PDF scan interrupted by user (Ctrl+C). Shutting down...") + self.stop() + sys.exit(0) def _register_signal_handlers(self) -> None: def signal_handler_callback(signal_number: int, frame_object: Any) -> None: - logger.info(f"Received shutdown signal ({signal_number}). Initiating graceful server exit...") + logger.info( + f"Received shutdown signal ({signal_number}). Initiating graceful server exit..." + ) self.stop() sys.exit(0) @@ -184,9 +258,13 @@ def stop(self) -> None: try: reset_count: int = self.tracker.reset_in_progress_processing_to_pending() if reset_count > 0: - logger.info(f"Reset {reset_count} in-progress processing file record(s) back to pending state.") - except Exception: - pass + logger.info( + f"Reset {reset_count} in-progress processing file record(s) back to pending state." + ) + except Exception as exc: + logger.debug( + f"Non-critical error resetting db status during shutdown: {exc}" + ) logger.info("Persistent document indexing server shutdown complete.") @@ -194,28 +272,21 @@ def is_ready(self) -> bool: return self._is_ready -def execute_background_server_pipeline( - config_filepath: str, - config: Config, - tracker: SqlliteDB, - console: Console, - host: str = "0.0.0.0", - port: int = 8000, -) -> Tuple[Any, Any]: - """Helper runner for pipeline execution.""" - server = BBQServer(config_filepath=config_filepath, host=host, port=port) - server.start() - return server.observer, server.engine - - def start_document_indexing_server( config: Optional[Config] = None, config_filepath: str = "config.yaml", host: str = "0.0.0.0", port: int = 8000, + without_logo: bool = False, ) -> None: """Main entrypoint for starting the document indexing server.""" - server = BBQServer(config=config, config_filepath=config_filepath, host=host, port=port) + server = BBQServer( + config=config, + config_filepath=config_filepath, + host=host, + port=port, + without_logo=without_logo, + ) try: server.start() while True: diff --git a/bbq/src/terminal/logger.py b/bbq/src/terminal/logger.py index 040423c..0b63008 100644 --- a/bbq/src/terminal/logger.py +++ b/bbq/src/terminal/logger.py @@ -1,15 +1,75 @@ -import os +from __future__ import annotations + import logging +import os +import sys from typing import Optional -from rich.console import Console -from rich.logging import RichHandler from bbq.src.utils.futils import get_system_cache_dir + +class BBQConsoleLogFormatter(logging.Formatter): + """ + Formats console log records in the format: + [LEVEL] message [Date] + without leading column spaces or omitted timestamp alignment. + """ + + LEVEL_COLORS = { + logging.DEBUG: "\033[36m", # Cyan + logging.INFO: "\033[32m", # Green + logging.WARNING: "\033[33m", # Yellow + logging.ERROR: "\033[31m", # Red + logging.CRITICAL: "\033[1;31m", # Bold Red + } + RESET = "\033[0m" + DIM = "\033[90m" + + def format(self, record: logging.LogRecord) -> str: + date_str = self.formatTime(record, self.datefmt or "%Y-%m-%d %H:%M:%S") + msg = record.getMessage() + + if sys.stdout.isatty(): + color = self.LEVEL_COLORS.get(record.levelno, "") + line = f"{color}[{record.levelname}]{self.RESET} {msg} {self.DIM}[{date_str}]{self.RESET}" + else: + line = f"[{record.levelname}] {msg} [{date_str}]" + + if record.exc_info: + if not record.exc_text: + record.exc_text = self.formatException(record.exc_info) + if record.exc_text: + line = f"{line}\n{record.exc_text}" + if record.stack_info: + line = f"{line}\n{self.formatStack(record.stack_info)}" + return line + + +class BBQFileLogFormatter(logging.Formatter): + """ + Formats file log records in plain text format: + [LEVEL] message [Date] + """ + + def format(self, record: logging.LogRecord) -> str: + date_str = self.formatTime(record, self.datefmt or "%Y-%m-%d %H:%M:%S") + msg = record.getMessage() + line = f"[{record.levelname}] {msg} [{date_str}]" + + if record.exc_info: + if not record.exc_text: + record.exc_text = self.formatException(record.exc_info) + if record.exc_text: + line = f"{line}\n{record.exc_text}" + if record.stack_info: + line = f"{line}\n{self.formatStack(record.stack_info)}" + return line + + class BBQLogger: """ Logger manager class that sets up dual logging: - - Interactive Rich console output + - Clean console output formatted as '[LEVEL] message [Date]' - Persistent file logging saved to disk """ @@ -36,27 +96,19 @@ def configure_logger(self) -> logging.Logger: logger.setLevel(self.level) logger.handlers.clear() - # 1. Rich Console Handler - rich_handler = RichHandler( - console=Console(), - show_path=False, - rich_tracebacks=True, - markup=True, + console_handler = logging.StreamHandler(sys.stdout) + console_handler.setFormatter( + BBQConsoleLogFormatter(datefmt="%Y-%m-%d %H:%M:%S") ) - rich_handler.setLevel(self.level) - logger.addHandler(rich_handler) + console_handler.setLevel(self.level) + logger.addHandler(console_handler) - # 2. Disk File Handler log_dir = os.path.dirname(self.log_file_path) if log_dir: os.makedirs(log_dir, exist_ok=True) file_handler = logging.FileHandler(self.log_file_path, encoding="utf-8") - file_formatter = logging.Formatter( - "[%(asctime)s] [%(levelname)s] [%(name)s]: %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - ) - file_handler.setFormatter(file_formatter) + file_handler.setFormatter(BBQFileLogFormatter(datefmt="%Y-%m-%d %H:%M:%S")) file_handler.setLevel(self.level) logger.addHandler(file_handler) diff --git a/bbq/src/terminal/tui.py b/bbq/src/terminal/tui.py index 7ab616e..7b227e6 100644 --- a/bbq/src/terminal/tui.py +++ b/bbq/src/terminal/tui.py @@ -18,13 +18,13 @@ BOLD = "\033[1m" -def create_rich_console_logging_handler() -> RichHandler: - return RichHandler( - console=Console(), - show_path=False, - rich_tracebacks=True, - markup=True, - ) +def create_rich_console_logging_handler() -> logging.Handler: + import sys + from bbq.src.terminal.logger import BBQConsoleLogFormatter + + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter(BBQConsoleLogFormatter(datefmt="%Y-%m-%d %H:%M:%S")) + return handler def render_server_status_rich_panel(config: Config) -> Panel: diff --git a/bbq/src/utils/__init__.py b/bbq/src/utils/__init__.py index f1a6cf4..3f3f533 100644 --- a/bbq/src/utils/__init__.py +++ b/bbq/src/utils/__init__.py @@ -1,5 +1,10 @@ import importlib -from typing import Any +from typing import Any, TYPE_CHECKING + +if TYPE_CHECKING: + from bbq.src.common.base import BaseModelLoader, BaseProcessor, BaseEngineWrapper + + from bbq.src.utils.model_loader import initialize_engine from bbq.src.utils.futils import get_system_cache_dir from bbq.src.utils.pdf_utils import ( @@ -16,14 +21,17 @@ "EngineWrapper": "bbq.src.utils.model_loader", "initialize_engine": "bbq.src.utils.model_loader", "initialize_engine_from_yaml_config": "bbq.src.utils.model_loader", - "PDFWatchHandler": "bbq.src.utils.watcher", - "start_pdf_folder_watcher": "bbq.src.utils.watcher", + "model_loader": "bbq.src.utils.model_loader", + "watcher": "bbq.src.utils.watcher", } def __getattr__(name: str) -> Any: if name in _LAZY_IMPORTS: - module = importlib.import_module(_LAZY_IMPORTS[name]) + target_path = _LAZY_IMPORTS[name] + module = importlib.import_module(target_path) + if target_path.endswith(f".{name}"): + return module return getattr(module, name) raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/bbq/src/utils/model_loader.py b/bbq/src/utils/model_loader.py index 8d40178..4ab3b10 100644 --- a/bbq/src/utils/model_loader.py +++ b/bbq/src/utils/model_loader.py @@ -1,10 +1,23 @@ from typing import Any, List, Optional, Tuple, Type, Union - import torch -from peft import PeftConfig, PeftModel -from transformers import logging as tf_logging -from bbq.src.common.base import BaseEngineWrapper, BaseModel, BaseModelLoader, BaseProcessor +try: + from peft import PeftConfig, PeftModel +except ImportError: + PeftConfig = Any # type: ignore + PeftModel = Any # type: ignore + +try: + from transformers import logging as tf_logging +except ImportError: + tf_logging = None # type: ignore + +from bbq.src.common.base import ( + BaseEngineWrapper, + BaseModel, + BaseModelLoader, + BaseProcessor, +) from bbq.src.common.errors import ( BaseModelInstantiateError, LoRAAdapterLoadError, @@ -13,7 +26,26 @@ from bbq.src.config import Config from bbq.src.models.registry import ModelRegistry -tf_logging.set_verbosity_warning() +import os +import logging +import warnings + +# Disable progress bars and warnings during weight loading +os.environ["TQDM_DISABLE"] = "1" +os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" +os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1" + +# Silence bitsandbytes BNB_CUDA_VERSION override warning +logging.getLogger("bitsandbytes").setLevel(logging.ERROR) +warnings.filterwarnings("ignore", message=".*BNB_CUDA_VERSION.*") +warnings.filterwarnings("ignore", module=".*bitsandbytes.*") + +if tf_logging is not None: + tf_logging.set_verbosity_error() + try: + tf_logging.disable_progress_bar() + except Exception: + pass def determine_target_torch_device(device_preference: str = "auto") -> str: @@ -23,7 +55,9 @@ def determine_target_torch_device(device_preference: str = "auto") -> str: return device_preference -def resolve_torch_data_type(dtype_name: str = "bfloat16", target_device: str = "cpu") -> torch.dtype: +def resolve_torch_data_type( + dtype_name: str = "bfloat16", target_device: str = "cpu" +) -> torch.dtype: """Resolve the torch data type based on name and target device.""" if target_device == "cpu": return torch.float32 @@ -34,13 +68,27 @@ def resolve_torch_data_type(dtype_name: str = "bfloat16", target_device: str = " return torch.float32 +def _from_pretrained_fast(loader_cls_or_fn: Any, model_id_or_obj: Any, **kwargs: Any) -> Any: + """Load model/adapter/processor from local huggingface cache first to avoid network latency.""" + try: + return loader_cls_or_fn.from_pretrained( + model_id_or_obj, local_files_only=True, **kwargs + ) + except Exception: + return loader_cls_or_fn.from_pretrained( + model_id_or_obj, local_files_only=False, **kwargs + ) + + class EngineModelLoader(BaseModelLoader): """ Class for loading engine models """ @classmethod - def load_model_and_processor(cls, config_input: Any) -> Tuple[Union[BaseModel, PeftModel], BaseProcessor]: + def load_model_and_processor( + cls, config_input: Any + ) -> Tuple[Union[BaseModel, PeftModel], BaseProcessor]: return cls.load_model_with_lora_adapters(config_input) @staticmethod @@ -64,12 +112,15 @@ def load_model_with_lora_adapters( config = config_input target_device = determine_target_torch_device(config.device) - resolved_dtype: torch.dtype = resolve_torch_data_type(config.torch_dtype, target_device) + resolved_dtype: torch.dtype = resolve_torch_data_type( + config.torch_dtype, target_device + ) model_cls, processor_cls = ModelRegistry.get_for_model(config.base_model_id) try: - base_model: Any = model_cls.from_pretrained( + base_model: Any = _from_pretrained_fast( + model_cls, config.base_model_id, torch_dtype=resolved_dtype, ) @@ -82,12 +133,24 @@ def load_model_with_lora_adapters( if config.lora_adapter_id: try: - lora_adapter_config: PeftConfig = PeftConfig.from_pretrained(config.lora_adapter_id) - model_to_use = PeftModel.from_pretrained( - base_model, + lora_adapter_config: PeftConfig = _from_pretrained_fast( + PeftConfig, config.lora_adapter_id, - config=lora_adapter_config, ) + try: + model_to_use = PeftModel.from_pretrained( + base_model, + config.lora_adapter_id, + config=lora_adapter_config, + local_files_only=True, + ) + except Exception: + model_to_use = PeftModel.from_pretrained( + base_model, + config.lora_adapter_id, + config=lora_adapter_config, + local_files_only=False, + ) except Exception as exception_instance: raise LoRAAdapterLoadError( f"Failed to load and apply LoRA adapter weights from {config.lora_adapter_id}: {exception_instance}" @@ -96,10 +159,16 @@ def load_model_with_lora_adapters( model_to_use = model_to_use.to(target_device).eval() try: - processor: BaseProcessor = processor_cls.from_pretrained(config.base_model_id) + processor: BaseProcessor = _from_pretrained_fast( + processor_cls, + config.base_model_id, + ) except Exception: try: - processor = processor_cls.from_pretrained(config.lora_adapter_id) + processor = _from_pretrained_fast( + processor_cls, + config.lora_adapter_id, + ) except Exception as exception_instance: raise ProcessorLoadError( f"Failed to load processor from {config.base_model_id} or {config.lora_adapter_id}: {exception_instance}" @@ -124,7 +193,9 @@ def __init__( self.processor = processor self.config: Config = config - def encode_multimodal_document_images(self, images: List[Any], batch_size: int = 4) -> torch.Tensor: + def encode_multimodal_document_images( + self, images: List[Any], batch_size: int = 4 + ) -> torch.Tensor: """ Encodes a list of images into embeddings using sub-batching to prevent OOM errors on large documents. """ @@ -138,7 +209,8 @@ def encode_multimodal_document_images(self, images: List[Any], batch_size: int = batch_images = images[i : i + batch_size] processed_inputs = self.processor.process_images(batch_images) processed_inputs = { - k: v.to(target_device) if isinstance(v, torch.Tensor) else v for k, v in processed_inputs.items() + k: v.to(target_device) if isinstance(v, torch.Tensor) else v + for k, v in processed_inputs.items() } with torch.inference_mode(): @@ -160,7 +232,8 @@ def encode_query_text_inputs(self, texts: List[str]) -> torch.Tensor: processed_inputs = self.processor.process_texts(texts) target_device = next(self.model.parameters()).device processed_inputs = { - k: v.to(target_device) if isinstance(v, torch.Tensor) else v for k, v in processed_inputs.items() + k: v.to(target_device) if isinstance(v, torch.Tensor) else v + for k, v in processed_inputs.items() } with torch.inference_mode(): diff --git a/build.rs b/build.rs new file mode 100644 index 0000000..f433532 --- /dev/null +++ b/build.rs @@ -0,0 +1,29 @@ +fn main() { + println!("cargo:rerun-if-changed=build.rs"); + println!("cargo:rerun-if-env-changed=CARGO_CFG_TARGET_ARCH"); + + let target_arch = std::env::var("CARGO_CFG_TARGET_ARCH").unwrap_or_default(); + if target_arch != "x86_64" { + eprintln!("ERROR: BBQ-RAG (maxsimd) is ONLY supported on x86_64 architecture with FMA."); + eprintln!("Detected target architecture: '{}'.", target_arch); + eprintln!("BBQ-RAG cannot be installed or run on this machine."); + std::process::exit(1); + } + + #[cfg(target_arch = "x86_64")] + { + let target = std::env::var("TARGET").unwrap_or_default(); + let host = std::env::var("HOST").unwrap_or_default(); + if target == host || target.is_empty() { + if is_x86_feature_detected!("fma") { + } else { + eprintln!( + "ERROR: BBQ-RAG (maxsimd) requires FMA (Fused Multiply-Add) CPU instruction support." + ); + eprintln!("Your CPU does not support FMA instructions."); + eprintln!("BBQ-RAG cannot be installed or run on this machine."); + std::process::exit(1); + } + } + } +} diff --git a/sample.config.yaml b/sample.config.yaml index 80c8548..4662aae 100644 --- a/sample.config.yaml +++ b/sample.config.yaml @@ -26,6 +26,9 @@ watch_folder_path: "data/watch" # pdf image dpi will be send to the model pdf_render_dpi: 150 +# Output directory for saved retrieved page images +images_output_dir: "data/rr" + # Google Gemini multimodal RAG settings # If gemini_api_key is omitted/empty, GEMINI_API_KEY / GOOGLE_API_KEY env vars are used. gemini_api_key: "" diff --git a/tests/test_client.py b/tests/test_client.py index 9dabffa..490b908 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -416,7 +416,6 @@ def test_interactive_page_picker_non_tty(capsys): assert "design.pdf" in out - def test_without_opener_flag(monkeypatch, capsys): """Verify that --without-opener bypasses interactive_page_picker and prints static results.""" from bbq.src.main import run_client_query_command @@ -444,3 +443,191 @@ def test_without_opener_flag(monkeypatch, capsys): assert len(mock_picker_called) == 0 out = capsys.readouterr().out assert "Top 1 Matching PDF Pages" in out + + +@patch.object(BBQClient, "get_page_image") +def test_client_save_page_images(mock_get_img, tmp_path): + """Verify BBQClient.save_page_images saves page images with proper filenames to disk.""" + dummy_img = Image.new("RGB", (20, 20), color="red") + + def fake_get_page_image(file_path, page_number=1, dpi=None, save_path=None): + if save_path: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + dummy_img.save(save_path, format="PNG") + return dummy_img + + mock_get_img.side_effect = fake_get_page_image + + client = BBQClient() + results = [ + {"file_path": "docA.pdf", "page_number": 12, "score": 4.1}, + {"file_path": "docB.pdf", "page_number": 45, "score": 3.2}, + ] + out_dir = str(tmp_path / "saved_rr") + saved_paths = client.save_page_images(results, images_output_dir=out_dir) + + assert len(saved_paths) == 2 + assert os.path.exists(os.path.join(out_dir, "rr_1_12.png")) + assert os.path.exists(os.path.join(out_dir, "rr_2_45.png")) + assert results[0]["image_available"] is True + assert results[0]["saved_image_path"] == os.path.join(out_dir, "rr_1_12.png") + assert results[1]["saved_image_path"] == os.path.join(out_dir, "rr_2_45.png") + + +@patch.object(BBQClient, "query") +@patch.object(BBQClient, "get_page_image") +def test_client_query_and_answer_save_images_without_llm( + mock_get_img, mock_query, tmp_path +): + """Verify query_and_answer with use_llm=False and save_images=True saves images without calling LLM.""" + dummy_img = Image.new("RGB", (10, 10), color="green") + + def fake_get_page_image(file_path, page_number=1, dpi=None, save_path=None): + if save_path: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + dummy_img.save(save_path, format="PNG") + return dummy_img + + mock_get_img.side_effect = fake_get_page_image + mock_query.return_value = [ + {"file_path": "manual.pdf", "page_number": 7, "score": 2.9} + ] + + out_dir = str(tmp_path / "retrieval_images") + client = BBQClient() + resp = client.query_and_answer( + "search terms", + top_k=5, + use_llm=False, + save_images=True, + images_output_dir=out_dir, + ) + + assert resp["status"] == "retrieval_only" + assert resp["answer"] is None + assert os.path.exists(os.path.join(out_dir, "rr_1_7.png")) + + +def test_cli_save_images_flags_parsing(): + """Verify parsing of --save-images, -i, and --images-output-dir / -o flags.""" + parser = build_argument_parser() + + args1 = parser.parse_args(["query", "search", "--save-images"]) + assert args1.save_images is True + + args2 = parser.parse_args(["client", "search", "-i"]) + assert args2.save_images is True + + args3 = parser.parse_args(["query", "search", "-i", "-o", "custom/dir"]) + assert args3.save_images is True + assert args3.images_output_dir == "custom/dir" + + client_parser = build_client_main_parser() + c_args1 = client_parser.parse_args(["search", "-i"]) + assert c_args1.save_images is True + + c_args2 = client_parser.parse_args(["search", "--images-output-dir", "custom/path"]) + assert c_args2.images_output_dir == "custom/path" + + +@patch.object(BBQClient, "query") +@patch.object(BBQClient, "get_page_image") +def test_cli_save_images_execution(mock_get_img, mock_query, tmp_path, capsys): + """Verify CLI saves images to specified directory when -i is passed without --use-llm.""" + from bbq.src.main import run_client_query_command + + dummy_img = Image.new("RGB", (10, 10), color="yellow") + + def fake_get_page_image(file_path, page_number=1, dpi=None, save_path=None): + if save_path: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + dummy_img.save(save_path, format="PNG") + return dummy_img + + mock_get_img.side_effect = fake_get_page_image + mock_query.return_value = [ + {"file_path": "doc.pdf", "page_number": 3, "total_pages": 10, "score": 3.0} + ] + + out_dir = str(tmp_path / "cli_rr") + parser = build_argument_parser() + args = parser.parse_args( + [ + "query", + "test save", + "-i", + "-o", + out_dir, + "--without-opener", + "--without-logo", + ] + ) + run_client_query_command(args) + + assert os.path.exists(os.path.join(out_dir, "rr_1_3.png")) + captured = capsys.readouterr().out + assert "Saved 1 page image(s)" in captured + + +@patch.object(BBQClient, "get_page_image") +def test_client_save_page_images_threaded(mock_get_img, tmp_path): + """Verify BBQClient.save_page_images_threaded executes in a background thread and fires callbacks.""" + import threading + + dummy_img = Image.new("RGB", (10, 10), color="purple") + + def fake_get_page_image(file_path, page_number=1, dpi=None, save_path=None): + if save_path: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + dummy_img.save(save_path, format="PNG") + return dummy_img + + mock_get_img.side_effect = fake_get_page_image + + progress_events = [] + complete_events = [] + + def on_prog(idx, total, path): + progress_events.append((idx, total, path)) + + def on_comp(paths): + complete_events.append(paths) + + client = BBQClient() + results = [ + {"file_path": "ch1.pdf", "page_number": 5, "score": 4.0}, + {"file_path": "ch2.pdf", "page_number": 9, "score": 3.8}, + ] + out_dir = str(tmp_path / "threaded_rr") + thread = client.save_page_images_threaded( + results=results, + images_output_dir=out_dir, + on_progress=on_prog, + on_complete=on_comp, + ) + + assert isinstance(thread, threading.Thread) + thread.join(timeout=5.0) + + assert os.path.exists(os.path.join(out_dir, "rr_1_5.png")) + assert os.path.exists(os.path.join(out_dir, "rr_2_9.png")) + assert len(progress_events) == 2 + assert len(complete_events) == 1 + assert len(complete_events[0]) == 2 + + +def test_render_query_results_rich_clean(capsys): + """Verify render_query_results_rich displays top matching pages cleanly without background banner.""" + from bbq.src.terminal import render_query_results_rich + + results = [ + {"file_path": "p1.pdf", "page_number": 10}, + {"file_path": "p2.pdf", "page_number": 20}, + ] + + render_query_results_rich(results) + out = capsys.readouterr().out + assert "Top 2 Matching PDF Pages" in out + assert "Saving in bg" not in out + assert "p1.pdf" in out + assert "p2.pdf" in out diff --git a/tests/test_config.py b/tests/test_config.py index c377b4d..91a3593 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -31,6 +31,7 @@ def test_config_defaults(): assert config.gemini_model == "gemini-3.6-flash" assert config.rag_top_k == 3 assert config.quantization == "f32" + assert config.images_output_dir == "data/rr" def test_config_custom_initialization(): diff --git a/tests/test_hardware.py b/tests/test_hardware.py new file mode 100644 index 0000000..b900337 --- /dev/null +++ b/tests/test_hardware.py @@ -0,0 +1,57 @@ +""" +Unit tests for hardware and CPU architecture compatibility checks. +""" + +from unittest.mock import mock_open, patch +import pytest + +from bbq.src.common.hardware import check_cpu_support, verify_hardware_or_exit + + +def test_check_cpu_support_current_machine(): + """Verify that the current development/test machine passes the check.""" + is_supported, reason = check_cpu_support() + assert is_supported is True + assert reason == "" + + +@patch("platform.machine", return_value="aarch64") +def test_check_cpu_support_unsupported_arch(mock_arch): + """Verify that non-x86_64 architectures (e.g., aarch64, arm64) are rejected.""" + is_supported, reason = check_cpu_support() + assert is_supported is False + assert "Unsupported CPU architecture: 'aarch64'" in reason + assert "x86_64" in reason + + +@patch("platform.machine", return_value="x86_64") +@patch("sys.platform", "linux") +def test_check_cpu_support_missing_fma_on_linux(mock_arch): + """Verify that an x86_64 CPU without the FMA flag on Linux is rejected.""" + cpuinfo_without_fma = "flags\t\t: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov avx2\n" + with patch("builtins.open", mock_open(read_data=cpuinfo_without_fma)): + is_supported, reason = check_cpu_support() + assert is_supported is False + assert "FMA" in reason + + +@patch( + "bbq.src.common.hardware.check_cpu_support", + return_value=(False, "Mocked unsupported CPU"), +) +def test_verify_hardware_or_exit_terminates(mock_check, capsys): + """Verify that verify_hardware_or_exit exits with code 1 and writes to stderr on unsupported hardware.""" + with pytest.raises(SystemExit) as exc_info: + verify_hardware_or_exit(exit_code=42) + + assert exc_info.value.code == 42 + err = capsys.readouterr().err + assert "ERROR: BBQ-RAG Hardware Compatibility Check Failed" in err + assert "Mocked unsupported CPU" in err + + +@patch("bbq.src.common.hardware.check_cpu_support", return_value=(True, "")) +def test_verify_hardware_or_exit_passes(mock_check): + """Verify that verify_hardware_or_exit does not exit when hardware is supported.""" + # Should complete without raising SystemExit + verify_hardware_or_exit() diff --git a/tests/test_server.py b/tests/test_server.py new file mode 100644 index 0000000..a7f8afe --- /dev/null +++ b/tests/test_server.py @@ -0,0 +1,624 @@ +import os +import threading +from typing import Any, List +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch +from fastapi.testclient import TestClient +from PIL import Image + +import bbq.src.server as server_module +from bbq.src.config import Config +from bbq.src.server.app import create_bbq_fastapi_app, run_http_server_in_thread +from bbq.src.server.ingestion import ( + process_single_pdf_file_deletion, + process_single_pdf_file_ingestion, + scan_and_ingest_existing_pdf_folder, + sync_and_cleanup_deleted_files, +) +from bbq.src.server.retrieval import query_indexed_documents +from bbq.src.server.server import BBQServer, start_document_indexing_server +from bbq.src.storage.sql import SqlliteDB + + +@pytest.fixture +def mock_config(tmp_path): + """Provides a lightweight Config pointing to temporary directories.""" + watch_dir = tmp_path / "watch" + watch_dir.mkdir(parents=True, exist_ok=True) + emb_dir = tmp_path / "embeddings" + emb_dir.mkdir(parents=True, exist_ok=True) + db_file = tmp_path / "bbq_test.db" + + cfg = Config() + cfg.watch_folder_path = str(watch_dir) + cfg.embeddings_output_path = str(emb_dir) + cfg.sqlite_db_path = str(db_file) + cfg.base_model_id = "test/mock-model" + cfg.embedding_dim = 128 + cfg.rag_top_k = 3 + cfg.quantization = "f32" + cfg.pdf_render_dpi = 72 + return cfg + + +@pytest.fixture +def mock_tracker(tmp_path): + """Provides a fresh SqlliteDB instance using a temporary file.""" + db_file = str(tmp_path / "tracker_test.db") + return SqlliteDB(db_filepath=db_file) + + +class DummyEngine: + """Mock multimodal engine for testing retrieval and ingestion.""" + + def __init__(self, config: Config, dim: int = 128): + self.config = config + self.dim = dim + # Mock parameter to report device + mock_param = torch.nn.Parameter(torch.zeros(1, device="cpu")) + self.model = torch.nn.Module() + self.model.register_parameter("dummy_weight", mock_param) + + def encode_query_text_inputs(self, queries: List[str]) -> torch.Tensor: + # Return shape [len(queries), q_tokens=4, dim=128] + return torch.randn(len(queries), 4, self.dim, dtype=torch.float32) + + def encode_multimodal_document_images( + self, images: List[Any], batch_size: int = 4 + ) -> torch.Tensor: + # Return shape [len(images), doc_tokens=8, dim=128] + return torch.randn(len(images), 8, self.dim, dtype=torch.float32) + + +def test_server_module_dynamic_exports(): + """Verify lazy exports via bbq.src.server.__getattr__.""" + assert hasattr(server_module, "BBQServer") + assert hasattr(server_module, "create_bbq_fastapi_app") + assert hasattr(server_module, "query_indexed_documents") + assert hasattr(server_module, "process_single_pdf_file_ingestion") + assert hasattr(server_module, "process_single_pdf_file_deletion") + assert hasattr(server_module, "scan_and_ingest_existing_pdf_folder") + assert hasattr(server_module, "sync_and_cleanup_deleted_files") + assert hasattr(server_module, "run_http_server_in_thread") + assert hasattr(server_module, "start_document_indexing_server") + + with pytest.raises(AttributeError, match="has no attribute 'non_existent_symbol'"): + _ = getattr(server_module, "non_existent_symbol") + + +def test_fastapi_root_endpoint(): + """Verify GET / returns expected online payload.""" + app = create_bbq_fastapi_app() + client = TestClient(app) + response = client.get("/") + assert response.status_code == 200 + data = response.json() + assert data["status"] == "online" + assert "BBQ RAG" in data["message"] + + +def test_fastapi_status_engine_loading(mock_tracker, mock_config): + """Verify GET /status when engine is not yet loaded.""" + app = create_bbq_fastapi_app(engine=None, tracker=mock_tracker, config=mock_config) + client = TestClient(app) + response = client.get("/status") + assert response.status_code == 200 + data = response.json() + assert data["status"] == "loading_engine" + assert data["engine_ready"] is False + assert data["base_model_id"] == "test/mock-model" + assert data["device"] == "Loading..." + assert data["total_documents"] == 0 + assert data["indexed_documents"] == 0 + + +def test_fastapi_status_engine_ready(mock_tracker, mock_config): + """Verify GET /status when engine is active.""" + engine = DummyEngine(mock_config) + # Add a mock record in tracker + mock_tracker.update_file_status_to_processing("hash123", "/path/to/doc.pdf") + mock_tracker.update_file_status_to_done( + "hash123", "/path/to/doc.pdf", 5, "/path/to/emb.npy" + ) + + app = create_bbq_fastapi_app( + engine=engine, tracker=mock_tracker, config=mock_config + ) + client = TestClient(app) + response = client.get("/status") + assert response.status_code == 200 + data = response.json() + assert data["status"] == "online" + assert data["engine_ready"] is True + assert data["device"] == "cpu" + assert data["quantization"] == "f32" + assert data["total_documents"] == 1 + assert data["indexed_documents"] == 1 + + +def test_fastapi_documents_endpoint(mock_tracker): + """Verify GET /documents returns list of tracked documents.""" + mock_tracker.update_file_status_to_processing("hash_a", "/docA.pdf") + app = create_bbq_fastapi_app(tracker=mock_tracker) + client = TestClient(app) + response = client.get("/documents") + assert response.status_code == 200 + docs = response.json().get("documents", []) + assert len(docs) == 1 + assert docs[0]["file_hash"] == "hash_a" + + +def test_fastapi_query_empty_rejects(mock_tracker, mock_config): + """Verify query endpoints reject empty or whitespace queries with 400.""" + engine = DummyEngine(mock_config) + app = create_bbq_fastapi_app( + engine=engine, tracker=mock_tracker, config=mock_config + ) + client = TestClient(app) + + # POST empty query + res_post = client.post("/query", json={"query": " ", "top_k": 5}) + assert res_post.status_code == 400 + assert "Query cannot be empty" in res_post.json()["detail"] + + # GET empty query + res_get = client.get("/query?q=") + assert res_get.status_code == 400 + assert "Query cannot be empty" in res_get.json()["detail"] + + +def test_fastapi_query_engine_not_ready(mock_tracker): + """Verify query endpoints return 503 when engine is None.""" + app = create_bbq_fastapi_app(engine=None, tracker=mock_tracker) + client = TestClient(app) + + res_post = client.post("/query", json={"query": "test query"}) + assert res_post.status_code == 503 + assert "initializing" in res_post.json()["detail"] + + res_get = client.get("/query?q=test") + assert res_get.status_code == 503 + assert "initializing" in res_get.json()["detail"] + + +@patch("bbq.src.server.app.query_indexed_documents") +def test_fastapi_query_success(mock_query_func, mock_tracker, mock_config): + """Verify POST and GET /query return structured retrieval results.""" + mock_query_func.return_value = [ + {"filename": "doc.pdf", "page_number": 1, "score": 25.5} + ] + engine = DummyEngine(mock_config) + app = create_bbq_fastapi_app( + engine=engine, tracker=mock_tracker, config=mock_config + ) + client = TestClient(app) + + # POST /query + res_post = client.post("/query", json={"query": "machine learning", "top_k": 3}) + assert res_post.status_code == 200 + post_data = res_post.json() + assert post_data["query"] == "machine learning" + assert post_data["top_k"] == 3 + assert len(post_data["results"]) == 1 + assert post_data["results"][0]["score"] == 25.5 + + # GET /query + res_get = client.get("/query?q=deep+learning&top_k=2") + assert res_get.status_code == 200 + get_data = res_get.json() + assert get_data["query"] == "deep learning" + assert get_data["top_k"] == 2 + + +@patch("bbq.src.server.app.extract_single_pdf_page_image") +def test_fastapi_page_image_endpoint(mock_extract_img): + """Verify GET /page_image renders PNG bytes or returns 400 on error.""" + # Success case: returns 10x10 PIL image + img = Image.new("RGB", (10, 10), color="blue") + mock_extract_img.return_value = img + + app = create_bbq_fastapi_app() + client = TestClient(app) + res = client.get("/page_image?file_path=/fake/doc.pdf&page_number=1&dpi=72") + assert res.status_code == 200 + assert res.headers["content-type"] == "image/png" + assert len(res.content) > 0 + + # Error case: exception raised + mock_extract_img.side_effect = RuntimeError("PDF page does not exist") + res_err = client.get("/page_image?file_path=/fake/doc.pdf&page_number=99") + assert res_err.status_code == 400 + assert "PDF page does not exist" in res_err.json()["detail"] + + +@patch("uvicorn.run") +def test_run_http_server_in_thread(mock_uvicorn_run): + """Verify run_http_server_in_thread spawns a daemon thread running uvicorn.""" + app = create_bbq_fastapi_app() + thread = run_http_server_in_thread(app, host="127.0.0.1", port=9999) + assert isinstance(thread, threading.Thread) + assert thread.daemon is True + thread.join(timeout=1.0) + mock_uvicorn_run.assert_called_once_with( + app, host="127.0.0.1", port=9999, log_level="warning" + ) + + +def test_query_indexed_documents_empty_inputs(mock_tracker, mock_config): + """Verify empty query or empty database returns empty list.""" + engine = DummyEngine(mock_config) + # Empty query + assert query_indexed_documents("", engine, mock_tracker) == [] + assert query_indexed_documents(" ", engine, mock_tracker) == [] + + # Non-empty query, but no indexed documents in tracker + assert query_indexed_documents("hello", engine, mock_tracker) == [] + + +def test_query_indexed_documents_float32(mock_tracker, mock_config, tmp_path): + """Verify retrieval against float32 .npy embeddings.""" + engine = DummyEngine(mock_config, dim=128) + + # Create dummy float32 document embeddings: 2 pages, 8 tokens per page, 128 dim + doc_emb = np.random.randn(2, 8, 128).astype(np.float32) + emb_file = str(tmp_path / "doc1.npy") + np.save(emb_file, doc_emb) + + mock_tracker.update_file_status_to_done( + file_hash="hash1", + file_path="/watch/doc1.pdf", + num_pages=2, + embedding_path=emb_file, + ) + + results = query_indexed_documents( + query_text="query test", + engine=engine, + tracker=mock_tracker, + top_k=5, + ) + assert len(results) == 2 + assert results[0]["filename"] == "doc1.pdf" + assert results[0]["file_hash"] == "hash1" + assert results[0]["total_pages"] == 2 + assert "score" in results[0] + assert results[0]["score"] >= results[1]["score"] + + +def test_query_indexed_documents_quantized_qi8(mock_tracker, mock_config, tmp_path): + """Verify retrieval against quantized qi8 .npz embeddings.""" + mock_config.quantization = "qi8" + engine = DummyEngine(mock_config, dim=128) + + from bbq.maxsimd.quantization import qi8 + + # Create 3 pages of document embeddings, quantize to qi8 + doc_f32 = np.random.randn(3, 8, 128).astype(np.float32) + val, scale = qi8(doc_f32, dim=128) + + emb_file = str(tmp_path / "quant_doc.npz") + np.savez(emb_file, values=val, scales=scale) + + mock_tracker.update_file_status_to_done( + file_hash="hash_quant", + file_path="/watch/quant_doc.pdf", + num_pages=3, + embedding_path=emb_file, + ) + + results = query_indexed_documents( + query_text="quantum computing", + engine=engine, + tracker=mock_tracker, + top_k=2, + ) + assert len(results) == 2 + assert results[0]["filename"] == "quant_doc.pdf" + assert results[0]["score"] >= results[1]["score"] + + +def test_query_indexed_documents_mixed_modes(mock_tracker, mock_config, tmp_path): + """Verify query handles cross-mode matching (float32 query against qi8 doc, and qi8 query against f32 doc).""" + from bbq.maxsimd.quantization import qi8 + + # Doc 1: quantized .npz + doc1_f32 = np.random.randn(2, 6, 128).astype(np.float32) + v1, s1 = qi8(doc1_f32, dim=128) + file_npz = str(tmp_path / "doc1.npz") + np.savez(file_npz, values=v1, scales=s1) + mock_tracker.update_file_status_to_done("h1", "/watch/doc1.pdf", 2, file_npz) + + # Doc 2: unquantized .npy + doc2_f32 = np.random.randn(2, 6, 128).astype(np.float32) + file_npy = str(tmp_path / "doc2.npy") + np.save(file_npy, doc2_f32) + mock_tracker.update_file_status_to_done("h2", "/watch/doc2.pdf", 2, file_npy) + + # Case A: Engine is configured as f32 -> scores against both doc1 (.npz) and doc2 (.npy) + mock_config.quantization = "f32" + engine_f32 = DummyEngine(mock_config, dim=128) + results_f32 = query_indexed_documents("hello", engine_f32, mock_tracker, top_k=10) + assert len(results_f32) == 4 + + # Case B: Engine is configured as qi8 -> scores against both doc1 (.npz) and doc2 (.npy) + mock_config.quantization = "qi8" + engine_qi8 = DummyEngine(mock_config, dim=128) + results_qi8 = query_indexed_documents("hello", engine_qi8, mock_tracker, top_k=10) + assert len(results_qi8) == 4 + + +def test_query_indexed_documents_corrupted_file_handling( + mock_tracker, mock_config, tmp_path +): + """Verify retrieval gracefully catches corrupt or missing embedding files without failing.""" + engine = DummyEngine(mock_config, dim=128) + + # Create a corrupted file + corrupt_file = str(tmp_path / "corrupt.npy") + with open(corrupt_file, "wb") as f: + f.write(b"not a valid numpy file") + + mock_tracker.update_file_status_to_done( + "h_corrupt", "/watch/corrupt.pdf", 1, corrupt_file + ) + mock_tracker.update_file_status_to_done( + "h_missing", "/watch/missing.pdf", 1, "/non/existent/path.npy" + ) + + # Should not raise exception and return empty + results = query_indexed_documents("test", engine, mock_tracker) + assert results == [] + +def test_process_single_pdf_file_ingestion_not_found(mock_tracker, mock_config): + """Verify process_single_pdf_file_ingestion returns early if file does not exist.""" + engine = DummyEngine(mock_config) + process_single_pdf_file_ingestion("/nonexistent/file.pdf", engine, mock_tracker) + assert len(mock_tracker.fetch_all_records()) == 0 + + +def test_process_single_pdf_file_ingestion_skips_done( + mock_tracker, mock_config, tmp_path +): + """Verify process_single_pdf_file_ingestion skips already processed PDFs.""" + dummy_pdf = tmp_path / "already_done.pdf" + dummy_pdf.write_bytes(b"%PDF-1.4 mock content") + + from bbq.src.utils.pdf_utils import compute_file_sha256_hash + + pdf_hash = compute_file_sha256_hash(str(dummy_pdf)) + mock_tracker.update_file_status_to_done( + pdf_hash, str(dummy_pdf), 1, "/fake/emb.npy" + ) + + engine = DummyEngine(mock_config) + with patch("bbq.src.server.ingestion.get_pdf_total_pages") as mock_pages: + process_single_pdf_file_ingestion(str(dummy_pdf), engine, mock_tracker) + mock_pages.assert_not_called() + + +@patch("bbq.src.server.ingestion.get_pdf_total_pages", return_value=0) +def test_process_single_pdf_file_ingestion_zero_pages( + mock_pages, mock_tracker, mock_config, tmp_path +): + """Verify zero-page PDF marks status as failed.""" + dummy_pdf = tmp_path / "zero_pages.pdf" + dummy_pdf.write_bytes(b"%PDF-1.4 mock") + + engine = DummyEngine(mock_config) + process_single_pdf_file_ingestion(str(dummy_pdf), engine, mock_tracker) + + records = mock_tracker.fetch_all_records() + assert len(records) == 1 + assert records[0]["status"] == "failed" + assert "zero pages" in records[0]["error_message"] + + +@patch("bbq.src.server.ingestion.extract_pdf_page_range_to_pil_images") +@patch("bbq.src.server.ingestion.get_pdf_total_pages", return_value=2) +def test_process_single_pdf_file_ingestion_f32_success( + mock_pages, mock_extract, mock_tracker, mock_config, tmp_path +): + """Verify successful ingestion in float32 mode produces .npy embeddings.""" + mock_config.quantization = "f32" + engine = DummyEngine(mock_config, dim=128) + + mock_extract.return_value = [Image.new("RGB", (10, 10)), Image.new("RGB", (10, 10))] + + dummy_pdf = tmp_path / "doc.pdf" + dummy_pdf.write_bytes(b"%PDF-1.4 test document content") + + process_single_pdf_file_ingestion(str(dummy_pdf), engine, mock_tracker) + + record = mock_tracker.fetch_file_record_by_hash( + mock_tracker.fetch_all_records()[0]["file_hash"] + ) + assert record["status"] == "done" + assert record["num_pages"] == 2 + assert record["embedding_path"].endswith(".npy") + assert os.path.exists(record["embedding_path"]) + + +@patch("bbq.src.server.ingestion.extract_pdf_page_range_to_pil_images") +@patch("bbq.src.server.ingestion.get_pdf_total_pages", return_value=1) +def test_process_single_pdf_file_ingestion_qi8_success( + mock_pages, mock_extract, mock_tracker, mock_config, tmp_path +): + """Verify successful ingestion in qi8 mode produces .npz embeddings.""" + mock_config.quantization = "qi8" + engine = DummyEngine(mock_config, dim=128) + + mock_extract.return_value = [Image.new("RGB", (10, 10))] + dummy_pdf = tmp_path / "doc_qi8.pdf" + dummy_pdf.write_bytes(b"%PDF-1.4 test qi8 document") + + process_single_pdf_file_ingestion(str(dummy_pdf), engine, mock_tracker) + + record = mock_tracker.fetch_all_records()[0] + assert record["status"] == "done" + assert record["embedding_path"].endswith(".npz") + assert os.path.exists(record["embedding_path"]) + + loaded = np.load(record["embedding_path"]) + assert "values" in loaded.files + assert "scales" in loaded.files + + +def test_process_single_pdf_file_deletion(mock_tracker, tmp_path): + """Verify deletion removes file and associated embeddings from tracker and filesystem.""" + emb_file = tmp_path / "test_emb.npy" + emb_file.write_bytes(b"dummy embedding data") + pdf_file = str(tmp_path / "sample.pdf") + + mock_tracker.update_file_status_to_done("hash_del", pdf_file, 1, str(emb_file)) + assert os.path.exists(str(emb_file)) + + # Existing file deletion + deleted = process_single_pdf_file_deletion(pdf_file, mock_tracker) + assert deleted is True + assert not os.path.exists(str(emb_file)) + assert len(mock_tracker.fetch_all_records()) == 0 + + # Non-existent record deletion returns False + assert process_single_pdf_file_deletion("/unknown.pdf", mock_tracker) is False + + +def test_sync_and_cleanup_deleted_files(mock_tracker, tmp_path): + """Verify cleanup purges missing files inside watch folder but preserves existing ones.""" + watch_dir = tmp_path / "watch" + watch_dir.mkdir(parents=True, exist_ok=True) + + # File A exists + file_a = watch_dir / "existing.pdf" + file_a.write_bytes(b"pdf data") + mock_tracker.update_file_status_to_done( + "ha", str(file_a), 1, str(tmp_path / "a.npy") + ) + + # File B deleted from disk + file_b = watch_dir / "removed.pdf" + mock_tracker.update_file_status_to_done( + "hb", str(file_b), 1, str(tmp_path / "b.npy") + ) + + # File C outside watch folder + outside_file = tmp_path / "outside.pdf" + mock_tracker.update_file_status_to_done( + "hc", str(outside_file), 1, str(tmp_path / "c.npy") + ) + + sync_and_cleanup_deleted_files(str(watch_dir), mock_tracker) + + remaining_hashes = [r["file_hash"] for r in mock_tracker.fetch_all_records()] + assert "ha" in remaining_hashes + assert "hb" not in remaining_hashes + assert "hc" in remaining_hashes + + +@patch("bbq.src.server.ingestion.process_single_pdf_file_ingestion") +def test_scan_and_ingest_existing_pdf_folder( + mock_ingest, mock_tracker, mock_config, tmp_path +): + """Verify scan_and_ingest_existing_pdf_folder processes all .pdf files.""" + watch_dir = tmp_path / "watch_scan" + watch_dir.mkdir(parents=True, exist_ok=True) + mock_config.watch_folder_path = str(watch_dir) + + (watch_dir / "file1.pdf").write_bytes(b"pdf1") + (watch_dir / "file2.PDF").write_bytes(b"pdf2") + (watch_dir / "file3.txt").write_bytes(b"not a pdf") + + engine = DummyEngine(mock_config) + scan_and_ingest_existing_pdf_folder(engine, mock_tracker) + + assert mock_ingest.call_count == 2 + + +def test_bbq_server_initialization(mock_config): + """Verify BBQServer initializes attributes correctly.""" + server = BBQServer( + config=mock_config, host="127.0.0.1", port=8001, without_logo=True + ) + assert server.host == "127.0.0.1" + assert server.port == 8001 + assert server.without_logo is True + assert server.is_ready() is False + assert server.engine is None + + +def test_bbq_server_print_logo(capsys, mock_config): + """Verify print_logo honors without_logo flag.""" + server_no_logo = BBQServer(config=mock_config, without_logo=True) + server_no_logo.print_logo() + assert capsys.readouterr().out == "" + + server_with_logo = BBQServer(config=mock_config, without_logo=False) + server_with_logo.print_logo() + out = capsys.readouterr().out + assert "SERVER" in out + + +@pytest.mark.skipif( + bool(os.environ.get("CI")) or bool(os.environ.get("GITHUB_ACTIONS")), + reason="Skipped in CI: requires full model environment with peft", +) +@patch("bbq.src.utils.model_loader.initialize_engine") +def test_bbq_server_load_engine_threaded(mock_init_engine, mock_config): + """Verify load_engine_threaded initializes engine and sets _is_ready.""" + mock_init_engine.return_value = DummyEngine(mock_config) + server = BBQServer(config=mock_config, without_logo=True) + + server.load_engine_threaded() + assert server.is_ready() is True + assert server.engine is not None + + +@pytest.mark.skipif( + bool(os.environ.get("CI")) or bool(os.environ.get("GITHUB_ACTIONS")), + reason="Skipped in CI: requires full model environment with peft", +) +@patch("bbq.src.server.server.start_pdf_folder_watcher") +@patch("bbq.src.server.server.run_http_server_in_thread") +@patch("bbq.src.server.server.scan_and_ingest_existing_pdf_folder") +@patch("bbq.src.utils.model_loader.initialize_engine") +def test_bbq_server_start_and_stop_lifecycle( + mock_init_engine, + mock_scan, + mock_run_http, + mock_watcher, + mock_config, +): + """Verify BBQServer.start runs the full server pipeline and stop() cleans up.""" + mock_init_engine.return_value = DummyEngine(mock_config) + mock_obs = MagicMock() + mock_watcher.return_value = (mock_obs, None) + mock_run_http.return_value = MagicMock() + + server = BBQServer(config=mock_config, without_logo=True) + with patch.object(server, "_register_signal_handlers"): + server.start() + + assert server.is_ready() is True + mock_scan.assert_called_once() + mock_watcher.assert_called_once() + mock_run_http.assert_called_once() + + server.stop() + mock_obs.stop.assert_called_once() + + +@patch("bbq.src.server.server.BBQServer") +def test_start_document_indexing_server(mock_server_cls, mock_config): + """Verify start_document_indexing_server instantiates and starts BBQServer.""" + mock_server_instance = MagicMock() + # Let start() run and stop immediately by raising KeyboardInterrupt on join + mock_server_instance.observer.join.side_effect = KeyboardInterrupt + mock_server_cls.return_value = mock_server_instance + + with pytest.raises(SystemExit): + start_document_indexing_server(config=mock_config, without_logo=True) + + mock_server_instance.start.assert_called_once() + mock_server_instance.stop.assert_called_once()