Search Your Docs Folder with pplx-embed-v2, Dense and Contextual

httpx==0.28.1
numpy==2.5.4
pytest==9.1.1

{"question": "How do I run my tests?", "file": "tutorial/testing.md", "contains": "$ pytest"}
{"question": "How do I make a path parameter accept values with slashes in them?", "file": "tutorial/path-params.md", "contains": ":path"}
{"question": "How do I run a function after returning a response, like sending an email?", "file": "tutorial/background-tasks.md", "contains": ".add_task("}
{"question": "How do I run the development server with auto reload?", "file": "fastapi-cli.md", "contains": "fastapi dev"}
{"question": "What status code is returned when request validation fails?", "file": "tutorial/handling-errors.md", "contains": "422"}
{"question": "How do I stop underscores in header names from being converted to hyphens?", "file": "tutorial/header-params.md", "contains": "convert_underscores"}
{"question": "How do I leave default values out of the response body?", "file": "tutorial/response-model.md", "contains": ["response_model_exclude_unset", "response_model_exclude_defaults"]}
{"question": "How do I apply one dependency to every route in a router?", "file": "tutorial/bigger-applications.md", "contains": ["we can add it to the `APIRouter`", "dependencies=[Depends(get_token_header)]", "all the *path operations* in the router", "`dependencies` in the `APIRouter`"]}
{"question": "How do I mark an endpoint as deprecated in the docs?", "file": "tutorial/path-operation-configuration.md", "contains": "deprecated"}
{"question": "How do I turn off the automatic documentation pages?", "file": "tutorial/metadata.md", "contains": "docs_url"}
{"question": "How do I serve a folder of images at a URL path?", "file": "tutorial/static-files.md", "contains": "StaticFiles(directory"}
{"question": "How do I allow requests from a frontend running on another origin?", "file": "tutorial/cors.md", "contains": "allow_origins"}
{"question": "How do I log in with a username and password from the docs UI?", "file": "tutorial/security/simple-oauth2.md", "contains": "\"Authorize\" button"}
{"question": "How do I run FastAPI behind a proxy that strips a path prefix?", "file": "advanced/behind-a-proxy.md", "contains": "--root-path"}
{"question": "Where do I put code that runs once when the server starts?", "file": "advanced/events.md", "contains": ["`lifespan` parameter", "@app.on_event(\"startup\")"]}
{"question": "How do I convert a datetime to something JSON can store?", "file": "tutorial/encoder.md", "contains": "jsonable_encoder"}
{"question": "How do I see the request body that failed validation?", "file": "tutorial/handling-errors.md", "contains": "exc.body"}
{"question": "How do I upload several files at once?", "file": "tutorial/request-files.md", "contains": "[bytes]"}
{"question": "How do I read settings from environment variables?", "file": "advanced/settings.md", "contains": "BaseSettings"}
{"question": "How do I replace a dependency with a fake one in tests?", "file": "advanced/testing-dependencies.md", "contains": ["app.dependency_overrides[", "as a key the original dependency"]}
{"question": "How do I set a cookie on the response?", "file": "advanced/response-cookies.md", "contains": "set cookies"}
{"question": "How do I stream a large file back to the client?", "file": "advanced/custom-response.md", "contains": "StreamingResponse"}
{"question": "How do I show example values for a request body in the docs?", "file": "tutorial/schema-extra-example.md", "contains": "examples=["}
{"question": "How do I accept the same query parameter several times?", "file": "tutorial/query-params-str-validations.md", "contains": "list[str]"}
{"question": "How do I return 201 when an item is created?", "file": "tutorial/response-status-code.md", "contains": "status_code=201"}
{"question": "How do I group query parameters into a Pydantic model?", "file": "tutorial/query-param-models.md", "contains": ["FilterParams, Query()", "FilterParams = Query()"]}
{"question": "How do I get the current user from a token in every route?", "file": "tutorial/security/get-current-user.md", "contains": "Depends(get_current_user)"}
{"question": "How do I return a file from disk as the response?", "file": "advanced/custom-response.md", "contains": "FileResponse"}
{"question": "How do I validate that a number path parameter is at least 1?", "file": "tutorial/path-params-numeric-validations.md", "contains": "ge=1"}
{"question": "How do I open a websocket connection?", "file": "advanced/websockets.md", "contains": "@app.websocket"}

"""Search a folder of Markdown files with pplx-embed-v2: the same model in dense and contextual mode, side by side.

Usage:
  python docsearch.py index <folder>              chunk every .md file, embed it three ways, write index/
  python docsearch.py ask "<question>" [-k 3] [--configs NAME ...]   top chunks per configuration
  python docsearch.py eval questions.jsonl         recall@k and MRR@100 for each configuration on labeled questions
"""

