diff --git a/lexical-graph/tests/unit/conftest.py b/lexical-graph/tests/unit/conftest.py index 444ae1ef..049289da 100644 --- a/lexical-graph/tests/unit/conftest.py +++ b/lexical-graph/tests/unit/conftest.py @@ -3,8 +3,10 @@ import pytest from unittest.mock import Mock, patch, MagicMock +from llama_index.core.schema import NodeRelationship, RelatedNodeInfo, TextNode from graphrag_toolkit.lexical_graph.tenant_id import TenantId from graphrag_toolkit.lexical_graph.indexing.id_generator import IdGenerator +from graphrag_toolkit.lexical_graph.indexing.load.s3_based_docs import S3DocDownloader @pytest.fixture @@ -75,3 +77,46 @@ def custom_id_gen(custom_tenant): Fixture for custom ID generator (backward compatible mode, no delimiter). ''' return IdGenerator(tenant_id=custom_tenant, include_classification_in_entity_id=True, use_chunk_id_delimiter=False) + + +@pytest.fixture +def chunk_node(): + ''' + Factory for a chunk node belonging to a source document. + + The SOURCE relationship is what SourceDocument.source_id() reads, so a node + built without it has no document identity. + ''' + def _chunk_node(node_id, source_id): + node = TextNode(text=f'text for {node_id}', id_=node_id) + node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(node_id=source_id) + return node + + return _chunk_node + + +@pytest.fixture +def download_source_prefix(): + ''' + Factory that reads one source prefix through S3DocDownloader against a mock + client, given a mapping of object key to the nodes that object holds. + + Every object under the prefix merges into one SourceDocument, so this is how + a test observes what a prefix reads back as. + ''' + def _download_source_prefix(objects): + downloader = S3DocDownloader( + key_prefix='p', collection_id='c', bucket_name='b', fn=lambda n: n + ) + s3_client = Mock() + s3_client.get_paginator.return_value.paginate.return_value = [ + {'Contents': [{'Key': key} for key in objects]} + ] + + def download_fileobj(bucket, key, stream): + stream.write('\n'.join(n.to_json() for n in objects[key]).encode('UTF-8')) + + s3_client.download_fileobj.side_effect = download_fileobj + return downloader._download_doc('prefix', s3_client) + + return _download_source_prefix diff --git a/lexical-graph/tests/unit/indexing/load/test_s3_based_docs.py b/lexical-graph/tests/unit/indexing/load/test_s3_based_docs.py index d695401a..f18d27ca 100644 --- a/lexical-graph/tests/unit/indexing/load/test_s3_based_docs.py +++ b/lexical-graph/tests/unit/indexing/load/test_s3_based_docs.py @@ -1023,40 +1023,22 @@ class TestDownloadDeduplicatesNodeIds: the same node id can arrive twice. """ - def _download(self, objects): - downloader = S3DocDownloader( - key_prefix='p', collection_id='c', bucket_name='b', fn=lambda n: n - ) - s3_client = Mock() - s3_client.get_paginator.return_value.paginate.return_value = [ - {'Contents': [{'Key': key} for key in objects]} - ] - - def download_fileobj(bucket, key, stream): - stream.write('\n'.join(n.to_json() for n in objects[key]).encode('UTF-8')) + SOURCE_ID = 'aws::dead:beef' - s3_client.download_fileobj.side_effect = download_fileobj - return downloader._download_doc('prefix', s3_client) + def test_a_node_id_in_two_objects_is_returned_once(self, download_source_prefix, chunk_node): + node = lambda node_id: chunk_node(node_id, self.SOURCE_ID) - def _node(self, node_id): - node = TextNode(text=f'text for {node_id}', id_=node_id) - node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(node_id='aws::dead:beef') - return node - - def test_a_node_id_in_two_objects_is_returned_once(self): - overlapping = { - 'round-a.jsonl': [self._node('c1'), self._node('c2')], - 'round-b.jsonl': [self._node('c2'), self._node('c3')], - } - - doc = self._download(overlapping) + doc = download_source_prefix({ + 'round-a.jsonl': [node('c1'), node('c2')], + 'round-b.jsonl': [node('c2'), node('c3')], + }) assert [n.node_id for n in doc.nodes] == ['c1', 'c2', 'c3'] - def test_distinct_objects_are_all_kept(self): - doc = self._download({ - 'a.jsonl': [self._node('c1')], - 'b.jsonl': [self._node('c2')], + def test_distinct_objects_are_all_kept(self, download_source_prefix, chunk_node): + doc = download_source_prefix({ + 'a.jsonl': [chunk_node('c1', self.SOURCE_ID)], + 'b.jsonl': [chunk_node('c2', self.SOURCE_ID)], }) assert len(doc.nodes) == 2 diff --git a/lexical-graph/tests/unit/indexing/test_source_id_collision.py b/lexical-graph/tests/unit/indexing/test_source_id_collision.py new file mode 100644 index 00000000..80c6fd7d --- /dev/null +++ b/lexical-graph/tests/unit/indexing/test_source_id_collision.py @@ -0,0 +1,136 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Two distinct documents whose source ids collide are indistinguishable downstream. + +The pair below was found by hashing sequentially numbered documents until two +shared the first eight characters of their md5 digest, which took 30,059 of them. +That is the problem in one number: the default width discriminates on 32 bits, so +a corpus reaches even odds of a collision far below the scale anyone designs for. + +`IdRewriter` passes `''` for the metadata component when a node carries no +metadata, which makes that component constant and leaves only the text digest to +separate two documents. That is the case measured in the collision spike and the +case these tests use. +""" + +import pytest + +from graphrag_toolkit.lexical_graph.indexing.id_generator import IdGenerator + +# md5(TEXT_A) and md5(TEXT_B) agree on their first eight hex characters, a4439cdb. +TEXT_A = 'document 27347 body text' +TEXT_B = 'document 30059 body text' + +COLLIDING_SOURCE_ID = 'aws::a4439cdb:d41d' + +NO_METADATA = '' + +# The width every existing graph was written with, pinned rather than read from +# config so the facts below hold whatever a run is configured to use. +LEGACY_WIDTH = 8 + +# The width under consideration for the fix. 64 discriminating bits. +CANDIDATE_WIDTH = 16 + + +def source_id_at(text, width, metadata_str=NO_METADATA): + """The id this text would get at an explicit width.""" + return IdGenerator(source_id_hash_length=width).create_source_id(text, metadata_str) + + +def source_id_as_configured(text, metadata_str=NO_METADATA): + """The id this text gets at whatever width the run is configured to use.""" + return IdGenerator().create_source_id(text, metadata_str) + + +class TestCollidingPair: + """ + Preconditions. These are arithmetic about md5, not statements about any + configuration. If one stops holding, the pair needs regenerating. + """ + + def test_the_two_documents_are_different(self): + # Guards the rest: every assertion below is worthless if these converge. + assert TEXT_A != TEXT_B + + def test_they_share_one_source_id_at_the_legacy_width(self): + assert (source_id_at(TEXT_A, LEGACY_WIDTH) + == source_id_at(TEXT_B, LEGACY_WIDTH) + == COLLIDING_SOURCE_ID) + + def test_a_wider_text_digest_separates_them(self): + assert source_id_at(TEXT_A, CANDIDATE_WIDTH) != source_id_at(TEXT_B, CANDIDATE_WIDTH) + + def test_metadata_separates_them_only_when_it_differs(self): + # The second component is a digest of the metadata, so it discriminates + # only across documents whose metadata is not identical. A corpus loaded + # without metadata, or with the same metadata throughout, gets no help + # from it however wide it is. + shared = 'file_path:corpus.txt' + assert (source_id_at(TEXT_A, LEGACY_WIDTH, shared) + == source_id_at(TEXT_B, LEGACY_WIDTH, shared)) + assert (source_id_at(TEXT_A, LEGACY_WIDTH, 'file_path:a.txt') + != source_id_at(TEXT_B, LEGACY_WIDTH, 'file_path:b.txt')) + + +class TestCollisionConsequences: + """ + What the shared id costs at the storage layer. These pass today: they + characterise the damage rather than assert the fix. + """ + + def test_one_prefix_reads_back_as_one_document_holding_both( + self, download_source_prefix, chunk_node + ): + # The S3 prefix is the bare source id, so a shared id is a shared prefix + # and each document's object lands beside the other's. + doc = download_source_prefix({ + f'{COLLIDING_SOURCE_ID}-aaaaa.jsonl': [chunk_node('a1', COLLIDING_SOURCE_ID)], + f'{COLLIDING_SOURCE_ID}-bbbbb.jsonl': [chunk_node('b1', COLLIDING_SOURCE_ID)], + }) + + # Two documents went in; one comes out, carrying a chunk from each. + assert {n.node_id for n in doc.nodes} == {'a1', 'b1'} + assert doc.source_id() == COLLIDING_SOURCE_ID + + def test_nothing_reports_the_collision(self, download_source_prefix, chunk_node): + # No error, no warning, no marker. A reader cannot tell this document + # from one that genuinely had two chunks, which is what makes the + # failure silent rather than something a run surfaces. + doc = download_source_prefix({ + f'{COLLIDING_SOURCE_ID}-aaaaa.jsonl': [chunk_node('a1', COLLIDING_SOURCE_ID)], + f'{COLLIDING_SOURCE_ID}-bbbbb.jsonl': [chunk_node('b1', COLLIDING_SOURCE_ID)], + }) + + assert len(doc.nodes) == 2 + + +@pytest.mark.xfail( + strict=True, + reason='Source ids discriminate on 32 bits at the default width. Remove this ' + 'marker when the default is widened; strict=True fails the run once the ' + 'assertions start passing, so the marker cannot outlive the fix.', +) +class TestSourceIdUniqueness: + """ + RED. Two distinct documents must be distinguishable by id alone, because + every downstream identity derives from it: the S3 prefix above, the + `__Source__` node the graph MERGEs on, and every chunk, topic, statement + and fact id. + + These read the configured width rather than a pinned one, so they assert + that the default is wide enough rather than anything about a given width. + """ + + def test_distinct_documents_get_distinct_source_ids(self): + assert source_id_as_configured(TEXT_A) != source_id_as_configured(TEXT_B) + + def test_distinct_documents_get_distinct_chunk_id_prefixes(self): + generator = IdGenerator() + + chunk_a = generator.create_chunk_id(source_id_as_configured(TEXT_A), TEXT_A, NO_METADATA) + chunk_b = generator.create_chunk_id(source_id_as_configured(TEXT_B), TEXT_B, NO_METADATA) + + assert chunk_a.rsplit(':', 1)[0] != chunk_b.rsplit(':', 1)[0]