Skip to content

[FIX] Preserve configured LLM in non-batch fallback extractors - #535

Open
gillespied wants to merge 1 commit into
awslabs:mainfrom
gillespied:fix/non-batch-fallback-preserves-configured-llm
Open

[FIX] Preserve configured LLM in non-batch fallback extractors#535
gillespied wants to merge 1 commit into
awslabs:mainfrom
gillespied:fix/non-batch-fallback-preserves-configured-llm

Conversation

@gillespied

Copy link
Copy Markdown

Description

_run_non_batch_extractor in BatchLLMPropositionExtractorSync and BatchTopicExtractorSync builds its fallback extractor without passing the configured llm, so it silently defaults to GraphRAGConfig.extraction_llm (us.anthropic.claude-sonnet-4-6) instead of whatever LLM was explicitly configured.

Fixes #534.

Changes

  • batch_llm_proposition_extractor_sync.py: pass llm=self.llm to the fallback LLMPropositionExtractor
  • batch_topic_extractor_sync.py: pass llm=self.llm to the fallback TopicExtractor
  • Regression tests for both classes asserting the fallback extractor is constructed with the extractor's configured llm

Problem

Related issue: #534

When a node set falls below Bedrock's minimum batch size (100 records), both batch extractors fall back to a non-batch extractor. That fallback ignores self.llm entirely, so it uses GraphRAGConfig.extraction_llm (DEFAULT_EXTRACTION_MODEL = 'us.anthropic.claude-sonnet-4-6') regardless of what LLM was configured. Outside a region where that us.* profile resolves, this raises ValidationException: The provided model identifier is invalid. We hit this in eu-west-2 with an explicitly configured eu.anthropic.claude-sonnet-4-6 LLM.

This is a different mechanism from #344/#347 (pickled BedrockConverse losing region_name in a spawned worker) — here the configured LLM is never passed to the fallback at all, so the global.* inference profile workaround for #344 doesn't apply.

Testing

  • Unit tests added/updated — one regression test per class, patching the target extractor class and asserting it's constructed with llm=<the configured llm>
  • Integration tests added (as appropriate) — not applicable, this only affects extractor construction, no AWS call needed to reproduce
  • Existing tests pass (pytest) — ran tests/unit/indexing/extract/test_batch_llm_proposition_extractor_sync.py and test_batch_topic_extractor_sync.py (11 passed), and the full tests/unit suite (2137 passed; one pre-existing, unrelated error in test_integ_dependency_compatibility.py caused by no pip binary in my venv, reproduced identically on main without this change)
  • Tested manually (describe below)

Manually reproduced against a real eu-west-2 account: a BatchConfig-configured eu.anthropic.claude-sonnet-4-6 extraction, where one ProcessPoolExecutor worker's chunk count landed under 100 and hit this exact fallback path, raising the ValidationException quoted in #534.

Checklist

  • Code follows existing style and conventions
  • License headers present on new files (no new files; existing headers preserved)
  • Documentation updated (if applicable) — not applicable, internal fallback behavior only
  • No breaking changes (or clearly documented) — additive llm= kwarg passthrough only

BatchLLMPropositionExtractorSync and BatchTopicExtractorSync accept an
explicit llm and store it as self.llm. When a node set falls below
Bedrock's minimum batch size (100 records), _run_non_batch_extractor
falls back to LLMPropositionExtractor / TopicExtractor respectively,
but constructs them without passing llm=self.llm.

Both target classes default to GraphRAGConfig.extraction_llm when no
llm is given, which resolves to DEFAULT_EXTRACTION_MODEL
('us.anthropic.claude-sonnet-4-6'). Outside a region where that us.*
inference profile resolves, this raises:

  ValidationException: The provided model identifier is invalid.

Even where it doesn't error, it silently uses a different
model/max_tokens/temperature than configured.

This is a different failure mode from awslabs#344 (LLMCache losing
region_name when BedrockConverse is pickled to a spawned worker) and
isn't fixed by awslabs#347's workaround of using global.* inference
profiles, since here the configured LLM is never passed to the
fallback extractor at all - the fallback ignores it outright.

Fix: pass llm=self.llm through to both fallback extractors.

Added regression tests for both classes asserting the fallback
extractor is constructed with the configured llm.

@noel-improv noel-improv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@gillespied thanks for your contribution so far these looks good. Both extractors take LLMCacheType, so self.llm is the right thing to pass here.

@mykola-pereyma mykola-pereyma left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the contribution! Clean, well-scoped fix — passing llm=self.llm into both non-batch fallbacks correctly preserves the configured inference profile, and the per-class regression tests pin the behavior nicely. Verified these are the only two fallback sinks (no non-sync siblings share the omission). LGTM.

@oussamahansal oussamahansal left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Non-batch fallback in BatchLLMPropositionExtractorSync/BatchTopicExtractorSync ignores configured LLM

4 participants