""" Thin wrapper around Neo4j that exposes the same queries as test_db.py but as Python methods usable by LangChain tools. """ from __future__ import annotations from neo4j import GraphDatabase from src.config import NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD class GraphClient: _instance: "GraphClient | None" = None def __init__(self): self.driver = self._new_driver() @staticmethod def _new_driver(): return GraphDatabase.driver( NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD), max_connection_lifetime=200, ) @classmethod def get(cls) -> "GraphClient": if cls._instance is None: cls._instance = cls() return cls._instance def _run(self, cypher: str, **params) -> list[dict]: for attempt in range(2): try: with self.driver.session() as s: return [dict(r) for r in s.run(cypher, **params)] except Exception as e: if attempt == 0 and "defunct" in str(e).lower(): self.driver.close() self.driver = self._new_driver() else: raise # ── Concept search ──────────────────────────────────────────────────────── def search_concepts(self, keyword: str) -> list[dict]: rows = self._run(""" CALL db.index.fulltext.queryNodes('concept_search', $kw) YIELD node, score OPTIONAL MATCH (node)-[:DEFINED_IN]->(s:Section) RETURN node.name AS name, node.definition AS definition, node.example AS example, node.page AS page, s.id AS section_id, s.title AS section_title, score ORDER BY score DESC LIMIT 8 """, kw=keyword) if not rows: # fallback: CONTAINS rows = self._run(""" MATCH (c:Concept) WHERE toLower(c.name) CONTAINS toLower($kw) OR toLower(c.definition) CONTAINS toLower($kw) OPTIONAL MATCH (c)-[:DEFINED_IN]->(s:Section) RETURN c.name AS name, c.definition AS definition, c.example AS example, c.page AS page, s.id AS section_id, s.title AS section_title, 0.0 AS score LIMIT 8 """, kw=keyword) return rows # ── Section queries ─────────────────────────────────────────────────────── def chapter_info(self, chapter_id: str) -> dict | None: rows = self._run(""" MATCH (ch:Chapter {id: $id}) OPTIONAL MATCH (pt:Part)-[:CONTAINS]->(ch) OPTIONAL MATCH (ch)-[:CONTAINS]->(s:Section) RETURN ch.id AS chapter_id, ch.title AS chapter_title, pt.id AS part_id, pt.name AS part_name, s.id AS section_id, s.title AS section_title, s.page AS page ORDER BY s.id """, id=chapter_id) if not rows: return None first = rows[0] sections = [ {"section_id": r["section_id"], "title": r["section_title"], "page": r["page"]} for r in rows if r["section_id"] ] return { "chapter_id": first["chapter_id"], "chapter_title": first["chapter_title"], "part_id": first["part_id"], "part_name": first["part_name"], "first_page": sections[0]["page"] if sections else None, "sections": sections, } def section_info(self, section_id: str) -> dict | None: rows = self._run(""" MATCH (s:Section {id: $id}) OPTIONAL MATCH (ch:Chapter)-[:CONTAINS]->(s) OPTIONAL MATCH (pt:Part)-[:CONTAINS]->(ch) RETURN s.id AS section_id, s.title AS title, s.page AS page, ch.id AS chapter_id, ch.title AS chapter_title, pt.id AS part_id, pt.name AS part_name """, id=section_id) return rows[0] if rows else None def section_prereqs(self, section_id: str) -> list[dict]: return self._run(""" MATCH path = (t:Section {id: $id})-[:REQUIRES*1..10]->(p:Section) WITH DISTINCT p, min(length(path)) AS depth RETURN p.id AS section_id, p.title AS title, p.page AS page, depth ORDER BY depth, p.id """, id=section_id) def section_concepts(self, section_id: str) -> list[dict]: """Concepts defined IN a section.""" return self._run(""" MATCH (s:Section {id: $id})-[:DEFINES]->(c:Concept) RETURN c.name AS name, c.definition AS definition, c.example AS example, c.page AS page ORDER BY c.name """, id=section_id) def concepts_needed_for(self, section_id: str) -> list[dict]: """Prereq concepts for a section (via CONCEPT_REQUIRES, deduped, excl. same section).""" return self._run(""" MATCH (s:Section {id: $id})-[:DEFINES]->(c:Concept) MATCH path = (c)-[:CONCEPT_REQUIRES*1..10]->(prereq:Concept) WHERE NOT (s)-[:DEFINES]->(prereq) WITH DISTINCT prereq, min(length(path)) AS depth OPTIONAL MATCH (prereq)-[:DEFINED_IN]->(ps:Section) RETURN prereq.name AS name, prereq.definition AS definition, prereq.example AS example, prereq.page AS page, ps.id AS section_id, depth ORDER BY ps.id, prereq.name """, id=section_id) # ── Concept queries ─────────────────────────────────────────────────────── def concept_prereqs(self, concept_name: str) -> list[dict]: return self._run(""" MATCH (t:Concept) WHERE t.name = $name OR (t)-[:ALSO_KNOWN_AS]->(:Concept {name: $name}) WITH t MATCH path = (t)-[:CONCEPT_REQUIRES*1..10]->(p:Concept) WITH DISTINCT p, min(length(path)) AS depth OPTIONAL MATCH (p)-[:DEFINED_IN]->(s:Section) RETURN p.name AS name, p.definition AS definition, p.example AS example, p.page AS page, s.id AS section_id, s.title AS section_title, depth ORDER BY depth, s.id, p.name """, name=concept_name) def concept_subtypes(self, concept_name: str) -> list[dict]: return self._run(""" MATCH (parent:Concept) WHERE parent.name = $name OR (parent)-[:ALSO_KNOWN_AS]->(:Concept {name: $name}) WITH parent MATCH (child:Concept)-[:IS_TYPE_OF]->(parent) RETURN child.name AS name, child.definition AS definition, child.example AS example ORDER BY child.name """, name=concept_name)