From ed2f002c42bcf0ff22e82d6e2b75e3faaf4fcd85 Mon Sep 17 00:00:00 2001
From: nelson <1475262689@qq.com>
Date: Wed, 10 Jun 2026 19:41:56 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E5=9B=BE=E7=89=87=E5=8D=A1?=
 =?UTF-8?q?=E4=BD=8Fagent=E7=9A=84bug?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 agent_base/agent_base.py | 102 ++++++++++++++++++++++++++++++++++++---
 1 file changed, 94 insertions(+), 8 deletions(-)

diff --git a/agent_base/agent_base.py b/agent_base/agent_base.py
index 3c100a3..46b4298 100644
--- a/agent_base/agent_base.py
+++ b/agent_base/agent_base.py
@@ -4,6 +4,7 @@
 提供所有 Agent 的基类，内置：
 - 防提示词注入检测
 - 模型调用耗时日志
+- 图片/非文本内容过滤（防止 token 爆炸和模型不支持报错）
 """
 
 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):
-    """在模型调用前检测提示词注入，并记录开始时间"""
-    # 1. 防提示词注入检测：检查最新的用户消息
+    """在模型调用前：
+    1. 过滤历史消息中的非文本内容（防止图片累积导致 token 爆炸）
+    2. 检测当前消息是否包含图片（不支持则直接回复提示）
+    3. 防提示词注入检测
+    4. 记录开始时间
+    """
     contents = getattr(llm_request, "contents", None) or []
-    # 找到最后一条 role="user" 的消息
+
+    # ── 1. 找到当前（最后一条）用户消息 ──
     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":
             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:
         text = ""
         parts = getattr(last_user_content, "parts", None) or []
@@ -56,7 +141,7 @@ async def _before_model_callback(callback_context, llm_request):
                 turn_complete=True,
             )
 
-    # 2. 记录开始时间（供 after_model_callback 使用）
+    # ── 5. 记录开始时间 ──
     callback_context.state["_log_model_call_start_time"] = time.perf_counter()
     return None
 
@@ -69,6 +154,7 @@ class HuiYuBaseAgent(LlmAgent):
     所有 Agent 的基类，自动集成：
     - 防提示词注入检测
     - 模型调用耗时日志
+    - 图片/非文本内容过滤
 
     用法：
         agent = HuiYuBaseAgent(
@@ -86,7 +172,7 @@ class HuiYuBaseAgent(LlmAgent):
         user_before = kwargs.pop("before_model_callback", None)
         user_after = kwargs.pop("after_model_callback", None)
 
-        # 构建回调链：防注入 → 用户回调 → 日志
+        # 构建回调链：防注入+图片过滤 → 用户回调 → 日志
         before_chain = [_before_model_callback]
         if user_before:
             before_chain.append(user_before)