Browse Source

feat: Paper III §4.3 graded cost prediction — static cost estimation (P1-2)

Implements graded type system for cost prediction:
- CostGrade = (probability, tokens, latency, money)
- grade_serial(): p1*p2, t1+t2, l1+l2, m1+m2
- grade_parallel(): p1*p2, t1+t2, max(l1,l2), m1+m2
- grade_iterate(): p^n, n*t, n*l, n*m
- grade_guard(): 1-(1-p)^k, k*t, k*l, k*m
- estimate_cost() for all 11 constructs with model-specific pricing
- format_cost_estimate() for human-readable output
- 18 passing tests

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
kenny67nju 5 months ago
parent
commit
9586ae39d8
2 changed files with 512 additions and 0 deletions
  1. 295 0
      lambdagent/cost_grade.py
  2. 217 0
      lambdagent/tests/test_cost_grade.py

+ 295 - 0
lambdagent/cost_grade.py

@@ -0,0 +1,295 @@
+"""
+lambdagent.cost_grade — Paper III §4.3 分级类型用于成本预测
+
+实现论文 III 的分级类型系统 (Definitions 11-12):
+  - CostGrade: 静态成本上界 (p, t, l, m)
+      p = 成功概率
+      t = token 数上界
+      l = 延迟上界 (秒)
+      m = 成本上界 (USD)
+  - 分级组合规则:
+      串行  g1 · g2 = (p1*p2, t1+t2, l1+l2, m1+m2)
+      并行  g1 ∥ g2 = (p1*p2, t1+t2, max(l1,l2), m1+m2)
+      迭代  g^n     = (p^n, n*t, n*l, n*m)
+      Guard g(k)    = (1-(1-p)^k, k*t, k*l, k*m)
+
+核心方程:
+    estimate_cost(agent) → CostGrade
+
+依赖:
+    types.py (AgentType), effects.py (Effect)
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+from .core import Term
+
+
+# ============================================================
+# CostGrade (Paper III Definition 11)
+# ============================================================
+
+@dataclass(frozen=True)
+class CostGrade:
+    """
+    分级类型: 静态成本上界。
+
+    Paper III Definition 11:
+        g = (p, t, l, m) where:
+        p ∈ [0, 1]  — 成功概率
+        t ∈ ℕ       — token 数上界
+        l ∈ ℝ⁺      — 延迟上界 (秒)
+        m ∈ ℝ⁺      — 成本上界 (USD)
+    """
+    probability: float = 1.0    # p: 成功概率
+    tokens: int = 0             # t: token 数上界
+    latency: float = 0.0        # l: 延迟上界 (秒)
+    money: float = 0.0          # m: 成本上界 (USD)
+
+    def __repr__(self) -> str:
+        return (
+            f"CostGrade(p={self.probability:.2%}, "
+            f"t={self.tokens}, "
+            f"l={self.latency:.1f}s, "
+            f"m=${self.money:.4f})"
+        )
+
+    @property
+    def is_free(self) -> bool:
+        """是否零成本(纯计算)"""
+        return self.tokens == 0 and self.money == 0.0
+
+
+# ============================================================
+# 分级组合规则 (Paper III Definition 12)
+# ============================================================
+
+def grade_serial(g1: CostGrade, g2: CostGrade) -> CostGrade:
+    """
+    串行组合: g1 · g2
+
+    Paper III Definition 12:
+        p = p1 × p2
+        t = t1 + t2
+        l = l1 + l2
+        m = m1 + m2
+    """
+    return CostGrade(
+        probability=g1.probability * g2.probability,
+        tokens=g1.tokens + g2.tokens,
+        latency=g1.latency + g2.latency,
+        money=g1.money + g2.money,
+    )
+
+
+def grade_parallel(g1: CostGrade, g2: CostGrade) -> CostGrade:
+    """
+    并行组合: g1 ∥ g2
+
+    Paper III Definition 12:
+        p = p1 × p2
+        t = t1 + t2
+        l = max(l1, l2)
+        m = m1 + m2
+    """
+    return CostGrade(
+        probability=g1.probability * g2.probability,
+        tokens=g1.tokens + g2.tokens,
+        latency=max(g1.latency, g2.latency),
+        money=g1.money + g2.money,
+    )
+
+
+def grade_iterate(g: CostGrade, n: int) -> CostGrade:
+    """
+    迭代: g^n
+
+    Paper III Definition 12:
+        p = p^n
+        t = n × t
+        l = n × l
+        m = n × m
+    """
+    return CostGrade(
+        probability=g.probability ** n,
+        tokens=n * g.tokens,
+        latency=n * g.latency,
+        money=n * g.money,
+    )
+
+
+def grade_guard(g: CostGrade, retries: int) -> CostGrade:
+    """
+    Guard 重试: (1+k) 次尝试
+
+    Paper III:
+        p = 1 - (1-p)^k  (至少一次成功的概率)
+        t = k × t
+        l = k × l
+        m = k × m
+    """
+    k = 1 + retries
+    return CostGrade(
+        probability=1.0 - (1.0 - g.probability) ** k,
+        tokens=k * g.tokens,
+        latency=k * g.latency,
+        money=k * g.money,
+    )
+
+
+# ============================================================
+# 模型成本配置
+# ============================================================
+
+# 每个模型的默认成本参数
+_MODEL_COSTS: Dict[str, Dict[str, float]] = {
+    "claude-sonnet-4-20250514": {"tokens_per_call": 800, "latency": 2.0, "price_per_1k": 0.003},
+    "claude-opus-4-20250514": {"tokens_per_call": 1200, "latency": 5.0, "price_per_1k": 0.015},
+    "claude-haiku-4-5-20251001": {"tokens_per_call": 500, "latency": 0.5, "price_per_1k": 0.00025},
+    "gpt-4": {"tokens_per_call": 1000, "latency": 3.0, "price_per_1k": 0.03},
+    "gpt-4o": {"tokens_per_call": 800, "latency": 1.5, "price_per_1k": 0.0025},
+    "qwen3-max": {"tokens_per_call": 600, "latency": 1.0, "price_per_1k": 0.001},
+}
+
+_DEFAULT_MODEL_COST = {"tokens_per_call": 800, "latency": 2.0, "price_per_1k": 0.003}
+
+
+def _get_model_cost(model: str) -> Dict[str, float]:
+    """获取模型的成本参数"""
+    for key, cost in _MODEL_COSTS.items():
+        if key in model.lower():
+            return cost
+    return _DEFAULT_MODEL_COST
+
+
+# ============================================================
+# 成本估算: estimate_cost(term) → CostGrade
+# ============================================================
+
+# 默认 LLM 成功概率(基于经验)
+_DEFAULT_LLM_SUCCESS_PROB = 0.95
+_DEFAULT_TOOL_SUCCESS_PROB = 0.98
+
+
+def estimate_cost(term: Term, model_costs: Dict[str, Dict[str, float]] | None = None) -> CostGrade:
+    """
+    静态估算 Agent 的最坏情况成本。
+
+    Paper III §4.3: 编译时为每个 agent 流水线计算成本上界。
+
+    Args:
+        term: Agent term
+        model_costs: 自定义模型成本参数 (可选)
+
+    Returns:
+        CostGrade: 成本上界 (p, t, l, m)
+    """
+    from .primitives import Lam, Compose, If, Loop, Pair, Fst, Snd, Tool
+    from .extensions import Par, Route, Memory, Guard
+
+    if model_costs:
+        _MODEL_COSTS.update(model_costs)
+
+    if isinstance(term, Lam):
+        mc = _get_model_cost(term.model)
+        tokens = mc["tokens_per_call"]
+        return CostGrade(
+            probability=_DEFAULT_LLM_SUCCESS_PROB,
+            tokens=tokens,
+            latency=mc["latency"],
+            money=tokens / 1000.0 * mc["price_per_1k"],
+        )
+
+    elif isinstance(term, Tool):
+        return CostGrade(
+            probability=_DEFAULT_TOOL_SUCCESS_PROB,
+            tokens=0,
+            latency=0.1,  # 100ms 默认工具延迟
+            money=0.0,
+        )
+
+    elif isinstance(term, (Fst, Snd)):
+        return CostGrade()  # 零成本
+
+    elif isinstance(term, Compose):
+        grades = [estimate_cost(s) for s in term.stages]
+        result = grades[0]
+        for g in grades[1:]:
+            result = grade_serial(result, g)
+        return result
+
+    elif isinstance(term, Pair):
+        g1 = estimate_cost(term.first)
+        g2 = estimate_cost(term.second)
+        return grade_parallel(g1, g2)
+
+    elif isinstance(term, Par):
+        grades = [estimate_cost(a) for a in term.agents]
+        result = grades[0]
+        for g in grades[1:]:
+            result = grade_parallel(result, g)
+        return result
+
+    elif isinstance(term, If):
+        g_then = estimate_cost(term.then_)
+        g_else = estimate_cost(term.else_)
+        # 最坏情况: 取成本更高的分支
+        cond_cost = CostGrade()
+        if isinstance(term.cond, Term):
+            cond_cost = estimate_cost(term.cond)
+        worst = g_then if g_then.money >= g_else.money else g_else
+        return grade_serial(cond_cost, worst)
+
+    elif isinstance(term, Loop):
+        body_cost = estimate_cost(term.body)
+        return grade_iterate(body_cost, term.max_steps)
+
+    elif isinstance(term, Guard):
+        inner_cost = estimate_cost(term.agent)
+        return grade_guard(inner_cost, term.retry)
+
+    elif isinstance(term, Memory):
+        inner_cost = estimate_cost(term.agent)
+        # Memory 本身几乎零成本
+        return grade_serial(CostGrade(latency=0.001), inner_cost)
+
+    elif isinstance(term, Route):
+        cls_cost = estimate_cost(term.classifier)
+        # 最坏情况: 选成本最高的路由
+        route_costs = [estimate_cost(r) for r in term.routes.values()]
+        if route_costs:
+            worst_route = max(route_costs, key=lambda g: g.money)
+            return grade_serial(cls_cost, worst_route)
+        return cls_cost
+
+    else:
+        # 多智能体扩展等 — 尝试分析
+        try:
+            from .multiagent import AsyncPar
+            if isinstance(term, AsyncPar):
+                grades = [estimate_cost(a) for a in term.agents]
+                result = grades[0]
+                for g in grades[1:]:
+                    result = grade_parallel(result, g)
+                return result
+        except ImportError:
+            pass
+
+    return CostGrade()  # 未知 term → 零成本(保守下界)
+
+
+def format_cost_estimate(grade: CostGrade) -> str:
+    """格式化成本估算为人类可读字符串"""
+    lines = [
+        f"┌─ Cost Estimate (Paper III §4.3) ─────────────┐",
+        f"│ Success probability: {grade.probability:.1%}",
+        f"│ Max tokens:         {grade.tokens:,}",
+        f"│ Max latency:        {grade.latency:.1f}s",
+        f"│ Max cost:           ${grade.money:.4f}",
+        f"└──────────────────────────────────────────────┘",
+    ]
+    return "\n".join(lines)

+ 217 - 0
lambdagent/tests/test_cost_grade.py

@@ -0,0 +1,217 @@
+"""
+Tests for Paper III §4.3 Graded Cost Prediction.
+
+Tests cover:
+  1. CostGrade construction and properties
+  2. Grade composition rules (serial, parallel, iterate, guard)
+  3. estimate_cost() for all constructs
+  4. Pipeline cost estimation
+  5. format_cost_estimate()
+"""
+
+import pytest
+from lambdagent.cost_grade import (
+    CostGrade,
+    grade_serial, grade_parallel, grade_iterate, grade_guard,
+    estimate_cost, format_cost_estimate,
+)
+from lambdagent.primitives import Lam, Compose, If, Loop, Pair, Fst, Snd, Tool
+from lambdagent.extensions import Par, Route, Memory, Guard
+
+
+# ============================================================
+# 1. CostGrade Construction
+# ============================================================
+
+class TestCostGrade:
+
+    def test_default_grade(self):
+        g = CostGrade()
+        assert g.probability == 1.0
+        assert g.tokens == 0
+        assert g.is_free
+
+    def test_repr(self):
+        g = CostGrade(probability=0.95, tokens=800, latency=2.0, money=0.0024)
+        r = repr(g)
+        assert "95.00%" in r
+        assert "800" in r
+
+
+# ============================================================
+# 2. Grade Composition Rules (Paper III Definition 12)
+# ============================================================
+
+class TestGradeComposition:
+
+    def test_serial(self):
+        """g1 · g2: probabilities multiply, everything else adds"""
+        g1 = CostGrade(0.9, 100, 1.0, 0.01)
+        g2 = CostGrade(0.8, 200, 2.0, 0.02)
+        r = grade_serial(g1, g2)
+        assert r.probability == pytest.approx(0.72)
+        assert r.tokens == 300
+        assert r.latency == pytest.approx(3.0)
+        assert r.money == pytest.approx(0.03)
+
+    def test_parallel(self):
+        """g1 ∥ g2: latency is max, rest same as serial"""
+        g1 = CostGrade(0.9, 100, 1.0, 0.01)
+        g2 = CostGrade(0.8, 200, 2.0, 0.02)
+        r = grade_parallel(g1, g2)
+        assert r.probability == pytest.approx(0.72)
+        assert r.tokens == 300
+        assert r.latency == pytest.approx(2.0)  # max(1, 2)
+        assert r.money == pytest.approx(0.03)
+
+    def test_iterate(self):
+        """g^n: probability is p^n, rest multiplied by n"""
+        g = CostGrade(0.9, 100, 1.0, 0.01)
+        r = grade_iterate(g, 5)
+        assert r.probability == pytest.approx(0.9**5)
+        assert r.tokens == 500
+        assert r.latency == pytest.approx(5.0)
+        assert r.money == pytest.approx(0.05)
+
+    def test_guard(self):
+        """Guard: p = 1-(1-p)^k, rest multiplied by k"""
+        g = CostGrade(0.8, 100, 1.0, 0.01)
+        r = grade_guard(g, retries=2)  # 3 total attempts
+        assert r.probability == pytest.approx(1.0 - 0.2**3)
+        assert r.tokens == 300
+        assert r.latency == pytest.approx(3.0)
+        assert r.money == pytest.approx(0.03)
+
+
+# ============================================================
+# 3. estimate_cost() for All Constructs
+# ============================================================
+
+class TestEstimateCost:
+
+    def test_lam(self):
+        """Lam has LLM cost"""
+        lam = Lam("test", "prompt", model="qwen3-max")
+        g = estimate_cost(lam)
+        assert g.tokens > 0
+        assert g.latency > 0
+        assert g.money > 0
+        assert 0 < g.probability < 1
+
+    def test_tool(self):
+        """Tool has zero token cost"""
+        tool = Tool("t", lambda x: x)
+        g = estimate_cost(tool)
+        assert g.tokens == 0
+        assert g.money == 0.0
+        assert g.latency > 0  # still has latency
+
+    def test_fst_snd(self):
+        """Fst/Snd are free"""
+        assert estimate_cost(Fst()).is_free
+        assert estimate_cost(Snd()).is_free
+
+    def test_compose(self):
+        """Compose costs add up"""
+        a = Lam("a", "p", model="qwen3-max")
+        b = Tool("b", lambda x: x)
+        g = estimate_cost(Compose(a, b))
+        ga = estimate_cost(a)
+        gb = estimate_cost(b)
+        assert g.tokens == ga.tokens + gb.tokens
+        assert g.latency == pytest.approx(ga.latency + gb.latency)
+
+    def test_pair(self):
+        """Pair: latency is max"""
+        a = Lam("a", "p", model="qwen3-max")
+        b = Lam("b", "p", model="qwen3-max")
+        g = estimate_cost(Pair(a, b))
+        ga = estimate_cost(a)
+        assert g.tokens == ga.tokens * 2
+        assert g.latency == pytest.approx(ga.latency)  # max(l, l) = l
+
+    def test_loop(self):
+        """Loop: cost * max_steps"""
+        body = Lam("b", "p", model="qwen3-max")
+        loop = Loop(body, lambda r, s: False, max_steps=5)
+        g = estimate_cost(loop)
+        gb = estimate_cost(body)
+        assert g.tokens == gb.tokens * 5
+        assert g.latency == pytest.approx(gb.latency * 5)
+
+    def test_guard(self):
+        """Guard: retries increase cost but improve probability"""
+        inner = Lam("a", "p", model="qwen3-max")
+        guard = Guard(inner, lambda x: True, retry=2)
+        g = estimate_cost(guard)
+        gi = estimate_cost(inner)
+        assert g.probability > gi.probability  # retries help
+        assert g.tokens == gi.tokens * 3
+
+    def test_route(self):
+        """Route: classifier + worst route"""
+        cls = Lam("cls", "classify", model="qwen3-max")
+        routes = {
+            "cheap": Tool("t", lambda x: x),
+            "expensive": Lam("exp", "p", model="claude-opus-4-20250514"),
+        }
+        route = Route(cls, routes)
+        g = estimate_cost(route)
+        gc = estimate_cost(cls)
+        assert g.tokens > gc.tokens  # classifier + route
+
+    def test_memory(self):
+        """Memory wraps inner cost"""
+        inner = Tool("t", lambda x: x)
+        mem = Memory(inner)
+        g = estimate_cost(mem)
+        gi = estimate_cost(inner)
+        assert g.tokens == gi.tokens
+
+
+# ============================================================
+# 4. Pipeline Cost Estimation
+# ============================================================
+
+class TestPipelineCost:
+
+    def test_realistic_pipeline(self):
+        """Realistic pipeline: summarize >> validate >> format"""
+        summarizer = Lam("summarize", "Summarize", model="qwen3-max")
+        validator = Guard(
+            Lam("validate", "Validate", model="qwen3-max"),
+            lambda x: len(x) > 10,
+            retry=2,
+        )
+        formatter = Tool("format", lambda x: f"<p>{x}</p>")
+        pipeline = Compose(summarizer, validator, formatter)
+
+        g = estimate_cost(pipeline)
+        assert g.tokens > 0
+        assert g.money > 0
+        assert g.probability < 1.0
+        # Should be sum of components
+        print(format_cost_estimate(g))
+
+    def test_parallel_reduces_latency(self):
+        """Parallel should have lower latency than serial"""
+        a = Lam("a", "p", model="qwen3-max")
+        b = Lam("b", "p", model="qwen3-max")
+        serial = grade_serial(estimate_cost(a), estimate_cost(b))
+        parallel = grade_parallel(estimate_cost(a), estimate_cost(b))
+        assert parallel.latency < serial.latency
+        assert parallel.money == serial.money  # same total cost
+
+
+# ============================================================
+# 5. format_cost_estimate()
+# ============================================================
+
+class TestFormatCostEstimate:
+
+    def test_format(self):
+        g = CostGrade(0.95, 800, 2.0, 0.0024)
+        s = format_cost_estimate(g)
+        assert "95.0%" in s
+        assert "800" in s
+        assert "$0.0024" in s