Ingestion Pipeline (LangGraph)
Relevant source files
The following files were used as context for generating this wiki page:
The Ingestion Pipeline is the core data processing engine of the Agentic Medical Diagnosis system. It transforms raw medical literature into a structured, schema-driven Knowledge Graph stored in Neo4j and a vector store for semantic retrieval. The system leverages LangGraph to orchestrate graph construction components, providing a linear but conditional flow for extraction, normalization, aggregation, and resolution.
The pipeline operates by streaming documents from various corpora, chunking them, and then passing them through a StateGraph. This graph-based architecture allows the system to maintain state—encapsulated in the IngestionState—as it progresses from raw text to refined clinical entities and relations.
Sources: README.md:95-100, am_diag/ingestion/extraction_pipeline.py:1-10
Pipeline Architecture and Orchestration
The ingestion process is fanned out across three primary sub-pipelines: Extraction, Embedding, and Graph Writing. The high-level orchestration is handled by run_corpus_ingestion, which manages the transition between raw document streaming and graph-based processing.
Data Flow Overview
The following diagram illustrates the flow from a corpus loader through the internal LangGraph extraction nodes to final storage.
The pipeline uses a StateGraph to manage IngestionState, which includes the batch_id, lists of entities, relations, and the final KnowledgeGraph object.
Sources: am_diag/ingestion/extraction_pipeline.py:90-135, am_diag/ingestion/corpus_ingestion.py:90-140
Extraction Nodes (LangGraph)
The LangGraph defined in extraction_pipeline.py consists of several discrete nodes that transform raw extracted data into normalized knowledge.
Processing Nodes
| Node Name | Description | Key Components Used |
|---|---|---|
prepare | Initializes batch_id and empty working lists for state accumulation. | IngestionState |
normalize | Deduplicates entities by (name, label) and relations by (head, tail, type) within a batch. | EntityRelationNormalizer |
aggregate | Collapses entries across all chunks into a unified candidate set. | KGAggregator |
resolve | Resolves entity and relation-type names using SemHash and LLM-based deduplication. | EntityEdgeResolver |
build_report | Generates ExtractionReport and IngestionReport; writes to Neo4j if a client is provided. | Neo4jClient |
Sources: am_diag/ingestion/extraction_pipeline.py:35-85, am_diag/graph_construction/normalize.py:15-30
Normalization Logic
The normalize node utilizes the EntityRelationNormalizer to filter entities based on min_entity_score and min_entity_length. It merges provenance (source extractors and chunk IDs) and surface forms for colliding entities.
Sources: am_diag/graph_construction/normalize.py:38-60
Knowledge Graph Construction
The construction phase follows a "Three-extractor fusion" strategy to ensure high-fidelity clinical knowledge capture. While the LangGraph orchestrates the flow, the content is derived from:
- GLiNER: Batch zero-shot Named Entity Recognition (NER).
- GLiREL: Zero-shot relation extraction over identified entity spans.
- LLM: BAML-driven extraction using dynamic enums constrained by the medical schema.
Schema Injection
The GraphSchema is injected into every node of the pipeline. It defines valid EntityType and RelationType objects, which steer the zero-shot extractors and LLM prompts.
# From am_diag/common/data_models/schema.py
class GraphSchema(BaseModel):
"""Runtime-injectable KG schema."""
entity_types: list[EntityType]
relation_types: list[RelationType]
def as_text(self) -> str:
"""Full schema string for prompt injection."""
# Generates (:Entity:Label) and (:Entity)-[:RELATES_TO {type: '...'}]->(:Entity)
Sources: README.md:105-115, am_diag/common/data_models/schema.py:40-90
Vector Embedding Pipeline
Once the knowledge graph is constructed and normalized, the system vectorizes entities, relations, and chunks. This occurs in the _embed_and_upsert phase of corpus_ingestion.py.
Embedding Representations
Entities are embedded by their canonical_name, and chunks are embedded by their raw text. Relations are unique in that they are embedded in three ways to capture different semantic aspects:
- Edge fact: The full description of the relation.
- Edge type: The predicate alone (e.g.,
TREATED_BY). - Full SPO: The Subject-Predicate-Object triple sentence.
Sources: README.md:145-165, am_diag/ingestion/corpus_ingestion.py:30-65
Multi-Corpus Ingestion (CLI)
The am-diag-ingest CLI tool provides the entry point for running the pipeline across multiple datasets. It supports a registry of medical corpora including PubMed, StatPearls, and Clinical Guidelines.
| Corpus Name | Loader Class | Description |
|---|---|---|
textbooks | TextbooksCorpusLoader | USMLE preparation textbooks. |
statpearls | StatPearlsCorpusLoader | Peer-reviewed point-of-care clinical articles. |
pubmed | PubMedCorpusLoader | Biomedical literature titles and abstracts. |
clinical_guidelines | ClinicalGuidelinesCorpusLoader | Authoritative health organization guidelines. |
Sources: am_diag/pipelines/ingestion.py:35-50, README.md:275-285
Summary
The LangGraph-based Ingestion Pipeline provides a structured and extensible framework for medical knowledge extraction. By separating the extraction, normalization, and resolution into discrete graph nodes, the system ensures that medical data is progressively refined before being committed to Neo4j. This architecture supports high-concurrency processing of clinical literature while maintaining strict adherence to a medically-informed schema.
Sources: am_diag/ingestion/extraction_pipeline.py:1-10, README.md:95-110