diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 68b80f8..39d9d83 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1,4 +1,4 @@ -name: maxsimd +name: CI on: push: @@ -7,7 +7,7 @@ on: branches: [ main ] jobs: - maxsimd: + test: name: Build & Test (Rust + Python) runs-on: ubuntu-latest @@ -32,7 +32,7 @@ jobs: with: python-version: "3.13" - # 5. Run Rust Cargo Tests (with dev features enabled for MKL/BLAS) + # 5. Run Rust Cargo Tests - name: Run Cargo Tests run: | cargo test --verbose @@ -43,36 +43,26 @@ jobs: python -m pip install --upgrade pip pip install maturin numpy pytest - # 7. Build and install maxsimd CPython extension via Maturin - - name: Build and install maxsimd (Release) + # 7. Build and install BBQ package via pip + - name: Build and install BBQ (Release) run: | pip install . - # 8. Check version and exported functions - - name: Verify maxsimd version and exports + # 8. Check BBQ version and CLI + - name: Verify BBQ version and CLI run: | python -c " - import maxsimd - version = getattr(maxsimd, '__version__', 'N/A') + import bbq + version = getattr(bbq, '__version__', 'N/A') print('=' * 50) - print(f'MaxSimD Version: {version}') - print('Exported functions:', [item for item in dir(maxsimd) if not item.startswith('__')]) + print(f'BBQ Version: {version}') print('=' * 50) + assert version != 'N/A' " + bbq --version - # 9. Run a fast MaxSim computation smoke test - - name: Run MaxSim Python smoke test + # 9. Run Python unit tests + - name: Run Pytest test suite run: | - python -c " - import maxsimd - import numpy as np + pytest - # 32 query tokens, 256 document tokens, 128 dimension - q = np.random.randn(32, 128).astype(np.float32) - d = np.random.randn(256, 128).astype(np.float32) - - score = maxsimd.maxsim(q, d) - print('Calculated MaxSim score successfully:', score) - assert len(score) == 1 - assert isinstance(score[0], float) - " diff --git a/Cargo.toml b/Cargo.toml index 70c515e..c3e453b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "maxsimd" -version = "0.1.3-alpha" +version = "0.1.4-alpha" edition = "2024" homepage = "https://github.com/cmd-HMN" description = "This packages is the simd implementations of max sim function for faster calculation" @@ -19,20 +19,15 @@ rust.unsafe_code = "allow" [dependencies] pyo3 = "0.29.0" -rayon = "1.10" +rayon = "1.10" pyo3-stub-gen = "0.23.0" intel-mkl-src = { version = "0.8.1", optional = true, features = ["mkl-static-lp64-seq"] } -libc = {version = "0.2", optional = true} -# [dependencies.intel-mkl-src] -# version = "0.8.1" -# optional = true +libc = { version = "0.2", optional = true } [dev-dependencies] criterion = "0.8.2" -approx = "0.5" rand = "0.10.2" - [[bin]] name = "stub" path = "rust-core/_bin/stub.rs" diff --git a/README.md b/README.md index f85157c..92fd0bb 100644 --- a/README.md +++ b/README.md @@ -107,6 +107,10 @@ rag_top_k: 3 Launch the document indexing server. It will monitor `data/watch/` for new PDF files and automatically compute embeddings: ```bash +# Using the bbq CLI command +bbq server --config config.yaml + +# Or via Python module python -m bbq.src.main server --config config.yaml ``` @@ -115,14 +119,32 @@ python -m bbq.src.main server --config config.yaml Search indexed documents from the command line: ```bash -# Query top 3 matching pages -python -m bbq.src.main query "What was the operating margin in Q3?" --top-k 3 +# Fetch top 10 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?" -# Query with Gemini Multimodal RAG (generates grounded answer from top 3 page images) +# Query with Gemini Multimodal LLM (displays live cooking spinner and multimodal answer) export GEMINI_API_KEY="your-gemini-api-key" -python -m bbq.src.main query "Summarize the revenue growth" --top-k 3 +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 + +# Run directly via the client module +python -m bbq.src.client --infinite --use-llm ``` +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. +- Type `clear` or `cls` to clear the terminal screen. +- Type `q`, `quit`, or `exit` (or `Ctrl+C`) to quit. + If no Gemini API key is provided or the API is unavailable, the client automatically displays the matching document pages without crashing. ### 4. Python API Usage @@ -207,7 +229,3 @@ We express our sincere thanks to the original authors and maintainers for their For full project lineage, paper citations, and intellectual property notices, see [ABOUT.md](ABOUT.md). --- - -## Disclaimer - -This project, its documentation, and parts of its codebase and benchmarks were developed with the assistance of AI tools. As an evolving early-stage project, this README and documentation may contain preliminary assumptions or specifications that are actively being refined. Future commits will continuously audit, validate, and update these details to ensure ongoing accuracy, correctness, and benchmarking rigor. No warranties or guarantees of fitness for a particular purpose are provided. diff --git a/bbq/__init__.py b/bbq/__init__.py new file mode 100644 index 0000000..8cae060 --- /dev/null +++ b/bbq/__init__.py @@ -0,0 +1,5 @@ +"""BBQ package.""" + +from bbq.src.common.version import __version__ + +__all__ = ["__version__"] diff --git a/bbq/maxsimd/__init__.pyi b/bbq/maxsimd/__init__.pyi index 507274d..aff3263 100644 --- a/bbq/maxsimd/__init__.pyi +++ b/bbq/maxsimd/__init__.pyi @@ -47,4 +47,4 @@ def maxsim( **kwargs: typing.Any, ) -> builtins.list[builtins.float]: ... -__version__ = "0.1.3-alpha" +__version__ = "0.1.4-alpha" diff --git a/bbq/src/client/__init__.py b/bbq/src/client/__init__.py index dff17b8..007aee0 100644 --- a/bbq/src/client/__init__.py +++ b/bbq/src/client/__init__.py @@ -1,5 +1,5 @@ from bbq.src.client.client import BBQClient, get_local_pdf_page_image from bbq.src.client.gemini import GeminiClient +from bbq.src.terminal.tui import print_bbq -__all__ = ["BBQClient", "GeminiClient", "get_local_pdf_page_image"] - +__all__ = ["BBQClient", "GeminiClient", "get_local_pdf_page_image", "print_bbq"] diff --git a/bbq/src/client/client.py b/bbq/src/client/client.py index 9963090..741df2f 100644 --- a/bbq/src/client/client.py +++ b/bbq/src/client/client.py @@ -1,144 +1,127 @@ """ -Client module for interacting with the ColPali Document Retrieval Server, -retrieving document matches, and fetching/saving page images. +BBQ RAG Client module """ +from __future__ import annotations + import os +import io import logging -import time -from typing import List, Dict, Any, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional import requests from PIL import Image -import io + +from bbq.src.terminal.tui import print_bbq + +if TYPE_CHECKING: + from bbq.src.config import Config logger = logging.getLogger("bbq.client") + class BBQClient: - """ - Client for querying the running ColPali RAG indexing server, generating Gemini multimodal answers, - and fetching/saving page images. - """ + """Client for querying the BBQ RAG server and multimodal answer generation.""" def __init__( self, server_url: str = "http://localhost:8000", - config: Optional[Any] = None, + config: Optional[Config] = None, + without_logo: bool = False, ) -> None: self.server_url: str = server_url.rstrip("/") - self.config = config + self.config: Optional[Config] = config + self.without_logo: bool = without_logo + + def print_logo( + self, + name: Optional[str] = None, + tag: str = "CLIENT", + server: Optional[Any] = None, + top_k: Optional[int] = 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.server_url + print_bbq(name=name, tag=tag, server=srv, top_k=top_k) def get_status(self) -> Dict[str, Any]: """Retrieves server status information.""" - logger.info(f"Checking server status at {self.server_url}/status...") - start_time = time.time() - response = requests.get(f"{self.server_url}/status", timeout=10) - elapsed = time.time() - start_time - response.raise_for_status() - logger.info(f"Server status retrieved in {elapsed:.3f}s (HTTP {response.status_code})") - return response.json() + resp = requests.get(f"{self.server_url}/status", timeout=10) + resp.raise_for_status() + return resp.json() def list_documents(self) -> List[Dict[str, Any]]: """Retrieves indexed document metadata records from server.""" - logger.info(f"Fetching document list from {self.server_url}/documents...") - start_time = time.time() - response = requests.get(f"{self.server_url}/documents", timeout=10) - elapsed = time.time() - start_time - response.raise_for_status() - docs = response.json().get("documents", []) - logger.info(f"Retrieved {len(docs)} document record(s) in {elapsed:.3f}s") - return docs - - def query(self, query_text: str, top_k: int = 3) -> List[Dict[str, Any]]: - """ - Sends a retrieval query to the server and returns the top matching PDF page results. - - Args: - query_text (str): Search text query. - top_k (int): Number of top matching PDF pages to retrieve (default: 3). + resp = requests.get(f"{self.server_url}/documents", timeout=10) + resp.raise_for_status() + return resp.json().get("documents", []) + + def query(self, query_text: str, top_k: int = 10) -> List[Dict[str, Any]]: + """Sends a retrieval query to the server and returns the top matching results.""" + resp = requests.post( + f"{self.server_url}/query", + json={"query": query_text, "top_k": top_k}, + timeout=30, + ) + resp.raise_for_status() + return resp.json().get("results", []) - Returns: - List[Dict[str, Any]]: List of matching results containing score, file_path, page_number, etc. - """ - logger.info(f"Sending query request to {self.server_url}/query: '{query_text}' (top_k={top_k})") - payload = {"query": query_text, "top_k": top_k} - start_time = time.time() - response = requests.post(f"{self.server_url}/query", json=payload, timeout=30) - elapsed = time.time() - start_time - response.raise_for_status() - data = response.json() - results = data.get("results", []) - logger.info(f"Query completed in {elapsed:.3f}s (HTTP {response.status_code}): received {len(results)} match(es)") - return results + def get_page_image( + self, + file_path: str, + page_number: int = 1, + dpi: int = 150, + save_path: Optional[str] = None, + ) -> Image.Image: + """Fetches rendered PNG page image from the server and optionally saves it.""" + resp = requests.get( + f"{self.server_url}/page_image", + params={"file_path": file_path, "page_number": page_number, "dpi": dpi}, + timeout=15, + ) + resp.raise_for_status() + img = Image.open(io.BytesIO(resp.content)).convert("RGB") + if save_path: + os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True) + img.save(save_path, format="PNG") + return img - def query_and_answer( + def generate_answer_from_results( self, query_text: str, - top_k: int = 3, + results: List[Dict[str, Any]], gemini_api_key: Optional[str] = None, - gemini_model: str = "gemini-3.6-flash", + gemini_model: Optional[str] = None, save_images: bool = False, images_output_dir: Optional[str] = None, ) -> Dict[str, Any]: """ - Retrieves top_k relevant pages (default: 3) and uses Google Gemini to generate - a grounded answer from the page images. If no API key is given or Gemini fails, - gracefully returns the retrieved pages. - - Args: - query_text (str): The search query / question. - top_k (int): Number of top pages to retrieve (default: 3). - gemini_api_key (Optional[str]): Gemini API key. If omitted, uses GEMINI_API_KEY env. - gemini_model (str): Gemini model name (default: 'gemini-3.6-flash' free-tier default). - save_images (bool): Whether to save retrieved page images locally. - images_output_dir (Optional[str]): Directory to save page images if save_images is True. - - Returns: - Dict[str, Any]: Dictionary containing 'query', 'answer', 'sources', 'status', and 'engine'. + Fetches page images for retrieved results and generates a grounded multimodal answer with Gemini. """ - from bbq.src.client.gemini import GeminiClient - - # 1. Retrieve top-k matching pages - results = self.query(query_text=query_text, top_k=top_k) - if not results: - return { - "query": query_text, - "answer": None, - "sources": [], - "status": "no_results", - "engine": "none", - "message": "No matching document pages found for query.", - } - - # 2. Fetch page images for top retrieved matches page_images: List[Image.Image] = [] for i, res in enumerate(results, 1): + save_path = ( + os.path.join( + images_output_dir or ".", + f"retrieved_rank_{i}_page_{res['page_number']}.png", + ) + if save_images + else None + ) img = None - save_path = None - if save_images: - out_dir = images_output_dir or "." - os.makedirs(out_dir, exist_ok=True) - save_path = os.path.join(out_dir, f"retrieved_rank_{i}_page_{res['page_number']}.png") - try: img = self.get_page_image( - file_path=res["file_path"], - page_number=res["page_number"], - save_path=save_path, + res["file_path"], res["page_number"], save_path=save_path ) - except Exception as e: - logger.warning(f"Could not fetch server page image for {res['file_path']} page {res['page_number']}: {e}") - # Fallback to local rendering if file exists locally - if os.path.exists(res["file_path"]): + except Exception: + if os.path.exists(res.get("file_path", "")): try: img = get_local_pdf_page_image( - file_path=res["file_path"], - page_number=res["page_number"], - save_path=save_path, + res["file_path"], res["page_number"], save_path=save_path ) - except Exception as local_err: - logger.warning(f"Local PDF render failed: {local_err}") - - if img is not None: + except Exception: + pass + if img: page_images.append(img) res["image_available"] = True if save_path: @@ -146,110 +129,124 @@ def query_and_answer( else: res["image_available"] = False - # Pull defaults from config if available - if self.config is not None: - if gemini_api_key is None and hasattr(self.config, "gemini_api_key"): - gemini_api_key = self.config.gemini_api_key - if gemini_model == "gemini-3.6-flash" and hasattr(self.config, "gemini_model"): - gemini_model = self.config.gemini_model + if self.config: + gemini_api_key = gemini_api_key or getattr( + self.config, "gemini_api_key", None + ) + gemini_model = gemini_model or getattr( + self.config, "gemini_model", "gemini-3.6-flash" + ) - # 3. Check Gemini API client availability - gemini = GeminiClient(api_key=gemini_api_key, model=gemini_model) + from bbq.src.client.gemini import GeminiClient + gemini = GeminiClient( + api_key=gemini_api_key, model=gemini_model or "gemini-3.6-flash" + ) + + answer = None + fallback_reason = None if not gemini.is_available(): - logger.info("No Gemini API key provided. Returning retrieved pages only.") + fallback_reason = ( + "No Gemini API key provided. Returning retrieved pages of the book." + ) + elif not page_images: + fallback_reason = ( + "Could not load page images for Gemini multimodal analysis." + ) + else: + answer = gemini.generate_answer( + query=query_text, images=page_images, page_metadata=results + ) + if not answer: + fallback_reason = "Gemini API request failed or rate limited. Returning retrieved pages of the book." + + if answer: return { "query": query_text, - "answer": None, + "answer": answer, "sources": results, - "status": "fallback_pages_only", - "engine": "retrieval_only", - "fallback_reason": "No Gemini API key provided. Returning retrieved pages of the book.", + "status": "success", + "engine": f"gemini ({gemini.model})", } - - # 4. Attempt Gemini Multimodal Generation - if page_images: - answer = gemini.generate_answer( - query=query_text, - images=page_images, - page_metadata=results, - ) - if answer: - return { - "query": query_text, - "answer": answer, - "sources": results, - "status": "success", - "engine": f"gemini ({gemini.model})", - } - else: - logger.warning("Gemini API call failed. Falling back to returning retrieved pages.") - return { - "query": query_text, - "answer": None, - "sources": results, - "status": "fallback_pages_only", - "engine": "retrieval_only", - "fallback_reason": "Gemini API request failed or rate limited. Returning retrieved pages of the book.", - } - return { "query": query_text, "answer": None, "sources": results, "status": "fallback_pages_only", "engine": "retrieval_only", - "fallback_reason": "Could not load page images for Gemini multimodal analysis.", + "fallback_reason": fallback_reason, } - def get_page_image( + def query_and_answer( self, - file_path: str, - page_number: int = 1, - dpi: int = 150, - save_path: Optional[str] = None, - ) -> Image.Image: + query_text: str, + top_k: int = 10, + use_llm: bool = True, + gemini_api_key: Optional[str] = None, + gemini_model: Optional[str] = None, + save_images: bool = False, + images_output_dir: Optional[str] = None, + without_logo: Optional[bool] = None, + ) -> Dict[str, Any]: """ - Fetches the rendered PNG image for a given PDF file and page number from the server. - Optionally saves it to save_path. + Retrieves top_k pages. If use_llm is True, passes retrieved pages to Gemini LLM. + Otherwise, returns retrieved matches directly without calling LLM. """ - logger.info(f"Requesting page image from server: '{file_path}' (page={page_number}, dpi={dpi})") - params = {"file_path": file_path, "page_number": page_number, "dpi": dpi} - start_time = time.time() - response = requests.get(f"{self.server_url}/page_image", params=params, timeout=15) - elapsed = time.time() - start_time - response.raise_for_status() + suppress_logo = self.without_logo if without_logo is None else without_logo + if not suppress_logo: + self.print_logo(top_k=top_k) - img = Image.open(io.BytesIO(response.content)).convert("RGB") - logger.info(f"Page image received in {elapsed:.3f}s ({len(response.content)} bytes)") + results = self.query(query_text=query_text, top_k=top_k) + if not results: + return { + "query": query_text, + "answer": None, + "sources": [], + "status": "no_results", + "engine": "none", + "message": "No matching document pages found for query.", + } - if save_path: - out_dir = os.path.dirname(save_path) - if out_dir: - os.makedirs(out_dir, exist_ok=True) - img.save(save_path, format="PNG") - logger.info(f"Saved PDF page image to disk: {save_path}") - print(f"Saved PDF page image to {save_path}") + if not use_llm: + return { + "query": query_text, + "answer": None, + "sources": results, + "status": "retrieval_only", + "engine": "none", + } - return img + return self.generate_answer_from_results( + query_text=query_text, + results=results, + gemini_api_key=gemini_api_key, + gemini_model=gemini_model, + save_images=save_images, + images_output_dir=images_output_dir, + ) def get_local_pdf_page_image( - file_path: str, page_number: int = 1, dpi: int = 150, save_path: Optional[str] = None + file_path: str, + page_number: int = 1, + dpi: int = 150, + save_path: Optional[str] = None, ) -> Image.Image: - """ - Renders a PDF page image directly locally without needing an HTTP server. - """ + """Renders a PDF page image directly locally.""" from bbq.src.utils.pdf_utils import extract_single_pdf_page_image - logger.info(f"Rendering local PDF page image: '{file_path}' (page={page_number}, dpi={dpi})") - img = extract_single_pdf_page_image(pdf_filepath=file_path, page_number=page_number, dpi=dpi) + img = extract_single_pdf_page_image( + pdf_filepath=file_path, page_number=page_number, dpi=dpi + ) if save_path: - out_dir = os.path.dirname(save_path) - if out_dir: - os.makedirs(out_dir, exist_ok=True) + os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True) img.save(save_path, format="PNG") - logger.info(f"Saved local PDF page image to {save_path}") - print(f"Saved local PDF page image to {save_path}") - return img + + +if __name__ == "__main__": + from bbq.src.main import build_client_main_parser, run_client_query_command + + parser = build_client_main_parser() + args = parser.parse_args() + run_client_query_command(args) diff --git a/bbq/src/client/gemini.py b/bbq/src/client/gemini.py index 6619ba9..363af26 100644 --- a/bbq/src/client/gemini.py +++ b/bbq/src/client/gemini.py @@ -1,9 +1,9 @@ """ Google Gemini Multimodal API Client for Visual Document RAG. -Supports free-tier defaults (gemini-3.6-flash / gemini-2.0-flash) with multimodal image inputs. -Gracefully handles API failures and missing keys with zero-crash fallbacks. """ +from __future__ import annotations + import os import io import base64 @@ -13,14 +13,12 @@ from PIL import Image logger = logging.getLogger("bbq.client.gemini") - DEFAULT_GEMINI_MODEL = "gemini-3.6-flash" GEMINI_API_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/models" + class GeminiClient: - """ - Multimodal client for sending document page images and user queries to Google Gemini. - """ + """Multimodal client for document page question answering via Google Gemini.""" def __init__( self, @@ -37,23 +35,28 @@ def __init__( self.timeout: int = timeout def is_available(self) -> bool: - """Returns True if a valid API key is configured.""" return bool(self.api_key and self.api_key.strip()) + @classmethod + def from_config(cls, config: Any, timeout: int = 90) -> GeminiClient: + return cls( + api_key=getattr(config, "gemini_api_key", None), + model=getattr(config, "gemini_model", DEFAULT_GEMINI_MODEL), + timeout=timeout, + ) + @staticmethod - def _image_to_base64_part(image: Image.Image, format: str = "JPEG", quality: int = 85) -> Dict[str, Any]: - """Converts a PIL Image into Gemini's inlineData base64 part.""" - buffered = io.BytesIO() - # Convert RGBA to RGB for JPEG compatibility + def _image_to_base64_part( + image: Image.Image, format: str = "JPEG" + ) -> Dict[str, Any]: + buf = io.BytesIO() if image.mode in ("RGBA", "P"): image = image.convert("RGB") - image.save(buffered, format=format, quality=quality) - img_bytes = buffered.getvalue() - b64_str = base64.b64encode(img_bytes).decode("utf-8") + image.save(buf, format=format, quality=85) return { "inlineData": { "mimeType": f"image/{format.lower()}", - "data": b64_str, + "data": base64.b64encode(buf.getvalue()).decode("utf-8"), } } @@ -64,102 +67,51 @@ def generate_answer( page_metadata: Optional[List[Dict[str, Any]]] = None, system_instruction: Optional[str] = None, ) -> Optional[str]: - """ - Generates an answer to the query using the retrieved document page images. - - Args: - query (str): User's natural language question. - images (List[Image.Image]): List of top retrieved PDF page images. - page_metadata (Optional[List[Dict[str, Any]]]): Metadata for each page (filename, page_number). - system_instruction (Optional[str]): Custom system instruction. - - Returns: - Optional[str]: Generated text answer, or None if the API key is missing or call failed. - """ - if not self.is_available(): - logger.warning("No Gemini API key provided. Skipping Gemini answer generation.") - return None - - if not images: - logger.warning("No document images provided to Gemini.") + if not self.is_available() or not images: return None - endpoint = f"{GEMINI_API_BASE_URL}/{self.model}:generateContent?key={self.api_key}" - - # Construct multimodal parts - parts: List[Dict[str, Any]] = [] - - default_sys_prompt = ( + prompt = system_instruction or ( "You are an expert visual document comprehension assistant. " - "You are provided with images of relevant document pages from a book/document. " - "Carefully analyze the text, diagrams, tables, and visuals in the provided page images to answer the user's question accurately. " - "Always cite the relevant page numbers in your answer when referencing information." + "Carefully analyze the text, diagrams, tables, and visuals in the provided page images to answer accurately. " + "Always cite the relevant page numbers in your answer." ) - instruction_text = system_instruction or default_sys_prompt - parts.append({"text": instruction_text}) + parts: List[Dict[str, Any]] = [{"text": prompt}] - # Add each retrieved page image with contextual label for idx, img in enumerate(images): - meta_label = f"[Document Page {idx + 1}" - if page_metadata and idx < len(page_metadata): - meta = page_metadata[idx] - fname = meta.get("filename") or os.path.basename(meta.get("file_path", "")) - pnum = meta.get("page_number", idx + 1) - meta_label += f" | File: {fname} | Page: {pnum}" - meta_label += "]" - - parts.append({"text": meta_label}) + meta = ( + page_metadata[idx] if page_metadata and idx < len(page_metadata) else {} + ) + fname = meta.get("filename") or os.path.basename(meta.get("file_path", "")) + pnum = meta.get("page_number", idx + 1) + parts.append( + {"text": f"[Document Page {idx + 1} | File: {fname} | Page: {pnum}]"} + ) parts.append(self._image_to_base64_part(img)) - # Add the user query - parts.append({"text": f"User Question: {query}\n\nPlease provide a clear, concise, and grounded answer based solely on the document pages above."}) + parts.append( + { + "text": f"User Question: {query}\n\nPlease provide a clear, concise, and grounded answer based solely on the document pages above." + } + ) payload = { - "contents": [ - { - "role": "user", - "parts": parts, - } - ], - "generationConfig": { - "temperature": 0.2, - "maxOutputTokens": 1024, - }, + "contents": [{"role": "user", "parts": parts}], + "generationConfig": {"temperature": 0.2, "maxOutputTokens": 8192}, } try: - logger.info(f"Sending multimodal request to Gemini ({self.model}) with {len(images)} page image(s)...") - response = requests.post( - endpoint, + resp = requests.post( + f"{GEMINI_API_BASE_URL}/{self.model}:generateContent?key={self.api_key}", json=payload, headers={"Content-Type": "application/json"}, timeout=self.timeout, ) - - if response.status_code != 200: - logger.warning( - f"Gemini API returned HTTP {response.status_code}: {response.text[:300]}" - ) - return None - - resp_data = response.json() - candidates = resp_data.get("candidates", []) - if not candidates: - logger.warning("Gemini API returned no response candidates.") - return None - - content = candidates[0].get("content", {}) - resp_parts = content.get("parts", []) - if not resp_parts: - return None - - answer_text = "".join(part.get("text", "") for part in resp_parts).strip() - logger.info(f"Successfully generated answer from Gemini ({len(answer_text)} chars).") - return answer_text - - except requests.exceptions.RequestException as req_err: - logger.warning(f"Network error communicating with Gemini API: {req_err}") - return None + if resp.status_code == 200: + candidates = resp.json().get("candidates", []) + if candidates: + parts = candidates[0].get("content", {}).get("parts", []) + return "".join(p.get("text", "") for p in parts).strip() or None + logger.warning(f"Gemini API error {resp.status_code}: {resp.text[:200]}") except Exception as err: - logger.warning(f"Unexpected error during Gemini generation: {err}") - return None + logger.warning(f"Gemini request failed: {err}") + return None diff --git a/bbq/src/common/__init__.py b/bbq/src/common/__init__.py index 7752581..e42778c 100644 --- a/bbq/src/common/__init__.py +++ b/bbq/src/common/__init__.py @@ -1,6 +1,26 @@ -from typing import Any +from typing import TYPE_CHECKING, Any +from bbq.src.common.version import __version__, _read_project_version + +if TYPE_CHECKING: + from bbq.src.common.base import ( + BaseModel, + BaseModelLoader, + BaseProcessor, + BaseEngineWrapper, + ) + from bbq.src.common.errors import ( + BaseBBQEngineException, + ConfigFNFWarning, + ConfigParseError, + BaseModelConfigLoadError, + BaseModelInstantiateError, + LoRAAdapterLoadError, + ProcessorLoadError, + ) __all__ = [ + "__version__", + "_read_project_version", "BaseModel", "BaseModelLoader", "BaseProcessor", @@ -14,9 +34,11 @@ "ProcessorLoadError", ] + def __getattr__(name: str) -> Any: if name in ("BaseModel", "BaseModelLoader", "BaseProcessor", "BaseEngineWrapper"): import bbq.src.common.base as base_mod + return getattr(base_mod, name) elif name in ( "BaseBBQEngineException", @@ -28,5 +50,10 @@ def __getattr__(name: str) -> Any: "ProcessorLoadError", ): import bbq.src.common.errors as err_mod + return getattr(err_mod, name) + elif name in ("__version__", "_read_project_version"): + import bbq.src.common.version as ver_mod + + return getattr(ver_mod, name) raise AttributeError(f"module '{__name__}' has no attribute '{name}'") diff --git a/bbq/src/common/base.py b/bbq/src/common/base.py index 5c8eb29..f02e49d 100644 --- a/bbq/src/common/base.py +++ b/bbq/src/common/base.py @@ -2,10 +2,12 @@ from typing import Tuple, List, Any, Mapping from transformers import PreTrainedModel + class BaseModel(PreTrainedModel, ABC): """ Abstract base class for models in BBQ RAG. """ + @classmethod @abstractmethod def supports_model(cls, base_model_id: str) -> bool: @@ -21,42 +23,33 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: """ pass + class BaseModelLoader(ABC): @classmethod @abstractmethod - def load_model_and_processor( - cls, config_input: Any - ) -> Tuple[Any, Any]: + def load_model_and_processor(cls, config_input: Any) -> Tuple[Any, Any]: pass + class BaseProcessor(ABC, PreTrainedModel): @abstractmethod - def process_images( - self, images: List[Any], **kwargs: Any - ) -> Mapping[str, Any]: + def process_images(self, images: List[Any], **kwargs: Any) -> Mapping[str, Any]: pass @abstractmethod - def process_texts( - self, texts: List[str], **kwargs: Any - ) -> Mapping[str, Any]: + def process_texts(self, texts: List[str], **kwargs: Any) -> Mapping[str, Any]: pass @abstractmethod - def score( - self, qs: Any, ps: Any, batch_size: int = 128 - ) -> Any: + def score(self, qs: Any, ps: Any, batch_size: int = 128) -> Any: pass + class BaseEngineWrapper(ABC): @abstractmethod - def encode_multimodal_document_images( - self, images: List[Any] - ) -> Any: + def encode_multimodal_document_images(self, images: List[Any]) -> Any: pass @abstractmethod - def encode_query_text_inputs( - self, texts: List[str] - ) -> Any: + def encode_query_text_inputs(self, texts: List[str]) -> Any: pass diff --git a/bbq/src/common/errors.py b/bbq/src/common/errors.py index 65f91ea..7de3403 100644 --- a/bbq/src/common/errors.py +++ b/bbq/src/common/errors.py @@ -2,64 +2,87 @@ class BaseBBQEngineException(Exception): """ Base class for all exceptions raised by the BBQ Engine. """ + def __init__(self, message: str = "An exception occurred in BBQ Engine.") -> None: super().__init__(message) self.message: str = message + class ConfigFNFWarning(UserWarning): """ Warning raised when a configuration file is not found. """ + def __init__(self, message: str = "Configuration file was not found.") -> None: super().__init__(message) self.message: str = message + class ConfigParseError(BaseBBQEngineException): """ Exception raised when a configuration file cannot be parsed. """ - def __init__(self, message: str = "Failed to parse YAML configuration file.") -> None: + + def __init__( + self, message: str = "Failed to parse YAML configuration file." + ) -> None: super().__init__(message) + class BaseModelConfigLoadError(BaseBBQEngineException): """ Exception raised when the base model configuration fails to load. """ - def __init__(self, message: str = "Failed to load explicit base configuration.") -> None: + + def __init__( + self, message: str = "Failed to load explicit base configuration." + ) -> None: super().__init__(message) + class BaseModelInstantiateError(BaseBBQEngineException): """ Exception raised when the base model fails to instantiate. """ + def __init__(self, message: str = "Failed to instantiate base model.") -> None: super().__init__(message) + class LoRAAdapterLoadError(BaseBBQEngineException): """ Exception raised when LoRA adapter weights fail to load or apply. """ - def __init__(self, message: str = "Failed to load and apply LoRA adapter weights.") -> None: + + def __init__( + self, message: str = "Failed to load and apply LoRA adapter weights." + ) -> None: super().__init__(message) + class ProcessorLoadError(BaseBBQEngineException): """ Exception raised when the model processor fails to load. """ + def __init__(self, message: str = "Failed to load processor.") -> None: super().__init__(message) + class ModelNotFound(UserWarning): """ Warning raised when the model is not found. Especially in Model Registry """ + def __init__(self, message: str = "Failed to load model.") -> None: super().__init__(message) self.message: str = message + class PdfNotFoundError(BaseBBQEngineException): """ Exception raised when a PDF file is not found. """ + def __init__(self, message: str = "PDF file was not found.") -> None: super().__init__(message) diff --git a/bbq/src/common/version.py b/bbq/src/common/version.py new file mode 100644 index 0000000..f20d3b6 --- /dev/null +++ b/bbq/src/common/version.py @@ -0,0 +1,41 @@ +import tomllib +from pathlib import Path + + +def _read_project_version() -> str: + bases = [ + Path.cwd(), + Path(__file__).resolve().parents[2], + Path(__file__).resolve().parents[3], + ] + for base in bases: + pyproject = base / "pyproject.toml" + if pyproject.exists(): + try: + with open(pyproject, "rb") as f: + v = tomllib.load(f).get("project", {}).get("version") + if v: + return str(v) + except Exception: + pass + cargo = base / "Cargo.toml" + if cargo.exists(): + try: + with open(cargo, "rb") as f: + cv = tomllib.load(f).get("package", {}).get("version") + if cv: + return str(cv) + except Exception: + pass + try: + from importlib.metadata import version + + return version("bbq") + except Exception: + return "0.1.4-alpha" + + +# Global package version directly read from pyproject.toml / Cargo.toml +__version__ = _read_project_version() + +__all__ = ["_read_project_version", "__version__"] diff --git a/bbq/src/config/__init__.py b/bbq/src/config/__init__.py index 1f7e50b..96269da 100644 --- a/bbq/src/config/__init__.py +++ b/bbq/src/config/__init__.py @@ -1,9 +1,4 @@ -from .config import ModelConfigWrapper, get_system_cache_dir, load_configuration_from_yaml_file, determine_target_torch_device, resolve_torch_data_type +from .config import Config + +__all__ = ["Config"] -__all__ = [ - "ModelConfigWrapper", - "get_system_cache_dir", - "load_configuration_from_yaml_file", - "determine_target_torch_device", - "resolve_torch_data_type", -] diff --git a/bbq/src/config/config.py b/bbq/src/config/config.py index ff0eb3a..0ac9e5c 100644 --- a/bbq/src/config/config.py +++ b/bbq/src/config/config.py @@ -1,79 +1,62 @@ +from __future__ import annotations + import os -import yaml import warnings -import traceback -from typing import Dict, Any, Optional -from bbq.src.common.errors import ConfigFNFWarning, ConfigParseError +from dataclasses import asdict, dataclass, fields +from typing import Any, Dict, Optional -def get_system_cache_dir(subfolder: str = "") -> str: - """ - Returns the system default user cache directory for bbq. - Respects XDG_CACHE_HOME if set, defaulting to ~/.cache/bbq. - """ - cache_base = os.environ.get("XDG_CACHE_HOME") - if not cache_base: - cache_base = os.path.expanduser("~/.cache") - bbq_cache_dir = os.path.join(cache_base, "bbq") - if subfolder: - return os.path.join(bbq_cache_dir, subfolder) - return bbq_cache_dir +import yaml + +from bbq.src.common.errors import ConfigFNFWarning, ConfigParseError +from bbq.src.utils.futils import get_system_cache_dir -class ModelConfigWrapper: +@dataclass +class Config: """ - Wrapper class for model configuration. - - Args: - base_model_id (str): The base model ID. - lora_adapter_id (str): The LoRA adapter ID. - embedding_dim (int): The embedding dimension. - device (str): The device to use. - torch_dtype (str): The data type. - mask_non_image_embeddings (bool): Whether to mask non-image embeddings. - visual_prompt_command (str): The visual prompt command. + BBQ configuration class that encapsulates various settings for the application. """ - def __init__( - self, - base_model_id: str = "HuggingFaceTB/SmolVLM-256M-Instruct", - lora_adapter_id: str = "vidore/colSmol-256M", - embedding_dim: int = 128, - device: str = "auto", - torch_dtype: str = "bfloat16", - mask_non_image_embeddings: bool = False, - visual_prompt_command: str = "What is written on the image.", - watch_folder_path: str = "data/watch", - embeddings_output_path: Optional[str] = None, - sqlite_db_path: Optional[str] = None, - pdf_render_dpi: int = 150, - gemini_api_key: Optional[str] = None, - gemini_model: str = "gemini-3.6-flash", - rag_top_k: int = 3, - ) -> None: - self.base_model_id: str = base_model_id - self.lora_adapter_id: str = lora_adapter_id - self.embedding_dim: int = embedding_dim - self.device: str = device - self.torch_dtype: str = torch_dtype - self.mask_non_image_embeddings: bool = mask_non_image_embeddings - self.visual_prompt_command: str = visual_prompt_command - self.watch_folder_path: str = watch_folder_path - self.embeddings_output_path: str = embeddings_output_path or get_system_cache_dir("embeddings") - self.sqlite_db_path: str = sqlite_db_path or get_system_cache_dir("tracker.db") - self.pdf_render_dpi: int = pdf_render_dpi - self.gemini_api_key: Optional[str] = ( - gemini_api_key - or os.environ.get("GEMINI_API_KEY") - or os.environ.get("GOOGLE_API_KEY") - ) - self.gemini_model: str = gemini_model - self.rag_top_k: int = rag_top_k - def get_gemini_client(self): - """ - Instantiates and returns a GeminiClient configured with this config's settings. - """ - from bbq.src.client.gemini import GeminiClient - return GeminiClient(api_key=self.gemini_api_key, model=self.gemini_model) + base_model_id: str = "HuggingFaceTB/SmolVLM-256M-Instruct" + lora_adapter_id: str = "vidore/colSmol-256M" + embedding_dim: int = 128 + device: str = "auto" + torch_dtype: str = "bfloat16" + mask_non_image_embeddings: bool = False + visual_prompt_command: str = "Explaing the image." + watch_folder_path: str = "data/watch" + embeddings_output_path: Optional[str] = None + sqlite_db_path: Optional[str] = None + pdf_render_dpi: int = 150 + gemini_api_key: Optional[str] = None + gemini_model: str = "gemini-3.6-flash" + rag_top_k: int = 3 + quantization: Optional[str] = "f32" + + def __post_init__(self) -> None: + if not self.base_model_id: + raise ValueError("base_model_id cannot be empty in configuration.") + if self.embedding_dim is not None: + self.embedding_dim = int(self.embedding_dim) + if self.pdf_render_dpi is not None: + self.pdf_render_dpi = int(self.pdf_render_dpi) + if self.rag_top_k is not None: + self.rag_top_k = int(self.rag_top_k) + + if not self.embeddings_output_path: + self.embeddings_output_path = get_system_cache_dir("embeddings") + if not self.sqlite_db_path: + self.sqlite_db_path = get_system_cache_dir("tracker.db") + + if self.quantization: + self.quantization = str(self.quantization).strip().lower() + + if self.gemini_api_key: + self.gemini_api_key = str(self.gemini_api_key).strip() or None + if not self.gemini_api_key: + self.gemini_api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get( + "GOOGLE_API_KEY" + ) def format_visual_prompt_prefix(self) -> str: """ @@ -85,157 +68,61 @@ def format_visual_prompt_prefix(self) -> str: def convert_to_dictionary_format(self) -> Dict[str, Any]: """ - Converts the configuration to a dictionary format. + Converts the configuration to a dictionary format without exposing secret API keys. """ - return { - "base_model_id": self.base_model_id, - "lora_adapter_id": self.lora_adapter_id, - "embedding_dim": self.embedding_dim, - "device": self.device, - "torch_dtype": self.torch_dtype, - "mask_non_image_embeddings": self.mask_non_image_embeddings, - "visual_prompt_command": self.visual_prompt_command, - "visual_prompt_prefix": self.format_visual_prompt_prefix(), - "watch_folder_path": self.watch_folder_path, - "embeddings_output_path": self.embeddings_output_path, - "sqlite_db_path": self.sqlite_db_path, - "pdf_render_dpi": self.pdf_render_dpi, - "gemini_api_key_configured": bool(self.gemini_api_key and self.gemini_api_key.strip()), - "gemini_model": self.gemini_model, - "rag_top_k": self.rag_top_k, - } - -def load_configuration_from_yaml_file( - config_filepath: str = "config.yaml", -) -> ModelConfigWrapper: - """ - Load configuration from a YAML file. - - Args: - config_filepath (str): The path to the YAML configuration file. - - Returns: - ModelConfigWrapper: An instance of ModelConfigWrapper containing the loaded configuration. - """ - - # Check if the configuration file exists - - if not os.path.exists(config_filepath): - warning_msg: str = ( - f"Configuration file not found at path: '{config_filepath}'. " - "Proceeding with default settings." + data = asdict(self) + data.pop("gemini_api_key", None) + data["gemini_api_key_configured"] = bool( + self.gemini_api_key and self.gemini_api_key.strip() ) - traceback.print_stack(limit=3) - warnings.warn(warning_msg, category=ConfigFNFWarning, stacklevel=2) - return ModelConfigWrapper() - - try: - with open(config_filepath, "r", encoding="utf-8") as file_stream: - parsed_yaml_data: Dict[str, Any] = yaml.safe_load(file_stream) or {} - except Exception as exception_instance: - traceback.print_exc() - raise ConfigParseError( - f"Failed to parse YAML configuration file: {exception_instance}" - ) from exception_instance - - base_model_id: str = str( - parsed_yaml_data.get("base_model_id", "HuggingFaceTB/SmolVLM-256M-Instruct") - ) - - lora_adapter_id: str = str( - parsed_yaml_data.get("lora_adapter_id", "vidore/colSmol-256M") - ) - - embedding_dim: int = int(parsed_yaml_data.get("embedding_dim", 128)) - - device: str = str(parsed_yaml_data.get("device", "auto")) - - torch_dtype: str = str(parsed_yaml_data.get("torch_dtype", "bfloat16")) - - mask_non_image_embeddings: bool = bool( - parsed_yaml_data.get("mask_non_image_embeddings", False) - ) - - visual_prompt_command: str = str( - parsed_yaml_data.get("visual_prompt_command", "What is written on the image.") - ) - - watch_folder_path: str = str( - parsed_yaml_data.get("watch_folder_path", "data/watch") - ) - - embeddings_output_path: str = str( - parsed_yaml_data.get("embeddings_output_path", get_system_cache_dir("embeddings")) - ) - - sqlite_db_path: str = str( - parsed_yaml_data.get("sqlite_db_path", get_system_cache_dir("tracker.db")) - ) - - pdf_render_dpi: int = int(parsed_yaml_data.get("pdf_render_dpi", 150)) - - gemini_api_key: Optional[str] = parsed_yaml_data.get("gemini_api_key") - if gemini_api_key is not None: - gemini_api_key = str(gemini_api_key).strip() or None - else: - gemini_api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY") - - gemini_model: str = str(parsed_yaml_data.get("gemini_model", "gemini-3.6-flash")) - rag_top_k: int = int(parsed_yaml_data.get("rag_top_k", 3)) - - if not base_model_id: - raise ValueError("base_model_id cannot be empty in configuration.") - - return ModelConfigWrapper( - base_model_id=base_model_id, - lora_adapter_id=lora_adapter_id, - embedding_dim=embedding_dim, - device=device, - torch_dtype=torch_dtype, - mask_non_image_embeddings=mask_non_image_embeddings, - visual_prompt_command=visual_prompt_command, - watch_folder_path=watch_folder_path, - embeddings_output_path=embeddings_output_path, - sqlite_db_path=sqlite_db_path, - pdf_render_dpi=pdf_render_dpi, - gemini_api_key=gemini_api_key, - gemini_model=gemini_model, - rag_top_k=rag_top_k, - ) - -def determine_target_torch_device(device_preference: str = "auto") -> str: - """ - Determine the target torch device based on the provided device preference. - - Args: - device_preference (str): The device preference, either "auto" or a specific device name. - - Returns: - str: The target torch device. - """ - import torch - if device_preference == "auto": - return "cuda" if torch.cuda.is_available() else "cpu" - return device_preference - -def resolve_torch_data_type( - dtype_name: str = "bfloat16", target_device: str = "cpu" -) -> Any: - """ - Resolve the torch data type based on the provided dtype name and target device. + data["visual_prompt_prefix"] = self.format_visual_prompt_prefix() + return data - Args: - dtype_name (str): The data type name, either "bfloat16" or "float16". - target_device (str): The target device, either "cpu" or "cuda". - - Returns: - torch.dtype: The resolved torch data type. - """ - import torch - if target_device == "cpu": - return torch.float32 - if dtype_name == "bfloat16": - return torch.bfloat16 - elif dtype_name == "float16": - return torch.float16 - return torch.float32 + @classmethod + def from_yaml(cls, config_filepath: str = "config.yaml") -> Config: + """ + Load configuration from a YAML file. + """ + resolved_path = config_filepath + if not os.path.exists(resolved_path): + candidates = [ + os.path.join(os.getcwd(), config_filepath), + os.path.join( + os.path.dirname( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) + ), + config_filepath, + ), + ] + for candidate in candidates: + if os.path.exists(candidate): + resolved_path = candidate + break + + if not os.path.exists(resolved_path): + warnings.warn( + f"Configuration file not found at path: '{config_filepath}'. Proceeding with default settings.", + category=ConfigFNFWarning, + stacklevel=2, + ) + return cls() + + config_filepath = resolved_path + + try: + with open(config_filepath, "r", encoding="utf-8") as file_stream: + parsed_yaml_data: Dict[str, Any] = yaml.safe_load(file_stream) or {} + except Exception as exception_instance: + raise ConfigParseError( + f"Failed to parse YAML configuration file: {exception_instance}" + ) from exception_instance + + valid_keys = {f.name for f in fields(cls)} + filtered_kwargs = { + k: v + for k, v in parsed_yaml_data.items() + if k in valid_keys and v is not None + } + return cls(**filtered_kwargs) diff --git a/bbq/src/main.py b/bbq/src/main.py index 7ee02f6..c961af3 100644 --- a/bbq/src/main.py +++ b/bbq/src/main.py @@ -1,17 +1,21 @@ import sys -import os import argparse import logging -from pathlib import Path - -# Ensure workspace root is in sys.path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) +from bbq.src.common.version import __version__, _read_project_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)) + + config = Config.from_yaml(config_filepath=args.config) start_document_indexing_server( + config=config, config_filepath=args.config, host=args.host, port=args.port, @@ -20,76 +24,157 @@ def run_server_command(args: argparse.Namespace) -> None: def run_client_query_command(args: argparse.Namespace) -> None: """Launches a client query against the running server with optional Gemini multimodal RAG.""" - from bbq.src.client import BBQClient - from bbq.src.config import load_configuration_from_yaml_file + # Support checking server status or documents via client query command + if getattr(args, "status", False): + run_status_command(args) + return + if getattr(args, "documents", False): + run_documents_command(args) + return - log_level = logging.INFO if args.verbose else logging.WARNING - logging.basicConfig(level=log_level, format="[%(asctime)s] [%(levelname)s] [%(name)s]: %(message)s") + from rich.console import Console + from bbq.src.client import BBQClient + from bbq.src.config import Config + from bbq.src.terminal import ( + print_bbq, + render_query_results_rich, + render_llm_answer_rich, + print_interactive_help, + interactive_page_picker, + ) # Load configuration from config.py / config.yaml - config = load_configuration_from_yaml_file(config_filepath=args.config) + config = Config.from_yaml(config_filepath=getattr(args, "config", "config.yaml")) - top_k = args.top_k if args.top_k is not None else config.rag_top_k + top_k = args.top_k if args.top_k is not None else (config.rag_top_k or 10) gemini_key = args.gemini_api_key or config.gemini_api_key gemini_model = args.gemini_model or config.gemini_model + use_llm = getattr(args, "use_llm", False) + is_infinite = getattr(args, "infinite", False) - client = BBQClient(server_url=args.server, config=config) - try: - print(f"Sending query to {args.server}: '{args.query}' (top_k={top_k})...\n") - - # Execute unified query + multimodal answer generation - rag_response = client.query_and_answer( - query_text=args.query, - top_k=top_k, - gemini_api_key=gemini_key, - gemini_model=gemini_model, - save_images=args.save_images, + if not getattr(args, "without_logo", False): + print_bbq(tag="CLIENT", server=args.server, top_k=top_k) + + log_level = logging.INFO if getattr(args, "verbose", False) else logging.WARNING + logging.basicConfig( + level=log_level, format="[%(asctime)s] [%(levelname)s] [%(name)s]: %(message)s" + ) + + console = Console() + + if not args.query and not is_infinite: + console.print( + "[red]Error: query text string is required (or run with --infinite for interactive mode).[/red]" ) + sys.exit(1) - results = rag_response.get("sources", []) - answer = rag_response.get("answer") - status = rag_response.get("status") + client = BBQClient( + server_url=args.server, + config=config, + without_logo=getattr(args, "without_logo", False), + ) - if not results: - print("No matching PDF pages found.") - return + current_query = args.query - # 1. If Gemini generated an answer, display it prominently - if answer: - print("=" * 60) - print(f"GEMINI MULTIMODAL ANSWER ({rag_response.get('engine', 'gemini')}):") - print("=" * 60) - print(answer) - print("\n" + "=" * 60) - print(f"GROUNDED RETRIEVED SOURCES (Top {len(results)} Pages):") - print("=" * 60) - else: - if status == "fallback_pages_only": - reason = rag_response.get("fallback_reason", "No API key or API call failed") - print(f"[Note: {reason}]") - print(f"Top {len(results)} Matching PDF Pages (Book Results):\n" + "=" * 60) - - # 2. Display the retrieved book pages - for i, res in enumerate(results, 1): - print(f"Rank {i}:") - print(f" Score : {res['score']:.4f}") - print(f" PDF File : {res['file_path']}") - print(f" Page : Page {res['page_number']} of {res['total_pages']}") - print(f" File Hash : {res['file_hash'][:12]}") - if res.get("saved_image_path"): - print(f" Saved Image : {res['saved_image_path']}") - print("-" * 60) + if is_infinite and not current_query: + console.print( + "[dim]Interactive mode active. Enter any query, 'help' for instructions, or 'q' to exit.[/dim]\n" + ) - except Exception as err: - print(f"Error querying BBQ server: {err}", file=sys.stderr) - sys.exit(1) + while True: + if not current_query: + try: + current_query = console.input( + r"[bold red]bbq[/bold red][bold yellow]\[query][/bold yellow] [bold white]>>[/bold white] " + ).strip() + except (KeyboardInterrupt, EOFError): + console.print("\n[dim]Exiting interactive query mode. Bye![/dim]") + break + + if not current_query: + if not is_infinite: + break + continue + + if current_query.lower() in ("help", "?", ":help"): + print_interactive_help(console) + current_query = None + continue + + if current_query.lower() in ("clear", "cls", ":clear"): + console.clear() + print_bbq(tag="CLIENT", server=args.server, top_k=top_k) + current_query = None + continue + + if current_query.lower() in ("exit", "quit", ":q", "q"): + console.print("[dim]Exiting interactive query mode. Bye![/dim]") + break + + try: + # 1. Fetch search matches from server + results = client.query(query_text=current_query, top_k=top_k) + + if not results: + console.print("[yellow]No matching PDF pages found.[/yellow]") + else: + # 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) + + # 3. If --use-llm is specified, pass top-k to LLM with cooking spinner + if use_llm: + with console.status( + "[bold cyan]LLM model is viewing the query and retrieved pages, cooking the meal...[/bold cyan]", + spinner="dots", + ): + rag_response = client.generate_answer_from_results( + query_text=current_query, + results=results, + gemini_api_key=gemini_key, + gemini_model=gemini_model, + save_images=args.save_images, + ) + + answer = rag_response.get("answer") + if answer: + render_llm_answer_rich( + answer=answer, + engine=rag_response.get("engine", "gemini"), + console=console, + ) + else: + fallback_reason = rag_response.get( + "fallback_reason", "No API key or API call failed" + ) + console.print( + f"\n[yellow][Note: LLM answer unavailable: {fallback_reason}][/yellow]" + ) + + except Exception as err: + console.print(f"[red]Error querying BBQ server:[/red] {err}") + if not is_infinite: + sys.exit(1) + + if not is_infinite: + break + + current_query = None def run_status_command(args: argparse.Namespace) -> None: """Fetches and displays running server status.""" from bbq.src.client import BBQClient + from bbq.src.terminal.tui import print_bbq + + if not getattr(args, "without_logo", False): + print_bbq(tag="CLIENT", server=args.server) - client = BBQClient(server_url=args.server) + client = BBQClient( + server_url=args.server, without_logo=getattr(args, "without_logo", False) + ) try: status = client.get_status() print("Server Status:") @@ -104,8 +189,14 @@ def run_status_command(args: argparse.Namespace) -> None: def run_documents_command(args: argparse.Namespace) -> None: """Lists all indexed document records.""" from bbq.src.client import BBQClient + from bbq.src.terminal.tui import print_bbq - client = BBQClient(server_url=args.server) + if not getattr(args, "without_logo", False): + print_bbq(tag="CLIENT", server=args.server) + + client = BBQClient( + server_url=args.server, without_logo=getattr(args, "without_logo", False) + ) try: documents = client.list_documents() if not documents: @@ -129,59 +220,258 @@ def build_argument_parser() -> argparse.ArgumentParser: prog="bbq", description="BBQ RAG - Unified Server & Client CLI Engine", ) + parser.add_argument( + "--version", + "-V", + action="version", + version=f"bbq {__version__}", + ) + parser.add_argument( + "--without-logo", + "--no-logo", + action="store_true", + default=False, + help="Do not print BBQ logo banner on run", + ) + + client_parent = argparse.ArgumentParser(add_help=False) + client_parent.add_argument( + "--without-logo", + "--no-logo", + action="store_true", + default=argparse.SUPPRESS, + help="Do not print BBQ logo banner on run", + ) subparsers = parser.add_subparsers(dest="command", help="Available commands") # --- Server Command --- - server_parser = subparsers.add_parser("server", help="Start persistent document indexing HTTP server") + server_parser = subparsers.add_parser( + "server", help="Start persistent document indexing HTTP server" + ) + server_parser.add_argument( + "--config", + "-c", + type=str, + default="config.yaml", + help="Path to YAML config file (default: config.yaml)", + ) + server_parser.add_argument( + "--host", + type=str, + default="0.0.0.0", + help="HTTP server bind host (default: 0.0.0.0)", + ) server_parser.add_argument( - "--config", "-c", type=str, default="config.yaml", help="Path to YAML config file (default: config.yaml)" + "--port", + "-p", + type=int, + default=8000, + help="HTTP server bind port (default: 8000)", ) - server_parser.add_argument("--host", type=str, default="0.0.0.0", help="HTTP server bind host (default: 0.0.0.0)") - server_parser.add_argument("--port", "-p", type=int, default=8000, help="HTTP server bind port (default: 8000)") server_parser.set_defaults(func=run_server_command) # --- Query / Client Command --- - query_parser = subparsers.add_parser("query", help="Query indexed documents via client API with optional Gemini RAG") - query_parser.add_argument("query", type=str, help="Search query text string") + query_parser = subparsers.add_parser( + "query", + aliases=["client"], + parents=[client_parent], + help="Query indexed documents via client API with optional Gemini RAG", + ) + query_parser.add_argument( + "query", + type=str, + nargs="?", + default=None, + help="Search query text string (optional when using --infinite)", + ) + query_parser.add_argument( + "--infinite", + "-inf", + action="store_true", + default=False, + help="Run interactive continuous query prompt loop", + ) + query_parser.add_argument( + "--config", + "-c", + type=str, + default="config.yaml", + help="Path to YAML config file (default: config.yaml)", + ) + query_parser.add_argument( + "--server", + "-s", + type=str, + default="http://localhost:8000", + help="Server URL (default: http://localhost:8000)", + ) query_parser.add_argument( - "--config", "-c", type=str, default="config.yaml", help="Path to YAML config file (default: config.yaml)" + "--top-k", + "-k", + type=int, + default=None, + help="Top K results to retrieve (default from config.yaml rag_top_k: 20)", ) - query_parser.add_argument("--server", "-s", type=str, default="http://localhost:8000", help="Server URL (default: http://localhost:8000)") - query_parser.add_argument("--top-k", "-k", type=int, default=None, help="Top K results to retrieve (default from config: 3)") query_parser.add_argument( - "--gemini-api-key", "-g", type=str, default=None, help="Google Gemini API key (or set in config.yaml / GEMINI_API_KEY env var)" + "--use-llm", + action="store_true", + default=False, + help="Pass top-k retrieved pages to LLM for multimodal answer generation", ) query_parser.add_argument( - "--gemini-model", type=str, default=None, help="Gemini model name (default from config: gemini-3.6-flash)" + "--gemini-api-key", + "-g", + type=str, + default=None, + help="Google Gemini API key (or set in config.yaml / GEMINI_API_KEY env var)", + ) + query_parser.add_argument( + "--gemini-model", + type=str, + default=None, + 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" + ) + query_parser.add_argument( + "--verbose", + "-v", + action="store_true", + help="Enable verbose client debug logging", + ) + query_parser.add_argument( + "--without-opener", + "--no-opener", + action="store_true", + default=False, + help="Do not launch interactive document page opener after search query", ) - query_parser.add_argument("--save-images", "-i", action="store_true", help="Save page images to disk") - query_parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose client debug logging") query_parser.set_defaults(func=run_client_query_command) # --- Status Command --- - status_parser = subparsers.add_parser("status", help="Check running server status") - status_parser.add_argument("--server", "-s", type=str, default="http://localhost:8000", help="Server URL") + status_parser = subparsers.add_parser( + "status", parents=[client_parent], help="Check running server status" + ) + status_parser.add_argument( + "--server", "-s", type=str, default="http://localhost:8000", help="Server URL" + ) status_parser.set_defaults(func=run_status_command) # --- Documents Command --- - docs_parser = subparsers.add_parser("documents", help="List indexed document metadata") - docs_parser.add_argument("--server", "-s", type=str, default="http://localhost:8000", help="Server URL") + docs_parser = subparsers.add_parser( + "documents", parents=[client_parent], help="List indexed document metadata" + ) + docs_parser.add_argument( + "--server", "-s", type=str, default="http://localhost:8000", help="Server URL" + ) docs_parser.set_defaults(func=run_documents_command) return parser +def build_client_main_parser() -> argparse.ArgumentParser: + """Builds the standalone parser used by python -m bbq.src.client.""" + parser = argparse.ArgumentParser( + prog="python -m bbq.src.client", + description="BBQ RAG - Client CLI Query & Management Engine", + ) + parser.add_argument( + "--version", + "-V", + action="version", + version=f"bbq {__version__}", + ) + parser.add_argument( + "query", type=str, nargs="?", default=None, help="Search query text string" + ) + parser.add_argument( + "--config", + "-c", + type=str, + default="config.yaml", + help="Path to YAML config file", + ) + parser.add_argument( + "--server", + "-s", + type=str, + default="http://localhost:8000", + help="Server URL (default: http://localhost:8000)", + ) + parser.add_argument( + "--top-k", + "-k", + type=int, + default=None, + help="Top K results to retrieve (default from config.yaml rag_top_k: 20)", + ) + parser.add_argument( + "--use-llm", + action="store_true", + default=False, + help="Pass top-k retrieved pages to LLM for multimodal answer generation", + ) + parser.add_argument( + "--infinite", + "-inf", + action="store_true", + default=False, + help="Run interactive continuous query prompt loop", + ) + parser.add_argument( + "--gemini-api-key", "-g", type=str, default=None, help="Google Gemini API key" + ) + parser.add_argument( + "--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" + ) + parser.add_argument( + "--without-logo", + "--no-logo", + action="store_true", + default=False, + help="Do not print BBQ logo banner on run", + ) + parser.add_argument( + "--without-opener", + "--no-opener", + action="store_true", + default=False, + help="Do not launch interactive document page opener", + ) + parser.add_argument("--status", action="store_true", help="Check server status") + parser.add_argument( + "--documents", action="store_true", help="List indexed documents" + ) + return parser + + def main() -> None: # Backward compatibility fallback: if run as 'python -m bbq.src.main [config.yaml]' without subcommands if ( len(sys.argv) > 1 and not sys.argv[1].startswith("-") - and sys.argv[1] not in ["server", "query", "status", "documents", "-h", "--help"] + and sys.argv[1] + not in [ + "server", + "query", + "client", + "status", + "documents", + "-h", + "--help", + "-V", + "--version", + ] ): config_path = sys.argv[1] print(f"Launching server mode with config: {config_path}") from bbq.src.server import start_document_indexing_server + start_document_indexing_server(config_filepath=config_path) return @@ -190,6 +480,7 @@ def main() -> None: # Default to server mode if no arguments provided if len(sys.argv) == 1: from bbq.src.server import start_document_indexing_server + start_document_indexing_server(config_filepath="config.yaml") return diff --git a/bbq/src/server/retrieval.py b/bbq/src/server/retrieval.py index 0241cf2..ed1c9dc 100644 --- a/bbq/src/server/retrieval.py +++ b/bbq/src/server/retrieval.py @@ -25,7 +25,9 @@ def query_indexed_documents( from maxsimd import maxsim records = tracker.fetch_all_records() - completed_records = [r for r in records if r.get("status") == "done" and r.get("embedding_path")] + completed_records = [ + r for r in records if r.get("status") == "done" and r.get("embedding_path") + ] if not completed_records: logger.info("No completed PDF embeddings found in tracker database.") @@ -35,12 +37,11 @@ def query_indexed_documents( q_tensor: torch.Tensor = engine.encode_query_text_inputs([query_text]) q_mat = q_tensor[0].cpu().float().numpy() - page_metadata = [] - page_embeddings_list = [] + all_page_results: List[Dict[str, Any]] = [] for record in completed_records: - emb_path = record["embedding_path"] - if not os.path.exists(emb_path): + emb_path = record.get("embedding_path") + if not emb_path or not os.path.exists(emb_path): continue doc_emb = np.load(emb_path) @@ -48,36 +49,23 @@ def query_indexed_documents( doc_emb = np.expand_dims(doc_emb, axis=0) num_pages = doc_emb.shape[0] - for page_idx in range(num_pages): - page_metadata.append( + + # 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] + + filename = os.path.basename(record.get("file_path", "")) + for page_idx, score in enumerate(scores_list): + all_page_results.append( { "file_path": record["file_path"], - "filename": os.path.basename(record["file_path"]), + "filename": filename, "file_hash": record["file_hash"], "page_number": page_idx + 1, "total_pages": num_pages, + "score": score, } ) - page_embeddings_list.append(doc_emb) - - if not page_embeddings_list: - return [] - - # Batch all document pages into a single contiguous 3D array: (total_pages, tokens_per_page, dim) - batched_docs_3d = ( - page_embeddings_list[0] - if len(page_embeddings_list) == 1 - else np.concatenate(page_embeddings_list, axis=0) - ) - - # Multi-core Rayon parallelized MaxSim scoring across all document pages - scores = maxsim(q_mat, batched_docs_3d) - - all_page_results = [] - for meta, score in zip(page_metadata, scores): - meta_dict = meta.copy() - meta_dict["score"] = float(score) - all_page_results.append(meta_dict) all_page_results.sort(key=lambda x: x["score"], reverse=True) return all_page_results[:top_k] diff --git a/bbq/src/server/server.py b/bbq/src/server/server.py index b64ed75..9bf0973 100644 --- a/bbq/src/server/server.py +++ b/bbq/src/server/server.py @@ -2,15 +2,14 @@ import os import signal import sys -import time import threading +import time from typing import Any, Optional, Tuple from rich.console import Console from bbq.src.config import ( - ModelConfigWrapper, - load_configuration_from_yaml_file, + Config, ) from bbq.src.server.app import create_bbq_fastapi_app, run_http_server_in_thread from bbq.src.server.ingestion import ( @@ -41,6 +40,7 @@ class BBQServer: def __init__( self, + config: Optional[Config] = None, config_filepath: str = "config.yaml", host: str = "0.0.0.0", port: int = 8000, @@ -51,7 +51,7 @@ def __init__( self.console = Console() self.bbq_logger = configure_server_logging() - self.config: ModelConfigWrapper = load_configuration_from_yaml_file(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 @@ -70,9 +70,9 @@ def load_engine_threaded(self) -> None: def _loader(): try: - from bbq.src.utils.model_loader import initialize_engine_from_yaml_config + from bbq.src.utils.model_loader import initialize_engine - self.engine = initialize_engine_from_yaml_config(self.config_filepath) + self.engine = initialize_engine(self.config) self._is_ready = True logger.info("Engine model and processor loaded successfully.") except Exception as exc: @@ -196,7 +196,7 @@ def is_ready(self) -> bool: def execute_background_server_pipeline( config_filepath: str, - config: ModelConfigWrapper, + config: Config, tracker: SqlliteDB, console: Console, host: str = "0.0.0.0", @@ -209,12 +209,13 @@ def execute_background_server_pipeline( def start_document_indexing_server( + config: Optional[Config] = None, config_filepath: str = "config.yaml", host: str = "0.0.0.0", port: int = 8000, ) -> None: """Main entrypoint for starting the document indexing server.""" - server = BBQServer(config_filepath=config_filepath, host=host, port=port) + server = BBQServer(config=config, config_filepath=config_filepath, host=host, port=port) try: server.start() while True: diff --git a/bbq/src/storage/sql.py b/bbq/src/storage/sql.py index 3e42324..c33ba8c 100644 --- a/bbq/src/storage/sql.py +++ b/bbq/src/storage/sql.py @@ -3,7 +3,7 @@ from datetime import datetime, timezone from typing import Any, Dict, Optional -from bbq.src.config import get_system_cache_dir +from bbq.src.utils.futils import get_system_cache_dir class SqlliteDB: diff --git a/bbq/src/terminal/__init__.py b/bbq/src/terminal/__init__.py index 5e9ff09..29d0eb6 100644 --- a/bbq/src/terminal/__init__.py +++ b/bbq/src/terminal/__init__.py @@ -2,6 +2,12 @@ create_rich_console_logging_handler, render_server_status_rich_panel, configure_rich_logging_for_server, + print_bbq, + render_query_results_rich, + render_llm_answer_rich, + print_interactive_help, + open_document_page, + interactive_page_picker, ) from bbq.src.terminal.logger import BBQLogger, configure_server_logging @@ -9,6 +15,12 @@ "create_rich_console_logging_handler", "render_server_status_rich_panel", "configure_rich_logging_for_server", + "print_bbq", + "render_query_results_rich", + "render_llm_answer_rich", + "print_interactive_help", + "open_document_page", + "interactive_page_picker", "BBQLogger", "configure_server_logging", ] diff --git a/bbq/src/terminal/logger.py b/bbq/src/terminal/logger.py index a99d35f..040423c 100644 --- a/bbq/src/terminal/logger.py +++ b/bbq/src/terminal/logger.py @@ -4,7 +4,7 @@ from rich.console import Console from rich.logging import RichHandler -from bbq.src.config import get_system_cache_dir +from bbq.src.utils.futils import get_system_cache_dir class BBQLogger: """ diff --git a/bbq/src/terminal/tui.py b/bbq/src/terminal/tui.py index 05e82db..7ab616e 100644 --- a/bbq/src/terminal/tui.py +++ b/bbq/src/terminal/tui.py @@ -1,10 +1,21 @@ +import os +import re import logging +from typing import Optional, Any + from rich.console import Console +from rich.logging import RichHandler from rich.panel import Panel +from rich.rule import Rule +from rich.markdown import Markdown from rich.table import Table -from rich.logging import RichHandler -from bbq.src.config import ModelConfigWrapper +from bbq.src.config import Config + +CYAN = "\033[1;36m" +RESET = "\033[0m" +DIM = "\033[90m" +BOLD = "\033[1m" def create_rich_console_logging_handler() -> RichHandler: @@ -16,7 +27,7 @@ def create_rich_console_logging_handler() -> RichHandler: ) -def render_server_status_rich_panel(config: ModelConfigWrapper) -> Panel: +def render_server_status_rich_panel(config: Config) -> Panel: table = Table(show_header=False, expand=True, box=None) table.add_column("Key", style="bold green", width=22) table.add_column("Value", style="bold white") @@ -41,3 +52,432 @@ def configure_rich_logging_for_server() -> None: root_logger.handlers.clear() rich_handler = create_rich_console_logging_handler() root_logger.addHandler(rich_handler) + + +def _extract_server_label(server: Optional[Any]) -> Optional[str]: + if server is None: + return None + s = str(server).strip() + if not s: + return None + if s.isdigit(): + return s + try: + from urllib.parse import urlparse + + parsed = urlparse(s if "://" in s else f"//{s}") + if parsed.port: + return str(parsed.port) + if parsed.hostname: + return parsed.hostname + except Exception: + pass + return s + + +def print_bbq( + logo_path: Optional[str] = None, + name: Optional[str] = None, + tag: str = "CLIENT", + without_logo: bool = False, + server: Optional[Any] = None, + top_k: Optional[int] = None, +) -> None: + """ + Prints the stylized ASCII BBQ logo banner with tag, server, and top-k annotations. + Format: BBQ [CLIENT] [8000] [top-10] + If without_logo is True, suppression is respected and nothing is printed. + """ + if without_logo: + return + + path = logo_path + if not path or not os.path.exists(path): + for candidate in [ + os.path.join(os.getcwd(), "assets", "logo.txt"), + os.path.join( + os.path.dirname( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) + ), + "assets", + "logo.txt", + ), + os.path.join( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ), + "assets", + "logo.txt", + ), + ]: + if os.path.exists(candidate): + path = candidate + break + + server_label = _extract_server_label(server) + plain_tokens = ["BBQ", f"[{tag.upper()}]"] + styled_tokens = [f"{BOLD}BBQ{RESET}", f"{CYAN}[{tag.upper()}]{RESET}"] + + if server_label: + plain_tokens.append(f"[{server_label}]") + styled_tokens.append(f"{CYAN}[{server_label}]{RESET}") + + if top_k is not None: + plain_tokens.append(f"[top-{top_k}]") + styled_tokens.append(f"{CYAN}[top-{top_k}]{RESET}") + + if name and name.upper() not in ("BBQ", "BBQ-RAG", "CLIENT", "SERVER", tag.upper()): + plain_tokens.append(f"[{name}]") + styled_tokens.append(f"{CYAN}[{name}]{RESET}") + + plain_header = " ".join(plain_tokens) + styled_header = " ".join(styled_tokens) + + if not path or not os.path.exists(path): + print(f"\n{'=' * 55}\n {plain_header}\n{'=' * 55}\n") + return + + try: + with open(path, "r", encoding="utf-8") as f: + data = [line.rstrip("\r\n") for line in f if line.strip()] + except Exception: + return + + ansi = lambda s: re.sub(r"\x1b\[[0-9;]*m", "", s) + vlen = max((len(ansi(line)) for line in data), default=0) + mid = len(data) // 2 + + divider_len = len(plain_header) + divider = "─" * divider_len + + print() + for i, line in enumerate(data): + vslen = len(ansi(line)) + spacing = " " * (vlen - vslen + 4) + if i == mid - 1: + print(f"{line}{spacing}{DIM}│{RESET} {styled_header}") + elif i == mid: + print(f"{line}{spacing}{DIM}│{RESET} {DIM}{divider}{RESET}") + else: + print(line) + print() + + +def render_query_results_rich( + results: list[dict], + console: Optional[Console] = None, +) -> None: + """ + Renders top retrieved PDF document pages using a clean list showing page number and document name. + """ + console = console or Console() + console.print( + Rule( + title=f"[bold yellow]Top {len(results)} Matching PDF Pages[/bold yellow]", + style="yellow", + ) + ) + for i, res in enumerate(results, 1): + book_name = res.get("filename") or os.path.basename(res.get("file_path", "")) + page_num = res.get("page_number", "?") + total_pages = res.get("total_pages") + if total_pages is not None and total_pages != "?": + page_info = f"Page [bold yellow]{page_num}[/bold yellow] of {total_pages}" + else: + page_info = f"Page [bold yellow]{page_num}[/bold yellow]" + + console.print( + f" [bold magenta]{i:2d}.[/bold magenta] {page_info} — [bold white]{book_name}[/bold white]" + ) + console.print() + + +def open_document_page(file_path: str, page_number: int = 1) -> bool: + """ + Opens the PDF document at the specified page using evince, zathura, okular, or xdg-open. + """ + import shutil + import subprocess + + if not file_path: + return False + + target_path = file_path + if not os.path.exists(target_path): + candidate = os.path.join(os.getcwd(), file_path) + if os.path.exists(candidate): + target_path = candidate + + target_path = os.path.abspath(target_path) + if not os.path.exists(target_path): + return False + + # 1. Evince (supports -p PAGE) + if shutil.which("evince"): + try: + subprocess.Popen( + ["evince", "-p", str(page_number), target_path], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + return True + except Exception: + pass + + # 2. Zathura (supports -P PAGE) + if shutil.which("zathura"): + try: + subprocess.Popen( + ["zathura", "-P", str(page_number), target_path], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + return True + except Exception: + pass + + # 3. Okular (supports -p PAGE) + if shutil.which("okular"): + try: + subprocess.Popen( + ["okular", "-p", str(page_number), target_path], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + return True + except Exception: + pass + + # 4. Fallback to xdg-open + if shutil.which("xdg-open"): + try: + subprocess.Popen( + ["xdg-open", target_path], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + return True + except Exception: + pass + + return False + + +def _read_terminal_key() -> str: + """Reads a single keypress or ANSI escape sequence from stdin in raw mode.""" + import sys + import os + import select + + if not sys.stdin.isatty(): + return "ESC" + + import tty + import termios + + fd = sys.stdin.fileno() + old_settings = termios.tcgetattr(fd) + try: + tty.setraw(fd) + raw = os.read(fd, 32) + if not raw: + return "ESC" + + if raw == b"\x1b": + r, _, _ = select.select([fd], [], [], 0.05) + if r: + raw += os.read(fd, 16) + + # Arrow keys: support both CSI (\x1b[) and SS3 (\x1bO) sequences across all terminals + if raw in (b"\x1b[A", b"\x1bOA", b"\x1b[1;2A", b"\x1b[1;5A"): + return "UP" + elif raw in (b"\x1b[B", b"\x1bOB", b"\x1b[1;2B", b"\x1b[1;5B"): + return "DOWN" + elif raw in (b"\x1b[C", b"\x1bOC"): + return "RIGHT" + elif raw in (b"\x1b[D", b"\x1bOD"): + return "LEFT" + elif raw == b"\x1b": + return "ESC" + elif raw in (b"\r", b"\n"): + return "ENTER" + elif raw in (b"\x03", b"\x04"): # Ctrl+C or Ctrl+D + return "CTRL_C" + elif raw in (b"k", b"K"): + return "UP" + elif raw in (b"j", b"J"): + return "DOWN" + elif raw in (b"q", b"Q"): + return "ESC" + return "ESC" + except Exception: + return "ESC" + finally: + termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) + + + + +def interactive_page_picker( + results: list[dict], + console: Optional[Console] = None, +) -> None: + """ + Launches an interactive arrow-key picker (Option C) allowing the user + to navigate retrieved pages, press Enter to open in PDF viewer, and Esc/q to proceed. + """ + import sys + + console = console or Console() + + if not results or not sys.stdin.isatty(): + render_query_results_rich(results, console=console) + return + + from rich.live import Live + from rich.console import Group + from rich.text import Text + + selected_index = 0 + total_options = len(results) + 1 + status_msg: Optional[str] = None + + def build_picker_renderable() -> Group: + items = [] + items.append( + Rule( + title=f"[bold yellow]Top {len(results)} Matching PDF Pages[/bold yellow]", + style="yellow", + ) + ) + items.append( + Text.from_markup( + "[dim cyan] Use ↑/↓ or j/k to navigate • Enter to open page in viewer • Esc/q to proceed[/dim cyan]\n" + ) + ) + + for i, res in enumerate(results): + book_name = res.get("filename") or os.path.basename( + res.get("file_path", "") + ) + page_num = res.get("page_number", "?") + total_pages = res.get("total_pages") + if total_pages is not None and total_pages != "?": + page_str = f"Page {page_num} of {total_pages}" + else: + page_str = f"Page {page_num}" + + if i == selected_index: + line = Text.from_markup( + f" [bold cyan]▶ {i + 1:2d}.[/bold cyan] [bold yellow]{page_str}[/bold yellow] — [bold white]{book_name}[/bold white]" + ) + else: + line = Text.from_markup( + f" [dim]{i + 1:2d}.[/dim] {page_str} — [dim]{book_name}[/dim]" + ) + items.append(line) + + done_idx = len(results) + if selected_index == done_idx: + items.append( + Text.from_markup( + "\n [bold green]▶ ✔ Done / Proceed to query prompt (Esc/q)[/bold green]" + ) + ) + else: + items.append( + Text.from_markup( + "\n [dim]✔ Done / Proceed to query prompt (Esc/q)[/dim]" + ) + ) + + if status_msg: + items.append(Text.from_markup(f"\n {status_msg}")) + + return Group(*items) + + try: + with Live( + build_picker_renderable(), + console=console, + auto_refresh=False, + transient=True, + ) as live: + while True: + live.update(build_picker_renderable()) + live.refresh() + + key = _read_terminal_key() + + if key == "UP": + selected_index = (selected_index - 1) % total_options + elif key == "DOWN": + selected_index = (selected_index + 1) % total_options + elif key == "ENTER": + if selected_index == len(results): + break + res = results[selected_index] + fpath = res.get("file_path", "") + pnum = res.get("page_number", 1) + fname = res.get("filename") or os.path.basename(fpath) + opened = open_document_page(fpath, page_number=pnum) + if opened: + status_msg = f"[bold green]✓ Opened Page {pnum} of {fname} in PDF viewer[/bold green]" + else: + status_msg = f"[bold red]✗ Could not open {fname} (viewer not found)[/bold red]" + elif key in ("ESC", "CTRL_C", "CTRL_D"): + break + except Exception: + pass + + # Finally print clean static list into terminal scrollback history + render_query_results_rich(results, console=console) + + +def render_llm_answer_rich( + answer: str, + engine: str = "gemini", + console: Optional[Console] = None, +) -> None: + """ + Renders LLM multimodal answer in a rich Panel using Markdown formatting. + """ + console = console or Console() + console.print( + Panel( + Markdown(answer), + title=f"[bold green]GEMINI MULTIMODAL ANSWER ({engine})[/bold green]", + title_align="left", + border_style="green", + padding=(1, 2), + ) + ) + + +def print_interactive_help(console: Optional[Console] = None) -> None: + """ + Displays the help instructions for the interactive query prompt loop. + """ + console = console or Console() + help_content = ( + "[bold white]• Enter any query[/bold white] Search across indexed PDF documents\n" + "[bold white]• help[/bold white] or [bold white]?[/bold white] Display this interactive commands help\n" + "[bold white]• clear[/bold white] or [bold white]cls[/bold white] Clear the terminal screen\n" + "[bold white]• q[/bold white], [bold white]quit[/bold white], or [bold white]exit[/bold white] Quit and exit interactive prompt\n" + "[bold white]• Ctrl+C[/bold white] Cancel query or terminate session" + ) + console.print( + Panel( + help_content, + title=r"[bold yellow]bbq\[query] Interactive Commands[/bold yellow]", + title_align="left", + border_style="yellow", + padding=(0, 1), + ) + ) diff --git a/bbq/src/utils/__init__.py b/bbq/src/utils/__init__.py index 7de67f2..f1a6cf4 100644 --- a/bbq/src/utils/__init__.py +++ b/bbq/src/utils/__init__.py @@ -1,6 +1,7 @@ import importlib from typing import Any +from bbq.src.utils.futils import get_system_cache_dir from bbq.src.utils.pdf_utils import ( compute_file_sha256_hash, extract_pdf_pages_to_pil_images, @@ -13,6 +14,7 @@ "BaseEngineWrapper": "bbq.src.common.base", "EngineModelLoader": "bbq.src.utils.model_loader", "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", @@ -33,4 +35,6 @@ def __getattr__(name: str) -> Any: "compute_file_sha256_hash", "extract_pdf_pages_to_pil_images", "extract_single_pdf_page_image", + "get_system_cache_dir", + "initialize_engine", ] diff --git a/bbq/src/utils/futils.py b/bbq/src/utils/futils.py new file mode 100644 index 0000000..be5ed17 --- /dev/null +++ b/bbq/src/utils/futils.py @@ -0,0 +1,16 @@ +""" +Utils regarding file management +""" +from __future__ import annotations + +import os + + +def get_system_cache_dir(subfolder: str = "") -> str: + """ + Returns the system default user cache directory for bbq. + Respects XDG_CACHE_HOME if set, defaulting to ~/.cache/bbq. + """ + cache_base = os.environ.get("XDG_CACHE_HOME") or os.path.expanduser("~/.cache") + bbq_cache_dir = os.path.join(cache_base, "bbq") + return os.path.join(bbq_cache_dir, subfolder) if subfolder else bbq_cache_dir diff --git a/bbq/src/utils/model_loader.py b/bbq/src/utils/model_loader.py index 27320c5..8d40178 100644 --- a/bbq/src/utils/model_loader.py +++ b/bbq/src/utils/model_loader.py @@ -1,38 +1,51 @@ from typing import Any, List, Optional, Tuple, Type, Union -from peft import PeftConfig, PeftModel import torch +from peft import PeftConfig, PeftModel from transformers import logging as tf_logging from bbq.src.common.base import BaseEngineWrapper, BaseModel, BaseModelLoader, BaseProcessor -from bbq.src.models.registry import ModelRegistry -from bbq.src.config import ( - ModelConfigWrapper, - determine_target_torch_device, - load_configuration_from_yaml_file, - resolve_torch_data_type, -) from bbq.src.common.errors import ( BaseModelInstantiateError, LoRAAdapterLoadError, ProcessorLoadError, ) +from bbq.src.config import Config +from bbq.src.models.registry import ModelRegistry tf_logging.set_verbosity_warning() + +def determine_target_torch_device(device_preference: str = "auto") -> str: + """Determine the target torch device based on preference and CUDA availability.""" + if device_preference == "auto": + return "cuda" if torch.cuda.is_available() else "cpu" + return device_preference + + +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 + if dtype_name == "bfloat16": + return torch.bfloat16 + elif dtype_name == "float16": + return torch.float16 + return torch.float32 + + 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 def load_model_with_lora_adapters( - config_input: Union[str, ModelConfigWrapper], + config_input: Union[str, Config], ) -> Tuple[Union[BaseModel, PeftModel], BaseProcessor]: """ Loads a model with LoRA adapters @@ -46,7 +59,7 @@ def load_model_with_lora_adapters( tf_logging.set_verbosity_error() if isinstance(config_input, str): - config: ModelConfigWrapper = load_configuration_from_yaml_file(config_input) + config: Config = Config.from_yaml(config_input) else: config = config_input @@ -94,25 +107,24 @@ def load_model_with_lora_adapters( return model_to_use, processor -class EngineWrapper(BaseEngineWrapper): +class EngineWrapper(BaseEngineWrapper): """ Class for wrapping engine models """ + def __init__( self, model: Any, processor: BaseProcessor, - config: ModelConfigWrapper, + config: Config, ) -> None: self.model = model self.processor = processor - self.config: ModelConfigWrapper = config + 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. """ @@ -126,8 +138,7 @@ def encode_multimodal_document_images( 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(): @@ -145,30 +156,37 @@ def encode_query_text_inputs(self, texts: List[str]) -> torch.Tensor: """ if not texts: raise ValueError("Input text list cannot be empty.") - + 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(): query_embeddings: torch.Tensor = self.model(**processed_inputs) - + return query_embeddings -def initialize_engine_from_yaml_config( - config_filepath: str = "config.yaml", + +def initialize_engine( + config: Config, loader_class: Optional[Type[BaseModelLoader]] = None, ) -> EngineWrapper: """ - Initializes an engine from a YAML configuration + Initializes an engine directly from an injected Config instance. """ - config: ModelConfigWrapper = load_configuration_from_yaml_file(config_filepath) - target_loader = loader_class or EngineModelLoader - model, processor = target_loader.load_model_and_processor(config) - return EngineWrapper(model=model, processor=processor, config=config) + + +def initialize_engine_from_yaml_config( + config_filepath: str = "config.yaml", + loader_class: Optional[Type[BaseModelLoader]] = None, +) -> EngineWrapper: + """ + Initializes an engine from a YAML configuration file path. + """ + config: Config = Config.from_yaml(config_filepath) + return initialize_engine(config=config, loader_class=loader_class) diff --git a/benchmarks/README.md b/benchmarks/README.md index bfc571e..5ee27e9 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -51,7 +51,7 @@ pip install -e ".[benchmark]" maturin develop --release ``` -Benchmark dependencies include `tqdm`, `matplotlib`, `jinja2`, `datasets`, and `plotly`. +Benchmark dependencies include `matplotlib`, `jinja2`, `datasets`, and `plotly`. --- diff --git a/benchmarks/cli.py b/benchmarks/cli.py index 3b04d96..73c45b7 100644 --- a/benchmarks/cli.py +++ b/benchmarks/cli.py @@ -120,6 +120,14 @@ def bm_parser() -> argparse.ArgumentParser: help="Save b1/b2 benchmark images to asset directory (default: None, only saves in bench/ folder. Pass --asset to save to assets/)", ) + parser.add_argument( + "--without-logo", + "--no-logo", + action="store_true", + default=False, + help="Do not print BBQ logo banner on run", + ) + return parser @@ -129,8 +137,10 @@ def main(): configure_global_threads(args.jobs) - if args.suite == "all": + if not args.without_logo: print_bbq(name=args.suite) + + if args.suite == "all": from benchmarks.maxsimd.run import run as run_maxsimd run_maxsimd( @@ -156,7 +166,6 @@ def main(): ) elif args.suite == "maxsimd": - print_bbq(name=args.suite) from benchmarks.maxsimd.run import run run( @@ -170,7 +179,6 @@ def main(): asset_dir=args.asset_dir, ) elif args.suite == "vidore": - print_bbq(name=args.suite) from benchmarks.vidore.run import run run( diff --git a/benchmarks/ui.py b/benchmarks/ui.py index 07f36cc..26cb86e 100644 --- a/benchmarks/ui.py +++ b/benchmarks/ui.py @@ -1,51 +1,25 @@ import os -from warnings import warn -import re +from typing import Optional +from bbq.src.terminal.tui import print_bbq as _bbq_print_bbq ASSETS_DIR = os.path.join( os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "assets" ) -CYAN = "\033[1;36m" -RESET = "\033[0m" -DIM = "\033[90m" -BOLD = "\033[1m" def print_bbq( - logo_path: str = os.path.join(ASSETS_DIR, "logo.txt"), name: str = "BBQ-RAG" + logo_path: Optional[str] = None, + name: str = "BB", + tag: str = "BENCHMARK", + without_logo: bool = False, ) -> None: path = logo_path or os.path.join(ASSETS_DIR, "logo.txt") - - if not os.path.exists(path): - print(f"\n{'=' * 55}\n BENCHMARK [{name.upper()}]\n{'=' * 55}\n") - warn(f"Logo file not found: {path}") - return - - with open(path, "r", encoding="utf-8") as f: - data = [l_.rstrip("\r\n") for l_ in f if l_.strip()] - - ansi = lambda s: re.sub(r"\x1b\[[0-9;]*m", "", s) - vlen = max(len(ansi(_l)) for _l in data) - - mid = len(data) // 2 - - print() - for i, _l in enumerate(data): - vslen = len(ansi(_l)) - - spacing = " " * (vlen - vslen + 4) - - if i == mid - 1: - print( - f"{_l}{spacing}{DIM}│{RESET} {BOLD}BENCHMARK{RESET} {CYAN}[{name}]{RESET}" - ) - elif i == mid: - divider = "─" * (len(name) + 14) - print(f"{_l}{spacing}{DIM}│{RESET} {DIM}{divider}{RESET}") - else: - print(_l) - - print() + _bbq_print_bbq( + logo_path=path, + name=name, + tag=tag, + without_logo=without_logo, + ) __all__ = [ diff --git a/pyproject.toml b/pyproject.toml index 53d08c6..6a73263 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["maturin>=1.15.0,<2.0"] build-backend = "maturin" [project] -name = "bbq-rag" +name = "bbq" requires-python = ">=3.8" license = { text = "MIT" } classifiers = [ @@ -18,9 +18,11 @@ dependencies = [ ] dynamic = ["version"] +[project.scripts] +bbq = "bbq.src.main:main" + [project.optional-dependencies] benchmark = [ - "tqdm>=4.70.1", "matplotlib>=3.10.9", "jinja2>=3.1.6", "datasets>=5.0.1", @@ -31,7 +33,6 @@ dev = ["pytest>=9.1.1"] [dependency-groups] dev = ["pytest>=9.1.1"] benchmark = [ - "tqdm>=4.70.1", "matplotlib>=3.10.9", "jinja2>=3.1.6", "datasets>=5.0.1", @@ -43,8 +44,8 @@ testpaths = ["tests"] pythonpath = ["."] [tool.maturin] -module-name = "maxsimd" -python-source = "bbq" +module-name = "bbq.maxsimd.maxsimd" +python-source = "." [tool.pyright] extraPaths = ["bbq"] diff --git a/sample.config.yaml b/sample.config.yaml index c5df862..80c8548 100644 --- a/sample.config.yaml +++ b/sample.config.yaml @@ -31,3 +31,6 @@ pdf_render_dpi: 150 gemini_api_key: "" gemini_model: "gemini-3.6-flash" rag_top_k: 3 + +# Quantization precision mode: "f32" (default) or "qi8" (quantized INT8) +quantization: "f32" diff --git a/tests/test_all_features.py b/tests/test_all_features.py index 06a4c49..e7e5de0 100644 --- a/tests/test_all_features.py +++ b/tests/test_all_features.py @@ -8,56 +8,14 @@ import os import shutil import tempfile + +import maxsimd import numpy as np import pytest import torch -from bbq.src.config import ModelConfigWrapper, get_system_cache_dir +from bbq.src.server import create_bbq_fastapi_app, query_indexed_documents from bbq.src.storage.sql import SqlliteDB -from bbq.src.server import query_indexed_documents, create_bbq_fastapi_app -from bbq.src.client import BBQClient -import maxsimd - - -def test_system_cache_path(): - """Verify that default cache paths use user system cache directory (~/.cache/bbq).""" - system_cache = get_system_cache_dir() - assert ".cache/bbq" in system_cache or "bbq" in system_cache - assert not system_cache.startswith("./.cache") - - config = ModelConfigWrapper() - assert config.embeddings_output_path == get_system_cache_dir("embeddings") - assert config.sqlite_db_path == get_system_cache_dir("tracker.db") - assert config.gemini_model == "gemini-3.6-flash" - assert config.rag_top_k == 3 - - tracker = SqlliteDB() - assert tracker.db_filepath == get_system_cache_dir("tracker.db") - - -def test_config_gemini_settings_and_client(): - """Verify Gemini configuration in ModelConfigWrapper and client initialization.""" - from bbq.src.config import load_configuration_from_yaml_file - - config = ModelConfigWrapper( - gemini_api_key="TEST_API_KEY_ABC", - gemini_model="gemini-2.0-flash", - rag_top_k=3, - ) - assert config.gemini_api_key == "TEST_API_KEY_ABC" - assert config.gemini_model == "gemini-2.0-flash" - assert config.rag_top_k == 3 - - dict_format = config.convert_to_dictionary_format() - assert dict_format["gemini_api_key_configured"] is True - assert dict_format["gemini_model"] == "gemini-2.0-flash" - assert dict_format["rag_top_k"] == 3 - - gemini_client = config.get_gemini_client() - assert gemini_client.api_key == "TEST_API_KEY_ABC" - assert gemini_client.model == "gemini-2.0-flash" - assert gemini_client.is_available() is True - def test_maxsim_ragged_flat_functionality(): @@ -143,6 +101,7 @@ class MockEngine: class MockConfig: base_model_id = "test-model" watch_folder_path = temp_dir + config = MockConfig() model = torch.nn.Linear(10, 10) @@ -151,7 +110,9 @@ def encode_query_text_inputs(self, texts): return torch.randn(1, 4, 128) mock_engine = MockEngine() - results = query_indexed_documents("test search query", mock_engine, tracker, top_k=2) + results = query_indexed_documents( + "test search query", mock_engine, tracker, top_k=2 + ) assert len(results) == 2 assert results[0]["file_path"] == pdf_path @@ -166,9 +127,10 @@ def encode_query_text_inputs(self, texts): def test_single_page_image_extraction(): """Verify extracting a single PDF page to a PIL Image.""" import pymupdf - from bbq.src.utils.pdf_utils import extract_single_pdf_page_image from PIL import Image + from bbq.src.utils.pdf_utils import extract_single_pdf_page_image + temp_dir = tempfile.mkdtemp() try: pdf_path = os.path.join(temp_dir, "test_doc.pdf") @@ -187,9 +149,10 @@ def test_single_page_image_extraction(): def test_gemini_client_and_fallback(): """Verify GeminiClient behavior with no key, invalid key, and image encoding.""" - from bbq.src.client.gemini import GeminiClient from PIL import Image + from bbq.src.client.gemini import GeminiClient + # 1. No key client_no_key = GeminiClient(api_key="") assert not client_no_key.is_available() @@ -280,7 +243,9 @@ def test_quantization_qi8_and_quantized_maxsim(): assert d_scale.dtype == torch.float32 # Call maxsim with quantized tensors - scores_torch = maxsimd.maxsim(q_val, d_val, q_scale=q_scale, d_scale=d_scale, jobs=-1) + scores_torch = maxsimd.maxsim( + q_val, d_val, q_scale=q_scale, d_scale=d_scale, jobs=-1 + ) assert len(scores_torch) == 4 for s in scores_torch: assert isinstance(s, float) @@ -303,11 +268,9 @@ def test_quantization_qi8_and_quantized_maxsim(): assert d_np_scale.shape == (20, dim // 32) assert d_np_scale.dtype == np.float32 - scores_np = maxsimd.maxsim(q_np_val, d_np_val, q_scale=q_np_scale, d_scale=d_np_scale) + scores_np = maxsimd.maxsim( + q_np_val, d_np_val, q_scale=q_np_scale, d_scale=d_np_scale + ) assert len(scores_np) == 1 assert isinstance(scores_np[0], float) assert np.isfinite(scores_np[0]) - - - - diff --git a/tests/test_client.py b/tests/test_client.py new file mode 100644 index 0000000..f5b1a39 --- /dev/null +++ b/tests/test_client.py @@ -0,0 +1,440 @@ +""" +Unit tests for BBQ RAG Client module, Gemini client, logo banner printing, +and --without-logo suppression behavior. +""" + +import io +from unittest.mock import patch, MagicMock +from PIL import Image + +import pytest +from bbq.src.terminal.tui import print_bbq +from bbq.src.client import BBQClient, GeminiClient +from bbq.src.main import ( + build_argument_parser, + build_client_main_parser, + __version__, +) + + +def test_print_bbq_banner_output(capsys): + """Verify that print_bbq prints the ASCII banner with BBQ, tag, server, and top-k annotations.""" + print_bbq(name="BBQ", tag="CLIENT", without_logo=False) + captured = capsys.readouterr().out + assert "BBQ" in captured + assert "[CLIENT]" in captured + + print_bbq(tag="CLIENT", server="http://localhost:8080", top_k=4) + captured2 = capsys.readouterr().out + assert "BBQ" in captured2 + assert "[CLIENT]" in captured2 + assert "[8080]" in captured2 + assert "[top-4]" in captured2 + + +def test_print_bbq_without_logo_suppresses_output(capsys): + """Verify that print_bbq prints nothing when without_logo is True.""" + print_bbq(name="BBQ", tag="CLIENT", without_logo=True) + captured = capsys.readouterr().out + assert captured == "" + + +def test_client_init_and_attributes(): + """Verify BBQClient attributes and initialization.""" + client = BBQClient(server_url="http://test-server:9000/", without_logo=True) + assert client.server_url == "http://test-server:9000" + assert client.without_logo is True + + default_client = BBQClient() + assert default_client.server_url == "http://localhost:8000" + assert default_client.without_logo is False + + +def test_client_print_logo(capsys): + """Verify BBQClient.print_logo respects without_logo.""" + client_with_logo = BBQClient(without_logo=False) + client_with_logo.print_logo() + out = capsys.readouterr().out + assert "CLIENT" in out + + client_without_logo = BBQClient(without_logo=True) + client_without_logo.print_logo() + assert capsys.readouterr().out == "" + + +@patch("requests.post") +def test_client_query_mock(mock_post): + """Verify client.query posts to server endpoint.""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "results": [{"file_path": "doc.pdf", "page_number": 1, "score": 2.5}] + } + mock_post.return_value = mock_response + + client = BBQClient(server_url="http://localhost:8000") + results = client.query("test query", top_k=5) + + assert len(results) == 1 + assert results[0]["file_path"] == "doc.pdf" + assert results[0]["page_number"] == 1 + mock_post.assert_called_once_with( + "http://localhost:8000/query", + json={"query": "test query", "top_k": 5}, + timeout=30, + ) + + +@patch("requests.get") +def test_client_status_and_documents_mock(mock_get): + """Verify client.get_status and client.list_documents.""" + mock_status_resp = MagicMock() + mock_status_resp.json.return_value = {"status": "online", "indexed_documents": 10} + mock_docs_resp = MagicMock() + mock_docs_resp.json.return_value = {"documents": [{"file_hash": "abc"}]} + + mock_get.side_effect = [mock_status_resp, mock_docs_resp] + + client = BBQClient() + status = client.get_status() + assert status["status"] == "online" + + docs = client.list_documents() + assert len(docs) == 1 + assert docs[0]["file_hash"] == "abc" + + +@patch("requests.get") +def test_client_get_page_image_mock(mock_get, tmp_path): + """Verify client.get_page_image retrieves image bytes and optionally saves to disk.""" + # Create simple dummy PNG image bytes + img = Image.new("RGB", (10, 10), color="blue") + buf = io.BytesIO() + img.save(buf, format="PNG") + png_bytes = buf.getvalue() + + mock_resp = MagicMock() + mock_resp.content = png_bytes + mock_get.return_value = mock_resp + + client = BBQClient() + save_file = str(tmp_path / "page_1.png") + ret_img = client.get_page_image("test.pdf", page_number=1, save_path=save_file) + + assert isinstance(ret_img, Image.Image) + assert (tmp_path / "page_1.png").exists() + + +def test_gemini_client_availability(): + """Verify GeminiClient availability check.""" + client_no_key = GeminiClient(api_key="") + assert client_no_key.is_available() is False + + client_with_key = GeminiClient(api_key="valid-key-xyz") + assert client_with_key.is_available() is True + assert client_with_key.model == "gemini-3.6-flash" + + +def test_main_cli_without_logo_options(): + """Verify that CLI parser flags --without-logo and --no-logo are correctly recognized.""" + parser = build_argument_parser() + + # Flag before query command + args1 = parser.parse_args(["--without-logo", "query", "what is this?"]) + assert args1.without_logo is True + assert args1.query == "what is this?" + + # Flag after query command + args2 = parser.parse_args(["query", "--without-logo", "what is this?"]) + assert args2.without_logo is True + + # --no-logo alias + args3 = parser.parse_args(["query", "--no-logo", "what is this?"]) + assert args3.without_logo is True + + # Default without flag + args4 = parser.parse_args(["query", "what is this?"]) + assert args4.without_logo is False + + # Status command with --without-logo + args5 = parser.parse_args(["status", "--without-logo"]) + assert args5.without_logo is True + + # Client alias command + args6 = parser.parse_args(["client", "find documents"]) + assert args6.query == "find documents" + + +def test_client_main_cli_parser(): + """Verify that bbq.src.client.__main__ parser recognizes flags.""" + parser = build_client_main_parser() + args = parser.parse_args(["my search", "--without-logo", "--top-k", "4"]) + assert args.query == "my search" + assert args.without_logo is True + assert args.top_k == 4 + + +def test_cli_use_llm_and_default_top_k(): + """Verify --use-llm flag and default top_k=None in parsers (deferring to config.rag_top_k).""" + # Main query parser + main_parser = build_argument_parser() + args_default = main_parser.parse_args(["query", "search text"]) + assert args_default.top_k is None + assert args_default.use_llm is False + + args_explicit = main_parser.parse_args(["query", "--top-k", "15", "search text"]) + assert args_explicit.top_k == 15 + + args_llm = main_parser.parse_args(["query", "--use-llm", "search text"]) + assert args_llm.use_llm is True + + # Client __main__ parser + client_parser = build_client_main_parser() + c_args_default = client_parser.parse_args(["search text"]) + assert c_args_default.top_k is None + assert c_args_default.use_llm is False + + c_args_explicit = client_parser.parse_args(["search text", "--top-k", "5"]) + assert c_args_explicit.top_k == 5 + + c_args_llm = client_parser.parse_args(["search text", "--use-llm"]) + assert c_args_llm.use_llm is True + + +@patch.object(BBQClient, "query") +def test_client_query_and_answer_without_llm(mock_query): + """Verify that query_and_answer with use_llm=False returns retrieved results without calling Gemini.""" + mock_query.return_value = [{"file_path": "doc.pdf", "page_number": 1, "score": 3.0}] + client = BBQClient() + + resp = client.query_and_answer("query text", top_k=10, use_llm=False) + assert resp["status"] == "retrieval_only" + assert resp["answer"] is None + assert len(resp["sources"]) == 1 + mock_query.assert_called_once_with(query_text="query text", top_k=10) + + +@patch.object(BBQClient, "generate_answer_from_results") +@patch.object(BBQClient, "query") +def test_client_query_and_answer_with_llm(mock_query, mock_generate): + """Verify that query_and_answer with use_llm=True calls generate_answer_from_results.""" + mock_query.return_value = [{"file_path": "doc.pdf", "page_number": 1, "score": 3.0}] + mock_generate.return_value = { + "query": "query text", + "answer": "Generated multimodal answer", + "sources": [{"file_path": "doc.pdf", "page_number": 1, "score": 3.0}], + "status": "success", + "engine": "gemini", + } + client = BBQClient() + + resp = client.query_and_answer("query text", top_k=10, use_llm=True) + assert resp["status"] == "success" + assert resp["answer"] == "Generated multimodal answer" + mock_generate.assert_called_once() + + +def test_cli_infinite_flags_parsing(): + """Verify --infinite and -inf flag parsing in both main and client __main__ parsers.""" + main_parser = build_argument_parser() + + args1 = main_parser.parse_args(["query", "--infinite"]) + assert args1.infinite is True + assert args1.query is None + + args2 = main_parser.parse_args(["client", "-inf", "initial query"]) + assert args2.infinite is True + assert args2.query == "initial query" + + client_parser = build_client_main_parser() + c_args1 = client_parser.parse_args(["--infinite"]) + assert c_args1.infinite is True + assert c_args1.query is None + + c_args2 = client_parser.parse_args(["-inf", "first query"]) + assert c_args2.infinite is True + assert c_args2.query == "first query" + + +@patch("builtins.input") +@patch.object(BBQClient, "query") +def test_main_cli_infinite_mode_interactive_loop(mock_query, mock_input, capsys): + """Verify interactive loop runs with initial query and prompts for next until exit.""" + from bbq.src.main import run_client_query_command + + mock_query.return_value = [ + {"file_path": "doc.pdf", "page_number": 1, "score": 2.5, "file_hash": "hash123"} + ] + # Simulate user input after the initial query finishes: first another query, then 'exit' + mock_input.side_effect = ["second query", "exit"] + + parser = build_argument_parser() + args = parser.parse_args( + ["query", "first query", "--top-k", "10", "--infinite", "--without-logo"] + ) + run_client_query_command(args) + + # query should have been called twice: first query, then second query + assert mock_query.call_count == 2 + mock_query.assert_any_call(query_text="first query", top_k=10) + mock_query.assert_any_call(query_text="second query", top_k=10) + + out = capsys.readouterr().out + assert "Exiting interactive query mode." in out + + +@patch("builtins.input") +@patch.object(BBQClient, "query") +def test_main_cli_infinite_mode_keyboard_interrupt(mock_query, mock_input, capsys): + """Verify interactive loop terminates gracefully on KeyboardInterrupt.""" + from bbq.src.main import run_client_query_command + + mock_input.side_effect = KeyboardInterrupt() + + parser = build_argument_parser() + # No initial query provided, goes straight to input prompt + args = parser.parse_args(["query", "--infinite", "--without-logo"]) + run_client_query_command(args) + + mock_query.assert_not_called() + out = capsys.readouterr().out + assert "Exiting interactive query mode." in out + + +def test_cli_version_flag(capsys): + """Verify that --version dynamically outputs the package version and exits.""" + parser = build_argument_parser() + with pytest.raises(SystemExit) as exc_info: + parser.parse_args(["--version"]) + assert exc_info.value.code == 0 + captured = capsys.readouterr().out + assert __version__ in captured + assert "bbq" in captured + + +@patch("builtins.input") +def test_main_cli_interactive_help_command(mock_input, capsys): + """Verify that entering 'help' displays the interactive commands menu.""" + from bbq.src.main import run_client_query_command + + mock_input.side_effect = ["help", "q"] + + parser = build_argument_parser() + args = parser.parse_args(["query", "--infinite", "--without-logo"]) + run_client_query_command(args) + + out = capsys.readouterr().out + assert "bbq[query] Interactive Commands" in out + assert "Enter any query" in out + assert "q to exit" in out or "quit" in out + + +def test_render_rich_components_no_tables(capsys): + """Verify render_query_results_rich and render_llm_answer_rich render without crashing.""" + from bbq.src.terminal import render_query_results_rich, render_llm_answer_rich + + sample_results = [ + { + "score": 5.7152, + "file_path": "data/watch/puzzles.pdf", + "page_number": 7, + "total_pages": 98, + "file_hash": "0cf7e26873a3", + } + ] + render_query_results_rich(sample_results) + out = capsys.readouterr().out + assert "Top 1 Matching PDF Pages" in out + assert "Page" in out and "7" in out + assert "puzzles.pdf" in out + + render_llm_answer_rich("This is a synthesized multimodal answer.") + llm_out = capsys.readouterr().out + assert "GEMINI MULTIMODAL ANSWER" in llm_out + assert "synthesized multimodal answer" in llm_out + + +@patch("builtins.input") +@patch("rich.console.Console.clear") +def test_main_cli_interactive_clear_command(mock_clear, mock_input, capsys): + """Verify that entering 'clear' or 'cls' triggers console.clear().""" + from bbq.src.main import run_client_query_command + + mock_input.side_effect = ["clear", "cls", "q"] + + parser = build_argument_parser() + args = parser.parse_args(["query", "--infinite", "--without-logo"]) + run_client_query_command(args) + + assert mock_clear.call_count == 2 + out = capsys.readouterr().out + assert "Exiting interactive query mode. Bye!" in out + + +def test_open_document_page(monkeypatch): + """Verify open_document_page calls subprocess.Popen with page arguments.""" + from bbq.src.terminal import open_document_page + + mock_calls = [] + + def mock_popen(args, **kwargs): + mock_calls.append(args) + return MagicMock() + + monkeypatch.setattr("subprocess.Popen", mock_popen) + monkeypatch.setattr("os.path.exists", lambda p: True) + + success = open_document_page("data/watch/test.pdf", page_number=42) + assert success is True + assert len(mock_calls) == 1 + # Check that page argument was passed + call_args = mock_calls[0] + assert "42" in call_args + + +def test_interactive_page_picker_non_tty(capsys): + """Verify interactive_page_picker renders cleanly when stdin is not a tty.""" + from bbq.src.terminal import interactive_page_picker + + sample_results = [ + { + "score": 4.5, + "file_path": "data/watch/design.pdf", + "page_number": 128, + "total_pages": 241, + } + ] + interactive_page_picker(sample_results) + out = capsys.readouterr().out + assert "Top 1 Matching PDF Pages" in out + 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 + from bbq.src.client import BBQClient + + mock_picker_called = [] + + monkeypatch.setattr( + BBQClient, + "query", + lambda self, query_text, top_k=10: [ + {"file_path": "doc.pdf", "page_number": 5, "total_pages": 10, "score": 3.0} + ], + ) + monkeypatch.setattr( + "bbq.src.terminal.interactive_page_picker", + lambda results, console=None: mock_picker_called.append(True), + ) + + parser = build_argument_parser() + args = parser.parse_args(["query", "search", "--without-opener", "--without-logo"]) + assert args.without_opener is True + run_client_query_command(args) + + assert len(mock_picker_called) == 0 + out = capsys.readouterr().out + assert "Top 1 Matching PDF Pages" in out diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..c377b4d --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,163 @@ +import os +import pytest +from bbq.src.config import Config +from bbq.src.client import GeminiClient +from bbq.src.utils.futils import get_system_cache_dir +from bbq.src.common.errors import ConfigFNFWarning, ConfigParseError + + +def test_system_cache_dir(): + """Verify system cache path computation and subfolder appending.""" + base_dir = get_system_cache_dir() + assert "bbq" in base_dir + assert not base_dir.startswith("./.cache") + + sub_dir = get_system_cache_dir("test_folder") + assert sub_dir == os.path.join(base_dir, "test_folder") + + +def test_config_defaults(): + """Verify Config defaults and dynamic post-init paths.""" + config = Config() + assert config.base_model_id == "HuggingFaceTB/SmolVLM-256M-Instruct" + assert config.lora_adapter_id == "vidore/colSmol-256M" + assert config.embedding_dim == 128 + assert config.device == "auto" + assert config.torch_dtype == "bfloat16" + assert config.mask_non_image_embeddings is False + assert config.pdf_render_dpi == 150 + assert config.embeddings_output_path == get_system_cache_dir("embeddings") + assert config.sqlite_db_path == get_system_cache_dir("tracker.db") + assert config.gemini_model == "gemini-3.6-flash" + assert config.rag_top_k == 3 + assert config.quantization == "f32" + + +def test_config_custom_initialization(): + """Verify custom values and type casting in Config.""" + config = Config( + base_model_id="custom/model", + embedding_dim=256, + pdf_render_dpi=200, + rag_top_k=5, + gemini_api_key="TEST_API_KEY_XYZ", + gemini_model="gemini-2.0-flash", + ) + assert config.base_model_id == "custom/model" + assert config.embedding_dim == 256 + assert config.pdf_render_dpi == 200 + assert config.rag_top_k == 5 + assert config.gemini_api_key == "TEST_API_KEY_XYZ" + assert config.gemini_model == "gemini-2.0-flash" + + +def test_config_empty_base_model_raises(): + """Verify that an empty base_model_id raises a ValueError.""" + with pytest.raises(ValueError, match="base_model_id cannot be empty"): + Config(base_model_id="") + + +def test_config_visual_prompt_prefix(): + """Verify visual prompt formatting.""" + config = Config(visual_prompt_command="Describe this document.") + prefix = config.format_visual_prompt_prefix() + assert ( + prefix + == "<|im_start|>User:Describe this document.\nAssistant:" + ) + + # If already formatted, keep as is + already_formatted = "<|im_start|>custom" + config_custom = Config(visual_prompt_command=already_formatted) + assert config_custom.format_visual_prompt_prefix() == already_formatted + + +def test_config_dictionary_format(): + """Verify dictionary serialization masks API key and includes helper keys.""" + config = Config( + gemini_api_key="SECRET_KEY_123", + gemini_model="gemini-2.0-flash", + rag_top_k=5, + ) + d = config.convert_to_dictionary_format() + assert "gemini_api_key" not in d + assert d["gemini_api_key_configured"] is True + assert d["gemini_model"] == "gemini-2.0-flash" + assert d["rag_top_k"] == 5 + assert "visual_prompt_prefix" in d + + +def test_gemini_client_from_config(): + """Verify GeminiClient.from_config factory method.""" + config = Config( + gemini_api_key="TEST_KEY_FOR_CLIENT", + gemini_model="gemini-2.0-flash", + ) + client = GeminiClient.from_config(config) + assert client.api_key == "TEST_KEY_FOR_CLIENT" + assert client.model == "gemini-2.0-flash" + assert client.is_available() is True + + +def test_config_from_yaml(tmp_path): + """Verify loading configuration from a YAML file.""" + yaml_file = tmp_path / "custom_config.yaml" + yaml_file.write_text(""" + base_model_id: "test/model" + embedding_dim: 64 + rag_top_k: 10 + gemini_model: "gemini-2.5-flash" + """) + + config = Config.from_yaml(str(yaml_file)) + assert config.base_model_id == "test/model" + assert config.embedding_dim == 64 + assert config.rag_top_k == 10 + assert config.gemini_model == "gemini-2.5-flash" + # Unspecified fields use dataclass defaults + assert config.device == "auto" + assert config.torch_dtype == "bfloat16" + + +def test_config_from_yaml_missing_file(): + """Verify that a non-existent YAML file triggers a warning and returns defaults.""" + with pytest.warns(ConfigFNFWarning, match="Configuration file not found"): + config = Config.from_yaml("non_existent_config_file_xyz.yaml") + assert isinstance(config, Config) + assert config.base_model_id == "HuggingFaceTB/SmolVLM-256M-Instruct" + + +def test_config_from_yaml_invalid_syntax(tmp_path): + """Verify that invalid YAML syntax raises ConfigParseError.""" + invalid_yaml = tmp_path / "invalid.yaml" + invalid_yaml.write_text("::: invalid yaml syntax :::") + + with pytest.raises(ConfigParseError): + Config.from_yaml(str(invalid_yaml)) + + +def test_config_quantization(): + """Verify default and customized quantization handling.""" + config_default = Config() + assert config_default.quantization == "f32" + + config_qi8 = Config(quantization="QI8") + assert config_qi8.quantization == "qi8" + + +def test_dependency_injection_server_and_client(): + """Verify that components accept an injected Config without reloading from YAML.""" + from bbq.src.client import BBQClient + from bbq.src.server import BBQServer + + custom_config = Config(embedding_dim=64, rag_top_k=9) + + server = BBQServer(config=custom_config) + assert server.config is custom_config + assert server.config.embedding_dim == 64 + assert server.config.rag_top_k == 9 + + client = BBQClient(config=custom_config) + assert client.config is custom_config + assert client.config is not None + assert client.config.embedding_dim == 64