import argparse
import base64
import json
import os
import re
import sys
import time
from dataclasses import asdict, dataclass
from pathlib import Path

import httpx
import numpy as np

API_URL = "https://api.perplexity.ai/v1/embeddings"
# One checkpoint per size serves both input modes. The index is always built with the 9b model; the 0.6b model shares
# its vector space at 1024 dimensions, so it can embed questions against a 9b index built at 1024.
INDEX_MODEL = "pplx-embed-v2-context-9b"
SMALL_MODEL = "pplx-embed-v2-context-0.6b"


@dataclass(frozen=True)
class Config:
    """One column of the evaluation: how the corpus was embedded and which model embeds the questions."""

    name: str
    mode: str  # "dense": every chunk embedded alone. "context": each document's chunks embedded together.
    dimensions: int
    query_model: str

    @property
    def index_key(self) -> str:
        return f"{self.mode}-{self.dimensions}"


CONFIGS = [
    Config("dense-2048", "dense", 2048, INDEX_MODEL),
    Config("context-2048", "context", 2048, INDEX_MODEL),
    Config("context-1024", "context", 1024, INDEX_MODEL),
    Config("context-1024-small", "context", 1024, SMALL_MODEL),
]

# A chunk is one heading section. Longer sections are split at blank lines into pieces of about this size, and a single
# line longer than this (an unwrapped paragraph, a minified file, an inline image) is cut into pieces of this size.
CHUNK_CHARS = 1500
# The API's caps are in tokens: 16,384 per text or per contextual document for the 9b model (32,768 for the 0.6b),
# counting one extra token per chunk in contextual mode, and 120,000 per request. The two character budgets below are
# heuristics tuned for this corpus, which runs about 3.6 characters per token overall and close to 3 in its link-heavy
# release notes. A corpus near 2 characters per token (CJK text, minified code, long link lists) can exceed either cap
# with these values; lower both, or count tokens with the published tokenizer.
# A document is one file, unless the file is longer than this or has more than REQUEST_CHUNKS chunks, in which case its
# sections are split into consecutive documents.
DOC_CHARS = 40_000
# Per request: 512 texts or documents, 16,000 chunks in a contextual request (the documented limit), and this many
# characters of text as the budget for the 120,000-token cap.
REQUEST_TEXTS = 512
REQUEST_CHUNKS = 16_000
REQUEST_CHARS = 240_000
INDEX_DIR = Path("index")
# How deep eval looks for the answer. A question whose answer ranks below this counts as not found.
EVAL_DEPTH = 100


# --- chunk
@dataclass
class Chunk:
    doc: str  # document name: the file path, with #2, #3 for the later parts of a split file
    heading: str
    text: str


def track_fence(line: str, fence: str) -> str:
    """The code fence that is open after this line, or "" outside code. A fence opens with three or more backticks or
    tildes and closes only on a line of at least as many of the same character, so a shorter fence nested inside a
    longer one (a three-backtick Python block inside a four-backtick Markdown block) does not close the outer block."""
    match = re.match(r" {0,3}(`{3,}|~{3,})(.*)", line)
    if not match:
        return fence
    marker, rest = match.groups()
    if not fence:
        return marker
    if marker[0] == fence[0] and len(marker) >= len(fence) and not rest.strip():
        return ""
    return fence


def split_long(section: str) -> list[str]:
    """Split one section at blank lines outside code fences into pieces of about CHUNK_CHARS. A paragraph or code block
    longer than that on its own is split at line breaks, and a single line longer than that is cut into CHUNK_CHARS pieces."""
    paragraphs, current, fence = [], "", ""
    for line in section.splitlines(keepends=True):
        fence = track_fence(line, fence)
        if not line.strip() and not fence and current:
            paragraphs.append(current.strip("\n"))
            current = ""
        else:
            current += line
    if current.strip():
        paragraphs.append(current.strip("\n"))
    pieces = []
    for paragraph in paragraphs:
        if len(paragraph) <= CHUNK_CHARS:
            pieces.append(paragraph)
        else:
            lines = [line[i : i + CHUNK_CHARS] for line in paragraph.splitlines() for i in range(0, len(line) or 1, CHUNK_CHARS)]
            pieces.extend(_group(lines, "\n"))
    return _group(pieces, "\n\n")


def _group(parts: list[str], sep: str) -> list[str]:
    pieces, current = [], ""
    for part in parts:
        if current and len(current) + len(part) > CHUNK_CHARS:
            pieces.append(current.strip())
            current = ""
        current += part + sep
    if current.strip():
        pieces.append(current.strip())
    return pieces


