247 lines
8.7 KiB
Python
247 lines
8.7 KiB
Python
"""Environment probe — readiness report shown at startup.
|
|
|
|
Checks (all cheap, all non-fatal):
|
|
|
|
* Python module availability for every pipeline dependency
|
|
(importlib.util.find_spec — no heavy imports).
|
|
* Qdrant service reachability (get_collections with a short timeout).
|
|
* Ollama service reachability (GET /api/tags, also lists installed
|
|
models so the UI can pre-fill the model combo).
|
|
|
|
The GUI runs this in a background thread because the FastEmbed model
|
|
check and service pings can take seconds on a cold start.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import platform
|
|
import urllib.request
|
|
from dataclasses import dataclass, field
|
|
|
|
# module name -> role in the pipeline
|
|
PROBED_MODULES: dict[str, str] = {
|
|
"bs4": "EPUB parsing",
|
|
"ebooklib": "EPUB parsing",
|
|
"pypdf": "PDF parsing",
|
|
"slugify": "vault note filenames",
|
|
"fastembed": "local embeddings",
|
|
"qdrant_client": "vector target: qdrant",
|
|
"chromadb": "vector target: chroma (optional)",
|
|
"lancedb": "vector target: lancedb (optional)",
|
|
"faiss": "vector target: faiss (optional)",
|
|
"pymilvus": "vector target: milvus lite (optional)",
|
|
"weaviate": "vector target: weaviate (optional)",
|
|
"psycopg": "vector target: pgvector (optional)",
|
|
"pgvector": "vector target: pgvector (optional)",
|
|
"duckdb": "vector target: duckdb (optional)",
|
|
"sqlite_vec": "vector target: sqlite-vec (optional)",
|
|
"pymysql": "vector target: mariadb (optional)",
|
|
"graphify": "graph engine (optional)",
|
|
"vanna": "ask interaction (optional)",
|
|
"minio": "object archive (optional)",
|
|
"lightrag": "knowledge graph (optional)",
|
|
}
|
|
|
|
|
|
def _module_available(name: str) -> bool:
|
|
try:
|
|
return importlib.util.find_spec(name) is not None
|
|
except (ImportError, ValueError):
|
|
return False
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class EnvProbe:
|
|
python_version: str = ""
|
|
modules: dict[str, bool] = field(default_factory=dict)
|
|
qdrant_up: bool = False
|
|
qdrant_error: str | None = None
|
|
qdrant_collections: list[str] = field(default_factory=list)
|
|
ollama_up: bool = False
|
|
ollama_error: str | None = None
|
|
ollama_models: list[str] = field(default_factory=list)
|
|
pg_up: bool = False
|
|
pg_error: str | None = None
|
|
weaviate_up: bool = False
|
|
weaviate_error: str | None = None
|
|
mariadb_up: bool = False
|
|
mariadb_error: str | None = None
|
|
minio_up: bool = False
|
|
minio_error: str | None = None
|
|
|
|
@property
|
|
def vault_ready(self) -> bool:
|
|
"""Vault-note stage: only needs slugify."""
|
|
return self.modules.get("slugify", False)
|
|
|
|
@property
|
|
def qdrant_ready(self) -> bool:
|
|
"""Vector stage: client + embedder libs AND a reachable service."""
|
|
return (
|
|
self.modules.get("fastembed", False)
|
|
and self.modules.get("qdrant_client", False)
|
|
and self.qdrant_up
|
|
)
|
|
|
|
@property
|
|
def graph_ready(self) -> bool:
|
|
"""Graph stage: lightrag installed AND Ollama reachable."""
|
|
return self.modules.get("lightrag", False) and self.ollama_up
|
|
|
|
# service key -> EnvProbe attribute carrying its liveness
|
|
SERVICE_LIVENESS = {"qdrant": "qdrant_up", "postgres": "pg_up",
|
|
"weaviate": "weaviate_up", "mariadb": "mariadb_up"}
|
|
|
|
def vector_ready(self, target: str) -> bool:
|
|
"""Readiness for a specific vector target: modules present,
|
|
plus the live service when the target is service-backed."""
|
|
from .vector_stores import TARGETS
|
|
spec = TARGETS.get(target)
|
|
if spec is None:
|
|
return False
|
|
if not all(self.modules.get(m, False) for m in spec.modules):
|
|
return False
|
|
live = self.SERVICE_LIVENESS.get(spec.service) if spec.service else None
|
|
return live is None or getattr(self, live, False)
|
|
|
|
@property
|
|
def ingest_ready(self) -> bool:
|
|
"""At least one parse path for every declared extension."""
|
|
return (
|
|
self.modules.get("pypdf", False)
|
|
and self.modules.get("bs4", False)
|
|
and self.modules.get("ebooklib", False)
|
|
and self.modules.get("slugify", False)
|
|
)
|
|
|
|
|
|
def _probe_qdrant(host: str, port: int, env: EnvProbe) -> None:
|
|
try:
|
|
from qdrant_client import QdrantClient
|
|
except ImportError:
|
|
env.qdrant_error = "qdrant-client not installed"
|
|
return
|
|
try:
|
|
# Connectivity ping only — no client/server version negotiation.
|
|
client = QdrantClient(url=f"http://{host}:{port}", timeout=3,
|
|
check_compatibility=False)
|
|
env.qdrant_collections = sorted(
|
|
c.name for c in client.get_collections().collections
|
|
)
|
|
env.qdrant_up = True
|
|
except Exception as e: # noqa: BLE001 — any failure means "down"
|
|
env.qdrant_error = str(e).splitlines()[0][:100]
|
|
|
|
|
|
def _probe_ollama(host: str, env: EnvProbe, port: int = 11434) -> None:
|
|
url = f"http://{host}:{port}/api/tags"
|
|
try:
|
|
with urllib.request.urlopen(url, timeout=3) as resp:
|
|
data = json.loads(resp.read().decode("utf-8", errors="ignore"))
|
|
env.ollama_models = sorted(
|
|
m.get("name", "?") for m in data.get("models", [])
|
|
)
|
|
env.ollama_up = True
|
|
except Exception as e: # noqa: BLE001
|
|
env.ollama_error = str(e)[:100]
|
|
|
|
|
|
def _probe_postgres(env: EnvProbe) -> None:
|
|
"""Env-driven like the MariaDB probe: PGHOST/PGPORT/PGUSER/
|
|
PGPASSWORD/PGDATABASE (or PGDSN) — the same variables the
|
|
pgvector target documents."""
|
|
import os
|
|
|
|
try:
|
|
import psycopg
|
|
except ImportError:
|
|
env.pg_error = "psycopg not installed"
|
|
return
|
|
kwargs = {"connect_timeout": 3}
|
|
if os.environ.get("PGDSN"):
|
|
conn_kwargs = {"dsn": os.environ["PGDSN"], **kwargs}
|
|
else:
|
|
conn_kwargs = {**kwargs, "host": os.environ.get("PGHOST", "localhost"),
|
|
"port": int(os.environ.get("PGPORT", "5432")),
|
|
"user": os.environ.get("PGUSER"),
|
|
"password": os.environ.get("PGPASSWORD") or None,
|
|
"dbname": os.environ.get("PGDATABASE")}
|
|
try:
|
|
with psycopg.connect(**conn_kwargs) as conn:
|
|
conn.execute("SELECT 1")
|
|
env.pg_up = True
|
|
except Exception as e: # noqa: BLE001 — any failure means "down"
|
|
env.pg_error = str(e).splitlines()[0][:100]
|
|
|
|
|
|
def _probe_minio(env: EnvProbe) -> None:
|
|
"""Ping the S3-compatible health endpoint — no client library needed."""
|
|
import urllib.request
|
|
|
|
endpoint = os.environ.get("MINIO_ENDPOINT", "localhost:9000")
|
|
host = endpoint if endpoint.startswith("http") else f"http://{endpoint}"
|
|
try:
|
|
with urllib.request.urlopen(f"{host}/minio/health/live", timeout=2):
|
|
env.minio_up = True
|
|
except Exception as e: # noqa: BLE001
|
|
env.minio_error = str(e)[:80]
|
|
|
|
|
|
def _probe_weaviate(env: EnvProbe, ports: tuple[int, ...] = (8080,)) -> None:
|
|
"""Ping /v1/.well-known/ready over HTTP — no client library needed."""
|
|
import urllib.request
|
|
|
|
for port in ports:
|
|
url = f"http://localhost:{port}/v1/.well-known/ready"
|
|
try:
|
|
with urllib.request.urlopen(url, timeout=2) as resp:
|
|
if resp.status == 200:
|
|
env.weaviate_up = True
|
|
return
|
|
except Exception:
|
|
continue
|
|
env.weaviate_error = "no response on ports " + ", ".join(map(str, ports))
|
|
|
|
|
|
def _probe_mariadb(env: EnvProbe) -> None:
|
|
try:
|
|
import pymysql
|
|
except ImportError:
|
|
env.mariadb_error = "pymysql not installed"
|
|
return
|
|
import os
|
|
|
|
try:
|
|
conn = pymysql.connect(
|
|
host=os.environ.get("MARIADB_HOST", "127.0.0.1"),
|
|
port=int(os.environ.get("MARIADB_PORT", "3306")),
|
|
user=os.environ.get("MARIADB_USER", "root"),
|
|
password=os.environ.get("MARIADB_PASSWORD", ""),
|
|
unix_socket=os.environ.get("MARIADB_UNIX_SOCKET") or None,
|
|
connect_timeout=2,
|
|
)
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT 1")
|
|
conn.close()
|
|
env.mariadb_up = True
|
|
except Exception as e: # noqa: BLE001
|
|
env.mariadb_error = str(e).splitlines()[0][:100]
|
|
|
|
|
|
def probe_environment(qdrant_host: str = "localhost",
|
|
qdrant_port: int = 6333,
|
|
ollama_host: str = "localhost") -> EnvProbe:
|
|
"""Gather the full readiness report. Never raises."""
|
|
env = EnvProbe(python_version=platform.python_version())
|
|
env.modules = {name: _module_available(name) for name in PROBED_MODULES}
|
|
_probe_qdrant(qdrant_host, qdrant_port, env)
|
|
_probe_ollama(ollama_host, env)
|
|
_probe_postgres(env)
|
|
_probe_weaviate(env)
|
|
_probe_mariadb(env)
|
|
_probe_minio(env)
|
|
return env
|