atif-sql 0.2.0__py3-none-any.whl

atif_embed/infrastructure/corpus_text_rows.py

@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Contract-layout corpus reader behind :class:`TextRowsPort`.
+
+atif-embed may never import atif-duck (import-linter independence contract),
+so this adapter reads the CONTRACT corpus layout
+(``<corpus_root>/sessions/<id>/trajectory.json``) directly with its OWN
+DuckDB connection + ``read_json`` — the same move atif-corpus made with
+``ConverterPort``: depend on the contract's artifact shapes, not on a
+sibling package's registry. The trajectory paths reach ``read_json`` as a
+bound list parameter, never as statement text.
+
+Selection semantics (CONTRACT-V2 §VSS):
+
+* the embeddable unit is a step's flattened text (``str | ContentPart[]``,
+  the ARRAY branch joined with blank lines — mirrors atif-duck's ``steps``
+  view and harbor's own text bundling);
+* main chain AND sidechain steps are included (no ``is_sidechain`` filter);
+* only texts of >= 32 characters qualify — shorter step texts are acks and
+  noise;
+* the row key is the step's PRIMARY uuid: the FIRST ``extra.source_uuids``
+  entry. Harbor bundles all raw records sharing an assistant ``message.id``
+  into one step, so a step maps to 1..n raw uuids; the first entry is the
+  bundle's head record, which is the uuid the ``messages`` view (edges.jsonl)
+  can join back to. Steps without source uuids are skipped — they cannot be
+  keyed against the raw-record surface.
+
+Torn-set guard: like atif-duck's raw readers, only session dirs with a
+``meta.json`` (written last by atif-corpus) contribute rows.
+
+Stored names: materialize stores the trajectory compressed, as
+``trajectory.json.zst``; a corpus written before that keeps the plain
+``trajectory.json`` until ``atif-sql corpus slim`` compresses it. Each
+session's stored file is resolved once (compressed first) and DuckDB
+decompresses it, so the rows are the same either way. Batches are sized by
+the decompressed bytes, which is what a statement holds.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from loguru import logger
+
+from atif_embed.domain.discovery import PendingSelection
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from pathlib import Path
+
+    from atif_embed.domain.text_stamp import PendingText
+
+#: Minimum characters for a step text to be worth embedding.
+MIN_TEXT_CHARS = 32
+
+#: ``read_json`` upper bound — live trajectory.json files reach 85 MB
+#: (harbor inlines subagent sidechains); 1 GiB gives ~10x headroom.
+_MAX_OBJECT_SIZE = 1_073_741_824
+
+#: The trajectory's name, and the suffix it carries when stored compressed
+#: (twins of atif-corpus's, pinned in atif-cli's tests).
+TRAJECTORY_FILENAME = "trajectory.json"
+COMPRESSED_SUFFIX = ".zst"
+
+#: The largest zstd frame header: enough bytes to read any frame's content size.
+_FRAME_HEADER_MAX_BYTES = 18
+
+#: Rows pulled off a DuckDB result per ``fetchmany``.
+_FETCH_PAGE_ROWS = 512
+
+#: Trajectory BYTES a single DuckDB statement may read. Residency scales with
+#: the bytes a statement materializes, not with the file count, so batching by
+#: bytes is what makes both ends bounded: thousands of small sessions are read
+#: together (few statements) while a session above the budget gets a statement
+#: to itself (peak stays near that one session). See this package's README for
+#: the measured curve across corpus shapes.
+_BATCH_MAX_BYTES = 4 * 1024 * 1024
+
+
+def _step_texts_sql() -> str:
+    """SQL selecting ``(uuid, text_content)`` from ONE BATCH of trajectories.
+
+    The batch's paths are NOT in this text: ``read_json(?)`` takes them as one
+    bound list parameter at execute time, so a path is data to DuckDB and
+    never part of the statement. The only interpolations are the two integer
+    constants below.
+
+    One trajectory document per file -> ``format='auto'`` (NOT NDJSON). The
+    explicit ``columns`` projection is a strict filter AND it skips JSON
+    schema inference, the dominant cost on a large corpus. The flattened-text
+    CASE lives once in a CTE so the >=32-char floor and the projection read
+    the same expression.
+
+    ``filename=true`` plus ``ORDER BY filename, step_id`` keeps row order
+    identical to reading the batch's files one at a time, so batch boundaries
+    cannot change which rows a ``--limit`` run picks. The two spellings of a
+    session's file differ only after its directory name, so a corpus mixing
+    them sorts in the same session order.
+    """
+    return f"""
+        WITH step_texts AS (
+            SELECT
+                t.filename                                           AS filename,
+                json_extract(step, '$.step_id')::BIGINT              AS step_id,
+                json_extract_string(step, '$.extra.source_uuids[0]') AS uuid,
+                CASE WHEN json_type(step, '$.message') = 'ARRAY'
+                     THEN array_to_string(
+                              list_transform(
+                                  json_extract(step, '$.message[*].text'),
+                                  part -> json_extract_string(part, '$')
+                              ),
+                              '\n\n'
+                          )
+                     ELSE json_extract_string(step, '$.message')
+                END                                                  AS text_content
+            FROM read_json(
+                     ?,
+                     format='auto',
+                     filename=true,
+                     columns={{steps: 'JSON[]'}},
+                     maximum_object_size={_MAX_OBJECT_SIZE}
+                 ) t,
+                 UNNEST(t.steps) AS s(step)
+        )
+        SELECT uuid, text_content
+        FROM step_texts
+        WHERE uuid IS NOT NULL
+          AND text_content IS NOT NULL
+          AND length(text_content) >= {MIN_TEXT_CHARS}
+        ORDER BY filename, step_id
+        """  # noqa: S608  # nosec B608 - paths are a bound parameter; both limits are int constants
+
+
+def _decoded_size(path: Path) -> int:
+    """The trajectory's size once decompressed: what a statement reading it holds."""
+    if not path.name.endswith(COMPRESSED_SUFFIX):
+        return path.stat().st_size
+    import zstandard
+
+    with path.open("rb") as handle:
+        size = zstandard.frame_content_size(handle.read(_FRAME_HEADER_MAX_BYTES))
+    # A frame without a recorded size (not written by atif-sql): its stored
+    # size is the only cheap stand-in, and it merely under-fills a batch.
+    return int(size) if size >= 0 else path.stat().st_size
+
+
+def _stored_trajectory(session_dir: Path) -> Path | None:
+    for name in (f"{TRAJECTORY_FILENAME}{COMPRESSED_SUFFIX}", TRAJECTORY_FILENAME):
+        path = session_dir / name
+        if path.is_file():
+            return path
+    return None
+
+
+def _complete_trajectory_paths(corpus_root: Path) -> list[tuple[str, int]]:
+    """``(path, decoded size)`` for every COMPLETE session dir (meta.json present)."""
+    sessions_dir = corpus_root / "sessions"
+    if not sessions_dir.is_dir():
+        return []
+    paths: list[tuple[str, int]] = []
+    for session_dir in sorted(sessions_dir.iterdir()):
+        if not session_dir.is_dir():
+            continue
+        if not (session_dir / "meta.json").is_file():
+            logger.warning(
+                "Skipping incomplete session dir {} (no meta.json); excluded from embed",
+                session_dir,
+            )
+            continue
+        trajectory = _stored_trajectory(session_dir)
+        if trajectory is not None:
+            paths.append((str(trajectory), _decoded_size(trajectory)))
+    return paths
+
+
+def _batch_by_bytes(paths: list[tuple[str, int]], *, max_bytes: int) -> list[list[str]]:
+    """Group ``(path, size)`` pairs into batches of at most ``max_bytes``.
+
+    Order is preserved, so batching cannot reorder rows. A file at or above
+    the budget forms a batch of one rather than being dropped or merged: the
+    budget is a target, and the true residency floor is the largest single
+    session, which no grouping can lower.
+    """
+    batches: list[list[str]] = []
+    current: list[str] = []
+    current_bytes = 0
+    for path, size in paths:
+        if current and current_bytes + size > max_bytes:
+            batches.append(current)
+            current = []
+            current_bytes = 0
+        current.append(path)
+        current_bytes += size
+    if current:
+        batches.append(current)
+    return batches
+
+
+class DuckDbTextRows:
+    """:class:`~atif_embed.domain.ports.TextRowsPort` over the contract corpus."""
+
+    #: This reader has one path: every complete session's trajectory.
+    discovery = "corpus"
+
+    def iter_unembedded(
+        self,
+        corpus_root: Path,
+        *,
+        embedded: dict[str, str] | None = None,
+        limit: int | None = None,
+    ) -> Iterator[PendingText]:
+        """Yield the texts needing embedding, one byte-bounded batch at a time.
+
+        DuckDB materializes a result set fully inside ``execute()``, so a
+        single query over every trajectory path holds the whole corpus's text
+        regardless of how it is paged off the result; splitting into
+        statements is what bounds the peak, and ``fetchmany`` bounds the
+        Python-side row buffer within a batch.
+
+        The split is by total file BYTES
+        (:data:`_BATCH_MAX_BYTES`), not per file. Residency tracks the bytes
+        one statement reads, while wall time tracks the statement COUNT, so
+        bytes is the axis that bounds both: a corpus of thousands of small
+        sessions becomes a few statements instead of thousands, and a session
+        larger than the budget still gets a statement to itself. Peak resident
+        text is therefore the batch budget or the largest single session,
+        whichever is larger.
+
+        The staleness join against ``embedded`` (the store's
+        ``{uuid: text_hash}`` map) happens in Python BEFORE ``limit`` is
+        applied. Ordering matters: filtering first guarantees ``--limit N``
+        always makes N rows of forward progress, whereas a SQL LIMIT ahead
+        of the filter spends the cap on rows that are already embedded and a
+        bounded run can make ZERO progress. At corpus scale the comparison
+        is cheap. A uuid whose stored hash differs from the current text is
+        STALE and is re-yielded with ``replaces_existing=True``.
+
+        Rows are deterministic: ordered by trajectory path then step_id, and
+        de-duplicated on uuid (first occurrence wins) so the Lance store
+        never receives two vectors for one key.
+        """
+        paths = _complete_trajectory_paths(corpus_root)
+        if not paths:
+            logger.info("No complete sessions under {} — nothing to embed", corpus_root)
+            return
+        selection = PendingSelection(embedded=embedded, limit=limit)
+        yield from selection.select(_step_rows(paths))
+
+    def commit(self, *, stored_rows: int) -> None:
+        """Nothing to record: this reader has no watermark and re-reads every session."""
+        del stored_rows
+
+
+def _step_rows(paths: list[tuple[str, int]]) -> Iterator[tuple[str, str]]:
+    """``(uuid, text)`` for every qualifying step, one byte-bounded batch per statement.
+
+    A generator, so a consumer that stops early (``--limit``) stops the
+    reads too: no batch past the one it's in is ever executed.
+    """
+    import duckdb
+
+    con = duckdb.connect(":memory:")
+    try:
+        for batch in _batch_by_bytes(paths, max_bytes=_BATCH_MAX_BYTES):
+            result = con.execute(_step_texts_sql(), [batch])
+            while page := result.fetchmany(_FETCH_PAGE_ROWS):
+                for row in page:
+                    yield str(row[0]), str(row[1])
+    finally:
+        con.close()
+
+
+__all__ = ["MIN_TEXT_CHARS", "DuckDbTextRows"]