def chunk_markdown(text: str) -> list[tuple[str, str]]:
    """Return (heading, chunk text) pairs. Each heading starts a chunk and stays in it; headings inside code fences do not count."""
    sections: list[str] = []
    fence = ""
    for line in text.splitlines(keepends=True):
        fence = track_fence(line, fence)
        if sections and (fence or not re.match(r"#{1,6} ", line)):
            sections[-1] += line
        else:
            sections.append(line)
    chunks = []
    heading = ""
    for section in sections:
        if not section.strip():
            continue
        if section.startswith("#"):
            heading = section.splitlines()[0].lstrip("#").strip()
        for piece in split_long(section):
            chunks.append((heading, piece))
    return chunks


def load_folder(root: str) -> list[list[Chunk]]:
    """One document per Markdown file, in path order. A file longer than DOC_CHARS or with more than REQUEST_CHUNKS
    chunks becomes several documents, so that a document stays within the character budget for the per-document cap."""
    documents = []
    for path in sorted(Path(root).rglob("*.md")):
        name = str(path.relative_to(root))
        parts: list[list[Chunk]] = [[]]
        for heading, text in chunk_markdown(path.read_text(encoding="utf-8")):
            if parts[-1] and (len(parts[-1]) == REQUEST_CHUNKS or sum(len(c.text) for c in parts[-1]) + len(text) > DOC_CHARS):
                parts.append([])
            parts[-1].append(Chunk(name, heading or path.name, text))
        for i, part in enumerate(parts):
            if part:
                if i:
                    for chunk in part:
                        chunk.doc = f"{name}#{i + 1}"
                documents.append(part)
    return documents


# --- embed
def batches(items: list, size_of, limits: tuple[int, ...]):
    """Yield consecutive slices whose summed sizes stay within limits. size_of returns one count per limit for an item.
    An item that is over a limit on its own cannot be sent, so that is an error here rather than a rejected request."""
    start, totals = 0, [0] * len(limits)
    for i, item in enumerate(items):
        sizes = size_of(item)
        if any(s > lim for s, lim in zip(sizes, limits, strict=True)):
            raise ValueError(f"item {i} alone measures {sizes}; the request limits are {limits}")
        if i > start and any(t + s > lim for t, s, lim in zip(totals, sizes, limits)):
            yield items[start:i]
            start, totals = i, [0] * len(limits)
        totals = [t + s for t, s in zip(totals, sizes)]
    if start < len(items):
        yield items[start:]


def post(client: httpx.Client, body: dict) -> dict:
    """One POST to /v1/embeddings. Retries 429 and 5xx with backoff; any other error stops with the API's message."""
    for attempt in range(5):
        response = client.post(API_URL, json=body)
        if response.status_code == 429 or response.status_code >= 500:
            time.sleep(2**attempt)
            continue
        if response.status_code != 200:
            raise RuntimeError(f"HTTP {response.status_code} from the Embeddings API: {response.text}")
        return response.json()
    raise RuntimeError(f"HTTP {response.status_code} from the Embeddings API after 5 attempts: {response.text}")


def decode(b64: str, dimensions: int) -> np.ndarray:
    """A base64_int8 embedding as a signed int8 vector, checked against the dimensions the request asked for."""
    vector = np.frombuffer(base64.b64decode(b64), dtype=np.int8)
    if len(vector) != dimensions:
        raise RuntimeError(f"the API returned a {len(vector)}-dimensional vector for a {dimensions}-dimensional request")
    return vector


def embed(client: httpx.Client, model: str, input_type: str, inputs: list, dimensions: int, usage: dict) -> np.ndarray:
    """Int8 vectors, one row per text or per chunk, in input order. A flat list of strings is dense input: each text
    is embedded alone. A list of lists is contextual input: each inner list is one document, and every chunk vector
    reflects the rest of its document. input_type is "query" for questions and "document" for everything indexed."""
    contextual = bool(inputs) and isinstance(inputs[0], list)
    if contextual:
        size_of = lambda doc: (1, len(doc), sum(len(c) for c in doc))
        limits = (REQUEST_TEXTS, REQUEST_CHUNKS, REQUEST_CHARS)
    else:
        size_of = lambda text: (1, len(text))
        limits = (REQUEST_TEXTS, REQUEST_CHARS)
    rows = []
    for batch in batches(inputs, size_of, limits):
        body = {"model": model, "input_type": input_type, "input": batch, "dimensions": dimensions, "encoding_format": "base64_int8"}
        response = post(client, body)
        record(usage, model, response)
        for item in in_order(response["data"]):
            if contextual:
                rows.extend(decode(c["embedding"], dimensions) for c in in_order(item["data"]))
            else:
                rows.append(decode(item["embedding"], dimensions))
    return np.stack(rows)


