init agent project
This commit is contained in:
15
.env.example
Normal file
15
.env.example
Normal file
@@ -0,0 +1,15 @@
|
|||||||
|
# ============================================================
|
||||||
|
# Root Agent 配置文件
|
||||||
|
# ============================================================
|
||||||
|
# 使用方法:复制本文件为 .env 并填入你的真实配置
|
||||||
|
# cp .env.example .env
|
||||||
|
# ============================================================
|
||||||
|
|
||||||
|
# MiniMax API Key(请替换为你自己的 API Key)
|
||||||
|
MINIMAX_API_KEY=YOUR_MINIMAX_API_KEY
|
||||||
|
|
||||||
|
# MiniMax API 地址(OpenAI 兼容接口)
|
||||||
|
MINIMAX_API_BASE=https://api.minimaxi.com/v1
|
||||||
|
|
||||||
|
# MiniMax 模型名称(例如: MiniMax-Text-01, abab-6.5s-chat, MiniMax-M1 等)
|
||||||
|
MINIMAX_MODEL=openai/MiniMax-Text-01
|
||||||
6
.gitignore
vendored
6
.gitignore
vendored
@@ -1,3 +1,9 @@
|
|||||||
|
# ---> macOS
|
||||||
|
.DS_Store
|
||||||
|
|
||||||
|
# ---> ADK
|
||||||
|
.adk/
|
||||||
|
|
||||||
# ---> Python
|
# ---> Python
|
||||||
# Byte-compiled / optimized / DLL files
|
# Byte-compiled / optimized / DLL files
|
||||||
__pycache__/
|
__pycache__/
|
||||||
|
|||||||
146
README.md
146
README.md
@@ -1,3 +1,145 @@
|
|||||||
# Agents
|
# 慧遇 Agent
|
||||||
|
|
||||||
慧遇所有的agent
|
慧遇智能体项目,基于 Google ADK (Agent Development Kit) 构建。
|
||||||
|
|
||||||
|
## 项目结构
|
||||||
|
|
||||||
|
```
|
||||||
|
workspace/
|
||||||
|
├── agents/ # 所有 Agent 模块
|
||||||
|
│ ├── __init__.py
|
||||||
|
│ └── root_agent/ # Root Agent
|
||||||
|
│ ├── __init__.py
|
||||||
|
│ └── agent.py
|
||||||
|
├── base/ # 基础工具类
|
||||||
|
│ ├── __init__.py
|
||||||
|
│ └── agent.py # HuiYuBaseAgent 基类
|
||||||
|
├── common/ # 公共模块
|
||||||
|
│ ├── __init__.py # 环境初始化(.env 加载、日志配置)
|
||||||
|
│ ├── logger.py # 模型调用耗时日志
|
||||||
|
│ └── prompt_guard.py # 防提示词注入检测
|
||||||
|
├── .env # 环境变量配置(不提交到 Git)
|
||||||
|
├── .env.example # 配置示例(可提交到 Git)
|
||||||
|
├── .gitignore
|
||||||
|
└── README.md
|
||||||
|
```
|
||||||
|
|
||||||
|
## 环境要求
|
||||||
|
|
||||||
|
- Python >= 3.10
|
||||||
|
- pip
|
||||||
|
|
||||||
|
## 环境搭建
|
||||||
|
|
||||||
|
### 1. 安装依赖
|
||||||
|
|
||||||
|
```bash
|
||||||
|
pip install "google-adk[extensions]" python-dotenv
|
||||||
|
```
|
||||||
|
|
||||||
|
### 2. 配置环境变量
|
||||||
|
|
||||||
|
复制示例配置文件并填入你的 API Key:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cp .env.example .env
|
||||||
|
```
|
||||||
|
|
||||||
|
编辑 `.env` 文件:
|
||||||
|
|
||||||
|
```env
|
||||||
|
# MiniMax API Key(替换为你自己的)
|
||||||
|
MINIMAX_API_KEY=your_api_key_here
|
||||||
|
|
||||||
|
# MiniMax API 地址
|
||||||
|
MINIMAX_API_BASE=https://api.minimaxi.com/v1
|
||||||
|
|
||||||
|
# 模型名称
|
||||||
|
MINIMAX_MODEL=openai/MiniMax-M2.7
|
||||||
|
```
|
||||||
|
|
||||||
|
### 3. 启动服务
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd agents
|
||||||
|
adk web
|
||||||
|
```
|
||||||
|
|
||||||
|
启动后访问 http://127.0.0.1:8000,在下拉菜单中选择 Agent 即可开始对话。
|
||||||
|
|
||||||
|
## 内置能力
|
||||||
|
|
||||||
|
所有继承 `HuiYuBaseAgent` 的 Agent 自动获得以下能力:
|
||||||
|
|
||||||
|
### 防提示词注入
|
||||||
|
|
||||||
|
在模型调用前自动检测用户输入中的提示词注入攻击,包括:
|
||||||
|
- 角色扮演 / 身份覆盖
|
||||||
|
- 指令泄露尝试
|
||||||
|
- 分隔符注入
|
||||||
|
- 越狱尝试
|
||||||
|
|
||||||
|
检测到注入时会拒绝请求并返回提示。
|
||||||
|
|
||||||
|
### 模型调用日志
|
||||||
|
|
||||||
|
每次模型调用完成后自动记录:
|
||||||
|
- Agent 名称
|
||||||
|
- 模型版本
|
||||||
|
- 调用耗时
|
||||||
|
- Token 使用量(prompt / completion / total)
|
||||||
|
|
||||||
|
日志示例:
|
||||||
|
|
||||||
|
```
|
||||||
|
2026-04-06 07:54:59 [adk.agent] INFO - 模型调用完成 | agent=minimax_agent model=MiniMax-M2.7 latency=11.701s prompt_tokens=78 completion_tokens=31 total_tokens=109
|
||||||
|
```
|
||||||
|
|
||||||
|
## 新增 Agent
|
||||||
|
|
||||||
|
在 `agents/` 目录下创建新的子目录即可:
|
||||||
|
|
||||||
|
```
|
||||||
|
agents/
|
||||||
|
├── root_agent/
|
||||||
|
│ ├── __init__.py
|
||||||
|
│ └── agent.py
|
||||||
|
└── my_new_agent/ # 新增 Agent
|
||||||
|
├── __init__.py # 内容同 root_agent/__init__.py
|
||||||
|
└── agent.py
|
||||||
|
```
|
||||||
|
|
||||||
|
`__init__.py` 内容(用于设置 Python 路径):
|
||||||
|
|
||||||
|
```python
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
|
||||||
|
|
||||||
|
from . import agent
|
||||||
|
|
||||||
|
__all__ = ["agent"]
|
||||||
|
```
|
||||||
|
|
||||||
|
`agent.py` 中继承 `HuiYuBaseAgent`:
|
||||||
|
|
||||||
|
```python
|
||||||
|
from base.agent import HuiYuBaseAgent
|
||||||
|
from google.adk.models.lite_llm import LiteLlm
|
||||||
|
import os
|
||||||
|
|
||||||
|
model = LiteLlm(
|
||||||
|
model=os.getenv("MINIMAX_MODEL"),
|
||||||
|
api_base=os.getenv("MINIMAX_API_BASE"),
|
||||||
|
api_key=os.getenv("MINIMAX_API_KEY"),
|
||||||
|
)
|
||||||
|
|
||||||
|
my_agent = HuiYuBaseAgent(
|
||||||
|
name="my_agent",
|
||||||
|
model=model,
|
||||||
|
instruction="你是一个专业的助手。",
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
重启 `adk web` 后,新 Agent 会自动出现在下拉菜单中。
|
||||||
|
|||||||
5
agents/__init__.py
Normal file
5
agents/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
|||||||
|
"""
|
||||||
|
Agent 模块
|
||||||
|
|
||||||
|
每个 Agent 是一个独立的子目录,包含 agent.py 和 __init__.py。
|
||||||
|
"""
|
||||||
10
agents/root_agent/__init__.py
Normal file
10
agents/root_agent/__init__.py
Normal file
@@ -0,0 +1,10 @@
|
|||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# 将项目根目录加入 sys.path(adk web 从 agents/ 启动)
|
||||||
|
# __file__ = agents/root_agent/__init__.py → parent.parent.parent = workspace/
|
||||||
|
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
|
||||||
|
|
||||||
|
from . import agent
|
||||||
|
|
||||||
|
__all__ = ["agent"]
|
||||||
93
agents/root_agent/agent.py
Normal file
93
agents/root_agent/agent.py
Normal file
@@ -0,0 +1,93 @@
|
|||||||
|
import os
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
from base.agent import HuiYuBaseAgent
|
||||||
|
from google.adk.models.lite_llm import LiteLlm
|
||||||
|
from google.adk.sessions import InMemorySessionService
|
||||||
|
from google.adk.runners import Runner
|
||||||
|
from google.genai import types
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# MiniMax 模型配置(从 .env 文件读取)
|
||||||
|
# ============================================================
|
||||||
|
MINIMAX_API_KEY = os.getenv("MINIMAX_API_KEY")
|
||||||
|
MINIMAX_API_BASE = os.getenv("MINIMAX_API_BASE")
|
||||||
|
MINIMAX_MODEL = os.getenv("MINIMAX_MODEL")
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 创建 LiteLlm 模型适配器
|
||||||
|
# ============================================================
|
||||||
|
model = LiteLlm(
|
||||||
|
model=MINIMAX_MODEL,
|
||||||
|
api_base=MINIMAX_API_BASE,
|
||||||
|
api_key=MINIMAX_API_KEY,
|
||||||
|
)
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 定义 Agent(继承 HuiYuBaseAgent,自动获得防注入 + 日志能力)
|
||||||
|
# ============================================================
|
||||||
|
root_agent = HuiYuBaseAgent(
|
||||||
|
name="minimax_agent",
|
||||||
|
model=model,
|
||||||
|
description="一个使用 MiniMax 模型的智能助手,能够回答用户的各种问题。",
|
||||||
|
instruction="你是一个乐于助人的中文 AI 助手,由 MiniMax 模型驱动。请用中文回答用户的问题,回答要准确、简洁、友好。",
|
||||||
|
)
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 运行入口(用于 adk run / adk web)
|
||||||
|
# ============================================================
|
||||||
|
APP_NAME = "慧遇agent app"
|
||||||
|
USER_ID = "user_001"
|
||||||
|
SESSION_ID = "session_001"
|
||||||
|
|
||||||
|
|
||||||
|
async def main():
|
||||||
|
"""主函数:创建会话并运行 Agent 对话"""
|
||||||
|
# 1. 创建会话服务
|
||||||
|
session_service = InMemorySessionService()
|
||||||
|
|
||||||
|
# 2. 创建会话
|
||||||
|
session = await session_service.create_session(
|
||||||
|
app_name=APP_NAME,
|
||||||
|
user_id=USER_ID,
|
||||||
|
session_id=SESSION_ID,
|
||||||
|
)
|
||||||
|
print(f"会话已创建: App='{APP_NAME}', User='{USER_ID}', Session='{SESSION_ID}'")
|
||||||
|
|
||||||
|
# 3. 创建 Runner
|
||||||
|
runner = Runner(
|
||||||
|
agent=root_agent,
|
||||||
|
app_name=APP_NAME,
|
||||||
|
session_service=session_service,
|
||||||
|
)
|
||||||
|
print(f"Runner 已创建,Agent: '{runner.agent.name}'")
|
||||||
|
|
||||||
|
# 4. 交互式对话循环
|
||||||
|
print("\n===== MiniMax Agent 对话 (输入 'exit' 退出) =====")
|
||||||
|
while True:
|
||||||
|
query = input("\n[user]: ").strip()
|
||||||
|
if query.lower() in ("exit", "quit", "退出"):
|
||||||
|
print("再见!")
|
||||||
|
break
|
||||||
|
if not query:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 构造 ADK 消息
|
||||||
|
content = types.Content(role="user", parts=[types.Part(text=query)])
|
||||||
|
|
||||||
|
# 调用 Agent
|
||||||
|
final_response_text = ""
|
||||||
|
async for event in runner.run_async(
|
||||||
|
user_id=USER_ID,
|
||||||
|
session_id=SESSION_ID,
|
||||||
|
new_message=content,
|
||||||
|
):
|
||||||
|
if event.is_final_response():
|
||||||
|
if event.content and event.content.parts:
|
||||||
|
final_response_text = event.content.parts[0].text
|
||||||
|
|
||||||
|
print(f"\n[agent]: {final_response_text}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
asyncio.run(main())
|
||||||
5
base/__init__.py
Normal file
5
base/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
|||||||
|
"""
|
||||||
|
基础工具模块
|
||||||
|
|
||||||
|
存放与模型、配置、工具等相关的底层工具类和函数。
|
||||||
|
"""
|
||||||
102
base/agent.py
Normal file
102
base/agent.py
Normal file
@@ -0,0 +1,102 @@
|
|||||||
|
"""
|
||||||
|
基础 Agent 模块
|
||||||
|
|
||||||
|
提供所有 Agent 的基类,内置:
|
||||||
|
- 防提示词注入检测
|
||||||
|
- 模型调用耗时日志
|
||||||
|
"""
|
||||||
|
|
||||||
|
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")
|
||||||
|
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 防注入的 before_model_callback
|
||||||
|
# ============================================================
|
||||||
|
async def _before_model_callback(callback_context, llm_request):
|
||||||
|
"""在模型调用前检测提示词注入,并记录开始时间"""
|
||||||
|
# 1. 防提示词注入检测:检查最新的用户消息
|
||||||
|
contents = getattr(llm_request, "contents", None) or []
|
||||||
|
# 找到最后一条 role="user" 的消息
|
||||||
|
last_user_content = None
|
||||||
|
for content in reversed(contents):
|
||||||
|
if getattr(content, "role", None) == "user":
|
||||||
|
last_user_content = content
|
||||||
|
break
|
||||||
|
|
||||||
|
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,
|
||||||
|
)
|
||||||
|
|
||||||
|
# 2. 记录开始时间(供 after_model_callback 使用)
|
||||||
|
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
|
||||||
|
|
||||||
|
super().__init__(**kwargs)
|
||||||
35
common/__init__.py
Normal file
35
common/__init__.py
Normal file
@@ -0,0 +1,35 @@
|
|||||||
|
"""
|
||||||
|
公共模块
|
||||||
|
|
||||||
|
提供所有 Agent 共用的基础配置和工具:
|
||||||
|
- 环境变量加载(.env)
|
||||||
|
- 日志配置
|
||||||
|
- 模型调用日志回调
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
import logging
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 项目根目录 & 环境初始化
|
||||||
|
# ============================================================
|
||||||
|
_PROJECT_ROOT = Path(__file__).parent.parent
|
||||||
|
|
||||||
|
# 加载 .env 文件
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
|
||||||
|
load_dotenv(_PROJECT_ROOT / ".env")
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 日志配置
|
||||||
|
# ============================================================
|
||||||
|
logging.basicConfig(
|
||||||
|
level=logging.INFO,
|
||||||
|
format="%(asctime)s [%(name)s] %(levelname)s - %(message)s",
|
||||||
|
datefmt="%Y-%m-%d %H:%M:%S",
|
||||||
|
)
|
||||||
|
|
||||||
|
from .logger import get_model_callbacks # noqa: E402
|
||||||
|
|
||||||
|
__all__ = ["get_model_callbacks"]
|
||||||
101
common/logger.py
Normal file
101
common/logger.py
Normal file
@@ -0,0 +1,101 @@
|
|||||||
|
"""
|
||||||
|
统一日志模块
|
||||||
|
|
||||||
|
为所有 Agent 提供一致的模型调用日志记录,包括:
|
||||||
|
- 模型调用耗时
|
||||||
|
- Token 使用量(prompt / completion)
|
||||||
|
- 模型版本信息
|
||||||
|
- Agent 名称
|
||||||
|
|
||||||
|
使用方式:
|
||||||
|
from common.logger import get_model_callbacks
|
||||||
|
|
||||||
|
agent = Agent(
|
||||||
|
name="my_agent",
|
||||||
|
model=model,
|
||||||
|
instruction="...",
|
||||||
|
**get_model_callbacks(),
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
|
||||||
|
import time
|
||||||
|
import logging
|
||||||
|
|
||||||
|
logger = logging.getLogger("adk.agent")
|
||||||
|
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 模型调用耗时记录的 state key
|
||||||
|
# ============================================================
|
||||||
|
_STATE_KEY_START_TIME = "_log_model_call_start_time"
|
||||||
|
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# before_model_callback:记录模型调用开始时间
|
||||||
|
# ============================================================
|
||||||
|
async def before_model_callback(callback_context, llm_request):
|
||||||
|
"""在模型调用前记录开始时间到 session state"""
|
||||||
|
callback_context.state[_STATE_KEY_START_TIME] = time.perf_counter()
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# after_model_callback:计算耗时并输出日志
|
||||||
|
# ============================================================
|
||||||
|
async def after_model_callback(callback_context, llm_response):
|
||||||
|
"""在模型调用完成后计算耗时,输出结构化日志"""
|
||||||
|
# 流式模式下只在最终响应时记录(避免重复打印)
|
||||||
|
if getattr(llm_response, "partial", None) and not getattr(
|
||||||
|
llm_response, "turn_complete", None
|
||||||
|
):
|
||||||
|
return None
|
||||||
|
|
||||||
|
start_time = callback_context.state.get(_STATE_KEY_START_TIME)
|
||||||
|
if start_time is None:
|
||||||
|
return None
|
||||||
|
callback_context.state[_STATE_KEY_START_TIME] = None
|
||||||
|
|
||||||
|
elapsed = time.perf_counter() - start_time
|
||||||
|
agent_name = callback_context.agent_name
|
||||||
|
model_version = getattr(llm_response, "model_version", None) or "unknown"
|
||||||
|
usage = getattr(llm_response, "usage_metadata", None)
|
||||||
|
|
||||||
|
prompt_tokens = getattr(usage, "prompt_token_count", None) if usage else None
|
||||||
|
candidates_tokens = (
|
||||||
|
getattr(usage, "candidates_token_count", None) if usage else None
|
||||||
|
)
|
||||||
|
total_tokens = getattr(usage, "total_token_count", None) if usage else None
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"模型调用完成 | agent=%s model=%s latency=%.3fs "
|
||||||
|
"prompt_tokens=%s completion_tokens=%s total_tokens=%s",
|
||||||
|
agent_name,
|
||||||
|
model_version,
|
||||||
|
elapsed,
|
||||||
|
prompt_tokens,
|
||||||
|
candidates_tokens,
|
||||||
|
total_tokens,
|
||||||
|
)
|
||||||
|
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 便捷函数:获取回调字典,用于展开到 Agent 构造参数
|
||||||
|
# ============================================================
|
||||||
|
def get_model_callbacks():
|
||||||
|
"""
|
||||||
|
返回 before_model_callback 和 after_model_callback 的字典。
|
||||||
|
|
||||||
|
用法:
|
||||||
|
agent = Agent(
|
||||||
|
name="my_agent",
|
||||||
|
model=model,
|
||||||
|
instruction="...",
|
||||||
|
**get_model_callbacks(),
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
return {
|
||||||
|
"before_model_callback": before_model_callback,
|
||||||
|
"after_model_callback": after_model_callback,
|
||||||
|
}
|
||||||
66
common/prompt_guard.py
Normal file
66
common/prompt_guard.py
Normal file
@@ -0,0 +1,66 @@
|
|||||||
|
"""
|
||||||
|
防提示词注入模块
|
||||||
|
|
||||||
|
在模型调用前检测并拦截潜在的提示词注入攻击。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import re
|
||||||
|
import logging
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
logger = logging.getLogger("adk.agent")
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 常见提示词注入模式
|
||||||
|
# ============================================================
|
||||||
|
_INJECTION_PATTERNS = [
|
||||||
|
# 角色扮演/身份覆盖
|
||||||
|
r"ignore\s+(all\s+)?previous\s+(instructions?|prompts?|rules?)",
|
||||||
|
r"forget\s+(all\s+)?previous\s+(instructions?|prompts?|rules?)",
|
||||||
|
r"disregard\s+(all\s+)?previous\s+(instructions?|prompts?|rules?)",
|
||||||
|
r"you\s+are\s+now\s+a",
|
||||||
|
r"pretend\s+(you\s+are|to\s+be)",
|
||||||
|
r"act\s+as\s+(if\s+you\s+(are|were)|a|an)",
|
||||||
|
r"roleplay\s+as",
|
||||||
|
r"你是一个",
|
||||||
|
r"假装你是",
|
||||||
|
r"扮演一个",
|
||||||
|
# 指令泄露
|
||||||
|
r"(show|reveal|display|print|output|dump)\s+(me\s+)?(your|the)?\s*system\s*(prompt|instructions?|rules?|config)",
|
||||||
|
r"(show|reveal|display|print|output|dump)\s+(me\s+)?your\s+(prompt|instructions?|rules?|config)",
|
||||||
|
r"what\s+(are|is)\s+your\s+(instructions?|prompts?|rules?|system\s*prompt)",
|
||||||
|
r"repeat\s+(your|the|back)\s+(instructions?|prompts?|system)",
|
||||||
|
r"(显示|输出|打印|告诉我|泄露)\s*(你的|系统|隐藏)",
|
||||||
|
# 分隔符注入
|
||||||
|
r"---\s*(system|instruction|prompt)",
|
||||||
|
r"###\s*(system|instruction|prompt)",
|
||||||
|
r"<\|im_start\|>",
|
||||||
|
r"<\|im_end\|>",
|
||||||
|
# 越狱尝试
|
||||||
|
r"(jailbreak|dan\s+mode|developer\s+mode)\b",
|
||||||
|
r"bypass\s+(safety|filter|security|restriction)",
|
||||||
|
r"no\s+(safety|ethical|moral|content)\s+(filter|restriction|limit)",
|
||||||
|
]
|
||||||
|
|
||||||
|
_COMPILED_PATTERNS = [re.compile(p, re.IGNORECASE) for p in _INJECTION_PATTERNS]
|
||||||
|
|
||||||
|
|
||||||
|
def check_prompt_injection(text: str) -> Optional[str]:
|
||||||
|
"""
|
||||||
|
检测文本中是否包含提示词注入。
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: 待检测的用户输入文本
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
如果检测到注入,返回匹配到的模式描述;否则返回 None
|
||||||
|
"""
|
||||||
|
if not text:
|
||||||
|
return None
|
||||||
|
|
||||||
|
for pattern in _COMPILED_PATTERNS:
|
||||||
|
match = pattern.search(text)
|
||||||
|
if match:
|
||||||
|
return match.group(0)
|
||||||
|
|
||||||
|
return None
|
||||||
Reference in New Issue
Block a user