tools/query_knowledge.py

from smolagents.tools import Tool
import asyncio
from collections import Counter
import hashlib
import logging
import math
import os
from pathlib import Path
from typing import Iterable, List, Optional
import re
from dotenv import load_dotenv
import chromadb
from chromadb.errors import NotFoundError
from pypdf import PdfReader

from llama_index.core import StorageContext, VectorStoreIndex
from llama_index.core.schema import Document, BaseNode, NodeWithScore, TextNode
from llama_index.core.node_parser import SentenceSplitter
from llama_index.vector_stores.chroma import ChromaVectorStore


load_dotenv()
BASE_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = BASE_DIR.parent
KNOWLEDGE_BASE_DIR = BASE_DIR / "knowledge_base"
LEGACY_KNOWLEDGE_BASE_DIR = BASE_DIR / "knowledge_base"
KNOWLEDGE_BASE_DIR = PROJECT_ROOT / "knowledge_base"
RAW_DIR = KNOWLEDGE_BASE_DIR / "raw"
CHROMA_DB_DIR = KNOWLEDGE_BASE_DIR / "chroma_db"
HF_CACHE_DIR = PROJECT_ROOT / "hf_cache"
COLLECTION_NAME = "options_knowledge"

EMBED_MODEL_NAME = os.getenv("RAG_EMBED_MODEL", "BAAI/bge-small-en-v1.5")
RERANKER_MODEL_NAME = os.getenv(
    "RAG_RERANKER_MODEL", "cross-encoder/ms-marco-MiniLM-L-6-v2")
RERANKER_BATCH_SIZE = int(os.getenv("RAG_RERANKER_BATCH_SIZE", "16"))
EMBED_MODEL_METADATA_KEY = "embedding_model"
BM25_METADATA_KEY = "bm25_score"
VECTOR_METADATA_KEY = "vector_score"
CHUNK_SIZE = 1000
CHUNK_OVERLAP = 150
PDF_REPEATED_LINE_MIN_PAGES = 3
PDF_BOUNDARY_LINE_COUNT = 4
PDF_EXTRACTION_METHOD = "pymupdf_formula_blocks_v5"
PDF_LINE_Y_TOLERANCE = 3.0
PDF_MIN_SECTION_CHARS = 240
PDF_STRONG_MATH_SYMBOLS = set("=∂∫∑∏√∞≈≠≤≥±×÷^_σΣΔδθΘλΛμρπΠφΦτν𝜎𝜇𝜌𝜃𝜕")
PDF_WEAK_MATH_SYMBOLS = set("+-−*/∕<>")
PDF_MATH_SYMBOLS = PDF_STRONG_MATH_SYMBOLS | PDF_WEAK_MATH_SYMBOLS
PDF_OPERATOR_MATH_SYMBOLS = set("=∂∫∑∏√∞≈≠≤≥±×÷^_+-−*/∕<>")
PDF_FORMULA_TRIGGER_SYMBOLS = set("=∂∫∑∏√∞≈≠≤≥±×÷^_∕<>")

logging.getLogger("pypdf").setLevel(logging.ERROR)


def load_pymupdf():
    try:
        import fitz
    except ImportError:
        return None

    return fitz


REQUIRED_METADATA = [
    "source_file",
    "file_name",
    "file_type",
    "document_title",
    "file_hash",
    "chunk_id",
    "chunk_index",
]


def configure_model_cache() -> None:
    HF_CACHE_DIR.mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("HF_HOME", str(HF_CACHE_DIR))
    os.environ.setdefault("SENTENCE_TRANSFORMERS_HOME", str(
        HF_CACHE_DIR / "sentence_transformers"))
    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
    if local_model_snapshot(EMBED_MODEL_NAME):
        os.environ.setdefault("HF_HUB_OFFLINE", "1")
        os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")


def local_model_snapshot(model_name: str) -> Optional[Path]:
    cached_model_dir = (
        HF_CACHE_DIR
        / "sentence_transformers"
        / f"models--{model_name.replace('/', '--')}"
    )
    snapshots_dir = cached_model_dir / "snapshots"
    if snapshots_dir.exists():
        snapshots = sorted(path for path in snapshots_dir.iterdir() if path.is_dir())
        for snapshot in reversed(snapshots):
            if (snapshot / "config.json").exists():
                return snapshot

    return None


def resolve_embed_model_name() -> str:
    snapshot = local_model_snapshot(EMBED_MODEL_NAME)
    if snapshot:
        return str(snapshot)

    return EMBED_MODEL_NAME


def resolve_reranker_model_name(model_name: str = RERANKER_MODEL_NAME) -> str:
    snapshot = local_model_snapshot(model_name)
    if snapshot:
        return str(snapshot)

    return model_name


def env_flag(name: str, default: bool = False) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


def effective_raw_dir(raw_dir: Path = RAW_DIR) -> Path:
    if any(iter_source_files(raw_dir)):
        return raw_dir

    legacy_raw_dir = LEGACY_KNOWLEDGE_BASE_DIR / "raw"
    if any(iter_source_files(legacy_raw_dir)):
        logging.warning(
            "Using legacy knowledge base path %s. Move files to %s when convenient.",
            legacy_raw_dir,
            raw_dir,
        )
        return legacy_raw_dir

    return raw_dir


class CrossEncoderReranker:
    def __init__(
        self,
        model_name: str = RERANKER_MODEL_NAME,
        batch_size: int = RERANKER_BATCH_SIZE,
    ):
        self.model_name = model_name
        self.batch_size = batch_size
        self._model = None

    def _load_model(self):
        if self._model is not None:
            return self._model

        from sentence_transformers import CrossEncoder

        self._model = CrossEncoder(
            resolve_reranker_model_name(self.model_name),
            max_length=512,
            cache_folder=str(HF_CACHE_DIR / "sentence_transformers"),
        )
        return self._model

    def rerank(
        self,
        query: str,
        results: list[NodeWithScore],
        top_n: Optional[int] = None,
    ) -> list[NodeWithScore]:
        if not results:
            return []

        pairs = [
            (query, result.node.get_content())
            for result in results
        ]
        model = self._load_model()
        scores = model.predict(
            pairs,
            batch_size=self.batch_size,
            show_progress_bar=False,
        )

        reranked = [
            NodeWithScore(node=result.node, score=float(score))
            for result, score in zip(results, scores)
        ]
        reranked.sort(key=lambda item: item.score or float("-inf"), reverse=True)
        return reranked[:top_n] if top_n else reranked