def in_order(items: list[dict]) -> list[dict]:
    """Response items sorted by index, after checking that the indices are exactly 0..n-1, so a response with a missing
    or repeated item cannot shift every row after it."""
    ordered = sorted(items, key=lambda d: d["index"])
    if [d["index"] for d in ordered] != list(range(len(ordered))):
        raise RuntimeError(f"the API returned indices {[d['index'] for d in ordered]}; expected 0 to {len(ordered) - 1}")
    return ordered

def record(usage: dict, model: str, response: dict) -> None:
    """Add up requests, tokens, and cost per model. Cost is summed only when the response reports it."""
    entry = usage.setdefault(model, {"requests": 0, "tokens": 0, "cost_usd": None})
    entry["requests"] += 1
    entry["tokens"] += response.get("usage", {}).get("total_tokens", 0)
    cost = response.get("usage", {}).get("cost", {}).get("total_cost")
    if cost is not None:
        entry["cost_usd"] = (entry["cost_usd"] or 0.0) + cost


# --- index
@dataclass
class Index:
    """One int8 matrix (a row per chunk) and the float32 norm of each row, so cosine similarity is an integer dot
    product divided by two norms. Storing int8 keeps 2,048 dimensions in 2 KB per chunk and 1,024 in 1 KB."""

    matrix: np.ndarray
    norms: np.ndarray

    @classmethod
    def of(cls, matrix: np.ndarray) -> "Index":
        return cls(matrix, np.linalg.norm(matrix.astype(np.float32), axis=1))


def build_index(client: httpx.Client, root: str, out: Path = INDEX_DIR) -> dict:
    """Embed every chunk under root once per distinct (mode, dimensions) in CONFIGS, with INDEX_MODEL, and write the
    int8 matrices, the chunk metadata, and a description of each matrix to out."""
    documents = load_folder(root)
    chunks = [chunk for doc in documents for chunk in doc]
    usage: dict = {}
    started = time.time()
    matrices: dict[str, np.ndarray] = {}
    for config in CONFIGS:
        if config.index_key in matrices:
            continue
        if config.mode == "dense":
            inputs: list = [c.text for c in chunks]
        else:
            inputs = [[c.text for c in doc] for doc in documents]
        matrix = embed(client, INDEX_MODEL, "document", inputs, config.dimensions, usage)
        if len(matrix) != len(chunks):
            raise RuntimeError(f"{config.index_key}: embedding count {len(matrix)} does not match chunk count {len(chunks)}")
        matrices[config.index_key] = matrix
    out.mkdir(exist_ok=True)
    np.savez(out / "vectors.npz", **matrices)
    (out / "chunks.json").write_text(json.dumps([asdict(c) for c in chunks]))
    described = {c.index_key: {"mode": c.mode, "dimensions": c.dimensions} for c in CONFIGS}
    (out / "models.json").write_text(json.dumps({"model": INDEX_MODEL, "indexes": described}))
    return {"documents": len(documents), "chunks": len(chunks), "seconds": time.time() - started, "usage": usage}


def load_index(path: Path = INDEX_DIR) -> tuple[list[Chunk], dict[str, Index]]:
    """Read the index back. Refuses an index built with another model, one that lacks a matrix a config needs, or one
    whose saved vectors do not have one row per chunk and the dimensions recorded when it was built."""
    models = json.loads((path / "models.json").read_text())
    if models["model"] != INDEX_MODEL:
        raise RuntimeError(f"index was built with {models['model']}; run index again")
    missing = sorted({c.index_key for c in CONFIGS} - set(models["indexes"]))
    if missing:
        raise RuntimeError(f"index has no vectors for {', '.join(missing)}; run index again")
    chunks = [Chunk(**c) for c in json.loads((path / "chunks.json").read_text())]
    vectors = np.load(path / "vectors.npz")
    indexes = {}
    for key, described in models["indexes"].items():
        matrix = vectors[key]
        if matrix.dtype != np.int8 or matrix.ndim != 2 or matrix.shape[1] != described["dimensions"]:
            raise RuntimeError(f"{key} vectors are {matrix.dtype} with shape {matrix.shape}, models.json says int8 with {described['dimensions']} dimensions; run index again")
        if matrix.shape[0] != len(chunks):
            raise RuntimeError(f"{key} has {matrix.shape[0]} vectors for {len(chunks)} chunks; run index again")
        indexes[key] = Index.of(matrix)
    return chunks, indexes


def embed_queries(client: httpx.Client, questions: list[str]) -> dict[str, np.ndarray]:
    """One query matrix per config, embedded with the config's query model at the config's dimensions. Questions go
    as flat strings with input_type "query": a question has no surrounding document."""
    usage: dict = {}
    queries = {}
    for config in CONFIGS:
        queries[config.name] = embed(client, config.query_model, "query", questions, config.dimensions, usage)
    return queries


