init agent project

2026-04-06 17:51:06 +08:00
parent 8daeef22e7
commit 7b7a2af4c6
11 changed files with 582 additions and 2 deletions
--- a/.env.example
+++ 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
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,9 @@
+# ---> macOS
+.DS_Store
+
+# ---> ADK
+.adk/
+
 # ---> Python
 # Byte-compiled / optimized / DLL files
 __pycache__/
--- a/README.md
+++ b/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 会自动出现在下拉菜单中。
--- a/agents/init.py
+++ b/agents/init.py
@@ -0,0 +1,5 @@
+"""
+Agent 模块
+
+每个 Agent 是一个独立的子目录，包含 agent.py 和 __init__.py。
+"""
--- a/agents/root_agent/init.py
+++ 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"]
--- a/agents/root_agent/agent.py
+++ 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())
--- a/base/init.py
+++ b/base/init.py
@@ -0,0 +1,5 @@
+"""
+基础工具模块
+
+存放与模型、配置、工具等相关的底层工具类和函数。
+"""
--- a/base/agent.py
+++ 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)
--- a/common/init.py
+++ 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"]
--- a/common/logger.py
+++ 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,
+    }
--- a/common/prompt_guard.py
+++ 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