awslabs / awslabs/graphrag-toolkit

Non-batch fallback in BatchLLMPropositionExtractorSync/BatchTopicExtractorSync ignores configured LLM

Open Beginner friendly
#534 0 comments 0 reactions 0 assignees View on GitHub
triage_needed
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.