fix: harden graph retrieval retries

src/civicsetu/agent/nodes.py

@@ -519,7 +519,12 @@ async def graph_retrieval_node(state: CivicSetuState) -> dict:
     top_k = state.get("top_k", 5)
     node_start = time.perf_counter()
 
-    log.info("graph_retrieval_node", query=query[:80])
+    log.info(
+        "graph_retrieval_node",
+        query=query[:80],
+        jurisdiction=str(jurisdiction) if jurisdiction else None,
+        source_section_id=state.get("source_section_id"),
+    )
 
     async def _retrieve():
         return await GraphRetriever.retrieve(
@@ -530,9 +535,25 @@ async def graph_retrieval_node(state: CivicSetuState) -> dict:
         )
 
     retrieve_start = time.perf_counter()
-    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))
+    try:
+        chunks = await _retrieve()
+        log.info("graph_retrieval_complete", results=len(chunks))
+    except Exception as e:
+        log.error(
+            "graph_retrieval_failed",
+            query=query[:80],
+            jurisdiction=str(jurisdiction) if jurisdiction else None,
+            source_section_id=state.get("source_section_id"),
+            error_type=type(e).__name__,
+            error=str(e),
+        )
+        raise
+    log.info(
+        "stage_timing",
+        node="graph_retrieval",
+        stage="neo4j_postgres_hydration",
+        duration_ms=round((time.perf_counter() - retrieve_start) * 1000, 2),
+    )
     
     # Fallback: if graph found nothing (no explicit section in query),
     # run RRF hybrid retrieval instead of pure vector
@@ -627,7 +648,12 @@ async def graph_retrieval_node(state: CivicSetuState) -> dict:
         )
         chunks = chunks[:_MAX_GRAPH_CHUNKS]
 
-    log.info("node_timing", node="graph_retrieval", duration_ms=round((time.perf_counter() - node_start) * 1000, 2), results=len(chunks))
+    log.info(
+        "node_timing",
+        node="graph_retrieval",
+        duration_ms=round((time.perf_counter() - node_start) * 1000, 2),
+        results=len(chunks),
+    )
     return {"retrieved_chunks": chunks}
 
 

src/civicsetu/retrieval/graph_retriever.py

@@ -12,9 +12,15 @@ from civicsetu.stores.vector_store import VectorStore
 
 log = structlog.get_logger(__name__)
 
