init agent project

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"]

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,
+    }

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