186 lines
6.4 KiB
Python
186 lines
6.4 KiB
Python
"""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": {}}
|