awslabs / awslabs/graphrag-toolkit
[FEATURE] Support custom/external retrievers in benchmarking harness
- Dominant language
- Python
- Stars
- 442
- Forks
- 106
- Avg merge
- 2d 3h
- Merged PRs (30d)
- 41
Description
### Package
lexical-graph
### Problem statement
Allow users to benchmark their own custom-designed retrievers alongside the built-in ones (traversal, agentic, semantic_guided, etc.) by providing either:
1. A Python module path to a script/class that implements the retriever interface, or
2. A formal interface/protocol that custom retrievers must implement
Related to PR #345 (see :https://github.com/awslabs/graphrag-toolkit/pull/345#discussion_r3462744491) by @aghassel.
### Proposed solution
1. Define a CustomRetriever protocol/interface — A minimal contract requiring a query(question: str) -> Response method (matching what run_benchmark_query expects), plus an optional get_config() -> dict for recording hyperparameters in metrics_summary.json.
2. Accept a custom: retriever ID — e.g. BENCHMARK_RETRIEVER=custom:my_module.MyRetriever. The factory would dynamically import the class and instantiate it with graph_store and vector_store.
3. Record custom retriever config — get_retriever_config() should handle the custom case by calling the retriever's get_config() if available, or recording the module path and any constructor kwargs.
### Alternatives considered
Acceptance Criteria
- [ ] A documented protocol/ABC that custom retrievers must satisfy
- [ ] retriever_factory.py supports a custom: (or similar) retriever ID
- [ ] Custom retriever hyperparameters are recorded in metrics_summary.json
- [ ] Example showing how to write and benchmark a custom retriever
- [ ] Existing built-in retrievers continue to work unchanged
Contributor guide
Research direction
Start by tracing retriever_factory.py, run_benchmark_query, and get_retriever_config() to understand the existing retriever contract and metrics flow. Define the documented custom-retriever interface, dynamic module loading, configuration recording in metrics_summary.json, and an example; verify that built-in retrievers still work and the acceptance criteria are covered.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, tooling
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100