mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 14:17:08 +00:00
fix(graphrag): stop losing chunks silently during a graph build
Three ways a chunk disappeared from a graph with no way to tell: A build checks out one connection and then spends minutes per chunk waiting on the model, so the connection idles long enough for the server or a pooler to drop it. The pool only validates a connection when it hands one out, and this one was handed out at the start of the build, so the next write raised "the connection is lost", the chunk was marked failed, and the build carried on a chunk short. apply_chunk and mark_chunk now reconnect and retry once; every statement they run is an idempotent upsert, so a replay cannot double-write. Only connection loss retries — a bad statement still surfaces. An unparseable model response marked the chunk failed and logged nothing at all, so failed_chunks was the only evidence and it named no chunk. Both failure modes now log the chunk id. The summary's node count summed per-chunk upserts, so an entity appearing in ten chunks counted ten times: it reported writes, not graph size. It now reports the distinct node count, falling back to the write count only if the count query fails.
This commit is contained in:
1 parent
92e19ac177
commit
e3d819d9fd
4 files changed
+379
-73
No files matched your search
@@ -890,3 +890,115 @@ class TestCountNodesMany:
|
||||
store, _, _ = self._store_with_mock_conn([(source_id.lower(), 3)])
|
||||
|
||||
assert store.count_nodes_many([source_id]) == {source_id: 3}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestWritesSurviveALostConnection:
|
||||
"""A graph build holds one pooled connection across its LLM calls.
|
||||
|
||||
Extraction spends minutes per chunk waiting on a model, so the connection
|
||||
sits idle between writes and the server (or a pooler) can drop it. The pool
|
||||
only validates a connection at checkout, and this one was checked out once
|
||||
at the start of the build, so the next write raises and the chunk is marked
|
||||
``failed`` — silently losing it from the graph. The write reconnects and
|
||||
retries once instead; the statements are idempotent upserts, so a retry
|
||||
cannot double-write.
|
||||
"""
|
||||
|
||||
def _store_with_connections(self, conns):
|
||||
"""Store that hands out ``conns`` in order, one per (re)connect."""
|
||||
store = GraphStore.__new__(GraphStore)
|
||||
store._tables_ensured = True
|
||||
store._connection = None
|
||||
handed = []
|
||||
closed = []
|
||||
|
||||
def _get_connection():
|
||||
if store._connection is None:
|
||||
store._connection = conns[len(handed)]
|
||||
handed.append(store._connection)
|
||||
return store._connection
|
||||
|
||||
def _close():
|
||||
if store._connection is not None:
|
||||
closed.append(store._connection)
|
||||
store._connection = None
|
||||
|
||||
store._get_connection = _get_connection
|
||||
store.close = _close
|
||||
return store, handed, closed
|
||||
|
||||
@staticmethod
|
||||
def _conn(execute_error=None):
|
||||
cursor = MagicMock()
|
||||
cursor.fetchone.return_value = [str(uuid.uuid4())]
|
||||
cursor.fetchall.return_value = []
|
||||
if execute_error is not None:
|
||||
cursor.execute.side_effect = execute_error
|
||||
conn = MagicMock()
|
||||
conn.cursor.return_value = cursor
|
||||
return conn
|
||||
|
||||
def test_mark_chunk_retries_on_a_dropped_connection(self):
|
||||
import psycopg
|
||||
|
||||
dead = self._conn(psycopg.OperationalError("the connection is lost"))
|
||||
alive = self._conn()
|
||||
store, handed, closed = self._store_with_connections([dead, alive])
|
||||
|
||||
store.mark_chunk(str(uuid.uuid4()), "c1", "done")
|
||||
|
||||
assert handed == [dead, alive]
|
||||
assert closed == [dead]
|
||||
alive.commit.assert_called_once()
|
||||
|
||||
def test_apply_chunk_retries_on_a_dropped_connection(self):
|
||||
import psycopg
|
||||
|
||||
dead = self._conn(psycopg.OperationalError("the connection is lost"))
|
||||
alive = self._conn()
|
||||
store, handed, closed = self._store_with_connections([dead, alive])
|
||||
entities = [
|
||||
{
|
||||
"name": "Ada",
|
||||
"normalized_name": "ada",
|
||||
"type": "person",
|
||||
"description": "d",
|
||||
}
|
||||
]
|
||||
|
||||
nodes, edges = store.apply_chunk(
|
||||
str(uuid.uuid4()), "c1", entities, [], {"ada": _embedding(0.5)}
|
||||
)
|
||||
|
||||
assert (nodes, edges) == (1, 0)
|
||||
assert handed == [dead, alive]
|
||||
assert closed == [dead]
|
||||
alive.commit.assert_called_once()
|
||||
|
||||
def test_a_second_connection_failure_is_not_retried_again(self):
|
||||
"""One retry, not a loop: a genuinely unreachable DB still fails."""
|
||||
import psycopg
|
||||
|
||||
dead = self._conn(psycopg.OperationalError("the connection is lost"))
|
||||
also_dead = self._conn(psycopg.OperationalError("the connection is lost"))
|
||||
store, handed, _ = self._store_with_connections([dead, also_dead])
|
||||
|
||||
with pytest.raises(psycopg.OperationalError):
|
||||
store.mark_chunk(str(uuid.uuid4()), "c1", "done")
|
||||
|
||||
assert handed == [dead, also_dead]
|
||||
|
||||
def test_a_query_error_is_not_retried(self):
|
||||
"""Only connection loss is retryable; a bad statement must surface."""
|
||||
import psycopg
|
||||
|
||||
broken = self._conn(psycopg.ProgrammingError("syntax error"))
|
||||
spare = self._conn()
|
||||
store, handed, _ = self._store_with_connections([broken, spare])
|
||||
|
||||
with pytest.raises(psycopg.ProgrammingError):
|
||||
store.mark_chunk(str(uuid.uuid4()), "c1", "done")
|
||||
|
||||
assert handed == [broken]
|
||||
broken.rollback.assert_called_once()
|
||||
Reference in new issue
Block a user