forked from xiaohei/taiji-AI-PAD
498 lines
14 KiB
Python
498 lines
14 KiB
Python
"""
|
|
Agent轨迹追踪模块
|
|
执行轨迹记录与查询
|
|
"""
|
|
|
|
import uuid
|
|
from datetime import datetime, timedelta
|
|
from decimal import Decimal
|
|
from typing import Dict, List, Optional, Any
|
|
from sqlalchemy import select, func, and_, desc
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
import structlog
|
|
|
|
from models import AgentTrace, Execution, Agent, User
|
|
|
|
logger = structlog.get_logger(__name__)
|
|
|
|
|
|
class TraceContext:
|
|
"""追踪上下文"""
|
|
|
|
def __init__(self, trace_id: Optional[str] = None):
|
|
self.trace_id = trace_id or str(uuid.uuid4())
|
|
self.spans: List[Dict] = []
|
|
self._current_span_id: Optional[str] = None
|
|
|
|
def start_span(
|
|
self,
|
|
operation_name: str,
|
|
operation_type: str,
|
|
parent_span_id: Optional[str] = None
|
|
) -> str:
|
|
"""开始一个新的span"""
|
|
span_id = str(uuid.uuid4())
|
|
span = {
|
|
"span_id": span_id,
|
|
"parent_span_id": parent_span_id or self._current_span_id,
|
|
"operation_name": operation_name,
|
|
"operation_type": operation_type,
|
|
"started_at": datetime.utcnow(),
|
|
"ended_at": None,
|
|
"status": "running",
|
|
"input_data": None,
|
|
"output_data": None,
|
|
"error_type": None,
|
|
"error_message": None,
|
|
}
|
|
self.spans.append(span)
|
|
self._current_span_id = span_id
|
|
return span_id
|
|
|
|
def end_span(
|
|
self,
|
|
span_id: str,
|
|
status: str = "success",
|
|
output_data: Optional[Dict] = None,
|
|
error_type: Optional[str] = None,
|
|
error_message: Optional[str] = None
|
|
):
|
|
"""结束一个span"""
|
|
for span in self.spans:
|
|
if span["span_id"] == span_id:
|
|
span["ended_at"] = datetime.utcnow()
|
|
span["status"] = status
|
|
span["output_data"] = output_data
|
|
span["error_type"] = error_type
|
|
span["error_message"] = error_message
|
|
break
|
|
|
|
def set_span_input(self, span_id: str, input_data: Dict):
|
|
"""设置span的输入数据"""
|
|
for span in self.spans:
|
|
if span["span_id"] == span_id:
|
|
span["input_data"] = input_data
|
|
break
|
|
|
|
|
|
async def create_trace_record(
|
|
execution_id: str,
|
|
agent_id: str,
|
|
user_id: str,
|
|
trace_id: str,
|
|
span_id: str,
|
|
parent_span_id: Optional[str],
|
|
operation_name: str,
|
|
operation_type: str,
|
|
started_at: datetime,
|
|
ended_at: Optional[datetime],
|
|
status: str,
|
|
input_data: Optional[Dict],
|
|
output_data: Optional[Dict],
|
|
error_type: Optional[str],
|
|
error_message: Optional[str],
|
|
tokens_used: int,
|
|
eu_consumed: Decimal,
|
|
db: AsyncSession
|
|
) -> AgentTrace:
|
|
"""
|
|
创建追踪记录
|
|
|
|
Args:
|
|
execution_id: 执行ID
|
|
agent_id: Agent ID
|
|
user_id: 用户ID
|
|
trace_id: 追踪ID
|
|
span_id: Span ID
|
|
parent_span_id: 父Span ID
|
|
operation_name: 操作名称
|
|
operation_type: 操作类型
|
|
started_at: 开始时间
|
|
ended_at: 结束时间
|
|
status: 状态
|
|
input_data: 输入数据
|
|
output_data: 输出数据
|
|
error_type: 错误类型
|
|
error_message: 错误消息
|
|
tokens_used: 使用的Token数
|
|
eu_consumed: 消耗的EU
|
|
db: 数据库会话
|
|
|
|
Returns:
|
|
AgentTrace记录
|
|
"""
|
|
# 计算持续时间
|
|
duration_ms = None
|
|
if ended_at and started_at:
|
|
duration_ms = int((ended_at - started_at).total_seconds() * 1000)
|
|
|
|
# 对输入输出数据进行脱敏
|
|
sanitized_input = _sanitize_data(input_data) if input_data else None
|
|
sanitized_output = _sanitize_data(output_data) if output_data else None
|
|
|
|
trace = AgentTrace(
|
|
execution_id=execution_id,
|
|
agent_id=agent_id,
|
|
user_id=user_id,
|
|
trace_id=trace_id,
|
|
span_id=span_id,
|
|
parent_span_id=parent_span_id,
|
|
operation_name=operation_name,
|
|
operation_type=operation_type,
|
|
started_at=started_at,
|
|
ended_at=ended_at,
|
|
duration_ms=duration_ms,
|
|
input_data=sanitized_input,
|
|
output_data=sanitized_output,
|
|
status=status,
|
|
error_type=error_type,
|
|
error_message=error_message,
|
|
tokens_used=tokens_used,
|
|
eu_consumed=eu_consumed,
|
|
)
|
|
|
|
db.add(trace)
|
|
await db.commit()
|
|
await db.refresh(trace)
|
|
|
|
return trace
|
|
|
|
|
|
async def save_trace_context(
|
|
trace_context: TraceContext,
|
|
execution_id: str,
|
|
agent_id: str,
|
|
user_id: str,
|
|
db: AsyncSession
|
|
) -> int:
|
|
"""
|
|
保存完整的追踪上下文
|
|
|
|
Args:
|
|
trace_context: 追踪上下文
|
|
execution_id: 执行ID
|
|
agent_id: Agent ID
|
|
user_id: 用户ID
|
|
db: 数据库会话
|
|
|
|
Returns:
|
|
保存的span数量
|
|
"""
|
|
count = 0
|
|
for span in trace_context.spans:
|
|
await create_trace_record(
|
|
execution_id=execution_id,
|
|
agent_id=agent_id,
|
|
user_id=user_id,
|
|
trace_id=trace_context.trace_id,
|
|
span_id=span["span_id"],
|
|
parent_span_id=span["parent_span_id"],
|
|
operation_name=span["operation_name"],
|
|
operation_type=span["operation_type"],
|
|
started_at=span["started_at"],
|
|
ended_at=span["ended_at"],
|
|
status=span["status"],
|
|
input_data=span["input_data"],
|
|
output_data=span["output_data"],
|
|
error_type=span["error_type"],
|
|
error_message=span["error_message"],
|
|
tokens_used=0,
|
|
eu_consumed=Decimal(0),
|
|
db=db,
|
|
)
|
|
count += 1
|
|
|
|
return count
|
|
|
|
|
|
async def get_execution_trace(
|
|
execution_id: str,
|
|
db: AsyncSession
|
|
) -> Dict:
|
|
"""
|
|
获取执行的完整追踪
|
|
|
|
Args:
|
|
execution_id: 执行ID
|
|
db: 数据库会话
|
|
|
|
Returns:
|
|
执行追踪详情
|
|
"""
|
|
result = await db.execute(
|
|
select(AgentTrace)
|
|
.where(AgentTrace.execution_id == execution_id)
|
|
.order_by(AgentTrace.started_at)
|
|
)
|
|
traces = result.scalars().all()
|
|
|
|
if not traces:
|
|
return {"executionId": execution_id, "spans": []}
|
|
|
|
trace_id = traces[0].trace_id if traces else None
|
|
|
|
# 构建span树
|
|
spans = []
|
|
for trace in traces:
|
|
spans.append({
|
|
"spanId": trace.span_id,
|
|
"parentSpanId": trace.parent_span_id,
|
|
"operationName": trace.operation_name,
|
|
"operationType": trace.operation_type,
|
|
"startedAt": trace.started_at.isoformat(),
|
|
"endedAt": trace.ended_at.isoformat() if trace.ended_at else None,
|
|
"durationMs": trace.duration_ms,
|
|
"status": trace.status,
|
|
"inputData": trace.input_data,
|
|
"outputData": trace.output_data,
|
|
"errorType": trace.error_type,
|
|
"errorMessage": trace.error_message,
|
|
"tokensUsed": trace.tokens_used,
|
|
"euConsumed": float(trace.eu_consumed) if trace.eu_consumed else 0,
|
|
})
|
|
|
|
# 计算总持续时间和EU
|
|
total_duration = sum(s["durationMs"] or 0 for s in spans)
|
|
total_eu = sum(s["euConsumed"] for s in spans)
|
|
|
|
return {
|
|
"executionId": execution_id,
|
|
"traceId": trace_id,
|
|
"spans": spans,
|
|
"totalDurationMs": total_duration,
|
|
"totalEuConsumed": total_eu,
|
|
"spanCount": len(spans),
|
|
}
|
|
|
|
|
|
async def query_traces(
|
|
user_id: Optional[str] = None,
|
|
agent_id: Optional[str] = None,
|
|
status: Optional[str] = None,
|
|
start_date: Optional[datetime] = None,
|
|
end_date: Optional[datetime] = None,
|
|
page: int = 1,
|
|
page_size: int = 20,
|
|
db: AsyncSession = None
|
|
) -> Dict:
|
|
"""
|
|
查询追踪记录
|
|
|
|
Args:
|
|
user_id: 用户ID过滤
|
|
agent_id: Agent ID过滤
|
|
status: 状态过滤
|
|
start_date: 开始日期
|
|
end_date: 结束日期
|
|
page: 页码
|
|
page_size: 每页大小
|
|
db: 数据库会话
|
|
|
|
Returns:
|
|
查询结果
|
|
"""
|
|
# 按trace_id分组查询
|
|
query = select(
|
|
AgentTrace.trace_id,
|
|
AgentTrace.execution_id,
|
|
AgentTrace.agent_id,
|
|
AgentTrace.user_id,
|
|
func.min(AgentTrace.started_at).label("started_at"),
|
|
func.max(AgentTrace.ended_at).label("ended_at"),
|
|
func.count(AgentTrace.id).label("span_count"),
|
|
func.sum(AgentTrace.tokens_used).label("total_tokens"),
|
|
func.sum(AgentTrace.eu_consumed).label("total_eu"),
|
|
)
|
|
|
|
# 应用过滤条件
|
|
conditions = []
|
|
if user_id:
|
|
conditions.append(AgentTrace.user_id == user_id)
|
|
if agent_id:
|
|
conditions.append(AgentTrace.agent_id == agent_id)
|
|
if status:
|
|
conditions.append(AgentTrace.status == status)
|
|
if start_date:
|
|
conditions.append(AgentTrace.started_at >= start_date)
|
|
if end_date:
|
|
conditions.append(AgentTrace.started_at <= end_date)
|
|
|
|
if conditions:
|
|
query = query.where(and_(*conditions))
|
|
|
|
query = query.group_by(
|
|
AgentTrace.trace_id,
|
|
AgentTrace.execution_id,
|
|
AgentTrace.agent_id,
|
|
AgentTrace.user_id,
|
|
)
|
|
|
|
# 计算总数
|
|
count_query = select(func.count(func.distinct(AgentTrace.trace_id)))
|
|
if conditions:
|
|
count_query = count_query.where(and_(*conditions))
|
|
count_result = await db.execute(count_query)
|
|
total = count_result.scalar() or 0
|
|
|
|
# 分页
|
|
query = query.order_by(desc(func.min(AgentTrace.started_at)))
|
|
query = query.offset((page - 1) * page_size).limit(page_size)
|
|
|
|
result = await db.execute(query)
|
|
rows = result.all()
|
|
|
|
# 获取Agent名称
|
|
agent_ids = list(set(row.agent_id for row in rows if row.agent_id))
|
|
agents_map = {}
|
|
if agent_ids:
|
|
agents_result = await db.execute(
|
|
select(Agent).where(Agent.id.in_(agent_ids))
|
|
)
|
|
for agent in agents_result.scalars().all():
|
|
agents_map[str(agent.id)] = agent.name
|
|
|
|
traces = []
|
|
for row in rows:
|
|
duration_ms = None
|
|
if row.ended_at and row.started_at:
|
|
duration_ms = int((row.ended_at - row.started_at).total_seconds() * 1000)
|
|
|
|
traces.append({
|
|
"traceId": row.trace_id,
|
|
"executionId": str(row.execution_id),
|
|
"agentId": str(row.agent_id),
|
|
"agentName": agents_map.get(str(row.agent_id), "Unknown"),
|
|
"userId": str(row.user_id),
|
|
"startedAt": row.started_at.isoformat(),
|
|
"endedAt": row.ended_at.isoformat() if row.ended_at else None,
|
|
"durationMs": duration_ms,
|
|
"spanCount": row.span_count,
|
|
"totalTokens": int(row.total_tokens or 0),
|
|
"totalEu": float(row.total_eu or 0),
|
|
})
|
|
|
|
return {
|
|
"total": total,
|
|
"page": page,
|
|
"pageSize": page_size,
|
|
"totalPages": (total + page_size - 1) // page_size,
|
|
"traces": traces,
|
|
}
|
|
|
|
|
|
async def get_trace_stats(
|
|
user_id: Optional[str],
|
|
start_date: datetime,
|
|
end_date: datetime,
|
|
db: AsyncSession
|
|
) -> Dict:
|
|
"""
|
|
获取追踪统计
|
|
|
|
Args:
|
|
user_id: 用户ID
|
|
start_date: 开始日期
|
|
end_date: 结束日期
|
|
db: 数据库会话
|
|
|
|
Returns:
|
|
追踪统计
|
|
"""
|
|
conditions = [
|
|
AgentTrace.started_at >= start_date,
|
|
AgentTrace.started_at <= end_date,
|
|
]
|
|
if user_id:
|
|
conditions.append(AgentTrace.user_id == user_id)
|
|
|
|
# 总体统计
|
|
stats_result = await db.execute(
|
|
select(
|
|
func.count(func.distinct(AgentTrace.trace_id)).label("total_traces"),
|
|
func.count(AgentTrace.id).label("total_spans"),
|
|
func.sum(AgentTrace.duration_ms).label("total_duration"),
|
|
func.avg(AgentTrace.duration_ms).label("avg_duration"),
|
|
func.sum(AgentTrace.tokens_used).label("total_tokens"),
|
|
func.sum(AgentTrace.eu_consumed).label("total_eu"),
|
|
)
|
|
.where(and_(*conditions))
|
|
)
|
|
stats = stats_result.first()
|
|
|
|
# 按状态统计
|
|
status_stats = await db.execute(
|
|
select(
|
|
AgentTrace.status,
|
|
func.count(AgentTrace.id).label("count"),
|
|
)
|
|
.where(and_(*conditions))
|
|
.group_by(AgentTrace.status)
|
|
)
|
|
status_breakdown = {row.status: row.count for row in status_stats.all()}
|
|
|
|
# 按操作类型统计
|
|
type_stats = await db.execute(
|
|
select(
|
|
AgentTrace.operation_type,
|
|
func.count(AgentTrace.id).label("count"),
|
|
func.avg(AgentTrace.duration_ms).label("avg_duration"),
|
|
)
|
|
.where(and_(*conditions))
|
|
.group_by(AgentTrace.operation_type)
|
|
)
|
|
type_breakdown = [
|
|
{
|
|
"operationType": row.operation_type,
|
|
"count": row.count,
|
|
"avgDurationMs": float(row.avg_duration or 0),
|
|
}
|
|
for row in type_stats.all()
|
|
]
|
|
|
|
return {
|
|
"startDate": start_date.isoformat(),
|
|
"endDate": end_date.isoformat(),
|
|
"totalTraces": stats.total_traces or 0,
|
|
"totalSpans": stats.total_spans or 0,
|
|
"totalDurationMs": int(stats.total_duration or 0),
|
|
"avgDurationMs": float(stats.avg_duration or 0),
|
|
"totalTokens": int(stats.total_tokens or 0),
|
|
"totalEu": float(stats.total_eu or 0),
|
|
"byStatus": status_breakdown,
|
|
"byOperationType": type_breakdown,
|
|
}
|
|
|
|
|
|
def _sanitize_data(data: Dict) -> Dict:
|
|
"""
|
|
对敏感数据进行脱敏处理
|
|
|
|
Args:
|
|
data: 原始数据
|
|
|
|
Returns:
|
|
脱敏后的数据
|
|
"""
|
|
if not data:
|
|
return data
|
|
|
|
# 敏感字段列表
|
|
sensitive_fields = {
|
|
"password", "secret", "token", "api_key", "apikey",
|
|
"authorization", "auth", "credential", "key",
|
|
}
|
|
|
|
def _sanitize_value(key: str, value: Any) -> Any:
|
|
if isinstance(value, dict):
|
|
return {k: _sanitize_value(k, v) for k, v in value.items()}
|
|
elif isinstance(value, list):
|
|
return [_sanitize_value(key, v) for v in value]
|
|
elif key.lower() in sensitive_fields:
|
|
return "***REDACTED***"
|
|
elif isinstance(value, str) and len(value) > 1000:
|
|
return value[:1000] + "...[truncated]"
|
|
return value
|
|
|
|
return {k: _sanitize_value(k, v) for k, v in data.items()}
|
|
|