单细胞基因组分析代理

┌─────────────────────────────────────────────────────────────────┐
│                      Pydantic AI Agent                          │
│                    (Orchestration Layer)                        │
└─────────────────────────┬───────────────────────────────────────┘
                          │ MCP Protocol
                          ▼
┌─────────────────────────────────────────────────────────────────┐
│                    FastMCP Server                               │
│  ┌─────────────┐  ┌─────────────┐  ┌─────────────────────────┐ │
│  │ Cell Type   │  │ Batch       │  │ Gene Network            │ │
│  │ Annotation  │  │ Integration │  │ Inference               │ │
│  └─────────────┘  └─────────────┘  └─────────────────────────┘ │
└─────────────────────────┬───────────────────────────────────────┘
                          │
                          ▼
┌─────────────────────────────────────────────────────────────────┐
│                   Temporal Workflows                            │
│              (Retry Logic & Orchestration)                      │
└─────────────────────────┬───────────────────────────────────────┘
                          │
                          ▼
┌─────────────────────────────────────────────────────────────────┐
│                      scGPT Model                                │
│              (GPU-accelerated inference)                        │
└─────────────────────────────────────────────────────────────────┘

@mcp.tool
async def annotate_cell_types(
    expression_data: str,  # Path to h5ad file or CSV
    reference_dataset: str = "cellxgene",
    batch_size: int = 64
) -> dict:
    """
    Annotate cell types from single-cell RNA-seq data.

    Returns predicted cell types with confidence scores.
    """

@mcp.tool
async def integrate_batches(
    dataset_paths: list[str],  # Paths to multiple h5ad files
    batch_key: str = "batch",
    n_hvg: int = 2000
) -> dict:
    """
    Integrate multiple single-cell datasets with batch correction.

    Returns path to integrated dataset and integration metrics.
    """

@mcp.tool
async def get_gene_embeddings(
    gene_list: list[str],
    model_checkpoint: str = "scGPT_human"
) -> dict:
    """
    Get scGPT embeddings for a list of genes.

    Returns gene embeddings that can be used for similarity analysis.
    """

# Clone the repository
git clone <repo-url>
cd agentic_scgpt

# Create virtual environment and install dependencies
uv sync

# Download scGPT checkpoints
python -m scripts.download_checkpoints

# Start Temporal (if using retry logic)
temporal server start-dev

# Run the MCP server
uv run python -m src.server

agentic_scgpt/
├── src/
│   ├── __init__.py
│   ├── server.py           # FastMCP server with scGPT tools
│   ├── agent.py            # Pydantic AI agent
│   ├── tools/
│   │   ├── __init__.py
│   │   ├── annotate.py     # Cell type annotation tool
│   │   ├── integrate.py    # Batch integration tool
│   │   └── embeddings.py   # Gene embeddings tool
│   ├── workflows/
│   │   ├── __init__.py
│   │   └── activities.py   # Temporal activities
│   └── models/
│       ├── __init__.py
│       └── schemas.py      # Pydantic schemas for tool I/O
├── scripts/
│   ├── check_gpu.py        # GPU diagnostics
│   └── download_checkpoints.py
├── tests/
├── checkpoints/            # scGPT model weights
├── pyproject.toml
└── README.md

# src/server.py
from fastmcp import FastMCP

mcp = FastMCP("scGPT Tools")

@mcp.tool
async def annotate_cell_types(expression_data: str) -> dict:
    """Annotate cell types from scRNA-seq data."""
    # Implementation here
    pass

if __name__ == "__main__":
    mcp.run(transport="streamable-http", port=8000)

# src/agent.py
from pydantic_ai import Agent
from pydantic_ai.mcp import MCPServerStreamableHTTP

server = MCPServerStreamableHTTP("http://localhost:8000/mcp")

agent = Agent(
    "anthropic:claude-sonnet-4-20250514",
    toolsets=[server],
    system_prompt="""You are a single-cell genomics expert assistant.
    You help researchers analyze scRNA-seq data using scGPT tools.
    Always explain your analysis steps and interpret results."""
)

async def main():
    async with agent:
        result = await agent.run(
            "Annotate the cell types in my dataset at data/pbmc.h5ad"
        )
        print(result.output)

# src/workflows/activities.py
from temporalio import activity
from datetime import timedelta

@activity.defn
async def annotate_cells_activity(data_path: str) -> dict:
    """Activity with automatic retry on GPU OOM errors."""
    # scGPT annotation logic
    pass

# In workflow
await workflow.execute_activity(
    annotate_cells_activity,
    "data/sample.h5ad",
    start_to_close_timeout=timedelta(minutes=10),
    retry_policy=RetryPolicy(
        maximum_attempts=3,
        initial_interval=timedelta(seconds=5),
        non_retryable_error_types=["InvalidDataError"]
    )
)

uv run python scripts/check_gpu.py