awslabs / awslabs/graphrag-toolkit
[FEATURE] Composable extraction pipeline with stage registry and typed contracts
- Dominant language
- Python
- Stars
- 442
- Forks
- 106
- Avg merge
- 2d 3h
- Merged PRs (30d)
- 41
Description
## Summary
The current extraction pipeline includes two LLM stages: extract propositions, and extract topics/statements/facts/entities. Users can lightly guide the topic/statement/fact/entity extraction. They can also replace the default prompts with their own domain-specific prompts.
While the output of the pipeline remains the same – a JSON representation of extracted topics, statements, facts and entities – we should allow for a more composable set of stages that can mix local extraction (using e.g. CPU-based NER extraction models such as GLiNER) and LLM-based extraction, with more configuration and filtering options that allow for supplying strict schema or ontology.
## Proposed Solution
Introduce a stage registry with typed contracts (Option 2 from design doc):
- `ExtractionStage` ABC with `input_keys()`/`output_keys()`/`as_transform()` contracts
- `PipelineBuilder` with build-time validation of stage compatibility
- `ExtractionConfig.from_stages()` factory method for custom pipelines
- `ExtractionSchema` + `EntityTypeConfig` for ontology constraints
- Built-in stages: LLMProposition, LocalProposition, LLMTopicExtraction, NERExtraction, EntityMerge, SchemaFilter, batch variants
## Key Requirements
- Pipeline output must remain TopicCollection format (backward compatible)
- `ExtractionConfig()` without stages preserves existing behavior
- Build-time validation catches wiring errors before runtime
- Support mixing local (CPU-based NER) and LLM-based extraction
- Schema/ontology can guide LLM extraction AND filter results post-extraction
## Acceptance Criteria
- [ ] ExtractionStage ABC with typed input/output key contracts
- [ ] PipelineBuilder validates stage compatibility at build time
- [ ] built-in stages implemented and tested
- [ ] ExtractionConfig.from_stages() factory method
- [ ] Schema auto-injection into LLM topic extraction prompt
- [ ] All existing tests pass (no regression)
- [ ] Extraction-specific tests
- [ ] Usage documentation and example notebook
Contributor guide
Assessment
This issue has not been assessed yet.