-# All jurisdictions searched when no filter is provided.
-# Order matters: CENTRAL first so DERIVED_FROM incoming edges surface
-# all state rules in one traversal pass.
+_TRANSIENT_CONNECTION_MARKERS = (
+    "defunct connection",
+    "the connection is closed",
+    "connection reset by peer",
+    "unable to retrieve routing information",
+    "service unavailable",
+    "terminating connection due to administrator command",
+)
+
 _ALL_JURISDICTIONS = [
     "CENTRAL",
     "MAHARASHTRA",
@@ -23,27 +29,10 @@ _ALL_JURISDICTIONS = [
     "TAMIL_NADU",
 ]
 
-# Max chunks returned to reranker — prevents FlashRank serial inference blowup
-# on high-connectivity sections (e.g. §9 has DERIVED_FROM edges from 8 state rules)
 _MAX_GRAPH_CHUNKS = 20
 
 
 class GraphRetriever:
-    """
-    Retrieves chunks via Neo4j graph traversal rather than vector similarity.
-
-    Used for:
-    - cross_reference queries: "What does Section 18 reference?"
-    - penalty_lookup queries: "What are the penalties under Section 59?"
-    - temporal queries: "Which sections were amended?"
-
-    Strategy:
-    1. Extract section_id from query (regex)
-    2. Traverse REFERENCES + DERIVED_FROM edges in Neo4j (all 5 jurisdictions in parallel)
-    3. Hydrate section_ids into LegalChunk objects via VectorStore.get_by_section()
-    4. Dedup + cap at _MAX_GRAPH_CHUNKS before returning
-    """
-
     @staticmethod
     async def retrieve(
         query: str,
@@ -69,13 +58,7 @@ class GraphRetriever:
             )
             return cached
 
-        # Explicit filter → search only that jurisdiction.
-        # No filter → search all jurisdictions (CENTRAL first).
-        jurisdictions_to_search = (
-            [jurisdiction.value] if jurisdiction
-            else _ALL_JURISDICTIONS
-        )
-
+        jurisdictions_to_search = [jurisdiction.value] if jurisdiction else _ALL_JURISDICTIONS
         log.info(
             "graph_retriever_traversing",
             section_id=section_id,
@@ -83,124 +66,155 @@ class GraphRetriever:
             jurisdictions=jurisdictions_to_search,
         )
 
-        chunks: list[RetrievedChunk] = []
-        seen_chunk_ids: set[str] = set()
-
-        async def _fetch_jurisdiction(jur_str: str) -> list[RetrievedChunk]:
-            async with AsyncSessionLocal() as session:  # ← session scoped per coroutine
-                jur_enum = Jurisdiction(jur_str) if not jurisdiction else jurisdiction
-                jur_chunks: list[RetrievedChunk] = []
-
-                source = await VectorStore.get_by_section(
-                    session=session, section_id=section_id, jurisdiction=jur_enum,
-                )
-                for rc in source:
-                    rc.retrieval_source = "graph"
-                    rc.graph_path = f"source:{section_id}@{jur_str}"
-                    rc.is_pinned = rc.chunk.section_id == section_id
-                jur_chunks.extend(source)
-
-                outgoing = await GraphStore.get_referenced_sections(section_id, jur_str, depth)
-                for node in outgoing:
-                    hydrated = await VectorStore.get_by_section(
-                        session=session,
-                        section_id=node["section_id"],
-                        jurisdiction=Jurisdiction(node["jurisdiction"]),
-                    )
-                    for rc in hydrated:
-                        rc.retrieval_source = "graph"
-                        rc.graph_path = f"{section_id} →[REFERENCES]→ {node['section_id']}"
-                    jur_chunks.extend(hydrated)
+        async def _fetch_once() -> list[RetrievedChunk]:
+            chunks: list[RetrievedChunk] = []
+            seen_chunk_ids: set[str] = set()
 
-                incoming = await GraphStore.get_sections_referencing(section_id, jur_str)
-                for node in incoming:
-                    hydrated = await VectorStore.get_by_section(
-                        session=session,
-                        section_id=node["section_id"],
-                        jurisdiction=Jurisdiction(node["jurisdiction"]),
-                    )
-                    for rc in hydrated:
-                        rc.retrieval_source = "graph"
-                        rc.graph_path = f"{node['section_id']} →[REFERENCES]→ {section_id}"
-                    jur_chunks.extend(hydrated)
+            async def _fetch_jurisdiction_once(jur_str: str) -> list[RetrievedChunk]:
+                async with AsyncSessionLocal() as session:
+                    jur_enum = Jurisdiction(jur_str) if not jurisdiction else jurisdiction
+                    jur_chunks: list[RetrievedChunk] = []
 
-                derived_act = await GraphStore.get_derived_act_sections(section_id, jur_str)
-                for node in derived_act:
-                    hydrated = await VectorStore.get_by_section(
+                    source = await VectorStore.get_by_section(
                         session=session,
-                        section_id=node["section_id"],
-                        jurisdiction=Jurisdiction(node["jurisdiction"]),
+                        section_id=section_id,
+                        jurisdiction=jur_enum,
                     )
-                    for rc in hydrated:
+                    for rc in source:
                         rc.retrieval_source = "graph"
-                        rc.graph_path = f"{section_id}@{jur_str} →[DERIVED_FROM]→ {node['section_id']}@{node['jurisdiction']}"
-                    jur_chunks.extend(hydrated)
-
-                deriving = await GraphStore.get_deriving_rule_sections(section_id, jur_str)
-                for node in deriving:
-                    hydrated = await VectorStore.get_by_section(
-                        session=session,
-                        section_id=node["section_id"],
-                        jurisdiction=Jurisdiction(node["jurisdiction"]),
+                        rc.graph_path = f"source:{section_id}@{jur_str}"
+                        rc.is_pinned = rc.chunk.section_id == section_id
+                    jur_chunks.extend(source)
+
+                    outgoing = await GraphStore.get_referenced_sections(section_id, jur_str, depth)
+                    for node in outgoing:
+                        hydrated = await VectorStore.get_by_section(
+                            session=session,
+                            section_id=node["section_id"],
+                            jurisdiction=Jurisdiction(node["jurisdiction"]),
+                        )
+                        for rc in hydrated:
+                            rc.retrieval_source = "graph"
+                            rc.graph_path = f"{section_id} ->[REFERENCES]-> {node['section_id']}"
+                        jur_chunks.extend(hydrated)
+
+                    incoming = await GraphStore.get_sections_referencing(section_id, jur_str)
+                    for node in incoming:
+                        hydrated = await VectorStore.get_by_section(
+                            session=session,
+                            section_id=node["section_id"],
+                            jurisdiction=Jurisdiction(node["jurisdiction"]),
+                        )
+                        for rc in hydrated:
+                            rc.retrieval_source = "graph"
+                            rc.graph_path = f"{node['section_id']} ->[REFERENCES]-> {section_id}"
+                        jur_chunks.extend(hydrated)
+
+                    derived_act = await GraphStore.get_derived_act_sections(section_id, jur_str)
+                    for node in derived_act:
+                        hydrated = await VectorStore.get_by_section(
+                            session=session,
+                            section_id=node["section_id"],
+                            jurisdiction=Jurisdiction(node["jurisdiction"]),
+                        )
+                        for rc in hydrated:
+                            rc.retrieval_source = "graph"
+                            rc.graph_path = f"{section_id}@{jur_str} ->[DERIVED_FROM]-> {node['section_id']}@{node['jurisdiction']}"
+                        jur_chunks.extend(hydrated)
+
+                    deriving = await GraphStore.get_deriving_rule_sections(section_id, jur_str)
+                    for node in deriving:
+                        hydrated = await VectorStore.get_by_section(
+                            session=session,
+                            section_id=node["section_id"],
+                            jurisdiction=Jurisdiction(node["jurisdiction"]),
+                        )
+                        for rc in hydrated:
+                            rc.retrieval_source = "graph"
+                            rc.graph_path = f"{node['section_id']}@{node['jurisdiction']} ->[DERIVED_FROM]-> {section_id}@{jur_str}"
+                        jur_chunks.extend(hydrated)
+
+                    return jur_chunks
+
+            async def _fetch_jurisdiction(jur_str: str) -> list[RetrievedChunk]:
+                try:
+                    return await _fetch_jurisdiction_once(jur_str)
+                except Exception as first_error:
+                    if not any(marker in str(first_error).lower() for marker in _TRANSIENT_CONNECTION_MARKERS):
+                        raise
+                    log.warning(
+                        "graph_jurisdiction_retrying",
+                        section_id=section_id,
+                        jurisdiction=jur_str,
+                        depth=depth,
+                        error=str(first_error),
                     )
-                    for rc in hydrated:
-                        rc.retrieval_source = "graph"
-                        rc.graph_path = f"{node['section_id']}@{node['jurisdiction']} →[DERIVED_FROM]→ {section_id}@{jur_str}"
-                    jur_chunks.extend(hydrated)
+                    from civicsetu.stores.graph_store import _reset_driver
 
-                return jur_chunks
+                    await _reset_driver()
+                    return await _fetch_jurisdiction_once(jur_str)
 
-        # gather is now OUTSIDE the session context
-        all_results = await asyncio.gather(
-            *[_fetch_jurisdiction(j) for j in jurisdictions_to_search],
-            return_exceptions=True,
-        )
-        for result in all_results:
-            if isinstance(result, Exception):
-                log.warning("graph_jurisdiction_fetch_failed", error=str(result))
-            else:
-                chunks.extend(result)
-
-        # Dedup by chunk_id (same chunk can arrive via multiple traversal paths)
-        deduped: list[RetrievedChunk] = []
-        for rc in chunks:
-            cid = str(rc.chunk.chunk_id)
-            if cid not in seen_chunk_ids:
-                seen_chunk_ids.add(cid)
-                deduped.append(rc)
-
-        # Pinned chunks always included; fill remaining slots from deduped
-        pinned = [rc for rc in deduped if rc.is_pinned]
-        rest = [rc for rc in deduped if not rc.is_pinned]
-        final = (pinned + rest)[:_MAX_GRAPH_CHUNKS]
+            all_results = await asyncio.gather(
+                *[_fetch_jurisdiction(j) for j in jurisdictions_to_search],
+                return_exceptions=True,
+            )
+            for result in all_results:
+                if isinstance(result, Exception):
+                    log.warning(
+                        "graph_jurisdiction_fetch_failed",
+                        section_id=section_id,
+                        error_type=type(result).__name__,
+                        error=str(result),
+                    )
+                else:
+                    chunks.extend(result)
+
+            deduped: list[RetrievedChunk] = []
+            for rc in chunks:
+                cid = str(rc.chunk.chunk_id)
+                if cid not in seen_chunk_ids:
+                    seen_chunk_ids.add(cid)
+                    deduped.append(rc)
+
+            pinned = [rc for rc in deduped if rc.is_pinned]
+            rest = [rc for rc in deduped if not rc.is_pinned]
+            final = (pinned + rest)[:_MAX_GRAPH_CHUNKS]
+            log.info(
+                "graph_retriever_complete",
+                section_id=section_id,
+                jurisdictions_searched=jurisdictions_to_search,
+                chunks_hydrated=len(final),
+                pinned=len(pinned),
+                deduped_total=len(deduped),
+            )
+            return final
+
+        try:
+            final = await _fetch_once()
+        except Exception as first_error:
+            if not any(marker in str(first_error).lower() for marker in _TRANSIENT_CONNECTION_MARKERS):
+                raise
+            log.warning(
+                "graph_retriever_retrying",
+                section_id=section_id,
+                jurisdiction=jurisdiction_key,
+                depth=depth,
+                error=str(first_error),
+            )
+            from civicsetu.stores.graph_store import _reset_driver
+
+            await _reset_driver()
+            final = await _fetch_once()
 
-        log.info(
-            "graph_retriever_complete",
-            section_id=section_id,
-            jurisdictions_searched=jurisdictions_to_search,
-            chunks_hydrated=len(final),
-            pinned=len(pinned),
-            deduped_total=len(deduped),
-        )
         graph_cache[cache_key] = final
         return final
 
     @staticmethod
     def _extract_section_id(query: str) -> str | None:
-        """
-        Extract a section number from natural language query.
-        Handles: "Section 18", "section 18A", "s. 18", "sec 18", "Rule 3"
-        Returns normalized section_id as stored in DB (e.g. "18", "18A").
-        """
         import re
-        section_pattern = re.compile(
-            r'\b(?:section|sec\.?|s\.)\s*(\d+[A-Z]?)\b',
-            re.IGNORECASE,
-        )
-        rule_pattern = re.compile(
-            r'\bRule\s+(\d+[A-Z]?)\b',
-            re.IGNORECASE,
-        )
+
+        section_pattern = re.compile(r"\b(?:section|sec\.?|s\.)\s*(\d+[A-Z]?)\b", re.IGNORECASE)
+        rule_pattern = re.compile(r"\bRule\s+(\d+[A-Z]?)\b", re.IGNORECASE)
         m = section_pattern.search(query)
         if m:
             return m.group(1)

tests/unit/retrieval/test_graph_retriever.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from tests.conftest import _make_rc
+
+
+@pytest.mark.asyncio
+async def test_graph_retriever_retries_transient_driver_failure():
+    from civicsetu.models.enums import Jurisdiction
+    from civicsetu.retrieval.graph_retriever import GraphRetriever
+
+    rc = _make_rc(section_id="18")
+
+    with patch("civicsetu.retrieval.graph_retriever.AsyncSessionLocal") as mock_scls, patch(
+        "civicsetu.retrieval.graph_retriever.VectorStore"
+    ) as mock_vs, patch("civicsetu.retrieval.graph_retriever.GraphStore") as mock_graph_store, patch(
+        "civicsetu.stores.graph_store._reset_driver", new=AsyncMock()
+    ) as reset_driver:
+        mock_session = AsyncMock()
+        mock_scls.return_value.__aenter__ = AsyncMock(return_value=mock_session)
+        mock_scls.return_value.__aexit__ = AsyncMock(return_value=False)
+        mock_vs.get_by_section = AsyncMock(return_value=[rc])
+        mock_graph_store.get_referenced_sections = AsyncMock(side_effect=[
+            RuntimeError("the connection is closed"),
+            [],
+        ])
+        mock_graph_store.get_sections_referencing = AsyncMock(return_value=[])
+        mock_graph_store.get_derived_act_sections = AsyncMock(return_value=[])
+        mock_graph_store.get_deriving_rule_sections = AsyncMock(return_value=[])
+
+        result = await GraphRetriever.retrieve(
+            query="What does Section 18 say?",
+            jurisdiction=Jurisdiction.CENTRAL,
+            depth=1,
+        )
+
+    assert result == [rc]
+    reset_driver.assert_awaited_once()