244 lines
7.8 KiB
Python
244 lines
7.8 KiB
Python
"""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_structured_output_threads_usage() -> None:
|
||
# 结构化输出必须记真实 usage（CR-H3）：instructor 的 create_with_completion 回
|
||
# (parsed, raw)，raw.usage 带真实 token 数；网关据此记账（§4.8）。
|
||
class _RawUsage:
|
||
input_tokens = 120
|
||
output_tokens = 80
|
||
cache_read_input_tokens = 30
|
||
|
||
class _Raw:
|
||
usage = _RawUsage()
|
||
|
||
class _FakeStructured:
|
||
def __init__(self) -> None:
|
||
self.last_kwargs: dict[str, Any] = {}
|
||
|
||
async def create_with_completion(
|
||
self, *, response_model: type[BaseModel], **kwargs: Any
|
||
) -> tuple[BaseModel, Any]:
|
||
self.last_kwargs = kwargs
|
||
return response_model(score=7), _Raw()
|
||
|
||
# 修前路径调 `.create`（仅回 parsed）——保留它，令 RED 是值不符（0!=120）
|
||
# 而非 AttributeError，坐实旧代码硬编码零 usage 的缺陷。
|
||
async def create(self, *, response_model: type[BaseModel], **kwargs: Any) -> BaseModel:
|
||
return response_model(score=7)
|
||
|
||
fake = _FakeStructured()
|
||
client = _FakeAnthropicClient(_FakeAnthropicMessages("", 0, 0, 0))
|
||
adapter = AnthropicAdapter("anthropic", client, structured_client=fake)
|
||
req = LlmRequest(
|
||
tier="analyst",
|
||
input="打分",
|
||
output_schema=Out,
|
||
scope=Scope(user_id=uuid.UUID(int=1)),
|
||
)
|
||
|
||
result = await adapter.complete(req, "claude-3-5-sonnet")
|
||
|
||
assert isinstance(result.parsed, Out)
|
||
assert result.parsed.score == 7
|
||
assert result.usage.input_tokens == 120
|
||
assert result.usage.output_tokens == 80
|
||
assert result.usage.cache_read_tokens == 30
|
||
|
||
|
||
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"
|