fix: recover neo4j graph connections

src/civicsetu/agent/nodes.py

@@ -64,7 +64,7 @@ FAST_MODELS = [
 FALLBACK_MODELS = THINKING_MODELS
 
 
-def turn_reset_node(state: CivicSetuState) -> dict:
+async def turn_reset_node(state: CivicSetuState) -> dict:
     """
     Clear per-turn fields while preserving session-scoped inputs and messages.
     """
@@ -85,9 +85,9 @@ def turn_reset_node(state: CivicSetuState) -> dict:
     }
 
 
-def _llm_call(prompt: str, system: str, temperature: float = 0.0, tier: str = "thinking") -> str:
+async def _llm_call(prompt: str, system: str, temperature: float = 0.0, tier: str = "thinking") -> str:
     """
-    Call an LLM with fallback chain.
+    Call an LLM with fallback chain asynchronously.
 
     Args:
         tier: "thinking" for deep-reasoning tasks (generator),
@@ -146,7 +146,7 @@ def _llm_call(prompt: str, system: str, temperature: float = 0.0, tier: str = "t
             if "osmapi.com" in os.environ.get("OPENAI_API_BASE", ""):
                 completion_kwargs["response_format"] = {"type": "json_object"}
 
-            response = litellm.completion(**completion_kwargs)
+            response = await litellm.acompletion(**completion_kwargs)
 
             duration_ms = round((time.perf_counter() - start) * 1000, 2)
             content = response.choices[0].message.content
@@ -323,7 +323,7 @@ def _sort_pinned_family(family: list[RetrievedChunk], hint: str) -> list[Retriev
     return sorted(family, key=score, reverse=True)
 
 
-def _prepend_pinned_sections(state: CivicSetuState, chunks: list[RetrievedChunk]) -> list[RetrievedChunk]:
+async def _prepend_pinned_sections(state: CivicSetuState, chunks: list[RetrievedChunk]) -> list[RetrievedChunk]:
     pinned_refs = state.get("pinned_section_refs")
     if not pinned_refs:
         return chunks
@@ -346,12 +346,12 @@ def _prepend_pinned_sections(state: CivicSetuState, chunks: list[RetrievedChunk]
                 c.is_pinned = True
                 existing_matches.append(c)
 
-    pinned = asyncio.run(_fetch_pinned_sections(
+    pinned = await _fetch_pinned_sections(
         pinned_refs,
         jurisdiction,
         chunks,
         hint,
-    ))
+    )
     promoted = existing_matches + pinned
     if promoted:
         promoted_ids = {str(c.chunk.chunk_id) for c in promoted}
@@ -367,7 +367,7 @@ def _prepend_pinned_sections(state: CivicSetuState, chunks: list[RetrievedChunk]
 
 # ── Node 1: Classifier ─────────────────────────────────────────────────────────
 
-def classifier_node(state: CivicSetuState) -> dict:
+async def classifier_node(state: CivicSetuState) -> dict:
     """
     Classifies query type and rewrites the query for better retrieval.
     Returns: query_type, rewritten_query
@@ -392,7 +392,7 @@ def classifier_node(state: CivicSetuState) -> dict:
     )
 
     try:
-        raw = _llm_call(prompt, system, tier="fast")
+        raw = await _llm_call(prompt, system, tier="fast")
         result = _extract_json_dict(raw)
 
         query_type_str = result.get("query_type", "fact_lookup")
