Customizing the Medical Schema
Relevant source files
The following files were used as context for generating this wiki page:
The medical schema in this system is a first-class, runtime-injectable value that defines the structure and constraints of the knowledge graph. It acts as the "source of truth" for every component in the pipeline, including extractors (GLiNER, GLiREL, LLM), resolvers, community summarizers, and the reasoning agent. By centralizing entity and relation definitions, the system ensures that all extracted knowledge adheres to a consistent medical taxonomy.
The schema is primarily defined using the GraphSchema class, which holds collections of EntityType and RelationType objects. These definitions include natural-language hints and constraints that steer zero-shot models and LLM prompts, allowing the system to be adapted to different medical domains or specific clinical focus areas without changing the underlying code logic.
Sources: am_diag/common/data_models/schema.py:41-47, README.md:104-108
Core Schema Components
The schema architecture is built on three primary data structures that define the nodes and edges of the medical knowledge graph.
GraphSchema Architecture
The GraphSchema serves as the container for the entire taxonomy. It provides utility methods to transform these definitions into formats required by various pipeline stages, such as labels for GLiNER/GLiREL and prompt blocks for LLMs.
The diagram above shows the relationship between the main schema container and its constituent type definitions. Sources: am_diag/common/data_models/schema.py:12-65
Entity Definitions
Each EntityType defines a specific category of medical concept. It includes a gliner_label, which is a natural-language description used to steer the GLiNER model for zero-shot named entity recognition.
| Field | Type | Description |
|---|---|---|
label | str | The internal unique identifier for the entity type (e.g., "Disease"). |
gliner_label | str | Natural language hint for the GLiNER extractor. |
description | str | Detailed description used in LLM prompts to explain the entity's scope. |
key_property | str | The primary attribute name (e.g., "name", "symbol", "genericName"). |
Sources: am_diag/common/data_models/schema.py:12-23, am_diag/common/schema/medical_graphrag.py:10-120
Relation Definitions
RelationType defines the directed edges between entities. Crucially, it includes head_labels and tail_labels which act as semantic constraints, preventing the extraction of implausible clinical relationships (e.g., a Drug cannot have a Symptom as a child in a HAS_SYMPTOM relation).
Sources: am_diag/common/data_models/schema.py:26-38, README.md:92-95
The Default Medical Schema
The project provides a comprehensive default schema tailored for general clinical diagnosis, consisting of 13 entity types and 25 relation types.
Standard Entity Types
The default schema includes entities ranging from macro-level clinical concepts to molecular targets:
- Clinical:
Disease,Symptom,RiskFactor - Pharmacological:
Drug,DrugClass,MechanismOfAction - Molecular:
Gene,Protein,Pathway - Structural:
AnatomicalStructure - Interventional:
Procedure,DiagnosticTest,Pathogen
Sources: am_diag/common/schema/medical_graphrag.py:10-142, README.md:110-124
Relation Groups
Relations are categorized by their clinical role to provide a structured view of medical knowledge:
| Group | Examples |
|---|---|
| Clinical | HAS_SYMPTOM, TREATED_BY, DIAGNOSED_BY, CAUSED_BY |
| Pharmacological | TARGETS, INHIBITS, METABOLIZED_BY, CONTRAINDICATED_IN |
| Structural/Molecular | ENCODES, PARTICIPATES_IN, PART_OF, IS_A |
Sources: am_diag/common/schema/medical_graphrag.py:146-311, README.md:126-135
Schema Injection and Data Flow
The schema is injected at runtime into various components of the ingestion and retrieval pipelines.
The schema acts as a central configuration that propagates through both the ingestion and reasoning layers. Sources: am_diag/ingestion/extraction_pipeline.py:53-61, README.md:104-108
BAML Dynamic Enums
For LLM-based extraction, the schema is injected into BAML as dynamic enums. This forces the LLM to output only valid EntityType and RelationType values defined in the schema, rather than relying on the model's internal training data.
// From am_diag/llm/baml_client/inlinedbaml.py
enum EntityType { @@dynamic }
enum RelationType { @@dynamic }
class ExtractedEntity {
name string
entity_type EntityType
// ...
}
Sources: am_diag/llm/baml_client/inlinedbaml.py:214-230
Prompt Engineering Integration
The GraphSchema provides methods like entity_prompt_block() and relation_prompt_block() that automatically format the schema for inclusion in system prompts for agents and researchers.
- Planner Agent: Uses the schema to map sub-questions to retrieval recipes.
- Researcher Agent: Uses schema definitions to identify which tools (entity search vs. relation search) are most relevant.
- Text-to-Cypher: Uses the
as_text()representation to understand the property graph structure (e.g., that all relations use the:RELATES_TOedge type with atypeproperty).
Sources: am_diag/common/data_models/schema.py:67-94, am_diag/agents/prompts.py:24-34, am_diag/llm/baml_client/inlinedbaml.py:408-420
Implementation Example: Adding a New Type
To customize the schema, a developer modifies the MEDICAL_GRAPHRAG_SCHEMA or provides a custom GraphSchema instance to the ExtractionPipelineConfig.
# Conceptual customization
new_entity = EntityType(
label="Biomarker",
gliner_label="clinical biomarker or molecular marker",
description="A measurable indicator of a biological state or condition."
)
custom_schema = GraphSchema(
entity_types=ENTITY_TYPES + [new_entity],
relation_types=RELATION_TYPES
)
Sources: am_diag/common/data_models/schema.py:41-47, am_diag/common/schema/medical_graphrag.py:313-316
Summary
Customizing the medical schema is the primary method for adapting the Agentic GraphRAG system to specific medical use cases. By defining EntityType and RelationType objects, developers control the extraction boundaries and semantic consistency of the resulting knowledge graph. The system's architecture ensures that these changes are automatically reflected across extraction, storage, and agentic reasoning components.