init agent project

base/__init__.py

@@ -0,0 +1,5 @@
+"""
+基础工具模块
+
+存放与模型、配置、工具等相关的底层工具类和函数。
+"""

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)