pingtai_agent/gcp_batch_agent

GCP Batch Agent

基于 Vertex AI 的批量推理 Agent，支持 Flex PayGo 模式。用户传入 Cloud Storage 地址，Agent 自动提交批量作业并监控状态。

功能特性

• 提交批量作业 - 支持 Gemini 系列模型的批量推理
• 监控作业状态 - 实时查询作业进度和状态
• 列出作业 - 查看所有批量作业
• 取消作业 - 取消正在运行的作业

快速开始

1. 环境配置

# 必需环境变量
export GCP_PROJECT_ID="your-project-id"
export GCP_LOCATION="us-central1"

# 服务账号凭证（二选一）
export GCP_SERVICE_ACCOUNT_JSON='{"type":"service_account",...}'  # JSON 字符串
# 或
export GCP_SERVICE_ACCOUNT_FILE="/path/to/service-account.json"  # 文件路径

# 可选配置
export DEFAULT_MODEL="gemini-2.0-flash"
export API_PORT="8000"

2. 用户授权存储桶

用户需要给 Agent 服务账号授权访问其 Cloud Storage 存储桶：

# 获取 Agent 服务账号邮箱
AGENT_SA="your-agent-sa@your-project.iam.gserviceaccount.com"

# 授权存储桶访问权限
gsutil iam ch serviceAccount:${AGENT_SA}:objectViewer gs://your-bucket
gsutil iam ch serviceAccount:${AGENT_SA}:objectCreator gs://your-bucket

3. 准备输入文件

创建 JSONL 格式的输入文件：

{"request":{"contents":[{"role":"user","parts":[{"text":"What is AI?"}]}]}}
{"request":{"contents":[{"role":"user","parts":[{"text":"Explain quantum computing."}]}]}}

上传到 Cloud Storage：

gsutil cp input.jsonl gs://your-bucket/input.jsonl

4. 启动服务

# 本地运行
pip install -r requirements.txt
python run_api_server.py

# Docker 运行
docker build -t gcp-batch-agent .
docker run -p 8000:8000 \
  -e GCP_PROJECT_ID=your-project \
  -e GCP_SERVICE_ACCOUNT_FILE=/app/sa.json \
  -v /path/to/sa.json:/app/sa.json \
  gcp-batch-agent

API 使用

提交批量作业

curl -X POST http://localhost:8000/api/v1/batch/submit \
  -H "Content-Type: application/json" \
  -H "api-key: your-api-key" \
  -d '{
    "input_uri": "gs://your-bucket/input.jsonl",
    "output_uri": "gs://your-bucket/output/",
    "model": "gemini-2.0-flash"
  }'

查询作业状态

curl http://localhost:8000/api/v1/batch/status/123456 \
  -H "api-key: your-api-key"

列出作业

curl "http://localhost:8000/api/v1/batch/list?page_size=10" \
  -H "api-key: your-api-key"

取消作业

curl -X POST http://localhost:8000/api/v1/batch/cancel/123456 \
  -H "api-key: your-api-key"

MCP 工具

工具	说明
submit_batch_job	提交批量推理作业
get_job_status	查询作业状态
list_jobs	列出批量作业
cancel_job	取消作业

支持的模型

• gemini-2.0-flash（默认）
• gemini-2.0-flash-lite
• gemini-1.5-flash
• gemini-1.5-pro
• gemini-1.0-pro

服务端点

端点	方法	说明
/	GET	服务状态
/health	GET	健康检查
/mcp	POST	MCP JSON-RPC
/mcp/sse	GET/POST	MCP SSE 流式
/api/v1/batch/submit	POST	提交批量作业
/api/v1/batch/status/{job_id}	GET	查询作业状态
/api/v1/batch/list	GET	列出作业
/api/v1/batch/cancel/{job_id}	POST	取消作业

项目结构

gcp_batch_agent/
├── Dockerfile
├── README.md
├── requirements.txt
├── run_api_server.py
├── USAGE.md
└── src/
    ├── __init__.py
    └── server/
        ├── __init__.py
        ├── api_server.py      # FastAPI + MCP HTTP 端点
        ├── mcp_server.py      # MCP 工具定义
        └── core/
            ├── __init__.py
            ├── config.py      # 配置管理
            ├── credentials.py # GCP 凭证管理
            ├── batch_client.py # Vertex AI 批量推理客户端
            └── schemas.py     # 数据模型

注意事项

1. Flex PayGo 模式：作业可能需要排队等待资源
2. 存储桶权限：用户必须提前授权 Agent 服务账号
3. 文件格式：输入必须是有效的 JSONL 格式
4. 计费：按实际处理的行数计费，中途取消只收取已完成部分