合规审计（Compliance Audit）

# 合规审计核心机制演示
# 基于 langfuse==2.x（可选），无 Langfuse 时自动降级为本地日志
import uuid
import json
from datetime import datetime, timezone
from enum import Enum
from typing import Dict, List


# ========== 第一部分：操作日志与追踪 ==========

class AuditLogger:
    """可审计的操作日志记录器"""

    def __init__(self):
        self.logs: List[Dict] = []  # 生产环境应替换为 append-only 存储
        self.trace_id: str = ""

    def start_trace(self, user_id: str, request_content: str) -> str:
        """为一次用户请求创建追踪会话"""
        self.trace_id = str(uuid.uuid4())
        self._append_log(
            operation="user_request",
            input_data={"user_id": user_id, "content": request_content},
            output_data={},
            status="initiated"
        )
        return self.trace_id

    def log_llm_call(self, prompt: str, model: str, response: str,
                     tokens_used: int, duration_ms: int):
        """记录一次 LLM 调用"""
        self._append_log(
            operation="llm_call",
            input_data={"prompt": prompt, "model": model},
            output_data={"response": response, "tokens": tokens_used},
            status="success",
            duration_ms=duration_ms
        )

    def log_tool_call(self, tool_name: str, params: dict,
                      result: dict, status: str, duration_ms: int):
        """记录一次工具调用"""
        self._append_log(
            operation="tool_call",
            input_data={"tool": tool_name, "params": params},
            output_data=result,
            status=status,
            duration_ms=duration_ms
        )

    def log_final_response(self, response: str):
        """记录最终输出"""
        self._append_log(
            operation="final_response",
            input_data={},
            output_data={"response": response},
            status="success"
        )

    def _append_log(self, operation: str, input_data: dict,
                    output_data: dict, status: str, duration_ms: int = 0):
        """写入一条不可修改的日志（生产环境应写入 append-only 存储）"""
        entry = {
            "log_id": str(uuid.uuid4()),
            "trace_id": self.trace_id,
            "operation": operation,
            "input": input_data,
            "output": output_data,
            "status": status,
            "duration_ms": duration_ms,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }
        self.logs.append(entry)

    def get_trace_logs(self) -> List[Dict]:
        """获取当前 Trace 的所有日志"""
        return [log for log in self.logs if log["trace_id"] == self.trace_id]


# ========== 第二部分：合规检查引擎 ==========

class CheckResult(Enum):
    PASS = "pass"
    WARNING = "warning"
    FAIL = "fail"

class ComplianceChecker:
    """合规检查引擎：对照法规标准逐项验证"""

    def check_audit_completeness(self, logs: List[Dict]) -> Dict:
        """检查审计日志完整性——所有关键操作是否都被记录"""
        required_ops = {"user_request", "llm_call", "tool_call", "final_response"}
        recorded_ops = {log["operation"] for log in logs}
        missing = required_ops - recorded_ops
        if missing:
            return {"check": "audit_completeness", "result": CheckResult.FAIL.value,
                    "detail": f"缺少操作日志: {missing}"}
        return {"check": "audit_completeness", "result": CheckResult.PASS.value,
                "detail": "所有关键操作均已记录"}

    def check_data_minimization(self, accessed: List[str],
                                required: List[str]) -> Dict:
        """检查数据最小化——是否只访问了必要字段"""
        unnecessary = set(accessed) - set(required)
        if unnecessary:
            return {"check": "data_minimization", "result": CheckResult.WARNING.value,
                    "detail": f"访问了不必要的字段: {unnecessary}"}
        return {"check": "data_minimization", "result": CheckResult.PASS.value,
                "detail": "符合数据最小化原则"}

    def check_access_control(self, accessed_resources: List[str],
                             user_permissions: List[str]) -> Dict:
        """检查访问控制——操作是否在授权范围内"""
        unauthorized = set(accessed_resources) - set(user_permissions)
        if unauthorized:
            return {"check": "access_control", "result": CheckResult.FAIL.value,
                    "detail": f"越权访问: {unauthorized}"}
        return {"check": "access_control", "result": CheckResult.PASS.value,
                "detail": "访问权限检查通过"}

    def run_audit(self, logs: List[Dict], context: Dict) -> Dict:
        """执行完整合规审计，返回审计报告"""
        checks = [
            self.check_audit_completeness(logs),
            self.check_data_minimization(
                context.get("accessed_fields", []),
                context.get("required_fields", [])
            ),
            self.check_access_control(
                context.get("accessed_resources", []),
                context.get("user_permissions", [])
            ),
        ]
        results = [c["result"] for c in checks]
        if CheckResult.FAIL.value in results:
            overall = "FAILED"
        elif CheckResult.WARNING.value in results:
            overall = "WARNING"
        else:
            overall = "PASSED"

        return {
            "audit_id": str(uuid.uuid4()),
            "trace_id": logs[0]["trace_id"] if logs else "",
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "checks": checks,
            "overall": overall,
        }


# ========== 使用示例 ==========

# 1. 记录操作
logger = AuditLogger()
trace_id = logger.start_trace("user-001", "查询我的账户余额")
logger.log_llm_call(
    prompt="用户想查余额，判断是否需要调用数据库",
    model="gpt-4o", response="需要调用 query_balance 工具",
    tokens_used=120, duration_ms=350
)
logger.log_tool_call(
    tool_name="query_balance", params={"user_id": "user-001"},
    result={"balance": 5000.0, "currency": "CNY"},
    status="success", duration_ms=80
)
logger.log_final_response("您的账户余额为 5000.00 元")

# 2. 执行合规检查
checker = ComplianceChecker()
report = checker.run_audit(
    logs=logger.get_trace_logs(),
    context={
        "accessed_fields": ["user_id", "balance"],
        "required_fields": ["user_id", "balance"],
        "accessed_resources": ["read_balance"],
        "user_permissions": ["read_balance"],
    }
)
print(json.dumps(report, indent=2, ensure_ascii=False))
# overall: "PASSED" —— 所有检查项通过