def rank(index: Index, query: np.ndarray) -> np.ndarray:
    """Row indices from most to least similar, ranked by cosine: the dot product over the query's norm times the row's
    norm, with the query's norm left out because it is the same for every row. The dot product runs in int32 (int8
    times int8 summed over 2,048 dimensions cannot overflow it) and the row norms divide it in float32."""
    if query.shape[-1] != index.matrix.shape[1]:
        raise RuntimeError(f"query vector has {query.shape[-1]} dimensions but the index has {index.matrix.shape[1]}; run index again")
    scores = (index.matrix @ query.astype(np.int32)) / index.norms
    return np.argsort(-scores, kind="stable")


# --- ask
def ask(client: httpx.Client, question: str, k: int, names: list[str]) -> None:
    chunks, indexes = load_index()
    queries = embed_queries(client, [question])
    for config in [c for c in CONFIGS if c.name in names]:
        print(f"\n{config.name}: {config.mode} index at {config.dimensions} dimensions, question embedded by {config.query_model}")
        for place, row in enumerate(rank(indexes[config.index_key], queries[config.name][0])[:k], start=1):
            chunk = chunks[row]
            preview = " ".join(chunk.text.split())[:160].rstrip()
            print(f"  {place}. {chunk.doc} > {chunk.heading}\n     {preview}")


# --- eval
def phrases(label: dict) -> list[str]:
    """The `contains` field is one phrase or a list of phrases; a chunk that has any of them is accepted."""
    value = label["contains"]
    return [value] if isinstance(value, str) else list(value)


def hit(chunk: Chunk, label: dict) -> bool:
    """A chunk answers a question when it comes from the labeled file and contains one of the labeled phrases. Only the
    #2, #3 part suffix of a split file is removed before comparing, so a file named `C# guide.md` keeps its name."""
    text = chunk.text.lower()
    return re.sub(r"#\d+$", "", chunk.doc) == label["file"] and any(p.lower() in text for p in phrases(label) if p.strip())


def check_labels(chunks: list[Chunk], labels: list[dict]) -> None:
    """Stop before spending any tokens on a label file that cannot be scored: an empty phrase would accept every chunk
    of a file, a repeated question would count twice, and a label that matches no chunk is a typo in the file or phrase,
    or a corpus that changed since the label was written."""
    if any(not phrases(label) or not all(p.strip() for p in phrases(label)) for label in labels):
        sys.exit("every label needs at least one non-empty phrase in 'contains'")
    questions = [label["question"] for label in labels]
    if len(set(questions)) != len(questions):
        sys.exit("duplicate question in the label file: " + next(q for q in questions if questions.count(q) > 1))
    unmatched = [label for label in labels if not any(hit(chunk, label) for chunk in chunks)]
    if unmatched:
        lines = [f"  {label['file']} containing {label['contains']!r}: {label['question']}" for label in unmatched]
        sys.exit(f"{len(unmatched)} label(s) match no chunk in the index. Fix the label or rebuild the index.\n" + "\n".join(lines))


def evaluate(chunks: list[Chunk], indexes: dict[str, Index], queries: dict, labels: list[dict]) -> dict[str, list[int]]:
    """For each config, the rank of the first chunk that answers each question (0 when none in the top EVAL_DEPTH)."""
    ranks: dict[str, list[int]] = {}
    for config in CONFIGS:
        ranks[config.name] = []
        for label, query in zip(labels, queries[config.name], strict=True):
            order = rank(indexes[config.index_key], query)[:EVAL_DEPTH]
            found = next((i + 1 for i, row in enumerate(order) if hit(chunks[row], label)), 0)
            ranks[config.name].append(found)
    return ranks


def summarize(ranks: list[int]) -> dict[str, float]:
    """recall@1, 3, 5 and MRR over the top EVAL_DEPTH, so the metric is MRR@100 and an answer below that counts as zero."""
    found = [r for r in ranks if r]
    result = {f"recall@{k}": sum(1 for r in found if r <= k) / len(ranks) for k in (1, 3, 5)}
    result[f"MRR@{EVAL_DEPTH}"] = sum(1 / r for r in found) / len(ranks)
    return result