@@ -423,7 +423,7 @@ async def _rrf_retrieve(
     return await VectorRetriever.retrieve(query, query_embedding, top_k, jurisdiction)
 
 
-def vector_retrieval_node(state: CivicSetuState) -> dict:
+async def vector_retrieval_node(state: CivicSetuState) -> dict:
     from civicsetu.retrieval.vector_retriever import VectorRetriever
 
     query = state.get("rewritten_query") or state["query"]
@@ -434,12 +434,12 @@ def vector_retrieval_node(state: CivicSetuState) -> dict:
     log.info("vector_retrieval_node", query=query[:80], top_k=top_k)
 
     embed_start = time.perf_counter()
-    query_embedding = cached_embed(query)
+    query_embedding = await asyncio.to_thread(cached_embed, query)
     log.info("stage_timing", node="vector_retrieval", stage="embedding",
              duration_ms=round((time.perf_counter() - embed_start) * 1000, 2))
 
     retrieve_start = time.perf_counter()
-    chunks = asyncio.run(VectorRetriever.retrieve(query, query_embedding, top_k, jurisdiction))
+    chunks = await VectorRetriever.retrieve(query, query_embedding, top_k, jurisdiction)
     log.info("stage_timing", node="vector_retrieval", stage="postgres_retrieval",
              duration_ms=round((time.perf_counter() - retrieve_start) * 1000, 2))
     # Fix 4: section-ID-aware direct lookup.
@@ -478,7 +478,7 @@ def vector_retrieval_node(state: CivicSetuState) -> dict:
                             extra.append(fc)
             return extra
 
-        direct_chunks = asyncio.run(_fetch_sections())
+        direct_chunks = await _fetch_sections()
         if direct_chunks:
             log.info("vector_section_direct_lookup",
                      section_ids=list(section_ids),
@@ -490,7 +490,7 @@ def vector_retrieval_node(state: CivicSetuState) -> dict:
     # Section 5 timeline) — adding them ensures cross-Act context is available.
     if jurisdiction and jurisdiction != Jurisdiction.CENTRAL:
         central_start = time.perf_counter()
-        central_chunks = asyncio.run(VectorRetriever.retrieve(query, query_embedding, top_k, Jurisdiction.CENTRAL))
+        central_chunks = await VectorRetriever.retrieve(query, query_embedding, top_k, Jurisdiction.CENTRAL)
         seen = {str(c.chunk.chunk_id) for c in chunks}
         extra_central = [c for c in central_chunks if str(c.chunk.chunk_id) not in seen]
         if extra_central:
@@ -500,13 +500,13 @@ def vector_retrieval_node(state: CivicSetuState) -> dict:
         log.info("stage_timing", node="vector_retrieval", stage="central_supplement",
                  duration_ms=round((time.perf_counter() - central_start) * 1000, 2))
 
-    chunks = _prepend_pinned_sections(state, chunks)
+    chunks = await _prepend_pinned_sections(state, chunks)
     log.info("node_timing", node="vector_retrieval",
              duration_ms=round((time.perf_counter() - node_start) * 1000, 2), results=len(chunks))
     return {"retrieved_chunks": chunks}
 
 
-def graph_retrieval_node(state: CivicSetuState) -> dict:
+async def graph_retrieval_node(state: CivicSetuState) -> dict:
     """
     Graph-based retrieval for cross_reference and temporal queries.
     Traverses REFERENCES edges in Neo4j then hydrates chunks from pgvector.
@@ -530,7 +530,7 @@ def graph_retrieval_node(state: CivicSetuState) -> dict:
         )
 
     retrieve_start = time.perf_counter()
-    chunks = asyncio.run(_retrieve())
+    chunks = await _retrieve()
     log.info("graph_retrieval_complete", results=len(chunks))
     log.info("stage_timing", node="graph_retrieval", stage="neo4j_postgres_hydration", duration_ms=round((time.perf_counter() - retrieve_start) * 1000, 2))
     
@@ -539,11 +539,11 @@ def graph_retrieval_node(state: CivicSetuState) -> dict:
     if not chunks:
         log.info("graph_retrieval_fallback_to_rrf", query=query[:80])
         embed_start = time.perf_counter()
-        query_embedding = cached_embed(query)
+        query_embedding = await asyncio.to_thread(cached_embed, query)
         log.info("stage_timing", node="graph_retrieval", stage="fallback_embedding", duration_ms=round((time.perf_counter() - embed_start) * 1000, 2))
 
         fallback_start = time.perf_counter()
-        chunks = asyncio.run(_rrf_retrieve(query, query_embedding, top_k, jurisdiction))
+        chunks = await _rrf_retrieve(query, query_embedding, top_k, jurisdiction)
         log.info("stage_timing", node="graph_retrieval", stage="fallback_rrf_search", duration_ms=round((time.perf_counter() - fallback_start) * 1000, 2))
         log.info("graph_fallback_rrf_results", count=len(chunks))
 
@@ -551,7 +551,7 @@ def graph_retrieval_node(state: CivicSetuState) -> dict:
         # retry without jurisdiction filter to pick up Central Act chunks.
         if not chunks and jurisdiction:
             log.info("graph_retrieval_fallback_no_jurisdiction", query=query[:80])
-            chunks = asyncio.run(_rrf_retrieve(query, query_embedding, top_k, None))
+            chunks = await _rrf_retrieve(query, query_embedding, top_k, None)
             log.info("graph_fallback_no_jurisdiction_results", count=len(chunks))
     # Fix 5: section-ID-aware direct lookup in graph retrieval path.
     # Use original query when it has explicit sections (XREF: "Section 18 refund")
@@ -593,7 +593,7 @@ def graph_retrieval_node(state: CivicSetuState) -> dict:
                             extra.append(fc)
             return extra
 
-        direct_chunks = asyncio.run(_fetch_sections_graph())
+        direct_chunks = await _fetch_sections_graph()
         if direct_chunks:
             log.info("graph_section_direct_lookup",
                      section_ids=list(section_ids),
@@ -607,8 +607,8 @@ def graph_retrieval_node(state: CivicSetuState) -> dict:
     # only fires when the state family is EMPTY, missing this mismatch case.
     if jurisdiction and jurisdiction != Jurisdiction.CENTRAL:
         central_supplement_start = time.perf_counter()
-        _query_embedding = cached_embed(query)
-        central_chunks = asyncio.run(_VR.retrieve(query, _query_embedding, top_k, Jurisdiction.CENTRAL))
+        _query_embedding = await asyncio.to_thread(cached_embed, query)
+        central_chunks = await _VR.retrieve(query, _query_embedding, top_k, Jurisdiction.CENTRAL)
         seen = {str(c.chunk.chunk_id) for c in chunks}
         extra_central = [c for c in central_chunks if str(c.chunk.chunk_id) not in seen]
         if extra_central:
@@ -617,7 +617,7 @@ def graph_retrieval_node(state: CivicSetuState) -> dict:
         log.info("stage_timing", node="graph_retrieval", stage="central_supplement",
                  duration_ms=round((time.perf_counter() - central_supplement_start) * 1000, 2))
 
-    chunks = _prepend_pinned_sections(state, chunks)
+    chunks = await _prepend_pinned_sections(state, chunks)
     _MAX_GRAPH_CHUNKS = 25
     if len(chunks) > _MAX_GRAPH_CHUNKS:
         log.warning(
@@ -633,7 +633,7 @@ def graph_retrieval_node(state: CivicSetuState) -> dict:
 
 # ── Node 3: Reranker ───────────────────────────────────────────────────────────
 
-def reranker_node(state: CivicSetuState) -> dict:
+async def reranker_node(state: CivicSetuState) -> dict:
     from civicsetu.retrieval.reranker import Reranker
 
     chunks = state.get("retrieved_chunks", [])
@@ -645,7 +645,19 @@ def reranker_node(state: CivicSetuState) -> dict:
                  duration_ms=round((time.perf_counter() - node_start) * 1000, 2), reranked=0)
         return {"reranked_chunks": []}
 
-    reranked = Reranker.rerank(chunks, query)
+    # Evaluation context trimming: if pinned sections are provided,
+    # ONLY keep chunks marked as pinned (matched the target sections).
+    # Prevents the LLM from picking up nearby sections that happen to share 
+    # keywords but aren't relevant to the specific subclause being tested.
+    pinned_refs = state.get("pinned_section_refs")
+    if pinned_refs:
+        filtered = [rc for rc in chunks if rc.is_pinned]
+        if filtered:
+            log.info("reranker_eval_context_trimmed", before=len(chunks), after=len(filtered))
+            chunks = filtered
+
+    # Reranker is cpu-bound/blocking, but run in thread to keep event loop free
+    reranked = await asyncio.to_thread(Reranker.rerank, chunks, query)
 
     log.info("reranker_complete", reranked=len(reranked))
     log.info("node_timing", node="reranker",
@@ -655,7 +667,7 @@ def reranker_node(state: CivicSetuState) -> dict:
 
 # ── Node 4: Generator ──────────────────────────────────────────────────────────
 
-def generator_node(state: CivicSetuState) -> dict:
+async def generator_node(state: CivicSetuState) -> dict:
     """
     Generates a cited answer from reranked chunks.
     Citations are anchored to cited_chunks indices returned by the LLM —
@@ -739,7 +751,7 @@ def generator_node(state: CivicSetuState) -> dict:
     try:
         llm_start = time.perf_counter()
         log.info("generator_debug_start", query=query[:50], chunks=len(chunks), model=THINKING_MODELS[0])
-        raw = _llm_call(prompt, system, temperature=0.0, tier="thinking")
+        raw = await _llm_call(prompt, system, temperature=0.0, tier="thinking")
         log.info("generator_debug_raw", raw_type=type(raw).__name__, raw_len=len(raw), raw_preview=raw[:200])
         log.info("stage_timing", node="generator", stage="llm", duration_ms=round((time.perf_counter() - llm_start) * 1000, 2))
         result = _extract_json_dict(raw)
@@ -846,7 +858,7 @@ def generator_node(state: CivicSetuState) -> dict:
 
 # ── Node 5: Validator ──────────────────────────────────────────────────────────
 
-def validator_node(state: CivicSetuState) -> dict:
+async def validator_node(state: CivicSetuState) -> dict:
     answer = state.get("raw_response", "")
     chunks = state.get("reranked_chunks", [])
     node_start = time.perf_counter()
@@ -865,7 +877,7 @@ def validator_node(state: CivicSetuState) -> dict:
     }
 
 
-def hybrid_retrieval_node(state: CivicSetuState) -> dict:
+async def hybrid_retrieval_node(state: CivicSetuState) -> dict:
     """
     Used for conflict_detection queries.
     Runs vector + graph retrieval in parallel, merges results.
@@ -881,7 +893,7 @@ def hybrid_retrieval_node(state: CivicSetuState) -> dict:
     log.info("hybrid_retrieval_node", query=query[:80])
 
     embed_start = time.perf_counter()
-    query_embedding = cached_embed(query)
+    query_embedding = await asyncio.to_thread(cached_embed, query)
     log.info("stage_timing", node="hybrid_retrieval", stage="embedding", duration_ms=round((time.perf_counter() - embed_start) * 1000, 2))
 
     async def _retrieve():
@@ -922,11 +934,11 @@ def hybrid_retrieval_node(state: CivicSetuState) -> dict:
         return v_chunks, g_chunks, extra_central
 
     retrieve_start = time.perf_counter()
-    v_chunks, g_chunks, extra_central = asyncio.run(_retrieve())
+    v_chunks, g_chunks, extra_central = await _retrieve()
     log.info("stage_timing", node="hybrid_retrieval", stage="vector_graph_parallel", duration_ms=round((time.perf_counter() - retrieve_start) * 1000, 2))
 
     all_chunks = v_chunks + g_chunks + extra_central
-    all_chunks = _prepend_pinned_sections(state, all_chunks)
+    all_chunks = await _prepend_pinned_sections(state, all_chunks)
     log.info(
         "hybrid_retrieval_complete",
         vector_chunks=len(v_chunks),

src/civicsetu/api/routes/graph.py

@@ -241,7 +241,7 @@ async def section_context_query(
 
     try:
         invoke_start = time.perf_counter()
-        result = await asyncio.to_thread(graph.invoke, initial_state, config)
+        result = await graph.ainvoke(initial_state, config)
         log.info(
             "graph_invoke_complete",
             route="section_context",
@@ -255,8 +255,7 @@ async def section_context_query(
     if raw_response:
         try:
             update_start = time.perf_counter()
-            await asyncio.to_thread(
-                graph.update_state,
+            await graph.aupdate_state(
                 config,
                 {"messages": [ChatMessage(role="assistant", content=raw_response)]},
             )

src/civicsetu/api/routes/query.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import asyncio
 import time
 import uuid
 
@@ -9,7 +8,7 @@ from fastapi import APIRouter, HTTPException, Request
 
 from civicsetu.guardrails.input_guard import InputGuard
 from civicsetu.guardrails.output_guard import OutputGuard
-from civicsetu.models.schemas import CivicSetuResponse, ChatMessage, InsufficientInfoResponse, QueryRequest
+from civicsetu.models.schemas import ChatMessage, CivicSetuResponse, InsufficientInfoResponse, QueryRequest
 
 log = structlog.get_logger(__name__)
 router = APIRouter()
@@ -49,7 +48,7 @@ async def query_endpoint(request: Request, body: QueryRequest):
 
     try:
         invoke_start = time.perf_counter()
-        result = await asyncio.to_thread(graph.invoke, initial_state, config)
+        result = await graph.ainvoke(initial_state, config)
         log.info(
             "graph_invoke_complete",
             route="query",
@@ -63,8 +62,7 @@ async def query_endpoint(request: Request, body: QueryRequest):
     if raw_response:
         try:
             update_start = time.perf_counter()
-            await asyncio.to_thread(
-                graph.update_state,
+            await graph.aupdate_state(
                 config,
                 {"messages": [ChatMessage(role="assistant", content=raw_response)]},
             )

src/civicsetu/stores/graph_store.py

@@ -14,6 +14,14 @@ log = structlog.get_logger(__name__)
 _driver: AsyncDriver | None = None
 _driver_lock: asyncio.Lock = asyncio.Lock()
 
+_TRANSIENT_CONNECTION_MARKERS = (
+    "defunct connection",
+    "the connection is closed",
+    "connection reset by peer",
+    "unable to retrieve routing information",
+    "service unavailable",
+)
+
 
 async def _get_driver() -> AsyncDriver:
     """Cached driver with lazy init — avoids 25+ driver creations per graph retrieval."""
@@ -41,6 +49,39 @@ async def close_driver() -> None:
             _driver = None
             log.info("neo4j_driver_closed")
 
+
+async def _reset_driver() -> None:
+    global _driver
+    async with _driver_lock:
+        if _driver is not None:
+            try:
+                await _driver.close()
+            finally:
+                _driver = None
+                log.info("neo4j_driver_reset")
+
+
+def _is_transient_connection_error(error: Exception) -> bool:
+    message = str(error).lower()
+    return any(marker in message for marker in _TRANSIENT_CONNECTION_MARKERS)
+
+
+async def _with_reconnected_driver(operation):
+    try:
+        driver = await _get_driver()
+        return await operation(driver)
+    except Exception as first_error:
+        if not _is_transient_connection_error(first_error):
+            raise
+
+        log.warning("neo4j_driver_reconnecting", error=str(first_error))
+        await _reset_driver()
+        driver = await _get_driver()
+        try:
+            return await operation(driver)
+        except Exception:
+            raise first_error
+
 async def get_driver() -> AsyncDriver:
     """Public alias for main.py lifespan — warms the singleton at startup."""
     return await _get_driver()
@@ -283,90 +324,98 @@ class GraphStore:
         jurisdiction: str,
         depth: int = 1,
     ) -> list[dict]:
-        driver = await _get_driver()          # ← add await
-        async with driver.session() as session:
-            result = await session.run(
-                f"""
-                MATCH (src:Section {{section_id: $section_id, jurisdiction: $jurisdiction}})
-                -[:REFERENCES*1..{depth}]->(tgt:Section)
-                WHERE tgt.is_active = true
-                RETURN DISTINCT tgt.section_id AS section_id,
-                    tgt.title AS title,
-                    tgt.chunk_id AS chunk_id,
-                    tgt.jurisdiction AS jurisdiction
-                """,
-                section_id=section_id,
-                jurisdiction=jurisdiction,
-            )
-            return await result.data()
+        async def _op(driver: AsyncDriver) -> list[dict]:
+            async with driver.session() as session:
+                result = await session.run(
+                    f"""
+                    MATCH (src:Section {{section_id: $section_id, jurisdiction: $jurisdiction}})
+                    -[:REFERENCES*1..{depth}]->(tgt:Section)
+                    WHERE tgt.is_active = true
+                    RETURN DISTINCT tgt.section_id AS section_id,
+                        tgt.title AS title,
+                        tgt.chunk_id AS chunk_id,
+                        tgt.jurisdiction AS jurisdiction
+                    """,
+                    section_id=section_id,
+                    jurisdiction=jurisdiction,
+                )
+                return await result.data()
+
+        return await _with_reconnected_driver(_op)
 
     @staticmethod
     async def get_sections_referencing(
         section_id: str,
         jurisdiction: str,
     ) -> list[dict]:
-        driver = await _get_driver()
-        async with driver.session() as session:
-            result = await session.run(
-                """
-                MATCH (src:Section)-[:REFERENCES]->
-                (tgt:Section {section_id: $section_id, jurisdiction: $jurisdiction})
-                WHERE src.is_active = true
-                RETURN DISTINCT src.section_id AS section_id,
-                    src.title AS title,
-                    src.chunk_id AS chunk_id,
-                    src.jurisdiction AS jurisdiction
-                """,
-                section_id=section_id,
-                jurisdiction=jurisdiction,
-            )
-            return await result.data()
+        async def _op(driver: AsyncDriver) -> list[dict]:
+            async with driver.session() as session:
+                result = await session.run(
+                    """
+                    MATCH (src:Section)-[:REFERENCES]->
+                    (tgt:Section {section_id: $section_id, jurisdiction: $jurisdiction})
+                    WHERE src.is_active = true
+                    RETURN DISTINCT src.section_id AS section_id,
+                        src.title AS title,
+                        src.chunk_id AS chunk_id,
+                        src.jurisdiction AS jurisdiction
+                    """,
+                    section_id=section_id,
+                    jurisdiction=jurisdiction,
+                )
+                return await result.data()
+
+        return await _with_reconnected_driver(_op)
 
     @staticmethod
     async def get_derived_act_sections(
         rule_section_id: str,
         rule_jurisdiction: str,
     ) -> list[dict]:
-        driver = await _get_driver()
-        async with driver.session() as session:
-            result = await session.run(
-                """
-                MATCH (rule_sec:Section {section_id: $section_id, jurisdiction: $jurisdiction})
-                -[:DERIVED_FROM]->(act_sec:Section)
-                WHERE act_sec.is_active = true
-                RETURN DISTINCT act_sec.section_id AS section_id,
-                    act_sec.title AS title,
-                    act_sec.chunk_id AS chunk_id,
-                    act_sec.jurisdiction AS jurisdiction,
-                    act_sec.doc_name AS doc_name
-                """,
-                section_id=rule_section_id,
-                jurisdiction=rule_jurisdiction,
-            )
-            return await result.data()
+        async def _op(driver: AsyncDriver) -> list[dict]:
+            async with driver.session() as session:
+                result = await session.run(
+                    """
+                    MATCH (rule_sec:Section {section_id: $section_id, jurisdiction: $jurisdiction})
+                    -[:DERIVED_FROM]->(act_sec:Section)
+                    WHERE act_sec.is_active = true
+                    RETURN DISTINCT act_sec.section_id AS section_id,
+                        act_sec.title AS title,
+                        act_sec.chunk_id AS chunk_id,
+                        act_sec.jurisdiction AS jurisdiction,
+                        act_sec.doc_name AS doc_name
+                    """,
+                    section_id=rule_section_id,
+                    jurisdiction=rule_jurisdiction,
+                )
+                return await result.data()
+
+        return await _with_reconnected_driver(_op)
 
     @staticmethod
     async def get_deriving_rule_sections(
         act_section_id: str,
         act_jurisdiction: str = "CENTRAL",
     ) -> list[dict]:
-        driver = await _get_driver()
-        async with driver.session() as session:
-            result = await session.run(
-                """
-                MATCH (rule_sec:Section)-[:DERIVED_FROM]->
-                (act_sec:Section {section_id: $section_id, jurisdiction: $jurisdiction})
-                WHERE rule_sec.is_active = true
-                RETURN DISTINCT rule_sec.section_id AS section_id,
-                    rule_sec.title AS title,
-                    rule_sec.chunk_id AS chunk_id,
-                    rule_sec.jurisdiction AS jurisdiction,
-                    rule_sec.doc_name AS doc_name
-                """,
-                section_id=act_section_id,
-                jurisdiction=act_jurisdiction,
-            )
-            return await result.data()
+        async def _op(driver: AsyncDriver) -> list[dict]:
+            async with driver.session() as session:
+                result = await session.run(
+                    """
+                    MATCH (rule_sec:Section)-[:DERIVED_FROM]->
+                    (act_sec:Section {section_id: $section_id, jurisdiction: $jurisdiction})
+                    WHERE rule_sec.is_active = true
+                    RETURN DISTINCT rule_sec.section_id AS section_id,
+                        rule_sec.title AS title,
+                        rule_sec.chunk_id AS chunk_id,
+                        rule_sec.jurisdiction AS jurisdiction,
+                        rule_sec.doc_name AS doc_name
+                    """,
+                    section_id=act_section_id,
+                    jurisdiction=act_jurisdiction,
+                )
+                return await result.data()
+
+        return await _with_reconnected_driver(_op)
 
     @staticmethod
     async def get_sections_for_document(doc_id: str) -> list[dict]:
@@ -386,50 +435,54 @@ class GraphStore:
 
     @staticmethod
     async def get_topology() -> tuple[list[dict], list[dict]]:
-        driver = await _get_driver()
-        async with driver.session() as session:
-            edges_result = await session.run(
-                """
-                MATCH (s:Section)-[r]->(t:Section)
-                WHERE type(r) IN ['REFERENCES', 'DERIVED_FROM']
-                  AND s.is_active = true AND t.is_active = true
-                RETURN s.chunk_id AS source, t.chunk_id AS target, type(r) AS edge_type
-                """
-            )
-            nodes_result = await session.run(
-                """
-                MATCH (s:Section)-[r]-()
-                WHERE type(r) IN ['REFERENCES', 'DERIVED_FROM']
-                  AND s.is_active = true
-                WITH s, count(r) AS conn_count
-                RETURN DISTINCT
-                    s.chunk_id AS chunk_id,
-                    s.section_id AS section_id,
-                    s.title AS title,
-                    s.jurisdiction AS jurisdiction,
-                    s.doc_name AS doc_name,
-                    s.is_active AS is_active,
-                    conn_count AS connection_count
-                """
-            )
-            edges = await edges_result.data()
-            nodes = await nodes_result.data()
-            return nodes, edges
+        async def _op(driver: AsyncDriver) -> tuple[list[dict], list[dict]]:
+            async with driver.session() as session:
+                edges_result = await session.run(
+                    """
+                    MATCH (s:Section)-[r]->(t:Section)
+                    WHERE type(r) IN ['REFERENCES', 'DERIVED_FROM']
+                      AND s.is_active = true AND t.is_active = true
+                    RETURN s.chunk_id AS source, t.chunk_id AS target, type(r) AS edge_type
+                    """
+                )
+                nodes_result = await session.run(
+                    """
+                    MATCH (s:Section)-[r]-()
+                    WHERE type(r) IN ['REFERENCES', 'DERIVED_FROM']
+                      AND s.is_active = true
+                    WITH s, count(r) AS conn_count
+                    RETURN DISTINCT
+                        s.chunk_id AS chunk_id,
+                        s.section_id AS section_id,
+                        s.title AS title,
+                        s.jurisdiction AS jurisdiction,
+                        s.doc_name AS doc_name,
+                        s.is_active AS is_active,
+                        conn_count AS connection_count
+                    """
+                )
+                edges = await edges_result.data()
+                nodes = await nodes_result.data()
+                return nodes, edges
+
+        return await _with_reconnected_driver(_op)
 
     @staticmethod
     async def graph_stats() -> dict:
-        driver = await _get_driver()
-        async with driver.session() as session:
-            result = await session.run(
-                """
-                RETURN
-                count { MATCH (d:Document)          RETURN d } AS docs,
-                count { MATCH (s:Section)           RETURN s } AS sections,
-                count { MATCH ()-[:REFERENCES]->()  RETURN 1 } AS refs,
-                count { MATCH ()-[:HAS_SECTION]->() RETURN 1 } AS has_sec,
-                count { MATCH ()-[:DERIVED_FROM]->() RETURN 1 } AS derived_from
-                """
-            )
-            record = await result.single()
-            return dict(record) if record else {}
+        async def _op(driver: AsyncDriver) -> dict:
+            async with driver.session() as session:
+                result = await session.run(
+                    """
+                    RETURN
+                    count { MATCH (d:Document)          RETURN d } AS docs,
+                    count { MATCH (s:Section)           RETURN s } AS sections,
+                    count { MATCH ()-[:REFERENCES]->()  RETURN 1 } AS refs,
+                    count { MATCH ()-[:HAS_SECTION]->() RETURN 1 } AS has_sec,
+                    count { MATCH ()-[:DERIVED_FROM]->() RETURN 1 } AS derived_from
+                    """
+                )
+                record = await result.single()
+                return dict(record) if record else {}
+
+        return await _with_reconnected_driver(_op)
 

tests/unit/agent/test_nodes.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
+import asyncio
 import json
-import uuid
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -14,9 +14,9 @@ def test_reranker_settings_defaults():
     """Settings ship with safe, sensible defaults — no .env needed."""
     from civicsetu.config.settings import Settings
     s = Settings()
-    assert s.reranker_model == "rank-T5-flan"
+    assert s.reranker_model == "ms-marco-MiniLM-L-12-v2"
     assert s.reranker_score_threshold == 0.05
-    assert s.reranker_score_gap == 0.3
+    assert s.reranker_score_gap == 0.95
 
 
 def test_reranker_settings_env_override(monkeypatch):
@@ -51,13 +51,15 @@ def test_neo4j_username_env_alias_override(monkeypatch):
     get_settings.cache_clear()
 
 
-def test_reranker_empty_chunks():
+@pytest.mark.asyncio
+async def test_reranker_empty_chunks():
     from civicsetu.agent.nodes import reranker_node
-    result = reranker_node(_base_state(retrieved_chunks=[], reranked_chunks=[]))
+    result = await reranker_node(_base_state(retrieved_chunks=[], reranked_chunks=[]))
     assert result["reranked_chunks"] == []
 
 
-def test_reranker_pinned_chunks_always_first():
+@pytest.mark.asyncio
+async def test_reranker_pinned_chunks_always_first():
     from civicsetu.agent.nodes import reranker_node
 
     pinned = _make_rc(section_id="18", is_pinned=True)
@@ -74,14 +76,15 @@ def test_reranker_pinned_chunks_always_first():
             reranked_chunks=[],
             query="test query",
         )
-        result = reranker_node(state)
+        result = await reranker_node(state)
 
     reranked = result["reranked_chunks"]
     assert reranked[0].is_pinned is True
     assert reranked[0].chunk.section_id == "18"
 
 
-def test_reranker_keeps_pinned_chunks_up_to_context_limit():
+@pytest.mark.asyncio
+async def test_reranker_keeps_pinned_chunks_up_to_context_limit():
     from civicsetu.agent.nodes import reranker_node
 
     pinned_chunks = [_make_rc(section_id=str(i), is_pinned=True) for i in range(4)]
@@ -95,13 +98,14 @@ def test_reranker_keeps_pinned_chunks_up_to_context_limit():
             reranked_chunks=[],
             query="test query",
         )
-        result = reranker_node(state)
+        result = await reranker_node(state)
 
     pinned_in_result = [c for c in result["reranked_chunks"] if c.is_pinned]
     assert len(pinned_in_result) == 4
 
 
-def test_reranker_deduplicates_by_chunk_id():
+@pytest.mark.asyncio
+async def test_reranker_deduplicates_by_chunk_id():
     from civicsetu.agent.nodes import reranker_node
 
     chunk = _make_rc(section_id="18")
@@ -118,7 +122,7 @@ def test_reranker_deduplicates_by_chunk_id():
             reranked_chunks=[],
             query="test query",
         )
-        result = reranker_node(state)
+        result = await reranker_node(state)
 
     all_ids = [str(c.chunk.chunk_id) for c in result["reranked_chunks"]]
     assert len(all_ids) == len(set(all_ids))
@@ -178,7 +182,8 @@ def test_pin_relevance_prefers_specific_subclauses_over_base_header_on_tie():
     assert [c.chunk.section_id for c in ranked] == ["7(6)", "7"]
 
 
-def test_prepend_pinned_sections_promotes_existing_matches(monkeypatch):
+@pytest.mark.asyncio
+async def test_prepend_pinned_sections_promotes_existing_matches(monkeypatch):
     from civicsetu.agent.nodes import _prepend_pinned_sections
     from civicsetu.models.enums import Jurisdiction
 
@@ -186,18 +191,17 @@ def test_prepend_pinned_sections_promotes_existing_matches(monkeypatch):
     rule_5 = _make_rc(section_id="5", jurisdiction=Jurisdiction.KARNATAKA)
     section_4 = _make_rc(section_id="4(14)", jurisdiction=Jurisdiction.CENTRAL)
 
-    def fake_run(coro):
-        coro.close()
+    async def fake_fetch(refs, jur, existing, hint=""):
         return []
 
-    monkeypatch.setattr("civicsetu.agent.nodes.asyncio.run", fake_run)
+    monkeypatch.setattr("civicsetu.agent.nodes._fetch_pinned_sections", fake_fetch)
 
-    result = _prepend_pinned_sections(
-        {
-            "pinned_section_refs": ["Rule 5", "Section 4"],
-            "pinned_section_jurisdiction": Jurisdiction.KARNATAKA,
-            "pinned_section_hint": "separate bank account seventy per cent",
-        },
+    result = await _prepend_pinned_sections(
+        _base_state(
+            pinned_section_refs=["Rule 5", "Section 4"],
+            pinned_section_jurisdiction=Jurisdiction.KARNATAKA,
+            pinned_section_hint="separate bank account seventy per cent",
+        ),
         [noisy, rule_5, section_4],
     )
 
@@ -206,7 +210,8 @@ def test_prepend_pinned_sections_promotes_existing_matches(monkeypatch):
     assert result[1].is_pinned is True
 
 
-def test_reranker_node_trims_eval_context_to_pinned_families():
+@pytest.mark.asyncio
+async def test_reranker_node_trims_eval_context_to_pinned_families():
     from civicsetu.agent.nodes import reranker_node
     from civicsetu.models.enums import Jurisdiction
 
@@ -215,11 +220,12 @@ def test_reranker_node_trims_eval_context_to_pinned_families():
     noisy_central = _make_rc(section_id="4(10)", jurisdiction=Jurisdiction.CENTRAL)
     noisy_state = _make_rc(section_id="3", jurisdiction=Jurisdiction.UTTAR_PRADESH)
 
-    with patch(
-        "civicsetu.retrieval.reranker.Reranker.rerank",
-        return_value=[matched_rule, matched_section, noisy_central, noisy_state],
-    ):
-        result = reranker_node(
+    # Use a side effect to return only the input chunks that were not filtered out
+    def mock_rerank(chunks, query):
+        return chunks
+
+    with patch("civicsetu.retrieval.reranker.Reranker.rerank", side_effect=mock_rerank):
+        result = await reranker_node(
             _base_state(
                 query="Which Karnataka rule implements project account maintenance?",
                 rewritten_query="Which Karnataka rule implements project account maintenance?",
@@ -244,7 +250,8 @@ def _llm_response(**kwargs) -> str:
     return json.dumps(payload)
 
 
-def test_generator_cites_only_referenced_chunks():
+@pytest.mark.asyncio
+async def test_generator_cites_only_referenced_chunks():
     from civicsetu.agent.nodes import generator_node
 
     chunks = [_make_rc(section_id=str(i)) for i in range(1, 6)]
@@ -256,14 +263,15 @@ def test_generator_cites_only_referenced_chunks():
             query="What does Section 1 say?",
             reranked_chunks=chunks,
         )
-        result = generator_node(state)
+        result = await generator_node(state)
 
     assert len(result["citations"]) == 2
     cited_ids = {c.section_id for c in result["citations"]}
     assert cited_ids == {"1", "3"}
 
 
-def test_generator_fallback_when_cited_chunks_empty():
+@pytest.mark.asyncio
+async def test_generator_fallback_when_cited_chunks_empty():
     from civicsetu.agent.nodes import generator_node
 
     chunks = [_make_rc(section_id=str(i)) for i in range(1, 4)]
@@ -271,13 +279,14 @@ def test_generator_fallback_when_cited_chunks_empty():
 
     with patch("civicsetu.agent.nodes._llm_call", return_value=llm_out):
         state = _base_state(reranked_chunks=chunks)
-        result = generator_node(state)
+        result = await generator_node(state)
 
     # Fallback: all 3 chunks cited
     assert len(result["citations"]) == 3
 
 
-def test_generator_filters_invalid_indices():
+@pytest.mark.asyncio
+async def test_generator_filters_invalid_indices():
     from civicsetu.agent.nodes import generator_node
 
     chunks = [_make_rc(section_id=str(i)) for i in range(1, 4)]
@@ -286,48 +295,52 @@ def test_generator_filters_invalid_indices():
 
     with patch("civicsetu.agent.nodes._llm_call", return_value=llm_out):
         state = _base_state(reranked_chunks=chunks)
-        result = generator_node(state)
+        result = await generator_node(state)
 
     # Only index 2 is valid → 1 citation
     assert len(result["citations"]) == 1
     assert result["citations"][0].section_id == "2"
 
 
-def test_generator_returns_empty_on_no_chunks():
+@pytest.mark.asyncio
+async def test_generator_returns_empty_on_no_chunks():
     from civicsetu.agent.nodes import generator_node
 
-    result = generator_node(_base_state(reranked_chunks=[]))
+    result = await generator_node(_base_state(reranked_chunks=[]))
     assert result["citations"] == []
     assert result["confidence_score"] == 0.0
 
 
-def test_generator_handles_malformed_llm_json():
+@pytest.mark.asyncio
+async def test_generator_handles_malformed_llm_json():
     from civicsetu.agent.nodes import generator_node
 
     chunks = [_make_rc(section_id="18")]
     with patch("civicsetu.agent.nodes._llm_call", return_value="not json {{{{"):
         state = _base_state(reranked_chunks=chunks)
-        result = generator_node(state)
+        result = await generator_node(state)
     # Salvage path: non-empty malformed text is returned as answer with 0.3 confidence
     assert result["confidence_score"] == 0.3
     assert len(result["citations"]) == 1  # salvage cites all provided chunks
 
 
-def test_generator_salvages_plain_text_when_json_missing():
+@pytest.mark.asyncio
+async def test_generator_salvages_plain_text_when_json_missing():
     from civicsetu.agent.nodes import generator_node
 
     chunks = [_make_rc(section_id="18"), _make_rc(section_id="31")]
     raw = "Promoter must register project, disclose details, and honor timelines."
 
     with patch("civicsetu.agent.nodes._llm_call", return_value=raw):
-        result = generator_node(_base_state(reranked_chunks=chunks))
+        result = await generator_node(_base_state(reranked_chunks=chunks))
 
     assert result["raw_response"] == raw
     assert result["confidence_score"] == 0.3
     assert len(result["citations"]) == 2
 
 
-def test_classifier_extracts_json_from_reasoning_wrapper():
+@pytest.mark.asyncio
+async def test_classifier_extracts_json_from_reasoning_wrapper():
     from civicsetu.agent.nodes import classifier_node
     from civicsetu.models.enums import QueryType
 
@@ -341,13 +354,14 @@ Thinking:
 """
 
     with patch("civicsetu.agent.nodes._llm_call", return_value=wrapped):
-        result = classifier_node(_base_state(query="What does Section 18 say?"))
+        result = await classifier_node(_base_state(query="What does Section 18 say?"))
 
     assert result["query_type"] == QueryType.CROSS_REFERENCE
     assert result["rewritten_query"] == "Section 18 refund obligations under RERA Act"
 
 
-def test_generator_extracts_json_from_reasoning_wrapper():
+@pytest.mark.asyncio
+async def test_generator_extracts_json_from_reasoning_wrapper():
     from civicsetu.agent.nodes import generator_node
 
     chunks = [_make_rc(section_id="18")]
@@ -360,14 +374,15 @@ Here is structured answer.
 """
 
     with patch("civicsetu.agent.nodes._llm_call", return_value=wrapped):
-        result = generator_node(_base_state(reranked_chunks=chunks))
+        result = await generator_node(_base_state(reranked_chunks=chunks))
 
     assert result["raw_response"] == "Refund due under Section 18."
     assert result["confidence_score"] == 0.8
     assert len(result["citations"]) == 1
 
 
-def test_llm_call_uses_json_mode_for_osmapi(monkeypatch):
+@pytest.mark.asyncio
+async def test_llm_call_uses_json_mode_for_osmapi(monkeypatch):
     import civicsetu.agent.nodes as nodes_mod
 
     monkeypatch.setenv("OPENAI_API_BASE", "https://api.osmapi.com/v1")
@@ -376,14 +391,16 @@ def test_llm_call_uses_json_mode_for_osmapi(monkeypatch):
     fake_response.choices = [MagicMock(message=MagicMock(content='{"ok": true}'))]
     fake_response.usage = None
 
-    with patch("civicsetu.agent.nodes.litellm.completion", return_value=fake_response) as completion:
-        result = nodes_mod._llm_call("prompt", "system")
+    with patch("civicsetu.agent.nodes.litellm.acompletion", new=MagicMock(return_value=asyncio.Future())) as completion:
+        completion.return_value.set_result(fake_response)
+        result = await nodes_mod._llm_call("prompt", "system")
 
     assert result == '{"ok": true}'
     assert completion.call_args.kwargs["response_format"] == {"type": "json_object"}
 
 
-def test_llm_call_does_not_force_no_reasoning_for_osmapi(monkeypatch):
+@pytest.mark.asyncio
+async def test_llm_call_does_not_force_no_reasoning_for_osmapi(monkeypatch):
     import civicsetu.agent.nodes as nodes_mod
 
     monkeypatch.setenv("OPENAI_API_BASE", "https://api.osmapi.com/v1")
@@ -395,16 +412,18 @@ def test_llm_call_does_not_force_no_reasoning_for_osmapi(monkeypatch):
     original_models = nodes_mod.THINKING_MODELS[:]
     nodes_mod.THINKING_MODELS[:] = ["openai/gpt-4o-mini"]
 
-    with patch("civicsetu.agent.nodes.litellm.completion", return_value=fake_response) as completion:
+    with patch("civicsetu.agent.nodes.litellm.acompletion", new=MagicMock(return_value=asyncio.Future())) as completion:
+        completion.return_value.set_result(fake_response)
         try:
-            nodes_mod._llm_call("prompt", "system")
+            await nodes_mod._llm_call("prompt", "system")
         finally:
             nodes_mod.THINKING_MODELS[:] = original_models
 
     assert "extra_body" not in completion.call_args.kwargs
 
 
-def test_llm_call_does_not_attach_deepseek_kwargs_to_non_deepseek_models():
+@pytest.mark.asyncio
+async def test_llm_call_does_not_attach_deepseek_kwargs_to_non_deepseek_models():
     import civicsetu.agent.nodes as nodes_mod
 
     fake_response = MagicMock()
@@ -414,9 +433,10 @@ def test_llm_call_does_not_attach_deepseek_kwargs_to_non_deepseek_models():
     original_models = nodes_mod.THINKING_MODELS[:]
     nodes_mod.THINKING_MODELS[:] = ["groq/llama-3.3-70b-versatile"]
 
-    with patch("civicsetu.agent.nodes.litellm.completion", return_value=fake_response) as completion:
+    with patch("civicsetu.agent.nodes.litellm.acompletion", new=MagicMock(return_value=asyncio.Future())) as completion:
+        completion.return_value.set_result(fake_response)
         try:
-            nodes_mod._llm_call("prompt", "system", tier="thinking")
+            await nodes_mod._llm_call("prompt", "system", tier="thinking")
         finally:
             nodes_mod.THINKING_MODELS[:] = original_models
 
@@ -424,7 +444,8 @@ def test_llm_call_does_not_attach_deepseek_kwargs_to_non_deepseek_models():
     assert "extra_body" not in completion.call_args.kwargs
 
 
-def test_generator_deduplicates_citations():
+@pytest.mark.asyncio
+async def test_generator_deduplicates_citations():
     from civicsetu.agent.nodes import generator_node
 
     # Two chunks with same (section_id, doc_name) — different chunk_ids
@@ -434,18 +455,19 @@ def test_generator_deduplicates_citations():
 
     with patch("civicsetu.agent.nodes._llm_call", return_value=llm_out):
         state = _base_state(reranked_chunks=[chunk_a, chunk_b])
-        result = generator_node(state)
+        result = await generator_node(state)
 
     assert len(result["citations"]) == 1
 
 
-def test_generator_system_prompt_uses_plain_language_persona():
+@pytest.mark.asyncio
+async def test_generator_system_prompt_uses_plain_language_persona():
     from civicsetu.agent.nodes import generator_node
     from civicsetu.models.enums import QueryType
 
     captured = {}
 
-    def fake_llm_call(prompt: str, system: str, temperature: float = 0.0) -> str:
+    async def fake_llm_call(prompt: str, system: str, temperature: float = 0.0, tier="thinking") -> str:
         captured["prompt"] = prompt
         captured["system"] = system
         return _llm_response(answer="If a builder misses the deadline, the buyer gets a remedy.")
@@ -455,7 +477,7 @@ def test_generator_system_prompt_uses_plain_language_persona():
             query_type=QueryType.FACT_LOOKUP,
             reranked_chunks=[_make_rc(section_id="18")],
         )
-        generator_node(state)
+        await generator_node(state)
 
     assert "plain-language guide to Indian RERA laws" in captured["system"]
     assert "explain what the law means in practice" in captured["system"]
@@ -474,12 +496,13 @@ def test_generator_system_prompt_uses_plain_language_persona():
         ("temporal", "Lead with the specific time period or deadline"),
     ],
 )
-def test_generator_system_prompt_includes_query_type_tone_hint(query_type, expected_hint):
+@pytest.mark.asyncio
+async def test_generator_system_prompt_includes_query_type_tone_hint(query_type, expected_hint):
     from civicsetu.agent.nodes import generator_node
 
     captured = {}
 
-    def fake_llm_call(prompt: str, system: str, temperature: float = 0.0) -> str:
+    async def fake_llm_call(prompt: str, system: str, temperature: float = 0.0, tier="thinking") -> str:
         captured["system"] = system
         return _llm_response()
 
@@ -488,7 +511,7 @@ def test_generator_system_prompt_includes_query_type_tone_hint(query_type, expec
             query_type=query_type,
             reranked_chunks=[_make_rc(section_id="18")],
         )
-        generator_node(state)
+        await generator_node(state)
 
     assert expected_hint in captured["system"]
 
@@ -502,12 +525,12 @@ def test_get_ranker_uses_settings_model():
     reranker_mod._ranker = None  # clear module-level cache
 
     with patch("civicsetu.retrieval.reranker.settings") as mock_settings:
-        mock_settings.reranker_model = "rank-T5-flan"
+        mock_settings.reranker_model = "ms-marco-MiniLM-L-12-v2"
         with patch("flashrank.Ranker") as MockRanker:
             MockRanker.return_value = MagicMock()
             reranker_mod._get_ranker()
             MockRanker.assert_called_once_with(
-                model_name="rank-T5-flan", cache_dir=".cache/flashrank"
+                model_name="ms-marco-MiniLM-L-12-v2", cache_dir=".cache/flashrank"
             )
     reranker_mod._ranker = None  # clean up
 
@@ -602,7 +625,8 @@ def test_score_gap_new_threshold_still_cuts_on_cliff():
 
 # ── reranker_node threshold + gap filtering ───────────────────────────────────
 
-def test_reranker_drops_below_threshold():
+@pytest.mark.asyncio
+async def test_reranker_drops_below_threshold():
     """Chunks scoring below reranker_score_threshold must not appear in output."""
     from civicsetu.agent.nodes import reranker_node
     from unittest.mock import patch, MagicMock
@@ -626,14 +650,15 @@ def test_reranker_drops_below_threshold():
                 reranked_chunks=[],
                 query="test query",
             )
-            result = reranker_node(state)
+            result = await reranker_node(state)
 
     section_ids = [c.chunk.section_id for c in result["reranked_chunks"]]
     assert "1" in section_ids
     assert "2" not in section_ids
 
 
-def test_reranker_applies_score_gap():
+@pytest.mark.asyncio
+async def test_reranker_applies_score_gap():
     """A large score cliff stops inclusion even if remaining chunks exceed threshold."""
     from civicsetu.agent.nodes import reranker_node
     from unittest.mock import patch, MagicMock
@@ -660,7 +685,7 @@ def test_reranker_applies_score_gap():
                 reranked_chunks=[],
                 query="test query",
             )
-            result = reranker_node(state)
+            result = await reranker_node(state)
 
     section_ids = [c.chunk.section_id for c in result["reranked_chunks"]]
     assert "1" in section_ids
@@ -668,7 +693,8 @@ def test_reranker_applies_score_gap():
     assert "3" not in section_ids
 
 
-def test_reranker_filtered_count_logged():
+@pytest.mark.asyncio
+async def test_reranker_filtered_count_logged():
     """reranker_filtered log event must report correct before/after/dropped counts."""
     from civicsetu.agent.nodes import reranker_node
     from unittest.mock import patch, MagicMock
@@ -693,7 +719,7 @@ def test_reranker_filtered_count_logged():
                     reranked_chunks=[],
                     query="test query",
                 )
-                result = reranker_node(state)
+                result = await reranker_node(state)
 
     # Verify reranked_chunks output
     assert len(result["reranked_chunks"]) == 1

tests/unit/api/test_query_route.py

@@ -142,14 +142,14 @@ def test_graph_topology_returns_empty_payload_when_neo4j_auth_fails(client):
 
 def test_query_returns_200_with_citations(client):
     test_client, mock_graph = client
-    mock_graph.invoke.return_value = {
+    mock_graph.ainvoke = AsyncMock(return_value={
         "raw_response": "Under Section 18, the promoter must...",
         "citations": [_make_citation("18")],
         "confidence_score": 0.9,
         "query_type": QueryType.CROSS_REFERENCE,
         "conflict_warnings": [],
         "amendment_notice": None,
-    }
+    })
 
     response = test_client.post("/api/v1/query", json={"query": "What does Section 18 say?"})
     assert response.status_code == 200
@@ -162,14 +162,14 @@ def test_query_returns_200_with_citations(client):
 
 def test_query_returns_insufficient_when_no_citations(client):
     test_client, mock_graph = client
-    mock_graph.invoke.return_value = {
+    mock_graph.ainvoke = AsyncMock(return_value={
         "raw_response": "I don't know",
         "citations": [],
         "confidence_score": 0.2,
         "query_type": QueryType.FACT_LOOKUP,
         "conflict_warnings": [],
         "amendment_notice": None,
-    }
+    })
 
     response = test_client.post("/api/v1/query", json={"query": "Some obscure question here"})
     assert response.status_code == 200
@@ -194,14 +194,14 @@ def test_query_rejects_top_k_out_of_range(client):
 
 def test_query_response_always_has_disclaimer(client):
     test_client, mock_graph = client
-    mock_graph.invoke.return_value = {
+    mock_graph.ainvoke = AsyncMock(return_value={
         "raw_response": "Under Section 18...",
         "citations": [_make_citation()],
         "confidence_score": 0.8,
         "query_type": QueryType.FACT_LOOKUP,
         "conflict_warnings": [],
         "amendment_notice": None,
-    }
+    })
 
     response = test_client.post("/api/v1/query", json={"query": "What does Section 18 say?"})
     body = response.json()
@@ -211,14 +211,14 @@ def test_query_response_always_has_disclaimer(client):
 
 def test_query_reuses_provided_session_id(client):
     test_client, mock_graph = client
-    mock_graph.invoke.return_value = {
+    mock_graph.ainvoke = AsyncMock(return_value={
         "raw_response": "Under Section 18...",
         "citations": [_make_citation()],
         "confidence_score": 0.8,
         "query_type": QueryType.FACT_LOOKUP,
         "conflict_warnings": [],
         "amendment_notice": None,
-    }
+    })
 
     response = test_client.post(
         "/api/v1/query",
@@ -228,7 +228,7 @@ def test_query_reuses_provided_session_id(client):
     body = response.json()
     assert response.status_code == 200
     assert body["session_id"] == "my-session-abc"
-    _, config = mock_graph.invoke.call_args.args
+    _, config = mock_graph.ainvoke.call_args.args
     assert config["configurable"]["thread_id"] == "my-session-abc"
 
 
@@ -393,14 +393,14 @@ def test_graph_section_content_resolves_graph_chunk_id_case_insensitive_status(c
 
 def test_section_context_query_sets_skip_classifier_and_source_section(client):
     test_client, mock_graph = client
-    mock_graph.invoke.return_value = {
+    mock_graph.ainvoke = AsyncMock(return_value={
         "raw_response": "Section 18 requires refund with interest.",
         "citations": [_make_citation("18")],
         "confidence_score": 0.88,
         "query_type": QueryType.CROSS_REFERENCE,
         "conflict_warnings": [],
         "amendment_notice": None,
-    }
+    })
 
     response = test_client.post(
         "/api/v1/query/section-context",
@@ -415,7 +415,7 @@ def test_section_context_query_sets_skip_classifier_and_source_section(client):
     assert response.status_code == 200
     body = response.json()
     assert body["session_id"] == "section-thread-1"
-    initial_state, config = mock_graph.invoke.call_args.args
+    initial_state, config = mock_graph.ainvoke.call_args.args
     assert initial_state["skip_classifier"] is True
     assert initial_state["source_section_id"] == "18"
     assert initial_state["query_type"] == QueryType.CROSS_REFERENCE

tests/unit/stores/test_graph_store.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from unittest.mock import AsyncMock
+
+import pytest
+
+
+def _make_driver(session):
+    driver = AsyncMock()
+
+    @asynccontextmanager
+    async def _session_cm():
+        yield session
+
+    driver.session = _session_cm
+    return driver
+
+
+@pytest.mark.asyncio
+async def test_graph_topology_retries_after_transient_driver_error(monkeypatch):
+    from civicsetu.stores import graph_store
+
+    first_session = AsyncMock()
+    first_session.run.side_effect = RuntimeError("the connection is closed")
+
+    edges_result = AsyncMock()
+    edges_result.data.return_value = [{"source": "a", "target": "b", "edge_type": "REFERENCES"}]
+    nodes_result = AsyncMock()
+    nodes_result.data.return_value = [{"chunk_id": "a", "section_id": "18"}]
+
+    second_session = AsyncMock()
+    second_session.run.side_effect = [edges_result, nodes_result]
+
+    first_driver = _make_driver(first_session)
+    second_driver = _make_driver(second_session)
+    get_driver = AsyncMock(side_effect=[first_driver, second_driver])
+    reset_driver = AsyncMock()
+
+    monkeypatch.setattr(graph_store, "_get_driver", get_driver)
+    monkeypatch.setattr(graph_store, "_reset_driver", reset_driver)
+
+    nodes, edges = await graph_store.GraphStore.get_topology()
+
+    assert nodes == [{"chunk_id": "a", "section_id": "18"}]
+    assert edges == [{"source": "a", "target": "b", "edge_type": "REFERENCES"}]
+    assert get_driver.await_count == 2
+    reset_driver.assert_awaited_once()