writer-work-flow/packages/llm_gateway/tests/test_new_adapters.py

"""T5.4 新适配器：Anthropic + Gemini（ARCH §4.2/§4.4）。

网络客户端经注入的 Protocol 替身——绝不联网、绝不依赖真实 SDK。
断言：能力矩阵声明正确、complete/stream 归一化、结构化输出、usage 提取。
"""

from __future__ import annotations

import uuid
from collections.abc import AsyncIterator
from typing import Any

from pydantic import BaseModel
from ww_llm_gateway.adapters.anthropic import AnthropicAdapter
from ww_llm_gateway.adapters.gemini import GeminiAdapter
from ww_llm_gateway.types import Block, LlmRequest, Scope


class Out(BaseModel):
    score: int


# ---- Anthropic ----


class _FakeAnthropicMessages:
    def __init__(self, text: str, in_tok: int, out_tok: int, cache_read: int) -> None:
        self._text = text
        self._in = in_tok
        self._out = out_tok
        self._cache_read = cache_read
        self.last_kwargs: dict[str, Any] = {}

    async def create(self, **kwargs: Any) -> Any:
        self.last_kwargs = kwargs

        class _Block:
            type = "text"
            text = self._text

        class _Usage:
            input_tokens = self._in
            output_tokens = self._out
            cache_read_input_tokens = self._cache_read

        class _Resp:
            content = [_Block()]
            usage = _Usage()

        return _Resp()


class _FakeAnthropicClient:
    def __init__(self, msgs: _FakeAnthropicMessages) -> None:
        self.messages = msgs


def _req(text: str = "写第一章", *, cache: bool = False) -> LlmRequest:
    sys = [Block(text="世界观", cache=cache)] if cache else []
    return LlmRequest(tier="writer", input=text, system=sys, scope=Scope(user_id=uuid.UUID(int=1)))


def test_anthropic_capabilities() -> None:
    client = _FakeAnthropicClient(_FakeAnthropicMessages("", 0, 0, 0))
    adapter = AnthropicAdapter("anthropic", client)
    caps = adapter.capabilities()
    assert caps.structured_output is True
    assert caps.prefix_cache is True
    assert caps.thinking is True


async def test_anthropic_complete_text_and_usage() -> None:
    msgs = _FakeAnthropicMessages("生成正文", in_tok=120, out_tok=80, cache_read=30)
    adapter = AnthropicAdapter("anthropic", _FakeAnthropicClient(msgs))

    result = await adapter.complete(_req(), "claude-3-5-sonnet")

    assert result.text == "生成正文"
    assert result.usage.input_tokens == 120
    assert result.usage.output_tokens == 80
    assert result.usage.cache_read_tokens == 30


async def test_anthropic_marks_cache_breakpoint() -> None:
    msgs = _FakeAnthropicMessages("x", 1, 1, 0)
    adapter = AnthropicAdapter("anthropic", _FakeAnthropicClient(msgs))

    await adapter.complete(_req(cache=True), "claude-3-5-sonnet")

    # 稳定块应带 cache_control 断点（Claude 显式缓存，§4.6）。
    system = msgs.last_kwargs["system"]
    assert any(blk.get("cache_control") for blk in system)


# ---- Gemini ----


class _FakeGeminiModels:
    def __init__(self, text: str, in_tok: int, out_tok: int) -> None:
        self._text = text
        self._in = in_tok
        self._out = out_tok
        self.last_kwargs: dict[str, Any] = {}

    async def generate_content(self, **kwargs: Any) -> Any:
        self.last_kwargs = kwargs

        class _Usage:
            prompt_token_count = self._in
            candidates_token_count = self._out
            cached_content_token_count = 0

        class _Resp:
            text = self._text
            usage_metadata = _Usage()

        return _Resp()


class _FakeGeminiAio:
    def __init__(self, models: _FakeGeminiModels) -> None:
        self.models = models


class _FakeGeminiClient:
    def __init__(self, models: _FakeGeminiModels) -> None:
        self.aio = _FakeGeminiAio(models)


def test_gemini_capabilities() -> None:
    adapter = GeminiAdapter("gemini", _FakeGeminiClient(_FakeGeminiModels("", 0, 0)))  # type: ignore[arg-type]
    caps = adapter.capabilities()
    assert caps.structured_output is True
    assert caps.thinking is True


async def test_gemini_complete_text_and_usage() -> None:
    models = _FakeGeminiModels("双子座正文", in_tok=200, out_tok=90)
    adapter = GeminiAdapter("gemini", _FakeGeminiClient(models))  # type: ignore[arg-type]

    result = await adapter.complete(_req(), "gemini-2.0-flash")

    assert result.text == "双子座正文"
    assert result.usage.input_tokens == 200
    assert result.usage.output_tokens == 90


async def test_gemini_structured_output() -> None:
    models = _FakeGeminiModels('{"score": 7}', in_tok=10, out_tok=5)
    adapter = GeminiAdapter("gemini", _FakeGeminiClient(models))  # type: ignore[arg-type]
    req = LlmRequest(
        tier="analyst",
        input="打分",
        output_schema=Out,
        scope=Scope(user_id=uuid.UUID(int=1)),
    )

    result = await adapter.complete(req, "gemini-2.0-flash")

    assert result.parsed is not None
    assert isinstance(result.parsed, Out)
    assert result.parsed.score == 7
    # 结构化请求应带 response schema 配置。
    assert "config" in models.last_kwargs


async def test_anthropic_stream_yields_text_then_usage() -> None:
    class _StreamEvent:
        def __init__(self, **kw: Any) -> None:
            self.__dict__.update(kw)

    class _StreamMessages:
        def stream(self, **kwargs: Any) -> Any:
            class _Ctx:
                async def __aenter__(self_inner) -> AsyncIterator[Any]:
                    async def _gen() -> AsyncIterator[Any]:
                        yield _StreamEvent(
                            type="content_block_delta", delta=_StreamEvent(text="hel")
                        )
                        yield _StreamEvent(
                            type="content_block_delta", delta=_StreamEvent(text="lo")
                        )

                    return _gen()

                async def __aexit__(self_inner, *a: Any) -> None:
                    return None

            return _Ctx()

    class _Client:
        def __init__(self) -> None:
            self.messages = _StreamMessages()

    adapter = AnthropicAdapter("anthropic", _Client())
    chunks = [c async for c in adapter.stream(_req(), "claude-3-5-sonnet")]
    text = "".join(c.text for c in chunks if c.text)
    assert text == "hello"