Knowledge Graph Schema & Data Models
Relevant source files
The following files were used as context for generating this wiki page:
The Knowledge Graph (KG) system in this project is built on a schema-driven architecture designed to structure clinical knowledge extracted from medical literature. The schema acts as a "first-class runtime value," meaning it is injected into every stage of the pipeline—including extraction, resolution, summarization, and agentic reasoning—to constrain and guide the Large Language Model (LLM) and local extractors like GLiNER and GLiREL.
The system utilizes a typed property graph model, primarily persisted in Neo4j. It encompasses five node labels (Document, Chunk, Entity, Community, and CommunityReport) and seven specific edge types to represent hierarchical medical relationships and data provenance. Sources: README.md:31-41, am_diag/common/data_models/schema.py:34-37
Graph Architecture & Relationship Model
The project employs a structured hierarchy where raw text is transformed into a web of interconnected medical concepts.
The diagram above illustrates the structural relationships between the primary data models in the system.
Sources: README.md:43-50
Node and Edge Definitions
| Label | Description |
|---|---|
| Document | The primary source of medical information (e.g., a textbook or article). |
| Chunk | A granular segment of text derived from a Document, used for local extraction. |
| Entity | A medical concept (Disease, Drug, etc.) extracted from a Chunk. |
| Community | A cluster of related entities and relations detected via hierarchical Leiden. |
| CommunityReport | An LLM-generated summary of a specific Community. |
All relationships between entities are stored as generic :RELATES_TO edges in the graph database, with the specific medical relationship type stored in a type property. This allows the schema to be updated dynamically without altering the underlying database structure. Sources: README.md:43-50, README.md:73-75
Schema Definition (GraphSchema)
The GraphSchema class is the central authority for entity and relation types. It provides natural-language hints and constraints used to steer zero-shot extractors and LLM prompts. Sources: am_diag/common/data_models/schema.py:34-37
Entity Types
The system defines 13 specialized medical entity types. Each type includes a gliner_label for local NER and a description for LLM instruction. Sources: am_diag/common/schema/medical_graphrag.py:10-143
| Entity Type | Example | Key Property |
|---|---|---|
Disease | type 2 diabetes mellitus | name |
Drug | metformin | genericName |
DrugClass | ACE inhibitors | name |
Symptom | dyspnea | name |
Pathogen | Streptococcus pneumoniae | name |
AnatomicalStructure | left ventricle | name |
Gene | CFTR | symbol |
Relation Types
The default schema includes 25 relation types, grouped into clinical, pharmacological, molecular, and structural categories. Each RelationType specifies allowed head_labels and tail_labels to prevent implausible extractions (e.g., TREATED_BY only connects Disease to Drug, DrugClass, or Procedure). Sources: am_diag/common/schema/medical_graphrag.py:147-313
Core Data Containers
KnowledgeGraph Model
The KnowledgeGraph class serves as the assembled output of a construction pipeline run. It is not used as an input to individual components but acts as a container for serialization and storage operations. Sources: am_diag/common/data_models/knowledge_graph.py:16-25
class KnowledgeGraph(BaseModel):
documents: list[Document]
chunks: list[Chunk]
entities: list[Entity]
relations: list[Relation]
communities: list[Community]
reports: list[CommunityReport]
graph_schema: GraphSchema | None
Sources: am_diag/common/data_models/knowledge_graph.py:27-35
GraphStore Interface
The GraphStore abstract base class defines the contract for persisting these models. It enforces a strict ordering guarantee for ingestion to maintain data integrity. Sources: am_diag/db/graph/base.py:24-34
This flowchart depicts the dependency order required when writing a KnowledgeGraph to a persistent store.
Sources: am_diag/db/graph/base.py:151-163
Summary of Clinical Schemas
The system provides a default implementation called MEDICAL_GRAPHRAG_SCHEMA which combines the ENTITY_TYPES and RELATION_TYPES. This schema is inspired by standard medical ontologies like SNOMED CT and the UMLS semantic network to ensure clinical relevance in the resulting knowledge graph. Sources: am_diag/common/schema/medical_graphrag.py:315-320, README.md:120-122