345 lines
14 KiB
Python
345 lines
14 KiB
Python
"""Agent Loop 主循环：LLM 对话 -> 工具调用 -> 结果回传 -> 继续。"""
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import re
|
||
import os
|
||
from typing import AsyncGenerator, Optional
|
||
|
||
from openai import AsyncOpenAI
|
||
|
||
from app.agent.tools import TOOL_DEFINITIONS, execute_tool
|
||
from app.config import (
|
||
get_llm_model,
|
||
get_llm_max_iterations,
|
||
get_image_model_config,
|
||
get_max_recent_turns,
|
||
)
|
||
from app.memory import get_memory
|
||
from app.services.image_gen import to_data_uri
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def _get_client() -> AsyncOpenAI:
|
||
base_url = os.getenv("OPENAI_BASE_URL")
|
||
return AsyncOpenAI(base_url=base_url) if base_url else AsyncOpenAI()
|
||
|
||
SYSTEM_PROMPT = """\
|
||
你是一个专业的游戏美术 AI 助手。你的工作是帮助美术人员通过对话生成游戏美术资源。
|
||
|
||
## 你的能力
|
||
- 根据用户的文字描述生成图片（UI图标、按钮、插画、立绘、概念图等）
|
||
- 理解用户的审美意图，将中文描述转化为高质量的英文生成 prompt
|
||
- 根据用户反馈迭代修改（调整颜色、风格、构图等）
|
||
- 如果用户提供了参考图，将参考图的风格元素融入生成 prompt
|
||
- 理解用户在图片上的标注（框选区域 + 文字批注），精准定位需要修改的部分
|
||
|
||
## 工作流程
|
||
1. 理解用户需求，必要时追问细节（尺寸、风格、用途等）
|
||
2. 将需求转化为详细的英文 prompt，调用 generate_image 工具生成图片
|
||
3. 向用户展示结果并询问反馈
|
||
4. 根据反馈调整 prompt 并重新生成
|
||
|
||
## 标注理解
|
||
当用户发送带有「图片标注」的消息时，表示用户在之前生成的图片上做了标注。
|
||
标注格式为：`[区域 (x%, y%) 大小 w%×h%]: 修改意见`
|
||
- 区域坐标表示标注框在图片上的相对位置
|
||
- 你需要理解标注区域所指的图片内容，并将修改意见融入新的 prompt
|
||
- 如果附带了标注截图（参考图），仔细观察红色标注框和文字来理解用户意图
|
||
- 在调整 prompt 时，保持原图整体风格不变，只针对标注区域做修改
|
||
|
||
## 生成 prompt 要求
|
||
- 必须使用英文
|
||
- 尽量详细描述：主体内容、颜色方案、光照、构图、材质等
|
||
- 如果用户要求游戏 UI 元素，添加相关关键词如 "game UI", "icon", "button" 等
|
||
- 如果你能直接看到参考图（图片内容），可以在 prompt 中描述参考图的风格特征
|
||
- 如果你无法看到参考图（只收到了文字提示说有参考图），参考图会由生图工具的 IP-Adapter 自动处理风格融合。此时你不要自行猜测画风/艺术风格关键词（如 pixel art、watercolor、oil painting 等），把风格交给参考图来决定。但如果用户在消息中明确指定了风格（如"赛博朋克风"、"水彩风"等），应保留并翻译到 prompt 中——尊重用户的主动意图
|
||
|
||
## 注意事项
|
||
- 用中文和用户交流
|
||
- 生成图片后简要说明你使用的 prompt 思路
|
||
- 主动建议迭代方向
|
||
- **禁止模拟工具调用**：生成图片时必须实际调用 generate_image 工具，绝不能用文字描述"已生成"或假装工具已执行。如果需要生成多张图片，每张都必须单独调用工具
|
||
"""
|
||
|
||
|
||
async def run_agent_loop(
|
||
messages: list[dict],
|
||
ref_image_url: Optional[str] = None,
|
||
image_model: Optional[str] = None,
|
||
session_id: Optional[str] = None,
|
||
user_id: str = "default_user",
|
||
) -> AsyncGenerator[dict, None]:
|
||
"""
|
||
运行 Agent Loop，以 SSE 事件流形式 yield 结果。
|
||
|
||
事件类型：
|
||
- text_delta: LLM 文字增量
|
||
- tool_start: 开始执行工具
|
||
- image_result: 图片生成结果
|
||
- memory_warning: 记忆存储异常（不中断对话）
|
||
- done: 完成
|
||
- error: 错误
|
||
"""
|
||
# ── 滑动窗口：截断过早的消息 ──
|
||
max_turns = get_max_recent_turns()
|
||
if len(messages) > max_turns:
|
||
messages = messages[-max_turns:]
|
||
|
||
# ── 记忆检索：用最新 user 消息语义检索相关记忆 ──
|
||
memory_block = ""
|
||
try:
|
||
memory = get_memory()
|
||
last_user_content = ""
|
||
for m in reversed(messages):
|
||
if m["role"] == "user":
|
||
last_user_content = m["content"]
|
||
break
|
||
|
||
if last_user_content:
|
||
search_kwargs = {"query": last_user_content, "user_id": user_id, "limit": 10}
|
||
if session_id:
|
||
search_kwargs["run_id"] = session_id
|
||
relevant = memory.search(**search_kwargs)
|
||
|
||
results = relevant.get("results", []) if isinstance(relevant, dict) else relevant
|
||
if results:
|
||
items = "\n".join(f"- {m['memory']}" for m in results if m.get("memory"))
|
||
if items:
|
||
memory_block = f"\n\n## 用户记忆（来自历史对话）\n{items}"
|
||
except Exception as e:
|
||
logger.error("Mem0 记忆检索失败: %s", e, exc_info=True)
|
||
yield {"type": "error", "data": {"message": f"记忆系统检索失败: {e}"}}
|
||
return
|
||
|
||
# ── 构建 system prompt ──
|
||
model_config = get_image_model_config(image_model)
|
||
current_model_name = model_config.get('name', '未知')
|
||
current_model_id = model_config.get('id', '未知')
|
||
model_hint = (
|
||
f"\n\n## 当前生图模型（重要）\n"
|
||
f"本次对话用户选择的生图模型是 **{current_model_name}**"
|
||
f"（model_id: {current_model_id}）。\n"
|
||
f"**注意**：对话历史中可能包含之前使用其他模型的记录，忽略那些旧模型名。"
|
||
f"本次生成使用的是 {current_model_name}，在回复中只能使用这个名称。"
|
||
)
|
||
api_messages = [{"role": "system", "content": SYSTEM_PROMPT + model_hint + memory_block}]
|
||
|
||
# 检测当前 LLM 是否支持 vision（多模态图片输入）
|
||
llm_model = get_llm_model().lower()
|
||
vision_capable = any(kw in llm_model for kw in ("gpt-4o", "gpt-4-vision", "claude"))
|
||
|
||
for msg in messages:
|
||
if msg["role"] == "user" and ref_image_url and msg is messages[-1]:
|
||
if vision_capable:
|
||
api_messages.append({
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": msg["content"]},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {"url": to_data_uri(ref_image_url)},
|
||
},
|
||
],
|
||
})
|
||
else:
|
||
hint = (
|
||
f"{msg['content']}\n\n"
|
||
"【系统提示：用户上传了一张参考图，已自动传递给图片生成工具的 IP-Adapter。"
|
||
"IP-Adapter 会从参考图中提取风格并融合到生成结果中。"
|
||
"你无法看到这张参考图，因此在生成 prompt 时：\n"
|
||
"1. 描述画面内容（主体、构图、光照、材质等）\n"
|
||
"2. 不要自行猜测画风/艺术风格——但如果用户明确指定了风格，保留到 prompt 中\n"
|
||
"3. 用户未指定风格时，风格完全由参考图通过 IP-Adapter 决定】"
|
||
)
|
||
api_messages.append({"role": "user", "content": hint})
|
||
continue
|
||
api_messages.append({"role": msg["role"], "content": msg["content"]})
|
||
|
||
client = _get_client()
|
||
|
||
for _ in range(get_llm_max_iterations()):
|
||
try:
|
||
response = await client.chat.completions.create(
|
||
model=get_llm_model(),
|
||
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
|
||
|
||
# 如果没有工具调用：检测 LLM 是否在用文字模拟生图
|
||
if not tool_calls_data:
|
||
if _looks_like_fake_generation(collected_text):
|
||
# LLM 用文字假装调用了工具，丢弃这段文字，注入纠正消息强制重试
|
||
yield {
|
||
"type": "text_delta",
|
||
"data": {"text": "\n\n[系统：检测到未调用生图工具，正在重试...]\n"},
|
||
}
|
||
api_messages.append({"role": "assistant", "content": collected_text})
|
||
api_messages.append({
|
||
"role": "user",
|
||
"content": (
|
||
"你刚才没有调用 generate_image 工具，只是用文字描述了生成过程。"
|
||
"请立即调用 generate_image 工具来实际生成图片。"
|
||
"不要解释，直接调用工具。"
|
||
f"当前使用的生图模型是 {model_config.get('name', '未知')}。"
|
||
),
|
||
})
|
||
continue
|
||
# 对话正常结束，异步存储记忆
|
||
async for evt in _store_and_done(messages, session_id, user_id):
|
||
yield evt
|
||
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, image_model)
|
||
|
||
used_model = result.get("model_name", "")
|
||
|
||
if result.get("errors"):
|
||
yield {
|
||
"type": "tool_error",
|
||
"data": {
|
||
"tool": tool_name,
|
||
"errors": result["errors"],
|
||
"model_name": used_model,
|
||
},
|
||
}
|
||
|
||
if result.get("images"):
|
||
yield {
|
||
"type": "image_result",
|
||
"data": {
|
||
"images": result["images"],
|
||
"prompt_used": result.get("prompt_used", ""),
|
||
"model_name": used_model,
|
||
},
|
||
}
|
||
|
||
# 工具结果回传给 LLM
|
||
api_messages.append({
|
||
"role": "tool",
|
||
"tool_call_id": tc_data["id"],
|
||
"content": json.dumps(result, ensure_ascii=False),
|
||
})
|
||
|
||
# ref_image_url 不清除：消息构建只在循环外执行一次，
|
||
# 后续工具调用仍需参考图（InstantStyle 等模型必须有 style_image）
|
||
|
||
# 迭代次数用尽，存储记忆后结束
|
||
async for evt in _store_and_done(messages, session_id, user_id):
|
||
yield evt
|
||
|
||
|
||
# ─── 异步记忆存储 ─────────────────────────────────────────
|
||
|
||
|
||
async def _store_and_done(
|
||
messages: list[dict], session_id: Optional[str], user_id: str = "default_user"
|
||
) -> AsyncGenerator[dict, None]:
|
||
"""触发 Mem0 异步存储后 yield done。存储失败时 yield warning 但不中断。"""
|
||
try:
|
||
memory = get_memory()
|
||
recent = messages[-4:] if len(messages) >= 4 else messages
|
||
add_kwargs: dict = {"user_id": user_id}
|
||
if session_id:
|
||
add_kwargs["run_id"] = session_id
|
||
|
||
loop = asyncio.get_running_loop()
|
||
await loop.run_in_executor(None, lambda: memory.add(recent, **add_kwargs))
|
||
except Exception as e:
|
||
logger.error("Mem0 记忆存储失败: %s", e, exc_info=True)
|
||
yield {
|
||
"type": "memory_warning",
|
||
"data": {"message": f"记忆存储失败: {e}"},
|
||
}
|
||
|
||
yield {"type": "done", "data": {}}
|
||
|
||
|
||
# ─── 假生成检测 ──────────────────────────────────────────
|
||
|
||
_FAKE_GEN_PATTERNS = re.compile(
|
||
r"已生成|生成完成|开始生成|正在生成|图片已|prompt.*?设计思路",
|
||
re.IGNORECASE,
|
||
)
|
||
|
||
|
||
def _looks_like_fake_generation(text: str) -> bool:
|
||
"""判断 LLM 的纯文字回复是否在假装已经调用了生图工具。"""
|
||
if not text or len(text) < 50:
|
||
return False
|
||
matches = _FAKE_GEN_PATTERNS.findall(text)
|
||
return len(matches) >= 2
|