diff --git a/.env.example b/.env.example index 59154cf..da4fa8b 100644 --- a/.env.example +++ b/.env.example @@ -8,3 +8,9 @@ AZURE_AI_PROJECT_ENDPOINT=https://.services.ai.azure.com/api/proj # The model deployment name from your Foundry project's model catalog. # Examples: mistral-large, gpt-4o, DeepSeek-R1, Meta-Llama-3-70B AZURE_AI_MODEL_DEPLOYMENT_NAME=gpt-4o + +# Optional NCBI E-utilities / GEO courtesy parameters. NCBI asks callers to +# identify themselves; supplying these raises rate limits. Both are optional — +# GEO fetches work without them. +# NCBI_EMAIL=you@example.com +# NCBI_API_KEY= diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 9edffd2..35889ea 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -82,7 +82,14 @@ modality is wanted before proteomics.) - **The LLM is boxed in.** It lives only inside unstructured normalizers and is constrained by `response_format` to fill canonical fields — it cannot emit prose. This keeps the system testable and reproducible, and makes the agent's - output directly comparable to the deterministic path. + output directly comparable to the deterministic path. Concretely (PR 5), a + normalizer reaches the LLM only through the narrow **synchronous** + `StructuredExtractor` contract (`parce.agent.base`): `extract(instructions, + content, response_model) -> response_model`. The real implementation + (`parce.agent.extraction.AzureExtractionAgent`) is the *only* Azure-touching + module; it bridges the async `agent-framework` API to that sync seam internally, + so normalizers stay synchronous and are unit-tested by injecting a deterministic + fake extractor (no Azure in CI). - **Ontology resolution is a shared stage, not per-adapter.** All sources must 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. @@ -110,7 +117,11 @@ Source-agnostic nodes (Pydantic v2, `extra="forbid"`). Implemented in PR 2 in `perturbation`, `timepoint`, `subject`, `organism`. (Reintroduced; the prior 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`). + edge (e.g. `HAS_SAMPLE`). *(First populated by GEO in PR 5: one `SampleNode` per + `GSM`, with `organism`/`data_uri` read deterministically from structured SOFT + fields and `condition`/`perturbation`/`timepoint`/`subject` extracted by the LLM + from `characteristics_ch1`. CELLxGENE still emits none — Census is dataset-level, + §7.)* - `BiologicalEntityNode` — `entity_type` ∈ {Disease, Tissue, Species, Perturbation, Assay}, `ontology_id`, `name`. **CellType is intentionally absent.** @@ -119,6 +130,16 @@ Source-agnostic nodes (Pydantic v2, `extra="forbid"`). Implemented in PR 2 in `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. + - *(Decision, PR 5: a source's design-context and `HAS_SAMPLE` edges originate + at whichever node is the natural containment root for that source. CELLxGENE is + dataset-centric (Census ships datasets), so its edges originate at the + `DatasetNode`. **GEO has no distinct dataset artifact** — a series *is* the + study, its data lives in per-sample supplementary files — so GEO emits no + `DatasetNode` and its `HAS_TISSUE`/`HAS_CONDITION`/`MEASURED_WITH`/`HAS_SAMPLE` + edges originate at the `StudyNode`, with `assay`/`molecular_layer` carried on + the study. This is deliberate: cross-source linking flows through shared entity + `ontology_id` **targets**, not the originating node, so a mixed origin does not + break the merge (PR 6).)* 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 0ffcb7e..cd55d81 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -8,22 +8,19 @@ protocol and [ARCHITECTURE.md](ARCHITECTURE.md) for the design. ## ▶ Next up -**PR 5 — GEO extraction agent (vertical slice).** First source whose metadata is -*unstructured* free text, so the first to use the LLM. Add a GEO adapter -(NCBI E-utilities / GEOparse) emitting a `RawRecord`, and an **Azure extraction -normalizer** that fills the canonical schema via `response_format` (structured -output only — never prose). Extract sample-level **design** covariates from -`characteristics_ch1` (`condition`, `perturbation`, `timepoint`, `subject`, -`organism`) into `SampleNode`s — never data-inferred annotations. Ground the -extracted free-text facets through the **existing `OntologyResolver`** (organism -→ NCBITaxon, assay → EFO + `molecular_layer`, tissue → UBERON, disease → MONDO), -and **supply the agent as the resolver's LLM-fallback callback** for strings OLS -can't map (the hook already exists, default off). Mark live tests -`@pytest.mark.integration`; keep unit tests offline by mocking the Azure client. -**Remove the `parce.agent.*` mypy exemption** in `pyproject.toml` once the agent -moves to the normalizer interface. **Blocker risk:** needs Azure creds + an -`az login` session; if absent, build/unit-test the deterministic scaffolding and -log the integration boundary as a blocker rather than working around it. +**PR 6 — Cross-source KG merge.** Merge per-study subgraphs from *different* +sources (CELLxGENE + GEO) into one knowledge graph, deduped by ontology entity ID, +with provenance preserved on edges. The shared-entity machinery already exists: +both normalizers register `BiologicalEntityNode`s keyed by `ontology_id` and emit +edges whose **targets** are those IDs (the originating node differs by source — see +ARCHITECTURE §4 — but the merge keys on targets). Build the merger in `graph/` +(reserved for exactly this since PR 3), take a list of `KnowledgeGraphOutput` +subgraphs → one merged graph, dedup entities by `ontology_id`, keep all +study/dataset/sample nodes, and carry source provenance so a shared entity records +which studies touch it. **Assert a cross-source edge exists in tests** — e.g. a +CELLxGENE study and a GEO study that both touch `UBERON:0002048` (lung) or a shared +`MONDO:` disease become connected through that one entity node. Offline unit tests +only (assemble two canned subgraphs and merge); no network. --- @@ -64,14 +61,26 @@ Each PR is one branch, one focused scope, green CI, and a roadmap update. exact EFO labels to ordered substring keywords** after validating against live EFO (the 10x family never reaches `RNA assay`); ambiguous lineages (bare mass-spec, multi-omic terms) stay `UNKNOWN` by design. -- [ ] **PR 5 — GEO extraction agent (vertical slice).** GEO adapter - (E-utilities/GEOparse) + Azure extraction normalizer emitting the canonical - schema via `response_format`; extract sample covariates from - `characteristics_ch1`. Integration test (marked). This is the agent's real - job; remove the `parce.agent.*` mypy exemption. *(Next up — see top of file.)* +- [x] **PR 5 — GEO extraction agent (vertical slice).** Deterministic `GeoAdapter` + (`sources/geo.py`): fetches GEO Series+Sample SOFT text from the GEO accession + endpoint, parses it (no GEOparse dep), carries `characteristics_ch1` **verbatim** + in the `RawRecord`. Agent-backed `GeoNormalizer` (`normalize/geo.py`): an LLM + (boxed behind the narrow sync `StructuredExtractor` seam, `agent/base.py`) fills + the `GeoExtraction` schema via `response_format` — design covariates only, no + field for any data-inferred annotation. `SampleNode`s now populated (organism + + data_uri read deterministically from structured SOFT fields; condition/ + perturbation/timepoint/subject from the LLM). Facets grounded through the existing + `OntologyResolver`; the agent is **wired as the resolver's LLM fallback** + (`make_ontology_fallback`, opt-in). The concrete Azure agent + (`agent/extraction.py`) bridges the async `agent-framework` API to the sync seam. + **`parce.agent.*` mypy exemption removed** — the whole of `src/parce` is now + type-checked. **Blocker:** live Azure extraction round-trip unverified (no + `AZURE_AI_PROJECT_ENDPOINT` in the headless env); deterministic GEO fetch/parse + verified live. *(GEO keyword `discover` via Entrez deferred to backlog — adapter + `discover` is identity on a `GSEnnnnn`, mirroring CELLxGENE's DOI identity.)* - [ ] **PR 6 — Cross-source KG merge.** Merge CELLxGENE + GEO into one graph linked through shared ontology entities; dedup; provenance on edges. Assert a - cross-source edge exists in tests. + cross-source edge exists in tests. *(Next up — see top of file.)* - [ ] **PR 7 — PRIDE proteomics adapter.** Second modality; prove the interface is modality-general. Adapter + extraction normalizer + integration test. - [ ] **PR 8 — KG export for modeling.** Serialize per-study context + sample @@ -85,6 +94,9 @@ Each PR is one branch, one focused scope, green CI, and a roadmap update. - Graph database backend (Neo4j) vs. flat JSON export — revisit at PR 8. - Discovery agent: given a research theme, propose seed DOIs/accessions across repositories. +- GEO keyword `discover` via Entrez `esearch`+`esummary` (the adapter's `discover` + is currently the identity on a `GSEnnnnn` accession). Pairs with the discovery + agent above. --- @@ -93,6 +105,63 @@ 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-28 — PR 5: GEO extraction agent (vertical slice) + +- Branch `pr5-geo-extraction-agent` off **`origin/main`** (4dcd5b7). +- **Stale-base catch (heeded the memory hazard):** the routine worktree's local + `main` was `64ed4f4`, two merges behind `origin/main` (PR 4 #7 + PR 4b #8). The + worktree's roadmap therefore showed PR 4 as "▶ Next up" — already merged. + `git fetch` + compare to `origin/main` caught it; rebased onto origin and did the + *real* next item (PR 5). Did **not** rebuild PR 4. +- **New files.** `sources/geo.py` (`GeoAdapter` + SOFT parser), `normalize/geo.py` + (`GeoNormalizer` + `GeoExtraction`/`SampleExtraction` schemas), `agent/base.py` + (`StructuredExtractor` Protocol), `agent/extraction.py` (`AzureExtractionAgent` + + `make_ontology_fallback`). Tests: `test_geo_adapter.py`, `test_geo_normalize.py`, + `test_geo_integration.py` (marked). +- **Design decisions (rationale):** + - **Deterministic vs LLM split.** GEO ships some fields structured (per-sample + `organism`, `supplementary_file`) — those are read straight from the record; the + LLM only parses the genuinely free-text `characteristics_ch1` into design + covariates and reads study-level assay/tissue/disease from the prose. Follows + "could a deterministic step do this? then do it" (CLAUDE.md). + - **Sample set is the record's, not the LLM's.** One `SampleNode` per real `GSM`; + the extraction is matched in by `sample_id`, so a dropped/hallucinated sample + can't change graph shape. Extraction failure degrades to samples-without- + covariates (logged), never a crash. + - **No `DatasetNode` for GEO.** A series *is* the study (data is per-sample suppl + files), so `assay`/`molecular_layer` live on `StudyNode` and design-context + + `HAS_SAMPLE` edges originate at the study. Merge (PR 6) keys on entity + `ontology_id` **targets**, so the differing origin vs CELLxGENE is fine. + Recorded in ARCHITECTURE §4. + - **No GEOparse dependency.** The fields needed are a handful of `!`-keys in SOFT + text; a ~40-line parser keeps deps minimal and the parse unit-testable. No dep + changes, so `uv.lock` untouched. + - **`discover` = identity on a `GSEnnnnn`** (mirrors CELLxGENE's DOI identity); + Entrez keyword search → backlog. + - **LLM boxed behind a sync `StructuredExtractor` seam** (`agent/base.py`); the + async `agent-framework` bridge lives only in `agent/extraction.py`. Normalizers + stay sync + offline-testable with a fake extractor. The agent is also wired as + the resolver's LLM fallback (`make_ontology_fallback`, accepts a result only if + the CURIE prefix matches the facet's ontology). ARCHITECTURE §3 updated. + - Added optional `ncbi_email`/`ncbi_api_key` settings (+ `.env.example`); passed + to the adapter, never read by it directly (keeps unit tests Settings-free). +- **mypy:** removed the `parce.agent.*` override — **all of `src/parce` now + type-checked** (28 files; agent-framework/azure are untyped so the glue is `Any` + at the boundary, which is sound here). +- **Gates green (hermetic — no `.env` in the worktree):** ruff check, ruff format + --check (47 files), mypy (28 files), **168 unit tests** (13 integration + deselected). Live `TestLiveGeoFetch` run against the real GEO endpoint — passes + (SOFT parser validated on `GSE10072`). +- **BLOCKER (integration boundary, per protocol):** the live **Azure extraction** + round-trip is **unverified** — this headless env has `az login` but no + `AZURE_AI_PROJECT_ENDPOINT` configured (no worktree `.env`), so + `TestLiveGeoExtraction` skips. The Azure call shape mirrors the previously-working + `agent/curator.py` (`agent.run(prompt, response_format=Model)` → `result.value`). + **Next session with Azure creds:** run `uv run pytest -m integration + tests/test_geo_integration.py` to confirm the live extraction, before relying on + the GEO path in PR 6's cross-source merge. +- **Next session:** PR 6 (cross-source KG merge) — see top of file. + ### 2026-06-26 — PR 4b: Schema refinement (EFO assay term + stored molecular_layer) - Branch `pr4b-schema-refinement` off `main` (35a28ce, the PR 4 merge). Note: the diff --git a/pyproject.toml b/pyproject.toml index 48930b2..4db82fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,11 +63,7 @@ warn_redundant_casts = true warn_unused_ignores = true check_untyped_defs = true -# The Azure agent glue is slated for rewrite as an extraction normalizer (PR 5, -# see docs/ROADMAP.md). Until then it is exempt from type checking so CI tracks -# the stable core. Remove this entry when the agent moves to the normalizer -# interface. (The former ``parce.tools.*`` fetchers were migrated to -# ``parce.sources.*`` in PR 3 and are now fully type-checked.) -[[tool.mypy.overrides]] -module = ["parce.agent.*"] -ignore_errors = true +# No module-level type-check exemptions remain. ``parce.tools.*`` was migrated to +# ``parce.sources.*`` in PR 3; ``parce.agent.*`` moved to the StructuredExtractor +# interface in PR 5 (the Azure glue is now a typed extraction agent). The whole of +# ``src/parce`` is type-checked. diff --git a/src/parce/agent/base.py b/src/parce/agent/base.py new file mode 100644 index 0000000..ab5cfc8 --- /dev/null +++ b/src/parce/agent/base.py @@ -0,0 +1,39 @@ +"""The :class:`StructuredExtractor` contract. + +An extractor is the single boundary between an agent-backed normalizer and a Large +Language Model. It does exactly one thing: take free-text metadata plus a target +Pydantic schema and return a populated, validated instance of that schema — the +*structured extraction* job the LLM is boxed into (docs/ARCHITECTURE.md §3). It +never returns prose. + +Normalizers depend on this narrow, **synchronous** Protocol rather than on a +concrete Azure client, so they can be driven entirely offline in unit tests by +injecting a deterministic fake. The real implementation +(:class:`~parce.agent.extraction.AzureExtractionAgent`) bridges to the async +``agent-framework`` API internally; that async complexity never leaks past this +contract. +""" + +from __future__ import annotations + +from typing import Protocol, TypeVar, runtime_checkable + +from pydantic import BaseModel + +#: The schema an extraction call fills. Bound to ``BaseModel`` so the extractor +#: can validate the model's output against it via ``response_format``. +SchemaT = TypeVar("SchemaT", bound=BaseModel) + + +@runtime_checkable +class StructuredExtractor(Protocol): + """Fills a Pydantic schema from free text via an LLM (structured output only).""" + + def extract(self, instructions: str, content: str, response_model: type[SchemaT]) -> SchemaT: + """Return ``response_model`` populated from ``content`` under ``instructions``. + + ``instructions`` is the system prompt (what to extract and the design-only + constraints); ``content`` is the free-text metadata to read. The result is + a validated ``response_model`` instance — never narrative text. + """ + ... diff --git a/src/parce/agent/extraction.py b/src/parce/agent/extraction.py new file mode 100644 index 0000000..422e9ad --- /dev/null +++ b/src/parce/agent/extraction.py @@ -0,0 +1,126 @@ +"""Azure-backed structured-extraction agent (the only Azure-touching module). + +:class:`AzureExtractionAgent` is the concrete +:class:`~parce.agent.base.StructuredExtractor`: it sends free-text metadata to a +model deployed in an Azure AI Foundry project and constrains the reply to a +Pydantic schema via ``response_format`` — structured output only, never prose +(docs/ARCHITECTURE.md §3). It is used by the agent-backed normalizers (GEO, later +PRIDE) and, optionally, as the ontology resolver's LLM fallback. + +The ``agent-framework`` API is async; this module bridges it to the **synchronous** +:class:`~parce.agent.base.StructuredExtractor` contract with ``asyncio.run`` so the +normalizers and the resolver stay synchronous. Credentials come from +``AzureCliCredential`` (an ``az login`` session) plus the project endpoint / +deployment in :class:`~parce.config.settings.Settings`. Nothing here runs at +import time, so importing the module needs no Azure session — only calling +:meth:`AzureExtractionAgent.extract` does (covered by the marked integration +tests, not unit CI). +""" + +from __future__ import annotations + +import asyncio +import logging + +from agent_framework.azure import AzureAIAgentClient +from azure.identity.aio import AzureCliCredential +from pydantic import BaseModel, ConfigDict + +from parce.agent.base import SchemaT, StructuredExtractor +from parce.config.settings import Settings +from parce.ontology import FACET_ONTOLOGY, Facet, ResolvedTerm +from parce.ontology.resolver import LlmFallback + +logger = logging.getLogger(__name__) + + +class AzureExtractionAgent: + """Structured extractor backed by an Azure AI Foundry model deployment. + + ``settings`` defaults to ``Settings()`` (loaded from the environment / ``.env``); + pass an explicit instance in tests. ``temperature`` defaults to 0 for + reproducible extraction. + """ + + def __init__(self, settings: Settings | None = None, *, temperature: float = 0.0) -> None: + self._settings = settings if settings is not None else Settings() + self._temperature = temperature + + def extract(self, instructions: str, content: str, response_model: type[SchemaT]) -> SchemaT: + """Fill ``response_model`` from ``content`` (synchronous façade over async).""" + return asyncio.run(self._aextract(instructions, content, response_model)) + + async def _aextract( + self, instructions: str, content: str, response_model: type[SchemaT] + ) -> SchemaT: + async with ( + AzureCliCredential() as credential, + AzureAIAgentClient( + project_endpoint=self._settings.azure_ai_project_endpoint, + model_deployment_name=self._settings.azure_ai_model_deployment_name, + credential=credential, + ).as_agent(name="PARCE-extractor", instructions=instructions) as agent, + ): + result = await agent.run( + content, + response_format=response_model, + options={"temperature": self._temperature}, + ) + + # ``result.value`` is the parsed model when the backend honours + # ``response_format``; otherwise fall back to validating ``result.text``. + value = getattr(result, "value", None) + if isinstance(value, response_model): + return value + text = (getattr(result, "text", None) or "").strip() + if not text: + raise ValueError("Extraction agent returned neither a structured value nor text") + return response_model.model_validate_json(text) + + +# -- ontology resolver LLM fallback ------------------------------------------- + + +class _FallbackTerm(BaseModel): + """The shape the LLM returns when asked to ground one hard free-text string.""" + + model_config = ConfigDict(extra="ignore") + + ontology_id: str | None = None + name: str | None = None + + +_FALLBACK_INSTRUCTIONS = ( + "You map a single free-text experiment-design term to one ontology term. " + "Return the term's CURIE (e.g. 'UBERON:0002048') in ontology_id and its label " + "in name, using ONLY the requested ontology. If you are not confident the term " + "exists in that ontology, return null for both fields — never invent an ID." +) + + +def make_ontology_fallback(extractor: StructuredExtractor) -> LlmFallback: + """Adapt a :class:`StructuredExtractor` into the resolver's LLM-fallback hook. + + The resolver calls this only for strings its deterministic OLS lookups could + not map (``llm_fallback`` is ``None`` by default, so wiring this in is opt-in). + The result is accepted only when the returned CURIE carries the facet's + expected ontology prefix, so a wrong-ontology or hallucinated ID is dropped + rather than grounded. + """ + + def _fallback(text: str, facet: Facet) -> ResolvedTerm | None: + primary = FACET_ONTOLOGY[facet].primary + content = ( + f"Ontology: {primary.prefix} ({primary.title})\nFacet: {facet.value}\nTerm: {text}" + ) + try: + result = extractor.extract(_FALLBACK_INSTRUCTIONS, content, _FallbackTerm) + except Exception as exc: + logger.warning("LLM ontology fallback failed for %r (%s): %s", text, facet, exc) + return None + if result.ontology_id and result.ontology_id.startswith(f"{primary.prefix}:"): + return ResolvedTerm(ontology_id=result.ontology_id, name=result.name or text) + logger.info("LLM fallback returned no usable term for %r (%s)", text, facet) + return None + + return _fallback diff --git a/src/parce/config/settings.py b/src/parce/config/settings.py index 5f506c9..f5d0102 100644 --- a/src/parce/config/settings.py +++ b/src/parce/config/settings.py @@ -17,3 +17,9 @@ class Settings(BaseSettings): azure_ai_project_endpoint: str azure_ai_model_deployment_name: str = "mistral-large" max_retries: int = 3 + + # NCBI E-utilities / GEO courtesy parameters (optional). NCBI asks callers to + # identify themselves; supplying these raises rate limits. Neither is required + # for GEO fetches to work. + ncbi_email: str | None = None + ncbi_api_key: str | None = None diff --git a/src/parce/main.py b/src/parce/main.py index 54c891f..0d83b72 100644 --- a/src/parce/main.py +++ b/src/parce/main.py @@ -1,21 +1,25 @@ -"""PARCE entry point — deterministic CELLxGENE ingestion. +"""PARCE entry point — multi-source omics ingestion. -Drives one source through the adapter → normalizer pipeline and writes the -canonical knowledge graph to disk. CELLxGENE Census already ships -ontology-grounded metadata, so this path is fully deterministic — there is no -LLM here (the extraction agent enters with GEO in PR 5). +Drives a source through its adapter → normalizer pipeline and writes the canonical +knowledge graph to disk. Two source paths exist: -Run with:: - - python -m parce.main +* ``cellxgene`` (default) — fully **deterministic**: CELLxGENE Census already ships + ontology-grounded metadata, so there is no LLM here. +* ``geo`` — the first **agent-backed** path: GEO's free-text sample metadata is + extracted into the canonical schema by the Azure extraction agent, then grounded + through the shared ontology resolver (docs/ARCHITECTURE.md §2-3). This path needs + Azure credentials (an ``az login`` session) at runtime. -or, after ``pip install -e .``:: +Run with:: - parce + parce # CELLxGENE, default DOI + parce cellxgene + parce geo GSE10072 # GEO series (requires Azure creds) """ from __future__ import annotations +import argparse import json import logging from pathlib import Path @@ -30,13 +34,32 @@ _DEFAULT_DOI = "10.1038/s41586-023-05869-0" -def run(doi: str = _DEFAULT_DOI) -> None: - """Fetch, normalize and persist the KG for a CELLxGENE collection DOI.""" +def _configure_logging() -> None: logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s - %(message)s", ) + +def _persist(kg: KnowledgeGraphOutput, *, filename: str = "output.json") -> None: + """Print a summary of ``kg`` and write it to the output directory.""" + print("Knowledge graph constructed successfully:") + 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)}") + + _OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + out_path = _OUTPUT_DIR / filename + out_path.write_text(json.dumps(kg.model_dump(mode="json"), indent=2)) + print(f"\nSaved to {out_path}") + + +def run(doi: str = _DEFAULT_DOI) -> None: + """Fetch, normalize and persist the KG for a CELLxGENE collection DOI.""" + _configure_logging() + adapter = CellxgeneAdapter() normalizer = CellxgeneNormalizer() @@ -75,23 +98,83 @@ def run(doi: str = _DEFAULT_DOI) -> None: # PR 3 is single-source/single-study, so there is exactly one subgraph here. # Merging multiple subgraphs into one graph (deduped by ontology ID) is PR 6. logger.info("Step 3/3: Writing knowledge graph") - kg = subgraphs[0] + _persist(subgraphs[0]) + + +def run_geo(accession: str) -> None: + """Fetch, extract, ground and persist the KG for a GEO series accession. + + This is the agent-backed path: it builds the Azure extraction agent (needs an + ``az login`` session + the Azure settings) and wires it both as the GEO + normalizer's extractor *and* as the ontology resolver's LLM fallback for hard + free-text terms. Imports of the Azure glue are deferred to here so the default + CELLxGENE path and ``import parce.main`` need no Azure session. + """ + _configure_logging() + + # Deferred imports: keep the Azure dependency off the default path / module import. + from parce.agent.extraction import AzureExtractionAgent, make_ontology_fallback + from parce.config.settings import Settings + from parce.normalize.geo import GeoNormalizer + from parce.ontology import OntologyResolver + from parce.sources.geo import GeoAdapter + + settings = Settings() + adapter = GeoAdapter(email=settings.ncbi_email, api_key=settings.ncbi_api_key) + extractor = AzureExtractionAgent(settings) + resolver = OntologyResolver(llm_fallback=make_ontology_fallback(extractor)) + normalizer = GeoNormalizer(extractor, resolver=resolver) + + refs = adapter.discover(accession) + logger.info("Step 1/3: Discovered %d GEO reference(s) for query=%s", len(refs), accession) + if not refs: + logger.error("No GEO series found for query=%s", accession) + return - print("Knowledge graph constructed successfully:") - 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)}") + logger.info("Step 2/3: Fetching, extracting and normalizing %d reference(s)", len(refs)) + subgraphs: list[KnowledgeGraphOutput] = [] + for ref in refs: + record = adapter.fetch(ref) + if not record.payload.get("samples"): + logger.warning("No samples for GEO ref=%s; skipping", ref) + continue + subgraphs.append(normalizer.normalize(record)) - _OUTPUT_DIR.mkdir(parents=True, exist_ok=True) - out_path = _OUTPUT_DIR / "output.json" - out_path.write_text(json.dumps(kg.model_dump(mode="json"), indent=2)) - print(f"\nSaved to {out_path}") + if not subgraphs: + logger.error("No samples fetched for query=%s", accession) + return + + logger.info("Step 3/3: Writing knowledge graph") + _persist(subgraphs[0], filename=f"{accession.upper()}.json") def main() -> None: - run() + parser = argparse.ArgumentParser( + prog="parce", description="Harvest omics experiments into a KG." + ) + parser.add_argument( + "source", + nargs="?", + default="cellxgene", + choices=["cellxgene", "geo"], + help="Which source to ingest (default: cellxgene).", + ) + parser.add_argument( + "query", + nargs="?", + default=None, + help="CELLxGENE collection DOI or GEO series accession (GSEnnnnn).", + ) + args = parser.parse_args() + + if args.source == "geo": + if not args.query: + parser.error( + "the 'geo' source requires a GEO series accession, e.g. parce geo GSE10072" + ) + run_geo(args.query) + else: + run(args.query or _DEFAULT_DOI) if __name__ == "__main__": diff --git a/src/parce/normalize/geo.py b/src/parce/normalize/geo.py new file mode 100644 index 0000000..879530b --- /dev/null +++ b/src/parce/normalize/geo.py @@ -0,0 +1,310 @@ +"""Agent-backed normalizer: a GEO ``RawRecord`` → canonical KG nodes. + +GEO is the project's first *unstructured* source, so this is the first normalizer +with an LLM in its path. The division of labour is deliberate and follows the +"could a deterministic step do this?" rule (CLAUDE.md): + +* **Deterministic** — fields GEO already ships structured are read straight from + the record: each sample's ``organism`` (``!Sample_organism_ch1``) and its raw + data URI (``!Sample_supplementary_file``). The set of samples (one ``SampleNode`` + per real ``GSM``) is the record's, never the LLM's, so a hallucinated or dropped + sample cannot change the graph's shape. +* **LLM (structured extraction only)** — the genuinely free-text job: turning the + ``!Sample_characteristics_ch1`` lines into the canonical **design** covariates + (``condition``/``perturbation``/``timepoint``/``subject``) and reading the + study-level ``assay``/``tissue``/``disease`` out of the series prose. The agent + is constrained by ``response_format`` to the :class:`GeoExtraction` schema — it + cannot emit prose, and the schema has **no field for any data-inferred + annotation** (cell type, clusters), so leakage is structurally impossible + (docs/ARCHITECTURE.md §1). +* **Shared ontology stage** — every extracted free-text facet is grounded through + the same :class:`~parce.ontology.base.OntologyService` the CELLxGENE path uses + (organism→NCBITaxon, assay→EFO + ``molecular_layer``, tissue→UBERON, + disease→MONDO), so GEO studies land on the *same* IDs and link to CELLxGENE + through shared entity nodes. + +GEO has no distinct dataset artifact (unlike CELLxGENE's Census datasets): the +series *is* the study and its data lives in per-sample supplementary files. So no +``DatasetNode`` is emitted; ``assay``/``molecular_layer`` live on the ``StudyNode`` +and the design-context + ``HAS_SAMPLE`` edges originate at the study. Cross-source +linking is unaffected — it flows through shared entity ``ontology_id`` targets, +not the originating node (docs/ARCHITECTURE.md §4). +""" + +from __future__ import annotations + +import logging +from collections import Counter + +from pydantic import BaseModel, ConfigDict, Field + +from parce.agent.base import StructuredExtractor +from parce.models.graph_schema import ( + BiologicalEntityNode, + EntityType, + GraphEdge, + KnowledgeGraphOutput, + MolecularLayer, + SampleNode, + StudyNode, +) +from parce.models.raw_record import RawRecord +from parce.ontology import Facet, OntologyResolver, OntologyService + +logger = logging.getLogger(__name__) + +_UNKNOWN = "unknown" + + +# -- the agent's response_format schema --------------------------------------- +# These are agent IO models, not canonical KG nodes, so they use ``extra="ignore"`` +# (be lenient with the model's output) rather than the KG models' ``extra="forbid"``. +# They carry only experiment-*design* fields — there is deliberately no field for a +# data-inferred annotation, so the schema itself forbids leakage. + + +class SampleExtraction(BaseModel): + """Per-sample design covariates the LLM reads from ``characteristics_ch1``.""" + + model_config = ConfigDict(extra="ignore") + + sample_id: str = Field(..., description="The GEO sample accession (GSM...), echoed verbatim.") + condition: str | None = Field( + default=None, description="Experimental condition as designed (e.g. 'tumor', 'control')." + ) + perturbation: str | None = Field( + default=None, description="Designed perturbation: a drug, dose, genetic knockout, etc." + ) + timepoint: str | None = Field( + default=None, description="Designed sampling timepoint (e.g. '0h', 'day 7')." + ) + subject: str | None = Field( + default=None, description="Subject/donor/patient/replicate identifier." + ) + + +class GeoExtraction(BaseModel): + """Structured view of one GEO series the extraction agent must return.""" + + model_config = ConfigDict(extra="ignore") + + assay: str | None = Field( + default=None, + description="The assay/technology in plain words (e.g. 'RNA-seq', 'microarray').", + ) + tissue: str | None = Field( + default=None, description="The tissue/anatomical source studied (e.g. 'lung')." + ) + disease: str | None = Field( + default=None, description="The disease/condition studied (e.g. 'lung adenocarcinoma')." + ) + samples: list[SampleExtraction] = Field(default_factory=list) + + +GEO_EXTRACTION_INSTRUCTIONS = ( + "You extract experiment-DESIGN metadata from a GEO series into the given schema. " + "Use only what the text states about how the experiment was DESIGNED: the assay/" + "technology, the tissue/anatomy sampled, the disease/condition studied, and per-" + "sample condition, perturbation (drug/dose/genetic), timepoint, and subject/donor. " + "NEVER infer or output cell types, cluster labels, or anything derived from the " + "measured data — only design variables. Echo each sample_id (GSM...) exactly. " + "Leave any field null when the text does not state it; do not guess." +) + + +def _serialise_for_agent(record: RawRecord) -> str: + """Render the record as the plain-text prompt content the agent reads.""" + series = record.payload.get("series", {}) + lines: list[str] = [ + f"GEO Series: {record.study_id}", + f"Title: {record.title}", + f"Type: {'; '.join(series.get('type', [])) or _UNKNOWN}", + f"Summary: {' '.join(series.get('summary', [])) or _UNKNOWN}", + f"Overall design: {series.get('overall_design', _UNKNOWN)}", + "", + "Samples:", + ] + for s in record.payload.get("samples", []): + chars = "; ".join(s.get("characteristics", [])) or "(none)" + lines.append( + f"- {s.get('sample_id', _UNKNOWN)} | title: {s.get('title', '')} | " + f"source: {s.get('source_name', '')} | characteristics: {chars}" + ) + return "\n".join(lines) + + +class GeoNormalizer: + """:class:`~parce.normalize.base.Normalizer` for GEO ``RawRecord``s (LLM-backed). + + Requires a :class:`~parce.agent.base.StructuredExtractor` (the LLM seam — inject + a fake to test offline). The :class:`~parce.ontology.base.OntologyService` + defaults to a real :class:`~parce.ontology.resolver.OntologyResolver`; inject a + fake to keep grounding offline too. + """ + + def __init__( + self, extractor: StructuredExtractor, *, resolver: OntologyService | None = None + ) -> None: + self._extractor = extractor + self._resolver: OntologyService = resolver if resolver is not None else OntologyResolver() + + def normalize(self, record: RawRecord) -> KnowledgeGraphOutput: + """Assemble the canonical single-study subgraph for one GEO series.""" + study_id = record.study_id + raw_samples = record.payload.get("samples", []) + + extraction = self._extract(record) + covariates = {s.sample_id: s for s in extraction.samples} + + entity_registry: dict[str, BiologicalEntityNode] = {} + edges: list[GraphEdge] = [] + samples: list[SampleNode] = [] + + # -- samples: structured fields from the record, design covariates from + # the LLM. The record's sample set is authoritative. + organism_counts: Counter[str] = Counter() + for raw in raw_samples: + sample_id = raw.get("sample_id") + if not sample_id: + continue + organism = raw.get("organism") or None + if organism: + organism_counts[organism] += 1 + + cov = covariates.get(sample_id) + samples.append( + SampleNode( + sample_id=sample_id, + data_uri=raw.get("supplementary_file") or None, + organism=organism, + condition=cov.condition if cov else None, + perturbation=cov.perturbation if cov else None, + timepoint=cov.timepoint if cov else None, + subject=cov.subject if cov else None, + ) + ) + edges.append( + GraphEdge(source_id=study_id, target_id=sample_id, relation_type="HAS_SAMPLE") + ) + + # -- study-level assay → EFO term + molecular_layer. + assay_id, assay_label = self._ground(extraction.assay, Facet.ASSAY) + layer = self._molecular_layer(assay_id, assay_label) + if assay_id != _UNKNOWN: + self._register( + entity_registry, + EntityType.ASSAY, + assay_id, + assay_label, + study_id, + "MEASURED_WITH", + edges, + ) + + # -- study-level tissue / disease. + tissue_id, tissue_label = self._ground(extraction.tissue, Facet.TISSUE) + if tissue_id != _UNKNOWN: + self._register( + entity_registry, + EntityType.TISSUE, + tissue_id, + tissue_label, + study_id, + "HAS_TISSUE", + edges, + ) + disease_id, disease_label = self._ground(extraction.disease, Facet.DISEASE) + if disease_id != _UNKNOWN: + self._register( + entity_registry, + EntityType.DISEASE, + disease_id, + disease_label, + study_id, + "HAS_CONDITION", + edges, + ) + + # -- dominant organism → Species entity + STUDIES edge. + if organism_counts: + dominant_organism = organism_counts.most_common(1)[0][0] + species = self._resolver.resolve_term(dominant_organism, Facet.ORGANISM) + if species is not None: + self._register( + entity_registry, + EntityType.SPECIES, + species.ontology_id, + species.name, + study_id, + "STUDIES", + edges, + ) + + study = StudyNode( + study_id=study_id, + title=record.title, + source=record.source, + assay=assay_id, + molecular_layer=layer, + ) + + kg = KnowledgeGraphOutput( + studies=[study], + datasets=[], + samples=samples, + biological_entities=list(entity_registry.values()), + edges=edges, + ) + logger.info( + "Normalized GEO study=%s: samples=%d entities=%d edges=%d assay=%s layer=%s", + study_id, + len(kg.samples), + len(kg.biological_entities), + len(kg.edges), + assay_id, + layer, + ) + return kg + + # -- helpers --------------------------------------------------------------- + def _extract(self, record: RawRecord) -> GeoExtraction: + """Run the extraction agent; degrade to an empty extraction on failure.""" + try: + return self._extractor.extract( + GEO_EXTRACTION_INSTRUCTIONS, _serialise_for_agent(record), GeoExtraction + ) + except Exception as exc: + logger.warning("GEO extraction failed for %s: %s", record.study_id, exc) + return GeoExtraction() + + def _ground(self, text: str | None, facet: Facet) -> tuple[str, str]: + """Ground free text to ``(ontology_id, label)``; ``('unknown', text)`` if not.""" + if not text or not text.strip(): + return _UNKNOWN, text or _UNKNOWN + term = self._resolver.resolve_term(text, facet) + if term is None: + return _UNKNOWN, text + return term.ontology_id, term.name + + def _molecular_layer(self, assay_id: str, assay_label: str) -> MolecularLayer: + """Derive the molecular layer; only EFO assay terms carry a walkable lineage.""" + if assay_id.startswith("EFO:"): + return self._resolver.molecular_layer(assay_id, assay_label=assay_label) + return MolecularLayer.UNKNOWN + + @staticmethod + def _register( + registry: dict[str, BiologicalEntityNode], + entity_type: EntityType, + ontology_id: str, + name: str, + source_id: str, + relation: str, + edges: list[GraphEdge], + ) -> None: + """Register an entity (deduped by ontology ID) and add its edge from source.""" + if ontology_id not in registry: + registry[ontology_id] = BiologicalEntityNode( + entity_type=entity_type, ontology_id=ontology_id, name=name + ) + edges.append(GraphEdge(source_id=source_id, target_id=ontology_id, relation_type=relation)) diff --git a/src/parce/sources/geo.py b/src/parce/sources/geo.py new file mode 100644 index 0000000..b86309b --- /dev/null +++ b/src/parce/sources/geo.py @@ -0,0 +1,226 @@ +"""GEO (NCBI Gene Expression Omnibus) source adapter — deterministic fetch only. + +The adapter pulls a GEO **Series** (``GSEnnnnn``) and its **Samples** (``GSMnnnnn``) +as SOFT-format text from the public GEO accession endpoint, parses the fields a +normalizer needs, and emits a source-shaped +:class:`~parce.models.raw_record.RawRecord`. It performs **no interpretation**: +the messy, semi-structured ``!Sample_characteristics_ch1`` lines are carried +verbatim in the payload. Turning that free text into canonical design covariates +is the matching agent-backed :class:`~parce.normalize.geo.GeoNormalizer`'s job — +this is GEO's reason for being the first LLM-using source (docs/ARCHITECTURE.md §2). + +Why a direct SOFT fetch + a small parser rather than GEOparse: the only fields we +need (series title/type/summary, and per-sample title/organism/source/ +characteristics/supplementary-file) are a handful of ``!``-prefixed keys in the +SOFT text. A focused ~40-line parser keeps the dependency surface minimal and the +parse fully unit-testable against a captured fixture, with no large data-table +download (``view=quick`` omits the matrices). + +The GEO endpoint is the only network IO here; it is wrapped in the shared +:func:`parce.sources._retry.with_retries` helper. +""" + +from __future__ import annotations + +import logging +import re +from typing import Any + +import requests + +from parce.models.raw_record import RawRecord +from parce.sources._retry import with_retries + +logger = logging.getLogger(__name__) + +SOURCE_NAME = "GEO" + +_GEO_ACC_URL = "https://www.ncbi.nlm.nih.gov/geo/query/acc.cgi" +_DEFAULT_TIMEOUT = 60 +#: Bound the per-series sample count handed downstream so a 1000-sample series +#: does not balloon a single LLM extraction call. Truncation is flagged on the +#: payload and logged (never silent). +_DEFAULT_MAX_SAMPLES = 200 + +_GSE_RE = re.compile(r"^GSE\d+$", re.IGNORECASE) + + +def _looks_like_series(query: str) -> bool: + """True if ``query`` is a GEO Series accession (``GSEnnnnn``).""" + return bool(_GSE_RE.match(query.strip())) + + +class GeoAdapter: + """:class:`~parce.sources.base.SourceAdapter` for NCBI GEO Series. + + Deterministic: it fetches and parses metadata only. ``email``/``api_key`` are + optional NCBI E-utilities courtesy parameters (higher rate limits, contact on + file); pass them from settings in production, omit them in tests. + """ + + source_name = SOURCE_NAME + + def __init__( + self, + *, + email: str | None = None, + api_key: str | None = None, + base_url: str = _GEO_ACC_URL, + timeout: int = _DEFAULT_TIMEOUT, + ) -> None: + self._email = email + self._api_key = api_key + self._base_url = base_url + self._timeout = timeout + + def discover(self, query: str) -> list[str]: + """Resolve ``query`` to GEO Series references the adapter can ``fetch``. + + A GEO Series accession is the reference, so ``discover`` is the identity on + a ``GSEnnnnn`` string (mirroring the CELLxGENE adapter's DOI identity). + Free-text keyword search via Entrez ``esearch`` is backlog + (see docs/ROADMAP.md); a non-accession query is rejected here rather than + silently returning nothing useful. + """ + q = query.strip() + if not _looks_like_series(q): + logger.warning( + "GEO discover only supports a Series accession (GSEnnnnn); got %r", query + ) + return [] + return [q.upper()] + + def fetch(self, ref: str, *, max_samples: int = _DEFAULT_MAX_SAMPLES) -> RawRecord: + """Fetch one GEO Series accession into a source-shaped ``RawRecord``. + + The payload carries the parsed ``series`` fields and a ``samples`` list + (each with its raw ``characteristics`` lines). When more than + ``max_samples`` samples exist, the list is truncated and ``truncated`` is + set on the payload (and logged). + """ + soft = self._fetch_soft(ref) + series, samples = _parse_soft(soft) + + truncated = len(samples) > max_samples + if truncated: + logger.warning( + "GEO series %s has %d samples; truncating to %d for extraction", + ref, + len(samples), + max_samples, + ) + samples = samples[:max_samples] + + payload: dict[str, Any] = { + "series": series, + "samples": samples, + "truncated": truncated, + } + logger.info("Fetched GEO series %s: samples=%d truncated=%s", ref, len(samples), truncated) + + return RawRecord( + source=self.source_name, + study_id=ref.upper(), + title=series.get("title", ""), + payload=payload, + ) + + def _fetch_soft(self, accession: str) -> str: + """Fetch the full SOFT text (series + samples) for ``accession``.""" + params: dict[str, str] = { + "acc": accession, + "targ": "all", # series + platform + all samples in one document + "form": "text", + "view": "quick", # metadata headers only; omit the big data matrices + } + if self._email: + params["email"] = self._email + if self._api_key: + params["api_key"] = self._api_key + + def _request() -> requests.Response: + resp = requests.get(self._base_url, params=params, timeout=self._timeout) + resp.raise_for_status() + return resp + + resp = with_retries(_request, description=f"GEO SOFT fetch acc={accession}") + return resp.text + + +# -- SOFT parsing -------------------------------------------------------------- +# SOFT is a flat, line-oriented format. Entity blocks open with a ``^`` marker +# (``^SERIES``, ``^PLATFORM``, ``^SAMPLE``); within a block, metadata lines are +# ``!Key = value`` and a key may repeat (summary, type, characteristics). + +_SERIES_SINGLE = { + "!Series_title": "title", + "!Series_overall_design": "overall_design", + "!Series_geo_accession": "geo_accession", +} +_SERIES_MULTI = { + "!Series_summary": "summary", + "!Series_type": "type", + "!Series_pubmed_id": "pubmed_id", +} + +_SAMPLE_SINGLE = { + "!Sample_title": "title", + "!Sample_geo_accession": "sample_id", + "!Sample_source_name_ch1": "source_name", + "!Sample_organism_ch1": "organism", + "!Sample_supplementary_file": "supplementary_file", +} + + +def _split_kv(line: str) -> tuple[str, str] | None: + """Split a SOFT ``!Key = value`` (or ``!Key=value``) line; ``None`` if not one.""" + if not line.startswith("!"): + return None + key, sep, value = line.partition("=") + if not sep: + return None + return key.strip(), value.strip() + + +def _parse_soft(text: str) -> tuple[dict[str, Any], list[dict[str, Any]]]: + """Parse SOFT text into ``(series_fields, samples)``. + + ``series_fields`` collapses single-valued keys to a string and multi-valued + keys (``summary``/``type``/``pubmed_id``) to a list. Each sample dict carries + its single-valued fields plus a ``characteristics`` list of the raw + ``key: value`` strings (left for the LLM to normalize). + """ + series: dict[str, Any] = {} + samples: list[dict[str, Any]] = [] + current: dict[str, Any] | None = None # the SAMPLE block being filled, if any + + for raw_line in text.splitlines(): + line = raw_line.rstrip("\n") + if line.startswith("^"): + marker = line[1:].split("=", 1)[0].strip().upper() + if marker == "SAMPLE": + current = {"characteristics": []} + samples.append(current) + else: # ^SERIES / ^PLATFORM — leave the sample context + current = None + continue + + kv = _split_kv(line) + if kv is None: + continue + key, value = kv + + if current is not None: # inside a SAMPLE block + if key == "!Sample_characteristics_ch1": + current["characteristics"].append(value) + elif key in _SAMPLE_SINGLE: + current.setdefault(_SAMPLE_SINGLE[key], value) + continue + + # SERIES (or PLATFORM, which we ignore) context. + if key in _SERIES_SINGLE: + series.setdefault(_SERIES_SINGLE[key], value) + elif key in _SERIES_MULTI: + series.setdefault(_SERIES_MULTI[key], []).append(value) + + return series, samples diff --git a/tests/test_geo_adapter.py b/tests/test_geo_adapter.py new file mode 100644 index 0000000..71a6e86 --- /dev/null +++ b/tests/test_geo_adapter.py @@ -0,0 +1,156 @@ +"""Unit tests for the GEO source adapter and its SOFT parser. + +All GEO network IO is mocked — these tests stay offline. The SOFT fixture mirrors +the real ``form=text&view=quick`` layout (a ``^SERIES`` block, a ``^PLATFORM`` +block to be ignored, then ``^SAMPLE`` blocks). +""" + +from __future__ import annotations + +from unittest.mock import patch + +from parce.models.raw_record import RawRecord +from parce.sources.base import SourceAdapter +from parce.sources.geo import GeoAdapter, _parse_soft + +# A compact but realistic SOFT document: 2 samples, repeated summary/type keys, +# a PLATFORM block whose !Platform_* keys must not leak into the series fields. +_SOFT = """^SERIES = GSE99999 +!Series_title = Smoking and lung adenocarcinoma +!Series_geo_accession = GSE99999 +!Series_summary = We profiled tumor and normal lung tissue. +!Series_summary = Keywords: comparative genomics +!Series_overall_design = 2 tumor and 2 normal samples from 2 subjects. +!Series_type = Expression profiling by array +!Series_pubmed_id = 18297132 +!Series_sample_id = GSM000001 +!Series_sample_id = GSM000002 +^PLATFORM = GPL96 +!Platform_title = Affymetrix HG-U133A +!Platform_organism = Homo sapiens +^SAMPLE = GSM000001 +!Sample_title = Lung Tumor A +!Sample_geo_accession = GSM000001 +!Sample_type = RNA +!Sample_source_name_ch1 = Adenocarcinoma of the Lung +!Sample_organism_ch1 = Homo sapiens +!Sample_characteristics_ch1 = gender: Male +!Sample_characteristics_ch1 = tissue: tumor +!Sample_supplementary_file = ftp://ftp.ncbi.nlm.nih.gov/geo/samples/GSM000001/suppl/GSM000001.CEL.gz +^SAMPLE = GSM000002 +!Sample_title = Lung Normal A +!Sample_geo_accession = GSM000002 +!Sample_type = RNA +!Sample_source_name_ch1 = Noninvolved Lung +!Sample_organism_ch1 = Homo sapiens +!Sample_characteristics_ch1 = gender: Female +!Sample_characteristics_ch1 = tissue: normal +!Sample_supplementary_file = ftp://ftp.ncbi.nlm.nih.gov/geo/samples/GSM000002/suppl/GSM000002.CEL.gz +""" + + +class TestParseSoft: + def test_series_single_and_multi_fields(self): + series, _ = _parse_soft(_SOFT) + assert series["title"] == "Smoking and lung adenocarcinoma" + assert series["geo_accession"] == "GSE99999" + assert series["overall_design"] == "2 tumor and 2 normal samples from 2 subjects." + # Repeated keys collect into lists. + assert series["summary"] == [ + "We profiled tumor and normal lung tissue.", + "Keywords: comparative genomics", + ] + assert series["type"] == ["Expression profiling by array"] + assert series["pubmed_id"] == ["18297132"] + + def test_platform_block_does_not_pollute_series(self): + series, _ = _parse_soft(_SOFT) + # !Platform_organism must not become a series/sample field. + assert "organism" not in series + + def test_samples_parsed_with_characteristics(self): + _, samples = _parse_soft(_SOFT) + assert [s["sample_id"] for s in samples] == ["GSM000001", "GSM000002"] + + first = samples[0] + assert first["title"] == "Lung Tumor A" + assert first["source_name"] == "Adenocarcinoma of the Lung" + assert first["organism"] == "Homo sapiens" + assert first["characteristics"] == ["gender: Male", "tissue: tumor"] + assert first["supplementary_file"].endswith("GSM000001.CEL.gz") + + def test_characteristics_kept_verbatim(self): + """The adapter must not interpret characteristics — they pass through raw.""" + _, samples = _parse_soft(_SOFT) + assert samples[1]["characteristics"] == ["gender: Female", "tissue: normal"] + + def test_handles_no_space_around_equals(self): + soft = "^SAMPLE = GSM1\n!Sample_geo_accession=GSM1\n!Sample_characteristics_ch1=stage:IIB\n" + _, samples = _parse_soft(soft) + assert samples[0]["sample_id"] == "GSM1" + assert samples[0]["characteristics"] == ["stage:IIB"] + + +class TestGeoAdapterDiscover: + def test_discover_identity_on_accession(self): + assert GeoAdapter().discover("GSE99999") == ["GSE99999"] + + def test_discover_uppercases_and_strips(self): + assert GeoAdapter().discover(" gse123 ") == ["GSE123"] + + def test_discover_rejects_non_accession(self): + assert GeoAdapter().discover("lung cancer") == [] + + +class TestGeoAdapterFetch: + def test_fetch_builds_raw_record(self): + with patch.object(GeoAdapter, "_fetch_soft", return_value=_SOFT): + record = GeoAdapter().fetch("GSE99999") + + assert isinstance(record, RawRecord) + assert record.source == "GEO" + assert record.study_id == "GSE99999" + assert record.title == "Smoking and lung adenocarcinoma" + assert len(record.payload["samples"]) == 2 + assert record.payload["truncated"] is False + assert record.payload["series"]["type"] == ["Expression profiling by array"] + + def test_fetch_uppercases_accession(self): + with patch.object(GeoAdapter, "_fetch_soft", return_value=_SOFT): + record = GeoAdapter().fetch("gse99999") + assert record.study_id == "GSE99999" + + def test_fetch_truncates_to_max_samples(self): + with patch.object(GeoAdapter, "_fetch_soft", return_value=_SOFT): + record = GeoAdapter().fetch("GSE99999", max_samples=1) + assert len(record.payload["samples"]) == 1 + assert record.payload["truncated"] is True + assert record.payload["samples"][0]["sample_id"] == "GSM000001" + + def test_fetch_passes_courtesy_params(self): + captured: dict[str, object] = {} + + class _Resp: + text = _SOFT + + def raise_for_status(self) -> None: + return None + + def _fake_get(url, params, timeout): + captured["url"] = url + captured["params"] = params + return _Resp() + + with patch("parce.sources.geo.requests.get", side_effect=_fake_get): + GeoAdapter(email="me@example.com", api_key="KEY").fetch("GSE99999") + + params = captured["params"] + assert params["acc"] == "GSE99999" + assert params["targ"] == "all" + assert params["email"] == "me@example.com" + assert params["api_key"] == "KEY" + + +class TestProtocolConformance: + def test_adapter_satisfies_source_adapter(self): + assert isinstance(GeoAdapter(), SourceAdapter) diff --git a/tests/test_geo_integration.py b/tests/test_geo_integration.py new file mode 100644 index 0000000..78dbcc4 --- /dev/null +++ b/tests/test_geo_integration.py @@ -0,0 +1,61 @@ +"""Live integration tests for the GEO vertical slice. Excluded from CI. + +Two tiers, both marked ``integration``: + +* the GEO **fetch/parse** test hits only the public GEO endpoint (no credentials); +* the **extraction** test additionally needs Azure (an ``az login`` session + the + ``AZURE_AI_*`` settings) and is skipped when those are absent. + +Run with:: + + uv run pytest -m integration tests/test_geo_integration.py +""" + +from __future__ import annotations + +import os + +import pytest + +from parce.normalize.geo import GeoNormalizer +from parce.sources.geo import GeoAdapter + +pytestmark = pytest.mark.integration + +# A small, stable, public GEO series used as the live fixture. +_GSE = "GSE10072" + + +class TestLiveGeoFetch: + def test_fetch_real_series(self): + record = GeoAdapter().fetch(_GSE, max_samples=5) + assert record.source == "GEO" + assert record.study_id == _GSE + assert record.title # series has a title + samples = record.payload["samples"] + assert 1 <= len(samples) <= 5 + # Real GEO samples carry organism + characteristics. + assert samples[0]["organism"] == "Homo sapiens" + assert samples[0]["characteristics"] + + +def _azure_configured() -> bool: + return bool(os.environ.get("AZURE_AI_PROJECT_ENDPOINT")) + + +@pytest.mark.skipif(not _azure_configured(), reason="Azure settings/credentials not configured") +class TestLiveGeoExtraction: + def test_fetch_extract_normalize(self): + # Deferred import so collection does not require the Azure deps to resolve. + from parce.agent.extraction import AzureExtractionAgent + from parce.ontology import OntologyResolver + + record = GeoAdapter().fetch(_GSE, max_samples=5) + normalizer = GeoNormalizer(AzureExtractionAgent(), resolver=OntologyResolver()) + kg = normalizer.normalize(record) + + assert len(kg.studies) == 1 + assert kg.studies[0].source == "GEO" + assert len(kg.samples) >= 1 + # No data-inferred annotations ever appear. + assert all(not e.ontology_id.startswith("CL:") for e in kg.biological_entities) diff --git a/tests/test_geo_normalize.py b/tests/test_geo_normalize.py new file mode 100644 index 0000000..effd670 --- /dev/null +++ b/tests/test_geo_normalize.py @@ -0,0 +1,256 @@ +"""Unit tests for the agent-backed GEO normalizer. + +Offline: the LLM is a deterministic fake :class:`StructuredExtractor` returning a +canned :class:`GeoExtraction`, and ontology grounding is a fake resolver. No +network, no Azure. +""" + +from __future__ import annotations + +import pytest + +from parce.agent.base import SchemaT, StructuredExtractor +from parce.models.graph_schema import EntityType, KnowledgeGraphOutput, MolecularLayer +from parce.models.raw_record import RawRecord +from parce.normalize.base import Normalizer +from parce.normalize.geo import GeoExtraction, GeoNormalizer, SampleExtraction +from parce.ontology import Facet, ResolvedTerm + +# Grounding table the fake resolver uses (organism/assay/tissue/disease). +_TERMS = { + (Facet.ORGANISM, "Homo sapiens"): ResolvedTerm("NCBITaxon:9606", "Homo sapiens"), + (Facet.ASSAY, "microarray"): ResolvedTerm("EFO:0002772", "microarray"), + (Facet.TISSUE, "lung"): ResolvedTerm("UBERON:0002048", "lung"), + (Facet.DISEASE, "lung adenocarcinoma"): ResolvedTerm("MONDO:0005061", "lung adenocarcinoma"), +} + + +class _FakeResolver: + """Offline OntologyService stand-in.""" + + def resolve_term(self, text: str, facet: Facet) -> ResolvedTerm | None: + return _TERMS.get((facet, text)) + + def molecular_layer(self, assay_id: str, *, assay_label: str | None = None) -> MolecularLayer: + return ( + MolecularLayer.TRANSCRIPTOME if assay_id.startswith("EFO:") else MolecularLayer.UNKNOWN + ) + + +class _FakeExtractor: + """Returns a preset GeoExtraction, recording the call for assertions.""" + + def __init__(self, extraction: GeoExtraction) -> None: + self._extraction = extraction + self.calls: list[tuple[str, str]] = [] + + def extract(self, instructions: str, content: str, response_model: type[SchemaT]) -> SchemaT: + self.calls.append((instructions, content)) + assert response_model is GeoExtraction + return self._extraction # type: ignore[return-value] + + +class _RaisingExtractor: + def extract(self, instructions: str, content: str, response_model: type[SchemaT]) -> SchemaT: + raise RuntimeError("LLM unavailable") + + +def _record() -> RawRecord: + return RawRecord( + source="GEO", + study_id="GSE99999", + title="Smoking and lung adenocarcinoma", + payload={ + "series": { + "type": ["Expression profiling by array"], + "summary": ["We profiled tumor and normal lung tissue."], + "overall_design": "2 tumor and 2 normal.", + }, + "samples": [ + { + "sample_id": "GSM000001", + "title": "Lung Tumor A", + "source_name": "Adenocarcinoma of the Lung", + "organism": "Homo sapiens", + "characteristics": ["gender: Male", "tissue: tumor"], + "supplementary_file": "ftp://host/GSM000001.CEL.gz", + }, + { + "sample_id": "GSM000002", + "title": "Lung Normal A", + "source_name": "Noninvolved Lung", + "organism": "Homo sapiens", + "characteristics": ["gender: Female", "tissue: normal"], + "supplementary_file": "ftp://host/GSM000002.CEL.gz", + }, + ], + "truncated": False, + }, + ) + + +def _extraction() -> GeoExtraction: + return GeoExtraction( + assay="microarray", + tissue="lung", + disease="lung adenocarcinoma", + samples=[ + SampleExtraction(sample_id="GSM000001", condition="tumor", subject="P1"), + SampleExtraction(sample_id="GSM000002", condition="normal", subject="P1"), + ], + ) + + +def _normalizer(extraction: GeoExtraction | None = None) -> GeoNormalizer: + return GeoNormalizer(_FakeExtractor(extraction or _extraction()), resolver=_FakeResolver()) + + +class TestGeoNormalizer: + def test_study_node(self): + kg = _normalizer().normalize(_record()) + assert len(kg.studies) == 1 + study = kg.studies[0] + assert study.study_id == "GSE99999" + assert study.source == "GEO" + assert study.assay == "EFO:0002772" + assert study.molecular_layer is MolecularLayer.TRANSCRIPTOME + + def test_no_dataset_node_for_geo(self): + """GEO has no distinct dataset artifact — the series is the study.""" + assert _normalizer().normalize(_record()).datasets == [] + + def test_one_sample_node_per_record_sample(self): + kg = _normalizer().normalize(_record()) + assert [s.sample_id for s in kg.samples] == ["GSM000001", "GSM000002"] + + def test_sample_covariates_from_llm_structured_from_record(self): + kg = _normalizer().normalize(_record()) + s1 = kg.samples[0] + # Design covariates come from the LLM extraction. + assert s1.condition == "tumor" + assert s1.subject == "P1" + # organism + data_uri are read deterministically from the record, not the LLM. + assert s1.organism == "Homo sapiens" + assert s1.data_uri == "ftp://host/GSM000001.CEL.gz" + + def test_has_sample_edges(self): + kg = _normalizer().normalize(_record()) + has_sample = [e for e in kg.edges if e.relation_type == "HAS_SAMPLE"] + assert {(e.source_id, e.target_id) for e in has_sample} == { + ("GSE99999", "GSM000001"), + ("GSE99999", "GSM000002"), + } + + def test_study_level_entities_grounded(self): + kg = _normalizer().normalize(_record()) + by_type = {e.entity_type: e for e in kg.biological_entities} + assert by_type[EntityType.ASSAY].ontology_id == "EFO:0002772" + assert by_type[EntityType.TISSUE].ontology_id == "UBERON:0002048" + assert by_type[EntityType.DISEASE].ontology_id == "MONDO:0005061" + assert by_type[EntityType.SPECIES].ontology_id == "NCBITaxon:9606" + + def test_design_context_edges_from_study(self): + kg = _normalizer().normalize(_record()) + rels = {(e.relation_type, e.target_id) for e in kg.edges if e.source_id == "GSE99999"} + assert ("MEASURED_WITH", "EFO:0002772") in rels + assert ("HAS_TISSUE", "UBERON:0002048") in rels + assert ("HAS_CONDITION", "MONDO:0005061") in rels + assert ("STUDIES", "NCBITaxon:9606") in rels + + def test_no_cell_type_entities(self): + """Even if covariates mention cells, no CL entity is ever produced.""" + kg = _normalizer().normalize(_record()) + assert all(not e.ontology_id.startswith("CL:") for e in kg.biological_entities) + + def test_unresolved_facets_skipped(self): + """Facets the resolver can't ground produce no entity; assay falls to unknown.""" + extraction = GeoExtraction(assay="hand-wavy assay", tissue=None, disease=None, samples=[]) + kg = _normalizer(extraction).normalize(_record()) + assert kg.studies[0].assay == "unknown" + assert kg.studies[0].molecular_layer is MolecularLayer.UNKNOWN + # Species still grounds from the record's structured organism field. + assert any(e.entity_type == EntityType.SPECIES for e in kg.biological_entities) + assert not any(e.entity_type == EntityType.ASSAY for e in kg.biological_entities) + + def test_extraction_failure_degrades_gracefully(self): + """An LLM error yields samples (from the record) but no extracted covariates.""" + normalizer = GeoNormalizer(_RaisingExtractor(), resolver=_FakeResolver()) + kg = normalizer.normalize(_record()) + assert len(kg.samples) == 2 + assert all(s.condition is None for s in kg.samples) + assert kg.studies[0].assay == "unknown" + # Organism is structured → species still resolves despite the LLM failure. + assert any(e.entity_type == EntityType.SPECIES for e in kg.biological_entities) + + def test_agent_receives_characteristics_in_prompt(self): + extractor = _FakeExtractor(_extraction()) + GeoNormalizer(extractor, resolver=_FakeResolver()).normalize(_record()) + _, content = extractor.calls[0] + assert "GSM000001" in content + assert "gender: Male" in content + + def test_roundtrip_json(self): + kg = _normalizer().normalize(_record()) + restored = KnowledgeGraphOutput.model_validate_json(kg.model_dump_json()) + assert restored == kg + + +class TestProtocolConformance: + def test_normalizer_satisfies_protocol(self): + assert isinstance(GeoNormalizer(_FakeExtractor(_extraction())), Normalizer) + + def test_fake_extractor_satisfies_structured_extractor(self): + assert isinstance(_FakeExtractor(_extraction()), StructuredExtractor) + + +# -- LLM ontology fallback factory (uses a fake extractor, no Azure) ---------- + + +class _FallbackExtractor: + """Returns a preset object for the fallback's private _FallbackTerm schema.""" + + def __init__(self, ontology_id: str | None, name: str | None) -> None: + self._ontology_id = ontology_id + self._name = name + + def extract(self, instructions: str, content: str, response_model: type[SchemaT]) -> SchemaT: + return response_model(ontology_id=self._ontology_id, name=self._name) # type: ignore[call-arg] + + +class TestOntologyFallbackFactory: + def test_accepts_correct_prefix(self): + from parce.agent.extraction import make_ontology_fallback + + fallback = make_ontology_fallback(_FallbackExtractor("UBERON:0002048", "lung")) + term = fallback("pulmonary tissue", Facet.TISSUE) + assert term == ResolvedTerm("UBERON:0002048", "lung") + + def test_rejects_wrong_prefix(self): + from parce.agent.extraction import make_ontology_fallback + + # A MONDO id returned for a TISSUE facet (expects UBERON) is dropped. + fallback = make_ontology_fallback(_FallbackExtractor("MONDO:0005061", "x")) + assert fallback("pulmonary tissue", Facet.TISSUE) is None + + def test_rejects_null(self): + from parce.agent.extraction import make_ontology_fallback + + fallback = make_ontology_fallback(_FallbackExtractor(None, None)) + assert fallback("???", Facet.TISSUE) is None + + def test_extractor_error_returns_none(self): + from parce.agent.extraction import make_ontology_fallback + + fallback = make_ontology_fallback(_RaisingExtractor()) + assert fallback("lung", Facet.TISSUE) is None + + +@pytest.mark.parametrize( + "field", + ["condition", "perturbation", "timepoint", "subject"], +) +def test_sample_extraction_design_only(field): + """The extraction schema exposes only design covariates (no cell-type field).""" + assert field in SampleExtraction.model_fields + assert "cell_type" not in SampleExtraction.model_fields + assert "cell_type" not in GeoExtraction.model_fields