LiteLLMとLangGraphで構築したAIエージェントをMicrosoft製agent-sreでSRE運用してみる

Python 3.13.9
agent-sre 3.5.0
langgraph 1.1.8
langchain-litellm 0.6.4
litellm 1.83.9

$ uv pip install \
    "agent-sre==3.5.0" \
    "langgraph==1.1.8" \
    "langchain-litellm==0.6.4" \
    "litellm==1.83.9"

$ export OPENAI_API_KEY="sk-..."

from langchain_core.tools import tool

@tool
def web_search(query: str) -> str:
    """Web で情報を検索する."""
    data = {
        "東京タワー": "東京タワーの高さは333メートルです。1958年に完成しました。",
        "富士山": "富士山の標高は3,776メートルで、日本最高峰です。",
    }
    for key, val in data.items():
        if key in query:
            return val
    return f"'{query}' の検索結果: 関連情報が見つかりました。"

@tool
def calculator(expression: str) -> str:
    """数式を計算する（四則演算のみ対応）."""
    allowed = set("0123456789+-*/.() ")
    if not all(c in allowed for c in expression):
        return "エラー: サポートされない式です"
    return str(eval(expression))

@tool
def get_weather(location: str) -> str:
    """指定した場所の現在の天気を取得する."""
    data = {
        "東京": "晴れ、気温 25°C、湿度 60%",
        "大阪": "曇り、気温 23°C、湿度 70%",
    }
    return f"{location}: {data.get(location, '晴れ、気温 20°C')}"

from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
from langchain_litellm import ChatLiteLLM
from langgraph.graph import StateGraph, MessagesState, START, END

TOOLS = [web_search, calculator, get_weather]
TOOLS_BY_NAME = {t.name: t for t in TOOLS}

# LiteLLM 経由で LLM を初期化。モデル名を変えるだけでプロバイダー切替が可能
LLM = ChatLiteLLM(model="gpt-4o-mini", temperature=0)
LLM_WITH_TOOLS = LLM.bind_tools(TOOLS)

SYSTEM_PROMPT = (
    "あなたはツールを使ってユーザーの質問に回答するアシスタントです。"
    "必要に応じてツールを呼び出してください。"
)