# --- output
def print_eval(labels: list[dict], ranks: dict[str, list[int]]) -> None:
    names = [c.name for c in CONFIGS]
    width = max(len(n) for n in names)
    metrics = {name: summarize(ranks[name]) for name in names}
    print(f"{len(labels)} questions\n")
    print(f"{'configuration':<{width}}  " + "  ".join(f"{k:>8}" for k in metrics[names[0]]))
    for name in names:
        print(f"{name:<{width}}  " + "  ".join(f"{v:>8.3f}" for v in metrics[name].values()))
    print(f"\nRank of the first chunk that answers each question (1 is best, - means not in the top {EVAL_DEPTH})")
    print(f"{'question':<62}  " + "  ".join(f"{n:>{len(n)}}" for n in names))
    for i, label in enumerate(labels):
        print(f"{label['question'][:62]:<62}  " + "  ".join(f"{ranks[n][i] or '-':>{len(n)}}" for n in names))


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="command", required=True)
    sub.add_parser("index").add_argument("folder")
    p_ask = sub.add_parser("ask")
    p_ask.add_argument("question")
    p_ask.add_argument("-k", type=int, default=3)
    p_ask.add_argument("--configs", nargs="+", choices=[c.name for c in CONFIGS], default=["dense-2048", "context-2048"])
    sub.add_parser("eval").add_argument("questions")
    args = parser.parse_args(argv)
    key = os.environ.get("PERPLEXITY_API_KEY")
    if not key:
        sys.exit("export PERPLEXITY_API_KEY before running this script")
    client = httpx.Client(headers={"Authorization": f"Bearer {key}"}, timeout=120)

    if args.command == "index":
        report = build_index(client, args.folder)
        print(f"{report['documents']} documents, {report['chunks']} chunks, indexed in {report['seconds']:.0f} s")
        for model, entry in report["usage"].items():
            cost = f"${entry['cost_usd']:.4f}" if entry["cost_usd"] is not None else "cost not reported by the API"
            print(f"  {model}: {entry['requests']} requests, {entry['tokens']:,} tokens, {cost}")
    elif args.command == "ask":
        ask(client, args.question, args.k, args.configs)
    else:
        labels = [json.loads(line) for line in Path(args.questions).read_text().splitlines() if line.strip()]
        chunks, indexes = load_index()
        check_labels(chunks, labels)
        queries = embed_queries(client, [label["question"] for label in labels])
        print_eval(labels, evaluate(chunks, indexes, queries, labels))


if __name__ == "__main__":
    main()

"""Offline tests for docsearch.py. A fake Embeddings API behind an httpx transport stands in for the service, so no
key is needed. The fake reproduces the behavior recorded from the live API: input_type is required, flat input is dense
and nested input is contextual, each model accepts its own dimensions, vectors come back as base64 int8, and the error
wording is the recorded one. Any other check it makes is the recipe's own."""

import base64
import json
import random

import httpx
import numpy as np
import pytest

import docsearch as ds


def test_chunk_markdown_starts_a_chunk_at_every_heading_and_keeps_the_heading_line():
    text = "# Title\n\nIntro.\n\n## Run it\n\nRun `pytest`.\n\n### Details\n\nMore.\n"
    chunks = ds.chunk_markdown(text)
    assert [h for h, _ in chunks] == ["Title", "Run it", "Details"]
    assert chunks[1][1] == "## Run it\n\nRun `pytest`."


def test_chunk_markdown_ignores_headings_and_blank_lines_inside_code_fences(monkeypatch):
    fence = "`" * 3
    chunks = ds.chunk_markdown(f"## Setup\n\n{fence}python\n# not a heading\nx = 1\n{fence}\n\n## Next\n\nText.\n")
    assert [h for h, _ in chunks] == ["Setup", "Next"]
    assert "# not a heading" in chunks[0][1]
    monkeypatch.setattr(ds, "CHUNK_CHARS", 40)
    chunks = ds.chunk_markdown("## Code\n\n" + "p" * 30 + f"\n\n{fence}\na\n\nb\n{fence}\n")
    assert [text for _, text in chunks] == ["## Code\n\n" + "p" * 30, f"{fence}\na\n\nb\n{fence}"]


def test_chunk_markdown_tracks_the_fence_character_and_length():
    outer, inner = "`" * 4, "`" * 3
    text = f"# Writing Markdown\n\n{outer}markdown\n{inner}python\n# a comment, not a heading\nprint(1)\n{inner}\n{outer}\n\n## Next\n\nBody.\n"
    assert [h for h, _ in ds.chunk_markdown(text)] == ["Writing Markdown", "Next"]
    tildes = "~" * 3
    text = f"# Mixed\n\n{tildes}\n{inner}\n# still code\n{tildes}\n\n## After\n\nBody.\n"
    assert [h for h, _ in ds.chunk_markdown(text)] == ["Mixed", "After"]


