ai-resume/corpus/qdrant.py

"""Qdrant REST client — standard library only.

No `qdrant-client` dependency: the four calls this pipeline needs are HTTP
requests, and AGENTS.md §4.1 freezes the stack. `_request` is the single
transport hook, which is what the tests replace to run the whole pipeline
without a container.

Qdrant is bound to `127.0.0.1:6333` on the runtime host (`.13`) and is never
publicly reachable; Caddy is the only public listener.
"""

from __future__ import annotations

import http.client
import json
import time
import urllib.error
import urllib.request

DISTANCE_COSINE = "Cosine"

# The runtime host (.13) is memory-bound: 2.3 GiB free and 8 GiB of swap already
# in use, so Qdrant gets a 400 m cap and must live inside it
# (reference/stack.md). Measured on 2026-10-26: the *default* HNSW and payload
# settings are OOM-killed at the cap between 5,000 and 20,000 chunks of 1024
# dimensions. Keeping the HNSW graph and the payloads on disk, and using a
# smaller graph (m=8, ef_construct=64), fits the cap because the kernel can
# reclaim those page-cache pages instead of killing the process. Vectors stay in
# RAM because they are the part search reads on every query.
#
# The trade is build time and a little search latency; the corpus is ~50 chunks
# today and the public site has no query burst. See docs/phase-4-notes.md §7.
MEMORY_FRUGAL_DEFAULTS = {
    "on_disk_payload": True,
    "hnsw_m": 8,
    "hnsw_ef_construct": 64,
    "vectors_on_disk": False,
}

# Segment behaviour under a cap. `memmap_threshold` makes Qdrant mmap segments
# this size and larger (page cache the kernel can reclaim instead of OOM), and
# `indexing_threshold` is raised so a bulk ingest does not build the HNSW graph
# while it is still writing segments. Measured: 20,000 chunks at 1024
# dimensions peaked at 296 MiB during ingest and settled to 96 MiB.
MEMORY_FRUGAL_OPTIMIZERS = {
    "default_segment_number": 2,
    "memmap_threshold": 5000,
    "indexing_threshold": 20000,
}


class QdrantError(RuntimeError):
    """A Qdrant call failed, or Qdrant is not reachable."""


