lambdagent/demo/demo3_cost_predict.py

#!/usr/bin/env python3
"""
Demo 3: 端侧成本预测
=====================
输入鸿蒙 Agent 配置，输出端侧推理的 token / 延迟 / 电池消耗预估。

核心信息:
  - 端侧模型每次推理消耗算力和电量
  - 多 Agent 并行 → 手机发热 → 用户投诉
  - lambdagent 成本分级可在编译时预估算力消耗
"""

import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import yaml
from lambdagent.types import AgentType, LamType, TypeTag


# ── 端侧模型参数（盘古小模型估算）──
PANGU_LITE_PARAMS = {
    "tokens_per_second": 15,        # 端侧推理速度 (tok/s)
    "avg_input_tokens": 1500,       # 平均输入 token
    "avg_output_tokens": 400,       # 平均输出 token
    "power_per_inference_mah": 50,  # 每次推理电池消耗 (mAh)
    "battery_capacity_mah": 5000,   # 典型手机电池容量
}


def estimate_pipeline_cost(config_path: str):
    """从 YAML 配置预测端侧执行成本"""

    with open(config_path, "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f)

    agent_type = cfg.get("type", "simple")
    model_name = cfg.get("model", {}).get("name", "unknown")

    # 计算 Agent 数量和步数
    if agent_type == "chain":
        steps = cfg.get("chain", {}).get("steps", [])
        n_agents = len(steps)
        n_llm_calls = n_agents  # 顺序: 每个 Agent 一次 LLM 调用
        parallelism = 1
    elif agent_type == "parallel":
        agents = cfg.get("parallel", {}).get("agents", [])
        n_agents = len(agents)
        n_llm_calls = n_agents  # 并行: 同时调用
        parallelism = n_agents
    else:
        n_agents = 1
        n_llm_calls = 1
        parallelism = 1

    max_steps = cfg.get("runtime", {}).get("max_steps", n_llm_calls)

    # 成本估算
    p = PANGU_LITE_PARAMS
    total_input_tokens = p["avg_input_tokens"] * n_llm_calls
    total_output_tokens = p["avg_output_tokens"] * n_llm_calls
    total_tokens = total_input_tokens + total_output_tokens

    if parallelism > 1:
        # 并行: 延迟取最长的那个，但算力翻倍
        latency_seconds = (p["avg_input_tokens"] + p["avg_output_tokens"]) / p["tokens_per_second"]
        power_mah = p["power_per_inference_mah"] * n_llm_calls
    else:
        # 顺序: 延迟累加
        latency_seconds = total_tokens / p["tokens_per_second"]
        power_mah = p["power_per_inference_mah"] * n_llm_calls

    battery_percent = (power_mah / p["battery_capacity_mah"]) * 100

    return {
        "config": cfg.get("agentId", "unknown"),
        "model": model_name,
        "type": agent_type,
        "n_agents": n_agents,
        "n_llm_calls": n_llm_calls,
        "parallelism": parallelism,
        "total_input_tokens": total_input_tokens,
        "total_output_tokens": total_output_tokens,
        "latency_seconds": round(latency_seconds, 1),
        "power_mah": power_mah,
        "battery_percent": round(battery_percent, 1),
    }


def main():
    print("=" * 60)
    print("Demo 3: 端侧成本预测")
    print("场景: 鸿蒙 Agent 在手机上运行的算力/延迟/电池消耗")
    print("=" * 60)
    print()

    demo_dir = os.path.dirname(__file__)

    # ── 场景 A: 旅行规划 (顺序管道) ──
    travel_path = os.path.join(demo_dir, "harmony_travel.yaml")
    if os.path.exists(travel_path):
        print("[场景 A] 旅行规划 — 顺序管道 (Planner → Transport → Hotel)")
        cost = estimate_pipeline_cost(travel_path)
        print(f"  Agent:       {cost['config']}")
        print(f"  模型:        {cost['model']}")
        print(f"  架构:        {cost['type']} ({cost['n_agents']} Agents)")
        print(f"  LLM 调用次数: {cost['n_llm_calls']}")
        print(f"  预估 Token:  ~{cost['total_input_tokens']:,} (输入) + ~{cost['total_output_tokens']:,} (输出)")
        print(f"  端侧延迟:    ~{cost['latency_seconds']}s")
        print(f"  电池消耗:    ~{cost['battery_percent']}%")
        print()

    # ── 场景 B: IoT 智能家居 (并行) ──
    iot_path = os.path.join(demo_dir, "harmony_iot.yaml")
    if os.path.exists(iot_path):
        print("[场景 B] IoT 智能家居 — 并行 (Climate ∥ Light ∥ Security)")
        cost = estimate_pipeline_cost(iot_path)
        print(f"  Agent:       {cost['config']}")
        print(f"  模型:        {cost['model']}")
        print(f"  架构:        {cost['type']} ({cost['n_agents']} Agents, 并行度={cost['parallelism']})")
        print(f"  LLM 调用次数: {cost['n_llm_calls']} (同时)")
        print(f"  预估 Token:  ~{cost['total_input_tokens']:,} (输入) + ~{cost['total_output_tokens']:,} (输出)")
        print(f"  端侧延迟:    ~{cost['latency_seconds']}s (并行取最长)")
        print(f"  电池消耗:    ~{cost['battery_percent']}% (算力翻倍)")
        print()
        print("  ⚠️  3 个 Agent 并行 → CPU 满载 → 手机发热")
        print(f"  💡 建议: 电量 < 20% 时自动降级为单 Agent 模式")
    print()

    # ── 对比 ──
    print("[对比] 仓颉编译器 vs lambdagent")
    print("  仓颉:      编译通过 → 运行 → 手机发热 → 用户投诉 → 才发现")
    print("  lambdagent: 编译时预估成本 → 超阈值告警 → 自动降级策略")
    print()

    print("=" * 60)
    print("结论: 端侧算力消耗可预测 → 优化电池/发热体验")
    print("手机快没电时自动降级 → 用户体验不降级")
    print("=" * 60)


if __name__ == "__main__":
    main()