┌──────────┐    ┌─────────────┐    ┌──────────┐
│ Agent A  │───▶│  AI Gateway │───▶│ LLM API  │
│ project:X │    │ inject:     │    │          │
│ task:123  │    │ metadata    │    │          │
└──────────┘    │ {project,   │    └──────────┘
                │  task,agent} │
                └─────────────┘

from langchain_community.callbacks import UsageMetadataCallbackHandler

# 按Agent实例绑定回调
agent_researcher = create_react_agent(
    llm=gpt4o,
    tools=search_tools,
    callbacks=[UsageMetadataCallbackHandler(
        agent_name="researcher",
        project="market_analysis"
    )]
)

# LangGraph 原生Token追踪
from langgraph.graph import StateGraph

graph = StateGraph(AgentState)
graph.add_node("researcher", researcher_node)
graph.add_node("writer", writer_node)

# 执行后自动获取每个节点的Token消耗
result = graph.invoke(initial_state)
token_report = result["__token_usage__"]
# {'researcher': {'input': 4200, 'output': 1800},
#  'writer': {'input': 2100, 'output': 3500}}

# LiteLLM 按Agent设置预算
from litellm import completion_with_budgets

budget_config = {
    "researcher_agent": {
        "max_tokens_per_hour": 500000,
        "max_cost_per_day": 50.0,
        "model": "gpt-4o-mini"  # 自动降级模型
    },
    "reviewer_agent": {
        "max_tokens_per_hour": 200000,
        "max_cost_per_day": 20.0,
        "model": "gpt-4o"
    }
}

用户请求
    │
    ▼
┌─────────────┐
│  路由层      │ ← 判断任务复杂度
└─────┬───────┘
      │
      ├── 简单任务（格式化、提取）→ GPT-4.1-nano / Claude Haiku
      │   成本：$0.10/M input
      │
      ├── 中等任务（总结、分类）→ GPT-4o-mini / Gemini Flash
      │   成本：$0.15/M input
      │
      └── 复杂任务（推理、代码）→ GPT-4o / Claude Sonnet
          成本：$2.50/M input

def route_model(task: AgentTask) -> str:
    """基于五个维度决定使用哪个模型"""
    score = 0

    # 维度1：工具复杂度
    if task.requires_code_execution: score += 3
    if task.requires_multi_step_reasoning: score += 2

    # 维度2：输出精度要求
    if task.error_tolerance < 0.01: score += 3

    # 维度3：上下文长度
    if task.context_tokens > 10000: score += 2

    # 维度4：延迟容忍度
    if task.max_latency_ms < 1000: score -= 2

    # 维度5：成本预算
    if task.remaining_budget < task.avg_cost * 2: score -= 3

    # 路由到对应层级
    if score >= 6: return "gpt-4o"         # 顶级
    elif score >= 3: return "gpt-4o-mini"   # 中级
    else: return "gpt-4.1-nano"             # 经济级

Level 0: 无压缩（基线）
  每请求Token：~25,000
  成本：$0.0625/请求

Level 1: 摘要压缩
  对话历史 → 递归摘要 → 注入摘要
  压缩比：3:1
  每请求Token：~10,000
  成本：$0.025/请求（↓60%）

Level 2: 检索增强记忆
  记忆文件 → 向量索引 → Top-K召回
  压缩比：8:1
  每请求Token：~4,000
  成本：$0.010/请求（↓84%）

Level 3: 多信号融合压缩
  语义相似度 + 时序衰减 + 任务相关性
  压缩比：12:1
  每请求Token：~2,500
  成本：$0.006/请求（↓90%）

# 反模式：传递原始输出
researcher_output = """
[5000 tokens of raw search results...]
"""
writer_agent.invoke({"research": researcher_output})

# 优化：传递结构化摘要
structured_summary = summarize_with_schema(
    researcher_output,
    schema={
        "key_findings": "list[str]",
        "data_points": "list[tuple[str, float]]",
        "recommendations": "list[str]"
    }
)  # 约500 tokens
writer_agent.invoke({"research": structured_summary})

# 使用Redis/共享存储替代消息传递
from langgraph.graph import StateGraph

class SharedBlackboard(TypedDict):
    findings: Annotated[list, add_messages]
    shared_context: str  # 压缩后的共享上下文

# 每个Agent只读取blackboard中的摘要
# 而非完整的历史对话

graph TB
    subgraph "应用层"
        A[用户请求] --> B[Orchestrator]
        B --> C[Agent A: 研究员]
        B --> D[Agent B: 分析师]
        B --> E[Agent C: 撰稿人]
    end

    subgraph "网关层"
        C --> F[LiteLLM Gateway]
        D --> F
        E --> F
        F --> G[模型路由决策]
        G --> H[缓存检查]
        H --> I[Token计量]
    end

    subgraph "模型层"
        I --> J[GPT-4o]
        I --> K[GPT-4o-mini]
        I --> L[Claude Haiku]
    end

    subgraph "可观测性层"
        I --> M[Helicone Dashboard]
        I --> N[OpenTelemetry]
        N --> O[Granafa]
        M --> P[成本告警]
        O --> Q[Token热力图]
    end

