dp-zp-agent/tools/runtime/tool_health.py

109 lines
4.4 KiB
Python

"""Independent, sanitized readiness results for the two retrieval backends."""
from __future__ import annotations
import json
import sqlite3
from contextlib import closing
from pathlib import Path
from neo4j import Query
from tools.graphrag.graph_db import create_neo4j_driver, get_neo4j_settings
from tools.graphrag.graphrag import GRAPH_SCOPE
from tools.sync.index_state import IndexUnavailableError, index_snapshot, sqlite_version
from tools.retrieval.publication import PUBLICATION_POLICY_VERSION
from tools.rag.search_core import verify_search_schema, embedding_index_info
PROBE_TIMEOUT_SECONDS = 2.0
def readiness(ready: bool, reason: str | None = None, **details) -> dict:
return {
"ready": ready,
"status": "ready" if ready else "unavailable",
"reason": reason,
**details,
}
def rag_readiness(db_file: Path, version: dict | None, *, embedding_ready: bool) -> dict:
if not db_file.is_file():
return readiness(False, "index_missing")
try:
if version is not None and sqlite_version(db_file) != version:
return readiness(False, "version_mismatch")
with closing(
sqlite3.connect(db_file.resolve().as_uri() + "?mode=ro", uri=True, timeout=1)
) as conn:
conn.row_factory = sqlite3.Row
verify_search_schema(conn)
model, dimensions = embedding_index_info(conn)
conn.execute("SELECT rowid FROM chunks_fts LIMIT 1").fetchone()
if not embedding_ready:
return readiness(False, "embedding_not_ready", embedding_model=model)
return readiness(True, embedding_model=model, dimensions=dimensions, index_version=version)
except (sqlite3.Error, RuntimeError, ValueError, OSError):
return readiness(False, "index_invalid")
def neo4j_readiness(version: dict | None) -> tuple[dict, dict]:
try:
settings = get_neo4j_settings()
except RuntimeError:
return readiness(False, "not_configured"), readiness(False, "neo4j_not_configured")
try:
with create_neo4j_driver(
settings,
connection_timeout=PROBE_TIMEOUT_SECONDS,
connection_acquisition_timeout=PROBE_TIMEOUT_SECONDS,
connection_write_timeout=PROBE_TIMEOUT_SECONDS,
max_transaction_retry_time=0,
max_connection_pool_size=1,
) as driver:
driver.verify_connectivity()
with driver.session(database=settings.database) as session:
record = session.run(
Query(
"""MATCH (n) WHERE n.graph_scope=$scope
WITH count(n) AS nodes
OPTIONAL MATCH (state:IndexState {id:$scope})
RETURN nodes, state.metadata_json AS metadata""",
timeout=PROBE_TIMEOUT_SECONDS,
),
scope=GRAPH_SCOPE,
).single()
except Exception:
# Credentials, connection strings and raw driver errors stay server-side.
return readiness(False, "unreachable"), readiness(False, "neo4j_unavailable")
neo4j = readiness(True)
if not record or record["nodes"] == 0:
return neo4j, readiness(False, "graph_empty")
try:
actual = json.loads(record["metadata"]) if record["metadata"] else None
except (ValueError, TypeError):
return neo4j, readiness(False, "graph_version_invalid")
if version is not None and actual != version:
return neo4j, readiness(False, "version_mismatch")
if version is None or version.get("publication_policy") != PUBLICATION_POLICY_VERSION:
return neo4j, readiness(False, "publication_sync_required")
return neo4j, readiness(True, index_version=actual, nodes=record["nodes"])
def tools_readiness(
db_file: Path, *, security_configured: bool, embedding_ready: bool
) -> tuple[dict, dict | None]:
version = None
try:
with index_snapshot(db_file.parent) as version:
rag = rag_readiness(db_file, version, embedding_ready=embedding_ready)
neo4j, graph = neo4j_readiness(version)
except IndexUnavailableError:
neo4j, _ = neo4j_readiness(None)
rag = readiness(False, "snapshot_unavailable")
graph = readiness(False, "snapshot_unavailable")
if not security_configured:
rag = readiness(False, "api_not_configured")
graph = readiness(False, "api_not_configured")
return {"rag": rag, "neo4j": neo4j, "graphrag": graph}, version