diff --git a/Dockerfile b/Dockerfile index 29a5b99..ebc3856 100644 --- a/Dockerfile +++ b/Dockerfile @@ -8,14 +8,10 @@ RUN apt-get update \ WORKDIR /opt/llmize -ARG INSTALL_ENRICH=false -RUN python3 -m pip install --no-cache-dir --break-system-packages "ollama>=0.5" \ - && if [ "$INSTALL_ENRICH" = "true" ]; then \ - python3 -m pip install --no-cache-dir --break-system-packages tooluniverse "PyYAML>=6" ; \ - fi +RUN python3 -m pip install --no-cache-dir --break-system-packages "ollama>=0.5" -# TEST_MODEL is baked in for offline runs; LLMIZE_MODEL is the runtime default. -# Set BAKE_MODEL=true to also bake LLMIZE_MODEL for air-gapped use. +# TEST_MODEL is baked in for offline runs, LLMIZE_MODEL is the runtime default. +# Set BAKE_MODEL=true to also bake LLMIZE_MODEL for isolated use. ARG TEST_MODEL=smollm2:135m ARG LLMIZE_MODEL=gemma4 ENV LLMIZE_MODEL=${LLMIZE_MODEL} diff --git a/README.md b/README.md index 9749ce6..b5d9cf6 100644 --- a/README.md +++ b/README.md @@ -25,9 +25,8 @@ source .venv/bin/activate python3 -m pip install -r requirements.txt ``` -This installs `ollama` plus the optional enrich stack -(`tooluniverse`, `PyYAML`). Enrich is optional, the pipeline can run without -it but enrich gives more context. +This installs `ollama` plus the optional enrich stack (`tooluniverse`, `PyYAML`). +Enrich is optional, the pipeline can run withoutit but enrich gives more context. ### 3. Run the pipeline ```bash diff --git a/check_env.py b/check_env.py index d56ed63..95d69e8 100644 --- a/check_env.py +++ b/check_env.py @@ -73,32 +73,6 @@ def _ollama_server_and_models() -> list: return [server_ok, model_check] -def _tooluniverse() -> Check: - if importlib.util.find_spec("tooluniverse") is None: - return Check( - "ToolUniverse (optional, --enrich)", WARN, - "not installed; gene enrichment will be skipped. " - "Install with: python3 -m pip install tooluniverse PyYAML", - required=False, - ) - try: - from enrich import EntityEnricher - enricher = EntityEnricher() - except Exception as exc: - return Check( - "ToolUniverse (optional, --enrich)", WARN, - f"installed but failed to initialise ({type(exc).__name__}); enrichment will be skipped.", - required=False, - ) - if enricher.available: - return Check("ToolUniverse (optional, --enrich)", OK, "ready (gene-lookup tool resolved)", required=False) - return Check( - "ToolUniverse (optional, --enrich)", WARN, - "installed but no gene-lookup tool matched. Set LLMIZE_GENE_TOOL to a valid tool name.", - required=False, - ) - - def _descriptor_schema() -> Check: path = os.path.join(PROJECT_ROOT, "json_reduction", "descriptor_schema.json") if os.path.exists(path): @@ -126,7 +100,6 @@ def _data_dir_writable() -> Check: def run_checks() -> list: checks = [_python_version(), _ollama_package()] checks += _ollama_server_and_models() - checks.append(_tooluniverse()) checks.append(_descriptor_schema()) checks.append(_data_dir_writable()) return checks diff --git a/docker/boot_ollama.sh b/docker/boot_ollama.sh index 9ed1d57..ea76e66 100755 --- a/docker/boot_ollama.sh +++ b/docker/boot_ollama.sh @@ -1,19 +1,12 @@ #!/usr/bin/env bash -# Start the in-container Ollama server, wait for it, and pull the model. -# -# Reusable so both the image ENTRYPOINT (docker run) and the Nextflow process -# script can call it. The server is started with nohup + disown so it survives -# this script exiting and is reachable by later commands in the same task. -# -# Honours: -# LLMIZE_MODEL - model tag to pull/run (default: gemma4) -# OLLAMA_HOST - server/client endpoint (default: 127.0.0.1:11434) +# Start the in-container Ollama server and pull the model. nohup + disown let the +# server survive this script exiting so later commands in the task can reach it. +# Honours LLMIZE_MODEL (default gemma4) and OLLAMA_HOST (default 127.0.0.1:11434). set -euo pipefail MODEL="${LLMIZE_MODEL:-gemma4}" ENDPOINT="http://${OLLAMA_HOST:-127.0.0.1:11434}" -# Start the server only if one isn't already answering. if ! curl -sf "${ENDPOINT}/api/tags" >/dev/null 2>&1; then echo "[boot] starting 'ollama serve'..." nohup ollama serve >/tmp/ollama.log 2>&1 & @@ -33,9 +26,8 @@ for i in $(seq 1 60); do fi done -# Skip the pull if the model is already present (baked into the image, or cached -# on a mounted volume). This is what makes an offline / air-gapped image work: -# 'ollama pull' would otherwise contact the registry even for an existing model. +# Skip the pull when the model is already present (baked in or volume-cached) so the +# image works offline; 'ollama pull' would otherwise contact the registry regardless. if ollama list 2>/dev/null | grep -qF "${MODEL}"; then echo "[boot] model '${MODEL}' already present; skipping pull (offline-safe)." else diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index dac8c4e..81ad4ab 100755 --- a/docker/entrypoint.sh +++ b/docker/entrypoint.sh @@ -5,7 +5,6 @@ set -euo pipefail source /opt/llmize/docker/boot_ollama.sh -# Default to the env preflight if no command was given. if [ "$#" -eq 0 ]; then exec python3 /opt/llmize/pipeline.py --check fi diff --git a/enrich.py b/enrich.py deleted file mode 100644 index 963adbb..0000000 --- a/enrich.py +++ /dev/null @@ -1,351 +0,0 @@ -"""ToolUniverse gene enrichment: extract entities, look up annotations, build a -reference block for the system prompt. Optional, cached, and fail-safe. - - python3 enrich.py --input data/annotated_report.json -""" - -from __future__ import annotations - -import argparse -import json -import os -import sys - -# cache ToolUniverse lookups -CACHE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data", "enrichment_cache.json") -DEFAULT_MAX_GENES = 25 - -# Candidate ToolUniverse tool names for gene lookup -GENE_TOOL_CANDIDATES = [ - "get_target_id_description_by_name", - "get_target_synonyms_by_ensemblID", -] - -# Argument name the gene tool expects. -GENE_TOOL_ARG = "targetName" - -def _collect_cell_types(report: dict) -> set: - """Gather cell-type names from the cell-type and spatial-neighbor sections.""" - cell_types: set = set() - skip = {"unknown", "na"} - - for name, section in report.items(): - if not isinstance(section, dict): - continue - data = section.get("data", {}) - if not isinstance(data, dict): - continue - - if name == "multiqc_spatial_neighbors": - for sub in data.values(): - if not isinstance(sub, dict): - continue - focal = sub.get("focal_cell_type") - if isinstance(focal, str) and focal.lower() not in skip: - cell_types.add(focal) - for sample in sub.get("data", {}).values(): - if isinstance(sample, dict): - cell_types.update(sample.keys()) - elif name.endswith("_ct") or "deconvolved" in name: - for sample in data.values(): - if isinstance(sample, dict): - cell_types.update(sample.keys()) - - return {c for c in cell_types if isinstance(c, str) and c.lower() not in skip} - - -def _collect_moran_genes(report: dict) -> list: - section = report.get("multiqc_Moran_I_interactions") - if not isinstance(section, dict): - return [] - - best: dict = {} - for sample in section.get("data", {}).values(): - if not isinstance(sample, dict): - continue - for key, score in sample.items(): - gene = key[:-2] if key.endswith("-I") else key - if isinstance(score, (int, float)): - best[gene] = max(best.get(gene, float("-inf")), score) - - return [g for g, _ in sorted(best.items(), key=lambda kv: kv[1], reverse=True)] - - -def _parse_ligrec_key(key: str, cell_types: set) -> dict | None: - remainder = key - trailing = [] - for _ in range(2): - match = None - for ct in cell_types: - suffix = "-" + ct - if remainder.endswith(suffix) and (match is None or len(ct) > len(match)): - match = ct - if match is None: - break - trailing.insert(0, match) - remainder = remainder[: -(len(match) + 1)] - - if len(trailing) < 2 or not remainder: - return None - - return { - "ligand_receptor": remainder, - "sender": trailing[0], - "receiver": trailing[1], - "raw": key, - } - - -def extract_entities(report: dict, max_genes: int = DEFAULT_MAX_GENES) -> dict: - """Extract genes, cell types, and ligand-receptor pairs from an annotated report.""" - cell_types = _collect_cell_types(report) - genes = _collect_moran_genes(report) - - ligrec = [] - section = report.get("multiqc_squidpy_ligrec_interactions") - if isinstance(section, dict): - seen = set() - for sample in section.get("data", {}).values(): - if not isinstance(sample, dict): - continue - for key in sample: - parsed = _parse_ligrec_key(key, cell_types) - if parsed and parsed["ligand_receptor"] not in seen: - seen.add(parsed["ligand_receptor"]) - ligrec.append(parsed) - - return { - "genes": genes[:max_genes], - "all_genes_count": len(genes), - "cell_types": sorted(cell_types), - "ligand_receptor_pairs": ligrec, - } - - -# ToolUniverse lookups - -def _load_cache() -> dict: - if os.path.exists(CACHE_PATH): - try: - with open(CACHE_PATH, encoding="utf-8") as f: - return json.load(f) - except Exception: - return {} - return {} - - -def _save_cache(cache: dict) -> None: - os.makedirs(os.path.dirname(CACHE_PATH), exist_ok=True) - with open(CACHE_PATH, "w", encoding="utf-8") as f: - json.dump(cache, f, indent=2) - - -class EntityEnricher: - - def __init__(self): - self._tu = None - self._gene_tool = None - self._gene_arg = GENE_TOOL_ARG - self._cache = _load_cache() - self._init_tooluniverse() - - @property - def available(self) -> bool: - return self._tu is not None and self._gene_tool is not None - - def _init_tooluniverse(self) -> None: - try: - from tooluniverse import ToolUniverse - except Exception as exc: - print(f"[enrich] tooluniverse not available ({exc}); skipping live enrichment.") - return - - try: - tu = ToolUniverse() - tu.load_tools() - except Exception as exc: - print(f"[enrich] Could not initialise ToolUniverse ({exc}); skipping live enrichment.") - return - - self._tu = tu - self._gene_tool = self._select_gene_tool(tu) - if self._gene_tool is None: - print( - "[enrich] No gene-lookup tool matched. Set LLMIZE_GENE_TOOL to a valid " - "ToolUniverse tool name (find one with `tu list | grep -i gene`)." - ) - - @staticmethod - def _known_tool_names(tu) -> set: - names = set() - tools = getattr(tu, "all_tools", None) - if isinstance(tools, (list, tuple)): - for t in tools: - if isinstance(t, dict) and "name" in t: - names.add(t["name"]) - elif isinstance(t, str): - names.add(t) - tool_dict = getattr(tu, "all_tool_dict", None) - if isinstance(tool_dict, dict): - names.update(tool_dict.keys()) - return names - - def _select_gene_tool(self, tu): - known = self._known_tool_names(tu) - env_choice = os.environ.get("LLMIZE_GENE_TOOL") - candidates = ([env_choice] if env_choice else []) + GENE_TOOL_CANDIDATES - self._gene_arg = os.environ.get("LLMIZE_GENE_ARG", GENE_TOOL_ARG) - for name in candidates: - if not known or name in known: - return name - return None - - def _run_tool(self, gene: str) -> str | None: - spec = {"name": self._gene_tool, "arguments": {self._gene_arg: gene}} - try: - runner = getattr(self._tu, "run", None) - if callable(runner): - result = runner(spec) - else: - result = self._tu.run_tool(self._gene_tool, arguments={self._gene_arg: gene}) - except Exception as exc: - print(f"[enrich] lookup failed for {gene}: {exc}") - return None - return _extract_gene_description(result, gene) or _summarise_result(result) - - def enrich_genes(self, genes: list) -> dict: - out: dict = {} - dirty = False - for gene in genes: - cache_key = f"{self._gene_tool}:{gene}" - if cache_key in self._cache: - if self._cache[cache_key]: - out[gene] = self._cache[cache_key] - continue - if not self.available: - continue - text = self._run_tool(gene) - self._cache[cache_key] = text or "" - dirty = True - if text: - out[gene] = text - if dirty: - _save_cache(self._cache) - return out - - -def _extract_gene_description(result, gene: str) -> str | None: - if isinstance(result, str): - try: - result = json.loads(result) - except Exception: - return None - if not isinstance(result, dict): - return None - - hits = result.get("data", {}).get("search", {}).get("hits") - if not isinstance(hits, list) or not hits: - return None - - exact = next( - (h for h in hits if isinstance(h, dict) and str(h.get("name", "")).upper() == gene.upper()), - None, - ) - hit = exact or hits[0] - if not isinstance(hit, dict): - return None - desc = hit.get("description") or hit.get("name") - return " ".join(str(desc).split()) if desc else None - - -def _summarise_result(result, limit: int = 280) -> str | None: - """Reduce an arbitrary ToolUniverse result to a short text snippet.""" - if result is None: - return None - if isinstance(result, str): - text = result - elif isinstance(result, dict): - for key in ("function", "description", "summary", "comment", "text"): - if isinstance(result.get(key), str) and result[key].strip(): - text = result[key] - break - else: - text = json.dumps(result, ensure_ascii=False) - else: - text = str(result) - - text = " ".join(text.split()) - if not text: - return None - return text[:limit] + ("..." if len(text) > limit else "") - - -# Reference block assembly - -def build_reference_block(entities: dict, gene_annotations: dict) -> str: - """Assemble the markdown reference block for the system prompt.""" - if not entities.get("genes") and not entities.get("cell_types"): - return "" - - lines = ["\n\n--- BIOLOGICAL REFERENCE (auto-generated context) ---"] - - if gene_annotations: - lines.append("Gene annotations (from ToolUniverse):") - for gene, text in gene_annotations.items(): - lines.append(f"- {gene}: {text}") - elif entities.get("genes"): - top = ", ".join(entities["genes"]) - lines.append( - f"Top spatially-variable genes (no external annotations available): {top}" - ) - - if entities.get("cell_types"): - lines.append("Cell types present: " + ", ".join(entities["cell_types"]) + ".") - - lines.append( - "Use these annotations to ground your interpretation; do not invent gene functions " - "beyond what is stated here and your own established knowledge." - ) - return "\n".join(lines) - - -def enrich_report(report: dict, max_genes: int = DEFAULT_MAX_GENES) -> str: - """Top-level entry point: extract entities, look them up, return a reference block.""" - entities = extract_entities(report, max_genes=max_genes) - genes = entities.get("genes", []) - print( - f"[enrich] Extracted {entities.get('all_genes_count', 0)} genes " - f"(enriching top {len(genes)}), {len(entities.get('cell_types', []))} cell types, " - f"{len(entities.get('ligand_receptor_pairs', []))} ligand-receptor pairs." - ) - - enricher = EntityEnricher() - gene_annotations = enricher.enrich_genes(genes) if genes else {} - if gene_annotations: - print(f"[enrich] Retrieved annotations for {len(gene_annotations)} genes.") - - return build_reference_block(entities, gene_annotations) - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Extract and enrich entities from an annotated report.") - parser.add_argument("--input", "-i", required=True, help="Path to annotated_report.json") - parser.add_argument("--max-genes", type=int, default=DEFAULT_MAX_GENES) - return parser.parse_args() - - -def main() -> None: - args = _parse_args() - if not os.path.exists(args.input): - sys.exit(f"[enrich] Input not found: {args.input}") - with open(args.input, encoding="utf-8") as f: - report = json.load(f) - - entities = extract_entities(report, max_genes=args.max_genes) - print(json.dumps(entities, indent=2)) - print("\n--- Reference block preview ---") - print(enrich_report(report, max_genes=args.max_genes)) - - -if __name__ == "__main__": - main() diff --git a/interpret.py b/interpret.py index ae85475..6b76128 100644 --- a/interpret.py +++ b/interpret.py @@ -1,6 +1,5 @@ from __future__ import annotations -import argparse import json import os import sys @@ -11,14 +10,10 @@ import ollama -from verify import deterministic_findings - - -# Prompt builder +from verify import deterministic_findings, extract_entities SAMPLESHEET_KEY = "multiqc_samplesheet" -# Tone rules appended to every prompt. STYLE_GUIDE = """ Style and register (apply throughout): - Write in the neutral, descriptive register of a peer-reviewed research paper. @@ -41,13 +36,13 @@ expand what it stands for. """ -# Evidence rules to curb over-interpretation. Applied to every call. EVIDENCE_RULES = """ Evidence and interpretation rules: - Tie every claim to numeric evidence. Do not infer causality; describe associations only. -- State a responder vs non-responder difference ONLY if it is consistent across at least - 2 samples per group, OR both the mean and the median support the same direction. - Otherwise state plainly: "No consistent group-level difference detected." +- State a between-group difference (for any sample-metadata grouping — e.g. response, + timepoint, region) ONLY if it is consistent across at least 2 samples per group, OR both + the mean and the median support the same direction. Otherwise state plainly: + "No consistent group-level difference detected." - When you say a group is higher or lower, check that the direction matches the numbers (e.g. do not call the smaller mean "higher"). - Do not generalize a pattern driven by a single sample. If one sample drives a group's @@ -59,7 +54,7 @@ """ SYSTEM_PROMPT = """You are an expert bioinformatician. You are given a report generated -by a bioinformatics workflow. The report incluced outputs from a number of bioinformatics +by a bioinformatics workflow. The report includes outputs from a number of bioinformatics tools and include quality control metrics, as well as actual aggregated analysis results. The report includes the samplesheet table that may contain important clinical metadata. Different report sections may have different structures and different biological meaning, @@ -71,7 +66,7 @@ def build_prompt(report: dict) -> str: report_str = json.dumps(report, indent=2) return ( - "Here is the annotated MultiQC spatial transcriptomics report.\n" + "Here is an annotated spatial transcriptomics report.\n" "Please analyze it and provide a structured biological interpretation.\n\n" f"```json\n{report_str}\n```" ) @@ -85,15 +80,30 @@ def _split_descriptor(section_obj: dict) -> tuple: def build_samplesheet_context(report: dict) -> str: - """Sample sheet block for the shared system prompt.""" + """Sample sheet block for the shared system prompt, naming the grouping columns.""" samplesheet = report.get(SAMPLESHEET_KEY) if not isinstance(samplesheet, dict): return "" + groups = sample_groups(report) + if groups: + cols = "; ".join( + f"{col} ({', '.join(sorted(set(mapping.values())))})" + for col, mapping in groups.items() + ) + grouping_note = ( + "Group samples using the sample-metadata columns that vary across samples: " + f"{cols}. Compare along whichever of these axes are biologically relevant, and " + "carry them into every section." + ) + else: + grouping_note = ( + "No sample-metadata column meaningfully separates the samples, so do not force " + "group comparisons — describe patterns across samples directly." + ) return ( "\n\n--- SAMPLE SHEET (shared context for every section below) ---\n" f"```json\n{json.dumps(samplesheet, indent=2)}\n```\n" - "Carry these per-sample labels (e.g. responder vs non-responder) into your " - "interpretation of every section." + + grouping_note ) @@ -216,54 +226,120 @@ def compute_percentages(section_name: str, data: dict) -> dict: return {} -def response_groups(report: dict) -> dict: - """Map sample_id -> response label (e.g. 'responder') from the sample sheet.""" +_STRUCTURAL_META = {"sample_id", "file", "data_directory", "expression_profile", "ref_scrna"} + + +def _is_number(s) -> bool: + try: + float(s) + return True + except (TypeError, ValueError): + return False + + +def sample_groups(report: dict, max_group_values: int = 10) -> dict: + """Auto-detect categorical grouping columns from the sample sheet. + + Returns {column: {sample_id: value}} for each sample-sheet column usable to group + samples — 2+ distinct categorical values, not all-identical, not near-unique, not + continuous. Duplicate partitions (e.g. timepoint vs timepoint_date) are collapsed. + Metadata-agnostic: works for response, timepoint, region, sample_type, etc. + """ samplesheet = report.get(SAMPLESHEET_KEY, {}) data = samplesheet.get("data", {}) if isinstance(samplesheet, dict) else {} + samples = {sid: meta for sid, meta in data.items() if isinstance(meta, dict)} + n = len(samples) + if n < 2: + return {} + + columns = {} + for sid, meta in samples.items(): + for col, val in meta.items(): + if col in _STRUCTURAL_META or val in (None, ""): + continue + columns.setdefault(col, {})[sid] = str(val) + groups = {} - for sample, meta in data.items(): - if isinstance(meta, dict): - label = meta.get("responce") or meta.get("response") - if label: - groups[sample] = label + seen_partitions = set() + for col, mapping in columns.items(): + distinct = set(mapping.values()) + if len(distinct) < 2 or len(distinct) >= n or len(distinct) > max_group_values: + continue # all-same, near-unique identifier, or too many groups + if len(distinct) > 3 and all(_is_number(v) for v in distinct): + continue # continuous (e.g. age) + partition = frozenset( + frozenset(sid for sid, v in mapping.items() if v == val) for val in distinct + ) + if partition in seen_partitions: + continue # same split as an already-kept column + seen_partitions.add(partition) + groups[col] = mapping return groups -def compute_group_stats(data: dict, groups: dict, max_metrics: int = 20) -> dict: - """Per-group mean/median for sample-keyed numeric sections (bounded metric count).""" +def compute_group_stats(data: dict, groups: dict, max_metrics: int = 20, max_groups: int = 6) -> dict: + """Per-group mean/median for sample-keyed numeric sections, for each grouping column. + + `groups` is {column: {sample_id: value}}. Returns + {column: {metric: {group_value: {mean, median}}}}. Columns with more than + `max_groups` distinct values are skipped (too many groups to compare meaningfully). + """ if not isinstance(data, dict) or not groups: return {} - samples = [s for s in data if s in groups and isinstance(data[s], dict)] - if len(samples) < 2: + numeric_samples = [s for s in data if isinstance(data.get(s), dict)] + if len(numeric_samples) < 2: return {} metrics = [] - for s in samples: + for s in numeric_samples: for k, v in data[s].items(): if isinstance(v, (int, float)) and k not in metrics: metrics.append(k) if not metrics or len(metrics) > max_metrics: return {} - by_group = {} - for s in samples: - by_group.setdefault(groups[s], []).append(s) - out = {} - for metric in metrics: - per_group = {} - for group, gsamples in by_group.items(): - vals = [data[s][metric] for s in gsamples if isinstance(data[s].get(metric), (int, float))] - if vals: - per_group[group] = { - "mean": round(statistics.mean(vals), 3), - "median": round(statistics.median(vals), 3), - } - if per_group: - out[metric] = per_group + for col, mapping in groups.items(): + by_group = {} + for s in numeric_samples: + if s in mapping: + by_group.setdefault(mapping[s], []).append(s) + if len(by_group) < 2 or len(by_group) > max_groups: + continue + col_stats = {} + for metric in metrics: + per_group = {} + for group, gsamples in by_group.items(): + vals = [data[s][metric] for s in gsamples if isinstance(data[s].get(metric), (int, float))] + if vals: + per_group[group] = { + "mean": round(statistics.mean(vals), 3), + "median": round(statistics.median(vals), 3), + } + if len(per_group) >= 2: + col_stats[metric] = per_group + if col_stats: + out[col] = col_stats return out +def _round_floats(obj, sig: int = 4): + """Recursively round floats to `sig` significant figures to drop spurious precision. + + High-precision floats (e.g. 0.28376598223386557) add no information and can send + small models into digit-echoing repetition loops; 4 significant figures is enough. + """ + if isinstance(obj, bool): + return obj + if isinstance(obj, float): + return float(f"{obj:.{sig}g}") + if isinstance(obj, dict): + return {k: _round_floats(v, sig) for k, v in obj.items()} + if isinstance(obj, list): + return [_round_floats(v, sig) for v in obj] + return obj + + def build_section_prompt(section_name: str, section_obj: dict, groups: dict = None) -> str: """Build the prompt analysing one MultiQC section.""" descriptor, data = _split_descriptor(section_obj) @@ -274,7 +350,7 @@ def build_section_prompt(section_name: str, section_obj: dict, groups: dict = No "Section descriptor (defines what each field/key/value means):\n" f"```json\n{json.dumps(descriptor, indent=2)}\n```\n\n" "Section data:\n" - f"```json\n{json.dumps(data, indent=2)}\n```\n\n" + f"```json\n{json.dumps(_round_floats(data), indent=2)}\n```\n\n" ) percentages = compute_percentages(section_name, data) @@ -288,16 +364,18 @@ def build_section_prompt(section_name: str, section_obj: dict, groups: dict = No group_stats = compute_group_stats(data, groups) if groups else {} if group_stats: prompt += ( - "Per-group summary statistics (mean and median by response group). Compare " - "these group centers; do NOT infer a group difference from the min/max range " - "alone, and explicitly flag any single-sample outlier that skews a group:\n" - f"```json\n{json.dumps(group_stats, indent=2)}\n```\n\n" + "Per-group summary statistics (mean and median), computed for each sample-metadata " + "grouping present (e.g. by response, timepoint, region). Compare group centers within " + "a grouping; do NOT infer a group difference from the min/max range alone, and " + "explicitly flag any single-sample outlier that skews a group:\n" + f"```json\n{json.dumps(_round_floats(group_stats), indent=2)}\n```\n\n" ) prompt += ( "Give a concise, analytical interpretation of THIS section only: the main " - "patterns, notable or outlier values, differences between samples (relate them to " - "the responder / non-responder labels from the sample sheet), and the biological " + "patterns, notable or outlier values, and differences between samples — relating them to " + "whichever sample-metadata groupings are present in the sample sheet (if none are " + "informative, describe patterns without forcing a group comparison) — plus the biological " "meaning. Cite specific numbers/percentages. Do not speculate about sections you were not shown." ) @@ -352,7 +430,7 @@ def synthesize_sections( """Final call condensing the section analyses into a summary; returns (content, thinking).""" combined = "\n\n".join(f"### {name}\n{text.strip()}" for name, text in responses) prompt = ( - "Here are the per-section analyses of one MultiQC spatial transcriptomics report. " + "Here are the per-section analyses of one report section. " "Write the executive summary as instructed.\n\n" f"{combined}" ) @@ -420,7 +498,6 @@ def review_interpretation(text, model, report, num_ctx=32768, passes=2, print(f"[review] pass {i + 1}: {len(findings)} unsupported symbol(s): {', '.join(findings)}") else: print(f"[review] pass {i + 1}: no unsupported symbols detected; evidence and consistency review") - from enrich import extract_entities entities = extract_entities(report, max_genes=50) review_prompt = ( "Deterministic entities extracted from the report (use these as the allowed symbol set):\n" @@ -448,45 +525,53 @@ def print_thinking(label: str, thinking: str) -> None: print(f"[think] ===== end {label} =====\n") -def interpret_per_section( +def build_chunks(report: dict, whole_report: bool, groups: dict) -> list: + """Return the (name, prompt) units to interpret. Whole-report mode is a single chunk.""" + if whole_report: + return [("Whole report", build_prompt(report))] + return [(name, build_section_prompt(name, obj, groups=groups)) + for name, obj in iter_analysis_sections(report)] + + +def interpret_report( report: dict, model: str, num_ctx: int = 32768, - extra_system_context: str = "", + whole_report: bool = False, synthesize_final: bool = True, think: bool = True, gen_options: dict = None, user_instruction: str = "", ) -> str: - """Analyse each section in its own call, then combine into one document.""" + """Interpret the report as one or more chunks; a single chunk is the whole report.""" samplesheet_context = build_samplesheet_context(report) glossary_context = build_glossary_context(report) system = (SYSTEM_PROMPT + samplesheet_context + glossary_context - + extra_system_context + build_user_instruction(user_instruction)) - groups = response_groups(report) + + build_user_instruction(user_instruction)) + groups = sample_groups(report) - # Resolve the model once model = ensure_model(model) + chunks = build_chunks(report, whole_report, groups) + total = len(chunks) + print(f"[interpret] Analyzing {total} chunk(s) with model '{model}'" + f"{' (thinking)' if think else ''}.", flush=True) responses = [] - sections = list(iter_analysis_sections(report)) - total = len(sections) - print(f"[interpret] Analyzing {total} sections one-by-one with model '{model}'" - f"{' (thinking)' if think else ''}.", flush=True) - for idx, (name, section_obj) in enumerate(sections, 1): + for idx, (name, prompt) in enumerate(chunks, 1): print(f"[interpret] [{idx}/{total}] -> {name}", flush=True) started = time.monotonic() - prompt = build_section_prompt(name, section_obj, groups=groups) content, thinking = chat_ollama( prompt, model=model, system=system, num_ctx=num_ctx, think=think, gen_options=gen_options ) - elapsed = time.monotonic() - started - print(f"[interpret] [{idx}/{total}] {name} done in {elapsed:.1f}s", flush=True) + print(f"[interpret] [{idx}/{total}] {name} done in {time.monotonic() - started:.1f}s", flush=True) print_thinking(name, thinking) responses.append((name, content)) + if total == 1: + return responses[0][1] + summary = None - if synthesize_final and responses: + if synthesize_final: print("[interpret] Synthesizing executive summary from per-section analyses...", flush=True) started = time.monotonic() summary, summary_thinking = synthesize_sections( @@ -573,7 +658,6 @@ def _available_models() -> list: return [] -# Ollama call def ensure_model(model: str) -> str: """Resolve a usable Ollama model name, prompting the user if needed.""" @@ -684,35 +768,6 @@ def chat_ollama( return message.get("content", ""), (message.get("thinking") or "") -def call_ollama( - prompt: str, - model: str, - system: str = SYSTEM_PROMPT, - num_ctx: int = 32768, - think: bool = True, - gen_options: dict = None, -) -> tuple: - """Resolve the model then send one prompt; returns (content, thinking).""" - print(f"[interpret] Preparing to call model '{model}' via Ollama with num_ctx={num_ctx}...") - model = ensure_model(model) - print(f"[interpret] Calling model '{model}' via Ollama...") - return chat_ollama(prompt, model=model, system=system, num_ctx=num_ctx, think=think, gen_options=gen_options) - - -def load_report(path: str) -> dict: - if not os.path.exists(path): - sys.exit(f"[error] Input file not found: {path}") - with open(path, encoding="utf-8") as f: - return json.load(f) - - -def save_output(text: str, path: str) -> None: - os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) - with open(path, "w", encoding="utf-8") as f: - f.write(text) - print(f"[interpret] Response saved to: {path}") - - def _git_short_sha() -> str: """Return the short commit SHA (+'-dirty' if the tree has changes), or '' if unavailable.""" root = os.path.dirname(os.path.abspath(__file__)) @@ -738,7 +793,6 @@ def build_run_footer( think: bool, gen_options: dict = None, mode: str = "section-by-section", - enrich: bool = False, user_instruction: str = "", input_path: str = None, ) -> str: @@ -758,10 +812,7 @@ def g(key): if input_path: lines.append(f"- input: {input_path}") lines.append(f"- model: {model}") - lines.append( - f"- mode: {mode} | thinking: {'on' if think else 'off'} | " - f"enrichment: {'on' if enrich else 'off'}" - ) + lines.append(f"- mode: {mode} | thinking: {'on' if think else 'off'}") lines.append( f"- sampling: temperature={g('temperature')}, seed={g('seed')}, " f"top_p={g('top_p')}, top_k={g('top_k')}, num_predict={g('num_predict')}" @@ -772,138 +823,4 @@ def g(key): sha = _git_short_sha() if sha: lines.append(f"- code: {sha}") - return "\n".join(lines) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Pass an annotated MultiQC JSON report to a local Ollama model." - ) - parser.add_argument( - "--input", "-i", - required=True, - help="Path to the annotated JSON report (output of main.py)", - ) - parser.add_argument( - "--model", "-m", - default="gemma4", - help="Ollama model name to use (default: llama3)", - ) - parser.add_argument( - "--output", "-o", - default=None, - help="Path to save the LLM response. Defaults to _interpretation.txt", - ) - parser.add_argument( - "--num_ctx", - type=int, - default=32768, - help="Context window size (default: 32768)", - ) - parser.add_argument( - "--whole-report", - action="store_true", - help="Send the entire report in a single call instead of analysing each section one-by-one.", - ) - parser.add_argument( - "--enrich", - action="store_true", - help="Enrich the system prompt with ToolUniverse gene annotations (requires the 'tooluniverse' package).", - ) - parser.add_argument( - "--no-synthesis", - action="store_true", - help="Skip the final executive-summary synthesis pass (section-by-section mode only).", - ) - parser.add_argument( - "--think", - action=argparse.BooleanOptionalAction, - default=True, - help="Use the model's native thinking mode; the reasoning is printed to the terminal " - "(not saved) and kept out of the interpretation. On by default; --no-think disables it (faster).", - ) - parser.add_argument( - "--prompt", - default=None, - help="Extra instruction appended to the system prompt for every section " - "(e.g. 'Focus on immune cell types' or 'Answer in bullet points').", - ) - parser.add_argument("--temperature", type=float, default=None, - help="Sampling temperature (model default if unset; gemma4 defaults to 1).") - parser.add_argument("--top_p", type=float, default=None, help="Nucleus sampling top_p.") - parser.add_argument("--top_k", type=int, default=None, help="Top-k sampling.") - parser.add_argument("--seed", type=int, default=None, - help="Sampling seed for reproducible runs.") - parser.add_argument("--num_predict", type=int, default=None, - help="Max tokens to generate per call (model default if unset).") - return parser.parse_args() - - -def _build_enrichment_context(report: dict, enabled: bool) -> str: - if not enabled: - return "" - try: - from enrich import enrich_report - except Exception as exc: - print(f"[interpret] Enrichment unavailable ({exc}); continuing without it.") - return "" - return enrich_report(report) - - -def main() -> None: - args = parse_args() - - if args.output is None: - stem = os.path.splitext(os.path.basename(args.input))[0] - out_dir = os.path.dirname(os.path.abspath(args.input)) - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - args.output = os.path.join(out_dir, f"{stem}_interpretation_{timestamp}.md") - - print(f"[interpret] Loading report: {args.input}") - report = load_report(args.input) - - section_count = sum(1 for v in report.values() if isinstance(v, dict) and "data" in v) - print(f"[interpret] Loaded {section_count} data sections.") - - enrichment = _build_enrichment_context(report, args.enrich) - gen_options = build_gen_options( - temperature=args.temperature, top_p=args.top_p, top_k=args.top_k, - seed=args.seed, num_predict=args.num_predict, - ) - - if args.whole_report: - prompt = build_prompt(report) - system = (SYSTEM_PROMPT + enrichment + build_glossary_context(report) - + build_user_instruction(args.prompt)) - response_text, thinking = call_ollama( - prompt, model=args.model, system=system, - num_ctx=args.num_ctx, think=args.think, gen_options=gen_options, - ) - print_thinking("Whole report", thinking) - else: - response_text = interpret_per_section( - report, - model=args.model, - num_ctx=args.num_ctx, - extra_system_context=enrichment, - synthesize_final=not args.no_synthesis, - think=args.think, - gen_options=gen_options, - user_instruction=args.prompt, - ) - - footer = build_run_footer( - model=args.model, num_ctx=args.num_ctx, think=args.think, gen_options=gen_options, - mode="whole-report" if args.whole_report else "section-by-section", - enrich=args.enrich, user_instruction=args.prompt, input_path=args.input, - ) - save_output(response_text + footer, args.output) - - print("LLM RESPONSE PREVIEW (first 1000 chars)") - print(response_text[:1000]) - if len(response_text) > 1000: - print(f"\n... [{len(response_text) - 1000} more characters in saved file]") - - -if __name__ == "__main__": - main() \ No newline at end of file + return "\n".join(lines) \ No newline at end of file diff --git a/json_reduction/__init__.py b/json_reduction/__init__.py index 7e5e69d..4d5218f 100644 --- a/json_reduction/__init__.py +++ b/json_reduction/__init__.py @@ -1,6 +1,21 @@ -"""json_reduction: load, clean, and annotate MultiQC JSON reports. +"""json_reduction: load, reduce, and annotate MultiQC JSON reports.""" -Marking this directory as a package so `from json_reduction.X import ...` -(used by pipeline.py) resolves reliably from a fresh clone, regardless of the -current working directory. -""" +from .reduction import ( + DATA_DIR, + resolve_path, + load_json, + save_json, + extract_report_saved_raw_data, + extract_focal_labels, + annotate, +) + +__all__ = [ + "DATA_DIR", + "resolve_path", + "load_json", + "save_json", + "extract_report_saved_raw_data", + "extract_focal_labels", + "annotate", +] diff --git a/json_reduction/descriptor_schema.json b/json_reduction/descriptor_schema.json index f8edfe7..5acc0eb 100644 --- a/json_reduction/descriptor_schema.json +++ b/json_reduction/descriptor_schema.json @@ -1,10 +1,37 @@ { + "multiqc_degree_centrality": { + "description": "Squidpy degree centrality of each cell type on the spatial-neighbor graph: the (normalized) average number of neighbors a cell of that type has. Higher = that cell type sits in more densely connected neighborhoods.", + "structure": "sample_id to { cell_type_name: score }", + "key_type": "Cell type names. 'NA' = spots with no valid cell-type/neighborhood assignment.", + "value_type": "float, ~0-1 (normalized degree centrality) - higher = more connected" + }, + "multiqc_closeness_centrality": { + "description": "Squidpy closeness centrality of each cell type on the spatial-neighbor graph: how central a cell type is, i.e. the inverse of its average graph distance to other cells. Higher = more centrally embedded in the tissue neighborhood graph.", + "structure": "sample_id to { cell_type_name: score }", + "key_type": "Cell type names. 'NA' = spots with no valid assignment.", + "value_type": "float, ~0-1 - higher = more central" + }, + "multiqc_average_clustering": { + "description": "Squidpy average clustering coefficient of each cell type on the spatial-neighbor graph: how tightly cells of that type cluster together (the fraction of a cell's neighbors that are themselves neighbors). Higher = more spatial self-aggregation.", + "structure": "sample_id to { cell_type_name: score }", + "key_type": "Cell type names. 'NA' = spots with no valid assignment.", + "value_type": "float, ~0-1 - higher = more clustered" + }, + "multiqc_spacemarkers_LRscores_interactions": { + "description": "SpaceMarkers ligand-receptor interaction scores: the strength of inferred, spatially co-localized signaling for each ligand-receptor pair between a sender and a receiver cell type, per sample. Higher = stronger inferred interaction.", + "structure": "sample_id to { 'LIGAND(_RECEPTOR parts)-SENDER_to_RECEIVER': score }", + "key_type": "Interaction keys of the form -_to_ (e.g. 'IGFBP3_TMEM219-FIBROBLASTS_to_PDAC').", + "value_type": "float - higher = stronger inferred ligand-receptor interaction" + }, + "multiqc_co_occurrence": { + "description": "Mean squidpy co-occurrence: how the probability of each cell type changes with spatial distance from a focal cell type, as a ratio to that cell type's overall frequency. Each section (multiqc_co_occurrence, _1, _2, ...) is conditioned on one focal cell type. Ratio > 1 at a distance = enriched near the focal type there; < 1 = depleted; ~1 = as expected by chance. Report the overall trend (enriched/depleted, and at short vs long range) rather than every distance value.", + "structure": "focal_cell_type to [ [distance_bin, ratio], ... ] (a curve over increasing distance)", + "key_type": "Focal cell type names. 'NA' = spots with no valid assignment.", + "value_type": "list of [distance_bin, ratio]; ratio > 1 = enriched at that distance, < 1 = depleted" + }, "multiqc_samplesheet": { - "description": "The sample sheet provided as input to the pipeline. Contains per-sample metadata columns. Columns containing file paths (data_directory, expression_profile) are excluded.", - "structure": "sample_id to { metadata fields }", - "properties": { - "responce": "Treatment response label (example: responder / non-responder)." - } + "description": "The sample sheet provided as input to the pipeline. Contains per-sample metadata columns that vary by study (e.g. treatment response, timepoint, tissue region, sample type, case/control). Grouping columns are detected automatically from whatever columns are present; columns containing file paths (data_directory, expression_profile) are excluded.", + "structure": "sample_id to { metadata fields }" }, "multiqc_atlas": { "description": "Summary of the cells and genes in the reference atlas. Transcription profiles from the reference atlas are used to infer cell types in the samples.", @@ -62,7 +89,7 @@ "structure": "sample_id to { GENE-I: score }", "key_format": "GENE_NAME-I (example: 'CXCL14-I', 'HLA-DRB1-I')", "value_type": "float, range 0-1 - higher = expression more spatially clustered across tissue spots", - "analysis_instructions": "Assess whether spatial clustering differs systematically between responders and non-responders. If there is no clear, consistent group-level difference, state that plainly (e.g. 'Moran's I scores do not differ systematically between response groups') and do NOT enumerate per-sample scores gene by gene. Name only the few genes (if any) with a notable and consistent group difference, and otherwise describe the overall pattern briefly." + "analysis_instructions": "Assess whether spatial clustering differs systematically between the sample-metadata groups present. If there is no clear, consistent group-level difference, state that plainly (e.g. 'Moran's I scores do not differ systematically between groups') and do NOT enumerate per-sample scores gene by gene. Name only the few genes (if any) with a notable and consistent group difference, and otherwise describe the overall pattern briefly." }, "multiqc_spatial_neighbors": { "description": "Cell type immediate neighborhood across samples. Each section (multiqc_spatial_neighbors, _1, _2, ...) represents one focal cell type and counts how many neighbors of each type surround it. Focal cell types are dynamically generated during merge based on actual data.", diff --git a/json_reduction/json_clean.py b/json_reduction/json_clean.py deleted file mode 100644 index faaf400..0000000 --- a/json_reduction/json_clean.py +++ /dev/null @@ -1,36 +0,0 @@ -TARGET_KEYS = ["report_saved_raw_data", "report_raw_saved_data"] -IGNORE_KEYS = { - "multiqc_samplesheet": {"data_directory", "expression_profile"}, -} - -def extract_report_saved_raw_data(data : dict) -> dict: - for key in TARGET_KEYS: - if key in data: - cleaned = _strip_ignored_keys(data[key]) - return {key: cleaned} - - raise KeyError( - "Neither {} found in JSON. Available keys: {}".format( - TARGET_KEYS, list(data.keys()) - ) - ) - -def _strip_ignored_keys(raw): - result = {} - for section, samples in raw.items(): - ignore = IGNORE_KEYS.get(section, set()) - if not ignore or not isinstance(samples, dict): - result[section] = samples - continue - - cleaned_samples = {} - for sample_id, metrics in samples.items(): - if isinstance(metrics, dict): - cleaned_samples[sample_id] = { - k: v for k, v in metrics.items() if k not in ignore - } - else: - cleaned_samples[sample_id] = metrics - result[section] = cleaned_samples - - return result \ No newline at end of file diff --git a/json_reduction/json_load.py b/json_reduction/json_load.py deleted file mode 100644 index 92bce67..0000000 --- a/json_reduction/json_load.py +++ /dev/null @@ -1,32 +0,0 @@ -import json -import os - -PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -DATA_DIR = os.path.join(PROJECT_ROOT, "data") - -def resolve_path(user_input: str) -> str: - if os.path.isabs(user_input): - return user_input - - cwd_path = os.path.join(os.getcwd(), user_input) - if os.path.exists(cwd_path): - return cwd_path - - data_path = os.path.join(DATA_DIR, user_input) - return data_path - -def load_json(filepath: str) -> dict: - with open(filepath, "r") as f: - return json.load(f) - -def fetch_file() -> str: - print(f"(JSON files should be placed in: {DATA_DIR})") - while True: - user_input = input("Enter JSON filename or path: ").strip() - resolved = resolve_path(user_input) - - if os.path.exists(resolved): - print(f"Loading: {resolved}") - return resolved - - print(f" File not found: '{resolved}'. Please try again.") \ No newline at end of file diff --git a/json_reduction/json_topKeys.py b/json_reduction/json_topKeys.py deleted file mode 100644 index 85ec4c2..0000000 --- a/json_reduction/json_topKeys.py +++ /dev/null @@ -1,8 +0,0 @@ -def print_top_keys (data : dict) -> None: - print("\nTop Level Keys Found: ") - for key in data.keys(): - print(f" - {key}") - - -def has_key(data : dict, key: str) -> bool: - return key in data \ No newline at end of file diff --git a/json_reduction/json_write.py b/json_reduction/json_write.py deleted file mode 100644 index c57af98..0000000 --- a/json_reduction/json_write.py +++ /dev/null @@ -1,21 +0,0 @@ -import json -import os - -def save_json(data: dict, data_dir: str, filename: str, indent: int = 2) -> str: - bare_name = os.path.basename(filename) - output_path = os.path.join(data_dir, bare_name) - os.makedirs(data_dir, exist_ok=True) - - with open(output_path, "w") as f: - json.dump(data, f, indent=indent) - - abs_path = os.path.abspath(output_path) - print(f"\nFile saved successfully.") - print(f"Location : {abs_path}") - - return abs_path - -def req_filename(source_filename: str) -> str: - default = "extracted_" + os.path.basename(source_filename) - user_input = input(f"\nEnter output filename (*Enter* = {default}): ").strip() - return user_input if user_input else default \ No newline at end of file diff --git a/json_reduction/main.py b/json_reduction/main.py deleted file mode 100644 index f9c190d..0000000 --- a/json_reduction/main.py +++ /dev/null @@ -1,69 +0,0 @@ -import os -import sys - -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -from json_load import load_json, fetch_file, PROJECT_ROOT, DATA_DIR -from json_topKeys import print_top_keys, has_key -from json_clean import extract_report_saved_raw_data -from json_write import save_json, req_filename -from json_merger import merge, req_annotated_filename - -SCHEMA_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "descriptor_schema.json") - - -def main(): - print("Project root : {}".format(PROJECT_ROOT)) - print("Data folder : {}".format(DATA_DIR)) - print("Schema : {}".format(SCHEMA_PATH)) - - if not os.path.exists(SCHEMA_PATH): - print("\nERROR: descriptor_schema.json not found at:") - print("{}".format(SCHEMA_PATH)) - print("Place it in the json_reduction/ folder.") - sys.exit(1) - - json_path = fetch_file() - - try: - data = load_json(json_path) - except Exception as e: - print("\nFailed to load JSON: {}".format(e)) - sys.exit(1) - - print_top_keys(data) - - try: - reduced = extract_report_saved_raw_data(data) - found_key = next(iter(reduced)) - section_count = len(next(iter(reduced.values()))) - print("\nExtracted '{}' with {} sections".format(found_key, section_count)) - except KeyError as e: - print("\nExtraction failed: {}".format(e)) - sys.exit(1) - - output_filename = req_filename(source_filename=json_path) - try: - reduced_path = save_json(reduced, DATA_DIR, output_filename) - except Exception as e: - print("\nFailed to save: {}".format(e)) - sys.exit(1) - - try: - annotated_filename = req_annotated_filename() - annotated_path = merge( - data_path=reduced_path, - descriptor_path=SCHEMA_PATH, - output_dir=DATA_DIR, - output_filename=annotated_filename, - ) - except Exception as e: - print("\nFailed to merge: {}".format(e)) - sys.exit(1) - - print("Extracted data : {}".format(reduced_path)) - print("Annotated report : {}".format(annotated_path)) - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/json_reduction/json_merger.py b/json_reduction/reduction.py similarity index 60% rename from json_reduction/json_merger.py rename to json_reduction/reduction.py index d004846..07c23e9 100644 --- a/json_reduction/json_merger.py +++ b/json_reduction/reduction.py @@ -1,11 +1,69 @@ +"""Load, reduce, and annotate MultiQC JSON reports. + +Public API: DATA_DIR, resolve_path, load_json, save_json, +extract_report_saved_raw_data, extract_focal_labels, annotate. +""" + import json import os +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA_DIR = os.path.join(PROJECT_ROOT, "data") + +TARGET_KEYS = ["report_saved_raw_data", "report_raw_saved_data"] +IGNORE_KEYS = { + "multiqc_samplesheet": {"data_directory", "expression_profile"}, +} + + +def resolve_path(user_input: str) -> str: + if os.path.isabs(user_input): + return user_input + cwd_path = os.path.join(os.getcwd(), user_input) + if os.path.exists(cwd_path): + return cwd_path + return os.path.join(DATA_DIR, user_input) + + +def load_json(filepath: str) -> dict: + with open(filepath, "r", encoding="utf-8") as f: + return json.load(f) -def req_annotated_filename() -> str: - default = "annotated_report.json" - user_input = input(f"\nEnter annotated report filename (*Enter* = {default}): ").strip() - return user_input if user_input else default + +def save_json(data: dict, data_dir: str, filename: str, indent: int = 2) -> str: + output_path = os.path.join(data_dir, os.path.basename(filename)) + os.makedirs(data_dir, exist_ok=True) + with open(output_path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=indent) + abs_path = os.path.abspath(output_path) + print(f"[reduction] Saved: {abs_path}") + return abs_path + + +def extract_report_saved_raw_data(data: dict) -> dict: + for key in TARGET_KEYS: + if key in data: + return {key: _strip_ignored_keys(data[key])} + raise KeyError( + "Neither {} found in JSON. Available keys: {}".format(TARGET_KEYS, list(data.keys())) + ) + + +def _strip_ignored_keys(raw): + result = {} + for section, samples in raw.items(): + ignore = IGNORE_KEYS.get(section, set()) + if not ignore or not isinstance(samples, dict): + result[section] = samples + continue + cleaned_samples = {} + for sample_id, metrics in samples.items(): + if isinstance(metrics, dict): + cleaned_samples[sample_id] = {k: v for k, v in metrics.items() if k not in ignore} + else: + cleaned_samples[sample_id] = metrics + result[section] = cleaned_samples + return result def _section_index(section_name): @@ -39,7 +97,6 @@ def extract_focal_labels(full_data): if names: return names - # Fallback: dataset-level labels. for ds in plot.get("datasets", []) or []: if isinstance(ds, dict) and ds.get("label"): names.append(ds["label"]) @@ -54,37 +111,24 @@ def _get_focal_cell_types(data, focal_labels=None): [k for k in data.keys() if "spatial_neighbors" in k], key=_section_index, ) - for idx, section in enumerate(spatial_sections): if idx < len(focal_labels): label = focal_labels[idx] else: label = f"focal cell type {idx + 1} (unlabeled)" focal_types[section] = {"index": idx, "focal_cell_type": label} - return focal_types -def merge(data_path, descriptor_path, output_dir, output_filename="annotated_report.json", - focal_labels=None): - - with open(data_path, encoding="utf-8") as f: - data = json.load(f) - - with open(descriptor_path, encoding="utf-8") as f: - descriptor = json.load(f) - - top_key = next( - (k for k in data if "saved" in k and "raw" in k), - None - ) +def annotate(reduced: dict, descriptor: dict, focal_labels=None) -> dict: + """Overlay descriptor metadata onto the reduced report; return the annotated report dict.""" + top_key = next((k for k in reduced if "saved" in k and "raw" in k), None) if top_key is None: raise KeyError( - "Could not find report_saved_raw_data in data file. " - "Available keys: {}".format(list(data.keys())) + "Could not find report_saved_raw_data. Available keys: {}".format(list(reduced.keys())) ) - raw = data[top_key] + raw = reduced[top_key] focal_cell_types = _get_focal_cell_types(raw, focal_labels=focal_labels) annotated = {} spatial_neighbors_parent = None @@ -113,7 +157,7 @@ def merge(data_path, descriptor_path, output_dir, output_filename="annotated_rep spatial_children[child_key] = { "focal_cell_type": focal_cell_types[section]["focal_cell_type"], - "data": samples + "data": samples, } else: section_schema = descriptor.get(section, {}) @@ -122,7 +166,6 @@ def merge(data_path, descriptor_path, output_dir, output_filename="annotated_rep if isinstance(v, dict) and k == "sections": continue entry[k] = v - entry["data"] = samples annotated[section] = entry @@ -130,12 +173,4 @@ def merge(data_path, descriptor_path, output_dir, output_filename="annotated_rep spatial_neighbors_parent["data"] = spatial_children annotated["multiqc_spatial_neighbors"] = spatial_neighbors_parent - os.makedirs(output_dir, exist_ok=True) - output_path = os.path.join(output_dir, output_filename) - with open(output_path, "w", encoding="utf-8") as f: - json.dump(annotated, f, indent=2) - - abs_path = os.path.abspath(output_path) - print("Annotated report saved.") - print("Location : {}".format(abs_path)) - return abs_path \ No newline at end of file + return annotated diff --git a/main.nf b/main.nf index 954ab2b..3e5baac 100644 --- a/main.nf +++ b/main.nf @@ -13,7 +13,6 @@ process INTERPRET { def home = workflow.containerEngine ? '/opt/llmize' : "${projectDir}" def boot = workflow.containerEngine ? "export LLMIZE_MODEL='${params.model}'\n bash ${home}/docker/boot_ollama.sh" : '' def think_flag = "${params.think}".toBoolean() ? '--think' : '--no-think' - def enrich_flag = "${params.enrich}".toBoolean() ? '--enrich' : '' def review_flag = "${params.review}".toBoolean() ? "--review --review-passes ${params.review_passes}" : '' def whole_flag = "${params.whole_report}".toBoolean() ? '--whole-report' : '' def synth_flag = "${params.synthesis}".toBoolean() ? '' : '--no-synthesis' @@ -31,7 +30,7 @@ process INTERPRET { --input '${report}' \\ --model '${params.model}' \\ --num_ctx ${params.num_ctx} \\ - ${think_flag} ${enrich_flag} ${review_flag} ${whole_flag} ${synth_flag} \\ + ${think_flag} ${review_flag} ${whole_flag} ${synth_flag} \\ ${prompt_flag} ${temp_flag} ${seed_flag} ${top_p_flag} ${top_k_flag} ${numpred_flag} \\ --output "${report.baseName}_interpretation_\${STAMP}.md" """ diff --git a/nextflow.config b/nextflow.config index 6004263..994d09e 100644 --- a/nextflow.config +++ b/nextflow.config @@ -4,7 +4,6 @@ params { model = 'gemma4' think = true - enrich = false review = false review_passes = 2 whole_report = false diff --git a/pipeline.py b/pipeline.py index bcfc5a0..6495686 100644 --- a/pipeline.py +++ b/pipeline.py @@ -1,8 +1,6 @@ -#!/usr/bin/env python3 from __future__ import annotations import argparse -import json import os import sys from datetime import datetime @@ -11,21 +9,20 @@ if PROJECT_ROOT not in sys.path: sys.path.insert(0, PROJECT_ROOT) -from json_reduction.json_load import resolve_path, load_json, DATA_DIR -from json_reduction.json_clean import extract_report_saved_raw_data -from json_reduction.json_write import save_json -from json_reduction.json_merger import merge, extract_focal_labels +from json_reduction import ( + resolve_path, + load_json, + save_json, + DATA_DIR, + extract_report_saved_raw_data, + extract_focal_labels, + annotate, +) from interpret import ( - build_prompt, - call_ollama, - interpret_per_section, - print_thinking, + interpret_report, build_gen_options, - build_user_instruction, - build_glossary_context, build_run_footer, review_interpretation, - SYSTEM_PROMPT, ) DEFAULT_DESCRIPTOR = os.path.join(PROJECT_ROOT, "json_reduction", "descriptor_schema.json") @@ -52,15 +49,20 @@ def parse_args() -> argparse.Namespace: default=DEFAULT_DESCRIPTOR, help="Path to the descriptor schema JSON file.", ) + parser.add_argument( + "--save-intermediates", + action="store_true", + help="Write the reduced and annotated JSON to data/ (off by default; run is in-memory).", + ) parser.add_argument( "--extracted-output", default=None, - help="Filename for the extracted intermediate JSON saved in data/.", + help="Filename for the extracted intermediate JSON (only with --save-intermediates).", ) parser.add_argument( "--annotated-output", default=None, - help="Filename for the annotated report JSON saved in data/.", + help="Filename for the annotated report JSON (only with --save-intermediates).", ) parser.add_argument( "--output", "-o", @@ -78,11 +80,6 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Interpret the whole report in one call instead of section-by-section.", ) - parser.add_argument( - "--enrich", - action="store_true", - help="Enrich the system prompt with ToolUniverse gene annotations (requires 'tooluniverse').", - ) parser.add_argument( "--no-synthesis", action="store_true", @@ -123,23 +120,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--check", action="store_true", - help="Run environment preflight checks (Ollama, models, ToolUniverse, schema) and exit.", + help="Run environment preflight checks (Ollama, models, schema) and exit.", ) return parser.parse_args() -def _build_enrichment_context(report: dict, enabled: bool) -> str: - """Return a ToolUniverse reference block for the system prompt, or '' if disabled/unavailable.""" - if not enabled: - return "" - try: - from enrich import enrich_report - except Exception as exc: - print(f"[pipeline] Enrichment unavailable ({exc}); continuing without it.") - return "" - return enrich_report(report) - - def resolve_input_path(path: str) -> str: resolved = resolve_path(path) if not os.path.exists(resolved): @@ -147,9 +132,9 @@ def resolve_input_path(path: str) -> str: return resolved -def default_output_name(input_path: str, prefix: str, suffix: str = "") -> str: +def default_output_name(input_path: str, prefix: str) -> str: stem = os.path.splitext(os.path.basename(input_path))[0] - return f"{prefix}{stem}{suffix}" + return f"{prefix}{stem}" def save_text(text: str, path: str) -> None: @@ -168,64 +153,43 @@ def run_pipeline( output_path: str | None, num_ctx: int, whole_report: bool = False, - enrich: bool = False, synthesize_final: bool = True, think: bool = True, gen_options: dict = None, user_instruction: str = "", review: bool = False, review_passes: int = 2, + save_intermediates: bool = False, ) -> str: input_path = resolve_input_path(input_json) print(f"[pipeline] Loading raw JSON: {input_path}") data = load_json(input_path) reduced = extract_report_saved_raw_data(data) - - if extracted_filename is None: - extracted_filename = default_output_name(input_path, prefix="extracted_") - extracted_path = save_json(reduced, DATA_DIR, extracted_filename) - print(f"[pipeline] Extracted JSON saved: {extracted_path}") - - if annotated_filename is None: - annotated_filename = "annotated_report.json" # Recover real focal cell types from the raw report's plot metadata. focal_labels = extract_focal_labels(data) if focal_labels: print(f"[pipeline] Recovered spatial-neighbors focal cell types: {focal_labels}") - annotated_path = merge( - data_path=extracted_path, - descriptor_path=descriptor_path, - output_dir=DATA_DIR, - output_filename=annotated_filename, - focal_labels=focal_labels, + descriptor = load_json(descriptor_path) + report = annotate(reduced, descriptor, focal_labels=focal_labels) + print(f"[pipeline] Reduced and annotated {len(report)} sections in memory.") + + if save_intermediates: + save_json(reduced, DATA_DIR, extracted_filename or default_output_name(input_path, prefix="extracted_")) + save_json(report, DATA_DIR, annotated_filename or "annotated_report.json") + + mode = "whole report" if whole_report else "section-by-section" + print(f"[pipeline] Calling Ollama model '{model}' ({mode})...") + response = interpret_report( + report, + model=model, + num_ctx=num_ctx, + whole_report=whole_report, + synthesize_final=synthesize_final, + think=think, + gen_options=gen_options, + user_instruction=user_instruction, ) - print(f"[pipeline] Annotated report saved: {annotated_path}") - - report = load_json(annotated_path) - enrichment = _build_enrichment_context(report, enrich) - - if whole_report: - prompt = build_prompt(report) - system = (SYSTEM_PROMPT + enrichment + build_glossary_context(report) - + build_user_instruction(user_instruction)) - print(f"[pipeline] Calling Ollama model '{model}' (whole report)...") - response, thinking = call_ollama( - prompt, model=model, system=system, num_ctx=num_ctx, think=think, gen_options=gen_options - ) - print_thinking("Whole report", thinking) - else: - print(f"[pipeline] Calling Ollama model '{model}' (section-by-section)...") - response = interpret_per_section( - report, - model=model, - num_ctx=num_ctx, - extra_system_context=enrichment, - synthesize_final=synthesize_final, - think=think, - gen_options=gen_options, - user_instruction=user_instruction, - ) if review: print(f"[pipeline] Reviewing interpretation (up to {review_passes} pass(es))...") @@ -243,7 +207,7 @@ def run_pipeline( footer = build_run_footer( model=model, num_ctx=num_ctx, think=think, gen_options=gen_options, mode="whole-report" if whole_report else "section-by-section", - enrich=enrich, user_instruction=user_instruction, input_path=input_path, + user_instruction=user_instruction, input_path=input_path, ) save_text(response + footer, output_path) return output_path @@ -269,13 +233,13 @@ def main() -> None: output_path=args.output, num_ctx=args.num_ctx, whole_report=args.whole_report, - enrich=args.enrich, synthesize_final=not args.no_synthesis, think=args.think, gen_options=gen_options, user_instruction=args.prompt, review=args.review, review_passes=args.review_passes, + save_intermediates=args.save_intermediates, ) print(f"[pipeline] Completed. Final report: {final_path}") diff --git a/verify.py b/verify.py index 02c70bf..62bf5e2 100644 --- a/verify.py +++ b/verify.py @@ -2,7 +2,106 @@ import re -from enrich import extract_entities +DEFAULT_MAX_GENES = 25 + + +def _collect_cell_types(report: dict) -> set: + """Gather cell-type names from the cell-type and spatial-neighbor sections.""" + cell_types: set = set() + skip = {"unknown", "na"} + + for name, section in report.items(): + if not isinstance(section, dict): + continue + data = section.get("data", {}) + if not isinstance(data, dict): + continue + + if name == "multiqc_spatial_neighbors": + for sub in data.values(): + if not isinstance(sub, dict): + continue + focal = sub.get("focal_cell_type") + if isinstance(focal, str) and focal.lower() not in skip: + cell_types.add(focal) + for sample in sub.get("data", {}).values(): + if isinstance(sample, dict): + cell_types.update(sample.keys()) + elif name.endswith("_ct") or "deconvolved" in name: + for sample in data.values(): + if isinstance(sample, dict): + cell_types.update(sample.keys()) + + return {c for c in cell_types if isinstance(c, str) and c.lower() not in skip} + + +def _collect_moran_genes(report: dict) -> list: + section = report.get("multiqc_Moran_I_interactions") + if not isinstance(section, dict): + return [] + + best: dict = {} + for sample in section.get("data", {}).values(): + if not isinstance(sample, dict): + continue + for key, score in sample.items(): + gene = key[:-2] if key.endswith("-I") else key + if isinstance(score, (int, float)): + best[gene] = max(best.get(gene, float("-inf")), score) + + return [g for g, _ in sorted(best.items(), key=lambda kv: kv[1], reverse=True)] + + +def _parse_ligrec_key(key: str, cell_types: set) -> dict | None: + remainder = key + trailing = [] + for _ in range(2): + match = None + for ct in cell_types: + suffix = "-" + ct + if remainder.endswith(suffix) and (match is None or len(ct) > len(match)): + match = ct + if match is None: + break + trailing.insert(0, match) + remainder = remainder[: -(len(match) + 1)] + + if len(trailing) < 2 or not remainder: + return None + + return { + "ligand_receptor": remainder, + "sender": trailing[0], + "receiver": trailing[1], + "raw": key, + } + + +def extract_entities(report: dict, max_genes: int = DEFAULT_MAX_GENES) -> dict: + """Extract genes, cell types, and ligand-receptor pairs from an annotated report.""" + cell_types = _collect_cell_types(report) + genes = _collect_moran_genes(report) + + ligrec = [] + section = report.get("multiqc_squidpy_ligrec_interactions") + if isinstance(section, dict): + seen = set() + for sample in section.get("data", {}).values(): + if not isinstance(sample, dict): + continue + for key in sample: + parsed = _parse_ligrec_key(key, cell_types) + if parsed and parsed["ligand_receptor"] not in seen: + seen.add(parsed["ligand_receptor"]) + ligrec.append(parsed) + + return { + "genes": genes[:max_genes], + "all_genes_count": len(genes), + "cell_types": sorted(cell_types), + "ligand_receptor_pairs": ligrec, + } + _TOKEN = re.compile(r"[A-Za-z][A-Za-z0-9]{1,14}") _SPLIT = re.compile(r"[^A-Za-z0-9]+")