diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..871cc97 --- /dev/null +++ b/.env.example @@ -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 diff --git a/.gitignore b/.gitignore index 36b13f1..6113ee5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,9 @@ +# ---> macOS +.DS_Store + +# ---> ADK +.adk/ + # ---> Python # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/README.md b/README.md index fc115c4..751c454 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,145 @@ -# Agents +# 慧遇 Agent -慧遇所有的agent \ No newline at end of file +慧遇智能体项目，基于 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 会自动出现在下拉菜单中。 diff --git a/agents/__init__.py b/agents/__init__.py new file mode 100644 index 0000000..bbddce9 --- /dev/null +++ b/agents/__init__.py @@ -0,0 +1,5 @@ +""" +Agent 模块 + +每个 Agent 是一个独立的子目录，包含 agent.py 和 __init__.py。 +""" diff --git a/agents/root_agent/__init__.py b/agents/root_agent/__init__.py new file mode 100644 index 0000000..d10156d --- /dev/null +++ b/agents/root_agent/__init__.py @@ -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"] diff --git a/agents/root_agent/agent.py b/agents/root_agent/agent.py new file mode 100644 index 0000000..fc34ee8 --- /dev/null +++ b/agents/root_agent/agent.py @@ -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()) diff --git a/base/__init__.py b/base/__init__.py new file mode 100644 index 0000000..fba4eba --- /dev/null +++ b/base/__init__.py @@ -0,0 +1,5 @@ +""" +基础工具模块 + +存放与模型、配置、工具等相关的底层工具类和函数。 +""" diff --git a/base/agent.py b/base/agent.py new file mode 100644 index 0000000..e68a32b --- /dev/null +++ b/base/agent.py @@ -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) diff --git a/common/__init__.py b/common/__init__.py new file mode 100644 index 0000000..2594368 --- /dev/null +++ b/common/__init__.py @@ -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"] diff --git a/common/logger.py b/common/logger.py new file mode 100644 index 0000000..9f510f0 --- /dev/null +++ b/common/logger.py @@ -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, + } diff --git a/common/prompt_guard.py b/common/prompt_guard.py new file mode 100644 index 0000000..d12206a --- /dev/null +++ b/common/prompt_guard.py @@ -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