atif_embed/infrastructure/gemma_local.py

@@ -0,0 +1,355 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""EmbeddingGemma 2 on this machine: the local ``EmbeddingProvider`` adapter.
+
+Loads ``google/embeddinggemma-2`` through sentence-transformers in its
+text-only setup (``config_kwargs`` drops the vision and audio encoders, so
+270M of the model's 740M parameters load), at a pinned repository revision.
+Documents encode under the model's ``Document`` prompt (``title: none |
+text: ``) and queries under ``SearchQuery`` (``task: search result | query:
+``), truncated to the configured Matryoshka width and re-normalized.
+
+Numerical precision follows the model card: bfloat16 on a CUDA device that
+supports it, float32 everywhere else, and never float16, whose range the
+model's activations exceed (it returns NaN or degraded vectors without
+raising). A non-finite vector is dropped as a failed slot rather than stored.
+
+The heavy libraries (torch, transformers, sentence-transformers) come from
+the optional ``local`` extra and are imported by name at load time, never at
+module scope: the default Cohere path, CI, and both type checkers run without
+them. Selecting this provider without the extra raises
+:class:`~atif_embed.domain.errors.EmbeddingProviderNotInstalled`.
+
+The model downloads once into the Hugging Face cache (about 1.5 GB; the
+checkpoint is one file even when only the text tower loads) and later runs
+read it from there; ``HF_HUB_OFFLINE=1`` keeps every run off the network.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import importlib
+import math
+import threading
+import time
+from typing import TYPE_CHECKING, Any
+
+from loguru import logger
+
+from atif_embed.domain.errors import (
+    EmbeddingProviderNotInstalled,
+    EmbeddingProviderUnavailable,
+    EmbeddingResponseInvalid,
+)
+from atif_embed.domain.text_stamp import MAX_EMBEDDABLE_CHARS, clip_text
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from atif_embed.infrastructure.settings import EmbedSettings
+
+#: sentence-transformers prompt names from the model's own
+#: ``config_sentence_transformers.json``: the corpus side and the query side
+#: of its asymmetric search task.
+DOCUMENT_PROMPT = "Document"
+QUERY_PROMPT = "SearchQuery"
+
+#: ``AutoConfig`` overrides that load the text tower alone. Text-only and
+#: full-model embeddings share one vector space, so the stamped model id
+#: carries no mode suffix.
+TEXT_ONLY_CONFIG: dict[str, None] = {"vision_config": None, "audio_config": None}
+
+#: The model's context window. The checkpoint's tokenizer declares no
+#: maximum (sentence-transformers reads it as 10**30), so without this cap a
+#: 50,000-character step would run past the window the model was trained on.
+MAX_TOKENS = 8192
+
+#: How to get the extra, named in the error a missing install raises.
+INSTALL_HINT = (
+    "install the local extra: `uv sync --all-packages --extra local` in a checkout, or "
+    "`uv tool install 'atif-sql[local]' --torch-backend cpu`; or unset "
+    "ATIF_SQL_EMBED_PROVIDER to embed with Cohere on Bedrock"
+)
+
+#: What a forward pass can raise that one batch's loss should absorb: torch
+#: raises RuntimeError (its out-of-memory error included), tokenizers raise
+#: ValueError, and a device fault surfaces as OSError.
+_BATCH_ERRORS = (RuntimeError, ValueError, OSError)
+
+#: Padded characters one forward pass may hold: its text count times its
+#: longest text. A batch pads every text to its longest, and attention memory
+#: grows with the square of that length, so 32 texts near the model's
+#: 8,192-token window took 26 GB of RSS on CPU while short texts batch freely.
+#: At about 4 characters a token this is one full-window text per pass.
+BATCH_CHAR_BUDGET = 32_000
+
+
+def plan_batches(lengths: list[int], *, max_texts: int) -> list[list[int]]:
+    """Group text positions shortest first into batches a forward pass can hold.
+
+    A batch closes at ``max_texts`` texts, or when the next (longer) text
+    would push ``texts x longest`` past :data:`BATCH_CHAR_BUDGET`. A text
+    over the budget on its own still gets a batch of one.
+    """
+    batches: list[list[int]] = []
+    current: list[int] = []
+    for position in sorted(range(len(lengths)), key=lambda i: lengths[i]):
+        if current and (
+            len(current) >= max_texts or (len(current) + 1) * lengths[position] > BATCH_CHAR_BUDGET
+        ):
+            batches.append(current)
+            current = []
+        current.append(position)
+    if current:
+        batches.append(current)
+    return batches
+
+
+def _import(name: str) -> Any:
+    """Import one module of the ``local`` extra, or say how to install it.
+
+    By name rather than as a statement: the extra is optional, and the type
+    checkers and CI run without it, so a static import would not resolve.
+    """
+    try:
+        return importlib.import_module(name)
+    except ImportError as exc:
+        msg = f"the gemma embedding provider needs {name}, which is not installed; {INSTALL_HINT}"
+        raise EmbeddingProviderNotInstalled(msg) from exc
+
+
+def pick_device(requested: str, torch: Any) -> str:
+    """The device to load on: ``requested``, or for ``auto`` CUDA, then MPS, then CPU."""
+    if requested != "auto":
+        return requested
+    if torch.cuda.is_available():
+        return "cuda"
+    mps = getattr(torch.backends, "mps", None)
+    if mps is not None and mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+def pick_dtype(device: str, torch: Any) -> Any:
+    """bfloat16 on a CUDA device that supports it, else float32; never float16."""
+    if device.startswith("cuda") and torch.cuda.is_bf16_supported():
+        return torch.bfloat16
+    return torch.float32
+
+
+def load_sentence_transformer(settings: EmbedSettings) -> Any:
+    """Load EmbeddingGemma 2 text-only at the pinned revision, on the chosen device.
+
+    Raises
+    ------
+    EmbeddingProviderNotInstalled
+        torch or sentence-transformers is not installed.
+    EmbeddingProviderUnavailable
+        The model could not be fetched or loaded (no network and no cached
+        copy, a revision the hub does not have, a device that is not there).
+    """
+    torch = _import("torch")
+    sentence_transformers = _import("sentence_transformers")
+    device = pick_device(settings.gemma_device, torch)
+    dtype = pick_dtype(device, torch)
+    logger.info(
+        "Loading {} at {} text-only on {} ({})",
+        settings.gemma_model_id,
+        settings.gemma_revision[:12],
+        device,
+        str(dtype).removeprefix("torch."),
+    )
+    try:
+        model = sentence_transformers.SentenceTransformer(
+            settings.gemma_model_id,
+            revision=settings.gemma_revision,
+            device=device,
+            config_kwargs=dict(TEXT_ONLY_CONFIG),
+            model_kwargs={"dtype": dtype},
+        )
+    except (OSError, RuntimeError, ValueError) as exc:
+        msg = (
+            f"could not load {settings.gemma_model_id} at revision {settings.gemma_revision} "
+            f"on {device}: {type(exc).__name__}: {exc}"
+        )
+        raise EmbeddingProviderUnavailable(msg) from exc
+    model.max_seq_length = MAX_TOKENS
+    return model
+
+
+def _finite_rows(encoded: Any, expected: int, dim: int) -> list[list[float] | None]:
+    """The encoder's rows as float lists, a non-finite or mis-sized row as ``None``."""
+    rows = encoded.tolist() if hasattr(encoded, "tolist") else list(encoded)
+    if len(rows) != expected:
+        msg = f"the local model returned {len(rows)} vectors for {expected} texts"
+        raise EmbeddingResponseInvalid(msg)
+    out: list[list[float] | None] = []
+    for row in rows:
+        vector = [float(x) for x in row]
+        out.append(vector if len(vector) == dim and all(map(math.isfinite, vector)) else None)
+    return out
+
+
+class EmbeddingGemmaLocalEmbedder:
+    """``EmbeddingProvider`` over EmbeddingGemma 2, run on this machine.
+
+    The model loads once, lazily, on the first call that needs it, so
+    constructing the embedder costs nothing and a run with nothing pending
+    never touches torch. ``model_factory`` replaces the loader (tests pass a
+    fake; production uses :func:`load_sentence_transformer`).
+    """
+
+    provider = "gemma-local"
+
+    def __init__(
+        self,
+        settings: EmbedSettings,
+        *,
+        model_factory: Callable[[EmbedSettings], Any] | None = None,
+    ) -> None:
+        self._settings = settings
+        self._factory = model_factory or load_sentence_transformer
+        self._model: Any = None
+        self._lock = threading.Lock()
+
+    @property
+    def model_id(self) -> str:
+        """The Hugging Face repository this embedder loads, stamped on every row."""
+        return self._settings.gemma_model_id
+
+    @property
+    def dimension(self) -> int:
+        """The Matryoshka width every vector is truncated to (768, 512, 256 or 128)."""
+        return self._settings.embedding_dim
+
+    def _loaded(self) -> Any:
+        """The model, loaded on first use and shared by every later call."""
+        with self._lock:
+            if self._model is None:
+                self._model = self._factory(self._settings)
+            return self._model
+
+    def _encode(self, model: Any, texts: list[str], *, prompt_name: str) -> Any:
+        """One forward pass: prompt, truncate to the width, re-normalize."""
+        return model.encode(
+            texts,
+            prompt_name=prompt_name,
+            batch_size=len(texts),
+            truncate_dim=self.dimension,
+            normalize_embeddings=True,
+            convert_to_numpy=True,
+            show_progress_bar=False,
+        )
+
+    async def embed_documents(self, texts: list[str]) -> list[list[float] | None]:
+        """Embed corpus documents one batch at a time; one slot per input text, in order.
+
+        Texts are batched shortest first under :func:`plan_batches`, so each
+        forward pass pads to a similar length and a long text never pads a
+        whole batch to its length, then put back in input order. A batch that
+        raises is
+        logged and left as ``None`` slots, and so is any single non-finite
+        vector; the other batches are kept, and the next run's staleness
+        anti-join re-picks the rows that are missing.
+        """
+        if not texts:
+            return []
+        model = await asyncio.to_thread(self._loaded)
+
+        clipped: list[str] = []
+        for position, text in enumerate(texts):
+            sent, was_truncated = clip_text(text)
+            if was_truncated:
+                logger.warning(
+                    "Clipping text at position {} for embedding: {} chars -> {} "
+                    "(content past the cap will not match a search)",
+                    position,
+                    len(text),
+                    MAX_EMBEDDABLE_CHARS,
+                )
+            clipped.append(sent)
+
+        batches = plan_batches([len(t) for t in clipped], max_texts=self._settings.gemma_batch_size)
+        vectors: list[list[float] | None] = [None] * len(texts)
+        failed_batches = 0
+        n_batches = len(batches)
+        t0 = time.monotonic()
+        for index, positions in enumerate(batches):
+            batch = [clipped[i] for i in positions]
+            try:
+                encoded = await asyncio.to_thread(
+                    self._encode, model, batch, prompt_name=DOCUMENT_PROMPT
+                )
+                rows = _finite_rows(encoded, len(batch), self.dimension)
+            except (*_BATCH_ERRORS, EmbeddingResponseInvalid) as exc:
+                failed_batches += 1
+                logger.error(
+                    "Local batch {} ({} texts) failed ({}: {}); keeping the batches that "
+                    "succeeded — the next run re-picks these rows",
+                    index,
+                    len(batch),
+                    type(exc).__name__,
+                    exc,
+                )
+                continue
+            for position, row in zip(positions, rows, strict=True):
+                vectors[position] = row
+        elapsed = time.monotonic() - t0
+
+        embedded = sum(1 for v in vectors if v is not None)
+        dropped = len(texts) - embedded
+        if dropped and not failed_batches:
+            logger.error("Dropped {} non-finite vectors; those rows stay unembedded", dropped)
+        logger.info(
+            "Embedded {} vectors locally across {}/{} batches in {:.2f}s ({:.1f} vec/s)",
+            embedded,
+            n_batches - failed_batches,
+            n_batches,
+            elapsed,
+            embedded / elapsed if elapsed > 0 else 0.0,
+        )
+        return vectors
+
+    def embed_query(self, text: str) -> list[float]:
+        """Embed one query under the ``SearchQuery`` prompt, at :attr:`dimension`.
+
+        Raises
+        ------
+        EmbeddingProviderUnavailable
+            The forward pass failed.
+        EmbeddingResponseInvalid
+            The vector came back non-finite or at the wrong width.
+        """
+        sent, was_truncated = clip_text(text)
+        if was_truncated:
+            logger.warning(
+                "Clipping query for embedding: {} chars -> {}", len(text), MAX_EMBEDDABLE_CHARS
+            )
+        model = self._loaded()
+        try:
+            encoded = self._encode(model, [sent], prompt_name=QUERY_PROMPT)
+        except _BATCH_ERRORS as exc:
+            msg = f"the local model failed on the query: {type(exc).__name__}: {exc}"
+            raise EmbeddingProviderUnavailable(msg) from exc
+        vector = _finite_rows(encoded, 1, self.dimension)[0]
+        if vector is None:
+            msg = (
+                f"the local model returned a non-finite or {self.dimension}-mismatched query vector"
+            )
+            raise EmbeddingResponseInvalid(msg)
+        return vector
+
+
+__all__ = [
+    "BATCH_CHAR_BUDGET",
+    "DOCUMENT_PROMPT",
+    "INSTALL_HINT",
+    "MAX_TOKENS",
+    "QUERY_PROMPT",
+    "TEXT_ONLY_CONFIG",
+    "EmbeddingGemmaLocalEmbedder",
+    "load_sentence_transformer",
+    "pick_device",
+    "pick_dtype",
+    "plan_batches",
+]