Spaces:
Running
Running
File size: 9,470 Bytes
39ff632 1ab43d7 39ff632 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 | """LLM client wire-protocol tests with a mocked httpx transport."""
from __future__ import annotations
import json
import httpx
import pytest
from app.llm import (LLMClient, LLMError, _chat_completions_url, _messages_url,
extract_json_object, user_turn, assistant_turn)
PNG = b"\x89PNG\r\n\x1a\nfake"
def make_settings(make_settings, **overrides):
defaults = {"LLM_MAX_RETRIES": 0, "LLM_TIMEOUT_S": 5}
defaults.update(overrides)
return make_settings(**defaults)
def anthropic_ok(text: str) -> httpx.Response:
return httpx.Response(200, json={
"content": [{"type": "text", "text": text}],
"stop_reason": "end_turn",
})
def openai_ok(text: str) -> httpx.Response:
return httpx.Response(200, json={
"choices": [{"message": {"content": text}, "finish_reason": "stop"}],
})
class TestUrlJoin:
def test_base_without_v1(self):
assert _messages_url("https://api.anthropic.com") == \
"https://api.anthropic.com/v1/messages"
assert _messages_url("https://openrouter.ai/api") == \
"https://openrouter.ai/api/v1/messages"
def test_base_with_v1(self):
assert _messages_url("https://x.test/v1") == "https://x.test/v1/messages"
assert _messages_url("https://x.test/v1/") == "https://x.test/v1/messages"
def test_chat_completions_join(self):
assert _chat_completions_url("https://openrouter.ai/api") == \
"https://openrouter.ai/api/v1/chat/completions"
assert _chat_completions_url("https://x.test/v1") == \
"https://x.test/v1/chat/completions"
class TestAnthropicCall:
@pytest.mark.asyncio
async def test_happy_path_headers_and_body(self, make_settings):
seen = {}
def handler(request: httpx.Request) -> httpx.Response:
seen["url"] = str(request.url)
seen["x-api-key"] = request.headers.get("x-api-key")
seen["authorization"] = request.headers.get("authorization")
seen["anthropic-version"] = request.headers.get("anthropic-version")
seen["body"] = json.loads(request.content)
return anthropic_ok('{"ok": true}')
settings = make_settings()
transport = httpx.MockTransport(handler)
async with httpx.AsyncClient(transport=transport) as http:
client = LLMClient(settings, http=http)
reply = await client.complete_vision(
system="sys", messages=[user_turn("hello", image_png=PNG)])
assert reply.text == '{"ok": true}'
assert reply.style == "anthropic"
assert seen["url"] == "https://llm.test/api/v1/messages"
assert seen["x-api-key"] == "test-key-not-real"
assert seen["authorization"] == "Bearer test-key-not-real"
assert seen["anthropic-version"] == "2023-06-01"
assert seen["body"]["model"] == "test-model"
assert seen["body"]["system"] == "sys"
content = seen["body"]["messages"][0]["content"]
assert content[0]["type"] == "image"
assert content[0]["source"]["media_type"] == "image/png"
assert content[1] == {"type": "text", "text": "hello"}
@pytest.mark.asyncio
async def test_no_retry_on_401(self, make_settings):
calls = []
def handler(request):
calls.append(request)
return httpx.Response(401, json={"error": {"message": "bad key"}})
settings = make_settings(LLM_MAX_RETRIES=3)
transport = httpx.MockTransport(handler)
async with httpx.AsyncClient(transport=transport) as http:
client = LLMClient(settings, http=http)
with pytest.raises(LLMError) as exc:
await client.complete_vision(system="s", messages=[user_turn("hi")])
assert exc.value.code == "llm_rejected"
assert len(calls) == 1 # never retried
assert "test-key-not-real" not in str(exc.value) # key never leaks
@pytest.mark.asyncio
async def test_retry_on_429_then_success(self, make_settings, monkeypatch):
calls = []
async def no_sleep(_):
return None
monkeypatch.setattr("app.llm.asyncio.sleep", no_sleep)
def handler(request):
calls.append(request)
if len(calls) < 3:
return httpx.Response(429, headers={"Retry-After": "0"})
return anthropic_ok("done")
settings = make_settings(LLM_MAX_RETRIES=3)
transport = httpx.MockTransport(handler)
async with httpx.AsyncClient(transport=transport) as http:
client = LLMClient(settings, http=http)
reply = await client.complete_vision(system="s", messages=[user_turn("hi")])
assert reply.text == "done"
assert len(calls) == 3
@pytest.mark.asyncio
async def test_retry_on_non_json_success_then_success(
self, make_settings, monkeypatch
):
calls = []
async def no_sleep(_):
return None
monkeypatch.setattr("app.llm.asyncio.sleep", no_sleep)
def handler(request):
calls.append(request)
if len(calls) < 3:
return httpx.Response(
200,
text="<html>temporary upstream gateway response</html>",
headers={"content-type": "text/html"},
)
return anthropic_ok("done")
settings = make_settings(LLM_MAX_RETRIES=2)
transport = httpx.MockTransport(handler)
async with httpx.AsyncClient(transport=transport) as http:
client = LLMClient(settings, http=http)
reply = await client.complete_vision(
system="s", messages=[user_turn("hi")]
)
assert reply.text == "done"
assert len(calls) == 3
@pytest.mark.asyncio
async def test_truncation_raises(self, make_settings):
def handler(request):
return httpx.Response(200, json={
"content": [{"type": "text", "text": "partial"}],
"stop_reason": "max_tokens",
})
settings = make_settings()
transport = httpx.MockTransport(handler)
async with httpx.AsyncClient(transport=transport) as http:
client = LLMClient(settings, http=http)
with pytest.raises(LLMError) as exc:
await client.complete_vision(system="s", messages=[user_turn("hi")])
assert exc.value.code == "llm_truncated"
class TestOpenAIFallback:
@pytest.mark.asyncio
async def test_404_falls_back_to_chat_completions(self, make_settings):
urls = []
def handler(request: httpx.Request) -> httpx.Response:
urls.append(str(request.url))
if request.url.path.endswith("/v1/messages"):
return httpx.Response(404, json={"error": "not found"})
return openai_ok("fallback-ok")
settings = make_settings() # LLM_API_STYLE default auto
transport = httpx.MockTransport(handler)
async with httpx.AsyncClient(transport=transport) as http:
client = LLMClient(settings, http=http)
reply = await client.complete_vision(
system="sys", messages=[user_turn("hello", image_png=PNG)])
assert reply.text == "fallback-ok"
assert reply.style == "openai"
assert urls == [
"https://llm.test/api/v1/messages",
"https://llm.test/api/v1/chat/completions",
]
@pytest.mark.asyncio
async def test_openai_body_shape(self, make_settings):
seen = {}
def handler(request):
seen["body"] = json.loads(request.content)
return openai_ok("ok")
settings = make_settings(LLM_API_STYLE="openai")
transport = httpx.MockTransport(handler)
async with httpx.AsyncClient(transport=transport) as http:
client = LLMClient(settings, http=http)
await client.complete_vision(
system="SYS", messages=[
user_turn("first", image_png=PNG),
assistant_turn("reply-one"),
user_turn("second"),
])
body = seen["body"]
assert body["messages"][0] == {"role": "system", "content": "SYS"}
user1 = body["messages"][1]
assert user1["role"] == "user"
types = [part["type"] for part in user1["content"]]
assert types == ["image_url", "text"]
image_part = user1["content"][0]
assert image_part["image_url"]["url"].startswith("data:image/png;base64,")
assert user1["content"][1] == {"type": "text", "text": "first"}
assert body["messages"][2] == {"role": "assistant", "content": "reply-one"}
assert body["messages"][3] == {"role": "user", "content": "second"}
class TestJsonExtraction:
def test_plain(self):
assert extract_json_object('{"a": 1}') == {"a": 1}
def test_fenced(self):
assert extract_json_object('```json\n{"a": 2}\n```') == {"a": 2}
def test_prose_wrapped(self):
assert extract_json_object('Here you go:\n{"a": 3}\nHope that helps!') == {"a": 3}
def test_malformed(self):
with pytest.raises(LLMError) as exc:
extract_json_object("no json at all")
assert exc.value.code == "llm_bad_json"
def test_non_object(self):
with pytest.raises(LLMError):
extract_json_object("[1, 2, 3]")
|