PR 2: Canonical KG schema refactor

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.

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

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,16 +80,16 @@ 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"],
             )
         )
 
         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),
     )

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)}")