def llm_node(state: MessagesState):
    msgs = [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
    return {"messages": [LLM_WITH_TOOLS.invoke(msgs)]}

def tool_node(state: MessagesState):
    results = []
    for tc in state["messages"][-1].tool_calls:
        fn = TOOLS_BY_NAME[tc["name"]]
        out = fn.invoke(tc["args"])
        results.append(ToolMessage(content=str(out), tool_call_id=tc["id"]))
    return {"messages": results}

def should_continue(state: MessagesState):
    if state["messages"][-1].tool_calls:
        return "tool_node"
    return END

def build_agent():
    graph = StateGraph(MessagesState)
    graph.add_node("llm", llm_node)
    graph.add_node("tool_node", tool_node)

    graph.add_edge(START, "llm")
    graph.add_conditional_edges("llm", should_continue, ["tool_node", END])
    graph.add_edge("tool_node", "llm")
    return graph.compile()

AGENT = build_agent()

from agent_sre.anomaly import (
    RiskLevel,
    RogueAgentDetector,
    RogueDetectorConfig,
)
from agent_sre.cascade.circuit_breaker import (
    CircuitBreaker,
    CircuitBreakerConfig,
    CircuitOpenError,
)
from agent_sre.cost import CostGuard
from agent_sre.slo import SLI, SLO, ErrorBudget, SLIValue
from agent_sre.slo.objectives import ExhaustionAction

AGENT_ID = "research-agent"

breaker = CircuitBreaker(
    agent_id=AGENT_ID,
    config=CircuitBreakerConfig(
        failure_threshold=3,
        recovery_timeout_seconds=10.0,
        half_open_max_calls=1,
    ),
)

class TaskSuccessRateSLI(SLI):
    """タスク成功率 SLI."""

    def collect(self) -> SLIValue:
        values = self.values_in_window()
        if not values:
            return self.record(1.0)
        good = sum(1 for v in values if v.is_good)
        return self.record(good / len(values))

slo = SLO(
    name=f"{AGENT_ID}-reliability",
    indicators=[
        TaskSuccessRateSLI(name="task_success_rate", target=0.95, window="1h"),
    ],
    error_budget=ErrorBudget(
        total=0.05,
        window_seconds=3600,
        burn_rate_alert=2.0,
        burn_rate_critical=10.0,
        exhaustion_action=ExhaustionAction.THROTTLE,
    ),
    agent_id=AGENT_ID,
)

detector = RogueAgentDetector(
    config=RogueDetectorConfig(
        frequency_z_threshold=2.0,
        quarantine_risk_level=RiskLevel.HIGH,
    ),
)
detector.register_capability_profile(
    agent_id=AGENT_ID,
    allowed_tools=["web_search", "calculator", "get_weather"],
)

cost_guard = CostGuard(
    per_task_limit=0.10,
    per_agent_daily_limit=1.00,
    org_monthly_budget=100.00,
    auto_throttle=True,
    kill_switch_threshold=0.95,
)

def init_sre_components(
    failure_threshold: int = 3,
    recovery_timeout: float = 10.0,
) -> tuple[CircuitBreaker, SLO, RogueAgentDetector, CostGuard]:
    breaker = CircuitBreaker(
        agent_id=AGENT_ID,
        config=CircuitBreakerConfig(
            failure_threshold=failure_threshold,
            recovery_timeout_seconds=recovery_timeout,
            half_open_max_calls=1,
        ),
    )
    slo = SLO(
        name=f"{AGENT_ID}-reliability",
        indicators=[
            TaskSuccessRateSLI(name="task_success_rate", target=0.95, window="1h"),
        ],
        error_budget=ErrorBudget(
            total=0.05,
            window_seconds=3600,
            burn_rate_alert=2.0,
            burn_rate_critical=10.0,
            exhaustion_action=ExhaustionAction.THROTTLE,
        ),
        agent_id=AGENT_ID,
    )
    detector = RogueAgentDetector(
        config=RogueDetectorConfig(
            frequency_z_threshold=2.0,
            quarantine_risk_level=RiskLevel.HIGH,
        ),
    )
    detector.register_capability_profile(
        agent_id=AGENT_ID,
        allowed_tools=["web_search", "calculator", "get_weather"],
    )
    cost_guard = CostGuard(
        per_task_limit=0.10,
        per_agent_daily_limit=1.00,
        org_monthly_budget=100.00,
        auto_throttle=True,
        kill_switch_threshold=0.95,
    )
    return breaker, slo, detector, cost_guard

import time

import litellm
from langchain_core.messages import AIMessage, HumanMessage

MODEL = "gpt-4o-mini"

def run_with_sre(
    question: str,
    *,
    breaker: CircuitBreaker,
    slo: SLO,
    detector: RogueAgentDetector,
    cost_guard: CostGuard,
    estimated_cost: float = 0.005,
    simulate_failure: bool = False,
) -> dict:
    # 1. コストチェック（プリフライト）
    allowed, reason = cost_guard.check_task(AGENT_ID, estimated_cost=estimated_cost)
    if not allowed:
        return {"status": "cost_blocked", "reason": reason}

    # 2. サーキットブレーカーでエージェント実行を包む
    def _invoke():
        if simulate_failure:
            raise RuntimeError("simulated_fault")
        return AGENT.invoke({"messages": [HumanMessage(content=question)]})

    try:
        result = breaker.call(_invoke)
    except CircuitOpenError as e:
        slo.indicators[0].record(0.0)
        slo.record_event(good=False)
        return {"status": "circuit_open", "retry_after": round(e.retry_after, 1),
                "cb_state": breaker.state}
    except Exception as e:  # CircuitBreaker.call() は任意の例外を失敗として再送出する
        slo.indicators[0].record(0.0)
        slo.record_event(good=False)
        return {"status": "failed", "error": str(e), "cb_state": breaker.state}

    # 3. ツール使用状況を収集して不正検出へ
    tool_calls = []
    for msg in result["messages"]:
        if hasattr(msg, "tool_calls") and msg.tool_calls:
            tool_calls.extend(tc["name"] for tc in msg.tool_calls)
    for name in tool_calls:
        detector.record_action(
            agent_id=AGENT_ID, action="tool_call", tool_name=name,
            timestamp=time.time(),
        )
    assessment = detector.assess(AGENT_ID)

    # 4. SLO 記録（SLI に値 → SLO 経由で error_budget と evaluate をまとめて更新）
    slo.indicators[0].record(1.0)
    slo.record_event(good=True)

    # 5. LLM のトークン使用量から実コストを算出して CostGuard に記録
    input_tokens = 0
    output_tokens = 0
    for msg in result["messages"]:
        if isinstance(msg, AIMessage) and msg.usage_metadata:
            input_tokens += msg.usage_metadata.get("input_tokens", 0)
            output_tokens += msg.usage_metadata.get("output_tokens", 0)
    prompt_cost, completion_cost = litellm.cost_per_token(
        model=MODEL,
        prompt_tokens=input_tokens,
        completion_tokens=output_tokens,
    )
    actual_cost = prompt_cost + completion_cost
    cost_guard.record_cost(
        agent_id=AGENT_ID,
        task_id=f"task-{int(time.time()*1000)}",
        cost_usd=actual_cost,
        breakdown={"input": prompt_cost, "output": completion_cost},
    )

    return {
        "status": "success",
        "response": result["messages"][-1].content,
        "tool_calls": tool_calls,
        "risk_level": assessment.risk_level.value,
    }

$ python demo.py

============================================================
デモ 1: エージェント実行と Agent SRE 監視
============================================================

  質問 1: 東京タワーの高さは何メートルですか？
  回答: 東京タワーの高さは333メートルです。1958年に完成しました。
  レイテンシ=3.47s  ツール=['web_search']  リスク=low
  SLO バジェット残=100.0%  累積コスト=$0.0001

  質問 2: 123 * 456 を計算してください。
  回答: 123 * 456 の計算結果は 56088 です。
  レイテンシ=1.72s  ツール=['calculator']  リスク=low
  SLO バジェット残=100.0%  累積コスト=$0.0001

  質問 3: 東京の天気を教えてください。
  回答: 東京の天気は晴れで、気温は25°C、湿度は60%です。
  レイテンシ=1.78s  ツール=['get_weather']  リスク=low
  SLO バジェット残=100.0%  累積コスト=$0.0002

breaker, slo, detector, cost_guard = init_sre_components(
    failure_threshold=3, recovery_timeout=2.0,
)

============================================================
デモ 2: サーキットブレーカー (失敗シミュレーション)
============================================================
  リクエスト 1: failed          CB=CLOSED
  リクエスト 2: failed          CB=CLOSED
  リクエスト 3: failed          CB=OPEN
  リクエスト 4: circuit_open    CB=OPEN
  リクエスト 5: circuit_open    CB=OPEN

  (recovery_timeout=2秒 待機...)
  リクエスト 6 (復帰): success         CB=CLOSED

============================================================
デモ 3: SLO トラッキングとエラーバジェット
============================================================
  SLO: タスク成功率 >= 95%  エラーバジェット: 5%

  [1] success         バジェット残=100.0%  バーンレート= 0.00  SLO=healthy    アラート=[]
  [2] success         バジェット残=100.0%  バーンレート= 0.00  SLO=healthy    アラート=[]
  [3] failed          バジェット残=  0.0%  バーンレート=24000.00  SLO=exhausted  アラート=['burn_rate_warning', 'burn_rate_critical']
  [4] success         バジェット残=  0.0%  バーンレート=18000.00  SLO=exhausted  アラート=['burn_rate_warning', 'burn_rate_critical']
  [5] failed          バジェット残=  0.0%  バーンレート=28800.00  SLO=exhausted  アラート=['burn_rate_warning', 'burn_rate_critical']
  [6] success         バジェット残=  0.0%  バーンレート=24000.00  SLO=exhausted  アラート=['burn_rate_warning', 'burn_rate_critical']
  [7] success         バジェット残=  0.0%  バーンレート=20571.43  SLO=exhausted  アラート=['burn_rate_warning', 'burn_rate_critical']
  [8] success         バジェット残=  0.0%  バーンレート=18000.00  SLO=exhausted  アラート=['burn_rate_warning', 'burn_rate_critical']

============================================================
デモ 4: 不正エージェント検出（3シグナル別）
============================================================
  A) 正常              リスク=low       スコア=0.00  隔離推奨=False
    capability=0.00  frequency=0.00  entropy=0.00
  B) 許可外ツール多用  リスク=low       スコア=0.50  隔離推奨=False
    capability=0.50  frequency=0.00  entropy=0.00
  C) ループ (低entropy)  リスク=medium    スコア=1.00  隔離推奨=False
    capability=0.00  frequency=0.00  entropy=1.00
  D) バースト攻撃     リスク=critical  スコア=18.30  隔離推奨=True
    capability=0.00  frequency=18.30  entropy=0.00

