198 lines
7.2 KiB
Python
198 lines
7.2 KiB
Python
"""
|
||
基础 Agent 模块
|
||
|
||
提供所有 Agent 的基类，内置：
|
||
- 防提示词注入检测
|
||
- 模型调用耗时日志
|
||
- 图片/非文本内容过滤（防止 token 爆炸和模型不支持报错）
|
||
"""
|
||
|
||
import time
|
||
import logging
|
||
|
||
from google.adk.agents import LlmAgent
|
||
from google.adk.models.llm_response import LlmResponse
|
||
from google.genai import types
|
||
|
||
from common.logger import after_model_callback
|
||
from common.prompt_guard import check_prompt_injection
|
||
|
||
logger = logging.getLogger("adk.agent")
|
||
|
||
|
||
# ============================================================
|
||
# 工具函数：判断一个 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. 过滤历史消息中的非文本内容（防止图片累积导致 token 爆炸）
|
||
2. 检测当前消息是否包含图片（不支持则直接回复提示）
|
||
3. 防提示词注入检测
|
||
4. 记录开始时间
|
||
"""
|
||
contents = getattr(llm_request, "contents", None) or []
|
||
|
||
# ── 1. 找到当前（最后一条）用户消息 ──
|
||
last_user_content = None
|
||
last_user_index = -1
|
||
for i, content in enumerate(contents):
|
||
if getattr(content, "role", None) == "user":
|
||
last_user_content = content
|
||
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 []
|
||
for part in parts:
|
||
part_text = getattr(part, "text", None)
|
||
if part_text:
|
||
text += part_text
|
||
|
||
injection = check_prompt_injection(text)
|
||
if injection:
|
||
logger.warning(
|
||
"检测到提示词注入 | agent=%s pattern='%s'",
|
||
callback_context.agent_name,
|
||
injection,
|
||
)
|
||
return LlmResponse(
|
||
content=types.Content(
|
||
role="model",
|
||
parts=[types.Part(text="抱歉，我无法处理该请求。")],
|
||
),
|
||
turn_complete=True,
|
||
)
|
||
|
||
# ── 5. 记录开始时间 ──
|
||
callback_context.state["_log_model_call_start_time"] = time.perf_counter()
|
||
return None
|
||
|
||
|
||
# ============================================================
|
||
# HuiYuBaseAgent 基类
|
||
# ============================================================
|
||
class HuiYuBaseAgent(LlmAgent):
|
||
"""
|
||
所有 Agent 的基类，自动集成：
|
||
- 防提示词注入检测
|
||
- 模型调用耗时日志
|
||
- 图片/非文本内容过滤
|
||
|
||
用法：
|
||
agent = HuiYuBaseAgent(
|
||
name="my_agent",
|
||
model=model,
|
||
instruction="...",
|
||
)
|
||
|
||
如果需要自定义回调，仍然可以传入 before_model_callback / after_model_callback，
|
||
它们会在防注入检测和日志记录之间执行。
|
||
"""
|
||
|
||
def __init__(self, **kwargs):
|
||
# 获取用户传入的回调
|
||
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)
|
||
|
||
after_chain = []
|
||
if user_after:
|
||
after_chain.append(user_after)
|
||
after_chain.append(after_model_callback)
|
||
|
||
kwargs["before_model_callback"] = before_chain
|
||
kwargs["after_model_callback"] = after_chain
|
||
|
||
# 强制模型用中文思考：在 instruction 前附加系统级要求
|
||
instruction = kwargs.get("instruction", "")
|
||
if instruction:
|
||
kwargs["instruction"] = (
|
||
"【系统要求：你的所有思考过程、分析推理、内心独白必须使用中文，"
|
||
"禁止在思考中使用英文。你的最终回复也使用中文。】\n\n"
|
||
+ instruction
|
||
)
|
||
|
||
super().__init__(**kwargs)
|