AI Agent 基础设施：可观测性、数据管理与运行时治理

from opentelemetry import trace, context
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
import time
import json

# 初始化 Agent 专用的 Tracer
trace.set_tracer_provider(TracerProvider())
tracer = trace.get_tracer("agent-runtime")
exporter = OTLPSpanExporter(endpoint="http://otel-collector:4317")
trace.get_tracer_provider().add_span_processor(BatchSpanProcessor(exporter))

class AgentTracer:
    """Agent 执行追踪器：记录完整的 Thought-Action-Observation 链路"""
    
    def __init__(self, agent_id: str, session_id: str):
        self.agent_id = agent_id
        self.session_id = session_id
        self.step_count = 0
    
    def trace_tool_call(self, tool_name: str, tool_input: dict, tool_output: dict):
        """追踪一次工具调用"""
        self.step_count += 1
        with tracer.start_as_current_span(
            f"agent.{tool_name}",
            attributes={
                "agent.id": self.agent_id,
                "agent.session_id": self.session_id,
                "agent.step": self.step_count,
                "agent.tool.name": tool_name,
                "agent.tool.input_size": len(json.dumps(tool_input)),
                "agent.tool.output_size": len(json.dumps(tool_output)),
                "agent.type": "tool_call",
            }
        ) as span:
            # 记录决策上下文
            span.set_attribute("agent.thought", "检索相关文档后决定调用搜索工具")
            span.set_attribute("agent.alternatives", 
                             json.dumps(["直接回答", "调用知识库", "请求用户澄清"]))
            return tool_output
    
    def trace_llm_call(self, prompt: str, response: str, model: str, 
                       latency_ms: float, tokens: dict):
        """追踪一次 LLM 调用"""
        with tracer.start_as_current_span(
            f"agent.llm.{model}",
            attributes={
                "agent.id": self.agent_id,
                "agent.llm.model": model,
                "agent.llm.latency_ms": latency_ms,
                "agent.llm.prompt_tokens": tokens.get("prompt", 0),
                "agent.llm.completion_tokens": tokens.get("completion", 0),
                "agent.llm.total_cost": tokens.get("cost", 0),
                "agent.type": "llm_call",
            }
        ) as span:
            return response

# 使用示例
tr = AgentTracer(agent_id="research-agent-01", session_id="session-abc123")
tool_output = tr.trace_tool_call(
    tool_name="web_search",
    tool_input={"query": "2026 AI Agent 基础设施趋势"},
    tool_output={"results": 15, "top_sources": ["arxiv", "github", "blog"]}
)
tr.trace_llm_call(
    prompt="根据搜索结果分析...",
    response="基础设施的核心趋势是...",
    model="gpt-4o",
    latency_ms=1230,
    tokens={"prompt": 4500, "completion": 800, "cost": 0.012}
)

import redis
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct, Filter, FieldCondition
import hashlib
import json
import time
from typing import Optional

