diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index b090857..01ad55c 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -87,22 +87,33 @@ modality is wanted before proteomics.) land on the same IDs or the graph won't link. Deterministic resolvers are tried first (OLS, text2term); the LLM is a fallback for ambiguous strings. -## 4. Canonical KG schema (target) +## 4. Canonical KG schema -Source-agnostic nodes (Pydantic v2, `extra="forbid"`): +Source-agnostic nodes (Pydantic v2, `extra="forbid"`). Implemented in PR 2 in +`models/graph_schema.py`: - `StudyNode` — `study_id` (DOI/accession), `title`, `source` (provenance), - `modality`. *(No `experimental_narrative`; that field is removed.)* -- `DatasetNode` — `dataset_id`, `data_uri`, `assay`, `cell_count`/size, - parent study. + `modality`. *(No `experimental_narrative`; that field is removed.)* Raw free + text (abstracts, full descriptions) is **not** stored on the node — it belongs + to the per-source `RawRecord`; the canonical node holds only normalized, + design-describing fields. +- `DatasetNode` — `dataset_id`, `data_uri`, `assay`, `cell_count`/size. Its + parent study is a typed `EXTRACTED_FROM` **edge**, not a stored foreign-key + field. *(Decision, PR 2: containment/relationships live on edges only; + duplicating them as node fields invites drift and gives two sources of truth.)* - `SampleNode` — `sample_id`, `data_uri`, and **design covariates**: `condition`, `perturbation`, `timepoint`, `subject`, `organism`. (Reintroduced; the prior - schema was dataset-level only.) + schema was dataset-level only.) All covariates are optional — different + sources populate different subsets. Linked to its dataset/study via a typed + edge (e.g. `HAS_SAMPLE`). - `BiologicalEntityNode` — `entity_type` ∈ {Disease, Tissue, Species, Perturbation, Assay}, `ontology_id`, `name`. **CellType is intentionally absent.** -- `GraphEdge` — typed, directed: `EXTRACTED_FROM`, `MEASURED_WITH`, - `HAS_CONDITION`, `STUDIES`, `HAS_SAMPLE`, etc. +- `GraphEdge` — typed, directed: `EXTRACTED_FROM` (Dataset→Study), `HAS_SAMPLE` + (Dataset→Sample), `HAS_TISSUE` / `HAS_CONDITION` (Dataset→BiologicalEntity), + `MEASURED_WITH` (Dataset→Assay), `STUDIES` (Study→Species), etc. `relation_type` + is a free `str` for now (the vocabulary still grows as GEO/PRIDE land); it may + become a `StrEnum` once the set stabilizes. Cross-source links are *emergent*: two studies share an edge target (`ontology_id`) rather than any source-specific key. diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index 5136d7b..a020165 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -8,10 +8,14 @@ protocol and [ARCHITECTURE.md](ARCHITECTURE.md) for the design. ## ▶ Next up -**PR 2 — Canonical KG schema refactor.** Introduce the source-agnostic schema -(`StudyNode`, `DatasetNode`, `SampleNode`, `BiologicalEntityNode`, `GraphEdge`), -reintroduce sample-level covariates, and **remove the narrative field and the -CellType entity**. Migrate `graph/builder.py` and tests. No new source yet. +**PR 3 — Source-adapter interface + cheap CELLxGENE adapter.** Define +`SourceAdapter` / `Normalizer` protocols in `sources/` + `normalize/`. Refactor +CELLxGENE into a deterministic adapter. **Remove the LLM/Azure narrative path +entirely** (delete `agent/prompts.py` narrative role, `models/narrative.py`, the +narrative `NarrativeOutput` schema in `models/graph_schema.py`, and step 2 of +`main.py` — `_build_narrative_prompt`, the agent call, and the now-unused +`narrative` variable). Drop cell-type extraction. Remove the `parce.tools.*` +mypy exemption as those modules move under `sources/`. --- @@ -22,15 +26,16 @@ Each PR is one branch, one focused scope, green CI, and a roadmap update. - [x] **PR 1 — Foundations & tooling.** CLAUDE.md, docs (architecture, roadmap, session prompt), GitHub Actions CI (ruff, mypy, pytest), ruff rule set + mypy config in `pyproject.toml`, code formatted to baseline. No behavior change. -- [ ] **PR 2 — Canonical KG schema.** Source-agnostic nodes/edges; add +- [x] **PR 2 — Canonical KG schema.** Source-agnostic nodes/edges; add `SampleNode` with design covariates; drop `experimental_narrative` and - `CellType`. Migrate builder + tests. *(Next up.)* + `CellType`. Migrate builder + tests. - [ ] **PR 3 — Source-adapter interface + cheap CELLxGENE adapter.** Define `SourceAdapter` / `Normalizer` protocols in `sources/` + `normalize/`. Refactor CELLxGENE into a deterministic adapter. **Remove the LLM/Azure narrative path entirely** (delete `agent/prompts.py` narrative role, `models/narrative.py`, - the step-2 block in `main.py`). Drop cell-type extraction. Remove the - `parce.tools.*` mypy exemption as those modules move under `sources/`. + the `NarrativeOutput` schema, and the step-2 block in `main.py`). Drop + cell-type extraction. Remove the `parce.tools.*` mypy exemption as those + modules move under `sources/`. *(Next up.)* - [ ] **PR 4 — Ontology resolver.** Shared `ontology/` stage: free-text → UBERON/MONDO/assay IDs (deterministic OLS/text2term + on-disk cache), LLM fallback for hard cases. Wire into normalizers. @@ -63,6 +68,40 @@ Each PR is one branch, one focused scope, green CI, and a roadmap update. Newest first. One entry per working session: what changed, decisions made, and what the next session should know. Keep entries short and factual. +### 2026-06-23 — PR 2: Canonical KG schema + +- Branch `pr2-canonical-kg-schema` off `main`. +- Rewrote `models/graph_schema.py` to the source-agnostic canonical schema: + - `PublicationNode` → **`StudyNode`** (`study_id`, `title`, `source`, + `modality`); dropped `abstract` and `experimental_narrative`. Raw free text + belongs to the future `RawRecord`, not the canonical node. + - `DatasetNode`: `uri`→`data_uri`, `modality`→`assay`; no parent-study field. + - **`SampleNode`** added (design covariates only: `condition`, `perturbation`, + `timepoint`, `subject`, `organism`, `data_uri`; all optional). Not populated + by the CELLxGENE path yet (Census is dataset-level — see ARCHITECTURE §6). + - `EntityType`: **`CellType` removed** (data-inferred → leakage). + - `KnowledgeGraphOutput.publications` → `studies`; added `samples`. +- Migrated `graph/builder.py`: signature is now + `build_knowledge_graph(paper_data, cellxgene_data)` (no `narrative`); emits + `StudyNode`/`DatasetNode`; ignores input `cell_types`; tissue→`HAS_TISSUE`, + disease→`HAS_CONDITION`, assay→`MEASURED_WITH`, study→species `STUDIES`. + `source="CELLxGENE"`, study `modality="scRNA-seq"` (constants for this path). +- Decision: **containment is edge-only.** `DatasetNode` does not store its parent + `study_id`; the `EXTRACTED_FROM` edge is the single source of truth (avoids a + denormalized FK that can drift). Recorded in ARCHITECTURE §4. +- `main.py`: step 3 drops the `narrative` arg; summary prints `Studies`/`Samples`. + **Deferred to PR 3 (not done here):** step 2 still generates the narrative via + Azure, but its output is now discarded (a comment marks this). `NarrativeOutput` + stays in `graph_schema.py` and `_build_narrative_prompt` still references + `cell_types` — both are part of the narrative path PR 3 deletes wholesale. +- Updated tests: `test_graph_schema.py`, `test_builder.py`, `test_orchestration.py` + (asserts `studies`/`study_id`, no narrative, cell-type exclusion, sample + covariates, tissue dedup). `models/narrative.py` + `test_models.py` untouched + (legacy GEO agent schema; PR 3/PR 5 territory). +- Gates green locally (incl. hermetic run with no `.env`): ruff check, ruff + format --check, mypy (16 files), **52 unit tests** pass. No dep changes. +- **Next session:** PR 3 (source-adapter interface + rip out the narrative path). + ### 2026-06-23 — PR 1: Foundations & tooling - Branch `restructure-context-metadata` off `main` (post-cxg-merge). - Decision: the LLM is repurposed from **narrative writing** to **structured diff --git a/src/parce/graph/builder.py b/src/parce/graph/builder.py index 10cec84..be3b4c7 100644 --- a/src/parce/graph/builder.py +++ b/src/parce/graph/builder.py @@ -1,8 +1,14 @@ -"""Deterministic Knowledge Graph construction from tool outputs + LLM narrative. +"""Deterministic Knowledge Graph construction from CELLxGENE + paper metadata. -The LLM only generates the ``experimental_narrative``. All other nodes -and edges are assembled programmatically from the structured data returned -by the CELLxGENE and EuropePMC tools. +This is the CELLxGENE ingestion path: it assembles canonical nodes and edges +programmatically from the structured data returned by the CELLxGENE and +EuropePMC tools. No LLM is involved. (It is slated to become a proper +``SourceAdapter``/``Normalizer`` in PR 3 — see docs/ROADMAP.md.) + +``CellType`` is intentionally not extracted: it is a data-inferred annotation, +not an experiment-design variable. CELLxGENE Census is dataset-level, so no +``SampleNode`` records are emitted here yet (see ARCHITECTURE.md, open question +on sample granularity). """ from __future__ import annotations @@ -15,22 +21,26 @@ EntityType, GraphEdge, KnowledgeGraphOutput, - PublicationNode, + StudyNode, ) from parce.tools.cellxgene_fetcher import _ORGANISM_ONTOLOGY logger = logging.getLogger(__name__) +# Provenance + high-level modality for everything built by this path. +_SOURCE = "CELLxGENE" +_STUDY_MODALITY = "scRNA-seq" + +# Ontology categories from CELLxGENE that become design-context entities. Cell +# types are deliberately omitted (data-inferred → leakage). _CATEGORY_TO_ENTITY_TYPE: dict[str, EntityType] = { - "cell_types": EntityType.CELL_TYPE, "tissues": EntityType.TISSUE, "diseases": EntityType.DISEASE, "assays": EntityType.ASSAY, } _CATEGORY_TO_RELATION: dict[str, str] = { - "cell_types": "MEASURES", - "tissues": "MEASURES", + "tissues": "HAS_TISSUE", "diseases": "HAS_CONDITION", "assays": "MEASURED_WITH", } @@ -39,27 +49,24 @@ def build_knowledge_graph( paper_data: dict, cellxgene_data: dict, - narrative: str, ) -> KnowledgeGraphOutput: - """Assemble a complete ``KnowledgeGraphOutput`` from structured data. + """Assemble a canonical ``KnowledgeGraphOutput`` from structured data. Parameters ---------- paper_data: - Dict with keys ``doi``, ``title``, ``abstract`` (from ``fetch_paper_metadata``). + Dict with keys ``doi``, ``title`` (from ``fetch_paper_metadata``). cellxgene_data: Dict with key ``datasets`` containing per-dataset metadata and ontology summaries (from ``fetch_cellxgene_datasets``). - narrative: - The LLM-generated experimental narrative string. """ - doi = paper_data["doi"] + study_id = paper_data["doi"] - publication = PublicationNode( - doi=doi, + study = StudyNode( + study_id=study_id, title=paper_data.get("title", ""), - abstract=paper_data.get("abstract", ""), - experimental_narrative=narrative, + source=_SOURCE, + modality=_STUDY_MODALITY, ) datasets: list[DatasetNode] = [] @@ -73,8 +80,8 @@ def build_knowledge_graph( datasets.append( DatasetNode( dataset_id=dataset_id, - uri=ds["h5ad_uri"], - modality=ds.get("modality", "unknown"), + data_uri=ds["h5ad_uri"], + assay=ds.get("modality", "unknown"), cell_count=ds["cell_count"], ) ) @@ -82,7 +89,7 @@ def build_knowledge_graph( edges.append( GraphEdge( source_id=dataset_id, - target_id=doi, + target_id=study_id, relation_type="EXTRACTED_FROM", ) ) @@ -100,7 +107,7 @@ def build_knowledge_graph( name=name, ) - # Register entities and create edges per category + # Register entities and create edges per design-context category for category, entity_type in _CATEGORY_TO_ENTITY_TYPE.items(): relation = _CATEGORY_TO_RELATION[category] for term in ontology.get(category, []): @@ -122,28 +129,30 @@ def build_knowledge_graph( ) ) - # Publication -> Species edges - for ont_id in species_seen: - resolved_id = _ORGANISM_ONTOLOGY[ont_id][0] if ont_id in _ORGANISM_ONTOLOGY else ont_id + # Study -> Species edges + for organism_key in species_seen: + species_id = _ORGANISM_ONTOLOGY[organism_key][0] edges.append( GraphEdge( - source_id=doi, - target_id=resolved_id, + source_id=study_id, + target_id=species_id, relation_type="STUDIES", ) ) kg = KnowledgeGraphOutput( - publications=[publication], + studies=[study], datasets=datasets, + samples=[], biological_entities=list(entity_registry.values()), edges=edges, ) logger.info( - "Built KG: publications=%d datasets=%d entities=%d edges=%d", - len(kg.publications), + "Built KG: studies=%d datasets=%d samples=%d entities=%d edges=%d", + len(kg.studies), len(kg.datasets), + len(kg.samples), len(kg.biological_entities), len(kg.edges), ) diff --git a/src/parce/main.py b/src/parce/main.py index 13513f1..ea46c54 100644 --- a/src/parce/main.py +++ b/src/parce/main.py @@ -193,14 +193,18 @@ async def run(doi: str = _DEFAULT_DOI) -> None: return # ------------------------------------------------------------------ - # Step 3: Build the Knowledge Graph programmatically + # Step 3: Build the canonical Knowledge Graph programmatically # ------------------------------------------------------------------ + # NOTE: the canonical KG no longer stores a narrative. The narrative step + # above is retained only until PR 3 removes the LLM/Azure path entirely + # (see docs/ROADMAP.md); its output is intentionally not fed into the KG. logger.info("Step 3/3: Assembling knowledge graph") - kg = build_knowledge_graph(paper_data, cellxgene_data, narrative) + kg = build_knowledge_graph(paper_data, cellxgene_data) print("Knowledge graph constructed successfully:") - print(f" Publications: {len(kg.publications)}") + print(f" Studies: {len(kg.studies)}") print(f" Datasets: {len(kg.datasets)}") + print(f" Samples: {len(kg.samples)}") print(f" Biological entities: {len(kg.biological_entities)}") print(f" Edges: {len(kg.edges)}") diff --git a/src/parce/models/graph_schema.py b/src/parce/models/graph_schema.py index bdf316e..40ccfaa 100644 --- a/src/parce/models/graph_schema.py +++ b/src/parce/models/graph_schema.py @@ -1,8 +1,18 @@ -"""Pydantic schemas for the Knowledge Graph output. +"""Pydantic schemas for the canonical Knowledge Graph. -Defines node and edge types for an entity-centric graph that links -publications, datasets, and biological entities (cell types, tissues, -diseases, species, perturbations) via typed edges. +Source-agnostic node and edge types for an entity-centric graph that links +studies, datasets, and samples to biological entities (tissues, diseases, +species, perturbations, assays) via typed edges. Every source — structured or +agent-extracted — emits these same models; per-source variation lives entirely +in the adapters/normalizers, not here. + +Two design rules are enforced by this schema: + +- **Context is design, not data-inferred outcome.** ``CellType`` is intentionally + absent from :class:`EntityType`: cell type is called from expression and would + leak the signal the downstream model must learn. +- **Cross-source links are emergent.** Studies connect through shared + ``ontology_id`` edge targets, never through any source-specific key. """ from __future__ import annotations @@ -13,51 +23,93 @@ class EntityType(StrEnum): - """Controlled vocabulary for biological entity categories.""" + """Controlled vocabulary for biological entity categories. + + ``CellType`` is deliberately excluded — it is a data-inferred annotation + (called from expression), not an experiment-design variable, and storing it + would leak the downstream learning target. + """ DISEASE = "Disease" - CELL_TYPE = "CellType" TISSUE = "Tissue" SPECIES = "Species" PERTURBATION = "Perturbation" ASSAY = "Assay" -class PublicationNode(BaseModel): - """A published study identified by DOI.""" +class StudyNode(BaseModel): + """A study (publication or repository accession), source-agnostic. + + Holds only normalized, design-describing fields. Raw free text (abstracts, + full descriptions) lives in the per-source ``RawRecord``, not here. + """ model_config = ConfigDict(extra="forbid") - doi: str = Field( - ..., description="Digital Object Identifier (e.g. 10.1038/s41586-023-05869-0)." + study_id: str = Field( + ..., + description="Stable study identifier: a DOI or repository accession (e.g. GSE164378).", ) - title: str = Field(..., description="Publication title.") - abstract: str = Field(..., description="Full abstract text.") - experimental_narrative: str = Field( + title: str = Field(..., description="Study title.") + source: str = Field( ..., - description=( - "LLM-generated concise narrative synthesising the abstract and " - "the structured CELLxGENE ontology data into a description of " - "how the data was obtained." - ), + description="Provenance of the study record (e.g. 'CELLxGENE', 'GEO', 'PRIDE').", + ) + modality: str = Field( + ..., + description="High-level assay modality of the study (e.g. 'scRNA-seq', 'proteomics').", ) class DatasetNode(BaseModel): - """A single-cell dataset hosted on CELLxGENE.""" + """A dataset belonging to a study; the link to its study is a typed edge.""" model_config = ConfigDict(extra="forbid") - dataset_id: str = Field(..., description="CELLxGENE dataset UUID.") - uri: str = Field( + dataset_id: str = Field(..., description="Source dataset identifier (e.g. CELLxGENE UUID).") + data_uri: str = Field( ..., - description="Remote URI for the H5AD file (e.g. s3://cellxgene-data-public/...).", + description="Remote URI for the dataset payload (e.g. s3://cellxgene-data-public/...).", ) - modality: str = Field( + assay: str = Field( ..., - description="Assay modality (e.g. 'scRNA-seq', '10x 3' v3', 'Smart-seq2').", + description="Specific assay/technology (e.g. \"10x 3' v3\", 'Smart-seq2').", + ) + cell_count: int = Field(..., description="Number of cells (or rows) in the dataset.") + + +class SampleNode(BaseModel): + """A biological sample with experiment-*design* covariates only. + + Different sources populate different subsets of the covariates, so all of + them are optional. Data-inferred annotations (cell type, cluster labels) are + never recorded here. The link to a parent dataset/study is a typed edge. + """ + + model_config = ConfigDict(extra="forbid") + + sample_id: str = Field(..., description="Repository sample accession (e.g. GSM1234567).") + data_uri: str | None = Field( + default=None, + description="URI to this sample's raw data, if the source exposes one.", + ) + organism: str | None = Field(default=None, description="Species name (e.g. 'Mus musculus').") + condition: str | None = Field( + default=None, + description="Experimental condition as designed (e.g. 'control', 'stimulated').", + ) + perturbation: str | None = Field( + default=None, + description="Designed perturbation (e.g. a gene knockout, drug, dose).", + ) + timepoint: str | None = Field( + default=None, + description="Designed sampling timepoint (e.g. '0h', 'day 7').", + ) + subject: str | None = Field( + default=None, + description="Subject/donor/replicate identifier the sample belongs to.", ) - cell_count: int = Field(..., description="Total number of cells in the dataset.") class BiologicalEntityNode(BaseModel): @@ -69,11 +121,11 @@ class BiologicalEntityNode(BaseModel): ontology_id: str = Field( ..., description=( - "Ontology identifier (e.g. 'CL:0000084' for T cell, " + "Ontology identifier (e.g. 'UBERON:0000178' for blood, " "'MONDO:0005061' for lung disease, or 'unknown')." ), ) - name: str = Field(..., description="Human-readable name (e.g. 'T cell').") + name: str = Field(..., description="Human-readable name (e.g. 'blood').") class GraphEdge(BaseModel): @@ -82,36 +134,39 @@ class GraphEdge(BaseModel): model_config = ConfigDict(extra="forbid") source_id: str = Field( - ..., description="Identifier of the source node (DOI or dataset_id or ontology_id)." + ..., + description="Source node id (study_id, dataset_id, sample_id, or ontology_id).", ) target_id: str = Field(..., description="Identifier of the target node.") relation_type: str = Field( ..., description=( - "Relationship label, e.g. 'EXTRACTED_FROM' (Dataset -> Publication), " - "'MEASURES' (Dataset -> BiologicalEntity), " - "'HAS_CONDITION' (Dataset -> BiologicalEntity)." + "Relationship label, e.g. 'EXTRACTED_FROM' (Dataset -> Study), " + "'HAS_SAMPLE' (Dataset -> Sample), 'HAS_TISSUE' / 'HAS_CONDITION' " + "(Dataset -> BiologicalEntity), 'MEASURED_WITH' (Dataset -> Assay), " + "'STUDIES' (Study -> Species)." ), ) class KnowledgeGraphOutput(BaseModel): - """Top-level wrapper returned by the agent: a complete knowledge graph - for one or more publications and their associated datasets.""" + """Top-level wrapper: a complete knowledge graph for one or more studies.""" model_config = ConfigDict(extra="forbid") - publications: list[PublicationNode] = Field(default_factory=list) + studies: list[StudyNode] = Field(default_factory=list) datasets: list[DatasetNode] = Field(default_factory=list) + samples: list[SampleNode] = Field(default_factory=list) biological_entities: list[BiologicalEntityNode] = Field(default_factory=list) edges: list[GraphEdge] = Field(default_factory=list) class NarrativeOutput(BaseModel): - """Minimal schema for the LLM narrative generation step. + """Transitional schema for the legacy LLM narrative step. - Only the experimental_narrative requires LLM intelligence; - all other KG fields are assembled programmatically. + The canonical KG no longer stores a narrative; this remains only as the + ``response_format`` for the narrative agent until the whole narrative path is + removed in PR 3 (see docs/ROADMAP.md). """ experimental_narrative: str = Field( diff --git a/tests/test_builder.py b/tests/test_builder.py index fd81939..21d369e 100644 --- a/tests/test_builder.py +++ b/tests/test_builder.py @@ -11,6 +11,9 @@ "abstract": "We profiled T cells.", } +# Note: ``cell_types`` are present in the input but must be ignored by the +# builder (CellType is a data-inferred annotation, intentionally excluded). +# ``blood`` (UBERON:0000178) appears in both datasets to exercise dedup. _CELLXGENE_DATA = { "doi": "10.1234/test", "datasets": [ @@ -51,6 +54,7 @@ {"name": "T cell", "ontology_id": "CL:0000084"}, ], "tissues": [ + {"name": "blood", "ontology_id": "UBERON:0000178"}, {"name": "lung", "ontology_id": "UBERON:0002048"}, ], "diseases": [], @@ -62,73 +66,90 @@ ], } -_NARRATIVE = "This study profiled T and B cells from blood and lung tissue." - class TestBuildKnowledgeGraph: def test_basic_structure(self): - kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA, _NARRATIVE) + kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA) - assert len(kg.publications) == 1 - assert kg.publications[0].doi == "10.1234/test" - assert kg.publications[0].experimental_narrative == _NARRATIVE - assert kg.publications[0].title == "Test Study" - assert kg.publications[0].abstract == "We profiled T cells." + assert len(kg.studies) == 1 + assert kg.studies[0].study_id == "10.1234/test" + assert kg.studies[0].title == "Test Study" + assert kg.studies[0].source == "CELLxGENE" + assert kg.studies[0].modality == "scRNA-seq" assert len(kg.datasets) == 2 assert kg.datasets[0].dataset_id == "ds-001" + assert kg.datasets[0].data_uri == "s3://bucket/ds-001.h5ad" + assert kg.datasets[0].assay == "10x 3' v3" assert kg.datasets[1].dataset_id == "ds-002" - def test_entity_deduplication(self): - """T cell (CL:0000084) appears in both datasets but should be one entity.""" - kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA, _NARRATIVE) + def test_cell_type_excluded(self): + """Cell types in the input must not produce entities or edges.""" + kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA) + + names = {e.name for e in kg.biological_entities} + assert "T cell" not in names + assert "B cell" not in names + + ontology_ids = {e.ontology_id for e in kg.biological_entities} + assert "CL:0000084" not in ontology_ids + assert "CL:0000236" not in ontology_ids + + edge_targets = {e.target_id for e in kg.edges} + assert "CL:0000084" not in edge_targets + + def test_tissue_entity_deduplication(self): + """blood (UBERON:0000178) appears in both datasets but is one entity.""" + kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA) entity_ids = [e.ontology_id for e in kg.biological_entities] - assert entity_ids.count("CL:0000084") == 1 + assert entity_ids.count("UBERON:0000178") == 1 def test_species_entity_created(self): - kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA, _NARRATIVE) + kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA) species = [e for e in kg.biological_entities if e.entity_type == EntityType.SPECIES] assert len(species) == 1 assert species[0].ontology_id == "NCBITaxon:9606" assert species[0].name == "Homo sapiens" + def test_no_samples_for_cellxgene(self): + """Census is dataset-level; no SampleNode records are emitted (yet).""" + kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA) + assert kg.samples == [] + def test_extracted_from_edges(self): - kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA, _NARRATIVE) + kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA) extracted = [e for e in kg.edges if e.relation_type == "EXTRACTED_FROM"] assert len(extracted) == 2 assert {e.source_id for e in extracted} == {"ds-001", "ds-002"} assert all(e.target_id == "10.1234/test" for e in extracted) - def test_measures_edges(self): - kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA, _NARRATIVE) - - measures = [e for e in kg.edges if e.relation_type == "MEASURES"] - source_target_pairs = {(e.source_id, e.target_id) for e in measures} + def test_has_tissue_edges(self): + kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA) - assert ("ds-001", "CL:0000084") in source_target_pairs - assert ("ds-001", "CL:0000236") in source_target_pairs - assert ("ds-001", "UBERON:0000178") in source_target_pairs - assert ("ds-002", "CL:0000084") in source_target_pairs - assert ("ds-002", "UBERON:0002048") in source_target_pairs + tissue_edges = [e for e in kg.edges if e.relation_type == "HAS_TISSUE"] + pairs = {(e.source_id, e.target_id) for e in tissue_edges} + assert ("ds-001", "UBERON:0000178") in pairs + assert ("ds-002", "UBERON:0000178") in pairs + assert ("ds-002", "UBERON:0002048") in pairs def test_has_condition_edges(self): - kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA, _NARRATIVE) + kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA) conditions = [e for e in kg.edges if e.relation_type == "HAS_CONDITION"] assert any(e.target_id == "PATO:0000461" for e in conditions) def test_measured_with_edges(self): - kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA, _NARRATIVE) + kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA) assay_edges = [e for e in kg.edges if e.relation_type == "MEASURED_WITH"] assert any(e.source_id == "ds-001" and e.target_id == "EFO:0009922" for e in assay_edges) assert any(e.source_id == "ds-002" and e.target_id == "EFO:0008931" for e in assay_edges) def test_studies_edge(self): - kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA, _NARRATIVE) + kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA) studies = [e for e in kg.edges if e.relation_type == "STUDIES"] assert len(studies) == 1 @@ -136,15 +157,16 @@ def test_studies_edge(self): assert studies[0].target_id == "NCBITaxon:9606" def test_empty_datasets(self): - kg = build_knowledge_graph(_PAPER_DATA, {"doi": "10.1234/test", "datasets": []}, _NARRATIVE) + kg = build_knowledge_graph(_PAPER_DATA, {"doi": "10.1234/test", "datasets": []}) - assert len(kg.publications) == 1 + assert len(kg.studies) == 1 assert len(kg.datasets) == 0 + assert len(kg.samples) == 0 assert len(kg.biological_entities) == 0 assert len(kg.edges) == 0 def test_roundtrip_json(self): - kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA, _NARRATIVE) + kg = build_knowledge_graph(_PAPER_DATA, _CELLXGENE_DATA) json_str = kg.model_dump_json() from parce.models.graph_schema import KnowledgeGraphOutput diff --git a/tests/test_graph_schema.py b/tests/test_graph_schema.py index 39d75b2..5c0ceea 100644 --- a/tests/test_graph_schema.py +++ b/tests/test_graph_schema.py @@ -1,4 +1,4 @@ -"""Unit tests for the Knowledge Graph Pydantic models.""" +"""Unit tests for the canonical Knowledge Graph Pydantic models.""" from __future__ import annotations @@ -12,41 +12,52 @@ GraphEdge, KnowledgeGraphOutput, NarrativeOutput, - PublicationNode, + SampleNode, + StudyNode, ) class TestEntityType: def test_values(self): assert EntityType.DISEASE == "Disease" - assert EntityType.CELL_TYPE == "CellType" assert EntityType.TISSUE == "Tissue" assert EntityType.SPECIES == "Species" assert EntityType.PERTURBATION == "Perturbation" assert EntityType.ASSAY == "Assay" + def test_celltype_removed(self): + """CellType is intentionally absent (data-inferred → leakage).""" + assert not hasattr(EntityType, "CELL_TYPE") + with pytest.raises(ValueError): + EntityType("CellType") + def test_from_string(self): assert EntityType("Disease") is EntityType.DISEASE -class TestPublicationNode: +class TestStudyNode: def test_valid(self): - pub = PublicationNode( - doi="10.1038/s41586-023-05869-0", + study = StudyNode( + study_id="10.1038/s41586-023-05869-0", title="A study", - abstract="An abstract.", - experimental_narrative="Narrative text.", + source="CELLxGENE", + modality="scRNA-seq", ) - assert pub.doi == "10.1038/s41586-023-05869-0" + assert study.study_id == "10.1038/s41586-023-05869-0" + assert study.source == "CELLxGENE" + + def test_no_narrative_field(self): + """The narrative field was removed from the canonical schema.""" + assert "experimental_narrative" not in StudyNode.model_fields def test_extra_forbidden(self): with pytest.raises(ValidationError): - PublicationNode( - doi="10.1234/test", + StudyNode( + study_id="10.1234/test", title="T", - abstract="A", - experimental_narrative="N", - extra="bad", + source="CELLxGENE", + modality="scRNA-seq", + abstract="leftover", ) @@ -54,25 +65,59 @@ class TestDatasetNode: def test_valid(self): ds = DatasetNode( dataset_id="abc-123", - uri="s3://cellxgene-data-public/cell-census/h5ads/abc-123.h5ad", - modality="scRNA-seq", + data_uri="s3://cellxgene-data-public/cell-census/h5ads/abc-123.h5ad", + assay="10x 3' v3", cell_count=50000, ) assert ds.cell_count == 50000 + assert ds.assay == "10x 3' v3" def test_extra_forbidden(self): with pytest.raises(ValidationError): - DatasetNode(dataset_id="x", uri="s3://x", modality="m", cell_count=1, oops=True) + DatasetNode(dataset_id="x", data_uri="s3://x", assay="m", cell_count=1, oops=True) + + +class TestSampleNode: + def test_minimal(self): + sample = SampleNode(sample_id="GSM0001") + assert sample.sample_id == "GSM0001" + assert sample.data_uri is None + assert sample.organism is None + assert sample.condition is None + assert sample.perturbation is None + assert sample.timepoint is None + assert sample.subject is None + + def test_full_design_covariates(self): + sample = SampleNode( + sample_id="GSM0002", + data_uri="https://sra/SRR001.fastq", + organism="Mus musculus", + condition="stimulated", + perturbation="Pdcd1 knockout", + timepoint="day 7", + subject="donor-3", + ) + assert sample.perturbation == "Pdcd1 knockout" + assert sample.timepoint == "day 7" + + def test_no_data_inferred_fields(self): + """Design covariates only — no cell_type / cluster annotations.""" + assert "cell_type" not in SampleNode.model_fields + + def test_extra_forbidden(self): + with pytest.raises(ValidationError): + SampleNode(sample_id="GSM0003", cell_type="CD8+ T cell") class TestBiologicalEntityNode: def test_valid(self): entity = BiologicalEntityNode( - entity_type=EntityType.CELL_TYPE, - ontology_id="CL:0000084", - name="T cell", + entity_type=EntityType.TISSUE, + ontology_id="UBERON:0000178", + name="blood", ) - assert entity.entity_type == EntityType.CELL_TYPE + assert entity.entity_type == EntityType.TISSUE def test_string_coercion(self): entity = BiologicalEntityNode( @@ -86,6 +131,10 @@ def test_invalid_entity_type(self): with pytest.raises(ValidationError): BiologicalEntityNode(entity_type="NotAType", ontology_id="X", name="bad") + def test_celltype_rejected(self): + with pytest.raises(ValidationError): + BiologicalEntityNode(entity_type="CellType", ontology_id="CL:0000084", name="T cell") + def test_extra_forbidden(self): with pytest.raises(ValidationError): BiologicalEntityNode( @@ -110,32 +159,40 @@ def test_extra_forbidden(self): class TestKnowledgeGraphOutput: def test_empty(self): kg = KnowledgeGraphOutput() - assert kg.publications == [] + assert kg.studies == [] + assert kg.samples == [] assert kg.edges == [] def test_roundtrip_json(self): kg = KnowledgeGraphOutput( - publications=[ - PublicationNode( - doi="10.1234/test", + studies=[ + StudyNode( + study_id="10.1234/test", title="Test", - abstract="Abstract.", - experimental_narrative="Narrative.", + source="CELLxGENE", + modality="scRNA-seq", ) ], datasets=[ DatasetNode( dataset_id="ds-1", - uri="s3://bucket/ds-1.h5ad", - modality="scRNA-seq", + data_uri="s3://bucket/ds-1.h5ad", + assay="10x 3' v3", cell_count=1000, ) ], + samples=[ + SampleNode( + sample_id="GSM1", + organism="Homo sapiens", + condition="control", + ) + ], biological_entities=[ BiologicalEntityNode( - entity_type=EntityType.CELL_TYPE, - ontology_id="CL:0000084", - name="T cell", + entity_type=EntityType.TISSUE, + ontology_id="UBERON:0000178", + name="blood", ) ], edges=[ diff --git a/tests/test_orchestration.py b/tests/test_orchestration.py index c03b10d..39efaa9 100644 --- a/tests/test_orchestration.py +++ b/tests/test_orchestration.py @@ -109,9 +109,10 @@ async def test_full_pipeline(self, _mock_tools, _mock_agent, tmp_path): import json kg = json.loads(out_file.read_text()) - assert len(kg["publications"]) == 1 - assert kg["publications"][0]["doi"] == "10.1234/mock" - assert kg["publications"][0]["experimental_narrative"] == "A mock narrative about T cells." + assert len(kg["studies"]) == 1 + assert kg["studies"][0]["study_id"] == "10.1234/mock" + # The canonical KG no longer stores a narrative. + assert "experimental_narrative" not in kg["studies"][0] assert len(kg["datasets"]) == 1 assert len(kg["biological_entities"]) > 0 assert len(kg["edges"]) > 0