def test_chunk_markdown_splits_long_sections_at_blank_lines(monkeypatch):
    monkeypatch.setattr(ds, "CHUNK_CHARS", 60)
    paragraphs = [f"Paragraph {i} " + "x" * 30 for i in range(4)]
    chunks = ds.chunk_markdown("## Long\n\n" + "\n\n".join(paragraphs))
    assert len(chunks) == 4
    assert all(heading == "Long" for heading, _ in chunks)
    assert chunks[0][1].startswith("## Long\n\nParagraph 0")
    assert all("\n\nParagraph" not in text or len(text) <= 60 for _, text in chunks)


def test_chunk_markdown_splits_an_oversized_paragraph_at_line_breaks(monkeypatch):
    monkeypatch.setattr(ds, "CHUNK_CHARS", 50)
    lines = "\n".join(f"* item {i} " + "y" * 20 for i in range(6))
    chunks = ds.chunk_markdown("## List\n\n" + lines)
    assert len(chunks) > 1
    assert "".join(text for _, text in chunks).count("* item") == 6


def test_chunk_markdown_cuts_a_single_line_longer_than_the_chunk_size(monkeypatch):
    monkeypatch.setattr(ds, "CHUNK_CHARS", 1000)
    chunks = ds.chunk_markdown("# Overview\n\n" + "operation " * 40_000)
    assert all(len(text) <= 1000 for _, text in chunks)
    assert "".join(text for _, text in chunks).count("operation") == 40_000
    image = "![screenshot](data:image/png;base64," + "aB9+" * 100_000 + ")"
    chunks = ds.chunk_markdown("# Screenshot\n\n" + image + "\n\n## Next\n\nBody.\n")
    assert all(len(text) <= 1000 for _, text in chunks)
    assert chunks[-1] == ("Next", "## Next\n\nBody.")


def test_load_folder_makes_one_document_per_file_and_splits_long_files(tmp_path, monkeypatch):
    monkeypatch.setattr(ds, "DOC_CHARS", 120)
    (tmp_path / "a.md").write_text("# A\n\nShort file.\n")
    (tmp_path / "sub").mkdir()
    (tmp_path / "sub" / "b.md").write_text("".join(f"## Part {i}\n\n{'z' * 40}\n\n" for i in range(4)))
    documents = ds.load_folder(str(tmp_path))
    names = [doc[0].doc for doc in documents]
    assert names == ["a.md", "sub/b.md", "sub/b.md#2"]
    assert all(sum(len(c.text) for c in doc) <= 120 for doc in documents[1:])
    assert sum(len(doc) for doc in documents) == 5


def test_load_folder_splits_a_file_with_many_short_sections_so_a_document_fits_one_request(tmp_path, monkeypatch):
    monkeypatch.setattr(ds, "REQUEST_CHUNKS", 200)
    (tmp_path / "options.md").write_text("".join(f"## Option {i}\n\nValue.\n\n" for i in range(450)))
    documents = ds.load_folder(str(tmp_path))
    assert [len(doc) for doc in documents] == [200, 200, 50]
    assert [doc[0].doc for doc in documents] == ["options.md", "options.md#2", "options.md#3"]
    size_of = lambda doc: (1, len(doc), sum(len(c.text) for c in doc))
    assert all(len(b) == 1 for b in ds.batches(documents, size_of, (ds.REQUEST_TEXTS, 200, ds.REQUEST_CHARS)))


def test_batches_respect_every_limit_and_refuse_an_item_that_cannot_fit():
    items = ["a" * 10, "b" * 10, "c" * 25, "d" * 10, "e" * 10, "f" * 10]
    by_count = list(ds.batches(items, lambda t: (1, len(t)), (2, 1000)))
    assert [len(b) for b in by_count] == [2, 2, 2]
    by_chars = list(ds.batches(items, lambda t: (1, len(t)), (100, 25)))
    assert by_chars == [["a" * 10, "b" * 10], ["c" * 25], ["d" * 10, "e" * 10], ["f" * 10]]
    with pytest.raises(ValueError, match=r"item 2 alone measures \(1, 25\); the request limits are \(100, 20\)"):
        list(ds.batches(items, lambda t: (1, len(t)), (100, 20)))
    documents = [["x"], ["y", "z"], ["w"]]
    size_of = lambda d: (1, len(d), sum(len(c) for c in d))
    assert list(ds.batches(documents, size_of, (10, 2, 100))) == [[["x"]], [["y", "z"]], [["w"]]]
    with pytest.raises(ValueError, match=r"measures \(1, 2, 2\)"):
        list(ds.batches(documents, size_of, (10, 1, 100)))