class AgentDataManager:
    """Agent 数据管理器：三层数据架构（热/温/冷）"""
    
    def __init__(self, redis_url: str, qdrant_url: str, embedding_model):
        # 热数据层：Redis 缓存
        self.cache = redis.from_url(redis_url, decode_responses=True)
        # 温数据层：Qdrant 向量数据库
        self.vector_db = QdrantClient(url=qdrant_url)
        # 嵌入模型
        self.embedding_model = embedding_model
        self.default_ttl = 3600  # 缓存默认 TTL: 1 小时
    
    def store_context(self, session_id: str, context: dict) -> str:
        """存储会话上下文到热数据层"""
        key = f"agent:context:{session_id}"
        version = int(time.time())
        payload = {
            "data": context,
            "version": version,
            "timestamp": time.time(),
        }
        self.cache.setex(key, self.default_ttl, json.dumps(payload))
        return f"v{version}"
    
    def retrieve_context(self, session_id: str) -> Optional[dict]:
        """从热数据层检索会话上下文"""
        key = f"agent:context:{session_id}"
        raw = self.cache.get(key)
        if raw is None:
            return None
        payload = json.loads(raw)
        # 检查 TTL 是否接近过期
        age = time.time() - payload["timestamp"]
        if age > self.default_ttl * 0.8:
            print(f"Warning: context for {session_id} is {age:.0f}s old, "
                  f"consider refreshing")
        return payload["data"]
    
    def search_knowledge(self, query: str, collection: str = "knowledge",
                        top_k: int = 5, min_score: float = 0.7) -> list:
        """从向量数据库检索相关知识"""
        # 生成查询向量
        query_vector = self.embedding_model.encode(query).tolist()
        # 语义搜索
        results = self.vector_db.search(
            collection_name=collection,
            query_vector=query_vector,
            limit=top_k,
        )
        # 过滤低分结果
        relevant = [
            {
                "content": hit.payload.get("content", ""),
                "source": hit.payload.get("source", ""),
                "score": hit.score,
                "version": hit.payload.get("version", "unknown"),
            }
            for hit in results if hit.score >= min_score
        ]
        return relevant
    
    def cache_tool_response(self, tool_name: str, input_hash: str,
                           response: dict, ttl: int = 1800):
        """缓存工具响应（避免重复调用相同工具）"""
        key = f"agent:tool_cache:{tool_name}:{input_hash}"
        self.cache.setex(key, ttl, json.dumps({
            "response": response,
            "cached_at": time.time(),
        }))
    
    def get_cached_tool_response(self, tool_name: str, 
                                 input_hash: str) -> Optional[dict]:
        """获取缓存的工具响应"""
        key = f"agent:tool_cache:{tool_name}:{input_hash}"
        raw = self.cache.get(key)
        return json.loads(raw) if raw else None

# 使用示例
data_mgr = AgentDataManager(
    redis_url="redis://localhost:6379",
    qdrant_url="http://localhost:6334",
    embedding_model=None  # 实际使用 sentence-transformers
)

# 缓存会话上下文
version = data_mgr.store_context("session-001", {
    "user_goal": "分析 2026 年 AI Agent 基础设施趋势",
    "constraints": ["只用中文回复", "包含代码示例"],
})

# 检索知识库
results = data_mgr.search_knowledge(
    query="Agent 可观测性最佳实践",
    top_k=3,
    min_score=0.75,
)
print(f"找到 {len(results)} 条相关知识")

import json
import time
import math
from dataclasses import dataclass, field
from typing import Optional
from collections import deque

@dataclass
class MemoryItem:
    """单条记忆的标准化结构"""
    content: str
    source: str           # "conversation" | "tool_result" | "user_preference"
    timestamp: float
    importance: float     # 0.0 ~ 1.0，重要性评分
    access_count: int = 0
    last_accessed: float = 0.0
    tags: list = field(default_factory=list)
    
    def relevance_score(self, current_time: float, decay_rate: float = 0.0001) -> float:
        """计算当前时间点的相关性分数（含时间衰减）"""
        age = current_time - self.timestamp
        # 时间衰减 + 使用频率加成
        time_decay = math.exp(-decay_rate * age)
        usage_bonus = min(self.access_count * 0.05, 0.3)
        return self.importance * time_decay + usage_bonus


class AgentMemoryManager:
    """Agent 三层记忆管理器"""
    
    def __init__(self, context_window_limit: int = 128000,
                 short_term_tokens: int = 8000,
                 working_memory_limit: int = 50):
        self.context_limit = context_window_limit
        self.st_limit = short_term_tokens
        # 短期记忆：双端队列，FIFO 淘汰
        self.short_term = deque(maxlen=20)
        # 工作记忆：结构化存储
        self.working_memory: dict[str, MemoryItem] = {}
        self.wm_limit = working_memory_limit
        # 长期记忆索引（简化版，实际对接向量数据库）
        self.long_term_index: list[MemoryItem] = []
    
    def add_to_short_term(self, message: str, token_count: int):
        """添加消息到短期记忆（对话历史）"""
        self.short_term.append({
            "content": message,
            "tokens": token_count,
        })
    
    def compress_short_term(self) -> str:
        """压缩短期记忆：将最近对话摘要化"""
        if len(self.short_term) <= 3:
            # 对话不多时保留原文
            return "
