awslabs / awslabs/graphrag-toolkit

[FEATURE] Composable extraction pipeline with stage registry and typed contracts

Open
#240 1 comment 0 reactions 1 assignee Claimed by @mykola-pereyma View on GitHub
enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.