实时 Agent 架构设计：低延迟智能体

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import asyncio
import json

app = FastAPI()

class ChatRequest(BaseModel):
    session_id: str
    message: str

async def stream_agent_response(session_id: str, message: str):
    """流式生成 Agent 响应"""
    # 1. 发送会话元信息
    yield f"data: {json.dumps({'type': 'meta', 'session_id': session_id})}\n\n"
    
    # 2. 发送思考状态
    yield f"data: {json.dumps({'type': 'status', 'status': 'thinking'})}\n\n"
    await asyncio.sleep(0.1)
    
    # 3. 检索知识（如果有 RAG）
    yield f"data: {json.dumps({'type': 'status', 'status': 'retrieving'})}\n\n"
    context = await retrieve_knowledge(message)
    await asyncio.sleep(0.05)
    
    # 4. 流式生成回复
    yield f"data: {json.dumps({'type': 'status', 'status': 'generating'})}\n\n"
    async for token in llm_stream_generate(message, context):
        yield f"data: {json.dumps({'type': 'token', 'content': token})}\n\n"
    
    # 5. 发送完成信号
    yield f"data: {json.dumps({'type': 'done'})}\n\n"


async def llm_stream_generate(prompt: str, context: str = ""):
    """调用 LLM 流式 API"""
    # 模拟流式 token 生成
    full_response = f"基于您的问题「{prompt}」，这是流式生成的回答。"
    for char in full_response:
        await asyncio.sleep(0.03)  # 模拟推理延迟
        yield char


async def retrieve_knowledge(query: str) -> str:
    """检索知识库"""
    await asyncio.sleep(0.1)
    return f"相关知识: {query}"


@app.post("/chat/stream")
async def chat_stream(req: ChatRequest):
    return StreamingResponse(
        stream_agent_response(req.session_id, req.message),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",  # 禁用 Nginx 缓冲
        }
    )

from fastapi import WebSocket, WebSocketDisconnect
from typing import Dict
import asyncio

class ConnectionManager:
    """WebSocket 连接管理器"""
    
    def __init__(self):
        self.active: Dict[str, WebSocket] = {}
    
    async def connect(self, session_id: str, ws: WebSocket):
        await ws.accept()
        self.active[session_id] = ws
    
    def disconnect(self, session_id: str):
        self.active.pop(session_id, None)
    
    async def send_token(self, session_id: str, token: str):
        ws = self.active.get(session_id)
        if ws:
            await ws.send_json({"type": "token", "content": token})

manager = ConnectionManager()

@app.websocket("/ws/chat/{session_id}")
async def ws_chat(ws: WebSocket, session_id: str):
    await manager.connect(session_id, ws)
    try:
        while True:
            data = await ws.receive_json()
            
            if data.get("type") == "message":
                message = data["content"]
                
                # 流式输出
                async for token in llm_stream_generate(message):
                    await manager.send_token(session_id, token)
                
                await ws.send_json({"type": "done"})
            
            elif data.get("type") == "interrupt":
                # 用户中断生成
                await ws.send_json({"type": "interrupted"})
    
    except WebSocketDisconnect:
        manager.disconnect(session_id)

import httpx
from typing import AsyncIterator

class StreamingLLMClient:
    """流式 LLM 推理客户端"""
    
    def __init__(self, api_base: str, api_key: str):
        self.api_base = api_base
        self.headers = {"Authorization": f"Bearer {api_key}"}
    
    async def stream_chat(
        self,
        messages: list[dict],
        model: str = "gpt-4",
        temperature: float = 0.7,
        max_tokens: int = 2048
    ) -> AsyncIterator[str]:
        """流式调用 LLM Chat API"""
        async with httpx.AsyncClient() as client:
            async with client.stream(
                "POST",
                f"{self.api_base}/v1/chat/completions",
                json={
                    "model": model,
                    "messages": messages,
                    "temperature": temperature,
                    "max_tokens": max_tokens,
                    "stream": True  # 关键：开启流式
                },
                headers=self.headers,
                timeout=60.0
            ) as response:
                async for line in response.aiter_lines():
                    if line.startswith("data: "):
                        data = line[6:]
                        if data == "[DONE]":
                            break
                        
                        chunk = json.loads(data)
                        delta = chunk["choices"][0].get("delta", {})
                        if "content" in delta:
                            yield delta["content"]


