diff --git a/.gitignore b/.gitignore index f308f225..cecf5c63 100644 --- a/.gitignore +++ b/.gitignore @@ -71,4 +71,10 @@ scarf/_version.py # Local working notes for countsT strip-shard experiments (tmp Modal hacks). /performance_reconfig.md .env -temp/ \ No newline at end of file +temp/ +# Dev testing scard reports +modal_scarf/ +notebook/ +PR/ +scarf/agent/*.md +.cache/ diff --git a/.readthedocs.yml b/.readthedocs.yml index bc5daf61..894d98cd 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -16,6 +16,7 @@ python: - method: uv command: sync extras: + - agent - docs - extra diff --git a/docs/.jupyter_cache/executed/30e76fa5af053e197ea52e41c1409264/base.ipynb b/docs/.jupyter_cache/executed/30e76fa5af053e197ea52e41c1409264/base.ipynb deleted file mode 100644 index f60d58d3..00000000 --- a/docs/.jupyter_cache/executed/30e76fa5af053e197ea52e41c1409264/base.ipynb +++ /dev/null @@ -1,532 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 1, - "id": "998beee3", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{'active_cells': 3940, 'total_cells': 5025, 'assays': ['RNA']}" - ] - }, - "execution_count": 1, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "import scarf\n", - "from scarf.agent import (\n", - " BiologicalContext,\n", - " BiologicalInterpretationAgent,\n", - " DataEnrichmentAgent,\n", - " DataEnrichmentContext,\n", - " ExperimentalContextAgent,\n", - " ParameterCandidate,\n", - " ParameterTuningAgent,\n", - ")\n", - "\n", - "scarf.configure_output(level=\"WARNING\", progress=False)\n", - "\n", - "dataset = scarf.cytebase.connect(\"scarf_docs\").download_dataset(\n", - " \"tenx_5K_pbmc_rnaseq\",\n", - " destination=\"scarf_datasets\",\n", - " zarr=True,\n", - ")\n", - "ds = scarf.DataStore(\n", - " f\"{dataset}/data.zarr\",\n", - " default_assay=\"RNA\",\n", - " nthreads=2,\n", - ")\n", - "\n", - "{\n", - " \"active_cells\": int(ds.cells.fetch_all(\"I\").sum()),\n", - " \"total_cells\": ds.cells.N,\n", - " \"assays\": ds.assay_names,\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "ea9415a9", - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "from pydantic_ai.messages import (\n", - " ModelMessage,\n", - " ModelResponse,\n", - " ToolCallPart,\n", - " ToolReturnPart,\n", - ")\n", - "from pydantic_ai.models.function import AgentInfo, FunctionModel\n", - "\n", - "from scarf.agent.biological_interpretation import (\n", - " ClusterCompositionEvidence,\n", - " ClusterMarkerBatchEvidence,\n", - ")\n", - "from scarf.agent.data_enrichment import AssayFeatureInspectionBatch\n", - "from scarf.agent.experimental_context import CovariateEvidence\n", - "\n", - "\n", - "def _tool_returns(messages: list[ModelMessage]) -> list[ToolReturnPart]:\n", - " return [\n", - " part\n", - " for message in messages\n", - " for part in message.parts\n", - " if isinstance(part, ToolReturnPart)\n", - " ]\n", - "\n", - "\n", - "def _tool_call(name: str, args: dict | None = None) -> ModelResponse:\n", - " return ModelResponse(parts=[ToolCallPart(tool_name=name, args=args or {})])\n", - "\n", - "\n", - "def _structured_output(info: AgentInfo, payload: dict) -> ModelResponse:\n", - " return _tool_call(info.output_tools[0].name, payload)\n", - "\n", - "\n", - "async def _enrichment_reply(\n", - " messages: list[ModelMessage],\n", - " info: AgentInfo,\n", - ") -> ModelResponse:\n", - " returns = _tool_returns(messages)\n", - " if not returns:\n", - " return _tool_call(\"inspect_assay_features_batch\")\n", - "\n", - " batch = AssayFeatureInspectionBatch.model_validate(returns[-1].content)\n", - " inspection = batch.inspections[0]\n", - " species_observed = inspection.species != \"unknown\"\n", - " species = inspection.species if species_observed else \"homo_sapiens\"\n", - " evidence_ids = list(inspection.evidenceIds)\n", - " if not species_observed:\n", - " evidence_ids.append(\"context:organism\")\n", - " policy = {\n", - " \"assay\": inspection.assay,\n", - " \"species\": species,\n", - " \"speciesConfidence\": \"high\" if species_observed else \"medium\",\n", - " \"speciesRationale\": (\n", - " inspection.speciesReason\n", - " or \"The inspected features and caller context support this species.\"\n", - " ),\n", - " \"excludeFamilies\": [\n", - " family.family\n", - " for family in inspection.families\n", - " if family.count > 0 and family.defaultExclude is True\n", - " ],\n", - " \"protectFamilies\": [\n", - " family.family\n", - " for family in inspection.families\n", - " if family.count > 0 and family.defaultExclude is False\n", - " ],\n", - " \"rationale\": \"Exclude observed technical families and preserve protected ones.\",\n", - " \"evidenceIds\": evidence_ids,\n", - " }\n", - " return _structured_output(info, {\"status\": \"done\", \"policies\": [policy]})\n", - "\n", - "\n", - "async def _experimental_reply(\n", - " messages: list[ModelMessage],\n", - " info: AgentInfo,\n", - ") -> ModelResponse:\n", - " returns = _tool_returns(messages)\n", - " if not returns:\n", - " return _tool_call(\"inspect_cell_covariates\")\n", - " if len(returns) == 1:\n", - " return _tool_call(\n", - " \"analyze_experimental_design\",\n", - " {\n", - " \"column_domains\": {},\n", - " \"coefficients_of_interest\": [],\n", - " \"units_of_inference\": {},\n", - " \"batch_columns\": [],\n", - " },\n", - " )\n", - "\n", - " design = CovariateEvidence.model_validate(returns[-1].content)\n", - " evidence_id = design.evidenceIds[0]\n", - " return _structured_output(\n", - " info,\n", - " {\n", - " \"batchCorrection\": {\n", - " \"action\": \"skip\",\n", - " \"rationale\": (\n", - " \"No explicit technical batch or biological contrast was supplied.\"\n", - " ),\n", - " \"evidenceIds\": [evidence_id],\n", - " },\n", - " \"rationale\": \"Continue with an uncorrected baseline.\",\n", - " \"evidenceIds\": [evidence_id],\n", - " },\n", - " )\n", - "\n", - "\n", - "async def _tuning_reply(\n", - " _messages: list[ModelMessage],\n", - " info: AgentInfo,\n", - ") -> ModelResponse:\n", - " return _structured_output(\n", - " info,\n", - " {\n", - " \"status\": \"done\",\n", - " \"recommendedCandidateId\": \"baseline\",\n", - " \"confidence\": \"medium\",\n", - " \"rationale\": \"The single authorized baseline completed successfully.\",\n", - " \"evidenceIds\": [\"candidate:baseline:clusters\"],\n", - " \"stopReason\": \"The authorized candidate was evaluated.\",\n", - " },\n", - " )\n", - "\n", - "\n", - "async def _biology_reply(\n", - " messages: list[ModelMessage],\n", - " info: AgentInfo,\n", - ") -> ModelResponse:\n", - " returns = _tool_returns(messages)\n", - " if not returns:\n", - " return _tool_call(\"inspect_cluster_composition\")\n", - " if len(returns) == 1:\n", - " composition = ClusterCompositionEvidence.model_validate(returns[-1].content)\n", - " cluster_id = sorted(\n", - " composition.clusterCounts,\n", - " key=lambda value: (-composition.clusterCounts[value], value),\n", - " )[0]\n", - " return _tool_call(\n", - " \"inspect_cluster_markers_batch\",\n", - " {\"cluster_ids\": [cluster_id]},\n", - " )\n", - "\n", - " marker_batch = ClusterMarkerBatchEvidence.model_validate(returns[-1].content)\n", - " marker = marker_batch.clusters[0]\n", - " if not marker.evidenceId:\n", - " return _structured_output(\n", - " info,\n", - " {\n", - " \"status\": \"needsInput\",\n", - " \"needsInput\": {\n", - " \"question\": \"No markers passed the bounded search thresholds.\",\n", - " \"requiredInputs\": [\"markerArtifact\"],\n", - " },\n", - " \"limitations\": marker.warnings,\n", - " \"stopReason\": \"Marker evidence was unavailable.\",\n", - " },\n", - " )\n", - " names = [item.featureName or item.featureId for item in marker.markers[:3]]\n", - " return _structured_output(\n", - " info,\n", - " {\n", - " \"status\": \"done\",\n", - " \"clusterInterpretations\": [\n", - " {\n", - " \"clusterId\": marker.clusterId,\n", - " \"proposedIdentity\": \"unresolved marker-defined cluster\",\n", - " \"identityIsHypothesis\": True,\n", - " \"confidence\": \"low\",\n", - " \"rationale\": f\"Top returned marker features: {', '.join(names)}.\",\n", - " \"evidenceIds\": [marker.evidenceId],\n", - " }\n", - " ],\n", - " \"evidenceIds\": [marker.evidenceId],\n", - " \"limitations\": [\n", - " \"The scripted documentation model does not assign cell identities.\"\n", - " ],\n", - " \"stopReason\": \"One bounded cluster was reviewed.\",\n", - " },\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "0370ebfc", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{'status': 'done',\n", - " 'species': 'homo_sapiens',\n", - " 'exclude_families': ['ribosomal', 'histone'],\n", - " 'protect_families': ['cellCycle'],\n", - " 'tool_calls': ['inspect_assay_features_batch']}" - ] - }, - "execution_count": 3, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "enrichment = DataEnrichmentAgent(FunctionModel(_enrichment_reply)).run(\n", - " ds,\n", - " context=DataEnrichmentContext(\n", - " studyContext=(\n", - " \"10x 5K PBMC RNA-seq from peripheral blood of a healthy human donor.\"\n", - " ),\n", - " organismHint=\"human\",\n", - " tissueReferences=[\"peripheral blood\"],\n", - " cellTypeReferences=[\"T cell\", \"B cell\", \"NK cell\", \"monocyte\"],\n", - " experimentalDetails=[\"10x 3 prime RNA-seq\", \"single donor\"],\n", - " ),\n", - " assays=[\"RNA\"],\n", - ")\n", - "\n", - "policy = enrichment.policies[0]\n", - "{\n", - " \"status\": enrichment.status,\n", - " \"species\": policy.species,\n", - " \"exclude_families\": policy.excludeFamilies,\n", - " \"protect_families\": policy.protectFamilies,\n", - " \"tool_calls\": [call.name for call in enrichment.toolCalls],\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "25e8b1bf", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{'run_id': 'a502c6ba45de187eb8aa73bca96bbd09fa2c5722976aaf534d2dce9207ffafab',\n", - " 'active_cells': 3940,\n", - " 'feature_selection': 'adeede36cdca822fde8cf2e62bb430843f28056a4e34414f753607c97a69240d',\n", - " 'normalized': '24cfc31e5c045f8066a38de7ce90d4ca1c6fe7d9c140e6b6686c77e731fd75d0'}" - ] - }, - "execution_count": 4, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "run = ds.pipeline.open(label=\"docs_default\")\n", - "normalized = run[\"normalized\"]\n", - "hvg_ref = run[\"highly_variable_features\"]\n", - "\n", - "{\n", - " \"run_id\": run.run_id,\n", - " \"active_cells\": int(run.cells.fetch_all(\"I\").sum()),\n", - " \"feature_selection\": hvg_ref.artifact_id,\n", - " \"normalized\": normalized.artifact_id,\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "80c798f4", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{'status': 'done',\n", - " 'batch_action': 'skip',\n", - " 'batch_columns': [],\n", - " 'coefficients': []}" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "experimental = ExperimentalContextAgent(FunctionModel(_experimental_reply)).run(\n", - " ds,\n", - " study_context=(\n", - " \"Healthy-donor 5K PBMC. No treatment or batch labels are available. \"\n", - " \"Do not invent a technical batch or biological contrast.\"\n", - " ),\n", - " run=run,\n", - ")\n", - "\n", - "{\n", - " \"status\": experimental.status,\n", - " \"batch_action\": experimental.decision.batchCorrection.action,\n", - " \"batch_columns\": experimental.decision.batchCorrection.batchColumns,\n", - " \"coefficients\": experimental.decision.coefficientsOfInterest,\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "50ca3712", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{'status': 'done',\n", - " 'recommended_candidate': 'baseline',\n", - " 'eligible': True,\n", - " 'clusters': 12,\n", - " 'smallest_cluster': 12,\n", - " 'cluster_artifact': 'c9ec9a51d0f0ea779459e20c1d030262aa50c6d766ad0e9b7dd02c09b423a195',\n", - " 'cell_selection': 'a4137c5ada6d3933e211d551dd636a097d6f0ed8103afde75fb3b2b05d67c16b'}" - ] - }, - "execution_count": 6, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "if experimental.status != \"done\":\n", - " raise RuntimeError(f\"Experimental Context stopped with {experimental.status!r}\")\n", - "\n", - "tuning_handoff = experimental.to_parameter_tuning_handoff()\n", - "candidate = ParameterCandidate(\n", - " candidateId=\"baseline\",\n", - " dimensions=15,\n", - " leidenResolution=0.5,\n", - " neighborsK=11,\n", - " useHarmony=False,\n", - ")\n", - "tuning = ParameterTuningAgent(FunctionModel(_tuning_reply)).run(\n", - " ds,\n", - " normalized=normalized,\n", - " candidates=[candidate],\n", - " experimental_handoff=tuning_handoff,\n", - " max_candidates=1,\n", - " max_refined_candidates=0,\n", - " min_cluster_cells=10,\n", - ")\n", - "\n", - "evaluation = tuning.evaluations[0]\n", - "{\n", - " \"status\": tuning.status,\n", - " \"recommended_candidate\": tuning.recommendedCandidateId,\n", - " \"eligible\": evaluation.eligible,\n", - " \"clusters\": evaluation.metrics.nClusters,\n", - " \"smallest_cluster\": evaluation.metrics.minClusterCells,\n", - " \"cluster_artifact\": evaluation.artifacts[\"clusters\"].artifactId,\n", - " \"cell_selection\": evaluation.cellSelection.artifactId,\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "799168ba", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{'status': 'done',\n", - " 'interpretations': [{'cluster': '4',\n", - " 'identity': 'unresolved marker-defined cluster',\n", - " 'rationale': 'Top returned marker features: SOCS3, TCF7, TRAC.',\n", - " 'evidence': ['markers:4482f78c6a16d49ce8d19ec1de75cc13f76d43cce8b82d06509240d52937a762:clusters:c9ec9a51d0f0ea779459e20c1d030262aa50c6d766ad0e9b7dd02c09b423a195:cluster:4']}],\n", - " 'tool_calls': ['inspect_cluster_composition',\n", - " 'inspect_cluster_markers_batch'],\n", - " 'treatment_observations': 0,\n", - " 'limitations': ['The scripted documentation model does not assign cell identities.']}" - ] - }, - "execution_count": 7, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "if tuning.status != \"done\":\n", - " raise RuntimeError(f\"Parameter Tuning stopped with {tuning.status!r}\")\n", - "\n", - "biology_handoff = tuning.to_biological_handoff()\n", - "biology = BiologicalInterpretationAgent(FunctionModel(_biology_reply)).run(\n", - " ds,\n", - " tuning_handoff=biology_handoff,\n", - " biological_context=BiologicalContext(\n", - " organism=\"Homo sapiens\",\n", - " tissue=\"peripheral blood\",\n", - " cellTypeReferences=[\"T cell\", \"B cell\", \"NK cell\", \"monocyte\"],\n", - " experimentalDetails=[\"healthy donor PBMC\", \"no treatment contrast\"],\n", - " ),\n", - " allow_marker_search=True,\n", - " marker_features=hvg_ref,\n", - " max_clusters=1,\n", - " max_markers=5,\n", - " marker_min_score=0.01,\n", - " marker_min_fraction=0.0,\n", - ")\n", - "\n", - "{\n", - " \"status\": biology.status,\n", - " \"interpretations\": [\n", - " {\n", - " \"cluster\": item.clusterId,\n", - " \"identity\": item.proposedIdentity,\n", - " \"rationale\": item.rationale,\n", - " \"evidence\": item.evidenceIds,\n", - " }\n", - " for item in biology.clusterInterpretations\n", - " ],\n", - " \"tool_calls\": [call.toolName for call in biology.runInfo.toolCalls],\n", - " \"treatment_observations\": len(biology.treatmentObservations),\n", - " \"limitations\": biology.limitations,\n", - "}" - ] - } - ], - "metadata": { - "description": "Run Scarf's four grounded analysis agents from a rebuilt 5K PBMC baseline.", - "jupytext": { - "cell_metadata_filter": "tags", - "text_representation": { - "extension": ".md", - "format_name": "myst", - "format_version": 0.13, - "jupytext_version": "1.14.1" - } - }, - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.13" - }, - "source_map": [ - 14, - 37, - 67, - 75, - 259, - 267, - 290, - 302, - 313, - 322, - 338, - 350, - 382, - 390, - 427 - ] - }, - "nbformat": 4, - "nbformat_minor": 5 -} \ No newline at end of file diff --git a/docs/.jupyter_cache/executed/a9d25b8976d878b82c36688d01bd384b/base.ipynb b/docs/.jupyter_cache/executed/a9d25b8976d878b82c36688d01bd384b/base.ipynb new file mode 100644 index 00000000..6169db3b --- /dev/null +++ b/docs/.jupyter_cache/executed/a9d25b8976d878b82c36688d01bd384b/base.ipynb @@ -0,0 +1,976 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "26f9edf6", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
Downloading bucket files18098007 / 18098007 complete18098007 / 18098007 complete
" + ], + "text/plain": [ + "Downloading bucket files: 18098007 / 18098007 complete" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
Downloading bytes18098007 / 18098007 complete18098007 / 18098007 complete
" + ], + "text/plain": [ + "Downloading bytes: 18098007 / 18098007 complete" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/plain": [ + "'data.h5'" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from pathlib import Path\n", + "from tempfile import TemporaryDirectory\n", + "\n", + "import pandas as pd\n", + "import scarf\n", + "from scarf.agent import analyze_rna\n", + "\n", + "scarf.configure_output(level=\"WARNING\", progress=False)\n", + "source_path = scarf.cytebase.connect(\"scarf_docs\").download(\n", + " \"tenx_5K_pbmc_rnaseq/data.h5\", destination=\"scarf_datasets\",\n", + ")[0]\n", + "teaching_directory = TemporaryDirectory(prefix=\"scarf-agent-teaching-\")\n", + "zarr_path = Path(teaching_directory.name) / \"analysis.zarr\"\n", + "study_context = (\n", + " \"Human 10x Genomics 5K PBMC 3-prime gene expression from peripheral blood, \"\n", + " \"collected from one healthy donor. The teaching cohort is a deterministic random \"\n", + " \"subset of 1,000 cells (seed 42), not the full public dataset. \"\n", + " \"No treatment comparison, trusted technical \"\n", + " \"batch column, paired modality, or independent replication metadata is available. \"\n", + " \"Do not invent missing design variables or report treatment effects.\"\n", + ")\n", + "source_path.name" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "75fcef71", + "metadata": { + "tags": [ + "remove-cell" + ] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "WARNING: Minimum cell count (502) is lower than size factor multiplier (1000)\n" + ] + } + ], + "source": [ + "import json\n", + "from typing import Any\n", + "\n", + "import numpy as np\n", + "from IPython import get_ipython\n", + "from pydantic_ai.messages import (\n", + " ModelMessage,\n", + " ModelResponse,\n", + " ToolCallPart,\n", + " ToolReturnPart,\n", + ")\n", + "from pydantic_ai.models.function import AgentInfo, FunctionModel\n", + "\n", + "from scarf.agent.data_enrichment import (\n", + " AssayFeatureInspectionBatch,\n", + " DataEnrichmentReport,\n", + " FeatureSelectionPolicy,\n", + " StudyContextSummary,\n", + ")\n", + "from scarf.agent.experimental_context import (\n", + " BatchCorrectionPlan,\n", + " ExperimentalContextDecision,\n", + ")\n", + "from scarf.agent.ingest import ingest\n", + "\n", + "prepared_input = ingest(path=source_path, zarrPath=zarr_path)\n", + "if prepared_input.status != \"done\":\n", + " raise RuntimeError(f\"Teaching dataset import failed: {prepared_input.notes}\")\n", + "teaching_store = scarf.DataStore(\n", + " str(zarr_path), min_features_per_cell=-1, mito_pattern=\"\", ribo_pattern=\"\",\n", + ")\n", + "teaching_store.cells.reset_key(\"I\")\n", + "teaching_cells = np.zeros(teaching_store.cells.N, dtype=bool)\n", + "teaching_cells[np.random.default_rng(42).choice(teaching_store.cells.N, 1000, replace=False)] = True\n", + "teaching_store.cells.update_key(teaching_cells, \"I\")\n", + "source_path = zarr_path\n", + "del teaching_store\n", + "\n", + "notebook_shell = get_ipython()\n", + "if notebook_shell is not None:\n", + " notebook_shell.run_line_magic(\"matplotlib\", \"inline\")\n", + "\n", + "def _prompt_text(messages: list[ModelMessage]) -> str:\n", + " values = []\n", + " for message in messages:\n", + " for part in message.parts:\n", + " content = getattr(part, \"content\", None)\n", + " if isinstance(content, str):\n", + " values.append(content)\n", + " elif isinstance(content, tuple):\n", + " values.extend(item for item in content if isinstance(item, str))\n", + " return \"\\n\".join(values)\n", + "\n", + "\n", + "def _tool_result(\n", + " messages: list[ModelMessage],\n", + " tool_name: str,\n", + " model_type: Any,\n", + ") -> Any:\n", + " for message in reversed(messages):\n", + " for part in reversed(message.parts):\n", + " if isinstance(part, ToolReturnPart) and part.tool_name == tool_name:\n", + " if isinstance(part.content, model_type):\n", + " return part.content\n", + " if model_type is dict:\n", + " return json.loads(part.content) if isinstance(part.content, str) else part.content\n", + " if isinstance(part.content, str):\n", + " return model_type.model_validate_json(part.content)\n", + " return model_type.model_validate(part.content)\n", + " raise AssertionError(f\"Missing tool return {tool_name!r}\")\n", + "\n", + "\n", + "def _tool_call(name: str, args: dict[str, Any] | None = None) -> ModelResponse:\n", + " return ModelResponse(parts=[ToolCallPart(tool_name=name, args=args or {})])\n", + "\n", + "\n", + "def _structured_output(info: AgentInfo, value: Any) -> ModelResponse:\n", + " payload = value.model_dump() if hasattr(value, \"model_dump\") else value\n", + " return _tool_call(info.output_tools[0].name, payload)\n", + "\n", + "\n", + "def _scripted_workflow_model() -> tuple[FunctionModel, dict[str, Any]]:\n", + " state = {\n", + " \"enrichment\": 0,\n", + " \"context\": 0,\n", + " \"parameter\": 0,\n", + " \"assessments\": [],\n", + " \"requests\": 0,\n", + " }\n", + "\n", + " async def reply(\n", + " messages: list[ModelMessage],\n", + " info: AgentInfo,\n", + " ) -> ModelResponse:\n", + " state[\"requests\"] += 1\n", + " tools = {tool.name for tool in info.function_tools}\n", + "\n", + " if (\n", + " \"inspect_assay_features_batch\" in tools\n", + " or state[\"enrichment\"] == 1\n", + " or any(\n", + " tool.parameters_json_schema.get(\"title\") == \"DataEnrichmentReport\"\n", + " for tool in info.output_tools\n", + " )\n", + " ):\n", + " if state[\"enrichment\"] == 0:\n", + " state[\"enrichment\"] = 1\n", + " return _tool_call(\"inspect_assay_features_batch\")\n", + "\n", + " batch = _tool_result(\n", + " messages,\n", + " \"inspect_assay_features_batch\",\n", + " AssayFeatureInspectionBatch,\n", + " )\n", + " policies = []\n", + " for inspection in batch.inspections:\n", + " species_observed = inspection.species != \"unknown\"\n", + " policy_evidence = list(inspection.evidenceIds)\n", + " if not species_observed:\n", + " policy_evidence.append(\"context:study\")\n", + " policies.append(\n", + " FeatureSelectionPolicy(\n", + " assay=inspection.assay,\n", + " species=(\n", + " inspection.species\n", + " if species_observed\n", + " else \"homo_sapiens\"\n", + " ),\n", + " speciesConfidence=\"high\" if species_observed else \"medium\",\n", + " speciesRationale=(\n", + " inspection.speciesReason\n", + " or \"The exact study paragraph identifies a human sample.\"\n", + " ),\n", + " excludeFamilies=[\n", + " family.family\n", + " for family in inspection.families\n", + " if family.count > 0 and family.defaultExclude is True\n", + " ],\n", + " protectFamilies=[\n", + " family.family\n", + " for family in inspection.families\n", + " if family.count > 0 and family.defaultExclude is False\n", + " ],\n", + " rationale=(\n", + " \"Exclude observed technical families and preserve \"\n", + " \"observed protected families.\"\n", + " ),\n", + " evidenceIds=list(dict.fromkeys(policy_evidence)),\n", + " )\n", + " )\n", + " state[\"enrichment\"] = 2\n", + " return _structured_output(\n", + " info,\n", + " DataEnrichmentReport(\n", + " status=\"done\",\n", + " studyContextSummary=StudyContextSummary(\n", + " organismReferences=[\"Human\"],\n", + " tissueReferences=[\"peripheral blood\"],\n", + " experimentalReferences=[\n", + " \"10x Genomics 5K PBMC 3-prime gene expression\"\n", + " ],\n", + " analysisIntentReferences=[\n", + " \"Discover stable major immune-cell populations.\"\n", + " ],\n", + " ),\n", + " policies=policies,\n", + " ),\n", + " )\n", + "\n", + " if tools.intersection(\n", + " {\n", + " \"inspect_cell_covariates\",\n", + " \"analyze_experimental_design\",\n", + " \"score_current_representation\",\n", + " }\n", + " ) or state[\"context\"] in {1, 2} or any(\n", + " tool.parameters_json_schema.get(\"title\") == \"ExperimentalContextDecision\"\n", + " for tool in info.output_tools\n", + " ):\n", + " if state[\"context\"] == 0:\n", + " state[\"context\"] = 1\n", + " return _tool_call(\"inspect_cell_covariates\")\n", + " if state[\"context\"] == 1:\n", + " state[\"context\"] = 2\n", + " return _tool_call(\n", + " \"analyze_experimental_design\",\n", + " {\n", + " \"column_domains\": {},\n", + " \"coefficients_of_interest\": [],\n", + " \"units_of_inference\": {},\n", + " \"batch_columns\": [],\n", + " },\n", + " )\n", + "\n", + " design = _tool_result(\n", + " messages,\n", + " \"analyze_experimental_design\",\n", + " dict,\n", + " )\n", + " profile = next(\n", + " value\n", + " for value in design[\"qcProfiles\"]\n", + " if value[\"action\"] == \"skip\"\n", + " )\n", + " evidence_id = profile[\"evidenceId\"]\n", + " state[\"context\"] = 3\n", + " return _structured_output(\n", + " info,\n", + " ExperimentalContextDecision(\n", + " batchCorrection=BatchCorrectionPlan(\n", + " action=\"skip\",\n", + " rationale=\"No trusted technical batch column was supplied.\",\n", + " evidenceIds=[evidence_id],\n", + " ),\n", + " rationale=\"No experimental covariates were supplied.\",\n", + " evidenceIds=[evidence_id],\n", + " ),\n", + " )\n", + "\n", + " prompt = _prompt_text(messages)\n", + " if any(\n", + " {\"selectedCandidateId\", \"comparisonConclusions\"}.issubset(\n", + " tool.parameters_json_schema.get(\"properties\", {})\n", + " )\n", + " for tool in info.output_tools\n", + " ):\n", + " evidence, _ = json.JSONDecoder().raw_decode(prompt[prompt.index(\"{\") :])\n", + " coverage = evidence[\"comparisonCoverage\"]\n", + " settings = coverage[\"candidateSettings\"]\n", + " candidates = {\n", + " item[\"candidateId\"]: item for item in evidence[\"candidates\"]\n", + " if item[\"status\"] == \"done\" and item[\"eligible\"]\n", + " }\n", + " if not candidates:\n", + " raise AssertionError(\"The teaching run has no supported partition\")\n", + "\n", + " def rank(identity):\n", + " metrics = settings[identity][\"metrics\"]\n", + " return tuple(\n", + " float(metrics[name]) if metrics.get(name) is not None else -1.0\n", + " for name in (\"seedStability\", \"markerCoherence\")\n", + " )\n", + "\n", + " axis_names = {\n", + " \"hvgCount\": \"variable-gene count\", \"hvgRanking\": \"gene ranking\",\n", + " \"featurePolicy\": \"gene-family policy\", \"pca\": \"PCA dimensions\",\n", + " \"neighbors\": \"neighbor count\", \"partition\": \"clustering resolution\",\n", + " }\n", + " metric_names = {\n", + " \"seedStability\": \"agreement across clustering runs\",\n", + " \"subsampleStability\": \"agreement after resampling\",\n", + " \"markerCoherence\": \"the fraction of clusters with qualifying markers\",\n", + " \"markerSpecificityMedian\": \"median marker specificity\",\n", + " \"macroF1\": \"classification agreement\",\n", + " }\n", + " axis_ids = {}\n", + " for row in coverage[\"comparisons\"]:\n", + " identities = axis_ids.setdefault(row[\"axis\"], [])\n", + " for identity in (row[\"baselineCandidateId\"], row[\"alternativeCandidateId\"]):\n", + " if identity is not None and identity not in identities:\n", + " identities.append(identity)\n", + " comparison_ids = {axis: list(identities) for axis, identities in axis_ids.items()}\n", + " if coverage[\"phase\"] != \"sensitivity\":\n", + " axis_ids[\"partition\"] = list(dict.fromkeys(\n", + " [*axis_ids[\"partition\"], *coverage[\"resolutionCandidateIds\"]]\n", + " ))\n", + " preferences = {axis: max(identities, key=rank) for axis, identities in axis_ids.items()}\n", + " pending_policy = any(row[\"status\"] == \"pending\" for row in coverage[\"comparisons\"])\n", + " experiment_id = None\n", + " if coverage[\"phase\"] == \"sensitivity\":\n", + " selected_id = evidence[\"currentCandidateId\"]\n", + " action_name = \"combine\"\n", + " else:\n", + " eligible_resolutions = [\n", + " identity for identity in coverage[\"resolutionCandidateIds\"]\n", + " if identity in candidates\n", + " ]\n", + " if not eligible_resolutions:\n", + " raise AssertionError(\"The combined representation has no supported partition\")\n", + " selected_id = max(eligible_resolutions, key=rank)\n", + " preferences[\"partition\"] = selected_id\n", + " action_name = \"accept\"\n", + " if pending_policy:\n", + " families = evidence[\"featureEvidence\"][selected_id][\"families\"]\n", + " supported = [\n", + " (key, option) for key, option in evidence[\"experiments\"].items()\n", + " if option[\"parameter\"] in {\"includeFamily\", \"excludeFamily\"}\n", + " and families.get(option[\"value\"], {}).get(\n", + " \"selectedExamples\" if option[\"parameter\"] == \"excludeFamily\"\n", + " else \"excludedExamples\"\n", + " )\n", + " ]\n", + " if not supported:\n", + " raise AssertionError(\"The teaching policy has no observed family program to nominate\")\n", + " experiment_id, option = max(supported, key=lambda item: (\n", + " families[item[1][\"value\"]].get(\"selectedGenes\", 0),\n", + " item[1][\"affectedEligibleGenes\"],\n", + " ))\n", + " action_name = \"experiment\"\n", + " selected = candidates[selected_id]\n", + " metrics = selected[\"metrics\"]\n", + " genes = list(dict.fromkeys(\n", + " gene for names in metrics.get(\"topMarkerGenes\", {}).values()\n", + " for gene in names\n", + " ))[:8]\n", + " qualitative = (\n", + " \"The saved marker preview contains \" + \", \".join(genes) + \".\"\n", + " if genes else \"The saved marker preview is empty; cell identities remain unresolved.\"\n", + " )\n", + " conclusions = []\n", + " for axis, identities in axis_ids.items():\n", + " preferred = preferences[axis]\n", + " score = settings[preferred][\"metrics\"]\n", + " conclusions.append({\n", + " \"axis\": axis,\n", + " \"candidateIds\": identities,\n", + " \"preferredCandidateId\": preferred,\n", + " \"quantitativeReason\": \"; \".join(\n", + " f\"Observed alternative {index + 1}: repeat agreement \"\n", + " f\"{settings[identity]['metrics'].get('seedStability')}, \"\n", + " f\"marker coverage {settings[identity]['metrics'].get('markerCoherence')}\"\n", + " for index, identity in enumerate(identities)\n", + " ),\n", + " \"biologicalReason\": (\n", + " \"This scripted teaching policy does not establish cell identities or \"\n", + " \"infer that a gene program is a technical artifact. \" + qualitative\n", + " ),\n", + " \"plainLanguageSummary\": (\n", + " f\"The teaching policy compared {len(identities)} observed {axis_names[axis]} \"\n", + " f\"settings and preferred repeat agreement {score.get('seedStability')}, \"\n", + " f\"using marker coverage {score.get('markerCoherence')} to break ties.\"\n", + " ),\n", + " \"tradeoffs\": [{\n", + " \"alternativeCandidateId\": identity,\n", + " \"metric\": metric,\n", + " \"preferredValue\": score[metric],\n", + " \"alternativeValue\": settings[identity][\"metrics\"][metric],\n", + " \"interpretation\": (\n", + " f\"An alternative has higher {metric_names[metric]} \"\n", + " f\"({settings[identity]['metrics'][metric]} versus {score[metric]}). \"\n", + " \"The teaching policy prioritizes repeat agreement, then marker coverage; \"\n", + " \"this loss remains an explicit limit of its choice.\"\n", + " ),\n", + " } for identity in comparison_ids[axis] if identity != preferred\n", + " for metric in (\"seedStability\", \"subsampleStability\", \"markerCoherence\",\n", + " \"markerSpecificityMedian\", \"macroF1\")\n", + " if isinstance(score.get(metric), (int, float))\n", + " and isinstance(settings[identity][\"metrics\"].get(metric), (int, float))\n", + " and settings[identity][\"metrics\"][metric] > score[metric]],\n", + " })\n", + " quantitative = (\n", + " f\"Observed resolution {selected['parameters']['leidenResolution']} has \"\n", + " f\"repeat agreement {metrics.get('seedStability')} and marker coverage \"\n", + " f\"{metrics.get('markerCoherence')}.\"\n", + " )\n", + " summary = (\n", + " \"The teaching policy proposes testing a represented gene family; \"\n", + " \"its contribution must be measured before retaining or changing the gene selection.\"\n", + " if action_name == \"experiment\" else\n", + " \"The teaching policy proposes a combination of the observed settings; \"\n", + " \"Scarf must execute that combination and compare its four resolutions.\"\n", + " if action_name == \"combine\" else\n", + " \"The teaching policy selected the measured combined representation and \"\n", + " \"its most repeatable eligible partition. Marker identities remain unvalidated.\"\n", + " )\n", + " action = {\n", + " \"action\": action_name,\n", + " \"selectedCandidateId\": selected_id,\n", + " \"experimentId\": experiment_id,\n", + " \"correctionNeed\": \"notApplicable\",\n", + " \"comparisonConclusions\": conclusions,\n", + " \"plainLanguageSummary\": summary,\n", + " \"evidenceIds\": [\n", + " f\"candidate:{selected_id}\",\n", + " *list(evidence[\"imageHashes\"])[:1],\n", + " \"studyContract\", \"qcPolicy\", \"samplingCoverage\", \"featureEvidence\",\n", + " ],\n", + " \"quantitativeFindings\": [quantitative],\n", + " \"qualitativeFindings\": [qualitative],\n", + " \"objectivePreservation\": (\n", + " \"Preserve the single-donor population structure and retain marker \"\n", + " \"uncertainty; no batch or treatment comparison is supported.\"\n", + " ),\n", + " \"rationale\": summary + \" \" + quantitative,\n", + " \"concern\": (\n", + " f\"Observed family {option['value']} contains \"\n", + " f\"{families[option['value']].get('selectedGenes', 0)} selected genes. \"\n", + " \"Test sensitivity to this program without assuming that it is technical.\"\n", + " ) if pending_policy else \"\",\n", + " \"expectedImprovement\": (\n", + " \"Measure whether changing this gene-family selection preserves the major \"\n", + " \"marker programs and improves repeat agreement.\"\n", + " ) if pending_policy else \"\",\n", + " \"populationConcerns\": [{\n", + " \"candidateId\": selected_id,\n", + " \"clusterId\": cluster,\n", + " \"status\": \"nonEssentialLimitation\",\n", + " \"evidenceIds\": [f\"candidate:{selected_id}\"],\n", + " \"explanation\": (\n", + " f\"Population {cluster} has no qualifying marker genes and remains unclassified. \"\n", + " \"This tutorial demonstrates selecting analysis settings; it does not validate \"\n", + " \"cell identities or claim that every population has been biologically resolved.\"\n", + " ),\n", + " } for cluster, names in metrics.get(\"topMarkerGenes\", {}).items() if not names],\n", + " }\n", + " if action_name == \"combine\":\n", + " action[\"combinedSettings\"] = {\n", + " field: preferences[axis] for field, axis in (\n", + " (\"hvgCountCandidateId\", \"hvgCount\"),\n", + " (\"hvgRankingCandidateId\", \"hvgRanking\"),\n", + " (\"featurePolicyCandidateId\", \"featurePolicy\"),\n", + " (\"pcaCandidateId\", \"pca\"),\n", + " (\"neighborsCandidateId\", \"neighbors\"),\n", + " )\n", + " }\n", + " state[\"assessments\"].append({\n", + " \"selection\": action,\n", + " \"alternatives\": [{\n", + " \"resolution\": settings[identity][\"parameters\"][\"leidenResolution\"],\n", + " \"clusters\": settings[identity][\"metrics\"].get(\"nClusters\"),\n", + " \"repeat_agreement\": settings[identity][\"metrics\"].get(\"seedStability\"),\n", + " \"clusters_with_markers\": settings[identity][\"metrics\"].get(\"markerCoherence\"),\n", + " \"selected\": identity == selected_id,\n", + " } for identity in (\n", + " coverage[\"resolutionCandidateIds\"] if action_name == \"accept\"\n", + " else list(candidates)\n", + " )],\n", + " })\n", + " return _structured_output(info, action)\n", + "\n", + "\n", + " payload, _ = json.JSONDecoder().raw_decode(prompt[prompt.index(\"{\") :])\n", + " decision = payload[\"spec\"]\n", + " evidence_by_class = {}\n", + " evidence_class_by_id = {}\n", + " for item in payload[\"evidence\"][\"evidence\"]:\n", + " evidence_by_class.setdefault(\n", + " item[\"evidenceClass\"],\n", + " item[\"evidenceId\"],\n", + " )\n", + " evidence_class_by_id[item[\"evidenceId\"]] = item[\"evidenceClass\"]\n", + " preferred = decision.get(\"metricPreferredOptionId\")\n", + " selected = (\n", + " next(\n", + " option\n", + " for option in decision[\"options\"]\n", + " if option[\"optionId\"] == preferred\n", + " )\n", + " if preferred is not None\n", + " else next(\n", + " option\n", + " for option in decision[\"options\"]\n", + " if option[\"status\"] in {\"apply\", \"skip\"}\n", + " )\n", + " )\n", + " evidence_ids = list(selected.get(\"requiredEvidenceIds\", []))\n", + " cited_classes = {\n", + " evidence_class_by_id[evidence_id] for evidence_id in evidence_ids\n", + " }\n", + " for evidence_class in selected[\"requiredEvidenceClasses\"]:\n", + " if evidence_class not in cited_classes:\n", + " evidence_ids.append(evidence_by_class[evidence_class])\n", + " state[\"parameter\"] += 1\n", + " return _structured_output(\n", + " info,\n", + " dict(\n", + " selectedOptionId=selected[\"optionId\"],\n", + " evidenceIds=evidence_ids,\n", + " rationale=f\"Use the offered {selected['label']} policy with its required observed evidence.\",\n", + " confidence=\"high\",\n", + " ),\n", + " )\n", + "\n", + " return FunctionModel(reply), state\n", + "\n", + "model, model_state = _scripted_workflow_model()" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "137f5014", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "WARNING: WARNING: Number of valid features is less than value of parameter `top_n`: 33538. Resetting `top_n` to 10744\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "WARNING: WARNING: Number of valid features is less than value of parameter `top_n`: 33538. Resetting `top_n` to 10958\n" + ] + }, + { + "data": { + "text/plain": [ + "{'status': 'completed'}" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "result = analyze_rna(\n", + " source_path,\n", + " model=model,\n", + " study_context=study_context,\n", + " study_objective=\"Discover stable major immune-cell populations.\",\n", + " score_doublets=False,\n", + ")\n", + "{\"status\": result.status}" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "fc2ac52e", + "metadata": { + "tags": [ + "remove-input" + ] + }, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
resolutionclustersrepeat_agreementclusters_with_markersselected
00.5061.0000001.0True
10.7560.9900961.0False
21.0061.0000001.0False
31.2571.0000001.0False
\n", + "
" + ], + "text/plain": [ + " resolution clusters repeat_agreement clusters_with_markers selected\n", + "0 0.50 6 1.000000 1.0 True\n", + "1 0.75 6 0.990096 1.0 False\n", + "2 1.00 6 1.000000 1.0 False\n", + "3 1.25 7 1.000000 1.0 False" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "assessment = model_state[\"assessments\"][-1]\n", + "pd.DataFrame(assessment[\"alternatives\"])" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "9c555061", + "metadata": { + "tags": [ + "remove-input" + ] + }, + "outputs": [ + { + "data": { + "text/plain": [ + "{'Why this setting': 'The teaching policy selected the measured combined representation and its most repeatable eligible partition. Marker identities remain unvalidated. Observed resolution 0.5 has repeat agreement 1.0 and marker coverage 1.0.',\n", + " 'Marker evidence': ['The saved marker preview contains GZMK, NSG1, AQP3, TNFRSF4, CD40LG, DPP4, HNRNPLL, LYAR.'],\n", + " 'Biology to preserve': 'Preserve the single-donor population structure and retain marker uncertainty; no batch or treatment comparison is supported.'}" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "selection = assessment[\"selection\"]\n", + "{\n", + " \"Why this setting\": selection[\"rationale\"],\n", + " \"Marker evidence\": selection[\"qualitativeFindings\"],\n", + " \"Biology to preserve\": selection[\"objectivePreservation\"],\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "9e67a867", + "metadata": {}, + "outputs": [ + { + "data": { + "image/png": "iVBORw0KGgoAAAANSUhEUgAAA48AAAJjCAYAAACsmCRCAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjExLjEsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvctoD+AAAAAlwSFlzAAAPYQAAD2EBqD+naQAA7iFJREFUeJzs3Xd4VFX6B/DvtExmUia990AgBAg19F6liV3sil3XXnd1Leu66/pb17J2FBvqrmsXQVCqUqSXENJI73UmySRT7++PKwMhk0xNAb+f58kDuXPuvefeIWTee855X4kgCAKIiIiIiIiIeiDt7w4QERERERHRwMfgkYiIiIiIiBxi8EhEREREREQOMXgkIiIiIiIihxg8EhERERERkUMMHomIiIiIiMghBo9ERERERETkEINHIiIiIiIicojBIxF1UVRUhAceeACXXnop1qxZ09/d8brt27fjpptu6vZ7b9m4cSP+8Ic/eP24A0Fv3TM6e/XVz1V/n5OI6PdM3t8dIPo9qqysxKpVq1BQUICIiAjceOONGDp0aKc2+/fvx+eff47y8nKkpqbixhtvRExMTKdj3HXXXV2O/cQTT2DEiBFu900QBMybNw9TpkzBJZdcgvT09LPuGhwpKSnBt99+2+333lJYWIh169Z5/bgDgbfv2ZYtW/DJJ5/gzTff9NoxB6pz9Vr76ueqv89JRPR7xpFHoj527NgxDB06FNnZ2ZgzZw4EQUBWVhZ2795ta/Pcc8/hlltugVqtxuzZs3H48GEMGTIEhw8ftrXR6XT4/PPPMX/+fFx++eW2r8jISI/6V11djcLCQjz22GO45JJLMHz48LPuGujsU1xcjLVr1/Z3N/rE7+laiYjo3MKRR6I+9re//Q1Dhw7Ff/7zH9s2o9GI++67D7/88gsA4KqrrsLDDz9se/3aa6/FxIkT8Y9//AMfffRRp+MtWrQIcXFxTp+/uroab7zxBvLz8xEZGYnrr7/eNsr3yy+/4JlnngEA3H333VCr1fi///s/JCUlDahraG1txerVq7Fnzx5ERERg5cqVnUZIW1tb8c4772Dv3r0ICQnB0qVLMXfuXKePb7FY8Omnn2Lr1q2QyWRYsGABli9f7nbbnJwcfPTRR6ioqMCECRNw8803QyaT2V7v6T2xWCy48sor8dhjj9kC+T/96U9obGzE66+/DgCoqqrCH/7wB7z99tvYu3cvvvnmG9x+++09ntMZju7z6T7//HPs27cPzz77rG3brl278Prrr+P999/v8V7t27cPr776KhobG3HxxRcDAFasWIGLLrrI4Xu5ceNGfPPNN7j11lvx1ltvoaqqCu+88w4CAgJcvp6e3gdnr/Fkf7q7/91d6/Lly53+N+fJvT79XrjyM3LkyBE89dRTAACVSoW0tDTcdtttCAsL63YfZwykn2UiInKMI49Efay8vBxDhgzptG3IkCHYsWMHGhsbAQCxsbFd9ouNjUVTU1OX7X/6059w7bXX4tlnn0VNTU2P566rq8Po0aOxc+dOzJ49Gy0tLRg/frwt4EtMTMTixYsBAEuXLsXll1+O4ODgAXUNjY2NGD9+PN577z1kZWUhJiYGV111FcrLywEAzc3NmDBhAjZt2oTZs2cjISEBV111FV588cUej3u6xx9/HI899hhGjhyJsWPH4qOPPsILL7zgVtuqqipccskliIqKwvjx4/H00093CqodvScymQwnTpywjVSZTCa8+OKLWLVqFaqrqwEAP/30E3bt2oXg4GAUFhbi3Xff7fGcJpMJF198MTZu3Oj2fT5TTk4ONm3a1GlbeXk5vv76a4f3KjY2FhMmTIBKpbKNPmdkZDj1XhYWFuKdd97BhRdeiJSUFFx66aVQKpUuX4+j98HZa3R0/7u7Vlf+zXlyrwH3fkYiIyNt/Z07dy4OHz6MESNG2P15dtZA+1kmIiInCETUpx588EEhJiZGqK+vFwRBEEwmkzBnzhwBgLBnzx67+xQUFAgqlUp46aWXbNtycnKEjIwM4R//+Ifw1ltvCXPnzhU0Go1w8ODBbs99zz33CBkZGYLFYrFtu+qqq4SJEyfavs/PzxcACEVFRQPyGu69914hOTlZ6OjosG3T6/VCa2urrW9z5szptM/69esFlUolGAwGQRAE4cMPPxQiIyNtr5/5fWZmpvDaa691OkZdXZ3d/vTU9vXXXxckEolQUFBge+2tt94SoqKibN878548+OCDwoIFCwRBEITt27cLiYmJwrRp04RPPvlEEARBWLlypXDllVc6fc729nYBgPD666/bvSZBcHyfz7xnf/nLX4QJEyZ0OsZnn30maDQap+7V6tWrhdjY2E6vOfNevv766wIA4dChQ91eizPX48z74Mw1OnP/7V2rK//mPL3XztxXZ2RlZQn//Oc/bd87+rk600D7WSYiIsc4bZWojz322GPYvXs3hgwZgqysLOTm5mLs2LEAAIPB0KV9U1MTli9fjqysLNx+++227QkJCdi/fz98fHwAADfddBNmz56Ne+65B5s3b7Z77p9//hnLly+HVHpq0sEll1yCCy+8ECaTCQqFYsBfw4YNG7qMLqlUKtvf169fD0EQcPnll0MQBAiCgPb2drS3t6OwsLDbaZenGzNmDF599VVERkZi1qxZCA4O7nZ6nqO2sbGxSE1NtX2fmpqK6upqWK1WSKVSp96TmTNn4vXXX4fZbMaWLVswa9YsxMfHY8uWLbj88suxZcsWPPLII06f08fHB5999hnGjBnT7T1wdJ/d4cp9BZx/L8PDwzFy5Mgez+3oerz1swE4vv/2uHpvHOnpeO7+jOzYsQNffPEFKisrYTQaUVNTg7y8PLf7ONB+lomIyDFOWyXqY4GBgdi6dSs2b96M22+/Hd999x1uuOEGAOiSKEar1WLBggXw8/PDN998A7n81PMetVptC7pOOv/887F3795uz93Q0NBlGmpISAgsFguam5vPimvQarUIDQ3t9vXm5mZkZmbi4osvxiWXXIJLL70U1157LT777DNERUU5dX2vvfYarrvuOvzrX/9CTEwMZsyY0SnRjyttz5xCeTJ4sFqtAJx7T6ZNm4b29nbs2bMHW7ZswcyZMzFz5kxs3rwZFRUVKCwsxMyZM50+p1QqxcUXX4yUlJRu74Gj++wOV+4r4Px7GRgY6PDcjq7HWz8bgOP7b4+r98aRno7nzs/IBx98gPnz50OhUGDBggW4/PLLkZCQgNbWVrf7ONB+lomIyDGOPBL1kxEjRtiScbz99tuIi4vrNFqh0+mwYMECAMAPP/zg9Afk7kY2ACApKQmFhYWdthUWFkKtViM8PPysuIbExMQeRzvi4+MhCIItGYk7fH198cADD+CBBx5AS0sLbrrpJlx33XXYv3+/R23tceY9CQgIwJgxY7Bhwwbs3LkT7777LiIiIlBSUoI1a9YgNjYWgwYNcvt67XF0n8/k6+vbZdS5oaGhS5vu7pVEIulyTG+8lyc5uh5n3gdnrtEZ9q7V1X9zntxrd+7r6tWrcd999+Hpp5+2bXNl7aE9A+1nmYiIHOPII1Efq62txY4dO2zf79u3D2+++Sb+9Kc/2T5UtrS0YMGCBRAEARs3boRGo+lynLVr16Kurs72fWlpKd58802cf/753Z57xYoV+OSTT1BcXAxAzGT40ksvYcWKFWfNNVx77bVYs2YNDhw4YNu2Y8cOW/KYlStX4r///S+2bNlie91oNGL16tVOX9+7775r+3AeEBCAjIwM6PV6j9va4+x7MnPmTLz66quIiopCQkICfH19kZWVhf/7v//rNOroDGcS5ji6z2caOnQojh8/jsrKSgCAXq/Hu+++26lNT/cqLCwMzc3NMJvNtvbeeC+dvR5n3gdnrtEZ9q7VlX9Hnt5rd+6rj4+P7XyA+DDo9P8D3NHfP8tlZWW4+OKLceTIEY+ug4jo94Qjj0R9TKVS4bHHHkN7eztUKhV2796NBx98ELfeequtzcMPP4xdu3Zh7ty5WLlypW17QkKCLVOgRCLB1KlTERYWBl9fX+zevRsLFizASy+91O25b7jhBmzatAmjRo1CVlYWjh07hqioKPztb387a67hxhtvxLFjxzB58mSMHz8eBoMBvr6++Oabb2zXWFhYiPPOOw+ZmZnw8/PD8ePHcd111zl9fYWFhUhJScGwYcPQ0dGBo0eP4r333vO4rT3OviczZ87E888/j6VLl3batn37dpeDx6amJnz++ee4/vrru23j6D6fadGiRZg6dSpGjx6NsWPH4vjx40hPT0dubq6tTU/3avr06QgODkZWVhZSUlKwYsUKr7yXzl6PM++DM9foDHvX6sq/I0/vtTv39dFHH8XSpUtx9OhRKJVK5ObmOlxn6kh//yxrtVp8/vnnuPXWWzuVZCEiou5JBEEQ+rsTRL83giBg7969aGhowLhx47okcNizZw9KSkq67KfRaDBv3jzb9waDAQcPHoROp0NaWhoSExOdOn9OTg7y8/MRFRWFcePGdZom2tbWhnXr1mHRokVQq9UD9hrKy8tx6NAhREdHY9SoUV2mutbV1WHfvn1QKpUYNWpUp/VspaWlOHz4MJYsWWL3e0AsI7B//34oFAqMGTPGbt1AR21PnDiBvLw8LFy4sFO/tm7diosuuqjT9MWe3hNAHF36/vvvkZmZicGDB9vuwa5duzBr1izb2jFnzrlmzRq88cYb2L59u9v32d49s1qt2Lt3L3Q6HTIzM2E0GrFv3z4sW7bMqfva1taGX3/9FY2NjcjIyMDQoUNt/e/uvbR3ve5cz0mO3gdH1+jse27vWl35N+fpvXZ0X+052d7HxwdZWVk4fvw4BEHA+PHjATj3c2VPf/0s63Q6bNiwAdOnT0dERESPfSQiIhGDRyKi35kjR44gICAASUlJ/d0VIiIiOosweCQiIiIiIiKHmDCHiIiIiIiIHGLwSERERERERA4xeCQiIiIiIiKHGDwSERERERGRQwweiYiIiIiIyKE+Dx4FQYBOpwOTvBIREREREZ09+jx4bGlpgUajQUtLS1+fmoiIiIiIiNzEaatERERERETkEINHIiIiIiIicojBIxERERERETnE4JGIiIiIiIgcYvBIREREREREDjF4JCIiIiIiIocYPBIREREREZFDDB6JiIiIiIjIIQaPRERERERE5BCDRyIiIiIiInKIwSMRERERERE5xOCRiIiIiIiIHGLwSERERERERA4xeCQiIiIiIiKHGDwSERERERGRQwweiYiIiIiIyCEGj0REREREROQQg0ciIiIiIiJyiMEjEREREREROcTgkYiIiIiIiBxi8EhEREREREQOMXgkIiIiIiIihxg8Ep1jrG1t0O/dC6vB0N9dcchqNMLc1OR0e0EQbH+3tLbBVFHRG90iIiIiIjsYPBKdY5q/+BJNn3wK3Xdr++R81vZ2NKxaheb//a9TcOeMhjfeQM1fn0VHTo7DtuaGBlT/+QnUvvQSBEFA/Ssvo/aFf6Hj+HF3u05ERERELmDwSHSO8UmIh0Quh09CfJ+cz1hWho6c42jbtRtCR4dL+woWKyAI4p8OWLRaWPV6mGtqAbMZEpUKkEkh8VG623UiIiIicoFEcHWowEM6nQ4ajQZarRaBgYF9eWoi6gWCIKB182bINBqox451aV9rRwcsWh0UkRFOte/Iy4MsIACK6GgIJhOsBgNk/v7udJuIiIiIXMTgkYiIiIiIiBzitFUi6pG1rQ11L7+MhtXvubym8WxgbmqCtb29v7tBRERENOAxeCQ6x7Vu3YqKhx5C265dbu1vqqmBsaQUHdnZEPo5yOo4dgx1r74KQ0GBV45nLClBzbN/Q+2//tXltdbt26Hfs8cr5yEiIiI6F8j7uwNE5Ji1vR11L74IAAi/915IfX2d3tdYXg5YrDCVl7t1bmVKCoIuvRSyAH9I1Wq3juEMwWKBRCbrsU3bzp0wniiCfs8eKAcN8sJJhd++xLWbEokEgBhUar/6GgCgHDIEMk6xJyIiImLwSNQfTFVVaP35Z/hPnQpFdLTD9tb2dpgbGm1/dyV4DLrgAvimp8M3I8Pt/vpNyHJ7X2cYCgrQ8PYq+A4fjpCrr+q2XeDChZAFh8B/2lSvnNcnKQmRf3wUUpXKFjgCgDwqCr7D0iHx9YXU3x+CxYLG9z+AtV2P0BtugFSl8sr5iYiIiM4mDB6J+kHLxh/RfugQhPZ2hFxzjcP28pAQhN1+OyAB5MHBLp1LqlZDPWaMu121y1xfj9atW6GeMAE+cXFeOF4DBLMZ5prqHtspYmMRdOEFTh9Xt349DCdOIHjFim7vmzwkpMs2qVKJ0JUrYW1rg0QqhUWnQ8exY4AgwFxXB5+EBKf7QERERHSuYPBI1A/8Jk+CVd8G9cSJTu+jTEn2ah9MNTWQh4ZCIu/634C5sRHaL7+E7/ARdkcdW37aBP2vv8Lc2Iiwm25y6bztBw9C4usL36FDbdvUE7IgCwyA3IlRWFe0/vwzhPYOGPLyIXdx9LRl82bovluLgHnzELhwAUKuuRpWfTsDRyIiIvrdYvBI1A98kpMRduutXjlW69ataD90GEGXXgJFVJRz+/z8C7RffglV5ki7I5/tBw+i41gOzLW1doNH9fjxsDQ2wG/SZJf6aiwvR+OHHwESCaKe+DNkAQEAAIlEAt9hw1w6ljNCrrwSxpJSqEaPcnlfq04HALBotQAA1ciR3uwaERER0VmHwSNRH9Pv34+mjz+B/4wZ0CxdAqBzshZXte3YCXN9PTqOHXM6eISky186UWdlwdLYCGV6ut3XlSnJUN52m8t9lYeEQBEbC6nKt0/WDfqmp8O3m2twJHDxYijT06FMSnLYtiMvDzI/PyhiY906FxEREdHZQCL0ceE2nU4HjUYDrVaLQGYwpN8h3fr1aNn4I3yHpSN05UqYGxtR9+JLkAUHI/zuuyCRulZBx1BYCENuLvxnzXIpIDPV1EIeGmJ32io5z1BUhPp/vwqJQoGop5+C1Menv7tERERE1Cv4qZGojwXMnQtFbJxtDaNFq4W1rQ2CyQSYzYCLwYc8MgqK+PgegxZrWxuM5RVQpg22jXAqIiPcv4g+0JGbi46jRxEwb57LpTIa3l0NU1UVwm69BfLQ0F7qoUgWFASpvz9kQUEMxImIiOicxk86RH1MIpdDNWK47XvBaILE1xf+M2dA4mLgaKqpQd0L/4I0MBCRjzzcbZ3Exg8/giE/H4FLFiNg1iyP+n+S9ptv0Pbrrwi55hr4pqV1ek0QBFjq6yELC3N7Oq72629grqmBNCAAgfPnO72fIAgw5OdDMBphqq6GPDQUpspK6NatgzprQqd77w3y4GBEPfmE09dpqqmFVOkDWVCQV/tBRERE1Ntcmx9HRF3oDxxA/dtvw1RV5db+7YcPQejogKmkxOV9BZMJgsUCwWgUi913QxYcDEgkLpf56ImhqAhCewdM5RVdXtN99x1q/v4cWn74we3jB8yaCd/0oVCNGuVU++YvvkTtiy/C0tyMsFtvQfCVV9iS8LT9+is6juWgdcsWAED7kSMwlpW53bczOR04VlSg9vnnUfvPF8T3jIiIiOgswpFHIg+1bt0KU1k59NEx0CxZbNvekZODxjVr4D9lCgLPO6/b/QPnzYMsIACq0a7XYvSJi0PEgw9CqvKFRC4Xp74KQpcRzKBLL4Fm+fmQKpUun6M7IVdfDeOJE3aDO8Fs6fSnO9Tjx0M9frzT7fV79kAwGmEsLoZ69Gj4JCbaXvOfNg1ChwHqcWPRkZeHxvfeh0ShQPRfnoZEoXCrf4LZDItWa3dabNNnn8FUWoqQ667r/LpCAYlcDolSCbi4tpWIiIiovzF4JPKQZvFi6A8ehN+UzmUrjKVlYo3BwhM97i8LCkLgwoVun//k2kWrwYDafzwPwWRCxAP3d1onKJFIxIDFA5bWVsj8/W3fy0NCIA8JsdtWc/4y+GWNhzwmxqNzuiLk+utgrqqyW1JDHhqK4MsvAyDWsJRpNJCHhwEerFFs/OBDtB8+DN+MYdAsWwZFZKTttfYDByEYDDDk5cMcVAtlWhokMhkUERGIevwxWxBJREREdDbhpxciD0l8fKAcNKjLlNCAWTMhCw6CcvBgl48pGI0ur38UTCZYW1shWK2wdnS4nGSmJ7ofNqBlwwYEnrcQAXPnOmwvkUo9LlshCAJaNmyE1FcJ/xkzHLb3TUsDzlh7aY88JARRf37c9r1u40a0bduG4BUrXKs1abXAXFODtuZmCB0GhP/hTttLoStvgKmyCu3Z2TDk5CBg3jwELlwAAJD6+Tl/DiIiIqIBhMEjkQcEiwX1r78BwWSCVKWC79ChttckPj7wy8py+Zhtu3ah+bP/IWDObAQuWuT0fjJ/f4Tfdx9gNkER4d1MqtbWFgCApaXVq8ftiam0FC0bNgAAfEeO9Op6zdMZCwpg1bfDWFzcbfBoLC8XRw6jo23bQq65Br7Dh6Nl0yb4Ds/o1F6ZmgplaiosjQ0wHJdAHmp/hJaIiIjobMLgkcgDEpkMykGpYs1ELwVs5vqG3/6sd3nf3iq/oVm2DKrMTPgkJTnVvnX7dkh8lPCb4HrwfJIiJgbqcWMh8VV5LTNp26+/QvftdwhcssTWt6DLL4chNw+qMaPt7mOur0fdSy9DIpUg8k9/OjWiK5dDOWgQ/CZO7PZ8mvPPR8B557H2IxEREZ0TGDwSeSj0xhu9erzAhQugTEmGT0qKw7YWrRZStbrHpC9WoxGW5maPRiMlCgWUgwY51dZYVgbtV18DAHyHpDkd+BlOFEEWHGQbYZQoFAhescKt/goWC9p27YIiJgbK5ORTfSsshFWvh6GwwBY8yoODIZ84odtjSXx9IfVTQyJXdFo3qv3yS7Tt2InAxYsQMHt2t/szcCQiIqJzBYNHIg8IggBzVRXkkZHd1lh0lUQud2rtXUdODhreeRc+KckIv/32bts1vPEmjCUlUGdlQSKTInDxYkhVKq/01R55ZCR8hw2DRKmENCDAqX06cnPR8NbbkPr7u1QzsTvthw5B+8WXkKh8EfPMM7btgUuXwScxEb52kup0R+bvj6jHHgMkkk7vsS2jrMnsUV+JiIiIzhYMHok80LJxI1p+2AD1xAnwnzED8qAglxLdtG7fDlN1NTTLlrlcRuNkWQ7B4KBe4G+BWOuWLZCq1VAkJLi1FtNZUh8fhK68waV9ZP7+kCgUkIUEexw4AoBPQgLkkZHwSU464zx+8Js82f5OPTgzM6pgNsM3YxhU48dDecY5iIiIiM5VEkHoobJ4L9DpdNBoNNBqtQj0YjZIov7Q8tNP0H2/DorYGJgqKqFMS0PYLTc7ta8gCKh88CEIVitCrrkaajv1Eh0xVVRAFhLS40iisboadf96EbCY4TsyE0EXXgiZ/8DL+Gk1GsUaiG7WPxQEAR1Hj0IeEdGpbIYnLC0tMNfWQpma2mm79tvv0LplC1SjRiHk6qu8ci4iIiKigY5VqoncIAgCWrdvhyw4GBEPPwS/6dNdPoZEIoH/jOkwV1WhZeNGOHqO0/Tpp6h98UVYmptt2xSxsQ6noAptbYDZDInCByErLvcocOw4dgz1b78NY3m528fojtTHB4LZDHNTk8O2ptpatGzeDKtef6pvhw6h8b33Uf/GG17rU8Obb6L+tdeh37ev0/aT2VOZRZWIiIh+TzhtlcgNxuJiW1KYqKeehN+4ceJUSRezgvpmZEAREwNLUzNgsXRbtF6wWKDfvx+CyQxjeQVULpxHmZqK0JtuhCwgoMfEOs5o3bYNhvwC6ENC4BMX59Gx7Kl/7TWYyisQetON8B0yBBadDub6eijPSB6k/fxzGAoKYW1tg2bpEgCAPDwcUrUKPgmJXuuPLCgIptraLms3/SZPhmrcOCbDISIiot8VBo9EblBER0M5ZAikfmpb0XdH2UwNhYVo3boNAXNmwydRDHCUKSkIu+1WSAMCuqyrO51EJoMiKgrth48AUtfXBJ5ef9ITAQsWQBYcAn8HI63m+nqY6+rgm57u2gksVkAQAKsVAFD/xpsw19Qg+MoroT6tlIbv8OEwNzVBOSTNtk0RG4vov/zFpdNZOzpgyM2FMj3dbiAYsnIlBKPR7npUBo5ERET0e8PgkcgNUl9fhN18U49tBJMJdf9+FUJHO8L+cBdat25DR3Y2JEofhCSeGh1ztgQGZDLIQ0NhPW3aal9TJid3Kn3RnfrXXodFq0XINVdDlZnp9PHD7rgdFp3OFojLgoJgrq+DLLDzyJ//tGnwnzbNtc7bof3qK+j37IXfpIkIuvhi2/aOnBx05BxHwPz5A3J9KBEREVF/YPBI1EusBgNMlZWA1QqrTouAObMhUfogYOZMt44XesMNMFWKSXncZaqshDw83OPpq44oYqJhbW+HLDTUpf2kvr6Q+vravg+96cZuR/68QR4ZCUgkkEdFd9re/MWXsDQ2QhYUhIDZs3rl3ERERERnG2ZbJepFhhMnIBiNXpk2KphMaD96FMrUVMjc+Nlp3b4d2q++hu+I4Qi97jqP++OIIAgel90w1dZC5u8PqVrtpV4BltZWaL/6GspBg+A3cQIEq7VLhte2nTvRkZ0NzQUXQH5aAGyqqUHje+/Dd+gQaM4/32t9IiIiIjobMNsqUS9SpqR4bb1hy6bNaPpoDZo+/gTmpianspKe7uSaSonM+xMO2o8cQf1bb8NYXnHqfE4Gjtb2djR+8AF069d32m4oKEDtP55H3Usvedw//f79qH/7bZhqatB+6BDaDxyAbu1aAIC5rg71b70N/f79tvZ+kyYh9MYbOwWOAGAsKoK5thbthw573CciIiKisw2nrRL1E8FoRNuuXfBJSYVPXKzD9oqYaEAmhSw4GLV/fw6QShH5x0chOyMTaHf8Jk2CcvBgyIKDPe16F63btsF4ogj68HCnruV0hsITtmAsYO7cU4mDpFJAIgG8EOy2bt4CU2Ul2uP2w2/adJjKymy1G9v37YMhNxfWlhaox4zp8TjqceMgmEzwSUjwuE9EREREZxtOWyXqJ63btkH79TeQh4Uh8tFHnNpHEARYW1pQ849/QCKRIvKRh23ZXvuT4UQR9Hv3IGDOnC6jdd0RBAGCyQSJRALd+h8gDw+H38QJndpYmpshUas9zmzakZeHjsOHETB3LmRnlDmxNDej5aefoBo5EsrBgz06DxEREdG5jMEjUT+wtrXB1NCA5jUfw3fkCGgWL3Z5f0ilkKpUXu2X/sABWBob4T9rVpd1gN7W8N576Dh2DKErb4TvEPeTABERERFR3+C0VaJeZCgoQOvWrQiYMwc+SUnithMnUP/GG/CJT3B6xPFMvTHaKBiNaFrzMSAIUMTEuF6j0Ummyko0/fe/MBYVQyKXw9KPpUeIiIiIyHlMmEPUi1q3bUfHsRy0/vKLbZtV3w5YrOLo4WksOh0aP/wIbbt/7XIc/Z490O/b12mbsbwcVr3ea32V+PhANXo0fFJSul3TJ1itaNm0Ge2H3U8Y03H8OExl5ZCFhSH0lpvhNyHL7WMB4rRTq9Ho0THcPrdWC0trm+OGREREROcAjjwS9aKAuXMg8VUiYMYM2zbV8AyE330XZCEhndq2HzqE9oMHYThR2CmgMtXUoOnT/wAAfBITIQ8Lg37/ATStWQNFfBwi7rnHK30119ej48gRSJRKQG7/vwbD8eNillKJBDF/e9bpepGG/HyYKivhN3Uq/CZNgmAyw3dImm001l2G/HzUv/kWFDExiLjvXnTk5kEeGgJ5WJhHx3WGubERtc/9AxIfH0T+8VEAQPNn/4M8OgqB8+b1+vmJiIiI+hqDR6Je5JOQgJArrrC7/Uyq0aPR8tMmWBoa0ZGXB980cR2gPDgYyiFDIJFKINNoAABSXyUgkUCq8l79Q8FqBQQrYLUCAFp//gW6b7+FZvn58Js0CQCgSEiEcvAgyMPCnA4cAaBh9XsQDAZIVCqox49H4IL53umzxQIIAgSzGR3HjqHhnXch9fdH9FNPeuX4J7UfOQJ5SAgUsadlkj1jubghLw/thw4BvyXm8bTGJREREdFAw4Q5RANI3Sv/hrG4GAFz58B/1iyYysvhk5pqNxCxtLRAqlZDIpN57fzmpiZI5HLIAgLQ9Mkn0O/dB3VWFoIvu9Sj4zZ/8SX0Bw/CqtVCNXIEQleu9FKPAVNNLWSBAbA0NaHu369CERuD8Dvu8NrxO3Jy0LDqHUiUSkT/9ZlO74WluRmQyyHz94fVaITu+++hiIqC38SJXjs/ERER0UDB4JFoADHX1aHj+HGos7LQ9OGH6Mg5jsDFixEwe1af98Wq16MjOxu+w4d7Jaurft8+NH38CeSRkYh86EEv9LArwWLxWjDdkZMDwWyGIioKda++CkV0DMJuudkrxyYiIiI6G3HaKtEAIg8Ph394OABA+tvDFVlggFeObTUaXaqXKFWroR4/vst2wSpObZV0sy6yO+qxYyENCIAiMtKl/VzhrcDR0tyMhnfeBQQBEQ/cj+gnn/TKcYmIiIjOZgweiQaooEsugWbJEkjVnq9r1H63Fq2bNyPokos9mlIpmM2o/ecLsLa2IPy++yAPDu62rbmxEc3/+xy+Q4fAf/p0ALCt4/SUta0NLT/9BN9hw6AcNMgrxzyd1M8PPslJEEwmyIKCvH58IiIiorMRS3UQDVASicSjwNFqNKL9yBFYOzpgaWwAAFgaGz3qk2CxiKUx9O1dSo2cqSM7G4bcXLRs3uzROe1p27kTrVu3ofnzL7q81n7kCBo//hjmpia3jy9RKBB+xx2IuOcer0zZJSIiIjoXcOSR6BylW/s92n7+GaoxoxF06aVQjx8PpYcjf1KlEuH33A2hvR0+cXE9tlWPGwdzQwOUgwa7fT5DURGaP/0UqnHjOpW/8B0xAob8fPiOHNllH93362CurYU8LAyB872T1bU7giCgdetWyIODocrM7NVzEREREfU3Bo9E5yhFZAQgkUARFQWpry9809O9dFzn1ixKVSoELV/u0bmMhYUw1zeg48jRTsGjIjISYbfdZnefgPnz0HHkCNTjxnl0bmcY8vKg+/Y7QCJBdHq6S2tKiYiIiM42DB6JzlF+kydDPWGCV0t59DW/adMgUShcGjFVjx4N9ejRvdirU3zi4qAcPAiykBAGjkRERHTOY6kOIiIiIiIicogJc4iIiIiIiMghBo9ERERERETkEINHIiIiIiIicojBIxERERERETnE4JGIBgxTbS0aP/wIHbm5/d0VIiIiIjoDS3UQ0YDRtmMH2g8ehKW5Gb5DhvR3d4iIiIjoNAweiWjA8Js0CVadDurx4/u7K0RERER0BtZ5JCIiIiIiIoe45pGIiIiIiIgcYvBIREREREREDjF4JCIiIiIiIocYPBIREREREZFDDB6JiIiIiIjIIQaPRERERERE5BCDRyIiIiIiInKIwSMRERERERE5xOCRiIiIiIiIHGLwSERERERERA4xeCQiIiIiIiKHGDwSERERERGRQwweiYiIiIiIyCEGj0REREREROQQg0ciIiIiIiJyiMEjEREREREROcTgkYiIiIiIiBxi8EhEREREREQOMXgkIiIiIiIihxg8EhERERERkUMMHomIiIiIiMghBo9ERERERETkEINHIiIiIiIicojBIxERERERETnE4JGIiIiIiIgcYvBIREREREREDjF4JCIiIiIiIocYPBIREREREZFDDB6JiIiIiIjIIQaPRERERERE5BCDRyIiIiIiInKIwSMRERERERE5xOCRiIiIiIiIHGLwSERERERERA4xeCQiIiIiIiKHGDwSERERERGRQwweiYiIiIiIyCEGj0REREREROQQg0ciIiIiIiJyiMEjEREREREROcTgkYiIiIiIiBxi8EhEREREREQOMXgkIiIiIiIihxg8EhERERERkUMMHomIiIiIiMghBo9ERERERETkEINHIiIiIiIicojBIxERERERETnE4JGIiIiIiIgcYvBIREREREREDjF4JCIiIiIiIocYPBIREREREZFDDB6JiIiIiIjIIQaPRERERERE5BCDRyIiIiIiInKIwSMRERERERE5xOCRiIiIiIiIHGLwSERERERERA4xeCQiIiIiIiKHGDwSERERERGRQwweiYiIiIiIyCEGj0REREREROQQg0ciIiIiIiJyiMEjEREREREROcTgkYiIiIiIiBxi8EhEREREREQOMXgkIiIiIiIihxg8EhERERERkUMMHomIiIiIiMghBo9ERERERETkEINHIiIiIiIicojBIxERERERETnE4JGIiIiIiIgcYvBIREREREREDjF4JCIiIiIiIocYPBIREREREZFDDB6JiIiIiIjIIQaPRERERERE5BCDRyIiIiIiInKIwSMRERERERE5xOCRiIiIiIiIHGLwSERERERERA4xeCQiIiIiIiKHGDwSERERERGRQwweiYiIiIiIyCEGj0REREREROQQg0ciIiIiIiJyiMEjEREREREROcTgkYiIiIiIiBxi8EhEREREREQOMXgkIiIiIiIihxg8EhERERERkUMMHomIiIiIiMghBo9ERERERETkEINHIiIiIiIicojBIxERERERETnE4JGIiIiIiIgcYvBIREREREREDjF4JCIiIiIiIocYPBIREREREZFDDB6JiIiIiIjIIQaPRERERERE5BCDRyIiIiIiInKIwSMRERERERE5xOCRiIiIiIiIHGLwSERERERERA4xeCQiIiIiIiKHGDwSERERERGRQwweiYiIiIiIyCEGj0REREREROQQg0ciIiIiIiJyiMEjEREREREROcTgkYiIiIiIiBxi8EhEREREREQOMXgkIiIiIiIihxg8EhERERERkUMMHomIiIiIiMghBo9ERERERETkEINHIiIiIiIicojBIxERERERETnE4JGIiIiIiIgcYvBI5xSL1YJ9NftQp6/r764QEREREZ1TGDzSWctgMWBHxQ7Ut9fbtu2u3o3/5P4H72W/1+vnb+powp7qPTBZTL1+LiIiIiKi/sbgkQaMbeXb8MLeF1CqK3Wq/ZayLfiq8Ct8lvuZbVusfyz8FH5IDUrtpV6e8mnup/gs7zNsKd/S6+ciIiIiIupvDB5pwDhYexDV+mocbzzuVPsUTQoCfQKRFpJm25YYmIjzks5DrH9sb3XTJlmTDJVchfiA+F4/FxERERFRf5MIgiD05Ql1Oh00Gg20Wi0CAwP78tQ0wFW2ViKnMQeTYyZDJVe5dYw6fR2e3/s8AOD+sfcj0i/Sm110SBAEFGmLEOUXBbVC3afnJiIiIiLqTfL+7gDRSTH+MYjxj/HoGEG+QcgIzQAAhKhCvNEtm1JdKXRGHYaHDe+2zc6qnfiq4Cska5JxW+ZtXj0/EREREVF/YvBI5xSFVIFrM671+nGNFiPePPwmTFYTbh15K1KCUuy2C/QJhAQSBPrYH1W3WC34vuh7+Mp9MS9xntf7SURERETUWxg80lnNaDFCgAClTNlju+ONx7G1bCvmJc7rNvDriUKqQEJAAho7GhHi2/2I5vCw4Xhy8pPwlfnafb28tRzbK7YDALKisqBRalzuCxERERFRf2DwSGctg8WA/9vzfzBajbhv7H09BmK7q3ajUFuIoOogt4JHiUSCWzJvcaptT+s14/zjMDlmMnzlvj321ypYsaF4A1RyFWbEz3C5v0RERP2qfC+gUAORw/q7J0TkRQweqVdYBSvaze3wU/h5dJz9NfvxVcFXmJ80H1Njp3Z6zWK1oN3cDotggcnac63FeYnzEOgT2OUYfU0mlWH5oOUO25XqSrGpbBMAYFTEqG4Dzfr2eryf/T6SApNwUdpF3uwqERGRexqLgF2vAxIpsOQFQBnQ3z0iIi9h8Ei94qNjHyG7IRvXDLsGGWEZbh+nRFeCDksHirRFXQI/tUKNe8beA4vVgjBVWI/HifGPwQWDL3DqnGW6MryX/R6Ghw13ep/TGS1G+Mh8XN7vdLEBsRgfNR4qucq2frK6rRrlLeUYHTEaMqkMgBhk1uhr0GJsYfBIREQDg184oIkFfALE0UezESjaBkQMBTRx/d07IvIAg0fqFW2mNggQ0GZq8+g4C5MXItY/Fumh6XZfPz1oNFlMUMgUHp0PAMpaytBiakFBc4HL+24s2YiNJRtxfur5mBI7xe0+KKQKjI0ci1+rfkWtvhaRfpFYfXQ1mgxNsApWZEVnAQAywzPRYe7wOEstERGR1yj9gXlPn/o+fyNw5DMxoDx9OxGddRg8klPKWsqgNWh7LFNxumszrkWtvhZJmiSPzquSq2yBUk+2lm3F2qK1WJS8CDPjZzpsb7KY0GxoxvdF32NUxChkhmfaXpsQPQEKmZggx1WN7Y0AgIaOBpf3PdNPJT8hvzkfCqkCF6VdhLTgNOQ05nQKFGVSGSbHTvb4XERE9DvQVg/kbwASJwPBSS7s1wD8+hYQNhgYcbHr5w0bDPiFATGjXd+XiAYUBo/kkMVqwRuH3oDJasKNI25EWnBap9ffPvw2qtqqcPuo220jgWqF2uPA0RX17fWd/uzJ//L+hz3VezAoaBDym/PR1NHUKXiUSWUYHzXerX4sH7wcI8NHYlDwIJf2EwQBhc2FiA2ItSXcmR43HT4yH0yKmQQAnJZKRESeyf0eOLEV0FUC0x9wfr/6XKChANCWuxc8hqYC5z3nuF3eD8DxtcDYa4HYsa6fh4h6HYNHckgmlSEpMAm1+touawsFQUBpSykMFgMa2hscrj3sLUtTl2JY6DCkBqU6bFvfXg8BAmL8Y6BRajoFjicZLAZ8lvsZNEoNlqYudbofSpmy2ym2PdlWvg1ri9YiLTgNN464EQAwJGQIhoQMcflYREREdiVMBnRVQMos1/aLGw/oG4HgxN7p10l1xwFjG1Cfz+CRaICSCIIg9OUJdTodNBoNtFotAgPtF1KngctgMcAqWDuVo6hsrUSToQkZofYT4wiCAIlE0ldddKjV2Iry1nKkBadBKpHabZPXlIdVR1YBAJ6a/FS35TfMVjPWFa1DoE+g3ZIaWoMW3xd9j/SQdIyKGNVtnw7WHsQnxz/BuKhxuCTtEtcvioiI6GynbwSqDwPxEwBF92WviKj/2P/kTGRHh7kDz+95Hn//9e/QGrS27TH+Md0Gjjsrd+LRnx/F1rKtTp3j/ez38ezuZ52afvpxzsd4ef/LaDG2OHcBv/H38cfQkKHdBo4AkKpJxZyEObho8EU91m0s0hZhe8V2rC1aC71J3+X1A7UHcKD2ANYXr++xT6MiRuHJyU/i4sFuTAdyU3ZDNp7c8SQ2lW7qs3MSEdFZKm8DsGcVYOz6u85r1CFAykwGjme5Tz/9FHfccUd/d8NlDzzwAH788Uen2595nX1x3Z6c8z//+Q+effZZj/vA4JGcZhWs6DB3wGgxOqyr2GHuQHlLOarbqmEVrKhsq3R4fEEQUNBcgGZDM2r0Nbbthc2FOFh7sFNbk9WEI/VHUN5ajspWx8d2lUwqw4KkBZgQPaHHdkmBSRgfNR5zEuZArVB3eX10xGiMjhiNhUkLHZ5TJVe5NEIrCAJOaE+gw9zh9D6nK9YWQ2/W44T2hFv7ExFRH7CYxKmm/UkQxGypJTuBmqPeO66+Edj6PHD0c+8dk3rd/v378eCDD2LZsmW4/PLL8dRTT6GsrMz2enFxMXbv3u21861ZswZ33XWX145nz+bNm/HZZ59h+vTpTu9z5nV6+7q9fc558+bhn//8J7Kzsz3qA9c8ktPUCjXuG3efU3UVV2evRpG2CBcMugBXDL3CqbV7EokEl6VdBrNgxrCQYQDEabKrjqyCRbAgUBmIFE0KALGUxXUZ16Gxo7FLAh97BEFAeWs5ovyioJC6V85jT/UelLWUYVHyIvjKfcV+yBQ9TjPVKDVYMXSFW+c7KbcxF2q5GvGB8Z22b6/Yju9OfNdpnaQr5ibORZgqDEOCua6SiGjA2vUaUHUYGHc9kDTVfpsOLaAMBHpriYhEAoxaATSXAVEje25bfRQ48l9gyCIgYWLPbeuOi19NRcBwJoU7GzzxxBP4xz/+gZtvvhnXX389LBYLcnJyMHfuXPz3v/9FZmbXPBKeKioqwq+//ur1457u2WefxS233AIfH8/qdA9kISEhuPjii/H888/jvffec/s4DB7JJSG+IU6185X5QgIJAnwCnC7vcbJG4oy4Gbb1gQqpAqlBqdAatF0CVleSyWwt34rvi77HyLCRuGrYVU7vd7qvCr6CyWpCuCock2MmQyaVuXUcV5TqSvHO0XegkCrw+MTHbUErAAT4BAAAAn3cWzuslCkdjqwSEdFA0U1gWLQd2PeeGFiOu967p2woFKeSqoKB1NnO7VOxF9BWAKW7HAePceOBtjogqJcT8ZBXfPnll3j66afx9ddfY9myZZ1ee+SRR2AwGOzu99JLL0Gr1eLPf/6zbdsXX3yBdevW4e233wYAtLW14dVXX8Xu3bvh7++P888/HxdeeCG+++47vPfee6irq8PcuXMBAA899BDmz5+P8vJyvPLKK8jOzkZUVBQuv/xyWxsAeO+993DkyBEsWLAA7733HnQ6Hb777rsu/SspKcGPP/6IN954w7Ztw4YN+Mc//gEA8Pf3R0ZGBu69916EhXmWGLKiogKvvvoqjhw5gvj4eNx9990YMuTU51lH1+RId/fxpEsvvRRLlizBa6+9BrW664w5ZzB4JABAQ3sD3jr8FqL8onD9cM9/+Vwz7Bq0mdtcCmxMFnEq7OlTYqUSabejaoIgwGAxdAqouuMjFZ8k+cjcf6K0MGkh9lTvwVcFXyG3Kdet0T5XaZQaaHw0CFQGdhkxHR0xGkOCh0AlV8FkNeHzvM+hkquwLHXZgEpQREREHph4O6BvAAKi7L9ubP3tzzbvnrfqEPDLy4A6FFj0D+f3S18G+GqAeAeBIwDIFMCw893vI/WpV155BZMnT+4SOAKAQqGAQmF/ZldOTg7q6zvnsigtLcXOnTtt31933XUoLS3FfffdB6vVik8//RStra2YM2cOpk2bht27d+ORRx4BAGRkZCA/Px9TpkzB5ZdfjptuugmlpaVYsWIFnnvuOdxwww0AgIKCArz55pv4+eefcc899yAqyv7P0ObNmxEeHo7U1FMZ+0eMGGE7n06nw5o1azB+/HgcO3YMKpV7a3Jzc3MxZcoUTJkyBVdddRWampqwYsUK/PLLL1CpVE5dkyPd3cdrrrkGADBx4kQYjUb8/PPPmD9/vlvXweDxd8hsNSO3KRfJgcm2dXq1+lo0GZrQZmrzSnZUmVTmdOB4uO4wchpyMD9pPkaEj0CMX4xT+63JWYMj9Udwbca1GBY6rMe2k2MnIz00HRqlxqlj2zMtbhr8ffzxyfFPbMlx8pvyEaGO8Oi4PdEoNfjTxD91+/rJ96+ypRL7a/cDAGbGz+yxP6W6Unxd+DUmRk90u54lERH1EZmi+8ARANIWAqGDgaD47tu4w8cfkMoBVZBr+6lDgIwLvNsXGhD2799vC0K8bfPmzVi9ejWWLhXLo61YsQItLS0ICAhAamqqbWrsSRdddBEuvPBCvPzyy7ZtQUFBePTRRzsFWoIgYO3atT2OGBYUFCA+vvPPT3R0NKKjo23fL1++HIMHD8aXX36JK664wq1rfOihhzBixAh8/fXXtm3XXHONbarsI4884tQ19aS7+3iSn58fQkJCUFBQwOCRnPdj6Y/YVLoJGaEZuDbjWgBAemg6rhh6BUJ9Q/t81Gpd0To0dDQgXB2O2QlOTosB0GxohgABOqOu0/amjiao5KouI5LBvsEe93V0xGiEqcIQ6huKg7UH8fHxjxHqG4qHsx7u1O6nkp9wtOEoLhtyGaL8evil7yUJAQmYmzAXKrnKYSD7c8XPyG/Kh0wiY/BIRNSfGk8AfhGA0t/9Y0gkQNgg7/XppNBUYOmLgAczdjyy7z2guRSYdKcYkFK/0+v1CAoK6pVjT5kyBY8//jhaW1sxa9YsREVFISAgoNv2P/30E+Lj47Fw4UIIggBBEKDValFRUQGtVguNRvwslJ6e7nCqaXt7O3x9O39mtFqt+Oyzz7Bu3TpUV1fDbDajqakJBQUFbl/jpk2b8Nxzz3XadvrUUWevqSfO3EeVSgW93v2syQwezwLbyrehuq0ay1KXOTVF05FwVTgkkHRZQ9hTHcKeNHU0obylHBlhGT2Wv+jO/KT5yK7PxuiI0S7td/3w61HdVo3UoFPTDEp0JXj90OsIVgZ3Cei8JT5AfDoVpAyCQqpAuDq8S5u9NXvR0NGAvKa8LsGjzqhDma4MQ0OGem3dpEQiwfwkx0+QdEYdDtQegNagxYy4rnUpvSG7IRv+Cn8kBnINCxFRt8r3ArteBzRxwLyn+rs3ouYyIP8HcX1jSEr/lcwQBHHNpMUkJtNh8DggxMbGori4uFeO/emnn+Ktt97C6tWrcdNNNyEzMxOrVq1Cenq63fYtLS1YvHix3fWAp08r9fd3/GAmPDwcDQ0NnbY9/vjjeP/99/HAAw8gJSUFarUa999/v0dBl6Pg29lr6okz97GxsREREREu9/8kBo8DnCAIWHtiLQQIGBoyFCPDHWQ5c8LYyLEYGT7SbtZRQRBQ3VaNSL9IpwPB97PfR2VbJc5PPR9TYqfYtm8q3YRDdYdw6ZBLEesf2+3+J8tZuMpP4dcpcDzZf0EQIEBw+XiuStIk4cnJT0Iu6fpjdNmQy1DYXGg3Ic2aY2tQpCvC4uTFmBHfNYBrMbbg16pfMTJ8pN3A1FWFzYXQm/QYET4CcokcAT4B8FP49UpwV6IrwfvZ70MukeOJyU9AKVN6/RxERGeV1lpg57+BsDRg9GkJ2xQqQCIVp4f2BosZqDoontfXyfwDeevFoM3UDkz+A1CTDegqxWCyD5LE2Ugk4vl1VUC0658PqHcsXrwYn376qdOjYCepVCp0dHQuK1ZbW9ulzd133427774ber0el1xyCe677z6sW7cOUmnXz6NJSUkwGAwuJZPpzrhx4/D4449Dr9fbRgI//fRTPP3007bpooIgoKmpyaPzJCcn49ixY92+7o1r6uk+AkBhYSHa2tqQlZXl9jlY53GAk0gkWJq6FFlRWS5lF3Wku3IV64rW4V/7/4W1J9Y6fazYgFhIJVIUNBfgeONx2/b9NftR1VaF/KZ8j/pa316PNw+9iW3l2xy2TdIk4aHxD+HuMXd7dE5nKaQKu9N8kzRJmJM4x27wFOkXCZlE1m1g+GPJj/ih5Ad8U/iNx/3Tm/RYdWQVPsz5EKW6UqgVajyS9QgenfAo/HvhA0uQMgghviGIC4hzuyQKEdE5paFAzD5adkYttsgMYMkLwLT7eue8uWvFkc09q5zfJ3U2ED0SGDRPHP3b8Qpw6FOg8kDv9LEnkRnA4LmAncDBbdoK4Ju7gJ9f9N4xf0f++Mc/Qi6X4/rrr+8USFksFnz00Uc4fvy43f0yMjKwa9cu2+heZWUlPvzwQ9vrVqsV//znP23ZWtVqNcLCwiAI4kBAeHg4qqurOx3z1ltvxVtvvYUdO3bYttXX13fKmOqsqVOnIjAwEFu2bLFt8/PzQ25uru37F198sVMtS3fceOONeO2113DkyBHbto8//tgWWHt6TY7uIyBOjU1LS8PQoUPdvg6OPJ4FpsZ2U9fJTe3mdrQYWxCh7jpkfXK00ZV1j5ekXYJ4/3h8UfAFirRFeHLykwDE0be85jxMipnkVj9zG3OhkqtQrCtGobYQjR2NmB7nuHhrqCrUrfP1lQsHX4jzU8/vNGW1VFeKL/K/wPio8RgWOgwFzQVeGWVWypRICkxCi7HFVmalp4yzRosROqPOYR3P7miUGjyS9Yhb+xIRnZPiJwBGPRBsZ7aHsvs1XR4LjBVHCzUuJNIJTQWmnPbwNWES0FwCBCd7v3/9obVGzErbWNjfPTkrxcTE4Oeff8Ztt92G2NhYjBw5EmazGUVFRVi6dCnOO+88u/tdffXV+OCDDzBkyBAMHjwY1dXVGD16NE6cOAFA/MxZVVWFuLg4pKamoqmpCUajEV9++SUAYOnSpfjzn/+MjIwMREdH46GHHsL999+Puro6zJkzB4MGDYJUKkVDQwOefvppl6/L19cXN998M1avXo1FixYBEOs+rlixAj/88AMMBgMkEkm3U2id9eCDD6K8vBxZWVlIT09HU1MTpkyZgssuuwwAPL4mR/cRAN5//33ceeedHl2HRDg9HO0DOp0OGo0GWq0WgYHu1af7Papvr8faE2sxMnykW1M8T/fC3hdQra/GyuErbaOZJzOs7q3eC6PViEnRk1wKIBs7GvGf3P9gUNAgzEuc51H/AKC8pRwvH3gZcokc94+7H7uqdmFQ0CC3R1/NVjMO1x1Gsia5U+Ico8WIcl05TIIJqUGpkEv753nK+qL12FS2CfEB8fjD6D/0Sx8A4JUDr6CspQzXZ1yP9FDP/pMkIqJ+JgjiFNDfi6YSoD4fSJkhZqm1p2I/4B8JaLpfTkOO1dbWIj8/H2q1GoMHD+60trCkpAQ1NTWdpkZarVbk5OTAaDRi2LBhqK2tRUVFBSZOPFXSpbW1FdnZ2fDz80N6ejpkslMP2dvb25GdnQ2tVothw4bZMqFqtVpkZ2cjICAAQ4cO7VQupLCwEM3NzRg7dqzD62lubsbQoUOxadMmDBsmZvBvamrCsWPHEBAQgOHDh+PQoUMICAjAoEGD7F6nveu2p66uDrm5uUhKSkJcXFyX13u6JmfO2d193LJlC2666SYcPXoUSqX7y4oYPPYzq2BFfXu9mMSmh//gN5ZsxMaSjYhUR+L+cfd7dM6X97+MitYK3DTiJgwKHoRVR1ahRFeCZanL8FneZ5BAgicmPWErA+GOT45/goqWCtww4gbbiJcrdEYd/n3g3/BX+OP2Ubd7HNRtLt2MdcXrkBCQgDtHn3ri8n72+/i+6Huo5WpclHYRlqQs8eg87tKb9Pil8hekh6QjLqDrfyS9oVhbDItg6bRu1FHwaBWs2FezD1F+UbbEQa74peIXGC1GzIyfyVqURETnusoDgKkDSHRvBpLL1j8qrjHNXCFOeSVyQU5ODnx9fZGcfI6MtJ8hNzcXCoUCKSkpHh2H01b72TeF32BH5Q7MS5zX44hdVlQWmjqaMCJshMfnvDXzVrSb220lHSpbK2GwGCCBBPEB8fBX+Huc1TW7PhtGqxGVrZVuBY+BPoH444Q/etSH08X4x0ApU3ZJEuMr84WP1AcyiQzBSvdKefxU8hN2V+/GZUMu65LAx1lqhdorI7bOaOpowqojq3Co7hBi/WPxwLgHbOsvbx55M1qMLd1OWz1Udwif5X0GtVxtm57synm/LhRrGw0KHuRW8ElE1O8sZsBqBhRO/p40tYtr7OQ+wOS7AZmHH70EATB39H0m1O3/BFqqgRkPA35OLG1obwZ2vir2NzAaCE7q7R4CsWOBin29U7qEznmeTksd6IYM8U7uFAaP/ezkwK9VsPbYTqPU4NIhl3rlnD4yn07r3m7LvA0NHQ0YGjIU46LGeeUcK0esRK2+FhmhGV45njM6zB149+i78JH54LqM6zqNVg4JGYK/TPlLl30uGXIJFiYvhFKmdDtgzmnMQbOhGSe0J+wGj+Ut5dhStgVTY6ciSZPk1jkc2VezD79U/IIlqUuQoun5iVKxrhjV+mq0GdsQ4hsCP4Wf7TWlTAmlqvupDJHqSLQZ2xCpinS5j1qDFsHKYMQFxPVJ7UsiIq+zWoAfnwDam4A5fwYCnPi/rK1OTJojkQAGnedlJ/a+I2ZEnXArEOed39kOWS1AfQFgMYoBpDPBozIAiBgmBs9+7pcFAAA0FgEKNRDw2++e3HVA6U5g3A2dg9IRF4tfRNRrGDz2s2WpyzApZhIi1T1/GLdYLV6rCXimcHW4V0pCnC5Zk4xkTd8O+zd2NKJYVwwJJGgztdlGVnsilUidateTy4ZchrymvG4D7+0V23G4/jAsgqVT8Kg36fHygZchk8hw15i7PCprsbd6L8pby3Gk7ojD4HFk2EgsS1mGaP9opAWnuXSeZkMz/Hz8UNteC5PV5FJG1Y9yPoLOqMPshNl449Ab0Bl0uHP0nR7ffyKiPiNYgQ4tYDaIiVecEZQAZN0EyHy8U69Q3ySO5ukbnWvfVAwoAz07t1QGzHhQPGfUcOf36SmTbEsNULoDSJ7Rc9+aSoDNfwVkSmDxP8UR39KdYubU6qPeGdG0mIHdr4tB8sTbxVFiIrKLwWM/k0llDkdhjjcex/vZ72N0xGivjT4ONCarCeuL1iPEN6RTrUhXxPjHYMXQFfCR+UCj1KBUV4oWYwsywnp39NNR8D0tdhosVkuXrLltpjY0dTRBIpGgVl8LCSRur3dcmroUB2sPOnXvZFKZ3fqSzkgMTERacBrCVGHdBo717fWQSqRdpiuPDB+J/KZ8xPjFoLqtGiarCVqDlsEjEZ09ZApg7lOAscW1oCVhouM2zpp4G9BcCkQ4McWuLg/Y+pw4Crj4BftlL0p3i3UgHR0vJEX88pYj/wUqD4rTW8dd3307Hz9A7gv4aoCTM4rGXg/UHAUGzenaXhCAX98G2mrFWpG+TvyOaW8S+wKI+2lc/F2c8y1QfQQYt/LU6CjROYrB41mgvr0eFsGCWn2t48ZeZrKacKj2EFKCUmzBgMVqwef5nwMALhp8kVdGRAubC/F1wdfQGrUYFDQIkX7u/ed7MhOt0WLEm4ffhMlqwm2Zt/X5KOjp4gLicNWwq7psD1eH45bMWyCTyPB+9vvQGXW4acRNGBw8uFM7vUmPj3I+gsZHnLpsL9FMjH8MYvxjeu0aTvJT+OHGETd2+3pzRzP+te9fAICHsx5GoM+ppFjLUpfZ/n7ryFuhN+uREJjQe50lIuoNfqHiV39R+gORw5xr66MWRzyVAfazrtblAb++JY4SLn2pb9dRRqSL008djeD6hQFL/gVIpGI/ASAkWfyyx2ICyveIo8TNZUCUE8Gjf7g4BVawuB44AsCJrWIAWnOUwSOd8xg8ngUmx0y2FV7va9vLt2N98fpOWUrr2+uxt2YvAGBG3Ay3A73TJWuSYREsUMlVON54vNMxjzUcw47KHViQuADxgV2TrOhNelTrq5EcmGwLrORSOeID4tFsaHYrYU9fOTnFNMAnAO3mdrvrLitbK1HQVICqtioE+gTivBT7dZQGAplUBoVUAYlEArmk+/9e7L2PRERnlbYGYMuzgDoMmPmImETn0CeAOhQYuri/eyfSxAFLXgCkCvvBY0AkEBgDqILFaaHe0NYgjhY6SiikUAMB0WLA5Uh3ZTfskfuII476BiDSyZlHbfUABCDezRHi8SuBulwg0b2ZU0RnEwaPZwGpRIphoU4+ZfSyuIA4qOSqTiN3kX6RWJy8GAKEHgPH443HcbD2IOYnzXcYwCllStw39j4cbzqOMZFjOr32S8UvyG/OR7Ay2G7Q8cGxD3BCewLLU5djcuxkAOI9uzXzVlcutV/dMeoOmKwmqORdn/qmBqViQvQErC9ej83lmzExZmKnWpWeEgQBbx5+E3X6Otw+6naEqtx/oh7gE4CHsx6GBBKPM/YSEQ1o+gZxyqVRL4521eeJI1AAkDJLHPVzhqFVDI7kXgreztTTaKKvBpjfNZmc2+oLxGmy/hHAgr/23DZ2rLhuMTix53buiB7Z8+uGVvG+nBzJ/PVtMamRvhEYtqznfQFxlDHnOyBmtLgGNCLduWnEROcABo9nqcaORrx39D1E+0djxdAV3bb7PO9zHKk/ghuG3+DWFMG04DQ8NfmpLtudWTO3oXgDylvLEegTiEUpixy2HxE+AiPCxVIkBU0FONZ4DLPjZ2NB0gIEVwd3e84AnwBIIIG/j7/d109qN7fji/wvEKWOwpxEO+sknLSpdBNajC1YnLLY4/qTJ8ml8m6PJZFIsHzQckgkEihlSgQpg7xyzpOsghXlLeUwWo1o7Gj0KHgEYDcAJiI654SnAVPuEgMwuQ8QPgRImSmOPDobOGrLgU3PiCN/8/9qf01iXyrdDVTuB0Zc6t7UXMEKQPjtTwfkSmDkJa6fw1WCABz8GDC0iGsrGwqBn/8FhKWJSYAA8b1sqXR+TeeJrcCJLUBDvvMJhMgjFquAX4saUdvSgYgAX2Qlh0Am7Z960f/5z3+g1Wpx880398rxH374YVx88cUYP358rxzfUwwez1KVrZWo1lejoaMBgiB0W3D9hPYE9GY9Klor+nx92ZyEOdhXsw/jo1z/x/914deo0ddAJVdhXuK8Hvt+xdArcMGgC6BW9PzLuqCpAIfqDuEwDmNWwixIJV1/Sec15WF90XrMiJ+BzPDMLq+3mdqwvng9AGB42HC36zq6SiaV4cLBF/basW8bdRuaDc1d1lsSEVEPok/7PSFTAGOudm1/i1Gc7mpqByA4t09dHqCJFaeGdsdsAI78T5y2muJCgrRjXwOtNeJ+6Uud3++k8DRgwbOAg4e5vc7UIQaIUrk4pbRwk7g9dTZg0ovBrUF3qv3wi8QvZ8VPAJqKxD+p160/WoWnvj2GKm2HbVu0xhdPLB2GhcOjPT7+Sy+9hNdffx0AcOedd+LOO+/stq3FYsETTzyBHTt2ONy3trYWjz/+OPbu3Yu4uDj87W9/w7Bhwxzud+mll+KRRx7Bxo0bPb623sDgsRftrNyJ3MZcnD/ofK9OMwSAjNAMXDT4IoSrwrsNHAHguozrUNZSZjcQ6m0ZYRluZzqdGjsVB2sPYmSYg6knEEfmHAWOgDiKOi12GqL8ouwGjgBwqO4QylvLsb9mv9175qfww7zEeWgxtpxTyV5i/WMR6x/r8n45DTn4NPdTTImZgvlJ890+v96kx4HaA9hUugmjI0djScoSt49FROQ0i0mcfupMvcbeEJICzHtaXP/nTPK5E1uA/R8CYYPFdZbdqT4iBkwSCZA83f56R3tGXARUHvBs7Z6/hzUdvUHfIE5DBcSRx9FXiSOPYYPFezH7T4BfDyXKTB3iMXa/DoSkds0GGxgNTL239/pPNuuPVuG2j/Z3ebRSre3AbR/tx+tXjfE4gLzyyiuxYMECPPPMM6ivr++5P+vXY+jQoQgJCXG479VXX42oqCi8/fbb2LlzJ5YsWYK8vDzI5fIe9xs7dizKy8tRWFiI1NS+GaRwBYPHXrSpdBO0Ri1S61MxLW6aV48tkUgwIdrxE6/eqOHY2zrMHchrykNsQKxXkvEAQK2+Fq8efBXhqnAsTe3+aerchLnwV/hjTMSYbtvMS5znlT6dC0pbStFubkeRtsij47x79F3srdkLuVSOAJ8AAOJazB9Lf4Rarna7fAsRUY92vALUZAPjbwQSJ/VPHwJdyJR9Mmuqo/ITEcPEUTFNnPOBIyCuQ4wd63x7dxnbxCBYHSwmqwkf2n3ZjfYm12tUamKBCbeKAblfKJA6q/PrPU1Pba0FfnxKHKEExLqaPZUSoV5jsQp46ttjdsfkBQASAE99ewzzhkV5NIU1LCwMYWFhCAoKcth2w4YNmDbt1Gf67vYVBAGbN29GUVERYmNjMWbMGLz99tvYsGEDFi1a5PCc06dPxw8//IDbb7/d7evqLQwee9H5g85HflM+xkY6/o+4WFuMX6t/hclqwoSoCRgUPKgPejgwlehKcKT+CAAxUFN6IQOc1qBFu7kdde11sApWSCVSbCjegKaOJlww+AL4yMSCwMG+wTgveeBmMx1I/nP8PzhcfxjTYqd1qWHpKj+FHyJUEciKzsKMOHGKVYmuBBtLNsJgMWBL2RYMDh58ztY5JSJySuxYYMmL4khlT3zUwITeWY/ltLo8YOe/gYRJwKgzcjPUZIvlNFprxZHK+jz7wePBNUDhZmDkZUCai7Nb4t1cL2ZqBywGcVpw+rLuS4JQr/u1qLHTVNUzCQCqtB34tagRk1L7pnxOXl4epk+f7rCdRCJBfHw8tmzZgiuvvBIlJSU4ceIECgsLnTpPcnIycnNzPe1ur2Dw2IuGhw3H8LDh0Bl1eO7X5xDgE4BbRt5ity7ilwVfYl/NPkglUjR2NOIPwX/o076arWbkN+UjSZPUawlPGtobsK18G8ZFjuuxVMOgoEGYnTAbwcpgrwSOADA4eDBuGnETNEoNpBIpDBYDfiz9EQCQGZGJoSFDXT7mCe0J5DXlYUbcjAGVJKZMV4aSlhJMiJ4AhdRxenOdUYechhyMDB/p0nWUtZTBZDUhQh3h8bTsa4Zdg3Zze6ekR7H+sRgVPgr17fUoaylDdkO2R+cgIupi0p1Ae2P/TVt1h7Kf1hOaDeL6zJ7WWp6uqVgcYaz/7QOwvhE48hkQNRKIHQMMmiuWOWmtFpMN2WNq/+3c7R5332nBicCcP4sBul+Y948vCOJ0abmP9499jqlt6T5wdKedN5hMJigUzpWOeeONN3DllVfi0Ucfhb+/PzIzM2EymZza18fHBwaDwZOu9hoGj32gqaMJDR0N0Bq0MFqNUEm7fkCfGD0RJosJPjIfTI9z/ETD234s+RGbyjYhIzQD12Zc67D91rKtaDO1YUHSArvBsD3byrdhZ9VO1LfX46aRN3XbTiaVYWHSQqf77qzTk8EoZUosS12Gpo4ml5LeWKwWrC1aC3+FP7aXb0deUx6MFiOWpTqR2tuOnZU7UaQtwrLUZZ0Cp4rWCoSpwuAj9cGuql3wV/jbMtE68lHOR2gyNEEmkWFE2AgcqT+CzPDMbteFfpn/JbIbslHdVo3zB53vdN+vG34dKlsrMTzM80xzMqmsS7ZchUyBK9KvgFWwYk/1HkT5nUUf7ojo7CD36dvA0ah3PhOrK2qygcP/AdIWAomTvX98iwnY8JgYDM76k5iIJsDBspLU2WKgG/rbTKqyX8WvhkJxivCZo5H2jL1OPI6zWVDtKdkJdGiBtAXOT+ENcjGnQfaXQNUhIOtmx9OQtz4nBtYzHvLsun4HIgKcK/nlbDtviI2NRU1NjVNt582bh6qqKlRUVCAqKgqZmZkYNMi5mYVVVVWIjx+YNbEZPPaBxMBEXJtxLfzkft2O7EyKmYRJMf203gJAqCoUEkgQpnL8lK3F2IK1RWsBAENDh9oK3TsyNnIsavW1/Xqdp3NnqmWJrgQ/V/wMAJBL5NCb9ahuq3a7D+uL16Pd3G6r5QgA+2r24T+5/0GyJhkLkhbgy4IvIYEETwQ94VRioOFhw3G88TiSNcn4pvAbHKw7iLKWsm6nfCZrknFCewKJga7V2gpThTn178VTUonUqfW9RDTAtVQDEhng/9s6fEEQE8D4RwKR/VPLuE/l/wgc+gRIXwJkXODdY1fsE2smlv3qWfCobxSnikYM6zyN1Gr5bTqnCdj7DtBUAoy7AUjqYT26TN65LwkTAV25OPLoLJkCCO3mAe/J/vgGdr+/oRXYs0r8e0hy9yOczqjLE0ci7a2/LNkh3rvaHMfBY1u92O8Orft9+Z3ISg5BtMYX1doOu+seJQCiNGLZjr4ybdo07N27Fzfd1P0gyOlkMhkSEhLw73//Gx0dHVi40LnBkT179uAvf/FiDVYvYvDYRzJC3cs66k2NHY043ngcYyPHdpkOOj5qPEZFjHJqmmOATwBmxM1Am6kN8QHOPxVJCEzALZm3ONXWZDVhd9VuJAYk9jjFta/FB8ZjQtQE+Pv4Iy4gDptLN2Nm/Ey3j7c0ZSmKdEUYEXZqVNFHKk5lUcqUiFJHIVmTjECfQKenlC5NXWpLCjQoaBDymvJ6HF2dHjcdU2KmYFfVLhRri5GkSXL7egDgcN1hBPoEenwcIjqHtNYBG/8sjlgt/Lv4gb/qEHDgI3Hb8tecyzbqil2vA81lwLT73J9+aLWIgVlw8qmg1136hs5/eoPZIAYsaQsBZaAYoHmi8gBQeRBoPNE5eFT4illhzQbgyH/Fbc7UcjydKkhMTOQNZgPww5/EAHLeU0BzqfgeByd1bufjJwawHVpA48FniapDwC8vi9ew+J9dX8+6WVy3meTEQ+mZjwL6es8C2d8JmVSCJ5YOw20f7YcEnYvZnBxDfmLpMI/rPX777bd48MEHUVNTA5lMhk8//RSrV6/GpEldBzsuvvhiPPvss7BYLJDJZD3ue7IcR0NDAwICAvDpp5/Cx8fH4Tlra2tRVVWFqVM9yyfRWySCIDhZWMg7dDodNBoNtFotAgN7eFp0Fqhuq8ahukOYFDMJgT4D/1rePPQmCrWFmB0/GwuTvT8t1Jt2Ve3CF/lfIEgZhD9O+GN/d8dlB2sP4tvCbzE/ab5bo2ZagxZ+Cj/IpX3zfOfkaKdarsaTk590+zgnmk/gjcNvQC6R48nJT9oSETnrYO1B+Mp93VqDSkQDTF2eOJ0vbb74oX7D44DMB5j/jDh1U98oZjsNjAGynHuK7zRBAL6+QwwyJt8JxIx27ziFm4ADa8TsnfOedr8/rXVA9WFAFQxEZgBy76znx/4PxdHbpKneyQhqaBXfs4h0IG6c/TYWkzh6Fuh5fT23mTqAdQ+K7+/IS4GDnwByX+D8f7uWXdZZjUXidNOgRGDWo945ZmudGPD2Rn/PMb1d51Gn06GysrLTtoSEBKjV9md7PfLIIxg+fDiuuuqqHvetr69HfX09VCoVEhISOpXW62m/xx57DKmpqbj++oGZ5Zcjjx74pvAbFDQXoMPc0WWtmCAIqG+vR5gqrMc6jH1pcPBg1LXXIVnjOHNYTkMOylvLMTN+pt3RSIPF0G0yG6PF6HLQcKYw3zCEKEPcrhN5kiAIKG0pRbRftMd9ckVBcwFaTC3Ib8p3K3jUKB2kYfeyhIAE2yinJ0JVoQj1DUWIb4hTo9inK9OV4ePjH0MCCR6f+HiXNZBEdJYp+UUcjZErgamjgUXPAxLpqcBJHQLMfaJ3zi2RAFPvEwveR49y/ziBcWJQ4unatD2rxLqD6UvFZDHe4v/bukNvrdtU+gNjru65jUzRNXCsOQYc+woYcp77gborFL7AvL8AFqP4oEC1ThxZ7K3PWyHJwNKXxIcfzjK0Ajnfig8Los+Yqnt8LXD0C3F0d9QV3u3rOWjh8GjMGxaFX4saUdvSgYgAcaqqpyOOJwUGBro0oPXkk0+iqanJ4b4ny3G4es4bb7wRCQkDt5Y4g0cPjI4YDb1JbzdhyA/FP2BT2SbMjJuJRSmLXD52sbYYHZYOr47AzE6YjdkJs51q+/Hxj2GwGBDiG9Kl1Mim0k1YX7wei5MXY0b8jE6v7a7ajc/zP+/xukt0JThYexAz42faDZIqWyvxztF3oFFqsDhlcafXBEHA3pq9CFYGO1XOZFv5NqwtWut0IiBvWZS8CNF+0V5JJuNNJqsJ28u3I9Y/FkNCTk2ZCVeH475x93l8fI1Sg4ezHna6vSAItocroapQxPnHQa1QD6jstUTkpqGLxQ/byb/VRFPY+blubxYL2seNFwMCbwobJH55IjwNWP6q532JHCYGsqEe9udMafPFGoYy1x7W2TSeAHRV4tROTwKv0p1iIpzin/smeATEKaQn2ZtK6m2ujhYX/wwU/AhUHewaPFqMnf8kh2RSSZ+V43DE19cX0dG9N/KelJTUa8f2BgaPHhgfNR7jo+zXEbIIlk5/ukJv0uPNw2/CIlhw1+i7EBcQ51E/3TElZgqKdcVI1XRdK9fY0djpT2dfO2ntibUo1hVDKpHa1uadzmK1wCpYYbaaIQjCqYntAPKb8/FZ3meQSWR4esrTDke3TiaYcScYqWqtwq6qXZgcMxmRfg6yytk5b18Xtm81tmJb+TYMDxuOhED7T6wO1x3G+uL1UMqU+MsU9xZinx7wObK7ajd+Kv0Jy1KXdQmkD9QewH9y/4NZ8bOwIGkB1Ao17hpzl1t9IqIBQhBOBSH+EcDoK3tuv/99oOowoKsEMi/r/f71l2Hni1+OtDcD+1YDIanAMCezeLsbOALA9hd+SzxjFNcHxozuerzmUnEtZOqc7suEpC8V90+aZv91R6xWYMdL4trEqff1nASnv9QeBwo2ig9FnBmJjh0D1B4DYkaJU7RVwad+NoYtF++1J+swifoJg0c7chtz8VPpT5iTMKfT6IwrFiUvwuiI0W6VF/CR+SDGPwZ6kx5ByqBOr1W3VUMhVSBU1btPX7pbEykIApaliIGAvSQs8xLnIVmTjKTApG6PPSlmEqQSKUZH2H86GR8YjwfHPwhfmW+XMiCR6khE+0Uj1DfUqWmR46PGY0jIEAQoAnC0/ij+l/c/zIqf1WXE1J4NJRuwr2Yfvir4Cucln4fLh17ucB9v2F+zH8G+wS5PIf254mdsKd+CguaCboOwFE0KEgMTXc6setLR+qNYk7MGU2OndhkVtud443E0G5qR25jbJXisbquGVbCiqrXKrb4QUR+zWsT1bt2NEOauE6fijbkaSHay5FRYmrg2kiULRDXZQPVRoD7f+eDRE1EjgcZCoHi7mEF16GJg+IWd2+z/UByhtJiAERfbP45/BJDpwe9Ic7s49VWwAq213g0eBUFcxyn3BYa6PhPMJv8H8UGH3BfIcuLfq3+EmKwpdx3w/YNiqZCRv2U9l0i6JvchOku4FDyWlpbi0UcfRX5+PrKysvDYY48hKupUcPTmm2+irq4Ojz32mNc72pf21uxFsa4Ye2v2uh08SiQSRPtFo93c7lR5hdPJpXL8YfQfumyv09fhxf0vQi6R448T/ujycT2lN+nx4v4XIZVIcc+Ye+wGb3Kp3OFU29ERo7sNHE/qrgSERqnBvWPvxfHG49hTvafbkd/TnUxmVKQtgt6sR6G20KngcWL0RJzQnoDJakJBc4HD9idZrBana1+eKbcxFx/lfASFVIG/Tv2rSwlzhocNR0FzAcZHjcfbh99GRWsFbsu8rdOoabBvMO4YdYdbfTNbzShoLoBFsKCitcKpfZalLkOyJhljIrqu8ZmbOBex/rEu1dk8kyujoETkoc1/FUcIZz5i/4Nvc6n44b+5zPljDjlP/PI2i0nM4nq2/f8QNw5oqRKzu/aFCTeLfx79HNCWAxo7M50SJgJWk2vTUZuKAUiAYCceVFpM4qjllLsBQ0vPU40tJiBvvdhPZ/vTVCSuMQTExEdSubgO0VmttUDhZiAuC1CoxSDQFca2zn8SneWc/mRqMpkwa9YsBAYGYtKkSfjuu+/wv//9D+vXr8eoUaMAAC0tLWhubu6lrvadeYnzEOATgInRnqW8/jT3UxysPYgr0q9AZnimx/3ykflAKVPCR+rTZ1k4T9dubofWoIVUIoXBYoCvvO+Ksp5Ob9Jj9dHVECAgTBXm9Ajd/KT5iFBHOL2OdEjIEDw56UkcrDuIaD/n5rZ/d+I7bC/fjsuHXu4wQLZHLVfjhPYE/OR+MFqMLr3PcQFxuHP0nRAEAd8XfY92czsaOxpdnnLbnfey30NeYx4mRk/E3KS5Tu0T7BuM6XGnRiAO1B6ASq7C0JChUEgVGBnuQr2vM2wp24J1ReuwfNDyAVM7lOic1t4sfng3tNp/fdSVACTi64bW7qc49ra6PGD7P8V1hlPutt9m77viVMwp93RfR7A/yJXdj+45q6lYLC3R0zTTMw2/CMi40H6wPWgOkDBJzJDrjLZ6YNNfxWMt/Lv9uogn7XtfHPWceBsQO7b7didV7AeyvxKn1l7whnP90SSI624lMmDn62IgPPNR59fD5nwDlOwUp59O7vpg3y59I3BiM5A4RbyvMaPFTK1E5wCnP5lu2LABKpUKe/bsgVwuR3t7O2655RbMmTMHGzduxJgxXswe1kd+rvgZ64rWYVnqsk4ZMSPUEViWKk4XMVqMWHVkFSQSCVYOX+lSxs4WYwsECGg1nvpFW99ejyJtkdM1FU+nUWrwxwl/hBRSKOyscdAatJBIJC6XDdGb9E6NYoaqQnHnqDsBSd9nAz2dr9wXw0KHQWfUobGjER8c+wAz42Y6HE1UypROZz5tNbaisaMRCYEJLgWBtfpaCBBQp69zep/TBfgEICM0A1KJFFZXa2j9RiKR4LbM29DY0Yj00HS3jmGP0WIEJEByULJbpWnKdGX45PgnXsumWtVWBQECavQ1Hh2HiJw0+zGgvan7YEvpL5ajMLUDEUOBFMczPHqFvgGwmsURo+7UFwBGvTja1hvBo6EVqNgLxIzp+/V7Bz4SS0tYLV2noPaku1Haws3iMQfPd25dqtwX8NX8llnXwUPm1lpxWmlbvXN9DE8Dwoe6NuXTYhRrYIYPFQNrfYO4/tBZCZOAlmoxEHTWsa/FhDnaCmDKXaf+jRnbgLY6Tlmls5rTwWNJSQmmT58OuVzcRaVS4f3338fdd9+NuXPn4scff+y1TvaWspYymKwmlLeUdxtUaA1aFOuKAQA6o67b6ZT2XD3salS1VSE58NTI2Mc5H6O8tRytxlbMSpjlcp+7K4/R3NGM5/c+D5lEhkeyHnF6SutPJT/hh5IfsDBpoVOZWOMDnV/cbRWskEqksFgtWFu0Fj4yHyxM8ry+pFQitWVOXXtiLdpMbchvzndqKqqz3jz8Jmr0Nbg6/WqMCB/h9H6XD7kcRboiDAnuebpzTVsNNEpNl9HbIN8gPDDuAUglUreDK51RB7VC7dZ6257cMPwGNHY0IsY/xq39vZ1Ndfmg5cgIzWBNSKK+og7peRQJEEfN6vO9W47CVYmTxIAtoIcZI1PuEgOJuKze6cORz8TgoS7v1NTQk5pLgd1viqNRno4y2pMwSQyeoz2f8QRAfGAAAO3dJ8LrROkPnPccAAkglfbcduKtYqAb5eTvWVUwMONBsc6jqcO5DL1F28Rpq2W/Auf93bnznC4yw7VproA4/bi5RJzye7rt/xTXlk64BYjvpX975yqrBSjZIWYt9o8UMwS7uUTIU3v27EFpaSkuuuiiXjn+v//9byxfvhxxcX2fMNMZTgePSUlJWLt2badtEokEL7/8MmQyGebOnYvzzjuvV1PXetv5qecjLTgNw0KHddsmXB2Oq9OvhkQicSlwBMQMnymazouqU4NS0WRo6jYbZnfaTG1QypTdTmOUSCSQQAKpxMF/1GfQGrXinwatS/s5Uqevw78P/hthqjCcn3o+fq74GQCQFZWFEF8HHz5cMDdxLkJVoRga7J0AIq8pDyW6EqhkKsgkMpfXlaoVamSE9vxLJrshG+9nv484/zi7iW2CfV14InoGrUGL5/c8D6lEiofGP+TVWom+cl+XAsfylnIcrT+KqbFT4e/j71I2Va1Bi+0V2zEqfFS32YZVcpVH016JyAM534k17MZeJwZrJ6XMFL/62+kf9tvqxayX8RMB+W+zhwKivFcb0Z7wIUDlfnGk7Ex1eeJIVsW+zsGjxQTsfkNcNzrhtlN9ddWgOeKXtww7X7yeEBdGaJ39UK8M6FrGwhFDK7DhT+KI5fxnHI/sRo8Upyj3VQkRQAyG7QXECrU4wttPy37OWse+AdY/LK65PikwBlj4nFcSS1ksFqxZswb79+9HWloarrvuOqjV3X/+u+eee7B69WoAwFdffYXvvvsOALBs2TIsW3aqP1arFWvWrMHevXsRFxeHW2+9FQEBAQ5fS0xMxJ///Ge8++67Hl9bb3A60pg1axZycnJQW9t1Gsi//vUvXH/99fj444+92rneplaoMTZyrMNRkBHhI7xWr29h0kL8ccIfXUoScqL5BJ7Z9QzeOvxWt21O1td7aPxDdgOePdV7sPbEWpispk7bl6YuxY0jbsSS1CXOX4QTdEYd2s3tqNXXIkodhamxUzE7YbZHgaO9AFcpU2Ji9EQE+QZ50NtTPj3+KTaWbMSoiFF4fOLjHiVz6Y5cIocEki4PAr7M/xJ///XvqG6r9uj4EonE4yQyrcZWnNCe8OgYXxd+jU1lm7CtfJvL+24t34pt5dvw7YlvPeoDEXmB1QKU7hKnGNbnA2ajOGpnNYujaAPdr2+La+vy1vXdORMnA8tesR9Ip8wQM5NOuKXz9vYmoPKgmNFT7+Q0ztPVHAP2rOp5uq47pDIxGPd2HU53WU2A2SBOR7WaHbcPjAFmPQoM8Xzmk8em3ivWpHQ1YP49O/YN8N9rOgeOgFif9L/XiK97aPHixdi0aROSk5Px3//+FwsWdJ8Uac+ePZBKpUhLEx8MxcfHY+LEiThx4gT279/fqe0dd9yBF198EampqThw4ACWLFni1GuLFi3CTz/9hIaGBo+vrTc4PfKoUqnw8ssvo6ioCBEREV1e/+c//4nExEQkJLg2onau+rbwW+yr2Yerh11tC0BMVhNe2PsC2s3tuGfMPU4HPO2WdlgEC9pMPWfqCvAJsLvdYrXgf3n/gwABSYFJyAg79URWIVUgLdjOk1EPpQal4qYRNyHQJxA+ch/bGlJ3rS9aj01lm7AgcQHmJIpPVA0WA+QSuduZTe2ZGD0RBc0FSAtJg1wqR5muzKWpus4YEjIEj2Y9Cj8fv06ZWY81HIPWqEWprtTtKacapQYPjX8IEkg8GnV8+8jbqGqrwmVDLsPYSCeSGNgxLnIcLFaLWw9eMsMzUdZShglRzq1RJaJeVPAjcPi/4gd2uVKcijf2WnFqnjdHc4p/Fkt9jLy063Q/T0QMFTOYujJy1ptkCmDwvK7b/SOAcTeII4+BbiwPOPY10FAAKPyAUSs6v1a+DyjdISbGcefY3lbwk5gFNXOFmGm1Jw2F4ijt4PmAKkicujrvafE1R9OoBxqpTFwPSs6xWsQRRwh2XhQASID1j4glZjz4LPjmm28iMVFMaHTjjTciMDAQtbW1duOdb775BnPnnkoaOHbsWIwdOxYHDx7s3DtBwIcffmgbzbzrrruQlpaGXbt2YcKECd2+NnHiRMhkMkyZMgXr1q3DVVdd5fZ19RaXUnaeHhXbc9ddLPB90smyEOUt5bbg0Ww1o8XYApPVhA5Lh9PHGhYyDHMS5nSZAussmVSG2QmzUauvRUpQ79fSMllMgAQYHDzYa8fUm/Wd/qxorcBrB19DhDoCd4/pJpueG+Ynzcd8zAcArD66GjmNOViWugxTY6d67RyAuLZxTc4aZNdnY+WIlUgNSsW1GdeivLUcGaEZyG3MxaCgQW4Fxt09RHCpf8og1Opr3UqMc9KE6AndriV2VM4kMTDR7ZIiRORlmnhxmp1/pLjeyMdPnG7oTIBn6hCDlsgRgH94z21rssUi8TXZ3g0eMy4Qv84GSS4kZTnTkPPEzKX2amzmrhXX2vlHASMvcf8czmqtBQ58KCYMSrWT3+HIZ+I03fB08ZpPbBEzmo6+Egg6YxDi0CfiukiJ9NQ0X/+uH+o9pm8Ug++YMYCMZdAHhJIdXUccOxEAXYXYLnma26c5GTgCwPHjxxEREYGwMPtL1Q4fPowrrrjC4TElEgmUSiUaG8W1wgaDAa2trTh06BAmTpzY42sAkJ6ejoMHD579weO5xGK14Iv8L7C7ejdmxc/qUvB8X80+rC9aj8UpizEqYpTLx78y/UpbVtWTVHIV7h5zNwwWg0sjS4frD+On0p+wU74TT05+0uW+AMCCpK5D8NkN2WgxtnhckuR0OqMOL+x9AQqpAg+Mf6DbBD+uWpa6DGMixiA+QBwF1Jv0MFlN0Bl0vVbr72QyG29dw5lq9DUwC2Y0tDcgNSgVcQFxiAuIw/vZ7yO7IRvzEudhXqKdp9N94LqM63qtHMvOyp34suDLTqPIRDSARQ4Dlr8q/r292bWRk9y1wPHvxTVzMx7que3Iy8TRQUeJRHLXix/uBzlXMsjrao6JZRjSlwFBXpyZYrWK004NOmDi7T2Xxjg5NfX0ICpmlPhlT8aFQNlu+4Fcb6jYJ96ntnr758xcATSeODVyfWKLWB+0Yl/X4DF5hhg4OlPKwxO73xSDx+EXiiNZztJWiOtn+yl5yzmt1cls6s62c6CmpgbXXHMN3n33XUi7SfbU0tICPz8Ho+W/eeyxx7B06VLMmzcP+fn5CAkJsZU07Ok1APD390dZmQs1c/vQ7zZ4PFJ/BGuL1qK6rRpx/l0TchxvPA6tUYu8pjy3gsdQVShCVaFdtoerHTx5tSNSHYlAn0AkBSa5vG939CY9Psj+AAIERKgj3B7VPJPRYkSHpQMmqwkmi8lrgZdcKkeSJsn2/eDgwbhz1J0I9AlEdkM2tpRtwYKkBW6PdmoNWqwvWo/00HRbEpbLhlyGRcmLeq0sycrhK1HZWtklW+jJhDlByqBeOa8zJBJJr9XxrGsXy5g4U2LjSN0RCBCYGIdooFAFudY+bAjg+4u4Zq69WUxcEp9lf6qiKggY7CAgbC4VR60AICrT8Whmdyxm4OAasR+uZjzNXQvUHgd8/MXpu95ibBUDPEAcbemuDmF7E7Dxz+LfFzzr3NTNqOHiV19Jmia+391lUU2Z0bmUy6irxH8bqXYeKCZP82hUyWkhKWKGVI0LGS5z1wFH/ieucR2/svf69nvl72Sdamfb9aC8vBznnXcennnmGZx33nndtgsLC0NTU5NTx7z33nsxf/58HDp0CGPHjsX1119vG+Xs6TUAaGpqQni4m/+/9bLfbfCYEJiAzPBMjAgbgSvTr+zy+pKUJYgPiHer0Lu3RflF4bGJj3n1mCq5CiPCRkBn1CFK7b2Mc2GqMNw9+m7IpfJOa+4sVgskEtezwfbkZMbaLwu+RGlLKQ7UHnA7eNxfsx/7avehWFdsC1SkEqlTgaMgCMhpzEGsf6xLgaZGqbHbflnqMixIWtBrI579bV7iPKQEpmBQcM8Fmuvb6/FhzocAgAf9HnTrwQsRuam+QBxJCUl23LYnUcOBJS+If9/5qljkXVsOjLnaveMFRIvBp1Tm2Xq3pmKxhAMApM527VhDl4p1A70x8mnqENeRSiRi1tAJtwCGlp4L2EvlwMma03ZqPjtU/AtQulNM2uNKoNQdQ6uYuOb0BwtK/67rLnsSNqjna+4LmZc5V8eyk99mPfXC7CeCGJQHxojJceyue5SIrydO9ug0xcXFWLhwIf7+979j+fLlPbYdP348jhw54vSxMzIykJGRgR07diA7O7vTesmeXjt48CBuuukml6+lL0gEQbD3bnQrJycHubm5SEpKwqhRo1w+oU6ng0ajgVarRWBgHxfOPUs5Wh92ugO1B3C88Xivjpi5qtXYihf3vwi5VI57x97rclDUamxFRWsF0oLT7E5PrW6rxp7qPZgcM9nuaO/pchpysKV8CxYkLui0/rO5oxlri9ZiWOgwlx8Y7KjYga8Kv+pUeqPF2II91XuQGZ7psE+/J/Xt9XjlwCvQKDW4Z8w9PT5MMFlMePfou7DCipXDV8JH5mbaeiJyja4K2Pi4OFXwvH+4PtrYnYKfxKQuo64EEvo5GZbVAhz9XBx5PHOKYks1sOs1IGKYGGD1xGwQk/xo4lwfHavYL54nfgKQ5eKHRFMHAAFQuFEz98cnxSmiQxeLUzQ9YeoQE5qYjWISG3dHgk9n1IvTV8PSxEDZz8u/QwVBnJ6qiet6/4xt4nviGwSMv9G5oFBXKY58nf45zdQBHP6POK051XENberByWyrADoHkL+9N5d+4HG5juHDh8NqtWLy5FNB6NNPP42YmK7JpYqLi7FkyRIcPXoUALBr1y6sWrUKO3fuhFKpxJgxY3D//fcjPT3dVsajsbERmzZtslWnANDja3q9HkOHDkVeXh58fQdIluPTuDTy+MADD+Cf//yn7furrroKH374odc7dTaqb69Hq7G109RKb6huq8arB19FjF8Mbht1m8P264vWo8nQhBi/GMyIn+GwfV/osHSgxdgCmUQGg8XgcvD4/rH3cazhGFICU3DjyBu7BMVRflFYmrrUqWPtqtqFIm0R9tbs7RQ8BvkG2R2BdkaoKhQyiQyR6lPTJjYUb8Du6t0o1ZXiuuHXuXXcc1GbqQ3t5nZYBSssgqXH4FEhU+CWzFu6fZ2IeomPn/jhWa50Lzg5088viuvbZj7i3fqDnpDKxKyu9tTni+vYOnSOg8eqQ2I2WokUSJrq2ghUe5MYyOgbnd/nJE/KZoy8XAzOnA1qBEGcUqqJs5OoRhCzw0KA/ZEhN2R/ARRuBloqgcA4scxGSA9LawQB2PEy0FYHTL3P8ShywY/AoU/FqdTT7uv8WnOZOCVZIgFGXeE4EyxgP3tt9WFxZFsiBVJmcWTSE8OWiQGi3TqPf/dKncc//vGP0Ov1nbapVPb/70tKSkJGRga2bt2KGTNmIDw8HBMnTrQlugEAjUb8nHqyjIdKpcILL7yApKQkW5ueXvvoo49w3XXXDcjAEXAheDx69CheeeUVvPfee5g2bRr27duHG264AT/99BPmzBkgvwxclNeUh4KmAsxKmOWw1mNPzFYzXjnwCtrN7bgt8zYkazyc5nOaZkMzDBaDU+vDAGBh8kIcbzzu1jpNAGg3t2N7+XakBad5LRAOU4XhztF3Qi6Ru5W9M8Q3BFVtVbBYLdhavtWjsh8LkhYgSBmEaXHeWz8xJGQInpnyTKfR4WGhw1CoLfRafdDe1mpsxc6qnRgRNsLtMiHOSAxMxG2Zt8Ff4Q+F1I3pVkTU+3wDgUXPd/3Au/tNMWPn5D8AgdHOH6/xhDii01Ldf2Ui2hrEgNGZUdSEiYCp3bkpu5EZQNx4cYTJUYBgMYnHPVnUPnW2eD80TiTdKdkhjnJ6I+FNxFDxy1mlu8REPn5hwHnPdX5NoQLmPyOuIXVnhLBDJ/7bOP3fU0iKeL1KDcTgVBCP31ptf5qtxSRm6LVaxH9jjoLHk0tqlAHiGlq/iFPBePgQMZmPKsi5wLE7kcOBxEmAJoGBozcMWyaOlJfsEJPj+EeKU1W9lKTImeypp3vhhRdQUlICAEhNTUVqqv0yQCfLeLj6WlxcnMt96ktOT1t977338MMPP+CTTz6xbbv77rsRExODhx9+2OkTDqRpq3/b/Tc0GZqwJGUJJsVMwv6a/UjWJCNC7TgFtMFiwJf5XyLYNxjzE+fjxf0voqmjCX8Y/Qevr83Ka8pDiG8IwlT20wZ70+bSzVhXvA4Rqgg8MP6BXj+fMwRBwN6avdhbsxdLkpd4ve5iT/QmPVRyVa9kcz3T0fqjqGitwOyE2X0eWH134jtsK9+GZE0ybst0PMJNRL9DX90ujpIJViB+PDBupTiy4igg01aIH/x7O1tmd9rqgQ2PiVMgz3uuc1Bg6vBsFM8Vm54Rg5VpDwDhLtRXbmsA1v2WqXbO40BwUs/td7xy6jwBPSQSKf4FOPixWMakp0RFjSeAbf8nBsqTvFhCyWoFvn9AzC4781Eg9IwP4MY2Mdj2CwN+fVsMYkdeCqTZKeBelyv+20yc5Ny5O3Ti6ODe1c5lAiYiG6dHHhsaGhAbG9tpW3x8PGpqvJMetz9MiZ2C7IZspIekY2flTnx34jtE+UXhvrH3Ody3SFuE/bX7AQAz42finjH3wCyYe+VDf1qwC79kPJQemo6jDUcxMmzgZLeUSCQYHzUe46PGw2K1oERXgjj/OLdqILricN1hrMlZg8zwTFyR3vtPgD45/glMVhMi1BFeT9TUbm7Hdye+Q0JAgt3ai8NCh4mZhcNHeXSeWn0tCpoKMC5qHNcoEp1rpt0PFG0Fin4GGoqAH/4ESKXAwufEBCnd0cSKX97WeALwCxdHkHokEYNciRS2dVKAGDgVbgLG3eBxwg2nGPXi6JjZ+TrPAMTgPHYsYDGICYN6IgjitEtzhzjts6fgsaFAbNdYCKCH4NEvAjj/Ve+PoEkk4tRoo0wM7M/k43cq0D957u6WOoQPce3cvoFi7VJA7AMROc3p4FEQBDQ3N6OgoMC2raGhAU1NTZ22BQcHIzR04CcI2Vu9F0aLEbeMvAVSiRRtpjYE+gRiSLBz/wENChqE6XHTEaIMsa3hU0i6DxxLdaXYU70HM+JndBpB7K0ahQCwrmgdKlorcNmQy5wuHB/lF4U/jP5Dp21F2iKYrWa3M5l60/dF32N7xXZMip6ECwZ7r+izIAjYWbUTQcogDAsdBkCcyilAQIuxxWvn6cm0uGko05UhNcj+9AdPZNdnY0/1HhysPWg3eEzRpDj10MSRT49/ivLWcrSb21nDkehcE5oqfsVPAGRK4OcXuv8w39vK94mJTTSxYqKWnviFAgv/BkhknWsnttX9tuawoXf6aLV0nlY361FxneOZdQwdkcqASbc711YiAWY8KNaBjB7Vc9sRl4hTRKMzu2+T/yNw6BNgyHmulzRxRCIB5j4pTsf1dTATbez1wNAlYj1Fb4kbJ07PdqVuKRG5ljDnnXfewTvvvGN3+0n3338//u///s/znvUig8WA/+b9FwAQHxCPISFDkKRJcqkchlwqx5KUJU63X1+8HgXNBZBKpLagp769Hv8+8G+Eq8Nxxyjnp4IYLUa0mdps9QC7s718O8yCGYXNhS6tgaxuq8b3Rd9jTMQYpASl4M3Db8IqWHH/2PsR6ed5LR1PnFybeuYa1YKmAhisBmSEZrh13PzmfHxV8BWkEimenvw0fGQ+mBQzCVH+UYj2c2FtjwcWJi3stWMPCx2GMRFjkBDg4ocWFw0NGYoWY4tX1/0S0QAT+dv/sycDstNHHQs3i0FZxgXulZHoSWutmK01YaIYBEqkgJMPRu0GCONvEkcvI4b1vK8giKUt/CO7Tq3szpa/i8lXZj16ap2eMsCJUVIvCE5yPLUVEO/hyQyxVguw/QXA2AJMf+jUe9rRLP7Z7lxdO5fJlc6N/Ell3g0cT/Kk3AvR75TTweOKFSs6ZRLqTlycF2oG9TKlTIkpMVPQ1NFkqxXY26bEToFUIsW4yHG2bU0dTdCb9ahuq3apHMdrB19DVVsVbh55c4+jVCuGrkBlWyUywlwLqPbX7MfxxuNoNbZiWOgwRKgjYLKYnB69PF1zRzM2lW1CZnimV0bU5ibOxfio8Z0yrrYaW7Hq6CpYBSvuHnM3Yv1dnx4V4xeDxMBEBCuDbVOPJRIJUjQ9ZHhzk86ow87KncgMz+zV5DSnUyvUuHyog6yBXjA/aT7mJ83v9fMQ0QBwZkBmNgIHPhL/Hj6k5xEtdxRuBkp2inUi5z0FLPmXZ9lgfdRiHUpHqg4Ce94Rgxxnp2/qKsUpofpG79RS7G0mPVCfJ65nbas7FTxmXCA+LOgp26kzLCbxHBzlIzrrOR08xsbGdlnzeDY7f9D5fXq+jNCMLqNig4MHY+XwlQhSBrm0fs8iWDr92Z0R4SMwInyEy32dFDMJerMemeGZ8JH5eDSd8ZfKX7CrahcqWiu6TId115mlOnzlvkgISEC7uR3Byp5HY7vj7+Pv0uivJzaVbsKOyh0o1ZXippEDswAsEZHL5D5A+hKxfIXZ4P3jJ08Hjn8H1BwVp63GeSkBj6FFTGYjVwKzHhOv43SBv5Wo0DiRUfWk2Y+JiXoiHYxq9sRsBH59S+zXuBu8llnSLmUAMPXerllmpTIgIr1r+7o8oPqIeE/ixzu+L5ufFYP+GQ8BYV5eAmM2ALvfENcwjr+xd+8TEbk2bbUndXV1+OCDD6BWq3HbbczW6KwhIS4u8gZw+6jb0Wps9XpW15OCfYNxcZp31jaMiRiDqraqTiOu3iaXynH7KCfXgwwAw0OH44T2hNvlVIiIBqzkGcDxtWJZD01c59IcZb8CFfvFtXbulHUIjAaiRgJ1x4H232ojCoLniVw6dGKgJ5UB5vauwaN/uDhF1xX+EXZqIrqopVKsrwgAGRe6ds/am4CyPUDCBOdH+1wJdH95Ucx+GpwoTnP1UYvTerubWmruEEc1e+OhQksVUHVY/PuISzgVlXpdeXk5SktLMXly7yTa+v777zF9+nT4+/eQiKwfeRQ8Wq1WbNiwAatWrcI333yDsLAw/OMf//BW36gbKrnKo7qUfSnaPxo3jrixv7thU9laiSBlENQKtePGvWRQ8CCvJKchIhpwfPyB4GTAYgR8gzq/lv2luG5REyeOULpj4m1iGYqIdODEFuDAGmD4hWJCF3dpYoHpDwAyH+9MqzwzUY67gpPEmoNyH9eD7cP/FYN1bak4GudtsWMBbRmgDhMTDu1ZJd67JS/Ybz/rj+IU3pqjYiAZ58ID5fYm4OcXgaZiYMgiIPOyzq8HJwGjrhBHHhk4nrMsVgv21+5Hnb4O4epwjIkY49Ws++Xl5fjll18QERGBGTNmQCrtPhnY/fffjzvuODVbLTc3FwcOHEBaWhrGjBnTqa3VasWWLVtQXFwMq9WKa665Bj4+Pti3bx8OHBAfDo0ZM6bTfmVlZXj++efx1FNPee36vMmt4LGkpATvvvsuVq9ejfLycsyYMQPbtm3DhAkT+qQeXn+p1dfi+xPfY3TkaGSGe3kth4fq2+txtP4osqKy+jUw6klFawUCfQLdWjvpDTkNOVidvRrhqnA8OP7BfunDuS6nIQf7a/djfuL8XhsZJ6IBTO4DzP6T/deGXySOpHlSFkPpf2qErLlMHM3Slrt/vJPsTc1sKhZrIQ6eL44+OmPX6+K03an3OVfLsUMrBnlx4+3Xy+yp/mJPIoeLaxgjnVjT6Y7xK8UvQAzm89b1nEVWGQA0FAJHvxBHiqPfAGROfgRtKgFqsoH6XDEoT1/aOWsuAAxidu9z2Y8lP+Lvv/4dNfpT5QEj1ZF4JOsRzE1082fkNO+++y6ef/55jB49GkePHoVarcb27duhUHRN+lVcXIzs7GxMnz4dAHDvvffihx9+QGZmJrZu3Yrly5fjtddeAwC0tbVh/vz5aGxsxIQJEyCXy3HllVcCEAPEXbt2YefOnbjkkks6BY9XXXUV0tPT8eijj8LXt4/q0LrA6eDRaDTiq6++wqpVq/DTTz9hypQp+Nvf/obi4mI0NTU5lUznbJTflI/C5kLMjJ+JfTX7cKzxGJoNzV4NHndX7ca28m1YPmi52+Uwvsz/EvnN+WgxtmBp6lKv9c1b8pvy8faRtxGkDMIfJ/yxX/rgI/OBVCJ1a9RWb9JDLpW7VbvQZDFhf+1+pAaldirTci76sfRHlLWUIdAncED+OySifhQ3zrURJ0dGXiom5nEmQHJneuvRL8SgxWoGxl7r3D66CjE5TFutc8Hjkc/EJECNRcCEm13rX0+SpohffSEoAVj2b8f3NzRVDNIDop0PHAEx8VLWzUBtDhCT2TVwdFfhZvHBw4hLAMXA+4BOoh9LfsR9W+6DAKHT9lp9Le7bch9emPmCxwHkkCFDcPToUchkMpjNZoSFhSE/Px/DhnWdyv3ZZ5/hvPNOzXSYO3cuXnjhBUgkEpSVlSEpKQl//etfERwcjH/961/w9fXF4cOHuwSiy5cvx/Lly3HnnXd2OYefnx9GjRqFDRs2YNmyZR5dW29w+qf3lVdewcMPP4w777wTL7/8MoYOHQoAA74sh6f+l/c/NBma4Kfww8ToiWgxtnh91PFw3WHUtdchpzHH7eAxIzQDjR2NSAt24pdVP/CV+UIukcNP4dcrx283t+Oz3M8QpgrDopRFdtukBqXiTxP+BF+5a78kGjsa8cLeF6CSq/DQ+IegcDH9/PaK7VhfvB7xAfFeSxo0UM1LnId9NfswKWZSf3eFiHpT7XGxfEX6UsCvnx6KyZVAfJbjdpUHxZqQKTPF6Y3OSpkhBoKJLgRhU+4FdOXi2kxnRA4Hao6dKn9ytnImMFcGiNOD3Tl22nzxy1sEATi4RvwzfKiY9IcGHIvVgr//+vcugSMACBAggQTP/focZsXP8mgK65QpU7B//37s3bsX+/fvx9SpU5GWZv/z9J49e7B06amH44sXL7b9XSKRwNfXF35+4mfdL774Ao8++ijWr18Pq9WK2bNnIyDAudl3o0aNwu7du8/u4DElJQV+fn744osvEBQUhBtuuAEJCX1T5qI/TYqZhJyGHKSHpiPYNxiXDrnU6+dYPmg5DtcfxoSorsXbnTU5djImx/bOwl1viA+Mx2MTH3Nr5M4ZxdpiHG04CkAMYLoL8NyZMmuymKAz6pDTmIOfK3/GrPhZ3bbdWrYVxxqO4eK0i23TNhMDExGgCMDgIC9nmBuAhoYMxdCQof3dDSLqbUc/F2skKtRd16DZ01AojhwNmtv3ozy6SnEdYnOZa/vFjhW/upO/ESjaCoy5DggbJG7zC3VtfWLCRPGrL2jLgV9eFkf/xl3fN+fsTtF2ccpxyoz+Ob9EIo44asvO/sD9HLa/dn+nqapnEiCgWl+N/bX7MT7KswcAJ9c8HjhwANOmTet2zWN9fT2Cg7tm9jeZTFi5ciX+8pe/wMdH/KxbVlaGF154AZGRkdDpdLjjjjuwZ88eREc7rh8eHByMnJwcj66ptzgdPF5wwQVYsGABPvvsM7zzzjt45plnsGDBAqhUKiQlJfViF/vXzPiZmBk/s1fPEa4Ox5wEz+brGy1GvHPkHQgQcOOIG3stSPNEb67FHBw8GLPjZyNUFeryyKAjkX6RmB43Hb9U/IJDtYd6DB53Vu1EY0cjchpzbMFjalAqHp/0uFf7RN0TBAHZ2yvR0WZC5ux4KJRM207kdWkLgeLtjqdGCgJQlwvsexdoawCkcmDIQvfOaTaIawRdzWKatkDcJ3SQe+ftTtluQFcFVB86FTz2B7MByP4KCIrveT1pU4mY3Kb6SJ91za6WGmDfe+LfwwZ3zsjbl9IW9M95yWl1+jqvtuvJsmXLsGzZMphMJgwfPhzr16/HokVdZ7IFBASgra2t0zaDwYBLLrkEWVlZuO++UwkRNRoNFi1ahMcfFz8DXnLJJVi1apXt+560tbUhMDDQw6vqHS4lzFGr1bj22mtx7bXXIi8vD6tWrcIHH3yAbdu2oa2tDcuXL8esWbNsETf1HZ1RhyJdEQBAa9D+7pKVyKVyLEx28wOJExYlL0KATwCGhfScyvyiwRehoLnA4ydg5D6zyYrK/GYAgK6hHaExAzPVNdFZLW6s4zqLLdVA+V4xy6pMIQY3ntQ9/PlfQH0+MOkOIHaM4/YnSWX211paTGJZEU2ce2sxx1wrJsdJ/e2BotUirqPTxNpPwNNbKg8C+RvE6+wpeEyYKI729ZTYxl1WK3D8WzHjasrMntuqQ8X3T7ACfr+vzyrkGmc/y3r6mbeqqso2GqhQKBAYGIiOjg67bUeOHInc3Fzb93q9HsuXL8fEiRPx9NNPd2qbmZmJ0NBTMxHCwsK6Pe6Zjh071mlK7EAiEQSh60RiF5jNZnz77bd45513sH79etx77714/vnnu22v0+mg0Wig1WoHbEQ9kJW3lKPZ0IzhYV0TBBytF6dt2nuN6Pek+oQWHf/f3n3Hx1nd+eL/PNM1I82M+qj34l5xrxgwzRRDCDcQSIAQyELCheWGkLub5N4ky27y2yU3m2yShSQkWZbQMQFTbMDGNjbutixZXVbv0vT+PL8/hA2yRtKMpkny5/165eXoec5zzlFlvnPO+X7tXhTMT53VGaCJpq3ancDpl0dWlSydQNGG4JPOjOfDfwIGGoDV35p4O2mw2o+MZEeVyYGbfxt+vci2T0fqWyrUwE2/Dn9+wXJbgWN/GilZURmnF5u9Z4G9n7322/b0yPlGojD5RT+2vrIVvY7egOceBQjI1GbinVveCevM47XXXovKykoUFRVhz549OH78OI4dOwaDYWzpniNHjuDRRx/F3r17AQDXXHMNOjs78fDDn+e0uOWWW5CcnIz9+/fj1ltvxXe+8x1YLBb86le/wv79+zF//nw0Njbiww8/xPPPP4/ExETccMMNuPbaa5GdnQ2/34+SkhIcOXIEaWnTL9FiWHUeAUChUODmm2/GzTffjI6ODrS3RyBlNgXkF/34j5P/Aa/oxTcWfGNMch0GjUQjTMURqNVGRFMniSP/phQDG78LqCKQLG39oyM1/8YrRD8Rl2WkjuMXz1umlY+cdzPmhx84AkBKCZBcMBLExZI6aWQ1NlpEP3DglyP1Gdd+B1AGyFieXDiyeqvWj9T6JIoAuUyOJ1Y8gUc/ehQChFEBpICR39nvrvhu2PUeX331Vfz+979HbW0tNm7ciGeffTZg4AgAy5cvhyiKqKurQ3l5OebOnYucnBwcPHjwQptrr70WycnJWLt2LV555RW89NJL0Gq1OHjwIObMGdmV0NfXh4MHD6K4uBgAcPDgQaxduxbZ2dnYuXMnNm/ePC0DRyCElUebzYbh4eFJ2yUlJY37BQe48hiu/zz1n+h19OLBxQ8iRTO6GK7T58SQawjZiWPPD/Q6elE7WIsVWSuglqtjNd2oEyURA86BS26bLhHRtCZJI2UrkrJGVvbiydwBfPB/R7ZUbv1p/OcTa45BoKcKyF0xtWRFzmHg7b8f+Z5u+YfYB8d0yQtU59GkNeG7K74bkTqPoTp8+DBaW1txyy23RKX/X/3qV7jxxhuRm5sblf7DFXTw+POf/xyPPz55YfXHHntswvIdDB6j5xfHfoEOWwe+OuerWJC+YNS9fz/+72i1tmJL/hZsLZw9h8RfqXsFh7oP4dqia6Oe2IiIiGagoRbggx+PrIhd+/PQagyeN9gEJCSP/C8ckgQc//NIQLbi/thknt3785FMtxXXAAtuDf15SQI++AngtQFX/t+R86tEMeYX/TjWewx9jj6ka9OxNGNp2CuONDVB/wUVBAFyuRxbt27FXXfdhczMzIDt8vLyIjY5Co1KroIAIWCm1bmpc2H1WFFqjGNGuCjwS/5R/xIR0SXAZRlJdpO1aPIkPMmFwFU/HtluOZXAsbtqJFlPQjJwXZi1rT12oGnPyP8fagEyYlDaKL0S6Dj6+VbiUAy1ACf/G+g6DmiMwHArkFoS6RkSTUoukzMZ4TQR9Mqj1+vFG2+8gWeffRYffvghrrnmGtx777245pprIJcHH/lz5TF6vKIXTp8TetWl83X1i370Onph0pmYGIXgdvrQd86CzGIDlCq+I0k0a1XvAKrfAPRZI4FhNA02AR/980hW1i0RKLvUsm+k5EjFtZE5azkZcwfw/j+OjHXNvwDalMmfOe/oc0Dz3pHVxsrrgMrrYzNnIpq2ppRtta2tDb///e/xhz/8AV6vF3fddRe++c1vBlXvkcEjEUXLyd1t6GmxILcyGXPXxql2GBFFn6UTOPE8kHtZ5AvNN+8Fjv8FmLf985qUHsdIFtWZuE3O6xzJhCpTAhv+PrRtp9bukRXeoo2xrWXpdQEHfz2yWrzi/pn5dSeapcIq1SGKIv7zP/8T3/nOd/DQQw9NeNbxPAaP8ef2u9Fj70G+Pgq1noji6NyZATQe60Xl6ixklxrjPR0imomO/2WkXmPeSmDl/fGeTfhEPyDIZtaK4WDzyDlVALj2Z6GtlhJRVE2pVIfL5cJrr72GZ599Fvv27cP111+P22+/PdJzoyj5S/VfUDtUi5tKbsKanAkKCtO080nnJ6gbqsNNpTfBoGY5iosVzEtFwbzUyRsSEY1n/q1AailgWjB52+nO0gl8+BMgKRu4/Pvxnk3wUoqAJXcCCg0DR6JpJqTg8fjx43j22Wfx/PPPIzs7G/feey9eeOGFaVuHhALTKDSj/qWZY9e5XbB6RxIfrc1ZG+/pxFXVnnaY+11YcmU+tPqxSaIizev2w+PyQWeYPaVuiCgApQbIXxXvWUSGc3hkC6i1CxBFQCaL94yCV7I53jMgCsjpdGJgYCBqpTQaGxtRVFQE2TT9fQ162+qvf/1rPPzww9i6dSvuvfderFy5MmA71nmc/kRJhM1ru6QS68STKIl4r+U9aBSasMuJnOw7icbhRlxdeDW0Sm1kJjhD7X6uGn6fhIWX58JUFP1V2H0v1sNh9WD5tYVIyYpAwXMioljorRnJFJtkivdMiKZM8vvhOHIUvr4+KNLToV2+DEIICTuDVVtbi+TkZGRkZIzb5vHHH8eSJUvwla98Bb29vejt7QUAZGRkjHpuonuNjY1wOp2j+s3Ly4PBYMA//MM/oKioCPfcc08kP7WIYZ1HmhJJkqZddlOv6IVSNv3qT52znMOvTvwKAPD9ld/ndtMIGeq2wzbsRm55MgRZaD+LPq8fZ/Z2Qq1ToHJVVlDP7H+lAQ6zG5ddVwRj5qUduBMREcWK5b330PPTf4Kvu/vCNYXJhMwnvwf9VVdFbJy33noL27dvx+OPP44f/zhwFuehoSEsW7YM9fX1kMvl+OUvf4nf/va36OrqwsMPP4wf/vCHF9pOdO/WW2/F2bNnL3x85swZ7NmzBxs2bEBvby/WrVuH2traafdaGwgheOzo6EBzc/Ok7XJzcyfMusrgMb4GnAN47sxzyNPn4UvlX5pSH/9x8j/Qbe/G3y3+O2Rox39nJpbO9J/Bn6r/hGWZy3BbxW1j7nfYOiAX5DDpYv/Oq1f04rX615CgSMD1xddPyz8E4fC4fBAEAUr15+8AiqKEE++3wuv2Y+nWglH3Yj4/pw8Nx3qRlpuIjIKRvzkDnTYc3XkOALD5zsqg5ufz+uHziNDopt8bFERERLOR5b330PGdR4CLw5XPXkvl/OLpiASQQ0NDuO666zB37lyYTKZxg8ff/e53OHbsGH7zm9+Muv7QQw8hLS1tVIAYzD0A2LdvH+6++240NDRceI24adMm/J//83+wYcOGsD6vaAj6zGNOTg5ycnKiOReKgTZrG7od3RhyD00peJQkCd32bjh9Tgy5hqZN8DjgGoAECQPOgTH3+p39+OXxX0IuyPG9Fd9DoipxyuPUDtbi9YbXsT5nfdDJhpQyZcCAdjZwO33Y/3I9BABrv1QGlWbkT4rP7Ud/uw0A4LC4YUiP/UpdW80g2moGoU/VoLPBjIF224XgMcWkQ/HidGh0iqADW4VSDoWS6eKJiIhiQfL70fPTfxobOAIj1wQBPT/9JyRt2RL2FtZvf/vb+PGPf4wdO3ZM2G7fvn0RD+ieeeYZ3HPPPaMWF1asWIG9e/fO7ODxxIkT2Ldv35jrMpkMubm5WL58ObKzL826an7Rj0+6PkFOYg6KDEXxns6EFqYvhNPnRKYuc0rPC4KAby36Fobdw6hIqYjw7KZufc56ZGozkZM49g0OtVwNrUILpUwJZSj1rQKoG6rDgGsAZwbOMFMtANugE/3tNugMakji53/cVQkKLLkyH16PPy6BIwB0N5lhG3JDn6pBel4S0guSLtwTZAJKl038xocoSuhpMkOfnsAkOURERDHmOHJ01FbVMSQJvu5uOI4chW7liimP8/rrryMpKQmXX375pMFjR0cHMjOn9ho6EIvFgldffRU1NTWjrptMJtTV1UVsnEgKOnj86KOP8MQTT4y5LooivF4vNBoN/vVf/xUPPvhgRCc4ExzvPY4djTugVWjxwzU/jPd0JiQTZFidvTqsPjJ1mVMOPqNFEIRxg9kkVRKeXPkkBAiQh1loeEv+FiSpkjAvdV5Y/cwWfW02JCVrYEjXQK0dHZin5yeN81RsVK7OQu85C/LmpFxYEQ1FR90QavZ3QWtQYd2tZVGYIREREY3H19cX0XaBDA8P44knnsAf//hHVFVVYWBgAE6nE+3t7QGzqapUKng8nimPd7Hnn38e69evH7O70+12Q62enm9cB/2K6pFHHsEjjzwS8N7g4CBefPFFPProo7jiiitQVnZpvdAqNBTCpDOhxFAS87EPdBxA9WA1tpdtR4qGtZDGo5BNqaTpGFqlNuyMqbNJ/txU+Nx+ZJUa4z2VMZJSNEhKmbgcTXvtEGoPdaP8skzkzRn9+5OUrIFSJUMyk+MQERHFnCI9PaLtAunp6YFCocB9990HAOju7oZMJoNer8fPfvazMe0rKirQ0tIy5fEu9swzz+DJJ58cc725uRlLliyJ2DiRFJFX1CkpKXjggQdw4MABfPDBB5dc8JiWkIZHlz0as/HMbjNeq38Npcml+Lj9Ywy5h1DVX4UNudNvX3S4LC4v3jndjXk5eszLZpbS6UarV2H+xujUOYoFc68Dfq+I4R7HmODRmKnF5q/OidPMiIiILm3a5cugMJng6+kJfO5REKDIzIR2+bIpj1FRUYGqqqoLHz/yyCNITEwcN2HO1q1b8dvf/haPPfYYAMBsNqOtrQ2Dg4MQRRFVVVUoKiqCTqeb8B4AnDx5Eq2trdi2bduYcT7++OOAOz6ng4hWnywtLUVHR0cku6QAzvSfQfVgNT5o/QA3ld6EdTnrsDxzebynFRWfNA5gX0M/XjvGn6uZxuPyYf8rDTi0owl+vxjv6QRUvsKEeeuyUbEqcBberoZhNB7rhSgGlZSaiIiIIkSQy5H55Pc+++CiTPWffZz55PciWu8xKytrwjONW7duRW1tLfr7+wGMBHm33347Tp06hb179+L222/H6dOnJ70HALt27cJ3vvMdKJWjj/0cOXIEeXl5KC4ujtjnFUlBl+oIxh133IGNGzfi/vvvH7cNS3WEz+lz4t2Wd1FsKMbC9IXxnk5U9VpdePVYBxblGrG6JDXe06EQWAdd+OS1RgDAxv9RPuZM5HTicfkAYNTZSL9fxO4/1sDnFVEwLwWVa7Igl0f0/TYiIiKaRKzqPAbr5ZdfxtDQEL7xjW9Epf8nn3wSt9xyC5Ytm/qKajSFHTyKooju7m689NJL+N73voeqqqoJI2UGj7Ex5BrCM6efgUlnwlfnfjXe06FLVF+bFXK5DCnZurD78nn8qDvcA0NaAnIqkiMwuxHny40AwLpby6BK+DyAPHuwC2c/6YZGp0T5ikwUL576uQoiIiKaGsnvH8m+2tcHRXo6tMuXRXTFkYIX9JnHn//853j88cfHva/T6fDLX/5y2i6xTneSJKG2x4q8ZC106vCPonbbu9Hn7MOwexiiJEImcMVkumoYasDOlp3YlLsJC9IXxHs6EZWeF7mMq73nLGg/O4RO+fC4waPD4sGxd89Bn6bBws15QfUridKFMiMXv5dWuSoLMrmAjrNDMKQnhPcJEBER0ZQIcnlY5TgocoKOUrZu3Qqj0TjmukwmQ05ODpYuXYr0MLIdXeo+quvDa8c6UJqRiG9vCT/hUGVKJW4rvw2pCakMHKe5433H0WZtw5GeI7MueIyktLwkmIr10KeNH8RZB11wWDxwO7xB96vRKbH2lpHfuUBba8svM6H8ssBnIomIiIguJUEHjwsWLMCCBXxhGy0pWhUEAUjRqSLSnyAIWG6aXkl0mvps2F3Tiy1zMlCcnhjv6QRkdnrh9vmRkTRxiYdI2pK/BVqFFssyp+fe9ulCpVFMupqYUZCEeeuyoTOGVhtJkzh9z2MSERERTRcRTZgTDJ55HJ/L64daIYNwcUapWeL3+5pxom0Yi/OMuGddUbynM4bHJ+IHO6rg9Pjx+NWVyDFOn22Kww4PkjRKyGWz82cjliz9TuiMasgVoa3Iu+xeHH+vFTqjGgs3By5P0t9uRX+7DcWL00cl3yEiIiKaDfjqZhrRKGf3wd8r5mRCLhOwuSIj3lMJSBAApVwGt0yEchoFaSfbhvH7/c2Yk6XHAxtL4j2diHE7vBAEYVSCmmg7VzWA2kPdyChIwuIr8kN61jrggnXQBfuwG6IoQRbgZ6RmfxecNi/UWiWKFqZFatpERERE0wKDx1nq6LlBdA67cPV8E5TTpLxAfqoWd68pjPc0xqWUy/DktXPg9YtI0kyfbYxunwhJGlmZni1cdi/2v1wPQSZg3ZfKprxK57B40N1oRk6FMahSIDLFSMAX6qojAKTlJWLO2ixok1QBA0cAKFyQht5WCzILuauCiIiIZh8Gj7OQJEn48yfnYHZ6seNkB65fmI0bF+fEe1ozgkYpn3YrwCuKUpBt1CAtMbRzfJeCswe70N9mg8vhxdy12ZO2z6tMQVpOItS60N8cEAQBeZUpE/c/NwV5cyduQ0RERKERRQld9cOwW9zQ6dXIKjOO+0ZutP32t7+F1+vFQw89NOU+vv3tb+Pmm2/G5s2bIziz2JgeS1IUUYIgYMucTKToVFDIZKjrscZ7ShSm3GTttAtqw6HRKbHuS2UjdRXDOBtoKtJDZ1AhPX/ykiB1h7vx6ZtNgIC4/QeHiIiIQtN4vBd/evIAXv+343j/2Wq8/m/H8acnD6DxeG9E+v/Rj34Eo9F44X833XTTuG09Hg9+9rOf4Wtf+9qFay+99BKWLl2KlJQUGI1GmM1mAMBPfvKTC30+9dRTo/p54IEH8P3vfz8i8481Bo8zTK/FhY/r++D2TbyFcduibDx1y0LctboQX1sz/ZLTzCS9FhecntmzZXS6UGuVo8472s1umPscIfWRXZaMtbeWBVVPsr1mCMO9Tgx12UOea7AGO+2wD7uj1j8REdGlpPF4L975bdWY/7bah91457dVEQkgnU4nHnnkEbS0tKClpQX/9V//NW7bN998E8uXL0di4kjVgL179+LBBx/EP/7jP6Kurg4tLS0XEoI+9thjaGlpwW233QaXyzWqn7lz58LhcODMmTNhzz/WuG11hvnLwXNoGXDA7vbj6vkT155TymVYXZIao5lNL5IkweMXoVaEt1pX3WnBb/Y0IsuowfeumROh2dHF/F4RB99ogt8rYuUNxTCkRz7T7cLLc2HpcyGz2DDlPiRJGjcb8lC3HUd2tkCulGHzHRWQTZOzxkRERDORKEr4+K/1E7bZ92I9ihalh72jSKPRBKxnf7Hdu3djzZo1Fz7+53/+Z/yv//W/Aq5WajQaaDQaqFSBy/CtW7cO77//PubNmzfVaccFX91EyTtV3fjhjjNo6LVFtN85WXoYEpQoSddFtN/Z5j8/bsITr5xGfZhbduUyAYIAyGdp+ZRQnM8yep7fK+LIzhYce+8c/H4xrL4FuQCNVgGFSgalOjrbc9Nyk1C8JB3ySYI6v1+Ey+4dc72vzYpdf6xG9f7OgM+ptUooVDLo9CoI3BZLREQUlq764Ul389iG3OiqHw57rH/7t39DVlYWNm3ahL17947brrGxEXl5n9ecrqqqgtfrRUVFBcrKyvCjH/0o6DHz8vLQ2NgY1rzjgSuPUXKm04xBuwcNvVaUZiRGrN9VxamwunyzshakxeXFM3ubkJ6kxldXF4bV14DNA78oYcgxNggIRYUpCT/YNg+6KAU0M0XLqX7UHe5BTrkR89aPJF9yWD0Y7BzZAupx+JCQFPidtWDIZAJWby+FJEmTBnfj8Th96KgfRmahHlp9aHNxO7yo2tMBfXoCzL1ODHbZseTK/FFnKe3DbkgiYB10BexDq1dh0x2VEATMyt9PIiKiWLJbgjsGEmy78fzwhz/EE088AafTiddffx033HAD6uvrkZ6ePqatKIqQyT5/neJ2u7F792689dZbsFgsuPXWW1FZWYkvf/nLk44rl8vh98+8Y1FceYySO1YWYPvSHGyujGxNw731fdjX0I8dJzoi2u900DboQMuAA8dah+ENcyXrW5tL8XebS7GiaPLMlxaXF/vq+2F3+wLeT9Gpxt3+2m12odcSOJiYbiRJwtmDXTj1YRv83tC+vufXG6XPFx6RlKLBvPXZWLAxJ6zA8TyZTJhy4AiMnIuoP9yDs590hfxsT7MFNQe6cOiNJrg+e8PB7xv9Ncqfl4pFW/Kw6PK8QF0AGPkcGDgSERGFT6cPLst8sO3Gc37LalZWFh588EHk5+fj1KlTAdsWFBSgq+vz1xl5eXn42te+htLSUixduhS33HILDh8+HNS4HR0dKCwsDGvu8cDgMUpMBg02VWSEfebuYssKklFhSsLGirHvhsx0c7P0uHVZLu5dVxR2bUpDghIVpsmTqADAq0fb8eKRNrxxIvB2xPEM2j3453fO4ql3zsLsDG+FMxa8Lj9azwyiu8kCy4AzpGeLFqZhzfYSzF03uhxGTnkyskqNEZzlaA6LByd2taK72Txp27S8JGiTVMgoCO77/kXphXpojSoYMhKwcFMuVt1UDNNFZyNlMgGZhXpoplDmg4iIiEKTVWaEzjhxYJiYPFK2IxIkScLOnTvR3NyMioqKgG02btyIQ4cOXfj4tttuw+uvvw6Xy4X+/n689957mD9/flDjHTx4cEaW6uC21RkmN1mLv9tcGu9pRIUgCNhQHvuguDQjCbU9VpRlhra9WCkXoFHKIZcBqhmQHEWVoEDlahM8Tj+MGdqQn09M1kRhVhPrqBtC7zkrHBYPTEUTJ7pJz0sKKutqIAk6Ja57cCFEUUJCYvirqERERBQemUzA+i+X4Z3fVo3bZt1tZWEny1m4cCFaW1vhdruRk5ODP/zhD8jNzQ3Ydvv27fjBD34At9sNtVqNhx56CJ9++imMRiMUCgXuvvtu3HXXXQCAF198Effffz+cTicEQcDTTz+NN954Axs3bkRbWxtsNhsuu+yysOYeD4IkfXEjWvRZLBYYDAaYzeYLqWyJZiq3zw8BAlSK6R88AsC5MwPwe0QULU6bEdsrnTYPGo/2IrPIcOH8YU+LBR11QyhbnomklNACWrfTB6/Lj8Tk4La4DHba0XK6H8WL02HMDD3gJiIiovA0Hu/Fx3+tH5U8JzFZjXW3laFkSfjHwywWC0RRREJCAtTqyV8f/OQnP0FqaioeeOCBC9c8Hg8UCsWo85Berxd2++jyYImJiVAoFHjkkUewfv163HLLLWHPP9YYPNKs4fWLYW93nc1cNi/2/rUOALDyhiIY0qdnMOT3iaj7tBtavRoF88eWmjm0ownmPicK5qeiYqUJkiSh/kgPfG4RlatN45bIEEUJe1+og8fpC/rzP/5+K/parTAVG7Bwc+B3IQOp3t+J/lYrFl+ZD31a5MuOEBERXUpEURrJvmpxQ6cf2aoa7orjVPn9frjdbmi1U38dZbVakZQ0td1S8cZtqzQr7KruwY6Tnbh5SeSTFMXLsMODBJU8Yudm1ToFciuT4fOISPxsxU70i2g41gttkgq5lZMnF4qFwU472mqGAAC5c5LHJNEpXZ6BrgYz8ueOzNft8KHl1AAAILvMOO4KoQBAoZDBJxeCrsFYvDgdSrV8TBDbUTsEUZKQN87XbKDdBpfDB0u/k8EjERFRmGQyATkVyfGeBoCRLKnhBI4AZmzgCDB4pFmi1+r+7N+Zkfl0IkN2D37+Xi2qOy1YVZKK715dGZF+BUHA3LWjE94MdNgvBF5ZZcawsp2GQpIk9LVaoU9LGJOAJiVLh+wyI7R6VcD5pGYnIjX78/OpkiTBmJmApFQNDOnjB2qCTMDqm0vg94lQJQT3p8+QngBDes6Fj9uqB+H1+tFwpBcAkJypDXgWdPGV+TD3OZEdoUP8RERERNMBg8c4GbR70G12YW529Lbuenwi3j3TjfwULRblGaM2znRwy7IcLMgxoNwUuZqaAPD68Q70Wd24Y1U+tKrY/Lo09dvRMeREv80Nry+8kiWTSTZpR+oiGgIHatHSXjOEmk+6kJSqweqbSkbdkytlmL8hZ5wnx6rZ34XhHicMaQkQJtnCIlfKIFdO7fM09zlQ80kXJElCanYi5EoZEsapJ5mUogn5PCYRERHRdMfgMU7+/YMG9NvcuGt1AZYXRme74Im2Ybxf3QOVQjbrg0e1Qo4FuRNn4wyV1y/iw9peSBLQ1GfH/JzI9j+exXlG3LWmAEq5DMsKortFQ6GSY9GW8esWBjLc60B/mw0F81OhVE9tS63WoIJMJiApyMQ1E0nJ1sE64EJyli7svi7m8/ph6Xch2aSFzqhGao4OCpUcCzblxu2sBREREVG8MHiMk0y9GsNOD1J0Uy8L4Pb5caJ1GJVZehgSxtaeq8hMQoUpCUVpkX9RfSlQymW4Y2UB+qxuzMmKXXInuUzA5ZWZMRsvVNX7OmEbckOQYcpZzlJzErHl7jmTrhQGo3BBGgoXpIXdTyBVezrQe86K0mUZKF6cjmVXF0ZlHCIiIqKZgMFjnHxzYwn8ogR5GC+e36nqxu6aXszL1uObG0vG3DdolUHXhHzteDv2NwzgnrVFIW+ldXn9kMuEWZnpdEXR9EgiM134PH7YhtywD7tHnTucikgEjqEa7nFAo1NCkzj2zZZAzp+NDPaMJBEREdFsNvte7c8g4QSOAJCfooVKIYvIymJzvwMen4iOYWdIz/VYXPjfr1fhn3eehV+MadUXigOH1QMA0BnVQQdg4RJFCR6nL+x++ttt+PRvzTi0o2ncNp31w+hqGL7w8Zw1Wdh0RwVyp0mGNyIiokuRKPrRduYUavbvQduZUxBFf9zmsmvXLvz5z38Oq4+f/vSnqK2tjdCMYotvp8dAj8UFlVyG5DC2qAayJD8ZS/Ij86L2a2sK0dhnw5IQz0a6vH54/SJsbh9ESYIcM/8c2JDdA1GSkJoY/nm8eKvttmLQ7sGq4hQIQvjfG31qAhZszIFcKRuTJTVaTu5qQ1+bFYuvyENGwdS3D6s0csjkwrhBr23Ijaq9HQAAQ4YWWr0KgiBApeGfSSIionipP3QAH/zxd7AN9l+4lpiShsu/dj/KVq6JyBg1NTX45S9/iZaWFoiiiFdeeQU6XeDFmccffxxvvvkmAMDhcOAXv/gFjhw5gtzcXDzxxBPIysqatN/Vq1fj+9//Pl5++eWIzD+W+KooSt482Ymz3RZcPd+EZz9uhlohx49unAeNMjI1+yItRadCii70LZoFqTo8vrUCWpViVmxbtbq8+OnbNfBLEv73dXPDOpPaa3FhyOFFhSkJb5/uQlOfDXeuKoBRG9k3Ecbj9Yv4zZ5G+EUJRq0yrHOb5j4HTn7QjsxCPSpWmiI4y8l53COrjj7P2HcZLf1OCDIhqMym+rQEbLqjYtwajwmJSqTm6CAIAtQ6/mkkIiKKt/pDB7DjX3865rptsB87/vWnuOHRJ8MOIOvr67FmzRp885vfxLe+9S0oFAqo1YEXEPbs2YP09HTk5uYCAO655x4MDg7igQcewCeffIKrr74aJ06cgCAIE/a7adMm3Hfffejs7ER2dnbAsaYrvkKKkkPNA7A4fWgdcEApl0GjlEEWgZWf6Sg3ObxCqdOJXCZAIRcgiIA8jO+XJEn41/fr4PD48eCmEnx4thdun4jabitWFqdO3kEEKGQC5ucY0Gd1I9sQXqH64V4nXDYv+tqsMQ8el15VALvZDWPG6J8zp9UzsgVVANbfVg6NTgnRL8Lvl6BUBX6TRjHBmzdypYwJcYiIiKYJUfTjgz/+bsI2Hz73O5RcthIy2dQXZ3784x/j61//Op566qlJ27711lvYvHkzgJHXem+88QbOnj2LgoICbN++HTt37sRHH32EzZs3T9ivIAhYv349du7ciXvvvXfKc48HBo9R8rU1hWjqs2NTRQYur8yETAaoFDN/ZW4ivRYXXjzShsV5yVhXNnn2y3equjBo9+LWZblR/dp4/SJ+v68ZfknCveuKoFaM/wdGq1LgH6+fB1GSoFNP/ddDEARk6DXoGnYiWavCXasL0TroiNg242DncO+6ooj0lVuZDJlMgDEz9m8UKNXyMYEjMBLsqTQKCAIu1G48tKMZdrMbl11XBEN6eAEzERERxU9HzZlRW1UDsQ70o6PmDPLmLZzyOAcPHsQTTzyBu+++G6Io4o477sDVV18dsG11dTXuvvtuACOvs5KTk1FbW4uCggIMDw+ju7sbNTU12Lx586T9lpeXo6qqasrzjhcGj1FSmpGE0oykiPXn9vnx+30tUClk+NqawrCT7UTDibZh1PXYMOzwTho8urx+vH26GwCwJN8YcEul2enFmyc7sSDHEFadymGHF2c6LQCAQbsHWZOswiWMs2oVqkevLL+QUddk0ES8DmUsyeUy5M2JfuZZc58DzSf7UTA/FcmmiRNBqTQKrP9yGSAIF2ouel0+iH4p4BZXIiIimjlsw0MRbTeewcFB/L//9//w2GOPwWKx4Pbbb8eOHTuwYcOGMW0dDgcSEj5/Hfmzn/0Mt912GxYuXIj+/n4UFhbCbrcH1W9CQgIcDkdYc48HBo8zRK/FjZqukQDI7PSGdRYvWtaUpsHm9mFe9sRBUnO/HX/6pAV5KQkoSktEaUbgkg+HmgbwafMgmvrsQQePkiThlWMd8PhE3LY8Fwq5DOlJaty5qgCiJE0aOJ7v481TXfD5Rdy0OAcymQCzw4tOsxOVpqSQEs9MxyB/KppO9KGv1Yr5G3KgM47s13dYPBjssiOrxAB5hFaOW88MovecFZIEJJt0aK8dAiQJuZWBA9eLzy+uvLEYbocP+jSuOhIREc1kicbgdmsF2248JpMJDz30EO68804AQF1dHd56662AwaPJZEJ//+eroXfccQeuvPJKnD17FnPnzsVVV12FoqKioPrt6+sblVxnppjd+yhnCEmSsLumB0daBsdtk5eixW3L83DnqoIxgaMkSTjSMojOEMtsTKSqw4w/fdKCAZs76GcS1QpsX5qLCtPEK671PVYM2DwQJeDWZbnjJtpZXpiChbkGXD0/+DN2g3YP9tb14WDTADqHXReuryhKwaogzxr22zzYVd2Dj2r7LpQu+fWeBvzHR4042DT+92g2a68dgrnPif4O24Vrp/e0o3pfJ86dHojYOAXzU2Eq1qNoURrsZjeq93Wien8XbEPB/RyqtUoGjkRERLNAzpx5SEyZeCdbUmoacubMC2ucq666CjU1NQBGXlOfPXsWGRkZAduuWrUKx48fH3UtIyMDGzZswL59+9De3o6tW7cG1e+xY8ewZk1kssXGkiBJUkyL81ksFhgMBpjNZuj1U8/+OJvU9Vjx7x80QBCAp7YvDHnb5JGWQfzpk3PQJyjw45sWhDUXn1/EHw+0YHdND1IT1dg6z4Rti4LLArW3rg9vnuzE9qW5WF0yfqDm9vlxoHEAFZlJyDZG/oX+ruoeuH0irl1gCrhK6PWL+PWHjfD4/Xhoc9mYr7ckSfjbqS74RBE3LhpZeXzm4yac6bTggY0lkwbHs9Fgpx1D3XYULEi9kHSm/kgP2muHsHBTLlJzAq8eh8PvF3H6w3ZIErDw8lzIZ0E2XyIiIgreeNlWz4tEttWuri6sX78eRqMRNpsNer0eu3fvRlLS2Nd73d3d2LBhA2prayEIAv70pz/h+eefx+DgIBoaGvDcc89h27Ztk/ZrNpuxePFi1NfXQ6GYWRtBZ9ZsZ6m8ZC0qs5JgTFBBowz9BXK2MQH6BAXKInDGcsjhxal2MxRyGUrSE4NerQNGtqO6fSJaBuwTBo9qhRybKwK/oxMJV8zNnPC+3e1DU78NkgQMOTxIUI0OYAVBGBMw37e+GF6/GHCV1OsXUdNlQVlGUsTOSwajqc8GuUxAQerEZwMjISVbh5Ts0eOULc9E2fKJv9bhkMtlWHxFftT6JyIioumtbOUa3PDok2PqPCalpmHz3ZGp85iVlYXq6mocPnwYWq0WixYtgkwW+PW4yWTCpk2b8NZbb+H666/H6tWrkZGRgYSEBCxdunRUwDlRv8888wwefvjhGRc4Alx5pAA+aRyAKElYWzp5xtQvsrl9ON1uxqI8A7Sq+PwyDNo92HGiA4vyjBNmNq3utMDrF8NKxHPeGyc6sLumF4vzjLgnQtlNJ9NrdeEnb9VAJgj4wba5MasdOdsM9zrgsnlhKp65yYyIiIhmO1H0j2RfHR5CojEZOXPmhVWeIxyDg4Po7OzE/Pnzp9zHoUOHsGTJEqhUM+/128wLdynqJlo1nEiiWjHlZyPl0+ZBHGsdRvuQc8LgcW52eG9cSJIEi8sHQ4ISGUkaCAKQqZ+8UH2k6FQKJGtVkAkCNBPULqTxiaKEo2+3wO+XoFDJkZYb+a23REREFD6ZTB5WOY5ISklJQUpKeBnoV65cGaHZxB6Dx2lu0O7BvoZ+rCpKQUYMg5NIOtIyiFePd+D6BVlYE+JqZqhWFqWg1+rCwlxjVMd59VgH9tT1YfvSHGyqyMCKopSYZlbVqRX4wba5IWV+pdFkMgHJWTrYh93QGWbeO39EREREscbgcZp7+3QXPm0eRI/ZhW9sKI5o326fH2pF9Fet6nttsLl8ONttjXrwmKxT4a7VhVEdAwAcHh8AwO4eqScYj5IcDBzDt3RrQbynQERERDRjMHiMIafHj4ZeGyqzksYtT3GxpfnJ6Bh2YnlheDVsLvbXw6040DiAu1YXYFlB4KX3U+3DONg0gOsXZo+bFVWSJNjcPiRplOOOdcOibOQYEyJyvvBifz54Du1DDnxzQ0lMa19++bJ8rC1NQ2EMktWcJ4oS3qvuRrJWhZUhJDIiIiIiIooEBo8x9MLhVhxvHcbllRm4aUlOUM/MzdaHfT4vkAG7B5IEDNg847bZVd2DlgEH0hLV2L40N2Cbl4+24+P6fty2PA/rygKvKurUCmwoT4/IvL9IkiQcbx2Czy+hbdAR0+BRpZChOD28M3JvnOjAvvp+3LOuCHOyJv8e1/Va8fbpbgDAojwjzzoSERERUUwxeIwh02dnFjP06jjPBLh7dSFaBx2oyBy/vMd1C7NxuGVwwsDP5h7Zvml1eSM6P0mS8OuPGtFndeORK8oCZhMVBAHf2lSCHosbC3NnXrbMxr6R0iZtg46ggseCFB3m5+iRrFUxcCQiIiKimGOpjhjzi1JMz8c5PD7squnF3Cw9SjMin03S7fOjbdCJknTdhTN4x1uHkKCSo9I09e+vxyfiiVdPweeX8PDlpSibIMidzo63DuFIyxBuXJw9JuHRoN2Dhl4bluQbg97GTEREREQUL3zFGmOxTqyyr74fu6p78MKnrVHpX62QozQj8ULg2DrgwB/2t+A3HzVOeTVSFCWoFDJ8+/Iy3LOuaMYGjgDwXnUPTneYcah5cMy9FJ0KK4pSGDgSERER0YzAbauz3IJcA850WrAk3xiT8VITVcgyapCoUsDrk9A+5EBusjbo518/3oEPa3tx1+pCLCsIPklQx7ATv/6wAWUZifja2qKpTD0qti3MxvG2IayNcpZZIiIiIqJoY/AYB3vr+vDa8Q7cvCQnKolkvijLkID/eWV5VMf4Ip1age9dMweSJOEf3qiCxenDt7eUojQjuNXDTrMTkgT0Wlwhjds17ITV5UN9r20q046aaCU8IiIiIiKKNQaPcdA25IBflNA25Ij3VCLi1WPt6La4cNfqQiSqR36kBEGATq2AwxNaLcmvripAU58d80IMuJYVJEMQRoLlqTjUNACry4ctczLGrZ94ptOMd6q6sXWeCfNzQk/Q09RnwydNA7hyTuaY849ERERERNMdg8c42L4kFxWZSZiXPfMyhF7ML0rYU9cHSQKa++xY8IWsp39/VQW8fhFaVfA/Zkka5ZTrQR5oGECfzY1HrigPqWyHze3Dfx1qhdnphU4jx+riwFtMDzcP4tyAA582D04peNxZ1Y3abivkgoDbV+SH/DwRERERUTwxU0ccJKjkWF6YggRV5MotuLx+vHSkDYeaBiLW53jeqerGf3zUCLPTC7lMwN1rCnHtAtOY7ZlKuSykwDEcPlFC84Adww4vBmzukJ7VqeQoTNOiz+rGy0faIYqBExBfuzALl1dm4PqFWeP21WtxwecXA97bWJ6Oudl6rCmJ/PlHSZLQ1GeDy+uPeN9ERERERABXHmeNU+1mfFzfD4VcwMri1KCf+/BsLwbtHtywODvorJ+7anrg8Yk422XByuJULM0PPrFNuFoHHHj9RAfWlqZiWUHKhetKuQzf2VIGs9MbcnZWQRBw58oC9FrcyNRrIBsnI25GkgY3LckZc93i8uKXu+sxZPfA7ROxJD8Z96wbm7Rnfo5hSiuWwfiwthevH+/EvGw9vrmxJCpjEBEREdGljcFjDImiBEHAuGfqwjE3W4/FeUYUpgWf2dTl9eO14x0AgDlZwSd2uXNVAc4N2LH4ogyuZqcXVpc36OyqXr+IQ02DKM1IhMkQ3BnAI+cG0dBrg1+URgWPAFCQqguqj0Ay9Br80/YFU/reDNg86LG4MWBzw6hVIaaFUz+j+2yFV6fmrzQRERERRYcgSVJMX+taLBYYDAaYzWbo9ZdOFkqzw4t/fvcstEo5vntNZcxq+7UNOvDvHzSgMisJXw9QwmLn6S4M2D340vLckBLbfNGJtmHY3T68U9UNs9OLhy4vRXkQq3+7qnuw42QncpIT8N2rK4Maa8juwe6zvVhekIzCtKkHi19kdnrx2rF2VJj0WF0S/KrtFx1vHUKSRoFEtRKpiaqQvr+SJMHm9iFJo5zS2OdZXF4kqRVReXOCiIiIiIjLFDFi9/hgd/vg8vrh8YkxCx67zC44vX6cGwic2fWaBWPP771T1Y2GXiu+srJg0sQzVpcXf9jfDEkCtCo5FDIBmiCD0OJ0HYxaJSpNwW8zTdapcOuy3KDbj6e224pPGvtx9fwsnO224FjrMBr6bFMOHpeEsXX3zwfP4UjLEO5aXYDlhSmTPzAOfZjBJxERERHRRBg8xki2MQHf2VIGtVIe062FlxUmQykXkGMMvoTFnrpe2N1+nO2yYM0kxe11KgUW5hphc/nw9bWFkMuEoD6/1gEHWgcd+Ifr58YskP6id890o6HXBq1agWvmm9A+5AwpiI0km9s36t+LOTw+yAQBGmXkEiwREREREYWK21bjrLbbigGbG6tLUmO23VCSJIgSIB8nMUxVhxlN/XZcNTdzwoDl4/o+DNg8uH5hFhQhBoA/3HEGg3YPbl2Wiw3l6SE9Gwk1XRYcaBzAtQtMU64NGSkurx8dw04Up+nG/AwM2Nx4audZJKjk+P51c6a8tZiIiIiIKFxceYwjn1/Eb/Y0wi9KMGiVMan7KEkSfvZuLQZsHvz91gqkJ6nHtAkmK6jb58dLR9oBAGWZiSHPfUm+EafbzShOH31u0ezwAgJgSIjuFsw5WXrMyZoeb15olHKUpCcGvOcTJfhECW6viNi+zUNERERENBqDxzhSyGVYmGtAj8WNXGPwWVID+fPBc2gdsOObG0uQljg2IDzPJ0ros7rh9omwuLwBg8fzJEnCrz9qRJ/VjW9vKRt1/lGtkGPrPBMG7O4xgU91pwW7anpw9XzTuIlzblycgxsWZY9aaTM7vPi/b1VDAPCP2+bC65cgEwCjduJzl/EwaPdAIReifs4wU6/B9z5LsMRtq0REREQUTwwe4yxQBtSpONU2DLdPRNugY8LgUSmX4e+3VsDq8o672nWeT5TQ2GeDzz8ScF6cPOe6hWOT7QDAgcZ+NPTacLBpYNzgsdvswr+9X4fc5AQ8vKVs5KIACABkgoAhhxdP76qDXBDwj9vmhp2JFBgJho+1DiPLoEF2CGdAL9ZrdeGpt89CrZThB9vmhR3UiaI0bm1JYKSMCBERERFRvDF4jDOvPzKZVx/cVIJuswuLco2Tts3Ua5AZRECilMvw7cvLMOTwoCKEZDLXLsiCPkGJjROcZey3ueH0+tE+5IQkSRAEAYYEJf5x21wAgCgCckGAQi6MOZtZ32NFepI65BXJY61DeO7AOegTFPjxTQtCevaLBAgQBCGoM6pvnuzE8dZhfH1tIfJSxq4uH2jox1+PtOGa+Vm4er5pynMiIiIiIoo2Bo9x9NKRNuxr6MddqwvGFLwPxOnx450zXSjLSBpzJrE4PRHFk6wkBqO+x4rdZ3tx1dxMFKcnojBNh0KEVk8x25iA25bnTdhmfo4B928oRlqielQQ9sUVxh/eMA+CAGhVn/+YHm8dwh/2tyAtUX0h0AxWpl6DLrMTDo8SLq9/yiuG6UkjYytkk2dAPdU+jH6bGw29toDBY5fZBUkCus3OKc2FiIiIiChWGDzGUZ/NDUkC+qyeoNofah7Ah2f7cLx1eNKENqH68GwvPH4RrQN2VHdaoFXJUZyeiENNA9hb34ftS3Mn3eYaqsk+h0AlPwwJSijkAlITQz8HqU9QXsis2jnsDCvYDjahz12rC9HQZ8PaksAlT7YtykZJRiIqxtneS0REREQ0XTB4jKO7Vxeiud8edNbPBTkGVHdZMDdA+3fPdKOqw4w7VhbAZAjtjNyAzY3XjncAAO5YmY8ElQKbKzMAAIeaB9E26MSJ1uGIB49TUZyeiKe2L4RSHnpZE71GiTtW5sPu8aEoLbTV1KnKS9EGXHE8T6WQYXGeMeR+/aI0bqkVIiIiIqJoYJ3HWeJ83cSbl+RcCPyCJYoSXjzSBq9fxO0r8kedwWwfcuB46zA2lKcHvdp2ttuCvx5uw8bydGyqCG0ul5L6HivSEtVI1oW2itrQa8OvP2rA/GwD7lkXmYRLREREREST4crjDNFrcWHI4R03cc2dqwpQ32PF6pLUkPuWyQTcviI/4L3cZC1yk0MrI3K2y4oBmwen280MHsdxut2M//y4CclaJX504/wx9xt6rWjpd2BjRfqYhEp9Vjd8fgndFlespktERERExOAxnlxePw63DKLSpJ+03uL/914dnF4/vrW5BJWmsSu2pRmJKM2I/7ZSALhybiaSNAos+MKZxgGbG2qlHIkBzjFG0plOM5K1qrBKcYRryO7By0fbMSdLj3Vlgc866hMUn53dDPx9//3+FthcPmhVcqwpHd3HquIUJGkUcf0ciYiIiOjSw+AxjnbX9OLdM90oStPhf15ZPm47QRCQqVej2+KCMSH0RDGxplMrsGVO5oWP24cc+Pm7tUjUKPDDbfOgiEBpkkBqu6347Z4mqBUyPHXLwridCTzRPozTHWa0DTnGDR4LUnUTnt1cWZSCuh5rwDcEBEGIeMIkIiIiIqLJMHiMo5IMHfRNiqAS5jx6VcWMTZIi+6wmoizI2ohTlaJTQaeWI1OvifrXyS9KkAkI+PmsKExBn8U9aW1MlWL8IPrGxTlhz5GIiIiIKJKYMIeixuMToZSPBIzDDg/UCjkSVKHXVuwcdqKqw4x1ZWmjaj4GMmBzo9viwrzs6K3MDTs8+Od3aqFTyfHdayrHnEmMhl6LC2+e6sKKwhQsyOWqIxERERHFHlceKSraBh14elc9cowaPHpVBYzaqW+3feFwG1r67XD7RGxblD1h219+0IBBuwd3rynEsoLkKY85EbvHD4fHB49PhM8vQfmFeLjb7IJRq4RGGXqQPJFPmgZwsm0YQ3YPg0ciIiIiigsGjxTQoaYBvHumG0vyjLhuYTZkIW4DfeFwKz5tHsCiKdQwvNiSfCMcbl9Q23tNBg2sLh9SQyx/EYocYwL+5xXlUCtlo1ZST7YN49l9zShI1eKxqyoiOuaakjSYnd6oBcRERERERJNh8EgB7Wvox+6zvTjRNowsYwKWF6aE9LxPlFCakYjrFmSFPZfNFRnYHGTJjwc2lsTkbGhhmm7MNblMgCCMnPGMtPQkNe5aXRjxfomIiIiIgsXg8RJjdnihVsom3Va5fWkuOs0uqOQCcpJDLwlx//pitA46sCjXOMWZTl28kgrNzzHgB9vmRb0cCRERERFRPDBhzjTi8vqxr74fFaYk5KVoI95/26AD//p+HQwJSvxg29ygM59+cLYHkoRR5TdmouOtQ9hxshPXL8zCsoLQVlKJiIiIiC510U8TSUHbW9eHHSc78ZdD56LSvyhJECUJflFCsG8Z9FhceP14J9440YkuszMq84qVqg4zBmwenG43B9Xe6fFDFGP63goRERER0bTF/XXTSIUpCUfODWFZfnSSohSk6vC/r5sLrUoedAKctEQ1LisaWaXLSNJEZV6xcsOiHJgMCbiscPKvb3WnBb/b24g5WXp8c2NJDGY3on3IAasruORARERERESxxOBxGilI1eHJa+dEdYz0JHVI7eUyAV9dVRCl2cSWQavElXPHbr0902nGK0c7cMXcDKwpSQMA2Nw+iBJgcXnDHrfb7IIoScg2Tnx21OsX8fSuenh8Ir69pQylGYlhj01EREREFCkMHumSd6bDgn6bG6fazReCxxVFKUhNVCFTH95qq9npxb+8exaSBHz/ujlISxw/eFfIBGQZNBiwe5CsVYY1LhERERFRpDF4pGlHkiT8+wcNGHR48J0tZTBqp1azUZIkPHegBUMOL+7fUAzdOFlQr1lgglGrxJKLtguXpIe/8qeSy5CoVsAnSlArJj5iLAhCxOtDEhERERFFCoPHGe7lo+040TaE+9YVB6w9OBN5/CKa++3wiRL6bZ4pB49un4jjbcOQJKBz2ImyzKSA7ZI0Slw1zxTOlMeVoJLjf11dAYVs8vIoRERERETTGYPHGe5slwUWpw/nBh1RDR4PNPRDpZBheWFkSlz4RWnceoxqhRwPXV4Ki8sX1rk/jVKO+9YVw+z0TqmfLrMTu2t6sb4sDQWpU/va9lnd+Jd3zsKQoMT3rp0TtxqUREREREThYvA4w923vhhN/TasiFBQF0jboAMvHG4DMLKVM1k3tZXA8w41DeD5T1txxZxMbFuUHbBNcQS2jALAglzDlJ/dVdOLw82DsLt9U8646vD44PaJsLi8EwbMRERERETTHYPHGc5k0MBkiG4JjQy9GpVZSVDJZdAnjE3k0j7kwNkuK9aVpQXcmmlz+6BTySEII4FTt8UFSRqpITmdrS9Ng93tw6aKjCn3UZCqw2NXlUOrUkA1yZlHIiIiIqLpTJCkYMvFR4bFYoHBYIDZbIZez1p20TRo9+A3exph0mtwz7qiqI3z1M6z6Bx24rqFWdh60dnBj+v78NKRdmwsT8cty3IBjJSkqOmyoDQjEVoV378gIiIiIpoJ+Mp9FuscdqLb7EK/zQ1RlCCL0pbJpflGiJKEStPYhDRWl++zfz+vl6iUy7Aw1xhU39zqSUREREQ0PXDlcRaTJAmHmgeRlqiOScF5l9ePd890ozBVh0V5RgCAKEpo7LMhP1ULtSK0bKPtQw48vaseBSlaLM43otfixrZF2TNi+2ePxYVj54awpjQNhgBbfYmIiIiIZhquPM5igiBgVXFqzMY73jqM3TW9SFDKLwSPMpkwbomMyQw7vPD4RHSanag/YgMAlGUmBr1qGWsurx/HW4cxN1uPV46142yXFVa3D7ctz4v31IiIiIiIwsbgkSJmTlYSKrOSUBqhTKnzcwz41uYSpOrUOHpuCL1WF8qnGIjGws6qLnx4tg/zsvVYlp8Ms8OLhWFkeyUiIiIimk64bZXCcqChH6+f6MBNi3OwpjQt3tOZsg/O9uBwyxD+x2X5yE/VTqmP461D+O9PW3HVXBOumJsZ4RkSEREREcUXVx4pLI39dri8Ihr77TM6eDzcMoSOISequyxTDh6X5CdjSX5yhGdGRERERDQ9cOWRwmJ3+3CyfRiLco3QqT9/L8LrF6GUT//ENue1DjhQ3WXBxvJ0JKhCS+wTCcdbh3CgcQA3LMpGXsrUglciIiIiomjiyiOFRadWYE3J6BXHN0924v3qHty+Im/MvekqP1U75RXHSNhT14emPjs+1Q8yeCQiIiKiaWnmLA3RjNFndY/6N1QWlxc7TnaibdARyWlNazcsysa60jRsrsyI91SIiIiIiALittVZ6HjrEHbV9OCGRTmoMIWWndTt8+OP+1ugUshw1+pCyGVCyOM7PX409NpQmZU0pa2rb5zowO6aXhSm6fDoleUhPx8NO0934b3qHty5qgDLCniukYiIiIguPVx5nIU+bR5E26ATR88Nhfxsr8WNM50WHG8dhtnpndL4CSo5FuQapnzmcVGuEfkpWqwqTpnS89HQOuiAX5QuqdVQIiIiIqIv4srjLNRtduFQ8wDWl6UjRacK+fm9dX1QKWRYVZwahdmFT5IkCEJoK6LDDg/8ooTURPWUxrS6vDjUPIB3TvfAZNDgsavKQ54DEREREdFMxoQ5s5DJoMGNi3Om/PyG8vQIziay3jvTjbdOd+HLy/OCLg1id/vw07dr4BMlfP/aOVMKIJM0ShSk6ODxi+i1uiBKgJyxIxERERFdQhg8EoCRs45vn+5CjlGLFUWx2y7q9vkhigi6PEbnsBOSBHSaXUGPIRMEyGUySJI4pTOc55VlJuGhy0thSFCG1Q8RERER0UzEbasEADh6bhDPHTgHuUzAv315cUzG9PhE/PTtGtjdPjxxTWVQK4JOjx9nuy2Ym62HWhF8PUanxw+/JCFRPf3eLzE7vTjcPIilBclT2mZMRERERBQL0++VNMVFhUmPhbkG5CbHrsagKElweHzw+kV4/GJQzySo5FiSH3q202BXNuPhrVNdONg0gHODDty7rmjMfbvbh6PnhrAw1wCjlsElEREREcUHg0cCACSqFbhvfXFMx9Qo5Xjimjnw+ESYDJqYjj2dzMvWo67HigU5hoD33zrVhX0N/ajrscb8e0REREREdB5LdVBU/PmTFvxwxxl0T3I2MUWninrgeLhlEL/Z04heS/DnJMMlSRIONg2gpd8+YbtPmwehUsjwwxvmjXvWtDQzEYkaBcozQ6vZSUREREQUSVx5pIir7bZix8lOpOhU6Bh2xn1VcVdND7qGXchP0eLaBVkxGfNkuxnPH2qFWiHDz760KGCbhl4r/nLwHGQC8NQtC6FRBt5auzQ/GUunsFWXiIiIiCiSGDxSxP3tVCcMCUqUZyZhab4x3tPBTYtzcLJtGGtLgivtEYqdp7twqsOMr64qQLYx4cL1HGMC0hLVyE9JGPfZTL0GOckJMCYooVZwEwARERERTW/MtkoRd6JtGIeaBnDj4py4rzpOhSRJePloO2xuH76yMn/CrK4/3HEGg3YPbl6ag80VGTGcJRERERFRbHHlkSJucZ4Ri/OM8Z7GlNncPnxc3w8A2FCejpL0xHHb3rW6AA29NqwpSY3V9IiIiIiI4oLBI9FFkjRK3LY8Dza3F8VpugnbFqcnoniC4JKIiIiIaLZg8EgUwLqy0M9H1vVYkanXwJCgjMKMiIiIiIjii1k6iCLg0+ZB/PsHDfiPjxqn3MeBhn48tfMsGvtsEZwZEREREVFkMHgkigBDghJymYBUnSqo9l6/OObapy2D6Bx2oqrDHPCZmi4LemJYq5KIiIiI6IuYbZUoQtw+P1RyGQRBmLDd6XYznt3XhBVFqfjKyvwL1zuGnTjZNoz1ZWlI0oze+lrTZcF/fNSIBKUcT92yYNIxiIiIiIgijWceaVb6tHkQrxxtx7ZF2VM6vzgVE5X0+KJ+mxuiBPRZ3aOu5xgTkGMMXBcyWatCglIOk0HDwJGIiIiI4oLBI81KjX02OL1+NPRaYxY8BmtjeToy9Grkp2iDfsZk0OCfti+ATMbAkYiIiIjig8EjRVWvxYW/HGrFvGw9ts4zRaRPUZQmDaJuXJyNghQtFuQaIjJmJMlkAuZlTzyvbrMLJ9qGsbY09cIWVgaORERERBRPTJhDUVXdZUFLvx0HGvoj0l9ttxV//9JJ/OmTlgnbaVUKrCkde3Zwpnj5aBvePt2F9870xHsqREREREQAuPJIUbaqOBV2tx9lmYkR6a/P6oZPlNBlnl5ZR9882YkPa3tx9+pCLMozht3f0vxkWFy+ablySkRERESXJgaPFFUapRzXLcyKWH9rS1NhSFAiLyVwYplo6rW48G+76pBlSMC3t5SNutc66IDPL6F9yBl28FjVYcaJ9mF8ZUU+CtN0YfVFRERERBQpDB5pRhEEIW6rcQN2D+xuP9qHHGPOXd65qgD1PdaIrDrure/D2S4rkrWqcYPHqg4z3j7dhWvmZ3F1koiIiIhigsEjUZDmZOlx/4ZipOhUY5LXGBKUWF6YEpFxrp2fhWStClsqM8Ztc6RlEO1DThw5N8jgkYiIiIhiQpAkSYrlgBaLBQaDAWazGXq9PpZDE00rbp8f++r7UZ6ZhLwQynYAI2c/DzT2Y01JGtKT1FGaIRERERHR55htlShO9jf0440TnXjuQEvIz6YnqXHj4hwGjkREREQUMwweiaKkbdCBj+v74POLAe+XZSQhU6/G0oLkgPdbBxw40jKIGG8OICIiIiIKiGceaVZ540QHarut+NqaQmToNXGdy7P7mjFo90CAgHVlaWPu56Vo8f3r5gZ8VpIk/PKDerh9ItQKOc81EhEREVHcceWRZpVPm0cSyTT02uI9FSzIMSA1UYXCtNDOMwIjWWUrTElI0algMsQ3CCYiIiIiApgwh2aZhl4bWvrt2FiRDqU8Ou+NdJmd+LR5EOvL0pGiU0VlDCIiIiKi6YbbVmlWKc1IRGlGYlTH2HGiE2c6LbC7/fjKyvyojkVERERENF0weCQK0YqiFNjcPiwvDJzohoiIiIhoNuK2VSIiIiIiIpoUE+YQxYHLbkPziaNwO+zxngoRERERUVAYPBLFQc2+j1C99wOc2vUO7MND8Z4OEREREdGkGDwSxUF6XiGUag1az5zCnr/8PqQAsvHoIRx963W4HY4ozpCIiIiIaDQGj0RxkDt3Pi7/+jdhSM+ETC6HIAv+V7H+0AF0N9aj71xTFGdIRERERDQaE+YQBTDc3QV1YiISEpOiOo7P44Eo+qHSJAT9TMfZagz3dKFi9XooVKwzSURERESxweCR6CL9rS049PpLUGt1uOK+b43bbqC9FVqDEQlJ/DkmIiIiotmP21aJLqLUaCCTy6FJTBy3TXdDHQ6++ld88sp/x3BmRERERETxo4j3BIimG0OGCVd84+8gV4z/66FOTIRMLoc2yRCVOXQ31sM+PITiJctDOg9JRERERBQtDB6JAlCq1BPeTzZl46pvPgyZPPK/Qn6fD0ffeh0AkJiSisyikoiPQUREREQUKgaPRFMkVyij1K8CefMWwj48CGNmVlTGICIiIiIKFRPmEBERERER0aR4mIqIiIiIiIgmxeCRLmletwt1h/ZjuLsr7L6sg/04s2c3bEODEZgZEREREdH0wuCRLmktJ4+h/tABnNr9Tth91R/cj5aTx1B3cH8EZkZERERENL0weKRLWkZhMfTpGcidMz/svnLnLoDRlI28ueH3RUREREQ03TBhDhEREREREU2KK49EcSJJEo7t3IH9f/0z3A5HvKdDRERERDQhBo9EceL3etHdUIfhnm5Y+3vjPR0iIiIiogkp4j0BounO5/Wi7pOPoc/IRG7lvKn14fHgwMvPQwCw+tavQKFSQaFSYdn1N8NlsyI1ryCykyYiIiIiijAGj0ST6G6sQ/OJo5DJ5VMOHt0OO6z9fQAAj8sJhUoFAMgsKonYPImIiIiIoonBI12S7MNDqDu0H7mV85BeUDRh2/T8IphKymDIME15PJ0xGStuvBUQBGj1hnHbdTXUwmW1onDxMgiCMOXxAvE4HfB5vROOT0REREQ0HgaPdEk6d/oEOmtr4DAPTxo8qrVaLLvuprDHnGwcn8eDY2/vAAAkpqYhPb8w7DHP8/t82Ptff4DH5cK62++CPi09Yn0TERER0aWBwSNdkvLmLYDTYo5IfccvGuzsQNWH7yF/3kIULl4W0rNypRK5c+bBYTHDkJF54bro90MQBAiyMPNbCTIIQMRXNImIiIjo0sA6j0QRdPbAXjQeOQSjKQtrb7sz7P4cFjM+fv45aBITsf4rd0Mmk0+5L6/LBZ/Pi4TEpLDnRURERESXHq48En1GFP0Qff4LyWymonjpZZArFMgsKo3InDwOB3weN5wWEaJfDCt4VGo0UEITkXkRERER0aWHK49En/n4+T/CPjyENV+6A/r0jHhP54LBzg6oNBokpqRGbYz26ioMdLRhzvpNUGkSojYOEREREc1cXHmkS1rVR7vQ39qCpdfdBJfdBr/PB4/LGe9pjZKSnRP1Mao//gBetxtGUxYKFiyO+nhERERENPMweKRLWndDHdwOO8zdnVh3+11w2+0wmrIi0rckSUEnp2mtOgVBJiBv7oKIjB2qyrUbMdDRBlNJeVzGJyIiIqLpj8EjXdKWXX8TzD3dyJkzDzKZHAlJkdlKffL9neisrcbyG26ZtOSGdbAfpz94FwCQkp0LnTE5InMIRf78Rcifvyjm4xIRERHRzMHgkS5pyaZsJJuyI96vdaAPoijCPjgwafCo1RuRUVgMQRCgSUwKacWSiIiIiChWmDCHKAqcVgvMPd3ILC69UJ/RZbdBrlRCqVKP+1ztJx+j4fBBLLrqWuRWzovVdC+wDvaj/1wL8uYtDCvrLBERERHNPlx5JIogv8+Hs/v3QGswomjxsgvXLf192P/Cn6DS6rD5a98Yt+SGpb8PAGAd6I/JfC926v2dGO7phtfjRvnKtXGZAxERERFNTwweiSKov+0cWk4eA4CR1TulEgAgiSIkSYIk+oEJ1voXXXkNBtpbkVFYEovpjpFRXAq3wwF9egYajx6CqaR8zBnM9poqVH34PspWrkXJshVxmScRERERxR6DR6IISs3NR+6cedAajBcCRwAwZGRi41fvhUKlhkweeNURAFSaBGSVVsRiqgGVXbYaZZetxpk9u9Fy8hj621qx8qYvjWoz3NMNv88Hc09XnGZJRERERPHA4JEoghRKJRZdeW3Ae/HIojpV6QVF6G1uhKm4bMy9yjUbYMw0Ib2gOA4zIyIiIqJ4YcIcokuEz+NB84kjSMsrQHJWTrynQ0REREQzjCzeEyCiqeuorUHN/j3w+7yTtj13+gTqDu7HyfffjsHMiIiIiGi24bZVohns5PtvQxJF6NMykFMxZ8K26QVF6Kyrgam0PKJzaD5xFIIgoHDR0oj2S0RERETTC4NHohms9LJVMPd0Iy2vYNK2+rR0rP8fd0/YxtzbDY/TifSCoqDGt/T3oXrvBwCAtPxCJCanBPUcEREREc08DB6Jphmf1wtLXy+Ss7IhCMKEbSNZi9Hn8eDAS89D9Pux5rY7kGzKnvQZnTEZWWWVEGQCtAZDxOZCRERERNMPg0eiaebke2+hu7EeFWs2oHT5ypiNK1PIkZSaDpfNCk1iUlDPyBUKLL1mW5RnRkRERETTAYNHomlGqdYAAFQaTUzHlcnkWHf7V2M6JhERERHNHCzVQTTNSJIEt8MOjS4xJuNZ+vuQmJwCmVwek/GIiIiIaGZiqQ6iaUYQhJgFjk3HD+Pj5/84bvkOl90GS39vRMe09PVi/1//jOYTRyPaLxERERFFF4NHohjqrKvBqd3vwuNyjttGkqQJ7werr7UFn77xMoa6OsZtIwiyUf9ePI99//0nfPz8cxjoaAt7Puf1NDdguKcbrVUnI9YnEREREUUfzzwSxdCZPbvhcTqhT0sfty7i8XfeRFd9LZZeewOySiumPFbLiaPoO9cMVYIWyVk5AdsULV6G9IKigJlSBUGAUqOB1+2CQqma8jwuVrhwKfw+HzIKSyLWJxERERFFH4NHohgqX7kO/W0tMJWUj9vG7bADADyO8FYfy1asgSpBi+JllwEAhro70XbmFEqWrYTOmHyh3US1Gdfdfhf8Pi9UmoSw5vJFSo0GlWs2RKw/IiIiIooNJswhmma8bhesA/1Iyc6NaL+HXnsR/W3nkD9/ERZcflVE+yYiIiKi2Y8rj0TTjFKtiXjgCGBkm6wgIG/ugoj1KUkSTu16By6bFUuvveFCmREiIiIimn0YPBJdIjKLS5FZXBrRPn0eN9prqgAAwz3dMGaaGEASERERzVLMtko0A1j6+7D/xb+g6djhqI3RfvYM6g8dgCj6g35GqdZg8dbrUbl2I86dOo73fvtL9DQ3jmoj+v3oaWqA1+WK9JSJiIiIKIYYPBLNAH3nmjDc3RW18hY+rxcn33sbdYf2o7/1XEjP5lTMQcmyFXDbbQAAz2cJf86rO7QfR/722ri1JImIiIhoZuC2VaIZIH/+Ivg8XqQXFEalf4VSicJFS+EwDyPZlD2lPi674RZYBweQmpN3Ud8q2IYGkL9gMY6/+zcMd3dhxY23jsr4SkRERETTH4NHohlAqdagYvW6qI4xb+OWoNs6rRZU7/0A6QXFyJ+/EACgStAiNUc7pu1gZzt0hmTIZDL0NjfC5/HA0tfL4JGIiIhohmHwSDTLtZ05haGuTsxZtwlKTWSS2XQ11KK7sR7DPV0XgsfxZBQUwTrQh7S8AmRXzIFtcACm0vHrXBIRERHR9MTgkWgGsw8PwdLfC1NJOQRBCNjmzN4P4Pd6kZyVjbx5Ewd6F3M77LD0jwR+X+w/p2IebIMDSM8vGvfZvnPNEP1+FC5ehsLFyy5cn+q2WCIiIiKKLwaPRDPYoddfhNNiwaIrr0HunPkB21Su3Yihro4Jy3QMdXVAlaAds5X0yJuvYrinG/M3X4mCBYsvXFdrtVi45epx+3ParPj0jZcBABvu+DqSUtNC+KyIiIiIaDpi8Eg0gxkzs+B1uZCYnDpum8KFS1C4cMm494e6OnDgpeehUKlwxX1/B7ni8z8LCXoDzL09SEjShzQvlSYBRlMWRJ8PmsTEkJ4lIiIioumJwSNRjPU0N6LlxFGUr14X9hbOpdfcEPZ8VAlaKFQqJCTpIZONrt6z5Opt8Pt8UCiVIfUpVyiw9rY7w54bEREREU0fDB6JYqzlxFH0t51Dgl4/Lc7/6YzJuOK+v4NMJoNwUfAoCMKkgeNgZwckUURqbt6E7YiIiIhoZmPwSBRj5avXIUGvR/HSFfGeygWCIODsJx9DqzeMOts4GafNioOv/DckScKmu+5j+Q0iIiKiWYzBI1GMJZuyp8WK4xf1t51D09FPAQA5FXOhUKmCek6pUkNrTIYkilBpEqI5RSIiIiKKMwaPRNPUUHcnDr32IjIKiyNytnEiKdm5yCotR4LBGHTgCAAKlQqbvnpvFGdGRERERNMFg0eiaco+NAi/1wtzT3fUx1KoVFh67Y1RH+c8r8sFpUYzYZvmE0fhdthRsWrdmLOYRERERBR7DB6JpqmcynlQKFVISkuP91Qiqv7QAdQd2o/KtRtQsmxlwDZuhwPVez8AAKTnFzEZDxEREdE0wOCRaJoSBAGm0vJ4TyNsbocDzccPI6usAoYME5w2KwDAabWO+4wqIQFFS5bD7bDDkGmK1VSJiIiIaAKCJElSLAe0WCwwGAwwm83Q60MrPE5EM0/Nvo/QdOwwkrOyseZLd8Dv82Gwsx0p2bmQK/j+FREREdFMwYNERDOc1+3CiXffQtPxw1EdRxT9aD5xFAPtrSE9ZyothyHDhMziMkiiCLlCgfT8QgaORERERDMMg0eiGa63uQkdtdU4u39vVMfpbqhD9d4PcPjNV0N6LtmUjZJlK3B2/x4c27kjSrMjIiIiomhj8Eg0w2UWlyJv3kLM27AlpOcG2lsx3N0VdHujKRv6tHTkVMwNdYrwul2wDw+h/ewZ+H3ekJ8nIiIiovjjmUeiS5Clvw8fP/9HCDIZttzzINRabVTHE0U/3vjZT6BQqXDZtu2zIhEQERER0aWGh46IZgnb0CA0ukQoVKpJ26q1OmgSE6FQqYNqH4r+1hYcf/dvyJu3EJVrNgAAZDI55l9+JWwD/UjNy4/oeEREREQUGwweiWaBnqYGHPnbazBkZGLd7XdN2l6t1WLLPQ9GZS5D3Z3wOJ1jEuuUXbY6KuMRERERUWwweCSaBQSZbNS/0SZJEgRBCHiveOllUGt1SM3lCiMRERHRbMIzj0SzhMNihipBC4VSOep6b0sTqj56H8VLLkPhoqVhjzPU3YlDr76I9IIiLLvuxrD7+6LuxnpU7/0ApZetQv78RRHtm4iIiIjCw2yrRLOEVm8YEzgCQN+5ZjgtFvQ0NQAAWqtOomb/Hoh+/5TGsQ8Nwu/zwtLXE9Z8A+lrbYbTakFPc2PE+yYiIiKi8HDbKtEsV7ZyDTSJiTCVlMPv8+Lkrndg7umGfXgIy6+7KeT+cirnQaFSIyk1DUNdHdAakkdla/W6XWg4cgjpeQVIyy8EAJw9sBeDHe1YvPU6aPWGcfuuWLUOOkMys7ESERERTUNceSSa5VSaBJQsWwmdMRlyhRIZBUWQJBFddWfh9/lC7k8QBJhKymDu7cGBl57HwVdfGHW/rfo0mo5+ilO7371w7dypExjq6sBge9vEc03QonjpZRMGmEREREQUH1x5JLrELLv+ZmiNRiQmp0KumPqfAJVGA0EQIFcqUX/4E+RUzIVWb0BGYQl6GhuQWVx6oe2Sq6+HubcbWeWVkfgUiIiIiCgOmDCHiCbldtjRVV+LrLLK0VtUXS5UfbQLnXU1yCotx9JrI5tAh4iIiIimD25bJSIAI4l1Tu1+F06bdcy96o8/xJk9u1Gz76NR15UaDbJKy6E1GJFZXBajmRIRERFRPHDbKtEsI0kSnFZLyOcGa/Z9BOtAPzQ6HcpXrRt1Ly2vAP2tLUgLULvRVFrOBDdERERElwAGj0SzzNl9H6Hp+BGUr16HsstWX7judbnQ29KEzOJSKFSqMc+VLFuJzroaZFfMHXMvb+4C5M1dENV5ExEREdH0xuCRaJbxeb0AAL/HM+p61Z5d6KytQf6CRViw+aoxz+VUzkVO5djAkYiIiIgIYPBINOvM27QFeXMXwJBpGnXdkJGJrvpaGNIz4zSzsVqrTuLMR7tQuW4TihYvi/d0iIiIiGgCDB6JZhmZTA6jKWvM9eIll6F4yWVxmNH4hnu6IYoizD3d8Z4KEREREU2CwSMRxc2c9ZuQkp2DzKLSyRsTERERUVyxVAfRJcTc243qjz+E02qJ2hh955pxatc7QY2hVKmRO2c+lBpN1OZDRERERJHB4JHoElLz8UdoPn4EjUc/Dbsvr9uFU7vfRWvVyVHXz+7fg7bq0zh3+kTYYxARERHR9MHgkegSkr9gEZKzspFdPifsvnqaGtB25hTO7Nk96nrJ8pXIKCpB7pz5Ifc51N2J6o8/hMtuC3t+RERERBRZPPNIdAnJLp8TkcARADKKSpBTMXdMVtdwxqje+wGGu7sgCALmrNsUgVkSERERUaQweCSiKVFpErB463UR7bNg/mIIgoCsssqI9ktERERE4RMkSZJiOaDFYoHBYIDZbIZer4/l0EQ0Rd0NddAajdCnZcR7KkREREQUJzzzSEQT6mlqwNG338AnL78QVj/t1VXY9cyv0V5TBWAk4U7ziaNwWMyRmCYRERERRRmDRyKakNZohCohAclZ2WPuiaIfLSeOYqirY9J++lqb4XbY0d/WCgCoP3QA1Xs/QNWH70d8zkREREQUeTzzSEQAAI/LiY6z1cgsLoVWb7hwPSklDVd+46GAz3ScrcaZvR9AlZAwbpvz5qzfDKNpJNPrcE836g4dgMtmxdwNlwMAGo4cgstqwZz1myFX8E8TERER0XTDlUciAjC1lcBkUzZ0ySlBJbjR6BJRtHgZ1FotbIP9UKrVyCgqRtHiZfC6Xag9sBfnTp/AYGd7OJ8GEREREUUJ394nIgBAWl4BuhrqkFFYHPQziSmp2PTVewEATccPo+noYSzYshWZRSUTPpdTOQ9yhRJJaekAAKVag/JVa+G0WpGSnTv1T4KIiIiIoobBIxEBADKLS5FZXDrl5/vOtcDtsGOgvXVM8ChJEhqPHIJMIUfxkss+K8dRMapN2Yo1Ux6biIiIiKKPwSMRRcSCy69Cb1MDcubMG3PP0teL2k8+BgCYisugNRhjPDsiIiIiCheDRyKalCRJ6Kithj4tA/rPtppeTKs3oHDxsoD3ElNSkVMxFzKFHAlJwdV3PbX7HZh7urHs+ptHJfAhIiIiovhg8EhEk+o4ewYn398JtS4RV9z7YEjPDnV1oPaTfShavCykbbGdtTXw+3wY7u5i8EhEREQ0DTB4JKJJ6dMzoNbqkJabH/KzbdVVGGhvhSCThRQ8Lt+2HbaBfphKy0Iek4iIiIgiT5AkSYrlgBaLBQaDAWazGXp9cNvXiGjmsg8Poen4YeTNXQhjpine0yEiIiKiKWLwSERERERERJOSxXsCRERERERENP0xeCQiIiIiIqJJMXgkIiIiIiKiSTF4JCIiIiIiokkxeCQiIiIiIqJJMXgkIiIiIiKiSTF4JCIiIiIiokkxeCQiIiIiIqJJMXgkIiIiIiKiSTF4JCIiIiIiokkxeCQiIiIiIqJJMXgkIiIiIiKiSTF4JCIiIiIiokkxeCQiIiIiIqJJMXgkIiIiIiKiSTF4JCIiIiIiokkxeCQiIiIiIqJJMXgkIiIiIiKiSTF4JCIiIiIiokkxeCQiIiIiIqJJMXgkIiIiIiKiSTF4JCIiIiIiokkpYj2gJEkAAIvFEuuhiYiIiIjoC5KSkiAIQrynQTNEzINHq9UKAMjLy4v10ERERERE9AVmsxl6vT7e06AZQpDOLwXGiCiK6Ozs5LscRERERERxxtfkFIqYB49EREREREQ08zBhDhEREREREU2KwSMRERERERFNKuYJc4iIKLpqa2tx5swZbN++fdR1URTx4osvYvXq1SgoKEBnZyf27t0Lg8GAa665ZlRbt9uN1157DQBw6623QqEY/Z+LDz/8EP39/fjSl740Zvzz/QKAIAjIyMjAwoULkZqaOuG8zWYzdu3ahczMTKxbty7kz5uIiIiii2ceiYhmmZ///Of48Y9/jOHh4VHXXS4XEhIS8Oc//xl33nkn/va3v2Hbtm1QKBQ4d+4csrOzL7T97//+b3zlK18BMJIlOzEx8cI9m82GrKws2Gw27Nu3D2vXrh01zvl+b775ZqhUKjQ3N6OqqgpPP/00vvGNb4yZr9VqxWOPPYa//e1vEEURq1atwuuvvx65LwgRERFFBLetEhFd4tasWYPnnntu1LVnn30WGzduDNj+hRdeQFpaGm699VY888wz4/b7u9/9Di+88AIOHTqERx99FA899BD6+vrGtHM6nVi+fDnq6+u54khERDSNMXgkIrrE3Xvvvfj9739/4ePm5mbs378fd955Z8D2zzzzDO677z488MADeOmll2CxWCYd46abboLH40F1dfWYexkZGbj//vuh0+mm/kkQERFR1DF4JCK6xG3btg0Wi+XCOcVnn30W27ZtQ1pa2pi2Z86cwdGjR/H1r38dl19+ObKysvDCCy9MOkZ9fT0AICsrK7KTJyIiophh8EhEdIlTKpX46le/imeffRZ+vx9//OMfce+99wZs+8wzz+C6665DdnY2BEHAfffdN+7W1ddeew0vvPACnnrqKTz88MO46aabUF5eHs1PhYiIiKKI2VaJiGYZQRAQKBfa+WuCIIy5d++992L58uXYunUrZDIZrrzySuzYsWNUG4/Hg7/85S+47bbbLqw2JiQk4PDhwzh9+jQWLFgwqv3OnTuhVquRlpaGp59+GrfffnukPkUiIiKKAwaPRESzjMlkgsVigdvthlqtvnC9t7cXQOCto3PmzMGiRYvwrW99C9/5zncgk43dmPL6669DkiQMDAyMyoY6Z84cPPPMM/jFL34xqv3vfve7gFtfiYiIaGZi8EhENMssX74cMpkM7777Lm644YYL1999911oNJoxK4TnPfnkk/jLX/6Ce+65J+D9Z555BnfeeSeefvrpUddfe+013HffffiXf/mXUcEqERERzS4MHomIZpmKigo8/vjjuOuuu/Dtb38bRUVFqKmpwa9+9Sv85Cc/QXp6esDnrr/+elx//fUB7507dw67d+/Gk08+Oebe1q1b4XQ68dprr015a+qrr74Kj8eD9vZ2+P1+vPDCC1CpVNi+ffuU+iMiIqLIY/BIRDQLPfXUU7j66qvx/vvvY+/evcjLy8NHH32Eyy677EKbnJwcfPnLX4ZSqQzYR25u7oX71dXV+MpXvoL169ePaafVavG9730P3d3do/oNZRXyrbfegt1uR2FhIYCRLbI6nY7BIxER0TQiSIGyKhARERERERF9AUt1EBERERER0aQYPBIREREREdGkGDwSERERERHRpBg8EhERERER0aQYPBIREREREdGkGDwSERERERHRpBg8EhERERER0aQYPBIREREREdGkGDwSERERERHRpBg8EhERERER0aQYPBIREREREdGkGDwSERERERHRpP5/1+DhT9XXwmcAAAAASUVORK5CYII=", + "text/plain": [ + "
" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "result.plot_embedding(figsize=(9, 6))" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "d16de7a0", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
group_idfeature_namescorefrac_exp
01GZMK0.645710.50262
11NSG10.552760.37696
2222CSTA0.949480.96985
2232CLEC12A0.946990.93970
18823LRRN30.745560.54406
18833FHIT0.596580.57088
22034BANK10.915660.94017
22044CD79A0.903530.97436
28315HIST1H2BG0.505420.31148
28325AC092683.10.443790.34426
29276FGFBP20.893070.90625
29286GZMH0.881870.91667
\n", + "
" + ], + "text/plain": [ + " group_id feature_name score frac_exp\n", + "0 1 GZMK 0.64571 0.50262\n", + "1 1 NSG1 0.55276 0.37696\n", + "222 2 CSTA 0.94948 0.96985\n", + "223 2 CLEC12A 0.94699 0.93970\n", + "1882 3 LRRN3 0.74556 0.54406\n", + "1883 3 FHIT 0.59658 0.57088\n", + "2203 4 BANK1 0.91566 0.94017\n", + "2204 4 CD79A 0.90353 0.97436\n", + "2831 5 HIST1H2BG 0.50542 0.31148\n", + "2832 5 AC092683.1 0.44379 0.34426\n", + "2927 6 FGFBP2 0.89307 0.90625\n", + "2928 6 GZMH 0.88187 0.91667" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "marker_table = result.get_markers()\n", + "marker_table.sort_values(\n", + " [\"group_id\", \"score\"], ascending=[True, False],\n", + ").groupby(\"group_id\", sort=True).head(2)[\n", + " [\"group_id\", \"feature_name\", \"score\", \"frac_exp\"]\n", + "].head(12)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "906a71c9", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'report': 'index.html', 'exists': True}" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "report_path = result.report()\n", + "{\"report\": report_path.name, \"exists\": report_path.is_file()}" + ] + } + ], + "metadata": { + "description": "Choose, explain, and execute RNA analysis settings with Scarf agents.", + "jupytext": { + "cell_metadata_filter": "tags", + "text_representation": { + "extension": ".md", + "format_name": "myst", + "format_version": 0.13, + "jupytext_version": "1.14.1" + } + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.0" + }, + "source_map": [ + 14, + 169, + 192, + 197, + 679, + 688, + 696, + 703, + 712, + 722, + 726, + 733, + 738, + 741 + ] + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/docs/.jupyter_cache/global.db b/docs/.jupyter_cache/global.db index fd036171..65df934b 100644 Binary files a/docs/.jupyter_cache/global.db and b/docs/.jupyter_cache/global.db differ diff --git a/docs/source/analysis_with_agents.md b/docs/source/analysis_with_agents.md index 2e84aea4..82628577 100644 --- a/docs/source/analysis_with_agents.md +++ b/docs/source/analysis_with_agents.md @@ -8,7 +8,8 @@ description: Use Scarf safely in an autonomous or AI-assisted single-cell analys This page is a routing and reasoning guide for an AI agent that uses Scarf to analyse data. It does not replace the workflow tutorials or define one correct analysis. The study question, experimental design, and user instructions remain authoritative. -For an executable example of the four bounded Scarf agents and their validated handoffs, see {doc}`tutorials/agent_workflow`. +For an executable ingest-to-finalization example with persisted decisions and report generation, +see {doc}`tutorials/agent_workflow`. ## Scope and authority @@ -96,6 +97,123 @@ Before the first mutating operation, make a short execution record containing: Update this record before changing the cohort, inputs, or decision criteria. This prospective boundary makes unintended writes and retrospective justifications visible. +`AgentOrchestrator` keeps one authoritative stage history containing the immutable request, +effective configuration, evidence, decisions, checks, and final artifact references. The caller still owns the scientific question and +unit of inference. + +### When to use the automated agent workflow + +Use `analyze_rna` with a dataset, a configured model, study context, and a study objective. +The workflow analyzes one RNA assay, even when other modalities coexist in the store. Pass +`assay` explicitly when several RNA assays exist. Automated integration, HTO assignment, and +biological significance or differential-expression hypothesis execution are outside this +workflow; ordinary Scarf APIs remain available for them. Experimental Context still explores +individual and joint covariate patterns and possible explanations of the study design. + +```python +from scarf.agent import analyze_rna + +result = analyze_rna( + "study.h5ad", + model=model, + study_context="One paragraph describing the study design and metadata roles.", + study_objective="Discover stable populations relevant to the study.", + zarr_path="study.zarr", +) +result.plot_embedding() +markers = result.get_markers() +report_path = result.report() +``` + +The beginner call returns a completed result or raises `AnalysisError`. Its result address is +available as `error.result` when work remains unresolved. There is no separate candidate-budget +argument on this interface. Advanced callers import the orchestrator and configuration from +`scarf.agent.orchestrator` for explicit workspaces, numerical limits, and resumable pauses. +See {doc}`reference/api/agent` for that boundary. + +The model first interprets observed study metadata and proposes at most eight objective-led +design comparisons, with one follow-up round of at most four. Comparisons may use a single +explanatory variable, a joint categorical group, or conditioning within categorical strata. +Continuous conditioning bins and regression adjustments are not invented. Continuous biological +variables without a supported preservation measure remain explicitly unresolved. + +Scarf starts from its RNA settings and four partitions of the same graph. The model reviews +quantitative diagnostics, marker and loading-gene evidence, and supplied images before accepting +or requesting one registered experiment. The model cannot generate executable analysis code or +arbitrary `DataStore` calls. When the design permits correction, a matched native/Harmony +evaluation is required even when correction initially appears unnecessary. Accepting correction +requires measured improvement while preserving protected biology, including supported joint +groups, and passing the existing doublet checks. Unsafe or unknown designs cannot license +correction. + +New workflows use an immutable uniform screening cohort containing 10% of retained cells, +rounded up and bounded to 10,000–100,000 cells, never exceeding the retained population. +One larger nested screen may use up to 100,000 cells when the first is smaller. Existing +explicit integer screening sizes retain their exact meaning on resume. A compatible interrupted +run configured for 50,000 screening cells keeps that setting and its matching admitted work; +it does not switch to the new fractional default. Coverage and +rare-population concerns can require targeted full-cohort recovery of measured settings; +a missing screening comparison does not authorize an unbounded full-cohort search. +Screening selects settings; it does not replace the final QC-retained cohort. Selected settings +are executed and assessed on the full cohort. Each screening population must compare the +baseline against 2,000 and 4,000 variable genes, 10 and 30 PCA dimensions, and 21 and 41 +neighbors, changing one setting at a time. The four baseline resolutions share one graph. +Supported batch-aware ranking and an evidence-nominated feature policy provide additional +comparisons. The agent interprets these results, proposes combined settings, and assesses their +actual execution and resolution alternatives before accepting them. A list of reviewed domains +or a general preference for defaults does not establish sufficient evidence. + +The default limits allow 24 screening evaluations per population and 48 overall, with four +additional final-validation graphs, eight additional partitions, and one targeted repair. +For small cohorts, discovery uses every retained cell; these are full-sized comparisons counted +in the screening allowance. Exact completed artifacts are reused for final validation without +another admission. Report the analyzed population and diagnostic operations separately; the +candidate allowance does not bound doublet calculations, elapsed time, or provider cost. +A proposed recovery panel and its matched native controls must fit together before execution. +Four corrected resolutions and four native controls consume all eight additional partitions; +the one-repair limit does not reserve a ninth partition. + +Advanced history records attempted, completed and failed operation calls separately from metric +cache hits, saved-evidence restores and confirmed artifact reuse, for each invocation. These +observed calls differ from counts of unique saved artifacts, and a core operation may itself +reuse earlier work. Older histories without operation records have unknown counts, not zero. + +Experimental Context records objective evidence requirements before tuning. Explicit joint +or conditional covariate requests remain unresolved when only marginal comparisons were +nominated. Full study text is retained; compact model views deduplicate shared sources and +capture-design evidence while complete measurements remain in the stage journal. These +summaries retain adverse findings, missingness, protected-group loss and design constraints. +The agent can retrieve one exact saved policy/capture or design record when detailed thresholds +or donor examples are needed. This lookup performs no scientific recomputation. Completed +metadata inspection and design rounds are checkpointed before further model requests, so +interruption does not reset the eight-initial/four-follow-up allowance. Repeated donors and +incomplete pairing receive descriptive counts and support summaries without treating cells or +repeated samples as independent replicates. A method that cannot compute an association does +not establish that an effect is absent or unidentifiable. Essential unresolved evidence blocks +a consequential decision; measured confounding may prohibit correction while descriptive +population discovery remains possible. + +RNA percentages use explicit gene-selection artifacts. Imported percentage columns remain +available for comparison but do not override a validated definition. In particular, a +mitochondrial symbol definition matches `MT-` rather than every gene beginning with `MT`. +Changing the metric definition requires new QC thresholds and dependent evidence. + +All saved execution and decisions belong to the orchestration stage history. Identical calls +reuse completed work or resume matching interrupted work; changed inputs and identity checks +prevent silent reuse of stale evidence. The result's plot and marker methods use the exact saved +workspace and artifacts read-only. The cluster map displays at most 50,000 cells while retaining +full counts and provenance. `report()` returns or regenerates one local analysis summary from +saved evidence, with no new model calls or numerical analysis. + +This release deliberately breaks the earlier agent imports and persistence contracts. The root +agent facade exports only `analyze_rna`, `AutomatedWorkflowResult`, and `AnalysisError`. Runs using removed configuration or incompatible saved contracts must be restarted; their +numerical artifacts remain readable through ordinary Scarf APIs. +There are no implicit migrations. Compatible histories with newly uncovered objective questions +receive an explicit context-evidence revision, preserving prior records and artifacts; essential +unanswered questions still prevent completion. Model attempt records include known provider +usage, output-validation feedback and failures, with unavailable usage labeled explicitly. +Standalone scientific agent APIs remain in their concrete +packages, such as `scarf.agent.biological_interpretation`. ### When to use the pipeline @@ -220,7 +338,11 @@ Classify the problem before retrying: - **Scientific ambiguity:** preserve branches, seek another independent form of evidence, narrow the claim, or report that the available design does not resolve the alternatives. In a granular workflow, retry the lowest failed stage. A failed pipeline run is not resumable; -start a new run, which can reuse matching complete artifacts from the earlier attempt. +start a new run, which can reuse matching complete artifacts from the earlier attempt. An automated +agent workflow validates its exact request and stage history before reusing completed work or +resuming an interruption. Explicit questions require grounded answers through the advanced +resume interface. A changed dataset, model identity, or configuration cannot be silently attached +to an older history. ## Progress and deterministic comparisons @@ -249,4 +371,8 @@ A useful handoff reports: Artifact provenance records how Scarf produced a result. It does not replace this study-level reasoning record. +For an automated run, the result's plotting, markers, and report helpers resolve the exact final +artifacts from the authoritative stage history. The result is a small address, not a second saved +copy of requests, reports, decisions, or finalization state. The HTML report is a replaceable +presentation of that history. See {doc}`index` for the implemented methods and current boundaries. diff --git a/docs/source/developers/architecture.md b/docs/source/developers/architecture.md index 8c831ce0..10d8d16b 100644 --- a/docs/source/developers/architecture.md +++ b/docs/source/developers/architecture.md @@ -134,13 +134,30 @@ ledger under `pipeline/runs`; it does not write live metadata. DataStore-owned p loading, and export consume narrow frozen-run views. Completed runs can be reopened by their immutable label or exact run ID. +`agent/` owns the optional single-RNA workflow. Its lazy root facade exposes `analyze_rna`, +`AutomatedWorkflowResult`, and `AnalysisError`. Standalone scientific agent contracts and runners +remain in their concrete packages: `data_enrichment`, `experimental_context`, `parameter_tuning`, +and `biological_interpretation`. Internal modules import concrete owners rather than the root +facade. `agent/tools/` contains only infrastructure shared by more than one agent. + +The orchestration stage history is the sole owner of the immutable request, scientific evidence, +choices, checks, work reservations, recovery, and final artifact references. Checkpoints belong +to their exact stage inputs; they do not create a second workflow lifecycle. The result is a small +address that resolves this history. `report/` renders one replaceable analysis page from saved +evidence. It does not call a provider or recompute scientific results. + ### Presentation -`plotting/` is the only plotting package. +`plotting/` owns the reusable plotting APIs. It has no import dependency on `datastore`. The removed `scarf.plots`, `scarf.plotting._legacy`, and `DataStore.plot_*` APIs must not be restored. New plots should return the established plotting result types, accept documented data contracts, and use narrow adapters instead of adding storage-path knowledge. +The single-RNA agent keeps its bounded final-map display in `agent/_plots.py`. This limited report +view reads exact final artifact references, samples only displayed coordinates and labels, and +returns the existing public `PlotResult` and provenance types. It introduces no core plotting API +or module-load dependency from plotting to datastore. + `DataStore.plots` is a thin, store-bound accessor over the canonical store-first functions in `scarf.plotting`. The accessor imports concrete plot implementations only when a method is called, so this convenience namespace does not reverse the dependency from plotting to datastore. diff --git a/docs/source/index.md b/docs/source/index.md index 569d2a19..92717f73 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -100,7 +100,7 @@ Execution choices such as thread count or local scratch are recorded separately This record is useful whenever the analysis is long-running, revisited after a gap, or executed through a pipeline or software agent, because the dependency chain can be inspected independently of the code or description that produced it. See {doc}`analysis_with_agents` for the scientific decision and troubleshooting framework, -{doc}`tutorials/agent_workflow` for its executable four-stage example, +{doc}`tutorials/agent_workflow` for its executable automated workflow, {doc}`concepts/provenance` for the data model, and {doc}`tutorials/reuse_and_tracing` for an executable branching example. diff --git a/docs/source/llms.txt b/docs/source/llms.txt index 6f03108e..fdea4434 100644 --- a/docs/source/llms.txt +++ b/docs/source/llms.txt @@ -57,8 +57,8 @@ - [Provenance](concepts/provenance.html): artifacts, reuse, lineage, and limits - [Reuse and tracing](tutorials/reuse_and_tracing.html): executable branching and inspection example -- [Analysis with AI agents](analysis_with_agents.html): scientific decision loop, task routing, troubleshooting, and handoff -- [Grounded agent workflow](tutorials/agent_workflow.html): executable enrichment, design, tuning, and interpretation stages +- [Analysis with AI agents](analysis_with_agents.html): scientific decision loop, run and artifact inspection, task routing, troubleshooting, and handoff +- [Automated agent workflow](tutorials/agent_workflow.html): executable ingest, preprocessing, tuning, finalization, interpretation, resume, and local report ## Advanced examples diff --git a/docs/source/reference/api.md b/docs/source/reference/api.md index 9ac38105..72c2ccb1 100644 --- a/docs/source/reference/api.md +++ b/docs/source/reference/api.md @@ -11,6 +11,7 @@ Public Scarf surfaces for analysts: - `scarf.plotting` - Documented integration metrics (`DataStore.metric_*`; `scarf.metrics` holds the underlying functions) - `MappingReference` / `MappingResult` for atlas-style mapping +- `scarf.agent.analyze_rna` and its completed result, with the optional agent dependency Inheritance helpers (`BaseDataStore`, `GraphDataStore`, `MappingDatastore`) are listed under {doc}`api/datastore` for completeness. Prefer calling methods on `DataStore`. @@ -24,6 +25,7 @@ Prefer calling methods on `DataStore`. | Graph construction | {doc}`api/graph_construction` | | Artifacts, lineage, and summaries | {doc}`api/artifacts` | | Analysis pipeline | {doc}`api/pipeline` | +| Agent analysis | {doc}`api/agent` | | Assays and metadata | {doc}`api/assays` | | Integration and metrics | {doc}`api/integration` | | Mapping | {doc}`api/mapping` | diff --git a/docs/source/reference/api/agent.md b/docs/source/reference/api/agent.md new file mode 100644 index 00000000..0e90497a --- /dev/null +++ b/docs/source/reference/api/agent.md @@ -0,0 +1,154 @@ +# Agent analysis API reference + +The optional `scarf[agent]` dependency provides a small interface for selecting and explaining +RNA analysis settings. The root `scarf.agent` facade exports three objects: + +```{eval-rst} +.. autofunction:: scarf.agent.analyze_rna + +.. autoclass:: scarf.agent.AutomatedWorkflowResult + :members: plot_embedding, get_markers, report + +.. autoexception:: scarf.agent.AnalysisError +``` + +## Start an analysis + +For model setup and a beginner walkthrough, see {doc}`../../tutorials/agent_workflow`. + +```python +from scarf.agent import analyze_rna + +result = analyze_rna( + "study.zarr", + model=model, + study_context="Human blood from one healthy donor, with no treatment comparison.", + study_objective="Identify stable major immune-cell populations.", + score_doublets=False, +) +result.plot_embedding() +markers = result.get_markers() +report_path = result.report() +``` + +The four required inputs are the source, configured Pydantic AI model, study context, and study +objective. `assay` selects an RNA assay when several exist. `zarr_path` selects the destination +when converting a supported input file. Other modalities in an existing store are ignored. +Automated integration, HTO assignment, and biological significance or differential-expression +hypothesis execution are outside this workflow. Experimental Context still explores individual +and joint covariate patterns and possible explanations of the study design. + +The beginner call returns only after completion. If it cannot complete, it raises `AnalysisError`, +whose `result` provides the status, notes, and exact saved address for investigation or advanced +resume. An identical repeated call reuses a completed analysis or resumes matching interrupted +work. Input and model identity checks prevent attaching changed analysis intent to saved work. + +## Doublet scoring and correction + +`analyze_rna` defaults to `score_doublets=False`. This disables optional advisory doublet +scoring when Harmony is unavailable or prohibited. When the assessed design permits Harmony, +the workflow still computes the doublet diagnostics required to compare native and corrected +representations. Set `score_doublets=True` to request advisory scoring as well. Neither option +automatically removes cells. + +Correction requires an assessed design that protects relevant biology. When that design permits +correction, the workflow evaluates Harmony rather than treating uncertain benefit as a reason +to skip it. Demonstrated confounding can prohibit correction; improved mixing does not override +that constraint or the preservation checks. + +## Use the result + +`plot_embedding()` displays the exact final UMAP colored by its clusters and returns Scarf's +established `PlotResult`. It accepts `figsize`, `show`, `seed`, and `max_points` from 1 to 50,000. +Sampling affects only display; cluster counts and final marker statistics describe the complete +QC-retained population. The image provenance records the workspace, artifact references, +population counts, and display sampling. No live metadata is created or overwritten. + +`get_markers()` returns the saved marker table using Scarf's standard `group_id`, `min_score`, +and `min_frac_exp` filters. `report()` returns the local `index.html` path and can regenerate the +single-page summary from saved evidence. It does not open a browser or rerun the analysis. + +The result stores a small address and outcome, including `zarrPath`, `workspace`, and +`workflowRunId`. Requests, scientific reports, decision histories, and final artifacts remain +owned by the orchestration journal. They are not repeated as public result fields. + +## Advanced control + +```python +from scarf.agent.orchestrator import ( + AgentOrchestrator, + AutomatedWorkflowConfig, + AutomatedWorkflowRequest, +) + +runner = AgentOrchestrator(model, config=AutomatedWorkflowConfig( + inputPolicy="pause", scoreDoublets=False, +)) +result = runner.run(AutomatedWorkflowRequest( + sourcePath="study.zarr", + workspace="analysis", + studyContext="The observed study design and metadata roles.", + studyObjective="The biological structure that must be preserved.", +)) +``` + +The advanced interface exposes numerical limits, provider limits, existing-store workspaces, and +explicit pauses. Inspect its returned status and questions before continuing. The example sets +`scoreDoublets=False` to match the beginner call. The advanced configuration retains +`scoreDoublets=True` as its default, including for compatible older saved configurations. + +### Screening and work limits + +`screeningCells=None` selects automatic sampling: 10% of the QC-retained cohort, rounded up, +bounded to 10,000–100,000 cells, and capped by the retained population. An integer of at least 20 sets +a fixed-size override. A compatible saved integer keeps its exact meaning on resume; the new +default does not resize an existing screening population. + +| Default allowance | Limit | +|---|---:| +| Initial screening population | Automatic sampling as above | +| Evidence-triggered enlargement | One, up to 100,000 cells | +| Candidate evaluations per screening population | 24 | +| Candidate evaluations across screening populations | 48 | +| Additional full-cohort validation graphs | 4 | +| Additional full-cohort validation partitions | 8 | +| Targeted full-cohort repair | 1, within those graph and partition allowances | + +The selected settings are executed and validated on all retained cells. Screening includes +every retained cell in small datasets; those comparisons count in the screening allowance, and +exact artifacts can be reused for final validation. The additional-validation allowance is not +a cap on every graph built during all-cell screening. A recovery comparison and its required +controls must fit the remaining allowance before execution. Four corrected resolutions plus +four matched native controls, for example, use all eight additional partitions. + +These limits count distinct admitted work, including failed attempts. Exact completed reuse +does not spend another slot. They do not bound wall time, retries, diagnostic suboperations, +or provider spend. QC, markers, stability, doublet diagnostics, I/O, and the final UMAP also +have costs. + +### Failures and saved history + +Beginner failures raise `AnalysisError`; its message includes the stage, reason, and available +resume location. Its `result` retains the structured outcome used by advanced callers. Keep the +store and repeat the same call, including model configuration and doublet setting, to reuse +matching work. An unresolved scientific requirement may need investigation before a repeat can +succeed. Changed data, metadata, study text, model configuration, or execution settings cannot +silently reinterpret an existing history. + +The stage journal owns requests, committed evidence, validated decisions, rationales, and artifact +references. Failed model attempts retain available usage and validation feedback. History +separates attempted and completed operation calls from restored evidence and confirmed reuse. +Counts of saved artifacts do not establish how many computations ran: a core call may itself +reuse work, and older histories without operation records have unknown counts, not zero. + +Compatible histories can append a context-evidence revision when a requested joint or +conditional question was unanswered. Previous records remain immutable; changed scientific +evidence must be reassessed. + +The previous agent workflow records, result fields, candidate-budget aliases, and root imports +are unsupported. Histories without mandatory objective requirements and completed comparison +coverage cannot be resumed or used to regenerate reports under this contract. Start a new +workflow; existing historical HTML remains readable. Numerical artifacts remain accessible through +the ordinary Scarf artifact APIs; no saved records are silently migrated. + +For a complete executable example, see {doc}`../../tutorials/agent_workflow`. diff --git a/docs/source/toctree.yml b/docs/source/toctree.yml index f0490ae1..c711f64f 100644 --- a/docs/source/toctree.yml +++ b/docs/source/toctree.yml @@ -86,7 +86,7 @@ subtrees: - file: analysis_with_agents title: Analysis with AI agents - file: tutorials/agent_workflow - title: Grounded agent workflow + title: Automated agent workflow - caption: Advanced examples entries: - file: tutorials/plotting @@ -112,6 +112,8 @@ subtrees: title: Artifacts and lineage - file: reference/api/pipeline title: Pipeline + - file: reference/api/agent + title: Agent analysis - file: reference/api/import_export title: Import and export - file: reference/api/assays diff --git a/docs/source/tutorials/agent_workflow.md b/docs/source/tutorials/agent_workflow.md index 94bc963a..e7deb7cd 100644 --- a/docs/source/tutorials/agent_workflow.md +++ b/docs/source/tutorials/agent_workflow.md @@ -1,5 +1,5 @@ --- -description: Run Scarf's four grounded analysis agents from a rebuilt 5K PBMC baseline. +description: Choose, explain, and execute RNA analysis settings with Scarf agents. jupytext: cell_metadata_filter: tags text_representation: @@ -15,66 +15,193 @@ kernelspec: (agent_workflow)= -# Run a grounded agent workflow +# Automate an RNA analysis -This tutorial runs Data Enrichment, Experimental Context, Parameter Tuning, and Biological -Interpretation against a current 5K PBMC store. -Each stage can inspect only its bounded tools, and each recommendation must cite evidence returned by those tools. +Give Scarf your data, study context, objective, and a language model. Scarf measures the evidence, +the agent compares analysis settings, and Scarf runs the selected analysis. You receive a cluster +map, descriptive markers, and a report explaining the choices and limitations. -The committed build uses small scripted `FunctionModel` callbacks. -They exercise the real tool calls, validators, artifact creation, and handoffs without sending data to an external model or requiring an API key. -The scripted Biological Interpretation callback deliberately leaves the selected cluster unresolved: it demonstrates evidence grounding, not biological expertise. -A live-provider configuration is shown at the end. +## Before you start -## 1. Open the rebuilt teaching store +Install the optional agent package in the Python environment used by your script or notebook: -Install the optional agent dependencies before running this workflow outside the documentation environment: +```bash +uv pip install "scarf[agent]" +``` + +You need an RNA count file, such as H5AD, or an existing Scarf Zarr store. Keep the study metadata +with the cells. The workflow selects one RNA assay; if your store has several, pass +`assay="RNA2"` with its actual name. Other modalities receive no automated processing. + +Scarf uses a model you configure through +[Pydantic AI's provider setup](https://pydantic.dev/docs/ai/models/overview/). +The agent extra includes OpenAI-compatible provider support; other providers may require the +optional dependencies described in their setup instructions. +Configure your provider credentials, then set `SCARF_AGENT_MODEL` in your environment to its +supported `provider:model-name` identifier. This variable is just a convenient way for the example +below to read your choice. If you already have a configured Pydantic AI model object in your +notebook, use that object as `model` instead. Keep credentials out of study descriptions and +shared notebooks. + +Image support is optional. Models without it receive numerical, marker, and loading-gene evidence +instead. Provider calls can incur charges; Scarf's analysis limits do not set a provider spending cap. + +## Run your data + +Replace the paths and study description with your own. Choose a destination you can keep: it +stores the completed work and lets you resume an interrupted run. + +```python +import os + +from scarf.agent import analyze_rna + +model = os.environ["SCARF_AGENT_MODEL"] +result = analyze_rna( + "study.h5ad", + model=model, + study_context="Human blood from one healthy donor; no treatment comparison.", + study_objective="Identify stable major immune-cell populations.", + zarr_path="study.zarr", + score_doublets=False, +) +result.plot_embedding() +markers = result.get_markers() +report_path = result.report() +print(report_path) +``` + +The call runs unattended. It returns a completed result or raises `AnalysisError` with the failing +stage and available resume address. It does not return silently after an unsuccessful analysis. + +Describe the experiment rather than the settings you expect the agent to choose: + +| Input | What to include | +|---|---| +| Study context | Species, tissue, assay, and the actual metadata columns for donors, samples, captures, conditions, and batches. State repeated donors, pairing, and known limitations. | +| Study objective | Which populations or structure you want to investigate, and which biological differences must remain interpretable. | + +A donor, a sample, and a physical capture can be different units. Name their columns explicitly +when you know them. Missing replication or confounded groups cannot be repaired by a language model. + +## Read the result + +- `plot_embedding()` displays the final UMAP colored by cluster. +- `get_markers()` returns a DataFrame of saved descriptive markers. For one cluster, use + `result.get_markers(group_id="0")` with its actual label. +- `report()` returns the local `index.html` path. Open that file in your browser to inspect + the cohort, populations, decisions, and limitations. Regenerating it uses saved evidence. + +Review the explanations alongside the executed settings and measurements. Completion does not +validate every biological interpretation or establish a cell identity. Differential-expression +testing, causal claims, automated multimodal integration, and HTO assignment are outside this workflow. + +## What runs automatically? -```console -pip install "scarf[agent]" +```{mermaid} +flowchart TD + A[Import RNA and inspect study metadata] --> B[Assess cell QC and retain a cohort] + B --> C[Compare genes, PCs, neighbors and clustering settings] + C --> D[Agent weighs evidence and proposes settings] + D --> E[Execute and validate on all retained cells] + E --> F[Final UMAP, markers and report] ``` +Scarf checks defaults and supported alternatives, including variable-gene counts of 1,000, 2,000, +and 4,000, PCA dimensions of 10, 21, and 30, and neighbor counts of 11, 21, and 41. Infeasible +values and identical gene selections are recorded. Batch-aware ranking and gene-family changes +depend on the available evidence. Several clustering resolutions can share one graph. + +When the assessed design permits correction, the workflow evaluates native and Harmony-corrected +representations before deciding whether to retain correction. Confounding between batch and +protected biology can prohibit correction. Better mixing alone does not establish a better analysis. + +`score_doublets=False` is the beginner default. It disables optional advisory doublet scoring +when Harmony is unavailable or prohibited. **Harmony-eligible runs still perform the doublet +diagnostics required for their correction checks.** Set `score_doublets=True` to request advisory +scoring as well. Scoring does not remove cells; doublet removal remains a separate analysis decision. + +## Which cells are analyzed and displayed? + +| Step | Cells used | +|---|---| +| QC | The input cohort, producing the retained cohort. | +| Screening settings | A deterministic 10% sample of retained cells, rounded up, with a minimum of 10,000 and maximum of 100,000; never more than the retained cohort. | +| Final execution and validation | All retained cells. Exact results can be reused when screening already included all of them. | +| UMAP display | At most 50,000 cells, sampled proportionally by cluster with every cluster represented. | + +For 62,721 retained cells, initial screening uses 10,000. The final analysis and UMAP use all +62,721, while the plot heading says **50,000 of 62,721** because it displays fewer points. +Cluster counts and marker statistics still describe the full retained cohort. Display sampling +does not change the analysis. Insufficient screening support can trigger a larger sample or a +bounded additional comparison; it must not be interpreted as evidence that a small population is absent. + +## Continue after an interruption + +Keep the Zarr store and rerun the **same call**, including the same source, model configuration, +study text, destination, and `score_doublets` value. Matching completed stages and evidence are +reused. After upgrading Scarf in a notebook, restart the kernel first so it imports the updated code. +If only report generation failed, the completed analysis can be reused to generate the report. + +An error describing unresolved evidence needs investigation; repeating the same call does not +guarantee that the model can resolve it. Read the reported stage, reason, and resume address. +For a deliberately different analysis, choose another destination rather than deleting your work. +An older run that enabled advisory scoring needs `score_doublets=True` when resumed. + +Work limits, explicit resume, workspaces, and saved-history compatibility are documented in +{doc}`../reference/api/agent`. You do not need those interfaces for the basic call above. + +## Worked example without an API key + +The executable example uses a deterministic 1,000-cell teaching cohort drawn without replacement +from the public 10x Genomics 5K PBMC dataset (random seed 42). Preparation imports the public +file into a temporary Scarf store and marks those 1,000 cells as the active input; the downloaded +dataset is unchanged. Quality filtering and every required comparison still run through the +agent workflow. This small cohort demonstrates the workflow and does not represent an analysis +of all cells in the public dataset. + +The example uses the real analysis operations and a local scripted `FunctionModel`. +The script chooses among observed partitions by agreement across clustering runs, then the +fraction of clusters with qualifying markers. Marker coverage alone does not establish biological +coherence. This makes the example reproducible without an API key. The scripted model below is +only for this demonstration; use your configured model for your own analysis as shown above. + ```{code-cell} ipython3 +from pathlib import Path +from tempfile import TemporaryDirectory + +import pandas as pd import scarf -from scarf.agent import ( - BiologicalContext, - BiologicalInterpretationAgent, - DataEnrichmentAgent, - DataEnrichmentContext, - ExperimentalContextAgent, - ParameterCandidate, - ParameterTuningAgent, -) +from scarf.agent import analyze_rna scarf.configure_output(level="WARNING", progress=False) - -dataset = scarf.cytebase.connect("scarf_docs").download_dataset( - "tenx_5K_pbmc_rnaseq", - destination="scarf_datasets", - zarr=True, -) -ds = scarf.DataStore( - f"{dataset}/data.zarr", - default_assay="RNA", - nthreads=2, +source_path = scarf.cytebase.connect("scarf_docs").download( + "tenx_5K_pbmc_rnaseq/data.h5", destination="scarf_datasets", +)[0] +teaching_directory = TemporaryDirectory(prefix="scarf-agent-teaching-") +zarr_path = Path(teaching_directory.name) / "analysis.zarr" +study_context = ( + "Human 10x Genomics 5K PBMC 3-prime gene expression from peripheral blood, " + "collected from one healthy donor. The teaching cohort is a deterministic random " + "subset of 1,000 cells (seed 42), not the full public dataset. " + "No treatment comparison, trusted technical " + "batch column, paired modality, or independent replication metadata is available. " + "Do not invent missing design variables or report treatment effects." ) - -{ - "active_cells": int(ds.cells.fetch_all("I").sum()), - "total_cells": ds.cells.N, - "assays": ds.assay_names, -} +source_path.name ``` -The downloaded store was rebuilt with this version of Scarf. Its labelled pipeline run supplies -the frozen baseline below, while agent-created candidates remain separate immutable artifacts. - -The hidden setup below defines deterministic model responses for this documentation build. -Every response is assembled from the actual tool return, so fabricated cluster IDs, feature families, or evidence IDs still fail the production validators. +The hidden provider fixture assembles responses from actual tool results. Invented assay names, +feature families, candidates, and evidence identifiers still fail the production validators. ```{code-cell} ipython3 :tags: [remove-cell] +import json +from typing import Any + +import numpy as np +from IPython import get_ipython from pydantic_ai.messages import ( ModelMessage, ModelResponse, @@ -83,376 +210,537 @@ from pydantic_ai.messages import ( ) from pydantic_ai.models.function import AgentInfo, FunctionModel -from scarf.agent.biological_interpretation import ( - ClusterCompositionEvidence, - ClusterMarkerBatchEvidence, +from scarf.agent.data_enrichment import ( + AssayFeatureInspectionBatch, + DataEnrichmentReport, + FeatureSelectionPolicy, + StudyContextSummary, ) -from scarf.agent.data_enrichment import AssayFeatureInspectionBatch -from scarf.agent.experimental_context import CovariateEvidence - - -def _tool_returns(messages: list[ModelMessage]) -> list[ToolReturnPart]: - return [ - part - for message in messages - for part in message.parts - if isinstance(part, ToolReturnPart) - ] - +from scarf.agent.experimental_context import ( + BatchCorrectionPlan, + ExperimentalContextDecision, +) +from scarf.agent.ingest import ingest -def _tool_call(name: str, args: dict | None = None) -> ModelResponse: +prepared_input = ingest(path=source_path, zarrPath=zarr_path) +if prepared_input.status != "done": + raise RuntimeError(f"Teaching dataset import failed: {prepared_input.notes}") +teaching_store = scarf.DataStore( + str(zarr_path), min_features_per_cell=-1, mito_pattern="", ribo_pattern="", +) +teaching_store.cells.reset_key("I") +teaching_cells = np.zeros(teaching_store.cells.N, dtype=bool) +teaching_cells[np.random.default_rng(42).choice(teaching_store.cells.N, 1000, replace=False)] = True +teaching_store.cells.update_key(teaching_cells, "I") +source_path = zarr_path +del teaching_store + +notebook_shell = get_ipython() +if notebook_shell is not None: + notebook_shell.run_line_magic("matplotlib", "inline") + +def _prompt_text(messages: list[ModelMessage]) -> str: + values = [] + for message in messages: + for part in message.parts: + content = getattr(part, "content", None) + if isinstance(content, str): + values.append(content) + elif isinstance(content, tuple): + values.extend(item for item in content if isinstance(item, str)) + return "\n".join(values) + + +def _tool_result( + messages: list[ModelMessage], + tool_name: str, + model_type: Any, +) -> Any: + for message in reversed(messages): + for part in reversed(message.parts): + if isinstance(part, ToolReturnPart) and part.tool_name == tool_name: + if isinstance(part.content, model_type): + return part.content + if model_type is dict: + return json.loads(part.content) if isinstance(part.content, str) else part.content + if isinstance(part.content, str): + return model_type.model_validate_json(part.content) + return model_type.model_validate(part.content) + raise AssertionError(f"Missing tool return {tool_name!r}") + + +def _tool_call(name: str, args: dict[str, Any] | None = None) -> ModelResponse: return ModelResponse(parts=[ToolCallPart(tool_name=name, args=args or {})]) -def _structured_output(info: AgentInfo, payload: dict) -> ModelResponse: +def _structured_output(info: AgentInfo, value: Any) -> ModelResponse: + payload = value.model_dump() if hasattr(value, "model_dump") else value return _tool_call(info.output_tools[0].name, payload) -async def _enrichment_reply( - messages: list[ModelMessage], - info: AgentInfo, -) -> ModelResponse: - returns = _tool_returns(messages) - if not returns: - return _tool_call("inspect_assay_features_batch") - - batch = AssayFeatureInspectionBatch.model_validate(returns[-1].content) - inspection = batch.inspections[0] - species_observed = inspection.species != "unknown" - species = inspection.species if species_observed else "homo_sapiens" - evidence_ids = list(inspection.evidenceIds) - if not species_observed: - evidence_ids.append("context:organism") - policy = { - "assay": inspection.assay, - "species": species, - "speciesConfidence": "high" if species_observed else "medium", - "speciesRationale": ( - inspection.speciesReason - or "The inspected features and caller context support this species." - ), - "excludeFamilies": [ - family.family - for family in inspection.families - if family.count > 0 and family.defaultExclude is True - ], - "protectFamilies": [ - family.family - for family in inspection.families - if family.count > 0 and family.defaultExclude is False - ], - "rationale": "Exclude observed technical families and preserve protected ones.", - "evidenceIds": evidence_ids, +def _scripted_workflow_model() -> tuple[FunctionModel, dict[str, Any]]: + state = { + "enrichment": 0, + "context": 0, + "parameter": 0, + "assessments": [], + "requests": 0, } - return _structured_output(info, {"status": "done", "policies": [policy]}) + async def reply( + messages: list[ModelMessage], + info: AgentInfo, + ) -> ModelResponse: + state["requests"] += 1 + tools = {tool.name for tool in info.function_tools} + + if ( + "inspect_assay_features_batch" in tools + or state["enrichment"] == 1 + or any( + tool.parameters_json_schema.get("title") == "DataEnrichmentReport" + for tool in info.output_tools + ) + ): + if state["enrichment"] == 0: + state["enrichment"] = 1 + return _tool_call("inspect_assay_features_batch") + + batch = _tool_result( + messages, + "inspect_assay_features_batch", + AssayFeatureInspectionBatch, + ) + policies = [] + for inspection in batch.inspections: + species_observed = inspection.species != "unknown" + policy_evidence = list(inspection.evidenceIds) + if not species_observed: + policy_evidence.append("context:study") + policies.append( + FeatureSelectionPolicy( + assay=inspection.assay, + species=( + inspection.species + if species_observed + else "homo_sapiens" + ), + speciesConfidence="high" if species_observed else "medium", + speciesRationale=( + inspection.speciesReason + or "The exact study paragraph identifies a human sample." + ), + excludeFamilies=[ + family.family + for family in inspection.families + if family.count > 0 and family.defaultExclude is True + ], + protectFamilies=[ + family.family + for family in inspection.families + if family.count > 0 and family.defaultExclude is False + ], + rationale=( + "Exclude observed technical families and preserve " + "observed protected families." + ), + evidenceIds=list(dict.fromkeys(policy_evidence)), + ) + ) + state["enrichment"] = 2 + return _structured_output( + info, + DataEnrichmentReport( + status="done", + studyContextSummary=StudyContextSummary( + organismReferences=["Human"], + tissueReferences=["peripheral blood"], + experimentalReferences=[ + "10x Genomics 5K PBMC 3-prime gene expression" + ], + analysisIntentReferences=[ + "Discover stable major immune-cell populations." + ], + ), + policies=policies, + ), + ) -async def _experimental_reply( - messages: list[ModelMessage], - info: AgentInfo, -) -> ModelResponse: - returns = _tool_returns(messages) - if not returns: - return _tool_call("inspect_cell_covariates") - if len(returns) == 1: - return _tool_call( - "analyze_experimental_design", + if tools.intersection( { - "column_domains": {}, - "coefficients_of_interest": [], - "units_of_inference": {}, - "batch_columns": [], - }, - ) - - design = CovariateEvidence.model_validate(returns[-1].content) - evidence_id = design.evidenceIds[0] - return _structured_output( - info, - { - "batchCorrection": { - "action": "skip", - "rationale": ( - "No explicit technical batch or biological contrast was supplied." + "inspect_cell_covariates", + "analyze_experimental_design", + "score_current_representation", + } + ) or state["context"] in {1, 2} or any( + tool.parameters_json_schema.get("title") == "ExperimentalContextDecision" + for tool in info.output_tools + ): + if state["context"] == 0: + state["context"] = 1 + return _tool_call("inspect_cell_covariates") + if state["context"] == 1: + state["context"] = 2 + return _tool_call( + "analyze_experimental_design", + { + "column_domains": {}, + "coefficients_of_interest": [], + "units_of_inference": {}, + "batch_columns": [], + }, + ) + + design = _tool_result( + messages, + "analyze_experimental_design", + dict, + ) + profile = next( + value + for value in design["qcProfiles"] + if value["action"] == "skip" + ) + evidence_id = profile["evidenceId"] + state["context"] = 3 + return _structured_output( + info, + ExperimentalContextDecision( + batchCorrection=BatchCorrectionPlan( + action="skip", + rationale="No trusted technical batch column was supplied.", + evidenceIds=[evidence_id], + ), + rationale="No experimental covariates were supplied.", + evidenceIds=[evidence_id], ), - "evidenceIds": [evidence_id], - }, - "rationale": "Continue with an uncorrected baseline.", - "evidenceIds": [evidence_id], - }, - ) - - -async def _tuning_reply( - _messages: list[ModelMessage], - info: AgentInfo, -) -> ModelResponse: - return _structured_output( - info, - { - "status": "done", - "recommendedCandidateId": "baseline", - "confidence": "medium", - "rationale": "The single authorized baseline completed successfully.", - "evidenceIds": ["candidate:baseline:clusters"], - "stopReason": "The authorized candidate was evaluated.", - }, - ) - - -async def _biology_reply( - messages: list[ModelMessage], - info: AgentInfo, -) -> ModelResponse: - returns = _tool_returns(messages) - if not returns: - return _tool_call("inspect_cluster_composition") - if len(returns) == 1: - composition = ClusterCompositionEvidence.model_validate(returns[-1].content) - cluster_id = sorted( - composition.clusterCounts, - key=lambda value: (-composition.clusterCounts[value], value), - )[0] - return _tool_call( - "inspect_cluster_markers_batch", - {"cluster_ids": [cluster_id]}, + ) + + prompt = _prompt_text(messages) + if any( + {"selectedCandidateId", "comparisonConclusions"}.issubset( + tool.parameters_json_schema.get("properties", {}) + ) + for tool in info.output_tools + ): + evidence, _ = json.JSONDecoder().raw_decode(prompt[prompt.index("{") :]) + coverage = evidence["comparisonCoverage"] + settings = coverage["candidateSettings"] + candidates = { + item["candidateId"]: item for item in evidence["candidates"] + if item["status"] == "done" and item["eligible"] + } + if not candidates: + raise AssertionError("The teaching run has no supported partition") + + def rank(identity): + metrics = settings[identity]["metrics"] + return tuple( + float(metrics[name]) if metrics.get(name) is not None else -1.0 + for name in ("seedStability", "markerCoherence") + ) + + axis_names = { + "hvgCount": "variable-gene count", "hvgRanking": "gene ranking", + "featurePolicy": "gene-family policy", "pca": "PCA dimensions", + "neighbors": "neighbor count", "partition": "clustering resolution", + } + metric_names = { + "seedStability": "agreement across clustering runs", + "subsampleStability": "agreement after resampling", + "markerCoherence": "the fraction of clusters with qualifying markers", + "markerSpecificityMedian": "median marker specificity", + "macroF1": "classification agreement", + } + axis_ids = {} + for row in coverage["comparisons"]: + identities = axis_ids.setdefault(row["axis"], []) + for identity in (row["baselineCandidateId"], row["alternativeCandidateId"]): + if identity is not None and identity not in identities: + identities.append(identity) + comparison_ids = {axis: list(identities) for axis, identities in axis_ids.items()} + if coverage["phase"] != "sensitivity": + axis_ids["partition"] = list(dict.fromkeys( + [*axis_ids["partition"], *coverage["resolutionCandidateIds"]] + )) + preferences = {axis: max(identities, key=rank) for axis, identities in axis_ids.items()} + pending_policy = any(row["status"] == "pending" for row in coverage["comparisons"]) + experiment_id = None + if coverage["phase"] == "sensitivity": + selected_id = evidence["currentCandidateId"] + action_name = "combine" + else: + eligible_resolutions = [ + identity for identity in coverage["resolutionCandidateIds"] + if identity in candidates + ] + if not eligible_resolutions: + raise AssertionError("The combined representation has no supported partition") + selected_id = max(eligible_resolutions, key=rank) + preferences["partition"] = selected_id + action_name = "accept" + if pending_policy: + families = evidence["featureEvidence"][selected_id]["families"] + supported = [ + (key, option) for key, option in evidence["experiments"].items() + if option["parameter"] in {"includeFamily", "excludeFamily"} + and families.get(option["value"], {}).get( + "selectedExamples" if option["parameter"] == "excludeFamily" + else "excludedExamples" + ) + ] + if not supported: + raise AssertionError("The teaching policy has no observed family program to nominate") + experiment_id, option = max(supported, key=lambda item: ( + families[item[1]["value"]].get("selectedGenes", 0), + item[1]["affectedEligibleGenes"], + )) + action_name = "experiment" + selected = candidates[selected_id] + metrics = selected["metrics"] + genes = list(dict.fromkeys( + gene for names in metrics.get("topMarkerGenes", {}).values() + for gene in names + ))[:8] + qualitative = ( + "The saved marker preview contains " + ", ".join(genes) + "." + if genes else "The saved marker preview is empty; cell identities remain unresolved." + ) + conclusions = [] + for axis, identities in axis_ids.items(): + preferred = preferences[axis] + score = settings[preferred]["metrics"] + conclusions.append({ + "axis": axis, + "candidateIds": identities, + "preferredCandidateId": preferred, + "quantitativeReason": "; ".join( + f"Observed alternative {index + 1}: repeat agreement " + f"{settings[identity]['metrics'].get('seedStability')}, " + f"marker coverage {settings[identity]['metrics'].get('markerCoherence')}" + for index, identity in enumerate(identities) + ), + "biologicalReason": ( + "This scripted teaching policy does not establish cell identities or " + "infer that a gene program is a technical artifact. " + qualitative + ), + "plainLanguageSummary": ( + f"The teaching policy compared {len(identities)} observed {axis_names[axis]} " + f"settings and preferred repeat agreement {score.get('seedStability')}, " + f"using marker coverage {score.get('markerCoherence')} to break ties." + ), + "tradeoffs": [{ + "alternativeCandidateId": identity, + "metric": metric, + "preferredValue": score[metric], + "alternativeValue": settings[identity]["metrics"][metric], + "interpretation": ( + f"An alternative has higher {metric_names[metric]} " + f"({settings[identity]['metrics'][metric]} versus {score[metric]}). " + "The teaching policy prioritizes repeat agreement, then marker coverage; " + "this loss remains an explicit limit of its choice." + ), + } for identity in comparison_ids[axis] if identity != preferred + for metric in ("seedStability", "subsampleStability", "markerCoherence", + "markerSpecificityMedian", "macroF1") + if isinstance(score.get(metric), (int, float)) + and isinstance(settings[identity]["metrics"].get(metric), (int, float)) + and settings[identity]["metrics"][metric] > score[metric]], + }) + quantitative = ( + f"Observed resolution {selected['parameters']['leidenResolution']} has " + f"repeat agreement {metrics.get('seedStability')} and marker coverage " + f"{metrics.get('markerCoherence')}." + ) + summary = ( + "The teaching policy proposes testing a represented gene family; " + "its contribution must be measured before retaining or changing the gene selection." + if action_name == "experiment" else + "The teaching policy proposes a combination of the observed settings; " + "Scarf must execute that combination and compare its four resolutions." + if action_name == "combine" else + "The teaching policy selected the measured combined representation and " + "its most repeatable eligible partition. Marker identities remain unvalidated." + ) + action = { + "action": action_name, + "selectedCandidateId": selected_id, + "experimentId": experiment_id, + "correctionNeed": "notApplicable", + "comparisonConclusions": conclusions, + "plainLanguageSummary": summary, + "evidenceIds": [ + f"candidate:{selected_id}", + *list(evidence["imageHashes"])[:1], + "studyContract", "qcPolicy", "samplingCoverage", "featureEvidence", + ], + "quantitativeFindings": [quantitative], + "qualitativeFindings": [qualitative], + "objectivePreservation": ( + "Preserve the single-donor population structure and retain marker " + "uncertainty; no batch or treatment comparison is supported." + ), + "rationale": summary + " " + quantitative, + "concern": ( + f"Observed family {option['value']} contains " + f"{families[option['value']].get('selectedGenes', 0)} selected genes. " + "Test sensitivity to this program without assuming that it is technical." + ) if pending_policy else "", + "expectedImprovement": ( + "Measure whether changing this gene-family selection preserves the major " + "marker programs and improves repeat agreement." + ) if pending_policy else "", + "populationConcerns": [{ + "candidateId": selected_id, + "clusterId": cluster, + "status": "nonEssentialLimitation", + "evidenceIds": [f"candidate:{selected_id}"], + "explanation": ( + f"Population {cluster} has no qualifying marker genes and remains unclassified. " + "This tutorial demonstrates selecting analysis settings; it does not validate " + "cell identities or claim that every population has been biologically resolved." + ), + } for cluster, names in metrics.get("topMarkerGenes", {}).items() if not names], + } + if action_name == "combine": + action["combinedSettings"] = { + field: preferences[axis] for field, axis in ( + ("hvgCountCandidateId", "hvgCount"), + ("hvgRankingCandidateId", "hvgRanking"), + ("featurePolicyCandidateId", "featurePolicy"), + ("pcaCandidateId", "pca"), + ("neighborsCandidateId", "neighbors"), + ) + } + state["assessments"].append({ + "selection": action, + "alternatives": [{ + "resolution": settings[identity]["parameters"]["leidenResolution"], + "clusters": settings[identity]["metrics"].get("nClusters"), + "repeat_agreement": settings[identity]["metrics"].get("seedStability"), + "clusters_with_markers": settings[identity]["metrics"].get("markerCoherence"), + "selected": identity == selected_id, + } for identity in ( + coverage["resolutionCandidateIds"] if action_name == "accept" + else list(candidates) + )], + }) + return _structured_output(info, action) + + + payload, _ = json.JSONDecoder().raw_decode(prompt[prompt.index("{") :]) + decision = payload["spec"] + evidence_by_class = {} + evidence_class_by_id = {} + for item in payload["evidence"]["evidence"]: + evidence_by_class.setdefault( + item["evidenceClass"], + item["evidenceId"], + ) + evidence_class_by_id[item["evidenceId"]] = item["evidenceClass"] + preferred = decision.get("metricPreferredOptionId") + selected = ( + next( + option + for option in decision["options"] + if option["optionId"] == preferred + ) + if preferred is not None + else next( + option + for option in decision["options"] + if option["status"] in {"apply", "skip"} + ) ) - - marker_batch = ClusterMarkerBatchEvidence.model_validate(returns[-1].content) - marker = marker_batch.clusters[0] - if not marker.evidenceId: + evidence_ids = list(selected.get("requiredEvidenceIds", [])) + cited_classes = { + evidence_class_by_id[evidence_id] for evidence_id in evidence_ids + } + for evidence_class in selected["requiredEvidenceClasses"]: + if evidence_class not in cited_classes: + evidence_ids.append(evidence_by_class[evidence_class]) + state["parameter"] += 1 return _structured_output( info, - { - "status": "needsInput", - "needsInput": { - "question": "No markers passed the bounded search thresholds.", - "requiredInputs": ["markerArtifact"], - }, - "limitations": marker.warnings, - "stopReason": "Marker evidence was unavailable.", - }, + dict( + selectedOptionId=selected["optionId"], + evidenceIds=evidence_ids, + rationale=f"Use the offered {selected['label']} policy with its required observed evidence.", + confidence="high", + ), ) - names = [item.featureName or item.featureId for item in marker.markers[:3]] - return _structured_output( - info, - { - "status": "done", - "clusterInterpretations": [ - { - "clusterId": marker.clusterId, - "proposedIdentity": "unresolved marker-defined cluster", - "identityIsHypothesis": True, - "confidence": "low", - "rationale": f"Top returned marker features: {', '.join(names)}.", - "evidenceIds": [marker.evidenceId], - } - ], - "evidenceIds": [marker.evidenceId], - "limitations": [ - "The scripted documentation model does not assign cell identities." - ], - "stopReason": "One bounded cluster was reviewed.", - }, - ) -``` + return FunctionModel(reply), state -## 2. Inspect species and feature families +model, model_state = _scripted_workflow_model() -Data Enrichment is read-only. -It inspects the requested assays and returns a policy rather than changing the datastore. -The context below is caller-supplied study evidence, not a label inferred from expression. +``` ```{code-cell} ipython3 -enrichment = DataEnrichmentAgent(FunctionModel(_enrichment_reply)).run( - ds, - context=DataEnrichmentContext( - studyContext=( - "10x 5K PBMC RNA-seq from peripheral blood of a healthy human donor." - ), - organismHint="human", - tissueReferences=["peripheral blood"], - cellTypeReferences=["T cell", "B cell", "NK cell", "monocyte"], - experimentalDetails=["10x 3 prime RNA-seq", "single donor"], - ), - assays=["RNA"], +result = analyze_rna( + source_path, + model=model, + study_context=study_context, + study_objective="Discover stable major immune-cell populations.", + score_doublets=False, ) - -policy = enrichment.policies[0] -{ - "status": enrichment.status, - "species": policy.species, - "exclude_families": policy.excludeFamilies, - "protect_families": policy.protectFamilies, - "tool_calls": [call.name for call in enrichment.toolCalls], -} +{"status": result.status} ``` -The family policy is advisory. -The feature-selection API accepts an exact selection or regular-expression blacklist, so do not silently translate family names into guessed feature names. - -## 3. Prepare a frozen baseline +### What did the teaching policy choose? -Open the durable pipeline run built with this dataset. It supplies frozen metadata, -normalization, and a current graph to Experimental Context. Parameter Tuning consumes the exact -normalized artifact and creates its own PCA, neighbour, graph, and clustering candidate without -changing the run. +The starting graph is compared at four clustering resolutions: 0.5, 0.75, 1.0, and 1.25. +These partitions share the same cells, features, and graph. The table below contains the exact +observations offered to the scripted provider, followed by its recorded explanation. ```{code-cell} ipython3 -run = ds.pipeline.open(label="docs_default") -normalized = run["normalized"] -hvg_ref = run["highly_variable_features"] +:tags: [remove-input] -{ - "run_id": run.run_id, - "active_cells": int(run.cells.fetch_all("I").sum()), - "feature_selection": hvg_ref.artifact_id, - "normalized": normalized.artifact_id, -} +assessment = model_state["assessments"][-1] +pd.DataFrame(assessment["alternatives"]) ``` -## 4. Check the experimental context - -Experimental Context classifies metadata and authorizes an exact Harmony batch-column set only when the design supports it. -Passing the completed run binds metadata and integration metrics to its frozen cell selection and -graph artifact. -This single-donor store has no treatment or batch labels, so the grounded action is to keep an uncorrected baseline. - ```{code-cell} ipython3 -experimental = ExperimentalContextAgent(FunctionModel(_experimental_reply)).run( - ds, - study_context=( - "Healthy-donor 5K PBMC. No treatment or batch labels are available. " - "Do not invent a technical batch or biological contrast." - ), - run=run, -) +:tags: [remove-input] +selection = assessment["selection"] { - "status": experimental.status, - "batch_action": experimental.decision.batchCorrection.action, - "batch_columns": experimental.decision.batchCorrection.batchColumns, - "coefficients": experimental.decision.coefficientsOfInterest, + "Why this setting": selection["rationale"], + "Marker evidence": selection["qualitativeFindings"], + "Biology to preserve": selection["objectivePreservation"], } ``` -The validated result becomes a narrow handoff. -Downstream tuning receives the exact batch action and columns, rather than reparsing prose. +The table shows measured comparisons, not independently validated cell identities. A live model +must weigh marker programs, study design, and protected biology alongside these metrics. -## 5. Evaluate one authorized parameter branch +### Inspect the example's map, markers and report -The original notebook screens five defaults. -For a bounded documentation run, authorize one explicit candidate and disable refinement. -The candidate executes PCA, neighbours, graph construction, Leiden clustering, and its available -diagnostics against the baseline run's exact normalized artifact. +The result fixes the saved layout and cluster labels. Display options include `figsize`, `show`, +`seed`, and a lower `max_points` display cap; they change only the picture. ```{code-cell} ipython3 -if experimental.status != "done": - raise RuntimeError(f"Experimental Context stopped with {experimental.status!r}") - -tuning_handoff = experimental.to_parameter_tuning_handoff() -candidate = ParameterCandidate( - candidateId="baseline", - dimensions=15, - leidenResolution=0.5, - neighborsK=11, - useHarmony=False, -) -tuning = ParameterTuningAgent(FunctionModel(_tuning_reply)).run( - ds, - normalized=normalized, - candidates=[candidate], - experimental_handoff=tuning_handoff, - max_candidates=1, - max_refined_candidates=0, - min_cluster_cells=10, -) - -evaluation = tuning.evaluations[0] -{ - "status": tuning.status, - "recommended_candidate": tuning.recommendedCandidateId, - "eligible": evaluation.eligible, - "clusters": evaluation.metrics.nClusters, - "smallest_cluster": evaluation.metrics.minClusterCells, - "cluster_artifact": evaluation.artifacts["clusters"].artifactId, - "cell_selection": evaluation.cellSelection.artifactId, -} +result.plot_embedding(figsize=(9, 6)) ``` -## 6. Inspect one cluster with marker evidence - -Biological Interpretation consumes the exact selected cluster artifact. -Marker search is explicitly authorized because Parameter Tuning does not create a marker table. -The build reviews only the largest cluster and uses relaxed retrieval thresholds so the example remains bounded; the returned identity remains an unresolved hypothesis. - ```{code-cell} ipython3 -if tuning.status != "done": - raise RuntimeError(f"Parameter Tuning stopped with {tuning.status!r}") - -biology_handoff = tuning.to_biological_handoff() -biology = BiologicalInterpretationAgent(FunctionModel(_biology_reply)).run( - ds, - tuning_handoff=biology_handoff, - biological_context=BiologicalContext( - organism="Homo sapiens", - tissue="peripheral blood", - cellTypeReferences=["T cell", "B cell", "NK cell", "monocyte"], - experimentalDetails=["healthy donor PBMC", "no treatment contrast"], - ), - allow_marker_search=True, - marker_features=hvg_ref, - max_clusters=1, - max_markers=5, - marker_min_score=0.01, - marker_min_fraction=0.0, -) - -{ - "status": biology.status, - "interpretations": [ - { - "cluster": item.clusterId, - "identity": item.proposedIdentity, - "rationale": item.rationale, - "evidence": item.evidenceIds, - } - for item in biology.clusterInterpretations - ], - "tool_calls": [call.toolName for call in biology.runInfo.toolCalls], - "treatment_observations": len(biology.treatmentObservations), - "limitations": biology.limitations, -} +marker_table = result.get_markers() +marker_table.sort_values( + ["group_id", "score"], ascending=[True, False], +).groupby("group_id", sort=True).head(2)[ + ["group_id", "feature_name", "score", "frac_exp"] +].head(12) ``` -There are no replicated condition labels, so treatment observations remain empty. -Artifact provenance supports the computational chain, but a real biological identity still requires an appropriate model, study context, and independent validation. +These markers describe clusters. Replicate-aware differential expression and validated cell +identities require additional analysis. -## Use a live model - -Replace the four scripted models with one supported Pydantic AI model in an interactive analysis. -For an OpenAI-compatible endpoint, keep credentials in environment variables and never place them in a notebook or datastore: - -```python -import os - -from pydantic_ai.models.openai import OpenAIChatModel -from pydantic_ai.providers.openai import OpenAIProvider - -model = OpenAIChatModel( - os.environ["SCARF_AGENT_MODEL"], - provider=OpenAIProvider( - base_url=os.environ["SCARF_AGENT_BASE_URL"], - api_key=os.environ["SCARF_AGENT_API_KEY"], - ), -) - -enrichment = DataEnrichmentAgent(model).run( - ds, - context=DataEnrichmentContext(organismHint="human"), -) +```{code-cell} ipython3 +report_path = result.report() +{"report": report_path.name, "exists": report_path.is_file()} ``` -Use the same `model` for the other stages, retain the explicit handoffs, and set execution limits appropriate to the provider. -Model output remains provisional: Scarf's validators reject unknown evidence, but they cannot establish that a biologically plausible interpretation is true. +Open the returned HTML file to see the map, population counts, recorded comparisons, and +limitations together. For your own data, return to the model setup and `analyze_rna` call at the +top of this page. See {doc}`../reference/api/agent` for configuration and saved-history details, +or {doc}`../analysis_with_agents` for scientific reasoning guidance. diff --git a/pyproject.toml b/pyproject.toml index 0203f487..76af8f9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,7 @@ dependencies = [ [project.optional-dependencies] agent = [ - "pydantic-ai-slim[openai]>=2.0", + "pydantic-ai-slim[openai]>=2.19.0", ] extra = [ "anndata>=0.12", @@ -76,7 +76,7 @@ test = [ "anndata>=0.12", "scikit-network>=0.33.1", "rdata>=1.1.0", - "pydantic-ai-slim[openai]>=2.0", + "pydantic-ai-slim[openai]>=2.19.0", ] docs = [ "jinja2", diff --git a/scarf/agent/.env.example b/scarf/agent/.env.example deleted file mode 100644 index 9ead698e..00000000 --- a/scarf/agent/.env.example +++ /dev/null @@ -1,15 +0,0 @@ -# Copy to .env and fill in values: -# cp scarf/agent/.env.example scarf/agent/.env -# -# Local Ollama (no API key needed): -# OLLAMA_BASE_URL=http://localhost:11434/v1 -# OLLAMA_MODEL=qwen3.5:4b -# -# Ollama Cloud: -# OLLAMA_BASE_URL=https://ollama.com/v1 -# OLLAMA_API_KEY=your-key-here -# OLLAMA_MODEL=qwen3.5:4b - -OLLAMA_BASE_URL=http://localhost:11434/v1 -OLLAMA_MODEL=qwen3.5:4b -# OLLAMA_API_KEY= diff --git a/scarf/agent/__init__.py b/scarf/agent/__init__.py index b3aec4b9..27c425e4 100644 --- a/scarf/agent/__init__.py +++ b/scarf/agent/__init__.py @@ -1,173 +1,22 @@ -"""Optional grounded decision helpers for Scarf workflows.""" +"""Optional automated RNA analysis with a small, lazy public interface.""" -from .biological_interpretation import ( - BiologicalContext, - BiologicalInterpretationAgent, - BiologicalInterpretationReport, -) -from .characterize_covariates import ( - CovariateCharacterization, - characterize_covariates, -) -from .characterize_features import ( - FeatureCharacterization, - characterize_features, -) -from .config import AgentRunConfig -from .config import _deps as _deps -from .config.agent_exec import run_agent, run_agent_sync -from .data_enrichment import ( - DataEnrichmentAgent, - DataEnrichmentContext, - DataEnrichmentReport, - StudyContextSummary, -) -from .decide import DecisionValidationError, decide -from .experimental_context import ( - CellQcPlan, - ExperimentalContextAgent, - ExperimentalContextResult, - NamedArtifactSource, -) -from .ingest import IngestResult, detect_format, ingest -from .orchestrator import ( - AgentOrchestrator, - AssayPreprocessingPlan, - AutomatedPreprocessingPlan, - AutomatedWorkflowConfig, - AutomatedWorkflowRequest, - AutomatedWorkflowResult, - AutomatedWorkflowResumeRequest, - FinalAnalysisHandoff, - NativeAnalysisHandoff, - PreprocessedAssayHandoff, - WorkflowNeedsInput, - WorkflowQuestion, - WorkflowStageAttempt, - WorkflowStageLink, -) -from .parameter_tuning import ( - FinalGraphSelection, - IntegrationCandidateEvaluation, - IntegrationMetrics, - ParameterCandidate, - ParameterSearchPlan, - ParameterTuningAgent, - ParameterTuningAssayInput, - ParameterTuningReport, - get_default_parameter_candidates, - tune_parameters, -) -from .persistence import ( - AgentInvocation, - AgentName, - AgentPersistenceTarget, - AgentReport, - AgentReportLink, - AgentReportRecord, - AgentReportReference, - AgentReportType, - AgentTerminalStatus, - AgentWorkflowRun, - AgentWorkflowStatus, - create_agent_workflow, - finalize_agent_workflow, - list_agent_reports, - list_agent_workflows, - load_agent_record, - load_agent_report, - load_agent_workflow, - save_agent_report, -) -from .runtime import check_runtime, load_env -from .types import ( - BatchSafetyEvidence, - Decision, - EvidenceItem, - ExperimentalBiologyHandoff, - ExperimentalTuningHandoff, - NeedsInput, - StageResult, - StageStatus, - TuningBiologyHandoff, -) +from importlib import import_module +from typing import Any -__all__ = [ - "AgentInvocation", - "AgentOrchestrator", - "AgentRunConfig", - "AgentName", - "AgentPersistenceTarget", - "AgentReport", - "AgentReportLink", - "AgentReportRecord", - "AgentReportReference", - "AgentReportType", - "AgentTerminalStatus", - "AgentWorkflowRun", - "AgentWorkflowStatus", - "AssayPreprocessingPlan", - "AutomatedPreprocessingPlan", - "AutomatedWorkflowConfig", - "AutomatedWorkflowRequest", - "AutomatedWorkflowResult", - "AutomatedWorkflowResumeRequest", - "BatchSafetyEvidence", - "BiologicalContext", - "BiologicalInterpretationAgent", - "BiologicalInterpretationReport", - "CellQcPlan", - "CovariateCharacterization", - "DataEnrichmentAgent", - "DataEnrichmentContext", - "DataEnrichmentReport", - "Decision", - "DecisionValidationError", - "EvidenceItem", - "ExperimentalBiologyHandoff", - "ExperimentalContextAgent", - "ExperimentalContextResult", - "ExperimentalTuningHandoff", - "FinalAnalysisHandoff", - "FinalGraphSelection", - "FeatureCharacterization", - "IngestResult", - "IntegrationCandidateEvaluation", - "IntegrationMetrics", - "NativeAnalysisHandoff", - "NamedArtifactSource", - "NeedsInput", - "ParameterCandidate", - "ParameterSearchPlan", - "ParameterTuningAssayInput", - "ParameterTuningAgent", - "ParameterTuningReport", - "PreprocessedAssayHandoff", - "StageResult", - "StageStatus", - "StudyContextSummary", - "TuningBiologyHandoff", - "WorkflowNeedsInput", - "WorkflowQuestion", - "WorkflowStageAttempt", - "WorkflowStageLink", - "characterize_covariates", - "characterize_features", - "check_runtime", - "create_agent_workflow", - "decide", - "detect_format", - "get_default_parameter_candidates", - "ingest", - "load_env", - "finalize_agent_workflow", - "list_agent_reports", - "list_agent_workflows", - "load_agent_record", - "load_agent_report", - "load_agent_workflow", - "run_agent", - "run_agent_sync", - "save_agent_report", - "tune_parameters", -] +__all__ = ["analyze_rna", "AutomatedWorkflowResult", "AnalysisError"] + +for _export in __all__: + globals().pop(_export, None) + + +def __getattr__(name: str) -> Any: + if name not in __all__: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + module = ".orchestrator.api" if name == "analyze_rna" else ".orchestrator.models" + value = getattr(import_module(module, __name__), name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(__all__)) diff --git a/scarf/agent/config/_deps.py b/scarf/agent/_deps.py similarity index 100% rename from scarf/agent/config/_deps.py rename to scarf/agent/_deps.py diff --git a/scarf/agent/_plots.py b/scarf/agent/_plots.py new file mode 100644 index 00000000..1714e39c --- /dev/null +++ b/scarf/agent/_plots.py @@ -0,0 +1,282 @@ +"""Bounded views of the exact artifacts selected by an agent analysis.""" + +import hashlib +from collections import Counter +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any + +import numpy as np + +from ..graph.feature_projection import graph_cell_selection +from ..storage.refs import ArtifactRef +from ..storage.selections import validate_stored_selection_integrity +from ..storage.types import as_zarr_array + +if TYPE_CHECKING: + from ..datastore.datastore import DataStore + from ..plotting import PlotResult + + +DISPLAY_CELLS = 50_000 +DISPLAY_BLOCK_ROWS = 100_000 + + +def cluster_counts(store: "DataStore", clusters: ArtifactRef) -> dict[str, int]: + """Count every saved cluster label while retaining only one row block.""" + values = as_zarr_array(store.load_artifact(clusters)["values"], name="values") + if values.ndim != 1: + raise ValueError("Cluster labels must be one-dimensional") + counts: Counter[str] = Counter() + for start in range(0, values.shape[0], DISPLAY_BLOCK_ROWS): + labels, frequencies = np.unique( + np.asarray(values[start : start + DISPLAY_BLOCK_ROWS]).astype(str), + return_counts=True, + ) + counts.update(dict(zip(labels.tolist(), frequencies.tolist(), strict=True))) + return dict(sorted(counts.items())) + + +def _selection_input(store: "DataStore", ref: ArtifactRef) -> ArtifactRef: + status = store.inspect_artifact(ref) + if not status.complete: + raise ValueError("The selected analysis artifact is incomplete") + raw = (status.inputs or {}).get("cell_selection") + if not isinstance(raw, Mapping): + raise ValueError("The selected artifact has no frozen cell selection") + return ArtifactRef.from_dict(raw) + + +def _sample_quotas(counts: Mapping[str, int], maximum: int) -> dict[str, int]: + """Allocate a proportional display, keeping at least one cell per cluster.""" + total = sum(counts.values()) + if total <= maximum: + return dict(counts) + if len(counts) > maximum: + raise RuntimeError("There are more clusters than the display cell limit") + remaining = maximum - len(counts) + available = total - len(counts) + shares = { + label: remaining * (count - 1) / available for label, count in counts.items() + } + quotas = {label: 1 + int(share) for label, share in shares.items()} + remainder = maximum - sum(quotas.values()) + for label in sorted(shares, key=lambda key: (-(shares[key] % 1), key))[:remainder]: + quotas[label] += 1 + return quotas + + +def _sample_cluster_rows( + values: Any, + counts: Mapping[str, int], + *, + maximum: int, + seed: int, + block_rows: int = DISPLAY_BLOCK_ROWS, +) -> tuple[np.ndarray, np.ndarray]: + """Keep bounded per-cluster priority samples, independent of read boundaries.""" + quotas = _sample_quotas(counts, maximum) + rng = np.random.Generator(np.random.PCG64(seed)) + priorities = {label: np.empty(0, dtype=np.float64) for label in counts} + rows = {label: np.empty(0, dtype=np.int64) for label in counts} + for start in range(0, int(values.shape[0]), block_rows): + labels = np.asarray(values[start : start + block_rows]).astype(str) + keys = rng.random(len(labels)) + for label in np.unique(labels): + positions = np.flatnonzero(labels == label) + combined_keys = np.concatenate((priorities[label], keys[positions])) + combined_rows = np.concatenate((rows[label], positions + start)) + quota = quotas[label] + if len(combined_rows) > quota: + keep = np.argpartition(combined_keys, quota - 1)[:quota] + combined_keys = combined_keys[keep] + combined_rows = combined_rows[keep] + priorities[label] = combined_keys + rows[label] = combined_rows + selected = np.concatenate(list(rows.values())) + labels = np.concatenate( + [np.full(len(rows[label]), label, dtype=object) for label in counts] + ) + order = np.argsort(selected) + return selected[order], labels[order] + + +def plot_final_umap( + store: "DataStore", + *, + umap: ArtifactRef, + clusters: ArtifactRef, + cell_selection: ArtifactRef, + graph: ArtifactRef, + max_points: int = DISPLAY_CELLS, + seed: int = 0, + figsize: tuple[float, float] = (9, 6), + show: bool = True, +) -> "PlotResult": + """Draw a bounded categorical map from one completed analysis. + + Population counts always include every cell. Display sampling changes no + saved selection, coordinates, labels, markers, or other analysis result. + """ + from ..plotting import CategoricalScale, LegendSpec, PlotProvenance, PlotResult + + if ( + isinstance(max_points, bool) + or not isinstance(max_points, int) + or not 1 <= max_points <= DISPLAY_CELLS + ): + raise ValueError(f"max_points must be an integer from 1 to {DISPLAY_CELLS}") + if umap.kind != "embedding" or clusters.kind != "cluster_labels": + raise ValueError("The final map requires embedding and cluster-label artifacts") + if umap.assay != clusters.assay or umap.scope != clusters.scope: + raise ValueError("The final UMAP and clusters must belong to the same assay") + for ref in (umap, clusters): + if _selection_input(store, ref) != cell_selection: + raise ValueError( + "Final artifacts must share the exact frozen cell selection" + ) + if graph_cell_selection(store.zw, graph) != cell_selection: + raise ValueError("The final graph must use the exact frozen cell selection") + embedding_status = store.inspect_artifact(umap) + if embedding_status.operation != "run_umap": + raise ValueError("The final UMAP must be a saved run_umap artifact") + for ref in (umap, clusters): + raw_graph = (store.inspect_artifact(ref).inputs or {}).get("graph") + if ( + not isinstance(raw_graph, Mapping) + or ArtifactRef.from_dict(raw_graph) != graph + ): + raise ValueError("The final UMAP and clusters must use the selected graph") + selection = validate_stored_selection_integrity( + store.zw, + cell_selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ) + coordinates = as_zarr_array(store.load_artifact(umap)["values"], name="values") + labels = as_zarr_array(store.load_artifact(clusters)["values"], name="values") + total = selection.selected_count + if coordinates.shape != (total, 2) or labels.shape != (total,): + raise ValueError( + "Final coordinates and labels must align with the frozen selection" + ) + if np.dtype(coordinates.dtype).kind not in "fiu": + raise TypeError("Final UMAP coordinates must be numeric") + if total == 0: + raise ValueError("The final selection contains no cells") + counts = cluster_counts(store, clusters) + rows, sampled_labels = _sample_cluster_rows( + labels, counts, maximum=max_points, seed=seed + ) + sampled_coordinates = np.empty((len(rows), 2), dtype=np.float64) + for start in range(0, total, DISPLAY_BLOCK_ROWS): + left, right = np.searchsorted(rows, [start, start + DISPLAY_BLOCK_ROWS]) + if left == right: + continue + block = np.asarray(coordinates[start : start + DISPLAY_BLOCK_ROWS]) + if not np.isfinite(block).all(): + raise ValueError("Final UMAP coordinates must be finite") + sampled_coordinates[left:right] = block[rows[left:right] - start] + + import matplotlib.pyplot as plt + import pandas as pd + from matplotlib.colors import to_hex + from matplotlib.lines import Line2D + + categories = tuple(counts) + cmap = plt.get_cmap( + "tab10" + if len(categories) <= 10 + else "tab20" + if len(categories) <= 20 + else "gist_rainbow" + ) + palette = { + label: to_hex( + cmap(index if len(categories) <= 20 else index / (len(categories) - 1)) + ) + for index, label in enumerate(categories) + } + figure, axis = plt.subplots(figsize=figsize, constrained_layout=True) + try: + order = np.random.Generator(np.random.PCG64(seed)).permutation(len(rows)) + axis.scatter( + sampled_coordinates[order, 0], + sampled_coordinates[order, 1], + color=[palette[label] for label in sampled_labels[order]], + s=3, + alpha=0.65, + linewidths=0, + rasterized=True, + ) + axis.set(xlabel="UMAP 1", ylabel="UMAP 2", xticks=[], yticks=[]) + axis.spines[["top", "right"]].set_visible(False) + if len(categories) <= 40: + axis.legend( + handles=[ + Line2D( + [], + [], + color=palette[label], + marker="o", + linestyle="", + label=f"{label} ({counts[label]:,})", + ) + for label in categories + ], + title="Cluster (all cells)", + loc="center left", + bbox_to_anchor=(1, 0.5), + frameon=False, + fontsize=8, + ) + else: + for label in categories: + center = np.median(sampled_coordinates[sampled_labels == label], axis=0) + axis.annotate(str(label), center, fontsize=7) + caption = ( + f"{len(rows):,} of {total:,} cells shown; cluster counts use all cells." + ) + axis.set_title(caption, fontsize=10) + result = PlotResult( + figure=figure, + axes={"clusters": axis}, + tables={ + "cluster_counts": pd.DataFrame( + {"cluster": list(counts), "cells": list(counts.values())} + ) + }, + legends=(LegendSpec(kind="categorical", label="Cluster (all cells)"),), + scales=(CategoricalScale(order=categories, palette=palette),), + provenance=PlotProvenance( + assay=umap.assay, + n_cells=len(rows), + notes=("Final saved UMAP; display sampling only.",), + extras={ + "layout": umap.to_dict(), + "color_artifacts": [clusters.to_dict()], + "cell_selection": cell_selection.to_dict(), + "graph": graph.to_dict(), + "workspace": store.workspace, + "input_n_cells": total, + "cluster_counts": counts, + "display_sampling": { + "method": "proportional_clusters_with_minimum_one", + "max_points": max_points, + "seed": seed, + "generator": "PCG64", + "row_sha256": hashlib.sha256( + rows.astype(" "BiologicalContext": - return cls() - - @classmethod - def get_example(cls) -> "BiologicalContext": - return cls( - organism="Homo sapiens", - studyContext=( - "Human lung samples were profiled after drug or vehicle treatment." - ), - tissue="lung", - cellTypeReferences=["alveolar macrophage", "T cell"], - experimentalDetails=["drug and vehicle groups"], - treatmentQuestion="Which populations respond selectively to treatment?", - ) - - -class ConditionClusterSummary(AgentDataModel): - """Aggregate cluster abundance for one condition without sample identifiers.""" - - condition: str = "" - clusterId: str = "" - nSamples: int = 0 - meanFraction: float = 0.0 - minFraction: float = 0.0 - maxFraction: float = 0.0 - cellCount: int = 0 - evidenceId: str = "" - - @classmethod - def get_example(cls) -> "ConditionClusterSummary": - return cls( - condition="treated", - clusterId="3", - nSamples=4, - meanFraction=0.18, - minFraction=0.12, - maxFraction=0.25, - cellCount=180, - evidenceId="composition:RNA_cluster:condition:treated:cluster:3", - ) - - -class ClusterCompositionEvidence(AgentDataModel): - """Bounded deterministic evidence about cluster sizes and conditions.""" - - clusterArtifact: ArtifactReferenceModel | None = None - cellSelection: ArtifactReferenceModel | None = None - totalCells: int = 0 - clusterCounts: dict[str, int] = Field(default_factory=dict) - sampleColumn: str | None = None - conditionColumn: str | None = None - conditionSummaries: list[ConditionClusterSummary] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - warnings: list[str] = Field(default_factory=list) - - @classmethod - def get_example(cls) -> "ClusterCompositionEvidence": - summary = ConditionClusterSummary.get_example() - reference_summary = ConditionClusterSummary( - condition="control", - clusterId=summary.clusterId, - nSamples=4, - meanFraction=0.11, - minFraction=0.08, - maxFraction=0.15, - cellCount=110, - evidenceId="composition:RNA_cluster:condition:control:cluster:3", - ) - return cls( - clusterArtifact=ArtifactReferenceModel( - assay="RNA", - kind="cluster_labels", - artifactId="b" * 64, - ), - cellSelection=ArtifactReferenceModel( - scope="datastore", - assay=None, - kind="cell_selection", - artifactId="c" * 64, - ), - totalCells=1000, - clusterCounts={"0": 520, "1": 300, "3": 180}, - sampleColumn="sample", - conditionColumn="treatment", - conditionSummaries=[reference_summary, summary], - evidenceIds=[ - "composition:RNA_cluster:counts", - reference_summary.evidenceId, - summary.evidenceId, - ], - ) - - -class MarkerFeature(AgentDataModel): - """One observed marker feature and its available Scarf statistics.""" - - featureId: str = "" - featureName: str = "" - featureIndex: int | None = None - score: float | None = None - foldChange: float | None = None - fractionExpressed: float | None = None - fractionExpressedRest: float | None = None - mean: float | None = None - meanRest: float | None = None - auc: float | None = None - adjustedPvalue: float | None = None - - @classmethod - def get_example(cls) -> "MarkerFeature": - return cls( - featureId="ENSG00000173372", - featureName="C1QA", - featureIndex=123, - score=0.83, - foldChange=3.4, - fractionExpressed=0.76, - fractionExpressedRest=0.18, - auc=0.91, - adjustedPvalue=0.001, - ) - - -class ClusterMarkerEvidence(AgentDataModel): - """Bounded markers for one exact cluster label.""" - - clusterId: str = "" - markers: list[MarkerFeature] = Field(default_factory=list) - markerArtifact: ArtifactReferenceModel | None = None - evidenceId: str = "" - warnings: list[str] = Field(default_factory=list) - - @classmethod - def get_example(cls) -> "ClusterMarkerEvidence": - return cls( - clusterId="3", - markers=[MarkerFeature.get_example()], - markerArtifact=ArtifactReferenceModel( - assay="RNA", - kind="marker_table", - artifactId="a" * 64, - ), - evidenceId="markers:RNA_cluster:cluster:3", - ) - - -class ClusterMarkerBatchEvidence(AgentDataModel): - """Markers for all model-selected clusters returned by one tool call.""" - - clusters: list[ClusterMarkerEvidence] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - warnings: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "ClusterMarkerBatchEvidence": - return cls() - - @classmethod - def get_example(cls) -> "ClusterMarkerBatchEvidence": - cluster = ClusterMarkerEvidence.get_example() - return cls(clusters=[cluster], evidenceIds=[cluster.evidenceId]) - - -class ClusterInterpretation(AgentDataModel): - """One evidence-linked cluster interpretation or hypothesis.""" - - clusterId: str = "" - proposedIdentity: str = "unresolved" - identityIsHypothesis: bool = True - confidence: InterpretationConfidence = "low" - rationale: str = "" - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_example(cls) -> "ClusterInterpretation": - return cls( - clusterId="3", - proposedIdentity="alveolar macrophage-like", - identityIsHypothesis=True, - confidence="medium", - rationale="Observed marker pattern is consistent with the proposed identity.", - evidenceIds=["markers:RNA_cluster:cluster:3"], - ) - - -class TreatmentObservation(AgentDataModel): - """Descriptive treatment observation with no unsupported causal claim.""" - - clusterId: str = "" - referenceCondition: str = "" - comparisonCondition: str = "" - direction: TreatmentDirection = "equal" - observation: str = "" - isDescriptiveOnly: Literal[True] = True - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_example(cls) -> "TreatmentObservation": - return cls( - clusterId="3", - referenceCondition="control", - comparisonCondition="treated", - direction="higher", - observation="Cluster 3 has a higher mean fraction in treated samples.", - evidenceIds=[ - "composition:RNA_cluster:condition:control:cluster:3", - "composition:RNA_cluster:condition:treated:cluster:3", - ], - ) - - -class FollowUpRecommendation(AgentDataModel): - """A bounded next analysis tied to an observed uncertainty.""" - - question: str = "" - operation: str = "" - rationale: str = "" - requiredInputs: list[str] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_example(cls) -> "FollowUpRecommendation": - return cls( - question="Is the abundance difference reproducible across donors?", - operation="sample-level differential abundance", - rationale="Current evidence is descriptive and requires independent replicates.", - requiredInputs=["sample", "condition", "donor"], - evidenceIds=[ - "composition:RNA_cluster:condition:control:cluster:3", - "composition:RNA_cluster:condition:treated:cluster:3", - ], - ) - - -class BiologicalInterpretationNeedsInput(AgentDataModel): - question: str = "" - requiredInputs: list[str] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_example(cls) -> "BiologicalInterpretationNeedsInput": - return cls( - question="Provide an exact marker artifact or authorize marker search.", - requiredInputs=["markerArtifact"], - ) - - -class BiologicalInterpretationReport(AgentDataModel): - """Structured, evidence-grounded biological review.""" - - status: StageStatus = "needsInput" - clusterInterpretations: list[ClusterInterpretation] = Field(default_factory=list) - treatmentObservations: list[TreatmentObservation] = Field(default_factory=list) - followUps: list[FollowUpRecommendation] = Field(default_factory=list) - clusterArtifact: ArtifactReferenceModel | None = None - markerArtifact: ArtifactReferenceModel | None = None - graphAssay: str | None = None - markerAssay: str | None = None - evidenceIds: list[str] = Field(default_factory=list) - limitations: list[str] = Field(default_factory=list) - stopReason: str = "" - needsInput: BiologicalInterpretationNeedsInput | None = None - runInfo: AgentRunInfo = Field(default_factory=AgentRunInfo) - - @classmethod - def get_example(cls) -> "BiologicalInterpretationReport": - interpretation = ClusterInterpretation.get_example() - observation = TreatmentObservation.get_example() - follow_up = FollowUpRecommendation.get_example() - return cls( - status="done", - clusterInterpretations=[interpretation], - treatmentObservations=[observation], - followUps=[follow_up], - clusterArtifact=ClusterCompositionEvidence.get_example().clusterArtifact, - markerArtifact=ClusterMarkerEvidence.get_example().markerArtifact, - graphAssay="RNA", - markerAssay="RNA", - evidenceIds=sorted( - { - *interpretation.evidenceIds, - *observation.evidenceIds, - *follow_up.evidenceIds, - } - ), - limitations=[ - "Cell identities remain hypotheses until independently validated." - ], - stopReason="The requested clusters were reviewed.", - ) - - -class BiologicalInterpretationDependencies(AgentDataModel): - """Runtime state available only to biological interpretation tools.""" - - store: Any = Field(default=None, exclude=True) - cluster: Any = Field(default=None, exclude=True) - cellSelection: Any = Field(default=None, exclude=True) - cellIndices: Any = Field(default=None, exclude=True) - fromAssay: str | None = None - graphAssay: str | None = None - markerAssay: str | None = None - markerAssayType: str | None = None - sampleColumn: str | None = None - conditionColumn: str | None = None - marker: Any = Field(default=None, exclude=True) - markerFeatures: Any = Field(default=None, exclude=True) - allowMarkerSearch: bool = False - maxClusters: int = 12 - maxMarkers: int = 10 - markerMinScore: float = 0.25 - markerMinFraction: float = 0.2 - evidenceIds: set[str] = Field(default_factory=set, exclude=True) - clusterValues: dict[str, Any] = Field(default_factory=dict, exclude=True) - markerEvidenceIds: dict[str, str] = Field(default_factory=dict, exclude=True) - toolCalls: list[str] = Field(default_factory=list, exclude=True) - conditionEvidence: dict[str, ConditionClusterSummary] = Field( - default_factory=dict, - exclude=True, - ) - designHandoff: ExperimentalBiologyHandoff | None = Field( - default=None, - exclude=True, - ) - - @classmethod - def get_example(cls) -> "BiologicalInterpretationDependencies": - return cls( - cluster=object(), - fromAssay="RNA", - graphAssay="RNA", - markerAssay="RNA", - markerAssayType="RNA", - sampleColumn="sample", - conditionColumn="treatment", - ) - - -def _prepare_biological_interpretation_tool( - ctx: RunContext[BiologicalInterpretationDependencies], - tool_definition: ToolDefinition, -) -> ToolDefinition | None: - """Expose composition and marker batches once and in the required order.""" - completed_calls = set(ctx.deps.toolCalls) - if tool_definition.name == "inspect_cluster_composition": - return None if tool_definition.name in completed_calls else tool_definition - if tool_definition.name == "inspect_cluster_markers_batch": - if ( - "inspect_cluster_composition" not in completed_calls - or tool_definition.name in completed_calls - ): - return None - return tool_definition - return tool_definition - - -_SYSTEM_PROMPT = dedent( - """ - You are Scarf's Biological Interpretation Agent. Use only the supplied - tools and caller context. Call inspect_cluster_composition exactly once. - Then select every cluster you intend to interpret and call - inspect_cluster_markers_batch exactly once with all selected cluster IDs. - Tool calls execute Scarf operations, so wait for their results before - drawing conclusions. Do not split marker inspection across calls. - - Treat cell identities as hypotheses unless the caller supplied a trusted - label. Do not invent genes, cell types, statistics, artifact identifiers, - or evidence identifiers. Cite only evidenceIds returned by tools. For each - cluster interpretation, copy the exact non-empty marker evidenceId returned - for that cluster into its evidenceIds. Do not interpret a cluster whose - marker evidenceId is empty. Cluster abundance summaries are descriptive, - not tests of significance or causal effects. Treatment observations must - compare two returned independent-unit condition summaries for the same - cluster. Independent units may occur in more than one condition in paired - or repeated-measure designs. Return treatmentObservations empty unless the - exact experimental handoff confirms independent-unit aggregation, a - between-unit coefficient, an estimable coefficient, and at least two - independent units in each cited condition. - Marker p-values describe cluster-versus-rest marker specificity, not - condition effects. Keep treatment content out of cluster identity - interpretations. Recommend a named follow-up operation when replication, a - covariate, or an exact artifact is missing. Do not write exploratory code, - use a shell, access files, or call arbitrary Scarf methods. Return only - fields defined by the structured output schema. - """ -).strip() - - -def _string_value(value: Any) -> str: - return str(value.item() if isinstance(value, np.generic) else value) - - -def _finite_float(value: Any) -> float | None: - if value is None: - return None - try: - number = float(value) - except (TypeError, ValueError): - return None - return number if np.isfinite(number) else None - - -def _check_column(store: Any, column: str, label: str) -> None: - if column not in set(store.cells.columns): - raise ValueError(f"{label} {column!r} is not present in cell metadata") - - -async def inspect_cluster_composition( - ctx: RunContext[BiologicalInterpretationDependencies], -) -> ClusterCompositionEvidence: - """Inspect bounded cluster and condition composition without identifiers.""" - deps = ctx.deps - logger.info( - f"Inspecting cluster composition from artifact " - f"{getattr(deps.cluster, 'artifact_id', '')!r}; " - f"max_clusters={deps.maxClusters}" - ) - if deps.sampleColumn is not None: - _check_column(deps.store, deps.sampleColumn, "sample column") - if deps.conditionColumn is not None: - _check_column(deps.store, deps.conditionColumn, "condition column") - - if deps.cluster is None: - raise ValueError("An exact cluster artifact is required") - deps.cluster = core_artifact_reference(deps.cluster) - cluster_artifact = artifact_reference(deps.cluster) - if cluster_artifact.kind not in {"cluster_labels", "cluster_cut"}: - raise ValueError( - "cluster must identify a cluster_labels or cluster_cut artifact" - ) - if ( - deps.graphAssay is not None - and cluster_artifact.scope == "assay" - and cluster_artifact.assay != deps.graphAssay - ): - raise ValueError("cluster artifact belongs to a different assay") - if cluster_artifact.scope == "datastore" and cluster_artifact.assay is not None: - raise ValueError("datastore-scoped cluster artifacts must not name an assay") - status = deps.store.inspect_artifact(deps.cluster) - if not getattr(status, "exists", True): - raise ValueError("cluster artifact does not exist") - if not getattr(status, "complete", False): - raise ValueError("cluster artifact is incomplete") - inputs = getattr(status, "inputs", None) or {} - raw_selection = inputs.get("cell_selection") - if not isinstance(raw_selection, Mapping): - raise ValueError("cluster artifact has no cell-selection input") - cell_selection = ArtifactRef.from_dict(dict(raw_selection)) - if ( - cell_selection.scope != "datastore" - or cell_selection.kind != "cell_selection" - or cell_selection.assay is not None - ): - raise ValueError("cluster artifact has an invalid cell-selection input") - if ( - deps.cellSelection is not None - and core_artifact_reference(deps.cellSelection) != cell_selection - ): - raise ValueError( - "cluster artifact cell selection conflicts with the prepared selection" - ) - cell_indices = read_stored_selection_indices( - deps.store.zw, - cell_selection, - kind="cell_selection", - scope="datastore", - assay=None, - table_path="cellData", - ).astype(np.int64, copy=False) - deps.cellSelection = cell_selection - deps.cellIndices = cell_indices - cluster_group = deps.store.load_artifact(deps.cluster) - value_name = "labels" if cluster_artifact.kind == "cluster_cut" else "values" - if value_name not in cluster_group: - raise ValueError( - f"cluster artifact does not contain its {value_name!r} label vector" - ) - cluster_values = np.asarray(cluster_group[value_name][:]) - if cluster_values.ndim != 1 or len(cluster_values) != len(cell_indices): - raise ValueError("cluster artifact labels do not align with its cell selection") - if len(cluster_values) == 0: - raise ValueError("cluster artifact selects no cells") - counts = Counter(_string_value(value) for value in cluster_values) - ordered_clusters = sorted(counts, key=lambda value: (-counts[value], value)) - retained_clusters = ordered_clusters[: deps.maxClusters] - deps.clusterValues = { - _string_value(value): value.item() if isinstance(value, np.generic) else value - for value in cluster_values - if _string_value(value) in retained_clusters - } - evidence_prefix = f"composition:{cluster_artifact.artifactId}" - count_evidence = f"{evidence_prefix}:counts" - deps.evidenceIds.add(count_evidence) - warnings: list[str] = [] - if len(ordered_clusters) > deps.maxClusters: - warnings.append(f"Only the {deps.maxClusters} largest clusters were returned.") - - condition_summaries: list[ConditionClusterSummary] = [] - if deps.conditionColumn is not None: - condition_values = read_metadata_rows( - deps.store.cells, - deps.conditionColumn, - cell_indices, - ) - if len(condition_values) != len(cluster_values): - raise ValueError("condition and cluster columns are not aligned") - condition_missing = read_metadata_missing_rows( - deps.store.cells, - deps.conditionColumn, - cell_indices, - ) - if condition_missing is not None and np.any(condition_missing): - raise ValueError("condition column contains missing selected values") - n_conditions = len({_string_value(value) for value in condition_values}) - if n_conditions > CONFIG._MAX_CONDITIONS: - warnings.append( - f"Only the first {CONFIG._MAX_CONDITIONS} conditions were returned." - ) - if deps.sampleColumn is not None: - sample_values = read_metadata_rows( - deps.store.cells, - deps.sampleColumn, - cell_indices, - ) - if len(sample_values) != len(cluster_values): - raise ValueError("sample and cluster columns are not aligned") - sample_missing = read_metadata_missing_rows( - deps.store.cells, - deps.sampleColumn, - cell_indices, - ) - if sample_missing is not None and np.any(sample_missing): - raise ValueError("sample column contains missing selected values") - summaries = _sample_condition_summaries( - sample_values=sample_values, - condition_values=condition_values, - cluster_values=cluster_values, - retained_clusters=retained_clusters, - evidence_prefix=evidence_prefix, - ) - else: - summaries = _cell_condition_summaries( - condition_values=condition_values, - cluster_values=cluster_values, - retained_clusters=retained_clusters, - evidence_prefix=evidence_prefix, - ) - warnings.append( - "No sample column was supplied; condition fractions are cell-level summaries." - ) - condition_summaries = summaries[ - : CONFIG._MAX_CONDITIONS * len(retained_clusters) - ] - deps.evidenceIds.update(summary.evidenceId for summary in condition_summaries) - deps.conditionEvidence.update( - {summary.evidenceId: summary for summary in condition_summaries} - ) - - deps.toolCalls.append("inspect_cluster_composition") - evidence = ClusterCompositionEvidence( - clusterArtifact=cluster_artifact, - cellSelection=artifact_reference(cell_selection), - totalCells=len(cluster_values), - clusterCounts={cluster: counts[cluster] for cluster in retained_clusters}, - sampleColumn=deps.sampleColumn, - conditionColumn=deps.conditionColumn, - conditionSummaries=condition_summaries, - evidenceIds=sorted(deps.evidenceIds), - warnings=warnings, - ) - logger.info( - f"Completed cluster composition inspection: cells={evidence.totalCells}, " - f"clusters={len(evidence.clusterCounts)}, " - f"condition_summaries={len(evidence.conditionSummaries)}, " - f"warnings={len(evidence.warnings)}" - ) - return evidence - - -def _sample_condition_summaries( - *, - sample_values: np.ndarray, - condition_values: np.ndarray, - cluster_values: np.ndarray, - retained_clusters: list[str], - evidence_prefix: str, -) -> list[ConditionClusterSummary]: - """Aggregate each condition-unit pair without exposing unit identifiers.""" - sample_counts: dict[tuple[str, str], Counter[str]] = defaultdict(Counter) - sample_totals: Counter[tuple[str, str]] = Counter() - for sample, condition, cluster in zip( - sample_values, - condition_values, - cluster_values, - strict=True, - ): - condition_label = _string_value(condition) - sample_label = _string_value(sample) - key = (condition_label, sample_label) - sample_totals[key] += 1 - sample_counts[key][_string_value(cluster)] += 1 - - fractions: dict[tuple[str, str], list[float]] = defaultdict(list) - cell_counts: Counter[tuple[str, str]] = Counter() - for key, total in sample_totals.items(): - condition, _sample = key - for cluster in retained_clusters: - count = sample_counts[key][cluster] - fractions[(condition, cluster)].append(count / total) - cell_counts[(condition, cluster)] += count - - output: list[ConditionClusterSummary] = [] - for condition, cluster in sorted(fractions): - values = fractions[(condition, cluster)] - evidence_id = f"{evidence_prefix}:condition:{condition}:cluster:{cluster}" - output.append( - ConditionClusterSummary( - condition=condition, - clusterId=cluster, - nSamples=len(values), - meanFraction=float(np.mean(values)), - minFraction=float(np.min(values)), - maxFraction=float(np.max(values)), - cellCount=cell_counts[(condition, cluster)], - evidenceId=evidence_id, - ) - ) - return output - - -def _cell_condition_summaries( - *, - condition_values: np.ndarray, - cluster_values: np.ndarray, - retained_clusters: list[str], - evidence_prefix: str, -) -> list[ConditionClusterSummary]: - totals: Counter[str] = Counter(_string_value(value) for value in condition_values) - counts: Counter[tuple[str, str]] = Counter() - for condition, cluster in zip(condition_values, cluster_values, strict=True): - counts[(_string_value(condition), _string_value(cluster))] += 1 - output: list[ConditionClusterSummary] = [] - for condition in sorted(totals): - for cluster in retained_clusters: - count = counts[(condition, cluster)] - fraction = count / totals[condition] - evidence_id = f"{evidence_prefix}:condition:{condition}:cluster:{cluster}" - output.append( - ConditionClusterSummary( - condition=condition, - clusterId=cluster, - meanFraction=fraction, - minFraction=fraction, - maxFraction=fraction, - cellCount=count, - evidenceId=evidence_id, - ) - ) - return output - - -async def inspect_cluster_markers( - ctx: RunContext[BiologicalInterpretationDependencies], - cluster_id: str, -) -> ClusterMarkerEvidence: - """Load markers for one observed cluster, optionally creating one artifact.""" - deps = ctx.deps - logger.debug(f"Inspecting markers for cluster {cluster_id!r}") - if not deps.clusterValues: - raise ModelRetry("Call inspect_cluster_composition before inspecting markers.") - if cluster_id not in deps.clusterValues: - raise ModelRetry(f"cluster_id must be one of {sorted(deps.clusterValues)}") - if deps.marker is None: - if not deps.allowMarkerSearch: - logger.warning( - f"Markers for cluster {cluster_id!r} are unavailable because no " - "marker artifact was supplied or authorized" - ) - return ClusterMarkerEvidence( - clusterId=cluster_id, - evidenceId="", - warnings=[ - "No exact marker artifact was supplied and marker search was not authorized." - ], - ) - if deps.markerFeatures is None: - logger.warning( - f"Markers for cluster {cluster_id!r} are unavailable because no " - "feature selection was supplied" - ) - return ClusterMarkerEvidence( - clusterId=cluster_id, - evidenceId="", - warnings=["Marker search requires an exact feature selection."], - ) - logger.info( - f"Creating one marker artifact for assay {deps.markerAssay!r} " - f"from cluster artifact {deps.cluster.artifact_id!r}" - ) - deps.marker = deps.store.run_marker_search( - deps.cluster, - features=deps.markerFeatures, - ) - if not hasattr(deps.marker, "artifact_id"): - raise RuntimeError("marker search did not return an artifact reference") - deps.marker = core_artifact_reference(deps.marker) - - marker_artifact = artifact_reference(deps.marker) - if marker_artifact.kind != "marker_table": - raise ModelRetry("marker must identify a marker_table artifact") - if deps.markerAssay is not None and marker_artifact.assay != deps.markerAssay: - raise ModelRetry("marker artifact belongs to a different assay") - if hasattr(deps.store, "inspect_artifact"): - marker_status = deps.store.inspect_artifact(deps.marker) - if not getattr(marker_status, "exists", True): - raise ModelRetry("marker artifact does not exist") - if not getattr(marker_status, "complete", False): - raise ModelRetry("marker artifact is incomplete") - marker_inputs = getattr(marker_status, "inputs", None) or {} - stored_clusters = marker_inputs.get("clusters") - expected_cluster = artifact_reference(deps.cluster) - if ( - not isinstance(stored_clusters, Mapping) - or stored_clusters.get("artifact_id") != expected_cluster.artifactId - or stored_clusters.get("kind") != expected_cluster.kind - or stored_clusters.get("scope") != expected_cluster.scope - or stored_clusters.get("assay") != expected_cluster.assay - ): - raise ModelRetry( - "marker artifact is not linked to the exact cluster artifact" - ) - - frame = deps.store.get_markers( - deps.marker, - group_id=deps.clusterValues[cluster_id], - min_score=deps.markerMinScore, - min_frac_exp=deps.markerMinFraction, - ) - if "score" in frame.columns: - frame = frame.sort_values("score", ascending=False, na_position="last") - markers = [ - _marker_feature(row) - for row in frame.head(min(deps.maxMarkers, CONFIG._MAX_MARKERS)).to_dict( - "records" - ) - ] - cluster_artifact = artifact_reference(deps.cluster) - evidence_id = ( - f"markers:{marker_artifact.artifactId}:clusters:" - f"{cluster_artifact.artifactId}:cluster:{cluster_id}" - ) - if markers: - deps.evidenceIds.add(evidence_id) - deps.markerEvidenceIds[cluster_id] = evidence_id - evidence = ClusterMarkerEvidence( - clusterId=cluster_id, - markers=markers, - markerArtifact=marker_artifact, - evidenceId=evidence_id if markers else "", - warnings=[] if markers else ["No markers passed the requested thresholds."], - ) - logger.debug( - f"Completed marker inspection for cluster {cluster_id!r}: " - f"markers={len(markers)}" - ) - return evidence - - -async def inspect_cluster_markers_batch( - ctx: RunContext[BiologicalInterpretationDependencies], - cluster_ids: list[str], -) -> ClusterMarkerBatchEvidence: - """Inspect every selected cluster in one bounded model tool call.""" - if not ctx.deps.clusterValues: - raise ModelRetry("Call inspect_cluster_composition before inspecting markers.") - if not cluster_ids: - raise ModelRetry("cluster_ids must contain at least one observed cluster") - if len(cluster_ids) > ctx.deps.maxClusters: - raise ModelRetry( - f"cluster_ids may contain at most {ctx.deps.maxClusters} values" - ) - if len(set(cluster_ids)) != len(cluster_ids): - raise ModelRetry("cluster_ids must not contain duplicates") - - logger.info(f"Inspecting markers for {len(cluster_ids)} cluster(s) in one batch") - clusters = [ - await inspect_cluster_markers(ctx, cluster_id=cluster_id) - for cluster_id in cluster_ids - ] - evidence_ids = [cluster.evidenceId for cluster in clusters if cluster.evidenceId] - warnings = [ - f"Cluster {cluster.clusterId}: {warning}" - for cluster in clusters - for warning in cluster.warnings - ] - ctx.deps.toolCalls.append("inspect_cluster_markers_batch") - evidence = ClusterMarkerBatchEvidence( - clusters=clusters, - evidenceIds=evidence_ids, - warnings=warnings, - ) - logger.info( - f"Completed marker batch inspection: clusters={len(clusters)}, " - f"clusters_with_markers={sum(bool(cluster.markers) for cluster in clusters)}, " - f"evidence_records={len(evidence_ids)}" - ) - return evidence - - -def _marker_feature(row: dict[str, Any]) -> MarkerFeature: - raw_index = _finite_float(row.get("feature_index")) - return MarkerFeature( - featureId=str(row.get("feature_id", "")), - featureName=str(row.get("feature_name", "")), - featureIndex=int(raw_index) if raw_index is not None else None, - score=_finite_float(row.get("score")), - foldChange=_finite_float(row.get("fold_change")), - fractionExpressed=_finite_float(row.get("frac_exp")), - fractionExpressedRest=_finite_float(row.get("frac_exp_rest")), - mean=_finite_float(row.get("mean")), - meanRest=_finite_float(row.get("mean_rest")), - auc=_finite_float(row.get("auc")), - adjustedPvalue=_finite_float(row.get("p_value_adjusted")), - ) - - -def _canonicalize_cluster_interpretations( - report: BiologicalInterpretationReport, - deps: BiologicalInterpretationDependencies, -) -> tuple[list[ClusterInterpretation], list[str]]: - canonical_interpretations: list[ClusterInterpretation] = [] - omitted_interpretation_clusters: list[str] = [] - for interpretation in report.clusterInterpretations: - marker_id = deps.markerEvidenceIds.get(interpretation.clusterId) - if marker_id is None: - omitted_interpretation_clusters.append(interpretation.clusterId) - continue - non_marker_evidence = sorted(set(interpretation.evidenceIds) - {marker_id}) - if non_marker_evidence: - raise ModelRetry( - "Cluster identity interpretations may cite only their exact marker " - f"evidence: {non_marker_evidence}" - ) - canonical_interpretations.append( - interpretation.model_copy( - update={ - "evidenceIds": [marker_id], - **( - { - "identityIsHypothesis": True, - "confidence": "low", - } - if deps.markerAssayType == "ATAC" - else {} - ), - } - ) - ) - return canonical_interpretations, omitted_interpretation_clusters - - -def _canonicalize_treatment_observations( - report: BiologicalInterpretationReport, - deps: BiologicalInterpretationDependencies, -) -> list[TreatmentObservation]: - if report.treatmentObservations and deps.conditionColumn is None: - raise ModelRetry("Treatment observations require a condition column.") - if report.treatmentObservations and deps.sampleColumn is None: - raise ModelRetry( - "Treatment observations require independent-unit composition summaries." - ) - if report.treatmentObservations: - handoff = deps.designHandoff - if ( - handoff is None - or not handoff.conditionColumn - or handoff.conditionColumn != deps.conditionColumn - or not handoff.independentUnit - or handoff.independentUnit != deps.sampleColumn - or handoff.coefficientScope != "betweenUnit" - or handoff.estimability.get("status") != "ok" - or handoff.estimability.get("coefficientEstimable") is not True - ): - raise ModelRetry( - "Treatment observations require an explicit condition, aggregation " - "at the independent unit, a between-unit coefficient, and an " - "estimable experimental contrast." - ) - - canonical_observations: list[TreatmentObservation] = [] - for observation in report.treatmentObservations: - if not observation.isDescriptiveOnly: - raise ModelRetry("Treatment observations must remain descriptive.") - if len(observation.evidenceIds) != 2 or len(set(observation.evidenceIds)) != 2: - raise ModelRetry( - "Every treatment observation must cite exactly two distinct " - "condition summaries." - ) - if any( - evidence_id not in deps.conditionEvidence - for evidence_id in observation.evidenceIds - ): - raise ModelRetry( - "Treatment observations may cite only condition composition evidence." - ) - summaries = [ - deps.conditionEvidence[evidence_id] - for evidence_id in observation.evidenceIds - ] - if any(summary.clusterId != observation.clusterId for summary in summaries): - raise ModelRetry( - "Every treatment observation must cite condition summaries for " - "its exact cluster." - ) - if ( - not observation.referenceCondition - or not observation.comparisonCondition - or observation.referenceCondition == observation.comparisonCondition - ): - raise ModelRetry( - "Treatment observations require two distinct named conditions." - ) - summaries_by_condition = {summary.condition: summary for summary in summaries} - expected_conditions = { - observation.referenceCondition, - observation.comparisonCondition, - } - if set(summaries_by_condition) != expected_conditions: - raise ModelRetry( - "Treatment observation conditions must match the two cited " - "condition summaries." - ) - if any(summary.nSamples < 2 for summary in summaries): - raise ModelRetry( - "Sample-level treatment observations require at least two samples " - "in every cited condition." - ) - reference = summaries_by_condition[observation.referenceCondition] - comparison = summaries_by_condition[observation.comparisonCondition] - if math.isclose( - comparison.meanFraction, - reference.meanFraction, - rel_tol=1e-9, - abs_tol=1e-12, - ): - expected_direction: TreatmentDirection = "equal" - elif comparison.meanFraction > reference.meanFraction: - expected_direction = "higher" - else: - expected_direction = "lower" - if observation.direction != expected_direction: - raise ModelRetry( - "Treatment observation direction does not match the cited mean " - "independent-unit fractions." - ) - if expected_direction == "equal": - canonical_text = ( - f"Cluster {observation.clusterId} has equal mean independent-unit " - f"fractions in {comparison.condition} and {reference.condition} " - f"({comparison.meanFraction:.6g}); this is descriptive only." - ) - else: - canonical_text = ( - f"Cluster {observation.clusterId} has a {expected_direction} mean " - f"independent-unit fraction in {comparison.condition} " - f"({comparison.meanFraction:.6g}) than in {reference.condition} " - f"({reference.meanFraction:.6g}); this is descriptive only." - ) - canonical_observations.append( - observation.model_copy(update={"observation": canonical_text}) - ) - return canonical_observations - - -def validate_biological_interpretation_report( - report: BiologicalInterpretationReport, - deps: BiologicalInterpretationDependencies, -) -> BiologicalInterpretationReport: - """Reject invented evidence, clusters, or completed marker-free reviews.""" - if not deps.clusterValues: - raise ModelRetry("Call inspect_cluster_composition before returning a report.") - expected_cluster_artifact = artifact_reference(deps.cluster) - if ( - report.clusterArtifact is not None - and report.clusterArtifact != expected_cluster_artifact - ): - raise ModelRetry("Report clusterArtifact does not match the inspected artifact") - if deps.marker is not None: - expected_marker_artifact = artifact_reference(deps.marker) - if ( - report.markerArtifact is not None - and report.markerArtifact != expected_marker_artifact - ): - raise ModelRetry( - "Report markerArtifact does not match the inspected artifact" - ) - - cited = set(report.evidenceIds) - for interpretation in report.clusterInterpretations: - cited.update(interpretation.evidenceIds) - for observation in report.treatmentObservations: - cited.update(observation.evidenceIds) - for follow_up in report.followUps: - cited.update(follow_up.evidenceIds) - if report.needsInput is not None: - cited.update(report.needsInput.evidenceIds) - unknown = cited.difference(deps.evidenceIds) - if unknown: - raise ModelRetry(f"Unknown evidenceIds: {sorted(unknown)}") - interpreted_clusters = {item.clusterId for item in report.clusterInterpretations} - observed_clusters = {item.clusterId for item in report.treatmentObservations} - unknown_clusters = (interpreted_clusters | observed_clusters).difference( - deps.clusterValues - ) - if unknown_clusters: - raise ModelRetry(f"Unknown cluster ids: {sorted(unknown_clusters)}") - canonical_interpretations, omitted_interpretation_clusters = ( - _canonicalize_cluster_interpretations(report, deps) - ) - canonical_observations = _canonicalize_treatment_observations(report, deps) - if report.status == "done" and not canonical_interpretations: - raise ModelRetry( - "A done report must contain at least one cluster interpretation with " - "non-empty marker evidence." - ) - limitations = list(report.limitations) - if deps.markerAssayType == "ATAC": - atac_limitation = ( - "ATAC peak markers are descriptive, so all cell identities remain " - "low-confidence hypotheses." - ) - if atac_limitation not in limitations: - limitations.append(atac_limitation) - if omitted_interpretation_clusters: - omitted_clusters = ", ".join(sorted(set(omitted_interpretation_clusters))) - marker_limitation = ( - "Cluster identity interpretations without non-empty marker evidence " - f"were omitted for clusters: {omitted_clusters}." - ) - if marker_limitation not in limitations: - limitations.append(marker_limitation) - if canonical_observations: - descriptive_limitation = ( - "Independent-unit cluster fractions are descriptive summaries, not " - "tests of significance or causal treatment effects." - ) - if descriptive_limitation not in limitations: - limitations.append(descriptive_limitation) - validated = report.model_copy( - update={ - "clusterInterpretations": canonical_interpretations, - "treatmentObservations": canonical_observations, - "evidenceIds": sorted( - { - *report.evidenceIds, - *( - evidence_id - for interpretation in canonical_interpretations - for evidence_id in interpretation.evidenceIds - ), - } - ), - "limitations": limitations, - "clusterArtifact": expected_cluster_artifact, - "markerArtifact": ( - artifact_reference(deps.marker) if deps.marker is not None else None - ), - "graphAssay": deps.graphAssay, - "markerAssay": deps.markerAssay, - } - ) - logger.debug( - f"Validated biological interpretation report: status={validated.status}, " - f"cluster_interpretations={len(validated.clusterInterpretations)}, " - f"treatment_observations={len(validated.treatmentObservations)}, " - f"omitted_interpretations={len(omitted_interpretation_clusters)}" - ) - return validated - - -def _prepare_biological_interpretation_dependencies( - store: Any, - *, - cluster: Any, - from_assay: str | None, - graph_assay: str | None, - marker_assay_type: str | None, - sample_column: str | None, - condition_column: str | None, - tuning_handoff: TuningBiologyHandoff | None, - experimental_handoff: ExperimentalBiologyHandoff | None, - marker: Any, - marker_features: Any, - allow_marker_search: bool, - max_clusters: int, - max_markers: int, - marker_min_score: float, - marker_min_fraction: float, -) -> BiologicalInterpretationDependencies: - expected_selections: list[ArtifactRef] = [] - if tuning_handoff is not None: - if tuning_handoff.clusterArtifact is None: - raise ValueError("tuning_handoff lacks a cluster artifact") - tuning_selection = core_artifact_reference(tuning_handoff.cellSelection) - if not isinstance(tuning_selection, ArtifactRef): - raise ValueError("tuning_handoff lacks an exact cell selection") - expected_selections.append(tuning_selection) - if cluster is not None and ( - artifact_reference(cluster) != tuning_handoff.clusterArtifact - ): - raise ValueError("cluster conflicts with tuning_handoff") - if from_assay is not None and from_assay != tuning_handoff.fromAssay: - raise ValueError("from_assay conflicts with tuning_handoff") - if graph_assay is not None and graph_assay != tuning_handoff.graphAssay: - raise ValueError("graph_assay conflicts with tuning_handoff") - cluster = tuning_handoff.clusterArtifact - from_assay = tuning_handoff.fromAssay - graph_assay = tuning_handoff.graphAssay - if experimental_handoff is not None: - experimental_selection = core_artifact_reference( - experimental_handoff.cellSelection - ) - if not isinstance(experimental_selection, ArtifactRef): - raise ValueError("experimental_handoff lacks an exact cell selection") - expected_selections.append(experimental_selection) - if len(expected_selections) == 2 and ( - expected_selections[0] != expected_selections[1] - ): - raise ValueError( - "Experimental and tuning handoffs use different cell selections" - ) - if ( - condition_column is not None - and condition_column != experimental_handoff.conditionColumn - ): - raise ValueError("condition_column conflicts with experimental_handoff") - aggregation_unit = ( - experimental_handoff.independentUnit or experimental_handoff.observationUnit - ) - if sample_column is not None and sample_column != aggregation_unit: - raise ValueError("sample_column conflicts with experimental_handoff") - condition_column = experimental_handoff.conditionColumn - sample_column = aggregation_unit - if cluster is None: - raise ValueError("cluster must identify an exact cluster artifact") - if not 1 <= max_clusters <= CONFIG._MAX_CLUSTERS: - raise ValueError(f"max_clusters must be between 1 and {CONFIG._MAX_CLUSTERS}") - if not 1 <= max_markers <= CONFIG._MAX_MARKERS: - raise ValueError(f"max_markers must be between 1 and {CONFIG._MAX_MARKERS}") - if not 0 < marker_min_score <= 1: - raise ValueError("marker_min_score must be greater than 0 and at most 1") - if not 0 <= marker_min_fraction <= 1: - raise ValueError("marker_min_fraction must be between 0 and 1") - if allow_marker_search and marker is None and marker_features is None: - raise ValueError("marker_features is required when marker search is authorized") - - cluster = core_artifact_reference(cluster) - marker = core_artifact_reference(marker) - marker_features = core_artifact_reference(marker_features) - if not isinstance(cluster, ArtifactRef): - raise TypeError("cluster must be an ArtifactRef") - if marker is not None and ( - not isinstance(marker, ArtifactRef) or marker.kind != "marker_table" - ): - raise TypeError("marker must be a marker_table ArtifactRef") - if marker_features is not None and ( - not isinstance(marker_features, ArtifactRef) - or marker_features.kind != "feature_selection" - ): - raise TypeError("marker_features must be a feature_selection ArtifactRef") - cluster_artifact = artifact_reference(cluster) - if cluster_artifact.kind not in {"cluster_labels", "cluster_cut"}: - raise ValueError( - "cluster must identify a cluster_labels or cluster_cut artifact" - ) - if cluster_artifact.scope == "datastore" and cluster_artifact.assay is not None: - raise ValueError("datastore-scoped cluster artifacts must not name an assay") - if ( - tuning_handoff is not None - and cluster_artifact.scope == "datastore" - and not tuning_handoff.markerAssay - ): - raise ValueError( - "Integrated tuning handoffs must explicitly identify markerAssay" - ) - resolved_graph_assay = graph_assay or cluster_artifact.assay - if ( - resolved_graph_assay is not None - and cluster_artifact.scope == "assay" - and cluster_artifact.assay != resolved_graph_assay - ): - raise ValueError("cluster belongs to a different assay") - resolved_marker_assay = ( - tuning_handoff.markerAssay or cluster_artifact.assay - if tuning_handoff is not None - else ( - marker.assay - if isinstance(marker, ArtifactRef) - else ( - marker_features.assay - if isinstance(marker_features, ArtifactRef) - else from_assay or cluster_artifact.assay - ) - ) - ) - if cluster_artifact.scope == "datastore" and not resolved_marker_assay: - raise ValueError( - "from_assay is required to resolve markers for integrated clusters" - ) - if isinstance(marker, ArtifactRef) and marker.assay != resolved_marker_assay: - raise ValueError("marker artifact belongs to a different marker assay") - if ( - isinstance(marker_features, ArtifactRef) - and marker_features.assay != resolved_marker_assay - ): - raise ValueError("marker feature selection belongs to a different assay") - - cluster_status = store.inspect_artifact(cluster) - if not getattr(cluster_status, "exists", True): - raise ValueError("cluster artifact does not exist") - if not getattr(cluster_status, "complete", False): - raise ValueError("cluster artifact is incomplete") - raw_selection = (getattr(cluster_status, "inputs", None) or {}).get( - "cell_selection" - ) - if not isinstance(raw_selection, Mapping): - raise ValueError("cluster artifact has no cell-selection input") - cell_selection = ArtifactRef.from_dict(dict(raw_selection)) - if ( - cell_selection.scope != "datastore" - or cell_selection.kind != "cell_selection" - or cell_selection.assay is not None - ): - raise ValueError("cluster artifact has an invalid cell-selection input") - if any(selection != cell_selection for selection in expected_selections): - raise ValueError("handoff cell selection conflicts with cluster") - cell_indices = read_stored_selection_indices( - store.zw, - cell_selection, - kind="cell_selection", - scope="datastore", - assay=None, - table_path="cellData", - ).astype(np.int64, copy=False) - return BiologicalInterpretationDependencies( - store=store, - cluster=cluster, - cellSelection=cell_selection, - cellIndices=cell_indices, - fromAssay=from_assay or cluster_artifact.assay or resolved_marker_assay, - graphAssay=resolved_graph_assay, - markerAssay=resolved_marker_assay, - markerAssayType=marker_assay_type, - sampleColumn=sample_column, - conditionColumn=condition_column, - designHandoff=experimental_handoff, - marker=marker, - markerFeatures=marker_features, - allowMarkerSearch=allow_marker_search, - maxClusters=max_clusters, - maxMarkers=max_markers, - markerMinScore=marker_min_score, - markerMinFraction=marker_min_fraction, - ) - - -class BiologicalInterpretationAgent: - """Run a bounded biological review through explicit Scarf tools.""" - - def __init__( - self, - model: Any, - *, - config: AgentRunConfig | None = None, - ) -> None: - self.model = model - self.config = (config or AgentRunConfig()).with_limits( - request_limit=8, - tool_call_limit=5, - output_token_limit=32768, - timeout_seconds=600.0, - ) - - def run( - self, - store: Any, - *, - cluster: ArtifactRef | ArtifactReferenceModel | None = None, - biological_context: BiologicalContext | None = None, - from_assay: str | None = None, - graph_assay: str | None = None, - marker_assay_type: str | None = None, - sample_column: str | None = None, - condition_column: str | None = None, - tuning_handoff: TuningBiologyHandoff | None = None, - experimental_handoff: ExperimentalBiologyHandoff | None = None, - marker: ArtifactRef | ArtifactReferenceModel | None = None, - marker_features: ArtifactRef | ArtifactReferenceModel | None = None, - allow_marker_search: bool = False, - max_clusters: int = 12, - max_markers: int = 10, - marker_min_score: float = 0.25, - marker_min_fraction: float = 0.2, - ) -> BiologicalInterpretationReport: - """Interpret cluster results while exposing only bounded tools to the model.""" - deps = _prepare_biological_interpretation_dependencies( - store, - cluster=cluster, - from_assay=from_assay, - graph_assay=graph_assay, - marker_assay_type=marker_assay_type, - sample_column=sample_column, - condition_column=condition_column, - tuning_handoff=tuning_handoff, - experimental_handoff=experimental_handoff, - marker=marker, - marker_features=marker_features, - allow_marker_search=allow_marker_search, - max_clusters=max_clusters, - max_markers=max_markers, - marker_min_score=marker_min_score, - marker_min_fraction=marker_min_fraction, - ) - logger.info( - f"Starting biological interpretation: " - f"cluster_artifact={deps.cluster.artifact_id!r}, " - f"graph_assay={deps.graphAssay!r}, marker_assay={deps.markerAssay!r}, " - f"marker_artifact_supplied={deps.marker is not None}, " - f"marker_search_authorized={deps.allowMarkerSearch}" - ) - context = biological_context or BiologicalContext() - cluster_artifact = artifact_reference(deps.cluster) - marker_state = "provided" if deps.marker is not None else "not provided" - user_prompt = ( - dedent( - """ - Review the exact cluster artifact {cluster_artifact} over - cell-selection artifact {cell_selection}. The graph owner is - {graph_assay}; markers are resolved from {marker_assay}. The exact - marker artifact is {marker_state}; creating a marker artifact is - authorized={allow_marker_search}. Review no more - than {max_clusters} clusters and return no more than {max_markers} - markers per tool call. - - Caller biological context: - {biological_context} - - Experimental design context: - {experimental_context} - - Call inspect_cluster_composition once. Then send every cluster you - intend to interpret in one inspect_cluster_markers_batch call. If - markers cannot be inspected, return needsInput and state the exact - missing input. Each tool is removed after it succeeds, so request - every required cluster in that one marker batch. - """ - ) - .strip() - .format( - graph_assay=deps.graphAssay or "datastore integration", - marker_assay=deps.markerAssay, - cluster_artifact=cluster_artifact.model_dump_json(), - cell_selection=deps.cellSelection.artifact_id, - marker_state=marker_state, - allow_marker_search=allow_marker_search, - max_clusters=max_clusters, - max_markers=max_markers, - biological_context=context.model_dump_json(), - experimental_context=( - experimental_handoff.model_dump_json() - if experimental_handoff is not None - else "not provided" - ), - ) - ) - logger.info( - f"Requesting biological interpretation for at most " - f"{deps.maxClusters} clusters" - ) - execution = run_agent_sync( - model=self.model, - output_type=BiologicalInterpretationReport, - system_prompt=_SYSTEM_PROMPT, - user_prompt=user_prompt, - tools=( - Tool( - inspect_cluster_composition, - prepare=_prepare_biological_interpretation_tool, - sequential=self.config.sequentialTools, - timeout=self.config.timeoutSeconds, - ), - Tool( - inspect_cluster_markers_batch, - prepare=_prepare_biological_interpretation_tool, - sequential=self.config.sequentialTools, - timeout=self.config.timeoutSeconds, - ), - ), - deps_type=BiologicalInterpretationDependencies, - deps=deps, - config=self.config, - name="biological_interpretation", - output_validator=lambda report: validate_biological_interpretation_report( - report, - deps, - ), - ) - report = validate_biological_interpretation_report(execution.output, deps) - report.runInfo = execution.runInfo - logger.info( - f"Completed biological interpretation: status={report.status}, " - f"interpreted_clusters={len(report.clusterInterpretations)}, " - f"treatment_observations={len(report.treatmentObservations)}, " - f"follow_ups={len(report.followUps)}, tool_calls={len(deps.toolCalls)}" - ) - return report diff --git a/scarf/agent/biological_interpretation/__init__.py b/scarf/agent/biological_interpretation/__init__.py new file mode 100644 index 00000000..a9c449a0 --- /dev/null +++ b/scarf/agent/biological_interpretation/__init__.py @@ -0,0 +1,41 @@ +"""Grounded biological interpretation of Scarf cluster results.""" + +from .agent import BiologicalInterpretationAgent +from .contracts import ( + BiologicalContext, + BiologicalInterpretationNeedsInput, + BiologicalInterpretationReport, + ClusterCompositionEvidence, + ClusterInterpretation, + ClusterMarkerBatchEvidence, + ClusterMarkerEvidence, + ConditionClusterSummary, + FollowUpRecommendation, + MarkerFeature, + TreatmentObservation, +) +from .tools import ( + inspect_cluster_composition, + inspect_cluster_markers, + inspect_cluster_markers_batch, +) +from .validation import validate_biological_interpretation_report + +__all__ = [ + "BiologicalContext", + "BiologicalInterpretationAgent", + "BiologicalInterpretationNeedsInput", + "BiologicalInterpretationReport", + "ClusterCompositionEvidence", + "ClusterInterpretation", + "ClusterMarkerBatchEvidence", + "ClusterMarkerEvidence", + "ConditionClusterSummary", + "FollowUpRecommendation", + "MarkerFeature", + "TreatmentObservation", + "inspect_cluster_composition", + "inspect_cluster_markers_batch", + "inspect_cluster_markers", + "validate_biological_interpretation_report", +] diff --git a/scarf/agent/biological_interpretation/agent.py b/scarf/agent/biological_interpretation/agent.py new file mode 100644 index 00000000..6c6b18ef --- /dev/null +++ b/scarf/agent/biological_interpretation/agent.py @@ -0,0 +1,255 @@ +"""Biological interpretation prompt and agent runner.""" + +from textwrap import dedent +from typing import Any + +from ...storage.refs import ArtifactRef +from ...utils.logging import logger +from ..config import AgentRunConfig +from ..config.agent_exec import run_agent_sync +from ..tools import artifact_reference +from ..types import ( + ArtifactReferenceModel, + ExperimentalBiologyHandoff, + TuningBiologyHandoff, +) +from .contracts import ( + BiologicalContext, + BiologicalInterpretationDependencies, + BiologicalInterpretationReport, +) +from .tools import inspect_cluster_composition, inspect_cluster_markers_batch +from .validation import ( + _prepare_biological_interpretation_dependencies, + _prepare_biological_interpretation_tool, + fallback_biological_interpretation_report, + validate_biological_interpretation_report, +) + +try: + from pydantic_ai import Tool, UnexpectedModelBehavior, UsageLimitExceeded +except ImportError as exc: + from .._deps import AGENT_INSTALL_HINT + + raise ImportError(AGENT_INSTALL_HINT) from exc + + +_SYSTEM_PROMPT = dedent( + """ + You are Scarf's Biological Interpretation Agent. Use only the supplied + tools and caller context. Call inspect_cluster_composition exactly once. + Then select every cluster you intend to interpret and call + inspect_cluster_markers_batch exactly once with all selected cluster IDs. + Tool calls execute Scarf operations, so wait for their results before + drawing conclusions. Do not split marker inspection across calls. + + This API does not provide trusted per-cluster identities. Treat every cell + identity as a hypothesis and always set identityIsHypothesis=true. Caller + cell-type references are context, not assignments to clusters. Prefer + proposedIdentity="unresolved" when the returned markers do not ground a + specific hypothesis. Do not invent genes, cell types, statistics, artifact + identifiers, or evidence identifiers. Cite only evidenceIds returned by + tools. For each cluster interpretation, copy the exact non-empty marker + evidenceId returned for that cluster into its evidenceIds. Do not interpret + a cluster whose marker evidenceId is empty. Cluster abundance summaries are + descriptive, not tests of significance or causal effects. Treatment observations must + compare two returned independent-unit condition summaries for the same + cluster. Independent units may occur in more than one condition in paired + or repeated-measure designs. Return treatmentObservations empty unless the + exact experimental handoff confirms independent-unit aggregation, a + between-unit coefficient, an estimable coefficient, and at least two + independent units in each cited condition. + Marker p-values describe cluster-versus-rest marker specificity, not + condition effects. Keep treatment content out of cluster identity + interpretations. Never return status=failed for uncertainty or weak + evidence. Return needsInput with one concrete question instead. Recommend a + named follow-up operation when replication, a covariate, or an exact + artifact is missing. Do not write exploratory code, use a shell, access + files, or call arbitrary Scarf methods. Return only fields defined by the + structured output schema. + """ +).strip() + + +class BiologicalInterpretationAgent: + """Run a bounded biological review through explicit Scarf tools.""" + + def __init__( + self, + model: Any, + *, + config: AgentRunConfig | None = None, + ) -> None: + self.model = model + self.config = (config or AgentRunConfig()).with_limits( + request_limit=8, + tool_call_limit=5, + output_token_limit=32768, + timeout_seconds=600.0, + ) + + def run( + self, + store: Any, + *, + cluster: ArtifactRef | ArtifactReferenceModel | None = None, + biological_context: BiologicalContext | None = None, + from_assay: str | None = None, + graph_assay: str | None = None, + marker_assay_type: str | None = None, + sample_column: str | None = None, + condition_column: str | None = None, + tuning_handoff: TuningBiologyHandoff | None = None, + experimental_handoff: ExperimentalBiologyHandoff | None = None, + marker: ArtifactRef | ArtifactReferenceModel | None = None, + marker_features: ArtifactRef | ArtifactReferenceModel | None = None, + allow_marker_search: bool = False, + max_clusters: int = 12, + max_markers: int = 10, + marker_min_score: float = 0.25, + marker_min_fraction: float = 0.2, + ) -> BiologicalInterpretationReport: + """Interpret cluster results while exposing only bounded tools to the model.""" + deps = _prepare_biological_interpretation_dependencies( + store, + cluster=cluster, + from_assay=from_assay, + graph_assay=graph_assay, + marker_assay_type=marker_assay_type, + sample_column=sample_column, + condition_column=condition_column, + tuning_handoff=tuning_handoff, + experimental_handoff=experimental_handoff, + marker=marker, + marker_features=marker_features, + allow_marker_search=allow_marker_search, + max_clusters=max_clusters, + max_markers=max_markers, + marker_min_score=marker_min_score, + marker_min_fraction=marker_min_fraction, + ) + logger.info( + f"Starting biological interpretation: " + f"cluster_artifact={deps.cluster.artifact_id!r}, " + f"graph_assay={deps.graphAssay!r}, marker_assay={deps.markerAssay!r}, " + f"marker_artifact_supplied={deps.marker is not None}, " + f"marker_search_authorized={deps.allowMarkerSearch}" + ) + context = biological_context or BiologicalContext() + cluster_artifact = artifact_reference(deps.cluster) + marker_state = "provided" if deps.marker is not None else "not provided" + treatment_eligible = bool( + experimental_handoff is not None + and experimental_handoff.conditionColumn + and experimental_handoff.independentUnit + and experimental_handoff.coefficientScope == "betweenUnit" + and experimental_handoff.estimability.get("status") == "ok" + and experimental_handoff.estimability.get("coefficientEstimable") is True + ) + user_prompt = ( + dedent( + """ + Review the exact cluster artifact {cluster_artifact} over + cell-selection artifact {cell_selection}. The graph owner is + {graph_assay}; markers are resolved from {marker_assay}. The exact + marker artifact is {marker_state}; creating a marker artifact is + authorized={allow_marker_search}. Review no more + than {max_clusters} clusters. The tool returns no more than + {max_markers} markers per cluster. Treatment observations are + eligible from the supplied design={treatment_eligible}. + + Caller biological context: + {biological_context} + + Experimental design context: + {experimental_context} + + Call inspect_cluster_composition once. Copy every returned cluster + ID exactly and send the complete unique list in one + inspect_cluster_markers_batch call. Interpret only clusters with a + non-empty returned marker evidenceId. Use proposedIdentity="unresolved" + when markers do not ground a specific hypothesis, and always set + identityIsHypothesis=true. Caller cell-type references are not + cluster labels. If marker evidence is empty, return needsInput with + one populated question. If treatment eligibility is false, return + treatmentObservations=[]. Never return status=failed for biological + uncertainty; use needsInput with a concrete question. Leave artifact + fields null or copy only exact tool-returned values. Each tool is + removed after it succeeds, so request every cluster in that one + marker batch. + """ + ) + .strip() + .format( + graph_assay=deps.graphAssay or "datastore integration", + marker_assay=deps.markerAssay, + cluster_artifact=cluster_artifact.model_dump_json(), + cell_selection=deps.cellSelection.artifact_id, + marker_state=marker_state, + allow_marker_search=allow_marker_search, + max_clusters=max_clusters, + max_markers=max_markers, + treatment_eligible=str(treatment_eligible).lower(), + biological_context=context.model_dump_json(), + experimental_context=( + experimental_handoff.model_dump_json() + if experimental_handoff is not None + else "not provided" + ), + ) + ) + logger.info( + f"Requesting biological interpretation for at most " + f"{deps.maxClusters} clusters" + ) + try: + execution = run_agent_sync( + model=self.model, + output_type=BiologicalInterpretationReport, + system_prompt=_SYSTEM_PROMPT, + user_prompt=user_prompt, + tools=( + Tool( + inspect_cluster_composition, + prepare=_prepare_biological_interpretation_tool, + sequential=self.config.sequentialTools, + timeout=self.config.timeoutSeconds, + ), + Tool( + inspect_cluster_markers_batch, + max_retries=1, + prepare=_prepare_biological_interpretation_tool, + sequential=self.config.sequentialTools, + timeout=self.config.timeoutSeconds, + ), + ), + deps_type=BiologicalInterpretationDependencies, + deps=deps, + config=self.config, + name="biological_interpretation", + output_validator=lambda report: ( + validate_biological_interpretation_report( + report, + deps, + ) + ), + ) + except (UnexpectedModelBehavior, UsageLimitExceeded) as exc: + if not deps.clusterValues: + raise + model_name = getattr(self.model, "model_name", type(self.model).__name__) + return fallback_biological_interpretation_report( + deps, + error=exc, + model_name=str(model_name), + ) + report = BiologicalInterpretationReport.model_validate(execution.output) + report = validate_biological_interpretation_report(report, deps) + report.runInfo = execution.runInfo + logger.info( + f"Completed biological interpretation: status={report.status}, " + f"interpreted_clusters={len(report.clusterInterpretations)}, " + f"treatment_observations={len(report.treatmentObservations)}, " + f"follow_ups={len(report.followUps)}, tool_calls={len(deps.toolCalls)}" + ) + return report diff --git a/scarf/agent/biological_interpretation/contracts.py b/scarf/agent/biological_interpretation/contracts.py new file mode 100644 index 00000000..e154d853 --- /dev/null +++ b/scarf/agent/biological_interpretation/contracts.py @@ -0,0 +1,205 @@ +"""Serializable contracts for biological interpretation.""" + +from typing import Any, Literal + +from pydantic import Field +from pydantic.json_schema import SkipJsonSchema + +from ..types import ( + AgentDataModel, + AgentRunInfo, + ArtifactReferenceModel, + ExperimentalBiologyHandoff, + StageStatus, +) + +type InterpretationConfidence = Literal["low", "medium", "high"] +type TreatmentDirection = Literal["higher", "lower", "equal"] + +_MAX_CLUSTERS = 20 +_MAX_CONDITIONS = 30 +_MAX_MARKERS = 25 + + +class BiologicalContext(AgentDataModel): + """Caller-supplied facts that constrain biological interpretation.""" + + organism: str = "" + studyContext: str = "" + tissue: str = "" + cellTypeReferences: list[str] = Field(default_factory=list) + experimentalDetails: list[str] = Field(default_factory=list) + treatmentQuestion: str = "" + + @classmethod + def get_blank(cls) -> "BiologicalContext": + return cls() + + +class ConditionClusterSummary(AgentDataModel): + """Aggregate cluster abundance for one condition without sample identifiers.""" + + condition: str = "" + clusterId: str = "" + nSamples: int = 0 + meanFraction: float = 0.0 + minFraction: float = 0.0 + maxFraction: float = 0.0 + cellCount: int = 0 + evidenceId: str = "" + + +class ClusterCompositionEvidence(AgentDataModel): + """Bounded deterministic evidence about cluster sizes and conditions.""" + + clusterArtifact: ArtifactReferenceModel | None = None + cellSelection: ArtifactReferenceModel | None = None + totalCells: int = 0 + clusterCounts: dict[str, int] = Field(default_factory=dict) + sampleColumn: str | None = None + conditionColumn: str | None = None + conditionSummaries: list[ConditionClusterSummary] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + warnings: list[str] = Field(default_factory=list) + + +class MarkerFeature(AgentDataModel): + """One observed marker feature and its available Scarf statistics.""" + + featureId: str = "" + featureName: str = "" + featureIndex: int | None = None + score: float | None = None + foldChange: float | None = None + fractionExpressed: float | None = None + fractionExpressedRest: float | None = None + mean: float | None = None + meanRest: float | None = None + auc: float | None = None + adjustedPvalue: float | None = None + + +class ClusterMarkerEvidence(AgentDataModel): + """Bounded markers for one exact cluster label.""" + + clusterId: str = "" + markers: list[MarkerFeature] = Field(default_factory=list) + markerArtifact: ArtifactReferenceModel | None = None + evidenceId: str = "" + warnings: list[str] = Field(default_factory=list) + + +class ClusterMarkerBatchEvidence(AgentDataModel): + """Markers for all model-selected clusters returned by one tool call.""" + + clusters: list[ClusterMarkerEvidence] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + warnings: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "ClusterMarkerBatchEvidence": + return cls() + + +class ClusterInterpretation(AgentDataModel): + """One evidence-linked cluster interpretation or hypothesis.""" + + clusterId: str = "" + proposedIdentity: str = "unresolved" + identityIsHypothesis: bool = True + confidence: InterpretationConfidence = "low" + rationale: str = "" + evidenceIds: list[str] = Field(default_factory=list) + + +class TreatmentObservation(AgentDataModel): + """Descriptive treatment observation with no unsupported causal claim.""" + + clusterId: str = "" + referenceCondition: str = "" + comparisonCondition: str = "" + direction: TreatmentDirection = "equal" + observation: str = "" + isDescriptiveOnly: Literal[True] = True + evidenceIds: list[str] = Field(default_factory=list) + + +class FollowUpRecommendation(AgentDataModel): + """A bounded next analysis tied to an observed uncertainty.""" + + question: str = "" + operation: str = "" + rationale: str = "" + requiredInputs: list[str] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + + +class BiologicalInterpretationNeedsInput(AgentDataModel): + question: str = "" + requiredInputs: list[str] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + + +class BiologicalInterpretationReport(AgentDataModel): + """Structured, evidence-grounded biological review.""" + + status: StageStatus = "needsInput" + clusterInterpretations: list[ClusterInterpretation] = Field(default_factory=list) + treatmentObservations: list[TreatmentObservation] = Field(default_factory=list) + followUps: list[FollowUpRecommendation] = Field(default_factory=list) + clusterArtifact: SkipJsonSchema[ArtifactReferenceModel | None] = None + markerArtifact: SkipJsonSchema[ArtifactReferenceModel | None] = None + graphAssay: SkipJsonSchema[str | None] = None + markerAssay: SkipJsonSchema[str | None] = None + evidenceIds: list[str] = Field(default_factory=list) + limitations: list[str] = Field(default_factory=list) + stopReason: str = "" + needsInput: BiologicalInterpretationNeedsInput | None = None + runInfo: SkipJsonSchema[AgentRunInfo] = Field(default_factory=AgentRunInfo) + + +class BiologicalInterpretationDependencies(AgentDataModel): + """Runtime state available only to biological interpretation tools.""" + + store: Any = Field(default=None, exclude=True) + cluster: Any = Field(default=None, exclude=True) + cellSelection: Any = Field(default=None, exclude=True) + cellIndices: Any = Field(default=None, exclude=True) + fromAssay: str | None = None + graphAssay: str | None = None + markerAssay: str | None = None + markerAssayType: str | None = None + sampleColumn: str | None = None + conditionColumn: str | None = None + marker: Any = Field(default=None, exclude=True) + markerFeatures: Any = Field(default=None, exclude=True) + allowMarkerSearch: bool = False + maxClusters: int = 12 + maxMarkers: int = 10 + markerMinScore: float = 0.25 + markerMinFraction: float = 0.2 + evidenceIds: set[str] = Field(default_factory=set, exclude=True) + clusterValues: dict[str, Any] = Field(default_factory=dict, exclude=True) + markerEvidenceIds: dict[str, str] = Field(default_factory=dict, exclude=True) + markerEvidence: dict[str, ClusterMarkerEvidence] = Field( + default_factory=dict, + exclude=True, + ) + compositionEvidence: ClusterCompositionEvidence | None = Field( + default=None, + exclude=True, + ) + markerBatch: ClusterMarkerBatchEvidence | None = Field( + default=None, + exclude=True, + ) + markerBatchClusterIds: list[str] = Field(default_factory=list, exclude=True) + toolCalls: list[str] = Field(default_factory=list, exclude=True) + conditionEvidence: dict[str, ConditionClusterSummary] = Field( + default_factory=dict, + exclude=True, + ) + designHandoff: ExperimentalBiologyHandoff | None = Field( + default=None, + exclude=True, + ) diff --git a/scarf/agent/biological_interpretation/tools.py b/scarf/agent/biological_interpretation/tools.py new file mode 100644 index 00000000..b5d44de5 --- /dev/null +++ b/scarf/agent/biological_interpretation/tools.py @@ -0,0 +1,483 @@ +"""Bounded evidence tools for biological interpretation.""" + +from collections import Counter, defaultdict +from collections.abc import Mapping +from typing import Any + +import numpy as np + +from ...metadata.rows import read_metadata_missing_rows, read_metadata_rows +from ...storage.refs import ArtifactRef +from ...storage.selections import read_stored_selection_indices +from ...utils.logging import logger +from ..tools import artifact_reference, core_artifact_reference +from .contracts import ( + _MAX_CONDITIONS, + _MAX_MARKERS, + BiologicalInterpretationDependencies, + ClusterCompositionEvidence, + ClusterMarkerBatchEvidence, + ClusterMarkerEvidence, + ConditionClusterSummary, + MarkerFeature, +) + +try: + from pydantic_ai import ModelRetry, RunContext +except ImportError as exc: + from .._deps import AGENT_INSTALL_HINT + + raise ImportError(AGENT_INSTALL_HINT) from exc + + +def _string_value(value: Any) -> str: + return str(value.item() if isinstance(value, np.generic) else value) + + +def _finite_float(value: Any) -> float | None: + if value is None: + return None + try: + number = float(value) + except (TypeError, ValueError): + return None + return number if np.isfinite(number) else None + + +def _check_column(store: Any, column: str, label: str) -> None: + if column not in set(store.cells.columns): + raise ValueError(f"{label} {column!r} is not present in cell metadata") + + +async def inspect_cluster_composition( + ctx: RunContext[BiologicalInterpretationDependencies], +) -> ClusterCompositionEvidence: + """Inspect bounded cluster and condition composition without identifiers.""" + deps = ctx.deps + if deps.compositionEvidence is not None: + logger.info("Reused completed cluster composition inspection") + return deps.compositionEvidence + logger.info( + f"Inspecting cluster composition from artifact " + f"{getattr(deps.cluster, 'artifact_id', '')!r}; " + f"max_clusters={deps.maxClusters}" + ) + if deps.sampleColumn is not None: + _check_column(deps.store, deps.sampleColumn, "sample column") + if deps.conditionColumn is not None: + _check_column(deps.store, deps.conditionColumn, "condition column") + + if deps.cluster is None: + raise ValueError("An exact cluster artifact is required") + deps.cluster = core_artifact_reference(deps.cluster) + cluster_artifact = artifact_reference(deps.cluster) + if cluster_artifact.kind not in {"cluster_labels", "cluster_cut"}: + raise ValueError( + "cluster must identify a cluster_labels or cluster_cut artifact" + ) + if ( + deps.graphAssay is not None + and cluster_artifact.scope == "assay" + and cluster_artifact.assay != deps.graphAssay + ): + raise ValueError("cluster artifact belongs to a different assay") + if cluster_artifact.scope == "datastore" and cluster_artifact.assay is not None: + raise ValueError("datastore-scoped cluster artifacts must not name an assay") + status = deps.store.inspect_artifact(deps.cluster) + if not getattr(status, "exists", True): + raise ValueError("cluster artifact does not exist") + if not getattr(status, "complete", False): + raise ValueError("cluster artifact is incomplete") + inputs = getattr(status, "inputs", None) or {} + raw_selection = inputs.get("cell_selection") + if not isinstance(raw_selection, Mapping): + raise ValueError("cluster artifact has no cell-selection input") + cell_selection = ArtifactRef.from_dict(dict(raw_selection)) + if ( + cell_selection.scope != "datastore" + or cell_selection.kind != "cell_selection" + or cell_selection.assay is not None + ): + raise ValueError("cluster artifact has an invalid cell-selection input") + if ( + deps.cellSelection is not None + and core_artifact_reference(deps.cellSelection) != cell_selection + ): + raise ValueError( + "cluster artifact cell selection conflicts with the prepared selection" + ) + cell_indices = read_stored_selection_indices( + deps.store.zw, + cell_selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + deps.cellSelection = cell_selection + deps.cellIndices = cell_indices + cluster_group = deps.store.load_artifact(deps.cluster) + value_name = "labels" if cluster_artifact.kind == "cluster_cut" else "values" + if value_name not in cluster_group: + raise ValueError( + f"cluster artifact does not contain its {value_name!r} label vector" + ) + cluster_values = np.asarray(cluster_group[value_name][:]) + if cluster_values.ndim != 1 or len(cluster_values) != len(cell_indices): + raise ValueError("cluster artifact labels do not align with its cell selection") + if len(cluster_values) == 0: + raise ValueError("cluster artifact selects no cells") + counts = Counter(_string_value(value) for value in cluster_values) + ordered_clusters = sorted(counts, key=lambda value: (-counts[value], value)) + retained_clusters = ordered_clusters[: deps.maxClusters] + deps.clusterValues = { + _string_value(value): value.item() if isinstance(value, np.generic) else value + for value in cluster_values + if _string_value(value) in retained_clusters + } + evidence_prefix = f"composition:{cluster_artifact.artifactId}" + count_evidence = f"{evidence_prefix}:counts" + deps.evidenceIds.add(count_evidence) + warnings: list[str] = [] + if len(ordered_clusters) > deps.maxClusters: + warnings.append(f"Only the {deps.maxClusters} largest clusters were returned.") + + condition_summaries: list[ConditionClusterSummary] = [] + if deps.conditionColumn is not None: + condition_values = read_metadata_rows( + deps.store.cells, + deps.conditionColumn, + cell_indices, + ) + if len(condition_values) != len(cluster_values): + raise ValueError("condition and cluster columns are not aligned") + condition_missing = read_metadata_missing_rows( + deps.store.cells, + deps.conditionColumn, + cell_indices, + ) + if condition_missing is not None and np.any(condition_missing): + raise ValueError("condition column contains missing selected values") + n_conditions = len({_string_value(value) for value in condition_values}) + if n_conditions > _MAX_CONDITIONS: + warnings.append( + f"Only the first {_MAX_CONDITIONS} conditions were returned." + ) + if deps.sampleColumn is not None: + sample_values = read_metadata_rows( + deps.store.cells, + deps.sampleColumn, + cell_indices, + ) + if len(sample_values) != len(cluster_values): + raise ValueError("sample and cluster columns are not aligned") + sample_missing = read_metadata_missing_rows( + deps.store.cells, + deps.sampleColumn, + cell_indices, + ) + if sample_missing is not None and np.any(sample_missing): + raise ValueError("sample column contains missing selected values") + summaries = _sample_condition_summaries( + sample_values=sample_values, + condition_values=condition_values, + cluster_values=cluster_values, + retained_clusters=retained_clusters, + evidence_prefix=evidence_prefix, + ) + else: + summaries = _cell_condition_summaries( + condition_values=condition_values, + cluster_values=cluster_values, + retained_clusters=retained_clusters, + evidence_prefix=evidence_prefix, + ) + warnings.append( + "No sample column was supplied; condition fractions are cell-level summaries." + ) + condition_summaries = summaries[: _MAX_CONDITIONS * len(retained_clusters)] + deps.evidenceIds.update(summary.evidenceId for summary in condition_summaries) + deps.conditionEvidence.update( + {summary.evidenceId: summary for summary in condition_summaries} + ) + + deps.toolCalls.append("inspect_cluster_composition") + evidence = ClusterCompositionEvidence( + clusterArtifact=cluster_artifact, + cellSelection=artifact_reference(cell_selection), + totalCells=len(cluster_values), + clusterCounts={cluster: counts[cluster] for cluster in retained_clusters}, + sampleColumn=deps.sampleColumn, + conditionColumn=deps.conditionColumn, + conditionSummaries=condition_summaries, + evidenceIds=sorted(deps.evidenceIds), + warnings=warnings, + ) + deps.compositionEvidence = evidence + logger.info( + f"Completed cluster composition inspection: cells={evidence.totalCells}, " + f"clusters={len(evidence.clusterCounts)}, " + f"condition_summaries={len(evidence.conditionSummaries)}, " + f"warnings={len(evidence.warnings)}" + ) + return evidence + + +def _sample_condition_summaries( + *, + sample_values: np.ndarray, + condition_values: np.ndarray, + cluster_values: np.ndarray, + retained_clusters: list[str], + evidence_prefix: str, +) -> list[ConditionClusterSummary]: + """Aggregate each condition-unit pair without exposing unit identifiers.""" + sample_counts: dict[tuple[str, str], Counter[str]] = defaultdict(Counter) + sample_totals: Counter[tuple[str, str]] = Counter() + for sample, condition, cluster in zip( + sample_values, + condition_values, + cluster_values, + strict=True, + ): + condition_label = _string_value(condition) + sample_label = _string_value(sample) + key = (condition_label, sample_label) + sample_totals[key] += 1 + sample_counts[key][_string_value(cluster)] += 1 + + fractions: dict[tuple[str, str], list[float]] = defaultdict(list) + cell_counts: Counter[tuple[str, str]] = Counter() + for key, total in sample_totals.items(): + condition, _sample = key + for cluster in retained_clusters: + count = sample_counts[key][cluster] + fractions[(condition, cluster)].append(count / total) + cell_counts[(condition, cluster)] += count + + output: list[ConditionClusterSummary] = [] + for condition, cluster in sorted(fractions): + values = fractions[(condition, cluster)] + evidence_id = f"{evidence_prefix}:condition:{condition}:cluster:{cluster}" + output.append( + ConditionClusterSummary( + condition=condition, + clusterId=cluster, + nSamples=len(values), + meanFraction=float(np.mean(values)), + minFraction=float(np.min(values)), + maxFraction=float(np.max(values)), + cellCount=cell_counts[(condition, cluster)], + evidenceId=evidence_id, + ) + ) + return output + + +def _cell_condition_summaries( + *, + condition_values: np.ndarray, + cluster_values: np.ndarray, + retained_clusters: list[str], + evidence_prefix: str, +) -> list[ConditionClusterSummary]: + totals: Counter[str] = Counter(_string_value(value) for value in condition_values) + counts: Counter[tuple[str, str]] = Counter() + for condition, cluster in zip(condition_values, cluster_values, strict=True): + counts[(_string_value(condition), _string_value(cluster))] += 1 + output: list[ConditionClusterSummary] = [] + for condition in sorted(totals): + for cluster in retained_clusters: + count = counts[(condition, cluster)] + fraction = count / totals[condition] + evidence_id = f"{evidence_prefix}:condition:{condition}:cluster:{cluster}" + output.append( + ConditionClusterSummary( + condition=condition, + clusterId=cluster, + meanFraction=fraction, + minFraction=fraction, + maxFraction=fraction, + cellCount=count, + evidenceId=evidence_id, + ) + ) + return output + + +async def inspect_cluster_markers( + ctx: RunContext[BiologicalInterpretationDependencies], + cluster_id: str, +) -> ClusterMarkerEvidence: + """Load markers for one observed cluster, optionally creating one artifact.""" + deps = ctx.deps + logger.debug(f"Inspecting markers for cluster {cluster_id!r}") + if not deps.clusterValues: + raise ModelRetry("Call inspect_cluster_composition before inspecting markers.") + if cluster_id not in deps.clusterValues: + raise ModelRetry(f"cluster_id must be one of {sorted(deps.clusterValues)}") + cached = deps.markerEvidence.get(cluster_id) + if cached is not None: + logger.debug(f"Reused cached markers for cluster {cluster_id!r}") + return cached + if deps.marker is None: + if not deps.allowMarkerSearch: + logger.warning( + f"Markers for cluster {cluster_id!r} are unavailable because no " + "marker artifact was supplied or authorized" + ) + return ClusterMarkerEvidence( + clusterId=cluster_id, + evidenceId="", + warnings=[ + "No exact marker artifact was supplied and marker search was not authorized." + ], + ) + if deps.markerFeatures is None: + logger.warning( + f"Markers for cluster {cluster_id!r} are unavailable because no " + "feature selection was supplied" + ) + return ClusterMarkerEvidence( + clusterId=cluster_id, + evidenceId="", + warnings=["Marker search requires an exact feature selection."], + ) + logger.info( + f"Creating one marker artifact for assay {deps.markerAssay!r} " + f"from cluster artifact {deps.cluster.artifact_id!r}" + ) + deps.marker = deps.store.run_marker_search( + deps.cluster, + features=deps.markerFeatures, + ) + if not hasattr(deps.marker, "artifact_id"): + raise RuntimeError("marker search did not return an artifact reference") + deps.marker = core_artifact_reference(deps.marker) + + marker_artifact = artifact_reference(deps.marker) + if marker_artifact.kind != "marker_table": + raise ModelRetry("marker must identify a marker_table artifact") + if deps.markerAssay is not None and marker_artifact.assay != deps.markerAssay: + raise ModelRetry("marker artifact belongs to a different assay") + if hasattr(deps.store, "inspect_artifact"): + marker_status = deps.store.inspect_artifact(deps.marker) + if not getattr(marker_status, "exists", True): + raise ModelRetry("marker artifact does not exist") + if not getattr(marker_status, "complete", False): + raise ModelRetry("marker artifact is incomplete") + marker_inputs = getattr(marker_status, "inputs", None) or {} + stored_clusters = marker_inputs.get("clusters") + expected_cluster = artifact_reference(deps.cluster) + if ( + not isinstance(stored_clusters, Mapping) + or stored_clusters.get("artifact_id") != expected_cluster.artifactId + or stored_clusters.get("kind") != expected_cluster.kind + or stored_clusters.get("scope") != expected_cluster.scope + or stored_clusters.get("assay") != expected_cluster.assay + ): + raise ModelRetry( + "marker artifact is not linked to the exact cluster artifact" + ) + + frame = deps.store.get_markers( + deps.marker, + group_id=deps.clusterValues[cluster_id], + min_score=deps.markerMinScore, + min_frac_exp=deps.markerMinFraction, + ) + if "score" in frame.columns: + frame = frame.sort_values("score", ascending=False, na_position="last") + markers = [ + _marker_feature(row) + for row in frame.head(min(deps.maxMarkers, _MAX_MARKERS)).to_dict("records") + ] + cluster_artifact = artifact_reference(deps.cluster) + evidence_id = ( + f"markers:{marker_artifact.artifactId}:clusters:" + f"{cluster_artifact.artifactId}:cluster:{cluster_id}" + ) + if markers: + deps.evidenceIds.add(evidence_id) + deps.markerEvidenceIds[cluster_id] = evidence_id + evidence = ClusterMarkerEvidence( + clusterId=cluster_id, + markers=markers, + markerArtifact=marker_artifact, + evidenceId=evidence_id if markers else "", + warnings=[] if markers else ["No markers passed the requested thresholds."], + ) + deps.markerEvidence[cluster_id] = evidence + logger.debug( + f"Completed marker inspection for cluster {cluster_id!r}: " + f"markers={len(markers)}" + ) + return evidence + + +async def inspect_cluster_markers_batch( + ctx: RunContext[BiologicalInterpretationDependencies], + cluster_ids: list[str], +) -> ClusterMarkerBatchEvidence: + """Inspect every selected cluster in one bounded model tool call.""" + if not ctx.deps.clusterValues: + raise ModelRetry("Call inspect_cluster_composition before inspecting markers.") + if not cluster_ids: + raise ModelRetry("cluster_ids must contain at least one observed cluster") + if len(cluster_ids) > ctx.deps.maxClusters: + raise ModelRetry( + f"cluster_ids may contain at most {ctx.deps.maxClusters} values" + ) + if len(set(cluster_ids)) != len(cluster_ids): + raise ModelRetry("cluster_ids must not contain duplicates") + if ctx.deps.markerBatch is not None: + if cluster_ids != ctx.deps.markerBatchClusterIds: + raise ModelRetry( + "Marker inspection already completed. Use the returned evidence " + "and do not request a different cluster batch." + ) + logger.info("Reused completed cluster marker batch") + return ctx.deps.markerBatch + + logger.info(f"Inspecting markers for {len(cluster_ids)} cluster(s) in one batch") + clusters = [ + await inspect_cluster_markers(ctx, cluster_id=cluster_id) + for cluster_id in cluster_ids + ] + evidence_ids = [cluster.evidenceId for cluster in clusters if cluster.evidenceId] + warnings = [ + f"Cluster {cluster.clusterId}: {warning}" + for cluster in clusters + for warning in cluster.warnings + ] + ctx.deps.toolCalls.append("inspect_cluster_markers_batch") + evidence = ClusterMarkerBatchEvidence( + clusters=clusters, + evidenceIds=evidence_ids, + warnings=warnings, + ) + ctx.deps.markerBatch = evidence + ctx.deps.markerBatchClusterIds = list(cluster_ids) + logger.info( + f"Completed marker batch inspection: clusters={len(clusters)}, " + f"clusters_with_markers={sum(bool(cluster.markers) for cluster in clusters)}, " + f"evidence_records={len(evidence_ids)}" + ) + return evidence + + +def _marker_feature(row: dict[str, Any]) -> MarkerFeature: + raw_index = _finite_float(row.get("feature_index")) + return MarkerFeature( + featureId=str(row.get("feature_id", "")), + featureName=str(row.get("feature_name", "")), + featureIndex=int(raw_index) if raw_index is not None else None, + score=_finite_float(row.get("score")), + foldChange=_finite_float(row.get("fold_change")), + fractionExpressed=_finite_float(row.get("frac_exp")), + fractionExpressedRest=_finite_float(row.get("frac_exp_rest")), + mean=_finite_float(row.get("mean")), + meanRest=_finite_float(row.get("mean_rest")), + auc=_finite_float(row.get("auc")), + adjustedPvalue=_finite_float(row.get("p_value_adjusted")), + ) diff --git a/scarf/agent/biological_interpretation/validation.py b/scarf/agent/biological_interpretation/validation.py new file mode 100644 index 00000000..4d160a72 --- /dev/null +++ b/scarf/agent/biological_interpretation/validation.py @@ -0,0 +1,566 @@ +"""Validation and dependency preparation for biological interpretation.""" + +import math +from collections.abc import Mapping +from typing import Any + +import numpy as np + +from ...storage.refs import ArtifactRef +from ...storage.selections import read_stored_selection_indices +from ...utils.logging import logger +from ..tools import artifact_reference, core_artifact_reference +from ..types import AgentRunInfo, ExperimentalBiologyHandoff, TuningBiologyHandoff +from .contracts import ( + _MAX_CLUSTERS, + _MAX_MARKERS, + BiologicalInterpretationDependencies, + BiologicalInterpretationReport, + ClusterInterpretation, + TreatmentDirection, + TreatmentObservation, +) + +try: + from pydantic_ai import ( + ModelRetry, + RunContext, + UnexpectedModelBehavior, + UsageLimitExceeded, + ) + from pydantic_ai.tools import ToolDefinition +except ImportError as exc: + from .._deps import AGENT_INSTALL_HINT + + raise ImportError(AGENT_INSTALL_HINT) from exc + + +def _prepare_biological_interpretation_tool( + ctx: RunContext[BiologicalInterpretationDependencies], + tool_definition: ToolDefinition, +) -> ToolDefinition | None: + """Expose composition and marker batches once and in the required order.""" + completed_calls = set(ctx.deps.toolCalls) + if tool_definition.name == "inspect_cluster_composition": + return None if tool_definition.name in completed_calls else tool_definition + if tool_definition.name == "inspect_cluster_markers_batch": + if ( + "inspect_cluster_composition" not in completed_calls + or tool_definition.name in completed_calls + ): + return None + return tool_definition + return tool_definition + + +def _canonicalize_cluster_interpretations( + report: BiologicalInterpretationReport, + deps: BiologicalInterpretationDependencies, +) -> tuple[list[ClusterInterpretation], list[str]]: + canonical_interpretations: list[ClusterInterpretation] = [] + omitted_interpretation_clusters: list[str] = [] + for interpretation in report.clusterInterpretations: + marker_id = deps.markerEvidenceIds.get(interpretation.clusterId) + if marker_id is None: + omitted_interpretation_clusters.append(interpretation.clusterId) + continue + non_marker_evidence = sorted(set(interpretation.evidenceIds) - {marker_id}) + if non_marker_evidence: + raise ModelRetry( + "Cluster identity interpretations may cite only their exact marker " + f"evidence: {non_marker_evidence}" + ) + canonical_interpretations.append( + interpretation.model_copy( + update={ + "evidenceIds": [marker_id], + "identityIsHypothesis": True, + **( + { + "confidence": "low", + } + if deps.markerAssayType == "ATAC" + else {} + ), + } + ) + ) + return canonical_interpretations, omitted_interpretation_clusters + + +def _canonicalize_treatment_observations( + report: BiologicalInterpretationReport, + deps: BiologicalInterpretationDependencies, +) -> list[TreatmentObservation]: + if report.treatmentObservations and deps.conditionColumn is None: + raise ModelRetry("Treatment observations require a condition column.") + if report.treatmentObservations and deps.sampleColumn is None: + raise ModelRetry( + "Treatment observations require independent-unit composition summaries." + ) + if report.treatmentObservations: + handoff = deps.designHandoff + if ( + handoff is None + or not handoff.conditionColumn + or handoff.conditionColumn != deps.conditionColumn + or not handoff.independentUnit + or handoff.independentUnit != deps.sampleColumn + or handoff.coefficientScope != "betweenUnit" + or handoff.estimability.get("status") != "ok" + or handoff.estimability.get("coefficientEstimable") is not True + ): + raise ModelRetry( + "Treatment observations require an explicit condition, aggregation " + "at the independent unit, a between-unit coefficient, and an " + "estimable experimental contrast." + ) + + canonical_observations: list[TreatmentObservation] = [] + for observation in report.treatmentObservations: + if not observation.isDescriptiveOnly: + raise ModelRetry("Treatment observations must remain descriptive.") + if len(observation.evidenceIds) != 2 or len(set(observation.evidenceIds)) != 2: + raise ModelRetry( + "Every treatment observation must cite exactly two distinct " + "condition summaries." + ) + if any( + evidence_id not in deps.conditionEvidence + for evidence_id in observation.evidenceIds + ): + raise ModelRetry( + "Treatment observations may cite only condition composition evidence." + ) + summaries = [ + deps.conditionEvidence[evidence_id] + for evidence_id in observation.evidenceIds + ] + if any(summary.clusterId != observation.clusterId for summary in summaries): + raise ModelRetry( + "Every treatment observation must cite condition summaries for " + "its exact cluster." + ) + if ( + not observation.referenceCondition + or not observation.comparisonCondition + or observation.referenceCondition == observation.comparisonCondition + ): + raise ModelRetry( + "Treatment observations require two distinct named conditions." + ) + summaries_by_condition = {summary.condition: summary for summary in summaries} + expected_conditions = { + observation.referenceCondition, + observation.comparisonCondition, + } + if set(summaries_by_condition) != expected_conditions: + raise ModelRetry( + "Treatment observation conditions must match the two cited " + "condition summaries." + ) + if any(summary.nSamples < 2 for summary in summaries): + raise ModelRetry( + "Sample-level treatment observations require at least two samples " + "in every cited condition." + ) + reference = summaries_by_condition[observation.referenceCondition] + comparison = summaries_by_condition[observation.comparisonCondition] + if math.isclose( + comparison.meanFraction, + reference.meanFraction, + rel_tol=1e-9, + abs_tol=1e-12, + ): + expected_direction: TreatmentDirection = "equal" + elif comparison.meanFraction > reference.meanFraction: + expected_direction = "higher" + else: + expected_direction = "lower" + if observation.direction != expected_direction: + raise ModelRetry( + "Treatment observation direction does not match the cited mean " + "independent-unit fractions." + ) + if expected_direction == "equal": + canonical_text = ( + f"Cluster {observation.clusterId} has equal mean independent-unit " + f"fractions in {comparison.condition} and {reference.condition} " + f"({comparison.meanFraction:.6g}); this is descriptive only." + ) + else: + canonical_text = ( + f"Cluster {observation.clusterId} has a {expected_direction} mean " + f"independent-unit fraction in {comparison.condition} " + f"({comparison.meanFraction:.6g}) than in {reference.condition} " + f"({reference.meanFraction:.6g}); this is descriptive only." + ) + canonical_observations.append( + observation.model_copy(update={"observation": canonical_text}) + ) + return canonical_observations + + +def validate_biological_interpretation_report( + report: BiologicalInterpretationReport, + deps: BiologicalInterpretationDependencies, +) -> BiologicalInterpretationReport: + """Reject invented evidence, clusters, or completed marker-free reviews.""" + if not deps.clusterValues: + raise ModelRetry("Call inspect_cluster_composition before returning a report.") + if report.status == "failed": + raise ModelRetry( + "Do not return failed for biological uncertainty; return needsInput " + "with one concrete question instead." + ) + if report.status == "needsInput" and ( + report.needsInput is None or not report.needsInput.question.strip() + ): + raise ModelRetry("A needsInput report requires one concrete input question.") + if report.status != "needsInput" and report.needsInput is not None: + raise ModelRetry("Only a needsInput report may include an input question.") + expected_cluster_artifact = artifact_reference(deps.cluster) + if ( + report.clusterArtifact is not None + and report.clusterArtifact != expected_cluster_artifact + ): + raise ModelRetry("Report clusterArtifact does not match the inspected artifact") + if deps.marker is not None: + expected_marker_artifact = artifact_reference(deps.marker) + if ( + report.markerArtifact is not None + and report.markerArtifact != expected_marker_artifact + ): + raise ModelRetry( + "Report markerArtifact does not match the inspected artifact" + ) + + cited = set(report.evidenceIds) + for interpretation in report.clusterInterpretations: + cited.update(interpretation.evidenceIds) + for observation in report.treatmentObservations: + cited.update(observation.evidenceIds) + for follow_up in report.followUps: + cited.update(follow_up.evidenceIds) + if report.needsInput is not None: + cited.update(report.needsInput.evidenceIds) + unknown = cited.difference(deps.evidenceIds) + if unknown: + raise ModelRetry(f"Unknown evidenceIds: {sorted(unknown)}") + interpreted_clusters = {item.clusterId for item in report.clusterInterpretations} + observed_clusters = {item.clusterId for item in report.treatmentObservations} + unknown_clusters = (interpreted_clusters | observed_clusters).difference( + deps.clusterValues + ) + if unknown_clusters: + raise ModelRetry(f"Unknown cluster ids: {sorted(unknown_clusters)}") + canonical_interpretations, omitted_interpretation_clusters = ( + _canonicalize_cluster_interpretations(report, deps) + ) + canonical_observations = _canonicalize_treatment_observations(report, deps) + if report.status == "done" and not canonical_interpretations: + raise ModelRetry( + "A done report must contain at least one cluster interpretation with " + "non-empty marker evidence." + ) + limitations = list(report.limitations) + if deps.markerAssayType == "ATAC": + atac_limitation = ( + "ATAC peak markers are descriptive, so all cell identities remain " + "low-confidence hypotheses." + ) + if atac_limitation not in limitations: + limitations.append(atac_limitation) + if omitted_interpretation_clusters: + omitted_clusters = ", ".join(sorted(set(omitted_interpretation_clusters))) + marker_limitation = ( + "Cluster identity interpretations without non-empty marker evidence " + f"were omitted for clusters: {omitted_clusters}." + ) + if marker_limitation not in limitations: + limitations.append(marker_limitation) + if canonical_observations: + descriptive_limitation = ( + "Independent-unit cluster fractions are descriptive summaries, not " + "tests of significance or causal treatment effects." + ) + if descriptive_limitation not in limitations: + limitations.append(descriptive_limitation) + validated = report.model_copy( + update={ + "clusterInterpretations": canonical_interpretations, + "treatmentObservations": canonical_observations, + "evidenceIds": sorted( + { + *report.evidenceIds, + *( + evidence_id + for interpretation in canonical_interpretations + for evidence_id in interpretation.evidenceIds + ), + } + ), + "limitations": limitations, + "clusterArtifact": expected_cluster_artifact, + "markerArtifact": ( + artifact_reference(deps.marker) if deps.marker is not None else None + ), + "graphAssay": deps.graphAssay, + "markerAssay": deps.markerAssay, + } + ) + logger.debug( + f"Validated biological interpretation report: status={validated.status}, " + f"cluster_interpretations={len(validated.clusterInterpretations)}, " + f"treatment_observations={len(validated.treatmentObservations)}, " + f"omitted_interpretations={len(omitted_interpretation_clusters)}" + ) + return validated + + +def fallback_biological_interpretation_report( + deps: BiologicalInterpretationDependencies, + *, + error: UnexpectedModelBehavior | UsageLimitExceeded, + model_name: str, +) -> BiologicalInterpretationReport: + """Keep measured evidence without claiming completed model interpretation.""" + if not deps.clusterValues: + raise error + from ..config.agent_exec import describe_agent_error + + error_detail = describe_agent_error(error) + report = BiologicalInterpretationReport( + status="failed", + clusterArtifact=artifact_reference(deps.cluster) + if deps.cluster is not None + else None, + markerArtifact=artifact_reference(deps.marker) + if deps.marker is not None + else None, + graphAssay=deps.graphAssay, + markerAssay=deps.markerAssay, + evidenceIds=sorted({*deps.evidenceIds, *deps.markerEvidenceIds.values()}), + limitations=[ + "Measured composition and marker evidence remain available. No biological " + "interpretation or treatment observation was accepted after model failure.", + error_detail, + ], + stopReason="The model's biological interpretation could not be validated.", + runInfo=getattr( + error, + "agent_run_info", + AgentRunInfo( + agentName="biological_interpretation_failed", modelName=model_name + ), + ), + ) + logger.warning( + f"Biological interpretation failed; observed evidence was retained: {error_detail}" + ) + return report + + +def _prepare_biological_interpretation_dependencies( + store: Any, + *, + cluster: Any, + from_assay: str | None, + graph_assay: str | None, + marker_assay_type: str | None, + sample_column: str | None, + condition_column: str | None, + tuning_handoff: TuningBiologyHandoff | None, + experimental_handoff: ExperimentalBiologyHandoff | None, + marker: Any, + marker_features: Any, + allow_marker_search: bool, + max_clusters: int, + max_markers: int, + marker_min_score: float, + marker_min_fraction: float, +) -> BiologicalInterpretationDependencies: + expected_selections: list[ArtifactRef] = [] + if tuning_handoff is not None: + if tuning_handoff.clusterArtifact is None: + raise ValueError("tuning_handoff lacks a cluster artifact") + tuning_selection = core_artifact_reference(tuning_handoff.cellSelection) + if not isinstance(tuning_selection, ArtifactRef): + raise ValueError("tuning_handoff lacks an exact cell selection") + expected_selections.append(tuning_selection) + if cluster is not None and ( + artifact_reference(cluster) != tuning_handoff.clusterArtifact + ): + raise ValueError("cluster conflicts with tuning_handoff") + if from_assay is not None and from_assay != tuning_handoff.fromAssay: + raise ValueError("from_assay conflicts with tuning_handoff") + if graph_assay is not None and graph_assay != tuning_handoff.graphAssay: + raise ValueError("graph_assay conflicts with tuning_handoff") + cluster = tuning_handoff.clusterArtifact + from_assay = tuning_handoff.fromAssay + graph_assay = tuning_handoff.graphAssay + if experimental_handoff is not None: + experimental_selection = core_artifact_reference( + experimental_handoff.cellSelection + ) + if not isinstance(experimental_selection, ArtifactRef): + raise ValueError("experimental_handoff lacks an exact cell selection") + expected_selections.append(experimental_selection) + if len(expected_selections) == 2 and ( + expected_selections[0] != expected_selections[1] + ): + raise ValueError( + "Experimental and tuning handoffs use different cell selections" + ) + if ( + condition_column is not None + and condition_column != experimental_handoff.conditionColumn + ): + raise ValueError("condition_column conflicts with experimental_handoff") + aggregation_unit = ( + experimental_handoff.independentUnit or experimental_handoff.observationUnit + ) + if sample_column is not None and sample_column != aggregation_unit: + raise ValueError("sample_column conflicts with experimental_handoff") + condition_column = experimental_handoff.conditionColumn + sample_column = aggregation_unit + if cluster is None: + raise ValueError("cluster must identify an exact cluster artifact") + if not 1 <= max_clusters <= _MAX_CLUSTERS: + raise ValueError(f"max_clusters must be between 1 and {_MAX_CLUSTERS}") + if not 1 <= max_markers <= _MAX_MARKERS: + raise ValueError(f"max_markers must be between 1 and {_MAX_MARKERS}") + if not 0 < marker_min_score <= 1: + raise ValueError("marker_min_score must be greater than 0 and at most 1") + if not 0 <= marker_min_fraction <= 1: + raise ValueError("marker_min_fraction must be between 0 and 1") + if allow_marker_search and marker is None and marker_features is None: + raise ValueError("marker_features is required when marker search is authorized") + + cluster = core_artifact_reference(cluster) + marker = core_artifact_reference(marker) + marker_features = core_artifact_reference(marker_features) + if not isinstance(cluster, ArtifactRef): + raise TypeError("cluster must be an ArtifactRef") + if marker is not None and ( + not isinstance(marker, ArtifactRef) or marker.kind != "marker_table" + ): + raise TypeError("marker must be a marker_table ArtifactRef") + if marker_features is not None and ( + not isinstance(marker_features, ArtifactRef) + or marker_features.kind != "feature_selection" + ): + raise TypeError("marker_features must be a feature_selection ArtifactRef") + cluster_artifact = artifact_reference(cluster) + if cluster_artifact.kind not in {"cluster_labels", "cluster_cut"}: + raise ValueError( + "cluster must identify a cluster_labels or cluster_cut artifact" + ) + if cluster_artifact.scope == "datastore" and cluster_artifact.assay is not None: + raise ValueError("datastore-scoped cluster artifacts must not name an assay") + if ( + tuning_handoff is not None + and cluster_artifact.scope == "datastore" + and not tuning_handoff.markerAssay + ): + raise ValueError( + "Integrated tuning handoffs must explicitly identify markerAssay" + ) + resolved_graph_assay = graph_assay or cluster_artifact.assay + if ( + resolved_graph_assay is not None + and cluster_artifact.scope == "assay" + and cluster_artifact.assay != resolved_graph_assay + ): + raise ValueError("cluster belongs to a different assay") + resolved_marker_assay = ( + tuning_handoff.markerAssay or cluster_artifact.assay + if tuning_handoff is not None + else ( + marker.assay + if isinstance(marker, ArtifactRef) + else ( + marker_features.assay + if isinstance(marker_features, ArtifactRef) + else from_assay or cluster_artifact.assay + ) + ) + ) + if cluster_artifact.scope == "datastore" and not resolved_marker_assay: + raise ValueError( + "from_assay is required to resolve markers for integrated clusters" + ) + if isinstance(marker, ArtifactRef) and marker.assay != resolved_marker_assay: + raise ValueError("marker artifact belongs to a different marker assay") + if ( + isinstance(marker_features, ArtifactRef) + and marker_features.assay != resolved_marker_assay + ): + raise ValueError("marker feature selection belongs to a different assay") + + cluster_status = store.inspect_artifact(cluster) + if not getattr(cluster_status, "exists", True): + raise ValueError("cluster artifact does not exist") + if not getattr(cluster_status, "complete", False): + raise ValueError("cluster artifact is incomplete") + raw_selection = (getattr(cluster_status, "inputs", None) or {}).get( + "cell_selection" + ) + if not isinstance(raw_selection, Mapping): + raise ValueError("cluster artifact has no cell-selection input") + cell_selection = ArtifactRef.from_dict(dict(raw_selection)) + if ( + cell_selection.scope != "datastore" + or cell_selection.kind != "cell_selection" + or cell_selection.assay is not None + ): + raise ValueError("cluster artifact has an invalid cell-selection input") + if any(selection != cell_selection for selection in expected_selections): + raise ValueError("handoff cell selection conflicts with cluster") + cell_indices = read_stored_selection_indices( + store.zw, + cell_selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + if isinstance(marker, ArtifactRef): + marker_status = store.inspect_artifact(marker) + if not getattr(marker_status, "exists", True): + raise ValueError("marker artifact does not exist") + if not getattr(marker_status, "complete", False): + raise ValueError("marker artifact is incomplete") + marker_inputs = getattr(marker_status, "inputs", None) or {} + stored_clusters = marker_inputs.get("clusters") + expected_cluster = artifact_reference(cluster) + if ( + not isinstance(stored_clusters, Mapping) + or stored_clusters.get("artifact_id") != expected_cluster.artifactId + or stored_clusters.get("kind") != expected_cluster.kind + or stored_clusters.get("scope") != expected_cluster.scope + or stored_clusters.get("assay") != expected_cluster.assay + ): + raise ValueError( + "marker artifact is not linked to the exact cluster artifact" + ) + return BiologicalInterpretationDependencies( + store=store, + cluster=cluster, + cellSelection=cell_selection, + cellIndices=cell_indices, + fromAssay=from_assay or cluster_artifact.assay or resolved_marker_assay, + graphAssay=resolved_graph_assay, + markerAssay=resolved_marker_assay, + markerAssayType=marker_assay_type, + sampleColumn=sample_column, + conditionColumn=condition_column, + designHandoff=experimental_handoff, + marker=marker, + markerFeatures=marker_features, + allowMarkerSearch=allow_marker_search, + maxClusters=max_clusters, + maxMarkers=max_markers, + markerMinScore=marker_min_score, + markerMinFraction=marker_min_fraction, + ) diff --git a/scarf/agent/cell_quality/__init__.py b/scarf/agent/cell_quality/__init__.py new file mode 100644 index 00000000..54a6eb78 --- /dev/null +++ b/scarf/agent/cell_quality/__init__.py @@ -0,0 +1 @@ +"""Registered cell-quality profiles and execution.""" diff --git a/scarf/agent/cell_quality/execution.py b/scarf/agent/cell_quality/execution.py new file mode 100644 index 00000000..856485f8 --- /dev/null +++ b/scarf/agent/cell_quality/execution.py @@ -0,0 +1,714 @@ +"""Persist exact outputs from a registered agent cell-quality decision.""" + +from collections.abc import Iterable, Mapping +from numbers import Real +from typing import Any, cast + +import numpy as np + +from ...metadata.artifacts import ( + plan_cell_data_artifact, + write_cell_data_artifact, +) +from ...metadata.rows import read_metadata_rows_chunkwise +from ...metadata.selection import NamedCellArtifact, resolve_cell_aligned_artifact +from ...storage.artifacts import canonical_bytes, fingerprint_array, fingerprint_strings +from ...storage.refs import ArtifactRef +from ...storage.selections import ( + read_stored_selection_mask, + resolve_generated_selection_artifact, +) +from ...utils.logging import logger +from .profiles import ( + REGISTERED_CELL_QC_PROFILES, + AutoFilterAction, + RegisteredCellQcProfile, + project_auto_filter_profile, + project_registered_qc_profile, + qc_metric_execution_name, +) + + +def _validated_named_cell_artifacts( + values: Iterable[NamedCellArtifact] | None, + *, + expected_kind: str, + label: str, +) -> list[NamedCellArtifact]: + sources = list(values or ()) + names: set[str] = set() + for source in sources: + if not isinstance(source, NamedCellArtifact): + raise TypeError(f"{label} must contain NamedCellArtifact values") + if source.artifact.kind != expected_kind: + raise ValueError(f"{label} must reference {expected_kind!r} artifacts") + if source.name in names: + raise ValueError(f"{label} must use unique semantic names") + names.add(source.name) + return sources + + +def execute_registered_cell_qc( + store: Any, + profile: RegisteredCellQcProfile, + *, + profile_parameters: Mapping[str, Any], + expected_active_cells: int, + expected_retained_cells: int, + expected_flag_counts: Mapping[str, int], + attrs: Iterable[str] | None = None, + artifact_metrics: Iterable[NamedCellArtifact] | None = None, + cell_selection: ArtifactRef | None = None, + sample_column: str | None = None, + sample_artifact: NamedCellArtifact | None = None, + invalidate_cache: bool = False, +) -> tuple[ArtifactRef, ArtifactRef | None]: + """Recompute, verify, and persist one registered cell-QC decision.""" + if profile not in REGISTERED_CELL_QC_PROFILES: + raise ValueError(f"Unknown registered cell-QC profile {profile!r}") + if ( + isinstance(expected_active_cells, bool) + or not isinstance(expected_active_cells, int) + or expected_active_cells < 1 + ): + raise ValueError("expected_active_cells must be a positive integer") + if ( + isinstance(expected_retained_cells, bool) + or not isinstance(expected_retained_cells, int) + or expected_retained_cells < 0 + ): + raise ValueError("expected_retained_cells must be a non-negative integer") + flag_counts = dict(expected_flag_counts) + if any( + not isinstance(name, str) + or not name + or isinstance(count, bool) + or not isinstance(count, int) + or count < 0 + for name, count in flag_counts.items() + ): + raise ValueError( + "expected_flag_counts must map non-empty names to non-negative integers" + ) + + parameters = dict(profile_parameters) + required_parameter_keys = { + "policyVersion", + "profile", + "nMads", + "boundPolicy", + "resolvedBounds", + "captureSizes", + "captureComparisons", + "captureComparisonSource", + "pooledReferenceCaptures", + } + if set(parameters) != required_parameter_keys: + raise ValueError("Registered cell-QC parameters do not match policy version 1") + if ( + isinstance(parameters["policyVersion"], bool) + or parameters["policyVersion"] != 1 + ): + raise ValueError("Registered cell-QC policyVersion must be 1") + if parameters["profile"] != profile: + raise ValueError("Registered cell-QC profile and parameters disagree") + expected_n_mads = 3.0 if profile == "captureMad3Sensitivity" else 5.0 + if ( + isinstance(parameters["nMads"], bool) + or not isinstance(parameters["nMads"], Real) + or float(parameters["nMads"]) != expected_n_mads + ): + raise ValueError( + f"Registered cell-QC profile {profile!r} requires nMads={expected_n_mads}" + ) + expected_bound_policy = { + "count": {"remove": "lower", "flag": "upper"}, + "feature": {"remove": "lower", "flag": "upper"}, + "mitochondrial": {"remove": "upper", "fixedCutoff": None}, + "diagnostic": {"remove": "none"}, + } + if parameters["boundPolicy"] != expected_bound_policy: + raise ValueError("Registered cell-QC boundPolicy is not supported") + raw_bounds = parameters["resolvedBounds"] + if not isinstance(raw_bounds, list) or any( + not isinstance(value, Mapping) for value in raw_bounds + ): + raise TypeError("Registered cell-QC resolvedBounds must be a list of maps") + pooled_references = parameters["pooledReferenceCaptures"] + if not isinstance(pooled_references, list) or any( + not isinstance(value, str) for value in pooled_references + ): + raise TypeError( + "Registered cell-QC pooledReferenceCaptures must be a list of strings" + ) + canonical_bytes(parameters) + + attrs_list = list(attrs or ()) + if any(not isinstance(attr, str) for attr in attrs_list): + raise TypeError("attrs must contain only column names") + metric_artifacts = _validated_named_cell_artifacts( + artifact_metrics, + expected_kind="quality_metric", + label="artifact_metrics", + ) + resolved_sample_artifact: NamedCellArtifact | None = None + if sample_artifact is not None: + resolved_sample_artifact = _validated_named_cell_artifacts( + [sample_artifact], + expected_kind="hto_identity", + label="sample_artifact", + )[0] + if sample_column is not None and resolved_sample_artifact is not None: + raise ValueError("sample_column and sample_artifact are mutually exclusive") + capture_profile = profile in { + "captureMad5", + "captureMad3Sensitivity", + "pooledReferenceMad5", + } + has_capture_source = (sample_column is None) != (resolved_sample_artifact is None) + if capture_profile and not has_capture_source: + raise ValueError( + f"Registered cell-QC profile {profile!r} requires one capture source" + ) + if not capture_profile and has_capture_source: + raise ValueError( + f"Registered cell-QC profile {profile!r} cannot use a capture source" + ) + if sample_column is not None and sample_column not in store.cells.columns: + raise ValueError(f"sample_column '{sample_column}' not found in cell metadata") + missing = [attr for attr in attrs_list if attr not in store.cells.columns] + if missing: + joined = ", ".join(repr(attr) for attr in missing) + raise KeyError(f"Cell metadata columns not found: {joined}") + artifact_names = {source.name for source in metric_artifacts} + metadata_collisions = set(attrs_list).intersection(artifact_names) + execution_artifacts = [ + NamedCellArtifact( + name=qc_metric_execution_name( + source.name, + artifact_id=source.artifact.artifact_id, + collides_with_metadata=source.name in metadata_collisions, + ), + artifact=source.artifact, + ) + for source in metric_artifacts + ] + + prior = store._filter_input_selection(cell_selection) + active = read_stored_selection_mask( + store.zw, + prior, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ) + active_idx = np.flatnonzero(active).astype(np.int64, copy=False) + if len(active_idx) != expected_active_cells: + raise ValueError( + "Registered cell-QC active-cell count differs from its evidence" + ) + + values_by_name: dict[str, np.ndarray] = {} + metadata_fingerprints: dict[str, str] = {} + for attr in attrs_list: + values = np.asarray( + read_metadata_rows_chunkwise(store.cells, attr, active_idx), + dtype=float, + ) + if values.shape != (len(active_idx),): + raise ValueError( + f"QC metadata column {attr!r} does not align with cell_selection" + ) + if not np.isfinite(values).all(): + raise ValueError(f"QC values in {attr!r} contain non-finite entries") + values_by_name[attr] = values + metadata_fingerprints[attr] = fingerprint_array(values) + for source, execution_source in zip( + metric_artifacts, + execution_artifacts, + strict=True, + ): + resolved = resolve_cell_aligned_artifact( + store.zw, + source.artifact, + cell_selection=prior, + expected_kind="quality_metric", + ) + values = np.asarray(resolved.values, dtype=float) + if not np.isfinite(values).all(): + raise ValueError( + f"QC artifact values in {source.name!r} contain non-finite entries" + ) + values_by_name[execution_source.name] = values + + sample_labels: np.ndarray | None = None + sample_inputs: dict[str, Any] = {} + if sample_column is not None: + sample_labels = np.asarray( + read_metadata_rows_chunkwise(store.cells, sample_column, active_idx) + ) + sample_inputs["capture_assignments_fingerprint"] = fingerprint_strings( + sample_labels + ) + elif resolved_sample_artifact is not None: + resolved_sample = resolve_cell_aligned_artifact( + store.zw, + resolved_sample_artifact.artifact, + cell_selection=prior, + expected_kind="hto_identity", + ) + sample_labels = np.asarray(resolved_sample.values) + sample_inputs["capture_artifact"] = resolved_sample_artifact.artifact + + projection = project_registered_qc_profile( + profile, + values_by_metric=values_by_name, + active=np.ones(len(active_idx), dtype=bool), + capture_labels=sample_labels, + grouping_proven=capture_profile, + min_cells_per_capture=20, + pooled_reference_captures=tuple(pooled_references), + ) + recomputed_bounds = [value.to_dict() for value in projection.thresholds] + if canonical_bytes(recomputed_bounds) != canonical_bytes(raw_bounds): + raise ValueError("Registered cell-QC resolved bounds do not match the inputs") + if capture_profile: + if projection.captureSizes != parameters["captureSizes"]: + raise ValueError( + "Registered cell-QC capture sizes do not match the exact inputs" + ) + comparisons = [ + comparison.to_dict() for comparison in projection.captureComparisons + ] + if canonical_bytes(comparisons) != canonical_bytes( + parameters["captureComparisons"] + ): + raise ValueError( + "Registered cell-QC capture comparisons do not match the inputs" + ) + if projection.retainedCells != expected_retained_cells: + raise ValueError( + "Registered cell-QC retained-cell count differs from its evidence" + ) + if projection.flagCounts != flag_counts: + raise ValueError( + "Registered cell-QC diagnostic-flag counts differ from their evidence" + ) + + metric_sources: list[dict[str, Any]] = [ + { + "name": attr, + "executionName": attr, + "source": "metadataColumn", + "column": attr, + } + for attr in attrs_list + ] + metric_sources.extend( + { + "name": source.name, + "executionName": execution_source.name, + "source": "artifact", + "artifact": source.artifact.to_dict(), + } + for source, execution_source in zip( + metric_artifacts, + execution_artifacts, + strict=True, + ) + ) + source_inputs: dict[str, Any] = { + "metadata_fingerprints": metadata_fingerprints, + "artifact_metrics": { + execution_source.name: source.artifact + for source, execution_source in zip( + metric_artifacts, + execution_artifacts, + strict=True, + ) + }, + **sample_inputs, + } + execution_parameters: dict[str, Any] = { + "profile": profile, + "metricSources": metric_sources, + "profileParameters": parameters, + "flagCounts": projection.flagCounts, + } + canonical_bytes( + { + "operation": "run_registered_cell_qc", + "parameters": execution_parameters, + "inputs": {"prior_cell_selection": prior, **source_inputs}, + } + ) + + flag_names = tuple(sorted(projection.flags)) + flag_ref: ArtifactRef | None = None + if flag_names: + flag_values = np.column_stack( + [projection.flags[name] for name in flag_names] + ).astype(bool, copy=False) + planned_flags = plan_cell_data_artifact( + store.zw, + scope="datastore", + kind="metadata_snapshot", + operation="run_registered_cell_qc_flags", + parameters={ + **execution_parameters, + "flagNames": list(flag_names), + }, + inputs=source_inputs, + execution_options={}, + cell_selection=prior, + arrays={"values": (flag_values.shape, "b")}, + invalidate_cache=invalidate_cache, + ) + write_cell_data_artifact( + store.zw, + planned_flags, + {"values": flag_values}, + fingerprint_payload=True, + ) + flag_ref = planned_flags.ref + + keep = np.zeros(store.cells.N, dtype=bool) + keep[active_idx] = projection.keep + selection_inputs: dict[str, Any] = { + "prior_cell_selection": prior, + **source_inputs, + } + if flag_ref is not None: + selection_inputs["diagnostic_flags"] = flag_ref + ref, stored = resolve_generated_selection_artifact( + store.zw, + scope="datastore", + kind="cell_selection", + values=keep, + row_ids=np.asarray(store.cells.fetch_all("ids")), + operation="run_registered_cell_qc", + parameters=execution_parameters, + inputs=selection_inputs, + source_column="artifact", + invalidate_cache=invalidate_cache, + ) + logger.info( + f"Registered cell QC {profile!r} retained " + f"{int(stored.sum())}/{store.cells.N} cells" + ) + return ref, flag_ref + + +def execute_auto_cell_qc( + store: Any, + action: AutoFilterAction, + *, + profile_parameters: Mapping[str, Any], + expected_active_cells: int, + expected_retained_cells: int, + expected_flag_counts: Mapping[str, int], + expected_resolved_bounds: Mapping[str, Any], + attrs: Iterable[str] | None = None, + artifact_metrics: Iterable[NamedCellArtifact] | None = None, + cell_selection: ArtifactRef | None = None, + sample_column: str | None = None, + sample_artifact: NamedCellArtifact | None = None, + capture_column: str | None = None, + capture_artifact: NamedCellArtifact | None = None, + invalidate_cache: bool = False, +) -> tuple[ArtifactRef, ArtifactRef | None]: + """Verify evidence, call core auto-filtering, and persist its exact flags.""" + if action not in {"globalGaussian", "sampleMad"}: + raise ValueError(f"Unknown automatic cell-QC action {action!r}") + attrs_list = list(attrs or ()) + if any(not isinstance(attr, str) for attr in attrs_list): + raise TypeError("attrs must contain only column names") + metric_artifacts = _validated_named_cell_artifacts( + artifact_metrics, + expected_kind="quality_metric", + label="artifact_metrics", + ) + sample_sources = _validated_named_cell_artifacts( + [sample_artifact] if sample_artifact is not None else [], + expected_kind="hto_identity", + label="sample_artifact", + ) + capture_sources = _validated_named_cell_artifacts( + [capture_artifact] if capture_artifact is not None else [], + expected_kind="hto_identity", + label="capture_artifact", + ) + resolved_sample_artifact = sample_sources[0] if sample_sources else None + resolved_capture_artifact = capture_sources[0] if capture_sources else None + if sample_column is not None and resolved_sample_artifact is not None: + raise ValueError("sample_column and sample_artifact are mutually exclusive") + if capture_column is not None and resolved_capture_artifact is not None: + raise ValueError("capture_column and capture_artifact are mutually exclusive") + if action == "sampleMad" and ( + (sample_column is None) == (resolved_sample_artifact is None) + ): + raise ValueError("sampleMad requires exactly one sample source") + if action == "globalGaussian" and ( + sample_column is not None or resolved_sample_artifact is not None + ): + raise ValueError("globalGaussian cannot use a core sample source") + for column in (sample_column, capture_column): + if column is not None and column not in store.cells.columns: + raise ValueError(f"QC grouping column {column!r} was not found") + missing = [attr for attr in attrs_list if attr not in store.cells.columns] + if missing: + raise KeyError(f"Cell metadata columns not found: {missing}") + + parameters = dict(profile_parameters) + canonical_bytes(parameters) + if ( + isinstance(expected_active_cells, bool) + or not isinstance(expected_active_cells, int) + or expected_active_cells < 1 + ): + raise ValueError("expected_active_cells must be a positive integer") + if ( + isinstance(expected_retained_cells, bool) + or not isinstance(expected_retained_cells, int) + or expected_retained_cells < 0 + ): + raise ValueError("expected_retained_cells must be a non-negative integer") + flag_counts = dict(expected_flag_counts) + if any( + not isinstance(name, str) + or not name + or isinstance(count, bool) + or not isinstance(count, int) + or count < 0 + for name, count in flag_counts.items() + ): + raise ValueError( + "expected_flag_counts must map non-empty names to non-negative integers" + ) + + prior = store._filter_input_selection(cell_selection) + active = read_stored_selection_mask( + store.zw, + prior, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ) + active_idx = np.flatnonzero(active).astype(np.int64, copy=False) + if len(active_idx) != expected_active_cells: + raise ValueError( + "Automatic cell-QC active-cell count differs from its evidence" + ) + + artifact_names = {source.name for source in metric_artifacts} + metadata_collisions = set(attrs_list).intersection(artifact_names) + execution_artifacts = [ + NamedCellArtifact( + name=qc_metric_execution_name( + source.name, + artifact_id=source.artifact.artifact_id, + collides_with_metadata=source.name in metadata_collisions, + ), + artifact=source.artifact, + ) + for source in metric_artifacts + ] + values_by_name: dict[str, np.ndarray] = {} + metadata_fingerprints: dict[str, str] = {} + for attr in attrs_list: + values = np.asarray( + read_metadata_rows_chunkwise(store.cells, attr, active_idx), + dtype=float, + ) + if values.shape != (len(active_idx),) or not np.isfinite(values).all(): + raise ValueError( + f"QC metadata column {attr!r} is not a finite aligned vector" + ) + values_by_name[attr] = values + metadata_fingerprints[attr] = fingerprint_array(values) + for source, execution_source in zip( + metric_artifacts, + execution_artifacts, + strict=True, + ): + resolved = resolve_cell_aligned_artifact( + store.zw, + source.artifact, + cell_selection=prior, + expected_kind="quality_metric", + ) + values = np.asarray(resolved.values, dtype=float) + if values.shape != (len(active_idx),) or not np.isfinite(values).all(): + raise ValueError( + f"QC artifact {source.name!r} is not a finite aligned vector" + ) + values_by_name[execution_source.name] = values + + projection_labels: np.ndarray | None = None + grouping_source: dict[str, Any] = {} + core_sample_column = sample_column + core_sample_artifact = resolved_sample_artifact + if action == "sampleMad": + if sample_column is not None: + projection_labels = np.asarray( + read_metadata_rows_chunkwise( + store.cells, + sample_column, + active_idx, + ) + ) + grouping_source = { + "source": "metadataColumn", + "column": sample_column, + "fingerprint": fingerprint_strings(projection_labels), + } + else: + assert resolved_sample_artifact is not None + resolved = resolve_cell_aligned_artifact( + store.zw, + resolved_sample_artifact.artifact, + cell_selection=prior, + expected_kind="hto_identity", + ) + projection_labels = np.asarray(resolved.values) + grouping_source = { + "source": "artifact", + "artifact": resolved_sample_artifact.artifact, + } + elif capture_column is not None: + projection_labels = np.asarray( + read_metadata_rows_chunkwise(store.cells, capture_column, active_idx) + ) + grouping_source = { + "source": "metadataColumn", + "column": capture_column, + "fingerprint": fingerprint_strings(projection_labels), + } + elif resolved_capture_artifact is not None: + resolved = resolve_cell_aligned_artifact( + store.zw, + resolved_capture_artifact.artifact, + cell_selection=prior, + expected_kind="hto_identity", + ) + projection_labels = np.asarray(resolved.values) + grouping_source = { + "source": "artifact", + "artifact": resolved_capture_artifact.artifact, + } + + n_mads = float(parameters.get("nMads", 3.0)) + min_cells = int(parameters.get("minCellsPerSample", 20)) + min_p = float(parameters.get("minP", 0.01)) + max_p = float(parameters.get("maxP", 0.99)) + projection = project_auto_filter_profile( + action, + values_by_metric=values_by_name, + active=np.ones(len(active_idx), dtype=bool), + sample_labels=projection_labels, + grouping_proven=projection_labels is not None, + min_p=min_p, + max_p=max_p, + n_mads=n_mads, + min_cells_per_sample=min_cells, + ) + expected_summary: dict[str, Any] + if action == "globalGaussian": + expected_summary = projection.parameters + else: + expected_summary = { + "nMads": n_mads, + "minCellsPerSample": min_cells, + "nSamples": len(projection.captureSizes), + "nSkippedSamples": len( + cast(dict[str, object], projection.parameters["skipReasons"]) + ), + } + if canonical_bytes(parameters) != canonical_bytes(expected_summary): + raise ValueError("Automatic cell-QC parameters do not match the exact inputs") + if canonical_bytes(dict(expected_resolved_bounds)) != canonical_bytes( + projection.parameters["resolvedBounds"] + ): + raise ValueError("Automatic cell-QC resolved bounds differ from its evidence") + if projection.retainedCells != expected_retained_cells: + raise ValueError( + "Automatic cell-QC retained-cell count differs from its evidence" + ) + if projection.flagCounts != flag_counts: + raise ValueError("Automatic cell-QC flag counts differ from its evidence") + + result = store.auto_filter_cells( + attrs_list, + min_p=min_p, + max_p=max_p, + cell_selection=prior, + artifact_metrics=execution_artifacts, + invalidate_cache=invalidate_cache, + sample_column=core_sample_column, + sample_artifact=core_sample_artifact, + n_mads=n_mads, + min_cells_per_sample=min_cells, + ) + actual = read_stored_selection_mask( + store.zw, + result, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ) + expected = np.zeros(store.cells.N, dtype=bool) + expected[active_idx] = projection.keep + if not np.array_equal(actual, expected): + raise RuntimeError( + "DataStore.auto_filter_cells output differs from its projected evidence" + ) + + flag_names = tuple(sorted(projection.flags)) + flag_ref: ArtifactRef | None = None + if flag_names: + values = np.column_stack([projection.flags[name] for name in flag_names]) + source_inputs: dict[str, Any] = { + "prior_cell_selection": prior, + "metadata_fingerprints": metadata_fingerprints, + "artifact_metrics": { + execution_source.name: source.artifact + for source, execution_source in zip( + metric_artifacts, + execution_artifacts, + strict=True, + ) + }, + "grouping_source": grouping_source, + "auto_filter_selection": result, + } + planned = plan_cell_data_artifact( + store.zw, + scope="datastore", + kind="metadata_snapshot", + operation="run_auto_cell_qc_flags", + parameters={ + "action": action, + "profileParameters": parameters, + "resolvedBounds": projection.parameters["resolvedBounds"], + "flagNames": list(flag_names), + }, + inputs=source_inputs, + execution_options={}, + cell_selection=prior, + arrays={"values": (values.shape, "b")}, + invalidate_cache=invalidate_cache, + ) + write_cell_data_artifact( + store.zw, + planned, + {"values": values.astype(bool, copy=False)}, + fingerprint_payload=True, + ) + flag_ref = planned.ref + return result, flag_ref + + +__all__ = ["execute_auto_cell_qc", "execute_registered_cell_qc"] diff --git a/scarf/agent/cell_quality/profiles.py b/scarf/agent/cell_quality/profiles.py new file mode 100644 index 00000000..1a4cddd3 --- /dev/null +++ b/scarf/agent/cell_quality/profiles.py @@ -0,0 +1,909 @@ +"""Registered and core-parity cell-quality profile projections.""" + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Literal, cast + +import numpy as np + +from ...quality_control.filtering import ( + _apply_bounds, + _clamp_metric_bound, + _from_work_scale, + _mad_bounds, + _sample_aware_mad_mask, + _validated_sample_labels, + _validated_work_scale, + gaussian_quantile_bounds, +) + +type RegisteredCellQcProfile = Literal[ + "retainWithFlags", + "globalMad5", + "captureMad5", + "captureMad3Sensitivity", + "pooledReferenceMad5", +] +type CellQualityProfile = ( + RegisteredCellQcProfile | Literal["coreGlobalGaussian", "coreSampleMad3"] +) + + +def cell_qc_policy( + action: str, registered_profile: RegisteredCellQcProfile | None +) -> CellQualityProfile | None: + """Name the exact filtering route shared by selection, execution, and resume.""" + if registered_profile is not None: + return registered_profile + if action == "globalGaussian": + return "coreGlobalGaussian" + if action == "sampleMad": + return "coreSampleMad3" + return None + + +type AutoFilterAction = Literal["globalGaussian", "sampleMad"] +type QcMetricRole = Literal[ + "count", + "feature", + "mitochondrial", + "ribosomal", + "diagnostic", +] +type QcRemovalDirection = Literal["lower", "upper", "none"] + +REGISTERED_CELL_QC_PROFILES: tuple[RegisteredCellQcProfile, ...] = ( + "retainWithFlags", + "globalMad5", + "captureMad5", + "captureMad3Sensitivity", + "pooledReferenceMad5", +) + + +def qc_metric_execution_name( + name: str, + *, + artifact_id: str | None = None, + collides_with_metadata: bool = False, +) -> str: + """Return the deterministic metric name passed to core filtering helpers.""" + if not isinstance(name, str) or not name.strip() or name != name.strip(): + raise ValueError("QC metric names must be non-empty trimmed strings") + if not collides_with_metadata: + return name + if not isinstance(artifact_id, str) or not artifact_id: + raise ValueError("A colliding artifact metric requires an artifact id") + return f"artifact_{artifact_id[:16]}_{name}" + + +@dataclass(frozen=True, slots=True) +class RegisteredQcThreshold: + """One data-derived threshold for one metric and reference population.""" + + metric: str + role: QcMetricRole + group: str + nMads: float + transform: Literal["identity", "log1p"] + removalDirection: QcRemovalDirection + median: float + scaledMad: float + lowerRemoval: float | None + upperRemoval: float | None + upperFlag: float | None + skipReason: Literal["zeroMad"] | None = None + + def to_dict(self) -> dict[str, str | float | None]: + """Return JSON-safe threshold evidence.""" + return { + "metric": self.metric, + "role": self.role, + "group": self.group, + "nMads": self.nMads, + "transform": self.transform, + "removalDirection": self.removalDirection, + "median": self.median, + "scaledMad": self.scaledMad, + "lowerRemoval": self.lowerRemoval, + "upperRemoval": self.upperRemoval, + "upperFlag": self.upperFlag, + "skipReason": self.skipReason, + } + + +@dataclass(frozen=True, slots=True) +class CaptureQcComparison: + """Comparison of one capture median with the complete active population.""" + + capture: str + cells: int + adverseGlobalOutlier: bool + adverseAxes: tuple[QcMetricRole, ...] = () + independentAdverseAxes: int = 0 + wholeCaptureFailure: bool = False + retainedCells: int | None = None + retainedFraction: float | None = None + reasons: tuple[str, ...] = () + metricComparisons: dict[str, dict[str, float | str | None]] = field( + default_factory=dict + ) + + def to_dict(self) -> dict[str, object]: + """Return JSON-safe failed-capture evidence.""" + return { + "capture": self.capture, + "cells": self.cells, + "adverseGlobalOutlier": self.adverseGlobalOutlier, + "adverseAxes": list(self.adverseAxes), + "independentAdverseAxes": self.independentAdverseAxes, + "wholeCaptureFailure": self.wholeCaptureFailure, + "retainedCells": self.retainedCells, + "retainedFraction": self.retainedFraction, + "reasons": list(self.reasons), + "metricComparisons": self.metricComparisons, + } + + +@dataclass(frozen=True, slots=True) +class RegisteredQcProjection: + """Exact masks and evidence for one registered cell-QC profile.""" + + profile: RegisteredCellQcProfile + keep: np.ndarray + flags: dict[str, np.ndarray] + thresholds: tuple[RegisteredQcThreshold, ...] + captureSizes: dict[str, int] = field(default_factory=dict) + retainedByCapture: dict[str, int] = field(default_factory=dict) + captureComparisons: tuple[CaptureQcComparison, ...] = () + failedCaptureCandidates: tuple[str, ...] = () + warnings: tuple[str, ...] = () + + @property + def retainedCells(self) -> int: + """Number of active cells retained by the projection.""" + return int(self.keep.sum()) + + @property + def flagCounts(self) -> dict[str, int]: + """Counts for each non-removal diagnostic flag.""" + return {name: int(mask.sum()) for name, mask in self.flags.items()} + + @property + def metricFlagCounts(self) -> dict[str, dict[str, int]]: + """Counts grouped by exact metric and flag name.""" + grouped: dict[str, dict[str, int]] = {} + for name, mask in self.flags.items(): + metric, flag = name.rsplit(":", 1) + grouped.setdefault(metric, {})[flag] = int(mask.sum()) + return grouped + + +@dataclass(frozen=True, slots=True) +class AutoFilterProjection: + """Exact in-memory projection of one core ``auto_filter_cells`` path.""" + + action: AutoFilterAction + keep: np.ndarray + flags: dict[str, np.ndarray] + parameters: dict[str, object] + captureSizes: dict[str, int] = field(default_factory=dict) + retainedByCapture: dict[str, int] = field(default_factory=dict) + captureComparisons: tuple[CaptureQcComparison, ...] = () + failedCaptureCandidates: tuple[str, ...] = () + warnings: tuple[str, ...] = () + + @property + def retainedCells(self) -> int: + """Number of active cells retained by the projection.""" + return int(self.keep.sum()) + + @property + def flagCounts(self) -> dict[str, int]: + """Counts for each exact metric-bound flag.""" + return {name: int(mask.sum()) for name, mask in self.flags.items()} + + @property + def metricFlagCounts(self) -> dict[str, dict[str, int]]: + """Counts grouped by exact metric and bound side.""" + grouped: dict[str, dict[str, int]] = {} + for name, mask in self.flags.items(): + metric, flag = name.rsplit(":", 1) + grouped.setdefault(metric, {})[flag] = int(mask.sum()) + return grouped + + +def registered_qc_metric_role(metric: str) -> QcMetricRole: + """Classify one conventional cell-quality metric without gene inspection.""" + normalized = metric.replace("_", "").lower() + if normalized.endswith("ncounts") or normalized.endswith("totalcounts"): + return "count" + if normalized.endswith("nfeatures") or normalized.endswith("ngenesbycounts"): + return "feature" + if ( + normalized.endswith("percentmito") + or normalized.endswith("pctcountsmt") + or normalized.endswith("mitochondrialpercent") + ): + return "mitochondrial" + if ( + normalized.endswith("percentribo") + or normalized.endswith("pctcountsribo") + or normalized.endswith("ribosomalpercent") + ): + return "ribosomal" + return "diagnostic" + + +def _metric_policy( + role: QcMetricRole, +) -> tuple[Literal["identity", "log1p"], QcRemovalDirection]: + if role in {"count", "feature"}: + return "log1p", "lower" + if role in {"mitochondrial", "ribosomal"}: + return "identity", "upper" + return "identity", "none" + + +def _validated_inputs( + values_by_metric: dict[str, np.ndarray], + active: np.ndarray, +) -> tuple[dict[str, np.ndarray], np.ndarray]: + active_mask = np.asarray(active, dtype=bool) + if active_mask.ndim != 1: + raise ValueError("active must be a one-dimensional boolean vector") + if not active_mask.any(): + raise ValueError("active must select at least one cell") + values: dict[str, np.ndarray] = {} + for metric, raw in values_by_metric.items(): + if ( + not isinstance(metric, str) + or not metric.strip() + or metric != metric.strip() + ): + raise ValueError("QC metric names must be non-empty trimmed strings") + array = np.asarray(raw, dtype=float) + if array.ndim != 1 or array.shape != active_mask.shape: + raise ValueError(f"QC metric {metric!r} must align with active") + if not np.isfinite(array[active_mask]).all(): + raise ValueError(f"QC metric {metric!r} contains non-finite active values") + values[metric] = array + return values, active_mask + + +def _threshold( + metric: str, + values: np.ndarray, + reference: np.ndarray, + *, + group: str, + n_mads: float, +) -> RegisteredQcThreshold: + role = registered_qc_metric_role(metric) + transform, removal_direction = _metric_policy(role) + work = _validated_work_scale(values[reference], attr=metric, transform=transform) + median_work = float(np.median(work)) + low_work, high_work, scaled_mad = _mad_bounds(work, n_mads) + if scaled_mad == 0.0: + median = _clamp_metric_bound( + _from_work_scale(median_work, transform), + transform=transform, + is_percent=role in {"mitochondrial", "ribosomal"}, + ) + return RegisteredQcThreshold( + metric=metric, + role=role, + group=group, + nMads=n_mads, + transform=transform, + removalDirection=removal_direction, + median=median, + scaledMad=0.0, + lowerRemoval=None, + upperRemoval=None, + upperFlag=None, + skipReason="zeroMad", + ) + + median = _clamp_metric_bound( + _from_work_scale(median_work, transform), + transform=transform, + is_percent=role in {"mitochondrial", "ribosomal"}, + ) + low = _clamp_metric_bound( + _from_work_scale(low_work, transform), + transform=transform, + is_percent=role in {"mitochondrial", "ribosomal"}, + ) + high = _clamp_metric_bound( + _from_work_scale(high_work, transform), + transform=transform, + is_percent=role in {"mitochondrial", "ribosomal"}, + ) + return RegisteredQcThreshold( + metric=metric, + role=role, + group=group, + nMads=n_mads, + transform=transform, + removalDirection=removal_direction, + median=median, + scaledMad=scaled_mad, + lowerRemoval=low if removal_direction == "lower" else None, + upperRemoval=high if removal_direction == "upper" else None, + upperFlag=high if role in {"count", "feature"} else None, + ) + + +def _apply_threshold( + values: np.ndarray, + target: np.ndarray, + threshold: RegisteredQcThreshold, + keep: np.ndarray, + flags: dict[str, np.ndarray], + *, + apply_removal: bool, +) -> None: + if threshold.lowerRemoval is not None: + removed = target & (values < threshold.lowerRemoval) + low_flag = flags.setdefault( + f"{threshold.metric}:lowQuality", + np.zeros(target.shape[0], dtype=bool), + ) + low_flag[removed] = True + if apply_removal: + keep[removed] = False + if threshold.upperRemoval is not None: + removed = target & (values > threshold.upperRemoval) + mito_flag = flags.setdefault( + f"{threshold.metric}:highMito", + np.zeros(target.shape[0], dtype=bool), + ) + mito_flag[removed] = True + if apply_removal: + keep[removed] = False + if threshold.upperFlag is not None: + flag_name = f"{threshold.metric}:high" + flag = flags.setdefault(flag_name, np.zeros(target.shape[0], dtype=bool)) + flag[target & (values > threshold.upperFlag)] = True + + +def _ordered_capture_masks( + capture_labels: np.ndarray, + active: np.ndarray, +) -> list[tuple[str, np.ndarray]]: + labels = _validated_sample_labels( + capture_labels, + active, + label_name="physical capture labels", + ) + captures: list[tuple[str, np.ndarray]] = [] + seen_values: list[object] = [] + seen_keys: set[str] = set() + for raw in labels[active]: + value = raw.item() if isinstance(raw, np.generic) else raw + if any(value == seen for seen in seen_values): + continue + key = value.decode("utf-8") if isinstance(value, bytes) else str(value) + if key in seen_keys: + raise ValueError( + "Physical capture labels collide after provenance encoding" + ) + seen_values.append(value) + seen_keys.add(key) + captures.append((key, active & (labels == raw))) + return captures + + +def _global_capture_comparisons( + values_by_metric: dict[str, np.ndarray], + active: np.ndarray, + captures: list[tuple[str, np.ndarray]], +) -> tuple[CaptureQcComparison, ...]: + global_thresholds = { + metric: _threshold( + metric, + values, + active, + group="globalComparison", + n_mads=5.0, + ) + for metric, values in values_by_metric.items() + if registered_qc_metric_role(metric) != "diagnostic" + } + comparisons: list[CaptureQcComparison] = [] + for capture, mask in captures: + metric_comparisons: dict[str, dict[str, float | str | None]] = {} + reasons: list[str] = [] + adverse_axes: set[QcMetricRole] = set() + for metric, threshold in global_thresholds.items(): + role = threshold.role + transform, _ = _metric_policy(role) + capture_work = _validated_work_scale( + values_by_metric[metric][mask], + attr=metric, + transform=transform, + ) + capture_median_work = float(np.median(capture_work)) + capture_median = _clamp_metric_bound( + _from_work_scale(capture_median_work, transform), + transform=transform, + is_percent=role in {"mitochondrial", "ribosomal"}, + ) + global_mad = threshold.scaledMad + standardized_shift = ( + None + if global_mad == 0.0 + else (capture_median_work - np.log1p(threshold.median)) / global_mad + if transform == "log1p" + else (capture_median - threshold.median) / global_mad + ) + adverse = ( + threshold.lowerRemoval is not None + and capture_median < threshold.lowerRemoval + ) or ( + threshold.upperRemoval is not None + and capture_median > threshold.upperRemoval + ) + if adverse: + reasons.append(f"{metric}:{role}:adverseGlobalMedian") + adverse_axes.add(role) + metric_comparisons[metric] = { + "role": role, + "captureMedian": capture_median, + "globalMedian": threshold.median, + "globalLower": threshold.lowerRemoval, + "globalUpper": threshold.upperRemoval, + "standardizedMedianShift": standardized_shift, + } + comparisons.append( + CaptureQcComparison( + capture=capture, + cells=int(mask.sum()), + adverseGlobalOutlier=bool(reasons), + adverseAxes=tuple(sorted(adverse_axes)), + independentAdverseAxes=len(adverse_axes), + wholeCaptureFailure=len(adverse_axes) >= 2, + reasons=tuple(reasons), + metricComparisons=metric_comparisons, + ) + ) + return tuple(comparisons) + + +def _with_capture_retention( + comparisons: tuple[CaptureQcComparison, ...], + captures: list[tuple[str, np.ndarray]], + keep: np.ndarray, +) -> tuple[CaptureQcComparison, ...]: + masks = dict(captures) + output: list[CaptureQcComparison] = [] + for comparison in comparisons: + mask = masks[comparison.capture] + retained = int((mask & keep).sum()) + fraction = retained / comparison.cells if comparison.cells else 0.0 + output.append( + CaptureQcComparison( + capture=comparison.capture, + cells=comparison.cells, + adverseGlobalOutlier=comparison.adverseGlobalOutlier, + adverseAxes=comparison.adverseAxes, + independentAdverseAxes=comparison.independentAdverseAxes, + wholeCaptureFailure=comparison.wholeCaptureFailure, + retainedCells=retained, + retainedFraction=fraction, + reasons=comparison.reasons, + metricComparisons=comparison.metricComparisons, + ) + ) + return tuple(output) + + +def project_registered_qc_profile( + profile: RegisteredCellQcProfile, + *, + values_by_metric: dict[str, np.ndarray], + active: np.ndarray, + capture_labels: np.ndarray | None = None, + grouping_proven: bool = False, + min_cells_per_capture: int = 20, + pooled_reference_captures: tuple[str, ...] | None = None, +) -> RegisteredQcProjection: + """Project one registered QC profile without reading or writing a matrix.""" + if profile not in REGISTERED_CELL_QC_PROFILES: + raise ValueError(f"Unknown registered cell-QC profile {profile!r}") + if min_cells_per_capture < 2: + raise ValueError("min_cells_per_capture must be at least 2") + values, active_mask = _validated_inputs(values_by_metric, active) + filtering_values = { + metric: metric_values + for metric, metric_values in values.items() + if registered_qc_metric_role(metric) in {"count", "feature", "mitochondrial"} + } + if not filtering_values: + if profile != "retainWithFlags": + raise ValueError("MAD profiles require count, feature, or mito metrics") + return RegisteredQcProjection( + profile=profile, + keep=active_mask.copy(), + flags={}, + thresholds=(), + warnings=("No registered count, feature, or mito metrics were available",), + ) + + capture_profile = profile in { + "captureMad5", + "captureMad3Sensitivity", + "pooledReferenceMad5", + } + if capture_profile and (capture_labels is None or not grouping_proven): + raise ValueError( + f"{profile} requires an explicitly proven physical capture grouping" + ) + captures = ( + _ordered_capture_masks(np.asarray(capture_labels), active_mask) + if capture_labels is not None and grouping_proven + else [] + ) + capture_sizes = {name: int(mask.sum()) for name, mask in captures} + if profile in {"captureMad5", "captureMad3Sensitivity"}: + undersized = [ + name for name, size in capture_sizes.items() if size < min_cells_per_capture + ] + if undersized: + raise ValueError( + f"{profile} requires at least {min_cells_per_capture} cells in every " + f"capture; undersized={undersized}" + ) + if profile == "pooledReferenceMad5": + references = tuple(pooled_reference_captures or ()) + if len(references) < 2 or len(references) != len(set(references)): + raise ValueError( + "pooledReferenceMad5 requires at least two unique reference captures" + ) + unknown = sorted(set(references).difference(capture_sizes)) + if unknown: + raise ValueError(f"Unknown pooled reference captures: {unknown}") + pooled_cells = sum(capture_sizes[name] for name in references) + if pooled_cells < min_cells_per_capture: + raise ValueError( + "Pooled reference captures do not contain enough active cells" + ) + + n_mads = 3.0 if profile == "captureMad3Sensitivity" else 5.0 + keep = active_mask.copy() + flags: dict[str, np.ndarray] = {} + thresholds: list[RegisteredQcThreshold] = [] + warnings: list[str] = [] + + reference_groups: list[tuple[str, np.ndarray, np.ndarray]] + if profile in {"retainWithFlags", "globalMad5"}: + reference_groups = [("global", active_mask, active_mask)] + elif profile == "pooledReferenceMad5": + reference_names = set(pooled_reference_captures or ()) + reference = np.zeros(active_mask.shape[0], dtype=bool) + for name, mask in captures: + if name in reference_names: + reference |= mask + reference_groups = [("pooledReference", active_mask, reference)] + else: + reference_groups = [(name, mask, mask) for name, mask in captures] + + for group, target, reference in reference_groups: + for metric, metric_values in filtering_values.items(): + threshold = _threshold( + metric, + metric_values, + reference, + group=group, + n_mads=n_mads, + ) + thresholds.append(threshold) + if threshold.skipReason is not None: + warnings.append( + f"Ignored {metric!r} for {group!r} because its MAD was zero" + ) + continue + _apply_threshold( + metric_values, + target, + threshold, + keep, + flags, + apply_removal=profile != "retainWithFlags", + ) + + comparison_values = { + metric: metric_values + for metric, metric_values in values.items() + if registered_qc_metric_role(metric) != "diagnostic" + } + comparisons = ( + _with_capture_retention( + _global_capture_comparisons( + comparison_values, + active_mask, + captures, + ), + captures, + keep, + ) + if captures + else () + ) + failed = tuple( + comparison.capture + for comparison in comparisons + if comparison.wholeCaptureFailure + ) + if failed: + warnings.append( + "Captures failed at least two independent global QC axes: " + + ", ".join(failed) + ) + retained_by_capture = {name: int((mask & keep).sum()) for name, mask in captures} + return RegisteredQcProjection( + profile=profile, + keep=keep, + flags=flags, + thresholds=tuple(thresholds), + captureSizes=capture_sizes, + retainedByCapture=retained_by_capture, + captureComparisons=comparisons, + failedCaptureCandidates=failed, + warnings=tuple(warnings), + ) + + +def _auto_bound_flags( + *, + metric: str, + values: np.ndarray, + target: np.ndarray, + low: float | None, + high: float | None, + flags: dict[str, np.ndarray], +) -> None: + if low is not None: + flags.setdefault( + f"{metric}:low", + np.zeros(target.shape[0], dtype=bool), + )[target & (values <= low)] = True + if high is not None: + flags.setdefault( + f"{metric}:high", + np.zeros(target.shape[0], dtype=bool), + )[target & (values >= high)] = True + + +def project_auto_filter_profile( + action: AutoFilterAction, + *, + values_by_metric: dict[str, np.ndarray], + active: np.ndarray, + sample_labels: np.ndarray | None = None, + grouping_proven: bool = False, + min_p: float = 0.01, + max_p: float = 0.99, + n_mads: float = 3.0, + min_cells_per_sample: int = 20, +) -> AutoFilterProjection: + """Project one existing core auto-filter path without writing data. + + The projection calls the same filtering helpers as + :meth:`DataStore.auto_filter_cells`. It rejects any input for which the + core operation would not produce finite bounds. + """ + if action not in {"globalGaussian", "sampleMad"}: + raise ValueError(f"Unknown automatic cell-QC action {action!r}") + values, active_mask = _validated_inputs(values_by_metric, active) + if not values: + raise ValueError("Automatic cell-QC profiles require at least one metric") + + flags: dict[str, np.ndarray] = {} + warnings: list[str] = [] + captures: list[tuple[str, np.ndarray]] = [] + parameters: dict[str, object] + if action == "globalGaussian": + keep = active_mask.copy() + resolved_bounds: dict[str, dict[str, float]] = {} + for metric, metric_values in values.items(): + low, high = gaussian_quantile_bounds( + metric_values[active_mask], + min_p, + max_p, + ) + if not np.isfinite([low, high]).all(): + raise ValueError( + f"QC metric {metric!r} produced non-finite Gaussian bounds" + ) + resolved_bounds[metric] = {"low": low, "high": high} + keep &= _apply_bounds(metric_values, low, high) + _auto_bound_flags( + metric=metric, + values=metric_values, + target=active_mask, + low=low, + high=high, + flags=flags, + ) + parameters = { + "minP": float(min_p), + "maxP": float(max_p), + "resolvedBounds": resolved_bounds, + } + if sample_labels is not None and grouping_proven: + captures = _ordered_capture_masks( + np.asarray(sample_labels), + active_mask, + ) + else: + if sample_labels is None or not grouping_proven: + raise ValueError( + "sampleMad requires an explicitly proven physical capture grouping" + ) + if min_p != 0.01 or max_p != 0.99: + raise ValueError( + "sampleMad requires the core Gaussian probabilities to remain " + "at 0.01 and 0.99" + ) + keep_from_core, provenance = _sample_aware_mad_mask( + values_by_attr=values, + sample_labels=np.asarray(sample_labels), + active=active_mask, + n_mads=n_mads, + min_cells_per_sample=min_cells_per_sample, + attrs=list(values), + ) + keep = active_mask & keep_from_core + captures = _ordered_capture_masks( + np.asarray(sample_labels), + active_mask, + ) + capture_masks = dict(captures) + raw_bounds = provenance["resolved_bounds"] + for capture, bounds_by_metric in raw_bounds.items(): + target = capture_masks[capture] + for metric, raw_bound in bounds_by_metric.items(): + bound = cast(Mapping[str, object], raw_bound) + low_value = bound.get("low") + high_value = bound.get("high") + sample_low = ( + float(cast(float, low_value)) if low_value is not None else None + ) + sample_high = ( + float(cast(float, high_value)) if high_value is not None else None + ) + _auto_bound_flags( + metric=metric, + values=values[metric], + target=target, + low=sample_low, + high=sample_high, + flags=flags, + ) + warnings.extend(provenance["warnings"]) + parameters = { + "minP": 0.01, + "maxP": 0.99, + "nMads": float(n_mads), + "minCellsPerSample": int(min_cells_per_sample), + "madScale": float(provenance["mad_scale"]), + "metricPolicies": provenance["metric_policies"], + "sampleSizes": provenance["sample_sizes"], + "skipReasons": provenance["skip_reasons"], + "resolvedBounds": provenance["resolved_bounds"], + } + + capture_sizes = {name: int(mask.sum()) for name, mask in captures} + retained_by_capture = {name: int((mask & keep).sum()) for name, mask in captures} + comparisons = ( + _with_capture_retention( + _global_capture_comparisons(values, active_mask, captures), + captures, + keep, + ) + if captures + else () + ) + failed = tuple( + comparison.capture + for comparison in comparisons + if comparison.wholeCaptureFailure + ) + if failed: + warnings.append( + "Captures failed at least two independent global QC axes: " + + ", ".join(failed) + ) + return AutoFilterProjection( + action=action, + keep=keep, + flags=flags, + parameters=parameters, + captureSizes=capture_sizes, + retainedByCapture=retained_by_capture, + captureComparisons=comparisons, + failedCaptureCandidates=failed, + warnings=tuple(warnings), + ) + + +def offered_registered_qc_profiles( + *, + values_by_metric: dict[str, np.ndarray], + active: np.ndarray, + capture_labels: np.ndarray | None = None, + grouping_proven: bool = False, + min_cells_per_capture: int = 20, + pooled_reference_captures: tuple[str, ...] | None = None, +) -> list[RegisteredQcProjection]: + """Build only profiles whose deterministic eligibility gates pass.""" + projections = [ + project_registered_qc_profile( + "retainWithFlags", + values_by_metric=values_by_metric, + active=active, + capture_labels=capture_labels, + grouping_proven=grouping_proven, + ) + ] + try: + global_projection = project_registered_qc_profile( + "globalMad5", + values_by_metric=values_by_metric, + active=active, + capture_labels=capture_labels, + grouping_proven=grouping_proven, + ) + except ValueError: + return projections + projections.append(global_projection) + if capture_labels is None or not grouping_proven: + return projections + for profile in ("captureMad5", "captureMad3Sensitivity"): + try: + projection = project_registered_qc_profile( + profile, + values_by_metric=values_by_metric, + active=active, + capture_labels=capture_labels, + grouping_proven=True, + min_cells_per_capture=min_cells_per_capture, + ) + except ValueError: + continue + projections.append(projection) + if pooled_reference_captures: + try: + projection = project_registered_qc_profile( + "pooledReferenceMad5", + values_by_metric=values_by_metric, + active=active, + capture_labels=capture_labels, + grouping_proven=True, + min_cells_per_capture=min_cells_per_capture, + pooled_reference_captures=pooled_reference_captures, + ) + except ValueError: + pass + else: + projections.append(projection) + return projections + + +__all__ = [ + "REGISTERED_CELL_QC_PROFILES", + "AutoFilterAction", + "AutoFilterProjection", + "CaptureQcComparison", + "QcMetricRole", + "RegisteredCellQcProfile", + "RegisteredQcProjection", + "RegisteredQcThreshold", + "offered_registered_qc_profiles", + "project_auto_filter_profile", + "project_registered_qc_profile", + "qc_metric_execution_name", + "registered_qc_metric_role", +] diff --git a/scarf/agent/config/__init__.py b/scarf/agent/config/__init__.py index 1b7b9f9e..a491577c 100644 --- a/scarf/agent/config/__init__.py +++ b/scarf/agent/config/__init__.py @@ -1,6 +1,5 @@ """Configuration shared by the four Scarf domain agents.""" -import re from typing import Any, Literal from urllib.parse import urlparse @@ -12,67 +11,14 @@ "AgentRunConfig", "get_model_settings", "get_usage_limits", - "CONFIG", ] -class Config: - """Shared configuration for the four Scarf domain agents.""" - - # BiologicalInterpretation - _MAX_CLUSTERS: int = 20 - _MAX_CONDITIONS: int = 30 - _MAX_MARKERS: int = 25 - # CharacterizeFeatures - _MAX_EXOGENOUS: int = 25 - _CONTEXT_LIMIT: int = 1200 - _AUTO_DOWNLOAD_SPECIES: frozenset[str] = frozenset({"homo_sapiens", "mus_musculus"}) - # DataEnrichment - _MAX_FEATURE_QUERIES: int = 50 - # ParameterTuning - _MAX_CANDIDATES_OFFERED: int = 25 - _CANDIDATE_ID: re.Pattern[str] = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_]{0,63}$") - _RANDOM_SEED: int = 4444 - _PCA_RANDOM_SEED: int = 4466 - # CharacterizeCovariates - _CATEGORICAL_MAX_LEVELS: int = 50 - _EMBEDDING_TOKENS: tuple[str, ...] = ( - "umap", - "pca", - "tsne", - "scvi", - "latent", - "phate", - "forceatlas", - "diffmap", - "diffusionmap", - "diffusion", - ) - _DOMAINS: frozenset[str] = frozenset( - {"biological", "technical", "design", "ignore", "unknown"} - ) # Only these domains reach the design table, so only they are worth collapsing. - _ANALYSED: frozenset[str] = frozenset({"biological", "technical", "design"}) - _KINDS: frozenset[str] = frozenset({"categorical", "continuous"}) - _RESERVED_COLUMNS = frozenset({"I", "ids", "names"}) - - _SHORT_EMBEDDING_PARTS: frozenset[str] = frozenset({"fa", "dm", "pc"}) - _INDEXED_NAME: re.Pattern[str] = re.compile(r"(?P.+?)[-_]?(?P\d+)") - _ONTOLOGY_SUFFIX: str = "_ontology_term_id" - _SAMPLE_LEVELS: int = 8 - _ASSOCIATION_FLOOR: float = 0.1 - _DROP_REASONS: dict[str, str] = { - "dropAssayStat": "Scarf assay statistic column", - "dropProvenance": "analysis-linked column", - "dropEmbedding": "embedding-style column", - "dropConstant": "single-level column", - } - - class AgentRunConfig(AgentDataModel): """Bound one agent run without selecting a scientific workflow.""" - requestLimit: int = 9 - toolCallLimit: int = 5 + requestLimit: int = 10 + toolCallLimit: int = 10 inputTokenLimit: int | None = None outputTokenLimit: int | None = 32768 # Per provider response. totalTokenLimit: int | None = None @@ -143,10 +89,6 @@ def with_limits( ) return type(self).model_validate(values) - @classmethod - def get_example(cls) -> "AgentRunConfig": - return cls(requestLimit=9, toolCallLimit=5, outputTokenLimit=2048) - def get_model_settings( config: AgentRunConfig | None = None, @@ -202,6 +144,3 @@ def get_usage_limits(config: AgentRunConfig | None = None) -> Any: output_tokens_limit=output_tokens_limit, total_tokens_limit=run_config.totalTokenLimit, ) - - -CONFIG = Config() diff --git a/scarf/agent/config/agent_exec.py b/scarf/agent/config/agent_exec.py index 0467dfde..93f41ed6 100644 --- a/scarf/agent/config/agent_exec.py +++ b/scarf/agent/config/agent_exec.py @@ -1,24 +1,280 @@ """Common bounded execution for the four Scarf domain agents.""" import asyncio +import json +import math import sys import time +import uuid from collections.abc import Callable, Sequence from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field from inspect import isawaitable, iscoroutinefunction -from typing import Any +from threading import Event, Lock +from typing import TYPE_CHECKING, Any, Literal from ...utils.logging import logger +from .._deps import require_pydantic_ai from ..types import ( AgentExecutionResult, + AgentProviderFailure, AgentRunInfo, AgentUsageInfo, + AgentValidationRetry, ToolCallInfo, ) from . import AgentRunConfig, get_model_settings, get_usage_limits -from ._deps import require_pydantic_ai -__all__ = ["run_agent", "run_agent_sync"] +if TYPE_CHECKING: + from pydantic_ai.messages import UserContent +else: + UserContent = Any + +type AgentUserPrompt = str | Sequence[UserContent] +type ImageMediaType = Literal["image/png", "image/jpeg", "image/webp"] + +_MAX_VISUAL_EVIDENCE_ITEMS = 8 +_MAX_VISUAL_EVIDENCE_ITEM_BYTES = 4 * 1024 * 1024 +_MAX_VISUAL_EVIDENCE_TOTAL_BYTES = 16 * 1024 * 1024 +_RATE_LIMIT_DELAYS = (15.0, 30.0, 60.0) +_RATE_LIMIT_WAIT_LIMIT = 120.0 + +__all__ = [ + "AgentUserPrompt", + "build_visual_evidence_prompt", + "ImageEvidence", + "ImageInputUnsupportedError", + "ImageMediaType", + "run_agent", + "run_agent_async", + "run_agent_sync", +] + + +@dataclass(frozen=True, slots=True) +class ImageEvidence: + """One in-memory image supplied to a bounded agent comparison.""" + + identifier: str + data: bytes + media_type: ImageMediaType = "image/png" + + +class ImageInputUnsupportedError(RuntimeError): + """The configured model or provider rejected image input.""" + + +@dataclass +class _ProviderRequestAudit: + requests: int = 0 + waited: float = 0.0 + failures: list[AgentProviderFailure] = field(default_factory=list) + response_requests: list[int] = field(default_factory=list) + + +def _rate_limit_delay(error: Exception, retry: int) -> float | None: + """Retry temporary throttling, preserving provider timing and hard quotas.""" + from pydantic_ai.exceptions import ModelHTTPError + + if not isinstance(error, ModelHTTPError) or error.status_code != 429: + return None + detail = describe_agent_error(error).casefold() + if any( + marker in detail + for marker in ( + "insufficient_quota", + "insufficient quota", + "exceeded your current quota", + "insufficient credit", + "insufficient balance", + "billing", + "payment required", + "spending limit", + "daily limit", + "monthly limit", + ) + ): + return None + if retry >= len(_RATE_LIMIT_DELAYS): + return None + current: BaseException | None = error + seen: set[int] = set() + while current is not None and id(current) not in seen and len(seen) < 8: + seen.add(id(current)) + headers = getattr(current, "headers", None) or getattr( + getattr(current, "response", None), "headers", None + ) + if headers: + parsed = ModelHTTPError(429, error.model_name, headers=headers).retry_after + if parsed is not None and math.isfinite(parsed): + return parsed + current = current.__cause__ or ( + current.__context__ if not current.__suppress_context__ else None + ) + return _RATE_LIMIT_DELAYS[retry] + + +def _rate_limited_model( + model: Any, + *, + config: AgentRunConfig, + audit: _ProviderRequestAudit, + name: str | None, +) -> Any: + """Retry only a failed request; completed agent tools are never replayed.""" + from pydantic_ai.exceptions import UsageLimitExceeded + from pydantic_ai.models.wrapper import WrapperModel + + class RateLimitedModel(WrapperModel): + async def request( + self, messages: Any, model_settings: Any, model_request_parameters: Any + ) -> Any: + retry = 0 + while True: + if audit.requests >= config.requestLimit: + raise UsageLimitExceeded( + "The model request limit includes failed provider requests; " + f"all {config.requestLimit} available requests were used" + ) + audit.requests += 1 + try: + response = await self.wrapped.request( + messages, model_settings, model_request_parameters + ) + audit.response_requests.append(audit.requests) + return response + except Exception as exc: + delay = _rate_limit_delay(exc, retry) + if delay is not None and ( + audit.requests >= config.requestLimit + or audit.waited + delay > _RATE_LIMIT_WAIT_LIMIT + ): + delay = None + audit.failures.append( + AgentProviderFailure( + requestIndex=audit.requests, + statusCode=getattr(exc, "status_code", None), + error=describe_agent_error(exc), + retryDelaySeconds=delay, + ) + ) + if delay is None: + raise + logger.warning( + f"Agent {name or 'unnamed'}: provider rate limit; retrying " + f"the same request in {delay:g}s " + f"({retry + 1}/{len(_RATE_LIMIT_DELAYS)} retries). " + "Completed tools and scientific evidence are preserved." + ) + audit.waited += delay + await asyncio.sleep(delay) + retry += 1 + + return RateLimitedModel(model) + + +def describe_agent_error(error: BaseException, *, limit: int = 8000) -> str: + """Describe the concrete cause chain without traceback or response payloads.""" + parts: list[str] = [] + seen: set[int] = set() + current: BaseException | None = error + while current is not None and id(current) not in seen and len(parts) < 8: + seen.add(id(current)) + detail = " ".join(str(current).split()) + notes = getattr(current, "__notes__", ()) + if notes: + detail += "; " + "; ".join(str(note) for note in notes) + parts.append(f"{type(current).__name__}: {detail[:2000]}") + current = current.__cause__ or ( + current.__context__ if not current.__suppress_context__ else None + ) + return "; caused by ".join(parts)[:limit] + + +def _image_input_is_unsupported(exc: Exception) -> bool: + from pydantic_ai.exceptions import ModelHTTPError, UserError + + if isinstance(exc, UserError): + detail = str(exc).casefold() + elif isinstance(exc, ModelHTTPError) and exc.status_code in {400, 415, 422}: + detail = str(exc.body).casefold() + else: + return False + return any( + marker in detail + for marker in ( + "binary content is not supported", + "binary input is not supported", + "does not support binary content", + "does not support multimodal", + "doesn't support multimodal", + "image content is not supported", + "multimodal input is not supported", + "image input is not supported", + "image inputs are not supported", + "images are not supported", + "does not support image", + "unsupported image input", + "only text input", + ) + ) + + +def build_visual_evidence_prompt( + prompt: str, + images: Sequence[ImageEvidence], +) -> tuple[UserContent, ...]: + """Build bounded Pydantic AI image content without creating report files.""" + + if not isinstance(prompt, str) or not prompt.strip(): + raise ValueError("Visual evidence prompt must be a non-empty string") + values = tuple(images) + if not values: + raise ValueError("Visual evidence requires at least one image") + if len(values) > _MAX_VISUAL_EVIDENCE_ITEMS: + raise ValueError( + "Visual evidence exceeds the maximum of " + f"{_MAX_VISUAL_EVIDENCE_ITEMS} images" + ) + identifiers = [value.identifier for value in values] + if any(not value or value != value.strip() for value in identifiers): + raise ValueError("Visual evidence identifiers must be non-empty and trimmed") + if len(identifiers) != len(set(identifiers)): + raise ValueError("Visual evidence identifiers must be unique") + total_bytes = 0 + for value in values: + if value.media_type not in {"image/png", "image/jpeg", "image/webp"}: + raise ValueError( + f"Visual evidence image {value.identifier!r} has unsupported media type" + ) + if not isinstance(value.data, bytes) or not value.data: + raise ValueError("Visual evidence images must contain non-empty bytes") + if len(value.data) > _MAX_VISUAL_EVIDENCE_ITEM_BYTES: + raise ValueError( + f"Visual evidence image {value.identifier!r} exceeds " + f"{_MAX_VISUAL_EVIDENCE_ITEM_BYTES} bytes" + ) + total_bytes += len(value.data) + if total_bytes > _MAX_VISUAL_EVIDENCE_TOTAL_BYTES: + raise ValueError( + "Visual evidence exceeds the total byte limit of " + f"{_MAX_VISUAL_EVIDENCE_TOTAL_BYTES}" + ) + + require_pydantic_ai() + from pydantic_ai.messages import BinaryContent + + return ( + prompt, + *( + BinaryContent( + data=value.data, + media_type=value.media_type, + identifier=value.identifier, + ) + for value in values + ), + ) def _tool_definitions( @@ -109,11 +365,15 @@ def _tool_calls( continue if part.tool_name not in allowed_names: continue + try: + arguments = part.args_as_dict() + except (TypeError, ValueError): + arguments = {"unparsedArguments": str(part.args)} calls.append( ToolCallInfo( toolName=part.tool_name, callId=part.tool_call_id or "", - arguments=part.args_as_dict(), + arguments=arguments, ) ) return calls @@ -146,12 +406,21 @@ def _build_agent( name: str | None, output_validator: Callable[[Any], Any] | None, normalize_sync_function_model: bool, + validation_failures: list[AgentValidationRetry] | None = None, + provider_audit: _ProviderRequestAudit | None = None, ) -> Any: require_pydantic_ai() - from pydantic_ai import Agent + from pydantic_ai import Agent, RunContext + normalized_model = ( + _normalize_model(model) if normalize_sync_function_model else model + ) + if provider_audit is not None: + normalized_model = _rate_limited_model( + normalized_model, config=config, audit=provider_audit, name=name + ) agent = Agent( - _normalize_model(model) if normalize_sync_function_model else model, + normalized_model, output_type=output_type, system_prompt=system_prompt, deps_type=deps_type or object, @@ -164,21 +433,52 @@ def _build_agent( if output_validator is not None: @agent.output_validator - async def validate_output(output: Any) -> Any: + async def validate_output(context: RunContext[Any], output: Any) -> Any: from pydantic_ai import ModelRetry + submitted = ( + output.model_dump(mode="json") + if hasattr(output, "model_dump") + else str(output) + ) try: validated = output_validator(output) if isawaitable(validated): return await validated return validated except ModelRetry as exc: + if validation_failures is not None: + validation_failures.append( + AgentValidationRetry( + source="output", + requestIndex=( + provider_audit.requests + if provider_audit is not None + else context.usage.requests + ), + message=str(exc), + response=submitted, + ) + ) logger.warning( f"Agent {name or 'unnamed'} requested a structured-output " f"retry: {str(exc)[:500]}" ) raise except (TypeError, ValueError) as exc: + if validation_failures is not None: + validation_failures.append( + AgentValidationRetry( + source="output", + requestIndex=( + provider_audit.requests + if provider_audit is not None + else context.usage.requests + ), + message=str(exc), + response=submitted, + ) + ) logger.warning( f"Agent {name or 'unnamed'} rejected structured output: " f"{str(exc)[:500]}" @@ -188,35 +488,241 @@ async def validate_output(output: Any) -> Any: return agent -def _execution_result( +def _run_info( *, - result: Any, + messages: Sequence[Any], + usage: Any, model: Any, name: str | None, started: float, tools: Sequence[Callable[..., Any] | Any], -) -> AgentExecutionResult: - messages = result.new_messages() + validation_failures: Sequence[AgentValidationRetry], + error: BaseException | None = None, + provider_audit: _ProviderRequestAudit | None = None, +) -> AgentRunInfo: + from pydantic import ValidationError + from pydantic_ai.messages import ModelResponse, RetryPromptPart, ToolCallPart + calls = _tool_calls(messages, allowed_names=_tool_names(tools)) - execution = AgentExecutionResult( - output=result.output, - runInfo=AgentRunInfo( - agentName=name or "", - modelName=_model_name(model), - runId=str(getattr(result, "run_id", "")), - durationSeconds=time.monotonic() - started, - usage=_usage_info(result.usage, tool_calls=len(calls)), - toolCalls=calls, + reported_usage = _usage_info(usage, tool_calls=len(calls)) + if provider_audit is not None: + reported_usage.requests = provider_audit.requests + responses = [message for message in messages if isinstance(message, ModelResponse)] + measured = [message for message in responses if message.usage.has_values()] + reported_usage.availability = ( + "reported" + if measured + and not (provider_audit is not None and provider_audit.failures) + and len(measured) == len(responses) == reported_usage.requests + and not ( + error is not None + and messages + and not isinstance(messages[-1], ModelResponse) + ) + and all( + getattr(message, "state", "complete") == "complete" for message in responses + ) + else "partial" + if measured + else "unavailable" + ) + retries = list(validation_failures) + semantic_messages = {retry.message for retry in retries} + pending: dict[str, Any] = {} + tool_names = _tool_names(tools) + request_index = 0 + response_index = 0 + for message in messages: + if isinstance(message, ModelResponse): + request_index = ( + provider_audit.response_requests[response_index] + if provider_audit is not None + and response_index < len(provider_audit.response_requests) + else request_index + 1 + ) + response_index += 1 + for part in getattr(message, "parts", ()): + if isinstance(part, ToolCallPart): + pending[part.tool_call_id] = part.args + elif isinstance(part, RetryPromptPart): + detail = ( + part.content + if isinstance(part.content, str) + else json.dumps(part.content, default=str, sort_keys=True) + ) + if detail not in semantic_messages: + retries.append( + AgentValidationRetry( + source="tool" if part.tool_name in tool_names else "schema", + requestIndex=request_index, + message=detail, + response=pending.get(part.tool_call_id), + ) + ) + error_detail = describe_agent_error(error) if error is not None else None + if ( + error is not None + and isinstance(error.__cause__, ValidationError) + and messages + and isinstance(messages[-1], ModelResponse) + ): + # The SDK does not append feedback after the last schema retry is exhausted. + calls_in_response = [ + part for part in messages[-1].parts if isinstance(part, ToolCallPart) + ] + if len(calls_in_response) == 1: + rejected = calls_in_response[0] + retries.append( + AgentValidationRetry( + source="tool" if rejected.tool_name in tool_names else "schema", + requestIndex=request_index, + message=str(error.__cause__), + response=rejected.args, + ) + ) + return AgentRunInfo( + agentName=name or "unnamed", + modelName=_model_name(model), + runId=next( + ( + str(message.run_id) + for message in reversed(messages) + if getattr(message, "run_id", None) + ), + uuid.uuid4().hex, ), + durationSeconds=time.monotonic() - started, + usage=reported_usage, + toolCalls=calls, + status="failed" if error is not None else "done", + validationRetries=sorted(retries, key=lambda retry: retry.requestIndex), + providerFailures=provider_audit.failures if provider_audit is not None else [], + errorType=type(error).__name__ if error is not None else None, + error=error_detail, ) - usage = execution.runInfo.usage - logger.info( - f"Agent {name or 'unnamed'} completed in " - f"{execution.runInfo.durationSeconds:.2f}s: requests={usage.requests}, " - f"tool_calls={usage.toolCalls}, input_tokens={usage.inputTokens}, " - f"output_tokens={usage.outputTokens}" + + +async def _execute_agent( + *, + model: Any, + output_type: Any, + system_prompt: str, + user_prompt: AgentUserPrompt, + tools: Sequence[Callable[..., Any] | Any], + deps_type: type[Any] | None, + deps: Any, + config: AgentRunConfig | None, + name: str | None, + output_validator: Callable[[Any], Any] | None, + message_history: Sequence[Any], + on_attempt: Callable[[AgentRunInfo], None] | None, + normalize_sync_function_model: bool, + cancel_requested: Event | None = None, +) -> AgentExecutionResult: + from pydantic_ai import capture_run_messages + from pydantic_ai.usage import RunUsage + + run_config = config or AgentRunConfig() + agent_name = name or "unnamed" + usage_limits = get_usage_limits(run_config) + usage = RunUsage() + failures: list[AgentValidationRetry] = [] + provider_audit = _ProviderRequestAudit() + started = time.monotonic() + logger.debug( + f"Starting agent {agent_name}: model={_model_name(model)}, " + f"tools={len(tools)}, request_limit={run_config.requestLimit}, " + f"tool_call_limit={run_config.toolCallLimit}, retries={run_config.retries}, " + f"per_response_output_limit={run_config.outputTokenLimit}, " + f"run_output_limit={usage_limits.output_tokens_limit}" ) - return execution + with capture_run_messages() as messages: + try: + if cancel_requested is not None and cancel_requested.is_set(): + raise asyncio.CancelledError( + "Analysis interrupted before model execution" + ) + agent = _build_agent( + model=model, + output_type=output_type, + system_prompt=system_prompt, + tools=tools, + deps_type=deps_type, + config=run_config, + name=name, + output_validator=output_validator, + normalize_sync_function_model=normalize_sync_function_model, + validation_failures=failures, + provider_audit=provider_audit, + ) + async with agent: + try: + result = await agent.run( + user_prompt, + deps=deps, + message_history=message_history, + usage_limits=usage_limits, + usage=usage, + ) + except Exception as exc: + if not isinstance(user_prompt, str) and _image_input_is_unsupported( + exc + ): + raise ImageInputUnsupportedError( + "The configured model does not accept image input" + ) from exc + raise + except (Exception, asyncio.CancelledError) as exc: + info = _run_info( + messages=messages[len(message_history) :], + usage=usage, + model=model, + name=name, + started=started, + tools=tools, + validation_failures=failures, + error=exc, + provider_audit=provider_audit, + ) + setattr(exc, "agent_run_info", info) + logger.error( + f"Agent {agent_name} failed after {info.durationSeconds:.2f}s: " + f"{info.error}" + ) + if on_attempt is not None: + try: + on_attempt(info) + except Exception as callback_error: + exc.add_note( + f"Saving agent execution evidence also failed: {type(callback_error).__name__}: {callback_error}" + ) + raise + info = _run_info( + messages=result.new_messages(), + usage=result.usage, + model=model, + name=name, + started=started, + tools=tools, + validation_failures=failures, + provider_audit=provider_audit, + ) + if on_attempt is not None: + try: + on_attempt(info) + except Exception as exc: + setattr(exc, "agent_run_info", info) + exc.add_note( + "The model completed, but saving its execution evidence failed." + ) + raise + logger.debug( + f"Agent {agent_name} completed in {info.durationSeconds:.2f}s: " + f"requests={info.usage.requests}, tool_calls={info.usage.toolCalls}, " + f"input_tokens={info.usage.inputTokens}, output_tokens={info.usage.outputTokens}, " + f"usage={info.usage.availability}" + ) + return AgentExecutionResult(output=result.output, runInfo=info) def run_agent_sync( @@ -224,7 +730,7 @@ def run_agent_sync( model: Any, output_type: Any, system_prompt: str, - user_prompt: str, + user_prompt: AgentUserPrompt, tools: Sequence[Callable[..., Any] | Any] = (), deps_type: type[Any] | None = None, deps: Any = None, @@ -232,6 +738,7 @@ def run_agent_sync( name: str | None = None, output_validator: Callable[[Any], Any] | None = None, message_history: Sequence[Any] = (), + on_attempt: Callable[[AgentRunInfo], None] | None = None, ) -> AgentExecutionResult: """Run one synchronous agent loop and return its bounded audit record. @@ -239,51 +746,38 @@ def run_agent_sync( asynchronous runner cannot drive that loop synchronously, so this hops to a worker thread in that case. Entering the agent context ensures provider HTTP clients are closed on that worker's event loop before it exits. + + ``on_attempt`` receives invocation evidence once, including failure usage and + rejected responses. It runs on the agent's event-loop thread. A failed run + re-raises its original exception with this evidence in ``agent_run_info``. """ - async def execute() -> AgentExecutionResult: - run_config = config or AgentRunConfig() - agent_name = name or "unnamed" - usage_limits = get_usage_limits(run_config) - logger.info( - f"Starting agent {agent_name}: model={_model_name(model)}, " - f"tools={len(tools)}, request_limit={run_config.requestLimit}, " - f"tool_call_limit={run_config.toolCallLimit}, retries={run_config.retries}, " - f"per_response_output_limit={run_config.outputTokenLimit}, " - f"run_output_limit={usage_limits.output_tokens_limit}" - ) - agent = _build_agent( + cancel_requested = Event() + worker_lock = Lock() + worker: tuple[asyncio.AbstractEventLoop, asyncio.Task[Any]] | None = None + + async def execute(*, notebook_worker: bool = False) -> AgentExecutionResult: + nonlocal worker + if notebook_worker: + task = asyncio.current_task() + assert task is not None + with worker_lock: + worker = (asyncio.get_running_loop(), task) + return await _execute_agent( model=model, output_type=output_type, system_prompt=system_prompt, + user_prompt=user_prompt, tools=tools, deps_type=deps_type, - config=run_config, + deps=deps, + config=config, name=name, output_validator=output_validator, normalize_sync_function_model=True, - ) - started = time.monotonic() - try: - async with agent: - result = await agent.run( - user_prompt, - deps=deps, - message_history=message_history, - usage_limits=usage_limits, - ) - except Exception as exc: - logger.error( - f"Agent {agent_name} failed after {time.monotonic() - started:.2f}s: " - f"{type(exc).__name__}" - ) - raise - return _execution_result( - result=result, - model=model, - name=name, - started=started, - tools=tools, + message_history=message_history, + on_attempt=on_attempt, + cancel_requested=cancel_requested if notebook_worker else None, ) try: @@ -292,15 +786,46 @@ async def execute() -> AgentExecutionResult: return asyncio.run(execute()) with ThreadPoolExecutor(max_workers=1) as pool: - return pool.submit(asyncio.run, execute()).result() + future = pool.submit(asyncio.run, execute(notebook_worker=True)) + try: + return future.result() + except KeyboardInterrupt as interrupted: + cancel_requested.set() + with worker_lock: + running = worker + if running is not None: + loop, task = running + try: + loop.call_soon_threadsafe(task.cancel, "Analysis interrupted") + except RuntimeError as cancellation_error: + # The worker may have finished and closed its loop already. + interrupted.add_note( + "Cancellation raced with worker shutdown: " + + describe_agent_error(cancellation_error) + ) + try: + completed = future.result() + except BaseException as worker_error: + info = getattr(worker_error, "agent_run_info", None) + if info is not None: + setattr(interrupted, "agent_run_info", info) + interrupted.add_note( + "Worker shutdown reported: " + describe_agent_error(worker_error) + ) + else: + setattr(interrupted, "agent_run_info", completed.runInfo) + interrupted.add_note( + "The model invocation completed before cancellation was delivered." + ) + raise -async def run_agent( +async def run_agent_async( *, model: Any, output_type: Any, system_prompt: str, - user_prompt: str, + user_prompt: AgentUserPrompt, tools: Sequence[Callable[..., Any] | Any] = (), deps_type: type[Any] | None = None, deps: Any = None, @@ -308,48 +833,24 @@ async def run_agent( name: str | None = None, output_validator: Callable[[Any], Any] | None = None, message_history: Sequence[Any] = (), + on_attempt: Callable[[AgentRunInfo], None] | None = None, ) -> AgentExecutionResult: """Run one asynchronous agent loop and return its bounded audit record.""" - run_config = config or AgentRunConfig() - agent_name = name or "unnamed" - usage_limits = get_usage_limits(run_config) - logger.info( - f"Starting agent {agent_name}: model={_model_name(model)}, " - f"tools={len(tools)}, request_limit={run_config.requestLimit}, " - f"tool_call_limit={run_config.toolCallLimit}, retries={run_config.retries}, " - f"per_response_output_limit={run_config.outputTokenLimit}, " - f"run_output_limit={usage_limits.output_tokens_limit}" - ) - agent = _build_agent( + return await _execute_agent( model=model, output_type=output_type, system_prompt=system_prompt, + user_prompt=user_prompt, tools=tools, deps_type=deps_type, - config=run_config, + deps=deps, + config=config, name=name, output_validator=output_validator, normalize_sync_function_model=False, + message_history=message_history, + on_attempt=on_attempt, ) - started = time.monotonic() - try: - async with agent: - result = await agent.run( - user_prompt, - deps=deps, - message_history=message_history, - usage_limits=usage_limits, - ) - except Exception as exc: - logger.error( - f"Agent {agent_name} failed after {time.monotonic() - started:.2f}s: " - f"{type(exc).__name__}" - ) - raise - return _execution_result( - result=result, - model=model, - name=name, - started=started, - tools=tools, - ) + + +run_agent = run_agent_async diff --git a/scarf/agent/data_enrichment.py b/scarf/agent/data_enrichment.py deleted file mode 100644 index 2149fc6e..00000000 --- a/scarf/agent/data_enrichment.py +++ /dev/null @@ -1,1563 +0,0 @@ -"""Read-only feature and organism enrichment agent.""" - -import re -from collections.abc import Sequence -from pathlib import Path -from textwrap import dedent -from typing import Any, Literal - -from ..features.gene_reference import species_registry -from ..utils.logging import logger -from .characterize_features import characterize_features -from .config import CONFIG, AgentRunConfig -from .config._deps import AGENT_INSTALL_HINT -from .config.agent_exec import run_agent_sync -from .tools import bounded_list -from .types import AgentDataModel, AgentRunInfo, StageStatus - -try: - from pydantic import ConfigDict, Field, model_validator - from pydantic_ai import ModelRetry, RunContext, Tool - from pydantic_ai.tools import ToolDefinition -except ImportError as exc: - raise ImportError(AGENT_INSTALL_HINT) from exc - -__all__ = [ - "AdtControlEvidence", - "AssayFeatureInspection", - "AssayFeatureInspectionBatch", - "AssayModalityEvidence", - "AtacCoordinateEvidence", - "DataEnrichmentAgent", - "DataEnrichmentContext", - "DataEnrichmentDependencies", - "DataEnrichmentReport", - "DataEnrichmentToolCall", - "ExogenousFeatureEvidence", - "FeatureFamilyEvidence", - "FeatureLookupResult", - "FeatureLookupBatch", - "FeatureMatch", - "FeatureReference", - "FeatureSelectionPolicy", - "HtoTagEvidence", - "StudyContextSummary", - "find_present_features", - "find_present_features_batch", - "inspect_assay_features", - "inspect_assay_features_batch", - "validate_data_enrichment_report", -] - -_SUPPORTED_SPECIES = species_registry() -_SYSTEM_PROMPT = ( - dedent( - """ - You are Scarf's Data Enrichment Agent. Work only through the supplied - read-only tools. Inspect every requested assay before making a decision. - Use gene identifiers and names together with the supplied organism hint, - tissue references, cell-type references, and experimental details when - species evidence is ambiguous. Supported species keys are: {supported_species}. - - Call inspect_assay_features_batch once for all requested assays. Never - invent a feature. If individual features are needed, collect all proposed - names across assays and call find_present_features_batch once before - placing them in a policy. Persisted assay types determine modality routes; - never infer a route from an assay label. Copy exact ADT controls, HTO tags, - and ATAC-coordinate status from inspection evidence. Treat Ensembl release - misses as unresolved, not artificial. Mitochondrial, ribosomal, and histone - families may be exclusion candidates. Sex-linked and cell-cycle families - are protected by default in this initial implementation. - - Structure studyContextSummary using only verbatim spans from the supplied - study paragraph or exact caller references. Do not paraphrase, infer, or - invent an organism, tissue, cell type, experiment, hypothesis, or analysis - intent. Empty optional hint lists do not mean that the paragraph lacks - those references. When a category is explicitly present in the paragraph, - include its exact span in the corresponding summary list. The validator - binds the original paragraph and exact caller references. Return a bounded - report with citations copied from tool or context evidence IDs. - Do not write code, mutate the datastore, or request arbitrary Scarf calls. - """ - ) - .strip() - .format(supported_species=", ".join(sorted(_SUPPORTED_SPECIES))) -) - - -class DataEnrichmentContext(AgentDataModel): - """Study evidence that may help resolve organism and feature policy.""" - - studyContext: str = "" - organismHint: str = "" - tissueReferences: list[str] = Field(default_factory=list) - cellTypeReferences: list[str] = Field(default_factory=list) - experimentalDetails: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "DataEnrichmentContext": - return cls() - - @classmethod - def get_example(cls) -> "DataEnrichmentContext": - return cls( - studyContext="Single-cell profiling of treated lung tissue", - organismHint="human", - tissueReferences=["lung"], - cellTypeReferences=["alveolar macrophage", "T cell"], - experimentalDetails=["CRISPR perturbation", "10x 3 prime RNA-seq"], - ) - - -class StudyContextSummary(AgentDataModel): - """Verbatim, evidence-backed references extracted from the study context.""" - - studyContext: str = "" - organismReferences: list[str] = Field(default_factory=list) - tissueReferences: list[str] = Field(default_factory=list) - cellTypeReferences: list[str] = Field(default_factory=list) - experimentalReferences: list[str] = Field(default_factory=list) - hypothesisReferences: list[str] = Field(default_factory=list) - analysisIntentReferences: list[str] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "StudyContextSummary": - return cls() - - @classmethod - def get_example(cls) -> "StudyContextSummary": - return cls( - studyContext=( - "Single-cell profiling of treated human lung tests whether " - "treatment changes alveolar macrophage states." - ), - organismReferences=["human"], - tissueReferences=["lung"], - cellTypeReferences=["alveolar macrophage"], - experimentalReferences=["treated"], - hypothesisReferences=["treatment changes alveolar macrophage states"], - analysisIntentReferences=["Single-cell profiling"], - evidenceIds=["context:study"], - ) - - -class AdtControlEvidence(AgentDataModel): - """One exact observed ADT feature carrying an explicit control token.""" - - featureId: str - featureName: str - matchedToken: Literal["control", "isotype"] - evidenceId: str - - @classmethod - def get_blank(cls) -> "AdtControlEvidence": - return cls( - featureId="", - featureName="", - matchedToken="control", - evidenceId="", - ) - - @classmethod - def get_example(cls) -> "AdtControlEvidence": - return cls( - featureId="Mouse-IgG1-Control", - featureName="Mouse IgG1 isotype control", - matchedToken="isotype", - evidenceId="assay:ADT:adtControl:Mouse-IgG1-Control", - ) - - -class HtoTagEvidence(AgentDataModel): - """One exact feature from an assay persisted with the HTO type.""" - - featureId: str - featureName: str - evidenceId: str - - @classmethod - def get_blank(cls) -> "HtoTagEvidence": - return cls(featureId="", featureName="", evidenceId="") - - @classmethod - def get_example(cls) -> "HtoTagEvidence": - return cls( - featureId="HTO-1", - featureName="Sample tag 1", - evidenceId="assay:HTO:htoTag:HTO-1", - ) - - -class AtacCoordinateEvidence(AgentDataModel): - """Validation evidence for exact ATAC feature IDs as genomic intervals.""" - - status: Literal["notApplicable", "valid", "partial", "invalid"] = "notApplicable" - coordinateColumn: Literal["ids"] = "ids" - coordinateFormat: str = "chrom:start-end" - totalFeatures: int = 0 - validFeatures: int = 0 - invalidExamples: list[str] = Field(default_factory=list) - validExamples: list[str] = Field(default_factory=list) - genomeBuild: Literal["unknown"] = "unknown" - evidenceId: str = "" - - @classmethod - def get_blank(cls) -> "AtacCoordinateEvidence": - return cls() - - @classmethod - def get_example(cls) -> "AtacCoordinateEvidence": - return cls( - status="valid", - totalFeatures=2, - validFeatures=2, - validExamples=["chr1:100-200", "chr2:300-450"], - evidenceId="assay:ATAC:atacCoordinates", - ) - - -class AssayModalityEvidence(AgentDataModel): - """Bounded deterministic routing evidence for one persisted assay type.""" - - assayType: str = "Assay" - modality: Literal["RNA", "ATAC", "ADT", "HTO", "unsupported"] = "unsupported" - typeSource: Literal["persisted", "assayClass", "unknown"] = "unknown" - graphEligible: bool = False - markerEligible: bool = False - demultiplexEligible: bool = False - adtControls: list[AdtControlEvidence] = Field(default_factory=list) - htoTags: list[HtoTagEvidence] = Field(default_factory=list) - atacCoordinates: AtacCoordinateEvidence = Field( - default_factory=AtacCoordinateEvidence.get_blank - ) - totalObservedFeatures: int = 0 - reportedFeatures: int = 0 - truncated: bool = False - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "AssayModalityEvidence": - return cls() - - @classmethod - def get_example(cls) -> "AssayModalityEvidence": - control = AdtControlEvidence.get_example() - return cls( - assayType="ADT", - modality="ADT", - typeSource="persisted", - graphEligible=True, - markerEligible=True, - adtControls=[control], - totalObservedFeatures=20, - reportedFeatures=1, - evidenceIds=["assay:ADT:modality", control.evidenceId], - ) - - -class FeatureFamilyEvidence(AgentDataModel): - """One observed feature family from deterministic Scarf analysis.""" - - family: str - species: str = "unknown" - method: str = "" - count: int = 0 - examples: list[str] = Field(default_factory=list) - defaultExclude: bool | None = None - skipped: str | None = None - catalogSuspect: str | None = None - catalogSize: int | None = None - catalogJoinRate: float | None = None - catalogJoined: int | None = None - evidenceId: str - - @classmethod - def get_blank(cls) -> "FeatureFamilyEvidence": - return cls(family="", evidenceId="") - - @classmethod - def get_example(cls) -> "FeatureFamilyEvidence": - return cls( - family="mitochondrial", - species="homo_sapiens", - method="chromosome", - count=2, - examples=["MT-CO1", "MT-CYB"], - defaultExclude=True, - evidenceId="assay:RNA:family:mitochondrial", - ) - - -class ExogenousFeatureEvidence(AgentDataModel): - """One bounded candidate for an artificial or exogenous feature.""" - - featureId: str - featureName: str - score: int = 0 - classification: str = "unresolved" - evidenceId: str - - @classmethod - def get_blank(cls) -> "ExogenousFeatureEvidence": - return cls(featureId="", featureName="", evidenceId="") - - @classmethod - def get_example(cls) -> "ExogenousFeatureEvidence": - return cls( - featureId="ERCC-00002", - featureName="ERCC-00002", - score=4, - classification="potentialExogenous", - evidenceId="assay:RNA:exogenous:ERCC-00002", - ) - - -class AssayFeatureInspection(AgentDataModel): - """Bounded read-only inspection returned to the model.""" - - assay: str - assayKind: str = "" - identity: dict[str, Any] = Field(default_factory=dict) - species: str = "unknown" - speciesMethod: str | None = None - speciesReason: str = "" - families: list[FeatureFamilyEvidence] = Field(default_factory=list) - exogenous: list[ExogenousFeatureEvidence] = Field(default_factory=list) - modalityEvidence: AssayModalityEvidence = Field( - default_factory=AssayModalityEvidence.get_blank - ) - notes: list[str] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "AssayFeatureInspection": - return cls(assay="") - - @classmethod - def get_example(cls) -> "AssayFeatureInspection": - family = FeatureFamilyEvidence.get_example() - modality = AssayModalityEvidence( - assayType="RNA", - modality="RNA", - typeSource="persisted", - graphEligible=True, - markerEligible=True, - totalObservedFeatures=20_000, - evidenceIds=["assay:RNA:modality"], - ) - return cls( - assay="RNA", - assayKind="RNAassay", - identity={"nFeatures": 20_000, "nDuplicateIds": 0}, - species="homo_sapiens", - speciesMethod="ensemblPrefix", - speciesReason="Most feature IDs carry the ENSG prefix", - families=[family], - modalityEvidence=modality, - evidenceIds=[ - "assay:RNA:identity", - "assay:RNA:species", - family.evidenceId, - *modality.evidenceIds, - ], - ) - - -class AssayFeatureInspectionBatch(AgentDataModel): - """All requested assay inspections returned by one model tool call.""" - - inspections: list[AssayFeatureInspection] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "AssayFeatureInspectionBatch": - return cls() - - @classmethod - def get_example(cls) -> "AssayFeatureInspectionBatch": - inspection = AssayFeatureInspection.get_example() - return cls( - inspections=[inspection], - evidenceIds=list(inspection.evidenceIds), - ) - - -class FeatureReference(AgentDataModel): - """An exact feature identifier and name observed in one assay.""" - - featureId: str - featureName: str - - @classmethod - def get_blank(cls) -> "FeatureReference": - return cls(featureId="", featureName="") - - @classmethod - def get_example(cls) -> "FeatureReference": - return cls(featureId="ENSG00000198727", featureName="MT-CYB") - - -class FeatureMatch(AgentDataModel): - """Resolution of one proposed feature against an assay.""" - - query: str - status: Literal["present", "ambiguous", "absent"] - matches: list[FeatureReference] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "FeatureMatch": - return cls(query="", status="absent") - - @classmethod - def get_example(cls) -> "FeatureMatch": - return cls( - query="MT-CYB", - status="present", - matches=[FeatureReference.get_example()], - evidenceIds=["assay:RNA:feature:ENSG00000198727"], - ) - - -class FeatureLookupResult(AgentDataModel): - """Bounded result from exact feature lookup.""" - - assay: str - results: list[FeatureMatch] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "FeatureLookupResult": - return cls(assay="") - - @classmethod - def get_example(cls) -> "FeatureLookupResult": - match = FeatureMatch.get_example() - return cls( - assay="RNA", - results=[match], - evidenceIds=list(match.evidenceIds), - ) - - -class FeatureLookupBatch(AgentDataModel): - """Exact feature lookups for every requested assay in one tool result.""" - - lookups: list[FeatureLookupResult] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "FeatureLookupBatch": - return cls() - - @classmethod - def get_example(cls) -> "FeatureLookupBatch": - lookup = FeatureLookupResult.get_example() - return cls(lookups=[lookup], evidenceIds=list(lookup.evidenceIds)) - - -class FeatureSelectionPolicy(AgentDataModel): - """Grounded feature policy proposed for one assay.""" - - assay: str - species: str = "unknown" - organismName: str = "unknown" - speciesConfidence: Literal["high", "medium", "low", "unknown"] = "unknown" - speciesRationale: str = "" - excludeFamilies: list[str] = Field(default_factory=list) - protectFamilies: list[str] = Field(default_factory=list) - excludeFeatures: list[str] = Field(default_factory=list) - protectFeatures: list[str] = Field(default_factory=list) - artificialFeatures: list[str] = Field(default_factory=list) - tissueReferences: list[str] = Field(default_factory=list) - cellTypeReferences: list[str] = Field(default_factory=list) - experimentalReferences: list[str] = Field(default_factory=list) - assayType: str = "Assay" - assayModality: Literal["RNA", "ATAC", "ADT", "HTO", "unsupported"] = "unsupported" - graphEligible: bool = False - markerEligible: bool = False - demultiplexEligible: bool = False - exactControlFeatures: list[FeatureReference] = Field(default_factory=list) - exactTagFeatures: list[FeatureReference] = Field(default_factory=list) - peakCoordinateStatus: Literal["notApplicable", "valid", "partial", "invalid"] = ( - "notApplicable" - ) - rationale: str = "" - evidenceIds: list[str] = Field(default_factory=list) - - @model_validator(mode="after") - def validate_non_conflicting_policy(self) -> "FeatureSelectionPolicy": - family_overlap = set(self.excludeFamilies) & set(self.protectFamilies) - feature_overlap = set(self.excludeFeatures) & set(self.protectFeatures) - if family_overlap: - raise ValueError( - "feature families cannot be both excluded and protected: " - f"{sorted(family_overlap)}" - ) - if feature_overlap: - raise ValueError( - "features cannot be both excluded and protected: " - f"{sorted(feature_overlap)}" - ) - return self - - @classmethod - def get_blank(cls) -> "FeatureSelectionPolicy": - return cls(assay="") - - @classmethod - def get_example(cls) -> "FeatureSelectionPolicy": - return cls( - assay="RNA", - species="homo_sapiens", - organismName="human", - speciesConfidence="high", - speciesRationale="Gene IDs and study context agree", - excludeFamilies=["mitochondrial", "ribosomal"], - protectFamilies=["cellCycle", "sex"], - artificialFeatures=["ERCC-00002"], - tissueReferences=["lung"], - cellTypeReferences=["alveolar macrophage"], - experimentalReferences=["ERCC spike-in"], - assayType="RNA", - assayModality="RNA", - graphEligible=True, - markerEligible=True, - rationale="Use technical families for feature-selection exclusions", - evidenceIds=["assay:RNA:species", "assay:RNA:family:mitochondrial"], - ) - - -class DataEnrichmentToolCall(AgentDataModel): - """Compact audit record for one read-only model tool call.""" - - name: str - assay: str - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "DataEnrichmentToolCall": - return cls(name="", assay="") - - @classmethod - def get_example(cls) -> "DataEnrichmentToolCall": - return cls( - name="inspect_assay_features", - assay="RNA", - evidenceIds=["assay:RNA:identity", "assay:RNA:species"], - ) - - -class DataEnrichmentReport(AgentDataModel): - """Final grounded report from :class:`DataEnrichmentAgent`.""" - - status: StageStatus - policies: list[FeatureSelectionPolicy] = Field(default_factory=list) - inspections: list[AssayFeatureInspection] = Field(default_factory=list) - studyContextSummary: StudyContextSummary = Field( - default_factory=StudyContextSummary.get_blank - ) - unresolvedQuestions: list[str] = Field(default_factory=list) - limitations: list[str] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - toolCalls: list[DataEnrichmentToolCall] = Field(default_factory=list) - runInfo: AgentRunInfo = Field(default_factory=AgentRunInfo) - - @model_validator(mode="after") - def validate_status(self) -> "DataEnrichmentReport": - if self.status == "done" and not self.policies: - raise ValueError("done reports require at least one feature policy") - if self.status == "needsInput" and not self.unresolvedQuestions: - raise ValueError("needsInput reports require an unresolved question") - if self.status == "failed" and not self.limitations: - raise ValueError("failed reports require a limitation") - return self - - @classmethod - def get_blank(cls) -> "DataEnrichmentReport": - return cls(status="failed", limitations=["No agent result was produced"]) - - @classmethod - def get_example(cls) -> "DataEnrichmentReport": - policy = FeatureSelectionPolicy.get_example() - inspection = AssayFeatureInspection.get_example() - return cls( - status="done", - policies=[policy], - inspections=[inspection], - studyContextSummary=StudyContextSummary.get_example(), - evidenceIds=list(policy.evidenceIds), - toolCalls=[DataEnrichmentToolCall.get_example()], - runInfo=AgentRunInfo.get_example(), - ) - - -class DataEnrichmentDependencies(AgentDataModel): - """Hidden runtime state supplied to read-only enrichment tools.""" - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - store: Any = Field(default=None, exclude=True) - context: DataEnrichmentContext = Field( - default_factory=DataEnrichmentContext.get_blank - ) - assays: list[str] = Field(default_factory=list) - assayTypes: dict[str, str] = Field(default_factory=dict) - cacheDir: Path | None = None - allowDownload: bool = False - evidenceIds: set[str] = Field(default_factory=set) - inspections: dict[str, AssayFeatureInspection] = Field(default_factory=dict) - confirmedFeatures: dict[str, set[str]] = Field(default_factory=dict) - toolCalls: list[DataEnrichmentToolCall] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "DataEnrichmentDependencies": - return cls() - - @classmethod - def get_example(cls) -> "DataEnrichmentDependencies": - return cls( - context=DataEnrichmentContext.get_example(), - assays=["RNA"], - cacheDir=Path("/tmp/scarf-gene-reference"), - allowDownload=False, - evidenceIds={"context:organism", "context:tissue:0"}, - ) - - -def _prepare_data_enrichment_tool( - ctx: RunContext[DataEnrichmentDependencies], - tool_definition: ToolDefinition, -) -> ToolDefinition | None: - """Expose each batched enrichment tool only while its work is pending.""" - deps = ctx.deps - completed_calls = {call.name for call in deps.toolCalls} - inspection_complete = "inspect_assay_features_batch" in completed_calls - - if tool_definition.name == "inspect_assay_features_batch": - return None if inspection_complete else tool_definition - if tool_definition.name == "find_present_features_batch": - if not inspection_complete or tool_definition.name in completed_calls: - return None - return tool_definition - return tool_definition - - -def _persisted_assay_types(store: Any, assays: Sequence[str]) -> dict[str, str]: - """Read exact persisted assay types through the public datastore summary.""" - summary_method = getattr(store, "summary", None) - if not callable(summary_method): - return {} - summary = summary_method() - requested = set(assays) - return { - str(item.name): str(item.assay_type) - for item in getattr(summary, "assays", ()) - if str(item.name) in requested - } - - -def _assay_modality( - assay_type: str | None, - assay_kind: str, -) -> tuple[ - Literal["RNA", "ATAC", "ADT", "HTO", "unsupported"], - str, - Literal["persisted", "assayClass", "unknown"], -]: - """Map persisted types to supported routes, with a mock-store class fallback.""" - if assay_type is not None: - if assay_type == "RNA": - return "RNA", assay_type, "persisted" - if assay_type == "ATAC": - return "ATAC", assay_type, "persisted" - if assay_type == "ADT": - return "ADT", assay_type, "persisted" - if assay_type == "HTO": - return "HTO", assay_type, "persisted" - return "unsupported", assay_type, "persisted" - if assay_kind == "RNAassay": - return "RNA", assay_kind, "assayClass" - if assay_kind == "ATACassay": - return "ATAC", assay_kind, "assayClass" - if assay_kind: - return "unsupported", assay_kind, "assayClass" - return "unsupported", "Assay", "unknown" - - -def _feature_tokens(*values: str) -> set[str]: - """Return literal alphanumeric tokens without accepting generated patterns.""" - text = " ".join(values).casefold() - normalized = "".join( - character if character.isalnum() else " " for character in text - ) - return set(normalized.split()) - - -def _valid_peak_coordinate(value: str) -> bool: - """Validate the documented ``chrom:start-end`` representation exactly.""" - chromosome, separator, interval = value.partition(":") - if not separator or not chromosome: - return False - start_text, separator, end_text = interval.partition("-") - if not separator or not start_text or not end_text: - return False - try: - start = int(start_text) - end = int(end_text) - except ValueError: - return False - return start >= 0 and end > start - - -def _inspect_adt_features( - assay_name: str, - feature_rows: list[tuple[str, str]], -) -> list[AdtControlEvidence]: - """Return control candidates from exact observed ADT features.""" - candidates: list[AdtControlEvidence] = [] - for feature_id, feature_name in feature_rows: - tokens = _feature_tokens(feature_id, feature_name) - matched_token: Literal["control", "isotype"] | None = None - if "isotype" in tokens: - matched_token = "isotype" - elif "control" in tokens: - matched_token = "control" - if matched_token is None: - continue - candidates.append( - AdtControlEvidence( - featureId=feature_id, - featureName=feature_name, - matchedToken=matched_token, - evidenceId=f"assay:{assay_name}:adtControl:{feature_id}", - ) - ) - return candidates - - -def _inspect_hto_features( - assay_name: str, - feature_rows: list[tuple[str, str]], -) -> list[HtoTagEvidence]: - """Return HTO tag evidence in exact observed order.""" - return [ - HtoTagEvidence( - featureId=feature_id, - featureName=feature_name, - evidenceId=f"assay:{assay_name}:htoTag:{feature_id}", - ) - for feature_id, feature_name in feature_rows - ] - - -def _inspect_atac_features( - assay_name: str, - feature_ids: list[str], -) -> AtacCoordinateEvidence: - """Validate exact observed ATAC coordinates without inferring a build.""" - valid_ids = [value for value in feature_ids if _valid_peak_coordinate(value)] - invalid_ids = [value for value in feature_ids if not _valid_peak_coordinate(value)] - if not feature_ids or not valid_ids: - coordinate_status: Literal["valid", "partial", "invalid"] = "invalid" - elif invalid_ids: - coordinate_status = "partial" - else: - coordinate_status = "valid" - return AtacCoordinateEvidence( - status=coordinate_status, - totalFeatures=len(feature_ids), - validFeatures=len(valid_ids), - validExamples=bounded_list(valid_ids, limit=5), - invalidExamples=bounded_list(invalid_ids, limit=5), - evidenceId=f"assay:{assay_name}:atacCoordinates", - ) - - -def _inspect_modality_features( - *, - assay_name: str, - assay: Any, - assay_type: str | None, - assay_kind: str, - identity: dict[str, Any], -) -> AssayModalityEvidence: - """Build bounded modality evidence from exact observed feature metadata.""" - modality, resolved_type, type_source = _assay_modality(assay_type, assay_kind) - modality_evidence_id = f"assay:{assay_name}:modality" - total_features = int(identity.get("nFeatures", 0)) - evidence_ids = [modality_evidence_id] - graph_eligible = modality in {"RNA", "ATAC", "ADT"} - marker_eligible = graph_eligible - - adt_controls: list[AdtControlEvidence] = [] - hto_tags: list[HtoTagEvidence] = [] - atac_coordinates = AtacCoordinateEvidence.get_blank() - reported_features = 0 - truncated = False - - if modality in {"ADT", "HTO", "ATAC"}: - feature_ids = [str(value) for value in assay.feats.fetch_all("ids")] - total_features = len(feature_ids) - else: - feature_ids = [] - - if modality in {"ADT", "HTO"}: - feature_names = [str(value) for value in assay.feats.fetch_all("names")] - feature_rows = list(zip(feature_ids, feature_names, strict=True)) - else: - feature_rows = [] - - if modality == "ADT": - control_candidates = _inspect_adt_features(assay_name, feature_rows) - adt_controls = bounded_list( - control_candidates, - limit=CONFIG._MAX_FEATURE_QUERIES, - ) - evidence_ids.extend(item.evidenceId for item in adt_controls) - reported_features = len(adt_controls) - truncated = len(control_candidates) > len(adt_controls) - - if modality == "HTO": - limited_rows = bounded_list(feature_rows, limit=CONFIG._MAX_FEATURE_QUERIES) - hto_tags = _inspect_hto_features(assay_name, limited_rows) - evidence_ids.extend(item.evidenceId for item in hto_tags) - reported_features = len(hto_tags) - truncated = len(feature_rows) > len(limited_rows) - - if modality == "ATAC": - atac_coordinates = _inspect_atac_features( - assay_name, - feature_ids, - ) - evidence_ids.append(atac_coordinates.evidenceId) - reported_features = len(atac_coordinates.validExamples) + len( - atac_coordinates.invalidExamples - ) - truncated = len(feature_ids) > reported_features - - return AssayModalityEvidence( - assayType=resolved_type, - modality=modality, - typeSource=type_source, - graphEligible=graph_eligible, - markerEligible=marker_eligible, - demultiplexEligible=modality == "HTO", - adtControls=adt_controls, - htoTags=hto_tags, - atacCoordinates=atac_coordinates, - totalObservedFeatures=total_features, - reportedFeatures=reported_features, - truncated=truncated, - evidenceIds=evidence_ids, - ) - - -async def inspect_assay_features( - ctx: RunContext[DataEnrichmentDependencies], - assay_name: str, -) -> AssayFeatureInspection: - """Inspect feature identity, species evidence, families, and exogenous cues.""" - deps = ctx.deps - if deps.store is None: - raise ModelRetry("The datastore is unavailable") - if assay_name not in deps.assays: - raise ModelRetry( - f"assay_name must be one of the requested assays: {deps.assays}" - ) - - characterization = characterize_features( - deps.store, - studyContext=deps.context.studyContext, - model=None, - assays=[assay_name], - cacheDir=deps.cacheDir, - allowDownload=deps.allowDownload, - ) - if characterization.status != "done" or not characterization.assays: - logger.warning(f"Data Enrichment could not characterize assay {assay_name!r}") - detail = "; ".join(characterization.notes) or "feature inspection failed" - raise ModelRetry(detail) - - record = characterization.assays[0] - family_evidence: list[FeatureFamilyEvidence] = [] - evidence_ids = [f"assay:{assay_name}:identity", f"assay:{assay_name}:species"] - for family in record.get("families", []): - family_name = str(family.get("family", "")) - evidence_id = f"assay:{assay_name}:family:{family_name}" - family_evidence.append( - FeatureFamilyEvidence( - family=family_name, - species=str(family.get("species", record.get("species", "unknown"))), - method=str(family.get("method", "")), - count=int(family.get("count", 0)), - examples=[str(value) for value in family.get("examples", [])], - defaultExclude=family.get("defaultExclude"), - skipped=family.get("skipped"), - catalogSuspect=family.get("catalogSuspect"), - catalogSize=family.get("catalogSize"), - catalogJoinRate=family.get("catalogJoinRate"), - catalogJoined=family.get("catalogJoined"), - evidenceId=evidence_id, - ) - ) - evidence_ids.append(evidence_id) - - exogenous_evidence: list[ExogenousFeatureEvidence] = [] - for item in record.get("exogenous", []): - feature_id = str(item.get("id", "")) - evidence_id = f"assay:{assay_name}:exogenous:{feature_id}" - exogenous_evidence.append( - ExogenousFeatureEvidence( - featureId=feature_id, - featureName=str(item.get("name", "")), - score=int(item.get("score", 0)), - classification=str(item.get("class", "unresolved")), - evidenceId=evidence_id, - ) - ) - evidence_ids.append(evidence_id) - - resolution = record.get("speciesResolution") or {} - assay = deps.store.get_assay(assay_name) - identity = dict(record.get("identity") or {}) - modality_evidence = _inspect_modality_features( - assay_name=assay_name, - assay=assay, - assay_type=deps.assayTypes.get(assay_name), - assay_kind=str(record.get("assayKind", "")), - identity=identity, - ) - evidence_ids.extend(modality_evidence.evidenceIds) - inspection = AssayFeatureInspection( - assay=assay_name, - assayKind=str(record.get("assayKind", "")), - identity=identity, - species=str(record.get("species", "unknown")), - speciesMethod=record.get("speciesMethod"), - speciesReason=str(resolution.get("reason", "")), - families=family_evidence, - exogenous=exogenous_evidence, - modalityEvidence=modality_evidence, - notes=[str(value) for value in record.get("notes", [])], - evidenceIds=evidence_ids, - ) - deps.inspections[assay_name] = inspection - deps.evidenceIds.update(evidence_ids) - deps.toolCalls.append( - DataEnrichmentToolCall( - name="inspect_assay_features", - assay=assay_name, - evidenceIds=evidence_ids, - ) - ) - logger.debug( - "Data Enrichment inspected " - f"assay={assay_name!r}, modality={modality_evidence.modality}, " - f"species={inspection.species}, families={len(family_evidence)}, " - f"exogenous={len(exogenous_evidence)}, evidence={len(evidence_ids)}" - ) - return inspection - - -async def inspect_assay_features_batch( - ctx: RunContext[DataEnrichmentDependencies], -) -> AssayFeatureInspectionBatch: - """Inspect every requested assay and return one bounded tool result.""" - deps = ctx.deps - if not deps.assays: - raise ModelRetry("No assays were requested") - logger.info( - f"Data Enrichment feature inspection started for {len(deps.assays)} assays" - ) - start = len(deps.toolCalls) - inspections = [ - await inspect_assay_features(ctx, assay_name=assay_name) - for assay_name in deps.assays - ] - del deps.toolCalls[start:] - evidence_ids = list( - dict.fromkeys( - evidence_id - for inspection in inspections - for evidence_id in inspection.evidenceIds - ) - ) - deps.toolCalls.append( - DataEnrichmentToolCall( - name="inspect_assay_features_batch", - assay=",".join(deps.assays), - evidenceIds=evidence_ids, - ) - ) - supported_routes = sum( - inspection.modalityEvidence.modality != "unsupported" - for inspection in inspections - ) - logger.info( - "Data Enrichment feature inspection completed: " - f"assays={len(inspections)}, supportedRoutes={supported_routes}, " - f"evidence={len(evidence_ids)}" - ) - return AssayFeatureInspectionBatch( - inspections=inspections, - evidenceIds=evidence_ids, - ) - - -async def find_present_features( - ctx: RunContext[DataEnrichmentDependencies], - assay_name: str, - queries: list[str], -) -> FeatureLookupResult: - """Resolve a bounded list of gene IDs or names against one exact assay.""" - deps = ctx.deps - if deps.store is None: - raise ModelRetry("The datastore is unavailable") - if assay_name not in deps.assays: - raise ModelRetry( - f"assay_name must be one of the requested assays: {deps.assays}" - ) - clean_queries = list( - dict.fromkeys(value.strip() for value in queries if value.strip()) - ) - if not clean_queries or len(clean_queries) > CONFIG._MAX_FEATURE_QUERIES: - raise ModelRetry( - f"queries must contain between 1 and {CONFIG._MAX_FEATURE_QUERIES} values" - ) - - assay = deps.store.get_assay(assay_name) - feature_ids = [str(value) for value in assay.feats.fetch_all("ids")] - feature_names = [str(value) for value in assay.feats.fetch_all("names")] - rows = list(zip(feature_ids, feature_names, strict=True)) - results: list[FeatureMatch] = [] - result_evidence_ids: list[str] = [] - confirmed = deps.confirmedFeatures.setdefault(assay_name, set()) - - for query in clean_queries: - exact = [row for row in rows if query in row] - candidates = exact - if not candidates: - folded = query.casefold() - candidates = [ - row - for row in rows - if folded == row[0].casefold() or folded == row[1].casefold() - ] - unique_candidates = bounded_list( - dict.fromkeys(candidates), - limit=10, - ) - references = [ - FeatureReference(featureId=feature_id, featureName=feature_name) - for feature_id, feature_name in unique_candidates - ] - evidence_ids = [ - f"assay:{assay_name}:feature:{reference.featureId}" - for reference in references - ] - if len(references) == 1: - status: Literal["present", "ambiguous", "absent"] = "present" - confirmed.update({references[0].featureId, references[0].featureName}) - deps.evidenceIds.update(evidence_ids) - result_evidence_ids.extend(evidence_ids) - elif references: - status = "ambiguous" - else: - status = "absent" - results.append( - FeatureMatch( - query=query, - status=status, - matches=references, - evidenceIds=evidence_ids if status == "present" else [], - ) - ) - - result = FeatureLookupResult( - assay=assay_name, - results=results, - evidenceIds=list(dict.fromkeys(result_evidence_ids)), - ) - deps.toolCalls.append( - DataEnrichmentToolCall( - name="find_present_features", - assay=assay_name, - evidenceIds=result.evidenceIds, - ) - ) - return result - - -async def find_present_features_batch( - ctx: RunContext[DataEnrichmentDependencies], - queries_by_assay: dict[str, list[str]], -) -> FeatureLookupBatch: - """Resolve all proposed individual features through one model tool call.""" - deps = ctx.deps - unknown_assays = sorted(set(queries_by_assay) - set(deps.assays)) - if unknown_assays: - raise ModelRetry(f"Unknown requested assays: {unknown_assays}") - if not queries_by_assay: - raise ModelRetry("queries_by_assay must contain at least one assay") - clean_queries_by_assay = { - assay_name: list( - dict.fromkeys(value.strip() for value in queries if value.strip()) - ) - for assay_name, queries in queries_by_assay.items() - } - empty_assays = sorted( - assay_name - for assay_name, queries in clean_queries_by_assay.items() - if not queries - ) - if empty_assays: - raise ModelRetry(f"Feature-query batches cannot be empty: {empty_assays}") - query_count = sum(len(queries) for queries in clean_queries_by_assay.values()) - if query_count > CONFIG._MAX_FEATURE_QUERIES: - raise ModelRetry( - "The batch may contain at most " - f"{CONFIG._MAX_FEATURE_QUERIES} feature queries in total" - ) - - logger.info( - "Data Enrichment feature lookup started: " - f"assays={len(clean_queries_by_assay)}, queries={query_count}" - ) - - start = len(deps.toolCalls) - lookups = [ - await find_present_features( - ctx, - assay_name=assay_name, - queries=clean_queries_by_assay[assay_name], - ) - for assay_name in deps.assays - if assay_name in clean_queries_by_assay - ] - del deps.toolCalls[start:] - evidence_ids = list( - dict.fromkeys( - evidence_id for lookup in lookups for evidence_id in lookup.evidenceIds - ) - ) - deps.toolCalls.append( - DataEnrichmentToolCall( - name="find_present_features_batch", - assay=",".join(queries_by_assay), - evidenceIds=evidence_ids, - ) - ) - result_counts = {"present": 0, "ambiguous": 0, "absent": 0} - for lookup in lookups: - for result in lookup.results: - result_counts[result.status] += 1 - logger.info( - "Data Enrichment feature lookup completed: " - f"present={result_counts['present']}, " - f"ambiguous={result_counts['ambiguous']}, " - f"absent={result_counts['absent']}, evidence={len(evidence_ids)}" - ) - return FeatureLookupBatch(lookups=lookups, evidenceIds=evidence_ids) - - -def _ground_study_context_summary( - context: DataEnrichmentContext, - proposed: StudyContextSummary, -) -> StudyContextSummary: - """Bind structured context references to exact caller text.""" - original_context = context.studyContext - organism_references = [context.organismHint] if context.organismHint else [] - for species in _SUPPORTED_SPECIES.values(): - match = re.search( - rf"\b{re.escape(species.label)}\b", - original_context, - flags=re.IGNORECASE, - ) - if match is not None: - organism_references.append(match.group(0)) - field_sources = { - "organismReferences": organism_references, - "tissueReferences": list(context.tissueReferences), - "cellTypeReferences": list(context.cellTypeReferences), - "experimentalReferences": list(context.experimentalDetails), - "hypothesisReferences": [], - "analysisIntentReferences": [], - } - grounded: dict[str, list[str]] = {} - for field_name, supplied_values in field_sources.items(): - exact_supplied = [value.strip() for value in supplied_values if value.strip()] - proposed_values = list(getattr(proposed, field_name)) - combined = list( - dict.fromkeys( - value.strip() - for value in [*exact_supplied, *proposed_values] - if value.strip() - ) - ) - if len(combined) > 12: - raise ValueError( - f"studyContextSummary.{field_name} may contain at most 12 values" - ) - supplied = set(exact_supplied) - invalid = [ - value - for value in combined - if value not in supplied and value not in original_context - ] - if invalid: - raise ValueError( - f"Study-context references must be verbatim caller text: {invalid}" - ) - oversized = [value for value in combined if len(value) > 240] - if oversized: - raise ValueError("Study-context references may not exceed 240 characters") - grounded[field_name] = combined - - evidence_ids: list[str] = [] - if original_context: - evidence_ids.append("context:study") - if context.organismHint: - evidence_ids.append("context:organism") - evidence_ids.extend( - f"context:tissue:{index}" - for index, _value in enumerate(context.tissueReferences) - ) - evidence_ids.extend( - f"context:cellType:{index}" - for index, _value in enumerate(context.cellTypeReferences) - ) - evidence_ids.extend( - f"context:experiment:{index}" - for index, _value in enumerate(context.experimentalDetails) - ) - return StudyContextSummary( - studyContext=original_context, - **grounded, - evidenceIds=evidence_ids, - ) - - -def _validate_feature_policy( - deps: DataEnrichmentDependencies, - policy: FeatureSelectionPolicy, - grounded_context: StudyContextSummary, -) -> None: - """Ground and validate one feature policy in deterministic order.""" - supported_species = {*_SUPPORTED_SPECIES, "unknown"} - if policy.species not in supported_species: - raise ValueError( - f"unsupported species {policy.species!r}; choose a supported key or unknown" - ) - if not policy.evidenceIds: - raise ValueError(f"policy for assay {policy.assay!r} requires evidence IDs") - policy.organismName = ( - _SUPPORTED_SPECIES[policy.species].label - if policy.species in _SUPPORTED_SPECIES - else "unknown" - ) - inspection = deps.inspections.get(policy.assay) - if inspection is None: - raise ValueError(f"assay {policy.assay!r} was not inspected") - modality = inspection.modalityEvidence - policy.assayType = modality.assayType - policy.assayModality = modality.modality - policy.graphEligible = modality.graphEligible - policy.markerEligible = modality.markerEligible - policy.demultiplexEligible = modality.demultiplexEligible - policy.exactControlFeatures = [ - FeatureReference( - featureId=item.featureId, - featureName=item.featureName, - ) - for item in modality.adtControls - ] - policy.exactTagFeatures = [ - FeatureReference( - featureId=item.featureId, - featureName=item.featureName, - ) - for item in modality.htoTags - ] - policy.peakCoordinateStatus = modality.atacCoordinates.status - policy.evidenceIds = list( - dict.fromkeys([*policy.evidenceIds, *modality.evidenceIds]) - ) - if ( - inspection.species in _SUPPORTED_SPECIES - and policy.species != inspection.species - ): - raise ValueError( - f"policy species {policy.species!r} conflicts with inspected " - f"species {inspection.species!r}" - ) - if ( - inspection.species == "unknown" - and policy.species != "unknown" - and not any( - evidence_id.startswith("context:") for evidence_id in policy.evidenceIds - ) - ): - raise ValueError( - "A context-derived species decision must cite context evidence" - ) - observed_families = {item.family for item in inspection.families} - cited_families = set(policy.excludeFamilies) | set(policy.protectFamilies) - unknown_families = cited_families - observed_families - if unknown_families: - raise ValueError( - f"policy cites unobserved families: {sorted(unknown_families)}" - ) - protected_defaults = { - item.family for item in inspection.families if item.defaultExclude is False - } - excluded_protected = sorted( - set(policy.excludeFamilies).intersection(protected_defaults) - ) - if excluded_protected: - raise ValueError( - "The initial enrichment policy cannot exclude families that " - f"deterministic evidence protects by default: {excluded_protected}" - ) - confirmed = deps.confirmedFeatures.get(policy.assay, set()) - cited_features = { - *policy.excludeFeatures, - *policy.protectFeatures, - *policy.artificialFeatures, - } - unknown_features = cited_features - confirmed - if unknown_features: - raise ValueError( - "Call find_present_features_batch before citing individual features: " - f"{sorted(unknown_features)}" - ) - exogenous_evidence = { - value: item.evidenceId - for item in inspection.exogenous - for value in (item.featureId, item.featureName) - } - unsupported_artificial: list[str] = [] - for feature in policy.artificialFeatures: - evidence_id = exogenous_evidence.get(feature) - if evidence_id is not None and evidence_id in policy.evidenceIds: - continue - matching_context_ids = { - f"context:experiment:{index}" - for index, detail in enumerate(deps.context.experimentalDetails) - if feature.casefold() in detail.casefold() - } - if matching_context_ids.intersection(policy.evidenceIds): - continue - unsupported_artificial.append(feature) - if unsupported_artificial: - raise ValueError( - "Artificial features require their exogenous evidence ID or a " - "feature-specific experimental-context evidence ID: " - f"{sorted(unsupported_artificial)}" - ) - unknown_evidence = set(policy.evidenceIds) - deps.evidenceIds - if unknown_evidence: - raise ValueError( - f"policy cites unknown evidence IDs: {sorted(unknown_evidence)}" - ) - policy.tissueReferences = list(grounded_context.tissueReferences) - policy.cellTypeReferences = list(grounded_context.cellTypeReferences) - policy.experimentalReferences = list(grounded_context.experimentalReferences) - - -def validate_data_enrichment_report( - deps: DataEnrichmentDependencies, - report: DataEnrichmentReport, -) -> DataEnrichmentReport: - """Ground an agent report in inspected assays, context, and exact lookups.""" - if not deps.inspections: - raise ValueError("Inspect every requested assay before returning the report") - - requested = set(deps.assays) - reported = {policy.assay for policy in report.policies} - if not reported.issubset(requested): - raise ValueError( - f"policies cite assays outside the requested set: {sorted(reported - requested)}" - ) - if report.status == "done" and reported != requested: - raise ValueError( - f"done reports require one policy for every requested assay: {deps.assays}" - ) - - grounded_context = _ground_study_context_summary( - deps.context, - report.studyContextSummary, - ) - for policy in report.policies: - _validate_feature_policy(deps, policy, grounded_context) - - report.studyContextSummary = grounded_context - report.inspections = [deps.inspections[name] for name in deps.assays] - report.toolCalls = list(deps.toolCalls) - report.evidenceIds = list( - dict.fromkeys( - evidence_id - for evidence_id in [ - *report.studyContextSummary.evidenceIds, - *( - evidence_id - for policy in report.policies - for evidence_id in policy.evidenceIds - ), - ] - ) - ) - logger.debug( - "Data Enrichment report validated: " - f"status={report.status}, policies={len(report.policies)}, " - f"inspections={len(report.inspections)}, " - f"toolCalls={len(report.toolCalls)}, evidence={len(report.evidenceIds)}" - ) - return report - - -class DataEnrichmentAgent: - """A small read-only tool agent for feature and organism enrichment.""" - - def __init__( - self, - model: Any, - *, - config: AgentRunConfig | None = None, - ) -> None: - self.model = model - self.config = (config or AgentRunConfig()).with_limits( - request_limit=8, - tool_call_limit=5, - output_token_limit=32768, - timeout_seconds=600.0, - ) - - def run( - self, - store: Any, - *, - context: DataEnrichmentContext | None = None, - assays: Sequence[str] | None = None, - cache_dir: Path | str | None = None, - allow_download: bool = False, - ) -> DataEnrichmentReport: - """Run the bounded tool loop without mutating the supplied datastore.""" - available_assays = [str(value) for value in store.assay_names] - selected_assays = ( - [str(value) for value in assays] if assays is not None else available_assays - ) - unknown_assays = sorted(set(selected_assays) - set(available_assays)) - if unknown_assays: - raise ValueError(f"unknown assays: {unknown_assays}") - if not selected_assays: - raise ValueError("at least one assay is required") - - logger.info( - "Data Enrichment Agent started: " - f"assays={len(selected_assays)}, allowDownload={allow_download}" - ) - - enrichment_context = context or DataEnrichmentContext.get_blank() - evidence_ids: set[str] = set() - if enrichment_context.studyContext: - evidence_ids.add("context:study") - if enrichment_context.organismHint: - evidence_ids.add("context:organism") - evidence_ids.update( - f"context:tissue:{index}" - for index, _value in enumerate(enrichment_context.tissueReferences) - ) - evidence_ids.update( - f"context:cellType:{index}" - for index, _value in enumerate(enrichment_context.cellTypeReferences) - ) - evidence_ids.update( - f"context:experiment:{index}" - for index, _value in enumerate(enrichment_context.experimentalDetails) - ) - deps = DataEnrichmentDependencies( - store=store, - context=enrichment_context, - assays=selected_assays, - assayTypes=_persisted_assay_types(store, selected_assays), - cacheDir=Path(cache_dir) if cache_dir is not None else None, - allowDownload=allow_download, - evidenceIds=evidence_ids, - ) - user_prompt = ( - dedent( - """ - Enrich the feature policy for assays: {assays}. - Study context: {study_context} - Organism hint: {organism_hint} - Tissue references: {tissue_references} - Cell-type references: {cell_type_references} - Experimental details: {experimental_details} - - Call inspect_assay_features_batch exactly once to inspect every - assay together. If a policy needs individual features, collect all - proposed names for all assays and call - find_present_features_batch exactly once. Do not call a singular - assay tool or split lookups across calls. A batched tool is removed - after it succeeds, so use each call to request all required data. - Populate studyContextSummary only with exact verbatim spans from - the paragraph or caller references. Empty optional hint fields do - not erase references present in the paragraph. Before returning, - verify that every explicit organism, tissue, cell population, - experiment, hypothesis, and analysis intent has been placed in its - corresponding summary list. Copy modality evidence from the - inspection without changing feature IDs or names. - """ - ) - .strip() - .format( - assays=", ".join(selected_assays), - study_context=enrichment_context.studyContext or "not provided", - organism_hint=enrichment_context.organismHint or "not provided", - tissue_references=", ".join(enrichment_context.tissueReferences) - or "not provided", - cell_type_references=", ".join(enrichment_context.cellTypeReferences) - or "not provided", - experimental_details=", ".join(enrichment_context.experimentalDetails) - or "not provided", - ) - ) - execution = run_agent_sync( - model=self.model, - output_type=DataEnrichmentReport, - system_prompt=_SYSTEM_PROMPT, - user_prompt=user_prompt, - tools=[ - Tool( - inspect_assay_features_batch, - prepare=_prepare_data_enrichment_tool, - sequential=self.config.sequentialTools, - timeout=self.config.timeoutSeconds, - ), - Tool( - find_present_features_batch, - prepare=_prepare_data_enrichment_tool, - sequential=self.config.sequentialTools, - timeout=self.config.timeoutSeconds, - ), - ], - deps_type=DataEnrichmentDependencies, - deps=deps, - config=self.config, - name="data_enrichment", - output_validator=lambda report: validate_data_enrichment_report( - deps, - report, - ), - ) - report = DataEnrichmentReport.model_validate(execution.output) - report = validate_data_enrichment_report(deps, report) - report.runInfo = execution.runInfo - logger.info( - "Data Enrichment Agent completed: " - f"status={report.status}, policies={len(report.policies)}, " - f"toolCalls={len(report.toolCalls)}, evidence={len(report.evidenceIds)}" - ) - return report diff --git a/scarf/agent/data_enrichment/__init__.py b/scarf/agent/data_enrichment/__init__.py new file mode 100644 index 00000000..3288727c --- /dev/null +++ b/scarf/agent/data_enrichment/__init__.py @@ -0,0 +1,61 @@ +"""Read-only feature and organism enrichment agent.""" + +from .agent import DataEnrichmentAgent +from .contracts import ( + AdtControlEvidence, + AssayFeatureInspection, + AssayFeatureInspectionBatch, + AssayModalityEvidence, + AtacCoordinateEvidence, + DataEnrichmentContext, + DataEnrichmentDependencies, + DataEnrichmentReport, + DataEnrichmentToolCall, + DefaultHvgFamilyEvidence, + ExogenousFeatureEvidence, + FeatureFamilyEvidence, + FeatureLookupBatch, + FeatureLookupResult, + FeatureMatch, + FeatureReference, + FeatureSelectionPolicy, + HtoTagEvidence, + RnaFeatureInventoryEvidence, + StudyContextSummary, +) +from .tools import ( + find_present_features, + find_present_features_batch, + inspect_assay_features, + inspect_assay_features_batch, +) +from .validation import validate_data_enrichment_report + +__all__ = [ + "AdtControlEvidence", + "AssayFeatureInspection", + "AssayFeatureInspectionBatch", + "AssayModalityEvidence", + "AtacCoordinateEvidence", + "DataEnrichmentAgent", + "DataEnrichmentContext", + "DataEnrichmentDependencies", + "DataEnrichmentReport", + "DataEnrichmentToolCall", + "DefaultHvgFamilyEvidence", + "ExogenousFeatureEvidence", + "FeatureFamilyEvidence", + "FeatureLookupResult", + "FeatureLookupBatch", + "FeatureMatch", + "FeatureReference", + "FeatureSelectionPolicy", + "HtoTagEvidence", + "RnaFeatureInventoryEvidence", + "StudyContextSummary", + "find_present_features", + "find_present_features_batch", + "inspect_assay_features", + "inspect_assay_features_batch", + "validate_data_enrichment_report", +] diff --git a/scarf/agent/data_enrichment/agent.py b/scarf/agent/data_enrichment/agent.py new file mode 100644 index 00000000..239c34a6 --- /dev/null +++ b/scarf/agent/data_enrichment/agent.py @@ -0,0 +1,266 @@ +"""Data enrichment prompt and agent runner.""" + +from collections.abc import Callable, Sequence +from pathlib import Path +from textwrap import dedent +from typing import Any + +from ...utils.logging import logger +from .._deps import AGENT_INSTALL_HINT +from ..config import AgentRunConfig +from ..config.agent_exec import run_agent_sync +from ..types import AgentRunInfo +from .contracts import ( + DataEnrichmentContext, + DataEnrichmentDependencies, + DataEnrichmentReport, +) +from .tools import ( + _prepare_data_enrichment_tool, + find_present_features_batch, + inspect_assay_features_batch, +) +from .validation import ( + _SUPPORTED_SPECIES, + failed_data_enrichment_report, + validate_data_enrichment_report, +) + +try: + from pydantic_ai import Tool, UnexpectedModelBehavior, UsageLimitExceeded +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + +_SYSTEM_PROMPT = ( + dedent( + """ + You are Scarf's Data Enrichment Agent. Work only through the supplied + read-only tools. Inspect every requested assay before making a decision. + Use gene identifiers and names together with the supplied organism hint, + tissue references, cell-type references, and experimental details when + species evidence is ambiguous. Supported species keys are: {supported_species}. + + Call inspect_assay_features_batch once for all requested assays. Never + invent a feature. If individual features are needed, collect all proposed + names across assays and call find_present_features_batch once before + placing them in a policy. Do not call feature lookup when no individual + feature decision is needed. Absent or ambiguous lookup results must never + enter a policy. If inspection resolves a supported species, copy that exact + species key. Use caller organism context only when inspection leaves the + species unknown. Use excludeFamilies only to nominate one conditional + representation-sensitivity bundle from observed families with + defaultExclude=true. It is not an instruction to remove those families. + Never nominate a family with defaultExclude=false. + The defaultFeatureInventory is separate deterministic evidence for Scarf's + exact default HVG blacklist. It is evidence only, not an automatic + exclusion or a source of policy nominations. Keep marker eligibility + broader than any graph-feature exclusion. + + Persisted assay types determine modality routes; never infer a route from + an assay label. The validator fills assay type, modality eligibility, ADT + controls, HTO tags, ATAC-coordinate status, inspections, tool calls, and + report-level evidence. Do not return those derived fields. + Treat Ensembl release misses as unresolved, + not artificial. Mitochondrial, ribosomal, and histone families may be + sensitivity candidates. Sex-linked and cell-cycle families are protected + by default. Marker testing retains conditional biological families. + + Structure studyContextSummary using only verbatim spans from the supplied + study paragraph, study objective, or exact caller references. Do not + paraphrase, infer, or + invent an organism, tissue, cell type, experiment, hypothesis, or analysis + intent. Empty optional hint lists do not mean that the paragraph lacks + those references. Select at most 12 objective-relevant spans per category, + each at most 240 characters. These lists are excerpts, not an exhaustive + replacement for the full context and objective that Scarf preserves and + passes downstream. The validator binds the original paragraph and exact + caller references. Return a bounded + report with citations copied from tool or context evidence IDs. + Do not write code, mutate the datastore, or request arbitrary Scarf calls. + """ + ) + .strip() + .format(supported_species=", ".join(sorted(_SUPPORTED_SPECIES))) +) + + +def _persisted_assay_types(store: Any, assays: Sequence[str]) -> dict[str, str]: + """Read exact persisted assay types through the public datastore summary.""" + summary_method = getattr(store, "summary", None) + if not callable(summary_method): + return {} + summary = summary_method() + requested = set(assays) + return { + str(item.name): str(item.assay_type) + for item in getattr(summary, "assays", ()) + if str(item.name) in requested + } + + +class DataEnrichmentAgent: + """A small read-only tool agent for feature and organism enrichment.""" + + def __init__( + self, + model: Any, + *, + config: AgentRunConfig | None = None, + ) -> None: + self.model = model + self.config = (config or AgentRunConfig()).with_limits( + request_limit=8, + tool_call_limit=5, + output_token_limit=32768, + timeout_seconds=600.0, + ) + + def run( + self, + store: Any, + *, + context: DataEnrichmentContext | None = None, + assays: Sequence[str] | None = None, + cache_dir: Path | str | None = None, + allow_download: bool = False, + on_attempt: Callable[[AgentRunInfo], None] | None = None, + ) -> DataEnrichmentReport: + """Run the bounded tool loop without mutating the supplied datastore.""" + available_assays = [str(value) for value in store.assay_names] + selected_assays = ( + [str(value) for value in assays] if assays is not None else available_assays + ) + unknown_assays = sorted(set(selected_assays) - set(available_assays)) + if unknown_assays: + raise ValueError(f"unknown assays: {unknown_assays}") + if not selected_assays: + raise ValueError("at least one assay is required") + + logger.info( + "Data Enrichment Agent started: " + f"assays={len(selected_assays)}, allowDownload={allow_download}" + ) + + enrichment_context = context or DataEnrichmentContext.get_blank() + evidence_ids: set[str] = set() + if enrichment_context.studyContext: + evidence_ids.add("context:study") + if enrichment_context.studyObjective: + evidence_ids.add("context:objective") + if enrichment_context.organismHint: + evidence_ids.add("context:organism") + evidence_ids.update( + f"context:tissue:{index}" + for index, _value in enumerate(enrichment_context.tissueReferences) + ) + evidence_ids.update( + f"context:cellType:{index}" + for index, _value in enumerate(enrichment_context.cellTypeReferences) + ) + evidence_ids.update( + f"context:experiment:{index}" + for index, _value in enumerate(enrichment_context.experimentalDetails) + ) + deps = DataEnrichmentDependencies( + store=store, + context=enrichment_context, + assays=selected_assays, + assayTypes=_persisted_assay_types(store, selected_assays), + cacheDir=Path(cache_dir) if cache_dir is not None else None, + allowDownload=allow_download, + evidenceIds=evidence_ids, + ) + user_prompt = ( + dedent( + """ + Enrich the feature policy for assays: {assays}. + Study context: {study_context} + Study objective: {study_objective} + Organism hint: {organism_hint} + Tissue references: {tissue_references} + Cell-type references: {cell_type_references} + Experimental details: {experimental_details} + + Call inspect_assay_features_batch exactly once to inspect every + assay together. If a policy needs individual features, collect all + proposed names for all assays and call + find_present_features_batch exactly once. Do not call a singular + assay tool or split lookups across calls. A batched tool is removed + after it succeeds, so use each call to request all required data. + If no policy needs an individual feature, do not call feature lookup + and keep excludeFeatures, protectFeatures, and artificialFeatures + empty. Return exactly one policy for every requested assay. Copy a + resolved inspection species exactly; otherwise use unknown unless + exact caller context supports a species. Exclude only observed + defaultExclude=true families and protect every observed + defaultExclude=false family. + Populate studyContextSummary only with exact verbatim spans from + the paragraph or caller references. Empty optional hint fields do + not erase references present in the paragraph. Before returning, + select at most 12 verbatim references per category, each no more + than 240 characters, prioritizing the objective. These bounded + excerpts do not replace the preserved full context. Do not return + inspections, modality-derived fields, exact controls and tags, + toolCalls, runInfo, or report evidence; validation attaches them. + """ + ) + .strip() + .format( + assays=", ".join(selected_assays), + study_context=enrichment_context.studyContext or "not provided", + study_objective=enrichment_context.studyObjective or "not provided", + organism_hint=enrichment_context.organismHint or "not provided", + tissue_references=", ".join(enrichment_context.tissueReferences) + or "not provided", + cell_type_references=", ".join(enrichment_context.cellTypeReferences) + or "not provided", + experimental_details=", ".join(enrichment_context.experimentalDetails) + or "not provided", + ) + ) + try: + execution = run_agent_sync( + model=self.model, + output_type=DataEnrichmentReport, + system_prompt=_SYSTEM_PROMPT, + user_prompt=user_prompt, + tools=[ + Tool( + inspect_assay_features_batch, + max_retries=1, + prepare=_prepare_data_enrichment_tool, + sequential=self.config.sequentialTools, + timeout=self.config.timeoutSeconds, + ), + Tool( + find_present_features_batch, + max_retries=1, + prepare=_prepare_data_enrichment_tool, + sequential=self.config.sequentialTools, + timeout=self.config.timeoutSeconds, + ), + ], + deps_type=DataEnrichmentDependencies, + deps=deps, + config=self.config, + name="data_enrichment", + on_attempt=on_attempt, + output_validator=lambda report: validate_data_enrichment_report( + deps, + report, + ), + ) + except (UnexpectedModelBehavior, UsageLimitExceeded) as exc: + model_name = getattr(self.model, "model_name", type(self.model).__name__) + return failed_data_enrichment_report( + deps, error=exc, model_name=str(model_name) + ) + report = DataEnrichmentReport.model_validate(execution.output) + report = validate_data_enrichment_report(deps, report) + report.runInfo = execution.runInfo + logger.info( + "Data Enrichment Agent completed: " + f"status={report.status}, policies={len(report.policies)}, " + f"toolCalls={len(report.toolCalls)}, evidence={len(report.evidenceIds)}" + ) + return report diff --git a/scarf/agent/characterize_features.py b/scarf/agent/data_enrichment/characterization.py similarity index 82% rename from scarf/agent/characterize_features.py rename to scarf/agent/data_enrichment/characterization.py index b86f9b6b..428c8faf 100644 --- a/scarf/agent/characterize_features.py +++ b/scarf/agent/data_enrichment/characterization.py @@ -1,18 +1,19 @@ """Characterize feature identity, species, families, and exogenous candidates.""" +import re from collections.abc import Mapping, Sequence from pathlib import Path from typing import Any -from ..assay import RNAassay -from ..features.gene_reference import ( +from ...assay import RNAassay +from ...features.gene_reference import ( GeneReference, default_cache_dir, ensure_reference, load_reference, species_registry, ) -from ..features.identity import ( +from ...features.identity import ( audit_feature_identity, backfill_symbols, exogenous_candidates, @@ -20,16 +21,16 @@ reference_misses, resolve_species, ) -from ..quality_control.cell_cycle_genes import ( +from ...features.variability import DEFAULT_HVG_BLACKLIST +from ...quality_control.cell_cycle_genes import ( g2m_phase_genes, g2m_phase_genes_mouse, s_phase_genes, s_phase_genes_mouse, ) -from .config import CONFIG -from .config._deps import AGENT_INSTALL_HINT -from .decide import DecisionValidationError, decide -from .types import AgentDataModel, Decision, EvidenceItem, StageStatus +from .._deps import AGENT_INSTALL_HINT +from ..decisions.selection import DecisionValidationError, decide +from ..types import AgentDataModel, Decision, EvidenceItem, StageStatus try: from pydantic import Field @@ -47,6 +48,26 @@ "mus_musculus": {"s": s_phase_genes_mouse, "g2m": g2m_phase_genes_mouse}, } _SEX_COEFFICIENT_TOKENS = frozenset({"sex", "gender", "Sex", "Gender"}) +_FEATURE_INVENTORY_EXAMPLE_LIMIT = 8 +_CONTEXT_LIMIT = 1200 +_AUTO_DOWNLOAD_SPECIES = frozenset({"homo_sapiens", "mus_musculus"}) +_MAX_EXOGENOUS = 25 +_DEFAULT_HVG_FAMILY_PATTERNS = ( + ("mitochondrial", r"^MT-"), + ("ribosomalProtein", r"^RPS|^RPL"), + ("mitoribosomal", r"^MRPS|^MRPL"), + ("cellCycleCcn", r"^CCN"), + ("hla", r"^HLA-"), + ("h2", r"^H2-"), + ("histone", r"^HIST"), + ( + "sexLinked", + ( + r"^XIST$|^DDX3Y$|^USP9Y$|^EIF1AY$|^KDM5D$|^SRY$|^ZFY$|^UTY$|" + r"^TMSB4Y$|^NLGN4Y$" + ), + ), +) class FeatureCharacterization(AgentDataModel): @@ -61,22 +82,58 @@ class FeatureCharacterization(AgentDataModel): def get_blank(cls) -> "FeatureCharacterization": return cls(status="failed") - @classmethod - def get_example(cls) -> "FeatureCharacterization": - return cls( - status="done", - notes=["Feature identity and families were characterized."], - assays=[{"assay": "RNA", "species": "homo_sapiens"}], - ) - def _bounded_context(study_context: str | None) -> str: text = (study_context or "").strip() - return ( - text - if len(text) <= CONFIG._CONTEXT_LIMIT - else text[: CONFIG._CONTEXT_LIMIT - 3] + "..." - ) + return text if len(text) <= _CONTEXT_LIMIT else text[: _CONTEXT_LIMIT - 3] + "..." + + +def _feature_pattern_matches(names: Sequence[str], pattern: str) -> list[str]: + compiled = re.compile(pattern.upper()) + return [name for name in names if compiled.match(name.upper()) is not None] + + +def _bounded_feature_examples(matches: Sequence[str]) -> list[str]: + return sorted(set(matches), key=lambda value: (value.casefold(), value))[ + :_FEATURE_INVENTORY_EXAMPLE_LIMIT + ] + + +def _scarf_default_feature_inventory( + assay_name: str, + names: Sequence[str], +) -> dict[str, Any]: + evidence_prefix = f"assay:{assay_name}:scarfDefaultHvg" + combined_matches = _feature_pattern_matches(names, DEFAULT_HVG_BLACKLIST) + families: list[dict[str, Any]] = [] + family_evidence_ids: list[str] = [] + for family, pattern in _DEFAULT_HVG_FAMILY_PATTERNS: + matches = _feature_pattern_matches(names, pattern) + evidence_id = f"{evidence_prefix}:family:{family}" + families.append( + { + "family": family, + "pattern": pattern, + "caseInsensitive": True, + "count": len(matches), + "examples": _bounded_feature_examples(matches), + "evidenceId": evidence_id, + } + ) + family_evidence_ids.append(evidence_id) + evidence_id = f"{evidence_prefix}:combined" + return { + "source": "scarfDefaultHvgBlacklist", + "policyEffect": "evidenceOnly", + "featureColumn": "names", + "totalFeatures": len(names), + "blacklist": DEFAULT_HVG_BLACKLIST, + "matchCount": len(combined_matches), + "examples": _bounded_feature_examples(combined_matches), + "families": families, + "evidenceId": evidence_id, + "evidenceIds": [evidence_id, *family_evidence_ids], + } def _audit( @@ -324,6 +381,19 @@ def _characterize_assay( ) return record + default_inventory = _scarf_default_feature_inventory(assay_name, names) + record["defaultFeatureInventory"] = default_inventory + _audit( + audit_log, + kind="scarfDefaultFeatureInventory", + detail=( + f"Scarf's default HVG blacklist matched " + f"{default_inventory['matchCount']} of {len(names)} RNA features" + ), + assay=assay_name, + evidenceIds=default_inventory["evidenceIds"], + ) + species_by_assay = dict(directions.get("speciesByAssay") or {}) directed_species = species_by_assay.get(assay_name) resolution = resolve_species( @@ -386,7 +456,7 @@ def _characterize_assay( # Only human/mouse auto-download; any other species needs an explicit direction. may_download = allow_download and ( - species in CONFIG._AUTO_DOWNLOAD_SPECIES or directed_species == species + species in _AUTO_DOWNLOAD_SPECIES or directed_species == species ) reference = _load_or_fetch_reference( species, @@ -445,17 +515,17 @@ def _characterize_assay( assay=assay_name, ) - raw_max = directions.get("maxExogenousCandidates", CONFIG._MAX_EXOGENOUS) + raw_max = directions.get("maxExogenousCandidates", _MAX_EXOGENOUS) try: max_exogenous = int(raw_max) except (TypeError, ValueError): _audit( audit_log, kind="invalidDirection", - detail=f"maxExogenousCandidates={raw_max!r}; using {CONFIG._MAX_EXOGENOUS}", + detail=f"maxExogenousCandidates={raw_max!r}; using {_MAX_EXOGENOUS}", assay=assay_name, ) - max_exogenous = CONFIG._MAX_EXOGENOUS + max_exogenous = _MAX_EXOGENOUS if reference is not None: misses = reference_misses(ids, symbols, reference) if misses["count"]: diff --git a/scarf/agent/data_enrichment/contracts.py b/scarf/agent/data_enrichment/contracts.py new file mode 100644 index 00000000..edabd158 --- /dev/null +++ b/scarf/agent/data_enrichment/contracts.py @@ -0,0 +1,395 @@ +"""Serializable contracts for data enrichment.""" + +from pathlib import Path +from typing import Any, Literal + +from .._deps import AGENT_INSTALL_HINT +from ..types import AgentDataModel, AgentRunInfo, StageStatus + +try: + from pydantic import ConfigDict, Field, model_validator + from pydantic.json_schema import SkipJsonSchema +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + + +class DataEnrichmentContext(AgentDataModel): + """Study evidence that may help resolve organism and feature policy.""" + + studyContext: str = "" + studyObjective: str = "" + organismHint: str = "" + tissueReferences: list[str] = Field(default_factory=list) + cellTypeReferences: list[str] = Field(default_factory=list) + experimentalDetails: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "DataEnrichmentContext": + return cls() + + +class StudyContextSummary(AgentDataModel): + """Verbatim, evidence-backed references extracted from the study context.""" + + studyContext: SkipJsonSchema[str] = "" + studyObjective: SkipJsonSchema[str] = "" + organismReferences: list[str] = Field(default_factory=list) + tissueReferences: list[str] = Field(default_factory=list) + cellTypeReferences: list[str] = Field(default_factory=list) + experimentalReferences: list[str] = Field(default_factory=list) + hypothesisReferences: list[str] = Field(default_factory=list) + analysisIntentReferences: list[str] = Field(default_factory=list) + evidenceIds: SkipJsonSchema[list[str]] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "StudyContextSummary": + return cls() + + +class AdtControlEvidence(AgentDataModel): + """One exact observed ADT feature carrying an explicit control token.""" + + featureId: str + featureName: str + matchedToken: Literal["control", "isotype"] + evidenceId: str + + @classmethod + def get_blank(cls) -> "AdtControlEvidence": + return cls( + featureId="", + featureName="", + matchedToken="control", + evidenceId="", + ) + + +class HtoTagEvidence(AgentDataModel): + """One exact feature from an assay persisted with the HTO type.""" + + featureId: str + featureName: str + evidenceId: str + + @classmethod + def get_blank(cls) -> "HtoTagEvidence": + return cls(featureId="", featureName="", evidenceId="") + + +class AtacCoordinateEvidence(AgentDataModel): + """Validation evidence for exact ATAC feature IDs as genomic intervals.""" + + status: Literal["notApplicable", "valid", "partial", "invalid"] = "notApplicable" + coordinateColumn: Literal["ids"] = "ids" + coordinateFormat: str = "chrom:start-end" + totalFeatures: int = 0 + validFeatures: int = 0 + invalidExamples: list[str] = Field(default_factory=list) + validExamples: list[str] = Field(default_factory=list) + genomeBuild: Literal["unknown"] = "unknown" + evidenceId: str = "" + + @classmethod + def get_blank(cls) -> "AtacCoordinateEvidence": + return cls() + + +class AssayModalityEvidence(AgentDataModel): + """Bounded deterministic routing evidence for one persisted assay type.""" + + assayType: str = "Assay" + modality: Literal["RNA", "ATAC", "ADT", "HTO", "unsupported"] = "unsupported" + typeSource: Literal["persisted", "assayClass", "unknown"] = "unknown" + graphEligible: bool = False + markerEligible: bool = False + demultiplexEligible: bool = False + adtControls: list[AdtControlEvidence] = Field(default_factory=list) + htoTags: list[HtoTagEvidence] = Field(default_factory=list) + atacCoordinates: AtacCoordinateEvidence = Field( + default_factory=AtacCoordinateEvidence.get_blank + ) + totalObservedFeatures: int = 0 + reportedFeatures: int = 0 + truncated: bool = False + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "AssayModalityEvidence": + return cls() + + +class FeatureFamilyEvidence(AgentDataModel): + """One observed feature family from deterministic Scarf analysis.""" + + family: str + species: str = "unknown" + method: str = "" + count: int = 0 + examples: list[str] = Field(default_factory=list) + defaultExclude: bool | None = None + skipped: str | None = None + catalogSuspect: str | None = None + catalogSize: int | None = None + catalogJoinRate: float | None = None + catalogJoined: int | None = None + evidenceId: str + + @classmethod + def get_blank(cls) -> "FeatureFamilyEvidence": + return cls(family="", evidenceId="") + + +class DefaultHvgFamilyEvidence(AgentDataModel): + """One case-insensitive family within Scarf's default HVG blacklist.""" + + family: str = "" + pattern: str = "" + caseInsensitive: Literal[True] = True + count: int = 0 + examples: list[str] = Field(default_factory=list) + evidenceId: str = "" + + @classmethod + def get_blank(cls) -> "DefaultHvgFamilyEvidence": + return cls() + + +class RnaFeatureInventoryEvidence(AgentDataModel): + """Exact name-column matches for Scarf's default HVG blacklist.""" + + source: Literal["scarfDefaultHvgBlacklist"] = "scarfDefaultHvgBlacklist" + policyEffect: Literal["evidenceOnly"] = "evidenceOnly" + featureColumn: Literal["names"] = "names" + totalFeatures: int = 0 + blacklist: str = "" + matchCount: int = 0 + examples: list[str] = Field(default_factory=list) + families: list[DefaultHvgFamilyEvidence] = Field(default_factory=list) + evidenceId: str = "" + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "RnaFeatureInventoryEvidence": + return cls() + + +class ExogenousFeatureEvidence(AgentDataModel): + """One bounded candidate for an artificial or exogenous feature.""" + + featureId: str + featureName: str + score: int = 0 + classification: str = "unresolved" + evidenceId: str + + @classmethod + def get_blank(cls) -> "ExogenousFeatureEvidence": + return cls(featureId="", featureName="", evidenceId="") + + +class AssayFeatureInspection(AgentDataModel): + """Bounded read-only inspection returned to the model.""" + + assay: str + assayKind: str = "" + identity: dict[str, Any] = Field(default_factory=dict) + species: str = "unknown" + speciesMethod: str | None = None + speciesReason: str = "" + families: list[FeatureFamilyEvidence] = Field(default_factory=list) + defaultFeatureInventory: RnaFeatureInventoryEvidence | None = None + exogenous: list[ExogenousFeatureEvidence] = Field(default_factory=list) + modalityEvidence: AssayModalityEvidence = Field( + default_factory=AssayModalityEvidence.get_blank + ) + notes: list[str] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "AssayFeatureInspection": + return cls(assay="") + + +class AssayFeatureInspectionBatch(AgentDataModel): + """All requested assay inspections returned by one model tool call.""" + + inspections: list[AssayFeatureInspection] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "AssayFeatureInspectionBatch": + return cls() + + +class FeatureReference(AgentDataModel): + """An exact feature identifier and name observed in one assay.""" + + featureId: str + featureName: str + + @classmethod + def get_blank(cls) -> "FeatureReference": + return cls(featureId="", featureName="") + + +class FeatureMatch(AgentDataModel): + """Resolution of one proposed feature against an assay.""" + + query: str + status: Literal["present", "ambiguous", "absent"] + matches: list[FeatureReference] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "FeatureMatch": + return cls(query="", status="absent") + + +class FeatureLookupResult(AgentDataModel): + """Bounded result from exact feature lookup.""" + + assay: str + results: list[FeatureMatch] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "FeatureLookupResult": + return cls(assay="") + + +class FeatureLookupBatch(AgentDataModel): + """Exact feature lookups for every requested assay in one tool result.""" + + lookups: list[FeatureLookupResult] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "FeatureLookupBatch": + return cls() + + +class FeatureSelectionPolicy(AgentDataModel): + """Grounded feature policy proposed for one assay.""" + + assay: str + species: str = "unknown" + organismName: SkipJsonSchema[str] = "unknown" + speciesConfidence: Literal["high", "medium", "low", "unknown"] = "unknown" + speciesRationale: str = "" + excludeFamilies: list[str] = Field(default_factory=list) + protectFamilies: list[str] = Field(default_factory=list) + excludeFeatures: list[str] = Field(default_factory=list) + protectFeatures: list[str] = Field(default_factory=list) + artificialFeatures: list[str] = Field(default_factory=list) + tissueReferences: SkipJsonSchema[list[str]] = Field(default_factory=list) + cellTypeReferences: SkipJsonSchema[list[str]] = Field(default_factory=list) + experimentalReferences: SkipJsonSchema[list[str]] = Field(default_factory=list) + assayType: SkipJsonSchema[str] = "Assay" + assayModality: SkipJsonSchema[ + Literal["RNA", "ATAC", "ADT", "HTO", "unsupported"] + ] = "unsupported" + graphEligible: SkipJsonSchema[bool] = False + markerEligible: SkipJsonSchema[bool] = False + demultiplexEligible: SkipJsonSchema[bool] = False + exactControlFeatures: SkipJsonSchema[list[FeatureReference]] = Field( + default_factory=list + ) + exactTagFeatures: SkipJsonSchema[list[FeatureReference]] = Field( + default_factory=list + ) + peakCoordinateStatus: SkipJsonSchema[ + Literal["notApplicable", "valid", "partial", "invalid"] + ] = "notApplicable" + rationale: str = "" + evidenceIds: list[str] = Field(default_factory=list) + + @model_validator(mode="after") + def validate_non_conflicting_policy(self) -> "FeatureSelectionPolicy": + family_overlap = set(self.excludeFamilies) & set(self.protectFamilies) + feature_overlap = set(self.excludeFeatures) & set(self.protectFeatures) + if family_overlap: + raise ValueError( + "feature families cannot be both excluded and protected: " + f"{sorted(family_overlap)}" + ) + if feature_overlap: + raise ValueError( + "features cannot be both excluded and protected: " + f"{sorted(feature_overlap)}" + ) + return self + + @classmethod + def get_blank(cls) -> "FeatureSelectionPolicy": + return cls(assay="") + + +class DataEnrichmentToolCall(AgentDataModel): + """Compact audit record for one read-only model tool call.""" + + name: str + assay: str + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "DataEnrichmentToolCall": + return cls(name="", assay="") + + +class DataEnrichmentReport(AgentDataModel): + """Final grounded report from :class:`DataEnrichmentAgent`.""" + + status: StageStatus + policies: list[FeatureSelectionPolicy] = Field(default_factory=list) + inspections: SkipJsonSchema[list[AssayFeatureInspection]] = Field( + default_factory=list + ) + studyContextSummary: StudyContextSummary = Field( + default_factory=StudyContextSummary.get_blank + ) + unresolvedQuestions: list[str] = Field(default_factory=list) + limitations: list[str] = Field(default_factory=list) + evidenceIds: SkipJsonSchema[list[str]] = Field(default_factory=list) + toolCalls: SkipJsonSchema[list[DataEnrichmentToolCall]] = Field( + default_factory=list + ) + runInfo: SkipJsonSchema[AgentRunInfo] = Field(default_factory=AgentRunInfo) + + @model_validator(mode="after") + def validate_status(self) -> "DataEnrichmentReport": + if self.status == "done" and not self.policies: + raise ValueError("done reports require at least one feature policy") + if self.status == "needsInput" and not self.unresolvedQuestions: + raise ValueError("needsInput reports require an unresolved question") + if self.status == "failed" and not self.limitations: + raise ValueError("failed reports require a limitation") + return self + + @classmethod + def get_blank(cls) -> "DataEnrichmentReport": + return cls(status="failed", limitations=["No agent result was produced"]) + + +class DataEnrichmentDependencies(AgentDataModel): + """Hidden runtime state supplied to read-only enrichment tools.""" + + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + + store: Any = Field(default=None, exclude=True) + context: DataEnrichmentContext = Field( + default_factory=DataEnrichmentContext.get_blank + ) + assays: list[str] = Field(default_factory=list) + assayTypes: dict[str, str] = Field(default_factory=dict) + cacheDir: Path | None = None + allowDownload: bool = False + evidenceIds: set[str] = Field(default_factory=set) + inspections: dict[str, AssayFeatureInspection] = Field(default_factory=dict) + confirmedFeatures: dict[str, set[str]] = Field(default_factory=dict) + lookupBatch: FeatureLookupBatch | None = Field(default=None, exclude=True) + lookupQueries: dict[str, list[str]] = Field(default_factory=dict, exclude=True) + toolCalls: list[DataEnrichmentToolCall] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "DataEnrichmentDependencies": + return cls() diff --git a/scarf/agent/data_enrichment/tools.py b/scarf/agent/data_enrichment/tools.py new file mode 100644 index 00000000..2e60046a --- /dev/null +++ b/scarf/agent/data_enrichment/tools.py @@ -0,0 +1,615 @@ +"""Read-only feature inspection and lookup tools.""" + +from typing import Any, Literal + +from ...features.variability import DEFAULT_HVG_BLACKLIST +from ...utils.logging import logger +from .._deps import AGENT_INSTALL_HINT +from ..tools import bounded_list +from .characterization import characterize_features +from .contracts import ( + AdtControlEvidence, + AssayFeatureInspection, + AssayFeatureInspectionBatch, + AssayModalityEvidence, + AtacCoordinateEvidence, + DataEnrichmentDependencies, + DataEnrichmentToolCall, + ExogenousFeatureEvidence, + FeatureFamilyEvidence, + FeatureLookupBatch, + FeatureLookupResult, + FeatureMatch, + FeatureReference, + HtoTagEvidence, + RnaFeatureInventoryEvidence, +) + +try: + from pydantic_ai import ModelRetry, RunContext + from pydantic_ai.tools import ToolDefinition +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + +_MAX_FEATURE_QUERIES = 50 + + +def _prepare_data_enrichment_tool( + ctx: RunContext[DataEnrichmentDependencies], + tool_definition: ToolDefinition, +) -> ToolDefinition | None: + """Expose each batched enrichment tool only while its work is pending.""" + deps = ctx.deps + completed_calls = {call.name for call in deps.toolCalls} + inspection_complete = "inspect_assay_features_batch" in completed_calls + + if tool_definition.name == "inspect_assay_features_batch": + return None if inspection_complete else tool_definition + if tool_definition.name == "find_present_features_batch": + if not inspection_complete or tool_definition.name in completed_calls: + return None + return tool_definition + return tool_definition + + +def _assay_modality( + assay_type: str | None, + assay_kind: str, +) -> tuple[ + Literal["RNA", "ATAC", "ADT", "HTO", "unsupported"], + str, + Literal["persisted", "assayClass", "unknown"], +]: + """Map persisted types to supported routes, with a mock-store class fallback.""" + if assay_type is not None: + if assay_type == "RNA": + return "RNA", assay_type, "persisted" + if assay_type == "ATAC": + return "ATAC", assay_type, "persisted" + if assay_type == "ADT": + return "ADT", assay_type, "persisted" + if assay_type == "HTO": + return "HTO", assay_type, "persisted" + return "unsupported", assay_type, "persisted" + if assay_kind == "RNAassay": + return "RNA", assay_kind, "assayClass" + if assay_kind == "ATACassay": + return "ATAC", assay_kind, "assayClass" + if assay_kind: + return "unsupported", assay_kind, "assayClass" + return "unsupported", "Assay", "unknown" + + +def _feature_tokens(*values: str) -> set[str]: + """Return literal alphanumeric tokens without accepting generated patterns.""" + text = " ".join(values).casefold() + normalized = "".join( + character if character.isalnum() else " " for character in text + ) + return set(normalized.split()) + + +def _valid_peak_coordinate(value: str) -> bool: + """Validate the documented ``chrom:start-end`` representation exactly.""" + chromosome, separator, interval = value.partition(":") + if not separator or not chromosome: + return False + start_text, separator, end_text = interval.partition("-") + if not separator or not start_text or not end_text: + return False + try: + start = int(start_text) + end = int(end_text) + except ValueError: + return False + return start >= 0 and end > start + + +def _inspect_adt_features( + assay_name: str, + feature_rows: list[tuple[str, str]], +) -> list[AdtControlEvidence]: + """Return control candidates from exact observed ADT features.""" + candidates: list[AdtControlEvidence] = [] + for feature_id, feature_name in feature_rows: + tokens = _feature_tokens(feature_id, feature_name) + matched_token: Literal["control", "isotype"] | None = None + if "isotype" in tokens: + matched_token = "isotype" + elif "control" in tokens: + matched_token = "control" + if matched_token is None: + continue + candidates.append( + AdtControlEvidence( + featureId=feature_id, + featureName=feature_name, + matchedToken=matched_token, + evidenceId=f"assay:{assay_name}:adtControl:{feature_id}", + ) + ) + return candidates + + +def _inspect_hto_features( + assay_name: str, + feature_rows: list[tuple[str, str]], +) -> list[HtoTagEvidence]: + """Return HTO tag evidence in exact observed order.""" + return [ + HtoTagEvidence( + featureId=feature_id, + featureName=feature_name, + evidenceId=f"assay:{assay_name}:htoTag:{feature_id}", + ) + for feature_id, feature_name in feature_rows + ] + + +def _inspect_atac_features( + assay_name: str, + feature_ids: list[str], +) -> AtacCoordinateEvidence: + """Validate exact observed ATAC coordinates without inferring a build.""" + valid_ids = [value for value in feature_ids if _valid_peak_coordinate(value)] + invalid_ids = [value for value in feature_ids if not _valid_peak_coordinate(value)] + if not feature_ids or not valid_ids: + coordinate_status: Literal["valid", "partial", "invalid"] = "invalid" + elif invalid_ids: + coordinate_status = "partial" + else: + coordinate_status = "valid" + return AtacCoordinateEvidence( + status=coordinate_status, + totalFeatures=len(feature_ids), + validFeatures=len(valid_ids), + validExamples=bounded_list(valid_ids, limit=5), + invalidExamples=bounded_list(invalid_ids, limit=5), + evidenceId=f"assay:{assay_name}:atacCoordinates", + ) + + +def _inspect_modality_features( + *, + assay_name: str, + assay: Any, + assay_type: str | None, + assay_kind: str, + identity: dict[str, Any], +) -> AssayModalityEvidence: + """Build bounded modality evidence from exact observed feature metadata.""" + modality, resolved_type, type_source = _assay_modality(assay_type, assay_kind) + modality_evidence_id = f"assay:{assay_name}:modality" + total_features = int(identity.get("nFeatures", 0)) + evidence_ids = [modality_evidence_id] + graph_eligible = modality in {"RNA", "ATAC", "ADT"} + marker_eligible = graph_eligible + + adt_controls: list[AdtControlEvidence] = [] + hto_tags: list[HtoTagEvidence] = [] + atac_coordinates = AtacCoordinateEvidence.get_blank() + reported_features = 0 + truncated = False + + if modality in {"ADT", "HTO", "ATAC"}: + feature_ids = [str(value) for value in assay.feats.fetch_all("ids")] + total_features = len(feature_ids) + else: + feature_ids = [] + + if modality in {"ADT", "HTO"}: + feature_names = [str(value) for value in assay.feats.fetch_all("names")] + feature_rows = list(zip(feature_ids, feature_names, strict=True)) + else: + feature_rows = [] + + if modality == "ADT": + control_candidates = _inspect_adt_features(assay_name, feature_rows) + adt_controls = bounded_list( + control_candidates, + limit=_MAX_FEATURE_QUERIES, + ) + evidence_ids.extend(item.evidenceId for item in adt_controls) + reported_features = len(adt_controls) + truncated = len(control_candidates) > len(adt_controls) + + if modality == "HTO": + limited_rows = bounded_list(feature_rows, limit=_MAX_FEATURE_QUERIES) + hto_tags = _inspect_hto_features(assay_name, limited_rows) + evidence_ids.extend(item.evidenceId for item in hto_tags) + reported_features = len(hto_tags) + truncated = len(feature_rows) > len(limited_rows) + + if modality == "ATAC": + atac_coordinates = _inspect_atac_features( + assay_name, + feature_ids, + ) + evidence_ids.append(atac_coordinates.evidenceId) + reported_features = len(atac_coordinates.validExamples) + len( + atac_coordinates.invalidExamples + ) + truncated = len(feature_ids) > reported_features + + return AssayModalityEvidence( + assayType=resolved_type, + modality=modality, + typeSource=type_source, + graphEligible=graph_eligible, + markerEligible=marker_eligible, + demultiplexEligible=modality == "HTO", + adtControls=adt_controls, + htoTags=hto_tags, + atacCoordinates=atac_coordinates, + totalObservedFeatures=total_features, + reportedFeatures=reported_features, + truncated=truncated, + evidenceIds=evidence_ids, + ) + + +async def inspect_assay_features( + ctx: RunContext[DataEnrichmentDependencies], + assay_name: str, +) -> AssayFeatureInspection: + """Inspect feature identity, species evidence, families, and exogenous cues.""" + deps = ctx.deps + if deps.store is None: + raise ModelRetry("The datastore is unavailable") + if assay_name not in deps.assays: + raise ModelRetry( + f"assay_name must be one of the requested assays: {deps.assays}" + ) + cached = deps.inspections.get(assay_name) + if cached is not None: + logger.debug( + f"Data Enrichment reused cached inspection for assay {assay_name!r}" + ) + return cached + + characterization = characterize_features( + deps.store, + studyContext=deps.context.studyContext, + model=None, + assays=[assay_name], + cacheDir=deps.cacheDir, + allowDownload=deps.allowDownload, + ) + if characterization.status != "done" or not characterization.assays: + logger.warning(f"Data Enrichment could not characterize assay {assay_name!r}") + detail = "; ".join(characterization.notes) or "feature inspection failed" + raise ModelRetry(detail) + + record = characterization.assays[0] + family_evidence: list[FeatureFamilyEvidence] = [] + evidence_ids = [f"assay:{assay_name}:identity", f"assay:{assay_name}:species"] + raw_default_inventory = record.get("defaultFeatureInventory") + default_inventory: RnaFeatureInventoryEvidence | None = None + if raw_default_inventory is not None: + default_inventory = RnaFeatureInventoryEvidence.model_validate( + raw_default_inventory + ) + if default_inventory.blacklist != DEFAULT_HVG_BLACKLIST: + raise ModelRetry( + "RNA feature inventory does not use Scarf's exact default HVG blacklist" + ) + evidence_prefix = f"assay:{assay_name}:scarfDefaultHvg" + default_inventory.evidenceId = f"{evidence_prefix}:combined" + for family in default_inventory.families: + family.evidenceId = f"{evidence_prefix}:family:{family.family}" + default_inventory.evidenceIds = [ + default_inventory.evidenceId, + *(family.evidenceId for family in default_inventory.families), + ] + evidence_ids.extend(default_inventory.evidenceIds) + for family in record.get("families", []): + family_name = str(family.get("family", "")) + evidence_id = f"assay:{assay_name}:family:{family_name}" + family_evidence.append( + FeatureFamilyEvidence( + family=family_name, + species=str(family.get("species", record.get("species", "unknown"))), + method=str(family.get("method", "")), + count=int(family.get("count", 0)), + examples=[str(value) for value in family.get("examples", [])], + defaultExclude=family.get("defaultExclude"), + skipped=family.get("skipped"), + catalogSuspect=family.get("catalogSuspect"), + catalogSize=family.get("catalogSize"), + catalogJoinRate=family.get("catalogJoinRate"), + catalogJoined=family.get("catalogJoined"), + evidenceId=evidence_id, + ) + ) + evidence_ids.append(evidence_id) + + exogenous_evidence: list[ExogenousFeatureEvidence] = [] + for item in record.get("exogenous", []): + feature_id = str(item.get("id", "")) + evidence_id = f"assay:{assay_name}:exogenous:{feature_id}" + exogenous_evidence.append( + ExogenousFeatureEvidence( + featureId=feature_id, + featureName=str(item.get("name", "")), + score=int(item.get("score", 0)), + classification=str(item.get("class", "unresolved")), + evidenceId=evidence_id, + ) + ) + evidence_ids.append(evidence_id) + + resolution = record.get("speciesResolution") or {} + assay = deps.store.get_assay(assay_name) + identity = dict(record.get("identity") or {}) + modality_evidence = _inspect_modality_features( + assay_name=assay_name, + assay=assay, + assay_type=deps.assayTypes.get(assay_name), + assay_kind=str(record.get("assayKind", "")), + identity=identity, + ) + evidence_ids.extend(modality_evidence.evidenceIds) + inspection = AssayFeatureInspection( + assay=assay_name, + assayKind=str(record.get("assayKind", "")), + identity=identity, + species=str(record.get("species", "unknown")), + speciesMethod=record.get("speciesMethod"), + speciesReason=str(resolution.get("reason", "")), + families=family_evidence, + defaultFeatureInventory=default_inventory, + exogenous=exogenous_evidence, + modalityEvidence=modality_evidence, + notes=[str(value) for value in record.get("notes", [])], + evidenceIds=evidence_ids, + ) + deps.inspections[assay_name] = inspection + deps.evidenceIds.update(evidence_ids) + deps.toolCalls.append( + DataEnrichmentToolCall( + name="inspect_assay_features", + assay=assay_name, + evidenceIds=evidence_ids, + ) + ) + logger.debug( + "Data Enrichment inspected " + f"assay={assay_name!r}, modality={modality_evidence.modality}, " + f"species={inspection.species}, families={len(family_evidence)}, " + f"defaultHvgMatches=" + f"{default_inventory.matchCount if default_inventory is not None else 0}, " + f"exogenous={len(exogenous_evidence)}, evidence={len(evidence_ids)}" + ) + return inspection + + +async def inspect_assay_features_batch( + ctx: RunContext[DataEnrichmentDependencies], +) -> AssayFeatureInspectionBatch: + """Inspect every requested assay and return one bounded tool result.""" + deps = ctx.deps + if not deps.assays: + raise ModelRetry("No assays were requested") + if any( + call.name == "inspect_assay_features_batch" for call in deps.toolCalls + ) and all(assay_name in deps.inspections for assay_name in deps.assays): + inspections = [deps.inspections[assay_name] for assay_name in deps.assays] + evidence_ids = list( + dict.fromkeys( + evidence_id + for inspection in inspections + for evidence_id in inspection.evidenceIds + ) + ) + logger.info("Data Enrichment reused the completed feature inspection batch") + return AssayFeatureInspectionBatch( + inspections=inspections, + evidenceIds=evidence_ids, + ) + logger.info( + f"Data Enrichment feature inspection started for {len(deps.assays)} assays" + ) + start = len(deps.toolCalls) + try: + inspections = [ + await inspect_assay_features(ctx, assay_name=assay_name) + for assay_name in deps.assays + ] + except Exception: + del deps.toolCalls[start:] + raise + del deps.toolCalls[start:] + evidence_ids = list( + dict.fromkeys( + evidence_id + for inspection in inspections + for evidence_id in inspection.evidenceIds + ) + ) + deps.toolCalls.append( + DataEnrichmentToolCall( + name="inspect_assay_features_batch", + assay=",".join(deps.assays), + evidenceIds=evidence_ids, + ) + ) + supported_routes = sum( + inspection.modalityEvidence.modality != "unsupported" + for inspection in inspections + ) + logger.info( + "Data Enrichment feature inspection completed: " + f"assays={len(inspections)}, supportedRoutes={supported_routes}, " + f"evidence={len(evidence_ids)}" + ) + return AssayFeatureInspectionBatch( + inspections=inspections, + evidenceIds=evidence_ids, + ) + + +async def find_present_features( + ctx: RunContext[DataEnrichmentDependencies], + assay_name: str, + queries: list[str], +) -> FeatureLookupResult: + """Resolve a bounded list of gene IDs or names against one exact assay.""" + deps = ctx.deps + if deps.store is None: + raise ModelRetry("The datastore is unavailable") + if assay_name not in deps.assays: + raise ModelRetry( + f"assay_name must be one of the requested assays: {deps.assays}" + ) + clean_queries = list( + dict.fromkeys(value.strip() for value in queries if value.strip()) + ) + if not clean_queries or len(clean_queries) > _MAX_FEATURE_QUERIES: + raise ModelRetry( + f"queries must contain between 1 and {_MAX_FEATURE_QUERIES} values" + ) + + assay = deps.store.get_assay(assay_name) + feature_ids = [str(value) for value in assay.feats.fetch_all("ids")] + feature_names = [str(value) for value in assay.feats.fetch_all("names")] + rows = list(zip(feature_ids, feature_names, strict=True)) + results: list[FeatureMatch] = [] + result_evidence_ids: list[str] = [] + confirmed = deps.confirmedFeatures.setdefault(assay_name, set()) + + for query in clean_queries: + exact = [row for row in rows if query in row] + candidates = exact + if not candidates: + folded = query.casefold() + candidates = [ + row + for row in rows + if folded == row[0].casefold() or folded == row[1].casefold() + ] + unique_candidates = bounded_list( + dict.fromkeys(candidates), + limit=10, + ) + references = [ + FeatureReference(featureId=feature_id, featureName=feature_name) + for feature_id, feature_name in unique_candidates + ] + evidence_ids = [ + f"assay:{assay_name}:feature:{reference.featureId}" + for reference in references + ] + if len(references) == 1: + status: Literal["present", "ambiguous", "absent"] = "present" + confirmed.update({references[0].featureId, references[0].featureName}) + deps.evidenceIds.update(evidence_ids) + result_evidence_ids.extend(evidence_ids) + elif references: + status = "ambiguous" + else: + status = "absent" + results.append( + FeatureMatch( + query=query, + status=status, + matches=references, + evidenceIds=evidence_ids if status == "present" else [], + ) + ) + + result = FeatureLookupResult( + assay=assay_name, + results=results, + evidenceIds=list(dict.fromkeys(result_evidence_ids)), + ) + deps.toolCalls.append( + DataEnrichmentToolCall( + name="find_present_features", + assay=assay_name, + evidenceIds=result.evidenceIds, + ) + ) + return result + + +async def find_present_features_batch( + ctx: RunContext[DataEnrichmentDependencies], + queries_by_assay: dict[str, list[str]], +) -> FeatureLookupBatch: + """Resolve all proposed individual features through one model tool call.""" + deps = ctx.deps + unknown_assays = sorted(set(queries_by_assay) - set(deps.assays)) + if unknown_assays: + raise ModelRetry(f"Unknown requested assays: {unknown_assays}") + if not queries_by_assay: + raise ModelRetry("queries_by_assay must contain at least one assay") + clean_queries_by_assay = { + assay_name: list( + dict.fromkeys(value.strip() for value in queries if value.strip()) + ) + for assay_name, queries in queries_by_assay.items() + } + empty_assays = sorted( + assay_name + for assay_name, queries in clean_queries_by_assay.items() + if not queries + ) + if empty_assays: + raise ModelRetry(f"Feature-query batches cannot be empty: {empty_assays}") + query_count = sum(len(queries) for queries in clean_queries_by_assay.values()) + if query_count > _MAX_FEATURE_QUERIES: + raise ModelRetry( + "The batch may contain at most " + f"{_MAX_FEATURE_QUERIES} feature queries in total" + ) + if deps.lookupBatch is not None: + if clean_queries_by_assay != deps.lookupQueries: + raise ModelRetry( + "Feature lookup already completed. Use only the returned lookup " + "evidence and do not request a different batch." + ) + logger.info("Data Enrichment reused the completed feature lookup batch") + return deps.lookupBatch + + logger.info( + "Data Enrichment feature lookup started: " + f"assays={len(clean_queries_by_assay)}, queries={query_count}" + ) + + start = len(deps.toolCalls) + lookups = [ + await find_present_features( + ctx, + assay_name=assay_name, + queries=clean_queries_by_assay[assay_name], + ) + for assay_name in deps.assays + if assay_name in clean_queries_by_assay + ] + del deps.toolCalls[start:] + evidence_ids = list( + dict.fromkeys( + evidence_id for lookup in lookups for evidence_id in lookup.evidenceIds + ) + ) + deps.toolCalls.append( + DataEnrichmentToolCall( + name="find_present_features_batch", + assay=",".join(queries_by_assay), + evidenceIds=evidence_ids, + ) + ) + result_counts = {"present": 0, "ambiguous": 0, "absent": 0} + for lookup in lookups: + for result in lookup.results: + result_counts[result.status] += 1 + logger.info( + "Data Enrichment feature lookup completed: " + f"present={result_counts['present']}, " + f"ambiguous={result_counts['ambiguous']}, " + f"absent={result_counts['absent']}, evidence={len(evidence_ids)}" + ) + batch = FeatureLookupBatch(lookups=lookups, evidenceIds=evidence_ids) + deps.lookupBatch = batch + deps.lookupQueries = clean_queries_by_assay + return batch diff --git a/scarf/agent/data_enrichment/validation.py b/scarf/agent/data_enrichment/validation.py new file mode 100644 index 00000000..ff791696 --- /dev/null +++ b/scarf/agent/data_enrichment/validation.py @@ -0,0 +1,317 @@ +"""Ground and validate data enrichment reports.""" + +import re + +from ...features.gene_reference import species_registry +from ...utils.logging import logger +from ..config.agent_exec import describe_agent_error +from ..types import AgentRunInfo +from .contracts import ( + DataEnrichmentContext, + DataEnrichmentDependencies, + DataEnrichmentReport, + FeatureReference, + FeatureSelectionPolicy, + StudyContextSummary, +) + +_SUPPORTED_SPECIES = species_registry() + + +def _ground_study_context_summary( + context: DataEnrichmentContext, + proposed: StudyContextSummary, +) -> StudyContextSummary: + """Bind structured context references to exact caller text.""" + original_context = context.studyContext + original_objective = context.studyObjective + grounded_text = f"{original_context}\n{original_objective}" + organism_references = [context.organismHint] if context.organismHint else [] + for species in _SUPPORTED_SPECIES.values(): + match = re.search( + rf"\b{re.escape(species.label)}\b", + grounded_text, + flags=re.IGNORECASE, + ) + if match is not None: + organism_references.append(match.group(0)) + field_sources = { + "organismReferences": organism_references, + "tissueReferences": list(context.tissueReferences), + "cellTypeReferences": list(context.cellTypeReferences), + "experimentalReferences": list(context.experimentalDetails), + "hypothesisReferences": [], + "analysisIntentReferences": [], + } + grounded: dict[str, list[str]] = {} + for field_name, supplied_values in field_sources.items(): + exact_supplied = [value.strip() for value in supplied_values if value.strip()] + proposed_values = list(getattr(proposed, field_name)) + combined = list( + dict.fromkeys( + value.strip() + for value in [*exact_supplied, *proposed_values] + if value.strip() + ) + ) + excerpts = [ + value for value in proposed_values if value.strip() not in exact_supplied + ] + if len(excerpts) > 12: + raise ValueError( + f"studyContextSummary.{field_name} may contain at most 12 proposed excerpts; " + f"received {len(excerpts)}. Keep objective-relevant verbatim spans; " + "the complete original context and exact caller references remain available." + ) + supplied = set(exact_supplied) + invalid = [ + value + for value in combined + if value not in supplied and value not in grounded_text + ] + if invalid: + raise ValueError( + f"Study-context references must be verbatim caller text: {invalid}" + ) + oversized = [ + value + for value in proposed_values + if value not in supplied and len(value) > 240 + ] + if oversized: + raise ValueError("Study-context references may not exceed 240 characters") + grounded[field_name] = combined + + evidence_ids: list[str] = [] + if original_context: + evidence_ids.append("context:study") + if original_objective: + evidence_ids.append("context:objective") + if context.organismHint: + evidence_ids.append("context:organism") + evidence_ids.extend( + f"context:tissue:{index}" + for index, _value in enumerate(context.tissueReferences) + ) + evidence_ids.extend( + f"context:cellType:{index}" + for index, _value in enumerate(context.cellTypeReferences) + ) + evidence_ids.extend( + f"context:experiment:{index}" + for index, _value in enumerate(context.experimentalDetails) + ) + return StudyContextSummary( + studyContext=original_context, + studyObjective=original_objective, + **grounded, + evidenceIds=evidence_ids, + ) + + +def _validate_feature_policy( + deps: DataEnrichmentDependencies, + policy: FeatureSelectionPolicy, + grounded_context: StudyContextSummary, +) -> None: + """Ground and validate one feature policy in deterministic order.""" + supported_species = {*_SUPPORTED_SPECIES, "unknown"} + if policy.species not in supported_species: + raise ValueError( + f"unsupported species {policy.species!r}; choose a supported key or unknown" + ) + if not policy.evidenceIds: + raise ValueError(f"policy for assay {policy.assay!r} requires evidence IDs") + policy.organismName = ( + _SUPPORTED_SPECIES[policy.species].label + if policy.species in _SUPPORTED_SPECIES + else "unknown" + ) + inspection = deps.inspections.get(policy.assay) + if inspection is None: + raise ValueError(f"assay {policy.assay!r} was not inspected") + modality = inspection.modalityEvidence + policy.assayType = modality.assayType + policy.assayModality = modality.modality + policy.graphEligible = modality.graphEligible + policy.markerEligible = modality.markerEligible + policy.demultiplexEligible = modality.demultiplexEligible + policy.exactControlFeatures = [ + FeatureReference( + featureId=item.featureId, + featureName=item.featureName, + ) + for item in modality.adtControls + ] + policy.exactTagFeatures = [ + FeatureReference( + featureId=item.featureId, + featureName=item.featureName, + ) + for item in modality.htoTags + ] + policy.peakCoordinateStatus = modality.atacCoordinates.status + policy.evidenceIds = list( + dict.fromkeys([*policy.evidenceIds, *modality.evidenceIds]) + ) + if ( + inspection.species in _SUPPORTED_SPECIES + and policy.species != inspection.species + ): + raise ValueError( + f"policy species {policy.species!r} conflicts with inspected " + f"species {inspection.species!r}" + ) + if ( + inspection.species == "unknown" + and policy.species != "unknown" + and not any( + evidence_id.startswith("context:") for evidence_id in policy.evidenceIds + ) + ): + raise ValueError( + "A context-derived species decision must cite context evidence" + ) + observed_families = {item.family for item in inspection.families} + cited_families = set(policy.excludeFamilies) | set(policy.protectFamilies) + unknown_families = cited_families - observed_families + if unknown_families: + raise ValueError( + f"policy cites unobserved families: {sorted(unknown_families)}" + ) + protected_defaults = { + item.family for item in inspection.families if item.defaultExclude is False + } + excluded_protected = sorted( + set(policy.excludeFamilies).intersection(protected_defaults) + ) + if excluded_protected: + raise ValueError( + "The initial enrichment policy cannot exclude families that " + f"deterministic evidence protects by default: {excluded_protected}" + ) + confirmed = deps.confirmedFeatures.get(policy.assay, set()) + cited_features = { + *policy.excludeFeatures, + *policy.protectFeatures, + *policy.artificialFeatures, + } + unknown_features = cited_features - confirmed + if unknown_features: + raise ValueError( + "Call find_present_features_batch before citing individual features: " + f"{sorted(unknown_features)}" + ) + exogenous_evidence = { + value: item.evidenceId + for item in inspection.exogenous + for value in (item.featureId, item.featureName) + } + unsupported_artificial: list[str] = [] + for feature in policy.artificialFeatures: + evidence_id = exogenous_evidence.get(feature) + if evidence_id is not None and evidence_id in policy.evidenceIds: + continue + matching_context_ids = { + f"context:experiment:{index}" + for index, detail in enumerate(deps.context.experimentalDetails) + if feature.casefold() in detail.casefold() + } + if matching_context_ids.intersection(policy.evidenceIds): + continue + unsupported_artificial.append(feature) + if unsupported_artificial: + raise ValueError( + "Artificial features require their exogenous evidence ID or a " + "feature-specific experimental-context evidence ID: " + f"{sorted(unsupported_artificial)}" + ) + unknown_evidence = set(policy.evidenceIds) - deps.evidenceIds + if unknown_evidence: + raise ValueError( + f"policy cites unknown evidence IDs: {sorted(unknown_evidence)}" + ) + policy.tissueReferences = list(grounded_context.tissueReferences) + policy.cellTypeReferences = list(grounded_context.cellTypeReferences) + policy.experimentalReferences = list(grounded_context.experimentalReferences) + + +def validate_data_enrichment_report( + deps: DataEnrichmentDependencies, + report: DataEnrichmentReport, +) -> DataEnrichmentReport: + """Ground an agent report in inspected assays, context, and exact lookups.""" + if not deps.inspections: + raise ValueError("Inspect every requested assay before returning the report") + + requested = set(deps.assays) + reported = {policy.assay for policy in report.policies} + if len(reported) != len(report.policies): + raise ValueError("reports may contain only one policy for each assay") + if not reported.issubset(requested): + raise ValueError( + f"policies cite assays outside the requested set: {sorted(reported - requested)}" + ) + if report.status == "done" and reported != requested: + raise ValueError( + f"done reports require one policy for every requested assay: {deps.assays}" + ) + + grounded_context = _ground_study_context_summary( + deps.context, + report.studyContextSummary, + ) + for policy in report.policies: + _validate_feature_policy(deps, policy, grounded_context) + + report.studyContextSummary = grounded_context + report.inspections = [deps.inspections[name] for name in deps.assays] + report.toolCalls = list(deps.toolCalls) + report.evidenceIds = list( + dict.fromkeys( + evidence_id + for evidence_id in [ + *report.studyContextSummary.evidenceIds, + *( + evidence_id + for policy in report.policies + for evidence_id in policy.evidenceIds + ), + ] + ) + ) + logger.debug( + "Data Enrichment report validated: " + f"status={report.status}, policies={len(report.policies)}, " + f"inspections={len(report.inspections)}, " + f"toolCalls={len(report.toolCalls)}, evidence={len(report.evidenceIds)}" + ) + return report + + +def failed_data_enrichment_report( + deps: DataEnrichmentDependencies, + *, + error: Exception, + model_name: str, +) -> DataEnrichmentReport: + """Retain observed inspection evidence without selecting a policy after failure.""" + context = _ground_study_context_summary(deps.context, StudyContextSummary()) + return DataEnrichmentReport( + status="failed", + inspections=[ + deps.inspections[name] for name in deps.assays if name in deps.inspections + ], + studyContextSummary=context, + toolCalls=list(deps.toolCalls), + evidenceIds=sorted({*deps.evidenceIds, *context.evidenceIds}), + limitations=[ + "No scientific feature policy was selected after model failure.", + describe_agent_error(error), + ], + runInfo=getattr( + error, + "agent_run_info", + AgentRunInfo(agentName="data_enrichment_failed", modelName=model_name), + ), + ) diff --git a/scarf/agent/decisions/__init__.py b/scarf/agent/decisions/__init__.py new file mode 100644 index 00000000..a3b9fa0c --- /dev/null +++ b/scarf/agent/decisions/__init__.py @@ -0,0 +1 @@ +"""Decision contracts and registered RNA decision policies.""" diff --git a/scarf/agent/decisions/kernel.py b/scarf/agent/decisions/kernel.py new file mode 100644 index 00000000..6c682ccc --- /dev/null +++ b/scarf/agent/decisions/kernel.py @@ -0,0 +1,672 @@ +"""Typed decision contracts and deterministic validation for agent workflows. + +The models in this module deliberately separate deterministic option construction +from agent output. An agent may select one offered option and cite observed +evidence, but it cannot add operations or numeric execution parameters. +""" + +import hashlib +import re +from collections.abc import Iterable +from typing import Literal + +from pydantic import ConfigDict, Field, field_validator, model_validator + +from .. import record_io +from ..types import AgentDataModel, ArtifactReferenceModel + +type DecisionStatus = Literal["apply", "skip", "defer", "abstain"] +type DecisionSource = Literal["rule", "agent", "human"] +type DecisionConfidence = Literal["low", "medium", "high", "notApplicable"] +type EvidenceClass = Literal[ + "geometric", + "markerCoherence", + "resamplingStability", + "crossUnitSupport", + "protectedVariablePreservation", + "qualityControl", + "technical", + "design", + "provenance", + "batchRemoval", + "biologicalConservation", + "other", +] +type VerificationStatus = Literal["passed", "failed", "inconclusive"] +type DecisionWorkflowStatus = Literal[ + "running", + "completed", + "needsInput", + "abstained", + "failed", +] +type ProtectedVariableEffectStatus = Literal[ + "preserved", + "degraded", + "improved", + "notEvaluated", +] + +_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.:/-]{0,255}$") +_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$") +_NON_GEOMETRIC_OVERRIDE_CLASSES: frozenset[EvidenceClass] = frozenset( + { + "markerCoherence", + "resamplingStability", + "crossUnitSupport", + "protectedVariablePreservation", + } +) + + +def _validate_identifier(value: str, field_name: str) -> str: + if _ID_PATTERN.fullmatch(value) is None: + raise ValueError( + f"{field_name} must be a non-empty stable identifier containing only " + "letters, digits, '.', '_', ':', '/', or '-'" + ) + return value + + +def _validate_unique(values: list[str], field_name: str) -> list[str]: + if len(values) != len(set(values)): + raise ValueError(f"{field_name} must not contain duplicates") + return values + + +def _default_decision_sources() -> list[DecisionSource]: + return ["rule", "agent", "human"] + + +class DecisionKernelModel(AgentDataModel): + """Base for immutable, closed decision-kernel contracts.""" + + model_config = ConfigDict(extra="forbid", frozen=True, validate_default=True) + + +class DecisionEvidence(DecisionKernelModel): + """One bounded observed fact available to a decision.""" + + evidenceId: str + evidenceClass: EvidenceClass + summary: str = Field(min_length=1, max_length=2000) + artifactReferences: list[ArtifactReferenceModel] = Field(default_factory=list) + + @field_validator("evidenceId") + @classmethod + def validate_evidence_id(cls, value: str) -> str: + return _validate_identifier(value, "evidenceId") + + @field_validator("summary") + @classmethod + def validate_summary(cls, value: str) -> str: + if value != value.strip(): + raise ValueError("summary must not contain surrounding whitespace") + return value + + @model_validator(mode="after") + def validate_artifact_references(self) -> "DecisionEvidence": + identities: list[tuple[str, str | None, str, str]] = [] + for reference in self.artifactReferences: + if not reference.kind or not reference.artifactId: + raise ValueError("artifactReferences require kind and artifactId") + identities.append( + ( + reference.scope, + reference.assay, + reference.kind, + reference.artifactId, + ) + ) + if len(identities) != len(set(identities)): + raise ValueError("artifactReferences must not contain duplicates") + return self + + +class EvidenceBundle(DecisionKernelModel): + """The complete immutable evidence inventory offered for one decision.""" + + formatVersion: Literal[1] = 1 + bundleId: str + decisionId: str + evidence: list[DecisionEvidence] = Field(default_factory=list) + contentSha256: str | None = None + + @field_validator("bundleId", "decisionId") + @classmethod + def validate_ids(cls, value: str, info: object) -> str: + field_name = getattr(info, "field_name", "identifier") + return _validate_identifier(value, field_name) + + @field_validator("contentSha256") + @classmethod + def validate_content_sha256(cls, value: str | None) -> str | None: + if value is not None and _SHA256_PATTERN.fullmatch(value) is None: + raise ValueError("contentSha256 must be a lowercase SHA-256 digest") + return value + + @model_validator(mode="after") + def validate_evidence_ids(self) -> "EvidenceBundle": + _validate_unique( + [item.evidenceId for item in self.evidence], + "EvidenceBundle.evidence IDs", + ) + if self.contentSha256 is not None: + payload = self.model_dump(mode="json", exclude={"contentSha256"}) + expected = hashlib.sha256( + record_io.canonical_json_bytes(payload) + ).hexdigest() + if self.contentSha256 != expected: + raise ValueError("contentSha256 does not match the evidence bundle") + return self + + def evidence_by_id(self) -> dict[str, DecisionEvidence]: + """Return evidence indexed by its stable identifier.""" + return {item.evidenceId: item for item in self.evidence} + + def with_content_sha256(self) -> "EvidenceBundle": + """Return this bundle with its canonical content identity.""" + payload = self.model_dump(mode="json", exclude={"contentSha256"}) + digest = hashlib.sha256(record_io.canonical_json_bytes(payload)).hexdigest() + return EvidenceBundle.model_validate({**payload, "contentSha256": digest}) + + +class DecisionOption(DecisionKernelModel): + """One pre-registered option whose execution details live outside the model.""" + + optionId: str + status: DecisionStatus + label: str = Field(min_length=1, max_length=200) + description: str = Field(min_length=1, max_length=2000) + requiredEvidenceClasses: list[EvidenceClass] = Field(default_factory=list) + requiredEvidenceIds: list[str] = Field(default_factory=list) + + @field_validator("optionId") + @classmethod + def validate_option_id(cls, value: str) -> str: + return _validate_identifier(value, "optionId") + + @field_validator("label", "description") + @classmethod + def validate_text(cls, value: str, info: object) -> str: + if value != value.strip(): + field_name = getattr(info, "field_name", "text") + raise ValueError(f"{field_name} must not contain surrounding whitespace") + return value + + @field_validator("requiredEvidenceClasses") + @classmethod + def validate_required_classes( + cls, value: list[EvidenceClass] + ) -> list[EvidenceClass]: + if len(value) != len(set(value)): + raise ValueError("requiredEvidenceClasses must not contain duplicates") + return value + + @field_validator("requiredEvidenceIds") + @classmethod + def validate_required_evidence_ids(cls, value: list[str]) -> list[str]: + for evidence_id in value: + _validate_identifier(evidence_id, "requiredEvidenceIds item") + return _validate_unique(value, "requiredEvidenceIds") + + +class DecisionSpec(DecisionKernelModel): + """Authoritative closed option set for one atomic decision.""" + + decisionId: str + definitionVersion: int = Field(ge=1, strict=True) + checkpoint: str + question: str = Field(min_length=1, max_length=2000) + evidenceBundleId: str + options: list[DecisionOption] = Field(min_length=1) + baselineOptionId: str | None = None + metricPreferredOptionId: str | None = None + requireIndependentOverrideEvidence: bool = Field(default=False, strict=True) + allowedSources: list[DecisionSource] = Field( + default_factory=_default_decision_sources + ) + + @field_validator("decisionId", "checkpoint", "evidenceBundleId") + @classmethod + def validate_ids(cls, value: str, info: object) -> str: + field_name = getattr(info, "field_name", "identifier") + return _validate_identifier(value, field_name) + + @field_validator("question") + @classmethod + def validate_question(cls, value: str) -> str: + if value != value.strip(): + raise ValueError("question must not contain surrounding whitespace") + return value + + @field_validator("allowedSources") + @classmethod + def validate_allowed_sources( + cls, value: list[DecisionSource] + ) -> list[DecisionSource]: + if not value: + raise ValueError("allowedSources must not be empty") + if len(value) != len(set(value)): + raise ValueError("allowedSources must not contain duplicates") + return value + + @model_validator(mode="after") + def validate_option_set(self) -> "DecisionSpec": + option_ids = [option.optionId for option in self.options] + _validate_unique(option_ids, "DecisionSpec option IDs") + if ( + self.baselineOptionId is not None + and self.baselineOptionId not in option_ids + ): + raise ValueError("baselineOptionId must reference an offered option") + if ( + self.metricPreferredOptionId is not None + and self.metricPreferredOptionId not in option_ids + ): + raise ValueError("metricPreferredOptionId must reference an offered option") + if ( + self.requireIndependentOverrideEvidence + and self.metricPreferredOptionId is None + ): + raise ValueError( + "requireIndependentOverrideEvidence requires metricPreferredOptionId" + ) + return self + + def option_by_id(self) -> dict[str, DecisionOption]: + """Return offered options indexed by their stable identifier.""" + return {option.optionId: option for option in self.options} + + +class ProtectedVariableEffect(DecisionKernelModel): + """Observed effect of a choice on one objective-protected variable.""" + + variable: str + status: ProtectedVariableEffectStatus + evidenceIds: list[str] = Field(default_factory=list) + summary: str = Field(min_length=1, max_length=1000) + + @field_validator("variable") + @classmethod + def validate_variable(cls, value: str) -> str: + return _validate_identifier(value, "variable") + + @field_validator("evidenceIds") + @classmethod + def validate_evidence_ids(cls, value: list[str]) -> list[str]: + for evidence_id in value: + _validate_identifier(evidence_id, "evidenceIds item") + return _validate_unique(value, "ProtectedVariableEffect.evidenceIds") + + @field_validator("summary") + @classmethod + def validate_summary(cls, value: str) -> str: + if value != value.strip(): + raise ValueError("summary must not contain surrounding whitespace") + return value + + +class DecisionSelection(DecisionKernelModel): + """The bounded choice an agent or human may return.""" + + selectedOptionId: str + evidenceIds: list[str] = Field(default_factory=list) + rationale: str = Field(min_length=1, max_length=4000) + confidence: DecisionConfidence = "notApplicable" + protectedVariableEffects: list[ProtectedVariableEffect] = Field( + default_factory=list + ) + overrideOfOptionId: str | None = None + overrideEvidenceIds: list[str] = Field(default_factory=list) + + @field_validator("selectedOptionId", "overrideOfOptionId") + @classmethod + def validate_ids(cls, value: str | None, info: object) -> str | None: + if value is None: + return None + field_name = getattr(info, "field_name", "identifier") + return _validate_identifier(value, field_name) + + @field_validator("evidenceIds", "overrideEvidenceIds") + @classmethod + def validate_id_lists(cls, value: list[str], info: object) -> list[str]: + field_name = getattr(info, "field_name", "identifiers") + for item in value: + _validate_identifier(item, f"{field_name} item") + return _validate_unique(value, field_name) + + @field_validator("rationale") + @classmethod + def validate_rationale(cls, value: str) -> str: + if value != value.strip(): + raise ValueError("rationale must not contain surrounding whitespace") + return value + + @model_validator(mode="after") + def validate_override(self) -> "DecisionSelection": + if not set(self.overrideEvidenceIds).issubset(self.evidenceIds): + raise ValueError("overrideEvidenceIds must be included in evidenceIds") + if self.overrideOfOptionId is None and self.overrideEvidenceIds: + raise ValueError("overrideEvidenceIds require overrideOfOptionId") + if self.overrideOfOptionId == self.selectedOptionId: + raise ValueError("overrideOfOptionId must differ from selectedOptionId") + return self + + +class DecisionRecord(DecisionKernelModel): + """One durable rule, agent, or human selection from an exact option set.""" + + recordId: str + decisionId: str + definitionVersion: int = Field(ge=1, strict=True) + evidenceBundleId: str + evidenceBundleSha256: str + offeredOptionIds: list[str] = Field(min_length=1) + availableEvidenceIds: list[str] = Field(default_factory=list) + selectedOptionId: str + status: DecisionStatus + source: DecisionSource + evidenceIds: list[str] = Field(default_factory=list) + rationale: str = Field(min_length=1, max_length=4000) + confidence: DecisionConfidence = "notApplicable" + protectedVariableEffects: list[ProtectedVariableEffect] = Field( + default_factory=list + ) + overrideOfOptionId: str | None = None + overrideEvidenceIds: list[str] = Field(default_factory=list) + promptSha256: str | None = None + modelName: str | None = None + softwareSha256: str | None = None + createdAtNs: int = Field(default=0, ge=0, strict=True) + + @field_validator( + "recordId", + "decisionId", + "evidenceBundleId", + "selectedOptionId", + "overrideOfOptionId", + ) + @classmethod + def validate_ids(cls, value: str | None, info: object) -> str | None: + if value is None: + return None + field_name = getattr(info, "field_name", "identifier") + return _validate_identifier(value, field_name) + + @field_validator( + "offeredOptionIds", + "availableEvidenceIds", + "evidenceIds", + "overrideEvidenceIds", + ) + @classmethod + def validate_id_lists(cls, value: list[str], info: object) -> list[str]: + field_name = getattr(info, "field_name", "identifiers") + for item in value: + _validate_identifier(item, f"{field_name} item") + return _validate_unique(value, field_name) + + @field_validator("rationale") + @classmethod + def validate_rationale(cls, value: str) -> str: + if value != value.strip(): + raise ValueError("rationale must not contain surrounding whitespace") + return value + + @field_validator( + "evidenceBundleSha256", + "promptSha256", + "softwareSha256", + ) + @classmethod + def validate_sha256(cls, value: str | None, info: object) -> str | None: + if value is not None and _SHA256_PATTERN.fullmatch(value) is None: + field_name = getattr(info, "field_name", "digest") + raise ValueError(f"{field_name} must be a lowercase SHA-256 digest") + return value + + @field_validator("modelName") + @classmethod + def validate_model_name(cls, value: str | None) -> str | None: + if value is not None and (not value.strip() or value != value.strip()): + raise ValueError( + "modelName must be non-empty without surrounding whitespace" + ) + return value + + @model_validator(mode="after") + def validate_references(self) -> "DecisionRecord": + offered = set(self.offeredOptionIds) + available = set(self.availableEvidenceIds) + used = set(self.evidenceIds) + override = set(self.overrideEvidenceIds) + if self.selectedOptionId not in offered: + raise ValueError("selectedOptionId must reference an offered option") + if not used.issubset(available): + raise ValueError("evidenceIds must reference only available evidence") + if not override.issubset(used): + raise ValueError("overrideEvidenceIds must be included in evidenceIds") + protected_ids = { + evidence_id + for effect in self.protectedVariableEffects + for evidence_id in effect.evidenceIds + } + if not protected_ids.issubset(available): + raise ValueError( + "protectedVariableEffects must reference only available evidence" + ) + if self.overrideOfOptionId is None and self.overrideEvidenceIds: + raise ValueError("overrideEvidenceIds require overrideOfOptionId") + if self.overrideOfOptionId is not None: + if self.overrideOfOptionId not in offered: + raise ValueError("overrideOfOptionId must reference an offered option") + if self.overrideOfOptionId == self.selectedOptionId: + raise ValueError("overrideOfOptionId must differ from selectedOptionId") + return self + + +class VerificationCheck(DecisionKernelModel): + """One deterministic invariant evaluated by the auditor.""" + + checkId: str + status: VerificationStatus + summary: str = Field(min_length=1, max_length=2000) + evidenceIds: list[str] = Field(default_factory=list) + + @field_validator("checkId") + @classmethod + def validate_check_id(cls, value: str) -> str: + return _validate_identifier(value, "checkId") + + @field_validator("evidenceIds") + @classmethod + def validate_evidence_ids(cls, value: list[str]) -> list[str]: + for evidence_id in value: + _validate_identifier(evidence_id, "evidenceIds item") + return _validate_unique(value, "VerificationCheck.evidenceIds") + + @field_validator("summary") + @classmethod + def validate_summary(cls, value: str) -> str: + if value != value.strip(): + raise ValueError("summary must not contain surrounding whitespace") + return value + + +class DeterministicDecisionAuditor: + """Cross-check a decision against its authoritative spec and evidence bundle.""" + + @classmethod + def audit( + cls, + spec: DecisionSpec, + evidence: EvidenceBundle, + record: DecisionRecord, + *, + created_at_ns: int = 0, + ) -> list[VerificationCheck]: + """Return a deterministic verification without repairing invalid output.""" + checks: list[VerificationCheck] = [] + evidence_sha256 = ( + evidence.contentSha256 + if evidence.contentSha256 is not None + else evidence.with_content_sha256().contentSha256 + ) + + def add_check( + check_id: str, + passed: bool, + pass_summary: str, + fail_summary: str, + evidence_ids: Iterable[str] = (), + ) -> None: + checks.append( + VerificationCheck( + checkId=check_id, + status="passed" if passed else "failed", + summary=pass_summary if passed else fail_summary, + evidenceIds=list(evidence_ids), + ) + ) + + add_check( + "decisionIdentity", + record.decisionId == spec.decisionId + and record.definitionVersion == spec.definitionVersion + and evidence.decisionId == spec.decisionId + and record.evidenceBundleId == spec.evidenceBundleId + and evidence.bundleId == spec.evidenceBundleId + and record.evidenceBundleSha256 == evidence_sha256, + "Decision, definition, and evidence bundle identities agree.", + "Decision, definition, or evidence bundle identity does not agree.", + ) + + offered_option_ids = [option.optionId for option in spec.options] + add_check( + "exactOptionInventory", + record.offeredOptionIds == offered_option_ids, + "The durable record contains the exact offered option inventory.", + "The durable record does not contain the exact offered option inventory.", + ) + + available_evidence_ids = [item.evidenceId for item in evidence.evidence] + add_check( + "exactEvidenceInventory", + record.availableEvidenceIds == available_evidence_ids, + "The durable record contains the exact evidence inventory.", + "The durable record does not contain the exact evidence inventory.", + ) + + option = spec.option_by_id().get(record.selectedOptionId) + add_check( + "selectedOption", + option is not None and option.status == record.status, + "The selected option and decision status agree.", + "The selected option is unavailable or its status does not agree.", + ) + + add_check( + "decisionSource", + record.source in spec.allowedSources, + "The decision source is allowed by the definition.", + "The decision source is not allowed by the definition.", + ) + + evidence_by_id = evidence.evidence_by_id() + cited_ids = set(record.evidenceIds) + required_classes = set(option.requiredEvidenceClasses) if option else set() + required_ids = set(option.requiredEvidenceIds) if option else set() + cited_classes = { + evidence_by_id[evidence_id].evidenceClass + for evidence_id in cited_ids + if evidence_id in evidence_by_id + } + required_evidence_ok = ( + cited_ids.issubset(evidence_by_id) + and required_ids.issubset(cited_ids) + and required_classes.issubset(cited_classes) + ) + add_check( + "requiredEvidence", + required_evidence_ok, + "All cited and required evidence is present.", + "Cited evidence is unavailable or a required evidence class is missing.", + record.evidenceIds, + ) + + protected_evidence_ids = { + evidence_id + for effect in record.protectedVariableEffects + for evidence_id in effect.evidenceIds + } + protected_ok = not any( + effect.status == "degraded" for effect in record.protectedVariableEffects + ) and protected_evidence_ids.issubset(cited_ids) + add_check( + "protectedVariablePreservation", + protected_ok, + "No cited protected variable is degraded.", + "A protected variable degraded or its evidence was not cited.", + sorted(protected_evidence_ids), + ) + + metric_override = ( + spec.requireIndependentOverrideEvidence + and spec.metricPreferredOptionId is not None + and record.status in {"apply", "skip"} + and record.selectedOptionId != spec.metricPreferredOptionId + ) + if metric_override: + override_classes = { + evidence_by_id[evidence_id].evidenceClass + for evidence_id in record.overrideEvidenceIds + if evidence_id in evidence_by_id + } + qualifying_classes = override_classes & _NON_GEOMETRIC_OVERRIDE_CLASSES + override_ok = ( + record.overrideOfOptionId == spec.metricPreferredOptionId + and set(record.overrideEvidenceIds).issubset(record.evidenceIds) + and ( + not required_ids + or set(record.overrideEvidenceIds).issubset(required_ids) + ) + and len(qualifying_classes) >= 2 + ) + add_check( + "independentOverrideEvidence", + override_ok, + "The override cites at least two independent non-geometric evidence classes.", + "The override requires two independent non-geometric evidence classes.", + record.overrideEvidenceIds, + ) + else: + add_check( + "independentOverrideEvidence", + record.overrideOfOptionId is None and not record.overrideEvidenceIds, + "No metric override evidence is required.", + "Override fields were supplied without an eligible metric override.", + record.overrideEvidenceIds, + ) + + return checks + + +__all__ = [ + "DecisionConfidence", + "DecisionEvidence", + "DecisionOption", + "DecisionRecord", + "DecisionSelection", + "DecisionSource", + "DecisionSpec", + "DecisionStatus", + "DecisionWorkflowStatus", + "DeterministicDecisionAuditor", + "EvidenceBundle", + "EvidenceClass", + "ProtectedVariableEffect", + "ProtectedVariableEffectStatus", + "VerificationCheck", + "VerificationStatus", +] diff --git a/scarf/agent/decisions/rna.py b/scarf/agent/decisions/rna.py new file mode 100644 index 00000000..8866b7ea --- /dev/null +++ b/scarf/agent/decisions/rna.py @@ -0,0 +1,497 @@ +"""Deterministic RNA decision definitions and executor-owned option payloads.""" + +from collections.abc import Mapping, Sequence +from typing import Annotated, Literal + +from pydantic import ConfigDict, Field, field_validator, model_validator + +from ..cell_quality.profiles import CellQualityProfile +from ..types import AgentDataModel +from .kernel import ( + DecisionOption, + DecisionRecord, + DecisionSpec, + DecisionStatus, + DeterministicDecisionAuditor, + EvidenceBundle, + VerificationCheck, +) + +type RnaDecisionCheckpoint = Literal["qcGrouping", "cellQuality"] +type QcGroupingMode = Literal["global", "physicalCapture", "pooledReference"] + + +class RnaDecisionGateError(ValueError): + """Raised when deterministic evidence forbids constructing an option set.""" + + +class RnaDecisionCompilationError(ValueError): + """Raised when an unverified decision cannot compile to an executor payload.""" + + +class RnaRegistryModel(AgentDataModel): + """Base for immutable, closed RNA registry contracts.""" + + model_config = ConfigDict(extra="forbid", frozen=True, validate_default=True) + + +class CellQualityExecutorPayload(RnaRegistryModel): + """Exact cell-quality thresholds owned by the deterministic executor.""" + + operation: Literal["cellQualityProfile"] = "cellQualityProfile" + profile: CellQualityProfile + lowerCountMad: float | None = Field(default=None, ge=0, strict=True) + lowerFeatureMad: float | None = Field(default=None, ge=0, strict=True) + upperMitoMad: float | None = Field(default=None, ge=0, strict=True) + groupByCapture: bool = Field(strict=True) + pooledReference: bool = Field(strict=True) + sensitivityOnly: bool = Field(strict=True) + flagHighCounts: Literal[True] = True + flagHighFeatures: Literal[True] = True + + @model_validator(mode="after") + def validate_profile(self) -> "CellQualityExecutorPayload": + thresholds = (self.lowerCountMad, self.lowerFeatureMad, self.upperMitoMad) + if self.profile == "retainWithFlags": + if any(value is not None for value in thresholds): + raise ValueError("retainWithFlags cannot define removal thresholds") + if self.groupByCapture or self.pooledReference or self.sensitivityOnly: + raise ValueError("retainWithFlags cannot enable filtering modes") + return self + if self.profile in {"coreGlobalGaussian", "coreSampleMad3"}: + if any(value is not None for value in thresholds): + raise ValueError( + "Core profiles use exact core bounds, not one-sided MAD fields" + ) + if self.groupByCapture != (self.profile == "coreSampleMad3"): + raise ValueError("The core sample profile requires capture grouping") + if self.pooledReference or self.sensitivityOnly: + raise ValueError( + "Core profiles cannot enable registered filtering modes" + ) + return self + if any(value is None for value in thresholds): + raise ValueError("Filtering profiles require all three MAD thresholds") + if self.profile.startswith("captureMad") != self.groupByCapture: + raise ValueError("Capture profiles and groupByCapture must agree") + if (self.profile == "pooledReferenceMad5") != self.pooledReference: + raise ValueError("pooledReferenceMad5 and pooledReference must agree") + if (self.profile == "captureMad3Sensitivity") != self.sensitivityOnly: + raise ValueError("captureMad3Sensitivity and sensitivityOnly must agree") + return self + + +class QcGroupingExecutorPayload(RnaRegistryModel): + """Exact population used to estimate registered cell-quality thresholds.""" + + operation: Literal["qcGrouping"] = "qcGrouping" + groupingMode: QcGroupingMode + + +class NoExecutionPayload(RnaRegistryModel): + """Typed terminal or pause outcome with no analytical operation.""" + + operation: Literal["noExecution"] = "noExecution" + reasonCode: Literal[ + "needsInput", + "scientificAbstention", + ] + + +type RnaOptionPayload = Annotated[ + QcGroupingExecutorPayload | CellQualityExecutorPayload | NoExecutionPayload, + Field(discriminator="operation"), +] + + +class RnaExecutorOption(RnaRegistryModel): + """Executor-only payload keyed by the semantic option shown to an agent.""" + + checkpoint: RnaDecisionCheckpoint + optionId: str + payload: RnaOptionPayload + + @field_validator("optionId") + @classmethod + def validate_option_id(cls, value: str) -> str: + if not value or value != value.strip(): + raise ValueError( + "optionId must be non-empty without surrounding whitespace" + ) + return value + + +class RnaDecisionDefinition(RnaRegistryModel): + """Agent-visible decision spec paired with executor-only payloads.""" + + checkpoint: RnaDecisionCheckpoint + spec: DecisionSpec + executorOptions: list[RnaExecutorOption] = Field(min_length=1) + + @model_validator(mode="after") + def validate_exact_registry(self) -> "RnaDecisionDefinition": + if self.spec.checkpoint != self.checkpoint: + raise ValueError( + "DecisionSpec checkpoint must match the registry checkpoint" + ) + visible_ids = [option.optionId for option in self.spec.options] + executor_ids = [option.optionId for option in self.executorOptions] + if len(executor_ids) != len(set(executor_ids)): + raise ValueError("executorOptions must not contain duplicate option IDs") + if executor_ids != visible_ids: + raise ValueError( + "executorOptions must exactly match the ordered visible option IDs" + ) + if any(option.checkpoint != self.checkpoint for option in self.executorOptions): + raise ValueError("Every executor option must match the registry checkpoint") + return self + + def executor_option(self, option_id: str) -> RnaExecutorOption: + """Return one exact registered executor option.""" + for option in self.executorOptions: + if option.optionId == option_id: + return option + raise KeyError(f"Unknown option ID for {self.spec.decisionId}: {option_id}") + + +class CompiledRnaDecision(RnaRegistryModel): + """Verified executor handoff kept separate from agent-authored records.""" + + decisionRecordId: str + decisionId: str + selectedOptionId: str + status: DecisionStatus + executorPayload: RnaOptionPayload + checks: list[VerificationCheck] + + +def compile_rna_decision( + definition: RnaDecisionDefinition, + evidence: EvidenceBundle, + record: DecisionRecord, + *, + created_at_ns: int = 0, +) -> CompiledRnaDecision: + """Audit an exact selection and resolve its executor-owned payload.""" + verification = DeterministicDecisionAuditor.audit( + definition.spec, + evidence, + record, + created_at_ns=created_at_ns, + ) + failed_checks = [ + check.checkId for check in verification if check.status == "failed" + ] + if failed_checks: + raise RnaDecisionCompilationError( + "Decision failed deterministic verification: " + ", ".join(failed_checks) + ) + executor_option = definition.executor_option(record.selectedOptionId) + return CompiledRnaDecision( + decisionRecordId=record.recordId, + decisionId=record.decisionId, + selectedOptionId=record.selectedOptionId, + status=record.status, + executorPayload=executor_option.payload, + checks=verification, + ) + + +def _defer_option( + checkpoint: RnaDecisionCheckpoint, option_id: str +) -> tuple[DecisionOption, RnaExecutorOption]: + return ( + DecisionOption( + optionId=option_id, + status="defer", + label="Request input", + description="Pause because the available evidence cannot identify a choice.", + ), + RnaExecutorOption( + checkpoint=checkpoint, + optionId=option_id, + payload=NoExecutionPayload(reasonCode="needsInput"), + ), + ) + + +def _build_definition( + *, + checkpoint: RnaDecisionCheckpoint, + decision_id: str, + evidence_bundle_id: str, + question: str, + visible_options: list[DecisionOption], + executor_options: list[RnaExecutorOption], + baseline_option_id: str | None, +) -> RnaDecisionDefinition: + return RnaDecisionDefinition( + checkpoint=checkpoint, + spec=DecisionSpec( + decisionId=decision_id, + definitionVersion=1, + checkpoint=checkpoint, + question=question, + evidenceBundleId=evidence_bundle_id, + options=visible_options, + baselineOptionId=baseline_option_id, + ), + executorOptions=executor_options, + ) + + +def require_option_evidence( + definition: RnaDecisionDefinition, + requirements: Mapping[str, Sequence[str]], +) -> RnaDecisionDefinition: + """Bind exact evidence IDs to the registered options they support.""" + options = definition.spec.option_by_id() + unknown = sorted(set(requirements) - set(options)) + if unknown: + raise ValueError(f"Evidence requirements name unknown options: {unknown}") + updated: list[DecisionOption] = [] + for option in definition.spec.options: + raw_ids = requirements.get(option.optionId, ()) + if isinstance(raw_ids, str | bytes): + raise TypeError("Option evidence requirements must be sequences of IDs") + updated.append( + option.model_copy( + update={"requiredEvidenceIds": list(dict.fromkeys(raw_ids))} + ) + ) + spec = definition.spec.model_copy(update={"options": updated}) + return definition.model_copy(update={"spec": spec}) + + +def build_qc_grouping_decision( + *, + evidence_bundle_id: str, + physical_capture_eligible: bool, + pooled_reference_eligible: bool, +) -> RnaDecisionDefinition: + """Build only quality-threshold grouping modes licensed by design evidence.""" + if pooled_reference_eligible and not physical_capture_eligible: + raise RnaDecisionGateError( + "Pooled-reference QC requires eligible physical captures" + ) + rows: list[tuple[str, str, str, QcGroupingMode]] = [ + ( + "qcGrouping:global", + "Global reference", + "Estimate registered quality boundaries across the selected dataset.", + "global", + ) + ] + if physical_capture_eligible: + rows.append( + ( + "qcGrouping:physicalCapture", + "Physical-capture references", + "Estimate registered boundaries separately within proven captures.", + "physicalCapture", + ) + ) + if pooled_reference_eligible: + rows.append( + ( + "qcGrouping:pooledReference", + "Comparable pooled reference", + "Estimate registered boundaries from the exact comparable captures.", + "pooledReference", + ) + ) + visible = [ + DecisionOption( + optionId=option_id, + status="apply", + label=label, + description=description, + requiredEvidenceClasses=["qualityControl", "design"], + ) + for option_id, label, description, _mode in rows + ] + executor = [ + RnaExecutorOption( + checkpoint="qcGrouping", + optionId=option_id, + payload=QcGroupingExecutorPayload(groupingMode=mode), + ) + for option_id, _label, _description, mode in rows + ] + defer_visible, defer_executor = _defer_option("qcGrouping", "qcGrouping:defer") + visible.append(defer_visible) + executor.append(defer_executor) + return _build_definition( + checkpoint="qcGrouping", + decision_id="qcGrouping", + evidence_bundle_id=evidence_bundle_id, + question="Which reference population should define cell-quality boundaries?", + visible_options=visible, + executor_options=executor, + baseline_option_id="qcGrouping:global", + ) + + +def build_cell_quality_decision( + *, + evidence_bundle_id: str, + available_profiles: Sequence[CellQualityProfile], +) -> RnaDecisionDefinition: + """Build exactly the cell-quality profiles licensed by grouping evidence.""" + if len(available_profiles) != len(set(available_profiles)): + raise RnaDecisionGateError("Cell-quality profiles must not contain duplicates") + if not available_profiles: + raise RnaDecisionGateError("At least one cell-quality profile is required") + definitions: dict[ + CellQualityProfile, + tuple[str, DecisionStatus, str, str, CellQualityExecutorPayload], + ] = { + "retainWithFlags": ( + "cellQuality:retainWithFlags", + "skip", + "Retain with flags", + "Preserve the published cell set and retain diagnostic quality flags.", + CellQualityExecutorPayload( + profile="retainWithFlags", + groupByCapture=False, + pooledReference=False, + sensitivityOnly=False, + ), + ), + "coreGlobalGaussian": ( + "cellQuality:coreGlobalGaussian", + "apply", + "Scarf default global filter", + "Use Scarf's default global Gaussian quantiles (0.01 and 0.99).", + CellQualityExecutorPayload( + profile="coreGlobalGaussian", + groupByCapture=False, + pooledReference=False, + sensitivityOnly=False, + ), + ), + "coreSampleMad3": ( + "cellQuality:coreSampleMad3", + "apply", + "Scarf sample-aware default", + "Use Scarf's exact sample-aware filter with three scaled MADs within proven physical captures.", + CellQualityExecutorPayload( + profile="coreSampleMad3", + groupByCapture=True, + pooledReference=False, + sensitivityOnly=False, + ), + ), + "globalMad5": ( + "cellQuality:globalMad5", + "apply", + "Global lenient filter", + "Apply one-sided global five-MAD quality boundaries.", + CellQualityExecutorPayload( + profile="globalMad5", + lowerCountMad=5.0, + lowerFeatureMad=5.0, + upperMitoMad=5.0, + groupByCapture=False, + pooledReference=False, + sensitivityOnly=False, + ), + ), + "captureMad5": ( + "cellQuality:captureMad5", + "apply", + "Capture-aware lenient filter", + "Apply one-sided five-MAD boundaries within physical captures.", + CellQualityExecutorPayload( + profile="captureMad5", + lowerCountMad=5.0, + lowerFeatureMad=5.0, + upperMitoMad=5.0, + groupByCapture=True, + pooledReference=False, + sensitivityOnly=False, + ), + ), + "captureMad3Sensitivity": ( + "cellQuality:captureMad3Sensitivity", + "apply", + "Capture sensitivity branch", + "Evaluate stricter three-MAD capture boundaries as sensitivity evidence.", + CellQualityExecutorPayload( + profile="captureMad3Sensitivity", + lowerCountMad=3.0, + lowerFeatureMad=3.0, + upperMitoMad=3.0, + groupByCapture=True, + pooledReference=False, + sensitivityOnly=True, + ), + ), + "pooledReferenceMad5": ( + "cellQuality:pooledReferenceMad5", + "apply", + "Pooled-reference filter", + "Apply five-MAD boundaries from comparable pooled reference captures.", + CellQualityExecutorPayload( + profile="pooledReferenceMad5", + lowerCountMad=5.0, + lowerFeatureMad=5.0, + upperMitoMad=5.0, + groupByCapture=False, + pooledReference=True, + sensitivityOnly=False, + ), + ), + } + profiles = [definitions[profile] for profile in available_profiles] + + visible = [ + DecisionOption( + optionId=option_id, + status=status, + label=label, + description=description, + requiredEvidenceClasses=["qualityControl"], + ) + for option_id, status, label, description, _payload in profiles + ] + executor = [ + RnaExecutorOption(checkpoint="cellQuality", optionId=option_id, payload=payload) + for option_id, _status, _label, _description, payload in profiles + ] + defer_visible, defer_executor = _defer_option("cellQuality", "cellQuality:defer") + visible.append(defer_visible) + executor.append(defer_executor) + return _build_definition( + checkpoint="cellQuality", + decision_id="cellQuality", + evidence_bundle_id=evidence_bundle_id, + question="Which cell-quality policy preserves valid biology?", + visible_options=visible, + executor_options=executor, + baseline_option_id=( + "cellQuality:coreGlobalGaussian" + if "coreGlobalGaussian" in available_profiles + else "cellQuality:retainWithFlags" + if "retainWithFlags" in available_profiles + else profiles[0][0] + ), + ) + + +__all__ = [ + "CellQualityExecutorPayload", + "QcGroupingExecutorPayload", + "CompiledRnaDecision", + "NoExecutionPayload", + "RnaDecisionCheckpoint", + "RnaDecisionCompilationError", + "RnaDecisionDefinition", + "RnaDecisionGateError", + "RnaExecutorOption", + "RnaOptionPayload", + "build_cell_quality_decision", + "build_qc_grouping_decision", + "compile_rna_decision", + "require_option_evidence", +] diff --git a/scarf/agent/decide.py b/scarf/agent/decisions/selection.py similarity index 66% rename from scarf/agent/decide.py rename to scarf/agent/decisions/selection.py index fe4e7d97..067f386e 100644 --- a/scarf/agent/decide.py +++ b/scarf/agent/decisions/selection.py @@ -4,8 +4,8 @@ from textwrap import dedent from typing import Any -from .config.agent_exec import run_agent_sync -from .types import Decision, EvidenceItem +from ..config.agent_exec import run_agent_sync +from ..types import Decision, EvidenceItem _SYSTEM_PROMPT = dedent( """ @@ -22,25 +22,6 @@ class DecisionValidationError(ValueError): """Raised when a model decision cites unknown or invalid evidence.""" -def _coerce_evidence_id(evidence_id: str, allowed: set[str]) -> str: - """Map a model-emitted id onto an allowed evidence id when unambiguous. - - Live models often echo prompt scaffolding such as ``id=domain:biological`` - instead of the bare id. Accept that when exactly one allowed id is embedded. - """ - if evidence_id in allowed: - return evidence_id - stripped = evidence_id.strip() - if stripped.startswith("id="): - stripped = stripped[3:].strip() - if stripped in allowed: - return stripped - matches = [allowed_id for allowed_id in allowed if allowed_id in evidence_id] - if len(matches) == 1: - return matches[0] - return evidence_id - - def validate_decision( decision: Decision, evidence: Sequence[EvidenceItem], @@ -48,19 +29,6 @@ def validate_decision( allowed = {item.id for item in evidence} if not allowed: raise DecisionValidationError("evidence must contain at least one item") - selected_id = _coerce_evidence_id(decision.selectedId, allowed) - evidence_ids = [ - _coerce_evidence_id(evidence_id, allowed) - for evidence_id in decision.evidenceIds - ] - if selected_id not in evidence_ids and selected_id in allowed: - evidence_ids = [selected_id, *evidence_ids] - if selected_id != decision.selectedId or evidence_ids != list(decision.evidenceIds): - decision = Decision( - selectedId=selected_id, - rationale=decision.rationale, - evidenceIds=evidence_ids, - ) if decision.selectedId not in allowed: raise DecisionValidationError( f"selectedId {decision.selectedId!r} is not in evidence ids {sorted(allowed)}" @@ -130,11 +98,27 @@ def decide( f"evidence ids must be unique; duplicates: {sorted(duplicates)}" ) - execution = run_agent_sync( - model=model, - output_type=Decision, - system_prompt=system_prompt, - user_prompt=_format_user_prompt(question, evidence), - name="decision", - ) + from pydantic_ai import UnexpectedModelBehavior + + try: + execution = run_agent_sync( + model=model, + output_type=Decision, + system_prompt=system_prompt, + user_prompt=_format_user_prompt(question, evidence), + name="decision", + output_validator=lambda value: validate_decision(value, evidence), + ) + except UnexpectedModelBehavior as exc: + cause: BaseException | None = exc.__cause__ + seen_causes: set[int] = set() + while cause is not None and id(cause) not in seen_causes: + seen_causes.add(id(cause)) + if isinstance(cause, DecisionValidationError): + invalid = DecisionValidationError(str(cause)) + if hasattr(exc, "agent_run_info"): + setattr(invalid, "agent_run_info", exc.agent_run_info) + raise invalid from exc + cause = cause.__cause__ + raise return validate_decision(execution.output, evidence) diff --git a/scarf/agent/experimental_context.py b/scarf/agent/experimental_context.py deleted file mode 100644 index 6494fe99..00000000 --- a/scarf/agent/experimental_context.py +++ /dev/null @@ -1,2303 +0,0 @@ -"""Tool-driven experimental-design and batch-correction assessment.""" - -import json -import math -from collections.abc import Mapping, Sequence -from textwrap import dedent -from typing import TYPE_CHECKING, Any, Literal - -import numpy as np - -from ..graph.feature_projection import graph_cell_selection -from ..metadata.queries import reduce_observation_units -from ..metadata.selection import resolve_cell_aligned_artifact -from ..metrics.association import coefficient_estimability -from ..quality_control.filtering import ( - _sample_aware_mad_mask, - gaussian_quantile_bounds, -) -from ..storage.refs import ArtifactRef -from ..storage.selections import read_stored_selection_mask -from ..utils.logging import logger -from .characterize_covariates import ( - CovariateCharacterization, - _SelectionBoundCells, - characterize_covariates, -) -from .config import AgentRunConfig -from .config._deps import AGENT_INSTALL_HINT -from .config.agent_exec import run_agent_sync -from .tools import artifact_reference, core_artifact_reference -from .types import ( - AgentDataModel, - AgentRunInfo, - ArtifactReferenceModel, - BatchCorrectionAction, - BatchSafetyEvidence, - BatchSafetyStatus, - ExperimentalBiologyHandoff, - ExperimentalTuningHandoff, - StageStatus, -) - -if TYPE_CHECKING: - from ..datastore.pipeline_run import PipelineRun - -try: - from pydantic import ConfigDict, Field, model_validator - from pydantic_ai import ModelRetry, RunContext, Tool - from pydantic_ai.tools import ToolDefinition -except ImportError as exc: - raise ImportError(AGENT_INSTALL_HINT) from exc - -__all__ = [ - "BatchCorrectionPlan", - "BatchSafetyEvidence", - "CellQcPlan", - "CellQcProfileEvidence", - "CovariateEvidence", - "ExperimentalContextAgent", - "ExperimentalContextDecision", - "ExperimentalContextDependencies", - "ExperimentalContextResult", - "InferenceUnit", - "NamedArtifactSource", - "RepresentationEvaluation", - "analyze_experimental_design", - "inspect_cell_covariates", - "score_current_representation", - "validate_experimental_context", -] - -type ColumnDomain = Literal["biological", "technical", "design", "ignore", "unknown"] -type IntegrationMetric = Literal[ - "iLISI", - "cLISI", - "graphConnectivity", - "proportionalBatchMixing", -] -type CellQcAction = Literal["skip", "globalGaussian", "sampleMad"] -type CellQcDriverType = Literal["RNA", "ATAC"] - -_CONTEXT_LIMIT = 1200 -_MAX_QC_SAMPLE_PROFILES = 4 -_MAX_SAMPLE_RETENTION_ITEMS = 20 - - -class InferenceUnit(AgentDataModel): - """Observation and independent units for one biological coefficient.""" - - observationUnit: str | None = None - independentUnit: str | None = None - - @classmethod - def get_blank(cls) -> "InferenceUnit": - return cls() - - @classmethod - def get_example(cls) -> "InferenceUnit": - return cls(observationUnit="sample", independentUnit="donor") - - -class BatchCorrectionPlan(AgentDataModel): - """A grounded recommendation about whether Harmony should be evaluated.""" - - action: BatchCorrectionAction - batchColumns: list[str] = Field(default_factory=list) - preserveColumns: list[str] = Field(default_factory=list) - metricsRequired: list[IntegrationMetric] = Field(default_factory=list) - rationale: str = "" - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "BatchCorrectionPlan": - return cls(action="needsInput") - - @classmethod - def get_example(cls) -> "BatchCorrectionPlan": - return cls( - action="evaluateHarmony", - batchColumns=["batch"], - preserveColumns=["cell_type", "treatment"], - metricsRequired=[ - "iLISI", - "cLISI", - "graphConnectivity", - ], - rationale=( - "Batch is technical and crossed with treatment, so compare an exact " - "Harmony candidate while protecting biological labels." - ), - evidenceIds=[ - "column:batch", - "estimability:treatment", - "batchEstimability:treatment:batch", - ], - ) - - -class NamedArtifactSource(AgentDataModel): - """One semantic name bound to an exact immutable artifact.""" - - name: str = "" - artifact: ArtifactReferenceModel = Field(default_factory=ArtifactReferenceModel) - - @model_validator(mode="after") - def validate_source(self) -> "NamedArtifactSource": - if self.name != self.name.strip(): - raise ValueError("Artifact source names cannot have surrounding whitespace") - if bool(self.name.strip()) != bool(self.artifact.artifactId): - raise ValueError("A named artifact source requires both name and artifact") - return self - - @classmethod - def get_blank(cls) -> "NamedArtifactSource": - return cls() - - @classmethod - def get_example(cls) -> "NamedArtifactSource": - return cls( - name="RNA_percentMito", - artifact=ArtifactReferenceModel( - assay="RNA", - kind="quality_metric", - artifactId="1" * 64, - ), - ) - - -def _validate_qc_sources( - *, - action: CellQcAction, - attributes: list[str], - artifact_metrics: list[NamedArtifactSource], - sample_column: str | None, - sample_artifact: NamedArtifactSource | None, -) -> None: - if len(attributes) != len(set(attributes)): - raise ValueError("Cell-QC metadata attributes must be unique") - if any( - not attribute.strip() or attribute != attribute.strip() - for attribute in attributes - ): - raise ValueError( - "Cell-QC metadata attributes cannot be blank or have surrounding whitespace" - ) - artifact_names = [source.name for source in artifact_metrics] - if len(artifact_names) != len(set(artifact_names)): - raise ValueError("Cell-QC artifact metric names must be unique") - if any(source.artifact.kind != "quality_metric" for source in artifact_metrics): - raise ValueError( - "Cell-QC artifactMetrics must reference quality_metric artifacts" - ) - collisions = sorted(set(attributes).intersection(artifact_names)) - if collisions: - raise ValueError( - f"Cell-QC metadata and artifact metric names collide: {collisions}" - ) - if sample_column is not None and sample_artifact is not None: - raise ValueError( - "Cell-QC sampleColumn and sampleArtifact are mutually exclusive" - ) - if sample_column is not None and ( - not sample_column.strip() or sample_column != sample_column.strip() - ): - raise ValueError( - "Cell-QC sampleColumn cannot be blank or have surrounding whitespace" - ) - if sample_artifact is not None and sample_artifact.artifact.kind != "hto_identity": - raise ValueError( - "Cell-QC sampleArtifact must reference an hto_identity artifact" - ) - if sample_artifact is not None and sample_artifact.name in artifact_names: - raise ValueError("Cell-QC sample and metric artifact names must be distinct") - if action == "skip" and (attributes or artifact_metrics): - raise ValueError("skip cannot include Cell-QC metrics") - if action != "skip" and not attributes and not artifact_metrics: - raise ValueError("Cell-QC filtering requires at least one metric") - if action == "sampleMad" and (sample_column is None) == (sample_artifact is None): - raise ValueError( - "sampleMad requires exactly one sampleColumn or sampleArtifact" - ) - if action != "sampleMad" and ( - sample_column is not None or sample_artifact is not None - ): - raise ValueError("Only sampleMad can include a sample source") - - -class CellQcProfileEvidence(AgentDataModel): - """Projected retention for one executor-supported cell-QC profile.""" - - profileId: str = "" - action: CellQcAction = "skip" - driverAssay: str | None = None - driverAssayType: CellQcDriverType | None = None - sampleColumn: str | None = None - sampleArtifact: NamedArtifactSource | None = None - attributes: list[str] = Field(default_factory=list) - artifactMetrics: list[NamedArtifactSource] = Field(default_factory=list) - parameters: dict[str, Any] = Field(default_factory=dict) - activeCells: int = 0 - retainedCells: int = 0 - retainedFraction: float = 0.0 - sampleRetainedCells: dict[str, int] = Field(default_factory=dict) - notes: list[str] = Field(default_factory=list) - evidenceId: str = "" - - @model_validator(mode="after") - def validate_sources(self) -> "CellQcProfileEvidence": - _validate_qc_sources( - action=self.action, - attributes=self.attributes, - artifact_metrics=self.artifactMetrics, - sample_column=self.sampleColumn, - sample_artifact=self.sampleArtifact, - ) - return self - - @classmethod - def get_blank(cls) -> "CellQcProfileEvidence": - return cls() - - @classmethod - def get_example(cls) -> "CellQcProfileEvidence": - return cls( - profileId="cellQc:RNA:RNA:globalGaussian:0.01:0.99", - action="globalGaussian", - driverAssay="RNA", - driverAssayType="RNA", - attributes=["RNA_nCounts", "RNA_nFeatures"], - artifactMetrics=[NamedArtifactSource.get_example()], - parameters={"minP": 0.01, "maxP": 0.99}, - activeCells=100, - retainedCells=96, - retainedFraction=0.96, - evidenceId=("qcProfile:cellQc:RNA:RNA:globalGaussian:0.01:0.99"), - ) - - -class CellQcPlan(AgentDataModel): - """A validated selection from the bounded cell-QC profiles.""" - - action: CellQcAction = "skip" - profileId: str = "" - driverAssay: str | None = None - driverAssayType: CellQcDriverType | None = None - sampleColumn: str | None = None - sampleArtifact: NamedArtifactSource | None = None - attributes: list[str] = Field(default_factory=list) - artifactMetrics: list[NamedArtifactSource] = Field(default_factory=list) - rationale: str = "" - evidenceIds: list[str] = Field(default_factory=list) - - @model_validator(mode="after") - def validate_sources(self) -> "CellQcPlan": - _validate_qc_sources( - action=self.action, - attributes=self.attributes, - artifact_metrics=self.artifactMetrics, - sample_column=self.sampleColumn, - sample_artifact=self.sampleArtifact, - ) - return self - - @classmethod - def get_blank(cls) -> "CellQcPlan": - return cls() - - @classmethod - def get_example(cls) -> "CellQcPlan": - evidence = CellQcProfileEvidence.get_example() - return cls( - action=evidence.action, - profileId=evidence.profileId, - driverAssay=evidence.driverAssay, - driverAssayType=evidence.driverAssayType, - sampleColumn=evidence.sampleColumn, - sampleArtifact=evidence.sampleArtifact, - attributes=evidence.attributes, - artifactMetrics=evidence.artifactMetrics, - rationale="Use the bounded global profile for the RNA assay.", - evidenceIds=[evidence.evidenceId], - ) - - -class ExperimentalContextDecision(AgentDataModel): - """Model-authored choices that are revalidated against the datastore.""" - - columnDomains: dict[str, ColumnDomain] = Field(default_factory=dict) - coefficientsOfInterest: list[str] = Field(default_factory=list) - unitsOfInference: dict[str, InferenceUnit] = Field(default_factory=dict) - batchCorrection: BatchCorrectionPlan = Field( - default_factory=BatchCorrectionPlan.get_blank - ) - cellQc: CellQcPlan = Field(default_factory=CellQcPlan.get_blank) - rationale: str = "" - evidenceIds: list[str] = Field(default_factory=list) - needsInput: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "ExperimentalContextDecision": - return cls() - - @classmethod - def get_example(cls) -> "ExperimentalContextDecision": - return cls( - columnDomains={ - "batch": "technical", - "sample": "design", - "donor": "design", - "treatment": "biological", - }, - coefficientsOfInterest=["treatment"], - unitsOfInference={"treatment": InferenceUnit.get_example()}, - batchCorrection=BatchCorrectionPlan.get_example(), - cellQc=CellQcPlan.get_example(), - rationale="Treatment is the primary between-sample contrast.", - evidenceIds=[ - "column:batch", - "column:donor", - "column:sample", - "column:treatment", - ], - ) - - -class RepresentationEvaluation(AgentDataModel): - """Bounded integration metrics for one exact graph representation.""" - - available: bool = False - assay: str | None = None - cellSelection: ArtifactReferenceModel | None = None - neighbors: ArtifactReferenceModel | None = None - connectivityMap: ArtifactReferenceModel | None = None - metrics: dict[str, float] = Field(default_factory=dict) - notes: list[str] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "RepresentationEvaluation": - return cls() - - @classmethod - def get_example(cls) -> "RepresentationEvaluation": - return cls( - available=True, - assay="RNA", - cellSelection=ArtifactReferenceModel( - scope="datastore", - kind="cell_selection", - artifactId="c" * 64, - ), - neighbors=ArtifactReferenceModel( - assay="RNA", - kind="neighbors", - artifactId="a" * 64, - ), - connectivityMap=ArtifactReferenceModel( - assay="RNA", - kind="connectivity_map", - artifactId="b" * 64, - ), - metrics={"iLISI:batch": 0.71, "cLISI:cell_type": 0.94}, - evidenceIds=[ - "metric:iLISI:batch:assay:RNA:neighbors:example-neighbors", - "metric:cLISI:cell_type:assay:RNA:neighbors:example-neighbors", - ], - ) - - -class CovariateEvidence(AgentDataModel): - """One deterministic covariate characterization returned by a tool.""" - - characterization: CovariateCharacterization = Field( - default_factory=lambda: CovariateCharacterization(status="needsInput") - ) - batchSafety: list[BatchSafetyEvidence] = Field(default_factory=list) - qcProfiles: list[CellQcProfileEvidence] = Field(default_factory=list) - htoIdentityColumns: list[str] = Field(default_factory=list) - htoIdentityArtifacts: list[NamedArtifactSource] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_example(cls) -> "CovariateEvidence": - return cls( - characterization=CovariateCharacterization( - status="done", - notes=["Example deterministic covariate characterization"], - ), - qcProfiles=[CellQcProfileEvidence.get_example()], - htoIdentityColumns=["sample_id"], - htoIdentityArtifacts=[ - NamedArtifactSource( - name="HTO_htoIdentity", - artifact=ArtifactReferenceModel( - assay="HTO", - kind="hto_identity", - artifactId="2" * 64, - ), - ) - ], - evidenceIds=[ - "column:batch", - CellQcProfileEvidence.get_example().evidenceId, - "htoIdentity:sample_id", - f"htoIdentityArtifact:HTO_htoIdentity:{'2' * 64}", - ], - ) - - -class ExperimentalContextResult(AgentDataModel): - """Canonical experimental-context report returned to the caller.""" - - status: StageStatus - decision: ExperimentalContextDecision - characterization: CovariateCharacterization - cellSelection: ArtifactReferenceModel | None = None - cellQc: CellQcPlan = Field(default_factory=CellQcPlan.get_blank) - qcProfiles: list[CellQcProfileEvidence] = Field(default_factory=list) - qualityMetricArtifacts: list[NamedArtifactSource] = Field(default_factory=list) - htoIdentityColumns: list[str] = Field(default_factory=list) - htoIdentityArtifacts: list[NamedArtifactSource] = Field(default_factory=list) - batchSafety: list[BatchSafetyEvidence] = Field(default_factory=list) - currentRepresentation: RepresentationEvaluation = Field( - default_factory=RepresentationEvaluation.get_blank - ) - notes: list[str] = Field(default_factory=list) - runInfo: AgentRunInfo = Field(default_factory=AgentRunInfo) - - @classmethod - def get_blank(cls) -> "ExperimentalContextResult": - return cls( - status="needsInput", - decision=ExperimentalContextDecision.get_blank(), - characterization=CovariateCharacterization(status="needsInput"), - ) - - @classmethod - def get_example(cls) -> "ExperimentalContextResult": - representation = RepresentationEvaluation.get_example() - return cls( - status="done", - decision=ExperimentalContextDecision.get_example(), - characterization=CovariateCharacterization( - status="done", - notes=["Example deterministic design characterization"], - ), - cellSelection=representation.cellSelection, - cellQc=CellQcPlan.get_example(), - qcProfiles=[CellQcProfileEvidence.get_example()], - qualityMetricArtifacts=[NamedArtifactSource.get_example()], - htoIdentityColumns=["sample_id"], - htoIdentityArtifacts=[ - NamedArtifactSource( - name="HTO_htoIdentity", - artifact=ArtifactReferenceModel( - assay="HTO", - kind="hto_identity", - artifactId="2" * 64, - ), - ) - ], - batchSafety=[BatchSafetyEvidence.get_example()], - currentRepresentation=representation, - runInfo=AgentRunInfo.get_example(), - ) - - def to_parameter_tuning_handoff(self) -> ExperimentalTuningHandoff: - """Return validated integration inputs for Parameter Tuning.""" - if self.status != "done": - raise ValueError( - "Experimental Context must be done before creating a tuning handoff" - ) - if self.cellSelection is None: - raise ValueError("Experimental Context result lacks a cell selection") - plan = self.decision.batchCorrection - batch_columns = sorted(plan.batchColumns) - safety = sorted( - ( - item - for item in self.batchSafety - if item.batchColumns == batch_columns - and item.coefficient in self.decision.coefficientsOfInterest - ), - key=lambda item: item.coefficient, - ) - if plan.action in {"evaluateHarmony", "unsafe"}: - expected = set(self.decision.coefficientsOfInterest) - if {item.coefficient for item in safety} != expected: - raise ValueError( - "Experimental Context result lacks exact batch safety evidence" - ) - if any(item.evidenceId not in plan.evidenceIds for item in safety): - raise ValueError( - "Batch-correction plan does not cite its exact safety evidence" - ) - if plan.action == "evaluateHarmony" and any( - item.status != "safe" for item in safety - ): - raise ValueError("Harmony plan contains non-safe batch evidence") - if plan.action == "unsafe" and ( - any(item.status == "notComputed" for item in safety) - or not any(item.status == "unsafe" for item in safety) - ): - raise ValueError("Unsafe plan lacks exact unsafe batch evidence") - return ExperimentalTuningHandoff( - cellSelection=self.cellSelection, - batchAction=plan.action, - batchColumns=batch_columns, - preservationColumns=list(plan.preserveColumns), - coefficientsOfInterest=list(self.decision.coefficientsOfInterest), - batchSafety=safety, - evidenceIds=sorted({*self.decision.evidenceIds, *plan.evidenceIds}), - ) - - def to_biological_handoff( - self, - coefficient: str | None = None, - ) -> ExperimentalBiologyHandoff: - """Return one explicitly resolved biological coefficient.""" - if self.status != "done": - raise ValueError( - "Experimental Context must be done before creating a biology handoff" - ) - if self.cellSelection is None: - raise ValueError("Experimental Context result lacks a cell selection") - coefficients = list(self.decision.coefficientsOfInterest) - if coefficient is None: - if len(coefficients) != 1: - raise ValueError( - "Select one coefficient explicitly for biological interpretation" - ) - coefficient = coefficients[0] - if coefficient not in coefficients: - raise ValueError(f"Unknown coefficient of interest {coefficient!r}") - records = { - record.get("name"): record - for record in self.characterization.coefficients - if isinstance(record.get("name"), str) - } - record = records.get(coefficient) - if record is None: - raise ValueError(f"Missing characterization for {coefficient!r}") - reports = { - report.get("coefficient"): report - for report in self.characterization.confounding - if isinstance(report.get("coefficient"), str) - } - report = reports.get(coefficient) - known_evidence = characterization_evidence(self.characterization) - relevant_evidence = { - f"column:{coefficient}", - f"coefficient:{coefficient}", - f"estimability:{coefficient}", - *( - evidence_id - for evidence_id in known_evidence - if evidence_id.startswith(f"confounding:{coefficient}:") - ), - } - for unit_name in ( - record.get("observationUnit"), - record.get("independentUnit"), - ): - if isinstance(unit_name, str): - relevant_evidence.add(f"column:{unit_name}") - return ExperimentalBiologyHandoff( - cellSelection=self.cellSelection, - conditionColumn=coefficient, - observationUnit=record.get("observationUnit"), - independentUnit=record.get("independentUnit"), - coefficientScope=str(record.get("scope", "")), - estimability=dict(report.get("estimability") or {}) if report else {}, - evidenceIds=sorted(relevant_evidence.intersection(known_evidence)), - ) - - -class ExperimentalContextDependencies(AgentDataModel): - """Runtime-only state shared by the agent's read-only tools.""" - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - store: Any = Field(default=None, exclude=True) - cells: Any = Field(default=None, exclude=True) - neighbors: Any = Field(default=None, exclude=True) - connectivityMap: Any = Field(default=None, exclude=True) - cellSelection: Any = Field(default=None, exclude=True) - studyContext: str = "" - directions: dict[str, Any] = Field(default_factory=dict) - evidenceIds: set[str] = Field(default_factory=set) - characterization: CovariateCharacterization | None = None - batchSafety: dict[str, BatchSafetyEvidence] = Field(default_factory=dict) - qcProfiles: dict[str, CellQcProfileEvidence] = Field(default_factory=dict) - htoIdentityColumns: list[str] = Field(default_factory=list) - qualityMetricArtifacts: list[NamedArtifactSource] = Field(default_factory=list) - htoIdentityArtifacts: list[NamedArtifactSource] = Field(default_factory=list) - currentRepresentation: RepresentationEvaluation = Field( - default_factory=RepresentationEvaluation.get_blank - ) - toolCalls: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "ExperimentalContextDependencies": - return cls() - - @classmethod - def get_example(cls) -> "ExperimentalContextDependencies": - return cls( - studyContext="Case-control study with samples nested in donors.", - directions={"columnDomains": {"batch": "technical"}}, - ) - - -def _prepare_experimental_context_tool( - ctx: RunContext[ExperimentalContextDependencies], - tool_definition: ToolDefinition, -) -> ToolDefinition | None: - """Expose each context tool once and in its required dependency order.""" - completed_calls = set(ctx.deps.toolCalls) - if tool_definition.name == "inspect_cell_covariates": - return None if tool_definition.name in completed_calls else tool_definition - if tool_definition.name == "analyze_experimental_design": - if ( - "inspect_cell_covariates" not in completed_calls - or tool_definition.name in completed_calls - ): - return None - return tool_definition - if tool_definition.name == "score_current_representation": - if ( - "analyze_experimental_design" not in completed_calls - or tool_definition.name in completed_calls - ): - return None - characterization = ctx.deps.characterization - if characterization is not None and not any( - record.get("domain") == "technical" and record.get("kind") == "categorical" - for record in characterization.columns - ): - return None - return tool_definition - return tool_definition - - -def characterization_evidence( - characterization: CovariateCharacterization, -) -> set[str]: - """Build stable evidence IDs from one deterministic characterization.""" - evidence_ids = { - f"column:{record['name']}" - for record in characterization.columns - if isinstance(record.get("name"), str) - } - for record in characterization.coefficients: - coefficient = record.get("name") - if isinstance(coefficient, str): - evidence_ids.add(f"coefficient:{coefficient}") - for report in characterization.confounding: - coefficient = report.get("coefficient") - if not isinstance(coefficient, str): - continue - evidence_ids.add(f"estimability:{coefficient}") - for pair in report.get("pairs", []): - technical = pair.get("technical") - if isinstance(technical, str): - evidence_ids.add(f"confounding:{coefficient}:{technical}") - return evidence_ids - - -def _persisted_assay_type(store: Any, assay_name: str) -> str: - """Read one persisted assay type without inferring modality from features.""" - root = getattr(store, "zw", None) - attrs = getattr(root, "attrs", {}) - raw_types = attrs.get("assayTypes", {}) if isinstance(attrs, Mapping) else {} - if isinstance(raw_types, Mapping): - assay_type = raw_types.get(assay_name) - if isinstance(assay_type, str): - return assay_type - return assay_name if assay_name in {"RNA", "ATAC", "ADT", "HTO"} else "Assay" - - -def _qc_driver(store: Any) -> tuple[str, CellQcDriverType] | None: - """Choose the first RNA assay, otherwise the first ATAC assay.""" - assay_names = [str(name) for name in getattr(store, "assay_names", [])] - for assay_type in ("RNA", "ATAC"): - for assay_name in assay_names: - if _persisted_assay_type(store, assay_name) == assay_type: - return assay_name, assay_type - return None - - -def _hto_identity_columns(deps: ExperimentalContextDependencies) -> list[str]: - """Return explicitly supplied imported HTO identity metadata columns.""" - requested: list[str] = [] - directed_many = deps.directions.get("htoIdentityColumns") - if isinstance(directed_many, list | tuple): - requested.extend(str(value) for value in directed_many) - directed_one = deps.directions.get("htoIdentityColumn") - if isinstance(directed_one, str): - requested.append(directed_one) - available = set(deps.store.cells.columns) - return list(dict.fromkeys(name for name in requested if name in available)) - - -def _cell_selection_ref(deps: ExperimentalContextDependencies) -> ArtifactRef: - selection = core_artifact_reference(deps.cellSelection) - if not isinstance(selection, ArtifactRef): - raise ValueError("cellSelection must identify an exact artifact") - if selection.kind != "cell_selection" or selection.scope != "datastore": - raise ValueError("cellSelection must identify a datastore cell selection") - return selection - - -def _active_cell_count(deps: ExperimentalContextDependencies) -> int: - selection = _cell_selection_ref(deps) - active = read_stored_selection_mask( - deps.store.zw, - selection, - kind="cell_selection", - scope="datastore", - assay=None, - table_path="cellData", - ) - if active.ndim != 1 or active.shape[0] != deps.store.cells.N: - raise ValueError( - "cellSelection must contain an aligned boolean selection vector" - ) - return int(active.sum()) - - -def _source_ref( - source: NamedArtifactSource, - *, - expected_kind: str, -) -> ArtifactRef: - if not isinstance(source, NamedArtifactSource): - raise TypeError("Artifact sources must be NamedArtifactSource values") - if not source.name.strip(): - raise ValueError("Artifact sources require a non-empty semantic name") - artifact = core_artifact_reference(source.artifact) - if not isinstance(artifact, ArtifactRef) or artifact.kind != expected_kind: - raise ValueError( - f"Artifact source {source.name!r} must reference {expected_kind!r}" - ) - return artifact - - -def _artifact_evidence_id(source: NamedArtifactSource) -> str: - return f"htoIdentityArtifact:{source.name}:{source.artifact.artifactId}" - - -def _hto_artifact_map( - deps: ExperimentalContextDependencies, -) -> dict[str, ArtifactRef]: - artifacts: dict[str, ArtifactRef] = {} - for source in deps.htoIdentityArtifacts: - if source.name in artifacts: - raise ValueError("HTO identity artifact names must be unique") - artifacts[source.name] = _source_ref( - source, - expected_kind="hto_identity", - ) - return artifacts - - -def _resolved_artifact_values( - deps: ExperimentalContextDependencies, - source: NamedArtifactSource, - *, - expected_kind: str, -) -> np.ndarray: - resolved = resolve_cell_aligned_artifact( - deps.store.zw, - _source_ref(source, expected_kind=expected_kind), - cell_selection=_cell_selection_ref(deps), - expected_kind=expected_kind, - ) - return np.asarray(resolved.values) - - -def _qc_attributes(store: Any, assay_name: str, assay_type: str) -> list[str]: - del assay_type - suffixes = ["nCounts", "nFeatures"] - available = set(store.cells.columns) - return [ - f"{assay_name}_{suffix}" - for suffix in suffixes - if f"{assay_name}_{suffix}" in available - ] - - -def _qc_sample_columns( - deps: ExperimentalContextDependencies, - characterization: CovariateCharacterization | None, -) -> list[str]: - requested: list[str] = [] - directed = deps.directions.get("cellQc") - if isinstance(directed, Mapping): - sample_column = directed.get("sampleColumn") - if isinstance(sample_column, str): - requested.append(sample_column) - if characterization is not None: - for record in characterization.coefficients: - observation_unit = record.get("observationUnit") - if isinstance(observation_unit, str): - requested.append(observation_unit) - requested.extend(deps.htoIdentityColumns) - available = set(deps.store.cells.columns) - return list( - dict.fromkeys(name for name in requested if name in available and name != "I") - )[:_MAX_QC_SAMPLE_PROFILES] - - -def _qc_profile_id( - action: CellQcAction, - *, - driver: tuple[str, CellQcDriverType] | None, - sample_column: str | None = None, - sample_artifact: NamedArtifactSource | None = None, -) -> str: - assay_name, assay_type = driver or ("none", "none") - suffix = { - "skip": "skip", - "globalGaussian": "globalGaussian:0.01:0.99", - "sampleMad": ( - f"sampleMad:metadata:{sample_column}:3:20" - if sample_artifact is None - else ( - f"sampleMad:artifact:{sample_artifact.name}:" - f"{sample_artifact.artifact.artifactId}:3:20" - ) - ), - }[action] - return f"cellQc:{assay_type}:{assay_name}:{suffix}" - - -def _global_qc_profile( - deps: ExperimentalContextDependencies, - driver: tuple[str, CellQcDriverType], - active: np.ndarray, - active_cells: int, - values_by_attr: dict[str, np.ndarray], - metadata_attributes: list[str], - artifact_metrics: list[NamedArtifactSource], - attribute_notes: list[str], -) -> CellQcProfileEvidence | None: - """Build the bounded global Gaussian QC profile when bounds are valid.""" - resolved_bounds: dict[str, dict[str, float]] = {} - global_keep = active.copy() - selected_names: list[str] = [] - for attribute, values in values_by_attr.items(): - if float(np.std(values)) == 0.0: - attribute_notes.append(f"Ignored constant QC metric {attribute!r}") - continue - low, high = gaussian_quantile_bounds(values, 0.01, 0.99) - if not np.isfinite([low, high]).all(): - attribute_notes.append( - f"Ignored QC column {attribute!r} with non-finite Gaussian bounds" - ) - continue - resolved_bounds[attribute] = {"low": low, "high": high} - selected_names.append(attribute) - global_keep &= (values > low) & (values < high) - if not selected_names: - return None - retained_cells = int(global_keep.sum()) - profile_id = _qc_profile_id( - "globalGaussian", - driver=driver, - ) - selected = set(selected_names) - return CellQcProfileEvidence( - profileId=profile_id, - action="globalGaussian", - driverAssay=driver[0], - driverAssayType=driver[1], - attributes=[name for name in metadata_attributes if name in selected], - artifactMetrics=[ - source for source in artifact_metrics if source.name in selected - ], - parameters={ - "minP": 0.01, - "maxP": 0.99, - "resolvedBounds": resolved_bounds, - }, - activeCells=active_cells, - retainedCells=retained_cells, - retainedFraction=retained_cells / active_cells, - notes=attribute_notes, - evidenceId=f"qcProfile:{profile_id}", - ) - - -def _sample_qc_profiles( - deps: ExperimentalContextDependencies, - characterization: CovariateCharacterization | None, - driver: tuple[str, CellQcDriverType], - active: np.ndarray, - active_cells: int, - values_by_attr: dict[str, np.ndarray], - metadata_attributes: list[str], - artifact_metrics: list[NamedArtifactSource], -) -> list[CellQcProfileEvidence]: - """Build bounded sample-aware MAD profiles from exact sample sources.""" - attributes = list(values_by_attr) - profiles: list[CellQcProfileEvidence] = [] - sample_sources: list[tuple[str | None, NamedArtifactSource | None]] = [ - (None, source) for source in deps.htoIdentityArtifacts - ] - sample_sources.extend( - (column, None) for column in _qc_sample_columns(deps, characterization) - ) - for sample_column, sample_artifact in sample_sources[:_MAX_QC_SAMPLE_PROFILES]: - if not attributes: - break - artifact_labels = ( - None - if sample_artifact is None - else _resolved_artifact_values( - deps, - sample_artifact, - expected_kind="hto_identity", - ) - ) - try: - sample_labels = ( - np.asarray(deps.cells.fetch(sample_column)) - if sample_column is not None - else np.asarray(artifact_labels) - ) - keep, provenance = _sample_aware_mad_mask( - values_by_attr=values_by_attr, - sample_labels=sample_labels, - active=active, - n_mads=3.0, - min_cells_per_sample=20, - attrs=attributes, - ) - except (TypeError, ValueError): - continue - retained_mask = active & keep - retained_cells = int(retained_mask.sum()) - sample_retention: dict[str, int] = {} - seen: set[object] = set() - for label in sample_labels[active]: - value = label.item() if isinstance(label, np.generic) else label - if value in seen: - continue - seen.add(value) - key = value.decode("utf-8") if isinstance(value, bytes) else str(value) - sample_retention[key] = int( - (retained_mask & (sample_labels == label)).sum() - ) - notes = list(provenance["warnings"]) - if len(sample_retention) > _MAX_SAMPLE_RETENTION_ITEMS: - notes.append( - "Per-sample retention was truncated to the first " - f"{_MAX_SAMPLE_RETENTION_ITEMS} samples" - ) - sample_retention = dict( - list(sample_retention.items())[:_MAX_SAMPLE_RETENTION_ITEMS] - ) - profile_id = _qc_profile_id( - "sampleMad", - driver=driver, - sample_column=sample_column, - sample_artifact=sample_artifact, - ) - profiles.append( - CellQcProfileEvidence( - profileId=profile_id, - action="sampleMad", - driverAssay=driver[0], - driverAssayType=driver[1], - sampleColumn=sample_column, - sampleArtifact=sample_artifact, - attributes=list(metadata_attributes), - artifactMetrics=list(artifact_metrics), - parameters={ - "nMads": 3.0, - "minCellsPerSample": 20, - "nSamples": len(provenance["sample_sizes"]), - "nSkippedSamples": len(provenance["skip_reasons"]), - }, - activeCells=active_cells, - retainedCells=retained_cells, - retainedFraction=retained_cells / active_cells, - sampleRetainedCells=sample_retention, - notes=notes, - evidenceId=f"qcProfile:{profile_id}", - ) - ) - return profiles - - -def _offered_qc_profiles( - deps: ExperimentalContextDependencies, - characterization: CovariateCharacterization | None = None, -) -> list[CellQcProfileEvidence]: - """Project bounded QC profiles against the exact shared cell selection.""" - active_cells = _active_cell_count(deps) - active = np.ones(active_cells, dtype=bool) - driver = _qc_driver(deps.store) - driver_assay = driver[0] if driver is not None else None - driver_type = driver[1] if driver is not None else None - skip_id = _qc_profile_id( - "skip", - driver=driver, - ) - skip_notes = ( - [] - if driver is not None - else ["No RNA or ATAC assay is eligible to drive automatic cell QC"] - ) - profiles = [ - CellQcProfileEvidence( - profileId=skip_id, - action="skip", - driverAssay=driver_assay, - driverAssayType=driver_type, - activeCells=active_cells, - retainedCells=active_cells, - retainedFraction=1.0 if active_cells else 0.0, - notes=skip_notes, - evidenceId=f"qcProfile:{skip_id}", - ) - ] - if driver is None or active_cells == 0: - deps.qcProfiles = {profile.profileId: profile for profile in profiles} - return profiles - - driver_assay, driver_type = driver - metadata_attributes = _qc_attributes(deps.store, driver_assay, driver_type) - values_by_attr: dict[str, np.ndarray] = {} - attribute_notes: list[str] = [] - for attribute in metadata_attributes: - try: - values = np.asarray( - deps.cells.fetch(attribute), - dtype=float, - ) - except (TypeError, ValueError): - attribute_notes.append(f"Ignored non-numeric QC column {attribute!r}") - continue - if values.ndim != 1 or values.shape != active.shape: - attribute_notes.append(f"Ignored unaligned QC column {attribute!r}") - continue - if not np.isfinite(values).all(): - attribute_notes.append(f"Ignored non-finite QC column {attribute!r}") - continue - values_by_attr[attribute] = values - valid_metadata_attributes = [ - attribute for attribute in metadata_attributes if attribute in values_by_attr - ] - artifact_metrics: list[NamedArtifactSource] = [] - for source in deps.qualityMetricArtifacts: - artifact = _source_ref(source, expected_kind="quality_metric") - if artifact.assay != driver_assay: - continue - if source.name in values_by_attr: - raise ValueError( - f"QC artifact name {source.name!r} collides with a metadata metric" - ) - values = np.asarray( - _resolved_artifact_values( - deps, - source, - expected_kind="quality_metric", - ), - dtype=float, - ) - if values.ndim != 1 or values.shape != active.shape: - raise ValueError( - f"QC artifact {source.name!r} does not align with cellSelection" - ) - if not np.isfinite(values).all(): - raise ValueError(f"QC artifact {source.name!r} contains non-finite values") - values_by_attr[source.name] = values - artifact_metrics.append(source) - - global_profile = _global_qc_profile( - deps, - driver, - active, - active_cells, - values_by_attr, - valid_metadata_attributes, - artifact_metrics, - attribute_notes, - ) - if global_profile is not None: - profiles.append(global_profile) - profiles.extend( - _sample_qc_profiles( - deps, - characterization, - driver, - active, - active_cells, - values_by_attr, - valid_metadata_attributes, - artifact_metrics, - ) - ) - - deps.qcProfiles = {profile.profileId: profile for profile in profiles} - return profiles - - -async def inspect_cell_covariates( - ctx: RunContext[ExperimentalContextDependencies], -) -> CovariateEvidence: - """Inspect cell metadata without making model-driven choices or writing data.""" - logger.info( - "Experimental Context covariate inspection started: " - f"cellSelection={ctx.deps.cellSelection.artifact_id}" - ) - ctx.deps.htoIdentityColumns = _hto_identity_columns(ctx.deps) - characterization = characterize_covariates( - ctx.deps.store, - cellSelection=ctx.deps.cellSelection, - studyContext=ctx.deps.studyContext, - model=None, - directions=ctx.deps.directions, - groupingArtifacts=_hto_artifact_map(ctx.deps), - ) - ctx.deps.characterization = characterization - qc_profiles = _offered_qc_profiles(ctx.deps) - evidence_ids = characterization_evidence(characterization) - evidence_ids.update(profile.evidenceId for profile in qc_profiles) - evidence_ids.update( - f"htoIdentity:{column}" for column in ctx.deps.htoIdentityColumns - ) - evidence_ids.update( - _artifact_evidence_id(source) for source in ctx.deps.htoIdentityArtifacts - ) - ctx.deps.evidenceIds.update(evidence_ids) - ctx.deps.toolCalls.append("inspect_cell_covariates") - logger.info( - "Experimental Context covariate inspection completed: " - f"status={characterization.status}, " - f"columns={len(characterization.columns)}, " - f"coefficients={len(characterization.coefficients)}, " - f"qcProfiles={len(qc_profiles)}, " - f"htoIdentities={len(ctx.deps.htoIdentityColumns)}, " - f"evidence={len(evidence_ids)}" - ) - return CovariateEvidence( - characterization=characterization, - qcProfiles=qc_profiles, - htoIdentityColumns=ctx.deps.htoIdentityColumns, - htoIdentityArtifacts=ctx.deps.htoIdentityArtifacts, - evidenceIds=sorted(evidence_ids), - ) - - -async def analyze_experimental_design( - ctx: RunContext[ExperimentalContextDependencies], - column_domains: dict[str, ColumnDomain], - coefficients_of_interest: list[str], - units_of_inference: dict[str, InferenceUnit], - batch_columns: list[str] | str | None = None, -) -> CovariateEvidence: - """Validate proposed domains and inference units and compute confounding. - - Args: - ctx: Pydantic AI run context containing the existing datastore. - column_domains: Domain assignment for each metadata column under review. - coefficients_of_interest: Biological columns representing study contrasts. - units_of_inference: Observation and independent units for each coefficient. - batch_columns: Exact technical columns proposed for Harmony evaluation. - A single column may be supplied as either a string or a one-item list. - """ - proposed_batch_count = ( - 1 if isinstance(batch_columns, str) else len(batch_columns or []) - ) - logger.info( - "Experimental Context design analysis started: " - f"domains={len(column_domains)}, " - f"coefficients={len(coefficients_of_interest)}, " - f"inferenceUnits={len(units_of_inference)}, " - f"batchColumns={proposed_batch_count}" - ) - directions = dict(ctx.deps.directions) - directed_domains = dict(column_domains) - directed_domains.update(dict(directions.get("columnDomains") or {})) - directions["columnDomains"] = directed_domains - directed_coefficients = list( - dict.fromkeys( - [ - *coefficients_of_interest, - *(directions.get("coefficientsOfInterest") or []), - ] - ) - ) - directions["coefficientsOfInterest"] = directed_coefficients - directed_units = { - name: unit.model_dump(exclude_none=True) - for name, unit in units_of_inference.items() - } - directed_units.update(dict(directions.get("unitsOfInference") or {})) - directions["unitsOfInference"] = directed_units - - characterization = characterize_covariates( - ctx.deps.store, - cellSelection=ctx.deps.cellSelection, - studyContext=ctx.deps.studyContext, - model=None, - directions=directions, - groupingArtifacts=_hto_artifact_map(ctx.deps), - ) - if characterization.status == "failed": - logger.warning("Experimental Context design characterization failed") - raise ModelRetry("; ".join(characterization.notes)) - - proposed_batch_columns = ( - [batch_columns] if isinstance(batch_columns, str) else list(batch_columns or []) - ) - canonical_batch_columns = sorted(set(proposed_batch_columns)) - if len(canonical_batch_columns) != len(proposed_batch_columns): - raise ModelRetry("Proposed batch columns must be unique") - column_records = { - record.get("name"): record - for record in characterization.columns - if isinstance(record.get("name"), str) - } - for batch_column in canonical_batch_columns: - record = column_records.get(batch_column) - if record is None: - raise ModelRetry(f"Unknown batch column {batch_column!r}") - if record.get("domain") != "technical": - raise ModelRetry( - f"Batch column {batch_column!r} must be classified as technical" - ) - if record.get("kind") != "categorical": - raise ModelRetry( - f"Batch column {batch_column!r} must be categorical for Harmony" - ) - - coefficient_records = { - record.get("name"): record - for record in characterization.coefficients - if isinstance(record.get("name"), str) - } - confounding_reports = { - report.get("coefficient"): report - for report in characterization.confounding - if isinstance(report.get("coefficient"), str) - } - batch_safety: list[BatchSafetyEvidence] = [] - for coefficient in directed_coefficients: - if not canonical_batch_columns: - break - coefficient_record = coefficient_records.get(coefficient) - report = confounding_reports.get(coefficient) - coefficient_kind = ( - coefficient_record.get("kind") if coefficient_record is not None else None - ) - if coefficient_kind not in {"categorical", "continuous"}: - coefficient_kind = None - observation_unit = ( - report.get("observationUnit") - if report is not None - else ( - coefficient_record.get("observationUnit") - if coefficient_record is not None - else None - ) - ) - unit_constant = { - pair.get("technical") - for pair in (report.get("pairs", []) if report is not None else []) - if isinstance(pair.get("technical"), str) - } - effective_batch_columns = [ - name for name in canonical_batch_columns if name in unit_constant - ] - estimability: dict[str, Any] - if ( - coefficient_record is None - or coefficient_record.get("scope") != "betweenUnit" - or report is None - or not isinstance(observation_unit, str) - or coefficient_kind is None - ): - estimability = { - "status": "notComputed", - "reason": "unresolvedCoefficientDesign", - } - else: - try: - design = reduce_observation_units( - ctx.deps.cells, - observation_unit, - [coefficient, *effective_batch_columns], - cell_key="I", - ) - estimability = coefficient_estimability( - design[coefficient].to_numpy(), - coefficientKind=coefficient_kind, - technicals={ - name: design[name].to_numpy() - for name in effective_batch_columns - }, - technicalKinds={ - name: column_records[name]["kind"] - for name in effective_batch_columns - }, - ) - except (KeyError, TypeError, ValueError) as exc: - logger.debug( - "Experimental Context batch estimability was not computed: " - f"errorType={type(exc).__name__}" - ) - estimability = { - "status": "notComputed", - "reason": type(exc).__name__, - } - if estimability.get("status") != "ok": - safety_status: BatchSafetyStatus = "notComputed" - elif estimability.get("coefficientEstimable") is True and not bool( - estimability.get("rankDeficient") - ): - safety_status = "safe" - else: - safety_status = "unsafe" - batch_token = ",".join(canonical_batch_columns) - safety = BatchSafetyEvidence( - coefficient=coefficient, - coefficientKind=coefficient_kind, - observationUnit=( - observation_unit if isinstance(observation_unit, str) else None - ), - batchColumns=canonical_batch_columns, - unitConstantBatchColumns=effective_batch_columns, - status=safety_status, - estimability=estimability, - evidenceId=f"batchEstimability:{coefficient}:{batch_token}", - ) - batch_safety.append(safety) - ctx.deps.batchSafety[safety.evidenceId] = safety - - ctx.deps.characterization = characterization - if not ctx.deps.htoIdentityColumns: - ctx.deps.htoIdentityColumns = _hto_identity_columns(ctx.deps) - qc_profiles = _offered_qc_profiles(ctx.deps, characterization) - evidence_ids = characterization_evidence(characterization) - evidence_ids.update(item.evidenceId for item in batch_safety) - evidence_ids.update(profile.evidenceId for profile in qc_profiles) - evidence_ids.update( - f"htoIdentity:{column}" for column in ctx.deps.htoIdentityColumns - ) - evidence_ids.update( - _artifact_evidence_id(source) for source in ctx.deps.htoIdentityArtifacts - ) - ctx.deps.evidenceIds.update(evidence_ids) - ctx.deps.toolCalls.append("analyze_experimental_design") - safety_counts = { - status: sum(item.status == status for item in batch_safety) - for status in ("safe", "unsafe", "notComputed") - } - logger.info( - "Experimental Context design analysis completed: " - f"status={characterization.status}, " - f"batchSafetySafe={safety_counts['safe']}, " - f"batchSafetyUnsafe={safety_counts['unsafe']}, " - f"batchSafetyNotComputed={safety_counts['notComputed']}, " - f"qcProfiles={len(qc_profiles)}, evidence={len(evidence_ids)}" - ) - return CovariateEvidence( - characterization=characterization, - batchSafety=batch_safety, - qcProfiles=qc_profiles, - htoIdentityColumns=ctx.deps.htoIdentityColumns, - htoIdentityArtifacts=ctx.deps.htoIdentityArtifacts, - evidenceIds=sorted(evidence_ids), - ) - - -async def score_current_representation( - ctx: RunContext[ExperimentalContextDependencies], - batch_column: str, - biological_column: str | None = None, -) -> RepresentationEvaluation: - """Score one explicitly supplied graph without changing datastore state. - - Args: - ctx: Pydantic AI run context containing the existing datastore. - batch_column: Categorical technical column used to assess batch mixing. - biological_column: Optional biological label used to assess preservation. - """ - logger.info( - "Experimental Context representation scoring started: " - f"graphSupplied={ctx.deps.neighbors is not None}, " - f"biologicalLabelSpecified={biological_column is not None}" - ) - store = ctx.deps.store - available_columns = set(store.cells.columns) - if batch_column not in available_columns: - raise ModelRetry(f"Unknown batch column {batch_column!r}") - if biological_column is not None and biological_column not in available_columns: - raise ModelRetry(f"Unknown biological column {biological_column!r}") - characterization = ctx.deps.characterization - if characterization is not None: - batch_record = next( - ( - record - for record in characterization.columns - if record.get("name") == batch_column - ), - None, - ) - if ( - batch_record is None - or batch_record.get("domain") != "technical" - or batch_record.get("kind") != "categorical" - ): - raise ModelRetry( - "Representation scoring requires a characterized categorical " - "technical batch column" - ) - - neighbors = core_artifact_reference(ctx.deps.neighbors) - connectivity = core_artifact_reference(ctx.deps.connectivityMap) - if neighbors is None: - evaluation = RepresentationEvaluation( - cellSelection=( - artifact_reference(ctx.deps.cellSelection) - if ctx.deps.cellSelection is not None - else None - ), - notes=["No exact neighbors artifact was supplied"], - ) - ctx.deps.currentRepresentation = evaluation - ctx.deps.toolCalls.append("score_current_representation") - logger.info( - "Experimental Context representation scoring skipped: " - "no current neighbors artifact" - ) - return evaluation - if not isinstance(neighbors, ArtifactRef) or neighbors.kind != "neighbors": - raise ModelRetry("neighbors must identify an exact neighbors artifact") - if connectivity is not None and ( - not isinstance(connectivity, ArtifactRef) - or connectivity.kind not in {"connectivity_map", "integrated_graph"} - ): - raise ModelRetry( - "connectivity_map must identify an exact connectivity graph artifact" - ) - - metrics: dict[str, float] = {} - notes: list[str] = [] - evidence_ids: list[str] = [] - neighbor_route = f"assay:{neighbors.assay}:neighbors:{neighbors.artifact_id}" - try: - value = float(store.metric_ilisi(batch_column, neighbors)) - if math.isfinite(value): - metrics[f"iLISI:{batch_column}"] = value - evidence_ids.append(f"metric:iLISI:{batch_column}:{neighbor_route}") - except (KeyError, TypeError, ValueError, RuntimeError) as exc: - logger.debug( - "Experimental Context iLISI scoring was unavailable: " - f"errorType={type(exc).__name__}" - ) - notes.append(f"iLISI could not be scored: {exc}") - try: - value = float(store.metric_proportional_batch_mixing(batch_column, neighbors)) - if math.isfinite(value): - metrics[f"proportionalBatchMixing:{batch_column}"] = value - evidence_ids.append( - f"metric:proportionalBatchMixing:{batch_column}:{neighbor_route}" - ) - except (KeyError, TypeError, ValueError, RuntimeError) as exc: - logger.debug( - "Experimental Context batch-mixing scoring was unavailable: " - f"errorType={type(exc).__name__}" - ) - notes.append(f"Proportional batch mixing could not be scored: {exc}") - if biological_column is not None: - try: - value = float(store.metric_clisi(biological_column, neighbors)) - if math.isfinite(value): - metrics[f"cLISI:{biological_column}"] = value - evidence_ids.append( - f"metric:cLISI:{biological_column}:{neighbor_route}" - ) - except (KeyError, TypeError, ValueError, RuntimeError) as exc: - logger.debug( - "Experimental Context cLISI scoring was unavailable: " - f"errorType={type(exc).__name__}" - ) - notes.append(f"cLISI could not be scored: {exc}") - if connectivity is not None: - try: - value = float( - store.metric_graph_connectivity(biological_column, connectivity) - ) - if math.isfinite(value): - metrics[f"graphConnectivity:{biological_column}"] = value - evidence_ids.append( - "metric:graphConnectivity:" - f"{biological_column}:assay:{connectivity.assay}:connectivity:" - f"{connectivity.artifact_id}" - ) - except (KeyError, TypeError, ValueError, RuntimeError) as exc: - logger.debug( - "Experimental Context connectivity scoring was unavailable: " - f"errorType={type(exc).__name__}" - ) - notes.append(f"Graph connectivity could not be scored: {exc}") - - evaluation = RepresentationEvaluation( - available=bool(metrics), - assay=neighbors.assay, - cellSelection=( - artifact_reference(ctx.deps.cellSelection) - if ctx.deps.cellSelection is not None - else None - ), - neighbors=artifact_reference(neighbors), - connectivityMap=( - artifact_reference(connectivity) if connectivity is not None else None - ), - metrics=metrics, - notes=notes, - evidenceIds=evidence_ids, - ) - ctx.deps.currentRepresentation = evaluation - ctx.deps.evidenceIds.update(evidence_ids) - ctx.deps.toolCalls.append("score_current_representation") - logger.info( - "Experimental Context representation scoring completed: " - f"available={evaluation.available}, metrics={len(metrics)}, " - f"notes={len(notes)}, evidence={len(evidence_ids)}" - ) - return evaluation - - -def _canonical_cell_qc_plan( - plan: CellQcPlan, - deps: ExperimentalContextDependencies, - characterization: CovariateCharacterization, -) -> CellQcPlan: - """Resolve one exact offered profile and reject model-authored parameters.""" - if not deps.qcProfiles: - _offered_qc_profiles(deps, characterization) - directed = deps.directions.get("cellQc") - direction_map = dict(directed) if isinstance(directed, Mapping) else {} - directed_profile_id = direction_map.get("profileId") - if directed_profile_id is not None and not isinstance(directed_profile_id, str): - raise ModelRetry("cellQc.profileId direction must be a string") - - has_directed_selector = any( - key in direction_map - for key in ("profileId", "action", "sampleColumn", "sampleArtifactName") - ) - selected_id = directed_profile_id or ( - "" if has_directed_selector else plan.profileId - ) - if not selected_id: - requested_action = direction_map.get("action") - requested_sample = direction_map.get("sampleColumn") - requested_sample_artifact = direction_map.get("sampleArtifactName") - if requested_sample is not None and requested_sample_artifact is not None: - raise ModelRetry( - "cellQc directions cannot select both sampleColumn and " - "sampleArtifactName" - ) - if requested_sample_artifact is not None and not isinstance( - requested_sample_artifact, str - ): - raise ModelRetry("cellQc.sampleArtifactName must be a string") - if requested_action is not None and requested_action not in { - "skip", - "globalGaussian", - "sampleMad", - }: - raise ModelRetry(f"Unsupported cellQc.action {requested_action!r}") - matches = [ - profile - for profile in deps.qcProfiles.values() - if (requested_action is None or profile.action == requested_action) - and (requested_sample is None or profile.sampleColumn == requested_sample) - and ( - requested_sample_artifact is None - or ( - profile.sampleArtifact is not None - and profile.sampleArtifact.name == requested_sample_artifact - ) - ) - ] - if requested_action is not None: - if len(matches) != 1: - raise ModelRetry( - "cellQc directions must identify exactly one offered profile" - ) - selected_id = matches[0].profileId - else: - global_profiles = [ - profile - for profile in deps.qcProfiles.values() - if profile.action == "globalGaussian" - ] - if global_profiles: - selected_id = global_profiles[0].profileId - else: - selected_id = next( - profile.profileId - for profile in deps.qcProfiles.values() - if profile.action == "skip" - ) - - profile = deps.qcProfiles.get(selected_id) - if profile is None: - raise ModelRetry( - f"Cell-QC profile {selected_id!r} was not offered by the evidence tool" - ) - model_selected = bool(plan.profileId) and not has_directed_selector - if model_selected: - expected_fields = { - "action": profile.action, - "driverAssay": profile.driverAssay, - "driverAssayType": profile.driverAssayType, - "sampleColumn": profile.sampleColumn, - "sampleArtifact": profile.sampleArtifact, - "attributes": profile.attributes, - "artifactMetrics": profile.artifactMetrics, - } - mismatches = [ - name - for name, expected in expected_fields.items() - if getattr(plan, name) != expected - ] - if mismatches: - raise ModelRetry( - "Cell-QC plan must copy the selected offered profile exactly: " - f"{mismatches}" - ) - if profile.evidenceId not in plan.evidenceIds: - raise ModelRetry( - "Cell-QC plan must cite its exact profile retention evidence" - ) - rationale = plan.rationale.strip() - if not rationale: - rationale = ( - "Selected the caller-directed bounded cell-QC profile." - if direction_map - else "Selected the bounded default cell-QC profile." - ) - cited_evidence = plan.evidenceIds if model_selected else [] - return CellQcPlan( - action=profile.action, - profileId=profile.profileId, - driverAssay=profile.driverAssay, - driverAssayType=profile.driverAssayType, - sampleColumn=profile.sampleColumn, - sampleArtifact=profile.sampleArtifact, - attributes=profile.attributes, - artifactMetrics=profile.artifactMetrics, - rationale=rationale, - evidenceIds=sorted({*cited_evidence, profile.evidenceId}), - ) - - -def _validate_batch_correction_plan( - decision: ExperimentalContextDecision, - deps: ExperimentalContextDependencies, - characterization: CovariateCharacterization, - requested_coefficients: set[str], - units_of_inference: dict[str, dict[str, Any]], - cell_qc_plan: CellQcPlan, - records: dict[str, dict[str, Any]], - coefficient_records: dict[str, dict[str, Any]], -) -> None: - """Validate one batch plan against exact design, safety, and metric evidence.""" - confounding_reports = { - report.get("coefficient"): report - for report in characterization.confounding - if isinstance(report.get("coefficient"), str) - } - plan = decision.batchCorrection - unknown_columns = sorted(set(decision.columnDomains) - set(records)) - if unknown_columns: - raise ModelRetry(f"Unknown column domain assignments: {unknown_columns}") - unit_columns = { - unit_name - for unit in units_of_inference.values() - for unit_name in ( - unit.get("observationUnit"), - unit.get("independentUnit"), - ) - if isinstance(unit_name, str) - } - if plan.action == "evaluateHarmony" and not plan.batchColumns: - raise ModelRetry("evaluateHarmony requires at least one batch column") - if plan.action == "unsafe" and not plan.batchColumns: - raise ModelRetry("unsafe requires the exact batch columns that were assessed") - if plan.action == "skip" and plan.batchColumns: - raise ModelRetry("skip must not include batch columns") - if plan.action == "needsInput" and not decision.needsInput: - raise ModelRetry("needsInput action requires at least one concrete question") - if len(set(plan.batchColumns)) != len(plan.batchColumns): - raise ModelRetry("Batch columns must be unique") - - for batch_column in plan.batchColumns: - record = records.get(batch_column) - if record is None: - raise ModelRetry(f"Unknown batch column {batch_column!r}") - if record.get("domain") != "technical": - raise ModelRetry( - f"Batch column {batch_column!r} must be classified as technical" - ) - if record.get("kind") != "categorical": - raise ModelRetry( - f"Batch column {batch_column!r} must be categorical for Harmony" - ) - if batch_column in requested_coefficients or batch_column in unit_columns: - raise ModelRetry( - f"Batch column {batch_column!r} cannot be a coefficient or unit of inference" - ) - - if plan.action == "evaluateHarmony": - mixing_metrics = {"iLISI", "proportionalBatchMixing"} - preservation_metrics = {"cLISI", "graphConnectivity"} - if not mixing_metrics.intersection(plan.metricsRequired): - raise ModelRetry( - "evaluateHarmony requires iLISI or proportionalBatchMixing" - ) - if plan.preserveColumns and not preservation_metrics.intersection( - plan.metricsRequired - ): - raise ModelRetry( - "evaluateHarmony requires cLISI or graphConnectivity for preservation" - ) - missing_preserve = sorted(requested_coefficients - set(plan.preserveColumns)) - if missing_preserve: - raise ModelRetry( - "preserveColumns must include every coefficient of interest: " - f"{missing_preserve}" - ) - unresolved_coefficients = sorted( - coefficient - for coefficient in requested_coefficients - if coefficient_records[coefficient].get("scope") != "betweenUnit" - or coefficient not in confounding_reports - ) - if unresolved_coefficients: - raise ModelRetry( - "evaluateHarmony requires a between-unit coefficient with a " - "matching estimability report; use needsInput or unsafe for: " - f"{unresolved_coefficients}" - ) - for preserve_column in plan.preserveColumns: - record = records.get(preserve_column) - if record is None: - raise ModelRetry(f"Unknown preservation column {preserve_column!r}") - if record.get("domain") != "biological": - raise ModelRetry( - f"Preservation column {preserve_column!r} must be biological" - ) - if record.get("kind") != "categorical": - raise ModelRetry( - f"Preservation column {preserve_column!r} must be categorical" - ) - - matched_safety: list[BatchSafetyEvidence] = [] - if plan.action in {"evaluateHarmony", "unsafe"}: - canonical_batch_columns = sorted(plan.batchColumns) - for coefficient in sorted(requested_coefficients): - coefficient_record = coefficient_records[coefficient] - report = confounding_reports.get(coefficient) - observation_unit = ( - report.get("observationUnit") - if report is not None - else coefficient_record.get("observationUnit") - ) - unit_constant = { - pair.get("technical") - for pair in (report.get("pairs", []) if report is not None else []) - if isinstance(pair.get("technical"), str) - } - expected_effective = [ - name for name in canonical_batch_columns if name in unit_constant - ] - candidates = [ - item - for item in deps.batchSafety.values() - if item.coefficient == coefficient - and item.coefficientKind == coefficient_record.get("kind") - and item.observationUnit == observation_unit - and item.batchColumns == canonical_batch_columns - and item.unitConstantBatchColumns == expected_effective - ] - if len(candidates) != 1: - raise ModelRetry( - "Call analyze_experimental_design with the exact proposed batch " - f"columns before returning a recommendation for {coefficient!r}" - ) - matched_safety.append(candidates[0]) - missing_safety_evidence = sorted( - item.evidenceId - for item in matched_safety - if item.evidenceId not in plan.evidenceIds - ) - if missing_safety_evidence: - raise ModelRetry( - "Batch-correction recommendations must cite exact batch " - f"estimability evidence: {missing_safety_evidence}" - ) - not_computed = [ - item.coefficient for item in matched_safety if item.status == "notComputed" - ] - if not_computed: - raise ModelRetry( - "Batch estimability could not be computed; use action='needsInput' " - f"for: {sorted(not_computed)}" - ) - unsafe_coefficients = [ - item.coefficient for item in matched_safety if item.status == "unsafe" - ] - if plan.action == "evaluateHarmony" and unsafe_coefficients: - raise ModelRetry( - "Batch correction is unsafe because the biological coefficient is " - "not estimable after the exact proposed batch columns; use " - f"action='unsafe' for: {sorted(unsafe_coefficients)}" - ) - if plan.action == "unsafe" and not unsafe_coefficients: - raise ModelRetry( - "The exact proposed batch columns were estimable for every " - "coefficient; use action='evaluateHarmony' or 'skip'" - ) - - cited_ids = [ - *decision.evidenceIds, - *plan.evidenceIds, - *cell_qc_plan.evidenceIds, - ] - unknown_evidence = sorted(set(cited_ids) - deps.evidenceIds) - if unknown_evidence: - raise ModelRetry(f"Unknown evidence IDs: {unknown_evidence}") - if plan.action in {"evaluateHarmony", "skip", "unsafe"} and not plan.evidenceIds: - raise ModelRetry("Batch-correction recommendations require evidence IDs") - current_metric_evidence = set(deps.currentRepresentation.evidenceIds) - stale_metric_evidence = sorted( - evidence_id - for evidence_id in cited_ids - if evidence_id.startswith("metric:") - and evidence_id not in current_metric_evidence - ) - if stale_metric_evidence: - raise ModelRetry( - "Metric evidence must come from the returned exact representation: " - f"{stale_metric_evidence}" - ) - - -def validate_experimental_context( - decision: ExperimentalContextDecision, - deps: ExperimentalContextDependencies, -) -> ExperimentalContextDecision: - """Recompute and validate every model-authored design choice.""" - directions = dict(deps.directions) - column_domains = dict(decision.columnDomains) - column_domains.update(dict(directions.get("columnDomains") or {})) - directions["columnDomains"] = column_domains - directions["coefficientsOfInterest"] = list( - dict.fromkeys( - [ - *decision.coefficientsOfInterest, - *(directions.get("coefficientsOfInterest") or []), - ] - ) - ) - units_of_inference = { - name: unit.model_dump(exclude_none=True) - for name, unit in decision.unitsOfInference.items() - } - units_of_inference.update(dict(directions.get("unitsOfInference") or {})) - directions["unitsOfInference"] = units_of_inference - - characterization = characterize_covariates( - deps.store, - cellSelection=deps.cellSelection, - studyContext=deps.studyContext, - model=None, - directions=directions, - groupingArtifacts=_hto_artifact_map(deps), - ) - if characterization.status == "failed": - raise ModelRetry("; ".join(characterization.notes)) - deps.characterization = characterization - deps.evidenceIds.update(characterization_evidence(characterization)) - - if "inspect_cell_covariates" not in deps.toolCalls: - raise ModelRetry("Call inspect_cell_covariates before returning a decision") - if "analyze_experimental_design" not in deps.toolCalls: - raise ModelRetry("Call analyze_experimental_design before returning a decision") - - cell_qc_plan = _canonical_cell_qc_plan( - decision.cellQc, - deps, - characterization, - ) - deps.evidenceIds.update(profile.evidenceId for profile in deps.qcProfiles.values()) - - requested_coefficients = set(directions["coefficientsOfInterest"]) - characterized_coefficients = { - record.get("name") for record in characterization.coefficients - } - missing_coefficients = sorted( - name - for name in requested_coefficients - if name not in characterized_coefficients - ) - if missing_coefficients: - raise ModelRetry( - "Coefficients of interest must be classified as biological: " - f"{missing_coefficients}" - ) - - coefficient_records: dict[str, dict[str, Any]] = {} - for record in characterization.coefficients: - name = record.get("name") - if isinstance(name, str): - coefficient_records[name] = record - records: dict[str, dict[str, Any]] = {} - for record in characterization.columns: - name = record.get("name") - if isinstance(name, str): - records[name] = record - _validate_batch_correction_plan( - decision, - deps, - characterization, - requested_coefficients, - units_of_inference, - cell_qc_plan, - records, - coefficient_records, - ) - canonical_domains = { - name: records[name]["domain"] - for name in column_domains - if name in records - and records[name].get("domain") - in { - "biological", - "technical", - "design", - "ignore", - "unknown", - } - } - canonical_units = { - coefficient: InferenceUnit( - observationUnit=coefficient_records[coefficient].get("observationUnit"), - independentUnit=coefficient_records[coefficient].get("independentUnit"), - ) - for coefficient in directions["coefficientsOfInterest"] - if coefficient in coefficient_records - } - validated = decision.model_copy( - update={ - "columnDomains": canonical_domains, - "coefficientsOfInterest": list(directions["coefficientsOfInterest"]), - "unitsOfInference": canonical_units, - "cellQc": cell_qc_plan, - } - ) - logger.debug( - "Experimental Context decision validated: " - f"domains={len(validated.columnDomains)}, " - f"coefficients={len(validated.coefficientsOfInterest)}, " - f"cellQc={validated.cellQc.action}, " - f"batchCorrection={validated.batchCorrection.action}, " - f"needsInput={len(validated.needsInput)}" - ) - return validated - - -class ExperimentalContextAgent: - """A narrow agent for study design and batch-correction planning.""" - - def __init__( - self, - model: Any, - *, - config: AgentRunConfig | None = None, - ) -> None: - self.model = model - self.config = (config or AgentRunConfig()).with_limits( - request_limit=9, - tool_call_limit=5, - output_token_limit=32768, - timeout_seconds=600.0, - ) - self.system_prompt = ( - dedent( - """ - You are Scarf's Experimental Context Agent. Work only through the - provided read-only tools and return the structured decision schema. - - Call inspect_cell_covariates exactly once. Then call - analyze_experimental_design exactly once with all explicit domains, - all biological coefficients, every unit of inference, and the complete - exact batch-column set being considered. You may call - score_current_representation at most once when an exact supplied graph - can add evidence. Do not split metadata, coefficients, or batch columns across - calls, and do not repeat a tool call. Pass batch_columns as a JSON array, - including when the array contains exactly one column. Each tool is - removed after it succeeds, so include the complete decision context in - its single call. - - The tools return bounded cell-QC profiles projected against the exact - shared cell selection. Select one returned profileId, copy its action, - driver assay name and type, metadata attributes, artifact metrics, and - sample source exactly, and cite its evidenceId. RNA is the preferred - QC driver and ATAC is the fallback. ADT and HTO never drive automatic - cell filtering. An exact HTO identity artifact may be used as sample - or grouping evidence. It is not a live metadata column and does not - make HTO a QC driver. - - A batch column must be categorical and technical. Never use donor, - sample, observation-unit, independent-unit, biological, cluster, or - embedding columns as Harmony batch columns. A biological coefficient - that is not estimable with the exact proposed batch columns makes - correction unsafe. - LISI evaluates a representation; it does not identify which metadata - column is a batch. Recommend evaluateHarmony, not application, because - Parameter Tuning must compare exact uncorrected and corrected artifacts. - - Cite only evidenceIds returned by tools. Ask for input when study - design cannot be resolved. Never propose Python, shell commands, - direct Zarr access, or any datastore mutation. Return only fields - defined by the structured output schema. - """ - ) - .strip() - .format() - ) - - def run( - self, - store: Any, - *, - study_context: str | None = None, - cell_selection: ArtifactRef | None = None, - directions: Mapping[str, Any] | None = None, - run: "PipelineRun | None" = None, - neighbors: ArtifactRef | None = None, - connectivity_map: ArtifactRef | None = None, - quality_metric_artifacts: Sequence[NamedArtifactSource] = (), - hto_identity_artifacts: Sequence[NamedArtifactSource] = (), - ) -> ExperimentalContextResult: - """Inspect one datastore and return a validated experimental-context report.""" - study_context = (study_context or "").strip() - if len(study_context) > _CONTEXT_LIMIT: - study_context = study_context[: _CONTEXT_LIMIT - 3] + "..." - direction_map = dict(directions or {}) - if run is not None: - if ( - cell_selection is not None - or neighbors is not None - or connectivity_map is not None - ): - raise ValueError( - "run is mutually exclusive with explicit artifact inputs" - ) - if getattr(run, "_owner", store) is not store: - raise ValueError("run must be opened from this datastore") - neighbors = run["neighbors"] - cell_selection = run["analysis_cell_selection"] - connectivity_map = ( - run["connectivity_map"] if "connectivity_map" in run else None - ) - cell_selection = core_artifact_reference(cell_selection) - neighbors = core_artifact_reference(neighbors) - connectivity_map = core_artifact_reference(connectivity_map) - if not isinstance(cell_selection, ArtifactRef) or ( - cell_selection.kind != "cell_selection" - or cell_selection.scope != "datastore" - ): - raise TypeError( - "cell_selection must be a datastore cell_selection ArtifactRef" - ) - if neighbors is not None: - if not isinstance(neighbors, ArtifactRef) or neighbors.kind != "neighbors": - raise TypeError("neighbors must be a neighbors ArtifactRef") - if graph_cell_selection(store.zw, neighbors) != cell_selection: - raise ValueError( - "neighbors and metadata must use the same cell selection" - ) - if connectivity_map is not None: - if not isinstance( - connectivity_map, ArtifactRef - ) or connectivity_map.kind not in { - "connectivity_map", - "integrated_graph", - }: - raise TypeError( - "connectivity_map must be a connectivity graph ArtifactRef" - ) - if graph_cell_selection(store.zw, connectivity_map) != cell_selection: - raise ValueError( - "neighbors and connectivity_map must use the same cell selection" - ) - quality_sources = list(quality_metric_artifacts) - hto_sources = list(hto_identity_artifacts) - source_names: set[str] = set() - for sources, expected_kind in ( - (quality_sources, "quality_metric"), - (hto_sources, "hto_identity"), - ): - for source in sources: - artifact = _source_ref(source, expected_kind=expected_kind) - if source.name in source_names: - raise ValueError( - "Experimental Context artifact source names must be unique" - ) - source_names.add(source.name) - resolve_cell_aligned_artifact( - store.zw, - artifact, - cell_selection=cell_selection, - expected_kind=expected_kind, - ) - directed_qc = direction_map.get("cellQc") - directed_qc_map = dict(directed_qc) if isinstance(directed_qc, Mapping) else {} - if "cellKey" in directed_qc_map: - raise ValueError( - "cellQc.cellKey is unsupported; use the exact cell_selection input" - ) - logger.info( - "Experimental Context Agent started: " - f"cellSelection={cell_selection.artifact_id}, " - f"directions={len(direction_map)}, " - f"qualityMetrics={len(quality_sources)}, " - f"htoIdentities={len(hto_sources)}, " - f"studyContextProvided={bool(study_context)}" - ) - deps = ExperimentalContextDependencies( - store=store, - cells=_SelectionBoundCells( - store.zw, - store.cells, - cell_selection, - artifacts={ - source.name: _source_ref( - source, - expected_kind="hto_identity", - ) - for source in hto_sources - }, - ), - neighbors=neighbors, - connectivityMap=connectivity_map, - cellSelection=cell_selection, - studyContext=study_context, - directions=direction_map, - qualityMetricArtifacts=quality_sources, - htoIdentityArtifacts=hto_sources, - ) - user_prompt = ( - dedent( - """ - Characterize this experiment's metadata, select one offered cell-QC - profile, and decide whether Harmony should be evaluated. - - Study context: {study_context} - Exact cell-selection artifact: {cell_selection} - Exact quality-metric artifacts: {quality_metrics} - Exact HTO identity artifacts: {hto_identities} - Caller directions: {directions} - """ - ) - .strip() - .format( - study_context=study_context or "not provided", - cell_selection=cell_selection.artifact_id, - quality_metrics=json.dumps( - [source.model_dump(mode="json") for source in quality_sources], - sort_keys=True, - ), - hto_identities=json.dumps( - [source.model_dump(mode="json") for source in hto_sources], - sort_keys=True, - ), - directions=json.dumps(direction_map, sort_keys=True, default=str), - ) - ) - execution = run_agent_sync( - model=self.model, - output_type=ExperimentalContextDecision, - system_prompt=self.system_prompt, - user_prompt=user_prompt, - tools=( - Tool( - inspect_cell_covariates, - prepare=_prepare_experimental_context_tool, - sequential=self.config.sequentialTools, - timeout=self.config.timeoutSeconds, - ), - Tool( - analyze_experimental_design, - prepare=_prepare_experimental_context_tool, - sequential=self.config.sequentialTools, - timeout=self.config.timeoutSeconds, - ), - Tool( - score_current_representation, - prepare=_prepare_experimental_context_tool, - sequential=self.config.sequentialTools, - timeout=self.config.timeoutSeconds, - ), - ), - deps_type=ExperimentalContextDependencies, - deps=deps, - config=self.config, - name="experimental_context", - output_validator=lambda decision: validate_experimental_context( - decision, - deps, - ), - ) - decision = ExperimentalContextDecision.model_validate(execution.output) - decision = validate_experimental_context(decision, deps) - characterization = deps.characterization - if characterization is None: - characterization = characterize_covariates( - store, - cellSelection=cell_selection, - studyContext=study_context, - model=None, - directions=direction_map, - groupingArtifacts=_hto_artifact_map(deps), - ) - if characterization.status == "failed": - status: StageStatus = "failed" - elif decision.needsInput or decision.batchCorrection.action == "needsInput": - status = "needsInput" - else: - status = "done" - logger.info( - "Experimental Context Agent completed: " - f"status={status}, cellQc={decision.cellQc.action}, " - f"batchCorrection={decision.batchCorrection.action}, " - f"coefficients={len(decision.coefficientsOfInterest)}, " - f"toolCalls={len(deps.toolCalls)}, evidence={len(deps.evidenceIds)}" - ) - return ExperimentalContextResult( - status=status, - decision=decision, - characterization=characterization, - cellSelection=artifact_reference(cell_selection), - cellQc=decision.cellQc, - qcProfiles=list(deps.qcProfiles.values()), - qualityMetricArtifacts=deps.qualityMetricArtifacts, - htoIdentityColumns=deps.htoIdentityColumns, - htoIdentityArtifacts=deps.htoIdentityArtifacts, - batchSafety=list(deps.batchSafety.values()), - currentRepresentation=deps.currentRepresentation, - notes=[*characterization.notes, *decision.needsInput], - runInfo=execution.runInfo, - ) diff --git a/scarf/agent/experimental_context/__init__.py b/scarf/agent/experimental_context/__init__.py new file mode 100644 index 00000000..e5275bd9 --- /dev/null +++ b/scarf/agent/experimental_context/__init__.py @@ -0,0 +1,53 @@ +"""Tool-driven experimental-design and batch-correction assessment.""" + +from ..cell_quality.profiles import RegisteredCellQcProfile +from ..types import BatchSafetyEvidence +from .agent import ExperimentalContextAgent +from .contracts import ( + BatchCorrectionPlan, + CaptureFailureEvidence, + CellQcPlan, + CellQcProfileEvidence, + ContrastPlan, + CovariateEvidence, + ExperimentalContextDecision, + ExperimentalContextDependencies, + ExperimentalContextResult, + InferenceUnit, + NamedArtifactSource, + QcMetricSourceEvidence, + QcSourceConcordance, + RepresentationEvaluation, +) +from .tools import ( + analyze_experimental_design, + contrast_plans_from_characterization, + inspect_cell_covariates, + score_current_representation, +) +from .validation import validate_experimental_context + +__all__ = [ + "BatchCorrectionPlan", + "BatchSafetyEvidence", + "CellQcPlan", + "CellQcProfileEvidence", + "CaptureFailureEvidence", + "ContrastPlan", + "CovariateEvidence", + "ExperimentalContextAgent", + "ExperimentalContextDecision", + "ExperimentalContextDependencies", + "ExperimentalContextResult", + "InferenceUnit", + "NamedArtifactSource", + "QcMetricSourceEvidence", + "QcSourceConcordance", + "RepresentationEvaluation", + "RegisteredCellQcProfile", + "analyze_experimental_design", + "contrast_plans_from_characterization", + "inspect_cell_covariates", + "score_current_representation", + "validate_experimental_context", +] diff --git a/scarf/agent/experimental_context/agent.py b/scarf/agent/experimental_context/agent.py new file mode 100644 index 00000000..b4af3b86 --- /dev/null +++ b/scarf/agent/experimental_context/agent.py @@ -0,0 +1,605 @@ +"""Experimental-context agent prompts and execution.""" + +import json +from collections.abc import Mapping, Sequence +from textwrap import dedent +from typing import TYPE_CHECKING, Any + +from ...graph.feature_projection import graph_cell_selection +from ...metadata.selection import resolve_cell_aligned_artifact +from ...storage.refs import ArtifactRef +from ...utils.logging import logger +from .._deps import AGENT_INSTALL_HINT +from ..config import AgentRunConfig +from ..config.agent_exec import run_agent_sync +from ..tools import artifact_reference, core_artifact_reference +from ..types import StageStatus +from .characterization import _SelectionBoundCells, characterize_covariates +from .contracts import ( + CellQcPlan, + ExperimentalContextDecision, + ExperimentalContextDependencies, + ExperimentalContextResult, + NamedArtifactSource, +) +from .qc_evidence import ( + _derive_missing_percentage_artifacts, + _hto_artifact_map, + _qc_driver, + _source_ref, +) +from .tools import ( + _prepare_experimental_context_tool, + analyze_experimental_design, + capture_repair_inputs, + compact_context_evidence, + contrast_plans_from_characterization, + inspect_cell_covariates, + inspect_context_evidence, + model_evidence_tool, + restore_context_evidence, + score_current_representation, +) +from .validation import ( + failed_experimental_context_result, + validate_experimental_context, +) + +if TYPE_CHECKING: + from ...datastore.pipeline_run import PipelineRun + +try: + from pydantic_ai import Tool, UnexpectedModelBehavior +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + + +class ExperimentalContextAgent: + """A narrow agent for study design and batch-correction planning.""" + + def __init__( + self, + model: Any, + *, + config: AgentRunConfig | None = None, + ) -> None: + self.model = model + self.config = (config or AgentRunConfig()).with_limits( + request_limit=10, + tool_call_limit=10, + output_token_limit=32768, + timeout_seconds=600.0, + ) + self.system_prompt = ( + dedent( + """ + You are Scarf's Experimental Context Agent. Work only through the + provided read-only tools and return the structured decision schema. + + Call inspect_cell_covariates exactly once unless its committed evidence + is already supplied on resume. Then call + analyze_experimental_design with all explicit domains, all biological + coefficients, every unit of inference, and the complete exact batch + column set. Nominate up to eight comparisons that explain the study + objective: single variables, two-column joint effects, or associations + within categorical strata. Each comparison must have a distinct + response and either one or two explanatoryColumns with conditionedOn + null, or exactly one explanatory column and one distinct categorical + conditionedOn column. These are the at most three measured columns; + observationUnit and independentUnit are separate unit fields and do + not count toward that limit. Never repeat response among explanatory + or conditioning columns. Never append unit identifiers to explanatory + columns just to identify replication. Two explanatory columns plus a + conditioning column are unsupported. One response, one explanatory + column and one conditioning column are supported for descriptive + design coverage; association support depends on the measured units. + Separate simpler comparisons do not establish an unsupported joint + conditional finding. If a tool rejects a proposal, correct + the named fields while preserving its scientific question or record + the unsupported requirement explicitly. You may make one + follow-up call with at most four new or revised comparisons after + reading the first evidence. Never split the exact batch-column set. + A donor can carry biological variation and also be the explicitly + declared independent unit; these roles do not conflict. Keep its + biological protection when specifying its unit role. Do not replace + independent donors with cells or samples to obtain a supported result. + Protect objective-relevant pairs of categorical biological variables + with protectCombination. The tools report unsupported designs, + missingness, replication, and sparse strata; do not turn these into + negative findings. Explain unsupported requested comparisons in the + final rationale; a proposed comparison is not a completed analysis. + Set each proposal's purpose to designCoverage for measured counts, + crossing, replication, or pairing; association when an association + coefficient is essential; effectEstimation for an explicitly requested + biological effect, which this workflow cannot deliver. Copy an exact + objectiveQuote from the supplied study text and mark explicit objective + questions essential. Do not downgrade an essential association or effect + question to descriptive coverage to obtain completion. The returned + evidenceRequirements and evidenceCoverage enforce this distinction. + For repeated donors and incomplete pairing, inspect descriptiveDesign: + it retains observation counts, distinct donors, group support and paired + coverage without collapsing a donor to its first condition. A computed + descriptiveDesign answers a designCoverage question even if the separate + association method is unsupported. Preserve that association limitation + in the rationale; do not ask the user to resolve an optional association + before continuing descriptive population discovery. Unsupported + association methods do not establish non-identifiability. Only measured + rank and estimability for the exact tested design support that claim. + Use the single follow-up round to resolve missing design evidence. + If essential evidence remains unsupported, ask for clarification or + abstain. Optional questions must remain explicit limitations. + Continuous conditioning and expression hypothesis + testing are unsupported. You may call + score_current_representation at most once when an exact supplied graph + can add evidence. Pass batch_columns as a JSON array, including a + singleton. When the study explicitly identifies physical capture, pass + capture_proposal in analyze_experimental_design, preferably in its first + call. This existing tool validates capture provenance; no separate + capture tool or additional user confirmation is required for an explicit + supported study statement. Capture proposals must name an exact observed + column and quote the study statement identifying it as a physical capture. + An optional reference pool also needs an exact quote identifying the + observed reference captures. Sample uniqueness is not capture proof. + Leave unresolved capture provenance explicit. Validated tools own + capture identities and protected combinations; do not copy them into + the final decision. Nominate any new combination through a design tool. + If capture was omitted from two completed design rounds, the same tool + permits only a capture provenance repair with identical domains, + coefficients, units and batch columns, and no proposals. This does not + authorize another design comparison round. + + The tools return bounded cell-QC profiles projected against the exact + shared cell selection. Do not choose a profile or return cellQc. + Summaries retain adverse findings, missingness, protected group loss, + replication and correction constraints. Use inspect_context_evidence + to inspect one exact saved policy/capture or design record when its + details are needed. Copy the supplied section and record_id from + confounding details; do not construct an evidence ID for this lookup. + Reuse returned details instead of requesting the same record again. + Omitted donor examples and detailed thresholds + remain available; do not interpret their omission as passing evidence. + A later audited checkpoint compares the Scarf default with eligible + alternatives and selects one exact policy. Never author + or alter numeric quality bounds. RNA is the preferred QC driver and + ATAC is the fallback. ADT and HTO never drive automatic cell filtering. + An exact HTO identity artifact may be used as grouping evidence. It is + not a live metadata column and does not make HTO a QC driver. + + A batch column must be categorical and technical. Never use donor, + sample, observation-unit, independent-unit, biological, cluster, or + embedding columns as Harmony batch columns. A biological coefficient + that is not estimable with the exact proposed batch columns makes + correction unsafe. A sample or library identifier is not automatically + technical. When no exact observed column is both categorical and + technical, pass batch_columns=[] and recommend skipping Harmony. Every + observation and independent unit must be an exact observed column name + or null. + LISI evaluates a representation; it does not identify which metadata + column is a batch. Recommend evaluateHarmony, not application, because + Parameter Tuning must compare exact uncorrected and corrected artifacts. + Current Harmony preservation metrics support categorical biology. + A continuous protected variable has unavailable preservation evidence. + Keep this limitation explicit; it does not prove correction unnecessary. + Later matched acceptance must reject unavailable required protection, + and unresolved essential evidence requires input or abstention. + + Cite only evidenceIds returned by tools. Ask for input when study + design cannot be resolved. The study objective is authoritative: use + it to identify protected biological variables and the intended unit + of inference. Study-design explanations are in scope; biological + expression hypothesis testing, effect inference and causal claims are not. + Never propose Python, shell commands, + direct Zarr access, or any datastore mutation. Every rationale and + question must be plain prose. Never place serialized JSON, schema + field names, or sibling output fields inside a narrative string. + Return only fields defined by the structured output schema. + """ + ) + .strip() + .format() + ) + + def run( + self, + store: Any, + *, + study_context: str | None = None, + study_objective: str | None = None, + cell_selection: ArtifactRef | None = None, + directions: Mapping[str, Any] | None = None, + run: "PipelineRun | None" = None, + neighbors: ArtifactRef | None = None, + connectivity_map: ArtifactRef | None = None, + quality_metric_artifacts: Sequence[NamedArtifactSource] = (), + hto_identity_artifacts: Sequence[NamedArtifactSource] = (), + qc_assay: str | None = None, + checkpoint_read: Any = None, + checkpoint_write: Any = None, + on_attempt: Any = None, + previous_context: ExperimentalContextResult | None = None, + ) -> ExperimentalContextResult: + """Inspect one datastore and return a validated experimental-context report.""" + study_context = (study_context or "").strip() + study_objective = (study_objective or "").strip() + direction_map = dict(directions or {}) + if run is not None: + if ( + cell_selection is not None + or neighbors is not None + or connectivity_map is not None + ): + raise ValueError( + "run is mutually exclusive with explicit artifact inputs" + ) + if getattr(run, "_owner", store) is not store: + raise ValueError("run must be opened from this datastore") + neighbors = run["neighbors"] + cell_selection = run["analysis_cell_selection"] + connectivity_map = ( + run["connectivity_map"] if "connectivity_map" in run else None + ) + cell_selection = core_artifact_reference(cell_selection) + neighbors = core_artifact_reference(neighbors) + connectivity_map = core_artifact_reference(connectivity_map) + if not isinstance(cell_selection, ArtifactRef) or ( + cell_selection.kind != "cell_selection" + or cell_selection.scope != "datastore" + ): + raise TypeError( + "cell_selection must be a datastore cell_selection ArtifactRef" + ) + if neighbors is not None: + if not isinstance(neighbors, ArtifactRef) or neighbors.kind != "neighbors": + raise TypeError("neighbors must be a neighbors ArtifactRef") + if graph_cell_selection(store.zw, neighbors) != cell_selection: + raise ValueError( + "neighbors and metadata must use the same cell selection" + ) + if connectivity_map is not None: + if not isinstance( + connectivity_map, ArtifactRef + ) or connectivity_map.kind not in { + "connectivity_map", + "integrated_graph", + }: + raise TypeError( + "connectivity_map must be a connectivity graph ArtifactRef" + ) + if graph_cell_selection(store.zw, connectivity_map) != cell_selection: + raise ValueError( + "neighbors and connectivity_map must use the same cell selection" + ) + quality_sources = _derive_missing_percentage_artifacts( + store, + cell_selection=cell_selection, + driver=_qc_driver(store, qc_assay), + quality_sources=quality_metric_artifacts, + ) + hto_sources = list(hto_identity_artifacts) + source_names: set[str] = set() + for sources, expected_kind in ( + (quality_sources, "quality_metric"), + (hto_sources, "hto_identity"), + ): + for source in sources: + artifact = _source_ref(source, expected_kind=expected_kind) + if source.name in source_names: + raise ValueError( + "Experimental Context artifact source names must be unique" + ) + source_names.add(source.name) + resolve_cell_aligned_artifact( + store.zw, + artifact, + cell_selection=cell_selection, + expected_kind=expected_kind, + ) + directed_qc = direction_map.get("cellQc") + directed_qc_map = dict(directed_qc) if isinstance(directed_qc, Mapping) else {} + if "cellKey" in directed_qc_map: + raise ValueError( + "cellQc.cellKey is unsupported; use the exact cell_selection input" + ) + logger.info( + "Experimental Context Agent started: " + f"cellSelection={cell_selection.artifact_id}, " + f"directions={len(direction_map)}, " + f"qualityMetrics={len(quality_sources)}, " + f"htoIdentities={len(hto_sources)}, " + f"studyContextProvided={bool(study_context)}, " + f"studyObjectiveProvided={bool(study_objective)}" + ) + deps = ExperimentalContextDependencies( + store=store, + qcAssay=qc_assay, + cells=_SelectionBoundCells( + store.zw, + store.cells, + cell_selection, + artifacts={ + source.name: _source_ref( + source, + expected_kind="hto_identity", + ) + for source in hto_sources + }, + ), + neighbors=neighbors, + connectivityMap=connectivity_map, + cellSelection=cell_selection, + studyContext=study_context, + studyObjective=study_objective, + directions=direction_map, + qualityMetricArtifacts=quality_sources, + htoIdentityArtifacts=hto_sources, + checkpointRead=checkpoint_read, + checkpointWrite=checkpoint_write, + ) + if checkpoint_read is not None: + saved_result = checkpoint_read("result") + if saved_result is not None: + report = ExperimentalContextResult.model_validate( + saved_result["report"] + ) + if report.cellSelection != artifact_reference(cell_selection): + raise ValueError( + "Committed context decision has a different cell selection" + ) + return report + restored = restore_context_evidence(deps) + if not restored and previous_context is not None: + # A completed older stage remains immutable. Its measurements seed + # an explicitly requested evidence revision under new stage inputs. + deps.characterization = previous_context.characterization + deps.comparisons = list(previous_context.characterization.comparisons) + deps.captureProposal = previous_context.characterization.captureProvenance + deps.protectedCombinations = list( + previous_context.decision.protectedCombinations + ) + deps.qcProfiles = { + item.profileId: item for item in previous_context.qcProfiles + } + deps.qcMetricSources = list(previous_context.qcMetricSources) + deps.qcSourceConcordance = list(previous_context.qcSourceConcordance) + deps.batchSafety = { + item.evidenceId: item for item in previous_context.batchSafety + } + deps.contrastPlans = { + item.coefficient: item for item in previous_context.contrastPlans + } + deps.evidenceIds.update(previous_context.decision.evidenceIds) + deps.evidenceIds.update( + item.evidenceId for item in previous_context.batchSafety + ) + deps.toolCalls = ["inspect_cell_covariates", "analyze_experimental_design"] + deps.designRounds = min( + 2, + max( + 1, + sum( + item.toolName == "analyze_experimental_design" + for item in previous_context.runInfo.toolCalls + ), + ), + ) + restored = True + user_prompt = ( + dedent( + """ + Characterize this experiment's metadata and decide whether Harmony + should be evaluated. Return cell-QC candidates as tool evidence; + do not return cellQc; the later filtering checkpoint owns it. + + Study context: {study_context} + Study objective: {study_objective} + Exact cell-selection artifact: {cell_selection} + Exact quality-metric artifacts: {quality_metrics} + Exact HTO identity artifacts: {hto_identities} + Caller directions: {directions} + """ + ) + .strip() + .format( + study_context=study_context or "not provided", + study_objective=study_objective or "not provided", + cell_selection=cell_selection.artifact_id, + quality_metrics=json.dumps( + [source.model_dump(mode="json") for source in quality_sources], + sort_keys=True, + ), + hto_identities=json.dumps( + [source.model_dump(mode="json") for source in hto_sources], + sort_keys=True, + ), + directions=json.dumps(direction_map, sort_keys=True, default=str), + ) + ) + if restored: + from .contracts import CovariateEvidence + from .requirements import requested_design_questions + + assert deps.characterization is not None + + user_prompt += ( + "\nCommitted evidence already measured; do not repeat completed tools:\n" + + json.dumps( + compact_context_evidence( + CovariateEvidence( + characterization=deps.characterization, + batchSafety=list(deps.batchSafety.values()), + qcProfiles=list(deps.qcProfiles.values()), + qcMetricSources=deps.qcMetricSources, + qcSourceConcordance=deps.qcSourceConcordance, + contrastPlans=list(deps.contrastPlans.values()), + evidenceIds=sorted(deps.evidenceIds), + ) + ), + sort_keys=True, + ) + + f"\nCompleted design rounds: {deps.designRounds} of 2." + ) + user_prompt += ( + "\nExplicit requested questions requiring matched evidence: " + + json.dumps( + requested_design_questions( + study_context, + study_objective, + [row["name"] for row in deps.characterization.columns], + ) + ) + ) + repair_inputs = capture_repair_inputs(deps) + if repair_inputs is not None: + user_prompt += ( + "\nCapture provenance can be repaired with these exact saved " + "analyze_experimental_design inputs. Add only capture_proposal " + "with an exact supporting study quote: " + + json.dumps(repair_inputs, sort_keys=True) + ) + if previous_context is not None and previous_context.status == "needsInput": + from .requirements import objective_evidence + + requirements, coverage = ( + objective_evidence( + study_context=study_context, + study_objective=study_objective, + experimental_result=previous_context.model_copy( + update={ + "characterization": deps.characterization, + "batchSafety": list(deps.batchSafety.values()), + } + ), + ) + if study_objective + else ([], []) + ) + user_prompt += ( + "\nCurrent objective requirements and their measured support: " + + json.dumps( + { + "evidenceRequirements": [ + item.model_dump(mode="json") for item in requirements + ], + "evidenceCoverage": [ + item.model_dump(mode="json") for item in coverage + ], + } + ) + + "\nThe previous interpretation stopped with these unresolved questions: " + + json.dumps(previous_context.decision.needsInput) + + "\nReassess each against the original study text and committed evidence. " + "Use the existing design tool for an omitted capture proposal. " + "Distinguish computed descriptive support from unavailable association " + "methods. Keep genuinely essential unresolved questions; do not copy " + "a previous blocker when supplied provenance or measured evidence answers it." + ) + try: + execution = run_agent_sync( + model=self.model, + output_type=ExperimentalContextDecision, + system_prompt=self.system_prompt, + user_prompt=user_prompt, + tools=( + Tool( + model_evidence_tool(inspect_cell_covariates), + prepare=_prepare_experimental_context_tool, + sequential=self.config.sequentialTools, + timeout=self.config.timeoutSeconds, + ), + Tool( + model_evidence_tool(analyze_experimental_design), + max_retries=3, + prepare=_prepare_experimental_context_tool, + sequential=self.config.sequentialTools, + timeout=self.config.timeoutSeconds, + ), + Tool( + inspect_context_evidence, + sequential=self.config.sequentialTools, + timeout=self.config.timeoutSeconds, + ), + Tool( + score_current_representation, + prepare=_prepare_experimental_context_tool, + sequential=self.config.sequentialTools, + timeout=self.config.timeoutSeconds, + ), + ), + deps_type=ExperimentalContextDependencies, + deps=deps, + config=self.config, + name="experimental_context", + on_attempt=on_attempt, + output_validator=lambda decision: validate_experimental_context( + decision, + deps, + ), + ) + except UnexpectedModelBehavior as exc: + model_name = getattr(self.model, "model_name", type(self.model).__name__) + return failed_experimental_context_result( + deps, + error=exc, + model_name=str(model_name), + ) + decision = ExperimentalContextDecision.model_validate(execution.output) + run_info = execution.runInfo + characterization = deps.characterization + if characterization is None: + characterization = characterize_covariates( + store, + cellSelection=cell_selection, + studyContext=(f"{study_context}\nStudy objective: {study_objective}"), + model=None, + directions=direction_map, + groupingArtifacts=_hto_artifact_map(deps), + ) + if characterization.status == "failed": + status: StageStatus = "failed" + elif decision.needsInput or decision.batchCorrection.action == "needsInput": + status = "needsInput" + else: + status = "done" + logger.info( + "Experimental Context Agent completed: " + f"status={status}, qcProfiles={len(deps.qcProfiles)}, " + f"batchCorrection={decision.batchCorrection.action}, " + f"coefficients={len(decision.coefficientsOfInterest)}, " + f"toolCalls={len(deps.toolCalls)}, evidence={len(deps.evidenceIds)}" + ) + contrast_plans = list(deps.contrastPlans.values()) + if not contrast_plans: + contrast_plans = contrast_plans_from_characterization(characterization) + report = ExperimentalContextResult( + status=status, + decision=decision, + characterization=characterization, + cellSelection=artifact_reference(cell_selection), + cellQc=CellQcPlan.get_blank(), + qcProfiles=list(deps.qcProfiles.values()), + qcMetricSources=deps.qcMetricSources, + qcSourceConcordance=deps.qcSourceConcordance, + contrastPlans=contrast_plans, + qualityMetricArtifacts=deps.qualityMetricArtifacts, + htoIdentityColumns=deps.htoIdentityColumns, + htoIdentityArtifacts=deps.htoIdentityArtifacts, + batchSafety=list(deps.batchSafety.values()), + currentRepresentation=deps.currentRepresentation, + notes=[ + *characterization.notes, + *decision.needsInput, + *( + [ + "Matched preservation evidence is unavailable for continuous variables: " + + ", ".join(decision.unsupportedProtection) + + ". A safe design does not by itself authorize correction." + ] + if decision.unsupportedProtection + else [] + ), + ], + runInfo=run_info, + ) + if checkpoint_write is not None: + checkpoint_write("result", {"report": report.model_dump(mode="json")}) + return report diff --git a/scarf/agent/characterize_covariates.py b/scarf/agent/experimental_context/characterization.py similarity index 67% rename from scarf/agent/characterize_covariates.py rename to scarf/agent/experimental_context/characterization.py index e8a85366..682b705b 100644 --- a/scarf/agent/characterize_covariates.py +++ b/scarf/agent/experimental_context/characterization.py @@ -1,5 +1,6 @@ """Characterize cell covariates and study-design confounding.""" +import hashlib import re from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass, field @@ -8,38 +9,27 @@ import numpy as np import pandas as pd -from ..metadata.queries import ( +from ...metadata.queries import ( PartitionDigest, column_constant_within, column_partition_digest, columns_same_partition, reduce_observation_units, ) -from ..metadata.rows import ( +from ...metadata.rows import ( MetaDataRowBlock, read_metadata_missing_rows_chunkwise, read_metadata_rows_chunkwise, ) -from ..metadata.selection import resolve_cell_aligned_artifact -from ..metrics.association import directional_mapping, report_confounding -from ..storage.refs import ArtifactRef -from ..storage.selections import read_stored_selection_indices -from .config import CONFIG -from .config._deps import AGENT_INSTALL_HINT -from .decide import DecisionValidationError, decide -from .tools import artifact_reference -from .types import ( - AgentDataModel, - ArtifactReferenceModel, - Decision, - EvidenceItem, - StageStatus, -) - -try: - from pydantic import Field -except ImportError as exc: - raise ImportError(AGENT_INSTALL_HINT) from exc +from ...metadata.selection import resolve_cell_aligned_artifact +from ...metrics.association import directional_mapping, report_confounding +from ...storage.artifacts import fingerprint_array +from ...storage.refs import ArtifactRef +from ...storage.selections import read_stored_selection_indices +from ..decisions.selection import DecisionValidationError, decide +from ..tools import artifact_reference +from ..types import Decision, EvidenceItem +from .contracts import CovariateCharacterization __all__ = [ "CovariateCharacterization", @@ -49,6 +39,35 @@ Domain = Literal["biological", "technical", "design", "ignore", "unknown"] ColumnKind = Literal["categorical", "continuous"] +_CATEGORICAL_MAX_LEVELS = 50 +_EMBEDDING_TOKENS = ( + "umap", + "pca", + "tsne", + "scvi", + "latent", + "phate", + "forceatlas", + "diffmap", + "diffusionmap", + "diffusion", +) +_DOMAINS = frozenset({"biological", "technical", "design", "ignore", "unknown"}) +_ANALYSED = frozenset({"biological", "technical", "design"}) +_KINDS = frozenset({"categorical", "continuous"}) +_RESERVED_COLUMNS = frozenset({"I", "ids", "names"}) +_SHORT_EMBEDDING_PARTS = frozenset({"fa", "dm", "pc"}) +_INDEXED_NAME = re.compile(r"(?P.+?)[-_]?(?P\d+)") +_ONTOLOGY_SUFFIX = "_ontology_term_id" +_SAMPLE_LEVELS = 8 +_ASSOCIATION_FLOOR = 0.1 +_DROP_REASONS = { + "dropAssayStat": "Scarf assay statistic column", + "dropProvenance": "analysis-linked column", + "dropEmbedding": "embedding-style column", + "dropConstant": "single-level column", +} + _DOMAIN_EVIDENCE = [ EvidenceItem( @@ -91,48 +110,61 @@ ] -class CovariateCharacterization(AgentDataModel): - status: StageStatus - cellSelection: ArtifactReferenceModel | None = None - auditLog: list[dict[str, Any]] = Field(default_factory=list) - actions: list[str] = Field(default_factory=list) - notes: list[str] = Field(default_factory=list) - decisions: list[dict[str, Any]] = Field(default_factory=list) - columns: list[dict[str, Any]] = Field(default_factory=list) - coefficients: list[dict[str, Any]] = Field(default_factory=list) - technicalNesting: list[dict[str, Any]] = Field(default_factory=list) - confounding: list[dict[str, Any]] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "CovariateCharacterization": - return cls(status="failed") - - @classmethod - def get_example(cls) -> "CovariateCharacterization": - return cls( - status="done", - cellSelection=ArtifactReferenceModel( - scope="datastore", - kind="cell_selection", - artifactId="c" * 64, - ), - notes=["Cell covariates and confounding were characterized."], - columns=[{"name": "batch", "domain": "technical"}], - ) - - @dataclass(frozen=True, slots=True) class _ColumnProfile: kind: ColumnKind summary: str digest: PartitionDigest artifact: ArtifactRef | None = None + levelCounts: tuple[dict[str, Any], ...] = () + levelCountsTruncated: bool = False + + +_MAX_LEVEL_COUNT_ITEMS = 32 +_MAX_STRUCTURE_ITEMS = 32 + + +def _json_scalar(value: Any) -> str | int | float | bool | None: + if isinstance(value, np.generic): + value = value.item() + missing = pd.isna(value) + if isinstance(missing, bool | np.bool_) and bool(missing): + return None + if isinstance(value, bytes): + return value.decode("utf-8", errors="replace") + if isinstance(value, str | int | float | bool): + return value + return str(value) + + +def _bounded_level_counts( + values: Sequence[Any] | np.ndarray | pd.Series, + *, + limit: int = _MAX_LEVEL_COUNT_ITEMS, +) -> tuple[tuple[dict[str, Any], ...], bool]: + series = pd.Series(values) + counts = series.value_counts(dropna=False, sort=False) + output = tuple( + { + "value": _json_scalar(value), + "count": int(count), + } + for value, count in counts.iloc[:limit].items() + ) + return output, len(counts) > limit class _SelectionBoundCells: """Read metadata through one validated immutable cell selection.""" - __slots__ = ("_artifact_sources", "_artifact_values", "_indices", "_source") + __slots__ = ( + "_artifact_sources", + "_artifact_values", + "_indices", + "_source", + "_metadata_values", + "_cache_values", + ) def __init__( self, @@ -141,8 +173,11 @@ def __init__( selection: ArtifactRef, *, artifacts: Mapping[str, ArtifactRef] | None = None, + cache_values: bool = False, ) -> None: self._source = source + self._cache_values = cache_values + self._metadata_values: dict[str, np.ndarray] = {} self._indices = read_stored_selection_indices( root, selection, @@ -184,10 +219,16 @@ def artifact_source(self, column: str) -> ArtifactRef | None: def fetch(self, column: str, key: str = "I") -> np.ndarray: if key != "I": raise ValueError("A bound metadata view accepts only its stored selection") - return self._read(column, self._indices) + if not self._cache_values: + return self._read(column, self._indices) + if column not in self._metadata_values: + self._metadata_values[column] = self._read(column, self._indices) + return self._metadata_values[column] def _read(self, column: str, indices: np.ndarray) -> np.ndarray: - artifact_values = self._artifact_values.get(column) + artifact_values = self._metadata_values.get(column) + if artifact_values is None: + artifact_values = self._artifact_values.get(column) if artifact_values is not None: positions = np.searchsorted(self._indices, indices) if np.any(positions >= len(self._indices)) or not np.array_equal( @@ -259,6 +300,7 @@ def __init__( store.cells, selection, artifacts=artifacts, + cache_values=True, ) self.assay_names = store.assay_names @@ -324,14 +366,14 @@ def ask( def _is_embedding_column(name: str) -> bool: - match = CONFIG._INDEXED_NAME.fullmatch(name) + match = _INDEXED_NAME.fullmatch(name) if match is None: return False parts = [part for part in re.split(r"[-_]+", match.group("stem").lower()) if part] compact = "".join(parts) - if any(token in compact for token in CONFIG._EMBEDDING_TOKENS): + if any(token in compact for token in _EMBEDDING_TOKENS): return True - return any(part in CONFIG._SHORT_EMBEDDING_PARTS for part in parts) + return any(part in _SHORT_EMBEDDING_PARTS for part in parts) def _infer_kind(values: np.ndarray) -> ColumnKind: @@ -348,7 +390,7 @@ def _infer_kind(values: np.ndarray) -> ColumnKind: finite = numeric[np.isfinite(numeric)] if finite.size == 0 or not bool(np.all(np.mod(finite, 1) == 0)): return "continuous" - limit = min(CONFIG._CATEGORICAL_MAX_LEVELS, max(2, len(values) // 20)) + limit = min(_CATEGORICAL_MAX_LEVELS, max(2, len(values) // 20)) return "categorical" if int(np.unique(finite).size) <= limit else "continuous" @@ -359,7 +401,7 @@ def _summarize(values: np.ndarray, kind: ColumnKind) -> str: levels = series.dropna().astype(str).value_counts() top = ", ".join( f"{level}={int(count)}" - for level, count in levels.head(CONFIG._SAMPLE_LEVELS).items() + for level, count in levels.head(_SAMPLE_LEVELS).items() ) return f"categorical levels={levels.shape[0]} missing={missing} top=[{top}]" numeric = pd.to_numeric(series, errors="coerce").to_numpy(dtype=float, copy=False) @@ -382,18 +424,44 @@ def _profile_column( *, cell_key: str, kind: ColumnKind | None = None, + inventory: dict[str, Any] | None = None, ) -> _ColumnProfile: values = store.cells.fetch(name, key=cell_key) resolved_kind = kind or _infer_kind(values) + artifact_source = getattr(store.cells, "artifact_source", lambda _name: None)(name) + identity = None + if inventory is not None: + # Metadata are mutable even when the selected cell identities are frozen. + value_identity = ( + hashlib.sha256( + "\n".join( + repr((type(value).__name__, value)) for value in values.tolist() + ).encode() + ).hexdigest() + if values.dtype.hasobject + else fingerprint_array(values) + ) + identity = (str(values.dtype), value_identity, resolved_kind, artifact_source) + saved = inventory.get(name) + if saved is not None and saved[0] == identity: + return cast(_ColumnProfile, saved[1]) summary = _summarize(values, resolved_kind) digest = column_partition_digest(store.cells, name, cell_key=cell_key) - artifact_source = getattr(store.cells, "artifact_source", lambda _name: None)(name) - return _ColumnProfile( + level_counts: tuple[dict[str, Any], ...] = () + level_counts_truncated = False + if resolved_kind == "categorical": + level_counts, level_counts_truncated = _bounded_level_counts(values) + profile = _ColumnProfile( kind=resolved_kind, summary=summary, digest=digest, artifact=artifact_source, + levelCounts=level_counts, + levelCountsTruncated=level_counts_truncated, ) + if inventory is not None: + inventory[name] = (identity, profile) + return profile def _triage_columns( @@ -408,7 +476,7 @@ def _triage_columns( candidates: list[str] = [] dropped: list[tuple[str, str]] = [] for name in store.cells.columns: - if name in CONFIG._RESERVED_COLUMNS or name == cell_key or name in exclude: + if name in _RESERVED_COLUMNS or name == cell_key or name in exclude: continue if name in artifact_columns: candidates.append(name) @@ -438,9 +506,9 @@ def _collapse_ontology_aliases( dropped: set[str] = set() notes: list[dict[str, Any]] = [] for name in columns: - if not name.endswith(CONFIG._ONTOLOGY_SUFFIX): + if not name.endswith(_ONTOLOGY_SUFFIX): continue - base = name[: -len(CONFIG._ONTOLOGY_SUFFIX)] + base = name[: -len(_ONTOLOGY_SUFFIX)] if base not in present or {name, base} & dropped: continue if _digest_key(profiles[name].digest) != _digest_key(profiles[base].digest): @@ -467,15 +535,6 @@ def _collapse_ontology_aliases( return [name for name in columns if name not in dropped], aliases, notes -def _bounded_context(study_context: str | None) -> str: - text = (study_context or "").strip() - return ( - text - if len(text) <= CONFIG._CONTEXT_LIMIT - else text[: CONFIG._CONTEXT_LIMIT - 3] + "..." - ) - - def _validate_directions( directions: Mapping[str, Any], available: set[str], @@ -492,8 +551,8 @@ def check_names(key: str, names: Any) -> list[str] | None: return list(names) for key, allowed in ( - ("columnKinds", CONFIG._KINDS), - ("columnDomains", CONFIG._DOMAINS), + ("columnKinds", _KINDS), + ("columnDomains", _DOMAINS), ): mapping = directions.get(key) if mapping is None: @@ -594,7 +653,7 @@ class is only eligible when every member carries the same analysis domain, """ classes: dict[tuple[bytes, int, int], list[str]] = {} for name in candidates: - if run.kind(name) != "categorical" or run.domains[name] not in CONFIG._ANALYSED: + if run.kind(name) != "categorical" or run.domains[name] not in _ANALYSED: continue classes.setdefault(_digest_key(run.digest(name)), []).append(name) @@ -673,7 +732,7 @@ def _assign_domain(run: _Run, name: str, directed: Mapping[str, Domain]) -> Doma ) return "unknown" selected = decision.selectedId.removeprefix("domain:") - if selected not in CONFIG._DOMAINS: + if selected not in _DOMAINS: run.note( kind="domainUnknown", detail=f"Unsupported domain {selected!r} returned for {name}", @@ -773,6 +832,194 @@ def _independent_is_coarser( return nesting in {"leftInRight", "equivalent"} +def _valid_value_mask(values: Sequence[Any] | np.ndarray | pd.Series) -> np.ndarray: + missing = pd.isna(np.asarray(values, dtype=object)) + return np.asarray(~missing, dtype=bool) + + +def _ordered_group_values(values: pd.Series) -> list[str | int | float | bool]: + raw = values.to_numpy(dtype=object, copy=False) + valid = _valid_value_mask(raw) + return [ + cast(str | int | float | bool, _json_scalar(value)) + for value in pd.unique(raw[valid]) + ] + + +def _group_counts( + values: pd.Series, +) -> tuple[list[dict[str, Any]], bool]: + raw = values.to_numpy(dtype=object, copy=False) + valid = _valid_value_mask(raw) + counts, truncated = _bounded_level_counts(raw[valid]) + return list(counts), truncated + + +def _imbalance_from_counts(counts: Sequence[dict[str, Any]]) -> dict[str, Any]: + positive = [int(item["count"]) for item in counts if int(item["count"]) > 0] + if not positive: + return { + "minimum": 0, + "maximum": 0, + "maxToMinRatio": None, + "balanced": False, + } + minimum = min(positive) + maximum = max(positive) + return { + "minimum": minimum, + "maximum": maximum, + "maxToMinRatio": float(maximum / minimum), + "balanced": minimum == maximum, + } + + +def _distinct_units_by_group( + design: pd.DataFrame, + *, + coefficient: str, + unit: str, + group_order: Sequence[str | int | float | bool], +) -> list[dict[str, Any]]: + output: list[dict[str, Any]] = [] + coefficient_values = design[coefficient].to_numpy(dtype=object, copy=False) + unit_values = design[unit].to_numpy(dtype=object, copy=False) + for group in group_order: + mask = coefficient_values == group + valid_units = unit_values[mask & _valid_value_mask(unit_values)] + output.append( + { + "group": group, + "count": int(pd.Series(valid_units).nunique(dropna=True)), + } + ) + return output + + +def _design_missingness( + design: pd.DataFrame, + columns: Sequence[str], +) -> list[dict[str, Any]]: + output: list[dict[str, Any]] = [] + rows = int(len(design)) + for column in dict.fromkeys(columns): + values = design[column].to_numpy(dtype=object, copy=False) + missing = int((~_valid_value_mask(values)).sum()) + output.append( + { + "column": column, + "rows": rows, + "missing": missing, + "missingFraction": missing / rows if rows else 0.0, + } + ) + return output + + +def _categorical_structure( + design: pd.DataFrame, + left: str, + right: str, +) -> dict[str, Any]: + left_values = design[left].to_numpy(dtype=object, copy=False) + right_values = design[right].to_numpy(dtype=object, copy=False) + mapping = directional_mapping(left_values, right_values) + valid = _valid_value_mask(left_values) & _valid_value_mask(right_values) + fully_crossed = False + if valid.any(): + table = pd.crosstab(left_values[valid], right_values[valid]) + fully_crossed = bool(table.size and np.all(table.to_numpy() > 0)) + nesting = str(mapping.get("nesting", "none")) + relationship = ( + nesting + if nesting != "none" + else "crossed" + if fully_crossed + else "partiallyCrossed" + ) + return { + "left": left, + "right": right, + "relationship": relationship, + "fullyCrossed": fully_crossed, + "directionalMapping": mapping, + } + + +def _paired_coverage( + design: pd.DataFrame, + *, + coefficient: str, + pair_by: str, + group_order: Sequence[str | int | float | bool], +) -> dict[str, Any]: + groups = design[coefficient].to_numpy(dtype=object, copy=False) + pairs = design[pair_by].to_numpy(dtype=object, copy=False) + valid = _valid_value_mask(groups) & _valid_value_mask(pairs) + frame = pd.DataFrame({"group": groups[valid], "pair": pairs[valid]}) + required = set(group_order) + complete = 0 + incomplete = 0 + duplicate_pair_groups = 0 + single_group_pairs = 0 + incomplete_examples: list[dict[str, Any]] = [] + for pair, pair_frame in frame.groupby("pair", sort=False, observed=False): + counts = pair_frame["group"].value_counts(dropna=False) + present = {_json_scalar(value) for value in counts.index} + has_duplicate = bool((counts > 1).any()) + if len(present) == 1: + single_group_pairs += 1 + if present == required and not has_duplicate: + complete += 1 + continue + incomplete += 1 + duplicate_pair_groups += int((counts > 1).sum()) + if len(incomplete_examples) < _MAX_LEVEL_COUNT_ITEMS: + incomplete_examples.append( + { + "pair": _json_scalar(pair), + "presentGroups": [ + value for value in group_order if value in present + ], + "duplicateGroups": [ + _json_scalar(value) + for value, count in counts.items() + if int(count) > 1 + ], + } + ) + total_pairs = int(frame["pair"].nunique(dropna=True)) + return { + "pairBy": pair_by, + "requiredGroups": list(group_order), + "pairs": total_pairs, + "completePairs": complete, + "incompletePairs": incomplete, + "duplicatePairGroups": duplicate_pair_groups, + "betweenIndependentUnits": ( + total_pairs >= 2 and single_group_pairs == total_pairs + ), + "design": ( + "paired" + if total_pairs >= 2 + and complete == total_pairs + and incomplete == 0 + and duplicate_pair_groups == 0 + else "betweenIndependentUnits" + if total_pairs >= 2 and single_group_pairs == total_pairs + else "mixedOrIncomplete" + ), + "incompleteExamples": incomplete_examples, + "examplesTruncated": incomplete > len(incomplete_examples), + "complete": ( + total_pairs >= 2 + and complete == total_pairs + and incomplete == 0 + and duplicate_pair_groups == 0 + ), + } + + def _resolve_units( run: _Run, coefficient: str, @@ -941,6 +1188,7 @@ def _characterize_coefficient( observation_unit: str | None, independent_unit: str | None, technical: Sequence[str], + design_columns: Sequence[str] = (), ) -> tuple[dict[str, Any], dict[str, Any] | None]: record: dict[str, Any] = { "name": coefficient, @@ -948,6 +1196,16 @@ def _characterize_coefficient( "observationUnit": observation_unit, "independentUnit": independent_unit, "scope": "unresolvedUnit", + "groupOrder": [], + "groupCounts": [], + "groupCountsTruncated": False, + "groupImbalance": {}, + "unitLevelCounts": {}, + "replication": {}, + "missingness": [], + "designStructures": [], + "pairedCoverage": {}, + "estimability": {}, } if observation_unit is None or observation_unit not in run.profiles: run.note( @@ -1024,7 +1282,28 @@ def _characterize_coefficient( ) independent_unit = None record["independentUnit"] = None - columns = list(dict.fromkeys([*group_cols, coefficient, *unit_constant])) + design_constant = [ + name + for name in design_columns + if name not in {coefficient, observation_unit, independent_unit} + and name in run.profiles + and column_constant_within( + run.store.cells, + name, + observation_unit, + cell_key=run.cell_key, + ) + ] + columns = list( + dict.fromkeys( + [ + *group_cols, + coefficient, + *unit_constant, + *design_constant, + ] + ) + ) design = reduce_observation_units( run.store.cells, observation_unit, @@ -1049,6 +1328,94 @@ def _characterize_coefficient( record["scope"] = "betweenUnit" record["designRows"] = design_rows + group_order = _ordered_group_values(design[coefficient]) + group_counts, group_counts_truncated = _group_counts(design[coefficient]) + record["groupOrder"] = group_order + record["groupCounts"] = group_counts + record["groupCountsTruncated"] = group_counts_truncated + record["groupImbalance"] = _imbalance_from_counts(group_counts) + + unit_level_counts: dict[str, Any] = { + "observationUnit": { + "column": observation_unit, + "levels": design_rows, + }, + "coefficient": { + "column": coefficient, + "levels": len(group_order), + "counts": group_counts, + "truncated": group_counts_truncated, + }, + } + observation_by_group = [ + { + "group": item["value"], + "count": item["count"], + } + for item in group_counts + ] + independent_by_group: list[dict[str, Any]] = [] + if independent_unit is not None: + independent_counts, independent_counts_truncated = _bounded_level_counts( + design[independent_unit] + ) + independent_by_group = _distinct_units_by_group( + design, + coefficient=coefficient, + unit=independent_unit, + group_order=group_order, + ) + unit_level_counts["independentUnit"] = { + "column": independent_unit, + "levels": int(design[independent_unit].nunique(dropna=True)), + "counts": list(independent_counts), + "truncated": independent_counts_truncated, + "byGroup": independent_by_group, + } + record["unitLevelCounts"] = unit_level_counts + replication_counts = ( + independent_by_group if independent_unit is not None else observation_by_group + ) + minimum_replication = min( + (int(item["count"]) for item in replication_counts), + default=0, + ) + record["replication"] = { + "unit": independent_unit or observation_unit, + "observationUnitsByGroup": observation_by_group, + "independentUnitsByGroup": independent_by_group, + "minimumPerGroup": minimum_replication, + "sufficient": ( + len(group_order) >= 2 + and not group_counts_truncated + and minimum_replication >= 2 + ), + } + record["missingness"] = _design_missingness(design, columns) + + structural_columns = list( + dict.fromkeys( + [ + *([independent_unit] if independent_unit is not None else []), + *unit_constant, + *design_constant, + ] + ) + ) + structures: list[dict[str, Any]] = [] + if run.kind(coefficient) == "categorical": + for name in structural_columns[:_MAX_STRUCTURE_ITEMS]: + if run.kind(name) != "categorical": + continue + structures.append(_categorical_structure(design, coefficient, name)) + record["designStructures"] = structures + if independent_unit is not None and len(group_order) >= 2: + record["pairedCoverage"] = _paired_coverage( + design, + coefficient=coefficient, + pair_by=independent_unit, + group_order=group_order, + ) report = report_confounding( design, @@ -1058,10 +1425,20 @@ def _characterize_coefficient( coefficient: run.kind(coefficient), **{name: run.kind(name) for name in unit_constant}, }, - associationFloor=CONFIG._ASSOCIATION_FLOOR, + associationFloor=_ASSOCIATION_FLOOR, ) report["observationUnit"] = observation_unit report["independentUnit"] = independent_unit + report["groupOrder"] = group_order + report["groupCounts"] = group_counts + report["groupCountsTruncated"] = group_counts_truncated + report["groupImbalance"] = record["groupImbalance"] + report["unitLevelCounts"] = unit_level_counts + report["replication"] = record["replication"] + report["missingness"] = record["missingness"] + report["designStructures"] = structures + report["pairedCoverage"] = record["pairedCoverage"] + record["estimability"] = dict(report.get("estimability") or {}) run.actions.append(f"confounding:{coefficient}") return record, report @@ -1090,6 +1467,7 @@ def _characterize_coefficients( observation_unit=observation, independent_unit=independent, technical=technical, + design_columns=design_columns, ) records.append(record) if report is not None: @@ -1140,6 +1518,16 @@ def _column_records( "domain": run.domains[name], "summary": profile.summary, "aliases": list(aliases.get(name, [])), + "nRows": profile.digest.nRows, + "nLevels": profile.digest.nLevels, + "nMissing": profile.digest.nMissing, + "missingFraction": ( + profile.digest.nMissing / profile.digest.nRows + if profile.digest.nRows + else 0.0 + ), + "levelCounts": list(profile.levelCounts), + "levelCountsTruncated": profile.levelCountsTruncated, } if profile.artifact is None: record.update( @@ -1157,14 +1545,45 @@ def _column_records( ) records.append(record) for name, reason in dropped: + dropped_profile = run.profiles.get(name) record = { "name": name, - "kind": "continuous", + "kind": dropped_profile.kind + if dropped_profile is not None + else "continuous", "domain": "ignore", - "summary": f"dropped before triage ({CONFIG._DROP_REASONS[reason]})", + "summary": ( + f"excluded from design roles ({_DROP_REASONS[reason]})" + + ( + f"; {dropped_profile.summary}" + if dropped_profile is not None + else "" + ) + ), "aliases": [], + "nRows": ( + dropped_profile.digest.nRows if dropped_profile is not None else 0 + ), + "nLevels": ( + dropped_profile.digest.nLevels if dropped_profile is not None else 0 + ), + "nMissing": ( + dropped_profile.digest.nMissing if dropped_profile is not None else 0 + ), + "missingFraction": ( + dropped_profile.digest.nMissing / dropped_profile.digest.nRows + if dropped_profile is not None and dropped_profile.digest.nRows + else 0.0 + ), + "levelCounts": ( + list(dropped_profile.levelCounts) if dropped_profile is not None else [] + ), + "levelCountsTruncated": ( + dropped_profile.levelCountsTruncated + if dropped_profile is not None + else False + ), } - dropped_profile = run.profiles.get(name) if dropped_profile is not None and dropped_profile.artifact is not None: record.update( { @@ -1191,6 +1610,7 @@ def characterize_covariates( model: Any | None = None, directions: Mapping[str, Any] | None = None, groupingArtifacts: Mapping[str, ArtifactRef] | None = None, + inventory: dict[str, Any] | None = None, ) -> CovariateCharacterization: """Label cell covariates and record design-level confounding.""" if ( @@ -1204,6 +1624,12 @@ def characterize_covariates( cellSelection, artifacts=grouping_artifacts, ) + if inventory is not None and inventory.get("cellSelection") != cellSelection: + inventory.clear() + inventory["cellSelection"] = cellSelection + column_inventory = ( + inventory.setdefault("columns", {}) if inventory is not None else None + ) cell_key = "I" direction_map = dict(directions or {}) available = set(bound_store.cells.columns) @@ -1234,9 +1660,8 @@ def characterize_covariates( bound_store, name, cell_key=cell_key, - kind=cast(ColumnKind, directed_kind) - if directed_kind in CONFIG._KINDS - else None, + kind=cast(ColumnKind, directed_kind) if directed_kind in _KINDS else None, + inventory=column_inventory, ) profiles[name] = profile n_rows = profile.digest.nRows @@ -1245,6 +1670,15 @@ def characterize_covariates( continue varying.append(name) candidates = varying + for name, reason in dropped: + if reason == "dropAssayStat" and name not in profiles: + profiles[name] = _profile_column( + bound_store, + name, + cell_key=cell_key, + kind="continuous", + inventory=column_inventory, + ) candidates, aliases, alias_notes = _collapse_ontology_aliases( bound_store, candidates, @@ -1256,12 +1690,12 @@ def characterize_covariates( store=bound_store, cell_key=cell_key, n_rows=n_rows, - context=_bounded_context(studyContext), + context=(studyContext or "").strip(), model=model, profiles=profiles, ) for name, reason in dropped: - detail = f"Dropped {CONFIG._DROP_REASONS[reason]} {name}" + detail = f"Dropped {_DROP_REASONS[reason]} {name}" if reason == "dropConstant" and name in directed_coefficients: detail = ( f"{detail}; also listed in coefficientsOfInterest but has no variation" @@ -1300,6 +1734,12 @@ def characterize_covariates( categorical_technical = [ name for name in technical if run.kind(name) == "categorical" ] + column_records = _column_records( + run, + candidates, + aliases=aliases, + dropped=dropped, + ) return CovariateCharacterization( status="done", cellSelection=artifact_reference(cellSelection), @@ -1310,7 +1750,7 @@ def characterize_covariates( "deterministic drops and ontology alias collapse" ], decisions=run.decisions, - columns=_column_records(run, candidates, aliases=aliases, dropped=dropped), + columns=column_records, coefficients=records, technicalNesting=( _technical_nesting_reports( @@ -1322,4 +1762,53 @@ def characterize_covariates( else [] ), confounding=reports, + unitLevelCounts=[ + { + "coefficient": record["name"], + **dict(record["unitLevelCounts"]), + } + for record in records + if record.get("unitLevelCounts") + ], + groupImbalance=[ + { + "coefficient": record["name"], + **dict(record["groupImbalance"]), + } + for record in records + if record.get("groupImbalance") + ], + missingness=[ + { + "column": record["name"], + "rows": record.get("nRows", 0), + "missing": record.get("nMissing", 0), + "missingFraction": record.get("missingFraction", 0.0), + } + for record in column_records + ], + designStructures=[ + { + "coefficient": record["name"], + **structure, + } + for record in records + for structure in record.get("designStructures", []) + ], + pairedCoverage=[ + { + "coefficient": record["name"], + **dict(record["pairedCoverage"]), + } + for record in records + if record.get("pairedCoverage") + ], + coefficientEstimability=[ + { + "coefficient": record["name"], + **dict(record["estimability"]), + } + for record in records + if record.get("estimability") + ], ) diff --git a/scarf/agent/experimental_context/comparisons.py b/scarf/agent/experimental_context/comparisons.py new file mode 100644 index 00000000..d5ee77f2 --- /dev/null +++ b/scarf/agent/experimental_context/comparisons.py @@ -0,0 +1,633 @@ +"""Bounded, objective-led comparisons of metadata on independent study units.""" + +import hashlib +import json +import re +from collections.abc import Sequence +from typing import Any, Literal, cast + +import numpy as np +import pandas as pd + +from ...metrics.association import association_pair, coefficient_estimability +from .. import record_io +from .characterization import _json_scalar, _paired_coverage +from .contracts import ( + CaptureProposal, + CovariateCharacterization, + CovariateComparison, + CovariateProposal, + ExperimentalContextDependencies, +) + +DESIGN_ROUND_LIMITS = (8, 4) +MAX_COMBINATIONS = 32 +MAX_STRATA = 16 +MAX_DESCRIPTIVE_ROWS = 128 + + +def combination_labels(cells: Any, columns: Sequence[str]) -> np.ndarray: + """Encode exact, already selection-aligned metadata without writing columns.""" + if len(columns) != 2 or len(set(columns)) != 2: + raise ValueError("A protected combination requires two distinct columns") + arrays = [np.asarray(cells.fetch(column)) for column in columns] + if any(array.ndim != 1 or array.shape != arrays[0].shape for array in arrays): + raise ValueError("Combination columns must align to the same cell selection") + return _tuple_labels(arrays) + + +def _tuple_labels(arrays: Sequence[np.ndarray]) -> np.ndarray: + output: list[str] = [] + for row in zip(*arrays, strict=True): + encoded = [] + for value in row: + value = value.item() if isinstance(value, np.generic) else value + if pd.isna(value) or isinstance(value, float) and not np.isfinite(value): + raise ValueError("Combination columns contain missing values") + if isinstance(value, bytes): + value = value.decode("utf-8") + encoded.append([type(value).__name__, value]) + output.append(json.dumps(encoded, ensure_ascii=False, separators=(",", ":"))) + return np.asarray(output, dtype=str) + + +def _proposal_key(proposal: CovariateProposal) -> str: + return hashlib.sha256( + record_io.canonical_json_bytes(proposal.model_dump(exclude={"rationale"})) + ).hexdigest() + + +def _association( + frame: pd.DataFrame, response: str, explanatory: str, kinds: dict[str, Any] +) -> dict[str, Any]: + if len(frame) < 4: + return {"status": "notComputed", "reason": "fewerThanFourIndependentUnits"} + for column in (response, explanatory): + if kinds[column] == "categorical" and ( + frame[column].nunique() < 2 or frame[column].value_counts().min() < 2 + ): + return { + "status": "notComputed", + "reason": "categoricalGroupsRequireTwoIndependentUnits", + } + return association_pair( + frame[response].to_numpy(), + frame[explanatory].to_numpy(), + leftKind=kinds[response], + rightKind=kinds[explanatory], + ) + + +def _descriptive_design( + design: pd.DataFrame, + *, + columns: Sequence[str], + independent: str, + kinds: dict[str, Any], +) -> dict[str, Any]: + """Count observations and distinct units without treating repeats as replicates.""" + categorical = list( + dict.fromkeys(name for name in columns if kinds[name] == "categorical") + ) + output: dict[str, Any] = { + "status": "computed", + "observationUnits": len(design), + "independentUnits": int(design[independent].nunique()), + "independentUnit": independent, + "interpretation": ( + "Descriptive counts only. A unit can occur in multiple groups; group " + "memberships are not disjoint independent replicates or effect estimates." + ), + "groupSupport": {}, + "pairedCoverage": {}, + "sharedIndependentUnits": {}, + } + if any(design[name].nunique() > MAX_COMBINATIONS for name in categorical): + output.update(status="unsupported", reason="moreThanThirtyTwoCategoricalLevels") + return output + for name in categorical: + rows: list[dict[str, Any]] = [] + memberships: list[tuple[Any, set[Any]]] = [] + for label, subset in design.groupby(name, sort=False, observed=True): + memberships.append((_json_scalar(label), set(subset[independent]))) + rows.append( + { + "group": _json_scalar(label), + "observationUnits": len(subset), + "independentUnits": int(subset[independent].nunique()), + } + ) + output["groupSupport"][name] = rows + overlap = [ + { + "groups": [left, right], + "independentUnits": len(left_units & right_units), + } + for index, (left, left_units) in enumerate(memberships) + for right, right_units in memberships[index + 1 :] + ] + output["sharedIndependentUnits"][name] = { + "pairs": overlap[:MAX_DESCRIPTIVE_ROWS], + "truncated": len(overlap) > MAX_DESCRIPTIVE_ROWS, + } + if name != independent and len(rows) >= 2: + output["pairedCoverage"][name] = _paired_coverage( + design, + coefficient=name, + pair_by=independent, + group_order=list(design[name].unique()), + ) + if categorical: + groups = design.groupby(categorical, sort=False, observed=True) + if groups.ngroups > MAX_DESCRIPTIVE_ROWS: + output.update(status="unsupported", reason="moreThan128DescriptiveGroups") + else: + rows = [] + for labels, subset in groups: + if not isinstance(labels, tuple): + labels = (labels,) + rows.append( + { + "groups": dict( + zip(categorical, map(_json_scalar, labels), strict=True) + ), + "observationUnits": len(subset), + "independentUnits": int(subset[independent].nunique()), + } + ) + output["jointGroupSupport"] = rows + output["continuousSummaries"] = { + name: { + "minimum": float(design[name].min()), + "median": float(design[name].median()), + "maximum": float(design[name].max()), + "unit": "observationUnit", + } + for name in columns + if kinds[name] == "continuous" + } + return output + + +def compare_covariates( + cells: Any, + characterization: CovariateCharacterization, + proposal: CovariateProposal, + *, + selection_identity: dict[str, Any], +) -> CovariateComparison: + """Compute descriptive evidence; unsupported designs remain explicit.""" + evidence: dict[str, Any] = {} + reasons: list[str] = [] + records = {record["name"]: record for record in characterization.columns} + columns = [proposal.response, *proposal.explanatoryColumns] + if proposal.conditionedOn is not None: + columns.append(proposal.conditionedOn) + unit = proposal.observationUnit + independent = proposal.independentUnit or unit + requested = list(dict.fromkeys([*columns, unit, independent])) + evidence["columnKinds"] = { + name: records.get(name, {}).get("kind") for name in requested + } + evidence["columnDomains"] = { + name: records.get(name, {}).get("domain") for name in requested + } + declared_units = { + ( + record.get("observationUnit"), + record.get("independentUnit") or record.get("observationUnit"), + ) + for record in characterization.coefficients + } + declared_pair = (unit, independent) in declared_units + evidence["unitRoles"] = { + "observationUnit": unit, + "independentUnit": independent, + "declaredInCharacterization": declared_pair, + } + if any(name not in cells.columns or name not in records for name in requested): + reasons.append("unknownObservedColumn") + elif any( + records[name].get("kind") != "categorical" for name in {unit, independent} + ): + reasons.append("observationAndIndependentUnitsMustBeCategorical") + elif any( + records[name].get("domain") not in {"design", "technical"} + and not (declared_pair and records[name].get("domain") == "biological") + for name in {unit, independent} + ): + reasons.append("observationAndIndependentUnitsMustBeDesignOrTechnical") + if not reasons: + kinds = {name: records[name].get("kind", "categorical") for name in columns} + values = {name: np.asarray(cells.fetch(name)) for name in requested} + frame = pd.DataFrame(values) + evidence["cells"] = len(frame) + evidence["missingCellsByColumn"] = {} + complete = np.ones(len(frame), dtype=bool) + for name in requested: + valid = frame[name].notna().to_numpy() + if kinds.get(name) == "continuous": + numeric = pd.to_numeric(frame[name], errors="coerce") + valid = valid & np.isfinite(numeric.to_numpy(dtype=float)) + frame[name] = numeric + evidence["missingCellsByColumn"][name] = int((~valid).sum()) + complete &= valid + evidence["missingCells"] = int((~complete).sum()) + frame = frame.loc[complete] + if frame.empty: + reasons.append("noCompleteObservations") + else: + grouped = frame.groupby(unit, sort=False, observed=True) + constant = [*proposal.explanatoryColumns, independent] + if proposal.conditionedOn is not None: + constant.append(proposal.conditionedOn) + if kinds[proposal.response] == "categorical": + constant.append(proposal.response) + if any(grouped[name].nunique().gt(1).any() for name in set(constant)): + reasons.append("explanatoryColumnsMustBeConstantWithinObservationUnit") + if grouped.ngroups >= len(frame): + reasons.append("observationUnitIsCellIdentifier") + design = grouped[requested].first().reset_index(drop=True) + if kinds[proposal.response] == "continuous": + design[proposal.response] = ( + grouped[proposal.response].median().to_numpy() + ) + evidence["responseAggregation"] = "medianPerObservationUnit" + evidence["observationUnits"] = len(design) + if not reasons: + evidence["descriptiveDesign"] = _descriptive_design( + design, columns=columns, independent=independent, kinds=kinds + ) + if independent != unit: + grouped_independent = design.groupby( + independent, sort=False, observed=True + ) + constant = [*proposal.explanatoryColumns] + if proposal.conditionedOn is not None: + constant.append(proposal.conditionedOn) + if kinds[proposal.response] == "categorical": + constant.append(proposal.response) + if any( + grouped_independent[name].nunique().gt(1).any() for name in constant + ): + reasons.append("withinIndependentUnitComparisonsAreUnsupported") + reduced = grouped_independent[columns].first().reset_index(drop=True) + if kinds[proposal.response] == "continuous": + reduced[proposal.response] = ( + grouped_independent[proposal.response].median().to_numpy() + ) + evidence["independentAggregation"] = "medianOfObservationMedians" + design = reduced + evidence["independentUnits"] = len(design) + evidence["unitOfComparison"] = independent + if len(design) < 4: + reasons.append("fewerThanFourIndependentUnits") + if any( + kinds[name] == "categorical" + and design[name].nunique() > MAX_COMBINATIONS + for name in columns + ): + reasons.append("moreThanThirtyTwoCategoricalLevels") + if not reasons: + if proposal.conditionedOn is not None: + condition = proposal.conditionedOn + if kinds[condition] != "categorical": + reasons.append("continuousConditioningIsUnsupported") + elif design[condition].nunique() > MAX_STRATA: + reasons.append("moreThanSixteenConditioningStrata") + else: + strata: list[dict[str, Any]] = [] + for label, subset in design.groupby( + condition, sort=False, observed=True + ): + result = _association( + subset, + proposal.response, + proposal.explanatoryColumns[0], + kinds, + ) + strata.append( + { + "stratum": str(label), + "independentUnits": len(subset), + "association": result, + } + ) + evidence["strata"] = strata + if any( + row["association"].get("status") != "ok" for row in strata + ): + reasons.append("unsupportedConditioningStrata") + else: + evidence["singleAssociations"] = { + name: _association(design, proposal.response, name, kinds) + for name in proposal.explanatoryColumns + } + if any( + value.get("status") != "ok" + for value in evidence["singleAssociations"].values() + ): + reasons.append("unsupportedSingleAssociation") + if len(proposal.explanatoryColumns) == 2: + evidence["jointEstimability"] = coefficient_estimability( + design[proposal.response].to_numpy(), + coefficientKind=cast( + Literal["categorical", "continuous"], + kinds[proposal.response], + ), + technicals={ + name: design[name].to_numpy() + for name in proposal.explanatoryColumns + }, + technicalKinds={ + name: kinds[name] + for name in proposal.explanatoryColumns + }, + ) + if all( + kinds[name] == "categorical" + for name in proposal.explanatoryColumns + ): + labels = _tuple_labels( + [ + design[name].to_numpy() + for name in proposal.explanatoryColumns + ] + ) + if len(np.unique(labels)) > MAX_COMBINATIONS: + reasons.append("moreThanThirtyTwoJointGroups") + else: + design = design.assign(_joint=labels) + evidence["jointGroupCounts"] = { + str(key): int(value) + for key, value in design["_joint"] + .value_counts() + .items() + } + evidence["jointAssociation"] = _association( + design, + proposal.response, + "_joint", + {**kinds, "_joint": "categorical"}, + ) + evidence["jointGroupEstimability"] = ( + coefficient_estimability( + design[proposal.response].to_numpy(), + coefficientKind=cast( + Literal["categorical", "continuous"], + kinds[proposal.response], + ), + technicals={"joint": labels}, + technicalKinds={"joint": "categorical"}, + ) + ) + if evidence["jointAssociation"].get("status") != "ok": + reasons.append("unsupportedJointAssociation") + if proposal.protectCombination: + if any( + records.get(name, {}).get("domain") != "biological" + or records.get(name, {}).get("kind") != "categorical" + for name in proposal.explanatoryColumns + ): + reasons.append("protectedCombinationsMustBeCategoricalBiology") + evidence_id = ( + "designComparison:" + + hashlib.sha256( + record_io.canonical_json_bytes( + { + "selection": selection_identity, + "proposal": proposal.model_dump(), + "evidence": evidence, + "reasons": reasons, + } + ) + ).hexdigest() + ) + return CovariateComparison( + proposal=proposal, + status="unsupported" if reasons else "computed", + evidence=evidence, + reasons=list(dict.fromkeys(reasons)), + evidenceId=evidence_id, + ) + + +def evaluate_proposals( + deps: ExperimentalContextDependencies, + characterization: CovariateCharacterization, + proposals: Sequence[CovariateProposal], +) -> None: + """Consume one bounded round, retaining comparisons from earlier rounds.""" + if deps.designRounds >= len(DESIGN_ROUND_LIMITS): + raise ValueError("Design comparison permits at most two evidence rounds") + if len(proposals) > DESIGN_ROUND_LIMITS[deps.designRounds]: + raise ValueError( + "Design comparison permits eight initial and four follow-up proposals" + ) + for proposal in proposals: + if deps.studyObjective and ( + not proposal.objectiveQuote or "purpose" not in proposal.model_fields_set + ): + raise ValueError( + "Objective comparisons require an exact objectiveQuote and explicit purpose" + ) + if proposal.objectiveQuote and proposal.objectiveQuote not in ( + f"{deps.studyContext}\n{deps.studyObjective}" + ): + raise ValueError("Comparison objectiveQuote must copy exact study text") + if ( + proposal.objectiveQuote + and proposal.objectiveQuote in deps.studyObjective + and not proposal.essential + ): + raise ValueError( + "A question quoting the explicit study objective must remain essential" + ) + deps.designRounds += 1 + previous = {_proposal_key(item.proposal): item for item in deps.comparisons} + records = {record["name"]: record for record in characterization.columns} + for proposal in proposals: + key = _proposal_key(proposal) + prior = previous.get(key) + if prior is not None and any( + records.get(column, {}).get(field) != value + for field, recorded in ( + ("kind", "columnKinds"), + ("domain", "columnDomains"), + ) + for column, value in prior.evidence.get(recorded, {}).items() + ): + del previous[key] + if prior is not None and key in previous: + declared_pairs = { + ( + record.get("observationUnit"), + record.get("independentUnit") or record.get("observationUnit"), + ) + for record in characterization.coefficients + } + declared = ( + proposal.observationUnit, + proposal.independentUnit or proposal.observationUnit, + ) in declared_pairs + if ( + prior.evidence.get("unitRoles", {}).get("declaredInCharacterization") + != declared + ): + del previous[key] + if key not in previous: + if prior is not None: + deps.comparisons.remove(prior) + comparison = compare_covariates( + deps.cells, + characterization, + proposal, + selection_identity=deps.cellSelection.to_dict(), + ) + deps.comparisons.append(comparison) + previous[key] = comparison + deps.protectedCombinations = [] + for comparison in deps.comparisons: + if comparison.proposal.protectCombination: + columns = sorted(comparison.proposal.explanatoryColumns) + if any( + records.get(column, {}).get("domain") != "biological" + or records.get(column, {}).get("kind") != "categorical" + for column in columns + ): + continue + if columns not in deps.protectedCombinations: + deps.protectedCombinations.append(columns) + characterization.comparisons = list(deps.comparisons) + characterization.captureProvenance = deps.captureProposal + + +def accept_capture_proposal( + deps: ExperimentalContextDependencies, + characterization: CovariateCharacterization, + proposal: CaptureProposal, +) -> None: + """Require an observed capture and exact supporting study statements.""" + study = f"{deps.studyContext}\n{deps.studyObjective}" + quote = proposal.provenanceQuote + records = {record["name"]: record for record in characterization.columns} + record = records.get(proposal.column, {}) + if ( + proposal.column not in deps.cells.columns + or record.get("kind") != "categorical" + or record.get("domain") not in {"design", "technical"} + ): + raise ValueError( + "Capture proposal must name an observed categorical design column" + ) + if ( + quote not in study + or proposal.column not in quote + or "capture" not in quote.lower() + or re.search(r"\b(?:not|unknown|unresolved|uncertain)\b", quote, re.I) + ): + raise ValueError( + "Capture proposal requires an exact study quote identifying that column as a physical capture" + ) + directed = deps.directions.get("physicalCaptureColumn") + if directed is not None and directed != proposal.column: + raise ValueError( + "Capture proposal conflicts with the caller's physical capture" + ) + references = proposal.referenceCaptures + if references: + labels = np.asarray(deps.cells.fetch(proposal.column)).astype(str) + reference_quote = proposal.referenceProvenanceQuote + if ( + len(references) < 2 + or len(references) != len(set(references)) + or not set(references).issubset(set(labels)) + ): + raise ValueError( + "Reference pool requires at least two distinct observed captures" + ) + if ( + not reference_quote + or reference_quote not in study + or any(name not in reference_quote for name in references) + or not any( + word in reference_quote.lower() + for word in ("reference", "baseline", "control") + ) + or re.search( + r"\b(?:not|unknown|unresolved|uncertain)\b", reference_quote, re.I + ) + ): + raise ValueError( + "Reference captures require an exact study quote supporting their baseline role" + ) + from .qc_evidence import ( + _directed_capture_source, + _directed_pooled_reference_captures, + ) + + proposed_deps = deps.model_copy(update={"captureProposal": proposal}) + _directed_capture_source(proposed_deps) + _directed_pooled_reference_captures(proposed_deps) + deps.captureProposal = proposal + characterization.captureProvenance = proposal + + +def canonical_design_choices( + deps: ExperimentalContextDependencies, decision: Any +) -> dict[str, Any]: + """Expose only combinations and capture choices supported by tool evidence.""" + from .qc_evidence import ( + _directed_capture_source, + _directed_pooled_reference_captures, + ) + + combinations = sorted(deps.protectedCombinations) + records = { + record["name"]: record + for record in ( + deps.characterization.columns if deps.characterization is not None else [] + ) + } + if any( + records.get(column, {}).get("domain") != "biological" + or records.get(column, {}).get("kind") != "categorical" + for columns in combinations + for column in columns + ): + raise ValueError( + "Protected combinations must remain categorical biological columns" + ) + supplied = sorted(sorted(columns) for columns in decision.protectedCombinations) + if supplied and supplied != combinations: + raise ValueError( + "Protected combinations must match the evaluated objective-led proposals" + ) + capture = _directed_capture_source(deps) + capture_name = None + if capture is not None: + capture_name = capture[0] + if capture_name is None and capture[1] is not None: + capture_name = capture[1].name + references = list(_directed_pooled_reference_captures(deps) or ()) + if ( + decision.physicalCaptureColumn not in {None, capture_name} + or decision.pooledReferenceCaptures + and decision.pooledReferenceCaptures != references + ): + raise ValueError("Capture choices must match provenance-backed tool evidence") + return { + "protectedCombinations": combinations, + "physicalCaptureColumn": capture_name, + "pooledReferenceCaptures": references, + "unsupportedProtection": sorted( + { + column + for column in [ + *decision.coefficientsOfInterest, + *decision.batchCorrection.preserveColumns, + ] + if records.get(column, {}).get("kind") == "continuous" + } + ), + } diff --git a/scarf/agent/experimental_context/contracts.py b/scarf/agent/experimental_context/contracts.py new file mode 100644 index 00000000..11d88ef5 --- /dev/null +++ b/scarf/agent/experimental_context/contracts.py @@ -0,0 +1,915 @@ +"""Experimental-context contracts and handoffs.""" + +import math +from typing import Any, Literal + +from .._deps import AGENT_INSTALL_HINT +from ..cell_quality.profiles import QcMetricRole, RegisteredCellQcProfile +from ..types import ( + AgentDataModel, + AgentRunInfo, + ArtifactReferenceModel, + BatchCorrectionAction, + BatchSafetyEvidence, + ExperimentalBiologyHandoff, + ExperimentalTuningHandoff, + StageStatus, +) + +try: + from pydantic import ConfigDict, Field, model_validator + from pydantic.json_schema import SkipJsonSchema +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + +type ColumnDomain = Literal["biological", "technical", "design", "ignore", "unknown"] +type IntegrationMetric = Literal[ + "iLISI", + "cLISI", + "graphConnectivity", + "proportionalBatchMixing", +] +type CellQcAction = Literal[ + "skip", + "globalGaussian", + "sampleMad", + "registeredMad", +] +type LegacyCellQcAction = Literal["skip", "globalGaussian", "sampleMad"] +type CellQcDriverType = Literal["RNA", "ATAC"] + + +type ContrastTest = Literal["mann_whitney", "kruskal_wallis", "wilcoxon"] +type ContrastSampleStatistic = Literal["mean", "median", "fraction"] +type ContrastStatus = Literal["licensed", "blocked", "needsInput"] + + +class CovariateProposal(AgentDataModel): + """Compare at most three distinct measured columns, without expression tests. + + Use one response with either one or two explanatory columns and no conditioning, + or one response with one explanatory column and one categorical conditioning + column. Observation and independent units are separate and do not count toward + this limit. Two explanatory columns plus conditioning are unsupported; one + explanatory column within a categorical stratum is an offered comparison. + """ + + response: str = Field( + description=( + "Exact observed outcome column. It must differ from every explanatory " + "and conditioning column; do not compare a column with itself." + ) + ) + explanatoryColumns: list[str] = Field( + min_length=1, + max_length=2, + description=( + "One or two distinct observed columns, each different from response. " + "With two explanatory columns, conditionedOn must be null. With a " + "conditioning column, supply exactly one explanatory column. Unit " + "identifiers belong in observationUnit/independentUnit unless they " + "are themselves the explicitly requested scientific comparison." + ), + ) + conditionedOn: str | None = Field( + default=None, + description=( + "Optional exact categorical column defining strata, different from " + "response and the single explanatory column. Set null for a two-column " + "joint explanation. Continuous conditioning is unsupported." + ), + ) + observationUnit: str = Field( + description=( + "Exact observed column identifying the observation unit, such as a " + "sample. This unit field does not count toward the three measured " + "columns; it does not automatically belong in explanatoryColumns." + ) + ) + independentUnit: str | None = Field( + default=None, + description=( + "Exact observed independent-unit column, such as donor, or null to " + "use observationUnit. It may equal observationUnit and does not count " + "toward the three measured columns. Preserve repeated-donor identity." + ), + ) + rationale: str = Field(min_length=1) + protectCombination: bool = False + purpose: Literal["designCoverage", "association", "effectEstimation"] = ( + "designCoverage" + ) + essential: bool = True + objectiveQuote: str = "" + + @model_validator(mode="after") + def validate_columns(self) -> "CovariateProposal": + columns = [self.response, *self.explanatoryColumns] + if self.conditionedOn is not None: + columns.append(self.conditionedOn) + repeated = sorted({name for name in columns if columns.count(name) > 1}) + if repeated: + raise ValueError( + "Comparison columns must be distinct; repeated columns " + f"{repeated!r} occur in response={self.response!r}, " + f"explanatoryColumns={self.explanatoryColumns!r}, " + f"conditionedOn={self.conditionedOn!r}. A column cannot explain " + "itself or also define its conditioning strata. Choose the actual " + "distinct measured columns. Observation and independent units may " + "share names and are not part of this uniqueness check." + ) + if len(columns) > 3: + raise ValueError( + "A comparison requires at most three distinct measured columns; " + f"received response={self.response!r}, " + f"explanatoryColumns={self.explanatoryColumns!r}, " + f"conditionedOn={self.conditionedOn!r}. With two explanatory " + "columns set conditionedOn=null for a joint comparison. A " + "within-stratum comparison requires one explanatory column. " + "Observation and independent units do not count toward this " + "limit. A joint explanation within strata is unsupported; do not " + "claim that separate simpler comparisons answer that question." + ) + if any(not value.strip() for value in [*columns, self.observationUnit]): + raise ValueError( + "Comparison columns and observation unit must be non-empty" + ) + if self.protectCombination and len(self.explanatoryColumns) != 2: + raise ValueError("A protected combination requires two explanatory columns") + return self + + +class CovariateComparison(AgentDataModel): + proposal: CovariateProposal + status: Literal["computed", "unsupported"] + evidence: dict[str, Any] = Field(default_factory=dict) + reasons: list[str] = Field(default_factory=list) + evidenceId: str + + +class DesignEvidenceRequirement(AgentDataModel): + """One objective question whose completion is checked against measured evidence.""" + + requirementId: str = Field(min_length=1) + question: str = Field(min_length=1) + objectiveQuote: str = Field(min_length=1) + kind: Literal["studyDesign", "designCoverage", "association", "effectEstimation"] + columns: list[str] = Field(default_factory=list) + observationUnit: str | None = None + independentUnit: str | None = None + essential: bool = True + + +class DesignEvidenceCoverage(AgentDataModel): + """Measured answer to a requirement, without assigning a second workflow status.""" + + requirementId: str + status: Literal["computed", "nonIdentifiable", "unsupported", "failed"] + evidenceIds: list[str] = Field(default_factory=list) + reasons: list[str] = Field(default_factory=list) + + +class CaptureProposal(AgentDataModel): + """An exact capture column and optional references supported by study prose.""" + + column: str + provenanceQuote: str = Field(min_length=1) + referenceCaptures: list[str] = Field(default_factory=list, max_length=32) + referenceProvenanceQuote: str = "" + + +class CovariateCharacterization(AgentDataModel): + status: StageStatus + cellSelection: ArtifactReferenceModel | None = None + auditLog: list[dict[str, Any]] = Field(default_factory=list) + actions: list[str] = Field(default_factory=list) + notes: list[str] = Field(default_factory=list) + decisions: list[dict[str, Any]] = Field(default_factory=list) + columns: list[dict[str, Any]] = Field(default_factory=list) + coefficients: list[dict[str, Any]] = Field(default_factory=list) + technicalNesting: list[dict[str, Any]] = Field(default_factory=list) + confounding: list[dict[str, Any]] = Field(default_factory=list) + unitLevelCounts: list[dict[str, Any]] = Field(default_factory=list) + groupImbalance: list[dict[str, Any]] = Field(default_factory=list) + missingness: list[dict[str, Any]] = Field(default_factory=list) + designStructures: list[dict[str, Any]] = Field(default_factory=list) + pairedCoverage: list[dict[str, Any]] = Field(default_factory=list) + coefficientEstimability: list[dict[str, Any]] = Field(default_factory=list) + comparisons: list[CovariateComparison] = Field(default_factory=list) + captureProvenance: CaptureProposal | None = None + + @classmethod + def get_blank(cls) -> "CovariateCharacterization": + return cls(status="failed") + + +class InferenceUnit(AgentDataModel): + """Observation and independent units for one biological coefficient.""" + + observationUnit: str | None = None + independentUnit: str | None = None + + @classmethod + def get_blank(cls) -> "InferenceUnit": + return cls() + + +class BatchCorrectionPlan(AgentDataModel): + """A grounded recommendation about whether Harmony should be evaluated.""" + + action: BatchCorrectionAction + batchColumns: list[str] = Field(default_factory=list) + preserveColumns: list[str] = Field(default_factory=list) + metricsRequired: list[IntegrationMetric] = Field(default_factory=list) + rationale: str = "" + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "BatchCorrectionPlan": + return cls(action="needsInput") + + +class NamedArtifactSource(AgentDataModel): + """One semantic name bound to an exact immutable artifact.""" + + name: str = "" + artifact: ArtifactReferenceModel = Field(default_factory=ArtifactReferenceModel) + + @model_validator(mode="after") + def validate_source(self) -> "NamedArtifactSource": + if self.name != self.name.strip(): + raise ValueError("Artifact source names cannot have surrounding whitespace") + if bool(self.name.strip()) != bool(self.artifact.artifactId): + raise ValueError("A named artifact source requires both name and artifact") + return self + + @classmethod + def get_blank(cls) -> "NamedArtifactSource": + return cls() + + +class QcMetricSourceEvidence(AgentDataModel): + """One source-specific quality metric on the exact active cells.""" + + sourceId: str = "" + metricName: str = "" + metricRole: QcMetricRole = "diagnostic" + assay: str | None = None + sourceType: Literal["metadataColumn", "artifact"] = "metadataColumn" + origin: Literal[ + "ingestionMetadata", + "derivedArtifact", + "externalArtifact", + ] = "ingestionMetadata" + executionName: str = "" + metadataColumn: str | None = None + artifact: ArtifactReferenceModel | None = None + cellSelection: ArtifactReferenceModel | None = None + inputArtifacts: list[ArtifactReferenceModel] = Field(default_factory=list) + provenanceOperation: str | None = None + valuesFingerprint: str = "" + activeCells: int = 0 + missingCells: int = 0 + missingCellsByCapture: dict[str, int] = Field(default_factory=dict) + usableForFiltering: bool = False + notes: list[str] = Field(default_factory=list) + + @model_validator(mode="after") + def validate_source(self) -> "QcMetricSourceEvidence": + if ( + not self.sourceId + and not self.metricName + and self.metadataColumn is None + and self.artifact is None + ): + return self + if self.sourceType == "metadataColumn": + if self.metadataColumn is None or self.artifact is not None: + raise ValueError( + "A metadata QC source requires only metadataColumn provenance" + ) + if self.origin != "ingestionMetadata": + raise ValueError( + "Metadata QC sources must use ingestionMetadata origin" + ) + elif self.artifact is None or self.metadataColumn is not None: + raise ValueError("An artifact QC source requires only artifact provenance") + if self.activeCells < 0 or not 0 <= self.missingCells <= self.activeCells: + raise ValueError("QC source active and missing counts are inconsistent") + if self.usableForFiltering and self.missingCells: + raise ValueError("A QC source with missing values cannot drive filtering") + return self + + +class QcSourceConcordance(AgentDataModel): + """Observed agreement between imported and derived forms of one metric.""" + + metricRole: QcMetricRole = "diagnostic" + leftSourceId: str = "" + rightSourceId: str = "" + comparedCells: int = 0 + missingCells: int = 0 + meanAbsoluteDifference: float | None = None + maximumAbsoluteDifference: float | None = None + pearsonCorrelation: float | None = None + exactlyEqual: bool = False + numericallyClose: bool = False + evidenceId: str = "" + + +class CaptureFailureEvidence(AgentDataModel): + """Multi-axis capture anomaly plus exclusion-safety inputs.""" + + capture: str = "" + activeCells: int = 0 + retainedCells: int = 0 + retainedFraction: float = 0.0 + adverseAxes: list[QcMetricRole] = Field(default_factory=list) + independentAdverseAxes: int = 0 + metricMissingFractions: dict[str, float] = Field(default_factory=dict) + reasons: list[str] = Field(default_factory=list) + wholeCaptureFailure: bool = False + conditionAndUnitSafety: list[dict[str, Any]] = Field(default_factory=list) + preservesConditionCoverage: bool = False + preservesIndependentUnitCoverage: bool = False + exclusionEligible: bool = False + doubletEvidenceIds: list[str] = Field(default_factory=list) + evidenceId: str = "" + + @model_validator(mode="after") + def validate_failure(self) -> "CaptureFailureEvidence": + if self.independentAdverseAxes != len(set(self.adverseAxes)): + raise ValueError("Capture failure axis count must match its unique axes") + if self.wholeCaptureFailure != (self.independentAdverseAxes >= 2): + raise ValueError( + "Whole-capture failure requires at least two independent QC axes" + ) + if self.exclusionEligible and ( + not self.wholeCaptureFailure + or not self.preservesConditionCoverage + or not self.preservesIndependentUnitCoverage + ): + raise ValueError( + "Capture exclusion requires failure and preserved design coverage" + ) + return self + + +class ContrastPlan(AgentDataModel): + """One deterministic sample-aware statistical-testing license.""" + + coefficient: str = "" + groupOrder: list[str | int | float | bool] = Field(default_factory=list) + sampleBy: str | None = None + pairBy: str | None = None + test: ContrastTest | None = None + sampleStatistic: ContrastSampleStatistic = "mean" + expressionCutoff: float = 0.0 + status: ContrastStatus = "blocked" + betweenUnitDesign: bool = False + replicationPassed: bool = False + estimabilityPassed: bool = False + pairedCoveragePassed: bool | None = None + replication: dict[str, Any] = Field(default_factory=dict) + estimability: dict[str, Any] = Field(default_factory=dict) + pairedCoverage: dict[str, Any] = Field(default_factory=dict) + blockedReasons: list[str] = Field(default_factory=list) + evidenceId: str = "" + evidenceIds: list[str] = Field(default_factory=list) + + @model_validator(mode="after") + def validate_contrast(self) -> "ContrastPlan": + if self.coefficient != self.coefficient.strip(): + raise ValueError( + "Contrast coefficient cannot contain surrounding whitespace" + ) + if self.sampleBy is not None and ( + not self.sampleBy.strip() or self.sampleBy != self.sampleBy.strip() + ): + raise ValueError("Contrast sampleBy must be a non-empty trimmed name") + if self.pairBy is not None and ( + not self.pairBy.strip() or self.pairBy != self.pairBy.strip() + ): + raise ValueError("Contrast pairBy must be a non-empty trimmed name") + group_keys = [(type(value).__name__, repr(value)) for value in self.groupOrder] + if len(group_keys) != len(set(group_keys)): + raise ValueError("Contrast groupOrder must contain unique values") + if any( + isinstance(value, float) and not math.isfinite(value) + for value in self.groupOrder + ): + raise ValueError("Contrast groupOrder cannot contain non-finite values") + if not math.isfinite(self.expressionCutoff): + raise ValueError("Contrast expressionCutoff must be finite") + if self.sampleStatistic != "fraction" and self.expressionCutoff != 0.0: + raise ValueError( + "Contrast expressionCutoff is only used with fraction summaries" + ) + if self.test == "mann_whitney" and len(self.groupOrder) != 2: + raise ValueError("mann_whitney requires exactly two ordered groups") + if self.test == "kruskal_wallis" and len(self.groupOrder) < 3: + raise ValueError("kruskal_wallis requires at least three ordered groups") + if self.test == "wilcoxon": + if len(self.groupOrder) != 2 or self.pairBy is None: + raise ValueError( + "wilcoxon requires exactly two groups and an explicit pairBy" + ) + elif self.pairBy is not None and self.test is not None: + raise ValueError("A paired contrast must use the wilcoxon test") + if self.status == "licensed": + if ( + not self.coefficient + or self.sampleBy is None + or self.test is None + or self.blockedReasons + or not self.betweenUnitDesign + or not self.replicationPassed + or not self.estimabilityPassed + or (self.pairBy is not None and self.pairedCoveragePassed is not True) + ): + raise ValueError( + "A licensed contrast requires resolved design, replication, " + "estimability, and paired coverage" + ) + elif not self.blockedReasons: + raise ValueError("A non-licensed contrast requires blockedReasons") + return self + + @classmethod + def get_blank(cls) -> "ContrastPlan": + return cls(blockedReasons=["unresolvedContrast"]) + + +def _validate_qc_sources( + *, + action: CellQcAction, + attributes: list[str], + artifact_metrics: list[NamedArtifactSource], + sample_column: str | None, + sample_artifact: NamedArtifactSource | None, + registered_profile: RegisteredCellQcProfile | None = None, + allow_metric_name_collisions: bool = False, +) -> None: + if len(attributes) != len(set(attributes)): + raise ValueError("Cell-QC metadata attributes must be unique") + if any( + not attribute.strip() or attribute != attribute.strip() + for attribute in attributes + ): + raise ValueError( + "Cell-QC metadata attributes cannot be blank or have surrounding whitespace" + ) + artifact_names = [source.name for source in artifact_metrics] + if len(artifact_names) != len(set(artifact_names)): + raise ValueError("Cell-QC artifact metric names must be unique") + if not allow_metric_name_collisions and set(attributes) & set(artifact_names): + raise ValueError( + "Cell-QC metadata and artifact metric names collide; explicitly " + "validated multi-source evidence is required" + ) + if any(source.artifact.kind != "quality_metric" for source in artifact_metrics): + raise ValueError( + "Cell-QC artifactMetrics must reference quality_metric artifacts" + ) + if sample_column is not None and sample_artifact is not None: + raise ValueError( + "Cell-QC sampleColumn and sampleArtifact are mutually exclusive" + ) + if sample_column is not None and ( + not sample_column.strip() or sample_column != sample_column.strip() + ): + raise ValueError( + "Cell-QC sampleColumn cannot be blank or have surrounding whitespace" + ) + if sample_artifact is not None and sample_artifact.artifact.kind != "hto_identity": + raise ValueError( + "Cell-QC sampleArtifact must reference an hto_identity artifact" + ) + if sample_artifact is not None and sample_artifact.name in artifact_names: + raise ValueError("Cell-QC sample and metric artifact names must be distinct") + if registered_profile is not None: + if registered_profile == "retainWithFlags": + if action != "skip": + raise ValueError( + "retainWithFlags must use the non-filtering skip action" + ) + if sample_column is not None or sample_artifact is not None: + raise ValueError("retainWithFlags cannot include a capture source") + return + if action != "registeredMad": + raise ValueError(f"{registered_profile} must use the registeredMad action") + capture_profile = registered_profile in { + "captureMad5", + "captureMad3Sensitivity", + "pooledReferenceMad5", + } + has_one_capture_source = (sample_column is None) != (sample_artifact is None) + if capture_profile and not has_one_capture_source: + raise ValueError( + f"{registered_profile} requires exactly one proven capture source" + ) + if not capture_profile and ( + sample_column is not None or sample_artifact is not None + ): + raise ValueError(f"{registered_profile} cannot include a capture source") + if not attributes and not artifact_metrics: + raise ValueError("Registered MAD filtering requires at least one metric") + return + if action == "registeredMad": + raise ValueError("registeredMad requires a registeredProfile") + if action == "skip" and (attributes or artifact_metrics): + raise ValueError("skip cannot include Cell-QC metrics") + if action != "skip" and not attributes and not artifact_metrics: + raise ValueError("Cell-QC filtering requires at least one metric") + if action == "sampleMad" and (sample_column is None) == (sample_artifact is None): + raise ValueError( + "sampleMad requires exactly one sampleColumn or sampleArtifact" + ) + if action != "sampleMad" and ( + sample_column is not None or sample_artifact is not None + ): + raise ValueError("Only sampleMad can include a sample source") + + +class CellQcProfileEvidence(AgentDataModel): + """Projected retention for one registered or legacy cell-QC profile.""" + + profileId: str = "" + action: CellQcAction = "skip" + registeredProfile: RegisteredCellQcProfile | None = None + driverAssay: str | None = None + driverAssayType: CellQcDriverType | None = None + sampleColumn: str | None = None + sampleArtifact: NamedArtifactSource | None = None + captureColumn: str | None = None + captureArtifact: NamedArtifactSource | None = None + attributes: list[str] = Field(default_factory=list) + artifactMetrics: list[NamedArtifactSource] = Field(default_factory=list) + metricSources: list[QcMetricSourceEvidence] = Field(default_factory=list) + sourceConcordance: list[QcSourceConcordance] = Field(default_factory=list) + parameters: dict[str, Any] = Field(default_factory=dict) + resolvedBounds: dict[str, Any] | list[dict[str, Any]] = Field(default_factory=dict) + activeCells: int = 0 + retainedCells: int = 0 + retainedFraction: float = 0.0 + activeCellsByCapture: dict[str, int] = Field(default_factory=dict) + sampleRetainedCells: dict[str, int] = Field(default_factory=dict) + retainedCellsByColumn: dict[str, dict[str, int]] = Field(default_factory=dict) + retainedCellsByCombination: dict[str, dict[str, int]] = Field(default_factory=dict) + unsafeRetentionGroups: list[str] = Field(default_factory=list) + flaggedCells: dict[str, int] = Field(default_factory=dict) + metricFlaggedCells: dict[str, dict[str, int]] = Field(default_factory=dict) + failedCaptureCandidates: list[str] = Field(default_factory=list) + captureFailureEvidence: list[CaptureFailureEvidence] = Field(default_factory=list) + excludableCaptureCandidates: list[str] = Field(default_factory=list) + notes: list[str] = Field(default_factory=list) + evidenceId: str = "" + + @model_validator(mode="after") + def validate_sources(self) -> "CellQcProfileEvidence": + _validate_qc_sources( + action=self.action, + attributes=self.attributes, + artifact_metrics=self.artifactMetrics, + sample_column=self.sampleColumn, + sample_artifact=self.sampleArtifact, + registered_profile=self.registeredProfile, + allow_metric_name_collisions=True, + ) + if self.captureColumn is not None and self.captureArtifact is not None: + raise ValueError( + "Cell-QC captureColumn and captureArtifact are mutually exclusive" + ) + if ( + self.captureArtifact is not None + and self.captureArtifact.artifact.kind != "hto_identity" + ): + raise ValueError( + "Cell-QC captureArtifact must reference an hto_identity artifact" + ) + failures = {item.capture: item for item in self.captureFailureEvidence} + if len(failures) != len(self.captureFailureEvidence): + raise ValueError("Cell-QC capture failure evidence must be unique") + expected_failed = sorted( + capture for capture, item in failures.items() if item.wholeCaptureFailure + ) + if failures and sorted(self.failedCaptureCandidates) != expected_failed: + raise ValueError( + "Cell-QC failed captures must match their multi-axis evidence" + ) + expected_excludable = sorted( + capture for capture, item in failures.items() if item.exclusionEligible + ) + if failures and sorted(self.excludableCaptureCandidates) != expected_excludable: + raise ValueError( + "Cell-QC excludable captures must match design-safety evidence" + ) + return self + + @classmethod + def get_blank(cls) -> "CellQcProfileEvidence": + return cls() + + +class CellQcPlan(AgentDataModel): + """A validated selection from the bounded cell-QC profiles.""" + + action: CellQcAction = "skip" + registeredProfile: RegisteredCellQcProfile | None = None + profileId: str = "" + driverAssay: str | None = None + driverAssayType: CellQcDriverType | None = None + sampleColumn: str | None = None + sampleArtifact: NamedArtifactSource | None = None + attributes: list[str] = Field(default_factory=list) + artifactMetrics: list[NamedArtifactSource] = Field(default_factory=list) + rationale: str = "" + evidenceIds: list[str] = Field(default_factory=list) + + @model_validator(mode="after") + def validate_sources(self) -> "CellQcPlan": + _validate_qc_sources( + action=self.action, + attributes=self.attributes, + artifact_metrics=self.artifactMetrics, + sample_column=self.sampleColumn, + sample_artifact=self.sampleArtifact, + registered_profile=self.registeredProfile, + allow_metric_name_collisions=True, + ) + return self + + @classmethod + def get_blank(cls) -> "CellQcPlan": + return cls() + + +class ExperimentalContextDecision(AgentDataModel): + """Model-authored choices that are revalidated against the datastore.""" + + columnDomains: dict[str, ColumnDomain] = Field(default_factory=dict) + coefficientsOfInterest: list[str] = Field(default_factory=list) + unitsOfInference: dict[str, InferenceUnit] = Field(default_factory=dict) + protectedCombinations: SkipJsonSchema[list[list[str]]] = Field(default_factory=list) + physicalCaptureColumn: SkipJsonSchema[str | None] = None + pooledReferenceCaptures: SkipJsonSchema[list[str]] = Field(default_factory=list) + unsupportedProtection: SkipJsonSchema[list[str]] = Field(default_factory=list) + batchCorrection: BatchCorrectionPlan = Field( + default_factory=BatchCorrectionPlan.get_blank + ) + cellQc: SkipJsonSchema[CellQcPlan] = Field(default_factory=CellQcPlan.get_blank) + rationale: str = "" + evidenceIds: list[str] = Field(default_factory=list) + needsInput: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "ExperimentalContextDecision": + return cls() + + +class RepresentationEvaluation(AgentDataModel): + """Bounded integration metrics for one exact graph representation.""" + + available: bool = False + assay: str | None = None + cellSelection: ArtifactReferenceModel | None = None + neighbors: ArtifactReferenceModel | None = None + connectivityMap: ArtifactReferenceModel | None = None + metrics: dict[str, float] = Field(default_factory=dict) + notes: list[str] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "RepresentationEvaluation": + return cls() + + +class CovariateEvidence(AgentDataModel): + """One deterministic covariate characterization returned by a tool.""" + + characterization: CovariateCharacterization = Field( + default_factory=lambda: CovariateCharacterization(status="needsInput") + ) + batchSafety: list[BatchSafetyEvidence] = Field(default_factory=list) + qcProfiles: list[CellQcProfileEvidence] = Field(default_factory=list) + qcMetricSources: list[QcMetricSourceEvidence] = Field(default_factory=list) + qcSourceConcordance: list[QcSourceConcordance] = Field(default_factory=list) + contrastPlans: list[ContrastPlan] = Field(default_factory=list) + htoIdentityColumns: list[str] = Field(default_factory=list) + htoIdentityArtifacts: list[NamedArtifactSource] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + evidenceRequirements: list[DesignEvidenceRequirement] = Field(default_factory=list) + evidenceCoverage: list[DesignEvidenceCoverage] = Field(default_factory=list) + + +class ExperimentalContextResult(AgentDataModel): + """Canonical experimental-context report returned to the caller.""" + + status: StageStatus + decision: ExperimentalContextDecision + characterization: CovariateCharacterization + cellSelection: ArtifactReferenceModel | None = None + cellQc: CellQcPlan = Field(default_factory=CellQcPlan.get_blank) + qcProfiles: list[CellQcProfileEvidence] = Field(default_factory=list) + qcMetricSources: list[QcMetricSourceEvidence] = Field(default_factory=list) + qcSourceConcordance: list[QcSourceConcordance] = Field(default_factory=list) + contrastPlans: list[ContrastPlan] = Field(default_factory=list) + qualityMetricArtifacts: list[NamedArtifactSource] = Field(default_factory=list) + htoIdentityColumns: list[str] = Field(default_factory=list) + htoIdentityArtifacts: list[NamedArtifactSource] = Field(default_factory=list) + batchSafety: list[BatchSafetyEvidence] = Field(default_factory=list) + currentRepresentation: RepresentationEvaluation = Field( + default_factory=RepresentationEvaluation.get_blank + ) + notes: list[str] = Field(default_factory=list) + runInfo: AgentRunInfo = Field(default_factory=AgentRunInfo) + + @classmethod + def get_blank(cls) -> "ExperimentalContextResult": + return cls( + status="needsInput", + decision=ExperimentalContextDecision.get_blank(), + characterization=CovariateCharacterization(status="needsInput"), + ) + + def to_parameter_tuning_handoff(self) -> ExperimentalTuningHandoff: + """Return validated integration inputs for Parameter Tuning.""" + if self.status != "done": + raise ValueError( + "Experimental Context must be done before creating a tuning handoff" + ) + if self.cellSelection is None: + raise ValueError("Experimental Context result lacks a cell selection") + plan = self.decision.batchCorrection + batch_columns = sorted(plan.batchColumns) + safety = sorted( + ( + item + for item in self.batchSafety + if item.batchColumns == batch_columns + and item.coefficient in self.decision.coefficientsOfInterest + ), + key=lambda item: item.coefficient, + ) + if plan.action in {"evaluateHarmony", "unsafe"}: + expected = set(self.decision.coefficientsOfInterest) + if {item.coefficient for item in safety} != expected: + raise ValueError( + "Experimental Context result lacks exact batch safety evidence" + ) + if any(item.evidenceId not in plan.evidenceIds for item in safety): + raise ValueError( + "Batch-correction plan does not cite its exact safety evidence" + ) + if plan.action == "evaluateHarmony" and any( + item.status != "safe" for item in safety + ): + raise ValueError("Harmony plan contains non-safe batch evidence") + if plan.action == "unsafe" and ( + any(item.status == "notComputed" for item in safety) + or not any(item.status == "unsafe" for item in safety) + ): + raise ValueError("Unsafe plan lacks exact unsafe batch evidence") + return ExperimentalTuningHandoff( + cellSelection=self.cellSelection, + batchAction=plan.action, + batchColumns=batch_columns, + preservationColumns=list(plan.preserveColumns), + coefficientsOfInterest=list(self.decision.coefficientsOfInterest), + batchSafety=safety, + evidenceIds=sorted({*self.decision.evidenceIds, *plan.evidenceIds}), + ) + + def to_biological_handoff( + self, + coefficient: str | None = None, + ) -> ExperimentalBiologyHandoff: + """Return one explicitly resolved biological coefficient.""" + if self.status != "done": + raise ValueError( + "Experimental Context must be done before creating a biology handoff" + ) + if self.cellSelection is None: + raise ValueError("Experimental Context result lacks a cell selection") + coefficients = list(self.decision.coefficientsOfInterest) + if coefficient is None: + if len(coefficients) != 1: + raise ValueError( + "Select one coefficient explicitly for biological interpretation" + ) + coefficient = coefficients[0] + if coefficient not in coefficients: + raise ValueError(f"Unknown coefficient of interest {coefficient!r}") + records = { + record.get("name"): record + for record in self.characterization.coefficients + if isinstance(record.get("name"), str) + } + record = records.get(coefficient) + if record is None: + raise ValueError(f"Missing characterization for {coefficient!r}") + reports = { + report.get("coefficient"): report + for report in self.characterization.confounding + if isinstance(report.get("coefficient"), str) + } + report = reports.get(coefficient) + known_evidence = characterization_evidence(self.characterization) + relevant_evidence = { + f"column:{coefficient}", + f"coefficient:{coefficient}", + f"estimability:{coefficient}", + *( + evidence_id + for evidence_id in known_evidence + if evidence_id.startswith(f"confounding:{coefficient}:") + ), + } + for unit_name in ( + record.get("observationUnit"), + record.get("independentUnit"), + ): + if isinstance(unit_name, str): + relevant_evidence.add(f"column:{unit_name}") + return ExperimentalBiologyHandoff( + cellSelection=self.cellSelection, + conditionColumn=coefficient, + observationUnit=record.get("observationUnit"), + independentUnit=record.get("independentUnit"), + coefficientScope=str(record.get("scope", "")), + estimability=dict(report.get("estimability") or {}) if report else {}, + evidenceIds=sorted(relevant_evidence.intersection(known_evidence)), + ) + + +class ExperimentalContextDependencies(AgentDataModel): + """Runtime-only state shared by the agent's read-only tools.""" + + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + + store: Any = Field(default=None, exclude=True) + qcAssay: str | None = Field(default=None, exclude=True) + cells: Any = Field(default=None, exclude=True) + neighbors: Any = Field(default=None, exclude=True) + connectivityMap: Any = Field(default=None, exclude=True) + cellSelection: Any = Field(default=None, exclude=True) + studyContext: str = "" + studyObjective: str = "" + directions: dict[str, Any] = Field(default_factory=dict) + characterizationInputs: dict[str, Any] = Field(default_factory=dict, exclude=True) + inventoryData: dict[str, Any] = Field(default_factory=dict, exclude=True) + qcDesignData: Any = Field(default=None, exclude=True) + checkpointRead: Any = Field(default=None, exclude=True) + checkpointWrite: Any = Field(default=None, exclude=True) + evidenceIds: set[str] = Field(default_factory=set) + characterization: CovariateCharacterization | None = None + designRounds: int = 0 + comparisons: list[CovariateComparison] = Field(default_factory=list) + protectedCombinations: list[list[str]] = Field(default_factory=list) + captureProposal: CaptureProposal | None = None + batchSafety: dict[str, BatchSafetyEvidence] = Field(default_factory=dict) + qcProfiles: dict[str, CellQcProfileEvidence] = Field(default_factory=dict) + qcMetricSources: list[QcMetricSourceEvidence] = Field(default_factory=list) + qcSourceConcordance: list[QcSourceConcordance] = Field(default_factory=list) + contrastPlans: dict[str, ContrastPlan] = Field(default_factory=dict) + htoIdentityColumns: list[str] = Field(default_factory=list) + qualityMetricArtifacts: list[NamedArtifactSource] = Field(default_factory=list) + htoIdentityArtifacts: list[NamedArtifactSource] = Field(default_factory=list) + currentRepresentation: RepresentationEvaluation = Field( + default_factory=RepresentationEvaluation.get_blank + ) + toolCalls: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "ExperimentalContextDependencies": + return cls() + + +def characterization_evidence( + characterization: CovariateCharacterization, +) -> set[str]: + """Build stable evidence IDs from one deterministic characterization.""" + evidence_ids = { + f"column:{record['name']}" + for record in characterization.columns + if isinstance(record.get("name"), str) + } + for record in characterization.coefficients: + coefficient = record.get("name") + if isinstance(coefficient, str): + evidence_ids.add(f"coefficient:{coefficient}") + for report in characterization.confounding: + coefficient = report.get("coefficient") + if not isinstance(coefficient, str): + continue + evidence_ids.add(f"estimability:{coefficient}") + for pair in report.get("pairs", []): + technical = pair.get("technical") + if isinstance(technical, str): + evidence_ids.add(f"confounding:{coefficient}:{technical}") + evidence_ids.update(item.evidenceId for item in characterization.comparisons) + if characterization.captureProvenance is not None: + evidence_ids.add( + f"captureProvenance:{characterization.captureProvenance.column}" + ) + return evidence_ids diff --git a/scarf/agent/experimental_context/qc_evidence.py b/scarf/agent/experimental_context/qc_evidence.py new file mode 100644 index 00000000..cc321c6d --- /dev/null +++ b/scarf/agent/experimental_context/qc_evidence.py @@ -0,0 +1,1867 @@ +"""Experimental-context quality-control evidence assembly.""" + +import json +import math +import re +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Literal, cast + +import numpy as np +import pandas as pd + +from ...metadata.selection import resolve_cell_aligned_artifact +from ...quality_control.filtering import ( + _validated_sample_labels, + gaussian_quantile_bounds, +) +from ...storage.artifacts import ( + fingerprint_array, + fingerprint_strings, + inspect_artifact, +) +from ...storage.refs import ArtifactRef +from ...storage.selections import read_stored_selection_mask +from ..cell_quality.profiles import ( + AutoFilterProjection, + QcMetricRole, + RegisteredCellQcProfile, + RegisteredQcProjection, + offered_registered_qc_profiles, + project_auto_filter_profile, + qc_metric_execution_name, + registered_qc_metric_role, +) +from ..tools import artifact_reference, core_artifact_reference +from ..types import ArtifactReferenceModel +from .contracts import ( + CaptureFailureEvidence, + CellQcAction, + CellQcDriverType, + CellQcProfileEvidence, + CovariateCharacterization, + ExperimentalContextDependencies, + LegacyCellQcAction, + NamedArtifactSource, + QcMetricSourceEvidence, + QcSourceConcordance, +) + +_MAX_QC_SAMPLE_PROFILES = 4 + + +def _persisted_assay_type(store: Any, assay_name: str) -> str: + """Read one persisted assay type without inferring modality from features.""" + root = getattr(store, "zw", None) + attrs = getattr(root, "attrs", {}) + raw_types = attrs.get("assayTypes", {}) if isinstance(attrs, Mapping) else {} + if isinstance(raw_types, Mapping): + assay_type = raw_types.get(assay_name) + if isinstance(assay_type, str): + return assay_type + return assay_name if assay_name in {"RNA", "ATAC", "ADT", "HTO"} else "Assay" + + +def _qc_driver( + store: Any, selected_assay: str | None = None +) -> tuple[str, CellQcDriverType] | None: + """Use an explicit QC assay, otherwise the first RNA or ATAC assay.""" + assay_names = [str(name) for name in getattr(store, "assay_names", [])] + if selected_assay is not None: + if selected_assay not in assay_names: + raise ValueError(f"Unknown QC assay {selected_assay!r}") + selected_type = _persisted_assay_type(store, selected_assay) + if selected_type == "RNA": + return selected_assay, "RNA" + if selected_type == "ATAC": + return selected_assay, "ATAC" + raise ValueError("The selected QC assay must have persisted RNA or ATAC type") + for assay_type in ("RNA", "ATAC"): + for assay_name in assay_names: + if _persisted_assay_type(store, assay_name) == assay_type: + return assay_name, assay_type + return None + + +def _hto_identity_columns(deps: ExperimentalContextDependencies) -> list[str]: + """Return explicitly supplied imported HTO identity metadata columns.""" + requested: list[str] = [] + directed_many = deps.directions.get("htoIdentityColumns") + if isinstance(directed_many, list | tuple): + requested.extend(str(value) for value in directed_many) + directed_one = deps.directions.get("htoIdentityColumn") + if isinstance(directed_one, str): + requested.append(directed_one) + available = set(deps.store.cells.columns) + return list(dict.fromkeys(name for name in requested if name in available)) + + +def _cell_selection_ref(deps: ExperimentalContextDependencies) -> ArtifactRef: + selection = core_artifact_reference(deps.cellSelection) + if not isinstance(selection, ArtifactRef): + raise ValueError("cellSelection must identify an exact artifact") + if selection.kind != "cell_selection" or selection.scope != "datastore": + raise ValueError("cellSelection must identify a datastore cell selection") + return selection + + +def _active_cell_count(deps: ExperimentalContextDependencies) -> int: + selection = _cell_selection_ref(deps) + active = read_stored_selection_mask( + deps.store.zw, + selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ) + if active.ndim != 1 or active.shape[0] != deps.store.cells.N: + raise ValueError( + "cellSelection must contain an aligned boolean selection vector" + ) + return int(active.sum()) + + +def _source_ref( + source: NamedArtifactSource, + *, + expected_kind: str, +) -> ArtifactRef: + if not isinstance(source, NamedArtifactSource): + raise TypeError("Artifact sources must be NamedArtifactSource values") + if not source.name.strip(): + raise ValueError("Artifact sources require a non-empty semantic name") + artifact = core_artifact_reference(source.artifact) + if not isinstance(artifact, ArtifactRef) or artifact.kind != expected_kind: + raise ValueError( + f"Artifact source {source.name!r} must reference {expected_kind!r}" + ) + return artifact + + +def _artifact_evidence_id(source: NamedArtifactSource) -> str: + return f"htoIdentityArtifact:{source.name}:{source.artifact.artifactId}" + + +def _hto_artifact_map( + deps: ExperimentalContextDependencies, +) -> dict[str, ArtifactRef]: + artifacts: dict[str, ArtifactRef] = {} + for source in deps.htoIdentityArtifacts: + if source.name in artifacts: + raise ValueError("HTO identity artifact names must be unique") + artifacts[source.name] = _source_ref( + source, + expected_kind="hto_identity", + ) + return artifacts + + +def _resolved_artifact_values( + deps: ExperimentalContextDependencies, + source: NamedArtifactSource, + *, + expected_kind: str, +) -> np.ndarray: + resolved = resolve_cell_aligned_artifact( + deps.store.zw, + _source_ref(source, expected_kind=expected_kind), + cell_selection=_cell_selection_ref(deps), + expected_kind=expected_kind, + ) + return np.asarray(resolved.values) + + +def _artifact_input_references( + value: Any, + *, + limit: int = 16, +) -> list[ArtifactReferenceModel]: + refs: list[ArtifactReferenceModel] = [] + seen: set[tuple[str, str | None, str, str]] = set() + + def visit(item: Any) -> None: + if len(refs) >= limit: + return + if isinstance(item, ArtifactRef): + ref = item + elif isinstance(item, Mapping) and { + "scope", + "kind", + "artifact_id", + }.issubset(item): + try: + ref = ArtifactRef.from_dict(item) + except (KeyError, TypeError, ValueError): + ref = None + else: + ref = None + if ref is not None: + key = (ref.scope, ref.assay, ref.kind, ref.artifact_id) + if key not in seen: + seen.add(key) + refs.append(artifact_reference(ref)) + return + if isinstance(item, Mapping): + for nested in item.values(): + visit(nested) + elif isinstance(item, list | tuple): + for nested in item: + visit(nested) + + visit(value) + return refs + + +def _qc_metric_sources( + deps: ExperimentalContextDependencies, + driver: tuple[str, CellQcDriverType], +) -> tuple[ + dict[str, np.ndarray], + list[str], + list[NamedArtifactSource], + list[QcMetricSourceEvidence], + list[QcSourceConcordance], + list[str], + dict[str, np.ndarray], +]: + assay_name, assay_type = driver + del assay_type + selection = _cell_selection_ref(deps) + selection_model = artifact_reference(selection) + active_cells = _active_cell_count(deps) + metadata_names = _qc_attributes(deps.store, assay_name, driver[1]) + artifact_candidates: list[NamedArtifactSource] = [] + for source in deps.qualityMetricArtifacts: + artifact = _source_ref(source, expected_kind="quality_metric") + if artifact.assay == assay_name: + artifact_candidates.append(source) + metadata_collisions = set(metadata_names).intersection( + source.name for source in artifact_candidates + ) + + values_by_execution_name: dict[str, np.ndarray] = {} + values_by_source: dict[str, np.ndarray] = {} + sources: list[QcMetricSourceEvidence] = [] + valid_metadata: list[str] = [] + valid_artifacts: list[NamedArtifactSource] = [] + notes: list[str] = [] + + for name in metadata_names: + raw = np.asarray(deps.cells.fetch(name)) + try: + values = np.asarray(raw, dtype=float) + except (TypeError, ValueError): + fingerprint = fingerprint_strings(raw) + source_id = f"qcMetric:metadata:{assay_name}:{name}:{fingerprint}" + sources.append( + QcMetricSourceEvidence( + sourceId=source_id, + metricName=name, + metricRole=registered_qc_metric_role(name), + assay=assay_name, + sourceType="metadataColumn", + origin="ingestionMetadata", + executionName=name, + metadataColumn=name, + cellSelection=selection_model, + valuesFingerprint=fingerprint, + activeCells=active_cells, + missingCells=active_cells, + notes=["Metric is not numeric and cannot drive filtering"], + ) + ) + notes.append(f"QC metadata source {name!r} is not numeric") + continue + if values.ndim != 1 or values.shape != (active_cells,): + raise ValueError( + f"QC metadata source {name!r} does not align with cellSelection" + ) + fingerprint = fingerprint_array(values) + missing = int((~np.isfinite(values)).sum()) + source_id = f"qcMetric:metadata:{assay_name}:{name}:{fingerprint}" + usable = missing == 0 + source_notes = ( + [] if usable else [f"{missing} active cells have non-finite metric values"] + ) + sources.append( + QcMetricSourceEvidence( + sourceId=source_id, + metricName=name, + metricRole=registered_qc_metric_role(name), + assay=assay_name, + sourceType="metadataColumn", + origin="ingestionMetadata", + executionName=name, + metadataColumn=name, + cellSelection=selection_model, + valuesFingerprint=fingerprint, + activeCells=active_cells, + missingCells=missing, + usableForFiltering=usable, + notes=source_notes, + ) + ) + values_by_source[source_id] = values + if usable: + values_by_execution_name[name] = values + valid_metadata.append(name) + else: + notes.extend(source_notes) + + for source in artifact_candidates: + artifact = _source_ref(source, expected_kind="quality_metric") + values = np.asarray( + _resolved_artifact_values( + deps, + source, + expected_kind="quality_metric", + ), + dtype=float, + ) + if values.ndim != 1 or values.shape != (active_cells,): + raise ValueError( + f"QC artifact {source.name!r} does not align with cellSelection" + ) + execution_name = qc_metric_execution_name( + source.name, + artifact_id=artifact.artifact_id, + collides_with_metadata=source.name in metadata_collisions, + ) + if execution_name in values_by_execution_name: + raise ValueError( + f"QC execution metric name {execution_name!r} is not unique" + ) + fingerprint = fingerprint_array(values) + missing = int((~np.isfinite(values)).sum()) + status = inspect_artifact(deps.store.zw, artifact) + operation = status.operation + origin: Literal[ + "ingestionMetadata", + "derivedArtifact", + "externalArtifact", + ] = ( + "derivedArtifact" + if operation == "run_feature_percentage" + else "externalArtifact" + ) + source_id = ( + f"qcMetric:artifact:{artifact.assay}:{source.name}:{artifact.artifact_id}" + ) + usable = missing == 0 + source_notes = ( + [] if usable else [f"{missing} active cells have non-finite metric values"] + ) + sources.append( + QcMetricSourceEvidence( + sourceId=source_id, + metricName=source.name, + metricRole=registered_qc_metric_role(source.name), + assay=assay_name, + sourceType="artifact", + origin=origin, + executionName=execution_name, + artifact=artifact_reference(artifact), + cellSelection=selection_model, + inputArtifacts=_artifact_input_references(status.inputs or {}), + provenanceOperation=operation, + valuesFingerprint=fingerprint, + activeCells=active_cells, + missingCells=missing, + usableForFiltering=usable, + notes=source_notes, + ) + ) + values_by_source[source_id] = values + if usable: + values_by_execution_name[execution_name] = values + valid_artifacts.append(source) + else: + notes.extend(source_notes) + + concordance: list[QcSourceConcordance] = [] + metadata_sources = [ + source for source in sources if source.sourceType == "metadataColumn" + ] + artifact_sources = [source for source in sources if source.sourceType == "artifact"] + # An imported percentage has no frozen gene definition. Keep it visible for + # comparison, but let an exactly defined derived metric own filtering. + canonical_roles: set[QcMetricRole] = set() + if ( + driver[1] == "RNA" + and callable(getattr(deps.store, "get_assay", None)) + and callable(getattr(deps.store, "load_artifact", None)) + ): + expected_masks = { + role: mask + for role, _, _, mask in _rna_percentage_feature_masks( + deps.store, assay_name + ) + } + for metric_source in artifact_sources: + if metric_source.provenanceOperation != "run_feature_percentage": + continue + for reference in metric_source.inputArtifacts: + if ( + reference.kind != "feature_selection" + or reference.assay != assay_name + ): + continue + expected = expected_masks.get(metric_source.metricRole) + actual = np.asarray( + deps.store.load_artifact(core_artifact_reference(reference))[ + "values" + ][:], + dtype=bool, + ) + if expected is not None and np.array_equal(expected, actual): + canonical_roles.add(metric_source.metricRole) + for metric_source in metadata_sources: + if metric_source.metricRole not in {"mitochondrial", "ribosomal"}: + continue + metric_source.usableForFiltering = False + note = ( + f"Imported {metric_source.metricName} is retained for comparison; filtering " + "uses the derived percentage with an exact gene selection." + if metric_source.metricRole in canonical_roles + else f"Imported {metric_source.metricName} has no validated gene definition " + "and cannot drive filtering. This percentage QC axis remains unavailable." + ) + metric_source.notes.append(note) + notes.append(note) + values_by_execution_name.pop(metric_source.executionName, None) + if metric_source.metadataColumn in valid_metadata: + valid_metadata.remove(metric_source.metadataColumn) + for metric_source in artifact_sources: + if metric_source.artifact is None: + continue + execution_name = qc_metric_execution_name( + metric_source.metricName, + artifact_id=metric_source.artifact.artifactId, + collides_with_metadata=metric_source.metricName in valid_metadata, + ) + if execution_name != metric_source.executionName: + if metric_source.executionName in values_by_execution_name: + values_by_execution_name[execution_name] = ( + values_by_execution_name.pop(metric_source.executionName) + ) + metric_source.executionName = execution_name + for role in expected_masks: + if not any( + metric_source.metricRole == role and metric_source.usableForFiltering + for metric_source in artifact_sources + ): + notes.append( + f"The {role} percentage has no exact usable artifact; " + "QC conclusions cannot claim that this axis was evaluated." + ) + for left in metadata_sources: + for right in artifact_sources: + if left.metricRole != right.metricRole or left.metricRole == "diagnostic": + continue + if right.artifact is None: + raise ValueError("Artifact QC source lacks its exact reference") + left_values = values_by_source.get(left.sourceId) + right_values = values_by_source.get(right.sourceId) + if left_values is None or right_values is None: + continue + finite = np.isfinite(left_values) & np.isfinite(right_values) + compared = int(finite.sum()) + missing = int(len(finite) - compared) + mean_difference: float | None = None + maximum_difference: float | None = None + pearson: float | None = None + exactly_equal = False + numerically_close = False + if compared: + left_finite = left_values[finite] + right_finite = right_values[finite] + differences = np.abs(left_finite - right_finite) + mean_difference = float(differences.mean()) + maximum_difference = float(differences.max()) + exactly_equal = missing == 0 and bool( + np.array_equal(left_finite, right_finite) + ) + numerically_close = missing == 0 and bool( + np.allclose( + left_finite, + right_finite, + rtol=1e-6, + atol=1e-8, + ) + ) + if ( + compared >= 2 + and float(np.std(left_finite)) > 0.0 + and float(np.std(right_finite)) > 0.0 + ): + correlation = float(np.corrcoef(left_finite, right_finite)[0, 1]) + if math.isfinite(correlation): + pearson = correlation + evidence_id = ( + f"qcConcordance:{left.metricRole}:" + f"{left.valuesFingerprint}:{right.artifact.artifactId}" + ) + concordance.append( + QcSourceConcordance( + metricRole=left.metricRole, + leftSourceId=left.sourceId, + rightSourceId=right.sourceId, + comparedCells=compared, + missingCells=missing, + meanAbsoluteDifference=mean_difference, + maximumAbsoluteDifference=maximum_difference, + pearsonCorrelation=pearson, + exactlyEqual=exactly_equal, + numericallyClose=numerically_close, + evidenceId=evidence_id, + ) + ) + return ( + values_by_execution_name, + valid_metadata, + valid_artifacts, + sources, + concordance, + notes, + values_by_source, + ) + + +def _qc_attributes(store: Any, assay_name: str, assay_type: str) -> list[str]: + del assay_type + suffixes = ["nCounts", "nFeatures", "percentMito", "percentRibo"] + available = set(store.cells.columns) + return [ + f"{assay_name}_{suffix}" + for suffix in suffixes + if f"{assay_name}_{suffix}" in available + ] + + +def _rna_percentage_feature_masks( + store: Any, assay_name: str +) -> list[tuple[QcMetricRole, str, str, np.ndarray]]: + """Resolve symbol-defined RNA percentages without the ambiguous MT prefix.""" + assay = store.get_assay(assay_name) + feature_ids = np.asarray(assay.feats.fetch_all("ids")).astype(str) + feature_names = np.asarray(assay.feats.fetch_all("names")).astype(str) + specifications: tuple[tuple[QcMetricRole, str, str], ...] = ( + ("mitochondrial", "percentMito", r"(?i)^MT-"), + ("ribosomal", "percentRibo", r"(?i)^(RPS|RPL|MRPS|MRPL)"), + ) + resolved = [] + for role, suffix, pattern in specifications: + compiled = re.compile(pattern) + mask = np.fromiter( + ( + compiled.search(feature_id) is not None + or compiled.search(feature_name) is not None + for feature_id, feature_name in zip( + feature_ids, feature_names, strict=True + ) + ), + dtype=bool, + count=assay.feats.N, + ) + resolved.append((role, suffix, pattern, mask)) + return resolved + + +def _derive_missing_percentage_artifacts( + store: Any, + *, + cell_selection: ArtifactRef, + driver: tuple[str, CellQcDriverType] | None, + quality_sources: Sequence[NamedArtifactSource], +) -> list[NamedArtifactSource]: + """Derive RNA percentage artifacts even when unbound metadata is present.""" + sources = list(quality_sources) + if driver is None or driver[1] != "RNA": + return sources + if not callable(getattr(store, "set_feature_selection", None)) or not callable( + getattr(store, "run_feature_percentage", None) + ): + return sources + assay_name = driver[0] + supplied_roles = { + registered_qc_metric_role(source.name) + for source in sources + if source.artifact.assay == assay_name + } + existing_names = {source.name for source in sources} + for role, suffix, _, mask in _rna_percentage_feature_masks(store, assay_name): + metric_name = f"{assay_name}_{suffix}" + if role in supplied_roles: + continue + if not mask.any(): + continue + if metric_name in existing_names: + raise ValueError(f"Derived QC metric name {metric_name!r} is not unique") + feature_selection = store.set_feature_selection( + from_assay=assay_name, + mask=mask, + invalidate_cache=False, + ) + metric = store.run_feature_percentage( + cell_selection, + feature_selection, + invalidate_cache=False, + ) + sources.append( + NamedArtifactSource( + name=metric_name, + artifact=artifact_reference(metric), + ) + ) + existing_names.add(metric_name) + supplied_roles.add(role) + return sources + + +def _qc_sample_columns( + deps: ExperimentalContextDependencies, + characterization: CovariateCharacterization | None, +) -> list[str]: + requested: list[str] = [] + directed = deps.directions.get("cellQc") + if isinstance(directed, Mapping): + sample_column = directed.get("sampleColumn") + if isinstance(sample_column, str): + requested.append(sample_column) + if characterization is not None: + for record in characterization.coefficients: + observation_unit = record.get("observationUnit") + if isinstance(observation_unit, str): + requested.append(observation_unit) + requested.extend(deps.htoIdentityColumns) + available = set(deps.store.cells.columns) + return list( + dict.fromkeys(name for name in requested if name in available and name != "I") + )[:_MAX_QC_SAMPLE_PROFILES] + + +def _qc_profile_id( + action: LegacyCellQcAction, + *, + driver: tuple[str, CellQcDriverType] | None, + sample_column: str | None = None, + sample_artifact: NamedArtifactSource | None = None, +) -> str: + assay_name, assay_type = driver or ("none", "none") + suffix = { + "skip": "skip", + "globalGaussian": "globalGaussian:0.01:0.99", + "sampleMad": ( + f"sampleMad:metadata:{sample_column}:3:20" + if sample_artifact is None + else ( + f"sampleMad:artifact:{sample_artifact.name}:" + f"{sample_artifact.artifact.artifactId}:3:20" + ) + ), + }[action] + return f"cellQc:{assay_type}:{assay_name}:{suffix}" + + +def _registered_qc_profile_id( + profile: RegisteredCellQcProfile, + *, + driver: tuple[str, CellQcDriverType], + sample_column: str | None, + sample_artifact: NamedArtifactSource | None, +) -> str: + if sample_column is not None: + source = f"metadata:{sample_column}" + elif sample_artifact is not None: + source = ( + f"artifact:{sample_artifact.name}:{sample_artifact.artifact.artifactId}" + ) + else: + source = "global" + return f"cellQc:{driver[1]}:{driver[0]}:registered:{profile}:{source}" + + +def _directed_capture_source( + deps: ExperimentalContextDependencies, +) -> tuple[str | None, NamedArtifactSource | None, np.ndarray] | None: + directed_qc = deps.directions.get("cellQc") + qc_directions = dict(directed_qc) if isinstance(directed_qc, Mapping) else {} + candidates = [ + deps.directions.get("physicalCaptureColumn"), + qc_directions.get("physicalCaptureColumn"), + qc_directions.get("captureColumn"), + deps.captureProposal.column if deps.captureProposal is not None else None, + ] + specified = [value for value in candidates if value is not None] + if not specified: + return None + if any(not isinstance(value, str) or not value.strip() for value in specified): + raise ValueError("physicalCaptureColumn must be a non-empty string") + names = list(dict.fromkeys(str(value) for value in specified)) + if len(names) != 1: + raise ValueError("Conflicting physical capture columns were supplied") + name = names[0] + matching_artifacts = [ + source for source in deps.htoIdentityArtifacts if source.name == name + ] + if len(matching_artifacts) > 1: + raise ValueError(f"Physical capture artifact {name!r} is not unique") + if matching_artifacts: + source = matching_artifacts[0] + labels = _resolved_artifact_values( + deps, + source, + expected_kind="hto_identity", + ) + return None, source, np.asarray(labels) + if name not in deps.cells.columns: + raise ValueError( + f"physicalCaptureColumn {name!r} is not observed metadata or an " + "exact HTO identity artifact" + ) + return name, None, np.asarray(deps.cells.fetch(name)) + + +def _directed_pooled_reference_captures( + deps: ExperimentalContextDependencies, +) -> tuple[str, ...] | None: + directed_qc = deps.directions.get("cellQc") + qc_directions = dict(directed_qc) if isinstance(directed_qc, Mapping) else {} + raw = qc_directions.get( + "pooledReferenceCaptures", + deps.directions.get("pooledReferenceCaptures"), + ) + proposed = ( + deps.captureProposal.referenceCaptures + if deps.captureProposal is not None + else [] + ) + if proposed: + if raw is not None and ( + not isinstance(raw, list | tuple) or list(raw) != proposed + ): + raise ValueError( + "Reference proposal conflicts with caller reference captures" + ) + raw = proposed + if raw is None: + return None + if not isinstance(raw, list | tuple) or any( + not isinstance(value, str) or not value.strip() for value in raw + ): + raise ValueError("pooledReferenceCaptures must contain non-empty strings") + references = tuple(str(value) for value in raw) + if len(references) < 2 or len(references) != len(set(references)): + raise ValueError( + "pooledReferenceCaptures must contain at least two unique captures" + ) + return references + + +def _provenance_label(value: Any) -> str | None: + if isinstance(value, np.generic): + value = value.item() + if value is None or value is pd.NA or value is pd.NaT: + return None + if isinstance(value, float) and not math.isfinite(value): + return None + if isinstance(value, bytes): + try: + value = value.decode("utf-8") + except UnicodeDecodeError: + return None + if isinstance(value, str) and not value.strip(): + return None + return str(value) + + +@dataclass +class _QcDesignData: + """Frozen metadata shared only while projecting one set of QC policies.""" + + cells: Any + values: dict[str, np.ndarray] = field(default_factory=dict) + labels: dict[str, np.ndarray] = field(default_factory=dict) + captureSource: np.ndarray | None = None + captures: np.ndarray | None = None + units: dict[tuple[str, str, str | None], list[tuple[Any, ...]]] = field( + default_factory=dict + ) + exclusions: dict[str, tuple[list[dict[str, Any]], bool, bool]] = field( + default_factory=dict + ) + combinations: dict[tuple[str, ...], np.ndarray] = field(default_factory=dict) + + @property + def columns(self) -> list[str]: + return list(self.cells.columns) + + def fetch(self, column: str) -> np.ndarray: + if column not in self.values: + self.values[column] = np.asarray(self.cells.fetch(column)).copy() + return self.values[column] + + def encoded(self, column: str) -> np.ndarray: + if column not in self.labels: + self.labels[column] = np.asarray( + [_provenance_label(value) for value in self.fetch(column)], dtype=object + ) + return self.labels[column] + + def combined(self, columns: Sequence[str]) -> np.ndarray: + from .comparisons import combination_labels + + key = tuple(columns) + if key not in self.combinations: + self.combinations[key] = combination_labels(self, list(columns)) + return self.combinations[key] + + +def _capture_design_safety( + deps: ExperimentalContextDependencies, + characterization: CovariateCharacterization | None, + capture_labels: np.ndarray, + capture: str, +) -> tuple[list[dict[str, Any]], bool, bool]: + if characterization is None: + return [], False, False + data = deps.qcDesignData + if not isinstance(data, _QcDesignData): + data = _QcDesignData(deps.cells) + if data.captureSource is not capture_labels: + data = _QcDesignData(deps.cells) if data.captureSource is not None else data + active = np.ones(len(capture_labels), dtype=bool) + normalized = _validated_sample_labels( + capture_labels, active, label_name="physical capture labels" + ) + data.captures = np.asarray( + [ + value.decode("utf-8") if isinstance(value, bytes) else str(value) + for value in normalized + ], + dtype=object, + ) + data.captureSource = capture_labels + if capture not in data.exclusions: + data.exclusions[capture] = _compute_capture_design_safety( + data, characterization, deps.protectedCombinations, capture + ) + return data.exclusions[capture] + + +def _compute_capture_design_safety( + data: _QcDesignData, + characterization: CovariateCharacterization, + protected_combinations: Sequence[Sequence[str]], + capture: str, +) -> tuple[list[dict[str, Any]], bool, bool]: + assert data.captures is not None + encoded = data.captures + kinds = {record["name"]: record.get("kind") for record in characterization.columns} + after = encoded != capture + safety: list[dict[str, Any]] = [] + for record in characterization.coefficients: + coefficient = record.get("name") + observation = record.get("observationUnit") + independent = record.get("independentUnit") + if ( + not isinstance(coefficient, str) + or not isinstance(observation, str) + or record.get("scope") != "betweenUnit" + ): + continue + if any( + name not in data.columns + for name in (coefficient, observation, independent) + if isinstance(name, str) + ): + safety.append( + { + "coefficient": coefficient, + "reason": "missingDesignColumn", + "preservesConditionCoverage": False, + "preservesIndependentUnitCoverage": False, + } + ) + continue + condition_values = data.encoded(coefficient) + observation_values = data.encoded(observation) + if ( + condition_values.shape != after.shape + or observation_values.shape != after.shape + ): + raise ValueError("Capture safety columns do not align with cellSelection") + if kinds.get(coefficient) not in {"categorical", "continuous"}: + safety.append( + { + "coefficient": coefficient, + "reason": "unknownCovariateKind", + "preservesConditionCoverage": False, + "preservesIndependentUnitCoverage": False, + } + ) + continue + if kinds.get(coefficient) == "continuous": + values = np.asarray( + pd.to_numeric(data.fetch(coefficient), errors="coerce"), dtype=float + ) + observed = np.isfinite(values) + matched_after = observed & after + safety.append( + { + "coefficient": coefficient, + "conditionColumn": coefficient, + "kind": "continuous", + "observationUnit": observation, + "independentUnit": independent, + "matchedRowsBeforeExclusion": int(observed.sum()), + "matchedRowsAfterExclusion": int(matched_after.sum()), + "missingRowsAfterExclusion": int((after & ~observed).sum()), + "quantilesBeforeExclusion": np.quantile( + values[observed], [0, 0.25, 0.5, 0.75, 1] + ).tolist() + if observed.any() + else [], + "quantilesAfterExclusion": np.quantile( + values[matched_after], [0, 0.25, 0.5, 0.75, 1] + ).tolist() + if matched_after.any() + else [], + "independentUnitsAfterExclusion": len( + { + value + for value in data.encoded(independent or observation)[ + matched_after + ] + if value is not None + } + ), + "reason": "Continuous distributions and unit support are descriptive; preservation under capture exclusion has not been established.", + "preservesConditionCoverage": False, + "preservesIndependentUnitCoverage": False, + } + ) + continue + required_groups = [ + value for value in dict.fromkeys(condition_values) if value is not None + ] + remaining_groups = [ + value + for value in dict.fromkeys(condition_values[after]) + if value is not None + ] + preserves_conditions = set(remaining_groups) == set(required_groups) + + observation_counts: list[dict[str, Any]] = [] + independent_counts: list[dict[str, Any]] = [] + independent_values: np.ndarray | None = None + if isinstance(independent, str): + independent_values = data.encoded(independent) + if independent_values.shape != after.shape: + raise ValueError( + "Capture independent-unit column does not align with cellSelection" + ) + unit_key = (coefficient, observation, independent) + if unit_key not in data.units: + data.units[unit_key] = list( + dict.fromkeys( + zip( + encoded, + condition_values, + observation_values, + independent_values + if independent_values is not None + else np.full(len(after), None), + strict=True, + ) + ) + ) + remaining_units = [row for row in data.units[unit_key] if row[0] != capture] + for group in required_groups: + observation_levels = { + row[2] + for row in remaining_units + if row[1] == group and row[2] is not None + } + observation_counts.append( + {"group": group, "count": len(observation_levels)} + ) + if independent_values is not None: + independent_levels = { + row[3] + for row in remaining_units + if row[1] == group and row[3] is not None + } + independent_counts.append( + {"group": group, "count": len(independent_levels)} + ) + + replication_counts = ( + independent_counts if independent_values is not None else observation_counts + ) + minimum_units = min( + (int(item["count"]) for item in replication_counts), + default=0, + ) + complete_pairs = 0 + incomplete_pairs = 0 + duplicate_pair_groups = 0 + single_group_pairs = 0 + if independent_values is not None: + pair_groups: dict[str, dict[str, set[str]]] = {} + for _, pair_group, observation_value, pair in remaining_units: + if pair is None or pair_group is None or observation_value is None: + continue + pair_groups.setdefault(pair, {}).setdefault(pair_group, set()).add( + observation_value + ) + required_set = set(required_groups) + for groups in pair_groups.values(): + if len(groups) == 1: + single_group_pairs += 1 + duplicate_pair_groups += sum( + len(observations) > 1 for observations in groups.values() + ) + if set(groups) == required_set and all( + len(observations) == 1 for observations in groups.values() + ): + complete_pairs += 1 + else: + incomplete_pairs += 1 + original_pair_design = dict(record.get("pairedCoverage") or {}).get("design") + pair_structure_safe = ( + True + if independent_values is None + else ( + complete_pairs >= 2 + and incomplete_pairs == 0 + and duplicate_pair_groups == 0 + ) + if original_pair_design == "paired" + else (len(pair_groups) >= 2 and single_group_pairs == len(pair_groups)) + if original_pair_design == "betweenIndependentUnits" + else False + ) + preserves_units = ( + preserves_conditions and minimum_units >= 2 and pair_structure_safe + ) + safety.append( + { + "coefficient": coefficient, + "conditionColumn": coefficient, + "observationUnit": observation, + "independentUnit": independent, + "requiredGroups": required_groups, + "remainingGroups": remaining_groups, + "observationUnitsByGroup": observation_counts, + "independentUnitsByGroup": independent_counts, + "minimumIndependentUnitsAfterExclusion": minimum_units, + "completePairsAfterExclusion": complete_pairs, + "incompletePairsAfterExclusion": incomplete_pairs, + "duplicatePairGroupsAfterExclusion": duplicate_pair_groups, + "independentUnitDesign": original_pair_design, + "preservesConditionCoverage": preserves_conditions, + "preservesIndependentUnitCoverage": preserves_units, + } + ) + for columns in protected_combinations: + label = json.dumps(columns, separators=(",", ":")) + try: + combined = data.combined(columns) + except (KeyError, ValueError): + safety.append( + { + "conditionColumns": columns, + "preservesConditionCoverage": False, + "preservesIndependentUnitCoverage": False, + "reason": "missingProtectedCombination", + } + ) + continue + joint_groups = np.unique(combined) + coverage = set(np.unique(combined[after])) == set(joint_groups) + units = { + record.get("independentUnit") or record.get("observationUnit") + for record in characterization.coefficients + if record.get("name") in columns + } + units.discard(None) + independent_safe = coverage and bool(units) + for unit in units: + if not isinstance(unit, str) or unit not in data.columns: + independent_safe = False + continue + values = data.encoded(unit) + independent_safe = independent_safe and all( + len( + { + value + for value in values[after & (combined == group)] + if value is not None + } + ) + >= 2 + for group in joint_groups + ) + safety.append( + { + "conditionColumns": columns, + "combination": label, + "preservesConditionCoverage": coverage, + "preservesIndependentUnitCoverage": independent_safe, + } + ) + return ( + safety, + bool(safety) and all(item["preservesConditionCoverage"] for item in safety), + bool(safety) + and all(item["preservesIndependentUnitCoverage"] for item in safety), + ) + + +def _capture_source_missingness( + sources: Sequence[QcMetricSourceEvidence], + values_by_source: Mapping[str, np.ndarray], + capture_labels: np.ndarray | None, +) -> list[QcMetricSourceEvidence]: + if capture_labels is None: + return list(sources) + active = np.ones(len(capture_labels), dtype=bool) + normalized = _validated_sample_labels( + capture_labels, + active, + label_name="physical capture labels", + ) + captures: list[tuple[str, np.ndarray]] = [] + seen: set[str] = set() + for raw in normalized: + value = raw.item() if isinstance(raw, np.generic) else raw + key = value.decode("utf-8") if isinstance(value, bytes) else str(value) + if key in seen: + continue + seen.add(key) + captures.append((key, normalized == value)) + output: list[QcMetricSourceEvidence] = [] + for source in sources: + values = values_by_source.get(source.sourceId) + missing_by_capture: dict[str, int] = {} + if values is not None: + for capture, mask in captures: + missing_by_capture[capture] = int((~np.isfinite(values[mask])).sum()) + elif source.missingCells == source.activeCells: + missing_by_capture = { + capture: int(mask.sum()) for capture, mask in captures + } + output.append( + source.model_copy(update={"missingCellsByCapture": missing_by_capture}) + ) + return output + + +def _capture_failure_models( + projection: RegisteredQcProjection | AutoFilterProjection, + *, + deps: ExperimentalContextDependencies, + characterization: CovariateCharacterization | None, + capture_labels: np.ndarray | None, + metric_sources: Sequence[QcMetricSourceEvidence], +) -> list[CaptureFailureEvidence]: + if capture_labels is None: + return [] + output: list[CaptureFailureEvidence] = [] + source_by_id = {source.sourceId: source for source in metric_sources} + for comparison in projection.captureComparisons: + missing_fractions = { + source_id: ( + source.missingCellsByCapture.get(comparison.capture, 0) + / comparison.cells + if comparison.cells + else 0.0 + ) + for source_id, source in source_by_id.items() + } + safety, condition_safe, unit_safe = _capture_design_safety( + deps, + characterization, + capture_labels, + comparison.capture, + ) + failure = CaptureFailureEvidence( + capture=comparison.capture, + activeCells=comparison.cells, + retainedCells=comparison.retainedCells or 0, + retainedFraction=comparison.retainedFraction or 0.0, + adverseAxes=list(comparison.adverseAxes), + independentAdverseAxes=comparison.independentAdverseAxes, + metricMissingFractions=missing_fractions, + reasons=list(comparison.reasons), + wholeCaptureFailure=comparison.wholeCaptureFailure, + conditionAndUnitSafety=safety, + preservesConditionCoverage=condition_safe, + preservesIndependentUnitCoverage=unit_safe, + exclusionEligible=( + comparison.wholeCaptureFailure and condition_safe and unit_safe + ), + evidenceId=( + f"qcCapture:{comparison.capture}:" + f"{comparison.independentAdverseAxes}axes" + ), + ) + output.append(failure) + return output + + +def _design_retention( + deps: ExperimentalContextDependencies, + characterization: CovariateCharacterization | None, + active: np.ndarray, + keep: np.ndarray, +) -> dict[str, Any]: + """Check exact categorical conditions, units, and protected joint groups.""" + cells = ( + deps.qcDesignData + if isinstance(deps.qcDesignData, _QcDesignData) + else deps.cells + if deps.cells is not None + else deps.store.cells + ) + retention_columns: list[str] = [] + if characterization is not None: + kinds = { + record["name"]: record.get("kind") for record in characterization.columns + } + for coefficient in characterization.coefficients: + name = coefficient.get("name") + for value in ( + name if kinds.get(name) == "categorical" else None, + coefficient.get("observationUnit"), + coefficient.get("independentUnit"), + ): + if isinstance(value, str) and value in cells.columns: + retention_columns.append(value) + retained_by_column: dict[str, dict[str, int]] = {} + unsafe_groups: list[str] = [] + retained = np.asarray(keep, dtype=bool) & np.asarray(active, dtype=bool) + for column in dict.fromkeys(retention_columns): + labels = ( + cells.encoded(column) + if isinstance(cells, _QcDesignData) + else np.asarray( + [_provenance_label(value) for value in cells.fetch(column)], + dtype=object, + ) + ) + if labels.shape != retained.shape: + raise ValueError( + f"QC retention column {column!r} does not align with cellSelection" + ) + counts: dict[str, int] = {} + present = labels != None # noqa: E711 + if (np.asarray(active, dtype=bool) & ~present).any(): + unsafe_groups.append(f"{column}:missingValues") + for raw_label in np.unique(labels[np.asarray(active, dtype=bool) & present]): + label = raw_label.item() if isinstance(raw_label, np.generic) else raw_label + key = label.decode("utf-8") if isinstance(label, bytes) else str(label) + count = int((retained & (labels == raw_label)).sum()) + counts[key] = count + if count == 0: + unsafe_groups.append(f"{column}={key}") + retained_by_column[column] = counts + from .comparisons import combination_labels + + retained_by_combination: dict[str, dict[str, int]] = {} + for columns in deps.protectedCombinations: + key = json.dumps(columns, separators=(",", ":")) + try: + labels = ( + cells.combined(columns) + if isinstance(cells, _QcDesignData) + else combination_labels(cells, columns) + ) + except (KeyError, ValueError): + unsafe_groups.append(f"combination:{key}:missingValues") + continue + counts = { + str(label): int((retained & (labels == label)).sum()) + for label in np.unique(labels[np.asarray(active, dtype=bool)]) + } + retained_by_combination[key] = counts + unsafe_groups.extend( + f"combination:{key}={label}" + for label, count in counts.items() + if count == 0 + ) + return { + "retainedCellsByColumn": retained_by_column, + "retainedCellsByCombination": retained_by_combination, + "unsafeRetentionGroups": sorted(unsafe_groups), + } + + +def _registered_profile_evidence( + projection: RegisteredQcProjection, + *, + deps: ExperimentalContextDependencies, + characterization: CovariateCharacterization | None, + driver: tuple[str, CellQcDriverType], + active: np.ndarray, + values_by_attr: dict[str, np.ndarray], + metadata_attributes: list[str], + artifact_metrics: list[NamedArtifactSource], + metric_sources: list[QcMetricSourceEvidence], + source_concordance: list[QcSourceConcordance], + sample_column: str | None, + sample_artifact: NamedArtifactSource | None, + capture_column: str | None, + capture_artifact: NamedArtifactSource | None, + capture_labels: np.ndarray | None, + pooled_reference_captures: tuple[str, ...] | None, + active_cells: int, + comparison_source: str | None, +) -> CellQcProfileEvidence: + attributes = list(metadata_attributes) + metric_artifacts = list(artifact_metrics) + profile_id = _registered_qc_profile_id( + projection.profile, + driver=driver, + sample_column=sample_column, + sample_artifact=sample_artifact, + ) + n_mads = 3.0 if projection.profile == "captureMad3Sensitivity" else 5.0 + action: CellQcAction = ( + "skip" if projection.profile == "retainWithFlags" else "registeredMad" + ) + parameters: dict[str, Any] = { + "policyVersion": 1, + "profile": projection.profile, + "nMads": n_mads, + "boundPolicy": { + "count": {"remove": "lower", "flag": "upper"}, + "feature": {"remove": "lower", "flag": "upper"}, + "mitochondrial": {"remove": "upper", "fixedCutoff": None}, + "diagnostic": {"remove": "none"}, + }, + "resolvedBounds": [threshold.to_dict() for threshold in projection.thresholds], + "captureSizes": projection.captureSizes, + "captureComparisons": [ + comparison.to_dict() for comparison in projection.captureComparisons + ], + "captureComparisonSource": comparison_source, + "pooledReferenceCaptures": list(pooled_reference_captures or ()), + } + failure_evidence = _capture_failure_models( + projection, + deps=deps, + characterization=characterization, + capture_labels=capture_labels, + metric_sources=metric_sources, + ) + return CellQcProfileEvidence( + profileId=profile_id, + action=action, + registeredProfile=projection.profile, + driverAssay=driver[0], + driverAssayType=driver[1], + sampleColumn=sample_column, + sampleArtifact=sample_artifact, + captureColumn=capture_column, + captureArtifact=capture_artifact, + attributes=attributes, + artifactMetrics=metric_artifacts, + metricSources=metric_sources, + sourceConcordance=source_concordance, + parameters=parameters, + resolvedBounds=parameters["resolvedBounds"], + activeCells=active_cells, + retainedCells=projection.retainedCells, + retainedFraction=( + projection.retainedCells / active_cells if active_cells else 0.0 + ), + activeCellsByCapture=projection.captureSizes, + sampleRetainedCells=projection.retainedByCapture, + **_design_retention(deps, characterization, active, projection.keep), + flaggedCells=projection.flagCounts, + metricFlaggedCells=projection.metricFlagCounts, + failedCaptureCandidates=list(projection.failedCaptureCandidates), + captureFailureEvidence=failure_evidence, + excludableCaptureCandidates=[ + item.capture for item in failure_evidence if item.exclusionEligible + ], + notes=list(projection.warnings), + evidenceId=f"qcProfile:{profile_id}", + ) + + +def _registered_qc_profiles( + deps: ExperimentalContextDependencies, + *, + characterization: CovariateCharacterization | None, + driver: tuple[str, CellQcDriverType], + active: np.ndarray, + values_by_attr: dict[str, np.ndarray], + metadata_attributes: list[str], + artifact_metrics: list[NamedArtifactSource], + metric_sources: list[QcMetricSourceEvidence], + source_concordance: list[QcSourceConcordance], + capture: tuple[str | None, NamedArtifactSource | None, np.ndarray] | None = None, +) -> list[CellQcProfileEvidence]: + if capture is None: + capture = _directed_capture_source(deps) + sample_column: str | None = None + sample_artifact: NamedArtifactSource | None = None + capture_labels: np.ndarray | None = None + if capture is not None: + sample_column, sample_artifact, capture_labels = capture + if sample_column is not None: + comparison_source = f"metadata:{sample_column}" + elif sample_artifact is not None: + comparison_source = ( + f"artifact:{sample_artifact.name}:{sample_artifact.artifact.artifactId}" + ) + else: + comparison_source = None + pooled_references = _directed_pooled_reference_captures(deps) + if pooled_references is not None and capture is None: + raise ValueError( + "pooledReferenceCaptures requires an explicit physicalCaptureColumn" + ) + projections = offered_registered_qc_profiles( + values_by_metric=values_by_attr, + active=active, + capture_labels=capture_labels, + grouping_proven=capture is not None, + min_cells_per_capture=20, + pooled_reference_captures=pooled_references, + ) + profiles: list[CellQcProfileEvidence] = [] + for projection in projections: + uses_capture = projection.profile in { + "captureMad5", + "captureMad3Sensitivity", + "pooledReferenceMad5", + } + profiles.append( + _registered_profile_evidence( + projection, + deps=deps, + characterization=characterization, + driver=driver, + active=active, + values_by_attr=values_by_attr, + metadata_attributes=metadata_attributes, + artifact_metrics=artifact_metrics, + metric_sources=metric_sources, + source_concordance=source_concordance, + sample_column=sample_column if uses_capture else None, + sample_artifact=sample_artifact if uses_capture else None, + capture_column=sample_column, + capture_artifact=sample_artifact, + capture_labels=capture_labels, + pooled_reference_captures=( + pooled_references + if projection.profile == "pooledReferenceMad5" + else None + ), + active_cells=int(active.sum()), + comparison_source=comparison_source, + ) + ) + return profiles + + +def _global_qc_profile( + deps: ExperimentalContextDependencies, + driver: tuple[str, CellQcDriverType], + active: np.ndarray, + active_cells: int, + values_by_attr: dict[str, np.ndarray], + metadata_attributes: list[str], + artifact_metrics: list[NamedArtifactSource], + attribute_notes: list[str], + *, + characterization: CovariateCharacterization | None = None, + metric_sources: list[QcMetricSourceEvidence] | None = None, + source_concordance: list[QcSourceConcordance] | None = None, + capture: tuple[str | None, NamedArtifactSource | None, np.ndarray] | None = None, +) -> CellQcProfileEvidence | None: + """Build an execution-exact projection of core global auto-filtering.""" + if not values_by_attr: + return None + metric_sources = list(metric_sources or []) + source_concordance = list(source_concordance or []) + for name, values in values_by_attr.items(): + selected = np.asarray(values)[active] + if selected.size and np.all(selected == selected[0]): + attribute_notes.append( + f"Scarf default global QC is unavailable: constant metric {name!r} produces non-finite Gaussian bounds" + ) + return None + low, high = gaussian_quantile_bounds(selected, 0.01, 0.99) + if not np.isfinite([low, high]).all(): + attribute_notes.append( + f"Scarf default global QC is unavailable: metric {name!r} produces non-finite Gaussian bounds" + ) + return None + capture_column: str | None = None + capture_artifact: NamedArtifactSource | None = None + capture_labels: np.ndarray | None = None + if capture is not None: + capture_column, capture_artifact, capture_labels = capture + try: + projection = project_auto_filter_profile( + "globalGaussian", + values_by_metric=values_by_attr, + active=active, + sample_labels=capture_labels, + grouping_proven=capture is not None, + ) + except ValueError as exc: + attribute_notes.append(f"Global Gaussian QC is not executable: {exc}") + return None + profile_id = _qc_profile_id( + "globalGaussian", + driver=driver, + ) + failures = _capture_failure_models( + projection, + deps=deps, + characterization=characterization, + capture_labels=capture_labels, + metric_sources=metric_sources, + ) + return CellQcProfileEvidence( + profileId=profile_id, + action="globalGaussian", + driverAssay=driver[0], + driverAssayType=driver[1], + captureColumn=capture_column, + captureArtifact=capture_artifact, + attributes=list(metadata_attributes), + artifactMetrics=list(artifact_metrics), + metricSources=metric_sources, + sourceConcordance=source_concordance, + parameters=projection.parameters, + resolvedBounds=cast(dict[str, Any], projection.parameters["resolvedBounds"]), + activeCells=active_cells, + retainedCells=projection.retainedCells, + retainedFraction=projection.retainedCells / active_cells, + activeCellsByCapture=projection.captureSizes, + sampleRetainedCells=projection.retainedByCapture, + **_design_retention(deps, characterization, active, projection.keep), + flaggedCells=projection.flagCounts, + metricFlaggedCells=projection.metricFlagCounts, + failedCaptureCandidates=list(projection.failedCaptureCandidates), + captureFailureEvidence=failures, + excludableCaptureCandidates=[ + item.capture for item in failures if item.exclusionEligible + ], + notes=[*attribute_notes, *projection.warnings], + evidenceId=f"qcProfile:{profile_id}", + ) + + +def _sample_qc_profiles( + deps: ExperimentalContextDependencies, + characterization: CovariateCharacterization | None, + driver: tuple[str, CellQcDriverType], + active: np.ndarray, + active_cells: int, + values_by_attr: dict[str, np.ndarray], + metadata_attributes: list[str], + artifact_metrics: list[NamedArtifactSource], + metric_sources: list[QcMetricSourceEvidence], + source_concordance: list[QcSourceConcordance], + capture: tuple[str | None, NamedArtifactSource | None, np.ndarray] | None, +) -> list[CellQcProfileEvidence]: + """Build core-parity sample MAD profiles from exact grouping sources.""" + attributes = list(values_by_attr) + profiles: list[CellQcProfileEvidence] = [] + sample_sources: list[ + tuple[str | None, NamedArtifactSource | None, np.ndarray | None, bool] + ] = [] + if capture is not None: + sample_sources.append((*capture[:2], capture[2], True)) + sample_sources.extend( + (None, source, None, False) for source in deps.htoIdentityArtifacts + ) + sample_sources.extend( + (column, None, None, False) + for column in _qc_sample_columns(deps, characterization) + ) + seen_sources: set[str] = set() + for ( + sample_column, + sample_artifact, + supplied_labels, + is_physical_capture, + ) in sample_sources: + source_key = ( + f"metadata:{sample_column}" + if sample_column is not None + else ( + f"artifact:{sample_artifact.artifact.artifactId}" + if sample_artifact is not None + else "" + ) + ) + if not source_key or source_key in seen_sources: + continue + seen_sources.add(source_key) + if len(seen_sources) > _MAX_QC_SAMPLE_PROFILES: + break + if not attributes: + break + artifact_labels = ( + supplied_labels + if supplied_labels is not None + else None + if sample_artifact is None + else _resolved_artifact_values( + deps, + sample_artifact, + expected_kind="hto_identity", + ) + ) + try: + sample_labels = ( + np.asarray(supplied_labels) + if supplied_labels is not None + else np.asarray(deps.cells.fetch(sample_column)) + if sample_column is not None + else np.asarray(artifact_labels) + ) + projection = project_auto_filter_profile( + "sampleMad", + values_by_metric=values_by_attr, + sample_labels=sample_labels, + active=active, + grouping_proven=True, + n_mads=3.0, + min_cells_per_sample=20, + ) + except (TypeError, ValueError): + continue + profile_id = _qc_profile_id( + "sampleMad", + driver=driver, + sample_column=sample_column, + sample_artifact=sample_artifact, + ) + failures = ( + _capture_failure_models( + projection, + deps=deps, + characterization=characterization, + capture_labels=sample_labels, + metric_sources=metric_sources, + ) + if is_physical_capture + else [] + ) + skip_reasons = cast( + dict[str, object], + projection.parameters["skipReasons"], + ) + profiles.append( + CellQcProfileEvidence( + profileId=profile_id, + action="sampleMad", + driverAssay=driver[0], + driverAssayType=driver[1], + sampleColumn=sample_column, + sampleArtifact=sample_artifact, + captureColumn=sample_column if is_physical_capture else None, + captureArtifact=sample_artifact if is_physical_capture else None, + attributes=list(metadata_attributes), + artifactMetrics=list(artifact_metrics), + metricSources=metric_sources, + sourceConcordance=source_concordance, + parameters={ + "nMads": 3.0, + "minCellsPerSample": 20, + "nSamples": len(projection.captureSizes), + "nSkippedSamples": len(skip_reasons), + }, + resolvedBounds=cast( + dict[str, Any], + projection.parameters["resolvedBounds"], + ), + activeCells=active_cells, + retainedCells=projection.retainedCells, + retainedFraction=projection.retainedCells / active_cells, + activeCellsByCapture=projection.captureSizes, + sampleRetainedCells=projection.retainedByCapture, + **_design_retention(deps, characterization, active, projection.keep), + flaggedCells=projection.flagCounts, + metricFlaggedCells=projection.metricFlagCounts, + failedCaptureCandidates=( + list(projection.failedCaptureCandidates) + if is_physical_capture + else [] + ), + captureFailureEvidence=failures, + excludableCaptureCandidates=[ + item.capture for item in failures if item.exclusionEligible + ], + notes=list(projection.warnings), + evidenceId=f"qcProfile:{profile_id}", + ) + ) + return profiles + + +def _offered_qc_profiles( + deps: ExperimentalContextDependencies, + characterization: CovariateCharacterization | None = None, +) -> list[CellQcProfileEvidence]: + """Share frozen design summaries across policies, never across changed inputs.""" + previous = deps.qcDesignData + deps.qcDesignData = _QcDesignData( + deps.cells if deps.cells is not None else deps.store.cells + ) + try: + return _project_qc_profiles(deps, characterization) + finally: + deps.qcDesignData = previous + + +def _project_qc_profiles( + deps: ExperimentalContextDependencies, + characterization: CovariateCharacterization | None, +) -> list[CellQcProfileEvidence]: + """Project bounded QC profiles against the exact shared cell selection.""" + active_cells = _active_cell_count(deps) + active = np.ones(active_cells, dtype=bool) + driver = _qc_driver(deps.store, deps.qcAssay) + driver_assay = driver[0] if driver is not None else None + driver_type = driver[1] if driver is not None else None + skip_id = _qc_profile_id( + "skip", + driver=driver, + ) + skip_notes = ( + [] + if driver is not None + else ["No RNA or ATAC assay is eligible to drive automatic cell QC"] + ) + profiles: list[CellQcProfileEvidence] = [] + if driver is None or active_cells == 0: + profiles.append( + CellQcProfileEvidence( + profileId=skip_id, + action="skip", + driverAssay=driver_assay, + driverAssayType=driver_type, + activeCells=active_cells, + retainedCells=active_cells, + retainedFraction=1.0 if active_cells else 0.0, + notes=skip_notes, + evidenceId=f"qcProfile:{skip_id}", + ) + ) + deps.qcProfiles = {profile.profileId: profile for profile in profiles} + return profiles + + ( + values_by_attr, + valid_metadata_attributes, + artifact_metrics, + metric_sources, + source_concordance, + attribute_notes, + values_by_source, + ) = _qc_metric_sources(deps, driver) + capture = _directed_capture_source(deps) + capture_column: str | None = None + capture_artifact: NamedArtifactSource | None = None + capture_labels: np.ndarray | None = None + capture_sizes: dict[str, int] = {} + if capture is not None: + capture_column, capture_artifact, capture_labels = capture + normalized = _validated_sample_labels( + capture_labels, + active, + label_name="physical capture labels", + ) + for raw in normalized: + value = raw.item() if isinstance(raw, np.generic) else raw + key = value.decode("utf-8") if isinstance(value, bytes) else str(value) + capture_sizes[key] = capture_sizes.get(key, 0) + 1 + metric_sources = _capture_source_missingness( + metric_sources, + values_by_source, + capture_labels, + ) + deps.qcMetricSources = metric_sources + deps.qcSourceConcordance = source_concordance + profiles = [ + CellQcProfileEvidence( + profileId=skip_id, + action="skip", + driverAssay=driver_assay, + driverAssayType=driver_type, + captureColumn=capture_column, + captureArtifact=capture_artifact, + metricSources=metric_sources, + sourceConcordance=source_concordance, + activeCells=active_cells, + retainedCells=active_cells, + retainedFraction=1.0, + activeCellsByCapture=capture_sizes, + sampleRetainedCells=capture_sizes, + notes=[*skip_notes, *attribute_notes], + evidenceId=f"qcProfile:{skip_id}", + ) + ] + + global_profile = _global_qc_profile( + deps, + driver, + active, + active_cells, + values_by_attr, + valid_metadata_attributes, + artifact_metrics, + attribute_notes, + characterization=characterization, + metric_sources=metric_sources, + source_concordance=source_concordance, + capture=capture, + ) + if global_profile is not None: + profiles.append(global_profile) + profiles.extend( + _sample_qc_profiles( + deps, + characterization, + driver, + active, + active_cells, + values_by_attr, + valid_metadata_attributes, + artifact_metrics, + metric_sources, + source_concordance, + capture, + ) + ) + profiles.extend( + _registered_qc_profiles( + deps, + characterization=characterization, + driver=driver, + active=active, + values_by_attr=values_by_attr, + metadata_attributes=valid_metadata_attributes, + artifact_metrics=artifact_metrics, + metric_sources=metric_sources, + source_concordance=source_concordance, + capture=capture, + ) + ) + + for profile in profiles: + profile.notes = list(dict.fromkeys([*attribute_notes, *profile.notes])) + deps.qcProfiles = {profile.profileId: profile for profile in profiles} + return profiles diff --git a/scarf/agent/experimental_context/requirements.py b/scarf/agent/experimental_context/requirements.py new file mode 100644 index 00000000..7173b8fb --- /dev/null +++ b/scarf/agent/experimental_context/requirements.py @@ -0,0 +1,377 @@ +"""Objective evidence requirements derived from bounded, measured study designs.""" + +import hashlib +import re +from typing import Any, Literal + +from ..record_io import canonical_json_bytes +from .contracts import ( + DesignEvidenceCoverage, + DesignEvidenceRequirement, + characterization_evidence, +) + + +def active_batch_safety(result: Any) -> list[Any]: + """Select the exact final assessed design without unioning alternatives.""" + all_assessments = list(result.batchSafety) + columns = sorted(result.decision.batchCorrection.batchColumns) + if not columns: + tested = {tuple(sorted(item.batchColumns)) for item in all_assessments} + if len(tested) > 1: + raise ValueError( + "The final correction plan must identify one exact assessed batch set; " + "separate alternatives cannot be combined into an untested design" + ) + columns = list(next(iter(tested), ())) + return [item for item in all_assessments if sorted(item.batchColumns) == columns] + + +def requested_design_questions( + study_context: str, study_objective: str, columns: list[str] +) -> list[tuple[str, list[str], bool]]: + """Keep explicit named joint/conditional requests visible before proposals.""" + output: list[tuple[str, list[str], bool]] = [] + for text in (study_context, study_objective): + for raw in re.split(r"[\n.!?]+", text): + quote = raw.strip() + joint = bool( + re.search( + r"\b(joint|jointly|combined|combination|interaction)\b", quote, re.I + ) + ) + conditional = bool( + re.search(r"\b(within|conditioned|stratif\w*)\b", quote, re.I) + ) + if not (joint or conditional) or re.search( + r"\b(?:do not|not requested|outside scope)\b", quote, re.I + ): + continue + named = sorted( + name + for name in columns + if re.search(r"(?= 2 or not named and generic) and ( + quote, + named, + conditional, + ) not in output: + output.append((quote, named, conditional)) + return output + + +def _measured_joint_design(comparison: Any) -> bool: + """Recognize joint descriptive counts without claiming an association.""" + proposal = comparison.proposal + columns = { + proposal.response, + *proposal.explanatoryColumns, + *([proposal.conditionedOn] if proposal.conditionedOn else []), + } + descriptive = comparison.evidence.get("descriptiveDesign", {}) + rows = descriptive.get("jointGroupSupport", []) + if ( + len(columns) != 3 + or descriptive.get("status") != "computed" + or not isinstance(rows, list) + or not rows + or any( + comparison.evidence.get("columnKinds", {}).get(name) != "categorical" + for name in columns + ) + ): + return False + for row in rows: + if ( + not isinstance(row, dict) + or set(row.get("groups", {})) != columns + or type(row.get("observationUnits")) is not int + or type(row.get("independentUnits")) is not int + or not 0 < row["independentUnits"] <= row["observationUnits"] + ): + return False + return bool( + sum(row["observationUnits"] for row in rows) + == descriptive.get("observationUnits") + ) + + +def requested_design_purpose( + quote: str, +) -> Literal["effectEstimation", "association", "designCoverage"]: + """Keep model guidance and validation on the same requested evidence kind.""" + if re.search(r"\bestimat\w*.*\beffect", quote, re.I): + return "effectEstimation" + if re.search(r"\bassociat\w*", quote, re.I): + return "association" + return "designCoverage" + + +def objective_evidence( + *, study_context: str, study_objective: str, experimental_result: Any +) -> tuple[list[DesignEvidenceRequirement], list[DesignEvidenceCoverage]]: + """Require design coverage and retain the purpose of every proposed question.""" + result = experimental_result + characterization = result.characterization + decision = result.decision + records = {item["name"]: item for item in characterization.columns} + coefficients = {item["name"]: item for item in characterization.coefficients} + mentioned = { + name + for name, item in records.items() + if item.get("domain") == "biological" + and re.search(r"(?= 2: + paired = record.get("pairedCoverage", {}) + if not paired: + unavailable.append(f"{name}: repeated-unit coverage is unavailable") + elif paired.get("design") == "mixedOrIncomplete": + reasons.append( + f"{name}: measured pairing is mixed or incomplete; no paired effect is estimated" + ) + if batch_columns: + for name in conditions: + matched = [ + item + for item in batch_safety + if item.coefficient == name + and sorted(item.batchColumns) == batch_columns + ] + if not matched or any(item.status == "notComputed" for item in matched): + unavailable.append( + f"{name}: estimability for the complete exact batch set is unavailable" + ) + elif any(item.status == "unsafe" for item in matched): + non_identifiable = True + reasons.append( + f"{name}: not identifiable under the exact assessed batch design" + ) + if not conditions: + reasons.append("No coefficient-level effect inference is authorized") + if not batch_columns: + reasons.append("No technical batch correction is proposed") + if characterization.captureProvenance is None: + reasons.append( + "Physical capture identity is unresolved; capture-dependent decisions are not authorized" + ) + coverage = [ + DesignEvidenceCoverage( + requirementId="studyDesign", + status=( + "failed" + if characterization.status == "failed" + else "unsupported" + if unavailable + else "nonIdentifiable" + if non_identifiable + else "computed" + ), + evidenceIds=sorted(evidence_ids), + reasons=[*unavailable, *reasons], + ) + ] + seen: set[str] = set() + for comparison in characterization.comparisons: + proposal = comparison.proposal + identity = hashlib.sha256( + canonical_json_bytes(proposal.model_dump(exclude={"rationale"})) + ).hexdigest() + requirement_id = f"designQuestion:{identity}" + if requirement_id in seen: + raise ValueError("Objective comparisons must have unique current proposals") + seen.add(requirement_id) + quote = proposal.objectiveQuote or study_objective + if quote not in f"{study_context}\n{study_objective}": + raise ValueError( + "Objective evidence requirements must quote exact study text" + ) + if ( + proposal.objectiveQuote + and quote in study_objective + and not proposal.essential + ): + raise ValueError( + "A question quoting the explicit study objective must remain essential" + ) + requirements.append( + DesignEvidenceRequirement( + requirementId=requirement_id, + question=proposal.rationale, + objectiveQuote=quote, + kind=proposal.purpose, + columns=[ + proposal.response, + *proposal.explanatoryColumns, + *([proposal.conditionedOn] if proposal.conditionedOn else []), + ], + observationUnit=proposal.observationUnit, + independentUnit=proposal.independentUnit or proposal.observationUnit, + essential=proposal.essential, + ) + ) + descriptive = comparison.evidence.get("descriptiveDesign", {}) + computed = comparison.status == "computed" + answer_reasons = list(comparison.reasons) + if proposal.purpose == "designCoverage": + computed = descriptive.get("status") == "computed" + if computed and comparison.status == "unsupported": + answer_reasons.append( + "Descriptive support answers the design question; the association method remains unsupported" + ) + elif proposal.purpose == "effectEstimation": + computed = False + answer_reasons.append( + "This workflow does not estimate biological effects or test expression hypotheses" + ) + if any( + records.get(column, {}).get(field) != value + for field, recorded in ( + ("kind", "columnKinds"), + ("domain", "columnDomains"), + ) + for column, value in comparison.evidence.get(recorded, {}).items() + ): + computed = False + answer_reasons.append( + "Comparison evidence has different column roles or kinds from the final design" + ) + coverage.append( + DesignEvidenceCoverage( + requirementId=requirement_id, + status="computed" if computed else "unsupported", + evidenceIds=[comparison.evidenceId], + reasons=answer_reasons, + ) + ) + for quote, names, conditional in requested_design_questions( + study_context, study_objective, list(records) + ): + purpose = requested_design_purpose(quote) + matches = [ + index + for index, item in enumerate(characterization.comparisons, start=1) + if set(names).issubset( + { + item.proposal.response, + *item.proposal.explanatoryColumns, + item.proposal.conditionedOn, + } + ) + and item.proposal.purpose == purpose + and ( + ( + item.proposal.conditionedOn is not None + if conditional + else len(item.proposal.explanatoryColumns) == 2 + ) + and item.proposal.essential + or purpose == "designCoverage" + and coverage[index].status == "computed" + and _measured_joint_design(item) + ) + ] + if matches: + for index in matches: + # A model's optional label cannot erase an explicit requirement. + # Reuse its measured question rather than adding a duplicate that + # would consume another slot in the bounded study contract. + if not requirements[index].essential: + requirements[index] = requirements[index].model_copy( + update={"essential": True} + ) + continue + identifier = "requestedDesign:" + hashlib.sha256(quote.encode()).hexdigest() + requirements.append( + DesignEvidenceRequirement( + requirementId=identifier, + question=quote, + objectiveQuote=quote, + kind=purpose, + columns=names, + ) + ) + coverage.append( + DesignEvidenceCoverage( + requirementId=identifier, + status="unsupported", + reasons=[ + "The explicitly requested joint or conditional comparison has not been nominated and measured; marginal comparisons do not answer it." + ], + ) + ) + if len(requirements) > 13: + raise ValueError( + "Objective requirements permit one design summary and eight plus four questions" + ) + return requirements, coverage + + +def unmet_objective_requirements( + requirements: list[DesignEvidenceRequirement], + coverage: list[DesignEvidenceCoverage], +) -> list[str]: + """Return essential questions with no measured answer of the required kind.""" + measured = {item.requirementId: item for item in coverage} + unmet = [] + for requirement in requirements: + item = measured.get(requirement.requirementId) + satisfied = item is not None and ( + item.status == "computed" + or item.status == "nonIdentifiable" + and requirement.kind in {"studyDesign", "designCoverage"} + ) + if requirement.essential and not satisfied: + reasons = ( + "; ".join(item.reasons) if item is not None else "evidence is missing" + ) + unmet.append(f"{requirement.question} Unresolved: {reasons}") + return unmet diff --git a/scarf/agent/experimental_context/study.py b/scarf/agent/experimental_context/study.py new file mode 100644 index 00000000..840f4e84 --- /dev/null +++ b/scarf/agent/experimental_context/study.py @@ -0,0 +1,342 @@ +"""Grounded study contracts for decision-driven RNA analysis.""" + +from collections.abc import Iterable +from typing import Any, Literal + +from pydantic import Field, model_validator + +from ..types import AgentDataModel +from .contracts import ( + CovariateComparison, + DesignEvidenceCoverage, + DesignEvidenceRequirement, +) +from .requirements import ( + active_batch_safety, + objective_evidence, + unmet_objective_requirements, +) + +type AuthorLabelPolicy = Literal["holdout", "preservation"] +type ProcessingGoal = Literal[ + "populationDiscovery", + "conditionPreservingDiscovery", +] +type CorrectionLicense = Literal[ + "safe", + "unsafeConfounded", + "indeterminate", + "notApplicable", +] + + +class StudyContract(AgentDataModel): + """Scientific authority and design constraints for one analysis.""" + + studyContext: str = "" + studyObjective: str = "" + processingGoal: ProcessingGoal = "populationDiscovery" + scientificQuestions: list[str] = Field(default_factory=list) + targetCohort: list[str] = Field(default_factory=list) + physicalCaptureColumn: str | None = None + independentUnitColumns: list[str] = Field(default_factory=list) + conditionColumns: list[str] = Field(default_factory=list) + technicalBatchColumns: list[str] = Field(default_factory=list) + protectedColumns: list[str] = Field(default_factory=list) + protectedCombinations: list[list[str]] = Field(default_factory=list) + unsupportedProtection: list[str] = Field(default_factory=list) + columnKinds: dict[str, Literal["categorical", "continuous"]] = Field( + default_factory=dict + ) + authorLabelPolicy: AuthorLabelPolicy = "holdout" + correctionLicense: CorrectionLicense = "notApplicable" + allowedClaims: list[str] = Field(default_factory=list) + unsupportedClaims: list[str] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + limitations: list[str] = Field(default_factory=list) + evidenceRequirements: list[DesignEvidenceRequirement] = Field( + min_length=1, max_length=13 + ) + evidenceCoverage: list[DesignEvidenceCoverage] = Field(min_length=1, max_length=13) + + @model_validator(mode="after") + def validate_contract(self) -> "StudyContract": + if not self.studyContext.strip(): + raise ValueError("studyContext must be non-empty") + if not self.studyObjective.strip(): + raise ValueError("studyObjective must be non-empty") + for field_name in ( + "independentUnitColumns", + "conditionColumns", + "technicalBatchColumns", + "protectedColumns", + "evidenceIds", + ): + values = getattr(self, field_name) + if len(values) != len(set(values)): + raise ValueError(f"{field_name} must not contain duplicates") + if self.physicalCaptureColumn in self.conditionColumns: + raise ValueError("A condition column cannot be the physical capture") + if set(self.technicalBatchColumns).intersection(self.conditionColumns): + raise ValueError("Technical batch columns cannot also be condition columns") + if self.correctionLicense == "safe" and not self.technicalBatchColumns: + raise ValueError("A safe correction license requires batch columns") + if any( + len(columns) != 2 + or len(set(columns)) != 2 + or not set(columns).issubset(self.protectedColumns) + for columns in self.protectedCombinations + ): + raise ValueError( + "Protected combinations require two distinct protected columns" + ) + requirements = {item.requirementId: item for item in self.evidenceRequirements} + coverage = {item.requirementId: item for item in self.evidenceCoverage} + if ( + len(requirements) != len(self.evidenceRequirements) + or len(coverage) != len(self.evidenceCoverage) + or requirements.keys() != coverage.keys() + ): + raise ValueError( + "Objective requirements and measured coverage must match uniquely" + ) + mandatory = requirements.get("studyDesign") + if ( + mandatory is None + or mandatory.kind != "studyDesign" + or not mandatory.essential + ): + raise ValueError( + "The essential measured studyDesign requirement cannot be omitted" + ) + for item in self.evidenceRequirements: + if item.objectiveQuote not in f"{self.studyContext}\n{self.studyObjective}": + raise ValueError("Objective requirements must quote exact study text") + if any( + set(item.evidenceIds) - set(self.evidenceIds) + for item in self.evidenceCoverage + ): + raise ValueError( + "Objective coverage cites evidence outside the study contract" + ) + if any( + item.status in {"computed", "nonIdentifiable"} and not item.evidenceIds + for item in self.evidenceCoverage + ): + raise ValueError( + "Computed objective coverage requires measured evidence IDs" + ) + return self + + @classmethod + def get_blank(cls) -> "StudyContract": + return cls( + studyContext="Study context", + studyObjective="Study objective", + evidenceRequirements=[ + DesignEvidenceRequirement( + requirementId="studyDesign", + question="Resolve study design", + objectiveQuote="Study objective", + kind="studyDesign", + ) + ], + evidenceCoverage=[ + DesignEvidenceCoverage( + requirementId="studyDesign", + status="unsupported", + reasons=["Study evidence is unavailable"], + ) + ], + ) + + +def _unique(values: Iterable[str | None]) -> list[str]: + return list(dict.fromkeys(value for value in values if value)) + + +def unsupported_comparison_limitations( + comparisons: Iterable[CovariateComparison], +) -> list[str]: + """Describe unsupported comparisons without implying a scientific finding.""" + + limitations = [] + for comparison in comparisons: + if comparison.status == "unsupported": + proposal = comparison.proposal + limitations.append( + f"Unresolved design comparison {comparison.evidenceId}: " + f"{proposal.response} against {', '.join(proposal.explanatoryColumns)} " + f"using observation unit {proposal.observationUnit!r} and " + f"independent unit {proposal.independentUnit or proposal.observationUnit!r}; " + f"reasons={', '.join(comparison.reasons) or 'unsupported evidence'}. " + "This comparison provides no supported association or absence finding." + ) + return limitations + + +def validate_objective_evidence( + contract: StudyContract, experimental_result: Any | None = None +) -> None: + """Reject unresolved essential questions and mismatched saved measured coverage.""" + # model_copy can bypass Pydantic validators; authority checks cannot. + StudyContract.model_validate(contract.model_dump(mode="json")) + if experimental_result is not None: + requirements, coverage = objective_evidence( + study_context=contract.studyContext, + study_objective=contract.studyObjective, + experimental_result=experimental_result, + ) + if ( + requirements != contract.evidenceRequirements + or coverage != contract.evidenceCoverage + ): + raise ValueError( + "Study objective evidence differs from its measured context report" + ) + unmet = unmet_objective_requirements( + contract.evidenceRequirements, contract.evidenceCoverage + ) + if unmet: + raise ValueError( + "Essential objective evidence is unresolved: " + " | ".join(unmet) + ) + + +def build_study_contract( + *, + study_context: str, + study_objective: str, + experimental_result: Any, + author_label_policy: AuthorLabelPolicy = "holdout", + physical_capture_column: str | None = None, +) -> StudyContract: + """Build a strict contract from one validated Experimental Context result.""" + + if getattr(experimental_result, "status", None) != "done": + raise ValueError("Experimental Context must be done before contract creation") + decision = experimental_result.decision + batch_plan = decision.batchCorrection + batch_safety = active_batch_safety(experimental_result) + conditions = list(decision.coefficientsOfInterest) + independent_units = _unique( + unit.independentUnit for unit in decision.unitsOfInference.values() + ) + protected = _unique( + [ + *conditions, + *independent_units, + *batch_plan.preserveColumns, + *( + column + for columns in decision.protectedCombinations + for column in columns + ), + ] + ) + assessed_batch_columns = _unique( + [ + *batch_plan.batchColumns, + *( + column + for assessment in batch_safety + for column in assessment.batchColumns + ), + ] + ) + correction_license: CorrectionLicense + if any(assessment.status == "unsafe" for assessment in batch_safety): + correction_license = "unsafeConfounded" + elif batch_plan.action == "evaluateHarmony": + correction_license = "safe" + elif batch_plan.action == "unsafe": + correction_license = "unsafeConfounded" + elif batch_plan.action == "needsInput": + correction_license = "indeterminate" + elif assessed_batch_columns and all( + assessment.status == "safe" for assessment in batch_safety + ): + correction_license = "safe" + else: + correction_license = "notApplicable" + processing_goal: ProcessingGoal = ( + "conditionPreservingDiscovery" if conditions else "populationDiscovery" + ) + evidence_ids = _unique( + [ + *decision.evidenceIds, + *batch_plan.evidenceIds, + *(item.evidenceId for item in experimental_result.batchSafety), + ] + ) + requirements, coverage = objective_evidence( + study_context=study_context, + study_objective=study_objective, + experimental_result=experimental_result, + ) + evidence_ids = _unique( + [ + *evidence_ids, + *(evidence_id for item in coverage for evidence_id in item.evidenceIds), + ] + ) + limitations = [ + *experimental_result.notes, + *(reason for item in coverage for reason in item.reasons), + *unsupported_comparison_limitations( + experimental_result.characterization.comparisons + ), + ] + if physical_capture_column is None: + limitations.append( + "Physical capture identity is unresolved; capture-aware doublet removal " + "is not authorized." + ) + if author_label_policy == "preservation": + limitations.append( + "Author labels were available for preservation checks; this run is " + "ineligible for label-based benchmark scoring." + ) + return StudyContract( + studyContext=study_context, + studyObjective=study_objective, + processingGoal=processing_goal, + scientificQuestions=[study_objective], + physicalCaptureColumn=physical_capture_column, + independentUnitColumns=independent_units, + conditionColumns=conditions, + technicalBatchColumns=assessed_batch_columns, + protectedColumns=protected, + protectedCombinations=[ + list(columns) for columns in decision.protectedCombinations + ], + unsupportedProtection=list(decision.unsupportedProtection), + columnKinds={ + record["name"]: record["kind"] + for record in experimental_result.characterization.columns + if record.get("kind") in {"categorical", "continuous"} + }, + authorLabelPolicy=author_label_policy, + correctionLicense=correction_license, + allowedClaims=[ + "Describe population structure supported by the selected representation.", + "Compare preprocessing alternatives against explicit evidence.", + ], + unsupportedClaims=[ + "This workflow does not test differential-expression hypotheses.", + "Cells are not independent biological replicates.", + ], + evidenceIds=evidence_ids, + limitations=limitations, + evidenceRequirements=requirements, + evidenceCoverage=coverage, + ) + + +__all__ = [ + "AuthorLabelPolicy", + "ProcessingGoal", + "StudyContract", + "build_study_contract", + "validate_objective_evidence", +] diff --git a/scarf/agent/experimental_context/tools.py b/scarf/agent/experimental_context/tools.py new file mode 100644 index 00000000..4f868f5c --- /dev/null +++ b/scarf/agent/experimental_context/tools.py @@ -0,0 +1,1208 @@ +"""Read-only Pydantic AI tools for experimental context.""" + +import hashlib +import json +import math +from collections.abc import Sequence +from copy import deepcopy +from functools import wraps +from types import SimpleNamespace +from typing import Any, Literal + +import numpy as np + +from ...metadata.queries import reduce_observation_units +from ...metrics.association import coefficient_estimability +from ...storage.artifacts import fingerprint_array +from ...storage.refs import ArtifactRef +from ...utils.logging import logger +from .._deps import AGENT_INSTALL_HINT +from ..record_io import canonical_json_bytes +from ..tools import artifact_reference, core_artifact_reference +from ..types import BatchSafetyEvidence, BatchSafetyStatus +from .characterization import characterize_covariates +from .comparisons import ( + DESIGN_ROUND_LIMITS, + accept_capture_proposal, + evaluate_proposals, +) +from .contracts import ( + BatchCorrectionPlan, + CaptureProposal, + ColumnDomain, + ContrastPlan, + ContrastStatus, + ContrastTest, + CovariateCharacterization, + CovariateEvidence, + CovariateProposal, + ExperimentalContextDecision, + ExperimentalContextDependencies, + InferenceUnit, + RepresentationEvaluation, + characterization_evidence, +) +from .qc_evidence import ( + _artifact_evidence_id, + _hto_artifact_map, + _hto_identity_columns, + _offered_qc_profiles, +) +from .requirements import ( + objective_evidence, + requested_design_purpose, + requested_design_questions, +) + +try: + from pydantic_ai import ModelRetry, RunContext + from pydantic_ai.tools import ToolDefinition +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + + +def compact_context_evidence(evidence: CovariateEvidence) -> dict[str, Any]: + """Present design findings; retain detailed policy measurements in the journal.""" + payload = evidence.model_dump(mode="json") + characterization = payload["characterization"] + coefficients = { + record["name"]: record for record in characterization["coefficients"] + } + for index, report in enumerate(characterization["confounding"]): + source = coefficients.get(report.get("coefficient"), {}) + for name in list(report): + if name in source and report[name] == source[name]: + report.pop(name) + report["coefficientDetails"] = f"coefficient:{report.get('coefficient')}" + report["details"] = {"section": "confounding", "record_id": str(index)} + for plan in payload["contrastPlans"]: + source = coefficients.get(plan["coefficient"], {}) + for name in ("replication", "estimability", "pairedCoverage"): + if plan[name] == source.get(name): + plan.pop(name) + plan["coefficientDetails"] = f"coefficient:{plan['coefficient']}" + # These tables duplicate the named coefficient records exactly. + for name in ( + "unitLevelCounts", + "groupImbalance", + "missingness", + "designStructures", + "pairedCoverage", + "coefficientEstimability", + ): + characterization.pop(name, None) + for column in characterization["columns"]: + counts = column.get("levelCounts", []) + if len(counts) > 8: + column["levelCounts"] = counts[:8] + column["levelCountsOmitted"] = len(counts) - 8 + column["details"] = f"column:{column['name']}" + + def omit_examples(value: Any) -> None: + if isinstance(value, dict): + examples = value.pop("incompleteExamples", None) + if examples is not None: + value["incompleteExamplesInSavedDetails"] = len(examples) + for child in value.values(): + omit_examples(child) + elif isinstance(value, list): + for child in value: + omit_examples(child) + + omit_examples(characterization) + omit_examples(payload["contrastPlans"]) + for record in characterization["coefficients"]: + record.pop("unitLevelCounts", None) + record["details"] = f"coefficient:{record['name']}" + shared_safety: dict[str, Any] = {} + for profile in payload["qcProfiles"]: + # The later QC decision receives complete thresholds and retention tables. + # Context needs capture provenance, adverse evidence and design constraints. + profile.pop("metricSources", None) + profile.pop("sourceConcordance", None) + profile.pop("resolvedBounds", None) + profile["parameters"] = { + key: value + for key, value in profile["parameters"].items() + if key not in {"resolvedBounds", "captureComparisons", "captureSizes"} + } + profile["details"] = f"qcProfile:{profile['profileId']}" + for failure in profile.get("captureFailureEvidence", []): + safety = failure.pop("conditionAndUnitSafety", []) + for check in safety: + if "requiredGroups" in check and "remainingGroups" in check: + check["lostGroups"] = [ + group + for group in check["requiredGroups"] + if group not in check["remainingGroups"] + ] + for name in ( + "requiredGroups", + "remainingGroups", + "observationUnitsByGroup", + "independentUnitsByGroup", + "quantilesBeforeExclusion", + "quantilesAfterExclusion", + ): + check.pop(name, None) + identity = hashlib.sha256(canonical_json_bytes(safety)).hexdigest() + shared_safety[identity] = safety + failure["designSafetyRef"] = identity + missingness = failure.pop("metricMissingFractions", {}) + failure["metricMissingness"] = { + "measuredSources": len(missingness), + "nonzeroOrUnavailable": { + key: value for key, value in missingness.items() if value != 0 + }, + } + payload["captureDesignSafety"] = shared_safety + payload["savedDetails"] = ( + "inspect_context_evidence returns one exact saved column, coefficient, " + "comparison, confounding record, or QC policy/capture. Policy thresholds, " + "individual donor examples and complete distributions remain saved; " + "summaries are not a substitute for required evidence." + ) + return payload + + +async def inspect_context_evidence( + ctx: RunContext[ExperimentalContextDependencies], + section: Literal["column", "coefficient", "comparison", "confounding", "qcProfile"], + record_id: str, + capture: str | None = None, +) -> dict[str, Any]: + """Read one saved evidence record without model calls or recomputation. + + Use a column/coefficient name, comparison evidenceId, zero-based confounding + index, or profileId. For capture-level QC detail supply one exact capture. + """ + characterization = ctx.deps.characterization + if characterization is None: + raise ModelRetry("Inspect covariates before requesting saved evidence") + if section == "qcProfile": + profile = ctx.deps.qcProfiles.get(record_id) + if profile is None: + raise ModelRetry("Choose an offered profileId") + if capture is not None: + failure = next( + ( + row + for row in profile.captureFailureEvidence + if row.capture == capture + ), + None, + ) + if failure is None: + raise ModelRetry("Choose one capture recorded in this QC profile") + bounds = profile.resolvedBounds + return { + "profileId": record_id, + "capture": failure.model_dump(mode="json"), + "resolvedBounds": [row for row in bounds if row.get("group") == capture] + if isinstance(bounds, list) + else deepcopy(bounds), + } + result = profile.model_dump(mode="json") + result.pop("captureFailureEvidence", None) + result.pop("metricSources", None) + result.pop("sourceConcordance", None) + result["parameters"].pop("captureComparisons", None) + result["parameters"].pop("resolvedBounds", None) + if profile.sampleColumn is not None or profile.sampleArtifact is not None: + result.pop("resolvedBounds", None) + result["captureDetailsRequired"] = ( + "Supply one capture to retrieve its exact thresholds and exclusion safety" + ) + return result + if capture is not None: + raise ModelRetry("A capture can be requested only for a QC profile") + if section == "comparison": + for comparison in characterization.comparisons: + if comparison.evidenceId == record_id: + return comparison.model_dump(mode="json") + elif section == "confounding": + if record_id.isdecimal() and int(record_id) < len(characterization.confounding): + return deepcopy(characterization.confounding[int(record_id)]) + raise ModelRetry( + "For section='confounding', record_id must be a zero-based index " + "string, not an evidence ID. Available record_id values: " + + json.dumps( + [str(index) for index in range(len(characterization.confounding))] + ) + ) + else: + records = ( + characterization.columns + if section == "column" + else characterization.coefficients + ) + for record in records: + if record["name"] == record_id: + return deepcopy(record) + raise ModelRetry("Choose an exact record from the saved context summary") + + +def capture_repair_inputs( + deps: ExperimentalContextDependencies, +) -> dict[str, Any] | None: + """Offer the exact assessed design for one missing-provenance repair.""" + saved_directions = deps.characterizationInputs.get("directions") + tested_batch_sets = { + tuple(sorted(item.batchColumns)) for item in deps.batchSafety.values() + } + if ( + deps.designRounds < len(DESIGN_ROUND_LIMITS) + or deps.captureProposal is not None + or not isinstance(saved_directions, dict) + or len(tested_batch_sets) > 1 + ): + return None + return { + "column_domains": deepcopy(saved_directions.get("columnDomains", {})), + "coefficients_of_interest": list( + saved_directions.get("coefficientsOfInterest", []) + ), + "units_of_inference": deepcopy(saved_directions.get("unitsOfInference", {})), + "batch_columns": list(next(iter(tested_batch_sets), ())), + "proposals": [], + } + + +def model_evidence_tool(function: Any) -> Any: + """Keep complete tool results in state and send a deduplicated model view.""" + + @wraps(function) + async def invoke(*args: Any, **kwargs: Any) -> dict[str, Any]: + result = await function(*args, **kwargs) + payload = compact_context_evidence(result) + context = args[0] if args else kwargs["ctx"] + repair_inputs = capture_repair_inputs(context.deps) + payload["captureRepairAvailable"] = repair_inputs is not None + if payload["captureRepairAvailable"]: + payload["captureRepairInstructions"] = ( + "analyze_experimental_design can validate capture_proposal using an " + "exact supplied study quote. Keep the assessed domains, coefficients, " + "units and batch columns unchanged and pass proposals=[]. This only " + "adds capture provenance and dependent QC evidence; it does not " + "authorize another covariate comparison round." + ) + payload["captureRepairInputs"] = repair_inputs + payload["requestedComparisons"] = [ + { + "question": quote, + "columns": columns, + "conditional": conditional, + "purpose": requested_design_purpose(quote), + "essential": True, + "completionEvidence": "evidenceRequirements and evidenceCoverage", + } + for quote, columns, conditional in requested_design_questions( + context.deps.studyContext, + context.deps.studyObjective, + [row["name"] for row in result.characterization.columns], + ) + ] + return payload + + return invoke + + +def persist_context_evidence(deps: ExperimentalContextDependencies, key: str) -> None: + """Commit completed context measurements to the owning stage journal.""" + if deps.checkpointWrite is not None: + deps.checkpointWrite( + key, + { + "state": deps.model_dump(mode="json"), + "characterizationInputs": deps.characterizationInputs, + }, + ) + + +def restore_context_evidence(deps: ExperimentalContextDependencies) -> bool: + """Restore complete evidence rounds without resetting proposal allowances.""" + if deps.checkpointRead is None: + return False + restored = False + for key in ("inspection", "design1", "design2", "capture"): + saved = deps.checkpointRead(key) + if saved is None: + continue + state = ExperimentalContextDependencies.model_validate(saved["state"]) + for name, field in ExperimentalContextDependencies.model_fields.items(): + if not field.exclude: + setattr(deps, name, getattr(state, name)) + deps.characterizationInputs = saved["characterizationInputs"] + restored = True + return restored + + +def characterize_context( + deps: ExperimentalContextDependencies, directions: dict[str, Any] +) -> CovariateCharacterization: + """Reuse characterization for the exact frozen stage and declared design.""" + metadata = {} + for column in deps.cells.columns: + values = np.asarray(deps.cells.fetch(column)) + metadata[column] = ( + hashlib.sha256( + repr( + [(type(value).__name__, value) for value in values.tolist()] + ).encode() + ).hexdigest() + if values.dtype.hasobject + else fingerprint_array(values) + ) + inputs = { + "directions": deepcopy(directions), + "cellSelection": deps.cellSelection.to_dict(), + "studyContext": deps.studyContext, + "studyObjective": deps.studyObjective, + "metadata": metadata, + } + if deps.characterization is not None and deps.characterizationInputs == inputs: + return deps.characterization + result = characterize_covariates( + deps.store, + cellSelection=deps.cellSelection, + studyContext=f"{deps.studyContext}\nStudy objective: {deps.studyObjective}", + model=None, + directions=directions, + groupingArtifacts=_hto_artifact_map(deps), + inventory=deps.inventoryData, + ) + if result.status != "failed": + deps.characterization = result + deps.characterizationInputs = inputs + return result + + +def _prepare_experimental_context_tool( + ctx: RunContext[ExperimentalContextDependencies], + tool_definition: ToolDefinition, +) -> ToolDefinition | None: + """Expose inspection once and at most two ordered design evidence rounds.""" + completed_calls = set(ctx.deps.toolCalls) + if tool_definition.name == "inspect_cell_covariates": + return None if tool_definition.name in completed_calls else tool_definition + if tool_definition.name == "analyze_experimental_design": + if ( + "inspect_cell_covariates" not in completed_calls + or ctx.deps.designRounds >= len(DESIGN_ROUND_LIMITS) + and capture_repair_inputs(ctx.deps) is None + ): + return None + return tool_definition + if tool_definition.name == "score_current_representation": + if ( + "analyze_experimental_design" not in completed_calls + or tool_definition.name in completed_calls + ): + return None + characterization = ctx.deps.characterization + if characterization is not None and not any( + record.get("domain") == "technical" and record.get("kind") == "categorical" + for record in characterization.columns + ): + return None + return tool_definition + return tool_definition + + +def contrast_plans_from_characterization( + characterization: CovariateCharacterization, +) -> list[ContrastPlan]: + """Build deterministic test licenses from bounded coefficient evidence.""" + reports = { + report.get("coefficient"): report + for report in characterization.confounding + if isinstance(report.get("coefficient"), str) + } + plans: list[ContrastPlan] = [] + for record in characterization.coefficients: + coefficient = record.get("name") + if not isinstance(coefficient, str): + continue + report = reports.get(coefficient, {}) + raw_groups = record.get("groupOrder") + group_order = ( + list(raw_groups) + if isinstance(raw_groups, list) + and all( + isinstance(value, str | int | float | bool) + and not (isinstance(value, float) and not math.isfinite(value)) + for value in raw_groups + ) + else [] + ) + sample_by = record.get("observationUnit") + sample_by = sample_by if isinstance(sample_by, str) else None + independent_unit = record.get("independentUnit") + independent_unit = ( + independent_unit if isinstance(independent_unit, str) else None + ) + between_unit = record.get("scope") == "betweenUnit" + replication = dict(record.get("replication") or {}) + replication_passed = replication.get("sufficient") is True + estimability = dict( + record.get("estimability") or report.get("estimability") or {} + ) + estimability_passed = ( + estimability.get("status") == "ok" + and estimability.get("coefficientEstimable") is True + and estimability.get("rankDeficient") is not True + ) + paired_coverage = dict(record.get("pairedCoverage") or {}) + pair_by: str | None = None + paired_passed: bool | None = None + mixed_independent_design = False + if independent_unit is not None: + if paired_coverage.get("complete") is True: + pair_by = independent_unit + paired_passed = True + elif paired_coverage.get("betweenIndependentUnits") is True: + sample_by = independent_unit + else: + pair_by = independent_unit + paired_passed = False + mixed_independent_design = True + + reasons: list[str] = [] + needs_input = False + if record.get("kind") != "categorical": + reasons.append("coefficientRequiresExplicitCategoricalGroups") + needs_input = True + if not between_unit: + reasons.append("coefficientIsNotBetweenUnit") + if sample_by is None: + reasons.append("sampleByIsUnresolved") + needs_input = True + if len(group_order) < 2: + reasons.append("fewerThanTwoObservedGroups") + needs_input = True + if record.get("groupCountsTruncated") is True: + reasons.append("groupOrderIsTruncated") + needs_input = True + if not replication_passed: + reasons.append("insufficientIndependentReplication") + if not estimability_passed: + reasons.append("coefficientIsNotEstimable") + + test: ContrastTest | None = None + if pair_by is not None: + if len(group_order) != 2: + reasons.append("pairedTestsRequireExactlyTwoGroups") + else: + test = "wilcoxon" + if paired_passed is not True: + reasons.append("pairedCoverageIsIncomplete") + elif mixed_independent_design: + reasons.append("independentUnitStructureIsMixed") + elif len(group_order) == 2: + test = "mann_whitney" + elif len(group_order) >= 3: + test = "kruskal_wallis" + + reasons = list(dict.fromkeys(reasons)) + status: ContrastStatus = ( + "licensed" if not reasons else "needsInput" if needs_input else "blocked" + ) + evidence_ids = [ + f"column:{coefficient}", + f"coefficient:{coefficient}", + f"estimability:{coefficient}", + ] + plans.append( + ContrastPlan( + coefficient=coefficient, + groupOrder=group_order, + sampleBy=sample_by, + pairBy=pair_by, + test=test, + status=status, + betweenUnitDesign=between_unit, + replicationPassed=replication_passed, + estimabilityPassed=estimability_passed, + pairedCoveragePassed=paired_passed, + replication=replication, + estimability=estimability, + pairedCoverage=paired_coverage, + blockedReasons=reasons, + evidenceId=f"contrastPlan:{coefficient}:{status}", + evidenceIds=evidence_ids, + ) + ) + return plans + + +async def inspect_cell_covariates( + ctx: RunContext[ExperimentalContextDependencies], +) -> CovariateEvidence: + """Inspect cell metadata without making model-driven choices or writing data.""" + logger.info( + "Experimental Context covariate inspection started: " + f"cellSelection={ctx.deps.cellSelection.artifact_id}" + ) + ctx.deps.htoIdentityColumns = _hto_identity_columns(ctx.deps) + characterization = characterize_context(ctx.deps, ctx.deps.directions) + ctx.deps.characterization = characterization + qc_profiles = _offered_qc_profiles(ctx.deps) + contrast_plans = contrast_plans_from_characterization(characterization) + ctx.deps.contrastPlans = {plan.coefficient: plan for plan in contrast_plans} + evidence_ids = characterization_evidence(characterization) + evidence_ids.update(profile.evidenceId for profile in qc_profiles) + evidence_ids.update(source.sourceId for source in ctx.deps.qcMetricSources) + evidence_ids.update(item.evidenceId for item in ctx.deps.qcSourceConcordance) + evidence_ids.update(plan.evidenceId for plan in contrast_plans) + evidence_ids.update( + failure.evidenceId + for profile in qc_profiles + for failure in profile.captureFailureEvidence + ) + evidence_ids.update( + f"htoIdentity:{column}" for column in ctx.deps.htoIdentityColumns + ) + evidence_ids.update( + _artifact_evidence_id(source) for source in ctx.deps.htoIdentityArtifacts + ) + ctx.deps.evidenceIds.update(evidence_ids) + ctx.deps.toolCalls.append("inspect_cell_covariates") + persist_context_evidence(ctx.deps, "inspection") + logger.info( + "Experimental Context covariate inspection completed: " + f"status={characterization.status}, " + f"columns={len(characterization.columns)}, " + f"coefficients={len(characterization.coefficients)}, " + f"qcProfiles={len(qc_profiles)}, " + f"htoIdentities={len(ctx.deps.htoIdentityColumns)}, " + f"evidence={len(evidence_ids)}" + ) + return CovariateEvidence( + characterization=characterization, + qcProfiles=qc_profiles, + qcMetricSources=ctx.deps.qcMetricSources, + qcSourceConcordance=ctx.deps.qcSourceConcordance, + contrastPlans=contrast_plans, + htoIdentityColumns=ctx.deps.htoIdentityColumns, + htoIdentityArtifacts=ctx.deps.htoIdentityArtifacts, + evidenceIds=sorted(evidence_ids), + ) + + +def _batch_safety_evidence( + deps: ExperimentalContextDependencies, + characterization: CovariateCharacterization, + *, + coefficients: Sequence[str], + batch_columns: Sequence[str], +) -> list[BatchSafetyEvidence]: + column_records = { + record.get("name"): record + for record in characterization.columns + if isinstance(record.get("name"), str) + } + coefficient_records = { + record.get("name"): record + for record in characterization.coefficients + if isinstance(record.get("name"), str) + } + confounding_reports = { + report.get("coefficient"): report + for report in characterization.confounding + if isinstance(report.get("coefficient"), str) + } + canonical_batch_columns = sorted(batch_columns) + batch_safety: list[BatchSafetyEvidence] = [] + for coefficient in coefficients: + if not canonical_batch_columns: + break + coefficient_record = coefficient_records.get(coefficient) + report = confounding_reports.get(coefficient) + coefficient_kind = ( + coefficient_record.get("kind") if coefficient_record is not None else None + ) + if coefficient_kind not in {"categorical", "continuous"}: + coefficient_kind = None + observation_unit = ( + report.get("observationUnit") + if report is not None + else ( + coefficient_record.get("observationUnit") + if coefficient_record is not None + else None + ) + ) + unit_constant = { + pair.get("technical") + for pair in (report.get("pairs", []) if report is not None else []) + if isinstance(pair.get("technical"), str) + } + effective_batch_columns = [ + name for name in canonical_batch_columns if name in unit_constant + ] + estimability: dict[str, Any] + if ( + coefficient_record is None + or coefficient_record.get("scope") != "betweenUnit" + or report is None + or not isinstance(observation_unit, str) + or coefficient_kind is None + ): + estimability = { + "status": "notComputed", + "reason": "unresolvedCoefficientDesign", + } + else: + try: + design = reduce_observation_units( + deps.cells, + observation_unit, + [coefficient, *effective_batch_columns], + cell_key="I", + ) + estimability = coefficient_estimability( + design[coefficient].to_numpy(), + coefficientKind=coefficient_kind, + technicals={ + name: design[name].to_numpy() + for name in effective_batch_columns + }, + technicalKinds={ + name: column_records[name]["kind"] + for name in effective_batch_columns + }, + ) + except (KeyError, TypeError, ValueError) as exc: + logger.debug( + "Experimental Context batch estimability was not computed: " + f"errorType={type(exc).__name__}" + ) + estimability = { + "status": "notComputed", + "reason": type(exc).__name__, + } + if estimability.get("status") != "ok": + safety_status: BatchSafetyStatus = "notComputed" + elif estimability.get("coefficientEstimable") is True and not bool( + estimability.get("rankDeficient") + ): + safety_status = "safe" + else: + safety_status = "unsafe" + # Sparse descriptive associations cannot erase a computed design constraint. + joint_checks = [ + { + "evidenceId": comparison.evidenceId, + **comparison.evidence["jointGroupEstimability"], + } + for comparison in characterization.comparisons + if comparison.proposal.response == coefficient + and comparison.proposal.observationUnit == observation_unit + and comparison.proposal.conditionedOn is None + and all( + column_records.get(name, {}).get("kind") == kind + for name, kind in comparison.evidence.get("columnKinds", {}).items() + ) + and set(comparison.proposal.explanatoryColumns).issubset( + canonical_batch_columns + ) + and "jointGroupEstimability" in comparison.evidence + ] + if joint_checks: + estimability = {**estimability, "jointStratumChecks": joint_checks} + if any( + check.get("status") == "ok" + and check.get("coefficientEstimable") is False + for check in joint_checks + ): + safety_status = "unsafe" + batch_token = ",".join(canonical_batch_columns) + safety = BatchSafetyEvidence( + coefficient=coefficient, + coefficientKind=coefficient_kind, + observationUnit=( + observation_unit if isinstance(observation_unit, str) else None + ), + batchColumns=canonical_batch_columns, + unitConstantBatchColumns=effective_batch_columns, + status=safety_status, + estimability=estimability, + evidenceId=f"batchEstimability:{coefficient}:{batch_token}", + ) + batch_safety.append(safety) + deps.batchSafety[safety.evidenceId] = safety + return batch_safety + + +async def analyze_experimental_design( + ctx: RunContext[ExperimentalContextDependencies], + column_domains: dict[str, ColumnDomain], + coefficients_of_interest: list[str], + units_of_inference: dict[str, InferenceUnit], + batch_columns: list[str], + proposals: list[CovariateProposal] | None = None, + capture_proposal: CaptureProposal | None = None, +) -> CovariateEvidence: + """Validate proposed domains and inference units and compute confounding. + + Args: + ctx: Pydantic AI run context containing the existing datastore. + column_domains: Domain assignment for each metadata column under review. + coefficients_of_interest: Biological columns representing study contrasts. + units_of_inference: Observation and independent units for each coefficient. + batch_columns: Exact technical columns proposed for Harmony evaluation. + proposals: Up to eight initial or four follow-up objective-led comparisons. + Each uses distinct response/explanatory/conditioning columns: one or + two explanatory columns without conditioning, or one explanatory + column with one categorical conditioning column. Observation and + independent units do not count toward the three-column limit. + Two explanatory columns together with conditioning are unsupported. + One explanatory column within one categorical stratum is supported + subject to observed unit support; descriptive joint counts can remain + available when the association method is unsupported. Do not discard + a scientific question or a unit identity just to fit this schema. + capture_proposal: Exact capture and baseline identities supported by study prose. + After two comparison rounds, only this provenance repair is permitted: + keep the assessed design unchanged and pass no proposals. The tool + validates the supplied quote and refreshes dependent QC evidence. + """ + logger.info( + "Experimental Context design analysis started: " + f"domains={len(column_domains)}, " + f"coefficients={len(coefficients_of_interest)}, " + f"inferenceUnits={len(units_of_inference)}, " + f"batchColumns={len(batch_columns)}" + ) + capture_repair = ctx.deps.designRounds >= len(DESIGN_ROUND_LIMITS) + if capture_repair and (capture_proposal is None or proposals): + raise ModelRetry( + "Design comparison permits at most two evidence rounds. Only capture " + "provenance can now be repaired: supply capture_proposal, proposals=[], " + "and the exact unchanged assessed design." + ) + if ( + not capture_repair + and len(proposals or ()) > DESIGN_ROUND_LIMITS[ctx.deps.designRounds] + ): + raise ModelRetry( + "Design comparison permits eight initial and four follow-up proposals" + ) + directions = dict(ctx.deps.directions) + directed_domains = dict(column_domains) + directed_domains.update(dict(directions.get("columnDomains") or {})) + directions["columnDomains"] = directed_domains + directed_coefficients = list( + dict.fromkeys( + [ + *coefficients_of_interest, + *(directions.get("coefficientsOfInterest") or []), + ] + ) + ) + directions["coefficientsOfInterest"] = directed_coefficients + directed_units = { + name: unit.model_dump(exclude_none=True) + for name, unit in units_of_inference.items() + } + directed_units.update(dict(directions.get("unitsOfInference") or {})) + directions["unitsOfInference"] = directed_units + + proposed_batch_columns = list(batch_columns) + directed_batch_columns = directions.get("batchColumns") + if directed_batch_columns is not None: + if not isinstance(directed_batch_columns, list) or any( + not isinstance(value, str) or not value.strip() + for value in directed_batch_columns + ): + raise ModelRetry( + "directions.batchColumns must be a list of exact metadata columns" + ) + if len(set(directed_batch_columns)) != len(directed_batch_columns): + raise ModelRetry("directions.batchColumns must be unique") + if proposed_batch_columns != directed_batch_columns: + logger.info( + "Experimental Context replaced model-proposed batch columns with " + "the exact directed columns" + ) + proposed_batch_columns = list(directed_batch_columns) + canonical_batch_columns = sorted(set(proposed_batch_columns)) + if len(canonical_batch_columns) != len(proposed_batch_columns): + logger.warning( + "Experimental Context rejected duplicate proposed batch columns: " + f"{proposed_batch_columns[:20]}" + ) + raise ModelRetry("Proposed batch columns must be unique") + inspected_records = { + record.get("name"): record + for record in ( + ctx.deps.characterization.columns + if ctx.deps.characterization is not None + else [] + ) + if isinstance(record.get("name"), str) + } + if ctx.deps.characterization is not None: + for batch_column in canonical_batch_columns: + inspected = inspected_records.get(batch_column) + if inspected is None: + logger.warning( + "Experimental Context rejected unknown proposed batch column " + f"before design recomputation: {batch_column!r}" + ) + raise ModelRetry(f"Unknown batch column {batch_column!r}") + proposed_domain = directed_domains.get( + batch_column, + inspected.get("domain"), + ) + if proposed_domain != "technical": + logger.warning( + "Experimental Context rejected proposed batch column before " + f"design recomputation: {batch_column!r}, " + f"domain={proposed_domain!r}, required='technical'" + ) + raise ModelRetry( + f"Batch column {batch_column!r} must be classified as technical" + ) + if inspected.get("kind") != "categorical": + logger.warning( + "Experimental Context rejected proposed batch column before " + f"design recomputation: {batch_column!r}, " + f"kind={inspected.get('kind')!r}, required='categorical'" + ) + raise ModelRetry( + f"Batch column {batch_column!r} must be categorical for Harmony" + ) + + if capture_repair: + characterization = ctx.deps.characterization + if ( + characterization is None + or directions != ctx.deps.characterizationInputs.get("directions") + ): + raise ModelRetry( + "Capture provenance repair requires the exact saved domains, " + "coefficients and inference units; it cannot change the design." + ) + tested_batch_sets = { + tuple(sorted(item.batchColumns)) for item in ctx.deps.batchSafety.values() + } + if tested_batch_sets != {tuple(canonical_batch_columns)} and ( + tested_batch_sets or canonical_batch_columns + ): + raise ModelRetry( + "Capture provenance repair requires one unambiguous, unchanged " + "assessed batch-column set. Alternative designs cannot be combined." + ) + if ctx.deps.captureProposal is not None: + raise ModelRetry( + "Capture provenance is already committed; reuse its saved evidence." + ) + logger.info( + "Experimental Context capture provenance: validating the supplied study " + "statement and refreshing dependent QC; reusing measured study design" + ) + else: + characterization = characterize_context(ctx.deps, directions) + if characterization.status == "failed": + rejection = "; ".join(characterization.notes).strip() + logger.warning( + "Experimental Context design characterization rejected the proposed " + f"directions: {rejection[:1000]}; " + f"domainColumns={sorted(column_domains)[:50]}, " + f"coefficients={coefficients_of_interest[:50]}, " + f"inferenceUnits={sorted(units_of_inference)[:50]}" + ) + raise ModelRetry("; ".join(characterization.notes)) + + # Retain the validated deterministic work even when the proposed Harmony + # columns below are rejected. A bounded retry or resumed decision can reuse + # the evidence without rescanning metadata or accepting an unsafe choice. + ctx.deps.characterization = characterization + try: + if not capture_repair: + evaluate_proposals(ctx.deps, characterization, proposals or ()) + if capture_proposal is not None: + accept_capture_proposal(ctx.deps, characterization, capture_proposal) + except ValueError as exc: + raise ModelRetry(str(exc)) from exc + if not ctx.deps.htoIdentityColumns: + ctx.deps.htoIdentityColumns = _hto_identity_columns(ctx.deps) + try: + qc_profiles = _offered_qc_profiles(ctx.deps, characterization) + except Exception: + if capture_repair: + ctx.deps.captureProposal = None + characterization.captureProvenance = None + raise + contrast_plans = contrast_plans_from_characterization(characterization) + ctx.deps.contrastPlans = {plan.coefficient: plan for plan in contrast_plans} + evidence_ids = characterization_evidence(characterization) + evidence_ids.update(profile.evidenceId for profile in qc_profiles) + evidence_ids.update(source.sourceId for source in ctx.deps.qcMetricSources) + evidence_ids.update(item.evidenceId for item in ctx.deps.qcSourceConcordance) + evidence_ids.update(plan.evidenceId for plan in contrast_plans) + evidence_ids.update( + failure.evidenceId + for profile in qc_profiles + for failure in profile.captureFailureEvidence + ) + evidence_ids.update( + f"htoIdentity:{column}" for column in ctx.deps.htoIdentityColumns + ) + evidence_ids.update( + _artifact_evidence_id(source) for source in ctx.deps.htoIdentityArtifacts + ) + ctx.deps.evidenceIds.update(evidence_ids) + + column_records = { + record.get("name"): record + for record in characterization.columns + if isinstance(record.get("name"), str) + } + for batch_column in canonical_batch_columns: + record = column_records.get(batch_column) + if record is None: + logger.warning( + "Experimental Context rejected unknown proposed batch column: " + f"{batch_column!r}" + ) + raise ModelRetry(f"Unknown batch column {batch_column!r}") + if record.get("domain") != "technical": + logger.warning( + "Experimental Context rejected proposed batch column " + f"{batch_column!r}: domain={record.get('domain')!r}, " + "required='technical'" + ) + raise ModelRetry( + f"Batch column {batch_column!r} must be classified as technical" + ) + if record.get("kind") != "categorical": + logger.warning( + "Experimental Context rejected proposed batch column " + f"{batch_column!r}: kind={record.get('kind')!r}, " + "required='categorical'" + ) + raise ModelRetry( + f"Batch column {batch_column!r} must be categorical for Harmony" + ) + + batch_safety = ( + list(ctx.deps.batchSafety.values()) + if capture_repair + else _batch_safety_evidence( + ctx.deps, + characterization, + coefficients=directed_coefficients, + batch_columns=canonical_batch_columns, + ) + ) + + evidence_ids.update(item.evidenceId for item in batch_safety) + ctx.deps.evidenceIds.update(evidence_ids) + ctx.deps.toolCalls.append("analyze_experimental_design") + persist_context_evidence( + ctx.deps, "capture" if capture_repair else f"design{ctx.deps.designRounds}" + ) + safety_counts = { + status: sum(item.status == status for item in batch_safety) + for status in ("safe", "unsafe", "notComputed") + } + logger.info( + "Experimental Context design analysis completed: " + f"status={characterization.status}, " + f"batchSafetySafe={safety_counts['safe']}, " + f"batchSafetyUnsafe={safety_counts['unsafe']}, " + f"batchSafetyNotComputed={safety_counts['notComputed']}, " + f"qcProfiles={len(qc_profiles)}, evidence={len(evidence_ids)}" + ) + requirements, coverage = ( + objective_evidence( + study_context=ctx.deps.studyContext, + study_objective=ctx.deps.studyObjective, + experimental_result=SimpleNamespace( + characterization=characterization, + decision=ExperimentalContextDecision( + coefficientsOfInterest=directed_coefficients, + batchCorrection=BatchCorrectionPlan( + action="needsInput", + batchColumns=canonical_batch_columns, + ), + ), + batchSafety=batch_safety, + ), + ) + if ctx.deps.studyObjective + else ([], []) + ) + return CovariateEvidence( + characterization=characterization, + batchSafety=batch_safety, + evidenceRequirements=requirements, + evidenceCoverage=coverage, + qcProfiles=qc_profiles, + qcMetricSources=ctx.deps.qcMetricSources, + qcSourceConcordance=ctx.deps.qcSourceConcordance, + contrastPlans=contrast_plans, + htoIdentityColumns=ctx.deps.htoIdentityColumns, + htoIdentityArtifacts=ctx.deps.htoIdentityArtifacts, + evidenceIds=sorted(evidence_ids), + ) + + +async def score_current_representation( + ctx: RunContext[ExperimentalContextDependencies], + batch_column: str, + biological_column: str | None = None, +) -> RepresentationEvaluation: + """Score one explicitly supplied graph without changing datastore state. + + Args: + ctx: Pydantic AI run context containing the existing datastore. + batch_column: Categorical technical column used to assess batch mixing. + biological_column: Optional biological label used to assess preservation. + """ + logger.info( + "Experimental Context representation scoring started: " + f"graphSupplied={ctx.deps.neighbors is not None}, " + f"biologicalLabelSpecified={biological_column is not None}" + ) + store = ctx.deps.store + available_columns = set(store.cells.columns) + if batch_column not in available_columns: + raise ModelRetry(f"Unknown batch column {batch_column!r}") + if biological_column is not None and biological_column not in available_columns: + raise ModelRetry(f"Unknown biological column {biological_column!r}") + characterization = ctx.deps.characterization + if characterization is not None: + batch_record = next( + ( + record + for record in characterization.columns + if record.get("name") == batch_column + ), + None, + ) + if ( + batch_record is None + or batch_record.get("domain") != "technical" + or batch_record.get("kind") != "categorical" + ): + raise ModelRetry( + "Representation scoring requires a characterized categorical " + "technical batch column" + ) + + neighbors = core_artifact_reference(ctx.deps.neighbors) + connectivity = core_artifact_reference(ctx.deps.connectivityMap) + if neighbors is None: + evaluation = RepresentationEvaluation( + cellSelection=( + artifact_reference(ctx.deps.cellSelection) + if ctx.deps.cellSelection is not None + else None + ), + notes=["No exact neighbors artifact was supplied"], + ) + ctx.deps.currentRepresentation = evaluation + ctx.deps.toolCalls.append("score_current_representation") + logger.info( + "Experimental Context representation scoring skipped: " + "no current neighbors artifact" + ) + return evaluation + if not isinstance(neighbors, ArtifactRef) or neighbors.kind != "neighbors": + raise ModelRetry("neighbors must identify an exact neighbors artifact") + if connectivity is not None and ( + not isinstance(connectivity, ArtifactRef) + or connectivity.kind not in {"connectivity_map", "integrated_graph"} + ): + raise ModelRetry( + "connectivity_map must identify an exact connectivity graph artifact" + ) + + metrics: dict[str, float] = {} + notes: list[str] = [] + evidence_ids: list[str] = [] + neighbor_route = f"assay:{neighbors.assay}:neighbors:{neighbors.artifact_id}" + try: + value = float(store.metric_ilisi(batch_column, neighbors)) + if math.isfinite(value): + metrics[f"iLISI:{batch_column}"] = value + evidence_ids.append(f"metric:iLISI:{batch_column}:{neighbor_route}") + except (KeyError, TypeError, ValueError, RuntimeError) as exc: + logger.debug( + "Experimental Context iLISI scoring was unavailable: " + f"errorType={type(exc).__name__}" + ) + notes.append(f"iLISI could not be scored: {exc}") + try: + value = float(store.metric_proportional_batch_mixing(batch_column, neighbors)) + if math.isfinite(value): + metrics[f"proportionalBatchMixing:{batch_column}"] = value + evidence_ids.append( + f"metric:proportionalBatchMixing:{batch_column}:{neighbor_route}" + ) + except (KeyError, TypeError, ValueError, RuntimeError) as exc: + logger.debug( + "Experimental Context batch-mixing scoring was unavailable: " + f"errorType={type(exc).__name__}" + ) + notes.append(f"Proportional batch mixing could not be scored: {exc}") + if biological_column is not None: + try: + value = float(store.metric_clisi(biological_column, neighbors)) + if math.isfinite(value): + metrics[f"cLISI:{biological_column}"] = value + evidence_ids.append( + f"metric:cLISI:{biological_column}:{neighbor_route}" + ) + except (KeyError, TypeError, ValueError, RuntimeError) as exc: + logger.debug( + "Experimental Context cLISI scoring was unavailable: " + f"errorType={type(exc).__name__}" + ) + notes.append(f"cLISI could not be scored: {exc}") + if connectivity is not None: + try: + value = float( + store.metric_graph_connectivity(biological_column, connectivity) + ) + if math.isfinite(value): + metrics[f"graphConnectivity:{biological_column}"] = value + evidence_ids.append( + "metric:graphConnectivity:" + f"{biological_column}:assay:{connectivity.assay}:connectivity:" + f"{connectivity.artifact_id}" + ) + except (KeyError, TypeError, ValueError, RuntimeError) as exc: + logger.debug( + "Experimental Context connectivity scoring was unavailable: " + f"errorType={type(exc).__name__}" + ) + notes.append(f"Graph connectivity could not be scored: {exc}") + + evaluation = RepresentationEvaluation( + available=bool(metrics), + assay=neighbors.assay, + cellSelection=( + artifact_reference(ctx.deps.cellSelection) + if ctx.deps.cellSelection is not None + else None + ), + neighbors=artifact_reference(neighbors), + connectivityMap=( + artifact_reference(connectivity) if connectivity is not None else None + ), + metrics=metrics, + notes=notes, + evidenceIds=evidence_ids, + ) + ctx.deps.currentRepresentation = evaluation + ctx.deps.evidenceIds.update(evidence_ids) + ctx.deps.toolCalls.append("score_current_representation") + logger.info( + "Experimental Context representation scoring completed: " + f"available={evaluation.available}, metrics={len(metrics)}, " + f"notes={len(notes)}, evidence={len(evidence_ids)}" + ) + return evaluation diff --git a/scarf/agent/experimental_context/validation.py b/scarf/agent/experimental_context/validation.py new file mode 100644 index 00000000..ee882fc6 --- /dev/null +++ b/scarf/agent/experimental_context/validation.py @@ -0,0 +1,506 @@ +"""Experimental-context canonicalization and explicit model failures.""" + +from types import SimpleNamespace +from typing import Any + +from ...utils.logging import logger +from .._deps import AGENT_INSTALL_HINT +from ..tools import artifact_reference +from ..types import AgentRunInfo, BatchSafetyEvidence +from .comparisons import canonical_design_choices +from .contracts import ( + CellQcPlan, + CovariateCharacterization, + ExperimentalContextDecision, + ExperimentalContextDependencies, + ExperimentalContextResult, + InferenceUnit, + characterization_evidence, +) +from .qc_evidence import ( + _offered_qc_profiles, +) +from .requirements import ( + active_batch_safety, + objective_evidence, + unmet_objective_requirements, +) +from .tools import characterize_context, contrast_plans_from_characterization + +try: + from pydantic import ValidationError + from pydantic_ai import ModelRetry +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + + +def _validate_batch_correction_plan( + decision: ExperimentalContextDecision, + deps: ExperimentalContextDependencies, + characterization: CovariateCharacterization, + requested_coefficients: set[str], + units_of_inference: dict[str, dict[str, Any]], + records: dict[str, dict[str, Any]], + coefficient_records: dict[str, dict[str, Any]], +) -> None: + """Validate one batch plan against exact design, safety, and metric evidence.""" + confounding_reports = { + report.get("coefficient"): report + for report in characterization.confounding + if isinstance(report.get("coefficient"), str) + } + plan = decision.batchCorrection + try: + active_batch_safety( + SimpleNamespace( + decision=decision, batchSafety=list(deps.batchSafety.values()) + ) + ) + except ValueError as exc: + raise ModelRetry(str(exc)) from exc + directed_batch_columns = deps.directions.get("batchColumns") + if directed_batch_columns is not None: + if not isinstance(directed_batch_columns, list) or any( + not isinstance(value, str) or not value.strip() + for value in directed_batch_columns + ): + raise ModelRetry( + "directions.batchColumns must be a list of exact metadata columns" + ) + canonical_directed_batch = sorted(directed_batch_columns) + directed_plan_mismatch = ( + ( + plan.action not in {"evaluateHarmony", "unsafe"} + or sorted(plan.batchColumns) != canonical_directed_batch + ) + if canonical_directed_batch + else plan.action != "skip" or bool(plan.batchColumns) + ) + if directed_plan_mismatch: + raise ModelRetry( + "The batch-correction plan must assess the exact directed batch " + f"columns: {canonical_directed_batch}" + ) + unknown_columns = sorted(set(decision.columnDomains) - set(records)) + if unknown_columns: + raise ModelRetry(f"Unknown column domain assignments: {unknown_columns}") + unit_columns = { + unit_name + for unit in units_of_inference.values() + for unit_name in ( + unit.get("observationUnit"), + unit.get("independentUnit"), + ) + if isinstance(unit_name, str) + } + if plan.action == "evaluateHarmony" and not plan.batchColumns: + raise ModelRetry("evaluateHarmony requires at least one batch column") + if plan.action == "unsafe" and not plan.batchColumns: + raise ModelRetry("unsafe requires the exact batch columns that were assessed") + if plan.action == "skip" and plan.batchColumns: + raise ModelRetry("skip must not include batch columns") + if plan.action == "needsInput" and not decision.needsInput: + raise ModelRetry("needsInput action requires at least one concrete question") + if len(set(plan.batchColumns)) != len(plan.batchColumns): + raise ModelRetry("Batch columns must be unique") + + for batch_column in plan.batchColumns: + record = records.get(batch_column) + if record is None: + raise ModelRetry(f"Unknown batch column {batch_column!r}") + if record.get("domain") != "technical": + raise ModelRetry( + f"Batch column {batch_column!r} must be classified as technical" + ) + if record.get("kind") != "categorical": + raise ModelRetry( + f"Batch column {batch_column!r} must be categorical for Harmony" + ) + if batch_column in requested_coefficients or batch_column in unit_columns: + raise ModelRetry( + f"Batch column {batch_column!r} cannot be a coefficient or unit of inference" + ) + + if plan.action == "evaluateHarmony": + mixing_metrics = {"iLISI", "proportionalBatchMixing"} + preservation_metrics = {"cLISI", "graphConnectivity"} + if not mixing_metrics.intersection(plan.metricsRequired): + raise ModelRetry( + "evaluateHarmony requires iLISI or proportionalBatchMixing" + ) + if any( + records.get(column, {}).get("kind") == "categorical" + for column in plan.preserveColumns + ) and not preservation_metrics.intersection(plan.metricsRequired): + raise ModelRetry( + "evaluateHarmony requires cLISI or graphConnectivity for preservation" + ) + missing_preserve = sorted(requested_coefficients - set(plan.preserveColumns)) + if missing_preserve: + raise ModelRetry( + "preserveColumns must include every coefficient of interest: " + f"{missing_preserve}" + ) + unresolved_coefficients = sorted( + coefficient + for coefficient in requested_coefficients + if coefficient_records[coefficient].get("scope") != "betweenUnit" + or coefficient not in confounding_reports + ) + if unresolved_coefficients: + raise ModelRetry( + "evaluateHarmony requires a between-unit coefficient with a " + "matching estimability report; use needsInput or unsafe for: " + f"{unresolved_coefficients}" + ) + for preserve_column in plan.preserveColumns: + record = records.get(preserve_column) + if record is None: + raise ModelRetry(f"Unknown preservation column {preserve_column!r}") + if record.get("domain") != "biological": + raise ModelRetry( + f"Preservation column {preserve_column!r} must be biological" + ) + + matched_safety: list[BatchSafetyEvidence] = [] + if plan.action in {"evaluateHarmony", "unsafe"}: + canonical_batch_columns = sorted(plan.batchColumns) + for coefficient in sorted(requested_coefficients): + coefficient_record = coefficient_records[coefficient] + report = confounding_reports.get(coefficient) + observation_unit = ( + report.get("observationUnit") + if report is not None + else coefficient_record.get("observationUnit") + ) + unit_constant = { + pair.get("technical") + for pair in (report.get("pairs", []) if report is not None else []) + if isinstance(pair.get("technical"), str) + } + expected_effective = [ + name for name in canonical_batch_columns if name in unit_constant + ] + candidates = [ + item + for item in deps.batchSafety.values() + if item.coefficient == coefficient + and item.coefficientKind == coefficient_record.get("kind") + and item.observationUnit == observation_unit + and item.batchColumns == canonical_batch_columns + and item.unitConstantBatchColumns == expected_effective + ] + if len(candidates) != 1: + raise ModelRetry( + "Call analyze_experimental_design with the exact proposed batch " + f"columns before returning a recommendation for {coefficient!r}" + ) + matched_safety.append(candidates[0]) + missing_safety_evidence = sorted( + item.evidenceId + for item in matched_safety + if item.evidenceId not in plan.evidenceIds + ) + if missing_safety_evidence: + raise ModelRetry( + "Batch-correction recommendations must cite exact batch " + f"estimability evidence: {missing_safety_evidence}" + ) + not_computed = [ + item.coefficient for item in matched_safety if item.status == "notComputed" + ] + if not_computed: + raise ModelRetry( + "Batch estimability could not be computed; use action='needsInput' " + f"for: {sorted(not_computed)}" + ) + unsafe_coefficients = [ + item.coefficient for item in matched_safety if item.status == "unsafe" + ] + if plan.action == "evaluateHarmony" and unsafe_coefficients: + raise ModelRetry( + "Batch correction is unsafe because the biological coefficient is " + "not estimable after the exact proposed batch columns; use " + f"action='unsafe' for: {sorted(unsafe_coefficients)}" + ) + if plan.action == "unsafe" and not unsafe_coefficients: + raise ModelRetry( + "The exact proposed batch columns were estimable for every " + "coefficient; use action='evaluateHarmony' or 'skip'" + ) + + cited_ids = [ + *decision.evidenceIds, + *plan.evidenceIds, + ] + unknown_evidence = sorted(set(cited_ids) - deps.evidenceIds) + if unknown_evidence: + raise ModelRetry(f"Unknown evidence IDs: {unknown_evidence}") + if plan.action in {"evaluateHarmony", "skip", "unsafe"} and not plan.evidenceIds: + raise ModelRetry("Batch-correction recommendations require evidence IDs") + current_metric_evidence = set(deps.currentRepresentation.evidenceIds) + stale_metric_evidence = sorted( + evidence_id + for evidence_id in cited_ids + if evidence_id.startswith("metric:") + and evidence_id not in current_metric_evidence + ) + if stale_metric_evidence: + raise ModelRetry( + "Metric evidence must come from the returned exact representation: " + f"{stale_metric_evidence}" + ) + + +def validate_experimental_context( + decision: ExperimentalContextDecision, + deps: ExperimentalContextDependencies, +) -> ExperimentalContextDecision: + """Recompute and validate every model-authored design choice.""" + if decision.cellQc != CellQcPlan.get_blank(): + raise ModelRetry( + "Experimental Context must leave cellQc blank; the audited filtering " + "checkpoint selects from qcProfiles" + ) + narrative_fields = { + "rationale": decision.rationale, + "batchCorrection.rationale": decision.batchCorrection.rationale, + **{ + f"needsInput[{index}]": question + for index, question in enumerate(decision.needsInput) + }, + } + serialized_field_markers = ( + '"evidenceIds":', + '"needsInput":', + '"runInfo":', + '"batchCorrection":', + '"cellQc":', + ) + invalid_narratives = [ + name + for name, value in narrative_fields.items() + if any( + marker in value.replace('\\"', '"') for marker in serialized_field_markers + ) + ] + if invalid_narratives: + raise ModelRetry( + "Narrative fields must contain plain prose without serialized sibling " + f"fields: {invalid_narratives}" + ) + for required in ("inspect_cell_covariates", "analyze_experimental_design"): + if required not in deps.toolCalls: + raise ModelRetry(f"Call {required} before returning a decision") + try: + canonical_design_choices(deps, decision) + except ValueError as exc: + raise ModelRetry(str(exc)) from exc + directions = dict(deps.directions) + column_domains = dict(decision.columnDomains) + column_domains.update(dict(directions.get("columnDomains") or {})) + directions["columnDomains"] = column_domains + directions["coefficientsOfInterest"] = list( + dict.fromkeys( + [ + *decision.coefficientsOfInterest, + *(directions.get("coefficientsOfInterest") or []), + ] + ) + ) + units_of_inference = { + name: unit.model_dump(exclude_none=True) + for name, unit in decision.unitsOfInference.items() + } + units_of_inference.update(dict(directions.get("unitsOfInference") or {})) + directions["unitsOfInference"] = units_of_inference + + previous_inputs = deps.characterizationInputs + characterization = characterize_context(deps, directions) + if characterization.status == "failed": + raise ModelRetry("; ".join(characterization.notes)) + characterization.comparisons = list(deps.comparisons) + characterization.captureProvenance = deps.captureProposal + deps.characterization = characterization + deps.evidenceIds.update(characterization_evidence(characterization)) + contrast_plans = contrast_plans_from_characterization(characterization) + deps.contrastPlans = {plan.coefficient: plan for plan in contrast_plans} + deps.evidenceIds.update(plan.evidenceId for plan in contrast_plans) + + if not deps.qcProfiles or previous_inputs != deps.characterizationInputs: + _offered_qc_profiles(deps, characterization) + deps.evidenceIds.update(profile.evidenceId for profile in deps.qcProfiles.values()) + deps.evidenceIds.update(source.sourceId for source in deps.qcMetricSources) + deps.evidenceIds.update(item.evidenceId for item in deps.qcSourceConcordance) + + requested_coefficients = set(directions["coefficientsOfInterest"]) + characterized_coefficients = { + record.get("name") for record in characterization.coefficients + } + missing_coefficients = sorted( + name + for name in requested_coefficients + if name not in characterized_coefficients + ) + if missing_coefficients: + raise ModelRetry( + "Coefficients of interest must be classified as biological: " + f"{missing_coefficients}" + ) + + coefficient_records: dict[str, dict[str, Any]] = {} + for record in characterization.coefficients: + name = record.get("name") + if isinstance(name, str): + coefficient_records[name] = record + records: dict[str, dict[str, Any]] = {} + for record in characterization.columns: + name = record.get("name") + if isinstance(name, str): + records[name] = record + _validate_batch_correction_plan( + decision, + deps, + characterization, + requested_coefficients, + units_of_inference, + records, + coefficient_records, + ) + canonical_domains = { + name: records[name]["domain"] + for name in column_domains + if name in records + and records[name].get("domain") + in { + "biological", + "technical", + "design", + "ignore", + "unknown", + } + } + canonical_units = { + coefficient: InferenceUnit( + observationUnit=coefficient_records[coefficient].get("observationUnit"), + independentUnit=coefficient_records[coefficient].get("independentUnit"), + ) + for coefficient in directions["coefficientsOfInterest"] + if coefficient in coefficient_records + } + try: + design_choices = canonical_design_choices(deps, decision) + except ValueError as exc: + raise ModelRetry(str(exc)) from exc + validated = decision.model_copy( + update={ + "columnDomains": canonical_domains, + "coefficientsOfInterest": list(directions["coefficientsOfInterest"]), + "unitsOfInference": canonical_units, + "cellQc": CellQcPlan.get_blank(), + **design_choices, + } + ) + if deps.studyObjective: + requirements, coverage = objective_evidence( + study_context=deps.studyContext, + study_objective=deps.studyObjective, + experimental_result=SimpleNamespace( + decision=validated, + characterization=characterization, + batchSafety=list(deps.batchSafety.values()), + ), + ) + unanswered = unmet_objective_requirements(requirements, coverage) + if unanswered: + pending_comparisons = [ + item.question + for item in requirements + if item.requirementId.startswith("requestedDesign:") + and any( + row.requirementId == item.requirementId + and row.status == "unsupported" + for row in coverage + ) + ] + if pending_comparisons and deps.designRounds < 2: + raise ModelRetry( + "Explicit study comparisons have not been measured: " + + "; ".join(pending_comparisons) + + ". Use the remaining analyze_experimental_design round for " + "these questions before finalizing. Match their requested purpose: " + "descriptive crossing/replication is designCoverage, not an " + "association or effect estimate. Preserve unsupported methods " + "as limitations; do not substitute marginal comparisons." + ) + validated = validated.model_copy( + update={ + "needsInput": list( + dict.fromkeys([*validated.needsInput, *unanswered]) + ), + } + ) + logger.debug( + "Experimental Context decision validated: " + f"domains={len(validated.columnDomains)}, " + f"coefficients={len(validated.coefficientsOfInterest)}, " + f"qcProfiles={len(deps.qcProfiles)}, " + f"batchCorrection={validated.batchCorrection.action}, " + f"needsInput={len(validated.needsInput)}" + ) + return validated + + +def failed_experimental_context_result( + deps: ExperimentalContextDependencies, + *, + error: Exception, + model_name: str, +) -> ExperimentalContextResult: + """Keep a failed model run from selecting an unsupported scientific default.""" + characterization = deps.characterization or CovariateCharacterization( + status="failed", + notes=["Deterministic covariate characterization is unavailable."], + ) + model_detail = str(error).replace("\n", " ").strip()[:500] + failure_notes = [] + cause = error.__cause__ + if isinstance(cause, ValidationError): + details = [ + f"{'.'.join(str(part) for part in item['loc'])}: {item['msg']}" + for item in cause.errors( + include_url=False, include_input=False, include_context=False + )[:8] + ] + failure_notes.append("Invalid tool arguments: " + "; ".join(details)) + elif isinstance(cause, ModelRetry): + failure_notes.append("Rejected tool arguments: " + str(cause)[:1000]) + return ExperimentalContextResult( + status="failed", + decision=ExperimentalContextDecision( + rationale="No validated experimental-context decision was available.", + evidenceIds=sorted(deps.evidenceIds), + ), + characterization=characterization, + cellSelection=artifact_reference(deps.cellSelection), + cellQc=CellQcPlan.get_blank(), + qcProfiles=list(deps.qcProfiles.values()), + qcMetricSources=deps.qcMetricSources, + qcSourceConcordance=deps.qcSourceConcordance, + contrastPlans=list(deps.contrastPlans.values()), + qualityMetricArtifacts=deps.qualityMetricArtifacts, + htoIdentityColumns=deps.htoIdentityColumns, + htoIdentityArtifacts=deps.htoIdentityArtifacts, + batchSafety=list(deps.batchSafety.values()), + currentRepresentation=deps.currentRepresentation, + notes=[ + "The model did not produce a validated experimental-context decision.", + f"Model failure: {model_detail}", + *failure_notes, + ], + runInfo=getattr(error, "agent_run_info", None) + or AgentRunInfo( + agentName="experimental_context_failed", + modelName=model_name, + ), + ) diff --git a/scarf/agent/ingest/__init__.py b/scarf/agent/ingest/__init__.py index 31c90b48..34a52f86 100644 --- a/scarf/agent/ingest/__init__.py +++ b/scarf/agent/ingest/__init__.py @@ -9,6 +9,12 @@ from .detect import detect_format from .h5ad import ingest_h5ad from .loom import ingest_loom +from .manifest import ( + DatasetManifest, + DatasetManifestDecision, + inspect_h5ad_manifest, + is_author_label_column, +) from .mtx import ingest_mtx from .result import IngestResult, needs_input from .seurat import ingest_seurat @@ -16,8 +22,12 @@ __all__ = [ "IngestResult", + "DatasetManifest", + "DatasetManifestDecision", "detect_format", "ingest", + "inspect_h5ad_manifest", + "is_author_label_column", ] diff --git a/scarf/agent/ingest/common.py b/scarf/agent/ingest/common.py index 2df7f0b4..2d8f8078 100644 --- a/scarf/agent/ingest/common.py +++ b/scarf/agent/ingest/common.py @@ -7,7 +7,7 @@ from ...storage.profiles import is_local_zarr_path from ...storage.stores import zarr_location_has_content -from ..decide import DecisionValidationError, decide +from ..decisions.selection import DecisionValidationError, decide from ..types import Decision, EvidenceItem from .result import ( IngestResult, @@ -250,36 +250,6 @@ def finish( ) accepted_actions = [*convert_actions, action] resolved_action_labels = list(action_labels) - workflow_run = None - if summary_mode != "r": - from ..persistence import create_agent_workflow - - try: - workflow_run = create_agent_workflow(zarr_path) - except AGENT_PERSISTENCE_ERRORS as exc: - return IngestResult( - status="failed", - format=format_name, - zarrPath=zarr_path, - assayNames=assay_names, - summary=summary, - decision=decision, - actions=resolved_action_labels, - acceptedActions=accepted_actions, - notes=[ - *notes, - failure_note("create agent workflow", exc), - f"The converted Scarf store remains available at {zarr_path}", - ], - ) - resolved_action_labels.append("create_agent_workflow") - accepted_actions.append( - { - "op": "createAgentWorkflow", - "zarrPath": zarr_path, - "workflowRunId": workflow_run.workflowRunId, - } - ) return done( format_name=format_name, zarr_path=zarr_path, @@ -288,7 +258,6 @@ def finish( accepted_actions=accepted_actions, action_labels=resolved_action_labels, notes=notes, - workflow_run=workflow_run, decision=decision, ) diff --git a/scarf/agent/ingest/manifest.py b/scarf/agent/ingest/manifest.py new file mode 100644 index 00000000..7c6f84e5 --- /dev/null +++ b/scarf/agent/ingest/manifest.py @@ -0,0 +1,1036 @@ +"""Read-only H5AD inventory for decision-driven RNA workflows.""" + +from collections.abc import Iterator +from hashlib import sha256 +from pathlib import Path +from typing import Any, Literal, cast + +import h5py +import numpy as np + +from ...readers._h5ad_inspect import ( + _as_text, + _column_names, + _matrix_candidates, + _MatrixCandidate, + _node_length, + _select_matrix, + inspect_h5ad, +) +from .._deps import AGENT_INSTALL_HINT +from ..types import AgentDataModel + +try: + from pydantic import Field +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + + +type AuthorLabelPolicy = Literal["holdout", "preservation"] +type ManifestStatus = Literal["supported", "needsInput", "abstained"] +type MatrixCountSemantics = Literal["integerLikeCounts", "nonIntegerValues"] + +_AUTHOR_LABEL_EXACT = frozenset( + { + "annotation", + "annotations", + "author_annotation", + "author_cell_type", + "author_cell_type_ontology_term_id", + "cell_type", + "cell_type_ontology_term_id", + "cluster", + "clusters", + "clustering", + "leiden", + "louvain", + } +) +_AUTHOR_LABEL_FRAGMENTS = ( + "annotation", + "cell_label", + "cell_type", + "celltype", + "cluster", + "leiden", + "louvain", +) +_OBS_SCHEMA_FIELDS = frozenset( + { + "assay", + "assay_ontology_term_id", + "development_stage", + "development_stage_ontology_term_id", + "disease", + "disease_ontology_term_id", + "donor_id", + "is_primary_data", + "organism", + "organism_ontology_term_id", + "self_reported_ethnicity", + "self_reported_ethnicity_ontology_term_id", + "sex", + "suspension_type", + "tissue", + "tissue_ontology_term_id", + } +) +_VAR_SCHEMA_FIELDS = frozenset( + { + "feature_biotype", + "feature_is_filtered", + "feature_length", + "feature_name", + "feature_reference", + "feature_type", + "feature_types", + "gene_id", + "gene_ids", + "gene_name", + "gene_symbol", + } +) +_ANALYSIS_KEY_FRAGMENTS = ( + "cluster", + "leiden", + "louvain", + "marker", + "neighbors", + "pca", + "rank_gene", + "tsne", + "umap", +) +_DEFAULT_CHUNK_VALUES = 8_192 +_DEFAULT_MAX_COLUMNS = 256 +_DEFAULT_MAX_DOMAIN_VALUES = 16 +_DEFAULT_MAX_INVENTORY_ITEMS = 256 +_DIGEST_CHUNK_BYTES = 8 * 1024 * 1024 + + +class MatrixCandidateManifest(AgentDataModel): + """One dimension-compatible matrix observed in the source file.""" + + key: str + encoding: Literal["csr", "csc", "dense"] + nCells: int + nFeatures: int + integerLike: bool + countSemantics: MatrixCountSemantics + dimensionCompatible: bool + featureMetadataKey: str | None = None + selected: bool = False + + +class MetadataColumnSummary(AgentDataModel): + """Bounded summary of one H5AD dataframe column.""" + + name: str + dtype: str + storageKind: Literal["categorical", "boolean", "numeric", "string", "other"] + valueCount: int + missingCount: int + missingFraction: float + domainSize: int | None = None + domainValues: list[str] = Field(default_factory=list) + domainTruncated: bool = False + valueCounts: dict[str, int] = Field(default_factory=dict) + minimum: float | None = None + maximum: float | None = None + schemaField: bool = False + + +class MetadataTableSummary(AgentDataModel): + """Column inventory for an H5AD dataframe-like node.""" + + key: str + rowCount: int + columns: list[MetadataColumnSummary] = Field(default_factory=list) + schemaFields: list[str] = Field(default_factory=list) + identifierColumns: list[str] = Field(default_factory=list) + omittedColumnCount: int = 0 + heldOutAuthorColumnCount: int = 0 + + +class H5adInventory(AgentDataModel): + """Bounded top-level inventory of reusable and author-produced assets.""" + + layers: list[str] = Field(default_factory=list) + obsm: list[str] = Field(default_factory=list) + uns: list[str] = Field(default_factory=list) + priorEmbeddings: list[str] = Field(default_factory=list) + authorAnalysisArtifacts: list[str] = Field(default_factory=list) + omittedLayerCount: int = 0 + omittedObsmCount: int = 0 + omittedUnsCount: int = 0 + heldOutUnsItemCount: int = 0 + + +class PriorFilteringFacts(AgentDataModel): + """Facts and limitations imposed by the published H5AD cell universe.""" + + cellXGeneSchemaDetected: bool + cellXGeneSchemaVersion: str | None = None + cellXGeneSchemaReference: str | None = None + cellSetStatus: Literal["publishedCellsOnly", "unknown"] + rawCountsAvailable: bool + originalDropletsAvailable: bool = False + originalCellCallingSupported: bool = False + ambientCorrectionSupported: bool = False + featureFilteringFlagAvailable: bool = False + filteredFeatureCount: int | None = None + primaryDataFlagAvailable: bool = False + nonPrimaryCellCount: int | None = None + limitations: list[str] = Field(default_factory=list) + + +class DatasetManifestDecision(AgentDataModel): + """Eligibility result for the count-dependent RNA decision workflow.""" + + status: ManifestStatus + reasonCode: Literal[ + "rawRnaCountsAvailable", + "ambiguousCountMatrices", + "normalizedOnly", + "rnaModalityUnavailable", + ] + summary: str + selectedMatrixKey: str | None = None + options: list[str] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + + +class DatasetManifest(AgentDataModel): + """Read-only evidence collected before materializing an H5AD source.""" + + formatVersion: Literal[1] = 1 + sourcePath: str + sourceUri: str | None = None + sourceSha256: str + sourceSizeBytes: int + authorLabelPolicy: AuthorLabelPolicy + labelBenchmarkEligible: bool + nCells: int + nFeatures: int + selectedFeatureMetadataKey: str | None = None + matrixCandidates: list[MatrixCandidateManifest] = Field(default_factory=list) + obs: MetadataTableSummary + var: MetadataTableSummary + rawVar: MetadataTableSummary | None = None + assayMetadata: MetadataColumnSummary | None = None + suspensionMetadata: MetadataColumnSummary | None = None + organismMetadata: MetadataColumnSummary | None = None + declaredBatchColumns: list[str] = Field( + default_factory=list, + exclude_if=lambda value: not value, + ) + inventory: H5adInventory + priorFiltering: PriorFilteringFacts + decision: DatasetManifestDecision + + +def _source_sha256(path: Path) -> str: + digest = sha256() + with path.open("rb") as source: + while chunk := source.read(_DIGEST_CHUNK_BYTES): + digest.update(chunk) + return digest.hexdigest() + + +def _is_author_label(name: str) -> bool: + normalized = name.strip().lower() + return normalized in _AUTHOR_LABEL_EXACT or any( + fragment in normalized for fragment in _AUTHOR_LABEL_FRAGMENTS + ) + + +def is_author_label_column(name: str) -> bool: + """Return whether a metadata column is quarantined as an author label.""" + return _is_author_label(name) + + +def _is_missing(values: np.ndarray) -> np.ndarray: + if values.dtype.kind in {"f", "c"}: + return cast(np.ndarray, ~np.isfinite(values)) + if values.dtype.kind in {"S", "U"}: + return np.asarray([not _as_text(value) for value in values], dtype=bool) + if values.dtype.kind != "O": + return np.zeros(values.shape, dtype=bool) + return np.asarray( + [ + value is None + or (isinstance(value, float | np.floating) and not np.isfinite(value)) + or ( + isinstance(value, str | bytes | np.str_ | np.bytes_) + and not _as_text(value) + ) + for value in values + ], + dtype=bool, + ) + + +def _value_text(value: Any) -> str: + if isinstance(value, bool | np.bool_): + return "true" if bool(value) else "false" + if isinstance(value, bytes | np.bytes_): + return value.decode("utf-8", errors="replace") + if isinstance(value, np.integer): + return str(int(value)) + if isinstance(value, np.floating): + return repr(float(value)) + return str(value) + + +def _storage_kind( + dtype: np.dtype[Any], *, categorical: bool = False +) -> Literal["categorical", "boolean", "numeric", "string", "other"]: + if categorical: + return "categorical" + if np.issubdtype(dtype, np.bool_): + return "boolean" + if np.issubdtype(dtype, np.number): + return "numeric" + if dtype.kind in {"O", "S", "U"}: + return "string" + return "other" + + +def _dataset_value_chunks( + dataset: h5py.Dataset, + *, + field: str | None, + row_count: int, + chunk_values: int, +) -> Iterator[np.ndarray]: + for start in range(0, row_count, chunk_values): + stop = min(row_count, start + chunk_values) + raw = np.asarray(dataset[start:stop]) + yield np.asarray(raw[field] if field is not None else raw).reshape(-1) + + +def _column_source( + table: h5py.Group | h5py.Dataset, + name: str, + *, + row_count: int, + chunk_values: int, +) -> tuple[Iterator[np.ndarray], np.dtype[Any], h5py.Dataset | None]: + if isinstance(table, h5py.Dataset): + if table.dtype.names is None or name not in table.dtype.names: + raise ValueError(f"Column {name!r} is unavailable in {table.name}") + dtype = table.dtype.fields[name][0] + return ( + _dataset_value_chunks( + table, + field=name, + row_count=row_count, + chunk_values=chunk_values, + ), + dtype, + None, + ) + + node = table.get(name) + if isinstance(node, h5py.Dataset): + categories: h5py.Dataset | None = None + for category_group_name in ("__categories", "categories"): + category_group = table.get(category_group_name) + if isinstance(category_group, h5py.Group): + category_node = category_group.get(name) + if isinstance(category_node, h5py.Dataset): + categories = category_node + break + return ( + _dataset_value_chunks( + node, + field=None, + row_count=row_count, + chunk_values=chunk_values, + ), + node.dtype, + categories, + ) + if isinstance(node, h5py.Group): + codes = node.get("codes") + categories = node.get("categories") + if isinstance(codes, h5py.Dataset) and isinstance(categories, h5py.Dataset): + return ( + _dataset_value_chunks( + codes, + field=None, + row_count=row_count, + chunk_values=chunk_values, + ), + codes.dtype, + categories, + ) + raise ValueError(f"Unsupported H5AD column encoding for {table.name}/{name}") + + +def _categorical_summary( + *, + name: str, + chunks: Iterator[np.ndarray], + dtype: np.dtype[Any], + categories: h5py.Dataset, + row_count: int, + max_domain_values: int, + schema_field: bool, +) -> MetadataColumnSummary: + category_count = int(categories.shape[0]) + shown_count = min(category_count, max_domain_values) + shown = [_value_text(value) for value in np.asarray(categories[:shown_count])] + counts = np.zeros(shown_count, dtype=np.int64) + missing_count = 0 + for values in chunks: + codes = np.asarray(values, dtype=np.int64) + valid = (codes >= 0) & (codes < category_count) + missing_count += int((~valid).sum()) + shown_codes = codes[valid & (codes < shown_count)] + if shown_codes.size: + counts += np.bincount(shown_codes, minlength=shown_count) + value_counts = ( + {value: int(count) for value, count in zip(shown, counts, strict=True)} + if category_count <= max_domain_values + else {} + ) + return MetadataColumnSummary( + name=name, + dtype=str(dtype), + storageKind="categorical", + valueCount=row_count, + missingCount=missing_count, + missingFraction=missing_count / row_count if row_count else 0.0, + domainSize=category_count, + domainValues=shown, + domainTruncated=category_count > max_domain_values, + valueCounts=value_counts, + schemaField=schema_field, + ) + + +def _plain_summary( + *, + name: str, + chunks: Iterator[np.ndarray], + dtype: np.dtype[Any], + row_count: int, + max_domain_values: int, + schema_field: bool, +) -> MetadataColumnSummary: + missing_count = 0 + observed_counts: dict[str, int] = {} + domain_truncated = False + minimum: float | None = None + maximum: float | None = None + numeric = np.issubdtype(dtype, np.number) and not np.issubdtype(dtype, np.bool_) + for values in chunks: + missing = _is_missing(values) + missing_count += int(missing.sum()) + observed = values[~missing] + if numeric and observed.size: + observed_float = np.asarray(observed, dtype=np.float64) + block_minimum = float(observed_float.min()) + block_maximum = float(observed_float.max()) + minimum = block_minimum if minimum is None else min(minimum, block_minimum) + maximum = block_maximum if maximum is None else max(maximum, block_maximum) + if domain_truncated: + continue + for value in observed: + normalized = _value_text(value) + observed_counts[normalized] = observed_counts.get(normalized, 0) + 1 + if len(observed_counts) > max_domain_values: + observed_counts.clear() + domain_truncated = True + break + domain_values = sorted(observed_counts) + return MetadataColumnSummary( + name=name, + dtype=str(dtype), + storageKind=_storage_kind(dtype), + valueCount=row_count, + missingCount=missing_count, + missingFraction=missing_count / row_count if row_count else 0.0, + domainSize=None if domain_truncated else len(domain_values), + domainValues=domain_values, + domainTruncated=domain_truncated, + valueCounts={} if domain_truncated else observed_counts, + minimum=minimum, + maximum=maximum, + schemaField=schema_field, + ) + + +def _summarize_column( + table: h5py.Group | h5py.Dataset, + name: str, + *, + row_count: int, + chunk_values: int, + max_domain_values: int, + schema_fields: frozenset[str], +) -> MetadataColumnSummary: + chunks, dtype, categories = _column_source( + table, + name, + row_count=row_count, + chunk_values=chunk_values, + ) + if categories is not None: + return _categorical_summary( + name=name, + chunks=chunks, + dtype=dtype, + categories=categories, + row_count=row_count, + max_domain_values=max_domain_values, + schema_field=name.lower() in schema_fields, + ) + return _plain_summary( + name=name, + chunks=chunks, + dtype=dtype, + row_count=row_count, + max_domain_values=max_domain_values, + schema_field=name.lower() in schema_fields, + ) + + +def _index_columns(table: h5py.Group | h5py.Dataset) -> set[str]: + if isinstance(table, h5py.Dataset): + return {"_index", "index"} & set(table.dtype.names or ()) + names = {"_index", "index"} & set(table.keys()) + index_attribute = table.attrs.get("_index") + if index_attribute is not None: + names.add(_as_text(index_attribute)) + return names + + +def _summarize_table( + h5: h5py.File, + key: str, + *, + author_label_policy: AuthorLabelPolicy, + schema_fields: frozenset[str], + chunk_values: int, + max_columns: int, + max_domain_values: int, +) -> MetadataTableSummary: + node = h5.get(key) + if not isinstance(node, h5py.Group | h5py.Dataset): + return MetadataTableSummary(key=key, rowCount=0) + row_count = _node_length(node) or 0 + names = sorted(_column_names(node)) + identifiers = _index_columns(node) + held_out = { + name + for name in names + if author_label_policy == "holdout" and _is_author_label(name) + } + inspectable = [ + name for name in names if name not in identifiers and name not in held_out + ] + inspectable.sort(key=lambda name: (name.lower() not in schema_fields, name)) + selected = inspectable[:max_columns] + columns = [ + _summarize_column( + node, + name, + row_count=row_count, + chunk_values=chunk_values, + max_domain_values=max_domain_values, + schema_fields=schema_fields, + ) + for name in selected + ] + return MetadataTableSummary( + key=key, + rowCount=row_count, + columns=columns, + schemaFields=sorted(summary.name for summary in columns if summary.schemaField), + identifierColumns=sorted(identifiers), + omittedColumnCount=max(0, len(inspectable) - len(selected)), + heldOutAuthorColumnCount=len(held_out), + ) + + +def _inventory_keys( + h5: h5py.File, + key: str, + *, + max_items: int, + hide_author_labels: bool = False, +) -> tuple[list[str], int, int]: + node = h5.get(key) + if not isinstance(node, h5py.Group): + return [], 0, 0 + names = sorted(str(name) for name in node.keys()) + held_out = [name for name in names if hide_author_labels and _is_author_label(name)] + visible = [name for name in names if name not in held_out] + return visible[:max_items], max(0, len(visible) - max_items), len(held_out) + + +def _read_text_scalar( + h5: h5py.File, + paths: tuple[str, ...], + *, + max_length: int = 500, +) -> str | None: + for path in paths: + node = h5.get(path) + if not isinstance(node, h5py.Dataset) or node.shape not in {(), (1,)}: + continue + value = node[()] if node.shape == () else node[0] + return _as_text(value)[:max_length] + return None + + +def _read_text_vector( + h5: h5py.File, + paths: tuple[str, ...], + *, + max_items: int = 32, + max_length: int = 256, +) -> list[str]: + for path in paths: + node = h5.get(path) + if not isinstance(node, h5py.Dataset): + continue + if node.shape == (): + values = [node[()]] + elif len(node.shape) == 1: + if node.shape[0] > max_items: + raise ValueError(f"{path} contains too many values") + values = list(np.asarray(node[:]).reshape(-1)) + else: + continue + resolved = list( + dict.fromkeys( + text + for value in values + if (text := _as_text(value).strip()[:max_length]) + ) + ) + return resolved + return [] + + +def _column_by_name( + *tables: MetadataTableSummary | None, + names: tuple[str, ...], +) -> MetadataColumnSummary | None: + wanted = {name.lower() for name in names} + for table in tables: + if table is None: + continue + for column in table.columns: + if column.name.lower() in wanted: + return column + return None + + +def _count_value( + column: MetadataColumnSummary | None, + value: str, +) -> int | None: + if column is None or not column.valueCounts: + return None + return column.valueCounts.get(value, 0) + + +def _matrix_manifests( + h5: h5py.File, + *, + selected_key: str | None, +) -> tuple[list[MatrixCandidateManifest], list[_MatrixCandidate]]: + manifests: list[MatrixCandidateManifest] = [] + compatible: list[_MatrixCandidate] = [] + for candidate in _matrix_candidates(h5): + if candidate.encoding not in {"csr", "csc", "dense"}: + raise ValueError( + f"Unsupported matrix encoding for {candidate.key}: {candidate.encoding}" + ) + encoding = cast(Literal["csr", "csc", "dense"], candidate.encoding) + feature_key: str | None = None + dimension_compatible = False + try: + _, feature_key = _select_matrix(h5, [candidate]) + dimension_compatible = True + compatible.append(candidate) + except ValueError: + pass + manifests.append( + MatrixCandidateManifest( + key=candidate.key, + encoding=encoding, + nCells=candidate.shape[0], + nFeatures=candidate.shape[1], + integerLike=candidate.integerLike, + countSemantics=( + "integerLikeCounts" if candidate.integerLike else "nonIntegerValues" + ), + dimensionCompatible=dimension_compatible, + featureMetadataKey=feature_key, + selected=candidate.key == selected_key, + ) + ) + return manifests, compatible + + +def _matrix_decision( + compatible: list[_MatrixCandidate], + *, + matrix_key: str | None, +) -> tuple[DatasetManifestDecision, str | None]: + by_key = {candidate.key: candidate for candidate in compatible} + if matrix_key is not None: + candidate = by_key.get(matrix_key) + if candidate is None: + available = ", ".join(sorted(by_key)) + raise ValueError( + f"matrix_key {matrix_key!r} is not dimension-compatible. " + f"Available: {available}" + ) + if not candidate.integerLike: + return ( + DatasetManifestDecision( + status="abstained", + reasonCode="normalizedOnly", + summary=( + f"Selected matrix {matrix_key} is not integer-like and " + "cannot authorize count-dependent RNA analysis" + ), + options=[matrix_key], + evidenceIds=[f"matrix:{matrix_key}:nonIntegerValues"], + ), + None, + ) + return ( + DatasetManifestDecision( + status="supported", + reasonCode="rawRnaCountsAvailable", + summary=f"Selected integer-like count candidate {matrix_key}", + selectedMatrixKey=matrix_key, + options=[matrix_key], + evidenceIds=[f"matrix:{matrix_key}:integerLikeCounts"], + ), + matrix_key, + ) + + integer_keys = sorted( + candidate.key for candidate in compatible if candidate.integerLike + ) + if not integer_keys: + non_integer_keys = sorted(candidate.key for candidate in compatible) + return ( + DatasetManifestDecision( + status="abstained", + reasonCode="normalizedOnly", + summary=( + "No dimension-compatible integer-like count matrix is " + "available in this H5AD" + ), + options=non_integer_keys, + evidenceIds=[ + f"matrix:{key}:nonIntegerValues" for key in non_integer_keys + ], + ), + None, + ) + if "raw/X" in integer_keys: + conflicting_keys = [key for key in integer_keys if key not in {"raw/X", "X"}] + if not conflicting_keys: + return ( + DatasetManifestDecision( + status="supported", + reasonCode="rawRnaCountsAvailable", + summary=( + "Selected authoritative integer-like raw/X count matrix " + "over the visualization matrix X" + ), + selectedMatrixKey="raw/X", + options=["raw/X"], + evidenceIds=["matrix:raw/X:integerLikeCounts"], + ), + "raw/X", + ) + if len(integer_keys) > 1: + return ( + DatasetManifestDecision( + status="needsInput", + reasonCode="ambiguousCountMatrices", + summary=( + "Multiple dimension-compatible integer-like matrices are " + "available; select the intended raw count matrix" + ), + options=integer_keys, + evidenceIds=[f"matrix:{key}:integerLikeCounts" for key in integer_keys], + ), + None, + ) + selected_key = integer_keys[0] + return ( + DatasetManifestDecision( + status="supported", + reasonCode="rawRnaCountsAvailable", + summary=f"One integer-like count candidate is available: {selected_key}", + selectedMatrixKey=selected_key, + options=[selected_key], + evidenceIds=[f"matrix:{selected_key}:integerLikeCounts"], + ), + selected_key, + ) + + +def _inventory( + h5: h5py.File, + *, + author_label_policy: AuthorLabelPolicy, + max_items: int, +) -> H5adInventory: + layers, omitted_layers, _ = _inventory_keys(h5, "layers", max_items=max_items) + obsm, omitted_obsm, _ = _inventory_keys(h5, "obsm", max_items=max_items) + uns, omitted_uns, held_out_uns = _inventory_keys( + h5, + "uns", + max_items=max_items, + hide_author_labels=author_label_policy == "holdout", + ) + prior_embeddings = [ + f"obsm/{name}" + for name in obsm + if name.lower().startswith("x_") + or any(fragment in name.lower() for fragment in ("pca", "tsne", "umap")) + ] + analysis_artifacts = [ + f"{group}/{name}" + for group, names in (("obsm", obsm), ("uns", uns)) + for name in names + if any(fragment in name.lower() for fragment in _ANALYSIS_KEY_FRAGMENTS) + ] + if held_out_uns: + analysis_artifacts.append(f"uns/heldOutAuthorItems:{held_out_uns}") + return H5adInventory( + layers=layers, + obsm=obsm, + uns=uns, + priorEmbeddings=prior_embeddings, + authorAnalysisArtifacts=analysis_artifacts, + omittedLayerCount=omitted_layers, + omittedObsmCount=omitted_obsm, + omittedUnsCount=omitted_uns, + heldOutUnsItemCount=held_out_uns, + ) + + +def inspect_h5ad_manifest( + path: str | Path, + *, + source_uri: str | None = None, + author_label_policy: AuthorLabelPolicy = "holdout", + matrix_key: str | None = None, + chunk_values: int = _DEFAULT_CHUNK_VALUES, + max_columns: int = _DEFAULT_MAX_COLUMNS, + max_domain_values: int = _DEFAULT_MAX_DOMAIN_VALUES, + max_inventory_items: int = _DEFAULT_MAX_INVENTORY_ITEMS, +) -> DatasetManifest: + """Inspect an H5AD without converting or modifying it. + + Matrix values are sampled only through :func:`inspect_h5ad`; metadata is + summarized in bounded chunks. Author cell-type and clustering columns are + not read when ``author_label_policy`` is ``"holdout"``. + """ + source = Path(path) + if not source.is_file(): + raise ValueError(f"H5AD source is not a file: {source}") + if author_label_policy not in {"holdout", "preservation"}: + raise ValueError( + "author_label_policy must be either 'holdout' or 'preservation'" + ) + if chunk_values < 1: + raise ValueError("chunk_values must be at least 1") + if max_columns < 1: + raise ValueError("max_columns must be at least 1") + if max_domain_values < 1: + raise ValueError("max_domain_values must be at least 1") + if max_inventory_items < 1: + raise ValueError("max_inventory_items must be at least 1") + + with h5py.File(source, mode="r") as h5: + matrix_candidates, compatible = _matrix_manifests(h5, selected_key=None) + if not compatible: + raise ValueError("No matrix candidate matches the obs and var dimensions") + decision, selected_key = _matrix_decision(compatible, matrix_key=matrix_key) + if selected_key is not None: + inspection = inspect_h5ad(str(source), matrix_key=selected_key) + else: + inspection = inspect_h5ad(str(source)) + matrix_candidates, _ = _matrix_manifests(h5, selected_key=selected_key) + + obs = _summarize_table( + h5, + "obs", + author_label_policy=author_label_policy, + schema_fields=_OBS_SCHEMA_FIELDS, + chunk_values=chunk_values, + max_columns=max_columns, + max_domain_values=max_domain_values, + ) + var = _summarize_table( + h5, + "var", + author_label_policy="preservation", + schema_fields=_VAR_SCHEMA_FIELDS, + chunk_values=chunk_values, + max_columns=max_columns, + max_domain_values=max_domain_values, + ) + raw_var = ( + _summarize_table( + h5, + "raw/var", + author_label_policy="preservation", + schema_fields=_VAR_SCHEMA_FIELDS, + chunk_values=chunk_values, + max_columns=max_columns, + max_domain_values=max_domain_values, + ) + if isinstance(h5.get("raw/var"), h5py.Group | h5py.Dataset) + else None + ) + inventory = _inventory( + h5, + author_label_policy=author_label_policy, + max_items=max_inventory_items, + ) + schema_version = _read_text_scalar( + h5, + ("uns/schema_version", "uns/cellxgene_schema_version"), + ) + schema_reference = _read_text_scalar( + h5, + ("uns/schema_reference", "uns/cellxgene_schema_reference"), + ) + declared_batch_columns = _read_text_vector( + h5, + ("uns/batch_condition",), + ) + obs_node = h5.get("obs") + obs_columns = ( + set(_column_names(obs_node)) + if isinstance(obs_node, h5py.Group | h5py.Dataset) + else set() + ) + unknown_batch_columns = sorted(set(declared_batch_columns) - obs_columns) + if unknown_batch_columns: + raise ValueError( + "uns/batch_condition references unknown obs columns: " + f"{unknown_batch_columns}" + ) + + selected_table = raw_var if inspection.featureAttrsKey == "raw/var" else var + assay = _column_by_name( + obs, + names=("assay_ontology_term_id", "assay"), + ) + suspension = _column_by_name(obs, names=("suspension_type",)) + organism = _column_by_name( + obs, + names=("organism_ontology_term_id", "organism"), + ) + feature_filtered = _column_by_name( + selected_table, + var, + raw_var, + names=("feature_is_filtered",), + ) + is_primary = _column_by_name(obs, names=("is_primary_data",)) + cxg_detected = schema_version is not None or ( + assay is not None + and suspension is not None + and organism is not None + and feature_filtered is not None + ) + limitations = [ + "The H5AD observation table contains only the published cell universe; " + "discarded barcodes and empty droplets are unavailable.", + "Original cell calling and empty-droplet ambient correction cannot be " + "reconstructed from this file.", + ] + if decision.status == "abstained" and decision.reasonCode == "normalizedOnly": + limitations.append( + "Count-dependent quality control and feature selection require an " + "unambiguous integer-like source matrix." + ) + prior_filtering = PriorFilteringFacts( + cellXGeneSchemaDetected=cxg_detected, + cellXGeneSchemaVersion=schema_version, + cellXGeneSchemaReference=schema_reference, + cellSetStatus="publishedCellsOnly" if cxg_detected else "unknown", + rawCountsAvailable=any( + candidate.integerLike and candidate.dimensionCompatible + for candidate in matrix_candidates + ), + featureFilteringFlagAvailable=feature_filtered is not None, + filteredFeatureCount=_count_value(feature_filtered, "true"), + primaryDataFlagAvailable=is_primary is not None, + nonPrimaryCellCount=_count_value(is_primary, "false"), + limitations=limitations, + ) + + if decision.status == "supported" and inspection.suggestedAssays: + if "RNA" not in inspection.suggestedAssays: + decision = DatasetManifestDecision( + status="abstained", + reasonCode="rnaModalityUnavailable", + summary="The selected count matrix contains no RNA feature span", + options=sorted(inspection.suggestedAssays), + evidenceIds=[ + f"assay:{name}:{count}" + for name, count in sorted(inspection.suggestedAssays.items()) + ], + ) + selected_key = None + for candidate in matrix_candidates: + candidate.selected = False + + return DatasetManifest( + sourcePath=str(source.resolve()), + sourceUri=source_uri, + sourceSha256=_source_sha256(source), + sourceSizeBytes=source.stat().st_size, + authorLabelPolicy=author_label_policy, + labelBenchmarkEligible=author_label_policy == "holdout", + nCells=inspection.nCells, + nFeatures=inspection.nFeatures, + selectedFeatureMetadataKey=( + inspection.featureAttrsKey if selected_key is not None else None + ), + matrixCandidates=matrix_candidates, + obs=obs, + var=var, + rawVar=raw_var, + assayMetadata=assay, + suspensionMetadata=suspension, + organismMetadata=organism, + declaredBatchColumns=declared_batch_columns, + inventory=inventory, + priorFiltering=prior_filtering, + decision=decision, + ) + + +__all__ = [ + "AuthorLabelPolicy", + "DatasetManifest", + "DatasetManifestDecision", + "H5adInventory", + "ManifestStatus", + "MatrixCandidateManifest", + "MetadataColumnSummary", + "MetadataTableSummary", + "PriorFilteringFacts", + "inspect_h5ad_manifest", + "is_author_label_column", +] diff --git a/scarf/agent/ingest/result.py b/scarf/agent/ingest/result.py index d44f6db9..246719f8 100644 --- a/scarf/agent/ingest/result.py +++ b/scarf/agent/ingest/result.py @@ -3,8 +3,7 @@ from collections.abc import Sequence from typing import Any -from ..config._deps import AGENT_INSTALL_HINT -from ..persistence import AgentWorkflowRun +from .._deps import AGENT_INSTALL_HINT from ..types import AgentDataModel, Decision, NeedsInput, StageStatus try: @@ -17,7 +16,6 @@ class IngestResult(AgentDataModel): status: StageStatus format: str | None = None zarrPath: str | None = None - workflowRun: AgentWorkflowRun | None = None assayNames: list[str] = Field(default_factory=list) summary: dict[str, Any] | None = None decision: Decision | None = None @@ -30,16 +28,6 @@ class IngestResult(AgentDataModel): def get_blank(cls) -> "IngestResult": return cls(status="failed") - @classmethod - def get_example(cls) -> "IngestResult": - return cls( - status="done", - format="h5ad", - zarrPath="dataset.zarr", - workflowRun=AgentWorkflowRun.get_example(), - assayNames=["RNA"], - ) - def done( *, @@ -50,14 +38,12 @@ def done( accepted_actions: list[dict[str, Any]], action_labels: list[str], notes: list[str], - workflow_run: AgentWorkflowRun | None = None, decision: Decision | None = None, ) -> IngestResult: return IngestResult( status="done", format=format_name, zarrPath=zarr_path, - workflowRun=workflow_run, assayNames=assay_names, summary=summary, decision=decision, diff --git a/scarf/agent/orchestrator/__init__.py b/scarf/agent/orchestrator/__init__.py index f7ce55f1..e18236a0 100644 --- a/scarf/agent/orchestrator/__init__.py +++ b/scarf/agent/orchestrator/__init__.py @@ -1,37 +1,15 @@ -"""Public facade for automated Scarf agent orchestration.""" +"""Advanced RNA workflow configuration and explicit resume.""" from .main import AgentOrchestrator from .models import ( - AssayPreprocessingPlan, - AutomatedPreprocessingPlan, AutomatedWorkflowConfig, AutomatedWorkflowRequest, - AutomatedWorkflowResult, AutomatedWorkflowResumeRequest, - FinalAnalysisHandoff, - NativeAnalysisHandoff, - PreprocessedAssayHandoff, - WorkflowNeedsInput, - WorkflowQuestion, - WorkflowStageAttempt, - WorkflowStageLink, - artifact_model_to_ref, ) __all__ = [ "AgentOrchestrator", - "AssayPreprocessingPlan", - "AutomatedPreprocessingPlan", "AutomatedWorkflowConfig", "AutomatedWorkflowRequest", - "AutomatedWorkflowResult", "AutomatedWorkflowResumeRequest", - "FinalAnalysisHandoff", - "NativeAnalysisHandoff", - "PreprocessedAssayHandoff", - "WorkflowNeedsInput", - "WorkflowQuestion", - "WorkflowStageAttempt", - "WorkflowStageLink", - "artifact_model_to_ref", ] diff --git a/scarf/agent/orchestrator/api.py b/scarf/agent/orchestrator/api.py new file mode 100644 index 00000000..3b6c18a2 --- /dev/null +++ b/scarf/agent/orchestrator/api.py @@ -0,0 +1,64 @@ +"""Small entry point for the supported automated RNA analysis.""" + +from pathlib import Path +from typing import Any + +from .main import AgentOrchestrator +from .models import ( + AnalysisError, + AutomatedWorkflowConfig, + AutomatedWorkflowRequest, + AutomatedWorkflowResult, +) + + +def analyze_rna( + source: str | Path, + *, + model: Any, + study_context: str, + study_objective: str, + assay: str | None = None, + zarr_path: str | Path | None = None, + score_doublets: bool = False, +) -> AutomatedWorkflowResult: + """Choose and explain settings for one RNA assay, then execute them. + + ``source`` is a supported input file or an existing Zarr store. ``assay`` + selects the RNA assay when the input contains more than one. The workflow + runs unattended and raises ``AnalysisError`` if essential evidence remains + unresolved or execution fails. A completed result provides ``plot_embedding()``, + ``get_markers()``, and ``report()``. + + Work is bounded by the advanced orchestrator's screening and full-cohort + limits. Use that interface for explicit workspaces, execution limits, and + resumable pauses. An identical repeated call reuses or resumes exact work. + + ``score_doublets=False`` skips advisory doublet scoring when Harmony is + unavailable or prohibited. Harmony-eligible runs retain the matched doublet + diagnostics required by the correction acceptance gate. Scoring does not + remove cells. Changing this option requires a new workflow destination. + Advisory scoring defaults to disabled for new beginner calls. Pass + ``score_doublets=True`` explicitly to resume a run that enabled scoring. + """ + if model is None or isinstance(model, str) and not model.strip(): + raise ValueError( + "model must be a configured model or a non-empty model identifier" + ) + request = AutomatedWorkflowRequest( + sourcePath=str(source), + zarrPath=str(zarr_path) if zarr_path is not None else None, + studyContext=study_context, + studyObjective=study_objective, + primaryAssay=assay, + markerAssay=assay, + analysisAssays=[assay] if assay is not None else [], + ) + config = AutomatedWorkflowConfig( + inputPolicy="unattended", + scoreDoublets=score_doublets, + ) + result = AgentOrchestrator(model, config=config).run(request) + if result.status != "completed": + raise AnalysisError(result) + return result diff --git a/scarf/agent/orchestrator/budget.py b/scarf/agent/orchestrator/budget.py new file mode 100644 index 00000000..604cfd1a --- /dev/null +++ b/scarf/agent/orchestrator/budget.py @@ -0,0 +1,329 @@ +"""Write-ahead admissions for bounded RNA experiments in the workflow journal.""" + +import hashlib +from typing import Any, cast + +from .. import record_io +from . import journal +from .models import AutomatedWorkflowConfig + + +class CandidateBudgetExceeded(ValueError): + """The next scientific experiment exceeds an explicit execution limit.""" + + +def candidate_identity(inputs: dict[str, Any], *, graph: bool = False) -> str: + """Identify exact numerical inputs, omitting resolution for graph reuse.""" + values = dict(inputs) + parameters = dict(values["parameters"]) + parameters.pop("candidateId", None) + if graph: + parameters.pop("leidenResolution", None) + values["parameters"] = parameters + return hashlib.sha256(record_io.canonical_json_bytes(values)).hexdigest() + + +class CandidateBudget: + """Reconstruct bounded admissions from fixed immutable journal slots.""" + + def __init__( + self, + store: Any, + prefix: str, + workflow_run_id: str, + config: AutomatedWorkflowConfig, + provenance: dict[str, Any], + *, + previous_provenances: tuple[dict[str, Any], ...] = (), + ) -> None: + self.store = store + self.prefix = prefix + self.workflow_run_id = workflow_run_id + self.config = config + self.provenance = provenance + self.admissions: dict[str, list[dict[str, Any]]] = {} + for scope in ("sample0", "sample1", "full"): + limit = ( + config.maxFullPartitions + if scope == "full" + else config.maxScreeningEvaluations + ) + rows: list[dict[str, Any]] = [] + for slot in range(limit): + record = journal.read_checkpoint( + store, + prefix, + workflow_run_id, + self._key(scope, slot, "admission"), + ) + if record is None: + continue + if record["inputs"] not in (provenance, *previous_provenances): + raise ValueError( + "Candidate admission has changed checkpoint inputs" + ) + row = record["outputs"] + if ( + row.get("slot") != slot + or row.get("scope") != scope + or slot != len(rows) + ): + raise ValueError("Candidate admission history is inconsistent") + rows.append(row) + self.admissions[scope] = rows + self.repairs = self._read_repairs(previous_provenances) + + @staticmethod + def _repair_inputs( + cells: dict[str, Any], setting: dict[str, Any], experiment: dict[str, Any] + ) -> dict[str, Any]: + """Identify a scientific intervention independently of review attempts.""" + baseline = dict(setting) + baseline["parameters"] = { + key: value + for key, value in setting["parameters"].items() + if key != "candidateId" + } + return {"cells": cells, "baseline": baseline, "experiment": experiment} + + def _read_repairs( + self, previous_provenances: tuple[dict[str, Any], ...] + ) -> dict[str, dict[str, Any]]: + """Derive reserved repairs from this journal, including earlier decisions.""" + repairs: dict[str, dict[str, Any]] = {} + explicit_admission = False + repair_keys = [ + f"parameter_tuning/full/repairs/{slot}" + for slot in range(self.config.maxFullRepairs + 1) + ] + scopes = {"parameter_tuning/full"} | { + "parameter_tuning/evidence_revisions/" + + hashlib.sha256(record_io.canonical_json_bytes(provenance)).hexdigest() + + "/full" + for provenance in (self.provenance, *previous_provenances) + } + review_keys = [ + f"{scope}/review{index}{suffix}" + for scope in sorted(scopes) + for index in range(self.config.maxFullPartitions + 1) + for suffix in ("", "/answer") + ] + for key in [*repair_keys, *review_keys]: + record = journal.read_checkpoint( + self.store, self.prefix, self.workflow_run_id, key + ) + if record is None: + continue + if key in repair_keys: + inputs = record["inputs"] + identity = hashlib.sha256( + record_io.canonical_json_bytes(inputs) + ).hexdigest() + if record["outputs"] != {"identity": identity, "reserved": True}: + raise ValueError("Full-cohort repair admission is inconsistent") + repairs[identity] = inputs + explicit_admission = True + continue + # At most one repair is supported. Once explicitly admitted, later + # unadmitted proposals must not masquerade as additional work. + if explicit_admission: + continue + evidence, action = record["inputs"], record["outputs"]["action"] + if ( + action.get("action") != "experiment" + or evidence.get("comparisonCoverage", {}).get("phase") != "validation" + ): + continue + experiment = evidence["experiments"][action["experimentId"]] + if experiment["parameter"] == "useHarmony": + continue + candidate_id = action["selectedCandidateId"] + candidate = next( + row + for row in evidence["candidates"] + if row["candidateId"] == candidate_id + ) + inputs = self._repair_inputs( + candidate["cellSelection"], + evidence["settings"][candidate_id], + experiment, + ) + identity = hashlib.sha256( + record_io.canonical_json_bytes(inputs) + ).hexdigest() + repairs[identity] = inputs + if len(repairs) > self.config.maxFullRepairs: + raise CandidateBudgetExceeded( + "Saved full-cohort repairs exceed the configured allowance" + ) + return repairs + + def admit_repair( + self, cells: dict[str, Any], setting: dict[str, Any], experiment: dict[str, Any] + ) -> None: + """Reserve one global full-cohort repair before executing its intervention.""" + inputs = self._repair_inputs(cells, setting, experiment) + identity = hashlib.sha256(record_io.canonical_json_bytes(inputs)).hexdigest() + if identity in self.repairs: + return + if len(self.repairs) >= self.config.maxFullRepairs: + raise CandidateBudgetExceeded( + "The allowed full-cohort repair has been used; scientific acceptance remains unresolved" + ) + journal.save_checkpoint( + self.store, + self.prefix, + self.workflow_run_id, + f"parameter_tuning/full/repairs/{len(self.repairs)}", + inputs, + {"identity": identity, "reserved": True}, + ) + self.repairs[identity] = inputs + + def admission_provenance(self, admission: dict[str, Any]) -> dict[str, Any]: + """Return the verified original inputs without changing an admission.""" + record = journal.read_checkpoint( + self.store, + self.prefix, + self.workflow_run_id, + self._key(admission["scope"], admission["slot"], "admission"), + ) + if record is None or record["outputs"] != admission: + raise ValueError("Candidate admission is missing or changed") + return cast(dict[str, Any], record["inputs"]) + + @staticmethod + def _key(scope: str, slot: int, kind: str) -> str: + return f"parameter_tuning/{scope}/evaluation{slot}/{kind}" + + def admit(self, scope: str, inputs: dict[str, Any]) -> dict[str, Any]: + if scope not in self.admissions: + raise ValueError("Unknown candidate execution scope") + identity = candidate_identity(inputs) + rows = self.admissions[scope] + for row in rows: + if row["identity"] == identity: + return row + graph_identity = candidate_identity(inputs, graph=True) + if scope == "full": + if len(rows) >= self.config.maxFullPartitions: + raise CandidateBudgetExceeded("Full-cohort partition limit reached") + graphs = {row["graphIdentity"] for row in rows} | {graph_identity} + if len(graphs) > self.config.maxFullGraphs: + raise CandidateBudgetExceeded("Full-cohort graph limit reached") + else: + if len(rows) >= self.config.maxScreeningEvaluations: + raise CandidateBudgetExceeded("Screening candidate limit reached") + total = sum(len(self.admissions[name]) for name in ("sample0", "sample1")) + if total >= self.config.maxTotalScreeningEvaluations: + raise CandidateBudgetExceeded("Total screening candidate limit reached") + row = { + "slot": len(rows), + "scope": scope, + "identity": identity, + "graphIdentity": graph_identity, + "executionInputs": inputs, + } + journal.save_checkpoint( + self.store, + self.prefix, + self.workflow_run_id, + self._key(scope, len(rows), "admission"), + inputs=self.provenance, + outputs=row, + ) + rows.append(row) + return row + + def check_many( + self, scope: str, inputs: list[dict[str, Any]] + ) -> dict[str, dict[str, Any]]: + """Check room for a comparison without committing an admission.""" + existing = self.admissions[scope] + identities = {row["identity"] for row in existing} + pending = { + candidate_identity(value): value + for value in inputs + if candidate_identity(value) not in identities + } + total = len(existing) + len(pending) + if scope == "full": + graphs = {row["graphIdentity"] for row in existing} | { + candidate_identity(value, graph=True) for value in pending.values() + } + if ( + total > self.config.maxFullPartitions + or len(graphs) > self.config.maxFullGraphs + ): + raise CandidateBudgetExceeded( + "The full-cohort comparison exceeds the remaining graph/partition limit" + ) + elif ( + total > self.config.maxScreeningEvaluations + or sum(len(self.admissions[name]) for name in ("sample0", "sample1")) + + len(pending) + > self.config.maxTotalScreeningEvaluations + ): + raise CandidateBudgetExceeded( + "The screening comparison exceeds the remaining candidate limit" + ) + return pending + + def admit_many(self, scope: str, inputs: list[dict[str, Any]]) -> None: + """Check a baseline or matched pair in full before its first computation.""" + for value in self.check_many(scope, inputs).values(): + self.admit(scope, value) + + def completed(self, admission: dict[str, Any]) -> dict[str, Any] | None: + return journal.load_checkpoint( + self.store, + self.prefix, + self.workflow_run_id, + self._key(admission["scope"], admission["slot"], "complete"), + inputs=admission, + ) + + def completed_source( + self, inputs: dict[str, Any] + ) -> tuple[dict[str, Any], dict[str, Any]] | None: + """Find exact completed evidence before charging additional validation work.""" + identity = candidate_identity(inputs) + for rows in self.admissions.values(): + for admission in rows: + if admission["identity"] == identity: + completed = self.completed(admission) + if completed is not None: + return admission, completed + return None + + def complete(self, admission: dict[str, Any], output: dict[str, Any]) -> None: + journal.save_checkpoint( + self.store, + self.prefix, + self.workflow_run_id, + self._key(admission["scope"], admission["slot"], "complete"), + inputs=admission, + outputs=output, + ) + + def summary(self) -> dict[str, Any]: + """Count distinct reserved and completed comparisons, including reuse.""" + scopes = {} + for scope, rows in self.admissions.items(): + completed = [row for row in rows if self.completed(row) is not None] + scopes[scope] = { + state: { + "graphs": len({row["graphIdentity"] for row in entries}), + "partitions": len({row["identity"] for row in entries}), + } + for state, entries in (("reserved", rows), ("completed", completed)) + } + return { + "scopes": scopes, + "limits": { + "perScreen": self.config.maxScreeningEvaluations, + "totalScreens": self.config.maxTotalScreeningEvaluations, + "fullPartitions": self.config.maxFullPartitions, + "fullGraphs": self.config.maxFullGraphs, + }, + } diff --git a/scarf/agent/orchestrator/context.py b/scarf/agent/orchestrator/context.py index 74c96a39..c20d7b16 100644 --- a/scarf/agent/orchestrator/context.py +++ b/scarf/agent/orchestrator/context.py @@ -1,40 +1,65 @@ -"""Ingest, enrichment, HTO, and experimental-context workflow stages.""" +"""Ingest, RNA enrichment, quality metrics, and experimental-context stages.""" -import re +import hashlib from collections.abc import Mapping, Sequence from typing import Any, cast -import numpy as np - from ...datastore.datastore import DataStore from ...utils.logging import logger -from ..data_enrichment import ( - DataEnrichmentAgent, +from ..data_enrichment.agent import DataEnrichmentAgent +from ..data_enrichment.contracts import ( DataEnrichmentContext, DataEnrichmentReport, ) -from ..experimental_context import ( - ExperimentalContextAgent, +from ..experimental_context.agent import ExperimentalContextAgent +from ..experimental_context.contracts import ( ExperimentalContextResult, NamedArtifactSource, ) -from ..ingest import IngestResult -from ..persistence import ( - AgentInvocation, - AgentReportReference, - AgentWorkflowRun, +from ..experimental_context.requirements import objective_evidence +from ..experimental_context.study import ( + StudyContract, + build_study_contract, + validate_objective_evidence, ) -from ..types import ArtifactReferenceModel +from ..ingest import IngestResult +from ..ingest.manifest import DatasetManifest, is_author_label_column +from ..record_io import canonical_json_bytes +from ..types import AgentRunInfo, ArtifactReferenceModel from . import journal from .models import ( OrchestrationRequestRecord, OrchestrationResumeRecord, + StageEvidenceReference, + WorkflowIdentity, WorkflowNeedsInput, WorkflowQuestion, WorkflowStageAttempt, WorkflowStageLink, artifact_model_to_ref, ) +from .rna import ( + selected_store_rna_assay, + validate_rna_context, + validate_rna_directions, +) + + +def _context_metadata_identity( + store: DataStore, request_record: OrchestrationRequestRecord +) -> dict[str, str]: + """Bind added metadata without repeating the original resume fingerprint scan.""" + from ..parameter_tuning.execution import _metadata_column_fingerprint + + original = request_record.inputIdentity.get("data", {}).get("metadata", {}) + # Public resume already validates every original column against this identity. + # Enrichment may add columns afterward; those also bind committed context work. + return { + column: original[column] + if column in original + else _metadata_column_fingerprint(store.cells, column) + for column in sorted(store.cells.columns) + } class ContextStagesMixin: @@ -111,9 +136,10 @@ def record_ingest_stage( self, store: DataStore, prefix: str, - workflow: AgentWorkflowRun, + workflow: WorkflowIdentity, request_record: OrchestrationRequestRecord, ingest_result: IngestResult, + dataset_manifest: DatasetManifest | None = None, ) -> WorkflowStageAttempt: existing = journal._validated_done_outcome( store, @@ -124,7 +150,7 @@ def record_ingest_stage( [], ) if existing is not None: - logger.info( + logger.debug( f"Workflow {workflow.workflowRunId}: reusing persisted ingest stage" ) return existing @@ -159,6 +185,11 @@ def record_ingest_stage( else None ), "summary": ingest_result.summary, + "datasetManifest": ( + dataset_manifest.model_dump(mode="json") + if dataset_manifest is not None + else None + ), "operations": [ { "operation": "snapshot_cell_selection", @@ -171,7 +202,7 @@ def record_ingest_stage( notes=ingest_result.notes, ) journal._save_outcome(store.zw, prefix, outcome) - logger.info( + logger.debug( f"Workflow {workflow.workflowRunId}: ingest recorded " f"{len(ingest_result.assayNames)} assay(s)" ) @@ -180,7 +211,7 @@ def record_ingest_stage( def data_enrichment_stage( self, store: DataStore, - workflow: AgentWorkflowRun, + workflow: WorkflowIdentity, request_record: OrchestrationRequestRecord, parents: Sequence[WorkflowStageLink], cell_selection: ArtifactReferenceModel, @@ -188,6 +219,7 @@ def data_enrichment_stage( *, resume_record: OrchestrationResumeRecord | None = None, ) -> tuple[WorkflowStageAttempt, DataEnrichmentReport]: + selected = selected_store_rna_assay(store, request_record.request) prefix = journal._ensure_orchestration_store(store) existing = journal._validated_done_outcome( store, @@ -198,14 +230,23 @@ def data_enrichment_stage( parents, ) if existing is not None: - logger.info( + logger.debug( f"Workflow {workflow.workflowRunId}: reusing Data Enrichment report" ) report = journal.load_stage_report(store, existing, DataEnrichmentReport) - return existing, cast(DataEnrichmentReport, report) + report = cast(DataEnrichmentReport, report) + if ( + len(report.policies) != 1 + or report.policies[0].assay != selected + or report.policies[0].assayModality != "RNA" + ): + raise ValueError( + "Saved enrichment includes unsupported assays; start a new RNA workflow." + ) + return existing, report request = request_record.request - selected_assays = request.analysisAssays or list(store.assay_names) - logger.info( + selected_assays = [selected] + logger.debug( f"Workflow {workflow.workflowRunId}: Data Enrichment will inspect " f"{len(selected_assays)} assay(s)" ) @@ -230,6 +271,7 @@ def data_enrichment_stage( parents, inputs={ "studyContext": request.studyContext, + "studyObjective": request.studyObjective, "assays": selected_assays, "cellSelection": cell_selection.model_dump(mode="json"), "allowDownload": request_record.config.allowDownloads, @@ -240,7 +282,10 @@ def data_enrichment_stage( actions: list[str] = [] operations: list[dict[str, Any]] = [] try: - context_payload: dict[str, Any] = {"studyContext": request.studyContext} + context_payload: dict[str, Any] = { + "studyContext": request.studyContext, + "studyObjective": request.studyObjective, + } supplied_context = answers.get("dataEnrichmentContext") if isinstance(supplied_context, Mapping): context_payload.update(dict(supplied_context)) @@ -250,7 +295,6 @@ def data_enrichment_stage( recovered = journal._recover_persisted_stage_report( store, started, - agent_name="data_enrichment", expected_type=DataEnrichmentReport, ) if recovered is not None: @@ -258,7 +302,7 @@ def data_enrichment_stage( report = cast(DataEnrichmentReport, recovered_report) actions.append("recover_persisted_data_enrichment_report") else: - logger.info( + logger.debug( f"Workflow {workflow.workflowRunId}: invoking Data Enrichment" ) agent = DataEnrichmentAgent( @@ -271,32 +315,50 @@ def data_enrichment_stage( assays=selected_assays, cache_dir=request_record.config.cacheDir, allow_download=request_record.config.allowDownloads, + on_attempt=journal.model_attempt_callback( + store, + prefix, + workflow.workflowRunId, + "data_enrichment", + { + "requestSha256": request_record.requestSha256, + "configSha256": request_record.configSha256, + "parents": [ + parent.model_dump(mode="json") for parent in parents + ], + "inputs": started.inputs, + }, + ), ) saved_report, reference = journal._save_stage_report( store, started, report, - invocation=AgentInvocation( - agentName="data_enrichment", - inputs={ - "context": enrichment_context.model_dump(mode="json"), - "assays": selected_assays, - "cellSelection": cell_selection.model_dump(mode="json"), - "cacheDir": request_record.config.cacheDir, - "allowDownload": request_record.config.allowDownloads, - }, - artifacts={"cellSelection": cell_selection}, - runConfig=agent.config, - ), expected_type=DataEnrichmentReport, ) report = cast(DataEnrichmentReport, saved_report) - logger.info( + logger.debug( f"Workflow {workflow.workflowRunId}: Data Enrichment returned " f"status={report.status!r}, policies={len(report.policies)}, " f"inspections={len(report.inspections)}" ) if report.status == "needsInput": + if request_record.config.inputPolicy == "unattended": + outcome = journal._complete_attempt( + started, + status="failed", + report_references=[reference], + artifacts={"cellSelection": cell_selection}, + actions=actions, + outputs={"operations": operations}, + error=( + "The unattended Data Enrichment stage returned an " + "unresolved decision" + ), + notes=report.limitations, + ) + journal._save_outcome(store.zw, prefix, outcome) + return outcome, report questions = [ WorkflowQuestion( questionId="dataEnrichmentContext", @@ -343,10 +405,6 @@ def data_enrichment_stage( notes=report.limitations, ) journal._save_outcome(store.zw, prefix, outcome) - if outcome.status == "failed": - journal.finalize_failed( - store, workflow, outcome.error or "enrichment failed" - ) return outcome, report except Exception as exc: outcome = journal.finish_exception( @@ -361,10 +419,10 @@ def data_enrichment_stage( ) return outcome, DataEnrichmentReport.get_blank() - def _hto_stage( + def _rna_quality_metrics_stage( self, store: DataStore, - workflow: AgentWorkflowRun, + workflow: WorkflowIdentity, request_record: OrchestrationRequestRecord, parents: Sequence[WorkflowStageLink], enrichment: DataEnrichmentReport, @@ -372,44 +430,53 @@ def _hto_stage( *, resume_record: OrchestrationResumeRecord | None = None, ) -> WorkflowStageAttempt: + selected = selected_store_rna_assay(store, request_record.request) + if ( + len(enrichment.policies) != 1 + or enrichment.policies[0].assay != selected + or enrichment.policies[0].assayModality != "RNA" + ): + raise ValueError( + "Quality metrics require enrichment of only the selected RNA assay" + ) prefix = journal._ensure_orchestration_store(store) existing = journal._validated_done_outcome( store, prefix, workflow.workflowRunId, - "hto_demultiplexing", + "rna_quality_metrics", request_record, parents, ) if existing is not None: + if "percentageDefinitions" not in existing.outputs: + raise ValueError( + "Saved RNA quality metrics lack exact percentage definitions; " + "start a new workflow. Existing analysis artifacts remain accessible." + ) self._named_stage_artifacts( existing, "qualityMetricArtifacts", "quality_metric", ) - self._named_stage_artifacts( + hto_sources = self._named_stage_artifacts( existing, "htoIdentityArtifacts", "hto_identity", ) - logger.info( - f"Workflow {workflow.workflowRunId}: reusing HTO demultiplexing stage" - ) + if hto_sources: + raise ValueError( + "Saved automatic HTO processing is unsupported; start a new RNA workflow." + ) + logger.info("Reusing RNA quality metrics") return existing cell_selection_ref = artifact_model_to_ref(cell_selection) - eligible_hto = sum( - policy.assayModality == "HTO" and policy.demultiplexEligible - for policy in enrichment.policies - ) - logger.info( - f"Workflow {workflow.workflowRunId}: HTO stage found " - f"{eligible_hto} eligible assay(s)" - ) + logger.info("Computing RNA quality metrics") started = journal._start_attempt( store.zw, prefix, workflow.workflowRunId, - "hto_demultiplexing", + "rna_quality_metrics", request_record, parents, inputs={ @@ -428,138 +495,52 @@ def _hto_stage( } artifacts: dict[str, ArtifactReferenceModel] = {"cellSelection": cell_selection} try: - inspections = {value.assay: value for value in enrichment.inspections} - for policy in enrichment.policies: - if policy.assayModality == "RNA": - inspection = inspections.get(policy.assay) - observed_families = ( - { - value.family - for value in inspection.families - if value.count > 0 - } - if inspection is not None - else set() - ) - assay = store.get_assay(policy.assay) - feature_ids = np.asarray(assay.feats.fetch_all("ids")).astype(str) - feature_names = np.asarray(assay.feats.fetch_all("names")).astype( - str - ) - family_patterns = ( - ( - "mitochondrial", - r"^(MT-|mt-)", - "percentMito", - "percent_mito", - ), - ( - "ribosomal", - r"^(RPS|RPL|MRPS|MRPL|Rps|Rpl|Mrps|Mrpl)", - "percentRibo", - "percent_ribo", - ), + from ..experimental_context.qc_evidence import _rna_percentage_feature_masks + + definitions = [] + for family, suffix, pattern, mask in _rna_percentage_feature_masks( + store, selected + ): + definition = { + "family": family, + "pattern": pattern, + "matchedGenes": int(mask.sum()), + } + definitions.append(definition) + if not mask.any(): + definition["limitation"] = ( + "No genes match this symbol-based percentage definition" ) - for ( - family, - pattern, - artifact_suffix, - action_suffix, - ) in family_patterns: - if family not in observed_families: - continue - compiled = re.compile(pattern) - mask = np.fromiter( - ( - compiled.search(feature_id) is not None - or compiled.search(feature_name) is not None - for feature_id, feature_name in zip( - feature_ids, - feature_names, - strict=True, - ) - ), - dtype=bool, - count=assay.feats.N, - ) - if not mask.any(): - continue - features_ref = store.set_feature_selection( - from_assay=policy.assay, - mask=mask, - invalidate_cache=False, - ) - metric_ref = store.run_feature_percentage( - cell_selection_ref, - features_ref, - invalidate_cache=False, - ) - features_model = ArtifactReferenceModel.from_artifact_ref( - features_ref - ) - metric_model = ArtifactReferenceModel.from_artifact_ref( - metric_ref - ) - artifact_name = f"{policy.assay}_{artifact_suffix}" - if artifact_name in artifacts: - raise ValueError( - f"Duplicate generated artifact name {artifact_name!r}" - ) - source = NamedArtifactSource( - name=artifact_name, - artifact=metric_model, - ) - artifacts[f"{artifact_name}_features"] = features_model - artifacts[artifact_name] = metric_model - cast( - list[dict[str, Any]], - outputs["qualityMetricArtifacts"], - ).append(source.model_dump(mode="json")) - cast(list[dict[str, Any]], outputs["operations"]).append( - { - "operation": "run_feature_percentage", - "assay": policy.assay, - "family": family, - "pattern": pattern, - "cellSelection": cell_selection.model_dump(mode="json"), - "features": features_model.model_dump(mode="json"), - "artifact": metric_model.model_dump(mode="json"), - } - ) - actions.append(f"compute_{action_suffix}:{policy.assay}") - if policy.assayModality != "HTO" or not policy.demultiplexEligible: continue - identity_ref = store.run_hto_demultiplexing( - cell_selection_ref, - from_assay=policy.assay, - random_seed=0, - invalidate_cache=False, - ) - identity_model = ArtifactReferenceModel.from_artifact_ref(identity_ref) - artifact_name = f"{policy.assay}_htoIdentity" - if artifact_name in artifacts: - raise ValueError( - f"Duplicate generated artifact name {artifact_name!r}" - ) - source = NamedArtifactSource( - name=artifact_name, - artifact=identity_model, + features_ref = store.set_feature_selection( + from_assay=selected, mask=mask, invalidate_cache=False ) - artifacts[artifact_name] = identity_model - cast(list[dict[str, Any]], outputs["htoIdentityArtifacts"]).append( + metric_ref = store.run_feature_percentage( + cell_selection_ref, features_ref, invalidate_cache=False + ) + features_model = ArtifactReferenceModel.from_artifact_ref(features_ref) + metric_model = ArtifactReferenceModel.from_artifact_ref(metric_ref) + artifact_name = f"{selected}_{suffix}" + source = NamedArtifactSource(name=artifact_name, artifact=metric_model) + artifacts[f"{artifact_name}_features"] = features_model + artifacts[artifact_name] = metric_model + cast(list[dict[str, Any]], outputs["qualityMetricArtifacts"]).append( source.model_dump(mode="json") ) cast(list[dict[str, Any]], outputs["operations"]).append( { - "operation": "run_hto_demultiplexing", - "assay": policy.assay, + "operation": "run_feature_percentage", + "assay": selected, + **definition, "cellSelection": cell_selection.model_dump(mode="json"), - "randomSeed": 0, - "invalidateCache": False, - "artifact": identity_model.model_dump(mode="json"), + "features": features_model.model_dump(mode="json"), + "artifact": metric_model.model_dump(mode="json"), } ) - actions.append(f"demultiplex_hto:{policy.assay}") + actions.append( + f"compute_{'percent_mito' if family == 'mitochondrial' else 'percent_ribo'}:{selected}" + ) + outputs["percentageDefinitions"] = definitions outcome = journal._complete_attempt( started, status="done", @@ -568,11 +549,7 @@ def _hto_stage( actions=actions, ) journal._save_outcome(store.zw, prefix, outcome) - logger.info( - f"Workflow {workflow.workflowRunId}: HTO stage produced " - f"{len(cast(list[dict[str, Any]], outputs['htoIdentityArtifacts']))} " - "identity artifact(s)" - ) + logger.info("RNA quality metrics completed") return outcome except Exception as exc: return journal.finish_exception( @@ -589,23 +566,31 @@ def _hto_stage( def experimental_context_stage( self, store: DataStore, - workflow: AgentWorkflowRun, + workflow: WorkflowIdentity, request_record: OrchestrationRequestRecord, parents: Sequence[WorkflowStageLink], cell_selection: ArtifactReferenceModel, - enrichment_reference: AgentReportReference, + enrichment_reference: StageEvidenceReference, quality_metric_artifacts: Sequence[NamedArtifactSource], hto_identity_artifacts: Sequence[NamedArtifactSource], answers: Mapping[str, Any], *, resume_record: OrchestrationResumeRecord | None = None, ) -> tuple[WorkflowStageAttempt, ExperimentalContextResult]: + context_revision: dict[str, Any] = {} + prior_context: ExperimentalContextResult | None = None + selected = selected_store_rna_assay(store, request_record.request) + if hto_identity_artifacts: + raise ValueError( + "Automatic HTO identities are unsupported by the RNA workflow" + ) prefix = journal._ensure_orchestration_store(store) context_artifacts = self._experimental_context_artifacts( cell_selection, quality_metric_artifacts, hto_identity_artifacts, ) + metadata_identity = _context_metadata_identity(store, request_record) existing = journal._validated_done_outcome( store, prefix, @@ -615,7 +600,28 @@ def experimental_context_stage( parents, ) if existing is not None: - logger.info( + saved_metadata = existing.inputs.get("metadataFingerprints") + if saved_metadata is not None and saved_metadata != metadata_identity: + raise ValueError( + "Experimental Context metadata changed; start a new analysis" + ) + if saved_metadata is None: + # Older compatible stages may already have committed downstream + # metadata identities. Verify them; do not rewrite the old report. + for tuning in journal._stage_starts( + store.zw, prefix, workflow.workflowRunId, "parameter_tuning" + ): + if any( + metadata_identity.get(column) != digest + for column, digest in tuning.inputs.get( + "metadataFingerprints", {} + ).items() + ): + raise ValueError( + "Experimental Context metadata differs from saved tuning evidence; " + "restore the original inputs or start a new analysis" + ) + logger.debug( f"Workflow {workflow.workflowRunId}: reusing Experimental Context " "report" ) @@ -623,6 +629,39 @@ def experimental_context_stage( store, existing, ExperimentalContextResult ) resolved_report = cast(ExperimentalContextResult, report) + validate_rna_context(resolved_report, selected) + saved_contract = StudyContract.model_validate( + existing.outputs.get("studyContract") + ) + validate_objective_evidence(saved_contract) + current_requirements, current_coverage = objective_evidence( + study_context=request_record.request.studyContext, + study_objective=request_record.request.studyObjective, + experimental_result=resolved_report, + ) + missing_questions = [ + item.model_dump(mode="json") + for item in current_requirements + if item.requirementId.startswith("requestedDesign:") + and any( + row.requirementId == item.requirementId + and row.status == "unsupported" + for row in current_coverage + ) + ] + if missing_questions: + context_revision = { + "reassessContextReport": existing.reportReferences[0].model_dump( + mode="json" + ), + "requiredDesignQuestions": missing_questions, + } + prior_context = resolved_report + logger.info( + "Experimental context: reassessing explicit study questions missing from prior evidence" + ) + else: + validate_objective_evidence(saved_contract, resolved_report) if existing.artifacts != context_artifacts: raise ValueError( "Persisted Experimental Context stage artifacts are stale" @@ -639,15 +678,8 @@ def experimental_context_stage( raise ValueError( "Persisted Experimental Context HTO artifacts are stale" ) - record = journal.load_agent_record( - store, - existing.reportReferences[0], - ) - if record.invocation.artifacts != context_artifacts: - raise ValueError( - "Persisted Experimental Context invocation artifacts are stale" - ) - return existing, resolved_report + if not context_revision: + return existing, resolved_report cell_selection_ref = artifact_model_to_ref(cell_selection) paused = journal._validated_done_outcome( store, @@ -664,6 +696,88 @@ def experimental_context_stage( directions.update(dict(supplied_directions)) elif isinstance(supplied_directions, str) and supplied_directions.strip(): directions["callerAnswer"] = supplied_directions.strip() + validate_rna_directions(directions) + if request_record.request.authorLabelPolicy == "holdout": + held_out_columns = sorted( + column + for column in store.cells.columns + if is_author_label_column(column) + ) + existing_exclusions = directions.get("excludeColumns") + if existing_exclusions is None: + existing_exclusion_list: list[str] = [] + elif isinstance(existing_exclusions, list) and all( + isinstance(value, str) for value in existing_exclusions + ): + existing_exclusion_list = existing_exclusions + else: + raise ValueError("experimentalDirections.excludeColumns must be a list") + + referenced_held_out: set[str] = set() + + def find_held_out_references(value: Any) -> None: + if isinstance(value, str): + if value in held_out_columns: + referenced_held_out.add(value) + return + if isinstance(value, Mapping): + for nested in value.values(): + find_held_out_references(nested) + return + if isinstance(value, list | tuple | set): + for nested in value: + find_held_out_references(nested) + + for key, value in directions.items(): + if key != "excludeColumns": + find_held_out_references(value) + if referenced_held_out: + raise ValueError( + "authorLabelPolicy='holdout' forbids runtime use of author " + "annotation columns: " + ", ".join(sorted(referenced_held_out)) + ) + directions["excludeColumns"] = sorted( + { + *held_out_columns, + *existing_exclusion_list, + } + ) + retry_inputs: dict[str, Any] = {} + failed = journal._validated_done_outcome( + store, + prefix, + workflow.workflowRunId, + "experimental_context", + request_record, + parents, + required_status="failed", + ) + if failed is not None and "retryAfterFailedReport" in failed.inputs: + # A later persistence or validation error must keep the same retry + # identity so its already committed decision remains recoverable. + retry_inputs["retryAfterFailedReport"] = failed.inputs[ + "retryAfterFailedReport" + ] + if failed is not None and failed.reportReferences: + failed_report = journal.load_stage_report( + store, failed, ExperimentalContextResult + ) + if cast(ExperimentalContextResult, failed_report).status in { + "failed", + "needsInput", + }: + # An unresolved unattended answer is also an unsuccessful attempt. + # Keep it immutable while reusing its measured design evidence. + # A committed successful retry still has a stable recovery key. + retry_inputs["retryAfterFailedReport"] = failed.reportReferences[ + 0 + ].model_dump(mode="json") + if ( + cast(ExperimentalContextResult, failed_report).status + == "needsInput" + ): + prior_context = cast(ExperimentalContextResult, failed_report) + logger.info("Retrying experimental context after the previous failure") started = journal._start_attempt( store.zw, prefix, @@ -672,9 +786,13 @@ def experimental_context_stage( request_record, parents, inputs={ + **retry_inputs, + **context_revision, "studyContext": request_record.request.studyContext, + "studyObjective": request_record.request.studyObjective, "cellSelection": cell_selection.model_dump(mode="json"), "directions": directions, + "metadataFingerprints": metadata_identity, "qualityMetricArtifacts": [ source.model_dump(mode="json") for source in quality_metric_artifacts @@ -685,7 +803,7 @@ def experimental_context_stage( }, resume_record=resume_record, ) - logger.info( + logger.debug( f"Workflow {workflow.workflowRunId}: Experimental Context will evaluate " f"{len(quality_metric_artifacts)} quality metric artifact(s) and " f"{len(hto_identity_artifacts)} HTO identity artifact(s)" @@ -697,29 +815,31 @@ def experimental_context_stage( and paused.outputs.get("unsafeBatchCorrection") is True else None ) + no_inference_resolution = ( + paused is not None + and paused.outputs.get("unsafeBatchCorrection") is not True + and isinstance(supplied_directions, Mapping) + and supplied_directions.get("coefficientsOfInterest") == [] + and supplied_directions.get("unitsOfInference") == {} + and isinstance(supplied_directions.get("batchCorrection"), Mapping) + and supplied_directions["batchCorrection"].get("action") == "skip" + ) actions: list[str] = [] recovered = journal._recover_persisted_stage_report( store, started, - agent_name="experimental_context", expected_type=ExperimentalContextResult, ) if recovered is not None: recovered_report, reference = recovered report = cast(ExperimentalContextResult, recovered_report) - recovered_record = journal.load_agent_record(store, reference) - if recovered_record.invocation.artifacts != context_artifacts: - raise ValueError( - "Recovered Experimental Context invocation artifacts are stale" - ) actions.append("recover_persisted_experimental_context_report") else: - parent_reports = [journal._report_link(enrichment_reference)] - if unsafe_resolution == "skip": + if unsafe_resolution == "skip" or no_inference_resolution: assert paused is not None if not paused.reportReferences: raise ValueError( - "Unsafe Experimental Context pause has no persisted report" + "Experimental Context pause has no persisted report" ) prior_report = cast( ExperimentalContextResult, @@ -743,29 +863,49 @@ def experimental_context_stage( raise ValueError( "Paused Experimental Context exact inputs are stale" ) - paused_record = journal.load_agent_record( - store, - paused.reportReferences[0], - ) - if paused_record.invocation.artifacts != context_artifacts: - raise ValueError( - "Paused Experimental Context invocation artifacts are stale" - ) prior_plan = prior_report.decision.batchCorrection + if no_inference_resolution: + plan_updates: dict[str, Any] = { + "preserveColumns": [], + "rationale": ( + "The caller explicitly continued without " + "coefficient-level inference and skipped Harmony." + ), + } + decision_updates: dict[str, Any] = { + "coefficientsOfInterest": [], + "unitsOfInference": {}, + } + resolution_note = ( + "Caller explicitly continued without coefficient-level " + "inference and skipped Harmony." + ) + resolution_action = ( + "resolve_experimental_context:no_inference_skip_harmony" + ) + else: + plan_updates = { + "rationale": ( + "The caller explicitly skipped Harmony after reviewing " + "the persisted unsafe batch-correction evidence." + ) + } + decision_updates = {} + resolution_note = ( + "Caller explicitly skipped Harmony after an unsafe result." + ) + resolution_action = "resolve_unsafe_batch_correction:skip" skip_plan = prior_plan.model_copy( update={ "action": "skip", "batchColumns": [], "metricsRequired": [], - "rationale": ( - "The caller explicitly skipped Harmony after " - "reviewing the persisted unsafe batch-correction " - "evidence." - ), + **plan_updates, } ) decision = prior_report.decision.model_copy( update={ + **decision_updates, "batchCorrection": skip_plan, "needsInput": [], } @@ -774,20 +914,15 @@ def experimental_context_stage( update={ "status": "done", "decision": decision, - "notes": [ - *prior_report.notes, - "Caller explicitly skipped Harmony after an unsafe " - "result.", - ], + "notes": [*prior_report.notes, resolution_note], + "runInfo": AgentRunInfo( + agentName="experimental_context_resolution" + ), } ) - parent_reports.append( - journal._report_link(paused.reportReferences[0]) - ) - run_config = request_record.config.agentRunConfig - actions.append("resolve_unsafe_batch_correction:skip") + actions.append(resolution_action) else: - logger.info( + logger.debug( f"Workflow {workflow.workflowRunId}: invoking Experimental " "Context" ) @@ -795,39 +930,81 @@ def experimental_context_stage( self.model, config=request_record.config.agentRunConfig, ) + evidence_inputs = { + "request": request_record.model_dump(mode="json"), + "parents": [ + parent.model_dump(mode="json") for parent in parents + ], + "context": { + key: value + for key, value in started.inputs.items() + if key != "retryAfterFailedReport" + }, + } + evidence_key = ( + "experimental_context/evidence/" + + hashlib.sha256( + canonical_json_bytes(evidence_inputs) + ).hexdigest() + ) + result_key = "result" + if "retryAfterFailedReport" in retry_inputs: + result_key += ( + "/" + + hashlib.sha256( + canonical_json_bytes( + retry_inputs["retryAfterFailedReport"] + ) + ).hexdigest() + ) + + def read_evidence(key: str) -> dict[str, Any] | None: + return journal.load_checkpoint( + store, + prefix, + workflow.workflowRunId, + evidence_key + + "/" + + (result_key if key == "result" else key), + evidence_inputs, + ) + + def write_evidence(key: str, output: dict[str, Any]) -> None: + journal.save_checkpoint( + store, + prefix, + workflow.workflowRunId, + evidence_key + + "/" + + (result_key if key == "result" else key), + evidence_inputs, + output, + ) + report = agent.run( store, + qc_assay=selected, study_context=request_record.request.studyContext, + study_objective=request_record.request.studyObjective, cell_selection=cell_selection_ref, directions=directions, quality_metric_artifacts=quality_metric_artifacts, hto_identity_artifacts=hto_identity_artifacts, + checkpoint_read=read_evidence, + checkpoint_write=write_evidence, + on_attempt=journal.model_attempt_callback( + store, + prefix, + workflow.workflowRunId, + evidence_key, + evidence_inputs, + ), + previous_context=prior_context, ) - run_config = agent.config saved_report, reference = journal._save_stage_report( store, started, report, - invocation=AgentInvocation( - agentName="experimental_context", - parentReports=parent_reports, - inputs={ - "studyContext": request_record.request.studyContext, - "cellSelection": cell_selection.model_dump(mode="json"), - "directions": directions, - "qualityMetricArtifacts": [ - source.model_dump(mode="json") - for source in quality_metric_artifacts - ], - "htoIdentityArtifacts": [ - source.model_dump(mode="json") - for source in hto_identity_artifacts - ], - "unsafeResolution": unsafe_resolution, - }, - artifacts=context_artifacts, - runConfig=run_config, - ), expected_type=ExperimentalContextResult, ) report = cast(ExperimentalContextResult, saved_report) @@ -835,6 +1012,8 @@ def experimental_context_stage( raise ValueError( "Experimental Context returned a different cell selection" ) + if report.status == "done": + validate_rna_context(report, selected) if report.qualityMetricArtifacts != list(quality_metric_artifacts): raise ValueError( "Experimental Context returned different quality metric artifacts" @@ -843,12 +1022,27 @@ def experimental_context_stage( raise ValueError( "Experimental Context returned different HTO identity artifacts" ) - logger.info( + logger.debug( f"Workflow {workflow.workflowRunId}: Experimental Context returned " f"status={report.status!r}, batchAction=" f"{report.decision.batchCorrection.action!r}" ) if report.status == "needsInput": + if request_record.config.inputPolicy == "unattended": + outcome = journal._complete_attempt( + started, + status="failed", + report_references=[reference], + artifacts=context_artifacts, + error=( + "The unattended Experimental Context stage returned an " + "unresolved decision: " + + "; ".join(report.decision.needsInput or report.notes) + ), + notes=report.notes, + ) + journal._save_outcome(store.zw, prefix, outcome) + return outcome, report questions = [ WorkflowQuestion( questionId="experimentalDirections", @@ -875,7 +1069,10 @@ def experimental_context_stage( artifacts=context_artifacts, error="; ".join(report.notes) or "Experimental Context failed", ) - elif report.decision.batchCorrection.action == "unsafe": + elif ( + report.decision.batchCorrection.action == "unsafe" + and request_record.config.inputPolicy != "unattended" + ): batch_plan = report.decision.batchCorrection outcome = journal._complete_attempt( started, @@ -903,6 +1100,56 @@ def experimental_context_stage( notes=report.notes, ) else: + if report.decision.batchCorrection.action == "unsafe": + actions.append("evaluate_unsafe_harmony_for_diagnosis") + physical_capture = directions.get("physicalCaptureColumn") + if not isinstance(physical_capture, str) or not physical_capture: + physical_capture = None + elif physical_capture not in { + *store.cells.columns, + *report.htoIdentityColumns, + }: + raise ValueError( + "physicalCaptureColumn must identify observed metadata or " + "an exact HTO identity" + ) + study_contract = build_study_contract( + study_context=request_record.request.studyContext, + study_objective=request_record.request.studyObjective, + experimental_result=report, + author_label_policy=(request_record.request.authorLabelPolicy), + physical_capture_column=report.decision.physicalCaptureColumn + or physical_capture, + ) + try: + validate_objective_evidence(study_contract, report) + except ValueError as exc: + unattended = request_record.config.inputPolicy == "unattended" + outcome = journal._complete_attempt( + started, + status="failed" if unattended else "needsInput", + report_references=[reference], + artifacts=context_artifacts, + outputs={ + "studyContract": study_contract.model_dump(mode="json") + }, + actions=actions, + error=str(exc) if unattended else None, + needs_input=None + if unattended + else WorkflowNeedsInput( + questions=[ + WorkflowQuestion( + questionId="experimentalDirections", + question=str(exc), + evidenceIds=list(study_contract.evidenceIds), + ) + ] + ), + notes=[*report.notes, str(exc)], + ) + journal._save_outcome(store.zw, prefix, outcome) + return outcome, report outcome = journal._complete_attempt( started, status="done", @@ -923,15 +1170,12 @@ def experimental_context_stage( for source in quality_metric_artifacts ], "metadataColumns": report.htoIdentityColumns, + "studyContract": study_contract.model_dump(mode="json"), }, actions=actions, notes=report.notes, ) journal._save_outcome(store.zw, prefix, outcome) - if outcome.status == "failed": - journal.finalize_failed( - store, workflow, outcome.error or "context failed" - ) return outcome, report except Exception as exc: outcome = journal.finish_exception( diff --git a/scarf/agent/orchestrator/decisions.py b/scarf/agent/orchestrator/decisions.py new file mode 100644 index 00000000..1fded9cd --- /dev/null +++ b/scarf/agent/orchestrator/decisions.py @@ -0,0 +1,338 @@ +"""Constrained RNA choices persisted with their owning stage evidence.""" + +import hashlib +import json +import time +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from ...utils.logging import logger +from .. import record_io +from ..config.agent_exec import run_agent_sync +from ..decisions.kernel import ( + DecisionRecord, + DecisionSelection, + DecisionSource, + EvidenceBundle, +) +from ..decisions.rna import ( + CompiledRnaDecision, + RnaDecisionDefinition, + compile_rna_decision, +) +from . import journal +from .models import OrchestrationRequestRecord, WorkflowQuestion + + +@dataclass(frozen=True, slots=True) +class DecisionResolution: + """One validated choice or an explicit unresolved scientific question.""" + + record: DecisionRecord | None + compiled: CompiledRnaDecision | None + checkpointSha256: str + pending: WorkflowQuestion | None = None + + +def _sha256(value: object) -> str: + return hashlib.sha256(record_io.canonical_json_bytes(value)).hexdigest() + + +def _record_from_selection( + definition: RnaDecisionDefinition, + evidence: EvidenceBundle, + selection: DecisionSelection, + source: DecisionSource, + model_name: str | None, + prompt_sha256: str | None, + created_at_ns: int, +) -> DecisionRecord: + if evidence.contentSha256 is None: + raise ValueError("Decision evidence must have a content digest") + option = definition.spec.option_by_id().get(selection.selectedOptionId) + if option is None: + raise ValueError("Selected option is not offered") + return DecisionRecord( + recordId=f"decision:{definition.spec.decisionId}:{_sha256(selection.model_dump(mode='json'))[:24]}", + decisionId=definition.spec.decisionId, + definitionVersion=definition.spec.definitionVersion, + evidenceBundleId=evidence.bundleId, + evidenceBundleSha256=evidence.contentSha256, + offeredOptionIds=[v.optionId for v in definition.spec.options], + availableEvidenceIds=[v.evidenceId for v in evidence.evidence], + selectedOptionId=selection.selectedOptionId, + status=option.status, + source=source, + evidenceIds=list(selection.evidenceIds), + rationale=selection.rationale, + confidence=selection.confidence, + protectedVariableEffects=list(selection.protectedVariableEffects), + overrideOfOptionId=selection.overrideOfOptionId, + overrideEvidenceIds=list(selection.overrideEvidenceIds), + modelName=model_name, + promptSha256=prompt_sha256, + createdAtNs=created_at_ns, + softwareSha256=_sha256(definition.model_dump(mode="json")), + ) + + +def _validate_selection( + definition: RnaDecisionDefinition, + evidence: EvidenceBundle, + selection: DecisionSelection, +) -> DecisionSelection: + record = _record_from_selection( + definition, evidence, selection, "agent", None, None, 0 + ) + compile_rna_decision(definition, evidence, record) + return selection + + +def _human_selection( + definition: RnaDecisionDefinition, + evidence: EvidenceBundle, + answer: Mapping[str, Any], +) -> DecisionSelection: + if ( + set(answer) != {"decisionId", "optionId", "rationale"} + or answer["decisionId"] != definition.spec.decisionId + ): + raise ValueError( + "A decision answer must name this decisionId, optionId, and rationale" + ) + option = definition.spec.option_by_id().get(answer["optionId"]) + if option is None: + raise ValueError("Human answer does not select an offered option") + ids = [ + v.evidenceId + for v in evidence.evidence + if not option.requiredEvidenceClasses + or v.evidenceClass in option.requiredEvidenceClasses + ] + ids = list(dict.fromkeys([*option.requiredEvidenceIds, *ids])) + override = ( + definition.spec.requireIndependentOverrideEvidence + and definition.spec.metricPreferredOptionId is not None + and option.optionId != definition.spec.metricPreferredOptionId + and option.status in {"apply", "skip"} + ) + override_ids = ( + [ + v.evidenceId + for v in evidence.evidence + if v.evidenceClass + in { + "markerCoherence", + "resamplingStability", + "crossUnitSupport", + "protectedVariablePreservation", + } + and ( + not option.requiredEvidenceIds + or v.evidenceId in option.requiredEvidenceIds + ) + ] + if override + else [] + ) + return DecisionSelection( + selectedOptionId=option.optionId, + evidenceIds=list(dict.fromkeys([*ids, *override_ids])), + rationale=answer["rationale"], + confidence="notApplicable", + overrideOfOptionId=definition.spec.metricPreferredOptionId + if override + else None, + overrideEvidenceIds=override_ids, + ) + + +class DecisionStagesMixin: + """Resolve choices through scientific validation and one stage-owned checkpoint.""" + + model: Any + + def _resolve_rna_decision( + self, + store: Any, + request_record: OrchestrationRequestRecord, + definition: RnaDecisionDefinition, + evidence: EvidenceBundle, + answers: Mapping[str, Any], + *, + rule_selection: DecisionSelection | None = None, + agent_selection: DecisionSelection | None = None, + agent_model_name: str | None = None, + qc_evidence: Mapping[str, Any] | None = None, + ) -> DecisionResolution: + evidence = ( + evidence if evidence.contentSha256 else evidence.with_content_sha256() + ) + if ( + evidence.decisionId != definition.spec.decisionId + or evidence.bundleId != definition.spec.evidenceBundleId + ): + raise ValueError("Decision and exact evidence identities differ") + if rule_selection is not None and agent_selection is not None: + raise ValueError("A decision cannot have two supplied owners") + decision_id = definition.spec.decisionId + stage = ( + "preprocessing_plan" + if decision_id in {"qcGrouping", "cellQuality"} + else "preprocessing" + ) + question_id = f"decision:{decision_id}" + answer = answers.get(question_id) + identity = { + "requestSha256": request_record.requestSha256, + "configSha256": request_record.configSha256, + "definition": definition.model_dump(mode="json"), + "evidence": evidence.model_dump(mode="json"), + "answer": answer, + "ruleSelection": rule_selection.model_dump(mode="json") + if rule_selection + else None, + "agentSelection": agent_selection.model_dump(mode="json") + if agent_selection + else None, + } + if qc_evidence is not None: + identity["qcPolicyEvidence"] = dict(qc_evidence) + digest = _sha256(identity) + key = f"{stage}/decisions/{decision_id}/{digest}" + prefix = journal._ensure_orchestration_store(store) + saved = journal.load_checkpoint( + store, prefix, request_record.workflowRunId, key, identity + ) + if saved is not None: + record = DecisionRecord.model_validate(saved["record"]) + compiled = compile_rna_decision(definition, evidence, record) + if [v.model_dump(mode="json") for v in compiled.checks] != saved["checks"]: + raise ValueError("Saved decision checks differ from exact replay") + pending = ( + WorkflowQuestion.model_validate(saved["pending"]) + if saved.get("pending") + else None + ) + return DecisionResolution( + record, None if pending else compiled, digest, pending + ) + source: DecisionSource = "agent" + model_name = agent_model_name + prompt_hash = None + if rule_selection is not None: + if answer is not None: + raise ValueError("A rule-owned decision cannot accept a human answer") + selection = rule_selection + source = "rule" + elif agent_selection is not None: + if answer is not None: + raise ValueError( + "A supplied agent decision cannot accept a human answer" + ) + selection = agent_selection + elif answer is not None: + if not isinstance(answer, Mapping): + raise ValueError("Human decision answer must be a mapping") + selection = _human_selection(definition, evidence, answer) + source = "human" + else: + payload = { + "studyContext": request_record.request.studyContext, + "studyObjective": request_record.request.studyObjective, + "question": definition.spec.question, + "spec": definition.spec.model_dump(mode="json"), + "evidence": evidence.model_dump(mode="json"), + } + if qc_evidence is not None: + payload["qcPolicyEvidence"] = dict(qc_evidence) + prompt = json.dumps(payload, indent=2, sort_keys=True) + prompt_hash = hashlib.sha256(prompt.encode()).hexdigest() + execution = run_agent_sync( + model=self.model, + output_type=DecisionSelection, + system_prompt=( + "Assess the offered settings against the study objective using the observed quantitative and qualitative evidence. " + "The objective identifies questions and biology to protect; it does not predetermine the answer. " + "Select only an offered option and cite its required evidence. " + "Write the rationale as two or three plain-language sentences for the analysis report: " + "state the chosen outcome, the relevant measured comparison, and its scientific tradeoff or limitation. " + "Use readable study and measurement names. Keep option identifiers, artifact identifiers and " + "internal field names out of the rationale; cite identifiers in evidenceIds instead. " + "Do not infer nuisance from gene-family names alone. Retain defaults only when evidence supports them. " + "Defer essential unresolved questions. Model failure or a work limit never justifies an unsupported choice. " + "For QC, distinguish retained group coverage from preserved biological structure: " + "group counts and absence of unsafe flags do not establish balanced retention, " + "cell validity, or preservation of marker programs. Compare retention fractions " + "and metric-specific flags where supplied; unmeasured effects remain unknown. " + "Use qcPolicyEvidence for exact thresholds, distributions, retention " + "and limitations. Its Ref fields point to sharedMeasurements in the " + "same payload. Do not substitute the short evidence summaries for " + "these measurements or dismiss an unsafe rejected policy as untested. " + "Within-capture QC does not require an independent biological unit or a healthy reference. " + "An override needs the independent evidence required by the supplied specification." + ), + user_prompt=prompt, + config=request_record.config.agentRunConfig, + name=f"rna_{decision_id}_decision", + output_validator=lambda value: _validate_selection( + definition, evidence, value + ), + on_attempt=journal.model_attempt_callback( + store, prefix, request_record.workflowRunId, key, payload + ), + ) + selection = execution.output + if not isinstance(selection, DecisionSelection): + raise TypeError("Decision model returned an unexpected result") + model_name = execution.runInfo.modelName + record = _record_from_selection( + definition, + evidence, + selection, + source, + model_name, + prompt_hash, + time.time_ns(), + ) + compiled = compile_rna_decision(definition, evidence, record) + pending = ( + WorkflowQuestion( + questionId=question_id, + decisionId=decision_id, + question=f"{definition.spec.question} {record.rationale}", + options=[v.optionId for v in definition.spec.options], + evidenceIds=[v.evidenceId for v in evidence.evidence], + planChecksum=digest, + ) + if record.status in {"defer", "abstain"} + else None + ) + value = { + "stage": stage, + "checkpointSha256": digest, + "spec": definition.spec.model_dump(mode="json"), + "evidence": evidence.model_dump(mode="json"), + "record": record.model_dump(mode="json"), + "checks": [v.model_dump(mode="json") for v in compiled.checks], + "pending": pending.model_dump(mode="json") if pending else None, + } + journal.save_checkpoint( + store, prefix, request_record.workflowRunId, key, identity, value + ) + option = definition.spec.option_by_id()[record.selectedOptionId] + logger.info( + f"{definition.spec.question} Selected {option.label}: {record.rationale}" + ) + return DecisionResolution( + record, None if pending else compiled, digest, pending + ) + + @staticmethod + def _pending_decision_question( + resolution: DecisionResolution, definition: RnaDecisionDefinition + ) -> WorkflowQuestion: + if resolution.pending is None: + raise ValueError("Decision has no pending question") + return resolution.pending diff --git a/scarf/agent/orchestrator/finalization.py b/scarf/agent/orchestrator/finalization.py index 22ee5f4d..519c7472 100644 --- a/scarf/agent/orchestrator/finalization.py +++ b/scarf/agent/orchestrator/finalization.py @@ -1,58 +1,43 @@ -"""Final analysis and biological interpretation workflow stages.""" +"""Validate the selected full RNA analysis and compute its final layout.""" -import json -from collections.abc import Mapping, Sequence -from typing import Any, Literal, cast +from collections.abc import Sequence +from typing import Any from ...datastore.datastore import DataStore -from ...storage.refs import ArtifactRef +from ...graph.feature_projection import graph_cell_selection from ...utils.logging import logger -from ..biological_interpretation import ( - BiologicalContext, - BiologicalInterpretationAgent, - BiologicalInterpretationReport, -) -from ..data_enrichment import DataEnrichmentReport -from ..experimental_context import ExperimentalContextResult -from ..parameter_tuning import ParameterTuningAgent, ParameterTuningReport -from ..persistence import ( - AgentInvocation, - AgentReportReference, - AgentWorkflowRun, -) -from ..types import ArtifactReferenceModel, ExperimentalBiologyHandoff +from ..parameter_tuning.contracts import ParameterTuningReport +from ..parameter_tuning.selection import promote_parameter_candidate +from ..types import ArtifactReferenceModel from . import journal from .models import ( AutomatedPreprocessingPlan, FinalAnalysisHandoff, - NativeAnalysisHandoff, OrchestrationRequestRecord, OrchestrationResumeRecord, PreprocessedAssayHandoff, - ReductionMethod, - WorkflowNeedsInput, - WorkflowQuestion, + StageEvidenceReference, + WorkflowIdentity, WorkflowStageAttempt, WorkflowStageLink, artifact_model_to_ref, ) +from .rna import selected_store_rna_assay, validate_rna_handoffs, validate_rna_plan class FinalizationStagesMixin: - """Finalize layouts, clusters, markers, and biological interpretation.""" - - model: Any + """Finish one full-cohort RNA representation using its exact saved artifacts.""" def analysis_finalization_stage( self, store: DataStore, - workflow: AgentWorkflowRun, + workflow: WorkflowIdentity, request_record: OrchestrationRequestRecord, parents: Sequence[WorkflowStageLink], plan: AutomatedPreprocessingPlan, preprocessed: Sequence[PreprocessedAssayHandoff], tuning_report: ParameterTuningReport, - tuning_reference: AgentReportReference, + tuning_reference: StageEvidenceReference, *, resume_record: OrchestrationResumeRecord | None = None, ) -> tuple[WorkflowStageAttempt, FinalAnalysisHandoff]: @@ -66,19 +51,10 @@ def analysis_finalization_stage( parents, ) if existing is not None: - logger.info( - f"Workflow {workflow.workflowRunId}: reusing finalized analysis" - ) + logger.info("Reusing the validated final RNA analysis") return existing, FinalAnalysisHandoff.model_validate( existing.outputs["finalAnalysis"] ) - if tuning_report.cellSelection is None: - raise ValueError("Parameter Tuning lacks an exact cell selection") - cell_selection = tuning_report.cellSelection - if any(value.cellSelection != cell_selection for value in preprocessed): - raise ValueError( - "Finalization inputs do not share the selected tuning cells" - ) started = journal._start_attempt( store.zw, prefix, @@ -91,223 +67,94 @@ def analysis_finalization_stage( "preprocessedAssays": [ value.model_dump(mode="json") for value in preprocessed ], - "cellSelection": cell_selection.model_dump(mode="json"), - "finalClusters": ( - tuning_report.finalClusterArtifact.model_dump(mode="json") - if tuning_report.finalClusterArtifact is not None - else None - ), }, resume_record=resume_record, ) - artifacts: dict[str, ArtifactReferenceModel] = {"cellSelection": cell_selection} - actions: list[str] = [] + artifacts: dict[str, ArtifactReferenceModel] = {} operations: list[dict[str, Any]] = [] - logger.info( - f"Workflow {workflow.workflowRunId}: finalizing " - f"{len(preprocessed)} native analysis route(s) and markers from " - f"assay {plan.markerAssay!r}" - ) try: + assay_name = selected_store_rna_assay(store, request_record.request) + validate_rna_plan(plan, assay_name) + validate_rna_handoffs(preprocessed, assay_name) + handoff = preprocessed[0] + cells = handoff.cellSelection + if cells is None or cells != tuning_report.cellSelection: + raise ValueError( + "Finalization requires the original full-cohort selection, not sampled cells" + ) + if handoff.normalized is None or handoff.markerFeatures is None: + raise ValueError( + "Finalization requires exact normalized and marker features" + ) if ( tuning_report.status != "done" - or tuning_report.finalClusterArtifact is None + or tuning_report.recommendedIntegrationId is not None + or tuning_report.assayReports ): - raise ValueError("Parameter Tuning has no finalized cluster branch") - preprocessed_by_assay = {value.assay: value for value in preprocessed} - plan_by_assay = {value.assay: value for value in plan.assays} - agent = ParameterTuningAgent( - self.model, - config=request_record.config.agentRunConfig, - ) - native_handoffs, native_umaps = self.finalize_native_analyses( - store, - agent, - request_record, - tuning_report, - preprocessed_by_assay, - artifacts, - actions, - operations, - ) - ( - graph_method, - final_graph, - final_initialization, - final_umap, - ) = self.finalize_selected_graph( - store, - plan, - tuning_report, - native_handoffs, - native_umaps, - actions, - operations, - ) - final_clusters = ArtifactReferenceModel.model_validate( - tuning_report.finalClusterArtifact.model_dump() - ) - marker_handoff = preprocessed_by_assay[plan.markerAssay] - if marker_handoff.markerFeatures is None: - raise ValueError("Marker assay lacks an exact feature panel") - marker_plan = plan_by_assay[plan.markerAssay] - marker_ref = store.run_marker_search( - artifact_model_to_ref(final_clusters), - from_assay=plan.markerAssay, - features=artifact_model_to_ref(marker_handoff.markerFeatures), - invalidate_cache=False, - log_transform=bool( - marker_plan.normalizationParameters.get("logTransform", False) - ), - renormalize_subset=bool( - marker_plan.normalizationParameters.get("renormalizeSubset", False) - ), - ) - if not isinstance(marker_ref, ArtifactRef): - raise TypeError("Saved marker search did not return an artifact") - marker_model = ArtifactReferenceModel.from_artifact_ref(marker_ref) - artifacts.update( - { - "final_graph": final_graph, - "final_clusters": final_clusters, - "final_embedding_initialization": final_initialization, - "final_umap": final_umap, - "marker_features": marker_handoff.markerFeatures, - "markers": marker_model, - } - ) - limitations = list( - dict.fromkeys([*plan.limitations, *tuning_report.limitations]) - ) - if marker_plan.assayType == "ATAC": - limitations.append( - "ATAC peak markers are descriptive and cannot establish " - "confident cell identities alone" + raise ValueError( + "Finalization requires a completed single-RNA recommendation" ) - final_analysis = FinalAnalysisHandoff( - workflowRunId=workflow.workflowRunId, - primaryAssay=plan.primaryAssay, - markerAssay=plan.markerAssay, - cellSelection=cell_selection, - nativeAnalyses=native_handoffs, - graph=final_graph, - graphMethod=graph_method, - clusters=final_clusters, - embeddingInitialization=final_initialization, - umap=final_umap, - markerFeatures=marker_handoff.markerFeatures, - markers=marker_model, - parameterReport=tuning_reference, - limitations=list(dict.fromkeys(limitations)), - ) - actions.append(f"run_markers:{plan.markerAssay}") - operations.append( - { - "operation": "run_marker_search", - "assay": plan.markerAssay, - "clusters": final_clusters.model_dump(mode="json"), - "cellSelection": cell_selection.model_dump(mode="json"), - "features": marker_handoff.markerFeatures.model_dump(mode="json"), - "invalidateCache": False, - "logTransform": bool( - marker_plan.normalizationParameters.get("logTransform", False) - ), - "renormalizeSubset": bool( - marker_plan.normalizationParameters.get( - "renormalizeSubset", False - ) - ), - "artifact": marker_model.model_dump(mode="json"), - } - ) - outcome = journal._complete_attempt( - started, - status="done", - artifacts=artifacts, - outputs={ - "finalAnalysis": final_analysis.model_dump(mode="json"), - "operations": operations, - }, - actions=actions, - notes=final_analysis.limitations, - ) - journal._save_outcome(store.zw, prefix, outcome) - logger.info( - f"Workflow {workflow.workflowRunId}: finalized " - f"graphMethod={graph_method!r}, nativeLayouts=" - f"{len(native_handoffs)}, markerAssay={plan.markerAssay!r}" - ) - return outcome, final_analysis - except Exception as exc: - outcome = journal.finish_exception( + if journal.read_stage_evidence( + store, tuning_reference + ) != tuning_report.model_dump(mode="json"): + raise ValueError("Final tuning differs from its committed evidence") + selected = promote_parameter_candidate( store, - prefix, - workflow, - started, - exc, - artifacts=artifacts, - actions=actions, - outputs={ - "operations": operations, - }, - ) - return outcome, FinalAnalysisHandoff.get_blank() - - def finalize_native_analyses( - self, - store: DataStore, - agent: ParameterTuningAgent, - request_record: OrchestrationRequestRecord, - tuning_report: ParameterTuningReport, - preprocessed_by_assay: Mapping[str, PreprocessedAssayHandoff], - artifacts: dict[str, ArtifactReferenceModel], - actions: list[str], - operations: list[dict[str, Any]], - ) -> tuple[ - list[NativeAnalysisHandoff], - dict[ - str, - tuple[ArtifactReferenceModel, ArtifactReferenceModel], - ], - ]: - native_handoffs: list[NativeAnalysisHandoff] = [] - native_umaps: dict[ - str, - tuple[ArtifactReferenceModel, ArtifactReferenceModel], - ] = {} - for assay, assay_report in tuning_report.assayReports.items(): - logger.info( - f"Finalizing native analysis for assay {assay!r} with candidate " - f"{assay_report.recommendedCandidateId!r}" + report=tuning_report, + normalized=artifact_model_to_ref(handoff.normalized), ) - preprocessed_assay = preprocessed_by_assay[assay] - normalized = preprocessed_assay.normalized - if normalized is None or preprocessed_assay.cellSelection is None: + if selected.parameters.reductionMethod != "pca": raise ValueError( - f"Assay {assay!r} lacks normalization or cell selection" + "RNA finalization requires the selected PCA representation" ) - native_selection = preprocessed_assay.cellSelection.model_dump(mode="json") - promoted = agent.promote( - store, - report=assay_report, - normalized=artifact_model_to_ref(normalized), - identity_feature_limit=request_record.config.maxIdentityFeatures, - ) - reduction_key = promoted.parameters.reductionMethod - reduction_record = promoted.artifacts[reduction_key] - reduction_ref = artifact_model_to_ref( - ArtifactReferenceModel.model_validate(reduction_record.model_dump()) - ) + selected_artifacts = { + name: ArtifactReferenceModel.model_validate( + value.model_dump(mode="json") + ) + for name, value in selected.artifacts.items() + } + required = {"pca", "connectivityMap", "clusters", "markerTable"} + if not required.issubset(selected_artifacts): + raise ValueError( + f"Final candidate lacks required artifacts: {sorted(required - selected_artifacts.keys())}" + ) + graph = selected_artifacts["connectivityMap"] + clusters = selected_artifacts["clusters"] + markers = selected_artifacts["markerTable"] + if tuning_report.finalClusterArtifact is None or artifact_model_to_ref( + clusters + ) != artifact_model_to_ref(tuning_report.finalClusterArtifact): + raise ValueError("Finalization changed the selected cluster artifact") + cells_ref = artifact_model_to_ref(cells) + graph_ref = artifact_model_to_ref(graph) + if graph_cell_selection(store.zw, graph_ref) != cells_ref: + raise ValueError( + "Selected graph does not contain the full-cohort selection" + ) + for label, ref in (("clusters", clusters), ("markers", markers)): + status = store.inspect_artifact(artifact_model_to_ref(ref)) + if not status.complete or ref.assay != assay_name: + raise ValueError( + f"Final {label} are incomplete or belong to another assay" + ) + inputs = status.inputs or {} + if inputs.get("cell_selection") != cells_ref.to_dict(): + raise ValueError(f"Final {label} use a different cell selection") + parent_key, parent = ( + ("graph", graph) if label == "clusters" else ("clusters", clusters) + ) + if inputs.get(parent_key) != artifact_model_to_ref(parent).to_dict(): + raise ValueError( + f"Final {label} do not match the selected {parent_key}" + ) + for ref in selected_artifacts.values(): + store.load_artifact(artifact_model_to_ref(ref)) initialization_ref = store.build_embedding_initialization( - reduction_ref, - n_centroids=min(1000, preprocessed_assay.nCells), + artifact_model_to_ref(selected_artifacts["pca"]), + n_centroids=min(1000, handoff.nCells), rand_state=4466, invalidate_cache=False, ) - graph_record = promoted.artifacts["connectivityMap"] - graph_ref = artifact_model_to_ref( - ArtifactReferenceModel.model_validate(graph_record.model_dump()) - ) umap_ref = store.run_umap( graph_ref, initialization_ref, @@ -315,519 +162,98 @@ def finalize_native_analyses( random_seed=4444, invalidate_cache=False, ) - promoted_artifacts = { - name: ArtifactReferenceModel.model_validate(value.model_dump()) - for name, value in promoted.artifacts.items() - } - umap_model = ArtifactReferenceModel.from_artifact_ref(umap_ref) - initialization_model = ArtifactReferenceModel.from_artifact_ref( + initialization = ArtifactReferenceModel.from_artifact_ref( initialization_ref ) - native_umaps[assay] = (initialization_model, umap_model) + umap = ArtifactReferenceModel.from_artifact_ref(umap_ref) operations.extend( [ - { - "operation": "promote_parameter_candidate", - "assay": assay, - "candidate": promoted.parameters.model_dump(mode="json"), - "normalized": normalized.model_dump(mode="json"), - "cellSelection": native_selection, - "identityFeatureLimit": ( - request_record.config.maxIdentityFeatures - ), - "artifacts": { - name: value.model_dump(mode="json") - for name, value in promoted_artifacts.items() - }, - }, { "operation": "build_embedding_initialization", - "assay": assay, - "reduction": ArtifactReferenceModel.model_validate( - reduction_record.model_dump() - ).model_dump(mode="json"), - "cellSelection": native_selection, - "nCentroids": min(1000, preprocessed_assay.nCells), - "randomSeed": 4466, - "invalidateCache": False, - "artifact": initialization_model.model_dump(mode="json"), - }, - { - "operation": "run_umap", - "assay": assay, - "graph": promoted_artifacts["connectivityMap"].model_dump( - mode="json" - ), - "initialization": initialization_model.model_dump(mode="json"), - "cellSelection": native_selection, - "parallel": False, - "randomSeed": 4444, - "invalidateCache": False, - "artifact": umap_model.model_dump(mode="json"), + "artifact": initialization.model_dump(mode="json"), }, + {"operation": "run_umap", "artifact": umap.model_dump(mode="json")}, ] ) - native_handoffs.append( - NativeAnalysisHandoff( - assay=assay, - reductionMethod=cast( - ReductionMethod, promoted.parameters.reductionMethod - ), - featureSelection=preprocessed_assay.graphFeatures, - markerFeatures=preprocessed_assay.markerFeatures, - normalized=normalized, - reduction=promoted_artifacts[reduction_key], - batchCorrection=promoted_artifacts.get("harmony"), - annIndex=promoted_artifacts["annIndex"], - embeddingInitialization=initialization_model, - neighbors=promoted_artifacts["neighbors"], - graph=promoted_artifacts["connectivityMap"], - clusters=promoted_artifacts["clusters"], - umap=umap_model, - ) - ) - artifacts.update( - { - f"{assay}_{name}": value - for name, value in { - **promoted_artifacts, - "embeddingInitialization": initialization_model, - "umap": umap_model, - }.items() - } - ) - actions.extend([f"promote_native:{assay}", f"run_native_umap:{assay}"]) - logger.info(f"Finalized native UMAP and clusters for assay {assay!r}") - return native_handoffs, native_umaps - - def finalize_selected_graph( - self, - store: DataStore, - plan: AutomatedPreprocessingPlan, - tuning_report: ParameterTuningReport, - native_handoffs: Sequence[NativeAnalysisHandoff], - native_umaps: Mapping[ - str, - tuple[ArtifactReferenceModel, ArtifactReferenceModel], - ], - actions: list[str], - operations: list[dict[str, Any]], - ) -> tuple[ - Literal["native", "snn", "wnn"], - ArtifactReferenceModel, - ArtifactReferenceModel, - ArtifactReferenceModel, - ]: - if tuning_report.cellSelection is None: - raise ValueError("Final graph selection lacks an exact cell selection") - final_cell_selection = tuning_report.cellSelection.model_dump(mode="json") - if tuning_report.recommendedIntegrationId is not None: - selected_integration = next( + doublet_scores = [ value - for value in tuning_report.integrationEvaluations - if value.integrationId == tuning_report.recommendedIntegrationId - ) - if selected_integration.graphArtifact is None: - raise ValueError("Selected integration lacks its graph artifact") - final_graph = ArtifactReferenceModel.model_validate( - selected_integration.graphArtifact.model_dump() - ) - graph_method = selected_integration.method - logger.info( - f"Finalizing selected {graph_method.upper()} graph " - f"{selected_integration.integrationId!r}" - ) - graph_ref = artifact_model_to_ref(final_graph) - primary_native = next( - value for value in native_handoffs if value.assay == plan.primaryAssay + for name, value in sorted(selected_artifacts.items()) + if name.startswith("doubletScore:") + ] + doublet_selections = [ + value + for name, value in sorted(selected_artifacts.items()) + if name.startswith("doubletCellSelection:") + ] + if len(doublet_scores) != len(doublet_selections): + raise ValueError("Advisory doublet scores lack exact selection lineage") + limitations = list( + dict.fromkeys([*plan.limitations, *tuning_report.limitations]) ) - if primary_native.embeddingInitialization is None: + doublet_limitations = [ + warning + for warning in selected.warnings + if "doublet" in warning.lower() + or "physical capture identity" in warning.lower() + ] + if not doublet_scores and not any( + warning.startswith("Advisory doublet scoring was not run for assay ") + for warning in doublet_limitations + ): raise ValueError( - "Primary native analysis lacks embedding initialization" + "Selected cluster evidence lacks advisory doublet scores" ) - final_initialization = primary_native.embeddingInitialization - umap_ref = store.run_umap( - graph_ref, - artifact_model_to_ref(final_initialization), - parallel=False, - random_seed=4444, - invalidate_cache=False, - ) - final_umap = ArtifactReferenceModel.from_artifact_ref(umap_ref) - actions.append(f"run_final_umap:{graph_method}") - operations.append( - { - "operation": "run_umap", - "graphMethod": graph_method, - "graph": final_graph.model_dump(mode="json"), - "initialization": final_initialization.model_dump(mode="json"), - "cellSelection": final_cell_selection, - "parallel": False, - "randomSeed": 4444, - "invalidateCache": False, - "artifact": final_umap.model_dump(mode="json"), - } - ) - return graph_method, final_graph, final_initialization, final_umap - graph_assay = tuning_report.graphAssay - if graph_assay is None: - raise ValueError("Native final selection lacks graphAssay") - native = next(value for value in native_handoffs if value.assay == graph_assay) - if native.graph is None: - raise ValueError("Native final selection lacks graph artifact") - final_initialization, final_umap = native_umaps[graph_assay] - logger.info( - f"Reusing native graph and UMAP from assay {graph_assay!r} as final" - ) - return "native", native.graph, final_initialization, final_umap - - def biological_interpretation_stage( - self, - store: DataStore, - workflow: AgentWorkflowRun, - request_record: OrchestrationRequestRecord, - parents: Sequence[WorkflowStageLink], - enrichment: DataEnrichmentReport, - experimental: ExperimentalContextResult, - tuning_report: ParameterTuningReport, - final_analysis: FinalAnalysisHandoff, - enrichment_reference: AgentReportReference, - experimental_reference: AgentReportReference, - tuning_reference: AgentReportReference, - answers: Mapping[str, Any], - *, - resume_record: OrchestrationResumeRecord | None = None, - ) -> WorkflowStageAttempt: - prefix = journal._ensure_orchestration_store(store) - existing = journal._validated_done_outcome( - store, - prefix, - workflow.workflowRunId, - "biological_interpretation", - request_record, - parents, - ) - if existing is not None: - logger.info( - f"Workflow {workflow.workflowRunId}: reusing Biological " - "Interpretation report" - ) - journal.load_stage_report(store, existing, BiologicalInterpretationReport) - return existing - requested_coefficient = answers.get("primaryCoefficient") - if not isinstance(requested_coefficient, str) or not requested_coefficient: - directed = request_record.request.experimentalDirections.get( - "primaryCoefficient" - ) - requested_coefficient = directed if isinstance(directed, str) else None - coefficients = list(experimental.decision.coefficientsOfInterest) - logger.info( - f"Workflow {workflow.workflowRunId}: Biological Interpretation has " - f"{len(coefficients)} validated coefficient option(s)" - ) - if final_analysis.cellSelection is None: - raise ValueError("Final analysis lacks an exact cell selection") - started = journal._start_attempt( - store.zw, - prefix, - workflow.workflowRunId, - "biological_interpretation", - request_record, - parents, - inputs={ - "studyContext": request_record.request.studyContext, - "studyContextSummary": enrichment.studyContextSummary.model_dump( - mode="json" - ), - "experimentalContextReport": experimental_reference.model_dump( - mode="json" - ), - "finalAnalysis": final_analysis.model_dump(mode="json"), - "cellSelection": final_analysis.cellSelection.model_dump(mode="json"), - "primaryCoefficient": requested_coefficient, - "biologicalInterpretation": answers.get("biologicalInterpretation"), - }, - resume_record=resume_record, - ) - if requested_coefficient is None and len(coefficients) > 1: - logger.info( - f"Workflow {workflow.workflowRunId}: Biological Interpretation " - "requires a primary coefficient selection" + limitations.extend(doublet_limitations) + artifacts.update(selected_artifacts) + artifacts.update( + cellSelection=cells, + normalized=handoff.normalized, + graph=graph, + clusters=clusters, + markers=markers, + markerFeatures=handoff.markerFeatures, + embeddingInitialization=initialization, + umap=umap, + ) + final = FinalAnalysisHandoff( + workflowRunId=workflow.workflowRunId, + primaryAssay=assay_name, + markerAssay=assay_name, + cellSelection=cells, + graph=graph, + clusters=clusters, + embeddingInitialization=initialization, + umap=umap, + markerFeatures=handoff.markerFeatures, + markers=markers, + doubletScores=doublet_scores, + doubletScoreSelections=doublet_selections, + limitations=list(dict.fromkeys(limitations)), ) outcome = journal._complete_attempt( started, - status="needsInput", - artifacts={"cellSelection": final_analysis.cellSelection}, - needs_input=WorkflowNeedsInput( - questions=[ - WorkflowQuestion( - questionId="primaryCoefficient", - question=( - "Which validated coefficient should constrain " - "treatment observations?" - ), - options=coefficients, - evidenceIds=list(experimental.decision.evidenceIds), - ) - ] - ), - notes=[ - "Cluster identities can be interpreted after one treatment " - "coefficient is selected" - ], - ) - journal._save_outcome(store.zw, prefix, outcome) - return outcome - try: - experimental_handoff: ExperimentalBiologyHandoff | None = None - if coefficients: - experimental_handoff = experimental.to_biological_handoff( - requested_coefficient - ).model_copy(update={"cellSelection": final_analysis.cellSelection}) - tuning_handoff = tuning_report.to_biological_handoff() - if tuning_handoff.cellSelection != final_analysis.cellSelection: - raise ValueError( - "Biological handoffs do not share the final cell selection" - ) - marker_policy = next( - ( - value - for value in enrichment.policies - if value.assay == final_analysis.markerAssay - ), - None, - ) - summary = enrichment.studyContextSummary - biological_context = BiologicalContext( - organism=( - marker_policy.organismName - if marker_policy is not None - and marker_policy.organismName != "unknown" - else "" - ), - studyContext=request_record.request.studyContext, - tissue=", ".join(summary.tissueReferences), - cellTypeReferences=list(summary.cellTypeReferences), - experimentalDetails=[ - *summary.experimentalReferences, - *final_analysis.limitations, - ], - treatmentQuestion=str( - request_record.request.experimentalDirections.get( - "treatmentQuestion", "" - ) - ), - ) - biological_answer = answers.get("biologicalInterpretation") - if isinstance(biological_answer, str) and biological_answer.strip(): - biological_context = biological_context.model_copy( - update={ - "experimentalDetails": [ - *biological_context.experimentalDetails, - biological_answer.strip(), - ] - } - ) - elif isinstance(biological_answer, Mapping): - biological_context = biological_context.model_copy( - update={ - "experimentalDetails": [ - *biological_context.experimentalDetails, - json.dumps(dict(biological_answer), sort_keys=True), - ] - } - ) - if ( - final_analysis.clusters is None - or final_analysis.markers is None - or final_analysis.markerFeatures is None - ): - raise ValueError("Final analysis lacks clusters or marker artifacts") - recovered = journal._recover_persisted_stage_report( - store, - started, - agent_name="biological_interpretation", - expected_type=BiologicalInterpretationReport, - ) - if recovered is not None: - recovered_report, reference = recovered - report = cast(BiologicalInterpretationReport, recovered_report) - recovery_actions = [ - "recover_persisted_biological_interpretation_report" - ] - else: - recovery_actions = [] - logger.info( - f"Workflow {workflow.workflowRunId}: invoking Biological " - "Interpretation" - ) - agent = BiologicalInterpretationAgent( - self.model, - config=request_record.config.agentRunConfig, - ) - report = agent.run( - store, - cluster=artifact_model_to_ref(final_analysis.clusters), - biological_context=biological_context, - from_assay=final_analysis.markerAssay, - graph_assay=tuning_handoff.graphAssay, - marker_assay_type=( - marker_policy.assayModality - if marker_policy is not None - else None - ), - tuning_handoff=tuning_handoff, - experimental_handoff=experimental_handoff, - marker=artifact_model_to_ref(final_analysis.markers), - marker_features=artifact_model_to_ref( - final_analysis.markerFeatures - ), - allow_marker_search=False, - ) - if marker_policy is not None and marker_policy.assayModality == "ATAC": - atac_limitation = ( - "ATAC peak markers are descriptive, so all cell identities " - "remain low-confidence hypotheses." - ) - report = report.model_copy( - update={ - "clusterInterpretations": [ - value.model_copy( - update={ - "identityIsHypothesis": True, - "confidence": "low", - } - ) - for value in report.clusterInterpretations - ], - "limitations": list( - dict.fromkeys([*report.limitations, atac_limitation]) - ), - } - ) - parent_reports = [ - journal._report_link(enrichment_reference), - journal._report_link(experimental_reference), - journal._report_link(tuning_reference), - ] - saved_report, reference = journal._save_stage_report( - store, - started, - report, - invocation=AgentInvocation( - agentName="biological_interpretation", - parentReports=parent_reports, - inputs={ - "biologicalContext": biological_context.model_dump( - mode="json" - ), - "cellSelection": ( - final_analysis.cellSelection.model_dump(mode="json") - ), - "markerAssay": final_analysis.markerAssay, - "graphAssay": tuning_handoff.graphAssay, - "markerAssayType": ( - marker_policy.assayModality - if marker_policy is not None - else None - ), - "allowMarkerSearch": False, - "studyContextSummary": summary.model_dump(mode="json"), - "experimentalContextReport": ( - experimental_reference.model_dump(mode="json") - ), - }, - artifacts={ - "cellSelection": final_analysis.cellSelection, - "clusters": final_analysis.clusters, - "markers": final_analysis.markers, - "markerFeatures": final_analysis.markerFeatures, - }, - runConfig=agent.config, - experimentalBiologyHandoff=experimental_handoff, - tuningBiologyHandoff=tuning_handoff, - ), - expected_type=BiologicalInterpretationReport, - ) - report = cast(BiologicalInterpretationReport, saved_report) - logger.info( - f"Workflow {workflow.workflowRunId}: Biological Interpretation " - f"returned status={report.status!r}, clusters=" - f"{len(report.clusterInterpretations)}, treatments=" - f"{len(report.treatmentObservations)}" + status="done", + artifacts=artifacts, + outputs={ + "finalAnalysis": final.model_dump(mode="json"), + "operations": operations, + }, + actions=["reuse_validated_full_cohort", "run_final_umap"], + notes=final.limitations, ) - if report.status == "needsInput": - needs_input = report.needsInput - assert needs_input is not None - outcome = journal._complete_attempt( - started, - status="needsInput", - report_references=[reference], - artifacts={ - "cellSelection": final_analysis.cellSelection, - "clusters": final_analysis.clusters, - "markers": final_analysis.markers, - }, - actions=recovery_actions, - needs_input=WorkflowNeedsInput( - questions=[ - WorkflowQuestion( - questionId="biologicalInterpretation", - question=needs_input.question, - options=list(needs_input.requiredInputs), - evidenceIds=list(needs_input.evidenceIds), - ) - ] - ), - notes=report.limitations, - ) - elif report.status == "failed": - outcome = journal._complete_attempt( - started, - status="failed", - report_references=[reference], - artifacts={ - "cellSelection": final_analysis.cellSelection, - "clusters": final_analysis.clusters, - "markers": final_analysis.markers, - }, - actions=recovery_actions, - error="; ".join(report.limitations) - or "Biological Interpretation failed", - ) - else: - outcome = journal._complete_attempt( - started, - status="done", - report_references=[reference], - artifacts={ - "cellSelection": final_analysis.cellSelection, - "clusters": final_analysis.clusters, - "markers": final_analysis.markers, - }, - actions=recovery_actions, - outputs={ - "clusterCount": len(report.clusterInterpretations), - "treatmentObservationCount": len(report.treatmentObservations), - }, - notes=report.limitations, - ) journal._save_outcome(store.zw, prefix, outcome) logger.info( - f"Workflow {workflow.workflowRunId}: Biological Interpretation " - f"outcome status={outcome.status!r}" + f"Final RNA analysis: {selected.metrics.nClusters} populations; descriptive markers saved" ) - if outcome.status == "failed": - journal.finalize_failed( - store, workflow, outcome.error or "biology failed" - ) - return outcome + return outcome, final except Exception as exc: - return journal.finish_exception( + outcome = journal.finish_exception( store, prefix, workflow, started, exc, - artifacts={"cellSelection": final_analysis.cellSelection}, + artifacts=artifacts, + outputs={"operations": operations}, ) + return outcome, FinalAnalysisHandoff.get_blank() diff --git a/scarf/agent/orchestrator/journal.py b/scarf/agent/orchestrator/journal.py index 2ff3ca16..598714db 100644 --- a/scarf/agent/orchestrator/journal.py +++ b/scarf/agent/orchestrator/journal.py @@ -1,10 +1,13 @@ -"""Immutable orchestration journal and stage lifecycle operations.""" +"""One immutable RNA workflow history containing stage evidence and decisions.""" import hashlib +import json import re import time import uuid -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from pathlib import Path from typing import Any, Literal, cast import zarr @@ -14,31 +17,17 @@ from ...datastore.datastore import DataStore from ...utils.logging import logger from .. import record_io -from ..persistence import ( - AgentInvocation, - AgentName, - AgentReport, - AgentReportLink, - AgentReportReference, - AgentWorkflowRun, - finalize_agent_workflow, - list_agent_reports, - load_agent_record, - load_agent_report, - load_agent_workflow, - save_agent_report, -) -from ..types import AgentDataModel, ArtifactReferenceModel +from ..experimental_context.study import StudyContract +from ..types import AgentDataModel, AgentRunInfo, ArtifactReferenceModel from .models import ( - _ORCHESTRATION_FORMAT, - _ORCHESTRATION_VERSION, _STAGE_ORDER, - AutomatedPreprocessingPlan, + AutomatedWorkflowConfig, AutomatedWorkflowResult, - AutomatedWorkflowStatus, FinalAnalysisHandoff, OrchestrationRequestRecord, OrchestrationResumeRecord, + StageEvidenceReference, + WorkflowIdentity, WorkflowNeedsInput, WorkflowStageAttempt, WorkflowStageLink, @@ -46,6 +35,12 @@ artifact_model_to_ref, ) +_INCOMPATIBLE = ( + "Unsupported saved agent workflow. Start a new RNA workflow with this release; " + "older requests cannot be resumed or regenerated. Existing analysis artifacts " + "remain accessible through Scarf's artifact APIs." +) + def _sha256_model(value: AgentDataModel) -> str: return hashlib.sha256( @@ -53,6 +48,109 @@ def _sha256_model(value: AgentDataModel) -> str: ).hexdigest() +def model_attempt_callback( + store: Any, + prefix: str, + workflow_run_id: str, + key: str, + inputs: Mapping[str, Any], +) -> Callable[[AgentRunInfo], None]: + """Save provider attempts beside their exact owning scientific evidence.""" + identity = hashlib.sha256(record_io.canonical_json_bytes(dict(inputs))).hexdigest() + + def save(run_info: AgentRunInfo) -> None: + save_checkpoint( + store, + prefix, + workflow_run_id, + f"{key}/model_attempts/{uuid.uuid4().hex}", + inputs={"evidenceSha256": identity}, + outputs={ + "runInfo": run_info.model_dump(mode="json"), + "recordedAtNs": time.time_ns(), + }, + ) + + return save + + +@contextmanager +def diagnostic_attempt( + store: Any, + prefix: str, + workflow_run_id: str, + inputs: Mapping[str, Any], +) -> Iterator[dict[str, dict[str, int]]]: + """Record observed diagnostic calls, including interrupted and repeated work.""" + from ..parameter_tuning.execution import diagnostic_work + + key = f"parameter_tuning/diagnostic_attempts/{uuid.uuid4().hex}" + identity = dict(inputs) + started = time.time_ns() + save_checkpoint( + store, + prefix, + workflow_run_id, + key + "/started", + identity, + {"recordedAtNs": started}, + ) + with diagnostic_work() as counts: + status = "completed" + error = None + failure: BaseException | None = None + try: + yield counts + except BaseException as exc: + failure = exc + status = "failed" if isinstance(exc, Exception) else "interrupted" + error = f"{type(exc).__name__}: {exc}" + raise + finally: + try: + logger.info( + f"Diagnostic work: {sum(row['completed'] for row in counts.values())}/" + f"{sum(row['attempted'] for row in counts.values())} " + "observed calls completed/attempted, including " + f"{counts.get('core.doubletDetection', {}).get('attempted', 0)} " + "doublet-scoring calls; " + f"{sum(row['cacheHits'] for row in counts.values())} metric cache hits, " + f"{sum(row['restored'] for row in counts.values())} saved-evidence restores. " + "Composite diagnostics include their underlying calls; these counts " + "do not measure internal rebuilds or simulated-doublet writes." + ) + save_checkpoint( + store, + prefix, + workflow_run_id, + key + "/finished", + identity, + { + "recordedAtNs": time.time_ns(), + "elapsedSeconds": (time.time_ns() - started) / 1e9, + "status": status, + "error": error, + "operations": counts, + "interpretation": ( + "Counts are observed agent calls and explicit reuse events. " + "Composite diagnostics include underlying calls; elapsed time " + "covers the tuning invocation, including model waiting. " + "A core call may reuse artifacts or still repeat internal work; " + "internal numerical rebuilds and simulated-doublet writes are " + "not observable from returned references. Historical attempts " + "without these records have unknown operation counts." + ), + }, + ) + except BaseException as persistence_error: + if failure is None: + raise + failure.add_note( + "Saving diagnostic work also failed: " + f"{type(persistence_error).__name__}: {persistence_error}" + ) + + def _record_checksum(value: AgentDataModel) -> str: return hashlib.sha256( record_io.canonical_json_bytes( @@ -76,53 +174,100 @@ def _list_keys(group: zarr.Group, prefix: str) -> list[str]: return record_io.list_keys(group, prefix) -def _orchestration_prefix(store: DataStore) -> str: - root_path = str(getattr(store.zw, "path", "")).strip("/") - return record_io.join_key(root_path, "agents", "orchestrations") +def _checkpoint_key(prefix: str, workflow_run_id: str, key: str) -> str: + parts = key.split("/") + if not parts or any( + re.fullmatch(r"[A-Za-z0-9_.:-]+", part) is None or part in {".", ".."} + for part in parts + ): + raise ValueError("Checkpoint keys must contain safe, non-empty path components") + if re.fullmatch(r"[a-z0-9][a-z0-9_-]{0,127}", workflow_run_id) is None: + raise ValueError("Invalid workflow identifier") + return record_io.join_key(prefix, workflow_run_id, "checkpoints", key + ".json") -def _ensure_orchestration_store(store: DataStore) -> str: - if "agents" not in store.zw: - raise RuntimeError("Create the agent workflow before orchestration records") - agents = store.zw["agents"] - if not isinstance(agents, zarr.Group): - raise ValueError("The agents namespace must be a Zarr group") - if "orchestrations" not in agents: - candidate_prefix = _orchestration_prefix(store) - if _list_keys(store.zw, candidate_prefix): - raise ValueError( - "A non-Zarr object already occupies the orchestrations namespace" - ) - agents.create_group( - "orchestrations", - attributes={ - "format": _ORCHESTRATION_FORMAT, - "format_version": _ORCHESTRATION_VERSION, - }, +def read_checkpoint( + store: DataStore, + prefix: str, + workflow_run_id: str, + key: str, +) -> dict[str, Any] | None: + """Read an exact journal checkpoint with its validated inputs and outputs.""" + raw = record_io.read_key(store.zw, _checkpoint_key(prefix, workflow_run_id, key)) + if raw is None: + return None + value = json.loads(raw) + if not isinstance(value, dict) or set(value) != { + "inputs", + "outputs", + "contentSha256", + }: + raise ValueError( + "Unsupported RNA checkpoint contract; start a new workflow. Existing analysis artifacts remain accessible." ) - logger.info("Initialized the automated workflow orchestration journal") - node = agents["orchestrations"] - if not isinstance(node, zarr.Group): - raise ValueError("The orchestrations namespace must be a Zarr group") - if ( - node.attrs.get("format") != _ORCHESTRATION_FORMAT - or node.attrs.get("format_version") != _ORCHESTRATION_VERSION - ): - raise ValueError("Unrecognized orchestration persistence format") - return _orchestration_prefix(store) + payload = {"inputs": value["inputs"], "outputs": value["outputs"]} + digest = hashlib.sha256(record_io.canonical_json_bytes(payload)).hexdigest() + if value["contentSha256"] != digest: + raise ValueError("RNA checkpoint checksum does not match its contents") + if not isinstance(value["inputs"], dict) or not isinstance(value["outputs"], dict): + raise ValueError("RNA checkpoint inputs and outputs must be mappings") + return value -def _request_key(prefix: str, workflow_run_id: str) -> str: - return record_io.join_key(prefix, workflow_run_id, "request.json") +def load_checkpoint( + store: DataStore, + prefix: str, + workflow_run_id: str, + key: str, + inputs: Mapping[str, Any] | None, +) -> dict[str, Any] | None: + """Read committed outputs only when the exact scientific inputs agree.""" + value = read_checkpoint(store, prefix, workflow_run_id, key) + if value is None: + return None + if inputs is not None and record_io.canonical_json_bytes( + inputs + ) != record_io.canonical_json_bytes(value["inputs"]): + raise ValueError(f"Checkpoint {key!r} has different scientific inputs") + return cast(dict[str, Any], value["outputs"]) -def _resume_key(prefix: str, workflow_run_id: str, resume_id: str) -> str: - return record_io.join_key( - prefix, - workflow_run_id, - "resumes", - f"{resume_id}.json", - ) +def save_checkpoint( + store: DataStore, + prefix: str, + workflow_run_id: str, + key: str, + inputs: Mapping[str, Any], + outputs: Mapping[str, Any], +) -> dict[str, Any]: + """Commit evidence or admission before its dependent work; exact replay is idempotent.""" + payload = {"inputs": dict(inputs), "outputs": dict(outputs)} + value = { + **payload, + "contentSha256": hashlib.sha256( + record_io.canonical_json_bytes(payload) + ).hexdigest(), + } + path = _checkpoint_key(prefix, workflow_run_id, key) + try: + _write_key_once(store.zw, path, record_io.display_json_bytes(value)) + except FileExistsError: + existing = load_checkpoint(store, prefix, workflow_run_id, key, inputs) + if existing != outputs: + raise ValueError( + f"Checkpoint {key!r} already contains a different outcome" + ) from None + return existing + return dict(outputs) + + +def _orchestration_prefix(store: DataStore) -> str: + root_path = str(getattr(store.zw, "path", "")).strip("/") + return record_io.join_key(root_path, "agents", "orchestrations") + + +def _request_key(prefix: str, workflow_run_id: str) -> str: + return record_io.join_key(prefix, workflow_run_id, "request.json") def _stage_prefix( @@ -147,10 +292,6 @@ def _stage_key( ) -def _result_key(prefix: str, workflow_run_id: str) -> str: - return record_io.join_key(prefix, workflow_run_id, "result.json") - - def _read_model( group: zarr.Group, key: str, @@ -162,7 +303,14 @@ def _read_model( try: return model_type.model_validate_json(raw) except ValueError as exc: - raise ValueError(f"Malformed orchestration record {key!r}") from exc + hint = ( + "; recreate the request with the current AutomatedWorkflowConfig " + "and start a new workflow. Older saved request/config shapes are " + "unsupported and are not migrated" + if model_type is OrchestrationRequestRecord + else "" + ) + raise ValueError(f"Malformed orchestration record {key!r}{hint}") from exc def _write_model_once(group: zarr.Group, key: str, value: AgentDataModel) -> None: @@ -180,8 +328,8 @@ def _stage_checksum(attempt: WorkflowStageAttempt) -> str: def _complete_attempt( started: WorkflowStageAttempt, *, - status: Literal["done", "needsInput", "failed"], - report_references: Sequence[AgentReportReference] = (), + status: Literal["done", "needsInput", "abstained", "failed"], + report_references: Sequence[StageEvidenceReference] = (), artifacts: Mapping[str, ArtifactReferenceModel] | None = None, outputs: Mapping[str, Any] | None = None, actions: Sequence[str] = (), @@ -218,15 +366,12 @@ def _start_attempt( ) -> WorkflowStageAttempt: attempt_inputs = dict(inputs or {}) if resume_record is not None: - attempt_inputs["resumeLineage"] = { - "resumeId": resume_record.resumeId, - "answeredAttempt": ( - resume_record.answeredAttempt.model_dump(mode="json") - if resume_record.answeredAttempt is not None - else None - ), - "questionIds": list(resume_record.questionIds), - } + attempt_inputs["resumeAnswers"] = dict(resume_record.answers) + attempt_inputs["answeredAttempt"] = ( + resume_record.answeredAttempt.model_dump(mode="json") + if resume_record.answeredAttempt is not None + else None + ) attempt = WorkflowStageAttempt( workflowRunId=workflow_run_id, stage=stage, @@ -250,10 +395,7 @@ def _start_attempt( ), attempt, ) - logger.info( - f"Workflow {workflow_run_id}: started stage={stage!r} " - f"attempt={attempt.attemptId}" - ) + logger.info(f"{stage.replace('_', ' ').capitalize()}: started") return attempt @@ -278,30 +420,17 @@ def _save_outcome( if outcome.completedAtNs is not None else 0.0 ) - details = ( - f"reports={len(outcome.reportReferences)}, " - f"artifacts={len(outcome.artifacts)}, actions={len(outcome.actions)}" - ) + label = outcome.stage.replace("_", " ").capitalize() if outcome.status == "failed": - error_kind = (outcome.error or "unknown error").partition(":")[0] - logger.error( - f"Workflow {outcome.workflowRunId}: stage={outcome.stage!r} " - f"failed ({error_kind}; {details}; {elapsed_seconds:.1f}s)" - ) - elif outcome.status == "needsInput": - question_count = ( - len(outcome.needsInput.questions) if outcome.needsInput is not None else 0 - ) - logger.info( - f"Workflow {outcome.workflowRunId}: stage={outcome.stage!r} paused " - f"for {question_count} input question(s) ({details}; " - f"{elapsed_seconds:.1f}s)" - ) + logger.error(f"{label}: {outcome.error} ({elapsed_seconds:.1f}s)") + elif outcome.status in {"needsInput", "abstained"}: + reasons = "; ".join(outcome.notes) + if outcome.needsInput is not None: + reasons = "; ".join(q.question for q in outcome.needsInput.questions) + logger.warning(f"{label}: {outcome.status}: {reasons}") else: - logger.info( - f"Workflow {outcome.workflowRunId}: completed stage={outcome.stage!r} " - f"({details}; {elapsed_seconds:.1f}s)" - ) + logger.info(f"{label}: completed ({elapsed_seconds:.1f}s)") + logger.debug(f"Workflow {outcome.workflowRunId}, attempt {outcome.attemptId}") def _stage_outcomes( @@ -423,6 +552,36 @@ def _resume_answer_errors( for question_id in sorted(expected_ids & supplied_ids): question = questions[question_id] answer = answers[question_id] + if question.decisionId is not None: + if not isinstance(answer, Mapping): + errors.append( + f"Resume answer for {question_id!r} must contain decisionId, " + "optionId, and rationale" + ) + continue + if set(answer) != {"decisionId", "optionId", "rationale"}: + errors.append( + f"Resume answer for {question_id!r} must contain exactly " + "decisionId, optionId, and rationale" + ) + continue + if answer.get("decisionId") != question.decisionId: + errors.append( + f"Resume answer for {question_id!r} does not match decision " + f"{question.decisionId!r}" + ) + option_id = answer.get("optionId") + if not isinstance(option_id, str) or option_id not in question.options: + errors.append( + f"Resume answer for {question_id!r} must select one persisted " + f"option {question.options!r}" + ) + rationale = answer.get("rationale") + if not isinstance(rationale, str) or not rationale.strip(): + errors.append( + f"Resume answer for {question_id!r} requires a non-empty rationale" + ) + continue if unsafe_context and question_id == "experimentalDirections": if _unsafe_context_resolution(answer) is None: errors.append( @@ -452,61 +611,6 @@ def _resume_answer_errors( return errors -def _validated_resume_record( - store: DataStore, - prefix: str, - workflow_run_id: str, - resume_id: str, -) -> OrchestrationResumeRecord: - record = cast( - OrchestrationResumeRecord, - _read_model( - store.zw, - _resume_key(prefix, workflow_run_id, resume_id), - OrchestrationResumeRecord, - ), - ) - if record.workflowRunId != workflow_run_id or record.resumeId != resume_id: - raise ValueError("Persisted resume record identity does not match its path") - if record.contentSha256 != _record_checksum(record): - raise ValueError("Persisted resume record checksum does not match its content") - if record.answeredAttempt is None: - if record.questionIds or record.answers: - raise ValueError( - "A resume without an answered attempt cannot contain question answers" - ) - return record - matches = [ - outcome - for outcome in _stage_outcomes( - store.zw, - prefix, - workflow_run_id, - record.answeredAttempt.stage, - ) - if outcome.attemptId == record.answeredAttempt.attemptId - and outcome.contentSha256 == record.answeredAttempt.contentSha256 - ] - if len(matches) != 1 or matches[0].status != "needsInput": - raise ValueError( - "Persisted resume record does not cite one paused stage attempt" - ) - answered_outcome = matches[0] - assert answered_outcome.needsInput is not None - expected_question_ids = [ - question.questionId for question in answered_outcome.needsInput.questions - ] - if record.questionIds != expected_question_ids: - raise ValueError("Persisted resume record question IDs are stale") - answer_errors = _resume_answer_errors(answered_outcome, record.answers) - if answer_errors: - raise ValueError( - "Persisted resume record contains invalid answers: " - + "; ".join(answer_errors) - ) - return record - - def _validated_done_outcome( store: DataStore, prefix: str, @@ -515,7 +619,7 @@ def _validated_done_outcome( request_record: OrchestrationRequestRecord, parent_attempts: Sequence[WorkflowStageLink], *, - required_status: Literal["done", "needsInput"] = "done", + required_status: Literal["done", "needsInput", "failed"] = "done", ) -> WorkflowStageAttempt | None: """Return the newest lineage-matching stage whose persisted outputs resolve.""" @@ -607,7 +711,7 @@ def _stage_outcome_resolves( for reference in outcome.reportReferences: if reference.workflowRunId != workflow_run_id: raise ValueError("Stage report belongs to a different workflow") - load_agent_report(store, reference) + read_stage_evidence(store, reference) for artifact_reference in outcome.artifacts.values(): store.load_artifact(artifact_model_to_ref(artifact_reference)) metadata_columns = outcome.outputs.get("metadataColumns", []) @@ -629,19 +733,10 @@ def _parent_link(outcome: WorkflowStageAttempt) -> WorkflowStageLink: ) -def _safe_label(value: str) -> str: - clean = re.sub(r"[^A-Za-z0-9_]+", "_", value).strip("_") - return clean or "assay" - - -def _report_link(reference: AgentReportReference) -> AgentReportLink: - return AgentReportLink.from_reference(reference) - - def _stage_execution_id(started: WorkflowStageAttempt) -> str: """Return the stable identity of one logical agent-stage invocation.""" inputs = dict(started.inputs) - inputs.pop("resumeLineage", None) + inputs.pop("answeredAttempt", None) payload = { "workflowRunId": started.workflowRunId, "stage": started.stage, @@ -656,55 +751,106 @@ def _stage_execution_id(started: WorkflowStageAttempt) -> str: return f"orchestrator_{started.stage}_{digest[:40]}" -def _stage_invocation( - started: WorkflowStageAttempt, - invocation: AgentInvocation, -) -> AgentInvocation: - execution_id = _stage_execution_id(started) - inputs = dict(invocation.inputs) - observed = inputs.get("orchestrationExecutionId") - if observed is not None and observed != execution_id: - raise ValueError("Agent invocation has a conflicting orchestration identity") - inputs["orchestrationExecutionId"] = execution_id - return invocation.model_copy(update={"inputs": inputs}) +def _ensure_orchestration_store(store: DataStore) -> str: + """Initialize only the journal namespace, never a second workflow ledger.""" + agents = store.zw.require_group("agents") + if "orchestrations" not in agents: + agents.create_group( + "orchestrations", + attributes={"format": "scarf_agent_orchestrations", "format_version": 2}, + ) + node = agents["orchestrations"] + if not isinstance(node, zarr.Group) or dict(node.attrs) != { + "format": "scarf_agent_orchestrations", + "format_version": 2, + }: + raise ValueError(_INCOMPATIBLE) + return _orchestration_prefix(store) + + +def read_request( + group: zarr.Group, prefix: str, workflow_run_id: str +) -> OrchestrationRequestRecord: + try: + value = cast( + OrchestrationRequestRecord, + _read_model( + group, _request_key(prefix, workflow_run_id), OrchestrationRequestRecord + ), + ) + except (ValueError, TypeError) as exc: + raise ValueError(_INCOMPATIBLE) from exc + if ( + value.workflowRunId != workflow_run_id + or value.requestSha256 != _sha256_model(value.request) + or value.configSha256 != _sha256_model(value.config) + or value.contentSha256 != _record_checksum(value) + ): + raise ValueError("RNA workflow request identity or checksum is invalid") + return value + + +def open_analysis_store( + target: str | Path, workflow_run_id: str, *, workspace: str | None = None +) -> DataStore: + """Open the journal's selected RNA assay read-only, before resolving any result.""" + from ...storage.schema import validate_workspace_name + + validate_workspace_name(workspace) + root = zarr.open_group(str(target), mode="r") + active = root if workspace is None else root[workspace] + if not isinstance(active, zarr.Group): + raise ValueError("Analysis workspace is not a group") + prefix = record_io.join_key(active.path, "agents", "orchestrations") + record = read_request(active, prefix, workflow_run_id) + if record.request.zarrPath is None: + raise ValueError("Saved workflow has no store path") + if ( + record.request.workspace != workspace + or Path(record.request.zarrPath).resolve() != Path(target).resolve() + ): + raise ValueError("Analysis address does not match the saved request") + if not record.request.primaryAssay: + raise ValueError("Saved analysis has no selected RNA assay") + return DataStore( + str(target), + workspace=workspace, + default_assay=record.request.primaryAssay, + zarr_mode="r", + min_features_per_cell=-1, + mito_pattern="", + ribo_pattern="", + ) + + +def _report_checkpoint(started: WorkflowStageAttempt) -> str: + return f"{started.stage}/report/{_stage_execution_id(started)}" def _recover_persisted_stage_report( store: DataStore, started: WorkflowStageAttempt, *, - agent_name: AgentName, expected_type: type[AgentDataModel], -) -> tuple[AgentDataModel, AgentReportReference] | None: - """Recover a report committed before its stage outcome was persisted.""" - execution_id = _stage_execution_id(started) - matches = [ - reference - for reference in list_agent_reports( - store, - started.workflowRunId, - agent_name=agent_name, - ) - if reference.agentRunId == execution_id - ] - if not matches: +) -> tuple[AgentDataModel, StageEvidenceReference] | None: + key = _report_checkpoint(started) + inputs = {"execution": _stage_execution_id(started)} + data = load_checkpoint( + store, _orchestration_prefix(store), started.workflowRunId, key, inputs + ) + if data is None: return None - if len(matches) != 1: - raise ValueError("A logical stage execution has multiple persisted reports") - reference = matches[0] - record = load_agent_record(store, reference) - if record.invocation.inputs.get("orchestrationExecutionId") != execution_id: - raise ValueError("Persisted stage report has a stale execution identity") - report = load_agent_report(store, reference) - if not isinstance(report, expected_type): - raise TypeError( - f"Persisted {agent_name!r} report is not {expected_type.__name__}" - ) - for artifact in record.invocation.artifacts.values(): - store.load_artifact(artifact_model_to_ref(artifact)) - logger.info( - f"Workflow {started.workflowRunId}: recovered {agent_name!r} report for " - f"stage={started.stage!r}" + if ( + set(data) != {"reportType", "report"} + or data["reportType"] != expected_type.__name__ + ): + raise ValueError("Stage evidence has a different scientific result type") + report = expected_type.model_validate(data["report"]) + reference = StageEvidenceReference( + workflowRunId=started.workflowRunId, + stage=started.stage, + key=key, + contentSha256=_sha256_model(report), ) return report, reference @@ -714,142 +860,74 @@ def _save_stage_report( started: WorkflowStageAttempt, report: AgentDataModel, *, - invocation: AgentInvocation, expected_type: type[AgentDataModel], -) -> tuple[AgentDataModel, AgentReportReference]: - """Persist a stage report under its stable logical execution identity.""" - tagged_invocation = _stage_invocation(started, invocation) - try: - reference = save_agent_report( - store, - started.workflowRunId, - cast(AgentReport, report), - invocation=tagged_invocation, - agent_run_id=_stage_execution_id(started), - ) - logger.info( - f"Workflow {started.workflowRunId}: persisted " - f"{tagged_invocation.agentName!r} report for stage={started.stage!r}" - ) - return report, reference - except FileExistsError: - recovered = _recover_persisted_stage_report( - store, - started, - agent_name=tagged_invocation.agentName, - expected_type=expected_type, - ) - if recovered is None: - raise - logger.debug( - f"Workflow {started.workflowRunId}: reused concurrently persisted " - f"{tagged_invocation.agentName!r} report" - ) - return recovered - - -def _load_terminal_result( - store: DataStore, - prefix: str, - workflow: AgentWorkflowRun, -) -> AutomatedWorkflowResult | None: - raw = record_io.read_key( - store.zw, - _result_key(prefix, workflow.workflowRunId), + attempt_owned: bool = False, +) -> tuple[AgentDataModel, StageEvidenceReference]: + report = expected_type.model_validate(report.model_dump(mode="json")) + # Revisable stages recover individual evidence checkpoints across attempts; + # their aggregate reports can change as additional evidence is completed. + key = ( + f"{started.stage}/report/attempt_{started.attemptId}" + if attempt_owned + else _report_checkpoint(started) ) - if raw is None: - return None - try: - result = AutomatedWorkflowResult.model_validate_json(raw) - except ValueError as exc: - raise ValueError("Malformed automated workflow result") from exc - if result.contentSha256 != _record_checksum(result): - raise ValueError("Automated workflow result checksum is invalid") - if result.workflowRun is None: - raise ValueError("Terminal workflow result is missing its workflow identity") - stored_workflow = result.workflowRun - identity_fields = ( - "workflowRunId", - "workspace", - "status", - "finalizedAtNs", - "finalizationMessage", - "analysisStore", - "datasetFingerprints", + data = { + "reportType": expected_type.__name__, + "report": report.model_dump(mode="json"), + } + save_checkpoint( + store, + _orchestration_prefix(store), + started.workflowRunId, + key, + {"execution": _stage_execution_id(started)}, + data, + ) + return report, StageEvidenceReference( + workflowRunId=started.workflowRunId, + stage=started.stage, + key=key, + contentSha256=_sha256_model(report), ) - if any( - getattr(stored_workflow, field) != getattr(workflow, field) - for field in identity_fields - ): - raise ValueError("Automated workflow result has stale workflow metadata") - if result.status != workflow.status: - raise ValueError("Automated workflow result has a stale terminal status") - if result.reportReferences != workflow.reports: - raise ValueError("Automated workflow result has stale report references") - return result -def _persist_terminal_result( - store: DataStore, - prefix: str, - workflow: AgentWorkflowRun, - result: AutomatedWorkflowResult, -) -> AutomatedWorkflowResult: - existing = _load_terminal_result(store, prefix, workflow) - if existing is not None: - logger.debug( - f"Workflow {workflow.workflowRunId}: reused terminal result " - f"status={workflow.status!r}" - ) - return existing - if result.contentSha256 != _record_checksum(result): - raise ValueError("Terminal result must carry its exact content checksum") - try: - _write_model_once( - store.zw, - _result_key(prefix, workflow.workflowRunId), - result, - ) - except FileExistsError: - pass - persisted = _load_terminal_result(store, prefix, workflow) - if persisted is None: - raise RuntimeError("Terminal workflow result was not persisted") - logger.info( - f"Workflow {workflow.workflowRunId}: persisted terminal result " - f"status={workflow.status!r}" +def read_stage_evidence( + store: DataStore, reference: StageEvidenceReference +) -> dict[str, Any]: + if not reference.key.startswith(reference.stage + "/report/"): + raise ValueError("Evidence checkpoint is not owned by its declared stage") + value = load_checkpoint( + store, + _orchestration_prefix(store), + reference.workflowRunId, + reference.key, + None, ) - return persisted + if value is None or set(value) != {"reportType", "report"}: + raise ValueError("Stage evidence checkpoint is missing or malformed") + report = value["report"] + if ( + not isinstance(report, dict) + or hashlib.sha256(record_io.canonical_json_bytes(report)).hexdigest() + != reference.contentSha256 + ): + raise ValueError("Stage evidence does not match its exact reference") + return report def load_stage_report( - store: DataStore, - outcome: WorkflowStageAttempt, - expected_type: type[AgentDataModel], + store: DataStore, outcome: WorkflowStageAttempt, expected_type: type[AgentDataModel] ) -> AgentDataModel: - references = list(outcome.reportReferences) - if outcome.stage == "parameter_tuning" and len(references) > 1: - execution_id = _stage_execution_id(outcome) - references = [ - reference - for reference in references - if reference.agentRunId == execution_id - ] - if len(references) != 1: - raise ValueError( - f"Stage {outcome.stage!r} must have exactly one report reference" - ) - report = load_agent_report(store, references[0]) - if not isinstance(report, expected_type): - raise TypeError( - f"Stage {outcome.stage!r} report is not {expected_type.__name__}" - ) - return report + if len(outcome.reportReferences) != 1: + raise ValueError("A scientific stage must own exactly one evidence report") + return expected_type.model_validate( + read_stage_evidence(store, outcome.reportReferences[0]) + ) def failed_stage( store: DataStore, - workflow: AgentWorkflowRun, + workflow: WorkflowIdentity, request_record: OrchestrationRequestRecord, stage: WorkflowStageName, parents: Sequence[WorkflowStageLink], @@ -858,9 +936,6 @@ def failed_stage( artifacts: Mapping[str, ArtifactReferenceModel] | None = None, resume_record: OrchestrationResumeRecord | None = None, ) -> WorkflowStageAttempt: - logger.error( - f"Workflow {workflow.workflowRunId}: stage={stage!r} failed validation" - ) prefix = _ensure_orchestration_store(store) started = _start_attempt( store.zw, @@ -872,20 +947,16 @@ def failed_stage( resume_record=resume_record, ) outcome = _complete_attempt( - started, - status="failed", - artifacts=artifacts, - error=error, + started, status="failed", artifacts=artifacts, error=error ) _save_outcome(store.zw, prefix, outcome) - finalize_failed(store, workflow, error) return outcome def finish_exception( store: DataStore, prefix: str, - workflow: AgentWorkflowRun, + workflow: WorkflowIdentity, started: WorkflowStageAttempt, exc: BaseException, *, @@ -894,98 +965,461 @@ def finish_exception( outputs: Mapping[str, Any] | None = None, notes: Sequence[str] = (), ) -> WorkflowStageAttempt: - error = f"{type(exc).__name__}: {exc}" - logger.error( - f"Workflow {workflow.workflowRunId}: stage={started.stage!r} raised " - f"{type(exc).__name__}; details were persisted in the stage outcome" - ) - execution_id = _stage_execution_id(started) - report_references = [ - reference - for reference in list_agent_reports(store, workflow.workflowRunId) - if reference.agentRunId == execution_id - or reference.agentRunId.startswith(f"{execution_id}_integration_") - ] + from ..config.agent_exec import describe_agent_error + + error = describe_agent_error(exc) + saved_outputs = dict(outputs or {}) + run_info = getattr(exc, "agent_run_info", None) + if isinstance(run_info, AgentRunInfo): + saved_outputs["runInfo"] = run_info.model_dump(mode="json") outcome = _complete_attempt( started, status="failed", - report_references=report_references, artifacts=artifacts, - outputs=outputs, + outputs=saved_outputs, actions=actions, notes=notes, error=error, ) _save_outcome(store.zw, prefix, outcome) - finalize_failed(store, workflow, error) return outcome -def finalize_failed( - store: DataStore, - workflow: AgentWorkflowRun, - message: str, -) -> None: - current = load_agent_workflow(store, workflow.workflowRunId) - if current.status == "running": - logger.warning(f"Finalizing workflow {workflow.workflowRunId} as failed") - finalize_agent_workflow( - store, - workflow.workflowRunId, - status="failed", - message=message, - ) - - -def all_report_references( - store: DataStore, - prefix: str, - workflow_run_id: str, -) -> list[AgentReportReference]: - references: dict[tuple[str, str], AgentReportReference] = {} - for stage in _STAGE_ORDER: - for outcome in _stage_outcomes(store.zw, prefix, workflow_run_id, stage): - for reference in outcome.reportReferences: - if reference.workflowRunId != workflow_run_id: - raise ValueError("Stage report belongs to a different workflow") - load_agent_report(store, reference) - references[(reference.agentName, reference.agentRunId)] = reference - return sorted( - references.values(), - key=lambda value: (value.createdAtNs, value.agentName, value.agentRunId), - ) - - def paused_or_failed_result( store: DataStore, - workflow: AgentWorkflowRun, + workflow: WorkflowIdentity, request_record: OrchestrationRequestRecord, outcome: WorkflowStageAttempt, *, - preprocessing_plan: AutomatedPreprocessingPlan | None = None, - final_analysis: FinalAnalysisHandoff | None = None, + study_contract: StudyContract | None = None, ) -> AutomatedWorkflowResult: - prefix = _ensure_orchestration_store(store) - current = load_agent_workflow(store, workflow.workflowRunId) - status: AutomatedWorkflowStatus = ( - "needsInput" if outcome.status == "needsInput" else "failed" + status: Literal["needsInput", "abstained", "failed"] = ( + "needsInput" + if outcome.status == "needsInput" + else "abstained" + if outcome.status == "abstained" + else "failed" ) - result = AutomatedWorkflowResult( + questions = ( + [q.question for q in outcome.needsInput.questions] if outcome.needsInput else [] + ) + return AutomatedWorkflowResult( status=status, currentStage=outcome.stage, zarrPath=str(store.zarr_loc), - workflowRun=current, - reportReferences=list(current.reports), - preprocessingPlan=preprocessing_plan, - finalAnalysis=final_analysis, + workspace=workflow.workspace, + workflowRunId=workflow.workflowRunId, needsInput=outcome.needsInput, - notes=[*outcome.notes, *([outcome.error] if outcome.error else [])], + notes=[*outcome.notes, *([outcome.error] if outcome.error else []), *questions], + limitations=list(study_contract.limitations) if study_contract else [], + unresolvedClaims=questions, ) - result = result.model_copy(update={"contentSha256": _record_checksum(result)}) - if status == "failed": - return _persist_terminal_result(store, prefix, current, result) - logger.info( - f"Workflow {workflow.workflowRunId}: returning needsInput at " - f"stage={outcome.stage!r}" + + +def _analysis_review_views( + store: DataStore, + prefix: str, + workflow_run_id: str, + stages: list[dict[str, Any]], + config: AutomatedWorkflowConfig, +) -> list[dict[str, Any]]: + """Derive bounded scientific views from the active exact review checkpoints.""" + views: list[dict[str, Any]] = [] + seen: set[str] = set() + for stage in stages: + if stage["stage"] != "parameter_tuning": + continue + history = stage["outputs"].get("tuningEvidence", {}).get("history", []) + for entry in history: + if "review" not in entry: + continue + key = entry.get("checkpointKey", "") + match = re.fullmatch( + r"parameter_tuning/(?:evidence_revisions/[a-f0-9]{64}/)?" + r"(sample0|sample1|full)/review([0-9]+)(?:/answer)?", + key, + ) + if match is None or match[1] != entry.get("scope"): + raise ValueError("Analysis review has an invalid checkpoint address") + if key in seen: + raise ValueError("Analysis review repeats a checkpoint") + seen.add(key) + limit = ( + config.maxFullPartitions + if match[1] == "full" + else config.maxScreeningEvaluations + ) + if int(match[2]) > limit: + raise ValueError("Analysis review exceeds its declared work allowance") + value = read_checkpoint(store, prefix, workflow_run_id, key) + if value is None: + raise ValueError("Analysis review checkpoint is missing") + if value["contentSha256"] != entry.get("checkpointSha256"): + raise ValueError("Analysis review does not match its exact checkpoint") + inputs, outputs = value["inputs"], value["outputs"] + if ( + inputs.get("scope") != entry["scope"] + or outputs.get("action") != entry["review"] + or inputs.get("imageHashes") != entry.get("imageHashes") + or inputs.get("evidenceMode") != entry.get("evidenceMode") + or inputs.get("visualInspection") != entry.get("visualInspection") + ): + raise ValueError("Analysis review evidence bindings differ") + mode, inspection = ( + inputs.get("evidenceMode"), + inputs.get("visualInspection"), + ) + if ( + mode not in {"visual", "structured"} + or inspection != ("available" if mode == "visual" else "unavailable") + or bool(inputs.get("imageHashes")) != (mode == "visual") + ): + raise ValueError( + "Analysis review has inconsistent evidence availability" + ) + candidates = inputs.get("candidates", []) + settings = inputs.get("settings", {}) + features = inputs.get("featureEvidence", {}) + candidate_ids = [item["candidateId"] for item in candidates] + if ( + not 0 < len(candidates) <= limit + or len(set(candidate_ids)) != len(candidate_ids) + or set(settings) != set(candidate_ids) + or set(features) != set(candidate_ids) + or entry["review"].get("selectedCandidateId") not in candidate_ids + ): + raise ValueError("Analysis review candidate evidence does not align") + for candidate in candidates: + setting = settings[candidate["candidateId"]] + if setting.get("parameters") != candidate.get( + "parameters" + ) or setting.get("features") != candidate.get("artifacts", {}).get( + "graphFeatures" + ): + raise ValueError( + "Analysis review settings do not match its artifacts" + ) + views.append( + { + "scope": entry["scope"], + "evidenceMode": mode, + "visualInspection": inspection, + **entry["review"], + "coverage": inputs.get("coverage", {}), + "comparisonCoverage": inputs.get("comparisonCoverage"), + "populationSupport": inputs.get("assessmentContext", {}).get( + "populationSupport", {} + ), + "harmonyGates": inputs.get("harmonyGates", {}), + "candidates": [ + { + name: item[name] + for name in ( + "candidateId", + "parameters", + "metrics", + "cellSelection", + "artifacts", + ) + } + for item in candidates + ], + "settings": { + identity: { + name: setting.get(name) + for name in ( + "hvgCount", + "ranking", + "rankingColumn", + "features", + "eligibleFeatures", + ) + } + for identity, setting in settings.items() + }, + "featureEvidence": { + identity: { + name: evidence.get(name) + for name in ( + "selectedGenes", + "eligibleGenes", + "families", + "topSelectedGenes", + ) + } + for identity, evidence in features.items() + }, + } + ) + return views + + +def _model_usage( + model_attempts: Sequence[Mapping[str, Any]], + saved_views: Sequence[Mapping[str, Any]], +) -> dict[str, Any]: + """Count measured invocations once across callbacks and copied report views.""" + invocations: dict[str, AgentRunInfo] = {} + + def collect(value: Any) -> None: + if isinstance(value, Mapping): + raw = value.get("runInfo") + if isinstance(raw, Mapping): + info = AgentRunInfo.model_validate(raw) + # Empty, code-created outcomes are not model invocations. + if ( + info.runId + or info.modelName + or info.status + or any( + ( + info.usage.requests, + info.usage.inputTokens, + info.usage.outputTokens, + info.usage.totalTokens, + info.usage.toolCalls, + ) + ) + ): + identity = info.runId or _sha256_model(info) + invocations.setdefault(identity, info) + for key, child in value.items(): + if key != "runInfo": + collect(child) + elif isinstance(value, list): + for child in value: + collect(child) + + # Invocation callbacks own measured usage; reports can repeat these records. + collect(list(model_attempts)) + collect(list(saved_views)) + infos = list(invocations.values()) + partial = sum(info.usage.availability == "partial" for info in infos) + unavailable = sum(info.usage.availability == "unavailable" for info in infos) + unspecified = sum(info.usage.availability is None for info in infos) + known = any( + info.usage.inputTokens or info.usage.outputTokens or info.usage.totalTokens + for info in infos ) - return result + return { + "invocations": len(infos), + "failedInvocations": sum(info.status == "failed" for info in infos), + "validationRetries": sum(len(info.validationRetries) for info in infos), + "durationSeconds": sum(info.durationSeconds for info in infos), + **{ + field: sum(getattr(info.usage, field) for info in infos) + for field in ( + "requests", + "toolCalls", + "inputTokens", + "outputTokens", + "totalTokens", + ) + }, + "availability": ( + "reported" + if infos and not (partial or unavailable or unspecified) + else "partial" + if known + else "unavailable" + ), + "partialUsageInvocations": partial, + "unavailableUsageInvocations": unavailable, + "unspecifiedUsageInvocations": unspecified, + } + + +def analysis_snapshot(store: DataStore, workflow_run_id: str) -> dict[str, Any]: + """Validate one journal and derive its status, final artifacts, and report view.""" + prefix = _orchestration_prefix(store) + request = read_request(store.zw, prefix, workflow_run_id) + outcomes_by_stage = { + stage: _stage_outcomes(store.zw, prefix, workflow_run_id, stage) + for stage in _STAGE_ORDER + } + usage_views: list[dict[str, Any]] = [ + outcome.outputs + for outcomes in outcomes_by_stage.values() + for outcome in outcomes + ] + stages: list[dict[str, Any]] = [] + parents: list[WorkflowStageLink] = [] + final: dict[str, Any] | None = None + status = "running" + for stage in _STAGE_ORDER: + outcomes = outcomes_by_stage[stage] + matching = [v for v in outcomes if v.parentAttempts == parents] + if not matching: + break + outcome = matching[-1] + if not _stage_outcome_resolves( + store, prefix, workflow_run_id, request, outcome + ): + raise ValueError( + f"Stage {stage!r} contains unresolved artifact or evidence references" + ) + stage_view = outcome.model_dump(mode="json") + stage_view["report"] = ( + read_stage_evidence(store, outcome.reportReferences[0]) + if outcome.reportReferences + else None + ) + stage_view["decisions"] = [] + stages.append(stage_view) + if outcome.status != "done": + status = outcome.status + break + parents = [_parent_link(outcome)] + if stage == "analysis_finalization": + resolved = FinalAnalysisHandoff.model_validate( + outcome.outputs["finalAnalysis"] + ) + if ( + resolved.workflowRunId != workflow_run_id + or resolved.primaryAssay != request.request.primaryAssay + or resolved.markerAssay != resolved.primaryAssay + ): + raise ValueError("Final analysis belongs to a different RNA workflow") + for name, kind in { + "cellSelection": "cell_selection", + "graph": "connectivity_map", + "clusters": "cluster_labels", + "umap": "embedding", + "embeddingInitialization": "embedding_initialization", + "markerFeatures": "feature_selection", + "markers": "marker_table", + }.items(): + ref = getattr(resolved, name) + if ( + ref is None + or ref.kind != kind + or ref != outcome.artifacts.get(name) + ): + raise ValueError( + f"Final analysis lacks its validated {name} artifact" + ) + if name != "cellSelection" and ref.assay != resolved.primaryAssay: + raise ValueError(f"Final {name} belongs to another assay") + inputs = outcome.inputs.get("preprocessedAssays", []) + assert resolved.cellSelection is not None + if len(inputs) != 1 or inputs[0].get( + "cellSelection" + ) != resolved.cellSelection.model_dump(mode="json"): + raise ValueError( + "Final analysis does not use the full preprocessing cohort" + ) + final = resolved.model_dump(mode="json") + status = "completed" + # Decisions belong to this same history and retain their offered evidence. + checkpoint_prefix = record_io.join_key(prefix, workflow_run_id, "checkpoints") + decisions: list[dict[str, Any]] = [] + model_attempts: list[dict[str, Any]] = [] + diagnostic_attempts: dict[str, dict[str, Any]] = {} + for path in _list_keys(store.zw, checkpoint_prefix): + if not path.endswith(".json") or not any( + part in path + for part in ( + "/decisions/", + "/model_attempts/", + "/report/", + "/diagnostic_attempts/", + ) + ): + continue + key = path[len(checkpoint_prefix) + 1 : -5] + checkpoint = read_checkpoint(store, prefix, workflow_run_id, key) + value = checkpoint["outputs"] if checkpoint is not None else None + if value is not None: + usage_views.append(value) + if value is not None and "record" in value: + decisions.append(value) + if value is not None and "/model_attempts/" in path: + run_info = AgentRunInfo.model_validate(value["runInfo"]) + model_attempts.append( + { + "checkpoint": key, + "recordedAtNs": value.get("recordedAtNs"), + "runInfo": run_info.model_dump(mode="json"), + } + ) + if value is not None and "/diagnostic_attempts/" in path: + assert checkpoint is not None + attempt, _, phase = key.rpartition("/") + if phase not in {"started", "finished"}: + raise ValueError("Unknown diagnostic attempt checkpoint") + entry = diagnostic_attempts.setdefault( + attempt, {"checkpoint": attempt, "inputs": checkpoint["inputs"]} + ) + if entry["inputs"] != checkpoint["inputs"]: + raise ValueError( + "Diagnostic attempt inputs changed between start and finish" + ) + entry[phase] = value + + if any( + "finished" in row and "started" not in row + for row in diagnostic_attempts.values() + ): + raise ValueError("Completed diagnostic attempt has no matching start record") + + def contains(value: Any, digest: str) -> bool: + if isinstance(value, Mapping): + return any(contains(v, digest) for v in value.values()) + if isinstance(value, list): + return any(contains(v, digest) for v in value) + return bool(value == digest) + + for value in decisions: + digest = value.get("checkpointSha256") + if not isinstance(digest, str): + raise ValueError("Decision checkpoint has no input identity") + owner = next( + ( + stage + for stage in stages + if stage["stage"] == value.get("stage") + and ( + contains(stage["inputs"], digest) + or contains(stage["outputs"], digest) + ) + ), + None, + ) + if owner is not None: + owner["decisions"].append(value) + reviews = _analysis_review_views( + store, prefix, workflow_run_id, stages, request.config + ) + if status == "completed": + assert final is not None + tuning = next(stage for stage in stages if stage["stage"] == "parameter_tuning") + report = tuning["report"] + full_reviews = [value for value in reviews if value["scope"] == "full"] + if ( + not full_reviews + or full_reviews[-1]["action"] != "accept" + or full_reviews[-1]["selectedCandidateId"] + != report.get("recommendedCandidateId") + or report.get("finalClusterArtifact") != final["clusters"] + ): + raise ValueError( + "Final analysis lacks its exact full-cohort acceptance evidence" + ) + return { + "runId": workflow_run_id, + "status": status, + "request": request.request.model_dump(mode="json"), + "stages": stages, + "modelAttempts": sorted( + model_attempts, key=lambda row: row["recordedAtNs"] or 0 + ), + "modelUsage": _model_usage(model_attempts, usage_views), + "diagnosticAttempts": sorted( + diagnostic_attempts.values(), + key=lambda row: row.get("started", {}).get("recordedAtNs", 0), + ), + "finalAnalysis": final, + "modelIdentity": request.modelIdentity, + "analysisReviews": reviews, + "config": request.config.model_dump(mode="json"), + } diff --git a/scarf/agent/orchestrator/main.py b/scarf/agent/orchestrator/main.py index 65a153c0..2f30f5dc 100644 --- a/scarf/agent/orchestrator/main.py +++ b/scarf/agent/orchestrator/main.py @@ -1,5 +1,6 @@ -"""Public controller for resumable automated Scarf agent workflows.""" +"""Controller for one resumable RNA analysis with checkpoint-owned state.""" +import hashlib import time import uuid from collections.abc import Mapping @@ -9,60 +10,160 @@ import zarr from ...datastore.datastore import DataStore +from ...datastore.summary import summarize_zarr_readonly from ...utils.logging import logger from .. import record_io +from ..config.agent_exec import describe_agent_error +from ..experimental_context.study import StudyContract, validate_objective_evidence from ..ingest import IngestResult, detect_format, ingest -from ..persistence import ( - AgentWorkflowRun, - create_agent_workflow, - finalize_agent_workflow, - load_agent_report, - load_agent_workflow, -) +from ..ingest.manifest import DatasetManifest, inspect_h5ad_manifest from . import journal from .context import ContextStagesMixin from .finalization import FinalizationStagesMixin from .models import ( - _RUN_ID_PATTERN, _STAGE_ORDER, - AutomatedPreprocessingPlan, AutomatedWorkflowConfig, AutomatedWorkflowRequest, AutomatedWorkflowResult, AutomatedWorkflowResumeRequest, - AutomatedWorkflowStatus, - FinalAnalysisHandoff, OrchestrationRequestRecord, OrchestrationResumeRecord, + WorkflowIdentity, WorkflowNeedsInput, WorkflowQuestion, - WorkflowStageAttempt, - WorkflowStageName, ) from .preprocessing import PreprocessingStagesMixin +from .rna import ( + selected_rna_assay, + selected_store_rna_assay, + validate_rna_directions, + validate_rna_request_fields, + validate_saved_rna_history, +) from .tuning import TuningStagesMixin +def _model_identity(model: Any) -> str: + from ..config.agent_exec import _model_name + + settings = getattr(model, "settings", None) or {} + provider = getattr(model, "provider", None) + profile = getattr(model, "profile", None) + image_input = getattr(model, "supports_image_input", None) + if not isinstance(image_input, bool) and isinstance(profile, Mapping): + image_input = profile.get("supports_image_input") + identity = { + "settings": settings, + "system": getattr(model, "system", None), + "provider": getattr(provider, "name", None), + "baseUrl": str(getattr(provider, "base_url", "")), + "supportsImageInput": image_input if isinstance(image_input, bool) else None, + } + digest = hashlib.sha256(record_io.canonical_json_bytes(identity)).hexdigest() + return f"{type(model).__module__}.{type(model).__qualname__}:{_model_name(model)}:{digest}" + + +def _submitted_identity(request: AutomatedWorkflowRequest) -> str: + value = request.model_dump(mode="json") + value["sourcePath"] = str(Path(request.sourcePath).resolve()) + if request.zarrPath is not None: + value["zarrPath"] = str(Path(request.zarrPath).resolve()) + return hashlib.sha256(record_io.canonical_json_bytes(value)).hexdigest() + + +def _source_identity(path: str) -> dict[str, Any]: + source = Path(path).resolve() + if source.is_file(): + stat = source.stat() + return { + "path": str(source), + "bytes": stat.st_size, + "modifiedNs": stat.st_mtime_ns, + } + return {"path": str(source)} + + +def _data_identity( + store: DataStore, + assay_name: str, + *, + columns: list[str] | None = None, + feature_columns: list[str] | None = None, +) -> dict[str, Any]: + """Fingerprint the selected assay and original metadata in bounded blocks.""" + from ..parameter_tuning.execution import _metadata_column_fingerprint + + assay = store.get_assay(assay_name) + digest = hashlib.sha256() + digest.update(str(assay.rawData.shape).encode()) + digest.update(str(assay.rawData.dtype).encode()) + for block in assay.rawData.stream_blocks(nthreads=1, prefetch=1): + digest.update(block.tobytes(order="C")) + names = sorted(columns if columns is not None else store.cells.columns) + feature_names = sorted( + feature_columns if feature_columns is not None else assay.feats.columns + ) + return { + "assay": assay_name, + "countsSha256": digest.hexdigest(), + "featureMetadata": { + name: _metadata_column_fingerprint(assay.feats, name) + for name in feature_names + }, + "metadata": { + name: _metadata_column_fingerprint(store.cells, name) for name in names + }, + } + + class AgentOrchestrator( ContextStagesMixin, PreprocessingStagesMixin, TuningStagesMixin, FinalizationStagesMixin, ): - """Run one bounded, persisted workflow through the four Scarf agents.""" + """Run bounded RNA analysis; the journal owns all durable state.""" def __init__( - self, - model: Any, - *, - config: AutomatedWorkflowConfig | None = None, + self, model: Any, *, config: AutomatedWorkflowConfig | None = None ) -> None: self.model = model self.config = config or AutomatedWorkflowConfig() + self._explicit_config_fields = ( + set(config.model_fields_set) if config is not None else set() + ) + + def _validate_resume_config(self, saved: AutomatedWorkflowConfig) -> None: + """Saved defaults remain authoritative unless a caller overrides them.""" + if any( + getattr(saved, field) != getattr(self.config, field) + for field in self._explicit_config_fields + ): + raise ValueError("Resume execution settings differ from the saved workflow") def run(self, request: AutomatedWorkflowRequest) -> AutomatedWorkflowResult: - """Ingest the request and continue until completion or a persisted pause.""" + try: + result = self._run(request) + except Exception as exc: + result = AutomatedWorkflowResult(notes=[describe_agent_error(exc)]) + if result.status != "completed": + logger.error( + f"RNA analysis {result.status} during {result.currentStage}: " + + "; ".join(result.notes) + ) + return result + + def _run(self, request: AutomatedWorkflowRequest) -> AutomatedWorkflowResult: + submitted = request + try: + validate_rna_request_fields(request) + except ValueError as exc: + return AutomatedWorkflowResult(notes=[describe_agent_error(exc)]) + reused = self._reuse_or_resume(request) + if reused is not None: + return reused format_name = detect_format(request.sourcePath) + dataset_manifest: DatasetManifest | None = None logger.info( f"Starting automated agent workflow from {format_name!r} input " f"(workspace={request.workspace is not None})" @@ -91,29 +192,136 @@ def run(self, request: AutomatedWorkflowRequest) -> AutomatedWorkflowResult: currentStage="ingest", notes=["An existing Zarr input cannot be copied implicitly"], ) - if format_name == "zarr" and request.workspace is not None: + if format_name == "h5ad": + matrix_key = request.ingestDirections.get("matrixKey") + try: + dataset_manifest = inspect_h5ad_manifest( + request.sourcePath, + source_uri=( + str(request.ingestDirections["sourceUri"]) + if request.ingestDirections.get("sourceUri") is not None + else request.sourcePath + ), + author_label_policy=request.authorLabelPolicy, + matrix_key=str(matrix_key) if matrix_key is not None else None, + ) + except (OSError, RuntimeError, TypeError, ValueError) as exc: + return AutomatedWorkflowResult( + status="failed", + currentStage="ingest", + notes=[ + f"CELLxGENE manifest inspection failed: {describe_agent_error(exc)}" + ], + ) + if dataset_manifest.declaredBatchColumns: + experimental_directions = dict(request.experimentalDirections) + raw_batch_columns = experimental_directions.get("batchColumns") + if raw_batch_columns is None: + experimental_directions["batchColumns"] = list( + dataset_manifest.declaredBatchColumns + ) + elif not isinstance(raw_batch_columns, list) or any( + not isinstance(value, str) or not value.strip() + for value in raw_batch_columns + ): + return AutomatedWorkflowResult( + status="failed", + currentStage="ingest", + notes=[ + "experimentalDirections.batchColumns must be a list " + "of exact observation-column names" + ], + ) + elif not set(dataset_manifest.declaredBatchColumns).issubset( + raw_batch_columns + ): + return AutomatedWorkflowResult( + status="failed", + currentStage="ingest", + notes=[ + "experimentalDirections.batchColumns must include the " + "CELLxGENE uns/batch_condition columns" + ], + ) + request = request.model_copy( + update={"experimentalDirections": experimental_directions} + ) + manifest_decision = dataset_manifest.decision + if manifest_decision.status == "needsInput": + if self.config.inputPolicy == "unattended": + return AutomatedWorkflowResult( + status="abstained", + currentStage="ingest", + limitations=list(dataset_manifest.priorFiltering.limitations), + unresolvedClaims=[manifest_decision.summary], + notes=[ + "The unattended workflow abstained because the count " + "matrix was ambiguous." + ], + ) + return AutomatedWorkflowResult( + status="needsInput", + currentStage="ingest", + needsInput=WorkflowNeedsInput( + questions=[ + WorkflowQuestion( + questionId="datasetMatrixKey", + question=( + manifest_decision.summary + + ". Rerun with ingestDirections.matrixKey set " + "to the selected option." + ), + options=list(manifest_decision.options), + evidenceIds=list(manifest_decision.evidenceIds), + ) + ] + ), + limitations=list(dataset_manifest.priorFiltering.limitations), + ) + if manifest_decision.status == "abstained": + return AutomatedWorkflowResult( + status="abstained", + currentStage="ingest", + limitations=list(dataset_manifest.priorFiltering.limitations), + unresolvedClaims=[manifest_decision.summary], + notes=[ + "The count-dependent RNA workflow did not run because its " + "input contract is not satisfied." + ], + ) + selected_matrix = manifest_decision.selectedMatrixKey + if selected_matrix is None: + raise RuntimeError("Supported manifest lacks a selected matrix") + ingest_directions = { + **request.ingestDirections, + "matrixKey": selected_matrix, + } + request = request.model_copy(update={"ingestDirections": ingest_directions}) + + if format_name == "zarr": zarr_path = str(Path(request.sourcePath).resolve()) effective_request = request.model_copy(update={"zarrPath": zarr_path}) try: - store = self.open_store(zarr_path, effective_request) - except (KeyError, RuntimeError, TypeError, ValueError) as exc: - logger.error( - "Opening the automated workflow workspace failed " - f"({type(exc).__name__})" + summary = summarize_zarr_readonly( + zarr_path, + workspace=request.workspace, ) + except (OSError, KeyError, RuntimeError, TypeError, ValueError) as exc: return AutomatedWorkflowResult( status="failed", currentStage="ingest", zarrPath=zarr_path, - notes=[f"Opening the requested workspace failed: {exc}"], + notes=[ + f"Opening the requested RNA store failed: {describe_agent_error(exc)}" + ], ) ingest_result = IngestResult( status="done", format="zarr", zarrPath=zarr_path, - assayNames=list(store.assay_names), - summary=store.summary().to_dict(), - actions=["summarize_zarr_workspace"], + assayNames=[assay.name for assay in summary.assays], + summary=summary.to_dict(), + actions=["summarize_zarr"], ) else: ingest_result = ingest( @@ -145,6 +353,20 @@ def run(self, request: AutomatedWorkflowRequest) -> AutomatedWorkflowResult: ) ] ) + if needs_input is not None and self.config.inputPolicy == "unattended": + return AutomatedWorkflowResult( + status="abstained", + currentStage="ingest", + zarrPath=ingest_result.zarrPath, + unresolvedClaims=[ + question.question for question in needs_input.questions + ], + notes=[ + *ingest_result.notes, + "The unattended workflow abstained instead of waiting for " + "an ingest decision.", + ], + ) return AutomatedWorkflowResult( status=("needsInput" if needs_input is not None else "failed"), currentStage="ingest", @@ -153,14 +375,35 @@ def run(self, request: AutomatedWorkflowRequest) -> AutomatedWorkflowResult: notes=list(ingest_result.notes), ) - if not (format_name == "zarr" and request.workspace is not None): + try: + selected = selected_rna_assay( + effective_request, + { + value["name"]: value["assay_type"] + for value in (ingest_result.summary or {}).get("assays", []) + }, + ) + effective_request = effective_request.model_copy( + update={ + "primaryAssay": selected, + "markerAssay": selected, + "analysisAssays": [selected], + } + ) store = self.open_store(ingest_result.zarrPath, effective_request) - workflow = ingest_result.workflowRun or create_agent_workflow(store) + except (OSError, KeyError, RuntimeError, TypeError, ValueError) as exc: + return AutomatedWorkflowResult( + zarrPath=ingest_result.zarrPath, notes=[describe_agent_error(exc)] + ) + ignored = [name for name in store.assay_names if name != selected] logger.info( - f"Continuing automated workflow {workflow.workflowRunId} with " - f"{len(store.assay_names)} datastore assays" + f"RNA analysis: selected assay {selected!r}" + + (f"; ignored other assays {ignored}" if ignored else "") + ) + workflow = WorkflowIdentity(uuid.uuid4().hex, effective_request.workspace) + request_record = self.initialize_request( + store, workflow, effective_request, submitted ) - request_record = self.initialize_request(store, workflow, effective_request) prefix = journal._ensure_orchestration_store(store) self.record_ingest_stage( store, @@ -168,6 +411,7 @@ def run(self, request: AutomatedWorkflowRequest) -> AutomatedWorkflowResult: workflow, request_record, ingest_result, + dataset_manifest, ) return self._continue( store, @@ -176,282 +420,249 @@ def run(self, request: AutomatedWorkflowRequest) -> AutomatedWorkflowResult: answers={}, ) - def resume( - self, - request: AutomatedWorkflowResumeRequest, - ) -> AutomatedWorkflowResult: - """Resume a running workflow after validating its immutable request.""" - logger.info( - f"Resuming automated workflow {request.workflowRunId} with " - f"{len(request.answers)} answer field(s)" - ) - record, store = self.load_request_for_resume(request) - workflow = load_agent_workflow( - store, - request.workflowRunId, - workspace=request.workspace, - ) - if workflow.status != "running": - logger.warning( - f"Automated workflow {workflow.workflowRunId} cannot resume from " - f"status={workflow.status!r}" - ) - prefix = journal._ensure_orchestration_store(store) - if journal._load_terminal_result(store, prefix, workflow) is not None: - raise RuntimeError( - f"Cannot resume a workflow with status {workflow.status!r}" - ) - return self.repair_terminal_result(store, workflow, record) - prefix = journal._ensure_orchestration_store(store) - outcomes = [ - outcome - for stage in _STAGE_ORDER - for outcome in journal._stage_outcomes( - store.zw, prefix, workflow.workflowRunId, stage - ) - ] - starts = [ - started - for stage in _STAGE_ORDER - for started in journal._stage_starts( - store.zw, prefix, workflow.workflowRunId, stage - ) - ] - completed_attempt_ids = { - (outcome.stage, outcome.attemptId) for outcome in outcomes - } - interrupted_starts = [ - started - for started in starts - if (started.stage, started.attemptId) not in completed_attempt_ids - ] - logger.info( - f"Workflow {workflow.workflowRunId} resume scan found " - f"{len(outcomes)} outcome(s) and {len(interrupted_starts)} " - "interrupted attempt(s)" - ) - latest_outcome = ( - max( - outcomes, - key=lambda value: ( - _STAGE_ORDER.index(value.stage), - value.startedAtNs, - value.attemptId, - ), - ) - if outcomes - else None - ) - latest_interrupted = ( - max( - interrupted_starts, - key=lambda value: ( - _STAGE_ORDER.index(value.stage), - value.startedAtNs, - value.attemptId, - ), + def _reuse_or_resume( + self, request: AutomatedWorkflowRequest + ) -> AutomatedWorkflowResult | None: + source = Path(request.sourcePath) + fmt = detect_format(request.sourcePath) + if request.zarrPath is not None: + destination = Path(request.zarrPath) + elif fmt == "zarr": + destination = source + elif source.is_dir(): + destination = source.with_name(source.name + ".zarr") + else: + source_stem = ( + source.with_suffix("") if source.suffix.lower() == ".gz" else source ) - if interrupted_starts - else None - ) - interrupted_lineage = ( - latest_interrupted.inputs.get("resumeLineage") - if latest_interrupted is not None - else None - ) - interrupted_answers_latest_pause = bool( - latest_interrupted is not None - and latest_outcome is not None - and latest_interrupted.stage == latest_outcome.stage - and isinstance(interrupted_lineage, Mapping) - and interrupted_lineage.get("answeredAttempt") - == journal._parent_link(latest_outcome).model_dump(mode="json") - ) - active_interrupted = ( - latest_interrupted - if latest_interrupted is not None - and ( - latest_outcome is None - or _STAGE_ORDER.index(latest_interrupted.stage) - > _STAGE_ORDER.index(latest_outcome.stage) - or interrupted_answers_latest_pause - or ( - latest_interrupted.stage == latest_outcome.stage - and ( - latest_interrupted.startedAtNs, - latest_interrupted.attemptId, - ) - > (latest_outcome.startedAtNs, latest_outcome.attemptId) - ) + destination = source_stem.with_suffix(".zarr") + if not destination.exists(): + return None + root = zarr.open_group(str(destination), mode="r") + active = root if request.workspace is None else root[request.workspace] + if not isinstance(active, zarr.Group): + raise ValueError("Requested workspace is not a group") + prefix = record_io.join_key(active.path, "agents", "orchestrations") + matches = [] + for key in journal._list_keys(active, prefix): + if not key.endswith("/request.json"): + continue + identifier = key.rsplit("/", 2)[-2] + try: + saved = journal.read_request(active, prefix, identifier) + except ValueError: + continue + if saved.inputIdentity.get("userRequestSha256") == _submitted_identity( + request + ): + matches.append(saved) + if len(matches) > 1: + raise ValueError( + "Several workflows match this request; use an exact advanced resume identifier" ) - else None - ) - latest_paused = ( - latest_outcome - if latest_outcome is not None - and latest_outcome.status == "needsInput" - and active_interrupted is None - else None - ) - effective_answers = dict(request.answers) - inherited_resume: OrchestrationResumeRecord | None = None - if active_interrupted is not None: - lineage = active_interrupted.inputs.get("resumeLineage") - if lineage is not None: - if not isinstance(lineage, Mapping): - raise ValueError("Interrupted stage resumeLineage is malformed") - inherited_resume_id = lineage.get("resumeId") - if ( - not isinstance(inherited_resume_id, str) - or _RUN_ID_PATTERN.fullmatch(inherited_resume_id) is None - ): - raise ValueError("Interrupted stage resumeId is malformed") - inherited_resume = journal._validated_resume_record( - store, - prefix, - workflow.workflowRunId, - inherited_resume_id, + if matches: + saved = matches[0] + self._validate_resume_config(saved.config) + if saved.modelIdentity != _model_identity(self.model): + raise ValueError( + "The destination contains this request with different model or execution settings; use a new destination or exact advanced workflow" ) - expected_answered_attempt = ( - inherited_resume.answeredAttempt.model_dump(mode="json") - if inherited_resume.answeredAttempt is not None - else None + return self.resume( + AutomatedWorkflowResumeRequest( + zarrPath=str(destination.resolve()), + workspace=request.workspace, + workflowRunId=saved.workflowRunId, ) - if ( - lineage.get("answeredAttempt") != expected_answered_attempt - or lineage.get("questionIds") != inherited_resume.questionIds - ): - raise ValueError( - "Interrupted stage resumeLineage does not match its resume record" - ) - if request.answers and request.answers != inherited_resume.answers: - raise ValueError( - "Cannot change answers for an in-flight logical invocation" - ) - effective_answers = dict(inherited_resume.answers) - answered_attempt = ( - journal._parent_link(latest_paused) if latest_paused is not None else None - ) - question_ids = ( - [question.questionId for question in latest_paused.needsInput.questions] - if latest_paused is not None and latest_paused.needsInput is not None - else [] - ) - if inherited_resume is not None: - answered_attempt = inherited_resume.answeredAttempt - question_ids = list(inherited_resume.questionIds) - elif latest_paused is None and request.answers: - raise ValueError( - "Cannot provide resume answers: no active persisted questions" ) - resume_record = OrchestrationResumeRecord( + if fmt != "zarr": + raise FileExistsError( + "The destination exists without an exactly matching RNA request; choose a different destination" + ) + return None + + def initialize_request( + self, + store: DataStore, + workflow: WorkflowIdentity, + request: AutomatedWorkflowRequest, + submitted: AutomatedWorkflowRequest | None = None, + ) -> OrchestrationRequestRecord: + prefix = journal._ensure_orchestration_store(store) + if request.primaryAssay is None: + raise ValueError("RNA selection must be resolved before saving the request") + identity = { + "userRequestSha256": _submitted_identity(submitted or request), + "source": _source_identity(request.sourcePath), + "data": _data_identity(store, request.primaryAssay), + } + record = OrchestrationRequestRecord( workflowRunId=workflow.workflowRunId, - resumeId=uuid.uuid4().hex, createdAtNs=time.time_ns(), - answeredAttempt=answered_attempt, - questionIds=question_ids, - answers=effective_answers, + request=request, + config=self.config, + requestSha256=journal._sha256_model(request), + configSha256=journal._sha256_model(self.config), + modelIdentity=_model_identity(self.model), + inputIdentity=identity, ) - resume_record = resume_record.model_copy( - update={"contentSha256": journal._record_checksum(resume_record)} + record = record.model_copy( + update={"contentSha256": journal._record_checksum(record)} ) journal._write_model_once( - store.zw, - journal._resume_key(prefix, workflow.workflowRunId, resume_record.resumeId), - resume_record, - ) - logger.info( - f"Persisted resume {resume_record.resumeId} for workflow " - f"{workflow.workflowRunId} (questions={len(question_ids)})" + store.zw, journal._request_key(prefix, workflow.workflowRunId), record ) - if latest_paused is not None: - answer_errors = journal._resume_answer_errors( - latest_paused, effective_answers - ) - if answer_errors: - logger.warning( - f"Resume answers for workflow {workflow.workflowRunId} did not " - "satisfy the persisted questions" - ) - result = journal.paused_or_failed_result( - store, - workflow, - record, - latest_paused, - ) - result = result.model_copy( - update={"notes": [*result.notes, *answer_errors]} - ) - return result.model_copy( - update={"contentSha256": journal._record_checksum(result)} - ) - return self._continue( + return record + + def load_request_for_resume( + self, request: AutomatedWorkflowResumeRequest + ) -> tuple[OrchestrationRequestRecord, DataStore]: + store = journal.open_analysis_store( + request.zarrPath, request.workflowRunId, workspace=request.workspace + ) + prefix = journal._orchestration_prefix(store) + record = journal.read_request(store.zw, prefix, request.workflowRunId) + if record.modelIdentity != _model_identity(self.model): + raise ValueError("Resume model differs from the saved workflow") + self._validate_resume_config(record.config) + selected = selected_store_rna_assay(store, record.request) + validate_saved_rna_history(store, prefix, request.workflowRunId, selected) + expected = record.inputIdentity + if expected["source"] != _source_identity(record.request.sourcePath): + raise ValueError("Source input has changed since this workflow was started") + observed = _data_identity( store, - workflow, - record, - answers=effective_answers, - resume_record=resume_record, + selected, + columns=list(expected["data"]["metadata"]), + feature_columns=list(expected["data"]["featureMetadata"]), ) + if observed != expected["data"]: + raise ValueError( + "Selected RNA data or relevant metadata changed; start a new analysis" + ) + return record, self.open_store(request.zarrPath, record.request) - def cancel( - self, - request: AutomatedWorkflowResumeRequest, - *, - message: str = "Automated workflow cancelled by the caller", + def resume( + self, request: AutomatedWorkflowResumeRequest + ) -> AutomatedWorkflowResult: + original_config = self.config + try: + result = self._resume(request) + except Exception as exc: + result = AutomatedWorkflowResult( + zarrPath=request.zarrPath, + workspace=request.workspace, + workflowRunId=request.workflowRunId, + notes=[describe_agent_error(exc)], + ) + finally: + self.config = original_config + if result.status != "completed": + logger.error( + f"RNA analysis {result.status} during {result.currentStage}: " + + "; ".join(result.notes) + ) + return result + + def _resume( + self, request: AutomatedWorkflowResumeRequest ) -> AutomatedWorkflowResult: - """Finalize one running automated workflow as abandoned.""" - logger.info(f"Cancelling automated workflow {request.workflowRunId}") record, store = self.load_request_for_resume(request) - workflow = load_agent_workflow( - store, - request.workflowRunId, - workspace=request.workspace, - ) - if workflow.status != "running": - prefix = journal._ensure_orchestration_store(store) - if ( - workflow.status == "abandoned" - and journal._load_terminal_result(store, prefix, workflow) is None - ): - return self.repair_terminal_result(store, workflow, record) - raise RuntimeError( - f"Cannot cancel a workflow with status {workflow.status!r}" + self.config = record.config + workflow = WorkflowIdentity(record.workflowRunId, record.request.workspace) + snapshot = journal.analysis_snapshot(store, workflow.workflowRunId) + if snapshot["status"] == "completed": + if request.answers: + raise ValueError( + "A completed analysis cannot accept new decision answers" + ) + result = AutomatedWorkflowResult( + status="completed", + currentStage="analysis_finalization", + zarrPath=request.zarrPath, + workspace=request.workspace, + workflowRunId=request.workflowRunId, ) - terminal = finalize_agent_workflow( - store, - workflow.workflowRunId, - status="abandoned", - message=message, - ) - prefix = journal._ensure_orchestration_store(store) - observed = [ - outcome + final = snapshot["finalAnalysis"] + result = result.model_copy( + update={"limitations": list(final.get("limitations", []))} + ) + try: + result.report() + except Exception as exc: + return result.model_copy( + update={ + "status": "failed", + "currentStage": "report", + "notes": [describe_agent_error(exc)], + } + ) + return result + stages = snapshot["stages"] + latest = stages[-1] if stages else None + prefix = journal._orchestration_prefix(store) + starts = [ + value for stage in _STAGE_ORDER - for outcome in journal._stage_outcomes( + for value in journal._stage_starts( store.zw, prefix, workflow.workflowRunId, stage ) ] - current_stage: WorkflowStageName = ( - max(observed, key=lambda value: value.startedAtNs).stage - if observed - else "ingest" - ) - result = AutomatedWorkflowResult( - status="abandoned", - currentStage=current_stage, - zarrPath=str(store.zarr_loc), - workflowRun=terminal, - reportReferences=list(terminal.reports), - notes=[message], - ) - result = result.model_copy( - update={"contentSha256": journal._record_checksum(result)} + latest_start = ( + max(starts, key=lambda value: value.startedAtNs) if starts else None ) - logger.info( - f"Automated workflow {workflow.workflowRunId} was abandoned at " - f"stage={current_stage!r}" + answers = dict(request.answers) + resume_record = None + if latest is not None and latest["status"] == "needsInput": + from .models import WorkflowStageAttempt + + outcome = WorkflowStageAttempt.model_validate( + {k: v for k, v in latest.items() if k not in {"report", "decisions"}} + ) + answered = journal._parent_link(outcome) + if ( + not answers + and latest_start is not None + and latest_start.startedAtNs > outcome.startedAtNs + ): + if latest_start.inputs.get("answeredAttempt") == answered.model_dump( + mode="json" + ): + answers = dict(latest_start.inputs.get("resumeAnswers", {})) + if not answers and outcome.stage != "parameter_tuning": + return journal.paused_or_failed_result(store, workflow, record, outcome) + if answers: + errors = journal._resume_answer_errors(outcome, answers) + if errors: + raise ValueError("; ".join(errors)) + assert outcome.needsInput is not None + resume_record = OrchestrationResumeRecord( + workflowRunId=workflow.workflowRunId, + answeredAttempt=answered, + answers=answers, + questionIds=[q.questionId for q in outcome.needsInput.questions], + ) + # Tuning replays committed actions and budgets. With no answer it + # preserves a scientific defer, but retries an uncommitted assessment. + elif answers: + raise ValueError("Resume answers require an exact pending stage") + elif latest_start is not None and latest_start.inputs.get("resumeAnswers"): + from .models import WorkflowStageLink + + answers = dict(latest_start.inputs["resumeAnswers"]) + resume_record = OrchestrationResumeRecord( + workflowRunId=workflow.workflowRunId, + answeredAttempt=WorkflowStageLink.model_validate( + latest_start.inputs["answeredAttempt"] + ), + answers=answers, + questionIds=list(answers), + ) + directions = answers.get("experimentalDirections") + if isinstance(directions, Mapping): + validate_rna_directions(directions) + return self._continue( + store, workflow, record, answers=answers, resume_record=resume_record ) - return journal._persist_terminal_result(store, prefix, terminal, result) def open_store( self, @@ -472,229 +683,67 @@ def open_store( workspace=request.workspace, ) - def initialize_request( - self, - store: DataStore, - workflow: AgentWorkflowRun, - request: AutomatedWorkflowRequest, - ) -> OrchestrationRequestRecord: - prefix = journal._ensure_orchestration_store(store) - request_checksum = journal._sha256_model(request) - config_checksum = journal._sha256_model(self.config) - record = OrchestrationRequestRecord( - workflowRunId=workflow.workflowRunId, - createdAtNs=time.time_ns(), - request=request, - config=self.config, - requestSha256=request_checksum, - configSha256=config_checksum, - ) - record = record.model_copy( - update={"contentSha256": journal._record_checksum(record)} - ) - journal._write_model_once( - store.zw, - journal._request_key(prefix, workflow.workflowRunId), - record, - ) - logger.debug( - f"Persisted immutable request for workflow {workflow.workflowRunId}" - ) - return record - - def load_request_for_resume( - self, - request: AutomatedWorkflowResumeRequest, - ) -> tuple[OrchestrationRequestRecord, DataStore]: - logger.debug(f"Loading immutable request for workflow {request.workflowRunId}") - root = zarr.open_group(request.zarrPath, mode="r") - active = root if request.workspace is None else root[request.workspace] - if not isinstance(active, zarr.Group): - raise ValueError("The requested workspace is not a Zarr group") - root_path = str(getattr(active, "path", "")).strip("/") - prefix = record_io.join_key(root_path, "agents", "orchestrations") - record = cast( - OrchestrationRequestRecord, - journal._read_model( - active, - journal._request_key(prefix, request.workflowRunId), - OrchestrationRequestRecord, - ), - ) - if record.workflowRunId != request.workflowRunId: - raise ValueError("Stored orchestration request has a different workflow") - if ( - Path(cast(str, record.request.zarrPath)).resolve() - != Path(request.zarrPath).resolve() - ): - raise ValueError("Resume zarrPath does not match the stored request") - if record.request.workspace != request.workspace: - raise ValueError("Resume workspace does not match the stored request") - if record.requestSha256 != journal._sha256_model(record.request): - raise ValueError("Stored orchestration request checksum is invalid") - if record.configSha256 != journal._sha256_model(record.config): - raise ValueError("Stored orchestration config checksum is invalid") - if record.contentSha256 != journal._record_checksum(record): - raise ValueError("Stored orchestration request envelope is invalid") - store = self.open_store(request.zarrPath, record.request) - return record, store - - def repair_terminal_result( + def _continue( self, store: DataStore, - workflow: AgentWorkflowRun, + workflow: WorkflowIdentity, request_record: OrchestrationRequestRecord, + *, + answers: Mapping[str, Any], + resume_record: OrchestrationResumeRecord | None = None, ) -> AutomatedWorkflowResult: - """Load or reconstruct the JSON result after terminal finalization.""" - prefix = journal._ensure_orchestration_store(store) - existing = journal._load_terminal_result(store, prefix, workflow) - if existing is not None: - logger.debug( - f"Loaded terminal result for workflow {workflow.workflowRunId}" - ) - return existing - - logger.warning( - f"Repairing missing terminal result for workflow {workflow.workflowRunId}" - ) - - observed = [ - outcome - for stage in _STAGE_ORDER - for outcome in journal._stage_outcomes( - store.zw, - prefix, - workflow.workflowRunId, - stage, + try: + return self._execute_stages( + store, + workflow, + request_record, + answers=answers, + resume_record=resume_record, ) - ] - if not observed: - raise RuntimeError("Terminal workflow has no persisted stage outcomes") - by_identity = { - (outcome.stage, outcome.attemptId): outcome for outcome in observed - } - - def validated_chain( - terminal: WorkflowStageAttempt, - ) -> dict[WorkflowStageName, WorkflowStageAttempt] | None: - chain: dict[WorkflowStageName, WorkflowStageAttempt] = {} - current = terminal - while True: - stage_index = _STAGE_ORDER.index(current.stage) - if current.stage in chain: - raise ValueError("Stage lineage contains a cycle") - if not journal._stage_outcome_resolves( - store, - prefix, - workflow.workflowRunId, - request_record, - current, - ): - return None - chain[current.stage] = current - if stage_index == 0: - if current.parentAttempts: - raise ValueError("The ingest stage cannot have a parent") - return chain - if len(current.parentAttempts) != 1: - raise ValueError("Every post-ingest stage must have one parent") - parent_link = current.parentAttempts[0] - if ( - workflow.status != "abandoned" - and parent_link.stage != _STAGE_ORDER[stage_index - 1] - ): - raise ValueError("Terminal stage lineage skips a workflow stage") - if _STAGE_ORDER.index(parent_link.stage) >= stage_index: - raise ValueError("Stage lineage does not move toward ingest") - parent = by_identity.get((parent_link.stage, parent_link.attemptId)) - if ( - parent is None - or parent.status != "done" - or parent.contentSha256 != parent_link.contentSha256 - ): - return None - current = parent - - if workflow.status == "completed": - terminal_candidates = [ - outcome - for outcome in observed - if outcome.stage == "biological_interpretation" - and outcome.status == "done" - ] - elif workflow.status == "failed": - terminal_candidates = [ - outcome for outcome in observed if outcome.status == "failed" + except Exception as exc: + prefix = journal._orchestration_prefix(store) + starts = [ + value + for stage in _STAGE_ORDER + for value in journal._stage_starts( + store.zw, prefix, workflow.workflowRunId, stage + ) ] - else: - terminal_candidates = list(observed) - terminal_candidates.sort( - key=lambda value: (value.startedAtNs, value.attemptId), - reverse=True, - ) - terminal_outcome: WorkflowStageAttempt | None = None - validated_done: dict[WorkflowStageName, WorkflowStageAttempt] = {} - for candidate in terminal_candidates: - chain = validated_chain(candidate) - if chain is not None: - terminal_outcome = candidate - validated_done = { - stage: outcome - for stage, outcome in chain.items() - if outcome.status == "done" - } - break - if terminal_outcome is None: - raise RuntimeError( - "Terminal automated workflow lacks one valid persisted stage chain" - ) - - preprocessing_plan: AutomatedPreprocessingPlan | None = None - plan_outcome = validated_done.get("preprocessing_plan") - if plan_outcome is not None and "preprocessingPlan" in plan_outcome.outputs: - preprocessing_plan = AutomatedPreprocessingPlan.model_validate( - plan_outcome.outputs["preprocessingPlan"] + latest = ( + max(starts, key=lambda value: value.startedAtNs) if starts else None ) - - final_analysis: FinalAnalysisHandoff | None = None - finalization_outcome = validated_done.get("analysis_finalization") - if ( - finalization_outcome is not None - and "finalAnalysis" in finalization_outcome.outputs - ): - final_analysis = FinalAnalysisHandoff.model_validate( - finalization_outcome.outputs["finalAnalysis"] + if latest is not None and not any( + value.attemptId == latest.attemptId + for value in journal._stage_outcomes( + store.zw, prefix, workflow.workflowRunId, latest.stage + ) + ): + try: + journal.finish_exception(store, prefix, workflow, latest, exc) + except Exception as persistence_error: + exc.add_note( + "Saving the failed stage also failed: " + + describe_agent_error(persistence_error) + ) + return AutomatedWorkflowResult( + currentStage=latest.stage if latest else "ingest", + zarrPath=str(store.zarr_loc), + workspace=workflow.workspace, + workflowRunId=workflow.workflowRunId, + notes=[describe_agent_error(exc)], ) - for reference in workflow.reports: - load_agent_report(store, reference) - notes = [workflow.finalizationMessage] if workflow.finalizationMessage else [] - result = AutomatedWorkflowResult( - status=cast(AutomatedWorkflowStatus, workflow.status), - currentStage=terminal_outcome.stage, - zarrPath=str(store.zarr_loc), - workflowRun=workflow, - reportReferences=list(workflow.reports), - preprocessingPlan=preprocessing_plan, - finalAnalysis=final_analysis, - notes=notes, - ) - result = result.model_copy( - update={"contentSha256": journal._record_checksum(result)} - ) - return journal._persist_terminal_result(store, prefix, workflow, result) - - def _continue( + def _execute_stages( self, store: DataStore, - workflow: AgentWorkflowRun, + workflow: WorkflowIdentity, request_record: OrchestrationRequestRecord, *, answers: Mapping[str, Any], resume_record: OrchestrationResumeRecord | None = None, ) -> AutomatedWorkflowResult: """Continue the stage machine from the latest validated checkpoint.""" - logger.info(f"Running stage sequence for workflow {workflow.workflowRunId}") + logger.debug(f"Running stage sequence for workflow {workflow.workflowRunId}") prefix = journal._ensure_orchestration_store(store) ingest_outcome = journal._validated_done_outcome( store, @@ -731,7 +780,7 @@ def _continue( ) parents = [journal._parent_link(enrichment_outcome)] - hto_outcome = self._hto_stage( + quality_outcome = self._rna_quality_metrics_stage( store, workflow, request_record, @@ -740,24 +789,24 @@ def _continue( cell_selection, resume_record=resume_record, ) - if hto_outcome.status != "done": + if quality_outcome.status != "done": return journal.paused_or_failed_result( store, workflow, request_record, - hto_outcome, + quality_outcome, ) quality_metric_artifacts = self._named_stage_artifacts( - hto_outcome, + quality_outcome, "qualityMetricArtifacts", "quality_metric", ) hto_identity_artifacts = self._named_stage_artifacts( - hto_outcome, + quality_outcome, "htoIdentityArtifacts", "hto_identity", ) - parents = [journal._parent_link(hto_outcome)] + parents = [journal._parent_link(quality_outcome)] context_outcome, experimental = self.experimental_context_stage( store, @@ -778,6 +827,10 @@ def _continue( request_record, context_outcome, ) + study_contract = StudyContract.model_validate( + context_outcome.outputs["studyContract"] + ) + validate_objective_evidence(study_contract, experimental) parents = [journal._parent_link(context_outcome)] plan_outcome, preprocessing_plan = self.preprocessing_plan_stage( @@ -788,6 +841,7 @@ def _continue( enrichment, experimental, ingest_outcome, + study_contract, answers, resume_record=resume_record, ) @@ -797,17 +851,23 @@ def _continue( workflow, request_record, plan_outcome, - preprocessing_plan=preprocessing_plan, + study_contract=study_contract, ) parents = [journal._parent_link(plan_outcome)] - preprocessing_outcome, preprocessed = self.preprocessing_stage( + ( + preprocessing_outcome, + preprocessed, + preprocessing_plan, + ) = self.preprocessing_stage( store, workflow, request_record, parents, preprocessing_plan, experimental, + study_contract, + answers, resume_record=resume_record, ) if preprocessing_outcome.status != "done": @@ -816,7 +876,7 @@ def _continue( workflow, request_record, preprocessing_outcome, - preprocessing_plan=preprocessing_plan, + study_contract=study_contract, ) parents = [journal._parent_link(preprocessing_outcome)] @@ -831,6 +891,7 @@ def _continue( enrichment_outcome.reportReferences[0], context_outcome.reportReferences[0], answers, + study_contract=study_contract, resume_record=resume_record, ) if tuning_outcome.status != "done": @@ -839,9 +900,34 @@ def _continue( workflow, request_record, tuning_outcome, - preprocessing_plan=preprocessing_plan, + study_contract=study_contract, ) + validate_objective_evidence(study_contract, experimental) parents = [journal._parent_link(tuning_outcome)] + tuning_reference = tuning_outcome.reportReferences[0] + selected = next( + ( + value + for value in tuning_report.evaluations + if value.candidateId == tuning_report.recommendedCandidateId + ), + None, + ) + if selected is None: + raise ValueError("Full-cohort tuning did not select an evaluated candidate") + preprocessed = [ + value.model_copy( + update={ + "normalized": selected.artifacts.get( + "normalized", value.normalized + ), + "graphFeatures": selected.artifacts.get( + "graphFeatures", value.graphFeatures + ), + } + ) + for value in preprocessed + ] finalization_outcome, final_analysis = self.analysis_finalization_stage( store, @@ -851,7 +937,7 @@ def _continue( preprocessing_plan, preprocessed, tuning_report, - tuning_outcome.reportReferences[0], + tuning_reference, resume_record=resume_record, ) if finalization_outcome.status != "done": @@ -860,56 +946,32 @@ def _continue( workflow, request_record, finalization_outcome, - preprocessing_plan=preprocessing_plan, - ) - parents = [journal._parent_link(finalization_outcome)] - - biology_outcome = self.biological_interpretation_stage( - store, - workflow, - request_record, - parents, - enrichment, - experimental, - tuning_report, - final_analysis, - enrichment_outcome.reportReferences[0], - context_outcome.reportReferences[0], - tuning_outcome.reportReferences[0], - answers, - resume_record=resume_record, - ) - if biology_outcome.status != "done": - return journal.paused_or_failed_result( - store, - workflow, - request_record, - biology_outcome, - preprocessing_plan=preprocessing_plan, - final_analysis=final_analysis, + study_contract=study_contract, ) - terminal = finalize_agent_workflow( - store, - workflow.workflowRunId, - status="completed", - message="Automated Scarf agent workflow completed", - ) completed = AutomatedWorkflowResult( status="completed", - currentStage="biological_interpretation", + currentStage="analysis_finalization", zarrPath=str(store.zarr_loc), - workflowRun=terminal, - reportReferences=list(terminal.reports), - preprocessingPlan=preprocessing_plan, - finalAnalysis=final_analysis, - notes=["Automated analysis completed"], - ) - completed = completed.model_copy( - update={"contentSha256": journal._record_checksum(completed)} - ) - logger.info( - f"Automated workflow {workflow.workflowRunId} completed with " - f"{len(terminal.reports)} report(s)" + workspace=workflow.workspace, + workflowRunId=workflow.workflowRunId, + limitations=list(final_analysis.limitations), + notes=["RNA analysis completed"], ) - return journal._persist_terminal_result(store, prefix, terminal, completed) + from ..report.generator import generate_agent_report + + try: + path = generate_agent_report(store, workflow.workflowRunId) + except Exception as exc: + logger.error(f"Analysis report failed: {describe_agent_error(exc)}") + return completed.model_copy( + update={ + "status": "failed", + "currentStage": "report", + "notes": [ + f"Report generation failed: {describe_agent_error(exc)}; the validated analysis is saved and can be resumed." + ], + } + ) + logger.info(f"Completed RNA analysis. Report saved to {path}") + return completed diff --git a/scarf/agent/orchestrator/models.py b/scarf/agent/orchestrator/models.py index 812b66c3..9ef34336 100644 --- a/scarf/agent/orchestrator/models.py +++ b/scarf/agent/orchestrator/models.py @@ -1,53 +1,105 @@ """Public data models for resumable automated agent workflows.""" import re -from typing import Any, Literal +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal from pydantic import Field, field_validator, model_validator from ...storage.refs import ArtifactRef +from ..cell_quality.profiles import cell_qc_policy from ..config import AgentRunConfig -from ..experimental_context import CellQcPlan -from ..persistence import AgentReportReference, AgentWorkflowRun +from ..decisions.rna import CellQualityExecutorPayload +from ..experimental_context.contracts import CellQcPlan +from ..experimental_context.study import AuthorLabelPolicy from ..types import AgentDataModel, ArtifactReferenceModel -type AutomatedWorkflowStatus = Literal["completed", "needsInput", "failed", "abandoned"] -type WorkflowStageStatus = Literal["started", "done", "needsInput", "failed"] +if TYPE_CHECKING: + import pandas as pd + + from ...datastore.datastore import DataStore + from ...plotting._figure import PlotResult + +type AutomatedWorkflowStatus = Literal[ + "completed", + "needsInput", + "abstained", + "failed", + "abandoned", +] +type WorkflowInputPolicy = Literal["pause", "unattended"] +type WorkflowStageStatus = Literal[ + "started", "done", "needsInput", "abstained", "failed" +] type WorkflowStageName = Literal[ "ingest", "data_enrichment", - "hto_demultiplexing", + "rna_quality_metrics", "experimental_context", "preprocessing_plan", "preprocessing", "parameter_tuning", "analysis_finalization", - "biological_interpretation", + "report", ] type AssayRole = Literal["graph", "hto", "unsupported"] type ReductionMethod = Literal["pca", "lsi", "identity", "none"] _RUN_ID_PATTERN = re.compile(r"^[a-z0-9][a-z0-9_-]{0,127}$") _SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$") -_ORCHESTRATION_FORMAT = "scarf_agent_orchestrations" -_ORCHESTRATION_VERSION = 1 _STAGE_ORDER: tuple[WorkflowStageName, ...] = ( "ingest", "data_enrichment", - "hto_demultiplexing", + "rna_quality_metrics", "experimental_context", "preprocessing_plan", "preprocessing", "parameter_tuning", "analysis_finalization", - "biological_interpretation", ) +@dataclass(frozen=True) +class WorkflowIdentity: + """Runtime address of a workflow whose state belongs to its journal.""" + + workflowRunId: str + workspace: str | None = None + + +class StageEvidenceReference(AgentDataModel): + """Exact evidence checkpoint owned by an orchestration stage.""" + + workflowRunId: str + stage: WorkflowStageName + key: str + contentSha256: str + + +class AnalysisError(RuntimeError): + """An unattended analysis could not produce a supported completed result.""" + + def __init__(self, result: "AutomatedWorkflowResult") -> None: + self.result = result + detail = "; ".join(result.notes) or "Required scientific evidence is unresolved" + address = ( + f" Resume workflow {result.workflowRunId!r} in {result.zarrPath!r}" + f" (workspace={result.workspace!r})." + if result.workflowRunId and result.zarrPath + else "" + ) + super().__init__( + f"RNA analysis {result.status} during {result.currentStage}: {detail}.{address}" + ) + + class WorkflowQuestion(AgentDataModel): """One stable question that can be answered by a resume request.""" questionId: str = "" + decisionId: str | None = None question: str = "" options: list[str] = Field(default_factory=list) evidenceIds: list[str] = Field(default_factory=list) @@ -57,14 +109,6 @@ class WorkflowQuestion(AgentDataModel): def get_blank(cls) -> "WorkflowQuestion": return cls() - @classmethod - def get_example(cls) -> "WorkflowQuestion": - return cls( - questionId="approvePlanChecksum", - question="Approve this preprocessing plan?", - planChecksum="0" * 64, - ) - class WorkflowNeedsInput(AgentDataModel): """All questions blocking the next workflow stage.""" @@ -75,10 +119,6 @@ class WorkflowNeedsInput(AgentDataModel): def get_blank(cls) -> "WorkflowNeedsInput": return cls() - @classmethod - def get_example(cls) -> "WorkflowNeedsInput": - return cls(questions=[WorkflowQuestion.get_example()]) - class WorkflowStageLink(AgentDataModel): """Immutable identity of one completed parent stage attempt.""" @@ -105,10 +145,6 @@ def validate_checksum(cls, value: str) -> str: def get_blank(cls) -> "WorkflowStageLink": return cls() - @classmethod - def get_example(cls) -> "WorkflowStageLink": - return cls(stage="ingest", attemptId="attempt-1", contentSha256="0" * 64) - class WorkflowStageAttempt(AgentDataModel): """Append-only record for one orchestration stage attempt.""" @@ -122,7 +158,7 @@ class WorkflowStageAttempt(AgentDataModel): requestSha256: str = "" configSha256: str = "" parentAttempts: list[WorkflowStageLink] = Field(default_factory=list) - reportReferences: list[AgentReportReference] = Field(default_factory=list) + reportReferences: list[StageEvidenceReference] = Field(default_factory=list) artifacts: dict[str, ArtifactReferenceModel] = Field(default_factory=dict) inputs: dict[str, Any] = Field(default_factory=dict) outputs: dict[str, Any] = Field(default_factory=dict) @@ -148,20 +184,6 @@ def validate_lifecycle(self) -> "WorkflowStageAttempt": def get_blank(cls) -> "WorkflowStageAttempt": return cls() - @classmethod - def get_example(cls) -> "WorkflowStageAttempt": - return cls( - workflowRunId="workflow-1", - stage="ingest", - attemptId="attempt-1", - status="done", - startedAtNs=1, - completedAtNs=2, - requestSha256="0" * 64, - configSha256="1" * 64, - contentSha256="2" * 64, - ) - class AssayPreprocessingPlan(AgentDataModel): """Exact allowlisted preprocessing route for one assay.""" @@ -174,9 +196,6 @@ class AssayPreprocessingPlan(AgentDataModel): featureMethod: Literal["hvg", "prevalentPeaks", "panel", "none"] = "none" reductionMethod: ReductionMethod = "none" featureParameters: dict[str, Any] = Field(default_factory=dict) - normalizationParameters: dict[str, Any] = Field(default_factory=dict) - reductionParameters: dict[str, Any] = Field(default_factory=dict) - exactExcludedFeatures: list[str] = Field(default_factory=list) evidenceIds: list[str] = Field(default_factory=list) limitations: list[str] = Field(default_factory=list) @@ -184,19 +203,6 @@ class AssayPreprocessingPlan(AgentDataModel): def get_blank(cls) -> "AssayPreprocessingPlan": return cls() - @classmethod - def get_example(cls) -> "AssayPreprocessingPlan": - return cls( - assay="RNA", - assayType="RNA", - role="graph", - graphEligible=True, - markerEligible=True, - featureMethod="hvg", - reductionMethod="pca", - featureParameters={"topN": 1000, "minCells": 20}, - ) - class AutomatedPreprocessingPlan(AgentDataModel): """Dataset-wide preprocessing plan produced before selection changes.""" @@ -205,29 +211,28 @@ class AutomatedPreprocessingPlan(AgentDataModel): markerAssay: str = "" cellSelection: ArtifactReferenceModel | None = None cellQc: CellQcPlan = Field(default_factory=CellQcPlan.get_blank) + cellQualityPayload: CellQualityExecutorPayload | None = None assays: list[AssayPreprocessingPlan] = Field(default_factory=list) pairedAssays: list[str] = Field(default_factory=list) planChecksum: str = "" limitations: list[str] = Field(default_factory=list) + @model_validator(mode="after") + def validate_cell_quality_payload(self) -> "AutomatedPreprocessingPlan": + if ( + self.cellQualityPayload is not None + and cell_qc_policy(self.cellQc.action, self.cellQc.registeredProfile) + != self.cellQualityPayload.profile + ): + raise ValueError( + "cellQualityPayload must match the exact selected QC policy" + ) + return self + @classmethod def get_blank(cls) -> "AutomatedPreprocessingPlan": return cls() - @classmethod - def get_example(cls) -> "AutomatedPreprocessingPlan": - return cls( - primaryAssay="RNA", - markerAssay="RNA", - cellSelection=ArtifactReferenceModel( - scope="datastore", - kind="cell_selection", - artifactId="c" * 64, - ), - assays=[AssayPreprocessingPlan.get_example()], - planChecksum="0" * 64, - ) - class PreprocessedAssayHandoff(AgentDataModel): """Exact normalized input and selections handed to Parameter Tuning.""" @@ -239,6 +244,13 @@ class PreprocessedAssayHandoff(AgentDataModel): graphFeatures: ArtifactReferenceModel | None = None markerFeatures: ArtifactReferenceModel | None = None normalized: ArtifactReferenceModel | None = None + graphFeatureCandidates: dict[str, ArtifactReferenceModel] = Field( + default_factory=dict + ) + normalizedCandidates: dict[str, ArtifactReferenceModel] = Field( + default_factory=dict + ) + featureCandidateEvaluations: list[dict[str, Any]] = Field(default_factory=list) nCells: int = 0 nFeatures: int = 0 @@ -246,123 +258,115 @@ class PreprocessedAssayHandoff(AgentDataModel): def get_blank(cls) -> "PreprocessedAssayHandoff": return cls() - @classmethod - def get_example(cls) -> "PreprocessedAssayHandoff": - return cls( - assay="RNA", - assayType="RNA", - cellSelection=ArtifactReferenceModel( - scope="datastore", - kind="cell_selection", - artifactId="c" * 64, - ), - reductionMethod="pca", - graphFeatures=ArtifactReferenceModel.get_example(), - markerFeatures=ArtifactReferenceModel.get_example(), - normalized=ArtifactReferenceModel( - assay="RNA", kind="normalized", artifactId="1" * 64 - ), - nCells=100, - nFeatures=1000, - ) - - -class NativeAnalysisHandoff(AgentDataModel): - """Selected immutable native analysis chain for one assay.""" - - assay: str = "" - reductionMethod: ReductionMethod = "none" - featureSelection: ArtifactReferenceModel | None = None - markerFeatures: ArtifactReferenceModel | None = None - normalized: ArtifactReferenceModel | None = None - reduction: ArtifactReferenceModel | None = None - batchCorrection: ArtifactReferenceModel | None = None - annIndex: ArtifactReferenceModel | None = None - embeddingInitialization: ArtifactReferenceModel | None = None - neighbors: ArtifactReferenceModel | None = None - graph: ArtifactReferenceModel | None = None - clusters: ArtifactReferenceModel | None = None - umap: ArtifactReferenceModel | None = None - - @classmethod - def get_blank(cls) -> "NativeAnalysisHandoff": - return cls() - - @classmethod - def get_example(cls) -> "NativeAnalysisHandoff": - return cls(assay="RNA", reductionMethod="pca") - class FinalAnalysisHandoff(AgentDataModel): - """Replayable final analysis used by Biological Interpretation.""" + """Exact final RNA artifacts validated by the concluding journal checkpoint.""" workflowRunId: str = "" primaryAssay: str = "" markerAssay: str = "" cellSelection: ArtifactReferenceModel | None = None - nativeAnalyses: list[NativeAnalysisHandoff] = Field(default_factory=list) graph: ArtifactReferenceModel | None = None - graphMethod: Literal["native", "snn", "wnn"] = "native" + graphMethod: Literal["native"] = "native" clusters: ArtifactReferenceModel | None = None embeddingInitialization: ArtifactReferenceModel | None = None umap: ArtifactReferenceModel | None = None markerFeatures: ArtifactReferenceModel | None = None markers: ArtifactReferenceModel | None = None - parameterReport: AgentReportReference | None = None + doubletScores: list[ArtifactReferenceModel] = Field(default_factory=list) + doubletScoreSelections: list[ArtifactReferenceModel] = Field(default_factory=list) limitations: list[str] = Field(default_factory=list) @classmethod def get_blank(cls) -> "FinalAnalysisHandoff": return cls() - @classmethod - def get_example(cls) -> "FinalAnalysisHandoff": - return cls( - workflowRunId="workflow-1", - primaryAssay="RNA", - markerAssay="RNA", - cellSelection=ArtifactReferenceModel( - scope="datastore", - kind="cell_selection", - artifactId="c" * 64, - ), - nativeAnalyses=[NativeAnalysisHandoff.get_example()], - graph=ArtifactReferenceModel( - assay="RNA", kind="connectivity_map", artifactId="2" * 64 - ), - embeddingInitialization=ArtifactReferenceModel( - assay="RNA", - kind="embedding_initialization", - artifactId="5" * 64, - ), - clusters=ArtifactReferenceModel( - assay="RNA", kind="cluster_labels", artifactId="3" * 64 - ), - ) - class AutomatedWorkflowConfig(AgentDataModel): """Bounded execution policy for automated workflows.""" - primaryInitialCandidates: int = Field(default=5, ge=1) - secondaryInitialCandidates: int = Field(default=3, ge=1) - maxRefinedCandidatesPerAssay: int = Field(default=1, ge=0, le=1) - maxHarmonyCandidatesPerAssay: int = Field(default=1, ge=0, le=1) - integrationResolutionCandidates: int = Field(default=3, ge=1) - maxCandidateBranches: int = Field(default=24, ge=1) - minClusterCells: int = Field(default=20, ge=1) - maxIdentityFeatures: int = Field(default=64, ge=2) - maxGraphAssays: int = Field(default=3, ge=1) + inputPolicy: WorkflowInputPolicy = Field( + default="pause", + exclude_if=lambda value: value == "pause", + ) + # Older requests always scored doublets. Keep their serialized defaults exact. + scoreDoublets: bool = Field( + default=True, + strict=True, + exclude_if=lambda value: value is True, + description="Score advisory doublets; Harmony-eligible runs always retain required doublet diagnostics. The beginner API defaults to false; absent saved fields retain enabled scoring.", + ) + screeningCells: int | None = Field( + default=None, + ge=20, + description="Fixed initial screening size; null uses 10% of retained cells bounded to 10,000–100,000.", + ) + maxScreeningCells: int = Field(default=100_000, ge=20) + maxScreeningEvaluations: int = Field(default=24, ge=4) + maxTotalScreeningEvaluations: int = Field(default=48, ge=4) + maxFullGraphs: int = Field(default=4, ge=1) + maxFullPartitions: int = Field(default=8, ge=1) + maxFullRepairs: int = Field(default=1, ge=0, le=1) + randomSeed: int = Field(default=4444, ge=0) allowDownloads: bool = False cacheDir: str | None = None agentRunConfig: AgentRunConfig = Field(default_factory=AgentRunConfig) - @classmethod - def get_blank(cls) -> "AutomatedWorkflowConfig": - return cls() + @model_validator(mode="before") + @classmethod + def reject_obsolete_configuration(cls, value: Any) -> Any: + if isinstance(value, Mapping): + obsolete = sorted( + set(value) + & { + "maxRefinedCandidatesPerAssay", + "maxHarmonyCandidatesPerAssay", + "runConfoundedHarmonyDiagnostic", + "maxCandidateEvaluations", + "maxIdentityFeatures", + "minClusterCells", + "hvgCandidateCounts", + "pcaCandidateDimensions", + "graphNeighborCandidates", + "leidenResolutionCandidates", + "maxRevisions", + "primaryInitialCandidates", + "secondaryInitialCandidates", + "integrationResolutionCandidates", + "maxCandidateBranches", + "maxGraphAssays", + "leidenSeeds", + "clusterSubsamples", + "clusterSubsampleFraction", + } + ) + if obsolete: + raise ValueError( + "Unsupported legacy workflow configuration fields: " + + ", ".join(obsolete) + + ". Create a new single-RNA workflow configuration with " + "screening and full-cohort work limits. " + "Saved workflows using these fields cannot be resumed or " + "regenerated with this release; their analysis artifacts " + "remain available through Scarf's artifact APIs." + ) + return value + + @model_validator(mode="after") + def validate_work_limits(self) -> "AutomatedWorkflowConfig": + if ( + self.screeningCells is not None + and self.maxScreeningCells < self.screeningCells + ): + raise ValueError("maxScreeningCells must be at least screeningCells") + if self.maxTotalScreeningEvaluations < self.maxScreeningEvaluations: + raise ValueError( + "Whole-workflow screening allowance must cover one screening population" + ) + return self @classmethod - def get_example(cls) -> "AutomatedWorkflowConfig": + def get_blank(cls) -> "AutomatedWorkflowConfig": return cls() @@ -372,8 +376,9 @@ class AutomatedWorkflowRequest(AgentDataModel): sourcePath: str = "" zarrPath: str | None = None studyContext: str = "" + studyObjective: str = "" + authorLabelPolicy: AuthorLabelPolicy = "holdout" workspace: str | None = None - allowAssumptions: bool = False primaryAssay: str | None = None markerAssay: str | None = None analysisAssays: list[str] = Field(default_factory=list) @@ -387,24 +392,21 @@ def validate_request(self) -> "AutomatedWorkflowRequest": raise ValueError("sourcePath must be non-empty") if not self.studyContext.strip(): raise ValueError("studyContext must be non-empty") + if not self.studyObjective.strip(): + raise ValueError("studyObjective must be non-empty") if len(set(self.analysisAssays)) != len(self.analysisAssays): raise ValueError("analysisAssays must be unique") - if len(set(self.pairedAssays)) != len(self.pairedAssays): - raise ValueError("pairedAssays must be unique") - if self.pairedAssays and len(self.pairedAssays) < 2: - raise ValueError("pairedAssays must contain at least two assays") + from .rna import validate_rna_request_fields + + validate_rna_request_fields(self) return self @classmethod def get_blank(cls) -> "AutomatedWorkflowRequest": - return cls(sourcePath="dataset.h5", studyContext="Study context") - - @classmethod - def get_example(cls) -> "AutomatedWorkflowRequest": return cls( - sourcePath="dataset.h5ad", - zarrPath="dataset.zarr", - studyContext="Single-cell profiling of treated human blood.", + sourcePath="dataset.h5", + studyContext="Study context", + studyObjective="Discover stable population structure.", ) @@ -428,49 +430,114 @@ def validate_request(self) -> "AutomatedWorkflowResumeRequest": def get_blank(cls) -> "AutomatedWorkflowResumeRequest": return cls(zarrPath="dataset.zarr", workflowRunId="workflow-1") - @classmethod - def get_example(cls) -> "AutomatedWorkflowResumeRequest": - return cls( - zarrPath="dataset.zarr", - workflowRunId="workflow-1", - answers={"approvePlanChecksum": "0" * 64}, - ) - class AutomatedWorkflowResult(AgentDataModel): - """Bounded result of running or resuming an automated workflow.""" + """Small result address; scientific evidence remains in the stage journal.""" status: AutomatedWorkflowStatus = "failed" currentStage: WorkflowStageName = "ingest" zarrPath: str | None = None - workflowRun: AgentWorkflowRun | None = None - reportReferences: list[AgentReportReference] = Field(default_factory=list) - preprocessingPlan: AutomatedPreprocessingPlan | None = None - finalAnalysis: FinalAnalysisHandoff | None = None + workspace: str | None = None + workflowRunId: str | None = None needsInput: WorkflowNeedsInput | None = None + limitations: list[str] = Field(default_factory=list) + unresolvedClaims: list[str] = Field(default_factory=list) notes: list[str] = Field(default_factory=list) - contentSha256: str = "" + + def _analysis_store(self) -> "DataStore": + from .journal import open_analysis_store + + if self.status != "completed": + raise AnalysisError(self) + if self.zarrPath is None or self.workflowRunId is None: + raise RuntimeError("Completed analysis lacks its exact store and workflow") + return open_analysis_store( + self.zarrPath, self.workflowRunId, workspace=self.workspace + ) + + def _completed_analysis(self, store: "DataStore") -> FinalAnalysisHandoff: + from .journal import analysis_snapshot + + assert self.workflowRunId is not None + snapshot = analysis_snapshot(store, self.workflowRunId) + if snapshot["status"] != "completed": + raise RuntimeError( + "The referenced workflow has no validated final analysis" + ) + return FinalAnalysisHandoff.model_validate(snapshot["finalAnalysis"]) + + def plot_embedding(self, **kwargs: Any) -> "PlotResult": + """Display the saved cluster map with bounded reads and full-count provenance.""" + from .._plots import plot_final_umap + + store = self._analysis_store() + final = self._completed_analysis(store) + if ( + final.umap is None + or final.clusters is None + or final.cellSelection is None + or final.graph is None + ): + raise RuntimeError("Completed analysis lacks its final map artifacts") + for name in ( + "layout", + "run", + "color_by", + "umap", + "clusters", + "cell_selection", + "graph", + ): + if name in kwargs: + raise ValueError("plot_embedding uses the exact completed cluster map") + return plot_final_umap( + store, + umap=artifact_model_to_ref(final.umap), + clusters=artifact_model_to_ref(final.clusters), + cell_selection=artifact_model_to_ref(final.cellSelection), + graph=artifact_model_to_ref(final.graph), + **kwargs, + ) + + def get_markers( + self, + *, + group_id: str | int | None = None, + min_score: float = 0.25, + min_frac_exp: float = 0.2, + ) -> "pd.DataFrame": + """Load the final markers using Scarf's established marker filters.""" + store = self._analysis_store() + final = self._completed_analysis(store) + if final.markers is None: + raise RuntimeError("Completed analysis lacks its marker artifact") + return store.get_markers( + marker=artifact_model_to_ref(final.markers), + group_id=group_id, + min_score=min_score, + min_frac_exp=min_frac_exp, + ) + + def report(self) -> Path: + """Return or regenerate the compact report from saved evidence only.""" + from ..report.generator import generate_agent_report + + store = self._analysis_store() + self._completed_analysis(store) + assert self.workflowRunId is not None + return generate_agent_report(store, self.workflowRunId) @classmethod def get_blank(cls) -> "AutomatedWorkflowResult": return cls() - @classmethod - def get_example(cls) -> "AutomatedWorkflowResult": - return cls( - status="completed", - currentStage="biological_interpretation", - zarrPath="dataset.zarr", - workflowRun=AgentWorkflowRun.get_example(), - finalAnalysis=FinalAnalysisHandoff.get_example(), - ) - class OrchestrationRequestRecord(AgentDataModel): """Stored immutable request and effective configuration.""" recordType: Literal["automatedWorkflowRequest"] = "automatedWorkflowRequest" - formatVersion: Literal[1] = 1 + inputIdentity: dict[str, Any] + modelIdentity: str workflowRunId: str = "" createdAtNs: int = Field(default=0, ge=0) request: AutomatedWorkflowRequest = Field( @@ -483,47 +550,19 @@ class OrchestrationRequestRecord(AgentDataModel): configSha256: str = "" contentSha256: str = "" - @classmethod - def get_blank(cls) -> "OrchestrationRequestRecord": - return cls() - - @classmethod - def get_example(cls) -> "OrchestrationRequestRecord": - return cls( - workflowRunId="workflow-1", - createdAtNs=1, - request=AutomatedWorkflowRequest.get_example(), - config=AutomatedWorkflowConfig.get_example(), - requestSha256="0" * 64, - configSha256="1" * 64, - ) - class OrchestrationResumeRecord(AgentDataModel): - """One append-only set of answers supplied during resume.""" + """Runtime answers committed as inputs of their owning stage attempt.""" - recordType: Literal["automatedWorkflowResume"] = "automatedWorkflowResume" workflowRunId: str = "" - resumeId: str = "" - createdAtNs: int = Field(default=0, ge=0) answeredAttempt: WorkflowStageLink | None = None questionIds: list[str] = Field(default_factory=list) answers: dict[str, Any] = Field(default_factory=dict) - contentSha256: str = "" @classmethod def get_blank(cls) -> "OrchestrationResumeRecord": return cls() - @classmethod - def get_example(cls) -> "OrchestrationResumeRecord": - return cls( - workflowRunId="workflow-1", - resumeId="resume-1", - createdAtNs=1, - answers={"approvePlanChecksum": "0" * 64}, - ) - def artifact_model_to_ref(value: ArtifactReferenceModel) -> ArtifactRef: """Convert an agent artifact model to a validated core artifact reference.""" diff --git a/scarf/agent/orchestrator/preprocessing.py b/scarf/agent/orchestrator/preprocessing.py index de82328b..fd8181c8 100644 --- a/scarf/agent/orchestrator/preprocessing.py +++ b/scarf/agent/orchestrator/preprocessing.py @@ -1,44 +1,81 @@ """Preprocessing planning and execution stages.""" import hashlib -import re from collections.abc import Mapping, Sequence -from typing import Any, cast +from typing import Any import numpy as np +from ...assay import RNAassay from ...datastore.datastore import DataStore from ...datastore.summary import AssaySummary from ...metadata.selection import NamedCellArtifact from ...storage.refs import ArtifactRef from ...storage.selections import read_stored_selection_mask +from ...storage.types import as_zarr_array from ...utils.logging import logger from .. import record_io -from ..data_enrichment import ( +from ..cell_quality.execution import execute_auto_cell_qc, execute_registered_cell_qc +from ..cell_quality.profiles import CellQualityProfile, cell_qc_policy +from ..data_enrichment.contracts import ( AssayFeatureInspection, DataEnrichmentReport, FeatureSelectionPolicy, ) -from ..experimental_context import CellQcPlan, ExperimentalContextResult -from ..persistence import AgentWorkflowRun +from ..decisions.kernel import DecisionEvidence, DecisionSelection, EvidenceBundle +from ..decisions.rna import ( + CellQualityExecutorPayload, + QcGroupingExecutorPayload, + build_cell_quality_decision, + build_qc_grouping_decision, + require_option_evidence, +) +from ..experimental_context.contracts import ( + CellQcPlan, + CellQcProfileEvidence, + ExperimentalContextResult, +) +from ..experimental_context.study import StudyContract +from ..parameter_tuning.execution import ( + candidate_metric_cache, +) from ..types import ArtifactReferenceModel from . import journal +from .decisions import DecisionStagesMixin from .models import ( AssayPreprocessingPlan, AutomatedPreprocessingPlan, OrchestrationRequestRecord, OrchestrationResumeRecord, PreprocessedAssayHandoff, - ReductionMethod, + WorkflowIdentity, WorkflowNeedsInput, WorkflowQuestion, WorkflowStageAttempt, WorkflowStageLink, + WorkflowStageName, artifact_model_to_ref, ) +from .rna import ( + selected_store_rna_assay, + validate_rna_context, + validate_rna_handoffs, + validate_rna_plan, +) + + +class _DecisionNeedsInput(RuntimeError): + def __init__( + self, + question: WorkflowQuestion, + checkpoint_sha256: str, + ) -> None: + super().__init__("A registered RNA decision requires human input") + self.question = question + self.checkpointSha256 = checkpoint_sha256 -class PreprocessingStagesMixin: +class PreprocessingStagesMixin(DecisionStagesMixin): """Stages that plan and execute modality-specific preprocessing.""" @staticmethod @@ -60,19 +97,481 @@ def _cell_qc_stage_artifacts( artifacts[key] = plan.sampleArtifact.artifact return artifacts + @classmethod + def _cell_qc_candidate_artifacts( + cls, + profiles: Sequence[CellQcProfileEvidence], + ) -> dict[str, ArtifactReferenceModel]: + artifacts: dict[str, ArtifactReferenceModel] = {} + for profile in profiles: + plan = CellQcPlan( + action=profile.action, + registeredProfile=profile.registeredProfile, + profileId=profile.profileId, + driverAssay=profile.driverAssay, + driverAssayType=profile.driverAssayType, + sampleColumn=profile.sampleColumn, + sampleArtifact=profile.sampleArtifact, + attributes=profile.attributes, + artifactMetrics=profile.artifactMetrics, + evidenceIds=[profile.evidenceId], + ) + for key, artifact in cls._cell_qc_stage_artifacts(plan).items(): + existing = artifacts.get(key) + if existing is not None and existing != artifact: + raise ValueError( + f"Cell-QC candidate artifact key {key!r} is ambiguous" + ) + artifacts[key] = artifact + return artifacts + + @staticmethod + def _decision_evidence_bundle( + decision_id: str, + evidence: list[DecisionEvidence], + ) -> EvidenceBundle: + digest = hashlib.sha256( + record_io.canonical_json_bytes( + [item.model_dump(mode="json") for item in evidence] + ) + ).hexdigest() + return EvidenceBundle( + bundleId=f"bundle:{decision_id}:{digest[:24]}", + decisionId=decision_id, + evidence=evidence, + ).with_content_sha256() + + @staticmethod + def _profile_is_safe(profile: CellQcProfileEvidence) -> bool: + if profile.action == "sampleMad" and ( + profile.sampleColumn != profile.captureColumn + or profile.sampleArtifact != profile.captureArtifact + or (profile.captureColumn is None and profile.captureArtifact is None) + ): + return False + if profile.retainedCells == 0 or profile.unsafeRetentionGroups: + return False + if ( + profile.action == "sampleMad" + or profile.registeredProfile + in { + "captureMad5", + "captureMad3Sensitivity", + } + ) and profile.failedCaptureCandidates: + return False + if profile.registeredProfile == "pooledReferenceMad5" and set( + profile.failedCaptureCandidates + ).intersection(profile.parameters.get("pooledReferenceCaptures", [])): + return False + return True + + @staticmethod + def _qc_decision_evidence( + profiles: Sequence[CellQcProfileEvidence], + ) -> dict[str, Any]: + """Keep exact QC measurements while sharing repeated source/design evidence.""" + shared: dict[str, Any] = {} + + def reference(value: Any) -> str: + digest = hashlib.sha256(record_io.canonical_json_bytes(value)).hexdigest() + shared[digest] = value + return digest + + policies = [] + for profile in profiles: + policy = profile.model_dump(mode="json") + for field in ("metricSources", "sourceConcordance"): + policy[field + "Ref"] = reference(policy.pop(field)) + parameters = policy["parameters"] + if "captureComparisons" in parameters: + parameters["captureComparisonsRef"] = reference( + parameters.pop("captureComparisons") + ) + for duplicate, canonical in ( + ("resolvedBounds", "resolvedBounds"), + ("captureSizes", "activeCellsByCapture"), + ): + if ( + duplicate in parameters + and parameters[duplicate] == policy[canonical] + ): + parameters.pop(duplicate) + for capture in policy["captureFailureEvidence"]: + capture["conditionAndUnitSafetyRef"] = reference( + capture.pop("conditionAndUnitSafety") + ) + policies.append(policy) + return { + "policies": policies, + "sharedMeasurements": shared, + "interpretation": ( + "Each Ref identifies the exact record in sharedMeasurements. " + "Thresholds, distributions, flags and protected-group retention are " + "measured policy evidence; short evidence summaries do not replace them." + ), + } + + @staticmethod + def _profile_evidence( + profile: CellQcProfileEvidence, + retained_reference: CellQcProfileEvidence | None = None, + ) -> DecisionEvidence: + def retention_range(values: Sequence[int]) -> str: + if not values: + return "no groups" + ordered = sorted(int(value) for value in values) + return ( + f"{len(ordered)} groups, min/median/max=" + f"{ordered[0]}/{float(np.median(ordered)):g}/{ordered[-1]}" + ) + + capture_summary = retention_range(list(profile.sampleRetainedCells.values())) + capture_fractions = [ + row.retainedFraction for row in profile.captureFailureEvidence + ] + capture_fraction_summary = ( + f"min/median/max={min(capture_fractions):.1%}/" + f"{float(np.median(capture_fractions)):.1%}/{max(capture_fractions):.1%}" + if capture_fractions + else "not available" + ) + active_by_column = ( + retained_reference.retainedCellsByColumn + if retained_reference is not None + and retained_reference.retainedCells + == retained_reference.activeCells + == profile.activeCells + else {} + ) + + def design_retention(column: str, groups: Mapping[str, int]) -> str: + active = active_by_column.get(column, {}) + if not groups or any(not active.get(group) for group in groups): + return ( + retention_range(list(groups.values())) + "; fractions unavailable" + ) + if len(groups) <= 4: + return ", ".join( + f"{group}={retained}/{active[group]} ({retained / active[group]:.1%})" + for group, retained in sorted(groups.items()) + ) + fractions = [retained / active[group] for group, retained in groups.items()] + return ( + f"{retention_range(list(groups.values()))}, retained fractions " + f"min/median/max={min(fractions):.1%}/{float(np.median(fractions)):.1%}/" + f"{max(fractions):.1%}" + ) + + design_summary = "; ".join( + f"{column}: {design_retention(column, groups)}" + for column, groups in sorted(profile.retainedCellsByColumn.items()) + ) + summary = ( + f"{cell_qc_policy(profile.action, profile.registeredProfile) or profile.action} retains " + f"{profile.retainedCells}/{profile.activeCells} active cells " + f"({profile.retainedCells / profile.activeCells:.1%}); " + if profile.activeCells + else "No active cells; " + ) + ( + f"metric flags={profile.metricFlaggedCells or 'not available'} " + "(flags may overlap; high counts/features can be retained); " + f"capture retention={capture_summary}; capture retained fractions=" + f"{capture_fraction_summary}; design retention=" + f"{design_summary or 'no groups'}; failed capture candidates=" + f"{profile.failedCaptureCandidates}; unsafe retention groups=" + f"{profile.unsafeRetentionGroups}; limitations={profile.notes[:4]}." + ) + if len(summary) > 2_000: + summary = f"{summary[:1_997].rstrip()}..." + return DecisionEvidence( + evidenceId=profile.evidenceId, + evidenceClass="qualityControl", + summary=summary, + artifactReferences=[ + *[source.artifact for source in profile.artifactMetrics], + *( + [profile.sampleArtifact.artifact] + if profile.sampleArtifact is not None + else [] + ), + ], + ) + + def _resolve_qc_grouping_decision( + self, + store: DataStore, + request_record: OrchestrationRequestRecord, + experimental: ExperimentalContextResult, + study_contract: StudyContract, + answers: Mapping[str, Any], + ) -> tuple[QcGroupingExecutorPayload, str]: + profiles = [ + profile + for profile in experimental.qcProfiles + if cell_qc_policy(profile.action, profile.registeredProfile) is not None + ] + if not profiles: + raise ValueError("RNA decision workflow requires executable QC evidence") + safe_profiles = { + cell_qc_policy(profile.action, profile.registeredProfile): profile + for profile in profiles + if self._profile_is_safe(profile) + } + capture_eligible = bool( + study_contract.physicalCaptureColumn is not None + and any( + policy in safe_profiles for policy in ("coreSampleMad3", "captureMad5") + ) + ) + pooled_eligible = bool( + capture_eligible and "pooledReferenceMad5" in safe_profiles + ) + design_id = "evidence:qcGrouping:studyContract" + retained_reference = safe_profiles.get("retainWithFlags") + evidence = [ + DecisionEvidence( + evidenceId=design_id, + evidenceClass="design", + summary=( + "The validated physical capture is " + f"{study_contract.physicalCaptureColumn!r}; independent units=" + f"{study_contract.independentUnitColumns}; conditions=" + f"{study_contract.conditionColumns}. " + "Within-capture QC estimates technical quality boundaries; captures " + "need not be independent biological units or healthy references. " + "Healthy-reference provenance is required only for a pooled reference. " + "Compare global and capture five-MAD profiles when both are supplied; " + "different cutoff methods cannot isolate the effect of grouping." + ), + ) + ] + mode_profile: dict[str, CellQcProfileEvidence] = {} + matched_mad = "globalMad5" in safe_profiles and "captureMad5" in safe_profiles + global_profile = ( + safe_profiles["globalMad5"] + if matched_mad + else safe_profiles.get("coreGlobalGaussian") + or safe_profiles.get("globalMad5") + or safe_profiles.get("retainWithFlags") + ) + if global_profile is None: + raise ValueError("No safe global or retain-only QC profile is available") + mode_profile["qcGrouping:global"] = global_profile + baseline = next( + (profile for profile in profiles if profile.action == "globalGaussian"), + None, + ) + if baseline is not None and baseline != global_profile: + evidence.append(self._profile_evidence(baseline, retained_reference)) + evidence.append(self._profile_evidence(global_profile, retained_reference)) + if capture_eligible: + capture_profile = ( + safe_profiles["captureMad5"] + if matched_mad + else safe_profiles.get("coreSampleMad3") or safe_profiles["captureMad5"] + ) + mode_profile["qcGrouping:physicalCapture"] = capture_profile + evidence.append(self._profile_evidence(capture_profile, retained_reference)) + if pooled_eligible: + pooled_profile = safe_profiles["pooledReferenceMad5"] + mode_profile["qcGrouping:pooledReference"] = pooled_profile + evidence.append(self._profile_evidence(pooled_profile, retained_reference)) + bundle = self._decision_evidence_bundle("qcGrouping", evidence) + definition = build_qc_grouping_decision( + evidence_bundle_id=bundle.bundleId, + physical_capture_eligible=capture_eligible, + pooled_reference_eligible=pooled_eligible, + ) + definition = require_option_evidence( + definition, + { + option_id: [design_id, profile.evidenceId] + for option_id, profile in mode_profile.items() + }, + ) + selectable_ids = [ + option.optionId + for option in definition.spec.options + if option.status != "defer" + ] + rule_selection = ( + DecisionSelection( + selectedOptionId=selectable_ids[0], + evidenceIds=list( + definition.spec.option_by_id()[ + selectable_ids[0] + ].requiredEvidenceIds + ), + rationale=( + "Use the only grouping mode licensed by the validated design " + "and retention evidence." + ), + ) + if len(selectable_ids) == 1 + else None + ) + resolution = self._resolve_rna_decision( + store, + request_record, + definition, + bundle, + answers, + rule_selection=rule_selection, + qc_evidence=self._qc_decision_evidence(profiles), + ) + if resolution.compiled is None: + raise _DecisionNeedsInput( + self._pending_decision_question(resolution, definition), + resolution.checkpointSha256, + ) + payload = resolution.compiled.executorPayload + if not isinstance(payload, QcGroupingExecutorPayload): + raise TypeError("QC-grouping decision compiled an unexpected payload") + return payload, resolution.checkpointSha256 + + def _resolve_cell_quality_decision( + self, + store: DataStore, + request_record: OrchestrationRequestRecord, + experimental: ExperimentalContextResult, + grouping: QcGroupingExecutorPayload, + answers: Mapping[str, Any], + ) -> tuple[CellQualityExecutorPayload, CellQcPlan, str]: + all_profiles = [ + profile + for profile in experimental.qcProfiles + if cell_qc_policy(profile.action, profile.registeredProfile) is not None + ] + all_profiles.sort(key=lambda profile: profile.action != "globalGaussian") + allowed_by_grouping: dict[str, set[CellQualityProfile]] = { + "global": {"retainWithFlags", "coreGlobalGaussian", "globalMad5"}, + "physicalCapture": {"retainWithFlags", "coreSampleMad3", "captureMad5"}, + "pooledReference": {"retainWithFlags", "pooledReferenceMad5"}, + } + allowed = allowed_by_grouping[grouping.groupingMode] + profiles = [ + profile + for profile in all_profiles + if cell_qc_policy(profile.action, profile.registeredProfile) in allowed + and self._profile_is_safe(profile) + ] + if not profiles: + raise ValueError("QC grouping has no safe executable profile") + profiles.sort(key=lambda profile: profile.action != "globalGaussian") + retained_reference = next( + (p for p in all_profiles if p.registeredProfile == "retainWithFlags"), None + ) + evidence = [ + self._profile_evidence(profile, retained_reference) + for profile in all_profiles + ] + + bundle = self._decision_evidence_bundle("cellQuality", evidence) + available_profiles = [ + policy + for profile in profiles + if (policy := cell_qc_policy(profile.action, profile.registeredProfile)) + is not None + ] + definition = build_cell_quality_decision( + evidence_bundle_id=bundle.bundleId, + available_profiles=available_profiles, + ) + definition = require_option_evidence( + definition, + { + f"cellQuality:{cell_qc_policy(profile.action, profile.registeredProfile)}": [ + profile.evidenceId + ] + for profile in profiles + }, + ) + resolution = self._resolve_rna_decision( + store, + request_record, + definition, + bundle, + answers, + qc_evidence=self._qc_decision_evidence(all_profiles), + ) + if resolution.compiled is None: + raise _DecisionNeedsInput( + self._pending_decision_question(resolution, definition), + resolution.checkpointSha256, + ) + payload = resolution.compiled.executorPayload + if not isinstance(payload, CellQualityExecutorPayload): + raise TypeError("Cell-quality decision compiled an unexpected payload") + selected = next( + ( + profile + for profile in profiles + if cell_qc_policy(profile.action, profile.registeredProfile) + == payload.profile + ), + None, + ) + if selected is None or resolution.record is None: + raise ValueError("Audited cell-quality option lacks its exact profile") + plan = CellQcPlan( + action=selected.action, + registeredProfile=selected.registeredProfile, + profileId=selected.profileId, + driverAssay=selected.driverAssay, + driverAssayType=selected.driverAssayType, + sampleColumn=selected.sampleColumn, + sampleArtifact=selected.sampleArtifact, + attributes=selected.attributes, + artifactMetrics=selected.artifactMetrics, + rationale=resolution.record.rationale, + evidenceIds=list(resolution.record.evidenceIds), + ) + return payload, plan, resolution.checkpointSha256 + + @staticmethod + def _plan_with_selected_hvg_counts( + plan: AutomatedPreprocessingPlan, + handoffs: Sequence[PreprocessedAssayHandoff], + ) -> AutomatedPreprocessingPlan: + selected_counts = {handoff.assay: handoff.nFeatures for handoff in handoffs} + assays = [ + assay.model_copy( + update={ + "featureParameters": { + **assay.featureParameters, + "topN": selected_counts[assay.assay], + } + } + ) + if assay.featureMethod == "hvg" and assay.assay in selected_counts + else assay + for assay in plan.assays + ] + updated = plan.model_copy(update={"assays": assays, "planChecksum": ""}) + checksum = hashlib.sha256( + record_io.canonical_json_bytes( + updated.model_dump(mode="json", exclude={"planChecksum"}) + ) + ).hexdigest() + return updated.model_copy(update={"planChecksum": checksum}) + def preprocessing_plan_stage( self, store: DataStore, - workflow: AgentWorkflowRun, + workflow: WorkflowIdentity, request_record: OrchestrationRequestRecord, parents: Sequence[WorkflowStageLink], enrichment: DataEnrichmentReport, experimental: ExperimentalContextResult, ingest_outcome: WorkflowStageAttempt, + study_contract: StudyContract, answers: Mapping[str, Any], *, resume_record: OrchestrationResumeRecord | None = None, ) -> tuple[WorkflowStageAttempt, AutomatedPreprocessingPlan]: + selected = selected_store_rna_assay(store, request_record.request) + validate_rna_context(experimental, selected) prefix = journal._ensure_orchestration_store(store) existing = journal._validated_done_outcome( store, @@ -83,15 +582,17 @@ def preprocessing_plan_stage( parents, ) if existing is not None: - logger.info( + logger.debug( f"Workflow {workflow.workflowRunId}: reusing preprocessing plan" ) - return existing, AutomatedPreprocessingPlan.model_validate( + cached_plan = AutomatedPreprocessingPlan.model_validate( existing.outputs["preprocessingPlan"] ) + validate_rna_plan(cached_plan, selected) + return existing, cached_plan if experimental.cellSelection is None: raise ValueError("Experimental Context lacks an exact cell selection") - cell_qc_artifacts = self._cell_qc_stage_artifacts(experimental.cellQc) + cell_qc_artifacts = self._cell_qc_candidate_artifacts(experimental.qcProfiles) started = journal._start_attempt( store.zw, prefix, @@ -100,29 +601,81 @@ def preprocessing_plan_stage( request_record, parents, inputs={ - "approvalAnswer": answers.get("approvePlanChecksum"), - "allowAssumptions": request_record.request.allowAssumptions, "cellSelection": experimental.cellSelection.model_dump(mode="json"), + "decisionPolicy": "evidenceBoundedAutomaticExecution", + "studyContract": study_contract.model_dump(mode="json"), }, resume_record=resume_record, ) try: + grouping_payload, grouping_decision_snapshot = ( + self._resolve_qc_grouping_decision( + store, + request_record, + experimental, + study_contract, + answers, + ) + ) + cell_payload, cell_qc, cell_decision_snapshot = ( + self._resolve_cell_quality_decision( + store, + request_record, + experimental, + grouping_payload, + answers, + ) + ) plan = self.build_preprocessing_plan( store, request_record, enrichment, experimental, ingest_outcome, + cell_qc, ) + validate_rna_plan(plan, selected) + plan = plan.model_copy(update={"cellQualityPayload": cell_payload}) + # The exact core feature policy is a provisional baseline. Its + # scientific assessment follows observed PCA and marker evidence. route_summary = ", ".join( f"{value.assay}:{value.featureMethod}/{value.reductionMethod}" for value in plan.assays ) - logger.info( + logger.debug( f"Workflow {workflow.workflowRunId}: preprocessing plan built " f"(primary={plan.primaryAssay!r}, marker={plan.markerAssay!r}, " f"routes=[{route_summary}])" ) + except _DecisionNeedsInput as pending: + if request_record.config.inputPolicy == "unattended": + outcome = journal._complete_attempt( + started, + status="failed", + artifacts={ + "cellSelection": experimental.cellSelection, + **cell_qc_artifacts, + }, + outputs={"decisionCheckpointSha256": pending.checkpointSha256}, + error=( + "The unattended preprocessing plan returned an unresolved " + "registered decision" + ), + ) + else: + outcome = journal._complete_attempt( + started, + status="needsInput", + artifacts={ + "cellSelection": experimental.cellSelection, + **cell_qc_artifacts, + }, + outputs={"decisionCheckpointSha256": pending.checkpointSha256}, + needs_input=WorkflowNeedsInput(questions=[pending.question]), + notes=["A registered filtering decision requires input."], + ) + journal._save_outcome(store.zw, prefix, outcome) + return outcome, AutomatedPreprocessingPlan.get_blank() except Exception as exc: outcome = journal.finish_exception( store, @@ -136,57 +689,28 @@ def preprocessing_plan_stage( }, ) return outcome, AutomatedPreprocessingPlan.get_blank() - supplied_approval = answers.get("approvePlanChecksum") - preapproved = bool( - request_record.request.experimentalDirections.get( - "approveAutomatedAnalysis", False - ) + logger.debug( + f"Workflow {workflow.workflowRunId}: executing the evidence-bounded " + "preprocessing plan" ) - approved = ( - request_record.request.allowAssumptions - or preapproved - or supplied_approval == plan.planChecksum + outcome = journal._complete_attempt( + started, + status="done", + artifacts={ + "cellSelection": experimental.cellSelection, + **cell_qc_artifacts, + }, + outputs={ + "preprocessingPlan": plan.model_dump(mode="json"), + "qcGroupingDecisionCheckpoint": grouping_decision_snapshot, + "cellQualityDecisionCheckpoint": cell_decision_snapshot, + }, + actions=[ + "audit_qc_grouping_decision", + "audit_cell_quality_decision", + "accept_evidence_bounded_preprocessing_plan", + ], ) - if not approved: - logger.info( - f"Workflow {workflow.workflowRunId}: preprocessing plan requires " - "caller approval" - ) - outcome = journal._complete_attempt( - started, - status="needsInput", - artifacts={ - "cellSelection": experimental.cellSelection, - **cell_qc_artifacts, - }, - outputs={"preprocessingPlan": plan.model_dump(mode="json")}, - needs_input=WorkflowNeedsInput( - questions=[ - WorkflowQuestion( - questionId="approvePlanChecksum", - question=( - "Approve the persisted preprocessing and bounded " - "parameter-tuning plan?" - ), - planChecksum=plan.planChecksum, - ) - ] - ), - ) - else: - logger.info( - f"Workflow {workflow.workflowRunId}: preprocessing plan approved" - ) - outcome = journal._complete_attempt( - started, - status="done", - artifacts={ - "cellSelection": experimental.cellSelection, - **cell_qc_artifacts, - }, - outputs={"preprocessingPlan": plan.model_dump(mode="json")}, - actions=["approve_preprocessing_plan"], - ) journal._save_outcome(store.zw, prefix, outcome) return outcome, plan @@ -197,133 +721,38 @@ def build_preprocessing_plan( enrichment: DataEnrichmentReport, experimental: ExperimentalContextResult, ingest_outcome: WorkflowStageAttempt, + cell_qc: CellQcPlan, ) -> AutomatedPreprocessingPlan: + del ingest_outcome request = request_record.request store_summary = store.summary() - summaries = {value.name: value for value in store_summary.assays} - policies = {value.assay: value for value in enrichment.policies} - inspections = {value.assay: value for value in enrichment.inspections} - selected_names = request.analysisAssays or list(store.assay_names) - assay_plans: list[AssayPreprocessingPlan] = [] - graph_assays: list[str] = [] - limitations: list[str] = list(enrichment.limitations) - selected_qc_profile = next( - ( - value - for value in experimental.qcProfiles - if value.profileId == experimental.cellQc.profileId - ), - None, + selected = selected_store_rna_assay(store, request) + summary = next( + value for value in store_summary.assays if value.name == selected ) - projected_cells = ( - selected_qc_profile.retainedCells - if selected_qc_profile is not None and selected_qc_profile.retainedCells > 0 - else store_summary.active_cells + policy = next( + (value for value in enrichment.policies if value.assay == selected), None ) - effective_min_cells = min(20, max(1, projected_cells // 10)) - for assay_name in selected_names: - summary = summaries[assay_name] - policy = policies.get(assay_name) - modality = ( - policy.assayModality - if policy is not None - else ( - summary.assay_type - if summary.assay_type in {"RNA", "ATAC", "ADT", "HTO"} - else "unsupported" - ) - ) - plan = self.build_assay_preprocessing_plan( - store, - request_record, - assay_name, - summary, - policy, - inspections.get(assay_name), - modality, - effective_min_cells, - ) - if plan.graphEligible: - graph_assays.append(assay_name) - if modality == "unsupported": - limitations.extend(plan.limitations) - assay_plans.append(plan) - if not graph_assays: - raise ValueError("No supported graph-bearing assay remains") - if len(graph_assays) > request_record.config.maxGraphAssays: - raise ValueError( - "Too many graph-bearing assays; provide analysisAssays to select at " - f"most {request_record.config.maxGraphAssays}" - ) - modality_counts: dict[str, int] = {} - for name in graph_assays: - modality_counts[summaries[name].assay_type] = ( - modality_counts.get(summaries[name].assay_type, 0) + 1 - ) - duplicate_modalities = sorted( - name for name, count in modality_counts.items() if count > 1 + inspection = next( + (value for value in enrichment.inspections if value.assay == selected), None ) - if duplicate_modalities and not request.analysisAssays: - raise ValueError( - "Multiple same-kind biological assays require explicit " - f"analysisAssays selection: {duplicate_modalities}" - ) - primary = request.primaryAssay - if primary is not None and primary not in graph_assays: - raise ValueError("primaryAssay must name a graph-bearing selected assay") - if primary is None: - primary = next( - ( - name - for modality in ("RNA", "ADT", "ATAC") - for name in graph_assays - if summaries[name].assay_type == modality - ), - graph_assays[0], - ) - if request.markerAssay is not None: - if request.markerAssay not in graph_assays: - raise ValueError("markerAssay must name a graph-bearing selected assay") - marker_assay = request.markerAssay - else: - marker_assay = next( - ( - name - for modality in ("RNA", "ADT", "ATAC") - for name in graph_assays - if summaries[name].assay_type == modality - ), - primary, - ) - if request.pairedAssays: - paired = list(request.pairedAssays) - unknown_paired = sorted(set(paired) - set(graph_assays)) - if unknown_paired: - raise ValueError( - f"pairedAssays contains non-graph assays: {unknown_paired}" - ) - if primary not in paired: - raise ValueError("pairedAssays must include the primary assay") - elif ( - len(graph_assays) > 1 - and ingest_outcome.outputs.get("pairingProvenance") - == "singleSourceSharedCellAxis" - ): - paired = list(graph_assays) - else: - paired = [] - if len(graph_assays) > 1: - limitations.append( - "Multimodal integration skipped because pairing provenance was not supplied" - ) + if policy is not None and policy.assayModality != "RNA": + raise ValueError("Enrichment policy does not match the selected RNA assay") + assay_plan = self.build_assay_preprocessing_plan( + selected, + summary, + policy, + inspection, + ) + if not assay_plan.graphEligible: + raise ValueError("RNA requires at least three features for PCA") final_plan = AutomatedPreprocessingPlan( - primaryAssay=primary, - markerAssay=marker_assay, + primaryAssay=selected, + markerAssay=selected, cellSelection=experimental.cellSelection, - cellQc=experimental.cellQc, - assays=assay_plans, - pairedAssays=paired, - limitations=list(dict.fromkeys(limitations)), + cellQc=cell_qc, + assays=[assay_plan], + limitations=list(dict.fromkeys(enrichment.limitations)), ) checksum = hashlib.sha256( record_io.canonical_json_bytes( @@ -334,227 +763,113 @@ def build_preprocessing_plan( def build_assay_preprocessing_plan( self, - store: DataStore, - request_record: OrchestrationRequestRecord, assay_name: str, summary: AssaySummary, policy: FeatureSelectionPolicy | None, inspection: AssayFeatureInspection | None, - modality: str, - effective_min_cells: int, ) -> AssayPreprocessingPlan: - excluded: list[str] = [] - evidence_ids: list[str] = [] - if policy is not None: - excluded = list( - dict.fromkeys( - [ - *policy.excludeFeatures, - *policy.artificialFeatures, - *( - reference.featureId - for reference in policy.exactControlFeatures - ), - *( - reference.featureName - for reference in policy.exactControlFeatures - ), - ] - ) - ) - evidence_ids = list(policy.evidenceIds) - if modality == "RNA": - graph_eligible = summary.total_features >= 3 - return AssayPreprocessingPlan( - assay=assay_name, - assayType=summary.assay_type, - role="graph" if graph_eligible else "unsupported", - graphEligible=graph_eligible, - markerEligible=graph_eligible, - featureMethod="hvg" if graph_eligible else "none", - reductionMethod="pca" if graph_eligible else "none", - featureParameters={ - "topN": min(1000, summary.total_features), - "minCells": effective_min_cells, - "excludeFamilies": ( - list(policy.excludeFamilies) if policy is not None else [] - ), - }, - normalizationParameters={ - "logTransform": True, - "renormalizeSubset": True, - }, - reductionParameters={"dimensions": 21}, - exactExcludedFeatures=excluded, - evidenceIds=evidence_ids, - limitations=( - [] - if graph_eligible - else ["RNA requires at least three features for PCA"] - ), - ) - if modality == "ATAC": - graph_eligible = summary.total_features >= 3 - return AssayPreprocessingPlan( - assay=assay_name, - assayType=summary.assay_type, - role="graph" if graph_eligible else "unsupported", - graphEligible=graph_eligible, - markerEligible=graph_eligible, - featureMethod="prevalentPeaks" if graph_eligible else "none", - reductionMethod="lsi" if graph_eligible else "none", - featureParameters={"topN": min(25000, summary.total_features)}, - normalizationParameters={ - "logTransform": False, - "renormalizeSubset": False, - }, - reductionParameters={"dimensions": 50, "skipFirst": True}, - evidenceIds=evidence_ids, - limitations=list( - dict.fromkeys( - [ - *( - [] - if graph_eligible - else [ - "ATAC requires at least three peak features for LSI" - ] - ), - *( - [ - "ATAC feature coordinates are not uniformly " - "valid chrom:start-end intervals; the genome " - "build remains unknown" - ] - if policy is not None - and policy.peakCoordinateStatus - in {"partial", "invalid"} - else [] - ), - ] - ) + if summary.assay_type != "RNA": + raise ValueError("Automated preprocessing supports RNA only") + evidence_ids = list(policy.evidenceIds) if policy is not None else [] + graph_eligible = summary.total_features >= 3 + proposed_families = list(policy.excludeFamilies) if policy is not None else [] + return AssayPreprocessingPlan( + assay=assay_name, + assayType=summary.assay_type, + role="graph" if graph_eligible else "unsupported", + graphEligible=graph_eligible, + markerEligible=graph_eligible, + featureMethod="hvg" if graph_eligible else "none", + reductionMethod="pca" if graph_eligible else "none", + featureParameters={ + "excludeFamilies": [], + "useScarfDefaultBlacklist": True, + "proposedExcludeFamilies": proposed_families, + "protectFamilies": ( + list(policy.protectFamilies) if policy is not None else [] ), - ) - if modality == "ADT": - assay = store.get_assay(assay_name) - feature_ids = np.asarray(assay.feats.fetch_all("ids")).astype(str) - feature_names = np.asarray(assay.feats.fetch_all("names")).astype(str) - excluded = ( - list( - dict.fromkeys( - value - for reference in policy.exactControlFeatures - for value in (reference.featureId, reference.featureName) - if value - ) + "protectFeatures": list(policy.protectFeatures) + if policy is not None + else [], + "proposedExcludeFeatures": list( + dict.fromkeys([*policy.excludeFeatures, *policy.artificialFeatures]) ) if policy is not None - else [] - ) - excluded_values = {value for value in excluded if value} - panel_mask = ~np.isin(feature_ids, list(excluded_values)) - panel_mask &= ~np.isin(feature_names, list(excluded_values)) - selected_count = int(panel_mask.sum()) - graph_eligible = selected_count >= 2 - reduction = ( - "identity" - if graph_eligible - and selected_count <= request_record.config.maxIdentityFeatures - else ("pca" if graph_eligible else "none") - ) - return AssayPreprocessingPlan( - assay=assay_name, - assayType=summary.assay_type, - role="graph" if graph_eligible else "unsupported", - graphEligible=graph_eligible, - markerEligible=graph_eligible, - featureMethod="panel" if graph_eligible else "none", - reductionMethod=cast(ReductionMethod, reduction), - normalizationParameters={ - "logTransform": False, - "renormalizeSubset": False, - }, - reductionParameters={ - "dimensions": ( - selected_count - if reduction == "identity" - else min(15, max(2, selected_count - 1)) - ) - }, - exactExcludedFeatures=excluded, - evidenceIds=evidence_ids, - limitations=( - [ - "ADT control inventory was truncated; only exact observed " - "control features were excluded" - ] - if inspection is not None and inspection.modalityEvidence.truncated - else [] + else [], + "species": ( + inspection.species if inspection is not None else "unknown" ), - ) - if modality == "HTO": - return AssayPreprocessingPlan( - assay=assay_name, - assayType=summary.assay_type, - role="hto", - graphEligible=False, - markerEligible=False, - featureMethod="none", - reductionMethod="none", - evidenceIds=evidence_ids, - ) - message = f"Unsupported assay {assay_name!r} ({summary.assay_type})" - return AssayPreprocessingPlan( - assay=assay_name, - assayType=summary.assay_type, - role="unsupported", - limitations=[message], + "defaultFeatureInventory": ( + inspection.defaultFeatureInventory.model_dump(mode="json") + if inspection is not None + and inspection.defaultFeatureInventory is not None + else None + ), + }, + evidenceIds=evidence_ids, + limitations=( + [] + if graph_eligible + else ["RNA requires at least three features for PCA"] + ), ) def preprocessing_stage( self, store: DataStore, - workflow: AgentWorkflowRun, + workflow: WorkflowIdentity, request_record: OrchestrationRequestRecord, parents: Sequence[WorkflowStageLink], plan: AutomatedPreprocessingPlan, experimental: ExperimentalContextResult, + study_contract: StudyContract, + answers: Mapping[str, Any], *, resume_record: OrchestrationResumeRecord | None = None, - ) -> tuple[WorkflowStageAttempt, list[PreprocessedAssayHandoff]]: + stage_name: WorkflowStageName = "preprocessing", + ) -> tuple[ + WorkflowStageAttempt, + list[PreprocessedAssayHandoff], + AutomatedPreprocessingPlan, + ]: + selected = selected_store_rna_assay(store, request_record.request) + validate_rna_plan(plan, selected) + validate_rna_context(experimental, selected) prefix = journal._ensure_orchestration_store(store) existing = journal._validated_done_outcome( store, prefix, workflow.workflowRunId, - "preprocessing", + stage_name, request_record, parents, ) if existing is not None: - logger.info( + logger.debug( f"Workflow {workflow.workflowRunId}: reusing preprocessing artifacts" ) - return existing, [ + cached_handoffs = [ PreprocessedAssayHandoff.model_validate(value) for value in existing.outputs["assays"] ] + cached_plan = AutomatedPreprocessingPlan.model_validate( + existing.outputs["resolvedPreprocessingPlan"] + ) + validate_rna_plan(cached_plan, selected) + validate_rna_handoffs(cached_handoffs, selected) + return existing, cached_handoffs, cached_plan if plan.cellSelection is None: raise ValueError("Preprocessing plan lacks an exact cell selection") if experimental.cellSelection != plan.cellSelection: raise ValueError( "Preprocessing plan and Experimental Context selections differ" ) - if experimental.cellQc != plan.cellQc: - raise ValueError( - "Preprocessing plan and Experimental Context cell-QC plans differ" - ) input_cell_selection = artifact_model_to_ref(plan.cellSelection) started = journal._start_attempt( store.zw, prefix, workflow.workflowRunId, - "preprocessing", + stage_name, request_record, parents, inputs={ @@ -570,17 +885,30 @@ def preprocessing_stage( **self._cell_qc_stage_artifacts(plan.cellQc), } try: + if plan.cellQualityPayload is None: + raise ValueError( + "Decision-driven preprocessing requires an audited " + "cell-quality payload" + ) cell_selection = self.apply_cell_qc( store, experimental, input_cell_selection, actions, operations, + selected_plan=plan.cellQc, + decision_payload=plan.cellQualityPayload, ) cell_selection_model = ArtifactReferenceModel.from_artifact_ref( cell_selection ) artifacts["cellSelection"] = cell_selection_model + if operations: + diagnostic_flags = operations[-1].get("diagnosticFlags") + if diagnostic_flags is not None: + artifacts["cellQcDiagnosticFlags"] = ( + ArtifactReferenceModel.model_validate(diagnostic_flags) + ) active_cells = int( read_stored_selection_mask( store.zw, @@ -595,7 +923,7 @@ def preprocessing_stage( ( profile for profile in experimental.qcProfiles - if profile.profileId == experimental.cellQc.profileId + if profile.profileId == plan.cellQc.profileId ), None, ) @@ -606,32 +934,41 @@ def preprocessing_stage( "Executed cell-QC retention differs from the selected profile" ) logger.info( - f"Workflow {workflow.workflowRunId}: preprocessing retained " - f"{active_cells} active cell(s)" + f"QC: compared {len(experimental.qcProfiles)} policies; retained " + f"{active_cells:,}/{selected_profile.activeCells:,} cells " + f"({selected_profile.retainedFraction:.0%})." ) if active_cells < 3: raise ValueError("Preprocessing requires at least three active cells") handoffs: list[PreprocessedAssayHandoff] = [] - for assay_plan in plan.assays: - if not assay_plan.graphEligible: - continue - logger.info( - f"Workflow {workflow.workflowRunId}: preprocessing assay " - f"{assay_plan.assay!r} via {assay_plan.featureMethod}/" - f"{assay_plan.reductionMethod}" - ) - handoffs.append( - self.preprocess_assay( - store, - assay_plan, - cell_selection=cell_selection, - cell_selection_model=cell_selection_model, - active_cells=active_cells, - actions=actions, - operations=operations, - artifacts=artifacts, + with candidate_metric_cache(): + for assay_plan in plan.assays: + if not assay_plan.graphEligible: + continue + logger.debug( + f"Workflow {workflow.workflowRunId}: preprocessing assay " + f"{assay_plan.assay!r} via {assay_plan.featureMethod}/" + f"{assay_plan.reductionMethod}" ) - ) + handoffs.append( + self.preprocess_assay( + store, + assay_plan, + cell_selection=cell_selection, + cell_selection_model=cell_selection_model, + active_cells=active_cells, + request_record=request_record, + study_contract=study_contract, + answers=answers, + actions=actions, + operations=operations, + artifacts=artifacts, + ) + ) + resolved_plan = self._plan_with_selected_hvg_counts( + plan, + handoffs, + ) outcome = journal._complete_attempt( started, status="done", @@ -643,16 +980,56 @@ def preprocessing_stage( outputs={ "assays": [value.model_dump(mode="json") for value in handoffs], "cellSelection": cell_selection_model.model_dump(mode="json"), + "resolvedPreprocessingPlan": resolved_plan.model_dump(mode="json"), "operations": operations, }, actions=actions, ) journal._save_outcome(store.zw, prefix, outcome) - logger.info( + logger.debug( f"Workflow {workflow.workflowRunId}: preprocessing produced " f"{len(handoffs)} graph-ready assay handoff(s)" ) - return outcome, handoffs + return outcome, handoffs, resolved_plan + except _DecisionNeedsInput as pending: + if request_record.config.inputPolicy == "unattended": + outcome = journal._complete_attempt( + started, + status="failed", + artifacts={ + name: value + for name, value in artifacts.items() + if value is not None + }, + outputs={ + "operations": operations, + "decisionCheckpointSha256": pending.checkpointSha256, + }, + actions=actions, + error=( + "The unattended preprocessing stage returned an unresolved " + "registered decision" + ), + ) + else: + outcome = journal._complete_attempt( + started, + status="needsInput", + artifacts={ + name: value + for name, value in artifacts.items() + if value is not None + }, + outputs={ + "operations": operations, + "decisionCheckpointSha256": pending.checkpointSha256, + }, + needs_input=WorkflowNeedsInput(questions=[pending.question]), + actions=actions, + notes=["A registered RNA preprocessing decision requires input."], + ) + journal._save_outcome(store.zw, prefix, outcome) + return outcome, [], plan except Exception as exc: outcome = journal.finish_exception( store, @@ -664,7 +1041,7 @@ def preprocessing_stage( actions=actions, outputs={"operations": operations}, ) - return outcome, [] + return outcome, [], plan def preprocess_assay( self, @@ -674,233 +1051,74 @@ def preprocess_assay( cell_selection: ArtifactRef, cell_selection_model: ArtifactReferenceModel, active_cells: int, + request_record: OrchestrationRequestRecord | None = None, + study_contract: StudyContract | None = None, + answers: Mapping[str, Any] | None = None, actions: list[str], operations: list[dict[str, Any]], artifacts: dict[str, ArtifactReferenceModel], ) -> PreprocessedAssayHandoff: + """Prepare exact core defaults and feature evidence without a graph search.""" + del request_record, study_contract, answers assay = store.get_assay(assay_plan.assay) - min_cells = int(assay_plan.featureParameters.get("minCells", 1)) - marker_features: ArtifactRef - if assay_plan.featureMethod == "hvg": - blacklist = self.rna_blacklist(assay_plan) - actual_top_n = min( - int(assay_plan.featureParameters["topN"]), - max(1, assay.feats.N - 1), - ) - hvg_features = store.select_hvgs( - cell_selection, - from_assay=assay_plan.assay, - min_cells=min_cells, - top_n=actual_top_n, - blacklist=blacklist, - show_plot=False, - invalidate_cache=False, - ) - graph_features = self.exclude_exact_features( - store, - assay_plan, - hvg_features, - ) - detected = store.select_detected_features( - cell_selection, - from_assay=assay_plan.assay, - min_cells=min_cells, - invalidate_cache=False, - ) - marker_features = self.exclude_exact_features( - store, - assay_plan, - detected, - ) - actions.extend( - [ - f"select_hvgs:{assay_plan.assay}", - f"select_marker_features:{assay_plan.assay}", - ] - ) - operations.append( - { - "operation": "select_hvgs", - "assay": assay_plan.assay, - "cellSelection": cell_selection_model.model_dump(mode="json"), - "minCells": min_cells, - "topN": actual_top_n, - "blacklist": blacklist, - "showPlot": False, - "invalidateCache": False, - "artifact": ArtifactReferenceModel.from_artifact_ref( - hvg_features - ).model_dump(mode="json"), - } - ) - operations.extend( - [ - { - "operation": "set_feature_selection", - "assay": assay_plan.assay, - "source": ArtifactReferenceModel.from_artifact_ref( - hvg_features - ).model_dump(mode="json"), - "exactExcludedFeatures": list(assay_plan.exactExcludedFeatures), - "excludeFamilies": list( - assay_plan.featureParameters.get("excludeFamilies", []) - ), - "artifact": ArtifactReferenceModel.from_artifact_ref( - graph_features - ).model_dump(mode="json"), - }, - { - "operation": "select_detected_features", - "assay": assay_plan.assay, - "cellSelection": cell_selection_model.model_dump(mode="json"), - "minCells": min_cells, - "artifact": ArtifactReferenceModel.from_artifact_ref( - detected - ).model_dump(mode="json"), - }, - { - "operation": "set_feature_selection", - "assay": assay_plan.assay, - "source": ArtifactReferenceModel.from_artifact_ref( - detected - ).model_dump(mode="json"), - "exactExcludedFeatures": list(assay_plan.exactExcludedFeatures), - "excludeFamilies": list( - assay_plan.featureParameters.get("excludeFamilies", []) - ), - "artifact": ArtifactReferenceModel.from_artifact_ref( - marker_features - ).model_dump(mode="json"), - }, - ] - ) - artifacts.update( - { - f"{assay_plan.assay}_hvg_candidates": ( - ArtifactReferenceModel.from_artifact_ref(hvg_features) - ), - f"{assay_plan.assay}_detected_features": ( - ArtifactReferenceModel.from_artifact_ref(detected) - ), - } - ) - elif assay_plan.featureMethod == "prevalentPeaks": - actual_top_n = min( - int(assay_plan.featureParameters["topN"]), - assay.feats.N - 1, - ) - graph_features = store.select_prevalent_peaks( - cell_selection, - from_assay=assay_plan.assay, - top_n=actual_top_n, - invalidate_cache=False, - ) - marker_features = graph_features - actions.append(f"select_prevalent_peaks:{assay_plan.assay}") - operations.append( - { - "operation": "select_prevalent_peaks", - "assay": assay_plan.assay, - "cellSelection": cell_selection_model.model_dump(mode="json"), - "topN": actual_top_n, - "invalidateCache": False, - "artifact": ArtifactReferenceModel.from_artifact_ref( - graph_features - ).model_dump(mode="json"), - } - ) - elif assay_plan.featureMethod == "panel": - mask = np.ones(assay.feats.N, dtype=bool) - ids = np.asarray(assay.feats.fetch_all("ids")).astype(str) - names = np.asarray(assay.feats.fetch_all("names")).astype(str) - excluded = set(assay_plan.exactExcludedFeatures) - if excluded: - mask &= ~np.isin(ids, list(excluded)) - mask &= ~np.isin(names, list(excluded)) - if int(mask.sum()) < 2: - raise ValueError( - f"ADT assay {assay_plan.assay!r} has fewer than two non-control features" - ) - graph_features = store.set_feature_selection( - from_assay=assay_plan.assay, - mask=mask, - invalidate_cache=False, - ) - marker_features = graph_features - actions.append(f"select_adt_panel:{assay_plan.assay}") - operations.append( - { - "operation": "set_feature_selection", - "assay": assay_plan.assay, - "selectedFeatures": int(mask.sum()), - "exactExcludedFeatures": sorted(excluded), - "artifact": ArtifactReferenceModel.from_artifact_ref( - graph_features - ).model_dump(mode="json"), - } - ) - else: - raise ValueError(f"Unsupported feature route {assay_plan.featureMethod!r}") - normalized = store.run_normalization( - cell_selection, - features=graph_features, - log_transform=cast( - bool, - assay_plan.normalizationParameters.get("logTransform"), - ), - renormalize_subset=cast( - bool, - assay_plan.normalizationParameters.get("renormalizeSubset"), - ), - invalidate_cache=False, - ) - graph_feature_group = store.load_artifact(graph_features) - graph_feature_values = cast(Any, graph_feature_group["values"]) - selected_values = np.asarray(graph_feature_values[:], dtype=bool) - graph_features_model = ArtifactReferenceModel.from_artifact_ref(graph_features) - marker_features_model = ArtifactReferenceModel.from_artifact_ref( - marker_features + if not isinstance(assay, RNAassay): + raise TypeError("RNA preprocessing requires an RNAassay") + from ..parameter_tuning.hvg import core_hvg_evidence + + feature_refs = core_hvg_evidence( + store, assay=assay_plan.assay, cells=cell_selection ) - normalized_model = ArtifactReferenceModel.from_artifact_ref(normalized) - handoff = PreprocessedAssayHandoff( - assay=assay_plan.assay, - assayType=assay_plan.assayType, - cellSelection=cell_selection_model, - reductionMethod=assay_plan.reductionMethod, - graphFeatures=graph_features_model, - markerFeatures=marker_features_model, - normalized=normalized_model, - nCells=active_cells, - nFeatures=int(selected_values.sum()), + baseline = feature_refs["scarfDefault"] + selected_count = int( + np.asarray( + as_zarr_array(store.load_artifact(baseline)["values"], name="values")[ + : + ], + dtype=bool, + ).sum() ) + marker_features = store.select_all_features(from_assay=assay_plan.assay) + feature_models = { + key: ArtifactReferenceModel.from_artifact_ref(value) + for key, value in feature_refs.items() + } + marker_model = ArtifactReferenceModel.from_artifact_ref(marker_features) artifacts.update( { - f"{assay_plan.assay}_graph_features": graph_features_model, - f"{assay_plan.assay}_marker_features": marker_features_model, - f"{assay_plan.assay}_normalized": normalized_model, + f"{assay_plan.assay}_{key}": value + for key, value in feature_models.items() } ) - actions.append(f"normalize:{assay_plan.assay}") + artifacts[f"{assay_plan.assay}_marker_features"] = marker_model operations.append( { - "operation": "run_normalization", + "operation": "prepare_core_hvg_baseline", "assay": assay_plan.assay, "cellSelection": cell_selection_model.model_dump(mode="json"), - "features": graph_features_model.model_dump(mode="json"), - "logTransform": assay_plan.normalizationParameters.get("logTransform"), - "renormalizeSubset": assay_plan.normalizationParameters.get( - "renormalizeSubset" - ), - "invalidateCache": False, - "artifact": normalized_model.model_dump(mode="json"), + "requestedTopN": 1000, + "selectedTopN": selected_count, + "artifacts": { + key: value.model_dump(mode="json") + for key, value in feature_models.items() + }, } ) + actions.append(f"prepare_core_defaults:{assay_plan.assay}") logger.info( - f"Preprocessed assay {assay_plan.assay!r}: " - f"cells={handoff.nCells}, features={handoff.nFeatures}, " - f"reduction={handoff.reductionMethod!r}" + f"HVG baseline: core Scarf selected {selected_count:,} genes; " + "representation experiments will use the screening cells." + ) + return PreprocessedAssayHandoff( + assay=assay_plan.assay, + assayType="RNA", + cellSelection=cell_selection_model, + reductionMethod="pca", + graphFeatures=feature_models["scarfDefault"], + markerFeatures=marker_model, + graphFeatureCandidates=feature_models, + nCells=active_cells, + nFeatures=selected_count, ) - return handoff def apply_cell_qc( self, @@ -909,8 +1127,31 @@ def apply_cell_qc( cell_selection: ArtifactRef, actions: list[str], operations: list[dict[str, Any]], + *, + selected_plan: CellQcPlan | None = None, + decision_payload: CellQualityExecutorPayload | None = None, ) -> ArtifactRef: - plan = experimental.cellQc + plan = selected_plan or experimental.cellQc + if decision_payload is not None: + if ( + cell_qc_policy(plan.action, plan.registeredProfile) + != decision_payload.profile + ): + raise ValueError( + "Cell-QC execution plan differs from its audited payload" + ) + expected_capture = decision_payload.profile in { + "coreSampleMad3", + "captureMad5", + "captureMad3Sensitivity", + "pooledReferenceMad5", + } + if expected_capture != ( + plan.sampleColumn is not None or plan.sampleArtifact is not None + ): + raise ValueError( + "Cell-QC capture source differs from its audited payload" + ) input_model = ArtifactReferenceModel.from_artifact_ref(cell_selection) logger.info( f"Applying cell QC action={plan.action!r}, profile={plan.profileId!r}" @@ -927,6 +1168,7 @@ def apply_cell_qc( raise ValueError("Experimental Context selected an unknown QC profile") for field_name in ( "action", + "registeredProfile", "driverAssay", "driverAssayType", "sampleColumn", @@ -952,17 +1194,6 @@ def apply_cell_qc( raise ValueError( "Selected cell-QC sample artifact is absent from Experimental Context" ) - if plan.action == "skip": - actions.append("skip_cell_qc") - operations.append( - { - "operation": "skip_cell_qc", - "profileId": plan.profileId, - "cellSelection": input_model.model_dump(mode="json"), - "artifact": input_model.model_dump(mode="json"), - } - ) - return cell_selection artifact_metrics = [ NamedCellArtifact( name=source.name, @@ -978,57 +1209,100 @@ def apply_cell_qc( artifact=artifact_model_to_ref(plan.sampleArtifact.artifact), ) ) - if plan.action == "globalGaussian": - if plan.sampleColumn is not None or sample_artifact is not None: - raise ValueError("globalGaussian QC cannot include a sample source") - result = store.auto_filter_cells( + if plan.registeredProfile is not None: + result, diagnostic_flags = execute_registered_cell_qc( + store, + plan.registeredProfile, + profile_parameters=profile.parameters, + expected_active_cells=profile.activeCells, + expected_retained_cells=profile.retainedCells, + expected_flag_counts=profile.flaggedCells, attrs=plan.attributes, artifact_metrics=artifact_metrics, - min_p=float(profile.parameters.get("minP", 0.01)), - max_p=float(profile.parameters.get("maxP", 0.99)), cell_selection=cell_selection, + sample_column=plan.sampleColumn, + sample_artifact=sample_artifact, invalidate_cache=False, ) result_model = ArtifactReferenceModel.from_artifact_ref(result) - actions.append(f"cell_qc_global:{profile.profileId}") + flags_model = ( + None + if diagnostic_flags is None + else ArtifactReferenceModel.from_artifact_ref(diagnostic_flags) + ) + actions.append(f"cell_qc_registered:{profile.registeredProfile}") operations.append( { - "operation": "auto_filter_cells", + "operation": "run_registered_cell_qc", "profileId": profile.profileId, + "registeredProfile": profile.registeredProfile, "cellSelection": input_model.model_dump(mode="json"), "attrs": list(plan.attributes), "artifactMetrics": [ source.model_dump(mode="json") for source in plan.artifactMetrics ], - "minP": float(profile.parameters.get("minP", 0.01)), - "maxP": float(profile.parameters.get("maxP", 0.99)), - "sampleColumn": None, + "sampleColumn": plan.sampleColumn, + "sampleArtifact": ( + None + if plan.sampleArtifact is None + else plan.sampleArtifact.model_dump(mode="json") + ), + "profileParameters": profile.parameters, + "expectedActiveCells": profile.activeCells, + "expectedRetainedCells": profile.retainedCells, + "expectedFlagCounts": profile.flaggedCells, "invalidateCache": False, + "diagnosticFlags": ( + None + if flags_model is None + else flags_model.model_dump(mode="json") + ), "artifact": result_model.model_dump(mode="json"), } ) return result - if plan.action == "sampleMad": - if (plan.sampleColumn is None) == (sample_artifact is None): - raise ValueError( - "sampleMad QC requires exactly one metadata or artifact " - "sample source" + if plan.action == "skip": + actions.append("skip_cell_qc") + operations.append( + { + "operation": "skip_cell_qc", + "profileId": plan.profileId, + "cellSelection": input_model.model_dump(mode="json"), + "artifact": input_model.model_dump(mode="json"), + } + ) + return cell_selection + if plan.action in {"globalGaussian", "sampleMad"}: + if not isinstance(profile.resolvedBounds, dict): + raise ValueError("Core cell-QC evidence requires exact named bounds") + capture_artifact = ( + NamedCellArtifact( + name=profile.captureArtifact.name, + artifact=artifact_model_to_ref(profile.captureArtifact.artifact), ) - result = store.auto_filter_cells( + if profile.captureArtifact is not None + else None + ) + result, diagnostic_flags = execute_auto_cell_qc( + store, + plan.action, + profile_parameters=profile.parameters, + expected_active_cells=profile.activeCells, + expected_retained_cells=profile.retainedCells, + expected_flag_counts=profile.flaggedCells, + expected_resolved_bounds=profile.resolvedBounds, attrs=plan.attributes, artifact_metrics=artifact_metrics, cell_selection=cell_selection, sample_column=plan.sampleColumn, sample_artifact=sample_artifact, - n_mads=float(profile.parameters.get("nMads", 3.0)), - min_cells_per_sample=int( - profile.parameters.get("minCellsPerSample", 20) - ), + capture_column=profile.captureColumn, + capture_artifact=capture_artifact, invalidate_cache=False, ) result_model = ArtifactReferenceModel.from_artifact_ref(result) - actions.append(f"cell_qc_sample_mad:{profile.profileId}") + actions.append(f"cell_qc_{plan.action}:{profile.profileId}") operations.append( { "operation": "auto_filter_cells", @@ -1039,86 +1313,21 @@ def apply_cell_qc( source.model_dump(mode="json") for source in plan.artifactMetrics ], - "minP": 0.01, - "maxP": 0.99, "sampleColumn": plan.sampleColumn, - "sampleArtifact": ( - None - if plan.sampleArtifact is None - else plan.sampleArtifact.model_dump(mode="json") - ), - "nMads": float(profile.parameters.get("nMads", 3.0)), - "minCellsPerSample": int( - profile.parameters.get("minCellsPerSample", 20) - ), + "sampleArtifact": plan.sampleArtifact.model_dump(mode="json") + if plan.sampleArtifact + else None, + "profileParameters": profile.parameters, + "resolvedBounds": profile.resolvedBounds, + "expectedRetainedCells": profile.retainedCells, + "diagnosticFlags": ArtifactReferenceModel.from_artifact_ref( + diagnostic_flags + ).model_dump(mode="json") + if diagnostic_flags + else None, "invalidateCache": False, "artifact": result_model.model_dump(mode="json"), } ) return result raise ValueError(f"Unsupported cell QC action {plan.action!r}") - - def rna_blacklist(self, plan: AssayPreprocessingPlan) -> str: - patterns: list[str] = [] - families = set( - cast(list[str], plan.featureParameters.get("excludeFamilies", [])) - ) - if "mitochondrial" in families: - patterns.append(r"^(MT-|mt-)") - if "ribosomal" in families: - patterns.append(r"^(RPS|RPL|MRPS|MRPL|Rps|Rpl|Mrps|Mrpl)") - if "histone" in families: - patterns.append(r"^(HIST|Hist)") - patterns.extend( - rf"^{re.escape(value)}$" for value in plan.exactExcludedFeatures if value - ) - return "|".join(patterns) if patterns else r"(?!)" - - def exclude_exact_features( - self, - store: DataStore, - plan: AssayPreprocessingPlan, - source: ArtifactRef, - ) -> ArtifactRef: - families = set( - cast(list[str], plan.featureParameters.get("excludeFamilies", [])) - ) - if not plan.exactExcludedFeatures and not families: - return source - assay = store.get_assay(plan.assay) - source_group = store.load_artifact(source) - source_values = cast(Any, source_group["values"]) - mask = np.asarray(source_values[:], dtype=bool) - ids = np.asarray(assay.feats.fetch_all("ids")).astype(str) - names = np.asarray(assay.feats.fetch_all("names")).astype(str) - excluded = set(plan.exactExcludedFeatures) - mask &= ~np.isin(ids, list(excluded)) - mask &= ~np.isin(names, list(excluded)) - family_patterns: list[str] = [] - if "mitochondrial" in families: - family_patterns.append(r"^(MT-|mt-)") - if "ribosomal" in families: - family_patterns.append(r"^(RPS|RPL|MRPS|MRPL|Rps|Rpl|Mrps|Mrpl)") - if "histone" in families: - family_patterns.append(r"^(HIST|Hist)") - if family_patterns: - technical = np.zeros(len(mask), dtype=bool) - combined = re.compile("|".join(family_patterns)) - technical |= np.fromiter( - (combined.search(value) is not None for value in ids), - dtype=bool, - count=len(ids), - ) - technical |= np.fromiter( - (combined.search(value) is not None for value in names), - dtype=bool, - count=len(names), - ) - mask &= ~technical - if not mask.any(): - raise ValueError("Exact marker exclusions removed every feature") - return store.set_feature_selection( - from_assay=plan.assay, - mask=mask, - invalidate_cache=False, - ) diff --git a/scarf/agent/orchestrator/rna.py b/scarf/agent/orchestrator/rna.py new file mode 100644 index 00000000..a563699c --- /dev/null +++ b/scarf/agent/orchestrator/rna.py @@ -0,0 +1,245 @@ +"""Supported assay selection for the automated RNA workflow.""" + +from collections.abc import Mapping, Sequence +from typing import Any + + +def validate_rna_directions(directions: Mapping[str, Any]) -> None: + """Reject automated operations outside the RNA analysis contract.""" + if "hypothesisTesting" in directions: + raise ValueError( + "Automated hypothesis testing is not supported; run statistical " + "analysis separately after the RNA workflow." + ) + + +def validate_rna_request_fields(request: Any) -> None: + """Validate routing that does not require inspecting the input store.""" + if len(request.analysisAssays) > 1: + raise ValueError("analysisAssays must select at most one RNA assay") + if request.pairedAssays: + raise ValueError("pairedAssays is unsupported by the single-RNA workflow") + selected = request.analysisAssays[0] if request.analysisAssays else None + if selected is not None and request.primaryAssay not in {None, selected}: + raise ValueError("primaryAssay must match the selected analysisAssays entry") + selected = selected or request.primaryAssay + if selected is not None and request.markerAssay not in {None, selected}: + raise ValueError("markerAssay must match the selected RNA assay") + validate_rna_directions(request.experimentalDirections) + + +def selected_rna_assay(request: Any, assay_types: Mapping[str, str]) -> str: + """Resolve one explicit or uniquely available persisted RNA assay.""" + validate_rna_request_fields(request) + selected = ( + request.analysisAssays[0] if request.analysisAssays else request.primaryAssay + ) + if selected is None: + rna_assays = [ + name for name, assay_type in assay_types.items() if assay_type == "RNA" + ] + if len(rna_assays) != 1: + raise ValueError( + "The automated workflow requires one RNA assay; " + f"found {len(rna_assays)}. Select one with primaryAssay or " + "analysisAssays when multiple RNA assays are present." + ) + selected = rna_assays[0] + if selected not in assay_types: + raise ValueError(f"Unknown requested RNA assay {selected!r}") + if assay_types[selected] != "RNA": + raise ValueError( + f"Selected assay {selected!r} has type {assay_types[selected]!r}; " + "the automated workflow supports RNA only." + ) + if request.markerAssay not in {None, selected}: + raise ValueError("markerAssay must match the selected RNA assay") + return str(selected) + + +def selected_store_rna_assay(store: Any, request: Any) -> str: + """Resolve the workflow assay from the store's persisted summary.""" + return selected_rna_assay( + request, + {value.name: value.assay_type for value in store.summary().assays}, + ) + + +def validate_rna_plan(plan: Any, selected: str) -> None: + """Reject a stale or unsupported cached preprocessing route.""" + if ( + len(plan.assays) != 1 + or plan.assays[0].assay != selected + or plan.assays[0].assayType != "RNA" + or not plan.assays[0].graphEligible + or plan.primaryAssay != selected + or plan.markerAssay != selected + or plan.pairedAssays + ): + raise ValueError( + "Preprocessing must contain only the selected RNA assay. " + "Start a new workflow for incompatible saved analysis state." + ) + + +def validate_rna_context(report: Any, selected: str) -> None: + """Keep reused QC evidence on the selected RNA assay.""" + if report.htoIdentityArtifacts: + raise ValueError( + "Saved automatic HTO processing is unsupported; start a new RNA workflow." + ) + for profile in report.qcProfiles: + if profile.driverAssay != selected or profile.driverAssayType != "RNA": + raise ValueError( + "QC evidence must use the selected RNA assay; start a new workflow " + "for incompatible saved QC state." + ) + if report.cellQc.driverAssay not in {None, selected}: + raise ValueError("Cell QC must use the selected RNA assay") + + +def validate_rna_handoffs(handoffs: Sequence[Any], selected: str) -> None: + """Keep saved preprocessing on the selected persisted RNA modality.""" + if ( + len(handoffs) != 1 + or handoffs[0].assay != selected + or handoffs[0].assayType != "RNA" + ): + raise ValueError( + "Saved preprocessing must contain only the selected RNA assay; " + "start a new workflow for incompatible saved analysis state." + ) + + +def validate_saved_rna_history( + store: Any, prefix: str, workflow_run_id: str, selected: str +) -> None: + """Validate single-RNA ownership before opening resumed work for writes.""" + from ..data_enrichment.contracts import DataEnrichmentReport + from ..experimental_context.contracts import ExperimentalContextResult + from ..experimental_context.study import StudyContract, validate_objective_evidence + from ..parameter_tuning.contracts import ParameterTuningReport + from . import journal + from .models import ( + _STAGE_ORDER, + AutomatedPreprocessingPlan, + PreprocessedAssayHandoff, + ) + + for stage in _STAGE_ORDER: + for outcome in journal._stage_outcomes( + store.zw, prefix, workflow_run_id, stage + ): + if outcome.outputs.get("htoIdentityArtifacts"): + raise ValueError( + "Saved automatic HTO processing is unsupported; start a new RNA workflow" + ) + if stage == "rna_quality_metrics" and outcome.status == "done": + if "percentageDefinitions" not in outcome.outputs: + raise ValueError( + "Saved RNA quality metrics lack exact percentage definitions; " + "start a new workflow. Existing analysis artifacts remain accessible." + ) + for name in ("preprocessingPlan", "resolvedPreprocessingPlan"): + if outcome.outputs.get(name): + validate_rna_plan( + AutomatedPreprocessingPlan.model_validate( + outcome.outputs[name] + ), + selected, + ) + if stage == "preprocessing" and "assays" in outcome.outputs: + validate_rna_handoffs( + [ + PreprocessedAssayHandoff.model_validate(v) + for v in outcome.outputs["assays"] + ], + selected, + ) + if not outcome.reportReferences: + continue + if stage == "data_enrichment": + report = DataEnrichmentReport.model_validate( + journal.read_stage_evidence(store, outcome.reportReferences[0]) + ) + if report.status == "done" and ( + len(report.policies) != 1 + or report.policies[0].assay != selected + or report.policies[0].assayModality != "RNA" + ): + raise ValueError("Saved enrichment includes unsupported assays") + elif stage == "experimental_context": + context = ExperimentalContextResult.model_validate( + journal.read_stage_evidence(store, outcome.reportReferences[0]) + ) + if context.status == "done": + validate_rna_context(context, selected) + if outcome.status == "done": + raw_contract = outcome.outputs.get("studyContract", {}) + _require_objective_contract(raw_contract) + validate_objective_evidence( + StudyContract.model_validate(raw_contract), context + ) + elif stage == "parameter_tuning": + tuning = ParameterTuningReport.model_validate( + journal.read_stage_evidence(store, outcome.reportReferences[0]) + ) + if ( + tuning.recommendedIntegrationId is not None + or tuning.assayReports + or tuning.fromAssay != selected + ): + raise ValueError( + "Saved tuning includes unsupported assays or integration" + ) + validate_analysis_evidence(journal.analysis_snapshot(store, workflow_run_id)) + + +def _require_objective_contract(value: Any) -> None: + """Reject historical conclusions without rewriting their evidence.""" + if not isinstance(value, Mapping) or not { + "evidenceRequirements", + "evidenceCoverage", + }.issubset(value): + raise ValueError( + "Saved analysis lacks mandatory objective evidence requirements; " + "start a new workflow to resume or regenerate its report. " + "Existing analysis artifacts and historical HTML remain accessible." + ) + + +def validate_analysis_evidence(snapshot: Mapping[str, Any]) -> None: + """Check scientific completion from the authenticated journal view.""" + from ..experimental_context.contracts import ExperimentalContextResult + from ..experimental_context.study import StudyContract, validate_objective_evidence + from .rna_tuning import validate_completed_comparison_evidence + + context_validated = False + for stage in snapshot.get("stages", []): + if ( + stage.get("stage") != "experimental_context" + or stage.get("status") != "done" + ): + continue + raw_contract = stage.get("outputs", {}).get("studyContract", {}) + _require_objective_contract(raw_contract) + validate_objective_evidence( + StudyContract.model_validate(raw_contract), + ExperimentalContextResult.model_validate(stage["report"]), + ) + context_validated = True + accepted = [ + review + for review in snapshot.get("analysisReviews", []) + if review.get("action") == "accept" + ] + for review in accepted: + validate_completed_comparison_evidence(review) + if snapshot.get("status") == "completed": + if not context_validated or not any( + item.get("scope") == "full" for item in accepted + ): + raise ValueError( + "Completed analysis lacks its mandatory objective and comparison evidence; " + "start a new workflow. Existing analysis artifacts remain accessible." + ) diff --git a/scarf/agent/orchestrator/rna_tuning.py b/scarf/agent/orchestrator/rna_tuning.py new file mode 100644 index 00000000..7fdf0094 --- /dev/null +++ b/scarf/agent/orchestrator/rna_tuning.py @@ -0,0 +1,3162 @@ +"""Objective-led RNA experiments on frozen screening and full-cohort cells.""" + +import base64 +import hashlib +import json +import re +from collections.abc import Mapping, Sequence +from typing import Annotated, Any, Literal, cast + +import numpy as np +from pydantic import Field, create_model, model_validator +from pydantic.json_schema import SkipJsonSchema + +from ...metadata.rows import read_metadata_rows_chunkwise +from ...storage.refs import ArtifactRef +from ...storage.selections import ( + read_stored_selection_indices, + resolve_generated_selection_artifact, +) +from ...utils.logging import logger +from .. import record_io +from ..config.agent_exec import ( + ImageEvidence, + ImageInputUnsupportedError, + build_visual_evidence_prompt, + run_agent_sync, +) +from ..experimental_context.contracts import CovariateComparison +from ..experimental_context.study import ( + StudyContract, + unsupported_comparison_limitations, +) +from ..parameter_tuning.agent import prepare_parameter_tuning_dependencies +from ..parameter_tuning.comparisons import ( + CombinedSettings, + ComparisonConclusion, + ComparisonTradeoff, + PopulationConcern, + bind_comparison_measurements, + comparison_advantages, + partition_comparison_evidence, + setting_changes, + validate_comparison_review, +) +from ..parameter_tuning.contracts import ( + ArtifactRecord, + ParameterCandidate, + ParameterCandidateEvaluation, + ParameterTuningNeedsInput, + ParameterTuningReport, +) +from ..parameter_tuning.diagnostics import ( + SCARF_DEFAULT_DIAGNOSTIC_FAMILIES, + _family_mask, + _neighbor_overlap, + augment_cluster_evaluations, + augment_pca_evaluations, + population_support_evidence, + restore_advisory_doublets, + score_advisory_doublets, +) +from ..parameter_tuning.execution import ( + diagnostic_call, + diagnostic_reuse, + execute_parameter_candidate, + refresh_candidate_design_evidence, +) +from ..parameter_tuning.hvg import ( + HvgGroupVariability, + aggregate_hvg_rankings, + rank_core_hvgs, +) +from ..parameter_tuning.selection import ( + finalize_parameter_tuning_selection, + harmony_acceptance_gate, +) +from ..types import AgentDataModel, ArtifactReferenceModel +from . import journal +from .budget import CandidateBudget, CandidateBudgetExceeded, candidate_identity +from .models import ( + AutomatedPreprocessingPlan, + OrchestrationRequestRecord, + PreprocessedAssayHandoff, + artifact_model_to_ref, +) + +_STRUCTURED_VISUAL_LIMITATION = ( + "The model assessed structured marker, PCA loading and diagnostic evidence; " + "visual inspection was unavailable because the configured model does not accept images." +) + + +def _configured_image_input(model: Any) -> bool | None: + """Honor an explicit input capability without inferring it from image output.""" + declared = getattr(model, "supports_image_input", None) + if isinstance(declared, bool): + return declared + profile = getattr(model, "profile", None) + if isinstance(profile, Mapping): + declared = profile.get("supports_image_input") + if isinstance(declared, bool): + return declared + return None + + +def screening_sizes(n_cells: int, config: Any) -> tuple[int, ...]: + """Resolve the saved fixed policy or the bounded ten-percent default.""" + if n_cells < 1: + raise ValueError("Screening requires a non-empty retained cohort") + if config.screeningCells is None: + maximum = min(n_cells, config.maxScreeningCells, 100_000) + initial = min(maximum, max(10_000, (n_cells + 9) // 10)) + else: + initial = min(n_cells, config.screeningCells) + maximum = min(n_cells, config.maxScreeningCells) + return tuple(dict.fromkeys((initial, maximum))) + + +class TuningAction(AgentDataModel): + """An assessment of observed evidence and at most one registered experiment.""" + + action: Literal["accept", "combine", "experiment", "enlarge", "defer"] = Field( + description=( + "Accept supported observed evidence, request one next experiment, " + "enlarge a screening sample, or defer an unresolved essential question." + ) + ) + selectedCandidateId: str = Field( + description=( + "Copy an observed candidate ID. An experiment must keep " + "currentCandidateId as its fixed baseline." + ) + ) + experimentId: str | None = Field( + default=None, + description=( + "For action=experiment, copy one exact key from experiments. " + "This requests the next operation; it has not yet run. Otherwise null." + ), + ) + correctionNeed: Literal["needed", "notNeeded", "uncertain", "notApplicable"] + evidenceIds: list[str] = Field( + min_length=1, + description=( + "Copy exact IDs from availableEvidenceIds, including the selected " + "candidate's anchor or one of its supplied diagnostic evidence IDs." + ), + ) + quantitativeFindings: list[str] = Field( + min_length=1, + description="Describe supplied observed measurements, not predicted results.", + ) + qualitativeFindings: list[str] = Field(min_length=1) + comparisonConclusions: list[ComparisonConclusion] + populationConcerns: list[PopulationConcern] = Field(default_factory=list) + plainLanguageSummary: str = Field(min_length=1) + combinedSettings: CombinedSettings | None = None + concern: str = "" + expectedImprovement: str = Field( + default="", + description="Predict what the requested experiment should improve and why.", + ) + objectivePreservation: str = Field(min_length=1) + rationale: str = Field( + min_length=1, + description=( + "Explain why this action follows from observed evidence. " + "Do not describe an unexecuted experiment as a completed result." + ), + ) + + @model_validator(mode="after") + def validate_action(self) -> "TuningAction": + if (self.action == "combine") != (self.combinedSettings is not None): + raise ValueError("Only a combine action proposes combined settings") + if (self.action == "experiment") != (self.experimentId is not None): + raise ValueError("Only an experiment action names an experiment") + if self.action == "experiment" and ( + not self.concern.strip() or not self.expectedImprovement.strip() + ): + raise ValueError( + "An experiment needs an observed concern and expected improvement" + ) + return self + + +def _assessment_output_type( + candidate_ids: Sequence[str], + experiment_ids: Sequence[str], + *, + scope: str, + phase: str = "combined", +) -> type[TuningAction]: + """Constrain new model choices without changing the saved action contract.""" + if not candidate_ids: + raise ValueError("RNA assessment requires observed candidates") + tradeoff_type = create_model( + "ObservedTradeoffInterpretation", + __base__=ComparisonTradeoff, + preferredValue=(Annotated[float | None, SkipJsonSchema()], None), + alternativeValue=(Annotated[float | None, SkipJsonSchema()], None), + ) + conclusion_type = create_model( + "ObservedComparisonInterpretation", + __base__=ComparisonConclusion, + tradeoffs=( + cast(Any, list)[tradeoff_type], + Field( + description=( + "Required array: for each comparisonAdvantages row matching this " + "axis and preferredCandidateId, supply alternativeCandidateId, " + "metric and your interpretation. Explanations in quantitativeReason " + "or biologicalReason do not replace these entries. For combine or " + "accept, use [] only when this preference has no measured " + "alternative advantages. Include only matching inventory rows: " + "ties and measurements favoring your preference belong in " + "quantitativeReason, not tradeoffs. Scarf attaches the measured values." + ) + ), + ), + ) + actions = tuple( + action + for action in ("accept", "combine", "experiment", "enlarge", "defer") + if (action != "enlarge" or scope != "full") + and (action != "experiment" or experiment_ids) + and (action != "combine" or phase == "sensitivity") + and (action != "accept" or phase != "sensitivity") + ) + return create_model( + "ObservedRnaAssessment", + __base__=TuningAction, + comparisonConclusions=(cast(Any, list)[conclusion_type], ...), + action=( + Literal[actions], + Field(description=TuningAction.model_fields["action"].description), + ), + selectedCandidateId=( + Literal[tuple(dict.fromkeys(candidate_ids))], + Field( + description=TuningAction.model_fields["selectedCandidateId"].description + ), + ), + experimentId=( + Literal[tuple(dict.fromkeys(experiment_ids))] | None + if experiment_ids + else type(None), + Field( + default=None, + description=TuningAction.model_fields["experimentId"].description, + ), + ), + ) + + +class RnaSetting(AgentDataModel): + parameters: ParameterCandidate + features: ArtifactReferenceModel + eligibleFeatures: ArtifactReferenceModel + hvgCount: int = 1000 + ranking: Literal["global", "batchAware"] = "global" + rankingColumn: str | None = None + + +def validate_completed_comparison_evidence(review: Mapping[str, Any]) -> None: + """Check exact completed review coverage before finalization or report reuse.""" + if not { + "comparisonCoverage", + "comparisonConclusions", + "plainLanguageSummary", + }.issubset(review): + raise ValueError( + "Saved analysis lacks mandatory comparison evidence; start a new workflow" + ) + action = TuningAction.model_validate( + { + key: value + for key, value in review.items() + if key in TuningAction.model_fields + } + ) + coverage = review.get("comparisonCoverage") + if not isinstance(coverage, Mapping): + raise ValueError( + "Saved analysis lacks mandatory comparison evidence; start a new workflow" + ) + validate_comparison_review(coverage, action.model_dump(mode="json")) + for candidate in review.get("candidates", []): + setting = coverage["candidateSettings"].get(candidate["candidateId"]) + if ( + setting is None + or any( + setting[key] != candidate[key] + for key in ("parameters", "cellSelection") + ) + or setting["features"] + != candidate.get("artifacts", {}).get("graphFeatures") + ): + raise ValueError( + "Comparison evidence differs from the exact reviewed candidate" + ) + if any( + candidate["metrics"].get(key) != value + for key, value in setting.get("metrics", {}).items() + ): + raise ValueError( + "Comparison measurements differ from the reviewed candidate" + ) + + +def uniform_screening_selection( + store: Any, parent: ArtifactRef, *, size: int, seed: int +) -> ArtifactRef: + """Select nested uniform prefixes without weighting study groups differently.""" + indices = read_stored_selection_indices( + store.zw, + parent, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + if size >= len(indices): + return parent + if size < 3: + raise ValueError("Screening requires at least three cells") + order = np.random.Generator(np.random.PCG64(seed)).permutation(len(indices)) + mask = np.zeros(store.cells.N, dtype=bool) + mask[indices[order[:size]]] = True + selection, _ = resolve_generated_selection_artifact( + store.zw, + scope="datastore", + kind="cell_selection", + values=mask, + row_ids=np.asarray(store.cells.fetch_all("ids")), + operation="agent_uniform_rna_screen", + parameters={"size": size, "seed": seed, "generator": "PCG64"}, + inputs={"parent_selection": parent}, + source_column="agent_screening", + ) + return selection + + +def screening_coverage( + store: Any, + parent: ArtifactRef, + sample: ArtifactRef, + columns: Sequence[str], + combinations: Sequence[Sequence[str]] = (), +) -> tuple[dict[str, Any], list[str]]: + """Report population and sample proportions without oversampling rare groups.""" + + def indices(selection: ArtifactRef) -> np.ndarray: + return read_stored_selection_indices( + store.zw, + selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + + full_indices, sample_indices = indices(parent), indices(sample) + evidence: dict[str, Any] = { + "populationCells": len(full_indices), + "screeningCells": len(sample_indices), + "sampling": "Uniform without replacement; no group oversampling or weights.", + "groups": {}, + } + concerns: list[str] = [] + grouped = {} + for column in dict.fromkeys(columns): + full = np.asarray( + read_metadata_rows_chunkwise(store.cells, column, full_indices) + ).astype(str) + sampled = np.asarray( + read_metadata_rows_chunkwise(store.cells, column, sample_indices) + ).astype(str) + grouped[column] = (full, sampled) + if combinations: + from ..experimental_context.characterization import _SelectionBoundCells + from ..experimental_context.comparisons import combination_labels + + full_cells = _SelectionBoundCells(store.zw, store.cells, parent) + sample_cells = _SelectionBoundCells(store.zw, store.cells, sample) + for group_columns in combinations: + key = "joint:" + json.dumps(list(group_columns), separators=(",", ":")) + grouped[key] = ( + combination_labels(full_cells, group_columns), + combination_labels(sample_cells, group_columns), + ) + for column, (full, sampled) in grouped.items(): + values, counts = np.unique(full, return_counts=True) + sample_values, sample_counts = np.unique(sampled, return_counts=True) + lookup = dict(zip(sample_values, sample_counts, strict=True)) + rows = [] + for value, count in zip(values, counts, strict=True): + observed = int(lookup.get(value, 0)) + rows.append( + { + "value": str(value), + "populationCells": int(count), + "screeningCells": observed, + "populationFraction": float(count / len(full)), + "screeningFraction": observed / len(sampled), + } + ) + if observed < min(int(count), 20): + concerns.append( + f"{column}={value}: {observed} sampled of {int(count)} cells" + ) + evidence["groups"][column] = rows + return evidence, concerns + + +class RnaTuningRun: + """Execute evidence-requested comparisons with one authoritative history.""" + + def __init__( + self, + owner: Any, + store: Any, + workflow: Any, + request: OrchestrationRequestRecord, + plan: AutomatedPreprocessingPlan, + handoff: PreprocessedAssayHandoff, + study: StudyContract, + answers: Mapping[str, Any], + provenance: dict[str, Any], + *, + design_comparisons: Sequence[CovariateComparison] = (), + previous_provenances: Sequence[dict[str, Any]] = (), + ) -> None: + if ( + handoff.cellSelection is None + or handoff.graphFeatures is None + or handoff.markerFeatures is None + ): + raise ValueError( + "RNA tuning requires exact cells, graph genes and marker genes" + ) + self.owner, self.store, self.workflow = owner, store, workflow + self.request, self.plan, self.handoff, self.study = ( + request, + plan, + handoff, + study, + ) + self.answers, self.provenance = answers, provenance + self.previous_provenances = tuple(previous_provenances) + self.evidence_revision = ( + hashlib.sha256(record_io.canonical_json_bytes(provenance)).hexdigest() + if previous_provenances + else None + ) + self.design_comparisons = tuple(design_comparisons) + self.prefix = journal._ensure_orchestration_store(store) + self.cells = artifact_model_to_ref(handoff.cellSelection) + self.marker_features = artifact_model_to_ref(handoff.markerFeatures) + self.budget = CandidateBudget( + store, + self.prefix, + workflow.workflowRunId, + request.config, + provenance, + previous_provenances=self.previous_provenances, + ) + self.settings: dict[str, RnaSetting] = {} + self.history: list[dict[str, Any]] = [] + self.evaluations: dict[str, list[ParameterCandidateEvaluation]] = { + "sample0": [], + "sample1": [], + "full": [], + } + self.last_action: TuningAction | None = None + self.full_repairs = len(self.budget.repairs) + self.answer_consumed = False + self.scope_sizes: dict[str, int] = {} + self.feature_evidence_cache: dict[str, dict[str, Any]] = {} + self.neighbor_comparisons: dict[tuple[str, str], float] = {} + self.diagnostic_counts: dict[str, dict[str, int]] = {} + self.comparison_rows: dict[str, list[dict[str, Any]]] = {} + self.combined_candidates: dict[str, str] = {} + self.resolution_candidates: dict[str, list[str]] = {} + self.validation_sources: dict[str, dict[str, Any]] = {} + self.discovery_scope: str | None = None + self.recovery_scope: str | None = None + self.full_repair: dict[str, Any] | None = None + self.batch_columns = list(study.technicalBatchColumns) + self.coverage_columns = [ + value + for value in dict.fromkeys( + [ + study.physicalCaptureColumn, + *study.independentUnitColumns, + *study.conditionColumns, + *study.technicalBatchColumns, + *study.protectedColumns, + ] + ) + if value is not None and study.columnKinds.get(value) != "continuous" + ] + self.family_patterns = { + str(row["family"]): str(row["pattern"]) + for row in ( + plan.assays[0].featureParameters.get("defaultFeatureInventory") or {} + ).get("families", []) + if isinstance(row, Mapping) and "family" in row and "pattern" in row + } + + def baseline(self, resolution: float = 1.0) -> RnaSetting: + assert self.handoff.graphFeatures is not None + return RnaSetting( + parameters=ParameterCandidate( + candidateId="baseline", + dimensions=min(21, self.handoff.nFeatures - 1, self.handoff.nCells - 1), + neighborsK=min(11, self.handoff.nCells - 1), + leidenResolution=resolution, + useHarmony=False, + ), + features=self.handoff.graphFeatures, + eligibleFeatures=self.handoff.graphFeatureCandidates["eligibleDefault"], + hvgCount=self.handoff.nFeatures, + ) + + def checkpoint_scope(self, scope: str) -> str: + """Separate revised interpretations while retaining numerical admissions.""" + if self.evidence_revision is not None: + return ( + f"parameter_tuning/evidence_revisions/{self.evidence_revision}/{scope}" + ) + return f"parameter_tuning/{scope}" + + def execution_inputs( + self, cells: ArtifactRef, setting: RnaSetting + ) -> dict[str, Any]: + inputs: dict[str, Any] = { + "cells": cells.to_dict(), + "features": setting.features.model_dump(mode="json"), + "parameters": setting.parameters.model_dump(mode="json"), + } + if setting.parameters.useHarmony and any( + previous.get("studyContract", {}).get("technicalBatchColumns") + != self.batch_columns + for previous in self.previous_provenances + ): + inputs["harmonyBatchColumns"] = self.batch_columns + return inputs + + def execute( + self, scope: str, cells: ArtifactRef, setting: RnaSetting + ) -> ParameterCandidateEvaluation: + inputs = self.execution_inputs(cells, setting) + identity = candidate_identity(inputs) + parameters = setting.parameters.model_copy( + update={"candidateId": f"rna_{identity[:24]}"} + ) + setting = setting.model_copy(update={"parameters": parameters}) + self.settings[parameters.candidateId] = setting + saved: dict[str, Any] | None + source = self.budget.completed_source(inputs) if scope == "full" else None + if source is not None: + admission, saved = source + if admission["scope"] != "full": + self.validation_sources[parameters.candidateId] = { + "scope": admission["scope"], + "slot": admission["slot"], + "identity": admission["identity"], + } + else: + admission = self.budget.admit(scope, inputs) + saved = self.budget.completed(admission) + if saved is not None: + evaluation = ParameterCandidateEvaluation.model_validate( + saved["evaluation"] + ) + for artifact in evaluation.artifacts.values(): + status = self.store.inspect_artifact(artifact_model_to_ref(artifact)) + if not status.exists or not status.complete: + raise ValueError( + "Saved candidate evidence is unavailable or incomplete" + ) + diagnostic_reuse("diagnostic.primaryCandidateEvidence", "restored") + evaluation = self._reassess_saved_evaluation( + scope, setting, admission, evaluation + ) + else: + features = artifact_model_to_ref(setting.features) + normalized = diagnostic_call( + "core.primaryNormalization", + self.store.run_normalization, + cells, + features=features, + invalidate_cache=False, + ) + deps, ids = prepare_parameter_tuning_dependencies( + self.store, + normalized=normalized, + candidates=[parameters], + batch_columns=self.batch_columns, + preservation_columns=self.study.protectedColumns, + pair_harmony_candidates=False, + max_candidates=1, + min_cluster_cells=1, + ) + deps.protectedCombinations = tuple( + tuple(columns) for columns in self.study.protectedCombinations + ) + deps.columnKinds = self.study.columnKinds + evaluation = execute_parameter_candidate(deps, ids[0]) + if evaluation.status != "done": + raise RuntimeError( + evaluation.error + or "Candidate execution failed; its admitted work can be retried on resume" + ) + evaluation.artifacts.update( + { + "normalized": ArtifactRecord.from_ref(normalized), + "graphFeatures": ArtifactRecord.from_ref(features), + } + ) + evaluation = self._augment_evaluation(scope, setting, evaluation) + evaluation = ParameterCandidateEvaluation.model_validate_json( + record_io.canonical_json_bytes(evaluation.model_dump(mode="json")) + ) + self.budget.complete( + admission, {"evaluation": evaluation.model_dump(mode="json")} + ) + if evaluation.candidateId not in { + item.candidateId for item in self.evaluations[scope] + }: + self.evaluations[scope].append(evaluation) + return evaluation + + def _augment_evaluation( + self, + scope: str, + setting: RnaSetting, + evaluation: ParameterCandidateEvaluation, + *, + preserve_doublets: bool = False, + ) -> ParameterCandidateEvaluation: + """Attach required loading, population, stability and doublet evidence.""" + features = artifact_model_to_ref(setting.features) + parameters = setting.parameters + evaluation = augment_pca_evaluations( + self.store, + [evaluation], + feature_selection=features, + nominated_families=SCARF_DEFAULT_DIAGNOSTIC_FAMILIES, + protected_families=self.plan.assays[0].featureParameters.get( + "protectFamilies", [] + ), + technical_columns=self.batch_columns, + batch_columns=self.batch_columns, + protected_columns=self.study.protectedColumns, + qc_columns=[ + *self.plan.cellQc.attributes, + *(source.name for source in self.plan.cellQc.artifactMetrics), + ], + qc_artifacts={ + source.name: artifact_model_to_ref(source.artifact) + for source in self.plan.cellQc.artifactMetrics + }, + column_kinds={ + **self.study.columnKinds, + **{ + source.name: "continuous" + for source in self.plan.cellQc.artifactMetrics + }, + }, + )[0] + native = next( + ( + item + for item in self.evaluations[scope] + if not item.parameters.useHarmony + and self.settings[item.candidateId].features == setting.features + and item.parameters.model_dump(exclude={"candidateId", "useHarmony"}) + == parameters.model_dump(exclude={"candidateId", "useHarmony"}) + ), + None, + ) + doublets = None + if self.request.config.scoreDoublets or ( + self.study.correctionLicense == "safe" and self.batch_columns + ): + doublets = ( + restore_advisory_doublets( + evaluation, capture_column=self.study.physicalCaptureColumn + ) + if preserve_doublets and "doubletNativeGraph" in evaluation.artifacts + else score_advisory_doublets( + self.store, + native or evaluation, + [ + item + for item in [*self.evaluations[scope], evaluation] + if item.artifacts.get("graphFeatures") + == evaluation.artifacts["graphFeatures"] + and item.cellSelection == evaluation.cellSelection + ], + assay=self.handoff.assay, + feature_selection=features, + capture_column=self.study.physicalCaptureColumn, + ) + ) + else: + limitation = ( + f"Advisory doublet scoring was not run for assay {self.handoff.assay!r} " + "because score_doublets=False. Doublet contamination was not assessed." + ) + evaluation = evaluation.model_copy( + update={ + "warnings": list(dict.fromkeys([*evaluation.warnings, limitation])) + } + ) + logger.info(limitation) + evaluation = augment_cluster_evaluations( + self.store, + [evaluation], + marker_assay=self.handoff.assay, + marker_features=self.marker_features, + independent_unit_columns=self.study.independentUnitColumns, + technical_columns=self.batch_columns, + nominated_families=SCARF_DEFAULT_DIAGNOSTIC_FAMILIES, + protected_families=self.plan.assays[0].featureParameters.get( + "protectFamilies", [] + ), + doublet_evidence=doublets, + )[0] + return evaluation + + def _reassess_saved_evaluation( + self, + scope: str, + setting: RnaSetting, + admission: dict[str, Any], + evaluation: ParameterCandidateEvaluation, + ) -> ParameterCandidateEvaluation: + original = self.budget.admission_provenance(admission) + if original == self.provenance: + return evaluation + key = f"{self.checkpoint_scope(scope)}/evaluation{admission['slot']}/augmented" + inputs = {"admission": admission, "provenance": self.provenance} + saved = journal.load_checkpoint( + self.store, self.prefix, self.workflow.workflowRunId, key, inputs + ) + if saved is not None: + restored = ParameterCandidateEvaluation.model_validate(saved["evaluation"]) + for artifact in restored.artifacts.values(): + status = self.store.inspect_artifact(artifact_model_to_ref(artifact)) + if not status.exists or not status.complete: + raise ValueError( + "Revised candidate evidence is unavailable or incomplete" + ) + return restored + old = StudyContract.model_validate(original["studyContract"]) + fields = ( + "technicalBatchColumns", + "protectedColumns", + "protectedCombinations", + "columnKinds", + ) + if any(getattr(old, field) != getattr(self.study, field) for field in fields): + normalized = artifact_model_to_ref(evaluation.artifacts["normalized"]) + deps, _ = prepare_parameter_tuning_dependencies( + self.store, + normalized=normalized, + candidates=[evaluation.parameters], + batch_columns=self.batch_columns, + preservation_columns=self.study.protectedColumns, + pair_harmony_candidates=False, + max_candidates=1, + min_cluster_cells=1, + ) + deps.protectedCombinations = tuple( + tuple(columns) for columns in self.study.protectedCombinations + ) + deps.columnKinds = self.study.columnKinds + evaluation = refresh_candidate_design_evidence(deps, evaluation) + if any( + getattr(old, field) != getattr(self.study, field) + for field in ( + "technicalBatchColumns", + "protectedColumns", + "columnKinds", + "independentUnitColumns", + "physicalCaptureColumn", + ) + ): + evaluation = self._augment_evaluation( + scope, + setting, + evaluation, + preserve_doublets=old.physicalCaptureColumn + == self.study.physicalCaptureColumn, + ) + evaluation = ParameterCandidateEvaluation.model_validate_json( + record_io.canonical_json_bytes(evaluation.model_dump(mode="json")) + ) + journal.save_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + key, + inputs, + { + "evaluation": evaluation.model_dump(mode="json"), + "primaryArtifactsReused": True, + "priorProvenance": original, + }, + ) + return evaluation + + def execute_matched( + self, scope: str, cells: ArtifactRef, setting: RnaSetting + ) -> ParameterCandidateEvaluation: + if setting.parameters.useHarmony: + if self.study.correctionLicense != "safe" or not self.batch_columns: + raise ValueError( + "Harmony requires a safe design license and approved batch columns" + ) + native = setting.model_copy( + update={ + "parameters": setting.parameters.model_copy( + update={"useHarmony": False} + ) + } + ) + self.budget.admit_many( + scope, + [ + value + for value in ( + self.execution_inputs(cells, native), + self.execution_inputs(cells, setting), + ) + if scope != "full" or self.budget.completed_source(value) is None + ], + ) + self.execute(scope, cells, native) + return self.execute(scope, cells, setting) + + def harmony_gate( + self, scope: str, selected: ParameterCandidateEvaluation + ) -> tuple[bool, list[str]]: + if not selected.parameters.useHarmony: + return True, [] + setting = self.settings[selected.candidateId] + native = next( + ( + item + for item in self.evaluations[scope] + if not item.parameters.useHarmony + and self.settings[item.candidateId].features == setting.features + and item.parameters.model_dump(exclude={"candidateId", "useHarmony"}) + == selected.parameters.model_dump(exclude={"candidateId", "useHarmony"}) + ), + None, + ) + return harmony_acceptance_gate( + native, + selected, + batch_columns=self.batch_columns, + protected_columns=[ + *self.study.protectedColumns, + *( + "joint:" + json.dumps(columns, separators=(",", ":")) + for columns in self.study.protectedCombinations + ), + ], + independent_unit_columns=self.study.independentUnitColumns, + require_doublet_evidence=True, + ) + + def experiments( + self, selected: ParameterCandidateEvaluation + ) -> dict[str, dict[str, Any]]: + """Offer individual interventions, without executing their numerical work.""" + setting = self.settings[selected.candidateId] + n_cells = next( + ( + self.scope_sizes[name] + for name, rows in self.evaluations.items() + if selected in rows and name in self.scope_sizes + ), + self.handoff.nCells, + ) + options: dict[str, dict[str, Any]] = {} + for field, values in ( + ("dimensions", (10, 21, 30, 50)), + ("neighborsK", (11, 21, 41)), + ("leidenResolution", (0.25, 0.5, 0.75, 1.0, 1.25, 1.5)), + ): + for value in values: + if value != getattr(setting.parameters, field): + if field == "dimensions" and value >= min( + setting.hvgCount, n_cells + ): + continue + if field == "neighborsK" and value >= n_cells: + continue + options[f"{field}:{value}"] = {"parameter": field, "value": value} + for count in (1000, 2000, 4000): + if count != setting.hvgCount: + options[f"hvgCount:{count}"] = {"parameter": "hvgCount", "value": count} + for column in self.batch_columns: + if self.study.columnKinds.get(column) == "continuous": + continue + if setting.ranking != "batchAware" or setting.rankingColumn != column: + options[f"hvgRanking:batchAware:{column}"] = { + "parameter": "hvgRanking", + "value": "batchAware", + "column": column, + } + if setting.ranking == "batchAware": + options["hvgRanking:global"] = { + "parameter": "hvgRanking", + "value": "global", + } + options.update(self._feature_experiments(setting)) + if self.study.correctionLicense == "safe" and self.batch_columns: + options["useHarmony:true"] = {"parameter": "useHarmony", "value": True} + if setting.parameters.useHarmony: + options["useHarmony:false"] = {"parameter": "useHarmony", "value": False} + return options + + def _feature_experiments(self, setting: RnaSetting) -> dict[str, dict[str, Any]]: + """Offer policy changes only when exact eligible genes can change safely.""" + policy = self.plan.assays[0].featureParameters + assay = self.store.get_assay(self.handoff.assay) + names = np.asarray(assay.feats.fetch_all("names")).astype(str) + ids = np.asarray(assay.feats.fetch_all("ids")).astype(str) + eligible = np.asarray( + self.store.load_artifact(artifact_model_to_ref(setting.eligibleFeatures))[ + "values" + ][:], + dtype=bool, + ) + allowed = np.asarray( + self.store.load_artifact( + artifact_model_to_ref( + self.handoff.graphFeatureCandidates["eligibleAll"] + ) + )["values"][:], + dtype=bool, + ) + protected = np.isin(names, policy.get("protectFeatures", [])) | np.isin( + ids, policy.get("protectFeatures", []) + ) + for family in policy.get("protectFamilies", []): + mask = _family_mask(names, family) + if mask is not None: + protected |= mask + masks = { + ("Family", family): np.asarray( + [ + re.search(pattern, name, flags=re.IGNORECASE) is not None + for name in names + ] + ) + for family, pattern in self.family_patterns.items() + } + masks.update( + { + ("Feature", feature): (names == feature) | (ids == feature) + for feature in dict.fromkeys( + [ + *policy.get("proposedExcludeFeatures", []), + *policy.get("protectFeatures", []), + ] + ) + } + ) + options = {} + for (kind, value), mask in masks.items(): + for operation in ("include", "exclude"): + if operation == "exclude" and (mask & protected).any(): + continue + changed = ( + mask & allowed & ~eligible + if operation == "include" + else mask & eligible + ) + if changed.any(): + field = operation + kind + options[f"{field}:{value}"] = { + "parameter": field, + "value": value, + "affectedEligibleGenes": int(changed.sum()), + } + return options + + def batch_ranking( + self, eligible: ArtifactRef, count: int, column: str + ) -> np.ndarray: + """Use core per-group variability on the same globally eligible genes.""" + if column not in self.batch_columns: + raise ValueError("Batch-aware ranking needs an approved technical column") + from ..experimental_context.characterization import _SelectionBoundCells + + cells = _SelectionBoundCells(self.store.zw, self.store.cells, self.cells) + values, counts = np.unique(cells.fetch(column), return_counts=True) + groups = [] + for value, n_cells in zip(values, counts, strict=True): + if n_cells < 20: + continue + selection = self.store.filter_cells( + [column], + [value], + [value], + cell_selection=self.cells, + keep_bounds=True, + invalidate_cache=False, + ) + reference = self.store.select_hvgs( + selection, + from_assay=self.handoff.assay, + top_n=self.store.get_assay(self.handoff.assay).feats.N, + min_cells=1, + max_cells=np.inf, + blacklist="", + show_plot=False, + invalidate_cache=False, + ) + status = self.store.inspect_artifact(reference) + summary = ArtifactRef.from_dict(dict(status.inputs["feature_summary"])) + group = self.store.load_artifact(reference) + detected = ( + np.asarray(self.store.load_artifact(summary)["normed_n"][:]) >= 20 + ) + groups.append( + HvgGroupVariability( + group_id=str(value), + cell_count=int(n_cells), + corrected_variance=np.asarray( + group["corrected_variance"][:], dtype=np.float64 + ), + detected_features=detected, + ) + ) + if len(groups) < 2: + raise ValueError( + "Batch-aware ranking lacks two groups with sufficient cells" + ) + mask = np.asarray(self.store.load_artifact(eligible)["values"][:], dtype=bool) + statistics = artifact_model_to_ref( + self.handoff.graphFeatureCandidates["eligibleAll"] + ) + variance = np.asarray( + self.store.load_artifact(statistics)["corrected_variance"][:], + dtype=np.float64, + ) + ranking = aggregate_hvg_rankings( + variance, + mask, + groups, + valid_group_count=len(groups), + candidate_targets=(count,), + ) + return ranking.ranking + + def apply_experiment( + self, selected: ParameterCandidateEvaluation, experiment: dict[str, Any] + ) -> RnaSetting: + setting = self.settings[selected.candidateId] + field, value = experiment["parameter"], experiment["value"] + if field in {"dimensions", "neighborsK", "leidenResolution", "useHarmony"}: + return setting.model_copy( + update={ + "parameters": setting.parameters.model_copy(update={field: value}) + } + ) + eligible = artifact_model_to_ref(setting.eligibleFeatures) + count = int(value) if field == "hvgCount" else setting.hvgCount + ranking_mode = value if field == "hvgRanking" else setting.ranking + ranking_column = ( + experiment.get("column") if field == "hvgRanking" else setting.rankingColumn + ) + if field in { + "includeFamily", + "excludeFamily", + "includeFeature", + "excludeFeature", + }: + mask = np.asarray( + self.store.load_artifact(eligible)["values"][:], dtype=bool + ) + names = np.asarray( + self.store.get_assay(self.handoff.assay).feats.fetch_all("names") + ).astype(str) + feature_ids = np.asarray( + self.store.get_assay(self.handoff.assay).feats.fetch_all("ids") + ).astype(str) + if field.endswith("Family"): + pattern = re.compile(self.family_patterns[value], flags=re.IGNORECASE) + family_mask = np.asarray( + [pattern.search(name) is not None for name in names] + ) + else: + family_mask = (names == value) | (feature_ids == value) + if not family_mask.any(): + raise ValueError( + "The nominated exact feature is absent from the assay" + ) + if field.startswith("include"): + all_eligible = artifact_model_to_ref( + self.handoff.graphFeatureCandidates["eligibleAll"] + ) + allowed = np.asarray( + self.store.load_artifact(all_eligible)["values"][:], dtype=bool + ) + mask |= allowed & family_mask + else: + policy = self.plan.assays[0].featureParameters + protected_mask = np.isin( + names, policy.get("protectFeatures", []) + ) | np.isin(feature_ids, policy.get("protectFeatures", [])) + for family in policy.get("protectFamilies", []): + family_protection = _family_mask(names, family) + if family_protection is not None: + protected_mask |= family_protection + elif family in self.family_patterns: + pattern = re.compile( + self.family_patterns[family], flags=re.IGNORECASE + ) + protected_mask |= np.asarray( + [pattern.search(name) is not None for name in names] + ) + if np.any(family_mask & protected_mask): + raise ValueError( + "An objective-protected feature or family cannot be excluded" + ) + mask &= ~family_mask + eligible = self.store.set_feature_selection( + from_assay=self.handoff.assay, mask=mask, invalidate_cache=False + ) + if ranking_mode == "batchAware" and ranking_column is None: + raise ValueError( + "Batch-aware ranking requires an explicit technical column" + ) + indices = ( + self.batch_ranking(eligible, count, ranking_column) + if ranking_mode == "batchAware" and ranking_column is not None + else None + ) + features = rank_core_hvgs( + self.store, + eligible=eligible, + statistics=artifact_model_to_ref( + self.handoff.graphFeatureCandidates["eligibleAll"] + ), + top_n=count, + ranking=indices, + ) + n_features = int( + np.asarray( + self.store.load_artifact(features)["values"][:], dtype=bool + ).sum() + ) + if n_features <= setting.parameters.dimensions: + raise ValueError( + "This feature experiment cannot retain the fixed PCA dimension" + ) + return setting.model_copy( + update={ + "features": ArtifactReferenceModel.from_artifact_ref(features), + "eligibleFeatures": ArtifactReferenceModel.from_artifact_ref(eligible), + "hvgCount": n_features, + "ranking": ranking_mode, + "rankingColumn": ranking_column, + } + ) + + def _setting_checkpoint( + self, + scope: str, + key: str, + baseline: RnaSetting, + experiment: dict[str, Any], + ) -> tuple[str, dict[str, Any], dict[str, Any] | None]: + """Reuse exact feature work, with a new address for a revised intervention.""" + inputs: dict[str, Any] = { + "baseline": baseline.model_dump(mode="json"), + "experiment": experiment, + "cells": self.cells.to_dict(), + } + checkpoint = f"parameter_tuning/{scope}/{key}/setting" + previous = journal.read_checkpoint( + self.store, self.prefix, self.workflow.workflowRunId, checkpoint + ) + if self.evidence_revision is not None: + ranking_changed = experiment.get("parameter") == "hvgRanking" and any( + item["studyContract"]["technicalBatchColumns"] != self.batch_columns + for item in self.previous_provenances + ) + if previous is None or previous["inputs"] != inputs or ranking_changed: + checkpoint = f"{self.checkpoint_scope(scope)}/{key}/setting" + if ranking_changed: + inputs["rankingColumns"] = self.batch_columns + return ( + checkpoint, + inputs, + journal.load_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + checkpoint, + inputs, + ), + ) + + def _prepared_setting( + self, + scope: str, + key: str, + selected: ParameterCandidateEvaluation, + experiment: dict[str, Any], + cells: ArtifactRef, + ) -> RnaSetting: + """Commit feature work separately so interrupted reviews do not repeat it.""" + baseline = self.settings[selected.candidateId] + checkpoint, inputs, saved = self._setting_checkpoint( + scope, key, baseline, experiment + ) + if saved is not None: + return RnaSetting.model_validate(saved["setting"]) + proposed = self.execution_inputs(cells, baseline) + proposed["features"] = {"requestedExperiment": inputs} + self.budget.check_many(scope, [proposed]) + setting = self.apply_experiment(selected, experiment) + journal.save_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + checkpoint, + inputs, + {"setting": setting.model_dump(mode="json")}, + ) + return setting + + def _feature_nomination(self, baseline: RnaSetting) -> dict[str, Any] | None: + """Find one meaningful, previously justified policy intervention.""" + policy = self.plan.assays[0].featureParameters + eligible = np.asarray( + self.store.load_artifact(artifact_model_to_ref(baseline.eligibleFeatures))[ + "values" + ][:], + dtype=bool, + ) + all_eligible = np.asarray( + self.store.load_artifact( + artifact_model_to_ref( + self.handoff.graphFeatureCandidates["eligibleAll"] + ) + )["values"][:], + dtype=bool, + ) + names = np.asarray( + self.store.get_assay(self.handoff.assay).feats.fetch_all("names") + ).astype(str) + ids = np.asarray( + self.store.get_assay(self.handoff.assay).feats.fetch_all("ids") + ).astype(str) + nominations = [ + *(("includeFeature", value) for value in policy.get("protectFeatures", [])), + *(("includeFamily", value) for value in policy.get("protectFamilies", [])), + *( + ("excludeFeature", value) + for value in policy.get("proposedExcludeFeatures", []) + ), + *( + ("excludeFamily", value) + for value in policy.get("proposedExcludeFamilies", []) + ), + ] + protected = np.isin(names, policy.get("protectFeatures", [])) | np.isin( + ids, policy.get("protectFeatures", []) + ) + for family in policy.get("protectFamilies", []): + mask = _family_mask(names, family) + if mask is not None: + protected |= mask + for operation, value in nominations: + if operation.endswith("Family"): + mask = _family_mask(names, value) + if mask is None: + continue + if value not in self.family_patterns: + value = next( + ( + family + for family, pattern in self.family_patterns.items() + if np.array_equal( + mask, + np.asarray( + [ + re.search(pattern, name, flags=re.IGNORECASE) + is not None + for name in names + ] + ), + ) + ), + None, + ) + if value is None: + continue + else: + mask = (names == value) | (ids == value) + if operation.startswith("include"): + changed = mask & all_eligible & ~eligible + else: + if (mask & protected).any(): + continue + changed = mask & eligible + if changed.any(): + return {"parameter": operation, "value": value} + return None + + def _sensitivity_panel( + self, + scope: str, + cells: ArtifactRef, + baseline: ParameterCandidateEvaluation, + ) -> None: + """Execute a small single-axis floor, never a Cartesian product.""" + setting = self.settings[baseline.candidateId] + n_cells = self.scope_sizes[scope] + plans: list[tuple[str, str, dict[str, Any] | None, str]] = [] + for field, axis, values in ( + ("hvgCount", "hvgCount", (2000, 4000)), + ("dimensions", "pca", (10, 30)), + ("neighborsK", "neighbors", (21, 41)), + ): + for value in values: + reason = "" + if field == "dimensions" and value >= min(setting.hvgCount, n_cells): + reason = "The requested PCA dimension exceeds the observed cell or selected-feature rank." + if field == "neighborsK" and value >= n_cells: + reason = ( + "The requested neighbor count exceeds the observed cell count." + ) + plans.append( + ( + f"{field}:{value}", + axis, + None if reason else {"parameter": field, "value": value}, + reason, + ) + ) + columns = list( + dict.fromkeys( + [ + *( + [self.study.physicalCaptureColumn] + if self.study.physicalCaptureColumn in self.batch_columns + else [] + ), + *self.batch_columns, + ] + ) + ) + ranking = None + ranking_groups = {} + for column in columns: + if self.study.columnKinds.get(column) == "continuous": + continue + indices = read_stored_selection_indices( + self.store.zw, + self.cells, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ) + group_values = read_metadata_rows_chunkwise( + self.store.cells, column, indices + ) + _, counts = np.unique(group_values, return_counts=True) + ranking_groups[column] = int((counts >= 20).sum()) + if int((counts >= 20).sum()) >= 2: + ranking = { + "parameter": "hvgRanking", + "value": "batchAware", + "column": column, + } + break + plans.append( + ( + "hvgRanking", + "hvgRanking", + ranking, + "No approved categorical technical grouping has two groups with at least 20 cells." + if ranking is None + else "", + ) + ) + nomination = self._feature_nomination(setting) + pending_feature = nomination is None and bool( + self._feature_experiments(setting) + ) + plans.append( + ( + "featurePolicy", + "featurePolicy", + nomination, + "Baseline loading, marker and objective evidence must nominate one of the offered policy interventions before combining." + if pending_feature + else "Every registered inclusion/exclusion is a no-op on the exact eligible genes or would remove objective-protected genes." + if nomination is None + else "", + ) + ) + proposed: list[dict[str, Any]] = [] + for identifier, _, experiment, _ in plans: + if experiment is None: + continue + proposal = self.execution_inputs(cells, setting) + if experiment["parameter"] in {"dimensions", "neighborsK"}: + proposal["parameters"] = { + **proposal["parameters"], + experiment["parameter"]: experiment["value"], + } + else: + _, _, prepared = self._setting_checkpoint( + scope, f"sensitivity/{identifier}", setting, experiment + ) + if prepared is not None: + known = RnaSetting.model_validate(prepared["setting"]) + if ( + known.features == setting.features + or np.array_equal( + self.store.load_artifact( + artifact_model_to_ref(known.features) + )["values"][:], + self.store.load_artifact( + artifact_model_to_ref(setting.features) + )["values"][:], + ) + or ( + identifier == "featurePolicy" + and known.hvgCount != setting.hvgCount + ) + ): + continue + proposal = self.execution_inputs(cells, known) + else: + proposal["features"] = {"requiredSensitivity": identifier} + proposed.append(proposal) + self.budget.check_many(scope, proposed) + rows: list[dict[str, Any]] = [ + { + "comparisonId": f"defaultResolution:{item.parameters.leidenResolution}", + "axis": "partition", + "status": "completed", + "baselineCandidateId": baseline.candidateId, + "alternativeCandidateId": item.candidateId, + "reason": "", + } + for item in self.evaluations[scope] + if item.candidateId != baseline.candidateId + and item.parameters.leidenResolution in {0.5, 0.75, 1.25} + ] + for identifier, axis, experiment, reason in plans: + alternative = None + proof: dict[str, Any] = { + "baselineFeatures": setting.features.model_dump(mode="json") + } + if axis in {"pca", "neighbors"}: + proof["kind"] = "numericalBound" + elif axis == "hvgRanking": + proof.update( + kind="insufficientTechnicalGroups", + eligibleGroupsByColumn=ranking_groups, + ) + elif axis == "featurePolicy": + proof.update( + kind="noPermittedPolicy", + meaningfulPermittedInterventions=0, + registeredFamilies=list(self.family_patterns), + ) + if experiment is not None: + if experiment["parameter"] in {"dimensions", "neighborsK"}: + alternative_setting = self.apply_experiment(baseline, experiment) + else: + alternative_setting = self._prepared_setting( + scope, + f"sensitivity/{identifier}", + baseline, + experiment, + cells, + ) + same_genes = alternative_setting.features == setting.features + if not same_genes and axis in { + "hvgCount", + "hvgRanking", + "featurePolicy", + }: + same_genes = np.array_equal( + self.store.load_artifact( + artifact_model_to_ref(alternative_setting.features) + )["values"][:], + self.store.load_artifact( + artifact_model_to_ref(setting.features) + )["values"][:], + ) + if ( + axis == "featurePolicy" + and alternative_setting.hvgCount != setting.hvgCount + ): + reason = "The nominated policy leaves too few eligible genes to hold the baseline HVG count fixed." + proof.update( + kind="insufficientEligibleGenes", + eligibleFeatureCount=alternative_setting.hvgCount, + ) + elif same_genes and axis in {"hvgCount", "hvgRanking", "featurePolicy"}: + reason = "The requested intervention produces exactly the same selected genes as the baseline; its numerical representation is already evaluated." + proof.update( + kind="identicalSelectedGenes", + verifiedEqualMasks=True, + alternativeSetting=alternative_setting.model_dump(mode="json"), + ) + elif ( + axis in {"pca", "neighbors"} + and setting_changes( + setting.model_dump(mode="json"), + alternative_setting.model_dump(mode="json"), + ) + == {} + ): + reason = ( + "The requested value is already the exact baseline setting." + ) + else: + alternative = self.execute(scope, cells, alternative_setting) + rows.append( + { + "comparisonId": identifier, + "axis": axis, + "status": "completed" + if alternative is not None + else "pending" + if identifier == "featurePolicy" and pending_feature + else "notApplicable", + "baselineCandidateId": baseline.candidateId, + "alternativeCandidateId": alternative.candidateId + if alternative is not None + else None, + "reason": reason, + "observedProof": proof if alternative is None else None, + } + ) + self.comparison_rows[scope] = rows + + def comparison_coverage(self, scope: str, cells: ArtifactRef) -> dict[str, Any]: + source = ( + (self.discovery_scope or self.recovery_scope) if scope == "full" else scope + ) + source = source or scope + settings = { + item.candidateId: { + **self.settings[item.candidateId].model_dump(mode="json"), + "scope": name, + "status": item.status, + "nCells": self.scope_sizes.get(name, self.handoff.nCells), + "cellSelection": item.cellSelection.model_dump(mode="json") + if item.cellSelection is not None + else None, + "metrics": item.metrics.model_dump( + mode="json", + include={ + "nClusters", + "minClusterCells", + "seedStability", + "subsampleStability", + "markerCoherence", + "markerSpecificityMedian", + "macroF1", + "weightedF1", + "graphSilhouetteMedian", + "clusterConnectivity", + "crossUnitSupport", + "batchMixing", + "biologicalPreservation", + "topMarkerGenes", + }, + ), + } + for name, rows in self.evaluations.items() + for item in rows + } + rows = list(self.comparison_rows.get(source, [])) + panel_source = ( + "full" if scope == "full" and "full" in self.combined_candidates else source + ) + combined_id = self.combined_candidates.get(panel_source) + resolution_ids = self.resolution_candidates.get(panel_source, []) + if combined_id is not None: + rows.extend( + { + "comparisonId": f"combinedResolution:{settings[identifier]['parameters']['leidenResolution']}", + "axis": "partition", + "status": "completed", + "baselineCandidateId": combined_id, + "alternativeCandidateId": identifier, + "reason": "", + } + for identifier in resolution_ids + if identifier != combined_id + ) + return { + "phase": "validation" + if scope == "full" + and (self.discovery_scope is not None or self.recovery_scope is not None) + else "combined" + if source in self.combined_candidates + else "sensitivity", + "population": "allCells" if cells == self.cells else "subset", + "comparisons": rows, + "candidateSettings": settings, + "combinedCandidateId": combined_id, + "resolutionCandidateIds": resolution_ids, + "validationSources": dict(self.validation_sources), + "fullRepair": self.full_repair, + } + + def _combined_setting( + self, + scope: str, + review_index: int, + action: TuningAction, + cells: ArtifactRef, + ) -> RnaSetting: + assert action.combinedSettings is not None + choices = action.combinedSettings + policy = self.settings[choices.featurePolicyCandidateId] + count = self.settings[choices.hvgCountCandidateId].hvgCount + ranking = self.settings[choices.hvgRankingCandidateId] + parameters = policy.parameters.model_copy( + update={ + "dimensions": self.settings[ + choices.pcaCandidateId + ].parameters.dimensions, + "neighborsK": self.settings[ + choices.neighborsCandidateId + ].parameters.neighborsK, + "leidenResolution": 1.0, + "useHarmony": False, + } + ) + inputs = { + "choices": choices.model_dump(mode="json"), + "settings": { + identifier: self.settings[identifier].model_dump(mode="json") + for identifier in choices.model_dump().values() + }, + "cells": self.cells.to_dict(), + } + key = f"{self.checkpoint_scope(scope)}/review{review_index}/combined_setting" + saved = journal.load_checkpoint( + self.store, self.prefix, self.workflow.workflowRunId, key, inputs + ) + if saved is not None: + return RnaSetting.model_validate(saved["setting"]) + eligible = artifact_model_to_ref(policy.eligibleFeatures) + if ( + count == policy.hvgCount + and ranking.ranking == policy.ranking + and ranking.rankingColumn == policy.rankingColumn + ): + features = policy.features + else: + proposed = self.execution_inputs(cells, policy) + proposed["features"] = {"combinedSettings": inputs} + proposed["parameters"] = parameters.model_dump(mode="json") + self.budget.check_many(scope, [proposed]) + order = ( + self.batch_ranking(eligible, count, ranking.rankingColumn) + if ranking.ranking == "batchAware" and ranking.rankingColumn is not None + else None + ) + features = ArtifactReferenceModel.from_artifact_ref( + rank_core_hvgs( + self.store, + eligible=eligible, + statistics=artifact_model_to_ref( + self.handoff.graphFeatureCandidates["eligibleAll"] + ), + top_n=count, + ranking=order, + ) + ) + actual_count = int( + np.asarray( + self.store.load_artifact(artifact_model_to_ref(features))["values"][:], + dtype=bool, + ).sum() + ) + if parameters.dimensions >= min(actual_count, self.scope_sizes[scope]): + raise ValueError( + "Combined settings cannot support the selected PCA dimension" + ) + if actual_count != count: + raise ValueError( + "The proposed combined HVG count was not retained by the selected feature policy" + ) + setting = RnaSetting( + parameters=parameters, + features=features, + eligibleFeatures=policy.eligibleFeatures, + hvgCount=actual_count, + ranking=ranking.ranking, + rankingColumn=ranking.rankingColumn, + ) + journal.save_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + key, + inputs, + {"setting": setting.model_dump(mode="json")}, + ) + return setting + + def _resolution_panel( + self, + scope: str, + cells: ArtifactRef, + setting: RnaSetting, + ) -> ParameterCandidateEvaluation: + settings = [ + setting.model_copy( + update={ + "parameters": setting.parameters.model_copy( + update={"leidenResolution": resolution} + ) + } + ) + for resolution in (0.5, 0.75, 1.0, 1.25) + ] + proposals = [self.execution_inputs(cells, item) for item in settings] + if setting.parameters.useHarmony: + proposals += [ + self.execution_inputs( + cells, + item.model_copy( + update={ + "parameters": item.parameters.model_copy( + update={"useHarmony": False} + ) + } + ), + ) + for item in settings + ] + self.budget.check_many( + scope, + [ + row + for row in proposals + if scope != "full" or self.budget.completed_source(row) is None + ], + ) + if scope == "full" and self.recovery_scope is not None: + recovery_inputs = { + "sourceScope": self.recovery_scope, + "sourceReview": self.last_action.model_dump(mode="json") + if self.last_action is not None + else None, + "cells": cells.to_dict(), + "setting": setting.model_dump(mode="json"), + } + journal.save_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + "parameter_tuning/full/targeted_recovery/" + + hashlib.sha256( + record_io.canonical_json_bytes(recovery_inputs) + ).hexdigest(), + inputs=recovery_inputs, + outputs={ + "purpose": "Resolve the recorded screening concern using full-cohort population, marker, stability and preservation evidence for the proposed settings.", + "screeningAccepted": False, + "plannedEvaluations": proposals, + "limits": self.budget.summary()["limits"], + }, + ) + self.budget.admit_many( + scope, + [row for row in proposals if self.budget.completed_source(row) is None], + ) + evaluations = [self.execute_matched(scope, cells, item) for item in settings] + self.resolution_candidates[scope] = [item.candidateId for item in evaluations] + return next( + item for item in evaluations if item.parameters.leidenResolution == 1.0 + ) + + def feature_evidence( + self, selected: ParameterCandidateEvaluation + ) -> dict[str, Any]: + """Summarize frozen core statistics without rerunning feature variability.""" + setting = self.settings[selected.candidateId] + key = setting.model_dump_json(exclude={"parameters"}) + if key in self.feature_evidence_cache: + return self.feature_evidence_cache[key] + mask = np.asarray( + self.store.load_artifact(artifact_model_to_ref(setting.features))["values"][ + : + ], + dtype=bool, + ) + eligible = np.asarray( + self.store.load_artifact(artifact_model_to_ref(setting.eligibleFeatures))[ + "values" + ][:], + dtype=bool, + ) + reference = artifact_model_to_ref( + self.handoff.graphFeatureCandidates["eligibleAll"] + ) + variance = np.asarray( + self.store.load_artifact(reference)["corrected_variance"][:], + dtype=np.float64, + ) + names = np.asarray( + self.store.get_assay(self.handoff.assay).feats.fetch_all("names") + ).astype(str) + indices = np.flatnonzero(eligible) + ranked = indices[np.lexsort((indices, -variance[indices]))] + families = {} + for family, pattern in self.family_patterns.items(): + expression = re.compile(pattern, flags=re.IGNORECASE) + membership = np.asarray( + [expression.search(name) is not None for name in names] + ) + families[family] = { + "eligibleGenes": int((eligible & membership).sum()), + "selectedGenes": int((mask & membership).sum()), + "selectedExamples": names[ + np.flatnonzero(mask & membership)[:8] + ].tolist(), + "excludedExamples": names[ + np.flatnonzero(~eligible & membership)[:8] + ].tolist(), + } + evidence: dict[str, Any] = { + "statistics": reference.to_dict(), + "statisticsCells": self.cells.to_dict(), + "basis": "Core Scarf corrected variance on the full QC-retained cells; feature-axis summaries are descriptive, not a substitute for downstream comparisons.", + "selectedGenes": int(mask.sum()), + "eligibleGenes": int(eligible.sum()), + "ranking": setting.ranking, + "rankingColumn": setting.rankingColumn, + "globalRankLandmarks": [ + { + "rank": rank, + "gene": str(names[ranked[rank - 1]]), + "correctedVariance": float(variance[ranked[rank - 1]]), + } + for rank in (1, 500, 1000, 2000, 4000) + if rank <= len(ranked) + ], + "topSelectedGenes": names[ + [index for index in ranked if mask[index]][:20] + ].tolist(), + "families": families, + } + self.feature_evidence_cache[key] = evidence + return evidence + + def review( + self, + scope: str, + review_index: int, + selected: ParameterCandidateEvaluation, + coverage: dict[str, Any], + ) -> TuningAction: + from .tuning import _analysis_visual_content + + candidates = self.evaluations[scope] + key = f"{self.checkpoint_scope(scope)}/review{review_index}" + previous_review = journal.read_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + key, + ) + candidate_evidence = [item.model_dump(mode="json") for item in candidates] + setting_evidence = { + item.candidateId: self.settings[item.candidateId].model_dump(mode="json") + for item in candidates + } + committed_review = previous_review + prepared_inputs = { + "candidatesSha256": hashlib.sha256( + record_io.canonical_json_bytes(candidate_evidence) + ).hexdigest(), + "settingsSha256": hashlib.sha256( + record_io.canonical_json_bytes(setting_evidence) + ).hexdigest(), + "coverage": coverage, + "provenance": self.provenance, + } + prepared = None + if previous_review is None: + for evidence_mode in ("structured", "visual"): + prepared = journal.load_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + f"{key}/evidence/{evidence_mode}", + inputs=prepared_inputs, + ) + if prepared is not None: + previous_review = {"inputs": prepared["evidence"]} + break + if previous_review is not None and ( + previous_review["inputs"].get("candidates") != candidate_evidence + or previous_review["inputs"].get("settings") != setting_evidence + ): + raise ValueError( + "Saved review has different candidate evidence or settings. " + "If this review predates required Harmony comparisons for safe designs, " + "start a new workflow in a different destination; the saved history and artifacts remain intact." + ) + if ( + previous_review is not None + and "comparisonCoverage" not in previous_review["inputs"] + ): + raise ValueError( + "Saved analysis lacks mandatory sensitivity coverage; start a new workflow" + ) + if previous_review is not None: + diagnostic_reuse("diagnostic.reviewEvidence", "restored") + experiments = ( + previous_review["inputs"]["experiments"] + if previous_review is not None + else self.experiments(selected) + ) + completed_experiments: dict[str, str] = {} + matched_comparisons = [] + current_setting = self.settings[selected.candidateId] + for candidate in candidates: + other = self.settings[candidate.candidateId] + if ( + candidate.candidateId == selected.candidateId + or candidate.cellSelection != selected.cellSelection + or other.parameters.reductionMethod + != current_setting.parameters.reductionMethod + ): + continue + changes = setting_changes( + current_setting.model_dump(mode="json"), other.model_dump(mode="json") + ) + if len(changes) != 1: + continue + matched_comparisons.append( + { + "currentCandidateId": selected.candidateId, + "alternativeCandidateId": candidate.candidateId, + "changedParameter": changes, + "partitionEvidence": diagnostic_call( + "diagnostic.partitionComparison", + partition_comparison_evidence, + self.store, + selected, + candidate, + ) + if previous_review is None + else {}, + "basis": "Same frozen cells and every unaffected logical setting; the named intervention alone changes. Feature comparisons hold ranking/count/eligible genes fixed except for their declared axis.", + } + ) + if previous_review is None and candidate.status == "done": + field, values = next(iter(changes.items())) + field = { + "pca": "dimensions", + "neighbors": "neighborsK", + "partition": "leidenResolution", + "correction": "useHarmony", + }.get(field, field) + for experiment_id, experiment in experiments.items(): + requested_value = ( + [experiment["value"], experiment.get("column")] + if field == "hvgRanking" + else experiment["value"] + ) + if ( + experiment["parameter"] == field + and requested_value == values["alternative"] + ): + completed_experiments[experiment_id] = candidate.candidateId + if previous_review is None: + experiments = { + key: value + for key, value in experiments.items() + if key not in completed_experiments + } + comparisons = ( + previous_review["inputs"].get("neighborComparisons", []) + if previous_review is not None + else [] + ) + selected_neighbors = selected.artifacts.get("neighbors") + for alternative in [] if previous_review is not None else candidates: + other_neighbors = alternative.artifacts.get("neighbors") + if ( + selected_neighbors is None + or other_neighbors is None + or selected_neighbors == other_neighbors + or selected.parameters.neighborsK != alternative.parameters.neighborsK + ): + continue + left_id, right_id = sorted( + (selected_neighbors.artifactId, other_neighbors.artifactId) + ) + pair = (left_id, right_id) + if pair not in self.neighbor_comparisons: + self.neighbor_comparisons[pair] = diagnostic_call( + "diagnostic.neighborOverlap", + _neighbor_overlap, + self.store, + artifact_model_to_ref(selected_neighbors), + artifact_model_to_ref(other_neighbors), + ) + else: + diagnostic_reuse("diagnostic.neighborOverlap", "cacheHits") + comparisons.append( + { + "leftCandidateId": selected.candidateId, + "rightCandidateId": alternative.candidateId, + "meanNeighborJaccard": self.neighbor_comparisons[pair], + "basis": "Same frozen cells and k; descriptive response to settings, not independent stability or evidence that correction is beneficial.", + } + ) + declared_image_input = _configured_image_input(self.owner.model) + capability_key = "parameter_tuning/structured_evidence" + capability_inputs = { + **( + self.previous_provenances[0] + if self.previous_provenances + else self.provenance + ), + "configuredImageInput": declared_image_input, + } + capability = journal.load_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + capability_key, + inputs=capability_inputs, + ) + if capability is None and declared_image_input is False: + capability = journal.save_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + capability_key, + inputs=capability_inputs, + outputs={ + "evidenceMode": "structured", + "reason": "configuredImageInputUnsupported", + }, + ) + if capability is not None and capability.get("evidenceMode") != "structured": + raise ValueError("The recorded model capability is invalid") + mode = ( + committed_review["inputs"].get("evidenceMode") + if committed_review is not None + else "structured" + if capability is not None + else previous_review["inputs"].get("evidenceMode") + if previous_review is not None + else "visual" + ) + if mode not in {"visual", "structured"}: + raise ValueError( + "Saved review lacks an exact evidence mode; start a new workflow" + ) + visual_inspection = "available" if mode == "visual" else "unavailable" + images = [] + if ( + previous_review is not None + and previous_review["inputs"].get("evidenceMode") == mode + ): + image_hashes = previous_review["inputs"].get("imageHashes", {}) + if prepared is not None: + images = [ + ImageEvidence( + identifier=item["identifier"], + data=base64.b64decode(item["dataBase64"], validate=True), + media_type=item["mediaType"], + ) + for item in prepared["images"] + ] + if { + item.identifier: hashlib.sha256(item.data).hexdigest() + for item in images + } != image_hashes: + raise ValueError( + "Saved review images do not match their exact evidence hashes" + ) + elif mode == "visual": + images = _analysis_visual_content( + self.store, + selected, + candidates, + qc_columns=self.plan.cellQc.attributes, + qc_artifact_metrics=[ + (item.name, item.artifact) + for item in self.plan.cellQc.artifactMetrics + ], + ) + image_hashes = { + image.identifier: hashlib.sha256(image.data).hexdigest() + for image in images + } + else: + image_hashes = {} + if not isinstance(image_hashes, dict) or bool(image_hashes) != ( + mode == "visual" + ): + raise ValueError("Review images do not match its evidence mode") + evidence_ids = ( + [ + identifier + for identifier in previous_review["inputs"]["availableEvidenceIds"] + if identifier not in previous_review["inputs"].get("imageHashes", {}) + or identifier in image_hashes + ] + if previous_review is not None + else list( + dict.fromkeys( + [ + *(f"candidate:{item.candidateId}" for item in candidates), + *(key for item in candidates for key in item.evidenceIds), + *self.study.evidenceIds, + *self.plan.cellQc.evidenceIds, + *(item.evidenceId for item in self.design_comparisons), + *image_hashes, + "studyContract", + "qcPolicy", + "samplingCoverage", + "featureEvidence", + "neighborComparisons", + "assessmentContext", + "comparisonCoverage", + ] + ) + ) + ) + if not isinstance(evidence_ids, list) or any( + not isinstance(value, str) for value in evidence_ids + ): + raise ValueError("Review evidence IDs must be a list of strings") + evidence: dict[str, Any] = { + "studyContract": self.study.model_dump(mode="json"), + "qcPolicy": self.plan.cellQc.model_dump(mode="json"), + "scope": scope, + "evidenceMode": mode, + "visualInspection": visual_inspection, + "configuredImageInput": declared_image_input, + "coverage": coverage, + "currentCandidateId": selected.candidateId, + "candidates": candidate_evidence, + "settings": setting_evidence, + "featureEvidence": previous_review["inputs"]["featureEvidence"] + if previous_review is not None + else {item.candidateId: self.feature_evidence(item) for item in candidates}, + "neighborComparisons": comparisons, + "harmonyGates": { + item.candidateId: self.harmony_gate(scope, item) + for item in candidates + if item.parameters.useHarmony + }, + "experiments": experiments, + "availableEvidenceIds": evidence_ids, + "imageHashes": image_hashes, + "budget": { + "visibleEvaluations": { + name: len(rows) for name, rows in self.evaluations.items() + }, + "limits": self.budget.summary()["limits"], + }, + "fullRepairsUsed": previous_review["inputs"]["fullRepairsUsed"] + if previous_review is not None + else self.full_repairs, + "comparisonCoverage": self.comparison_coverage( + scope, artifact_model_to_ref(selected.cellSelection) + ) + if selected.cellSelection is not None + else {}, + "pilotPopulationWarnings": { + item.candidateId: "This sampled partition includes fewer than 20 cells in a population. Assess its relevance and support; it is not evidence of an invalid biological group. Accepting this partition requires a larger sample or full-cohort assessment." + for item in candidates + if scope != "full" + and selected.cellSelection != self.handoff.cellSelection + and item.metrics.minClusterCells is not None + and item.metrics.minClusterCells < 20 + }, + "smallPopulationPolicy": "Small full-cohort groups require explicit marker, stability, graph and applicable independent-unit support assessment against the objective. Their size alone is neither proof of biology nor grounds for rejection. Defer when essential evidence is insufficient.", + } + if previous_review is None: + support_columns = list( + dict.fromkeys( + column + for column in ( + self.study.physicalCaptureColumn, + *self.study.independentUnitColumns, + ) + if column is not None + ) + ) + evidence["assessmentContext"] = { + "correctionPolicy": { + "license": self.study.correctionLicense, + "harmonyPermitted": self.study.correctionLicense == "safe" + and bool(self.batch_columns), + "nativeAcceptance": "A native descriptive analysis may be accepted when its required evidence supports the objective, while explicitly retaining confounding limitations. A Harmony gate is required only for accepting Harmony. Defer if the objective requires effects that the design cannot separate.", + }, + "matchedComparisons": matched_comparisons, + "alreadyEvaluatedExperiments": completed_experiments, + "designComparisons": [ + item.model_dump(mode="json") for item in self.design_comparisons + ], + "populationSupport": { + item.candidateId: diagnostic_call( + "diagnostic.populationSupport", + population_support_evidence, + self.store, + item, + support_columns, + ) + for item in candidates + if scope == "full" or item.candidateId == selected.candidateId + } + if support_columns + else {}, + "previousActions": [ + { + "scope": row["scope"], + **{ + name: row["review"][name] + for name in ( + "action", + "selectedCandidateId", + "experimentId", + "correctionNeed", + ) + }, + } + for row in self.history + if "review" in row + ], + } + elif "assessmentContext" in previous_review["inputs"]: + evidence["assessmentContext"] = previous_review["inputs"][ + "assessmentContext" + ] + if committed_review is None: + evidence["comparisonAdvantages"] = comparison_advantages( + evidence["comparisonCoverage"] + ) + elif "comparisonAdvantages" in committed_review["inputs"]: + evidence["comparisonAdvantages"] = committed_review["inputs"][ + "comparisonAdvantages" + ] + + def validate(action: TuningAction, *, replay: bool = False) -> TuningAction: + if not replay: + action = TuningAction.model_validate( + bind_comparison_measurements( + evidence["comparisonCoverage"], action.model_dump(mode="json") + ) + ) + by_id = {item.candidateId: item for item in candidates} + if action.selectedCandidateId not in by_id: + raise ValueError("Choose an observed candidate from the current cells") + unknown_ids = sorted(set(action.evidenceIds).difference(evidence_ids)) + if unknown_ids: + raise ValueError( + f"Assessment cited unknown evidence: {unknown_ids[:8]!r}. " + "Copy exact IDs from availableEvidenceIds; do not invent " + "suffixes or use artifact IDs as evidence IDs. The selected " + f"candidate anchor is 'candidate:{action.selectedCandidateId}'." + ) + if mode == "visual" and not set(action.evidenceIds).intersection( + image_hashes + ): + raise ValueError("Assessment must cite its actual visual evidence") + chosen = by_id[action.selectedCandidateId] + if not set(action.evidenceIds).intersection( + {f"candidate:{chosen.candidateId}", *chosen.evidenceIds} + ): + raise ValueError( + "Assessment must cite the selected numerical evidence: " + f"'candidate:{chosen.candidateId}' or one of that candidate's " + "supplied evidenceIds in availableEvidenceIds." + ) + if ( + action.experimentId is not None + and action.experimentId not in experiments + ): + raise ValueError( + f"Unknown experiment ID {action.experimentId!r}. " + f"For current candidate {selected.candidateId!r}, copy one exact " + f"key from experiments: {list(experiments)!r}. " + "Only action='experiment' may name a next operation." + ) + if ( + action.action == "experiment" + and action.selectedCandidateId != selected.candidateId + ): + raise ValueError( + "An offered experiment must use the current candidate as its fixed baseline" + ) + validate_comparison_review( + evidence["comparisonCoverage"], action.model_dump(mode="json") + ) + if action.combinedSettings is not None: + choices = action.combinedSettings + count = self.settings[choices.hvgCountCandidateId].hvgCount + available = evidence["featureEvidence"][ + choices.featurePolicyCandidateId + ]["eligibleGenes"] + if count > available: + raise ValueError( + "The combined feature policy cannot supply the requested HVG count; choose a compatible observed count or policy" + ) + if ( + self.study.correctionLicense == "unsafeConfounded" + and action.correctionNeed in {"needed", "notNeeded"} + ): + message = ( + "The design confounds the batch columns with protected biology; their association cannot establish a removable technical effect. " + "Harmony is not permitted, and a PCA or feature experiment is not a substitute for Harmony. " + "Use notApplicable for the prohibited correction with an explicit confounding limitation, or uncertain and defer if an essential question cannot be resolved." + ) + if replay and action.action == "experiment": + self.history.append( + { + "scope": scope, + "reason": "A saved screening rationale claimed identifiable correction necessity despite the confounded design. Its numerical experiment remains in the audit history; that scientific claim must be reassessed from the supplied evidence.", + } + ) + else: + raise ValueError(message) + if ( + self.study.correctionLicense == "safe" + and action.correctionNeed == "notApplicable" + ): + raise ValueError( + "A safe correction design still requires an observed necessity assessment" + ) + if action.action == "accept": + if action.correctionNeed == "uncertain": + raise ValueError( + "Uncertain correction necessity requires further evidence or deferral before acceptance" + ) + if not chosen.eligible: + raise ValueError( + "The selected candidate failed required full-cell checks" + ) + accepted, reasons = self.harmony_gate(scope, chosen) + if not accepted: + raise ValueError("Harmony acceptance failed: " + "; ".join(reasons)) + required = ( + "seedStability", + "subsampleStability", + "markerCoherence", + "membershipStrengthMean", + "clusterConnectivity", + ) + if any(getattr(chosen.metrics, field) is None for field in required): + raise ValueError( + "Required stability, marker or graph evidence is missing" + ) + if ( + self.study.independentUnitColumns + and chosen.metrics.crossUnitSupport is None + ): + raise ValueError( + "Required independent-unit support evidence is missing" + ) + if chosen.parameters.useHarmony and action.correctionNeed != "needed": + raise ValueError( + "Accepting Harmony requires observed correction necessity" + ) + if self.study.unsupportedProtection and action.correctionNeed in { + "needed", + "uncertain", + }: + raise ValueError( + "Correction necessity remains unresolved because required matched biological protection is unsupported" + ) + if ( + not chosen.parameters.useHarmony + and action.correctionNeed == "needed" + ): + raise ValueError( + "Native acceptance leaves required correction unresolved; provide more evidence or defer" + ) + if self.study.correctionLicense == "safe": + chosen_setting = self.settings[chosen.candidateId] + has_comparison = any( + item.parameters.useHarmony + and item.status == "done" + and item.cellSelection == chosen.cellSelection + and item.harmonyBatchColumns == self.batch_columns + and self.settings[item.candidateId].features + == chosen_setting.features + and item.parameters.model_dump( + exclude={"candidateId", "useHarmony"} + ) + == chosen.parameters.model_dump( + exclude={"candidateId", "useHarmony"} + ) + for item in candidates + ) + if not has_comparison: + raise ValueError( + "A safe correction design requires a completed, matched Harmony experiment " + "before acceptance, including a notNeeded conclusion. " + "Evaluate the offered useHarmony:true experiment or defer; " + "the comparison must use the same cells, features, settings, and approved batch columns." + ) + if self.study.correctionLicense == "indeterminate": + raise ValueError( + "Correction design authorization remains indeterminate" + ) + return action + + saved = journal.load_checkpoint( + self.store, self.prefix, self.workflow.workflowRunId, key, inputs=evidence + ) + pending_saved = saved is not None and saved["action"]["action"] == "defer" + answer = self.answers.get(key) + if ( + answer is None + and not self.answer_consumed + and (saved is None or pending_saved) + ): + answer = self.answers.get("parameter_tuning") + if pending_saved: + assert saved is not None + answer_inputs = {**evidence, "deferredAction": saved["action"]} + prior_answer = journal.load_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + key + "/answer", + inputs=answer_inputs, + ) + if prior_answer is not None: + if ( + answer is not None + and TuningAction.model_validate(answer).model_dump(mode="json") + == prior_answer["action"] + ): + self.answer_consumed = True + key += "/answer" + evidence = answer_inputs + saved = prior_answer + answer = None + elif answer is not None: + key += "/answer" + evidence = answer_inputs + saved = None + if saved is not None: + action = validate(TuningAction.model_validate(saved["action"]), replay=True) + elif answer is not None: + action = validate(TuningAction.model_validate(answer)) + self.answer_consumed = True + else: + journal.save_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + f"{key}/evidence/{mode}", + inputs=prepared_inputs, + outputs={ + "evidence": evidence, + "images": [ + { + "identifier": item.identifier, + "mediaType": item.media_type, + "dataBase64": base64.b64encode(item.data).decode("ascii"), + } + for item in images + ], + }, + ) + prompt = ( + "Assess this RNA analysis as a computational biologist against the exact study objective. " + "Start from Scarf defaults; keep them when observed quantitative evidence and biological interpretation support them. " + "Do not execute a search grid or favor a default solely because it is a default. " + "Explain observed marker programs and relevant PCA loading genes, " + "The supplied comparisonCoverage records which sensitivity comparisons actually ran. If featurePolicy is pending, nominate an offered family/feature experiment from observed baseline loading, marker and objective evidence before combining; missing nomination is not evidence of inapplicability. " + "Conclude every comparisonCoverage axis using its exact completed baseline and alternatives, including candidates with better stability or marker coverage than your preference. " + "In phase=sensitivity, use action=combine and choose each combinedSettings field from an observed candidate on that axis. A combination is a proposed hypothesis, not an observed result. " + "Scarf will execute it and compare four resolutions on its exact graph before acceptance. Do not request already covered settings as experiments. " + "comparisonConclusions must include quantitativeReason, biologicalReason and a short plainLanguageSummary for each axis. Explain how split or merged marker programs serve the stated objective, not just larger clusters or a single numerical maximum. " + "comparisonAdvantages enumerates the exact measured advantages for each possible preference on every axis. For each row belonging to your preference, include a tradeoffs entry naming alternativeCandidateId, metric and interpretation. Scarf attaches the exact saved values; do not transcribe numbers into tradeoff fields. Address every listed advantage, including small differences; these require explanation, not automatic winner selection. " + "tradeoffs contains only counterevidence: an alternative's measurement is strictly better than your preference. Copy only inventory rows matching the exact axis and preferredCandidateId; general comparisons favoring your preference and ties belong in quantitativeReason. " + "If validation requests repair, fix every listed error in the same response using permittedTradeoffs for your current preferences. Preserve unaffected choices and explanations. Explicitly explain any changed preference, update its combinedSettings field when applicable, and use comparisonAdvantages for that new preference. " + "Every numerical explanation must agree with the saved measurements, including direction and candidate identity. Unsupported biological interpretations must remain hypotheses or limitations; do not call protected sex-linked or cell-cycle programs artifacts merely because they appear in PCA loadings. " + "For each selected cluster with empty topMarkerGenes, populationConcerns must name candidateId, clusterId, cited evidenceIds and explain whether it is a nonEssentialLimitation or unresolvedEssential. An unresolved essential population blocks acceptance; do not invent marker support. " + "The action plainLanguageSummary should state the selected settings, what evidence changed the choice, and any unresolved population interpretations without workflow jargon. " + "QC/capture retention, batch associations per PC and protected biological structure. " + "Family dominance alone never proves nuisance; inclusion, exclusion and HVG bans need evidence and objective justification. " + "If a specific concern warrants testing, choose exactly one offered experiment and state its expected improvement and " + "what objective-relevant biology must be preserved. Family policy remains revisable when later evidence warrants it. " + "PCA, HVG, k and resolution changes are separate comparisons. Do not equate a small sampled cluster with an artifact. " + "For a safe design license, a matched Harmony experiment is required before acceptance, even when correction initially appears unnecessary. " + "A notNeeded conclusion must follow the observed native/Harmony comparison, including batch/PC and biological evidence. " + "Keep the native representation when the comparison supports it; evaluating Harmony does not require selecting it. Unsafe or unknown design is never authorization. " + "Only accepting a Harmony representation requires a matched Harmony gate; a native representation does not require one. " + "When correctionLicense is unsafeConfounded, do not infer correction necessity from batch mixing or PCA association. " + "Use notApplicable for the prohibited correction, retain the confounding limitation, and assess whether native descriptive population discovery satisfies the objective. " + "If essential effects remain inseparable, defer. Never request a different parameter change as a proxy for unavailable Harmony. " + "Screening estimates do not prove full-cohort transfer. " + "Findings describe observed results. The experimentId names the exact next operation and expectedImprovement predicts only that operation's effect. " + "Already evaluated settings are existing alternatives, not new experiments; use the supplied matchedComparisons to avoid mixing PCA effects with resolution effects. " + "For example, compare 10 versus 21 PCs at the same resolution rather than quoting the stability of another resolution. " + "A high scaled cLISI is local purity of the supplied label, not proof that a clinical phenotype or all cell types are preserved. " + "Graph connectivity is within-label connectivity. Low library mixing and high PCA association can reflect donor biology or cell composition; neither proves technical causality. " + "A QC association is correlation. Check featureEvidence for actual selected genes: marker-family enrichment cannot show that an excluded family drives PCA. " + "More retained cells, balanced group counts, or cross-unit support alone do not prove healthy cells or biological preservation. " + "Check proposed cell identities against the tissue context. Unexpected marker programs require capture/donor and provenance investigation; do not declare them ordinary tissue populations or assert contamination without evidence. " + "Detailed populationSupport is supplied for currentCandidateId during screening and every final validation candidate; do not claim to have compared unprovided distributions for alternatives. Inspect its capture/donor distribution and missing metadata. Broad support does not prove a biological identity; concentration alone does not prove contamination. " + "Unsupported design comparisons establish neither association nor absence; keep their unresolved requirements visible. " + "Previous actions are history, not scientific authority. Reassess their claims against the exact current evidence. " + "Request enlarge when sample evidence is insufficient; on full cells there is one targeted repair, then defer. " + "Copy evidence IDs exactly from availableEvidenceIds, including the selected candidate's anchor " + "or one of its supplied diagnostic evidence IDs. Do not invent IDs, suffixes or substitute artifact IDs. " + "Never accept merely because work limits are exhausted. " + ) + prompt += ( + "The supplied images are available for visual assessment. Cite actual image evidence IDs and connect the observed plots to the numerical evidence." + if mode == "visual" + else "Visual inspection is unavailable: no images were supplied. Do not claim to have seen, inspected or compared plots or images, and do not cite image IDs. Use qualitativeFindings to interpret the reported marker identities, PCA loading genes, feature families and structured diagnostic tables. State limitations and defer if the supplied evidence cannot resolve an essential question." + ) + serialized_evidence = json.dumps( + evidence, sort_keys=True, separators=(",", ":") + ) + logger.debug( + f"Analysis assessment payload: {len(serialized_evidence.encode('utf-8')):,} " + f"text bytes, {len(candidates)} candidates, {len(images)} images; " + "all saved scientific evidence retained" + ) + try: + result = run_agent_sync( + model=self.owner.model, + output_type=_assessment_output_type( + [item.candidateId for item in candidates], + list(experiments), + scope=scope, + phase=evidence["comparisonCoverage"]["phase"], + ), + system_prompt=prompt, + user_prompt=build_visual_evidence_prompt( + serialized_evidence, images + ) + if mode == "visual" + else serialized_evidence, + config=self.request.config.agentRunConfig.model_copy( + update={ + "retries": min( + self.request.config.agentRunConfig.retries, 2 + ), + "requestLimit": min( + self.request.config.agentRunConfig.requestLimit, 3 + ), + } + ), + name=f"rna_{scope}_assessment", + output_validator=validate, + on_attempt=journal.model_attempt_callback( + self.store, + self.prefix, + self.workflow.workflowRunId, + key, + evidence, + ), + ) + except ImageInputUnsupportedError: + if mode != "visual": + raise + journal.save_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + capability_key, + inputs=capability_inputs, + outputs={ + "evidenceMode": "structured", + "reason": "providerRejectedImageInput", + }, + ) + logger.info( + "Analysis assessment: model rejected images; continuing with structured marker, PCA and diagnostic evidence." + ) + return self.review(scope, review_index, selected, coverage) + action = validate(result.output) + journal.save_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + key, + inputs=evidence, + outputs={"action": action.model_dump(mode="json")}, + ) + self.history.append( + { + "scope": scope, + "evidenceMode": mode, + "visualInspection": visual_inspection, + "review": action.model_dump(mode="json"), + "imageHashes": image_hashes, + "checkpointKey": key, + "checkpointSha256": hashlib.sha256( + record_io.canonical_json_bytes( + { + "inputs": evidence, + "outputs": {"action": action.model_dump(mode="json")}, + } + ) + ).hexdigest(), + } + ) + self.last_action = action + operation = ( + f"experiment {action.experimentId} from {action.selectedCandidateId}" + if action.action == "experiment" + else f"{action.action} {action.selectedCandidateId}" + ) + logger.info(f"Analysis assessment: {action.plainLanguageSummary}") + logger.debug(f"Analysis assessment ({operation}): {action.rationale}") + return action + + def assess_scope( + self, + scope: str, + cells: ArtifactRef, + initial: RnaSetting | None, + ) -> tuple[ + Literal["accept", "enlarge", "defer"], ParameterCandidateEvaluation | None + ]: + coverage_inputs = { + "parentCells": self.cells.to_dict(), + "cells": cells.to_dict(), + "columns": self.coverage_columns, + "protectedCombinations": self.study.protectedCombinations, + "columnKinds": self.study.columnKinds, + "provenance": self.provenance, + } + measured = journal.load_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + f"{self.checkpoint_scope(scope)}/coverage", + inputs=coverage_inputs, + ) + if measured is None: + coverage, insufficient = screening_coverage( + self.store, + self.cells, + cells, + self.coverage_columns, + self.study.protectedCombinations, + ) + measured = journal.save_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + f"{self.checkpoint_scope(scope)}/coverage", + inputs=coverage_inputs, + outputs={"coverage": coverage, "insufficient": insufficient}, + ) + coverage, insufficient = measured["coverage"], measured["insufficient"] + logger.info( + f"Analysis population: {coverage['screeningCells']:,} of {coverage.get('populationCells', self.handoff.nCells):,} retained cells; " + + ("all retained cells." if cells == self.cells else "screening subset.") + ) + self.scope_sizes[scope] = coverage["screeningCells"] + self.history.append( + {"scope": scope, "coverage": coverage, "coverageConcerns": insufficient} + ) + if cells != self.cells and insufficient: + return "enlarge", None + if initial is None: + if scope == "full": + raise CandidateBudgetExceeded( + "No supported screening proposal is available for targeted full-cohort recovery. " + "The mandatory discovery comparisons cannot be replaced by a full-cohort baseline; " + "saved results and unmet coverage requirements are preserved." + ) + baseline = self.baseline().model_copy( + update={ + "parameters": self.baseline().parameters.model_copy( + update={ + "dimensions": min( + self.baseline().parameters.dimensions, + coverage["screeningCells"] - 1, + ), + "neighborsK": min( + self.baseline().parameters.neighborsK, + coverage["screeningCells"] - 1, + ), + } + ) + } + ) + selected = self._resolution_panel(scope, cells, baseline) + # The first panel describes defaults, not an accepted combined recipe. + self.resolution_candidates.pop(scope, None) + self._sensitivity_panel(scope, cells, selected) + else: + requested = initial + if self.study.correctionLicense == "safe": + initial = initial.model_copy( + update={ + "parameters": initial.parameters.model_copy( + update={"useHarmony": True} + ) + } + ) + selected = ( + self._resolution_panel(scope, cells, initial) + if scope == "full" and self.recovery_scope is not None + else self.execute_matched(scope, cells, initial) + ) + if scope == "full" and self.recovery_scope is not None: + self.combined_candidates[scope] = selected.candidateId + if not requested.parameters.useHarmony and initial.parameters.useHarmony: + selected = self.execute(scope, cells, requested) + limit = ( + self.request.config.maxFullPartitions + if scope == "full" + else self.request.config.maxScreeningEvaluations + ) + for review_index in range(limit + 1): + action = self.review(scope, review_index, selected, coverage) + selected = next( + item + for item in self.evaluations[scope] + if item.candidateId == action.selectedCandidateId + ) + if action.action in {"accept", "enlarge", "defer"}: + if action.action == "enlarge" and cells == self.cells: + self.history.append( + { + "scope": scope, + "reason": "All retained cells were already assessed; more sampling cannot supply the unresolved evidence.", + } + ) + return "defer", selected + if ( + action.action == "accept" + and cells != self.cells + and selected.metrics.minClusterCells is not None + and selected.metrics.minClusterCells < 20 + ): + self.history.append( + { + "scope": scope, + "reason": "The selected sampled partition contains a small population requiring more cells for assessment; all its cells are retained.", + } + ) + return "enlarge", selected + if action.action == "accept" and scope != "full": + self.discovery_scope = scope + return action.action, selected + if action.action == "combine": + setting = self._combined_setting(scope, review_index, action, cells) + if self.study.correctionLicense == "safe": + setting = setting.model_copy( + update={ + "parameters": setting.parameters.model_copy( + update={"useHarmony": True} + ) + } + ) + else: + selected = self.execute(scope, cells, setting) + setting = self.settings[selected.candidateId] + selected = self._resolution_panel(scope, cells, setting) + self.combined_candidates[scope] = selected.candidateId + continue + assert action.experimentId is not None + experiment = self.experiments(selected)[action.experimentId] + if ( + scope == "full" + and ( + self.discovery_scope is not None + or scope in self.combined_candidates + ) + and experiment["parameter"] != "useHarmony" + ): + self.budget.admit_repair( + ArtifactReferenceModel.from_artifact_ref(cells).model_dump( + mode="json" + ), + self.settings[selected.candidateId].model_dump(mode="json"), + experiment, + ) + self.full_repairs = len(self.budget.repairs) + if experiment["parameter"] in { + "hvgRanking", + "hvgCount", + "includeFamily", + "excludeFamily", + "includeFeature", + "excludeFeature", + }: + setting = self._prepared_setting( + scope, + f"review{review_index}/feature_experiment", + selected, + experiment, + cells, + ) + else: + setting = self.apply_experiment(selected, experiment) + prior = selected + pending_policy = next( + ( + row + for row in self.comparison_rows.get(scope, []) + if row["axis"] == "featurePolicy" and row["status"] == "pending" + ), + None, + ) + policy_experiment = experiment["parameter"] in { + "includeFamily", + "excludeFamily", + "includeFeature", + "excludeFeature", + } + if pending_policy is not None and policy_experiment: + if prior.candidateId != pending_policy["baselineCandidateId"]: + raise ValueError( + "The required policy comparison must retain the exact default baseline" + ) + prior_setting = self.settings[prior.candidateId] + if setting.hvgCount != prior_setting.hvgCount or np.array_equal( + self.store.load_artifact(artifact_model_to_ref(setting.features))[ + "values" + ][:], + self.store.load_artifact( + artifact_model_to_ref(prior_setting.features) + )["values"][:], + ): + pending_policy.update( + status="notApplicable", + reason="The nominated policy cannot change selected genes while retaining the fixed baseline HVG count.", + observedProof={ + "baselineFeatures": prior_setting.features.model_dump( + mode="json" + ), + **( + { + "kind": "insufficientEligibleGenes", + "eligibleFeatureCount": setting.hvgCount, + } + if setting.hvgCount != prior_setting.hvgCount + else { + "kind": "identicalSelectedGenes", + "verifiedEqualMasks": True, + "alternativeSetting": setting.model_dump( + mode="json" + ), + } + ), + }, + ) + continue + selected = ( + self._resolution_panel(scope, cells, setting) + if scope == "full" + and self.recovery_scope is not None + and experiment["parameter"] != "leidenResolution" + else self.execute_matched(scope, cells, setting) + ) + if ( + scope == "full" + and self.recovery_scope is not None + and experiment["parameter"] != "leidenResolution" + ): + self.combined_candidates[scope] = selected.candidateId + if pending_policy is not None and policy_experiment: + pending_policy.update( + status="completed", + alternativeCandidateId=selected.candidateId, + reason="Baseline loading, marker and objective evidence nominated this exact policy comparison.", + ) + if ( + scope == "full" + and ( + self.discovery_scope is not None + or scope in self.combined_candidates + ) + and experiment["parameter"] != "useHarmony" + ): + self.full_repair = { + "baselineCandidateId": prior.candidateId, + "selectedCandidateId": selected.candidateId, + "experimentId": action.experimentId, + } + if scope != "full" and scope in self.combined_candidates: + selected = self._resolution_panel( + scope, cells, self.settings[selected.candidateId] + ) + self.combined_candidates[scope] = selected.candidateId + return "defer", selected + + def run(self) -> tuple[ParameterTuningReport, dict[str, Any]]: + with journal.diagnostic_attempt( + self.store, + self.prefix, + self.workflow.workflowRunId, + {"provenance": self.provenance, "cells": self.cells.to_dict()}, + ) as counts: + self.diagnostic_counts = counts + selected: ParameterCandidateEvaluation | None = None + final_status = "defer" + reason = "Required scientific evidence remains unresolved." + try: + initial: RnaSetting | None = None + recovery_candidate: ParameterCandidateEvaluation | None = None + previous_sample = None + sizes = screening_sizes(self.handoff.nCells, self.request.config) + journal.save_checkpoint( + self.store, + self.prefix, + self.workflow.workflowRunId, + "parameter_tuning/sampling_policy", + inputs={ + "cells": self.cells.to_dict(), + "config": self.request.config.model_dump(mode="json"), + }, + outputs={ + "policy": "fixed" + if self.request.config.screeningCells is not None + else "boundedTenPercent", + "retainedCells": self.handoff.nCells, + "populationSizes": list(sizes), + "seed": self.request.config.randomSeed, + }, + ) + for index, size in enumerate(sizes): + sample = uniform_screening_selection( + self.store, + self.cells, + size=size, + seed=self.request.config.randomSeed, + ) + if sample == previous_sample: + break + previous_sample = sample + status, screened = self.assess_scope(f"sample{index}", sample, None) + if ( + status == "enlarge" + and screened is not None + and f"sample{index}" in self.combined_candidates + ): + recovery_candidate = screened + self.recovery_scope = f"sample{index}" + if status == "accept" and screened is not None: + initial = self.settings[screened.candidateId] + self.discovery_scope = f"sample{index}" + self.recovery_scope = None + break + if status == "defer": + return self.report( + None, + self.last_action.rationale + if self.last_action is not None + else reason, + ), self.summary() + if sample == self.cells: + return self.report( + None, + "All retained cells were assessed, but the required scientific evidence remains unresolved; another sample cannot resolve this concern.", + ), self.summary() + if initial is None: + if recovery_candidate is None: + return self.report( + None, + "Screening coverage remains inadequate and no measured combined proposal is available for targeted recovery within the full-cohort allowance. Saved evidence is preserved; essential comparisons remain unresolved.", + ), self.summary() + initial = self.settings[recovery_candidate.candidateId] + logger.info( + "Screening evidence requires targeted full-cohort validation of the agent's proposed settings; previous comparisons remain available." + ) + final_status, selected = self.assess_scope("full", self.cells, initial) + if final_status != "accept" and self.last_action is not None: + reason = self.last_action.rationale + except CandidateBudgetExceeded as exc: + reason = f"Required comparison coverage is incomplete within the configured work limits: {exc}" + if final_status != "accept": + selected = None + return self.report(selected, reason), self.summary() + + def summary(self) -> dict[str, Any]: + budget = self.budget.summary() + diagnostic_artifacts: dict[str, set[ArtifactRef]] = { + "pcaDiagnostics": set(), + "stabilityClusters": set(), + "markerTable": set(), + "doubletScore": set(), + } + subsample_evaluations = {} + for scope, evaluations in self.evaluations.items(): + subsample_evaluations[scope] = len( + { + evaluation.candidateId + for evaluation in evaluations + if evaluation.status == "done" + and evaluation.metrics.subsampleStability is not None + } + ) + for evaluation in evaluations: + if evaluation.status != "done": + continue + for name, artifact in evaluation.artifacts.items(): + diagnostic = ( + "pcaDiagnostics" + if name == "representationDiagnostic" + else name.split(":", 1)[0] + ) + if diagnostic in diagnostic_artifacts: + diagnostic_artifacts[diagnostic].add( + artifact_model_to_ref(artifact) + ) + diagnostic_counts = { + name: len(refs) for name, refs in diagnostic_artifacts.items() + } + for scope, counts in budget["scopes"].items(): + if counts["reserved"]["partitions"]: + label = { + "sample0": "screening sample 1", + "sample1": "screening sample 2", + "full": "full cohort", + }[scope] + completed, reserved = counts["completed"], counts["reserved"] + logger.info( + f"Tuning {label}: {completed['graphs']}/{reserved['graphs']} graphs " + f"and {completed['partitions']}/{reserved['partitions']} partitions " + "completed/reserved." + ) + limits = budget["limits"] + logger.info( + f"Tuning limits: {limits['perScreen']} partitions per screen, " + f"{limits['totalScreens']} across screens; full cohort " + f"{limits['fullGraphs']} graphs and {limits['fullPartitions']} partitions." + ) + logger.info( + f"Diagnostic evidence: {diagnostic_counts['pcaDiagnostics']} PCA summaries, " + f"{diagnostic_counts['stabilityClusters']} alternate-seed partitions, " + f"{diagnostic_counts['markerTable']} marker tables, " + f"{diagnostic_counts['doubletScore']} doublet scores; " + f"{sum(subsample_evaluations.values())} subsample-stability evaluations. " + "Saved evidence may be reused; these are not computation counts." + ) + for operation, counts in sorted(self.diagnostic_counts.items()): + logger.debug( + f"Diagnostic work this invocation, {operation}: " + f"{counts['completed']}/{counts['attempted']} calls completed/attempted, " + f"{counts['failed']} failed; {counts['cacheHits']} metric cache hits, " + f"{counts['restored']} saved evidence restores, " + f"{counts['artifactReuses']} confirmed artifact reuses." + ) + return { + "history": self.history, + "budget": budget, + "diagnosticEvidence": { + "uniqueArtifacts": diagnostic_counts, + "subsampleStabilityEvaluations": subsample_evaluations, + "interpretation": ( + "Counts describe saved diagnostic evidence used by these evaluations " + "and may include reused artifacts or metrics, not new computations." + ), + }, + "fullRepairs": self.full_repairs, + "diagnosticOperations": { + "scope": "thisInvocation", + "operations": self.diagnostic_counts, + "interpretation": ( + "Observed agent calls and explicit reuse events, not inferred core " + "numerical rebuilds or simulated-doublet writes. Earlier invocations " + "are recorded separately in the stage journal; absent historical " + "operation records mean unknown counts." + ), + }, + } + + def report( + self, selected: ParameterCandidateEvaluation | None, reason: str + ) -> ParameterTuningReport: + evaluations = self.evaluations["full"] + common: dict[str, Any] = { + "fromAssay": self.handoff.assay, + "cellSelection": self.handoff.cellSelection, + "evaluations": evaluations, + "totalCandidates": len(evaluations), + "markerAssay": self.handoff.assay, + "limitations": [ + "Screening comparisons describe their exact sampled cells; final artifacts and validation use the full QC-retained cohort.", + *self.study.limitations, + *unsupported_comparison_limitations(self.design_comparisons), + *( + [_STRUCTURED_VISUAL_LIMITATION] + if any( + row.get("evidenceMode") == "structured" for row in self.history + ) + else [] + ), + *( + f"Unsupported matched protection: {column}." + for column in self.study.unsupportedProtection + ), + *(str(row["reason"]) for row in self.history if "reason" in row), + *( + "Screening coverage concern: " + concern + for row in self.history + for concern in row.get("coverageConcerns", []) + ), + ], + } + common["limitations"] = list(dict.fromkeys(common["limitations"])) + if selected is None: + return ParameterTuningReport( + **common, + status="needsInput", + rationale=reason, + stopReason=reason, + needsInput=ParameterTuningNeedsInput(question=reason), + ) + assert self.last_action is not None + common["limitations"].extend( + f"Population {row.clusterId}: {row.explanation}" + for row in self.last_action.populationConcerns + if row.candidateId == selected.candidateId + ) + report = ParameterTuningReport( + **common, + status="done", + recommendedCandidateId=selected.candidateId, + selectedArtifacts=selected.artifacts, + confidence="medium", + rationale=self.last_action.rationale, + evidenceIds=self.last_action.evidenceIds, + stopReason="The objective-driven assessment accepted the full-cohort evidence.", + recommendedByAssay={self.handoff.assay: selected.candidateId}, + ) + return finalize_parameter_tuning_selection( + report, marker_assay=self.handoff.assay, native_assay=self.handoff.assay + ) diff --git a/scarf/agent/orchestrator/tuning.py b/scarf/agent/orchestrator/tuning.py index cca7544b..47176363 100644 --- a/scarf/agent/orchestrator/tuning.py +++ b/scarf/agent/orchestrator/tuning.py @@ -1,1337 +1,973 @@ -"""Parameter tuning and multimodal integration workflow stages.""" +"""Sequential RNA parameter tuning and review stages.""" -import hashlib -import json +import io from collections.abc import Mapping, Sequence -from typing import Any, Literal, cast +from typing import Any, cast import numpy as np -from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score from ...datastore.datastore import DataStore -from ...utils.logging import logger -from ..experimental_context import ExperimentalContextResult -from ..parameter_tuning import ( - ArtifactRecord, - FinalGraphComparison, - FinalGraphSelection, - IntegrationCandidateEvaluation, - IntegrationMetrics, - ParameterCandidate, - ParameterTuningAgent, - ParameterTuningAssayInput, +from ...metadata.rows import read_metadata_rows_chunkwise +from ...storage.refs import ArtifactRef +from ...storage.selections import read_stored_selection_indices +from ...storage.types import as_zarr_array +from ..config.agent_exec import ( + ImageEvidence, +) +from ..experimental_context.contracts import ExperimentalContextResult +from ..experimental_context.study import StudyContract +from ..parameter_tuning.contracts import ( + ParameterCandidateEvaluation, ParameterTuningReport, - final_graph_options, - finalize_parameter_tuning_selection, - validate_final_graph_selection, ) -from ..persistence import ( - AgentInvocation, - AgentReportLink, - AgentReportReference, - AgentWorkflowRun, - list_agent_reports, - load_agent_record, - load_agent_report, - save_agent_report, +from ..parameter_tuning.execution import ( + _metadata_column_fingerprint, + candidate_metric_cache, ) -from ..types import ArtifactReferenceModel, ExperimentalTuningHandoff +from ..types import ArtifactReferenceModel from . import journal +from .decisions import DecisionStagesMixin from .models import ( AutomatedPreprocessingPlan, - AutomatedWorkflowConfig, OrchestrationRequestRecord, OrchestrationResumeRecord, PreprocessedAssayHandoff, + StageEvidenceReference, + WorkflowIdentity, WorkflowNeedsInput, WorkflowQuestion, WorkflowStageAttempt, WorkflowStageLink, + WorkflowStageName, artifact_model_to_ref, ) -class TuningStagesMixin: - """Execute parameter searches, integration comparisons, and graph selection.""" +def _analysis_visual_content( + store: DataStore, + selected: ParameterCandidateEvaluation, + candidates: Sequence[ParameterCandidateEvaluation], + *, + qc_columns: Sequence[str] = (), + qc_artifact_metrics: Sequence[tuple[str, ArtifactReferenceModel]] = (), +) -> list[ImageEvidence]: + """Render one bounded diagnostic board for multimodal review.""" + try: + import matplotlib.pyplot as plt + except ImportError as exc: + raise RuntimeError( + "Visual adjudication requires the installed plotting dependencies" + ) from exc + + def image_content(figure: Any, identifier: str) -> ImageEvidence: + buffer = io.BytesIO() + figure.savefig(buffer, format="png", dpi=120) + plt.close(figure) + return ImageEvidence( + identifier=identifier, + data=buffer.getvalue(), + media_type="image/png", + ) + + def sampled_values(array: Any, selection: tuple[Any, ...]) -> np.ndarray: + return np.asarray( + cast(Any, as_zarr_array(array, name="diagnostic")).get_orthogonal_selection( + selection + ) + ) + + coordinate_name = ( + "harmony" + if selected.parameters.useHarmony + else selected.parameters.reductionMethod + ) + coordinate_record = selected.artifacts.get(coordinate_name) + cluster_record = selected.artifacts.get("clusters") + if coordinate_record is None or cluster_record is None: + raise ValueError( + "Selected candidate lacks visualizable coordinates or clusters" + ) + coordinate_group = store.load_artifact( + artifact_model_to_ref( + ArtifactReferenceModel.model_validate(coordinate_record.model_dump()) + ) + ) + cluster_group = store.load_artifact( + artifact_model_to_ref( + ArtifactReferenceModel.model_validate(cluster_record.model_dump()) + ) + ) + coordinate_values = as_zarr_array(coordinate_group["data"], name="data") + cluster_values = as_zarr_array(cluster_group["values"], name="values") + if len(coordinate_values.shape) != 2 or coordinate_values.shape[1] < 2: + raise ValueError("Selected coordinates need at least two dimensions") + if coordinate_values.shape[0] != cluster_values.shape[0]: + raise ValueError("Selected coordinates and clusters do not align") + sample_count = min(5_000, coordinate_values.shape[0]) + sample_indices = np.linspace( + 0, + coordinate_values.shape[0] - 1, + sample_count, + dtype=np.int64, + ) + coordinates = sampled_values( + coordinate_values, + (sample_indices, slice(0, 2)), + ).astype(np.float32, copy=False) + labels = sampled_values(cluster_values, (sample_indices,)) + _, label_codes = np.unique(labels, return_inverse=True) + + comparison = max( + ( + candidate + for candidate in candidates + if candidate.candidateId != selected.candidateId + and candidate.status == "done" + and candidate.eligible + and candidate.parameters.dimensions == selected.parameters.dimensions + and candidate.parameters.neighborsK == selected.parameters.neighborsK + and candidate.parameters.useHarmony == selected.parameters.useHarmony + and candidate.artifacts.get("graphFeatures") + == selected.artifacts.get("graphFeatures") + and candidate.cellSelection == selected.cellSelection + ), + key=lambda value: ( + value.metrics.markerCoherence or 0.0, + value.metrics.seedStability or 0.0, + ), + default=None, + ) + + figure, axes = plt.subplots(2, 3, figsize=(15, 9), constrained_layout=True) + variance = np.asarray(selected.metrics.componentVariance, dtype=float) + if variance.size: + axes[0, 0].plot(np.arange(1, variance.size + 1), variance, marker=".") + axes[0, 0].set_title("PCA explained variance") + axes[0, 0].set_xlabel("Component") + else: + axes[0, 0].text(0.5, 0.5, "Variance unavailable", ha="center") + axes[0, 0].set_axis_off() + + axes[0, 1].scatter( + coordinates[:, 0], + coordinates[:, 1], + c=label_codes, + cmap="tab20", + s=2, + alpha=0.65, + rasterized=True, + ) + axes[0, 1].set_title(f"Selected partition: {selected.candidateId}") + + if comparison is not None and "clusters" in comparison.artifacts: + comparison_group = store.load_artifact( + artifact_model_to_ref( + ArtifactReferenceModel.model_validate( + comparison.artifacts["clusters"].model_dump() + ) + ) + ) + comparison_labels = sampled_values( + comparison_group["values"], + (sample_indices,), + ) + _, comparison_codes = np.unique( + comparison_labels, + return_inverse=True, + ) + axes[0, 2].scatter( + coordinates[:, 0], + coordinates[:, 1], + c=comparison_codes, + cmap="tab20", + s=2, + alpha=0.65, + rasterized=True, + ) + axes[0, 2].set_title(f"Matched alternative: {comparison.candidateId}") + else: + axes[0, 2].text(0.5, 0.5, "No matched alternative", ha="center") + axes[0, 2].set_axis_off() + + completed = [ + candidate + for candidate in candidates + if candidate.status == "done" + and candidate.eligible + and candidate.parameters.dimensions == selected.parameters.dimensions + and candidate.parameters.neighborsK == selected.parameters.neighborsK + and candidate.parameters.useHarmony == selected.parameters.useHarmony + ] + resolutions = [value.parameters.leidenResolution for value in completed] + axes[1, 0].plot( + resolutions, + [value.metrics.markerCoherence or 0.0 for value in completed], + marker="o", + label="marker coherence", + ) + axes[1, 0].plot( + resolutions, + [value.metrics.seedStability or 0.0 for value in completed], + marker="o", + label="seed stability", + ) + axes[1, 0].plot( + resolutions, + [value.metrics.crossUnitSupport or 0.0 for value in completed], + marker="o", + label="cross-unit support", + ) + axes[1, 0].set_title("Partition evidence") + axes[1, 0].set_xlabel("Leiden resolution") + axes[1, 0].legend(fontsize=8) + + doublet_records = { + name.removeprefix("doubletScore:"): value + for name, value in selected.artifacts.items() + if name.startswith("doubletScore:") + } + doublet_selections = { + name.removeprefix("doubletCellSelection:"): value + for name, value in selected.artifacts.items() + if name.startswith("doubletCellSelection:") + } + doublet_sample = np.full(sample_count, np.nan, dtype=np.float64) + doublet_hist_values: list[np.ndarray] = [] + if doublet_records: + if selected.cellSelection is None: + raise ValueError("Doublet visuals require an exact cell selection") + parent_selection = artifact_model_to_ref(selected.cellSelection) + parent_indices = read_stored_selection_indices( + store.zw, + parent_selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + sampled_global_indices = parent_indices[sample_indices] + for score_index, record in sorted(doublet_records.items()): + score_group = store.load_artifact( + artifact_model_to_ref( + ArtifactReferenceModel.model_validate(record.model_dump()) + ) + ) + score_values = as_zarr_array(score_group["values"], name="values") + selection_record = doublet_selections.get(score_index) + if selection_record is None: + if score_values.shape != parent_indices.shape: + raise ValueError( + "Doublet score artifact lacks its exact cell selection" + ) + local_positions = sample_indices + matched_positions = np.arange(sample_count, dtype=np.int64) + else: + score_selection = artifact_model_to_ref( + ArtifactReferenceModel.model_validate(selection_record.model_dump()) + ) + score_indices = read_stored_selection_indices( + store.zw, + score_selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + if score_values.shape != score_indices.shape: + raise ValueError( + "Doublet scores do not align with their cell selection" + ) + candidate_positions = np.searchsorted( + score_indices, + sampled_global_indices, + ) + within = candidate_positions < len(score_indices) + matched = np.zeros(sample_count, dtype=bool) + matched[within] = ( + score_indices[candidate_positions[within]] + == sampled_global_indices[within] + ) + matched_positions = np.flatnonzero(matched) + local_positions = candidate_positions[matched] + sampled_scores = sampled_values(score_values, (local_positions,)).astype( + np.float64, + copy=False, + ) + doublet_sample[matched_positions] = sampled_scores + doublet_hist_values.append(sampled_scores) + finite_doublets = np.concatenate(doublet_hist_values) + finite_doublets = finite_doublets[np.isfinite(finite_doublets)] + axes[1, 1].hist(finite_doublets, bins=40) + axes[1, 1].set_title("Advisory doublet scores") + else: + axes[1, 1].text(0.5, 0.5, "Doublet scores unavailable", ha="center") + axes[1, 1].set_axis_off() + + family_values = { + **selected.metrics.loadingFamilyEnrichment, + **selected.metrics.markerFamilyEnrichment, + } + if family_values: + ordered = sorted( + family_values.items(), + key=lambda item: (-item[1], item[0]), + )[:12] + axes[1, 2].barh( + [name for name, _value in reversed(ordered)], + [value for _name, value in reversed(ordered)], + ) + axes[1, 2].set_title("Feature-family enrichment") + else: + axes[1, 2].text(0.5, 0.5, "Family evidence unavailable", ha="center") + axes[1, 2].set_axis_off() + + content = [image_content(figure, "analysis-overview")] + + paired_corrections: list[ + tuple[ParameterCandidateEvaluation, ParameterCandidateEvaluation] + ] = [] + candidates_by_parameters: dict[ + tuple[str, int, int, float, str, str], + dict[bool, ParameterCandidateEvaluation], + ] = {} + for candidate in candidates: + if candidate.status != "done" or not candidate.eligible: + continue + key = ( + candidate.parameters.reductionMethod, + candidate.parameters.dimensions, + candidate.parameters.neighborsK, + candidate.parameters.leidenResolution, + candidate.artifacts["graphFeatures"].model_dump_json() + if "graphFeatures" in candidate.artifacts + else "", + candidate.cellSelection.model_dump_json() + if candidate.cellSelection is not None + else "", + ) + candidates_by_parameters.setdefault(key, {})[ + candidate.parameters.useHarmony + ] = candidate + for correction_pair in candidates_by_parameters.values(): + if False in correction_pair and True in correction_pair: + paired_corrections.append((correction_pair[False], correction_pair[True])) + if paired_corrections: + native, harmony = min( + paired_corrections, + key=lambda pair: ( + abs(pair[0].parameters.dimensions - selected.parameters.dimensions), + abs(pair[0].parameters.neighborsK - selected.parameters.neighborsK), + abs( + pair[0].parameters.leidenResolution + - selected.parameters.leidenResolution + ), + ), + ) + correction_figure, correction_axes = plt.subplots( + 2, + 2, + figsize=(11, 10), + constrained_layout=True, + ) + batch_columns = [ + column + for column in ( + *native.metrics.batchMixing, + *harmony.metrics.batchMixing, + ) + if column in store.cells.columns + ] + batch_codes: np.ndarray | None = None + batch_label = "Batch unavailable" + if batch_columns and selected.cellSelection is not None: + parent_selection = artifact_model_to_ref(selected.cellSelection) + parent_indices = read_stored_selection_indices( + store.zw, + parent_selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + batch_label = batch_columns[0] + batch_values = read_metadata_rows_chunkwise( + store.cells, + batch_label, + parent_indices[sample_indices], + ).astype(str) + _, batch_codes = np.unique(batch_values, return_inverse=True) + for column_index, candidate in enumerate((native, harmony)): + name = ( + "harmony" + if candidate.parameters.useHarmony + else candidate.parameters.reductionMethod + ) + candidate_coordinate = candidate.artifacts.get(name) + candidate_cluster = candidate.artifacts.get("clusters") + if candidate_coordinate is None or candidate_cluster is None: + raise ValueError("Matched correction candidate lacks visual artifacts") + candidate_coordinate_group = store.load_artifact( + artifact_model_to_ref( + ArtifactReferenceModel.model_validate( + candidate_coordinate.model_dump() + ) + ) + ) + candidate_cluster_group = store.load_artifact( + artifact_model_to_ref( + ArtifactReferenceModel.model_validate( + candidate_cluster.model_dump() + ) + ) + ) + candidate_coordinates = sampled_values( + candidate_coordinate_group["data"], + (sample_indices, slice(0, 2)), + ) + candidate_labels = sampled_values( + candidate_cluster_group["values"], + (sample_indices,), + ) + _, candidate_codes = np.unique( + candidate_labels, + return_inverse=True, + ) + title = "Harmony" if candidate.parameters.useHarmony else "Native PCA" + correction_axes[0, column_index].scatter( + candidate_coordinates[:, 0], + candidate_coordinates[:, 1], + c=candidate_codes, + cmap="tab20", + s=2, + alpha=0.65, + rasterized=True, + ) + correction_axes[0, column_index].set_title(f"{title}, partition colors") + if batch_codes is not None: + correction_axes[1, column_index].scatter( + candidate_coordinates[:, 0], + candidate_coordinates[:, 1], + c=batch_codes, + cmap="tab20", + s=2, + alpha=0.65, + rasterized=True, + ) + correction_axes[1, column_index].set_title( + f"{title}, {batch_label} colors" + ) + else: + correction_axes[1, column_index].text( + 0.5, + 0.5, + batch_label, + ha="center", + ) + correction_axes[1, column_index].set_axis_off() + correction_figure.suptitle( + "Parameter-matched native and Harmony representations" + ) + content.append(image_content(correction_figure, "native-harmony-comparison")) + + marker_record = selected.artifacts.get("markerTable") + marker_genes = list( + dict.fromkeys( + gene for genes in selected.metrics.topMarkerGenes.values() for gene in genes + ) + )[:24] + if marker_record is not None and marker_genes: + marker_ref = artifact_model_to_ref( + ArtifactReferenceModel.model_validate(marker_record.model_dump()) + ) + marker_table = store.get_markers( + marker_ref, + min_score=0.25, + min_frac_exp=0.2, + ) + required_marker_columns = {"group_id", "feature_name", "score"} + if required_marker_columns.issubset(marker_table.columns): + marker_groups = sorted( + selected.metrics.topMarkerGenes, + key=lambda value: ( + (0, int(value)) if value.lstrip("-").isdigit() else (1, value) + ), + ) + marker_scores = np.zeros( + (len(marker_groups), len(marker_genes)), + dtype=np.float64, + ) + group_positions = { + value: index for index, value in enumerate(marker_groups) + } + gene_positions = {value: index for index, value in enumerate(marker_genes)} + for row in marker_table.itertuples(index=False): + group = str(getattr(row, "group_id")) + gene = str(getattr(row, "feature_name")) + if group not in group_positions or gene not in gene_positions: + continue + score = float(getattr(row, "score")) + if np.isfinite(score): + marker_scores[group_positions[group], gene_positions[gene]] = max( + marker_scores[ + group_positions[group], + gene_positions[gene], + ], + score, + ) + + def tagged_gene(gene: str) -> str: + upper = gene.upper() + if upper.startswith("MT-"): + return f"{gene} [mitochondrial]" + if upper.startswith(("MRPS", "MRPL")): + return f"{gene} [mitoribosomal]" + if upper.startswith(("RPS", "RPL")): + return f"{gene} [ribosomal]" + if upper.startswith("CCN"): + return f"{gene} [CCN]" + if upper.startswith("HLA-"): + return f"{gene} [HLA]" + if upper.startswith("H2-"): + return f"{gene} [H2]" + if upper.startswith("HIST"): + return f"{gene} [histone]" + if upper in { + "XIST", + "DDX3Y", + "USP9Y", + "EIF1AY", + "KDM5D", + "SRY", + "ZFY", + "UTY", + "TMSB4Y", + "NLGN4Y", + }: + return f"{gene} [sex-linked]" + return gene + + marker_figure, marker_axis = plt.subplots( + figsize=(max(9, len(marker_genes) * 0.45), 6), + constrained_layout=True, + ) + marker_image = marker_axis.imshow( + marker_scores, + aspect="auto", + interpolation="nearest", + cmap="viridis", + ) + marker_axis.set_xticks( + np.arange(len(marker_genes)), + [tagged_gene(gene) for gene in marker_genes], + rotation=75, + ha="right", + fontsize=8, + ) + marker_axis.set_yticks( + np.arange(len(marker_groups)), + marker_groups, + ) + marker_axis.set_xlabel("Top marker feature") + marker_axis.set_ylabel("Cluster") + marker_axis.set_title("Feature-level marker score heatmap") + marker_figure.colorbar(marker_image, ax=marker_axis, label="marker score") + content.append(image_content(marker_figure, "marker-score-heatmap")) + + available_qc_columns = [ + column + for column in dict.fromkeys([*qc_columns, *selected.metrics.qcPcaAssociation]) + if column in store.cells.columns + ] + qc_sources: list[tuple[str, ArtifactReferenceModel | None]] = [ + (column, None) for column in available_qc_columns + ] + known_qc_names = set(available_qc_columns) + for qc_name, reference in qc_artifact_metrics: + if qc_name not in known_qc_names: + qc_sources.append((qc_name, reference)) + known_qc_names.add(qc_name) + qc_sources = qc_sources[:4] + if qc_sources or np.isfinite(doublet_sample).any(): + diagnostic_count = len(qc_sources) + 1 + qc_figure, qc_axes = plt.subplots( + 1, + diagnostic_count, + figsize=(max(5, diagnostic_count * 3.2), 4), + constrained_layout=True, + ) + qc_axis_list = np.atleast_1d(qc_axes).tolist() + if selected.cellSelection is None: + raise ValueError("QC visuals require an exact cell selection") + parent_selection = artifact_model_to_ref(selected.cellSelection) + parent_indices = read_stored_selection_indices( + store.zw, + parent_selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + sampled_parent_indices = parent_indices[sample_indices] + for axis, (qc_name, artifact) in zip( + qc_axis_list, + qc_sources, + strict=False, + ): + if artifact is None: + values = np.asarray( + read_metadata_rows_chunkwise( + store.cells, + qc_name, + sampled_parent_indices, + ), + dtype=np.float64, + ) + else: + artifact_ref = artifact_model_to_ref(artifact) + artifact_group = store.load_artifact(artifact_ref) + artifact_values = as_zarr_array( + artifact_group["values"], + name="values", + ) + if artifact_values.shape == parent_indices.shape: + artifact_positions = sample_indices + else: + artifact_status = store.inspect_artifact(artifact_ref) + raw_selection = ( + getattr(artifact_status, "inputs", None) or {} + ).get("cell_selection") + if not isinstance(raw_selection, Mapping): + raise ValueError( + f"QC artifact {qc_name!r} lacks its cell selection" + ) + artifact_selection = ArtifactRef.from_dict(dict(raw_selection)) + artifact_indices = read_stored_selection_indices( + store.zw, + artifact_selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + candidate_positions = np.searchsorted( + artifact_indices, + sampled_parent_indices, + ) + if np.any( + candidate_positions >= len(artifact_indices) + ) or not np.array_equal( + artifact_indices[candidate_positions], + sampled_parent_indices, + ): + raise ValueError( + f"QC artifact {qc_name!r} does not cover selected cells" + ) + artifact_positions = candidate_positions + values = sampled_values( + artifact_values, + (artifact_positions,), + ).astype(np.float64, copy=False) + finite_values = values[np.isfinite(values)] + axis.violinplot(finite_values, showmedians=True) + jitter = ( + (np.arange(len(finite_values), dtype=np.float64) % 23.0) - 11.0 + ) / 115.0 + axis.scatter( + np.ones(len(finite_values)) + jitter, + finite_values, + s=1, + alpha=0.15, + rasterized=True, + ) + axis.set_xticks([]) + axis.set_title(qc_name) + doublet_axis = qc_axis_list[-1] + finite_doublet_mask = np.isfinite(doublet_sample) + if finite_doublet_mask.any(): + doublet_plot = doublet_axis.scatter( + coordinates[finite_doublet_mask, 0], + coordinates[finite_doublet_mask, 1], + c=doublet_sample[finite_doublet_mask], + cmap="magma", + s=2, + alpha=0.7, + rasterized=True, + ) + doublet_axis.set_title("Advisory doublet score") + qc_figure.colorbar( + doublet_plot, + ax=doublet_axis, + label="score", + ) + else: + doublet_axis.text( + 0.5, + 0.5, + "Doublet embedding unavailable", + ha="center", + ) + doublet_axis.set_axis_off() + qc_figure.suptitle("Selected-cell QC and non-removing doublet evidence") + content.append(image_content(qc_figure, "qc-doublet-diagnostics")) + return content + + +def _tuning_revision_provenances( + store: DataStore, + prefix: str, + workflow_run_id: str, + request: OrchestrationRequestRecord, + experimental_reference: StageEvidenceReference, + inputs: dict[str, Any], + previous_starts: Sequence[WorkflowStageAttempt], +) -> list[dict[str, Any]]: + """Authorize changed interpretation only along committed context revisions.""" + previous_inputs = [ + { + key: value + for key, value in previous.inputs.items() + if key not in {"resumeAnswers", "answeredAttempt"} + } + for previous in previous_starts + ] + changed = [value for value in previous_inputs if value != inputs] + if not changed: + return [] + contexts = journal._stage_outcomes( + store.zw, prefix, workflow_run_id, "experimental_context" + ) + current = next( + (row for row in contexts if experimental_reference in row.reportReferences), + None, + ) + if ( + current is None + or current.status != "done" + or current.outputs.get("studyContract") != inputs["studyContract"] + or current.requestSha256 != request.requestSha256 + or current.configSha256 != request.configSha256 + ): + raise ValueError( + "Tuning inputs changed; changed evidence requires its exact completed context revision" + ) + ancestors: list[dict[str, Any]] = [] + visited: set[str] = set() + while current is not None and "reassessContextReport" in current.inputs: + if current.attemptId in visited: + raise ValueError("Context revision ancestry contains a cycle") + visited.add(current.attemptId) + previous_ref = current.inputs["reassessContextReport"] + current = next( + ( + row + for row in contexts + if any( + ref.model_dump(mode="json") == previous_ref + for ref in row.reportReferences + ) + ), + None, + ) + if current is None or current.status != "done": + raise ValueError( + "Context revision is missing its committed previous evidence" + ) + if ( + current.requestSha256 != request.requestSha256 + or current.configSha256 != request.configSha256 + ): + raise ValueError( + "Context revision belongs to different request or configuration inputs" + ) + ancestors.append(current.outputs["studyContract"]) + authorized: list[dict[str, Any]] = [] + for previous in changed: + if ( + previous.get("studyContract") not in ancestors + or previous.get("preprocessedAssays") != inputs["preprocessedAssays"] + or previous.get("featureMetadataFingerprints") + != inputs["featureMetadataFingerprints"] + or any( + inputs["metadataFingerprints"].get(column) != fingerprint + for column, fingerprint in previous.get( + "metadataFingerprints", {} + ).items() + ) + or set(previous) != set(inputs) + ): + raise ValueError( + "Tuning inputs changed beyond an explicit context-evidence revision; " + "restore the original cohort, feature inputs and metadata or start a new workflow" + ) + provenance = { + **previous, + "requestSha256": request.requestSha256, + "configSha256": request.configSha256, + } + if provenance not in authorized: + authorized.append(provenance) + return authorized + - model: Any +class TuningStagesMixin(DecisionStagesMixin): + """Run bounded experiments and return validated full-cohort artifacts.""" def parameter_tuning_stage( self, store: DataStore, - workflow: AgentWorkflowRun, + workflow: WorkflowIdentity, request_record: OrchestrationRequestRecord, parents: Sequence[WorkflowStageLink], plan: AutomatedPreprocessingPlan, preprocessed: Sequence[PreprocessedAssayHandoff], experimental: ExperimentalContextResult, - enrichment_reference: AgentReportReference, - experimental_reference: AgentReportReference, + enrichment_reference: StageEvidenceReference, + experimental_reference: StageEvidenceReference, answers: Mapping[str, Any], *, + study_contract: StudyContract | None = None, resume_record: OrchestrationResumeRecord | None = None, + stage_name: WorkflowStageName = "parameter_tuning", ) -> tuple[WorkflowStageAttempt, ParameterTuningReport]: + from .rna_tuning import RnaTuningRun + + del enrichment_reference + if len(preprocessed) != 1 or study_contract is None: + raise ValueError("RNA tuning requires one assay and a study contract") + handoff = preprocessed[0] + if handoff.cellSelection is None: + raise ValueError("RNA tuning requires a frozen full-cohort selection") + columns = { + *study_contract.technicalBatchColumns, + *study_contract.protectedColumns, + *study_contract.independentUnitColumns, + *study_contract.conditionColumns, + *plan.cellQc.attributes, + } + if study_contract.physicalCaptureColumn is not None: + columns.add(study_contract.physicalCaptureColumn) + metadata_fingerprints = { + column: _metadata_column_fingerprint(store.cells, column) + for column in sorted(columns) + if column in store.cells.columns + } + feature_metadata = store.get_assay(plan.primaryAssay).feats + feature_fingerprints = { + column: _metadata_column_fingerprint(feature_metadata, column) + for column in ("ids", "names") + } + inputs = { + "preprocessedAssays": [handoff.model_dump(mode="json")], + "studyContract": study_contract.model_dump(mode="json"), + "metadataFingerprints": metadata_fingerprints, + "featureMetadataFingerprints": feature_fingerprints, + } prefix = journal._ensure_orchestration_store(store) - existing = journal._validated_done_outcome( + previous_provenances = _tuning_revision_provenances( store, prefix, workflow.workflowRunId, - "parameter_tuning", request_record, - parents, + experimental_reference, + inputs, + journal._stage_starts(store.zw, prefix, workflow.workflowRunId, stage_name), ) - if existing is not None: - logger.info( - f"Workflow {workflow.workflowRunId}: reusing Parameter Tuning report" - ) - report = journal.load_stage_report(store, existing, ParameterTuningReport) - return existing, cast(ParameterTuningReport, report) - paused = journal._validated_done_outcome( + existing = journal._validated_done_outcome( store, prefix, workflow.workflowRunId, - "parameter_tuning", + stage_name, request_record, parents, - required_status="needsInput", - ) - resumable_report: ParameterTuningReport | None = None - if paused is not None and paused.reportReferences: - loaded = journal.load_stage_report(store, paused, ParameterTuningReport) - candidate_report = cast(ParameterTuningReport, loaded) - if ( - candidate_report.finalSelection is not None - and candidate_report.finalSelection.status == "needsInput" - and candidate_report.assayReports - ): - resumable_report = candidate_report - cell_selection = preprocessed[0].cellSelection if preprocessed else None - if cell_selection is None or any( - value.cellSelection != cell_selection for value in preprocessed - ): - raise ValueError("Preprocessed assays must share one exact cell selection") - experimental_handoff = experimental.to_parameter_tuning_handoff().model_copy( - update={"cellSelection": cell_selection} ) - tuning_answer = answers.get("parameter_tuning") - if isinstance(tuning_answer, Mapping): - tuning_directions = json.dumps( - dict(tuning_answer), - sort_keys=True, + if existing is not None: + return existing, cast( + ParameterTuningReport, + journal.load_stage_report(store, existing, ParameterTuningReport), ) - elif isinstance(tuning_answer, str): - tuning_directions = tuning_answer.strip() - else: - tuning_directions = "" started = journal._start_attempt( store.zw, prefix, workflow.workflowRunId, - "parameter_tuning", + stage_name, request_record, parents, - inputs={ - "preprocessedAssays": [ - value.model_dump(mode="json") for value in preprocessed - ], - "experimentalTuningHandoff": experimental_handoff.model_dump( - mode="json" - ), - "cellSelection": cell_selection.model_dump(mode="json"), - "primaryAssay": plan.primaryAssay, - "markerAssay": plan.markerAssay, - "pairedAssays": plan.pairedAssays, - "finalGraphOptionId": answers.get("finalGraphOptionId"), - "parameterTuning": tuning_answer, - "resumeFromAttempt": ( - paused.attemptId - if paused is not None and resumable_report is not None - else None - ), - }, + inputs=inputs, resume_record=resume_record, ) - report = ParameterTuningReport.get_blank() - actions: list[str] = [] - integration_evaluations: list[IntegrationCandidateEvaluation] = [] - candidate_payload: dict[str, list[dict[str, Any]]] = {} - paired = list(plan.pairedAssays) - logger.info( - f"Workflow {workflow.workflowRunId}: Parameter Tuning started for " - f"{len(preprocessed)} assay(s), paired={len(paired)}" + runner = RnaTuningRun( + self, + store, + workflow, + request_record, + plan, + handoff, + study_contract, + answers, + { + **inputs, + "requestSha256": request_record.requestSha256, + "configSha256": request_record.configSha256, + }, + design_comparisons=experimental.characterization.comparisons, + previous_provenances=previous_provenances, ) try: - agent = ParameterTuningAgent( - self.model, - config=request_record.config.agentRunConfig, - ) - recovered = journal._recover_persisted_stage_report( - store, - started, - agent_name="parameter_tuning", - expected_type=ParameterTuningReport, - ) - if recovered is not None: - recovered_report, recovered_reference = recovered - report = cast(ParameterTuningReport, recovered_report) - integration_evaluations = list(report.integrationEvaluations) - candidate_payload = { - assay: [ - evaluation.parameters.model_dump(mode="json") - for evaluation in assay_report.evaluations - ] - for assay, assay_report in report.assayReports.items() - } - actions.append("recover_persisted_parameter_tuning_report") - logger.info( - f"Workflow {workflow.workflowRunId}: recovering completed " - "Parameter Tuning provider result" - ) - return self.save_parameter_tuning_outcome( - store, - prefix, - workflow, - request_record, - started, - report, - plan, - preprocessed, - integration_evaluations, - candidate_payload, - paired, - enrichment_reference, - experimental_reference, - experimental_handoff, - agent, - actions, - persisted_reference=recovered_reference, - ) - if resumable_report is not None: - assert paused is not None - resumed_integration_evaluations = list( - resumable_report.integrationEvaluations - ) - report = resumable_report.model_copy( - update={ - "status": "done", - "needsInput": None, - "finalSelection": None, - "recommendedIntegrationId": None, - "finalClusterColumn": None, - "finalClusterArtifact": None, - "graphAssay": None, - } - ) - report = self.select_final_graph( - agent, - report, - resumed_integration_evaluations, - marker_assay=plan.markerAssay, - answers=answers, - ) - logger.info( - f"Workflow {workflow.workflowRunId}: resumed final graph " - "selection without rerunning candidate evaluation" - ) - resumed_candidate_payload = { - assay: [ - evaluation.parameters.model_dump(mode="json") - for evaluation in assay_report.evaluations - ] - for assay, assay_report in report.assayReports.items() - } - return self.save_parameter_tuning_outcome( - store, - prefix, - workflow, - request_record, - started, - report, - plan, - preprocessed, - resumed_integration_evaluations, - resumed_candidate_payload, - list(plan.pairedAssays), - enrichment_reference, - experimental_reference, - experimental_handoff, - agent, - ["reuse_parameter_screen_and_integrations"], - prior_tuning_reference=paused.reportReferences[0], - ) - handoff_by_assay = {value.assay: value for value in preprocessed} - common_k = None - if paired: - common_k = min( - 21, - min(handoff_by_assay[assay].nCells - 1 for assay in paired), - ) - if common_k < 2: - raise ValueError("Paired integration requires at least three cells") - integration_budget = ( - 2 * request_record.config.integrationResolutionCandidates - if len(paired) >= 2 - else 0 - ) - assay_inputs: list[ParameterTuningAssayInput] = [] - for handoff in preprocessed: - if handoff.normalized is None: - raise ValueError(f"Assay {handoff.assay!r} lacks normalization") - initial_count = ( - request_record.config.primaryInitialCandidates - if handoff.assay == plan.primaryAssay - else request_record.config.secondaryInitialCandidates - ) - neighbors_k = common_k or min(11, handoff.nCells - 1) - candidates = self.initial_parameter_candidates( - workflow.workflowRunId, - handoff, - count=initial_count, - neighbors_k=neighbors_k, - ) - if ( - experimental_handoff.batchAction == "evaluateHarmony" - and request_record.config.maxHarmonyCandidatesPerAssay == 1 - ): - baseline = candidates[0] - candidates.append( - baseline.model_copy( - update={ - "candidateId": f"{baseline.candidateId}_harmony", - "useHarmony": True, - } - ) - ) - candidate_payload[handoff.assay] = [ - value.model_dump(mode="json") for value in candidates - ] - logger.info( - f"Workflow {workflow.workflowRunId}: planned " - f"{len(candidates)} native candidate(s) for assay " - f"{handoff.assay!r} (harmony=" - f"{sum(value.useHarmony for value in candidates)})" - ) - assay_inputs.append( - ParameterTuningAssayInput( - normalized=artifact_model_to_ref(handoff.normalized), - candidates=candidates, - batchColumns=( - list(experimental_handoff.batchColumns) - if experimental_handoff.batchAction == "evaluateHarmony" - else [] - ), - preservationColumns=list( - experimental_handoff.preservationColumns - ), - experimentalHandoff=( - None - if experimental_handoff.batchAction == "evaluateHarmony" - else experimental_handoff - ), - maxCandidates=( - len(candidates) - + request_record.config.maxRefinedCandidatesPerAssay - ), - maxRefinedCandidates=( - request_record.config.maxRefinedCandidatesPerAssay - ), - allowHarmonyRefinement=False, - minClusterCells=request_record.config.minClusterCells, - identityFeatureLimit=request_record.config.maxIdentityFeatures, - ) - ) - planned_native = sum(value.maxCandidates for value in assay_inputs) - if ( - planned_native + integration_budget - > request_record.config.maxCandidateBranches - ): - raise ValueError( - "The native and integrated candidate plan exceeds the global " - f"branch limit {request_record.config.maxCandidateBranches}" - ) - logger.info( - f"Workflow {workflow.workflowRunId}: executing " - f"{planned_native} native candidate branch(es) with " - f"{integration_budget} reserved integration branch(es)" - ) - report = agent.run_batch( - store, - assays=assay_inputs, - primary_assay=plan.primaryAssay, - max_total_candidates=( - request_record.config.maxCandidateBranches - integration_budget - ), - selection_directions=tuning_directions, - ) - logger.info( - f"Workflow {workflow.workflowRunId}: Parameter Tuning returned " - f"status={report.status!r}, evaluated={report.totalCandidates}" - ) - if report.status == "done": - for assay, assay_report in report.assayReports.items(): - normalized = handoff_by_assay[assay].normalized - assert normalized is not None - agent.promote( - store, - report=assay_report, - normalized=artifact_model_to_ref(normalized), - identity_feature_limit=request_record.config.maxIdentityFeatures, - ) - actions.append(f"promote_native:{assay}") - logger.info( - f"Workflow {workflow.workflowRunId}: promoted native " - f"candidate {assay_report.recommendedCandidateId!r} for " - f"assay {assay!r}" - ) - integration_evaluations = self.evaluate_integrations( - store, - workflow.workflowRunId, - plan, - report, - experimental_handoff, - request_record.config, - started=started, - parent_reports=[ - journal._report_link(enrichment_reference), - journal._report_link(experimental_reference), - ], - actions=actions, - ) - logger.info( - f"Workflow {workflow.workflowRunId}: evaluated " - f"{len(integration_evaluations)} integration candidate(s)" - ) - report = self.select_final_graph( - agent, - report, - integration_evaluations, - marker_assay=plan.markerAssay, - answers=answers, - ) - return self.save_parameter_tuning_outcome( - store, - prefix, - workflow, - request_record, - started, - report, - plan, - preprocessed, - integration_evaluations, - candidate_payload, - paired, - enrichment_reference, - experimental_reference, - experimental_handoff, - agent, - actions, - ) - except Exception as exc: - failure_artifacts: dict[str, ArtifactReferenceModel] = { - "cellSelection": cell_selection - } - for assay, assay_report in report.assayReports.items(): - for evaluation in assay_report.evaluations: - for name, artifact in evaluation.artifacts.items(): - failure_artifacts[ - f"{assay}_{evaluation.parameters.candidateId}_{name}" - ] = ArtifactReferenceModel.model_validate(artifact.model_dump()) - for integration_evaluation in integration_evaluations: - if integration_evaluation.graphArtifact is not None: - failure_artifacts[ - f"{integration_evaluation.integrationId}_graph" - ] = ArtifactReferenceModel.model_validate( - integration_evaluation.graphArtifact.model_dump() - ) - if integration_evaluation.clusterArtifact is not None: - failure_artifacts[ - f"{integration_evaluation.integrationId}_clusters" - ] = ArtifactReferenceModel.model_validate( - integration_evaluation.clusterArtifact.model_dump() - ) - outcome = journal.finish_exception( - store, - prefix, - workflow, - started, - exc, - artifacts=failure_artifacts, - actions=actions, - outputs={ - "candidatePlan": candidate_payload, - "integrationEvaluations": [ - value.model_dump(mode="json") - for value in integration_evaluations - ], - }, - ) - return outcome, ParameterTuningReport.get_blank() - - def select_final_graph( - self, - agent: ParameterTuningAgent, - report: ParameterTuningReport, - integration_evaluations: Sequence[IntegrationCandidateEvaluation], - *, - marker_assay: str, - answers: Mapping[str, Any], - ) -> ParameterTuningReport: - directed_option = answers.get("finalGraphOptionId") - if not isinstance(directed_option, str) or not directed_option: - logger.info( - f"Selecting final graph from native and " - f"{len(integration_evaluations)} integration evaluation(s)" - ) - return agent.select_final( - report=report, - integration_evaluations=integration_evaluations, - marker_assay=marker_assay, - ) - options = final_graph_options(report, integration_evaluations) - if directed_option not in options: - raise ValueError("finalGraphOptionId is not an eligible option") - logger.info(f"Applying caller-selected final graph {directed_option!r}") - selected_evidence = list(options[directed_option]["evidenceIds"]) - selection = FinalGraphSelection( - status="done", - selectedOptionId=directed_option, - markerAssay=marker_assay, - confidence="high", - rationale="The caller selected this persisted eligible option.", - evidenceIds=selected_evidence, - comparisons=[ - FinalGraphComparison( - optionId=option_id, - summary="The caller preferred the selected eligible option.", - evidenceIds=[ - *selected_evidence, - *cast(list[str], option["evidenceIds"]), - ], - ) - for option_id, option in options.items() - if option_id != directed_option - ], - ) - selection = validate_final_graph_selection( - selection, - report, - integration_evaluations=integration_evaluations, - marker_assay=marker_assay, - ) - return finalize_parameter_tuning_selection( - report, - marker_assay=marker_assay, - integration_evaluations=integration_evaluations, - recommended_integration_id=selection.integrationId, - native_assay=selection.nativeAssay, - final_selection=selection, - ) - - def save_parameter_tuning_outcome( - self, - store: DataStore, - prefix: str, - workflow: AgentWorkflowRun, - request_record: OrchestrationRequestRecord, - started: WorkflowStageAttempt, - report: ParameterTuningReport, - plan: AutomatedPreprocessingPlan, - preprocessed: Sequence[PreprocessedAssayHandoff], - integration_evaluations: Sequence[IntegrationCandidateEvaluation], - candidate_payload: Mapping[str, list[dict[str, Any]]], - paired: Sequence[str], - enrichment_reference: AgentReportReference, - experimental_reference: AgentReportReference, - experimental_handoff: ExperimentalTuningHandoff, - agent: ParameterTuningAgent, - actions: Sequence[str], - *, - prior_tuning_reference: AgentReportReference | None = None, - persisted_reference: AgentReportReference | None = None, - ) -> tuple[WorkflowStageAttempt, ParameterTuningReport]: - if experimental_handoff.cellSelection is None: - raise ValueError("Parameter tuning handoff lacks an exact cell selection") - if report.cellSelection != experimental_handoff.cellSelection: - raise ValueError("Parameter tuning report uses a different cell selection") - invocation_artifacts: dict[str, ArtifactReferenceModel] = {} - for value in preprocessed: - if value.normalized is not None: - invocation_artifacts[f"{value.assay}_normalized"] = value.normalized - invocation_artifacts["cellSelection"] = experimental_handoff.cellSelection - for integration_evaluation in integration_evaluations: - if integration_evaluation.graphArtifact is not None: - invocation_artifacts[ - f"{integration_evaluation.integrationId}_graph" - ] = ArtifactReferenceModel.model_validate( - integration_evaluation.graphArtifact.model_dump() - ) - if integration_evaluation.clusterArtifact is not None: - invocation_artifacts[ - f"{integration_evaluation.integrationId}_clusters" - ] = ArtifactReferenceModel.model_validate( - integration_evaluation.clusterArtifact.model_dump() - ) - stage_artifacts = dict(invocation_artifacts) - for assay, assay_report in report.assayReports.items(): - for name, artifact in assay_report.selectedArtifacts.items(): - stage_artifacts[f"{assay}_{name}"] = ( - ArtifactReferenceModel.model_validate(artifact.model_dump()) - ) - if report.finalClusterArtifact is not None: - stage_artifacts["final_clusters"] = ArtifactReferenceModel.model_validate( - report.finalClusterArtifact.model_dump() - ) - if report.recommendedIntegrationId is not None: - selected_integration = next( - value - for value in integration_evaluations - if value.integrationId == report.recommendedIntegrationId - ) - if selected_integration.graphArtifact is not None: - stage_artifacts["final_graph"] = ArtifactReferenceModel.model_validate( - selected_integration.graphArtifact.model_dump() - ) - elif report.graphAssay is not None: - assay_reports = report.assayReports or {report.fromAssay: report} - graph_artifact = assay_reports[report.graphAssay].selectedArtifacts[ - "connectivityMap" - ] - stage_artifacts["final_graph"] = ArtifactReferenceModel.model_validate( - graph_artifact.model_dump() - ) - checkpoint_ids = { - f"{journal._stage_execution_id(started)}_integration_{method}" - for method in {evaluation.method for evaluation in integration_evaluations} - } - checkpoint_references = sorted( - ( - reference - for reference in list_agent_reports( - store, - started.workflowRunId, - agent_name="parameter_tuning", - ) - if reference.agentRunId in checkpoint_ids - ), - key=lambda value: value.agentRunId, - ) - if persisted_reference is None: + with candidate_metric_cache(): + report, evidence = runner.run() saved_report, reference = journal._save_stage_report( store, started, report, - invocation=AgentInvocation( - agentName="parameter_tuning", - parentReports=[ - journal._report_link(enrichment_reference), - journal._report_link(experimental_reference), - *( - [journal._report_link(prior_tuning_reference)] - if prior_tuning_reference is not None - else [] - ), - *[ - journal._report_link(value) - for value in checkpoint_references - ], - ], - inputs={ - "assays": dict(candidate_payload), - "primaryAssay": plan.primaryAssay, - "markerAssay": plan.markerAssay, - "pairedAssays": list(paired), - "cellSelection": ( - experimental_handoff.cellSelection.model_dump(mode="json") - if experimental_handoff.cellSelection is not None - else None - ), - "maxCandidateBranches": request_record.config.maxCandidateBranches, - }, - artifacts=stage_artifacts, - runConfig=agent.config, - experimentalTuningHandoff=experimental_handoff, - ), expected_type=ParameterTuningReport, + attempt_owned=True, ) report = cast(ParameterTuningReport, saved_report) - else: - reference = persisted_reference - stage_report_references = [reference, *checkpoint_references] - operations: list[dict[str, Any]] = [] - for assay, assay_report in report.assayReports.items(): - for candidate_evaluation in assay_report.evaluations: - operations.append( - { - "operation": "execute_parameter_candidate", - "assay": assay, - "candidate": candidate_evaluation.parameters.model_dump( - mode="json" - ), - "phase": candidate_evaluation.phase, - "cellSelection": ( - experimental_handoff.cellSelection.model_dump(mode="json") - ), - "harmonyBatchColumns": list( - candidate_evaluation.harmonyBatchColumns - ), - "identityFeatureLimit": ( - request_record.config.maxIdentityFeatures - ), - "status": candidate_evaluation.status, - "artifacts": { - name: value.model_dump(mode="json") - for name, value in candidate_evaluation.artifacts.items() - }, - } - ) - seen_integrated_graphs: set[tuple[str, str]] = set() - for integration_evaluation in integration_evaluations: - integration_graph = integration_evaluation.graphArtifact - graph_id = ( - integration_graph.artifactId if integration_graph is not None else "" - ) - graph_key = (integration_evaluation.method, graph_id) - if graph_id and graph_key not in seen_integrated_graphs: - assert integration_graph is not None - seen_integrated_graphs.add(graph_key) - source_key = ( - "connectivityMap" - if integration_evaluation.method == "snn" - else "neighbors" - ) - source_artifacts = [] - for assay in integration_evaluation.assays: - assay_report = report.assayReports[assay] - selected = next( - value - for value in assay_report.evaluations - if value.candidateId == assay_report.recommendedCandidateId - ) - source_artifacts.append( - selected.artifacts[source_key].model_dump(mode="json") - ) - operations.append( - { - "operation": "integrate_assays", - "method": integration_evaluation.method, - "sources": source_artifacts, - "invalidateCache": True, - "l2Normalize": True, - "artifact": integration_graph.model_dump(mode="json"), - } + artifacts = { + f"{candidate.candidateId}:{name}": ArtifactReferenceModel.model_validate( + value.model_dump() ) - if integration_graph is None: - continue - operations.append( - { - "operation": "run_leiden_clustering", - "integrationId": integration_evaluation.integrationId, - "status": integration_evaluation.status, - "resolution": integration_evaluation.resolution, - "graph": integration_graph.model_dump(mode="json"), - "cellSelection": ( - integration_evaluation.cellSelection.model_dump(mode="json") - if integration_evaluation.cellSelection is not None - else None - ), - "backend": "igraph", - "symmetricGraph": False, - "graphUpperOnly": False, - "randomSeed": 4444, - "invalidateCache": False, - "artifact": ( - integration_evaluation.clusterArtifact.model_dump(mode="json") - if integration_evaluation.clusterArtifact is not None - else None - ), - } - ) - if report.status == "needsInput": - needs_input = report.needsInput - assert needs_input is not None - outcome = journal._complete_attempt( - started, - status="needsInput", - report_references=stage_report_references, - artifacts=stage_artifacts, - outputs={ - "candidateCount": report.totalCandidates, - "operations": operations, - }, - actions=actions, - needs_input=WorkflowNeedsInput( + for candidate in report.evaluations + for name, value in candidate.artifacts.items() + } + pending = None + if report.status == "needsInput": + assert report.needsInput is not None + pending = WorkflowNeedsInput( questions=[ WorkflowQuestion( - questionId=( - "finalGraphOptionId" - if report.finalSelection is not None - and report.finalSelection.status == "needsInput" - else "parameter_tuning" - ), - question=needs_input.question, - options=list(needs_input.options), - evidenceIds=list(needs_input.evidenceIds), + questionId="parameter_tuning", + question=report.needsInput.question, + options=report.needsInput.options, + evidenceIds=report.needsInput.evidenceIds, ) ] - ), - notes=report.limitations, - ) - elif report.status == "failed": - outcome = journal._complete_attempt( - started, - status="failed", - report_references=stage_report_references, - artifacts=stage_artifacts, - outputs={ - "candidateCount": report.totalCandidates, - "operations": operations, - }, - actions=actions, - error="; ".join(report.limitations) or "Parameter Tuning failed", - ) - else: + ) outcome = journal._complete_attempt( started, - status="done", - report_references=stage_report_references, - artifacts=stage_artifacts, + status=report.status, + report_references=[reference], + artifacts=artifacts, outputs={ + "tuningEvidence": evidence, "candidateCount": report.totalCandidates, - "recommendedByAssay": report.recommendedByAssay, - "recommendedIntegrationId": report.recommendedIntegrationId, - "operations": operations, }, - actions=actions, - notes=[*report.tradeoffs, *report.limitations], - ) - journal._save_outcome(store.zw, prefix, outcome) - logger.info( - f"Workflow {workflow.workflowRunId}: Parameter Tuning outcome " - f"status={outcome.status!r}, candidates={report.totalCandidates}, " - f"integrations={len(integration_evaluations)}" - ) - if outcome.status == "failed": - journal.finalize_failed(store, workflow, outcome.error or "tuning failed") - return outcome, report - - def initial_parameter_candidates( - self, - workflow_run_id: str, - handoff: PreprocessedAssayHandoff, - *, - count: int, - neighbors_k: int, - ) -> list[ParameterCandidate]: - max_dimensions = min(handoff.nCells, handoff.nFeatures) - 1 - if neighbors_k < 2 or neighbors_k >= handoff.nCells: - raise ValueError( - f"Assay {handoff.assay!r} has no rank-valid graph candidate" - ) - if handoff.reductionMethod == "identity": - if handoff.nFeatures < 2: - raise ValueError( - f"Assay {handoff.assay!r} has no rank-valid graph candidate" - ) - dimensions = handoff.nFeatures - dimension_values = [dimensions] - elif max_dimensions < 2: - raise ValueError( - f"Assay {handoff.assay!r} has no rank-valid graph candidate" - ) - elif handoff.reductionMethod == "lsi": - dimensions = min(50, max_dimensions) - dimension_values = [ - dimensions, - min(30, max_dimensions), - min(70, max_dimensions), - ] - else: - dimensions = min(21, max_dimensions) - dimension_values = [ - dimensions, - min(15, max_dimensions), - min(30, max_dimensions), - ] - unique_dimensions = list( - dict.fromkeys(value for value in dimension_values if value >= 2) - ) - specifications: list[tuple[int, float]] = [(unique_dimensions[0], 1.0)] - specifications.extend((value, 1.0) for value in unique_dimensions[1:]) - for resolution in (0.5, 1.5, 0.75, 1.25, 0.35, 1.75): - if len(specifications) >= count: - break - specifications.append((unique_dimensions[0], resolution)) - token = workflow_run_id[:10] - assay_token = journal._safe_label(handoff.assay).lower() - if len(assay_token) > 32: - digest = hashlib.blake2b( - handoff.assay.encode("utf-8"), digest_size=4 - ).hexdigest() - assay_token = f"{assay_token[:23]}_{digest}" - if handoff.reductionMethod == "identity": - candidates = [ - ParameterCandidate( - candidateId=f"w_{token}_{assay_token}_0", - reductionMethod="identity", - dimensions=handoff.nFeatures, - leidenResolution=1.0, - neighborsK=neighbors_k, - ) - ] - if count > 1 and max_dimensions >= 2: - candidates.append( - ParameterCandidate( - candidateId=f"w_{token}_{assay_token}_1", - reductionMethod="pca", - dimensions=min(21, max_dimensions), - leidenResolution=1.0, - neighborsK=neighbors_k, - ) - ) - for resolution in (0.5, 1.5, 0.75, 1.25, 0.35, 1.75): - if len(candidates) >= count: - break - index = len(candidates) - candidates.append( - ParameterCandidate( - candidateId=f"w_{token}_{assay_token}_{index}", - reductionMethod="identity", - dimensions=handoff.nFeatures, - leidenResolution=resolution, - neighborsK=neighbors_k, - ) - ) - return candidates - return [ - ParameterCandidate( - candidateId=f"w_{token}_{assay_token}_{index}", - reductionMethod=cast(Any, handoff.reductionMethod), - dimensions=dimension, - leidenResolution=resolution, - neighborsK=neighbors_k, - ) - for index, (dimension, resolution) in enumerate(specifications[:count]) - ] - - def load_integration_checkpoint( - self, - store: DataStore, - started: WorkflowStageAttempt, - method: Literal["snn", "wnn"], - ) -> tuple[list[IntegrationCandidateEvaluation], AgentReportReference] | None: - checkpoint_id = f"{journal._stage_execution_id(started)}_integration_{method}" - matches = [ - reference - for reference in list_agent_reports( - store, - started.workflowRunId, - agent_name="parameter_tuning", - ) - if reference.agentRunId == checkpoint_id - ] - if not matches: - return None - if len(matches) != 1: - raise ValueError("An integration checkpoint has multiple reports") - reference = matches[0] - record = load_agent_record(store, reference) - if ( - record.invocation.inputs.get("orchestrationExecutionId") != checkpoint_id - or record.invocation.inputs.get("stageExecutionId") - != journal._stage_execution_id(started) - or record.invocation.inputs.get("method") != method - ): - raise ValueError("Integration checkpoint identity is stale") - for artifact in record.invocation.artifacts.values(): - store.load_artifact(artifact_model_to_ref(artifact)) - report = load_agent_report(store, reference) - if not isinstance(report, ParameterTuningReport): - raise TypeError("Integration checkpoint is not a Parameter Tuning report") - evaluations = list(report.integrationEvaluations) - if not evaluations or any(value.method != method for value in evaluations): - raise ValueError("Integration checkpoint contains the wrong method") - logger.info( - f"Workflow {started.workflowRunId}: recovered {method.upper()} " - f"checkpoint with {len(evaluations)} evaluation(s)" - ) - return evaluations, reference - - def save_integration_checkpoint( - self, - store: DataStore, - started: WorkflowStageAttempt, - report: ParameterTuningReport, - method: Literal["snn", "wnn"], - evaluations: Sequence[IntegrationCandidateEvaluation], - parent_reports: Sequence[AgentReportLink], - ) -> AgentReportReference: - checkpoint_id = f"{journal._stage_execution_id(started)}_integration_{method}" - checkpoint_report = report.model_copy( - update={ - "integrationEvaluations": list(evaluations), - "recommendedIntegrationId": None, - "finalClusterColumn": None, - "finalClusterArtifact": None, - "finalSelection": None, - } - ) - artifacts: dict[str, ArtifactReferenceModel] = {} - cell_selection = next( - ( - evaluation.cellSelection - for evaluation in evaluations - if evaluation.cellSelection is not None - ), - None, - ) - if cell_selection is not None: - artifacts["cellSelection"] = cell_selection - for evaluation in evaluations: - if evaluation.graphArtifact is not None: - artifacts[f"{evaluation.integrationId}_graph"] = ( - ArtifactReferenceModel.model_validate( - evaluation.graphArtifact.model_dump() - ) - ) - if evaluation.clusterArtifact is not None: - artifacts[f"{evaluation.integrationId}_clusters"] = ( - ArtifactReferenceModel.model_validate( - evaluation.clusterArtifact.model_dump() - ) - ) - invocation = AgentInvocation( - agentName="parameter_tuning", - parentReports=list(parent_reports), - inputs={ - "orchestrationExecutionId": checkpoint_id, - "stageExecutionId": journal._stage_execution_id(started), - "method": method, - "cellSelection": ( - cell_selection.model_dump(mode="json") - if cell_selection is not None - else None - ), - }, - artifacts=artifacts, - ) - try: - reference = save_agent_report( - store, - started.workflowRunId, - checkpoint_report, - invocation=invocation, - agent_run_id=checkpoint_id, - ) - logger.info( - f"Workflow {started.workflowRunId}: persisted {method.upper()} " - f"checkpoint with {len(evaluations)} evaluation(s)" - ) - return reference - except FileExistsError: - recovered = self.load_integration_checkpoint(store, started, method) - if recovered is None: - raise - return recovered[1] - - def evaluate_integrations( - self, - store: DataStore, - workflow_run_id: str, - plan: AutomatedPreprocessingPlan, - report: ParameterTuningReport, - experimental_handoff: ExperimentalTuningHandoff, - config: AutomatedWorkflowConfig, - *, - started: WorkflowStageAttempt | None = None, - parent_reports: Sequence[AgentReportLink] = (), - actions: list[str] | None = None, - ) -> list[IntegrationCandidateEvaluation]: - assays = list(plan.pairedAssays) - if len(assays) < 2: - logger.info("Skipping SNN/WNN evaluation: fewer than two paired assays") - return [] - selected_k = { - next( - evaluation.parameters.neighborsK - for evaluation in assay_report.evaluations - if evaluation.candidateId == assay_report.recommendedCandidateId - ) - for assay, assay_report in report.assayReports.items() - if assay in assays - } - if len(selected_k) != 1: - raise ValueError("SNN and WNN require one common selected neighborsK") - primary_report = report.assayReports[plan.primaryAssay] - primary_evaluation = next( - value - for value in primary_report.evaluations - if value.candidateId == primary_report.recommendedCandidateId - ) - center = primary_evaluation.parameters.leidenResolution - count = config.integrationResolutionCandidates - multipliers = [1.0] if count == 1 else np.linspace(0.5, 1.5, count).tolist() - resolutions = list( - dict.fromkeys(max(0.05, round(center * value, 6)) for value in multipliers) - ) - logger.info( - f"Evaluating SNN and WNN across {len(resolutions)} resolution(s) " - f"for {len(assays)} paired assay(s)" - ) - if report.cellSelection is None: - raise ValueError("Parameter tuning report lacks an exact cell selection") - cell_selection = report.cellSelection - native_labels: dict[str, np.ndarray[Any, Any]] = {} - for assay, assay_report in report.assayReports.items(): - if assay not in assays: - continue - selected = next( - value - for value in assay_report.evaluations - if value.candidateId == assay_report.recommendedCandidateId - ) - cluster_model = ArtifactReferenceModel.model_validate( - selected.artifacts["clusters"].model_dump() - ) - cluster_group = store.load_artifact(artifact_model_to_ref(cluster_model)) - cluster_values = cast(Any, cluster_group["values"]) - native_labels[assay] = np.asarray(cluster_values[:]) - token = workflow_run_id[:12] - evaluations: list[IntegrationCandidateEvaluation] = [] - integration_methods: tuple[Literal["snn", "wnn"], ...] = ("snn", "wnn") - for method in integration_methods: - evaluations.extend( - self.evaluate_integration_method( - store, - method, - token, - assays, - resolutions, - native_labels, - report, - experimental_handoff, - config, - cell_selection, - started=started, - parent_reports=parent_reports, - actions=actions, - ) - ) - return evaluations - - def evaluate_integration_method( - self, - store: DataStore, - method: Literal["snn", "wnn"], - token: str, - assays: list[str], - resolutions: Sequence[float], - native_labels: Mapping[str, np.ndarray[Any, Any]], - report: ParameterTuningReport, - experimental_handoff: ExperimentalTuningHandoff, - config: AutomatedWorkflowConfig, - cell_selection: ArtifactReferenceModel, - *, - started: WorkflowStageAttempt | None, - parent_reports: Sequence[AgentReportLink], - actions: list[str] | None, - ) -> list[IntegrationCandidateEvaluation]: - if started is not None: - recovered = self.load_integration_checkpoint(store, started, method) - if recovered is not None: - if actions is not None: - actions.append(f"recover_integration_checkpoint:{method}") - return recovered[0] - logger.info( - f"Evaluating {method.upper()} integration across " - f"{len(resolutions)} resolution(s)" - ) - evaluations: list[IntegrationCandidateEvaluation] = [] - source_key = "connectivityMap" if method == "snn" else "neighbors" - sources = [] - for assay in assays: - assay_report = report.assayReports[assay] - selected = next( - value - for value in assay_report.evaluations - if value.candidateId == assay_report.recommendedCandidateId - ) - source_model = ArtifactReferenceModel.model_validate( - selected.artifacts[source_key].model_dump() - ) - sources.append(artifact_model_to_ref(source_model)) - try: - graph_ref = store.integrate_assays( - sources, - method=method, - invalidate_cache=True, - l2_normalize=True, + needs_input=pending, + actions=[ + "assess_rna_defaults", + "execute_evidence_requested_experiments", + "validate_full_cohort", + ], ) - weights_valid: bool | None = None - if method == "wnn": - graph_group = store.load_artifact(graph_ref) - stored_weights = cast(Any, graph_group["modality_weights"]) - weights = np.asarray(stored_weights[:], dtype=float) - weights_valid = bool( - weights.shape - == (len(next(iter(native_labels.values()))), len(assays)) - and np.all(np.isfinite(weights)) - and np.all(weights >= 0) - and np.allclose(weights.sum(axis=1), 1.0, rtol=1e-5, atol=1e-6) - ) + journal._save_outcome(store.zw, prefix, outcome) + return outcome, report except Exception as exc: - logger.warning( - f"{method.upper()} graph construction failed " - f"({type(exc).__name__}); persisting failed evaluations" - ) - for index, resolution in enumerate(resolutions): - evaluations.append( - IntegrationCandidateEvaluation( - integrationId=f"{method}_{token}_{index}", - method=cast(Any, method), - assays=assays, - status="failed", - cellSelection=cell_selection, - resolution=resolution, - error=f"{type(exc).__name__}: {exc}", - ) - ) - if started is not None: - self.save_integration_checkpoint( - store, - started, - report, - method, - evaluations, - parent_reports, - ) - if actions is not None: - actions.append(f"checkpoint_integration:{method}") - return evaluations - for index, resolution in enumerate(resolutions): - integration_id = f"{method}_{token}_{index}" - warnings: list[str] = [] - evidence_ids = [f"integration:{integration_id}:clusters"] - try: - cluster_ref = store.run_leiden_clustering( - graph_ref, - resolution=resolution, - backend="igraph", - symmetric_graph=False, - graph_upper_only=False, - random_seed=4444, - invalidate_cache=False, - ) - cluster_group = store.load_artifact(cluster_ref) - cluster_values = cast(Any, cluster_group["values"]) - values = np.asarray(cluster_values[:]) - _labels, counts = np.unique(values, return_counts=True) - metrics = IntegrationMetrics( - nClusters=int(len(counts)), - minClusterCells=int(counts.min()), - minClusterFraction=float(counts.min() / len(values)), - modalityWeightsValid=weights_valid, - ) - for assay, native in native_labels.items(): - metrics.adjustedRandByAssay[assay] = float( - adjusted_rand_score(native, values) - ) - metrics.normalizedMutualInformationByAssay[assay] = float( - normalized_mutual_info_score(native, values) - ) - evidence_ids.extend( - [ - f"integration:{integration_id}:ari:{assay}", - f"integration:{integration_id}:nmi:{assay}", - ] - ) - for column in experimental_handoff.preservationColumns: - try: - value = float( - store.metric_graph_connectivity( - column, - graph_ref, - ) - ) - if np.isfinite(value): - metrics.biologicalConnectivity[column] = value - evidence_ids.append( - f"integration:{integration_id}:graphConnectivity:{column}" - ) - except (KeyError, RuntimeError, TypeError, ValueError) as exc: - warnings.append( - f"Graph connectivity for {column!r} unavailable: {exc}" - ) - if method == "wnn": - evidence_ids.append(f"integration:{integration_id}:modalityWeights") - reasons: list[str] = [] - missing_connectivity = sorted( - set(experimental_handoff.preservationColumns) - - set(metrics.biologicalConnectivity) - ) - if missing_connectivity: - reasons.append( - "trusted-label connectivity is unavailable for " - + ", ".join(missing_connectivity) - ) - if metrics.nClusters is None or metrics.nClusters < 2: - reasons.append("fewer than two clusters") - if ( - metrics.minClusterCells is None - or metrics.minClusterCells < config.minClusterCells - ): - reasons.append("smallest cluster is below the configured minimum") - if method == "wnn" and weights_valid is not True: - reasons.append("WNN modality weights are invalid") - evaluations.append( - IntegrationCandidateEvaluation( - integrationId=integration_id, - method=cast(Any, method), - assays=assays, - status="done", - eligible=not reasons, - cellSelection=cell_selection, - resolution=resolution, - graphArtifact=ArtifactRecord.from_ref(graph_ref), - clusterArtifact=ArtifactRecord.from_ref(cluster_ref), - metrics=metrics, - evidenceIds=evidence_ids, - eligibilityReasons=reasons, - warnings=warnings, - ) - ) - except Exception as exc: - evaluations.append( - IntegrationCandidateEvaluation( - integrationId=integration_id, - method=cast(Any, method), - assays=assays, - status="failed", - cellSelection=cell_selection, - resolution=resolution, - graphArtifact=ArtifactRecord.from_ref(graph_ref), - evidenceIds=evidence_ids, - warnings=warnings, - error=f"{type(exc).__name__}: {exc}", - ) - ) - if started is not None: - self.save_integration_checkpoint( + outcome = journal.finish_exception( store, + prefix, + workflow, started, - report, - method, - evaluations, - parent_reports, + exc, + outputs={"tuningEvidence": runner.summary()}, ) - if actions is not None: - actions.append(f"checkpoint_integration:{method}") - eligible_count = sum( - value.status == "done" and value.eligible for value in evaluations - ) - failed_count = sum(value.status == "failed" for value in evaluations) - logger.info( - f"Completed {method.upper()} integration evaluation: " - f"eligible={eligible_count}, failed={failed_count}, " - f"total={len(evaluations)}" - ) - return evaluations + return outcome, ParameterTuningReport.get_blank() diff --git a/scarf/agent/parameter_tuning.py b/scarf/agent/parameter_tuning.py deleted file mode 100644 index 88368f6a..00000000 --- a/scarf/agent/parameter_tuning.py +++ /dev/null @@ -1,3559 +0,0 @@ -"""Bounded parameter tuning over explicit Scarf analysis candidates.""" - -import json -from collections.abc import Mapping, Sequence -from textwrap import dedent -from threading import Lock -from typing import Any, Literal - -import numpy as np - -from ..storage.refs import ArtifactRef -from .config import CONFIG, AgentRunConfig -from .config._deps import AGENT_INSTALL_HINT -from .config.agent_exec import run_agent_sync -from .tools import artifact_reference, core_artifact_reference -from .types import ( - AgentDataModel, - AgentRunInfo, - ArtifactReferenceModel, - ExperimentalTuningHandoff, - StageStatus, - TuningBiologyHandoff, -) -from ..utils.logging import logger - -try: - from pydantic import Field -except ImportError as exc: - raise ImportError(AGENT_INSTALL_HINT) from exc - -try: - from pydantic_ai import RunContext, UnexpectedModelBehavior -except ImportError as exc: - raise ImportError(AGENT_INSTALL_HINT) from exc - - -type CandidateStatus = Literal["done", "failed"] -type CandidatePhase = Literal["initial", "refined"] -type ParameterSearchStatus = Literal["complete", "refine"] -type TuningConfidence = Literal["low", "medium", "high"] -type ReductionMethod = Literal["pca", "lsi", "identity"] -type IntegrationMethod = Literal["snn", "wnn"] - - -class ArtifactRecord(ArtifactReferenceModel): - """JSON-safe identity for one artifact returned by candidate execution.""" - - @classmethod - def from_ref(cls, ref: Any) -> "ArtifactRecord": - return cls( - scope=getattr(ref, "scope", "assay"), - kind=str(getattr(ref, "kind", "")), - artifactId=str(getattr(ref, "artifact_id", ref)), - assay=getattr(ref, "assay", None), - ) - - @classmethod - def get_blank(cls) -> "ArtifactRecord": - return cls() - - @classmethod - def get_example(cls) -> "ArtifactRecord": - return cls( - scope="assay", - kind="connectivity_map", - artifactId="a" * 64, - assay="RNA", - ) - - -class ParameterCandidate(AgentDataModel): - """One exact, caller-authorized parameter candidate.""" - - candidateId: str = Field( - default="", - description="Exact candidate id supplied to the evaluation tool", - ) - reductionMethod: ReductionMethod = "pca" - dimensions: int = Field(default=21, ge=2) - leidenResolution: float = Field(default=1.0, gt=0) - neighborsK: int = Field(default=11, ge=2) - useHarmony: bool = False - - @classmethod - def get_blank(cls) -> "ParameterCandidate": - return cls() - - @classmethod - def get_example(cls) -> "ParameterCandidate": - return cls( - candidateId="baseline", - reductionMethod="pca", - dimensions=21, - leidenResolution=1.0, - neighborsK=11, - useHarmony=False, - ) - - -class ParameterMetrics(AgentDataModel): - """Bounded quality metrics for one candidate branch.""" - - nClusters: int | None = None - minClusterCells: int | None = None - minClusterFraction: float | None = None - graphSilhouetteMedian: float | None = None - pcaSilhouette: float | None = None - macroF1: float | None = None - weightedF1: float | None = None - batchMixing: dict[str, float] = Field(default_factory=dict) - biologicalPreservation: dict[str, dict[str, float]] = Field(default_factory=dict) - - @classmethod - def get_blank(cls) -> "ParameterMetrics": - return cls() - - @classmethod - def get_example(cls) -> "ParameterMetrics": - return cls( - nClusters=8, - minClusterCells=42, - minClusterFraction=0.021, - graphSilhouetteMedian=0.41, - pcaSilhouette=0.36, - macroF1=0.82, - weightedF1=0.86, - batchMixing={"batch": 0.73}, - biologicalPreservation={ - "cell_type": {"clisi": 0.88, "graphConnectivity": 0.91} - }, - ) - - -class ParameterCandidateEvaluation(AgentDataModel): - """Execution record returned to the model for one candidate.""" - - candidateId: str = "" - phase: CandidatePhase = "initial" - harmonyBatchColumns: list[str] = Field(default_factory=list) - status: CandidateStatus = "failed" - eligible: bool = False - parameters: ParameterCandidate = Field(default_factory=ParameterCandidate.get_blank) - artifacts: dict[str, ArtifactRecord] = Field(default_factory=dict) - cellSelection: ArtifactReferenceModel | None = None - clusterColumn: str | None = None - clusterLabel: str | None = None - effectiveDimensions: int | None = None - metrics: ParameterMetrics = Field(default_factory=ParameterMetrics.get_blank) - evidenceIds: list[str] = Field(default_factory=list) - eligibilityReasons: list[str] = Field(default_factory=list) - warnings: list[str] = Field(default_factory=list) - error: str | None = None - - @classmethod - def get_blank(cls) -> "ParameterCandidateEvaluation": - return cls() - - @classmethod - def get_example(cls) -> "ParameterCandidateEvaluation": - candidate = ParameterCandidate.get_example() - return cls( - candidateId=candidate.candidateId, - status="done", - eligible=True, - parameters=candidate, - artifacts={ - "connectivityMap": ArtifactRecord.get_example(), - "clusters": ArtifactRecord( - assay="RNA", - kind="cluster_labels", - artifactId="b" * 64, - ), - }, - cellSelection=ArtifactReferenceModel( - scope="datastore", - assay=None, - kind="cell_selection", - artifactId="c" * 64, - ), - clusterColumn="RNA_agent_tuning_baseline", - clusterLabel="agent_tuning_baseline", - effectiveDimensions=21, - metrics=ParameterMetrics.get_example(), - evidenceIds=["candidate:baseline:clusters"], - ) - - -class IntegrationMetrics(AgentDataModel): - """Metrics that are valid for an integrated graph comparison.""" - - nClusters: int | None = None - minClusterCells: int | None = None - minClusterFraction: float | None = None - adjustedRandByAssay: dict[str, float] = Field(default_factory=dict) - normalizedMutualInformationByAssay: dict[str, float] = Field(default_factory=dict) - biologicalConnectivity: dict[str, float] = Field(default_factory=dict) - modalityWeightsValid: bool | None = None - - @classmethod - def get_blank(cls) -> "IntegrationMetrics": - return cls() - - @classmethod - def get_example(cls) -> "IntegrationMetrics": - return cls( - nClusters=8, - minClusterCells=37, - minClusterFraction=0.0185, - adjustedRandByAssay={"RNA": 0.71, "ADT": 0.63}, - normalizedMutualInformationByAssay={"RNA": 0.76, "ADT": 0.69}, - modalityWeightsValid=True, - ) - - -class IntegrationCandidateEvaluation(AgentDataModel): - """One executor-produced SNN or WNN graph and cluster evaluation.""" - - integrationId: str = "" - method: IntegrationMethod = "wnn" - assays: list[str] = Field(default_factory=list) - status: CandidateStatus = "failed" - eligible: bool = False - cellSelection: ArtifactReferenceModel | None = None - resolution: float = Field(default=1.0, gt=0) - graphArtifact: ArtifactRecord | None = None - clusterArtifact: ArtifactRecord | None = None - clusterColumn: str | None = None - metrics: IntegrationMetrics = Field(default_factory=IntegrationMetrics.get_blank) - evidenceIds: list[str] = Field(default_factory=list) - eligibilityReasons: list[str] = Field(default_factory=list) - warnings: list[str] = Field(default_factory=list) - error: str | None = None - - @classmethod - def get_blank(cls) -> "IntegrationCandidateEvaluation": - return cls() - - @classmethod - def get_example(cls) -> "IntegrationCandidateEvaluation": - return cls( - integrationId="wnn_resolution_1", - method="wnn", - assays=["RNA", "ADT"], - status="done", - eligible=True, - cellSelection=ArtifactReferenceModel( - scope="datastore", - assay=None, - kind="cell_selection", - artifactId="c" * 64, - ), - graphArtifact=ArtifactRecord( - scope="datastore", - kind="integrated_graph", - artifactId="2" * 64, - ), - clusterArtifact=ArtifactRecord( - scope="datastore", - kind="cluster_labels", - artifactId="3" * 64, - ), - clusterColumn="agent_wnn_cluster", - metrics=IntegrationMetrics.get_example(), - evidenceIds=["integration:wnn_resolution_1:clusters"], - ) - - -class FinalGraphComparison(AgentDataModel): - """Evidence-backed comparison against one eligible final graph option.""" - - optionId: str = "" - summary: str = "" - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "FinalGraphComparison": - return cls() - - @classmethod - def get_example(cls) -> "FinalGraphComparison": - return cls( - optionId="native:ADT:baseline", - summary="The RNA-native option better preserves the requested labels.", - evidenceIds=[ - "native:RNA:candidate:baseline:clusters", - "native:ADT:candidate:baseline:clusters", - ], - ) - - -class FinalGraphNeedsInput(AgentDataModel): - """Concrete input needed before a final graph can be selected.""" - - question: str = "" - options: list[str] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "FinalGraphNeedsInput": - return cls() - - @classmethod - def get_example(cls) -> "FinalGraphNeedsInput": - return cls( - question="Which biological signal must the final graph preserve?", - options=["cell_type", "condition"], - ) - - -class FinalGraphSelection(AgentDataModel): - """Grounded choice among selected native, SNN, and WNN graph options.""" - - status: StageStatus = "needsInput" - selectedOptionId: str | None = None - graphMethod: Literal["native", "snn", "wnn"] | None = None - nativeAssay: str | None = None - nativeCandidateId: str | None = None - integrationId: str | None = None - markerAssay: str = "" - confidence: TuningConfidence = "low" - rationale: str = "" - evidenceIds: list[str] = Field(default_factory=list) - comparisons: list[FinalGraphComparison] = Field(default_factory=list) - tradeoffs: list[str] = Field(default_factory=list) - limitations: list[str] = Field(default_factory=list) - needsInput: FinalGraphNeedsInput | None = None - runInfo: AgentRunInfo = Field(default_factory=AgentRunInfo) - - @classmethod - def get_blank(cls) -> "FinalGraphSelection": - return cls() - - @classmethod - def get_example(cls) -> "FinalGraphSelection": - return cls( - status="done", - selectedOptionId="native:RNA:baseline", - graphMethod="native", - nativeAssay="RNA", - nativeCandidateId="baseline", - markerAssay="RNA", - confidence="medium", - rationale="The selected native graph has the strongest supported balance.", - evidenceIds=["native:RNA:candidate:baseline:clusters"], - runInfo=AgentRunInfo.get_example(), - ) - - -class CandidateComparison(AgentDataModel): - """Evidence-backed comparison against one executed non-selected candidate.""" - - candidateId: str = "" - summary: str = "" - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "CandidateComparison": - return cls() - - @classmethod - def get_example(cls) -> "CandidateComparison": - return cls( - candidateId="pca_15", - summary="The selected baseline retains larger minimum clusters.", - evidenceIds=[ - "candidate:baseline:clusters", - "candidate:pca_15:clusters", - ], - ) - - -class ParameterSearchPlan(AgentDataModel): - """Validated proposal for one bounded refinement pass.""" - - status: ParameterSearchStatus = Field( - default="complete", - description=( - "Summary derived from candidates: refine when candidates is non-empty " - "and complete when it is empty" - ), - ) - candidates: list[ParameterCandidate] = Field( - default_factory=list, - description=( - "Bounded unexecuted refinement candidates, or an empty list when the " - "initial screen is complete" - ), - ) - basedOnCandidateIds: list[str] = Field(default_factory=list) - harmonyBatchColumns: list[str] = Field(default_factory=list) - objectives: list[str] = Field(default_factory=list) - rationale: str = "" - evidenceIds: list[str] = Field(default_factory=list) - stoppingCriteria: list[str] = Field(default_factory=list) - runInfo: AgentRunInfo = Field(default_factory=AgentRunInfo) - - @classmethod - def get_blank(cls) -> "ParameterSearchPlan": - return cls() - - @classmethod - def get_example(cls) -> "ParameterSearchPlan": - return cls( - status="refine", - candidates=[ - ParameterCandidate( - candidateId="refined_pca_18", - dimensions=18, - leidenResolution=1.0, - neighborsK=11, - useHarmony=False, - ) - ], - basedOnCandidateIds=["baseline", "pca_15"], - harmonyBatchColumns=[], - objectives=["Resolve the dimension tradeoff."], - rationale="The initial screen brackets a narrower dimension range.", - evidenceIds=[ - "candidate:baseline:clusters", - "candidate:pca_15:clusters", - ], - stoppingCriteria=["Run the proposed candidate once."], - runInfo=AgentRunInfo.get_example(), - ) - - -class ParameterTuningBatchSearchPlan(AgentDataModel): - """One bounded refinement plan for every assay in a batched screen.""" - - assayPlans: dict[str, ParameterSearchPlan] = Field(default_factory=dict) - runInfo: AgentRunInfo = Field(default_factory=AgentRunInfo) - - @classmethod - def get_blank(cls) -> "ParameterTuningBatchSearchPlan": - return cls() - - @classmethod - def get_example(cls) -> "ParameterTuningBatchSearchPlan": - return cls(assayPlans={"RNA": ParameterSearchPlan.get_example()}) - - -class ParameterTuningNeedsInput(AgentDataModel): - """User input required before tuning can produce a recommendation.""" - - question: str = "" - options: list[str] = Field(default_factory=list) - evidenceIds: list[str] = Field(default_factory=list) - - @classmethod - def get_blank(cls) -> "ParameterTuningNeedsInput": - return cls() - - @classmethod - def get_example(cls) -> "ParameterTuningNeedsInput": - return cls( - question="Which trusted biological label should be preserved?", - options=["cell_type", "none"], - evidenceIds=["candidate:baseline:batchMixing:batch"], - ) - - -class ParameterTuningReport(AgentDataModel): - """Grounded recommendation over candidate branches actually executed.""" - - status: StageStatus = "failed" - fromAssay: str = "" - cellSelection: ArtifactReferenceModel | None = None - evaluations: list[ParameterCandidateEvaluation] = Field(default_factory=list) - recommendedCandidateId: str | None = None - selectedArtifacts: dict[str, ArtifactRecord] = Field(default_factory=dict) - confidence: TuningConfidence = "low" - rationale: str = "" - evidenceIds: list[str] = Field(default_factory=list) - comparisons: list[CandidateComparison] = Field(default_factory=list) - tradeoffs: list[str] = Field(default_factory=list) - limitations: list[str] = Field(default_factory=list) - stopReason: str = "" - needsInput: ParameterTuningNeedsInput | None = None - searchPlan: ParameterSearchPlan | None = None - assayReports: dict[str, "ParameterTuningReport"] = Field(default_factory=dict) - recommendedByAssay: dict[str, str] = Field(default_factory=dict) - totalCandidates: int = 0 - integrationEvaluations: list[IntegrationCandidateEvaluation] = Field( - default_factory=list - ) - recommendedIntegrationId: str | None = None - finalClusterColumn: str | None = None - finalClusterArtifact: ArtifactRecord | None = None - graphAssay: str | None = None - markerAssay: str | None = None - finalSelection: FinalGraphSelection | None = None - runInfo: AgentRunInfo = Field(default_factory=AgentRunInfo) - - @classmethod - def get_blank(cls) -> "ParameterTuningReport": - return cls() - - @classmethod - def get_example(cls) -> "ParameterTuningReport": - evaluation = ParameterCandidateEvaluation.get_example() - return cls( - status="done", - fromAssay="RNA", - cellSelection=evaluation.cellSelection, - evaluations=[evaluation], - recommendedCandidateId=evaluation.candidateId, - selectedArtifacts=dict(evaluation.artifacts), - confidence="medium", - rationale="The baseline balances separation and cluster size.", - evidenceIds=["candidate:baseline:clusters"], - tradeoffs=["Higher resolutions produced smaller clusters."], - limitations=["No trusted biological preservation label was supplied."], - stopReason="All authorized candidates were evaluated.", - recommendedByAssay={"RNA": evaluation.candidateId}, - totalCandidates=1, - graphAssay="RNA", - markerAssay="RNA", - finalSelection=FinalGraphSelection.get_example(), - runInfo=AgentRunInfo.get_example(), - ) - - def to_biological_handoff( - self, - *, - marker_assay: str | None = None, - ) -> TuningBiologyHandoff: - """Return the exact selected clustering branch for interpretation.""" - if self.status != "done": - raise ValueError( - "Parameter Tuning must be done before creating a biology handoff" - ) - if self.finalClusterArtifact is not None: - if self.cellSelection is None: - raise ValueError("Final branch lacks an exact cell selection") - resolved_marker_assay = marker_assay or self.markerAssay - if not resolved_marker_assay: - raise ValueError( - "A marker assay is required for an integrated biology handoff" - ) - if self.finalClusterArtifact.scope == "datastore": - if self.finalClusterArtifact.assay is not None: - raise ValueError( - "A datastore-scoped cluster artifact must not name an assay" - ) - elif ( - self.graphAssay is not None - and self.finalClusterArtifact.assay != self.graphAssay - ): - raise ValueError("Final cluster artifact does not match graphAssay") - integration = next( - ( - item - for item in self.integrationEvaluations - if item.integrationId == self.recommendedIntegrationId - ), - None, - ) - if self.finalSelection is not None: - evidence_ids = self.finalSelection.evidenceIds - elif integration is not None: - evidence_ids = integration.evidenceIds - else: - prefix = f"candidate:{self.recommendedCandidateId}:" - evidence_ids = [ - evidence_id - for evidence_id in self.evidenceIds - if evidence_id.startswith(prefix) - ] - return TuningBiologyHandoff( - cellSelection=self.cellSelection, - fromAssay=self.fromAssay, - graphAssay=self.graphAssay, - markerAssay=resolved_marker_assay, - recommendedCandidateId=( - self.recommendedIntegrationId - or ( - self.finalSelection.nativeCandidateId - if self.finalSelection is not None - else None - ) - or self.recommendedCandidateId - or "final" - ), - clusterArtifact=ArtifactReferenceModel.model_validate( - self.finalClusterArtifact.model_dump() - ), - evidenceIds=sorted(evidence_ids), - ) - if self.recommendedCandidateId is None: - raise ValueError( - "Parameter Tuning must recommend a candidate before creating a " - "biology handoff" - ) - selected = next( - ( - item - for item in self.evaluations - if item.candidateId == self.recommendedCandidateId - ), - None, - ) - if selected is None or selected.status != "done" or not selected.eligible: - raise ValueError("Recommended candidate is not an eligible execution") - cluster_artifact = selected.artifacts.get("clusters") - if cluster_artifact is None or selected.cellSelection is None: - raise ValueError("Recommended candidate lacks an exact cluster artifact") - if not self.fromAssay or cluster_artifact.assay != self.fromAssay: - raise ValueError("Recommended cluster artifact does not match the assay") - prefix = f"candidate:{selected.candidateId}:" - return TuningBiologyHandoff( - cellSelection=selected.cellSelection, - fromAssay=self.fromAssay, - graphAssay=self.fromAssay, - markerAssay=marker_assay or self.markerAssay or self.fromAssay, - recommendedCandidateId=selected.candidateId, - clusterArtifact=ArtifactReferenceModel.model_validate( - cluster_artifact.model_dump() - ), - evidenceIds=sorted( - evidence_id - for evidence_id in self.evidenceIds - if evidence_id.startswith(prefix) - ), - ) - - -class ParameterTuningDependencies(AgentDataModel): - """Runtime-only state hidden from the model and shared by tuning tools.""" - - store: Any = Field(default=None, exclude=True) - normalized: Any = Field(default=None, exclude=True) - cellSelection: Any = Field(default=None, exclude=True) - normalizedShape: tuple[int, int] | None = None - fromAssay: str = "" - candidates: dict[str, ParameterCandidate] = Field(default_factory=dict) - candidatePhases: dict[str, CandidatePhase] = Field(default_factory=dict) - batchColumns: tuple[str, ...] = () - preservationColumns: tuple[str, ...] = () - harmonyAuthorized: bool = False - maxCandidates: int = 5 - minClusterCells: int = 20 - identityFeatureLimit: int = 64 - evaluations: dict[str, ParameterCandidateEvaluation] = Field(default_factory=dict) - executionOrder: list[str] = Field(default_factory=list) - executionLock: Any = Field(default_factory=Lock, exclude=True, repr=False) - - @classmethod - def get_blank(cls) -> "ParameterTuningDependencies": - return cls() - - @classmethod - def get_example(cls) -> "ParameterTuningDependencies": - candidate = ParameterCandidate.get_example() - return cls( - fromAssay="RNA", - normalizedShape=(1000, 2000), - candidates={candidate.candidateId: candidate}, - batchColumns=("batch",), - preservationColumns=("cell_type",), - ) - - -class ParameterTuningAssayInput(AgentDataModel): - """One assay branch supplied to batched parameter tuning.""" - - normalized: Any = Field(default=None, exclude=True) - candidates: list[ParameterCandidate] = Field(default_factory=list) - batchColumns: list[str] = Field(default_factory=list) - preservationColumns: list[str] = Field(default_factory=list) - experimentalHandoff: ExperimentalTuningHandoff | None = None - maxCandidates: int = Field(default=5, ge=1) - maxRefinedCandidates: int = Field(default=0, ge=0) - allowHarmonyRefinement: bool = True - minClusterCells: int = Field(default=20, ge=1) - identityFeatureLimit: int = Field(default=64, ge=2) - - @classmethod - def get_blank(cls) -> "ParameterTuningAssayInput": - return cls() - - @classmethod - def get_example(cls) -> "ParameterTuningAssayInput": - return cls( - normalized=ArtifactRecord( - assay="RNA", - kind="normalized", - artifactId="4" * 64, - ), - candidates=get_default_parameter_candidates(), - experimentalHandoff=ExperimentalTuningHandoff(batchAction="skip"), - ) - - -def get_default_parameter_candidates() -> list[ParameterCandidate]: - """Return a small one-factor candidate set around Scarf defaults.""" - - return [ - ParameterCandidate( - candidateId="baseline", - dimensions=21, - leidenResolution=1.0, - ), - ParameterCandidate( - candidateId="pca_15", - dimensions=15, - leidenResolution=1.0, - ), - ParameterCandidate( - candidateId="pca_30", - dimensions=30, - leidenResolution=1.0, - ), - ParameterCandidate( - candidateId="leiden_0_5", - dimensions=21, - leidenResolution=0.5, - ), - ParameterCandidate( - candidateId="leiden_1_5", - dimensions=21, - leidenResolution=1.5, - ), - ] - - -def build_initial_parameter_candidates( - candidates: Sequence[ParameterCandidate], - *, - pair_harmony: bool, -) -> list[ParameterCandidate]: - """Build deterministic initial branches from caller-authorized parameters.""" - - initial: list[ParameterCandidate] = [] - for candidate in candidates: - if pair_harmony and candidate.useHarmony: - raise ValueError( - "Initial seed candidates must not set useHarmony when the " - "experimental handoff controls Harmony pairing" - ) - initial.append(candidate) - if pair_harmony: - payload = candidate.model_dump() - payload.update( - { - "candidateId": f"{candidate.candidateId}_harmony", - "useHarmony": True, - } - ) - initial.append(ParameterCandidate.model_validate(payload)) - return initial - - -def parameter_search_system_prompt() -> str: - """Build the stable prompt for the bounded refinement-planning call.""" - - return dedent( - """ - You are planning one bounded refinement pass for Scarf parameter tuning. - The initial candidate screen has already finished. Do not request tools or - claim that additional candidates ran. - - Return exactly one of these two plan shapes: - 1. status=complete with candidates=[] when the initial screen is sufficient. - 2. status=refine with one or more candidates when an untested candidate - inside the initial numeric search envelope can resolve a specific - evidence-backed uncertainty. - Never return status=complete with candidates. A Harmony candidate always - uses the exact authorized batch columns supplied in the prompt. You may - choose between no correction and that approved Harmony configuration, but - you must not propose or modify batch columns. - - Cite only evidenceIds from the initial evaluations. Identify the successful - initial candidates that motivate refinement, state focused objectives, and - provide concrete stopping criteria. Do not invent metrics, artifacts, or - candidate ids. - """ - ).strip() - - -def parameter_search_prompt( - *, - from_assay: str, - cell_selection: ArtifactReferenceModel, - evaluations: Sequence[ParameterCandidateEvaluation], - batch_columns: Sequence[str], - preservation_columns: Sequence[str], - harmony_authorized: bool, - max_refined_candidates: int, -) -> str: - """Build the planning prompt from deterministic initial evaluations.""" - - evaluation_payload = [evaluation.model_dump() for evaluation in evaluations] - correction_modes = ["none", "harmony"] if harmony_authorized else ["none"] - return ( - dedent( - """ - Inspect the completed initial screen for assay {from_assay} and exact - cell-selection artifact {cell_selection}. - - Initial evaluations: - {evaluation_payload} - - Authorized correction modes: {correction_modes} - Exact Harmony batch columns: {batch_columns} - Trusted biological preservation columns: {preservation_columns} - Maximum refined candidates: {max_refined_candidates} - - Return one ParameterSearchPlan. Refinement is optional and is limited to - one deterministic follow-up pass. - """ - ) - .strip() - .format( - from_assay=from_assay, - cell_selection=cell_selection.artifactId, - evaluation_payload=json.dumps( - evaluation_payload, - indent=2, - sort_keys=True, - ), - correction_modes=json.dumps(correction_modes), - batch_columns=json.dumps(list(batch_columns)), - preservation_columns=json.dumps(list(preservation_columns)), - max_refined_candidates=max_refined_candidates, - ) - ) - - -def parameter_tuning_system_prompt(min_cluster_cells: int) -> str: - """Build the stable prompt for final candidate selection.""" - - return ( - dedent( - """ - You are Scarf's parameter tuning selection agent. Every candidate in the - prompt has already finished deterministic execution. Do not request tools - or claim that another candidate ran. - - Recommend only a candidate whose evaluation has status=done and - eligible=true. A candidate is ineligible when it creates fewer than two - clusters or a cluster with fewer than {min_cluster_cells} cells. Do not - invent artifact ids, metrics, candidate ids, or evidence ids. Cite only - evidenceIds recorded in the completed evaluations. - - Balance cluster separation, cluster sizes, batch mixing, and biological - preservation. High batch mixing alone can indicate overcorrection, so do - not collapse the metrics into an invented score. UMAP appearance is not - evidence for parameter quality. When multiple candidates complete, - return one comparison for every non-selected successful candidate. Each - comparison must cite evidence from both the selected candidate and that - comparator. Return a concise structured report. - """ - ) - .strip() - .format(min_cluster_cells=min_cluster_cells) - ) - - -def parameter_tuning_prompt( - *, - from_assay: str, - cell_selection: ArtifactReferenceModel, - evaluations: Sequence[ParameterCandidateEvaluation], - batch_columns: Sequence[str], - preservation_columns: Sequence[str], - search_plan: ParameterSearchPlan, -) -> str: - """Build the final selection prompt from completed evaluations.""" - - evaluation_payload = [evaluation.model_dump() for evaluation in evaluations] - return ( - dedent( - """ - Select a completed candidate for assay {from_assay} and exact - cell-selection artifact {cell_selection}. - - Completed evaluations: - {evaluation_payload} - - Validated refinement plan: - {search_plan} - - Exact Harmony batch columns: {batch_columns} - Trusted biological preservation columns: {preservation_columns} - - Recommend one eligible candidate or explain why user input is needed. - Compare the recommendation with every other successful candidate. High - batch mixing does not by itself justify correction when biological - preservation declines. - """ - ) - .strip() - .format( - from_assay=from_assay, - cell_selection=cell_selection.artifactId, - evaluation_payload=json.dumps( - evaluation_payload, - indent=2, - sort_keys=True, - ), - search_plan=json.dumps( - search_plan.model_dump(exclude={"runInfo"}), - indent=2, - sort_keys=True, - ), - batch_columns=json.dumps(list(batch_columns)), - preservation_columns=json.dumps(list(preservation_columns)), - ) - ) - - -def parameter_batch_search_prompt( - dependencies: Mapping[str, ParameterTuningDependencies], - max_refined_by_assay: Mapping[str, int], -) -> str: - """Build one refinement prompt for all modality-specific screens.""" - - payload = { - assay: { - "evaluations": [ - deps.evaluations[candidate_id].model_dump() - for candidate_id in deps.executionOrder - ], - "authorizedHarmony": deps.harmonyAuthorized, - "batchColumns": list(deps.batchColumns), - "preservationColumns": list(deps.preservationColumns), - "maxRefinedCandidates": max_refined_by_assay[assay], - } - for assay, deps in dependencies.items() - } - return ( - dedent( - """ - Plan one optional refinement pass for every assay in this completed - multimodal initial screen: - {payload} - - Return exactly one assayPlans entry for every assay. Each entry must - obey the single-assay ParameterSearchPlan rules. Candidate ids need - only be unique within their assay. Do not compare metric fields that - are absent for a modality, and do not request additional tool calls. - """ - ) - .strip() - .format(payload=json.dumps(payload, indent=2, sort_keys=True)) - ) - - -def parameter_batch_search_system_prompt() -> str: - """Build the stable system prompt for batched refinement planning.""" - - return ( - dedent( - """ - {single_assay_rules} - - Return the plans together in one assayPlans mapping. - """ - ) - .strip() - .format(single_assay_rules=parameter_search_system_prompt()) - ) - - -def parameter_batch_selection_system_prompt() -> str: - """Build the stable system prompt for batched native selection.""" - - return ( - dedent( - """ - You are Scarf's batched native parameter selection agent. Every branch - has already executed. Return one aggregate ParameterTuningReport with - exactly one grounded single-assay report in assayReports per assay. - Apply eligibility, evidence, and comparison requirements independently. - Do not invent joint scores, artifacts, candidates, or evidence. UMAP - appearance is not evidence. Leave integration fields empty. - """ - ) - .strip() - .format() - ) - - -def parameter_batch_selection_prompt( - dependencies: Mapping[str, ParameterTuningDependencies], - search_plans: Mapping[str, ParameterSearchPlan], - primary_assay: str, - selection_directions: str = "", -) -> str: - """Build one native-selection prompt for all executed assay screens.""" - - payload = { - assay: { - "evaluations": [ - deps.evaluations[candidate_id].model_dump() - for candidate_id in deps.executionOrder - ], - "searchPlan": search_plans[assay].model_dump(exclude={"runInfo"}), - "minClusterCells": deps.minClusterCells, - "batchColumns": list(deps.batchColumns), - "preservationColumns": list(deps.preservationColumns), - } - for assay, deps in dependencies.items() - } - return ( - dedent( - """ - Select one eligible native candidate independently for every assay in - this completed multimodal screen: - {payload} - - Return a ParameterTuningReport whose assayReports contains exactly one - single-assay report per assay. Apply the normal evidence and comparison - rules independently inside each report. The primary assay is - {primary_assay}. At the aggregate level, summarize cross-assay - limitations without inventing a joint score. Integration has not run, - so leave all integration and final-cluster fields empty. - - Caller selection directions, which cannot override eligibility or - evidence requirements: {selection_directions} - """ - ) - .strip() - .format( - payload=json.dumps(payload, indent=2, sort_keys=True), - primary_assay=primary_assay, - selection_directions=selection_directions or "not provided", - ) - ) - - -def final_graph_selection_system_prompt() -> str: - """Build stable instructions for the final native/SNN/WNN choice.""" - - return ( - dedent( - """ - You are Scarf's final graph selection agent. Native assay candidates - and integrated SNN/WNN candidates have already executed. Select only - an eligible option supplied in the prompt. Do not request tools or - invent graph options, artifacts, metrics, evidence, or a combined - score. Compare cluster viability and biological preservation evidence - that is actually present. ARI and NMI describe agreement, not quality. - WNN modality weights are usable only when modalityWeightsValid=true. - UMAP appearance, native-neighbor LISI on an integrated graph, and - absent metric fields are not evidence. Return one comparison for every - eligible non-selected option, citing evidence from both options. - """ - ) - .strip() - .format() - ) - - -def final_graph_options( - report: ParameterTuningReport, - integration_evaluations: Sequence[IntegrationCandidateEvaluation], -) -> dict[str, dict[str, Any]]: - """Return the exact eligible graph options and option-scoped evidence.""" - - assay_reports = report.assayReports or {report.fromAssay: report} - report_cell_selection = core_artifact_reference(report.cellSelection) - if not isinstance(report_cell_selection, ArtifactRef): - raise ValueError("Parameter tuning report lacks an exact cell selection") - options: dict[str, dict[str, Any]] = {} - for assay, assay_report in assay_reports.items(): - candidate_id = assay_report.recommendedCandidateId - if candidate_id is None: - continue - native_evaluation = next( - ( - item - for item in assay_report.evaluations - if item.candidateId == candidate_id - ), - None, - ) - if ( - native_evaluation is None - or native_evaluation.status != "done" - or not native_evaluation.eligible - or "clusters" not in native_evaluation.artifacts - or "connectivityMap" not in native_evaluation.artifacts - or not native_evaluation.evidenceIds - ): - continue - if ( - core_artifact_reference(native_evaluation.cellSelection) - != report_cell_selection - ): - raise ValueError("Native graph option uses a different cell selection") - option_id = f"native:{assay}:{candidate_id}" - option_evidence = [ - f"native:{assay}:{evidence_id}" - for evidence_id in native_evaluation.evidenceIds - ] - evaluation_payload = native_evaluation.model_dump() - evaluation_payload["evidenceIds"] = option_evidence - options[option_id] = { - "optionId": option_id, - "graphMethod": "native", - "nativeAssay": assay, - "nativeCandidateId": candidate_id, - "evaluation": evaluation_payload, - "evidenceIds": option_evidence, - } - for integration_evaluation in integration_evaluations: - if ( - integration_evaluation.status != "done" - or not integration_evaluation.eligible - ): - continue - if ( - integration_evaluation.clusterArtifact is None - or integration_evaluation.graphArtifact is None - ): - continue - if not integration_evaluation.evidenceIds: - continue - if ( - core_artifact_reference(integration_evaluation.cellSelection) - != report_cell_selection - ): - raise ValueError("Integrated graph option uses a different cell selection") - if not integration_evaluation.integrationId: - raise ValueError("Eligible integration evaluations require integrationId") - if ( - integration_evaluation.graphArtifact.scope != "datastore" - or integration_evaluation.graphArtifact.assay is not None - or integration_evaluation.clusterArtifact.scope != "datastore" - or integration_evaluation.clusterArtifact.assay is not None - or integration_evaluation.graphArtifact.kind != "integrated_graph" - or integration_evaluation.clusterArtifact.kind - not in {"cluster_labels", "cluster_cut"} - ): - raise ValueError( - "Integrated graph and cluster artifacts must be datastore-scoped " - "without an assay" - ) - if ( - integration_evaluation.method == "wnn" - and integration_evaluation.metrics.modalityWeightsValid is not True - ): - continue - option_id = f"integration:{integration_evaluation.integrationId}" - if option_id in options: - raise ValueError( - f"Duplicate integration id {integration_evaluation.integrationId!r}" - ) - options[option_id] = { - "optionId": option_id, - "graphMethod": integration_evaluation.method, - "integrationId": integration_evaluation.integrationId, - "evaluation": integration_evaluation.model_dump(), - "evidenceIds": list(integration_evaluation.evidenceIds), - } - return options - - -def final_graph_selection_prompt( - *, - report: ParameterTuningReport, - integration_evaluations: Sequence[IntegrationCandidateEvaluation], - marker_assay: str, -) -> str: - """Build the selection prompt from executor-grounded final graph options.""" - - options = final_graph_options(report, integration_evaluations) - return ( - dedent( - """ - Select the final graph from these eligible executed options: - {options} - - The fixed marker assay is {marker_assay}. It determines marker - extraction and does not imply ownership of an integrated graph. - Return needsInput only when the supplied evidence cannot resolve a - scientifically material tradeoff. - """ - ) - .strip() - .format( - options=json.dumps(options, indent=2, sort_keys=True), - marker_assay=marker_assay, - ) - ) - - -def normalized_artifact_shape(store: Any, normalized: Any) -> tuple[int, int]: - """Return the exact cell-by-feature shape of a normalized artifact.""" - - group = store.load_artifact(normalized) - if "data" not in group: - raise ValueError("Normalized artifact does not contain a data matrix") - shape = getattr(group["data"], "shape", None) - if not isinstance(shape, tuple | list) or len(shape) != 2: - raise ValueError("Normalized artifact data must be two-dimensional") - n_cells, n_features = map(int, shape) - if n_cells < 2 or n_features < 2: - raise ValueError( - "Parameter tuning requires at least two cells and two selected features" - ) - return n_cells, n_features - - -def validate_parameter_candidate_rank( - candidate: ParameterCandidate, - normalized_shape: tuple[int, int], - *, - identity_feature_limit: int = 64, -) -> int: - """Validate a candidate before any branch operation and return output rank.""" - - n_cells, n_features = normalized_shape - if candidate.neighborsK >= n_cells: - raise ValueError( - f"neighborsK={candidate.neighborsK} requires more than " - f"{candidate.neighborsK} selected cells; observed {n_cells}" - ) - if candidate.reductionMethod == "pca": - if candidate.dimensions + 1 > min(n_cells, n_features): - raise ValueError( - f"PCA dimensions={candidate.dimensions} requires at least " - f"{candidate.dimensions + 1} cells and selected features; " - f"observed shape {normalized_shape}" - ) - return candidate.dimensions - if candidate.reductionMethod == "lsi": - required_rank = candidate.dimensions + 1 - if required_rank > min(n_cells, n_features): - raise ValueError( - "LSI dimensions, including the skipped component, exceed the " - f"normalized matrix rank for shape {normalized_shape}" - ) - return candidate.dimensions - if n_features > identity_feature_limit: - raise ValueError( - f"Identity reduction supports at most {identity_feature_limit} selected " - f"features; observed {n_features}" - ) - if candidate.dimensions != n_features: - raise ValueError( - "Identity reduction dimensions must equal the exact normalized feature " - f"count {n_features}; received {candidate.dimensions}" - ) - return n_features - - -def run_candidate_reduction( - store: Any, - *, - normalized: Any, - candidate: ParameterCandidate, - normalized_shape: tuple[int, int], - identity_feature_limit: int = 64, -) -> tuple[Any, str, int]: - """Run one validated modality-aware reduction with public Scarf methods.""" - - effective_dimensions = validate_parameter_candidate_rank( - candidate, - normalized_shape, - identity_feature_limit=identity_feature_limit, - ) - if candidate.reductionMethod == "pca": - ref = store.run_pca( - normalized, - dims=candidate.dimensions, - feat_scaling=True, - show_elbow_plot=False, - invalidate_cache=False, - ) - return ref, "pca", effective_dimensions - if candidate.reductionMethod == "lsi": - ref = store.run_lsi( - normalized, - dims=candidate.dimensions, - skip_first=True, - rand_state=CONFIG._PCA_RANDOM_SEED, - invalidate_cache=False, - ) - return ref, "lsi", effective_dimensions - loadings = np.eye(normalized_shape[1], dtype=np.float64) - ref = store.run_custom_reduction( - loadings, - normalized, - invalidate_cache=False, - ) - return ref, "identity", effective_dimensions - - -def _collect_parameter_candidate_metrics( - deps: ParameterTuningDependencies, - *, - candidate: ParameterCandidate, - candidate_id: str, - reduction_ref: Any, - neighbors_ref: Any, - graph_ref: Any, - cluster_ref: Any, - cluster_column: str, - evidence_ids: list[str], - warnings: list[str], -) -> tuple[ParameterMetrics, list[str]]: - store = deps.store - cluster_group = store.load_artifact(cluster_ref) - cluster_data = cluster_group["values"] - cluster_values = np.asarray(cluster_data[:]) - if cluster_values.ndim != 1 or len(cluster_values) == 0: - raise ValueError("Cluster artifact must contain one non-empty label vector") - if np.any(cluster_values < 0): - raise ValueError("Cluster artifact contains invalid negative labels") - _, cluster_counts = np.unique(cluster_values, return_counts=True) - n_clusters = int(len(cluster_counts)) - min_cluster_cells = int(cluster_counts.min()) - min_cluster_fraction = float(min_cluster_cells / len(cluster_values)) - metrics = ParameterMetrics( - nClusters=n_clusters, - minClusterCells=min_cluster_cells, - minClusterFraction=min_cluster_fraction, - ) - evidence_ids.append(f"candidate:{candidate_id}:clusters") - - try: - graph_scores = store.metric_graph_silhouette( - neighbors_ref, - cluster_ref, - random_seed=CONFIG._RANDOM_SEED, - sample_size=11, - ) - if graph_scores is not None: - finite_scores = np.asarray(graph_scores, dtype=float) - finite_scores = finite_scores[np.isfinite(finite_scores)] - if len(finite_scores): - metrics.graphSilhouetteMedian = float(np.median(finite_scores)) - evidence_ids.append(f"candidate:{candidate_id}:graphSilhouette") - except (KeyError, TypeError, ValueError) as exc: - warnings.append(f"Graph silhouette unavailable: {exc}") - - if candidate.reductionMethod == "pca": - try: - separability = store.metric_cluster_separability( - reduction_ref, - {cluster_column: cluster_ref}, - random_seed=CONFIG._RANDOM_SEED, - ) - table = separability.clustering_scores - rows = table.loc[table["clustering"] == cluster_column] - if len(rows): - row = rows.iloc[0] - for field_name, column_name, evidence_name in ( - ("pcaSilhouette", "silhouette_score", "pcaSilhouette"), - ("macroF1", "macro_f1_mean", "macroF1"), - ("weightedF1", "weighted_f1_mean", "weightedF1"), - ): - value = row[column_name] - if value is not None and np.isfinite(float(value)): - setattr(metrics, field_name, float(value)) - evidence_ids.append(f"candidate:{candidate_id}:{evidence_name}") - except (KeyError, TypeError, ValueError) as exc: - warnings.append(f"PCA cluster separability unavailable: {exc}") - - perplexity = max(1.0, float(candidate.neighborsK // 3)) - for column in deps.batchColumns: - try: - score = float( - store.metric_proportional_batch_mixing( - column, - neighbors_ref, - perplexity=perplexity, - ) - ) - if np.isfinite(score): - metrics.batchMixing[column] = score - evidence_ids.append(f"candidate:{candidate_id}:batchMixing:{column}") - except (KeyError, TypeError, ValueError) as exc: - warnings.append(f"Batch mixing for {column!r} unavailable: {exc}") - - for column in deps.preservationColumns: - scores: dict[str, float] = {} - try: - clisi = float( - store.metric_clisi( - column, - neighbors_ref, - perplexity=None, - scale=True, - ) - ) - if np.isfinite(clisi): - scores["clisi"] = clisi - evidence_ids.append(f"candidate:{candidate_id}:clisi:{column}") - except (KeyError, TypeError, ValueError) as exc: - warnings.append(f"cLISI for {column!r} unavailable: {exc}") - try: - connectivity = float( - store.metric_graph_connectivity( - column, - graph_ref, - ) - ) - if np.isfinite(connectivity): - scores["graphConnectivity"] = connectivity - evidence_ids.append( - f"candidate:{candidate_id}:graphConnectivity:{column}" - ) - except (KeyError, TypeError, ValueError) as exc: - warnings.append(f"Graph connectivity for {column!r} unavailable: {exc}") - if scores: - metrics.biologicalPreservation[column] = scores - - eligibility_reasons: list[str] = [] - if n_clusters < 2: - eligibility_reasons.append("fewer than two clusters") - if min_cluster_cells < deps.minClusterCells: - eligibility_reasons.append( - f"smallest cluster has {min_cluster_cells} cells; " - f"minimum is {deps.minClusterCells}" - ) - return metrics, eligibility_reasons - - -def execute_parameter_candidate( - deps: ParameterTuningDependencies, - candidate_id: str, -) -> ParameterCandidateEvaluation: - """Execute one allowlisted candidate without model involvement.""" - - with deps.executionLock: - if candidate_id in deps.evaluations: - logger.debug( - f"Parameter candidate {candidate_id!r} for assay " - f"{deps.fromAssay!r} reused its completed evaluation" - ) - return deps.evaluations[candidate_id] - if candidate_id not in deps.candidates: - logger.warning( - f"Parameter candidate {candidate_id!r} is not authorized for " - f"assay {deps.fromAssay!r}" - ) - return ParameterCandidateEvaluation( - candidateId=candidate_id, - status="failed", - error=( - f"Unknown candidate id {candidate_id!r}; allowed ids are " - f"{sorted(deps.candidates)}" - ), - ) - if len(deps.executionOrder) >= deps.maxCandidates: - logger.warning( - f"Parameter candidate {candidate_id!r} was not executed because " - f"assay {deps.fromAssay!r} reached its limit of " - f"{deps.maxCandidates} candidates" - ) - return ParameterCandidateEvaluation( - candidateId=candidate_id, - phase=deps.candidatePhases.get(candidate_id, "initial"), - harmonyBatchColumns=( - list(deps.batchColumns) - if deps.candidates[candidate_id].useHarmony - else [] - ), - status="failed", - parameters=deps.candidates[candidate_id], - error=f"Candidate execution limit {deps.maxCandidates} reached", - ) - - candidate = deps.candidates[candidate_id] - deps.executionOrder.append(candidate_id) - logger.info( - f"Running parameter candidate {candidate_id!r} for assay " - f"{deps.fromAssay!r}: method={candidate.reductionMethod}, " - f"dimensions={candidate.dimensions}, k={candidate.neighborsK}, " - f"resolution={candidate.leidenResolution}, " - f"harmony={candidate.useHarmony}" - ) - if candidate.useHarmony and not deps.batchColumns: - logger.warning( - f"Parameter candidate {candidate_id!r} cannot run Harmony because " - "no batch columns were authorized" - ) - evaluation = ParameterCandidateEvaluation( - candidateId=candidate_id, - phase=deps.candidatePhases.get(candidate_id, "initial"), - harmonyBatchColumns=[], - status="failed", - parameters=candidate, - error="Harmony candidate requires at least one authorized batch column", - ) - deps.evaluations[candidate_id] = evaluation - return evaluation - - store = deps.store - artifacts: dict[str, ArtifactRecord] = {} - warnings: list[str] = [] - evidence_ids: list[str] = [] - cluster_label = f"agent_tuning_{candidate_id}" - cluster_column = f"{deps.fromAssay}_{cluster_label}" - - try: - normalized_shape = deps.normalizedShape or normalized_artifact_shape( - store, - deps.normalized, - ) - effective_dimensions = validate_parameter_candidate_rank( - candidate, - normalized_shape, - identity_feature_limit=deps.identityFeatureLimit, - ) - reduction_ref, reduction_key, _ = run_candidate_reduction( - store, - normalized=deps.normalized, - candidate=candidate, - normalized_shape=normalized_shape, - identity_feature_limit=deps.identityFeatureLimit, - ) - artifacts[reduction_key] = ArtifactRecord.from_ref(reduction_ref) - logger.debug( - f"Parameter candidate {candidate_id!r}: completed " - f"{reduction_key} reduction" - ) - - coordinates_ref = reduction_ref - if candidate.useHarmony: - coordinates_ref = store.run_harmony( - reduction_ref, - list(deps.batchColumns), - invalidate_cache=False, - ) - artifacts["harmony"] = ArtifactRecord.from_ref(coordinates_ref) - logger.debug( - f"Parameter candidate {candidate_id!r}: completed Harmony " - f"using {len(deps.batchColumns)} batch column(s)" - ) - - ann_ref = store.build_ann_index( - coordinates_ref, - ann_metric="l2", - ann_parallel=False, - rand_state=CONFIG._PCA_RANDOM_SEED, - invalidate_cache=False, - ) - artifacts["annIndex"] = ArtifactRecord.from_ref(ann_ref) - logger.debug( - f"Parameter candidate {candidate_id!r}: completed ANN indexing" - ) - - neighbors_ref = store.query_neighbors( - ann_ref, - coordinates=coordinates_ref, - k=candidate.neighborsK, - invalidate_cache=False, - ) - artifacts["neighbors"] = ArtifactRecord.from_ref(neighbors_ref) - logger.debug( - f"Parameter candidate {candidate_id!r}: completed neighbor query" - ) - - graph_ref = store.build_connectivity_map( - neighbors_ref, - local_connectivity=1.0, - bandwidth=1.5, - invalidate_cache=False, - ) - artifacts["connectivityMap"] = ArtifactRecord.from_ref(graph_ref) - logger.debug( - f"Parameter candidate {candidate_id!r}: completed connectivity map" - ) - - cluster_ref = store.run_leiden_clustering( - graph_ref, - resolution=candidate.leidenResolution, - backend="igraph", - symmetric_graph=False, - graph_upper_only=False, - random_seed=CONFIG._RANDOM_SEED, - invalidate_cache=False, - ) - artifacts["clusters"] = ArtifactRecord.from_ref(cluster_ref) - logger.debug( - f"Parameter candidate {candidate_id!r}: completed Leiden clustering" - ) - - metrics, eligibility_reasons = _collect_parameter_candidate_metrics( - deps, - candidate=candidate, - candidate_id=candidate_id, - reduction_ref=reduction_ref, - neighbors_ref=neighbors_ref, - graph_ref=graph_ref, - cluster_ref=cluster_ref, - cluster_column=cluster_column, - evidence_ids=evidence_ids, - warnings=warnings, - ) - - evaluation = ParameterCandidateEvaluation( - candidateId=candidate_id, - phase=deps.candidatePhases.get(candidate_id, "initial"), - harmonyBatchColumns=( - list(deps.batchColumns) if candidate.useHarmony else [] - ), - status="done", - eligible=not eligibility_reasons, - parameters=candidate, - artifacts=artifacts, - cellSelection=artifact_reference(deps.cellSelection), - clusterColumn=cluster_column, - clusterLabel=cluster_label, - effectiveDimensions=effective_dimensions, - metrics=metrics, - evidenceIds=evidence_ids, - eligibilityReasons=eligibility_reasons, - warnings=warnings, - ) - logger.info( - f"Completed parameter candidate {candidate_id!r} for assay " - f"{deps.fromAssay!r}: eligible={evaluation.eligible}, " - f"clusters={metrics.nClusters}, " - f"minimum_cluster_cells={metrics.minClusterCells}, " - f"warnings={len(warnings)}" - ) - except (KeyError, TypeError, ValueError, RuntimeError) as exc: - evaluation = ParameterCandidateEvaluation( - candidateId=candidate_id, - phase=deps.candidatePhases.get(candidate_id, "initial"), - harmonyBatchColumns=( - list(deps.batchColumns) if candidate.useHarmony else [] - ), - status="failed", - parameters=candidate, - artifacts=artifacts, - cellSelection=( - artifact_reference(deps.cellSelection) - if deps.cellSelection is not None - else None - ), - evidenceIds=evidence_ids, - warnings=warnings, - error=str(exc), - ) - logger.warning( - f"Parameter candidate {candidate_id!r} for assay " - f"{deps.fromAssay!r} failed: {exc}" - ) - - deps.evaluations[candidate_id] = evaluation - return evaluation - - -async def evaluate_parameter_candidate( - ctx: RunContext[ParameterTuningDependencies], - candidate_id: str, -) -> ParameterCandidateEvaluation: - """Expose deterministic candidate execution as a bounded agent tool.""" - - return execute_parameter_candidate(ctx.deps, candidate_id) - - -def validate_parameter_search_plan( - plan: ParameterSearchPlan, - deps: ParameterTuningDependencies, - *, - initial_candidate_ids: Sequence[str], - max_refined_candidates: int, -) -> ParameterSearchPlan: - """Validate one refinement proposal against the completed initial screen.""" - - initial_evaluations = [ - deps.evaluations[candidate_id] - for candidate_id in initial_candidate_ids - if candidate_id in deps.evaluations - ] - known_evidence = { - evidence_id - for evaluation in initial_evaluations - for evidence_id in evaluation.evidenceIds - } - unknown_evidence = sorted(set(plan.evidenceIds) - known_evidence) - if unknown_evidence: - raise ValueError( - f"Parameter search plan cites unknown evidence ids {unknown_evidence}" - ) - authorized_batch_columns = list(deps.batchColumns) if deps.harmonyAuthorized else [] - if ( - plan.harmonyBatchColumns - and plan.harmonyBatchColumns != authorized_batch_columns - ): - raise ValueError( - "Parameter search plan cannot modify the exact authorized Harmony " - "batch columns" - ) - canonical_status: ParameterSearchStatus = ( - "refine" if plan.candidates else "complete" - ) - plan = plan.model_copy( - update={ - "status": canonical_status, - "harmonyBatchColumns": authorized_batch_columns, - } - ) - if plan.status == "complete": - return plan - - if len(plan.candidates) > max_refined_candidates: - raise ValueError( - "Parameter search plan exceeds the refined candidate limit " - f"{max_refined_candidates}" - ) - if not plan.rationale.strip(): - raise ValueError("A refinement plan requires a rationale") - if not plan.objectives: - raise ValueError("A refinement plan requires focused objectives") - if not plan.stoppingCriteria: - raise ValueError("A refinement plan requires stopping criteria") - if not plan.evidenceIds: - raise ValueError("A refinement plan requires initial-screen evidence") - - successful_initial_ids = { - evaluation.candidateId - for evaluation in initial_evaluations - if evaluation.status == "done" - } - if not plan.basedOnCandidateIds: - raise ValueError("A refinement plan must identify its initial candidates") - duplicate_parents = sorted( - { - candidate_id - for candidate_id in plan.basedOnCandidateIds - if plan.basedOnCandidateIds.count(candidate_id) > 1 - } - ) - if duplicate_parents: - raise ValueError(f"Duplicate refinement parent ids {duplicate_parents}") - invalid_parents = sorted(set(plan.basedOnCandidateIds) - successful_initial_ids) - if invalid_parents: - raise ValueError( - "Refinement parents must be successful initial candidates: " - f"{invalid_parents}" - ) - for parent_id in plan.basedOnCandidateIds: - prefix = f"candidate:{parent_id}:" - if not any(evidence_id.startswith(prefix) for evidence_id in plan.evidenceIds): - raise ValueError( - f"Refinement evidence must cite every parent candidate: {parent_id!r}" - ) - if deps.harmonyAuthorized: - parent_candidates = [ - deps.candidates[candidate_id] for candidate_id in plan.basedOnCandidateIds - ] - paired_modes: dict[tuple[str, int, float, int], set[bool]] = {} - for candidate in parent_candidates: - parameter_key = ( - candidate.reductionMethod, - candidate.dimensions, - candidate.leidenResolution, - candidate.neighborsK, - ) - paired_modes.setdefault(parameter_key, set()).add(candidate.useHarmony) - if not any(modes == {False, True} for modes in paired_modes.values()): - raise ValueError( - "Harmony refinement requires evidence from one matched corrected " - "and uncorrected initial pair" - ) - - initial_candidates = [ - deps.candidates[candidate_id] for candidate_id in initial_candidate_ids - ] - known_signatures = { - ( - candidate.reductionMethod, - candidate.dimensions, - candidate.leidenResolution, - candidate.neighborsK, - candidate.useHarmony, - ) - for candidate in initial_candidates - } - proposed_ids: set[str] = set() - proposed_signatures: set[tuple[str, int, float, int, bool]] = set() - for candidate in plan.candidates: - if not CONFIG._CANDIDATE_ID.fullmatch(candidate.candidateId): - raise ValueError( - "Refined candidateId must contain only ASCII letters, numbers, " - "and underscores" - ) - if ( - candidate.candidateId in deps.candidates - or candidate.candidateId in proposed_ids - ): - raise ValueError(f"Duplicate refined candidateId {candidate.candidateId!r}") - proposed_ids.add(candidate.candidateId) - method_candidates = [ - item - for item in initial_candidates - if item.reductionMethod == candidate.reductionMethod - ] - if not method_candidates: - raise ValueError( - "Refined candidates cannot introduce an untested reduction method: " - f"{candidate.reductionMethod!r}" - ) - dimension_bounds = ( - min(item.dimensions for item in method_candidates), - max(item.dimensions for item in method_candidates), - ) - resolution_bounds = ( - min(item.leidenResolution for item in method_candidates), - max(item.leidenResolution for item in method_candidates), - ) - neighbor_bounds = ( - min(item.neighborsK for item in method_candidates), - max(item.neighborsK for item in method_candidates), - ) - if not dimension_bounds[0] <= candidate.dimensions <= dimension_bounds[1]: - raise ValueError( - "Refined dimensions must remain inside the initial search envelope " - f"{dimension_bounds}" - ) - if not ( - resolution_bounds[0] <= candidate.leidenResolution <= resolution_bounds[1] - ): - raise ValueError( - "Refined Leiden resolution must remain inside the initial search " - f"envelope {resolution_bounds}" - ) - if not neighbor_bounds[0] <= candidate.neighborsK <= neighbor_bounds[1]: - raise ValueError( - "Refined neighbor count must remain inside the initial search " - f"envelope {neighbor_bounds}" - ) - if candidate.useHarmony and ( - not deps.harmonyAuthorized or not deps.batchColumns - ): - raise ValueError( - f"Refined candidate {candidate.candidateId!r} is not authorized " - "for Harmony" - ) - signature = ( - candidate.reductionMethod, - candidate.dimensions, - candidate.leidenResolution, - candidate.neighborsK, - candidate.useHarmony, - ) - if signature in known_signatures or signature in proposed_signatures: - raise ValueError( - f"Refined candidate {candidate.candidateId!r} duplicates an " - "evaluated or proposed parameter branch" - ) - proposed_signatures.add(signature) - return plan - - -def validate_parameter_batch_search_plan( - plan: ParameterTuningBatchSearchPlan, - dependencies: Mapping[str, ParameterTuningDependencies], - *, - initial_candidate_ids: Mapping[str, Sequence[str]], - max_refined_by_assay: Mapping[str, int], -) -> ParameterTuningBatchSearchPlan: - """Validate every assay entry in one batched refinement response.""" - - expected = set(dependencies) - actual = set(plan.assayPlans) - if actual != expected: - raise ValueError( - "Batched refinement must contain exactly the requested assays: " - f"missing={sorted(expected - actual)}, unexpected={sorted(actual - expected)}" - ) - validated = { - assay: validate_parameter_search_plan( - plan.assayPlans[assay], - dependencies[assay], - initial_candidate_ids=initial_candidate_ids[assay], - max_refined_candidates=max_refined_by_assay[assay], - ) - for assay in dependencies - } - return plan.model_copy(update={"assayPlans": validated}) - - -def validate_parameter_tuning_report( - report: ParameterTuningReport, - deps: ParameterTuningDependencies, - *, - search_plan: ParameterSearchPlan | None = None, -) -> ParameterTuningReport: - """Ground the model report in candidate executions recorded by the tool.""" - - evaluations = [ - deps.evaluations[candidate_id] - for candidate_id in deps.executionOrder - if candidate_id in deps.evaluations - ] - known_evidence = { - evidence_id - for evaluation in evaluations - for evidence_id in evaluation.evidenceIds - } - cited_evidence = set(report.evidenceIds) - for comparison in report.comparisons: - cited_evidence.update(comparison.evidenceIds) - if report.needsInput is not None: - cited_evidence.update(report.needsInput.evidenceIds) - unknown_evidence = sorted(cited_evidence - known_evidence) - if unknown_evidence: - raise ValueError( - f"Parameter tuning report cites unknown evidence ids {unknown_evidence}" - ) - if report.status == "done" and report.recommendedCandidateId is None: - raise ValueError("A done tuning report must recommend an executed candidate") - if report.status == "needsInput" and report.needsInput is None: - raise ValueError("A needsInput tuning report must include a concrete question") - successful = [ - evaluation for evaluation in evaluations if evaluation.status == "done" - ] - comparison_required = len(deps.candidates) > 1 and deps.maxCandidates > 1 - if report.status == "done": - if not report.evidenceIds: - raise ValueError("A done tuning report requires recommendation evidence") - if comparison_required and len(successful) < 2: - raise ValueError( - "A completed tuning recommendation requires at least two successful " - "candidate executions" - ) - if ( - comparison_required - and "baseline" in deps.candidates - and not any(item.candidateId == "baseline" for item in successful) - ): - raise ValueError( - "Evaluate the baseline before completing a multi-candidate comparison" - ) - - selected_artifacts: dict[str, ArtifactRecord] = {} - if report.recommendedCandidateId is not None: - selected = deps.evaluations.get(report.recommendedCandidateId) - if selected is None: - raise ValueError("Recommended candidate was not executed") - if selected.status != "done": - raise ValueError("Recommended candidate execution failed") - if not selected.eligible: - raise ValueError("Recommended candidate is not eligible") - recommendation_prefix = f"candidate:{selected.candidateId}:" - if not any( - evidence_id.startswith(recommendation_prefix) - for evidence_id in report.evidenceIds - ): - raise ValueError( - "Recommendation evidence must include the selected candidate" - ) - selected_artifacts = dict(selected.artifacts) - - if report.status == "done" and not comparison_required and report.comparisons: - raise ValueError( - "Candidate comparisons require a completed multi-candidate evaluation" - ) - if report.status == "done" and comparison_required: - assert report.recommendedCandidateId is not None - successful_ids = {item.candidateId for item in successful} - expected_comparators = successful_ids - {report.recommendedCandidateId} - comparison_ids = [item.candidateId for item in report.comparisons] - duplicate_comparators = sorted( - { - candidate_id - for candidate_id in comparison_ids - if comparison_ids.count(candidate_id) > 1 - } - ) - if duplicate_comparators: - raise ValueError(f"Duplicate candidate comparisons {duplicate_comparators}") - actual_comparators = set(comparison_ids) - missing_comparators = sorted(expected_comparators - actual_comparators) - invalid_comparators = sorted(actual_comparators - expected_comparators) - if missing_comparators: - raise ValueError( - "Completed tuning reports require comparisons for every successful " - f"non-selected candidate: {missing_comparators}" - ) - if invalid_comparators: - raise ValueError( - "Candidate comparisons must identify successful non-selected " - f"candidates: {invalid_comparators}" - ) - selected_prefix = f"candidate:{report.recommendedCandidateId}:" - for comparison in report.comparisons: - comparator_prefix = f"candidate:{comparison.candidateId}:" - if not any( - evidence_id.startswith(selected_prefix) - for evidence_id in comparison.evidenceIds - ): - raise ValueError( - "Each candidate comparison must cite evidence from the " - "selected candidate" - ) - if not any( - evidence_id.startswith(comparator_prefix) - for evidence_id in comparison.evidenceIds - ): - raise ValueError( - "Each candidate comparison must cite evidence from its comparator" - ) - if not comparison.summary.strip(): - raise ValueError( - "Each candidate comparison requires a concise grounded summary" - ) - - return report.model_copy( - update={ - "fromAssay": deps.fromAssay, - "cellSelection": artifact_reference(deps.cellSelection), - "evaluations": evaluations, - "selectedArtifacts": selected_artifacts, - "searchPlan": search_plan, - "assayReports": {}, - "recommendedByAssay": ( - {deps.fromAssay: report.recommendedCandidateId} - if report.recommendedCandidateId is not None - else {} - ), - "totalCandidates": len(evaluations), - "integrationEvaluations": [], - "recommendedIntegrationId": None, - "finalClusterColumn": None, - "finalClusterArtifact": None, - "graphAssay": deps.fromAssay, - "markerAssay": deps.fromAssay, - "finalSelection": None, - } - ) - - -def validate_parameter_tuning_batch_report( - report: ParameterTuningReport, - dependencies: Mapping[str, ParameterTuningDependencies], - *, - search_plans: Mapping[str, ParameterSearchPlan], - primary_assay: str, -) -> ParameterTuningReport: - """Ground one aggregate response in every assay's executed branches.""" - - expected = set(dependencies) - actual = set(report.assayReports) - if actual != expected: - raise ValueError( - "Batched selection must contain exactly the requested assays: " - f"missing={sorted(expected - actual)}, unexpected={sorted(actual - expected)}" - ) - if primary_assay not in dependencies: - raise ValueError(f"Unknown primary assay {primary_assay!r}") - validated_reports = { - assay: validate_parameter_tuning_report( - report.assayReports[assay], - dependencies[assay], - search_plan=search_plans[assay], - ) - for assay in dependencies - } - known_evidence = { - evidence_id - for assay_report in validated_reports.values() - for evaluation in assay_report.evaluations - for evidence_id in evaluation.evidenceIds - } - unknown_evidence = sorted(set(report.evidenceIds) - known_evidence) - if unknown_evidence: - raise ValueError( - f"Batched tuning report cites unknown evidence ids {unknown_evidence}" - ) - statuses = {assay_report.status for assay_report in validated_reports.values()} - if statuses == {"done"}: - status: StageStatus = "done" - elif "needsInput" in statuses: - status = "needsInput" - else: - status = "failed" - primary = validated_reports[primary_assay] - recommended = { - assay: assay_report.recommendedCandidateId - for assay, assay_report in validated_reports.items() - if assay_report.recommendedCandidateId is not None - } - if status == "done" and len(recommended) != len(validated_reports): - raise ValueError("Every completed assay report must recommend a candidate") - cell_selections = [ - core_artifact_reference(assay_report.cellSelection) - for assay_report in validated_reports.values() - ] - if ( - not cell_selections - or not isinstance(cell_selections[0], ArtifactRef) - or any(selection != cell_selections[0] for selection in cell_selections[1:]) - ): - raise ValueError("Every assay report must use the same exact cell selection") - return report.model_copy( - update={ - "status": status, - "fromAssay": primary_assay, - "cellSelection": primary.cellSelection, - "evaluations": primary.evaluations, - "recommendedCandidateId": primary.recommendedCandidateId, - "selectedArtifacts": primary.selectedArtifacts, - "needsInput": primary.needsInput if status != "done" else None, - "searchPlan": primary.searchPlan, - "assayReports": validated_reports, - "recommendedByAssay": recommended, - "totalCandidates": sum( - len(item.evaluations) for item in validated_reports.values() - ), - "integrationEvaluations": [], - "recommendedIntegrationId": None, - "finalClusterColumn": None, - "finalClusterArtifact": None, - "graphAssay": primary_assay, - "markerAssay": primary_assay, - "finalSelection": None, - } - ) - - -def fallback_parameter_tuning_report( - deps: ParameterTuningDependencies, - *, - search_plan: ParameterSearchPlan, - agent_name: str, -) -> ParameterTuningReport: - """Retain the first eligible branch when structured selection is unavailable.""" - - evaluations = [ - deps.evaluations[candidate_id] - for candidate_id in deps.executionOrder - if candidate_id in deps.evaluations - ] - successful = [item for item in evaluations if item.status == "done"] - eligible = [item for item in successful if item.eligible] - comparison_required = len(deps.candidates) > 1 and deps.maxCandidates > 1 - evidence_by_candidate = { - item.candidateId: next( - ( - evidence_id - for evidence_id in item.evidenceIds - if evidence_id == f"candidate:{item.candidateId}:clusters" - ), - next(iter(item.evidenceIds), None), - ) - for item in successful - } - cannot_recommend = ( - not eligible - or (comparison_required and len(successful) < 2) - or any(evidence_by_candidate[item.candidateId] is None for item in successful) - ) - if cannot_recommend: - logger.warning( - f"Parameter tuning fallback for assay {deps.fromAssay!r} requires " - f"input: completed={len(successful)}, eligible={len(eligible)}" - ) - known_evidence = sorted( - { - evidence_id - for evaluation in evaluations - for evidence_id in evaluation.evidenceIds - } - ) - report = ParameterTuningReport( - status="needsInput", - confidence="low", - rationale=( - "Structured model selection was unavailable and the executed " - "screen does not support a conservative automatic fallback." - ), - limitations=[ - "No parameter branch was selected without complete eligible " - "executor evidence." - ], - stopReason="The bounded screen completed without an automatic choice.", - needsInput=ParameterTuningNeedsInput( - question="Select one eligible executed parameter candidate.", - options=[item.candidateId for item in eligible], - evidenceIds=known_evidence, - ), - runInfo=AgentRunInfo(agentName=agent_name), - ) - return validate_parameter_tuning_report( - report, - deps, - search_plan=search_plan, - ) - - selected = eligible[0] - logger.warning( - f"Parameter tuning fallback retained candidate " - f"{selected.candidateId!r} for assay {deps.fromAssay!r} from " - f"{len(eligible)} eligible candidate(s)" - ) - selected_evidence = evidence_by_candidate[selected.candidateId] - assert selected_evidence is not None - comparisons: list[CandidateComparison] = [] - if comparison_required: - for item in successful: - if item.candidateId == selected.candidateId: - continue - comparator_evidence = evidence_by_candidate[item.candidateId] - assert comparator_evidence is not None - comparisons.append( - CandidateComparison( - candidateId=item.candidateId, - summary=( - "This executed branch remains a grounded comparator to the " - "conservatively retained first eligible branch." - ), - evidenceIds=[selected_evidence, comparator_evidence], - ) - ) - report = ParameterTuningReport( - status="done", - recommendedCandidateId=selected.candidateId, - confidence="low", - rationale=( - "Structured model selection was unavailable after bounded retries; " - "the first eligible authorized branch was retained conservatively." - ), - evidenceIds=[selected_evidence], - comparisons=comparisons, - tradeoffs=["No model-authored metric trade-off ranking was available."], - limitations=[ - "The fallback does not claim that the retained branch is metric-optimal." - ], - stopReason=( - "The deterministic screen completed and retained its first eligible " - "authorized branch." - ), - runInfo=AgentRunInfo(agentName=agent_name), - ) - return validate_parameter_tuning_report( - report, - deps, - search_plan=search_plan, - ) - - -def fallback_parameter_tuning_batch_report( - dependencies: Mapping[str, ParameterTuningDependencies], - *, - search_plans: Mapping[str, ParameterSearchPlan], - primary_assay: str, -) -> ParameterTuningReport: - """Build one grounded aggregate fallback over completed assay screens.""" - - logger.warning( - f"Using parameter tuning batch fallback for {len(dependencies)} assay(s)" - ) - assay_reports = { - assay: fallback_parameter_tuning_report( - deps, - search_plan=search_plans[assay], - agent_name="parameter_tuning_batch_fallback", - ) - for assay, deps in dependencies.items() - } - if any(item.status != "done" for item in assay_reports.values()): - primary = assay_reports[primary_assay] - if primary.status == "done": - primary = ParameterTuningReport( - status="needsInput", - confidence="low", - rationale=( - "At least one assay lacks a conservative automatic parameter " - "selection." - ), - limitations=[ - "The multimodal native screen requires an explicit selection." - ], - stopReason="The bounded native screens completed without all choices.", - needsInput=ParameterTuningNeedsInput( - question="Select eligible parameter candidates for every assay.", - options=[], - evidenceIds=list(primary.evidenceIds), - ), - runInfo=AgentRunInfo(agentName="parameter_tuning_batch_fallback"), - ) - assay_reports[primary_assay] = validate_parameter_tuning_report( - primary, - dependencies[primary_assay], - search_plan=search_plans[primary_assay], - ) - aggregate = ParameterTuningReport( - status=( - "done" - if all(item.status == "done" for item in assay_reports.values()) - else "needsInput" - ), - assayReports=assay_reports, - rationale=( - "Structured model selection was unavailable; each completed native " - "screen used the conservative fallback policy." - ), - evidenceIds=list( - dict.fromkeys( - evidence_id - for assay_report in assay_reports.values() - for evidence_id in assay_report.evidenceIds - ) - ), - limitations=[ - "Fallback recommendations retain first eligible authorized branches " - "without claiming a metric-optimal ranking." - ], - stopReason="The bounded native screens completed.", - runInfo=AgentRunInfo(agentName="parameter_tuning_batch_fallback"), - ) - logger.warning( - f"Parameter tuning batch fallback status={aggregate.status}; " - f"completed_assays={sum(item.status == 'done' for item in assay_reports.values())}" - ) - return validate_parameter_tuning_batch_report( - aggregate, - dependencies, - search_plans=search_plans, - primary_assay=primary_assay, - ) - - -def validate_final_graph_selection( - selection: FinalGraphSelection, - report: ParameterTuningReport, - *, - integration_evaluations: Sequence[IntegrationCandidateEvaluation], - marker_assay: str, -) -> FinalGraphSelection: - """Ground a final graph choice in the exact eligible executor outputs.""" - - if report.status != "done": - raise ValueError("Native parameter tuning must finish before graph selection") - options = final_graph_options(report, integration_evaluations) - if not options: - raise ValueError("No eligible native or integrated graph options are available") - known_evidence = { - evidence_id - for option in options.values() - for evidence_id in option["evidenceIds"] - } - cited = set(selection.evidenceIds) - for comparison in selection.comparisons: - cited.update(comparison.evidenceIds) - if selection.needsInput is not None: - cited.update(selection.needsInput.evidenceIds) - unknown = sorted(cited - known_evidence) - if unknown: - raise ValueError(f"Final graph selection cites unknown evidence ids {unknown}") - if selection.status == "needsInput": - if selection.needsInput is None or not selection.needsInput.question.strip(): - raise ValueError( - "A needsInput graph selection requires a concrete question" - ) - return selection.model_copy( - update={ - "selectedOptionId": None, - "graphMethod": None, - "nativeAssay": None, - "nativeCandidateId": None, - "integrationId": None, - "markerAssay": marker_assay, - } - ) - if selection.status != "done": - raise ValueError("Final graph selection must be done or needsInput") - if selection.selectedOptionId not in options: - raise ValueError("Selected final graph option is not eligible") - assert selection.selectedOptionId is not None - selected = options[selection.selectedOptionId] - selected_evidence = set(selected["evidenceIds"]) - if not selected_evidence.intersection(selection.evidenceIds): - raise ValueError( - "Final graph recommendation must cite selected-option evidence" - ) - expected_comparators = set(options) - {selection.selectedOptionId} - comparison_ids = [item.optionId for item in selection.comparisons] - if len(set(comparison_ids)) != len(comparison_ids): - raise ValueError("Final graph comparisons must not contain duplicates") - if set(comparison_ids) != expected_comparators: - raise ValueError( - "Final graph selection requires one comparison for every eligible " - "non-selected option" - ) - for comparison in selection.comparisons: - comparator_evidence = set(options[comparison.optionId]["evidenceIds"]) - if not selected_evidence.intersection(comparison.evidenceIds): - raise ValueError( - "Every final graph comparison must cite selected-option evidence" - ) - if not comparator_evidence.intersection(comparison.evidenceIds): - raise ValueError( - "Every final graph comparison must cite comparator evidence" - ) - if not comparison.summary.strip(): - raise ValueError("Every final graph comparison requires a summary") - return selection.model_copy( - update={ - "graphMethod": selected["graphMethod"], - "nativeAssay": selected.get("nativeAssay"), - "nativeCandidateId": selected.get("nativeCandidateId"), - "integrationId": selected.get("integrationId"), - "markerAssay": marker_assay, - "needsInput": None, - } - ) - - -def finalize_parameter_tuning_selection( - report: ParameterTuningReport, - *, - marker_assay: str, - integration_evaluations: Sequence[IntegrationCandidateEvaluation] = (), - recommended_integration_id: str | None = None, - native_assay: str | None = None, - final_selection: FinalGraphSelection | None = None, -) -> ParameterTuningReport: - """Attach an executor-selected native or integrated final cluster branch.""" - - logger.debug( - f"Finalizing parameter graph selection: marker_assay={marker_assay!r}, " - f"integration_candidates={len(integration_evaluations)}" - ) - if report.status != "done": - raise ValueError("Parameter tuning must be done before final graph selection") - if not marker_assay: - raise ValueError("marker_assay must be non-empty") - report_cell_selection = core_artifact_reference(report.cellSelection) - if not isinstance(report_cell_selection, ArtifactRef): - raise ValueError("Parameter tuning report lacks an exact cell selection") - assay_reports = report.assayReports or {report.fromAssay: report} - if marker_assay not in assay_reports: - raise ValueError(f"Unknown marker assay {marker_assay!r}") - evaluations = list(integration_evaluations) - integration_ids = [item.integrationId for item in evaluations] - if len(set(integration_ids)) != len(integration_ids): - raise ValueError("Integration evaluation ids must be unique") - if recommended_integration_id is not None and native_assay is not None: - raise ValueError("Choose either an integrated graph or one native assay") - if recommended_integration_id is not None: - selected = next( - ( - item - for item in evaluations - if item.integrationId == recommended_integration_id - ), - None, - ) - if selected is None: - raise ValueError("Recommended integration candidate was not evaluated") - if selected.status != "done" or not selected.eligible: - raise ValueError("Recommended integration candidate is not eligible") - if selected.clusterArtifact is None: - raise ValueError("Recommended integration lacks an exact cluster artifact") - if core_artifact_reference(selected.cellSelection) != report_cell_selection: - raise ValueError("Recommended integration uses a different cell selection") - if ( - selected.clusterArtifact.scope != "datastore" - or selected.clusterArtifact.assay is not None - ): - raise ValueError( - "Integrated cluster artifacts must be datastore-scoped without assay" - ) - if ( - selected.graphArtifact is None - or selected.graphArtifact.scope != "datastore" - ): - raise ValueError("Integrated graph artifact must be datastore-scoped") - cluster_artifact = selected.clusterArtifact - cluster_column = selected.clusterColumn - graph_assay = None - else: - selected_assay = native_assay or report.fromAssay - primary = assay_reports.get(selected_assay) - if primary is None or primary.recommendedCandidateId is None: - raise ValueError("Selected assay lacks a native tuning recommendation") - selected_native = next( - ( - item - for item in primary.evaluations - if item.candidateId == primary.recommendedCandidateId - ), - None, - ) - if ( - selected_native is None - or selected_native.status != "done" - or not selected_native.eligible - or "clusters" not in selected_native.artifacts - ): - raise ValueError("Primary native recommendation lacks exact clusters") - if ( - core_artifact_reference(selected_native.cellSelection) - != report_cell_selection - ): - raise ValueError( - "Recommended native candidate uses a different cell selection" - ) - cluster_artifact = selected_native.artifacts["clusters"] - cluster_column = selected_native.clusterColumn - graph_assay = selected_assay - finalized = report.model_copy( - update={ - "totalCandidates": ( - sum(len(value.evaluations) for value in assay_reports.values()) - + len(evaluations) - ), - "integrationEvaluations": evaluations, - "recommendedIntegrationId": recommended_integration_id, - "finalClusterColumn": cluster_column, - "finalClusterArtifact": cluster_artifact, - "graphAssay": graph_assay, - "markerAssay": marker_assay, - "finalSelection": final_selection, - } - ) - selected_graph = recommended_integration_id or graph_assay - logger.info( - f"Finalized parameter graph selection: graph={selected_graph!r}, " - f"marker_assay={marker_assay!r}, cluster_column={cluster_column!r}" - ) - return finalized - - -def select_final_parameter_graph( - *, - model: Any, - report: ParameterTuningReport, - integration_evaluations: Sequence[IntegrationCandidateEvaluation], - marker_assay: str, - config: AgentRunConfig | None = None, -) -> ParameterTuningReport: - """Use one bounded provider call to select and attach the final graph.""" - - evaluations = list(integration_evaluations) - if not marker_assay: - raise ValueError("marker_assay must be non-empty") - assay_reports = report.assayReports or {report.fromAssay: report} - if marker_assay not in assay_reports: - raise ValueError(f"Unknown marker assay {marker_assay!r}") - options = final_graph_options(report, evaluations) - if not options: - raise ValueError("No eligible native or integrated graph options are available") - logger.info( - f"Selecting final parameter graph from {len(options)} eligible option(s); " - f"marker_assay={marker_assay!r}" - ) - if len(options) == 1: - option_id, option = next(iter(options.items())) - logger.info( - f"Selecting sole eligible final graph option {option_id!r} " - "without a provider request" - ) - selection = validate_final_graph_selection( - FinalGraphSelection( - status="done", - selectedOptionId=option_id, - markerAssay=marker_assay, - confidence="high", - rationale="The executor produced exactly one eligible graph option.", - evidenceIds=list(option["evidenceIds"]), - limitations=["No alternative eligible final graph required ranking."], - runInfo=AgentRunInfo( - agentName="parameter_tuning_final_graph_deterministic" - ), - ), - report, - integration_evaluations=evaluations, - marker_assay=marker_assay, - ) - else: - run_config = (config or AgentRunConfig()).with_limits( - request_limit=6, - tool_call_limit=5, - output_token_limit=32768, - timeout_seconds=600.0, - ) - try: - logger.info( - f"Requesting final graph selection across {len(options)} " - "eligible options" - ) - execution = run_agent_sync( - model=model, - output_type=FinalGraphSelection, - system_prompt=final_graph_selection_system_prompt(), - user_prompt=final_graph_selection_prompt( - report=report, - integration_evaluations=evaluations, - marker_assay=marker_assay, - ), - deps_type=ParameterTuningDependencies, - deps=ParameterTuningDependencies.get_blank(), - config=run_config, - name="parameter_tuning_final_graph", - output_validator=lambda proposed: validate_final_graph_selection( - proposed, - report, - integration_evaluations=evaluations, - marker_assay=marker_assay, - ), - ) - except UnexpectedModelBehavior: - option_ids = sorted(options) - logger.warning( - "Final graph selection exhausted structured-output retries; " - f"requesting input for {len(option_ids)} eligible options" - ) - selection = validate_final_graph_selection( - FinalGraphSelection( - status="needsInput", - markerAssay=marker_assay, - confidence="low", - rationale=( - "Structured final-graph selection was unavailable after " - "bounded retries." - ), - limitations=[ - "No ranking was invented across multiple eligible graphs." - ], - needsInput=FinalGraphNeedsInput( - question="Select one eligible final graph option.", - options=option_ids, - evidenceIds=sorted( - { - evidence_id - for option in options.values() - for evidence_id in option["evidenceIds"] - } - ), - ), - runInfo=AgentRunInfo( - agentName="parameter_tuning_final_graph_fallback" - ), - ), - report, - integration_evaluations=evaluations, - marker_assay=marker_assay, - ) - else: - if not isinstance(execution.output, FinalGraphSelection): - raise TypeError( - "Final graph selector returned an unexpected output type" - ) - selection = validate_final_graph_selection( - execution.output, - report, - integration_evaluations=evaluations, - marker_assay=marker_assay, - ).model_copy(update={"runInfo": execution.runInfo}) - logger.info( - f"Provider selected final graph option {selection.selectedOptionId!r}" - ) - if selection.status == "needsInput": - needs_input = selection.needsInput or FinalGraphNeedsInput.get_blank() - option_ids = sorted(options) - canonical_question = "Select one eligible final graph option." - canonical_needs_input = needs_input.model_copy( - update={"question": canonical_question, "options": option_ids} - ) - logger.warning( - f"Final parameter graph selection needs input; options={len(option_ids)}" - ) - return report.model_copy( - update={ - "status": "needsInput", - "totalCandidates": ( - sum(len(value.evaluations) for value in assay_reports.values()) - + len(evaluations) - ), - "integrationEvaluations": evaluations, - "markerAssay": marker_assay, - "finalSelection": selection.model_copy( - update={"needsInput": canonical_needs_input} - ), - "needsInput": ParameterTuningNeedsInput( - question=canonical_question, - options=option_ids, - evidenceIds=needs_input.evidenceIds, - ), - } - ) - return finalize_parameter_tuning_selection( - report, - marker_assay=marker_assay, - integration_evaluations=evaluations, - recommended_integration_id=selection.integrationId, - native_assay=selection.nativeAssay, - final_selection=selection, - ) - - -def promote_parameter_candidate( - store: Any, - *, - report: ParameterTuningReport, - normalized: Any, - identity_feature_limit: int = 64, -) -> ParameterCandidateEvaluation: - """Resolve and verify the exact selected native branch without replaying it.""" - - if report.status != "done" or report.recommendedCandidateId is None: - raise ValueError("A completed native tuning recommendation is required") - evaluation = next( - ( - item - for item in report.evaluations - if item.candidateId == report.recommendedCandidateId - ), - None, - ) - if evaluation is None or evaluation.status != "done" or not evaluation.eligible: - raise ValueError("Recommended candidate is not an eligible execution") - if "clusters" not in evaluation.artifacts: - raise ValueError("Recommended candidate lacks an exact cluster artifact") - normalized_ref = core_artifact_reference(normalized) - if ( - not isinstance(normalized_ref, ArtifactRef) - or normalized_ref.kind != "normalized" - or normalized_ref.assay != report.fromAssay - ): - raise ValueError( - "normalized must identify the report's exact normalized assay artifact" - ) - status = store.inspect_artifact(normalized_ref) - if not getattr(status, "exists", True) or not getattr(status, "complete", False): - raise ValueError("normalized artifact is unavailable or incomplete") - raw_selection = (getattr(status, "inputs", None) or {}).get("cell_selection") - if not isinstance(raw_selection, Mapping): - raise ValueError("normalized artifact has no cell-selection input") - normalized_selection = ArtifactRef.from_dict(dict(raw_selection)) - if normalized_selection != core_artifact_reference(evaluation.cellSelection): - raise ValueError( - "Recommended candidate does not match normalized artifact lineage" - ) - if normalized_selection != core_artifact_reference(report.cellSelection): - raise ValueError( - "Parameter tuning report does not match normalized artifact lineage" - ) - if identity_feature_limit < 2: - raise ValueError("identity_feature_limit must be at least two") - logger.info( - f"Resolved parameter candidate {evaluation.candidateId!r} for assay " - f"{report.fromAssay!r} without replay" - ) - return evaluation - - -def _resolve_experimental_tuning_handoff( - *, - normalized_cell_selection: ArtifactRef, - batch_columns: Sequence[str], - preservation_columns: Sequence[str], - experimental_handoff: ExperimentalTuningHandoff | None, -) -> tuple[ArtifactRef, list[str], list[str]]: - resolved_batch_columns = list(batch_columns) - resolved_preservation_columns = list(preservation_columns) - if experimental_handoff is None: - return ( - normalized_cell_selection, - resolved_batch_columns, - resolved_preservation_columns, - ) - - handoff_batch_columns = list(experimental_handoff.batchColumns) - canonical_batch_columns = sorted(set(handoff_batch_columns)) - if len(canonical_batch_columns) != len(handoff_batch_columns): - raise ValueError("experimental_handoff batch columns must be unique") - handoff_cell_selection = core_artifact_reference(experimental_handoff.cellSelection) - if not isinstance(handoff_cell_selection, ArtifactRef): - raise ValueError("experimental_handoff lacks an exact cell selection") - if handoff_cell_selection != normalized_cell_selection: - raise ValueError("normalized selection conflicts with experimental_handoff") - if resolved_batch_columns and sorted(resolved_batch_columns) != ( - canonical_batch_columns - ): - raise ValueError("batch_columns conflict with experimental_handoff") - if resolved_preservation_columns and resolved_preservation_columns != list( - experimental_handoff.preservationColumns - ): - raise ValueError("preservation_columns conflict with experimental_handoff") - if experimental_handoff.batchAction == "needsInput": - raise ValueError("Experimental Context requires input before tuning") - if experimental_handoff.batchAction == "skip" and experimental_handoff.batchColumns: - raise ValueError("A skip handoff must not contain batch columns") - if experimental_handoff.batchAction == "evaluateHarmony": - expected_coefficients = set(experimental_handoff.coefficientsOfInterest) - safe_coefficients = { - item.coefficient - for item in experimental_handoff.batchSafety - if item.status == "safe" and item.batchColumns == canonical_batch_columns - } - if ( - not expected_coefficients - or not canonical_batch_columns - or safe_coefficients != expected_coefficients - ): - raise ValueError( - "Harmony handoff lacks safe evidence for every coefficient" - ) - if experimental_handoff.batchAction == "unsafe": - expected_coefficients = set(experimental_handoff.coefficientsOfInterest) - exact_safety = [ - item - for item in experimental_handoff.batchSafety - if item.batchColumns == canonical_batch_columns - and item.coefficient in expected_coefficients - ] - if ( - not expected_coefficients - or {item.coefficient for item in exact_safety} != expected_coefficients - or any(item.status == "notComputed" for item in exact_safety) - or not any(item.status == "unsafe" for item in exact_safety) - ): - raise ValueError("Unsafe handoff lacks exact unsafe batch evidence") - if any( - item.evidenceId not in experimental_handoff.evidenceIds - for item in experimental_handoff.batchSafety - ): - raise ValueError("Experimental handoff does not cite its batch evidence") - return ( - normalized_cell_selection, - canonical_batch_columns, - list(experimental_handoff.preservationColumns), - ) - - -def prepare_parameter_tuning_dependencies( - store: Any, - *, - normalized: ArtifactRef, - candidates: Sequence[ParameterCandidate] | None = None, - batch_columns: Sequence[str] = (), - preservation_columns: Sequence[str] = (), - experimental_handoff: ExperimentalTuningHandoff | None = None, - max_candidates: int = 5, - max_refined_candidates: int = 0, - allow_harmony_refinement: bool = True, - min_cluster_cells: int = 20, - identity_feature_limit: int = 64, -) -> tuple[ParameterTuningDependencies, list[str]]: - """Validate one assay request and construct branch-safe dependencies.""" - - if max_candidates < 1: - raise ValueError("max_candidates must be at least one") - if max_refined_candidates < 0: - raise ValueError("max_refined_candidates must be non-negative") - if min_cluster_cells < 1: - raise ValueError("min_cluster_cells must be at least one") - if identity_feature_limit < 2: - raise ValueError("identity_feature_limit must be at least two") - normalized = core_artifact_reference(normalized) - if not isinstance(normalized, ArtifactRef) or normalized.kind != "normalized": - raise TypeError("normalized must be a normalized ArtifactRef") - if normalized.assay is None: - raise ValueError("normalized artifact has no assay") - normalized_status = store.inspect_artifact(normalized) - if not getattr(normalized_status, "exists", True): - raise ValueError("normalized artifact does not exist") - if not getattr(normalized_status, "complete", False): - raise ValueError("normalized artifact is incomplete") - raw_cell_selection = (getattr(normalized_status, "inputs", None) or {}).get( - "cell_selection" - ) - if not isinstance(raw_cell_selection, Mapping): - raise ValueError("normalized artifact has no cell-selection input") - normalized_cell_selection = ArtifactRef.from_dict(dict(raw_cell_selection)) - if ( - normalized_cell_selection.scope != "datastore" - or normalized_cell_selection.kind != "cell_selection" - or normalized_cell_selection.assay is not None - ): - raise ValueError("normalized artifact has an invalid cell-selection input") - from_assay = normalized.assay - ( - resolved_cell_selection, - resolved_batch_columns, - resolved_preservation_columns, - ) = _resolve_experimental_tuning_handoff( - normalized_cell_selection=normalized_cell_selection, - batch_columns=batch_columns, - preservation_columns=preservation_columns, - experimental_handoff=experimental_handoff, - ) - if len(set(resolved_batch_columns)) != len(resolved_batch_columns): - raise ValueError("batch_columns must be unique") - seed_candidates = ( - get_default_parameter_candidates() if candidates is None else list(candidates) - ) - if not seed_candidates: - raise ValueError("candidates must be non-empty") - if len(seed_candidates) > max_candidates: - raise ValueError( - f"Initial candidate count exceeds max_candidates={max_candidates}" - ) - pair_harmony = ( - experimental_handoff is not None - and experimental_handoff.batchAction == "evaluateHarmony" - ) - candidate_values = build_initial_parameter_candidates( - seed_candidates, - pair_harmony=pair_harmony, - ) - if len(candidate_values) + max_refined_candidates > CONFIG._MAX_CANDIDATES_OFFERED: - raise ValueError( - "Initial and refined candidates may contain at most " - f"{CONFIG._MAX_CANDIDATES_OFFERED} values" - ) - candidate_map: dict[str, ParameterCandidate] = {} - for candidate in candidate_values: - if not CONFIG._CANDIDATE_ID.fullmatch(candidate.candidateId): - raise ValueError( - "candidateId must contain only ASCII letters, numbers, and underscores" - ) - if candidate.candidateId in candidate_map: - raise ValueError(f"Duplicate candidateId {candidate.candidateId!r}") - if candidate.useHarmony and not resolved_batch_columns: - raise ValueError( - f"Candidate {candidate.candidateId!r} requires batch_columns" - ) - if ( - candidate.useHarmony - and experimental_handoff is not None - and experimental_handoff.batchAction != "evaluateHarmony" - ): - raise ValueError( - f"Candidate {candidate.candidateId!r} is not authorized for Harmony" - ) - candidate_map[candidate.candidateId] = candidate - harmony_authorized = ( - allow_harmony_refinement - and bool(resolved_batch_columns) - and ( - experimental_handoff is None - or experimental_handoff.batchAction == "evaluateHarmony" - ) - ) - normalized_shape = normalized_artifact_shape(store, normalized) - deps = ParameterTuningDependencies( - store=store, - normalized=normalized, - cellSelection=resolved_cell_selection, - normalizedShape=normalized_shape, - fromAssay=from_assay, - candidates=candidate_map, - candidatePhases={candidate_id: "initial" for candidate_id in candidate_map}, - batchColumns=tuple(resolved_batch_columns), - preservationColumns=tuple(resolved_preservation_columns), - harmonyAuthorized=harmony_authorized, - maxCandidates=len(candidate_values) + max_refined_candidates, - minClusterCells=min_cluster_cells, - identityFeatureLimit=identity_feature_limit, - ) - return deps, list(candidate_map) - - -class ParameterTuningAgent: - """Run bounded tuning over caller-authorized Scarf candidates.""" - - def __init__( - self, - model: Any, - *, - config: AgentRunConfig | None = None, - ) -> None: - self.model = model - self.config = (config or AgentRunConfig()).with_limits( - request_limit=6, - tool_call_limit=5, - output_token_limit=32768, - timeout_seconds=600.0, - ) - - def run( - self, - store: Any, - *, - normalized: Any, - candidates: Sequence[ParameterCandidate] | None = None, - batch_columns: Sequence[str] = (), - preservation_columns: Sequence[str] = (), - experimental_handoff: ExperimentalTuningHandoff | None = None, - max_candidates: int = 5, - max_refined_candidates: int = 0, - min_cluster_cells: int = 20, - identity_feature_limit: int = 64, - ) -> ParameterTuningReport: - """Run deterministic screening, optional refinement, and final selection.""" - return tune_parameters( - store, - model=self.model, - normalized=normalized, - candidates=candidates, - batch_columns=batch_columns, - preservation_columns=preservation_columns, - experimental_handoff=experimental_handoff, - max_candidates=max_candidates, - max_refined_candidates=max_refined_candidates, - min_cluster_cells=min_cluster_cells, - identity_feature_limit=identity_feature_limit, - config=self.config, - ) - - def promote( - self, - store: Any, - *, - report: ParameterTuningReport, - normalized: Any, - identity_feature_limit: int = 64, - ) -> ParameterCandidateEvaluation: - """Resolve the exact selected native branch without mutating state.""" - - return promote_parameter_candidate( - store, - report=report, - normalized=normalized, - identity_feature_limit=identity_feature_limit, - ) - - def run_batch( - self, - store: Any, - *, - assays: Sequence[ParameterTuningAssayInput], - primary_assay: str | None = None, - max_total_candidates: int = 24, - selection_directions: str = "", - ) -> ParameterTuningReport: - """Tune several assays with one planning and one selection request.""" - - return tune_parameters_batch( - store, - model=self.model, - assays=assays, - primary_assay=primary_assay, - max_total_candidates=max_total_candidates, - selection_directions=selection_directions, - config=self.config, - ) - - def select_final( - self, - *, - report: ParameterTuningReport, - integration_evaluations: Sequence[IntegrationCandidateEvaluation], - marker_assay: str, - ) -> ParameterTuningReport: - """Select native, SNN, or WNN once and attach the final branch.""" - - return select_final_parameter_graph( - model=self.model, - report=report, - integration_evaluations=integration_evaluations, - marker_assay=marker_assay, - config=self.config, - ) - - -def _execute_parameter_candidates( - deps: ParameterTuningDependencies, - candidate_ids: Sequence[str], -) -> None: - logger.info( - f"Executing {len(candidate_ids)} parameter candidate(s) for assay " - f"{deps.fromAssay!r}" - ) - for candidate_id in candidate_ids: - execute_parameter_candidate(deps, candidate_id) - - -def _register_refined_parameter_candidates( - deps: ParameterTuningDependencies, - candidates: Sequence[ParameterCandidate], -) -> None: - if candidates: - logger.info( - f"Executing {len(candidates)} refined parameter candidate(s) for " - f"assay {deps.fromAssay!r}" - ) - for candidate in candidates: - deps.candidates[candidate.candidateId] = candidate - deps.candidatePhases[candidate.candidateId] = "refined" - execute_parameter_candidate(deps, candidate.candidateId) - - -def tune_parameters_batch( - store: Any, - *, - model: Any, - assays: Sequence[ParameterTuningAssayInput], - primary_assay: str | None = None, - max_total_candidates: int = 24, - selection_directions: str = "", - config: AgentRunConfig | None = None, -) -> ParameterTuningReport: - """Execute and select modality-specific native branches in two model calls.""" - - assay_inputs = list(assays) - if not assay_inputs: - raise ValueError("assays must contain at least one tuning input") - if max_total_candidates < 1: - raise ValueError("max_total_candidates must be at least one") - planned_total = sum(item.maxCandidates for item in assay_inputs) - if planned_total > max_total_candidates: - raise ValueError( - f"Batched tuning requests {planned_total} candidate branches; " - f"the global limit is {max_total_candidates}" - ) - normalized_refs = [ - core_artifact_reference(item.normalized) for item in assay_inputs - ] - if any( - not isinstance(ref, ArtifactRef) - or ref.kind != "normalized" - or ref.assay is None - for ref in normalized_refs - ): - raise TypeError( - "Every batched tuning input requires an assay-scoped normalized artifact" - ) - assay_names = [ref.assay for ref in normalized_refs] - if len(set(assay_names)) != len(assay_names): - raise ValueError("Batched tuning assay names must be unique") - resolved_primary = primary_assay or assay_names[0] - if resolved_primary not in assay_names: - raise ValueError(f"Unknown primary assay {resolved_primary!r}") - logger.info( - f"Starting batched parameter tuning for {len(assay_names)} assay(s); " - f"primary_assay={resolved_primary!r}, " - f"candidate_limit={max_total_candidates}" - ) - - dependencies: dict[str, ParameterTuningDependencies] = {} - initial_ids: dict[str, list[str]] = {} - max_refined_by_assay: dict[str, int] = {} - for item, assay_name in zip(assay_inputs, assay_names, strict=True): - deps, candidate_ids = prepare_parameter_tuning_dependencies( - store, - normalized=item.normalized, - candidates=item.candidates or None, - batch_columns=item.batchColumns, - preservation_columns=item.preservationColumns, - experimental_handoff=item.experimentalHandoff, - max_candidates=item.maxCandidates, - max_refined_candidates=item.maxRefinedCandidates, - allow_harmony_refinement=item.allowHarmonyRefinement, - min_cluster_cells=item.minClusterCells, - identity_feature_limit=item.identityFeatureLimit, - ) - dependencies[assay_name] = deps - initial_ids[assay_name] = candidate_ids - max_refined_by_assay[assay_name] = item.maxRefinedCandidates - cell_selections = {deps.cellSelection for deps in dependencies.values()} - if len(cell_selections) != 1: - raise ValueError("Batched tuning inputs must use the same cell selection") - for assay in assay_names: - deps = dependencies[assay] - _execute_parameter_candidates(deps, initial_ids[assay]) - logger.info( - "Completed batched initial parameter screen: " - + ", ".join( - f"{assay}={len(dependencies[assay].evaluations)}" for assay in assay_names - ) - ) - - run_config = (config or AgentRunConfig()).with_limits( - request_limit=6, - tool_call_limit=5, - output_token_limit=32768, - timeout_seconds=600.0, - ) - if any(max_refined_by_assay.values()): - try: - logger.info( - "Requesting one batched parameter refinement plan for " - f"{len(assay_names)} assay(s)" - ) - planning_execution = run_agent_sync( - model=model, - output_type=ParameterTuningBatchSearchPlan, - system_prompt=parameter_batch_search_system_prompt(), - user_prompt=parameter_batch_search_prompt( - dependencies, - max_refined_by_assay, - ), - deps_type=ParameterTuningDependencies, - deps=dependencies[resolved_primary], - config=run_config, - name="parameter_batch_search_planning", - output_validator=( - lambda proposed: validate_parameter_batch_search_plan( - proposed, - dependencies, - initial_candidate_ids=initial_ids, - max_refined_by_assay=max_refined_by_assay, - ) - ), - ) - except UnexpectedModelBehavior: - logger.warning( - "Batched parameter refinement planning exhausted " - "structured-output retries; skipping optional refinement" - ) - batch_plan = ParameterTuningBatchSearchPlan( - assayPlans={ - assay: ParameterSearchPlan( - status="complete", - rationale=( - "Structured refinement planning was unavailable after " - "bounded retries; optional refinement was skipped." - ), - stoppingCriteria=[ - "Use the completed deterministic initial screen." - ], - runInfo=AgentRunInfo( - agentName="parameter_batch_search_planning_fallback" - ), - ) - for assay in assay_names - }, - runInfo=AgentRunInfo( - agentName="parameter_batch_search_planning_fallback" - ), - ) - else: - if not isinstance( - planning_execution.output, ParameterTuningBatchSearchPlan - ): - raise TypeError("Batched parameter planner returned an unexpected type") - batch_plan = validate_parameter_batch_search_plan( - planning_execution.output, - dependencies, - initial_candidate_ids=initial_ids, - max_refined_by_assay=max_refined_by_assay, - ).model_copy(update={"runInfo": planning_execution.runInfo}) - logger.info( - "Completed batched parameter refinement plan: " - + ", ".join( - f"{assay}={len(plan.candidates)}" - for assay, plan in batch_plan.assayPlans.items() - ) - ) - else: - logger.info( - "Skipping batched parameter refinement because it is not authorized" - ) - batch_plan = ParameterTuningBatchSearchPlan( - assayPlans={ - assay: ParameterSearchPlan( - status="complete", - rationale=( - "Refinement was not authorized because " - "maxRefinedCandidates is zero." - ), - stoppingCriteria=[ - "Use the completed initial screen without refinement." - ], - ) - for assay in assay_names - } - ) - for assay, plan in batch_plan.assayPlans.items(): - deps = dependencies[assay] - _register_refined_parameter_candidates(deps, plan.candidates) - - try: - logger.info( - f"Requesting batched parameter selection across " - f"{sum(len(deps.evaluations) for deps in dependencies.values())} " - "executed candidates" - ) - selection_execution = run_agent_sync( - model=model, - output_type=ParameterTuningReport, - system_prompt=parameter_batch_selection_system_prompt(), - user_prompt=parameter_batch_selection_prompt( - dependencies, - batch_plan.assayPlans, - resolved_primary, - selection_directions, - ), - deps_type=ParameterTuningDependencies, - deps=dependencies[resolved_primary], - config=run_config, - name="parameter_tuning_batch", - output_validator=lambda proposed: validate_parameter_tuning_batch_report( - proposed, - dependencies, - search_plans=batch_plan.assayPlans, - primary_assay=resolved_primary, - ), - ) - except UnexpectedModelBehavior: - logger.warning( - "Batched parameter selection exhausted structured-output retries; " - "using the conservative executor-evidence fallback" - ) - return fallback_parameter_tuning_batch_report( - dependencies, - search_plans=batch_plan.assayPlans, - primary_assay=resolved_primary, - ) - if not isinstance(selection_execution.output, ParameterTuningReport): - raise TypeError("Batched parameter tuning returned an unexpected type") - report = validate_parameter_tuning_batch_report( - selection_execution.output, - dependencies, - search_plans=batch_plan.assayPlans, - primary_assay=resolved_primary, - ) - completed_report = report.model_copy( - update={"runInfo": selection_execution.runInfo} - ) - logger.info( - f"Completed batched parameter tuning: status={completed_report.status}, " - f"assays={len(completed_report.assayReports)}, " - f"candidates={completed_report.totalCandidates}" - ) - return completed_report - - -def tune_parameters( - store: Any, - *, - model: Any, - normalized: Any, - candidates: Sequence[ParameterCandidate] | None = None, - batch_columns: Sequence[str] = (), - preservation_columns: Sequence[str] = (), - experimental_handoff: ExperimentalTuningHandoff | None = None, - max_candidates: int = 5, - max_refined_candidates: int = 0, - min_cluster_cells: int = 20, - identity_feature_limit: int = 64, - config: AgentRunConfig | None = None, -) -> ParameterTuningReport: - """Run the bounded parameter tuning agent against an existing DataStore.""" - - deps, initial_candidate_ids = prepare_parameter_tuning_dependencies( - store, - normalized=normalized, - candidates=candidates, - batch_columns=batch_columns, - preservation_columns=preservation_columns, - experimental_handoff=experimental_handoff, - max_candidates=max_candidates, - max_refined_candidates=max_refined_candidates, - min_cluster_cells=min_cluster_cells, - identity_feature_limit=identity_feature_limit, - ) - from_assay = deps.fromAssay - cell_selection = artifact_reference(deps.cellSelection) - logger.info( - f"Starting parameter tuning for assay {from_assay!r}; " - f"candidate_limit={max_candidates}, " - f"refinement_limit={max_refined_candidates}" - ) - run_config = (config or AgentRunConfig()).with_limits( - request_limit=6, - tool_call_limit=5, - output_token_limit=32768, - timeout_seconds=600.0, - ) - _execute_parameter_candidates(deps, initial_candidate_ids) - initial_evaluations = [ - deps.evaluations[candidate_id] for candidate_id in initial_candidate_ids - ] - - if max_refined_candidates == 0: - logger.info( - f"Skipping parameter refinement for assay {from_assay!r} because it " - "is not authorized" - ) - plan = ParameterSearchPlan( - status="complete", - rationale=( - "Refinement was not authorized because max_refined_candidates is zero." - ), - stoppingCriteria=[ - "Use the completed initial screen without a refinement pass." - ], - ) - else: - try: - logger.info( - f"Requesting parameter refinement plan for assay {from_assay!r} " - f"from {len(initial_evaluations)} initial evaluations" - ) - planning_execution = run_agent_sync( - model=model, - output_type=ParameterSearchPlan, - system_prompt=parameter_search_system_prompt(), - user_prompt=parameter_search_prompt( - from_assay=from_assay, - cell_selection=cell_selection, - evaluations=initial_evaluations, - batch_columns=deps.batchColumns, - preservation_columns=deps.preservationColumns, - harmony_authorized=deps.harmonyAuthorized, - max_refined_candidates=max_refined_candidates, - ), - deps_type=ParameterTuningDependencies, - deps=deps, - config=run_config, - name="parameter_search_planning", - output_validator=( - lambda proposed_plan: validate_parameter_search_plan( - proposed_plan, - deps, - initial_candidate_ids=initial_candidate_ids, - max_refined_candidates=max_refined_candidates, - ) - ), - ) - except UnexpectedModelBehavior: - logger.warning( - f"Parameter refinement planning for assay {from_assay!r} " - "exhausted structured-output retries; skipping optional refinement" - ) - plan = ParameterSearchPlan( - status="complete", - rationale=( - "Structured refinement planning was unavailable after bounded " - "retries; optional refinement was skipped." - ), - stoppingCriteria=["Use the completed deterministic initial screen."], - runInfo=AgentRunInfo(agentName="parameter_search_planning_fallback"), - ) - else: - if not isinstance(planning_execution.output, ParameterSearchPlan): - raise TypeError( - "Parameter search planner returned an unexpected output type" - ) - plan = validate_parameter_search_plan( - planning_execution.output, - deps, - initial_candidate_ids=initial_candidate_ids, - max_refined_candidates=max_refined_candidates, - ).model_copy(update={"runInfo": planning_execution.runInfo}) - logger.info( - f"Completed parameter refinement plan for assay " - f"{from_assay!r}: status={plan.status}, " - f"candidates={len(plan.candidates)}" - ) - - _register_refined_parameter_candidates(deps, plan.candidates) - - evaluations = [ - deps.evaluations[candidate_id] - for candidate_id in deps.executionOrder - if candidate_id in deps.evaluations - ] - try: - logger.info( - f"Requesting parameter selection for assay {from_assay!r} across " - f"{len(evaluations)} executed candidates" - ) - selection_execution = run_agent_sync( - model=model, - output_type=ParameterTuningReport, - system_prompt=parameter_tuning_system_prompt(min_cluster_cells), - user_prompt=parameter_tuning_prompt( - from_assay=from_assay, - cell_selection=cell_selection, - evaluations=evaluations, - batch_columns=deps.batchColumns, - preservation_columns=deps.preservationColumns, - search_plan=plan, - ), - deps_type=ParameterTuningDependencies, - deps=deps, - config=run_config, - name="parameter_tuning", - output_validator=lambda report: validate_parameter_tuning_report( - report, - deps, - search_plan=plan, - ), - ) - except UnexpectedModelBehavior: - logger.warning( - f"Parameter selection for assay {from_assay!r} exhausted " - "structured-output retries; using the conservative executor-evidence " - "fallback" - ) - return fallback_parameter_tuning_report( - deps, - search_plan=plan, - agent_name="parameter_tuning_fallback", - ) - if not isinstance(selection_execution.output, ParameterTuningReport): - raise TypeError("Parameter tuning agent returned an unexpected output type") - report = validate_parameter_tuning_report( - selection_execution.output, - deps, - search_plan=plan, - ) - completed_report = report.model_copy( - update={"runInfo": selection_execution.runInfo} - ) - logger.info( - f"Completed parameter tuning for assay {from_assay!r}: " - f"status={completed_report.status}, " - f"selected={completed_report.recommendedCandidateId!r}, " - f"candidates={len(completed_report.evaluations)}" - ) - return completed_report - - -__all__ = [ - "ArtifactRecord", - "build_initial_parameter_candidates", - "CandidateComparison", - "execute_parameter_candidate", - "final_graph_options", - "final_graph_selection_prompt", - "final_graph_selection_system_prompt", - "FinalGraphComparison", - "FinalGraphNeedsInput", - "FinalGraphSelection", - "finalize_parameter_tuning_selection", - "IntegrationCandidateEvaluation", - "IntegrationMetrics", - "normalized_artifact_shape", - "ParameterCandidate", - "ParameterCandidateEvaluation", - "ParameterMetrics", - "ParameterSearchPlan", - "ParameterTuningAssayInput", - "ParameterTuningAgent", - "ParameterTuningBatchSearchPlan", - "ParameterTuningDependencies", - "ParameterTuningNeedsInput", - "ParameterTuningReport", - "evaluate_parameter_candidate", - "get_default_parameter_candidates", - "parameter_batch_search_prompt", - "parameter_batch_search_system_prompt", - "parameter_batch_selection_prompt", - "parameter_batch_selection_system_prompt", - "parameter_search_prompt", - "parameter_search_system_prompt", - "parameter_tuning_prompt", - "parameter_tuning_system_prompt", - "prepare_parameter_tuning_dependencies", - "promote_parameter_candidate", - "run_candidate_reduction", - "select_final_parameter_graph", - "tune_parameters", - "tune_parameters_batch", - "validate_parameter_batch_search_plan", - "validate_parameter_candidate_rank", - "validate_final_graph_selection", - "validate_parameter_search_plan", - "validate_parameter_tuning_batch_report", - "validate_parameter_tuning_report", -] diff --git a/scarf/agent/parameter_tuning/__init__.py b/scarf/agent/parameter_tuning/__init__.py new file mode 100644 index 00000000..921713d6 --- /dev/null +++ b/scarf/agent/parameter_tuning/__init__.py @@ -0,0 +1,117 @@ +"""Bounded parameter tuning over explicit Scarf analysis candidates.""" + +from .agent import ( + ParameterTuningAgent, + execute_parameter_search_plan, + prepare_parameter_tuning_dependencies, + tune_parameters, + tune_parameters_batch, +) +from .contracts import ( + ArtifactRecord, + CandidateComparison, + FinalGraphComparison, + FinalGraphNeedsInput, + FinalGraphSelection, + IntegrationCandidateEvaluation, + IntegrationMetrics, + ParameterCandidate, + ParameterCandidateEvaluation, + ParameterMetrics, + ParameterSearchPlan, + ParameterTuningAssayInput, + ParameterTuningBatchSearchPlan, + ParameterTuningDependencies, + ParameterTuningNeedsInput, + ParameterTuningReport, +) +from .execution import ( + evaluate_parameter_candidate, + execute_parameter_candidate, + normalized_artifact_shape, + run_candidate_reduction, + validate_parameter_candidate_rank, +) +from .prompts import ( + build_initial_parameter_candidates, + final_graph_selection_prompt, + final_graph_selection_system_prompt, + get_default_parameter_candidates, + parameter_batch_search_prompt, + parameter_batch_search_system_prompt, + parameter_batch_selection_prompt, + parameter_batch_selection_system_prompt, + parameter_search_prompt, + parameter_search_system_prompt, + parameter_tuning_prompt, + parameter_tuning_system_prompt, +) +from .selection import ( + annotate_candidate_dominance, + final_graph_options, + finalize_parameter_tuning_selection, + harmony_acceptance_gate, + parameter_evidence_classes, + promote_parameter_candidate, + require_dominated_candidate_evidence, + select_final_parameter_graph, + validate_final_graph_selection, + validate_parameter_batch_search_plan, + validate_parameter_search_plan, + validate_parameter_tuning_batch_report, + validate_parameter_tuning_report, +) + +__all__ = [ + "annotate_candidate_dominance", + "ArtifactRecord", + "build_initial_parameter_candidates", + "CandidateComparison", + "execute_parameter_candidate", + "execute_parameter_search_plan", + "final_graph_options", + "final_graph_selection_prompt", + "final_graph_selection_system_prompt", + "FinalGraphComparison", + "FinalGraphNeedsInput", + "FinalGraphSelection", + "finalize_parameter_tuning_selection", + "IntegrationCandidateEvaluation", + "IntegrationMetrics", + "normalized_artifact_shape", + "ParameterCandidate", + "ParameterCandidateEvaluation", + "ParameterMetrics", + "ParameterSearchPlan", + "ParameterTuningAssayInput", + "ParameterTuningAgent", + "ParameterTuningBatchSearchPlan", + "ParameterTuningDependencies", + "ParameterTuningNeedsInput", + "ParameterTuningReport", + "evaluate_parameter_candidate", + "get_default_parameter_candidates", + "harmony_acceptance_gate", + "parameter_batch_search_prompt", + "parameter_batch_search_system_prompt", + "parameter_batch_selection_prompt", + "parameter_batch_selection_system_prompt", + "parameter_search_prompt", + "parameter_search_system_prompt", + "parameter_evidence_classes", + "parameter_tuning_prompt", + "parameter_tuning_system_prompt", + "prepare_parameter_tuning_dependencies", + "promote_parameter_candidate", + "run_candidate_reduction", + "require_dominated_candidate_evidence", + "select_final_parameter_graph", + "tune_parameters", + "tune_parameters_batch", + "validate_parameter_batch_search_plan", + "validate_parameter_candidate_rank", + "validate_final_graph_selection", + "validate_parameter_search_plan", + "validate_parameter_tuning_batch_report", + "validate_parameter_tuning_report", +] diff --git a/scarf/agent/parameter_tuning/agent.py b/scarf/agent/parameter_tuning/agent.py new file mode 100644 index 00000000..792ef0cb --- /dev/null +++ b/scarf/agent/parameter_tuning/agent.py @@ -0,0 +1,948 @@ +from collections.abc import Mapping, Sequence +from typing import Any + +from ...storage.refs import ArtifactRef +from ...utils.logging import logger +from .._deps import AGENT_INSTALL_HINT +from ..config import AgentRunConfig +from ..config.agent_exec import run_agent_sync +from ..tools import artifact_reference, core_artifact_reference +from ..types import AgentRunInfo, ExperimentalTuningHandoff +from .contracts import ( + _CANDIDATE_ID, + IntegrationCandidateEvaluation, + ParameterCandidate, + ParameterCandidateEvaluation, + ParameterSearchPlan, + ParameterTuningAssayInput, + ParameterTuningBatchSearchPlan, + ParameterTuningDependencies, + ParameterTuningReport, +) +from .execution import execute_parameter_candidate, normalized_artifact_shape +from .prompts import ( + build_initial_parameter_candidates, + get_default_parameter_candidates, + parameter_batch_search_prompt, + parameter_batch_search_system_prompt, + parameter_batch_selection_prompt, + parameter_batch_selection_system_prompt, + parameter_search_prompt, + parameter_search_system_prompt, + parameter_tuning_prompt, + parameter_tuning_system_prompt, +) +from .selection import ( + annotate_candidate_dominance, + pending_parameter_tuning_batch_report, + pending_parameter_tuning_report, + promote_parameter_candidate, + select_final_parameter_graph, + validate_parameter_batch_search_plan, + validate_parameter_search_plan, + validate_parameter_tuning_batch_report, + validate_parameter_tuning_report, +) + +try: + from pydantic_ai import UnexpectedModelBehavior, UsageLimitExceeded +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + + +_MAX_CANDIDATES_OFFERED = 25 + + +def _resolve_experimental_tuning_handoff( + *, + normalized_cell_selection: ArtifactRef, + batch_columns: Sequence[str], + preservation_columns: Sequence[str], + experimental_handoff: ExperimentalTuningHandoff | None, +) -> tuple[ArtifactRef, list[str], list[str]]: + resolved_batch_columns = list(batch_columns) + resolved_preservation_columns = list(preservation_columns) + if experimental_handoff is None: + return ( + normalized_cell_selection, + resolved_batch_columns, + resolved_preservation_columns, + ) + + handoff_batch_columns = list(experimental_handoff.batchColumns) + canonical_batch_columns = sorted(set(handoff_batch_columns)) + if len(canonical_batch_columns) != len(handoff_batch_columns): + raise ValueError("experimental_handoff batch columns must be unique") + handoff_cell_selection = core_artifact_reference(experimental_handoff.cellSelection) + if not isinstance(handoff_cell_selection, ArtifactRef): + raise ValueError("experimental_handoff lacks an exact cell selection") + if handoff_cell_selection != normalized_cell_selection: + raise ValueError("normalized selection conflicts with experimental_handoff") + if resolved_batch_columns and sorted(resolved_batch_columns) != ( + canonical_batch_columns + ): + raise ValueError("batch_columns conflict with experimental_handoff") + if resolved_preservation_columns and resolved_preservation_columns != list( + experimental_handoff.preservationColumns + ): + raise ValueError("preservation_columns conflict with experimental_handoff") + if experimental_handoff.batchAction == "needsInput": + raise ValueError("Experimental Context requires input before tuning") + if experimental_handoff.batchAction == "skip" and experimental_handoff.batchColumns: + raise ValueError("A skip handoff must not contain batch columns") + if experimental_handoff.batchAction == "evaluateHarmony": + expected_coefficients = set(experimental_handoff.coefficientsOfInterest) + safe_coefficients = { + item.coefficient + for item in experimental_handoff.batchSafety + if item.status == "safe" and item.batchColumns == canonical_batch_columns + } + if ( + not expected_coefficients + or not canonical_batch_columns + or safe_coefficients != expected_coefficients + ): + raise ValueError( + "Harmony handoff lacks safe evidence for every coefficient" + ) + if experimental_handoff.batchAction == "unsafe": + expected_coefficients = set(experimental_handoff.coefficientsOfInterest) + exact_safety = [ + item + for item in experimental_handoff.batchSafety + if item.batchColumns == canonical_batch_columns + and item.coefficient in expected_coefficients + ] + if ( + not expected_coefficients + or {item.coefficient for item in exact_safety} != expected_coefficients + or any(item.status == "notComputed" for item in exact_safety) + or not any(item.status == "unsafe" for item in exact_safety) + ): + raise ValueError("Unsafe handoff lacks exact unsafe batch evidence") + if any( + item.evidenceId not in experimental_handoff.evidenceIds + for item in experimental_handoff.batchSafety + ): + raise ValueError("Experimental handoff does not cite its batch evidence") + return ( + normalized_cell_selection, + canonical_batch_columns, + list(experimental_handoff.preservationColumns), + ) + + +def prepare_parameter_tuning_dependencies( + store: Any, + *, + normalized: ArtifactRef, + candidates: Sequence[ParameterCandidate] | None = None, + batch_columns: Sequence[str] = (), + preservation_columns: Sequence[str] = (), + experimental_handoff: ExperimentalTuningHandoff | None = None, + max_candidates: int = 5, + max_refined_candidates: int = 0, + allow_harmony_refinement: bool = True, + pair_harmony_candidates: bool | None = None, + min_cluster_cells: int = 20, + identity_feature_limit: int = 64, +) -> tuple[ParameterTuningDependencies, list[str]]: + """Validate one assay request and construct branch-safe dependencies.""" + + if max_candidates < 1: + raise ValueError("max_candidates must be at least one") + if max_refined_candidates < 0: + raise ValueError("max_refined_candidates must be non-negative") + if pair_harmony_candidates is not None and not isinstance( + pair_harmony_candidates, + bool, + ): + raise TypeError("pair_harmony_candidates must be a boolean or None") + if min_cluster_cells < 1: + raise ValueError("min_cluster_cells must be at least one") + if identity_feature_limit < 2: + raise ValueError("identity_feature_limit must be at least two") + normalized = core_artifact_reference(normalized) + if not isinstance(normalized, ArtifactRef) or normalized.kind != "normalized": + raise TypeError("normalized must be a normalized ArtifactRef") + if normalized.assay is None: + raise ValueError("normalized artifact has no assay") + normalized_status = store.inspect_artifact(normalized) + if not getattr(normalized_status, "exists", True): + raise ValueError("normalized artifact does not exist") + if not getattr(normalized_status, "complete", False): + raise ValueError("normalized artifact is incomplete") + raw_cell_selection = (getattr(normalized_status, "inputs", None) or {}).get( + "cell_selection" + ) + if not isinstance(raw_cell_selection, Mapping): + raise ValueError("normalized artifact has no cell-selection input") + normalized_cell_selection = ArtifactRef.from_dict(dict(raw_cell_selection)) + if ( + normalized_cell_selection.scope != "datastore" + or normalized_cell_selection.kind != "cell_selection" + or normalized_cell_selection.assay is not None + ): + raise ValueError("normalized artifact has an invalid cell-selection input") + from_assay = normalized.assay + ( + resolved_cell_selection, + resolved_batch_columns, + resolved_preservation_columns, + ) = _resolve_experimental_tuning_handoff( + normalized_cell_selection=normalized_cell_selection, + batch_columns=batch_columns, + preservation_columns=preservation_columns, + experimental_handoff=experimental_handoff, + ) + if len(set(resolved_batch_columns)) != len(resolved_batch_columns): + raise ValueError("batch_columns must be unique") + seed_candidates = ( + get_default_parameter_candidates() if candidates is None else list(candidates) + ) + if not seed_candidates: + raise ValueError("candidates must be non-empty") + if len(seed_candidates) > max_candidates: + raise ValueError( + f"Initial candidate count exceeds max_candidates={max_candidates}" + ) + pair_harmony = ( + ( + experimental_handoff is not None + and experimental_handoff.batchAction == "evaluateHarmony" + ) + if pair_harmony_candidates is None + else pair_harmony_candidates + ) + candidate_values = build_initial_parameter_candidates( + seed_candidates, + pair_harmony=pair_harmony, + ) + if len(candidate_values) + max_refined_candidates > _MAX_CANDIDATES_OFFERED: + raise ValueError( + "Initial and refined candidates may contain at most " + f"{_MAX_CANDIDATES_OFFERED} values" + ) + candidate_map: dict[str, ParameterCandidate] = {} + for candidate in candidate_values: + if not _CANDIDATE_ID.fullmatch(candidate.candidateId): + raise ValueError( + "candidateId must contain only ASCII letters, numbers, and underscores" + ) + if candidate.candidateId in candidate_map: + raise ValueError(f"Duplicate candidateId {candidate.candidateId!r}") + if candidate.useHarmony and not resolved_batch_columns: + raise ValueError( + f"Candidate {candidate.candidateId!r} requires batch_columns" + ) + if ( + candidate.useHarmony + and experimental_handoff is not None + and experimental_handoff.batchAction != "evaluateHarmony" + ): + raise ValueError( + f"Candidate {candidate.candidateId!r} is not authorized for Harmony" + ) + candidate_map[candidate.candidateId] = candidate + harmony_authorized = ( + allow_harmony_refinement + and bool(resolved_batch_columns) + and ( + experimental_handoff is None + or experimental_handoff.batchAction == "evaluateHarmony" + ) + ) + normalized_shape = normalized_artifact_shape(store, normalized) + deps = ParameterTuningDependencies( + store=store, + normalized=normalized, + cellSelection=resolved_cell_selection, + normalizedShape=normalized_shape, + fromAssay=from_assay, + candidates=candidate_map, + candidatePhases={candidate_id: "initial" for candidate_id in candidate_map}, + batchColumns=tuple(resolved_batch_columns), + preservationColumns=tuple(resolved_preservation_columns), + harmonyAuthorized=harmony_authorized, + maxCandidates=len(candidate_values) + max_refined_candidates, + minClusterCells=min_cluster_cells, + identityFeatureLimit=identity_feature_limit, + ) + return deps, list(candidate_map) + + +class ParameterTuningAgent: + """Run bounded tuning over caller-authorized Scarf candidates.""" + + def __init__( + self, + model: Any, + *, + config: AgentRunConfig | None = None, + ) -> None: + self.model = model + self.config = (config or AgentRunConfig()).with_limits( + request_limit=6, + tool_call_limit=5, + output_token_limit=32768, + timeout_seconds=600.0, + ) + + def run( + self, + store: Any, + *, + normalized: Any, + candidates: Sequence[ParameterCandidate] | None = None, + batch_columns: Sequence[str] = (), + preservation_columns: Sequence[str] = (), + experimental_handoff: ExperimentalTuningHandoff | None = None, + max_candidates: int = 5, + max_refined_candidates: int = 0, + min_cluster_cells: int = 20, + identity_feature_limit: int = 64, + ) -> ParameterTuningReport: + """Run deterministic screening, optional refinement, and final selection.""" + return tune_parameters( + store, + model=self.model, + normalized=normalized, + candidates=candidates, + batch_columns=batch_columns, + preservation_columns=preservation_columns, + experimental_handoff=experimental_handoff, + max_candidates=max_candidates, + max_refined_candidates=max_refined_candidates, + min_cluster_cells=min_cluster_cells, + identity_feature_limit=identity_feature_limit, + config=self.config, + ) + + def promote( + self, + store: Any, + *, + report: ParameterTuningReport, + normalized: Any, + identity_feature_limit: int = 64, + ) -> ParameterCandidateEvaluation: + """Resolve the exact selected native branch without mutating state.""" + + return promote_parameter_candidate( + store, + report=report, + normalized=normalized, + identity_feature_limit=identity_feature_limit, + ) + + def run_batch( + self, + store: Any, + *, + assays: Sequence[ParameterTuningAssayInput], + primary_assay: str | None = None, + max_total_candidates: int = 24, + selection_directions: str = "", + ) -> ParameterTuningReport: + """Tune several assays with one planning and one selection request.""" + + return tune_parameters_batch( + store, + model=self.model, + assays=assays, + primary_assay=primary_assay, + max_total_candidates=max_total_candidates, + selection_directions=selection_directions, + config=self.config, + ) + + def select_final( + self, + *, + report: ParameterTuningReport, + integration_evaluations: Sequence[IntegrationCandidateEvaluation], + marker_assay: str, + ) -> ParameterTuningReport: + """Select native, SNN, or WNN once and attach the final branch.""" + + return select_final_parameter_graph( + model=self.model, + report=report, + integration_evaluations=integration_evaluations, + marker_assay=marker_assay, + config=self.config, + ) + + +def _execute_parameter_candidates( + deps: ParameterTuningDependencies, + candidate_ids: Sequence[str], +) -> None: + logger.info( + f"Executing {len(candidate_ids)} parameter candidate(s) for assay " + f"{deps.fromAssay!r}" + ) + for candidate_id in candidate_ids: + execute_parameter_candidate(deps, candidate_id) + _refresh_candidate_dominance(deps) + + +def _refresh_candidate_dominance(deps: ParameterTuningDependencies) -> None: + ordered = [ + deps.evaluations[candidate_id] + for candidate_id in deps.executionOrder + if candidate_id in deps.evaluations + ] + deps.evaluations.update( + { + evaluation.candidateId: evaluation + for evaluation in annotate_candidate_dominance(ordered) + } + ) + + +def _register_refined_parameter_candidates( + deps: ParameterTuningDependencies, + candidates: Sequence[ParameterCandidate], +) -> None: + if candidates: + logger.info( + f"Executing {len(candidates)} refined parameter candidate(s) for " + f"assay {deps.fromAssay!r}" + ) + for candidate in candidates: + deps.candidates[candidate.candidateId] = candidate + deps.candidatePhases[candidate.candidateId] = "refined" + execute_parameter_candidate(deps, candidate.candidateId) + _refresh_candidate_dominance(deps) + + +def execute_parameter_search_plan( + deps: ParameterTuningDependencies, + plan: ParameterSearchPlan, + *, + initial_candidate_ids: Sequence[str], + max_refined_candidates: int, +) -> tuple[ParameterSearchPlan, tuple[ParameterCandidateEvaluation, ...]]: + """Validate and execute one already-proposed bounded refinement plan.""" + + validated = validate_parameter_search_plan( + plan, + deps, + initial_candidate_ids=initial_candidate_ids, + max_refined_candidates=max_refined_candidates, + ) + _register_refined_parameter_candidates(deps, validated.candidates) + return ( + validated, + tuple( + deps.evaluations[candidate.candidateId] + for candidate in validated.candidates + ), + ) + + +def tune_parameters_batch( + store: Any, + *, + model: Any, + assays: Sequence[ParameterTuningAssayInput], + primary_assay: str | None = None, + max_total_candidates: int = 24, + selection_directions: str = "", + config: AgentRunConfig | None = None, +) -> ParameterTuningReport: + """Execute and select modality-specific native branches in two model calls.""" + + assay_inputs = list(assays) + if not assay_inputs: + raise ValueError("assays must contain at least one tuning input") + if max_total_candidates < 1: + raise ValueError("max_total_candidates must be at least one") + planned_total = sum(item.maxCandidates for item in assay_inputs) + if planned_total > max_total_candidates: + raise ValueError( + f"Batched tuning requests {planned_total} candidate branches; " + f"the global limit is {max_total_candidates}" + ) + normalized_refs = [ + core_artifact_reference(item.normalized) for item in assay_inputs + ] + if any( + not isinstance(ref, ArtifactRef) + or ref.kind != "normalized" + or ref.assay is None + for ref in normalized_refs + ): + raise TypeError( + "Every batched tuning input requires an assay-scoped normalized artifact" + ) + assay_names = [ref.assay for ref in normalized_refs] + if len(set(assay_names)) != len(assay_names): + raise ValueError("Batched tuning assay names must be unique") + resolved_primary = primary_assay or assay_names[0] + if resolved_primary not in assay_names: + raise ValueError(f"Unknown primary assay {resolved_primary!r}") + logger.info( + f"Starting batched parameter tuning for {len(assay_names)} assay(s); " + f"primary_assay={resolved_primary!r}, " + f"candidate_limit={max_total_candidates}" + ) + + dependencies: dict[str, ParameterTuningDependencies] = {} + initial_ids: dict[str, list[str]] = {} + max_refined_by_assay: dict[str, int] = {} + for item, assay_name in zip(assay_inputs, assay_names, strict=True): + deps, candidate_ids = prepare_parameter_tuning_dependencies( + store, + normalized=item.normalized, + candidates=item.candidates or None, + batch_columns=item.batchColumns, + preservation_columns=item.preservationColumns, + experimental_handoff=item.experimentalHandoff, + max_candidates=item.maxCandidates, + max_refined_candidates=item.maxRefinedCandidates, + allow_harmony_refinement=item.allowHarmonyRefinement, + min_cluster_cells=item.minClusterCells, + identity_feature_limit=item.identityFeatureLimit, + ) + dependencies[assay_name] = deps + initial_ids[assay_name] = candidate_ids + max_refined_by_assay[assay_name] = item.maxRefinedCandidates + cell_selections = {deps.cellSelection for deps in dependencies.values()} + if len(cell_selections) != 1: + raise ValueError("Batched tuning inputs must use the same cell selection") + for assay in assay_names: + deps = dependencies[assay] + _execute_parameter_candidates(deps, initial_ids[assay]) + logger.info( + "Completed batched initial parameter screen: " + + ", ".join( + f"{assay}={len(dependencies[assay].evaluations)}" for assay in assay_names + ) + ) + + run_config = (config or AgentRunConfig()).with_limits( + request_limit=6, + tool_call_limit=5, + output_token_limit=32768, + timeout_seconds=600.0, + ) + if any(max_refined_by_assay.values()): + try: + logger.info( + "Requesting one batched parameter refinement plan for " + f"{len(assay_names)} assay(s)" + ) + planning_execution = run_agent_sync( + model=model, + output_type=ParameterTuningBatchSearchPlan, + system_prompt=parameter_batch_search_system_prompt(), + user_prompt=parameter_batch_search_prompt( + dependencies, + max_refined_by_assay, + ), + deps_type=ParameterTuningDependencies, + deps=dependencies[resolved_primary], + config=run_config, + name="parameter_batch_search_planning", + output_validator=( + lambda proposed: validate_parameter_batch_search_plan( + proposed, + dependencies, + initial_candidate_ids=initial_ids, + max_refined_by_assay=max_refined_by_assay, + ) + ), + ) + except (UnexpectedModelBehavior, UsageLimitExceeded) as exc: + logger.warning( + "Batched parameter refinement model run failed within its bounds " + f"({type(exc).__name__}); pausing without a refinement decision" + ) + failed_info = getattr( + exc, + "agent_run_info", + AgentRunInfo(agentName="parameter_batch_search_planning_needs_input"), + ) + failed_plans = { + assay: ParameterSearchPlan( + status="complete", + basedOnCandidateIds=[ + next( + ( + candidate_id + for candidate_id in initial_ids[assay] + if dependencies[assay].evaluations[candidate_id].status + == "done" + and dependencies[assay] + .evaluations[candidate_id] + .eligible + ), + initial_ids[assay][0], + ) + ], + rationale=( + "The required bounded refinement review was unavailable." + ), + evidenceIds=sorted( + { + evidence_id + for candidate_id in initial_ids[assay] + for evidence_id in dependencies[assay] + .evaluations[candidate_id] + .evidenceIds + } + ), + stoppingCriteria=[ + "Obtain a grounded refinement disposition before selection." + ], + runInfo=failed_info, + ) + for assay in assay_names + } + batch_plan = ParameterTuningBatchSearchPlan( + assayPlans=failed_plans, + runInfo=failed_info, + ) + return pending_parameter_tuning_batch_report( + dependencies, + search_plans=batch_plan.assayPlans, + primary_assay=resolved_primary, + ).model_copy(update={"runInfo": failed_info}) + else: + if not isinstance( + planning_execution.output, ParameterTuningBatchSearchPlan + ): + raise TypeError("Batched parameter planner returned an unexpected type") + validated_batch_plan = validate_parameter_batch_search_plan( + planning_execution.output, + dependencies, + initial_candidate_ids=initial_ids, + max_refined_by_assay=max_refined_by_assay, + ) + batch_plan = validated_batch_plan.model_copy( + update={ + "assayPlans": { + assay: plan.model_copy( + update={"runInfo": planning_execution.runInfo} + ) + for assay, plan in validated_batch_plan.assayPlans.items() + }, + "runInfo": planning_execution.runInfo, + } + ) + logger.info( + "Completed batched parameter refinement plan: " + + ", ".join( + f"{assay}={len(plan.candidates)}" + for assay, plan in batch_plan.assayPlans.items() + ) + ) + else: + logger.info( + "Skipping batched parameter refinement because it is not authorized" + ) + batch_plan = ParameterTuningBatchSearchPlan( + assayPlans={ + assay: ParameterSearchPlan( + status="complete", + rationale=( + "Refinement was not authorized because " + "maxRefinedCandidates is zero." + ), + stoppingCriteria=[ + "Use the completed initial screen without refinement." + ], + ) + for assay in assay_names + } + ) + for assay, plan in batch_plan.assayPlans.items(): + deps = dependencies[assay] + _register_refined_parameter_candidates(deps, plan.candidates) + + try: + logger.info( + f"Requesting batched parameter selection across " + f"{sum(len(deps.evaluations) for deps in dependencies.values())} " + "executed candidates" + ) + selection_execution = run_agent_sync( + model=model, + output_type=ParameterTuningReport, + system_prompt=parameter_batch_selection_system_prompt(), + user_prompt=parameter_batch_selection_prompt( + dependencies, + batch_plan.assayPlans, + resolved_primary, + selection_directions, + ), + deps_type=ParameterTuningDependencies, + deps=dependencies[resolved_primary], + config=run_config, + name="parameter_tuning_batch", + output_validator=lambda proposed: validate_parameter_tuning_batch_report( + proposed, + dependencies, + search_plans=batch_plan.assayPlans, + primary_assay=resolved_primary, + ), + ) + except (UnexpectedModelBehavior, UsageLimitExceeded) as exc: + logger.warning( + "Batched parameter selection model run failed within its bounds " + f"({type(exc).__name__}); " + "returning needsInput with the completed executor evidence" + ) + return pending_parameter_tuning_batch_report( + dependencies, + search_plans=batch_plan.assayPlans, + primary_assay=resolved_primary, + ).model_copy( + update={ + "runInfo": getattr( + exc, + "agent_run_info", + AgentRunInfo(agentName="parameter_tuning_batch_needs_input"), + ) + } + ) + if not isinstance(selection_execution.output, ParameterTuningReport): + raise TypeError("Batched parameter tuning returned an unexpected type") + report = validate_parameter_tuning_batch_report( + selection_execution.output, + dependencies, + search_plans=batch_plan.assayPlans, + primary_assay=resolved_primary, + ) + completed_report = report.model_copy( + update={"runInfo": selection_execution.runInfo} + ) + logger.info( + f"Completed batched parameter tuning: status={completed_report.status}, " + f"assays={len(completed_report.assayReports)}, " + f"candidates={completed_report.totalCandidates}" + ) + return completed_report + + +def tune_parameters( + store: Any, + *, + model: Any, + normalized: Any, + candidates: Sequence[ParameterCandidate] | None = None, + batch_columns: Sequence[str] = (), + preservation_columns: Sequence[str] = (), + experimental_handoff: ExperimentalTuningHandoff | None = None, + max_candidates: int = 5, + max_refined_candidates: int = 0, + min_cluster_cells: int = 20, + identity_feature_limit: int = 64, + config: AgentRunConfig | None = None, +) -> ParameterTuningReport: + """Run the bounded parameter tuning agent against an existing DataStore.""" + + deps, initial_candidate_ids = prepare_parameter_tuning_dependencies( + store, + normalized=normalized, + candidates=candidates, + batch_columns=batch_columns, + preservation_columns=preservation_columns, + experimental_handoff=experimental_handoff, + max_candidates=max_candidates, + max_refined_candidates=max_refined_candidates, + min_cluster_cells=min_cluster_cells, + identity_feature_limit=identity_feature_limit, + ) + from_assay = deps.fromAssay + cell_selection = artifact_reference(deps.cellSelection) + logger.info( + f"Starting parameter tuning for assay {from_assay!r}; " + f"candidate_limit={max_candidates}, " + f"refinement_limit={max_refined_candidates}" + ) + run_config = (config or AgentRunConfig()).with_limits( + request_limit=6, + tool_call_limit=5, + output_token_limit=32768, + timeout_seconds=600.0, + ) + _execute_parameter_candidates(deps, initial_candidate_ids) + initial_evaluations = [ + deps.evaluations[candidate_id] for candidate_id in initial_candidate_ids + ] + + if max_refined_candidates == 0: + logger.info( + f"Skipping parameter refinement for assay {from_assay!r} because it " + "is not authorized" + ) + plan = ParameterSearchPlan( + status="complete", + rationale=( + "Refinement was not authorized because max_refined_candidates is zero." + ), + stoppingCriteria=[ + "Use the completed initial screen without a refinement pass." + ], + ) + else: + try: + logger.info( + f"Requesting parameter refinement plan for assay {from_assay!r} " + f"from {len(initial_evaluations)} initial evaluations" + ) + planning_execution = run_agent_sync( + model=model, + output_type=ParameterSearchPlan, + system_prompt=parameter_search_system_prompt(), + user_prompt=parameter_search_prompt( + from_assay=from_assay, + cell_selection=cell_selection, + evaluations=initial_evaluations, + batch_columns=deps.batchColumns, + preservation_columns=deps.preservationColumns, + harmony_authorized=deps.harmonyAuthorized, + max_refined_candidates=max_refined_candidates, + ), + deps_type=ParameterTuningDependencies, + deps=deps, + config=run_config, + name="parameter_search_planning", + output_validator=( + lambda proposed_plan: validate_parameter_search_plan( + proposed_plan, + deps, + initial_candidate_ids=initial_candidate_ids, + max_refined_candidates=max_refined_candidates, + ) + ), + ) + except (UnexpectedModelBehavior, UsageLimitExceeded) as exc: + logger.warning( + f"Parameter refinement planning for assay {from_assay!r} " + f"failed within its model-run bounds ({type(exc).__name__}); " + "pausing without a refinement decision" + ) + successful_parent = next( + ( + evaluation + for evaluation in initial_evaluations + if evaluation.status == "done" and evaluation.eligible + ), + initial_evaluations[0], + ) + failed_plan = ParameterSearchPlan( + status="complete", + basedOnCandidateIds=[successful_parent.candidateId], + rationale=("The required bounded refinement review was unavailable."), + evidenceIds=sorted( + { + evidence_id + for evaluation in initial_evaluations + for evidence_id in evaluation.evidenceIds + } + ), + stoppingCriteria=[ + "Obtain a grounded refinement disposition before selection." + ], + runInfo=getattr( + exc, + "agent_run_info", + AgentRunInfo(agentName="parameter_search_planning_needs_input"), + ), + ) + return pending_parameter_tuning_report( + deps, + search_plan=failed_plan, + agent_name="parameter_search_planning_needs_input", + ).model_copy(update={"runInfo": failed_plan.runInfo}) + else: + if not isinstance(planning_execution.output, ParameterSearchPlan): + raise TypeError( + "Parameter search planner returned an unexpected output type" + ) + plan = validate_parameter_search_plan( + planning_execution.output, + deps, + initial_candidate_ids=initial_candidate_ids, + max_refined_candidates=max_refined_candidates, + ).model_copy(update={"runInfo": planning_execution.runInfo}) + logger.info( + f"Completed parameter refinement plan for assay " + f"{from_assay!r}: status={plan.status}, " + f"candidates={len(plan.candidates)}" + ) + + _register_refined_parameter_candidates(deps, plan.candidates) + + evaluations = [ + deps.evaluations[candidate_id] + for candidate_id in deps.executionOrder + if candidate_id in deps.evaluations + ] + try: + logger.info( + f"Requesting parameter selection for assay {from_assay!r} across " + f"{len(evaluations)} executed candidates" + ) + selection_execution = run_agent_sync( + model=model, + output_type=ParameterTuningReport, + system_prompt=parameter_tuning_system_prompt(min_cluster_cells), + user_prompt=parameter_tuning_prompt( + from_assay=from_assay, + cell_selection=cell_selection, + evaluations=evaluations, + batch_columns=deps.batchColumns, + preservation_columns=deps.preservationColumns, + search_plan=plan, + ), + deps_type=ParameterTuningDependencies, + deps=deps, + config=run_config, + name="parameter_tuning", + output_validator=lambda report: validate_parameter_tuning_report( + report, + deps, + search_plan=plan, + ), + ) + except (UnexpectedModelBehavior, UsageLimitExceeded) as exc: + logger.warning( + f"Parameter selection for assay {from_assay!r} failed within its " + f"model-run bounds ({type(exc).__name__}); returning needsInput " + "with the completed executor evidence" + ) + return pending_parameter_tuning_report( + deps, + search_plan=plan, + agent_name="parameter_tuning_needs_input", + ).model_copy( + update={ + "runInfo": getattr( + exc, + "agent_run_info", + AgentRunInfo(agentName="parameter_tuning_needs_input"), + ) + } + ) + if not isinstance(selection_execution.output, ParameterTuningReport): + raise TypeError("Parameter tuning agent returned an unexpected output type") + report = validate_parameter_tuning_report( + selection_execution.output, + deps, + search_plan=plan, + ) + completed_report = report.model_copy( + update={"runInfo": selection_execution.runInfo} + ) + logger.info( + f"Completed parameter tuning for assay {from_assay!r}: " + f"status={completed_report.status}, " + f"selected={completed_report.recommendedCandidateId!r}, " + f"candidates={len(completed_report.evaluations)}" + ) + return completed_report diff --git a/scarf/agent/parameter_tuning/comparisons.py b/scarf/agent/parameter_tuning/comparisons.py new file mode 100644 index 00000000..e283861e --- /dev/null +++ b/scarf/agent/parameter_tuning/comparisons.py @@ -0,0 +1,597 @@ +"""Exact RNA sensitivity comparisons and their review requirements.""" + +import json +from collections import Counter +from collections.abc import Mapping +from typing import Any, Literal, get_args + +import numpy as np +from pydantic import Field + +from ...storage.refs import ArtifactRef +from ..types import AgentDataModel +from .contracts import ParameterCandidateEvaluation + +type ComparisonAxis = Literal[ + "hvgCount", "hvgRanking", "featurePolicy", "pca", "neighbors", "partition" +] + + +class ComparisonConclusion(AgentDataModel): + """Explain an observed comparison against the scientific objective.""" + + axis: ComparisonAxis + candidateIds: list[str] = Field(min_length=1) + preferredCandidateId: str + quantitativeReason: str = Field(min_length=1) + biologicalReason: str = Field(min_length=1) + plainLanguageSummary: str = Field(min_length=1) + tradeoffs: list["ComparisonTradeoff"] = Field(default_factory=list) + + +class ComparisonTradeoff(AgentDataModel): + """A measured advantage of an alternative that the preference must explain.""" + + alternativeCandidateId: str + metric: Literal[ + "seedStability", + "subsampleStability", + "markerCoherence", + "markerSpecificityMedian", + "macroF1", + ] + preferredValue: float = Field(allow_inf_nan=False) + alternativeValue: float = Field(allow_inf_nan=False) + interpretation: str = Field(min_length=1) + + +def comparison_advantages(coverage: Mapping[str, Any]) -> list[dict[str, Any]]: + """Enumerate exact measured counterevidence for each observed preference.""" + settings = coverage["candidateSettings"] + axes: dict[str, set[str]] = {} + for comparison in coverage["comparisons"]: + identifiers = axes.setdefault(comparison["axis"], set()) + identifiers.add(comparison["baselineCandidateId"]) + if comparison["status"] == "completed": + identifiers.add(comparison["alternativeCandidateId"]) + rows = [] + metrics = get_args(ComparisonTradeoff.model_fields["metric"].annotation) + for axis, identifiers in sorted(axes.items()): + for preferred_id in sorted(identifiers): + preferred = settings[preferred_id]["metrics"] + for alternative_id in sorted(identifiers - {preferred_id}): + alternative = settings[alternative_id]["metrics"] + for metric in metrics: + left, right = preferred.get(metric), alternative.get(metric) + if ( + isinstance(left, (int, float)) + and not isinstance(left, bool) + and isinstance(right, (int, float)) + and not isinstance(right, bool) + and np.isfinite(left) + and np.isfinite(right) + and right > left + ): + rows.append( + { + "axis": axis, + "preferredCandidateId": preferred_id, + "alternativeCandidateId": alternative_id, + "metric": metric, + "preferredValue": float(left), + "alternativeValue": float(right), + "difference": float(right - left), + } + ) + return rows + + +def bind_comparison_measurements( + coverage: Mapping[str, Any], action: Mapping[str, Any] +) -> dict[str, Any]: + """Attach measured values to model-authored interpretations, never rationales.""" + inventory = { + ( + row["axis"], + row["preferredCandidateId"], + row["alternativeCandidateId"], + row["metric"], + ): row + for row in comparison_advantages(coverage) + } + conclusions = [] + errors = [] + permitted = {} + for conclusion in action.get("comparisonConclusions", []): + preference = (conclusion["axis"], conclusion["preferredCandidateId"]) + required = {key: row for key, row in inventory.items() if key[:2] == preference} + permitted.update(required) + prefix = ( + f"comparisonConclusions[axis={preference[0]!r}, " + f"preferredCandidateId={preference[1]!r}].tradeoffs" + ) + tradeoffs = [] + supplied = set() + for interpretation in conclusion.get("tradeoffs", []): + key = ( + *preference, + interpretation["alternativeCandidateId"], + interpretation["metric"], + ) + if key in supplied: + errors.append(f"{prefix}: duplicate tradeoff entry {key[2:]!r}") + supplied.add(key) + measured = inventory.get(key) + if measured is None: + settings = coverage["candidateSettings"] + left = settings.get(preference[1], {}).get("metrics", {}).get(key[3]) + right = settings.get(key[2], {}).get("metrics", {}).get(key[3]) + errors.append( + f"{prefix}: {key[2:]!r} does not identify an observed advantage " + f"(preferred={left!r}, alternative={right!r}). " + "Only permittedTradeoffs rows belong here; comparisons favoring " + "the preference and ties belong in quantitativeReason." + ) + continue + values = { + field: measured[field] + for field in ("preferredValue", "alternativeValue") + } + if any( + interpretation.get(field) is not None and interpretation[field] != value + for field, value in values.items() + ): + errors.append( + f"{prefix}: use exact preferred and alternative measurements " + f"for {key[2:]!r}: {values!r}" + ) + tradeoffs.append({**interpretation, **values}) + if action.get("action") in {"accept", "combine"}: + for key in sorted(required.keys() - supplied): + errors.append( + f"{prefix}: add alternativeCandidateId={key[2]!r}, " + f"metric={key[3]!r} and an interpretation of its measured advantage" + ) + conclusions.append({**conclusion, "tradeoffs": tradeoffs}) + if errors: + raise ValueError( + "Repair all comparisonConclusions tradeoffs errors together. " + "permittedTradeoffs contains the exact allowed rows for your current " + "axis and preferredCandidateId choices. For accept or combine, explain " + "every matching row using alternativeCandidateId, metric and interpretation. " + "Do not include other comparisons, worse alternatives or ties in tradeoffs; " + "put general comparisons in quantitativeReason. Expanding quantitativeReason " + "or biologicalReason does not replace required entries. Scarf supplies " + "preferredValue and alternativeValue; do not transcribe them. " + "Preserve unaffected choices and explanations. If you reconsider a " + "preference, update its conclusion and combinedSettings consistently and " + "use comparisonAdvantages for the new preference. Required repairs: " + + json.dumps( + {"errors": errors, "permittedTradeoffs": list(permitted.values())}, + sort_keys=True, + ) + ) + return {**action, "comparisonConclusions": conclusions} + + +class PopulationConcern(AgentDataModel): + """Keep unsupported population interpretations explicit in acceptance.""" + + candidateId: str + clusterId: str + status: Literal["nonEssentialLimitation", "unresolvedEssential"] + evidenceIds: list[str] = Field(min_length=1) + explanation: str = Field(min_length=1) + + +class CombinedSettings(AgentDataModel): + """Propose one combination using settings from actual completed candidates.""" + + hvgCountCandidateId: str + hvgRankingCandidateId: str + featurePolicyCandidateId: str + pcaCandidateId: str + neighborsCandidateId: str + + +_REQUIRED_COMPARISONS = { + "defaultResolution:0.5", + "defaultResolution:0.75", + "defaultResolution:1.25", + "hvgCount:2000", + "hvgCount:4000", + "dimensions:10", + "dimensions:30", + "neighborsK:21", + "neighborsK:41", + "hvgRanking", + "featurePolicy", +} +_CHOICE_AXES = { + "hvgCountCandidateId": "hvgCount", + "hvgRankingCandidateId": "hvgRanking", + "featurePolicyCandidateId": "featurePolicy", + "pcaCandidateId": "pca", + "neighborsCandidateId": "neighbors", +} + + +def partition_comparison_evidence( + store: Any, + left: ParameterCandidateEvaluation, + right: ParameterCandidateEvaluation, +) -> dict[str, Any]: + """Describe bounded observed splits and merges on exactly matched cells.""" + if left.cellSelection is None or left.cellSelection != right.cellSelection: + raise ValueError("Partition comparisons require the same frozen cells") + arrays = [] + for candidate in (left, right): + reference = candidate.artifacts["clusters"] + ref = ArtifactRef( + scope=reference.scope, + assay=reference.assay, + kind=reference.kind, + artifact_id=reference.artifactId, + ) + status = store.inspect_artifact(ref) + cells = left.cellSelection + cell_ref = ArtifactRef( + scope=cells.scope, + assay=cells.assay, + kind=cells.kind, + artifact_id=cells.artifactId, + ) + if ( + not status.complete + or status.inputs.get("cell_selection") != cell_ref.to_dict() + ): + raise ValueError( + "Partition evidence does not bind its exact selected cells" + ) + arrays.append(store.load_artifact(ref)["values"]) + if len(arrays[0].shape) != 1 or arrays[0].shape != arrays[1].shape: + raise ValueError("Partition labels do not align") + counts: Counter[tuple[str, str]] = Counter() + for start in range(0, arrays[0].shape[0], 65_536): + a, b = ( + np.asarray(values[start : start + 65_536]).astype(str) for values in arrays + ) + pairs, numbers = np.unique(np.column_stack((a, b)), axis=0, return_counts=True) + counts.update( + { + (str(pair[0]), str(pair[1])): int(number) + for pair, number in zip(pairs, numbers, strict=True) + } + ) + + def summarize(reverse: bool) -> list[dict[str, Any]]: + by_source: dict[str, Counter[str]] = {} + for (a, b), count in counts.items(): + source, target = (b, a) if reverse else (a, b) + by_source.setdefault(source, Counter())[target] += count + source_metrics = right.metrics if reverse else left.metrics + target_metrics = left.metrics if reverse else right.metrics + rows: list[dict[str, Any]] = [] + for source, targets in by_source.items(): + total = sum(targets.values()) + ordered = targets.most_common() + rows.append( + { + "cluster": source, + "cells": total, + "fractionOutsideLargestMatch": 1 - ordered[0][1] / total, + "markerGenes": source_metrics.topMarkerGenes.get(source, [])[:5], + "matches": [ + { + "cluster": target, + "cells": number, + "fractionOfSource": number / total, + "markerGenes": target_metrics.topMarkerGenes.get( + target, [] + )[:5], + } + for target, number in ordered[:3] + ], + "omittedMatches": max(0, len(ordered) - 3), + "omittedCells": sum(number for _, number in ordered[3:]), + } + ) + return sorted( + rows, key=lambda row: (-row["fractionOutsideLargestMatch"], row["cells"]) + )[:5] + + return { + "matchedCells": arrays[0].shape[0], + "splits": summarize(False), + "merges": summarize(True), + "interpretation": "Observed same-cell partition overlap, with up to five split and merge examples and three matches per example. Marker names are bounded summaries; this is not independent stability or a validated cell-type identity.", + } + + +def setting_changes( + left: Mapping[str, Any], right: Mapping[str, Any] +) -> dict[str, dict[str, Any]]: + """Describe logical interventions, allowing the resulting gene set to change.""" + changes = {} + if left["hvgCount"] != right["hvgCount"] and left["features"] == right["features"]: + raise ValueError( + "One frozen feature artifact cannot have different selected-gene counts" + ) + for field, axis in ( + ("dimensions", "pca"), + ("neighborsK", "neighbors"), + ("leidenResolution", "partition"), + ("useHarmony", "correction"), + ): + a, b = left["parameters"][field], right["parameters"][field] + if a != b: + changes[axis] = {"current": a, "alternative": b} + for field, axis in ( + ("hvgCount", "hvgCount"), + ("eligibleFeatures", "featurePolicy"), + ): + if left[field] != right[field]: + changes[axis] = {"current": left[field], "alternative": right[field]} + a = (left["ranking"], left["rankingColumn"]) + b = (right["ranking"], right["rankingColumn"]) + if a != b: + changes["hvgRanking"] = {"current": list(a), "alternative": list(b)} + if left["features"] != right["features"] and not { + "hvgCount", + "hvgRanking", + "featurePolicy", + }.intersection(changes): + changes["unexplainedFeatures"] = { + "current": left["features"], + "alternative": right["features"], + } + return changes + + +def validate_comparison_review( + coverage: Mapping[str, Any], action: Mapping[str, Any] +) -> None: + """Validate saved or new conclusions against exact completed comparison inputs.""" + required = { + "phase", + "population", + "comparisons", + "candidateSettings", + "combinedCandidateId", + "resolutionCandidateIds", + } + if not required.issubset(coverage): + raise ValueError( + "Saved analysis lacks mandatory comparison coverage; start a new workflow" + ) + if coverage["phase"] not in {"sensitivity", "combined", "validation"}: + raise ValueError("Unknown RNA comparison phase") + if coverage["population"] not in {"subset", "allCells"}: + raise ValueError("Comparison population must identify sampled or all cells") + settings = coverage["candidateSettings"] + rows = coverage["comparisons"] + if not isinstance(settings, Mapping) or not isinstance(rows, list): + raise ValueError("Comparison coverage must contain exact settings and rows") + if len({row["comparisonId"] for row in rows}) != len(rows): + raise ValueError("Comparison IDs must be unique") + if not _REQUIRED_COMPARISONS.issubset({row["comparisonId"] for row in rows}): + raise ValueError("Required RNA sensitivity comparisons are missing") + axis_candidates: dict[str, set[str]] = {} + for row in rows: + axis = row["axis"] + left_id, right_id = row["baselineCandidateId"], row["alternativeCandidateId"] + if left_id not in settings or settings[left_id]["status"] != "done": + raise ValueError("A comparison baseline is not completed evidence") + axis_candidates.setdefault(axis, set()).add(left_id) + if row["status"] == "pending" and axis == "featurePolicy": + if action["action"] in {"combine", "accept"}: + raise ValueError( + "A feature-policy comparison still needs an evidence-based nomination and execution" + ) + continue + if row["status"] == "notApplicable": + if not str(row.get("reason", "")).strip(): + raise ValueError( + "An unavailable comparison needs its observed eligibility reason" + ) + if right_id is not None and right_id != left_id: + raise ValueError( + "An unavailable comparison cannot claim a different evaluated alternative" + ) + proof = row.get("observedProof", {}) + baseline = settings[left_id] + kind = proof.get("kind") + if proof.get("baselineFeatures") != baseline["features"]: + raise ValueError( + "Unavailable comparisons require observed proof bound to the exact baseline features" + ) + valid = False + if axis in {"pca", "neighbors"}: + requested = int(row["comparisonId"].split(":")[1]) + field = "dimensions" if axis == "pca" else "neighborsK" + bound = ( + min(baseline["hvgCount"], baseline["nCells"]) + if axis == "pca" + else baseline["nCells"] + ) + valid = kind == "numericalBound" and ( + requested >= bound or requested == baseline["parameters"][field] + ) + elif kind == "identicalSelectedGenes": + alternative = proof.get("alternativeSetting", {}) + differences = setting_changes( + baseline, {**alternative, "features": baseline["features"]} + ) + valid = ( + set(differences) <= {axis} + and alternative["hvgCount"] == baseline["hvgCount"] + and proof.get("verifiedEqualMasks") is True + ) + elif axis == "hvgRanking" and kind == "insufficientTechnicalGroups": + groups = proof.get("eligibleGroupsByColumn") + valid = isinstance(groups, dict) and all( + isinstance(count, int) and count < 2 for count in groups.values() + ) + elif axis == "featurePolicy" and kind == "noPermittedPolicy": + valid = proof.get( + "meaningfulPermittedInterventions" + ) == 0 and isinstance(proof.get("registeredFamilies"), list) + elif axis == "featurePolicy" and kind == "insufficientEligibleGenes": + valid = ( + isinstance(proof.get("eligibleFeatureCount"), int) + and proof["eligibleFeatureCount"] < baseline["hvgCount"] + ) + if not valid: + raise ValueError( + "An unavailable comparison needs a valid observed equivalence or infeasibility proof for its exact axis" + ) + continue + if row["status"] != "completed" or right_id not in settings: + raise ValueError("A required comparison lacks a completed alternative") + left, right = settings[left_id], settings[right_id] + if right["status"] != "done" or left["cellSelection"] != right["cellSelection"]: + raise ValueError( + "Comparison candidates must be complete on the exact same cells" + ) + differences = setting_changes(left, right) + if set(differences) != {axis}: + raise ValueError( + "A sensitivity comparison must change only its declared setting" + ) + axis_candidates[axis].add(right_id) + if ( + "comparisonConclusions" not in action + or not str(action.get("plainLanguageSummary", "")).strip() + ): + raise ValueError( + "Analysis review lacks explicit comparison conclusions and reader summary" + ) + conclusions = [ + ComparisonConclusion.model_validate(row) + for row in action["comparisonConclusions"] + ] + if action["action"] in {"accept", "combine"}: + by_axis: dict[str, ComparisonConclusion] = { + row.axis: row for row in conclusions + } + if len(by_axis) != len(conclusions) or set(by_axis) != set(axis_candidates): + raise ValueError( + "Conclude every required comparison axis before combining or accepting" + ) + for axis, ids in axis_candidates.items(): + conclusion = by_axis[axis] + if not ids.issubset(conclusion.candidateIds) or not set( + conclusion.candidateIds + ).issubset(settings): + raise ValueError( + "A conclusion must address the actual baseline and all its observed alternatives" + ) + if conclusion.preferredCandidateId not in ids: + raise ValueError( + "A comparison preference must name its observed candidate" + ) + bind_comparison_measurements(coverage, action) + if action["action"] == "combine": + if coverage["phase"] != "sensitivity": + raise ValueError( + "Only the sensitivity assessment may propose combined settings" + ) + proposal = CombinedSettings.model_validate(action.get("combinedSettings")) + for field, axis in _CHOICE_AXES.items(): + if getattr(proposal, field) not in axis_candidates[axis]: + raise ValueError( + "Combined settings must use values from the matching observed sensitivity axis" + ) + if getattr(proposal, field) != by_axis[axis].preferredCandidateId: + raise ValueError( + "The proposed combination must agree with its comparison conclusions" + ) + if action["action"] == "accept": + if ( + coverage["phase"] == "sensitivity" + or coverage["combinedCandidateId"] not in settings + ): + raise ValueError( + "Acceptance requires execution of the proposed combined settings" + ) + combined = settings[coverage["combinedCandidateId"]] + partition_ids = coverage["resolutionCandidateIds"] + if not isinstance(partition_ids, list) or len(partition_ids) != 4: + raise ValueError( + "Acceptance requires four resolutions on the combined representation" + ) + resolutions = set() + for identifier in partition_ids: + if identifier not in settings: + raise ValueError("A final resolution candidate is unavailable") + row = settings[identifier] + if ( + row["status"] != "done" + or row["cellSelection"] != combined["cellSelection"] + ): + raise ValueError( + "Final resolutions must be completed on the same cells" + ) + if set(setting_changes(combined, row)) - {"partition", "correction"}: + raise ValueError( + "Final resolution evidence changed the combined representation" + ) + resolutions.add(row["parameters"]["leidenResolution"]) + if resolutions != {0.5, 0.75, 1.0, 1.25}: + raise ValueError("Final resolution coverage is incomplete") + selected = settings.get(action.get("selectedCandidateId")) + if selected is None: + raise ValueError("Accepted settings lack exact completed evidence") + concerns = [ + PopulationConcern.model_validate(row) + for row in action.get("populationConcerns", []) + ] + if any(row.status == "unresolvedEssential" for row in concerns): + raise ValueError( + "An essential population interpretation remains unresolved" + ) + missing_markers = { + str(cluster) + for cluster, genes in selected["metrics"].get("topMarkerGenes", {}).items() + if not genes + } + if selected["metrics"].get("nClusters") != len( + selected["metrics"].get("topMarkerGenes", {}) + ): + raise ValueError( + "Acceptance needs marker support or explicit missing-marker evidence for every selected cluster" + ) + addressed = { + row.clusterId + for row in concerns + if row.candidateId == action["selectedCandidateId"] + } + if not missing_markers <= addressed: + raise ValueError( + "Acceptance must explicitly resolve each selected population without qualifying markers against the objective" + ) + for concern in concerns: + if concern.candidateId not in settings or not set( + concern.evidenceIds + ) <= set(action["evidenceIds"]): + raise ValueError( + "Population concerns must cite supplied candidate evidence used by the assessment" + ) + if not any( + set(setting_changes(selected, settings[identifier])) <= {"correction"} + for identifier in partition_ids + ): + repair = coverage.get("fullRepair") + if ( + coverage["phase"] != "validation" + or not isinstance(repair, Mapping) + or repair.get("selectedCandidateId") != action["selectedCandidateId"] + or repair.get("baselineCandidateId") not in settings + or len( + setting_changes(settings[repair["baselineCandidateId"]], selected) + ) + != 1 + ): + raise ValueError( + "Accepted settings were not validated by the final comparison panel or its one targeted full repair" + ) diff --git a/scarf/agent/parameter_tuning/contracts.py b/scarf/agent/parameter_tuning/contracts.py new file mode 100644 index 00000000..9284798e --- /dev/null +++ b/scarf/agent/parameter_tuning/contracts.py @@ -0,0 +1,526 @@ +import re +from threading import Lock +from typing import Any, Literal + +from .._deps import AGENT_INSTALL_HINT +from ..types import ( + AgentDataModel, + AgentRunInfo, + ArtifactReferenceModel, + ExperimentalTuningHandoff, + StageStatus, + TuningBiologyHandoff, +) + +try: + from pydantic import Field + from pydantic.json_schema import SkipJsonSchema +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + + +type CandidateStatus = Literal["done", "failed"] +type CandidatePhase = Literal["initial", "refined"] +type ParameterSearchStatus = Literal["complete", "refine"] +type TuningConfidence = Literal["low", "medium", "high"] +type ReductionMethod = Literal["pca", "lsi", "identity"] +type IntegrationMethod = Literal["snn", "wnn"] + +_CANDIDATE_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_]{0,63}$") + + +class ArtifactRecord(ArtifactReferenceModel): + """JSON-safe identity for one artifact returned by candidate execution.""" + + @classmethod + def from_ref(cls, ref: Any) -> "ArtifactRecord": + return cls( + scope=getattr(ref, "scope", "assay"), + kind=str(getattr(ref, "kind", "")), + artifactId=str(getattr(ref, "artifact_id", ref)), + assay=getattr(ref, "assay", None), + ) + + @classmethod + def get_blank(cls) -> "ArtifactRecord": + return cls() + + +class ParameterCandidate(AgentDataModel): + """One exact, caller-authorized parameter candidate.""" + + candidateId: str = Field( + default="", + description="Exact candidate id supplied to the evaluation tool", + ) + reductionMethod: ReductionMethod = "pca" + dimensions: int = Field(default=21, ge=2) + leidenResolution: float = Field(default=1.0, gt=0) + neighborsK: int = Field(default=11, ge=2) + useHarmony: bool = False + + @classmethod + def get_blank(cls) -> "ParameterCandidate": + return cls() + + +class ParameterMetrics(AgentDataModel): + """Bounded quality metrics for one candidate branch.""" + + nClusters: int | None = None + minClusterCells: int | None = None + minClusterFraction: float | None = None + graphSilhouetteMedian: float | None = None + pcaSilhouette: float | None = None + macroF1: float | None = None + weightedF1: float | None = None + membershipStrengthMean: float | None = None + membershipStrengthMedian: float | None = None + membershipStrengthP10: float | None = None + membershipStrengthByCluster: dict[str, float] = Field(default_factory=dict) + membershipStrengthSampleSize: int | None = None + clusterConnectivity: float | None = None + seedStability: float | None = None + subsampleStability: float | None = None + markerCoherence: float | None = None + markerSpecificityMedian: float | None = None + markerSpecificityByCluster: dict[str, float] = Field(default_factory=dict) + markerAucByCluster: dict[str, float] = Field(default_factory=dict) + topMarkerGenes: dict[str, list[str]] = Field(default_factory=dict) + crossUnitSupport: float | None = None + technicalAssociation: dict[str, float] = Field(default_factory=dict) + componentVariance: list[float] = Field(default_factory=list) + pcaExplainedVarianceRatio: list[float] = Field(default_factory=list) + pcaCumulativeExplainedVarianceRatio: list[float] = Field(default_factory=list) + topLoadingGenes: dict[str, list[str]] = Field(default_factory=dict) + loadingFamilyEnrichment: dict[str, float] = Field(default_factory=dict) + loadingFamilyEnrichmentByComponent: dict[str, dict[str, float]] = Field( + default_factory=dict + ) + pcaComponentAssociations: dict[str, dict[str, list[float]]] = Field( + default_factory=dict + ) + batchPcaAssociation: dict[str, float] = Field(default_factory=dict) + technicalPcaAssociation: dict[str, float] = Field(default_factory=dict) + protectedPcaAssociation: dict[str, float] = Field(default_factory=dict) + qcPcaAssociation: dict[str, float] = Field(default_factory=dict) + neighborPrefixOverlap: float | None = None + markerFamilyEnrichment: dict[str, float] = Field(default_factory=dict) + protectedMarkerFamilies: list[str] = Field(default_factory=list) + doubletHighScoreConcentration: float | None = None + doubletScoreQuantiles: dict[str, float] = Field(default_factory=dict) + doubletScoreByCapture: dict[str, dict[str, float]] = Field(default_factory=dict) + doubletCaptureCoverage: float | None = None + batchMixing: dict[str, float] = Field(default_factory=dict) + biologicalPreservation: dict[str, dict[str, float]] = Field(default_factory=dict) + paretoOptimal: bool | None = None + dominatedByCandidateIds: list[str] = Field(default_factory=list) + dominatesCandidateIds: list[str] = Field(default_factory=list) + dominanceMetrics: dict[str, list[str]] = Field(default_factory=dict) + + @classmethod + def get_blank(cls) -> "ParameterMetrics": + return cls() + + +class ParameterCandidateEvaluation(AgentDataModel): + """Execution record returned to the model for one candidate.""" + + candidateId: str = "" + phase: CandidatePhase = "initial" + harmonyBatchColumns: list[str] = Field(default_factory=list) + status: CandidateStatus = "failed" + eligible: bool = False + parameters: ParameterCandidate = Field(default_factory=ParameterCandidate.get_blank) + artifacts: dict[str, ArtifactRecord] = Field(default_factory=dict) + cellSelection: ArtifactReferenceModel | None = None + clusterColumn: str | None = None + clusterLabel: str | None = None + effectiveDimensions: int | None = None + metrics: ParameterMetrics = Field(default_factory=ParameterMetrics.get_blank) + evidenceIds: list[str] = Field(default_factory=list) + eligibilityReasons: list[str] = Field(default_factory=list) + warnings: list[str] = Field(default_factory=list) + error: str | None = None + + @classmethod + def get_blank(cls) -> "ParameterCandidateEvaluation": + return cls() + + +class IntegrationMetrics(AgentDataModel): + """Metrics that are valid for an integrated graph comparison.""" + + nClusters: int | None = None + minClusterCells: int | None = None + minClusterFraction: float | None = None + adjustedRandByAssay: dict[str, float] = Field(default_factory=dict) + normalizedMutualInformationByAssay: dict[str, float] = Field(default_factory=dict) + biologicalConnectivity: dict[str, float] = Field(default_factory=dict) + modalityWeightsValid: bool | None = None + + @classmethod + def get_blank(cls) -> "IntegrationMetrics": + return cls() + + +class IntegrationCandidateEvaluation(AgentDataModel): + """One executor-produced SNN or WNN graph and cluster evaluation.""" + + integrationId: str = "" + method: IntegrationMethod = "wnn" + assays: list[str] = Field(default_factory=list) + status: CandidateStatus = "failed" + eligible: bool = False + cellSelection: ArtifactReferenceModel | None = None + resolution: float = Field(default=1.0, gt=0) + graphArtifact: ArtifactRecord | None = None + clusterArtifact: ArtifactRecord | None = None + clusterColumn: str | None = None + metrics: IntegrationMetrics = Field(default_factory=IntegrationMetrics.get_blank) + evidenceIds: list[str] = Field(default_factory=list) + eligibilityReasons: list[str] = Field(default_factory=list) + warnings: list[str] = Field(default_factory=list) + error: str | None = None + + @classmethod + def get_blank(cls) -> "IntegrationCandidateEvaluation": + return cls() + + +class FinalGraphComparison(AgentDataModel): + """Evidence-backed comparison against one eligible final graph option.""" + + optionId: str = "" + summary: str = "" + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "FinalGraphComparison": + return cls() + + +class FinalGraphNeedsInput(AgentDataModel): + """Concrete input needed before a final graph can be selected.""" + + question: str = "" + options: list[str] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "FinalGraphNeedsInput": + return cls() + + +class FinalGraphSelection(AgentDataModel): + """Grounded choice among selected native, SNN, and WNN graph options.""" + + status: StageStatus = "needsInput" + selectedOptionId: str | None = None + graphMethod: SkipJsonSchema[Literal["native", "snn", "wnn"] | None] = None + nativeAssay: SkipJsonSchema[str | None] = None + nativeCandidateId: SkipJsonSchema[str | None] = None + integrationId: SkipJsonSchema[str | None] = None + markerAssay: SkipJsonSchema[str] = "" + confidence: TuningConfidence = "low" + rationale: str = "" + evidenceIds: list[str] = Field(default_factory=list) + comparisons: list[FinalGraphComparison] = Field(default_factory=list) + tradeoffs: list[str] = Field(default_factory=list) + limitations: list[str] = Field(default_factory=list) + needsInput: FinalGraphNeedsInput | None = None + runInfo: SkipJsonSchema[AgentRunInfo] = Field(default_factory=AgentRunInfo) + + @classmethod + def get_blank(cls) -> "FinalGraphSelection": + return cls() + + +class CandidateComparison(AgentDataModel): + """Evidence-backed comparison against one executed non-selected candidate.""" + + candidateId: str = "" + summary: str = "" + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "CandidateComparison": + return cls() + + +class ParameterSearchPlan(AgentDataModel): + """Validated proposal for one bounded refinement pass.""" + + status: ParameterSearchStatus = Field( + default="complete", + description=( + "Summary derived from candidates: refine when candidates is non-empty " + "and complete when it is empty" + ), + ) + candidates: list[ParameterCandidate] = Field( + default_factory=list, + description=( + "Bounded unexecuted refinement candidates, or an empty list when the " + "initial screen is complete" + ), + ) + basedOnCandidateIds: list[str] = Field(default_factory=list) + harmonyBatchColumns: list[str] = Field(default_factory=list) + objectives: list[str] = Field(default_factory=list) + rationale: str = "" + evidenceIds: list[str] = Field(default_factory=list) + stoppingCriteria: list[str] = Field(default_factory=list) + runInfo: SkipJsonSchema[AgentRunInfo] = Field(default_factory=AgentRunInfo) + + @classmethod + def get_blank(cls) -> "ParameterSearchPlan": + return cls() + + +class ParameterTuningBatchSearchPlan(AgentDataModel): + """One bounded refinement plan for every assay in a batched screen.""" + + assayPlans: dict[str, ParameterSearchPlan] = Field(default_factory=dict) + runInfo: SkipJsonSchema[AgentRunInfo] = Field(default_factory=AgentRunInfo) + + @classmethod + def get_blank(cls) -> "ParameterTuningBatchSearchPlan": + return cls() + + +class ParameterTuningNeedsInput(AgentDataModel): + """User input required before tuning can produce a recommendation.""" + + question: str = "" + options: list[str] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + + @classmethod + def get_blank(cls) -> "ParameterTuningNeedsInput": + return cls() + + +class ParameterTuningReport(AgentDataModel): + """Grounded recommendation over candidate branches actually executed.""" + + status: StageStatus = "failed" + fromAssay: SkipJsonSchema[str] = "" + cellSelection: SkipJsonSchema[ArtifactReferenceModel | None] = None + evaluations: SkipJsonSchema[list[ParameterCandidateEvaluation]] = Field( + default_factory=list + ) + recommendedCandidateId: str | None = None + selectedArtifacts: SkipJsonSchema[dict[str, ArtifactRecord]] = Field( + default_factory=dict + ) + confidence: TuningConfidence = "low" + rationale: str = "" + evidenceIds: list[str] = Field(default_factory=list) + comparisons: list[CandidateComparison] = Field(default_factory=list) + tradeoffs: list[str] = Field(default_factory=list) + limitations: list[str] = Field(default_factory=list) + stopReason: str = "" + needsInput: ParameterTuningNeedsInput | None = None + searchPlan: SkipJsonSchema[ParameterSearchPlan | None] = None + assayReports: dict[str, "ParameterTuningReport"] = Field(default_factory=dict) + recommendedByAssay: SkipJsonSchema[dict[str, str]] = Field(default_factory=dict) + totalCandidates: SkipJsonSchema[int] = 0 + integrationEvaluations: SkipJsonSchema[list[IntegrationCandidateEvaluation]] = ( + Field(default_factory=list) + ) + recommendedIntegrationId: SkipJsonSchema[str | None] = None + finalClusterColumn: SkipJsonSchema[str | None] = None + finalClusterArtifact: SkipJsonSchema[ArtifactRecord | None] = None + graphAssay: SkipJsonSchema[str | None] = None + markerAssay: SkipJsonSchema[str | None] = None + finalSelection: SkipJsonSchema[FinalGraphSelection | None] = None + runInfo: SkipJsonSchema[AgentRunInfo] = Field(default_factory=AgentRunInfo) + + @classmethod + def get_blank(cls) -> "ParameterTuningReport": + return cls() + + def to_biological_handoff( + self, + *, + marker_assay: str | None = None, + ) -> TuningBiologyHandoff: + """Return the exact selected clustering branch for interpretation.""" + if self.status != "done": + raise ValueError( + "Parameter Tuning must be done before creating a biology handoff" + ) + if self.finalClusterArtifact is not None: + if self.cellSelection is None: + raise ValueError("Final branch lacks an exact cell selection") + resolved_marker_assay = marker_assay or self.markerAssay + if not resolved_marker_assay: + raise ValueError( + "A marker assay is required for an integrated biology handoff" + ) + if self.finalClusterArtifact.scope == "datastore": + if self.finalClusterArtifact.assay is not None: + raise ValueError( + "A datastore-scoped cluster artifact must not name an assay" + ) + elif ( + self.graphAssay is not None + and self.finalClusterArtifact.assay != self.graphAssay + ): + raise ValueError("Final cluster artifact does not match graphAssay") + integration = next( + ( + item + for item in self.integrationEvaluations + if item.integrationId == self.recommendedIntegrationId + ), + None, + ) + if self.finalSelection is not None: + evidence_ids = self.finalSelection.evidenceIds + elif integration is not None: + evidence_ids = integration.evidenceIds + else: + prefix = f"candidate:{self.recommendedCandidateId}:" + evidence_ids = [ + evidence_id + for evidence_id in self.evidenceIds + if evidence_id.startswith(prefix) + ] + return TuningBiologyHandoff( + cellSelection=self.cellSelection, + fromAssay=self.fromAssay, + graphAssay=self.graphAssay, + markerAssay=resolved_marker_assay, + recommendedCandidateId=( + self.recommendedIntegrationId + or ( + self.finalSelection.nativeCandidateId + if self.finalSelection is not None + else None + ) + or self.recommendedCandidateId + or "final" + ), + clusterArtifact=ArtifactReferenceModel.model_validate( + self.finalClusterArtifact.model_dump() + ), + evidenceIds=sorted(evidence_ids), + ) + if self.recommendedCandidateId is None: + raise ValueError( + "Parameter Tuning must recommend a candidate before creating a " + "biology handoff" + ) + selected = next( + ( + item + for item in self.evaluations + if item.candidateId == self.recommendedCandidateId + ), + None, + ) + if selected is None or selected.status != "done" or not selected.eligible: + raise ValueError("Recommended candidate is not an eligible execution") + cluster_artifact = selected.artifacts.get("clusters") + if cluster_artifact is None or selected.cellSelection is None: + raise ValueError("Recommended candidate lacks an exact cluster artifact") + if not self.fromAssay or cluster_artifact.assay != self.fromAssay: + raise ValueError("Recommended cluster artifact does not match the assay") + prefix = f"candidate:{selected.candidateId}:" + return TuningBiologyHandoff( + cellSelection=selected.cellSelection, + fromAssay=self.fromAssay, + graphAssay=self.fromAssay, + markerAssay=marker_assay or self.markerAssay or self.fromAssay, + recommendedCandidateId=selected.candidateId, + clusterArtifact=ArtifactReferenceModel.model_validate( + cluster_artifact.model_dump() + ), + evidenceIds=sorted( + evidence_id + for evidence_id in self.evidenceIds + if evidence_id.startswith(prefix) + ), + ) + + +class ParameterTuningDependencies(AgentDataModel): + """Runtime-only state hidden from the model and shared by tuning tools.""" + + store: Any = Field(default=None, exclude=True) + normalized: Any = Field(default=None, exclude=True) + cellSelection: Any = Field(default=None, exclude=True) + normalizedShape: tuple[int, int] | None = None + fromAssay: str = "" + candidates: dict[str, ParameterCandidate] = Field(default_factory=dict) + candidatePhases: dict[str, CandidatePhase] = Field(default_factory=dict) + batchColumns: tuple[str, ...] = () + preservationColumns: tuple[str, ...] = () + protectedCombinations: tuple[tuple[str, ...], ...] = () + columnKinds: dict[str, Literal["categorical", "continuous"]] = Field( + default_factory=dict + ) + harmonyAuthorized: bool = False + maxCandidates: int = 5 + minClusterCells: int = 20 + identityFeatureLimit: int = 64 + evaluations: dict[str, ParameterCandidateEvaluation] = Field(default_factory=dict) + executionOrder: list[str] = Field(default_factory=list) + executionLock: Any = Field(default_factory=Lock, exclude=True, repr=False) + + @classmethod + def get_blank(cls) -> "ParameterTuningDependencies": + return cls() + + +class ParameterTuningAssayInput(AgentDataModel): + """One assay branch supplied to batched parameter tuning.""" + + normalized: Any = Field(default=None, exclude=True) + candidates: list[ParameterCandidate] = Field(default_factory=list) + batchColumns: list[str] = Field(default_factory=list) + preservationColumns: list[str] = Field(default_factory=list) + experimentalHandoff: ExperimentalTuningHandoff | None = None + maxCandidates: int = Field(default=5, ge=1) + maxRefinedCandidates: int = Field(default=0, ge=0) + allowHarmonyRefinement: bool = True + minClusterCells: int = Field(default=20, ge=1) + identityFeatureLimit: int = Field(default=64, ge=2) + + @classmethod + def get_blank(cls) -> "ParameterTuningAssayInput": + return cls() + + +def _default_parameter_candidates() -> list[ParameterCandidate]: + """Return a small one-factor candidate set around Scarf defaults.""" + + return [ + ParameterCandidate( + candidateId="baseline", + dimensions=21, + leidenResolution=1.0, + ), + ParameterCandidate( + candidateId="pca_15", + dimensions=15, + leidenResolution=1.0, + ), + ParameterCandidate( + candidateId="pca_30", + dimensions=30, + leidenResolution=1.0, + ), + ParameterCandidate( + candidateId="leiden_0_5", + dimensions=21, + leidenResolution=0.5, + ), + ParameterCandidate( + candidateId="leiden_1_5", + dimensions=21, + leidenResolution=1.5, + ), + ] diff --git a/scarf/agent/parameter_tuning/diagnostics.py b/scarf/agent/parameter_tuning/diagnostics.py new file mode 100644 index 00000000..96ee5779 --- /dev/null +++ b/scarf/agent/parameter_tuning/diagnostics.py @@ -0,0 +1,2212 @@ +"""Deterministic representation and partition evidence for RNA decisions.""" + +import hashlib +from collections import Counter, defaultdict +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, cast + +import numpy as np +import pandas as pd +from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score + +from ...clustering.leiden import leiden_membership +from ...metadata.rows import ( + read_metadata_missing_rows_chunkwise, + read_metadata_rows_chunkwise, +) +from ...metadata.selection import resolve_cell_aligned_artifact +from ...quality_control.cell_cycle_genes import ( + g2m_phase_genes, + g2m_phase_genes_mouse, + s_phase_genes, + s_phase_genes_mouse, +) +from ...storage.arrays import create_zarr_dataset +from ...storage.artifact_writer import ( + ArrayRequirement, + AttributeRequirement, + finish_artifact, + plan_artifact, + start_artifact, +) +from ...storage.artifacts import fingerprint_stored_arrays +from ...storage.feature_selection import read_feature_selection_indices +from ...storage.refs import ArtifactRef +from ...storage.selections import read_stored_selection_indices +from ...storage.types import as_zarr_array +from ...utils.logging import logger +from .contracts import ArtifactRecord, ParameterCandidateEvaluation +from .execution import ( + _cached_candidate_metric, + _metadata_column_fingerprint, + diagnostic_call, + diagnostic_reuse, +) +from .selection import annotate_candidate_dominance + +_PCA_DIAGNOSTIC_ARRAYS = ( + "component_variance", + "explained_variance_ratio", + "top_loading_feature_indices", + "top_loading_values", + "family_enrichment", + "covariate_association", + "adjacent_neighbor_overlap", +) +_MAX_DOUBLET_CAPTURES = 512 +SCARF_DEFAULT_DIAGNOSTIC_FAMILIES = ( + "mitochondrial", + "ribosomal", + "mitoribosomal", + "cellCycleCcn", + "hla", + "h2", + "histone", + "sexLinked", +) + + +@dataclass(frozen=True, slots=True) +class AdvisoryDoubletScores: + """Advisory score artifacts and their exact cell-axis selections.""" + + scores: tuple[ArtifactRef, ...] + cell_selections: tuple[ArtifactRef, ...] + native_graph: ArtifactRef + native_clusters: ArtifactRef + capture_values: tuple[str, ...] = () + score_summaries: tuple[dict[str, float], ...] = () + score_quantiles: Mapping[str, float] = field(default_factory=dict) + capture_coverage: float | None = None + capture_column: str | None = None + limitations: tuple[str, ...] = () + + +def restore_advisory_doublets( + evaluation: ParameterCandidateEvaluation, + *, + capture_column: str | None, +) -> AdvisoryDoubletScores: + """Restore the exact doublet inputs from an augmented phase evaluation.""" + score_keys = sorted( + (key for key in evaluation.artifacts if key.startswith("doubletScore:")), + key=lambda key: int(key.split(":", 1)[1]), + ) + if score_keys != [f"doubletScore:{index}" for index in range(len(score_keys))]: + raise ValueError("Persisted doublet score inventory is incomplete") + captures = tuple(evaluation.metrics.doubletScoreByCapture) + if len(captures) != len(score_keys): + raise ValueError("Persisted doublet capture summaries do not align") + restored = AdvisoryDoubletScores( + scores=tuple(_artifact_ref(evaluation, key) for key in score_keys), + cell_selections=tuple( + _artifact_ref(evaluation, f"doubletCellSelection:{index}") + for index in range(len(score_keys)) + ), + native_graph=_artifact_ref(evaluation, "doubletNativeGraph"), + native_clusters=_artifact_ref(evaluation, "doubletNativeClusters"), + capture_values=captures, + score_summaries=tuple( + dict(evaluation.metrics.doubletScoreByCapture[capture]) + for capture in captures + ), + score_quantiles=dict(evaluation.metrics.doubletScoreQuantiles), + capture_coverage=evaluation.metrics.doubletCaptureCoverage, + capture_column=capture_column, + limitations=tuple( + warning for warning in evaluation.warnings if "doublet" in warning.lower() + ), + ) + diagnostic_reuse("diagnostic.advisoryDoublets", "restored") + return restored + + +def _artifact_ref( + evaluation: ParameterCandidateEvaluation, + name: str, +) -> ArtifactRef: + artifact = evaluation.artifacts.get(name) + if artifact is None: + raise ValueError( + f"Candidate {evaluation.candidateId!r} lacks {name!r} evidence" + ) + return ArtifactRef( + scope=artifact.scope, + assay=artifact.assay, + kind=artifact.kind, + artifact_id=artifact.artifactId, + ) + + +def _cluster_labels(store: Any, cluster_ref: ArtifactRef) -> np.ndarray: + group = store.load_artifact(cluster_ref) + values = as_zarr_array(group["values"], name="values") + labels = np.asarray(values[:]) + if labels.ndim != 1: + raise ValueError("Cluster evidence must be a one-dimensional label vector") + return labels + + +def _selected_feature_names( + store: Any, + feature_selection: ArtifactRef, +) -> tuple[np.ndarray, np.ndarray]: + if feature_selection.assay is None: + raise ValueError("PCA feature selection must belong to one assay") + indices = read_feature_selection_indices( + store.zw, + feature_selection.assay, + feature_selection, + ).astype(np.int64, copy=False) + names = np.asarray( + store.get_assay(feature_selection.assay).feats.fetch_all("names") + ).astype(str) + return indices, names[indices] + + +def _family_mask(names: np.ndarray, family: str) -> np.ndarray | None: + upper = np.char.upper(names.astype(str)) + if family == "mitochondrial": + return np.char.startswith(upper, "MT-") + if family in {"ribosomal", "ribosomalProtein"}: + return np.asarray( + np.logical_or.reduce( + [np.char.startswith(upper, prefix) for prefix in ("RPS", "RPL")] + ), + dtype=bool, + ) + if family == "mitoribosomal": + return np.asarray( + np.logical_or.reduce( + [np.char.startswith(upper, prefix) for prefix in ("MRPS", "MRPL")] + ), + dtype=bool, + ) + if family == "cellCycleCcn": + return np.char.startswith(upper, "CCN") + if family == "cellCycle": + cycle_genes = { + *s_phase_genes, + *g2m_phase_genes, + *s_phase_genes_mouse, + *g2m_phase_genes_mouse, + } + return np.asarray( + np.char.startswith(upper, "CCN") + | np.isin(upper, [value.upper() for value in cycle_genes]), + dtype=bool, + ) + if family in {"hla", "HLA"}: + return np.char.startswith(upper, "HLA-") + if family in {"h2", "H2"}: + return np.char.startswith(upper, "H2-") + if family == "histone": + return np.char.startswith(upper, "HIST") + if family in {"sex", "sexLinked"}: + return np.isin( + upper, + [ + "XIST", + "DDX3Y", + "USP9Y", + "EIF1AY", + "KDM5D", + "SRY", + "ZFY", + "UTY", + "TMSB4Y", + "NLGN4Y", + ], + ) + if family == "hemoglobin": + return np.char.startswith(upper, "HB") + if family == "immuneReceptor": + return np.asarray( + np.logical_or.reduce( + [ + np.char.startswith(upper, prefix) + for prefix in ("IGH", "IGK", "IGL", "TRA", "TRB", "TRD", "TRG") + ] + ), + dtype=bool, + ) + if family == "stress": + return np.asarray( + np.logical_or.reduce( + [ + np.char.startswith(upper, prefix) + for prefix in ("FOS", "JUN", "HSP", "DUSP", "EGR") + ] + ), + dtype=bool, + ) + if family == "dissociation": + return np.isin( + upper, + [ + "ATF3", + "BTG1", + "BTG2", + "DUSP1", + "EGR1", + "FOS", + "FOSB", + "IER2", + "JUN", + "JUNB", + "JUND", + "ZFP36", + ], + ) + return None + + +def _component_variance(values: Any) -> np.ndarray: + if len(values.shape) != 2: + raise ValueError("PCA coordinates must be a two-dimensional array") + n_rows, n_components = values.shape + if n_rows < 1 or n_components < 1: + raise ValueError("PCA coordinates cannot be empty") + totals = np.zeros(n_components, dtype=np.float64) + totals_squared = np.zeros(n_components, dtype=np.float64) + for start in range(0, n_rows, 65_536): + block = np.asarray(values[start : start + 65_536], dtype=np.float64) + if not np.isfinite(block).all(): + raise ValueError("PCA coordinates must be finite") + totals += block.sum(axis=0) + totals_squared += np.square(block).sum(axis=0) + variance = totals_squared / n_rows - np.square(totals / n_rows) + return np.asarray(np.maximum(variance, 0.0), dtype=np.float64) + + +def _scaled_total_variance( + store: Any, + reduction_status: Any, + *, + n_rows: int, + n_features: int, + feature_selection: ArtifactRef, +) -> float: + def input_ref(value: Any, label: str) -> ArtifactRef: + if isinstance(value, ArtifactRef): + return value + if isinstance(value, Mapping): + return ArtifactRef.from_dict(dict(value)) + raise ValueError(f"Candidate PCA lacks its {label} input") + + inputs = getattr(reduction_status, "inputs", None) or {} + normalized_ref = input_ref(inputs.get("normalized"), "normalized matrix") + pca_cell_selection = input_ref( + inputs.get("pca_cell_selection"), + "PCA cell-selection", + ) + normalized_status = store.inspect_artifact(normalized_ref) + normalized_inputs = getattr(normalized_status, "inputs", None) or {} + normalized_cell_selection = input_ref( + normalized_inputs.get("cell_selection"), + "normalized cell-selection", + ) + normalized_feature_selection = input_ref( + normalized_inputs.get("feature_selection"), + "normalized feature-selection", + ) + if pca_cell_selection != normalized_cell_selection: + raise ValueError( + "Explained-variance ratios require PCA fitted on all normalized cells" + ) + if normalized_feature_selection != feature_selection: + raise ValueError("PCA loading genes do not match the normalized features") + normalized_group = store.load_artifact(normalized_ref) + normalized = as_zarr_array(normalized_group["data"], name="data") + if normalized.shape != (n_rows, n_features): + raise ValueError("Candidate PCA and normalized matrix shapes do not align") + if "feature_sum" in normalized_group and "feature_squared_sum" in normalized_group: + totals = np.asarray( + as_zarr_array(normalized_group["feature_sum"], name="feature_sum")[:], + dtype=np.float64, + ) + totals_squared = np.asarray( + as_zarr_array( + normalized_group["feature_squared_sum"], + name="feature_squared_sum", + )[:], + dtype=np.float64, + ) + else: + totals = np.zeros(n_features, dtype=np.float64) + totals_squared = np.zeros(n_features, dtype=np.float64) + for start in range(0, n_rows, 8192): + block = np.asarray( + normalized[start : start + 8192], + dtype=np.float64, + ) + if not np.isfinite(block).all(): + raise ValueError("Normalized PCA input must be finite") + totals += block.sum(axis=0) + totals_squared += np.square(block).sum(axis=0) + if totals.shape != (n_features,) or totals_squared.shape != (n_features,): + raise ValueError("Normalized PCA feature summaries do not align") + variance = np.maximum( + totals_squared / n_rows - np.square(totals / n_rows), + 0.0, + ) + total_scaled_variance = float(np.count_nonzero(variance)) + if total_scaled_variance <= 0: + raise ValueError("Feature-scaled PCA input has no non-constant features") + return total_scaled_variance + + +def _top_loadings( + loadings: Any, + selected_indices: np.ndarray, + family_masks: Mapping[str, np.ndarray], +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + if len(loadings.shape) != 2 or loadings.shape[0] != len(selected_indices): + raise ValueError("PCA loadings do not align with selected features") + if loadings.shape[0] < 1 or loadings.shape[1] < 1: + raise ValueError("PCA loadings cannot be empty") + if any(mask.shape != selected_indices.shape for mask in family_masks.values()): + raise ValueError("PCA family masks must align with selected features") + top_n = min(20, loadings.shape[0]) + top_rows = np.empty((loadings.shape[1], 0), dtype=np.int64) + top_values = np.empty((loadings.shape[1], 0), dtype=np.float64) + for start in range(0, loadings.shape[0], 8192): + block = np.abs(np.asarray(loadings[start : start + 8192], dtype=np.float64)) + if not np.isfinite(block).all(): + raise ValueError("PCA loadings must be finite") + block_rows = np.arange(start, start + len(block), dtype=np.int64) + retained_rows = np.empty( + (loadings.shape[1], min(top_n, top_rows.shape[1] + len(block))), + dtype=np.int64, + ) + retained_values = np.empty(retained_rows.shape, dtype=np.float64) + for component in range(loadings.shape[1]): + candidate_rows = np.concatenate((top_rows[component], block_rows)) + candidate_values = np.concatenate( + (top_values[component], block[:, component]) + ) + order = np.lexsort((candidate_rows, -candidate_values))[:top_n] + retained_rows[component] = candidate_rows[order] + retained_values[component] = candidate_values[order] + top_rows = retained_rows + top_values = retained_values + top_indices = selected_indices[top_rows] + enrichment = np.zeros( + (len(family_masks), loadings.shape[1]), + dtype=np.float64, + ) + for family_index, mask in enumerate(family_masks.values()): + for component in range(loadings.shape[1]): + background = float(mask.mean()) + enrichment[family_index, component] = ( + float(mask[top_rows[component]].mean()) / background + if background > 0 + else 0.0 + ) + return top_indices, top_values, enrichment + + +def _aligned_metadata_values( + store: Any, + cell_selection: ArtifactRef, + column: str, +) -> np.ndarray: + indices = read_stored_selection_indices( + store.zw, + cell_selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + values = np.asarray(read_metadata_rows_chunkwise(store.cells, column, indices)) + if values.shape != (len(indices),): + raise ValueError(f"Metadata column {column!r} does not align with PCA") + missing = read_metadata_missing_rows_chunkwise(store.cells, column, indices) + if missing is not None and np.any(missing): + values = values.astype(object) + values[missing] = None + return values + + +def _numeric_association(coordinates: Any, values: np.ndarray) -> np.ndarray: + numeric = pd.to_numeric(pd.Series(values), errors="coerce").to_numpy( + dtype=np.float64 + ) + valid = np.isfinite(numeric) + n_rows, n_components = coordinates.shape + total_x = np.zeros(n_components, dtype=np.float64) + total_x2 = np.zeros(n_components, dtype=np.float64) + total_xy = np.zeros(n_components, dtype=np.float64) + total_y = float(numeric[valid].sum()) + total_y2 = float(np.square(numeric[valid]).sum()) + for start in range(0, n_rows, 65_536): + block = np.asarray(coordinates[start : start + 65_536], dtype=np.float64) + selected = valid[start : start + len(block)] + y = numeric[start : start + len(block)][selected] + block = block[selected] + total_x += block.sum(axis=0) + total_x2 += np.square(block).sum(axis=0) + total_xy += (block * y[:, None]).sum(axis=0) + n_rows = int(valid.sum()) + numerator = n_rows * total_xy - total_x * total_y + denominator = np.sqrt( + np.maximum(n_rows * total_x2 - np.square(total_x), 0.0) + * max(n_rows * total_y2 - total_y * total_y, 0.0) + ) + return np.asarray( + np.divide( + np.abs(numerator), + denominator, + out=np.zeros_like(numerator), + where=denominator > 0, + ), + dtype=np.float64, + ) + + +def _categorical_association(coordinates: Any, values: np.ndarray) -> np.ndarray: + valid = np.asarray(~pd.isna(values), dtype=bool) + labels = values[valid].astype(str) + _levels, codes = np.unique(labels, return_inverse=True) + n_rows, n_components = coordinates.shape + totals = np.zeros(n_components, dtype=np.float64) + totals_squared = np.zeros(n_components, dtype=np.float64) + if not len(codes): + return np.zeros(n_components, dtype=np.float64) + all_codes = np.full(len(values), -1, dtype=np.int64) + all_codes[valid] = codes + group_sums = np.zeros((int(codes.max()) + 1, n_components), dtype=np.float64) + group_counts = np.bincount(codes, minlength=group_sums.shape[0]).astype(np.float64) + for start in range(0, n_rows, 65_536): + block = np.asarray(coordinates[start : start + 65_536], dtype=np.float64) + selected = valid[start : start + len(block)] + block_codes = all_codes[start : start + len(block)][selected] + block = block[selected] + totals += block.sum(axis=0) + totals_squared += np.square(block).sum(axis=0) + np.add.at(group_sums, block_codes, block) + n_rows = int(valid.sum()) + grand_mean = totals / n_rows + group_means = np.divide( + group_sums, + group_counts[:, None], + out=np.zeros_like(group_sums), + where=group_counts[:, None] > 0, + ) + between = ( + group_counts[:, None] * np.square(group_means - grand_mean[None, :]) + ).sum(axis=0) + total = totals_squared - n_rows * np.square(grand_mean) + return np.asarray( + np.sqrt( + np.divide( + between, + total, + out=np.zeros_like(between), + where=total > 0, + ) + ), + dtype=np.float64, + ) + + +def _covariate_associations( + store: Any, + cell_selection: ArtifactRef, + coordinates: Any, + columns: Sequence[str], + roles: Sequence[str], + column_kinds: Mapping[str, str] | None = None, + support: dict[str, Any] | None = None, + column_artifacts: Mapping[str, ArtifactRef] | None = None, +) -> np.ndarray: + if len(columns) != len(roles): + raise ValueError("PCA covariate columns and roles must align") + associations = np.zeros((len(columns), coordinates.shape[1]), dtype=np.float64) + from ..experimental_context.characterization import _infer_kind + + for index, column in enumerate(columns): + artifact = (column_artifacts or {}).get(column) + values = ( + _aligned_metadata_values(store, cell_selection, column) + if artifact is None + else np.asarray( + resolve_cell_aligned_artifact( + store.zw, + artifact, + cell_selection=cell_selection, + expected_kind="quality_metric", + ).values + ) + ) + if values.shape != (coordinates.shape[0],): + raise ValueError(f"Covariate {column!r} does not align with PCA rows") + kind = (column_kinds or {}).get(column) or ( + "continuous" if artifact is not None else _infer_kind(values) + ) + if artifact is not None and kind != "continuous": + raise ValueError( + "QC percentage artifacts require continuous covariate evidence" + ) + if kind not in {"continuous", "categorical"}: + raise ValueError(f"Unknown covariate kind for {column!r}: {kind!r}") + valid = np.asarray(~pd.isna(values), dtype=bool) + if kind == "continuous": + numeric = pd.to_numeric(pd.Series(values), errors="coerce").to_numpy( + dtype=float + ) + valid &= np.isfinite(numeric) + associations[index] = _numeric_association(coordinates, values) + else: + if len(pd.unique(values[valid])) < int(valid.sum()): + associations[index] = _categorical_association(coordinates, values) + levels = len(pd.unique(values[valid])) + if support is not None: + support[column] = { + "kind": kind, + "method": "absolutePearson" + if kind == "continuous" + else "correlationRatio", + "completeRows": int(valid.sum()), + "missingRows": int((~valid).sum()), + "levels": levels, + "status": "computed" + if int(valid.sum()) >= 2 + and levels >= 2 + and (kind == "continuous" or levels < int(valid.sum())) + else "notComputed", + } + return associations + + +def _neighbor_overlap(store: Any, left: ArtifactRef, right: ArtifactRef) -> float: + left_values = as_zarr_array(store.load_artifact(left)["indices"], name="indices") + right_values = as_zarr_array(store.load_artifact(right)["indices"], name="indices") + if left_values.shape != right_values.shape or len(left_values.shape) != 2: + raise ValueError("Adjacent PCA neighbor artifacts must align exactly") + total = 0.0 + rows = left_values.shape[0] + for start in range(0, rows, 4096): + left_block = np.asarray(left_values[start : start + 4096]) + right_block = np.asarray(right_values[start : start + 4096]) + intersection = ( + (left_block[:, :, None] == right_block[:, None, :]).any(axis=2).sum(axis=1) + ) + union = left_block.shape[1] + right_block.shape[1] - intersection + total += float(np.divide(intersection, union).sum()) + return total / rows + + +def _write_pca_diagnostic( + store: Any, + evaluation: ParameterCandidateEvaluation, + *, + feature_selection: ArtifactRef, + selected_indices: np.ndarray, + family_masks: Mapping[str, np.ndarray], + covariate_columns: Sequence[str], + covariate_roles: Sequence[str], + adjacent_overlap: float | None, + column_kinds: Mapping[str, str] | None = None, + column_artifacts: Mapping[str, ArtifactRef] | None = None, +) -> tuple[ + ArtifactRef, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, +]: + if column_artifacts: + column_kinds = dict(column_kinds or {}) + for column, artifact in column_artifacts.items(): + if ( + artifact.kind != "quality_metric" + or column_kinds.get(column, "continuous") != "continuous" + ): + raise ValueError( + "QC percentage artifacts require continuous covariate evidence" + ) + column_kinds[column] = "continuous" + reduction = _artifact_ref(evaluation, "pca") + neighbors = _artifact_ref(evaluation, "neighbors") + reduction_status = store.inspect_artifact(reduction) + reduction_parameters = getattr(reduction_status, "parameters", None) or {} + if reduction_parameters.get("feat_scaling") is not True: + raise ValueError( + "Explained-variance ratios require the feature-scaled candidate PCA" + ) + reduction_group = store.load_artifact(reduction) + coordinates = as_zarr_array(reduction_group["data"], name="data") + loadings = as_zarr_array(reduction_group["loadings"], name="loadings") + dimensions = int(coordinates.shape[1]) + top_n = min(20, len(selected_indices)) + if loadings.shape != (len(selected_indices), dimensions): + raise ValueError("PCA loadings do not align with selected features") + planned = plan_artifact( + store.zw, + scope="assay", + assay=reduction.assay, + kind="feature_summary", + operation="diagnose_pca_representation", + parameters={ + "family_names": list(family_masks), + "covariate_columns": list(covariate_columns), + "covariate_roles": list(covariate_roles), + "covariate_kinds": { + column: (column_kinds or {}).get(column, "inferred") + for column in covariate_columns + }, + "covariate_method": "typedCompleteCaseAssociation", + "top_loading_count": top_n, + "family_mask_fingerprints": { + family: hashlib.sha256( + np.asarray(mask, dtype=bool).tobytes() + ).hexdigest() + for family, mask in family_masks.items() + }, + "covariate_fingerprints": { + column: ( + column_artifacts[column].to_dict() + if column_artifacts and column in column_artifacts + else _metadata_column_fingerprint(store.cells, column) + ) + for column in covariate_columns + }, + "adjacent_neighbor_overlap": adjacent_overlap, + "explained_variance_basis": "scaled_nonconstant_features", + }, + inputs={ + "reduction": reduction, + "neighbors": neighbors, + "feature_selection": feature_selection, + "covariate_artifacts": dict(column_artifacts or {}), + }, + execution_options={}, + invalidate_cache=False, + required_arrays=( + ArrayRequirement( + "component_variance", shape=(dimensions,), dtype=np.float64 + ), + ArrayRequirement( + "explained_variance_ratio", shape=(dimensions,), dtype=np.float64 + ), + ArrayRequirement( + "top_loading_feature_indices", shape=(dimensions, top_n), dtype=np.int64 + ), + ArrayRequirement( + "top_loading_values", shape=(dimensions, top_n), dtype=np.float64 + ), + ArrayRequirement( + "family_enrichment", + shape=(len(family_masks), dimensions), + dtype=np.float64, + ), + ArrayRequirement( + "covariate_association", + shape=(len(covariate_columns), dimensions), + dtype=np.float64, + ), + ArrayRequirement("adjacent_neighbor_overlap", shape=(1,), dtype=np.float64), + ), + required_attributes=( + AttributeRequirement("family_names", expected_types=(list,)), + AttributeRequirement("covariate_columns", expected_types=(list,)), + AttributeRequirement("covariate_roles", expected_types=(list,)), + AttributeRequirement("covariate_support", expected_types=(dict,)), + AttributeRequirement("payload_fingerprint", expected_types=(str,)), + ), + ) + if planned.reused: + group = store.load_artifact(planned.ref) + if ( + fingerprint_stored_arrays(group, _PCA_DIAGNOSTIC_ARRAYS) + != group.attrs["payload_fingerprint"] + ): + raise ValueError("Stored PCA diagnostic payload fingerprint does not match") + restored = ( + planned.ref, + np.asarray( + as_zarr_array(group["component_variance"], name="component_variance")[:] + ), + np.asarray( + as_zarr_array( + group["explained_variance_ratio"], name="explained_variance_ratio" + )[:] + ), + np.asarray( + as_zarr_array( + group["top_loading_feature_indices"], + name="top_loading_feature_indices", + )[:] + ), + np.asarray( + as_zarr_array(group["top_loading_values"], name="top_loading_values")[:] + ), + np.asarray( + as_zarr_array(group["family_enrichment"], name="family_enrichment")[:] + ), + np.asarray( + as_zarr_array( + group["covariate_association"], name="covariate_association" + )[:] + ), + ) + diagnostic_reuse("diagnostic.pca", "artifactReuses") + return restored + component_variance = diagnostic_call( + "metric.pcaVariance", _component_variance, coordinates + ) + total_scaled_variance = diagnostic_call( + "metric.pcaScaledVariance", + _scaled_total_variance, + store, + reduction_status, + n_rows=int(coordinates.shape[0]), + n_features=len(selected_indices), + feature_selection=feature_selection, + ) + explained_variance_ratio = np.clip( + component_variance / total_scaled_variance, + 0.0, + 1.0, + ) + top_indices, top_values, family_enrichment = diagnostic_call( + "metric.pcaLoadings", + _top_loadings, + loadings, + selected_indices, + family_masks, + ) + covariate_support: dict[str, Any] = {} + associations = ( + diagnostic_call( + "metric.pcaCovariates", + _covariate_associations, + store, + ArtifactRef( + scope=evaluation.cellSelection.scope, + assay=evaluation.cellSelection.assay, + kind=evaluation.cellSelection.kind, + artifact_id=evaluation.cellSelection.artifactId, + ), + coordinates, + covariate_columns, + covariate_roles, + column_kinds, + covariate_support, + column_artifacts, + ) + if evaluation.cellSelection is not None + else np.zeros((len(covariate_columns), coordinates.shape[1]), dtype=np.float64) + ) + overlap_array = np.asarray( + [np.nan if adjacent_overlap is None else adjacent_overlap], + dtype=np.float64, + ) + payload = { + "component_variance": component_variance, + "explained_variance_ratio": explained_variance_ratio, + "top_loading_feature_indices": top_indices, + "top_loading_values": top_values, + "family_enrichment": family_enrichment, + "covariate_association": associations, + "adjacent_neighbor_overlap": overlap_array, + } + if not planned.reused: + group = start_artifact(store.zw, planned) + for name, values in payload.items(): + chunks = tuple(max(1, min(size, 4096)) for size in values.shape) + array = create_zarr_dataset( + group, + name, + chunks, + values.dtype, + values.shape, + ) + array[:] = values + group.attrs["family_names"] = list(family_masks) + group.attrs["covariate_columns"] = list(covariate_columns) + group.attrs["covariate_roles"] = list(covariate_roles) + group.attrs["covariate_support"] = covariate_support + group.attrs["payload_fingerprint"] = fingerprint_stored_arrays( + group, + _PCA_DIAGNOSTIC_ARRAYS, + ) + finish_artifact(group, planned) + return ( + planned.ref, + component_variance, + explained_variance_ratio, + top_indices, + top_values, + family_enrichment, + associations, + ) + + +def augment_pca_evaluations( + store: Any, + evaluations: Sequence[ParameterCandidateEvaluation], + *, + feature_selection: ArtifactRef, + nominated_families: Sequence[str], + protected_families: Sequence[str], + technical_columns: Sequence[str], + protected_columns: Sequence[str], + qc_columns: Sequence[str], + batch_columns: Sequence[str] = (), + column_kinds: Mapping[str, str] | None = None, + qc_artifacts: Mapping[str, ArtifactRef] | None = None, +) -> tuple[ParameterCandidateEvaluation, ...]: + """Attach persisted PCA loading, variance, topology, and covariate evidence.""" + selected_indices, selected_names = _selected_feature_names( + store, + feature_selection, + ) + family_masks = { + family: mask + for family in dict.fromkeys([*nominated_families, *protected_families]) + for mask in [_family_mask(selected_names, family)] + if mask is not None and bool(mask.any()) + } + requested_role_columns = { + "technical": tuple(technical_columns), + "batch": tuple(batch_columns or technical_columns), + "protected": tuple(protected_columns), + "qc": tuple(qc_columns), + } + columns: list[str] = [] + roles: list[str] = [] + for role, values in ( + ("technical", technical_columns), + ("batch", batch_columns), + ("protected", protected_columns), + ("qc", qc_columns), + ): + for column in values: + if ( + column in store.cells.columns or column in (qc_artifacts or {}) + ) and column not in columns: + columns.append(column) + roles.append(role) + completed = [ + evaluation + for evaluation in evaluations + if evaluation.status == "done" + and evaluation.eligible + and evaluation.cellSelection is not None + and "pca" in evaluation.artifacts + and "neighbors" in evaluation.artifacts + ] + previous_by_id: dict[str, float | None] = {} + previous: ParameterCandidateEvaluation | None = None + for evaluation in sorted( + completed, + key=lambda value: value.parameters.dimensions, + ): + overlap = ( + diagnostic_call( + "metric.neighborOverlap", + _neighbor_overlap, + store, + _artifact_ref(previous, "neighbors"), + _artifact_ref(evaluation, "neighbors"), + ) + if previous is not None + else None + ) + previous_by_id[evaluation.candidateId] = overlap + previous = evaluation + + augmented: list[ParameterCandidateEvaluation] = [] + for evaluation in evaluations: + if evaluation.candidateId not in previous_by_id: + augmented.append(evaluation) + continue + ( + diagnostic, + variance, + explained_variance_ratio, + top_indices, + _top_values, + family_enrichment, + associations, + ) = diagnostic_call( + "diagnostic.pca", + _write_pca_diagnostic, + store, + evaluation, + feature_selection=feature_selection, + selected_indices=selected_indices, + family_masks=family_masks, + covariate_columns=columns, + covariate_roles=roles, + adjacent_overlap=previous_by_id[evaluation.candidateId], + column_kinds=column_kinds, + column_artifacts=qc_artifacts, + ) + family_maxima = { + family: float(family_enrichment[index].max(initial=0.0)) + for index, family in enumerate(family_masks) + } + family_by_component = { + f"PC{component + 1}": { + family: float(family_enrichment[family_index, component]) + for family_index, family in enumerate(family_masks) + } + for component in range(family_enrichment.shape[1]) + } + feature_name_by_index = dict( + zip(selected_indices.tolist(), selected_names.tolist(), strict=True) + ) + top_loading_genes = { + f"PC{component + 1}": [ + str(feature_name_by_index[int(index)]) + for index in top_indices[component] + ] + for component in range(top_indices.shape[0]) + } + support = dict(store.load_artifact(diagnostic).attrs["covariate_support"]) + column_index = { + column: index + for index, column in enumerate(columns) + if support.get(column, {}).get("status") == "computed" + } + unsupported_covariates = [ + column for column in columns if column not in column_index + ] + component_associations = { + role: { + column: associations[column_index[column]].tolist() + for column in role_columns + if column in column_index + } + for role, role_columns in requested_role_columns.items() + } + role_associations = { + role: { + column: float(associations[column_index[column]].max(initial=0.0)) + for column in role_columns + if column in column_index + } + for role, role_columns in requested_role_columns.items() + } + metrics = evaluation.metrics.model_copy( + update={ + "componentVariance": variance.tolist(), + "pcaExplainedVarianceRatio": explained_variance_ratio.tolist(), + "pcaCumulativeExplainedVarianceRatio": np.cumsum( + explained_variance_ratio + ) + .clip(max=1.0) + .tolist(), + "topLoadingGenes": top_loading_genes, + "loadingFamilyEnrichment": family_maxima, + "loadingFamilyEnrichmentByComponent": family_by_component, + "pcaComponentAssociations": component_associations, + "batchPcaAssociation": role_associations["batch"], + "technicalPcaAssociation": role_associations["technical"], + "protectedPcaAssociation": role_associations["protected"], + "qcPcaAssociation": role_associations["qc"], + "neighborPrefixOverlap": previous_by_id[evaluation.candidateId], + } + ) + artifact = ArtifactRecord.from_ref(diagnostic) + augmented.append( + evaluation.model_copy( + update={ + "metrics": metrics, + "warnings": [ + *evaluation.warnings, + *( + f"PCA association for {column!r} is unavailable with {support.get(column, {}).get('completeRows', 0)} complete rows and {support.get(column, {}).get('levels', 0)} distinct values." + for column in unsupported_covariates + ), + ], + "artifacts": { + **evaluation.artifacts, + "representationDiagnostic": artifact, + }, + "evidenceIds": list( + dict.fromkeys( + [ + *evaluation.evidenceIds, + f"candidate:{evaluation.candidateId}:pcaVariance", + ( + f"candidate:{evaluation.candidateId}:" + "pcaExplainedVariance" + ), + f"candidate:{evaluation.candidateId}:pcaLoadings", + f"candidate:{evaluation.candidateId}:pcaCovariates", + *( + [ + f"candidate:{evaluation.candidateId}:neighborPrefixOverlap" + ] + if previous_by_id[evaluation.candidateId] + is not None + else [] + ), + ] + ) + ), + } + ) + ) + return annotate_candidate_dominance(augmented) + + +def _bounded_score_summary( + values: Any, + *, + maximum_sample_size: int, +) -> tuple[dict[str, float], np.ndarray]: + if maximum_sample_size < 1: + raise ValueError("maximum_sample_size must be positive") + if len(values.shape) != 1 or values.shape[0] < 1: + raise ValueError("Doublet scores must be one non-empty vector") + n_values = int(values.shape[0]) + stride = max(1, (n_values + maximum_sample_size - 1) // maximum_sample_size) + sampled: list[np.ndarray] = [] + minimum = float("inf") + maximum = float("-inf") + for start in range(0, n_values, 65_536): + block = np.asarray(values[start : start + 65_536], dtype=np.float64) + if not np.isfinite(block).all(): + raise ValueError("Doublet scores must be finite") + minimum = min(minimum, float(block.min())) + maximum = max(maximum, float(block.max())) + offset = (-start) % stride + sampled.append(block[offset::stride]) + sample = np.concatenate(sampled) + return ( + { + "nCells": float(n_values), + "sampleSize": float(len(sample)), + "minimum": minimum, + "p50": float(np.quantile(sample, 0.5)), + "p90": float(np.quantile(sample, 0.9)), + "p95": float(np.quantile(sample, 0.95)), + "p99": float(np.quantile(sample, 0.99)), + "maximum": maximum, + }, + sample, + ) + + +def _build_advisory_doublet_scores( + store: Any, + *, + scores: Sequence[ArtifactRef], + cell_selections: Sequence[ArtifactRef], + native_graph: ArtifactRef, + native_clusters: ArtifactRef, + parent_selection: ArtifactRef, + capture_values: Sequence[str], + capture_column: str | None, + limitations: Sequence[str], +) -> AdvisoryDoubletScores: + score_refs = tuple(scores) + selections = tuple(cell_selections) + captures = tuple(capture_values) + if not score_refs or not (len(score_refs) == len(selections) == len(captures)): + raise ValueError("Doublet score artifacts require aligned capture summaries") + parent_indices = read_stored_selection_indices( + store.zw, + parent_selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ) + if len(parent_indices) < 1: + raise ValueError("Doublet evidence parent selection cannot be empty") + score_arrays = tuple( + as_zarr_array(store.load_artifact(score_ref)["values"], name="values") + for score_ref in score_refs + ) + total_score_cells = sum(int(values.shape[0]) for values in score_arrays) + if total_score_cells < 1: + raise ValueError("Doublet score artifacts cannot all be empty") + summaries: list[dict[str, float]] = [] + samples: list[np.ndarray] = [] + covered_cells = 0 + for values, selection_ref in zip(score_arrays, selections, strict=True): + selection_indices = read_stored_selection_indices( + store.zw, + selection_ref, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ) + if values.shape != selection_indices.shape: + raise ValueError("Doublet scores do not align with their cell selection") + summary, sample = diagnostic_call( + "metric.doubletScoreSummary", + _bounded_score_summary, + values, + maximum_sample_size=max( + 1, + int(65_536 * int(values.shape[0]) / total_score_cells), + ), + ) + summaries.append(summary) + samples.append(sample) + covered_cells += len(selection_indices) + combined = np.concatenate(samples) + aggregate = { + "p50": float(np.quantile(combined, 0.5)), + "p90": float(np.quantile(combined, 0.9)), + "p95": float(np.quantile(combined, 0.95)), + "p99": float(np.quantile(combined, 0.99)), + "maximum": max(summary["maximum"] for summary in summaries), + } + reported_limitations = list(limitations) + if total_score_cells > 65_536: + reported_limitations.append( + "Doublet score quantiles use deterministic bounded samples above " + "65,536 scored cells." + ) + return AdvisoryDoubletScores( + scores=score_refs, + cell_selections=selections, + native_graph=native_graph, + native_clusters=native_clusters, + capture_values=captures, + score_summaries=tuple(summaries), + score_quantiles=aggregate, + capture_coverage=min(1.0, covered_cells / len(parent_indices)), + capture_column=capture_column, + limitations=tuple(reported_limitations), + ) + + +def _select_capture_cells( + store: Any, + parent: ArtifactRef, + *, + column: str, + value: Any, + active_indices: np.ndarray, + active_values: np.ndarray, +) -> tuple[ArtifactRef, int]: + if active_indices.shape != active_values.shape: + raise ValueError("Capture values must align with the selected cells") + labels = active_values.astype(str) + selected = labels == str(value) + expected = np.zeros(store.cells.N, dtype=bool) + expected[active_indices] = selected + reference = diagnostic_call( + "core.captureSelection", + store.filter_cells, + [column], + [value], + [value], + cell_selection=parent, + keep_bounds=True, + invalidate_cache=False, + ) + stored_array = cast(Any, store.load_artifact(reference)["values"]) + stored = np.asarray(stored_array[:], dtype=bool) + if not np.array_equal(stored, expected): + raise RuntimeError("Capture selection does not match its validated evidence") + return reference, int(stored.sum()) + + +def resolve_native_doublet_inputs( + store: Any, + selected: ParameterCandidateEvaluation, + evaluations: Sequence[ParameterCandidateEvaluation], +) -> tuple[ArtifactRef, ArtifactRef]: + """Return or reconstruct the parameter-matched uncorrected graph and clusters.""" + parameters = selected.parameters + exact_native = next( + ( + evaluation + for evaluation in evaluations + if evaluation.status == "done" + and evaluation.eligible + and not evaluation.parameters.useHarmony + and evaluation.parameters.dimensions == parameters.dimensions + and evaluation.parameters.neighborsK == parameters.neighborsK + and evaluation.parameters.leidenResolution == parameters.leidenResolution + and "clusters" in evaluation.artifacts + and "connectivityMap" in evaluation.artifacts + ), + None, + ) + if exact_native is not None: + inputs = ( + _artifact_ref(exact_native, "clusters"), + _artifact_ref(exact_native, "connectivityMap"), + ) + diagnostic_reuse("diagnostic.nativeDoubletInputs", "artifactReuses") + return inputs + if not selected.parameters.useHarmony: + inputs = ( + _artifact_ref(selected, "clusters"), + _artifact_ref(selected, "connectivityMap"), + ) + diagnostic_reuse("diagnostic.nativeDoubletInputs", "artifactReuses") + return inputs + reduction = _artifact_ref(selected, "pca") + ann = diagnostic_call( + "core.nativeDoubletAnn", + store.build_ann_index, + reduction, + ann_metric="l2", + ann_parallel=False, + rand_state=4444, + invalidate_cache=False, + ) + neighbors = diagnostic_call( + "core.nativeDoubletNeighbors", + store.query_neighbors, + ann, + coordinates=reduction, + k=parameters.neighborsK, + invalidate_cache=False, + ) + graph = diagnostic_call( + "core.nativeDoubletGraph", + store.build_connectivity_map, + neighbors, + local_connectivity=1.0, + bandwidth=1.5, + invalidate_cache=False, + ) + clusters = diagnostic_call( + "core.nativeDoubletPartition", + store.run_leiden_clustering, + graph, + resolution=parameters.leidenResolution, + backend="igraph", + symmetric_graph=False, + graph_upper_only=False, + random_seed=4444, + invalidate_cache=False, + ) + return clusters, graph + + +def score_advisory_doublets( + store: Any, + selected: ParameterCandidateEvaluation, + evaluations: Sequence[ParameterCandidateEvaluation], + *, + assay: str, + feature_selection: ArtifactRef, + capture_column: str | None, +) -> AdvisoryDoubletScores: + """Score doublet evidence without making singlet or removal decisions.""" + if selected.cellSelection is None: + raise ValueError("Doublet scoring requires an exact selected cell axis") + parent_selection = ArtifactRef( + scope=selected.cellSelection.scope, + assay=selected.cellSelection.assay, + kind=selected.cellSelection.kind, + artifact_id=selected.cellSelection.artifactId, + ) + native_clusters, native_graph = resolve_native_doublet_inputs( + store, + selected, + evaluations, + ) + feature_ids = np.asarray(store.get_assay(assay).feats.fetch_all("ids")).astype(str) + unique_ids, id_counts = np.unique(feature_ids, return_counts=True) + duplicate_ids = unique_ids[id_counts > 1] + if duplicate_ids.size: + examples = duplicate_ids[:5].tolist() + limitation = ( + f"Advisory doublet scoring was not run for assay {assay!r} because " + f"{duplicate_ids.size} feature identifiers are duplicated. " + "Simulated-to-observed mapping requires unique identifiers. " + f"Examples: {examples}." + ) + logger.warning(limitation) + return AdvisoryDoubletScores( + scores=(), + cell_selections=(), + native_graph=native_graph, + native_clusters=native_clusters, + capture_column=capture_column, + limitations=(limitation,), + ) + limitations: list[str] = [] + if capture_column is None or capture_column not in store.cells.columns: + score = diagnostic_call( + "core.doubletDetection", + store.run_doublet_detection, + native_clusters, + native_graph, + from_assay=assay, + invalidate_cache=False, + ) + limitations.append( + "Physical capture identity was unavailable, so advisory doublet " + "scores were computed across the selected dataset." + ) + return _build_advisory_doublet_scores( + store, + scores=(score,), + cell_selections=(parent_selection,), + native_graph=native_graph, + native_clusters=native_clusters, + parent_selection=parent_selection, + capture_values=("allSelectedCells",), + capture_column=None, + limitations=limitations, + ) + + active_indices = read_stored_selection_indices( + store.zw, + parent_selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + capture_values = read_metadata_rows_chunkwise( + store.cells, + capture_column, + active_indices, + ) + capture_labels = capture_values.astype(str) + unique_capture_labels, first_capture_indices = np.unique( + capture_labels, + return_index=True, + ) + capture_groups = unique_capture_labels.tolist() + raw_capture_values = { + str(label): capture_values[int(index)] + for label, index in zip( + unique_capture_labels, + first_capture_indices, + strict=True, + ) + } + if len(capture_groups) > _MAX_DOUBLET_CAPTURES: + raise ValueError( + "Physical capture column exceeds the advisory doublet limit of " + f"{_MAX_DOUBLET_CAPTURES} values" + ) + if len(capture_groups) == 1: + score = diagnostic_call( + "core.doubletDetection", + store.run_doublet_detection, + native_clusters, + native_graph, + from_assay=assay, + invalidate_cache=False, + ) + return _build_advisory_doublet_scores( + store, + scores=(score,), + cell_selections=(parent_selection,), + native_graph=native_graph, + native_clusters=native_clusters, + parent_selection=parent_selection, + capture_values=(capture_groups[0],), + capture_column=capture_column, + limitations=limitations, + ) + + n_features = len(read_feature_selection_indices(store.zw, assay, feature_selection)) + scores: list[ArtifactRef] = [] + selections: list[ArtifactRef] = [] + scored_captures: list[str] = [] + for capture_value in capture_groups: + capture_selection, capture_cells = _select_capture_cells( + store, + parent_selection, + column=capture_column, + value=raw_capture_values[capture_value], + active_indices=active_indices, + active_values=capture_values, + ) + dimensions = min( + selected.parameters.dimensions, + capture_cells - 1, + n_features - 1, + ) + neighbors_k = min(selected.parameters.neighborsK, capture_cells - 1) + if dimensions < 2 or neighbors_k < 2: + limitations.append( + "Advisory doublet scores were not computed for capture " + f"{capture_value!r} because it contains only {capture_cells} " + "selected cells." + ) + continue + normalized = diagnostic_call( + "core.captureNormalization", + store.run_normalization, + capture_selection, + features=feature_selection, + log_transform=True, + renormalize_subset=True, + invalidate_cache=False, + ) + reduction = diagnostic_call( + "core.capturePca", + store.run_pca, + normalized, + dims=dimensions, + feat_scaling=True, + invalidate_cache=False, + ) + ann = diagnostic_call( + "core.captureAnn", + store.build_ann_index, + reduction, + ann_metric="l2", + ann_parallel=False, + rand_state=4444, + invalidate_cache=False, + ) + neighbors = diagnostic_call( + "core.captureNeighbors", + store.query_neighbors, + ann, + coordinates=reduction, + k=neighbors_k, + invalidate_cache=False, + ) + graph = diagnostic_call( + "core.captureGraph", + store.build_connectivity_map, + neighbors, + local_connectivity=1.0, + bandwidth=1.5, + invalidate_cache=False, + ) + clusters = diagnostic_call( + "core.capturePartition", + store.run_leiden_clustering, + graph, + resolution=selected.parameters.leidenResolution, + backend="igraph", + symmetric_graph=False, + graph_upper_only=False, + random_seed=4444, + invalidate_cache=False, + ) + scores.append( + diagnostic_call( + "core.doubletDetection", + store.run_doublet_detection, + clusters, + graph, + from_assay=assay, + invalidate_cache=False, + ) + ) + selections.append(capture_selection) + scored_captures.append(capture_value) + if not scores: + raise ValueError("No physical capture had enough cells for doublet scoring") + return _build_advisory_doublet_scores( + store, + scores=scores, + cell_selections=selections, + native_graph=native_graph, + native_clusters=native_clusters, + parent_selection=parent_selection, + capture_values=scored_captures, + capture_column=capture_column, + limitations=limitations, + ) + + +def _doublet_concentration( + store: Any, + labels: np.ndarray, + cell_selection: ArtifactRef, + evidence: AdvisoryDoubletScores, +) -> float | None: + parent_indices = read_stored_selection_indices( + store.zw, + cell_selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + if labels.shape != parent_indices.shape: + raise ValueError("Cluster labels do not align with advisory doublet evidence") + if len(parent_indices) > 1 and np.any(parent_indices[1:] <= parent_indices[:-1]): + raise ValueError("Doublet parent selection indices must be strictly increasing") + high_score = np.zeros(len(parent_indices), dtype=bool) + covered = np.zeros(len(parent_indices), dtype=bool) + if evidence.score_summaries and len(evidence.score_summaries) != len( + evidence.scores + ): + raise ValueError("Doublet score summaries do not align with score artifacts") + for score_index, (score_ref, selection_ref) in enumerate( + zip( + evidence.scores, + evidence.cell_selections, + strict=True, + ) + ): + score_values = as_zarr_array( + store.load_artifact(score_ref)["values"], + name="values", + ) + selection_indices = read_stored_selection_indices( + store.zw, + selection_ref, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + if score_values.shape != selection_indices.shape: + raise ValueError("Doublet scores do not align with their cell selection") + summary = ( + evidence.score_summaries[score_index] + if evidence.score_summaries + else diagnostic_call( + "metric.doubletScoreSummary", + _bounded_score_summary, + score_values, + maximum_sample_size=65_536, + )[0] + ) + threshold = float(summary["p90"]) + for start in range(0, len(selection_indices), 65_536): + local_indices = selection_indices[start : start + 65_536] + local_positions = np.searchsorted(parent_indices, local_indices) + if np.any(local_positions >= len(parent_indices)) or not np.array_equal( + parent_indices[local_positions], + local_indices, + ): + raise ValueError( + "Doublet score selection is outside its parent selection" + ) + if covered[local_positions].any(): + raise ValueError("Doublet score selections must not overlap") + local_scores = np.asarray( + score_values[start : start + len(local_indices)], + dtype=np.float64, + ) + if not np.isfinite(local_scores).all(): + raise ValueError("Doublet scores must be finite") + covered[local_positions] = True + high_score[local_positions] = local_scores >= threshold + if not covered.any() or not high_score[covered].any(): + return None + baseline = float(high_score[covered].mean()) + enrichments = [ + float(high_score[covered & (labels == cluster)].mean()) / baseline + for cluster in np.unique(labels[covered]) + if bool((covered & (labels == cluster)).any()) + ] + return max(enrichments, default=0.0) + + +def _aligned_metadata( + store: Any, + cell_selection: ArtifactRef, + column: str, +) -> np.ndarray: + indices = read_stored_selection_indices( + store.zw, + cell_selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ).astype(np.int64, copy=False) + values = np.asarray(read_metadata_rows_chunkwise(store.cells, column, indices)) + if values.shape != (len(indices),): + raise ValueError(f"Metadata column {column!r} does not align with clusters") + return values.astype(str) + + +def _cross_unit_support(labels: np.ndarray, units: np.ndarray) -> float | None: + available_units = np.unique(units) + if len(available_units) < 2: + return None + supported = [ + len(np.unique(units[labels == cluster])) >= 2 for cluster in np.unique(labels) + ] + return float(np.mean(supported)) if supported else None + + +def population_support_evidence( + store: Any, + evaluation: ParameterCandidateEvaluation, + columns: Sequence[str], +) -> dict[str, Any]: + """Describe observed population support using exact cells and requested units. + + This is descriptive support, not a validation of population identity or + independent replication. Display limits do not change any count or fraction. + """ + if evaluation.status != "done" or evaluation.cellSelection is None: + raise ValueError("Population support requires completed, cell-bound evidence") + selection = ArtifactRef( + scope=evaluation.cellSelection.scope, + assay=evaluation.cellSelection.assay, + kind=evaluation.cellSelection.kind, + artifact_id=evaluation.cellSelection.artifactId, + ) + clusters = _artifact_ref(evaluation, "clusters") + status = store.inspect_artifact(clusters) + if not status.exists or not status.complete: + raise ValueError("Population support requires complete cluster evidence") + raw_selection = (status.inputs or {}).get("cell_selection") + if ( + not isinstance(raw_selection, Mapping) + or ArtifactRef.from_dict(dict(raw_selection)) != selection + ): + raise ValueError("Population support cluster and candidate cells differ") + if clusters.kind not in {"cluster_labels", "cluster_cut"}: + raise ValueError("Population support requires a clustering artifact") + indices = read_stored_selection_indices( + store.zw, + selection, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ) + if not len(indices) or np.any(indices[1:] <= indices[:-1]): + raise ValueError("Population support requires distinct ordered selected cells") + labels = as_zarr_array( + store.load_artifact(clusters)[ + "values" if clusters.kind == "cluster_labels" else "labels" + ], + name="cluster labels", + ) + if labels.shape != indices.shape or labels.dtype.kind not in "iuf": + raise ValueError("Population labels do not align with the selected cells") + requested = list(dict.fromkeys(columns)) + selected_columns = requested[:2] + available = [column for column in selected_columns if column in store.cells.columns] + totals: Counter[int] = Counter() + group_totals: dict[str, Counter[tuple[str, Any]]] = { + column: Counter() for column in available + } + counts: dict[str, dict[int, Counter[tuple[str, Any]]]] = { + column: defaultdict(Counter) for column in available + } + for start in range(0, len(indices), 65_536): + stop = min(start + 65_536, len(indices)) + block = np.asarray(labels[start:stop]) + if not np.isfinite(block).all() or not np.equal(block, np.floor(block)).all(): + raise ValueError("Population labels must be finite integers") + population_ids = block.astype(np.int64) + unique, sizes = np.unique(population_ids, return_counts=True) + totals.update( + { + int(population): int(size) + for population, size in zip(unique, sizes, strict=True) + } + ) + rows = indices[start:stop] + for column in available: + values = np.asarray(read_metadata_rows_chunkwise(store.cells, column, rows)) + missing = read_metadata_missing_rows_chunkwise(store.cells, column, rows) + if ( + values.shape != rows.shape + or missing is not None + and missing.shape != rows.shape + ): + raise ValueError( + f"Population metadata {column!r} does not align with cells" + ) + for offset, (population, value) in enumerate( + zip(population_ids, values, strict=True) + ): + value = value.item() if isinstance(value, np.generic) else value + if ( + missing is not None + and missing[offset] + or pd.isna(value) + or isinstance(value, float) + and not np.isfinite(value) + ): + continue + if isinstance(value, bytes): + value = value.decode("utf-8") + if not isinstance(value, (str, int, float, bool)): + raise ValueError("Population unit labels must be scalar values") + if isinstance(value, str) and not value.strip(): + continue + key = (type(value).__name__, value) + group_totals[column][key] += 1 + counts[column][int(population)][key] += 1 + evidence_columns: dict[str, Any] = {} + for column in selected_columns: + if column not in available: + evidence_columns[column] = { + "status": "unavailable", + "reason": "Column is absent.", + } + continue + population_rows: list[dict[str, Any]] = [] + for population, size in totals.items(): + group_counts = counts[column][population] + ordered = sorted( + group_counts.items(), + key=lambda item: (-item[1], item[0][0], str(item[0][1])), + ) + displayed = ordered[:5] + covered = sum(group_counts.values()) + population_rows.append( + { + "cluster": str(population), + "cells": size, + "coveredCells": covered, + "missingCells": size - covered, + "coverageFraction": covered / size, + "supportingGroups": len(group_counts), + "groupsWithAtLeast5Cells": sum( + count >= 5 for count in group_counts.values() + ), + "largestGroupFraction": ordered[0][1] / size if ordered else None, + "topGroups": [ + { + "value": key[1], + "valueType": key[0], + "cells": count, + "fractionOfPopulation": count / size, + "fractionOfGroup": count / group_totals[column][key], + } + for key, count in displayed + ], + "omittedGroups": len(ordered) - len(displayed), + "omittedCells": sum(count for _, count in ordered[5:]), + } + ) + population_rows.sort( + key=lambda row: ( + -(row["largestGroupFraction"] or 0), + row["cells"], + row["cluster"], + ) + ) + covered = sum(group_totals[column].values()) + evidence_columns[column] = { + "status": "computed", + "observedGroups": len(group_totals[column]), + "coveredCells": covered, + "missingCells": len(indices) - covered, + "coverageFraction": covered / len(indices), + "populations": population_rows[:64], + "omittedPopulations": max(0, len(population_rows) - 64), + "omittedPopulationCells": sum(row["cells"] for row in population_rows[64:]), + } + return { + "candidateId": evaluation.candidateId, + "cellSelection": selection.to_dict(), + "clusters": clusters.to_dict(), + "selectedCells": len(indices), + "observedPopulations": len(totals), + "columns": evidence_columns, + "omittedColumns": requested[2:], + "displayLimits": { + "columns": 2, + "populationsPerColumn": 64, + "groupsPerPopulation": 5, + }, + "interpretation": ( + "Counts use all exact selected cells. Population fractions include cells with missing unit metadata; " + "group fractions use all selected cells with that unit value. Missing values are unassigned. " + "Displayed populations prioritize concentration in one group, then smaller size. " + "A group with at least five cells is a descriptive count, not a replication threshold. " + "Shared donor or capture support does not establish biological identity or rule out artifacts." + ), + } + + +def _subsample_partition_stability( + graph: Any, + labels: np.ndarray, + resolution: float, +) -> float: + if graph.shape != (len(labels), len(labels)): + raise ValueError("Candidate graph does not align with cluster labels") + selected = np.arange(len(labels)) % 5 != 0 + if int(selected.sum()) < 3: + selected = np.ones(len(labels), dtype=bool) + subsample_labels = diagnostic_call( + "core.subsamplePartition", + leiden_membership, + graph[selected][:, selected], + resolution, + 4444, + backend="igraph", + ) + return float( + diagnostic_call( + "metric.partitionAgreement", + adjusted_rand_score, + labels[selected], + subsample_labels, + ) + ) + + +def augment_cluster_evaluations( + store: Any, + evaluations: Sequence[ParameterCandidateEvaluation], + *, + marker_assay: str, + marker_features: ArtifactRef, + independent_unit_columns: Sequence[str], + technical_columns: Sequence[str], + nominated_families: Sequence[str] = (), + protected_families: Sequence[str] = (), + doublet_evidence: AdvisoryDoubletScores | None = None, +) -> tuple[ParameterCandidateEvaluation, ...]: + """Add seed, marker, unit-support, and technical-association evidence.""" + _marker_indices, marker_feature_names = _selected_feature_names( + store, + marker_features, + ) + family_masks = { + family: mask + for family in dict.fromkeys([*nominated_families, *protected_families]) + for mask in [_family_mask(marker_feature_names, family)] + if mask is not None and bool(mask.any()) + } + augmented: list[ParameterCandidateEvaluation] = [] + for evaluation in evaluations: + if evaluation.status != "done" or evaluation.cellSelection is None: + augmented.append(evaluation) + continue + graph_ref = _artifact_ref(evaluation, "connectivityMap") + clusters_ref = _artifact_ref(evaluation, "clusters") + labels = _cluster_labels(store, clusters_ref) + alternative_ref = diagnostic_call( + "core.alternateSeedPartition", + store.run_leiden_clustering, + graph_ref, + resolution=evaluation.parameters.leidenResolution, + backend="igraph", + symmetric_graph=False, + graph_upper_only=False, + random_seed=9173, + invalidate_cache=False, + ) + alternative = _cluster_labels(store, alternative_ref) + if alternative.shape != labels.shape: + raise ValueError("Alternate-seed clusters do not align with the candidate") + seed_stability = float( + diagnostic_call( + "metric.partitionAgreement", adjusted_rand_score, labels, alternative + ) + ) + subsample_stability = _cached_candidate_metric( + ( + id(store), + "subsample_stability", + graph_ref, + clusters_ref, + evaluation.parameters.leidenResolution, + ), + lambda: _subsample_partition_stability( + store.load_graph(graph_ref), + labels, + evaluation.parameters.leidenResolution, + ), + ) + + cluster_count = len(np.unique(labels)) + marker_ref = None + markers = pd.DataFrame() + if cluster_count >= 2: + marker_ref = diagnostic_call( + "core.markers", + store.run_marker_search, + clusters_ref, + from_assay=marker_assay, + features=marker_features, + invalidate_cache=False, + ) + markers = store.get_markers( + marker_ref, + min_score=0.25, + min_frac_exp=0.2, + ) + marker_groups = ( + set(markers["group_id"].astype(str)) + if "group_id" in markers.columns + else set() + ) + marker_coherence = ( + float(len(marker_groups) / cluster_count) + if marker_ref is not None + else None + ) + marker_names = ( + markers["feature_name"].astype(str).to_numpy() + if "feature_name" in markers.columns + else np.asarray([], dtype=str) + ) + marker_specificity: dict[str, float] = {} + marker_auc: dict[str, float] = {} + top_marker_genes: dict[str, list[str]] = {} + marker_group_values = ( + markers["group_id"].astype(str) if "group_id" in markers.columns else None + ) + for cluster in np.unique(labels): + cluster_id = str(cluster) + cluster_markers = ( + markers.loc[marker_group_values == cluster_id] + if marker_group_values is not None + else markers.iloc[0:0] + ) + if "score" in cluster_markers: + cluster_markers = cluster_markers.sort_values( + "score", + ascending=False, + kind="stable", + ) + top_scores = cluster_markers["score"].to_numpy( + dtype=np.float64, + )[:10] + top_scores = top_scores[np.isfinite(top_scores)] + if len(top_scores): + marker_specificity[cluster_id] = float(np.median(top_scores)) + if "auc" in cluster_markers: + top_auc = cluster_markers["auc"].to_numpy(dtype=np.float64)[:10] + top_auc = top_auc[np.isfinite(top_auc)] + if len(top_auc): + marker_auc[cluster_id] = float(np.median(top_auc)) + top_marker_genes[cluster_id] = ( + cluster_markers["feature_name"].astype(str).head(10).tolist() + if "feature_name" in cluster_markers + else [] + ) + marker_specificity_median = ( + float(np.median(list(marker_specificity.values()))) + if marker_specificity + else None + ) + marker_family_enrichment: dict[str, float] = {} + protected_marker_families: list[str] = [] + for family, mask in family_masks.items(): + marker_mask = _family_mask(marker_names, family) + marker_fraction = ( + float(marker_mask.mean()) + if marker_mask is not None and marker_mask.size + else 0.0 + ) + background = float(mask.mean()) + enrichment = marker_fraction / background if background > 0 else 0.0 + if marker_ref is not None and family in nominated_families: + marker_family_enrichment[family] = enrichment + if family in protected_families and marker_fraction > 0: + protected_marker_families.append(family) + + selection_ref = ArtifactRef( + scope=evaluation.cellSelection.scope, + assay=evaluation.cellSelection.assay, + kind=evaluation.cellSelection.kind, + artifact_id=evaluation.cellSelection.artifactId, + ) + doublet_concentration = ( + diagnostic_call( + "metric.doubletConcentration", + _doublet_concentration, + store, + labels, + selection_ref, + doublet_evidence, + ) + if doublet_evidence is not None and doublet_evidence.scores + else None + ) + unit_scores = [ + score + for column in independent_unit_columns + if column in store.cells.columns + for score in [ + diagnostic_call( + "metric.crossUnitSupport", + _cross_unit_support, + labels, + _aligned_metadata(store, selection_ref, column), + ) + ] + if score is not None + ] + cross_unit_support = min(unit_scores) if unit_scores else None + technical_association = { + column: float( + diagnostic_call( + "metric.technicalAssociation", + normalized_mutual_info_score, + labels, + _aligned_metadata(store, selection_ref, column), + ) + ) + for column in technical_columns + if column in store.cells.columns + } + + metrics = evaluation.metrics.model_copy( + update={ + "seedStability": seed_stability, + "subsampleStability": subsample_stability, + "markerCoherence": marker_coherence, + "markerSpecificityMedian": marker_specificity_median, + "markerSpecificityByCluster": marker_specificity, + "markerAucByCluster": marker_auc, + "topMarkerGenes": top_marker_genes, + "crossUnitSupport": cross_unit_support, + "technicalAssociation": technical_association, + "markerFamilyEnrichment": marker_family_enrichment, + "protectedMarkerFamilies": protected_marker_families, + "doubletHighScoreConcentration": doublet_concentration, + "doubletScoreQuantiles": ( + dict(doublet_evidence.score_quantiles) + if doublet_evidence is not None + else {} + ), + "doubletScoreByCapture": ( + { + capture: dict(summary) + for capture, summary in zip( + doublet_evidence.capture_values, + doublet_evidence.score_summaries, + strict=True, + ) + } + if doublet_evidence is not None + else {} + ), + "doubletCaptureCoverage": ( + doublet_evidence.capture_coverage + if doublet_evidence is not None + else None + ), + } + ) + evidence_ids = [ + *evaluation.evidenceIds, + f"candidate:{evaluation.candidateId}:seedStability", + f"candidate:{evaluation.candidateId}:subsampleStability", + *( + [ + f"candidate:{evaluation.candidateId}:markerCoherence", + f"candidate:{evaluation.candidateId}:markerSpecificity", + f"candidate:{evaluation.candidateId}:markerFamilies", + ] + if marker_ref is not None + else [] + ), + *( + [f"candidate:{evaluation.candidateId}:crossUnitSupport"] + if cross_unit_support is not None + else [] + ), + *[ + f"candidate:{evaluation.candidateId}:technicalAssociation:{column}" + for column in technical_association + ], + *( + [f"candidate:{evaluation.candidateId}:doubletConcentration"] + if doublet_concentration is not None + else [] + ), + *( + [ + f"candidate:{evaluation.candidateId}:doubletScoreTails", + f"candidate:{evaluation.candidateId}:doubletCaptureCoverage", + ] + if doublet_evidence is not None and doublet_evidence.scores + else [] + ), + ] + artifacts = { + **evaluation.artifacts, + "stabilityClusters": ArtifactRecord.from_ref(alternative_ref), + **( + {"markerTable": ArtifactRecord.from_ref(marker_ref)} + if marker_ref is not None + else {} + ), + **( + { + f"doubletScore:{index}": ArtifactRecord.from_ref(score) + for index, score in enumerate(doublet_evidence.scores) + } + if doublet_evidence is not None + else {} + ), + **( + { + f"doubletCellSelection:{index}": ArtifactRecord.from_ref(selection) + for index, selection in enumerate(doublet_evidence.cell_selections) + } + if doublet_evidence is not None + else {} + ), + **( + { + "doubletNativeGraph": ArtifactRecord.from_ref( + doublet_evidence.native_graph + ), + "doubletNativeClusters": ArtifactRecord.from_ref( + doublet_evidence.native_clusters + ), + } + if doublet_evidence is not None + else {} + ), + } + augmented.append( + evaluation.model_copy( + update={ + "metrics": metrics, + "eligible": evaluation.eligible and cluster_count >= 2, + "eligibilityReasons": list( + dict.fromkeys( + [ + *evaluation.eligibilityReasons, + *( + [ + "Marker contrasts require at least two populated clusters" + ] + if cluster_count < 2 + else [] + ), + ] + ) + ), + "evidenceIds": list(dict.fromkeys(evidence_ids)), + "artifacts": artifacts, + "warnings": list( + dict.fromkeys( + [ + *evaluation.warnings, + *( + doublet_evidence.limitations + if doublet_evidence is not None + else () + ), + ] + ) + ), + } + ) + ) + return annotate_candidate_dominance(augmented) + + +__all__ = [ + "AdvisoryDoubletScores", + "augment_cluster_evaluations", + "augment_pca_evaluations", + "population_support_evidence", + "resolve_native_doublet_inputs", + "score_advisory_doublets", +] diff --git a/scarf/agent/parameter_tuning/execution.py b/scarf/agent/parameter_tuning/execution.py new file mode 100644 index 00000000..7340a9c9 --- /dev/null +++ b/scarf/agent/parameter_tuning/execution.py @@ -0,0 +1,1028 @@ +import hashlib +import json +from collections.abc import Callable, Iterator, Sequence +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Any, Literal, cast + +import numpy as np + +from ...metadata.rows import iter_metadata_column_blocks, metadata_missing_mask +from ...metrics import graph_connectivity +from ...storage.refs import ArtifactRef +from ...storage.types import as_zarr_array +from ...utils.logging import logger +from .._deps import AGENT_INSTALL_HINT +from ..tools import artifact_reference, core_artifact_reference +from .contracts import ( + ArtifactRecord, + IntegrationCandidateEvaluation, + ParameterCandidate, + ParameterCandidateEvaluation, + ParameterMetrics, + ParameterTuningDependencies, + ParameterTuningReport, +) + +try: + from pydantic_ai import RunContext +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + + +_RANDOM_SEED = 4444 +_PCA_RANDOM_SEED = 4466 + +_METRIC_CACHE: ContextVar[dict[tuple[Any, ...], Any] | None] = ContextVar( + "scarf_candidate_metric_cache", default=None +) +_DIAGNOSTIC_WORK: ContextVar[dict[str, dict[str, int]] | None] = ContextVar( + "scarf_diagnostic_work", default=None +) + + +@contextmanager +def diagnostic_work() -> Iterator[dict[str, dict[str, int]]]: + """Count agent operation calls, not internal core numerical rebuilds.""" + counts: dict[str, dict[str, int]] = {} + token = _DIAGNOSTIC_WORK.set(counts) + try: + yield counts + finally: + _DIAGNOSTIC_WORK.reset(token) + + +def _diagnostic_count(name: str, event: str, count: int = 1) -> None: + counts = _DIAGNOSTIC_WORK.get() + if counts is not None: + row = counts.setdefault( + name, + dict.fromkeys( + ( + "attempted", + "completed", + "failed", + "cacheHits", + "restored", + "artifactReuses", + ), + 0, + ), + ) + row[event] += count + + +def diagnostic_call[**P, T]( + name: str, operation: Callable[P, T], *args: P.args, **kwargs: P.kwargs +) -> T: + """Record an invoked operation while preserving its result and exception.""" + _diagnostic_count(name, "attempted") + try: + result = operation(*args, **kwargs) + except BaseException: + _diagnostic_count(name, "failed") + raise + _diagnostic_count(name, "completed") + return result + + +def diagnostic_reuse( + name: str, + kind: Literal["cacheHits", "restored", "artifactReuses"] = "cacheHits", + count: int = 1, +) -> None: + """Record reuse only when the agent has directly established it.""" + _diagnostic_count(name, kind, count) + + +@contextmanager +def candidate_metric_cache() -> Iterator[None]: + """Reuse exact-input metrics only for the current orchestration stage.""" + token = _METRIC_CACHE.set({}) + try: + yield + finally: + _METRIC_CACHE.reset(token) + + +def _cached_candidate_metric[T](key: tuple[Any, ...], compute: Callable[[], T]) -> T: + name = f"metric.{key[1]}" + cache = _METRIC_CACHE.get() + if cache is None: + return diagnostic_call(name, compute) + if key not in cache: + cache[key] = diagnostic_call(name, compute) + else: + diagnostic_reuse(name) + return cast(T, cache[key]) + + +def _metric_metadata_key(store: Any, column: str) -> str | None: + """Fingerprint live metric inputs so edits cannot reuse stale evidence.""" + if _METRIC_CACHE.get() is None: + return None + return _metadata_column_fingerprint(store.cells, column) + + +def _metadata_column_fingerprint(metadata: Any, column: str) -> str: + """Hash metadata values in bounded blocks, including scalar type identity.""" + digest = hashlib.sha256() + for block in iter_metadata_column_blocks(metadata, column): + digest.update(str(block.dtype).encode()) + digest.update(str(block.shape).encode()) + digest.update( + json.dumps( + [(type(value).__name__, repr(value)) for value in block.tolist()] + ).encode() + if block.dtype.hasobject + else block.tobytes() + ) + missing = metadata_missing_mask(metadata, column) + digest.update(b"missing:none" if missing is None else b"missing:present") + if missing is not None: + for start in range(0, len(missing), 65_536): + digest.update( + np.asarray(missing[start : start + 65_536], dtype=bool).tobytes() + ) + return digest.hexdigest() + + +def _final_graph_options( + report: ParameterTuningReport, + integration_evaluations: Sequence[IntegrationCandidateEvaluation], +) -> dict[str, dict[str, Any]]: + """Return the exact eligible graph options and option-scoped evidence.""" + + assay_reports = report.assayReports or {report.fromAssay: report} + report_cell_selection = core_artifact_reference(report.cellSelection) + if not isinstance(report_cell_selection, ArtifactRef): + raise ValueError("Parameter tuning report lacks an exact cell selection") + options: dict[str, dict[str, Any]] = {} + for assay, assay_report in assay_reports.items(): + candidate_id = assay_report.recommendedCandidateId + if candidate_id is None: + continue + native_evaluation = next( + ( + item + for item in assay_report.evaluations + if item.candidateId == candidate_id + ), + None, + ) + if ( + native_evaluation is None + or native_evaluation.status != "done" + or not native_evaluation.eligible + or "clusters" not in native_evaluation.artifacts + or "connectivityMap" not in native_evaluation.artifacts + or not native_evaluation.evidenceIds + ): + continue + if ( + core_artifact_reference(native_evaluation.cellSelection) + != report_cell_selection + ): + raise ValueError("Native graph option uses a different cell selection") + option_id = f"native:{assay}:{candidate_id}" + option_evidence = [ + f"native:{assay}:{evidence_id}" + for evidence_id in native_evaluation.evidenceIds + ] + evaluation_payload = native_evaluation.model_dump() + evaluation_payload["evidenceIds"] = option_evidence + options[option_id] = { + "optionId": option_id, + "graphMethod": "native", + "nativeAssay": assay, + "nativeCandidateId": candidate_id, + "evaluation": evaluation_payload, + "evidenceIds": option_evidence, + } + for integration_evaluation in integration_evaluations: + if ( + integration_evaluation.status != "done" + or not integration_evaluation.eligible + ): + continue + if ( + integration_evaluation.clusterArtifact is None + or integration_evaluation.graphArtifact is None + ): + continue + if not integration_evaluation.evidenceIds: + continue + if ( + core_artifact_reference(integration_evaluation.cellSelection) + != report_cell_selection + ): + raise ValueError("Integrated graph option uses a different cell selection") + if not integration_evaluation.integrationId: + raise ValueError("Eligible integration evaluations require integrationId") + if ( + integration_evaluation.graphArtifact.scope != "datastore" + or integration_evaluation.graphArtifact.assay is not None + or integration_evaluation.clusterArtifact.scope != "datastore" + or integration_evaluation.clusterArtifact.assay is not None + or integration_evaluation.graphArtifact.kind != "integrated_graph" + or integration_evaluation.clusterArtifact.kind + not in {"cluster_labels", "cluster_cut"} + ): + raise ValueError( + "Integrated graph and cluster artifacts must be datastore-scoped " + "without an assay" + ) + if ( + integration_evaluation.method == "wnn" + and integration_evaluation.metrics.modalityWeightsValid is not True + ): + continue + option_id = f"integration:{integration_evaluation.integrationId}" + if option_id in options: + raise ValueError( + f"Duplicate integration id {integration_evaluation.integrationId!r}" + ) + options[option_id] = { + "optionId": option_id, + "graphMethod": integration_evaluation.method, + "integrationId": integration_evaluation.integrationId, + "evaluation": integration_evaluation.model_dump(), + "evidenceIds": list(integration_evaluation.evidenceIds), + } + return options + + +def normalized_artifact_shape(store: Any, normalized: Any) -> tuple[int, int]: + """Return the exact cell-by-feature shape of a normalized artifact.""" + + group = store.load_artifact(normalized) + if "data" not in group: + raise ValueError("Normalized artifact does not contain a data matrix") + shape = getattr(group["data"], "shape", None) + if not isinstance(shape, tuple | list) or len(shape) != 2: + raise ValueError("Normalized artifact data must be two-dimensional") + n_cells, n_features = map(int, shape) + if n_cells < 2 or n_features < 2: + raise ValueError( + "Parameter tuning requires at least two cells and two selected features" + ) + return n_cells, n_features + + +def validate_parameter_candidate_rank( + candidate: ParameterCandidate, + normalized_shape: tuple[int, int], + *, + identity_feature_limit: int = 64, +) -> int: + """Validate a candidate before any branch operation and return output rank.""" + + n_cells, n_features = normalized_shape + if candidate.neighborsK >= n_cells: + raise ValueError( + f"neighborsK={candidate.neighborsK} requires more than " + f"{candidate.neighborsK} selected cells; observed {n_cells}" + ) + if candidate.reductionMethod == "pca": + if candidate.dimensions + 1 > min(n_cells, n_features): + raise ValueError( + f"PCA dimensions={candidate.dimensions} requires at least " + f"{candidate.dimensions + 1} cells and selected features; " + f"observed shape {normalized_shape}" + ) + return candidate.dimensions + if candidate.reductionMethod == "lsi": + required_rank = candidate.dimensions + 1 + if required_rank > min(n_cells, n_features): + raise ValueError( + "LSI dimensions, including the skipped component, exceed the " + f"normalized matrix rank for shape {normalized_shape}" + ) + return candidate.dimensions + if n_features > identity_feature_limit: + raise ValueError( + f"Identity reduction supports at most {identity_feature_limit} selected " + f"features; observed {n_features}" + ) + if candidate.dimensions != n_features: + raise ValueError( + "Identity reduction dimensions must equal the exact normalized feature " + f"count {n_features}; received {candidate.dimensions}" + ) + return n_features + + +def run_candidate_reduction( + store: Any, + *, + normalized: Any, + candidate: ParameterCandidate, + normalized_shape: tuple[int, int], + identity_feature_limit: int = 64, +) -> tuple[Any, str, int]: + """Run one validated modality-aware reduction with public Scarf methods.""" + + effective_dimensions = validate_parameter_candidate_rank( + candidate, + normalized_shape, + identity_feature_limit=identity_feature_limit, + ) + if candidate.reductionMethod == "pca": + ref = diagnostic_call( + "core.pca", + store.run_pca, + normalized, + dims=candidate.dimensions, + feat_scaling=True, + show_elbow_plot=False, + invalidate_cache=False, + ) + return ref, "pca", effective_dimensions + if candidate.reductionMethod == "lsi": + ref = diagnostic_call( + "core.lsi", + store.run_lsi, + normalized, + dims=candidate.dimensions, + skip_first=True, + rand_state=_PCA_RANDOM_SEED, + invalidate_cache=False, + ) + return ref, "lsi", effective_dimensions + loadings = np.eye(normalized_shape[1], dtype=np.float64) + ref = diagnostic_call( + "core.customReduction", + store.run_custom_reduction, + loadings, + normalized, + invalidate_cache=False, + ) + return ref, "identity", effective_dimensions + + +def _bounded_membership_summary( + values: Any, + labels: np.ndarray, + *, + maximum_sample_size: int = 65_536, +) -> tuple[float, float, float, dict[str, float], int]: + if len(values.shape) != 1 or values.shape != labels.shape: + raise ValueError("Membership strengths must align with cluster labels") + n_values = int(values.shape[0]) + if n_values < 1: + raise ValueError("Membership strengths cannot be empty") + stride = max(1, (n_values + maximum_sample_size - 1) // maximum_sample_size) + total = 0.0 + sampled_values: list[np.ndarray] = [] + sampled_labels: list[np.ndarray] = [] + for start in range(0, n_values, 65_536): + block = np.asarray(values[start : start + 65_536], dtype=np.float64) + if not np.isfinite(block).all(): + raise ValueError("Membership strengths must be finite") + total += float(block.sum()) + offset = (-start) % stride + sampled_values.append(block[offset::stride]) + sampled_labels.append(labels[start + offset : start + len(block) : stride]) + sample = np.concatenate(sampled_values) + sample_labels = np.concatenate(sampled_labels) + by_cluster = { + str(cluster): float(np.median(sample[sample_labels == cluster])) + for cluster in np.unique(sample_labels) + } + return ( + total / n_values, + float(np.median(sample)), + float(np.quantile(sample, 0.1)), + by_cluster, + int(len(sample)), + ) + + +def _collect_cluster_structure_metrics( + store: Any, + *, + cluster_ref: Any, + graph_ref: Any, + cluster_values: np.ndarray, + candidate_id: str, + metrics: ParameterMetrics, + evidence_ids: list[str], + warnings: list[str], +) -> ArtifactRef | None: + calculate_membership = getattr(store, "calc_membership_strength", None) + if not callable(calculate_membership): + return None + membership_ref: ArtifactRef | None = None + try: + membership_ref = diagnostic_call( + "core.membershipStrength", + calculate_membership, + cluster_ref, + graph_ref, + invalidate_cache=False, + ) + membership_group = store.load_artifact(membership_ref) + membership_values = as_zarr_array( + membership_group["values"], + name="values", + ) + mean, median, p10, by_cluster, sample_size = _cached_candidate_metric( + (id(store), "membership_summary", membership_ref, cluster_ref), + lambda: _bounded_membership_summary(membership_values, cluster_values), + ) + metrics.membershipStrengthMean = mean + metrics.membershipStrengthMedian = median + metrics.membershipStrengthP10 = p10 + metrics.membershipStrengthByCluster = by_cluster + metrics.membershipStrengthSampleSize = sample_size + evidence_ids.append(f"candidate:{candidate_id}:membershipStrength") + except (KeyError, TypeError, ValueError, RuntimeError) as exc: + warnings.append(f"Cluster membership strength unavailable: {exc}") + membership_ref = None + + try: + graph_group = store.load_artifact(graph_ref) + graph_edges = as_zarr_array(graph_group["edges"], name="edges") + connectivity = _cached_candidate_metric( + (id(store), "cluster_connectivity", graph_ref, cluster_ref), + lambda: float(graph_connectivity(graph_edges, cluster_values)), + ) + if np.isfinite(connectivity): + metrics.clusterConnectivity = connectivity + evidence_ids.append(f"candidate:{candidate_id}:clusterConnectivity") + except (KeyError, TypeError, ValueError, RuntimeError) as exc: + warnings.append(f"Cluster connectivity unavailable: {exc}") + return membership_ref + + +def _collect_parameter_candidate_metrics( + deps: ParameterTuningDependencies, + *, + candidate: ParameterCandidate, + candidate_id: str, + reduction_ref: Any, + neighbors_ref: Any, + graph_ref: Any, + cluster_ref: Any, + cluster_column: str, + evidence_ids: list[str], + warnings: list[str], +) -> tuple[ParameterMetrics, list[str], ArtifactRef | None]: + store = deps.store + cluster_group = store.load_artifact(cluster_ref) + cluster_data = cluster_group["values"] + cluster_values = np.asarray(cluster_data[:]) + if cluster_values.ndim != 1 or len(cluster_values) == 0: + raise ValueError("Cluster artifact must contain one non-empty label vector") + if np.any(cluster_values < 0): + raise ValueError("Cluster artifact contains invalid negative labels") + _, cluster_counts = np.unique(cluster_values, return_counts=True) + n_clusters = int(len(cluster_counts)) + min_cluster_cells = int(cluster_counts.min()) + min_cluster_fraction = float(min_cluster_cells / len(cluster_values)) + metrics = ParameterMetrics( + nClusters=n_clusters, + minClusterCells=min_cluster_cells, + minClusterFraction=min_cluster_fraction, + ) + evidence_ids.append(f"candidate:{candidate_id}:clusters") + membership_ref = _collect_cluster_structure_metrics( + store, + cluster_ref=cluster_ref, + graph_ref=graph_ref, + cluster_values=cluster_values, + candidate_id=candidate_id, + metrics=metrics, + evidence_ids=evidence_ids, + warnings=warnings, + ) + + try: + graph_scores = _cached_candidate_metric( + ( + id(store), + "graph_silhouette", + neighbors_ref, + cluster_ref, + _RANDOM_SEED, + 11, + ), + lambda: diagnostic_call( + "core.graphSilhouette", + store.metric_graph_silhouette, + neighbors_ref, + cluster_ref, + random_seed=_RANDOM_SEED, + sample_size=11, + ), + ) + if graph_scores is not None: + finite_scores = np.asarray(graph_scores, dtype=float) + finite_scores = finite_scores[np.isfinite(finite_scores)] + if len(finite_scores): + metrics.graphSilhouetteMedian = float(np.median(finite_scores)) + evidence_ids.append(f"candidate:{candidate_id}:graphSilhouette") + except (KeyError, TypeError, ValueError) as exc: + warnings.append(f"Graph silhouette unavailable: {exc}") + + if candidate.reductionMethod == "pca": + try: + + def separability_values() -> dict[str, Any]: + separability = diagnostic_call( + "core.clusterSeparability", + store.metric_cluster_separability, + reduction_ref, + {cluster_column: cluster_ref}, + random_seed=_RANDOM_SEED, + ) + table = separability.clustering_scores + rows = table.loc[table["clustering"] == cluster_column] + return dict(rows.iloc[0]) if len(rows) else {} + + row = _cached_candidate_metric( + ( + id(store), + "cluster_separability", + reduction_ref, + cluster_ref, + _RANDOM_SEED, + ), + separability_values, + ) + if row: + for field_name, column_name, evidence_name in ( + ("pcaSilhouette", "silhouette_score", "pcaSilhouette"), + ("macroF1", "macro_f1_mean", "macroF1"), + ("weightedF1", "weighted_f1_mean", "weightedF1"), + ): + value = row[column_name] + if value is not None and np.isfinite(float(value)): + setattr(metrics, field_name, float(value)) + evidence_ids.append(f"candidate:{candidate_id}:{evidence_name}") + except (KeyError, TypeError, ValueError) as exc: + warnings.append(f"PCA cluster separability unavailable: {exc}") + + _collect_covariate_metrics( + deps, + candidate=candidate, + candidate_id=candidate_id, + neighbors_ref=neighbors_ref, + graph_ref=graph_ref, + metrics=metrics, + evidence_ids=evidence_ids, + warnings=warnings, + ) + + eligibility_reasons: list[str] = [] + if n_clusters < 2: + eligibility_reasons.append("fewer than two clusters") + if min_cluster_cells < deps.minClusterCells: + eligibility_reasons.append( + f"smallest cluster has {min_cluster_cells} cells; " + f"minimum is {deps.minClusterCells}" + ) + return metrics, eligibility_reasons, membership_ref + + +def _collect_covariate_metrics( + deps: ParameterTuningDependencies, + *, + candidate: ParameterCandidate, + candidate_id: str, + neighbors_ref: ArtifactRef, + graph_ref: ArtifactRef, + metrics: ParameterMetrics, + evidence_ids: list[str], + warnings: list[str], +) -> None: + """Measure exact typed design effects on an existing graph and neighborhood.""" + store = deps.store + perplexity = max(1.0, float(candidate.neighborsK // 3)) + for column in deps.batchColumns: + try: + score = float( + _cached_candidate_metric( + ( + id(store), + "batch_mixing", + neighbors_ref, + column, + _metric_metadata_key(store, column), + perplexity, + ), + lambda: diagnostic_call( + "core.batchMixing", + store.metric_proportional_batch_mixing, + column, + neighbors_ref, + perplexity=perplexity, + ), + ) + ) + if np.isfinite(score): + metrics.batchMixing[column] = score + evidence_ids.append(f"candidate:{candidate_id}:batchMixing:{column}") + except (KeyError, TypeError, ValueError) as exc: + warnings.append(f"Batch mixing for {column!r} unavailable: {exc}") + + for column in deps.preservationColumns: + if deps.columnKinds.get(column) == "continuous": + warnings.append( + f"Matched graph preservation for continuous column {column!r} is unsupported; " + "PCA association is descriptive evidence only." + ) + continue + scores: dict[str, float] = {} + try: + clisi = float( + _cached_candidate_metric( + ( + id(store), + "clisi", + neighbors_ref, + column, + _metric_metadata_key(store, column), + None, + True, + ), + lambda: diagnostic_call( + "core.clisi", + store.metric_clisi, + column, + neighbors_ref, + perplexity=None, + scale=True, + ), + ) + ) + if np.isfinite(clisi): + scores["clisi"] = clisi + evidence_ids.append(f"candidate:{candidate_id}:clisi:{column}") + except (KeyError, TypeError, ValueError) as exc: + warnings.append(f"cLISI for {column!r} unavailable: {exc}") + try: + connectivity = float( + _cached_candidate_metric( + ( + id(store), + "protected_connectivity", + graph_ref, + column, + _metric_metadata_key(store, column), + ), + lambda: diagnostic_call( + "core.graphConnectivity", + store.metric_graph_connectivity, + column, + graph_ref, + ), + ) + ) + if np.isfinite(connectivity): + scores["graphConnectivity"] = connectivity + evidence_ids.append( + f"candidate:{candidate_id}:graphConnectivity:{column}" + ) + except (KeyError, TypeError, ValueError) as exc: + warnings.append(f"Graph connectivity for {column!r} unavailable: {exc}") + if scores: + metrics.biologicalPreservation[column] = scores + + if deps.protectedCombinations: + from ...metrics import clisi_knn + from ..experimental_context.characterization import _SelectionBoundCells + from ..experimental_context.comparisons import combination_labels + + bound_cells = _SelectionBoundCells(store.zw, store.cells, deps.cellSelection) + neighbor_group = store.load_artifact(neighbors_ref) + graph_group = store.load_artifact(graph_ref) + for columns in deps.protectedCombinations: + name = "joint:" + json.dumps(list(columns), separators=(",", ":")) + metadata_key = tuple( + _metric_metadata_key(store, column) for column in columns + ) + labels = combination_labels(bound_cells, columns) + scores = _cached_candidate_metric( + ( + id(store), + "protected_combination", + neighbors_ref, + graph_ref, + columns, + metadata_key, + ), + lambda: { + "clisi": float( + clisi_knn( + as_zarr_array( + neighbor_group["distances"], name="distances" + ), + as_zarr_array(neighbor_group["indices"], name="indices"), + labels, + perplexity=None, + scale=True, + ) + ), + "graphConnectivity": float( + graph_connectivity( + as_zarr_array(graph_group["edges"], name="edges"), + labels, + ) + ), + }, + ) + if not all(np.isfinite(value) for value in scores.values()): + raise ValueError("Protected combination metrics must be finite") + metrics.biologicalPreservation[name] = scores + evidence_ids.append(f"candidate:{candidate_id}:{name}") + + +def refresh_candidate_design_evidence( + deps: ParameterTuningDependencies, + evaluation: ParameterCandidateEvaluation, +) -> ParameterCandidateEvaluation: + """Reassess a revised design without repeating primary analysis or doublets.""" + evaluation = evaluation.model_copy(deep=True) + evaluation.metrics.batchMixing = {} + evaluation.metrics.biologicalPreservation = {} + identifiers = (":batchMixing:", ":clisi:", ":graphConnectivity:", ":joint:") + evaluation.evidenceIds = [ + item + for item in evaluation.evidenceIds + if not any(identifier in item for identifier in identifiers) + ] + prefixes = ( + "Batch mixing for ", + "cLISI for ", + "Graph connectivity for ", + "Matched graph preservation for continuous column ", + ) + evaluation.warnings = [ + item for item in evaluation.warnings if not item.startswith(prefixes) + ] + _collect_covariate_metrics( + deps, + candidate=evaluation.parameters, + candidate_id=evaluation.candidateId, + neighbors_ref=core_artifact_reference(evaluation.artifacts["neighbors"]), + graph_ref=core_artifact_reference(evaluation.artifacts["connectivityMap"]), + metrics=evaluation.metrics, + evidence_ids=evaluation.evidenceIds, + warnings=evaluation.warnings, + ) + return evaluation + + +def execute_parameter_candidate( + deps: ParameterTuningDependencies, + candidate_id: str, +) -> ParameterCandidateEvaluation: + """Execute one allowlisted candidate without model involvement.""" + + with deps.executionLock: + if candidate_id in deps.evaluations: + diagnostic_reuse("candidate.evaluation") + logger.debug( + f"Parameter candidate {candidate_id!r} for assay " + f"{deps.fromAssay!r} reused its completed evaluation" + ) + return deps.evaluations[candidate_id] + if candidate_id not in deps.candidates: + logger.warning( + f"Parameter candidate {candidate_id!r} is not authorized for " + f"assay {deps.fromAssay!r}" + ) + return ParameterCandidateEvaluation( + candidateId=candidate_id, + status="failed", + error=( + f"Unknown candidate id {candidate_id!r}; allowed ids are " + f"{sorted(deps.candidates)}" + ), + ) + if len(deps.executionOrder) >= deps.maxCandidates: + logger.warning( + f"Parameter candidate {candidate_id!r} was not executed because " + f"assay {deps.fromAssay!r} reached its limit of " + f"{deps.maxCandidates} candidates" + ) + return ParameterCandidateEvaluation( + candidateId=candidate_id, + phase=deps.candidatePhases.get(candidate_id, "initial"), + harmonyBatchColumns=( + list(deps.batchColumns) + if deps.candidates[candidate_id].useHarmony + else [] + ), + status="failed", + parameters=deps.candidates[candidate_id], + error=f"Candidate execution limit {deps.maxCandidates} reached", + ) + + candidate = deps.candidates[candidate_id] + deps.executionOrder.append(candidate_id) + logger.debug(f"Executing candidate {candidate_id!r} for {deps.fromAssay!r}") + logger.info( + f"Comparing settings for {deps.fromAssay}: {candidate.reductionMethod.upper()} " + f"dimensions={candidate.dimensions}, neighbors={candidate.neighborsK}, " + f"resolution={candidate.leidenResolution}, " + f"harmony={candidate.useHarmony}" + ) + if candidate.useHarmony and not deps.batchColumns: + logger.warning( + f"Parameter candidate {candidate_id!r} cannot run Harmony because " + "no batch columns were authorized" + ) + evaluation = ParameterCandidateEvaluation( + candidateId=candidate_id, + phase=deps.candidatePhases.get(candidate_id, "initial"), + harmonyBatchColumns=[], + status="failed", + parameters=candidate, + error="Harmony candidate requires at least one authorized batch column", + ) + deps.evaluations[candidate_id] = evaluation + return evaluation + + store = deps.store + artifacts: dict[str, ArtifactRecord] = {} + warnings: list[str] = [] + evidence_ids: list[str] = [] + cluster_label = f"agent_tuning_{candidate_id}" + cluster_column = f"{deps.fromAssay}_{cluster_label}" + + try: + normalized_shape = deps.normalizedShape or normalized_artifact_shape( + store, + deps.normalized, + ) + effective_dimensions = validate_parameter_candidate_rank( + candidate, + normalized_shape, + identity_feature_limit=deps.identityFeatureLimit, + ) + reduction_ref, reduction_key, _ = run_candidate_reduction( + store, + normalized=deps.normalized, + candidate=candidate, + normalized_shape=normalized_shape, + identity_feature_limit=deps.identityFeatureLimit, + ) + artifacts[reduction_key] = ArtifactRecord.from_ref(reduction_ref) + logger.debug( + f"Parameter candidate {candidate_id!r}: completed " + f"{reduction_key} reduction" + ) + + coordinates_ref = reduction_ref + if candidate.useHarmony: + coordinates_ref = diagnostic_call( + "core.harmony", + store.run_harmony, + reduction_ref, + list(deps.batchColumns), + invalidate_cache=False, + ) + artifacts["harmony"] = ArtifactRecord.from_ref(coordinates_ref) + logger.debug( + f"Parameter candidate {candidate_id!r}: completed Harmony " + f"using {len(deps.batchColumns)} batch column(s)" + ) + + ann_ref = diagnostic_call( + "core.ann", + store.build_ann_index, + coordinates_ref, + ann_metric="l2", + ann_parallel=False, + rand_state=_PCA_RANDOM_SEED, + invalidate_cache=False, + ) + artifacts["annIndex"] = ArtifactRecord.from_ref(ann_ref) + logger.debug( + f"Parameter candidate {candidate_id!r}: completed ANN indexing" + ) + + neighbors_ref = diagnostic_call( + "core.neighbors", + store.query_neighbors, + ann_ref, + coordinates=coordinates_ref, + k=candidate.neighborsK, + invalidate_cache=False, + ) + artifacts["neighbors"] = ArtifactRecord.from_ref(neighbors_ref) + logger.debug( + f"Parameter candidate {candidate_id!r}: completed neighbor query" + ) + + graph_ref = diagnostic_call( + "core.graph", + store.build_connectivity_map, + neighbors_ref, + local_connectivity=1.0, + bandwidth=1.5, + invalidate_cache=False, + ) + artifacts["connectivityMap"] = ArtifactRecord.from_ref(graph_ref) + logger.debug( + f"Parameter candidate {candidate_id!r}: completed connectivity map" + ) + + cluster_ref = diagnostic_call( + "core.partition", + store.run_leiden_clustering, + graph_ref, + resolution=candidate.leidenResolution, + backend="igraph", + symmetric_graph=False, + graph_upper_only=False, + random_seed=_RANDOM_SEED, + invalidate_cache=False, + ) + artifacts["clusters"] = ArtifactRecord.from_ref(cluster_ref) + logger.debug( + f"Parameter candidate {candidate_id!r}: completed Leiden clustering" + ) + + ( + metrics, + eligibility_reasons, + membership_ref, + ) = _collect_parameter_candidate_metrics( + deps, + candidate=candidate, + candidate_id=candidate_id, + reduction_ref=reduction_ref, + neighbors_ref=neighbors_ref, + graph_ref=graph_ref, + cluster_ref=cluster_ref, + cluster_column=cluster_column, + evidence_ids=evidence_ids, + warnings=warnings, + ) + if membership_ref is not None: + artifacts["membershipStrength"] = ArtifactRecord.from_ref( + membership_ref + ) + + evaluation = ParameterCandidateEvaluation( + candidateId=candidate_id, + phase=deps.candidatePhases.get(candidate_id, "initial"), + harmonyBatchColumns=( + list(deps.batchColumns) if candidate.useHarmony else [] + ), + status="done", + eligible=not eligibility_reasons, + parameters=candidate, + artifacts=artifacts, + cellSelection=artifact_reference(deps.cellSelection), + clusterColumn=cluster_column, + clusterLabel=cluster_label, + effectiveDimensions=effective_dimensions, + metrics=metrics, + evidenceIds=evidence_ids, + eligibilityReasons=eligibility_reasons, + warnings=warnings, + ) + logger.info( + f"Compared settings: {metrics.nClusters} clusters; " + f"smallest population has {metrics.minClusterCells} cells; " + f"{'ready for assessment' if evaluation.eligible else 'candidate checks unresolved'}." + ) + except (KeyError, TypeError, ValueError, RuntimeError) as exc: + evaluation = ParameterCandidateEvaluation( + candidateId=candidate_id, + phase=deps.candidatePhases.get(candidate_id, "initial"), + harmonyBatchColumns=( + list(deps.batchColumns) if candidate.useHarmony else [] + ), + status="failed", + parameters=candidate, + artifacts=artifacts, + cellSelection=( + artifact_reference(deps.cellSelection) + if deps.cellSelection is not None + else None + ), + evidenceIds=evidence_ids, + warnings=warnings, + error=str(exc), + ) + logger.warning( + f"Parameter candidate {candidate_id!r} for assay " + f"{deps.fromAssay!r} failed: {exc}" + ) + + deps.evaluations[candidate_id] = evaluation + return evaluation + + +async def evaluate_parameter_candidate( + ctx: RunContext[ParameterTuningDependencies], + candidate_id: str, +) -> ParameterCandidateEvaluation: + """Expose deterministic candidate execution as a bounded agent tool.""" + + return execute_parameter_candidate(ctx.deps, candidate_id) diff --git a/scarf/agent/parameter_tuning/hvg.py b/scarf/agent/parameter_tuning/hvg.py new file mode 100644 index 00000000..3a181708 --- /dev/null +++ b/scarf/agent/parameter_tuning/hvg.py @@ -0,0 +1,248 @@ +"""Core-backed feature selection and bounded technical-group rank aggregation.""" + +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from typing import Any, Literal, cast + +import numpy as np + +from ...storage.refs import ArtifactRef + +HVG_CANDIDATE_TARGETS = (1000, 2000, 4000) + + +def core_hvg_evidence( + store: Any, + *, + assay: str, + cells: ArtifactRef, +) -> dict[str, ArtifactRef]: + """Obtain baseline, eligible universes and variability from core Scarf. + + These are feature-axis computations. Requesting all eligible genes does + not normalize, reduce or construct a graph for another candidate. + """ + feature_count = int(store.get_assay(assay).feats.N) + options = {"from_assay": assay, "show_plot": False, "invalidate_cache": False} + return { + "scarfDefault": store.select_hvgs(cells, top_n=1000, **options), + "eligibleDefault": store.select_hvgs(cells, top_n=feature_count, **options), + "eligibleAll": store.select_hvgs( + cells, top_n=feature_count, blacklist="", **options + ), + } + + +def rank_core_hvgs( + store: Any, + *, + eligible: ArtifactRef, + statistics: ArtifactRef, + top_n: int, + ranking: np.ndarray | None = None, +) -> ArtifactRef: + """Select a count on a frozen universe using the core variability payload.""" + mask = np.asarray(store.load_artifact(eligible)["values"][:], dtype=bool) + variance = np.asarray( + store.load_artifact(statistics)["corrected_variance"][:], dtype=np.float64 + ) + if mask.shape != variance.shape or not np.isfinite(variance).all(): + raise ValueError("Core HVG statistics do not align with eligible genes") + indices = np.flatnonzero(mask) + if ranking is None: + indices = indices[np.lexsort((indices, -variance[indices]))] + else: + ordered = np.asarray(ranking, dtype=np.int64) + if ordered.ndim != 1 or len(np.unique(ordered)) != len(ordered): + raise ValueError("HVG ranking must contain unique feature indices") + if np.any(ordered < 0) or np.any(ordered >= len(mask)): + raise ValueError("HVG ranking contains invalid feature indices") + indices = ordered[mask[ordered]] + if len(indices) != int(mask.sum()): + raise ValueError("HVG ranking must cover the exact eligible universe") + if isinstance(top_n, bool) or top_n < 3: + raise ValueError("RNA representation needs at least three requested genes") + selected = np.zeros(mask.shape, dtype=bool) + selected[indices[:top_n]] = True + if int(selected.sum()) < 3: + raise ValueError("Fewer than three eligible genes remain") + return cast( + ArtifactRef, + store.set_feature_selection( + from_assay=eligible.assay, mask=selected, invalidate_cache=False + ), + ) + + +@dataclass(frozen=True, slots=True) +class HvgGroupVariability: + """One technical group's feature variability, streamed into aggregation.""" + + group_id: str + cell_count: int + corrected_variance: np.ndarray + detected_features: np.ndarray + + +@dataclass(frozen=True, slots=True) +class HvgRanking: + """Bounded feature-axis output of global or batch-aware HVG ranking.""" + + ranking_mode: Literal["global", "batchAware"] + eligible: np.ndarray + global_corrected_variance: np.ndarray + recurrence: np.ndarray + mean_within_group_rank: np.ndarray + ranking: np.ndarray + valid_group_count: int + candidate_counts: tuple[int, ...] + + @property + def eligible_feature_count(self) -> int: + return int(self.eligible.sum()) + + def candidate_mask(self, top_n: int) -> np.ndarray: + """Return one registered nested candidate, never an arbitrary count.""" + if top_n not in self.candidate_counts: + raise ValueError( + f"top_n must be one of the registered counts {self.candidate_counts}" + ) + values = np.zeros(self.eligible.shape, dtype=bool) + values[self.ranking[:top_n]] = True + return values + + +def effective_hvg_candidate_counts( + eligible_feature_count: int, + targets: Sequence[int] = HVG_CANDIDATE_TARGETS, +) -> tuple[int, ...]: + """Cap registered HVG counts by eligibility and remove capped duplicates.""" + if isinstance(eligible_feature_count, bool) or not isinstance( + eligible_feature_count, int + ): + raise TypeError("eligible_feature_count must be an integer") + if eligible_feature_count < 1: + raise ValueError("eligible_feature_count must be greater than 0") + if isinstance(targets, str | bytes): + raise TypeError("targets must be a sequence of positive integers") + resolved: list[int] = [] + for target in targets: + if isinstance(target, bool) or not isinstance(target, int): + raise TypeError("HVG candidate targets must be integers") + if target < 1: + raise ValueError("HVG candidate targets must be greater than 0") + effective = min(target, eligible_feature_count) + if effective not in resolved: + resolved.append(effective) + if not resolved: + raise ValueError("At least one HVG candidate target is required") + return tuple(resolved) + + +def aggregate_hvg_rankings( + global_corrected_variance: np.ndarray, + eligible_features: np.ndarray, + group_variability: Iterable[HvgGroupVariability], + *, + valid_group_count: int, + candidate_targets: Sequence[int] = HVG_CANDIDATE_TARGETS, +) -> HvgRanking: + """Aggregate global and optional batch-aware feature rankings.""" + corrected = np.asarray(global_corrected_variance, dtype=np.float64) + eligible = np.asarray(eligible_features, dtype=bool) + if corrected.ndim != 1 or eligible.shape != corrected.shape: + raise ValueError("Global variability and eligibility must be aligned vectors") + if not np.isfinite(corrected).all() or (corrected < 0).any(): + raise ValueError("Global corrected variability must be finite and non-negative") + if isinstance(valid_group_count, bool) or not isinstance(valid_group_count, int): + raise TypeError("valid_group_count must be an integer") + if valid_group_count < 0: + raise ValueError("valid_group_count must be non-negative") + eligible_count = int(eligible.sum()) + counts = effective_hvg_candidate_counts(eligible_count, candidate_targets) + indices = np.flatnonzero(eligible) + global_order = indices[np.lexsort((indices, -corrected[indices]))].astype( + np.int64, copy=False + ) + recurrence = np.zeros(corrected.shape, dtype=np.int32) + mean_rank = np.full(corrected.shape, np.inf, dtype=np.float64) + + if valid_group_count < 2: + denominator = max(1, len(global_order)) + mean_rank[global_order] = ( + np.arange(1, len(global_order) + 1, dtype=np.float64) / denominator + ) + return HvgRanking( + ranking_mode="global", + eligible=eligible.copy(), + global_corrected_variance=corrected.copy(), + recurrence=recurrence, + mean_within_group_rank=mean_rank, + ranking=global_order.copy(), + valid_group_count=valid_group_count, + candidate_counts=counts, + ) + + rank_sum = np.zeros(corrected.shape, dtype=np.float64) + broad_count = max(counts) + received = 0 + for group in group_variability: + received += 1 + if received > valid_group_count: + raise ValueError("More group summaries were supplied than declared") + if not isinstance(group.group_id, str) or not group.group_id: + raise ValueError("Every valid technical group needs a non-empty ID") + if isinstance(group.cell_count, bool) or not isinstance(group.cell_count, int): + raise TypeError("Technical-group cell counts must be integers") + if group.cell_count < 1: + raise ValueError("Technical-group cell counts must be greater than 0") + group_corrected = np.asarray(group.corrected_variance, dtype=np.float64) + detected = np.asarray(group.detected_features, dtype=bool) + if ( + group_corrected.shape != corrected.shape + or detected.shape != corrected.shape + ): + raise ValueError("Technical-group feature arrays must align globally") + if not np.isfinite(group_corrected).all() or (group_corrected < 0).any(): + raise ValueError( + "Technical-group variability must be finite and non-negative" + ) + group_candidates = np.flatnonzero(eligible & detected) + if group_candidates.size == 0: + raise ValueError( + f"Valid technical group {group.group_id!r} has no rankable features" + ) + order = group_candidates[ + np.lexsort((group_candidates, -group_corrected[group_candidates])) + ] + selected = order[: min(broad_count, len(order))] + recurrence[selected] += 1 + rank_sum[selected] += np.arange(1, len(selected) + 1, dtype=np.float64) / len( + order + ) + if received != valid_group_count: + raise ValueError( + f"Expected {valid_group_count} group summaries but received {received}" + ) + observed = recurrence > 0 + mean_rank[observed] = rank_sum[observed] / recurrence[observed] + ranking = indices[ + np.lexsort( + ( + indices, + -corrected[indices], + mean_rank[indices], + -recurrence[indices], + ) + ) + ].astype(np.int64, copy=False) + return HvgRanking( + ranking_mode="batchAware", + eligible=eligible.copy(), + global_corrected_variance=corrected.copy(), + recurrence=recurrence, + mean_within_group_rank=mean_rank, + ranking=ranking.copy(), + valid_group_count=valid_group_count, + candidate_counts=counts, + ) diff --git a/scarf/agent/parameter_tuning/prompts.py b/scarf/agent/parameter_tuning/prompts.py new file mode 100644 index 00000000..f3ac9bc4 --- /dev/null +++ b/scarf/agent/parameter_tuning/prompts.py @@ -0,0 +1,447 @@ +import json +from collections.abc import Mapping, Sequence +from textwrap import dedent +from typing import Any + +from ..types import ArtifactReferenceModel +from .contracts import ( + IntegrationCandidateEvaluation, + ParameterCandidate, + ParameterCandidateEvaluation, + ParameterSearchPlan, + ParameterTuningDependencies, + ParameterTuningReport, + _default_parameter_candidates, +) +from .execution import _final_graph_options + + +def get_default_parameter_candidates() -> list[ParameterCandidate]: + """Return a small one-factor candidate set around Scarf defaults.""" + + return _default_parameter_candidates() + + +def build_initial_parameter_candidates( + candidates: Sequence[ParameterCandidate], + *, + pair_harmony: bool, +) -> list[ParameterCandidate]: + """Build deterministic initial branches from caller-authorized parameters.""" + + initial: list[ParameterCandidate] = [] + for candidate in candidates: + if pair_harmony and candidate.useHarmony: + raise ValueError( + "Initial seed candidates must not set useHarmony when the " + "experimental handoff controls Harmony pairing" + ) + initial.append(candidate) + if pair_harmony: + payload = candidate.model_dump() + payload.update( + { + "candidateId": f"{candidate.candidateId}_harmony", + "useHarmony": True, + } + ) + initial.append(ParameterCandidate.model_validate(payload)) + return initial + + +def parameter_search_system_prompt() -> str: + """Build the stable prompt for the bounded refinement-planning call.""" + + return dedent( + """ + You are planning one bounded refinement pass for Scarf parameter tuning. + The initial candidate screen has already finished. Do not request tools or + claim that additional candidates ran. + + Return exactly one of these two plan shapes: + 1. status=complete with candidates=[] when the initial screen is sufficient. + 2. status=refine with one or more candidates when an untested candidate + inside the initial numeric search envelope can resolve a specific + evidence-backed uncertainty. + Never return status=complete with candidates. A Harmony candidate always + uses the exact authorized batch columns supplied in the prompt. You may + choose between no correction and that approved Harmony configuration, but + you must not propose or modify batch columns. When proposing any Harmony + refinement, base it on one matched corrected and uncorrected initial pair + with otherwise identical parameters. + + Cite only evidenceIds from the initial evaluations. Identify the successful + initial candidates that motivate refinement, state focused objectives, and + provide concrete stopping criteria. Do not invent metrics, artifacts, or + candidate ids. Treat pcaSilhouette, macroF1, and weightedF1 only as PCA + cluster-separability metrics. Biological preservation evidence exists only + in a non-empty biologicalPreservation map. Check every exact value before + stating a ranking or trend, and keep narrative fields as plain prose + without serialized JSON. + """ + ).strip() + + +def parameter_evaluation_payload( + evaluation: ParameterCandidateEvaluation, +) -> dict[str, Any]: + """Return only candidate evidence needed for planning and selection.""" + metrics = evaluation.metrics.model_dump(mode="json") + loading_items = list(evaluation.metrics.topLoadingGenes.items()) + bounded_loading_items = [ + *loading_items[:10], + *(loading_items[-3:] if len(loading_items) > 13 else loading_items[10:]), + ] + metrics["topLoadingGenes"] = { + component: genes[:10] for component, genes in bounded_loading_items + } + metrics["topMarkerGenes"] = { + cluster: genes[:10] + for cluster, genes in list(evaluation.metrics.topMarkerGenes.items())[:30] + } + return { + "candidateId": evaluation.candidateId, + "phase": evaluation.phase, + "harmonyBatchColumns": evaluation.harmonyBatchColumns, + "status": evaluation.status, + "eligible": evaluation.eligible, + "parameters": evaluation.parameters.model_dump(mode="json"), + "effectiveDimensions": evaluation.effectiveDimensions, + "metrics": metrics, + "evidenceIds": evaluation.evidenceIds, + "eligibilityReasons": evaluation.eligibilityReasons, + "warnings": [warning[:500] for warning in evaluation.warnings[:10]], + "error": evaluation.error[:500] if evaluation.error is not None else None, + } + + +def parameter_search_prompt( + *, + from_assay: str, + cell_selection: ArtifactReferenceModel, + evaluations: Sequence[ParameterCandidateEvaluation], + batch_columns: Sequence[str], + preservation_columns: Sequence[str], + harmony_authorized: bool, + max_refined_candidates: int, +) -> str: + """Build the planning prompt from deterministic initial evaluations.""" + + evaluation_payload = [ + parameter_evaluation_payload(evaluation) for evaluation in evaluations + ] + correction_modes = ["none", "harmony"] if harmony_authorized else ["none"] + return ( + dedent( + """ + Inspect the completed initial screen for assay {from_assay} and exact + cell-selection artifact {cell_selection}. + + Initial evaluations: + {evaluation_payload} + + Authorized correction modes: {correction_modes} + Exact Harmony batch columns: {batch_columns} + Trusted biological preservation columns: {preservation_columns} + Maximum refined candidates: {max_refined_candidates} + + Return one ParameterSearchPlan. Refinement is optional and is limited to + one deterministic follow-up pass. + """ + ) + .strip() + .format( + from_assay=from_assay, + cell_selection=cell_selection.artifactId, + evaluation_payload=json.dumps( + evaluation_payload, + indent=2, + sort_keys=True, + ), + correction_modes=json.dumps(correction_modes), + batch_columns=json.dumps(list(batch_columns)), + preservation_columns=json.dumps(list(preservation_columns)), + max_refined_candidates=max_refined_candidates, + ) + ) + + +def parameter_tuning_system_prompt(min_cluster_cells: int) -> str: + """Build the stable prompt for final candidate selection.""" + + return ( + dedent( + """ + You are Scarf's parameter tuning selection agent. Every candidate in the + prompt has already finished deterministic execution. Do not request tools + or claim that another candidate ran. + + Recommend only a candidate whose evaluation has status=done and + eligible=true. A candidate is ineligible when it creates fewer than two + clusters or a cluster with fewer than {min_cluster_cells} cells. Do not + invent artifact ids, metrics, candidate ids, or evidence ids. Cite only + evidenceIds recorded in the completed evaluations. + + Balance cluster separation, cluster sizes, batch mixing, and biological + preservation. High batch mixing alone can indicate overcorrection, so do + not collapse the metrics into an invented score. UMAP appearance is not + evidence for parameter quality. Treat pcaSilhouette, macroF1, and + weightedF1 only as PCA cluster-separability metrics. Biological + preservation evidence exists only in a non-empty biologicalPreservation + map. A candidate with non-empty dominatedByCandidateIds is Pareto + dominated. Selecting a dominated graph or resolution requires at least + two independent non-geometric evidence classes that explain the + tradeoff. Do not call any metric highest, lowest, improved, degraded, or + monotonic without checking its exact value across every relevant + candidate. Narrative fields contain plain prose only and must not contain + serialized JSON keys or objects. When multiple candidates complete, + return one comparison for every non-selected successful candidate. Each + comparison must cite evidence from both the selected candidate and that + comparator. Return only model-owned selection fields. Leave evaluations, + selectedArtifacts, searchPlan, assayReports, integration fields, final + graph fields, and runInfo out of the response because validation fills + them from executor state. Return a concise structured report. + """ + ) + .strip() + .format(min_cluster_cells=min_cluster_cells) + ) + + +def parameter_tuning_prompt( + *, + from_assay: str, + cell_selection: ArtifactReferenceModel, + evaluations: Sequence[ParameterCandidateEvaluation], + batch_columns: Sequence[str], + preservation_columns: Sequence[str], + search_plan: ParameterSearchPlan, +) -> str: + """Build the final selection prompt from completed evaluations.""" + + evaluation_payload = [ + parameter_evaluation_payload(evaluation) for evaluation in evaluations + ] + return ( + dedent( + """ + Select a completed candidate for assay {from_assay} and exact + cell-selection artifact {cell_selection}. + + Completed evaluations: + {evaluation_payload} + + Validated refinement plan: + {search_plan} + + Exact Harmony batch columns: {batch_columns} + Trusted biological preservation columns: {preservation_columns} + + Recommend one eligible candidate or explain why user input is needed. + Compare the recommendation with every other successful candidate. High + batch mixing does not by itself justify correction when biological + preservation declines. + """ + ) + .strip() + .format( + from_assay=from_assay, + cell_selection=cell_selection.artifactId, + evaluation_payload=json.dumps( + evaluation_payload, + indent=2, + sort_keys=True, + ), + search_plan=json.dumps( + search_plan.model_dump(exclude={"runInfo"}), + indent=2, + sort_keys=True, + ), + batch_columns=json.dumps(list(batch_columns)), + preservation_columns=json.dumps(list(preservation_columns)), + ) + ) + + +def parameter_batch_search_prompt( + dependencies: Mapping[str, ParameterTuningDependencies], + max_refined_by_assay: Mapping[str, int], +) -> str: + """Build one refinement prompt for all modality-specific screens.""" + + payload = { + assay: { + "evaluations": [ + parameter_evaluation_payload(deps.evaluations[candidate_id]) + for candidate_id in deps.executionOrder + ], + "authorizedHarmony": deps.harmonyAuthorized, + "batchColumns": list(deps.batchColumns), + "preservationColumns": list(deps.preservationColumns), + "maxRefinedCandidates": max_refined_by_assay[assay], + } + for assay, deps in dependencies.items() + } + return ( + dedent( + """ + Plan one optional refinement pass for every assay in this completed + multimodal initial screen: + {payload} + + Return exactly one assayPlans entry for every assay. Each entry must + obey the single-assay ParameterSearchPlan rules. Candidate ids need + only be unique within their assay. Do not compare metric fields that + are absent for a modality, and do not request additional tool calls. + """ + ) + .strip() + .format(payload=json.dumps(payload, indent=2, sort_keys=True)) + ) + + +def parameter_batch_search_system_prompt() -> str: + """Build the stable system prompt for batched refinement planning.""" + + return ( + dedent( + """ + {single_assay_rules} + + Return the plans together in one assayPlans mapping. + """ + ) + .strip() + .format(single_assay_rules=parameter_search_system_prompt()) + ) + + +def parameter_batch_selection_system_prompt() -> str: + """Build the stable system prompt for batched native selection.""" + + return ( + dedent( + """ + You are Scarf's batched native parameter selection agent. Every branch + has already executed. Return one aggregate ParameterTuningReport with + exactly one grounded single-assay report in assayReports per assay. + Apply eligibility, evidence, and comparison requirements independently. + Do not invent joint scores, artifacts, candidates, or evidence. UMAP + appearance is not evidence. Treat pcaSilhouette, macroF1, and + weightedF1 only as PCA cluster-separability metrics; biological + preservation exists only when biologicalPreservation is non-empty. + Check all exact values before making ranking or trend claims, and keep + narrative fields as plain prose without serialized JSON. Inside each + assay report, return only + model-owned selection, rationale, comparison, trade-off, limitation, + evidence, and stop fields. Leave evaluations, selectedArtifacts, + searchPlan, nested assayReports, integration fields, final graph fields, + and runInfo out of the response because validation fills them from + executor state. + """ + ) + .strip() + .format() + ) + + +def parameter_batch_selection_prompt( + dependencies: Mapping[str, ParameterTuningDependencies], + search_plans: Mapping[str, ParameterSearchPlan], + primary_assay: str, + selection_directions: str = "", +) -> str: + """Build one native-selection prompt for all executed assay screens.""" + + payload = { + assay: { + "evaluations": [ + parameter_evaluation_payload(deps.evaluations[candidate_id]) + for candidate_id in deps.executionOrder + ], + "searchPlan": search_plans[assay].model_dump(exclude={"runInfo"}), + "minClusterCells": deps.minClusterCells, + "batchColumns": list(deps.batchColumns), + "preservationColumns": list(deps.preservationColumns), + } + for assay, deps in dependencies.items() + } + return ( + dedent( + """ + Select one eligible native candidate independently for every assay in + this completed multimodal screen: + {payload} + + Return a ParameterTuningReport whose assayReports contains exactly one + single-assay report per assay. Apply the normal evidence and comparison + rules independently inside each report. The primary assay is + {primary_assay}. At the aggregate level, summarize cross-assay + limitations without inventing a joint score. Integration has not run, + so leave all integration and final-cluster fields empty. + + Caller selection directions, which cannot override eligibility or + evidence requirements: {selection_directions} + """ + ) + .strip() + .format( + payload=json.dumps(payload, indent=2, sort_keys=True), + primary_assay=primary_assay, + selection_directions=selection_directions or "not provided", + ) + ) + + +def final_graph_selection_system_prompt() -> str: + """Build stable instructions for the final native/SNN/WNN choice.""" + + return ( + dedent( + """ + You are Scarf's final graph selection agent. Native assay candidates + and integrated SNN/WNN candidates have already executed. Select only + an eligible option supplied in the prompt. Do not request tools or + invent graph options, artifacts, metrics, evidence, or a combined + score. Compare cluster viability and biological preservation evidence + that is actually present. ARI and NMI describe agreement, not quality. + WNN modality weights are usable only when modalityWeightsValid=true. + UMAP appearance, native-neighbor LISI on an integrated graph, and + absent metric fields are not evidence. Return one comparison for every + eligible non-selected option, citing evidence from both options. + Select an option and explain it; Scarf attaches its exact graph, + assay, and execution identities. Do not return those derived fields. + """ + ) + .strip() + .format() + ) + + +def final_graph_selection_prompt( + *, + report: ParameterTuningReport, + integration_evaluations: Sequence[IntegrationCandidateEvaluation], + marker_assay: str, +) -> str: + """Build the selection prompt from executor-grounded final graph options.""" + + options = _final_graph_options(report, integration_evaluations) + return ( + dedent( + """ + Select the final graph from these eligible executed options: + {options} + + The fixed marker assay is {marker_assay}. It determines marker + extraction and does not imply ownership of an integrated graph. + Return needsInput only when the supplied evidence cannot resolve a + scientifically material tradeoff. + """ + ) + .strip() + .format( + options=json.dumps(options, indent=2, sort_keys=True), + marker_assay=marker_assay, + ) + ) diff --git a/scarf/agent/parameter_tuning/selection.py b/scarf/agent/parameter_tuning/selection.py new file mode 100644 index 00000000..b854cc5f --- /dev/null +++ b/scarf/agent/parameter_tuning/selection.py @@ -0,0 +1,1534 @@ +from collections.abc import Mapping, Sequence +from typing import Any + +import numpy as np + +from ...storage.refs import ArtifactRef +from ...utils.logging import logger +from .._deps import AGENT_INSTALL_HINT +from ..config import AgentRunConfig +from ..config.agent_exec import run_agent_sync +from ..tools import artifact_reference, core_artifact_reference +from ..types import AgentRunInfo, StageStatus +from .contracts import ( + _CANDIDATE_ID, + ArtifactRecord, + FinalGraphNeedsInput, + FinalGraphSelection, + IntegrationCandidateEvaluation, + ParameterCandidate, + ParameterCandidateEvaluation, + ParameterMetrics, + ParameterSearchPlan, + ParameterSearchStatus, + ParameterTuningBatchSearchPlan, + ParameterTuningDependencies, + ParameterTuningNeedsInput, + ParameterTuningReport, +) +from .execution import _final_graph_options +from .prompts import final_graph_selection_prompt, final_graph_selection_system_prompt + +try: + from pydantic_ai import UnexpectedModelBehavior, UsageLimitExceeded +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + + +def final_graph_options( + report: ParameterTuningReport, + integration_evaluations: Sequence[IntegrationCandidateEvaluation], +) -> dict[str, dict[str, Any]]: + """Return the exact eligible graph options and option-scoped evidence.""" + + return _final_graph_options(report, integration_evaluations) + + +def _finite_metric( + objectives: dict[str, tuple[float, int, str]], + name: str, + value: float | None, + *, + direction: int, + evidence_class: str, +) -> None: + if value is not None and np.isfinite(value): + objectives[name] = (float(value), direction, evidence_class) + + +def _candidate_objectives( + metrics: ParameterMetrics, +) -> dict[str, tuple[float, int, str]]: + objectives: dict[str, tuple[float, int, str]] = {} + for name, value in ( + ("minClusterFraction", metrics.minClusterFraction), + ("graphSilhouetteMedian", metrics.graphSilhouetteMedian), + ("membershipStrengthMean", metrics.membershipStrengthMean), + ("membershipStrengthP10", metrics.membershipStrengthP10), + ("clusterConnectivity", metrics.clusterConnectivity), + ): + _finite_metric( + objectives, + name, + value, + direction=1, + evidence_class="geometric", + ) + for name, value in ( + ("seedStability", metrics.seedStability), + ("subsampleStability", metrics.subsampleStability), + ): + _finite_metric( + objectives, + name, + value, + direction=1, + evidence_class="resamplingStability", + ) + for name, value in ( + ("markerCoherence", metrics.markerCoherence), + ("markerSpecificityMedian", metrics.markerSpecificityMedian), + ): + _finite_metric( + objectives, + name, + value, + direction=1, + evidence_class="markerCoherence", + ) + _finite_metric( + objectives, + "crossUnitSupport", + metrics.crossUnitSupport, + direction=1, + evidence_class="crossUnitSupport", + ) + _finite_metric( + objectives, + "doubletHighScoreConcentration", + metrics.doubletHighScoreConcentration, + direction=-1, + evidence_class="qualityControl", + ) + for column, value in metrics.technicalAssociation.items(): + _finite_metric( + objectives, + f"technicalAssociation:{column}", + value, + direction=-1, + evidence_class="technical", + ) + for column, value in metrics.batchMixing.items(): + _finite_metric( + objectives, + f"batchMixing:{column}", + value, + direction=1, + evidence_class="batchRemoval", + ) + for column, values in metrics.biologicalPreservation.items(): + for name, value in values.items(): + _finite_metric( + objectives, + f"biologicalPreservation:{column}:{name}", + value, + direction=1, + evidence_class="protectedVariablePreservation", + ) + return objectives + + +def _single_varied_parameter( + left: ParameterCandidate, + right: ParameterCandidate, +) -> str | None: + if ( + left.reductionMethod != right.reductionMethod + or left.useHarmony != right.useHarmony + ): + return None + varied = [ + name + for name in ("dimensions", "neighborsK", "leidenResolution") + if getattr(left, name) != getattr(right, name) + ] + return varied[0] if len(varied) == 1 else None + + +def _dominance_metrics( + left: ParameterMetrics, + right: ParameterMetrics, + *, + tolerance: float, +) -> list[str]: + left_objectives = _candidate_objectives(left) + right_objectives = _candidate_objectives(right) + if not left_objectives or set(left_objectives) != set(right_objectives): + return [] + classes = {value[2] for value in left_objectives.values()} + if len(classes) < 2: + return [] + strict: list[str] = [] + for name in sorted(left_objectives): + left_value, direction, _evidence_class = left_objectives[name] + right_value = right_objectives[name][0] + difference = direction * (left_value - right_value) + if difference < -tolerance: + return [] + if difference > tolerance: + strict.append(name) + return strict + + +def annotate_candidate_dominance( + evaluations: Sequence[ParameterCandidateEvaluation], + *, + tolerance: float = 0.02, +) -> tuple[ParameterCandidateEvaluation, ...]: + """Attach conservative pairwise Pareto evidence to comparable candidates.""" + + values = list(evaluations) + if tolerance < 0 or not np.isfinite(tolerance): + raise ValueError("Dominance tolerance must be finite and non-negative") + completed = [value for value in values if value.status == "done" and value.eligible] + dominated_by: dict[str, list[str]] = {value.candidateId: [] for value in completed} + dominates: dict[str, list[str]] = {value.candidateId: [] for value in completed} + metrics_by_id: dict[str, dict[str, list[str]]] = { + value.candidateId: {} for value in completed + } + comparable: set[str] = set() + for left in completed: + for right in completed: + if left.candidateId == right.candidateId or ( + _single_varied_parameter(left.parameters, right.parameters) is None + ): + continue + comparable.add(left.candidateId) + strict = _dominance_metrics( + left.metrics, + right.metrics, + tolerance=tolerance, + ) + if not strict: + continue + dominates[left.candidateId].append(right.candidateId) + dominated_by[right.candidateId].append(left.candidateId) + metrics_by_id[left.candidateId][f"dominates:{right.candidateId}"] = strict + metrics_by_id[right.candidateId][f"dominatedBy:{left.candidateId}"] = strict + + annotated: list[ParameterCandidateEvaluation] = [] + for evaluation in values: + if evaluation.candidateId not in dominated_by: + annotated.append(evaluation) + continue + candidate_id = evaluation.candidateId + candidate_dominators = sorted(set(dominated_by[candidate_id])) + candidate_dominates = sorted(set(dominates[candidate_id])) + updated_metrics = evaluation.metrics.model_copy( + update={ + "paretoOptimal": ( + not candidate_dominators if candidate_id in comparable else None + ), + "dominatedByCandidateIds": candidate_dominators, + "dominatesCandidateIds": candidate_dominates, + "dominanceMetrics": metrics_by_id[candidate_id], + } + ) + prefix = f"candidate:{candidate_id}:" + retained_evidence = [ + evidence_id + for evidence_id in evaluation.evidenceIds + if not ( + evidence_id == f"{prefix}paretoDominance" + or evidence_id.startswith(f"{prefix}dominatedBy:") + or evidence_id.startswith(f"{prefix}dominates:") + ) + ] + dominance_evidence = ( + [f"{prefix}paretoDominance"] if candidate_id in comparable else [] + ) + dominance_evidence.extend( + f"{prefix}dominatedBy:{other}" for other in candidate_dominators + ) + dominance_evidence.extend( + f"{prefix}dominates:{other}" for other in candidate_dominates + ) + annotated.append( + evaluation.model_copy( + update={ + "metrics": updated_metrics, + "evidenceIds": [ + *retained_evidence, + *dominance_evidence, + ], + } + ) + ) + return tuple(annotated) + + +def harmony_acceptance_gate( + native: ParameterCandidateEvaluation | None, + harmony: ParameterCandidateEvaluation | None, + *, + batch_columns: Sequence[str], + protected_columns: Sequence[str], + independent_unit_columns: Sequence[str] = (), + tolerance: float = 0.05, + require_doublet_evidence: bool = False, +) -> tuple[bool, list[str]]: + """Require matched batch improvement without material biological loss.""" + + if tolerance < 0 or not np.isfinite(tolerance): + raise ValueError("Harmony gate tolerance must be finite and non-negative") + reasons: list[str] = [] + if native is None or harmony is None: + return False, ["Matched native and Harmony candidates are unavailable."] + if native.status != "done" or not native.eligible: + reasons.append("The matched native candidate is not an eligible execution.") + if harmony.status != "done" or not harmony.eligible: + reasons.append("The matched Harmony candidate is not an eligible execution.") + if native.parameters.useHarmony or not harmony.parameters.useHarmony: + reasons.append("Candidates do not have native and Harmony correction modes.") + native_parameters = native.parameters.model_dump( + mode="json", + exclude={"candidateId", "useHarmony"}, + ) + harmony_parameters = harmony.parameters.model_dump( + mode="json", + exclude={"candidateId", "useHarmony"}, + ) + if native_parameters != harmony_parameters: + reasons.append("Native and Harmony candidate parameters are not matched.") + if core_artifact_reference(native.cellSelection) != core_artifact_reference( + harmony.cellSelection + ): + reasons.append("Native and Harmony candidates use different cell selections.") + + columns = list(dict.fromkeys(batch_columns)) + if not columns: + reasons.append("No approved batch metric was supplied.") + batch_deltas: dict[str, float] = {} + for column in columns: + native_score = native.metrics.batchMixing.get(column) + harmony_score = harmony.metrics.batchMixing.get(column) + if native_score is None or harmony_score is None: + reasons.append(f"Batch comparison is missing for {column!r}.") + continue + batch_deltas[column] = harmony_score - native_score + if columns and len(batch_deltas) == len(columns): + if not any(delta > tolerance for delta in batch_deltas.values()): + reasons.append( + "Harmony did not improve an approved batch metric beyond tolerance." + ) + if any(delta < -tolerance for delta in batch_deltas.values()): + reasons.append("Harmony materially worsened an approved batch metric.") + + for column in dict.fromkeys(protected_columns): + native_scores = native.metrics.biologicalPreservation.get(column) + harmony_scores = harmony.metrics.biologicalPreservation.get(column) + if not native_scores or not harmony_scores: + reasons.append(f"Protected comparison is missing for {column!r}.") + continue + missing_metrics = sorted( + {"clisi", "graphConnectivity"} - (set(native_scores) & set(harmony_scores)) + ) + if missing_metrics: + reasons.append( + f"Required protected metrics {missing_metrics} are missing for {column!r}." + ) + continue + if set(native_scores) != set(harmony_scores): + reasons.append(f"Protected metrics do not align for {column!r}.") + continue + if any( + harmony_scores[name] < native_scores[name] - tolerance + for name in native_scores + ): + reasons.append( + f"Harmony materially degraded protected evidence for {column!r}." + ) + + if independent_unit_columns: + if ( + native.metrics.crossUnitSupport is None + or harmony.metrics.crossUnitSupport is None + ): + reasons.append("Cross-unit support comparison is missing.") + elif ( + harmony.metrics.crossUnitSupport + < native.metrics.crossUnitSupport - tolerance + ): + reasons.append("Harmony materially degraded cross-unit support.") + + if ( + native.metrics.markerCoherence is None + or harmony.metrics.markerCoherence is None + ): + reasons.append("Marker-coherence comparison is missing.") + elif harmony.metrics.markerCoherence < native.metrics.markerCoherence - tolerance: + reasons.append("Harmony materially degraded marker coherence.") + + for label, native_value, harmony_value in ( + ( + "marker specificity", + native.metrics.markerSpecificityMedian, + harmony.metrics.markerSpecificityMedian, + ), + ( + "cluster connectivity", + native.metrics.clusterConnectivity, + harmony.metrics.clusterConnectivity, + ), + ( + "membership strength", + native.metrics.membershipStrengthMean, + harmony.metrics.membershipStrengthMean, + ), + ): + if native_value is None and harmony_value is None: + continue + if native_value is None or harmony_value is None: + reasons.append(f"Matched {label} comparison is missing.") + elif harmony_value < native_value - tolerance: + reasons.append(f"Harmony materially degraded {label}.") + + native_doublet = native.metrics.doubletHighScoreConcentration + harmony_doublet = harmony.metrics.doubletHighScoreConcentration + if ( + require_doublet_evidence + or native_doublet is not None + or harmony_doublet is not None + ): + if native_doublet is None or harmony_doublet is None: + reasons.append("Matched doublet-concentration comparison is missing.") + elif harmony_doublet > native_doublet + tolerance: + reasons.append("Harmony materially increased doublet concentration.") + return not reasons, reasons + + +def validate_parameter_search_plan( + plan: ParameterSearchPlan, + deps: ParameterTuningDependencies, + *, + initial_candidate_ids: Sequence[str], + max_refined_candidates: int, +) -> ParameterSearchPlan: + """Validate one refinement proposal against the completed initial screen.""" + + initial_evaluations = [ + deps.evaluations[candidate_id] + for candidate_id in initial_candidate_ids + if candidate_id in deps.evaluations + ] + known_evidence = { + evidence_id + for evaluation in initial_evaluations + for evidence_id in evaluation.evidenceIds + } + unknown_evidence = sorted(set(plan.evidenceIds) - known_evidence) + if unknown_evidence: + raise ValueError( + f"Parameter search plan cites unknown evidence ids {unknown_evidence}" + ) + authorized_batch_columns = list(deps.batchColumns) if deps.harmonyAuthorized else [] + if ( + plan.harmonyBatchColumns + and plan.harmonyBatchColumns != authorized_batch_columns + ): + raise ValueError( + "Parameter search plan cannot modify the exact authorized Harmony " + "batch columns" + ) + canonical_status: ParameterSearchStatus = ( + "refine" if plan.candidates else "complete" + ) + plan = plan.model_copy( + update={ + "status": canonical_status, + "harmonyBatchColumns": authorized_batch_columns, + } + ) + if plan.status == "complete": + return plan + + if len(plan.candidates) > max_refined_candidates: + raise ValueError( + "Parameter search plan exceeds the refined candidate limit " + f"{max_refined_candidates}" + ) + if not plan.rationale.strip(): + raise ValueError("A refinement plan requires a rationale") + if not plan.objectives: + raise ValueError("A refinement plan requires focused objectives") + if not plan.stoppingCriteria: + raise ValueError("A refinement plan requires stopping criteria") + if not plan.evidenceIds: + raise ValueError("A refinement plan requires initial-screen evidence") + + successful_initial_ids = { + evaluation.candidateId + for evaluation in initial_evaluations + if evaluation.status == "done" + } + if not plan.basedOnCandidateIds: + raise ValueError("A refinement plan must identify its initial candidates") + duplicate_parents = sorted( + { + candidate_id + for candidate_id in plan.basedOnCandidateIds + if plan.basedOnCandidateIds.count(candidate_id) > 1 + } + ) + if duplicate_parents: + raise ValueError(f"Duplicate refinement parent ids {duplicate_parents}") + invalid_parents = sorted(set(plan.basedOnCandidateIds) - successful_initial_ids) + if invalid_parents: + raise ValueError( + "Refinement parents must be successful initial candidates: " + f"{invalid_parents}" + ) + for parent_id in plan.basedOnCandidateIds: + prefix = f"candidate:{parent_id}:" + if not any(evidence_id.startswith(prefix) for evidence_id in plan.evidenceIds): + raise ValueError( + f"Refinement evidence must cite every parent candidate: {parent_id!r}" + ) + if deps.harmonyAuthorized and any( + candidate.useHarmony for candidate in plan.candidates + ): + parent_candidates = [ + deps.candidates[candidate_id] for candidate_id in plan.basedOnCandidateIds + ] + paired_modes: dict[tuple[str, int, float, int], set[bool]] = {} + for candidate in parent_candidates: + parameter_key = ( + candidate.reductionMethod, + candidate.dimensions, + candidate.leidenResolution, + candidate.neighborsK, + ) + paired_modes.setdefault(parameter_key, set()).add(candidate.useHarmony) + if not any(modes == {False, True} for modes in paired_modes.values()): + raise ValueError( + "Harmony refinement requires evidence from one matched corrected " + "and uncorrected initial pair" + ) + + initial_candidates = [ + deps.candidates[candidate_id] for candidate_id in initial_candidate_ids + ] + known_signatures = { + ( + candidate.reductionMethod, + candidate.dimensions, + candidate.leidenResolution, + candidate.neighborsK, + candidate.useHarmony, + ) + for candidate in initial_candidates + } + proposed_ids: set[str] = set() + proposed_signatures: set[tuple[str, int, float, int, bool]] = set() + for candidate in plan.candidates: + if not _CANDIDATE_ID.fullmatch(candidate.candidateId): + raise ValueError( + "Refined candidateId must contain only ASCII letters, numbers, " + "and underscores" + ) + if ( + candidate.candidateId in deps.candidates + or candidate.candidateId in proposed_ids + ): + raise ValueError(f"Duplicate refined candidateId {candidate.candidateId!r}") + proposed_ids.add(candidate.candidateId) + method_candidates = [ + item + for item in initial_candidates + if item.reductionMethod == candidate.reductionMethod + ] + if not method_candidates: + raise ValueError( + "Refined candidates cannot introduce an untested reduction method: " + f"{candidate.reductionMethod!r}" + ) + dimension_bounds = ( + min(item.dimensions for item in method_candidates), + max(item.dimensions for item in method_candidates), + ) + resolution_bounds = ( + min(item.leidenResolution for item in method_candidates), + max(item.leidenResolution for item in method_candidates), + ) + neighbor_bounds = ( + min(item.neighborsK for item in method_candidates), + max(item.neighborsK for item in method_candidates), + ) + if not dimension_bounds[0] <= candidate.dimensions <= dimension_bounds[1]: + raise ValueError( + "Refined dimensions must remain inside the initial search envelope " + f"{dimension_bounds}" + ) + if not ( + resolution_bounds[0] <= candidate.leidenResolution <= resolution_bounds[1] + ): + raise ValueError( + "Refined Leiden resolution must remain inside the initial search " + f"envelope {resolution_bounds}" + ) + if not neighbor_bounds[0] <= candidate.neighborsK <= neighbor_bounds[1]: + raise ValueError( + "Refined neighbor count must remain inside the initial search " + f"envelope {neighbor_bounds}" + ) + if candidate.useHarmony and ( + not deps.harmonyAuthorized or not deps.batchColumns + ): + raise ValueError( + f"Refined candidate {candidate.candidateId!r} is not authorized " + "for Harmony" + ) + signature = ( + candidate.reductionMethod, + candidate.dimensions, + candidate.leidenResolution, + candidate.neighborsK, + candidate.useHarmony, + ) + if signature in known_signatures or signature in proposed_signatures: + raise ValueError( + f"Refined candidate {candidate.candidateId!r} duplicates an " + "evaluated or proposed parameter branch" + ) + proposed_signatures.add(signature) + return plan + + +def validate_parameter_batch_search_plan( + plan: ParameterTuningBatchSearchPlan, + dependencies: Mapping[str, ParameterTuningDependencies], + *, + initial_candidate_ids: Mapping[str, Sequence[str]], + max_refined_by_assay: Mapping[str, int], +) -> ParameterTuningBatchSearchPlan: + """Validate every assay entry in one batched refinement response.""" + + expected = set(dependencies) + actual = set(plan.assayPlans) + if actual != expected: + raise ValueError( + "Batched refinement must contain exactly the requested assays: " + f"missing={sorted(expected - actual)}, unexpected={sorted(actual - expected)}" + ) + validated = { + assay: validate_parameter_search_plan( + plan.assayPlans[assay], + dependencies[assay], + initial_candidate_ids=initial_candidate_ids[assay], + max_refined_candidates=max_refined_by_assay[assay], + ) + for assay in dependencies + } + return plan.model_copy(update={"assayPlans": validated}) + + +def parameter_evidence_classes(evidence_ids: Sequence[str]) -> frozenset[str]: + """Infer stable scientific evidence classes from executor evidence IDs.""" + + classes: set[str] = set() + for evidence_id in evidence_ids: + token = evidence_id.casefold() + if ( + "seedstability" in token + or "subsamplestability" in token + or token.endswith(":stability") + ): + classes.add("resamplingStability") + elif "marker" in token: + classes.add("markerCoherence") + elif "crossunitsupport" in token or "unitsupport" in token: + classes.add("crossUnitSupport") + elif ( + "protected" in token + or "clisi" in token + or ("graphconnectivity" in token and "clusterconnectivity" not in token) + ): + classes.add("protectedVariablePreservation") + elif "doublet" in token: + classes.add("qualityControl") + elif "technical" in token or "batchmixing" in token: + classes.add("technical") + elif any( + value in token + for value in ( + "clusterconnectivity", + "clusters", + "geometry", + "membershipstrength", + "neighbor", + "paretodominance", + "silhouette", + ) + ): + classes.add("geometric") + return frozenset(classes) + + +def require_dominated_candidate_evidence( + selected: ParameterCandidateEvaluation, + evidence_ids: Sequence[str], + *, + context: str, +) -> None: + """Require two independent non-geometric classes for a dominated choice.""" + + if not selected.metrics.dominatedByCandidateIds: + return + independent = parameter_evidence_classes(evidence_ids).intersection( + { + "markerCoherence", + "resamplingStability", + "crossUnitSupport", + "protectedVariablePreservation", + "qualityControl", + } + ) + if len(independent) < 2: + raise ValueError( + f"{context} selects a Pareto-dominated candidate and must cite at " + "least two independent non-geometric evidence classes" + ) + + +def validate_parameter_tuning_report( + report: ParameterTuningReport, + deps: ParameterTuningDependencies, + *, + search_plan: ParameterSearchPlan | None = None, +) -> ParameterTuningReport: + """Ground the model report in candidate executions recorded by the tool.""" + + evaluations = list( + annotate_candidate_dominance( + [ + deps.evaluations[candidate_id] + for candidate_id in deps.executionOrder + if candidate_id in deps.evaluations + ] + ) + ) + evaluations_by_id = { + evaluation.candidateId: evaluation for evaluation in evaluations + } + known_evidence = { + evidence_id + for evaluation in evaluations + for evidence_id in evaluation.evidenceIds + } + cited_evidence = set(report.evidenceIds) + for comparison in report.comparisons: + cited_evidence.update(comparison.evidenceIds) + if report.needsInput is not None: + cited_evidence.update(report.needsInput.evidenceIds) + unknown_evidence = sorted(cited_evidence - known_evidence) + if unknown_evidence: + raise ValueError( + f"Parameter tuning report cites unknown evidence ids {unknown_evidence}" + ) + if report.status == "done" and report.recommendedCandidateId is None: + raise ValueError("A done tuning report must recommend an executed candidate") + if report.status == "needsInput" and report.needsInput is None: + raise ValueError("A needsInput tuning report must include a concrete question") + successful = [ + evaluation for evaluation in evaluations if evaluation.status == "done" + ] + comparison_required = len(deps.candidates) > 1 and deps.maxCandidates > 1 + if report.status == "done": + if not report.evidenceIds: + raise ValueError("A done tuning report requires recommendation evidence") + if comparison_required and len(successful) < 2: + raise ValueError( + "A completed tuning recommendation requires at least two successful " + "candidate executions" + ) + if ( + comparison_required + and "baseline" in deps.candidates + and not any(item.candidateId == "baseline" for item in successful) + ): + raise ValueError( + "Evaluate the baseline before completing a multi-candidate comparison" + ) + + selected_artifacts: dict[str, ArtifactRecord] = {} + selected_evaluation: ParameterCandidateEvaluation | None = None + if report.recommendedCandidateId is not None: + selected = evaluations_by_id.get(report.recommendedCandidateId) + if selected is None: + raise ValueError("Recommended candidate was not executed") + if selected.status != "done": + raise ValueError("Recommended candidate execution failed") + if not selected.eligible: + raise ValueError("Recommended candidate is not eligible") + recommendation_prefix = f"candidate:{selected.candidateId}:" + if not any( + evidence_id.startswith(recommendation_prefix) + for evidence_id in report.evidenceIds + ): + raise ValueError( + "Recommendation evidence must include the selected candidate" + ) + selected_evaluation = selected + selected_artifacts = dict(selected.artifacts) + + if report.status == "done" and not comparison_required and report.comparisons: + raise ValueError( + "Candidate comparisons require a completed multi-candidate evaluation" + ) + if report.status == "done" and comparison_required: + assert report.recommendedCandidateId is not None + successful_ids = {item.candidateId for item in successful} + expected_comparators = successful_ids - {report.recommendedCandidateId} + comparison_ids = [item.candidateId for item in report.comparisons] + duplicate_comparators = sorted( + { + candidate_id + for candidate_id in comparison_ids + if comparison_ids.count(candidate_id) > 1 + } + ) + if duplicate_comparators: + raise ValueError(f"Duplicate candidate comparisons {duplicate_comparators}") + actual_comparators = set(comparison_ids) + missing_comparators = sorted(expected_comparators - actual_comparators) + invalid_comparators = sorted(actual_comparators - expected_comparators) + if missing_comparators: + raise ValueError( + "Completed tuning reports require comparisons for every successful " + f"non-selected candidate: {missing_comparators}" + ) + if invalid_comparators: + raise ValueError( + "Candidate comparisons must identify successful non-selected " + f"candidates: {invalid_comparators}" + ) + selected_prefix = f"candidate:{report.recommendedCandidateId}:" + for comparison in report.comparisons: + comparator_prefix = f"candidate:{comparison.candidateId}:" + if not any( + evidence_id.startswith(selected_prefix) + for evidence_id in comparison.evidenceIds + ): + raise ValueError( + "Each candidate comparison must cite evidence from the " + "selected candidate" + ) + if not any( + evidence_id.startswith(comparator_prefix) + for evidence_id in comparison.evidenceIds + ): + raise ValueError( + "Each candidate comparison must cite evidence from its comparator" + ) + if not comparison.summary.strip(): + raise ValueError( + "Each candidate comparison requires a concise grounded summary" + ) + if ( + selected_evaluation is not None + and comparison.candidateId + in selected_evaluation.metrics.dominatedByCandidateIds + and _single_varied_parameter( + selected_evaluation.parameters, + evaluations_by_id[comparison.candidateId].parameters, + ) + in {"neighborsK", "leidenResolution"} + ): + require_dominated_candidate_evidence( + selected_evaluation, + comparison.evidenceIds, + context=(f"The comparison with {comparison.candidateId!r}"), + ) + + graph_partition_dominators = ( + [ + candidate_id + for candidate_id in selected_evaluation.metrics.dominatedByCandidateIds + if candidate_id in evaluations_by_id + and _single_varied_parameter( + selected_evaluation.parameters, + evaluations_by_id[candidate_id].parameters, + ) + in {"neighborsK", "leidenResolution"} + ] + if selected_evaluation is not None + else [] + ) + if ( + report.status == "done" + and selected_evaluation is not None + and graph_partition_dominators + ): + selection_evidence = [ + *report.evidenceIds, + *( + evidence_id + for comparison in report.comparisons + for evidence_id in comparison.evidenceIds + ), + ] + require_dominated_candidate_evidence( + selected_evaluation, + selection_evidence, + context="The tuning recommendation", + ) + + return report.model_copy( + update={ + "fromAssay": deps.fromAssay, + "cellSelection": artifact_reference(deps.cellSelection), + "evaluations": evaluations, + "selectedArtifacts": selected_artifacts, + "searchPlan": search_plan, + "assayReports": {}, + "recommendedByAssay": ( + {deps.fromAssay: report.recommendedCandidateId} + if report.recommendedCandidateId is not None + else {} + ), + "totalCandidates": len(evaluations), + "integrationEvaluations": [], + "recommendedIntegrationId": None, + "finalClusterColumn": None, + "finalClusterArtifact": None, + "graphAssay": deps.fromAssay, + "markerAssay": deps.fromAssay, + "finalSelection": None, + } + ) + + +def validate_parameter_tuning_batch_report( + report: ParameterTuningReport, + dependencies: Mapping[str, ParameterTuningDependencies], + *, + search_plans: Mapping[str, ParameterSearchPlan], + primary_assay: str, +) -> ParameterTuningReport: + """Ground one aggregate response in every assay's executed branches.""" + + expected = set(dependencies) + actual = set(report.assayReports) + if actual != expected: + raise ValueError( + "Batched selection must contain exactly the requested assays: " + f"missing={sorted(expected - actual)}, unexpected={sorted(actual - expected)}" + ) + if primary_assay not in dependencies: + raise ValueError(f"Unknown primary assay {primary_assay!r}") + validated_reports = { + assay: validate_parameter_tuning_report( + report.assayReports[assay], + dependencies[assay], + search_plan=search_plans[assay], + ) + for assay in dependencies + } + known_evidence = { + evidence_id + for assay_report in validated_reports.values() + for evaluation in assay_report.evaluations + for evidence_id in evaluation.evidenceIds + } + unknown_evidence = sorted(set(report.evidenceIds) - known_evidence) + if unknown_evidence: + raise ValueError( + f"Batched tuning report cites unknown evidence ids {unknown_evidence}" + ) + statuses = {assay_report.status for assay_report in validated_reports.values()} + if statuses == {"done"}: + status: StageStatus = "done" + elif "needsInput" in statuses: + status = "needsInput" + else: + status = "failed" + primary = validated_reports[primary_assay] + recommended = { + assay: assay_report.recommendedCandidateId + for assay, assay_report in validated_reports.items() + if assay_report.recommendedCandidateId is not None + } + if status == "done" and len(recommended) != len(validated_reports): + raise ValueError("Every completed assay report must recommend a candidate") + cell_selections = [ + core_artifact_reference(assay_report.cellSelection) + for assay_report in validated_reports.values() + ] + if ( + not cell_selections + or not isinstance(cell_selections[0], ArtifactRef) + or any(selection != cell_selections[0] for selection in cell_selections[1:]) + ): + raise ValueError("Every assay report must use the same exact cell selection") + return report.model_copy( + update={ + "status": status, + "fromAssay": primary_assay, + "cellSelection": primary.cellSelection, + "evaluations": primary.evaluations, + "recommendedCandidateId": primary.recommendedCandidateId, + "selectedArtifacts": primary.selectedArtifacts, + "needsInput": primary.needsInput if status != "done" else None, + "searchPlan": primary.searchPlan, + "assayReports": validated_reports, + "recommendedByAssay": recommended, + "totalCandidates": sum( + len(item.evaluations) for item in validated_reports.values() + ), + "integrationEvaluations": [], + "recommendedIntegrationId": None, + "finalClusterColumn": None, + "finalClusterArtifact": None, + "graphAssay": primary_assay, + "markerAssay": primary_assay, + "finalSelection": None, + } + ) + + +def pending_parameter_tuning_report( + deps: ParameterTuningDependencies, + *, + search_plan: ParameterSearchPlan, + agent_name: str, +) -> ParameterTuningReport: + """Pause when structured selection is unavailable. + + Completed executor evidence is retained for an exact human resume, but it is + never converted into an implicit scientific recommendation. + """ + + evaluations = [ + deps.evaluations[candidate_id] + for candidate_id in deps.executionOrder + if candidate_id in deps.evaluations + ] + successful = [item for item in evaluations if item.status == "done"] + eligible = [item for item in successful if item.eligible] + logger.warning( + f"Parameter tuning for assay {deps.fromAssay!r} requires input after " + f"model exhaustion: completed={len(successful)}, eligible={len(eligible)}" + ) + known_evidence = sorted( + { + evidence_id + for evaluation in evaluations + for evidence_id in evaluation.evidenceIds + } + ) + report = ParameterTuningReport( + status="needsInput", + confidence="low", + rationale=( + "The bounded structured selection was unavailable. Executor evidence " + "is complete enough to resume, but it cannot choose a scientific " + "alternative by itself." + ), + evidenceIds=known_evidence, + limitations=["No candidate was selected merely to complete the workflow."], + stopReason="The bounded screen completed without a valid decision.", + needsInput=ParameterTuningNeedsInput( + question=( + "Select one eligible executed candidate and provide a scientific " + "rationale tied to the cited evidence." + ), + options=[item.candidateId for item in eligible], + evidenceIds=known_evidence, + ), + runInfo=AgentRunInfo(agentName=agent_name), + ) + return validate_parameter_tuning_report( + report, + deps, + search_plan=search_plan, + ) + + +def pending_parameter_tuning_batch_report( + dependencies: Mapping[str, ParameterTuningDependencies], + *, + search_plans: Mapping[str, ParameterSearchPlan], + primary_assay: str, +) -> ParameterTuningReport: + """Build one grounded pause over completed assay screens.""" + + logger.warning( + f"Pausing parameter tuning after model exhaustion for " + f"{len(dependencies)} assay(s)" + ) + assay_reports = { + assay: pending_parameter_tuning_report( + deps, + search_plan=search_plans[assay], + agent_name="parameter_tuning_batch_needs_input", + ) + for assay, deps in dependencies.items() + } + aggregate = ParameterTuningReport( + status="needsInput", + assayReports=assay_reports, + rationale=( + "Structured model selection was unavailable. All completed evidence " + "was retained without choosing a branch." + ), + evidenceIds=list( + dict.fromkeys( + evidence_id + for assay_report in assay_reports.values() + for evidence_id in assay_report.evidenceIds + ) + ), + limitations=["No assay candidate was selected merely to finish the workflow."], + stopReason="The bounded native screens completed without valid decisions.", + runInfo=AgentRunInfo(agentName="parameter_tuning_batch_needs_input"), + ) + logger.warning( + f"Parameter tuning batch pause status={aggregate.status}; " + f"completed_assays={sum(item.status == 'done' for item in assay_reports.values())}" + ) + return validate_parameter_tuning_batch_report( + aggregate, + dependencies, + search_plans=search_plans, + primary_assay=primary_assay, + ) + + +def validate_final_graph_selection( + selection: FinalGraphSelection, + report: ParameterTuningReport, + *, + integration_evaluations: Sequence[IntegrationCandidateEvaluation], + marker_assay: str, +) -> FinalGraphSelection: + """Ground a final graph choice in the exact eligible executor outputs.""" + + if report.status != "done": + raise ValueError("Native parameter tuning must finish before graph selection") + options = final_graph_options(report, integration_evaluations) + if not options: + raise ValueError("No eligible native or integrated graph options are available") + known_evidence = { + evidence_id + for option in options.values() + for evidence_id in option["evidenceIds"] + } + cited = set(selection.evidenceIds) + for comparison in selection.comparisons: + cited.update(comparison.evidenceIds) + if selection.needsInput is not None: + cited.update(selection.needsInput.evidenceIds) + unknown = sorted(cited - known_evidence) + if unknown: + raise ValueError(f"Final graph selection cites unknown evidence ids {unknown}") + if selection.status == "needsInput": + if selection.needsInput is None or not selection.needsInput.question.strip(): + raise ValueError( + "A needsInput graph selection requires a concrete question" + ) + return selection.model_copy( + update={ + "selectedOptionId": None, + "graphMethod": None, + "nativeAssay": None, + "nativeCandidateId": None, + "integrationId": None, + "markerAssay": marker_assay, + } + ) + if selection.status != "done": + raise ValueError("Final graph selection must be done or needsInput") + if selection.selectedOptionId not in options: + raise ValueError("Selected final graph option is not eligible") + assert selection.selectedOptionId is not None + selected = options[selection.selectedOptionId] + selected_evidence = set(selected["evidenceIds"]) + if not selected_evidence.intersection(selection.evidenceIds): + raise ValueError( + "Final graph recommendation must cite selected-option evidence" + ) + expected_comparators = set(options) - {selection.selectedOptionId} + comparison_ids = [item.optionId for item in selection.comparisons] + if len(set(comparison_ids)) != len(comparison_ids): + raise ValueError("Final graph comparisons must not contain duplicates") + if set(comparison_ids) != expected_comparators: + raise ValueError( + "Final graph selection requires one comparison for every eligible " + "non-selected option" + ) + for comparison in selection.comparisons: + comparator_evidence = set(options[comparison.optionId]["evidenceIds"]) + if not selected_evidence.intersection(comparison.evidenceIds): + raise ValueError( + "Every final graph comparison must cite selected-option evidence" + ) + if not comparator_evidence.intersection(comparison.evidenceIds): + raise ValueError( + "Every final graph comparison must cite comparator evidence" + ) + if not comparison.summary.strip(): + raise ValueError("Every final graph comparison requires a summary") + return selection.model_copy( + update={ + "graphMethod": selected["graphMethod"], + "nativeAssay": selected.get("nativeAssay"), + "nativeCandidateId": selected.get("nativeCandidateId"), + "integrationId": selected.get("integrationId"), + "markerAssay": marker_assay, + "needsInput": None, + } + ) + + +def finalize_parameter_tuning_selection( + report: ParameterTuningReport, + *, + marker_assay: str, + integration_evaluations: Sequence[IntegrationCandidateEvaluation] = (), + recommended_integration_id: str | None = None, + native_assay: str | None = None, + final_selection: FinalGraphSelection | None = None, +) -> ParameterTuningReport: + """Attach an executor-selected native or integrated final cluster branch.""" + + logger.debug( + f"Finalizing parameter graph selection: marker_assay={marker_assay!r}, " + f"integration_candidates={len(integration_evaluations)}" + ) + if report.status != "done": + raise ValueError("Parameter tuning must be done before final graph selection") + if not marker_assay: + raise ValueError("marker_assay must be non-empty") + report_cell_selection = core_artifact_reference(report.cellSelection) + if not isinstance(report_cell_selection, ArtifactRef): + raise ValueError("Parameter tuning report lacks an exact cell selection") + assay_reports = report.assayReports or {report.fromAssay: report} + if marker_assay not in assay_reports: + raise ValueError(f"Unknown marker assay {marker_assay!r}") + evaluations = list(integration_evaluations) + integration_ids = [item.integrationId for item in evaluations] + if len(set(integration_ids)) != len(integration_ids): + raise ValueError("Integration evaluation ids must be unique") + if recommended_integration_id is not None and native_assay is not None: + raise ValueError("Choose either an integrated graph or one native assay") + if recommended_integration_id is not None: + selected = next( + ( + item + for item in evaluations + if item.integrationId == recommended_integration_id + ), + None, + ) + if selected is None: + raise ValueError("Recommended integration candidate was not evaluated") + if selected.status != "done" or not selected.eligible: + raise ValueError("Recommended integration candidate is not eligible") + if selected.clusterArtifact is None: + raise ValueError("Recommended integration lacks an exact cluster artifact") + if core_artifact_reference(selected.cellSelection) != report_cell_selection: + raise ValueError("Recommended integration uses a different cell selection") + if ( + selected.clusterArtifact.scope != "datastore" + or selected.clusterArtifact.assay is not None + ): + raise ValueError( + "Integrated cluster artifacts must be datastore-scoped without assay" + ) + if ( + selected.graphArtifact is None + or selected.graphArtifact.scope != "datastore" + ): + raise ValueError("Integrated graph artifact must be datastore-scoped") + cluster_artifact = selected.clusterArtifact + cluster_column = selected.clusterColumn + graph_assay = None + else: + selected_assay = native_assay or report.fromAssay + primary = assay_reports.get(selected_assay) + if primary is None or primary.recommendedCandidateId is None: + raise ValueError("Selected assay lacks a native tuning recommendation") + selected_native = next( + ( + item + for item in primary.evaluations + if item.candidateId == primary.recommendedCandidateId + ), + None, + ) + if ( + selected_native is None + or selected_native.status != "done" + or not selected_native.eligible + or "clusters" not in selected_native.artifacts + ): + raise ValueError("Primary native recommendation lacks exact clusters") + if ( + core_artifact_reference(selected_native.cellSelection) + != report_cell_selection + ): + raise ValueError( + "Recommended native candidate uses a different cell selection" + ) + cluster_artifact = selected_native.artifacts["clusters"] + cluster_column = selected_native.clusterColumn + graph_assay = selected_assay + finalized = report.model_copy( + update={ + "totalCandidates": ( + sum(len(value.evaluations) for value in assay_reports.values()) + + len(evaluations) + ), + "integrationEvaluations": evaluations, + "recommendedIntegrationId": recommended_integration_id, + "finalClusterColumn": cluster_column, + "finalClusterArtifact": cluster_artifact, + "graphAssay": graph_assay, + "markerAssay": marker_assay, + "finalSelection": final_selection, + } + ) + selected_graph = recommended_integration_id or graph_assay + logger.info( + f"Finalized parameter graph selection: graph={selected_graph!r}, " + f"marker_assay={marker_assay!r}, cluster_column={cluster_column!r}" + ) + return finalized + + +def select_final_parameter_graph( + *, + model: Any, + report: ParameterTuningReport, + integration_evaluations: Sequence[IntegrationCandidateEvaluation], + marker_assay: str, + config: AgentRunConfig | None = None, +) -> ParameterTuningReport: + """Use one bounded provider call to select and attach the final graph.""" + + evaluations = list(integration_evaluations) + if not marker_assay: + raise ValueError("marker_assay must be non-empty") + assay_reports = report.assayReports or {report.fromAssay: report} + if marker_assay not in assay_reports: + raise ValueError(f"Unknown marker assay {marker_assay!r}") + options = final_graph_options(report, evaluations) + if not options: + raise ValueError("No eligible native or integrated graph options are available") + logger.info( + f"Selecting final parameter graph from {len(options)} eligible option(s); " + f"marker_assay={marker_assay!r}" + ) + if len(options) == 1: + option_id, option = next(iter(options.items())) + logger.info( + f"Selecting sole eligible final graph option {option_id!r} " + "without a provider request" + ) + selection = validate_final_graph_selection( + FinalGraphSelection( + status="done", + selectedOptionId=option_id, + markerAssay=marker_assay, + confidence="high", + rationale="The executor produced exactly one eligible graph option.", + evidenceIds=list(option["evidenceIds"]), + limitations=["No alternative eligible final graph required ranking."], + runInfo=AgentRunInfo( + agentName="parameter_tuning_final_graph_deterministic" + ), + ), + report, + integration_evaluations=evaluations, + marker_assay=marker_assay, + ) + else: + run_config = (config or AgentRunConfig()).with_limits( + request_limit=6, + tool_call_limit=5, + output_token_limit=32768, + timeout_seconds=600.0, + ) + try: + logger.info( + f"Requesting final graph selection across {len(options)} " + "eligible options" + ) + execution = run_agent_sync( + model=model, + output_type=FinalGraphSelection, + system_prompt=final_graph_selection_system_prompt(), + user_prompt=final_graph_selection_prompt( + report=report, + integration_evaluations=evaluations, + marker_assay=marker_assay, + ), + deps_type=ParameterTuningDependencies, + deps=ParameterTuningDependencies.get_blank(), + config=run_config, + name="parameter_tuning_final_graph", + output_validator=lambda proposed: validate_final_graph_selection( + proposed, + report, + integration_evaluations=evaluations, + marker_assay=marker_assay, + ), + ) + except (UnexpectedModelBehavior, UsageLimitExceeded) as exc: + option_ids = sorted(options) + logger.warning( + "Final graph selection model run failed within its bounds " + f"({type(exc).__name__}); " + f"requesting input for {len(option_ids)} eligible options" + ) + selection = validate_final_graph_selection( + FinalGraphSelection( + status="needsInput", + markerAssay=marker_assay, + confidence="low", + rationale=( + "The bounded structured final-graph selection was unavailable." + ), + limitations=[ + "No ranking was invented across multiple eligible graphs." + ], + needsInput=FinalGraphNeedsInput( + question="Select one eligible final graph option.", + options=option_ids, + evidenceIds=sorted( + { + evidence_id + for option in options.values() + for evidence_id in option["evidenceIds"] + } + ), + ), + runInfo=getattr( + exc, + "agent_run_info", + AgentRunInfo( + agentName="parameter_tuning_final_graph_needs_input" + ), + ), + ), + report, + integration_evaluations=evaluations, + marker_assay=marker_assay, + ) + else: + if not isinstance(execution.output, FinalGraphSelection): + raise TypeError( + "Final graph selector returned an unexpected output type" + ) + selection = validate_final_graph_selection( + execution.output, + report, + integration_evaluations=evaluations, + marker_assay=marker_assay, + ).model_copy(update={"runInfo": execution.runInfo}) + logger.info( + f"Provider selected final graph option {selection.selectedOptionId!r}" + ) + if selection.status == "needsInput": + needs_input = selection.needsInput or FinalGraphNeedsInput.get_blank() + option_ids = sorted(options) + canonical_question = "Select one eligible final graph option." + canonical_needs_input = needs_input.model_copy( + update={"question": canonical_question, "options": option_ids} + ) + logger.warning( + f"Final parameter graph selection needs input; options={len(option_ids)}" + ) + return report.model_copy( + update={ + "status": "needsInput", + "totalCandidates": ( + sum(len(value.evaluations) for value in assay_reports.values()) + + len(evaluations) + ), + "integrationEvaluations": evaluations, + "markerAssay": marker_assay, + "finalSelection": selection.model_copy( + update={"needsInput": canonical_needs_input} + ), + "needsInput": ParameterTuningNeedsInput( + question=canonical_question, + options=option_ids, + evidenceIds=needs_input.evidenceIds, + ), + } + ) + return finalize_parameter_tuning_selection( + report, + marker_assay=marker_assay, + integration_evaluations=evaluations, + recommended_integration_id=selection.integrationId, + native_assay=selection.nativeAssay, + final_selection=selection, + ) + + +def promote_parameter_candidate( + store: Any, + *, + report: ParameterTuningReport, + normalized: Any, + identity_feature_limit: int = 64, +) -> ParameterCandidateEvaluation: + """Resolve and verify the exact selected native branch without replaying it.""" + + if report.status != "done" or report.recommendedCandidateId is None: + raise ValueError("A completed native tuning recommendation is required") + evaluation = next( + ( + item + for item in report.evaluations + if item.candidateId == report.recommendedCandidateId + ), + None, + ) + if evaluation is None or evaluation.status != "done" or not evaluation.eligible: + raise ValueError("Recommended candidate is not an eligible execution") + if "clusters" not in evaluation.artifacts: + raise ValueError("Recommended candidate lacks an exact cluster artifact") + normalized_ref = core_artifact_reference(normalized) + if ( + not isinstance(normalized_ref, ArtifactRef) + or normalized_ref.kind != "normalized" + or normalized_ref.assay != report.fromAssay + ): + raise ValueError( + "normalized must identify the report's exact normalized assay artifact" + ) + status = store.inspect_artifact(normalized_ref) + if not getattr(status, "exists", True) or not getattr(status, "complete", False): + raise ValueError("normalized artifact is unavailable or incomplete") + raw_selection = (getattr(status, "inputs", None) or {}).get("cell_selection") + if not isinstance(raw_selection, Mapping): + raise ValueError("normalized artifact has no cell-selection input") + normalized_selection = ArtifactRef.from_dict(dict(raw_selection)) + if normalized_selection != core_artifact_reference(evaluation.cellSelection): + raise ValueError( + "Recommended candidate does not match normalized artifact lineage" + ) + if normalized_selection != core_artifact_reference(report.cellSelection): + raise ValueError( + "Parameter tuning report does not match normalized artifact lineage" + ) + if identity_feature_limit < 2: + raise ValueError("identity_feature_limit must be at least two") + logger.info( + f"Resolved parameter candidate {evaluation.candidateId!r} for assay " + f"{report.fromAssay!r} without replay" + ) + return evaluation diff --git a/scarf/agent/persistence.py b/scarf/agent/persistence.py deleted file mode 100644 index 8cfcd463..00000000 --- a/scarf/agent/persistence.py +++ /dev/null @@ -1,1588 +0,0 @@ -"""Immutable JSON persistence for structured Scarf agent workflows. - -Agent records are stored as plain JSON keys beneath an ``agents`` Zarr group in -the active Scarf data group. The namespace group contains metadata only; report -models are not stored as Zarr arrays or encoded chunks. Format version 2 is a -single-writer format with no implicit migration from earlier layouts. -""" - -import hashlib -import json -import re -import time -import uuid -from collections.abc import Mapping -from pathlib import Path -from typing import Any, Literal, cast - -import zarr -from pydantic import Field, field_validator, model_validator -from zarr.core.buffer import default_buffer_prototype -from zarr.core.sync import sync - -from ..datastore.datastore import DataStore -from ..storage.artifacts import inspect_artifact -from ..storage.refs import ArtifactRef -from ..storage.schema import validate_workspace_name -from ..utils.logging import logger -from . import record_io -from .biological_interpretation import BiologicalInterpretationReport -from .config import AgentRunConfig -from .data_enrichment import DataEnrichmentReport -from .experimental_context import ExperimentalContextResult -from .parameter_tuning import ParameterTuningReport -from .types import ( - AgentDataModel, - ArtifactReferenceModel, - ExperimentalBiologyHandoff, - ExperimentalTuningHandoff, - TuningBiologyHandoff, -) - -type AgentName = Literal[ - "data_enrichment", - "experimental_context", - "parameter_tuning", - "biological_interpretation", -] -type AgentReportType = Literal[ - "", - "DataEnrichmentReport", - "ExperimentalContextResult", - "ParameterTuningReport", - "BiologicalInterpretationReport", -] -type AgentReport = ( - DataEnrichmentReport - | ExperimentalContextResult - | ParameterTuningReport - | BiologicalInterpretationReport -) -type AgentPersistenceTarget = str | Path | zarr.Group | DataStore -type AgentWorkflowStatus = Literal["running", "completed", "failed", "abandoned"] -type AgentTerminalStatus = Literal["completed", "failed", "abandoned"] - -_FORMAT = "scarf_agent_reports" -_RUN_ID_PATTERN = re.compile(r"^[a-z0-9][a-z0-9_-]{0,127}$") -_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$") -_REPORT_TYPES: dict[AgentName, type[AgentDataModel]] = { - "data_enrichment": DataEnrichmentReport, - "experimental_context": ExperimentalContextResult, - "parameter_tuning": ParameterTuningReport, - "biological_interpretation": BiologicalInterpretationReport, -} -_AGENT_NAMES: dict[type[AgentDataModel], AgentName] = { - report_type: agent_name for agent_name, report_type in _REPORT_TYPES.items() -} - - -class AgentReportLink(AgentDataModel): - """Immutable identity of one report used as an invocation parent.""" - - type: Literal["agentReportLink"] = "agentReportLink" - workflowRunId: str = "" - workspace: str | None = None - agentName: AgentName = "data_enrichment" - agentRunId: str = "" - contentSha256: str = "" - - @field_validator("workflowRunId", "agentRunId") - @classmethod - def validate_run_ids(cls, value: str) -> str: - if value: - _validate_run_id(value, "run ID") - return value - - @field_validator("workspace") - @classmethod - def validate_workspace(cls, value: str | None) -> str | None: - validate_workspace_name(value) - return value - - @field_validator("contentSha256") - @classmethod - def validate_content_sha256(cls, value: str) -> str: - if value and _SHA256_PATTERN.fullmatch(value) is None: - raise ValueError("contentSha256 must be a lowercase SHA-256 digest") - return value - - @model_validator(mode="after") - def validate_complete_identity(self) -> "AgentReportLink": - if self.workflowRunId or self.agentRunId or self.contentSha256: - if not self.workflowRunId or not self.agentRunId or not self.contentSha256: - raise ValueError("A parent report link requires a complete identity") - return self - - @classmethod - def from_reference(cls, reference: "AgentReportReference") -> "AgentReportLink": - return cls( - workflowRunId=reference.workflowRunId, - workspace=reference.workspace, - agentName=reference.agentName, - agentRunId=reference.agentRunId, - contentSha256=reference.contentSha256, - ) - - @classmethod - def get_blank(cls) -> "AgentReportLink": - return cls() - - @classmethod - def get_example(cls) -> "AgentReportLink": - return cls( - workflowRunId="workflow-1", - agentName="experimental_context", - agentRunId="experimental-run-1", - contentSha256="0" * 64, - ) - - -class AgentInvocation(AgentDataModel): - """Replay-relevant inputs and typed handoffs for one agent invocation.""" - - agentName: AgentName = "data_enrichment" - parentReports: list[AgentReportLink] = Field(default_factory=list) - inputs: dict[str, Any] = Field(default_factory=dict) - artifacts: dict[str, ArtifactReferenceModel] = Field(default_factory=dict) - runConfig: AgentRunConfig = Field(default_factory=AgentRunConfig) - experimentalTuningHandoff: ExperimentalTuningHandoff | None = None - experimentalBiologyHandoff: ExperimentalBiologyHandoff | None = None - tuningBiologyHandoff: TuningBiologyHandoff | None = None - - @model_validator(mode="after") - def validate_parent_reports(self) -> "AgentInvocation": - identities = [ - (parent.workflowRunId, parent.agentName, parent.agentRunId) - for parent in self.parentReports - ] - if len(identities) != len(set(identities)): - raise ValueError("parentReports must not contain duplicate reports") - return self - - @classmethod - def get_blank(cls) -> "AgentInvocation": - return cls() - - @classmethod - def get_example(cls) -> "AgentInvocation": - cell_selection = ArtifactReferenceModel( - scope="datastore", - kind="cell_selection", - artifactId="c" * 64, - ) - return cls( - agentName="parameter_tuning", - parentReports=[AgentReportLink.get_example()], - inputs={ - "fromAssay": "RNA", - "cellSelection": cell_selection.model_dump(mode="json"), - }, - artifacts={"cellSelection": cell_selection}, - runConfig=AgentRunConfig.get_example(), - experimentalTuningHandoff=ExperimentalTuningHandoff( - cellSelection=cell_selection, - batchAction="skip", - ), - ) - - -class AgentReportReference(AgentDataModel): - """Stable identity for one immutable agent report.""" - - type: Literal["agentReport"] = "agentReport" - workflowRunId: str = "" - workspace: str | None = None - agentName: AgentName = "data_enrichment" - agentRunId: str = "" - reportType: AgentReportType = "" - executionRunId: str = "" - createdAtNs: int = Field(default=0, ge=0, strict=True) - complete: bool = Field(default=False, strict=True) - parentReports: list[AgentReportLink] = Field(default_factory=list) - contentSha256: str = "" - - @field_validator("workflowRunId", "agentRunId") - @classmethod - def validate_run_ids(cls, value: str) -> str: - if value: - _validate_run_id(value, "run ID") - return value - - @field_validator("workspace") - @classmethod - def validate_workspace(cls, value: str | None) -> str | None: - validate_workspace_name(value) - return value - - @field_validator("contentSha256") - @classmethod - def validate_content_sha256(cls, value: str) -> str: - if value and _SHA256_PATTERN.fullmatch(value) is None: - raise ValueError("contentSha256 must be a lowercase SHA-256 digest") - return value - - @model_validator(mode="after") - def validate_complete_identity(self) -> "AgentReportReference": - has_identity = bool( - self.workflowRunId - or self.agentRunId - or self.reportType - or self.createdAtNs - or self.complete - or self.contentSha256 - ) - if has_identity and ( - not self.workflowRunId - or not self.agentRunId - or not self.reportType - or self.createdAtNs < 1 - or not self.complete - or not self.contentSha256 - ): - raise ValueError("An agent report reference requires a complete identity") - return self - - @classmethod - def get_blank(cls) -> "AgentReportReference": - return cls() - - @classmethod - def get_example(cls) -> "AgentReportReference": - return cls( - workflowRunId="workflow-1", - agentName="data_enrichment", - agentRunId="agent-run-1", - reportType="DataEnrichmentReport", - executionRunId="provider-run-1", - createdAtNs=1, - complete=True, - contentSha256="0" * 64, - ) - - -class AgentReportRecord(AgentDataModel): - """Complete JSON envelope for one immutable report and its invocation.""" - - recordType: Literal["agentReport"] = "agentReport" - formatVersion: Literal[2] = 2 - reference: AgentReportReference = Field(default_factory=AgentReportReference) - invocation: AgentInvocation = Field(default_factory=AgentInvocation) - report: dict[str, Any] = Field(default_factory=dict) - - @model_validator(mode="after") - def validate_identity(self) -> "AgentReportRecord": - if self.reference.agentName != self.invocation.agentName: - raise ValueError("Report reference and invocation agent names differ") - if self.reference.parentReports != self.invocation.parentReports: - raise ValueError("Report reference and invocation parents differ") - if any( - parent.workflowRunId == self.reference.workflowRunId - and parent.agentName == self.reference.agentName - and parent.agentRunId == self.reference.agentRunId - for parent in self.invocation.parentReports - ): - raise ValueError("An agent report cannot cite itself as a parent") - return self - - @classmethod - def get_blank(cls) -> "AgentReportRecord": - return cls() - - @classmethod - def get_example(cls) -> "AgentReportRecord": - report = DataEnrichmentReport.get_example() - return cls( - reference=AgentReportReference.get_example(), - invocation=AgentInvocation(agentName="data_enrichment"), - report=report.model_dump(mode="json"), - ) - - -class AgentWorkflowRun(AgentDataModel): - """One dataset-bound workflow and its immutable report records.""" - - type: Literal["agentWorkflowRun"] = "agentWorkflowRun" - formatVersion: Literal[2] = 2 - workflowRunId: str = "" - workspace: str | None = None - createdAtNs: int = Field(default=0, ge=0, strict=True) - finalizedAtNs: int = Field(default=0, ge=0, strict=True) - status: AgentWorkflowStatus = "running" - finalizationMessage: str = "" - analysisStore: str = "" - datasetFingerprints: dict[str, str] = Field(default_factory=dict) - reports: list[AgentReportReference] = Field(default_factory=list) - - @field_validator("workflowRunId") - @classmethod - def validate_workflow_run_id(cls, value: str) -> str: - if value: - _validate_run_id(value, "workflowRunId") - return value - - @field_validator("workspace") - @classmethod - def validate_workspace(cls, value: str | None) -> str | None: - validate_workspace_name(value) - return value - - @field_validator("datasetFingerprints") - @classmethod - def validate_dataset_fingerprints(cls, value: dict[str, str]) -> dict[str, str]: - if any(not assay or not fingerprint for assay, fingerprint in value.items()): - raise ValueError("Dataset fingerprint names and values must be non-empty") - return dict(sorted(value.items())) - - @model_validator(mode="after") - def validate_lifecycle(self) -> "AgentWorkflowRun": - if self.workflowRunId and self.createdAtNs < 1: - raise ValueError("A workflow requires a positive createdAtNs") - if self.workflowRunId and not self.datasetFingerprints: - raise ValueError("A workflow requires exact dataset fingerprints") - if self.status == "running" and self.finalizedAtNs != 0: - raise ValueError("A running workflow cannot have finalizedAtNs") - if self.status == "running" and self.finalizationMessage: - raise ValueError("A running workflow cannot have a finalizationMessage") - if self.status != "running" and self.finalizedAtNs < 1: - raise ValueError("A terminal workflow requires finalizedAtNs") - if ( - self.status != "running" - and self.createdAtNs - and self.finalizedAtNs < self.createdAtNs - ): - raise ValueError("finalizedAtNs cannot precede createdAtNs") - return self - - @classmethod - def get_blank(cls) -> "AgentWorkflowRun": - return cls() - - @classmethod - def get_example(cls) -> "AgentWorkflowRun": - return cls( - workflowRunId="workflow-1", - createdAtNs=1, - analysisStore="analysis.zarr", - datasetFingerprints={"RNA": "dataset-1"}, - reports=[AgentReportReference.get_example()], - ) - - -class AgentStoreManifest(AgentDataModel): - """Identity document for one workspace-local agent JSON store.""" - - type: Literal["agentReportStore"] = "agentReportStore" - format: Literal["scarf_agent_reports"] = "scarf_agent_reports" - formatVersion: Literal[2] = 2 - workspace: str | None = None - - @field_validator("workspace") - @classmethod - def validate_workspace(cls, value: str | None) -> str | None: - validate_workspace_name(value) - return value - - @classmethod - def get_blank(cls) -> "AgentStoreManifest": - return cls() - - @classmethod - def get_example(cls) -> "AgentStoreManifest": - return cls(workspace="analysis") - - -class AgentWorkflowFinalization(AgentDataModel): - """Immutable terminal event for a workflow.""" - - recordType: Literal["agentWorkflowFinalization"] = "agentWorkflowFinalization" - formatVersion: Literal[2] = 2 - workflowRunId: str = "" - workspace: str | None = None - status: AgentTerminalStatus = "completed" - finalizedAtNs: int = Field(default=0, ge=0, strict=True) - message: str = "" - - @field_validator("workflowRunId") - @classmethod - def validate_workflow_run_id(cls, value: str) -> str: - if value: - _validate_run_id(value, "workflowRunId") - return value - - @field_validator("workspace") - @classmethod - def validate_workspace(cls, value: str | None) -> str | None: - validate_workspace_name(value) - return value - - @model_validator(mode="after") - def validate_finalization(self) -> "AgentWorkflowFinalization": - if self.workflowRunId and self.finalizedAtNs < 1: - raise ValueError("A finalization requires a positive finalizedAtNs") - return self - - @classmethod - def get_blank(cls) -> "AgentWorkflowFinalization": - return cls() - - @classmethod - def get_example(cls) -> "AgentWorkflowFinalization": - return cls( - workflowRunId="workflow-1", - status="completed", - finalizedAtNs=2, - ) - - -def _validate_run_id(value: str, label: str) -> str: - if _RUN_ID_PATTERN.fullmatch(value) is None: - raise ValueError( - f"{label} must be one safe path component containing 1-128 ASCII " - "lowercase letters, numbers, underscores, or hyphens" - ) - return value - - -def _key_exists(group: zarr.Group, key: str) -> bool: - return record_io.read_key(group, key) is not None - - -def _list_keys(group: zarr.Group, prefix: str) -> list[str]: - if not group.store.supports_listing: - raise NotImplementedError("Agent persistence requires a listable Zarr store") - return record_io.list_keys(group, prefix) - - -def _write_key_once(group: zarr.Group, key: str, payload: bytes) -> None: - store = group.store - if bool(getattr(store, "read_only", False)) or not bool( - getattr(store, "supports_writes", True) - ): - raise PermissionError("Agent persistence target is read-only") - if _key_exists(group, key): - raise FileExistsError(f"Immutable agent record {key!r} already exists") - buffer = default_buffer_prototype().buffer.from_bytes(payload) - sync(store.set_if_not_exists(key, buffer)) - stored = record_io.read_key(group, key) - if stored is None: - raise RuntimeError(f"Agent record {key!r} was not stored") - if stored != payload: - raise FileExistsError( - f"Immutable agent record {key!r} was written by another writer" - ) - - -def _read_json_model( - group: zarr.Group, - key: str, - model_type: type[AgentDataModel], -) -> AgentDataModel: - payload = record_io.read_key(group, key) - if payload is None: - raise FileNotFoundError(key) - try: - decoded = json.loads(payload) - except (UnicodeDecodeError, json.JSONDecodeError) as exc: - raise ValueError(f"Agent JSON record {key!r} is malformed") from exc - try: - return model_type.model_validate(decoded) - except ValueError as exc: - raise ValueError( - f"Agent JSON record {key!r} does not match its Pydantic schema" - ) from exc - - -def _validate_scarf_data_group(group: zarr.Group) -> None: - import_source = group.attrs.get("scarf:import_source") - if import_source is not None and not bool( - group.attrs.get("scarf:import_complete", False) - ): - raise RuntimeError(f"{import_source} import is incomplete") - if group.attrs.get("format") == _FORMAT: - raise ValueError( - "Standalone agent-report sidecars require an explicit migration" - ) - if "cellData" not in group or not isinstance(group["cellData"], zarr.Group): - raise ValueError("Agent persistence requires an existing Scarf data group") - assay_names = [] - for name in sorted(dict.fromkeys(group.group_keys())): - child = group[name] - if isinstance(child, zarr.Group) and "is_assay" in child.attrs: - assay_names.append(name) - if not assay_names: - raise ValueError("Agent persistence requires at least one Scarf assay") - - -def _resolve_target( - target: AgentPersistenceTarget, - *, - write: bool, - workspace: str | None, -) -> tuple[zarr.Group, DataStore | None, str | None, str]: - validate_workspace_name(workspace) - datastore: DataStore | None = None - analysis_store = "" - if isinstance(target, DataStore): - datastore = target - if workspace is not None and workspace != target.workspace: - raise ValueError("workspace does not match the DataStore workspace") - if write and target.zarr_mode != "r+": - raise PermissionError("Agent persistence requires a writable DataStore") - group = target.zw - resolved_workspace = target.workspace - analysis_store = str(target.zarr_loc) - elif isinstance(target, zarr.Group): - target_path = str(getattr(target, "path", "")).strip("/") - if workspace is None: - group = target - resolved_workspace = target_path or None - if resolved_workspace is not None: - validate_workspace_name(resolved_workspace) - elif target_path == workspace: - group = target - resolved_workspace = workspace - elif not target_path and workspace in target: - child = target[workspace] - if not isinstance(child, zarr.Group): - raise TypeError(f"Workspace {workspace!r} is not a Zarr group") - group = child - resolved_workspace = workspace - else: - raise ValueError("workspace does not match the supplied Zarr group") - else: - location = str(target) - if isinstance(target, Path) and not target.exists(): - raise FileNotFoundError(target) - root = zarr.open_group(location, mode="r+" if write else "r") - analysis_store = location - if workspace is None: - group = root - resolved_workspace = None - else: - if workspace not in root: - raise KeyError(f"Unknown Scarf workspace {workspace!r}") - child = root[workspace] - if not isinstance(child, zarr.Group): - raise TypeError(f"Workspace {workspace!r} is not a Zarr group") - group = child - resolved_workspace = workspace - _validate_scarf_data_group(group) - return group, datastore, resolved_workspace, analysis_store - - -def _live_dataset_fingerprints( - group: zarr.Group, - datastore: DataStore | None, - *, - ensure: bool, -) -> dict[str, str]: - if datastore is not None: - assay_names = list(datastore.assay_names) - fingerprints = { - assay_name: ( - datastore._ensure_dataset_fingerprint(assay_name) - if ensure - else str( - datastore._get_assay(assay_name).attrs.get("dataset_fingerprint") - or "" - ) - ) - for assay_name in assay_names - } - else: - assay_names = [] - fingerprints = {} - for name in sorted(dict.fromkeys(group.group_keys())): - child = group[name] - if isinstance(child, zarr.Group) and "is_assay" in child.attrs: - assay_names.append(name) - fingerprints[name] = str(child.attrs.get("dataset_fingerprint") or "") - if not assay_names: - raise ValueError("Dataset binding requires at least one assay") - missing = [name for name in assay_names if not fingerprints[name]] - if missing: - raise ValueError( - "Dataset fingerprints are missing for assays " + repr(sorted(missing)) - ) - return dict(sorted(fingerprints.items())) - - -def _validate_dataset_binding( - stored: dict[str, str], - observed: dict[str, str], -) -> None: - if stored != observed: - raise ValueError( - "Agent workflow dataset fingerprints do not match the current store: " - f"stored={stored!r}, observed={observed!r}" - ) - - -def _agents_prefix(group: zarr.Group) -> str: - return record_io.join_key(str(getattr(group, "path", "")), "agents") - - -def _manifest_key(prefix: str) -> str: - return record_io.join_key(prefix, "store.json") - - -def _workflow_prefix(prefix: str, workflow_run_id: str) -> str: - return record_io.join_key(prefix, "runs", workflow_run_id) - - -def _workflow_key(prefix: str, workflow_run_id: str) -> str: - return record_io.join_key( - _workflow_prefix(prefix, workflow_run_id), - "workflow.json", - ) - - -def _finalization_key(prefix: str, workflow_run_id: str) -> str: - return record_io.join_key( - _workflow_prefix(prefix, workflow_run_id), - "finalization.json", - ) - - -def _report_key( - prefix: str, - workflow_run_id: str, - agent_name: AgentName, - agent_run_id: str, -) -> str: - return record_io.join_key( - _workflow_prefix(prefix, workflow_run_id), - agent_name, - agent_run_id, - "report.json", - ) - - -def _open_agent_store( - group: zarr.Group, - *, - workspace: str | None, - initialize: bool, -) -> str: - prefix = _agents_prefix(group) - if "agents" in group: - node = group["agents"] - if ( - isinstance(node, zarr.Group) - and node.attrs.get("format") == _FORMAT - and node.attrs.get("format_version") == 1 - ): - raise ValueError( - "Zarr-backed agent report format version 1 requires an explicit " - "migration" - ) - if not isinstance(node, zarr.Group): - raise ValueError("The agents namespace must be a Zarr group") - if node.attrs.get("format") != _FORMAT or node.attrs.get("format_version") != 2: - raise ValueError( - "The agents namespace collides with an unrecognized Zarr group" - ) - else: - existing_keys = _list_keys(group, prefix) - if existing_keys: - raise ValueError( - "A plain-JSON agents hierarchy without Zarr group metadata " - "requires an explicit migration" - ) - if not initialize: - raise FileNotFoundError(_manifest_key(prefix)) - if bool(getattr(group.store, "read_only", False)) or not bool( - getattr(group.store, "supports_writes", True) - ): - raise PermissionError("Agent persistence target is read-only") - group.create_group( - "agents", - attributes={"format": _FORMAT, "format_version": 2}, - ) - - manifest_key = _manifest_key(prefix) - payload = record_io.read_key(group, manifest_key) - if payload is None: - if not initialize: - raise FileNotFoundError(manifest_key) - metadata_keys = { - record_io.join_key(prefix, "zarr.json"), - record_io.join_key(prefix, ".zgroup"), - record_io.join_key(prefix, ".zattrs"), - } - if any(key not in metadata_keys for key in _list_keys(group, prefix)): - raise ValueError("Refusing to initialize a non-empty agents hierarchy") - manifest = AgentStoreManifest(workspace=workspace) - _write_key_once( - group, - manifest_key, - record_io.display_json_bytes(manifest.model_dump(mode="json")), - ) - manifest = cast( - AgentStoreManifest, - _read_json_model(group, manifest_key, AgentStoreManifest), - ) - if manifest.workspace != workspace: - raise ValueError("Agent store workspace does not match the active data group") - return prefix - - -def _load_workflow_record( - group: zarr.Group, - prefix: str, - workflow_run_id: str, - workspace: str | None, -) -> AgentWorkflowRun: - workflow_run_id = _validate_run_id(workflow_run_id, "workflow_run_id") - key = _workflow_key(prefix, workflow_run_id) - try: - workflow = cast( - AgentWorkflowRun, - _read_json_model(group, key, AgentWorkflowRun), - ) - except FileNotFoundError as exc: - raise KeyError(f"Unknown agent workflow {workflow_run_id!r}") from exc - if workflow.workflowRunId != workflow_run_id: - raise ValueError("Agent workflow identity does not match its path") - if workflow.workspace != workspace: - raise ValueError("Agent workflow workspace does not match its path") - if ( - workflow.status != "running" - or workflow.finalizedAtNs != 0 - or workflow.finalizationMessage - or workflow.reports - ): - raise ValueError( - "Immutable workflow.json must contain the running identity only" - ) - finalization_payload = record_io.read_key( - group, - _finalization_key(prefix, workflow_run_id), - ) - if finalization_payload is None: - return workflow - finalization = cast( - AgentWorkflowFinalization, - _read_json_model( - group, - _finalization_key(prefix, workflow_run_id), - AgentWorkflowFinalization, - ), - ) - if finalization.workflowRunId != workflow_run_id: - raise ValueError("Agent workflow finalization identity does not match its path") - if finalization.workspace != workspace: - raise ValueError( - "Agent workflow finalization workspace does not match its path" - ) - return AgentWorkflowRun.model_validate( - { - **workflow.model_dump(mode="json"), - "status": finalization.status, - "finalizedAtNs": finalization.finalizedAtNs, - "finalizationMessage": finalization.message, - } - ) - - -def _report_checksum(record: AgentReportRecord) -> str: - reference = record.reference.model_dump( - mode="json", - exclude={"contentSha256"}, - ) - content = { - "recordType": record.recordType, - "formatVersion": record.formatVersion, - "reference": reference, - "invocation": record.invocation.model_dump(mode="json"), - "report": record.report, - } - return hashlib.sha256(record_io.canonical_json_bytes(content)).hexdigest() - - -def _load_report_record_at( - group: zarr.Group, - key: str, - *, - workflow_run_id: str, - agent_name: AgentName, - agent_run_id: str, - workspace: str | None, -) -> AgentReportRecord: - record = cast( - AgentReportRecord, - _read_json_model(group, key, AgentReportRecord), - ) - reference = record.reference - expected_type = cast(AgentReportType, _REPORT_TYPES[agent_name].__name__) - if ( - reference.workflowRunId != workflow_run_id - or reference.agentName != agent_name - or reference.agentRunId != agent_run_id - ): - raise ValueError("Agent report identity does not match its path") - if reference.workspace != workspace: - raise ValueError("Agent report workspace does not match its path") - if reference.reportType != expected_type: - raise ValueError( - f"Agent report type must be {expected_type!r} for {agent_name!r}" - ) - if not reference.complete: - raise ValueError("Atomic agent report records must be complete") - if _report_checksum(record) != reference.contentSha256: - raise ValueError( - "Agent report content checksum does not match its JSON payload" - ) - report_type = _REPORT_TYPES[agent_name] - try: - report = cast(AgentReport, report_type.model_validate(record.report)) - except ValueError as exc: - raise ValueError( - "Agent report payload does not match its Pydantic schema" - ) from exc - if report.runInfo.runId != reference.executionRunId: - raise ValueError("Agent report execution ID does not match its JSON payload") - return record - - -def _report_references( - group: zarr.Group, - prefix: str, - workflow_run_id: str, - workspace: str | None, -) -> list[AgentReportReference]: - run_prefix = _workflow_prefix(prefix, workflow_run_id) - references: list[AgentReportReference] = [] - for key in _list_keys(group, run_prefix): - relative = key.removeprefix(f"{run_prefix}/") - if relative in {"workflow.json", "finalization.json"}: - continue - parts = relative.split("/") - if len(parts) != 3 or parts[2] != "report.json": - raise ValueError(f"Unexpected agent workflow record {key!r}") - raw_agent_name, agent_run_id, _filename = parts - if raw_agent_name not in _REPORT_TYPES: - raise ValueError(f"Unknown agent name {raw_agent_name!r}") - agent_name = raw_agent_name - _validate_run_id(agent_run_id, "agent_run_id") - record = _load_report_record_at( - group, - key, - workflow_run_id=workflow_run_id, - agent_name=agent_name, - agent_run_id=agent_run_id, - workspace=workspace, - ) - _validate_invocation( - group, - prefix, - workflow_run_id, - workspace, - agent_name, - record.invocation, - ) - references.append(record.reference) - return sorted( - references, - key=lambda item: (item.createdAtNs, item.agentName, item.agentRunId), - ) - - -def _load_parent_records( - group: zarr.Group, - prefix: str, - workflow_run_id: str, - workspace: str | None, - invocation: AgentInvocation, -) -> dict[AgentName, list[AgentReportRecord]]: - parents: dict[AgentName, list[AgentReportRecord]] = { - name: [] for name in _REPORT_TYPES - } - for link in invocation.parentReports: - if link.workflowRunId != workflow_run_id: - raise ValueError("Parent reports must belong to the same workflow") - if link.workspace != workspace: - raise ValueError("Parent reports must belong to the same workspace") - key = _report_key( - prefix, - workflow_run_id, - link.agentName, - link.agentRunId, - ) - try: - record = _load_report_record_at( - group, - key, - workflow_run_id=workflow_run_id, - agent_name=link.agentName, - agent_run_id=link.agentRunId, - workspace=workspace, - ) - except FileNotFoundError as exc: - raise ValueError( - f"Unknown parent agent report {link.agentRunId!r}" - ) from exc - if AgentReportLink.from_reference(record.reference) != link: - raise ValueError("Parent report link does not match the stored report") - parents[link.agentName].append(record) - return parents - - -def _has_handoff_parent( - *, - label: str, - supplied: AgentDataModel | None, - parent_records: list[AgentReportRecord], -) -> bool: - if supplied is None and parent_records: - raise ValueError(f"{label} is required when its parent report is cited") - if supplied is not None and len(parent_records) != 1: - raise ValueError(f"{label} requires exactly one matching parent report") - return supplied is not None - - -def _selection_descends_from( - group: zarr.Group, - supplied: ArtifactReferenceModel | None, - expected: ArtifactReferenceModel | None, -) -> bool: - if supplied is None or expected is None: - return supplied == expected - current = ArtifactRef( - scope=supplied.scope, - assay=supplied.assay, - kind=supplied.kind, - artifact_id=supplied.artifactId, - ) - ancestor = ArtifactRef( - scope=expected.scope, - assay=expected.assay, - kind=expected.kind, - artifact_id=expected.artifactId, - ) - if current.kind != "cell_selection" or ancestor.kind != "cell_selection": - return False - visited: set[ArtifactRef] = set() - while current not in visited: - if current == ancestor: - return True - visited.add(current) - status = inspect_artifact(group, current) - if not status.exists or not status.complete: - return False - raw_parent = (status.inputs or {}).get("prior_cell_selection") - if not isinstance(raw_parent, Mapping): - return False - try: - current = ArtifactRef.from_dict(dict(raw_parent)) - except (KeyError, TypeError, ValueError): - return False - return False - - -def _validate_projected_experimental_handoff( - group: zarr.Group, - invocation: AgentInvocation, - *, - label: str, - supplied: ExperimentalTuningHandoff | ExperimentalBiologyHandoff, - expected: ExperimentalTuningHandoff | ExperimentalBiologyHandoff, -) -> None: - supplied_selection = supplied.cellSelection - if invocation.artifacts.get("cellSelection") != supplied_selection: - raise ValueError(f"{label} cellSelection must be an invocation artifact") - if not _selection_descends_from( - group, - supplied_selection, - expected.cellSelection, - ): - raise ValueError( - f"{label} cellSelection does not descend from the cited parent report" - ) - if supplied != expected.model_copy(update={"cellSelection": supplied_selection}): - raise ValueError(f"{label} does not match the cited parent report") - - -def _validate_invocation( - group: zarr.Group, - prefix: str, - workflow_run_id: str, - workspace: str | None, - agent_name: AgentName, - invocation: AgentInvocation, -) -> None: - if invocation.agentName != agent_name: - raise ValueError("Invocation agentName does not match the report type") - if not invocation.inputs: - raise ValueError("Invocation inputs must record the agent call arguments") - parents = _load_parent_records( - group, - prefix, - workflow_run_id, - workspace, - invocation, - ) - if agent_name == "parameter_tuning": - experimental_parents = parents["experimental_context"] - if _has_handoff_parent( - label="experimentalTuningHandoff", - supplied=invocation.experimentalTuningHandoff, - parent_records=experimental_parents, - ): - expected = ExperimentalContextResult.model_validate( - experimental_parents[0].report - ).to_parameter_tuning_handoff() - assert invocation.experimentalTuningHandoff is not None - _validate_projected_experimental_handoff( - group, - invocation, - label="experimentalTuningHandoff", - supplied=invocation.experimentalTuningHandoff, - expected=expected, - ) - elif invocation.experimentalTuningHandoff is not None: - raise ValueError("experimentalTuningHandoff is only valid for parameter_tuning") - - if agent_name == "biological_interpretation": - experimental_parents = parents["experimental_context"] - if len(experimental_parents) > 1: - raise ValueError( - "Biological Interpretation accepts at most one Experimental " - "Context parent report" - ) - if invocation.experimentalBiologyHandoff is not None: - if len(experimental_parents) != 1: - raise ValueError( - "experimentalBiologyHandoff requires exactly one matching " - "parent report" - ) - assert invocation.experimentalBiologyHandoff is not None - expected_experimental = ExperimentalContextResult.model_validate( - experimental_parents[0].report - ).to_biological_handoff( - invocation.experimentalBiologyHandoff.conditionColumn - ) - _validate_projected_experimental_handoff( - group, - invocation, - label="experimentalBiologyHandoff", - supplied=invocation.experimentalBiologyHandoff, - expected=expected_experimental, - ) - tuning_parents = parents["parameter_tuning"] - if _has_handoff_parent( - label="tuningBiologyHandoff", - supplied=invocation.tuningBiologyHandoff, - parent_records=tuning_parents, - ): - expected_tuning = ParameterTuningReport.model_validate( - tuning_parents[0].report - ).to_biological_handoff() - if invocation.tuningBiologyHandoff != expected_tuning: - raise ValueError( - "tuningBiologyHandoff does not match the cited parent report" - ) - elif ( - invocation.experimentalBiologyHandoff is not None - or invocation.tuningBiologyHandoff is not None - ): - raise ValueError( - "Biology handoffs are only valid for biological_interpretation" - ) - - -def _resolved_workflow( - group: zarr.Group, - datastore: DataStore | None, - prefix: str, - workflow_run_id: str, - workspace: str | None, - *, - ensure_fingerprints: bool, -) -> AgentWorkflowRun: - workflow = _load_workflow_record( - group, - prefix, - workflow_run_id, - workspace, - ) - observed = _live_dataset_fingerprints( - group, - datastore, - ensure=ensure_fingerprints, - ) - _validate_dataset_binding(workflow.datasetFingerprints, observed) - return workflow - - -def create_agent_workflow( - target: AgentPersistenceTarget, - *, - workflow_run_id: str | None = None, - analysis_store: str = "", - dataset_fingerprints: dict[str, str] | None = None, - workspace: str | None = None, -) -> AgentWorkflowRun: - """Create an immutable, dataset-bound workflow in the active data group.""" - if not isinstance(analysis_store, str): - raise TypeError("analysis_store must be a string") - group, datastore, resolved_workspace, inferred_store = _resolve_target( - target, - write=True, - workspace=workspace, - ) - observed = _live_dataset_fingerprints(group, datastore, ensure=True) - if dataset_fingerprints is not None: - if not isinstance(dataset_fingerprints, dict) or any( - not isinstance(assay, str) or not isinstance(fingerprint, str) - for assay, fingerprint in dataset_fingerprints.items() - ): - raise TypeError("dataset_fingerprints must map assay names to strings") - supplied = dict(sorted(dataset_fingerprints.items())) - _validate_dataset_binding(supplied, observed) - resolved_run_id = _validate_run_id( - uuid.uuid4().hex if workflow_run_id is None else workflow_run_id, - "workflow_run_id", - ) - prefix = _open_agent_store( - group, - workspace=resolved_workspace, - initialize=True, - ) - key = _workflow_key(prefix, resolved_run_id) - if _list_keys(group, _workflow_prefix(prefix, resolved_run_id)): - raise FileExistsError(f"Agent workflow {resolved_run_id!r} already exists") - workflow = AgentWorkflowRun( - workflowRunId=resolved_run_id, - workspace=resolved_workspace, - createdAtNs=time.time_ns(), - analysisStore=analysis_store or inferred_store, - datasetFingerprints=observed, - ) - _write_key_once( - group, - key, - record_io.display_json_bytes(workflow.model_dump(mode="json")), - ) - logger.info( - f"Created agent workflow {resolved_run_id}: workspace=" - f"{resolved_workspace or 'root'}, assays={len(observed)}" - ) - return workflow - - -def save_agent_report( - target: AgentPersistenceTarget, - workflow_run_id: str, - report: AgentReport, - *, - invocation: AgentInvocation, - agent_run_id: str | None = None, - workspace: str | None = None, -) -> AgentReportReference: - """Persist one immutable report together with its replay-relevant inputs.""" - agent_name = _AGENT_NAMES.get(type(report)) - if agent_name is None: - raise TypeError("report must be one of the four Scarf agent report models") - if not isinstance(invocation, AgentInvocation): - raise TypeError("invocation must be an AgentInvocation") - resolved_workflow_id = _validate_run_id(workflow_run_id, "workflow_run_id") - resolved_agent_run_id = _validate_run_id( - uuid.uuid4().hex if agent_run_id is None else agent_run_id, - "agent_run_id", - ) - group, datastore, resolved_workspace, _analysis_store = _resolve_target( - target, - write=True, - workspace=workspace, - ) - prefix = _open_agent_store( - group, - workspace=resolved_workspace, - initialize=False, - ) - workflow = _resolved_workflow( - group, - datastore, - prefix, - resolved_workflow_id, - resolved_workspace, - ensure_fingerprints=True, - ) - if workflow.status != "running": - raise RuntimeError( - f"Cannot save a report to a {workflow.status!r} agent workflow" - ) - _validate_invocation( - group, - prefix, - resolved_workflow_id, - resolved_workspace, - agent_name, - invocation, - ) - key = _report_key( - prefix, - resolved_workflow_id, - agent_name, - resolved_agent_run_id, - ) - if _key_exists(group, key): - raise FileExistsError( - f"Agent report {resolved_agent_run_id!r} already exists for {agent_name!r}" - ) - reference = AgentReportReference( - workflowRunId=resolved_workflow_id, - workspace=resolved_workspace, - agentName=agent_name, - agentRunId=resolved_agent_run_id, - reportType=cast(AgentReportType, type(report).__name__), - executionRunId=report.runInfo.runId, - createdAtNs=time.time_ns(), - complete=True, - parentReports=list(invocation.parentReports), - contentSha256="0" * 64, - ) - record = AgentReportRecord( - reference=reference, - invocation=invocation, - report=report.model_dump(mode="json"), - ) - checksum = _report_checksum(record) - reference = reference.model_copy(update={"contentSha256": checksum}) - record = record.model_copy(update={"reference": reference}) - _write_key_once( - group, - key, - record_io.display_json_bytes(record.model_dump(mode="json")), - ) - stored = _load_report_record_at( - group, - key, - workflow_run_id=resolved_workflow_id, - agent_name=agent_name, - agent_run_id=resolved_agent_run_id, - workspace=resolved_workspace, - ) - logger.info( - f"Saved {agent_name} report {resolved_agent_run_id} for workflow " - f"{resolved_workflow_id}: status={getattr(report, 'status', 'done')}, " - f"parents={len(invocation.parentReports)}" - ) - return stored.reference - - -def _validate_supplied_reference( - supplied: AgentReportReference, - stored: AgentReportReference, -) -> None: - if supplied.workflowRunId != stored.workflowRunId: - raise ValueError( - "Agent report reference workflow does not match stored metadata" - ) - if supplied.workspace != stored.workspace: - raise ValueError( - "Agent report reference workspace does not match stored metadata" - ) - if ( - supplied.agentName != stored.agentName - or supplied.agentRunId != stored.agentRunId - ): - raise ValueError( - "Agent report reference identity does not match stored metadata" - ) - if supplied.reportType and supplied.reportType != stored.reportType: - raise ValueError("Agent report reference type does not match stored metadata") - if supplied.executionRunId and supplied.executionRunId != stored.executionRunId: - raise ValueError("Agent report reference execution ID does not match metadata") - if supplied.createdAtNs and supplied.createdAtNs != stored.createdAtNs: - raise ValueError( - "Agent report reference timestamp does not match stored metadata" - ) - if supplied.parentReports and supplied.parentReports != stored.parentReports: - raise ValueError("Agent report reference parents do not match stored metadata") - if supplied.contentSha256 and supplied.contentSha256 != stored.contentSha256: - raise ValueError( - "Agent report reference checksum does not match stored metadata" - ) - - -def load_agent_record( - target: AgentPersistenceTarget, - reference: AgentReportReference, - *, - workspace: str | None = None, -) -> AgentReportRecord: - """Load and validate a report envelope, including lineage and inputs.""" - group, datastore, resolved_workspace, _analysis_store = _resolve_target( - target, - write=False, - workspace=workspace, - ) - prefix = _open_agent_store( - group, - workspace=resolved_workspace, - initialize=False, - ) - _resolved_workflow( - group, - datastore, - prefix, - reference.workflowRunId, - resolved_workspace, - ensure_fingerprints=False, - ) - key = _report_key( - prefix, - reference.workflowRunId, - reference.agentName, - reference.agentRunId, - ) - try: - record = _load_report_record_at( - group, - key, - workflow_run_id=reference.workflowRunId, - agent_name=reference.agentName, - agent_run_id=reference.agentRunId, - workspace=resolved_workspace, - ) - except FileNotFoundError as exc: - raise KeyError(f"Unknown agent report at {key!r}") from exc - _validate_invocation( - group, - prefix, - reference.workflowRunId, - resolved_workspace, - reference.agentName, - record.invocation, - ) - _validate_supplied_reference(reference, record.reference) - return record - - -def load_agent_report( - target: AgentPersistenceTarget, - reference: AgentReportReference, - *, - workspace: str | None = None, -) -> AgentReport: - """Load and strictly revalidate one persisted agent report.""" - record = load_agent_record(target, reference, workspace=workspace) - report_type = _REPORT_TYPES[record.reference.agentName] - return cast(AgentReport, report_type.model_validate(record.report)) - - -def list_agent_reports( - target: AgentPersistenceTarget, - workflow_run_id: str, - *, - agent_name: AgentName | None = None, - include_incomplete: bool = False, - workspace: str | None = None, -) -> list[AgentReportReference]: - """List atomic report references for one dataset-bound workflow.""" - del include_incomplete - if agent_name is not None and agent_name not in _REPORT_TYPES: - raise ValueError(f"Unknown agent name {agent_name!r}") - group, datastore, resolved_workspace, _analysis_store = _resolve_target( - target, - write=False, - workspace=workspace, - ) - prefix = _open_agent_store( - group, - workspace=resolved_workspace, - initialize=False, - ) - _resolved_workflow( - group, - datastore, - prefix, - workflow_run_id, - resolved_workspace, - ensure_fingerprints=False, - ) - references = _report_references( - group, - prefix, - workflow_run_id, - resolved_workspace, - ) - if agent_name is not None: - references = [item for item in references if item.agentName == agent_name] - return references - - -def load_agent_workflow( - target: AgentPersistenceTarget, - workflow_run_id: str, - *, - include_incomplete: bool = False, - workspace: str | None = None, -) -> AgentWorkflowRun: - """Load one workflow, its lifecycle state, and all report references.""" - del include_incomplete - group, datastore, resolved_workspace, _analysis_store = _resolve_target( - target, - write=False, - workspace=workspace, - ) - prefix = _open_agent_store( - group, - workspace=resolved_workspace, - initialize=False, - ) - workflow = _resolved_workflow( - group, - datastore, - prefix, - workflow_run_id, - resolved_workspace, - ensure_fingerprints=False, - ) - reports = _report_references( - group, - prefix, - workflow_run_id, - resolved_workspace, - ) - if workflow.status == "completed" and not reports: - raise ValueError("A completed workflow must contain at least one report") - return AgentWorkflowRun.model_validate( - { - **workflow.model_dump(mode="json"), - "reports": reports, - } - ) - - -def list_agent_workflows( - target: AgentPersistenceTarget, - *, - include_incomplete: bool = False, - workspace: str | None = None, -) -> list[AgentWorkflowRun]: - """List terminal workflows, optionally including workflows still running.""" - group, datastore, resolved_workspace, _analysis_store = _resolve_target( - target, - write=False, - workspace=workspace, - ) - prefix = _open_agent_store( - group, - workspace=resolved_workspace, - initialize=False, - ) - runs_prefix = record_io.join_key(prefix, "runs") - workflow_ids: set[str] = set() - for key in _list_keys(group, runs_prefix): - relative = key.removeprefix(f"{runs_prefix}/") - parts = relative.split("/") - if len(parts) < 2: - raise ValueError(f"Unexpected agent workflow key {key!r}") - workflow_ids.add(_validate_run_id(parts[0], "workflow_run_id")) - workflows = [ - load_agent_workflow( - target, - workflow_run_id, - workspace=workspace, - ) - for workflow_run_id in sorted(workflow_ids) - ] - if not include_incomplete: - workflows = [item for item in workflows if item.status != "running"] - return sorted( - workflows, - key=lambda item: (item.createdAtNs, item.workflowRunId), - ) - - -def finalize_agent_workflow( - target: AgentPersistenceTarget, - workflow_run_id: str, - *, - status: AgentTerminalStatus, - message: str = "", - workspace: str | None = None, -) -> AgentWorkflowRun: - """Write the one terminal event allowed for a running workflow.""" - if status not in {"completed", "failed", "abandoned"}: - raise ValueError("status must be completed, failed, or abandoned") - if not isinstance(message, str): - raise TypeError("message must be a string") - group, datastore, resolved_workspace, _analysis_store = _resolve_target( - target, - write=True, - workspace=workspace, - ) - prefix = _open_agent_store( - group, - workspace=resolved_workspace, - initialize=False, - ) - workflow = _resolved_workflow( - group, - datastore, - prefix, - workflow_run_id, - resolved_workspace, - ensure_fingerprints=True, - ) - if workflow.status != "running": - raise FileExistsError( - f"Agent workflow {workflow_run_id!r} is already {workflow.status!r}" - ) - reports = _report_references( - group, - prefix, - workflow_run_id, - resolved_workspace, - ) - if status == "completed" and not reports: - raise ValueError("A completed workflow must contain at least one report") - finalization = AgentWorkflowFinalization( - workflowRunId=workflow_run_id, - workspace=resolved_workspace, - status=status, - finalizedAtNs=time.time_ns(), - message=message, - ) - _write_key_once( - group, - _finalization_key(prefix, workflow_run_id), - record_io.display_json_bytes(finalization.model_dump(mode="json")), - ) - finalized = load_agent_workflow( - target, - workflow_run_id, - workspace=workspace, - ) - logger.info( - f"Finalized agent workflow {workflow_run_id}: status={status}, " - f"reports={len(finalized.reports)}" - ) - return finalized - - -__all__ = [ - "AgentInvocation", - "AgentName", - "AgentPersistenceTarget", - "AgentReport", - "AgentReportLink", - "AgentReportRecord", - "AgentReportReference", - "AgentReportType", - "AgentTerminalStatus", - "AgentWorkflowRun", - "AgentWorkflowStatus", - "create_agent_workflow", - "finalize_agent_workflow", - "list_agent_reports", - "list_agent_workflows", - "load_agent_record", - "load_agent_report", - "load_agent_workflow", - "save_agent_report", -] diff --git a/scarf/agent/report/__init__.py b/scarf/agent/report/__init__.py new file mode 100644 index 00000000..befa0df2 --- /dev/null +++ b/scarf/agent/report/__init__.py @@ -0,0 +1,5 @@ +"""Local HTML reports for completed automated Scarf agent workflows.""" + +from .generator import generate_agent_report + +__all__ = ["generate_agent_report"] diff --git a/scarf/agent/report/artifacts.py b/scarf/agent/report/artifacts.py new file mode 100644 index 00000000..4467f5a1 --- /dev/null +++ b/scarf/agent/report/artifacts.py @@ -0,0 +1,190 @@ +"""Read-only adapters from the authoritative stage journal to a local report.""" + +from collections.abc import Mapping +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from ...storage.refs import ArtifactRef +from ...storage.stores import zarr_root_path +from ..types import ArtifactReferenceModel +from .contracts import label, mapping, mappings, texts + +if TYPE_CHECKING: + from ...datastore.datastore import DataStore + + +def _local_root(target: "str | Path | DataStore") -> Path: + """Resolve a local store without accepting remote locations.""" + from ...datastore.datastore import DataStore + + if isinstance(target, DataStore): + location = zarr_root_path(target.z) + if location is None: + raise ValueError("Agent HTML reports require a local filesystem store") + path = Path(location) + elif isinstance(target, Path): + path = target + elif isinstance(target, str): + if "://" in target and not target.startswith("file://"): + raise ValueError("Agent HTML reports require a local filesystem store") + path = Path(target.removeprefix("file://")) + else: + raise TypeError("Agent HTML reports require a local filesystem store") + root = path.expanduser().resolve() + if not root.is_dir(): + raise FileNotFoundError(root) + return root + + +def report_directory(root: Path, run_id: str, workspace: str | None) -> Path: + """Keep derived pages within the exact journal owner's local directory.""" + if not run_id or run_id in {".", ".."} or "/" in run_id or "\\" in run_id: + raise ValueError("Invalid analysis run identifier") + active = root if workspace is None else (root / workspace).resolve() + if not active.is_relative_to(root): + raise ValueError("Analysis workspace resolves outside the store") + destination = (active / "agents" / "orchestrations" / run_id / "report").resolve() + if not destination.is_relative_to(active): + raise ValueError("Analysis report resolves outside the store") + return destination + + +def artifact_ref(value: Any) -> ArtifactRef: + ref = ArtifactReferenceModel.model_validate(value) + return ArtifactRef( + scope=ref.scope, assay=ref.assay, kind=ref.kind, artifact_id=ref.artifactId + ) + + +def scientific_summary(snapshot: Mapping[str, Any]) -> dict[str, Any]: + """Select recorded scientific values without inferring decision rationales.""" + from ..orchestrator.rna import validate_analysis_evidence + + validate_analysis_evidence(snapshot) + final = mapping(snapshot.get("finalAnalysis")) + stages = mappings(snapshot.get("stages")) + decisions = [ + decision for stage in stages for decision in mappings(stage.get("decisions")) + ] + limitations = texts(final.get("limitations")) + texts(snapshot.get("limitations")) + assessments = mappings(snapshot.get("analysisReviews")) + full_assessments = [ + item + for item in assessments + if item["scope"] == "full" and item.get("action") == "accept" + ] + accepted = full_assessments[-1] if full_assessments else {} + selected: dict[str, Any] = {} + qc_profiles: list[dict[str, Any]] = [] + qc_profile_id: Any = None + context: dict[str, Any] = {} + study: dict[str, Any] = {} + for stage in stages: + report = mapping(stage.get("report")) + stage_name = str(stage.get("stage", "")) + if stage_name.startswith("parameter_tuning"): + evaluations = mappings(report.get("evaluations")) + recommended = report.get("recommendedCandidateId") + selected = next( + ( + item + for item in evaluations + if item.get("candidateId") == recommended + ), + selected, + ) + if stage_name == "experimental_context": + context = report + study = mapping(mapping(stage.get("outputs")).get("studyContract")) + qc_profile_id = mapping(report.get("cellQc")).get("profileId") + qc_profiles = mappings(report.get("qcProfiles")) + outputs = mapping(stage.get("outputs")) + for plan_key in ("preprocessingPlan", "resolvedPreprocessingPlan"): + plan = mapping(outputs.get(plan_key)) + if plan: + qc_profile_id = mapping(plan.get("cellQc")).get("profileId") + qc = next( + (item for item in qc_profiles if item.get("profileId") == qc_profile_id), {} + ) + limitations.extend(texts(study.get("limitations"))) + limitations.extend( + str(item["explanation"]) + for item in mappings(accepted.get("populationConcerns")) + if item.get("explanation") + ) + comparison_limits = {} + for comparison in mappings( + mapping(context.get("characterization")).get("comparisons") + ): + if comparison.get("status") != "unsupported": + continue + proposal = mapping(comparison.get("proposal")) + question = ( + f"{label(str(proposal.get('response', 'Study factor')))} against " + + ", ".join( + label(name) for name in texts(proposal.get("explanatoryColumns")) + ) + ) + if proposal.get("conditionedOn"): + question += f" within {label(str(proposal['conditionedOn']))} groups" + reasons = "; ".join( + label(reason) for reason in texts(comparison.get("reasons")) + ) + comparison_limits[str(comparison.get("evidenceId"))] = ( + f"{question}: the requested association was not computed. {reasons}." + ) + displayed_limits = [] + for limitation in limitations: + match = next( + ( + text + for identity, text in comparison_limits.items() + if identity and identity in limitation + ), + None, + ) + displayed_limits.append(match or limitation) + population = mapping( + mapping(accepted.get("populationSupport")).get( + str(accepted.get("selectedCandidateId")) + ) + ) + if population: + for key in ("cellSelection", "clusters"): + reference = population.get(key) + if not isinstance(reference, Mapping): + raise ValueError( + f"Reported population support lacks its {key} reference" + ) + if ArtifactRef.from_dict(reference) != artifact_ref(final.get(key)): + raise ValueError( + "Reported population support differs from the final analysis" + ) + if population.get("candidateId") != accepted.get("selectedCandidateId"): + raise ValueError("Reported population support belongs to another candidate") + return { + "request": mapping(snapshot.get("request")), + "modelUsage": mapping(snapshot.get("modelUsage")), + "finalAnalysis": final, + "decisions": decisions, + "assessments": assessments, + "accepted": accepted, + "context": context, + "study": study, + "populationSupport": population, + "limitations": list(dict.fromkeys(displayed_limits)), + "selectedParameters": mapping(selected.get("parameters")), + "selectedMetrics": mapping(selected.get("metrics")), + "selectedSetting": mapping( + mapping(accepted.get("settings")).get( + str(accepted.get("selectedCandidateId")) + ) + ), + "selectedFeatures": mapping( + mapping(accepted.get("featureEvidence")).get( + str(accepted.get("selectedCandidateId")) + ) + ), + "qc": qc, + "qcProfiles": qc_profiles, + } diff --git a/scarf/agent/report/contracts.py b/scarf/agent/report/contracts.py new file mode 100644 index 00000000..1b920ee2 --- /dev/null +++ b/scarf/agent/report/contracts.py @@ -0,0 +1,56 @@ +"""Small value adapters for the saved analysis summary.""" + +import re +from collections.abc import Mapping, Sequence +from typing import Any + + +def mapping(value: Any) -> dict[str, Any]: + return dict(value) if isinstance(value, Mapping) else {} + + +def mappings(value: Any) -> list[dict[str, Any]]: + if not isinstance(value, Sequence) or isinstance(value, str | bytes): + return [] + return [dict(item) for item in value if isinstance(item, Mapping)] + + +def texts(value: Any) -> list[str]: + if not isinstance(value, Sequence) or isinstance(value, str | bytes): + return [] + return list( + dict.fromkeys( + item.strip() for item in value if isinstance(item, str) and item.strip() + ) + ) + + +def label(value: str) -> str: + names = { + "donor_id": "Donor", + "sample_id": "Sample", + "library_id": "Capture", + "T2D": "T2D status", + "hvgCount": "Number of variable genes", + "hvgRanking": "Variable-gene ranking", + "observationAndIndependentUnitsMustBeDesignOrTechnical": "The saved analysis did not support the declared study-unit roles", + "withinIndependentUnitComparisonsAreUnsupported": "Association testing for repeated observations from one study unit is unavailable", + "continuousConditioningIsUnsupported": "Conditioning on a continuous variable is unavailable", + "fewerThanFourIndependentUnits": "Fewer than four independent study units", + "unsupportedConditioningStrata": "Some conditioned groups do not support the requested comparison", + } + if value in names: + return names[value] + return re.sub(r"(?<=[a-z])(?=[A-Z])", " ", value.replace("_", " ")).capitalize() + + +def scalar(value: Any) -> str: + if value is None: + return "Unavailable" + if isinstance(value, bool): + return "Yes" if value else "No" + if isinstance(value, int): + return f"{value:,}" + if isinstance(value, float): + return f"{value:.3g}" + return str(value) diff --git a/scarf/agent/report/generator.py b/scarf/agent/report/generator.py new file mode 100644 index 00000000..88cb5542 --- /dev/null +++ b/scarf/agent/report/generator.py @@ -0,0 +1,63 @@ +"""Generate one local analysis summary from the authoritative stage history.""" + +import os +import uuid +from collections.abc import Mapping +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from .artifacts import _local_root, report_directory, scientific_summary +from .plots import collect_analysis_artifacts +from .rendering import render_analysis_document + +if TYPE_CHECKING: + from ...datastore.datastore import DataStore + + +def render_analysis_report( + store: "DataStore", snapshot: Mapping[str, Any], output_dir: Path +) -> Path: + """Replace derived display files without changing the saved analysis.""" + if snapshot.get("status") != "completed" or not snapshot.get("finalAnalysis"): + raise ValueError("Reports require a completed analysis with final artifacts") + payload = scientific_summary(snapshot) + output_dir.mkdir(parents=True, exist_ok=True) + payload.update( + collect_analysis_artifacts(store, payload["finalAnalysis"], output_dir) + ) + destination = output_dir / "index.html" + temporary = output_dir / f".index.{uuid.uuid4().hex}.tmp" + try: + temporary.write_text(render_analysis_document(payload), encoding="utf-8") + os.replace(temporary, destination) + finally: + temporary.unlink(missing_ok=True) + return destination + + +def generate_agent_report( + target: "str | Path | DataStore", + workflow_run_id: str, + *, + workspace: str | None = None, +) -> Path: + """Return one concise local report for a completed analysis. + + Decisions, their evidence, and numerical artifacts come from the saved + stage history. Regeneration makes no model or scientific computation calls. + """ + from ...datastore.datastore import DataStore + from ..orchestrator import journal + + root = _local_root(target) + if isinstance(target, DataStore): + if workspace is not None and workspace != target.workspace: + raise ValueError("workspace does not match the DataStore workspace") + store = target + workspace = store.workspace + else: + store = journal.open_analysis_store(root, workflow_run_id, workspace=workspace) + snapshot = journal.analysis_snapshot(store, workflow_run_id) + return render_analysis_report( + store, snapshot, report_directory(root, workflow_run_id, workspace) + ) diff --git a/scarf/agent/report/plots.py b/scarf/agent/report/plots.py new file mode 100644 index 00000000..c2d15cfa --- /dev/null +++ b/scarf/agent/report/plots.py @@ -0,0 +1,105 @@ +"""One saved map and a compact marker table for the analysis report.""" + +import os +import uuid +from collections.abc import Mapping +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from .._plots import cluster_counts, plot_final_umap +from .artifacts import artifact_ref + +if TYPE_CHECKING: + from ...datastore.datastore import DataStore + from ...plotting import PlotResult + + +def _save_plot(plot: "PlotResult", path: Path) -> None: + """Atomically replace the derived image and its provenance sidecar.""" + token = uuid.uuid4().hex + temporary = path.with_name(f".{path.stem}.{token}{path.suffix}") + sidecar = path.with_suffix(path.suffix + ".json") + temporary_sidecar = sidecar.with_name(f".{sidecar.stem}.{token}{sidecar.suffix}") + try: + plot.save(temporary, dpi=150) + plot.save_provenance(temporary_sidecar, figure_path=path, dpi=150) + os.replace(temporary, path) + os.replace(temporary_sidecar, sidecar) + finally: + temporary.unlink(missing_ok=True) + temporary_sidecar.unlink(missing_ok=True) + plot.close() + + +def collect_analysis_artifacts( + store: "DataStore", final: Mapping[str, Any], report_dir: Path +) -> dict[str, Any]: + """Read final numerical artifacts without launching any analysis operation.""" + clusters = artifact_ref(final.get("clusters")) + notes: list[str] = [] + plot_path: str | None = None + displayed: int | None = None + counts: dict[str, int] + try: + plot = plot_final_umap( + store, + umap=artifact_ref(final.get("umap")), + clusters=clusters, + cell_selection=artifact_ref(final.get("cellSelection")), + graph=artifact_ref(final.get("graph")), + show=False, + ) + counts = dict( + zip( + plot.tables["cluster_counts"]["cluster"], + plot.tables["cluster_counts"]["cells"], + strict=True, + ) + ) + displayed = plot.provenance.n_cells + plot_dir = report_dir / "plots" + plot_dir.mkdir(parents=True, exist_ok=True) + _save_plot(plot, plot_dir / "final_umap.png") + plot_path = "plots/final_umap.png" + except (ImportError, OSError, RuntimeError) as exc: + counts = cluster_counts(store, clusters) + notes.append(f"UMAP display unavailable: {type(exc).__name__}: {exc}") + marker_rows: list[dict[str, Any]] = [] + if final.get("markers") is not None: + marker = artifact_ref(final["markers"]) + marker_inputs = store.inspect_artifact(marker).inputs or {} + if marker_inputs.get("clusters") != clusters.to_dict(): + raise ValueError("Final marker statistics must use the selected clusters") + if ( + marker_inputs.get("cell_selection") + != artifact_ref(final.get("cellSelection")).to_dict() + ): + raise ValueError("Final markers must use the selected cells") + if marker.assay != clusters.assay: + raise ValueError("Final markers must use the selected RNA assay") + for cluster in counts: + try: + table = store.get_markers(marker, group_id=cluster) + if "score" in table: + table = table.sort_values("score", ascending=False, kind="stable") + for row in table.head(3).to_dict(orient="records"): + marker_rows.append( + { + "cluster": cluster, + "feature": row.get( + "feature_name", row.get("feature_id", "") + ), + "score": row.get("score"), + } + ) + except Exception as exc: + notes.append( + f"Markers unavailable for cluster {cluster}: {type(exc).__name__}: {exc}" + ) + return { + "clusterCounts": counts, + "markers": marker_rows, + "umap": plot_path, + "displayedCells": displayed, + "displayNotes": notes, + } diff --git a/scarf/agent/report/rendering.py b/scarf/agent/report/rendering.py new file mode 100644 index 00000000..558024f6 --- /dev/null +++ b/scarf/agent/report/rendering.py @@ -0,0 +1,489 @@ +"""A single readable analysis page, using only recorded scientific evidence.""" + +import html +import math +from collections.abc import Mapping, Sequence +from typing import Any + +from .contracts import label, mapping, mappings, scalar, texts + +_STYLES = """ +:root{color-scheme:light;font-family:Inter,sans-serif;font-size:16px;font-weight:300;line-height:1.2;color:#000000;background:#ffffff;letter-spacing:-.04em} +*{box-sizing:border-box}body{margin:0;background:#ffffff}main{max-width:1120px;margin:auto;padding:48px 32px 72px} +header{border-bottom:1px solid #000000;padding-bottom:32px}h1,h2,h3{line-height:1.2;color:#000000;font-weight:400} +h1{font-size:3.5rem;letter-spacing:0;margin:12px 0 20px}h2{font-size:2rem;letter-spacing:-.04em;margin:56px 0 20px}h3{font-size:1.25rem;letter-spacing:-.04em;margin:32px 0 12px} +p{max-width:82ch}header h1+p{font-size:1.5rem;font-weight:300;margin:0 0 24px}a{color:#0077fc} +small{display:block;color:#b4b4b4;font-size:.75rem;font-weight:400;letter-spacing:-.04em;text-transform:uppercase}.muted{color:#b4b4b4} +.header-links{display:flex;flex-wrap:wrap;justify-content:space-between;align-items:baseline;gap:8px 24px;margin-bottom:40px;font-size:.9rem}.header-links a{font-weight:400} +.numbers{display:inline-block;margin:0 0 20px;padding:10px 18px;border-radius:999px;background:#0077fc;color:#ffffff;font-size:1rem;font-weight:400} +.fraction{white-space:nowrap}progress{width:86px;height:8px;border:0;border-radius:999px;overflow:hidden;accent-color:#0077fc;background:#b4b4b4} +progress::-webkit-progress-bar{background:#b4b4b4;border-radius:999px}progress::-webkit-progress-value{background:#0077fc;border-radius:999px}progress::-moz-progress-bar{background:#0077fc;border-radius:999px} +section{margin-top:64px}section>h2:first-child{margin-top:0}.limitations::before{content:"";display:block;width:72px;height:8px;margin-bottom:24px;border-radius:999px;background:#0077fc} +.population-overview{display:grid;grid-template-columns:minmax(0,1fr) minmax(0,1.3fr);gap:20px;align-items:start}.population-overview figure{position:sticky;top:20px}.population-overview>div{min-width:0} +figure{margin:24px 0;background:#ffffff;padding:0}figure img{width:100%;height:auto} +figcaption{color:#b4b4b4;font-size:.9rem;text-align:center}.decision{border-top:1px solid #b4b4b4;padding:14px 0} +.decision h3{margin:0}.decision p{margin:8px 0}details{margin:20px 0}summary{display:inline-block;cursor:pointer;padding:10px 18px;border-radius:999px;box-shadow:inset 0 0 0 2px #0077fc;color:#0077fc;font-weight:400} +details[open] summary{background:#0077fc;color:#ffffff;box-shadow:none}summary:focus-visible{outline:2px solid #000000;outline-offset:3px} +.table-wrap{overflow-x:auto}table{border-collapse:collapse;width:100%;font-size:.92rem;margin:14px 0} +th,td{text-align:left;vertical-align:top;padding:10px 12px;border-bottom:1px solid #b4b4b4} +th,strong{font-weight:400}th{background:#ffffff}td p{margin:0}li{margin:8px 0} +footer{margin-top:64px;border-top:1px solid #000000;padding-top:20px;color:#b4b4b4;font-size:.85rem} +@media(max-width:800px){.population-overview{display:block}.population-overview figure{position:static}} +@media(max-width:600px){main{padding:32px 16px 48px}h1{font-size:2.5rem}h2{font-size:1.75rem}header h1+p{font-size:1.25rem}th,td{padding:8px}} +@media print{body{background:#ffffff}main{padding:0}details{break-inside:avoid}} +""" + + +def _escape(value: Any) -> str: + return html.escape(scalar(value)) + + +def _list(items: Sequence[str]) -> str: + return ( + "
    " + "".join(f"
  • {html.escape(item)}
  • " for item in items) + "
" + if items + else "" + ) + + +def _table(headers: Sequence[str], rows: Sequence[Sequence[Any]]) -> str: + if not rows: + return "" + head = "".join(f"{html.escape(value)}" for value in headers) + body = "".join( + "" + "".join(f"{_escape(value)}" for value in row) + "" + for row in rows + ) + return f'
{head}{body}
' + + +def _fraction_bar(value: Any) -> str: + if ( + not isinstance(value, int | float) + or not math.isfinite(value) + or not 0 <= value <= 1 + ): + return "Unavailable" + return f' {value:.1%}' + + +def _percentage(value: Any) -> str: + return ( + f"{value:.1%}" + if isinstance(value, int | float) and math.isfinite(value) and 0 <= value <= 1 + else "Unavailable" + ) + + +def _scope(assessment: Mapping[str, Any]) -> str: + coverage = mapping(assessment.get("coverage")) + all_cells = ( + assessment.get("scope") == "full" + or mapping(assessment.get("comparisonCoverage")).get("population") == "allCells" + ) + name = "Full cohort" if all_cells else "Screening sample" + size = coverage.get("screeningCells") + return f"{name} ({size:,} cells)" if isinstance(size, int) else name + + +def _population_table(payload: Mapping[str, Any]) -> str: + counts = mapping(payload.get("clusterCounts")) + total = sum(counts.values()) + markers = mappings(payload.get("markers")) + support = mapping(payload.get("populationSupport")) + units = texts(mapping(payload.get("study")).get("independentUnitColumns")) + columns = mapping(support.get("columns")) + unit = next((name for name in units if name in columns), None) + evidence = mapping(columns.get(unit)) if unit is not None else {} + populations = { + str(row["cluster"]): row for row in mappings(evidence.get("populations")) + } + rows = [] + for cluster, count in counts.items(): + row = populations.get(str(cluster), {}) + genes = ", ".join( + str(item["feature"]) + for item in markers + if str(item.get("cluster")) == str(cluster) + ) + observed = row.get("groupsWithAtLeast5Cells") + unit_count = str(observed) if isinstance(observed, int) else "Unavailable" + rows.append( + f"{_escape(cluster)}{count:,}
{_fraction_bar(count / total if total else None)}" + f"{_escape(genes or 'No marker preview available')}" + f"{unit_count}{_fraction_bar(row.get('largestGroupFraction'))}" + ) + unit_label = label(unit) if unit else "Study unit" + omitted = evidence.get("omittedPopulations", 0) + notes = [] + if omitted: + notes.append( + f"Support details were not saved for {omitted} populations; their support is unavailable here." + ) + if evidence.get("missingCells"): + notes.append( + f"{evidence['missingCells']:,} cells lack the recorded study-unit metadata." + ) + if not evidence: + notes.append("No per-population study-unit evidence was saved.") + return ( + '
' + f"" + + "".join(rows) + + "
PopulationCellsTop marker genes{_escape(unit_label)} groups with ≥5 cellsLargest group contribution
" + + '

Markers describe gene programs, not validated cell identities. Study-unit counts and concentration describe observed support; five cells is not a replication threshold.

' + + _list(notes) + ) + + +def _qc_section(payload: Mapping[str, Any]) -> str: + qc = mapping(payload.get("qc")) + profiles = mappings(payload.get("qcProfiles")) + names = { + "coreGlobalGaussian": "Scarf default global filter", + "globalGaussian": "Scarf default global filter", + "coreSampleMad3": "Scarf filter within samples", + "sampleMad": "Scarf filter within samples", + "retainWithFlags": "Retain cells with quality flags", + "skip": "Retain cells without filtering", + "globalMad5": "Lenient global filter", + "captureMad5": "Lenient filter within captures", + "captureMad3Sensitivity": "Stricter filter within captures", + "pooledReferenceMad5": "Filter using reference captures", + } + rows = [] + for profile in profiles: + name = names.get( + str(profile.get("registeredProfile") or profile.get("action")), + "Recorded quality filter", + ) + if profile.get("sampleColumn"): + name += f" ({label(str(profile['sampleColumn']))})" + chosen = profile.get("profileId") == qc.get("profileId") + rows.append( + f"{_escape(name)}{' (selected)' if chosen else ''}" + f"{_escape(profile.get('retainedCells'))}" + f"{_fraction_bar(profile.get('retainedFraction'))}" + ) + table = ( + '
' + + "".join(rows) + + "
Compared policyProjected cells retainedRetention
" + if rows + else "

Quality-policy comparisons are unavailable.

" + ) + grouped = [] + for column, groups in mapping(qc.get("retainedCellsByColumn")).items(): + values = list(mapping(groups).values()) + if not values: + continue + description = ( + "; ".join(f"{name}: {count:,}" for name, count in groups.items()) + if len(values) <= 8 + else f"{len(values)} groups; smallest {min(values):,}, largest {max(values):,} cells" + ) + grouped.append([label(column), description]) + decisions = [ + mapping(item.get("record")) + for item in mappings(payload.get("decisions")) + if mapping(item.get("record")).get("decisionId") + in {"cellQuality", "qcGrouping"} + ] + quality = [item for item in decisions if item.get("decisionId") == "cellQuality"] + chosen_reason = f"

{_escape(quality[-1].get('rationale'))}

" if quality else "" + explanations = "".join( + f"

{_escape(item.get('rationale'))}

" + for item in decisions + if item.get("decisionId") == "qcGrouping" + ) + if explanations: + explanations = f"
Study grouping for quality filtering{explanations}
" + return f'

Cell quality

{chosen_reason}{table}{_table(("Retained study groups", "Cells"), grouped)}{explanations}
' + + +def _design_section(payload: Mapping[str, Any]) -> str: + study = mapping(payload.get("study")) + context = mapping(payload.get("context")) + characterization = mapping(context.get("characterization")) + correction = mapping(payload.get("selectedParameters")).get("useHarmony") + license = study.get("correctionLicense") + if correction is True: + correction_text = "Batch correction was applied to the selected representation." + elif license == "unsafeConfounded": + correction_text = "Batch correction was not applied: the recorded design cannot separate the proposed batch effects from protected biology." + elif correction is False: + correction_text = "The selected representation uses no batch correction." + else: + correction_text = "The correction decision is unavailable." + rows = [] + for coefficient in mappings(characterization.get("coefficients")): + replication = mapping(coefficient.get("replication")) + groups = mappings(replication.get("independentUnitsByGroup")) + counts = "; ".join( + f"{item.get('group')}: {item.get('count')}" for item in groups + ) + paired = mapping(coefficient.get("pairedCoverage")) + pairing = ( + f"{paired['completePairs']} of {paired.get('pairs', 'unavailable')} complete pairs" + if paired.get("design") == "mixedOrIncomplete" + else "Between study units" + if paired.get("betweenIndependentUnits") + else "" + ) + rows.append([label(str(coefficient.get("name", ""))), counts, pairing]) + requirements = { + item["requirementId"]: item + for item in mappings(study.get("evidenceRequirements")) + } + questions = [] + for item in mappings(study.get("evidenceCoverage")): + requirement = requirements.get(item.get("requirementId"), {}) + status = { + "computed": "Assessed", + "nonIdentifiable": "Not identifiable", + "unsupported": "Not assessed", + "failed": "Failed", + }.get(str(item.get("status")), "Unavailable") + questions.append( + [ + requirement.get("question"), + status, + "; ".join(label(reason) for reason in texts(item.get("reasons"))), + ] + ) + claims = texts(study.get("unsupportedClaims")) + return f'

Study design and correction

{_escape(correction_text)}

{_table(("Study factor", "Independent units per group", "Design"), rows)}{_table(("Objective question", "Evidence", "Limit"), questions)}{_list(claims)}
' + + +_AXIS_LABELS = { + "hvgCount": "Number of variable genes", + "hvgRanking": "Variable-gene ranking", + "featurePolicy": "Gene families", + "pca": "PCA dimensions", + "dimensions": "PCA dimensions", + "neighbors": "Neighbors", + "resolution": "Clustering resolution", + "partition": "Clustering resolution", + "batchCorrection": "Batch correction", + "harmony": "Batch correction", +} + + +def _comparison_sections(payload: Mapping[str, Any]) -> str: + assessments = mappings(payload.get("assessments")) + scope_evidence = {str(item["scope"]): item for item in assessments} + sections = [] + for title, partition in (("Genes and representation", False), ("Clustering", True)): + scopes: dict[str, dict[str, Any]] = {} + for assessment in assessments: + coverage = mapping(assessment.get("comparisonCoverage")) + settings = mapping(coverage.get("candidateSettings")) + for conclusion in mappings(assessment.get("comparisonConclusions")): + if (conclusion.get("axis") == "partition") != partition: + continue + identities = texts(conclusion.get("candidateIds")) + for identity in identities: + candidate = mapping(settings.get(identity)) + scope = str(candidate.get("scope", "")) + if not candidate or scope not in {"sample0", "sample1", "full"}: + raise ValueError( + "Reported comparison lacks exact candidate scope and settings" + ) + group = scopes.setdefault( + scope, + { + "assessment": scope_evidence.get(scope, {"scope": scope}), + "conclusions": {}, + "candidates": {}, + "unavailable": {}, + }, + ) + key = (conclusion["axis"], tuple(sorted(identities))) + group["conclusions"][key] = conclusion + group["candidates"][identity] = candidate + for item in mappings(coverage.get("comparisons")): + if ( + item.get("status") != "notApplicable" + or (item.get("axis") == "partition") != partition + ): + continue + candidate = mapping(settings.get(str(item.get("baselineCandidateId")))) + scope = str(candidate.get("scope", "")) + if scope in scopes: + scopes[scope]["unavailable"][ + (item.get("axis"), item.get("reason")) + ] = item + content = [] + for group in scopes.values(): + content.append(f"

{_escape(_scope(group['assessment']))}

") + for conclusion in group["conclusions"].values(): + axis = str(conclusion["axis"]) + content.append( + f"

{_escape(_AXIS_LABELS.get(axis, label(axis)))}. {_escape(conclusion.get('plainLanguageSummary'))}

" + ) + content.append( + f"
Evidence behind this choice

{_escape(conclusion.get('quantitativeReason'))}

{_escape(conclusion.get('biologicalReason'))}

" + ) + content.append( + _list( + texts( + [ + item.get("interpretation") + for item in mappings(conclusion.get("tradeoffs")) + ] + ) + ) + ) + rows = [] + for identity, candidate in group["candidates"].items(): + setting = candidate + parameters = mapping(candidate.get("parameters")) + metrics = mapping(candidate.get("metrics")) + preferred = [ + _AXIS_LABELS.get(item["axis"], label(item["axis"])) + for item in group["conclusions"].values() + if item.get("preferredCandidateId") == identity + ] + choice = ( + "Preferred: " + ", ".join(dict.fromkeys(preferred)) + if preferred + else "Compared" + ) + if identity == mapping(payload.get("accepted")).get( + "selectedCandidateId" + ): + choice = "Selected final settings" + if partition: + rows.append( + [ + parameters.get("leidenResolution"), + metrics.get("nClusters"), + metrics.get("minClusterCells"), + metrics.get("seedStability"), + metrics.get("subsampleStability"), + _percentage(metrics.get("markerCoherence")), + choice, + ] + ) + else: + ranking = {"batchAware": "Within batches", "global": "Global"}.get( + str(setting.get("ranking")), "Unavailable" + ) + rows.append( + [ + setting.get("hvgCount"), + ranking, + parameters.get("dimensions"), + parameters.get("neighborsK"), + metrics.get("nClusters"), + metrics.get("seedStability"), + _percentage(metrics.get("markerCoherence")), + choice, + ] + ) + headers = ( + ( + "Resolution", + "Populations", + "Smallest population", + "Repeat agreement", + "Subsample agreement", + "Clusters with qualifying markers", + "Choice", + ) + if partition + else ( + "Variable genes", + "Ranking", + "PCA dimensions", + "Neighbors", + "Populations", + "Repeat agreement", + "Clusters with qualifying markers", + "Choice", + ) + ) + content.append(_table(headers, rows)) + for item in group["unavailable"].values(): + axis = str(item["axis"]) + content.append( + f"

{_escape(_AXIS_LABELS.get(axis, label(axis)))}: not compared. {_escape(item.get('reason'))}

" + ) + if content: + sections.append(f"

{title}

{''.join(content)}
") + return "".join(sections) + + +def render_analysis_document(payload: Mapping[str, Any]) -> str: + final = mapping(payload.get("finalAnalysis")) + request = mapping(payload.get("request")) + counts = mapping(payload.get("clusterCounts")) + total = sum(int(value) for value in counts.values()) + objective = request.get("studyObjective") or request.get("studyContext") or "" + assay = final.get("primaryAssay") or request.get("primaryAssay") or "RNA" + qc = mapping(payload.get("qc")) + qc_text = "" + if isinstance(qc.get("retainedCells"), int) and isinstance( + qc.get("retainedFraction"), int | float + ): + qc_text = f"

QC retained {_escape(qc['retainedCells'])} cells ({float(qc['retainedFraction']):.1%}).

" + accepted = mapping(payload.get("accepted")) + outcome = str( + accepted.get("plainLanguageSummary") + or "The selected populations and their marker programs are recorded below." + ) + map_markup = "" + if payload.get("umap"): + display = int(payload.get("displayedCells") or total) + map_markup = f'
Final UMAP colored by saved population labels
{display:,} of {total:,} cells shown. Counts and markers use the complete selection.
' + limitations = _list(texts(payload.get("limitations"))) + display_notes = _list(texts(payload.get("displayNotes"))) + mode_note = ( + "

The model assessed structured evidence. No plots were supplied for visual inspection.

" + if accepted.get("evidenceMode") == "structured" + else "" + ) + parameters = mapping(payload.get("selectedParameters")) + setting = mapping(payload.get("selectedSetting")) + methods = _table( + ("Selected setting", "Value"), + [ + [name, value] + for name, value in ( + ("Variable genes", setting.get("hvgCount")), + ("PCA dimensions", parameters.get("dimensions")), + ("Neighbors", parameters.get("neighborsK")), + ("Clustering resolution", parameters.get("leidenResolution")), + ("Batch correction", parameters.get("useHarmony")), + ) + ], + ) + usage = mapping(payload.get("modelUsage")) + usage_note = "" + if usage.get("invocations"): + usage_note = ( + f"

Recorded model work: {int(usage['invocations']):,} invocations, " + f"{int(usage.get('failedInvocations', 0)):,} failed; " + f"{int(usage.get('requests', 0)):,} completed responses and " + f"{int(usage.get('validationRetries', 0)):,} validation corrections. " + f"Reported tokens: {int(usage.get('inputTokens', 0)):,} input and " + f"{int(usage.get('outputTokens', 0)):,} output, including failed invocations.

" + ) + if usage.get("availability") != "reported": + usage_note += ( + "

Provider usage is incomplete or unavailable for some invocations. " + "Reported totals are known usage only; missing usage is not zero. " + "Failed requests without a response are not included in the response count.

" + ) + return f""" + +Scarf analysis summary
+
Scarf analysis

Analysis summary

{_escape(objective)}

+

{total:,} cells · {len(counts):,} clusters · {_escape(assay)}

{qc_text}

{_escape(outcome)}

+

Populations and markers

{map_markup}
{_population_table(payload)}
+{_qc_section(payload)}{_design_section(payload)}{_comparison_sections(payload)} +
Selected methods and evidence{methods}{mode_note}{usage_note}

Repeat and subsample agreement use adjusted Rand index. Marker coverage is the fraction of clusters with qualifying markers. These describe the selected analysis; they are not probabilities of biological correctness.

+{"
Unavailable displays" + display_notes + "
" if display_notes else ""} +{'

Limits of this analysis

' + limitations + "
" if limitations else ""} +
Generated locally by Scarf. All numerical evidence is read from the saved analysis; report generation makes no analysis or model calls.
+
""" diff --git a/scarf/agent/runtime.py b/scarf/agent/runtime.py deleted file mode 100644 index 0f5d8616..00000000 --- a/scarf/agent/runtime.py +++ /dev/null @@ -1,116 +0,0 @@ -"""Runtime checks and local env loading for LLM endpoints.""" - -import json -import os -import urllib.error -import urllib.request -from pathlib import Path -from typing import Any - -_ENV_PATH = Path(__file__).resolve().parent / ".env" - - -def load_env(path: Path | None = None) -> Path | None: - """Load scarf/agent/.env into os.environ without overriding existing vars. - - Returns the path loaded, or None if no file was found. - """ - env_path = path or _ENV_PATH - if not env_path.is_file(): - return None - for line in env_path.read_text(encoding="utf-8").splitlines(): - line = line.strip() - if not line or line.startswith("#") or "=" not in line: - continue - key, value = line.split("=", 1) - key = key.strip() - value = value.strip().strip('"').strip("'") - if key and key not in os.environ: - os.environ[key] = value - return env_path - - -def _unreachable_message(baseUrl: str, model: str) -> str: - return ( - f"OpenAI-compatible endpoint unreachable at {baseUrl!r} " - f"(requested model {model!r})." - ) - - -def _missing_model_message(baseUrl: str, model: str, available: set[str]) -> str: - listed = ", ".join(sorted(available)) if available else "(none)" - return f"Model {model!r} not found at {baseUrl!r}. Available models: {listed}." - - -def _missing_config_message() -> str: - example = _ENV_PATH.with_name(".env.example") - return ( - "baseUrl and model are required. Pass them as arguments or set " - "OLLAMA_BASE_URL and OLLAMA_MODEL in the environment or in " - f"{_ENV_PATH} (see {example})." - ) - - -def _request_json( - url: str, - *, - timeout: float, - api_key: str | None, -) -> Any: - headers = {"Accept": "application/json"} - if api_key: - headers["Authorization"] = f"Bearer {api_key}" - request = urllib.request.Request(url, headers=headers, method="GET") - with urllib.request.urlopen(request, timeout=timeout) as response: - return json.loads(response.read().decode("utf-8")) - - -def _listed_model_ids(payload: Any) -> set[str]: - ids: set[str] = set() - if not isinstance(payload, dict): - return ids - data = payload.get("data") - if isinstance(data, list): - for item in data: - if isinstance(item, dict) and isinstance(item.get("id"), str): - ids.add(item["id"]) - models = payload.get("models") - if isinstance(models, list): - for item in models: - if not isinstance(item, dict): - continue - name = item.get("name") or item.get("model") - if isinstance(name, str): - ids.add(name) - return ids - - -def check_runtime( - *, - baseUrl: str | None = None, - model: str | None = None, - timeout: float = 5.0, -) -> None: - """Fail fast if the OpenAI-compatible endpoint or model is unavailable. - - Loads `scarf/agent/.env` first (without overriding existing env vars). - Uses `OLLAMA_BASE_URL`, `OLLAMA_MODEL`, and `OLLAMA_API_KEY` when args are omitted. - """ - load_env() - resolved_base = baseUrl or os.environ.get("OLLAMA_BASE_URL") - resolved_model = model or os.environ.get("OLLAMA_MODEL") - if not resolved_base or not resolved_model: - raise ValueError(_missing_config_message()) - - api_key = os.environ.get("OLLAMA_API_KEY") or None - models_url = f"{resolved_base.rstrip('/')}/models" - try: - payload = _request_json(models_url, timeout=timeout, api_key=api_key) - except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as exc: - raise RuntimeError(_unreachable_message(resolved_base, resolved_model)) from exc - - available = _listed_model_ids(payload) - if resolved_model not in available: - raise RuntimeError( - _missing_model_message(resolved_base, resolved_model, available) - ) diff --git a/scarf/agent/types.py b/scarf/agent/types.py index 824fbd2c..fa878719 100644 --- a/scarf/agent/types.py +++ b/scarf/agent/types.py @@ -2,7 +2,7 @@ from typing import Any, Literal -from .config._deps import AGENT_INSTALL_HINT +from ._deps import AGENT_INSTALL_HINT try: from pydantic import BaseModel, ConfigDict, Field @@ -10,7 +10,7 @@ raise ImportError(AGENT_INSTALL_HINT) from exc -type StageStatus = Literal["done", "needsInput", "failed"] +type StageStatus = Literal["done", "needsInput", "abstained", "failed"] type BatchCorrectionAction = Literal[ "skip", "evaluateHarmony", @@ -30,11 +30,6 @@ def get_blank(cls) -> "AgentDataModel": """Return an empty but valid value for fallback paths.""" return cls() - @classmethod - def get_example(cls) -> "AgentDataModel": - """Return a small representative value for tests and fixtures.""" - return cls.get_blank() - class ArtifactReferenceModel(AgentDataModel): type: Literal["artifact"] = "artifact" @@ -53,14 +48,6 @@ def from_artifact_ref(cls, ref: Any) -> "ArtifactReferenceModel": artifactId=str(getattr(ref, "artifact_id", "")), ) - @classmethod - def get_example(cls) -> "ArtifactReferenceModel": - return cls( - assay="RNA", - kind="reduction", - artifactId="0" * 64, - ) - class BatchSafetyEvidence(AgentDataModel): """Estimability for one coefficient and exact proposed batch-column set.""" @@ -78,23 +65,6 @@ class BatchSafetyEvidence(AgentDataModel): def get_blank(cls) -> "BatchSafetyEvidence": return cls() - @classmethod - def get_example(cls) -> "BatchSafetyEvidence": - return cls( - coefficient="treatment", - coefficientKind="categorical", - observationUnit="sample", - batchColumns=["batch"], - unitConstantBatchColumns=["batch"], - status="safe", - estimability={ - "status": "ok", - "coefficientEstimable": True, - "rankDeficient": False, - }, - evidenceId="batchEstimability:treatment:batch", - ) - class ExperimentalTuningHandoff(AgentDataModel): """Validated Experimental Context inputs for Parameter Tuning.""" @@ -135,37 +105,12 @@ class TuningBiologyHandoff(AgentDataModel): def get_blank(cls) -> "TuningBiologyHandoff": return cls() - @classmethod - def get_example(cls) -> "TuningBiologyHandoff": - return cls( - cellSelection=ArtifactReferenceModel( - scope="datastore", - assay=None, - kind="cell_selection", - artifactId="c" * 64, - ), - fromAssay="RNA", - graphAssay="RNA", - markerAssay="RNA", - recommendedCandidateId="baseline", - clusterArtifact=ArtifactReferenceModel( - assay="RNA", - kind="cluster_labels", - artifactId="1" * 64, - ), - evidenceIds=["candidate:baseline:clusters"], - ) - class ToolCallInfo(AgentDataModel): toolName: str = "" callId: str = "" arguments: dict[str, Any] = Field(default_factory=dict) - @classmethod - def get_example(cls) -> "ToolCallInfo": - return cls(toolName="inspect_store", callId="tool-call-1") - class AgentUsageInfo(AgentDataModel): inputTokens: int = 0 @@ -173,16 +118,27 @@ class AgentUsageInfo(AgentDataModel): totalTokens: int = 0 requests: int = 0 toolCalls: int = 0 + availability: Literal["reported", "partial", "unavailable"] | None = Field( + default=None, exclude_if=lambda value: value is None + ) - @classmethod - def get_example(cls) -> "AgentUsageInfo": - return cls( - inputTokens=100, - outputTokens=50, - totalTokens=150, - requests=2, - toolCalls=1, - ) + +class AgentValidationRetry(AgentDataModel): + """A rejected response or tool call and the feedback supplied for repair.""" + + source: Literal["output", "tool", "schema"] + requestIndex: int = 0 + message: str + response: dict[str, Any] | str | None = None + + +class AgentProviderFailure(AgentDataModel): + """One failed model request, distinct from rejected scientific output.""" + + requestIndex: int + statusCode: int | None = None + error: str + retryDelaySeconds: float | None = None class AgentRunInfo(AgentDataModel): @@ -192,27 +148,23 @@ class AgentRunInfo(AgentDataModel): durationSeconds: float = 0.0 usage: AgentUsageInfo = Field(default_factory=AgentUsageInfo) toolCalls: list[ToolCallInfo] = Field(default_factory=list) - - @classmethod - def get_example(cls) -> "AgentRunInfo": - return cls( - agentName="data_enrichment", - modelName="example-model", - runId="example-run", - durationSeconds=0.1, - usage=AgentUsageInfo.get_example(), - toolCalls=[ToolCallInfo.get_example()], - ) + status: Literal["done", "failed"] | None = Field( + default=None, exclude_if=lambda value: value is None + ) + validationRetries: list[AgentValidationRetry] = Field( + default_factory=list, exclude_if=lambda value: not value + ) + providerFailures: list[AgentProviderFailure] = Field( + default_factory=list, exclude_if=lambda value: not value + ) + errorType: str | None = Field(default=None, exclude_if=lambda value: value is None) + error: str | None = Field(default=None, exclude_if=lambda value: value is None) class AgentExecutionResult(AgentDataModel): output: Any = None runInfo: AgentRunInfo = Field(default_factory=AgentRunInfo) - @classmethod - def get_example(cls) -> "AgentExecutionResult": - return cls(output={}, runInfo=AgentRunInfo.get_example()) - class EvidenceItem(AgentDataModel): id: str @@ -223,14 +175,6 @@ class EvidenceItem(AgentDataModel): def get_blank(cls) -> "EvidenceItem": return cls(id="", label="", summary="") - @classmethod - def get_example(cls) -> "EvidenceItem": - return cls( - id="evidence:example", - label="example", - summary="A bounded observed fact.", - ) - class Decision(AgentDataModel): selectedId: str = Field( @@ -246,14 +190,6 @@ class Decision(AgentDataModel): def get_blank(cls) -> "Decision": return cls(selectedId="", rationale="") - @classmethod - def get_example(cls) -> "Decision": - return cls( - selectedId="evidence:example", - rationale="The evidence directly answers the question.", - evidenceIds=["evidence:example"], - ) - class NeedsInput(AgentDataModel): question: str @@ -264,13 +200,6 @@ class NeedsInput(AgentDataModel): def get_blank(cls) -> "NeedsInput": return cls(question="") - @classmethod - def get_example(cls) -> "NeedsInput": - return cls( - question="Which condition column should be used?", - options=["condition", "treatment"], - ) - class StageResult(AgentDataModel): status: StageStatus @@ -282,7 +211,3 @@ class StageResult(AgentDataModel): @classmethod def get_blank(cls) -> "StageResult": return cls(status="needsInput") - - @classmethod - def get_example(cls) -> "StageResult": - return cls(status="done", decision=Decision.get_example()) diff --git a/tests/agent_comparison_examples.py b/tests/agent_comparison_examples.py new file mode 100644 index 00000000..f503434c --- /dev/null +++ b/tests/agent_comparison_examples.py @@ -0,0 +1,341 @@ +"""Exact, completed RNA comparison evidence for contract and report tests.""" + +from copy import deepcopy + + +def observed_action(evidence: dict, *, selected: str | None = None) -> dict: + """A deterministic synthetic-study reviewer grounded in its supplied rows.""" + coverage = evidence["comparisonCoverage"] + rows = coverage["comparisons"] + settings = coverage["candidateSettings"] + baseline = rows[0]["baselineCandidateId"] + chosen = selected or evidence["currentCandidateId"] + combining = coverage["phase"] == "sensitivity" + pending = next((row for row in rows if row["status"] == "pending"), None) + experiment = ( + next( + ( + key + for key, value in evidence.get("experiments", {}).items() + if value["parameter"] + in { + "includeFamily", + "excludeFamily", + "includeFeature", + "excludeFeature", + } + ), + None, + ) + if pending + else None + ) + conclusions = [] + for axis in dict.fromkeys(row["axis"] for row in rows): + ids = list( + dict.fromkeys( + identifier + for row in rows + if row["axis"] == axis + for identifier in ( + row["baselineCandidateId"], + row["alternativeCandidateId"], + ) + if identifier + ) + ) + preferred = ( + chosen + if axis == "partition" and chosen in ids and not combining + else baseline + ) + metrics = settings[preferred]["metrics"] + tradeoffs = [] + for identifier in ids: + if identifier == preferred: + continue + for metric in ( + "seedStability", + "subsampleStability", + "markerCoherence", + "markerSpecificityMedian", + "macroF1", + ): + value = settings[identifier]["metrics"].get(metric) + current = metrics.get(metric) + if ( + isinstance(value, (int, float)) + and isinstance(current, (int, float)) + and value > current + ): + tradeoffs.append( + { + "alternativeCandidateId": identifier, + "metric": metric, + "preferredValue": current, + "alternativeValue": value, + "interpretation": "The synthetic reference's marker program and exact population partition remain the objective; this measured improvement alone does not justify replacing that partition.", + } + ) + conclusions.append( + { + "axis": axis, + "candidateIds": ids, + "preferredCandidateId": preferred, + "quantitativeReason": f"Compare the actual stability and marker measurements for {len(ids)} observed settings.", + "biologicalReason": "Preserve the synthetic reference marker programs and population membership.", + "plainLanguageSummary": "The reference population remains represented with its marker program.", + "tradeoffs": tradeoffs, + } + ) + evidence_ids = [f"candidate:{chosen}", *evidence.get("imageHashes", {})] + return { + "action": "experiment" if pending else "combine" if combining else "accept", + "selectedCandidateId": chosen, + "experimentId": experiment, + "combinedSettings": { + field: baseline + for field in ( + "hvgCountCandidateId", + "hvgRankingCandidateId", + "featurePolicyCandidateId", + "pcaCandidateId", + "neighborsCandidateId", + ) + } + if combining and not pending + else None, + "correctionNeed": "notApplicable", + "evidenceIds": evidence_ids, + "quantitativeFindings": [ + "Use the supplied same-cell marker and stability measurements." + ], + "qualitativeFindings": [ + "Interpret the observed marker names against the known synthetic population program." + ], + "comparisonConclusions": conclusions, + "populationConcerns": [ + { + "candidateId": chosen, + "clusterId": cluster, + "status": "nonEssentialLimitation", + "evidenceIds": evidence_ids, + "explanation": "No cell identity is assigned to this marker-poor synthetic background group; the reference marker-defined group remains supported.", + } + for cluster, genes in settings[chosen]["metrics"] + .get("topMarkerGenes", {}) + .items() + if not genes + ], + "plainLanguageSummary": "The observed reference population and its marker genes are retained.", + "concern": "Test the nominated family contribution to the selected population program." + if pending + else "", + "expectedImprovement": "Determine whether this exact gene policy changes marker-supported separation." + if pending + else "", + "objectivePreservation": "Retain the known synthetic population and marker program.", + "rationale": "Use the measured comparisons and preserve the explicitly defined synthetic reference population.", + } + + +def comparison_review(scope: str = "full") -> dict: + cells = { + "scope": "datastore", + "assay": None, + "kind": "cell_selection", + "artifactId": "c" * 64, + "type": "artifact", + } + features = { + "scope": "assay", + "assay": "RNA2", + "kind": "feature_selection", + "artifactId": "4" * 64, + "type": "artifact", + } + baseline = { + "scope": scope, + "status": "done", + "nCells": 621200, + "cellSelection": cells, + "features": features, + "eligibleFeatures": {**features, "artifactId": "6" * 64}, + "hvgCount": 1000, + "ranking": "global", + "rankingColumn": None, + "parameters": { + "dimensions": 20, + "neighborsK": 15, + "leidenResolution": 1.0, + "useHarmony": False, + }, + "metrics": { + "nClusters": 2, + "minClusterCells": 1200, + "seedStability": 0.92, + "subsampleStability": 0.88, + "markerCoherence": 0.84, + "topMarkerGenes": {"0": ["MS4A1"], "1": ["CD3D"]}, + }, + } + settings = {"baseline": baseline} + rows = [] + for comparison_id, axis, identity, field, value in ( + ( + "defaultResolution:0.5", + "partition", + "resolution-half", + "leidenResolution", + 0.5, + ), + ( + "defaultResolution:0.75", + "partition", + "candidate-two", + "leidenResolution", + 0.75, + ), + ( + "defaultResolution:1.25", + "partition", + "resolution-high", + "leidenResolution", + 1.25, + ), + ("hvgCount:2000", "hvgCount", "genes-two", "hvgCount", 2000), + ("hvgCount:4000", "hvgCount", "genes-four", "hvgCount", 4000), + ("dimensions:10", "pca", "dimensions-ten", "dimensions", 10), + ("dimensions:30", "pca", "dimensions-thirty", "dimensions", 30), + ("neighborsK:21", "neighbors", "neighbors-twenty-one", "neighborsK", 21), + ("neighborsK:41", "neighbors", "neighbors-forty-one", "neighborsK", 41), + ): + setting = deepcopy(baseline) + target = setting if field == "hvgCount" else setting["parameters"] + target[field] = value + if field == "hvgCount": + setting["features"] = { + **features, + "artifactId": ("8" if value == 2000 else "9") * 64, + } + settings[identity] = setting + rows.append( + { + "comparisonId": comparison_id, + "axis": axis, + "status": "completed", + "baselineCandidateId": "baseline", + "alternativeCandidateId": identity, + "reason": "Observed sensitivity comparison", + } + ) + for axis, reason in ( + ("hvgRanking", "No technical grouping supports a batch-specific ranking."), + ( + "featurePolicy", + "No observed unprotected gene family supports an exclusion comparison.", + ), + ): + rows.append( + { + "comparisonId": axis, + "axis": axis, + "status": "notApplicable", + "baselineCandidateId": "baseline", + "alternativeCandidateId": None, + "reason": reason, + "observedProof": { + "baselineFeatures": features, + **( + { + "kind": "insufficientTechnicalGroups", + "eligibleGroupsByColumn": {}, + } + if axis == "hvgRanking" + else { + "kind": "noPermittedPolicy", + "meaningfulPermittedInterventions": 0, + "registeredFamilies": [], + } + ), + }, + } + ) + conclusions = [] + for axis in dict.fromkeys(row["axis"] for row in rows): + identities = [ + "baseline", + *[ + row["alternativeCandidateId"] + for row in rows + if row["axis"] == axis and row["alternativeCandidateId"] + ], + ] + conclusions.append( + { + "axis": axis, + "candidateIds": identities, + "preferredCandidateId": "candidate-two" + if axis == "partition" + else "baseline", + "quantitativeReason": "Repeat agreement was 0.92 and 84% of clusters had qualifying markers.", + "biologicalReason": "The recorded marker programs remain available for interpretation.", + "plainLanguageSummary": "Resolution 0.75 retains a small population with clear markers." + if axis == "partition" + else f"The observed {axis} alternatives did not justify changing this setting.", + } + ) + return { + "scope": scope, + "action": "accept", + "selectedCandidateId": "candidate-two", + "experimentId": None, + "correctionNeed": "notApplicable", + "evidenceIds": ["candidate:candidate-two:clusters"], + "quantitativeFindings": ["Repeat agreement was 0.92."], + "qualitativeFindings": ["MS4A1 and CD79A support the same marker program."], + "objectivePreservation": "Retain the small marker-supported population.", + "rationale": "Original detailed model reasoning retained in the journal.", + "plainLanguageSummary": "The selected settings retain a small population with clear marker genes.", + "comparisonConclusions": conclusions, + "evidenceMode": "structured", + "visualInspection": "unavailable", + "coverage": { + "populationCells": 621200, + "screeningCells": 621200 if scope == "full" else 50000, + }, + "comparisonCoverage": { + "phase": "validation" if scope == "full" else "combined", + "population": "allCells" if scope == "full" else "subset", + "comparisons": rows, + "candidateSettings": settings, + "combinedCandidateId": "baseline", + "resolutionCandidateIds": [ + "resolution-half", + "candidate-two", + "baseline", + "resolution-high", + ], + }, + "settings": { + identity: deepcopy(setting) for identity, setting in settings.items() + }, + "featureEvidence": { + identity: { + "selectedGenes": setting["hvgCount"], + "eligibleGenes": 12000, + "families": {"hla": {"eligibleGenes": 12, "selectedGenes": 6}}, + } + for identity, setting in settings.items() + }, + "candidates": [ + { + "candidateId": identity, + "parameters": setting["parameters"], + "metrics": setting["metrics"], + "cellSelection": cells, + "artifacts": {"graphFeatures": setting["features"]}, + } + for identity, setting in settings.items() + ], + } diff --git a/tests/agent_examples.py b/tests/agent_examples.py new file mode 100644 index 00000000..1fbf4e4b --- /dev/null +++ b/tests/agent_examples.py @@ -0,0 +1,1281 @@ +"""Representative agent values belong to test fixtures, not production APIs.""" + + +def example(model): + """Construct a fixture through its closest registered model factory.""" + for base in model.__mro__: + factory = _FACTORIES.get(f"{base.__module__}.{base.__name__}") + if factory is not None: + return factory(model) + return model.get_blank() + + +def _example_0_BiologicalContext(cls): + return cls( + organism="Homo sapiens", + studyContext="Human lung samples were profiled after drug or vehicle treatment.", + tissue="lung", + cellTypeReferences=["alveolar macrophage", "T cell"], + experimentalDetails=["drug and vehicle groups"], + treatmentQuestion="Which populations respond selectively to treatment?", + ) + + +def _example_1_ConditionClusterSummary(cls): + return cls( + condition="treated", + clusterId="3", + nSamples=4, + meanFraction=0.18, + minFraction=0.12, + maxFraction=0.25, + cellCount=180, + evidenceId="composition:RNA_cluster:condition:treated:cluster:3", + ) + + +def _example_2_ClusterCompositionEvidence(cls): + from scarf.agent.biological_interpretation.contracts import ( + ArtifactReferenceModel, + ConditionClusterSummary, + ) + + summary = example(ConditionClusterSummary) + reference_summary = ConditionClusterSummary( + condition="control", + clusterId=summary.clusterId, + nSamples=4, + meanFraction=0.11, + minFraction=0.08, + maxFraction=0.15, + cellCount=110, + evidenceId="composition:RNA_cluster:condition:control:cluster:3", + ) + return cls( + clusterArtifact=ArtifactReferenceModel( + assay="RNA", kind="cluster_labels", artifactId="b" * 64 + ), + cellSelection=ArtifactReferenceModel( + scope="datastore", assay=None, kind="cell_selection", artifactId="c" * 64 + ), + totalCells=1000, + clusterCounts={"0": 520, "1": 300, "3": 180}, + sampleColumn="sample", + conditionColumn="treatment", + conditionSummaries=[reference_summary, summary], + evidenceIds=[ + "composition:RNA_cluster:counts", + reference_summary.evidenceId, + summary.evidenceId, + ], + ) + + +def _example_3_MarkerFeature(cls): + return cls( + featureId="ENSG00000173372", + featureName="C1QA", + featureIndex=123, + score=0.83, + foldChange=3.4, + fractionExpressed=0.76, + fractionExpressedRest=0.18, + auc=0.91, + adjustedPvalue=0.001, + ) + + +def _example_4_ClusterMarkerEvidence(cls): + from scarf.agent.biological_interpretation.contracts import ( + ArtifactReferenceModel, + MarkerFeature, + ) + + return cls( + clusterId="3", + markers=[example(MarkerFeature)], + markerArtifact=ArtifactReferenceModel( + assay="RNA", kind="marker_table", artifactId="a" * 64 + ), + evidenceId="markers:RNA_cluster:cluster:3", + ) + + +def _example_5_ClusterMarkerBatchEvidence(cls): + from scarf.agent.biological_interpretation.contracts import ClusterMarkerEvidence + + cluster = example(ClusterMarkerEvidence) + return cls(clusters=[cluster], evidenceIds=[cluster.evidenceId]) + + +def _example_6_ClusterInterpretation(cls): + return cls( + clusterId="3", + proposedIdentity="alveolar macrophage-like", + identityIsHypothesis=True, + confidence="medium", + rationale="Observed marker pattern is consistent with the proposed identity.", + evidenceIds=["markers:RNA_cluster:cluster:3"], + ) + + +def _example_7_TreatmentObservation(cls): + return cls( + clusterId="3", + referenceCondition="control", + comparisonCondition="treated", + direction="higher", + observation="Cluster 3 has a higher mean fraction in treated samples.", + evidenceIds=[ + "composition:RNA_cluster:condition:control:cluster:3", + "composition:RNA_cluster:condition:treated:cluster:3", + ], + ) + + +def _example_8_FollowUpRecommendation(cls): + return cls( + question="Is the abundance difference reproducible across donors?", + operation="sample-level differential abundance", + rationale="Current evidence is descriptive and requires independent replicates.", + requiredInputs=["sample", "condition", "donor"], + evidenceIds=[ + "composition:RNA_cluster:condition:control:cluster:3", + "composition:RNA_cluster:condition:treated:cluster:3", + ], + ) + + +def _example_9_BiologicalInterpretationNeedsInput(cls): + return cls( + question="Provide an exact marker artifact or authorize marker search.", + requiredInputs=["markerArtifact"], + ) + + +def _example_10_BiologicalInterpretationReport(cls): + from scarf.agent.biological_interpretation.contracts import ( + ClusterCompositionEvidence, + ClusterInterpretation, + ClusterMarkerEvidence, + FollowUpRecommendation, + TreatmentObservation, + ) + + interpretation = example(ClusterInterpretation) + observation = example(TreatmentObservation) + follow_up = example(FollowUpRecommendation) + return cls( + status="done", + clusterInterpretations=[interpretation], + treatmentObservations=[observation], + followUps=[follow_up], + clusterArtifact=example(ClusterCompositionEvidence).clusterArtifact, + markerArtifact=example(ClusterMarkerEvidence).markerArtifact, + graphAssay="RNA", + markerAssay="RNA", + evidenceIds=sorted( + { + *interpretation.evidenceIds, + *observation.evidenceIds, + *follow_up.evidenceIds, + } + ), + limitations=[ + "Cell identities remain hypotheses until independently validated." + ], + stopReason="The requested clusters were reviewed.", + ) + + +def _example_11_BiologicalInterpretationDependencies(cls): + return cls( + cluster=object(), + fromAssay="RNA", + graphAssay="RNA", + markerAssay="RNA", + markerAssayType="RNA", + sampleColumn="sample", + conditionColumn="treatment", + ) + + +def _example_12_AgentRunConfig(cls): + return cls(requestLimit=9, toolCallLimit=5, outputTokenLimit=2048) + + +def _example_13_FeatureCharacterization(cls): + return cls( + status="done", + notes=["Feature identity and families were characterized."], + assays=[{"assay": "RNA", "species": "homo_sapiens"}], + ) + + +def _example_14_DataEnrichmentContext(cls): + return cls( + studyContext="Single-cell profiling of treated lung tissue", + studyObjective="Discover stable populations while preserving treatment effects.", + organismHint="human", + tissueReferences=["lung"], + cellTypeReferences=["alveolar macrophage", "T cell"], + experimentalDetails=["CRISPR perturbation", "10x 3 prime RNA-seq"], + ) + + +def _example_15_StudyContextSummary(cls): + return cls( + studyContext="Single-cell profiling of treated human lung tests whether treatment changes alveolar macrophage states.", + studyObjective="Discover populations while preserving the treatment comparison.", + organismReferences=["human"], + tissueReferences=["lung"], + cellTypeReferences=["alveolar macrophage"], + experimentalReferences=["treated"], + hypothesisReferences=["treatment changes alveolar macrophage states"], + analysisIntentReferences=["Single-cell profiling"], + evidenceIds=["context:study"], + ) + + +def _example_16_AdtControlEvidence(cls): + return cls( + featureId="Mouse-IgG1-Control", + featureName="Mouse IgG1 isotype control", + matchedToken="isotype", + evidenceId="assay:ADT:adtControl:Mouse-IgG1-Control", + ) + + +def _example_17_HtoTagEvidence(cls): + return cls( + featureId="HTO-1", + featureName="Sample tag 1", + evidenceId="assay:HTO:htoTag:HTO-1", + ) + + +def _example_18_AtacCoordinateEvidence(cls): + return cls( + status="valid", + totalFeatures=2, + validFeatures=2, + validExamples=["chr1:100-200", "chr2:300-450"], + evidenceId="assay:ATAC:atacCoordinates", + ) + + +def _example_19_AssayModalityEvidence(cls): + from scarf.agent.data_enrichment.contracts import AdtControlEvidence + + control = example(AdtControlEvidence) + return cls( + assayType="ADT", + modality="ADT", + typeSource="persisted", + graphEligible=True, + markerEligible=True, + adtControls=[control], + totalObservedFeatures=20, + reportedFeatures=1, + evidenceIds=["assay:ADT:modality", control.evidenceId], + ) + + +def _example_20_FeatureFamilyEvidence(cls): + return cls( + family="mitochondrial", + species="homo_sapiens", + method="chromosome", + count=2, + examples=["MT-CO1", "MT-CYB"], + defaultExclude=True, + evidenceId="assay:RNA:family:mitochondrial", + ) + + +def _example_21_DefaultHvgFamilyEvidence(cls): + return cls( + family="mitochondrial", + pattern="^MT-", + count=2, + examples=["MT-CO1", "MT-CYB"], + evidenceId="assay:RNA:scarfDefaultHvg:family:mitochondrial", + ) + + +def _example_22_RnaFeatureInventoryEvidence(cls): + from scarf.features.variability import DEFAULT_HVG_BLACKLIST + from scarf.agent.data_enrichment.contracts import DefaultHvgFamilyEvidence + + family = example(DefaultHvgFamilyEvidence) + evidence_id = "assay:RNA:scarfDefaultHvg:combined" + return cls( + totalFeatures=20000, + blacklist=DEFAULT_HVG_BLACKLIST, + matchCount=2, + examples=["MT-CO1", "MT-CYB"], + families=[family], + evidenceId=evidence_id, + evidenceIds=[evidence_id, family.evidenceId], + ) + + +def _example_23_ExogenousFeatureEvidence(cls): + return cls( + featureId="ERCC-00002", + featureName="ERCC-00002", + score=4, + classification="potentialExogenous", + evidenceId="assay:RNA:exogenous:ERCC-00002", + ) + + +def _example_24_AssayFeatureInspection(cls): + from scarf.agent.data_enrichment.contracts import ( + AssayModalityEvidence, + FeatureFamilyEvidence, + RnaFeatureInventoryEvidence, + ) + + family = example(FeatureFamilyEvidence) + default_inventory = example(RnaFeatureInventoryEvidence) + modality = AssayModalityEvidence( + assayType="RNA", + modality="RNA", + typeSource="persisted", + graphEligible=True, + markerEligible=True, + totalObservedFeatures=20000, + evidenceIds=["assay:RNA:modality"], + ) + return cls( + assay="RNA", + assayKind="RNAassay", + identity={"nFeatures": 20000, "nDuplicateIds": 0}, + species="homo_sapiens", + speciesMethod="ensemblPrefix", + speciesReason="Most feature IDs carry the ENSG prefix", + families=[family], + defaultFeatureInventory=default_inventory, + modalityEvidence=modality, + evidenceIds=[ + "assay:RNA:identity", + "assay:RNA:species", + family.evidenceId, + *default_inventory.evidenceIds, + *modality.evidenceIds, + ], + ) + + +def _example_25_AssayFeatureInspectionBatch(cls): + from scarf.agent.data_enrichment.contracts import AssayFeatureInspection + + inspection = example(AssayFeatureInspection) + return cls(inspections=[inspection], evidenceIds=list(inspection.evidenceIds)) + + +def _example_26_FeatureReference(cls): + return cls(featureId="ENSG00000198727", featureName="MT-CYB") + + +def _example_27_FeatureMatch(cls): + from scarf.agent.data_enrichment.contracts import FeatureReference + + return cls( + query="MT-CYB", + status="present", + matches=[example(FeatureReference)], + evidenceIds=["assay:RNA:feature:ENSG00000198727"], + ) + + +def _example_28_FeatureLookupResult(cls): + from scarf.agent.data_enrichment.contracts import FeatureMatch + + match = example(FeatureMatch) + return cls(assay="RNA", results=[match], evidenceIds=list(match.evidenceIds)) + + +def _example_29_FeatureLookupBatch(cls): + from scarf.agent.data_enrichment.contracts import FeatureLookupResult + + lookup = example(FeatureLookupResult) + return cls(lookups=[lookup], evidenceIds=list(lookup.evidenceIds)) + + +def _example_30_FeatureSelectionPolicy(cls): + return cls( + assay="RNA", + species="homo_sapiens", + organismName="human", + speciesConfidence="high", + speciesRationale="Gene IDs and study context agree", + excludeFamilies=["mitochondrial", "ribosomal"], + protectFamilies=["cellCycle", "sex"], + artificialFeatures=["ERCC-00002"], + tissueReferences=["lung"], + cellTypeReferences=["alveolar macrophage"], + experimentalReferences=["ERCC spike-in"], + assayType="RNA", + assayModality="RNA", + graphEligible=True, + markerEligible=True, + rationale="Use technical families for feature-selection exclusions", + evidenceIds=["assay:RNA:species", "assay:RNA:family:mitochondrial"], + ) + + +def _example_31_DataEnrichmentToolCall(cls): + return cls( + name="inspect_assay_features", + assay="RNA", + evidenceIds=["assay:RNA:identity", "assay:RNA:species"], + ) + + +def _example_32_DataEnrichmentReport(cls): + from scarf.agent.data_enrichment.contracts import ( + AgentRunInfo, + AssayFeatureInspection, + DataEnrichmentToolCall, + FeatureSelectionPolicy, + StudyContextSummary, + ) + + policy = example(FeatureSelectionPolicy) + inspection = example(AssayFeatureInspection) + return cls( + status="done", + policies=[policy], + inspections=[inspection], + studyContextSummary=example(StudyContextSummary), + evidenceIds=list(policy.evidenceIds), + toolCalls=[example(DataEnrichmentToolCall)], + runInfo=example(AgentRunInfo), + ) + + +def _example_33_DataEnrichmentDependencies(cls): + from scarf.agent.data_enrichment.contracts import DataEnrichmentContext, Path + + return cls( + context=example(DataEnrichmentContext), + assays=["RNA"], + cacheDir=Path("/tmp/scarf-gene-reference"), + allowDownload=False, + evidenceIds={"context:organism", "context:tissue:0"}, + ) + + +def _example_34_CovariateCharacterization(cls): + from scarf.agent.experimental_context.contracts import ArtifactReferenceModel + + return cls( + status="done", + cellSelection=ArtifactReferenceModel( + scope="datastore", kind="cell_selection", artifactId="c" * 64 + ), + notes=["Cell covariates and confounding were characterized."], + columns=[{"name": "batch", "domain": "technical"}], + ) + + +def _example_35_InferenceUnit(cls): + return cls(observationUnit="sample", independentUnit="donor") + + +def _example_36_BatchCorrectionPlan(cls): + return cls( + action="evaluateHarmony", + batchColumns=["batch"], + preserveColumns=["cell_type", "treatment"], + metricsRequired=["iLISI", "cLISI", "graphConnectivity"], + rationale="Batch is technical and crossed with treatment, so compare an exact Harmony candidate while protecting biological labels.", + evidenceIds=[ + "column:batch", + "estimability:treatment", + "batchEstimability:treatment:batch", + ], + ) + + +def _example_37_NamedArtifactSource(cls): + from scarf.agent.experimental_context.contracts import ArtifactReferenceModel + + return cls( + name="RNA_percentMito", + artifact=ArtifactReferenceModel( + assay="RNA", kind="quality_metric", artifactId="1" * 64 + ), + ) + + +def _example_38_CellQcProfileEvidence(cls): + from scarf.agent.experimental_context.contracts import NamedArtifactSource + + return cls( + profileId="cellQc:RNA:globalMad5", + action="registeredMad", + registeredProfile="globalMad5", + driverAssay="RNA", + driverAssayType="RNA", + attributes=["RNA_nCounts", "RNA_nFeatures"], + artifactMetrics=[example(NamedArtifactSource)], + parameters={"nMads": 5.0}, + activeCells=100, + retainedCells=96, + retainedFraction=0.96, + evidenceId="qcProfile:cellQc:RNA:globalMad5", + ) + + +def _example_39_CellQcPlan(cls): + from scarf.agent.experimental_context.contracts import CellQcProfileEvidence + + evidence = example(CellQcProfileEvidence) + return cls( + action=evidence.action, + registeredProfile=evidence.registeredProfile, + profileId=evidence.profileId, + driverAssay=evidence.driverAssay, + driverAssayType=evidence.driverAssayType, + sampleColumn=evidence.sampleColumn, + sampleArtifact=evidence.sampleArtifact, + attributes=evidence.attributes, + artifactMetrics=evidence.artifactMetrics, + rationale="Use the bounded global profile for the RNA assay.", + evidenceIds=[evidence.evidenceId], + ) + + +def _example_40_ExperimentalContextDecision(cls): + from scarf.agent.experimental_context.contracts import ( + BatchCorrectionPlan, + InferenceUnit, + ) + + return cls( + columnDomains={ + "batch": "technical", + "sample": "design", + "donor": "design", + "treatment": "biological", + }, + coefficientsOfInterest=["treatment"], + unitsOfInference={"treatment": example(InferenceUnit)}, + batchCorrection=example(BatchCorrectionPlan), + rationale="Treatment is the primary between-sample contrast.", + evidenceIds=[ + "column:batch", + "column:donor", + "column:sample", + "column:treatment", + ], + ) + + +def _example_41_RepresentationEvaluation(cls): + from scarf.agent.experimental_context.contracts import ArtifactReferenceModel + + return cls( + available=True, + assay="RNA", + cellSelection=ArtifactReferenceModel( + scope="datastore", kind="cell_selection", artifactId="c" * 64 + ), + neighbors=ArtifactReferenceModel( + assay="RNA", kind="neighbors", artifactId="a" * 64 + ), + connectivityMap=ArtifactReferenceModel( + assay="RNA", kind="connectivity_map", artifactId="b" * 64 + ), + metrics={"iLISI:batch": 0.71, "cLISI:cell_type": 0.94}, + evidenceIds=[ + "metric:iLISI:batch:assay:RNA:neighbors:example-neighbors", + "metric:cLISI:cell_type:assay:RNA:neighbors:example-neighbors", + ], + ) + + +def _example_42_CovariateEvidence(cls): + from scarf.agent.experimental_context.contracts import ( + ArtifactReferenceModel, + CellQcProfileEvidence, + CovariateCharacterization, + NamedArtifactSource, + ) + + return cls( + characterization=CovariateCharacterization( + status="done", notes=["Example deterministic covariate characterization"] + ), + qcProfiles=[example(CellQcProfileEvidence)], + htoIdentityColumns=["sample_id"], + htoIdentityArtifacts=[ + NamedArtifactSource( + name="HTO_htoIdentity", + artifact=ArtifactReferenceModel( + assay="HTO", kind="hto_identity", artifactId="2" * 64 + ), + ) + ], + evidenceIds=[ + "column:batch", + example(CellQcProfileEvidence).evidenceId, + "htoIdentity:sample_id", + f"htoIdentityArtifact:HTO_htoIdentity:{'2' * 64}", + ], + ) + + +def _example_43_ExperimentalContextResult(cls): + from scarf.agent.experimental_context.contracts import ( + AgentRunInfo, + ArtifactReferenceModel, + BatchSafetyEvidence, + CellQcProfileEvidence, + CovariateCharacterization, + ExperimentalContextDecision, + NamedArtifactSource, + RepresentationEvaluation, + ) + + representation = example(RepresentationEvaluation) + return cls( + status="done", + decision=example(ExperimentalContextDecision), + characterization=CovariateCharacterization( + status="done", notes=["Example deterministic design characterization"] + ), + cellSelection=representation.cellSelection, + qcProfiles=[example(CellQcProfileEvidence)], + qualityMetricArtifacts=[example(NamedArtifactSource)], + htoIdentityColumns=["sample_id"], + htoIdentityArtifacts=[ + NamedArtifactSource( + name="HTO_htoIdentity", + artifact=ArtifactReferenceModel( + assay="HTO", kind="hto_identity", artifactId="2" * 64 + ), + ) + ], + batchSafety=[example(BatchSafetyEvidence)], + currentRepresentation=representation, + runInfo=example(AgentRunInfo), + ) + + +def _example_44_ExperimentalContextDependencies(cls): + return cls( + studyContext="Case-control study with samples nested in donors.", + studyObjective="Discover populations while preserving the case-control contrast.", + directions={"columnDomains": {"batch": "technical"}}, + ) + + +def _example_45_StudyContract(cls): + return cls( + studyContext="Treated and control blood samples from multiple donors.", + studyObjective="Discover stable populations while preserving treatment-associated structure.", + processingGoal="conditionPreservingDiscovery", + scientificQuestions=[ + "Discover stable populations while preserving treatment-associated structure." + ], + physicalCaptureColumn="sample", + independentUnitColumns=["donor"], + conditionColumns=["treatment"], + technicalBatchColumns=["batch"], + protectedColumns=["treatment", "donor"], + correctionLicense="safe", + allowedClaims=["Describe reproducible population structure."], + unsupportedClaims=[ + "This workflow does not test differential-expression hypotheses." + ], + evidenceIds=["column:batch", "column:donor", "column:treatment"], + evidenceRequirements=[ + { + "requirementId": "studyDesign", + "question": "Which independent units and design constraints support population discovery?", + "objectiveQuote": "Discover stable populations while preserving treatment-associated structure.", + "kind": "studyDesign", + "columns": ["batch", "donor", "treatment"], + "observationUnit": "sample", + "independentUnit": "donor", + "essential": True, + } + ], + evidenceCoverage=[ + { + "requirementId": "studyDesign", + "status": "computed", + "evidenceIds": ["column:batch", "column:donor", "column:treatment"], + "reasons": [], + } + ], + ) + + +def _example_46_IngestResult(cls): + return cls( + status="done", format="h5ad", zarrPath="dataset.zarr", assayNames=["RNA"] + ) + + +def _example_47_WorkflowQuestion(cls): + return cls( + questionId="approvePlanChecksum", + question="Approve this preprocessing plan?", + planChecksum="0" * 64, + ) + + +def _example_48_WorkflowNeedsInput(cls): + from scarf.agent.orchestrator.models import WorkflowQuestion + + return cls(questions=[example(WorkflowQuestion)]) + + +def _example_49_WorkflowStageLink(cls): + return cls(stage="ingest", attemptId="attempt-1", contentSha256="0" * 64) + + +def _example_50_WorkflowStageAttempt(cls): + return cls( + workflowRunId="workflow-1", + stage="ingest", + attemptId="attempt-1", + status="done", + startedAtNs=1, + completedAtNs=2, + requestSha256="0" * 64, + configSha256="1" * 64, + contentSha256="2" * 64, + ) + + +def _example_51_AssayPreprocessingPlan(cls): + return cls( + assay="RNA", + assayType="RNA", + role="graph", + graphEligible=True, + markerEligible=True, + featureMethod="hvg", + reductionMethod="pca", + featureParameters={"topN": 1000, "minCells": 20}, + ) + + +def _example_52_AutomatedPreprocessingPlan(cls): + from scarf.agent.orchestrator.models import ( + ArtifactReferenceModel, + AssayPreprocessingPlan, + ) + + return cls( + primaryAssay="RNA", + markerAssay="RNA", + cellSelection=ArtifactReferenceModel( + scope="datastore", kind="cell_selection", artifactId="c" * 64 + ), + assays=[example(AssayPreprocessingPlan)], + planChecksum="0" * 64, + ) + + +def _example_53_PreprocessedAssayHandoff(cls): + from scarf.agent.orchestrator.models import ArtifactReferenceModel + + return cls( + assay="RNA", + assayType="RNA", + cellSelection=ArtifactReferenceModel( + scope="datastore", kind="cell_selection", artifactId="c" * 64 + ), + reductionMethod="pca", + graphFeatures=example(ArtifactReferenceModel), + markerFeatures=example(ArtifactReferenceModel), + normalized=ArtifactReferenceModel( + assay="RNA", kind="normalized", artifactId="1" * 64 + ), + nCells=100, + nFeatures=1000, + ) + + +def _example_55_FinalAnalysisHandoff(cls): + from scarf.agent.orchestrator.models import ArtifactReferenceModel + + return cls( + workflowRunId="workflow-1", + primaryAssay="RNA", + markerAssay="RNA", + cellSelection=ArtifactReferenceModel( + scope="datastore", kind="cell_selection", artifactId="c" * 64 + ), + graph=ArtifactReferenceModel( + assay="RNA", kind="connectivity_map", artifactId="2" * 64 + ), + embeddingInitialization=ArtifactReferenceModel( + assay="RNA", kind="embedding_initialization", artifactId="5" * 64 + ), + clusters=ArtifactReferenceModel( + assay="RNA", kind="cluster_labels", artifactId="3" * 64 + ), + ) + + +def _example_56_AutomatedWorkflowConfig(cls): + return cls() + + +def _example_57_AutomatedWorkflowRequest(cls): + return cls( + sourcePath="dataset.h5ad", + zarrPath="dataset.zarr", + studyContext="Single-cell profiling of treated human blood.", + studyObjective="Discover stable populations while preserving treatment structure.", + ) + + +def _example_58_AutomatedWorkflowResumeRequest(cls): + return cls( + zarrPath="dataset.zarr", + workflowRunId="workflow-1", + answers={"approvePlanChecksum": "0" * 64}, + ) + + +def _example_59_OrchestrationResumeRecord(cls): + return cls(workflowRunId="workflow-1", answers={"approvePlanChecksum": "0" * 64}) + + +def _example_60_ArtifactRecord(cls): + return cls(scope="assay", kind="connectivity_map", artifactId="a" * 64, assay="RNA") + + +def _example_61_ParameterCandidate(cls): + return cls( + candidateId="baseline", + reductionMethod="pca", + dimensions=21, + leidenResolution=1.0, + neighborsK=11, + useHarmony=False, + ) + + +def _example_62_ParameterMetrics(cls): + return cls( + nClusters=8, + minClusterCells=42, + minClusterFraction=0.021, + graphSilhouetteMedian=0.41, + pcaSilhouette=0.36, + macroF1=0.82, + weightedF1=0.86, + batchMixing={"batch": 0.73}, + biologicalPreservation={ + "cell_type": {"clisi": 0.88, "graphConnectivity": 0.91} + }, + ) + + +def _example_63_ParameterCandidateEvaluation(cls): + from scarf.agent.parameter_tuning.contracts import ( + ArtifactRecord, + ArtifactReferenceModel, + ParameterCandidate, + ParameterMetrics, + ) + + candidate = example(ParameterCandidate) + return cls( + candidateId=candidate.candidateId, + status="done", + eligible=True, + parameters=candidate, + artifacts={ + "connectivityMap": example(ArtifactRecord), + "clusters": ArtifactRecord( + assay="RNA", kind="cluster_labels", artifactId="b" * 64 + ), + }, + cellSelection=ArtifactReferenceModel( + scope="datastore", assay=None, kind="cell_selection", artifactId="c" * 64 + ), + clusterColumn="RNA_agent_tuning_baseline", + clusterLabel="agent_tuning_baseline", + effectiveDimensions=21, + metrics=example(ParameterMetrics), + evidenceIds=["candidate:baseline:clusters"], + ) + + +def _example_64_IntegrationMetrics(cls): + return cls( + nClusters=8, + minClusterCells=37, + minClusterFraction=0.0185, + adjustedRandByAssay={"RNA": 0.71, "ADT": 0.63}, + normalizedMutualInformationByAssay={"RNA": 0.76, "ADT": 0.69}, + modalityWeightsValid=True, + ) + + +def _example_65_IntegrationCandidateEvaluation(cls): + from scarf.agent.parameter_tuning.contracts import ( + ArtifactRecord, + ArtifactReferenceModel, + IntegrationMetrics, + ) + + return cls( + integrationId="wnn_resolution_1", + method="wnn", + assays=["RNA", "ADT"], + status="done", + eligible=True, + cellSelection=ArtifactReferenceModel( + scope="datastore", assay=None, kind="cell_selection", artifactId="c" * 64 + ), + graphArtifact=ArtifactRecord( + scope="datastore", kind="integrated_graph", artifactId="2" * 64 + ), + clusterArtifact=ArtifactRecord( + scope="datastore", kind="cluster_labels", artifactId="3" * 64 + ), + clusterColumn="agent_wnn_cluster", + metrics=example(IntegrationMetrics), + evidenceIds=["integration:wnn_resolution_1:clusters"], + ) + + +def _example_66_FinalGraphComparison(cls): + return cls( + optionId="native:ADT:baseline", + summary="The RNA-native option better preserves the requested labels.", + evidenceIds=[ + "native:RNA:candidate:baseline:clusters", + "native:ADT:candidate:baseline:clusters", + ], + ) + + +def _example_67_FinalGraphNeedsInput(cls): + return cls( + question="Which biological signal must the final graph preserve?", + options=["cell_type", "condition"], + ) + + +def _example_68_FinalGraphSelection(cls): + from scarf.agent.parameter_tuning.contracts import AgentRunInfo + + return cls( + status="done", + selectedOptionId="native:RNA:baseline", + graphMethod="native", + nativeAssay="RNA", + nativeCandidateId="baseline", + markerAssay="RNA", + confidence="medium", + rationale="The selected native graph has the strongest supported balance.", + evidenceIds=["native:RNA:candidate:baseline:clusters"], + runInfo=example(AgentRunInfo), + ) + + +def _example_69_CandidateComparison(cls): + return cls( + candidateId="pca_15", + summary="The selected baseline retains larger minimum clusters.", + evidenceIds=["candidate:baseline:clusters", "candidate:pca_15:clusters"], + ) + + +def _example_70_ParameterSearchPlan(cls): + from scarf.agent.parameter_tuning.contracts import AgentRunInfo, ParameterCandidate + + return cls( + status="refine", + candidates=[ + ParameterCandidate( + candidateId="refined_pca_18", + dimensions=18, + leidenResolution=1.0, + neighborsK=11, + useHarmony=False, + ) + ], + basedOnCandidateIds=["baseline", "pca_15"], + harmonyBatchColumns=[], + objectives=["Resolve the dimension tradeoff."], + rationale="The initial screen brackets a narrower dimension range.", + evidenceIds=["candidate:baseline:clusters", "candidate:pca_15:clusters"], + stoppingCriteria=["Run the proposed candidate once."], + runInfo=example(AgentRunInfo), + ) + + +def _example_71_ParameterTuningBatchSearchPlan(cls): + from scarf.agent.parameter_tuning.contracts import ParameterSearchPlan + + return cls(assayPlans={"RNA": example(ParameterSearchPlan)}) + + +def _example_72_ParameterTuningNeedsInput(cls): + return cls( + question="Which trusted biological label should be preserved?", + options=["cell_type", "none"], + evidenceIds=["candidate:baseline:batchMixing:batch"], + ) + + +def _example_73_ParameterTuningReport(cls): + from scarf.agent.parameter_tuning.contracts import ( + AgentRunInfo, + FinalGraphSelection, + ParameterCandidateEvaluation, + ) + + evaluation = example(ParameterCandidateEvaluation) + return cls( + status="done", + fromAssay="RNA", + cellSelection=evaluation.cellSelection, + evaluations=[evaluation], + recommendedCandidateId=evaluation.candidateId, + selectedArtifacts=dict(evaluation.artifacts), + confidence="medium", + rationale="The baseline balances separation and cluster size.", + evidenceIds=["candidate:baseline:clusters"], + tradeoffs=["Higher resolutions produced smaller clusters."], + limitations=["No trusted biological preservation label was supplied."], + stopReason="All authorized candidates were evaluated.", + recommendedByAssay={"RNA": evaluation.candidateId}, + totalCandidates=1, + graphAssay="RNA", + markerAssay="RNA", + finalSelection=example(FinalGraphSelection), + runInfo=example(AgentRunInfo), + ) + + +def _example_74_ParameterTuningDependencies(cls): + from scarf.agent.parameter_tuning.contracts import ParameterCandidate + + candidate = example(ParameterCandidate) + return cls( + fromAssay="RNA", + normalizedShape=(1000, 2000), + candidates={candidate.candidateId: candidate}, + batchColumns=("batch",), + preservationColumns=("cell_type",), + ) + + +def _example_75_ParameterTuningAssayInput(cls): + from scarf.agent.parameter_tuning.contracts import ( + ArtifactRecord, + ExperimentalTuningHandoff, + _default_parameter_candidates, + ) + + return cls( + normalized=ArtifactRecord(assay="RNA", kind="normalized", artifactId="4" * 64), + candidates=_default_parameter_candidates(), + experimentalHandoff=ExperimentalTuningHandoff(batchAction="skip"), + ) + + +def _example_76_AgentDataModel(cls): + """Return a small representative value for tests and fixtures.""" + return cls.get_blank() + + +def _example_77_ArtifactReferenceModel(cls): + return cls(assay="RNA", kind="reduction", artifactId="0" * 64) + + +def _example_78_BatchSafetyEvidence(cls): + return cls( + coefficient="treatment", + coefficientKind="categorical", + observationUnit="sample", + batchColumns=["batch"], + unitConstantBatchColumns=["batch"], + status="safe", + estimability={ + "status": "ok", + "coefficientEstimable": True, + "rankDeficient": False, + }, + evidenceId="batchEstimability:treatment:batch", + ) + + +def _example_79_TuningBiologyHandoff(cls): + from scarf.agent.types import ArtifactReferenceModel + + return cls( + cellSelection=ArtifactReferenceModel( + scope="datastore", assay=None, kind="cell_selection", artifactId="c" * 64 + ), + fromAssay="RNA", + graphAssay="RNA", + markerAssay="RNA", + recommendedCandidateId="baseline", + clusterArtifact=ArtifactReferenceModel( + assay="RNA", kind="cluster_labels", artifactId="1" * 64 + ), + evidenceIds=["candidate:baseline:clusters"], + ) + + +def _example_80_ToolCallInfo(cls): + return cls(toolName="inspect_store", callId="tool-call-1") + + +def _example_81_AgentUsageInfo(cls): + return cls( + inputTokens=100, outputTokens=50, totalTokens=150, requests=2, toolCalls=1 + ) + + +def _example_82_AgentRunInfo(cls): + from scarf.agent.types import AgentUsageInfo, ToolCallInfo + + return cls( + agentName="data_enrichment", + modelName="example-model", + runId="example-run", + durationSeconds=0.1, + usage=example(AgentUsageInfo), + toolCalls=[example(ToolCallInfo)], + ) + + +def _example_83_AgentExecutionResult(cls): + from scarf.agent.types import AgentRunInfo + + return cls(output={}, runInfo=example(AgentRunInfo)) + + +def _example_84_EvidenceItem(cls): + return cls( + id="evidence:example", label="example", summary="A bounded observed fact." + ) + + +def _example_85_Decision(cls): + return cls( + selectedId="evidence:example", + rationale="The evidence directly answers the question.", + evidenceIds=["evidence:example"], + ) + + +def _example_86_NeedsInput(cls): + return cls( + question="Which condition column should be used?", + options=["condition", "treatment"], + ) + + +def _example_87_StageResult(cls): + from scarf.agent.types import Decision + + return cls(status="done", decision=example(Decision)) + + +_FACTORIES = { + "scarf.agent.biological_interpretation.contracts.BiologicalContext": _example_0_BiologicalContext, + "scarf.agent.biological_interpretation.contracts.ConditionClusterSummary": _example_1_ConditionClusterSummary, + "scarf.agent.biological_interpretation.contracts.ClusterCompositionEvidence": _example_2_ClusterCompositionEvidence, + "scarf.agent.biological_interpretation.contracts.MarkerFeature": _example_3_MarkerFeature, + "scarf.agent.biological_interpretation.contracts.ClusterMarkerEvidence": _example_4_ClusterMarkerEvidence, + "scarf.agent.biological_interpretation.contracts.ClusterMarkerBatchEvidence": _example_5_ClusterMarkerBatchEvidence, + "scarf.agent.biological_interpretation.contracts.ClusterInterpretation": _example_6_ClusterInterpretation, + "scarf.agent.biological_interpretation.contracts.TreatmentObservation": _example_7_TreatmentObservation, + "scarf.agent.biological_interpretation.contracts.FollowUpRecommendation": _example_8_FollowUpRecommendation, + "scarf.agent.biological_interpretation.contracts.BiologicalInterpretationNeedsInput": _example_9_BiologicalInterpretationNeedsInput, + "scarf.agent.biological_interpretation.contracts.BiologicalInterpretationReport": _example_10_BiologicalInterpretationReport, + "scarf.agent.biological_interpretation.contracts.BiologicalInterpretationDependencies": _example_11_BiologicalInterpretationDependencies, + "scarf.agent.config.AgentRunConfig": _example_12_AgentRunConfig, + "scarf.agent.data_enrichment.characterization.FeatureCharacterization": _example_13_FeatureCharacterization, + "scarf.agent.data_enrichment.contracts.DataEnrichmentContext": _example_14_DataEnrichmentContext, + "scarf.agent.data_enrichment.contracts.StudyContextSummary": _example_15_StudyContextSummary, + "scarf.agent.data_enrichment.contracts.AdtControlEvidence": _example_16_AdtControlEvidence, + "scarf.agent.data_enrichment.contracts.HtoTagEvidence": _example_17_HtoTagEvidence, + "scarf.agent.data_enrichment.contracts.AtacCoordinateEvidence": _example_18_AtacCoordinateEvidence, + "scarf.agent.data_enrichment.contracts.AssayModalityEvidence": _example_19_AssayModalityEvidence, + "scarf.agent.data_enrichment.contracts.FeatureFamilyEvidence": _example_20_FeatureFamilyEvidence, + "scarf.agent.data_enrichment.contracts.DefaultHvgFamilyEvidence": _example_21_DefaultHvgFamilyEvidence, + "scarf.agent.data_enrichment.contracts.RnaFeatureInventoryEvidence": _example_22_RnaFeatureInventoryEvidence, + "scarf.agent.data_enrichment.contracts.ExogenousFeatureEvidence": _example_23_ExogenousFeatureEvidence, + "scarf.agent.data_enrichment.contracts.AssayFeatureInspection": _example_24_AssayFeatureInspection, + "scarf.agent.data_enrichment.contracts.AssayFeatureInspectionBatch": _example_25_AssayFeatureInspectionBatch, + "scarf.agent.data_enrichment.contracts.FeatureReference": _example_26_FeatureReference, + "scarf.agent.data_enrichment.contracts.FeatureMatch": _example_27_FeatureMatch, + "scarf.agent.data_enrichment.contracts.FeatureLookupResult": _example_28_FeatureLookupResult, + "scarf.agent.data_enrichment.contracts.FeatureLookupBatch": _example_29_FeatureLookupBatch, + "scarf.agent.data_enrichment.contracts.FeatureSelectionPolicy": _example_30_FeatureSelectionPolicy, + "scarf.agent.data_enrichment.contracts.DataEnrichmentToolCall": _example_31_DataEnrichmentToolCall, + "scarf.agent.data_enrichment.contracts.DataEnrichmentReport": _example_32_DataEnrichmentReport, + "scarf.agent.data_enrichment.contracts.DataEnrichmentDependencies": _example_33_DataEnrichmentDependencies, + "scarf.agent.experimental_context.contracts.CovariateCharacterization": _example_34_CovariateCharacterization, + "scarf.agent.experimental_context.contracts.InferenceUnit": _example_35_InferenceUnit, + "scarf.agent.experimental_context.contracts.BatchCorrectionPlan": _example_36_BatchCorrectionPlan, + "scarf.agent.experimental_context.contracts.NamedArtifactSource": _example_37_NamedArtifactSource, + "scarf.agent.experimental_context.contracts.CellQcProfileEvidence": _example_38_CellQcProfileEvidence, + "scarf.agent.experimental_context.contracts.CellQcPlan": _example_39_CellQcPlan, + "scarf.agent.experimental_context.contracts.ExperimentalContextDecision": _example_40_ExperimentalContextDecision, + "scarf.agent.experimental_context.contracts.RepresentationEvaluation": _example_41_RepresentationEvaluation, + "scarf.agent.experimental_context.contracts.CovariateEvidence": _example_42_CovariateEvidence, + "scarf.agent.experimental_context.contracts.ExperimentalContextResult": _example_43_ExperimentalContextResult, + "scarf.agent.experimental_context.contracts.ExperimentalContextDependencies": _example_44_ExperimentalContextDependencies, + "scarf.agent.experimental_context.study.StudyContract": _example_45_StudyContract, + "scarf.agent.ingest.result.IngestResult": _example_46_IngestResult, + "scarf.agent.orchestrator.models.WorkflowQuestion": _example_47_WorkflowQuestion, + "scarf.agent.orchestrator.models.WorkflowNeedsInput": _example_48_WorkflowNeedsInput, + "scarf.agent.orchestrator.models.WorkflowStageLink": _example_49_WorkflowStageLink, + "scarf.agent.orchestrator.models.WorkflowStageAttempt": _example_50_WorkflowStageAttempt, + "scarf.agent.orchestrator.models.AssayPreprocessingPlan": _example_51_AssayPreprocessingPlan, + "scarf.agent.orchestrator.models.AutomatedPreprocessingPlan": _example_52_AutomatedPreprocessingPlan, + "scarf.agent.orchestrator.models.PreprocessedAssayHandoff": _example_53_PreprocessedAssayHandoff, + "scarf.agent.orchestrator.models.FinalAnalysisHandoff": _example_55_FinalAnalysisHandoff, + "scarf.agent.orchestrator.models.AutomatedWorkflowConfig": _example_56_AutomatedWorkflowConfig, + "scarf.agent.orchestrator.models.AutomatedWorkflowRequest": _example_57_AutomatedWorkflowRequest, + "scarf.agent.orchestrator.models.AutomatedWorkflowResumeRequest": _example_58_AutomatedWorkflowResumeRequest, + "scarf.agent.orchestrator.models.OrchestrationResumeRecord": _example_59_OrchestrationResumeRecord, + "scarf.agent.parameter_tuning.contracts.ArtifactRecord": _example_60_ArtifactRecord, + "scarf.agent.parameter_tuning.contracts.ParameterCandidate": _example_61_ParameterCandidate, + "scarf.agent.parameter_tuning.contracts.ParameterMetrics": _example_62_ParameterMetrics, + "scarf.agent.parameter_tuning.contracts.ParameterCandidateEvaluation": _example_63_ParameterCandidateEvaluation, + "scarf.agent.parameter_tuning.contracts.IntegrationMetrics": _example_64_IntegrationMetrics, + "scarf.agent.parameter_tuning.contracts.IntegrationCandidateEvaluation": _example_65_IntegrationCandidateEvaluation, + "scarf.agent.parameter_tuning.contracts.FinalGraphComparison": _example_66_FinalGraphComparison, + "scarf.agent.parameter_tuning.contracts.FinalGraphNeedsInput": _example_67_FinalGraphNeedsInput, + "scarf.agent.parameter_tuning.contracts.FinalGraphSelection": _example_68_FinalGraphSelection, + "scarf.agent.parameter_tuning.contracts.CandidateComparison": _example_69_CandidateComparison, + "scarf.agent.parameter_tuning.contracts.ParameterSearchPlan": _example_70_ParameterSearchPlan, + "scarf.agent.parameter_tuning.contracts.ParameterTuningBatchSearchPlan": _example_71_ParameterTuningBatchSearchPlan, + "scarf.agent.parameter_tuning.contracts.ParameterTuningNeedsInput": _example_72_ParameterTuningNeedsInput, + "scarf.agent.parameter_tuning.contracts.ParameterTuningReport": _example_73_ParameterTuningReport, + "scarf.agent.parameter_tuning.contracts.ParameterTuningDependencies": _example_74_ParameterTuningDependencies, + "scarf.agent.parameter_tuning.contracts.ParameterTuningAssayInput": _example_75_ParameterTuningAssayInput, + "scarf.agent.types.AgentDataModel": _example_76_AgentDataModel, + "scarf.agent.types.ArtifactReferenceModel": _example_77_ArtifactReferenceModel, + "scarf.agent.types.BatchSafetyEvidence": _example_78_BatchSafetyEvidence, + "scarf.agent.types.TuningBiologyHandoff": _example_79_TuningBiologyHandoff, + "scarf.agent.types.ToolCallInfo": _example_80_ToolCallInfo, + "scarf.agent.types.AgentUsageInfo": _example_81_AgentUsageInfo, + "scarf.agent.types.AgentRunInfo": _example_82_AgentRunInfo, + "scarf.agent.types.AgentExecutionResult": _example_83_AgentExecutionResult, + "scarf.agent.types.EvidenceItem": _example_84_EvidenceItem, + "scarf.agent.types.Decision": _example_85_Decision, + "scarf.agent.types.NeedsInput": _example_86_NeedsInput, + "scarf.agent.types.StageResult": _example_87_StageResult, +} diff --git a/tests/agent_journal_store.py b/tests/agent_journal_store.py new file mode 100644 index 00000000..86ea6cfa --- /dev/null +++ b/tests/agent_journal_store.py @@ -0,0 +1,51 @@ +"""In-memory journal fixture without numerical work or a provider.""" + +from types import SimpleNamespace + +import zarr +from zarr.storage import MemoryStore + +from scarf.agent.orchestrator import AutomatedWorkflowConfig, AutomatedWorkflowRequest +from scarf.agent.orchestrator import journal +from scarf.agent.orchestrator.models import OrchestrationRequestRecord + + +def memory_journal(workspace: str | None = None): + root = zarr.open_group(store=MemoryStore(), mode="w") + active = root if workspace is None else root.create_group(workspace) + store = SimpleNamespace( + zw=active, + z=root, + zarr_loc="analysis.zarr", + workspace=workspace, + cells=SimpleNamespace(columns=[]), + load_artifact=lambda ref: ref, + ) + prefix = journal._ensure_orchestration_store(store) + request = AutomatedWorkflowRequest( + sourcePath="analysis.zarr", + zarrPath="analysis.zarr", + workspace=workspace, + primaryAssay="RNA", + markerAssay="RNA", + analysisAssays=["RNA"], + studyContext="Two replicated conditions", + studyObjective="Resolve stable populations", + ) + config = AutomatedWorkflowConfig() + record = OrchestrationRequestRecord( + workflowRunId="workflow-1", + request=request, + config=config, + requestSha256=journal._sha256_model(request), + configSha256=journal._sha256_model(config), + modelIdentity="test-model", + inputIdentity={"data": "selected-rna"}, + ) + record = record.model_copy( + update={"contentSha256": journal._record_checksum(record)} + ) + journal._write_model_once( + active, journal._request_key(prefix, record.workflowRunId), record + ) + return store, prefix, record diff --git a/tests/test_agent_analysis_plots.py b/tests/test_agent_analysis_plots.py new file mode 100644 index 00000000..e2ed3dfc --- /dev/null +++ b/tests/test_agent_analysis_plots.py @@ -0,0 +1,171 @@ +"""Bounded, immutable cluster displays for agent results and reports.""" + +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +from scarf.agent import _plots +from scarf.storage.refs import ArtifactRef + + +class VirtualArray: + """A large immutable array that refuses unbounded or whole-column reads.""" + + def __init__(self, rows: int, *, coordinates: bool = False) -> None: + self.shape = (rows, 2) if coordinates else (rows,) + self.ndim = len(self.shape) + self.dtype = np.dtype(float if coordinates else int) + self.coordinates = coordinates + self.read_sizes: list[int] = [] + + def __getitem__(self, selection: slice) -> np.ndarray: + assert isinstance(selection, slice) + assert selection.start is not None and selection.stop is not None + assert selection.stop - selection.start <= _plots.DISPLAY_BLOCK_ROWS + rows = np.arange(selection.start, min(selection.stop, self.shape[0])) + self.read_sizes.append(len(rows)) + if self.coordinates: + return np.column_stack((rows, rows % 11)).astype(float) + return np.where(rows == self.shape[0] - 1, 7, rows % 3) + + +def display_store( + monkeypatch: pytest.MonkeyPatch, n: int = 621_200 +) -> tuple[Any, dict[str, ArtifactRef], dict[str, VirtualArray]]: + refs = { + "cell_selection": ArtifactRef("datastore", "cell_selection", "1" * 64), + "graph": ArtifactRef("assay", "connectivity_map", "2" * 64, "RNA2"), + "clusters": ArtifactRef("assay", "cluster_labels", "3" * 64, "RNA2"), + "umap": ArtifactRef("assay", "embedding", "4" * 64, "RNA2"), + } + arrays = { + "clusters": VirtualArray(n), + "umap": VirtualArray(n, coordinates=True), + } + store = SimpleNamespace(zw=object(), workspace="selected-workspace") + store.load_artifact = lambda ref: { + "values": arrays[next(name for name in arrays if refs[name] == ref)] + } + store.inspect_artifact = lambda ref: SimpleNamespace( + complete=True, + operation="run_umap" if ref == refs["umap"] else "run_leiden_clustering", + inputs={ + "cell_selection": refs["cell_selection"].to_dict(), + "graph": refs["graph"].to_dict(), + }, + ) + monkeypatch.setattr(_plots, "as_zarr_array", lambda value, **_: value) + monkeypatch.setattr( + _plots, "graph_cell_selection", lambda root, ref: refs["cell_selection"] + ) + + def validate(root: Any, ref: ArtifactRef, **kwargs: Any) -> Any: + assert root is store.zw and ref == refs["cell_selection"] + assert kwargs == { + "kind": "cell_selection", + "scope": "datastore", + "assay": None, + "table_path": "cellData", + } + return SimpleNamespace(selected_count=n) + + monkeypatch.setattr(_plots, "validate_stored_selection_integrity", validate) + return store, refs, arrays + + +def test_large_final_map_is_bounded_cluster_aware_and_preserves_counts( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + store, refs, arrays = display_store(monkeypatch) + plot = _plots.plot_final_umap(store, **refs, show=False) + try: + assert plot.provenance.n_cells == 50_000 + assert plot.provenance.extras["input_n_cells"] == 621_200 + assert plot.provenance.extras["workspace"] == "selected-workspace" + assert plot.provenance.extras["layout"] == refs["umap"].to_dict() + counts = plot.tables["cluster_counts"].set_index("cluster")["cells"].to_dict() + assert sum(counts.values()) == 621_200 and counts["7"] == 1 + axes = plot.axes["clusters"] + assert sum(len(artist.get_offsets()) for artist in axes.collections) == 50_000 + from matplotlib.colors import to_rgb + + palette = plot.scales[0].palette + for artist in axes.collections: + positions = np.asarray(artist.get_offsets())[:, 0].astype(int) + expected = np.where(positions == 621_199, 7, positions % 3) + np.testing.assert_allclose( + artist.get_facecolors()[:, :3], + np.asarray([to_rgb(palette[str(label)]) for label in expected]), + ) + assert 621_199 in positions + assert all(max(array.read_sizes) <= 100_000 for array in arrays.values()) + sidecar = plot.save_provenance(tmp_path / "map.json") + assert ( + json.loads(sidecar.read_text())["provenance"]["extras"]["cluster_counts"] + == counts + ) + finally: + plot.close() + + +def test_sampling_is_reproducible_across_read_boundaries_and_keeps_rare_cells() -> None: + values = np.asarray(["common"] * 999 + ["rare"]) + counts = {"common": 999, "rare": 1} + first = _plots._sample_cluster_rows( + values, counts, maximum=50, seed=5, block_rows=73 + ) + second = _plots._sample_cluster_rows( + values, counts, maximum=50, seed=5, block_rows=221 + ) + np.testing.assert_array_equal(first[0], second[0]) + np.testing.assert_array_equal(first[1], second[1]) + assert first[0][-1] == 999 and list(first[1]).count("rare") == 1 + third = _plots._sample_cluster_rows(values, counts, maximum=50, seed=7) + assert not np.array_equal(first[0], third[0]) + + +@pytest.mark.parametrize("mismatch", ["selection", "graph", "producer", "shape"]) +def test_final_map_rejects_incompatible_saved_artifacts( + monkeypatch: pytest.MonkeyPatch, mismatch: str +) -> None: + store, refs, arrays = display_store(monkeypatch, n=12) + inspect = store.inspect_artifact + + def changed(ref: ArtifactRef) -> Any: + status = inspect(ref) + if ref == refs["umap"]: + if mismatch == "selection": + status.inputs["cell_selection"] = ArtifactRef( + "datastore", "cell_selection", "a" * 64 + ).to_dict() + elif mismatch == "graph": + status.inputs["graph"] = ArtifactRef( + "assay", "connectivity_map", "a" * 64, "RNA2" + ).to_dict() + elif mismatch == "producer": + status.operation = "import_dimreduc" + return status + + store.inspect_artifact = changed + if mismatch == "shape": + arrays["umap"].shape = (13, 2) + with pytest.raises(ValueError): + _plots.plot_final_umap(store, **refs, show=False) + assert not any(array.read_sizes for array in arrays.values()) + + +def test_small_map_draws_every_frozen_cell(monkeypatch: pytest.MonkeyPatch) -> None: + store, refs, _ = display_store(monkeypatch, n=12) + plot = _plots.plot_final_umap(store, **refs, show=False) + try: + assert plot.provenance.n_cells == 12 + offsets = np.concatenate( + [artist.get_offsets()[:, 0] for artist in plot.axes["clusters"].collections] + ) + np.testing.assert_array_equal(np.sort(offsets), np.arange(12)) + finally: + plot.close() diff --git a/tests/test_agent_assessment_contract_repair.py b/tests/test_agent_assessment_contract_repair.py new file mode 100644 index 00000000..a78f2307 --- /dev/null +++ b/tests/test_agent_assessment_contract_repair.py @@ -0,0 +1,419 @@ +"""Structured tradeoff repair keeps scientific counterevidence explicit.""" + +from copy import deepcopy +import json +from typing import Any + +from pydantic import ValidationError +from pydantic_ai.exceptions import UnexpectedModelBehavior +from pydantic_ai.messages import ModelResponse, RetryPromptPart, ToolCallPart +from pydantic_ai.models.function import FunctionModel +import pytest + +from scarf.agent.config import AgentRunConfig +from scarf.agent.config.agent_exec import run_agent_sync +from scarf.agent.orchestrator import rna_tuning +from scarf.agent.parameter_tuning.comparisons import ( + bind_comparison_measurements, + comparison_advantages, + validate_comparison_review, +) +from tests.agent_comparison_examples import comparison_review, observed_action + + +def _competing_review() -> tuple[dict[str, Any], dict[str, Any]]: + review = comparison_review() + coverage = review["comparisonCoverage"] + settings = coverage["candidateSettings"] + for candidate in settings.values(): + candidate["metrics"]["macroF1"] = 0.8 + for identifier in ( + "resolution-half", + "genes-two", + "dimensions-thirty", + "neighbors-forty-one", + ): + settings[identifier]["metrics"].update(seedStability=0.97, macroF1=0.8 + 1e-10) + review["currentCandidateId"] = review["selectedCandidateId"] + action = observed_action(review) + for conclusion in action["comparisonConclusions"]: + explanations = [ + f"{row['alternativeCandidateId']} improves {row['metric']} from " + f"{row['preferredValue']!r} to {row['alternativeValue']!r}; " + "the reference marker program still motivates this preference." + for row in conclusion["tradeoffs"] + ] + conclusion["quantitativeReason"] = " ".join(explanations) or ( + "No measured advantage requires explanation on this axis." + ) + conclusion["biologicalReason"] = ( + "Retain the reference marker program while acknowledging every " + "alternative's measured improvements in the quantitative explanation." + ) + conclusion["tradeoffs"] = [] + return coverage, action + + +def test_provider_requires_explicit_tradeoff_field_without_numeric_authorship() -> None: + coverage, action = _competing_review() + output_type = rna_tuning._assessment_output_type( + list(coverage["candidateSettings"]), [], scope="full" + ) + schema = output_type.model_json_schema() + conclusion_schema = schema["$defs"]["ObservedComparisonInterpretation"] + assert "tradeoffs" in conclusion_schema["required"] + assert ( + "comparisonAdvantages" + in conclusion_schema["properties"]["tradeoffs"]["description"] + ) + tradeoff_fields = schema["$defs"]["ObservedTradeoffInterpretation"]["properties"] + assert set(tradeoff_fields) == { + "alternativeCandidateId", + "metric", + "interpretation", + } + del action["comparisonConclusions"][0]["tradeoffs"] + with pytest.raises(ValidationError, match="tradeoffs"): + output_type.model_validate(action) + # The provider contract does not rewrite the existing persisted action schema. + saved = rna_tuning.TuningAction.model_validate(action) + assert saved.comparisonConclusions[0].tradeoffs == [] + + +def test_detailed_prose_cannot_replace_structured_tradeoff_explanations() -> None: + coverage, action = _competing_review() + original = deepcopy(action) + with pytest.raises(ValueError) as failure: + validate_comparison_review(coverage, action) + feedback = str(failure.value) + for axis in ("partition", "hvgCount", "pca", "neighbors"): + assert axis in feedback + for required in ( + "comparisonConclusions", + "tradeoffs", + "alternativeCandidateId", + "metric", + "interpretation", + "quantitativeReason", + "biologicalReason", + "resolution-half", + "genes-two", + "dimensions-thirty", + "neighbors-forty-one", + "0.8000000001", + ): + assert required in feedback + assert action == original + + +@pytest.mark.parametrize("repair", [True, False]) +def test_model_repairs_exact_fields_or_fails_without_accepting_prose( + repair: bool, +) -> None: + coverage, action = _competing_review() + inventory = comparison_advantages(coverage) + offered = json.dumps({"comparisonAdvantages": inventory}, sort_keys=True) + output_type = rna_tuning._assessment_output_type( + list(coverage["candidateSettings"]), [], scope="full" + ) + responses = [] + attempts = [] + + def provider(messages: Any, info: Any) -> ModelResponse: + proposed = deepcopy(action) + if responses: + feedback = " ".join( + str(part.content) + for message in messages + for part in message.parts + if isinstance(part, RetryPromptPart) + ) + for name in ( + "comparisonConclusions", + "tradeoffs", + "alternativeCandidateId", + "metric", + "interpretation", + ): + assert name in feedback + if repair: + for conclusion in proposed["comparisonConclusions"]: + conclusion["tradeoffs"] = [ + { + "alternativeCandidateId": row["alternativeCandidateId"], + "metric": row["metric"], + "interpretation": ( + f"The observed {row['metric']} advantage must be " + "weighed against preserving the reference marker " + "program, rather than asserting numerical superiority." + ), + } + for row in inventory + if row["axis"] == conclusion["axis"] + and row["preferredCandidateId"] + == conclusion["preferredCandidateId"] + ] + responses.append(proposed) + return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, proposed)]) + + def validate(proposed: Any) -> rna_tuning.TuningAction: + bound = bind_comparison_measurements(coverage, proposed.model_dump(mode="json")) + validate_comparison_review(coverage, bound) + return rna_tuning.TuningAction.model_validate(bound) + + def run(): + return run_agent_sync( + model=FunctionModel(provider), + output_type=output_type, + system_prompt="Adjudicate the supplied measured tradeoffs.", + user_prompt=offered, + config=AgentRunConfig(retries=2), + output_validator=validate, + on_attempt=attempts.append, + ) + + if repair: + result = run() + completed = result.output.model_dump(mode="json") + validate_comparison_review(coverage, completed) + measured = { + ( + row["axis"], + row["preferredCandidateId"], + row["alternativeCandidateId"], + row["metric"], + ): row + for row in inventory + } + checked = set() + for conclusion in completed["comparisonConclusions"]: + for tradeoff in conclusion["tradeoffs"]: + expected = measured[ + ( + conclusion["axis"], + conclusion["preferredCandidateId"], + tradeoff["alternativeCandidateId"], + tradeoff["metric"], + ) + ] + assert tradeoff["preferredValue"] == expected["preferredValue"] + assert tradeoff["alternativeValue"] == expected["alternativeValue"] + checked.add(conclusion["axis"]) + assert checked == {"partition", "hvgCount", "pca", "neighbors"} + assert attempts[0].status == "done" + assert len(attempts[0].validationRetries) == 1 + else: + with pytest.raises(UnexpectedModelBehavior, match="maximum output retries"): + run() + assert attempts[0].status == "failed" + assert len(attempts[0].validationRetries) == 3 + assert len(responses) == (2 if repair else 3) + assert len(attempts) == 1 + assert attempts[0].usage.requests == len(responses) + assert json.dumps({"comparisonAdvantages": inventory}, sort_keys=True) == offered + assert all( + "preferredValue" not in tradeoff and "alternativeValue" not in tradeoff + for response in responses + for conclusion in response["comparisonConclusions"] + for tradeoff in conclusion["tradeoffs"] + ) + + +def _permitted_feedback(feedback: str) -> list[dict[str, Any]]: + decoder = json.JSONDecoder() + for start, char in enumerate(feedback): + if char != "{": + continue + try: + value, _ = decoder.raw_decode(feedback[start:]) + except json.JSONDecodeError: + continue + if isinstance(value, dict) and "permittedTradeoffs" in value: + return value["permittedTradeoffs"] + pytest.fail("Repair feedback did not supply exact permittedTradeoffs as JSON") + + +def _apply_permitted(action: dict[str, Any], permitted: list[dict[str, Any]]) -> None: + for conclusion in action["comparisonConclusions"]: + conclusion["tradeoffs"] = [ + { + "alternativeCandidateId": row["alternativeCandidateId"], + "metric": row["metric"], + "interpretation": ( + "This measured advantage is counterevidence to the preference; " + "the recorded marker program motivates retaining the reference " + "while acknowledging the alternative's quantitative improvement." + ), + } + for row in permitted + if row["axis"] == conclusion["axis"] + and row["preferredCandidateId"] == conclusion["preferredCandidateId"] + ] + + +def test_one_retry_repairs_all_axes_without_pruning_or_rewriting_evidence() -> None: + coverage, action = _competing_review() + settings = coverage["candidateSettings"] + settings["resolution-high"]["metrics"]["subsampleStability"] = 0.7 + settings["genes-four"]["metrics"]["seedStability"] = 0.8 + inventory = comparison_advantages(coverage) + _apply_permitted(action, inventory) + by_axis = {row["axis"]: row for row in action["comparisonConclusions"]} + for axis, alternative, metric in ( + ("partition", "resolution-high", "subsampleStability"), + ("partition", "resolution-half", "markerCoherence"), + ("hvgCount", "genes-four", "seedStability"), + ("pca", "dimensions-thirty", "markerCoherence"), + ): + by_axis[axis]["tradeoffs"].append( + { + "alternativeCandidateId": alternative, + "metric": metric, + "interpretation": ( + "The alternative is tied or worse on this metric, which belongs " + "in the general quantitative explanation instead of counterevidence." + ), + } + ) + by_axis["partition"]["tradeoffs"].append( + deepcopy(by_axis["partition"]["tradeoffs"][0]) + ) + by_axis["neighbors"]["tradeoffs"] = [ + row for row in by_axis["neighbors"]["tradeoffs"] if row["metric"] != "macroF1" + ] + frozen_coverage, frozen_action = deepcopy(coverage), deepcopy(action) + responses: list[dict[str, Any]] = [] + attempts = [] + output_type = rna_tuning._assessment_output_type(list(settings), [], scope="full") + + def provider(messages: Any, info: Any) -> ModelResponse: + proposed = deepcopy(action) + if responses: + feedback = next( + str(part.content) + for message in reversed(messages) + for part in message.parts + if isinstance(part, RetryPromptPart) + ) + for identifier in ( + "resolution-high", + "resolution-half", + "genes-four", + "dimensions-thirty", + "neighbors-forty-one", + "duplicate", + "macroF1", + "markerCoherence", + "subsampleStability", + ): + assert identifier in feedback + permitted = _permitted_feedback(feedback) + preferences = { + (row["axis"], row["preferredCandidateId"]) + for row in proposed["comparisonConclusions"] + } + assert {json.dumps(row, sort_keys=True) for row in permitted} == { + json.dumps(row, sort_keys=True) + for row in inventory + if (row["axis"], row["preferredCandidateId"]) in preferences + } + _apply_permitted(proposed, permitted) + responses.append(deepcopy(proposed)) + return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, proposed)]) + + def validate(proposed: Any) -> rna_tuning.TuningAction: + bound = bind_comparison_measurements(coverage, proposed.model_dump(mode="json")) + validate_comparison_review(coverage, bound) + return rna_tuning.TuningAction.model_validate(bound) + + result = run_agent_sync( + model=FunctionModel(provider), + output_type=output_type, + system_prompt="Explain only the observed advantages using exact feedback.", + user_prompt=json.dumps({"comparisonAdvantages": inventory}), + config=AgentRunConfig(retries=2), + output_validator=validate, + on_attempt=attempts.append, + ) + completed = result.output.model_dump(mode="json") + validate_comparison_review(coverage, completed) + assert len(responses) == 2 + assert attempts[0].usage.requests == 2 + assert len(attempts[0].validationRetries) == 1 + for conclusion in completed["comparisonConclusions"]: + assert ( + conclusion["preferredCandidateId"] + == by_axis[conclusion["axis"]]["preferredCandidateId"] + ) + assert ( + conclusion["quantitativeReason"] + == by_axis[conclusion["axis"]]["quantitativeReason"] + ) + for tradeoff in conclusion["tradeoffs"]: + measured = next( + row + for row in inventory + if row["axis"] == conclusion["axis"] + and row["preferredCandidateId"] == conclusion["preferredCandidateId"] + and row["alternativeCandidateId"] == tradeoff["alternativeCandidateId"] + and row["metric"] == tradeoff["metric"] + ) + assert tradeoff["preferredValue"] == measured["preferredValue"] + assert tradeoff["alternativeValue"] == measured["alternativeValue"] + assert action == frozen_action + assert coverage == frozen_coverage + assert all( + "preferredValue" not in row and "alternativeValue" not in row + for response in responses + for conclusion in response["comparisonConclusions"] + for row in conclusion["tradeoffs"] + ) + + +def test_repair_rejects_unknown_and_fabricated_measurements_together() -> None: + coverage, action = _competing_review() + _apply_permitted(action, comparison_advantages(coverage)) + partition = next( + row for row in action["comparisonConclusions"] if row["axis"] == "partition" + ) + partition["tradeoffs"][0]["alternativeValue"] = 100.0 + partition["tradeoffs"].append( + { + "alternativeCandidateId": "unmeasured-candidate", + "metric": "seedStability", + "interpretation": "This candidate does not exist in the supplied evidence.", + } + ) + frozen = deepcopy(action) + with pytest.raises(ValueError) as failure: + bind_comparison_measurements(coverage, action) + feedback = str(failure.value) + assert "unmeasured-candidate" in feedback + assert "exact" in feedback + assert "resolution-half" in feedback + assert _permitted_feedback(feedback) + assert action == frozen + + +def test_changed_preference_recalculates_counterevidence_and_allows_no_advantages() -> ( + None +): + coverage, action = _competing_review() + _apply_permitted(action, comparison_advantages(coverage)) + partition = next( + row for row in action["comparisonConclusions"] if row["axis"] == "partition" + ) + partition["preferredCandidateId"] = "resolution-half" + action["selectedCandidateId"] = "resolution-half" + # Old obligations cannot be reused after a preference changes. + with pytest.raises(ValueError) as failure: + bind_comparison_measurements(coverage, action) + permitted = _permitted_feedback(str(failure.value)) + assert not any(row["axis"] == "partition" for row in permitted) + assert all(row["preferredCandidateId"] != "candidate-two" for row in permitted) + _apply_permitted(action, permitted) + assert partition["tradeoffs"] == [] + bound = bind_comparison_measurements(coverage, action) + validate_comparison_review(coverage, bound) + # General quantitative evidence remains available even without counterevidence. + assert bound["comparisonConclusions"][0]["quantitativeReason"] diff --git a/tests/test_agent_attempt_audit.py b/tests/test_agent_attempt_audit.py new file mode 100644 index 00000000..348e3e1f --- /dev/null +++ b/tests/test_agent_attempt_audit.py @@ -0,0 +1,396 @@ +"""Model attempts retain measured usage and rejected decisions without execution.""" + +import asyncio +from typing import Any + +import pytest +from pydantic_ai import UnexpectedModelBehavior +from pydantic_ai.messages import ( + ModelResponse, + RetryPromptPart, + TextPart, + ThinkingPart, + ToolCallPart, +) +from pydantic_ai.models.function import FunctionModel +from pydantic_ai.usage import RequestUsage + +from scarf.agent.config import AgentRunConfig +from scarf.agent.config.agent_exec import run_agent_async, run_agent_sync +from scarf.agent.decisions.selection import decide +from scarf.agent.types import AgentDataModel, AgentRunInfo, AgentUsageInfo, EvidenceItem + + +class CountDecision(AgentDataModel): + value: int + + +@pytest.mark.parametrize("host", ["sync", "async", "notebook"]) +def test_repaired_attempt_preserves_usage_and_rejection(host: str) -> None: + requests = 0 + attempts: list[AgentRunInfo] = [] + + async def respond(messages: Any, info: Any) -> ModelResponse: + nonlocal requests + requests += 1 + if requests == 2: + assert any( + isinstance(part, RetryPromptPart) + and "exact measured count is 2" in str(part.content) + for message in messages + for part in message.parts + ) + return ModelResponse( + parts=[ToolCallPart(info.output_tools[0].name, {"value": requests})], + usage=RequestUsage(input_tokens=11, output_tokens=3), + ) + + def validate(output: CountDecision) -> CountDecision: + if output.value != 2: + raise ValueError("The exact measured count is 2") + return output + + kwargs = dict( + model=FunctionModel(respond), + output_type=CountDecision, + system_prompt="Use measured counts.", + user_prompt="Choose the measured count.", + output_validator=validate, + on_attempt=attempts.append, + ) + if host == "async": + result = asyncio.run(run_agent_async(**kwargs)) + elif host == "notebook": + + async def notebook() -> Any: + return run_agent_sync(**kwargs) + + result = asyncio.run(notebook()) + else: + result = run_agent_sync(**kwargs) + assert requests == 2 + assert attempts == [result.runInfo] + assert result.runInfo.status == "done" + assert result.runInfo.usage.inputTokens == 22 + assert result.runInfo.usage.outputTokens == 6 + assert result.runInfo.usage.availability == "reported" + assert len(result.runInfo.validationRetries) == 1 + assert result.runInfo.validationRetries[0].response == {"value": 1} + assert result.runInfo.validationRetries[0].requestIndex == 1 + + +def test_exhausted_attempt_keeps_exact_error_usage_and_every_rejection() -> None: + attempts: list[AgentRunInfo] = [] + + async def respond(_messages: Any, info: Any) -> ModelResponse: + return ModelResponse( + parts=[ToolCallPart(info.output_tools[0].name, {"value": 1})], + usage=RequestUsage(input_tokens=7, output_tokens=2), + ) + + def reject(_output: CountDecision) -> CountDecision: + raise ValueError("Missing the measured alternative's marker advantage") + + with pytest.raises(UnexpectedModelBehavior) as caught: + run_agent_sync( + model=FunctionModel(respond), + output_type=CountDecision, + system_prompt="Assess evidence.", + user_prompt="Select only supported settings.", + config=AgentRunConfig(retries=1), + output_validator=reject, + on_attempt=attempts.append, + ) + info = caught.value.agent_run_info + assert attempts == [info] + assert info.status == "failed" + assert info.usage.requests == 2 + assert info.usage.inputTokens == 14 + assert info.usage.outputTokens == 4 + assert info.usage.availability == "reported" + assert len(info.validationRetries) == 2 + assert [row.requestIndex for row in info.validationRetries] == [1, 2] + assert all(item.response == {"value": 1} for item in info.validationRetries) + assert "marker advantage" in info.error + + +def test_failed_request_does_not_invent_provider_usage() -> None: + attempts: list[AgentRunInfo] = [] + + async def fail(_messages: Any, _info: Any) -> ModelResponse: + raise RuntimeError("Provider unavailable") + + with pytest.raises(RuntimeError, match="Provider unavailable") as caught: + run_agent_sync( + model=FunctionModel(fail), + output_type=CountDecision, + system_prompt="Assess evidence.", + user_prompt="Choose a count.", + on_attempt=attempts.append, + ) + assert attempts == [caught.value.agent_run_info] + assert attempts[0].usage.availability == "unavailable" + assert attempts[0].validationRetries == [] + + +@pytest.mark.parametrize("provider_fails", [False, True]) +def test_journal_callback_failure_preserves_model_outcome(provider_fails: bool) -> None: + from scarf.agent.config.agent_exec import describe_agent_error + + calls = [] + + async def respond(_messages: Any, info: Any) -> ModelResponse: + if provider_fails: + raise RuntimeError("Original provider outage") + return ModelResponse( + parts=[ToolCallPart(info.output_tools[0].name, {"value": 2})], + usage=RequestUsage(input_tokens=11, output_tokens=2), + ) + + def failed_journal(info: AgentRunInfo) -> None: + calls.append(info) + raise OSError("Test history store is unavailable") + + with pytest.raises(RuntimeError if provider_fails else OSError) as caught: + run_agent_sync( + model=FunctionModel(respond), + output_type=CountDecision, + system_prompt="Interpret saved evidence.", + user_prompt="Choose the measured count.", + on_attempt=failed_journal, + ) + assert len(calls) == 1 + assert caught.value.agent_run_info == calls[0] + assert calls[0].status == ("failed" if provider_fails else "done") + assert calls[0].runId + detail = describe_agent_error(caught.value) + assert "history store is unavailable" in detail + if provider_fails: + assert "Original provider outage" in detail + else: + assert "model completed" in detail + + +def test_provider_outage_retains_partial_usage_and_completed_tool_evidence() -> None: + requests = 0 + operations = 0 + + async def measure() -> int: + nonlocal operations + operations += 1 + return 2 + + async def respond(_messages: Any, _info: Any) -> ModelResponse: + nonlocal requests + requests += 1 + if requests > 1: + raise RuntimeError("Provider unavailable after the measured tool result") + return ModelResponse( + parts=[ToolCallPart("measure", {})], + usage=RequestUsage(input_tokens=11, output_tokens=2), + ) + + with pytest.raises(RuntimeError, match="after the measured") as caught: + run_agent_sync( + model=FunctionModel(respond), + output_type=CountDecision, + system_prompt="Measure once before interpreting.", + user_prompt="Assess the measured count.", + tools=[measure], + ) + info = caught.value.agent_run_info + assert operations == 1 + assert requests == 2 + assert info.usage.requests == 2 # Include the failed provider request. + assert info.usage.inputTokens == 11 + assert info.usage.availability == "partial" + assert len(info.providerFailures) == 1 + assert info.providerFailures[0].requestIndex == 2 + assert info.providerFailures[0].retryDelaySeconds is None + assert [call.toolName for call in info.toolCalls] == ["measure"] + + +def test_cancelled_attempt_retains_known_execution_information() -> None: + attempts: list[AgentRunInfo] = [] + + async def respond(_messages: Any, _info: Any) -> ModelResponse: + raise asyncio.CancelledError("Analysis interrupted") + + with pytest.raises(asyncio.CancelledError) as caught: + asyncio.run( + run_agent_async( + model=FunctionModel(respond), + output_type=CountDecision, + system_prompt="Assess evidence.", + user_prompt="Select a count.", + on_attempt=attempts.append, + ) + ) + assert attempts == [caught.value.agent_run_info] + assert attempts[0].errorType == "CancelledError" + assert attempts[0].usage.availability == "unavailable" + + +def test_old_run_information_serializes_without_added_default_fields() -> None: + payload = { + "agentName": "prior", + "modelName": "model", + "runId": "id", + "durationSeconds": 12.0, + "usage": { + "inputTokens": 1, + "outputTokens": 2, + "totalTokens": 3, + "requests": 1, + "toolCalls": 0, + }, + "toolCalls": [], + } + assert AgentRunInfo.model_validate(payload).model_dump(mode="json") == payload + assert "availability" not in AgentUsageInfo().model_dump(mode="json") + + +def test_schema_repair_does_not_repeat_a_successful_tool_operation() -> None: + operations = 0 + requests = 0 + + async def measure(value: int) -> int: + nonlocal operations + operations += 1 + return value + + async def respond(_messages: Any, info: Any) -> ModelResponse: + nonlocal requests + requests += 1 + if requests <= 2: + part = ToolCallPart("measure", {"value": "invalid" if requests == 1 else 2}) + else: + part = ToolCallPart(info.output_tools[0].name, {"value": 2}) + return ModelResponse( + parts=[part], usage=RequestUsage(input_tokens=5, output_tokens=1) + ) + + result = run_agent_sync( + model=FunctionModel(respond), + output_type=CountDecision, + system_prompt="Use one successful measured value.", + user_prompt="Measure the value.", + tools=[measure], + ) + assert operations == 1 + assert result.output.value == 2 + assert requests == 3 + assert len(result.runInfo.validationRetries) == 1 + assert result.runInfo.validationRetries[0].source == "tool" + assert result.runInfo.validationRetries[0].response == {"value": "invalid"} + + +@pytest.mark.parametrize("content", ["none", "text", "thinking", "both"]) +def test_failed_schema_attempt_retains_the_last_invalid_response(content: str) -> None: + attempts: list[AgentRunInfo] = [] + + async def respond(_messages: Any, info: Any) -> ModelResponse: + parts: list[Any] = [] + if content in {"text", "both"}: + parts.append(TextPart("Here is the requested structured result.")) + if content in {"thinking", "both"}: + parts.append(ThinkingPart("Checking the observed count.")) + parts.append(ToolCallPart(info.output_tools[0].name, {"value": "invalid"})) + return ModelResponse( + parts=parts, usage=RequestUsage(input_tokens=5, output_tokens=2) + ) + + with pytest.raises(UnexpectedModelBehavior) as caught: + run_agent_sync( + model=FunctionModel(respond), + output_type=CountDecision, + system_prompt="Use measured counts.", + user_prompt="Choose a count.", + config=AgentRunConfig(retries=1), + on_attempt=attempts.append, + ) + info = caught.value.agent_run_info + assert attempts == [info] + assert len(info.validationRetries) == 2 + assert all(row.response == {"value": "invalid"} for row in info.validationRetries) + assert [row.requestIndex for row in info.validationRetries] == [1, 2] + assert info.usage.inputTokens == 10 + assert info.usage.requests == 2 + assert info.status == "failed" + + +def test_exhausted_schema_with_multiple_calls_keeps_error_without_guessing_attribution() -> ( + None +): + async def respond(_messages: Any, info: Any) -> ModelResponse: + return ModelResponse( + parts=[ + TextPart("Two proposed results."), + ToolCallPart( + info.output_tools[0].name, + {"value": "first-invalid"}, + tool_call_id="first", + ), + ToolCallPart( + info.output_tools[0].name, + {"value": "second-invalid"}, + tool_call_id="second", + ), + ], + usage=RequestUsage(input_tokens=5, output_tokens=2), + ) + + with pytest.raises(UnexpectedModelBehavior) as caught: + run_agent_sync( + model=FunctionModel(respond), + output_type=CountDecision, + system_prompt="Use one measured count.", + user_prompt="Choose a count.", + config=AgentRunConfig(retries=0), + ) + info = caught.value.agent_run_info + assert "ValidationError" in info.error + assert info.validationRetries == [] + assert info.status == "failed" + assert info.usage.requests == 1 + assert info.usage.inputTokens == 5 + + +def test_exact_evidence_id_errors_repair_inside_the_model_loop() -> None: + requests = 0 + + async def respond(messages: Any, info: Any) -> ModelResponse: + nonlocal requests + requests += 1 + if requests == 2: + assert any( + isinstance(part, RetryPromptPart) + and "not in evidence ids" in str(part.content) + for message in messages + for part in message.parts + ) + identifier = "unrelated:matrix:raw/X" if requests == 1 else "matrix:raw/X" + return ModelResponse( + parts=[ + ToolCallPart( + info.output_tools[0].name, + { + "selectedId": identifier, + "rationale": "The observed raw matrix contains integer counts.", + "evidenceIds": [identifier], + }, + ) + ] + ) + + result = decide( + model=FunctionModel(respond), + question="Which measured matrix contains raw counts?", + evidence=[ + EvidenceItem( + id="matrix:raw/X", label="Raw counts", summary="Integer counts" + ) + ], + ) + assert requests == 2 + assert result.selectedId == "matrix:raw/X" diff --git a/tests/test_agent_beginner.py b/tests/test_agent_beginner.py new file mode 100644 index 00000000..02744368 --- /dev/null +++ b/tests/test_agent_beginner.py @@ -0,0 +1,508 @@ +"""Beginner RNA entry point and exact completed-result access.""" + +from tests.agent_examples import example + +import hashlib +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + +pytest.importorskip("pydantic_ai") + +from scarf.agent import AnalysisError, AutomatedWorkflowResult, analyze_rna +from scarf.agent.orchestrator import api, journal +from scarf.agent.orchestrator.main import AgentOrchestrator +from scarf.agent.orchestrator.models import ( + AutomatedWorkflowConfig, + AutomatedWorkflowRequest, + FinalAnalysisHandoff, + OrchestrationRequestRecord, + artifact_model_to_ref, +) +from scarf.agent.types import ArtifactReferenceModel +from scarf.agent.record_io import canonical_json_bytes + + +def _completed_result(root: Path) -> AutomatedWorkflowResult: + return AutomatedWorkflowResult( + status="completed", + currentStage="analysis_finalization", + zarrPath=str(root), + workspace="analysis", + workflowRunId="workflow-1", + ) + + +def _final_analysis() -> FinalAnalysisHandoff: + final = example(FinalAnalysisHandoff) + final.umap = ArtifactReferenceModel( + assay="RNA", kind="embedding", artifactId="6" * 64 + ) + final.markers = ArtifactReferenceModel( + assay="RNA", kind="marker_table", artifactId="7" * 64 + ) + return final + + +def test_analyze_rna_passes_one_request_and_bounded_defaults( + monkeypatch: pytest.MonkeyPatch, +) -> None: + called: dict[str, Any] = {} + outcome = _completed_result(Path("study.zarr")) + + class Orchestrator: + def __init__(self, model: Any, *, config: AutomatedWorkflowConfig) -> None: + called.update(model=model, config=config) + + def run(self, request: AutomatedWorkflowRequest) -> AutomatedWorkflowResult: + called["request"] = request + return outcome + + monkeypatch.setattr(api, "AgentOrchestrator", Orchestrator) + model = object() + result = analyze_rna( + Path("study.h5ad"), + model=model, + study_context="Human blood, one donor.", + study_objective="Identify stable populations.", + assay="counts", + zarr_path=Path("study.zarr"), + ) + assert result is outcome + assert called["model"] is model + config = called["config"] + assert config.inputPolicy == "unattended" + assert config.scoreDoublets is False + assert config.model_dump(mode="json")["scoreDoublets"] is False + assert config.screeningCells is None + assert config.maxScreeningCells == 100_000 + assert config.maxScreeningEvaluations == 24 + assert config.maxTotalScreeningEvaluations == 48 + assert config.maxFullGraphs == 4 + assert config.maxFullPartitions == 8 + assert config.maxFullRepairs == 1 + request = called["request"] + assert request.sourcePath == "study.h5ad" + assert request.zarrPath == "study.zarr" + assert request.primaryAssay == request.markerAssay == "counts" + assert request.analysisAssays == ["counts"] + assert request.ingestDirections == {} + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_beginner_can_control_advisory_doublet_scoring(monkeypatch, enabled): + observed = [] + outcome = _completed_result(Path("study.zarr")) + + def orchestrator(model, *, config): + observed.append(config) + return SimpleNamespace(run=lambda request: outcome) + + monkeypatch.setattr(api, "AgentOrchestrator", orchestrator) + assert ( + analyze_rna( + "study.h5ad", + model=object(), + study_context="Human blood from independent physical captures.", + study_objective="Discover stable populations.", + score_doublets=enabled, + ) + is outcome + ) + assert observed[0].scoreDoublets is enabled + if not enabled: + assert observed[0].model_dump(mode="json")["scoreDoublets"] is False + assert "scoreDoublets" in observed[0].model_fields_set + + +@pytest.mark.parametrize("value", [None, 0, 1, "false", "true", [], {}]) +def test_doublet_scoring_flag_requires_a_boolean_before_pipeline_work( + monkeypatch, value +): + def forbidden(*args, **kwargs): + pytest.fail("Invalid scoring policy must not start an orchestrator") + + monkeypatch.setattr(api, "AgentOrchestrator", forbidden) + with pytest.raises(ValueError, match="scoreDoublets"): + analyze_rna( + "study.h5ad", + model=object(), + study_context="Human blood.", + study_objective="Discover stable populations.", + score_doublets=value, + ) + with pytest.raises(ValueError, match="scoreDoublets"): + AutomatedWorkflowConfig(scoreDoublets=value) + + +def test_legacy_doublet_policy_preserves_saved_request_bytes_and_checksums(): + legacy_config = AutomatedWorkflowConfig(screeningCells=50_000).model_dump( + mode="json", exclude={"scoreDoublets"} + ) + restored = AutomatedWorkflowConfig.model_validate(legacy_config) + assert restored.scoreDoublets is True + assert restored.model_dump(mode="json") == legacy_config + assert "scoreDoublets" not in AutomatedWorkflowConfig().model_dump(mode="json") + assert "scoreDoublets" not in AutomatedWorkflowConfig( + scoreDoublets=True + ).model_dump(mode="json") + request = example(AutomatedWorkflowRequest).model_dump(mode="json") + saved = { + "recordType": "automatedWorkflowRequest", + "inputIdentity": {"source": "study.h5ad"}, + "modelIdentity": "test-model", + "workflowRunId": "legacy-workflow", + "createdAtNs": 1, + "request": request, + "config": legacy_config, + "requestSha256": hashlib.sha256(canonical_json_bytes(request)).hexdigest(), + "configSha256": hashlib.sha256(canonical_json_bytes(legacy_config)).hexdigest(), + } + saved["contentSha256"] = hashlib.sha256(canonical_json_bytes(saved)).hexdigest() + original = canonical_json_bytes(saved) + loaded = OrchestrationRequestRecord.model_validate_json(original) + assert loaded.config.scoreDoublets is True + assert canonical_json_bytes(loaded.model_dump(mode="json")) == original + assert journal._record_checksum(loaded) == saved["contentSha256"] + assert ( + hashlib.sha256( + canonical_json_bytes(loaded.config.model_dump(mode="json")) + ).hexdigest() + == saved["configSha256"] + ) + + +def test_doublet_policy_cannot_be_changed_while_resuming_saved_work(): + legacy = AutomatedWorkflowConfig.model_validate({"screeningCells": 50_000}) + AgentOrchestrator(object())._validate_resume_config(legacy) + AgentOrchestrator( + object(), config=AutomatedWorkflowConfig(scoreDoublets=True) + )._validate_resume_config(legacy) + with pytest.raises(ValueError, match="execution settings differ"): + AgentOrchestrator( + object(), config=AutomatedWorkflowConfig(scoreDoublets=False) + )._validate_resume_config(legacy) + + disabled = AutomatedWorkflowConfig(scoreDoublets=False) + AgentOrchestrator(object())._validate_resume_config(disabled) + AgentOrchestrator( + object(), config=AutomatedWorkflowConfig(scoreDoublets=False) + )._validate_resume_config(disabled) + with pytest.raises(ValueError, match="execution settings differ"): + AgentOrchestrator( + object(), config=AutomatedWorkflowConfig(scoreDoublets=True) + )._validate_resume_config(disabled) + + +@pytest.mark.parametrize("status", ["failed", "abstained", "needsInput"]) +def test_beginner_failure_raises_with_resumable_result( + monkeypatch: pytest.MonkeyPatch, status: str +) -> None: + outcome = AutomatedWorkflowResult.model_validate( + { + "status": status, + "currentStage": "experimental_context", + "zarrPath": "study.zarr", + "workflowRunId": "workflow-1", + "notes": ["Capture identity is unresolved"], + } + ) + monkeypatch.setattr( + api, + "AgentOrchestrator", + lambda *_args, **_kwargs: SimpleNamespace(run=lambda _request: outcome), + ) + with pytest.raises(AnalysisError, match="Capture identity is unresolved") as error: + analyze_rna( + "study.zarr", + model=object(), + study_context="Human blood.", + study_objective="Identify stable populations.", + ) + assert error.value.result is outcome + assert "Resume workflow 'workflow-1' in 'study.zarr'" in str(error.value) + + +def test_beginner_rejects_removed_candidate_control() -> None: + with pytest.raises(TypeError, match="max_candidates"): + analyze_rna( + "study.zarr", + model=object(), + study_context="Human blood.", + study_objective="Identify stable populations.", + **{"max_candidates": 1}, + ) + + +@pytest.mark.parametrize( + "obsolete", + [ + "maxRefinedCandidatesPerAssay", + "maxHarmonyCandidatesPerAssay", + "runConfoundedHarmonyDiagnostic", + "maxCandidateEvaluations", + "maxIdentityFeatures", + "hvgCandidateCounts", + "pcaCandidateDimensions", + "graphNeighborCandidates", + "leidenResolutionCandidates", + "maxRevisions", + "maxCandidateBranches", + "primaryInitialCandidates", + "secondaryInitialCandidates", + "integrationResolutionCandidates", + "maxGraphAssays", + "leidenSeeds", + "clusterSubsamples", + "clusterSubsampleFraction", + ], +) +def test_legacy_config_and_saved_requests_fail_explicitly(obsolete: str) -> None: + old_config = {obsolete: 1} + with pytest.raises(ValueError, match="Create a new single-RNA workflow"): + AutomatedWorkflowConfig.model_validate(old_config) + saved = OrchestrationRequestRecord( + inputIdentity={}, + modelIdentity="test", + request=example(AutomatedWorkflowRequest), + ).model_dump(mode="json") + saved["config"].update(old_config) + with pytest.raises(ValueError, match="cannot be resumed or regenerated"): + OrchestrationRequestRecord.model_validate(saved) + + +@pytest.mark.parametrize( + "limits", + [ + {"screeningCells": 19}, + {"screeningCells": 100, "maxScreeningCells": 99}, + {"maxScreeningEvaluations": 3}, + {"maxScreeningEvaluations": 12, "maxTotalScreeningEvaluations": 11}, + {"maxFullGraphs": 0}, + {"maxFullPartitions": 0}, + {"maxFullRepairs": 2}, + ], +) +def test_impossible_work_limits_fail_before_execution(limits: dict[str, int]) -> None: + with pytest.raises(ValueError): + AutomatedWorkflowConfig.model_validate(limits) + + +@pytest.mark.parametrize( + "routing", + [ + {"analysisAssays": ["RNA", "ADT"]}, + {"pairedAssays": ["RNA", "ADT"]}, + {"primaryAssay": "RNA", "analysisAssays": ["counts"]}, + {"primaryAssay": "RNA", "markerAssay": "ADT"}, + {"experimentalDirections": {"hypothesisTesting": {}}}, + ], +) +def test_request_rejects_unsupported_routing(routing: dict[str, Any]) -> None: + values = example(AutomatedWorkflowRequest).model_dump(mode="json") + with pytest.raises(ValueError): + AutomatedWorkflowRequest.model_validate({**values, **routing}) + + +def test_result_helpers_resolve_exact_journal_refs_and_workspace( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from scarf.agent import _plots + + result = _completed_result(tmp_path) + final = _final_analysis() + original = result.model_dump(mode="json") + opened: list[tuple[str, str, str | None]] = [] + snapshots: list[tuple[Any, str]] = [] + plotted: list[dict[str, Any]] = [] + markers: list[dict[str, Any]] = [] + plot_result, marker_table = object(), object() + store = SimpleNamespace( + get_markers=lambda **options: markers.append(options) or marker_table + ) + + def open_store(path: str, run_id: str, *, workspace: str | None) -> Any: + opened.append((path, run_id, workspace)) + return store + + def snapshot(target: Any, run_id: str) -> dict[str, Any]: + snapshots.append((target, run_id)) + return {"status": "completed", "finalAnalysis": final.model_dump(mode="json")} + + def plot(target: Any, **options: Any) -> Any: + assert target is store + plotted.append(options) + return plot_result + + monkeypatch.setattr(journal, "open_analysis_store", open_store) + monkeypatch.setattr(journal, "analysis_snapshot", snapshot) + monkeypatch.setattr(_plots, "plot_final_umap", plot) + assert result.plot_embedding(figsize=(8, 5)) is plot_result + assert result.get_markers(group_id="2", min_score=0.5) is marker_table + assert opened == [(str(tmp_path), "workflow-1", "analysis")] * 2 + assert snapshots == [(store, "workflow-1")] * 2 + assert final.umap is not None and final.clusters is not None + assert final.cellSelection is not None and final.graph is not None + assert final.markers is not None + assert plotted == [ + { + "umap": artifact_model_to_ref(final.umap), + "clusters": artifact_model_to_ref(final.clusters), + "cell_selection": artifact_model_to_ref(final.cellSelection), + "graph": artifact_model_to_ref(final.graph), + "figsize": (8, 5), + } + ] + assert markers == [ + { + "marker": artifact_model_to_ref(final.markers), + "group_id": "2", + "min_score": 0.5, + "min_frac_exp": 0.2, + } + ] + assert result.model_dump(mode="json") == original + assert "finalAnalysis" not in original + assert "workflowRun" not in original + with pytest.raises(ValueError, match="exact completed cluster map"): + result.plot_embedding(layout=artifact_model_to_ref(final.umap)) + with pytest.raises(ValueError, match="exact completed cluster map"): + result.plot_embedding(color_by="condition") + + +@pytest.mark.parametrize("method", ["plot_embedding", "get_markers", "report"]) +def test_result_helpers_explain_noncompleted_outcome(method: str) -> None: + result = AutomatedWorkflowResult(notes=["Input file is missing"]) + with pytest.raises( + AnalysisError, match="failed during ingest.*Input file is missing" + ): + getattr(result, method)() + + +def test_result_report_regenerates_from_exact_saved_analysis( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import scarf.agent.report.generator as generator + + result = _completed_result(tmp_path) + final = _final_analysis() + store = object() + expected = tmp_path / "analysis/agents/orchestrations/workflow-1/report/index.html" + generated: list[tuple[Any, str]] = [] + + def open_store(path: str, run_id: str, *, workspace: str | None) -> Any: + assert (path, run_id, workspace) == (str(tmp_path), "workflow-1", "analysis") + return store + + def generate(target: Any, run_id: str) -> Path: + generated.append((target, run_id)) + expected.parent.mkdir(parents=True, exist_ok=True) + expected.write_text("Analysis", encoding="utf-8") + return expected + + monkeypatch.setattr(journal, "open_analysis_store", open_store) + monkeypatch.setattr( + journal, + "analysis_snapshot", + lambda *_args: { + "status": "completed", + "finalAnalysis": final.model_dump(mode="json"), + }, + ) + monkeypatch.setattr(generator, "generate_agent_report", generate) + assert result.report() == expected + assert result.report() == expected + assert generated == [(store, "workflow-1")] * 2 + + +@pytest.mark.parametrize("workspace", [None, "analysis"]) +def test_legacy_saved_config_blocks_resume_and_report_without_changing_artifacts( + tmp_path: Path, workspace: str | None +) -> None: + import hashlib + + import numpy as np + + from scarf.agent import record_io + from scarf.agent.orchestrator import AgentOrchestrator + from scarf.agent.orchestrator.models import AutomatedWorkflowResumeRequest + from scarf.agent.report import generate_agent_report + from scarf.datastore.datastore import DataStore + from tests.agent_orchestrator_store import create_store + + path = create_store(tmp_path / "legacy.zarr", workspace=workspace) + store = DataStore( + str(path), + default_assay="RNA", + min_features_per_cell=-1, + mito_pattern="", + ribo_pattern="", + nthreads=2, + workspace=workspace, + ) + cells = store.snapshot_cell_selection() + features = store.select_all_features(from_assay="RNA") + normalized = store.run_normalization(cells, features) + original_values = np.asarray(store.load_artifact(normalized)["data"][:]) + prefix = journal._ensure_orchestration_store(store) + request = AutomatedWorkflowRequest( + sourcePath=str(path), + zarrPath=str(path), + studyContext="One human RNA sample.", + studyObjective="Inspect stable cell populations.", + workspace=workspace, + primaryAssay="RNA", + ) + old_config = AutomatedWorkflowConfig().model_dump(mode="json") + old_config["maxCandidateBranches"] = 24 + payload = OrchestrationRequestRecord( + inputIdentity={}, + modelIdentity="test", + workflowRunId="legacy-config", + request=request, + requestSha256=journal._sha256_model(request), + ).model_dump(mode="json") + payload["config"] = old_config + payload["configSha256"] = hashlib.sha256( + record_io.canonical_json_bytes(old_config) + ).hexdigest() + payload["contentSha256"] = hashlib.sha256( + record_io.canonical_json_bytes( + {key: value for key, value in payload.items() if key != "contentSha256"} + ) + ).hexdigest() + request_key = journal._request_key(prefix, "legacy-config") + original_request = record_io.display_json_bytes(payload) + journal._write_key_once(store.zw, request_key, original_request) + + message = "Unsupported saved agent workflow.*cannot be resumed or regenerated" + resume_request = AutomatedWorkflowResumeRequest( + zarrPath=str(path), workflowRunId="legacy-config", workspace=workspace + ) + orchestrator = AgentOrchestrator(object()) + with pytest.raises(ValueError, match=message): + orchestrator.load_request_for_resume(resume_request) + failed = orchestrator.resume(resume_request) + assert failed.status == "failed" + assert failed.workflowRunId == "legacy-config" + assert any("Unsupported saved agent workflow" in note for note in failed.notes) + + with pytest.raises(ValueError, match=message): + generate_agent_report(path, "legacy-config", workspace=workspace) + + reopened = DataStore( + str(path), + default_assay="RNA", + min_features_per_cell=-1, + mito_pattern="", + ribo_pattern="", + nthreads=2, + zarr_mode="r", + workspace=workspace, + ) + assert record_io.read_key(reopened.zw, request_key) == original_request + np.testing.assert_array_equal( + reopened.load_artifact(normalized)["data"][:], original_values + ) diff --git a/tests/test_agent_biological_interpretation.py b/tests/test_agent_biological_interpretation.py index 50bca054..8afa9470 100644 --- a/tests/test_agent_biological_interpretation.py +++ b/tests/test_agent_biological_interpretation.py @@ -1,5 +1,7 @@ """Tests for the bounded Biological Interpretation Agent.""" +from tests.agent_examples import example + import asyncio from types import SimpleNamespace @@ -7,7 +9,7 @@ import pandas as pd import pytest import zarr -from pydantic_ai import ModelRetry, RunContext +from pydantic_ai import ModelRetry, RunContext, UnexpectedModelBehavior from pydantic_ai.messages import ( ModelMessage, ModelResponse, @@ -19,12 +21,13 @@ from pydantic_ai.usage import RunUsage from zarr.storage import MemoryStore -import scarf.agent.biological_interpretation as biological_module +import scarf.agent.biological_interpretation.agent as biological_agent +import scarf.agent.biological_interpretation.tools as biological_tools +import scarf.agent.biological_interpretation.validation as biological_validation +from scarf.agent.biological_interpretation.agent import _SYSTEM_PROMPT from scarf.agent.biological_interpretation import ( - _SYSTEM_PROMPT, BiologicalContext, BiologicalInterpretationAgent, - BiologicalInterpretationDependencies, BiologicalInterpretationNeedsInput, BiologicalInterpretationReport, ClusterCompositionEvidence, @@ -40,6 +43,9 @@ inspect_cluster_markers, validate_biological_interpretation_report, ) +from scarf.agent.biological_interpretation.contracts import ( + BiologicalInterpretationDependencies, +) from scarf.agent.types import ( AgentRunInfo, ArtifactReferenceModel, @@ -262,7 +268,7 @@ def test_models_have_blank_and_example_constructors() -> None: ) for model in models: assert isinstance(model.get_blank(), model) - assert isinstance(model.get_example(), model) + assert isinstance(example(model), model) assert all("_" not in field_name for field_name in model.model_fields) @@ -536,6 +542,29 @@ def test_validator_rejects_unobserved_evidence() -> None: validate_biological_interpretation_report(report, run_context.deps) +def test_validator_rejects_empty_input_questions_and_model_authored_failure() -> None: + store = FakeStore() + run_context = context(store, marker=store.marker) + asyncio.run(inspect_cluster_composition(run_context)) + + with pytest.raises(ModelRetry, match="concrete input question"): + validate_biological_interpretation_report( + BiologicalInterpretationReport( + status="needsInput", + needsInput=BiologicalInterpretationNeedsInput(), + ), + run_context.deps, + ) + with pytest.raises(ModelRetry, match="Do not return failed"): + validate_biological_interpretation_report( + BiologicalInterpretationReport( + status="failed", + limitations=["The model was uncertain."], + ), + run_context.deps, + ) + + def test_treatment_observation_requires_evidence_for_its_exact_cluster() -> None: store = FakeStore() run_context = context(store, marker=store.marker) @@ -873,6 +902,67 @@ async def reply(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: } +def test_biological_interpretation_preserves_evidence_without_completed_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = FakeStore() + marker_retries: list[int] = [] + + def unavailable_structured_output(**kwargs: object) -> None: + deps = kwargs["deps"] + assert isinstance(deps, BiologicalInterpretationDependencies) + marker_tool = next( + tool + for tool in kwargs["tools"] + if tool.name == "inspect_cluster_markers_batch" + ) + marker_retries.append(marker_tool.max_retries) + run_context = SimpleNamespace(deps=deps) + asyncio.run(inspect_cluster_composition(run_context)) + asyncio.run( + inspect_cluster_markers_batch( + run_context, + cluster_ids=list(deps.clusterValues), + ) + ) + raise UnexpectedModelBehavior("structured output unavailable") + + monkeypatch.setattr( + biological_agent, + "run_agent_sync", + unavailable_structured_output, + ) + result = BiologicalInterpretationAgent(object()).run( + store, + cluster=store.cluster, + marker=store.marker, + ) + + assert result.status == "failed" + assert result.runInfo.agentName == "biological_interpretation_failed" + assert result.clusterInterpretations == [] + assert result.clusterArtifact == artifact_model(store.cluster) + assert result.markerArtifact == artifact_model(store.marker) + assert result.evidenceIds + assert result.treatmentObservations == [] + assert marker_retries == [1] + + +def test_marker_batch_cache_rejects_different_clusters() -> None: + store = FakeStore() + run_context = context(store, marker=store.marker) + asyncio.run(inspect_cluster_composition(run_context)) + + first = asyncio.run(inspect_cluster_markers_batch(run_context, cluster_ids=["0"])) + repeated = asyncio.run( + inspect_cluster_markers_batch(run_context, cluster_ids=["0"]) + ) + + assert repeated is first + with pytest.raises(ModelRetry, match="already completed"): + asyncio.run(inspect_cluster_markers_batch(run_context, cluster_ids=["1"])) + + def test_exact_cluster_artifact_conflict_is_rejected_before_model_execution() -> None: store = FakeStore() tuning_handoff = TuningBiologyHandoff( @@ -929,7 +1019,7 @@ def test_handoff_selection_must_match_exact_cluster_selection() -> None: def test_integrated_handoff_uses_marker_assay_without_claiming_graph_ownership( monkeypatch: pytest.MonkeyPatch, ) -> None: - from scarf.agent import biological_interpretation as module + from scarf.agent.biological_interpretation import agent as module store = FakeStore() store.cluster = ArtifactRef( @@ -1009,17 +1099,17 @@ def test_integrated_handoff_requires_explicit_marker_assay() -> None: def test_biological_scalar_and_column_validation_edges() -> None: - assert biological_module._finite_float(None) is None - assert biological_module._finite_float("not-a-number") is None - assert biological_module._finite_float(float("nan")) is None - assert biological_module._string_value(np.int64(3)) == "3" + assert biological_tools._finite_float(None) is None + assert biological_tools._finite_float("not-a-number") is None + assert biological_tools._finite_float(float("nan")) is None + assert biological_tools._string_value(np.int64(3)) == "3" with pytest.raises(ValueError, match="not present in cell metadata"): - biological_module._check_column(FakeStore(), "missing", "sample column") + biological_tools._check_column(FakeStore(), "missing", "sample column") tool = SimpleNamespace(name="future_tool") deps = context(FakeStore()).deps assert ( - biological_module._prepare_biological_interpretation_tool( + biological_validation._prepare_biological_interpretation_tool( SimpleNamespace(deps=deps), tool ) is tool @@ -1110,7 +1200,7 @@ def run(deps: BiologicalInterpretationDependencies) -> None: lambda _ref: {"values": np.asarray([], dtype=int)}, ) monkeypatch.setattr( - biological_module, + biological_tools, "read_stored_selection_indices", lambda *_args, **_kwargs: np.asarray([], dtype=np.int64), ) @@ -1122,20 +1212,20 @@ def test_cluster_composition_metadata_alignment_edges( monkeypatch: pytest.MonkeyPatch, ) -> None: store = FakeStore() - original_rows = biological_module.read_metadata_rows - original_missing = biological_module.read_metadata_missing_rows + original_rows = biological_tools.read_metadata_rows + original_missing = biological_tools.read_metadata_missing_rows monkeypatch.setattr( - biological_module, + biological_tools, "read_metadata_rows", lambda *_args, **_kwargs: np.asarray(["control"]), ) with pytest.raises(ValueError, match="condition and cluster"): asyncio.run(inspect_cluster_composition(context(store))) - monkeypatch.setattr(biological_module, "read_metadata_rows", original_rows) + monkeypatch.setattr(biological_tools, "read_metadata_rows", original_rows) monkeypatch.setattr( - biological_module, + biological_tools, "read_metadata_missing_rows", lambda *_args, **_kwargs: np.asarray( [True] + [False] * (len(store.cells.cluster_values) - 1) @@ -1145,7 +1235,7 @@ def test_cluster_composition_metadata_alignment_edges( asyncio.run(inspect_cluster_composition(context(store))) monkeypatch.setattr( - biological_module, "read_metadata_missing_rows", original_missing + biological_tools, "read_metadata_missing_rows", original_missing ) def sample_misaligned( @@ -1154,18 +1244,18 @@ def sample_misaligned( values = original_rows(table, column, indices) return values[:-1] if column == "sample" else values - monkeypatch.setattr(biological_module, "read_metadata_rows", sample_misaligned) + monkeypatch.setattr(biological_tools, "read_metadata_rows", sample_misaligned) with pytest.raises(ValueError, match="sample and cluster"): asyncio.run(inspect_cluster_composition(context(store))) - monkeypatch.setattr(biological_module, "read_metadata_rows", original_rows) + monkeypatch.setattr(biological_tools, "read_metadata_rows", original_rows) def sample_missing(_table: object, column: str, indices: np.ndarray) -> np.ndarray: if column == "sample": return np.asarray([True] + [False] * (len(indices) - 1)) return np.zeros(len(indices), dtype=bool) - monkeypatch.setattr(biological_module, "read_metadata_missing_rows", sample_missing) + monkeypatch.setattr(biological_tools, "read_metadata_missing_rows", sample_missing) with pytest.raises(ValueError, match="sample column contains missing"): asyncio.run(inspect_cluster_composition(context(store))) @@ -1283,16 +1373,16 @@ def report(value: TreatmentObservation) -> BiologicalInterpretationReport: without_condition = deps.model_copy(update={"conditionColumn": None}) with pytest.raises(ModelRetry, match="condition column"): - biological_module._canonicalize_treatment_observations( + biological_validation._canonicalize_treatment_observations( report(observation), without_condition ) with pytest.raises(ModelRetry, match="remain descriptive"): - biological_module._canonicalize_treatment_observations( + biological_validation._canonicalize_treatment_observations( report(observation.model_copy(update={"isDescriptiveOnly": False})), deps, ) with pytest.raises(ModelRetry, match="exactly two distinct"): - biological_module._canonicalize_treatment_observations( + biological_validation._canonicalize_treatment_observations( report( observation.model_copy( update={"evidenceIds": [control.evidenceId, control.evidenceId]} @@ -1301,7 +1391,7 @@ def report(value: TreatmentObservation) -> BiologicalInterpretationReport: deps, ) with pytest.raises(ModelRetry, match="condition composition evidence"): - biological_module._canonicalize_treatment_observations( + biological_validation._canonicalize_treatment_observations( report( observation.model_copy( update={"evidenceIds": [control.evidenceId, "unknown"]} @@ -1310,18 +1400,18 @@ def report(value: TreatmentObservation) -> BiologicalInterpretationReport: deps, ) with pytest.raises(ModelRetry, match="distinct named conditions"): - biological_module._canonicalize_treatment_observations( + biological_validation._canonicalize_treatment_observations( report(observation.model_copy(update={"comparisonCondition": "control"})), deps, ) - canonical = biological_module._canonicalize_treatment_observations( + canonical = biological_validation._canonicalize_treatment_observations( report(observation), deps ) assert "equal mean independent-unit fractions" in canonical[0].observation higher = treated.model_copy(update={"meanFraction": 0.75}) deps.conditionEvidence[treated.evidenceId] = higher - canonical = biological_module._canonicalize_treatment_observations( + canonical = biological_validation._canonicalize_treatment_observations( report(observation.model_copy(update={"direction": "higher"})), deps ) assert "higher mean independent-unit fraction" in canonical[0].observation diff --git a/tests/test_agent_characterize_covariates.py b/tests/test_agent_characterize_covariates.py index fd7b1369..0afed291 100644 --- a/tests/test_agent_characterize_covariates.py +++ b/tests/test_agent_characterize_covariates.py @@ -11,8 +11,9 @@ from pydantic_ai.models.function import AgentInfo, FunctionModel from scipy.sparse import csr_matrix -from scarf.agent import CovariateCharacterization, characterize_covariates -from scarf.agent.characterize_covariates import ( +from scarf.agent.experimental_context.contracts import CovariateCharacterization +from scarf.agent.experimental_context.characterization import characterize_covariates +from scarf.agent.experimental_context.characterization import ( _Run, _assign_domain, _characterize_coefficient, @@ -25,7 +26,7 @@ _triage_columns, _validate_directions, ) -from scarf.agent.decide import DecisionValidationError +from scarf.agent.decisions.selection import DecisionValidationError from scarf.agent.types import ArtifactReferenceModel, Decision, EvidenceItem from scarf.datastore.datastore import DataStore from scarf.storage import ArtifactRef, ArtifactResolutionError @@ -528,7 +529,7 @@ def test_characterize_covariates_avoids_bulk_metadata_loads( store = _store_with_design(tmp_path) cell_selection = store.snapshot_cell_selection("I") characterize_covariates_module = import_module( - "scarf.agent.characterize_covariates" + "scarf.agent.experimental_context.characterization" ) original_fetch = store.cells.fetch original_fetch_all = store.cells.fetch_all @@ -630,7 +631,7 @@ def test_covariate_direction_validation_reports_structural_errors( def test_run_ask_audits_invalid_mocked_decision(monkeypatch) -> None: characterize_covariates_module = import_module( - "scarf.agent.characterize_covariates" + "scarf.agent.experimental_context.characterization" ) run = _Run( store=object(), diff --git a/tests/test_agent_characterize_features.py b/tests/test_agent_characterize_features.py index c564733a..116c58e8 100644 --- a/tests/test_agent_characterize_features.py +++ b/tests/test_agent_characterize_features.py @@ -7,8 +7,11 @@ import numpy as np from scipy.sparse import csr_matrix -from scarf.agent import FeatureCharacterization, characterize_features -from scarf.agent.characterize_features import ( +from scarf.agent.data_enrichment.characterization import ( + FeatureCharacterization, + characterize_features, +) +from scarf.agent.data_enrichment.characterization import ( _assist_species, _load_or_fetch_reference, _sex_coefficient_note, @@ -333,7 +336,9 @@ def test_reference_loading_audits_mocked_download_success_and_failure( tmp_path: Path, monkeypatch, ) -> None: - characterize_features_module = import_module("scarf.agent.characterize_features") + characterize_features_module = import_module( + "scarf.agent.data_enrichment.characterization" + ) reference = GeneReference( species="homo_sapiens", release="test", diff --git a/tests/test_agent_comparison_boundaries.py b/tests/test_agent_comparison_boundaries.py new file mode 100644 index 00000000..c43c7a36 --- /dev/null +++ b/tests/test_agent_comparison_boundaries.py @@ -0,0 +1,176 @@ +"""Corrupt, incomplete and mismatched measurements cannot authorize acceptance.""" + +from copy import deepcopy +from typing import Any + +import pytest + +from scarf.agent.orchestrator.rna_tuning import validate_completed_comparison_evidence +from scarf.agent.parameter_tuning.comparisons import validate_comparison_review +from tests.agent_comparison_examples import comparison_review + + +@pytest.mark.parametrize( + ("field", "value", "reason"), + [ + ("phase", "unmeasured", "Unknown RNA comparison phase"), + ("population", "unknown", "population must identify"), + ("candidateSettings", [], "exact settings and rows"), + ("comparisons", {}, "exact settings and rows"), + ("resolutionCandidateIds", [], "four resolutions"), + ], +) +def test_review_rejects_incomplete_coverage( + field: str, value: Any, reason: str +) -> None: + review = comparison_review() + coverage = review["comparisonCoverage"] + coverage[field] = value + with pytest.raises(ValueError, match=reason): + validate_comparison_review(coverage, review) + + +@pytest.mark.parametrize( + ("failure", "reason"), + [ + ("duplicate", "IDs must be unique"), + ("missingBaseline", "baseline is not completed"), + ("unfinishedBaseline", "baseline is not completed"), + ("missingAlternative", "lacks a completed alternative"), + ("unfinishedAlternative", "complete on the exact same cells"), + ("differentCells", "complete on the exact same cells"), + ("missingReason", "observed eligibility reason"), + ("inventedAlternative", "cannot claim a different"), + ("unsupportedEquivalence", "valid observed equivalence"), + ("missingConclusion", "explicit comparison conclusions"), + ("unaddressedAlternative", "all its observed alternatives"), + ], +) +def test_each_comparison_requires_complete_attributable_evidence( + failure: str, reason: str +) -> None: + review = comparison_review() + coverage = review["comparisonCoverage"] + rows, settings = coverage["comparisons"], coverage["candidateSettings"] + if failure == "duplicate": + rows.append(deepcopy(rows[0])) + elif failure == "missingBaseline": + rows[0]["baselineCandidateId"] = "absent" + elif failure == "unfinishedBaseline": + settings["baseline"]["status"] = "failed" + elif failure == "missingAlternative": + rows[0]["alternativeCandidateId"] = "absent" + elif failure == "unfinishedAlternative": + settings["resolution-half"]["status"] = "failed" + elif failure == "differentCells": + settings["resolution-half"]["cellSelection"]["artifactId"] = "d" * 64 + elif failure == "missingReason": + rows[-1]["reason"] = " " + elif failure == "inventedAlternative": + rows[-1]["alternativeCandidateId"] = "genes-two" + elif failure == "unsupportedEquivalence": + rows[-1]["observedProof"]["meaningfulPermittedInterventions"] = 1 + elif failure == "missingConclusion": + review.pop("comparisonConclusions") + elif failure == "unaddressedAlternative": + review["comparisonConclusions"][0]["candidateIds"] = ["baseline"] + with pytest.raises(ValueError, match=reason): + validate_comparison_review(coverage, review) + + +@pytest.mark.parametrize( + ("failure", "reason"), + [ + ("absent", "candidate is unavailable"), + ("failed", "completed on the same cells"), + ("cells", "completed on the same cells"), + ("representation", "changed the combined representation"), + ("resolution", "coverage is incomplete"), + ("missingSelected", "lack exact completed evidence"), + ("missingClusterMarkers", "every selected cluster"), + ("inventedConcern", "must cite supplied candidate evidence"), + ("unvalidatedRepair", "one targeted full repair"), + ], +) +def test_final_panel_requires_exact_selected_combination( + failure: str, reason: str +) -> None: + review = comparison_review() + coverage = review["comparisonCoverage"] + settings = coverage["candidateSettings"] + replacement = deepcopy(settings["resolution-half"]) + settings["final-half"] = replacement + coverage["resolutionCandidateIds"][0] = "final-half" + if failure == "absent": + del settings["final-half"] + elif failure == "failed": + replacement["status"] = "failed" + elif failure == "cells": + replacement["cellSelection"]["artifactId"] = "a" * 64 + elif failure == "representation": + replacement["parameters"]["dimensions"] = 30 + elif failure == "resolution": + replacement["parameters"]["leidenResolution"] = 0.9 + elif failure == "missingSelected": + review["selectedCandidateId"] = "missing" + elif failure == "missingClusterMarkers": + settings["candidate-two"]["metrics"]["nClusters"] = 3 + elif failure == "inventedConcern": + review["populationConcerns"] = [ + { + "candidateId": "candidate-two", + "clusterId": "0", + "status": "nonEssentialLimitation", + "evidenceIds": ["invented-marker-proof"], + "explanation": "A proposed interpretation needs observed support.", + } + ] + elif failure == "unvalidatedRepair": + repaired = deepcopy(settings["baseline"]) + repaired["parameters"]["dimensions"] = 30 + repaired["parameters"]["neighborsK"] = 41 + settings["unvalidated"] = repaired + review["selectedCandidateId"] = "unvalidated" + with pytest.raises(ValueError, match=reason): + validate_comparison_review(coverage, review) + + +@pytest.mark.parametrize("failure", ["coverage", "candidate", "measurement"]) +def test_report_cannot_detach_accepted_review_from_saved_measurements( + failure: str, +) -> None: + review = comparison_review() + if failure == "coverage": + review["comparisonCoverage"] = None + reason = "mandatory comparison evidence" + elif failure == "candidate": + review["candidates"][0]["cellSelection"] = None + reason = "exact reviewed candidate" + else: + review["candidates"][0] = deepcopy(review["candidates"][0]) + review["candidates"][0]["metrics"]["seedStability"] = 1.0 + reason = "measurements differ" + with pytest.raises(ValueError, match=reason): + validate_completed_comparison_evidence(review) + + +def test_duplicate_or_fabricated_tradeoffs_cannot_justify_preference() -> None: + review = comparison_review() + coverage = review["comparisonCoverage"] + coverage["candidateSettings"]["genes-two"]["metrics"]["seedStability"] = 0.99 + conclusion = next( + c for c in review["comparisonConclusions"] if c["axis"] == "hvgCount" + ) + tradeoff = { + "alternativeCandidateId": "genes-two", + "metric": "seedStability", + "preferredValue": 0.92, + "alternativeValue": 0.99, + "interpretation": "The measured stability advantage is weighed against marker preservation.", + } + conclusion["tradeoffs"] = [tradeoff, deepcopy(tradeoff)] + with pytest.raises(ValueError, match="duplicate"): + validate_comparison_review(coverage, review) + conclusion["tradeoffs"] = [{**tradeoff, "alternativeValue": 1.0}] + with pytest.raises(ValueError, match="exact preferred and alternative"): + validate_comparison_review(coverage, review) diff --git a/tests/test_agent_context_capture_repair.py b/tests/test_agent_context_capture_repair.py new file mode 100644 index 00000000..3c7fca14 --- /dev/null +++ b/tests/test_agent_context_capture_repair.py @@ -0,0 +1,247 @@ +"""Capture provenance repair reuses measured designs without reopening comparisons.""" + +import asyncio +from copy import deepcopy +from types import SimpleNamespace + +import pytest +from pydantic_ai import ModelRetry +from pydantic_ai.tools import ToolDefinition + +from scarf.agent.experimental_context import tools +from scarf.agent.experimental_context.contracts import ( + CaptureProposal, + CellQcProfileEvidence, + CovariateComparison, + InferenceUnit, +) +from scarf.agent.types import BatchSafetyEvidence +from tests.test_agent_design_comparisons import _deps, _design, _proposal + + +def _repair_context(monkeypatch): + cells, characterization = _design() + characterization.columns.append( + {"name": "batch", "kind": "categorical", "domain": "technical"} + ) + deps = _deps(cells) + deps.studyContext = "sample identifies the physical capture." + deps.characterization = characterization + deps.designRounds = 2 + deps.toolCalls = [ + "inspect_cell_covariates", + "analyze_experimental_design", + "analyze_experimental_design", + ] + measured = CovariateComparison( + proposal=_proposal(), status="computed", evidenceId="measured-comparison" + ) + deps.comparisons = [measured] + characterization.comparisons = [measured] + deps.batchSafety = { + "measured-safety": BatchSafetyEvidence( + coefficient="response", + batchColumns=["batch"], + status="unsafe", + evidenceId="measured-safety", + ) + } + arguments = { + "column_domains": {"sample": "design", "batch": "technical"}, + "coefficients_of_interest": [], + "units_of_inference": {}, + "batch_columns": ["batch"], + "proposals": [], + "capture_proposal": CaptureProposal( + column="sample", provenanceQuote=deps.studyContext + ), + } + deps.characterizationInputs = { + "directions": { + "columnDomains": dict(arguments["column_domains"]), + "coefficientsOfInterest": [], + "unitsOfInference": {}, + }, + "metadata": {"sample": "unchanged"}, + } + saved = {} + + def write(key, value): + assert key not in saved, "Committed measurements must remain immutable" + saved[key] = deepcopy(value) + + deps.checkpointWrite = write + deps.checkpointRead = saved.get + tools.persist_context_evidence(deps, "design2") + qc_calls = [] + + def qc(current, design): + qc_calls.append(current.captureProposal) + assert design is current.characterization + profile = CellQcProfileEvidence( + profileId="capture-supported", + action="sampleMad", + sampleColumn="sample", + attributes=["counts"], + ) + current.qcProfiles = {profile.profileId: profile} + return [profile] + + monkeypatch.setattr(tools, "_offered_qc_profiles", qc) + for name in ( + "characterize_context", + "evaluate_proposals", + "_batch_safety_evidence", + ): + monkeypatch.setattr( + tools, + name, + lambda *a, **k: pytest.fail("Capture repair must reuse measured design"), + ) + return SimpleNamespace(deps=deps), arguments, saved, qc_calls + + +def test_capture_repair_preserves_design_and_immutable_checkpoints(monkeypatch): + context, arguments, saved, qc_calls = _repair_context(monkeypatch) + before = deepcopy(saved["design2"]) + result = asyncio.run(tools.analyze_experimental_design(context, **arguments)) + assert context.deps.designRounds == 2 + assert len(context.deps.comparisons) == 1 + assert result.characterization.captureProvenance == arguments["capture_proposal"] + assert result.batchSafety[0].status == "unsafe" + assert result.batchSafety[0].batchColumns == ["batch"] + assert len(qc_calls) == 1 + assert set(saved) == {"design2", "capture"} + assert saved["design2"] == before + assert ( + saved["capture"]["characterizationInputs"] == before["characterizationInputs"] + ) + restored = _deps(context.deps.cells) + restored.checkpointRead = saved.get + assert tools.restore_context_evidence(restored) + assert restored.captureProposal == arguments["capture_proposal"] + assert restored.designRounds == 2 + assert restored.comparisons == context.deps.comparisons + assert list(restored.qcProfiles) == ["capture-supported"] + assert len(qc_calls) == 1 + definition = ToolDefinition(name="analyze_experimental_design") + assert ( + tools._prepare_experimental_context_tool( + SimpleNamespace(deps=restored), definition + ) + is None + ) + + +@pytest.mark.parametrize( + "change,reason", + [ + ({"capture_proposal": None}, "two evidence rounds"), + ({"proposals": [_proposal()]}, "two evidence rounds"), + ({"column_domains": {"sample": "technical"}}, "exact saved domains"), + ({"coefficients_of_interest": ["response"]}, "exact saved domains"), + ( + { + "units_of_inference": { + "response": InferenceUnit(observationUnit="sample") + } + }, + "exact saved domains", + ), + ({"batch_columns": []}, "unchanged assessed batch"), + ( + { + "capture_proposal": CaptureProposal( + column="sample", provenanceQuote="sample is an inferred capture" + ) + }, + "exact study quote", + ), + ], +) +def test_capture_repair_rejects_new_work_or_unverified_provenance( + monkeypatch, change, reason +): + context, arguments, saved, qc_calls = _repair_context(monkeypatch) + arguments.update(change) + with pytest.raises(ModelRetry, match=reason): + asyncio.run(tools.analyze_experimental_design(context, **arguments)) + assert context.deps.designRounds == 2 + assert context.deps.captureProposal is None + assert set(saved) == {"design2"} + assert not qc_calls + + +def test_capture_repair_rejects_ambiguous_historical_batch_designs(monkeypatch): + context, arguments, saved, qc_calls = _repair_context(monkeypatch) + context.deps.batchSafety["alternative"] = BatchSafetyEvidence( + coefficient="response", batchColumns=["other"], status="safe" + ) + with pytest.raises(ModelRetry, match="Alternative designs cannot be combined"): + asyncio.run(tools.analyze_experimental_design(context, **arguments)) + assert not qc_calls + assert set(saved) == {"design2"} + + +def test_failed_capture_qc_refresh_does_not_authorize_incomplete_evidence(monkeypatch): + context, arguments, saved, _ = _repair_context(monkeypatch) + + def failed(*args): + raise RuntimeError("Capture quality measurements unavailable") + + monkeypatch.setattr(tools, "_offered_qc_profiles", failed) + with pytest.raises(RuntimeError, match="quality measurements unavailable"): + asyncio.run(tools.analyze_experimental_design(context, **arguments)) + assert context.deps.captureProposal is None + assert context.deps.characterization.captureProvenance is None + assert context.deps.designRounds == 2 + assert set(saved) == {"design2"} + + +def test_model_summary_exposes_provenance_repair_after_exhausted_design_rounds( + monkeypatch, +): + context, arguments, _, _ = _repair_context(monkeypatch) + definition = ToolDefinition(name="analyze_experimental_design") + assert tools._prepare_experimental_context_tool(context, definition) is definition + from scarf.agent.experimental_context.contracts import CovariateEvidence + + async def measured(ctx): + return CovariateEvidence(characterization=ctx.deps.characterization) + + before = asyncio.run(tools.model_evidence_tool(measured)(context)) + assert before["captureRepairAvailable"] is True + assert "proposals=[]" in before["captureRepairInstructions"] + assert before["captureRepairInputs"] == { + key: value for key, value in arguments.items() if key != "capture_proposal" + } + asyncio.run(tools.analyze_experimental_design(context, **arguments)) + after = asyncio.run(tools.model_evidence_tool(measured)(context)) + assert after["captureRepairAvailable"] is False + assert "captureRepairInstructions" not in after + + +@pytest.mark.parametrize( + "quote,purpose", + [ + ("Assess combined covariates.", "designCoverage"), + ("Assess joint association of covariates.", "association"), + ("Estimate joint effects of covariates.", "effectEstimation"), + ], +) +def test_requested_question_summary_preserves_its_required_purpose( + monkeypatch, quote, purpose +): + from scarf.agent.experimental_context.contracts import CovariateEvidence + + context, _, _, _ = _repair_context(monkeypatch) + context.deps.studyObjective = quote + + async def measured(ctx): + return CovariateEvidence(characterization=ctx.deps.characterization) + + payload = asyncio.run(tools.model_evidence_tool(measured)(context)) + question = payload["requestedComparisons"][0] + assert question["purpose"] == purpose + assert question["essential"] is True + assert question["completionEvidence"] == "evidenceRequirements and evidenceCoverage" diff --git a/tests/test_agent_context_computation_reuse.py b/tests/test_agent_context_computation_reuse.py new file mode 100644 index 00000000..6d80474b --- /dev/null +++ b/tests/test_agent_context_computation_reuse.py @@ -0,0 +1,310 @@ +"""Scientific metadata reuse keeps typed inputs and QC policies attributable.""" + +from collections import Counter +from copy import deepcopy +from typing import Any + +import numpy as np +import pandas as pd +import pytest + +from scarf.agent.experimental_context import characterization as characterization_module +from scarf.agent.experimental_context import qc_evidence +from scarf.agent.experimental_context.contracts import CovariateCharacterization +from tests.test_agent_experimental_context import _Store, _context, _replace_store_cells + + +def _capture_context() -> tuple[Any, CovariateCharacterization]: + store = _Store() + donors = np.repeat(["d1", "d2", "d3"], 8) + condition = np.tile(np.repeat(["case", "control"], 4), 3) + values = { + "I": np.ones(24, dtype=bool), + "ids": np.asarray([f"cell{i}" for i in range(24)]), + "names": np.asarray([f"cell{i}" for i in range(24)]), + "capture": donors.copy(), + "donor": donors, + "condition": condition, + "sample": np.asarray( + [f"{d}:{c}" for d, c in zip(donors, condition, strict=True)] + ), + "RNA_nCounts": np.tile(np.arange(8, dtype=float) + 100, 3), + "RNA_nFeatures": np.tile(np.arange(8, dtype=float) + 50, 3), + } + _replace_store_cells(store, values) + deps = _context(store, directions={"physicalCaptureColumn": "capture"}).deps + characterized = CovariateCharacterization( + status="done", + columns=[ + {"name": name, "kind": "categorical"} + for name in ("condition", "sample", "donor") + ], + coefficients=[ + { + "name": "condition", + "scope": "betweenUnit", + "observationUnit": "sample", + "independentUnit": "donor", + "pairedCoverage": {"design": "paired"}, + } + ], + ) + return deps, characterized + + +def test_capture_design_is_computed_once_across_real_policy_projections( + monkeypatch: pytest.MonkeyPatch, +) -> None: + deps, characterized = _capture_context() + captures = deps.cells.fetch("capture") + expected = { + label: qc_evidence._capture_design_safety(deps, characterized, captures, label) + for label in ("d1", "d2", "d3") + } + calls: Counter[str] = Counter() + reads: Counter[str] = Counter() + compute = qc_evidence._compute_capture_design_safety + fetch = qc_evidence._QcDesignData.fetch + + def counted(*args: Any) -> Any: + calls[args[-1]] += 1 + return compute(*args) + + def fetched(self: Any, column: str) -> Any: + if column not in self.values: + reads[column] += 1 + return fetch(self, column) + + monkeypatch.setattr(qc_evidence, "_compute_capture_design_safety", counted) + monkeypatch.setattr(qc_evidence._QcDesignData, "fetch", fetched) + profiles = qc_evidence._offered_qc_profiles(deps, characterized) + assert sum(bool(profile.captureFailureEvidence) for profile in profiles) >= 4 + assert calls == {"d1": 1, "d2": 1, "d3": 1} + assert reads == {"condition": 1, "sample": 1, "donor": 1} + assert deps.qcDesignData is None + for profile in profiles: + for failure in profile.captureFailureEvidence: + rows, condition_safe, unit_safe = expected[failure.capture] + assert failure.conditionAndUnitSafety == rows + assert failure.preservesConditionCoverage == condition_safe + assert failure.preservesIndependentUnitCoverage == unit_safe + # A later projection sees changed metadata instead of stale cached safety. + deps.store.cells._values["condition"][:8] = "only_here" + updated = qc_evidence._offered_qc_profiles(deps, characterized) + assert calls == {"d1": 2, "d2": 2, "d3": 2} + assert all( + not failure.preservesConditionCoverage + for profile in updated + for failure in profile.captureFailureEvidence + if failure.capture == "d1" + ) + + +@pytest.mark.parametrize("kind", ["continuous", "unknown"]) +def test_continuous_or_unknown_capture_protection_cannot_become_categorical( + kind: str, +) -> None: + deps, characterized = _capture_context() + deps.store.cells._values["condition"] = np.linspace(20, 80, 24) + deps.store.cells._values["condition"][0] = np.nan + characterized.columns[0]["kind"] = kind + rows, preserves_conditions, preserves_units = qc_evidence._capture_design_safety( + deps, characterized, deps.cells.fetch("capture"), "d1" + ) + assert not preserves_conditions and not preserves_units + assert "requiredGroups" not in rows[0] + if kind == "continuous": + assert rows[0]["matchedRowsBeforeExclusion"] == 23 + assert rows[0]["matchedRowsAfterExclusion"] == 16 + assert rows[0]["independentUnitsAfterExclusion"] == 2 + assert len(rows[0]["quantilesAfterExclusion"]) == 5 + assert "not been established" in rows[0]["reason"] + else: + assert rows[0]["reason"] == "unknownCovariateKind" + + +def test_column_inventory_reuse_checks_values_kinds_missingness_and_selection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _Store() + inventory: dict[str, Any] = {} + calls: Counter[str] = Counter() + digest = characterization_module.column_partition_digest + + def counted(cells: Any, column: str, **kwargs: Any) -> Any: + calls[column] += 1 + return digest(cells, column, **kwargs) + + monkeypatch.setattr(characterization_module, "column_partition_digest", counted) + + def characterize(**directions: Any) -> Any: + return characterization_module.characterize_covariates( + store, + cellSelection=store.cell_selection, + directions=directions, + inventory=inventory, + ) + + first = characterize() + baseline = calls.copy() + assert first.status == "done" + characterize(columnDomains={"sequencing_depth": "technical"}) + assert calls == baseline + store.cells._values["sequencing_depth"][0] += 1 + characterize() + assert calls == baseline + Counter({"sequencing_depth": 1}) + characterize(columnKinds={"sequencing_depth": "categorical"}) + assert calls == baseline + Counter({"sequencing_depth": 2}) + store.cells._values["sequencing_depth"][0] = np.nan + characterize(columnKinds={"sequencing_depth": "categorical"}) + assert calls == baseline + Counter({"sequencing_depth": 3}) + store.cells._values["I"][0] = False + store.refresh_cell_selection() + before = calls.copy() + characterize() + assert all(calls[column] > before[column] for column in baseline) + + +def test_qc_design_scope_is_cleared_after_projection_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + deps, characterized = _capture_context() + + def fail(*args: Any, **kwargs: Any) -> Any: + assert isinstance(deps.qcDesignData, qc_evidence._QcDesignData) + raise ValueError("invalid quality metric") + + monkeypatch.setattr(qc_evidence, "_project_qc_profiles", fail) + with pytest.raises(ValueError, match="invalid quality metric"): + qc_evidence._offered_qc_profiles(deps, characterized) + assert deps.qcDesignData is None + + +def test_nullable_design_metadata_retains_missingness_as_unresolved_protection() -> ( + None +): + deps, characterized = _capture_context() + deps.store.cells._values["condition"] = np.linspace(20, 80, 24).astype(object) + deps.store.cells._values["condition"][0] = pd.NA + characterized.columns[0]["kind"] = "continuous" + rows, protected, _ = qc_evidence._capture_design_safety( + deps, characterized, deps.cells.fetch("capture"), "d3" + ) + assert not protected and rows[0]["missingRowsAfterExclusion"] == 1 + deps.store.cells._values["donor"] = deps.store.cells._values["donor"].astype(object) + deps.store.cells._values["donor"][0] = None + retention = qc_evidence._design_retention( + deps, characterized, np.ones(24, dtype=bool), np.ones(24, dtype=bool) + ) + assert "donor:missingValues" in retention["unsafeRetentionGroups"] + assert "None" not in retention["retainedCellsByColumn"]["donor"] + + +@pytest.mark.parametrize("missing", [pd.NaT, b"", b" ", b"\xff"]) +def test_unusable_imported_unit_labels_remain_missing_in_retention( + missing: Any, +) -> None: + deps, characterized = _capture_context() + deps.store.cells._values["donor"] = deps.store.cells._values["donor"].astype(object) + deps.store.cells._values["donor"][0] = missing + retention = qc_evidence._design_retention( + deps, characterized, np.ones(24, dtype=bool), np.ones(24, dtype=bool) + ) + assert "donor:missingValues" in retention["unsafeRetentionGroups"] + assert retention["retainedCellsByColumn"]["donor"] == {"d1": 7, "d2": 8, "d3": 8} + + +def _joint_capture_context() -> tuple[Any, CovariateCharacterization]: + deps, characterized = _capture_context() + deps.store.cells._values["condition"][:8] = "case" + deps.store.cells._values["time"] = np.concatenate( + [np.repeat("late", 8), np.tile(np.repeat(["early", "late"], 4), 2)] + ) + characterized.columns.append({"name": "time", "kind": "categorical"}) + time = deepcopy(characterized.coefficients[0]) + time["name"] = "time" + characterized.coefficients.append(time) + deps.protectedCombinations = [["condition", "time"]] + return deps, characterized + + +def test_joint_population_loss_is_visible_despite_preserved_marginal_conditions() -> ( + None +): + deps, characterized = _joint_capture_context() + captures = deps.cells.fetch("capture") + rows, condition_safe, unit_safe = qc_evidence._capture_design_safety( + deps, characterized, captures, "d1" + ) + assert all(row["preservesConditionCoverage"] for row in rows[:-1]) + assert all(row["preservesIndependentUnitCoverage"] for row in rows[:-1]) + assert not rows[-1]["preservesConditionCoverage"] + assert not condition_safe and not unit_safe + retention = qc_evidence._design_retention( + deps, characterized, np.ones(24, dtype=bool), captures != "d1" + ) + for column in ("condition", "time"): + assert all(retention["retainedCellsByColumn"][column].values()) + assert any( + value.startswith("combination:") for value in retention["unsafeRetentionGroups"] + ) + deps.qcDesignData = qc_evidence._QcDesignData(deps.cells) + cached = qc_evidence._design_retention( + deps, characterized, np.ones(24, dtype=bool), captures != "d1" + ) + assert cached == retention + + +@pytest.mark.parametrize("missing", ["value", "column", "independentColumn"]) +def test_missing_joint_design_inputs_cannot_establish_safe_capture_exclusion( + missing: str, +) -> None: + deps, characterized = _joint_capture_context() + if missing == "value": + deps.store.cells._values["time"] = deps.store.cells._values["time"].astype( + object + ) + deps.store.cells._values["time"][0] = None + elif missing == "column": + del deps.store.cells._values["time"] + else: + del deps.store.cells._values["donor"] + rows, condition_safe, unit_safe = qc_evidence._capture_design_safety( + deps, characterized, deps.cells.fetch("capture"), "d1" + ) + assert not unit_safe + if missing != "independentColumn": + assert not condition_safe + assert rows[-1]["reason"] == "missingProtectedCombination" + retention = qc_evidence._design_retention( + deps, characterized, np.ones(24, dtype=bool), np.ones(24, dtype=bool) + ) + assert any( + entry.startswith("combination:") and entry.endswith(":missingValues") + for entry in retention["unsafeRetentionGroups"] + ) + else: + assert all(row["reason"] == "missingDesignColumn" for row in rows[:-1]) + + +def test_joint_replication_uses_independent_donors_and_scoped_capture_identity() -> ( + None +): + deps, characterized = _capture_context() + deps.store.cells._values["time"] = np.tile(["early", "late"], 12) + characterized.columns.append({"name": "time", "kind": "categorical"}) + deps.protectedCombinations = [["condition", "time"]] + deps.qcDesignData = qc_evidence._QcDesignData(deps.cells) + captures = deps.cells.fetch("capture") + rows, condition_safe, unit_safe = qc_evidence._capture_design_safety( + deps, characterized, captures, "d1" + ) + assert condition_safe and unit_safe + assert rows[-1]["preservesIndependentUnitCoverage"] + changed_captures = captures.copy() + changed_captures[changed_captures == "d2"] = "d1" + rows, condition_safe, unit_safe = qc_evidence._capture_design_safety( + deps, characterized, changed_captures, "d1" + ) + assert condition_safe and not unit_safe + assert not rows[-1]["preservesIndependentUnitCoverage"] diff --git a/tests/test_agent_context_design_limits.py b/tests/test_agent_context_design_limits.py new file mode 100644 index 00000000..8431a760 --- /dev/null +++ b/tests/test_agent_context_design_limits.py @@ -0,0 +1,178 @@ +"""Repeated-unit aggregation and sparse design limits remain explicit evidence.""" + +import numpy as np +import pandas as pd +import pytest + +from scarf.agent.experimental_context.comparisons import ( + compare_covariates, + combination_labels, +) +from scarf.agent.experimental_context.contracts import ( + CovariateCharacterization, + CovariateProposal, +) +from tests.test_agent_design_comparisons import _Cells + + +def _comparison(frame, *, conditioned=False, joint=False): + kinds = { + name: "continuous" + if name == "response" and frame[name].dtype.kind == "f" + else "categorical" + for name in frame.columns + } + known = CovariateCharacterization( + status="done", + columns=[ + { + "name": name, + "kind": kinds[name], + "domain": "design" if name in {"sample", "donor"} else "biological", + } + for name in frame.columns + ], + coefficients=[ + { + "name": "response", + "observationUnit": "sample", + "independentUnit": "donor", + } + ], + ) + proposal = CovariateProposal( + response="response", + explanatoryColumns=["treatment", "stratum"] if joint else ["treatment"], + conditionedOn="stratum" if conditioned else None, + observationUnit="sample", + independentUnit="donor", + rationale="Check the supported comparison across independent donors.", + ) + return compare_covariates( + _Cells(frame), known, proposal, selection_identity={"cells": "frozen"} + ) + + +def test_continuous_response_uses_observation_medians_then_independent_donors(): + frame = pd.DataFrame( + [ + { + "sample": f"s{i}-{j}", + "donor": f"d{i}", + "treatment": str(i % 2), + "response": float(i * 10 + j + c), + } + for i in range(8) + for j in range(2) + for c in range(3) + ] + ) + result = _comparison(frame) + assert result.status == "computed" + assert result.evidence["responseAggregation"] == "medianPerObservationUnit" + assert result.evidence["independentAggregation"] == "medianOfObservationMedians" + assert result.evidence["observationUnits"] == 16 + assert result.evidence["independentUnits"] == 8 + assert ( + result.evidence["descriptiveDesign"]["continuousSummaries"]["response"][ + "minimum" + ] + == 1.0 + ) + + +@pytest.mark.parametrize( + "damage,reason", + [ + ("missing", "noCompleteObservations"), + ("cellIdentity", "observationUnitIsCellIdentifier"), + ( + "inconsistentObservation", + "explanatoryColumnsMustBeConstantWithinObservationUnit", + ), + ("fewDonors", "fewerThanFourIndependentUnits"), + ], +) +def test_unsupported_designs_cannot_be_relabelled_negative_associations(damage, reason): + frame = pd.DataFrame( + [ + { + "sample": f"s{i}", + "donor": f"d{i}", + "treatment": str(i % 2), + "response": float(i + c), + } + for i in range(8) + for c in range(3) + ] + ) + if damage == "missing": + frame["response"] = np.nan + elif damage == "cellIdentity": + frame["sample"] = [str(i) for i in range(len(frame))] + elif damage == "inconsistentObservation": + frame.loc[0, "treatment"] = "different" + else: + frame = frame.iloc[:9] + result = _comparison(frame) + assert result.status == "unsupported" + assert reason in result.reasons + assert "singleAssociations" not in result.evidence + + +@pytest.mark.parametrize( + "case,reason", + [ + ("levels", "moreThanThirtyTwoCategoricalLevels"), + ("strata", "moreThanSixteenConditioningStrata"), + ("sparseStrata", "unsupportedConditioningStrata"), + ("joint", "moreThanThirtyTwoJointGroups"), + ], +) +def test_large_or_sparse_combinations_preserve_the_unsupported_reason(case, reason): + n = ( + 68 + if case == "strata" + else 49 + if case == "joint" + else 34 + if case == "levels" + else 8 + ) + rows = [ + { + "sample": f"s{i}", + "donor": f"d{i}", + "treatment": str( + i if case == "levels" else i % 7 if case == "joint" else i % 2 + ), + "stratum": str( + i // 7 if case == "joint" else i // 4 if case == "strata" else i // 2 + ), + "response": str(i % 2), + } + for i in range(n) + ] + frame = pd.DataFrame(rows).loc[np.repeat(np.arange(n), 3)].reset_index(drop=True) + result = _comparison( + frame, conditioned=case in {"strata", "sparseStrata"}, joint=case == "joint" + ) + assert result.status == "unsupported" + assert reason in result.reasons + assert result.evidence["independentUnits"] == n + if case == "sparseStrata": + assert all( + row["association"]["reason"] == "fewerThanFourIndependentUnits" + for row in result.evidence["strata"] + ) + + +def test_combination_identity_accepts_utf8_and_rejects_misaligned_columns(): + cells = _Cells(pd.DataFrame({"one": [b"alpha", b"beta"], "two": ["yes", "no"]})) + labels = combination_labels(cells, ["one", "two"]) + assert '"alpha"' in labels[0] + with pytest.raises(ValueError, match="two distinct"): + combination_labels(cells, ["one", "one"]) + cells.fetch = lambda name: np.asarray(["a", "b"] if name == "one" else ["c"]) + with pytest.raises(ValueError, match="align"): + combination_labels(cells, ["one", "two"]) diff --git a/tests/test_agent_context_efficiency.py b/tests/test_agent_context_efficiency.py new file mode 100644 index 00000000..5abae578 --- /dev/null +++ b/tests/test_agent_context_efficiency.py @@ -0,0 +1,376 @@ +"""Context retries retain measured evidence and cannot erase study questions.""" + +from types import SimpleNamespace + +import pytest +from pydantic_ai import ModelRetry + +from scarf.agent.experimental_context import tools, validation +from scarf.agent.experimental_context.contracts import ( + CaptureFailureEvidence, + CellQcPlan, + CellQcProfileEvidence, + CovariateCharacterization, + CovariateEvidence, + ExperimentalContextDecision, + ExperimentalContextDependencies, +) +from scarf.agent.experimental_context.requirements import ( + active_batch_safety, + objective_evidence, +) +from tests.test_agent_objective_requirements import _repeated_design + + +def test_provider_schema_omits_derived_fields_without_changing_serialization(): + decision = ExperimentalContextDecision() + before = decision.model_dump(mode="json") + schema = ExperimentalContextDecision.model_json_schema()["properties"] + derived = { + "cellQc", + "protectedCombinations", + "physicalCaptureColumn", + "pooledReferenceCaptures", + "unsupportedProtection", + } + assert not derived.intersection(schema) + assert derived <= before.keys() + assert ( + ExperimentalContextDecision.model_validate(before).model_dump(mode="json") + == before + ) + + +def test_illegal_qc_is_rejected_before_characterization(monkeypatch): + def forbidden(*args, **kwargs): + pytest.fail("An illegal field must be rejected before metadata work") + + monkeypatch.setattr(validation, "characterize_context", forbidden) + with pytest.raises(ModelRetry, match="cellQc blank"): + validation.validate_experimental_context( + ExperimentalContextDecision(cellQc=CellQcPlan(rationale="skip")), + ExperimentalContextDependencies(), + ) + + +def test_model_evidence_deduplicates_design_without_dropping_contradictions(): + safety = [ + { + "coefficient": "condition", + "preservesConditionCoverage": False, + "reason": "A protected group would disappear", + } + ] + failure = CaptureFailureEvidence(capture="library", conditionAndUnitSafety=safety) + profiles = [ + CellQcProfileEvidence(profileId=name, captureFailureEvidence=[failure]) + for name in ("global", "capture") + ] + full = CovariateEvidence(qcProfiles=profiles) + original = full.model_dump(mode="json") + compact = tools.compact_context_evidence(full) + assert len(compact["captureDesignSafety"]) == 1 + for profile in compact["qcProfiles"]: + ref = profile["captureFailureEvidence"][0]["designSafetyRef"] + assert compact["captureDesignSafety"][ref] == safety + assert full.model_dump(mode="json") == original + + +def test_context_details_restore_exact_capture_bounds_and_missingness(monkeypatch): + import asyncio + + failure = CaptureFailureEvidence( + capture="capture-a", + metricMissingFractions={"mitochondrial": 0.2, "counts": 0.0}, + conditionAndUnitSafety=[ + { + "coefficient": "condition", + "requiredGroups": ["rare", "common"], + "remainingGroups": ["common"], + "preservesConditionCoverage": False, + "preservesIndependentUnitCoverage": False, + } + ], + ) + profile = CellQcProfileEvidence( + profileId="capture-policy", + action="sampleMad", + attributes=["counts"], + sampleColumn="capture", + resolvedBounds=[ + {"group": "capture-a", "upperRemoval": 12.0}, + {"group": "capture-b", "upperRemoval": 15.0}, + ], + captureFailureEvidence=[failure], + ) + full = CovariateEvidence(qcProfiles=[profile]) + view = tools.compact_context_evidence(full) + capture = view["qcProfiles"][0]["captureFailureEvidence"][0] + assert capture["metricMissingness"] == { + "measuredSources": 2, + "nonzeroOrUnavailable": {"mitochondrial": 0.2}, + } + assert view["captureDesignSafety"][capture["designSafetyRef"]][0]["lostGroups"] == [ + "rare" + ] + deps = ExperimentalContextDependencies( + characterization=CovariateCharacterization(status="done"), + qcProfiles={profile.profileId: profile}, + ) + monkeypatch.setattr( + tools, + "characterize_covariates", + lambda *a, **k: pytest.fail("Details must not recompute"), + ) + detail = asyncio.run( + tools.inspect_context_evidence( + SimpleNamespace(deps=deps), "qcProfile", profile.profileId, "capture-a" + ) + ) + assert detail["capture"] == failure.model_dump(mode="json") + assert detail["resolvedBounds"] == [profile.resolvedBounds[0]] + with pytest.raises(ModelRetry, match="Choose one capture"): + asyncio.run( + tools.inspect_context_evidence( + SimpleNamespace(deps=deps), "qcProfile", profile.profileId, "unknown" + ) + ) + + +def test_compact_context_keeps_repeated_donor_failures_and_full_details(): + import asyncio + + record = { + "name": "condition", + "kind": "categorical", + "missingness": [{"missing": 3}], + "pairedCoverage": { + "complete": False, + "incompletePairs": 2, + "incompleteExamples": [{"pair": "donor-a"}, {"pair": "donor-b"}], + }, + } + characterization = CovariateCharacterization(status="done", coefficients=[record]) + view = tools.compact_context_evidence( + CovariateEvidence(characterization=characterization) + ) + coefficient = view["characterization"]["coefficients"][0] + assert coefficient["pairedCoverage"]["complete"] is False + assert coefficient["pairedCoverage"]["incompletePairs"] == 2 + assert coefficient["pairedCoverage"]["incompleteExamplesInSavedDetails"] == 2 + assert coefficient["missingness"] == [{"missing": 3}] + detail = asyncio.run( + tools.inspect_context_evidence( + SimpleNamespace( + deps=ExperimentalContextDependencies(characterization=characterization) + ), + "coefficient", + "condition", + ) + ) + assert detail == record + + +def test_completed_context_round_resumes_without_resetting_allowance(): + saved = {} + deps = ExperimentalContextDependencies( + characterization=CovariateCharacterization(status="done"), + characterizationInputs={"exact": "inputs"}, + designRounds=1, + toolCalls=["inspect_cell_covariates", "analyze_experimental_design"], + checkpointWrite=lambda key, value: saved.update({key: value}), + ) + tools.persist_context_evidence(deps, "design1") + resumed = ExperimentalContextDependencies(checkpointRead=saved.get) + assert tools.restore_context_evidence(resumed) + assert resumed.designRounds == 1 + assert resumed.characterizationInputs == deps.characterizationInputs + assert resumed.toolCalls == deps.toolCalls + assert resumed.characterization == deps.characterization + + +def test_explicit_joint_request_cannot_be_satisfied_by_marginal_proposal(): + result = _repeated_design() + objective = "Describe supported populations. Assess tissue and condition jointly" + requirements, coverage = objective_evidence( + study_context="The observations contain repeated donors and incomplete pairing.", + study_objective=objective, + experimental_result=result, + ) + missing = [ + item + for item in requirements + if item.requirementId.startswith("requestedDesign:") + ] + assert len(missing) == 1 + assert missing[0].columns == ["condition", "tissue"] + assert ( + next( + item for item in coverage if item.requirementId == missing[0].requirementId + ).status + == "unsupported" + ) + + +def test_batch_alternatives_do_not_create_untested_union_or_stale_license(): + first = SimpleNamespace(batchColumns=["library"], status="unsafe") + second = SimpleNamespace(batchColumns=["chemistry"], status="safe") + result = SimpleNamespace( + batchSafety=[first, second], decision=ExperimentalContextDecision() + ) + with pytest.raises(ValueError, match="exact assessed batch set"): + active_batch_safety(result) + result.decision.batchCorrection.batchColumns = ["chemistry"] + assert active_batch_safety(result) == [second] + result.decision.batchCorrection.batchColumns = ["library"] + assert active_batch_safety(result) == [first] + + +def test_exact_context_characterization_reuses_work_and_invalidates_changes( + monkeypatch, +): + from tests.test_agent_experimental_context import _Store, _context + + store = _Store() + deps = _context(store).deps + calls = [] + + def measured(*args, **kwargs): + calls.append(kwargs["directions"]) + return CovariateCharacterization(status="done") + + monkeypatch.setattr(tools, "characterize_covariates", measured) + first = tools.characterize_context(deps, {"columnDomains": {"batch": "technical"}}) + assert ( + tools.characterize_context(deps, {"columnDomains": {"batch": "technical"}}) + is first + ) + assert len(calls) == 1 + store.cells._values["batch"][0] = "b2" + tools.characterize_context(deps, {"columnDomains": {"batch": "technical"}}) + assert len(calls) == 2 + tools.characterize_context(deps, {"columnDomains": {"batch": "design"}}) + assert len(calls) == 3 + + +def test_context_identity_binds_added_columns_without_repeating_resume_scan( + monkeypatch, +): + from scarf.agent.orchestrator.context import _context_metadata_identity + from scarf.agent.parameter_tuning import execution + + values = {"capture": "first"} + measured = [] + + def fingerprint(metadata, column): + measured.append(column) + return values[column] + + monkeypatch.setattr(execution, "_metadata_column_fingerprint", fingerprint) + store = SimpleNamespace(cells=SimpleNamespace(columns=["original", "capture"])) + request = SimpleNamespace( + inputIdentity={"data": {"metadata": {"original": "validated"}}} + ) + first = _context_metadata_identity(store, request) + assert first == {"original": "validated", "capture": "first"} + assert measured == ["capture"] + values["capture"] = "changed" + assert _context_metadata_identity(store, request) != first + + +def test_generic_joint_request_also_requires_a_joint_proposal(): + result = _repeated_design() + requirements, coverage = objective_evidence( + study_context="The observations contain repeated donors and incomplete pairing.", + study_objective="Describe supported populations. Assess individual and combined covariates", + experimental_result=result, + ) + assert any( + item.requirementId.startswith("requestedDesign:") for item in requirements + ) + assert any( + item.requirementId.startswith("requestedDesign:") + and item.status == "unsupported" + for item in coverage + ) + + +@pytest.mark.parametrize( + "change", [None, "metadata", "features", "cohort", "noRevision"] +) +def test_tuning_context_revision_preserves_only_exact_numerical_inputs( + monkeypatch, change +): + from copy import deepcopy + from scarf.agent.orchestrator import journal + from scarf.agent.orchestrator.models import StageEvidenceReference + from scarf.agent.orchestrator.tuning import _tuning_revision_provenances + + request = SimpleNamespace(requestSha256="request", configSha256="config") + + def reference(name): + return StageEvidenceReference( + workflowRunId="workflow", + stage="experimental_context", + key=name, + contentSha256=name, + ) + + old_ref, new_ref = reference("old"), reference("new") + previous = { + "preprocessedAssays": [{"cells": "frozen", "features": "genes"}], + "featureMetadataFingerprints": {"names": "same"}, + "metadataFingerprints": {"condition": "same"}, + "studyContract": {"protectedCombinations": []}, + } + current = deepcopy(previous) + current["studyContract"] = {"protectedCombinations": [["condition", "tissue"]]} + current["metadataFingerprints"]["tissue"] = "newly measured" + old = SimpleNamespace( + attemptId="old", + status="done", + reportReferences=[old_ref], + inputs={}, + outputs={"studyContract": previous["studyContract"]}, + requestSha256="request", + configSha256="config", + ) + new = SimpleNamespace( + attemptId="new", + status="done", + reportReferences=[new_ref], + inputs={"reassessContextReport": old_ref.model_dump(mode="json")}, + outputs={"studyContract": current["studyContract"]}, + requestSha256="request", + configSha256="config", + ) + if change == "metadata": + current["metadataFingerprints"]["condition"] = "changed" + elif change == "features": + current["featureMetadataFingerprints"]["names"] = "changed" + elif change == "cohort": + current["preprocessedAssays"][0]["cells"] = "different" + elif change == "noRevision": + new.inputs = {} + monkeypatch.setattr(journal, "_stage_outcomes", lambda *args: [old, new]) + before = deepcopy(previous) + + def invoke(): + return _tuning_revision_provenances( + SimpleNamespace(zw=None), + "prefix", + "workflow", + request, + new_ref, + current, + [SimpleNamespace(inputs=previous)], + ) + + if change is not None: + with pytest.raises(ValueError, match="changed beyond"): + invoke() + else: + assert invoke() == [ + {**previous, "requestSha256": "request", "configSha256": "config"} + ] + assert previous == before diff --git a/tests/test_agent_context_failure_repair.py b/tests/test_agent_context_failure_repair.py new file mode 100644 index 00000000..5f0ac21c --- /dev/null +++ b/tests/test_agent_context_failure_repair.py @@ -0,0 +1,40 @@ +"""Unmeasured study questions are repaired before the design allowance closes.""" + +import asyncio + +import pytest +from pydantic_ai import ModelRetry + +from scarf.agent.experimental_context import tools, validation +from tests.test_agent_experimental_context import _Store, _context, _design_decision + + +def test_pending_joint_question_requests_followup_then_remains_unresolved(monkeypatch): + context = _context(_Store()) + context.deps.studyContext = "Assess individual and combined covariates." + context.deps.studyObjective = "Describe disease populations." + decision = _design_decision() + asyncio.run(tools.inspect_cell_covariates(context)) + arguments = { + "column_domains": decision.columnDomains, + "coefficients_of_interest": decision.coefficientsOfInterest, + "units_of_inference": decision.unitsOfInference, + "batch_columns": decision.batchCorrection.batchColumns, + } + asyncio.run(tools.analyze_experimental_design(context, **arguments)) + monkeypatch.setattr( + tools, + "characterize_covariates", + lambda *a, **k: pytest.fail( + "Output repair must reuse measured characterization" + ), + ) + with pytest.raises(ModelRetry, match="remaining analyze_experimental_design round"): + validation.validate_experimental_context(decision, context.deps) + assert context.deps.designRounds == 1 + # Repeating marginal evidence cannot answer the pending combined question. + asyncio.run(tools.analyze_experimental_design(context, **arguments)) + result = validation.validate_experimental_context(decision, context.deps) + assert context.deps.designRounds == 2 + assert any("combined covariates" in question for question in result.needsInput) + assert result.batchCorrection.action == "unsafe" diff --git a/tests/test_agent_context_journal_revision.py b/tests/test_agent_context_journal_revision.py new file mode 100644 index 00000000..d1e3524f --- /dev/null +++ b/tests/test_agent_context_journal_revision.py @@ -0,0 +1,208 @@ +"""Context evidence revisions append history and reject changed measured metadata.""" + +from types import SimpleNamespace + +import pytest + +from scarf.agent.experimental_context import agent as context_agent +from scarf.agent.experimental_context.contracts import ( + ExperimentalContextDecision, + ExperimentalContextResult, +) +from scarf.agent.experimental_context.characterization import characterize_covariates +from scarf.agent.experimental_context.study import build_study_contract +from scarf.agent.orchestrator import ( + AgentOrchestrator, + AutomatedWorkflowConfig, + AutomatedWorkflowRequest, +) +from scarf.agent.orchestrator import context, journal +from scarf.agent.orchestrator.models import ( + OrchestrationRequestRecord, + StageEvidenceReference, + WorkflowIdentity, +) +from scarf.agent.types import ArtifactReferenceModel +from scarf.datastore.datastore import DataStore +from tests.agent_orchestrator_store import create_store + + +@pytest.mark.parametrize("damage", [None, "currentMetadata", "historicalMetadata"]) +def test_missing_joint_question_appends_revision_without_replacing_completed_context( + tmp_path, monkeypatch, damage +): + path = create_store(tmp_path / "context-revision.zarr") + store = DataStore(str(path), default_assay="RNA", min_features_per_cell=0) + selected = store.snapshot_cell_selection("I") + selection = ArtifactReferenceModel.from_artifact_ref(selected) + known = characterize_covariates( + store, cellSelection=selected, model=None, directions={} + ) + report = ExperimentalContextResult( + status="done", + decision=ExperimentalContextDecision(), + characterization=known, + cellSelection=selection, + ) + contract = build_study_contract( + study_context="Known observed covariates.", + study_objective="Describe populations.", + experimental_result=report, + ) + request = OrchestrationRequestRecord( + workflowRunId="context-revision", + modelIdentity="test", + inputIdentity={}, + config=AutomatedWorkflowConfig(inputPolicy="unattended"), + request=AutomatedWorkflowRequest( + sourcePath=str(path), + zarrPath=str(path), + studyContext="Known observed covariates. Assess individual and combined covariates.", + studyObjective="Describe populations.", + ), + ) + workflow = WorkflowIdentity(request.workflowRunId) + prefix = journal._ensure_orchestration_store(store) + old_inputs = ( + {"metadataFingerprints": {"ids": "changed"}} + if damage == "currentMetadata" + else {} + ) + started = journal._start_attempt( + store.zw, + prefix, + workflow.workflowRunId, + "experimental_context", + request, + [], + inputs=old_inputs, + ) + _, reference = journal._save_stage_report( + store, started, report, expected_type=ExperimentalContextResult + ) + old = journal._complete_attempt( + started, + status="done", + report_references=[reference], + artifacts={"cellSelection": selection}, + outputs={"studyContract": contract.model_dump(mode="json")}, + ) + journal._save_outcome(store.zw, prefix, old) + original = journal.read_stage_evidence(store, reference) + if damage == "historicalMetadata": + journal._start_attempt( + store.zw, + prefix, + workflow.workflowRunId, + "parameter_tuning", + request, + [], + inputs={"metadataFingerprints": {"ids": "changed"}}, + ) + observed = [] + + class Review: + def __init__(self, *args, **kwargs): + pass + + def run(self, *args, **kwargs): + observed.append(kwargs["previous_context"]) + return report.model_copy( + update={ + "status": "needsInput", + "notes": ["Joint evidence remains unresolved"], + } + ) + + monkeypatch.setattr(context, "ExperimentalContextAgent", Review) + + def execute(): + return AgentOrchestrator(object()).experimental_context_stage( + store, + workflow, + request, + [], + selection, + StageEvidenceReference( + workflowRunId=workflow.workflowRunId, + stage="data_enrichment", + key="unused", + contentSha256="a" * 64, + ), + [], + [], + {}, + ) + + if damage: + with pytest.raises(ValueError, match="metadata"): + execute() + assert not observed + else: + revised, result = execute() + assert revised.status == "failed" + assert result.status == "needsInput" + assert revised.attemptId != old.attemptId + assert revised.inputs["reassessContextReport"] == reference.model_dump( + mode="json" + ) + assert ( + revised.inputs["requiredDesignQuestions"][0]["question"] + == "Assess individual and combined covariates" + ) + assert observed == [report] + assert journal.read_stage_evidence(store, reference) == original + + +@pytest.mark.parametrize( + "directions,error", + [ + ({"excludeColumns": "cell_type"}, "excludeColumns must be a list"), + ( + {"coefficientsOfInterest": ["cell_type"], "excludeColumns": ["cell_type"]}, + "forbids runtime use", + ), + ({"nested": {"protected": ["cell_type"]}}, "forbids runtime use"), + ], +) +def test_held_out_author_annotations_cannot_reenter_context_through_nested_directions( + tmp_path, monkeypatch, directions, error +): + import numpy as np + + path = create_store(tmp_path / "holdout.zarr") + store = DataStore(str(path), default_assay="RNA", min_features_per_cell=0) + store.cells.insert("cell_type", np.asarray(["A", "A", "B", "B"])) + request = OrchestrationRequestRecord( + workflowRunId="holdout", + modelIdentity="test", + inputIdentity={}, + request=AutomatedWorkflowRequest( + sourcePath=str(path), + zarrPath=str(path), + studyContext="Observed samples.", + studyObjective="Describe populations.", + experimentalDirections=directions, + ), + ) + monkeypatch.setattr( + context_agent, + "run_agent_sync", + lambda **k: pytest.fail( + "Held-out annotations must fail before model execution" + ), + ) + with pytest.raises(ValueError, match=error): + AgentOrchestrator(object()).experimental_context_stage( + store, + WorkflowIdentity("holdout"), + request, + [], + ArtifactReferenceModel.from_artifact_ref( + store.snapshot_cell_selection("I") + ), + SimpleNamespace(), + [], + [], + {}, + ) diff --git a/tests/test_agent_context_resume_boundaries.py b/tests/test_agent_context_resume_boundaries.py new file mode 100644 index 00000000..2c17e8b9 --- /dev/null +++ b/tests/test_agent_context_resume_boundaries.py @@ -0,0 +1,327 @@ +"""Committed context decisions and bounded details retain exact scientific inputs.""" + +import asyncio +import json +from types import SimpleNamespace + +import pytest +from pydantic import ValidationError +from pydantic_ai import ModelRetry, UnexpectedModelBehavior + +from scarf.agent.experimental_context import agent as context_agent +from scarf.agent.experimental_context import tools +from scarf.agent.experimental_context.contracts import ( + CellQcProfileEvidence, + CovariateCharacterization, + CovariateEvidence, + ExperimentalContextDependencies, + ExperimentalContextResult, + QcMetricSourceEvidence, + NamedArtifactSource, +) +from scarf.storage.refs import ArtifactRef +from scarf.agent.types import AgentRunInfo, ArtifactReferenceModel, ToolCallInfo +from tests.test_agent_design_comparisons import _design, _proposal +from tests.test_agent_experimental_context import _Store + + +def test_committed_context_result_replays_without_a_model_and_rejects_other_selection( + monkeypatch, +): + store = _Store() + selected = ArtifactReferenceModel.from_artifact_ref(store.cell_selection) + report = ExperimentalContextResult.get_blank().model_copy( + update={"cellSelection": selected} + ) + monkeypatch.setattr( + context_agent, "_derive_missing_percentage_artifacts", lambda *a, **k: [] + ) + monkeypatch.setattr( + context_agent, + "run_agent_sync", + lambda **k: pytest.fail("Committed result must not invoke a model"), + ) + agent = context_agent.ExperimentalContextAgent(object()) + result = agent.run( + store, + cell_selection=store.cell_selection, + checkpoint_read=lambda key: ( + {"report": report.model_dump(mode="json")} if key == "result" else None + ), + ) + assert result == report + report.cellSelection = selected.model_copy(update={"artifactId": "f" * 64}) + with pytest.raises(ValueError, match="different cell selection"): + agent.run( + store, + cell_selection=store.cell_selection, + checkpoint_read=lambda key: ( + {"report": report.model_dump(mode="json")} if key == "result" else None + ), + ) + + +@pytest.mark.parametrize("rounds", [1, 2]) +@pytest.mark.parametrize("previous_status", ["done", "needsInput"]) +@pytest.mark.parametrize( + "objective", + ["", "Discover populations while preserving condition."], + ids=["without-objective", "with-objective"], +) +def test_explicit_context_revision_preserves_rounds_and_complete_study_text( + monkeypatch, rounds, previous_status, objective +): + store = _Store() + previous = ExperimentalContextResult.get_blank().model_copy( + update={ + "status": previous_status, + "cellSelection": ArtifactReferenceModel.from_artifact_ref( + store.cell_selection + ), + "characterization": CovariateCharacterization( + status="done", columns=[{"name": "condition"}, {"name": "batch"}] + ), + "runInfo": AgentRunInfo( + agentName="context", + modelName="test", + toolCalls=[ + ToolCallInfo(toolName="analyze_experimental_design") + for _ in range(rounds) + ], + ), + } + ) + blocker = "Confirm whether the supplied batch identifies physical capture." + if previous_status == "needsInput": + previous.decision.needsInput = [blocker] + saved_previous = previous.model_dump(mode="json") + text = "Study provenance. " * 150 + "Assess condition and batch jointly." + monkeypatch.setattr( + context_agent, "_derive_missing_percentage_artifacts", lambda *a, **k: [] + ) + monkeypatch.setattr( + context_agent, + "characterize_covariates", + lambda *a, **k: pytest.fail("A revision must reuse measured characterization"), + ) + seen = [] + + def inspect(**kwargs): + deps = kwargs["deps"] + seen.append(deps.designRounds) + assert deps.characterization == previous.characterization + assert text in kwargs["user_prompt"] + assert "Committed evidence already measured" in kwargs["user_prompt"] + assert "Explicit requested questions" in kwargs["user_prompt"] + assert ( + f"Study objective: {objective or 'not provided'}" in kwargs["user_prompt"] + ) + if previous_status == "needsInput": + prompt = kwargs["user_prompt"] + assert blocker in prompt + assert "Keep genuinely essential unresolved questions" in prompt + evidence_json = prompt.split( + "Current objective requirements and their measured support: ", 1 + )[1].split("\nThe previous interpretation stopped", 1)[0] + evidence = json.loads(evidence_json) + if objective: + joint = next( + item + for item in evidence["evidenceRequirements"] + if item["question"] == "Assess condition and batch jointly" + ) + assert joint["columns"] == ["batch", "condition"] + assert joint["essential"] is True + measured = { + item["requirementId"]: item for item in evidence["evidenceCoverage"] + } + assert measured[joint["requirementId"]]["status"] == "unsupported" + assert "marginal comparisons do not answer" in " ".join( + measured[joint["requirementId"]]["reasons"] + ) + assert measured["studyDesign"]["status"] == "computed" + else: + assert evidence == {"evidenceRequirements": [], "evidenceCoverage": []} + raise UnexpectedModelBehavior("No additional provider attempt is available") + + monkeypatch.setattr(context_agent, "run_agent_sync", inspect) + result = context_agent.ExperimentalContextAgent(object()).run( + store, + cell_selection=store.cell_selection, + study_context=text, + study_objective=objective, + previous_context=previous, + ) + assert result.status == "failed" + assert seen == [rounds] + assert previous.model_dump(mode="json") == saved_previous + + +def test_saved_details_return_complete_inventory_and_policy_without_computation(): + cells, characterization = _design() + from scarf.agent.experimental_context.comparisons import compare_covariates + + measured = compare_covariates( + cells, characterization, _proposal(), selection_identity={} + ) + characterization.comparisons = [measured] + characterization.confounding = [ + {"coefficient": "response", "reason": "exact joint alias"} + ] + characterization.columns[0]["levelCounts"] = [ + {"value": str(i), "count": 2} for i in range(12) + ] + profile = CellQcProfileEvidence( + profileId="capture", + action="sampleMad", + sampleColumn="sample", + attributes=["counts"], + parameters={ + "resolvedBounds": [{"upper": 8}], + "captureComparisons": [{"unsafe": True}], + }, + resolvedBounds=[{"upper": 8}], + ) + deps = ExperimentalContextDependencies( + characterization=characterization, qcProfiles={"capture": profile} + ) + ctx = SimpleNamespace(deps=deps) + compact = tools.compact_context_evidence( + CovariateEvidence(characterization=characterization, qcProfiles=[profile]) + ) + assert compact["characterization"]["columns"][0]["levelCountsOmitted"] == 4 + detail = asyncio.run(tools.inspect_context_evidence(ctx, "column", "sample")) + assert len(detail["levelCounts"]) == 12 + assert asyncio.run( + tools.inspect_context_evidence(ctx, "comparison", measured.evidenceId) + ) == measured.model_dump(mode="json") + assert ( + asyncio.run(tools.inspect_context_evidence(ctx, "confounding", "0")) + == characterization.confounding[0] + ) + policy = asyncio.run(tools.inspect_context_evidence(ctx, "qcProfile", "capture")) + assert "captureDetailsRequired" in policy + assert "resolvedBounds" not in policy + assert profile.resolvedBounds == [{"upper": 8}] + for section, key, capture in [ + ("qcProfile", "missing", None), + ("column", "missing", None), + ("confounding", "99", None), + ("column", "sample", "capture"), + ]: + with pytest.raises(ModelRetry): + asyncio.run(tools.inspect_context_evidence(ctx, section, key, capture)) + with pytest.raises(ModelRetry, match="Inspect covariates"): + asyncio.run( + tools.inspect_context_evidence( + SimpleNamespace(deps=ExperimentalContextDependencies()), + "column", + "sample", + ) + ) + + +@pytest.mark.parametrize( + "change,reason", + [ + ({"metadataColumn": None}, "only metadataColumn"), + ({"origin": "derivedArtifact"}, "ingestionMetadata"), + ({"sourceType": "artifact", "metadataColumn": None}, "only artifact"), + ({"activeCells": 2, "missingCells": 3}, "counts are inconsistent"), + ( + {"activeCells": 3, "missingCells": 1, "usableForFiltering": True}, + "cannot drive filtering", + ), + ], +) +def test_qc_metric_source_cannot_claim_valid_filtering_with_inconsistent_provenance( + change, reason +): + with pytest.raises(ValidationError, match=reason): + QcMetricSourceEvidence.model_validate( + { + "sourceId": "source", + "metricName": "counts", + "metadataColumn": "counts", + **change, + } + ) + + +@pytest.mark.parametrize( + "change,reason", + [ + ({"response": ""}, "non-empty"), + ( + {"explanatoryColumns": ["treatment"], "protectCombination": True}, + "two explanatory", + ), + ], +) +def test_proposal_cannot_omit_its_question_or_invent_combination_protection( + change, reason +): + with pytest.raises(ValidationError, match=reason): + _proposal(**change) + + +@pytest.mark.parametrize( + "damage,reason", + [ + ("runAndSelection", "mutually exclusive"), + ("foreignRun", "opened from this datastore"), + ("notSelection", "datastore cell_selection"), + ("wrongNeighbors", "neighbors ArtifactRef"), + ("foreignNeighbors", "same cell selection"), + ("wrongConnectivity", "connectivity graph"), + ("foreignConnectivity", "same cell selection"), + ("cellKey", "cellQc.cellKey"), + ("duplicateMetric", "source names must be unique"), + ], +) +def test_context_rejects_foreign_artifacts_and_ambiguous_sources_before_model_execution( + monkeypatch, damage, reason +): + store = _Store() + wrong = ArtifactRef("datastore", "cell_selection", "f" * 64) + metric = NamedArtifactSource( + name="mitochondrial", + artifact=ArtifactReferenceModel( + scope="assay", assay="RNA", kind="quality_metric", artifactId="c" * 64 + ), + ) + monkeypatch.setattr( + context_agent, + "_derive_missing_percentage_artifacts", + lambda *a, **k: [metric, metric] if damage == "duplicateMetric" else [], + ) + monkeypatch.setattr( + context_agent, "resolve_cell_aligned_artifact", lambda *a, **k: None + ) + monkeypatch.setattr(context_agent, "graph_cell_selection", lambda *a: wrong) + monkeypatch.setattr( + context_agent, + "run_agent_sync", + lambda **k: pytest.fail("Invalid artifacts cannot reach the model"), + ) + kwargs = {"cell_selection": store.cell_selection} + if damage == "runAndSelection": + kwargs["run"] = SimpleNamespace(_owner=store) + elif damage == "foreignRun": + kwargs = {"run": SimpleNamespace(_owner=object())} + elif damage == "notSelection": + kwargs["cell_selection"] = ArtifactRef("assay", "neighbors", "a" * 64, "RNA") + elif damage == "wrongNeighbors": + kwargs["neighbors"] = store.cell_selection + elif damage == "foreignNeighbors": + kwargs["neighbors"] = ArtifactRef("assay", "neighbors", "a" * 64, "RNA") + elif damage == "wrongConnectivity": + kwargs["connectivity_map"] = store.cell_selection + elif damage == "foreignConnectivity": + kwargs["connectivity_map"] = ArtifactRef( + "assay", "connectivity_map", "a" * 64, "RNA" + ) + elif damage == "cellKey": + kwargs["directions"] = {"cellQc": {"cellKey": "I"}} + with pytest.raises((ValueError, TypeError), match=reason): + context_agent.ExperimentalContextAgent(object()).run(store, **kwargs) diff --git a/tests/test_agent_context_tool_budget.py b/tests/test_agent_context_tool_budget.py new file mode 100644 index 00000000..f9214ffb --- /dev/null +++ b/tests/test_agent_context_tool_budget.py @@ -0,0 +1,171 @@ +"""Context evidence reads fit the bounded run without expanding scientific work.""" + +import asyncio +from copy import deepcopy +from types import SimpleNamespace + +import pytest +from pydantic_ai import ModelRetry, UsageLimitExceeded +from pydantic_ai.messages import ModelResponse, ToolCallPart +from pydantic_ai.models.function import FunctionModel + +from scarf.agent.config import AgentRunConfig +from scarf.agent.experimental_context import ExperimentalContextAgent, tools +from scarf.agent.experimental_context.contracts import ( + CovariateCharacterization, + CovariateEvidence, + CovariateProposal, + ExperimentalContextDependencies, +) +from tests.test_agent_experimental_context import _Store, _design_decision + + +@pytest.mark.parametrize("tool_limit", [None, 6]) +def test_context_can_read_required_saved_details_with_bounded_usage( + monkeypatch, tool_limit +): + store = _Store() + decision = _design_decision() + calls = [] + snapshots = {} + attempts = [] + characterize = tools.characterize_covariates + + def measured(*args, **kwargs): + calls.append(kwargs["directions"]) + return characterize(*args, **kwargs) + + monkeypatch.setattr(tools, "characterize_covariates", measured) + actions = [ + ("inspect_cell_covariates", {}), + ( + "analyze_experimental_design", + { + "column_domains": decision.columnDomains, + "coefficients_of_interest": decision.coefficientsOfInterest, + "units_of_inference": { + name: unit.model_dump() + for name, unit in decision.unitsOfInference.items() + }, + "batch_columns": decision.batchCorrection.batchColumns, + }, + ), + # Reproduce the notebook's malformed lookup and bounded tool repair. + ( + "inspect_context_evidence", + {"section": "confounding", "record_id": "confounding:disease:batch"}, + ), + ( + "inspect_context_evidence", + {"section": "confounding", "record_id": "0"}, + ), + ( + "inspect_context_evidence", + {"section": "coefficient", "record_id": "disease"}, + ), + *[ + ("inspect_context_evidence", {"section": "column", "record_id": name}) + for name in ("batch", "sample", "donor") + ], + ] + requests = 0 + measured_before_details = None + + async def reply(_messages, info): + nonlocal requests, measured_before_details + index = requests + requests += 1 + if index == 2: + measured_before_details = len(calls) + assert measured_before_details > 0 + if index >= 2: + assert len(calls) == measured_before_details + if index < len(actions): + name, args = actions[index] + else: + name, args = info.output_tools[0].name, decision.model_dump() + return ModelResponse(parts=[ToolCallPart(tool_name=name, args=args)]) + + def checkpoint_write(key, value): + snapshots[key] = deepcopy(value) + + agent = ExperimentalContextAgent( + FunctionModel(reply), + config=AgentRunConfig(toolCallLimit=tool_limit) if tool_limit else None, + ) + kwargs = { + "study_context": "Case-control study with samples nested in donors.", + "cell_selection": store.cell_selection, + "checkpoint_write": checkpoint_write, + "on_attempt": attempts.append, + } + if tool_limit is not None: + with pytest.raises(UsageLimitExceeded, match="tool_calls_limit of 6"): + agent.run(store, **kwargs) + assert attempts[-1].status == "failed" + assert "result" not in snapshots + else: + result = agent.run(store, **kwargs) + assert result.status == "done" + assert result.decision.batchCorrection.action == "unsafe" + assert len(result.runInfo.toolCalls) == 8 + assert len(result.runInfo.validationRetries) == 1 + assert result.runInfo.usage.requests == 9 + assert len(calls) == measured_before_details + assert requests <= 10 + assert snapshots["design1"]["state"]["designRounds"] == 1 + assert "design2" not in snapshots + + # More room to inspect evidence does not permit a fifth follow-up proposal. + deps = ExperimentalContextDependencies.model_validate(snapshots["design1"]["state"]) + saved = deps.model_dump(mode="json") + with pytest.raises(ModelRetry, match="eight initial and four follow-up"): + asyncio.run( + tools.analyze_experimental_design( + SimpleNamespace(deps=deps), + column_domains=decision.columnDomains, + coefficients_of_interest=decision.coefficientsOfInterest, + units_of_inference=decision.unitsOfInference, + batch_columns=decision.batchCorrection.batchColumns, + proposals=[ + CovariateProposal( + response="disease", + explanatoryColumns=["batch"], + observationUnit="sample", + rationale="Check batch confounding across observation units.", + ) + ] + * 5, + ) + ) + assert deps.model_dump(mode="json") == saved + + +def test_context_summary_offers_exact_confounding_lookup_without_mutation(): + characterization = CovariateCharacterization( + status="done", + confounding=[ + {"coefficient": "disease", "reason": "Disease aliases the batch."}, + {"coefficient": "sex", "reason": "Sex is incompletely crossed."}, + ], + ) + evidence = CovariateEvidence(characterization=characterization) + original = evidence.model_dump(mode="json") + view = tools.compact_context_evidence(evidence) + ctx = SimpleNamespace( + deps=ExperimentalContextDependencies(characterization=characterization) + ) + for index, row in enumerate(view["characterization"]["confounding"]): + assert row["details"] == {"section": "confounding", "record_id": str(index)} + detail = asyncio.run(tools.inspect_context_evidence(ctx, **row["details"])) + assert detail == characterization.confounding[index] + detail["reason"] = "Returned copies cannot rewrite committed evidence." + with pytest.raises(ModelRetry) as exc_info: + asyncio.run( + tools.inspect_context_evidence( + ctx, "confounding", "confounding:disease:batch" + ) + ) + assert "record_id" in str(exc_info.value) + assert '["0", "1"]' in str(exc_info.value) + assert evidence.model_dump(mode="json") == original diff --git a/tests/test_agent_data_enrichment.py b/tests/test_agent_data_enrichment.py index 09b2eab4..c421757b 100644 --- a/tests/test_agent_data_enrichment.py +++ b/tests/test_agent_data_enrichment.py @@ -1,13 +1,20 @@ """Tests for the read-only data enrichment agent.""" +from tests.agent_examples import example + +import asyncio from types import SimpleNamespace import pytest from pydantic import BaseModel +from pydantic_ai import ModelRetry, UnexpectedModelBehavior from pydantic_ai.messages import ModelMessage, ModelResponse, ToolCallPart from pydantic_ai.models.function import AgentInfo, FunctionModel -from scarf.agent.characterize_features import FeatureCharacterization +import scarf.agent.data_enrichment.agent as data_enrichment_agent_module +import scarf.agent.data_enrichment.tools as data_enrichment_tools +import scarf.agent.data_enrichment.validation as data_enrichment_validation +from scarf.agent.data_enrichment.characterization import FeatureCharacterization from scarf.agent.data_enrichment import ( AdtControlEvidence, AssayFeatureInspection, @@ -28,6 +35,8 @@ FeatureSelectionPolicy, HtoTagEvidence, StudyContextSummary, + find_present_features, + find_present_features_batch, validate_data_enrichment_report, ) @@ -126,14 +135,14 @@ def test_data_enrichment_models_have_factories_and_camelcase_fields() -> None: for model_type in model_types: assert isinstance(model_type.get_blank(), model_type) - assert isinstance(model_type.get_example(), model_type) + assert isinstance(example(model_type), model_type) assert all("_" not in field_name for field_name in model_type.model_fields) def test_data_enrichment_agent_uses_only_read_tools_and_context( monkeypatch: pytest.MonkeyPatch, ) -> None: - from scarf.agent import data_enrichment as module + from scarf.agent.data_enrichment import tools as module store = ReadOnlyStore() tool_names: set[str] = set() @@ -253,7 +262,7 @@ async def reply( def test_data_enrichment_batches_grounded_multimodal_evidence( monkeypatch: pytest.MonkeyPatch, ) -> None: - from scarf.agent import data_enrichment as module + from scarf.agent.data_enrichment import tools as module assay_features = { "RNA": ( @@ -407,7 +416,7 @@ async def reply( def test_data_enrichment_retries_hallucinated_features( monkeypatch: pytest.MonkeyPatch, ) -> None: - from scarf.agent import data_enrichment as module + from scarf.agent.data_enrichment import tools as module store = ReadOnlyStore() monkeypatch.setattr( @@ -483,6 +492,308 @@ async def reply( assert state["request"] == 3 +def test_data_enrichment_fails_after_completed_inspection_without_selection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from scarf.agent.data_enrichment import tools as module + + store = ReadOnlyStore() + monkeypatch.setattr( + module, + "characterize_features", + lambda *_args, **_kwargs: characterization(), + ) + tool_retries: dict[str, int] = {} + + def unavailable_structured_output(**kwargs: object) -> None: + deps = kwargs["deps"] + assert isinstance(deps, DataEnrichmentDependencies) + for tool in kwargs["tools"]: + tool_retries[tool.name] = tool.max_retries + asyncio.run(module.inspect_assay_features_batch(SimpleNamespace(deps=deps))) + raise UnexpectedModelBehavior("structured output unavailable") + + monkeypatch.setattr( + data_enrichment_agent_module, + "run_agent_sync", + unavailable_structured_output, + ) + result = DataEnrichmentAgent(object()).run( + store, + context=DataEnrichmentContext(organismHint="human"), + ) + + assert result.status == "failed" + assert result.runInfo.agentName == "data_enrichment_failed" + assert result.policies == [] + assert result.unresolvedQuestions == [] + assert result.inspections[0].species == "unknown" + assert "No scientific feature policy was selected" in result.limitations[0] + assert tool_retries == { + "inspect_assay_features_batch": 1, + "find_present_features_batch": 1, + } + + +def test_data_enrichment_preserves_validated_policy_uncertainty( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from scarf.agent.data_enrichment import tools as module + from scarf.agent.types import AgentRunInfo + + store = ReadOnlyStore() + monkeypatch.setattr( + module, + "characterize_features", + lambda *_args, **_kwargs: characterization(), + ) + + def unresolved_structured_output(**kwargs: object) -> SimpleNamespace: + deps = kwargs["deps"] + assert isinstance(deps, DataEnrichmentDependencies) + asyncio.run(module.inspect_assay_features_batch(SimpleNamespace(deps=deps))) + return SimpleNamespace( + output=DataEnrichmentReport( + status="needsInput", + unresolvedQuestions=[ + "The objective does not resolve which response genes must be protected." + ], + ), + runInfo=AgentRunInfo(agentName="data_enrichment", modelName="test-model"), + ) + + monkeypatch.setattr( + data_enrichment_agent_module, + "run_agent_sync", + unresolved_structured_output, + ) + result = DataEnrichmentAgent(object()).run( + store, + context=DataEnrichmentContext(organismHint="human"), + ) + + assert result.status == "needsInput" + assert result.policies == [] + assert result.unresolvedQuestions == [ + "The objective does not resolve which response genes must be protected." + ] + assert result.runInfo.agentName == "data_enrichment" + + +def test_feature_lookup_cache_rejects_different_arguments() -> None: + deps = DataEnrichmentDependencies( + store=ReadOnlyStore(), + assays=["RNA"], + ) + run_context = SimpleNamespace(deps=deps) + + first = asyncio.run( + find_present_features_batch( + run_context, + queries_by_assay={"RNA": ["MT-CYB"]}, + ) + ) + repeated = asyncio.run( + find_present_features_batch( + run_context, + queries_by_assay={"RNA": ["MT-CYB"]}, + ) + ) + + assert repeated is first + with pytest.raises(ModelRetry, match="already completed"): + asyncio.run( + find_present_features_batch( + run_context, + queries_by_assay={"RNA": ["GAPDH"]}, + ) + ) + + +def test_data_enrichment_tool_helpers_cover_resolution_edges() -> None: + assert data_enrichment_tools._assay_modality(None, "RNAassay") == ( + "RNA", + "RNAassay", + "assayClass", + ) + assert data_enrichment_tools._assay_modality(None, "ATACassay") == ( + "ATAC", + "ATACassay", + "assayClass", + ) + assert data_enrichment_tools._assay_modality(None, "CustomAssay") == ( + "unsupported", + "CustomAssay", + "assayClass", + ) + assert data_enrichment_tools._assay_modality(None, "") == ( + "unsupported", + "Assay", + "unknown", + ) + assert data_enrichment_tools._valid_peak_coordinate("chr1:10") is False + assert data_enrichment_tools._valid_peak_coordinate("chr1:start-20") is False + controls = data_enrichment_tools._inspect_adt_features( + "ADT", + [("control-1", "IgG control")], + ) + assert controls[0].matchedToken == "control" + assert data_enrichment_tools._inspect_atac_features("ATAC", []).status == "invalid" + assert ( + data_enrichment_tools._inspect_atac_features( + "ATAC", + ["chr1:10-20"], + ).status + == "valid" + ) + future_tool = SimpleNamespace(name="future_tool") + assert ( + data_enrichment_tools._prepare_data_enrichment_tool( + SimpleNamespace(deps=DataEnrichmentDependencies()), + future_tool, + ) + is future_tool + ) + + +def test_find_present_features_reports_casefold_ambiguity_and_absence() -> None: + class LookupFeatures: + @staticmethod + def fetch_all(column: str) -> list[str]: + return { + "ids": ["GENE1", "GENE2", "GAPDH"], + "names": ["shared", "shared", "GAPDH"], + }[column] + + store = SimpleNamespace( + get_assay=lambda _name: SimpleNamespace(feats=LookupFeatures()) + ) + deps = DataEnrichmentDependencies(store=store, assays=["RNA"]) + + result = asyncio.run( + find_present_features( + SimpleNamespace(deps=deps), + assay_name="RNA", + queries=["gapdh", "shared", "missing"], + ) + ) + + assert [item.status for item in result.results] == [ + "present", + "ambiguous", + "absent", + ] + assert deps.confirmedFeatures["RNA"] == {"GAPDH"} + assert result.results[1].evidenceIds == [] + assert result.results[2].matches == [] + + +def test_feature_lookup_tools_reject_invalid_requests() -> None: + with pytest.raises(ModelRetry, match="datastore is unavailable"): + asyncio.run( + find_present_features( + SimpleNamespace( + deps=DataEnrichmentDependencies(store=None, assays=["RNA"]) + ), + assay_name="RNA", + queries=["GAPDH"], + ) + ) + deps = DataEnrichmentDependencies(store=ReadOnlyStore(), assays=["RNA"]) + context = SimpleNamespace(deps=deps) + with pytest.raises(ModelRetry, match="requested assays"): + asyncio.run( + find_present_features( + context, + assay_name="ADT", + queries=["CD3"], + ) + ) + with pytest.raises(ModelRetry, match="between 1 and 50"): + asyncio.run( + find_present_features( + context, + assay_name="RNA", + queries=[], + ) + ) + with pytest.raises(ModelRetry, match="Unknown requested assays"): + asyncio.run( + find_present_features_batch( + context, + queries_by_assay={"ADT": ["CD3"]}, + ) + ) + with pytest.raises(ModelRetry, match="at least one assay"): + asyncio.run(find_present_features_batch(context, queries_by_assay={})) + with pytest.raises(ModelRetry, match="cannot be empty"): + asyncio.run( + find_present_features_batch( + context, + queries_by_assay={"RNA": []}, + ) + ) + with pytest.raises(ModelRetry, match="at most 50"): + asyncio.run( + find_present_features_batch( + context, + queries_by_assay={ + "RNA": [f"gene-{index}" for index in range(51)], + }, + ) + ) + with pytest.raises(ModelRetry, match="requested assays"): + asyncio.run( + data_enrichment_tools.inspect_assay_features( + context, + assay_name="ADT", + ) + ) + empty = SimpleNamespace( + deps=DataEnrichmentDependencies(store=ReadOnlyStore(), assays=[]) + ) + with pytest.raises(ModelRetry, match="No assays were requested"): + asyncio.run(data_enrichment_tools.inspect_assay_features_batch(empty)) + + +def test_assay_inspection_rejects_failed_or_modified_characterization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + context = SimpleNamespace( + deps=DataEnrichmentDependencies(store=ReadOnlyStore(), assays=["RNA"]) + ) + monkeypatch.setattr( + data_enrichment_tools, + "characterize_features", + lambda *_args, **_kwargs: FeatureCharacterization( + status="failed", + notes=["feature inspection failed"], + ), + ) + with pytest.raises(ModelRetry, match="feature inspection failed"): + asyncio.run( + data_enrichment_tools.inspect_assay_features( + context, + assay_name="RNA", + ) + ) + + modified = characterization() + modified.assays[0]["defaultFeatureInventory"] = {"blacklist": "modified"} + monkeypatch.setattr( + data_enrichment_tools, + "characterize_features", + lambda *_args, **_kwargs: modified, + ) + with pytest.raises(ModelRetry, match="exact default HVG blacklist"): + asyncio.run( + data_enrichment_tools.inspect_assay_features( + context, + assay_name="RNA", + ) + ) + + def test_data_enrichment_validates_policy_and_assay() -> None: with pytest.raises(ValueError, match="both excluded and protected"): FeatureSelectionPolicy( @@ -536,6 +847,30 @@ def test_enrichment_copies_caller_context_instead_of_model_paraphrases() -> None ] +def test_enrichment_rejects_duplicate_assay_policies() -> None: + inspection = AssayFeatureInspection( + assay="RNA", + species="unknown", + evidenceIds=["assay:RNA:species"], + ) + deps = DataEnrichmentDependencies( + store=ReadOnlyStore(), + assays=["RNA"], + inspections={"RNA": inspection}, + evidenceIds={"assay:RNA:species"}, + ) + policy = FeatureSelectionPolicy( + assay="RNA", + evidenceIds=["assay:RNA:species"], + ) + + with pytest.raises(ValueError, match="one policy for each assay"): + validate_data_enrichment_report( + deps, + DataEnrichmentReport(status="done", policies=[policy, policy.model_copy()]), + ) + + def test_enrichment_rejects_protected_family_exclusion() -> None: family = FeatureFamilyEvidence( family="cellCycle", @@ -596,3 +931,210 @@ def test_artificial_feature_requires_feature_specific_evidence() -> None: with pytest.raises(ValueError, match="feature-specific"): validate_data_enrichment_report(deps, report) + + +def test_study_context_summary_rejects_unbounded_or_ungrounded_references() -> None: + context = DataEnrichmentContext(studyContext="Human lung study") + with pytest.raises(ValueError, match="verbatim caller text"): + data_enrichment_validation._ground_study_context_summary( + context, + StudyContextSummary(hypothesisReferences=["invented hypothesis"]), + ) + with pytest.raises(ValueError, match="at most 12"): + data_enrichment_validation._ground_study_context_summary( + context, + StudyContextSummary( + tissueReferences=[f"tissue-{index}" for index in range(13)] + ), + ) + long_reference = "x" * 241 + with pytest.raises(ValueError, match="may not exceed 240"): + data_enrichment_validation._ground_study_context_summary( + DataEnrichmentContext(studyContext=long_reference), + StudyContextSummary(analysisIntentReferences=[long_reference]), + ) + + +def test_feature_policy_validation_rejects_ungrounded_contracts() -> None: + inspection = AssayFeatureInspection( + assay="RNA", + species="unknown", + families=[ + FeatureFamilyEvidence( + family="mitochondrial", + defaultExclude=True, + evidenceId="assay:RNA:family:mitochondrial", + ) + ], + evidenceIds=["assay:RNA:species"], + ) + deps = DataEnrichmentDependencies( + store=ReadOnlyStore(), + context=DataEnrichmentContext(), + assays=["RNA"], + inspections={"RNA": inspection}, + evidenceIds={"assay:RNA:species"}, + ) + grounded = StudyContextSummary() + validate = data_enrichment_validation._validate_feature_policy + + unsupported = FeatureSelectionPolicy.model_construct( + assay="RNA", + species="unsupported", + evidenceIds=["assay:RNA:species"], + ) + with pytest.raises(ValueError, match="unsupported species"): + validate(deps, unsupported, grounded) + with pytest.raises(ValueError, match="requires evidence IDs"): + validate(deps, FeatureSelectionPolicy(assay="RNA"), grounded) + with pytest.raises(ValueError, match="was not inspected"): + validate( + deps.model_copy(update={"inspections": {}}), + FeatureSelectionPolicy( + assay="RNA", + evidenceIds=["assay:RNA:species"], + ), + grounded, + ) + with pytest.raises(ValueError, match="must cite context evidence"): + validate( + deps, + FeatureSelectionPolicy( + assay="RNA", + species="homo_sapiens", + evidenceIds=["assay:RNA:species"], + ), + grounded, + ) + with pytest.raises(ValueError, match="conflicts with inspected"): + validate( + deps.model_copy( + update={ + "inspections": { + "RNA": inspection.model_copy(update={"species": "homo_sapiens"}) + } + } + ), + FeatureSelectionPolicy( + assay="RNA", + species="mus_musculus", + evidenceIds=["assay:RNA:species"], + ), + grounded, + ) + with pytest.raises(ValueError, match="unobserved families"): + validate( + deps, + FeatureSelectionPolicy( + assay="RNA", + protectFamilies=["cellCycle"], + evidenceIds=["assay:RNA:species"], + ), + grounded, + ) + with pytest.raises(ValueError, match="unknown evidence IDs"): + validate( + deps, + FeatureSelectionPolicy( + assay="RNA", + evidenceIds=["evidence:unknown"], + ), + grounded, + ) + + +def test_feature_policy_accepts_feature_specific_context_evidence() -> None: + inspection = AssayFeatureInspection( + assay="RNA", + species="unknown", + evidenceIds=["assay:RNA:species"], + ) + deps = DataEnrichmentDependencies( + store=ReadOnlyStore(), + context=DataEnrichmentContext( + experimentalDetails=["ERCC-00002 spike-in"], + ), + assays=["RNA"], + inspections={"RNA": inspection}, + confirmedFeatures={"RNA": {"ERCC-00002"}}, + evidenceIds={"assay:RNA:species", "context:experiment:0"}, + ) + policy = FeatureSelectionPolicy( + assay="RNA", + artificialFeatures=["ERCC-00002"], + evidenceIds=["assay:RNA:species", "context:experiment:0"], + ) + + data_enrichment_validation._validate_feature_policy( + deps, + policy, + StudyContextSummary(), + ) + + assert policy.artificialFeatures == ["ERCC-00002"] + + +def test_data_enrichment_report_rejects_incomplete_assay_inventory() -> None: + inspection = AssayFeatureInspection( + assay="RNA", + species="unknown", + evidenceIds=["assay:RNA:species"], + ) + with pytest.raises(ValueError, match="Inspect every requested assay"): + validate_data_enrichment_report( + DataEnrichmentDependencies(store=ReadOnlyStore(), assays=["RNA"]), + DataEnrichmentReport( + status="needsInput", + unresolvedQuestions=["Inspect the requested assays."], + ), + ) + deps = DataEnrichmentDependencies( + store=ReadOnlyStore(), + assays=["RNA"], + inspections={"RNA": inspection}, + evidenceIds={"assay:RNA:species"}, + ) + with pytest.raises(ValueError, match="outside the requested set"): + validate_data_enrichment_report( + deps, + DataEnrichmentReport( + status="needsInput", + unresolvedQuestions=["Resolve the assay mismatch."], + policies=[ + FeatureSelectionPolicy( + assay="ADT", + evidenceIds=["assay:RNA:species"], + ) + ], + ), + ) + with pytest.raises(ValueError, match="one policy for every requested assay"): + validate_data_enrichment_report( + deps.model_copy(update={"assays": ["RNA", "ADT"]}), + DataEnrichmentReport( + status="done", + policies=[ + FeatureSelectionPolicy( + assay="RNA", + evidenceIds=["assay:RNA:species"], + ) + ], + ), + ) + + +def test_failed_enrichment_retains_partial_evidence_without_inventing_policy() -> None: + incomplete = DataEnrichmentDependencies(store=ReadOnlyStore(), assays=["RNA"]) + for inspections in ( + {}, + {"RNA": AssayFeatureInspection(assay="RNA", species="unknown")}, + ): + failed = data_enrichment_validation.failed_data_enrichment_report( + incomplete.model_copy(update={"inspections": inspections}), + error=RuntimeError("model failed"), + model_name="test", + ) + assert failed.status == "failed" + assert failed.policies == [] + assert failed.inspections == list(inspections.values()) + assert "RuntimeError: model failed" in failed.limitations diff --git a/tests/test_agent_decide.py b/tests/test_agent_decide.py index 2fbf84ea..1ef7d5ea 100644 --- a/tests/test_agent_decide.py +++ b/tests/test_agent_decide.py @@ -8,8 +8,9 @@ from pydantic_ai.models.function import AgentInfo from pydantic_ai.models.test import TestModel -from scarf.agent import DecisionValidationError, EvidenceItem, decide -from scarf.agent.decide import _SYSTEM_PROMPT, validate_decision +from scarf.agent.decisions.selection import DecisionValidationError, decide +from scarf.agent.types import EvidenceItem +from scarf.agent.decisions.selection import _SYSTEM_PROMPT, validate_decision from scarf.agent.types import Decision @@ -71,26 +72,24 @@ def test_decide_with_test_model_returns_schema_valid_decision() -> None: validate_decision(result, _evidence()) -def test_validate_decision_coerces_selected_id_embedded_in_line() -> None: +def test_validate_decision_rejects_selected_id_embedded_in_line() -> None: decision = Decision( selectedId="id=matrix:raw/X | label=raw/X | summary=integer-like", rationale="integer-like", evidenceIds=[], ) - result = validate_decision(decision, _evidence()) - assert result.selectedId == "matrix:raw/X" - assert result.evidenceIds == ["matrix:raw/X"] + with pytest.raises(DecisionValidationError, match="not in evidence ids"): + validate_decision(decision, _evidence()) -def test_validate_decision_coerces_id_equals_prefix_in_evidence_ids() -> None: +def test_validate_decision_rejects_id_equals_prefix_in_evidence_ids() -> None: decision = Decision( selectedId="id=matrix:raw/X", rationale="echoed prompt scaffolding", evidenceIds=["id=matrix:raw/X"], ) - result = validate_decision(decision, _evidence()) - assert result.selectedId == "matrix:raw/X" - assert result.evidenceIds == ["matrix:raw/X"] + with pytest.raises(DecisionValidationError, match="not in evidence ids"): + validate_decision(decision, _evidence()) def test_decide_rejects_unknown_selected_id() -> None: @@ -99,12 +98,13 @@ def test_decide_rejects_unknown_selected_id() -> None: rationale="guess", evidenceIds=["matrix:missing"], ) - with pytest.raises(DecisionValidationError, match="selectedId"): + with pytest.raises(DecisionValidationError) as caught: decide( model=_function_model(bad), question="Which matrix looks like raw counts?", evidence=_evidence(), ) + assert "selectedId" in caught.value.agent_run_info.error def test_validate_decision_rejects_unknown_evidence_ids() -> None: diff --git a/tests/test_agent_decision_kernel.py b/tests/test_agent_decision_kernel.py new file mode 100644 index 00000000..a3cd735c --- /dev/null +++ b/tests/test_agent_decision_kernel.py @@ -0,0 +1,382 @@ +"""Contract tests for the decision kernel and deterministic auditor.""" + +import pytest +from pydantic import ValidationError + +from scarf.agent.decisions.kernel import ( + DecisionEvidence, + DecisionOption, + DecisionRecord, + DecisionSpec, + DeterministicDecisionAuditor, + EvidenceBundle, + ProtectedVariableEffect, +) +from scarf.agent.types import ArtifactReferenceModel + + +def _evidence_bundle() -> EvidenceBundle: + return EvidenceBundle( + bundleId="bundle:clusterPartition", + decisionId="clusterPartition", + evidence=[ + DecisionEvidence( + evidenceId="evidence:silhouette", + evidenceClass="geometric", + summary="The coarse partition has the largest silhouette.", + ), + DecisionEvidence( + evidenceId="evidence:markers", + evidenceClass="markerCoherence", + summary="The finer partition has distinct marker programs.", + ), + DecisionEvidence( + evidenceId="evidence:stability", + evidenceClass="resamplingStability", + summary="The finer partition is stable under subsampling.", + ), + DecisionEvidence( + evidenceId="evidence:replicates", + evidenceClass="crossUnitSupport", + summary="The finer populations occur in independent units.", + ), + ], + ) + + +def _decision_spec() -> DecisionSpec: + return DecisionSpec( + decisionId="clusterPartition", + definitionVersion=1, + checkpoint="clustering", + question="Which registered partition is defensible?", + evidenceBundleId="bundle:clusterPartition", + options=[ + DecisionOption( + optionId="partition:coarse", + status="apply", + label="Coarse partition", + description="Use the metric-preferred coarse partition.", + ), + DecisionOption( + optionId="partition:fine", + status="apply", + label="Fine partition", + description="Use the finer marker-supported partition.", + requiredEvidenceClasses=["markerCoherence"], + ), + DecisionOption( + optionId="partition:abstain", + status="abstain", + label="No discrete partition", + description="Do not claim that a discrete partition is supported.", + ), + ], + baselineOptionId="partition:coarse", + metricPreferredOptionId="partition:coarse", + requireIndependentOverrideEvidence=True, + ) + + +def _decision_record( + *, + record_id: str = "decision:cluster:1", + selected_option_id: str = "partition:coarse", + status: str = "apply", + source: str = "agent", + evidence_ids: list[str] | None = None, + override_of: str | None = None, + override_evidence_ids: list[str] | None = None, + available_evidence_ids: list[str] | None = None, + supersedes: str | None = None, +) -> DecisionRecord: + bundle = _evidence_bundle().with_content_sha256() + assert bundle.contentSha256 is not None + return DecisionRecord( + recordId=record_id, + decisionId="clusterPartition", + definitionVersion=1, + evidenceBundleId=bundle.bundleId, + evidenceBundleSha256=bundle.contentSha256, + offeredOptionIds=[ + "partition:coarse", + "partition:fine", + "partition:abstain", + ], + availableEvidenceIds=available_evidence_ids + if available_evidence_ids is not None + else [item.evidenceId for item in bundle.evidence], + selectedOptionId=selected_option_id, + status=status, + source=source, + evidenceIds=evidence_ids or ["evidence:silhouette"], + rationale="The cited evidence supports this registered option.", + confidence="medium", + overrideOfOptionId=override_of, + overrideEvidenceIds=override_evidence_ids or [], + ) + + +def test_option_contract_rejects_freeform_execution_parameters() -> None: + with pytest.raises(ValidationError, match="Extra inputs are not permitted"): + DecisionOption.model_validate( + { + "optionId": "pca:30", + "status": "apply", + "label": "Thirty PCs", + "description": "Use the registered thirty-PC prefix.", + "dimensions": 30, + } + ) + + +def test_evidence_bundle_rejects_duplicate_evidence_ids() -> None: + item = DecisionEvidence( + evidenceId="evidence:qc", + evidenceClass="qualityControl", + summary="Observed cell quality evidence.", + ) + with pytest.raises(ValidationError, match="must not contain duplicates"): + EvidenceBundle( + bundleId="bundle:qc", + decisionId="cellQuality", + evidence=[item, item], + ) + + +def test_evidence_contract_rejects_duplicate_artifacts_and_bad_checksums() -> None: + reference = ArtifactReferenceModel( + scope="assay", + assay="RNA", + kind="pca", + artifactId="a" * 64, + ) + with pytest.raises(ValidationError, match="must not contain duplicates"): + DecisionEvidence( + evidenceId="evidence:pca", + evidenceClass="geometric", + summary="Observed PCA evidence.", + artifactReferences=[reference, reference], + ) + + values = _evidence_bundle().model_dump() + values["contentSha256"] = "0" * 64 + with pytest.raises(ValidationError, match="does not match"): + EvidenceBundle.model_validate(values) + values["contentSha256"] = "INVALID" + with pytest.raises(ValidationError, match="lowercase SHA-256"): + EvidenceBundle.model_validate(values) + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ( + {"selectedOptionId": "partition:invented"}, + "selectedOptionId must reference an offered option", + ), + ( + {"evidenceIds": ["evidence:invented"]}, + "evidenceIds must reference only available evidence", + ), + ( + { + "overrideOfOptionId": None, + "evidenceIds": ["evidence:markers"], + "overrideEvidenceIds": ["evidence:markers"], + }, + "overrideEvidenceIds require overrideOfOptionId", + ), + ], +) +def test_decision_record_rejects_non_exact_references( + changes: dict[str, object], message: str +) -> None: + values = _decision_record().model_dump() + values.update(changes) + with pytest.raises(ValidationError, match=message): + DecisionRecord.model_validate(values) + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ( + { + "overrideOfOptionId": "partition:fine", + "overrideEvidenceIds": ["evidence:stability"], + }, + "overrideEvidenceIds must be included", + ), + ( + {"overrideOfOptionId": "partition:invented"}, + "overrideOfOptionId must reference", + ), + ( + {"overrideOfOptionId": "partition:coarse"}, + "overrideOfOptionId must differ", + ), + ( + {"supersedes": "decision:cluster:1"}, + "Extra inputs are not permitted", + ), + ( + { + "protectedVariableEffects": [ + ProtectedVariableEffect( + variable="disease", + status="preserved", + evidenceIds=["evidence:invented"], + summary="Disease structure is preserved.", + ).model_dump() + ] + }, + "protectedVariableEffects must reference", + ), + ], +) +def test_decision_record_rejects_invalid_override_and_lineage_references( + changes: dict[str, object], + message: str, +) -> None: + values = _decision_record().model_dump() + values.update(changes) + with pytest.raises(ValidationError, match=message): + DecisionRecord.model_validate(values) + + +def test_auditor_accepts_an_exact_metric_preferred_decision() -> None: + record = _decision_record() + + verification = DeterministicDecisionAuditor.audit( + _decision_spec(), _evidence_bundle(), record, created_at_ns=10 + ) + + assert all(check.status == "passed" for check in verification) + assert {check.status for check in verification} == {"passed"} + + +def test_auditor_rejects_a_tampered_evidence_bundle_checksum() -> None: + values = _decision_record().model_dump() + values["evidenceBundleSha256"] = "f" * 64 + record = DecisionRecord.model_validate(values) + + verification = DeterministicDecisionAuditor.audit( + _decision_spec(), + _evidence_bundle(), + record, + ) + + assert any(check.status == "failed" for check in verification) + assert [check.checkId for check in verification if check.status == "failed"] == [ + "decisionIdentity" + ] + + +def test_auditor_requires_evidence_bound_to_the_selected_option() -> None: + spec_values = _decision_spec().model_dump() + spec_values["options"][0]["requiredEvidenceIds"] = ["evidence:markers"] + spec = DecisionSpec.model_validate(spec_values) + + verification = DeterministicDecisionAuditor.audit( + spec, + _evidence_bundle(), + _decision_record(), + ) + + assert any(check.status == "failed" for check in verification) + assert [check.checkId for check in verification if check.status == "failed"] == [ + "requiredEvidence" + ] + + +def test_auditor_rejects_a_tampered_option_or_evidence_inventory() -> None: + record = _decision_record( + available_evidence_ids=["evidence:silhouette"], + ) + values = record.model_dump() + values["offeredOptionIds"] = ["partition:coarse", "partition:abstain"] + record = DecisionRecord.model_validate(values) + + verification = DeterministicDecisionAuditor.audit( + _decision_spec(), _evidence_bundle(), record + ) + + assert any(check.status == "failed" for check in verification) + failed = {check.checkId for check in verification if check.status == "failed"} + assert failed == {"exactOptionInventory", "exactEvidenceInventory"} + + +def test_auditor_requires_two_independent_non_geometric_override_classes() -> None: + record = _decision_record( + selected_option_id="partition:fine", + evidence_ids=["evidence:silhouette", "evidence:markers"], + override_of="partition:coarse", + override_evidence_ids=["evidence:silhouette", "evidence:markers"], + ) + + verification = DeterministicDecisionAuditor.audit( + _decision_spec(), _evidence_bundle(), record + ) + + assert any(check.status == "failed" for check in verification) + failed = [check for check in verification if check.status == "failed"] + assert [check.checkId for check in failed] == ["independentOverrideEvidence"] + + +def test_auditor_accepts_two_independent_non_geometric_override_classes() -> None: + record = _decision_record( + selected_option_id="partition:fine", + evidence_ids=["evidence:markers", "evidence:stability"], + override_of="partition:coarse", + override_evidence_ids=["evidence:markers", "evidence:stability"], + ) + + verification = DeterministicDecisionAuditor.audit( + _decision_spec(), _evidence_bundle(), record + ) + + assert all(check.status == "passed" for check in verification) + + +def test_auditor_rejects_override_evidence_from_another_option() -> None: + spec_values = _decision_spec().model_dump() + spec_values["options"][1]["requiredEvidenceIds"] = [ + "evidence:markers", + "evidence:stability", + ] + spec = DecisionSpec.model_validate(spec_values) + record = _decision_record( + selected_option_id="partition:fine", + evidence_ids=[ + "evidence:markers", + "evidence:stability", + "evidence:replicates", + ], + override_of="partition:coarse", + override_evidence_ids=["evidence:markers", "evidence:replicates"], + ) + + verification = DeterministicDecisionAuditor.audit( + spec, + _evidence_bundle(), + record, + ) + + assert any(check.status == "failed" for check in verification) + assert [check.checkId for check in verification if check.status == "failed"] == [ + "independentOverrideEvidence" + ] + + +def test_human_choices_obey_the_same_source_and_status_contracts() -> None: + spec_values = _decision_spec().model_dump() + spec_values["allowedSources"] = ["rule", "agent"] + spec = DecisionSpec.model_validate(spec_values) + record = _decision_record(source="human", status="skip") + + verification = DeterministicDecisionAuditor.audit(spec, _evidence_bundle(), record) + + assert any(check.status == "failed" for check in verification) + failed = {check.checkId for check in verification if check.status == "failed"} + assert failed == {"selectedOption", "decisionSource"} diff --git a/tests/test_agent_decision_replay.py b/tests/test_agent_decision_replay.py new file mode 100644 index 00000000..436bd0aa --- /dev/null +++ b/tests/test_agent_decision_replay.py @@ -0,0 +1,203 @@ +"""QC decision ownership, human answers and replay retain exact verified evidence.""" + +from types import SimpleNamespace +from typing import Any +import json + +import pytest + +from scarf.agent.decisions.kernel import DecisionSelection +from scarf.agent.decisions.rna import build_qc_grouping_decision +from scarf.agent.orchestrator import decisions +from tests.test_agent_rna_adaptive import checkpoints as memory_checkpoints # noqa: F401 +from tests.test_agent_rna_decisions import _bundle + + +@pytest.fixture +def decision_case(request: pytest.FixtureRequest) -> Any: + request.getfixturevalue("memory_checkpoints") + definition = build_qc_grouping_decision( + evidence_bundle_id="bundle:qc", + physical_capture_eligible=True, + pooled_reference_eligible=False, + ) + evidence = _bundle("qcGrouping", "bundle:qc", ["qualityControl", "design"]) + selection = DecisionSelection( + selectedOptionId="qcGrouping:global", + evidenceIds=[item.evidenceId for item in evidence.evidence], + rationale="Global thresholds preserve the measured groups without losing a capture.", + confidence="medium", + ) + request = SimpleNamespace( + workflowRunId="workflow", requestSha256="a" * 64, configSha256="b" * 64 + ) + owner = decisions.DecisionStagesMixin() + return owner, request, definition, evidence, selection + + +@pytest.mark.parametrize("owner_name", ["rule", "agent", "human"]) +@pytest.mark.parametrize("defer", [False, True]) +def test_committed_qc_choice_replays_exact_owner_and_pending_outcome( + decision_case: Any, + request: pytest.FixtureRequest, + monkeypatch: pytest.MonkeyPatch, + owner_name: str, + defer: bool, +) -> None: + saved = request.getfixturevalue("memory_checkpoints") + owner, request, definition, evidence, selection = decision_case + if defer: + selection = selection.model_copy( + update={"selectedOptionId": "qcGrouping:defer"} + ) + answers, kwargs = {}, {} + if owner_name == "human": + answers["decision:qcGrouping"] = { + "decisionId": "qcGrouping", + "optionId": selection.selectedOptionId, + "rationale": selection.rationale, + } + else: + kwargs[f"{owner_name}_selection"] = selection + monkeypatch.setattr( + decisions, + "run_agent_sync", + lambda **kw: pytest.fail( + "Committed or explicitly owned decisions must not call a model" + ), + ) + result = owner._resolve_rna_decision( + None, request, definition, evidence, answers, **kwargs + ) + assert result.record.source == owner_name + assert (result.pending is not None) == defer + assert (result.compiled is None) == defer + replay = owner._resolve_rna_decision( + None, request, definition, evidence, answers, **kwargs + ) + assert result == replay + if defer: + assert owner._pending_decision_question(replay, definition).options == [ + item.optionId for item in definition.spec.options + ] + else: + with pytest.raises(ValueError, match="no pending question"): + owner._pending_decision_question(replay, definition) + next(iter(saved.values()))["outputs"]["checks"] = [] + with pytest.raises(ValueError, match="checks differ"): + owner._resolve_rna_decision( + None, request, definition, evidence, answers, **kwargs + ) + + +@pytest.mark.parametrize( + "failure", + [ + "twoOwners", + "ruleAnswer", + "agentAnswer", + "notMapping", + "wrongDecision", + "unoffered", + "wrongIdentity", + ], +) +def test_ambiguous_or_unbound_decision_requests_fail_before_persistence( + decision_case: Any, request: pytest.FixtureRequest, failure: str +) -> None: + saved = request.getfixturevalue("memory_checkpoints") + owner, request, definition, evidence, selection = decision_case + answer = { + "decisionId": "qcGrouping", + "optionId": "qcGrouping:global", + "rationale": selection.rationale, + } + answers, kwargs = {}, {} + if failure == "twoOwners": + kwargs = {"rule_selection": selection, "agent_selection": selection} + reason = "two supplied owners" + elif failure in {"ruleAnswer", "agentAnswer"}: + kwargs = {f"{failure[:-6]}_selection": selection} + answers["decision:qcGrouping"] = answer + reason = "cannot accept a human answer" + elif failure == "notMapping": + answers["decision:qcGrouping"] = "global" + reason = "must be a mapping" + elif failure == "wrongDecision": + answers["decision:qcGrouping"] = {**answer, "decisionId": "other"} + reason = "must name this decisionId" + elif failure == "unoffered": + answers["decision:qcGrouping"] = {**answer, "optionId": "global"} + reason = "offered option" + else: + evidence = evidence.model_copy(update={"decisionId": "other"}) + reason = "exact evidence identities differ" + with pytest.raises(ValueError, match=reason): + owner._resolve_rna_decision( + None, request, definition, evidence, answers, **kwargs + ) + assert not saved + + +def test_programmatic_record_creation_requires_digest_and_exact_option( + decision_case: Any, +) -> None: + _, _, definition, evidence, selection = decision_case + with pytest.raises(ValueError, match="content digest"): + decisions._record_from_selection( + definition, evidence, selection, "rule", None, None, 0 + ) + selection = selection.model_copy(update={"selectedOptionId": "global"}) + with pytest.raises(ValueError, match="not offered"): + decisions._record_from_selection( + definition, evidence.with_content_sha256(), selection, "rule", None, None, 0 + ) + + +def test_qc_model_and_checkpoint_receive_exact_policy_evidence( + decision_case: Any, request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch +) -> None: + from scarf.agent.config import AgentRunConfig + + saved = request.getfixturevalue("memory_checkpoints") + owner, request, definition, evidence, selection = decision_case + owner.model = object() + request.config = SimpleNamespace(agentRunConfig=AgentRunConfig()) + request.request = SimpleNamespace( + studyContext="Observed captures", studyObjective="Preserve the rare joint group" + ) + policy_evidence = { + "policies": [ + { + "evidenceId": evidence.evidence[0].evidenceId, + "resolvedBounds": {"mitochondrial": {"upper": 7.125}}, + "retainedCellsByCombination": {"condition/sex": {"treated/F": 12}}, + } + ] + } + calls = [] + + def provider(**kwargs: Any) -> Any: + payload = json.loads(kwargs["user_prompt"]) + assert payload["qcPolicyEvidence"] == policy_evidence + assert "exact thresholds" in kwargs["system_prompt"] + calls.append(payload) + return SimpleNamespace( + output=kwargs["output_validator"](selection), + runInfo=SimpleNamespace(modelName="offline-reviewer"), + ) + + monkeypatch.setattr(decisions, "run_agent_sync", provider) + result = owner._resolve_rna_decision( + None, request, definition, evidence, {}, qc_evidence=policy_evidence + ) + checkpoint = next(iter(saved.values())) + assert checkpoint["inputs"]["qcPolicyEvidence"] == policy_evidence + assert result.record.rationale == selection.rationale + assert ( + owner._resolve_rna_decision( + None, request, definition, evidence, {}, qc_evidence=policy_evidence + ) + == result + ) + assert len(calls) == 1 diff --git a/tests/test_agent_design_comparisons.py b/tests/test_agent_design_comparisons.py new file mode 100644 index 00000000..ca70993a --- /dev/null +++ b/tests/test_agent_design_comparisons.py @@ -0,0 +1,823 @@ +"""Objective-led metadata comparisons remain bounded and respect study units.""" + +import asyncio +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest +from pydantic import ValidationError +from pydantic_ai import Agent, ModelRetry +from pydantic_ai.messages import ( + ModelMessage, + ModelResponse, + RetryPromptPart, + TextPart, + ToolCallPart, +) +from pydantic_ai.models.function import AgentInfo, FunctionModel +from pydantic_ai.tools import Tool + +from scarf.agent.experimental_context import tools +from scarf.agent.experimental_context import qc_evidence +from scarf.agent.cell_quality.profiles import project_registered_qc_profile +from scarf.agent.experimental_context.comparisons import ( + accept_capture_proposal, + canonical_design_choices, + combination_labels, + compare_covariates, + evaluate_proposals, +) +from scarf.agent.experimental_context.contracts import ( + CaptureProposal, + CovariateCharacterization, + CovariateProposal, + ExperimentalContextDecision, + ExperimentalContextDependencies, +) +from scarf.agent.experimental_context.study import build_study_contract +from scarf.agent.parameter_tuning import diagnostics, execution +from scarf.storage.refs import ArtifactRef + + +class _Cells: + def __init__(self, frame: pd.DataFrame) -> None: + self.frame = frame + self.columns = list(frame.columns) + + def fetch(self, column: str) -> np.ndarray: + return self.frame[column].to_numpy() + + +def _design() -> tuple[_Cells, CovariateCharacterization]: + frame = pd.DataFrame( + { + "sample": [f"s{i}" for i in range(8)], + "treatment": ["a", "a", "b", "b"] * 2, + "time": ["early", "late", "early", "late"] * 2, + "response": ["same", "different", "different", "same"] * 2, + "age": np.arange(8, dtype=float), + } + ) + cells = _Cells(frame.loc[frame.index.repeat(3)].reset_index(drop=True)) + characterization = CovariateCharacterization( + status="done", + columns=[ + { + "name": name, + "kind": "continuous" if name == "age" else "categorical", + "domain": "design" if name == "sample" else "biological", + } + for name in cells.columns + ], + ) + return cells, characterization + + +def _proposal(**changes: object) -> CovariateProposal: + return CovariateProposal.model_validate( + { + "response": "response", + "explanatoryColumns": ["treatment", "time"], + "observationUnit": "sample", + "rationale": "Preserve treatment-by-time structure relevant to the objective.", + **changes, + } + ) + + +def _deps(cells: _Cells) -> ExperimentalContextDependencies: + return ExperimentalContextDependencies( + cells=cells, + store=SimpleNamespace(cells=cells), + cellSelection=ArtifactRef( + scope="datastore", kind="cell_selection", artifact_id="c" * 64 + ), + ) + + +def test_joint_explanation_detects_interaction_without_additive_alias() -> None: + cells, characterization = _design() + result = compare_covariates( + cells, characterization, _proposal(), selection_identity={"id": "one"} + ) + assert result.status == "computed" + assert result.evidence["independentUnits"] == 8 + assert result.evidence["jointEstimability"]["coefficientEstimable"] is True + assert result.evidence["jointGroupEstimability"]["coefficientEstimable"] is False + assert all( + item["value"] == 0 for item in result.evidence["singleAssociations"].values() + ) + assert ( + result.evidence["jointAssociation"]["directionalMapping"]["nesting"] != "none" + ) + + +def test_conditional_comparison_keeps_each_stratum_and_support() -> None: + cells, characterization = _design() + result = compare_covariates( + cells, + characterization, + _proposal(explanatoryColumns=["treatment"], conditionedOn="time"), + selection_identity={}, + ) + assert result.status == "computed" + assert len(result.evidence["strata"]) == 2 + assert all(row["independentUnits"] == 4 for row in result.evidence["strata"]) + + +@pytest.mark.parametrize( + "change,reason", + [ + ( + {"explanatoryColumns": ["treatment"], "conditionedOn": "age"}, + "continuousConditioningIsUnsupported", + ), + ( + {"observationUnit": "response"}, + "observationAndIndependentUnitsMustBeDesignOrTechnical", + ), + ({"response": "missing"}, "unknownObservedColumn"), + ], +) +def test_unsupported_explanations_are_explicit( + change: dict[str, object], reason: str +) -> None: + cells, characterization = _design() + result = compare_covariates( + cells, characterization, _proposal(**change), selection_identity={} + ) + assert result.status == "unsupported" + assert reason in result.reasons + + +def test_missingness_does_not_become_a_categorical_group() -> None: + cells, characterization = _design() + cells.frame.loc[:2, "response"] = None + result = compare_covariates( + cells, characterization, _proposal(), selection_identity={} + ) + assert result.evidence["missingCells"] == 3 + assert result.evidence["missingCellsByColumn"]["response"] == 3 + assert result.evidence["independentUnits"] == 7 + assert result.status == "unsupported" + + +@pytest.mark.parametrize("declared", [False, True]) +def test_biological_donor_can_be_an_explicit_independent_unit(declared: bool) -> None: + cells, characterization = _design() + cells.frame["donor"] = cells.frame["sample"].map( + {f"s{i}": f"d{i % 4}" for i in range(8)} + ) + cells.columns.append("donor") + characterization.columns.append( + {"name": "donor", "kind": "categorical", "domain": "biological"} + ) + if declared: + characterization.coefficients = [ + { + "name": "response", + "observationUnit": "sample", + "independentUnit": "donor", + } + ] + result = compare_covariates( + cells, + characterization, + _proposal(explanatoryColumns=["treatment"], independentUnit="donor"), + selection_identity={}, + ) + if declared: + assert result.status == "computed" + assert result.evidence["independentUnits"] == 4 + assert result.evidence["observationUnits"] == 8 + assert result.evidence["columnDomains"]["donor"] == "biological" + else: + assert result.status == "unsupported" + assert "observationAndIndependentUnitsMustBeDesignOrTechnical" in result.reasons + + +@pytest.mark.parametrize("invalid", ["continuous", "withinDonor"]) +def test_declared_biological_unit_keeps_kind_and_repeated_measure_guards( + invalid: str, +) -> None: + cells, characterization = _design() + cells.frame["donor"] = cells.frame["sample"].map( + {f"s{i}": f"d{i % 4}" for i in range(8)} + ) + cells.columns.append("donor") + characterization.columns.append( + { + "name": "donor", + "kind": "continuous" if invalid == "continuous" else "categorical", + "domain": "biological", + } + ) + characterization.coefficients = [ + {"name": "response", "observationUnit": "sample", "independentUnit": "donor"} + ] + if invalid == "withinDonor": + cells.frame.loc[cells.frame["sample"] == "s4", "response"] = "different" + result = compare_covariates( + cells, + characterization, + _proposal(explanatoryColumns=["treatment"], independentUnit="donor"), + selection_identity={}, + ) + assert result.status == "unsupported" + assert ( + "observationAndIndependentUnitsMustBeCategorical" + if invalid == "continuous" + else "withinIndependentUnitComparisonsAreUnsupported" + ) in result.reasons + + +def test_study_contract_preserves_unsupported_comparison_limitations() -> None: + cells, characterization = _design() + comparison = compare_covariates( + cells, + characterization, + _proposal(observationUnit="response"), + selection_identity={}, + ) + characterization.comparisons = [comparison] + original = characterization.model_dump(mode="json") + contract = build_study_contract( + study_context="Independent donors with technical captures.", + study_objective="Describe populations without unsupported associations.", + experimental_result=SimpleNamespace( + status="done", + decision=ExperimentalContextDecision.get_blank(), + batchSafety=[], + characterization=characterization, + notes=[], + ), + ) + limitation = next( + value for value in contract.limitations if comparison.evidenceId in value + ) + assert comparison.reasons[0] in limitation + assert "no supported association or absence finding" in limitation + assert characterization.model_dump(mode="json") == original + + +def test_two_round_limit_counts_retries_and_reuses_identical_proposals() -> None: + cells, characterization = _design() + deps = _deps(cells) + proposal = _proposal(protectCombination=True) + evaluate_proposals(deps, characterization, [proposal] * 8) + assert len(deps.comparisons) == 1 + assert deps.protectedCombinations == [["time", "treatment"]] + with pytest.raises(ValueError, match="eight initial and four"): + evaluate_proposals(deps, characterization, [proposal] * 5) + evaluate_proposals(deps, characterization, [proposal]) + assert len(deps.comparisons) == 1 + with pytest.raises(ValueError, match="two evidence rounds"): + evaluate_proposals(deps, characterization, []) + + +def test_objective_questions_require_explicit_grounded_purpose_before_a_round() -> None: + cells, characterization = _design() + deps = _deps(cells) + deps.studyObjective = "Explain treatment and time design coverage" + with pytest.raises(ValueError, match="explicit purpose"): + evaluate_proposals(deps, characterization, [_proposal()]) + assert deps.designRounds == 0 + proposal = _proposal(purpose="designCoverage", objectiveQuote=deps.studyObjective) + with pytest.raises(ValueError, match="must remain essential"): + evaluate_proposals( + deps, characterization, [proposal.model_copy(update={"essential": False})] + ) + assert deps.designRounds == 0 + evaluate_proposals(deps, characterization, [proposal]) + assert deps.designRounds == 1 + + +def test_changed_declared_unit_pair_recomputes_one_current_comparison() -> None: + cells, characterization = _design() + cells.frame["donor"] = cells.frame["sample"] + cells.columns.append("donor") + characterization.columns.append( + {"name": "donor", "domain": "biological", "kind": "categorical"} + ) + deps = _deps(cells) + proposal = _proposal(independentUnit="donor") + evaluate_proposals(deps, characterization, [proposal]) + assert deps.comparisons[0].status == "unsupported" + prior_id = deps.comparisons[0].evidenceId + characterization.coefficients = [ + {"name": "response", "observationUnit": "sample", "independentUnit": "donor"} + ] + evaluate_proposals(deps, characterization, [proposal]) + assert len(deps.comparisons) == 1 + assert deps.comparisons[0].status == "computed" + assert deps.comparisons[0].evidenceId != prior_id + + +def test_unsupported_explanation_does_not_discard_protected_biology() -> None: + cells, characterization = _design() + cells.frame.loc[:2, "response"] = None + deps = _deps(cells) + evaluate_proposals(deps, characterization, [_proposal(protectCombination=True)]) + assert deps.comparisons[0].status == "unsupported" + assert deps.protectedCombinations == [["time", "treatment"]] + + +def test_proposals_cannot_exceed_three_columns() -> None: + with pytest.raises(ValidationError, match="three distinct measured") as caught: + _proposal(conditionedOn="age") + message = str(caught.value) + assert "response='response'" in message + assert "explanatoryColumns=['treatment', 'time']" in message + assert "conditionedOn='age'" in message + assert "set conditionedOn=null" in message + assert "joint explanation within strata is unsupported" in message + + +@pytest.mark.parametrize( + ("explanatory", "condition"), + [(["treatment"], None), (["treatment", "time"], None), (["treatment"], "time")], +) +def test_proposal_measurement_limit_excludes_observation_and_independent_units( + explanatory: list[str], condition: str | None +) -> None: + proposal = _proposal( + explanatoryColumns=explanatory, + conditionedOn=condition, + observationUnit="sample", + independentUnit="donor", + ) + assert proposal.explanatoryColumns == explanatory + assert proposal.conditionedOn == condition + assert proposal.independentUnit == "donor" + assert _proposal(independentUnit="sample").independentUnit == "sample" + + +@pytest.mark.parametrize( + ("changes", "repeated"), + [ + ({"explanatoryColumns": ["treatment", "treatment"]}, "treatment"), + ({"explanatoryColumns": ["response"]}, "response"), + ( + {"explanatoryColumns": ["treatment"], "conditionedOn": "response"}, + "response", + ), + ( + {"explanatoryColumns": ["treatment"], "conditionedOn": "treatment"}, + "treatment", + ), + ], +) +def test_proposal_duplicate_errors_identify_the_repeated_measurement( + changes: dict[str, object], repeated: str +) -> None: + with pytest.raises( + ValidationError, match="Comparison columns must be distinct" + ) as caught: + _proposal(**changes) + message = str(caught.value) + assert f"repeated columns ['{repeated}']" in message + assert "A column cannot explain itself" in message + assert "at most three" not in message + + +@pytest.mark.parametrize( + ("invalid_changes", "correction_hint"), + [ + ({"conditionedOn": "age"}, "set conditionedOn=null"), + ({"explanatoryColumns": ["response", "time"]}, "repeated columns ['response']"), + ], +) +def test_design_tool_schema_and_retry_correct_proposals_before_computation( + monkeypatch: pytest.MonkeyPatch, + invalid_changes: dict[str, object], + correction_hint: str, +) -> None: + cells, characterization = _design() + batch_columns = ["batch_a", "batch_b"] + for batch, values in zip(batch_columns, ["treatment", "time"], strict=True): + cells.frame[batch] = cells.frame[values] + cells.columns.append(batch) + characterization.columns.append( + {"name": batch, "kind": "categorical", "domain": "technical"} + ) + deps = _deps(cells) + deps.characterization = characterization + scans: list[dict[str, object]] = [] + safety_columns: list[list[str]] = [] + batch_safety = tools._batch_safety_evidence + + def characterize(*_args: object, **kwargs: object) -> CovariateCharacterization: + scans.append(kwargs) + return characterization + + def record_batch_safety(*args: object, **kwargs: object) -> object: + safety_columns.append(kwargs["batch_columns"]) + return batch_safety(*args, **kwargs) + + monkeypatch.setattr(tools, "characterize_covariates", characterize) + monkeypatch.setattr(tools, "_offered_qc_profiles", lambda *_args: []) + monkeypatch.setattr(tools, "_batch_safety_evidence", record_batch_safety) + requests = 0 + + async def reply(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + nonlocal requests + schema = info.function_tools[0].parameters_json_schema + proposal_schema = schema["$defs"]["CovariateProposal"] + properties = proposal_schema["properties"] + assert ( + "Two explanatory columns plus conditioning are unsupported" + in proposal_schema["description"] + ) + assert ( + "conditionedOn must be null" + in properties["explanatoryColumns"]["description"] + ) + assert "does not count" in properties["observationUnit"]["description"] + assert "does not count" in properties["independentUnit"]["description"] + assert "must differ" in properties["response"]["description"] + request = requests + requests += 1 + if request == 1: + retry_parts = [ + part + for message in messages + for part in message.parts + if isinstance(part, RetryPromptPart) + ] + assert correction_hint in str(retry_parts[-1].content) + assert scans == [] + assert deps.designRounds == 0 + if request < 2: + proposal = _proposal().model_dump() + if request == 0: + proposal.update(invalid_changes) + return ModelResponse( + parts=[ + ToolCallPart( + tool_name="analyze_experimental_design", + args={ + "column_domains": { + batch: "technical" for batch in batch_columns + }, + "coefficients_of_interest": [], + "units_of_inference": {}, + "batch_columns": batch_columns, + "proposals": [proposal], + }, + ) + ] + ) + return ModelResponse(parts=[TextPart("Comparison evidence computed.")]) + + agent = Agent( + FunctionModel(reply), + deps_type=ExperimentalContextDependencies, + tools=[ + Tool( + tools.model_evidence_tool(tools.analyze_experimental_design), + max_retries=3, + ) + ], + ) + result = agent.run_sync("Compare the observed study design.", deps=deps) + assert result.output == "Comparison evidence computed." + assert requests == 3 + assert len(scans) == 1 + assert deps.designRounds == 1 + assert deps.toolCalls == ["analyze_experimental_design"] + assert safety_columns == [batch_columns] + assert len(deps.comparisons) == 1 + assert deps.comparisons[0].status == "computed" + assert deps.comparisons[0].proposal == _proposal() + + +def test_combinations_preserve_typed_values_and_reject_missing() -> None: + cells = _Cells( + pd.DataFrame({"a": ["x,y", "x", 1, "1"], "b": ["z", "y,z", "a", "a"]}) + ) + assert len(np.unique(combination_labels(cells, ["a", "b"]))) == 4 + cells.frame.loc[0, "b"] = None + with pytest.raises(ValueError, match="missing"): + combination_labels(cells, ["a", "b"]) + + +def test_capture_and_reference_selection_requires_verbatim_provenance() -> None: + cells, characterization = _design() + deps = _deps(cells) + deps.studyContext = ( + "sample identifies the physical capture. s0 and s1 are reference captures." + ) + proposal = CaptureProposal( + column="sample", + provenanceQuote="sample identifies the physical capture.", + referenceCaptures=["s0", "s1"], + referenceProvenanceQuote="s0 and s1 are reference captures.", + ) + accept_capture_proposal(deps, characterization, proposal) + assert deps.directions == {} + choices = canonical_design_choices(deps, ExperimentalContextDecision()) + assert choices["physicalCaptureColumn"] == "sample" + assert choices["pooledReferenceCaptures"] == ["s0", "s1"] + with pytest.raises(ValueError, match="exact study quote"): + accept_capture_proposal(_deps(cells), characterization, proposal) + with pytest.raises(ValueError, match="Reference captures require"): + accept_capture_proposal( + deps, + characterization, + proposal.model_copy( + update={"referenceProvenanceQuote": "Made up controls"} + ), + ) + + +def test_final_decision_cannot_invent_a_protected_combination() -> None: + cells, _ = _design() + deps = _deps(cells) + decision = ExperimentalContextDecision( + protectedCombinations=[["time", "treatment"]] + ) + with pytest.raises(ValueError, match="evaluated"): + canonical_design_choices(deps, decision) + + +def test_continuous_protection_is_explicit_without_changing_design_license() -> None: + cells, characterization = _design() + deps = _deps(cells) + deps.characterization = characterization + decision = ExperimentalContextDecision(coefficientsOfInterest=["age"]) + choices = canonical_design_choices(deps, decision) + assert choices["unsupportedProtection"] == ["age"] + assert "batchCorrection" not in choices + + +def test_qc_retention_checks_joint_groups_beside_marginal_groups() -> None: + cells, characterization = _design() + characterization.coefficients = [{"name": "treatment"}, {"name": "time"}] + deps = _deps(cells) + deps.protectedCombinations = [["treatment", "time"]] + keep = ~((cells.fetch("treatment") == "a") & (cells.fetch("time") == "early")) + mito = np.where(keep, np.linspace(1, 3, len(keep)), 80.0) + projection = project_registered_qc_profile( + "globalMad5", + values_by_metric={"RNA_percentMito": mito}, + active=np.ones(len(keep), dtype=bool), + ) + result = qc_evidence._registered_profile_evidence( + projection, + deps=deps, + characterization=characterization, + driver=("RNA", "RNA"), + active=np.ones(len(keep), dtype=bool), + values_by_attr={"RNA_percentMito": mito}, + metadata_attributes=["RNA_percentMito"], + artifact_metrics=[], + metric_sources=[], + source_concordance=[], + sample_column=None, + sample_artifact=None, + capture_column=None, + capture_artifact=None, + capture_labels=None, + pooled_reference_captures=None, + active_cells=len(keep), + comparison_source=None, + ) + assert all( + count > 0 + for counts in result.retainedCellsByColumn.values() + for count in counts.values() + ) + assert any( + count == 0 + for counts in result.retainedCellsByCombination.values() + for count in counts.values() + ) + assert any( + reason.startswith("combination:") for reason in result.unsafeRetentionGroups + ) + + +@pytest.mark.parametrize("independent_units", [5, 8]) +def test_interaction_confounding_blocks_only_the_matching_batch_set( + monkeypatch: pytest.MonkeyPatch, + independent_units: int, +) -> None: + cells, characterization = _design() + cells.frame = cells.frame.iloc[: independent_units * 3].copy() + for record in characterization.columns: + if record["name"] in {"treatment", "time"}: + record["domain"] = "technical" + characterization.coefficients = [ + { + "name": "response", + "kind": "categorical", + "observationUnit": "sample", + "scope": "betweenUnit", + } + ] + characterization.confounding = [ + { + "coefficient": "response", + "observationUnit": "sample", + "pairs": [{"technical": "treatment"}, {"technical": "time"}], + } + ] + deps = _deps(cells) + evaluate_proposals(deps, characterization, [_proposal()]) + if independent_units == 5: + assert deps.comparisons[0].status == "unsupported" + assert deps.comparisons[0].reasons == ["unsupportedJointAssociation"] + assert ( + deps.comparisons[0].evidence["jointGroupEstimability"][ + "coefficientEstimable" + ] + is False + ) + monkeypatch.setattr( + tools, + "reduce_observation_units", + lambda _cells, unit, columns, **_kwargs: cells.frame.groupby(unit).first()[ + columns + ], + ) + joint = tools._batch_safety_evidence( + deps, + characterization, + coefficients=["response"], + batch_columns=["treatment", "time"], + ) + single = tools._batch_safety_evidence( + deps, characterization, coefficients=["response"], batch_columns=["treatment"] + ) + assert joint[0].status == "unsafe" + assert single[0].status == "safe" + + +def test_tool_rejects_oversized_batch_before_scanning( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cells, _ = _design() + deps = _deps(cells) + monkeypatch.setattr( + tools, + "characterize_covariates", + lambda *_args, **_kwargs: pytest.fail("must reject before metadata scan"), + ) + with pytest.raises(ModelRetry, match="eight initial"): + asyncio.run( + tools.analyze_experimental_design( + SimpleNamespace(deps=deps), + column_domains={}, + coefficients_of_interest=[], + units_of_inference={}, + batch_columns=[], + proposals=[_proposal()] * 9, + ) + ) + + +def test_pca_associations_use_covariate_kind_independently_of_role( + monkeypatch: pytest.MonkeyPatch, +) -> None: + values = np.arange(12, dtype=float) + coordinates = np.column_stack([values, np.tile([1.0, -1.0], 6)]) + monkeypatch.setattr( + diagnostics, + "_aligned_metadata_values", + lambda _store, _selection, column: ( + values % 2 if column == "qc_code" else values + ), + ) + support: dict[str, object] = {} + scores = diagnostics._covariate_associations( + None, + None, + coordinates, + ["age", "technical_age", "qc_code"], + ["protected", "technical", "qc"], + {"age": "continuous", "technical_age": "continuous", "qc_code": "categorical"}, + support, + ) + assert scores[0, 0] == pytest.approx(1.0) + assert scores[0, 1] < 0.2 + np.testing.assert_array_equal(scores[0], scores[1]) + assert scores[2, 1] == pytest.approx(1.0) + assert support["age"]["kind"] == "continuous" + + +def test_pca_numeric_association_uses_complete_rows_without_loading_coordinates( + monkeypatch: pytest.MonkeyPatch, +) -> None: + coordinates = np.arange(10, dtype=float).reshape(5, 2) + + class BoundedCoordinates: + shape = coordinates.shape + + def __array__(self, *_args: object, **_kwargs: object) -> np.ndarray: + raise AssertionError("Coordinates must be read in bounded slices") + + def __getitem__(self, rows: slice) -> np.ndarray: + assert rows.stop - rows.start <= 65_536 + return coordinates[rows] + + values = np.asarray([0.0, 1.0, np.nan, 3.0, np.inf]) + monkeypatch.setattr(diagnostics, "_aligned_metadata_values", lambda *_args: values) + support: dict[str, object] = {} + scores = diagnostics._covariate_associations( + None, + None, + BoundedCoordinates(), + ["age"], + ["protected"], + {"age": "continuous"}, + support, + ) + np.testing.assert_allclose(scores, 1.0) + assert support["age"]["completeRows"] == 3 + assert support["age"]["missingRows"] == 2 + + +def test_metric_fingerprint_changes_when_only_missing_mask_changes() -> None: + values = np.asarray([1.0, 2.0, 3.0]) + missing = np.asarray([False, False, False]) + metadata = SimpleNamespace( + N=3, + _get_array=lambda _column: values, + default_block_rows=lambda _column: 2, + _get_missing_mask_array=lambda _column: missing, + ) + first = execution._metadata_column_fingerprint(metadata, "age") + missing[1] = True + assert execution._metadata_column_fingerprint(metadata, "age") != first + missing[1] = False + assert execution._metadata_column_fingerprint(metadata, "age") == first + + +def test_single_cluster_keeps_diagnostics_without_running_marker_contrasts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from scarf.agent.parameter_tuning.contracts import ( + ArtifactRecord, + ParameterCandidateEvaluation, + ) + from scarf.agent.types import ArtifactReferenceModel + + graph = ArtifactRef( + scope="assay", assay="RNA", kind="connectivity_map", artifact_id="a" * 64 + ) + clusters = ArtifactRef( + scope="assay", assay="RNA", kind="cluster_labels", artifact_id="b" * 64 + ) + selection = ArtifactRef( + scope="datastore", kind="cell_selection", artifact_id="c" * 64 + ) + labels = np.zeros(20, dtype=int) + evaluation = ParameterCandidateEvaluation( + candidateId="single", + status="done", + eligible=True, + cellSelection=ArtifactReferenceModel.from_artifact_ref(selection), + artifacts={ + "connectivityMap": ArtifactRecord.from_ref(graph), + "clusters": ArtifactRecord.from_ref(clusters), + }, + ) + monkeypatch.setattr( + diagnostics, + "_selected_feature_names", + lambda *_args: (np.arange(3), np.asarray(["A", "B", "C"])), + ) + monkeypatch.setattr(diagnostics, "_cluster_labels", lambda *_args: labels) + monkeypatch.setattr( + diagnostics, "_subsample_partition_stability", lambda *_args: 1.0 + ) + + def no_markers(*_args, **_kwargs): + raise AssertionError("A single cluster has no marker contrast") + + store = SimpleNamespace( + run_leiden_clustering=lambda *_args, **_kwargs: clusters, + load_graph=lambda *_args: object(), + run_marker_search=no_markers, + cells=SimpleNamespace(columns=[]), + ) + (result,) = diagnostics.augment_cluster_evaluations( + store, + [evaluation], + marker_assay="RNA", + marker_features=selection, + independent_unit_columns=[], + technical_columns=[], + ) + assert not result.eligible + assert result.metrics.markerCoherence is None + assert result.metrics.markerSpecificityMedian is None + assert result.metrics.markerFamilyEnrichment == {} + assert "markerTable" not in result.artifacts + assert result.metrics.seedStability == 1.0 + assert result.metrics.subsampleStability == 1.0 + assert ( + "Marker contrasts require at least two populated clusters" + in result.eligibilityReasons + ) diff --git a/tests/test_agent_diagnostic_accounting.py b/tests/test_agent_diagnostic_accounting.py new file mode 100644 index 00000000..11abf186 --- /dev/null +++ b/tests/test_agent_diagnostic_accounting.py @@ -0,0 +1,332 @@ +"""Diagnostic accounting records calls and verified reuse without inventing work.""" + +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +from scarf.agent.parameter_tuning import diagnostics as d +from scarf.agent.parameter_tuning import execution as e +from scarf.agent.parameter_tuning.contracts import ( + ArtifactRecord, + ParameterCandidateEvaluation, +) +from tests.test_agent_diagnostic_boundaries import _capture_scoring_store, _ref +from tests.test_agent_parameter_tuning import _FakeStore, _dependencies + + +def test_operation_counts_preserve_arguments_results_failures_and_scope() -> None: + value = object() + error = RuntimeError("diagnostic failed") + + def success(argument: Any, *, flag: bool) -> Any: + assert argument is value and flag is False + return value + + def failure() -> Any: + raise error + + assert e.diagnostic_call("core.operation", success, value, flag=False) is value + with e.diagnostic_work() as outer: + assert e.diagnostic_call("core.operation", success, value, flag=False) is value + with e.diagnostic_work() as inner: + with pytest.raises(RuntimeError) as caught: + e.diagnostic_call("core.operation", failure) + assert caught.value is error + e.diagnostic_call("core.operation", success, value, flag=False) + assert inner["core.operation"] == { + "attempted": 1, + "completed": 0, + "failed": 1, + "cacheHits": 0, + "restored": 0, + "artifactReuses": 0, + } + assert outer["core.operation"]["attempted"] == 2 + assert outer["core.operation"]["completed"] == 2 + assert outer["core.operation"]["failed"] == 0 + assert outer["core.operation"]["artifactReuses"] == 0 + e.diagnostic_reuse("outside", "restored") + assert "outside" not in outer + + +def test_failed_metric_is_retried_and_exact_cache_hits_are_counted() -> None: + attempts = [] + value = object() + + def calculate() -> Any: + attempts.append(1) + if len(attempts) == 1: + raise ValueError("metric unavailable") + return value + + with e.diagnostic_work() as counts, e.candidate_metric_cache(): + key = (123, "graph_connectivity", "exact graph", "exact labels") + with pytest.raises(ValueError, match="unavailable"): + e._cached_candidate_metric(key, calculate) + assert e._cached_candidate_metric(key, calculate) is value + assert e._cached_candidate_metric(key, calculate) is value + assert ( + e._cached_candidate_metric((*key, "changed selection"), calculate) is value + ) + assert attempts == [1, 1, 1] + assert counts["metric.graph_connectivity"] == { + "attempted": 3, + "completed": 2, + "failed": 1, + "cacheHits": 1, + "restored": 0, + "artifactReuses": 0, + } + with e.diagnostic_work() as uncached: + e._cached_candidate_metric(key, calculate) + e._cached_candidate_metric(key, calculate) + assert uncached["metric.graph_connectivity"]["completed"] == 2 + assert uncached["metric.graph_connectivity"]["cacheHits"] == 0 + + +@pytest.mark.parametrize("fails", [False, True]) +def test_capture_diagnostic_operations_are_counted_including_partial_failure( + monkeypatch: pytest.MonkeyPatch, fails: bool +) -> None: + labels = np.asarray(["tiny"] * 2 + ["small"] * 3 + ["large"] * 5) + store, selected, _, calls = _capture_scoring_store(monkeypatch, labels) + original = store.run_doublet_detection + attempted = [] + + def score(*args: Any, **kwargs: Any) -> Any: + attempted.append(1) + if fails and len(attempted) == 2: + raise RuntimeError("second capture interrupted") + return original(*args, **kwargs) + + store.run_doublet_detection = score + with e.diagnostic_work() as counts: + if fails: + with pytest.raises(RuntimeError, match="interrupted"): + d.score_advisory_doublets( + store, + selected, + [selected], + assay="RNA", + feature_selection=_ref("feature_selection"), + capture_column="capture", + ) + else: + evidence = d.score_advisory_doublets( + store, + selected, + [selected], + assay="RNA", + feature_selection=_ref("feature_selection"), + capture_column="capture", + ) + assert evidence.capture_coverage == 0.8 + assert counts["core.doubletDetection"] == { + "attempted": 2, + "completed": 1 if fails else 2, + "failed": int(fails), + "cacheHits": 0, + "restored": 0, + "artifactReuses": 0, + } + for name in ("Normalization", "Pca", "Ann", "Neighbors", "Graph", "Partition"): + assert counts[f"core.capture{name}"]["completed"] == 2 + assert counts["core.captureSelection"]["completed"] == (2 if fails else 3) + assert [row[2]["dims"] for row in calls if row[0] == "pca"] == [4, 2] + assert [row[2]["k"] for row in calls if row[0] == "neighbors"] == [4, 2] + + +def test_candidate_reuse_does_not_add_core_calls_or_claim_numerical_rebuilds() -> None: + store = _FakeStore() + deps = _dependencies(store) + identifier = next(iter(deps.candidates)) + with e.diagnostic_work() as counts: + first = e.execute_parameter_candidate(deps, identifier) + assert first.status == "done" + before = {name: dict(row) for name, row in counts.items()} + assert e.execute_parameter_candidate(deps, identifier) is first + assert counts["candidate.evaluation"]["cacheHits"] == 1 + for name, row in before.items(): + assert counts[name] == row + for name in ("pca", "ann", "neighbors", "graph", "partition"): + assert counts[f"core.{name}"]["completed"] == 1 + assert counts[f"core.{name}"]["artifactReuses"] == 0 + + +def test_doublet_restoration_is_recorded_only_after_complete_inventory_validation() -> ( + None +): + evaluation = ParameterCandidateEvaluation( + artifacts={ + "doubletScore:0": ArtifactRecord.from_ref(_ref("quality_metric")), + "doubletCellSelection:0": ArtifactRecord.from_ref(_ref("cell_selection")), + "doubletNativeGraph": ArtifactRecord.from_ref(_ref("connectivity_map")), + "doubletNativeClusters": ArtifactRecord.from_ref(_ref("cluster_labels")), + } + ) + evaluation.metrics.doubletScoreByCapture = {"capture": {"p90": 0.8}} + with e.diagnostic_work() as counts: + evidence = d.restore_advisory_doublets(evaluation, capture_column="capture") + assert len(evidence.scores) == 1 + del evaluation.artifacts["doubletCellSelection:0"] + with pytest.raises(ValueError, match="lacks"): + d.restore_advisory_doublets(evaluation, capture_column="capture") + assert counts == { + "diagnostic.advisoryDoublets": { + "attempted": 0, + "completed": 0, + "failed": 0, + "cacheHits": 0, + "restored": 1, + "artifactReuses": 0, + } + } + + +def test_validated_pca_artifact_reuse_precedes_numerical_arrays( + monkeypatch: pytest.MonkeyPatch, +) -> None: + payload = { + "component_variance": np.ones(2), + "explained_variance_ratio": np.ones(2), + "top_loading_feature_indices": np.tile(np.arange(3), (2, 1)), + "top_loading_values": np.ones((2, 3)), + "family_enrichment": np.ones((0, 2)), + "covariate_association": np.ones((0, 2)), + } + + class Saved(dict): + attrs = {"payload_fingerprint": "validated"} + + saved = Saved(payload) + store = SimpleNamespace( + zw=None, + cells=None, + inspect_artifact=lambda _: SimpleNamespace(parameters={"feat_scaling": True}), + load_artifact=lambda ref: ( + {"data": np.ones((4, 2)), "loadings": np.ones((3, 2))} + if ref.kind == "reduction" + else saved + ), + ) + evaluation = ParameterCandidateEvaluation( + artifacts={ + "pca": ArtifactRecord.from_ref(_ref("reduction")), + "neighbors": ArtifactRecord.from_ref(_ref("neighbors")), + } + ) + monkeypatch.setattr(d, "as_zarr_array", lambda value, **_: value) + monkeypatch.setattr(d, "fingerprint_stored_arrays", lambda *_: "validated") + monkeypatch.setattr( + d, + "plan_artifact", + lambda *_, **__: SimpleNamespace(reused=True, ref=_ref("feature_summary")), + ) + + def forbidden(*args: Any, **kwargs: Any) -> Any: + pytest.fail("Validated PCA diagnostics must bypass numerical arrays") + + for name in ("_component_variance", "_scaled_total_variance", "_top_loadings"): + monkeypatch.setattr(d, name, forbidden) + with e.diagnostic_work() as counts: + result = e.diagnostic_call( + "diagnostic.pca", + d._write_pca_diagnostic, + store, + evaluation, + feature_selection=_ref("feature_selection"), + selected_indices=np.arange(3), + family_masks={}, + covariate_columns=[], + covariate_roles=[], + adjacent_overlap=None, + ) + assert result[0] == _ref("feature_summary") + assert counts == { + "diagnostic.pca": { + "attempted": 1, + "completed": 1, + "failed": 0, + "cacheHits": 0, + "restored": 0, + "artifactReuses": 1, + } + } + + +def test_stability_markers_and_local_metric_reuse_are_separate_operations( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import pandas as pd + from scipy.sparse import block_diag, csr_matrix + from tests.agent_examples import example + + labels = np.asarray([1] * 8 + [2] * 8) + community = csr_matrix(np.ones((8, 8)) - np.eye(8)) + graph = block_diag((community, community), format="csr") + calls = [] + + def partition(*args: Any, **kwargs: Any) -> Any: + calls.append(("partition", kwargs)) + return _ref("cluster_labels", 2) + + def marker(*args: Any, **kwargs: Any) -> Any: + calls.append(("markers", kwargs)) + return _ref("marker_table") + + store = SimpleNamespace( + cells=SimpleNamespace(columns=["donor", "batch"]), + load_graph=lambda _: graph, + run_leiden_clustering=partition, + run_marker_search=marker, + get_markers=lambda *_, **__: pd.DataFrame( + { + "group_id": ["1", "2"], + "feature_name": ["A", "B"], + "score": [1.0, 1.0], + "auc": [0.9, 0.9], + } + ), + ) + monkeypatch.setattr(d, "_cluster_labels", lambda *_: labels) + monkeypatch.setattr( + d, "_selected_feature_names", lambda *_: (np.arange(2), np.asarray(["A", "B"])) + ) + monkeypatch.setattr(d, "_aligned_metadata", lambda *_: np.tile(["d1", "d2"], 8)) + evaluation = example(ParameterCandidateEvaluation) + evaluation.artifacts.update( + { + "clusters": ArtifactRecord.from_ref(_ref("cluster_labels")), + "connectivityMap": ArtifactRecord.from_ref(_ref("connectivity_map")), + } + ) + with e.diagnostic_work() as counts, e.candidate_metric_cache(): + results = [ + d.augment_cluster_evaluations( + store, + [evaluation], + marker_assay="RNA", + marker_features=_ref("feature_selection"), + independent_unit_columns=["donor"], + technical_columns=["batch"], + )[0] + for _ in range(2) + ] + assert results[0].metrics == results[1].metrics + assert results[0].metrics.subsampleStability == pytest.approx(1.0) + assert counts["metric.subsample_stability"]["completed"] == 1 + assert counts["metric.subsample_stability"]["cacheHits"] == 1 + assert counts["core.subsamplePartition"]["completed"] == 1 + assert counts["core.alternateSeedPartition"]["completed"] == 2 + assert counts["core.markers"]["completed"] == 2 + assert counts["core.markers"]["artifactReuses"] == 0 + assert counts["metric.crossUnitSupport"]["completed"] == 2 + assert counts["metric.technicalAssociation"]["completed"] == 2 + assert all(row[1]["random_seed"] == 9173 for row in calls if row[0] == "partition") + assert all( + row[1]["features"] == _ref("feature_selection") + for row in calls + if row[0] == "markers" + ) diff --git a/tests/test_agent_diagnostic_boundaries.py b/tests/test_agent_diagnostic_boundaries.py new file mode 100644 index 00000000..52036e26 --- /dev/null +++ b/tests/test_agent_diagnostic_boundaries.py @@ -0,0 +1,850 @@ +"""Numerical evidence remains bounded and rejects mismatched scientific inputs.""" + +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +from scarf.agent.parameter_tuning import diagnostics as d +from scarf.agent.parameter_tuning.contracts import ( + ArtifactRecord, + ParameterCandidateEvaluation, +) +from scarf.storage.refs import ArtifactRef +from tests.agent_examples import example + + +class _Blocks: + def __init__(self, values: np.ndarray, limit: int = 65_536) -> None: + self.values = values + self.shape, self.dtype = values.shape, values.dtype + self.limit = limit + self.reads: list[int] = [] + + def __getitem__(self, rows: slice) -> np.ndarray: + output = self.values[rows] + assert len(output) <= self.limit + self.reads.append(len(output)) + return output + + def __array__(self, *_args: Any, **_kwargs: Any) -> Any: + raise AssertionError("Scientific arrays must use bounded reads") + + +def _ref(kind: str, value: int = 1) -> ArtifactRef: + return ArtifactRef( + "datastore" if kind == "cell_selection" else "assay", + kind, + f"{value:064x}", + None if kind == "cell_selection" else "RNA", + ) + + +def test_streamed_component_and_loading_statistics_match_dense_reference() -> None: + rng = np.random.default_rng(7) + coordinates = _Blocks(rng.normal(size=(70_000, 3))) + np.testing.assert_allclose( + d._component_variance(coordinates), coordinates.values.var(axis=0), rtol=1e-12 + ) + assert coordinates.reads == [65_536, 4464] + loadings = _Blocks(rng.normal(size=(9000, 3)), 8192) + indices = np.arange(9000) * 2 + family = np.arange(9000) % 3 == 0 + top, magnitude, enrichment = d._top_loadings(loadings, indices, {"family": family}) + expected = np.argsort(-np.abs(loadings.values), axis=0, kind="stable")[:20].T + np.testing.assert_array_equal(top, indices[expected]) + for column in range(3): + np.testing.assert_allclose( + magnitude[column], np.abs(loadings.values[expected[column], column]) + ) + assert enrichment[0, column] == pytest.approx( + family[expected[column]].mean() / family.mean() + ) + assert loadings.reads == [8192, 808] + + +@pytest.mark.parametrize( + "values", [np.ones(3), np.ones((0, 2)), np.asarray([[1.0, np.nan]])] +) +def test_pca_variance_cannot_come_from_invalid_coordinates(values: np.ndarray) -> None: + with pytest.raises(ValueError): + d._component_variance(values) + + +@pytest.mark.parametrize("damage", ["shape", "empty", "familyMask", "nonfinite"]) +def test_loading_programs_require_exact_finite_feature_axis(damage: str) -> None: + values = np.ones((3, 2)) + indices = np.arange(3) + family = np.ones(3, dtype=bool) + if damage == "shape": + values = values[:2] + elif damage == "empty": + values = values[:, :0] + elif damage == "familyMask": + family = family[:2] + else: + values[0, 0] = np.inf + with pytest.raises(ValueError): + d._top_loadings(values, indices, {"family": family}) + + +def test_family_influence_distinguishes_representation_programs_without_removing_genes() -> ( + None +): + names = np.asarray( + [ + "MT-CO1", + "MTOR", + "RPS3", + "MRPS3", + "MKI67", + "HBA1", + "IGHM", + "FOS", + "ATF3", + "XIST", + "OTHER", + ] + ) + expected = { + "mitochondrial": {0}, + "ribosomal": {2}, + "mitoribosomal": {3}, + "cellCycle": {4}, + "hemoglobin": {5}, + "immuneReceptor": {6}, + "stress": {7}, + "dissociation": {7, 8}, + "sex": {9}, + } + for family, positions in expected.items(): + assert set(np.flatnonzero(d._family_mask(names, family))) == positions + assert d._family_mask(names, "unregisteredFamily") is None + np.testing.assert_array_equal( + names, + [ + "MT-CO1", + "MTOR", + "RPS3", + "MRPS3", + "MKI67", + "HBA1", + "IGHM", + "FOS", + "ATF3", + "XIST", + "OTHER", + ], + ) + + +def _variance_inputs( + monkeypatch: pytest.MonkeyPatch, +) -> tuple[Any, Any, ArtifactRef, dict[str, Any]]: + normalized, cells, features = ( + _ref("normalized"), + _ref("cell_selection"), + _ref("feature_selection"), + ) + data = _Blocks( + np.column_stack([np.arange(9000), np.ones(9000), np.arange(9000) % 3]).astype( + float + ), + 8192, + ) + payload: dict[str, Any] = {"data": data} + inputs = {"cell_selection": cells, "feature_selection": features} + store = SimpleNamespace( + inspect_artifact=lambda _: SimpleNamespace(inputs=inputs), + load_artifact=lambda _: payload, + ) + status = SimpleNamespace( + inputs={"normalized": normalized, "pca_cell_selection": cells} + ) + monkeypatch.setattr(d, "as_zarr_array", lambda value, **_: value) + return store, status, features, payload + + +def test_scaled_pca_variance_reuses_saved_feature_summaries( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store, status, features, payload = _variance_inputs(monkeypatch) + + def evaluate() -> float: + return d._scaled_total_variance( + store, status, n_rows=9000, n_features=3, feature_selection=features + ) + + assert evaluate() == 2 + assert payload["data"].reads == [8192, 808] + values = payload["data"].values + payload.update( + feature_sum=values.sum(axis=0), + feature_squared_sum=np.square(values).sum(axis=0), + ) + payload["data"].reads.clear() + assert evaluate() == 2 + assert payload["data"].reads == [] + + +@pytest.mark.parametrize( + "damage", + [ + "missingInput", + "differentCells", + "differentGenes", + "differentShape", + "nonfinite", + "invalidSummaries", + "constantGenes", + ], +) +def test_scaled_variance_denominator_requires_exact_normalization( + monkeypatch: pytest.MonkeyPatch, damage: str +) -> None: + store, status, features, payload = _variance_inputs(monkeypatch) + if damage == "missingInput": + status.inputs.pop("normalized") + elif damage == "differentCells": + status.inputs["pca_cell_selection"] = _ref("cell_selection", 2) + elif damage == "differentGenes": + features = _ref("feature_selection", 2) + elif damage == "differentShape": + payload["data"].shape = (9000, 4) + elif damage == "nonfinite": + payload["data"].values[0, 0] = np.nan + elif damage == "invalidSummaries": + payload.update(feature_sum=np.ones(2), feature_squared_sum=np.ones(2)) + else: + payload["data"].values[:] = 1 + with pytest.raises(ValueError): + d._scaled_total_variance( + store, status, n_rows=9000, n_features=3, feature_selection=features + ) + + +@pytest.mark.parametrize("damage", ["roles", "rows", "kind", "artifactKind"]) +def test_pca_covariate_evidence_rejects_wrong_type_or_selection( + monkeypatch: pytest.MonkeyPatch, damage: str +) -> None: + monkeypatch.setattr( + d, + "_aligned_metadata_values", + lambda *_: np.arange(3 if damage == "rows" else 4), + ) + monkeypatch.setattr( + d, + "resolve_cell_aligned_artifact", + lambda *_args, **_kwargs: SimpleNamespace(values=np.arange(4)), + ) + with pytest.raises(ValueError): + d._covariate_associations( + SimpleNamespace(zw=None), + _ref("cell_selection"), + np.ones((4, 2)), + ["covariate"], + [] if damage == "roles" else ["technical"], + column_kinds={ + "covariate": "unknown" if damage == "kind" else "categorical" + }, + column_artifacts={"covariate": _ref("quality_metric")} + if damage == "artifactKind" + else {}, + ) + + +def test_missing_categorical_covariates_have_zero_matched_support_without_coordinate_reads() -> ( + None +): + coordinates = _Blocks(np.ones((4, 2))) + np.testing.assert_array_equal( + d._categorical_association(coordinates, np.asarray([None] * 4)), [0, 0] + ) + assert coordinates.reads == [] + + +def test_doublet_summary_uses_bounded_scores_and_preserves_extremes() -> None: + values = _Blocks(np.linspace(0, 1, 70_000)) + summary, sample = d._bounded_score_summary(values, maximum_sample_size=1000) + assert len(sample) <= 1000 and summary["sampleSize"] == len(sample) + assert summary["minimum"] == 0 and summary["maximum"] == 1 + assert summary["p90"] == pytest.approx(np.quantile(sample, 0.9)) + assert values.reads == [65_536, 4464] + + +@pytest.mark.parametrize( + "values,limit", + [ + (np.ones(3), 0), + (np.ones((2, 2)), 10), + (np.ones(0), 10), + (np.asarray([1.0, np.nan]), 10), + ], +) +def test_doublet_summary_cannot_hide_invalid_scores( + values: np.ndarray, limit: int +) -> None: + with pytest.raises(ValueError): + d._bounded_score_summary(values, maximum_sample_size=limit) + + +@pytest.mark.parametrize("damage", ["gap", "captureSummary", "selection"]) +def test_resumed_doublets_require_complete_capture_inventory(damage: str) -> None: + evaluation = ParameterCandidateEvaluation( + artifacts={"doubletScore:0": ArtifactRecord.from_ref(_ref("quality_metric"))} + ) + evaluation.metrics.doubletScoreByCapture = {"capture": {"p90": 0.8}} + if damage == "gap": + evaluation.artifacts["doubletScore:1"] = evaluation.artifacts.pop( + "doubletScore:0" + ) + elif damage == "captureSummary": + evaluation.metrics.doubletScoreByCapture = {} + with pytest.raises(ValueError): + d.restore_advisory_doublets(evaluation, capture_column="capture") + + +def _score_evidence( + monkeypatch: pytest.MonkeyPatch, +) -> tuple[Any, Any, dict[ArtifactRef, np.ndarray]]: + parent, selected, score = ( + _ref("cell_selection"), + _ref("cell_selection", 2), + _ref("quality_metric"), + ) + selections = {parent: np.arange(4), selected: np.arange(4)} + store = SimpleNamespace( + zw=None, load_artifact=lambda _: {"values": np.asarray([0.1, 0.2, 0.8, 0.9])} + ) + monkeypatch.setattr(d, "as_zarr_array", lambda value, **_: value) + monkeypatch.setattr( + d, "read_stored_selection_indices", lambda _root, ref, **_: selections[ref] + ) + evidence = d.AdvisoryDoubletScores( + (score,), + (selected,), + _ref("connectivity_map"), + _ref("cluster_labels"), + score_summaries=({"p90": 0.75},), + ) + return store, evidence, selections + + +@pytest.mark.parametrize( + "damage", + [ + "labels", + "parentOrder", + "summaryCount", + "scoreShape", + "outsideParent", + "overlap", + "nonfinite", + ], +) +def test_doublet_concentration_cannot_compare_different_or_overlapping_cells( + monkeypatch: pytest.MonkeyPatch, damage: str +) -> None: + from dataclasses import replace + + store, evidence, selections = _score_evidence(monkeypatch) + labels = np.asarray([1, 1, 2, 2]) + if damage == "labels": + labels = labels[:3] + elif damage == "parentOrder": + selections[_ref("cell_selection")] = np.asarray([1, 0, 2, 3]) + elif damage == "summaryCount": + evidence = replace(evidence, score_summaries=({"p90": 0.5}, {"p90": 0.5})) + elif damage == "scoreShape": + store.load_artifact = lambda _: {"values": np.ones(3)} + elif damage == "outsideParent": + selections[_ref("cell_selection", 2)] = np.asarray([0, 1, 2, 4]) + elif damage == "overlap": + evidence = replace( + evidence, + scores=evidence.scores * 2, + cell_selections=evidence.cell_selections * 2, + score_summaries=evidence.score_summaries * 2, + ) + else: + store.load_artifact = lambda _: {"values": np.asarray([np.nan, 0, 0, 1])} + with pytest.raises(ValueError): + d._doublet_concentration(store, labels, _ref("cell_selection"), evidence) + + +def test_doublet_concentration_uses_matched_capture_thresholds_and_empty_evidence_is_unavailable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from dataclasses import replace + + store, evidence, _ = _score_evidence(monkeypatch) + assert ( + d._doublet_concentration( + store, np.asarray([1, 1, 2, 2]), _ref("cell_selection"), evidence + ) + == 2 + ) + missing = replace(evidence, scores=(), cell_selections=(), score_summaries=()) + assert ( + d._doublet_concentration( + store, np.asarray([1, 1, 2, 2]), _ref("cell_selection"), missing + ) + is None + ) + + +def test_native_doublet_reconstruction_uses_pca_and_preserves_selected_parameters() -> ( + None +): + selected = example(ParameterCandidateEvaluation) + selected.parameters.useHarmony = True + selected.artifacts["pca"] = ArtifactRecord.from_ref(_ref("reduction")) + pca = d._artifact_ref(selected, "pca") + calls = [] + + def operation(name: str, kind: str) -> Any: + def execute(*args: Any, **kwargs: Any) -> ArtifactRef: + calls.append((name, args, kwargs)) + return _ref(kind) + + return execute + + store = SimpleNamespace( + build_ann_index=operation("ann", "ann_index"), + query_neighbors=operation("neighbors", "neighbors"), + build_connectivity_map=operation("graph", "connectivity_map"), + run_leiden_clustering=operation("clusters", "cluster_labels"), + ) + clusters, graph = d.resolve_native_doublet_inputs(store, selected, []) + assert [row[0] for row in calls] == ["ann", "neighbors", "graph", "clusters"] + assert calls[0][1] == (pca,) + assert calls[1][2]["k"] == selected.parameters.neighborsK + assert calls[-1][2]["resolution"] == selected.parameters.leidenResolution + assert clusters == _ref("cluster_labels") and graph == _ref("connectivity_map") + + +def _capture_scoring_store( + monkeypatch: pytest.MonkeyPatch, labels: np.ndarray, *, capture_known: bool = True +) -> tuple[Any, Any, dict[ArtifactRef, np.ndarray], list[Any]]: + parent = _ref("cell_selection") + selections = {parent: np.arange(len(labels))} + payloads: dict[ArtifactRef, dict[str, np.ndarray]] = {} + calls = [] + current = [parent] + serial = [10] + + def reference(kind: str) -> ArtifactRef: + serial[0] += 1 + return _ref(kind, serial[0]) + + def record(name: str, kind: str) -> Any: + def run(*args: Any, **kwargs: Any) -> ArtifactRef: + calls.append((name, args, kwargs)) + if name == "normalize": + current[0] = args[0] + return reference(kind) + + return run + + def filter_cells( + columns: Any, lower: Any, upper: Any, **kwargs: Any + ) -> ArtifactRef: + assert columns == ["capture"] and lower == upper + assert kwargs["cell_selection"] == parent + mask = labels == lower[0] + selected = reference("cell_selection") + selections[selected] = np.flatnonzero(mask) + payloads[selected] = {"values": mask} + calls.append(("filter", (), {"capture": lower[0]})) + return selected + + def score(*args: Any, **kwargs: Any) -> ArtifactRef: + calls.append(("doublet", args, kwargs)) + result = reference("quality_metric") + payloads[result] = {"values": np.linspace(0, 1, len(selections[current[0]]))} + return result + + store = SimpleNamespace( + zw=None, + cells=SimpleNamespace( + N=len(labels), columns=["capture"] if capture_known else [] + ), + get_assay=lambda _: SimpleNamespace( + feats=SimpleNamespace( + fetch_all=lambda _: np.asarray(["g1", "g2", "g3", "g4", "g5"]) + ) + ), + load_artifact=lambda ref: payloads[ref], + filter_cells=filter_cells, + run_normalization=record("normalize", "normalized"), + run_pca=record("pca", "reduction"), + build_ann_index=record("ann", "ann_index"), + query_neighbors=record("neighbors", "neighbors"), + build_connectivity_map=record("graph", "connectivity_map"), + run_leiden_clustering=record("cluster", "cluster_labels"), + run_doublet_detection=score, + ) + monkeypatch.setattr(d, "as_zarr_array", lambda value, **_: value) + monkeypatch.setattr( + d, "read_stored_selection_indices", lambda _root, ref, **_: selections[ref] + ) + monkeypatch.setattr( + d, "read_metadata_rows_chunkwise", lambda _cells, _column, rows: labels[rows] + ) + monkeypatch.setattr(d, "read_feature_selection_indices", lambda *_: np.arange(5)) + monkeypatch.setattr( + d, + "resolve_native_doublet_inputs", + lambda *_: (_ref("cluster_labels"), _ref("connectivity_map")), + ) + selected = example(ParameterCandidateEvaluation) + from scarf.agent.tools import artifact_reference + + selected.cellSelection = artifact_reference(parent) + selected.parameters.dimensions = 21 + selected.parameters.neighborsK = 11 + return store, selected, selections, calls + + +def test_capture_doublet_routing_caps_rank_and_preserves_unscored_small_capture( + monkeypatch: pytest.MonkeyPatch, +) -> None: + labels = np.asarray(["tiny"] * 2 + ["small"] * 3 + ["large"] * 5) + store, selected, selections, calls = _capture_scoring_store(monkeypatch, labels) + evidence = d.score_advisory_doublets( + store, + selected, + [selected], + assay="RNA", + feature_selection=_ref("feature_selection"), + capture_column="capture", + ) + assert evidence.capture_values == ("large", "small") + assert evidence.capture_coverage == 0.8 + assert any("tiny" in note and "only 2" in note for note in evidence.limitations) + assert [row[2]["dims"] for row in calls if row[0] == "pca"] == [4, 2] + assert [row[2]["k"] for row in calls if row[0] == "neighbors"] == [4, 2] + assert sum(row[0] == "doublet" for row in calls) == 2 + np.testing.assert_array_equal(selections[_ref("cell_selection")], np.arange(10)) + for ref, capture in zip( + evidence.cell_selections, evidence.capture_values, strict=True + ): + assert np.all(labels[selections[ref]] == capture) + + +@pytest.mark.parametrize("capture_known", [False, True]) +def test_single_or_unknown_capture_uses_one_native_reference_with_visible_scope( + monkeypatch: pytest.MonkeyPatch, capture_known: bool +) -> None: + store, selected, _, calls = _capture_scoring_store( + monkeypatch, np.asarray(["one"] * 5), capture_known=capture_known + ) + evidence = d.score_advisory_doublets( + store, + selected, + [selected], + assay="RNA", + feature_selection=_ref("feature_selection"), + capture_column="capture" if capture_known else None, + ) + assert len(evidence.scores) == 1 and evidence.capture_coverage == 1 + assert [row[0] for row in calls] == ["doublet"] + assert evidence.capture_values == ("one" if capture_known else "allSelectedCells",) + assert bool(evidence.limitations) == (not capture_known) + + +@pytest.mark.parametrize( + "damage", ["missingSelection", "tooManyCaptures", "allCapturesTooSmall"] +) +def test_doublets_fail_visibly_when_scoring_scope_is_unavailable( + monkeypatch: pytest.MonkeyPatch, damage: str +) -> None: + labels = ( + np.arange(513).astype(str) + if damage == "tooManyCaptures" + else np.asarray(["a", "a", "b", "b"]) + ) + store, selected, _, calls = _capture_scoring_store(monkeypatch, labels) + if damage == "missingSelection": + selected.cellSelection = None + with pytest.raises(ValueError): + d.score_advisory_doublets( + store, + selected, + [selected], + assay="RNA", + feature_selection=_ref("feature_selection"), + capture_column="capture", + ) + assert not any(row[0] == "doublet" for row in calls) + + +@pytest.mark.parametrize( + "damage", ["inventory", "emptyParent", "emptyScores", "unalignedScores"] +) +def test_doublet_summary_requires_aligned_nonempty_capture_artifacts( + monkeypatch: pytest.MonkeyPatch, damage: str +) -> None: + store, evidence, selections = _score_evidence(monkeypatch) + scores = evidence.scores + if damage == "inventory": + scores = () + elif damage == "emptyParent": + selections[_ref("cell_selection")] = np.asarray([], dtype=int) + elif damage == "emptyScores": + store.load_artifact = lambda _: {"values": np.ones(0)} + else: + store.load_artifact = lambda _: {"values": np.ones(3)} + with pytest.raises(ValueError): + d._build_advisory_doublet_scores( + store, + scores=scores, + cell_selections=evidence.cell_selections, + native_graph=evidence.native_graph, + native_clusters=evidence.native_clusters, + parent_selection=_ref("cell_selection"), + capture_values=("one",), + capture_column="capture", + limitations=(), + ) + + +def test_large_doublet_summary_identifies_sampled_quantiles( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store, evidence, selections = _score_evidence(monkeypatch) + selections[_ref("cell_selection")] = np.arange(70_000) + selections[evidence.cell_selections[0]] = np.arange(70_000) + values = _Blocks(np.linspace(0, 1, 70_000)) + store.load_artifact = lambda _: {"values": values} + result = d._build_advisory_doublet_scores( + store, + scores=evidence.scores, + cell_selections=evidence.cell_selections, + native_graph=evidence.native_graph, + native_clusters=evidence.native_clusters, + parent_selection=_ref("cell_selection"), + capture_values=("one",), + capture_column="capture", + limitations=(), + ) + assert result.capture_coverage == 1 + assert result.score_summaries[0]["sampleSize"] < 70_000 + assert any("deterministic bounded samples" in note for note in result.limitations) + + +def test_metadata_diagnostics_preserve_missing_masks_and_reject_misalignment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + d, "read_stored_selection_indices", lambda *_args, **_kwargs: np.asarray([1, 3]) + ) + monkeypatch.setattr( + d, "read_metadata_rows_chunkwise", lambda *_: np.asarray([20.0, 40]) + ) + monkeypatch.setattr( + d, "read_metadata_missing_rows_chunkwise", lambda *_: np.asarray([False, True]) + ) + store = SimpleNamespace(zw=None, cells=None) + values = d._aligned_metadata_values(store, _ref("cell_selection"), "age") + assert values.tolist() == [20.0, None] + monkeypatch.setattr( + d, "read_metadata_rows_chunkwise", lambda *_: np.asarray([20.0]) + ) + for read in (d._aligned_metadata_values, d._aligned_metadata): + with pytest.raises(ValueError, match="align"): + read(store, _ref("cell_selection"), "age") + + +def test_topology_overlap_requires_same_neighborhood_axis( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(d, "as_zarr_array", lambda value, **_: value) + left = _ref("neighbors") + store = SimpleNamespace( + load_artifact=lambda ref: { + "indices": np.ones((4, 2) if ref == left else (3, 2), dtype=int) + } + ) + with pytest.raises(ValueError, match="align"): + d._neighbor_overlap(store, left, _ref("neighbors", 2)) + + +@pytest.mark.parametrize("damage", ["artifactKind", "scaling", "loadingShape"]) +def test_pca_diagnostic_rejects_unmatched_method_or_feature_axis_before_writing( + monkeypatch: pytest.MonkeyPatch, damage: str +) -> None: + evaluation = ParameterCandidateEvaluation( + artifacts={ + "pca": ArtifactRecord.from_ref(_ref("reduction")), + "neighbors": ArtifactRecord.from_ref(_ref("neighbors")), + } + ) + store = SimpleNamespace( + inspect_artifact=lambda _: SimpleNamespace( + parameters={"feat_scaling": damage != "scaling"} + ), + load_artifact=lambda _: {"data": np.ones((4, 2)), "loadings": np.ones((4, 2))}, + ) + monkeypatch.setattr(d, "as_zarr_array", lambda value, **_: value) + with pytest.raises(ValueError): + d._write_pca_diagnostic( + store, + evaluation, + feature_selection=_ref("feature_selection"), + selected_indices=np.arange(3), + family_masks={}, + covariate_columns=[], + covariate_roles=[], + adjacent_overlap=None, + column_artifacts={"counts": _ref("cluster_labels")} + if damage == "artifactKind" + else {}, + ) + + +def test_reused_pca_evidence_rejects_a_changed_saved_payload( + tmp_path: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + import zarr + from scarf.storage.artifacts import fingerprint_stored_arrays + + root = zarr.open_group(str(tmp_path / "diagnostics.zarr"), mode="w") + pca = root.create_group("pca") + pca.create_array("data", data=np.ones((4, 2))) + pca.create_array("loadings", data=np.ones((3, 2))) + saved = root.create_group("saved") + payload = { + "component_variance": np.ones(2), + "explained_variance_ratio": np.ones(2), + "top_loading_feature_indices": np.tile(np.arange(3), (2, 1)), + "top_loading_values": np.ones((2, 3)), + "family_enrichment": np.ones((0, 2)), + "covariate_association": np.ones((0, 2)), + "adjacent_neighbor_overlap": np.ones(1), + } + for name, values in payload.items(): + saved.create_array(name, data=values) + saved.attrs["payload_fingerprint"] = fingerprint_stored_arrays( + saved, d._PCA_DIAGNOSTIC_ARRAYS + ) + saved["component_variance"][0] = 2 + evaluation = ParameterCandidateEvaluation( + artifacts={ + "pca": ArtifactRecord.from_ref(_ref("reduction")), + "neighbors": ArtifactRecord.from_ref(_ref("neighbors")), + } + ) + store = SimpleNamespace( + zw=root, + cells=None, + inspect_artifact=lambda _: SimpleNamespace(parameters={"feat_scaling": True}), + load_artifact=lambda ref: pca if ref.kind == "reduction" else saved, + ) + monkeypatch.setattr( + d, + "plan_artifact", + lambda *_args, **_kwargs: SimpleNamespace( + reused=True, ref=_ref("feature_summary") + ), + ) + with pytest.raises(ValueError, match="fingerprint"): + d._write_pca_diagnostic( + store, + evaluation, + feature_selection=_ref("feature_selection"), + selected_indices=np.arange(3), + family_masks={}, + covariate_columns=[], + covariate_roles=[], + adjacent_overlap=None, + ) + + +def test_incomplete_candidates_keep_their_failure_without_running_new_diagnostics( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + d, + "_selected_feature_names", + lambda *_: (np.arange(3), np.asarray(["a", "b", "c"])), + ) + failed = ParameterCandidateEvaluation( + status="failed", error="PCA evidence unavailable" + ) + store = SimpleNamespace(cells=SimpleNamespace(columns=[])) + pca = d.augment_pca_evaluations( + store, + [failed], + feature_selection=_ref("feature_selection"), + nominated_families=[], + protected_families=[], + technical_columns=[], + protected_columns=[], + qc_columns=[], + ) + clusters = d.augment_cluster_evaluations( + store, + pca, + marker_assay="RNA", + marker_features=_ref("feature_selection"), + independent_unit_columns=[], + technical_columns=[], + ) + assert clusters[0].status == "failed" + assert clusters[0].error == "PCA evidence unavailable" + assert clusters[0].artifacts == {} + + +def test_tiny_population_stability_reclusters_the_entire_supported_graph() -> None: + from scipy.sparse import csr_matrix + + graph = csr_matrix(np.ones((3, 3)) - np.eye(3)) + assert d._subsample_partition_stability(graph, np.ones(3), 0.5) == 1 + + +@pytest.mark.parametrize("damage", ["labelMatrix", "seedSelection"]) +def test_cluster_diagnostics_reject_unmatched_partition_evidence_before_markers( + monkeypatch: pytest.MonkeyPatch, damage: str +) -> None: + from scarf.agent.tools import artifact_reference + + monkeypatch.setattr( + d, + "_selected_feature_names", + lambda *_: (np.arange(3), np.asarray(["a", "b", "c"])), + ) + monkeypatch.setattr(d, "as_zarr_array", lambda value, **_: value) + selected = ParameterCandidateEvaluation( + status="done", + eligible=True, + cellSelection=artifact_reference(_ref("cell_selection")), + artifacts={ + "clusters": ArtifactRecord.from_ref(_ref("cluster_labels")), + "connectivityMap": ArtifactRecord.from_ref(_ref("connectivity_map")), + }, + ) + initial = np.ones((3, 2)) if damage == "labelMatrix" else np.asarray([1, 1, 2]) + store = SimpleNamespace( + load_artifact=lambda ref: { + "values": initial if ref == _ref("cluster_labels") else np.asarray([1, 2]) + }, + run_leiden_clustering=lambda *_args, **_kwargs: _ref("cluster_labels", 2), + ) + with pytest.raises(ValueError, match="one-dimensional|align"): + d.augment_cluster_evaluations( + store, + [selected], + marker_assay="RNA", + marker_features=_ref("feature_selection"), + independent_unit_columns=[], + technical_columns=[], + ) + + +def test_pca_feature_evidence_requires_an_assay_owned_selection() -> None: + with pytest.raises(ValueError, match="belong to one assay"): + d._selected_feature_names( + None, ArtifactRef("datastore", "feature_selection", "a" * 64) + ) diff --git a/tests/test_agent_diagnostic_journal.py b/tests/test_agent_diagnostic_journal.py new file mode 100644 index 00000000..3260a7fe --- /dev/null +++ b/tests/test_agent_diagnostic_journal.py @@ -0,0 +1,183 @@ +"""Observed diagnostic calls and reuse remain honest across failed attempts.""" + +import pytest + +from scarf.agent.orchestrator import journal +from scarf.agent.parameter_tuning.execution import diagnostic_call, diagnostic_reuse +from tests.agent_journal_store import memory_journal + + +@pytest.mark.parametrize("interrupted", [False, True]) +def test_failed_diagnostic_attempt_preserves_actual_calls_and_partial_work( + interrupted: bool, +) -> None: + store, prefix, request = memory_journal("analysis") + error = ( + KeyboardInterrupt("Interrupted diagnostic") + if interrupted + else ValueError("Capture evidence failed") + ) + calls = [] + + def score(capture: str) -> str: + calls.append(capture) + if capture == "bad": + raise error + return "the-same-core-artifact" + + with pytest.raises(type(error)) as caught: + with journal.diagnostic_attempt( + store, prefix, request.workflowRunId, {"cells": "frozen"} + ) as counts: + assert ( + diagnostic_call("core.doubletDetection", score, "a") + == "the-same-core-artifact" + ) + assert ( + diagnostic_call("core.doubletDetection", score, "a") + == "the-same-core-artifact" + ) + diagnostic_reuse("metric.batch", "cacheHits") + diagnostic_call("core.doubletDetection", score, "bad") + assert caught.value is error + assert calls == ["a", "a", "bad"] + assert counts["core.doubletDetection"]["attempted"] == 3 + assert counts["core.doubletDetection"]["completed"] == 2 + assert counts["core.doubletDetection"]["failed"] == 1 + state = journal.analysis_snapshot(store, request.workflowRunId) + attempt = state["diagnosticAttempts"][0] + assert attempt["finished"]["status"] == ("interrupted" if interrupted else "failed") + assert attempt["finished"]["operations"] == counts + assert "not observable" in attempt["finished"]["interpretation"] + assert "unknown operation counts" in attempt["finished"]["interpretation"] + assert attempt["started"]["recordedAtNs"] <= attempt["finished"]["recordedAtNs"] + + with journal.diagnostic_attempt( + store, prefix, request.workflowRunId, {"cells": "frozen"} + ): + diagnostic_reuse("diagnostic.augmentedEvaluation", "restored") + resumed = journal.analysis_snapshot(store, request.workflowRunId) + assert resumed["diagnosticAttempts"][0] == attempt + assert len(resumed["diagnosticAttempts"]) == 2 + restored = resumed["diagnosticAttempts"][1]["finished"]["operations"] + assert "core.doubletDetection" not in restored + assert restored["diagnostic.augmentedEvaluation"]["restored"] == 1 + + +def test_unfinished_and_legacy_attempts_do_not_manufacture_zero_computation() -> None: + store, prefix, request = memory_journal("analysis") + assert ( + journal.analysis_snapshot(store, request.workflowRunId)["diagnosticAttempts"] + == [] + ) + journal.save_checkpoint( + store, + prefix, + request.workflowRunId, + "parameter_tuning/diagnostic_attempts/interrupted/started", + {"cells": "frozen"}, + {"recordedAtNs": 1}, + ) + attempt = journal.analysis_snapshot(store, request.workflowRunId)[ + "diagnosticAttempts" + ][0] + assert "finished" not in attempt + assert "operations" not in attempt["started"] + journal.save_checkpoint( + store, + prefix, + request.workflowRunId, + "parameter_tuning/diagnostic_attempts/interrupted/unrecognized", + {"cells": "frozen"}, + {}, + ) + with pytest.raises(ValueError, match="Unknown diagnostic attempt"): + journal.analysis_snapshot(store, request.workflowRunId) + + +@pytest.mark.parametrize("missing_start", [False, True]) +def test_diagnostic_finish_requires_the_exact_original_start( + missing_start: bool, +) -> None: + store, prefix, request = memory_journal("analysis") + key = "parameter_tuning/diagnostic_attempts/attempt" + if not missing_start: + journal.save_checkpoint( + store, + prefix, + request.workflowRunId, + key + "/started", + {"cells": "original"}, + {"recordedAtNs": 1}, + ) + journal.save_checkpoint( + store, + prefix, + request.workflowRunId, + key + "/finished", + {"cells": "different"}, + {"recordedAtNs": 2, "status": "completed", "operations": {}}, + ) + with pytest.raises( + ValueError, match="matching start" if missing_start else "inputs changed" + ): + journal.analysis_snapshot(store, request.workflowRunId) + + +def test_primary_evidence_restore_and_review_replay_count_no_new_diagnostic_calls( + monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> None: + from scarf.agent.parameter_tuning.execution import diagnostic_work + from scarf.agent.orchestrator import rna_tuning + from tests.test_agent_rna_evidence_mode import make_run, assess + + request.getfixturevalue("memory_checkpoints") + run, selected = make_run(monkeypatch, object()) + monkeypatch.setattr(rna_tuning, "run_agent_sync", assess) + with diagnostic_work() as first: + run.review("full", 0, selected, {}) + with diagnostic_work() as replay: + run.review("full", 0, selected, {}) + assert "diagnostic.reviewEvidence" not in first + assert replay["diagnostic.reviewEvidence"]["restored"] == 1 + assert all(row["attempted"] == 0 for row in replay.values()) + + +@pytest.mark.parametrize( + "failure", + [ + None, + ValueError("Scientific evidence failed"), + KeyboardInterrupt("Interrupted diagnostic"), + ], +) +def test_diagnostic_persistence_failure_cannot_replace_the_original_error( + monkeypatch: pytest.MonkeyPatch, failure: BaseException | None +) -> None: + store, prefix, request = memory_journal("analysis") + original_save = journal.save_checkpoint + + def save(store, prefix, workflow, key, inputs, outputs): + if key.endswith("/finished"): + raise OSError("Store is full") + return original_save(store, prefix, workflow, key, inputs, outputs) + + monkeypatch.setattr(journal, "save_checkpoint", save) + with pytest.raises(type(failure) if failure is not None else OSError) as caught: + with journal.diagnostic_attempt( + store, prefix, request.workflowRunId, {"cells": "frozen"} + ): + if failure is not None: + raise failure + if failure is not None: + assert caught.value is failure + assert "Store is full" in " ".join(failure.__notes__) + else: + assert str(caught.value) == "Store is full" + attempt = journal.analysis_snapshot(store, request.workflowRunId)[ + "diagnosticAttempts" + ][0] + assert "started" in attempt and "finished" not in attempt + + +from tests.test_agent_rna_adaptive import checkpoints as memory_checkpoints # noqa: E402, F401 diff --git a/tests/test_agent_exec.py b/tests/test_agent_exec.py index a0237efe..a9cde642 100644 --- a/tests/test_agent_exec.py +++ b/tests/test_agent_exec.py @@ -1,5 +1,7 @@ """Tests for shared Scarf agent execution and configuration.""" +from tests.agent_examples import example + import asyncio import json import threading @@ -20,9 +22,17 @@ from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.models.test import TestModel from pydantic_ai.providers.openai import OpenAIProvider - -from scarf.agent import CovariateCharacterization, FeatureCharacterization, IngestResult -from scarf.agent.config.agent_exec import _model_name, run_agent, run_agent_sync +from pydantic_ai.exceptions import ModelHTTPError, UserError + +from scarf.agent.experimental_context.contracts import CovariateCharacterization +from scarf.agent.data_enrichment.characterization import FeatureCharacterization +from scarf.agent.ingest import IngestResult +from scarf.agent.config.agent_exec import ( + _image_input_is_unsupported, + _model_name, + run_agent, + run_agent_sync, +) from scarf.agent.config import AgentRunConfig, get_model_settings, get_usage_limits from scarf.agent.tools import artifact_reference, core_artifact_reference from scarf.agent.types import ( @@ -47,6 +57,24 @@ def get_example(cls) -> "ExampleOutput": return cls(value="complete") +def test_image_input_rejection_is_distinguished_from_other_provider_errors() -> None: + unsupported = ModelHTTPError( + 400, + "text-model", + { + "message": "This model does not support multimodal (image/video/audio) inputs." + }, + ) + authentication = ModelHTTPError(401, "text-model", {"message": "Unauthorized"}) + unsupported_local = UserError("This model does not support binary content") + unrelated_local = UserError("max_retries must be greater than or equal to zero") + + assert _image_input_is_unsupported(unsupported) is True + assert _image_input_is_unsupported(authentication) is False + assert _image_input_is_unsupported(unsupported_local) is True + assert _image_input_is_unsupported(unrelated_local) is False + + def test_shared_models_have_blank_and_example_constructors() -> None: models = ( AgentRunConfig, @@ -66,19 +94,24 @@ def test_shared_models_have_blank_and_example_constructors() -> None: ) for model in models: assert isinstance(model.get_blank(), model) - assert isinstance(model.get_example(), model) + assert isinstance(example(model), model) assert all("_" not in field_name for field_name in model.model_fields) -def test_four_agent_objects_are_public() -> None: - import scarf.agent as agent_package +def test_focused_scientific_agents_remain_available_in_their_subpackages() -> None: + from scarf.agent.biological_interpretation import BiologicalInterpretationAgent + from scarf.agent.data_enrichment import DataEnrichmentAgent + from scarf.agent.experimental_context import ExperimentalContextAgent + from scarf.agent.parameter_tuning import ParameterTuningAgent - assert agent_package.DataEnrichmentAgent.__name__ == "DataEnrichmentAgent" - assert agent_package.ExperimentalContextAgent.__name__ == "ExperimentalContextAgent" - assert agent_package.ParameterTuningAgent.__name__ == "ParameterTuningAgent" - assert ( - agent_package.BiologicalInterpretationAgent.__name__ - == "BiologicalInterpretationAgent" + assert all( + callable(agent) + for agent in ( + DataEnrichmentAgent, + ExperimentalContextAgent, + ParameterTuningAgent, + BiologicalInterpretationAgent, + ) ) @@ -281,7 +314,7 @@ def reply( parts=[ ToolCallPart( tool_name=info.output_tools[0].name, - args=ExampleOutput.get_example().model_dump(), + args=example(ExampleOutput).model_dump(), ) ] ) @@ -295,7 +328,7 @@ def reply( name="example-agent", ) - assert result.output == ExampleOutput.get_example() + assert result.output == example(ExampleOutput) assert result.runInfo.agentName == "example-agent" assert result.runInfo.usage.toolCalls == 1 assert [call.toolName for call in result.runInfo.toolCalls] == ["inspect_value"] @@ -373,7 +406,7 @@ def reply( parts=[ ToolCallPart( tool_name=info.output_tools[0].name, - args=ExampleOutput.get_example().model_dump(), + args=example(ExampleOutput).model_dump(), ) ] ) @@ -390,7 +423,7 @@ async def call_from_running_loop() -> object: result = asyncio.run(call_from_running_loop()) - assert result.output == ExampleOutput.get_example() + assert result.output == example(ExampleOutput) assert result.runInfo.agentName == "notebook-host" assert [call.toolName for call in result.runInfo.toolCalls] == ["inspect_value"] @@ -425,7 +458,7 @@ async def respond(request: httpx.Request) -> httpx.Response: "function": { "name": output_tool_name, "arguments": json.dumps( - ExampleOutput.get_example().model_dump() + example(ExampleOutput).model_dump() ), }, } @@ -474,8 +507,8 @@ async def call_twice_from_running_loop() -> list[AgentExecutionResult]: results = asyncio.run(call_twice_from_running_loop()) assert [result.output for result in results] == [ - ExampleOutput.get_example(), - ExampleOutput.get_example(), + example(ExampleOutput), + example(ExampleOutput), ] assert len(clients) == 2 assert all(client.is_closed for client in clients) @@ -513,7 +546,7 @@ def reply( parts=[ ToolCallPart( tool_name=info.output_tools[0].name, - args=ExampleOutput.get_example().model_dump(), + args=example(ExampleOutput).model_dump(), ) ] ) @@ -527,7 +560,7 @@ def reply( ) ) - assert result.output == ExampleOutput.get_example() + assert result.output == example(ExampleOutput) assert callback_thread != event_loop_thread diff --git a/tests/test_agent_experimental_context.py b/tests/test_agent_experimental_context.py index 3c2db1bc..c8d5314e 100644 --- a/tests/test_agent_experimental_context.py +++ b/tests/test_agent_experimental_context.py @@ -1,5 +1,7 @@ """Tests for the tool-driven Experimental Context Agent.""" +from tests.agent_examples import example + import asyncio from types import SimpleNamespace from typing import Any, Literal @@ -8,19 +10,25 @@ import pytest import zarr from pydantic import ValidationError -from pydantic_ai import ModelRetry, RunContext +from pydantic_ai import ModelRetry, RunContext, UnexpectedModelBehavior from pydantic_ai.messages import ModelMessage, ModelResponse, ToolCallPart from pydantic_ai.models.function import AgentInfo, FunctionModel from pydantic_ai.models.test import TestModel from pydantic_ai.usage import RunUsage from zarr.storage import MemoryStore -import scarf.agent.experimental_context as experimental_context_module +import scarf.agent.experimental_context.agent as experimental_context_agent +import scarf.agent.experimental_context.contracts as experimental_context_contracts +import scarf.agent.experimental_context.qc_evidence as experimental_context_qc +import scarf.agent.experimental_context.tools as experimental_context_tools +import scarf.agent.experimental_context.validation as experimental_context_validation from scarf.agent.experimental_context import ( BatchCorrectionPlan, BatchSafetyEvidence, + CaptureFailureEvidence, CellQcPlan, CellQcProfileEvidence, + ContrastPlan, CovariateEvidence, ExperimentalContextAgent, ExperimentalContextDecision, @@ -34,10 +42,11 @@ score_current_representation, validate_experimental_context, ) -from scarf.agent.characterize_covariates import ( +from scarf.agent.experimental_context.characterization import ( CovariateCharacterization, _SelectionBoundCells, ) +from scarf.agent.experimental_context.study import StudyContract, build_study_contract from scarf.agent.types import ( ArtifactReferenceModel, ExperimentalBiologyHandoff, @@ -177,6 +186,15 @@ def inspect_artifact(self, ref: ArtifactRef) -> SimpleNamespace: ) +def _replace_store_cells(store: _Store, values: dict[str, np.ndarray]) -> None: + store.cells = _Cells(values) + store.zw = zarr.open_group(store=MemoryStore(), mode="w") + cell_data = store.zw.create_group("cellData") + cell_data.create_array("ids", data=np.asarray(values["ids"]).astype("U16")) + cell_data.create_array("I", data=np.asarray(values["I"], dtype=bool)) + store.refresh_cell_selection() + + def _write_cell_artifact( store: _Store, *, @@ -184,15 +202,17 @@ def _write_cell_artifact( kind: Literal["quality_metric", "hto_identity"], values: np.ndarray, assay: str, + operation: str | None = None, + inputs: dict[str, Any] | None = None, ) -> NamedArtifactSource: planned = plan_cell_data_artifact( store.zw, scope="assay", assay=assay, kind=kind, - operation=f"test_{kind}_source", + operation=operation or f"test_{kind}_source", parameters={"name": name}, - inputs={}, + inputs=dict(inputs or {}), execution_options={}, cell_selection=store.cell_selection, arrays={"values": ((len(values),), None)}, @@ -238,7 +258,7 @@ def metric_graph_connectivity(self, column: str, graph: ArtifactRef) -> float: @pytest.fixture(autouse=True) def _resolve_fake_graph_selection(monkeypatch: pytest.MonkeyPatch) -> None: - from scarf.agent import experimental_context as module + from scarf.agent.experimental_context import agent as module def resolve(root: zarr.Group, _graph: ArtifactRef) -> ArtifactRef: return ArtifactRef.from_dict(root.attrs["_test_cell_selection"]) @@ -407,7 +427,7 @@ def test_agent_models_have_blank_and_example_constructors() -> None: ) for model in models: assert isinstance(model.get_blank(), model) - assert isinstance(model.get_example(), model) + assert isinstance(example(model), model) assert all("_" not in field_name for field_name in model.model_fields) assert set(RepresentationEvaluation.model_fields) == { "available", @@ -429,6 +449,98 @@ def test_system_prompt_does_not_embed_fictional_output_values() -> None: assert "estimability:treatment" not in prompt +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"studyContext": ""}, "studyContext must be non-empty"), + ({"studyObjective": ""}, "studyObjective must be non-empty"), + ( + {"independentUnitColumns": ["donor", "donor"]}, + "must not contain duplicates", + ), + ( + { + "physicalCaptureColumn": "sample", + "conditionColumns": ["sample"], + }, + "cannot be the physical capture", + ), + ( + { + "technicalBatchColumns": ["batch"], + "conditionColumns": ["batch"], + }, + "cannot also be condition columns", + ), + ( + { + "correctionLicense": "safe", + "technicalBatchColumns": [], + }, + "requires batch columns", + ), + ], +) +def test_study_contract_rejects_inconsistent_design_authority( + changes: dict[str, object], + message: str, +) -> None: + values = StudyContract.get_blank().model_dump() + values.update(changes) + + with pytest.raises(ValidationError, match=message): + StudyContract.model_validate(values) + + +def test_study_contract_builder_records_correction_and_label_policies() -> None: + with pytest.raises(ValueError, match="must be done"): + build_study_contract( + study_context="Study context", + study_objective="Study objective", + experimental_result=SimpleNamespace(status="needsInput"), + ) + + def result(action: TestAction) -> SimpleNamespace: + return SimpleNamespace( + status="done", + decision=_design_decision(action=action), + characterization=CovariateCharacterization(status="done"), + batchSafety=[], + notes=[], + ) + + unsafe = build_study_contract( + study_context="Study context", + study_objective="Study objective", + experimental_result=result("unsafe"), + ) + unresolved = build_study_contract( + study_context="Study context", + study_objective="Study objective", + experimental_result=result("needsInput"), + ) + preservation = build_study_contract( + study_context="Study context", + study_objective="Study objective", + experimental_result=result("skip"), + author_label_policy="preservation", + ) + + assert unsafe.correctionLicense == "unsafeConfounded" + assert unresolved.correctionLicense == "indeterminate" + assert preservation.authorLabelPolicy == "preservation" + assert any("ineligible" in item for item in preservation.limitations) + + +def test_validator_rejects_serialized_fields_inside_narrative() -> None: + decision = ExperimentalContextDecision( + rationale='Study design is unresolved.", "evidenceIds": ["column:batch"]', + ) + + with pytest.raises(ModelRetry, match="plain prose"): + validate_experimental_context(decision, _context(_Store()).deps) + + def test_agent_runs_only_read_only_tools_and_returns_a_grounded_report() -> None: store = _Store() tool_names: set[str] = set() @@ -465,6 +577,26 @@ async def reply( ) ] ) + if request == 2: + decision = _design_decision() + return ModelResponse( + parts=[ + ToolCallPart( + tool_name="analyze_experimental_design", + args={ + "column_domains": decision.columnDomains, + "coefficients_of_interest": ( + decision.coefficientsOfInterest + ), + "units_of_inference": { + name: unit.model_dump() + for name, unit in decision.unitsOfInference.items() + }, + "batch_columns": decision.batchCorrection.batchColumns, + }, + ) + ] + ) return ModelResponse( parts=[ ToolCallPart( @@ -496,9 +628,11 @@ async def reply( assert [call.toolName for call in result.runInfo.toolCalls] == [ "inspect_cell_covariates", "analyze_experimental_design", + "analyze_experimental_design", ] assert tool_names == { "inspect_cell_covariates", + "inspect_context_evidence", "analyze_experimental_design", "score_current_representation", } @@ -511,12 +645,158 @@ async def reply( assert sorted(store.zw.group_keys()) == ["artifacts", "cellData"] +def test_agent_fails_after_design_tool_retry_exhaustion( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _Store() + analyze_retries: list[int | None] = [] + + def unavailable_design(**kwargs: Any) -> None: + deps = kwargs["deps"] + analyze_tool = next( + tool + for tool in kwargs["tools"] + if tool.name == "analyze_experimental_design" + ) + analyze_retries.append(analyze_tool.max_retries) + asyncio.run( + inspect_cell_covariates( + RunContext( + deps=deps, + model=TestModel(), + usage=RunUsage(), + ) + ) + ) + raise UnexpectedModelBehavior( + "Tool 'analyze_experimental_design' exceeded max retries count of 1" + ) + + monkeypatch.setattr( + experimental_context_agent, + "run_agent_sync", + unavailable_design, + ) + result = ExperimentalContextAgent(object()).run( + store, + study_context="Case-control study with samples nested in donors.", + cell_selection=store.cell_selection, + ) + + assert analyze_retries == [3] + assert result.status == "failed" + assert result.decision.batchCorrection.action == "needsInput" + assert result.decision.batchCorrection.batchColumns == [] + assert result.cellSelection is not None + assert result.cellSelection.artifactId == store.cell_selection.artifact_id + assert result.cellQc.profileId == "" + assert result.qcProfiles + assert result.runInfo.agentName == "experimental_context_failed" + with pytest.raises(ValueError, match="must be done"): + result.to_parameter_tuning_handoff() + assert any("did not produce" in note for note in result.notes) + + +def test_agent_rejects_malformed_batch_tool_call_without_default_selection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _Store() + + def unavailable_design(**kwargs: Any) -> None: + deps = kwargs["deps"] + asyncio.run( + inspect_cell_covariates( + RunContext( + deps=deps, + model=TestModel(), + usage=RunUsage(), + ) + ) + ) + raise UnexpectedModelBehavior( + "Tool 'analyze_experimental_design' exceeded max retries count of 3; " + "caused by ModelRetry: Unknown batch column 'batchassay'" + ) + + monkeypatch.setattr( + experimental_context_agent, + "run_agent_sync", + unavailable_design, + ) + result = ExperimentalContextAgent(object()).run( + store, + study_context="Population discovery across sequencing batches.", + study_objective="Discover stable cell populations.", + cell_selection=store.cell_selection, + directions={ + "columnDomains": { + "batch": "technical", + "donor": "design", + "sample": "design", + }, + "coefficientsOfInterest": [], + "unitsOfInference": {}, + "batchColumns": ["batch"], + }, + ) + + assert result.status == "failed" + assert result.decision.batchCorrection.batchColumns == [] + with pytest.raises(ValueError, match="must be done"): + result.to_parameter_tuning_handoff() + assert result.runInfo.agentName == "experimental_context_failed" + assert any("batchassay" in note for note in result.notes) + + +def test_agent_preserves_validated_design_uncertainty( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _Store() + + def unresolved_design(**kwargs: Any) -> SimpleNamespace: + deps = kwargs["deps"] + context = RunContext(deps=deps, model=TestModel(), usage=RunUsage()) + asyncio.run(inspect_cell_covariates(context)) + planned = _design_decision("needsInput") + asyncio.run( + analyze_experimental_design( + context, + column_domains=planned.columnDomains, + coefficients_of_interest=planned.coefficientsOfInterest, + units_of_inference=planned.unitsOfInference, + batch_columns=planned.batchCorrection.batchColumns, + ) + ) + planned.needsInput = ["Physical capture provenance remains unresolved."] + decision = kwargs["output_validator"](planned) + return SimpleNamespace( + output=decision, + runInfo=experimental_context_contracts.AgentRunInfo( + agentName="experimental_context", modelName="test" + ), + ) + + monkeypatch.setattr(experimental_context_agent, "run_agent_sync", unresolved_design) + result = ExperimentalContextAgent(object()).run( + store, + study_context="Case-control study with samples nested in donors.", + cell_selection=store.cell_selection, + ) + assert result.status == "needsInput" + assert result.decision.needsInput == [ + "Physical capture provenance remains unresolved." + ] + assert result.runInfo.agentName == "experimental_context" + with pytest.raises(ValueError, match="must be done"): + result.to_parameter_tuning_handoff() + + def test_handoff_builders_reject_incomplete_or_ambiguous_results() -> None: incomplete = ExperimentalContextResult.get_blank() with pytest.raises(ValueError, match="must be done"): incomplete.to_parameter_tuning_handoff() - ambiguous = ExperimentalContextResult.get_example() + ambiguous = example(ExperimentalContextResult) ambiguous.decision.coefficientsOfInterest.append("second_coefficient") with pytest.raises(ValueError, match="Select one coefficient explicitly"): ambiguous.to_biological_handoff() @@ -564,6 +844,36 @@ def test_tools_build_a_grounded_design_report_without_mutation() -> None: assert sorted(store.zw.group_keys()) == ["artifacts", "cellData"] +def test_exact_batch_direction_overrides_model_tool_arguments() -> None: + store = _Store() + context = _context(store, directions={"batchColumns": ["batch"]}) + decision = _design_decision() + + asyncio.run(inspect_cell_covariates(context)) + analyzed = asyncio.run( + analyze_experimental_design( + context, + column_domains=decision.columnDomains, + coefficients_of_interest=decision.coefficientsOfInterest, + units_of_inference=decision.unitsOfInference, + batch_columns=[], + ) + ) + + assert analyzed.batchSafety[0].batchColumns == ["batch"] + skip = decision.model_copy( + update={ + "batchCorrection": BatchCorrectionPlan( + action="skip", + rationale="Skip the declared batch condition.", + evidenceIds=["column:batch"], + ) + } + ) + with pytest.raises(ModelRetry, match="exact directed batch"): + validate_experimental_context(skip, context.deps) + + def test_qc_profiles_use_persisted_modality_and_shared_cell_selection() -> None: store = _Store() store.assay_names = ["protein", "peaks", "transcript"] @@ -590,6 +900,7 @@ def test_qc_profiles_use_persisted_modality_and_shared_cell_selection() -> None: assert {profile.action for profile in inspected.qcProfiles} == { "skip", "globalGaussian", + "registeredMad", } for profile in inspected.qcProfiles: assert profile.driverAssay == "transcript" @@ -659,23 +970,12 @@ def test_design_tool_offers_only_grounded_sample_mad_profiles() -> None: assert sample_profile.retainedCells == 12 assert sample_profile.evidenceId in analyzed.evidenceIds - decision.cellQc = CellQcPlan( - action=sample_profile.action, - profileId=sample_profile.profileId, - driverAssay=sample_profile.driverAssay, - driverAssayType=sample_profile.driverAssayType, - sampleColumn=sample_profile.sampleColumn, - sampleArtifact=sample_profile.sampleArtifact, - attributes=sample_profile.attributes, - artifactMetrics=sample_profile.artifactMetrics, - rationale="Use sample-aware retention evidence.", - evidenceIds=[sample_profile.evidenceId], - ) validated = validate_experimental_context(decision, context.deps) - assert validated.cellQc == decision.cellQc + assert validated.cellQc == CellQcPlan.get_blank() + assert sample_profile in context.deps.qcProfiles.values() -def test_caller_qc_direction_overrides_model_profile_selection() -> None: +def test_caller_qc_direction_shapes_evidence_without_preselection() -> None: store = _Store() store.cells._values["RNA_nCounts"] = np.arange(12, dtype=float) + 1 store.cells._values["RNA_nFeatures"] = np.arange(12, dtype=float) + 5 @@ -695,23 +995,13 @@ def test_caller_qc_direction_overrides_model_profile_selection() -> None: batch_columns=decision.batchCorrection.batchColumns, ) ) - decision.cellQc = CellQcPlan( - action="skip", - profileId="model-authored-profile", - evidenceIds=["model-authored-evidence"], - ) - validated = validate_experimental_context(decision, context.deps) - assert validated.cellQc.action == "sampleMad" - assert validated.cellQc.sampleColumn == "sample" - assert validated.cellQc.evidenceIds == [ - next( - profile.evidenceId - for profile in context.deps.qcProfiles.values() - if profile.action == "sampleMad" - ) - ] + assert validated.cellQc == CellQcPlan.get_blank() + assert any( + profile.action == "sampleMad" and profile.sampleColumn == "sample" + for profile in context.deps.qcProfiles.values() + ) def test_adt_and_hto_do_not_drive_qc_and_hto_identity_remains_metadata() -> None: @@ -820,6 +1110,348 @@ def test_artifact_metrics_and_hto_grouping_are_exact_context_evidence() -> None: assert "HTO_htoIdentity" not in store.cells.columns +def test_qc_metric_sources_report_metadata_artifact_concordance() -> None: + store = _Store() + counts = np.linspace(10.0, 100.0, store.cells.N) + features = np.linspace(5.0, 50.0, store.cells.N) + store.cells._values["RNA_nCounts"] = counts + store.cells._values["RNA_nFeatures"] = features + count_artifact = _write_cell_artifact( + store, + name="RNA_nCounts", + kind="quality_metric", + values=counts.copy(), + assay="RNA", + ) + feature_artifact = _write_cell_artifact( + store, + name="RNA_nFeatures", + kind="quality_metric", + values=features + 1e-9, + assay="RNA", + ) + context = _context( + store, + quality_metric_artifacts=[count_artifact, feature_artifact], + ) + + inspected = asyncio.run(inspect_cell_covariates(context)) + + concordance = {item.metricRole: item for item in inspected.qcSourceConcordance} + assert concordance["count"].exactlyEqual is True + assert concordance["count"].numericallyClose is True + assert concordance["count"].meanAbsoluteDifference == 0.0 + assert concordance["feature"].exactlyEqual is False + assert concordance["feature"].numericallyClose is True + assert concordance["feature"].pearsonCorrelation == pytest.approx(1.0) + assert {item.evidenceId for item in inspected.qcSourceConcordance}.issubset( + inspected.evidenceIds + ) + + +def test_capture_failure_evidence_preserves_paired_design_after_exclusion() -> None: + store = _Store() + captures = np.repeat(["capture-a", "capture-b", "failed"], 20) + donors = np.repeat(["donor-a", "donor-b", "donor-c"], 20) + disease = np.tile(np.repeat(["case", "control"], 10), 3) + samples = np.asarray( + [ + f"{donor}-{condition}" + for donor, condition in zip(donors, disease, strict=True) + ] + ) + counts = np.concatenate( + [ + np.linspace(90.0, 110.0, 20), + np.linspace(95.0, 115.0, 20), + np.linspace(3.0, 7.0, 20), + ] + ) + features = np.concatenate( + [ + np.linspace(45.0, 55.0, 20), + np.linspace(48.0, 58.0, 20), + np.linspace(2.0, 6.0, 20), + ] + ) + mito = np.concatenate( + [ + np.linspace(1.0, 3.0, 20), + np.linspace(1.0, 4.0, 20), + np.linspace(25.0, 35.0, 20), + ] + ) + n_cells = len(captures) + _replace_store_cells( + store, + { + "I": np.ones(n_cells, dtype=bool), + "ids": np.asarray([f"cell-{index}" for index in range(n_cells)]), + "names": np.asarray([f"cell-{index}" for index in range(n_cells)]), + "capture": captures, + "donor": donors, + "sample": samples, + "disease": disease, + "RNA_nCounts": counts, + "RNA_nFeatures": features, + "RNA_percentMito": mito, + }, + ) + context = _context( + store, + directions={"physicalCaptureColumn": "capture"}, + ) + domains = { + "capture": "design", + "donor": "design", + "sample": "design", + "disease": "biological", + } + units = { + "disease": InferenceUnit( + observationUnit="sample", + independentUnit="donor", + ) + } + + asyncio.run(inspect_cell_covariates(context)) + analyzed = asyncio.run( + analyze_experimental_design( + context, + column_domains=domains, + coefficients_of_interest=["disease"], + units_of_inference=units, + batch_columns=[], + ) + ) + + profile = next( + item for item in analyzed.qcProfiles if item.registeredProfile == "captureMad5" + ) + failure = next( + item for item in profile.captureFailureEvidence if item.capture == "failed" + ) + assert failure.wholeCaptureFailure is True + assert failure.independentAdverseAxes >= 2 + assert failure.preservesConditionCoverage is True + assert failure.preservesIndependentUnitCoverage is True + assert failure.exclusionEligible is True + assert failure.conditionAndUnitSafety[0]["completePairsAfterExclusion"] == 2 + assert failure.conditionAndUnitSafety[0]["incompletePairsAfterExclusion"] == 0 + assert "failed" in profile.failedCaptureCandidates + assert "failed" in profile.excludableCaptureCandidates + for source in profile.metricSources: + assert source.missingCellsByCapture == { + "capture-a": 0, + "capture-b": 0, + "failed": 0, + } + + +def test_unusable_qc_sources_preserve_provenance_and_degradation_evidence() -> None: + store = _Store() + store.cells._values["RNA_nCounts"] = np.asarray(["bad"] * store.cells.N) + store.cells._values["RNA_nFeatures"] = np.asarray( + [*np.linspace(10.0, 20.0, store.cells.N - 1), np.nan] + ) + artifact = _write_cell_artifact( + store, + name="RNA_percentMito", + kind="quality_metric", + values=np.asarray([*np.linspace(1.0, 2.0, store.cells.N - 1), np.inf]), + assay="RNA", + operation="run_feature_percentage", + inputs={ + "lineage": [ + { + "scope": "invalid", + "kind": "invalid", + "artifact_id": "invalid", + }, + store.cell_selection.to_dict(), + store.cell_selection.to_dict(), + ] + }, + ) + context = _context(store, quality_metric_artifacts=[artifact]) + + inspected = asyncio.run(inspect_cell_covariates(context)) + + sources = { + (source.sourceType, source.metricName): source + for source in inspected.qcMetricSources + } + nonnumeric = sources[("metadataColumn", "RNA_nCounts")] + assert nonnumeric.usableForFiltering is False + assert nonnumeric.missingCells == store.cells.N + assert nonnumeric.notes == ["Metric is not numeric and cannot drive filtering"] + nonfinite = sources[("metadataColumn", "RNA_nFeatures")] + assert nonfinite.usableForFiltering is False + assert nonfinite.missingCells == 1 + derived = sources[("artifact", "RNA_percentMito")] + assert derived.origin == "derivedArtifact" + assert derived.usableForFiltering is False + assert derived.missingCells == 1 + assert derived.inputArtifacts == [ + ArtifactReferenceModel.from_artifact_ref(store.cell_selection) + ] + skip = next(profile for profile in inspected.qcProfiles if profile.action == "skip") + assert any("not numeric" in note for note in skip.notes) + assert any("non-finite" in note for note in skip.notes) + + +def test_directed_capture_artifact_enables_pooled_reference_profile() -> None: + store = _Store() + labels = np.repeat(["capture-a", "capture-b", "capture-c"], 20) + counts = np.concatenate( + [ + np.linspace(90.0, 110.0, 20), + np.linspace(95.0, 115.0, 20), + np.linspace(100.0, 120.0, 20), + ] + ) + n_cells = len(labels) + _replace_store_cells( + store, + { + "I": np.ones(n_cells, dtype=bool), + "ids": np.asarray([f"cell-{index}" for index in range(n_cells)]), + "names": np.asarray([f"cell-{index}" for index in range(n_cells)]), + "RNA_nCounts": counts, + }, + ) + capture = _write_cell_artifact( + store, + name="capture", + kind="hto_identity", + values=labels, + assay="HTO", + ) + context = _context( + store, + directions={ + "physicalCaptureColumn": "capture", + "cellQc": { + "pooledReferenceCaptures": ["capture-a", "capture-b"], + }, + }, + hto_identity_artifacts=[capture], + ) + + inspected = asyncio.run(inspect_cell_covariates(context)) + + profile = next( + item + for item in inspected.qcProfiles + if item.registeredProfile == "pooledReferenceMad5" + ) + assert profile.captureColumn is None + assert profile.captureArtifact == capture + assert profile.sampleArtifact == capture + assert profile.parameters["pooledReferenceCaptures"] == [ + "capture-a", + "capture-b", + ] + assert profile.activeCellsByCapture == { + "capture-a": 20, + "capture-b": 20, + "capture-c": 20, + } + + +def test_registered_only_without_qc_driver_returns_explicit_skip() -> None: + store = _Store() + store.assay_names = [] + context = _context(store, directions={"registeredQcOnly": True}) + + inspected = asyncio.run(inspect_cell_covariates(context)) + + assert len(inspected.qcProfiles) == 1 + profile = inspected.qcProfiles[0] + assert profile.action == "skip" + assert profile.driverAssay is None + assert profile.driverAssayType is None + assert profile.notes == [ + "No RNA or ATAC assay is eligible to drive automatic cell QC" + ] + + +def test_missing_rna_percentage_metrics_are_derived_through_public_artifacts() -> None: + store = _Store() + feature_values = { + "ids": np.asarray(["MT-CO1", "RPS3", "GAPDH"]), + "names": np.asarray(["MT-CO1", "RPS3", "GAPDH"]), + } + assay = SimpleNamespace( + feats=SimpleNamespace( + N=3, + fetch_all=lambda column: feature_values[column], + ) + ) + selection_calls: list[dict[str, object]] = [] + percentage_calls: list[dict[str, object]] = [] + + def set_feature_selection(**kwargs: object) -> ArtifactRef: + selection_calls.append(kwargs) + return ArtifactRef( + scope="assay", + assay="RNA", + kind="feature_selection", + artifact_id=str(len(selection_calls)) * 64, + ) + + def run_feature_percentage( + cell_selection: ArtifactRef, + feature_selection: ArtifactRef, + *, + invalidate_cache: bool, + ) -> ArtifactRef: + percentage_calls.append( + { + "cellSelection": cell_selection, + "featureSelection": feature_selection, + "invalidateCache": invalidate_cache, + } + ) + return ArtifactRef( + scope="assay", + assay="RNA", + kind="quality_metric", + artifact_id=str(len(percentage_calls) + 2) * 64, + ) + + store.get_assay = lambda assay_name: assay + store.set_feature_selection = set_feature_selection + store.run_feature_percentage = run_feature_percentage + + sources = experimental_context_qc._derive_missing_percentage_artifacts( + store, + cell_selection=store.cell_selection, + driver=("RNA", "RNA"), + quality_sources=[], + ) + + assert [source.name for source in sources] == [ + "RNA_percentMito", + "RNA_percentRibo", + ] + np.testing.assert_array_equal( + selection_calls[0]["mask"], + [True, False, False], + ) + np.testing.assert_array_equal( + selection_calls[1]["mask"], + [False, True, False], + ) + assert all(call["from_assay"] == "RNA" for call in selection_calls) + assert all(call["invalidate_cache"] is False for call in selection_calls) + assert [call["cellSelection"] for call in percentage_calls] == [ + store.cell_selection, + store.cell_selection, + ] + assert all(call["invalidateCache"] is False for call in percentage_calls) + + def test_validator_rejects_harmony_when_batch_confounds_biology() -> None: store = _Store() context = _context(store) @@ -885,7 +1517,7 @@ def test_harmony_safety_uses_only_exact_proposed_batch_columns() -> None: def test_batch_safety_does_not_depend_on_pairwise_selected_flag( monkeypatch: pytest.MonkeyPatch, ) -> None: - from scarf.agent import experimental_context as module + from scarf.agent.experimental_context import tools as module store = _Store() decision = _design_decision(action="evaluateHarmony") @@ -968,6 +1600,8 @@ def test_batch_safety_checks_multiple_columns_jointly() -> None: batch_columns=["plate"], ) ) + context = _context(store, directions={"columnKinds": {"disease": "continuous"}}) + asyncio.run(inspect_cell_covariates(context)) joint = asyncio.run( analyze_experimental_design( context, @@ -1208,7 +1842,7 @@ async def reply( def test_harmony_requires_resolved_units_and_estimability( monkeypatch: pytest.MonkeyPatch, ) -> None: - from scarf.agent import experimental_context as module + from scarf.agent.experimental_context import validation as module characterization = CovariateCharacterization( status="needsInput", @@ -1228,7 +1862,7 @@ def test_harmony_requires_resolved_units_and_estimability( ) monkeypatch.setattr( module, - "characterize_covariates", + "characterize_context", lambda *_args, **_kwargs: characterization, ) store = _Store() @@ -1309,7 +1943,7 @@ def test_returned_decision_canonicalizes_caller_directions() -> None: def test_named_artifact_and_qc_source_validation_edges() -> None: - metric = NamedArtifactSource.get_example() + metric = example(NamedArtifactSource) identity = NamedArtifactSource( name="HTO_identity", artifact=ArtifactReferenceModel( @@ -1329,7 +1963,7 @@ def test_named_artifact_and_qc_source_validation_edges() -> None: with pytest.raises(ValidationError, match="requires both name and artifact"): NamedArtifactSource(artifact=metric.artifact) - validate = experimental_context_module._validate_qc_sources + validate = experimental_context_contracts._validate_qc_sources with pytest.raises(ValueError, match="metadata attributes must be unique"): validate( action="globalGaussian", @@ -1437,8 +2071,252 @@ def test_named_artifact_and_qc_source_validation_edges() -> None: ) +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"coefficient": " disease "}, "coefficient cannot contain"), + ({"sampleBy": " "}, "sampleBy must be"), + ({"pairBy": " donor "}, "pairBy must be"), + ({"groupOrder": ["case", "case"]}, "groupOrder must contain unique"), + ({"groupOrder": ["case", np.inf]}, "cannot contain non-finite"), + ({"expressionCutoff": np.nan}, "expressionCutoff must be finite"), + ( + {"expressionCutoff": 0.1}, + "expressionCutoff is only used with fraction", + ), + ( + {"groupOrder": ["a", "b", "c"], "test": "mann_whitney"}, + "mann_whitney requires exactly two", + ), + ( + {"groupOrder": ["a", "b"], "test": "kruskal_wallis"}, + "kruskal_wallis requires at least three", + ), + ({"test": "wilcoxon"}, "wilcoxon requires exactly two groups"), + ( + {"test": "mann_whitney", "pairBy": "donor"}, + "paired contrast must use the wilcoxon", + ), + ({"replicationPassed": False}, "licensed contrast requires resolved"), + ( + {"status": "blocked", "blockedReasons": []}, + "non-licensed contrast requires blockedReasons", + ), + ], +) +def test_contrast_plan_rejects_inconsistent_licenses( + changes: dict[str, object], + message: str, +) -> None: + values: dict[str, object] = { + "coefficient": "disease", + "groupOrder": ["case", "control"], + "sampleBy": "sample", + "test": "mann_whitney", + "status": "licensed", + "betweenUnitDesign": True, + "replicationPassed": True, + "estimabilityPassed": True, + } + values.update(changes) + with pytest.raises(ValidationError, match=message): + ContrastPlan.model_validate(values) + + +def test_contrast_plan_accepts_complete_paired_license() -> None: + plan = ContrastPlan( + coefficient="treatment", + groupOrder=["treated", "control"], + sampleBy="sample", + pairBy="donor", + test="wilcoxon", + status="licensed", + betweenUnitDesign=True, + replicationPassed=True, + estimabilityPassed=True, + pairedCoveragePassed=True, + ) + + assert plan.status == "licensed" + assert plan.test == "wilcoxon" + + +@pytest.mark.parametrize( + ("registered_profile", "action", "sample_column", "sample_artifact", "message"), + [ + ( + "retainWithFlags", + "registeredMad", + None, + None, + "non-filtering skip", + ), + ( + "retainWithFlags", + "skip", + "capture", + None, + "cannot include a capture source", + ), + ("globalMad5", "globalGaussian", None, None, "registeredMad action"), + ( + "captureMad5", + "registeredMad", + None, + None, + "requires exactly one proven capture source", + ), + ( + "globalMad5", + "registeredMad", + "capture", + None, + "cannot include a capture source", + ), + ], +) +def test_registered_qc_source_contract_rejects_inconsistent_modes( + registered_profile: str, + action: str, + sample_column: str | None, + sample_artifact: NamedArtifactSource | None, + message: str, +) -> None: + with pytest.raises(ValueError, match=message): + experimental_context_contracts._validate_qc_sources( + action=action, + attributes=["RNA_nCounts"], + artifact_metrics=[], + sample_column=sample_column, + sample_artifact=sample_artifact, + registered_profile=registered_profile, + ) + + +def test_registered_qc_source_contract_accepts_global_capture_and_retain_modes() -> ( + None +): + capture = NamedArtifactSource( + name="capture", + artifact=ArtifactReferenceModel( + assay="HTO", + kind="hto_identity", + artifactId="4" * 64, + ), + ) + validate = experimental_context_contracts._validate_qc_sources + + validate( + action="skip", + attributes=[], + artifact_metrics=[], + sample_column=None, + sample_artifact=None, + registered_profile="retainWithFlags", + ) + validate( + action="registeredMad", + attributes=["RNA_nCounts"], + artifact_metrics=[], + sample_column=None, + sample_artifact=None, + registered_profile="globalMad5", + ) + validate( + action="registeredMad", + attributes=["RNA_nCounts"], + artifact_metrics=[], + sample_column=None, + sample_artifact=capture, + registered_profile="captureMad5", + ) + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ( + {"independentAdverseAxes": 1}, + "axis count must match", + ), + ( + {"wholeCaptureFailure": False}, + "requires at least two independent", + ), + ( + {"preservesConditionCoverage": False}, + "exclusion requires failure and preserved", + ), + ], +) +def test_capture_failure_evidence_rejects_inconsistent_state( + changes: dict[str, object], + message: str, +) -> None: + values: dict[str, object] = { + "capture": "failed", + "activeCells": 20, + "retainedCells": 2, + "retainedFraction": 0.1, + "adverseAxes": ["count", "feature"], + "independentAdverseAxes": 2, + "wholeCaptureFailure": True, + "preservesConditionCoverage": True, + "preservesIndependentUnitCoverage": True, + "exclusionEligible": True, + } + values.update(changes) + with pytest.raises(ValidationError, match=message): + CaptureFailureEvidence.model_validate(values) + + +def test_cell_qc_profile_requires_exact_capture_failure_inventory() -> None: + capture = NamedArtifactSource( + name="capture", + artifact=ArtifactReferenceModel( + assay="HTO", + kind="hto_identity", + artifactId="5" * 64, + ), + ) + failure = CaptureFailureEvidence( + capture="failed", + activeCells=20, + retainedCells=2, + retainedFraction=0.1, + adverseAxes=["count", "feature"], + independentAdverseAxes=2, + wholeCaptureFailure=True, + preservesConditionCoverage=True, + preservesIndependentUnitCoverage=True, + exclusionEligible=True, + ) + with pytest.raises(ValidationError, match="mutually exclusive"): + CellQcProfileEvidence( + action="skip", + captureColumn="capture", + captureArtifact=capture, + ) + with pytest.raises(ValidationError, match="failure evidence must be unique"): + CellQcProfileEvidence( + action="skip", + captureFailureEvidence=[failure, failure], + ) + with pytest.raises(ValidationError, match="failed captures must match"): + CellQcProfileEvidence( + action="skip", + captureFailureEvidence=[failure], + ) + with pytest.raises(ValidationError, match="excludable captures must match"): + CellQcProfileEvidence( + action="skip", + captureFailureEvidence=[failure], + failedCaptureCandidates=["failed"], + ) + + def test_experimental_handoff_validation_edges() -> None: - result = ExperimentalContextResult.get_example() + result = example(ExperimentalContextResult) without_selection = result.model_copy(update={"cellSelection": None}) with pytest.raises(ValueError, match="lacks a cell selection"): without_selection.to_parameter_tuning_handoff() @@ -1508,14 +2386,14 @@ def test_experimental_context_private_input_guards( "htoIdentityColumn": "donor", }, ).deps - assert experimental_context_module._hto_identity_columns(deps) == [ + assert experimental_context_qc._hto_identity_columns(deps) == [ "sample", "donor", ] unknown_tool = SimpleNamespace(name="future_tool") assert ( - experimental_context_module._prepare_experimental_context_tool( + experimental_context_tools._prepare_experimental_context_tool( SimpleNamespace(deps=deps), unknown_tool, ) @@ -1526,12 +2404,13 @@ def test_experimental_context_private_input_guards( confounding=[{"coefficient": 3, "pairs": []}], ) assert ( - experimental_context_module.characterization_evidence(characterization) == set() + experimental_context_contracts.characterization_evidence(characterization) + == set() ) deps.cellSelection = None with pytest.raises(ValueError, match="exact artifact"): - experimental_context_module._cell_selection_ref(deps) + experimental_context_qc._cell_selection_ref(deps) deps.cellSelection = _artifact_ref = ArtifactRef( scope="assay", assay="RNA", @@ -1539,11 +2418,11 @@ def test_experimental_context_private_input_guards( artifact_id="4" * 64, ) with pytest.raises(ValueError, match="datastore cell selection"): - experimental_context_module._cell_selection_ref(deps) + experimental_context_qc._cell_selection_ref(deps) del _artifact_ref with pytest.raises(TypeError, match="NamedArtifactSource"): - experimental_context_module._source_ref( + experimental_context_qc._source_ref( object(), expected_kind="quality_metric", ) @@ -1552,12 +2431,12 @@ def test_experimental_context_private_input_guards( artifact=ArtifactReferenceModel(), ) with pytest.raises(ValueError, match="non-empty semantic name"): - experimental_context_module._source_ref( + experimental_context_qc._source_ref( blank_source, expected_kind="quality_metric", ) with pytest.raises(ValueError, match="quality_metric"): - experimental_context_module._source_ref( + experimental_context_qc._source_ref( NamedArtifactSource( name="identity", artifact=ArtifactReferenceModel( @@ -1579,19 +2458,19 @@ def test_experimental_context_private_input_guards( ) deps.htoIdentityArtifacts = [duplicate, duplicate] with pytest.raises(ValueError, match="names must be unique"): - experimental_context_module._hto_artifact_map(deps) + experimental_context_qc._hto_artifact_map(deps) deps.cellSelection = store.cell_selection monkeypatch.setattr( - experimental_context_module, + experimental_context_qc, "read_stored_selection_mask", lambda *_args, **_kwargs: np.ones(store.cells.N + 1, dtype=bool), ) with pytest.raises(ValueError, match="aligned boolean selection"): - experimental_context_module._active_cell_count(deps) + experimental_context_qc._active_cell_count(deps) -def test_qc_profile_degradation_and_selection_guards( +def test_core_qc_reports_unavailable_default_bounds( monkeypatch: pytest.MonkeyPatch, ) -> None: store = _Store() @@ -1599,7 +2478,7 @@ def test_qc_profile_degradation_and_selection_guards( deps = context.deps active = np.ones(store.cells.N, dtype=bool) notes: list[str] = [] - profile = experimental_context_module._global_qc_profile( + profile = experimental_context_qc._global_qc_profile( deps, ("RNA", "RNA"), active, @@ -1610,15 +2489,15 @@ def test_qc_profile_degradation_and_selection_guards( notes, ) assert profile is None - assert notes == ["Ignored constant QC metric 'constant'"] + assert "constant metric 'constant'" in notes[0] monkeypatch.setattr( - experimental_context_module, + experimental_context_qc, "gaussian_quantile_bounds", lambda *_args, **_kwargs: (float("nan"), float("nan")), ) notes = [] - profile = experimental_context_module._global_qc_profile( + profile = experimental_context_qc._global_qc_profile( deps, ("RNA", "RNA"), active, @@ -1631,87 +2510,6 @@ def test_qc_profile_degradation_and_selection_guards( assert profile is None assert "non-finite Gaussian bounds" in notes[0] - skip = CellQcProfileEvidence( - profileId="skip", - action="skip", - activeCells=store.cells.N, - retainedCells=store.cells.N, - retainedFraction=1.0, - evidenceId="qcProfile:skip", - ) - global_profile = CellQcProfileEvidence( - profileId="global", - action="globalGaussian", - driverAssay="RNA", - driverAssayType="RNA", - attributes=["RNA_nCounts"], - activeCells=store.cells.N, - retainedCells=store.cells.N - 1, - retainedFraction=(store.cells.N - 1) / store.cells.N, - evidenceId="qcProfile:global", - ) - deps.qcProfiles = {"skip": skip, "global": global_profile} - characterization = CovariateCharacterization(status="done") - - deps.directions = {"cellQc": {"profileId": 3}} - with pytest.raises(ModelRetry, match="profileId direction must be a string"): - experimental_context_module._canonical_cell_qc_plan( - CellQcPlan(), deps, characterization - ) - deps.directions = { - "cellQc": {"sampleColumn": "sample", "sampleArtifactName": "identity"} - } - with pytest.raises(ModelRetry, match="cannot select both"): - experimental_context_module._canonical_cell_qc_plan( - CellQcPlan(), deps, characterization - ) - deps.directions = {"cellQc": {"sampleArtifactName": 3}} - with pytest.raises(ModelRetry, match="sampleArtifactName must be a string"): - experimental_context_module._canonical_cell_qc_plan( - CellQcPlan(), deps, characterization - ) - deps.directions = {"cellQc": {"action": "unknown"}} - with pytest.raises(ModelRetry, match="Unsupported cellQc.action"): - experimental_context_module._canonical_cell_qc_plan( - CellQcPlan(), deps, characterization - ) - deps.directions = {"cellQc": {"action": "sampleMad"}} - with pytest.raises(ModelRetry, match="exactly one offered profile"): - experimental_context_module._canonical_cell_qc_plan( - CellQcPlan(), deps, characterization - ) - deps.directions = {"cellQc": {"profileId": "unknown"}} - with pytest.raises(ModelRetry, match="was not offered"): - experimental_context_module._canonical_cell_qc_plan( - CellQcPlan(), deps, characterization - ) - - deps.directions = {} - selected = experimental_context_module._canonical_cell_qc_plan( - CellQcPlan(), deps, characterization - ) - assert selected.profileId == "global" - - mismatched = CellQcPlan( - action="globalGaussian", - profileId="global", - driverAssay="RNA", - driverAssayType="RNA", - attributes=["different_metric"], - evidenceIds=["qcProfile:global"], - ) - with pytest.raises(ModelRetry, match="copy the selected offered profile"): - experimental_context_module._canonical_cell_qc_plan( - mismatched, deps, characterization - ) - missing_evidence = mismatched.model_copy( - update={"attributes": ["RNA_nCounts"], "evidenceIds": []} - ) - with pytest.raises(ModelRetry, match="cite its exact profile"): - experimental_context_module._canonical_cell_qc_plan( - missing_evidence, deps, characterization - ) - def test_design_analysis_rejects_invalid_batch_proposals( monkeypatch: pytest.MonkeyPatch, @@ -1742,7 +2540,7 @@ def test_design_analysis_rejects_invalid_batch_proposals( notes=["design failed"], ) monkeypatch.setattr( - experimental_context_module, + experimental_context_tools, "characterize_covariates", lambda *_args, **_kwargs: failed, ) @@ -1753,6 +2551,7 @@ def test_design_analysis_rejects_invalid_batch_proposals( column_domains={}, coefficients_of_interest=[], units_of_inference={}, + batch_columns=[], ) ) @@ -1778,7 +2577,7 @@ def test_design_analysis_records_not_computed_estimability( ) monkeypatch.setattr( - experimental_context_module, + experimental_context_tools, "reduce_observation_units", lambda *_args, **_kwargs: (_ for _ in ()).throw(ValueError("bad design")), ) @@ -1906,13 +2705,12 @@ def validate( candidate_coefficients: dict[str, dict[str, Any]] | None = None, requested: set[str] | None = None, ) -> None: - experimental_context_module._validate_batch_correction_plan( + experimental_context_validation._validate_batch_correction_plan( candidate, deps or context.deps, characterization, requested if requested is not None else {"disease"}, units, - candidate.cellQc, candidate_records or records, candidate_coefficients or coefficient_records, ) @@ -2010,7 +2808,7 @@ def changed_plan(**updates: Any) -> ExperimentalContextDecision: "kind": "continuous", }, } - with pytest.raises(ModelRetry, match="must be categorical"): + with pytest.raises(ModelRetry, match="Batch correction is unsafe"): validate( changed_plan( action="evaluateHarmony", diff --git a/tests/test_agent_feature_interventions.py b/tests/test_agent_feature_interventions.py new file mode 100644 index 00000000..be661a4b --- /dev/null +++ b/tests/test_agent_feature_interventions.py @@ -0,0 +1,237 @@ +"""Representation interventions use exact eligible genes and protected programs.""" + +import hashlib +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +from scarf.agent.orchestrator import rna_tuning +from scarf.agent.orchestrator.models import artifact_model_to_ref +from scarf.agent.types import ArtifactReferenceModel +from scarf.storage.refs import ArtifactRef +from tests.test_agent_rna_adaptive import checkpoints as memory_checkpoints # noqa: F401 +from tests.test_agent_rna_evidence_mode import make_run + + +@pytest.fixture +def feature_run(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest) -> Any: + request.getfixturevalue("memory_checkpoints") + run, selected = make_run(monkeypatch, object()) + names = np.array(["MT-CO1", "RPL3", "XIST", "CD3D", *[f"G{i}" for i in range(96)]]) + ids = np.array([f"ENSG{i}" for i in range(len(names))]) + arrays: dict[str, dict[str, np.ndarray]] = {} + + def selection(mask: Any, **kwargs: Any) -> ArtifactRef: + mask = np.asarray(mask, dtype=bool) + identifier = hashlib.sha256(mask.tobytes()).hexdigest() + arrays[identifier] = { + "values": mask.copy(), + "corrected_variance": np.arange(100, dtype=float), + } + return ArtifactRef( + scope="assay", assay="RNA", kind="feature_selection", artifact_id=identifier + ) + + allowed = selection(np.arange(100) != 99) + eligible = selection((np.arange(100) != 99) & (np.arange(100) != 0)) + setting = run.settings[selected.candidateId] + setting.features = ArtifactReferenceModel.from_artifact_ref(eligible) + setting.eligibleFeatures = setting.features + setting.hvgCount = 98 + run.handoff.graphFeatureCandidates["eligibleAll"] = ( + ArtifactReferenceModel.from_artifact_ref(allowed) + ) + run.plan.assays[0].featureParameters = {} + run.family_patterns = { + "mitochondrial": "^MT-", + "ribosomal": "^RP[LS]", + "sexLinked": "^XIST$", + } + run.store = SimpleNamespace( + load_artifact=lambda reference: { + key: value.copy() for key, value in arrays[reference.artifact_id].items() + }, + get_assay=lambda _: SimpleNamespace( + feats=SimpleNamespace( + N=100, fetch_all=lambda column: names if column == "names" else ids + ) + ), + set_feature_selection=selection, + ) + + def rank( + store: Any, *, eligible: ArtifactRef, top_n: int, **kwargs: Any + ) -> ArtifactRef: + mask = store.load_artifact(eligible)["values"].copy() + mask[np.flatnonzero(mask)[top_n:]] = False + return selection(mask) + + monkeypatch.setattr(rna_tuning, "rank_core_hvgs", rank) + return run, selected, arrays, selection + + +@pytest.mark.parametrize( + "intervention", + ["includeFamily", "includeFeature", "excludeFamily", "excludeFeature"], +) +def test_feature_policy_changes_only_representation_eligibility( + feature_run: Any, intervention: str +) -> None: + run, selected, arrays, _ = feature_run + marker_features = run.marker_features + previous = run.settings[selected.candidateId].model_copy(deep=True) + value = { + "includeFamily": "mitochondrial", + "includeFeature": "ENSG0", + "excludeFamily": "ribosomal", + "excludeFeature": "ENSG1", + }[intervention] + changed = run.apply_experiment( + selected, {"parameter": intervention, "value": value} + ) + old_mask = arrays[previous.eligibleFeatures.artifactId]["values"] + new_mask = arrays[changed.eligibleFeatures.artifactId]["values"] + expected = 0 if intervention.startswith("include") else 1 + assert np.flatnonzero(old_mask != new_mask).tolist() == [expected] + assert not new_mask[99], "Ineligible genes cannot be reinstated" + assert changed.parameters == previous.parameters + assert run.settings[selected.candidateId] == previous + assert run.marker_features == marker_features + + +@pytest.mark.parametrize( + "protection", + [ + {"protectFeatures": ["ENSG1"]}, + {"protectFamilies": ["ribosomal"]}, + {"protectFamilies": ["sexLinked"]}, + ], +) +def test_objective_protection_blocks_exclusion_before_feature_execution( + feature_run: Any, protection: dict[str, Any] +) -> None: + run, selected, arrays, _ = feature_run + run.plan.assays[0].featureParameters = protection + value = "XIST" if protection.get("protectFamilies") == ["sexLinked"] else "RPL3" + before = len(arrays) + with pytest.raises(ValueError, match="objective-protected"): + run.apply_experiment(selected, {"parameter": "excludeFeature", "value": value}) + assert len(arrays) == before + + +def test_missing_feature_or_underpowered_representation_cannot_be_executed( + feature_run: Any, +) -> None: + run, selected, _, _ = feature_run + with pytest.raises(ValueError, match="absent from the assay"): + run.apply_experiment( + selected, {"parameter": "excludeFeature", "value": "missing"} + ) + with pytest.raises(ValueError, match="fixed PCA dimension"): + run.apply_experiment(selected, {"parameter": "hvgCount", "value": 10}) + with pytest.raises(ValueError, match="explicit technical column"): + run.apply_experiment( + selected, {"parameter": "hvgRanking", "value": "batchAware"} + ) + + +def test_nominations_follow_eligible_protected_genes_and_skip_unknown_families( + feature_run: Any, +) -> None: + run, selected, _, _ = feature_run + setting = run.settings[selected.candidateId] + nominate = rna_tuning.RnaTuningRun._feature_nomination + run.plan.assays[0].featureParameters = { + "protectFamilies": ["notRegistered", "mitochondrial"] + } + assert nominate(run, setting) == { + "parameter": "includeFamily", + "value": "mitochondrial", + } + run.plan.assays[0].featureParameters = { + "protectFeatures": ["ENSG1"], + "proposedExcludeFeatures": ["RPL3", "missing"], + } + assert nominate(run, setting) is None + run.plan.assays[0].featureParameters = {"proposedExcludeFamilies": ["ribosomal"]} + assert nominate(run, setting) == { + "parameter": "excludeFamily", + "value": "ribosomal", + } + + +def test_offered_experiments_respect_population_rank_and_covariate_kind( + feature_run: Any, +) -> None: + run, selected, _, _ = feature_run + setting = run.settings[selected.candidateId] + setting.hvgCount = 22 + setting.ranking = "batchAware" + setting.rankingColumn = "capture" + setting.parameters.useHarmony = True + run.scope_sizes["full"] = 22 + run.batch_columns = ["capture", "age"] + run.study.columnKinds = {"capture": "categorical", "age": "continuous"} + options = run.experiments(selected) + assert "dimensions:30" not in options + assert "neighborsK:41" not in options + assert "hvgRanking:batchAware:age" not in options + assert "hvgRanking:batchAware:capture" not in options + assert "hvgRanking:global" in options and "useHarmony:false" in options + + +def test_batch_ranking_requires_supported_groups_and_reuses_core_variability( + feature_run: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + from scarf.agent.experimental_context import characterization + + run, selected, _, selection = feature_run + eligible = artifact_model_to_ref( + run.settings[selected.candidateId].eligibleFeatures + ) + labels = np.repeat(["a", "b", "small"], [25, 25, 5]) + calls: list[dict[str, Any]] = [] + run.batch_columns = ["capture"] + run.store.zw = None + run.store.cells = object() + monkeypatch.setattr( + characterization, + "_SelectionBoundCells", + lambda *a: SimpleNamespace(fetch=lambda _: labels), + ) + + def filter_cells(columns: Any, lower: Any, upper: Any, **kwargs: Any) -> str: + assert columns == ["capture"] and lower == upper + assert kwargs["cell_selection"] == run.cells + return lower[0] + + core_features = selection(np.ones(100, dtype=bool)) + summary = ArtifactRef( + scope="assay", assay="RNA", kind="feature_summary", artifact_id="b" * 64 + ) + original_load = run.store.load_artifact + run.store.load_artifact = lambda ref: ( + {"normed_n": np.repeat(25, 100)} if ref == summary else original_load(ref) + ) + run.store.filter_cells = filter_cells + run.store.inspect_artifact = lambda ref: SimpleNamespace( + inputs={"feature_summary": summary.to_dict()} + ) + + def select_hvgs(cells: Any, **kwargs: Any) -> ArtifactRef: + calls.append({"cells": cells, **kwargs}) + return core_features + + run.store.select_hvgs = select_hvgs + with pytest.raises(ValueError, match="approved technical column"): + run.batch_ranking(eligible, 40, "donor") + assert not calls + indices = run.batch_ranking(eligible, 40, "capture") + assert {call["cells"] for call in calls} == {"a", "b"} + assert all(call["blacklist"] == "" and call["top_n"] == 100 for call in calls) + assert len(indices) > 0 and 0 not in indices and 99 not in indices + labels[:] = "a" + with pytest.raises(ValueError, match="two groups"): + run.batch_ranking(eligible, 40, "capture") diff --git a/tests/test_agent_finalization_doublet_policy.py b/tests/test_agent_finalization_doublet_policy.py new file mode 100644 index 00000000..7c5f49f3 --- /dev/null +++ b/tests/test_agent_finalization_doublet_policy.py @@ -0,0 +1,217 @@ +"""Finalization distinguishes disabled scoring from missing required evidence.""" + +from types import SimpleNamespace + +from scarf.agent.orchestrator import finalization +from scarf.agent.orchestrator.models import ( + AssayPreprocessingPlan, + AutomatedPreprocessingPlan, + AutomatedWorkflowConfig, + AutomatedWorkflowRequest, + OrchestrationRequestRecord, + PreprocessedAssayHandoff, + StageEvidenceReference, + WorkflowIdentity, + WorkflowStageAttempt, + artifact_model_to_ref, +) +from scarf.agent.parameter_tuning.contracts import ( + ArtifactRecord, + ParameterCandidate, + ParameterCandidateEvaluation, + ParameterMetrics, + ParameterTuningReport, +) + +from tests.agent_examples import example + + +def _fixture(monkeypatch, *, scored=False, warning=None): + sequence = iter(range(1, 30)) + + def reference(kind, *, assay="RNA"): + return ArtifactRecord( + kind=kind, + assay=assay, + scope="datastore" if assay is None else "assay", + artifactId=f"{next(sequence):064x}", + ) + + cells = reference("cell_selection", assay=None) + normalized = reference("normalized") + marker_features = reference("feature_selection") + artifacts = { + "pca": reference("reduction"), + "connectivityMap": reference("connectivity_map"), + "clusters": reference("cluster_labels"), + "markerTable": reference("marker_table"), + } + selection = reference("cell_selection", assay=None) + score = reference("doublet_score") + if scored: + artifacts.update({"doubletScore:0": score, "doubletCellSelection:0": selection}) + selected = ParameterCandidateEvaluation( + candidateId="selected", + parameters=ParameterCandidate(candidateId="selected"), + status="done", + eligible=True, + cellSelection=cells, + artifacts=artifacts, + metrics=ParameterMetrics(nClusters=3, seedStability=0.95, markerCoherence=1), + warnings=[] if warning is None else [warning], + ) + report = ParameterTuningReport( + status="done", + fromAssay="RNA", + cellSelection=cells, + evaluations=[selected], + recommendedCandidateId=selected.candidateId, + finalClusterArtifact=artifacts["clusters"], + ) + statuses = {} + + def bind(ref, **inputs): + statuses[artifact_model_to_ref(ref)] = SimpleNamespace( + exists=True, + complete=True, + inputs={ + name: artifact_model_to_ref(value).to_dict() + for name, value in inputs.items() + }, + ) + + bind(normalized, cell_selection=cells) + bind( + artifacts["clusters"], cell_selection=cells, graph=artifacts["connectivityMap"] + ) + bind(artifacts["markerTable"], cell_selection=cells, clusters=artifacts["clusters"]) + bind(score, cell_selection=selection) + layouts = [] + + def initialize(ref, **kwargs): + layouts.append(("initialize", ref)) + return artifact_model_to_ref(reference("embedding_initialization")) + + def umap(ref, initialization, **kwargs): + layouts.append(("umap", ref)) + return artifact_model_to_ref(reference("embedding")) + + store = SimpleNamespace( + zw=None, + summary=lambda: SimpleNamespace( + assays=[SimpleNamespace(name="RNA", assay_type="RNA")] + ), + inspect_artifact=statuses.__getitem__, + load_artifact=lambda ref: {}, + build_embedding_initialization=initialize, + run_umap=umap, + ) + monkeypatch.setattr( + finalization, + "graph_cell_selection", + lambda root, graph: artifact_model_to_ref(cells), + ) + monkeypatch.setattr( + finalization.journal, "_ensure_orchestration_store", lambda store: "agents" + ) + monkeypatch.setattr( + finalization.journal, "_validated_done_outcome", lambda *args: None + ) + monkeypatch.setattr( + finalization.journal, + "_start_attempt", + lambda *args, **kwargs: WorkflowStageAttempt(stage="analysis_finalization"), + ) + monkeypatch.setattr(finalization.journal, "_save_outcome", lambda *args: None) + monkeypatch.setattr( + finalization.journal, + "read_stage_evidence", + lambda *args: report.model_dump(mode="json"), + ) + record = OrchestrationRequestRecord( + inputIdentity={}, + modelIdentity="test", + request=example(AutomatedWorkflowRequest).model_copy( + update={ + "primaryAssay": "RNA", + "markerAssay": "RNA", + "analysisAssays": ["RNA"], + } + ), + config=AutomatedWorkflowConfig(), + ) + plan = AutomatedPreprocessingPlan( + primaryAssay="RNA", + markerAssay="RNA", + cellSelection=cells, + assays=[ + AssayPreprocessingPlan(assay="RNA", assayType="RNA", graphEligible=True) + ], + ) + handoff = PreprocessedAssayHandoff( + assay="RNA", + assayType="RNA", + cellSelection=cells, + normalized=normalized, + markerFeatures=marker_features, + nCells=30, + ) + + def run(): + return finalization.FinalizationStagesMixin().analysis_finalization_stage( + store, + WorkflowIdentity(workflowRunId="test-workflow"), + record, + [], + plan, + [handoff], + report, + StageEvidenceReference( + workflowRunId="test-workflow", + stage="parameter_tuning", + key="report", + contentSha256="0" * 64, + ), + ) + + return SimpleNamespace( + run=run, + selected=selected, + score=score, + selection=selection, + layouts=layouts, + artifacts=artifacts, + ) + + +def test_explicitly_disabled_doublets_finalize_with_visible_limitation(monkeypatch): + warning = ( + "Advisory doublet scoring was not run for assay 'RNA' because " + "score_doublets=False. Doublet contamination was not assessed." + ) + fixture = _fixture(monkeypatch, warning=warning) + outcome, final = fixture.run() + assert outcome.status == "done", outcome.error + assert final.doubletScores == final.doubletScoreSelections == [] + assert warning in final.limitations + assert warning in outcome.notes + assert [name for name, ref in fixture.layouts] == ["initialize", "umap"] + assert final.markers.artifactId == fixture.artifacts["markerTable"].artifactId + assert fixture.selected.metrics.seedStability == 0.95 + + +def test_missing_doublets_without_a_recorded_limitation_prevent_completion(monkeypatch): + fixture = _fixture(monkeypatch) + outcome, final = fixture.run() + assert outcome.status == "failed" + assert "lacks advisory doublet scores" in outcome.error + assert final.umap is None + + +def test_existing_scores_retain_exact_capture_selection(monkeypatch): + fixture = _fixture(monkeypatch, scored=True) + outcome, final = fixture.run() + assert outcome.status == "done", outcome.error + assert final.doubletScores[0].artifactId == fixture.score.artifactId + assert final.doubletScoreSelections[0].artifactId == fixture.selection.artifactId + assert final.doubletScoreSelections[0] != final.cellSelection diff --git a/tests/test_agent_global_repairs.py b/tests/test_agent_global_repairs.py new file mode 100644 index 00000000..787d5804 --- /dev/null +++ b/tests/test_agent_global_repairs.py @@ -0,0 +1,237 @@ +"""Global full-cohort repair admissions survive retries and context revisions.""" + +import hashlib +from copy import deepcopy +from typing import Any + +import pytest + +from scarf.agent import record_io +from scarf.agent.orchestrator import journal, rna_tuning +from scarf.agent.orchestrator.budget import CandidateBudget, CandidateBudgetExceeded +from scarf.agent.orchestrator.models import AutomatedWorkflowConfig +from scarf.agent.types import ArtifactReferenceModel +from tests.agent_journal_store import memory_journal +from tests.test_agent_required_comparisons import panel_run # noqa: F401 + + +def repair_inputs() -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + return ( + ArtifactReferenceModel( + scope="datastore", kind="cell_selection", artifactId="a" * 64 + ).model_dump(mode="json"), + { + "features": "fixed", + "parameters": { + "candidateId": "baseline", + "dimensions": 21, + "neighborsK": 11, + "leidenResolution": 1.0, + }, + }, + {"parameter": "leidenResolution", "value": 1.5}, + ) + + +def save_review( + store: Any, prefix: str, run_id: str, key: str, *, harmony: bool = False +) -> None: + cells, setting, experiment = repair_inputs() + if harmony: + experiment = {"parameter": "useHarmony", "value": True} + journal.save_checkpoint( + store, + prefix, + run_id, + key, + { + "comparisonCoverage": {"phase": "validation"}, + "candidates": [{"candidateId": "baseline", "cellSelection": cells}], + "settings": {"baseline": setting}, + "experiments": {"target": experiment}, + }, + { + "action": { + "action": "experiment", + "selectedCandidateId": "baseline", + "experimentId": "target", + } + }, + ) + + +def test_repair_reservation_is_global_idempotent_and_does_not_charge_primary_work() -> ( + None +): + store, prefix, record = memory_journal() + original = {"study": "original"} + budget = CandidateBudget( + store, prefix, record.workflowRunId, record.config, original + ) + cells, setting, experiment = repair_inputs() + budget.admit_repair(cells, setting, experiment) + before = journal._list_keys(store.zw, prefix) + assert len(budget.repairs) == 1 + assert budget.summary()["scopes"]["full"]["reserved"] == { + "graphs": 0, + "partitions": 0, + } + resumed = CandidateBudget( + store, + prefix, + record.workflowRunId, + record.config, + {"study": "revised"}, + previous_provenances=(original,), + ) + renamed = deepcopy(setting) + renamed["parameters"]["candidateId"] = "replayed-label" + resumed.admit_repair(cells, renamed, experiment) + assert journal._list_keys(store.zw, prefix) == before + with pytest.raises(CandidateBudgetExceeded, match="repair has been used"): + resumed.admit_repair(cells, setting, {"parameter": "dimensions", "value": 30}) + assert len(resumed.repairs) == 1 + + +@pytest.mark.parametrize("revised", [False, True]) +def test_old_committed_repair_is_counted_without_rewriting_it(revised: bool) -> None: + store, prefix, record = memory_journal() + original = {"study": "original"} + digest = hashlib.sha256(record_io.canonical_json_bytes(original)).hexdigest() + key = ( + f"parameter_tuning/evidence_revisions/{digest}/full/review0" + if revised + else "parameter_tuning/full/review0" + ) + save_review(store, prefix, record.workflowRunId, key) + before = journal._list_keys(store.zw, prefix) + saved = journal.read_checkpoint(store, prefix, record.workflowRunId, key) + budget = CandidateBudget( + store, + prefix, + record.workflowRunId, + record.config, + {"study": "new"}, + previous_provenances=(original,), + ) + assert len(budget.repairs) == 1 + budget.admit_repair(*repair_inputs()) + assert journal._list_keys(store.zw, prefix) == before + assert journal.read_checkpoint(store, prefix, record.workflowRunId, key) == saved + with pytest.raises(CandidateBudgetExceeded): + cells, setting, _ = repair_inputs() + budget.admit_repair(cells, setting, {"parameter": "dimensions", "value": 30}) + + +def test_harmony_comparison_does_not_consume_a_targeted_repair() -> None: + store, prefix, record = memory_journal() + save_review( + store, + prefix, + record.workflowRunId, + "parameter_tuning/full/review0", + harmony=True, + ) + budget = CandidateBudget(store, prefix, record.workflowRunId, record.config, {}) + assert not budget.repairs + budget.admit_repair(*repair_inputs()) + assert len(budget.repairs) == 1 + + +def test_unadmitted_proposal_does_not_inflate_explicit_repair_accounting() -> None: + store, prefix, record = memory_journal() + budget = CandidateBudget(store, prefix, record.workflowRunId, record.config, {}) + cells, setting, _ = repair_inputs() + budget.admit_repair(cells, setting, {"parameter": "dimensions", "value": 30}) + save_review(store, prefix, record.workflowRunId, "parameter_tuning/full/review0") + restored = CandidateBudget(store, prefix, record.workflowRunId, record.config, {}) + assert restored.repairs == budget.repairs + + +def test_disabled_repair_limit_rejects_before_any_primary_admission() -> None: + store, prefix, record = memory_journal() + budget = CandidateBudget( + store, + prefix, + record.workflowRunId, + AutomatedWorkflowConfig(maxFullRepairs=0), + {}, + ) + with pytest.raises(CandidateBudgetExceeded): + budget.admit_repair(*repair_inputs()) + assert not budget.repairs + assert not budget.admissions["full"] + + +def test_full_recovery_executes_requested_resolution_and_preserves_baseline_panel( + request: pytest.FixtureRequest, + monkeypatch: Any, +) -> None: + run = request.getfixturevalue("panel_run") + run.recovery_scope = "sample1" + monkeypatch.setattr( + rna_tuning, + "screening_coverage", + lambda *args: ({"screeningCells": 100, "populationCells": 100}, []), + ) + monkeypatch.setattr( + run, + "experiments", + lambda selected: {"target": {"parameter": "leidenResolution", "value": 1.5}}, + ) + calls = [] + + def review(scope: str, index: int, selected: Any, coverage: Any) -> Any: + calls.append(selected.parameters.leidenResolution) + return rna_tuning.TuningAction( + action="experiment" if index == 0 else "accept", + selectedCandidateId=selected.candidateId, + experimentId="target" if index == 0 else None, + rationale="The observed population split needs a finer partition.", + plainLanguageSummary="Check the supported split.", + correctionNeed="notApplicable", + evidenceIds=[f"candidate:{selected.candidateId}"], + quantitativeFindings=["The supplied fixture has completed graph evidence."], + qualitativeFindings=["A supported population may contain a finer split."], + comparisonConclusions=[], + objectivePreservation="Retain the frozen cohort and gene selection.", + concern="The supported split requires resolution sensitivity evidence.", + expectedImprovement="Resolve the split without changing the graph.", + ) + + monkeypatch.setattr(run, "review", review) + status, selected = run.assess_scope("full", run.cells, run.baseline()) + assert status == "accept" + assert selected.parameters.leidenResolution == 1.5 + assert calls == [1.0, 1.5] + assert len(run.resolution_candidates["full"]) == 4 + assert run.budget.summary()["scopes"]["full"]["completed"] == { + "graphs": 1, + "partitions": 5, + } + assert len(run.budget.repairs) == run.full_repairs == 1 + assert run.full_repair["selectedCandidateId"] == selected.candidateId + + +def test_targeted_recovery_plans_do_not_collide_for_same_graph_different_question( + request: pytest.FixtureRequest, +) -> None: + run = request.getfixturevalue("panel_run") + run.recovery_scope = "sample1" + setting = run.baseline() + run._resolution_panel("full", run.cells, setting) + counts = run.budget.summary() + run.last_action = rna_tuning.TuningAction( + action="enlarge", + selectedCandidateId="observed", + rationale="The larger sample left donor support uncertain.", + plainLanguageSummary="Validate donor support with all cells.", + correctionNeed="notApplicable", + evidenceIds=["candidate:observed"], + quantitativeFindings=["Coverage was insufficient in the observed sample."], + qualitativeFindings=["Small populations require more donor support evidence."], + comparisonConclusions=[], + objectivePreservation="Keep rare populations eligible for assessment.", + ) + run._resolution_panel("full", run.cells, setting) + assert run.budget.summary() == counts diff --git a/tests/test_agent_harmony_reference.py b/tests/test_agent_harmony_reference.py new file mode 100644 index 00000000..b7446d66 --- /dev/null +++ b/tests/test_agent_harmony_reference.py @@ -0,0 +1,213 @@ +"""Matched numerical correction preserves biological programs and donor support.""" + +import json +from pathlib import Path + +import numpy as np +import pytest + +from scarf.agent.ingest import ingest +from scarf.agent.parameter_tuning.contracts import ( + ParameterCandidate, + ParameterTuningDependencies, +) +from scarf.agent.parameter_tuning.diagnostics import augment_cluster_evaluations +from scarf.agent.parameter_tuning.execution import execute_parameter_candidate +from scarf.agent.parameter_tuning.hvg import core_hvg_evidence +from scarf.agent.parameter_tuning.selection import harmony_acceptance_gate +from scarf.agent.tools import core_artifact_reference +from scarf.datastore.datastore import DataStore +from scarf.metrics.association import coefficient_estimability +from tests.test_agent_ingest import _write_h5ad + + +@pytest.mark.slow +@pytest.mark.parametrize( + ("batch_effect", "expected_benefit"), [(2.0, False), (0.5, True)] +) +def test_real_harmony_comparison_measures_batch_and_biological_preservation( + tmp_path: Path, + batch_effect: float, + expected_benefit: bool, +) -> None: + """A crossed technical effect is evaluated, with missing doublets still blocking.""" + rng = np.random.default_rng(4444) + n_cells, rare_start = 1200, 1152 + rows = np.arange(n_cells) + donor = rows % 6 + batch = (rows // 6) % 2 + labels = np.where(rows >= rare_start, "rare", np.where(rows < 576, "A", "B")) + design_rows = np.unique( + np.column_stack([donor.astype(str), batch.astype(str), labels]), axis=0 + ) + design = coefficient_estimability( + design_rows[:, 2], + coefficientKind="categorical", + technicals={"batch": design_rows[:, 1]}, + technicalKinds={"batch": "categorical"}, + ) + assert design["coefficientEstimable"] + confounded = coefficient_estimability( + design_rows[:, 2], + coefficientKind="categorical", + technicals={"batch": design_rows[:, 2]}, + technicalKinds={"batch": "categorical"}, + ) + assert not confounded["coefficientEstimable"] + counts = rng.poisson(0.2, (n_cells, 90)).astype(np.uint16) + for mask, columns in ( + (labels == "A", slice(0, 12)), + (labels == "B", slice(12, 24)), + (labels == "rare", slice(24, 36)), + ): + counts[mask, columns] += rng.poisson(6.0, (int(mask.sum()), 12)).astype( + np.uint16 + ) + for value, columns in ((0, slice(36, 48)), (1, slice(48, 60))): + mask = batch == value + counts[mask, columns] += rng.poisson( + batch_effect, (int(mask.sum()), 12) + ).astype(np.uint16) + names = ( + [f"A_{i}" for i in range(12)] + + [f"B_{i}" for i in range(12)] + + [f"RARE_{i}" for i in range(12)] + + [f"TECHNICAL_{i}" for i in range(24)] + + [f"BACKGROUND_{i}" for i in range(30)] + ) + source, target = tmp_path / "crossed.h5ad", tmp_path / "crossed.zarr" + _write_h5ad( + source, + counts, + feature_types=[b"Gene Expression"] * 90, + feature_names=[name.encode() for name in names], + ) + result = ingest(path=source, zarrPath=target, directions={"matrixKey": "X"}) + assert result.status == "done", result.notes + store = DataStore( + str(target), + default_assay="RNA", + min_features_per_cell=-1, + mito_pattern="", + ribo_pattern="", + nthreads=1, + ) + for column, values in ( + ("donor", donor.astype(str)), + ("batch", batch.astype(str)), + ("population", labels), + ): + store.cells.insert(column, values, overwrite=True) + cells = store.snapshot_cell_selection("I") + genes = core_hvg_evidence(store, assay="RNA", cells=cells)["scarfDefault"] + marker_genes = store.select_all_features(from_assay="RNA") + normalized = store.run_normalization(cells, features=genes) + candidates = { + name: ParameterCandidate( + candidateId=name, + dimensions=21, + neighborsK=21, + leidenResolution=0.5, + useHarmony=corrected, + ) + for name, corrected in (("native", False), ("harmony", True)) + } + deps = ParameterTuningDependencies( + store=store, + normalized=normalized, + cellSelection=cells, + fromAssay="RNA", + candidates=candidates, + batchColumns=("batch",), + preservationColumns=("population",), + columnKinds={"batch": "categorical", "population": "categorical"}, + harmonyAuthorized=True, + maxCandidates=2, + minClusterCells=10, + ) + measured = [execute_parameter_candidate(deps, name) for name in candidates] + assert all(row.status == "done" and row.eligible for row in measured), measured + native, corrected = augment_cluster_evaluations( + store, + measured, + marker_assay="RNA", + marker_features=marker_genes, + independent_unit_columns=("donor",), + technical_columns=("batch",), + ) + improvement = ( + corrected.metrics.batchMixing["batch"] - native.metrics.batchMixing["batch"] + ) + assert (improvement > 0.05) == expected_benefit, improvement + for metric, baseline in native.metrics.biologicalPreservation["population"].items(): + assert ( + corrected.metrics.biologicalPreservation["population"][metric] + >= baseline - 0.05 + ) + for row in (native, corrected): + clusters = np.asarray( + store.load_artifact(core_artifact_reference(row.artifacts["clusters"]))[ + "values" + ][:] + ) + rare_labels, sizes = np.unique(clusters[rare_start:], return_counts=True) + population = rare_labels[np.argmax(sizes)] + members = clusters == population + assert (members & (rows >= rare_start)).sum() / members.sum() >= 0.9 + assert len(np.unique(donor[members])) == 6 + assert ( + sum( + name.startswith("RARE_") + for name in row.metrics.topMarkerGenes[str(population)] + ) + >= 3 + ) + accepted, reasons = harmony_acceptance_gate( + native, + corrected, + batch_columns=("batch",), + protected_columns=("population",), + independent_unit_columns=("donor",), + require_doublet_evidence=True, + ) + assert not accepted + assert "Matched doublet-concentration comparison is missing." in reasons + other_checks, other_reasons = harmony_acceptance_gate( + native, + corrected, + batch_columns=("batch",), + protected_columns=("population",), + independent_unit_columns=("donor",), + ) + assert other_checks == expected_benefit, other_reasons + if not expected_benefit: + assert ( + "Harmony did not improve an approved batch metric beyond tolerance." + in reasons + ) + (tmp_path / "harmony_reference.json").write_text( + json.dumps( + { + "cells": n_cells, + "rareCells": n_cells - rare_start, + "nuisanceCountRate": batch_effect, + "batchMixingImprovement": improvement, + "batchMixing": { + row.candidateId: row.metrics.batchMixing + for row in (native, corrected) + }, + "preservation": { + row.candidateId: row.metrics.biologicalPreservation + for row in (native, corrected) + }, + "crossUnitSupport": { + row.candidateId: row.metrics.crossUnitSupport + for row in (native, corrected) + }, + "unresolvedAcceptance": reasons, + "primaryEvaluations": len(measured), + "limitation": "Synthetic measured correction, not model decision agreement or full workflow acceptance.", + }, + sort_keys=True, + ) + ) diff --git a/tests/test_agent_harmony_required.py b/tests/test_agent_harmony_required.py new file mode 100644 index 00000000..89c3f9e3 --- /dev/null +++ b/tests/test_agent_harmony_required.py @@ -0,0 +1,279 @@ +"""Safe correction designs require an executed matched Harmony comparison.""" + +from typing import Any + +import pytest + +from scarf.agent.orchestrator import rna_tuning +from scarf.agent.orchestrator.budget import CandidateBudgetExceeded +from scarf.agent.parameter_tuning.comparisons import ( + setting_changes, + validate_comparison_review, +) +from scarf.agent.parameter_tuning.contracts import ParameterCandidateEvaluation +from tests.agent_comparison_examples import observed_action +from tests.test_agent_required_comparisons import panel_run # noqa: F401 + + +def _set_design(run: rna_tuning.RnaTuningRun, *, safe: bool = True) -> None: + run.study = run.study.model_copy( + update={ + "correctionLicense": "safe" if safe else "unsafeConfounded", + "technicalBatchColumns": ["capture"], + } + ) + + +def _action( + selected: ParameterCandidateEvaluation, + *, + combining: bool = False, + choices: dict[str, str] | None = None, + safe: bool = True, +) -> rna_tuning.TuningAction: + return rna_tuning.TuningAction( + action="combine" if combining else "accept", + selectedCandidateId=selected.candidateId, + correctionNeed="notNeeded" if safe else "notApplicable", + combinedSettings=choices, + evidenceIds=[f"candidate:{selected.candidateId}"], + quantitativeFindings=["Compare the observed matched representations."], + qualitativeFindings=["Retain the reference population marker program."], + comparisonConclusions=[], + plainLanguageSummary="Prefer the native representation after comparison.", + objectivePreservation="Retain the reference population.", + rationale="The execution test nominates native; correction must still run.", + ) + + +def _assert_matched_pairs( + run: rna_tuning.RnaTuningRun, scope: str, *, partitions: int +) -> None: + corrected = [row for row in run.evaluations[scope] if row.parameters.useHarmony] + assert len(corrected) == partitions + for harmony in corrected: + setting = run.settings[harmony.candidateId] + native = [ + row + for row in run.evaluations[scope] + if not row.parameters.useHarmony + and row.parameters.model_dump(exclude={"candidateId", "useHarmony"}) + == harmony.parameters.model_dump(exclude={"candidateId", "useHarmony"}) + and run.settings[row.candidateId].features == setting.features + ] + assert len(native) == 1 + assert native[0].cellSelection == harmony.cellSelection + assert native[0].status == harmony.status == "done" + assert run.settings[native[0].candidateId].model_dump( + exclude={"parameters"} + ) == (setting.model_dump(exclude={"parameters"})) + + +@pytest.mark.parametrize("safe", [False, True]) +def test_combined_native_preference_still_executes_harmony_for_safe_design( + request: pytest.FixtureRequest, + monkeypatch: pytest.MonkeyPatch, + safe: bool, +) -> None: + run = request.getfixturevalue("panel_run") + _set_design(run, safe=safe) + reviews = [] + + def review(scope: str, index: int, selected: Any, coverage: Any) -> Any: + reviews.append(index) + if index == 0: + assert len(run.evaluations[scope]) == 12 + assert not any(row.parameters.useHarmony for row in run.evaluations[scope]) + return _action( + selected, + combining=True, + safe=safe, + choices={ + name: selected.candidateId + for name in ( + "hvgCountCandidateId", + "hvgRankingCandidateId", + "featurePolicyCandidateId", + "pcaCandidateId", + "neighborsCandidateId", + ) + }, + ) + _assert_matched_pairs(run, scope, partitions=4 if safe else 0) + native = next( + row + for row in run.evaluations[scope] + if not row.parameters.useHarmony + and row.parameters.dimensions == 21 + and row.parameters.neighborsK == 11 + and row.parameters.leidenResolution == 1.0 + ) + return _action(native, safe=safe) + + monkeypatch.setattr(run, "review", review) + status, selected = run.assess_scope("sample0", run.cells, None) + assert status == "accept" and selected is not None + assert not selected.parameters.useHarmony + assert reviews == [0, 1] + assert run.budget.summary()["scopes"]["sample0"]["completed"] == { + "graphs": 10 if safe else 9, + "partitions": 16 if safe else 12, + } + + +@pytest.mark.parametrize("safe", [False, True]) +@pytest.mark.parametrize("recovery", [False, True]) +def test_native_full_validation_and_recovery_admit_required_counterpart( + request: pytest.FixtureRequest, + monkeypatch: pytest.MonkeyPatch, + safe: bool, + recovery: bool, +) -> None: + run = request.getfixturevalue("panel_run") + _set_design(run, safe=safe) + if recovery: + run.recovery_scope = "sample1" + else: + run.discovery_scope = "sample0" + initial = run.baseline() + monkeypatch.setattr( + run, + "review", + lambda scope, index, selected, coverage: _action(selected, safe=safe), + ) + status, selected = run.assess_scope("full", run.cells, initial) + assert status == "accept" and selected is not None + assert not selected.parameters.useHarmony + assert selected.parameters.model_dump(exclude={"candidateId"}) == ( + initial.parameters.model_dump(exclude={"candidateId"}) + ) + _assert_matched_pairs(run, "full", partitions=(4 if recovery else 1) if safe else 0) + counts = { + "graphs": 2 if safe else 1, + "partitions": (4 if recovery else 1) * (2 if safe else 1), + } + assert run.budget.summary()["scopes"]["full"] == { + "reserved": counts, + "completed": counts, + } + if recovery: + coverage = run.comparison_coverage("full", run.cells) + for row in coverage["comparisons"]: + assert set( + setting_changes( + coverage["candidateSettings"][row["baselineCandidateId"]], + coverage["candidateSettings"][row["alternativeCandidateId"]], + ) + ) == {"partition"} + + +@pytest.mark.parametrize( + ("recovery", "limit", "value"), + [ + (False, "maxFullGraphs", 1), + (False, "maxFullPartitions", 1), + (True, "maxFullGraphs", 1), + (True, "maxFullPartitions", 7), + ], +) +def test_full_matched_plan_rejects_insufficient_capacity_before_execution( + request: pytest.FixtureRequest, + monkeypatch: pytest.MonkeyPatch, + recovery: bool, + limit: str, + value: int, +) -> None: + run = request.getfixturevalue("panel_run") + _set_design(run) + setattr(run.request.config, limit, value) + if recovery: + run.recovery_scope = "sample1" + monkeypatch.setattr( + run, "execute", lambda *args: pytest.fail("The whole pair must fit first") + ) + monkeypatch.setattr( + run, "review", lambda *args: pytest.fail("No incomplete pair may be reviewed") + ) + with pytest.raises(CandidateBudgetExceeded): + run.assess_scope("full", run.cells, run.baseline()) + assert run.evaluations["full"] == [] + assert run.budget.summary()["scopes"]["full"] == { + "reserved": {"graphs": 0, "partitions": 0}, + "completed": {"graphs": 0, "partitions": 0}, + } + + +def test_safe_combined_plan_admits_both_modes_before_new_native_partitions( + request: pytest.FixtureRequest, + monkeypatch: pytest.MonkeyPatch, +) -> None: + run = request.getfixturevalue("panel_run") + _set_design(run) + run.request.config.maxScreeningEvaluations = 18 + + def review(scope: str, index: int, selected: Any, coverage: Any) -> Any: + assert index == 0 + choices = { + name: selected.candidateId + for name in ( + "hvgCountCandidateId", + "hvgRankingCandidateId", + "featurePolicyCandidateId", + "pcaCandidateId", + "neighborsCandidateId", + ) + } + choices["pcaCandidateId"] = next( + row.candidateId + for row in run.evaluations[scope] + if row.parameters.dimensions == 10 + ) + return _action(selected, combining=True, choices=choices) + + monkeypatch.setattr(run, "review", review) + with pytest.raises(CandidateBudgetExceeded): + run.assess_scope("sample0", run.cells, None) + assert len(run.evaluations["sample0"]) == 12 + assert not any(row.parameters.useHarmony for row in run.evaluations["sample0"]) + assert run.budget.summary()["scopes"]["sample0"] == { + "reserved": {"graphs": 9, "partitions": 12}, + "completed": {"graphs": 9, "partitions": 12}, + } + + +def test_native_recovery_choice_validates_against_matched_resolution_panel( + request: pytest.FixtureRequest, + monkeypatch: pytest.MonkeyPatch, +) -> None: + run = request.getfixturevalue("panel_run") + _set_design(run) + sample = rna_tuning.artifact_model_to_ref( + run.handoff.cellSelection.model_copy(update={"artifactId": "e" * 64}) + ) + run.scope_sizes["sample1"] = 100 + baseline = run._resolution_panel("sample1", sample, run.baseline()) + run._sensitivity_panel("sample1", sample, baseline) + run.combined_candidates["sample1"] = baseline.candidateId + run.recovery_scope = "sample1" + + def review(scope: str, index: int, selected: Any, coverage: Any) -> Any: + assert scope == "full" and index == 0 + assert not selected.parameters.useHarmony + selected.metrics.topMarkerGenes = {"0": ["MS4A1"], "1": ["CD3D"]} + selected.metrics.nClusters = 2 + comparison = run.comparison_coverage(scope, run.cells) + action = observed_action( + { + "comparisonCoverage": comparison, + "currentCandidateId": selected.candidateId, + } + ) + action["correctionNeed"] = "notNeeded" + validate_comparison_review(comparison, action) + return rna_tuning.TuningAction.model_validate(action) + + monkeypatch.setattr(run, "review", review) + status, selected = run.assess_scope("full", run.cells, run.baseline()) + assert status == "accept" and selected is not None + assert not selected.parameters.useHarmony + _assert_matched_pairs(run, "full", partitions=4) diff --git a/tests/test_agent_hvg_boundaries.py b/tests/test_agent_hvg_boundaries.py new file mode 100644 index 00000000..6978c497 --- /dev/null +++ b/tests/test_agent_hvg_boundaries.py @@ -0,0 +1,211 @@ +"""Feature ranking respects frozen eligibility and measured group support.""" + +from dataclasses import replace +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +from scarf.agent.parameter_tuning.hvg import ( + HvgGroupVariability, + aggregate_hvg_rankings, + effective_hvg_candidate_counts, + rank_core_hvgs, +) +from scarf.storage.refs import ArtifactRef + + +def _ranking_store() -> tuple[ + Any, ArtifactRef, ArtifactRef, dict[str, np.ndarray], list[np.ndarray] +]: + eligible = ArtifactRef("assay", "feature_selection", "a" * 64, "RNA") + statistics = ArtifactRef("assay", "feature_selection", "b" * 64, "RNA") + data = { + "mask": np.asarray([True, False, True, True, True]), + "variance": np.asarray([1.0, 100, 2, 2, 0]), + } + saved = [] + + def load(ref: ArtifactRef) -> dict[str, np.ndarray]: + return ( + {"values": data["mask"]} + if ref == eligible + else {"corrected_variance": data["variance"]} + ) + + def save( + *, from_assay: str, mask: np.ndarray, invalidate_cache: bool + ) -> ArtifactRef: + assert from_assay == "RNA" and not invalidate_cache + saved.append(mask.copy()) + return eligible + + return ( + SimpleNamespace(load_artifact=load, set_feature_selection=save), + eligible, + statistics, + data, + saved, + ) + + +def test_ranked_feature_intervention_preserves_eligibility_and_tie_order() -> None: + store, eligible, statistics, _, saved = _ranking_store() + rank_core_hvgs(store, eligible=eligible, statistics=statistics, top_n=3) + np.testing.assert_array_equal(saved[-1], [True, False, True, True, False]) + rank_core_hvgs( + store, + eligible=eligible, + statistics=statistics, + top_n=3, + ranking=np.asarray([4, 3, 2, 1, 0]), + ) + np.testing.assert_array_equal(saved[-1], [False, False, True, True, True]) + assert len(saved) == 2 + + +@pytest.mark.parametrize( + "damage", + [ + "unaligned", + "nonfinite", + "duplicate", + "matrix", + "negativeIndex", + "largeIndex", + "incompleteUniverse", + "tooFewRequested", + "booleanCount", + "tooFewEligible", + ], +) +def test_ranked_feature_intervention_rejects_incomplete_evidence_before_persistence( + damage: str, +) -> None: + store, eligible, statistics, data, saved = _ranking_store() + ranking = None + top_n = 3 + if damage == "unaligned": + data["variance"] = data["variance"][:-1] + elif damage == "nonfinite": + data["variance"][0] = np.nan + elif damage == "duplicate": + ranking = np.asarray([0, 0, 2, 3, 4]) + elif damage == "matrix": + ranking = np.asarray([[0, 2, 3, 4]]) + elif damage == "negativeIndex": + ranking = np.asarray([-1, 0, 2, 3, 4]) + elif damage == "largeIndex": + ranking = np.asarray([0, 2, 3, 4, 5]) + elif damage == "incompleteUniverse": + ranking = np.asarray([0, 2, 3]) + elif damage == "tooFewRequested": + top_n = 2 + elif damage == "booleanCount": + top_n = True + else: + data["mask"][:] = [True, False, True, False, False] + with pytest.raises(ValueError): + rank_core_hvgs( + store, + eligible=eligible, + statistics=statistics, + top_n=top_n, + ranking=ranking, + ) + assert saved == [] + + +@pytest.mark.parametrize( + "count,targets", + [(True, (1,)), (0, (1,)), (10, "123"), (10, (True,)), (10, (0,)), (10, ())], +) +def test_hvg_search_requires_valid_registered_counts(count: Any, targets: Any) -> None: + with pytest.raises((TypeError, ValueError)): + effective_hvg_candidate_counts(count, targets) + + +def test_one_supported_group_is_explicit_global_evidence_and_masks_remain_nested() -> ( + None +): + def unavailable_groups() -> Any: + raise AssertionError( + "An unsupported batch ranking must not consume group summaries" + ) + yield + + ranking = aggregate_hvg_rankings( + np.asarray([3.0, 9, 2, 1]), + np.asarray([True, False, True, True]), + unavailable_groups(), + valid_group_count=1, + candidate_targets=(2, 3, 5), + ) + assert ranking.ranking_mode == "global" + assert ranking.eligible_feature_count == 3 + assert ranking.candidate_counts == (2, 3) + np.testing.assert_array_equal(ranking.candidate_mask(2), [True, False, True, False]) + assert np.all(ranking.candidate_mask(2) <= ranking.candidate_mask(3)) + with pytest.raises(ValueError, match="registered counts"): + ranking.candidate_mask(1) + + +@pytest.mark.parametrize( + "damage", + [ + "globalShape", + "globalNonfinite", + "globalNegative", + "groupCountType", + "negativeGroups", + "excessGroups", + "missingGroups", + "emptyGroupId", + "cellCountType", + "emptyGroup", + "groupShape", + "groupNonfinite", + "groupNegative", + "undetectedGroup", + ], +) +def test_batch_aware_ranking_cannot_claim_unavailable_group_evidence( + damage: str, +) -> None: + corrected = np.asarray([3.0, 2, 1]) + eligible = np.ones(3, dtype=bool) + first = HvgGroupVariability("batch_a", 100, corrected.copy(), eligible.copy()) + second = replace(first, group_id="batch_b") + groups = [first, second] + declared: Any = 2 + if damage == "globalShape": + corrected = corrected[:2] + elif damage == "globalNonfinite": + corrected[0] = np.nan + elif damage == "globalNegative": + corrected[0] = -1 + elif damage == "groupCountType": + declared = True + elif damage == "negativeGroups": + declared = -1 + elif damage == "excessGroups": + groups.append(replace(first, group_id="batch_c")) + elif damage == "missingGroups": + groups.pop() + elif damage == "emptyGroupId": + groups[0] = replace(first, group_id="") + elif damage == "cellCountType": + groups[0] = replace(first, cell_count=True) + elif damage == "emptyGroup": + groups[0] = replace(first, cell_count=0) + elif damage == "groupShape": + groups[0] = replace(first, detected_features=np.ones(2, dtype=bool)) + elif damage == "groupNonfinite": + groups[0] = replace(first, corrected_variance=np.asarray([np.nan, 2, 1])) + elif damage == "groupNegative": + groups[0] = replace(first, corrected_variance=np.asarray([-1.0, 2, 1])) + else: + groups[0] = replace(first, detected_features=np.zeros(3, dtype=bool)) + with pytest.raises((TypeError, ValueError)): + aggregate_hvg_rankings(corrected, eligible, groups, valid_group_count=declared) diff --git a/tests/test_agent_hvg_diagnostics.py b/tests/test_agent_hvg_diagnostics.py new file mode 100644 index 00000000..4d1b5293 --- /dev/null +++ b/tests/test_agent_hvg_diagnostics.py @@ -0,0 +1,50 @@ +import numpy as np +import pytest + +from scarf.agent.parameter_tuning.hvg import ( + HvgGroupVariability, + aggregate_hvg_rankings, + effective_hvg_candidate_counts, +) + + +def test_hvg_candidate_counts_are_capped_and_unique() -> None: + assert effective_hvg_candidate_counts(2500) == (1000, 2000, 2500) + assert effective_hvg_candidate_counts(8, (3, 5, 10, 10)) == (3, 5, 8) + with pytest.raises(ValueError, match="greater than 0"): + effective_hvg_candidate_counts(0) + + +def test_batch_aware_hvg_ranking_keeps_nested_registered_candidates() -> None: + corrected = np.asarray([5.0, 4.0, 3.0, 2.0, 1.0]) + eligible = np.asarray([True, True, True, True, False]) + groups = [ + HvgGroupVariability( + group_id="str:a", + cell_count=10, + corrected_variance=np.asarray([5.0, 1.0, 4.0, 3.0, 2.0]), + detected_features=np.ones(5, dtype=bool), + ), + HvgGroupVariability( + group_id="str:b", + cell_count=10, + corrected_variance=np.asarray([1.0, 5.0, 4.0, 3.0, 2.0]), + detected_features=np.ones(5, dtype=bool), + ), + ] + + ranking = aggregate_hvg_rankings( + corrected, + eligible, + groups, + valid_group_count=2, + candidate_targets=(2, 3), + ) + + assert ranking.ranking_mode == "batchAware" + assert ranking.recurrence.tolist() == [1, 1, 2, 2, 0] + narrow = ranking.candidate_mask(2) + broad = ranking.candidate_mask(3) + assert np.all(~narrow | broad) + assert int(narrow.sum()) == 2 + assert int(broad.sum()) == 3 diff --git a/tests/test_agent_ingest.py b/tests/test_agent_ingest.py index 5bc8e4c3..e0589c08 100644 --- a/tests/test_agent_ingest.py +++ b/tests/test_agent_ingest.py @@ -8,8 +8,7 @@ import zarr from scipy.sparse import csr_matrix -from scarf.agent import detect_format, ingest -from scarf.agent.persistence import AgentWorkflowRun +from scarf.agent.ingest import detect_format, ingest from scarf.agent.types import Decision from scarf.readers import inspect_h5ad @@ -86,7 +85,6 @@ def _patch_ingest_summary( import importlib ingest_common = importlib.import_module("scarf.agent.ingest.common") - persistence = importlib.import_module("scarf.agent.persistence") def summarize( _zarr_path: str, @@ -101,16 +99,6 @@ def summarize( ) monkeypatch.setattr(ingest_common, "open_summary", summarize) - monkeypatch.setattr( - persistence, - "create_agent_workflow", - lambda _path: AgentWorkflowRun( - workflowRunId="workflow-1", - createdAtNs=1, - analysisStore=str(_path), - datasetFingerprints={assay_name: "fingerprint-1"}, - ), - ) def test_detect_format_by_suffix(tmp_path: Path) -> None: @@ -174,24 +162,15 @@ def test_ingest_h5ad_prefers_raw_integer_matrix(tmp_path: Path) -> None: assert result.zarrPath is not None assert "RNA" in result.assayNames assert result.summary is not None - assert result.workflowRun is not None assert result.acceptedActions - assert result.acceptedActions[-2]["op"] == "DataStore" - assert result.acceptedActions[-1]["op"] == "createAgentWorkflow" + assert result.acceptedActions[-1]["op"] == "DataStore" root = zarr.open_group(result.zarrPath, mode="r") - assert result.workflowRun.datasetFingerprints == { - assay_name: str(root[assay_name].attrs["dataset_fingerprint"]) + assert all( + root[assay_name].attrs["dataset_fingerprint"] for assay_name in result.assayNames - } - assert isinstance(root["agents"], zarr.Group) - assert root["agents"].attrs["format"] == "scarf_agent_reports" - assert ( - Path(result.zarrPath) - / "agents" - / "runs" - / result.workflowRun.workflowRunId - / "workflow.json" - ).is_file() + ) + assert "agents" not in root + assert "workflowRun" not in result.model_dump() def test_ingest_h5ad_stops_on_prenormalized_only(tmp_path: Path) -> None: @@ -358,10 +337,8 @@ def test_ingest_10x_h5(tmp_path: Path) -> None: assert result.status == "done", result.notes assert result.format == "10x_h5" assert "RNA" in result.assayNames - assert result.workflowRun is not None assert result.acceptedActions - assert result.acceptedActions[-2]["op"] == "DataStore" - assert result.acceptedActions[-1]["op"] == "createAgentWorkflow" + assert result.acceptedActions[-1]["op"] == "DataStore" def _file_snapshot(location: Path) -> dict[str, tuple[int, int]]: @@ -438,29 +415,21 @@ def test_ingest_overwrite_true_replaces_destination(tmp_path: Path) -> None: assert not sentinel.exists() -def test_ingest_derives_destination_and_creates_workflow(tmp_path: Path) -> None: +def test_ingest_derives_destination_without_creating_workflow(tmp_path: Path) -> None: path = tmp_path / "counts.h5ad" _write_h5ad(path, np.array([[1, 0], [0, 2]], dtype=np.uint16)) result = ingest(path=path) assert result.status == "done" assert result.format == "h5ad" assert result.zarrPath == str(tmp_path / "counts.zarr") - assert result.workflowRun is not None assert any("Using derived Zarr destination" in note for note in result.notes) root = zarr.open_group(result.zarrPath, mode="r") - assert result.workflowRun.datasetFingerprints == { - assay_name: str(root[assay_name].attrs["dataset_fingerprint"]) + assert all( + root[assay_name].attrs["dataset_fingerprint"] for assay_name in result.assayNames - } - assert isinstance(root["agents"], zarr.Group) - assert root["agents"].attrs["format"] == "scarf_agent_reports" - assert ( - Path(result.zarrPath) - / "agents" - / "runs" - / result.workflowRun.workflowRunId - / "workflow.json" - ).is_file() + ) + assert "agents" not in root + assert "workflowRun" not in result.model_dump() def test_ingest_overlapping_source_destination_fails(tmp_path: Path) -> None: @@ -581,8 +550,7 @@ def make_writer(reader: FakeReader, *, zarr_loc: str) -> FakeWriter: assert created_writers[0].dump_calls == 1 assert result.acceptedActions[0]["op"] == "MtxToZarr" assert result.acceptedActions[0]["mtxIndex"] == 1 - assert result.acceptedActions[-2]["op"] == "DataStore" - assert result.acceptedActions[-1]["op"] == "createAgentWorkflow" + assert result.acceptedActions[-1]["op"] == "DataStore" assert destination.is_dir() @@ -764,8 +732,7 @@ def make_writer(reader: FakeReader, *, zarr_loc: str) -> FakeWriter: assert created_writers[0].reader is created_readers[0] assert created_writers[0].dump_calls == 1 assert result.acceptedActions[0]["op"] == "SeuratToZarr" - assert result.acceptedActions[-2]["op"] == "DataStore" - assert result.acceptedActions[-1]["op"] == "createAgentWorkflow" + assert result.acceptedActions[-1]["op"] == "DataStore" assert destination.is_dir() @@ -913,9 +880,8 @@ def make_writer( assert created_writers[0].assay_name == "ADT" assert created_writers[0].dump_calls == 1 assert result.acceptedActions[0]["op"] == "LoomToZarr" - assert result.acceptedActions[-2]["op"] == "DataStore" - assert result.acceptedActions[-2]["defaultAssay"] == "ADT" - assert result.acceptedActions[-1]["op"] == "createAgentWorkflow" + assert result.acceptedActions[-1]["op"] == "DataStore" + assert result.acceptedActions[-1]["defaultAssay"] == "ADT" assert destination.is_dir() @@ -1207,10 +1173,8 @@ def test_ingest_h5ad_still_initializes_qc(tmp_path: Path) -> None: dest = tmp_path / "out.zarr" result = ingest(path=path, zarrPath=dest) assert result.status == "done" - assert result.workflowRun is not None - assert result.acceptedActions[-2]["op"] == "DataStore" - assert result.acceptedActions[-2]["zarrMode"] == "r+" - assert result.acceptedActions[-1]["op"] == "createAgentWorkflow" + assert result.acceptedActions[-1]["op"] == "DataStore" + assert result.acceptedActions[-1]["zarrMode"] == "r+" from scarf.datastore.datastore import DataStore store = DataStore(str(dest), default_assay="RNA", nthreads=1) diff --git a/tests/test_agent_ingest_manifest.py b/tests/test_agent_ingest_manifest.py new file mode 100644 index 00000000..f5dd370b --- /dev/null +++ b/tests/test_agent_ingest_manifest.py @@ -0,0 +1,443 @@ +"""Tests for read-only H5AD decision manifests.""" + +from hashlib import sha256 +from pathlib import Path + +import h5py +import numpy as np +from scipy.sparse import csr_matrix + +from scarf.agent.ingest import ingest +from scarf.agent.ingest.manifest import inspect_h5ad_manifest + + +def _write_sparse_group( + h5: h5py.File | h5py.Group, + key: str, + values: np.ndarray, +) -> None: + matrix = csr_matrix(values) + group = h5.create_group(key) + group.attrs["encoding-type"] = "csr_matrix" + group.attrs["shape"] = values.shape + group.create_dataset("data", data=matrix.data) + group.create_dataset("indices", data=matrix.indices) + group.create_dataset("indptr", data=matrix.indptr) + + +def _write_categorical( + table: h5py.Group, + name: str, + *, + codes: list[int], + categories: list[str], +) -> None: + group = table.create_group(name) + group.attrs["encoding-type"] = "categorical" + group.create_dataset("codes", data=np.asarray(codes, dtype=np.int16)) + group.create_dataset( + "categories", + data=np.asarray([value.encode() for value in categories]), + ) + + +def _write_metadata_table( + h5: h5py.File, + key: str, + *, + n_rows: int, + feature_types: list[str] | None = None, + filtered: list[bool] | None = None, +) -> h5py.Group: + table = h5.create_group(key) + table.attrs["_index"] = "_index" + table.create_dataset( + "_index", + data=np.asarray( + [f"{key.replace('/', '_')}-{i}".encode() for i in range(n_rows)] + ), + ) + if feature_types is not None: + table.create_dataset( + "feature_name", + data=np.asarray([f"gene-{i}".encode() for i in range(n_rows)]), + ) + table.create_dataset( + "feature_types", + data=np.asarray([value.encode() for value in feature_types]), + ) + if filtered is not None: + table.create_dataset( + "feature_is_filtered", + data=np.asarray(filtered, dtype=bool), + ) + return table + + +def _write_cellxgene_h5ad(path: Path) -> None: + with h5py.File(path, mode="w") as h5: + _write_sparse_group( + h5, + "X", + np.asarray( + [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]], + dtype=np.float32, + ), + ) + _write_sparse_group( + h5, + "raw/X", + np.asarray( + [[1, 0, 3, 0], [0, 2, 4, 1], [5, 0, 0, 2]], + dtype=np.uint16, + ), + ) + layers = h5.create_group("layers") + _write_sparse_group( + layers, + "log1p", + np.asarray( + [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]], + dtype=np.float32, + ), + ) + + obs = _write_metadata_table(h5, "obs", n_rows=3) + _write_categorical( + obs, + "assay_ontology_term_id", + codes=[0, 0, 0], + categories=["EFO:0009922"], + ) + _write_categorical( + obs, + "suspension_type", + codes=[0, 0, 0], + categories=["cell"], + ) + _write_categorical( + obs, + "organism_ontology_term_id", + codes=[0, 0, 0], + categories=["NCBITaxon:9606"], + ) + _write_categorical( + obs, + "donor_id", + codes=[0, 0, 1], + categories=["donor-a", "donor-b"], + ) + _write_categorical( + obs, + "batch", + codes=[0, 0, 1], + categories=["batch-a", "batch-b"], + ) + _write_categorical( + obs, + "cell_type", + codes=[0, 1, 1], + categories=["secret-author-label-a", "secret-author-label-b"], + ) + _write_categorical( + obs, + "cell_type_ontology_term_id", + codes=[0, 1, 1], + categories=["CL:0000001", "CL:0000002"], + ) + _write_categorical( + obs, + "leiden", + codes=[0, 1, 1], + categories=["cluster-zero", "cluster-one"], + ) + obs.create_dataset( + "is_primary_data", + data=np.asarray([True, False, True], dtype=bool), + ) + + _write_metadata_table( + h5, + "var", + n_rows=3, + feature_types=["Gene Expression"] * 3, + filtered=[False, False, True], + ) + _write_metadata_table( + h5, + "raw/var", + n_rows=4, + feature_types=["Gene Expression"] * 4, + filtered=[False, True, False, True], + ) + obsm = h5.create_group("obsm") + obsm.create_dataset("X_umap", data=np.zeros((3, 2), dtype=np.float32)) + uns = h5.create_group("uns") + uns.create_dataset("schema_version", data=np.bytes_("7.1.0")) + uns.create_dataset( + "schema_reference", + data=np.bytes_("https://example.invalid/schema/7.1.0"), + ) + uns.create_dataset("batch_condition", data=np.asarray([b"batch"])) + uns.create_dataset( + "leiden_colors", + data=np.asarray([b"#000000", b"#ffffff"]), + ) + uns.create_group("neighbors") + + +def _write_simple_h5ad( + path: Path, + x: np.ndarray, + *, + raw: np.ndarray | None = None, + feature_type: str = "Gene Expression", +) -> None: + with h5py.File(path, mode="w") as h5: + _write_sparse_group(h5, "X", x) + obs = _write_metadata_table(h5, "obs", n_rows=x.shape[0]) + _write_categorical( + obs, + "donor_id", + codes=list(range(x.shape[0])), + categories=[f"d{i}" for i in range(x.shape[0])], + ) + _write_metadata_table( + h5, + "var", + n_rows=x.shape[1], + feature_types=[feature_type] * x.shape[1], + ) + if raw is not None: + _write_sparse_group(h5, "raw/X", raw) + _write_metadata_table( + h5, + "raw/var", + n_rows=raw.shape[1], + feature_types=[feature_type] * raw.shape[1], + ) + + +def test_h5ad_manifest_selects_one_raw_matrix_and_holds_out_labels( + tmp_path: Path, +) -> None: + path = tmp_path / "discover.h5ad" + _write_cellxgene_h5ad(path) + before = path.stat() + + manifest = inspect_h5ad_manifest( + path, + source_uri="cxg://collection/dataset@7.1.0", + chunk_values=2, + ) + + assert manifest.decision.status == "supported" + assert manifest.decision.selectedMatrixKey == "raw/X" + assert manifest.selectedFeatureMetadataKey == "raw/var" + assert manifest.nCells == 3 + assert manifest.nFeatures == 4 + assert manifest.sourceSha256 == sha256(path.read_bytes()).hexdigest() + assert manifest.sourceSizeBytes == path.stat().st_size + assert manifest.sourceUri == "cxg://collection/dataset@7.1.0" + assert [item.key for item in manifest.matrixCandidates if item.selected] == [ + "raw/X" + ] + assert {item.key for item in manifest.matrixCandidates} == { + "X", + "layers/log1p", + "raw/X", + } + + visible_obs = {column.name: column for column in manifest.obs.columns} + assert "cell_type" not in visible_obs + assert "cell_type_ontology_term_id" not in visible_obs + assert "leiden" not in visible_obs + assert manifest.obs.heldOutAuthorColumnCount == 3 + assert manifest.labelBenchmarkEligible is True + assert manifest.assayMetadata is not None + assert manifest.assayMetadata.domainValues == ["EFO:0009922"] + assert manifest.suspensionMetadata is not None + assert manifest.suspensionMetadata.domainValues == ["cell"] + assert manifest.organismMetadata is not None + assert manifest.organismMetadata.domainValues == ["NCBITaxon:9606"] + assert manifest.declaredBatchColumns == ["batch"] + serialized = manifest.model_dump_json() + assert "secret-author-label" not in serialized + assert "cluster-zero" not in serialized + + assert manifest.inventory.layers == ["log1p"] + assert manifest.inventory.obsm == ["X_umap"] + assert manifest.inventory.priorEmbeddings == ["obsm/X_umap"] + assert "leiden_colors" not in manifest.inventory.uns + assert manifest.inventory.heldOutUnsItemCount == 1 + assert manifest.priorFiltering.cellXGeneSchemaDetected is True + assert manifest.priorFiltering.cellXGeneSchemaVersion == "7.1.0" + assert manifest.priorFiltering.cellSetStatus == "publishedCellsOnly" + assert manifest.priorFiltering.rawCountsAvailable is True + assert manifest.priorFiltering.originalDropletsAvailable is False + assert manifest.priorFiltering.filteredFeatureCount == 2 + assert manifest.priorFiltering.nonPrimaryCellCount == 1 + assert path.stat().st_mtime_ns == before.st_mtime_ns + + +def test_h5ad_manifest_preservation_policy_exposes_labels_but_disables_benchmark( + tmp_path: Path, +) -> None: + path = tmp_path / "discover.h5ad" + _write_cellxgene_h5ad(path) + + manifest = inspect_h5ad_manifest(path, author_label_policy="preservation") + + columns = {column.name: column for column in manifest.obs.columns} + assert columns["cell_type"].domainValues == [ + "secret-author-label-a", + "secret-author-label-b", + ] + assert columns["leiden"].domainValues == ["cluster-zero", "cluster-one"] + assert manifest.obs.heldOutAuthorColumnCount == 0 + assert manifest.labelBenchmarkEligible is False + assert "leiden_colors" in manifest.inventory.uns + + +def test_h5ad_manifest_abstains_for_normalized_only_source(tmp_path: Path) -> None: + path = tmp_path / "normalized.h5ad" + _write_simple_h5ad( + path, + np.asarray([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32), + ) + + manifest = inspect_h5ad_manifest(path) + + assert manifest.decision.status == "abstained" + assert manifest.decision.reasonCode == "normalizedOnly" + assert manifest.decision.selectedMatrixKey is None + assert manifest.selectedFeatureMetadataKey is None + assert manifest.priorFiltering.rawCountsAvailable is False + assert not any(candidate.selected for candidate in manifest.matrixCandidates) + + +def test_h5ad_manifest_prefers_authoritative_raw_x_over_integer_x( + tmp_path: Path, +) -> None: + path = tmp_path / "raw-authoritative.h5ad" + _write_simple_h5ad( + path, + np.asarray([[1, 2], [3, 4]], dtype=np.uint16), + raw=np.asarray([[5, 6, 7], [8, 9, 10]], dtype=np.uint16), + ) + + manifest = inspect_h5ad_manifest(path) + + assert manifest.decision.status == "supported" + assert manifest.decision.selectedMatrixKey == "raw/X" + assert [ + candidate.key for candidate in manifest.matrixCandidates if candidate.selected + ] == ["raw/X"] + + +def test_h5ad_manifest_requires_selection_for_conflicting_count_layer( + tmp_path: Path, +) -> None: + path = tmp_path / "ambiguous.h5ad" + _write_simple_h5ad( + path, + np.asarray([[1, 2], [3, 4]], dtype=np.uint16), + raw=np.asarray([[5, 6, 7], [8, 9, 10]], dtype=np.uint16), + ) + with h5py.File(path, mode="a") as h5: + layers = h5.create_group("layers") + _write_sparse_group( + layers, + "counts", + np.asarray([[11, 12], [13, 14]], dtype=np.uint16), + ) + + unresolved = inspect_h5ad_manifest(path) + resolved = inspect_h5ad_manifest(path, matrix_key="raw/X") + + assert unresolved.decision.status == "needsInput" + assert unresolved.decision.reasonCode == "ambiguousCountMatrices" + assert unresolved.decision.options == ["X", "layers/counts", "raw/X"] + assert not any(candidate.selected for candidate in unresolved.matrixCandidates) + assert resolved.decision.status == "supported" + assert resolved.decision.selectedMatrixKey == "raw/X" + + +def test_manifest_selected_matrix_key_is_accepted_by_ingest(tmp_path: Path) -> None: + path = tmp_path / "raw-authoritative.h5ad" + _write_simple_h5ad( + path, + np.asarray([[1, 2], [3, 4]], dtype=np.uint16), + raw=np.asarray([[5, 6, 7], [8, 9, 10]], dtype=np.uint16), + ) + manifest = inspect_h5ad_manifest(path) + + result = ingest( + path=path, + zarrPath=tmp_path / "converted.zarr", + directions={"matrixKey": manifest.decision.selectedMatrixKey}, + ) + + assert result.status == "done" + assert result.acceptedActions[0] == { + "op": "inspect_h5ad", + "path": str(path), + "matrixKey": "raw/X", + } + + +def test_h5ad_manifest_abstains_when_selected_features_have_no_rna( + tmp_path: Path, +) -> None: + path = tmp_path / "adt.h5ad" + _write_simple_h5ad( + path, + np.asarray([[1, 2], [3, 4]], dtype=np.uint16), + feature_type="Antibody Capture", + ) + + manifest = inspect_h5ad_manifest(path) + + assert manifest.decision.status == "abstained" + assert manifest.decision.reasonCode == "rnaModalityUnavailable" + assert manifest.decision.options == ["ADT"] + assert not any(candidate.selected for candidate in manifest.matrixCandidates) + + +def test_h5ad_manifest_caps_domains_while_counting_missing_values( + tmp_path: Path, +) -> None: + path = tmp_path / "domains.h5ad" + with h5py.File(path, mode="w") as h5: + _write_sparse_group( + h5, + "X", + np.asarray([[1], [2], [3], [4]], dtype=np.uint16), + ) + obs = _write_metadata_table(h5, "obs", n_rows=4) + obs.create_dataset( + "free_text_group", + data=np.asarray([b"a", b"b", b"", b"c"]), + ) + _write_metadata_table( + h5, + "var", + n_rows=1, + feature_types=["Gene Expression"], + ) + + manifest = inspect_h5ad_manifest( + path, + chunk_values=1, + max_domain_values=2, + ) + + column = next( + column for column in manifest.obs.columns if column.name == "free_text_group" + ) + assert column.missingCount == 1 + assert column.missingFraction == 0.25 + assert column.domainTruncated is True + assert column.domainSize is None + assert column.domainValues == [] + assert column.valueCounts == {} diff --git a/tests/test_agent_interpretation_boundaries.py b/tests/test_agent_interpretation_boundaries.py new file mode 100644 index 00000000..aca0f1dc --- /dev/null +++ b/tests/test_agent_interpretation_boundaries.py @@ -0,0 +1,179 @@ +"""Standalone interpretation rejects stale artifacts before any model call.""" + +from types import SimpleNamespace +from typing import Any + +import pytest +from pydantic_ai import ModelRetry + +from scarf.agent.biological_interpretation.contracts import ( + BiologicalInterpretationReport, +) +from scarf.agent.biological_interpretation.validation import ( + _prepare_biological_interpretation_dependencies, + validate_biological_interpretation_report, +) +from scarf.storage.refs import ArtifactRef +from tests.test_agent_biological_interpretation import ( + FakeStore, + artifact_model, + context, +) + + +@pytest.mark.parametrize( + ("change", "reason"), + [ + ({"cluster": None}, "exact cluster artifact"), + ({"graph_assay": "Other"}, "different assay"), + ({"max_clusters": 0}, "max_clusters"), + ({"max_markers": 0}, "max_markers"), + ({"marker_min_score": 0}, "marker_min_score"), + ({"marker_min_fraction": 2}, "marker_min_fraction"), + ({"allow_marker_search": True, "marker": None}, "marker_features is required"), + ( + { + "marker": ArtifactRef( + scope="assay", + assay="RNA", + kind="feature_selection", + artifact_id="b" * 64, + ) + }, + "marker_table", + ), + ( + { + "marker_features": ArtifactRef( + scope="assay", + assay="RNA", + kind="marker_table", + artifact_id="b" * 64, + ) + }, + "feature_selection", + ), + ( + { + "cluster": ArtifactRef( + scope="assay", + assay="RNA", + kind="marker_table", + artifact_id="b" * 64, + ) + }, + "cluster_labels or cluster_cut", + ), + ( + { + "marker_features": ArtifactRef( + scope="assay", + assay="Other", + kind="feature_selection", + artifact_id="b" * 64, + ) + }, + "different assay", + ), + ], +) +def test_interpretation_input_contract_rejects_unsupported_work( + change: dict[str, Any], + reason: str, +) -> None: + store = FakeStore() + arguments = dict( + cluster=store.cluster, + from_assay=None, + graph_assay=None, + marker_assay_type=None, + sample_column=None, + condition_column=None, + tuning_handoff=None, + experimental_handoff=None, + marker=store.marker, + marker_features=None, + allow_marker_search=False, + max_clusters=8, + max_markers=10, + marker_min_score=0.1, + marker_min_fraction=0.1, + ) + with pytest.raises((TypeError, ValueError), match=reason): + _prepare_biological_interpretation_dependencies(store, **(arguments | change)) + assert store.marker_calls == 0 + + +@pytest.mark.parametrize( + ("artifact", "fields", "reason"), + [ + ("cluster", {"exists": False}, "cluster artifact does not exist"), + ("cluster", {"complete": False}, "cluster artifact is incomplete"), + ("cluster", {"inputs": {}}, "no cell-selection input"), + ("marker", {"exists": False}, "marker artifact does not exist"), + ("marker", {"complete": False}, "marker artifact is incomplete"), + ("marker", {"inputs": {}}, "exact cluster artifact"), + ], +) +def test_interpretation_artifact_status_is_validated_before_model_execution( + monkeypatch: Any, + artifact: str, + fields: dict[str, Any], + reason: str, +) -> None: + store = FakeStore() + inspect = store.inspect_artifact + damaged_ref = getattr(store, artifact) + + def damaged(ref: ArtifactRef) -> Any: + status = inspect(ref) + if ref != damaged_ref: + return status + return SimpleNamespace( + **( + { + "exists": status.exists, + "complete": status.complete, + "inputs": status.inputs, + } + | fields + ) + ) + + monkeypatch.setattr(store, "inspect_artifact", damaged) + with pytest.raises(ValueError, match=reason): + _prepare_biological_interpretation_dependencies( + store, + cluster=store.cluster, + from_assay=None, + graph_assay=None, + marker_assay_type=None, + sample_column=None, + condition_column=None, + tuning_handoff=None, + experimental_handoff=None, + marker=store.marker, + marker_features=None, + allow_marker_search=False, + max_clusters=8, + max_markers=10, + marker_min_score=0.1, + marker_min_fraction=0.1, + ) + assert store.marker_calls == 0 + + +@pytest.mark.parametrize("field", ["clusterArtifact", "markerArtifact"]) +def test_interpretation_model_cannot_replace_exact_artifact_bindings( + field: str, +) -> None: + store = FakeStore() + deps = context(store, marker=store.marker).deps + deps.clusterValues = ["0"] + reference = store.cluster if field == "clusterArtifact" else store.marker + changed = artifact_model(reference).model_copy(update={"artifactId": "f" * 64}) + report = BiologicalInterpretationReport.model_validate( + {"status": "done", field: changed} + ) + with pytest.raises(ModelRetry, match=f"{field} does not match"): + validate_biological_interpretation_report(report, deps) diff --git a/tests/test_agent_joint_design_recovery.py b/tests/test_agent_joint_design_recovery.py new file mode 100644 index 00000000..0509baf0 --- /dev/null +++ b/tests/test_agent_joint_design_recovery.py @@ -0,0 +1,211 @@ +"""Saved descriptive joint measurements answer design questions, not effects.""" + +from copy import deepcopy +from types import SimpleNamespace + +import pandas as pd +import pytest + +from scarf.agent.experimental_context.comparisons import compare_covariates +from scarf.agent.experimental_context.contracts import CovariateProposal +from scarf.agent.experimental_context.requirements import ( + objective_evidence, + unmet_objective_requirements, +) +from tests.test_agent_objective_requirements import _repeated_design + +CONTEXT = "The observations contain repeated donors and incomplete pairing." +REQUEST = "Assess individual and combined covariates, replication, and confounding" + + +def _joint_result(*, essential=False): + result = _repeated_design() + observations = pd.DataFrame( + { + "sample": ["a0", "b0", "a1", "b2", "a3", "b3"], + "donor": ["d0", "d0", "d1", "d2", "d3", "d3"], + "tissue": ["A", "B", "A", "B", "A", "B"], + "condition": ["yes", "yes", "no", "no", "yes", "yes"], + "sex": ["female", "female", "female", "male", "female", "female"], + } + ) + cells = observations.loc[observations.index.repeat(3)].reset_index(drop=True) + result.characterization.columns.append( + {"name": "sex", "kind": "categorical", "domain": "biological"} + ) + proposal = CovariateProposal( + response="tissue", + explanatoryColumns=["condition"], + conditionedOn="sex", + observationUnit="sample", + independentUnit="donor", + rationale="Assess joint tissue and condition support within sex groups.", + purpose="designCoverage", + essential=essential, + objectiveQuote=CONTEXT, + ) + result.characterization.comparisons = [ + compare_covariates( + SimpleNamespace( + columns=list(cells.columns), fetch=lambda name: cells[name].to_numpy() + ), + result.characterization, + proposal, + selection_identity={"cells": "fixed"}, + ) + ] + return result + + +def _evidence(result, request=REQUEST): + return objective_evidence( + study_context=f"{CONTEXT} {request}", + study_objective="Describe supported populations", + experimental_result=result, + ) + + +@pytest.mark.parametrize("essential", [False, True]) +def test_joint_counts_recover_descriptive_requirement_without_erasing_limitations( + essential, +): + result = _joint_result(essential=essential) + comparison = result.characterization.comparisons[0] + original = comparison.model_dump(mode="json") + assert comparison.status == "unsupported" + assert comparison.evidence["descriptiveDesign"]["observationUnits"] == 6 + assert comparison.evidence["descriptiveDesign"]["independentUnits"] == 4 + assert comparison.evidence["missingCells"] == 0 + assert ( + comparison.evidence["descriptiveDesign"]["pairedCoverage"]["tissue"][ + "completePairs" + ] + == 2 + ) + + requirements, coverage = _evidence(result) + + assert len(requirements) == 2 + assert requirements[1].essential + assert requirements[1].kind == "designCoverage" + assert coverage[1].status == "computed" + assert coverage[1].evidenceIds == [comparison.evidenceId] + assert "withinIndependentUnitComparisonsAreUnsupported" in coverage[1].reasons + assert any( + "association method remains unsupported" in x for x in coverage[1].reasons + ) + assert not unmet_objective_requirements(requirements, coverage) + assert comparison.model_dump(mode="json") == original + + +@pytest.mark.parametrize( + "damage", + [ + "missingTable", + "emptyTable", + "notComputed", + "missingColumn", + "missingGroup", + "invalidUnitCount", + "fractionalUnitCount", + "changedKind", + "changedDomain", + "missingKinds", + ], +) +def test_joint_descriptive_recovery_requires_exact_completed_measurements(damage): + result = _joint_result() + comparison = result.characterization.comparisons[0] + table = comparison.evidence["descriptiveDesign"] + if damage == "missingTable": + table.pop("jointGroupSupport") + elif damage == "emptyTable": + table["jointGroupSupport"] = [] + elif damage == "notComputed": + table["status"] = "unsupported" + elif damage == "missingColumn": + for row in table["jointGroupSupport"]: + row["groups"].pop("sex") + elif damage == "missingGroup": + table["jointGroupSupport"].pop() + elif damage == "invalidUnitCount": + table["jointGroupSupport"][0]["independentUnits"] = 100 + elif damage == "fractionalUnitCount": + table["jointGroupSupport"][0]["independentUnits"] = 1.5 + elif damage in {"changedKind", "changedDomain"}: + field, value = ( + ("kind", "continuous") if damage == "changedKind" else ("domain", "ignore") + ) + result.characterization.columns[-1][field] = value + else: + comparison.evidence.pop("columnKinds") + + requirements, coverage = _evidence(result) + + pending = [ + r for r in requirements if r.requirementId.startswith("requestedDesign:") + ] + assert len(pending) == 1 + assert pending[0].essential + assert unmet_objective_requirements(requirements, coverage) + + +def test_marginal_counts_cannot_replace_a_joint_design(): + requirements, coverage = _evidence(_repeated_design()) + assert any(r.requirementId.startswith("requestedDesign:") for r in requirements) + assert unmet_objective_requirements(requirements, coverage) + + +@pytest.mark.parametrize( + "question, purpose", + [ + ( + "Assess joint associations of tissue and condition within sex", + "designCoverage", + ), + ("Assess joint associations of tissue and condition within sex", "association"), + ( + "Estimate combined effects of tissue and condition within sex", + "designCoverage", + ), + ( + "Estimate combined effects of tissue and condition within sex", + "effectEstimation", + ), + ], +) +def test_descriptive_recovery_does_not_answer_an_association_or_effect( + question, purpose +): + result = _joint_result(essential=True) + result.characterization.comparisons[0].proposal.purpose = purpose + requirements, coverage = _evidence(result, question) + assert unmet_objective_requirements(requirements, coverage) + + +def test_joint_descriptive_recovery_does_not_ignore_a_requested_column(): + result = _joint_result() + result.characterization.columns.append( + {"name": "age", "kind": "continuous", "domain": "biological"} + ) + requirements, coverage = _evidence( + result, "Assess tissue and age jointly within sex groups" + ) + assert any(r.requirementId.startswith("requestedDesign:") for r in requirements) + assert unmet_objective_requirements(requirements, coverage) + + +def test_reusing_joint_measurements_does_not_duplicate_bounded_requirements(): + result = _joint_result() + comparison = result.characterization.comparisons[0] + result.characterization.comparisons = [] + quotes = [f"Optional observation {i}" for i in range(12)] + for index, quote in enumerate(quotes): + item = deepcopy(comparison) + item.proposal.objectiveQuote = quote + item.evidenceId = f"measured:{index}" + result.characterization.comparisons.append(item) + requirements, coverage = _evidence(result, ". ".join([*quotes, REQUEST])) + assert len(requirements) == len(coverage) == 13 + assert all(r.essential for r in requirements) + assert not unmet_objective_requirements(requirements, coverage) diff --git a/tests/test_agent_journal_recovery_boundaries.py b/tests/test_agent_journal_recovery_boundaries.py new file mode 100644 index 00000000..d3564bb1 --- /dev/null +++ b/tests/test_agent_journal_recovery_boundaries.py @@ -0,0 +1,366 @@ +"""Recovery requires the saved scientific identity, not merely a usable store.""" + +import hashlib +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import zarr + +from scarf.agent import record_io +from scarf.agent.decisions.kernel import DecisionSelection, DecisionSpec +from scarf.agent.orchestrator import AgentOrchestrator, journal +from scarf.agent.orchestrator import main as orchestrator_main +from scarf.agent.orchestrator.models import ( + AutomatedWorkflowRequest, + WorkflowNeedsInput, + WorkflowQuestion, + WorkflowStageAttempt, +) +from scarf.agent.types import AgentDataModel +from tests.agent_journal_store import memory_journal +from tests.test_agent_decision_kernel import _decision_spec + + +class ScientificEvidence(AgentDataModel): + conclusion: str = "Supported on the measured cohort" + + +class OtherEvidence(AgentDataModel): + conclusion: str = "A different scientific checkpoint" + + +@pytest.mark.parametrize( + ("payload", "reason"), + [ + ( + {"inputs": {}, "outputs": {}, "oldHistoryVersion": 1}, + "Unsupported RNA checkpoint contract", + ), + ({"inputs": [], "outputs": {}}, "inputs and outputs must be mappings"), + ({"inputs": {}, "outputs": []}, "inputs and outputs must be mappings"), + ], +) +def test_read_checkpoint_rejects_unsupported_history_without_rewriting_it( + payload: dict[str, Any], reason: str +) -> None: + store, prefix, record = memory_journal() + data = dict(payload) + if set(data) == {"inputs", "outputs"}: + data["contentSha256"] = hashlib.sha256( + record_io.canonical_json_bytes(data) + ).hexdigest() + key = journal._checkpoint_key(prefix, record.workflowRunId, "qc/evidence") + raw = record_io.canonical_json_bytes(data) + journal._write_key_once(store.zw, key, raw) + with pytest.raises(ValueError, match=reason): + journal.read_checkpoint(store, prefix, record.workflowRunId, "qc/evidence") + assert record_io.read_key(store.zw, key) == raw + + +def test_stage_recovery_rejects_wrong_report_type_checksum_and_ambiguous_ownership() -> ( + None +): + store, prefix, record = memory_journal() + started = journal._start_attempt( + store.zw, prefix, record.workflowRunId, "ingest", record, [] + ) + saved, reference = journal._save_stage_report( + store, started, ScientificEvidence(), expected_type=ScientificEvidence + ) + assert journal._recover_persisted_stage_report( + store, started, expected_type=ScientificEvidence + ) == (saved, reference) + with pytest.raises(ValueError, match="different scientific result type"): + journal._recover_persisted_stage_report( + store, started, expected_type=OtherEvidence + ) + with pytest.raises(ValueError, match="exact reference"): + journal.read_stage_evidence( + store, reference.model_copy(update={"contentSha256": "f" * 64}) + ) + duplicate = started.model_copy(update={"reportReferences": [reference, reference]}) + with pytest.raises(ValueError, match="exactly one evidence report"): + journal.load_stage_report(store, duplicate, ScientificEvidence) + assert journal.read_stage_evidence(store, reference) == saved.model_dump( + mode="json" + ) + + +@pytest.mark.parametrize( + ("answer", "reason"), + [ + ("native", "must contain decisionId"), + ( + { + "decisionId": "correction", + "optionId": "native", + "rationale": "Confounded", + "useHarmony": False, + }, + "contain exactly", + ), + ( + { + "decisionId": "some_other_decision", + "optionId": "native", + "rationale": "Confounded", + }, + "does not match decision", + ), + ( + {"decisionId": "correction", "optionId": "native", "rationale": " "}, + "non-empty rationale", + ), + ( + { + "decisionId": "correction", + "optionId": "invented", + "rationale": "Better mixing", + }, + "persisted option", + ), + ], +) +def test_resume_answer_is_bound_to_the_exact_saved_question_and_options( + answer: Any, reason: str +) -> None: + paused = WorkflowStageAttempt( + status="needsInput", + stage="experimental_context", + needsInput=WorkflowNeedsInput( + questions=[ + WorkflowQuestion( + questionId="correctionChoice", + decisionId="correction", + question="The design confounds batch with protected biology. Which offered action is justified?", + options=["native", "clarify"], + ) + ] + ), + ) + assert ( + journal._resume_answer_errors( + paused, + { + "correctionChoice": { + "decisionId": "correction", + "optionId": "native", + "rationale": "The design cannot separate batch from condition.", + } + }, + ) + == [] + ) + errors = journal._resume_answer_errors(paused, {"correctionChoice": answer}) + assert any(reason in error for error in errors) + + +@pytest.mark.parametrize( + "field", ["workflowRunId", "requestSha256", "configSha256", "contentSha256"] +) +def test_request_recovery_does_not_trust_a_corrupted_identity( + monkeypatch: pytest.MonkeyPatch, field: str +) -> None: + store, prefix, record = memory_journal() + altered = record.model_copy( + update={field: "wrong-workflow" if field == "workflowRunId" else "f" * 64} + ) + monkeypatch.setattr(journal, "_read_model", lambda *_args: altered) + with pytest.raises(ValueError, match="identity or checksum"): + journal.read_request(store.zw, prefix, record.workflowRunId) + + +@pytest.mark.parametrize( + ("change", "reason"), + [ + ({"zarrPath": None}, "no store path"), + ({"zarrPath": "other-study.zarr"}, "does not match the saved request"), + ({"workspace": "other_workspace"}, "does not match the saved request"), + ({"primaryAssay": ""}, "no selected RNA assay"), + ], +) +def test_result_store_is_validated_before_opening_datastore( + monkeypatch: pytest.MonkeyPatch, change: dict[str, Any], reason: str +) -> None: + store, _prefix, record = memory_journal() + altered = record.model_copy( + update={"request": record.request.model_copy(update=change)} + ) + calls: list[str] = [] + monkeypatch.setattr(journal.zarr, "open_group", lambda *_args, **_kwargs: store.zw) + monkeypatch.setattr(journal, "read_request", lambda *_args: altered) + monkeypatch.setattr( + journal, "DataStore", lambda *_args, **_kwargs: calls.append("open") + ) + with pytest.raises(ValueError, match=reason): + journal.open_analysis_store("analysis.zarr", record.workflowRunId) + assert calls == [] + + +@pytest.mark.parametrize("entry", ["artifact_array", "group"]) +def test_result_and_beginner_cannot_treat_an_array_as_a_workspace( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, entry: str +) -> None: + destination = tmp_path / "analysis.zarr" + root = zarr.open_group(str(destination), mode="w") + root.create_array("artifact_array", shape=(2,), dtype="i4") + root.create_group("group") + # The array is valid Zarr data, but cannot own the selected analysis history. + if entry == "artifact_array": + with pytest.raises(ValueError, match="workspace is not a group"): + journal.open_analysis_store(destination, "workflow", workspace=entry) + result = AgentOrchestrator("test-model").run( + AutomatedWorkflowRequest( + sourcePath=str(destination), + workspace=entry, + studyContext="RNA cells", + studyObjective="Characterize populations", + ) + ) + assert result.status == "failed" + assert "workspace is not a group" in " ".join(result.notes) + else: + # A real group proceeds as far as exact history lookup, not numerical work. + monkeypatch.setattr( + journal, + "read_request", + lambda *_args: (_ for _ in ()).throw(ValueError("Exact history absent")), + ) + with pytest.raises(ValueError, match="Exact history absent"): + journal.open_analysis_store(destination, "workflow", workspace=entry) + + +@pytest.mark.parametrize( + ("update", "reason"), + [ + ({"baselineOptionId": "invented"}, "baselineOptionId must reference"), + ( + {"metricPreferredOptionId": "invented"}, + "metricPreferredOptionId must reference", + ), + ({"metricPreferredOptionId": None}, "requires metricPreferredOptionId"), + ({"allowedSources": []}, "must not be empty"), + ({"allowedSources": ["agent", "agent"]}, "must not contain duplicates"), + ], +) +def test_closed_decision_spec_rejects_unregistered_authority( + update: dict[str, Any], reason: str +) -> None: + payload = _decision_spec().model_dump(mode="json") | update + with pytest.raises(ValueError, match=reason): + DecisionSpec.model_validate(payload) + + +@pytest.mark.parametrize( + ("update", "reason"), + [ + ( + {"overrideEvidenceIds": ["uncited"], "overrideOfOptionId": "baseline"}, + "must be included in evidenceIds", + ), + ({"overrideEvidenceIds": ["markers"]}, "require overrideOfOptionId"), + ({"overrideOfOptionId": "chosen"}, "must differ from selectedOptionId"), + ({"selectedOptionId": "invented option with spaces"}, "stable identifier"), + ], +) +def test_model_override_cannot_invent_a_comparator_or_uncited_support( + update: dict[str, Any], reason: str +) -> None: + payload = { + "selectedOptionId": "chosen", + "evidenceIds": ["markers"], + "rationale": "Marker programs support the selected populations.", + } | update + with pytest.raises(ValueError, match=reason): + DecisionSelection.model_validate(payload) + + +@pytest.mark.parametrize( + ("batch_columns", "reason"), + [ + ("library", "list of exact observation-column names"), + ([" "], "list of exact observation-column names"), + (["donor"], "must include the CELLxGENE"), + ], +) +def test_manifest_declared_batch_cannot_be_overridden_before_context_enrichment( + monkeypatch: pytest.MonkeyPatch, batch_columns: Any, reason: str +) -> None: + runner = AgentOrchestrator("test-model") + monkeypatch.setattr(runner, "_reuse_or_resume", lambda _request: None) + monkeypatch.setattr( + orchestrator_main, + "inspect_h5ad_manifest", + lambda *_args, **_kwargs: SimpleNamespace(declaredBatchColumns=["library"]), + ) + result = runner.run( + AutomatedWorkflowRequest( + sourcePath="study.h5ad", + studyContext="Two library preparations", + studyObjective="Compare populations", + experimentalDirections={"batchColumns": batch_columns}, + ) + ) + assert result.status == "failed" + assert result.currentStage == "ingest" + assert reason in " ".join(result.notes) + assert result.workflowRunId is None + + +@pytest.mark.parametrize( + ("manifest_status", "input_policy", "expected"), + [ + ("needsInput", "unattended", "abstained"), + ("needsInput", "pause", "needsInput"), + ("abstained", "unattended", "abstained"), + ], +) +def test_ambiguous_counts_never_reach_expensive_analysis( + monkeypatch: pytest.MonkeyPatch, + manifest_status: str, + input_policy: str, + expected: str, +) -> None: + from scarf.agent.orchestrator import AutomatedWorkflowConfig + + runner = AgentOrchestrator( + "test-model", config=AutomatedWorkflowConfig(inputPolicy=input_policy) + ) + monkeypatch.setattr(runner, "_reuse_or_resume", lambda _request: None) + manifest = SimpleNamespace( + declaredBatchColumns=["library"], + decision=SimpleNamespace( + status=manifest_status, + summary="The count matrix is ambiguous", + options=["X", "raw/X"], + evidenceIds=["counts:X", "counts:raw"], + ), + priorFiltering=SimpleNamespace( + limitations=["Only published cells are available"] + ), + ) + seen: list[Any] = [] + + def inspect(*args: Any, **kwargs: Any) -> Any: + seen.append((args, kwargs)) + return manifest + + monkeypatch.setattr(orchestrator_main, "inspect_h5ad_manifest", inspect) + result = runner.run( + AutomatedWorkflowRequest( + sourcePath="study.h5ad", + studyContext="RNA libraries", + studyObjective="Assess populations", + ) + ) + assert result.status == expected + assert result.currentStage == "ingest" + assert result.limitations == ["Only published cells are available"] + assert len(seen) == 1 + assert result.workflowRunId is None + if expected == "needsInput": + assert result.needsInput.questions[0].options == ["X", "raw/X"] + else: + assert "count" in " ".join(result.notes) diff --git a/tests/test_agent_journal_usage.py b/tests/test_agent_journal_usage.py new file mode 100644 index 00000000..b33ed88a --- /dev/null +++ b/tests/test_agent_journal_usage.py @@ -0,0 +1,181 @@ +"""Saved invocation usage survives failure and is counted once in derived views.""" + +from typing import Any + +import pytest + +from scarf.agent.orchestrator import AgentOrchestrator, journal +from scarf.agent.orchestrator.models import AutomatedWorkflowRequest, WorkflowIdentity +from scarf.agent.types import AgentDataModel, AgentRunInfo, AgentUsageInfo +from tests.agent_journal_store import memory_journal + + +class MeasuredReport(AgentDataModel): + runInfo: AgentRunInfo + + +def test_snapshot_counts_failed_and_completed_invocations_without_report_duplicates() -> ( + None +): + store, prefix, record = memory_journal("analysis") + inputs = {"selection": "cells", "metadata": {"age": "continuous"}} + save = journal.model_attempt_callback( + store, prefix, record.workflowRunId, "ingest/review", inputs + ) + complete = AgentRunInfo( + runId="complete", + agentName="test", + status="done", + usage=AgentUsageInfo( + requests=2, + inputTokens=20, + outputTokens=4, + totalTokens=24, + availability="reported", + ), + ) + failed = AgentRunInfo( + runId="failed", + agentName="test", + status="failed", + usage=AgentUsageInfo( + requests=1, + inputTokens=7, + outputTokens=1, + totalTokens=8, + availability="partial", + ), + error="RuntimeError: Provider outage after a measured response", + ) + unavailable = AgentRunInfo( + runId="no-response", + agentName="test", + status="failed", + usage=AgentUsageInfo(availability="unavailable"), + ) + for info in (failed, complete, unavailable, complete): + save(info) + started = journal._start_attempt( + store.zw, prefix, record.workflowRunId, "ingest", record, [], inputs=inputs + ) + _, reference = journal._save_stage_report( + store, started, MeasuredReport(runInfo=complete), expected_type=MeasuredReport + ) + outcome = journal._complete_attempt( + started, + status="done", + report_references=[reference], + outputs={"runInfo": complete.model_dump(mode="json")}, + ) + journal._save_outcome(store.zw, prefix, outcome) + snapshot = journal.analysis_snapshot(store, record.workflowRunId) + assert len(snapshot["modelAttempts"]) == 4 + usage = snapshot["modelUsage"] + assert usage["invocations"] == 3 + assert usage["failedInvocations"] == 2 + assert usage["requests"] == 3 + assert usage["inputTokens"] == 27 + assert usage["totalTokens"] == 32 + assert usage["availability"] == "partial" + assert usage["partialUsageInvocations"] == 1 + assert usage["unavailableUsageInvocations"] == 1 + assert snapshot == journal.analysis_snapshot(store, record.workflowRunId) + + +def test_usage_without_reported_availability_is_not_assumed_complete() -> None: + legacy = AgentRunInfo( + modelName="legacy-model", + usage=AgentUsageInfo(inputTokens=17, totalTokens=17), + ).model_dump(mode="json") + views = [{"runInfo": legacy}, {"report": {"runInfo": legacy}}, {"runInfo": {}}] + usage = journal._model_usage([], views) + assert usage["invocations"] == 1 + assert usage["inputTokens"] == 17 + assert usage["availability"] == "partial" + assert usage["unspecifiedUsageInvocations"] == 1 + assert journal._model_usage([], [])["availability"] == "unavailable" + + +def test_uncaught_stage_failure_preserves_cause_and_invocation_usage( + monkeypatch: Any, +) -> None: + store, prefix, record = memory_journal() + started = journal._start_attempt( + store.zw, prefix, record.workflowRunId, "ingest", record, [] + ) + info = AgentRunInfo( + runId="failed-call", + status="failed", + usage=AgentUsageInfo(inputTokens=9, totalTokens=9, availability="partial"), + ) + error = RuntimeError("The bounded decision could not be completed") + error.__cause__ = ValueError("Required marker comparison is missing") + error.agent_run_info = info + + def fail(*_args: Any, **_kwargs: Any) -> None: + raise error + + runner = AgentOrchestrator("test-model") + monkeypatch.setattr(runner, "_execute_stages", fail) + result = runner._continue( + store, WorkflowIdentity(record.workflowRunId, None), record, answers={} + ) + assert result.status == "failed" + assert result.currentStage == "ingest" + assert "Required marker comparison is missing" in " ".join(result.notes) + assert result.workflowRunId == record.workflowRunId + outcome = journal._stage_outcomes(store.zw, prefix, record.workflowRunId, "ingest")[ + 0 + ] + assert outcome.attemptId == started.attemptId + assert outcome.outputs["runInfo"] == info.model_dump(mode="json") + snapshot = journal.analysis_snapshot(store, record.workflowRunId) + assert snapshot["modelUsage"]["inputTokens"] == 9 + + +@pytest.mark.parametrize("persistence_fails", [False, True]) +def test_structured_failure_preserves_root_cause_before_persistence( + monkeypatch: Any, persistence_fails: bool +) -> None: + runner = AgentOrchestrator("test-model") + request = AutomatedWorkflowRequest( + sourcePath="missing.h5ad", + studyContext="Human cells", + studyObjective="Assess populations", + ) + error = RuntimeError("Could not open source") + error.__cause__ = OSError("Exact missing file") + if persistence_fails: + error.add_note("Saving failure also failed") + + def fail(_request: Any) -> None: + raise error + + monkeypatch.setattr(runner, "_run", fail) + result = runner.run(request) + assert result.status == "failed" + assert "OSError: Exact missing file" in result.notes[0] + if persistence_fails: + assert "Saving failure also failed" in result.notes[0] + + +def test_report_shows_saved_failed_usage_without_model_or_numerical_calls() -> None: + from scarf.agent.report.artifacts import scientific_summary + from scarf.agent.report.rendering import render_analysis_document + from tests.test_agent_report import display_payload, snapshot + + state = snapshot() + state["modelUsage"] = { + "invocations": 4, + "failedInvocations": 1, + "requests": 8, + "validationRetries": 2, + "inputTokens": 1234, + "outputTokens": 56, + "availability": "partial", + } + rendered = render_analysis_document(scientific_summary(state) | display_payload()) + assert "4 invocations, 1 failed" in rendered + assert "1,234 input and 56 output" in rendered + assert "missing usage is not zero" in rendered + assert "2 validation corrections" in rendered diff --git a/tests/test_agent_notebook_interrupt.py b/tests/test_agent_notebook_interrupt.py new file mode 100644 index 00000000..3174bd81 --- /dev/null +++ b/tests/test_agent_notebook_interrupt.py @@ -0,0 +1,215 @@ +"""Notebook interrupts cancel real worker tasks without waiting for model output.""" + +import asyncio +from concurrent.futures import ThreadPoolExecutor +from threading import Event +from typing import Any + +import pytest +from pydantic_ai.messages import ModelResponse, ToolCallPart +from pydantic_ai.models.function import FunctionModel +from pydantic_ai.usage import RequestUsage + +from scarf.agent.config import agent_exec +from scarf.agent.types import AgentDataModel + + +class Choice(AgentDataModel): + value: int + + +@pytest.mark.parametrize("callback_fails", [False, True]) +def test_notebook_interrupt_cancels_pending_request_and_keeps_partial_usage( + monkeypatch: Any, + callback_fails: bool, +) -> None: + waiting = Event() + cancelled = Event() + calls = {"requests": 0, "measurements": 0, "completedResponses": 0} + attempts = [] + original_interrupt = KeyboardInterrupt("Notebook interrupted") + + async def measure() -> int: + calls["measurements"] += 1 + return 2 + + async def respond(_messages: Any, _info: Any) -> ModelResponse: + calls["requests"] += 1 + if calls["requests"] == 1: + calls["completedResponses"] += 1 + return ModelResponse( + parts=[ToolCallPart("measure", {})], + usage=RequestUsage(input_tokens=11, output_tokens=2), + ) + waiting.set() + try: + await asyncio.Event().wait() + raise AssertionError("A pending request must not produce a completion") + except asyncio.CancelledError: + cancelled.set() + raise + + def save_attempt(info: Any) -> None: + attempts.append(info) + if callback_fails: + raise RuntimeError("Saving cancellation evidence failed") + + class InterruptingPool(ThreadPoolExecutor): + def submit(self, fn: Any, *args: Any, **kwargs: Any) -> Any: + future = super().submit(fn, *args, **kwargs) + original_result = future.result + first = True + + def result(timeout: float | None = None) -> Any: + nonlocal first + if first: + first = False + assert waiting.wait(10), "The simulated provider did not start" + raise original_interrupt + return original_result(timeout=timeout) + + future.result = result + return future + + monkeypatch.setattr(agent_exec, "ThreadPoolExecutor", InterruptingPool) + + async def notebook() -> BaseException: + with pytest.raises(KeyboardInterrupt) as caught: + agent_exec.run_agent_sync( + model=FunctionModel(respond), + output_type=Choice, + system_prompt="Interpret a measured count.", + user_prompt="Measure once, then decide.", + tools=[measure], + on_attempt=save_attempt, + ) + return caught.value + + interrupted = asyncio.run(notebook()) + assert interrupted is original_interrupt + assert cancelled.is_set() + assert calls == {"requests": 2, "measurements": 1, "completedResponses": 1} + assert len(attempts) == 1 + assert interrupted.agent_run_info is attempts[0] + assert attempts[0].status == "failed" + assert attempts[0].usage.availability == "partial" + assert attempts[0].usage.inputTokens == 11 + assert attempts[0].usage.requests == 2 + assert attempts[0].errorType == "CancelledError" + if callback_fails: + assert "Saving cancellation evidence failed" in agent_exec.describe_agent_error( + interrupted + ) + + +def test_notebook_interrupt_before_worker_ready_never_starts_a_provider_request( + monkeypatch: Any, +) -> None: + ready = Event() + release = Event() + attempts = [] + requests = [] + original_interrupt = KeyboardInterrupt("Interrupted before worker startup") + + async def respond(_messages: Any, _info: Any) -> ModelResponse: + requests.append("unexpected") + raise AssertionError("The interrupted invocation must not contact the provider") + + class StartingPool(ThreadPoolExecutor): + def submit(self, fn: Any, *args: Any, **kwargs: Any) -> Any: + def start() -> Any: + ready.set() + assert release.wait(10), "Worker startup was not released for cleanup" + return fn(*args, **kwargs) + + future = super().submit(start) + original_result = future.result + first = True + + def result(timeout: float | None = None) -> Any: + nonlocal first + if first: + first = False + assert ready.wait(10), "The worker thread did not start" + raise original_interrupt + release.set() + return original_result(timeout=timeout) + + future.result = result + return future + + monkeypatch.setattr(agent_exec, "ThreadPoolExecutor", StartingPool) + + async def notebook() -> BaseException: + with pytest.raises(KeyboardInterrupt) as caught: + agent_exec.run_agent_sync( + model=FunctionModel(respond), + output_type=Choice, + system_prompt="Interpret evidence.", + user_prompt="Assess the observed count.", + on_attempt=attempts.append, + ) + return caught.value + + interrupted = asyncio.run(notebook()) + assert interrupted is original_interrupt + assert requests == [] + assert len(attempts) == 1 + assert interrupted.agent_run_info is attempts[0] + assert attempts[0].status == "failed" + assert attempts[0].errorType == "CancelledError" + assert attempts[0].usage.availability == "unavailable" + assert "before model execution" in attempts[0].error + + +def test_notebook_interrupt_racing_with_completed_worker_preserves_actual_usage( + monkeypatch: Any, +) -> None: + attempts = [] + original_interrupt = KeyboardInterrupt("Interrupted after completion") + + async def respond(_messages: Any, info: Any) -> ModelResponse: + return ModelResponse( + parts=[ToolCallPart(info.output_tools[0].name, {"value": 1})], + usage=RequestUsage(input_tokens=7, output_tokens=2), + ) + + class CompletedPool(ThreadPoolExecutor): + def submit(self, fn: Any, *args: Any, **kwargs: Any) -> Any: + future = super().submit(fn, *args, **kwargs) + original_result = future.result + first = True + + def result(timeout: float | None = None) -> Any: + nonlocal first + completed = original_result(timeout=timeout) + if first: + first = False + raise original_interrupt + return completed + + future.result = result + return future + + monkeypatch.setattr(agent_exec, "ThreadPoolExecutor", CompletedPool) + + async def notebook() -> BaseException: + with pytest.raises(KeyboardInterrupt) as caught: + agent_exec.run_agent_sync( + model=FunctionModel(respond), + output_type=Choice, + system_prompt="Interpret evidence.", + user_prompt="Use the measured count.", + on_attempt=attempts.append, + ) + return caught.value + + interrupted = asyncio.run(notebook()) + assert interrupted is original_interrupt + assert len(attempts) == 1 + assert attempts[0].status == "done" + assert attempts[0].usage.inputTokens == 7 + assert interrupted.agent_run_info is attempts[0] + assert "completed before cancellation" in agent_exec.describe_agent_error( + interrupted + ) diff --git a/tests/test_agent_objective_requirements.py b/tests/test_agent_objective_requirements.py new file mode 100644 index 00000000..556a333b --- /dev/null +++ b/tests/test_agent_objective_requirements.py @@ -0,0 +1,226 @@ +"""Measured design coverage is required before an objective can be completed.""" + +from types import SimpleNamespace + +import pandas as pd +import pytest +from pydantic import ValidationError + +from scarf.agent.experimental_context.comparisons import compare_covariates +from scarf.agent.experimental_context.contracts import ( + CovariateCharacterization, + CovariateProposal, + ExperimentalContextDecision, + InferenceUnit, +) +from scarf.agent.experimental_context.study import ( + StudyContract, + build_study_contract, + validate_objective_evidence, +) + + +def _repeated_design(*, purpose: str = "designCoverage", essential: bool = True): + observations = pd.DataFrame( + { + "sample": ["a0", "b0", "a1", "b2", "a3", "b3"], + "donor": ["d0", "d0", "d1", "d2", "d3", "d3"], + "tissue": ["A", "B", "A", "B", "A", "B"], + "condition": ["yes", "yes", "no", "no", "yes", "yes"], + } + ) + cells = observations.loc[observations.index.repeat(3)].reset_index(drop=True) + characterization = CovariateCharacterization( + status="done", + columns=[ + { + "name": name, + "kind": "categorical", + "domain": "design" if name == "sample" else "biological", + } + for name in cells.columns + ], + coefficients=[ + { + "name": "tissue", + "scope": "betweenUnit", + "kind": "categorical", + "observationUnit": "sample", + "independentUnit": "donor", + "groupOrder": ["A", "B"], + "unitLevelCounts": { + "observationUnit": {"levels": 6}, + "independentUnit": {"levels": 4}, + }, + "replication": {"sufficient": True, "minimumPerGroup": 3}, + "pairedCoverage": { + "design": "mixedOrIncomplete", + "completePairs": 2, + "incompletePairs": 2, + }, + } + ], + ) + proposal = CovariateProposal.model_validate( + { + "response": "tissue", + "explanatoryColumns": ["condition"], + "observationUnit": "sample", + "independentUnit": "donor", + "rationale": "Assess observed tissue and condition support across donors.", + "purpose": purpose, + "essential": essential, + "objectiveQuote": "Describe supported populations", + } + ) + comparison = compare_covariates( + SimpleNamespace( + columns=list(cells.columns), fetch=lambda name: cells[name].to_numpy() + ), + characterization, + proposal, + selection_identity={"cells": "fixed"}, + ) + characterization.comparisons = [comparison] + result = SimpleNamespace( + status="done", + characterization=characterization, + batchSafety=[], + notes=[], + decision=ExperimentalContextDecision( + coefficientsOfInterest=["tissue"], + unitsOfInference={ + "tissue": InferenceUnit( + observationUnit="sample", independentUnit="donor" + ) + }, + ), + ) + return result + + +def _contract(result): + return build_study_contract( + study_context="The observations contain repeated donors and incomplete pairing.", + study_objective="Describe supported populations", + experimental_result=result, + ) + + +def test_repeated_donors_have_descriptive_coverage_without_an_association(): + result = _repeated_design() + comparison = result.characterization.comparisons[0] + assert comparison.status == "unsupported" + assert comparison.reasons == ["withinIndependentUnitComparisonsAreUnsupported"] + evidence = comparison.evidence["descriptiveDesign"] + assert evidence["observationUnits"] == 6 + assert evidence["independentUnits"] == 4 + assert evidence["groupSupport"]["tissue"] == [ + {"group": "A", "observationUnits": 3, "independentUnits": 3}, + {"group": "B", "observationUnits": 3, "independentUnits": 3}, + ] + assert evidence["pairedCoverage"]["tissue"]["completePairs"] == 2 + assert evidence["pairedCoverage"]["tissue"]["incompletePairs"] == 2 + assert evidence["sharedIndependentUnits"]["tissue"]["pairs"] == [ + {"groups": ["A", "B"], "independentUnits": 2}, + ] + assert "singleAssociations" not in comparison.evidence + contract = _contract(result) + validate_objective_evidence(contract, result) + assert [item.status for item in contract.evidenceCoverage] == [ + "computed", + "computed", + ] + assert any( + "association method remains unsupported" in text + for text in contract.limitations + ) + + +@pytest.mark.parametrize("purpose", ["association", "effectEstimation"]) +def test_descriptive_counts_cannot_satisfy_an_essential_effect_or_association(purpose): + result = _repeated_design(purpose=purpose) + contract = _contract(result) + assert contract.evidenceCoverage[1].status == "unsupported" + with pytest.raises(ValueError, match="Essential objective evidence is unresolved"): + validate_objective_evidence(contract, result) + + +def test_optional_unsupported_question_remains_an_explicit_limitation(): + result = _repeated_design(purpose="association", essential=False) + result.characterization.comparisons[ + 0 + ].proposal.objectiveQuote = ( + "The observations contain repeated donors and incomplete pairing." + ) + contract = _contract(result) + validate_objective_evidence(contract, result) + assert contract.evidenceCoverage[1].status == "unsupported" + assert any( + "no supported association or absence finding" in text + for text in contract.limitations + ) + + +def test_missing_replication_blocks_completion_even_without_proposals(): + result = _repeated_design() + result.characterization.comparisons = [] + result.characterization.coefficients[0]["replication"] = {} + contract = _contract(result) + with pytest.raises(ValueError, match="replication evidence is unavailable"): + validate_objective_evidence(contract, result) + + +def test_nonidentifiable_exact_batch_design_answers_design_question_only(): + result = _repeated_design() + result.batchSafety = [ + SimpleNamespace( + coefficient="tissue", + batchColumns=["batch"], + status="unsafe", + evidenceId="batchSafety:tissue:batch", + ) + ] + result.characterization.columns.append( + {"name": "batch", "domain": "technical", "kind": "categorical"} + ) + contract = _contract(result) + assert contract.evidenceCoverage[0].status == "nonIdentifiable" + validate_objective_evidence(contract, result) + assert contract.correctionLicense == "unsafeConfounded" + + +def test_recomputed_context_coverage_rejects_tampered_or_stale_answer(): + result = _repeated_design() + contract = _contract(result) + result.characterization.coefficients[0]["replication"] = {} + with pytest.raises(ValueError, match="differs from its measured context report"): + validate_objective_evidence(contract, result) + + +@pytest.mark.parametrize("field", ["evidenceRequirements", "evidenceCoverage"]) +def test_old_contracts_without_objective_evidence_are_incompatible(field): + values = _contract(_repeated_design()).model_dump(mode="json") + del values[field] + with pytest.raises(ValidationError, match=field): + StudyContract.model_validate(values) + + +def test_empty_requirements_and_fabricated_citations_cannot_bypass_gate(): + contract = _contract(_repeated_design()) + with pytest.raises(ValueError): + validate_objective_evidence( + contract.model_copy(update={"evidenceRequirements": []}) + ) + contract.evidenceCoverage[0].evidenceIds.append("unmeasured:claim") + with pytest.raises(ValueError, match="outside the study contract"): + validate_objective_evidence(contract) + + +def test_context_role_change_invalidates_previously_computed_design_question(): + result = _repeated_design() + result.characterization.columns[1]["domain"] = "technical" + contract = _contract(result) + assert contract.evidenceCoverage[1].status == "unsupported" + with pytest.raises(ValueError, match="different column roles"): + validate_objective_evidence(contract, result) diff --git a/tests/test_agent_orchestrator.py b/tests/test_agent_orchestrator.py index ac355a71..6cc7393d 100644 --- a/tests/test_agent_orchestrator.py +++ b/tests/test_agent_orchestrator.py @@ -1,6 +1,9 @@ """Public facade, model, and end-to-end orchestrator contracts.""" -import re +from tests.agent_examples import example +from tests.agent_comparison_examples import observed_action + +import json from pathlib import Path from typing import Any @@ -16,33 +19,36 @@ import scarf.agent as agent_module import scarf.agent.orchestrator as orchestrator_module -from scarf.agent.biological_interpretation import ( - BiologicalInterpretationReport, - ClusterCompositionEvidence, - ClusterInterpretation, - ClusterMarkerBatchEvidence, -) from scarf.agent.data_enrichment import ( DataEnrichmentReport, FeatureSelectionPolicy, StudyContextSummary, ) +from scarf.agent.decisions.kernel import DecisionSelection from scarf.agent.experimental_context import ( BatchCorrectionPlan, CellQcPlan, - CovariateEvidence, ExperimentalContextDecision, ) +from scarf.agent.parameter_tuning import ParameterTuningReport +from scarf.agent.parameter_tuning.contracts import ParameterCandidateEvaluation +from scarf.agent.orchestrator.journal import ( + analysis_snapshot, + load_checkpoint, + _ensure_orchestration_store, +) +from scarf.agent.orchestrator.rna_tuning import TuningAction from scarf.agent.orchestrator import ( AgentOrchestrator, - AssayPreprocessingPlan, - AutomatedPreprocessingPlan, AutomatedWorkflowConfig, AutomatedWorkflowRequest, - AutomatedWorkflowResult, AutomatedWorkflowResumeRequest, +) +from scarf.agent.orchestrator.models import ( + AssayPreprocessingPlan, + AutomatedPreprocessingPlan, + AutomatedWorkflowResult, FinalAnalysisHandoff, - NativeAnalysisHandoff, PreprocessedAssayHandoff, WorkflowNeedsInput, WorkflowQuestion, @@ -50,37 +56,36 @@ WorkflowStageLink, artifact_model_to_ref, ) -from scarf.agent.persistence import ( - load_agent_record, - load_agent_report, -) -from scarf.agent.parameter_tuning import ( - FinalGraphSelection, - ParameterTuningReport, -) from scarf.datastore.datastore import DataStore +from scarf.storage.refs import ArtifactRef +from scarf.storage.selections import read_stored_selection_indices from tests.test_agent_ingest import _write_h5ad _PLAN_CHECKSUM = "a" * 64 -def _rna_workflow_model() -> tuple[FunctionModel, dict[str, int]]: - state = { +def _rna_workflow_model() -> tuple[FunctionModel, dict[str, Any]]: + state: dict[str, Any] = { "enrichment": 0, "context": 0, "parameter": 0, + "pca_pauses": 0, + "pca_prompts": 0, "biology": 0, "requests": 0, } def prompt_text(messages: list[ModelMessage]) -> str: - return "\n".join( - part.content - for message in messages - for part in message.parts - if isinstance(getattr(part, "content", None), str) - ) + values: list[str] = [] + for message in messages: + for part in message.parts: + content = getattr(part, "content", None) + if isinstance(content, str): + values.append(content) + elif isinstance(content, tuple): + values.extend(item for item in content if isinstance(item, str)) + return "\n".join(values) def tool_result( messages: list[ModelMessage], @@ -93,6 +98,10 @@ def tool_result( content = part.content if isinstance(content, model_type): return content + if model_type is dict: + return ( + json.loads(content) if isinstance(content, str) else content + ) if isinstance(content, str): return model_type.model_validate_json(content) return model_type.model_validate(content) @@ -178,30 +187,20 @@ async def reply( context_evidence = tool_result( messages, "analyze_experimental_design", - CovariateEvidence, + dict, ) profile = next( value - for value in context_evidence.qcProfiles - if value.action == "globalGaussian" + for value in context_evidence["qcProfiles"] + if value["registeredProfile"] is not None ) - evidence_id = profile.evidenceId + evidence_id = profile["evidenceId"] decision = ExperimentalContextDecision( batchCorrection=BatchCorrectionPlan( action="skip", rationale="No trusted technical batch column was supplied.", evidenceIds=[evidence_id], ), - cellQc=CellQcPlan( - action=profile.action, - profileId=profile.profileId, - driverAssay=profile.driverAssay, - driverAssayType=profile.driverAssayType, - attributes=profile.attributes, - artifactMetrics=profile.artifactMetrics, - rationale="Apply the bounded global RNA QC profile.", - evidenceIds=[evidence_id], - ), rationale="No experimental covariates were supplied.", evidenceIds=[evidence_id], ) @@ -214,131 +213,91 @@ async def reply( ] ) - if ( - tools.intersection( - {"inspect_cluster_composition", "inspect_cluster_markers_batch"} + prompt = prompt_text(messages) + payload, _ = json.JSONDecoder().raw_decode(prompt[prompt.index("{") :]) + if any( + {"selectedCandidateId", "correctionNeed", "comparisonConclusions"}.issubset( + tool.parameters_json_schema.get("properties", {}) ) - or state["biology"] + for tool in info.output_tools ): - if state["biology"] == 0: - state["biology"] = 1 - return ModelResponse( - parts=[ - ToolCallPart( - tool_name="inspect_cluster_composition", - args={}, - ) - ] - ) - if state["biology"] == 1: - composition = tool_result( - messages, - "inspect_cluster_composition", - ClusterCompositionEvidence, - ) - state["biology"] = 2 - return ModelResponse( - parts=[ - ToolCallPart( - tool_name="inspect_cluster_markers_batch", - args={"cluster_ids": list(composition.clusterCounts)}, - ) - ] - ) - batch = tool_result( - messages, - "inspect_cluster_markers_batch", - ClusterMarkerBatchEvidence, - ) - interpretations = [] - for cluster in batch.clusters: - if cluster.evidenceId and cluster.markers: - marker = cluster.markers[0] - marker_name = marker.featureName or marker.featureId - interpretations.append( - ClusterInterpretation( - clusterId=cluster.clusterId, - proposedIdentity=f"{marker_name}-high RNA state", - identityIsHypothesis=True, - confidence="low", - rationale=( - f"The observed marker panel is led by {marker_name}." - ), - evidenceIds=[cluster.evidenceId], - ) - ) - state["biology"] = 3 - report = BiologicalInterpretationReport( - status="done", - clusterInterpretations=interpretations, - evidenceIds=[item.evidenceIds[0] for item in interpretations], - limitations=["Synthetic data supports marker-linked hypotheses only."], - stopReason=( - "Every cluster with returned marker evidence was reviewed." + state["pca_prompts"] += 1 + selected = next( + ( + row["candidateId"] + for row in payload["candidates"] + if row["parameters"]["leidenResolution"] == 0.5 ), + payload["currentCandidateId"], ) + if payload["comparisonCoverage"]["phase"] == "sensitivity": + selected = payload["currentCandidateId"] + action = TuningAction.model_validate( + observed_action(payload, selected=selected) + ) + if action.action == "accept" and state["pca_pauses"] == 0: + state["answer"] = action.model_dump(mode="json") + action = action.model_copy( + update={ + "action": "defer", + "rationale": "Review the completed comparisons before validating the selected combination on all retained cells.", + } + ) + state["pca_pauses"] += 1 return ModelResponse( parts=[ ToolCallPart( - tool_name=info.output_tools[0].name, - args=report.model_dump(), + tool_name=info.output_tools[0].name, args=action.model_dump() ) ] ) - - prompt = prompt_text(messages) - if state["parameter"] == 0: - match = re.search( - r'"candidateId"\s*:\s*"([A-Za-z0-9_]+)"', - prompt, + payload, _ = json.JSONDecoder().raw_decode(prompt[prompt.index("{") :]) + decision = payload["spec"] + decision_id = decision["decisionId"] + if decision_id == "pcaPrefix": + state["pca_prompts"] += 1 + evidence_by_class: dict[str, str] = {} + evidence_class_by_id: dict[str, str] = {} + for item in payload["evidence"]["evidence"]: + evidence_by_class.setdefault( + item["evidenceClass"], + item["evidenceId"], ) - assert match is not None - candidate_id = match.group(1) - evidence_id = f"candidate:{candidate_id}:clusters" - assay_report = ParameterTuningReport( - status="done", - recommendedCandidateId=candidate_id, - confidence="high", - rationale="The only authorized native branch is eligible.", - evidenceIds=[evidence_id], - stopReason="The bounded one-candidate screen completed.", + evidence_class_by_id[item["evidenceId"]] = item["evidenceClass"] + preferred = decision.get("metricPreferredOptionId") + selected = ( + next( + option for option in decision["options"] if option["status"] == "defer" ) - report = ParameterTuningReport( - status="done", - assayReports={"RNA": assay_report}, - rationale="The RNA native screen completed.", - evidenceIds=[evidence_id], - stopReason="Native selection completed.", + if decision_id == "pcaPrefix" and state["pca_pauses"] == 0 + else next( + option + for option in decision["options"] + if option["optionId"] == preferred ) - state["parameter"] = 1 - return ModelResponse( - parts=[ - ToolCallPart( - tool_name=info.output_tools[0].name, - args=report.model_dump(), - ) - ] + if preferred is not None + else next( + option + for option in decision["options"] + if option["status"] in {"apply", "skip"} ) - - match = re.search( - r'"optionId"\s*:\s*"(native:RNA:([A-Za-z0-9_]+))"', - prompt, ) - assert match is not None - option_id, candidate_id = match.groups() - evidence_id = f"native:RNA:candidate:{candidate_id}:clusters" - selection = FinalGraphSelection( - status="done", - selectedOptionId=option_id, - graphMethod="native", - nativeAssay="RNA", - nativeCandidateId=candidate_id, - markerAssay="RNA", + if decision_id == "pcaPrefix" and selected["status"] == "defer": + state["pca_pauses"] += 1 + evidence_ids = list(selected.get("requiredEvidenceIds", [])) + cited_classes = { + evidence_class_by_id[evidence_id] for evidence_id in evidence_ids + } + for evidence_class in selected["requiredEvidenceClasses"]: + if evidence_class not in cited_classes: + evidence_ids.append(evidence_by_class[evidence_class]) + selection = DecisionSelection( + selectedOptionId=selected["optionId"], + evidenceIds=evidence_ids, + rationale="Select the first eligible registered option for this test.", confidence="high", - rationale="The sole eligible native graph is selected.", - evidenceIds=[evidence_id], ) - state["parameter"] = 2 + state["parameter"] += 1 return ModelResponse( parts=[ ToolCallPart( @@ -360,7 +319,6 @@ def test_public_orchestrator_models_have_factories_and_camelcase_fields() -> Non AutomatedWorkflowResult, AutomatedWorkflowResumeRequest, FinalAnalysisHandoff, - NativeAnalysisHandoff, PreprocessedAssayHandoff, StudyContextSummary, CellQcPlan, @@ -372,7 +330,7 @@ def test_public_orchestrator_models_have_factories_and_camelcase_fields() -> Non for model in models: assert isinstance(model.get_blank(), model) - assert isinstance(model.get_example(), model) + assert isinstance(example(model), model) assert all("_" not in field_name for field_name in model.model_fields) for model in ( AutomatedPreprocessingPlan, @@ -381,36 +339,30 @@ def test_public_orchestrator_models_have_factories_and_camelcase_fields() -> Non ): assert "cellSelection" in model.model_fields assert "cellKey" not in model.model_fields - for model in (NativeAnalysisHandoff, FinalAnalysisHandoff): + for model in (FinalAnalysisHandoff,): assert "clusterColumn" not in model.model_fields assert "umapColumns" not in model.model_fields def test_orchestrator_package_preserves_the_public_facade() -> None: - assert agent_module.AgentOrchestrator is orchestrator_module.AgentOrchestrator + assert not hasattr(agent_module, "AgentOrchestrator") assert orchestrator_module.__all__ == [ "AgentOrchestrator", - "AssayPreprocessingPlan", - "AutomatedPreprocessingPlan", "AutomatedWorkflowConfig", "AutomatedWorkflowRequest", - "AutomatedWorkflowResult", "AutomatedWorkflowResumeRequest", - "FinalAnalysisHandoff", - "NativeAnalysisHandoff", - "PreprocessedAssayHandoff", - "WorkflowNeedsInput", - "WorkflowQuestion", - "WorkflowStageAttempt", - "WorkflowStageLink", - "artifact_model_to_ref", ] @pytest.mark.slow -def test_rna_h5ad_completes_public_automated_workflow(tmp_path: Path) -> None: +def test_rna_h5ad_completes_public_automated_workflow( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from scarf.agent.orchestrator import rna_tuning as tuning_module + rng = np.random.default_rng(4444) - values = rng.poisson(1.0, size=(80, 50)).astype(np.uint16) + values = rng.poisson(0.2, size=(80, 50)).astype(np.uint16) values[:40, :12] += rng.poisson(9.0, size=(40, 12)).astype(np.uint16) values[40:, 12:24] += rng.poisson(9.0, size=(40, 12)).astype(np.uint16) feature_names = [ @@ -449,61 +401,89 @@ def test_rna_h5ad_completes_public_automated_workflow(tmp_path: Path) -> None: feature_names=feature_names, ) model, state = _rna_workflow_model() + pca_diagnostic_calls: list[ArtifactRef] = [] + augment_pca = tuning_module.augment_pca_evaluations + + def track_pca_diagnostics(*args: Any, **kwargs: Any) -> Any: + pca_diagnostic_calls.append(kwargs["feature_selection"]) + return augment_pca(*args, **kwargs) + + monkeypatch.setattr(tuning_module, "augment_pca_evaluations", track_pca_diagnostics) orchestrator = AgentOrchestrator( model, - config=AutomatedWorkflowConfig( - primaryInitialCandidates=1, - secondaryInitialCandidates=1, - maxRefinedCandidatesPerAssay=0, - maxHarmonyCandidatesPerAssay=0, - integrationResolutionCandidates=1, - maxCandidateBranches=1, - minClusterCells=2, + config=AutomatedWorkflowConfig(screeningCells=60, maxScreeningCells=70), + ) + + request = AutomatedWorkflowRequest( + sourcePath=str(source), + zarrPath=str(target), + studyContext=( + "A human peripheral blood RNA study for a deterministic acceptance test." ), + studyObjective="Discover stable RNA populations.", + primaryAssay="RNA", + markerAssay="RNA", + analysisAssays=["RNA"], ) + paused = orchestrator.run(request) - result = orchestrator.run( - AutomatedWorkflowRequest( - sourcePath=str(source), - zarrPath=str(target), - studyContext=( - "A human peripheral blood RNA study for a deterministic " - "acceptance test." - ), - allowAssumptions=True, - primaryAssay="RNA", - markerAssay="RNA", - analysisAssays=["RNA"], - ) + assert paused.status == "needsInput" + assert paused.currentStage == "parameter_tuning" + assert paused.workflowRunId is not None + assert paused.needsInput is not None + assert len(paused.needsInput.questions) == 1 + question = paused.needsInput.questions[0] + assert question.questionId == "parameter_tuning" + pca_diagnostics_before_resume = list(pca_diagnostic_calls) + resume_request = AutomatedWorkflowResumeRequest( + zarrPath=str(target), + workflowRunId=paused.workflowRunId, + answers={"parameter_tuning": state["answer"]}, + ) + editable = DataStore( + str(target), + default_assay="RNA", + min_features_per_cell=-1, + mito_pattern="", + ribo_pattern="", ) + original_counts = editable.cells.fetch_all("RNA_nCounts") + changed_counts = original_counts.copy() + changed_counts[0] += 1 + editable.cells.insert("RNA_nCounts", changed_counts, overwrite=True) + rejected = orchestrator.resume(resume_request) + assert rejected.status == "failed" + assert any("metadata" in note for note in rejected.notes) + assert pca_diagnostic_calls == pca_diagnostics_before_resume + editable.cells.insert("RNA_nCounts", original_counts, overwrite=True) + execute_candidate = tuning_module.RnaTuningRun.execute + + def interrupt_before_full(self: Any, scope: str, *args: Any, **kwargs: Any) -> Any: + if scope == "full": + monkeypatch.setattr( + tuning_module.RnaTuningRun, "execute", execute_candidate + ) + raise KeyboardInterrupt( + "interrupt after the screening answer was committed" + ) + return execute_candidate(self, scope, *args, **kwargs) + + monkeypatch.setattr(tuning_module.RnaTuningRun, "execute", interrupt_before_full) + with pytest.raises(KeyboardInterrupt, match="screening answer"): + orchestrator.resume(resume_request) + assert pca_diagnostic_calls == pca_diagnostics_before_resume + result = orchestrator.resume(resume_request) assert result.status == "completed", result.notes - assert result.currentStage == "biological_interpretation" - assert result.workflowRun is not None - assert result.workflowRun.status == "completed" - assert state["requests"] == 9 - assert [reference.agentName for reference in result.reportReferences] == [ - "data_enrichment", - "experimental_context", - "parameter_tuning", - "biological_interpretation", - ] - assert result.finalAnalysis is not None - assert result.preprocessingPlan is not None - final = result.finalAnalysis - assert final.graphMethod == "native" - assert final.primaryAssay == final.markerAssay == "RNA" - assert final.graph is not None - assert final.clusters is not None - assert final.cellSelection is not None - assert final.embeddingInitialization is not None - assert final.umap is not None - assert final.markers is not None - assert final.cellSelection.kind == "cell_selection" - assert final.clusters.kind == "cluster_labels" - assert final.embeddingInitialization.kind == "embedding_initialization" - assert final.umap.kind == "embedding" - assert final.cellSelection != result.preprocessingPlan.cellSelection + assert len(pca_diagnostic_calls) == len(pca_diagnostics_before_resume) + 1 + assert state["pca_prompts"] >= 3 + assert result.currentStage == "analysis_finalization" + assert result.workflowRunId is not None + report_path = result.report() + assert report_path.is_file() + assert "Scarf analysis summary" in report_path.read_text(encoding="utf-8") + assert state["requests"] >= 8 + assert state["biology"] == 0 persisted = DataStore( str(target), default_assay="RNA", @@ -512,6 +492,39 @@ def test_rna_h5ad_completes_public_automated_workflow(tmp_path: Path) -> None: ribo_pattern="", zarr_mode="r", ) + snapshot = analysis_snapshot(persisted, result.workflowRunId) + stages = {stage["stage"]: stage for stage in snapshot["stages"]} + preprocessing_plan = AutomatedPreprocessingPlan.model_validate( + stages["preprocessing_plan"]["outputs"]["preprocessingPlan"] + ) + assert preprocessing_plan.cellQualityPayload is not None + final = FinalAnalysisHandoff.model_validate(snapshot["finalAnalysis"]) + assert final.primaryAssay == final.markerAssay == "RNA" + assert final.graph is not None and final.clusters is not None + assert final.cellSelection is not None and final.umap is not None + assert final.embeddingInitialization is not None and final.markers is not None + assert len(final.doubletScores) == 1 + assert final.cellSelection != preprocessing_plan.cellSelection + parameter_report = ParameterTuningReport.model_validate( + stages["parameter_tuning"]["report"] + ) + tuning_evaluations = parameter_report.evaluations + pca_evidence = next( + evaluation + for evaluation in tuning_evaluations + if evaluation.metrics.componentVariance + ) + assert "representationDiagnostic" in pca_evidence.artifacts + assert any( + evidence_id.endswith(":pcaLoadings") for evidence_id in pca_evidence.evidenceIds + ) + cluster_evidence = next( + evaluation + for evaluation in tuning_evaluations + if evaluation.metrics.doubletHighScoreConcentration is not None + ) + assert cluster_evidence.metrics.markerCoherence is not None + assert "doubletScore:0" in cluster_evidence.artifacts umap_inputs = persisted.inspect_artifact(artifact_model_to_ref(final.umap)).inputs assert umap_inputs is not None assert umap_inputs["graph"] == artifact_model_to_ref(final.graph).to_dict() @@ -528,24 +541,103 @@ def test_rna_h5ad_completes_public_automated_workflow(tmp_path: Path) -> None: marker_inputs["cell_selection"] == artifact_model_to_ref(final.cellSelection).to_dict() ) - - biology_reference = next( - reference - for reference in result.reportReferences - if reference.agentName == "biological_interpretation" + doublet_inputs = persisted.inspect_artifact( + artifact_model_to_ref(final.doubletScores[0]) + ).inputs + assert doublet_inputs is not None + assert ( + doublet_inputs["connectivity_map"] + == artifact_model_to_ref(final.graph).to_dict() ) - biology = load_agent_report(target, biology_reference) - assert isinstance(biology, BiologicalInterpretationReport) - assert biology.status == "done" - assert biology.clusterInterpretations - assert biology.markerArtifact == final.markers - record = load_agent_record(target, biology_reference) - assert len(record.invocation.parentReports) == 3 - assert record.invocation.tuningBiologyHandoff is not None - assert set(record.invocation.artifacts) == { - "cellSelection", - "clusters", - "markers", - "markerFeatures", + scored_partition = ArtifactRef.from_dict(doublet_inputs["clusters"]) + assert scored_partition.kind == "cluster_labels" + assert persisted.inspect_artifact(scored_partition).complete + selected_evaluation = tuning_evaluations[0] + assert ( + persisted.inspect_artifact( + artifact_model_to_ref(selected_evaluation.artifacts["normalized"]) + ).inputs["cell_selection"] + == artifact_model_to_ref(final.cellSelection).to_dict() + ) + assert len(tuning_evaluations) == 1 + snapshot = analysis_snapshot(persisted, result.workflowRunId) + evidence = next( + stage for stage in snapshot["stages"] if stage["stage"] == "parameter_tuning" + )["outputs"]["tuningEvidence"] + assert 4 < evidence["budget"]["scopes"]["sample0"]["reserved"]["partitions"] <= 24 + compared = { + row["comparisonId"] + for review in snapshot["analysisReviews"] + for row in review["comparisonCoverage"]["comparisons"] } - assert record.invocation.artifacts["cellSelection"] == final.cellSelection + assert { + "hvgCount:2000", + "hvgCount:4000", + "dimensions:10", + "dimensions:30", + "neighborsK:21", + "neighborsK:41", + "hvgRanking", + "featurePolicy", + } <= compared + assert evidence["budget"]["scopes"]["full"]["reserved"]["graphs"] == 1 + assert evidence["budget"]["scopes"]["full"]["reserved"]["partitions"] == 1 + sample_record = load_checkpoint( + persisted, + _ensure_orchestration_store(persisted), + result.workflowRunId, + "parameter_tuning/sample0/evaluation0/complete", + inputs=None, + ) + assert sample_record is not None + sample_evaluation = ParameterCandidateEvaluation.model_validate( + sample_record["evaluation"] + ) + assert ( + sample_evaluation.parameters.leidenResolution + == selected_evaluation.parameters.leidenResolution + == 0.5 + ) + assert sample_evaluation.cellSelection != final.cellSelection + + def selected_rows(reference: Any) -> np.ndarray: + return read_stored_selection_indices( + persisted.zw, + artifact_model_to_ref(reference), + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ) + + sample_rows = selected_rows(sample_evaluation.cellSelection) + full_rows = selected_rows(final.cellSelection) + sample_labels = np.asarray( + persisted.load_artifact( + artifact_model_to_ref(sample_evaluation.artifacts["clusters"]) + )["values"][:] + ) + full_labels = np.asarray( + persisted.load_artifact(artifact_model_to_ref(final.clusters))["values"][:] + ) + from sklearn.metrics import adjusted_rand_score + + assert adjusted_rand_score( + sample_labels, full_labels[np.searchsorted(full_rows, sample_rows)] + ) == pytest.approx(1.0) + markers = result.get_markers(min_score=0.0, min_frac_exp=0.0) + assert {"CD3D", "MS4A1"}.issubset(set(markers.feature_name)) + plotted = result.plot_embedding(show=False) + assert plotted.figure is not None + plotted.close() + requests_before = state["requests"] + diagnostics_before = list(pca_diagnostic_calls) + completed = orchestrator.run(request) + assert completed.status == "completed" + assert completed.workflowRunId == result.workflowRunId + assert analysis_snapshot(persisted, result.workflowRunId)[ + "finalAnalysis" + ] == final.model_dump(mode="json") + assert state["requests"] == requests_before + assert pca_diagnostic_calls == diagnostics_before + assert "pipeline" not in persisted.zw diff --git a/tests/test_agent_orchestrator_journal_edges.py b/tests/test_agent_orchestrator_journal_edges.py index 34e90889..4fb5e979 100644 --- a/tests/test_agent_orchestrator_journal_edges.py +++ b/tests/test_agent_orchestrator_journal_edges.py @@ -1,704 +1,386 @@ -"""Focused edge coverage for orchestration journal contracts.""" +"""Immutable evidence, integrity, and exact stage lineage contracts.""" +import json from types import SimpleNamespace -from typing import Any import pytest -import zarr -from pydantic import ValidationError -from zarr.storage import MemoryStore +from zarr.core.buffer import default_buffer_prototype +from zarr.core.sync import sync -import scarf.agent.orchestrator.journal as journal_module -from scarf.agent.orchestrator import ( - AutomatedWorkflowRequest, - AutomatedWorkflowResult, - AutomatedWorkflowResumeRequest, +from scarf.agent.orchestrator import journal +from scarf.agent.orchestrator.models import ( + OrchestrationResumeRecord, + StageEvidenceReference, WorkflowNeedsInput, WorkflowQuestion, WorkflowStageAttempt, - WorkflowStageLink, ) -from scarf.agent.orchestrator.models import ( - OrchestrationRequestRecord, - OrchestrationResumeRecord, -) -from scarf.agent.persistence import ( - AgentInvocation, - AgentReportReference, - AgentWorkflowRun, -) - - -def _started( - *, - stage: str = "ingest", - parents: list[WorkflowStageLink] | None = None, - inputs: dict[str, Any] | None = None, -) -> WorkflowStageAttempt: - attempt = WorkflowStageAttempt( - workflowRunId="workflow-1", - stage=stage, - attemptId="attempt-1", - status="started", - startedAtNs=1, - requestSha256="a" * 64, - configSha256="b" * 64, - parentAttempts=parents or [], - inputs=inputs or {}, - ) - return attempt.model_copy( - update={"contentSha256": journal_module._stage_checksum(attempt)} - ) - - -def _request_record() -> OrchestrationRequestRecord: - return OrchestrationRequestRecord( - workflowRunId="workflow-1", - requestSha256="a" * 64, - configSha256="b" * 64, - ) - - -def _with_checksum(record: Any) -> Any: - return record.model_copy( - update={"contentSha256": journal_module._record_checksum(record)} - ) - - -def _terminal_workflow() -> AgentWorkflowRun: - return AgentWorkflowRun( - workflowRunId="workflow-1", - createdAtNs=1, - finalizedAtNs=2, - status="completed", - finalizationMessage="completed", - analysisStore="analysis.zarr", - datasetFingerprints={"RNA": "fingerprint"}, - ) - - -def _terminal_result(workflow: AgentWorkflowRun) -> AutomatedWorkflowResult: - return _with_checksum( - AutomatedWorkflowResult( - status="completed", - currentStage="biological_interpretation", - zarrPath="analysis.zarr", - workflowRun=workflow, - reportReferences=list(workflow.reports), - ) - ) - - -def test_orchestration_model_validation_edges() -> None: - for field, value, message in ( - ("attemptId", "UPPER CASE", "lowercase run identifier"), - ("contentSha256", "not-a-checksum", "lowercase SHA-256"), - ): - with pytest.raises(ValidationError, match=message): - WorkflowStageLink(**{field: value}) - - invalid_attempts = ( - ({"status": "started", "completedAtNs": 1}, "cannot have completedAtNs"), - ( - {"status": "done", "startedAtNs": 2, "completedAtNs": 1}, - "must not precede", - ), - ({"status": "needsInput"}, "require questions"), - ({"status": "failed"}, "require an error"), - ) - for updates, message in invalid_attempts: - with pytest.raises(ValidationError, match=message): - WorkflowStageAttempt(**updates) - - invalid_requests = ( - ({"sourcePath": " ", "studyContext": "study"}, "sourcePath"), - ({"sourcePath": "data", "studyContext": " "}, "studyContext"), - ( - { - "sourcePath": "data", - "studyContext": "study", - "analysisAssays": ["RNA", "RNA"], - }, - "analysisAssays", - ), - ( - { - "sourcePath": "data", - "studyContext": "study", - "pairedAssays": ["RNA", "RNA"], - }, - "pairedAssays must be unique", - ), - ( - { - "sourcePath": "data", - "studyContext": "study", - "pairedAssays": ["RNA"], - }, - "at least two", - ), - ) - for values, message in invalid_requests: - with pytest.raises(ValidationError, match=message): - AutomatedWorkflowRequest(**values) - - with pytest.raises(ValidationError, match="zarrPath"): - AutomatedWorkflowResumeRequest(zarrPath=" ", workflowRunId="workflow-1") - with pytest.raises(ValidationError, match="workflowRunId"): - AutomatedWorkflowResumeRequest(zarrPath="data.zarr", workflowRunId="BAD") - - assert OrchestrationRequestRecord.get_blank().workflowRunId == "" - assert OrchestrationRequestRecord.get_example().workflowRunId == "workflow-1" - assert OrchestrationResumeRecord.get_blank().resumeId == "" - assert OrchestrationResumeRecord.get_example().resumeId == "resume-1" - - -def test_journal_storage_guards(monkeypatch: pytest.MonkeyPatch) -> None: - root = zarr.open_group(store=MemoryStore(), mode="w") - journal_module._write_key_once(root, "record.json", b"first") - with pytest.raises(FileExistsError, match="exists"): - journal_module._write_key_once(root, "record.json", b"second") - - with monkeypatch.context() as patch: - observed = iter((None, b"different")) - patch.setattr( - journal_module.record_io, - "read_key", - lambda *_args: next(observed), - ) - with pytest.raises(FileExistsError, match="raced"): - journal_module._write_key_once(root, "raced.json", b"payload") - - unsupported = SimpleNamespace(store=SimpleNamespace(supports_listing=False)) - with pytest.raises(NotImplementedError, match="requires listing"): - journal_module._list_keys(unsupported, "prefix") - - with monkeypatch.context() as patch: - patch.setattr(journal_module.record_io, "read_key", lambda *_args: b"{") - with pytest.raises(ValueError, match="Malformed orchestration record"): - journal_module._read_model(root, "bad.json", WorkflowQuestion) - - -def test_orchestration_namespace_validation(monkeypatch: pytest.MonkeyPatch) -> None: - root = zarr.open_group(store=MemoryStore(), mode="w") - with pytest.raises(RuntimeError, match="Create the agent workflow"): - journal_module._ensure_orchestration_store(SimpleNamespace(zw=root)) - - root = zarr.open_group(store=MemoryStore(), mode="w") - root.create_array("agents", shape=(1,), dtype="i1") - with pytest.raises(ValueError, match="agents namespace"): - journal_module._ensure_orchestration_store(SimpleNamespace(zw=root)) - - root = zarr.open_group(store=MemoryStore(), mode="w") - root.create_group("agents") - with monkeypatch.context() as patch: - patch.setattr(journal_module, "_list_keys", lambda *_args: ["occupied"]) - with pytest.raises(ValueError, match="non-Zarr object"): - journal_module._ensure_orchestration_store(SimpleNamespace(zw=root)) - - root = zarr.open_group(store=MemoryStore(), mode="w") - agents = root.create_group("agents") - agents.create_array("orchestrations", shape=(1,), dtype="i1") - with pytest.raises(ValueError, match="orchestrations namespace"): - journal_module._ensure_orchestration_store(SimpleNamespace(zw=root)) +from scarf.agent.types import AgentDataModel +from tests.agent_journal_store import memory_journal - root = zarr.open_group(store=MemoryStore(), mode="w") - agents = root.create_group("agents") - agents.create_group( - "orchestrations", - attributes={"format": "foreign", "format_version": 99}, - ) - with pytest.raises(ValueError, match="Unrecognized orchestration"): - journal_module._ensure_orchestration_store(SimpleNamespace(zw=root)) +class Evidence(AgentDataModel): + rationale: str + nCells: int -def test_stage_record_identity_and_checksum_guards( - monkeypatch: pytest.MonkeyPatch, -) -> None: - outcome_key = "root/workflow-1/stages/ingest/attempt-1/outcome.json" - start_key = "root/workflow-1/stages/ingest/attempt-1/started.json" - done = journal_module._complete_attempt(_started(), status="done") - with monkeypatch.context() as patch: - patch.setattr(journal_module, "_list_keys", lambda *_args: [outcome_key]) - patch.setattr( - journal_module, - "_read_model", - lambda *_args: done.model_copy(update={"attemptId": "other"}), - ) - with pytest.raises(ValueError, match="outcome identity"): - journal_module._stage_outcomes(object(), "root", "workflow-1", "ingest") - patch.setattr( - journal_module, - "_read_model", - lambda *_args: done.model_copy(update={"contentSha256": "0" * 64}), - ) - with pytest.raises(ValueError, match="outcome checksum"): - journal_module._stage_outcomes(object(), "root", "workflow-1", "ingest") - - started = _started() - with monkeypatch.context() as patch: - patch.setattr(journal_module, "_list_keys", lambda *_args: [start_key]) - patch.setattr( - journal_module, - "_read_model", - lambda *_args: started.model_copy(update={"status": "done"}), - ) - with pytest.raises(ValueError, match="start identity"): - journal_module._stage_starts(object(), "root", "workflow-1", "ingest") - patch.setattr( - journal_module, - "_read_model", - lambda *_args: started.model_copy(update={"contentSha256": "0" * 64}), - ) - with pytest.raises(ValueError, match="start checksum"): - journal_module._stage_starts(object(), "root", "workflow-1", "ingest") - - -def test_resume_answer_helper_edges() -> None: - assert not journal_module._has_resume_answer(None) - assert not journal_module._has_resume_answer(" ") - assert journal_module._has_resume_answer("answer") - assert not journal_module._has_resume_answer({}) - assert journal_module._has_resume_answer({"answer": 1}) - assert not journal_module._has_resume_answer([]) - assert journal_module._has_resume_answer([1]) - assert journal_module._has_resume_answer(0) - - assert journal_module._unsafe_context_resolution("skipHarmony") == "skip" +def test_checkpoint_is_immutable_idempotent_and_input_bound() -> None: + store, prefix, record = memory_journal("analysis") + inputs = {"selection": "cells-1", "metadata": {"age": "continuous"}} + value = {"decision": "retain", "evidence": [1, 2]} + key = "parameter_tuning/sample0/eval0" assert ( - journal_module._unsafe_context_resolution( - {"batchCorrection": {"action": "skip"}} - ) - == "skip" - ) - assert journal_module._unsafe_context_resolution({"selection": "other"}) is None - - assert journal_module._resume_answer_errors(_started(), {}) == [ - "The latest paused stage does not contain persisted questions" - ] - choice = journal_module._complete_attempt( - _started(), - status="needsInput", - needs_input=WorkflowNeedsInput( - questions=[ - WorkflowQuestion( - questionId="finalGraphOptionId", - options=["graph-a", "graph-b"], - ) - ] - ), + journal.save_checkpoint(store, prefix, record.workflowRunId, key, inputs, value) + == value ) + before = journal._list_keys(store.zw, prefix) assert ( - "must be one of" - in journal_module._resume_answer_errors( - choice, {"finalGraphOptionId": "graph-c"} - )[0] - ) - generic = journal_module._complete_attempt( - _started(), - status="needsInput", - needs_input=WorkflowNeedsInput( - questions=[WorkflowQuestion(questionId="freeText")] - ), + journal.save_checkpoint(store, prefix, record.workflowRunId, key, inputs, value) + == value ) + assert journal._list_keys(store.zw, prefix) == before assert ( - "must be non-empty" - in journal_module._resume_answer_errors(generic, {"freeText": " "})[0] + journal.load_checkpoint(store, prefix, record.workflowRunId, key, inputs) + == value ) - - -def test_persisted_resume_record_validation_edges( - monkeypatch: pytest.MonkeyPatch, -) -> None: - store = SimpleNamespace(zw=object()) - - def validate(record: OrchestrationResumeRecord) -> OrchestrationResumeRecord: - monkeypatch.setattr(journal_module, "_read_model", lambda *_args: record) - return journal_module._validated_resume_record( - store, "root", "workflow-1", "resume-1" + checkpoint = journal.read_checkpoint(store, prefix, record.workflowRunId, key) + assert checkpoint["inputs"] == inputs + assert checkpoint["outputs"] == value + with pytest.raises(ValueError, match="different scientific inputs"): + journal.load_checkpoint( + store, prefix, record.workflowRunId, key, {**inputs, "selection": "cells-2"} ) - - mismatched = _with_checksum( - OrchestrationResumeRecord(workflowRunId="other", resumeId="resume-1") - ) - with pytest.raises(ValueError, match="identity"): - validate(mismatched) - - with pytest.raises(ValueError, match="checksum"): - validate( - OrchestrationResumeRecord( - workflowRunId="workflow-1", - resumeId="resume-1", - contentSha256="0" * 64, - ) + with pytest.raises(ValueError, match="different outcome"): + journal.save_checkpoint( + store, prefix, record.workflowRunId, key, inputs, {"decision": "exclude"} ) + assert all(path.startswith("analysis/agents/orchestrations/") for path in before) - unanswered = _with_checksum( - OrchestrationResumeRecord( - workflowRunId="workflow-1", - resumeId="resume-1", - answers={"unexpected": "answer"}, - ) - ) - with pytest.raises(ValueError, match="without an answered attempt"): - validate(unanswered) - empty = _with_checksum( - OrchestrationResumeRecord(workflowRunId="workflow-1", resumeId="resume-1") +def test_checkpoint_corruption_is_not_silently_repaired() -> None: + store, prefix, record = memory_journal() + key = "preprocessing/decision" + journal.save_checkpoint( + store, prefix, record.workflowRunId, key, {"input": 1}, {"chosen": 2} ) - assert validate(empty) == empty - - paused = journal_module._complete_attempt( - _started(), - status="needsInput", - needs_input=WorkflowNeedsInput( - questions=[WorkflowQuestion(questionId="freeText")] - ), + path = journal._checkpoint_key(prefix, record.workflowRunId, key) + raw = journal.record_io.read_key(store.zw, path) + data = json.loads(raw) + data["outputs"]["chosen"] = 3 + altered = journal.record_io.display_json_bytes(data) + sync( + store.zw.store.set(path, default_buffer_prototype().buffer.from_bytes(altered)) ) - link = journal_module._parent_link(paused) - answered = _with_checksum( - OrchestrationResumeRecord( - workflowRunId="workflow-1", - resumeId="resume-1", - answeredAttempt=link, - questionIds=["freeText"], - answers={"freeText": "answer"}, - ) + with pytest.raises(ValueError, match="checksum"): + journal.load_checkpoint(store, prefix, record.workflowRunId, key, None) + assert journal.record_io.read_key(store.zw, path) == altered + + +@pytest.mark.parametrize("key", ["../escape", "x//y", "/root", "x/../y", "x\\y"]) +def test_checkpoint_path_cannot_escape_owner(key: str) -> None: + with pytest.raises(ValueError, match="path components"): + journal._checkpoint_key("agents/orchestrations", "workflow-1", key) + + +def test_committed_evidence_survives_interruption_before_stage_outcome() -> None: + store, prefix, request = memory_journal() + first = journal._start_attempt( + store.zw, + prefix, + request.workflowRunId, + "experimental_context", + request, + [], + inputs={"columns": ["condition"]}, + ) + report = Evidence(rationale="Condition crosses donors", nCells=100) + _, reference = journal._save_stage_report( + store, first, report, expected_type=Evidence + ) + restarted = journal._start_attempt( + store.zw, + prefix, + request.workflowRunId, + "experimental_context", + request, + [], + inputs=first.inputs, + ) + recovered, recovered_ref = journal._recover_persisted_stage_report( + store, restarted, expected_type=Evidence + ) + assert recovered == report + assert recovered_ref == reference + changed = journal._start_attempt( + store.zw, + prefix, + request.workflowRunId, + "experimental_context", + request, + [], + inputs={"columns": ["condition", "age"]}, ) - - monkeypatch.setattr(journal_module, "_stage_outcomes", lambda *_args: []) - with pytest.raises(ValueError, match="does not cite one paused"): - validate(answered) - - monkeypatch.setattr(journal_module, "_stage_outcomes", lambda *_args: [paused]) - stale = _with_checksum(answered.model_copy(update={"questionIds": ["stale"]})) - with pytest.raises(ValueError, match="question IDs are stale"): - validate(stale) - - invalid = _with_checksum(answered.model_copy(update={"answers": {"freeText": ""}})) - with pytest.raises(ValueError, match="invalid answers"): - validate(invalid) - - -def test_stage_outcome_resolution_edges(monkeypatch: pytest.MonkeyPatch) -> None: - request = _request_record() - store = SimpleNamespace( - zw=object(), - cells=SimpleNamespace(columns=[]), - load_artifact=lambda *_args: object(), + assert ( + journal._recover_persisted_stage_report(store, changed, expected_type=Evidence) + is None ) - done = journal_module._complete_attempt(_started(), status="done") - - with pytest.raises(ValueError, match="checksum is stale"): - journal_module._stage_outcome_resolves( - store, - "root", - "workflow-1", - request, - done.model_copy(update={"requestSha256": "c" * 64}), - ) - with pytest.raises(ValueError, match="cannot retain started"): - journal_module._stage_outcome_resolves( - store, "root", "workflow-1", request, _started() - ) - - monkeypatch.setattr( - journal_module, - "_read_model", - lambda *_args: _started().model_copy(update={"inputs": {"changed": True}}), + assert journal.read_stage_evidence(store, reference) == report.model_dump( + mode="json" ) - with pytest.raises(ValueError, match="do not match"): - journal_module._stage_outcome_resolves( - store, "root", "workflow-1", request, done + with pytest.raises(ValueError, match="owned"): + journal.read_stage_evidence( + store, + StageEvidenceReference.model_validate( + {**reference.model_dump(), "stage": "preprocessing"} + ), ) - - parent_outcome = journal_module._complete_attempt(_started(), status="done") - parent = journal_module._parent_link(parent_outcome) - child_started = _started(stage="data_enrichment", parents=[parent]) - child = journal_module._complete_attempt(child_started, status="done") - monkeypatch.setattr(journal_module, "_read_model", lambda *_args: child_started) - monkeypatch.setattr(journal_module, "_stage_outcomes", lambda *_args: []) - assert not journal_module._stage_outcome_resolves( - store, "root", "workflow-1", request, child - ) - - reference = AgentReportReference.get_example().model_copy( - update={"workflowRunId": "other-workflow"} - ) - report_outcome = journal_module._complete_attempt( - _started(), status="done", report_references=[reference] - ) - monkeypatch.setattr(journal_module, "_read_model", lambda *_args: _started()) - assert not journal_module._stage_outcome_resolves( - store, "root", "workflow-1", request, report_outcome - ) - - metadata_outcome = journal_module._complete_attempt( - _started(), status="done", outputs={"metadataColumns": ["missing"]} - ) - assert not journal_module._stage_outcome_resolves( - store, "root", "workflow-1", request, metadata_outcome + assert not any( + "/runs/" in key or "snapshot" in key or "verification" in key + for key in journal._list_keys(store.zw, "agents") ) - monkeypatch.setattr( - journal_module, - "_stage_outcomes", - lambda *_args: [done.model_copy(update={"configSha256": "c" * 64})], - ) - with pytest.raises(ValueError, match="checksum is stale"): - journal_module._validated_done_outcome( - store, "root", "workflow-1", "ingest", request, [] - ) - - -def test_stage_invocation_and_report_loading_edges( - monkeypatch: pytest.MonkeyPatch, -) -> None: - started = _started(stage="parameter_tuning") - with pytest.raises(ValueError, match="conflicting orchestration identity"): - journal_module._stage_invocation( - started, - AgentInvocation(inputs={"orchestrationExecutionId": "other"}), - ) - - with pytest.raises(ValueError, match="exactly one report reference"): - journal_module.load_stage_report(object(), started, WorkflowQuestion) - - execution_id = journal_module._stage_execution_id(started) - matching = AgentReportReference.get_example().model_copy( - update={"agentRunId": execution_id} - ) - other = AgentReportReference.get_example().model_copy( - update={"agentRunId": "other-report"} - ) - tuning = started.model_copy(update={"reportReferences": [matching, other]}) - monkeypatch.setattr( - journal_module, - "load_agent_report", - lambda *_args: WorkflowQuestion(questionId="answer"), - ) - assert isinstance( - journal_module.load_stage_report(object(), tuning, WorkflowQuestion), - WorkflowQuestion, - ) - wrong = started.model_copy(update={"reportReferences": [matching]}) - monkeypatch.setattr( - journal_module, - "load_agent_report", - lambda *_args: WorkflowNeedsInput(), +def test_stage_reuse_requires_exact_parent_and_resolving_evidence() -> None: + store, prefix, request = memory_journal() + start = journal._start_attempt( + store.zw, prefix, request.workflowRunId, "ingest", request, [] ) - with pytest.raises(TypeError, match="report is not WorkflowQuestion"): - journal_module.load_stage_report(object(), wrong, WorkflowQuestion) - - -def test_persisted_stage_report_recovery_edges( - monkeypatch: pytest.MonkeyPatch, -) -> None: - started = _started(stage="data_enrichment") - execution_id = journal_module._stage_execution_id(started) - reference = AgentReportReference.get_example().model_copy( - update={"agentRunId": execution_id} + parent = journal._complete_attempt(start, status="done") + journal._save_outcome(store.zw, prefix, parent) + link = journal._parent_link(parent) + child = journal._start_attempt( + store.zw, prefix, request.workflowRunId, "data_enrichment", request, [link] ) - store = SimpleNamespace(load_artifact=lambda *_args: object()) - - monkeypatch.setattr( - journal_module, - "list_agent_reports", - lambda *_args, **_kwargs: [reference, reference], + _, ref = journal._save_stage_report( + store, + child, + Evidence(rationale="RNA identified", nCells=100), + expected_type=Evidence, ) - with pytest.raises(ValueError, match="multiple persisted reports"): - journal_module._recover_persisted_stage_report( - store, - started, - agent_name="data_enrichment", - expected_type=WorkflowQuestion, + outcome = journal._complete_attempt(child, status="done", report_references=[ref]) + journal._save_outcome(store.zw, prefix, outcome) + assert ( + journal._validated_done_outcome( + store, prefix, request.workflowRunId, "data_enrichment", request, [link] ) - - monkeypatch.setattr( - journal_module, - "list_agent_reports", - lambda *_args, **_kwargs: [reference], + == outcome ) - monkeypatch.setattr( - journal_module, - "load_agent_record", - lambda *_args: SimpleNamespace( - invocation=AgentInvocation( - agentName="data_enrichment", - inputs={"orchestrationExecutionId": "stale"}, - ) - ), - ) - with pytest.raises(ValueError, match="stale execution identity"): - journal_module._recover_persisted_stage_report( - store, - started, - agent_name="data_enrichment", - expected_type=WorkflowQuestion, + assert ( + journal._validated_done_outcome( + store, prefix, request.workflowRunId, "data_enrichment", request, [] ) - - monkeypatch.setattr( - journal_module, - "load_agent_record", - lambda *_args: SimpleNamespace( - invocation=AgentInvocation( - agentName="data_enrichment", - inputs={"orchestrationExecutionId": execution_id}, - ) - ), - ) - monkeypatch.setattr( - journal_module, - "load_agent_report", - lambda *_args: WorkflowNeedsInput(), + is None ) - with pytest.raises(TypeError, match="is not WorkflowQuestion"): - journal_module._recover_persisted_stage_report( - store, - started, - agent_name="data_enrichment", - expected_type=WorkflowQuestion, + sync( + store.zw.store.delete( + journal._checkpoint_key(prefix, request.workflowRunId, ref.key) ) - - report = WorkflowQuestion(questionId="question") - invocation = AgentInvocation(agentName="data_enrichment") - monkeypatch.setattr( - journal_module, - "save_agent_report", - lambda *_args, **_kwargs: (_ for _ in ()).throw(FileExistsError), ) - monkeypatch.setattr( - journal_module, - "_recover_persisted_stage_report", - lambda *_args, **_kwargs: None, - ) - with pytest.raises(FileExistsError): - journal_module._save_stage_report( - store, - started, - report, - invocation=invocation, - expected_type=WorkflowQuestion, + assert ( + journal._validated_done_outcome( + store, prefix, request.workflowRunId, "data_enrichment", request, [link] ) - - monkeypatch.setattr( - journal_module, - "_recover_persisted_stage_report", - lambda *_args, **_kwargs: (report, reference), + is None ) - assert journal_module._save_stage_report( - store, - started, - report, - invocation=invocation, - expected_type=WorkflowQuestion, - ) == (report, reference) + with pytest.raises(ValueError, match="unresolved"): + journal.analysis_snapshot(store, request.workflowRunId) -def test_terminal_result_validation_and_persistence_edges( - monkeypatch: pytest.MonkeyPatch, -) -> None: - workflow = _terminal_workflow() - result = _terminal_result(workflow) - store = SimpleNamespace(zw=object()) - - def load(payload: bytes) -> AutomatedWorkflowResult | None: - monkeypatch.setattr( - journal_module.record_io, "read_key", lambda *_args: payload - ) - return journal_module._load_terminal_result(store, "root", workflow) - - with pytest.raises(ValueError, match="Malformed automated workflow result"): - load(b"{") - with pytest.raises(ValueError, match="checksum is invalid"): - load( - journal_module.record_io.display_json_bytes( - result.model_dump(mode="json") | {"contentSha256": "0" * 64} - ) - ) - - missing_workflow = _with_checksum( - AutomatedWorkflowResult( - status="completed", currentStage="biological_interpretation" - ) +def test_resume_answers_are_committed_in_stage_inputs_without_another_ledger() -> None: + store, prefix, request = memory_journal() + start = journal._start_attempt( + store.zw, prefix, request.workflowRunId, "preprocessing_plan", request, [] ) - with pytest.raises(ValueError, match="missing its workflow identity"): - load( - journal_module.record_io.display_json_bytes( - missing_workflow.model_dump(mode="json") - ) - ) - - stale = _with_checksum( - result.model_copy( - update={ - "workflowRun": workflow.model_copy( - update={"analysisStore": "other.zarr"} + pause = journal._complete_attempt( + start, + status="needsInput", + needs_input=WorkflowNeedsInput( + questions=[ + WorkflowQuestion( + questionId="decision:cellQuality", + decisionId="cellQuality", + options=["keep"], + question="Which supported QC policy?", ) - } - ) + ] + ), ) - with pytest.raises(ValueError, match="stale workflow metadata"): - load(journal_module.record_io.display_json_bytes(stale.model_dump(mode="json"))) - - wrong_status = _with_checksum(result.model_copy(update={"status": "failed"})) - with pytest.raises(ValueError, match="stale terminal status"): - load( - journal_module.record_io.display_json_bytes( - wrong_status.model_dump(mode="json") - ) + answers = { + "decision:cellQuality": { + "decisionId": "cellQuality", + "optionId": "keep", + "rationale": "Preserves replicated populations", + } + } + assert journal._resume_answer_errors(pause, answers) == [] + assert journal._resume_answer_errors(pause, {"unrelated": 1}) + bad = { + "decision:cellQuality": { + **answers["decision:cellQuality"], + "optionId": "invented", + } + } + assert journal._resume_answer_errors(pause, bad) + resume = OrchestrationResumeRecord( + workflowRunId=request.workflowRunId, + answeredAttempt=journal._parent_link(pause), + answers=answers, + questionIds=list(answers), + ) + answered = journal._start_attempt( + store.zw, + prefix, + request.workflowRunId, + "preprocessing_plan", + request, + [], + resume_record=resume, + ) + replay = journal._start_attempt( + store.zw, + prefix, + request.workflowRunId, + "preprocessing_plan", + request, + [], + resume_record=resume, + ) + assert answered.inputs["resumeAnswers"] == answers + assert journal._stage_execution_id(answered) == journal._stage_execution_id(replay) + assert not any("/resumes/" in key for key in journal._list_keys(store.zw, prefix)) + + +def test_stage_checksums_and_original_start_are_required() -> None: + store, prefix, request = memory_journal() + start = journal._start_attempt( + store.zw, prefix, request.workflowRunId, "ingest", request, [] + ) + outcome = journal._complete_attempt(start, status="done") + journal._save_outcome(store.zw, prefix, outcome) + altered = outcome.model_copy(update={"inputs": {"different": True}}) + with pytest.raises(ValueError, match="started and outcome"): + journal._stage_outcome_resolves( + store, prefix, request.workflowRunId, request, altered ) + with pytest.raises(FileExistsError, match="exists"): + journal._save_outcome(store.zw, prefix, outcome) + no_listing = SimpleNamespace(store=SimpleNamespace(supports_listing=False)) + with pytest.raises(NotImplementedError, match="listing"): + journal._list_keys(no_listing, prefix) + + +@pytest.mark.parametrize( + "values", + [ + {"status": "started", "completedAtNs": 1}, + {"status": "done", "startedAtNs": 2, "completedAtNs": 1}, + {"status": "needsInput"}, + {"status": "failed"}, + ], +) +def test_invalid_stage_lifecycle_is_rejected(values) -> None: + with pytest.raises(ValueError): + WorkflowStageAttempt(**values) - wrong_reports = _with_checksum( - result.model_copy( - update={"reportReferences": [AgentReportReference.get_example()]} - ) - ) - with pytest.raises(ValueError, match="stale report references"): - load( - journal_module.record_io.display_json_bytes( - wrong_reports.model_dump(mode="json") - ) - ) - with monkeypatch.context() as patch: - patch.setattr(journal_module, "_load_terminal_result", lambda *_args: result) - assert ( - journal_module._persist_terminal_result(store, "root", workflow, result) - == result +@pytest.mark.parametrize( + "mismatch", [None, "missingMarker", "differentCohort", "unaccepted"] +) +def test_final_checkpoint_requires_exact_artifacts_cohort_and_acceptance( + monkeypatch, mismatch +) -> None: + from scarf.agent.orchestrator.models import FinalAnalysisHandoff, _STAGE_ORDER + from scarf.agent.types import ArtifactReferenceModel + + class TuningProof(AgentDataModel): + recommendedCandidateId: str + finalClusterArtifact: ArtifactReferenceModel + + store, prefix, request = memory_journal() + artifacts = { + name: ArtifactReferenceModel( + scope="datastore" if name == "cellSelection" else "assay", + assay=None if name == "cellSelection" else "RNA", + kind=kind, + artifactId=f"{index:064x}", + ) + for index, (name, kind) in enumerate( + { + "cellSelection": "cell_selection", + "graph": "connectivity_map", + "clusters": "cluster_labels", + "umap": "embedding", + "embeddingInitialization": "embedding_initialization", + "markerFeatures": "feature_selection", + "markers": "marker_table", + }.items(), + 1, + ) + } + final = FinalAnalysisHandoff( + workflowRunId=request.workflowRunId, + primaryAssay="RNA", + markerAssay="RNA", + **artifacts, + ) + inputs = { + "preprocessedAssays": [ + {"cellSelection": artifacts["cellSelection"].model_dump(mode="json")} + ] + } + if mismatch == "missingMarker": + final = final.model_copy(update={"markers": None}) + if mismatch == "differentCohort": + inputs["preprocessedAssays"][0]["cellSelection"]["artifactId"] = "e" * 64 + parents = [] + for stage in _STAGE_ORDER: + start = journal._start_attempt( + store.zw, + prefix, + request.workflowRunId, + stage, + request, + parents, + inputs=inputs if stage == "analysis_finalization" else {}, ) - - with monkeypatch.context() as patch: - patch.setattr(journal_module, "_load_terminal_result", lambda *_args: None) - with pytest.raises(ValueError, match="exact content checksum"): - journal_module._persist_terminal_result( + references = [] + if stage == "parameter_tuning": + _, ref = journal._save_stage_report( store, - "root", - workflow, - result.model_copy(update={"contentSha256": ""}), + start, + TuningProof( + recommendedCandidateId="selected", + finalClusterArtifact=artifacts["clusters"], + ), + expected_type=TuningProof, ) - - with monkeypatch.context() as patch: - persisted = iter((None, result)) - patch.setattr( - journal_module, - "_load_terminal_result", - lambda *_args: next(persisted), - ) - patch.setattr( - journal_module, - "_write_model_once", - lambda *_args: (_ for _ in ()).throw(FileExistsError), - ) - assert ( - journal_module._persist_terminal_result(store, "root", workflow, result) - == result - ) - - with monkeypatch.context() as patch: - patch.setattr(journal_module, "_load_terminal_result", lambda *_args: None) - patch.setattr(journal_module, "_write_model_once", lambda *_args: None) - with pytest.raises(RuntimeError, match="was not persisted"): - journal_module._persist_terminal_result(store, "root", workflow, result) + references.append(ref) + outcome = journal._complete_attempt( + start, + status="done", + report_references=references, + artifacts=artifacts if stage == "analysis_finalization" else {}, + outputs={"finalAnalysis": final.model_dump(mode="json")} + if stage == "analysis_finalization" + else {}, + ) + journal._save_outcome(store.zw, prefix, outcome) + parents = [journal._parent_link(outcome)] + monkeypatch.setattr( + journal, + "_analysis_review_views", + lambda *args: [ + { + "scope": "full", + "action": "defer" if mismatch == "unaccepted" else "accept", + "selectedCandidateId": "selected", + } + ], + ) + if mismatch: + with pytest.raises(ValueError, match="Final analysis"): + journal.analysis_snapshot(store, request.workflowRunId) + else: + snapshot = journal.analysis_snapshot(store, request.workflowRunId) + assert snapshot["status"] == "completed" + assert snapshot["finalAnalysis"] == final.model_dump(mode="json") diff --git a/tests/test_agent_orchestrator_lifecycle.py b/tests/test_agent_orchestrator_lifecycle.py index eaa66862..00f602bf 100644 --- a/tests/test_agent_orchestrator_lifecycle.py +++ b/tests/test_agent_orchestrator_lifecycle.py @@ -1,981 +1,389 @@ -"""Lifecycle, persistence, resume, and recovery orchestrator contracts.""" +"""Journal-owned resume, input identity, and visible failure behavior.""" -import hashlib -import json -import shutil -import uuid -from collections.abc import Mapping from pathlib import Path -from typing import Any +from types import SimpleNamespace +import numpy as np import pytest import zarr -import scarf.agent.orchestrator.context as context_module -import scarf.agent.orchestrator.journal as journal_module -import scarf.agent.orchestrator.main as main_module -from scarf.agent.config import AgentRunConfig -from scarf.agent.data_enrichment import DataEnrichmentReport -from scarf.agent.ingest import IngestResult from scarf.agent.orchestrator import ( AgentOrchestrator, AutomatedWorkflowConfig, AutomatedWorkflowRequest, - AutomatedWorkflowResult, AutomatedWorkflowResumeRequest, - WorkflowNeedsInput, - WorkflowQuestion, - WorkflowStageAttempt, - WorkflowStageLink, + journal, ) from scarf.agent.orchestrator.models import ( - OrchestrationRequestRecord, + AutomatedWorkflowResult, OrchestrationResumeRecord, + WorkflowIdentity, + WorkflowNeedsInput, + WorkflowQuestion, ) -from scarf.agent.persistence import ( - AgentInvocation, - create_agent_workflow, - finalize_agent_workflow, - list_agent_reports, - load_agent_report, - load_agent_workflow, - save_agent_report, -) -from scarf.agent.types import AgentRunInfo, ArtifactReferenceModel -from scarf.datastore.datastore import DataStore +from tests.agent_journal_store import memory_journal from tests.agent_orchestrator_store import create_store -_PLAN_CHECKSUM = "a" * 64 - - -def _record_checksum(payload: dict[str, Any]) -> str: - content = dict(payload) - content.pop("contentSha256", None) - encoded = json.dumps( - content, - sort_keys=True, - separators=(",", ":"), - ensure_ascii=False, - allow_nan=False, - ).encode("utf-8") - return hashlib.sha256(encoded).hexdigest() - - -def _assert_record_checksum(path: Path) -> dict[str, Any]: - payload = json.loads(path.read_text()) - assert payload["contentSha256"] == _record_checksum(payload) - return payload - - -def _mock_ingest(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(main_module, "detect_format", lambda _path: "zarr") - - def fake_ingest( - path: str, - *, - zarrPath: str | None, - model: Any, - directions: Mapping[str, Any], - ) -> IngestResult: - del model, directions - resolved = str(Path(zarrPath or path).resolve()) - return IngestResult( - status="done", - format="zarr", - zarrPath=resolved, - assayNames=["RNA"], - summary={"assays": ["RNA"]}, - actions=["summarize_zarr"], - ) - - monkeypatch.setattr(main_module, "ingest", fake_ingest) - - -class _CheckpointOrchestrator(AgentOrchestrator): - """Minimal deterministic stage machine exercising persistence and resume.""" - - def __init__(self) -> None: - super().__init__(object()) - self.enrichmentExecutions = 0 - - def _continue( - self, - store: Any, - workflow: Any, - request_record: Any, - *, - answers: Mapping[str, Any], - resume_record: OrchestrationResumeRecord | None = None, - ) -> AutomatedWorkflowResult: - prefix = journal_module._ensure_orchestration_store(store) - ingest = journal_module._validated_done_outcome( - store, - prefix, - workflow.workflowRunId, - "ingest", - request_record, - [], - ) - assert ingest is not None - parents = [journal_module._parent_link(ingest)] - - enrichment = journal_module._validated_done_outcome( - store, - prefix, - workflow.workflowRunId, - "data_enrichment", - request_record, - parents, - ) - if enrichment is None: - self.enrichmentExecutions += 1 - report = DataEnrichmentReport.get_example() - report = report.model_copy( - update={ - "runInfo": report.runInfo.model_copy( - update={"runId": uuid.uuid4().hex} - ) - } - ) - reference = save_agent_report( - store, - workflow.workflowRunId, - report, - invocation=AgentInvocation( - agentName="data_enrichment", - inputs={"studyContext": request_record.request.studyContext}, - ), - ) - started = journal_module._start_attempt( - store.zw, - prefix, - workflow.workflowRunId, - "data_enrichment", - request_record, - parents, - inputs={"studyContext": request_record.request.studyContext}, - resume_record=resume_record, - ) - enrichment = journal_module._complete_attempt( - started, - status="done", - report_references=[reference], - ) - journal_module._save_outcome(store.zw, prefix, enrichment) - parents = [journal_module._parent_link(enrichment)] - - approved = answers.get("approvePlanChecksum") == _PLAN_CHECKSUM - completed_plan = journal_module._validated_done_outcome( - store, - prefix, - workflow.workflowRunId, - "preprocessing_plan", - request_record, - parents, - ) - if completed_plan is None: - started = journal_module._start_attempt( - store.zw, - prefix, - workflow.workflowRunId, - "preprocessing_plan", - request_record, - parents, - inputs={"planChecksum": _PLAN_CHECKSUM}, - resume_record=resume_record, - ) - if not approved: - paused = journal_module._complete_attempt( - started, - status="needsInput", - needs_input=WorkflowNeedsInput( - questions=[ - WorkflowQuestion( - questionId="approvePlanChecksum", - question="Approve the preprocessing plan?", - planChecksum=_PLAN_CHECKSUM, - ) - ] - ), - ) - journal_module._save_outcome(store.zw, prefix, paused) - return journal_module.paused_or_failed_result( - store, - workflow, - request_record, - paused, - ) - completed_plan = journal_module._complete_attempt( - started, - status="done", - outputs={"planChecksum": _PLAN_CHECKSUM}, - actions=["approve_preprocessing_plan"], - ) - journal_module._save_outcome(store.zw, prefix, completed_plan) - - terminal = finalize_agent_workflow( - store, - workflow.workflowRunId, - status="completed", - message="Checkpoint workflow completed", - ) - result = AutomatedWorkflowResult( - status="completed", - currentStage="preprocessing_plan", - zarrPath=str(Path(store.zarr_loc).resolve()), - workflowRun=terminal, - reportReferences=journal_module.all_report_references( - store, - prefix, - workflow.workflowRunId, - ), - ) - result = result.model_copy( - update={"contentSha256": journal_module._record_checksum(result)} - ) - journal_module._write_model_once( - store.zw, - journal_module._result_key(prefix, workflow.workflowRunId), - result, - ) - return result - - -def _start_paused_workflow( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - *, - workspace: str | None = None, -) -> tuple[_CheckpointOrchestrator, AutomatedWorkflowResult, Path]: - path = create_store(tmp_path / "data.zarr", workspace=workspace) - _mock_ingest(monkeypatch) - orchestrator = _CheckpointOrchestrator() - result = orchestrator.run( - AutomatedWorkflowRequest( - sourcePath=str(path), - zarrPath=str(path), - studyContext="A deterministic test study.", - workspace=workspace, - ) +def _saved_workflow(path: Path, *, workspace: str | None = None): + create_store(path, workspace=workspace) + orchestrator = AgentOrchestrator("test-model") + request = AutomatedWorkflowRequest( + sourcePath=str(path), + zarrPath=str(path), + workspace=workspace, + primaryAssay="RNA", + markerAssay="RNA", + analysisAssays=["RNA"], + studyContext="Independent donors in two conditions", + studyObjective="Resolve stable populations", ) - assert result.status == "needsInput" - assert result.workflowRun is not None - return orchestrator, result, path - - -def test_orchestration_records_are_plain_json_with_valid_checksums( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _orchestrator, result, path = _start_paused_workflow(tmp_path, monkeypatch) - workflow_id = result.workflowRun.workflowRunId - agents = path / "agents" - orchestration = agents / "orchestrations" - workflow_path = orchestration / workflow_id - - assert (agents / "zarr.json").is_file() - assert (agents / "store.json").is_file() - assert (orchestration / "zarr.json").is_file() - assert not (workflow_path / "zarr.json").exists() - assert not list(orchestration.rglob("c")) - - request = _assert_record_checksum(workflow_path / "request.json") - assert request["workflowRunId"] == workflow_id - records = sorted((workflow_path / "stages").rglob("*.json")) - assert records - for record_path in records: - _assert_record_checksum(record_path) - assert record_path.name in {"started.json", "outcome.json"} - - -def test_approval_resume_reuses_completed_stages_and_persists_answer_lineage( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - orchestrator, paused, path = _start_paused_workflow(tmp_path, monkeypatch) - workflow_id = paused.workflowRun.workflowRunId - workflow_path = path / "agents" / "orchestrations" / workflow_id - paused_outcome_path = next( - (workflow_path / "stages" / "preprocessing_plan").glob("*/outcome.json") + store = orchestrator.open_store(str(path), request) + store.cells.insert("condition", np.array(["a", "a", "b", "b"])) + record = orchestrator.initialize_request( + store, WorkflowIdentity("workflow-1", workspace), request ) - paused_outcome = json.loads(paused_outcome_path.read_text()) - - completed = orchestrator.resume( - AutomatedWorkflowResumeRequest( - zarrPath=str(path), - workflowRunId=workflow_id, - answers={"approvePlanChecksum": _PLAN_CHECKSUM}, - ) + resume = AutomatedWorkflowResumeRequest( + zarrPath=str(path), workspace=workspace, workflowRunId="workflow-1" ) - - assert completed.status == "completed" - assert completed.workflowRun is not None - assert completed.workflowRun.status == "completed" - assert orchestrator.enrichmentExecutions == 1 - assert len(list((workflow_path / "stages" / "ingest").glob("*/outcome.json"))) == 1 - assert ( - len(list((workflow_path / "stages" / "data_enrichment").glob("*/outcome.json"))) - == 1 - ) - assert ( - len( - list( - (workflow_path / "stages" / "preprocessing_plan").glob("*/outcome.json") - ) - ) - == 2 - ) - - resume_path = next((workflow_path / "resumes").glob("*.json")) - resume = _assert_record_checksum(resume_path) - assert resume["answeredAttempt"]["attemptId"] == paused_outcome["attemptId"] - assert resume["questionIds"] == ["approvePlanChecksum"] - assert resume["answers"] == {"approvePlanChecksum": _PLAN_CHECKSUM} - resumed_start = next( - payload - for payload in ( - json.loads(path.read_text()) - for path in (workflow_path / "stages" / "preprocessing_plan").glob( - "*/started.json" - ) - ) - if "resumeLineage" in payload["inputs"] - ) - assert resumed_start["inputs"]["resumeLineage"] == { - "resumeId": resume["resumeId"], - "answeredAttempt": resume["answeredAttempt"], - "questionIds": ["approvePlanChecksum"], - } - - persisted = _assert_record_checksum(workflow_path / "result.json") - assert persisted["status"] == "completed" - assert persisted["workflowRun"]["status"] == "completed" - assert load_agent_workflow(path, workflow_id).status == "completed" + return orchestrator, store, record, resume -def test_resume_does_not_mutate_constructor_configuration( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, +@pytest.mark.parametrize("workspace", [None, "analysis"]) +def test_resume_preserves_selected_rna_workspace_and_original_metadata( + tmp_path, workspace ) -> None: - orchestrator, paused, path = _start_paused_workflow(tmp_path, monkeypatch) - assert paused.workflowRun is not None - constructor_config = AutomatedWorkflowConfig(primaryInitialCandidates=2) - orchestrator.config = constructor_config - - completed = orchestrator.resume( - AutomatedWorkflowResumeRequest( - zarrPath=str(path), - workflowRunId=paused.workflowRun.workflowRunId, - answers={"approvePlanChecksum": _PLAN_CHECKSUM}, + orchestrator, store, record, resume = _saved_workflow( + tmp_path / "rna.zarr", workspace=workspace + ) + reopened_record, reopened = orchestrator.load_request_for_resume(resume) + assert reopened_record == record + assert reopened.workspace == workspace + assert reopened.summary().default_assay == "RNA" + store.cells.insert("derived_agent_column", np.array([1, 1, 2, 2])) + assert orchestrator.load_request_for_resume(resume)[0] == record + store.cells.insert("condition", np.array(["a", "b", "b", "b"]), overwrite=True) + with pytest.raises(ValueError, match="metadata changed"): + orchestrator.load_request_for_resume(resume) + + +def test_resume_rejects_changed_counts_model_and_execution_settings(tmp_path) -> None: + orchestrator, store, record, resume = _saved_workflow(tmp_path / "rna.zarr") + with pytest.raises(ValueError, match="model differs"): + AgentOrchestrator("another-model").load_request_for_resume(resume) + with pytest.raises(ValueError, match="execution settings differ"): + AgentOrchestrator( + "test-model", config=AutomatedWorkflowConfig(randomSeed=3) + ).load_request_for_resume(resume) + root = zarr.open_group(str(store.zarr_loc), mode="r+") + counts = root["RNA/counts"] + counts[0, 0] = int(counts[0, 0]) + 1 + with pytest.raises(ValueError, match="Selected RNA data"): + orchestrator.load_request_for_resume(resume) + assert ( + journal.read_request( + store.zw, journal._orchestration_prefix(store), record.workflowRunId ) + == record ) - assert completed.status == "completed" - assert orchestrator.config == constructor_config +def test_resume_never_reads_an_unrelated_latest_workflow(tmp_path, monkeypatch) -> None: + orchestrator, store, record, resume = _saved_workflow(tmp_path / "rna.zarr") + other_request = record.request.model_copy( + update={"studyObjective": "A different question"} + ) + orchestrator.initialize_request(store, WorkflowIdentity("other-run"), other_request) + seen = [] -@pytest.mark.parametrize( - ("answers", "expected_note"), - [ - ({}, "Missing resume answer"), - ({"approvePlanChecksum": "approve"}, "persisted plan checksum"), - ( - { - "approvePlanChecksum": _PLAN_CHECKSUM, - "unexpectedAnswer": "ignored", - }, - "Unknown resume answer key", - ), - ], -) -def test_invalid_resume_answers_retain_the_persisted_pause( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - answers: dict[str, Any], - expected_note: str, -) -> None: - orchestrator, paused, path = _start_paused_workflow(tmp_path, monkeypatch) - assert paused.workflowRun is not None - workflow_id = paused.workflowRun.workflowRunId - - result = orchestrator.resume( - AutomatedWorkflowResumeRequest( - zarrPath=str(path), - workflowRunId=workflow_id, - answers=answers, + def stop(request): + seen.append(request.workflowRunId) + return AutomatedWorkflowResult( + status="abstained", workflowRunId=request.workflowRunId ) - ) - assert result.status == "needsInput" - assert result.currentStage == "preprocessing_plan" - assert result.needsInput == paused.needsInput - assert result.workflowRun is not None - assert result.workflowRun.status == "running" - assert any(expected_note in note for note in result.notes) - workflow_path = path / "agents" / "orchestrations" / workflow_id + monkeypatch.setattr(orchestrator, "resume", stop) assert ( - len( - list( - (workflow_path / "stages" / "preprocessing_plan").glob("*/outcome.json") - ) - ) - == 1 + orchestrator._reuse_or_resume(record.request).workflowRunId + == record.workflowRunId ) - assert len(list((workflow_path / "resumes").glob("*.json"))) == 1 + assert seen == [record.workflowRunId] + conflicting = AgentOrchestrator("different-model") + with pytest.raises(ValueError, match="different model or execution settings"): + conflicting._reuse_or_resume(record.request) -def test_interrupted_attempt_supersedes_a_historical_pause_for_resume( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, +def test_older_request_reports_actionable_hard_break_without_rewriting( + tmp_path, ) -> None: - orchestrator, paused, path = _start_paused_workflow(tmp_path, monkeypatch) - assert paused.workflowRun is not None - workflow_id = paused.workflowRun.workflowRunId - record, store = orchestrator.load_request_for_resume( - AutomatedWorkflowResumeRequest( - zarrPath=str(path), - workflowRunId=workflow_id, - ) + orchestrator, store, record, resume = _saved_workflow(tmp_path / "rna.zarr") + path = Path(store.zarr_loc) / "agents/orchestrations/old/request.json" + path.parent.mkdir() + content = '{"formatVersion":2,"workflowRunId":"old"}' + path.write_text(content) + with pytest.raises( + ValueError, match="older requests cannot be resumed or regenerated" + ): + journal.open_analysis_store(store.zarr_loc, "old") + assert path.read_text() == content + assert store.get_assay("RNA").rawData.shape == (4, 4) + + +def _pause_with_interrupted_answer( + *, failed: bool = False, tuning: bool = False, interrupted: bool = True +): + store, prefix, record = memory_journal() + start = journal._start_attempt( + store.zw, prefix, record.workflowRunId, "ingest", record, [] + ) + ingest = journal._complete_attempt(start, status="done") + journal._save_outcome(store.zw, prefix, ingest) + parents = [journal._parent_link(ingest)] + if tuning: + for stage in ( + "data_enrichment", + "rna_quality_metrics", + "experimental_context", + "preprocessing_plan", + "preprocessing", + ): + started = journal._start_attempt( + store.zw, prefix, record.workflowRunId, stage, record, parents + ) + completed = journal._complete_attempt(started, status="done") + journal._save_outcome(store.zw, prefix, completed) + parents = [journal._parent_link(completed)] + stage = "parameter_tuning" if tuning else "data_enrichment" + question_id = "parameter_tuning" if tuning else "enrichmentDirections" + start = journal._start_attempt( + store.zw, prefix, record.workflowRunId, stage, record, parents + ) + paused = journal._complete_attempt( + start, + status="needsInput", + needs_input=WorkflowNeedsInput( + questions=[ + WorkflowQuestion( + questionId=question_id, + question="Inspect unresolved scientific evidence" + if tuning + else "Confirm RNA organism", + ) + ] + ), ) - prefix = journal_module._ensure_orchestration_store(store) - journal_module._start_attempt( + journal._save_outcome(store.zw, prefix, paused) + answers = { + question_id: {"action": "defer", "rationale": "More evidence is required"} + if tuning + else {"organism": "human"} + } + resume_record = OrchestrationResumeRecord( + workflowRunId=record.workflowRunId, + answeredAttempt=journal._parent_link(paused), + answers=answers, + questionIds=list(answers), + ) + if not interrupted: + return store, record, answers, resume_record + interrupted_start = journal._start_attempt( store.zw, prefix, - workflow_id, - "preprocessing", + record.workflowRunId, + stage, record, - [], - inputs={"simulatedCrash": True}, + parents, + resume_record=resume_record, ) - - with pytest.raises(ValueError, match="no active persisted questions"): - orchestrator.resume( - AutomatedWorkflowResumeRequest( - zarrPath=str(path), - workflowRunId=workflow_id, - answers={"approvePlanChecksum": _PLAN_CHECKSUM}, - ) + if failed: + outcome = journal._complete_attempt( + interrupted_start, status="failed", error="Provider unavailable" ) + journal._save_outcome(store.zw, prefix, outcome) + return store, record, answers, resume_record - resumed = orchestrator.resume( - AutomatedWorkflowResumeRequest( - zarrPath=str(path), - workflowRunId=workflow_id, - ) + +@pytest.mark.parametrize("failed", [False, True]) +@pytest.mark.parametrize("tuning", [False, True]) +def test_interrupted_answered_attempt_resumes_without_asking_or_spending_again( + monkeypatch, failed, tuning +) -> None: + store, record, answers, original = _pause_with_interrupted_answer( + failed=failed, tuning=tuning ) + orchestrator = AgentOrchestrator("test-model") + monkeypatch.setattr( + orchestrator, "load_request_for_resume", lambda request: (record, store) + ) + captured = {} - assert resumed.status == "needsInput" - assert not any("Missing resume answer" in note for note in resumed.notes) - workflow_path = path / "agents" / "orchestrations" / workflow_id - assert ( - len( - list( - (workflow_path / "stages" / "preprocessing_plan").glob("*/outcome.json") - ) + def continue_work(*args, **kwargs): + captured.update(kwargs) + return AutomatedWorkflowResult( + status="abstained", notes=["Stopped at the resumed boundary"] + ) + + monkeypatch.setattr(orchestrator, "_continue", continue_work) + result = orchestrator.resume( + AutomatedWorkflowResumeRequest( + zarrPath="analysis.zarr", workflowRunId=record.workflowRunId ) - == 2 ) + assert result.status == "abstained" + assert captured["answers"] == answers + assert captured["resume_record"].answeredAttempt == original.answeredAttempt -def test_interrupted_answered_attempt_inherits_exact_persisted_answers( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, +@pytest.mark.parametrize("tuning", [False, True]) +def test_empty_answer_reenters_only_the_checkpointed_tuning_stage( + monkeypatch, tuning ) -> None: - orchestrator, paused, path = _start_paused_workflow(tmp_path, monkeypatch) - assert paused.workflowRun is not None - workflow_id = paused.workflowRun.workflowRunId - record, store = orchestrator.load_request_for_resume( - AutomatedWorkflowResumeRequest( - zarrPath=str(path), - workflowRunId=workflow_id, - ) + store, record, _, _ = _pause_with_interrupted_answer( + tuning=tuning, interrupted=False ) - prefix = journal_module._ensure_orchestration_store(store) - paused_outcome = journal_module._stage_outcomes( - store.zw, - prefix, - workflow_id, - "preprocessing_plan", - )[0] - first_resume = OrchestrationResumeRecord( - workflowRunId=workflow_id, - resumeId="resume-before-crash", - createdAtNs=paused_outcome.completedAtNs + 1, - answeredAttempt=journal_module._parent_link(paused_outcome), - questionIds=["approvePlanChecksum"], - answers={"approvePlanChecksum": _PLAN_CHECKSUM}, - ) - first_resume = first_resume.model_copy( - update={"contentSha256": journal_module._record_checksum(first_resume)} - ) - journal_module._write_model_once( - store.zw, - journal_module._resume_key( - prefix, - workflow_id, - first_resume.resumeId, - ), - first_resume, + orchestrator = AgentOrchestrator("test-model") + monkeypatch.setattr( + orchestrator, "load_request_for_resume", lambda request: (record, store) ) - journal_module._start_attempt( - store.zw, - prefix, - workflow_id, - "preprocessing_plan", - record, - paused_outcome.parentAttempts, - inputs={"planChecksum": _PLAN_CHECKSUM}, - resume_record=first_resume, + before = journal.analysis_snapshot(store, record.workflowRunId) + assert before["status"] == "needsInput" + assert before["stages"][-1]["stage"] == ( + "parameter_tuning" if tuning else "data_enrichment" ) + captured = {} - with pytest.raises(ValueError, match="Cannot change answers"): - orchestrator.resume( - AutomatedWorkflowResumeRequest( - zarrPath=str(path), - workflowRunId=workflow_id, - answers={"approvePlanChecksum": "b" * 64}, - ) + def continue_work(*args, **kwargs): + captured.update(kwargs) + return AutomatedWorkflowResult( + status="abstained", currentStage="parameter_tuning" ) - completed = orchestrator.resume( + monkeypatch.setattr(orchestrator, "_continue", continue_work) + result = orchestrator.resume( AutomatedWorkflowResumeRequest( - zarrPath=str(path), - workflowRunId=workflow_id, + zarrPath="analysis.zarr", workflowRunId=record.workflowRunId ) ) + if tuning: + assert result.status == "abstained" + assert captured == {"answers": {}, "resume_record": None} + else: + assert result.status == "needsInput" + assert captured == {} + assert journal.analysis_snapshot(store, record.workflowRunId) == before - assert completed.status == "completed" - resume_records = sorted( - (path / "agents" / "orchestrations" / workflow_id / "resumes").glob("*.json") - ) - assert len(resume_records) == 2 - inherited = next( - json.loads(value.read_text()) - for value in resume_records - if value.stem != first_resume.resumeId + +def test_explicit_tuning_answer_keeps_its_exact_paused_attempt(monkeypatch) -> None: + store, record, answers, original = _pause_with_interrupted_answer( + tuning=True, interrupted=False ) - assert inherited["answers"] == first_resume.answers - assert inherited["answeredAttempt"] == first_resume.answeredAttempt.model_dump( - mode="json" + orchestrator = AgentOrchestrator("test-model") + monkeypatch.setattr( + orchestrator, "load_request_for_resume", lambda request: (record, store) ) - assert inherited["questionIds"] == first_resume.questionIds + captured = {} + def continue_work(*args, **kwargs): + captured.update(kwargs) + return AutomatedWorkflowResult(status="abstained") -def test_resume_rejects_changed_dataset_fingerprint( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - orchestrator, paused, path = _start_paused_workflow(tmp_path, monkeypatch) - workflow_id = paused.workflowRun.workflowRunId - root = zarr.open_group(str(path), mode="r+") - root["RNA"].attrs["dataset_fingerprint"] = "changed-dataset" - - with pytest.raises(ValueError, match="dataset fingerprints"): - orchestrator.resume( - AutomatedWorkflowResumeRequest( - zarrPath=str(path), - workflowRunId=workflow_id, - answers={"approvePlanChecksum": _PLAN_CHECKSUM}, - ) + monkeypatch.setattr(orchestrator, "_continue", continue_work) + result = orchestrator.resume( + AutomatedWorkflowResumeRequest( + zarrPath="analysis.zarr", + workflowRunId=record.workflowRunId, + answers=answers, ) + ) + assert result.status == "abstained" + assert captured["answers"] == answers + assert captured["resume_record"].answeredAttempt == original.answeredAttempt -def test_resume_rejects_request_envelope_and_destination_mismatch( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, +def test_completed_resume_regenerates_report_without_reentering_analysis( + monkeypatch, ) -> None: - orchestrator, paused, path = _start_paused_workflow(tmp_path, monkeypatch) - workflow_id = paused.workflowRun.workflowRunId - request_path = path / "agents" / "orchestrations" / workflow_id / "request.json" - payload = json.loads(request_path.read_text()) - payload["request"]["studyContext"] = "Tampered study context" - request_path.write_text(json.dumps(payload)) - - with pytest.raises(ValueError, match="request checksum"): - orchestrator.resume( - AutomatedWorkflowResumeRequest( - zarrPath=str(path), - workflowRunId=workflow_id, - ) - ) - - copied = tmp_path / "copied.zarr" - shutil.copytree(path, copied) - with pytest.raises(ValueError, match="zarrPath"): - orchestrator.resume( - AutomatedWorkflowResumeRequest( - zarrPath=str(copied), - workflowRunId=workflow_id, - ) + store, prefix, record = memory_journal() + orchestrator = AgentOrchestrator("test-model") + monkeypatch.setattr( + orchestrator, "load_request_for_resume", lambda request: (record, store) + ) + monkeypatch.setattr( + journal, + "analysis_snapshot", + lambda *_: { + "status": "completed", + "finalAnalysis": {"limitations": ["No donor replication"]}, + }, + ) + calls = [] + monkeypatch.setattr( + AutomatedWorkflowResult, "report", lambda self: calls.append(self.workflowRunId) + ) + monkeypatch.setattr( + orchestrator, + "_continue", + lambda *_a, **_k: pytest.fail("Numerical workflow repeated"), + ) + result = orchestrator.resume( + AutomatedWorkflowResumeRequest( + zarrPath="analysis.zarr", workflowRunId=record.workflowRunId ) + ) + assert result.status == "completed" + assert result.limitations == ["No donor replication"] + assert calls == [record.workflowRunId] -def test_workspace_records_are_isolated_and_wrong_workspace_cannot_resume( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - orchestrator, paused, path = _start_paused_workflow( - tmp_path, - monkeypatch, - workspace="workspace_a", +def test_nonstage_failure_keeps_last_stage_and_resume_address(monkeypatch) -> None: + store, prefix, record = memory_journal("analysis") + journal._start_attempt( + store.zw, prefix, record.workflowRunId, "parameter_tuning", record, [] ) - workflow_id = paused.workflowRun.workflowRunId - root = zarr.open_group(str(path), mode="r+") - root.create_group("workspace_b") - - assert ( - path - / "workspace_a" - / "agents" - / "orchestrations" - / workflow_id - / "request.json" - ).is_file() - assert not (path / "agents").exists() - assert not (path / "workspace_b" / "agents").exists() - - with pytest.raises(FileNotFoundError): - orchestrator.resume( - AutomatedWorkflowResumeRequest( - zarrPath=str(path), - workflowRunId=workflow_id, - workspace="workspace_b", - ) - ) + orchestrator = AgentOrchestrator("test-model") + def fail(*args, **kwargs): + raise RuntimeError("The selected full-cohort candidate lacks markers") -def test_cancel_finalizes_abandoned_and_prevents_future_resume( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - orchestrator, paused, path = _start_paused_workflow(tmp_path, monkeypatch) - workflow_id = paused.workflowRun.workflowRunId - request = AutomatedWorkflowResumeRequest( - zarrPath=str(path), - workflowRunId=workflow_id, - ) - - cancelled = orchestrator.cancel(request, message="Stop this test workflow") - - assert cancelled.status == "abandoned" - assert cancelled.workflowRun is not None - assert cancelled.workflowRun.status == "abandoned" - assert load_agent_workflow(path, workflow_id).status == "abandoned" - result_path = path / "agents" / "orchestrations" / workflow_id / "result.json" - persisted = _assert_record_checksum(result_path) - assert persisted["status"] == "abandoned" - assert persisted["notes"] == ["Stop this test workflow"] - with pytest.raises(RuntimeError, match="Cannot resume"): - orchestrator.resume(request) - - result_path.unlink() - repaired = orchestrator.resume(request) - assert repaired.status == "abandoned" - assert repaired.workflowRun is not None - assert repaired.workflowRun.status == "abandoned" - _assert_record_checksum(result_path) - with pytest.raises(RuntimeError, match="Cannot resume"): - orchestrator.resume(request) - - -def test_failed_stage_preserves_completed_operation_journal(tmp_path: Path) -> None: - path = create_store(tmp_path / "failure-journal.zarr") - store = DataStore( - str(path), - default_assay="RNA", - min_features_per_cell=-1, - mito_pattern="", - ribo_pattern="", - zarr_mode="r+", - ) - workflow = create_agent_workflow(store, workflow_run_id="failure-journal") - request_record = OrchestrationRequestRecord( - workflowRunId=workflow.workflowRunId, - request=AutomatedWorkflowRequest( - sourcePath=str(path), - zarrPath=str(path), - studyContext="A failed-stage journal test.", - allowAssumptions=True, - ), - config=AutomatedWorkflowConfig(), + monkeypatch.setattr(orchestrator, "_execute_stages", fail) + result = orchestrator._continue( + store, WorkflowIdentity(record.workflowRunId, "analysis"), record, answers={} ) - prefix = journal_module._ensure_orchestration_store(store) - started = journal_module._start_attempt( - store.zw, - prefix, - workflow.workflowRunId, - "preprocessing", - request_record, - [], - ) - operation = { - "operation": "filter_cells", - "attrs": ["RNA_nCounts"], - "resetPrevious": True, - } + assert result.currentStage == "parameter_tuning" + assert result.workflowRunId == record.workflowRunId + assert result.workspace == "analysis" + assert "lacks markers" in result.notes[0] - outcome = journal_module.finish_exception( - store, - prefix, - workflow, - started, - RuntimeError("failure after filtering"), - actions=["cell_qc_global:validated"], - outputs={"operations": [operation]}, - notes=["The completed selection mutation is retained for audit."], - ) - - assert outcome.status == "failed" - assert outcome.actions == ["cell_qc_global:validated"] - assert outcome.outputs["operations"] == [operation] - assert outcome.notes == ["The completed selection mutation is retained for audit."] - persisted = journal_module._stage_outcomes( - store.zw, - prefix, - workflow.workflowRunId, - "preprocessing", - ) - assert persisted == [outcome] - assert load_agent_workflow(store, workflow.workflowRunId).status == "failed" - - -def test_failed_stage_links_report_committed_before_exception(tmp_path: Path) -> None: - path = create_store(tmp_path / "failure-report-link.zarr") - store = DataStore( - str(path), - default_assay="RNA", - min_features_per_cell=-1, - mito_pattern="", - ribo_pattern="", - zarr_mode="r+", - ) - workflow = create_agent_workflow(store, workflow_run_id="failure-report-link") - request_record = OrchestrationRequestRecord( - workflowRunId=workflow.workflowRunId, - request=AutomatedWorkflowRequest( - sourcePath=str(path), - zarrPath=str(path), - studyContext="A report-link crash test.", - allowAssumptions=True, - ), - config=AutomatedWorkflowConfig(), - ) - prefix = journal_module._ensure_orchestration_store(store) - started = journal_module._start_attempt( - store.zw, - prefix, - workflow.workflowRunId, - "data_enrichment", - request_record, - [], - inputs={"studyContext": request_record.request.studyContext}, - ) - report = DataEnrichmentReport.get_example().model_copy( - update={ - "runInfo": AgentRunInfo( - agentName="data_enrichment", - runId=uuid.uuid4().hex, - ) - } - ) - _saved, reference = journal_module._save_stage_report( - store, - started, - report, - invocation=AgentInvocation( - agentName="data_enrichment", - inputs={"studyContext": request_record.request.studyContext}, - ), - expected_type=DataEnrichmentReport, - ) - outcome = journal_module.finish_exception( - store, - prefix, - workflow, - started, - RuntimeError("failure after report commit"), +def test_gene_annotations_are_part_of_completed_request_identity(tmp_path) -> None: + orchestrator, store, record, resume = _saved_workflow(tmp_path / "rna.zarr") + store.get_assay("RNA").feats.insert( + "names", np.array(["GENE9", "RPS3", "GENE1", "GENE2"]), overwrite=True ) + with pytest.raises(ValueError, match="data or relevant metadata changed"): + orchestrator.load_request_for_resume(resume) - assert outcome.reportReferences == [reference] - assert load_agent_report(store, reference) == report +def test_provider_configuration_distinguishes_identically_named_models() -> None: + from scarf.agent.orchestrator.main import _model_identity -def test_orphan_report_recovery_uses_semantic_stage_identity( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _orchestrator, paused, path = _start_paused_workflow(tmp_path, monkeypatch) - assert paused.workflowRun is not None - workflow_id = paused.workflowRun.workflowRunId - store = DataStore( - str(path), - min_features_per_cell=-1, - mito_pattern="", - ribo_pattern="", - zarr_mode="r+", - ) - started = WorkflowStageAttempt( - workflowRunId=workflow_id, - stage="data_enrichment", - attemptId="crash-attempt-one", - status="started", - startedAtNs=1, - requestSha256="1" * 64, - configSha256="2" * 64, - inputs={ - "effectiveContext": "same", - "resumeLineage": { - "resumeId": "first-resume", - "answeredAttempt": WorkflowStageLink.get_example().model_dump( - mode="json" - ), - "questionIds": ["approvePlanChecksum"], - }, - }, + first = SimpleNamespace( + model_name="rna-model", + system="openai", + settings={"temperature": 0}, + provider=SimpleNamespace(name="server", base_url="https://first.example/v1"), ) - report = DataEnrichmentReport.get_example().model_copy( - update={ - "runInfo": DataEnrichmentReport.get_example().runInfo.model_copy( - update={"runId": uuid.uuid4().hex} - ) - } - ) - persisted_report, reference = journal_module._save_stage_report( - store, - started, - report, - invocation=AgentInvocation( - agentName="data_enrichment", - inputs={"effectiveContext": "same"}, - ), - expected_type=DataEnrichmentReport, - ) - assert persisted_report == report - - retried = started.model_copy( - update={ - "attemptId": "crash-attempt-two", - "startedAtNs": 2, - "inputs": { - "effectiveContext": "same", - "resumeLineage": { - "resumeId": "second-resume", - "answeredAttempt": None, - "questionIds": [], - }, - }, + second = SimpleNamespace( + **{ + **vars(first), + "provider": SimpleNamespace( + name="server", base_url="https://second.example/v1" + ), } ) - assert journal_module._stage_execution_id(started) == ( - journal_module._stage_execution_id(retried) - ) - changed_context = retried.model_copy( - update={"inputs": {"effectiveContext": "different"}} + assert _model_identity(first) != _model_identity(second) + text_only = SimpleNamespace(**vars(first), supports_image_input=False) + assert _model_identity(first) != _model_identity(text_only) + assert _model_identity(text_only) == _model_identity( + SimpleNamespace(**vars(first), profile={"supports_image_input": False}) ) - assert journal_module._stage_execution_id(started) != ( - journal_module._stage_execution_id(changed_context) - ) - recovered = journal_module._recover_persisted_stage_report( - store, - retried, - agent_name="data_enrichment", - expected_type=DataEnrichmentReport, - ) - - assert recovered == (report, reference) - matching = [ - value - for value in list_agent_reports( - store, - workflow_id, - agent_name="data_enrichment", - ) - if value.agentRunId == reference.agentRunId - ] - assert matching == [reference] - - -def test_data_enrichment_recovers_report_after_outcome_write_crash( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - orchestrator, paused, path = _start_paused_workflow(tmp_path, monkeypatch) - assert paused.workflowRun is not None - workflow_id = paused.workflowRun.workflowRunId - store = DataStore( - str(path), - min_features_per_cell=-1, - mito_pattern="", - ribo_pattern="", - zarr_mode="r+", - ) - prefix = journal_module._ensure_orchestration_store(store) - request_record = journal_module._read_model( - store.zw, - journal_module._request_key(prefix, workflow_id), - OrchestrationRequestRecord, - ) - assert isinstance(request_record, OrchestrationRequestRecord) - workflow = load_agent_workflow(store, workflow_id) - calls = 0 - - class CountingAgent: - config = AgentRunConfig() - - def run(self, *_args: Any, **_kwargs: Any) -> DataEnrichmentReport: - nonlocal calls - calls += 1 - return DataEnrichmentReport.get_example().model_copy( - update={ - "runInfo": DataEnrichmentReport.get_example().runInfo.model_copy( - update={"runId": uuid.uuid4().hex} - ) - } - ) - - monkeypatch.setattr( - context_module, - "DataEnrichmentAgent", - lambda *_args, **_kwargs: CountingAgent(), - ) - save_outcome = journal_module._save_outcome - crashed = False - - def crash_after_report(*args: Any, **kwargs: Any) -> None: - nonlocal crashed - outcome = args[2] - if outcome.stage == "data_enrichment" and not crashed: - crashed = True - raise KeyboardInterrupt("simulated process crash") - save_outcome(*args, **kwargs) - - monkeypatch.setattr(journal_module, "_save_outcome", crash_after_report) - cell_selection = ArtifactReferenceModel.from_artifact_ref( - store.snapshot_cell_selection("I") - ) - with pytest.raises(KeyboardInterrupt, match="simulated process crash"): - orchestrator.data_enrichment_stage( - store, - workflow, - request_record, - [], - cell_selection, - {}, - ) - monkeypatch.setattr(journal_module, "_save_outcome", save_outcome) - - outcome, recovered = orchestrator.data_enrichment_stage( - store, - workflow, - request_record, - [], - cell_selection, - {}, - ) - - assert outcome.status == "done" - assert recovered.status == "done" - assert calls == 1 - assert "recover_persisted_data_enrichment_report" in outcome.actions diff --git a/tests/test_agent_orchestrator_stages.py b/tests/test_agent_orchestrator_stages.py index b0657397..45750d0d 100644 --- a/tests/test_agent_orchestrator_stages.py +++ b/tests/test_agent_orchestrator_stages.py @@ -1,5 +1,7 @@ """Context, preprocessing, tuning, integration, and finalization contracts.""" +from tests.agent_examples import example + import uuid from collections.abc import Mapping from pathlib import Path @@ -11,8 +13,8 @@ import scarf.agent.orchestrator.context as context_module import scarf.agent.orchestrator.journal as journal_module -import scarf.agent.orchestrator.tuning as tuning_module -import scarf.agent.parameter_tuning as parameter_tuning_module +from scarf.agent.parameter_tuning.selection import harmony_acceptance_gate +from scarf.agent.orchestrator.preprocessing import PreprocessingStagesMixin from scarf.agent.config import AgentRunConfig from scarf.agent.data_enrichment import ( AssayFeatureInspection, @@ -22,49 +24,39 @@ FeatureSelectionPolicy, ) from scarf.agent.experimental_context import ( - BatchCorrectionPlan, CellQcPlan, CellQcProfileEvidence, + ExperimentalContextDependencies, ExperimentalContextResult, NamedArtifactSource, ) +from scarf.agent.experimental_context.tools import ( + persist_context_evidence, + restore_context_evidence, +) from scarf.agent.orchestrator import ( AgentOrchestrator, - AssayPreprocessingPlan, - AutomatedPreprocessingPlan, AutomatedWorkflowConfig, AutomatedWorkflowRequest, - PreprocessedAssayHandoff, - WorkflowStageAttempt, - WorkflowStageLink, ) -from scarf.agent.orchestrator.models import OrchestrationRequestRecord -from scarf.agent.persistence import ( - AgentInvocation, - AgentReportReference, - create_agent_workflow, - list_agent_reports, - load_agent_record, - load_agent_workflow, - save_agent_report, +from scarf.agent.orchestrator.models import ( + AutomatedPreprocessingPlan, + WorkflowStageAttempt, ) +from scarf.agent.orchestrator.models import OrchestrationRequestRecord, WorkflowIdentity from scarf.agent.parameter_tuning import ( ArtifactRecord, - FinalGraphNeedsInput, - FinalGraphSelection, - IntegrationCandidateEvaluation, - IntegrationMetrics, ParameterCandidateEvaluation, - ParameterTuningReport, - final_graph_options, - finalize_parameter_tuning_selection, - select_final_parameter_graph, ) +from scarf.agent.cell_quality.profiles import RegisteredCellQcProfile from scarf.agent.types import ( AgentRunInfo, ArtifactReferenceModel, BatchSafetyEvidence, - ExperimentalTuningHandoff, +) +from scarf.agent.parameter_tuning.diagnostics import ( + _select_capture_cells, + resolve_native_doublet_inputs, ) from scarf.datastore.datastore import DataStore from scarf.storage.refs import ArtifactRef @@ -74,6 +66,51 @@ _PLAN_CHECKSUM = "a" * 64 +def _save_input_evidence(store, workflow, request, enrichment): + prefix = journal_module._ensure_orchestration_store(store) + started = journal_module._start_attempt( + store.zw, prefix, workflow.workflowRunId, "data_enrichment", request, [] + ) + return journal_module._save_stage_report( + store, started, enrichment, expected_type=DataEnrichmentReport + )[1] + + +def _measured_context_characterization(store, cell_selection): + from scarf.agent.experimental_context.characterization import ( + characterize_covariates, + ) + from scarf.agent.orchestrator.models import artifact_model_to_ref + + for name, values in { + "sample": ["s1", "s1", "s2", "s2"], + "donor": ["d1", "d1", "d2", "d2"], + "treatment": ["control", "control", "treated", "treated"], + "batch": ["b1", "b1", "b2", "b2"], + }.items(): + store.cells.insert(name, np.asarray(values)) + return characterize_covariates( + store, + cellSelection=artifact_model_to_ref(cell_selection), + model=None, + directions={ + "columnDomains": { + "sample": "design", + "donor": "design", + "treatment": "biological", + "batch": "technical", + }, + "coefficientsOfInterest": ["treatment"], + "unitsOfInference": { + "treatment": { + "observationUnit": "sample", + "independentUnit": "donor", + } + }, + }, + ) + + def _cell_selection_model() -> ArtifactReferenceModel: return ArtifactReferenceModel( scope="datastore", @@ -171,6 +208,7 @@ def _planning_inputs( DataEnrichmentReport, ExperimentalContextResult, WorkflowStageAttempt, + CellQcPlan, ]: store = _PlanningStore(assays) policies = [ @@ -189,17 +227,19 @@ def _planning_inputs( sourcePath="dataset.zarr", zarrPath="dataset.zarr", studyContext="A bounded plan-construction test.", - allowAssumptions=True, + studyObjective="Discover stable RNA populations.", primaryAssay=primary_assay, markerAssay=marker_assay, - analysisAssays=analysis_assays or list(assays), + analysisAssays=analysis_assays or [], ) request_record = OrchestrationRequestRecord( + inputIdentity={}, + modelIdentity="test-model", workflowRunId="planning-test", request=request, config=config or AutomatedWorkflowConfig(), ) - experimental = ExperimentalContextResult.get_example() + experimental = example(ExperimentalContextResult) ingest_outcome = WorkflowStageAttempt( workflowRunId="planning-test", stage="ingest", @@ -209,7 +249,14 @@ def _planning_inputs( completedAtNs=2, outputs={"format": "zarr"}, ) - return store, request_record, enrichment, experimental, ingest_outcome + return ( + store, + request_record, + enrichment, + experimental, + ingest_outcome, + example(CellQcPlan), + ) def _build_plan( @@ -220,111 +267,19 @@ def _build_plan( return AgentOrchestrator(object()).build_preprocessing_plan(*inputs) -def _native_assay_report(assay: str, token: int) -> ParameterTuningReport: - evaluation = ParameterCandidateEvaluation.get_example().model_copy( - update={ - "artifacts": { - "neighbors": ArtifactRecord( - assay=assay, - kind="neighbors", - artifactId=f"{token + 2:064x}", - ), - "connectivityMap": ArtifactRecord( - assay=assay, - kind="connectivity_map", - artifactId=f"{token:064x}", - ), - "clusters": ArtifactRecord( - assay=assay, - kind="cluster_labels", - artifactId=f"{token + 1:064x}", - ), - }, - "cellSelection": _cell_selection_model(), - "clusterColumn": f"{assay}_agent_clusters", - "evidenceIds": ["candidate:baseline:clusters"], - } - ) - return ParameterTuningReport( - status="done", - fromAssay=assay, - cellSelection=_cell_selection_model(), - evaluations=[evaluation], - recommendedCandidateId=evaluation.candidateId, - selectedArtifacts=dict(evaluation.artifacts), - evidenceIds=list(evaluation.evidenceIds), - stopReason="The bounded screen completed.", - ) - - -def _native_batch_report(*assays: str) -> ParameterTuningReport: - reports = { - assay: _native_assay_report(assay, index * 10 + 1) - for index, assay in enumerate(assays) - } - primary = reports[assays[0]] - return ParameterTuningReport( - status="done", - fromAssay=assays[0], - cellSelection=_cell_selection_model(), - evaluations=list(primary.evaluations), - recommendedCandidateId=primary.recommendedCandidateId, - selectedArtifacts=dict(primary.selectedArtifacts), - assayReports=reports, - recommendedByAssay={ - assay: report.recommendedCandidateId or "" - for assay, report in reports.items() - }, - totalCandidates=sum(len(report.evaluations) for report in reports.values()), - graphAssay=assays[0], - markerAssay=assays[0], - runInfo=AgentRunInfo( - agentName="parameter_tuning", - runId=uuid.uuid4().hex, - ), - ) - - -def _eligible_integration() -> IntegrationCandidateEvaluation: - return IntegrationCandidateEvaluation( - integrationId="wnn_resolution_1", - method="wnn", - assays=["RNA", "ADT"], - status="done", - eligible=True, - cellSelection=_cell_selection_model(), - resolution=1.0, - graphArtifact=ArtifactRecord( - scope="datastore", - kind="integrated_graph", - artifactId="8" * 64, - ), - clusterArtifact=ArtifactRecord( - scope="datastore", - kind="cluster_labels", - artifactId="9" * 64, - ), - clusterColumn="agent_wnn_cluster", - metrics=IntegrationMetrics( - nClusters=2, - minClusterCells=20, - modalityWeightsValid=True, - ), - evidenceIds=["integration:wnn_resolution_1:clusters"], - ) - - +@pytest.mark.parametrize("missing_replication", [False, True]) def test_unsafe_experimental_context_pauses_and_explicit_skip_reuses_evidence( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, + missing_replication: bool, ) -> None: path = create_store(tmp_path / "unsafe-context.zarr") store = DataStore(str(path), default_assay="RNA", min_features_per_cell=0) - workflow = create_agent_workflow(store, workflow_run_id="unsafe-context") + workflow = WorkflowIdentity("unsafe-context") cell_selection = ArtifactReferenceModel.from_artifact_ref( store.snapshot_cell_selection("I") ) - enrichment = DataEnrichmentReport.get_example().model_copy( + enrichment = example(DataEnrichmentReport).model_copy( update={ "runInfo": AgentRunInfo( agentName="data_enrichment", @@ -332,42 +287,41 @@ def test_unsafe_experimental_context_pauses_and_explicit_skip_reuses_evidence( ) } ) - enrichment_reference = save_agent_report( - store, - workflow.workflowRunId, - enrichment, - invocation=AgentInvocation( - agentName="data_enrichment", - inputs={"studyContext": "A deliberately confounded study."}, - ), - ) request_record = OrchestrationRequestRecord( + inputIdentity={}, + modelIdentity="test-model", workflowRunId=workflow.workflowRunId, request=AutomatedWorkflowRequest( sourcePath=str(path), zarrPath=str(path), studyContext="Treatment is confounded with batch.", - allowAssumptions=True, + studyObjective="Preserve treatment while discovering populations.", ), ) - example = ExperimentalContextResult.get_example() + enrichment_reference = _save_input_evidence( + store, workflow, request_record, enrichment + ) + sample_context = example(ExperimentalContextResult) evidence_id = "batchEstimability:treatment:batch" - unsafe_plan = example.decision.batchCorrection.model_copy( + unsafe_plan = sample_context.decision.batchCorrection.model_copy( update={ "action": "unsafe", "evidenceIds": [evidence_id], } ) - unsafe_report = example.model_copy( + unsafe_report = sample_context.model_copy( update={ + "characterization": _measured_context_characterization( + store, cell_selection + ), "cellSelection": cell_selection, "qualityMetricArtifacts": [], "htoIdentityArtifacts": [], - "decision": example.decision.model_copy( + "decision": sample_context.decision.model_copy( update={"batchCorrection": unsafe_plan} ), "batchSafety": [ - BatchSafetyEvidence.get_example().model_copy( + example(BatchSafetyEvidence).model_copy( update={ "status": "unsafe", "estimability": { @@ -384,6 +338,8 @@ def test_unsafe_experimental_context_pauses_and_explicit_skip_reuses_evidence( ), } ) + if missing_replication: + unsafe_report.characterization.coefficients[0]["replication"] = {} class UnsafeAgent: calls = 0 @@ -447,7 +403,15 @@ def run(self, *_args: Any, **_kwargs: Any) -> ExperimentalContextResult: {"experimentalDirections": "skipHarmony"}, ) - assert done_outcome.status == "done" + if missing_replication: + assert done_outcome.status == "needsInput" + assert done_outcome.needsInput is not None + assert ( + "replication evidence is unavailable" + in done_outcome.needsInput.questions[0].question + ) + else: + assert done_outcome.status == "done" assert done_outcome.actions == ["resolve_unsafe_batch_correction:skip"] assert resolved_report.decision.batchCorrection.action == "skip" assert resolved_report.decision.batchCorrection.batchColumns == [] @@ -457,7 +421,345 @@ def run(self, *_args: Any, **_kwargs: Any) -> ExperimentalContextResult: assert UnsafeAgent.calls == 1 -def test_preprocessing_plan_routes_supported_modalities_and_skips_others() -> None: +@pytest.mark.parametrize("after_commit", ["none", "interrupt", "exception"]) +@pytest.mark.parametrize("report_status", ["failed", "needsInput"]) +def test_failed_context_retries_without_overwriting_or_replaying_failure( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + after_commit: str, + report_status: str, +) -> None: + path = create_store(tmp_path / "retry-context.zarr") + store = DataStore(str(path), default_assay="RNA", min_features_per_cell=0) + workflow = WorkflowIdentity("retry-context") + selection = ArtifactReferenceModel.from_artifact_ref( + store.snapshot_cell_selection("I") + ) + request = OrchestrationRequestRecord( + inputIdentity={}, + modelIdentity="test-model", + workflowRunId=workflow.workflowRunId, + config=AutomatedWorkflowConfig(inputPolicy="unattended"), + request=AutomatedWorkflowRequest( + sourcePath=str(path), + zarrPath=str(path), + studyContext="Treatment is confounded with batch.", + studyObjective="Preserve treatment while discovering populations.", + ), + ) + enrichment_ref = _save_input_evidence( + store, workflow, request, example(DataEnrichmentReport) + ) + report = example(ExperimentalContextResult).model_copy( + update={ + "characterization": _measured_context_characterization(store, selection), + "cellSelection": selection, + "cellQc": CellQcPlan(), + "qcProfiles": [], + "qualityMetricArtifacts": [], + "htoIdentityColumns": [], + "htoIdentityArtifacts": [], + } + ) + report.decision.batchCorrection.action = "unsafe" + unresolved_question = "Assess the joint tissue and treatment contrast by donor." + failed_report = report.model_copy( + update={"status": report_status, "notes": ["Invalid design proposal"]}, + deep=True, + ) + failed_report.decision.batchCorrection.action = "needsInput" + failed_report.decision.needsInput = [unresolved_question] + evidence_readers = [] + + class RecoveringAgent: + calls = 0 + + def __init__(self, *_args: Any, **_kwargs: Any) -> None: + pass + + def run(self, *_args: Any, **kwargs: Any) -> ExperimentalContextResult: + type(self).calls += 1 + read = kwargs["checkpoint_read"] + evidence_readers.append(read) + # Every new interpretation has its own result, while the complete + # measured design and its consumed allowance remain shared. + assert read("result") is None + deps = ExperimentalContextDependencies( + checkpointRead=read, + checkpointWrite=kwargs["checkpoint_write"], + ) + if self.calls == 1: + assert not restore_context_evidence(deps) + deps.characterization = report.characterization + deps.toolCalls = ["inspect_cell_covariates"] + persist_context_evidence(deps, "inspection") + deps.designRounds = 1 + deps.toolCalls.append("analyze_experimental_design") + deps.evidenceIds.add("saved-design-evidence") + persist_context_evidence(deps, "design1") + else: + assert restore_context_evidence(deps) + assert deps.characterization == report.characterization + assert deps.designRounds == 1 + assert deps.evidenceIds == {"saved-design-evidence"} + assert deps.toolCalls.count("analyze_experimental_design") == 1 + current = failed_report if self.calls <= 2 else report + kwargs["checkpoint_write"]( + "result", {"report": current.model_dump(mode="json")} + ) + assert read("result")["report"]["status"] == current.status + return current + + monkeypatch.setattr(context_module, "ExperimentalContextAgent", RecoveringAgent) + orchestrator = AgentOrchestrator(object()) + + def execute(): + return orchestrator.experimental_context_stage( + store, workflow, request, [], selection, enrichment_ref, [], [], {} + ) + + first, _ = execute() + second, _ = execute() + assert first.status == second.status == "failed" + if report_status == "needsInput": + assert unresolved_question in first.error + assert unresolved_question in second.error + assert RecoveringAgent.calls == 2 + assert first.reportReferences != second.reportReferences + assert second.inputs["retryAfterFailedReport"] == first.reportReferences[ + 0 + ].model_dump(mode="json") + original_failure = journal_module.read_stage_evidence( + store, first.reportReferences[0] + ) + + if after_commit != "none": + save_outcome = journal_module._save_outcome + + def interrupt(_group, _prefix, outcome): + if outcome.status == "done": + if after_commit == "exception": + raise RuntimeError( + "persistence failed after committing the decision" + ) + raise KeyboardInterrupt("interrupted after committing the decision") + save_outcome(_group, _prefix, outcome) + + monkeypatch.setattr(journal_module, "_save_outcome", interrupt) + if after_commit == "interrupt": + with pytest.raises(KeyboardInterrupt, match="committing"): + execute() + else: + exception_outcome, _ = execute() + assert exception_outcome.status == "failed" + assert exception_outcome.reportReferences == [] + monkeypatch.setattr(journal_module, "_save_outcome", save_outcome) + else: + done, _ = execute() + assert done.status == "done" + + resumed, resumed_report = execute() + assert resumed.status == "done" + assert resumed_report.status == "done" + assert resumed_report.decision.batchCorrection.action == "unsafe" + assert RecoveringAgent.calls == 3 + assert [read("result")["report"]["status"] for read in evidence_readers] == [ + report_status, + report_status, + "done", + ] + for key in ("inspection", "design1"): + assert all( + read(key) == evidence_readers[0](key) for read in evidence_readers[1:] + ) + if after_commit != "none": + assert "recover_persisted_experimental_context_report" in resumed.actions + assert ( + journal_module.read_stage_evidence(store, first.reportReferences[0]) + == original_failure + ) + assert ( + journal_module.read_stage_evidence(store, second.reportReferences[0])["status"] + == report_status + ) + + +def test_explicit_no_inference_skip_resolves_context_without_provider_rerun( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + path = create_store(tmp_path / "no-inference-context.zarr") + store = DataStore(str(path), default_assay="RNA", min_features_per_cell=0) + workflow = WorkflowIdentity("no-inference-context") + cell_selection = ArtifactReferenceModel.from_artifact_ref( + store.snapshot_cell_selection("I") + ) + enrichment = example(DataEnrichmentReport).model_copy( + update={ + "runInfo": AgentRunInfo( + agentName="data_enrichment", + runId=uuid.uuid4().hex, + ) + } + ) + request_record = OrchestrationRequestRecord( + inputIdentity={}, + modelIdentity="test-model", + workflowRunId=workflow.workflowRunId, + request=AutomatedWorkflowRequest( + sourcePath=str(path), + zarrPath=str(path), + studyContext="A study with unresolved replication.", + studyObjective="Discover stable RNA populations.", + ), + ) + enrichment_reference = _save_input_evidence( + store, workflow, request_record, enrichment + ) + sample_context = example(ExperimentalContextResult) + needs_input_plan = sample_context.decision.batchCorrection.model_copy( + update={ + "action": "needsInput", + "batchColumns": [], + "preserveColumns": [], + "metricsRequired": [], + } + ) + needs_input_report = sample_context.model_copy( + update={ + "characterization": _measured_context_characterization( + store, cell_selection + ), + "status": "needsInput", + "cellSelection": cell_selection, + "cellQc": CellQcPlan(), + "qcProfiles": [], + "qualityMetricArtifacts": [], + "htoIdentityColumns": [], + "htoIdentityArtifacts": [], + "decision": sample_context.decision.model_copy( + update={ + "batchCorrection": needs_input_plan, + "cellQc": CellQcPlan(), + "needsInput": [ + "Provide replicated observation units or skip inference." + ], + } + ), + "runInfo": AgentRunInfo( + agentName="experimental_context", + runId=uuid.uuid4().hex, + ), + } + ) + + class NeedsInputAgent: + calls = 0 + + def __init__(self, *_args: Any, **_kwargs: Any) -> None: + self.config = AgentRunConfig() + + def run(self, *_args: Any, **_kwargs: Any) -> ExperimentalContextResult: + type(self).calls += 1 + return needs_input_report + + monkeypatch.setattr(context_module, "ExperimentalContextAgent", NeedsInputAgent) + orchestrator = AgentOrchestrator(object()) + paused_outcome, _ = orchestrator.experimental_context_stage( + store, + workflow, + request_record, + [], + cell_selection, + enrichment_reference, + [], + [], + {}, + ) + + assert paused_outcome.status == "needsInput" + resolved_outcome, resolved_report = orchestrator.experimental_context_stage( + store, + workflow, + request_record, + [], + cell_selection, + enrichment_reference, + [], + [], + { + "experimentalDirections": { + "coefficientsOfInterest": [], + "unitsOfInference": {}, + "batchCorrection": {"action": "skip"}, + } + }, + ) + + assert resolved_outcome.status == "done" + assert resolved_outcome.artifacts == paused_outcome.artifacts + assert resolved_outcome.actions == [ + "resolve_experimental_context:no_inference_skip_harmony" + ] + assert resolved_report.status == "done" + assert resolved_report.decision.coefficientsOfInterest == [] + assert resolved_report.decision.unitsOfInference == {} + assert resolved_report.decision.batchCorrection.action == "skip" + assert resolved_report.decision.batchCorrection.batchColumns == [] + assert resolved_report.decision.needsInput == [] + assert resolved_report.runInfo.agentName == "experimental_context_resolution" + assert resolved_report.runInfo.runId == "" + assert resolved_report.runInfo.usage.requests == 0 + assert NeedsInputAgent.calls == 1 + assert journal_module.read_stage_evidence( + store, resolved_outcome.reportReferences[0] + ) == resolved_report.model_dump(mode="json") + assert paused_outcome.reportReferences[0] != resolved_outcome.reportReferences[0] + + +def test_qc_profile_safety_rejects_self_normalizing_failed_captures() -> None: + def profile( + registered_profile: RegisteredCellQcProfile, + *, + failed: list[str], + references: list[str] | None = None, + ) -> CellQcProfileEvidence: + return CellQcProfileEvidence( + profileId=f"cellQc:RNA:{registered_profile}", + action="registeredMad", + registeredProfile=registered_profile, + driverAssay="RNA", + driverAssayType="RNA", + sampleColumn="capture", + attributes=["RNA_nCounts"], + parameters={"pooledReferenceCaptures": references or []}, + activeCells=100, + retainedCells=90, + retainedFraction=0.9, + failedCaptureCandidates=failed, + evidenceId=f"qcProfile:{registered_profile}", + ) + + assert not PreprocessingStagesMixin._profile_is_safe( + profile("captureMad5", failed=["capture-b"]) + ) + assert PreprocessingStagesMixin._profile_is_safe( + profile( + "pooledReferenceMad5", + failed=["capture-b"], + references=["capture-a"], + ) + ) + assert not PreprocessingStagesMixin._profile_is_safe( + profile( + "pooledReferenceMad5", + failed=["capture-a"], + references=["capture-a"], + ) + ) + + +def test_preprocessing_plan_selects_rna_and_ignores_other_modalities() -> None: assays = { "peaks": ( "ATAC", @@ -491,25 +793,8 @@ def test_preprocessing_plan_routes_supported_modalities_and_skips_others() -> No routes["transcriptome"].featureMethod, routes["transcriptome"].reductionMethod, ) == ("graph", "hvg", "pca") - assert ( - routes["peaks"].role, - routes["peaks"].featureMethod, - routes["peaks"].reductionMethod, - routes["peaks"].reductionParameters["skipFirst"], - ) == ("graph", "prevalentPeaks", "lsi", True) - assert ( - routes["proteins"].role, - routes["proteins"].featureMethod, - routes["proteins"].reductionMethod, - ) == ("graph", "panel", "identity") - assert routes["tags"].role == "hto" - assert not routes["tags"].graphEligible - assert not routes["tags"].markerEligible - assert routes["tags"].reductionMethod == "none" - assert routes["tags"].normalizationParameters == {} - assert routes["custom"].role == "unsupported" - assert not routes["custom"].graphEligible - assert any("Unsupported assay 'custom'" in value for value in plan.limitations) + assert set(routes) == {"transcriptome"} + assert plan.pairedAssays == [] def test_converted_input_preserves_exact_selection_and_typed_qc() -> None: @@ -535,59 +820,7 @@ def test_converted_input_preserves_exact_selection_and_typed_qc() -> None: assert isinstance(plan.cellQc, CellQcPlan) -@pytest.mark.parametrize( - ("pairing_provenance", "explicit_pairing", "expected"), - [ - (None, [], []), - ("singleSourceSharedCellAxis", [], ["RNA", "ADT"]), - (None, ["RNA", "ADT"], ["RNA", "ADT"]), - ], -) -def test_multimodal_pairing_requires_persisted_or_explicit_provenance( - pairing_provenance: str | None, - explicit_pairing: list[str], - expected: list[str], -) -> None: - inputs = list( - _planning_inputs( - { - "RNA": ( - "RNA", - ["gene-1", "gene-2", "gene-3"], - ["A", "B", "C"], - ), - "ADT": ( - "ADT", - ["adt-1", "adt-2", "adt-3"], - ["CD3", "CD19", "CD45"], - ), - }, - primary_assay="RNA", - ) - ) - request_record = inputs[1] - inputs[1] = request_record.model_copy( - update={ - "request": request_record.request.model_copy( - update={"pairedAssays": explicit_pairing} - ) - } - ) - inputs[4] = inputs[4].model_copy( - update={ - "outputs": { - "format": "h5ad" if pairing_provenance else "zarr", - "pairingProvenance": pairing_provenance, - } - } - ) - - plan = AgentOrchestrator(object()).build_preprocessing_plan(*inputs) - - assert plan.pairedAssays == expected - - -def test_percent_features_follow_deterministic_inspection_not_policy_lists( +def test_percent_features_use_exact_symbols_even_when_enrichment_omits_a_family( tmp_path: Path, ) -> None: path = create_store(tmp_path / "inspected-families.zarr") @@ -603,14 +836,16 @@ def test_percent_features_follow_deterministic_inspection_not_policy_lists( store.snapshot_cell_selection("I") ) assert "RNA_percentMito" not in store.cells.columns - workflow = create_agent_workflow(store, workflow_run_id="inspected-families") + workflow = WorkflowIdentity("inspected-families") request_record = OrchestrationRequestRecord( + inputIdentity={}, + modelIdentity="test-model", workflowRunId=workflow.workflowRunId, request=AutomatedWorkflowRequest( sourcePath=str(path), zarrPath=str(path), studyContext="A deterministic feature-family test.", - allowAssumptions=True, + studyObjective="Discover stable RNA populations.", ), config=AutomatedWorkflowConfig(), ) @@ -636,7 +871,7 @@ def test_percent_features_follow_deterministic_inspection_not_policy_lists( ) orchestrator = AgentOrchestrator(object()) - outcome = orchestrator._hto_stage( + outcome = orchestrator._rna_quality_metrics_stage( store, workflow, request_record, @@ -654,17 +889,58 @@ def test_percent_features_follow_deterministic_inspection_not_policy_lists( assert metric_source.name == "RNA_percentMito" assert metric_source.artifact.kind == "quality_metric" assert outcome.artifacts[metric_source.name] == metric_source.artifact - assert orchestrator._named_stage_artifacts( + metric_sources = orchestrator._named_stage_artifacts( outcome, "qualityMetricArtifacts", "quality_metric", - ) == [metric_source] + ) + assert metric_sources[0] == metric_source + assert [item.name for item in metric_sources] == [ + "RNA_percentMito", + "RNA_percentRibo", + ] + assert outcome.outputs["percentageDefinitions"] == [ + {"family": "mitochondrial", "pattern": r"(?i)^MT-", "matchedGenes": 1}, + { + "family": "ribosomal", + "pattern": r"(?i)^(RPS|RPL|MRPS|MRPL)", + "matchedGenes": 1, + }, + ] operation = outcome.outputs["operations"][0] assert operation["operation"] == "run_feature_percentage" assert operation["cellSelection"] == cell_selection.model_dump(mode="json") assert operation["features"]["kind"] == "feature_selection" assert operation["artifact"] == metric_source.artifact.model_dump(mode="json") + from scarf.agent.cell_quality.profiles import project_auto_filter_profile + from scarf.metadata.selection import resolve_cell_aligned_artifact + + mito_values = np.asarray( + resolve_cell_aligned_artifact( + store.zw, + ArtifactRef( + scope="assay", + assay="RNA", + kind="quality_metric", + artifact_id=metric_source.artifact.artifactId, + ), + cell_selection=ArtifactRef( + scope="datastore", + kind="cell_selection", + artifact_id=cell_selection.artifactId, + ), + expected_kind="quality_metric", + ).values + ) + projection = project_auto_filter_profile( + "globalGaussian", + values_by_metric={ + "RNA_nCounts": store.cells.fetch("RNA_nCounts"), + metric_source.name: mito_values, + }, + active=np.ones(len(mito_values), dtype=bool), + ) profile = CellQcProfileEvidence( profileId="cellQc:RNA:RNA:globalGaussian:0.01:0.99", action="globalGaussian", @@ -672,10 +948,12 @@ def test_percent_features_follow_deterministic_inspection_not_policy_lists( driverAssayType="RNA", attributes=["RNA_nCounts"], artifactMetrics=[metric_source], - parameters={"minP": 0.01, "maxP": 0.99}, - activeCells=4, - retainedCells=2, - retainedFraction=0.5, + parameters=projection.parameters, + resolvedBounds=projection.parameters["resolvedBounds"], + flaggedCells=projection.flagCounts, + activeCells=len(mito_values), + retainedCells=projection.retainedCells, + retainedFraction=projection.retainedCells / len(mito_values), evidenceId="qcProfile:artifact-backed-global", ) plan = CellQcPlan( @@ -725,7 +1003,7 @@ def test_percent_features_follow_deterministic_inspection_not_policy_lists( assert "RNA_percentMito" not in store.cells.columns -def test_hto_demultiplexing_is_checkpointed_once_and_never_graph_bearing( +def test_hto_processing_is_not_executed_by_rna_workflow( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -741,14 +1019,16 @@ def test_hto_demultiplexing_is_checkpointed_once_and_never_graph_bearing( cell_selection = ArtifactReferenceModel.from_artifact_ref( store.snapshot_cell_selection("I") ) - workflow = create_agent_workflow(store, workflow_run_id="hto-once") + workflow = WorkflowIdentity("hto-once") request_record = OrchestrationRequestRecord( + inputIdentity={}, + modelIdentity="test-model", workflowRunId=workflow.workflowRunId, request=AutomatedWorkflowRequest( sourcePath=str(path), zarrPath=str(path), studyContext="A deterministic HTO checkpoint test.", - allowAssumptions=True, + studyObjective="Discover stable RNA populations.", ), config=AutomatedWorkflowConfig(), ) @@ -803,7 +1083,14 @@ def run_hto( ) orchestrator = AgentOrchestrator(object()) - first = orchestrator._hto_stage( + with pytest.raises(ValueError, match="only the selected RNA assay"): + orchestrator._rna_quality_metrics_stage( + store, workflow, request_record, [], enrichment, cell_selection + ) + enrichment = DataEnrichmentReport( + status="done", policies=[_modality_policy("RNA", "RNA")] + ) + first = orchestrator._rna_quality_metrics_stage( store, workflow, request_record, @@ -811,7 +1098,7 @@ def run_hto( enrichment, cell_selection, ) - second = orchestrator._hto_stage( + second = orchestrator._rna_quality_metrics_stage( store, workflow, request_record, @@ -821,124 +1108,17 @@ def run_hto( ) assert first == second - assert calls == 1 - assert first.outputs["operations"][0]["operation"] == "run_hto_demultiplexing" - assert first.artifacts["HTO_htoIdentity"].artifactId == identity_ref.artifact_id - assert "htoIdentityColumns" not in first.outputs - identity_source = NamedArtifactSource.model_validate( - first.outputs["htoIdentityArtifacts"][0] - ) - assert identity_source.name == "HTO_htoIdentity" - assert identity_source.artifact == first.artifacts[identity_source.name] - assert orchestrator._named_stage_artifacts( - first, - "htoIdentityArtifacts", - "hto_identity", - ) == [identity_source] - missing_source = first.model_copy( - update={ - "outputs": { - **first.outputs, - "htoIdentityArtifacts": [], - } - } - ) - with pytest.raises(ValueError, match="must name every persisted"): - orchestrator._named_stage_artifacts( - missing_source, - "htoIdentityArtifacts", - "hto_identity", - ) - assert all("graph" not in action for action in first.actions) + assert first.status == "done" + assert calls == 0 + assert first.outputs["htoIdentityArtifacts"] == [] + assert all(ref.kind != "hto_identity" for ref in first.artifacts.values()) - enrichment_reference = save_agent_report( - store, - workflow.workflowRunId, - enrichment, - invocation=AgentInvocation( - agentName="data_enrichment", - inputs={"assays": ["HTO"]}, - ), - ) - captured_sources: list[NamedArtifactSource] = [] - context_example = ExperimentalContextResult.get_example() - context_report = context_example.model_copy( - update={ - "cellSelection": cell_selection, - "cellQc": CellQcPlan(), - "qcProfiles": [], - "qualityMetricArtifacts": [], - "htoIdentityColumns": [], - "htoIdentityArtifacts": [identity_source], - "decision": context_example.decision.model_copy( - update={"cellQc": CellQcPlan()} - ), - } - ) - - class CapturingContextAgent: - calls = 0 - - def __init__(self, *_args: Any, **_kwargs: Any) -> None: - self.config = AgentRunConfig() - - def run( - self, - *_args: Any, - quality_metric_artifacts: list[NamedArtifactSource], - hto_identity_artifacts: list[NamedArtifactSource], - **_kwargs: Any, - ) -> ExperimentalContextResult: - type(self).calls += 1 - assert quality_metric_artifacts == [] - captured_sources.extend(hto_identity_artifacts) - return context_report - - monkeypatch.setattr( - context_module, - "ExperimentalContextAgent", - CapturingContextAgent, - ) - context_outcome, _ = orchestrator.experimental_context_stage( - store, - workflow, - request_record, - [journal_module._parent_link(first)], - cell_selection, - enrichment_reference, - [], - [identity_source], - {}, - ) - - assert context_outcome.status == "done" - assert captured_sources == [identity_source] - assert context_outcome.artifacts[identity_source.name] == identity_source.artifact - context_record = load_agent_record( - store, - context_outcome.reportReferences[0], - ) - assert ( - context_record.invocation.artifacts[identity_source.name] - == identity_source.artifact - ) - reused_context, reused_report = orchestrator.experimental_context_stage( - store, - workflow, - request_record, - [journal_module._parent_link(first)], - cell_selection, - enrichment_reference, - [], - [identity_source], - {}, - ) - assert reused_context == context_outcome - assert reused_report == context_report - assert CapturingContextAgent.calls == 1 +def test_selected_sample_mad_qc_passes_exact_artifact_sources( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from scarf.agent.orchestrator import preprocessing as preprocessing_module -def test_selected_sample_mad_qc_passes_exact_artifact_sources() -> None: cell_selection = ArtifactReferenceModel( scope="datastore", kind="cell_selection", @@ -970,7 +1150,8 @@ def test_selected_sample_mad_qc_passes_exact_artifact_sources() -> None: sampleArtifact=sample, attributes=["RNA_nCounts"], artifactMetrics=[metric], - parameters={"nMads": 4.0, "minCellsPerSample": 12}, + parameters={"nMads": 3.0, "minCellsPerSample": 20}, + resolvedBounds={"capture-a": {"RNA_nCounts": {"low": 2.0, "high": 200.0}}}, activeCells=100, retainedCells=95, retainedFraction=0.95, @@ -1002,15 +1183,17 @@ def test_selected_sample_mad_qc_passes_exact_artifact_sources() -> None: ) captured: dict[str, Any] = {} - class Store: - def auto_filter_cells(self, **kwargs: Any) -> ArtifactRef: - captured.update(kwargs) - return output_selection + def execute(_store: Any, action: str, **kwargs: Any): + assert action == "sampleMad" + captured.update(kwargs) + return output_selection, None + + monkeypatch.setattr(preprocessing_module, "execute_auto_cell_qc", execute) actions: list[str] = [] operations: list[dict[str, Any]] = [] result = AgentOrchestrator(object()).apply_cell_qc( - Store(), + object(), experimental, ArtifactRef( scope="datastore", @@ -1032,1044 +1215,48 @@ def auto_filter_cells(self, **kwargs: Any) -> ArtifactRef: assert captured["sample_artifact"].artifact.artifact_id == ( sample.artifact.artifactId ) - assert captured["n_mads"] == 4.0 - assert captured["min_cells_per_sample"] == 12 + assert captured["profile_parameters"] == profile.parameters + assert captured["expected_resolved_bounds"] == profile.resolvedBounds + assert captured["expected_active_cells"] == 100 + assert captured["expected_retained_cells"] == 95 assert operations[0]["sampleArtifact"] == sample.model_dump(mode="json") assert operations[0]["artifactMetrics"] == [metric.model_dump(mode="json")] -@pytest.mark.parametrize( - ("assay_types", "explicit", "expected"), - [ - (["ATAC", "ADT", "RNA"], None, "RNA"), - (["ATAC", "ADT"], None, "ADT"), - (["ATAC"], None, "ATAC"), - (["RNA", "ADT", "ATAC"], "ATAC", "ATAC"), - ], -) -def test_marker_assay_precedence( - assay_types: list[str], - explicit: str | None, - expected: str, -) -> None: - assays = { - assay_type: ( - assay_type, - [f"{assay_type}-1", f"{assay_type}-2", f"{assay_type}-3"], - [f"{assay_type}-1", f"{assay_type}-2", f"{assay_type}-3"], - ) - for assay_type in assay_types - } - - plan = _build_plan( - assays, - primary_assay=assay_types[0], - marker_assay=explicit, - ) - - assert plan.primaryAssay == assay_types[0] - assert plan.markerAssay == expected - - -def test_adt_identity_limit_and_exact_observed_control_exclusion() -> None: - assays = { - "ADT": ( - "ADT", - ["adt-1", "control-id", "adt-2", "adt-3"], - [ - "CD3", - "Mouse IgG1 isotype control", - "control response protein", - "CD19", - ], - ) - } - controls = { - "ADT": [ - FeatureReference( - featureId="control-id", - featureName="Mouse IgG1 isotype control", - ) - ] - } - - identity = _build_plan( - assays, - controls=controls, - exclude_features={"ADT": ["adt-1"]}, - artificial_features={"ADT": ["adt-2"]}, - config=AutomatedWorkflowConfig(maxIdentityFeatures=3), - ).assays[0] - pca = _build_plan( - assays, - controls=controls, - exclude_features={"ADT": ["adt-1"]}, - artificial_features={"ADT": ["adt-2"]}, - config=AutomatedWorkflowConfig(maxIdentityFeatures=2), - ).assays[0] - - assert identity.exactExcludedFeatures == [ - "control-id", - "Mouse IgG1 isotype control", - ] - assert identity.reductionMethod == "identity" - assert identity.reductionParameters["dimensions"] == 3 - assert pca.reductionMethod == "pca" - assert pca.reductionParameters["dimensions"] == 2 - - -def test_atac_invalid_coordinates_are_limited_without_changing_lsi_route() -> None: - assays = { - "ATAC": ( - "ATAC", - ["chr1:1-20", "not-a-coordinate", "chr2:10-30"], - ["peak-1", "peak-2", "peak-3"], - ) - } - - plan = _build_plan(assays, peak_statuses={"ATAC": "invalid"}) - route = plan.assays[0] - - assert route.graphEligible - assert route.featureMethod == "prevalentPeaks" - assert route.reductionMethod == "lsi" - assert route.reductionParameters == {"dimensions": 50, "skipFirst": True} - assert route.limitations == [ - "ATAC feature coordinates are not uniformly valid chrom:start-end " - "intervals; the genome build remains unknown" - ] - - -@pytest.mark.parametrize( - ("method", "n_cells", "n_features", "expected_dimensions"), - [ - ("pca", 5, 3, {2}), - ("lsi", 6, 4, {3}), - ("identity", 4, 2, {2}), - ], -) -def test_initial_candidates_are_rank_valid( - method: str, - n_cells: int, - n_features: int, - expected_dimensions: set[int], -) -> None: - orchestrator = AgentOrchestrator(object()) - handoff = PreprocessedAssayHandoff( - assay="assay", - assayType="ADT" if method == "identity" else method.upper(), - reductionMethod=method, - normalized=ArtifactReferenceModel( - assay="assay", - kind="normalized", - artifactId="1" * 64, - ), - nCells=n_cells, - nFeatures=n_features, - ) - - candidates = orchestrator.initial_parameter_candidates( - "rank-test", - handoff, - count=5, - neighbors_k=n_cells - 1, - ) - - assert len(candidates) == 5 - assert {value.dimensions for value in candidates} == expected_dimensions - assert all(2 <= value.dimensions for value in candidates) - if method != "identity": - assert all(value.dimensions < min(n_cells, n_features) for value in candidates) - assert all(2 <= value.neighborsK < n_cells for value in candidates) - - -def test_initial_candidates_reject_fully_invalid_rank_or_neighbor_count() -> None: - orchestrator = AgentOrchestrator(object()) - rank_invalid = PreprocessedAssayHandoff( - assay="RNA", - assayType="RNA", - reductionMethod="pca", - normalized=ArtifactReferenceModel.get_example(), - nCells=4, - nFeatures=2, - ) - identity_invalid = rank_invalid.model_copy( - update={"assayType": "ADT", "reductionMethod": "identity", "nFeatures": 1} - ) - - with pytest.raises(ValueError, match="no rank-valid graph candidate"): - orchestrator.initial_parameter_candidates( - "rank-test", - rank_invalid, - count=3, - neighbors_k=3, - ) - with pytest.raises(ValueError, match="no rank-valid graph candidate"): - orchestrator.initial_parameter_candidates( - "rank-test", - identity_invalid, - count=3, - neighbors_k=3, - ) - with pytest.raises(ValueError, match="no rank-valid graph candidate"): - orchestrator.initial_parameter_candidates( - "rank-test", - rank_invalid.model_copy(update={"nFeatures": 3}), - count=3, - neighbors_k=4, - ) - - -def test_parameter_tuning_rejects_plan_above_global_branch_cap( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - path = create_store(tmp_path / "branch-cap.zarr") - store = DataStore(str(path), default_assay="RNA", min_features_per_cell=0) - workflow = create_agent_workflow(store, workflow_run_id="branch-cap") - config = AutomatedWorkflowConfig( - primaryInitialCandidates=5, - maxRefinedCandidatesPerAssay=1, - maxHarmonyCandidatesPerAssay=0, - maxCandidateBranches=5, - ) - request_record = OrchestrationRequestRecord( - workflowRunId=workflow.workflowRunId, - request=AutomatedWorkflowRequest( - sourcePath=str(path), - zarrPath=str(path), - studyContext="A branch-cap test.", - allowAssumptions=True, - ), - config=config, - ) - plan = AutomatedPreprocessingPlan( - primaryAssay="RNA", - markerAssay="RNA", - assays=[AssayPreprocessingPlan.get_example()], - ) - handoff = PreprocessedAssayHandoff.get_example() - experimental = ExperimentalContextResult.get_example() - decision = experimental.decision.model_copy( - update={"batchCorrection": BatchCorrectionPlan(action="skip")} - ) - experimental = experimental.model_copy( - update={"decision": decision, "batchSafety": []} - ) - - class UnusedAgent: - def run_batch(self, *_args: Any, **_kwargs: Any) -> None: - pytest.fail("The tuning agent must not run above the global branch cap") - - monkeypatch.setattr( - tuning_module, - "ParameterTuningAgent", - lambda *_args, **_kwargs: UnusedAgent(), - ) - outcome, report = AgentOrchestrator(object()).parameter_tuning_stage( - store, - workflow, - request_record, - [], - plan, - [handoff], - experimental, - AgentReportReference.get_example(), - AgentReportReference.get_example(), - {}, - ) - - assert outcome.status == "failed" - assert outcome.error is not None - assert "exceeds the global branch limit 5" in outcome.error - assert report.status == "failed" - assert load_agent_workflow(store, workflow.workflowRunId).status == "failed" - - -def test_final_selection_pause_exposes_exact_options_and_resumes_without_screen( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - path = create_store(tmp_path / "selection-resume.zarr") - store = DataStore(str(path), default_assay="RNA", min_features_per_cell=0) - workflow = create_agent_workflow(store, workflow_run_id="selection-resume") - enrichment = DataEnrichmentReport.get_example().model_copy( - update={ - "runInfo": AgentRunInfo( - agentName="data_enrichment", - runId=uuid.uuid4().hex, - ) - } - ) - experimental = ExperimentalContextResult.get_example().model_copy( - update={ - "runInfo": AgentRunInfo( - agentName="experimental_context", - runId=uuid.uuid4().hex, - ) - } - ) - enrichment_reference = save_agent_report( - store, - workflow.workflowRunId, - enrichment, - invocation=AgentInvocation( - agentName="data_enrichment", - inputs={"studyContext": "A final-selection resume test."}, - ), - ) - experimental_reference = save_agent_report( - store, - workflow.workflowRunId, - experimental, - invocation=AgentInvocation( - agentName="experimental_context", - inputs={"cellSelection": _cell_selection_model().model_dump(mode="json")}, - artifacts={"cellSelection": _cell_selection_model()}, - ), - ) - request_record = OrchestrationRequestRecord( - workflowRunId=workflow.workflowRunId, - request=AutomatedWorkflowRequest( - sourcePath=str(path), - zarrPath=str(path), - studyContext="A final-selection resume test.", - allowAssumptions=True, - ), - ) - plan = AutomatedPreprocessingPlan( - primaryAssay="RNA", - markerAssay="RNA", - assays=[ - AssayPreprocessingPlan.get_example(), - AssayPreprocessingPlan( - assay="ADT", - assayType="ADT", - role="graph", - graphEligible=True, - markerEligible=True, - featureMethod="panel", - reductionMethod="identity", - ), - ], - pairedAssays=["RNA", "ADT"], - ) - preprocessed = [ - PreprocessedAssayHandoff.get_example(), - PreprocessedAssayHandoff( - assay="ADT", - assayType="ADT", - cellSelection=_cell_selection_model(), - reductionMethod="identity", - normalized=ArtifactReferenceModel( - assay="ADT", - kind="normalized", - artifactId="2" * 64, - ), - nCells=100, - nFeatures=5, - ), - ] - integration = _eligible_integration() - calls = {"screen": 0, "promote": 0, "integrate": 0, "selection": 0} - - def selection_execution(**_kwargs: Any) -> Any: - return SimpleNamespace( - output=FinalGraphSelection( - status="needsInput", - rationale="The supplied evidence does not resolve the final graph.", - needsInput=FinalGraphNeedsInput( - question="An unconstrained provider question.", - options=["invented-option"], - ), - ), - runInfo=AgentRunInfo( - agentName="parameter_tuning_final_graph", - runId=uuid.uuid4().hex, - ), - ) - - monkeypatch.setattr(parameter_tuning_module, "run_agent_sync", selection_execution) - - class CountingAgent: - config = AgentRunConfig() - - def run_batch(self, *_args: Any, **_kwargs: Any) -> ParameterTuningReport: - calls["screen"] += 1 - return _native_batch_report("RNA", "ADT") - - def promote(self, *_args: Any, **_kwargs: Any) -> None: - calls["promote"] += 1 - - def select_final( - self, - *, - report: ParameterTuningReport, - integration_evaluations: list[IntegrationCandidateEvaluation], - marker_assay: str, - ) -> ParameterTuningReport: - calls["selection"] += 1 - return select_final_parameter_graph( - model=object(), - report=report, - integration_evaluations=integration_evaluations, - marker_assay=marker_assay, - config=self.config, - ) - - monkeypatch.setattr( - tuning_module, - "ParameterTuningAgent", - lambda *_args, **_kwargs: CountingAgent(), - ) - orchestrator = AgentOrchestrator(object()) - - def evaluate_integrations(*_args: Any, **_kwargs: Any) -> list[Any]: - calls["integrate"] += 1 - return [integration] - - monkeypatch.setattr(orchestrator, "evaluate_integrations", evaluate_integrations) - real_validated_outcome = journal_module._validated_done_outcome - paused_attempts: list[WorkflowStageAttempt] = [] - - def validated_outcome( - target_store: Any, - prefix: str, - workflow_run_id: str, - stage: str, - record: OrchestrationRequestRecord, - parents: list[WorkflowStageLink], - *, - required_status: str = "done", - ) -> WorkflowStageAttempt | None: - if stage == "parameter_tuning": - if required_status == "needsInput" and paused_attempts: - return paused_attempts[-1] - return None - return real_validated_outcome( - target_store, - prefix, - workflow_run_id, - stage, - record, - parents, - required_status=required_status, - ) - - monkeypatch.setattr( - journal_module, - "_validated_done_outcome", - validated_outcome, - ) - paused, paused_report = orchestrator.parameter_tuning_stage( - store, - workflow, - request_record, - [], - plan, - preprocessed, - experimental, - enrichment_reference, - experimental_reference, - {}, - ) - paused_attempts.append(paused) - expected_options = sorted(final_graph_options(paused_report, [integration])) - - assert paused.status == "needsInput" - assert paused.needsInput is not None - assert paused.needsInput.questions[0].questionId == "finalGraphOptionId" - assert paused.needsInput.questions[0].options == expected_options - assert paused_report.needsInput is not None - assert paused_report.needsInput.options == expected_options - assert paused_report.finalSelection is not None - assert paused_report.finalSelection.needsInput is not None - assert paused_report.finalSelection.needsInput.options == expected_options - assert paused_report.totalCandidates == 3 - - completed, completed_report = orchestrator.parameter_tuning_stage( - store, - workflow, - request_record, - [], - plan, - preprocessed, - experimental, - enrichment_reference, - experimental_reference, - {"finalGraphOptionId": "native:RNA:baseline"}, - ) - - assert completed.status == "done" - assert completed_report.status == "done" - assert completed_report.finalSelection is not None - assert completed_report.finalSelection.selectedOptionId == "native:RNA:baseline" - assert completed_report.totalCandidates == 3 - assert calls == {"screen": 1, "promote": 2, "integrate": 1, "selection": 1} - - -def test_integration_evaluations_contribute_to_final_candidate_count() -> None: - report = _native_batch_report("RNA", "ADT") - integration = _eligible_integration() - - finalized = finalize_parameter_tuning_selection( - report, - marker_assay="RNA", - integration_evaluations=[integration], - native_assay="RNA", - ) - - assert report.totalCandidates == 2 - assert finalized.totalCandidates == 3 - - -def test_long_assay_names_produce_bounded_unique_candidate_ids() -> None: - orchestrator = AgentOrchestrator(object()) - prefix = "Very long assay name with punctuation / and spaces " + "x" * 120 - first = PreprocessedAssayHandoff( - assay=f"{prefix} one", - assayType="RNA", - reductionMethod="pca", - normalized=ArtifactReferenceModel.get_example(), - nCells=100, - nFeatures=50, - ) - second = first.model_copy(update={"assay": f"{prefix} two"}) - - first_ids = { - candidate.candidateId - for candidate in orchestrator.initial_parameter_candidates( - "long-name-test", - first, - count=5, - neighbors_k=11, - ) - } - second_ids = { - candidate.candidateId - for candidate in orchestrator.initial_parameter_candidates( - "long-name-test", - second, - count=5, - neighbors_k=11, - ) - } - - assert first_ids.isdisjoint(second_ids) - assert all(len(candidate_id) <= 64 for candidate_id in first_ids | second_ids) - assert all( - all( - character.isdigit() or character.islower() or character in {"_", "-"} - for character in candidate_id - ) - for candidate_id in first_ids | second_ids - ) - assert all( - len(f"{candidate_id}_harmony") <= 64 for candidate_id in first_ids | second_ids - ) - - -def test_single_integration_resolution_is_centered_and_workflow_unique() -> None: - report = _native_batch_report("RNA", "ADT") - primary_report = report.assayReports["RNA"] - primary_evaluation = primary_report.evaluations[0] - centered_evaluation = primary_evaluation.model_copy( - update={ - "parameters": primary_evaluation.parameters.model_copy( - update={"leidenResolution": 1.25} - ) - } - ) - primary_report = primary_report.model_copy( - update={"evaluations": [centered_evaluation]} - ) - report = report.model_copy( - update={ - "evaluations": [centered_evaluation], - "assayReports": { - **report.assayReports, - "RNA": primary_report, - }, - } - ) - - class IntegrationStore: - def __init__(self) -> None: - self.integration_calls: list[dict[str, Any]] = [] - self.cluster_calls: list[dict[str, Any]] = [] - - def load_artifact(self, reference: ArtifactRef) -> dict[str, np.ndarray]: - if reference.kind == "integrated_graph": - return {"modality_weights": np.full((4, 2), 0.5)} - return {"values": np.asarray([0, 0, 1, 1])} - - def integrate_assays( - self, - sources: list[ArtifactRef], - **kwargs: Any, - ) -> ArtifactRef: - self.integration_calls.append({"sources": sources, **kwargs}) - token = len(self.integration_calls) - return ArtifactRef( - scope="datastore", - kind="integrated_graph", - artifact_id=f"{100 + token:064x}", - ) - - def run_leiden_clustering( - self, - graph: ArtifactRef, - **kwargs: Any, - ) -> ArtifactRef: - self.cluster_calls.append({"graph": graph, **kwargs}) - token = len(self.cluster_calls) - return ArtifactRef( - scope="datastore", - kind="cluster_labels", - artifact_id=f"{200 + token:064x}", - ) - - store = IntegrationStore() - plan = AutomatedPreprocessingPlan( - primaryAssay="RNA", - markerAssay="RNA", - pairedAssays=["RNA", "ADT"], - ) - evaluations = AgentOrchestrator(object()).evaluate_integrations( - store, - "integration-center", - plan, - report, - ExperimentalTuningHandoff( - cellSelection=_cell_selection_model(), - batchAction="skip", - ), - AutomatedWorkflowConfig( - integrationResolutionCandidates=1, - minClusterCells=1, - ), - ) - - assert len(evaluations) == 2 - assert {value.method for value in evaluations} == {"snn", "wnn"} - assert {value.resolution for value in evaluations} == {1.25} - assert len(store.integration_calls) == 2 - assert all(call["invalidate_cache"] is True for call in store.integration_calls) - assert {call["method"] for call in store.integration_calls} == {"snn", "wnn"} - for call in store.integration_calls: - expected_kind = "connectivity_map" if call["method"] == "snn" else "neighbors" - assert {source.kind for source in call["sources"]} == {expected_kind} - assert "label" not in call - assert len(store.cluster_calls) == 2 - assert {call["resolution"] for call in store.cluster_calls} == {1.25} - - -def test_integration_requires_trusted_label_connectivity() -> None: - class IntegrationStore: - def load_artifact(self, reference: ArtifactRef) -> dict[str, np.ndarray]: - if reference.kind == "integrated_graph": - return {"modality_weights": np.full((4, 2), 0.5)} - return {"values": np.asarray([0, 0, 1, 1])} - - def integrate_assays( - self, - _sources: list[ArtifactRef], - **_kwargs: Any, - ) -> ArtifactRef: - return ArtifactRef( - scope="datastore", - kind="integrated_graph", - artifact_id="7" * 64, - ) - - def run_leiden_clustering( - self, - _graph: ArtifactRef, - **_kwargs: Any, - ) -> ArtifactRef: - return ArtifactRef( - scope="datastore", - kind="cluster_labels", - artifact_id="8" * 64, - ) - - def metric_graph_connectivity(self, *_args: Any, **_kwargs: Any) -> float: - raise ValueError("trusted label is unavailable") - - evaluations = AgentOrchestrator(object()).evaluate_integrations( - IntegrationStore(), - "integration-connectivity", - AutomatedPreprocessingPlan( - primaryAssay="RNA", - markerAssay="RNA", - pairedAssays=["RNA", "ADT"], - ), - _native_batch_report("RNA", "ADT"), - ExperimentalTuningHandoff( - cellSelection=_cell_selection_model(), - batchAction="skip", - preservationColumns=["trusted_cell_type"], - ), - AutomatedWorkflowConfig( - integrationResolutionCandidates=1, - minClusterCells=1, - ), - ) - - assert evaluations - assert all(not value.eligible for value in evaluations) - assert all( - any( - "trusted-label connectivity" in reason - for reason in value.eligibilityReasons - ) - for value in evaluations - ) - - -def test_integration_checkpoints_prevent_retry_execution( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - path = create_store(tmp_path / "integration-checkpoint.zarr") +def test_capture_cell_selections_are_exact_and_idempotent(tmp_path: Path) -> None: + path = create_store(tmp_path / "capture-selections.zarr") store = DataStore( str(path), default_assay="RNA", - min_features_per_cell=-1, - mito_pattern="", - ribo_pattern="", + min_features_per_cell=0, zarr_mode="r+", ) - workflow = create_agent_workflow(store, workflow_run_id="integration-checkpoint") - request_record = OrchestrationRequestRecord( - workflowRunId=workflow.workflowRunId, - request=AutomatedWorkflowRequest( - sourcePath=str(path), - zarrPath=str(path), - studyContext="A deterministic integration retry test.", - allowAssumptions=True, - ), - config=AutomatedWorkflowConfig(), - ) - prefix = journal_module._ensure_orchestration_store(store) - started = journal_module._start_attempt( - store.zw, - prefix, - workflow.workflowRunId, - "parameter_tuning", - request_record, - [], - inputs={"semanticInput": "stable"}, - ) - integration_calls: list[dict[str, Any]] = [] - cluster_calls: list[dict[str, Any]] = [] - - def load_artifact(reference: ArtifactRef) -> dict[str, np.ndarray]: - if reference.kind == "integrated_graph": - return {"modality_weights": np.full((4, 2), 0.5)} - return {"values": np.asarray([0, 0, 1, 1])} - - def integrate_assays( - sources: list[ArtifactRef], - **kwargs: Any, - ) -> ArtifactRef: - integration_calls.append({"sources": sources, **kwargs}) - return ArtifactRef( - scope="datastore", - kind="integrated_graph", - artifact_id=f"{100 + len(integration_calls):064x}", - ) - - def cluster(graph: ArtifactRef, **kwargs: Any) -> ArtifactRef: - cluster_calls.append({"graph": graph, **kwargs}) - return ArtifactRef( - scope="datastore", - kind="cluster_labels", - artifact_id=f"{200 + len(cluster_calls):064x}", - ) - - monkeypatch.setattr(store, "load_artifact", load_artifact) - monkeypatch.setattr(store, "integrate_assays", integrate_assays) - monkeypatch.setattr(store, "run_leiden_clustering", cluster) - plan = AutomatedPreprocessingPlan( - primaryAssay="RNA", - markerAssay="RNA", - pairedAssays=["RNA", "ADT"], - ) - report = _native_batch_report("RNA", "ADT") - config = AutomatedWorkflowConfig( - integrationResolutionCandidates=1, - minClusterCells=1, - ) - first_actions: list[str] = [] - orchestrator = AgentOrchestrator(object()) + capture_values = np.asarray(["a", "a", "b", "b"]) + store.cells.insert("capture", capture_values, overwrite=True) + parent = store.snapshot_cell_selection("I") + active_indices = np.arange(store.cells.N, dtype=np.int64) - first = orchestrator.evaluate_integrations( + first_a, count_a = _select_capture_cells( store, - workflow.workflowRunId, - plan, - report, - ExperimentalTuningHandoff( - cellSelection=_cell_selection_model(), - batchAction="skip", - ), - config, - started=started, - actions=first_actions, - ) - retried = started.model_copy( - update={"attemptId": "retry-attempt", "startedAtNs": started.startedAtNs + 1} + parent, + column="capture", + value="a", + active_indices=active_indices, + active_values=capture_values, ) - retry_actions: list[str] = [] - second = orchestrator.evaluate_integrations( + second_a, repeated_count = _select_capture_cells( store, - workflow.workflowRunId, - plan, - report, - ExperimentalTuningHandoff( - cellSelection=_cell_selection_model(), - batchAction="skip", - ), - config, - started=retried, - actions=retry_actions, - ) - - assert second == first - assert len(integration_calls) == 2 - assert len(cluster_calls) == 2 - assert first_actions == [ - "checkpoint_integration:snn", - "checkpoint_integration:wnn", - ] - assert retry_actions == [ - "recover_integration_checkpoint:snn", - "recover_integration_checkpoint:wnn", - ] - checkpoint_reports = list_agent_reports( - store, - workflow.workflowRunId, - agent_name="parameter_tuning", + parent, + column="capture", + value="a", + active_indices=active_indices, + active_values=capture_values, ) - assert len(checkpoint_reports) == 2 - assert all("_integration_" in value.agentRunId for value in checkpoint_reports) - - -def test_preprocess_atac_and_adt_feature_routes() -> None: - class Features: - def __init__(self, ids: list[str], names: list[str]) -> None: - self.N = len(ids) - self.values = {"ids": np.asarray(ids), "names": np.asarray(names)} - - def fetch_all(self, column: str) -> np.ndarray: - return self.values[column] - - class Store: - def __init__(self) -> None: - self.assays = { - "ATAC": SimpleNamespace( - feats=Features( - ["chr1:1-10", "chr1:20-30", "chr2:1-10", "chr2:20-30"], - ["peak1", "peak2", "peak3", "peak4"], - ) - ), - "ADT": SimpleNamespace( - feats=Features( - ["adt1", "adt2", "adt3", "adt4"], - ["CD3", "control", "CD19", "CD45"], - ) - ), - } - self.masks: list[np.ndarray] = [] - - def get_assay(self, assay: str) -> Any: - return self.assays[assay] - - @staticmethod - def select_prevalent_peaks(*_args: Any, **_kwargs: Any) -> ArtifactRef: - return ArtifactRef( - scope="assay", - assay="ATAC", - kind="feature_selection", - artifact_id="1" * 64, - ) - - def set_feature_selection( - self, *, from_assay: str, mask: np.ndarray, **_kwargs: Any - ) -> ArtifactRef: - self.masks.append(mask.copy()) - return ArtifactRef( - scope="assay", - assay=from_assay, - kind="feature_selection", - artifact_id=("2" if from_assay == "ADT" else "3") * 64, - ) - - @staticmethod - def run_normalization( - _selection: ArtifactRef, *, features: ArtifactRef, **_kwargs: Any - ) -> ArtifactRef: - return ArtifactRef( - scope="assay", - assay=features.assay, - kind="normalized", - artifact_id=("4" if features.assay == "ATAC" else "5") * 64, - ) - - def load_artifact(self, ref: ArtifactRef) -> dict[str, np.ndarray]: - selected = 3 if ref.assay == "ATAC" else int(self.masks[-1].sum()) - return {"values": np.asarray([True] * selected + [False] * (4 - selected))} - - store = Store() - selection = ArtifactRef( - scope="datastore", - kind="cell_selection", - artifact_id="c" * 64, - ) - selection_model = ArtifactReferenceModel.from_artifact_ref(selection) - orchestrator = AgentOrchestrator(object()) - actions: list[str] = [] - operations: list[dict[str, Any]] = [] - artifacts: dict[str, ArtifactReferenceModel] = {} - atac = AssayPreprocessingPlan( - assay="ATAC", - assayType="ATAC", - role="graph", - graphEligible=True, - markerEligible=True, - featureMethod="prevalentPeaks", - reductionMethod="lsi", - featureParameters={"topN": 10, "minCells": 1}, - normalizationParameters={"logTransform": False, "renormalizeSubset": False}, - ) - atac_handoff = orchestrator.preprocess_assay( - store, - atac, - cell_selection=selection, - cell_selection_model=selection_model, - active_cells=12, - actions=actions, - operations=operations, - artifacts=artifacts, - ) - assert atac_handoff.nFeatures == 3 - assert operations[0]["topN"] == 3 - - adt = AssayPreprocessingPlan( - assay="ADT", - assayType="ADT", - role="graph", - graphEligible=True, - markerEligible=True, - featureMethod="panel", - reductionMethod="identity", - exactExcludedFeatures=["control", "adt4"], - normalizationParameters={"logTransform": True, "renormalizeSubset": True}, - ) - adt_handoff = orchestrator.preprocess_assay( - store, - adt, - cell_selection=selection, - cell_selection_model=selection_model, - active_cells=12, - actions=actions, - operations=operations, - artifacts=artifacts, - ) - assert adt_handoff.nFeatures == 2 - np.testing.assert_array_equal(store.masks[-1], [True, False, True, False]) - - with pytest.raises(ValueError, match="fewer than two non-control"): - orchestrator.preprocess_assay( - store, - adt.model_copy(update={"exactExcludedFeatures": ["adt2", "adt3", "adt4"]}), - cell_selection=selection, - cell_selection_model=selection_model, - active_cells=12, - actions=[], - operations=[], - artifacts={}, - ) - with pytest.raises(ValueError, match="Unsupported feature route"): - orchestrator.preprocess_assay( - store, - adt.model_copy(update={"featureMethod": "none"}), - cell_selection=selection, - cell_selection_model=selection_model, - active_cells=12, - actions=[], - operations=[], - artifacts={}, - ) - - -def test_exact_feature_exclusion_covers_all_supported_families() -> None: - class Features: - N = 4 - - @staticmethod - def fetch_all(column: str) -> np.ndarray: - values = { - "ids": np.asarray(["MT-CO1", "RPS3", "HIST1H1", "gene4"]), - "names": np.asarray(["mito", "ribo", "histone", "GENE4"]), - } - return values[column] - - class Store: - def __init__(self) -> None: - self.mask: np.ndarray | None = None + selected_a = np.asarray(store.load_artifact(first_a)["values"][:], dtype=bool) - @staticmethod - def get_assay(_assay: str) -> Any: - return SimpleNamespace(feats=Features()) - - @staticmethod - def load_artifact(_source: ArtifactRef) -> dict[str, np.ndarray]: - return {"values": np.ones(4, dtype=bool)} - - def set_feature_selection( - self, *, from_assay: str, mask: np.ndarray, **_kwargs: Any - ) -> ArtifactRef: - self.mask = mask.copy() - return ArtifactRef( - scope="assay", - assay=from_assay, - kind="feature_selection", - artifact_id="6" * 64, - ) - - plan = AssayPreprocessingPlan( - assay="RNA", - assayType="RNA", - role="graph", - graphEligible=True, - markerEligible=True, - featureMethod="hvg", - reductionMethod="pca", - featureParameters={ - "excludeFamilies": ["mitochondrial", "ribosomal", "histone"] - }, - ) - orchestrator = AgentOrchestrator(object()) - blacklist = orchestrator.rna_blacklist(plan) - assert all(token in blacklist for token in ("MT-", "RPS", "HIST")) - source = ArtifactRef( - scope="assay", - assay="RNA", - kind="feature_selection", - artifact_id="7" * 64, - ) - store = Store() - result = orchestrator.exclude_exact_features(store, plan, source) - assert result.kind == "feature_selection" - assert store.mask is not None - np.testing.assert_array_equal(store.mask, [False, False, False, True]) - with pytest.raises(ValueError, match="removed every feature"): - orchestrator.exclude_exact_features( - store, - plan.model_copy(update={"exactExcludedFeatures": ["gene4"]}), - source, - ) + assert first_a == second_a + assert count_a == repeated_count == 2 + assert selected_a.tolist() == [True, True, False, False] def test_cell_qc_artifact_and_execution_validation_edges() -> None: @@ -2237,7 +1424,7 @@ def report( invalid_global_profile = global_profile.model_copy( update={"sampleColumn": "sample"} ) - with pytest.raises(ValueError, match="cannot include a sample"): + with pytest.raises(ValueError, match="cannot use a core sample source"): orchestrator.apply_cell_qc( object(), report(invalid_global_plan, invalid_global_profile, quality=[metric]), @@ -2267,75 +1454,231 @@ def report( ) -def test_preprocessing_plan_rejects_invalid_assay_routing() -> None: - orchestrator = AgentOrchestrator(object()) - unsupported = _planning_inputs( - {"custom": ("CRISPR", ["guide"], ["guide"])}, - ) - with pytest.raises(ValueError, match="No supported graph-bearing"): - orchestrator.build_preprocessing_plan(*unsupported) +def test_harmony_doublet_graph_matches_selected_native_parameters() -> None: + class Store: + def __init__(self) -> None: + self.neighbors_k: int | None = None + self.resolution: float | None = None - too_many = list( - _planning_inputs( - { - "RNA": ("RNA", ["g1", "g2", "g3"], ["G1", "G2", "G3"]), - "ADT": ("ADT", ["a1", "a2", "a3"], ["A1", "A2", "A3"]), + def load_artifact(self, reference: ArtifactRef) -> dict[str, Any]: + assert reference.kind == "reduction" + return {} + + def build_ann_index( + self, coordinates: ArtifactRef, **_kwargs: Any + ) -> ArtifactRef: + assert coordinates.kind == "reduction" + return ArtifactRef( + scope="assay", + assay="RNA", + kind="ann_index", + artifact_id="2" * 64, + ) + + def query_neighbors( + self, + _ann: ArtifactRef, + *, + k: int, + **_kwargs: Any, + ) -> ArtifactRef: + self.neighbors_k = k + return ArtifactRef( + scope="assay", + assay="RNA", + kind="neighbors", + artifact_id="3" * 64, + ) + + def build_connectivity_map( + self, + _neighbors: ArtifactRef, + **_kwargs: Any, + ) -> ArtifactRef: + return ArtifactRef( + scope="assay", + assay="RNA", + kind="connectivity_map", + artifact_id="4" * 64, + ) + + def run_leiden_clustering( + self, + _graph: ArtifactRef, + *, + resolution: float, + **_kwargs: Any, + ) -> ArtifactRef: + self.resolution = resolution + return ArtifactRef( + scope="assay", + assay="RNA", + kind="cluster_labels", + artifact_id="5" * 64, + ) + + base = example(ParameterCandidateEvaluation) + native = base.model_copy( + update={ + "candidateId": "native", + "parameters": base.parameters.model_copy( + update={ + "candidateId": "native", + "dimensions": 20, + "neighborsK": 11, + "leidenResolution": 0.5, + "useHarmony": False, + } + ), + } + ) + selected = base.model_copy( + update={ + "candidateId": "harmony", + "parameters": base.parameters.model_copy( + update={ + "candidateId": "harmony", + "dimensions": 20, + "neighborsK": 41, + "leidenResolution": 1.5, + "useHarmony": True, + } + ), + "artifacts": { + **base.artifacts, + "pca": ArtifactRecord( + assay="RNA", + kind="reduction", + artifactId="1" * 64, + ), }, - config=AutomatedWorkflowConfig(maxGraphAssays=1), - ) + } + ) + store = Store() + clusters, graph = resolve_native_doublet_inputs( + store, + selected, + [native, selected], ) - with pytest.raises(ValueError, match="Too many graph-bearing"): - orchestrator.build_preprocessing_plan(*too_many) - duplicate = list( - _planning_inputs( - { - "RNA1": ("RNA", ["g1", "g2", "g3"], ["G1", "G2", "G3"]), - "RNA2": ("RNA", ["g4", "g5", "g6"], ["G4", "G5", "G6"]), - } - ) + assert store.neighbors_k == 41 + assert store.resolution == 1.5 + assert graph.artifact_id == "4" * 64 + assert clusters.artifact_id == "5" * 64 + + +@pytest.mark.parametrize( + ("batch_mixing", "marker_coherence", "expected"), + [ + (0.56, 0.80, True), + (0.54, 0.80, False), + (0.56, 0.74, False), + (None, 0.80, False), + ], +) +def test_harmony_acceptance_requires_improvement_without_biological_loss( + batch_mixing: float | None, + marker_coherence: float, + expected: bool, +) -> None: + base = example(ParameterCandidateEvaluation) + native_parameters = base.parameters.model_copy( + update={"candidateId": "native", "useHarmony": False} ) - request_record = duplicate[1] - duplicate[1] = request_record.model_copy( + harmony_parameters = base.parameters.model_copy( + update={"candidateId": "harmony", "useHarmony": True} + ) + native = base.model_copy( update={ - "request": request_record.request.model_copy(update={"analysisAssays": []}) + "candidateId": "native", + "parameters": native_parameters, + "metrics": base.metrics.model_copy( + update={ + "batchMixing": {"batch": 0.50}, + "biologicalPreservation": { + "condition": { + "clisi": 0.80, + "graphConnectivity": 0.80, + } + }, + "crossUnitSupport": 0.80, + "markerCoherence": 0.80, + } + ), + } + ) + harmony = base.model_copy( + update={ + "candidateId": "harmony", + "parameters": harmony_parameters, + "metrics": base.metrics.model_copy( + update={ + "batchMixing": ( + {} if batch_mixing is None else {"batch": batch_mixing} + ), + "biologicalPreservation": { + "condition": { + "clisi": 0.80, + "graphConnectivity": 0.80, + } + }, + "crossUnitSupport": 0.80, + "markerCoherence": marker_coherence, + } + ), } ) - with pytest.raises(ValueError, match="same-kind biological assays"): - orchestrator.build_preprocessing_plan(*duplicate) - - base = list( - _planning_inputs({"RNA": ("RNA", ["g1", "g2", "g3"], ["G1", "G2", "G3"])}) - ) - request_record = base[1] - for updates, message in ( - ({"primaryAssay": "missing"}, "primaryAssay"), - ({"markerAssay": "missing"}, "markerAssay"), - ({"pairedAssays": ["RNA", "missing"]}, "non-graph assays"), - ): - values = list(base) - values[1] = request_record.model_copy( - update={"request": request_record.request.model_copy(update=updates)} - ) - with pytest.raises(ValueError, match=message): - orchestrator.build_preprocessing_plan(*values) - paired = list( - _planning_inputs( + accepted, reasons = harmony_acceptance_gate( + native, + harmony, + batch_columns=["batch"], + protected_columns=["condition"], + independent_unit_columns=["donor"], + ) + + assert accepted is expected + assert bool(reasons) is not expected + + +@pytest.mark.parametrize( + "assays,updates,message", + [ + ({"custom": ("CRISPR", ["g"], ["G"])}, {}, "requires one RNA"), + ( { - "RNA": ("RNA", ["g1", "g2", "g3"], ["G1", "G2", "G3"]), - "ADT": ("ADT", ["a1", "a2", "a3"], ["A1", "A2", "A3"]), + "RNA1": ("RNA", ["a", "b", "c"], ["A", "B", "C"]), + "RNA2": ("RNA", ["d", "e", "f"], ["D", "E", "F"]), }, - primary_assay="RNA", - ) - ) - request_record = paired[1] - paired[1] = request_record.model_copy( - update={ - "request": request_record.request.model_copy( - update={"pairedAssays": ["ADT"]} - ) - } + {}, + "found 2", + ), + ( + {"RNA": ("RNA", ["a", "b", "c"], ["A", "B", "C"])}, + {"primaryAssay": "missing"}, + "Unknown requested RNA", + ), + ( + {"RNA": ("RNA", ["a", "b", "c"], ["A", "B", "C"])}, + {"markerAssay": "missing"}, + "markerAssay", + ), + ( + {"RNA": ("RNA", ["a", "b", "c"], ["A", "B", "C"])}, + {"pairedAssays": ["RNA", "other"]}, + "pairedAssays", + ), + ], +) +def test_preprocessing_plan_rejects_invalid_assay_routing( + assays: dict[str, Any], + updates: dict[str, Any], + message: str, +) -> None: + inputs = list(_planning_inputs(assays)) + record = inputs[1] + inputs[1] = record.model_copy( + update={"request": record.request.model_copy(update=updates)} ) - with pytest.raises(ValueError, match="must include the primary"): - orchestrator.build_preprocessing_plan(*paired) + with pytest.raises(ValueError, match=message): + AgentOrchestrator(object()).build_preprocessing_plan(*inputs) diff --git a/tests/test_agent_parameter_tuning.py b/tests/test_agent_parameter_tuning.py index 37edf941..10c3fa40 100644 --- a/tests/test_agent_parameter_tuning.py +++ b/tests/test_agent_parameter_tuning.py @@ -1,5 +1,7 @@ """Tests for bounded parameter tuning agent execution.""" +from tests.agent_examples import example + import asyncio from types import SimpleNamespace from typing import Any @@ -15,7 +17,9 @@ ) from pydantic_ai.models.function import AgentInfo, FunctionModel -import scarf.agent.parameter_tuning as parameter_tuning_module +import scarf.agent.parameter_tuning.agent as parameter_tuning_agent +import scarf.agent.parameter_tuning.execution as parameter_tuning_execution +import scarf.agent.parameter_tuning.selection as parameter_tuning_selection from scarf.agent.parameter_tuning import ( ArtifactRecord, CandidateComparison, @@ -52,6 +56,8 @@ validate_parameter_search_plan, validate_parameter_tuning_report, ) +from scarf.agent.parameter_tuning.prompts import parameter_evaluation_payload +from scarf.agent.parameter_tuning.selection import pending_parameter_tuning_report from scarf.agent.types import ( AgentDataModel, ArtifactReferenceModel, @@ -332,7 +338,7 @@ def _dependencies( max_candidates: int = 5, min_cluster_cells: int = 20, ) -> ParameterTuningDependencies: - candidate_values = candidates or [ParameterCandidate.get_example()] + candidate_values = candidates or [example(ParameterCandidate)] return ParameterTuningDependencies( store=store, normalized=store.normalized, @@ -383,7 +389,7 @@ def test_parameter_models_have_blank_and_example( model_type: type[AgentDataModel], ) -> None: assert isinstance(model_type.get_blank(), model_type) - assert isinstance(model_type.get_example(), model_type) + assert isinstance(example(model_type), model_type) assert all("_" not in field for field in model_type.model_fields) @@ -431,7 +437,7 @@ def test_harmony_pairing_covers_every_initial_candidate() -> None: def test_prompts_include_only_explicit_candidate_context() -> None: - evaluation = ParameterCandidateEvaluation.get_example() + evaluation = example(ParameterCandidateEvaluation) plan = ParameterSearchPlan(status="complete") cell_selection = ArtifactReferenceModel.from_artifact_ref(_cell_selection()) planning_system_prompt = parameter_search_system_prompt() @@ -712,7 +718,7 @@ def test_duplicate_candidate_returns_recorded_execution_without_rerun() -> None: def test_candidate_budget_prevents_another_execution() -> None: store = _FakeStore() candidates = [ - ParameterCandidate.get_example(), + example(ParameterCandidate), ParameterCandidate(candidateId="pca_15", dimensions=15), ] deps = _dependencies(store, candidates=candidates, max_candidates=1) @@ -729,7 +735,7 @@ def test_candidate_budget_prevents_another_execution() -> None: def test_refinement_plan_is_bounded_by_initial_evidence_and_envelope() -> None: store = _FakeStore() candidates = [ - ParameterCandidate.get_example(), + example(ParameterCandidate), ParameterCandidate(candidateId="pca_15", dimensions=15), ] deps = _dependencies(store, candidates=candidates, max_candidates=3) @@ -766,7 +772,7 @@ def test_refinement_plan_is_bounded_by_initial_evidence_and_envelope() -> None: def test_refinement_plan_canonicalizes_status_from_candidate_presence() -> None: store = _FakeStore() candidates = [ - ParameterCandidate.get_example(), + example(ParameterCandidate), ParameterCandidate(candidateId="pca_15", dimensions=15), ] deps = _dependencies(store, candidates=candidates, max_candidates=3) @@ -812,7 +818,7 @@ def test_refinement_plan_canonicalizes_status_from_candidate_presence() -> None: def test_refinement_plan_requires_authorized_matched_harmony_evidence() -> None: store = _FakeStore() candidates = [ - ParameterCandidate.get_example(), + example(ParameterCandidate), ParameterCandidate(candidateId="pca_15", dimensions=15), ] deps = _dependencies(store, candidates=candidates, max_candidates=3) @@ -913,7 +919,7 @@ def test_report_validation_uses_only_executed_results_and_artifacts() -> None: status="done", evaluations=[ParameterCandidateEvaluation.get_blank()], recommendedCandidateId="baseline", - selectedArtifacts={"invented": ArtifactRecord.get_example()}, + selectedArtifacts={"invented": example(ArtifactRecord)}, confidence="medium", rationale="Balanced candidate.", evidenceIds=[evaluation.evidenceIds[0]], @@ -988,7 +994,7 @@ def test_report_validation_rejects_unknown_evidence_and_ineligible_choice() -> N def test_done_report_requires_a_successful_comparator_when_available() -> None: store = _FakeStore() candidates = [ - ParameterCandidate.get_example(), + example(ParameterCandidate), ParameterCandidate(candidateId="pca_15", dimensions=15), ] deps = _dependencies(store, candidates=candidates, max_candidates=2) @@ -1029,7 +1035,7 @@ def test_done_report_requires_a_successful_comparator_when_available() -> None: def test_comparison_requires_selected_and_comparator_evidence() -> None: store = _FakeStore() candidates = [ - ParameterCandidate.get_example(), + example(ParameterCandidate), ParameterCandidate(candidateId="pca_15", dimensions=15), ] deps = _dependencies(store, candidates=candidates, max_candidates=2) @@ -1065,7 +1071,7 @@ def test_comparison_requires_selected_and_comparator_evidence() -> None: def test_comparison_rejects_unknown_or_duplicate_candidate_ids() -> None: store = _FakeStore() candidates = [ - ParameterCandidate.get_example(), + example(ParameterCandidate), ParameterCandidate(candidateId="pca_15", dimensions=15), ] deps = _dependencies(store, candidates=candidates, max_candidates=2) @@ -1207,7 +1213,7 @@ def test_biology_handoff_requires_selected_cluster_artifact() -> None: def test_integrated_final_selection_separates_graph_and_marker_assays() -> None: - native = ParameterTuningReport.get_example() + native = example(ParameterTuningReport) aggregate = ParameterTuningReport( status="done", fromAssay="RNA", @@ -1275,7 +1281,7 @@ def test_integrated_final_selection_requires_marker_assay() -> None: def test_final_graph_selector_uses_one_grounded_provider_request() -> None: - native_evaluation = ParameterCandidateEvaluation.get_example() + native_evaluation = example(ParameterCandidateEvaluation) native_evaluation.artifacts["clusters"] = ArtifactRecord( assay="RNA", kind="cluster_labels", @@ -1397,7 +1403,7 @@ async def reply( def test_parameter_tuning_agent_delegates_with_its_model_and_config( monkeypatch: pytest.MonkeyPatch, ) -> None: - from scarf.agent import parameter_tuning as module + from scarf.agent.parameter_tuning import agent as module model = object() expected = ParameterTuningReport.get_blank() @@ -1457,7 +1463,7 @@ async def reply( result = ParameterTuningAgent(FunctionModel(reply)).run( _FakeStore(), normalized=_artifact("normalized", 1), - candidates=[ParameterCandidate.get_example()], + candidates=[example(ParameterCandidate)], experimental_handoff=ExperimentalTuningHandoff( cellSelection=ArtifactReferenceModel.from_artifact_ref(_cell_selection()), batchAction="skip", @@ -1722,12 +1728,12 @@ async def reply( assays=[ ParameterTuningAssayInput( normalized=_artifact("normalized", 1, "RNA"), - candidates=[ParameterCandidate.get_example()], + candidates=[example(ParameterCandidate)], maxCandidates=1, ), ParameterTuningAssayInput( normalized=_artifact("normalized", 10, "ADT"), - candidates=[ParameterCandidate.get_example()], + candidates=[example(ParameterCandidate)], maxCandidates=1, ), ], @@ -1790,7 +1796,7 @@ async def reply( assays=[ ParameterTuningAssayInput( normalized=_artifact("normalized", 1), - candidates=[ParameterCandidate.get_example()], + candidates=[example(ParameterCandidate)], maxCandidates=2, maxRefinedCandidates=1, ) @@ -1801,12 +1807,20 @@ async def reply( assert model_calls == 3 assert result.status == "done" assert result.recommendedByAssay == {"RNA": "baseline"} + assert result.searchPlan is not None + assert result.searchPlan.runInfo.agentName == "parameter_batch_search_planning" + assert result.searchPlan.runInfo.usage.requests == 2 + assay_plan = result.assayReports["RNA"].searchPlan + assert assay_plan is not None + assert assay_plan.runInfo == result.searchPlan.runInfo + assert result.runInfo.agentName == "parameter_tuning_batch" + assert result.runInfo.usage.requests == 1 -def test_batched_tuning_falls_back_after_structured_output_exhaustion( +def test_batched_tuning_pauses_after_structured_output_exhaustion( monkeypatch: pytest.MonkeyPatch, ) -> None: - from scarf.agent import parameter_tuning as module + from scarf.agent.parameter_tuning import agent as module calls: list[str] = [] @@ -1822,7 +1836,7 @@ def unavailable_structured_output(**kwargs: Any) -> None: ParameterTuningAssayInput( normalized=_artifact("normalized", 1), candidates=[ - ParameterCandidate.get_example(), + example(ParameterCandidate), ParameterCandidate(candidateId="pca_15", dimensions=15), ], maxCandidates=3, @@ -1832,20 +1846,22 @@ def unavailable_structured_output(**kwargs: Any) -> None: primary_assay="RNA", ) - assert calls == ["parameter_batch_search_planning", "parameter_tuning_batch"] - assert result.status == "done" - assert result.recommendedByAssay == {"RNA": "baseline"} + assert calls == ["parameter_batch_search_planning"] + assert result.status == "needsInput" + assert result.recommendedByAssay == {} assert result.assayReports["RNA"].confidence == "low" - assert result.assayReports["RNA"].comparisons[0].candidateId == "pca_15" + assert result.assayReports["RNA"].recommendedCandidateId is None + assert result.assayReports["RNA"].needsInput is not None + assert result.assayReports["RNA"].needsInput.options == ["baseline", "pca_15"] assert result.searchPlan is not None assert result.searchPlan.status == "complete" - assert result.runInfo.agentName == "parameter_tuning_batch_fallback" + assert result.runInfo.agentName == "parameter_batch_search_planning_needs_input" -def test_single_tuning_falls_back_after_structured_output_exhaustion( +def test_single_tuning_pauses_after_structured_output_exhaustion( monkeypatch: pytest.MonkeyPatch, ) -> None: - from scarf.agent import parameter_tuning as module + from scarf.agent.parameter_tuning import agent as module def unavailable_structured_output(**_kwargs: Any) -> None: raise UnexpectedModelBehavior("structured output unavailable") @@ -1856,21 +1872,73 @@ def unavailable_structured_output(**_kwargs: Any) -> None: model=object(), normalized=_artifact("normalized", 1), candidates=[ - ParameterCandidate.get_example(), + example(ParameterCandidate), ParameterCandidate(candidateId="pca_15", dimensions=15), ], max_candidates=3, max_refined_candidates=1, ) - assert result.status == "done" - assert result.recommendedCandidateId == "baseline" + assert result.status == "needsInput" + assert result.recommendedCandidateId is None assert result.confidence == "low" - assert result.runInfo.agentName == "parameter_tuning_fallback" + assert result.needsInput is not None + assert result.needsInput.options == ["baseline", "pca_15"] + assert result.runInfo.agentName == "parameter_search_planning_needs_input" + + +def test_pending_parameter_report_does_not_select_without_successful_baseline() -> None: + candidates = [ + example(ParameterCandidate), + ParameterCandidate(candidateId="pca_15", dimensions=15), + ParameterCandidate(candidateId="pca_30", dimensions=30), + ] + deps = _dependencies( + _FakeStore(), + candidates=candidates, + max_candidates=3, + ) + for candidate in candidates: + execute_parameter_candidate(deps, candidate.candidateId) + deps.evaluations["baseline"] = deps.evaluations["baseline"].model_copy( + update={ + "status": "failed", + "eligible": False, + "error": "baseline failed", + } + ) + + result = pending_parameter_tuning_report( + deps, + search_plan=ParameterSearchPlan(status="complete"), + agent_name="parameter_tuning_needs_input", + ) + + assert result.status == "needsInput" + assert result.recommendedCandidateId is None + assert result.needsInput is not None + assert result.needsInput.options == ["pca_15", "pca_30"] + + +def test_parameter_prompt_payload_is_bounded_and_excludes_artifacts() -> None: + evaluation = example(ParameterCandidateEvaluation).model_copy( + update={ + "warnings": ["w" * 700 for _index in range(12)], + "error": "e" * 700, + } + ) + + payload = parameter_evaluation_payload(evaluation) + + assert "artifacts" not in payload + assert "cellSelection" not in payload + assert len(payload["warnings"]) == 10 + assert all(len(warning) == 500 for warning in payload["warnings"]) + assert len(payload["error"]) == 500 def test_single_eligible_final_graph_skips_provider_selection() -> None: - evaluation = ParameterCandidateEvaluation.get_example() + evaluation = example(ParameterCandidateEvaluation) evaluation.artifacts["clusters"] = ArtifactRecord( assay="RNA", kind="cluster_labels", @@ -1910,11 +1978,11 @@ def test_single_eligible_final_graph_skips_provider_selection() -> None: def test_final_graph_retry_exhaustion_pauses_when_multiple_options_exist( monkeypatch: pytest.MonkeyPatch, ) -> None: - from scarf.agent import parameter_tuning as module + from scarf.agent.parameter_tuning import selection as module reports: dict[str, ParameterTuningReport] = {} for assay, token in (("RNA", "7"), ("ADT", "8")): - evaluation = ParameterCandidateEvaluation.get_example().model_copy( + evaluation = example(ParameterCandidateEvaluation).model_copy( update={ "clusterColumn": f"{assay}_agent_tuning_baseline", "artifacts": { @@ -1966,7 +2034,7 @@ def unavailable_structured_output(**_kwargs: Any) -> None: ] assert selected.finalSelection is not None assert selected.finalSelection.runInfo.agentName == ( - "parameter_tuning_final_graph_fallback" + "parameter_tuning_final_graph_needs_input" ) @@ -1990,7 +2058,7 @@ def test_batched_tuning_enforces_global_candidate_limit_before_execution() -> No assays = [ ParameterTuningAssayInput( normalized=_artifact("normalized", 1, assay), - candidates=[ParameterCandidate.get_example()], + candidates=[example(ParameterCandidate)], maxCandidates=1, ) for assay in ("RNA", "ADT") @@ -2008,7 +2076,7 @@ def test_batched_tuning_enforces_global_candidate_limit_before_execution() -> No def test_parameter_tuning_handoff_validation_edges() -> None: - report = ParameterTuningReport.get_example() + report = example(ParameterTuningReport) with pytest.raises(ValueError, match="must be done"): report.model_copy(update={"status": "failed"}).to_biological_handoff() with pytest.raises(ValueError, match="must recommend"): @@ -2076,9 +2144,9 @@ def test_parameter_tuning_handoff_validation_edges() -> None: def test_final_graph_option_filtering_and_validation_edges() -> None: - report = ParameterTuningReport.get_example() + report = example(ParameterTuningReport) with pytest.raises(ValueError, match="lacks an exact cell selection"): - parameter_tuning_module.final_graph_options( + parameter_tuning_selection.final_graph_options( report.model_copy(update={"cellSelection": None}), [], ) @@ -2090,7 +2158,7 @@ def test_final_graph_option_filtering_and_validation_edges() -> None: "assayReports": {}, } ) - assert parameter_tuning_module.final_graph_options(skipped_report, []) == {} + assert parameter_tuning_selection.final_graph_options(skipped_report, []) == {} ineligible_report = report.model_copy( update={ "evaluations": [ @@ -2099,7 +2167,7 @@ def test_final_graph_option_filtering_and_validation_edges() -> None: "assayReports": {}, } ) - assert parameter_tuning_module.final_graph_options(ineligible_report, []) == {} + assert parameter_tuning_selection.final_graph_options(ineligible_report, []) == {} mismatched_native = report.evaluations[0].model_copy( update={ "cellSelection": ArtifactReferenceModel( @@ -2110,14 +2178,14 @@ def test_final_graph_option_filtering_and_validation_edges() -> None: } ) with pytest.raises(ValueError, match="Native graph option uses a different"): - parameter_tuning_module.final_graph_options( + parameter_tuning_selection.final_graph_options( report.model_copy( update={"evaluations": [mismatched_native], "assayReports": {}} ), [], ) - integration = IntegrationCandidateEvaluation.get_example() + integration = example(IntegrationCandidateEvaluation) for ignored in ( integration.model_copy(update={"status": "failed"}), integration.model_copy(update={"graphArtifact": None}), @@ -2131,11 +2199,11 @@ def test_final_graph_option_filtering_and_validation_edges() -> None: } ), ): - options = parameter_tuning_module.final_graph_options(report, [ignored]) + options = parameter_tuning_selection.final_graph_options(report, [ignored]) assert all(not key.startswith("integration:") for key in options) with pytest.raises(ValueError, match="Integrated graph option uses a different"): - parameter_tuning_module.final_graph_options( + parameter_tuning_selection.final_graph_options( report, [ integration.model_copy( @@ -2150,12 +2218,12 @@ def test_final_graph_option_filtering_and_validation_edges() -> None: ], ) with pytest.raises(ValueError, match="require integrationId"): - parameter_tuning_module.final_graph_options( + parameter_tuning_selection.final_graph_options( report, [integration.model_copy(update={"integrationId": ""})], ) with pytest.raises(ValueError, match="must be datastore-scoped"): - parameter_tuning_module.final_graph_options( + parameter_tuning_selection.final_graph_options( report, [ integration.model_copy( @@ -2171,7 +2239,9 @@ def test_final_graph_option_filtering_and_validation_edges() -> None: ], ) with pytest.raises(ValueError, match="Duplicate integration id"): - parameter_tuning_module.final_graph_options(report, [integration, integration]) + parameter_tuning_selection.final_graph_options( + report, [integration, integration] + ) def test_normalized_shape_and_candidate_metric_failure_edges() -> None: @@ -2183,14 +2253,14 @@ def load_artifact(self, _ref: Any) -> dict[str, Any]: return self.payload with pytest.raises(ValueError, match="does not contain"): - parameter_tuning_module.normalized_artifact_shape(ShapeStore({}), object()) + parameter_tuning_execution.normalized_artifact_shape(ShapeStore({}), object()) with pytest.raises(ValueError, match="two-dimensional"): - parameter_tuning_module.normalized_artifact_shape( + parameter_tuning_execution.normalized_artifact_shape( ShapeStore({"data": SimpleNamespace(shape=(4,))}), object(), ) with pytest.raises(ValueError, match="at least two cells"): - parameter_tuning_module.normalized_artifact_shape( + parameter_tuning_execution.normalized_artifact_shape( ShapeStore({"data": SimpleNamespace(shape=(1, 4))}), object(), ) @@ -2230,7 +2300,7 @@ def metric_graph_connectivity(self, *_args: Any, **_kwargs: Any) -> float: assert invalid.status == "failed" assert message in (invalid.error or "") - harmony = ParameterCandidate.get_example().model_copy( + harmony = example(ParameterCandidate).model_copy( update={"candidateId": "baseline_harmony", "useHarmony": True} ) harmony_deps = _dependencies(_FakeStore(), candidates=[harmony]) @@ -2242,7 +2312,7 @@ def metric_graph_connectivity(self, *_args: Any, **_kwargs: Any) -> float: def test_parameter_search_plan_validation_edges() -> None: candidates = [ - ParameterCandidate.get_example(), + example(ParameterCandidate), ParameterCandidate(candidateId="pca_15", dimensions=15), ] deps = _dependencies(_FakeStore(), candidates=candidates, max_candidates=4) @@ -2469,8 +2539,8 @@ def test_parameter_tuning_report_validation_status_edges() -> None: def test_final_graph_selection_validation_edges() -> None: - report = ParameterTuningReport.get_example() - integration = IntegrationCandidateEvaluation.get_example() + report = example(ParameterTuningReport) + integration = example(IntegrationCandidateEvaluation) native_evidence = "native:RNA:candidate:baseline:clusters" integration_evidence = integration.evidenceIds[0] comparison = FinalGraphComparison( @@ -2485,7 +2555,7 @@ def test_final_graph_selection_validation_edges() -> None: evidenceIds=[native_evidence], comparisons=[comparison], ) - validated = parameter_tuning_module.validate_final_graph_selection( + validated = parameter_tuning_selection.validate_final_graph_selection( valid, report, integration_evaluations=[integration], @@ -2494,7 +2564,7 @@ def test_final_graph_selection_validation_edges() -> None: assert validated.graphMethod == "native" with pytest.raises(ValueError, match="must finish"): - parameter_tuning_module.validate_final_graph_selection( + parameter_tuning_selection.validate_final_graph_selection( valid, report.model_copy(update={"status": "failed"}), integration_evaluations=[integration], @@ -2504,63 +2574,63 @@ def test_final_graph_selection_validation_edges() -> None: update={"evaluations": [], "recommendedCandidateId": None, "assayReports": {}} ) with pytest.raises(ValueError, match="No eligible"): - parameter_tuning_module.validate_final_graph_selection( + parameter_tuning_selection.validate_final_graph_selection( valid, no_options, integration_evaluations=[], marker_assay="RNA", ) with pytest.raises(ValueError, match="unknown evidence"): - parameter_tuning_module.validate_final_graph_selection( + parameter_tuning_selection.validate_final_graph_selection( valid.model_copy(update={"evidenceIds": ["unknown"]}), report, integration_evaluations=[integration], marker_assay="RNA", ) with pytest.raises(ValueError, match="concrete question"): - parameter_tuning_module.validate_final_graph_selection( + parameter_tuning_selection.validate_final_graph_selection( FinalGraphSelection(status="needsInput"), report, integration_evaluations=[integration], marker_assay="RNA", ) with pytest.raises(ValueError, match="done or needsInput"): - parameter_tuning_module.validate_final_graph_selection( + parameter_tuning_selection.validate_final_graph_selection( valid.model_copy(update={"status": "failed"}), report, integration_evaluations=[integration], marker_assay="RNA", ) with pytest.raises(ValueError, match="not eligible"): - parameter_tuning_module.validate_final_graph_selection( + parameter_tuning_selection.validate_final_graph_selection( valid.model_copy(update={"selectedOptionId": "missing"}), report, integration_evaluations=[integration], marker_assay="RNA", ) with pytest.raises(ValueError, match="cite selected-option"): - parameter_tuning_module.validate_final_graph_selection( + parameter_tuning_selection.validate_final_graph_selection( valid.model_copy(update={"evidenceIds": []}), report, integration_evaluations=[integration], marker_assay="RNA", ) with pytest.raises(ValueError, match="must not contain duplicates"): - parameter_tuning_module.validate_final_graph_selection( + parameter_tuning_selection.validate_final_graph_selection( valid.model_copy(update={"comparisons": [comparison, comparison]}), report, integration_evaluations=[integration], marker_assay="RNA", ) with pytest.raises(ValueError, match="one comparison for every"): - parameter_tuning_module.validate_final_graph_selection( + parameter_tuning_selection.validate_final_graph_selection( valid.model_copy(update={"comparisons": []}), report, integration_evaluations=[integration], marker_assay="RNA", ) with pytest.raises(ValueError, match="cite selected-option evidence"): - parameter_tuning_module.validate_final_graph_selection( + parameter_tuning_selection.validate_final_graph_selection( valid.model_copy( update={ "comparisons": [ @@ -2575,7 +2645,7 @@ def test_final_graph_selection_validation_edges() -> None: marker_assay="RNA", ) with pytest.raises(ValueError, match="cite comparator evidence"): - parameter_tuning_module.validate_final_graph_selection( + parameter_tuning_selection.validate_final_graph_selection( valid.model_copy( update={ "comparisons": [ @@ -2588,7 +2658,7 @@ def test_final_graph_selection_validation_edges() -> None: marker_assay="RNA", ) with pytest.raises(ValueError, match="requires a summary"): - parameter_tuning_module.validate_final_graph_selection( + parameter_tuning_selection.validate_final_graph_selection( valid.model_copy( update={"comparisons": [comparison.model_copy(update={"summary": ""})]} ), @@ -2617,7 +2687,7 @@ def test_experimental_tuning_handoff_resolution_edges() -> None: batchSafety=[safety], evidenceIds=[safety.evidenceId], ) - resolved = parameter_tuning_module._resolve_experimental_tuning_handoff( + resolved = parameter_tuning_agent._resolve_experimental_tuning_handoff( normalized_cell_selection=selection, batch_columns=[], preservation_columns=[], @@ -2650,21 +2720,21 @@ def test_experimental_tuning_handoff_resolution_edges() -> None: ] for updates, message in changes: with pytest.raises(ValueError, match=message): - parameter_tuning_module._resolve_experimental_tuning_handoff( + parameter_tuning_agent._resolve_experimental_tuning_handoff( normalized_cell_selection=selection, batch_columns=[], preservation_columns=[], experimental_handoff=handoff.model_copy(update=updates), ) with pytest.raises(ValueError, match="batch_columns conflict"): - parameter_tuning_module._resolve_experimental_tuning_handoff( + parameter_tuning_agent._resolve_experimental_tuning_handoff( normalized_cell_selection=selection, batch_columns=["other"], preservation_columns=[], experimental_handoff=handoff, ) with pytest.raises(ValueError, match="preservation_columns conflict"): - parameter_tuning_module._resolve_experimental_tuning_handoff( + parameter_tuning_agent._resolve_experimental_tuning_handoff( normalized_cell_selection=selection, batch_columns=[], preservation_columns=["other"], @@ -2675,7 +2745,7 @@ def test_experimental_tuning_handoff_resolution_edges() -> None: def test_prepare_parameter_tuning_dependencies_validation_edges() -> None: store = _FakeStore() normalized = store.normalized - candidate = ParameterCandidate.get_example() + candidate = example(ParameterCandidate) for kwargs, message in ( ({"max_candidates": 0}, "max_candidates"), @@ -2684,18 +2754,18 @@ def test_prepare_parameter_tuning_dependencies_validation_edges() -> None: ({"identity_feature_limit": 1}, "identity_feature_limit"), ): with pytest.raises(ValueError, match=message): - parameter_tuning_module.prepare_parameter_tuning_dependencies( + parameter_tuning_agent.prepare_parameter_tuning_dependencies( store, normalized=normalized, **kwargs, ) with pytest.raises(TypeError, match="normalized ArtifactRef"): - parameter_tuning_module.prepare_parameter_tuning_dependencies( + parameter_tuning_agent.prepare_parameter_tuning_dependencies( store, normalized=_artifact("reduction", 10), ) with pytest.raises(ValueError, match="has no assay"): - parameter_tuning_module.prepare_parameter_tuning_dependencies( + parameter_tuning_agent.prepare_parameter_tuning_dependencies( store, normalized=ArtifactRef( scope="datastore", @@ -2735,25 +2805,25 @@ def inspect_artifact(self, _normalized: ArtifactRef) -> Any: for status, message in statuses: with pytest.raises(ValueError, match=message): invalid_store = StatusStore(status) - parameter_tuning_module.prepare_parameter_tuning_dependencies( + parameter_tuning_agent.prepare_parameter_tuning_dependencies( invalid_store, normalized=invalid_store.normalized, ) with pytest.raises(ValueError, match="batch_columns must be unique"): - parameter_tuning_module.prepare_parameter_tuning_dependencies( + parameter_tuning_agent.prepare_parameter_tuning_dependencies( store, normalized=normalized, batch_columns=["batch", "batch"], ) with pytest.raises(ValueError, match="candidates must be non-empty"): - parameter_tuning_module.prepare_parameter_tuning_dependencies( + parameter_tuning_agent.prepare_parameter_tuning_dependencies( store, normalized=normalized, candidates=[], ) with pytest.raises(ValueError, match="exceeds max_candidates"): - parameter_tuning_module.prepare_parameter_tuning_dependencies( + parameter_tuning_agent.prepare_parameter_tuning_dependencies( store, normalized=normalized, candidates=[ @@ -2763,13 +2833,13 @@ def inspect_artifact(self, _normalized: ArtifactRef) -> Any: max_candidates=1, ) with pytest.raises(ValueError, match="only ASCII"): - parameter_tuning_module.prepare_parameter_tuning_dependencies( + parameter_tuning_agent.prepare_parameter_tuning_dependencies( store, normalized=normalized, candidates=[candidate.model_copy(update={"candidateId": "bad-id"})], ) with pytest.raises(ValueError, match="Duplicate candidateId"): - parameter_tuning_module.prepare_parameter_tuning_dependencies( + parameter_tuning_agent.prepare_parameter_tuning_dependencies( store, normalized=normalized, candidates=[candidate, candidate], diff --git a/tests/test_agent_plot_boundaries.py b/tests/test_agent_plot_boundaries.py new file mode 100644 index 00000000..878340ba --- /dev/null +++ b/tests/test_agent_plot_boundaries.py @@ -0,0 +1,139 @@ +"""Final maps fail visibly on invalid artifacts and retain dense population legends.""" + +from typing import Any + +import matplotlib.pyplot as plt +import numpy as np +import pytest + +from scarf.agent import _plots +from scarf.storage.refs import ArtifactRef +from tests.test_agent_analysis_plots import display_store + + +@pytest.mark.parametrize("limit", [False, 0, 50_001, 1.5]) +def test_invalid_display_limit_fails_before_artifact_reads( + monkeypatch: pytest.MonkeyPatch, limit: Any +) -> None: + store, refs, arrays = display_store(monkeypatch, n=12) + with pytest.raises(ValueError, match="max_points"): + _plots.plot_final_umap(store, **refs, max_points=limit, show=False) + assert not any(array.read_sizes for array in arrays.values()) + + +@pytest.mark.parametrize( + "damage", + [ + "incomplete", + "missingSelection", + "wrongAssay", + "wrongKind", + "nonNumeric", + "empty", + "graphSelection", + ], +) +def test_final_map_validates_all_artifact_boundaries( + monkeypatch: pytest.MonkeyPatch, damage: str +) -> None: + store, refs, arrays = display_store(monkeypatch, n=0 if damage == "empty" else 12) + inspect = store.inspect_artifact + + def changed(ref: ArtifactRef) -> Any: + status = inspect(ref) + if damage == "incomplete": + status.complete = False + elif damage == "missingSelection": + status.inputs.pop("cell_selection") + return status + + store.inspect_artifact = changed + if damage == "wrongAssay": + refs["umap"] = ArtifactRef("assay", "embedding", "4" * 64, "RNA") + elif damage == "wrongKind": + refs["umap"] = refs["graph"] + elif damage == "nonNumeric": + arrays["umap"].dtype = np.dtype("U8") + elif damage == "graphSelection": + monkeypatch.setattr( + _plots, + "graph_cell_selection", + lambda *_: ArtifactRef("datastore", "cell_selection", "f" * 64), + ) + with pytest.raises((ValueError, TypeError)): + _plots.plot_final_umap(store, **refs, show=False) + assert not any(array.read_sizes for array in arrays.values()) + + +def test_many_clusters_are_annotated_and_plot_failure_closes_figure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store, refs, arrays = display_store(monkeypatch, n=82) + arrays["clusters"] = np.repeat(np.arange(41), 2) + plot = _plots.plot_final_umap(store, **refs, show=False) + try: + assert len(plot.axes["clusters"].texts) == 41 + assert len(plot.scales[0].palette) == 41 + assert plot.tables["cluster_counts"]["cells"].sum() == 82 + finally: + plot.close() + before = set(plt.get_fignums()) + from matplotlib.axes import Axes + + def fail(*args: Any, **kwargs: Any) -> Any: + raise RuntimeError("rendering interrupted") + + monkeypatch.setattr(Axes, "scatter", fail) + with pytest.raises(RuntimeError, match="rendering interrupted"): + _plots.plot_final_umap(store, **refs, show=False) + assert set(plt.get_fignums()) == before + + +def test_display_cannot_hide_clusters_to_fit_an_impossible_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store, refs, arrays = display_store(monkeypatch, n=12) + arrays["clusters"] = np.arange(12) + with pytest.raises(RuntimeError, match="more clusters"): + _plots.plot_final_umap(store, **refs, max_points=10, show=False) + assert not arrays["umap"].read_sizes + + +def test_nonfinite_umap_fails_without_leaking_a_figure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store, refs, arrays = display_store(monkeypatch, n=12) + coordinates = np.zeros((12, 2)) + coordinates[3, 1] = np.nan + arrays["umap"] = coordinates + before = set(plt.get_fignums()) + with pytest.raises(ValueError, match="finite"): + _plots.plot_final_umap(store, **refs, show=False) + assert set(plt.get_fignums()) == before + + +def test_mid_sized_population_palette_and_default_display_use_existing_plot_result( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from scarf.plotting import PlotResult + + store, refs, arrays = display_store(monkeypatch, n=30) + arrays["clusters"] = np.repeat(np.arange(15), 2) + shown = [] + monkeypatch.setattr(PlotResult, "show", lambda self: shown.append(self)) + plot = _plots.plot_final_umap(store, **refs) + try: + assert shown == [plot] + assert len(plot.scales[0].palette) == 15 + assert len(plot.axes["clusters"].get_legend().get_texts()) == 15 + finally: + plot.close() + + +def test_cluster_count_summary_rejects_matrix_valued_labels( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store, refs, arrays = display_store(monkeypatch, n=12) + arrays["clusters"] = np.ones((12, 2), dtype=int) + with pytest.raises(ValueError, match="one-dimensional"): + _plots.cluster_counts(store, refs["clusters"]) diff --git a/tests/test_agent_population_support.py b/tests/test_agent_population_support.py new file mode 100644 index 00000000..ad51b8e3 --- /dev/null +++ b/tests/test_agent_population_support.py @@ -0,0 +1,249 @@ +"""Population support uses frozen cells, bounded reads and observed study units.""" + +from collections import Counter +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +from scarf.agent.parameter_tuning import diagnostics +from scarf.agent.parameter_tuning.contracts import ParameterCandidateEvaluation +from scarf.storage.refs import ArtifactRef +from tests.agent_examples import example + + +class _BoundedArray: + def __init__(self, values: np.ndarray) -> None: + self.values = values + self.shape, self.dtype = values.shape, values.dtype + self.reads: list[int] = [] + + def __getitem__(self, key: Any) -> np.ndarray: + value = self.values[key] + assert value.size <= 65_536 + self.reads.append(value.size) + return value + + +class _Metadata: + def __init__(self, values: dict[str, np.ndarray]) -> None: + self.arrays = {column: _BoundedArray(value) for column, value in values.items()} + self.columns = [*values, "author_cell_type"] + self.missing: dict[str, _BoundedArray] = {} + self.accesses: list[str] = [] + + def _get_array(self, column: str) -> _BoundedArray: + assert column != "author_cell_type" + self.accesses.append(column) + return self.arrays[column] + + def _get_missing_mask_array(self, column: str) -> _BoundedArray | None: + return self.missing.get(column) + + +def _setup( + monkeypatch: pytest.MonkeyPatch, + labels: np.ndarray, + metadata: dict[str, np.ndarray], + rows: np.ndarray | None = None, +) -> tuple[Any, ParameterCandidateEvaluation, _BoundedArray]: + evaluation = example(ParameterCandidateEvaluation) + selection = ArtifactRef( + scope="datastore", kind="cell_selection", artifact_id="c" * 64 + ) + evaluation.cellSelection = evaluation.cellSelection.model_copy( + update={ + "scope": "datastore", + "assay": None, + "kind": "cell_selection", + "artifactId": selection.artifact_id, + } + ) + status = SimpleNamespace( + exists=True, complete=True, inputs={"cell_selection": selection.to_dict()} + ) + calls: Counter[str] = Counter() + array = _BoundedArray(labels) + + def inspect(ref: ArtifactRef) -> Any: + calls["inspect_artifact"] += 1 + return status + + def load(ref: ArtifactRef) -> Any: + calls["load_artifact"] += 1 + return {"values": array} + + store = SimpleNamespace( + zw=object(), + cells=_Metadata(metadata), + inspect_artifact=inspect, + load_artifact=load, + calls=calls, + status=status, + ) + monkeypatch.setattr(diagnostics, "as_zarr_array", lambda value, **kwargs: value) + + def selected(*args: Any, **kwargs: Any) -> np.ndarray: + assert args[1] == selection + return np.arange(len(labels)) if rows is None else rows + + monkeypatch.setattr(diagnostics, "read_stored_selection_indices", selected) + return store, evaluation, array + + +def test_support_exposes_restricted_populations_and_missing_unit_coverage( + monkeypatch: pytest.MonkeyPatch, +) -> None: + rows = np.asarray([1, 3, 4, 6, 8, 9, 10, 11]) + donors = np.full(12, "unused", dtype=object) + captures = donors.copy() + donors[rows] = ["d1", "d1", "d1", "d2", "d2", "d2", "d2", "d3"] + captures[rows] = ["c1", "c1", "c1", "c2", "c3", "c3", "c3", "c4"] + store, evaluation, _ = _setup( + monkeypatch, + np.asarray([1, 1, 1, 1, 2, 2, 2, 3]), + {"donor": donors, "capture": captures}, + rows, + ) + mask = np.zeros(12, dtype=bool) + mask[10] = True + store.cells.missing["donor"] = _BoundedArray(mask) + evidence = diagnostics.population_support_evidence( + store, evaluation, ["donor", "capture"] + ) + donor = evidence["columns"]["donor"] + assert evidence["selectedCells"] == 8 + assert evidence["observedPopulations"] == 3 + assert donor["observedGroups"] == 3 + assert donor["coveredCells"] == 7 + assert donor["missingCells"] == 1 + assert donor["coverageFraction"] == 7 / 8 + assert donor["populations"][0]["cluster"] == "3" + populations = {row["cluster"]: row for row in donor["populations"]} + assert populations["1"]["supportingGroups"] == 2 + assert populations["1"]["largestGroupFraction"] == 3 / 4 + assert populations["1"]["topGroups"][1] == { + "value": "d2", + "valueType": "str", + "cells": 1, + "fractionOfPopulation": 1 / 4, + "fractionOfGroup": 1 / 3, + } + assert populations["2"]["supportingGroups"] == 1 + assert populations["2"]["coverageFraction"] == 2 / 3 + assert populations["2"]["largestGroupFraction"] == 2 / 3 + assert populations["3"]["groupsWithAtLeast5Cells"] == 0 + assert evidence["columns"]["capture"]["observedGroups"] == 4 + assert store.calls == {"inspect_artifact": 1, "load_artifact": 1} + assert set(store.cells.accesses) == {"donor", "capture"} + + +def test_support_reads_large_selected_axes_in_bounded_blocks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + n = 70_000 + labels = np.zeros(n, dtype=np.int32) + labels[-5:] = 1 + groups = np.full(n, "main", dtype=object) + groups[-5:] = "rare" + store, evaluation, array = _setup(monkeypatch, labels, {"donor": groups}) + evidence = diagnostics.population_support_evidence(store, evaluation, ["donor"]) + assert array.reads == [65_536, n - 65_536] + assert store.cells.arrays["donor"].reads == [65_536, n - 65_536] + rare = evidence["columns"]["donor"]["populations"][0] + assert rare["cluster"] == "1" + assert rare["cells"] == 5 + assert rare["supportingGroups"] == rare["groupsWithAtLeast5Cells"] == 1 + assert rare["largestGroupFraction"] == 1.0 + + +def test_group_display_truncation_keeps_exact_counts_and_fractions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + groups = np.asarray([f"d{i}" for i in range(70)]) + store, evaluation, _ = _setup( + monkeypatch, np.ones(70, dtype=int), {"donor": groups} + ) + evidence = diagnostics.population_support_evidence(store, evaluation, ["donor"]) + column = evidence["columns"]["donor"] + population = column["populations"][0] + assert column["observedGroups"] == population["supportingGroups"] == 70 + assert population["omittedGroups"] == population["omittedCells"] == 65 + assert len(population["topGroups"]) == 5 + assert population["largestGroupFraction"] == 1 / 70 + assert ( + sum(row["cells"] for row in population["topGroups"]) + + population["omittedCells"] + == population["cells"] + ) + + +def test_population_and_column_display_omissions_are_explicit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store, evaluation, _ = _setup( + monkeypatch, np.arange(70), {"donor": np.asarray(["one"] * 70)} + ) + evidence = diagnostics.population_support_evidence( + store, evaluation, ["donor", "absent", "unused"] + ) + column = evidence["columns"]["donor"] + assert evidence["observedPopulations"] == 70 + assert column["omittedPopulations"] == column["omittedPopulationCells"] == 6 + assert len(column["populations"]) == 64 + assert evidence["omittedColumns"] == ["unused"] + assert evidence["columns"]["absent"]["status"] == "unavailable" + assert store.cells.accesses == ["donor"] + + +def test_missing_nonfinite_and_typed_unit_values_are_not_silently_merged( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store, evaluation, _ = _setup( + monkeypatch, + np.ones(7, dtype=int), + { + "donor": np.asarray( + [1, "1", True, None, float("nan"), float("inf"), ""], dtype=object + ) + }, + ) + column = diagnostics.population_support_evidence(store, evaluation, ["donor"])[ + "columns" + ]["donor"] + assert column["observedGroups"] == 3 + assert column["coveredCells"] == 3 + assert column["missingCells"] == 4 + assert {row["valueType"] for row in column["populations"][0]["topGroups"]} == { + "str", + "int", + "bool", + } + + +@pytest.mark.parametrize( + "damage", ["failed", "incomplete", "wrongCells", "shape", "nonfinite", "fractional"] +) +def test_invalid_population_evidence_fails_before_metadata_reads( + monkeypatch: pytest.MonkeyPatch, damage: str +) -> None: + labels = np.asarray([1.0, 2.0]) + store, evaluation, array = _setup( + monkeypatch, labels, {"donor": np.asarray(["a", "b"])} + ) + if damage == "failed": + evaluation.status = "failed" + elif damage == "incomplete": + store.status.complete = False + elif damage == "wrongCells": + store.status.inputs["cell_selection"]["artifact_id"] = "d" * 64 + elif damage == "shape": + array.shape = (3,) + elif damage == "nonfinite": + labels[0] = np.nan + else: + labels[0] = 1.5 + with pytest.raises(ValueError): + diagnostics.population_support_evidence(store, evaluation, ["donor"]) + assert store.cells.accesses == [] diff --git a/tests/test_agent_provider_contracts.py b/tests/test_agent_provider_contracts.py new file mode 100644 index 00000000..50c3e9d9 --- /dev/null +++ b/tests/test_agent_provider_contracts.py @@ -0,0 +1,84 @@ +"""Provider schemas expose choices while exact evidence remains programmatic.""" + +from typing import Any + +import pytest + +from scarf.agent.biological_interpretation.contracts import ( + BiologicalInterpretationReport, +) +from scarf.agent.data_enrichment.contracts import ( + DataEnrichmentContext, + DataEnrichmentReport, + FeatureSelectionPolicy, + StudyContextSummary, +) +from scarf.agent.data_enrichment.validation import _ground_study_context_summary +from scarf.agent.parameter_tuning.contracts import ( + FinalGraphSelection, + ParameterTuningReport, +) +from tests.agent_examples import example + + +@pytest.mark.parametrize( + ("model", "derived", "authored"), + [ + ( + DataEnrichmentReport, + {"inspections", "evidenceIds", "runInfo", "toolCalls"}, + {"policies", "status"}, + ), + ( + FeatureSelectionPolicy, + {"assayType", "graphEligible", "organismName", "exactTagFeatures"}, + {"assay", "species", "excludeFamilies", "rationale"}, + ), + ( + ParameterTuningReport, + {"evaluations", "selectedArtifacts", "runInfo", "finalSelection"}, + {"recommendedCandidateId", "rationale", "comparisons", "assayReports"}, + ), + ( + FinalGraphSelection, + {"graphMethod", "nativeAssay", "integrationId", "runInfo"}, + {"selectedOptionId", "rationale", "comparisons"}, + ), + ( + BiologicalInterpretationReport, + {"clusterArtifact", "markerArtifact", "runInfo"}, + {"clusterInterpretations", "treatmentObservations", "status"}, + ), + ], +) +def test_derived_fields_are_not_requested_from_models( + model: Any, derived: set[str], authored: set[str] +) -> None: + schema = model.model_json_schema() + if "$ref" in schema: + schema = schema["$defs"][schema["$ref"].rsplit("/", 1)[-1]] + properties = set(schema["properties"]) + assert not properties.intersection(derived) + assert authored <= properties + payload = example(model).model_dump(mode="json") + assert derived <= payload.keys() + assert model.model_validate(payload).model_dump(mode="json") == payload + + +def test_bounded_model_excerpts_preserve_long_caller_context_and_references() -> None: + references = [ + (f"Experiment {index}: " + "method " * 40).strip() for index in range(14) + ] + context = DataEnrichmentContext( + studyContext="Human lung study with repeated donors and objective-specific observations.", + studyObjective="Compare the complete design, including all supplied experiments.", + experimentalDetails=references, + ) + summary = _ground_study_context_summary( + context, + StudyContextSummary(analysisIntentReferences=["Compare the complete design"]), + ) + assert summary.studyContext == context.studyContext + assert summary.studyObjective == context.studyObjective + assert summary.experimentalReferences == references + assert _ground_study_context_summary(context, summary) == summary diff --git a/tests/test_agent_provider_edges.py b/tests/test_agent_provider_edges.py new file mode 100644 index 00000000..b22d32dd --- /dev/null +++ b/tests/test_agent_provider_edges.py @@ -0,0 +1,317 @@ +"""Standalone scientific agent and provider failure edge cases.""" + +from tests.agent_examples import example + +import asyncio +from types import SimpleNamespace + +from scarf.storage.refs import ArtifactRef + +import pytest +from pydantic import ValidationError +from pydantic_ai import ModelRetry, UnexpectedModelBehavior + +import scarf.agent.biological_interpretation.tools as biological_tools +import scarf.agent.biological_interpretation.validation as biological_validation +import scarf.agent.config.agent_exec as agent_exec_module +import scarf.agent.data_enrichment.agent as enrichment_agent +import scarf.agent.data_enrichment.tools as enrichment_tools +import scarf.agent.data_enrichment.validation as enrichment_validation +import scarf.agent.experimental_context.tools as experimental_tools +import scarf.agent.experimental_context.validation as experimental_validation +from scarf.agent.biological_interpretation import ( + BiologicalInterpretationNeedsInput, + BiologicalInterpretationReport, + ClusterCompositionEvidence, + ClusterMarkerEvidence, +) +from scarf.agent.biological_interpretation.contracts import ( + BiologicalInterpretationDependencies, +) +from scarf.agent.data_enrichment import ( + AssayFeatureInspection, + DataEnrichmentAgent, + DataEnrichmentDependencies, + DataEnrichmentToolCall, +) +from scarf.agent.experimental_context import ( + CellQcProfileEvidence, + ExperimentalContextDependencies, +) +from scarf.agent.experimental_context.contracts import ( + CovariateCharacterization, + CovariateProposal, +) + + +def test_data_enrichment_cache_rollback_and_pending_branches( + monkeypatch: pytest.MonkeyPatch, +) -> None: + inspection = example(AssayFeatureInspection) + completed = DataEnrichmentDependencies( + store=object(), + assays=["RNA"], + inspections={"RNA": inspection}, + toolCalls=[ + DataEnrichmentToolCall( + name="inspect_assay_features_batch", + assay="all", + ) + ], + ) + completed_context = SimpleNamespace(deps=completed) + + assert ( + asyncio.run( + enrichment_tools.inspect_assay_features( + completed_context, + assay_name="RNA", + ) + ) + == inspection + ) + cached_batch = asyncio.run( + enrichment_tools.inspect_assay_features_batch(completed_context) + ) + assert cached_batch.inspections == [inspection] + assert cached_batch.evidenceIds == inspection.evidenceIds + + incomplete = DataEnrichmentDependencies( + assays=["RNA"], + toolCalls=[DataEnrichmentToolCall(name="sentinel", assay="RNA")], + ) + with pytest.raises(ModelRetry, match="datastore"): + asyncio.run( + enrichment_tools.inspect_assay_features_batch( + SimpleNamespace(deps=incomplete) + ) + ) + assert [call.name for call in incomplete.toolCalls] == ["sentinel"] + + provider_error = UnexpectedModelBehavior("provider output failed") + failed = enrichment_validation.failed_data_enrichment_report( + DataEnrichmentDependencies( + assays=["RNA"], + inspections={"RNA": inspection}, + evidenceIds=set(inspection.evidenceIds), + ), + error=provider_error, + model_name="test-model", + ) + assert failed.status == "failed" + assert failed.policies == [] + assert failed.inspections == [inspection] + + def fail_before_inspection(**_kwargs: object) -> object: + raise UnexpectedModelBehavior("no inspection completed") + + monkeypatch.setattr(enrichment_agent, "run_agent_sync", fail_before_inspection) + store = SimpleNamespace(assay_names=["RNA"]) + failed = DataEnrichmentAgent(object()).run(store) + assert failed.status == "failed" + assert failed.inspections == [] + assert failed.policies == [] + + +def test_biological_interpretation_cache_and_fallback_branches() -> None: + composition = example(ClusterCompositionEvidence) + composition_deps = BiologicalInterpretationDependencies( + compositionEvidence=composition + ) + assert ( + asyncio.run( + biological_tools.inspect_cluster_composition( + SimpleNamespace(deps=composition_deps) + ) + ) + == composition + ) + + marker = example(ClusterMarkerEvidence) + marker_deps = BiologicalInterpretationDependencies( + clusterValues={marker.clusterId: 0}, + markerEvidence={marker.clusterId: marker}, + ) + assert ( + asyncio.run( + biological_tools.inspect_cluster_markers( + SimpleNamespace(deps=marker_deps), + cluster_id=marker.clusterId, + ) + ) + == marker + ) + + invalid_report = BiologicalInterpretationReport( + status="done", + needsInput=BiologicalInterpretationNeedsInput(question="More context?"), + ) + with pytest.raises(ModelRetry, match="Only a needsInput"): + biological_validation.validate_biological_interpretation_report( + invalid_report, + BiologicalInterpretationDependencies(clusterValues={"0": 0}), + ) + + provider_error = UnexpectedModelBehavior("structured output failed") + with pytest.raises(UnexpectedModelBehavior, match="structured output failed"): + biological_validation.fallback_biological_interpretation_report( + BiologicalInterpretationDependencies(), + error=provider_error, + model_name="test-model", + ) + needs_markers = biological_validation.fallback_biological_interpretation_report( + BiologicalInterpretationDependencies( + clusterValues={"0": 0}, + evidenceIds={"composition:clusters"}, + ), + error=provider_error, + model_name="test-model", + ) + assert needs_markers.status == "failed" + assert needs_markers.needsInput is None + assert needs_markers.evidenceIds == ["composition:clusters"] + + +def test_experimental_context_rejects_invalid_batches_and_preserves_failed_evidence() -> ( + None +): + invalid_batches = ( + ( + [{"name": "batch", "domain": "technical", "kind": "categorical"}], + "missing", + "Unknown batch column", + ), + ( + [{"name": "condition", "domain": "biological", "kind": "categorical"}], + "condition", + "must be classified as technical", + ), + ( + [{"name": "depth", "domain": "technical", "kind": "continuous"}], + "depth", + "must be categorical", + ), + ) + for columns, batch_column, message in invalid_batches: + deps = ExperimentalContextDependencies( + characterization=CovariateCharacterization( + status="done", + columns=columns, + ) + ) + with pytest.raises(ModelRetry, match=message): + asyncio.run( + experimental_tools.analyze_experimental_design( + SimpleNamespace(deps=deps), + column_domains={}, + coefficients_of_interest=[], + units_of_inference={}, + batch_columns=[batch_column], + ) + ) + + characterization = CovariateCharacterization(status="done") + observed = example(CellQcProfileEvidence) + failed_deps = ExperimentalContextDependencies( + cellSelection=ArtifactRef( + scope="datastore", kind="cell_selection", artifact_id="c" * 64 + ), + characterization=characterization, + qcProfiles={observed.profileId: observed}, + htoIdentityColumns=["hto_identity"], + ) + failed = experimental_validation.failed_experimental_context_result( + failed_deps, + error=UnexpectedModelBehavior("design output failed"), + model_name="test-model", + ) + assert failed.status == "failed" + assert failed.characterization is characterization + assert failed.cellQc.profileId == "" + assert failed.qcProfiles == [observed] + assert failed.decision.batchCorrection.action == "needsInput" + assert failed.runInfo.agentName == "experimental_context_failed" + + +@pytest.mark.parametrize("validation_failure", [False, True]) +def test_context_failure_preserves_actionable_cause_without_argument_payload( + validation_failure: bool, +) -> None: + deps = ExperimentalContextDependencies( + cellSelection=ArtifactRef( + scope="datastore", kind="cell_selection", artifact_id="c" * 64 + ) + ) + try: + if validation_failure: + CovariateProposal( + response="tissue", + explanatoryColumns=["tissue", "condition"], + observationUnit="sample", + independentUnit="donor", + rationale="PRIVATE STUDY TEXT MUST NOT APPEAR IN ERROR NOTES", + ) + else: + raise ModelRetry("Unknown batch column 'missing_batch'") + except (ValidationError, ModelRetry) as cause: + error = UnexpectedModelBehavior("Design tool retry limit reached") + error.__cause__ = cause + result = experimental_validation.failed_experimental_context_result( + deps, error=error, model_name="test-model" + ) + notes = " ".join(result.notes) + assert result.status == "failed" + assert result.decision.batchCorrection.action == "needsInput" + assert "Design tool retry limit reached" in notes + if validation_failure: + assert "Invalid tool arguments" in notes + assert "tissue" in notes + else: + assert "Unknown batch column 'missing_batch'" in notes + assert "PRIVATE STUDY TEXT" not in notes + + +def test_agent_execution_logs_nested_failures_for_sync_and_async_runners( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class FailingAgent: + async def __aenter__(self) -> "FailingAgent": + return self + + async def __aexit__(self, *_args: object) -> bool: + return False + + async def run(self, *_args: object, **_kwargs: object) -> object: + try: + raise ValueError("inner failure") + except ValueError as cause: + raise RuntimeError("outer failure") from cause + + monkeypatch.setattr( + agent_exec_module, + "_build_agent", + lambda **_kwargs: FailingAgent(), + ) + messages: list[str] = [] + monkeypatch.setattr(agent_exec_module.logger, "error", messages.append) + with pytest.raises(RuntimeError, match="outer failure"): + agent_exec_module.run_agent_sync( + model=object(), + output_type=dict, + system_prompt="system", + user_prompt="user", + name="sync-failure", + ) + with pytest.raises(RuntimeError, match="outer failure"): + asyncio.run( + agent_exec_module.run_agent( + model=object(), + output_type=dict, + system_prompt="system", + user_prompt="user", + name="async-failure", + ) + ) + assert all("caused by ValueError: inner failure" in message for message in messages) + assert "sync-failure" in messages[0] + assert "async-failure" in messages[1] diff --git a/tests/test_agent_qc_decision_evidence.py b/tests/test_agent_qc_decision_evidence.py new file mode 100644 index 00000000..dc8a410b --- /dev/null +++ b/tests/test_agent_qc_decision_evidence.py @@ -0,0 +1,203 @@ +"""QC choices distinguish reference grouping, measured retention and biology.""" + +from types import SimpleNamespace +from copy import deepcopy + +import pytest + +from scarf.agent.experimental_context.contracts import ( + CaptureFailureEvidence, + CellQcProfileEvidence, +) +from scarf.agent.experimental_context.study import StudyContract +from scarf.agent.orchestrator import AgentOrchestrator +from scarf.agent.orchestrator.preprocessing import PreprocessingStagesMixin +from scarf.agent.decisions.rna import QcGroupingExecutorPayload + + +def profile(policy: str, retained: int) -> CellQcProfileEvidence: + return CellQcProfileEvidence.model_validate( + { + "action": "skip" if policy == "retainWithFlags" else "registeredMad", + "registeredProfile": policy, + "attributes": [] if policy == "retainWithFlags" else ["RNA_percentMito"], + "captureColumn": "library", + "sampleColumn": "library" if policy == "captureMad5" else None, + "activeCells": 100, + "retainedCells": retained, + "evidenceId": f"qc:{policy}", + } + ) + + +def test_qc_evidence_reports_true_median_fractions_and_metric_flags() -> None: + selected = profile("globalMad5", 40) + selected.sampleRetainedCells = {"a": 10, "b": 30} + selected.retainedCellsByColumn = {"condition": {"a": 10, "b": 30}} + selected.metricFlaggedCells = {"RNA_percentMito": {"highMito": 60}} + selected.captureFailureEvidence = [ + CaptureFailureEvidence( + capture="a", activeCells=50, retainedCells=10, retainedFraction=0.2 + ), + CaptureFailureEvidence( + capture="b", activeCells=50, retainedCells=30, retainedFraction=0.6 + ), + ] + reference = profile("retainWithFlags", 100) + reference.retainedCellsByColumn = {"condition": {"a": 50, "b": 50}} + summary = PreprocessingStagesMixin._profile_evidence(selected, reference).summary + assert "min/median/max=10/20/30" in summary + assert "min/median/max=20.0%/40.0%/60.0%" in summary + assert "a=10/50 (20.0%)" in summary + assert "b=30/50 (60.0%)" in summary + assert "'highMito': 60" in summary + assert "flags may overlap" in summary + assert ( + "fractions unavailable" + in PreprocessingStagesMixin._profile_evidence(selected).summary + ) + + +def test_qc_grouping_compares_the_same_cutoff_method( + monkeypatch: pytest.MonkeyPatch, +) -> None: + orchestrator = AgentOrchestrator("test-model") + profiles = [ + CellQcProfileEvidence( + action="globalGaussian", + attributes=["RNA_percentMito"], + activeCells=100, + retainedCells=60, + evidenceId="qc:globalGaussian", + ), + CellQcProfileEvidence( + action="sampleMad", + sampleColumn="library", + attributes=["RNA_percentMito"], + activeCells=100, + retainedCells=55, + evidenceId="qc:sampleMad3", + ), + profile("globalMad5", 90), + profile("captureMad5", 95), + profile("retainWithFlags", 100), + ] + seen = {} + + def resolve(_store, _request, definition, bundle, _answers, **kwargs): + seen["definition"] = definition + seen["bundle"] = bundle + seen["qcEvidence"] = kwargs["qc_evidence"] + return SimpleNamespace( + compiled=SimpleNamespace( + executorPayload=QcGroupingExecutorPayload(groupingMode="global") + ), + checkpointSha256="a" * 64, + ) + + monkeypatch.setattr(orchestrator, "_resolve_rna_decision", resolve) + orchestrator._resolve_qc_grouping_decision( + SimpleNamespace(), + SimpleNamespace(), + SimpleNamespace(qcProfiles=profiles), + StudyContract.get_blank().model_copy( + update={"physicalCaptureColumn": "library"} + ), + {}, + ) + options = seen["definition"].spec.option_by_id() + assert "qc:globalMad5" in options["qcGrouping:global"].requiredEvidenceIds + assert "qc:captureMad5" in options["qcGrouping:physicalCapture"].requiredEvidenceIds + assert ( + "qc:sampleMad3" not in options["qcGrouping:physicalCapture"].requiredEvidenceIds + ) + design = next( + row.summary for row in seen["bundle"].evidence if row.evidenceClass == "design" + ) + assert "need not be independent biological units or healthy references" in design + assert "different cutoff methods cannot isolate" in design + assert len(seen["qcEvidence"]["policies"]) == len(profiles) + + +def test_qc_handoff_deduplicates_without_losing_any_policy_measurements() -> None: + policies = [profile("globalMad5", 90), profile("captureMad5", 95)] + for item in policies: + item.resolvedBounds = [ + {"group": "a", "metric": "RNA_percentMito", "upper": 7.125} + ] + item.activeCellsByCapture = {"a": 100} + item.parameters = { + "resolvedBounds": deepcopy(item.resolvedBounds), + "captureSizes": {"a": 100}, + "captureComparisons": [ + { + "capture": "a", + "mitoQuantiles": [1.0, 3.1, 7.125], + "missingFraction": 0.05, + } + ], + } + item.retainedCellsByCombination = { + 'joint:["sex","condition"]': { + "F/treated": 0, + "M/control": item.retainedCells, + } + } + item.unsafeRetentionGroups = ["The protected joint group F/treated is absent"] + item.notes = ["A supported reference pool is unavailable"] + item.captureFailureEvidence = [ + CaptureFailureEvidence( + capture="a", + activeCells=100, + retainedCells=item.retainedCells, + retainedFraction=item.retainedCells / 100, + conditionAndUnitSafety=[ + { + "column": "age", + "kind": "continuous", + "missingFraction": 0.2, + "status": "unsupported", + "quantilesBeforeExclusion": [20, 50, 80], + } + ], + ) + ] + originals = [item.model_dump(mode="json") for item in policies] + payload = PreprocessingStagesMixin._qc_decision_evidence(policies) + shared = payload["sharedMeasurements"] + assert ( + payload["policies"][0]["captureFailureEvidence"][0]["conditionAndUnitSafetyRef"] + == payload["policies"][1]["captureFailureEvidence"][0][ + "conditionAndUnitSafetyRef" + ] + ) + assert ( + payload["policies"][0]["parameters"]["captureComparisonsRef"] + == payload["policies"][1]["parameters"]["captureComparisonsRef"] + ) + for saved, compact in zip(originals, payload["policies"], strict=True): + restored = deepcopy(compact) + for name in ("metricSources", "sourceConcordance"): + restored[name] = shared[restored.pop(name + "Ref")] + parameters = restored["parameters"] + parameters["captureComparisons"] = shared[ + parameters.pop("captureComparisonsRef") + ] + parameters["resolvedBounds"] = deepcopy(restored["resolvedBounds"]) + parameters["captureSizes"] = deepcopy(restored["activeCellsByCapture"]) + for capture in restored["captureFailureEvidence"]: + capture["conditionAndUnitSafety"] = shared[ + capture.pop("conditionAndUnitSafetyRef") + ] + assert restored == saved + assert [item.model_dump(mode="json") for item in policies] == originals + + +def test_qc_handoff_preserves_disagreeing_duplicate_thresholds_for_investigation() -> ( + None +): + selected = profile("globalMad5", 80) + selected.resolvedBounds = {"RNA_percentMito": [0.0, 7.0]} + selected.parameters["resolvedBounds"] = {"RNA_percentMito": [0.0, 10.0]} + payload = PreprocessingStagesMixin._qc_decision_evidence([selected])["policies"][0] + assert payload["resolvedBounds"] != payload["parameters"]["resolvedBounds"] diff --git a/tests/test_agent_qc_percentage_identity.py b/tests/test_agent_qc_percentage_identity.py new file mode 100644 index 00000000..a07ef9a3 --- /dev/null +++ b/tests/test_agent_qc_percentage_identity.py @@ -0,0 +1,219 @@ +"""Exact gene ownership for percentage evidence and QC execution.""" + +from pathlib import Path + +import numpy as np +import pytest + +from scarf.agent.experimental_context.characterization import _SelectionBoundCells +from scarf.agent.experimental_context.contracts import ExperimentalContextDependencies +from scarf.agent.experimental_context.qc_evidence import ( + _derive_missing_percentage_artifacts, + _offered_qc_profiles, + _qc_metric_sources, +) +from scarf.agent.cell_quality.execution import execute_registered_cell_qc +from scarf.agent.tools import core_artifact_reference +from scarf.agent.parameter_tuning.diagnostics import _covariate_associations +from scarf.datastore.datastore import DataStore +from scarf.metadata.selection import NamedCellArtifact +from tests.agent_orchestrator_store import create_store + + +def test_exact_mitochondrial_artifact_owns_filtering_without_overwriting_metadata( + tmp_path: Path, +) -> None: + path = create_store(tmp_path / "percentages.zarr") + store = DataStore( + str(path), + default_assay="RNA", + min_features_per_cell=-1, + mito_pattern="", + ribo_pattern="", + zarr_mode="r+", + ) + store.get_assay("RNA").feats.insert( + "names", + np.asarray(["MT-CO1", "MTAP", "MTOR", "RPS3"]), + overwrite=True, + ) + imported = np.asarray([100.0, 60.0, 100.0, 83.3333333333]) + store.cells.insert("RNA_percentMito", imported) + selected = store.snapshot_cell_selection("I") + sources = _derive_missing_percentage_artifacts( + store, + cell_selection=selected, + driver=("RNA", "RNA"), + quality_sources=[], + ) + deps = ExperimentalContextDependencies( + store=store, + cellSelection=selected, + cells=_SelectionBoundCells(store.zw, store.cells, selected), + qcAssay="RNA", + studyContext="Human nuclei with mitochondrial QC.", + qualityMetricArtifacts=sources, + ) + values, metadata, artifacts, evidence, concordance, _, _ = _qc_metric_sources( + deps, ("RNA", "RNA") + ) + np.testing.assert_array_equal(store.cells.fetch_all("RNA_percentMito"), imported) + np.testing.assert_allclose(values["RNA_percentMito"], [80, 0, 200 / 3, 0]) + assert "RNA_percentMito" not in metadata + assert {item.name for item in artifacts} == {"RNA_percentMito", "RNA_percentRibo"} + mito_sources = [item for item in evidence if item.metricRole == "mitochondrial"] + assert [(item.sourceType, item.usableForFiltering) for item in mito_sources] == [ + ("metadataColumn", False), + ("artifact", True), + ] + assert mito_sources[1].executionName == "RNA_percentMito" + assert concordance[0].exactlyEqual is False + support = {} + correlation = _covariate_associations( + store, + selected, + np.asarray(values["RNA_percentMito"])[:, None], + ["RNA_percentMito"], + ["qc"], + support=support, + column_artifacts={ + "RNA_percentMito": core_artifact_reference(sources[0].artifact) + }, + ) + np.testing.assert_allclose(correlation, [[1.0]]) + assert support["RNA_percentMito"]["kind"] == "continuous" + feature_ref = next( + ref for ref in mito_sources[1].inputArtifacts if ref.kind == "feature_selection" + ) + np.testing.assert_array_equal( + store.load_artifact(core_artifact_reference(feature_ref))["values"][:], + [True, False, False, False], + ) + profile = next( + value + for value in _offered_qc_profiles(deps) + if value.registeredProfile == "globalMad5" + ) + assert "RNA_percentMito" not in profile.attributes + assert any(bound["metric"] == "RNA_percentMito" for bound in profile.resolvedBounds) + # The projected thresholds must also execute under the same metric names. + filtered, _ = execute_registered_cell_qc( + store, + "globalMad5", + profile_parameters=profile.parameters, + expected_active_cells=profile.activeCells, + expected_retained_cells=profile.retainedCells, + expected_flag_counts=profile.flaggedCells, + attrs=profile.attributes, + artifact_metrics=[ + NamedCellArtifact( + name=item.name, artifact=core_artifact_reference(item.artifact) + ) + for item in profile.artifactMetrics + ], + cell_selection=selected, + ) + assert ( + int(np.asarray(store.load_artifact(filtered)["values"][:]).sum()) + == profile.retainedCells + ) + + +def test_changed_gene_definition_invalidates_percentage_artifact( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + path = create_store(tmp_path / "changed.zarr") + store = DataStore( + str(path), + default_assay="RNA", + min_features_per_cell=-1, + mito_pattern="", + ribo_pattern="", + zarr_mode="r+", + ) + selected = store.snapshot_cell_selection("I") + assay = store.get_assay("RNA") + compute = assay._compute_feature_percentage + calculated = [] + + def counted(cells, genes): + calculated.append(genes.copy()) + return compute(cells, genes) + + monkeypatch.setattr(assay, "_compute_feature_percentage", counted) + first = _derive_missing_percentage_artifacts( + store, + cell_selection=selected, + driver=("RNA", "RNA"), + quality_sources=[], + ) + repeated = _derive_missing_percentage_artifacts( + store, + cell_selection=selected, + driver=("RNA", "RNA"), + quality_sources=[], + ) + assert repeated == first + assert len(calculated) == 2 + store.get_assay("RNA").feats.insert( + "names", + np.asarray(["MT-CO1", "RPS3", "MT-ND1", "GENE2"]), + overwrite=True, + ) + changed = _derive_missing_percentage_artifacts( + store, + cell_selection=selected, + driver=("RNA", "RNA"), + quality_sources=[], + ) + assert changed[0].artifact != first[0].artifact + assert changed[1].artifact == first[1].artifact + assert len(calculated) == 3 + + +def test_unresolved_mitochondrial_definition_cannot_use_imported_percentages( + tmp_path: Path, +) -> None: + store = DataStore( + str(create_store(tmp_path / "unresolved.zarr")), + default_assay="RNA", + min_features_per_cell=-1, + mito_pattern="", + ribo_pattern="", + zarr_mode="r+", + ) + store.get_assay("RNA").feats.insert( + "names", np.asarray(["MTAP", "MTOR", "RPS3", "ACTB"]), overwrite=True + ) + imported = np.asarray([90.0, 0.0, 0.0, 0.0]) + store.cells.insert("RNA_percentMito", imported) + selected = store.snapshot_cell_selection("I") + sources = _derive_missing_percentage_artifacts( + store, + cell_selection=selected, + driver=("RNA", "RNA"), + quality_sources=[], + ) + deps = ExperimentalContextDependencies( + store=store, + cellSelection=selected, + cells=_SelectionBoundCells(store.zw, store.cells, selected), + qcAssay="RNA", + studyContext="Mitochondrial genes cannot be identified in these annotations.", + qualityMetricArtifacts=sources, + ) + values, metadata, _, evidence, _, notes, _ = _qc_metric_sources( + deps, ("RNA", "RNA") + ) + assert "RNA_percentMito" not in values + assert "RNA_percentMito" not in metadata + assert all( + not item.usableForFiltering + for item in evidence + if item.metricRole == "mitochondrial" + ) + assert any( + "mitochondrial percentage has no exact usable artifact" in n for n in notes + ) + np.testing.assert_array_equal(store.cells.fetch_all("RNA_percentMito"), imported) diff --git a/tests/test_agent_qc_policy_execution.py b/tests/test_agent_qc_policy_execution.py new file mode 100644 index 00000000..03d0f617 --- /dev/null +++ b/tests/test_agent_qc_policy_execution.py @@ -0,0 +1,209 @@ +"""Offered QC alternatives execute their exact frozen cohort and diagnostic flags.""" + +from copy import deepcopy + +import numpy as np +import pytest + +from scarf.agent.cell_quality.execution import execute_registered_cell_qc +from scarf.agent.cell_quality.profiles import project_registered_qc_profile +from scarf.storage.artifacts import artifact_group +from scarf.storage.selections import read_stored_selection_mask +from tests.test_registered_qc_profiles import ( + _memory_qc_store, + _profile_parameters, + _quality_values, +) + + +@pytest.mark.parametrize( + "policy", + [ + "retainWithFlags", + "globalMad5", + "captureMad5", + "captureMad3Sensitivity", + "pooledReferenceMad5", + ], +) +def test_qc_policy_projection_matches_execution_on_a_preselected_cohort( + policy: str, +) -> None: + values = _quality_values() + captures = np.asarray(["a"] * 21 + ["b"] * 21) + store, _ = _memory_qc_store({**values, "capture": captures}) + store.cells._get_array("I")[0] = False + source = store.snapshot_cell_selection("I") + active = store.cells.fetch_all("I") + grouped = policy in {"captureMad5", "captureMad3Sensitivity", "pooledReferenceMad5"} + references = ("a", "b") if policy == "pooledReferenceMad5" else () + projection = project_registered_qc_profile( + policy, + values_by_metric={name: value[active] for name, value in values.items()}, + active=np.ones(int(active.sum()), dtype=bool), + capture_labels=captures[active] if grouped else None, + grouping_proven=grouped, + pooled_reference_captures=references, + ) + parameters = _profile_parameters(projection) + parameters["pooledReferenceCaptures"] = list(references) + selected, flags = execute_registered_cell_qc( + store, + policy, + profile_parameters=parameters, + expected_active_cells=int(active.sum()), + expected_retained_cells=projection.retainedCells, + expected_flag_counts=projection.flagCounts, + attrs=list(values), + cell_selection=source, + sample_column="capture" if grouped else None, + ) + observed = read_stored_selection_mask( + store.zw, + selected, + kind="cell_selection", + scope="datastore", + assay=None, + table_path="cellData", + ) + expected = np.zeros(len(active), dtype=bool) + expected[active] = projection.keep + np.testing.assert_array_equal(observed, expected) + np.testing.assert_array_equal(store.cells.fetch_all("I"), active) + assert flags is not None + np.testing.assert_array_equal( + artifact_group(store.zw, flags)["values"][:], + np.column_stack([projection.flags[name] for name in sorted(projection.flags)]), + ) + if policy == "retainWithFlags": + np.testing.assert_array_equal(observed, active) + assert any(count > 0 for count in projection.flagCounts.values()) + for name, original in values.items(): + np.testing.assert_array_equal(store.cells.fetch_all(name), original) + + +@pytest.mark.parametrize("changed", ["captureSizes", "captureComparisons", "metric"]) +def test_capture_execution_rejects_stale_evidence_before_saving_a_selection( + changed: str, +) -> None: + values = _quality_values() + captures = np.asarray(["a"] * 21 + ["b"] * 21) + store, source = _memory_qc_store({**values, "capture": captures}) + projection = project_registered_qc_profile( + "captureMad5", + values_by_metric=values, + active=np.ones(42, dtype=bool), + capture_labels=captures, + grouping_proven=True, + ) + parameters = deepcopy(_profile_parameters(projection)) + if changed == "captureSizes": + parameters["captureSizes"]["a"] -= 1 + elif changed == "captureComparisons": + parameters["captureComparisons"][0]["retainedCells"] = -1 + else: + store.cells._get_array("RNA_nCounts")[0] = 0 + with pytest.raises(ValueError, match="do not match"): + execute_registered_cell_qc( + store, + "captureMad5", + profile_parameters=parameters, + expected_active_cells=42, + expected_retained_cells=projection.retainedCells, + expected_flag_counts=projection.flagCounts, + attrs=list(values), + cell_selection=source, + sample_column="capture", + ) + np.testing.assert_array_equal(store.cells.fetch_all("I"), np.ones(42, dtype=bool)) + + +@pytest.mark.parametrize( + "invalid", ["empty", "matrixSelection", "matrixMetric", "unaligned", "nonfinite"] +) +def test_qc_projection_rejects_unmeasurable_or_misaligned_selected_cells( + invalid: str, +) -> None: + active = np.ones(4, dtype=bool) + counts = np.asarray([10.0, 11, 12, 13]) + if invalid == "empty": + active[:] = False + elif invalid == "matrixSelection": + active = active.reshape(2, 2) + elif invalid == "matrixMetric": + counts = counts.reshape(2, 2) + elif invalid == "unaligned": + counts = counts[:-1] + else: + counts[0] = np.nan + with pytest.raises(ValueError): + project_registered_qc_profile( + "globalMad5", values_by_metric={"RNA_nCounts": counts}, active=active + ) + + +def test_qc_ignores_unselected_missing_values_and_diagnostic_only_covariates() -> None: + active = np.ones(40, dtype=bool) + active[0] = False + counts = np.arange(40, dtype=float) + 100 + counts[0] = np.nan + reference = project_registered_qc_profile( + "globalMad5", values_by_metric={"RNA_nCounts": counts}, active=active + ) + measured = project_registered_qc_profile( + "globalMad5", + values_by_metric={"RNA_nCounts": counts, "age": np.arange(40) ** 4}, + active=active, + ) + np.testing.assert_array_equal(measured.keep, reference.keep) + assert measured.thresholds == reference.thresholds + assert not measured.keep[0] + + +@pytest.mark.parametrize("references", [("a", "a"), ("a", "missing"), ("a", "b")]) +def test_invalid_reference_pool_is_unavailable_without_changing_global_evidence( + references: tuple[str, ...], +) -> None: + from scarf.agent.cell_quality.profiles import offered_registered_qc_profiles + + active = np.ones(30, dtype=bool) + captures = np.asarray(["a"] * 5 + ["b"] * 5 + ["c"] * 20) + values = {"RNA_nCounts": np.arange(30, dtype=float) + 100} + with pytest.raises( + ValueError, match="reference captures|reference captures do not" + ): + project_registered_qc_profile( + "pooledReferenceMad5", + values_by_metric=values, + active=active, + capture_labels=captures, + grouping_proven=True, + pooled_reference_captures=references, + ) + offered = offered_registered_qc_profiles( + values_by_metric=values, + active=active, + capture_labels=captures, + grouping_proven=True, + pooled_reference_captures=references, + ) + assert {profile.profile for profile in offered} == {"retainWithFlags", "globalMad5"} + global_profile = next( + profile for profile in offered if profile.profile == "globalMad5" + ) + reference = project_registered_qc_profile( + "globalMad5", values_by_metric=values, active=active + ) + np.testing.assert_array_equal(global_profile.keep, reference.keep) + + +def test_capture_provenance_cannot_merge_distinct_typed_labels() -> None: + captures = np.asarray([1] * 20 + ["1"] * 20, dtype=object) + with pytest.raises(ValueError, match="consistent label type"): + project_registered_qc_profile( + "captureMad5", + values_by_metric={"RNA_nCounts": np.arange(40, dtype=float) + 100}, + active=np.ones(40, dtype=bool), + capture_labels=captures, + grouping_proven=True, + ) diff --git a/tests/test_agent_qc_source_boundaries.py b/tests/test_agent_qc_source_boundaries.py new file mode 100644 index 00000000..7b412033 --- /dev/null +++ b/tests/test_agent_qc_source_boundaries.py @@ -0,0 +1,177 @@ +"""QC policies cannot acquire invented metric, capture, or replication support.""" + +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +from scarf.agent.experimental_context import qc_evidence as q +from scarf.agent.experimental_context.contracts import ( + CaptureProposal, + NamedArtifactSource, +) +from scarf.agent.tools import artifact_reference +from scarf.storage.refs import ArtifactRef +from tests.test_agent_context_computation_reuse import _capture_context + + +@pytest.mark.parametrize( + "directions", + [ + {"physicalCaptureColumn": 0}, + { + "physicalCaptureColumn": "capture", + "cellQc": {"physicalCaptureColumn": "donor"}, + }, + {"physicalCaptureColumn": "absent"}, + {"physicalCaptureColumn": "capture", "pooledReferenceCaptures": ["d1", "d1"]}, + {"physicalCaptureColumn": "capture", "pooledReferenceCaptures": [1, 2]}, + {"pooledReferenceCaptures": ["d1", "d2"]}, + ], +) +def test_unproven_or_conflicting_capture_scope_cannot_create_qc_policies( + directions: dict[str, Any], +) -> None: + deps, characterized = _capture_context() + deps.directions = directions + with pytest.raises(ValueError): + q._offered_qc_profiles(deps, characterized) + assert deps.qcDesignData is None + + +def test_reference_proposal_cannot_replace_explicit_caller_references() -> None: + deps, _ = _capture_context() + deps.directions["pooledReferenceCaptures"] = ["d1", "d2"] + deps.captureProposal = CaptureProposal( + column="capture", + provenanceQuote="capture is a physical capture", + referenceCaptures=["d2", "d3"], + referenceProvenanceQuote="d2 and d3 are reference captures", + ) + with pytest.raises(ValueError, match="conflicts"): + q._directed_pooled_reference_captures(deps) + + +def test_nonnumeric_counts_remain_missing_for_every_capture() -> None: + deps, characterized = _capture_context() + deps.store.cells._values["RNA_nCounts"] = np.asarray(["unavailable"] * 24) + profiles = q._offered_qc_profiles(deps, characterized) + assert profiles + for profile in profiles: + source = next( + item for item in profile.metricSources if item.metricName == "RNA_nCounts" + ) + assert not source.usableForFiltering + assert source.missingCellsByCapture == {"d1": 8, "d2": 8, "d3": 8} + assert "RNA_nCounts" not in profile.attributes + if isinstance(profile.resolvedBounds, dict): + assert "RNA_nCounts" not in profile.resolvedBounds + else: + assert all( + bound.get("metric") != "RNA_nCounts" for bound in profile.resolvedBounds + ) + + +@pytest.mark.parametrize( + "damage", ["metadataShape", "artifactShape", "duplicateMetric"] +) +def test_competing_qc_sources_require_exact_axes_and_unique_execution_names( + monkeypatch: pytest.MonkeyPatch, damage: str +) -> None: + deps, _ = _capture_context() + if damage == "metadataShape": + monkeypatch.setattr(q, "_active_cell_count", lambda _: 23) + else: + source = NamedArtifactSource( + name="externalCounts", + artifact=artifact_reference( + ArtifactRef("assay", "quality_metric", "a" * 64, "RNA") + ), + ) + deps.qualityMetricArtifacts = ( + [source] if damage == "artifactShape" else [source, source] + ) + monkeypatch.setattr( + q, + "_resolved_artifact_values", + lambda *_args, **_kwargs: np.ones(23 if damage == "artifactShape" else 24), + ) + monkeypatch.setattr( + q, + "inspect_artifact", + lambda *_: SimpleNamespace(operation="externalQuality", inputs={}), + ) + with pytest.raises(ValueError, match="align|not unique"): + q._qc_metric_sources(deps, ("RNA", "RNA")) + + +def test_qc_artifact_provenance_keeps_valid_refs_without_promoting_malformed_records() -> ( + None +): + valid = ArtifactRef("assay", "feature_selection", "a" * 64, "RNA") + malformed = { + "scope": "assay", + "kind": "feature_selection", + "artifact_id": "missing-fields", + } + refs = q._artifact_input_references( + {"untrusted": malformed, "valid": valid.to_dict(), "duplicate": [valid]} + ) + assert refs == [artifact_reference(valid)] + + +@pytest.mark.parametrize("column", ["condition", "donor"]) +def test_capture_safety_rejects_misaligned_condition_or_independent_unit_values( + column: str, +) -> None: + deps, characterized = _capture_context() + deps.qcDesignData = q._QcDesignData(deps.cells) + deps.qcDesignData.values[column] = np.asarray(["value"] * 23) + with pytest.raises(ValueError, match="align"): + q._capture_design_safety(deps, characterized, deps.cells.fetch("capture"), "d1") + with pytest.raises(ValueError, match="align"): + q._design_retention( + deps, characterized, np.ones(24, dtype=bool), np.ones(24, dtype=bool) + ) + + +def test_capture_exclusion_preserves_margins_but_cannot_claim_complete_repeated_donor_pairs() -> ( + None +): + deps, characterized = _capture_context() + deps.store.cells._values["condition"][16:] = "case" + rows, conditions, independent = q._capture_design_safety( + deps, characterized, deps.cells.fetch("capture"), "d1" + ) + assert conditions and not independent + assert rows[0]["completePairsAfterExclusion"] == 1 + assert rows[0]["incompletePairsAfterExclusion"] == 1 + characterized.coefficients[0]["scope"] = "withinCell" + rows, conditions, independent = q._capture_design_safety( + deps, characterized, deps.cells.fetch("capture"), "d1" + ) + assert rows == [] and not conditions and not independent + + +def test_core_default_failure_is_reported_without_claiming_that_policy_executed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + deps, characterized = _capture_context() + original = q.project_auto_filter_profile + + def reject(action: str, **kwargs: Any) -> Any: + if action == "globalGaussian": + raise ValueError("core projection unavailable") + return original(action, **kwargs) + + monkeypatch.setattr(q, "project_auto_filter_profile", reject) + profiles = q._offered_qc_profiles(deps, characterized) + assert profiles and not any( + profile.action == "globalGaussian" for profile in profiles + ) + assert any( + "core projection unavailable" in note + for profile in profiles + for note in profile.notes + ) diff --git a/tests/test_agent_rate_limit_recovery.py b/tests/test_agent_rate_limit_recovery.py new file mode 100644 index 00000000..e561b2cf --- /dev/null +++ b/tests/test_agent_rate_limit_recovery.py @@ -0,0 +1,399 @@ +"""Transport throttling replays one request without repeating scientific tools.""" + +import asyncio +from datetime import datetime, timedelta, timezone +from email.utils import format_datetime + +import httpx +import pytest +from pydantic_ai.exceptions import ModelHTTPError, UsageLimitExceeded +from pydantic_ai.messages import ModelResponse, ToolCallPart +from pydantic_ai.models.function import FunctionModel +from pydantic_ai.usage import RequestUsage + +from scarf.agent.config import AgentRunConfig, agent_exec +from scarf.agent.types import AgentDataModel, AgentRunInfo + + +class Choice(AgentDataModel): + value: int + + +@pytest.fixture +def recorded_waits(monkeypatch): + waits = [] + original_sleep = asyncio.sleep + + async def sleep(delay): + waits.append(delay) + await original_sleep(0) + + monkeypatch.setattr(agent_exec.asyncio, "sleep", sleep) + return waits + + +def _response(info, value=2): + return ModelResponse( + parts=[ToolCallPart(info.output_tools[0].name, {"value": value})], + usage=RequestUsage(input_tokens=11, output_tokens=3), + ) + + +@pytest.mark.parametrize("host", ["sync", "async", "notebook"]) +def test_rate_limit_retries_same_request_without_reexecuting_tools( + recorded_waits, host +): + calls = [] + tools = [] + attempts = [] + + async def measure(): + tools.append("measured") + return {"measured": 2} + + async def respond(messages, info): + calls.append(messages) + if len(calls) == 1: + return ModelResponse( + parts=[ToolCallPart("measure", {}, tool_call_id="measurement")], + usage=RequestUsage(input_tokens=7, output_tokens=2), + ) + if len(calls) in {2, 3}: + raise ModelHTTPError(429, "test", "Too Many Requests") + return _response(info) + + kwargs = dict( + model=FunctionModel(respond), + output_type=Choice, + system_prompt="Use the completed measurement.", + user_prompt="Measure once and select.", + tools=[measure], + on_attempt=attempts.append, + ) + if host == "sync": + result = agent_exec.run_agent_sync(**kwargs) + elif host == "async": + result = asyncio.run(agent_exec.run_agent_async(**kwargs)) + else: + + async def notebook(): + return agent_exec.run_agent_sync(**kwargs) + + result = asyncio.run(notebook()) + + assert result.output.value == 2 + assert len(calls) == 4 + assert calls[1] is calls[2] is calls[3] + assert tools == ["measured"] + assert recorded_waits == [15, 30] + assert attempts == [result.runInfo] + assert result.runInfo.usage.requests == 4 + assert result.runInfo.usage.toolCalls == 1 + assert result.runInfo.usage.inputTokens == 18 + assert result.runInfo.usage.outputTokens == 5 + assert result.runInfo.usage.availability == "partial" + assert [item.requestIndex for item in result.runInfo.providerFailures] == [2, 3] + assert all(item.statusCode == 429 for item in result.runInfo.providerFailures) + assert result.runInfo.validationRetries == [] + + +@pytest.mark.parametrize( + "status, body", + [ + (429, {"error": {"code": "insufficient_quota"}}), + (429, "You exceeded your current quota"), + (429, "Insufficient credits"), + (429, "Billing limit reached"), + (429, "Daily limit reached"), + (401, "Unauthorized"), + (400, "This model does not support multimodal inputs"), + (503, "Service unavailable"), + ], +) +def test_nontransient_provider_errors_are_recorded_without_retry( + recorded_waits, status, body +): + error = ModelHTTPError(status, "test", body) + requests = [] + + async def fail(messages, info): + requests.append(messages) + raise error + + with pytest.raises(ModelHTTPError) as caught: + agent_exec.run_agent_sync( + model=FunctionModel(fail), + output_type=Choice, + system_prompt="Use evidence.", + user_prompt="Assess evidence.", + ) + assert caught.value is error + assert len(requests) == 1 + assert recorded_waits == [] + saved = error.agent_run_info + assert saved.usage.requests == 1 + assert saved.usage.availability == "unavailable" + assert len(saved.providerFailures) == 1 + assert saved.providerFailures[0].retryDelaySeconds is None + assert saved.validationRetries == [] + + +@pytest.mark.parametrize( + "source", ["direct", "cause", "date", "past", "zero", "invalid"] +) +def test_retry_after_headers_are_preserved(recorded_waits, source): + calls = [] + wait = ( + format_datetime( + datetime.now(timezone.utc) + + timedelta(seconds=40 if source == "date" else -40) + ) + if source in {"date", "past"} + else "0" + if source == "zero" + else "invalid" + if source == "invalid" + else "7" + ) + error = ModelHTTPError( + 429, + "test", + "Too Many Requests", + headers={"Retry-After": wait} if source != "cause" else None, + ) + if source == "cause": + response = httpx.Response(429, headers={"Retry-After": wait}) + original = RuntimeError("Throttled request") + original.response = response + error.__cause__ = original + + async def respond(messages, info): + calls.append(messages) + if len(calls) == 1: + raise error + return _response(info) + + result = agent_exec.run_agent_sync( + model=FunctionModel(respond), + output_type=Choice, + system_prompt="Use evidence.", + user_prompt="Assess evidence.", + ) + assert len(calls) == 2 + if source == "date": + assert 30 < recorded_waits[0] <= 40 + elif source in {"past", "zero"}: + assert recorded_waits == [0] + else: + assert recorded_waits == [15 if source == "invalid" else 7] + assert result.runInfo.providerFailures[0].retryDelaySeconds == recorded_waits[0] + + +@pytest.mark.parametrize( + "header, request_limit, expected_calls, expected_waits", + [ + (None, 10, 4, [15, 30, 60]), + ("80", 10, 2, [80]), + ("121", 10, 1, []), + (None, 2, 2, [15]), + ], +) +def test_retries_respect_attempt_wait_and_request_limits( + recorded_waits, header, request_limit, expected_calls, expected_waits +): + calls = [] + error = ModelHTTPError( + 429, + "test", + "Too Many Requests", + headers={"Retry-After": header} if header else None, + ) + + async def fail(messages, info): + calls.append(messages) + raise error + + with pytest.raises(ModelHTTPError) as caught: + agent_exec.run_agent_sync( + model=FunctionModel(fail), + output_type=Choice, + system_prompt="Use evidence.", + user_prompt="Assess evidence.", + config=AgentRunConfig(requestLimit=request_limit), + ) + assert caught.value is error + assert len(calls) == expected_calls + assert recorded_waits == expected_waits + info = error.agent_run_info + assert info.usage.requests == expected_calls + assert len(info.providerFailures) == expected_calls + assert info.providerFailures[-1].retryDelaySeconds is None + assert info.usage.availability == "unavailable" + + +def test_rate_limit_wait_cancels_without_another_request(monkeypatch): + entered_wait = asyncio.Event() + attempts = [] + requests = [] + + async def wait(delay): + entered_wait.set() + await asyncio.Future() + + async def fail(messages, info): + requests.append(messages) + raise ModelHTTPError(429, "test", "Too Many Requests") + + monkeypatch.setattr(agent_exec.asyncio, "sleep", wait) + + async def interrupt(): + task = asyncio.create_task( + agent_exec.run_agent_async( + model=FunctionModel(fail), + output_type=Choice, + system_prompt="Use evidence.", + user_prompt="Assess evidence.", + on_attempt=attempts.append, + ) + ) + await entered_wait.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(interrupt()) + assert len(requests) == len(attempts) == 1 + assert attempts[0].status == "failed" + assert attempts[0].usage.requests == 1 + assert attempts[0].usage.availability == "unavailable" + assert attempts[0].providerFailures[0].statusCode == 429 + + +def test_transport_retries_do_not_replace_semantic_repair(recorded_waits): + calls = [] + + async def respond(messages, info): + calls.append(messages) + if len(calls) == 1: + raise ModelHTTPError(429, "test", "Too Many Requests") + return _response(info, 1 if len(calls) == 2 else 2) + + def validate(choice): + if choice.value != 2: + raise ValueError("Saved measurement requires value 2") + return choice + + result = agent_exec.run_agent_sync( + model=FunctionModel(respond), + output_type=Choice, + system_prompt="Use evidence.", + user_prompt="Assess evidence.", + output_validator=validate, + ) + assert result.output.value == 2 + assert result.runInfo.usage.requests == 3 + assert len(result.runInfo.validationRetries) == 1 + assert result.runInfo.validationRetries[0].requestIndex == 2 + assert len(result.runInfo.providerFailures) == 1 + assert result.runInfo.providerFailures[0].requestIndex == 1 + assert recorded_waits == [15] + + +def test_failed_requests_consume_allowance_before_the_next_agent_step(recorded_waits): + calls = [] + operations = [] + + async def measure(): + operations.append("measured") + return 2 + + async def respond(messages, info): + calls.append(messages) + if len(calls) == 1: + raise ModelHTTPError(429, "test", "Too Many Requests") + return ModelResponse( + parts=[ToolCallPart("measure", {})], + usage=RequestUsage(input_tokens=11, output_tokens=3), + ) + + with pytest.raises( + UsageLimitExceeded, match="includes failed provider requests" + ) as caught: + agent_exec.run_agent_sync( + model=FunctionModel(respond), + output_type=Choice, + system_prompt="Use evidence.", + user_prompt="Measure once, then decide.", + tools=[measure], + config=AgentRunConfig(requestLimit=2), + ) + assert len(calls) == 2 + assert operations == ["measured"] + assert caught.value.agent_run_info.usage.requests == 2 + assert len(caught.value.agent_run_info.providerFailures) == 1 + + +def test_total_wait_allowance_is_shared_across_agent_steps(recorded_waits): + calls = [] + operations = [] + + async def measure(): + operations.append("measured") + return 2 + + async def respond(messages, info): + calls.append(messages) + if len(calls) != 2: + raise ModelHTTPError( + 429, "test", "Too Many Requests", headers={"Retry-After": "80"} + ) + return ModelResponse( + parts=[ToolCallPart("measure", {})], + usage=RequestUsage(input_tokens=11, output_tokens=3), + ) + + with pytest.raises(ModelHTTPError) as caught: + agent_exec.run_agent_sync( + model=FunctionModel(respond), + output_type=Choice, + system_prompt="Use evidence.", + user_prompt="Measure once, then decide.", + tools=[measure], + ) + assert len(calls) == 3 + assert operations == ["measured"] + assert recorded_waits == [80] + assert caught.value.agent_run_info.usage.requests == 3 + assert [ + row.retryDelaySeconds for row in caught.value.agent_run_info.providerFailures + ] == [80, None] + + +def test_old_run_serialization_does_not_gain_empty_provider_records(): + saved = AgentRunInfo().model_dump(mode="json") + assert "providerFailures" not in saved + assert AgentRunInfo.model_validate(saved).model_dump(mode="json") == saved + + +def test_schema_retry_uses_actual_request_ordinal_after_transient_quota(recorded_waits): + calls = [] + + async def respond(messages, info): + calls.append(messages) + if len(calls) == 1: + raise ModelHTTPError(429, "test", "Per-minute quota exceeded") + return _response(info, "invalid" if len(calls) == 2 else 2) + + result = agent_exec.run_agent_sync( + model=FunctionModel(respond), + output_type=Choice, + system_prompt="Use evidence.", + user_prompt="Assess evidence.", + ) + assert result.output.value == 2 + assert result.runInfo.usage.requests == 3 + assert len(result.runInfo.validationRetries) == 1 + assert result.runInfo.validationRetries[0].source == "schema" + assert result.runInfo.validationRetries[0].requestIndex == 2 + assert result.runInfo.providerFailures[0].requestIndex == 1 + assert recorded_waits == [15] diff --git a/tests/test_agent_report.py b/tests/test_agent_report.py new file mode 100644 index 00000000..217343e0 --- /dev/null +++ b/tests/test_agent_report.py @@ -0,0 +1,784 @@ +"""One faithful, read-only analysis report from the authoritative journal.""" + +import copy +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pandas as pd +import pytest + +from scarf.agent.orchestrator import journal +from scarf.agent.report import generator, plots +from scarf.agent.report.artifacts import ( + _local_root, + report_directory, + scientific_summary, +) +from scarf.agent.report.rendering import render_analysis_document +from scarf.agent.types import ArtifactReferenceModel +from scarf.storage.refs import ArtifactRef +from tests.test_agent_analysis_plots import display_store + + +def snapshot() -> dict[str, Any]: + from scarf.agent.experimental_context import ExperimentalContextResult + from scarf.agent.experimental_context.study import build_study_contract + from tests.agent_comparison_examples import comparison_review + + request = { + "studyContext": "Human RNA cells", + "studyObjective": "Find stable populations batch artifacts.", + } + context = ExperimentalContextResult.get_blank().model_dump(mode="json") + context["status"] = "done" + context["characterization"].update( + status="done", + columns=[{"name": "sample", "kind": "categorical", "domain": "design"}], + ) + context["cellQc"]["profileId"] = "selected" + context["qcProfiles"] = [ + { + "profileId": "selected", + "action": "globalGaussian", + "driverAssay": "RNA2", + "driverAssayType": "RNA", + "attributes": ["RNA2_nCounts", "RNA2_nFeatures", "RNA2_percentMito"], + "activeCells": 675218, + "retainedCells": 621200, + "retainedFraction": 621200 / 675218, + "retainedCellsByColumn": { + "condition": {"control": 320000, "treated": 301200} + }, + } + ] + context = ExperimentalContextResult.model_validate(context) + study = build_study_contract( + study_context=request["studyContext"], + study_objective=request["studyObjective"], + experimental_result=context, + ) + review = comparison_review() + return { + "runId": "exact-analysis", + "status": "completed", + "request": request, + "finalAnalysis": { + "primaryAssay": "RNA2", + "limitations": ["Condition and batch are confounded."], + }, + "stages": [ + { + "stage": "parameter_tuning", + "status": "done", + "report": { + "recommendedCandidateId": "candidate-two", + "evaluations": review["candidates"], + }, + "decisions": [], + }, + { + "stage": "experimental_context", + "status": "done", + "report": context.model_dump(mode="json"), + "outputs": {"studyContract": study.model_dump(mode="json")}, + "decisions": [ + { + "record": { + "decisionId": "cellQuality", + "rationale": "The selected quality policy retains supported study groups.", + }, + } + ], + }, + ], + "analysisReviews": [review], + } + + +def display_payload() -> dict[str, Any]: + return { + "clusterCounts": {"0": 620_000, "1": 1_200}, + "markers": [{"cluster": "1", "feature": "MS4A1", "score": 0.84}], + "umap": "plots/final_umap.png", + "displayedCells": 50_000, + "displayNotes": [], + } + + +@pytest.mark.parametrize("mode", ["visual", "structured"]) +def test_one_page_shows_recorded_choices_evidence_and_qualitative_findings( + mode, +) -> None: + state = snapshot() + state["analysisReviews"][0]["evidenceMode"] = mode + original = copy.deepcopy(state) + payload = scientific_summary(state) | display_payload() + document = render_analysis_document(payload) + assert state == original + assert "621,200 cells" in document and "2 clusters" in document + assert "50,000 of 621,200" in document + assert "QC retained 621,200 cells (92.0%)" in document + assert "Resolution 0.75 retains a small population with clear markers." in document + assert ( + "Repeat agreement was 0.92 and 84% of clusters had qualifying markers." + in document + ) + assert "Original detailed model reasoning retained in the journal." not in document + assert "small population with clear markers" in document + assert "Condition and batch are confounded." in document + assert "The selected quality policy retains supported study groups." in document + assert document.index( + "The selected quality policy retains supported study groups." + ) < document.index("Compared policy") + assert "<without>" in document and "" not in document + assert "hidden-record-id" not in document and "not-in-report" not in document + assert "technical.html" not in document and "decision-tree" not in document + assert "added noise" not in document and "weaker" not in document + assert ("No plots were supplied for visual inspection" in document) == ( + mode == "structured" + ) + assert ( + document.index("final_umap.png") + < document.index("Cell quality") + < document.index("Selected methods and evidence") + < document.index("Limits of this analysis") + < document.index("