修复图片卡住agent的bug
This commit is contained in:
@@ -4,6 +4,7 @@
|
|||||||
提供所有 Agent 的基类,内置:
|
提供所有 Agent 的基类,内置:
|
||||||
- 防提示词注入检测
|
- 防提示词注入检测
|
||||||
- 模型调用耗时日志
|
- 模型调用耗时日志
|
||||||
|
- 图片/非文本内容过滤(防止 token 爆炸和模型不支持报错)
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import time
|
import time
|
||||||
@@ -20,19 +21,103 @@ logger = logging.getLogger("adk.agent")
|
|||||||
|
|
||||||
|
|
||||||
# ============================================================
|
# ============================================================
|
||||||
# 防注入的 before_model_callback
|
# 工具函数:判断一个 Part 是否为非文本内容(图片、视频、音频等)
|
||||||
|
# ============================================================
|
||||||
|
def _is_non_text_part(part) -> bool:
|
||||||
|
"""判断一个 Part 是否包含非文本内容(图片、视频、音频、文件等)"""
|
||||||
|
# inline_data: base64 编码的内联数据(图片、PDF 等)
|
||||||
|
if getattr(part, "inline_data", None) and getattr(part.inline_data, "data", None):
|
||||||
|
mime = getattr(part.inline_data, "mime_type", "") or ""
|
||||||
|
# text/* 类型不算非文本(如 inline text/csv)
|
||||||
|
if not mime.startswith("text/"):
|
||||||
|
return True
|
||||||
|
# file_data: 文件引用(GCS URI、文件 URI 等)
|
||||||
|
if getattr(part, "file_data", None) and getattr(part.file_data, "file_uri", None):
|
||||||
|
mime = getattr(part.file_data, "mime_type", "") or ""
|
||||||
|
if not mime.startswith("text/"):
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def _filter_non_text_parts(parts) -> list:
|
||||||
|
"""过滤掉非文本 part,只保留文本 part"""
|
||||||
|
return [p for p in parts if not _is_non_text_part(p)]
|
||||||
|
|
||||||
|
|
||||||
|
def _has_non_text_in_parts(parts) -> bool:
|
||||||
|
"""检查 parts 列表中是否包含非文本内容"""
|
||||||
|
return any(_is_non_text_part(p) for p in parts)
|
||||||
|
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 防注入 + 图片过滤的 before_model_callback
|
||||||
# ============================================================
|
# ============================================================
|
||||||
async def _before_model_callback(callback_context, llm_request):
|
async def _before_model_callback(callback_context, llm_request):
|
||||||
"""在模型调用前检测提示词注入,并记录开始时间"""
|
"""在模型调用前:
|
||||||
# 1. 防提示词注入检测:检查最新的用户消息
|
1. 过滤历史消息中的非文本内容(防止图片累积导致 token 爆炸)
|
||||||
|
2. 检测当前消息是否包含图片(不支持则直接回复提示)
|
||||||
|
3. 防提示词注入检测
|
||||||
|
4. 记录开始时间
|
||||||
|
"""
|
||||||
contents = getattr(llm_request, "contents", None) or []
|
contents = getattr(llm_request, "contents", None) or []
|
||||||
# 找到最后一条 role="user" 的消息
|
|
||||||
|
# ── 1. 找到当前(最后一条)用户消息 ──
|
||||||
last_user_content = None
|
last_user_content = None
|
||||||
for content in reversed(contents):
|
last_user_index = -1
|
||||||
|
for i, content in enumerate(contents):
|
||||||
if getattr(content, "role", None) == "user":
|
if getattr(content, "role", None) == "user":
|
||||||
last_user_content = content
|
last_user_content = content
|
||||||
break
|
last_user_index = i
|
||||||
|
|
||||||
|
# ── 2. 如果当前消息包含非文本内容,直接返回提示 ──
|
||||||
|
if last_user_content:
|
||||||
|
parts = getattr(last_user_content, "parts", None) or []
|
||||||
|
if _has_non_text_in_parts(parts):
|
||||||
|
# 提取文本部分作为上下文(如果有的话)
|
||||||
|
text_parts = _filter_non_text_parts(parts)
|
||||||
|
text_context = "".join(
|
||||||
|
getattr(p, "text", "") or "" for p in text_parts
|
||||||
|
).strip()
|
||||||
|
|
||||||
|
if text_context:
|
||||||
|
reply = (
|
||||||
|
f"我注意到你发送的内容中包含图片或文件,但目前我不支持处理图片输入。\n\n"
|
||||||
|
f"不过我看到了你附带的文字内容,你可以直接用文字描述你想说的,我会帮你处理:\n\n"
|
||||||
|
f"> {text_context[:500]}"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
reply = "抱歉,目前我不支持处理图片或文件输入。请用文字描述你的问题,我会尽力帮助你。"
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"拦截非文本消息 | agent=%s has_text=%s",
|
||||||
|
callback_context.agent_name,
|
||||||
|
bool(text_context),
|
||||||
|
)
|
||||||
|
return LlmResponse(
|
||||||
|
content=types.Content(
|
||||||
|
role="model",
|
||||||
|
parts=[types.Part(text=reply)],
|
||||||
|
),
|
||||||
|
turn_complete=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── 3. 过滤历史消息中的非文本 part(防止累积) ──
|
||||||
|
filtered = False
|
||||||
|
for content in contents:
|
||||||
|
parts = getattr(content, "parts", None) or []
|
||||||
|
if len(parts) > 1:
|
||||||
|
new_parts = _filter_non_text_parts(parts)
|
||||||
|
if len(new_parts) != len(parts):
|
||||||
|
content.parts = new_parts
|
||||||
|
filtered = True
|
||||||
|
|
||||||
|
if filtered:
|
||||||
|
logger.info(
|
||||||
|
"已过滤历史消息中的非文本内容 | agent=%s",
|
||||||
|
callback_context.agent_name,
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── 4. 防提示词注入检测(基于过滤后的内容) ──
|
||||||
if last_user_content:
|
if last_user_content:
|
||||||
text = ""
|
text = ""
|
||||||
parts = getattr(last_user_content, "parts", None) or []
|
parts = getattr(last_user_content, "parts", None) or []
|
||||||
@@ -56,7 +141,7 @@ async def _before_model_callback(callback_context, llm_request):
|
|||||||
turn_complete=True,
|
turn_complete=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
# 2. 记录开始时间(供 after_model_callback 使用)
|
# ── 5. 记录开始时间 ──
|
||||||
callback_context.state["_log_model_call_start_time"] = time.perf_counter()
|
callback_context.state["_log_model_call_start_time"] = time.perf_counter()
|
||||||
return None
|
return None
|
||||||
|
|
||||||
@@ -69,6 +154,7 @@ class HuiYuBaseAgent(LlmAgent):
|
|||||||
所有 Agent 的基类,自动集成:
|
所有 Agent 的基类,自动集成:
|
||||||
- 防提示词注入检测
|
- 防提示词注入检测
|
||||||
- 模型调用耗时日志
|
- 模型调用耗时日志
|
||||||
|
- 图片/非文本内容过滤
|
||||||
|
|
||||||
用法:
|
用法:
|
||||||
agent = HuiYuBaseAgent(
|
agent = HuiYuBaseAgent(
|
||||||
@@ -86,7 +172,7 @@ class HuiYuBaseAgent(LlmAgent):
|
|||||||
user_before = kwargs.pop("before_model_callback", None)
|
user_before = kwargs.pop("before_model_callback", None)
|
||||||
user_after = kwargs.pop("after_model_callback", None)
|
user_after = kwargs.pop("after_model_callback", None)
|
||||||
|
|
||||||
# 构建回调链:防注入 → 用户回调 → 日志
|
# 构建回调链:防注入+图片过滤 → 用户回调 → 日志
|
||||||
before_chain = [_before_model_callback]
|
before_chain = [_before_model_callback]
|
||||||
if user_before:
|
if user_before:
|
||||||
before_chain.append(user_before)
|
before_chain.append(user_before)
|
||||||
|
|||||||
Reference in New Issue
Block a user