class Qdrant:
    def __init__(self, url: str = "http://127.0.0.1:6333", timeout: float = 30.0, attempts: int = 3) -> None:
        self.url = url.rstrip("/")
        self.timeout = timeout
        self.attempts = max(1, attempts)

    # -- transport --------------------------------------------------------

    def _request(self, method: str, path: str, body=None) -> dict:
        """One HTTP call with retry/backoff. The only place that touches the network."""
        data = None if body is None else json.dumps(body).encode("utf-8")
        last_error = None
        for attempt in range(self.attempts):
            request = urllib.request.Request(
                f"{self.url}{path}",
                data=data,
                method=method,
                headers={"Content-Type": "application/json"} if data else {},
            )
            try:
                with urllib.request.urlopen(request, timeout=self.timeout) as response:
                    raw = response.read().decode("utf-8")
                return json.loads(raw) if raw.strip() else {}
            except urllib.error.HTTPError as exc:
                detail = exc.read().decode("utf-8", "replace")[:400]
                exc.close()
                if exc.code == 404:
                    raise QdrantNotFound(f"{method} {path}: {detail}") from None
                if exc.code != 429 and exc.code < 500:
                    raise QdrantError(f"{method} {path} failed: HTTP {exc.code} {detail}") from None
                last_error = f"HTTP {exc.code} {detail}"
            except (urllib.error.URLError, TimeoutError, http.client.HTTPException, ConnectionError) as exc:
                # RemoteDisconnected and friends arrive here when the server dies
                # mid-request (an OOM kill looks exactly like this).
                last_error = f"cannot reach {self.url}: {getattr(exc, 'reason', exc)}"
            if attempt + 1 < self.attempts:
                time.sleep(0.5 * (2 ** attempt))
        raise QdrantError(f"{method} {path} failed after {self.attempts} attempts: {last_error}")

    # -- cluster ----------------------------------------------------------

    def health(self) -> bool:
        """True when Qdrant answers. Never raises."""
        try:
            self._request("GET", "/")
            return True
        except (QdrantError, QdrantNotFound):
            return False

    def version(self) -> str:
        response = self._request("GET", "/")
        return str(response.get("version", "unknown"))

    # -- collections ------------------------------------------------------

    def collection_exists(self, name: str) -> bool:
        try:
            self._request("GET", f"/collections/{name}")
            return True
        except QdrantNotFound:
            return False

    def create_collection(
        self,
        name: str,
        dim: int,
        distance: str = DISTANCE_COSINE,
        on_disk_payload: bool = None,
        hnsw_m: int = None,
        hnsw_ef_construct: int = None,
        vectors_on_disk: bool = None,
    ) -> None:
        """Create a collection sized for a memory-bound host (see MEMORY_FRUGAL_DEFAULTS)."""
        defaults = MEMORY_FRUGAL_DEFAULTS
        body = {
            "vectors": {
                "size": int(dim),
                "distance": distance,
                "on_disk": bool(
                    defaults["vectors_on_disk"] if vectors_on_disk is None else vectors_on_disk
                ),
            },
            "hnsw_config": {
                "m": int(defaults["hnsw_m"] if hnsw_m is None else hnsw_m),
                "ef_construct": int(
                    defaults["hnsw_ef_construct"] if hnsw_ef_construct is None else hnsw_ef_construct
                ),
                "on_disk": True,
            },
            "on_disk_payload": bool(
                defaults["on_disk_payload"] if on_disk_payload is None else on_disk_payload
            ),
            "optimizers_config": dict(MEMORY_FRUGAL_OPTIMIZERS),
        }
        self._request("PUT", f"/collections/{name}", body)

    def drop_collection(self, name: str) -> None:
        try:
            self._request("DELETE", f"/collections/{name}")
        except QdrantNotFound:
            pass

    def collection_info(self, name: str) -> dict:
        response = self._request("GET", f"/collections/{name}")
        return dict(response.get("result") or {})

    def point_count(self, name: str) -> int:
        response = self._request("POST", f"/collections/{name}/points/count", {"exact": True})
        return int((response.get("result") or {}).get("count", 0))

    # -- points -----------------------------------------------------------

    def upsert(self, name: str, points: list, batch: int = 256) -> int:
        """Upsert points in batches. Returns the number of points sent."""
        sent = 0
        for start in range(0, len(points), max(1, batch)):
            chunk = points[start : start + max(1, batch)]
            self._request(
                "PUT",
                f"/collections/{name}/points?wait=true",
                {"points": chunk},
            )
            sent += len(chunk)
        return sent

    def delete_by_source(self, name: str, source_path: str) -> None:
        """Delete every point that came from one document.

        Re-chunking a document can produce fewer chunks than before, so the old
        points are removed rather than overwritten. Deterministic ids make
        duplicates impossible; this makes orphans impossible.
        """
        self._request(
            "POST",
            f"/collections/{name}/points/delete?wait=true",
            {"filter": {"must": [{"key": "source_path", "match": {"value": source_path}}]}},
        )

    def search(
        self,
        name: str,
        vector: list,
        limit: int = 5,
        area: str = None,
        confidence: str = None,
        with_payload: bool = True,
    ) -> list:
        """Nearest neighbours, optionally filtered *before* the search."""
        body = {
            "query": [float(value) for value in vector],
            "limit": int(limit),
            "with_payload": bool(with_payload),
        }
        filters = _filter_conditions(area, confidence)
        if filters:
            body["filter"] = {"must": filters}
        try:
            response = self._request("POST", f"/collections/{name}/points/query", body)
            points = ((response.get("result") or {}).get("points")) or []
        except QdrantNotFound:
            # Older Qdrant exposes the same operation as /points/search.
            response = self._request("POST", f"/collections/{name}/points/search", body)
            points = response.get("result") or []
        return [
            {
                "id": point.get("id"),
                "score": float(point.get("score") or 0.0),
                "payload": point.get("payload") or {},
            }
            for point in points
        ]

    def scroll_sources(self, name: str, limit: int = 1000) -> dict:
        """Map source_path -> number of vectors, for consistency checks."""
        counts = {}
        offset = None
        while True:
            body = {"limit": limit, "with_payload": ["source_path"], "with_vector": False}
            if offset is not None:
                body["offset"] = offset
            response = self._request("POST", f"/collections/{name}/points/scroll", body)
            result = response.get("result") or {}
            points = result.get("points") or []
            for point in points:
                source = (point.get("payload") or {}).get("source_path", "<missing>")
                counts[source] = counts.get(source, 0) + 1
            offset = result.get("next_page_offset")
            if not points or offset is None:
                return counts

    def metrics(self) -> str:
        """Prometheus text from Qdrant's /metrics route (used by the memory measurement).

        Returns "" when the route is not exposed, so callers can fall back to
        reading the container's own memory counters.
        """
        try:
            with urllib.request.urlopen(f"{self.url}/metrics", timeout=self.timeout) as response:
                return response.read().decode("utf-8", "replace")
        except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError):
            return ""


class QdrantNotFound(QdrantError):
    """A 404 from Qdrant: the collection or route does not exist."""


def _filter_conditions(area: str = None, confidence: str = None) -> list:
    conditions = []
    if area:
        conditions.append({"key": "area", "match": {"value": area}})
    if confidence:
        conditions.append({"key": "confidence", "match": {"value": confidence}})
    return conditions