awslabs / awslabs/graphrag-toolkit
Non-batch fallback in BatchLLMPropositionExtractorSync/BatchTopicExtractorSync ignores configured LLM
- Dominant language
- Python
- Stars
- 442
- Forks
- 106
- Avg merge
- 2d 3h
- Merged PRs (30d)
- 41
Description
## Description
When a node set passed to `BatchLLMPropositionExtractorSync` or `BatchTopicExtractorSync` falls below Bedrock's minimum batch size (100 records), `_process_nodes` in `batch_extractor_base.py` falls back to a non-batch extractor via `_run_non_batch_extractor`. Both implementations construct their fallback extractor without passing the configured `llm`, even though `self.llm` holds it correctly:
```python
# batch_llm_proposition_extractor_sync.py
def _run_non_batch_extractor(self, nodes):
all_nodes = [node for node in nodes]
extractor = LLMPropositionExtractor(
prompt_template=self.prompt_template,
source_metadata_field=self.source_metadata_field,
) # llm= omitted
```
```python
# batch_topic_extractor_sync.py
def _run_non_batch_extractor(self, nodes):
all_nodes = [node for node in nodes]
extractor = TopicExtractor(
prompt_template=self.prompt_template,
source_metadata_field=self.source_metadata_field,
entity_classification_provider=self.entity_classification_provider,
topic_provider=self.topic_provider,
) # llm= omitted
```
Both `LLMPropositionExtractor` and `TopicExtractor` accept `llm` and, when omitted, default to `GraphRAGConfig.extraction_llm`, which resolves to `DEFAULT_EXTRACTION_MODEL = 'us.anthropic.claude-sonnet-4-6'`.
## Impact
- **Hard failure outside a region where the `us.*` profile resolves.** We hit this in `eu-west-2`, with a `BatchConfig`-configured `eu.anthropic.claude-sonnet-4-6` LLM passed explicitly to the batch extractor. A worker whose chunk count landed under 100 fell back to the non-batch path, which then called Bedrock with `us.anthropic.claude-sonnet-4-6`:
```
ValidationException: The provided model identifier is invalid. [Model config: {..., "model": "us.anthropic.claude-sonnet-4-6", ...}]
```
Because the toolkit uses a plain `asyncio.gather` internally with no per-document fault isolation, this aborted our entire extraction run (`processed=0, failed=50`) despite the batch path having already completed successfully for the other 3 worker buckets.
- **Silent behaviour change even where it doesn't error** — any deployment with a non-default `extraction_llm` (different model, `max_tokens`, `temperature`, region, etc.) gets inconsistent behavior between the batch and non-batch paths for the same extraction run, with nothing logged to indicate the switch.
## Relationship to #344 / #347
This produces the same `ValidationException: The provided model identifier is invalid.` message as #344, but the mechanism is different:
- **#344** — `BedrockConverse` is pickled to a `ProcessPoolExecutor` worker; the recreated client loses the LLM's configured `region_name`. Fixed by #347 (pass `region_name=self.llm.region_name` when `LLMCache` recreates the client).
- **This issue** — the configured LLM is never passed to the fallback extractor at all, so `self.llm.region_name` is never even in play; the fallback ignores the caller's LLM outright and falls through to `GraphRAGConfig.extraction_llm`.
The `global.*` inference profile workaround suggested for #344 does not help here, since the fallback discards whatever LLM/profile was configured regardless.
## Reproduction
1. Configure `BatchLLMPropositionExtractorSync`/`BatchTopicExtractorSync` with an explicit `llm` pointing at a Bedrock inference profile only valid outside `us-*` (e.g. an `eu.*` cross-region profile), in a non-`us-*` region.
2. Run extraction with a node set whose per-worker chunk count for at least one `ProcessPoolExecutor` worker lands under `BEDROCK_MIN_BATCH_SIZE` (100).
3. That worker's `_process_nodes` takes the `_run_non_batch_extractor` branch, which calls Bedrock with `GraphRAGConfig.extraction_llm` (`us.anthropic.claude-sonnet-4-6`) instead of the configured LLM.
## Suggested fix
Pass `llm=self.llm` through in both `_run_non_batch_extractor` implementations. PR incoming: fixes this with a one-line change to each file plus regression tests asserting the fallback extractor is constructed with the configured LLM.
## Environment
- `graphrag-lexical-graph==3.19.1` (also present on `main`)
- Region: `eu-west-2`
- Model: `eu.anthropic.claude-sonnet-4-6` (cross-region inference profile) configured via `BatchConfig`
Contributor guide
Research direction
Start by reading _run_non_batch_extractor in batch_llm_proposition_extractor_sync.py and batch_topic_extractor_sync.py, then review the regression tests described in the issue. Reproduce the below-minimum-batch fallback and verify both fallback extractors retain the configured LLM, including its model and regional settings.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- aws, python
- Domain
- ai, backend, cloud
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 88/100