# 带首字延迟统计的包装器
class LatencyAwareStream:
    """追踪首字延迟和吞吐量的流式包装器"""
    
    def __init__(self, stream: AsyncIterator[str]):
        self.stream = stream
        self.first_token_time = None
        self.token_count = 0
        self.total_time = 0
        self._start = asyncio.get_event_loop().time()
    
    async def __aiter__(self):
        async for token in self.stream:
            if self.first_token_time is None:
                self.first_token_time = asyncio.get_event_loop().time() - self._start
            self.token_count += 1
            yield token
        self.total_time = asyncio.get_event_loop().time() - self._start
    
    def metrics(self) -> dict:
        return {
            "first_token_latency_ms": round(self.first_token_time * 1000, 1) if self.first_token_time else None,
            "total_time_ms": round(self.total_time * 1000, 1),
            "token_count": self.token_count,
            "tokens_per_second": round(self.token_count / self.total_time, 1) if self.total_time > 0 else 0,
        }

import asyncio
from typing import Any, AsyncIterator

class PipelineStage:
    """管道阶段"""
    
    def __init__(self, name: str, process_fn, min_delay: float = 0):
        self.name = name
        self.process_fn = process_fn
        self.min_delay = min_delay
    
    async def process(self, input_data: Any) -> Any:
        if self.min_delay:
            await asyncio.sleep(self.min_delay)
        return await self.process_fn(input_data)

class AgentPipeline:
    """Agent 处理管道，支持阶段并行重叠"""
    
    def __init__(self, stages: list[PipelineStage]):
        self.stages = stages
    
    async def execute(self, query: str) -> AsyncIterator[dict]:
        """管道执行，每个阶段完成后立即传递给下一阶段"""
        
        # 阶段1: 意图理解（快速完成）
        intent = await self.stages[0].process(query)
        yield {"stage": "intent", "result": intent}
        
        # 阶段2&3 并行: 知识检索 + 开始推理（不等检索完成就开始推理）
        retrieval_task = asyncio.create_task(self.stages[1].process(query))
        
        # 先用无增强上下文开始流式推理
        async for token in self._stream_with_early_start(query, retrieval_task):
            yield {"stage": "token", "result": token}
        
        # 如果检索结果在推理后完成，可作为"补充信息"
        retrieval_result = await retrieval_task
        if retrieval_result:
            yield {"stage": "retrieval", "result": retrieval_result}
        
        yield {"stage": "done"}
    
    async def _stream_with_early_start(self, query: str, retrieval_task):
        """在检索进行的同时开始推理"""
        # 策略：等待检索一小段时间，超时则直接开始推理
        try:
            context = await asyncio.wait_for(
                asyncio.shield(retrieval_task), 
                timeout=0.3  # 最多等 300ms
            )
        except asyncio.TimeoutError:
            context = None  # 检索未完成，先开始推理
        
        # 流式生成
        messages = [{"role": "user", "content": query}]
        if context:
            messages.insert(0, {"role": "system", "content": f"参考信息: {context}"})
        
        async for token in mock_stream_generate(query):
            yield token


async def mock_stream_generate(query: str):
    """模拟流式生成"""
    response = f"针对「{query}」的分析回答。"
    for char in response:
        await asyncio.sleep(0.02)
        yield char


# 构建管道
pipeline = AgentPipeline([
    PipelineStage("intent", lambda q: asyncio.sleep(0.05, result={"intent": "question"})),
    PipelineStage("retrieval", lambda q: asyncio.sleep(0.2, result="relevant docs")),
    PipelineStage("generation", lambda q: asyncio.sleep(0.3, result="answer")),
])

