From 7b121c513ede1555abe3cf2b557b953550b6f1bf Mon Sep 17 00:00:00 2001 From: noel-improv Date: Wed, 22 Jul 2026 10:20:11 -0600 Subject: [PATCH 1/2] feat(integration-tests): add S3 doc-store option to WikiHow extraction benchmark benchmark_extract hardcoded FileBasedDocs (local disk), so the benchmark could not measure S3-backed document storage. Add a BENCHMARK_DOC_STORE env switch: "file" (default, unchanged behaviour) keeps FileBasedDocs; "s3" uses S3BasedDocs, with BENCHMARK_S3_JSONL selecting the per-chunk or JSONL write path. S3BasedDocs uses collection_id=None (timestamp per run) so repeated runs do not accumulate into one collection. A log line records the resolved store, flag, and collection_id. build-tests.sh propagates the two new variables into .env.testing so they reach the notebook; without this they stay local to the caller's shell and the run silently falls back to FileBasedDocs. --- integration-tests/build-tests.sh | 6 +++++ .../benchmark_extract.py | 24 +++++++++++++++---- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/integration-tests/build-tests.sh b/integration-tests/build-tests.sh index a17cca6a..73fc8a2a 100644 --- a/integration-tests/build-tests.sh +++ b/integration-tests/build-tests.sh @@ -378,6 +378,12 @@ fi if [[ "$BENCHMARK_IS_PROTOTYPE" ]]; then echo "export BENCHMARK_IS_PROTOTYPE=$BENCHMARK_IS_PROTOTYPE" >> lexical-graph-examples/.env.testing fi +if [[ "$BENCHMARK_DOC_STORE" ]]; then + echo "export BENCHMARK_DOC_STORE=$BENCHMARK_DOC_STORE" >> lexical-graph-examples/.env.testing +fi +if [[ "$BENCHMARK_S3_JSONL" ]]; then + echo "export BENCHMARK_S3_JSONL=$BENCHMARK_S3_JSONL" >> lexical-graph-examples/.env.testing +fi zip -r graphrag-toolkit.zip graphrag-toolkit # zip under directory diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_extract.py b/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_extract.py index 510dd054..e502722c 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_extract.py +++ b/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_extract.py @@ -14,7 +14,7 @@ from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory from graphrag_toolkit.lexical_graph.storage import VectorStoreFactory from graphrag_toolkit.lexical_graph.storage.graph import NonRedactedGraphQueryLogFormatting -from graphrag_toolkit.lexical_graph.indexing.load import FileBasedDocs +from graphrag_toolkit.lexical_graph.indexing.load import FileBasedDocs, S3BasedDocs from graphrag_toolkit.lexical_graph.indexing.extract import BatchConfig from llama_index.core import SimpleDirectoryReader @@ -69,9 +69,25 @@ def run_benchmark_extract(handler: IntegrationTestHandler, ) indexing_config = IndexingConfig(batch_config=batch_config) - extracted_docs = FileBasedDocs( - docs_directory=os.path.join(data_dir, dataset_name, 'extracted'), - collection_id=dataset_name + doc_store = os.environ.get('BENCHMARK_DOC_STORE', 'file').lower() + if doc_store == 's3': + extracted_docs = S3BasedDocs( + region=os.environ['AWS_REGION_NAME'], + bucket_name=os.environ['S3_RESULTS_BUCKET'], + key_prefix=f'{os.environ["S3_RESULTS_PREFIX"]}/doc-store/{dataset_name}', + collection_id=None, + for_jsonl=os.environ.get('BENCHMARK_S3_JSONL', 'false').lower() == 'true' + ) + else: + extracted_docs = FileBasedDocs( + docs_directory=os.path.join(data_dir, dataset_name, 'extracted'), + collection_id=dataset_name + ) + + logger.info( + f'Doc store: {type(extracted_docs).__name__} ' + f'(for_jsonl={getattr(extracted_docs, "for_jsonl", "n/a")}) ' + f'collection_id={extracted_docs.collection_id}' ) with ( From 7042904e6c615b72e3d548fa6f8921b4c4b1a2ac Mon Sep 17 00:00:00 2001 From: noel-improv Date: Fri, 24 Jul 2026 15:25:01 -0600 Subject: [PATCH 2/2] fix(integration-tests): validate S3 env vars before S3 doc-store benchmark BENCHMARK_DOC_STORE=s3 read AWS_REGION_NAME/S3_RESULTS_BUCKET/S3_RESULTS_PREFIX via plain os.environ[...]. In non-batch (prototype) runs those vars are not otherwise required, so a missing one raised a bare KeyError and aborted the run before extraction. Check them up front and raise a clear ValueError naming the missing variables. --- .../graphrag_toolkit_tests/benchmark_extract.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_extract.py b/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_extract.py index e502722c..d618d9a5 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_extract.py +++ b/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_extract.py @@ -71,6 +71,14 @@ def run_benchmark_extract(handler: IntegrationTestHandler, doc_store = os.environ.get('BENCHMARK_DOC_STORE', 'file').lower() if doc_store == 's3': + missing = [ + var for var in ('AWS_REGION_NAME', 'S3_RESULTS_BUCKET', 'S3_RESULTS_PREFIX') + if not os.environ.get(var) + ] + if missing: + raise ValueError( + f"BENCHMARK_DOC_STORE=s3 requires {', '.join(missing)} to be set" + ) extracted_docs = S3BasedDocs( region=os.environ['AWS_REGION_NAME'], bucket_name=os.environ['S3_RESULTS_BUCKET'],