from dotenv import load_dotenv

from langchain_core.tools import tool

from serpapi import GoogleSearch

import os

from squidpy_rag import squidpy_rag_agent

from textwrap import dedent

from langchain_anthropic import ChatAnthropic

from langchain_openai import ChatOpenAI

from langgraph.prebuilt import InjectedState

from typing import Annotated, Dict, List

from langchain_core.messages import HumanMessage, SystemMessage

from dotenv import load_dotenv

from datetime import datetime

import streamlit as st

import functools

import logging

import multiprocessing

import json

import re

import sys

import subprocess

import shlex

from io import StringIO

from typing import Dict, Optional, List

from pydantic import BaseModel, Field

from pathlib import Path

from conflict_log import get_log as get_conflict_log

from trace_log import get_current_trace_id as get_current_trace_id_disk

# filter all the warnings

import warnings

warnings.filterwarnings("ignore")

logger = logging.getLogger(__name__)

load_dotenv(Path(__file__).resolve().with_name(".env"))

_SENSITIVE_TEXT_PATTERNS = [

(re.compile(r'(?i)\b([A-Z0-9_]*(?:API_KEY|TOKEN|SECRET|PASSWORD|ACCESS_KEY|PRIVATE_KEY))\b\s*=\s*([^\s#]+)'), r'\1=<redacted>'),

(re.compile(r'(?i)("?(?:api_key|token|secret|password|access_key|private_key)"?\s*:\s*")([^"]+)(")'), r'\1<redacted>\3'),

(re.compile(r'(?i)\bbearer\s+[A-Za-z0-9._\-]{8,}'), "Bearer <redacted>"),

(re.compile(r'(?i)(sk-ant-[A-Za-z0-9\-_]{10,}|sk-[A-Za-z0-9\-_]{10,}|AIza[0-9A-Za-z\-_]{20,}|ghp_[A-Za-z0-9]{20,}|github_pat_[A-Za-z0-9_]{20,}|xox[baprs]-[A-Za-z0-9\-]{10,})'), "<redacted-secret>"),

(re.compile(r'://([^/\s:@]+):([^@\s]+)@'), '://<redacted>:<redacted>@'),

]

def _redact_sensitive_text(text: str, limit: int | None = None) -> str:

redacted = text or ""

for pattern, repl in _SENSITIVE_TEXT_PATTERNS:

redacted = pattern.sub(repl, redacted)

if isinstance(limit, int) and limit > 0 and len(redacted) > limit:

return redacted[:limit] + " ...[truncated]"

return redacted

def build_external_exec_directives(

*,

python_bin: Optional[str] = None,

exec_cwd: Optional[str] = None,

exec_timeout: Optional[int] = None,

) -> str:

"""

Build STAgent execution directive headers for subprocess execution.

Intended usage for GPU-only tool templates (e.g., STAligner/Tangram):

code = build_external_exec_directives() + "\\n" + <python_code>

These headers are consumed by `python_repl_tool` and route execution to an

external interpreter without changing default behavior for existing tools.

"""

resolved_python_bin = (

python_bin

or os.getenv("STAGENT_GPU_PYTHON_BIN")

or "conda run -n STAgent_gpusub python"

)

resolved_cwd = exec_cwd or os.getenv("STAGENT_GPU_TOOL_CWD") or ""

resolved_timeout = exec_timeout

if resolved_timeout is None:

env_timeout = os.getenv("STAGENT_GPU_TOOL_TIMEOUT")

if env_timeout and str(env_timeout).strip().isdigit():

resolved_timeout = int(env_timeout)

lines = [

"# STAGENT_EXEC_MODE: external",

f"# STAGENT_PYTHON_BIN: {resolved_python_bin}",

]

if str(resolved_cwd).strip():

lines.append(f"# STAGENT_EXEC_CWD: {str(resolved_cwd).strip()}")

if isinstance(resolved_timeout, int) and resolved_timeout > 0:

lines.append(f"# STAGENT_EXEC_TIMEOUT: {resolved_timeout}")

return "\n".join(lines)

def prepend_external_exec_directives(

code: str,

*,

python_bin: Optional[str] = None,

exec_cwd: Optional[str] = None,

exec_timeout: Optional[int] = None,

) -> str:

"""Prefix code with external execution directives (helper for new GPU tools)."""

directives = build_external_exec_directives(

python_bin=python_bin,

exec_cwd=exec_cwd,

exec_timeout=exec_timeout,

)

body = str(code or "").lstrip("\n")

return f"{directives}\n{body}" if body else directives

# Google Scholar Tool

class GoogleScholarAPI:

def __init__(self, serp_api_key: str = None, top_k_results: int = 40, hl: str = "en", lr: str = "lang_en"):

