Tx Gemma MCP

🧬 TxGemma MCP

TxGemma MCP 是一个模型上下文协议（MCP）服务器，它公开了Google DeepMind的TxGemma治疗AI模型，用于代理药物发现工作流程。

特性 • 安装 • 配置 • 码头工人 • 测试 • 建筑

______________________________________________________________________

🚀 特性

• 双模型：用于快速TDC任务的预测模型+用于解释的聊天模型
• 配置驱动：通过以下方式控制模型、工具和行为 config.yaml
• 动态工具生成：根据TDC提示自动生成700多种工具
• 高级工具筛选：按占位符、复杂性和正则表达式模式过滤
• 延迟模型加载：首次使用时加载模型（快速启动）
• 安全强化：输入验证、SQL注入保护、XSS防御
• GPU优化：FP16的高效内存使用
• 双重运输：FastMCP同时支持stdio（MCP）和streamable-http（web API）模式
• 环境超越：用环境变量覆盖配置
• 生产就绪：447+次测试，96%的覆盖率，全面的CI/CD

______________________________________________________________________

📁 建筑

txgemma-mcp/
├── config.yaml               # Main configuration file
├── server.py                 # FastMCP entrypoint
├── scripts/
│   └── analyze_tools.py      # CLI tool for exploring available tools
├── txgemma/
│   ├── __init__.py           # Package exports with lazy loading
│   ├── config.py             # Configuration loader with env overrides
│   ├── model.py              # Predict + Chat model singletons
│   ├── chat_factory.py       # Chat tool registration
│   ├── tool_factory.py       # Auto-generate TDC tools from prompts
│   ├── executor.py           # Execute tool calls with models
│   ├── prompts.py            # Load TDC prompts from HuggingFace
│   ├── validation.py         # Input validation & security (NEW)
│   └── cache_utils.py        # Global state management (NEW)
├── tests/
│   ├── unit/                 # Unit tests (fast, mocked)
│   │   ├── test_validation.py      # 66 tests - Security
│   │   ├── test_executor.py        # 40 tests
│   │   ├── test_cache_utils.py     # 26 tests
│   │   ├── test_config.py          # 50+ tests
│   │   └── test_chat_factory.py    # 25+ tests
│   ├── integration/          # Integration tests (real components)
│   │   ├── test_tool_factory.py    # 52 tests
│   │   ├── test_prompts.py         # 60+ tests
│   │   ├── test_server.py          # 50+ tests (with security!)
│   │   └── test_analyze_tools.py   # 50+ tests - CLI
│   └── gpu/                  # GPU tests (optional)
│       └── test_model.py           # 28 tests
├── docs/
│   ├── TESTING.md            # Comprehensive testing guide
│   └── FILTERING.md          # Tool filtering guide
└── pyproject.toml

关键设计原则

1. 先配置：中的所有运行时设置 config.yaml
2. 双模型：快速预测+对话式解释
3. 智能过滤：仅加载您需要的内容（默认情况下为药物微笑）
4. 懒加载：模型仅在需要时加载
5. 设计安全：输入验证、注入保护、安全执行（）
6. 单例模式：每种模型类型一个实例
7. 环境超越：可以通过env变量重写配置
8. 双重运输：FastMCP提供stdio（MCP）和streamable-http（web API）
9. 测试驱动：96%的覆盖率，全面的测试套件

______________________________________________________________________

🧩 安装

先决条件

• Python 3.11或3.12 （均在CI中测试）
• GPU推荐 （CUDA或MPS）-型号为2B-27B参数
• 紫外线 （包管理器）
• HuggingFace帐户 （用于模型访问）

设置

# 1. Install uv
curl -LsSf https://astral.sh/uv/install.sh | sh

# 2. Clone repository
git clone https://github.com/t-neumann/TxGemma-MCP.git
cd TxGemma-MCP

# 3. Install dependencies
uv sync --all-extras

# 4. Login to HuggingFace
uv run huggingface-cli login

重要：在HuggingFace上接受TxGemma条款：\ https://huggingface.co/google/txgemma-2b-predict

______________________________________________________________________

⚙️ 配置

TxGemma MCP通过以下方式配置 config.yaml默认配置针对开发进行了优化（快速、低VRAM）。

默认配置

# Prediction Model (for TDC tasks)
predict:
  model: "google/txgemma-2b-predict"
  max_new_tokens: 64

# Chat Model (for explanations)
chat:
  model: "google/txgemma-9b-chat"
  max_new_tokens: 100

# Tool Loading
tools:
  prompts:
    filename: "tdc_prompts.json"
    # Prompts are auto-downloaded from predict model repo
  
  # Filter by placeholder 
  filter_placeholder: "Drug SMILES"  # Only load drug tools
  
  # Enable conversational chat tool
  enable_chat: true

配置预设

开发（默认-22GB VRAM）

predict:
  model: "google/txgemma-2b-predict"
chat:
  model: "google/txgemma-9b-chat"
  max_new_tokens: 100
tools:
  filter_placeholder: "Drug SMILES"

生产（36GB VRAM）

predict:
  model: "google/txgemma-9b-predict"
chat:
  model: "google/txgemma-9b-chat"
  max_new_tokens: 200
tools:
  filter_placeholder: "Drug SMILES"
  exclude_name_pattern: "^(ToxCast|Tox21)"  # Exclude ToxCast overload of tools for Agents

研究（54GB+VRAM）

predict:
  model: "google/txgemma-27b-predict"
chat:
  model: "google/txgemma-27b-chat"
  max_new_tokens: 500
tools:
  filter_placeholder: null  # Load all tools

环境变量覆盖

在不编辑文件的情况下覆盖配置：

# Override models
export TXGEMMA_PREDICT_MODEL=google/txgemma-9b-predict
export TXGEMMA_CHAT_MODEL=google/txgemma-27b-chat

# Override chat response length
export TXGEMMA_CHAT_MAX_TOKENS=500

# Load all tools instead of filtering
export TXGEMMA_FILTER_PLACEHOLDER=null

# Exclude tool patterns
export TXGEMMA_EXCLUDE_NAME_PATTERN="^ToxCast"

# Run server
uv run fastmcp run server.py

优先级：环境变量> config.yaml >默认值

可用模型

工具筛选选项

tools:
  # Option 1: Filter by single placeholder (most common)
  filter_placeholder: "Drug SMILES"  # Only drug-development tools
  
  # Option 2: Filter by multiple placeholders
  filter_placeholders: ["Drug SMILES", "Target sequence"]
  match_all: true  # Require ALL placeholders (AND logic)
  
  # Option 3: Fuzzy matching
  filter_placeholder: "sequence"
  exact_match: false  # Matches "Target sequence", "Protein sequence", etc.
  
  # Option 4: Limit complexity
  filter_placeholder: "Drug SMILES"
  max_placeholders: 2  # Only simple tools
  
  # Option 5: Exclude by regex pattern (NEW!)
  filter_placeholder: "Drug SMILES"
  exclude_name_pattern: "^ToxCast"  # Exclude ToxCast tools
  
  # Option 6: Complex combinations
  filter_placeholder: "Drug SMILES"
  max_placeholders: 2
  exclude_name_pattern: "^(ToxCast|Tox21)"
  
  # Option 7: Load all tools (slow, not recommended)
  filter_placeholder: null  # All 700+ tools

为什么要过滤？ 加载所有工具可能需要10-30秒，并且可能会让LLM代理有太多的选择。过滤到药物SMILES涵盖了大多数分子特性预测用例。

有关详细的筛选选项、示例和CLI用法，请参阅 过滤.md

______________________________________________________________________

🔍 探索可用工具

v0.1.1中的新功能: analyze_tools.py CLI用于探索工具目录！

# List all placeholders with usage counts
python scripts/analyze_tools.py --list-placeholders

# Show all Drug SMILES tools
python scripts/analyze_tools.py --placeholder "Drug SMILES"

# Fuzzy search for sequence-related tools
python scripts/analyze_tools.py --placeholder "sequence" --fuzzy

# Show simple tools only (≤2 parameters)
python scripts/analyze_tools.py --simple

# Exclude ToxCast tools
python scripts/analyze_tools.py --exclude "^ToxCast"

# Combine filters: Drug SMILES + simple + no ToxCast
python scripts/analyze_tools.py --placeholder "Drug SMILES" --simple --exclude "^ToxCast"

# Export to JSON
python scripts/analyze_tools.py --json > tools.json

# Show template details
python scripts/analyze_tools.py --template "tdc_ClinTox_predict"

# Get help
python scripts/analyze_tools.py --help

看 docs/FILTERING.md 获取完整的CLI文档和示例。

______________________________________________________________________

🧬 可用工具

预测工具

700+TDC预测工具 对于分子特性。确切的数量和类型取决于您的过滤配置。

默认：与 filter_placeholder: "Drug SMILES"，加载了约400-500个工具（不包括基于蛋白质/序列的工具）。

备注：使用 exclude_name_pattern: "^ToxCast" 以去除大部分监管检测，避免试剂过载。

聊天工具（可配置）

txgemma_chat -关于药物发现的对话问答

注： 默认情况下启用，通过禁用 tools.enable_chat: false 在config.yaml中。

示例查询：

{"question": "Why might aspirin cause stomach bleeding?"}
{"question": "What makes a good blood-brain barrier penetrant drug?"}
{"question": "Explain the mechanism of action for CC(=O)OC1=CC=CC=C1C(=O)O"}

______________________________________________________________________

🐳 Docker部署

构建

docker buildx build --platform linux/amd64 -t tobneu/txgemma-mcp:latest --push .

部署

# Create cache directory
mkdir -p ~/.cache/huggingface

docker run -d --gpus all \
  --restart unless-stopped \
  -e HF_TOKEN=$HF_TOKEN \
  -e HF_HOME=/root/.cache/huggingface \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  -p 8000:8000 \
  tobneu/txgemma-mcp:latest

# Check logs
docker logs -f 

# Verify config
docker logs  2>&1 | grep "configured"

在Docker中覆盖配置

# Override models and filtering at runtime
docker run -d --gpus all \
  -e HF_TOKEN=$HF_TOKEN \
  -e HF_HOME=/root/.cache/huggingface \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  -e TXGEMMA_PREDICT_MODEL=google/txgemma-27b-predict \
  -e TXGEMMA_CHAT_MODEL=google/txgemma-27b-chat \
  -e TXGEMMA_CHAT_MAX_TOKENS=500 \
  -e TXGEMMA_EXCLUDE_NAME_PATTERN="^ToxCast" \
  -p 8000:8000 \
  tobneu/txgemma-mcp:latest

______________________________________________________________________

🧪 示例用法

通过MCP协议

与Claude Desktop、Cline或任何MCP客户端一起使用：

{
  "mcpServers": {
    "txgemma": {
      "command": "docker",
      "args": [
        "run",
        "--rm",
        "-i",
        "--gpus",
        "all",
        "-e",
        "HF_TOKEN",
        "tobneu/txgemma-mcp:latest"
      ]
    }
  }
}

通过HTTP API

# Start server with streamable-http transport
uv run fastmcp run server.py --transport streamable-http

# Use MCP Inspector
npx @modelcontextprotocol/inspector --transport http --server-url http://localhost:8000/mcp

以编程方式

from txgemma import execute_tool, execute_chat

# Predict toxicity
result = execute_tool(
    "tdc_ClinTox_predict",
    {"Drug SMILES": "CC(=O)OC1=CC=CC=C1C(=O)O"}
)
print(f"Toxicity: {result}")

# Get explanation
explanation = execute_chat(
    "Why might aspirin (CC(=O)OC1=CC=CC=C1C(=O)O) cause stomach bleeding?"
)
print(f"Explanation: {explanation}")

______________________________________________________________________

测试

快速开始

# Run all fast tests (no GPU required)
uv run pytest -m "not gpu"

# With coverage report
uv run pytest -m "not gpu" --cov=txgemma --cov-report=html
open htmlcov/index.html

测试套件概述

运行时：约3-4秒（无GPU）

运行特定测试

# By category
pytest tests/unit/ -v              # Unit tests only
pytest tests/integration/ -v       # Integration tests only

# By module
pytest tests/unit/test_validation.py -v
pytest tests/integration/test_server.py -v

# Security tests
pytest -m security -v

# GPU tests (requires GPU)
pytest -m gpu -v

CI/CD

GitHub操作运行：

• 掉毛 （ruff）-Python 3.11和3.12
• 单元测试 -快速、模拟的依赖关系
• 集成测试 -真实组件，无GPU
• GPU测试 （可选）-在自托管EC2运行器上
• 覆盖率报告 -上传到Codecov

获取完整的测试文档，请参阅 测试.md

______________________________________________________________________

🧰 发展

棉绒和格式

# Check linting
uv run ruff check .

# Auto-fix issues
uv run ruff check --fix .

# Format code
uv run ruff format .

代码质量

• 拉夫 用于裁剪和格式化
• 类型提示 贯穿整个代码库
• 安全检查 （exec（）用法已验证）
• 导入排序 和组织
• 文档字符串 适用于所有公共API

______________________________________________________________________

🏗️ 建筑细部

模型加载策略

• 懒加载：模型首先加载 generate() 呼叫
• 单例：每种模型类型一个实例（预测/聊天）
• 配置：模型由以下因素确定 config.yaml 或环境变量
• 设备自动检测：CUDA>MPS>CPU

配置优先

1. 显式参数 （测试/覆盖）
2. 环境变量 (TXGEMMA_*)
3. 配置文件 (config.yaml)
4. 硬编码默认值 （回退）

即时流

Client Request
    ↓
server.py (FastMCP)
    ↓
validation.py (input validation) ← NEW!
    ↓
executor.py (execute_tool or execute_chat)
    ↓
prompts.py (load template) + model.py (generate)
    ↓
Result → Client

安全架构

1. 输入验证 (validation.py):

- SQL注入预防 - 路径遍历保护 - 命令注入预防 - XSS保护 - SMILES字符串验证

1. 服务器安全 (server.py):

- 安全 exec() 使用验证 - 防止代码注入 - 恶意输入拒绝

1. 参数安全 (executor.py):

- 参数名称规范化 - 空白剥离 - 类型验证

所有安全关键代码都有100%的测试覆盖率 🛡️

内存管理

发展（2b+9b）：

• 预测模型：~4GB
• 聊天模式：~18GB
• 总计：约22GB VRAM

产量（9b+9b）：

• 预测模型：~18GB
• 聊天模式：~18GB
• 总计：约36GB VRAM

第一代：

• 模型下载：~10-60秒（一次性）
• 模型加载：~10-30秒
• 生成时间：约1-5秒

后代：

• ~1-2秒（预测）
• 约2-5秒（聊天）

______________________________________________________________________

📝 添加自定义工具

选项1：等待官方更新（推荐）

TxGemma提示由谷歌维护。新任务添加到HuggingFace仓库时会自动出现。

选项2：本地覆盖

创建 custom_prompts.json:

{
  "your_tool_name": {
    "template": "Instruction: Your instruction.\nContext: Background.\nQuestion: {Your Placeholder}?\nAnswer:",
    "metadata": {
      "description": "Tool description",
      "category": "custom"
    }
  }
}

更新 config.yaml:

tools:
  prompts:
    local_override: "/path/to/custom_prompts.json"

该工具自动生成：

• JSON密钥中的名称
• 输入架构来自 {placeholders}
• 元数据描述

______________________________________________________________________

🚀 生产最佳实践

安全

# Use secrets manager for HF_TOKEN
docker run -d --gpus all \
  --restart unless-stopped \
  -e HF_TOKEN=$(aws secretsmanager get-secret-value ...) \
  -e HF_HOME=/root/.cache/huggingface \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  -e TXGEMMA_EXCLUDE_NAME_PATTERN="^(ToxCast|Tox21)" \
  -p 8000:8000 \
  tobneu/txgemma-mcp:latest

扩展

• 使用GPU实例（AWS g5、g4dn、p3）
• 实现高负载请求排队
• 考虑模型服务框架（vLLM、TGI）
• 缓存常用预测
• 使用工具筛选来减少初始化时间

监控

# Check loaded tools
docker logs  2>&1 | grep "Loaded.*tools"

# Check excluded tools
docker logs  2>&1 | grep "Excluded.*matching pattern"

# Verify configuration
docker logs  2>&1 | grep "configured"

______________________________________________________________________

📚 资源

• TxGemma文档
• TxGemma纸业（arXiv）
• 模型上下文协议
• 治疗数据共享
• FastMCP文档

项目文件：

• 测试指导 -全面的测试套件文档
• 过滤指南 -工具过滤选项和示例

______________________________________________________________________

⚠️ 局限性

• 需要GPU：型号需要8-64GB VRAM，具体取决于大小
• 第一负载：初始下载和加载需要时间
• 上下文长度：受模型上下文窗口的限制（约8K令牌）
• 速率限制：HuggingFace Hub有下载限制

______________________________________________________________________

🐛 故障排除

配置未加载

# Check config exists
ls -la config.yaml

# Verify environment variables
printenv | grep TXGEMMA

# Check Docker logs
docker logs  2>&1 | grep -i config

型号不变

# Environment variable names need TXGEMMA_ prefix
export TXGEMMA_PREDICT_MODEL=google/txgemma-9b-predict  # ✅ Correct
export PREDICT_MODEL=google/txgemma-9b-predict          # ❌ Wrong

# Verify config loaded
docker logs  2>&1 | grep "configured"

内存不足

# Use smaller models
export TXGEMMA_PREDICT_MODEL=google/txgemma-2b-predict
export TXGEMMA_CHAT_MODEL=google/txgemma-9b-chat

# Or reduce chat length
export TXGEMMA_CHAT_MAX_TOKENS=100

工具未加载

# Check filter setting
docker logs  2>&1 | grep "filter"

# Check exclusion pattern
docker logs  2>&1 | grep "Excluded"

# Load all tools (slower)
export TXGEMMA_FILTER_PLACEHOLDER=null

工具筛选不起作用

# Verify pattern is correct
python scripts/analyze_tools.py --exclude "^ToxCast" --json | jq '.[].name'

# Check logs for exclusion
docker logs  2>&1 | grep "Excluded.*tools matching pattern"

______________________________________________________________________

📋 更新日志

v0.1.1（2026-02-12）

🎉 主要版本：安全、测试和过滤改进

新功能：

• ✨ 使用正则表达式模式的高级工具过滤(exclude_name_pattern)
• ✨ 使用AND/OR逻辑进行多占位符过滤
• ✨ analyze_tools.py 用于浏览工具目录的CLI
• ✨ 全面的输入验证和安全强化

安全：

• 🛡️ SQL注入保护
• 🛡️ 路径遍历保护
• 🛡️ 命令注入预防
• 🛡️ XSS保护
• 🛡️ 安全 exec() 使用验证
• 🛡️ SMILES字符串验证

测试：

• ✅ 447+次综合测试（约170次）
• ✅ 平均覆盖率96%（约85%）
• ✅ 所有关键路径的安全测试
• ✅ Python 3.11和3.12 CI矩阵测试
• ✅ 支持自托管运行器的GPU测试套件

改进：

• ⚡ 更快的测试套件（所有快速测试约3-4s）
• 📝 综合文档（TESTING.md、FILTERING.md）
• 🔧 更好的错误消息和日志记录
• 🎯 改进了工具过滤性能
• 🐛 固定参数映射边缘情况

基础设施：

• 🔄 改进了带缓存的CI/CD
• 📊 向Codecov报告覆盖范围
• 🎨 Ruff linting和格式化
• 🐳 已更新Docker配置

v0.1.0（初始版本）

• 初始版本支持双模型
• 基本工具筛选
• 配置系统
• Docker部署
• FastMCP集成

______________________________________________________________________

🙏 致谢

• 谷歌DeepMind 对于TxGemma型号
• 治疗数据共享 用于训练数据和基准测试
• Anthropic 用于模型上下文协议规范和Claude
• FastMCP MCP服务器框架项目
• 星体 用于uv和ruff工具

______________________________________________________________________

📄 许可证

此项目根据MIT许可证获得许可-请参阅 许可证 文件以获取详细信息。

______________________________________________________________________

Author: Tobias Neumann

Version: 0.1.1