guard = CostGuard(
    per_task_limit=0.10,        # 1タスク上限 $0.10
    per_agent_daily_limit=0.50, # 1エージェントの1日上限 $0.50
    auto_throttle=True,
    kill_switch_threshold=0.95,
)

============================================================
デモ 5: コストガード
============================================================
  タスク1(通常): allowed=True  支出=$0.05/0.50  利用率=10%
  タスク2(通常): allowed=True  支出=$0.10/0.50  利用率=20%
  タスク3(通常): allowed=True  支出=$0.15/0.50  利用率=30%

  大型タスク: allowed=False  reason=Estimated cost $0.50 exceeds per-task limit $0.10
  タスク1: 支出=$0.20  利用率= 40%  状態=ok  アラート=[]
  タスク2: 支出=$0.25  利用率= 50%  状態=ok  アラート=['warning']
  タスク3: 支出=$0.30  利用率= 60%  状態=ok  アラート=[]
  タスク4: 支出=$0.35  利用率= 70%  状態=ok  アラート=[]
  タスク5: 支出=$0.40  利用率= 80%  状態=ok  アラート=['warning']
  タスク6: 支出=$0.45  利用率= 90%  状態=throttled  アラート=['warning']
  タスク7: 支出=$0.50  利用率=100%  状態=throttled,killed  アラート=['critical', 'critical', 'critical']
  タスク8: 支出=$0.55  利用率=110%  状態=throttled,killed  アラート=['critical']

