Architecture Overview
Relevant source files
The following files were used as context for generating this wiki page:
The Agentic GraphRAG system for Medical Diagnosis is designed to ingest unstructured medical literature and transform it into a structured, hierarchical knowledge graph. This graph serves as the foundation for an agentic reasoning loop that performs multi-strategy retrieval to answer complex clinical questions with high precision and grounding.
The architecture is divided into two primary data flows: the Ingestion Pipeline, which handles knowledge extraction and graph construction, and the Retrieval & Agentic Loop, which manages user queries through a plan-research-verify cycle.
Sources: README.md:21-40, am_diag/ingestion/extraction_pipeline.py:1-12
Knowledge Representation & Schema
The system utilizes a Schema-driven Knowledge Graph built in Neo4j. The schema is a first-class runtime value injected into every component, including LLM extractors via BAML dynamic enums. This ensures that every node and relationship adheres to predefined medical types.
Graph Data Model
The graph consists of five primary node labels and a specialized relationship model:
- Node Labels:
Document,Chunk,Entity,Community,CommunityReport. - Edge Types:
PART_OF,NEXT_CHUNK,HAS_ENTITY,RELATES_TO,IN_COMMUNITY,PARENT_COMMUNITY,HAS_REPORT.
Sources: README.md:46-56, am_diag/db/graph/base.py:155-185, am_diag/common/data_models/schema.py:112-126
Medical Schema Components
The system defines 13 default entity types and 25 relation types grouped by clinical role (Clinical, Pharmacological, Molecular, Structural).
| Component | Description | Identity Key |
|---|---|---|
| EntityType | Defines medical categories (e.g., Disease, Drug). | (label,) |
| RelationType | Defines valid triples (e.g., Disease TREATED_BY Drug). | (type,) |
| GraphSchema | Runtime object managing DDL and prompt generation. | N/A |
Sources: am_diag/common/data_models/schema.py:17-40, README.md:104-123
Ingestion Pipeline
The ingestion process is implemented as a LangGraph StateGraph, moving from raw text chunks to a resolved and summarized knowledge graph.
Ingestion Flow Diagram
The following diagram illustrates the linear progression of data through the ingestion nodes.
Explanation: The ingestion pipeline aggregates extracted entities and relations before resolving them into a canonical graph structure.
Sources: am_diag/ingestion/extraction_pipeline.py:15-115, README.md:131-155
Core Extraction Strategies
The system uses a Three-extractor fusion approach:
- GLiNER: Batch zero-shot Named Entity Recognition.
- GLiREL: Zero-shot relation extraction based on GLiNER spans.
- LLM (BAML): Type-safe extraction using dynamic enums to constrain output to the schema.
Sources: README.md:134-142, am_diag/llm/baml_client/inlinedbaml.py:270-305
Retrieval & Search Engine
The retrieval architecture utilizes layered retrievers and search "recipes" to gather evidence from multiple indices.
Vector Collections
Data is vectorized across four distinct collections to support semantic search:
- Entity: Embedded by canonical name.
- Relation: Embedded as Edge facts, Edge types, and Full SPO (Subject-Predicate-Object).
- Chunk: Embedded raw text.
- CommunityReport: Embedded report summaries.
Sources: README.md:166-175
Retrieval Recipes
Retrieval is orchestrated by a SearchEngine that applies specific fusion strategies, such as Reciprocal Rank Fusion (RRF) or Cross-encoders.
| Recipe | Methods Used | Reranker |
|---|---|---|
hybrid_rrf | Entity + Relation + Chunk + Community | RRF |
text2cypher | LLM-generated Cypher queries | RRF |
bfs_expand | Entity + Degree-aware BFS | RRF |
Sources: README.md:195-207, am_diag/llm/baml_client/inlinedbaml.py:465-495
Agentic Reasoning Loop
The system employs a multi-agent "Plan-Research-Verify" loop managed by the DeepAgents harness.
Agent Interactions Sequence
Explanation: The Orchestrator manages sub-agents to ensure the final answer is fully grounded in retrieved evidence.
Sources: README.md:213-228, am_diag/agents/prompts.py:10-55
Sub-agent Responsibilities
- Planner: Decomposes complex questions into recipes (e.g., mapping a symptom check to
entitysearch). - Researcher: Executes tools like
text2cypherorcommunity_map_reduce. - Verifier: Performs deterministic sufficiency checks based on coverage and evidence depth.
Sources: am_diag/agents/prompts.py:27-65, README.md:215-223
Database Implementation (Neo4j)
The Neo4jClient implements the GraphStore interface, abstracting Cypher execution. It enforces an ordering guarantee during writes: Constraints/Indexes -> Nodes -> Relationships.
DDL and Node Upserts
The client renders Cypher templates to handle dynamic secondary labels (e.g., (:Entity:Drug)) which cannot be passed as standard parameters.
# Sources: [am_diag/db/graph/neo4j.py:214-228]
async def upsert_entities(self, entities: list[Entity]) -> None:
by_label: dict[str, list[dict[str, Any]]] = {}
for e in entities:
by_label.setdefault(e.label, []).append(entity_record(e))
for label, records in by_label.items():
cypher = render("ingest/upsert_entities", Label=label)
await self.execute_batch(cypher, records)
Sources: am_diag/db/graph/neo4j.py:20-40, am_diag/db/graph/base.py:180-195
Summary
The Agentic Med-Diag architecture provides a robust framework for medical knowledge management. By combining a schema-constrained ingestion pipeline with a multi-agent reasoning loop, the system ensures that diagnostic answers are not only semantically relevant but also structurally grounded in medical literature and verified for completeness.