muse-agent-example/tests/skills/call-content-model/test_call_persistence.py

#!/usr/bin/env python3
"""模型调用持久化合同的离线测试。

这些测试不连网、不连库，只固定共享 LLM 入口必须向持久化适配器提供的证据形状。
"""
import json
import os
import pathlib
import sys
import types


import muse_llm as llm

llm.TOKEN = "test-token"

PROJECT_ROOT = pathlib.Path(__file__).resolve().parents[3]
EVIDENCE_SCRIPTS = PROJECT_ROOT / ".agent" / "skills" / "record-run-evidence" / "scripts"
sys.path.insert(0, str(EVIDENCE_SCRIPTS))
from persist_llm_call import (  # noqa: E402
    _event_model_match,
    _usage_input_tokens,
    _usage_output_tokens,
)

# 本文件直接测 chat() 传输与持久化形状，显式允许未治理直连。
os.environ["MUSE_LLM_ALLOW_UNGOVERNED"] = "1"


USAGE = {
    "prompt_tokens": 12,
    "completion_tokens": 7,
    "prompt_tokens_details": {"cached_tokens": 3},
}


class FakeResponse:
    status_code = 200
    text = ""

    def raise_for_status(self):
        return None

    def json(self):
        return {
            "choices": [{
                "message": {"content": "模型输出"},
                "finish_reason": "stop",
            }],
            "usage": USAGE,
            "id": "completion-1",
        }


def test_chat_emits_a_complete_persistence_event():
    """成功调用应把完整请求/响应和审计字段交给原子落库适配器。"""
    events = []

    class FakeSession:
        trust_env = True

        def post(self, url, headers=None, json=None, timeout=None):
            return FakeResponse()

    old_session, old_time = llm.requests.Session, llm.time
    try:
        llm.requests.Session = FakeSession
        llm.time = types.SimpleNamespace(time=lambda: 100.0, sleep=lambda _: None)
        content, usage = llm.chat(
            "用户提示",
            model="MiniMax-M3",
            system="系统提示",
            retries=0,
            run_id="run-1",
            caller="parse-book",
            persist_call=events.append,
        )
    finally:
        llm.requests.Session, llm.time = old_session, old_time

    assert (content, usage) == ("模型输出", USAGE)
    assert len(events) == 1
    event = events[0]
    assert event["run_id"] == "run-1"
    assert event["caller"] == "parse-book"
    assert event["requested_model_id"] == "MiniMax-M3"
    assert event["actual_model_id"] == "MiniMax-M3"
    assert event["model_match"] is True
    assert event["usage"] == USAGE
    assert event["duration_ms"] == 0
    assert event["stop_reason"] == "stop"
    assert json.loads(event["prompt"]) == {
        "messages": [
            {"role": "system", "content": "系统提示"},
            {"role": "user", "content": "用户提示"},
        ],
        "model": "MiniMax-M3",
        "max_tokens": 512000,
        "temperature": 0.2,
    }
    assert json.loads(event["response"])["choices"][0]["message"]["content"] == "模型输出"


def test_governed_forwards_persistence_context_to_actual_model_call():
    """额度治理选出的实际模型必须继续携带 run/caller/持久化适配器。"""
    calls = []

    def fake_chat(prompt, model=None, **kwargs):
        calls.append((prompt, model, kwargs))
        return "ok", {"prompt_tokens": 1, "completion_tokens": 1}

    old_now = llm._now
    old_read = llm._read_window
    old_bump = llm._bump_window
    old_chat = llm.chat
    old_pricing = llm._PRICING_CACHE
    try:
        llm._now = lambda: __import__("datetime").datetime(2026, 7, 16, 12, 0)
        llm._read_window = lambda _: (0.0, 0)
        llm._bump_window = lambda *_: (0.0, 1)
        llm._PRICING_CACHE = dict(llm.PRICING_FALLBACK)
        llm.chat = fake_chat
        marker = object()
        content, _, used = llm.chat_governed(
            "prompt",
            run_id="run-2",
            caller="extract-knowledge",
            persist_call=marker,
        )
    finally:
        llm._now = old_now
        llm._read_window = old_read
        llm._bump_window = old_bump
        llm.chat = old_chat
        llm._PRICING_CACHE = old_pricing

    assert (content, used) == ("ok", "MiniMax-M3")
    assert calls[0][2]["run_id"] == "run-2"
    assert calls[0][2]["caller"] == "extract-knowledge"
    assert calls[0][2]["persist_call"] is marker
    assert calls[0][2]["model_match"] is True


def test_anthropic_usage_and_policy_match_are_normalized():
    usage = {
        "input_tokens": 10,
        "cache_creation_input_tokens": 3,
        "cache_read_input_tokens": 7,
        "output_tokens": 5,
    }
    assert _usage_input_tokens(usage) == 20
    assert _usage_output_tokens(usage) == 5
    event = {"model_match": True}
    assert _event_model_match(event, "opus", "claude-opus-4-8[1M]") is True
    assert _event_model_match({}, "same", "same") is True
    assert _event_model_match({}, "opus", "claude-opus-4-8[1M]") is False


def test_chat_fails_closed_without_token():
    old = llm.TOKEN
    llm.TOKEN = ""
    try:
        try:
            llm.chat("x", retries=0, allow_ungoverned=True)
            raise AssertionError("缺令牌必须失败关闭")
        except RuntimeError as exc:
            assert "MUSE_LLM_TOKEN" in str(exc)
    finally:
        llm.TOKEN = old


def test_chat_fails_closed_without_ungoverned_permission():
    """未显式允许时 chat() 必须失败关闭，逼管线走 chat_governed。"""
    old = os.environ.pop("MUSE_LLM_ALLOW_UNGOVERNED", None)
    try:
        try:
            llm.chat("x", retries=0)
            raise AssertionError("未授权 chat() 应抛 RuntimeError")
        except RuntimeError as exc:
            assert "chat_governed" in str(exc)
    finally:
        if old is None:
            os.environ["MUSE_LLM_ALLOW_UNGOVERNED"] = "1"
        else:
            os.environ["MUSE_LLM_ALLOW_UNGOVERNED"] = old


def main():
    for test in (
        test_chat_emits_a_complete_persistence_event,
        test_governed_forwards_persistence_context_to_actual_model_call,
        test_anthropic_usage_and_policy_match_are_normalized,
        test_chat_fails_closed_without_token,
        test_chat_fails_closed_without_ungoverned_permission,
    ):
        test()
        print(f"  ✓ {test.__name__}")
    print("全部通过")


if __name__ == "__main__":
    main()