╔══════════════════════════════════════════════╗
║  Multi-Agent Token Cost Report - 2026-05-25 ║
╠══════════════════════════════════════════════╣
║  Task: market_analysis_v2                    ║
║  Duration: 47.3s                             ║
║  Total Tokens: 142,847                       ║
║  Total Cost: $0.427                          ║
╠══════════════════════════════════════════════╣
║  Agent Breakdown:                            ║
║  ├── researcher  78,200 tok  $0.234 (54.8%) ║
║  ├── analyst     41,650 tok  $0.125 (29.2%) ║
║  └── writer      22,997 tok  $0.068 (16.0%) ║
║                                              ║
║  Efficiency:                                 ║
║  ├── Input:     112,400 tok (78.7%)          ║
║  ├── Output:     30,447 tok (21.3%)          ║
║  ├── Wasted:     28,500 tok (20.0%) ← ⚠     ║
║  └── Cached:     31,200 tok (saved $0.094)   ║
║                                              ║
║  Budget Status:                              ║
║  ├── Daily:     $12.40 / $50.00 (24.8%)     ║
║  └── Hourly:    $2.10 / $10.00 (21.0%)      ║
╚══════════════════════════════════════════════╝

"""Token预算可控的多Agent系统"""
from litellm import completion
from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated
import hashlib

# === 1. Token预算管理器 ===
class TokenBudgetManager:
    def __init__(self, budgets: dict):
        self.budgets = budgets  # {"researcher": 500000, ...}
        self.usage = {k: 0 for k in budgets}

    def check(self, agent: str, estimated: int) -> str:
        """检查预算，返回推荐的模型"""
        remaining = self.budgets[agent] - self.usage[agent]
        if remaining < estimated * 0.2:
            return "gpt-4.1-nano"  # 紧急降级
        elif remaining < estimated * 0.5:
            return "gpt-4o-mini"   # 节约模式
        return None  # 使用默认模型

    def record(self, agent: str, tokens: int):
        self.usage[agent] += tokens

# === 2. 结构化摘要函数 ===
def compress_output(raw_output: str, schema: dict) -> str:
    """用便宜模型压缩输出"""
    prompt = f"Extract the following from the text:\n{schema}\n\nText:\n{raw_output}"
    resp = completion(
        model="gpt-4.1-nano",  # 用最便宜的模型压缩
        messages=[{"role": "user", "content": prompt}],
        max_tokens=500
    )
    return resp.choices[0].message.content

# === 3. 共享状态（黑板模式）===
class AgentState(TypedDict):
    task: str
    research_summary: str  # 压缩后的研究摘要
    analysis_summary: str  # 压缩后的分析摘要
    final_output: str
    token_log: Annotated[list, lambda x, y: x + y]

# === 4. Agent节点（带预算检查）===
budget_mgr = TokenBudgetManager({
    "researcher": 500000,
    "analyst": 300000,
    "writer": 200000
})

def researcher_node(state: AgentState) -> dict:
    model_override = budget_mgr.check("researcher", 5000)
    model = model_override or "gpt-4o-mini"

    resp = completion(
        model=model,
        messages=[{"role": "user", "content": state["task"]}],
        tools=search_tool_defs
    )

    # 压缩输出
    summary = compress_output(
        resp.choices[0].message.content,
        {"findings": "list", "sources": "list"}
    )
    budget_mgr.record("researcher", resp.usage.total_tokens)

    return {
        "research_summary": summary,
        "token_log": [{"agent": "researcher", "tokens": resp.usage.total_tokens}]
    }

# === 5. 构建图 ===
graph = StateGraph(AgentState)
graph.add_node("research", researcher_node)
graph.add_node("analyze", analyst_node)
graph.add_node("write", writer_node)
graph.add_edge("research", "analyze")
graph.add_edge("analyze", "write")
graph.add_edge("write", END)
graph.set_entry_point("research")

app = graph.compile()

graph LR
    subgraph "Token经济学六层栈"
        L1["第六层<br/>可观测性<br/>Helicone/Grafana"]
        L2["第五层<br/>预算引擎<br/>LiteLLM/Portkey"]
        L3["第四层<br/>模型路由<br/>三层分级策略"]
        L4["第三层<br/>上下文压缩<br/>摘要/检索/多信号"]
        L5["第二层<br/>成本归因<br/>网关/回调/原生"]
        L6["第一层<br/>Token计量<br/>Usage API"]
    end

    L1 --> L2 --> L3 --> L4 --> L5 --> L6

    style L1 fill:#e1f5fe
    style L2 fill:#b3e5fc
    style L3 fill:#81d4fa
    style L4 fill:#4fc3f7
    style L5 fill:#29b6f6
    style L6 fill:#03a9f4