90 lines
4.3 KiB
Python
90 lines
4.3 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.index_state import IndexUnavailableError, index_snapshot, sqlite_version
|
|
from tools.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
|