class BM25Retriever:
    def __init__(self, nodes: list[TextNode]):
        self.nodes = nodes
        self.tokenized_docs = [self.tokenize(node.get_content()) for node in nodes]
        self.doc_freqs: Counter[str] = Counter()
        for tokens in self.tokenized_docs:
            self.doc_freqs.update(set(tokens))
        self.avg_doc_len = (
            sum(len(tokens) for tokens in self.tokenized_docs) / len(self.tokenized_docs)
            if self.tokenized_docs
            else 0.0
        )

    @staticmethod
    def tokenize(text: str) -> list[str]:
        return [
            token.lower()
            for token in re.findall(r"[A-Za-z]+(?:[-'][A-Za-z]+)*|\d+(?:\.\d+)*|[^\sA-Za-z0-9]", text)
            if token.strip()
        ]

    def score(self, query_tokens: list[str], doc_tokens: list[str]) -> float:
        if not query_tokens or not doc_tokens:
            return 0.0

        token_counts = Counter(doc_tokens)
        doc_len = len(doc_tokens)
        total_docs = len(self.tokenized_docs)
        k1 = 1.5
        b = 0.75
        score = 0.0

        for token in query_tokens:
            term_freq = token_counts.get(token, 0)
            if term_freq == 0:
                continue
            doc_freq = self.doc_freqs.get(token, 0)
            idf = math.log(1 + (total_docs - doc_freq + 0.5) / (doc_freq + 0.5))
            denominator = term_freq + k1 * (
                1 - b + b * doc_len / max(self.avg_doc_len, 1.0)
            )
            score += idf * (term_freq * (k1 + 1)) / denominator

        return score

    def retrieve(self, query: str, top_k: int) -> list[NodeWithScore]:
        query_tokens = self.tokenize(query)
        scored: list[NodeWithScore] = []

        for node, doc_tokens in zip(self.nodes, self.tokenized_docs):
            score = self.score(query_tokens, doc_tokens)
            if score <= 0:
                continue
            node.metadata[BM25_METADATA_KEY] = score
            scored.append(NodeWithScore(node=node, score=score))

        scored.sort(key=lambda item: item.score or float("-inf"), reverse=True)
        return scored[:top_k]


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as file:
        for block in iter(lambda: file.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def load_md_file(path: Path) -> Document:
    text = path.read_text(encoding="utf-8")

    return Document(
        text=text,
        metadata={
            "source_file": str(path.resolve()),
            "file_name": path.name,
            "file_type": "md",
            "document_title": path.stem,
            "file_hash": file_sha256(path),
        },
    )


def load_md_documents(path: Path) -> List[Document]:
    text = path.read_text(encoding="utf-8")
    file_hash = file_sha256(path)
    documents: List[Document] = []
    current_heading = ""
    current_lines: List[str] = []

    def flush() -> None:
        nonlocal current_lines
        section_text = "\n".join(current_lines).strip()
        if not section_text:
            current_lines = []
            return
        documents.append(
            Document(
                text=section_text,
                metadata={
                    "source_file": str(path.resolve()),
                    "file_name": path.name,
                    "file_type": path.suffix.lower().lstrip("."),
                    "document_title": path.stem,
                    "file_hash": file_hash,
                    "content_type": "markdown_section",
                    "chapter_title": "",
                    "section_title": current_heading,
                    "section_path": current_heading,
                    "char_count": len(section_text),
                },
            )
        )
        current_lines = []

    for line in text.splitlines():
        heading_match = re.match(r"^(#{1,6})\s+(.+?)\s*$", line)
        if heading_match:
            flush()
            current_heading = heading_match.group(2).strip()
        current_lines.append(line)

    flush()
    return documents or [load_md_file(path)]


def append_visual_fragment(line_parts: List[str], text: str, baseline_y: float, item: dict) -> None:
    if not text:
        return

    stripped = text.strip()
    if not stripped:
        return

    font_size = item["font_size"]
    y_offset = item["y"] - baseline_y
    is_small = font_size < item["line_font_size"] * 0.82

    if is_small and y_offset > max(1.5, item["line_font_size"] * 0.18):
        line_parts.append(f"^{{{stripped}}}")
    elif is_small and y_offset < -max(1.5, item["line_font_size"] * 0.18):
        line_parts.append(f"_{{{stripped}}}")
    else:
        line_parts.append(stripped)


def join_visual_line(items: List[dict]) -> str:
    if not items:
        return ""

    items = sorted(items, key=lambda value: value["x"])
    baseline_y = sorted(item["y"] for item in items)[len(items) // 2]
    line_font_size = max(item["font_size"] for item in items)
    previous_right = None
    line_parts: List[str] = []

    for item in items:
        item["line_font_size"] = line_font_size
        if previous_right is not None:
            gap = item["x"] - previous_right
            if gap > max(2.5, line_font_size * 0.28):
                line_parts.append(" ")

        append_visual_fragment(line_parts, item["text"], baseline_y, item)
        previous_right = max(previous_right or item["x"], item["x"] + item["width"])

    return normalize_pdf_line("".join(line_parts))


def extract_pdf_text_by_position(page) -> str:
    fragments: List[dict] = []

    def visitor_text(text, cm, tm, font_dict, font_size):
        if not text or not text.strip():
            return

        x = float(tm[4])
        y = float(tm[5])
        width = max(len(text.strip()) * float(font_size) * 0.45, float(font_size))
        fragments.append(
            {
                "text": text,
                "x": x,
                "y": y,
                "width": width,
                "font_size": float(font_size or 1.0),
            }
        )

    try:
        page.extract_text(visitor_text=visitor_text)
    except Exception:
        return ""

    if not fragments:
        return ""

    lines: List[List[dict]] = []
    for fragment in sorted(fragments, key=lambda value: (-value["y"], value["x"])):
        for line in lines:
            if abs(line[0]["y"] - fragment["y"]) <= PDF_LINE_Y_TOLERANCE:
                line.append(fragment)
                break
        else:
            lines.append([fragment])

    return "\n".join(join_visual_line(line) for line in lines)


def math_text_score(text: str) -> float:
    if not text.strip():
        return 0.0

    lines = [line for line in text.splitlines() if line.strip()]
    compact_length = len(re.sub(r"\s+", "", text))
    math_symbol_count = sum(1 for char in text if char in PDF_MATH_SYMBOLS)
    superscript_markers = text.count("^{") + text.count("_{")
    multiline_bonus = sum(1 for line in lines if is_formula_like(line)) * 8
    equation_block_bonus = sum(
        1
        for index, line in enumerate(lines)
        if is_formula_like(line)
        and (
            index > 0
            and is_formula_like(lines[index - 1])
            or index + 1 < len(lines)
            and is_formula_like(lines[index + 1])
        )
    ) * 12
    return (
        compact_length
        + math_symbol_count * 12
        + superscript_markers * 20
        + multiline_bonus
        + equation_block_bonus
    )


def extract_pdf_text(page) -> str:
    positioned_text = extract_pdf_text_by_position(page)

    try:
        layout_text = page.extract_text(extraction_mode="layout") or ""
    except Exception:
        layout_text = ""

    try:
        plain_text = page.extract_text() or ""
    except Exception:
        plain_text = ""

    candidates = [positioned_text, layout_text, plain_text]
    candidates = [candidate for candidate in candidates if candidate.strip()]
    if not candidates:
        return ""

    return max(candidates, key=math_text_score)


def pymupdf_span_text(span: dict) -> str:
    return normalize_pdf_line(span.get("text", ""))


def pymupdf_line_text(line: dict) -> str:
    return normalize_pdf_line("".join(pymupdf_span_text(span) for span in line.get("spans", [])))


def pymupdf_block_text(block: dict) -> str:
    lines = [
        pymupdf_line_text(line)
        for line in block.get("lines", [])
    ]
    return "\n".join(line for line in lines if line)


def pymupdf_span_has_math_font(span: dict) -> bool:
    font_name = span.get("font", "").lower()
    return any(
        marker in font_name
        for marker in ("math", "symbol", "cmmi", "cmsy", "cmex", "stix")
    )


def is_formula_block_line(line: str) -> bool:
    stripped = line.strip()
    if not stripped:
        return False

    trigger_math_count = sum(1 for char in stripped if char in PDF_FORMULA_TRIGGER_SYMBOLS)
    digit_count = sum(1 for char in stripped if char.isdigit())
    alpha_count = sum(1 for char in stripped if char.isalpha())
    alpha_words = [
        word
        for word in re.findall(r"[A-Za-z]+", stripped)
        if word.lower() not in {"and", "or", "the", "where", "then", "with", "for"}
    ]
    compact_length = len(re.sub(r"\s+", "", stripped))

    if compact_length < 3:
        return False
    if re.fullmatch(r"\(?\d+(\.\d+)?\)?", stripped):
        return False
    if re.search(r"\(\d+(\.\d+)+[a-z]?\)$", stripped) and compact_length <= 240:
        return True
    if "=" in stripped and compact_length <= 260 and len(alpha_words) <= 12:
        return True
    if any(char in stripped for char in "∂∫∑∏√∞≈≠≤≥±×÷") and compact_length <= 220 and len(alpha_words) <= 10:
        return True
    if trigger_math_count >= 2 and compact_length <= 120 and len(alpha_words) <= 6:
        return True
    if trigger_math_count >= 1 and digit_count >= 1 and alpha_count <= 18 and compact_length <= 100:
        return True

    return False


def is_formula_block(block: dict) -> bool:
    text = pymupdf_block_text(block)
    if not text:
        return False

    lines = [line for line in text.splitlines() if line.strip()]
    if any(is_formula_block_line(line) for line in lines):
        return True

    spans = [
        span
        for line in block.get("lines", [])
        for span in line.get("spans", [])
        if pymupdf_span_text(span)
    ]
    if not spans:
        return False

    math_font_count = sum(1 for span in spans if pymupdf_span_has_math_font(span))
    strong_math_count = sum(1 for char in text if char in PDF_STRONG_MATH_SYMBOLS)
    alpha_count = sum(1 for char in text if char.isalpha())
    digit_count = sum(1 for char in text if char.isdigit())
    compact_length = len(re.sub(r"\s+", "", text))

    if math_font_count >= 2 and compact_length <= 220:
        return True
    if strong_math_count >= 3 and compact_length <= 260:
        return True
    if strong_math_count >= 1 and digit_count >= 1 and alpha_count <= 20 and compact_length <= 160:
        return True

    return False


def block_bbox_string(block: dict) -> str:
    bbox = block.get("bbox") or []
    if len(bbox) != 4:
        return ""
    return ",".join(f"{float(value):.2f}" for value in bbox)


def line_bbox_string(line: dict) -> str:
    bbox = line.get("bbox") or []
    if len(bbox) != 4:
        return ""
    return ",".join(f"{float(value):.2f}" for value in bbox)


def pymupdf_line_has_math_font(line: dict) -> bool:
    return any(
        pymupdf_span_has_math_font(span)
        for span in line.get("spans", [])
        if pymupdf_span_text(span)
    )


def should_extract_formula_line(line: dict) -> bool:
    text = pymupdf_line_text(line)
    if not text:
        return False

    if is_formula_block_line(text):
        return True

    compact_length = len(re.sub(r"\s+", "", text))
    trigger_math_count = sum(1 for char in text if char in PDF_FORMULA_TRIGGER_SYMBOLS)
    alpha_words = re.findall(r"[A-Za-z]+", text)
    if (
        pymupdf_line_has_math_font(line)
        and trigger_math_count >= 1
        and compact_length <= 180
        and len(alpha_words) <= 6
    ):
        return True

    return False


def is_formula_continuation_line(text: str) -> bool:
    stripped = text.strip()
    if not stripped:
        return False

    compact = re.sub(r"\s+", "", stripped)
    if len(compact) > 90:
        return False
    if compact in {"(", ")", "[", "]", "{", "}", "√"}:
        return True

    alpha_words = re.findall(r"[A-Za-z]+", stripped)
    math_count = sum(1 for char in stripped if char in PDF_MATH_SYMBOLS)
    digit_count = sum(1 for char in stripped if char.isdigit())

    if len(alpha_words) <= 4 and (math_count >= 1 or digit_count >= 1):
        return True

    return False


def append_formula_block(
    formula_blocks: List[dict],
    body_blocks: List[str],
    page_number: int,
    formula_index: int,
    formula_lines: List[str],
    formula_bboxes: List[str],
) -> int:
    formula_text = clean_formula_text("\n".join(formula_lines))
    if not is_useful_formula_text(formula_text):
        return formula_index

    formula_id = f"formula-{page_number}-{formula_index}"
    formula_bbox = merge_bbox_strings(formula_bboxes)
    formula_blocks.append(
        {
            "id": formula_id,
            "text": formula_text,
            "bbox": formula_bbox,
        }
    )
    body_blocks.append(f"[FORMULA id={formula_id}]\n{formula_text}\n[/FORMULA]")
    return formula_index + 1


def merge_bbox_strings(bbox_strings: List[str]) -> str:
    boxes = []
    for bbox_string in bbox_strings:
        if not bbox_string:
            continue
        values = bbox_string.split(",")
        if len(values) != 4:
            continue
        try:
            boxes.append([float(value) for value in values])
        except ValueError:
            continue

    if not boxes:
        return ""

    x0 = min(box[0] for box in boxes)
    y0 = min(box[1] for box in boxes)
    x1 = max(box[2] for box in boxes)
    y1 = max(box[3] for box in boxes)
    return f"{x0:.2f},{y0:.2f},{x1:.2f},{y1:.2f}"


def is_useful_formula_text(text: str) -> bool:
    stripped = text.strip()
    if not stripped:
        return False

    compact_length = len(re.sub(r"\s+", "", stripped))
    if compact_length < 6:
        return False

    lines = [line.strip() for line in stripped.splitlines() if line.strip()]
    if re.search(r"\(\d+(\.\d+)+[a-z]?\)", stripped):
        return True
    if any(char in stripped for char in "∂∫∑∏∞≈≠≤≥±×÷"):
        alpha_words = re.findall(r"[A-Za-z]+", stripped)
        return len(alpha_words) <= 12 or "=" in stripped

    for line in lines:
        if "=" not in line:
            continue

        alpha_words = [
            word
            for word in re.findall(r"[A-Za-z]+", line)
            if word.lower() not in {"and", "or", "the", "where", "then", "with", "for"}
        ]
        if len(alpha_words) <= 12 and len(line) <= 260:
            return True

    return False


def extract_pymupdf_page(page) -> dict:
    page_dict = page.get_text("dict", sort=True)
    body_blocks: List[str] = []
    formula_blocks: List[dict] = []
    formula_lines: List[str] = []
    formula_bboxes: List[str] = []
    formula_index = 0
    page_number = page.number + 1

    for block in page_dict.get("blocks", []):
        if block.get("type") != 0:
            continue

        normal_lines: List[str] = []

        for line in block.get("lines", []):
            line_text = pymupdf_line_text(line)
            if not line_text:
                continue

            if should_extract_formula_line(line) or (
                formula_lines and is_formula_continuation_line(line_text)
            ):
                if normal_lines:
                    body_blocks.append("\n".join(normal_lines))
                    normal_lines = []
                formula_lines.append(line_text)
                formula_bboxes.append(line_bbox_string(line))
            else:
                if formula_lines:
                    formula_index = append_formula_block(
                        formula_blocks=formula_blocks,
                        body_blocks=body_blocks,
                        page_number=page_number,
                        formula_index=formula_index,
                        formula_lines=formula_lines,
                        formula_bboxes=formula_bboxes,
                    )
                    formula_lines = []
                    formula_bboxes = []
                normal_lines.append(line_text)

        if normal_lines:
            body_blocks.append("\n".join(normal_lines))

    if formula_lines:
        append_formula_block(
            formula_blocks=formula_blocks,
            body_blocks=body_blocks,
            page_number=page_number,
            formula_index=formula_index,
            formula_lines=formula_lines,
            formula_bboxes=formula_bboxes,
        )

    return {
        "text": "\n".join(body_blocks),
        "formula_blocks": formula_blocks,
        "backend": "pymupdf",
    }


def extract_pdf_pages_with_pymupdf(path: Path) -> Optional[List[dict]]:
    fitz = load_pymupdf()
    if fitz is None:
        return None

    try:
        document = fitz.open(str(path))
    except Exception:
        return None

    try:
        return [extract_pymupdf_page(page) for page in document]
    finally:
        document.close()


def clean_formula_text(text: str) -> str:
    lines = page_lines(text)
    if not lines:
        return ""

    text = "\n".join(lines)
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def normalize_pdf_line(line: str) -> str:
    line = line.replace("\x00", " ")
    line = line.replace("\ufb00", "ff")
    line = line.replace("\ufb01", "fi")
    line = line.replace("\ufb02", "fl")
    line = line.replace("\ufb03", "ffi")
    line = line.replace("\ufb04", "ffl")
    line = re.sub(r"[ \t]+", " ", line)
    return line.strip()


def is_noise_line(line: str) -> bool:
    if not line:
        return True
    if re.fullmatch(r"\d+", line):
        return True
    if re.fullmatch(r"page\s+\d+(\s+of\s+\d+)?", line, flags=re.IGNORECASE):
        return True
    if re.fullmatch(r"[-_=\s]{3,}", line):
        return True
    return False


def is_formula_like(line: str) -> bool:
    stripped = line.strip()
    if not stripped:
        return False

    strong_math_count = sum(1 for char in stripped if char in PDF_STRONG_MATH_SYMBOLS)
    weak_math_count = sum(1 for char in stripped if char in PDF_WEAK_MATH_SYMBOLS)
    alpha_count = sum(1 for char in stripped if char.isalpha())
    digit_count = sum(1 for char in stripped if char.isdigit())
    compact = stripped.replace(" ", "")

    if "={" in compact or "^{" in compact or "_{" in compact:
        return True
    if compact in {"(", ")", "[", "]", "{", "}"}:
        return True
    if len(compact) <= 40 and any(char in compact for char in PDF_MATH_SYMBOLS):
        return True
    if strong_math_count >= 2 and len(stripped) <= 180:
        return True
    if strong_math_count >= 1 and weak_math_count >= 1 and len(stripped) <= 180:
        return True
    if "=" in stripped and (alpha_count + digit_count) >= 2 and len(stripped) <= 220:
        return True
    if re.search(r"\b(d|D|exp|ln|sqrt|max|min|var|cov)\s*[\(\[]", stripped):
        return True
    if alpha_count <= 4 and (strong_math_count + weak_math_count) >= 1 and digit_count >= 1:
        return True

    return False


def normalized_line_key(line: str) -> str:
    return re.sub(r"\d+", "#", line.lower()).strip()


def page_lines(text: str) -> List[str]:
    lines = []
    for line in text.replace("\r\n", "\n").replace("\r", "\n").split("\n"):
        normalized = normalize_pdf_line(line)
        if not is_noise_line(normalized):
            lines.append(normalized)
    return lines


def find_repeated_boundary_lines(raw_pages: List[str]) -> set[str]:
    counter: Counter[str] = Counter()

    for raw_text in raw_pages:
        lines = page_lines(raw_text)
        boundary_lines = lines[:PDF_BOUNDARY_LINE_COUNT] + lines[-PDF_BOUNDARY_LINE_COUNT:]
        counter.update(
            normalized_line_key(line)
            for line in boundary_lines
            if 3 <= len(line) <= 140
        )

    min_count = min(
        PDF_REPEATED_LINE_MIN_PAGES,
        max(2, len(raw_pages) // 3),
    )
    return {line for line, count in counter.items() if count >= min_count}


def clean_pdf_text(text: str, repeated_boundary_lines: set[str]) -> str:
    lines = page_lines(text)
    cleaned_lines = []

    for index, line in enumerate(lines):
        is_boundary = (
            index < PDF_BOUNDARY_LINE_COUNT
            or index >= len(lines) - PDF_BOUNDARY_LINE_COUNT
        )
        if is_boundary and normalized_line_key(line) in repeated_boundary_lines:
            continue
        cleaned_lines.append(line)

    merged_lines = []
    for line in cleaned_lines:
        if merged_lines and merged_lines[-1].endswith("-") and line[:1].islower():
            merged_lines[-1] = merged_lines[-1][:-1] + line
        else:
            merged_lines.append(line)

    text = "\n".join(merged_lines)
    text = preserve_math_line_breaks(text)
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def preserve_math_line_breaks(text: str) -> str:
    lines = text.split("\n")
    if not lines:
        return ""

    output = [lines[0]]
    in_formula_block = is_formula_like(lines[0])
    for line in lines[1:]:
        previous = output[-1]
        line_is_formula = is_formula_like(line)
        previous_is_formula = is_formula_like(previous)

        if previous_is_formula or line_is_formula or in_formula_block:
            output.append(line)
            in_formula_block = line_is_formula or (
                in_formula_block
                and not line.endswith((".", ";", ":", "?", "!"))
            )
        elif previous.endswith((".", ":", ";", "?", "!", ")")):
            output.append(line)
            in_formula_block = False
        else:
            output[-1] = f"{previous} {line}"
            in_formula_block = False

    return "\n".join(output)


def is_chapter_heading(line: str) -> bool:
    return bool(re.fullmatch(
        r"(chapter|appendix)\s+([0-9]+|[ivxlcdm]+|[a-z])",
        line.strip(),
        flags=re.IGNORECASE,
    ))


def titlecase_word_ratio(words: List[str]) -> float:
    candidate_words = [
        word.strip("()[]{}:;,.")
        for word in words
        if any(char.isalpha() for char in word)
    ]
    if not candidate_words:
        return 0.0

    titlecase_words = [
        word
        for word in candidate_words
        if word[:1].isupper()
        or word.lower() in {"a", "an", "and", "for", "in", "of", "on", "or", "the", "to", "with"}
    ]
    return len(titlecase_words) / len(candidate_words)


def uppercase_letter_ratio(text: str) -> float:
    letters = [char for char in text if char.isalpha()]
    if not letters:
        return 0.0
    return sum(1 for char in letters if char.isupper()) / len(letters)


def is_section_heading(line: str) -> bool:
    stripped = line.strip()
    if not 4 <= len(stripped) <= 150:
        return False
    letters = [char for char in stripped if char.isalpha()]
    digit_count = sum(1 for char in stripped if char.isdigit())
    alpha_words = [
        word.strip("()[]{}:;,.")
        for word in stripped.split()
        if any(char.isalpha() for char in word)
    ]
    if len(letters) < 6 or len(alpha_words) < 2:
        return False
    if digit_count > max(4, len(letters)):
        return False
    if "%" in stripped and digit_count >= len(letters) / 2:
        return False
    numbered_heading = bool(re.match(r"^\d+(\.\d+)+\s+", stripped))
    if stripped[:1].isdigit() and not numbered_heading:
        return False
    if re.match(
        r"^(in|from|where|thus|then|now|let|because|while|figure|table|for)\b",
        stripped,
        flags=re.IGNORECASE,
    ):
        return False
    if is_formula_like(stripped):
        return False
    if stripped.endswith((".", ",", ";")):
        return False
    if re.match(r"^(figure|table)\s+\d", stripped, flags=re.IGNORECASE):
        return False
    if numbered_heading:
        return True

    words = stripped.split()
    if len(words) > 16:
        return False
    if uppercase_letter_ratio(stripped) >= 0.72 and len(words) >= 2:
        return True
    if len(words) >= 4 and titlecase_word_ratio(words) >= 0.68:
        return True

    return False


def make_section_path(chapter_title: str, section_title: str) -> str:
    if chapter_title and section_title and section_title != chapter_title:
        return f"{chapter_title} > {section_title}"
    return section_title or chapter_title


def split_pdf_page_into_sections(
    path: Path,
    page_index: int,
    text: str,
    file_hash: str,
    section_state: dict,
    extraction_backend: str,
    formula_count: int,
) -> List[Document]:
    documents = []
    lines = text.splitlines()
    pending_lines: List[str] = []
    pending_metadata = {
        "chapter_title": section_state.get("chapter_title", ""),
        "section_title": section_state.get("section_title", ""),
    }

    def flush_pending() -> None:
        nonlocal pending_lines, pending_metadata
        section_text = "\n".join(line for line in pending_lines if line.strip()).strip()
        if not section_text:
            pending_lines = []
            return

        chapter_title = pending_metadata.get("chapter_title", "")
        section_title = pending_metadata.get("section_title", "")
        documents.append(
            Document(
                text=section_text,
                metadata={
                    "source_file": str(path.resolve()),
                    "file_name": path.name,
                    "file_type": "pdf",
                    "document_title": path.stem,
                    "file_hash": file_hash,
                    "page_number": page_index,
                    "extraction_method": PDF_EXTRACTION_METHOD,
                    "extraction_backend": extraction_backend,
                    "char_count": len(section_text),
                    "formula_count": formula_count,
                    "content_type": "text",
                    "chapter_title": chapter_title,
                    "section_title": section_title,
                    "section_path": make_section_path(chapter_title, section_title),
                },
            )
        )
        pending_lines = []

    for line in lines:
        stripped = line.strip()
        if not stripped:
            continue

        if is_chapter_heading(stripped):
            if len("\n".join(pending_lines)) >= PDF_MIN_SECTION_CHARS:
                flush_pending()
            section_state["pending_chapter_label"] = stripped.title()
            section_state["chapter_title"] = stripped.title()
            section_state["section_title"] = stripped.title()
            pending_metadata = {
                "chapter_title": section_state["chapter_title"],
                "section_title": section_state["section_title"],
            }
            pending_lines.append(stripped)
            continue

        if section_state.get("pending_chapter_label") and is_section_heading(stripped):
            if pending_lines == [section_state["pending_chapter_label"]]:
                pending_lines[0] = f"{section_state['pending_chapter_label']}: {stripped}"
            else:
                pending_lines.append(stripped)

            section_state["chapter_title"] = pending_lines[-1]
            section_state["section_title"] = pending_lines[-1]
            section_state["pending_chapter_label"] = ""
            pending_metadata = {
                "chapter_title": section_state["chapter_title"],
                "section_title": section_state["section_title"],
            }
            continue

        if is_section_heading(stripped):
            if len("\n".join(pending_lines)) >= PDF_MIN_SECTION_CHARS:
                flush_pending()
            section_state["section_title"] = stripped
            section_state["pending_chapter_label"] = ""
            pending_metadata = {
                "chapter_title": section_state.get("chapter_title", ""),
                "section_title": section_state["section_title"],
            }

        pending_lines.append(stripped)

    flush_pending()
    return documents


def make_formula_documents(
    path: Path,
    page_index: int,
    formula_blocks: List[dict],
    file_hash: str,
    extraction_backend: str,
) -> List[Document]:
    documents = []

    for formula_index, formula in enumerate(formula_blocks):
        formula_text = formula.get("text", "").strip()
        if not formula_text:
            continue

        documents.append(
            Document(
                text=f"[FORMULA]\n{formula_text}\n[/FORMULA]",
                metadata={
                    "source_file": str(path.resolve()),
                    "file_name": path.name,
                    "file_type": "pdf",
                    "document_title": path.stem,
                    "file_hash": file_hash,
                    "page_number": page_index,
                    "extraction_method": PDF_EXTRACTION_METHOD,
                    "extraction_backend": extraction_backend,
                    "char_count": len(formula_text),
                    "content_type": "formula",
                    "formula_id": formula.get("id", f"formula-{page_index}-{formula_index}"),
                    "formula_index": formula_index,
                    "formula_bbox": formula.get("bbox", ""),
                    "formula_count": 1,
                    "chapter_title": "",
                    "section_title": "",
                    "section_path": "",
                },
            )
        )

    return documents


def load_pdf_file(path: Path) -> List[Document]:
    reader = PdfReader(str(path))
    documents = []
    pymupdf_pages = extract_pdf_pages_with_pymupdf(path)

    if pymupdf_pages:
        page_payloads = pymupdf_pages
    else:
        page_payloads = [
            {
                "text": extract_pdf_text(page),
                "formula_blocks": [],
                "backend": "pypdf",
            }
            for page in reader.pages
        ]

    raw_pages = [payload["text"] for payload in page_payloads]
    repeated_boundary_lines = find_repeated_boundary_lines(raw_pages)
    file_hash = file_sha256(path)
    section_state: dict = {
        "chapter_title": "",
        "section_title": "",
        "pending_chapter_label": "",
    }

    for page_index, payload in enumerate(page_payloads, start=1):
        raw_text = payload["text"]
        text = clean_pdf_text(raw_text, repeated_boundary_lines)
        formula_blocks = payload.get("formula_blocks", [])
        extraction_backend = payload.get("backend", "pypdf")

        if not text.strip():
            documents.extend(
                make_formula_documents(
                    path=path,
                    page_index=page_index,
                    formula_blocks=formula_blocks,
                    file_hash=file_hash,
                    extraction_backend=extraction_backend,
                )
            )
            continue

        documents.extend(
            split_pdf_page_into_sections(
                path=path,
                page_index=page_index,
                text=text,
                file_hash=file_hash,
                section_state=section_state,
                extraction_backend=extraction_backend,
                formula_count=len(formula_blocks),
            )
        )
        documents.extend(
            make_formula_documents(
                path=path,
                page_index=page_index,
                formula_blocks=formula_blocks,
                file_hash=file_hash,
                extraction_backend=extraction_backend,
            )
        )

    return documents


def load_txt_file(path: Path) -> List[Document]:
    text = path.read_text(encoding="utf-8")
    return [
        Document(
            text=text,
            metadata={
                "source_file": str(path.resolve()),
                "file_name": path.name,
                "file_type": "txt",
                "document_title": path.stem,
                "file_hash": file_sha256(path),
                "content_type": "text",
                "chapter_title": "",
                "section_title": "",
                "section_path": "",
                "char_count": len(text),
            },
        )
    ]


def iter_source_files(raw_dir: Path) -> Iterable[Path]:
    supported_suffixes = {".md", ".markdown", ".pdf", ".txt"}
    for path in sorted(raw_dir.rglob("*")):
        if path.is_file() and path.suffix.lower() in supported_suffixes:
            yield path


def load_docs(raw_dir: Path = RAW_DIR) -> List[Document]:
    documents: List[Document] = []
    raw_dir = effective_raw_dir(raw_dir)

    for path in iter_source_files(raw_dir):
        suffix = path.suffix.lower()

        if suffix in {".md", ".markdown"}:
            documents.extend(load_md_documents(path))
        elif suffix == ".pdf":
            documents.extend(load_pdf_file(path))
        elif suffix == ".txt":
            documents.extend(load_txt_file(path))

    if not documents:
        raise ValueError(f"No supported documents found under {raw_dir}")

    return documents


def add_chunk_metadata(nodes: List[BaseNode]) -> List[BaseNode]:
    counters: dict[str, int] = {}

    for node in nodes:
        source_file = node.metadata["source_file"]
        chunk_index = counters.get(source_file, 0)
        counters[source_file] = chunk_index + 1

        file_hash = node.metadata["file_hash"][:12]
        page_number = node.metadata.get("page_number", "na")
        chunk_id = f"{Path(source_file).stem}-{file_hash}-p{page_number}-c{chunk_index}"

        node.metadata["chunk_id"] = chunk_id
        node.metadata["chunk_index"] = chunk_index
        node.metadata[EMBED_MODEL_METADATA_KEY] = EMBED_MODEL_NAME
        node.id_ = chunk_id

    return nodes


def validate_nodes(nodes: List[BaseNode]) -> None:
    if not nodes:
        raise ValueError("No chunks were created from the source documents.")

    for node in nodes:
        missing = [key for key in REQUIRED_METADATA if key not in node.metadata]
        if missing:
            raise ValueError(
                f"Node {node.node_id} is missing metadata fields: {missing}")

        if node.metadata["file_type"] == "pdf" and "page_number" not in node.metadata:
            raise ValueError(
                f"PDF node {node.node_id} is missing page_number metadata.")


def split_documents(documents: List[Document]) -> List[BaseNode]:
    splitter = SentenceSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
    )
    nodes = splitter.get_nodes_from_documents(documents)
    add_chunk_metadata(nodes)
    validate_nodes(nodes)
    return nodes


def build_nodes(raw_dir: Path = RAW_DIR) -> List[BaseNode]:
    documents = load_docs(raw_dir)
    return split_documents(documents)


def load_source_file(path: Path) -> List[Document]:
    suffix = path.suffix.lower()
    if suffix in {".md", ".markdown"}:
        return load_md_documents(path)
    if suffix == ".pdf":
        return load_pdf_file(path)
    if suffix == ".txt":
        return load_txt_file(path)
    return []


def list_current_sources(raw_dir: Path = RAW_DIR) -> dict[str, dict[str, str]]:
    raw_dir = effective_raw_dir(raw_dir)
    sources = {}
    for path in iter_source_files(raw_dir):
        resolved = str(path.resolve())
        sources[resolved] = {
            "file_hash": file_sha256(path),
            "file_type": path.suffix.lower().lstrip("."),
        }
    return sources


def existing_source_metadata(chroma_collection) -> dict[str, dict[str, str]]:
    existing: dict[str, dict[str, str]] = {}
    if chroma_collection.count() == 0:
        return existing

    offset = 0
    limit = 500
    while True:
        batch = chroma_collection.get(
            limit=limit,
            offset=offset,
            include=["metadatas"],
        )
        metadatas = batch.get("metadatas") or []
        if not metadatas:
            break
        for metadata in metadatas:
            source_file = metadata.get("source_file")
            if not source_file:
                continue
            existing[source_file] = {
                "file_hash": metadata.get("file_hash", ""),
                "file_type": metadata.get("file_type", ""),
                "embedding_model": metadata.get(EMBED_MODEL_METADATA_KEY, ""),
                "extraction_method": metadata.get("extraction_method", ""),
            }
        if len(metadatas) < limit:
            break
        offset += limit
    return existing


def source_needs_update(current: dict[str, str], existing: dict[str, str] | None) -> bool:
    if not existing:
        return True
    if existing.get("file_hash") != current["file_hash"]:
        return True
    if existing.get("embedding_model") != EMBED_MODEL_NAME:
        return True
    if current["file_type"] == "pdf" and existing.get("extraction_method") != PDF_EXTRACTION_METHOD:
        return True
    return False


def incremental_update_index(
    raw_dir: Path,
    chroma_collection,
    storage_context: StorageContext,
    embed_model,
) -> bool:
    current_sources = list_current_sources(raw_dir)
    existing_sources = existing_source_metadata(chroma_collection)

    deleted_sources = sorted(set(existing_sources) - set(current_sources))
    changed_sources = sorted(
        source_file
        for source_file, current in current_sources.items()
        if source_needs_update(current, existing_sources.get(source_file))
    )

    for source_file in deleted_sources + changed_sources:
        try:
            chroma_collection.delete(where={"source_file": source_file})
        except Exception as exc:
            logging.warning("Could not delete stale chunks for %s: %s", source_file, exc)

    if not changed_sources:
        if deleted_sources:
            print(f"Removed {len(deleted_sources)} stale source(s) from collection '{COLLECTION_NAME}'.")
        return bool(deleted_sources)

    documents: List[Document] = []
    for source_file in changed_sources:
        documents.extend(load_source_file(Path(source_file)))

    nodes = split_documents(documents)
    VectorStoreIndex(
        nodes,
        storage_context=storage_context,
        embed_model=embed_model,
        show_progress=True,
    )
    print(
        f"Incrementally indexed {len(nodes)} chunk(s) from {len(changed_sources)} source file(s)."
    )
    return True


def collection_needs_rebuild(chroma_collection) -> bool:
    if chroma_collection.count() == 0:
        return True

    try:
        sample = chroma_collection.peek(limit=min(chroma_collection.count(), 20))
    except Exception:
        return False

    for metadata in sample.get("metadatas") or []:
        if metadata.get(EMBED_MODEL_METADATA_KEY) != EMBED_MODEL_NAME:
            return True
        if metadata.get("file_type") == "pdf":
            return metadata.get("extraction_method") != PDF_EXTRACTION_METHOD

    return False


async def build_index(raw_dir: Path = RAW_DIR, rebuild: bool = False) -> VectorStoreIndex:
    configure_model_cache()

    from llama_index.embeddings.huggingface import HuggingFaceEmbedding

    load_dotenv()
    raw_dir = effective_raw_dir(raw_dir)
    CHROMA_DB_DIR.mkdir(parents=True, exist_ok=True)

    db = chromadb.PersistentClient(path=str(CHROMA_DB_DIR))

    if rebuild:
        try:
            db.delete_collection(COLLECTION_NAME)
        except (NotFoundError, ValueError):
            pass

    chroma_collection = db.get_or_create_collection(COLLECTION_NAME)

    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    embed_model = HuggingFaceEmbedding(
        model_name=resolve_embed_model_name(),
        cache_folder=str(HF_CACHE_DIR / "sentence_transformers"),
    )

    if rebuild or chroma_collection.count() == 0:
        nodes = build_nodes(raw_dir)
        index = VectorStoreIndex(
            nodes,
            storage_context=storage_context,
            embed_model=embed_model,
            show_progress=True,
        )
        print(
            f"Indexed {len(nodes)} chunks into collection '{COLLECTION_NAME}'")
        return index

    incremental_update_index(
        raw_dir=raw_dir,
        chroma_collection=chroma_collection,
        storage_context=storage_context,
        embed_model=embed_model,
    )

    print(
        f"Loaded existing collection '{COLLECTION_NAME}' with {chroma_collection.count()} chunks.")
    return VectorStoreIndex.from_vector_store(vector_store, embed_model=embed_model)


class QueryKnowledgeTool(Tool):
    name = "query_knowledge"
    description = (
        "Searches the local options trading knowledge base. Use this for option "
        "concepts, volatility, Greeks, strategies, formulas, equation numbers, "
        "and citations from reference books."
    )
    inputs = {'query': {'type': 'string',
                        'description': 'The search query to perform.'}}
    output_type = "string"

    @staticmethod
    def format_results(results, max_chars: int = 800):
        output = []

        for result in results:
            metadata = result.node.metadata
            source = metadata.get("file_name", "unknown")
            page = metadata.get("page_number", "n/a")
            section = metadata.get("section_path") or metadata.get("section_title") or "n/a"
            content_type = metadata.get("content_type", "text")
            formula_id = metadata.get("formula_id", "")
            score = result.score
            text = result.node.get_content()
            if len(text) > max_chars:
                text = f"{text[:max_chars].rstrip()}..."

            output.append(
                f"source：{source}\n"
                f"page：{page}\n"
                f"section：{section}\n"
                f"content_type：{content_type}\n"
                f"formula_id：{formula_id or 'n/a'}\n"
                f"score：{score:.4f}\n"
                f"vector_score：{metadata.get(VECTOR_METADATA_KEY, 'n/a')}\n"
                f"bm25_score：{metadata.get(BM25_METADATA_KEY, 'n/a')}\n"
                f"content：{text}"
            )

        return "\n\n---\n\n".join(output)

    @staticmethod
    def load_bm25_nodes(collection_name: str = COLLECTION_NAME) -> list[TextNode]:
        db = chromadb.PersistentClient(path=str(CHROMA_DB_DIR))
        try:
            collection = db.get_collection(collection_name)
        except Exception:
            return []

        nodes: list[TextNode] = []
        offset = 0
        limit = 500
        while True:
            batch = collection.get(
                limit=limit,
                offset=offset,
                include=["documents", "metadatas"],
            )
            documents = batch.get("documents") or []
            metadatas = batch.get("metadatas") or []
            ids = batch.get("ids") or []
            if not documents:
                break

            for index, text in enumerate(documents):
                metadata = dict(metadatas[index] or {})
                node_id = ids[index] if index < len(ids) else metadata.get("chunk_id", "")
                nodes.append(TextNode(id_=node_id, text=text or "", metadata=metadata))

            if len(documents) < limit:
                break
            offset += limit

        return nodes

    @staticmethod
    def merge_results(
        vector_results: list[NodeWithScore],
        bm25_results: list[NodeWithScore],
        top_k: int,
    ) -> list[NodeWithScore]:
        merged: dict[str, NodeWithScore] = {}

        for rank, result in enumerate(vector_results):
            node_id = result.node.node_id
            result.node.metadata[VECTOR_METADATA_KEY] = result.score
            merged[node_id] = NodeWithScore(
                node=result.node,
                score=1.0 / (rank + 1),
            )

        for rank, result in enumerate(bm25_results):
            node_id = result.node.node_id
            result.node.metadata[BM25_METADATA_KEY] = result.score
            reciprocal_rank_score = 1.0 / (rank + 1)
            if node_id in merged:
                merged[node_id].score = (merged[node_id].score or 0.0) + reciprocal_rank_score
                merged[node_id].node.metadata[BM25_METADATA_KEY] = result.score
            else:
                merged[node_id] = NodeWithScore(
                    node=result.node,
                    score=reciprocal_rank_score,
                )

        results = list(merged.values())
        results.sort(key=lambda item: item.score or float("-inf"), reverse=True)
        return results[:top_k]

    def __init__(
        self,
        max_results=20,
        top_k=5,
        use_reranker: Optional[bool] = None,
        use_hybrid: Optional[bool] = None,
        reranker_top_n: Optional[int] = None,
        reranker_model_name: Optional[str] = None,
        **kwargs,
    ):
        super().__init__()
        self.max_results = max_results
        self.top_k = top_k
        self.use_reranker = (
            env_flag("RAG_USE_RERANKER", True)
            if use_reranker is None
            else use_reranker
        )
        self.use_hybrid = (
            env_flag("RAG_USE_HYBRID", True)
            if use_hybrid is None
            else use_hybrid
        )
        self.reranker_top_n = reranker_top_n or top_k
        self.reranker = (
            CrossEncoderReranker(reranker_model_name or RERANKER_MODEL_NAME)
            if self.use_reranker
            else None
        )
        index = asyncio.run(build_index(rebuild=False))
        retrieve_top_k = max(max_results, top_k) if self.use_reranker else top_k
        self.retriever = index.as_retriever(similarity_top_k=retrieve_top_k)
        self.bm25_retriever = (
            BM25Retriever(self.load_bm25_nodes())
            if self.use_hybrid
            else None
        )

    def forward(self, query: str) -> str:
        vector_results = self.retriever.retrieve(query)
        results = vector_results
        if self.bm25_retriever:
            bm25_results = self.bm25_retriever.retrieve(query, self.max_results)
            results = self.merge_results(
                vector_results=vector_results,
                bm25_results=bm25_results,
                top_k=max(self.max_results, self.top_k),
            )
        if self.reranker:
            try:
                results = self.reranker.rerank(
                    query,
                    results,
                    top_n=self.reranker_top_n,
                )
            except Exception as exc:
                logging.warning("Reranker failed; falling back to vector ranking: %s", exc)
                results = results[:self.top_k]
        return QueryKnowledgeTool.format_results(results[:self.top_k])


if __name__ == "__main__":
    query_tool = QueryKnowledgeTool()
    res: str = query_tool.forward("What is option?")
    print(res)