Image Agent 架构设计文档

重构日期：2026-03-18 作者：Claude Code

1. 背景与目标

1.1 背景

当前 api/image.py 直接调用 model_integration.generate_image()，存在以下问题：

缺乏扩展性：无法轻松添加新的图片处理步骤（如压缩、水印）
缺乏统一架构：与 ChatAgent 不一致，不便于维护
缺乏工作流管理：预处理、生成、后处理逻辑混杂

1.2 目标

统一架构：参考 ChatAgent，使用 LangGraph StateGraph 管理图片生成流程
多模型路由：支持 Gemini、ByteDance 等多种图片生成模型
可扩展流水线：预处理 → 生成 → 后处理，便于插入新节点
轻量级状态管理：通过 thread_id 关联历史记录

2. 架构设计

2.1 目录结构

agent/
├── base.py                  # BaseAgent（已有，不修改）
├── chat/                    # Chat Agent（已有）
└── image/                   # Image Agent（新建）
    ├── __init__.py          # 模块导出
    ├── config.py            # ImageConfig 配置类
    ├── state.py             # ImageState 状态定义
    ├── graph.py             # ImageGraphBuilder
    ├── agent.py             # ImageAgent 实现
    ├── providers/           # 多模型提供商路由
    │   ├── __init__.py
    │   ├── base.py          # BaseImageProvider 抽象类
    │   ├── gemini.py        # GeminiProvider
    │   └── bytedance.py     # ByteDanceProvider
    └── nodes/               # LangGraph 节点
        ├── __init__.py
        ├── preprocess.py    # 预处理节点
        ├── generate.py      # 生成节点
        └── postprocess.py   # 后处理节点

2.2 LangGraph StateGraph 工作流

2.3 节点职责

节点	职责
`preprocess`	验证提示词、转换图片格式（URL/base64/COS → base64）、记录日志
`generate`	根据模型路由到提供商、调用 API、返回 base64 图片
`postprocess`	保存到 COS、生成签名 URL、返回结果

2.4 节点详解

2.4.1 preprocess_node - 预处理节点

职责：在调用 AI 模型之前，准备好所有输入数据

处理流程：

验证提示词
- 确保不为空
- 去除首尾空格
- 如果为空，设置 error 并终止流程
转换图片格式（图生图场景）
- 用户上传的图片可能是多种格式：
  - data:image/png;base64,xxx - base64 编码
  - cos://uploads/images/xxx - COS 存储键
  - https://example.com/xxx - 外部 URL
- 需要统一转换为 base64 格式供模型使用
- 调用 ImageProcessor.process_mixed_images() 完成转换

代码位置：agent/image/nodes/preprocess.py

2.4.2 postprocess_node - 后处理节点

职责：模型生成图片后，保存并返回给用户

处理流程：

保存到 COS - 将 base64 图片上传到腾讯云 COS
生成签名 URL - 生成有时效的访问链接供前端展示
错误处理 - 单张图片失败不阻塞整体流程

代码位置：agent/image/nodes/postprocess.py

2.4.3 _route_on_error - 错误路由

职责：在节点之间判断是否发生错误，决定是否跳过后续步骤

工作原理：

每个节点执行后，检查 state.get("error") 是否有值
如果有错误，跳过后续所有节点，直接返回错误结果
避免无意义的 API 调用（如提示词为空时不应调用模型）

示例场景：

场景	preprocess 结果	路由决策	后续行为
提示词为空	`error = "提示词不能为空"`	→ END	不调用 API
图片格式转换失败	记录警告，不阻塞	→ generate	继续生成
正常输入	`validated_prompt = "..."`	→ generate	调用 API

代码位置：agent/image/graph.py::_route_on_error()

3. 核心组件

3.1 ImageState

python

class ImageState(TypedDict, total=False):
    # 输入
    prompt: str
    model: str
    reference_images: Optional[List[str]]
    user_id: int
    thread_id: str
    metadata: Dict[str, Any]

    # 预处理结果
    preprocessed_images: Optional[List[str]]
    validated_prompt: str

    # 生成结果
    generated_base64_urls: Annotated[List[str], merge_lists]

    # 后处理结果
    cos_keys: Annotated[List[str], merge_lists]
    signed_urls: Annotated[List[str], merge_lists]

    # 错误处理
    error: Optional[str]

3.2 Provider 路由

Provider	模型前缀	特点
`GeminiProvider`	`google/`	支持 `modalities: ["image", "text"]`，支持图生图
`ByteDanceProvider`	`bytedance-seed/`	支持 `modalities: ["image"]`，仅文生图

3.3 ImageAgent API

python

class ImageAgent:
    def generate(
        self,
        prompt: str,
        model: str,
        user_id: int,
        reference_images: Optional[List[str]] = None,
        thread_id: Optional[str] = None,
        user_info: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """返回 {
            "success": bool,
            "image_urls": List[str],  # 签名 URL
            "cos_keys": List[str],    # COS 存储键
            "error": Optional[str]
        }"""

4. 数据流示例

4.1 文生图流程

4.2 图生图流程

4.3 错误处理流程

5. 使用示例

4.1 直接使用 ImageAgent

python

from agent.image import get_image_agent

agent = get_image_agent()
result = agent.generate(
    prompt="a cute cat playing with a ball",
    model="google/gemini-2.5-flash-image-preview",
    user_id=1,
    user_info={"user_id": 1, "nickname": "test"}
)

if result["success"]:
    for url in result["image_urls"]:
        print(f"图片 URL: {url}")
else:
    print(f"生成失败: {result['error']}")

4.2 通过 API 调用

bash

curl -X POST http://localhost:5000/api/generate-image \
  -H "Content-Type: application/json" \
  -H "Cookie: session=xxx" \
  -d '{
    "prompt": "a cute cat playing with a ball",
    "model": "google/gemini-2.5-flash-image-preview"
  }'

6. 扩展方向

6.1 多步骤工作流

未来可在 StateGraph 中添加更多节点：

实现步骤：

在 nodes/ 中添加新节点：

python

# nodes/enhance_prompt.py
def enhance_prompt_node(state: ImageState) -> Dict[str, Any]:
    """使用 LLM 增强提示词"""
    pass

# nodes/compress.py
def compress_node(state: ImageState) -> Dict[str, Any]:
    """压缩图片"""
    pass

# nodes/watermark.py
def watermark_node(state: ImageState) -> Dict[str, Any]:
    """添加水印"""
    pass

在 graph.py 中添加节点和边

5.2 新增图片生成提供商

创建新的 Provider 类：

python

# agent/image/providers/openai.py
class OpenAIImageProvider(BaseImageProvider):
    @property
    def name(self) -> str:
        return "openai"

    def generate(self, prompt, model, reference_images=None):
        # 实现 DALL-E 调用
        pass

注册到 Provider 注册表：

python

# agent/image/providers/__init__.py
from agent.image.providers.openai import OpenAIImageProvider
register_provider("openai", OpenAIImageProvider)

更新配置映射：

python

# agent/image/config.py
provider_mapping = {
    "google/": "gemini",
    "bytedance-seed/": "bytedance",
    "openai/": "openai",  # 新增
}

5.3 状态管理增强

如需完整状态管理（如恢复中断的任务），可参考 ChatAgent 添加 MySQL Checkpointer：

python

from services.checkpointer import get_checkpointer

class ImageAgent:
    def __init__(self, config: Optional[ImageConfig] = None):
        self.config = config or ImageConfig()
        self._checkpointer = get_checkpointer()
        self._graph_builder = ImageGraphBuilder(self.config)

    def generate(self, ...):
        graph = self._graph_builder.build(self._checkpointer)
        config = {"configurable": {"thread_id": thread_id}}
        result = graph.invoke(initial_state, config)

7. 与现有架构的关系

6.1 与 ChatAgent 的对比

特性	ChatAgent	ImageAgent
流式输出	✅ token 级别	❌ 同步返回
Checkpointer	✅ 完整状态管理	❌ 轻量级
多步骤工作流	❌ 单节点	✅ 预处理/生成/后处理
多模型路由	✅ 动态模型选择	✅ Provider 路由

6.2 迁移自 model_integration.py

原方法	新位置
`generate_image()`	`ImageAgent.generate()`
`_generate_image_gemini()`	`providers/gemini.py::GeminiProvider.generate()`
`_generate_image_bytedance()`	`providers/bytedance.py::ByteDanceProvider.generate()`

8. 测试方案

7.1 单元测试

python

# tests/agent/image/test_providers.py
def test_gemini_provider():
    provider = GeminiProvider(api_key="test", base_url="...")
    result = provider.generate("a cat", "google/gemini-2.5-flash-image")
    assert isinstance(result, list)

# tests/agent/image/test_nodes.py
def test_preprocess_node():
    state = {"prompt": "test", "model": "google/gemini-2.5-flash-image"}
    result = preprocess_node(state)
    assert "validated_prompt" in result

7.2 集成测试

python

# tests/agent/image/test_agent.py
def test_image_agent_generate():
    agent = get_image_agent()
    result = agent.generate(
        prompt="a cute cat",
        model="google/gemini-2.5-flash-image-preview",
        user_id=1
    )
    assert result["success"] == True
    assert len(result["image_urls"]) > 0

7.3 端到端测试

bash

# 启动服务
make run

# 调用 API
curl -X POST http://localhost:5000/api/generate-image \
  -H "Content-Type: application/json" \
  -H "Cookie: session=xxx" \
  -d '{"prompt": "a cute cat", "model": "google/gemini-2.5-flash-image-preview"}'

9. 参考

10. 总结

10.1 核心设计理念

Image Agent 采用 LangGraph StateGraph 实现图片生成工作流，核心理念：

单一职责：每个节点只负责一个明确的任务
可扩展性：通过添加新节点/Provider 扩展功能
错误隔离：错误不会传播，通过路由提前终止
类型安全：使用 TypedDict 定义状态结构

10.2 与原实现对比

方面	原 model_integration.py	新 ImageAgent
架构	单函数调用	StateGraph 工作流
扩展性	需修改原函数	添加新节点/Provider
错误处理	异常抛出	状态流转 + 路由
可测试性	需 mock API	可单独测试每个节点
可维护性	逻辑混杂	职责分离

10.3 快速导航

文件	说明
`agent/image/agent.py`	入口类，调用 `generate()` 开始
`agent/image/graph.py`	工作流定义，节点连接关系
`agent/image/state.py`	状态结构定义
`agent/image/nodes/preprocess.py`	输入验证和预处理
`agent/image/nodes/generate.py`	模型调用和路由
`agent/image/nodes/postprocess.py`	结果处理和存储
`agent/image/providers/`	多模型提供商实现

Image Agent 架构设计文档 ​

1. 背景与目标 ​

1.1 背景 ​

1.2 目标 ​

2. 架构设计 ​

2.1 目录结构 ​

2.2 LangGraph StateGraph 工作流 ​

2.3 节点职责 ​

2.4 节点详解 ​

2.4.1 preprocess_node - 预处理节点 ​

2.4.2 postprocess_node - 后处理节点 ​

2.4.3 _route_on_error - 错误路由 ​

3. 核心组件 ​

3.1 ImageState ​

3.2 Provider 路由 ​

3.3 ImageAgent API ​

4. 数据流示例 ​

4.1 文生图流程 ​

4.2 图生图流程 ​

4.3 错误处理流程 ​

5. 使用示例 ​

4.1 直接使用 ImageAgent ​

4.2 通过 API 调用 ​

6. 扩展方向 ​

6.1 多步骤工作流 ​

5.2 新增图片生成提供商 ​

5.3 状态管理增强 ​

7. 与现有架构的关系 ​

6.1 与 ChatAgent 的对比 ​

6.2 迁移自 model_integration.py ​

8. 测试方案 ​

7.1 单元测试 ​

7.2 集成测试 ​

7.3 端到端测试 ​

9. 参考 ​

10. 总结 ​

10.1 核心设计理念 ​

10.2 与原实现对比 ​

10.3 快速导航 ​

Image Agent 架构设计文档

1. 背景与目标

1.1 背景

1.2 目标

2. 架构设计

2.1 目录结构

2.2 LangGraph StateGraph 工作流

2.3 节点职责

2.4 节点详解

2.4.1 preprocess_node - 预处理节点

2.4.2 postprocess_node - 后处理节点

2.4.3 _route_on_error - 错误路由

3. 核心组件

3.1 ImageState

3.2 Provider 路由

3.3 ImageAgent API

4. 数据流示例

4.1 文生图流程

4.2 图生图流程

4.3 错误处理流程

5. 使用示例

4.1 直接使用 ImageAgent

4.2 通过 API 调用

6. 扩展方向

6.1 多步骤工作流

5.2 新增图片生成提供商

5.3 状态管理增强

7. 与现有架构的关系

6.1 与 ChatAgent 的对比

6.2 迁移自 model_integration.py

8. 测试方案

7.1 单元测试

7.2 集成测试

7.3 端到端测试

9. 参考

10. 总结

10.1 核心设计理念

10.2 与原实现对比

10.3 快速导航