".join(m["content"] for m in self.short_term)
        
        # 超过 3 条时：保留最后 2 条原文，中间部分摘要
        recent = list(self.short_term)[-2:]
        middle = list(self.short_term)[1:-2]
        summary = f"[{len(middle)} 轮对话摘要：讨论{middle[0]['content'][:50]}...]"
        return summary + "
" + "
".join(m["content"] for m in recent)
    
    def set_working_memory(self, key: str, content: str, importance: float = 0.8):
        """设置工作记忆条目"""
        if len(self.working_memory) >= self.wm_limit:
            # 淘汰最久未访问的条目
            oldest = min(self.working_memory.items(), 
                        key=lambda x: x[1].last_accessed)
            del self.working_memory[oldest[0]]
        
        self.working_memory[key] = MemoryItem(
            content=content,
            source="agent_state",
            timestamp=time.time(),
            importance=importance,
        )
    
    def retrieve_long_term(self, query: str, top_k: int = 3) -> list[str]:
        """从长期记忆检索（简化版——实际应使用向量数据库）"""
        current_time = time.time()
        # 按相关性排序
        scored = [
            (item, item.relevance_score(current_time))
            for item in self.long_term_index
            if query.lower() in item.content.lower() or 
               any(query.lower() in tag for tag in item.tags)
        ]
        scored.sort(key=lambda x: x[1], reverse=True)
        
        # 更新访问计数
        for item, _ in scored[:top_k]:
            item.access_count += 1
            item.last_accessed = current_time
        
        return [item.content for item, _ in scored[:top_k]]
    
    def get_context_prompt(self) -> str:
        """构建 LLM 的完整上下文 prompt"""
        parts = []
        
        # 工作记忆（结构化状态）
        if self.working_memory:
            wm_parts = []
            for key, item in self.working_memory.items():
                wm_parts.append(f"【{key}】{item.content}")
            parts.append("## 当前任务状态\n" + "\n".join(wm_parts))
        
        # 短期记忆（压缩后的对话历史）
        st = self.compress_short_term()
        if st:
            parts.append("## 对话历史\n" + st)
        
        return "\n\n".join(parts)

# 使用示例
mem = AgentMemoryManager()
mem.add_to_short_term("用户：帮我查一下 2026 年 AI Agent 的趋势", 20)
mem.add_to_short_term("Agent：好的，我来搜索相关信息", 15)
mem.add_to_short_term("Agent：根据搜索结果，主要趋势有...", 150)

mem.set_working_memory("current_task", "研究 2026 AI Agent 基础设施", 0.9)
mem.set_working_memory("step", "已完成信息检索，正在整理结果", 0.7)

print(mem.get_context_prompt())

import json
import time
import enum
from dataclasses import dataclass, asdict
from typing import Optional, Any
import sqlite3

class TaskStatus(enum.Enum):
    WAITING = "waiting"
    PLANNING = "planning"
    EXECUTING = "executing"
    VERIFYING = "verifying"
    COMPLETED = "completed"
    FAILED = "failed"
    RECOVERING = "recovering"

@dataclass
class Checkpoint:
    """Agent 执行检查点"""
    checkpoint_id: str
    timestamp: float
    task_status: str
    step_number: int
    tool_history: list[dict]
    working_memory: dict
    error: Optional[str] = None
    
    def save_to_db(self, db_path: str):
        """持久化检查点到 SQLite"""
        conn = sqlite3.connect(db_path)
        conn.execute("""
            INSERT OR REPLACE INTO checkpoints 
            (checkpoint_id, timestamp, task_status, step_number, 
             tool_history, working_memory, error)
            VALUES (?, ?, ?, ?, ?, ?, ?)
        """, (
            self.checkpoint_id, self.timestamp, self.task_status,
            self.step_number,
            json.dumps(self.tool_history),
            json.dumps(self.working_memory),
            self.error,
        ))
        conn.commit()
        conn.close()
    
    @classmethod
    def load_latest(cls, db_path: str) -> Optional['Checkpoint']:
        """加载最近的检查点用于恢复"""
        conn = sqlite3.connect(db_path)
        row = conn.execute(
            "SELECT * FROM checkpoints ORDER BY timestamp DESC LIMIT 1"
        ).fetchone()
        conn.close()
        if row is None:
            return None
        return cls(
            checkpoint_id=row[0], timestamp=row[1],
            task_status=row[2], step_number=row[3],
            tool_history=json.loads(row[4]),
            working_memory=json.loads(row[5]),
            error=row[6],
        )


