fix(graphrag): review fixes — replay-safe writes, strict count, zero weights

Three findings from review, all on code this branch introduced.

apply_chunk was not replay-safe. commit() can report connection loss *after*
Postgres committed, and the reconnect retry then replays the write: _upsert_node
bumps doc_freq a second time and _add_edge inserts another row, since
graph_edges has no uniqueness constraint for a logical edge. The chunk's
graph_ingest_progress row is now written in the same transaction as the rows it
describes, and a replay that finds it already "done" returns (0, 0) without
touching the graph. Extraction drops its separate mark_chunk("done"): the
checkpoint and the graph can no longer disagree.

count_nodes swallows every query failure and answers 0, so extraction's
"fall back to the write count" handler could never run — a failed count after a
successful build reported an empty graph. count_nodes grows a strict mode that
re-raises; retrieval keeps the swallow, which is what routes a source to
ClassicRAG.

A zero edge weight was read as a full-strength link: `or 1.0` rewrote an
explicit 0 before the <= 0 filter. Only missing and null weights default now.
The same coercion sat in _ppr_scores, where it would have kept the ranker's
rule unreachable from the product path, so it is fixed there too.

docsgpt/graphrag/extraction.py

@@ -290,7 +290,9 @@ def extract_graph_for_source(
             )
             node_upserts += chunk_nodes
             edges += chunk_edges
-            store.mark_chunk(source_id, chunk_id, "done")
+            # ``apply_chunk`` marks the chunk done inside the transaction that
+            # writes its rows, so the checkpoint cannot disagree with the graph
+            # and a replayed write cannot apply the chunk twice.
             chunks_processed += 1
         except Exception as exc:
             logger.warning(
@@ -311,9 +313,12 @@ def extract_graph_for_source(
     # upserts and a single node, so the old count overstated every graph whose
     # entities recur. Report what the graph holds, falling back to the write
     # count only if the count query itself fails.
+    # ``strict`` is what makes the fallback below reachable: the default
+    # count swallows query failures and answers 0, which would report a
+    # successful build as an empty graph.
     nodes = node_upserts
     try:
-        nodes = store.count_nodes(source_id)
+        nodes = store.count_nodes(source_id, strict=True)
     except Exception as exc:
         logger.warning(
             "count_nodes failed for source %s; reporting upserts instead: %s",

docsgpt/graphrag/store.py

@@ -556,8 +556,12 @@ class GraphStore:
         ``name_embeddings`` maps ``normalized_name`` to its embedding. Degrees
         are not bumped here — the caller runs ``set_node_degrees`` once at the
         end. Reconnects and retries once if the connection died while the
-        extraction was waiting on the model. Returns
-        ``(nodes_upserted, edges_added)``.
+        extraction was waiting on the model.
+
+        The chunk's ``graph_ingest_progress`` row is written in this same
+        transaction, so the checkpoint and the rows it describes commit
+        together and a replay of an already-applied chunk returns ``(0, 0)``
+        without touching the graph. Returns ``(nodes_upserted, edges_added)``.
         """
         self._ensure_tables_once()
 
@@ -566,6 +570,21 @@ class GraphStore:
             node_ids: Dict[str, str] = {}
             edges_added = 0
             try:
+                # ``commit()`` can report connection loss *after* the server
+                # committed, and the retry then replays this write: doc_freq
+                # would be bumped twice and a second logical edge inserted
+                # (graph_edges has no uniqueness constraint). The progress row
+                # below is written in this transaction, so a replay sees it.
+                cursor.execute(
+                    "SELECT status FROM graph_ingest_progress "
+                    "WHERE source_id = %s AND chunk_id = %s;",
+                    (source_id, str(chunk_id)),
+                )
+                applied = cursor.fetchone()
+                if applied is not None and applied[0] == "done":
+                    conn.rollback()
+                    return 0, 0
+
                 for entity in entities:
                     normalized_name = entity["normalized_name"]
                     node_id = self._upsert_node(
@@ -601,6 +620,15 @@ class GraphStore:
                     )
                     edges_added += 1
 
+                cursor.execute(
+                    """
+                    INSERT INTO graph_ingest_progress (source_id, chunk_id, status)
+                    VALUES (%s, %s, 'done')
+                    ON CONFLICT (source_id, chunk_id)
+                    DO UPDATE SET status = EXCLUDED.status;
+                    """,
+                    (source_id, str(chunk_id)),
+                )
                 conn.commit()
                 return len(entities), edges_added
             except Exception:
@@ -671,8 +699,22 @@ class GraphStore:
             cursor.close()
             conn.rollback()
 
-    def count_nodes(self, source_id: str) -> int:
-        """Number of nodes for a source. Zero drives the ClassicRAG fallback."""
+    def count_nodes(self, source_id: str, strict: bool = False) -> int:
+        """Number of nodes for a source. Zero drives the ClassicRAG fallback.
+
+        Args:
+            source_id: Source whose nodes to count.
+            strict: Re-raise a query failure instead of reporting ``0``.
+                Retrieval wants the swallow — a broken count there just routes
+                the source to ClassicRAG — but a caller reporting how big a
+                graph is must not read a failed query as "the graph is empty".
+
+        Returns:
+            int: The node count, or ``0`` when a query failure is swallowed.
+
+        Raises:
+            Exception: The underlying query failure, when ``strict`` is set.
+        """
         conn = self._get_connection()
         cursor = conn.cursor()
         try:
@@ -683,6 +725,8 @@ class GraphStore:
             return int(cursor.fetchone()[0])
         except Exception as e:
             logging.error(f"Error counting nodes: {e}")
+            if strict:
+                raise
             return 0
         finally:
             cursor.close()