Neo4j Graph Database Integration
Relevant source files
The following files were used as context for generating this wiki page:
The Neo4j Graph Database Integration provides the persistent storage layer for the Agentic GraphRAG system. It is responsible for storing structured clinical knowledge extracted from medical literature, including entities, their relationships, and hierarchical communities. The integration is built on an abstract GraphStore base class to ensure storage-agnosticism for upper-level components, with a concrete Neo4jClient implementation that utilizes the official Neo4j async driver.
Sources: am_diag/db/graph/base.py:1-15, README.md:27-31
Architecture and Components
The integration follows a decoupled architecture where data models, serialization logic, and database transport are separated. The Neo4jClient handles the execution of Cypher queries, which are externally loaded from .cypher files rather than being hardcoded in Python.
Core Classes
| Class | Description |
|---|---|
GraphStore | Abstract base class defining the interface for graph database operations, including low-level transport and high-level ingest methods. |
Neo4jClient | Concrete implementation of GraphStore using the neo4j async driver. |
Neo4jTransaction | A wrapper for AsyncTransaction that provides a simplified run interface within atomic blocks. |
Neo4jSettings | Pydantic-based configuration for connection details (URI, user, password, database). |
Sources: am_diag/db/graph/base.py:25-40, am_diag/db/graph/neo4j.py:47-75, am_diag/db/graph/settings.py
Class Diagram
The following diagram illustrates the relationship between the base abstraction and the Neo4j-specific implementation.
Schema and Data Model
The system utilizes a typed property graph model. The schema is runtime-injectable via the GraphSchema class, which defines medical entity types (e.g., Disease, Drug) and relation types (e.g., TREATED_BY).
Node Labels and Edge Types
The system persists five primary node labels and seven relationship types:
- Node Labels:
Document,Chunk,Entity,Community,CommunityReport. - Edge Types:
(:Document) <-[:PART_OF]- (:Chunk)(:Chunk) -[:NEXT_CHUNK]-> (:Chunk)(:Chunk) -[:HAS_ENTITY]-> (:Entity)(:Entity) -[:RELATES_TO]-> (:Entity)(:Entity) -[:IN_COMMUNITY]-> (:Community)(:Community) -[:PARENT_COMMUNITY]-> (:Community)(:Community) -[:HAS_REPORT]-> (:CommunityReport)
Sources: README.md:33-43, am_diag/common/data_models/schema.py:14-38
DDL and Constraints
The setup_schema method in Neo4jClient ensures that the database is properly initialized with uniqueness constraints on the uuid property for all labels and specialized indexes for search.
Sources: am_diag/db/graph/neo4j.py:157-175
Ingestion Pipeline
Data is written to Neo4j in a specific dependency order to maintain referential integrity. The write_knowledge_graph method orchestrates this flow, ensuring nodes exist before relationships are merged between them.
Data Flow Diagram
Batch Execution
To handle large datasets efficiently, the client implements execute_batch. This method splits a list of records into slices (default size 500) and executes them within explicit transactions. This prevents memory issues and ensures that a failure in one batch does not corrupt others.
Sources: am_diag/db/graph/neo4j.py:118-135, am_diag/db/graph/base.py:59-71
Serialization
The serialize.py module contains helper functions that transform Pydantic data models into flat dictionaries suitable for Cypher UNWIND parameters.
- Entity Serialization: Entities are grouped by their
label(e.g.,:Drug) during upsert because Neo4j secondary labels are Cypher identifiers that cannot be passed as parameters. - Relationship Serialization: Relations are only written if they have resolved
head_idandtail_id.
Sources: am_diag/db/graph/serialize.py, am_diag/db/graph/neo4j.py:204-230, tests/am_diag/db/graph/test_base_ingest.py:43-70
Configuration
Configuration is managed via the Neo4jSettings class, which reads from environment variables with a NEO4J_ prefix.
| Setting | Env Var | Default |
|---|---|---|
uri | NEO4J_URI | bolt://localhost:7687 |
user | NEO4J_USER | neo4j |
password | NEO4J_PASSWORD | password |
database | NEO4J_DATABASE | neo4j |
Sources: am_diag/db/graph/settings.py, am_diag/db/graph/neo4j.py:335-347
Conclusion
The Neo4j integration provides a robust, schema-driven storage system for medical knowledge. By utilizing an abstract base class and externalized Cypher queries, the system maintains high maintainability and testability while leveraging Neo4j's capabilities for representing complex clinical relationships and hierarchical community structures.