278 lines
11 KiB
Python
278 lines
11 KiB
Python
"""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,
|
|
platform: str = None,
|
|
tier: int = 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, platform, tier)
|
|
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, platform: str = None, tier: int = None
|
|
) -> list:
|
|
conditions = []
|
|
if area:
|
|
conditions.append({"key": "area", "match": {"value": area}})
|
|
if confidence:
|
|
conditions.append({"key": "confidence", "match": {"value": confidence}})
|
|
if platform:
|
|
conditions.append({"key": "platform", "match": {"value": platform}})
|
|
if tier is not None:
|
|
conditions.append({"key": "tier", "match": {"value": int(tier)}})
|
|
return conditions
|