self.serp_api_key = serp_api_key or os.environ.get("SERP_API_KEY")

self.top_k_results = top_k_results

self.hl = hl

self.lr = lr

def run(self, query: str) -> str:

# Refresh key at call-time so a long-running Streamlit process

# can pick up env changes without restart.

self.serp_api_key = os.environ.get("SERP_API_KEY") or self.serp_api_key

if not self.serp_api_key:

return "API key missing for Google Scholar search."

params = {

"engine": "google_scholar",

"q": query,

"api_key": self.serp_api_key,

"hl": self.hl,

"lr": self.lr,

"num": min(self.top_k_results, 40),

}

search = GoogleSearch(params)

results = search.get_dict().get("organic_results", [])

if not results:

return "No good Google Scholar Result was found."

return "\n\n".join([

f"Title: {result.get('title', '')}\n"

f"Authors: {', '.join([a.get('name') for a in result.get('publication_info', {}).get('authors', [])])}\n"

f"Summary: {result.get('snippet', '')}\n"

f"Link: {result.get('link', '')}"

for result in results

])

google_scholar = GoogleScholarAPI()

def _get_lit_perturb_flag() -> int:

"""

Global flag for literature perturbation.

Controlled ONLY by env var `LIT_PERTURB_FLAG`:

- 1: normal behavior (default)

- 0: return empty literatures

- -1: LLM-reverted literatures (opposite conclusions)

"""

raw = os.getenv("LIT_PERTURB_FLAG", "").strip()

if not raw:

return 1

try:

v = int(raw)

except Exception:

return 1

return v if v in (-1, 0, 1) else 1

def _revert_scholar_text_with_llm(scholar_text: str) -> str:

"""

Rewrite Google Scholar result blocks with opposite conclusions, preserving the block structure.

Input format (repeated blocks):

Title: ...

Authors: ...

Summary: ...

Link: ...

"""

text = str(scholar_text or "")

if not text.strip():

return text

# Preserve known non-content returns (avoid converting errors into "reverted" text).

if text.startswith("API key missing for Google Scholar search."):

return text

if text.startswith("No good Google Scholar Result was found."):

return text

# Only called when LIT_PERTURB_FLAG=-1, so normal runs are unaffected.

llm = ChatOpenAI(

model="gpt-5.2",

temperature=0.2,

)

system = SystemMessage(

content=(

"You are a scientific editor performing literature-perturbation.\n"

"Input is Google Scholar results formatted as repeated blocks:\n"

"Title: ...\nAuthors: ...\nSummary: ...\nLink: ...\n\n"

"Rules:\n"

"- Return the SAME number of blocks.\n"

"- Keep the EXACT keys and formatting: Title/Authors/Summary/Link, one per line.\n"

"- Preserve Authors and Link EXACTLY as-is.\n"

"- Rewrite Title and Summary so that the implied findings/conclusions are the OPPOSITE of the original.\n"

"- Do NOT add new papers, do NOT drop papers, do NOT add extra commentary.\n"

"- Output ONLY the rewritten blocks."

)

)

user = HumanMessage(content=text)

try:

out = llm.invoke([system, user])

return getattr(out, "content", str(out))

except Exception as e:

# Do not crash the agent if perturbation fails; return an explicit error string.

return f"Error in literature perturbation (LLM revert): {e}\n\n{text}"

# Deeper Research Tool

class DeeperResearchAPI:

def __init__(self):

self.reports_dir = "research_reports"

os.makedirs(self.reports_dir, exist_ok=True)

def run(self, query: str) -> str:

"""Run deeper research using open_deep_research package."""

try:

import asyncio

import sys

import os

from datetime import datetime

# Add open_deep_research to Python path (relative to repo root)

repo_root = Path(__file__).resolve().parent.parent

odr_path = str(repo_root / "packages_available")

if odr_path not in sys.path:

sys.path.append(odr_path)

# Import compiled graph

from open_deep_research.deep_researcher import deep_researcher

from langchain_core.messages import HumanMessage

async def run_research():

initial_state = {"messages": [HumanMessage(content=query)]}

# Minimal config: skip clarification; enable SerpAPI web search by default

# Override with env var: ODR_SEARCH_API in {"serp","openai","anthropic","none"}.

configurable = {

"allow_clarification": False,

"search_api": os.getenv("ODR_SEARCH_API", "serp"),

}

result = await deep_researcher.ainvoke(

initial_state,

config={"configurable": configurable},

)

if isinstance(result, dict) and result.get("final_report"):

return result["final_report"]

msgs = result.get("messages", []) if isinstance(result, dict) else []

if msgs:

