""" SpatialAgent - Deepagents-based implementation. Refactored SpatialAgent using the deepagents framework for: - Built-in task planning (write_todos) - Filesystem operations - Sub-agent delegation - Context management - Better scalability and maintainability """ from typing import Annotated, List, Dict, Any, TypedDict, Literal, Optional import os, re, uuid, logging from dataclasses import dataclass, field from deepagents import create_deep_agent from deepagents.backends.filesystem import FilesystemBackend from deepagents.backends.state import StateBackend from deepagents.middleware.skills import SkillsMiddleware from ..hooks_deep import HooksMiddleware from langchain_core.messages import BaseMessage, HumanMessage from langchain_core.tools import BaseTool from .make_prompt import AgentPrompts from .make_llm import make_llm, DEFAULT_CLAUDE_MODEL from .utils import load_all_tools @dataclass class SpatialAgentConfig: """Configuration for SpatialAgent using deepagents.""" llm: Any = None data_path: str = "./data" save_path: str = "./experiments" tool_retrieval: bool = False tool_retrieval_method: str = "llm" min_tools: int = 5 max_tools: int = 20 skill_retrieval: bool = True num_skills: int = 1 auto_interpret_figures: bool = True act_timeout: int = 1800 web_search_model: str = "gemini-3-flash-preview" backend_type: Literal["filesystem", "state"] = "filesystem" class SpatialAgent: """SpatialAgent refactored using deepagents framework.""" def __init__(self, config: SpatialAgentConfig = None, **kwargs): """ Initialize SpatialAgent with deepagents. Args: config: SpatialAgentConfig instance with configuration options **kwargs: Additional configuration options (for backward compatibility) """ if config is None: config = SpatialAgentConfig() for key, value in kwargs.items(): if hasattr(config, key): setattr(config, key, value) self.config = config if config.llm is None: print(f"No LLM provided, using default: {DEFAULT_CLAUDE_MODEL}", flush=True) config.llm = make_llm(DEFAULT_CLAUDE_MODEL) self.llm = config.llm self.tool_retrieval = config.tool_retrieval self.min_tools = config.min_tools self.max_tools = config.max_tools self.auto_interpret_figures = config.auto_interpret_figures self.act_timeout = config.act_timeout from . import set_agent_model model_name = None for attr in ['deployment_name', 'model_id', 'model_name', 'model']: val = getattr(config.llm, attr, None) if val and isinstance(val, str): model_name = val break if not model_name: model_name = "unknown" set_agent_model(model_name, config.llm) self.web_search_model = config.web_search_model if config.web_search_model else model_name self.observation_log = [] self._observation_log_path = os.path.join(config.save_path, "observation_log.jsonl") self.conversation_history = {} self.default_thread_id = str(uuid.uuid4()) data_path = os.path.abspath(config.data_path) save_path = os.path.abspath(config.save_path) self.save_path = save_path self.data_path = data_path print(f"Data path: {data_path}", flush=True) print(f"Save path: {save_path}", flush=True) os.makedirs(save_path, exist_ok=True) os.makedirs(data_path, exist_ok=True) self._build_backend() self._build_skills() self._load_tools() self._build_system_prompt() self._build_agent() def _load_tools(self): """Load all spatial transcriptomics tools.""" print("Loading tools from tool modules...", flush=True) self.tools = load_all_tools(save_path=self.save_path, data_path=self.data_path) print(f"Loaded {len(self.tools)} tools", flush=True) if self.config.skill_retrieval and hasattr(self, 'skills_metadata') and self.skills_metadata: self._add_skill_tool() def _add_skill_tool(self): """Add a use_skill tool to read skill details.""" from langchain_core.tools import tool skills_meta = self.skills_metadata @tool def use_skill(skill_name: str) -> str: """Read the full instructions for a spatial transcriptomics analysis skill. Args: skill_name: Name of the skill to read (e.g., "liana-analysis", "spatial-deconvolution") Returns: Full skill documentation with workflow steps, best practices, and examples. """ for skill in skills_meta: if skill['name'] == skill_name or skill['name'].replace('-', '_') == skill_name: try: with open(skill['path'], 'r') as f: return f.read() except Exception as e: return f"Error reading skill: {e}" available = [s['name'] for s in skills_meta] return f"Skill '{skill_name}' not found. Available skills: {', '.join(available)}" self.tools.append(use_skill) print(f"Added use_skill tool for {len(self.skills_metadata)} skills", flush=True) def _build_system_prompt(self): """Build system prompt with spatial transcriptomics domain knowledge.""" tool_list = [] for tool in self.tools: tool_name = getattr(tool, 'name', tool.__class__.__name__) tool_desc = getattr(tool, 'description', 'No description') tool_list.append((tool_name, tool_desc)) tool_names = [t[0] for t in tool_list] skills_section = "" if hasattr(self, 'skills_metadata') and self.skills_metadata: skill_list_str = "\n".join([f"- **{s['name']}**: {s['description'][:80]}" for s in self.skills_metadata]) skills_section = f""" ## SPECIALIZED SPATIAL TRANSCRIPTOMICS SKILLS You also have access to {len(self.skills_metadata)} curated workflow skills. Use the `use_skill` tool to read full instructions. Available skills: {skill_list_str} ## How to Use Skills 1. When a user's task matches a skill's domain, call `use_skill(skill_name="...")` to get detailed workflow instructions 2. Follow the skill's step-by-step guidance 3. Use the specialized tools listed above to execute each step ## TOOLS vs SKILLS - **TOOLS (工具)**: Individual callable functions (search_panglao, liana_inference, execute_python, etc.) - 72 total - **SKILLS (技能)**: Curated workflow guides that tell you WHICH tools to use and WHEN - {len(self.skills_metadata)} total - Use `use_skill` to read skill details, then use tools to execute """ self.system_prompt = f"""You are a helpful assistant specialized in spatial transcriptomics analysis. ## CONVERSATION CONTEXT IMPORTANT: You HAVE FULL ACCESS to conversation history. ALWAYS read and remember ALL previous messages in the conversation. Use this context to provide personalized responses. You can and should reference information from earlier messages. ## COMPLEX TASK PLANNING When you receive a complex task that requires multiple steps or tool calls, YOU MUST FIRST OUTPUT A PLAN (todo list) before executing any tools. This helps organize the workflow and ensures all necessary steps are completed. ### HOW TO DETERMINE IF A TASK IS COMPLEX: - Simple tasks: Can be answered in 1-2 sentences or with a single tool call (e.g., "What is my name?", "Query disease genes for prostate cancer") - Complex tasks: Require multiple steps, multiple tool calls, or synthesis of information (e.g., "Design a 50-gene panel", "Analyze spatial transcriptomics data", "Generate a comprehensive report") ### PLAN FORMAT (MUST FOLLOW): For complex tasks, output your plan FIRST in this format: ```plan ## Task: [Task Name] ### Steps: 1. [Step 1 description] - [Tool to use if applicable] 2. [Step 2 description] - [Tool to use if applicable] 3. [Step 3 description] - [Tool to use if applicable] ... ### Expected Output: [Brief description of what the final output will include] ``` After outputting the plan, you can start executing the steps one by one using tool calls. ## TOOL CALLING FORMAT (MUST FOLLOW) When you need to call a tool, output ONLY a JSON object with "name" and "arguments" fields. DO NOT use markdown code blocks. EXACT FORMAT: {{ "name": "TOOL_NAME", "arguments": {{ "PARAM1": "VALUE1", "PARAM2": VALUE2 }} }} ## CRITICAL RULES - NEVER use ```json or any markdown code blocks - NEVER use tags - ALWAYS use plain JSON format for tool calls - ALWAYS include both "name" and "arguments" fields - Parameter names must match exactly (e.g., "disease" not "diseases") - Execute ONE tool at a time, wait for results, then continue - ALWAYS remember and use conversation history for context - NEVER say you don't have access to conversation history - you ALWAYS have access - FOR SIMPLE QUESTIONS that can be answered from conversation history alone (like "what is my name?", "what did we talk about?"), DO NOT CALL TOOLS. Answer directly using the conversation history. - FOR COMPLEX TASKS, ALWAYS OUTPUT A PLAN FIRST before executing tools. ## TOOLS AVAILABLE Database: search_panglao, search_cellmarker2, search_czi_datasets, query_tissue_expression, query_disease_genes Literature: query_pubmed, search_semantic_scholar, web_search Analytics: liana_inference, squidpy_ligrec, cellphonedb_analysis, tangram_map_cells, spagcn_clustering, scanpy_score_genes Interpretation: annotate_cell_types, annotate_tissue_niches Coding: execute_python, execute_bash Subagent: report_subagent, verification_subagent ## WHEN TO USE SKILLS Use `use_skill(skill_name="...")` for complex workflows. For simple queries, call tools directly. """ def _build_backend(self): """Build the filesystem backend for the agent.""" if self.config.backend_type == "filesystem": self.backend = FilesystemBackend( root_dir=self.save_path, virtual_mode=True, max_file_size_mb=50 ) print(f"Using FilesystemBackend with root_dir: {self.save_path}", flush=True) else: self.backend = StateBackend() print("Using StateBackend (ephemeral storage)", flush=True) def _build_skills(self): """Load skill templates for common spatial transcriptomics workflows.""" if self.config.skill_retrieval: skills_source_dir = os.path.join(os.path.dirname(__file__), '..', 'skill') skills_source_dir = os.path.abspath(skills_source_dir) skills_target_dir = os.path.join(self.save_path, 'skills') if os.path.exists(skills_source_dir): self._copy_skills(skills_source_dir, skills_target_dir) self.skills_path = skills_target_dir print(f"Skills loaded from: {skills_source_dir}", flush=True) print(f"Skills copied to: {skills_target_dir}", flush=True) self._load_skills_metadata() else: self.skills_path = None self.skills_metadata = [] print("Skills source directory not found, skill retrieval disabled", flush=True) else: self.skills_path = None self.skills_metadata = [] def _load_skills_metadata(self): """Load skill metadata from skill directories.""" import re self.skills_metadata = [] if not self.skills_path or not os.path.exists(self.skills_path): return for item in sorted(os.listdir(self.skills_path)): item_path = os.path.join(self.skills_path, item) if not os.path.isdir(item_path): continue skill_md = os.path.join(item_path, 'SKILL.md') if not os.path.exists(skill_md): continue with open(skill_md, 'r') as f: content = f.read() name = item description = "" frontmatter_match = re.match(r'^---\n(.*?)\n---', content, re.DOTALL) if frontmatter_match: frontmatter = frontmatter_match.group(1) for line in frontmatter.split('\n'): if line.startswith('name:'): name = line.split(':', 1)[1].strip() elif line.startswith('description:'): description = line.split(':', 1)[1].strip() if not description: lines = content.strip().split('\n') for line in lines: if line.startswith('# '): continue if line.strip(): description = line.strip() break self.skills_metadata.append({ 'name': name, 'description': description, 'path': skill_md }) print(f"Loaded {len(self.skills_metadata)} skills", flush=True) def _copy_skills(self, source_dir: str, target_dir: str): """Copy skill directories from source to target.""" import shutil os.makedirs(target_dir, exist_ok=True) for item in os.listdir(source_dir): item_path = os.path.join(source_dir, item) if os.path.isdir(item_path): target_item = os.path.join(target_dir, item) if os.path.exists(target_item): shutil.rmtree(target_item) shutil.copytree(item_path, target_item) def _build_agent(self): """Build the deep agent using langgraph with proper tool binding and memory.""" print("Building deep agent...", flush=True) from langchain_core.messages import SystemMessage from langgraph.graph import StateGraph, END from langgraph.prebuilt import ToolNode from langgraph.checkpoint.memory import MemorySaver def should_continue(state): messages = state["messages"] last_message = messages[-1] if hasattr(last_message, 'tool_calls') and last_message.tool_calls: return "tools" return END def call_model(state): messages = state["messages"] system_message = SystemMessage(content=self.system_prompt) full_messages = [system_message] + messages content_to_send = "" for msg in full_messages: if isinstance(msg, SystemMessage): content_to_send += f"System: {msg.content}\n" elif hasattr(msg, 'content'): role = "User" if isinstance(msg, HumanMessage) else "Assistant" content_to_send += f"{role}: {msg.content}\n" response = self.llm.invoke(content_to_send) response_content = response.content if hasattr(response, 'content') else str(response) import json import re import uuid tool_calls = [] json_match = re.search(r'\{[\s\S]*"name":\s*"[^"]+"[\s\S]*\}', response_content) if json_match: try: tool_call_json = json.loads(json_match.group()) if "name" in tool_call_json and "arguments" in tool_call_json: tool_calls = [{ "id": str(uuid.uuid4()), "name": tool_call_json["name"], "args": tool_call_json["arguments"] }] except json.JSONDecodeError: pass from langchain_core.messages import AIMessage if tool_calls: ai_response = AIMessage(content=response_content, tool_calls=tool_calls) else: ai_response = AIMessage(content=response_content) return {"messages": messages + [ai_response]} tool_node = ToolNode(self.tools) workflow = StateGraph(dict) workflow.add_node("agent", call_model) workflow.add_node("tools", tool_node) workflow.set_entry_point("agent") workflow.add_conditional_edges("agent", should_continue) workflow.add_edge("tools", "agent") memory_saver = MemorySaver() self.agent = workflow.compile(checkpointer=memory_saver) print("Deep agent built successfully with memory!", flush=True) def _log_observation(self, step_number: int, code: str, result: str, figure_interpretations: str = ""): """Log an observation for later report generation.""" import json from datetime import datetime entry = { "step": step_number, "timestamp": datetime.now().isoformat(), "code_snippet": code, "result_summary": result, "figure_interpretations": figure_interpretations, } self.observation_log.append(entry) try: with open(self._observation_log_path, 'a') as f: f.write(json.dumps(entry) + '\n') except Exception as e: print(f"Warning: Could not write to observation log: {e}") def _display_figures(self, code_context: str = "", user_query: str = "") -> str: """Display any new image files and optionally interpret them using vision LLM.""" interpretations = [] try: from spatialagent.tool.coding import get_new_image_files image_files = get_new_image_files() if not image_files: return "" try: from IPython.display import display, Image, SVG import os print(f"📊 Displaying {len(image_files)} figure(s)...") for img_path in image_files: if not os.path.exists(img_path): print(f"⚠️ File not found: {img_path}") continue ext = os.path.splitext(img_path)[1].lower() if ext == '.svg': display(SVG(filename=img_path)) elif ext in ('.png', '.jpg', '.jpeg'): display(Image(filename=img_path)) elif ext == '.pdf': print(f"📄 Created: {os.path.basename(img_path)}") except ImportError: import os print(f"[{len(image_files)} figure(s) created: {', '.join(os.path.basename(f) for f in image_files)}]") if self.auto_interpret_figures and image_files: print(f"🔍 Interpreting {len(image_files)} figure(s)...") from spatialagent.tool.interpretation import interpret_figure for img_path in image_files: import os if not os.path.exists(img_path): continue ext = os.path.splitext(img_path)[1].lower() if ext == '.pdf': continue try: context = self._infer_figure_context(code_context, img_path, user_query) interpretation = interpret_figure.invoke({ "image_path": img_path, "context": context, "analysis_focus": "general" }) fig_name = os.path.basename(img_path) interpretations.append(f"\n### Figure Interpretation: {fig_name}\n{interpretation}") except Exception as e: print(f"⚠️ Could not interpret {os.path.basename(img_path)}: {e}") except Exception as e: print(f"⚠️ Error displaying/interpreting figures: {e}") return "\n".join(interpretations) if interpretations else "" def _infer_figure_context(self, code: str, img_path: str, user_query: str = "") -> str: """Infer the context/type of a figure from the code that generated it.""" import os fig_name = os.path.basename(img_path) context_parts = [f"Figure: {fig_name}"] code_lower = code.lower() plot_types = [] if "umap" in code_lower: plot_types.append("UMAP dimensionality reduction") if "tsne" in code_lower or "t-sne" in code_lower: plot_types.append("t-SNE dimensionality reduction") if "pca" in code_lower and "plot" in code_lower: plot_types.append("PCA plot") if "sc.pl.spatial" in code_lower or "sq.pl.spatial" in code_lower or "spatial_scatter" in code_lower: plot_types.append("Spatial plot showing tissue coordinates") if "heatmap" in code_lower or "sns.heatmap" in code_lower or "clustermap" in code_lower: plot_types.append("Heatmap visualization") if "violin" in code_lower: plot_types.append("Violin plot") if "dotplot" in code_lower or "dot_plot" in code_lower: plot_types.append("Dot plot") if "stacked_violin" in code_lower: plot_types.append("Stacked violin plot") if "matrixplot" in code_lower: plot_types.append("Matrix plot") if "rank_genes" in code_lower: plot_types.append("Ranked genes plot") if "barplot" in code_lower or "bar(" in code_lower or "barh(" in code_lower: plot_types.append("Bar plot") if "boxplot" in code_lower: plot_types.append("Box plot") if "scatter" in code_lower and "spatial" not in code_lower: plot_types.append("Scatter plot") if plot_types: context_parts.append(" + ".join(plot_types)) color_by = [] if "cell_type" in code_lower or "celltype" in code_lower or "tier3" in code_lower: color_by.append("cell type") if "leiden" in code_lower: color_by.append("Leiden clusters") if "louvain" in code_lower: color_by.append("Leiden clusters") if "batch" in code_lower or "sample" in code_lower: color_by.append("batch/sample") if "condition" in code_lower or "sample_type" in code_lower: color_by.append("condition/disease stage") if "leiden_neigh" in code_lower or "neighborhood" in code_lower or "neigh" in code_lower: color_by.append("spatial neighborhood") if "niche" in code_lower: color_by.append("tissue niche") if color_by: context_parts.append(f"colored/grouped by: {', '.join(color_by)}") title_patterns = [ r'plt\.title\s*\(\s*[\'"]([^\'"]+)[\'"]', r'\.set_title\s*\(\s*[\'"]([^\'"]+)[\'"]', r'title\s*=\s*[\'"]([^\'"]+)[\'"]', ] for pattern in title_patterns: match = re.search(pattern, code) if match: context_parts.append(f"Title: {match.group(1)}") break gene_patterns = [ r'var_names\s*=\s*\[([^\]]+)\]', r'genes\s*=\s*\[([^\]]+)\]', r"color\s*=\s*['\"]([A-Z][A-Z0-9]+)['\"]", ] for pattern in gene_patterns: match = re.search(pattern, code, re.IGNORECASE) if match: genes = match.group(1).strip() if len(genes) < 200: context_parts.append(f"Genes: {genes}") break comment_pattern = r'#\s*(.+?)$' comments = re.findall(comment_pattern, code, re.MULTILINE) relevant_comments = [c.strip() for c in comments if len(c.strip()) > 10 and len(c.strip()) < 100] if relevant_comments: context_parts.append(f"Code comments: {'; '.join(relevant_comments[:2])}") if "comparison" in code_lower or "vs" in code_lower or "versus" in code_lower: context_parts.append("Comparative analysis") if "composition" in code_lower: context_parts.append("Composition analysis") if "proportion" in code_lower or "percentage" in code_lower: context_parts.append("Proportion/percentage analysis") if "dynamics" in code_lower or "trajectory" in code_lower: context_parts.append("Dynamics/trajectory analysis") if "interaction" in code_lower: context_parts.append("Cell-cell interaction analysis") if user_query: query_truncated = user_query[:500] if len(user_query) > 500 else user_query context_parts.append(f"Biological context: {query_truncated}") return " | ".join(context_parts) def run(self, user_query: str, config: Dict[str, Any] = None) -> Dict[str, Any]: """ Run the agent with a user query. Args: user_query: The user's task/question config: Optional configuration dict Returns: Final agent state """ if config is None: config = {"recursion_limit": 50} elif "recursion_limit" not in config: config["recursion_limit"] = 50 thread_id = config.get("thread_id", self.default_thread_id) langgraph_config = { "recursion_limit": config.get("recursion_limit", 50), "configurable": { "thread_id": thread_id } } import sys print(f"\033[1m\033[0m\n{user_query.strip()}\n\033[1m\033[0m\n") sys.stdout.flush() try: existing_messages = self.conversation_history.get(thread_id, []) if existing_messages: print(f"✅ Restored {len(existing_messages)} messages from memory for thread: {thread_id}", flush=True) else: print(f"🔄 Starting new conversation thread: {thread_id}", flush=True) initial_state = { "messages": existing_messages + [HumanMessage(content=user_query)], } final_state = None step_count = 0 all_messages_accumulated = [] for state_update in self.agent.stream(initial_state, stream_mode="values", config=langgraph_config): step_count += 1 messages = state_update.get("messages", []) all_messages_accumulated = messages logging.debug(f"Stream step {step_count}: {len(messages)} messages") for msg in messages[-1:]: self._print_message(msg) final_state = state_update msg_content = messages[-1].content if messages and hasattr(messages[-1], 'content') else "" if isinstance(msg_content, str) and "" in msg_content: break if all_messages_accumulated: self.conversation_history[thread_id] = all_messages_accumulated logging.debug(f"Stream loop ended after {step_count} steps") return {"messages": all_messages_accumulated} except Exception as e: print(f"Error: {e}", flush=True) raise def _print_message(self, message: BaseMessage): """Print a message with appropriate formatting.""" import sys if isinstance(message, HumanMessage): return msg = message.content if not msg or not isinstance(msg, str): return msg_stripped = msg.strip() if not msg_stripped or msg_stripped == "[]" or msg_stripped == "{}": return if msg_stripped.startswith("[System]"): return def format_tag(text, tag, color_code): pattern = rf"<{tag}>(.*?)" def replacer(match): content = match.group(1).strip() return f"{color_code}<{tag}>\033[0m\n{content}\n{color_code}\033[0m" return re.sub(pattern, replacer, text, flags=re.DOTALL) conclude_match = re.search(r"(.*?)", msg, re.DOTALL) if conclude_match: try: from rich.console import Console from rich.markdown import Markdown from rich.theme import Theme custom_theme = Theme({ "markdown.code": "bold cyan", "markdown.code_block": "cyan on grey93", }) console = Console(theme=custom_theme, force_terminal=True) conclude_content = conclude_match.group(1).strip() if conclude_match.start() > 0: pre_conclude = msg[:conclude_match.start()].strip() pre_conclude = format_tag(pre_conclude, "act", "\033[91m") pre_conclude = format_tag(pre_conclude, "observation", "\033[94m") print(pre_conclude) print() sys.stdout.flush() print("\033[1m\033[0m") md = Markdown(conclude_content, code_theme="github-light", inline_code_theme="cyan") console.print(md) print("\033[1m\033[0m") print() sys.stdout.flush() except ImportError: display_msg = format_tag(msg_stripped, "act", "\033[91m") display_msg = format_tag(display_msg, "observation", "\033[94m") display_msg = format_tag(display_msg, "conclude", "\033[1m") print(display_msg) print() sys.stdout.flush() else: display_msg = format_tag(msg_stripped, "act", "\033[91m") display_msg = format_tag(display_msg, "observation", "\033[94m") print(display_msg) print() sys.stdout.flush()