AI Image Generation Pipeline with LangGraph and MCP

使用 LangGraph 和 MCP 的 AI 图像生成管道

本项目展示了如何使用模型上下文协议（MCP）和 LangGraph 创建工作流，根据给定的主题生成提示词和 AI 生成的图像。该项目由三个主要文件组成：app.py、graph.py 和 ai-image-gen-pipeline.py。每个文件展示了使用 MCP 与 LangGraph 的不同方面，包括 LangGraph 函数式 API、图 API 以及与 Open WebUI 管道 的集成。这些脚本使用 Comfy MCP 服务 来生成 AI 图像提示词和 AI 图像。

文件

app.py

此脚本演示了 LangGraph 函数式 API 与人工在环（HIL）交互的使用，根据给定的主题生成提示词和 AI 生成的图像。工作流包括用户反馈，在生成相应图像之前批准生成的提示词。

关键组件：

• 依赖项：aiosqlite、langgraph、langgraph-checkpoint-sqlite、mcp[cli]。
• 函数：
  • run_tool(tool: str, args: dict) -> str：使用 MCP 服务运行工具。
  • generate_prompt(topic: str) -> str：为给定的主题生成提示词。
  • generate_image(prompt: str) -> str：根据给定的提示词生成图像。
  • get_feedback(topic: str, prompt: str) -> str：收集用户对生成的提示词的反馈。
  • workflow_func(saver)：定义具有检查点功能的工作流函数。
• 主函数：
  • 解析命令行参数以获取线程 ID 以及可选的主题和反馈。
  • 初始化工作流并根据提供的输入运行它。

graph.py

此脚本演示了 LangGraph 图 API 与人工在环（HIL）交互的使用，根据给定的主题生成提示词和 AI 生成的图像。工作流包括用户反馈，在生成相应图像之前批准生成的提示词。

关键组件：

• 依赖项：aiosqlite、langgraph、langgraph-checkpoint-sqlite、mcp[cli]。
• 函数：
  • run_tool(tool: str, args: dict) -> str：使用 MCP 服务运行工具。
  • generate_prompt(state: State) -> State：为给定的主题生成提示词并更新状态。
  • generate_image(state: State) -> State：根据给定的提示词生成图像并更新状态。
  • prompt_feedback(state: State) -> State：收集用户对生成的提示词的反馈。
  • process_feedback(state: State) -> str：处理用户反馈以确定工作流中的下一步。
• 主函数：
  • 解析命令行参数以获取线程 ID、主题和反馈。
  • 初始化状态图并根据提供的输入运行它。

ai-image-gen-pipeline.py

此脚本演示了 LangGraph API 与人工在环（HIL）在 Open WebUI 管道 中的集成。它定义了一个使用 MCP 生成提示词和图像的管道，包括用于生成提示词、处理反馈和生成图像的节点。

关键组件：

• 依赖项：aiosqlite、langgraph、langgraph-checkpoint-sqlite、mcp[cli]。
• 类：
  • Pipeline：定义包含生成提示词、处理反馈和生成图像节点的工作流。
    • Valves(BaseModel)：包含 MCP 服务配置的环境变量。
• 函数：
  • inlet(body: dict, user: dict) -> dict：处理传入的消息。
  • outlet(body: dict, user: dict) -> dict：处理传出的消息。
  • pipe(user_message: str, model_id: str, messages: List[dict], body: dict) -> Union[str, Generator, Iterator]：定义主工作流逻辑。
  • run_tool(tool: str, args: dict) -> str：使用 MCP 服务运行工具。
  • generate_prompt(state: State) -> State：为给定主题生成提示词并更新状态。
  • generate_image(state: State) -> State：基于给定提示词生成图像并更新状态。
  • prompt_feedback(state: State) -> State：收集用户对生成的提示词的反馈。
  • process_feedback(state: State) -> str：处理用户反馈以确定工作流中的下一步。

使用方法

1. 安装依赖项：确保已安装所需的依赖项。

   pip install aiosqlite langgraph langgraph-checkpoint-sqlite mcp[cli] comfy-mcp-server
   

2. 运行应用程序：

   • 对于 app.py：

     python app.py --topic "Your topic here"
     

   • 对于 graph.py：

     python graph.py --thread_id "your-thread-id" --topic "Your topic here" 
     

     对于反馈：

     python graph.py --thread_id "your-thread-id" --feedback "y/n" 
     

3. 使用 uv 工具：您也可以使用 uv 工具启动 app.py 和 graph.py。该工具管理 Python 版本和依赖项，因此无需预先安装依赖项。

   • 对于 app.py：

     uv run app.py --topic "Your topic here"
     

   • 对于 graph.py：

     uv run graph.py --thread_id "your-thread-id" --topic "Your topic here" 
     

     对于反馈：

     uv run graph.py --thread_id "your-thread-id" --feedback "y/n" 
     

4. 环境变量：设置 MCP 服务配置所需的环境变量。

   export COMFY_URL="comfy-url"
   export COMFY_URL_EXTERNAL="comfy-url-external"
   export COMFY_WORKFLOW_JSON_FILE="path-to-workflow-json-file"
   export PROMPT_NODE_ID="prompt-node-id"
   export OUTPUT_NODE_ID="output-node-id"
   export OLLAMA_API_BASE="ollama-api-base"
   export PROMPT_LLM="prompt-llm"
   

贡献

欢迎通过提交拉取请求或问题来为该项目做出贡献。确保任何更改都有充分的文档和测试。

许可证

本项目在 MIT 许可证下获得许可。