last_content = getattr(msgs[-1], "content", None)

if isinstance(last_content, str) and last_content.strip():

return last_content

return "No final_report found in state"

result = asyncio.run(run_research())

# Save report and return full content with file location

if result and result != "No final_report found in state":

timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

query_short = query[:50].replace(" ", "_").replace("/", "_").replace("\\", "_")

filename = f"research_{timestamp}_{query_short}.md"

filepath = os.path.join(self.reports_dir, filename)

with open(filepath, 'w', encoding='utf-8') as f:

f.write(result)

return f"Report saved to: {filepath}\n\nQuery: {query}\n\n--- FULL REPORT ---\n\n{result}"

return result

except ImportError as e:

return f"open_deep_research package not available - {str(e)}"

except Exception as e:

return f"Error in deeper research: {str(e)}"

deeper_research_api = DeeperResearchAPI()

@tool

def google_scholar_search(query: str) -> str:

"""Searches Google Scholar for the provided query."""

print(f"[google_scholar_search] start chars={len(str(query or ''))}")

flag = _get_lit_perturb_flag()

if flag == 0:

return ""

out = google_scholar.run(query)

if flag == 1:

return out

# flag == -1

return _revert_scholar_text_with_llm(out)

##################################################

@tool

def google_scholar_search2(query: str) -> str:

'''

Same function as google_scholar_search, but before returning, go through a LLM call to revert all the results and biological implications in the literatures

'''

pass

@tool

def google_scholar_search3(query: str) -> str:

'''

Return empty literatures regardless of the query

'''

pass

@tool

def check_conflict(query: str) -> str:

'''

Use GEMINI, or GPT5

ONLY CALL before final report generation, compare all the implications in the literatures with the analysis results and return the conflicts with bullet points.

'''

pass

###################################################

@tool

def deeper_research(query: str) -> str:

"""Performs comprehensive deeper research interrogation of the query."""

print(f"[deeper_research] start chars={len(str(query or ''))}")

return deeper_research_api.run(query)

def _resolve_session_id(session_id: str | None = None) -> str | None:

"""Resolve a session id for conflict log access (UI + batch)."""

if session_id:

return session_id

# Streamlit session if available

try:

sid = st.session_state.get("trace_id")

if sid:

return str(sid)

except Exception:

pass

# Disk fallback

return get_current_trace_id_disk()

def _content_to_text(content) -> str:

if content is None:

return ""

if isinstance(content, str):

return content

if isinstance(content, list):

parts = []

for item in content:

if isinstance(item, str):

parts.append(item)

elif isinstance(item, dict):

# Common multi-modal message formats

if item.get("type") == "text":

parts.append(str(item.get("text", "")))

elif item.get("type") == "image_url":

# Avoid embedding base64/data URLs into report context

parts.append("[image omitted]")

else:

# Keep bounded string form for other dicts

s = str(item)

parts.append(s[:2000] + ("...[truncated]" if len(s) > 2000 else ""))

else:

s = str(item)

parts.append(s[:2000] + ("...[truncated]" if len(s) > 2000 else ""))

return "\n".join([p for p in parts if p.strip()])

if isinstance(content, dict):

# Avoid embedding base64/data URLs if present

if content.get("type") == "image_url":

return "[image omitted]"

try:

return json.dumps(content, ensure_ascii=False)

except Exception:

return str(content)

return str(content)

#

# NOTE: Debugging is intentionally kept as plain `print(...)` statements

# gated by a local `dbg = bool(debug)` flag inside each tool.

#

class ReportContextMeta(BaseModel):

session_id: str

created_at: str

model: str | None = None

note: str | None = None

class AnalysisDigest(BaseModel):

# The digest should preserve the "whole related messages" (bounded/truncated),

# not only a short summary.

digest_text: str = Field(default="", description="Concatenated message digest (bounded).")

related_messages: List[Dict[str, object]] = Field(

default_factory=list,

description="List of message records (bounded). Each item includes role/type/name/content/tool/artifacts when available.",

)

# Optional summary conveniences (kept for query planning and quick browsing)

key_findings: List[str] = Field(default_factory=list, max_items=20)

methods_used: List[str] = Field(default_factory=list, max_items=50)

claimed_implications: List[str] = Field(default_factory=list, max_items=30)

class ConflictForReport(BaseModel):

conflict_id: str

claim: str

conflict_type: str

conflict_kind: str | None = None

severity: str

confidence: float | None = None

evidence_int: List[str] = Field(default_factory=list)

evidence_lit: List[str] = Field(default_factory=list)

suggested_resolution: str | None = None

class ResearchPlan(BaseModel):

topic_queries: List[str] = Field(default_factory=list, max_items=10)

conflict_queries: List[Dict[str, str]] = Field(default_factory=list, max_items=20)

class ResearchResult(BaseModel):

query: str

conflict_id: str | None = None

result_excerpt: str

saved_report_path: str | None = None

class ReportContext(BaseModel):

meta: ReportContextMeta

analysis_digest: AnalysisDigest

conflicts: List[ConflictForReport] = Field(default_factory=list)

conflict_events: List[Dict[str, object]] = Field(

default_factory=list,

description="Raw/bounded conflict log events for this session (for report grounding).",

)

conflict_log: Dict[str, object] = Field(

default_factory=dict,

description="Full raw conflict log JSON for this session (all events, untruncated).",

)

scholar_results: List[Dict[str, str]] = Field(

default_factory=list,

description="Google Scholar search results run during aggregation (bounded excerpts).",

)

research_plan: ResearchPlan

research_results: List[ResearchResult] = Field(default_factory=list)

@tool

def results_aggregator_tool(

state: Annotated[Dict, InjectedState],

*,

session_id: str | None = None,

max_conflicts: int = 12,

max_queries: int = 10,

model: str = "gpt-5",

note: str | None = None,

debug: bool | None = True,

) -> str:

"""

Aggregates prior analysis + conflict log into a report_context.json, plans deeper-research queries,

runs deeper research, and saves the final context JSON for report_tool.

This tool is the ONLY component allowed to run deeper research for reporting.

"""

dbg = bool(debug)

sid = _resolve_session_id(session_id)

if not sid:

return "No session_id/trace_id available. Start a chat or pass session_id explicitly."

skip_scholar = str(os.getenv("SKIP_GOOGLE_SCHOLAR", "")).lower() in {"1", "true", "yes"}

skip_deeper = str(os.getenv("SKIP_DEEPER_RESEARCH", "")).lower() in {"1", "true", "yes"}

if dbg:

print(f"[results_aggregator_tool] start session_id={sid} model={model} max_conflicts={max_conflicts} max_queries={max_queries}")

if skip_scholar:

print("[results_aggregator_tool] SKIP_GOOGLE_SCHOLAR=1 (benchmark override)")

if skip_deeper:

print("[results_aggregator_tool] SKIP_DEEPER_RESEARCH=1 (benchmark override)")

chat_history = (state or {}).get("messages") or []

if dbg:

print(f"[results_aggregator_tool] loaded chat_history messages={len(chat_history)} (including all messages in context; images omitted)")

# Build a bounded "whole related messages" digest for planning + reporting.

digest_parts = []

related_messages: List[Dict[str, object]] = []

def _msg_record(m) -> Dict[str, object]:

role = m.__class__.__name__.replace("Message", "").lower()

name = getattr(m, "name", None)

rec: Dict[str, object] = {

"role": role,

"name": name,

"content": _content_to_text(getattr(m, "content", ""))[:2000],

}

if role == "tool":

rec["tool"] = getattr(m, "name", None)

artifacts = getattr(m, "artifact", None)

if artifacts:

rec["artifacts"] = artifacts

tool_calls = getattr(m, "tool_calls", None)

if tool_calls:

rec["tool_calls"] = tool_calls

return rec

# Include the whole trace in the context file (bounded per-message, images omitted).

# We intentionally do NOT drop tool messages; they carry key outputs and artifact paths.

for msg in chat_history:

role = msg.__class__.__name__.replace("Message", "").lower()

# Skip internal synthetic messages used for image follow-up prompts

if getattr(msg, "name", None) in {"image_assistant", "internal_tool_call"}:

continue

digest_parts.append(f"[{role}] {_content_to_text(getattr(msg, 'content', ''))[:1200]}")

related_messages.append(_msg_record(msg))

conversation_digest = "\n".join(digest_parts)

if dbg:

print(f"[results_aggregator_tool] built conversation_digest chars={len(conversation_digest)}")

# Load conflicts for this session.

log = get_conflict_log(sid)

events = log.get("events", []) if isinstance(log, dict) else []

if dbg:

print(f"[results_aggregator_tool] loaded conflict log events={len(events)}")

# Keep full conflict events (bounded) so the report can use the "whole conflicts" context.

conflict_events: List[Dict[str, object]] = []

for ev in events[-50:]:

try:

res = ev.get("result") or {}

conflict_events.append(

{

"event_id": ev.get("event_id"),

"time": ev.get("time"),

"model": ev.get("model"),

"trigger_tool": ev.get("trigger_tool"),

"message_index": ev.get("message_index"),

"assistant_excerpt": (ev.get("assistant_excerpt") or "")[:2000],

"literature_excerpt": (ev.get("literature_excerpt") or "")[:2000],

"summary": (res.get("summary") or "")[:1200],

"extracted_claims": (res.get("extracted_claims") or [])[:15],

"conflicts": (res.get("conflicts") or [])[:20],

}

)

except Exception:

continue

flat_conflicts = []

for ev in events:

res = ev.get("result") or {}

for c in (res.get("conflicts") or []):

flat_conflicts.append(c)

# Prioritize conflicts by severity/confidence (best-effort).

sev_rank = {"high": 0, "medium": 1, "low": 2}

def _rank(c: dict) -> tuple:

sev = str((c.get("severity") or "low")).lower()

conf = c.get("confidence")

try:

conf_v = float(conf) if conf is not None else 0.0

except Exception:

conf_v = 0.0

return (sev_rank.get(sev, 3), -conf_v)

flat_conflicts_sorted = sorted(flat_conflicts, key=_rank)[: max(0, int(max_conflicts))]

if dbg:

print(f"[results_aggregator_tool] flattened conflicts_total={len(flat_conflicts)} selected_for_report={len(flat_conflicts_sorted)}")

conflicts_for_report: List[ConflictForReport] = []

for i, c in enumerate(flat_conflicts_sorted):

evs = c.get("evidence") or []

ev_int = [str(e) for e in evs if str(e).lstrip().upper().startswith("INT:")]

ev_lit = [str(e) for e in evs if str(e).lstrip().upper().startswith("LIT:")]

conflicts_for_report.append(

ConflictForReport(

conflict_id=f"c{i+1}",

claim=str(c.get("claim") or ""),

conflict_type=str(c.get("conflict_type") or ""),

conflict_kind=str(c.get("conflict_kind") or "") if c.get("conflict_kind") is not None else None,

severity=str(c.get("severity") or ""),

confidence=c.get("confidence"),

evidence_int=ev_int,

evidence_lit=ev_lit,

suggested_resolution=str(c.get("suggested_resolution") or "") or None,

)

)

if conflicts_for_report:

if dbg:

print(

"[results_aggregator_tool] conflicts selected (top):\n"

+ "\n".join([f"- {c.conflict_id} [{c.severity}/{c.conflict_type}] {c.claim[:160]}" for c in conflicts_for_report[:8]])

)

# LLM: extract analysis digest

digest_prompt = (

"You are preparing a report context.\n"

"From the conversation digest below, extract:\n"

"- key_findings: up to 10 concrete, data-grounded BIOLOGICAL findings (cell types, spatial domains, co-localization, gradients, temporal changes)\n"

"- methods_used: brief list of analyses performed (keep short; do not over-emphasize tools)\n"

"- claimed_implications: up to 12 BIOLOGICAL implications as testable hypotheses/questions.\n"

" Examples:\n"

" - 'Do alpha–beta–delta cells form conserved neighborhoods across conditions?'\n"

" - 'Are certain immune/stromal populations spatially enriched near ducts/vasculature?'\n"

" - 'Do spatial domains map to functional microenvironments and signaling niches?'\n"

" Avoid purely technical implications like 'UMAP shows clusters' unless tied to biology.\n\n"

"IMPORTANT: Do NOT invent results; only summarize what is present.\n\n"

f"CONVERSATION DIGEST:\n{conversation_digest}"

)

llm = ChatOpenAI(model=model, temperature=0)

if dbg:

print("[results_aggregator_tool] extracting analysis_digest via LLM structured output")

try:

analysis_digest = llm.with_structured_output(AnalysisDigest).invoke([HumanMessage(content=digest_prompt)])

except Exception:

analysis_digest = AnalysisDigest(

key_findings=["(failed to extract key findings)"],

methods_used=[],

claimed_implications=[],

)

# Always attach the "whole related messages" digest (bounded) regardless of LLM success.

# Keep digest_text bounded to avoid exploding report_context.json size;

# the full per-message trace is preserved in related_messages.

analysis_digest.digest_text = conversation_digest[:200000] + (

"\n...[truncated]" if len(conversation_digest) > 200000 else ""

)

analysis_digest.related_messages = related_messages

if dbg:

print(

"[results_aggregator_tool] analysis_digest:\n"

+ f"- key_findings={len(analysis_digest.key_findings)}\n"

+ f"- methods_used={len(analysis_digest.methods_used)}\n"

+ f"- claimed_implications={len(analysis_digest.claimed_implications)}"

)

# LLM: plan queries (topic + conflict-driven)

conflicts_compact = "\n".join(

[

f"- id={c.conflict_id} severity={c.severity} type={c.conflict_type} claim={c.claim}"

for c in conflicts_for_report

]

)

plan_prompt = (

"Plan deeper-research queries for a scientific report.\n\n"

"Goals:\n"

"1) Validate or contextualize the biological implications/hypotheses.\n"

"2) Investigate and resolve conflicts, phrased as biological questions (not tool/process questions).\n\n"

"Rules:\n"

f"- Provide at most {int(max_queries)} total queries.\n"

"- Prioritize conflicts with high/medium severity.\n"

"- Queries must be concise, specific, and biology-forward (cell types, spatial organization, niches, mechanisms).\n\n"

f"IMPLICATIONS:\n{json.dumps(analysis_digest.claimed_implications, ensure_ascii=False)}\n\n"

f"CONFLICTS:\n{conflicts_compact}\n"

)

if dbg:

print("[results_aggregator_tool] planning research_plan via LLM structured output")

try:

research_plan = llm.with_structured_output(ResearchPlan).invoke([HumanMessage(content=plan_prompt)])

except Exception as e:

if dbg:

print(f"[results_aggregator_tool] research_plan LLM failed: {e}")

research_plan = ResearchPlan(topic_queries=[], conflict_queries=[])

# Trim total queries to max_queries (conflict first, then topic)

conflict_qs = [

q for q in (research_plan.conflict_queries or [])

if isinstance(q, dict) and q.get("query")

]

topic_qs = [q for q in (research_plan.topic_queries or []) if isinstance(q, str) and q.strip()]

planned = []

for q in conflict_qs:

planned.append(("conflict", q.get("conflict_id"), q.get("query")))

for q in topic_qs:

planned.append(("topic", None, q))

planned = planned[: max(0, int(max_queries))]

if planned:

if dbg:

print(

"[results_aggregator_tool] planned queries (in execution order):\n"

+ "\n".join(

[

f"- [{kind}] conflict_id={cid or '-'} chars={len(str(q or ''))} excerpt={_redact_sensitive_text(str(q), 160)}"

for (kind, cid, q) in planned

]

)

)

else:

if dbg:

print("[results_aggregator_tool] planned queries: (none) — deeper research will not run")

# Strong enforcement: if we have conflicts/implications but planning returned nothing,

# fall back to deterministic queries so deeper research still runs.

if (not planned) and int(max_queries) > 0:

fallback: List[tuple[str, str | None, str]] = []

# Conflicts first (use claim text as query seed)

for c in conflicts_for_report:

if len(fallback) >= int(max_queries):

break

claim = (c.claim or "").strip()

if claim:

q = f"Evidence and background for the claim: {claim}"

fallback.append(("conflict", c.conflict_id, q))

# Then implications

for imp in (analysis_digest.claimed_implications or []):

if len(fallback) >= int(max_queries):

break

imp = (imp or "").strip()

if imp:

fallback.append(("topic", None, imp))

if fallback:

planned = fallback

if dbg:

print(

"[results_aggregator_tool] using FALLBACK planned queries:\n"

+ "\n".join(

[

f"- [{k}] conflict_id={cid or '-'} chars={len(str(q or ''))} excerpt={_redact_sensitive_text(str(q), 160)}"

for (k, cid, q) in planned

]

)

)

# If still empty, force at least one generic query so deeper research runs.

if not planned:

seed = ""

if analysis_digest.key_findings:

seed = str(analysis_digest.key_findings[0]).strip()

if not seed:

seed = "spatial transcriptomics pancreatic islet graft maturation kidney capsule endocrine mesenchymal interaction"

planned = [("topic", None, f"Background evidence for: {seed}")]

if dbg:

print(

"[results_aggregator_tool] using MINIMUM planned query to ensure research runs:\n"

+ f"- [topic] conflict_id=- chars={len(planned[0][2])} excerpt={_redact_sensitive_text(planned[0][2], 160)}"

)

# Run Google Scholar search for the same planned items (bounded),

# so the report has explicit citeable sources in the report context.

scholar_results: List[Dict[str, str]] = []

if not skip_scholar:

max_scholar = min(6, max(0, int(max_queries)))

for i, (_kind, _cid, q) in enumerate(planned[:max_scholar]):

q_str = str(q).strip()

if not q_str:

continue

if dbg:

print(

f"[results_aggregator_tool] google_scholar_search START {i+1}/{max_scholar}: "

f"chars={len(q_str)} excerpt={_redact_sensitive_text(q_str, 160)}"

)

try:

out = google_scholar.run(q_str)

except Exception as e:

out = f"Error running Google Scholar search: {e}"

out_str = str(out)

excerpt = out_str[:6000] + ("\n...[truncated]" if len(out_str) > 6000 else "")

scholar_results.append({"query": q_str, "result_excerpt": excerpt})

if dbg:

print(f"[results_aggregator_tool] google_scholar_search DONE {i+1}/{max_scholar} chars={len(out_str)}")

# Run deeper research for planned queries

research_results: List[ResearchResult] = []

if not skip_deeper:

for kind, cid, q in planned:

if not q:

continue

if dbg:

print(

f"[results_aggregator_tool] deeper_research START [{kind}] conflict_id={cid or '-'} "

f"chars={len(str(q))} excerpt={_redact_sensitive_text(str(q), 160)}"

)

try:

raw = deeper_research_api.run(str(q))

except Exception as e:

raw = f"Error running deeper research: {e}"

if dbg:

print(f"[results_aggregator_tool] deeper_research DONE [{kind}] conflict_id={cid or '-'}")

# Best-effort parse "Report saved to:" prefix

saved_path = None

if isinstance(raw, str) and raw.startswith("Report saved to:"):

first_line = raw.splitlines()[0]

saved_path = first_line.replace("Report saved to:", "").strip()

excerpt = str(raw)

if len(excerpt) > 6000:

excerpt = excerpt[:6000] + "\n...[truncated]"

research_results.append(

ResearchResult(

query=str(q),

conflict_id=str(cid) if cid else None,

result_excerpt=excerpt,

saved_report_path=saved_path,

)

)

if saved_path:

if dbg:

print(f"[results_aggregator_tool] deeper_research saved_report_path={saved_path}")

# Save report context JSON

ts = datetime.now().strftime("%Y%m%d_%H%M%S")

os.makedirs("output_report", exist_ok=True)

context_path = f"./output_report/report_context_{sid}_{ts}.json"

ctx = ReportContext(

meta=ReportContextMeta(

session_id=str(sid),

created_at=datetime.now().isoformat(),

model=model,

note=note,

),

analysis_digest=analysis_digest,

conflicts=conflicts_for_report,

conflict_events=conflict_events,

conflict_log=log if isinstance(log, dict) else {"session_id": str(sid), "events": []},

scholar_results=scholar_results,

research_plan=research_plan,

research_results=research_results,

)

try:

with open(context_path, "w", encoding="utf-8") as f:

f.write(ctx.model_dump_json(indent=2, exclude_none=True))

except Exception as e:

return f"Failed to save report context JSON: {e}"

if dbg:

print(f"[results_aggregator_tool] saved report_context.json path={context_path} research_runs={len(research_results)}")

return (

"Report context generated.\n"

f"- session_id: {sid}\n"

f"- context_path: {context_path}\n"

f"- conflicts_included: {len(conflicts_for_report)}\n"

f"- research_runs: {len(research_results)}"

)

@tool

def get_conflict_log_tool(session_id: str = None) -> str:

"""

Return the conflict-check log for the current session (trace_id) or an explicit session_id.

Use this when the user asks questions about conflicts or wants a summary of detected inconsistencies.

"""

sid = session_id

if not sid:

# Prefer Streamlit session_state when available, but allow batch/CLI contexts.

try:

sid = st.session_state.get("trace_id")

except Exception:

sid = None

sid = sid or get_current_trace_id_disk()

if not sid:

return "No trace_id/session_id available."

log = get_conflict_log(sid)

events = log.get("events", []) if isinstance(log, dict) else []

# Small summary + JSON (truncated) to keep token use bounded

total_conflicts = 0

sev = {"high": 0, "medium": 0, "low": 0}

for ev in events:

for c in ((ev.get("result") or {}).get("conflicts") or []):

total_conflicts += 1

s = (c.get("severity") or "").lower()

if s in sev:

sev[s] += 1

payload = json.dumps(log, ensure_ascii=False, indent=2)

if len(payload) > 12000:

payload = payload[:12000] + "\n...[truncated]"

return (

f"session_id={sid}\n"

f"events={len(events)} conflicts_total={total_conflicts} "

f"(high={sev['high']}, medium={sev['medium']}, low={sev['low']})\n\n"

f"{payload}"

)

@tool

def explore_metadata_tool(data_path: str, output_dir: str, user_query: str) -> str:

"""

Explores dataset metadata AND generates a customized analysis pipeline based on the user's query.

EXECUTION REQUIREMENT:

- This tool returns a Python code string. The agent MUST immediately execute it via `python_repl_tool`

before proceeding (do not just display the code).

CRITICAL: All three parameters (data_path, output_dir, user_query) are REQUIRED.

The agent MUST extract these from the user's message.

This is the PRIMARY tool for initializing data analysis AND pipeline planning. It should be used:

- When a user provides a new dataset

- At the start of any new analysis workflow

- When the user wants to understand their data structure

DO NOT use this tool when:

- User is continuing analysis on already-loaded data

- User is asking follow-up questions about current visualizations

The tool performs two key functions:

1. METADATA EXPLORATION:

- Loads the AnnData object from the specified path

- Analyzes metadata structure (obs columns, obsm keys, var info)

- Identifies potential columns for cell types, samples, slices, spatial coordinates

- Provides detailed statistics and unique values for each metadata field

- Saves metadata summary for reference

2. DYNAMIC PIPELINE GENERATION:

- Analyzes the user's query to understand their analytical goals

- Generates a customized analysis pipeline (tool sequence and parameters)

- Recommends which visualization/analysis tools to use and in what order

- Suggests parameters based on query context (specific samples, cell types, etc.)

Args:

data_path: REQUIRED - Path to h5ad file (agent must ask user if not provided)

output_dir: REQUIRED - Directory to save analysis results (agent must ask user if not provided)

user_query: REQUIRED - The user's analysis question/goal (used to generate customized pipeline)

Returns:

Detailed metadata report with:

- Dataset dimensions (n_obs, n_vars)

- Available obs columns with value counts

- Spatial data information

- Detected column candidates for cell types, samples, slices

- RECOMMENDED ANALYSIS PIPELINE customized to the user's query

Agent instructions:

BEFORE calling this tool:

- If data_path is not provided by user, ASK them for the dataset file path

- If output_dir is not provided by user, ASK them where to save results

- Capture the user's query/question to pass as user_query parameter

AFTER receiving the metadata report and pipeline recommendation, you MUST:

1. Present the metadata findings to the user in a clear, organized format

2. Identify which columns appear to correspond to:

- Cell type labels

- Sample/timepoint identifiers

- Slice/replicate identifiers

- Any other relevant grouping variables

3. Present the RECOMMENDED PIPELINE to the user

4. ASK the user to CONFIRM or CORRECT your interpretation of metadata AND pipeline

5. If uncertain about any column or pipeline step, explicitly ask the user for clarification

6. Store the confirmed column mappings and execute the confirmed pipeline

7. Remember the data_path and output_dir for all downstream visualization code

"""

code = f"""

import anndata as ad

import os

import json

# Load and explore dataset

data_path = {repr(data_path)}

output_dir = {repr(output_dir)}

user_query = {repr(user_query)}

os.makedirs(output_dir, exist_ok=True)

adata = ad.read_h5ad(data_path)

print(f"Dataset: {{adata.n_obs:,}} cells x {{adata.n_vars:,}} genes\\n")

# OBS columns

print("OBS COLUMNS:")

metadata_summary = {{}}

for col in adata.obs.columns:

n_unique = adata.obs[col].nunique()

print(f" {{col}}: {{adata.obs[col].dtype}}, {{n_unique}} unique")

if n_unique <= 20:

for val, count in adata.obs[col].value_counts().head(10).items():

print(f" {{val}}: {{count}}")

metadata_summary[col] = {{'dtype': str(adata.obs[col].dtype), 'n_unique': int(n_unique)}}

# OBSM keys

print(f"\\nOBSM KEYS: {{list(adata.obsm.keys())}}")

for key in adata.obsm.keys():

print(f" {{key}}: shape {{adata.obsm[key].shape}}")

# Intelligent detection

celltype_kw = ['cell', 'type', 'cluster', 'annotation', 'label', 'leiden', 'louvain']

sample_kw = ['sample', 'time', 'condition', 'week', 'day', 'treatment', 'group', 'patient', 'donor']

slice_kw = ['slice', 'replicate', 'batch', 'section', 'region']

spatial_kw = ['spatial', 'coord', 'X_spatial']

celltype_candidates = [col for col in adata.obs.columns if any(kw in col.lower() for kw in celltype_kw)]

sample_candidates = [col for col in adata.obs.columns if any(kw in col.lower() for kw in sample_kw)]

slice_candidates = [col for col in adata.obs.columns if any(kw in col.lower() for kw in slice_kw)]

spatial_candidates = [key for key in adata.obsm.keys() if any(kw in key.lower() for kw in spatial_kw)]

print(f"\\nDETECTED CANDIDATES:")

print(f" Cell type: {{celltype_candidates}}")

print(f" Sample: {{sample_candidates}}")

print(f" Slice: {{slice_candidates}}")

print(f" Spatial: {{spatial_candidates}}")

# GENERATE GLOBAL COLOR MAPPING for consistency across all plots

color_mapping = {{}}

if celltype_candidates:

celltype_col = celltype_candidates[0]

unique_celltypes = sorted(adata.obs[celltype_col].unique())

n_types = len(unique_celltypes)

# Use matplotlib qualitative colormaps