大活

art-agent/backend/app/agent/loop.py

@@ -0,0 +1,185 @@
+"""Agent Loop 主循环：LLM 对话 -> 工具调用 -> 结果回传 -> 继续。"""
+
+import json
+from typing import AsyncGenerator, Optional
+
+from openai import AsyncOpenAI
+
+from app.agent.tools import TOOL_DEFINITIONS, execute_tool
+
+
+def _get_client() -> AsyncOpenAI:
+    return AsyncOpenAI()
+
+SYSTEM_PROMPT = """\
+你是一个专业的游戏美术 AI 助手。你的工作是帮助美术人员通过对话生成游戏美术资源。
+
+## 你的能力
+- 根据用户的文字描述生成图片（UI图标、按钮、插画、立绘、概念图等）
+- 理解用户的审美意图，将中文描述转化为高质量的英文生成 prompt
+- 根据用户反馈迭代修改（调整颜色、风格、构图等）
+- 如果用户提供了参考图，将参考图的风格元素融入生成 prompt
+
+## 工作流程
+1. 理解用户需求，必要时追问细节（尺寸、风格、用途等）
+2. 将需求转化为详细的英文 prompt，调用 generate_image 工具生成图片
+3. 向用户展示结果并询问反馈
+4. 根据反馈调整 prompt 并重新生成
+
+## 生成 prompt 要求
+- 必须使用英文
+- 尽量详细描述：主体内容、风格、颜色方案、光照、构图、材质等
+- 如果用户要求游戏 UI 元素，添加相关关键词如 "game UI", "icon", "button" 等
+- 如果用户提供了参考图，在 prompt 中描述参考图的风格特征
+
+## 注意事项
+- 用中文和用户交流
+- 生成图片后简要说明你使用的 prompt 思路
+- 主动建议迭代方向
+"""
+
+
+async def run_agent_loop(
+    messages: list[dict],
+    ref_image_url: Optional[str] = None,
+) -> AsyncGenerator[dict, None]:
+    """
+    运行 Agent Loop，以 SSE 事件流形式 yield 结果。
+
+    事件类型：
+    - text_delta: LLM 文字增量
+    - tool_start: 开始执行工具
+    - image_result: 图片生成结果
+    - done: 完成
+    - error: 错误
+    """
+    # 构建消息列表
+    api_messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+
+    for msg in messages:
+        if msg["role"] == "user" and ref_image_url:
+            # 最后一条用户消息附加参考图（GPT-4o vision）
+            if msg == messages[-1] or (
+                msg.get("role") == "user"
+                and messages.index(msg) == len(messages) - 1
+            ):
+                api_messages.append({
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": msg["content"]},
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": ref_image_url},
+                        },
+                    ],
+                })
+                continue
+        api_messages.append({"role": msg["role"], "content": msg["content"]})
+
+    client = _get_client()
+
+    max_iterations = 5
+    for _ in range(max_iterations):
+        try:
+            response = await client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=api_messages,
+                tools=TOOL_DEFINITIONS,
+                stream=True,
+            )
+        except Exception as e:
+            yield {"type": "error", "data": {"message": str(e)}}
+            return
+
+        collected_text = ""
+        tool_calls_data: dict[int, dict] = {}
+
+        async for chunk in response:
+            delta = chunk.choices[0].delta if chunk.choices else None
+            if not delta:
+                continue
+
+            # 文字内容
+            if delta.content:
+                collected_text += delta.content
+                yield {"type": "text_delta", "data": {"text": delta.content}}
+
+            # 工具调用（流式累积）
+            if delta.tool_calls:
+                for tc in delta.tool_calls:
+                    idx = tc.index
+                    if idx not in tool_calls_data:
+                        tool_calls_data[idx] = {
+                            "id": "",
+                            "name": "",
+                            "arguments": "",
+                        }
+                    if tc.id:
+                        tool_calls_data[idx]["id"] = tc.id
+                    if tc.function and tc.function.name:
+                        tool_calls_data[idx]["name"] = tc.function.name
+                    if tc.function and tc.function.arguments:
+                        tool_calls_data[idx]["arguments"] += tc.function.arguments
+
+        finish_reason = chunk.choices[0].finish_reason if chunk.choices else None
+
+        # 如果没有工具调用，对话结束
+        if not tool_calls_data:
+            yield {"type": "done", "data": {}}
+            return
+
+        # 将 assistant 消息（含 tool_calls）加入历史
+        assistant_msg: dict = {"role": "assistant"}
+        if collected_text:
+            assistant_msg["content"] = collected_text
+        else:
+            assistant_msg["content"] = None
+
+        assistant_msg["tool_calls"] = []
+        for idx in sorted(tool_calls_data.keys()):
+            tc_data = tool_calls_data[idx]
+            assistant_msg["tool_calls"].append({
+                "id": tc_data["id"],
+                "type": "function",
+                "function": {
+                    "name": tc_data["name"],
+                    "arguments": tc_data["arguments"],
+                },
+            })
+        api_messages.append(assistant_msg)
+
+        # 依次执行每个工具调用
+        for idx in sorted(tool_calls_data.keys()):
+            tc_data = tool_calls_data[idx]
+            tool_name = tc_data["name"]
+
+            yield {
+                "type": "tool_start",
+                "data": {"tool": tool_name, "message": "正在生成图片..."},
+            }
+
+            try:
+                arguments = json.loads(tc_data["arguments"])
+            except json.JSONDecodeError:
+                arguments = {}
+
+            result = await execute_tool(tool_name, arguments, ref_image_url)
+
+            # 如果有图片结果，推送给前端
+            if result.get("images"):
+                yield {
+                    "type": "image_result",
+                    "data": {"images": result["images"]},
+                }
+
+            # 工具结果回传给 LLM
+            api_messages.append({
+                "role": "tool",
+                "tool_call_id": tc_data["id"],
+                "content": json.dumps(result, ensure_ascii=False),
+            })
+
+        # 清除 ref_image_url，避免后续轮次重复附加
+        ref_image_url = None
+
+    yield {"type": "done", "data": {}}