Spaces:
Running
Running
| """Quick connectivity checks for Chroma and Neo4j data sources.""" | |
| import sys | |
| from pathlib import Path | |
| from typing import Optional | |
| from dotenv import load_dotenv | |
| try: | |
| from src.database.chroma_manager import get_vector_store, query_segments | |
| from src.database.neo4j_manager import get_graph | |
| from src.config import CHROMA_DIR, SEGMENTED_DIR | |
| except Exception as import_err: | |
| print(f"Import error: {import_err}") | |
| sys.exit(1) | |
| def check_chroma(sample_query: str = "alpha school") -> bool: | |
| """Validate Chroma is reachable and returns results for a sample query.""" | |
| try: | |
| vector_store = get_vector_store() | |
| results = query_segments(sample_query, k=1) | |
| print(f"Chroma OK: returned {len(results)} results for '{sample_query}'.") | |
| # Confirm the collection is readable | |
| try: | |
| stats: Optional[dict] = getattr(vector_store, "_collection", None) | |
| if stats and hasattr(stats, "count"): | |
| print(f"Chroma collection count: {stats.count()}") | |
| except Exception: | |
| # Collection stats are best-effort; ignore if unsupported | |
| pass | |
| return True | |
| except Exception as exc: # pragma: no cover - best-effort runtime check | |
| print(f"Chroma check failed: {exc}") | |
| return False | |
| def check_neo4j() -> bool: | |
| """Validate Neo4j is reachable and can execute a trivial query.""" | |
| try: | |
| graph = get_graph() | |
| res = graph.query("MATCH (n) RETURN count(n) AS cnt LIMIT 1") | |
| count = res[0].get("cnt", 0) if res else 0 | |
| print(f"Neo4j OK: node count sample={count}") | |
| return True | |
| except Exception as exc: # pragma: no cover - best-effort runtime check | |
| print(f"Neo4j check failed: {exc}") | |
| return False | |
| def check_filesystem() -> bool: | |
| """Ensure expected data directories exist and are readable.""" | |
| ok = True | |
| if not CHROMA_DIR.exists(): | |
| print(f"Chroma dir missing: {CHROMA_DIR}") | |
| ok = False | |
| else: | |
| print(f"Chroma dir present: {CHROMA_DIR}") | |
| if not SEGMENTED_DIR.exists(): | |
| print(f"Segmented transcripts dir missing: {SEGMENTED_DIR}") | |
| ok = False | |
| else: | |
| sample_files = list(Path(SEGMENTED_DIR).glob("*.json"))[:3] | |
| print(f"Segmented transcripts present: {len(sample_files)} sample(s) {sample_files}") | |
| return ok | |
| def main() -> int: | |
| """Run connectivity checks; exit 0 on success, 1 on any failure.""" | |
| load_dotenv() | |
| fs_ok = check_filesystem() | |
| chroma_ok = check_chroma() | |
| neo4j_ok = check_neo4j() | |
| return 0 if fs_ok and chroma_ok and neo4j_ok else 1 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |