Retrieval Recipes & Reranking
Relevant source files
The following files were used as context for generating this wiki page:
Retrieval Recipes & Reranking represent the layered retrieval architecture of the am_diag system. This module provides a flexible, data-driven framework for executing multi-strategy searches across various data collections, including medical entities, relations, text chunks, and community reports. By decoupling the retrieval logic into "recipes," the system allows for complex fan-out search strategies and sophisticated results fusion without requiring changes to the core orchestration code.
The primary entry point for this system is the SearchEngine, which coordinates retrievers for different data types and applies reranking strategies such as Reciprocal Rank Fusion (RRF) or Cross-Encoders to produce a final, ranked list of SearchResult items.
Sources: README.md:162-167, am_diag/retrieval/search.py:27-40
Architecture and Data Flow
The retrieval architecture follows a fan-out and fuse pattern. When a query is received, the SearchEngine identifies the requested recipe, executes multiple retrieval methods concurrently, and then passes the aggregated results to a reranker.
Retrieval Pipeline Flow
The diagram below illustrates the sequence from a natural language query to the final ranked results.
The SearchEngine manages the concurrent execution of these tasks using asyncio.gather.
Sources: am_diag/retrieval/search.py:88-124, README.md:162-175
Retrieval Recipes
Retrieval Recipes are defined as SearchConfig constants. They specify which retrieval methods to use, the reranking strategy, and parameters for graph traversal (BFS).
Core Recipe Definitions
The system includes several predefined recipes catering to different clinical question types:
| Recipe Name | Methods Included | Reranker | BFS | Description |
|---|---|---|---|---|
entity | Entity | RRF | No | Direct lookup for medical entities. |
relation | Relation | RRF | No | Search for predicates (treats, causes). |
chunk | Chunk | RRF | No | Passage-level semantic search. |
community | Community | RRF | No | Searches thematic community report summaries. |
hybrid_rrf | All (Entity, Relation, Chunk, Community) | RRF | No | General-purpose fused search. |
hybrid_cross_encoder | All | Cross-Encoder | Yes | High-precision search with graph expansion. |
bfs_expand | Entity | RRF | Yes | Seed entity search followed by BFS expansion. |
text2cypher | Text2Cypher | RRF | No | LLM-generated graph queries. |
Sources: am_diag/retrieval/recipes.py:11-66, am_diag/retrieval/search.py:157-175, tests/am_diag/retrieval/test_recipes.py:11-45
Data Structures
The SearchConfig class is the primary data model for defining these recipes.
# am_diag/retrieval/config.py (referenced in recipes.py)
class SearchConfig:
methods: list[str]
reranker: str = "rrf"
limit: int = 20
bfs: bool = False
bfs_depth: int = 2
Sources: am_diag/retrieval/recipes.py:8-10, am_diag/retrieval/search.py:75-80
Reranking Strategies
Once raw results are gathered from multiple sources, they must be unified into a single ranked list. The system supports two primary fusion mechanisms.
Reciprocal Rank Fusion (RRF)
RRF is the default strategy for most recipes. It calculates a score based on the rank of an item across multiple result sets, allowing for the combination of results with incompatible original scores (e.g., combining BM25 and vector similarity). The rrf_k parameter in RetrievalConfig controls the ranking sensitivity.
Sources: am_diag/retrieval/search.py:133-134, README.md:177-184
Cross-Encoder Reranking
The hybrid_cross_encoder recipe utilizes a secondary LLM or transformer-based reranker. Unlike RRF, which relies on prior ranks, the Cross-Encoder re-evaluates the query-document pairs to provide a more nuanced relevance score.
Sources: am_diag/retrieval/search.py:130-132, README.md:183-184
Result Filtering
After fusion, the SearchEngine applies a minimum score threshold defined by reranker_min_score. Any results falling below this threshold are discarded before the final list is truncated to the recipe's limit.
Sources: am_diag/retrieval/search.py:136-139, tests/am_diag/retrieval/test_search.py:86-101
Integration with Agentic Tools
The retrieval recipes are exposed to the clinical agents through the DeepAgents harness. Each tool in the agentic loop corresponds to a specific retrieval recipe.
Agent Tool Mapping
The make_retrieval_tools function wraps the SearchEngine.search() calls into LangChain-compatible tools.
entity_search→recipes.ENTITYrelation_search→recipes.RELATIONchunk_search→recipes.CHUNKcommunity_search→recipes.COMMUNITYhybrid_search→recipes.HYBRID_RRFcypher_search→recipes.TEXT2CYPHER
Sources: am_diag/agents/tools.py:82-125, tests/am_diag/agents/test_tools.py:33-47
Summary
The Retrieval Recipes & Reranking module provides the "research" capabilities of the medical diagnosis system. By standardizing retrieval into named recipes, the system enables an agentic planner to select the most effective search strategy based on the nature of a clinical sub-question. The combination of concurrent fan-out and sophisticated fusion ensures that the agent receives high-quality, grounded evidence from diverse sources within the knowledge graph and vector store.
Sources: README.md:162-184, am_diag/retrieval/search.py:27-40