app/api/metrics.py

# app/api/metrics.py

import time
import psutil  # 导入 psutil 库用于获取系统信息
import logging
from typing import Dict, Any, List, Optional # 导入 Optional
from fastapi import APIRouter, Depends, HTTPException, status # 导入 APIRouter, Depends, HTTPException, status
from fastapi.responses import JSONResponse # 导入 JSONResponse
from pydantic import BaseModel, Field # 导入 BaseModel 和 Field

# 导入自定义模块
from app.api.auth import get_auth_token, get_admin_api_key # 从 auth 模块导入认证依赖，更改为 get_auth_token

logger = logging.getLogger(__name__)

# 定义 Pydantic 模型
class MetricData(BaseModel):
    """服务指标数据模型"""
    total_requests: int = Field(0, description="总请求数")
    successful_requests: int = Field(0, description="成功请求数")
    failed_requests: int = Field(0, description="失败请求数")
    active_keys_count: int = Field(0, description="活跃认证令牌数量")
    average_response_time: float = Field(0.0, description="平均响应时间（秒）")
    memory_usage: str = Field("N/A", description="当前内存使用情况")

class CallRecord(BaseModel):
    """单次调用记录数据模型"""
    timestamp: float = Field(..., description="请求时间戳（Unix时间）")
    key: str = Field(..., description="认证令牌（部分隐藏）")
    ip: str = Field(..., description="客户端 IP 地址")
    type: str = Field(..., description="请求类型（如 chat_completion）")
    model: str = Field("未知模型", description="使用的模型名称")
    status: str = Field(..., description="请求状态（success 或 failed）")
    response_time: float = Field(..., description="响应时间（秒）")
    request_summary: str = Field("N/A", description="请求内容摘要")
    response_summary: str = Field("N/A", description="响应内容摘要")

# 内存存储：服务指标和调用记录
# 注意：这些数据在应用重启时会丢失
# 使用 MetricData 模型的默认值初始化 service_metrics
service_metrics: MetricData = MetricData()
call_records: List[CallRecord] = []  # 存储调用记录
MAX_RECORDS = 1000  # 最大记录数

# 简单的内存存储，用于跟踪活跃 key (与 metrics 模块共享或在此处独立管理)
# 为了简化，这里先独立管理，后续可以考虑统一
active_tokens_status: Dict[str, float] = {} # 从 auth 模块移动到这里，以便统一管理活跃令牌

# 获取当前进程
process = psutil.Process()

# 创建 APIRouter 实例
router = APIRouter()

# 获取内存使用情况的辅助函数
def get_memory_usage():
    """获取当前进程的内存使用情况百分比"""
    try:
        # 获取进程的内存信息，单位为字节
        memory_info = process.memory_info()
        # resident set size (RSS) 是进程实际使用的物理内存量
        rss_mb = memory_info.rss / (1024 * 1024) # 转换为 MB

        # 打印原始内存数据，用于调试
        logger.debug(f"原始内存数据: RSS={rss_mb:.2f} MB")

        # 直接返回进程的 RSS (MB)，因为在某些环境中获取系统总内存可能不准确或固定
        return f"{rss_mb:.2f} MB"
    except Exception as e:
        logger.error(f"获取内存使用情况失败: {e}", exc_info=True) # 打印完整异常信息
        return "N/A" # 简化错误显示


def update_metrics_on_request(auth_token: str, client_host: str):
    """更新请求相关的服务指标"""
    service_metrics.total_requests += 1
    # 记录活跃的认证令牌及其最后活跃时间
    active_tokens_status[auth_token] = time.time()