class AgentStateManager:
    """Agent 状态管理器：支持检查点和异常恢复"""
    
    def __init__(self, agent_id: str, db_path: str = "agent_state.db"):
        self.agent_id = agent_id
        self.db_path = db_path
        self.current_status = TaskStatus.WAITING
        self.step_number = 0
        self.tool_history: list[dict] = []
        self.working_memory: dict = {}
        self._init_db()
    
    def _init_db(self):
        conn = sqlite3.connect(self.db_path)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS checkpoints (
                checkpoint_id TEXT PRIMARY KEY,
                timestamp REAL,
                task_status TEXT,
                step_number INTEGER,
                tool_history TEXT,
                working_memory TEXT,
                error TEXT
            )
        """)
        conn.commit()
        conn.close()
    
    def create_checkpoint(self, error: Optional[str] = None) -> Checkpoint:
        """创建检查点"""
        cp = Checkpoint(
            checkpoint_id=f"{self.agent_id}-{self.step_number}-{int(time.time())}",
            timestamp=time.time(),
            task_status=self.current_status.value,
            step_number=self.step_number,
            tool_history=self.tool_history[-20:],  # 保留最近 20 条
            working_memory=dict(self.working_memory),
            error=error,
        )
        cp.save_to_db(self.db_path)
        return cp
    
    def recover(self) -> bool:
        """从最近检查点恢复"""
        latest = Checkpoint.load_latest(self.db_path)
        if latest is None:
            return False
        
        self.current_status = TaskStatus.RECOVERING
        self.step_number = latest.step_number
        self.tool_history = latest.tool_history
        self.working_memory = latest.working_memory
        
        print(f"Recovered from checkpoint {latest.checkpoint_id} "
              f"(step {latest.step_number}, status: {latest.task_status})")
        if latest.error:
            print(f"Last error: {latest.error}")
        
        self.current_status = TaskStatus.EXECUTING
        return True

# 使用示例
sm = AgentStateManager("agent-041-demo")
sm.current_status = TaskStatus.EXECUTING
sm.step_number = 3
sm.tool_history.append({
    "tool": "web_search",
    "input": "AI Agent 基础设施 2026",
    "output": "...",
    "status": "success",
})
sm.create_checkpoint()
print("Checkpoint created")

import time
import re
import enum
from dataclasses import dataclass
from typing import Optional

class RiskLevel(enum.Enum):
    LOW = "low"           # 只读操作
    MEDIUM = "medium"     # 写入但可恢复
    HIGH = "high"         # 破坏性操作
    CRITICAL = "critical" # 不可逆操作

@dataclass
class ToolPermission:
    """工具权限定义"""
    tool_name: str
    allowed: bool
    risk_level: RiskLevel
    requires_approval: bool
    rate_limit_per_min: int = 60
    data_scope: str = "all"  # "all" | "own" | "none"

@dataclass
class AuditLog:
    """行为审计日志条目"""
    timestamp: float
    agent_id: str
    session_id: str
    action: str
    tool_name: str
    input_params: dict
    output: str
    risk_level: str
    approved_by: Optional[str] = None
    blocked: bool = False
    block_reason: str = ""


class AgentGovernanceEngine:
    """Agent 治理引擎：权限控制 + 审计 + 输出验证"""
    
    def __init__(self, agent_id: str, session_id: str):
        self.agent_id = agent_id
        self.session_id = session_id
        self.permissions: dict[str, ToolPermission] = {}
        self.audit_log: list[AuditLog] = []
        self.token_budget = 100000
        self.tokens_used = 0
        self.cost_budget = 10.0  # 美元
        self.cost_used = 0.0
    
    def set_permission(self, perm: ToolPermission):
        """设置工具权限"""
        self.permissions[perm.tool_name] = perm
    
    def check_tool_permission(self, tool_name: str) -> tuple[bool, str]:
        """检查 Agent 是否有权限调用某个工具"""
        if tool_name not in self.permissions:
            return False, f"工具 {tool_name} 未授权"
        
        perm = self.permissions[tool_name]
        if not perm.allowed:
            return False, f"工具 {tool_name} 被禁用"
        
        if perm.requires_approval:
            return False, f"工具 {tool_name} 需要人工审批"
        
        return True, "已授权"
    
    def audit_tool_call(self, tool_name: str, input_params: dict, 
                       output: str, risk: RiskLevel, 
                       blocked: bool = False, block_reason: str = ""):
        """记录工具调用的审计日志"""
        log = AuditLog(
            timestamp=time.time(),
            agent_id=self.agent_id,
            session_id=self.session_id,
            action="tool_call",
            tool_name=tool_name,
            input_params=input_params,
            output=output[:500],  # 截断避免日志过大
            risk_level=risk.value,
            blocked=blocked,
            block_reason=block_reason,
        )
        self.audit_log.append(log)
    
    def check_budget(self, tokens: int, cost: float) -> tuple[bool, str]:
        """检查 token 和成本预算"""
        if self.tokens_used + tokens > self.token_budget:
            return False, f"Token 预算不足: 已用 {self.tokens_used}/{self.token_budget}"
        if self.cost_used + cost > self.cost_budget:
            return False, f"成本预算不足: 已用 ${self.cost_used:.2f}/${self.cost_budget:.2f}"
        return True, "预算充足"
    
    def validate_output(self, output: str) -> tuple[bool, list[str]]:
        """验证输出内容的合规性"""
        issues = []
        # 敏感信息扫描
        patterns = {
            "api_key": r"sk-[a-zA-Z0-9]{20,}",
            "password": r"password['\"]?\s*[:=]\s*['\"][^'\"]+['\"]",
            "email": r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}",
        }
        for name, pattern in patterns.items():
            matches = re.findall(pattern, output)
            if matches:
                issues.append(f"检测到疑似{name}泄露: {len(matches)} 处")
        # 长度检查
        if len(output) > 50000:
            issues.append(f"输出过长: {len(output)} 字符，建议截断")
        return len(issues) == 0, issues

# 使用示例
governance = AgentGovernanceEngine("agent-041", "session-xyz")
governance.set_permission(ToolPermission(
    tool_name="web_search", allowed=True,
    risk_level=RiskLevel.LOW, requires_approval=False,
    rate_limit_per_min=30,
))
governance.set_permission(ToolPermission(
    tool_name="delete_record", allowed=True,
    risk_level=RiskLevel.CRITICAL, requires_approval=True,
    rate_limit_per_min=5,
))

allowed, reason = governance.check_tool_permission("web_search")
print(f"web_search: {allowed} - {reason}")

allowed, reason = governance.check_tool_permission("delete_record")
print(f"delete_record: {allowed} - {reason}")

valid, issues = governance.validate_output(
    "这是正常输出内容，没有任何敏感信息。"
)
print(f"输出合规: {valid}")
if issues:
    for issue in issues:
        print(f"  ⚠ {issue}")

import asyncio
import time
from dataclasses import dataclass
from typing import Optional
from enum import Enum

class AgentRole(Enum):
    RESEARCHER = "researcher"
    WRITER = "writer"
    REVIEWER = "reviewer"

@dataclass
class TaskMessage:
    """Agent 间的任务消息"""
    task_id: str
    from_role: AgentRole
    to_role: AgentRole
    payload: dict
    timestamp: float = 0.0
    retry_count: int = 0
    max_retries: int = 3
    
    def __post_init__(self):
        if self.timestamp == 0:
            self.timestamp = time.time()


class MessageQueue:
    """简单的 Agent 间消息队列（生产环境应使用 RabbitMQ / Kafka）"""
    
    def __init__(self):
        self.queues: dict[AgentRole, list[TaskMessage]] = {
            role: [] for role in AgentRole
        }
        self.processed: list[TaskMessage] = []
        self.dead_letter: list[TaskMessage] = []
    
    async def publish(self, message: TaskMessage):
        """发布消息到目标 Agent 的队列"""
        self.queues[message.to_role].append(message)
    
    async def consume(self, role: AgentRole, 
                     timeout: float = 5.0) -> Optional[TaskMessage]:
        """从队列消费消息"""
        start = time.time()
        while time.time() - start < timeout:
            if self.queues[role]:
                return self.queues[role].pop(0)
            await asyncio.sleep(0.1)
        return None
    
    async def nack(self, message: TaskMessage):
        """消息处理失败，重新入队或进入死信队列"""
        message.retry_count += 1
        if message.retry_count >= message.max_retries:
            self.dead_letter.append(message)
            print(f"Message {message.task_id} moved to dead letter queue "
                  f"after {message.retry_count} retries")
        else:
            self.queues[message.to_role].insert(0, message)
            print(f"Message {message.task_id} requeued "
                  f"(retry {message.retry_count}/{message.max_retries})")
    
    async def ack(self, message: TaskMessage):
        """消息处理成功，标记为已处理"""
        self.processed.append(message)


class AgentOrchestrator:
    """Agent 编排器：协调多 Agent 协作流程"""
    
    def __init__(self):
        self.message_queue = MessageQueue()
        self.task_results: dict[str, dict] = {}
    
    async def run_pipeline(self, user_query: str) -> dict:
        """执行完整的多 Agent 协作流程"""
        task_id = f"task-{int(time.time())}"
        
        # 阶段 1: Researcher 收集信息
        research_msg = TaskMessage(
            task_id=task_id,
            from_role=AgentRole.RESEARCHER,
            to_role=AgentRole.RESEARCHER,
            payload={"query": user_query, "phase": "research"},
        )
        await self.message_queue.publish(research_msg)
        
        # 模拟 Researcher 工作
        research_result = {
            "findings": ["发现 15 篇相关文献", "3 个关键技术趋势"],
            "sources": ["arxiv", "github", "tech blog"],
        }
        self.task_results[task_id] = {"research": research_result}
        
        # 阶段 2: Writer 撰写报告
        writer_msg = TaskMessage(
            task_id=task_id,
            from_role=AgentRole.RESEARCHER,
            to_role=AgentRole.WRITER,
            payload={
                "query": user_query,
                "research_findings": research_result,
                "phase": "writing",
            },
        )
        await self.message_queue.publish(writer_msg)
        
        # 模拟 Writer 工作
        draft = "## 2026 AI Agent 基础设施趋势分析\n\n根据研究..."
        self.task_results[task_id]["draft"] = draft
        
        # 阶段 3: Reviewer 质量审核
        review_msg = TaskMessage(
            task_id=task_id,
            from_role=AgentRole.WRITER,
            to_role=AgentRole.REVIEWER,
            payload={"draft": draft, "phase": "review"},
        )
        await self.message_queue.publish(review_msg)
        
        # 模拟 Reviewer 工作
        review_result = {
            "approved": True,
            "suggestions": ["补充代码示例", "增加对比表格"],
        }
        self.task_results[task_id]["review"] = review_result
        
        return self.task_results[task_id]

# 使用示例
async def main():
    orchestrator = AgentOrchestrator()
    result = await orchestrator.run_pipeline(
        "分析 2026 年 AI Agent 基础设施的技术趋势"
    )
    print(f"Task completed: {list(result.keys())}")

asyncio.run(main())