class SpeculativeDecoder:
    """推测解码：小模型草稿 + 大模型验证"""
    
    def __init__(self, draft_model, target_model):
        self.draft = draft_model  # 小模型（快）
        self.target = target_model  # 大模型（准）
    
    async def generate(self, prompt: str, max_tokens: int = 200) -> AsyncIterator[str]:
        tokens = []
        while len(tokens) < max_tokens:
            # 1. 小模型生成 K 个草稿 token
            draft_tokens = await self.draft.generate(prompt, n=4)
            
            # 2. 大模型批量验证
            verified = await self.target.verify(prompt, draft_tokens)
            
            # 3. 输出通过的 token
            for token in verified.accepted:
                yield token
                tokens.append(token)
            
            # 4. 如果全部通过，继续；否则从拒绝点重新开始
            if not verified.all_accepted:
                yield verified.corrected_token
                tokens.append(verified.corrected_token)

class KVCacheManager:
    """KV Cache 管理：避免重复计算历史 token 的注意力"""
    
    def __init__(self, max_cache_entries: int = 100):
        self.cache: dict[str, dict] = {}
        self.max_entries = max_cache_entries
        self.lru: list[str] = []
    
    def get(self, session_id: str) -> dict | None:
        """获取会话的 KV Cache"""
        if session_id in self.cache:
            self.lru.remove(session_id)
            self.lru.insert(0, session_id)
            return self.cache[session_id]
        return None
    
    def set(self, session_id: str, kv_cache: dict):
        """保存 KV Cache"""
        if session_id in self.cache:
            self.lru.remove(session_id)
        
        self.cache[session_id] = kv_cache
        self.lru.insert(0, session_id)
        
        # LRU 淘汰
        while len(self.lru) > self.max_entries:
            evicted = self.lru.pop()
            del self.cache[evicted]
    
    def append_token(self, session_id: str, token_kv: dict):
        """增量更新 KV Cache（新 token）"""
        if session_id in self.cache:
            self.cache[session_id]["layers"].append(token_kv)

import hashlib
from typing import Optional

class ResponseCache:
    """语义感知的响应缓存"""
    
    def __init__(self, ttl: int = 3600, max_entries: int = 10000):
        self.cache: dict[str, dict] = {}
        self.ttl = ttl
        self.max_entries = max_entries
    
    def _key(self, query: str, session_context: str = "") -> str:
        """生成缓存键"""
        normalized = query.strip().lower()
        return hashlib.sha256(f"{normalized}:{session_context}".encode()).hexdigest()
    
    def get(self, query: str, context: str = "") -> Optional[dict]:
        import time
        key = self._key(query, context)
        entry = self.cache.get(key)
        if entry and time.time() - entry["timestamp"] < self.ttl:
            return entry["response"]
        return None
    
    def set(self, query: str, response: dict, context: str = ""):
        import time
        key = self._key(query, context)
        self.cache[key] = {
            "response": response,
            "timestamp": time.time()
        }
        # 简单容量控制
        if len(self.cache) > self.max_entries:
            oldest = min(self.cache.items(), key=lambda x: x[1]["timestamp"])
            del self.cache[oldest[0]]

class EdgeRouter:
    """边缘路由器：选择最近的推理节点"""
    
    def __init__(self):
        self.regions = {
            "cn-east": {"endpoint": "https://shanghai.inference.ai", "latency_base": 20},
            "cn-south": {"endpoint": "https://guangzhou.inference.ai", "latency_base": 25},
            "cn-north": {"endpoint": "https://beijing.inference.ai", "latency_base": 22},
            "cn-west": {"endpoint": "https://chengdu.inference.ai", "latency_base": 30},
        }
    
    def select_region(self, client_ip: str = None, client_region: str = None) -> str:
        """选择最佳推理区域"""
        if client_region and client_region in self.regions:
            return client_region
        
        # 基于 IP 地理定位
        detected = self._geo_lookup(client_ip) if client_ip else "cn-east"
        return detected if detected in self.regions else "cn-east"
    
    def _geo_lookup(self, ip: str) -> str:
        """IP 地理定位（简化版）"""
        # 实际中使用 IP 地理位置数据库
        return "cn-east"
    
    def get_endpoint(self, region: str) -> str:
        return self.regions[region]["endpoint"]

用户 (北京) → CDN 边缘节点 (北京) → 推理服务 (北京) → 中心知识库 (上海)
     ↑                                                          |
     └────────── 流式响应回流 ←──────────────────────────────────┘