102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
"""
|
||
统一日志模块
|
||
|
||
为所有 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,
|
||
}
|