awslabs / awslabs/graphrag-toolkit

[FEATURE] Support custom/external retrievers in benchmarking harness

Open
#346 0 comments 1 reaction 0 assignees View on GitHub
enhancement
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.