from agent_sre.slo.persistence import SQLiteMeasurementStore

store = SQLiteMeasurementStore(db_path="~/.agent_sre/sli.db")
sli = TaskSuccessRateSLI(
    name="task_success_rate", target=0.95, window="1h", store=store,
)

{
  "total": 5.0,
  "consumed": 1.0,
  "remaining_percent": 80.0,
  "is_exhausted": false,
  "burn_rate": 360.0,
  "exhaustion_action": "throttle",
  "firing_alerts": ["burn_rate_warning", "burn_rate_critical"]
}

{
  "agent_id": "ag1",
  "daily_limit_usd": 1.0,
  "per_task_limit_usd": 0.1,
  "spent_today_usd": 0.15,
  "remaining_today_usd": 0.85,
  "utilization_percent": 15.0,
  "task_count_today": 2,
  "avg_cost_per_task": 0.075,
  "throttled": false,
  "killed": false
}

# シャットダウン時 / 定期的にスナップショットを保存
snapshot = {
    "error_budget": slo.error_budget.to_dict(),
    "cost_budgets": {aid: cg.get_budget(aid).to_dict() for aid in agent_ids},
}
redis.set("agent-sre:state", json.dumps(snapshot))

# 起動時に復元（from_dict が無いので内部状態を直接書き戻す）
saved = json.loads(redis.get("agent-sre:state"))
# consumed は失敗イベント件数の累積（record_event ごとに +1.0 される値）
slo.error_budget.consumed = saved["error_budget"]["consumed"]
slo.evaluate()  # consumed を書き戻した後に SLOStatus を明示的に再評価しておく
for aid, b in saved["cost_budgets"].items():
    budget = cost_guard.get_budget(aid)
    budget.spent_today_usd = b["spent_today_usd"]
    budget.task_count_today = b["task_count_today"]
    budget.throttled = b["throttled"]
    budget.killed = b["killed"]
# 注: _cost_history（Z-scoreベースライン）は private 属性のため公式の書き戻し手段が無く、
# 再起動直後は履歴 10 件未満の間 Z-score 異常検知が一時的に無効化される点に留意。