def update_metrics_on_response(
    auth_token: str,
    status: str,
    response_time: float,
    client_host: str,
    request: Optional[Dict[str, Any]] = None, # 使用 Optional
    response: Optional[Dict[str, Any]] = None, # 使用 Optional
):
    """更新响应相关的服务指标和调用记录"""
    if status == "success":
        service_metrics.successful_requests += 1
        # 更新平均响应时间 (简单移动平均)
        # 避免除以零
        if service_metrics.total_requests > 0:
            service_metrics.average_response_time = (
                service_metrics.average_response_time
                * (service_metrics.total_requests - 1)
                + response_time
            ) / service_metrics.total_requests
        else:
            service_metrics.average_response_time = response_time
    else:
        service_metrics.failed_requests += 1

    # 获取模型名称
    model_name = request.get("model", "未知模型") if request else "未知模型"

    # 生成请求摘要
    request_summary = "N/A"
    if request and request.get("messages"):
        try:
            # 尝试获取最后一个消息的内容，并截取前 50 个字符
            last_message_content = request["messages"][-1]["content"]
            if isinstance(last_message_content, list): # 处理多模态输入
                text_parts = [p["text"] for p in last_message_content if p["type"] == "text"]
                request_summary = (text_parts[0][:50] + "...") if text_parts else "多模态请求"
            else:
                request_summary = last_message_content[:50]
        except Exception:
            request_summary = "请求解析错误"

    # 生成响应摘要
    response_summary = "N/A"
    if response and response.get("choices") and response["choices"][0].get("message"):
        try:
            response_summary = response["choices"][0]["message"]["content"][
                :50
            ]  # 截取前 50 个字符
        except Exception:
            response_summary = "响应解析错误"

    # 添加调用记录
    record = CallRecord(
        timestamp=time.time(),
        key=f"***{auth_token[-4:]}",  # 隐藏部分认证令牌
        ip=client_host,  # 从请求对象获取 IP 地址
        type="chat_completion",  # 或根据请求内容更具体
        model=model_name, # 添加模型名称
        status=status,
        response_time=response_time,
        request_summary=request_summary,  # 简要请求内容
        response_summary=response_summary,  # 简要响应内容
    )
    call_records.append(record)
    # 限制记录数
    if len(call_records) > MAX_RECORDS:
        call_records.pop(0)  # 移除最早的记录


def get_current_metrics() -> MetricData: # 返回类型改为 MetricData
    """获取当前服务指标数据"""
    current_time = time.time()
    # 计算活跃令牌数量
    service_metrics.active_keys_count = len(
        [
            token
            for token, last_active in active_tokens_status.items()
            if current_time - last_active < 300
        ]
    )  # 5分钟内活跃算活跃

    # 更新内存使用情况
    service_metrics.memory_usage = get_memory_usage()

    return service_metrics # 直接返回 MetricData 实例


def get_current_records() -> List[CallRecord]: # 返回类型改为 List[CallRecord]
    """获取最近的调用记录"""
    # 返回最近的记录，可以按时间排序
    return sorted(call_records, key=lambda x: x.timestamp, reverse=True) # 使用 x.timestamp

@router.get(
    "/metrics", dependencies=[Depends(get_auth_token)] # 更改为 get_auth_token
)  # 指标接口现在允许所有有效认证令牌访问
async def get_metrics_endpoint(auth_token: str = Depends(get_auth_token)): # 更改参数名为 auth_token
    """
    获取服务运行指标。
    """
    # 注意：auth_token 在这里主要用于权限校验，实际的指标收集发生在其他地方 (如 chat_completion)
    # 或者如果需要记录 metrics 接口本身的访问，可以在这里调用 update_metrics_on_request
    # logger.info(f"管理员令牌 '{auth_token}' 请求访问 /metrics 端点。") # 可选日志
    try:
        metrics_data = get_current_metrics()  # 从 metrics 模块获取指标
        return JSONResponse(content=metrics_data.model_dump()) # 将 Pydantic 模型转换为字典
    except Exception as e:
        logger.error(f"获取指标时发生错误: {e}", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"获取指标失败: {e}",
        )

@router.get(
    "/records", dependencies=[Depends(get_auth_token)] # 更改为 get_auth_token
)  # 记录接口现在允许所有有效认证令牌访问
async def get_records_endpoint(auth_token: str = Depends(get_auth_token)): # 更改参数名为 auth_token
    """
    获取服务调用记录。
    """
    try:
        records_data = get_current_records()  # 从 metrics 模块获取记录
        # 将 CallRecord 列表转换为字典列表
        return JSONResponse(content=[record.model_dump() for record in records_data])
    except Exception as e:
        logger.error(f"获取记录时发生错误: {e}", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"获取记录失败: {e}",
        )