class FakeEmbeddingsAPI:
    """POST /v1/embeddings with the request shape and response shape the recipe relies on. The error messages for
    input_type, dimensions, and encoding_format are the live API's wording; the other checks are this fake's own, not
    verified API behavior. Vectors are deterministic: the first value is the text length,
    the second is 1 for dense and 2 for contextual input, the rest are -1, at the requested dimensions. Both models
    return the same vector for the same text, which is the shared-space property the cross-size config relies on.
    Items are shuffled in the response so the code under test has to restore input order."""

    DIMENSIONS = {"pplx-embed-v2-context-9b": (1024, 2048), "pplx-embed-v2-context-0.6b": (512, 1024)}

    def __init__(self, fail_first_with: int | None = None):
        self.bodies: list[dict] = []
        self.fail_first_with = fail_first_with
        self.transport = httpx.MockTransport(self.handle)

    def client(self) -> httpx.Client:
        return httpx.Client(transport=self.transport, headers={"Authorization": "Bearer test-key"})

    def handle(self, request: httpx.Request) -> httpx.Response:
        if request.url.path != "/v1/embeddings" or not request.headers.get("Authorization", "").startswith("Bearer "):
            return httpx.Response(401, json={"error": "unauthorized"})
        if self.fail_first_with:
            status, self.fail_first_with = self.fail_first_with, None
            return httpx.Response(status, json={"error": "try again"})
        body = json.loads(request.content)
        self.bodies.append(body)
        error = self.validate(body)
        if error:
            return httpx.Response(400, json={"error": error})
        dims = body.get("dimensions") or self.DIMENSIONS[body["model"]][-1]
        inputs = body["input"]
        if isinstance(inputs[0], str):
            data = [{"index": i, "embedding": self.vector(t, 1, dims)} for i, t in enumerate(inputs)]
            tokens = sum(len(t) for t in inputs)
        else:
            data = []
            for i, doc in enumerate(inputs):
                inner = [{"index": j, "embedding": self.vector(c, 2, dims)} for j, c in enumerate(doc)]
                random.shuffle(inner)
                data.append({"index": i, "data": inner})
            tokens = sum(len(c) for doc in inputs for c in doc)
        random.shuffle(data)
        usage = {"prompt_tokens": tokens, "total_tokens": tokens, "cost": {"total_cost": tokens / 1e6}}
        return httpx.Response(200, json={"data": data, "model": body["model"], "usage": usage})

    def validate(self, body: dict) -> str | None:
        if body.get("model") not in self.DIMENSIONS:
            return "invalid model"
        if body.get("input_type") not in ("query", "document"):
            return f"input_type is required for model {body.get('model')}; use query or document"
        if "num_chunks" in body:
            return "num_chunks is only supported for image input"
        if body.get("encoding_format", "base64_int8") not in ("base64_int8", "base64_binary"):
            return "encoding format must be one of: base64_int8 base64_binary"
        if body.get("dimensions") is not None and body["dimensions"] not in self.DIMENSIONS[body["model"]]:
            return f"Output dimensions must be one of {list(self.DIMENSIONS[body['model']])}, got {body['dimensions']}."
        inputs = body.get("input")
        if not isinstance(inputs, list) or not inputs or len(inputs) > 512:
            return "input must be a non-empty list of at most 512 items"
        if all(isinstance(x, str) for x in inputs):
            return "empty text" if any(not x for x in inputs) else None
        if not all(isinstance(x, list) and x and all(isinstance(c, str) and c for c in x) for x in inputs):
            return "input must be a list of strings or a list of lists of strings"
        if sum(len(x) for x in inputs) > 16_000:
            return "too many chunks in one contextual request"
        return None

    @staticmethod
    def vector(text: str, mode: int, dims: int) -> str:
        values = np.full(dims, -1, dtype=np.int8)
        values[0], values[1] = min(len(text), 127), mode
        return base64.b64encode(values.tobytes()).decode()


def test_fake_api_rejects_malformed_requests_and_post_surfaces_the_message():
    api = FakeEmbeddingsAPI()
    good = {"model": ds.INDEX_MODEL, "input_type": "document", "input": ["a"], "dimensions": 1024, "encoding_format": "base64_int8"}
    assert len(ds.post(api.client(), good)["data"]) == 1
    for change, message in [
        ({"input_type": None}, "input_type is required"),
        ({"input": ["a", ["b"]]}, "list of strings or a list of lists"),
        ({"dimensions": 512}, r"Output dimensions must be one of \[1024, 2048\], got 512"),
        ({"model": ds.SMALL_MODEL, "dimensions": 2048}, r"Output dimensions must be one of \[512, 1024\], got 2048"),
        ({"encoding_format": "float"}, "encoding format must be one of: base64_int8 base64_binary"),
        ({"input": ["a", ""]}, "empty text"),
        ({"num_chunks": 4}, "only supported for image input"),
    ]:
        with pytest.raises(RuntimeError, match=f"HTTP 400 from the Embeddings API: .*{message}"):
            ds.post(api.client(), {**good, **change})