@wendellresearch-ai/benchwarmers 0.2.1 → 0.2.2
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/AGENTS.md +6 -1
- package/GATEWAY.md +149 -6
- package/README.md +42 -2
- package/USAGE.md +16 -0
- package/VERSION +1 -1
- package/cli.py +9 -1
- package/gateway.py +275 -91
- package/metrics/catalog.py +10 -1
- package/metrics/client.py +537 -127
- package/metrics/cost.py +38 -19
- package/metrics/formats.py +1021 -0
- package/metrics/guard.py +37 -7
- package/metrics/learn.py +19 -6
- package/metrics/ledger.py +148 -0
- package/metrics/model_roles.py +13 -3
- package/metrics/pacing.py +248 -0
- package/metrics/quota.py +383 -0
- package/metrics/route.py +5 -3
- package/package.json +2 -1
- package/perf/run.py +30 -7
- package/scripts/check.py +5 -0
- package/scripts/onboard.py +10 -1
- package/scripts/test_client.py +50 -6
- package/scripts/test_formats.py +549 -0
- package/scripts/test_merged.py +4 -3
- package/scripts/test_quota.py +334 -0
- package/scripts/test_serve_edit.py +10 -5
- package/scripts/test_server_json.py +79 -0
- package/scripts/test_verdict.py +5 -1
- package/scripts/test_vision.py +6 -3
- package/serve.py +104 -36
- package/ui/bench.js +3 -2
|
@@ -0,0 +1,549 @@
|
|
|
1
|
+
"""Other API formats on the gateway: the translators, and every route against fake seats.
|
|
2
|
+
|
|
3
|
+
python scripts/test_formats.py
|
|
4
|
+
|
|
5
|
+
metrics/formats.py both ways, without a server: OpenAI messages to each
|
|
6
|
+
native request and back (text, a system text, an assistant turn, an inline
|
|
7
|
+
and a linked image), and an answer to each native reply and stream and back
|
|
8
|
+
to its text, so the outbound half the built-in client will use is the
|
|
9
|
+
inverse of the inbound half the gateway uses. Then the real gateway handler,
|
|
10
|
+
its seats stubbed (no model is called), on every route: Anthropic Messages
|
|
11
|
+
and count_tokens, Responses, legacy completions, Ollama's chat, generate,
|
|
12
|
+
tags, version and show, Gemini's generateContent, streamGenerateContent (SSE
|
|
13
|
+
and JSON array), models and countTokens. Streamed and not; seat by seat as
|
|
14
|
+
each finishes, a merged reply at the end, keep-alives while the seats think,
|
|
15
|
+
a refusal before the panel keeps its status and an error after the 200 goes
|
|
16
|
+
into the stream in the format's own shape; tools refused; images arrive as
|
|
17
|
+
image_url parts; keys by x-api-key and a Gemini route's ?key=, never logged.
|
|
18
|
+
Temp data folder; BENCHWARMERS_HOME is set before anything is imported.
|
|
19
|
+
"""
|
|
20
|
+
|
|
21
|
+
from __future__ import annotations
|
|
22
|
+
|
|
23
|
+
import http.client
|
|
24
|
+
import io
|
|
25
|
+
import json
|
|
26
|
+
import os
|
|
27
|
+
import pathlib
|
|
28
|
+
import shutil
|
|
29
|
+
import socket
|
|
30
|
+
import sys
|
|
31
|
+
import tempfile
|
|
32
|
+
import threading
|
|
33
|
+
import time
|
|
34
|
+
from http.server import ThreadingHTTPServer
|
|
35
|
+
|
|
36
|
+
ROOT = pathlib.Path(__file__).resolve().parent.parent
|
|
37
|
+
TMP = pathlib.Path(tempfile.mkdtemp(prefix="bench-formats-"))
|
|
38
|
+
(TMP / "config").mkdir()
|
|
39
|
+
(TMP / "llm").mkdir()
|
|
40
|
+
os.environ.update(BENCHWARMERS_HOME=str(TMP), LLM_USER_PATH=str(TMP / "llm"), BENCH_CLIENT="builtin",
|
|
41
|
+
BENCHWARMERS_KEY_FAKE="sk-test")
|
|
42
|
+
for name in ("BENCHWARMERS_ADMIN_TOKEN", "BENCHWARMERS_REQUIRE_KEYS", "BENCH_ALLOWED_ORIGINS"):
|
|
43
|
+
os.environ.pop(name, None)
|
|
44
|
+
sys.path[:0] = [str(ROOT), str(ROOT / "metrics"), str(ROOT / "suite")]
|
|
45
|
+
|
|
46
|
+
import formats as F # noqa: E402
|
|
47
|
+
|
|
48
|
+
PIXEL = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
|
|
49
|
+
DATA_URL = "data:image/png;base64," + PIXEL
|
|
50
|
+
PHOTO = "https://photos.example/chair.png"
|
|
51
|
+
TALK = [{"role": "system", "content": "You price estate-sale items."},
|
|
52
|
+
{"role": "user", "content": [F.text_part("What is this chair worth?"), F.image(DATA_URL)]},
|
|
53
|
+
{"role": "assistant", "content": "About forty dollars."},
|
|
54
|
+
{"role": "user", "content": "And with the cushion?"}]
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def sse_events(raw: str) -> list[tuple[str, str]]:
|
|
58
|
+
"""(event name, data) per server-sent event; ("comment", line) for a comment."""
|
|
59
|
+
out = []
|
|
60
|
+
for block in raw.replace("\r\n", "\n").split("\n\n"):
|
|
61
|
+
name, data = "", []
|
|
62
|
+
for line in block.splitlines():
|
|
63
|
+
if line.startswith(":"):
|
|
64
|
+
out.append(("comment", line))
|
|
65
|
+
elif line.startswith("event: "):
|
|
66
|
+
name = line[7:]
|
|
67
|
+
elif line.startswith("data: "):
|
|
68
|
+
data.append(line[6:])
|
|
69
|
+
if data:
|
|
70
|
+
out.append((name, "\n".join(data)))
|
|
71
|
+
return out
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
def sse_json(raw: str) -> list[dict]:
|
|
75
|
+
return [json.loads(d) for _, d in sse_events(raw) if d.startswith("{")]
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
def refused(fn, *args, said: str = "") -> str:
|
|
79
|
+
try:
|
|
80
|
+
fn(*args)
|
|
81
|
+
except F.Unsupported as exc:
|
|
82
|
+
assert said in str(exc), (said, str(exc))
|
|
83
|
+
return str(exc)
|
|
84
|
+
raise AssertionError(f"{fn.__name__} took {args!r}")
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def translators() -> None:
|
|
88
|
+
# OpenAI messages -> each native request -> back: the same conversation.
|
|
89
|
+
assert F.anthropic_ask(F.anthropic_request(TALK, "bench/core"))["messages"] == TALK
|
|
90
|
+
assert F.gemini_ask(F.gemini_request(TALK), "bench/core")["messages"] == TALK
|
|
91
|
+
assert F.responses_ask(F.responses_request(TALK, "bench/core"))["messages"] == TALK
|
|
92
|
+
assert F.ollama_chat_ask(F.ollama_request(TALK, "bench/core"))["messages"] == TALK
|
|
93
|
+
linked = [{"role": "user", "content": [F.text_part("This one?"), F.image(PHOTO)]}]
|
|
94
|
+
for there_and_back in (lambda: F.anthropic_ask(F.anthropic_request(linked, "m")),
|
|
95
|
+
lambda: F.gemini_ask(F.gemini_request(linked), "m"),
|
|
96
|
+
lambda: F.responses_ask(F.responses_request(linked, "m"))):
|
|
97
|
+
assert there_and_back()["messages"] == linked
|
|
98
|
+
refused(F.ollama_request, linked, "m", said="base64 only") # Ollama takes no URLs
|
|
99
|
+
# The native shapes, as their APIs document them.
|
|
100
|
+
a = F.anthropic_request(TALK, "claude-x", {"max_tokens": 50, "temperature": 0.2, "stop": ["END"]}, stream=True)
|
|
101
|
+
assert a["system"] == "You price estate-sale items." and a["max_tokens"] == 50 and a["stream"] is True
|
|
102
|
+
assert a["stop_sequences"] == ["END"] and a["messages"][0]["content"][1] == {
|
|
103
|
+
"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PIXEL}}, a
|
|
104
|
+
assert F.anthropic_request(TALK, "m")["max_tokens"] == F.ANTHROPIC_MAX_TOKENS
|
|
105
|
+
g = F.gemini_request(TALK, {"max_tokens": 9})
|
|
106
|
+
assert g["systemInstruction"] == {"parts": [{"text": "You price estate-sale items."}]}
|
|
107
|
+
assert [c["role"] for c in g["contents"]] == ["user", "model", "user"] and g["generationConfig"] == {"maxOutputTokens": 9}
|
|
108
|
+
assert g["contents"][0]["parts"][1] == {"inlineData": {"mimeType": "image/png", "data": PIXEL}}
|
|
109
|
+
assert F.gemini_request(linked)["contents"][0]["parts"][1] == {"fileData": {"mimeType": "image/png", "fileUri": PHOTO}}
|
|
110
|
+
assert F.gemini_path("gemini-x") == "/v1beta/models/gemini-x:generateContent"
|
|
111
|
+
assert F.gemini_path("gemini-x", True).endswith(":streamGenerateContent?alt=sse")
|
|
112
|
+
r = F.responses_request(TALK, "gpt-x", {"max_tokens": 7})
|
|
113
|
+
assert r["store"] is False and r["instructions"] and r["max_output_tokens"] == 7
|
|
114
|
+
assert r["input"][1]["content"] == [{"type": "output_text", "text": "About forty dollars."}]
|
|
115
|
+
o = F.ollama_request(TALK, "llama", {"max_tokens": 5})
|
|
116
|
+
assert o["messages"][1] == {"role": "user", "content": "What is this chair worth?", "images": [PIXEL]}
|
|
117
|
+
assert o["options"] == {"num_predict": 5} and o["stream"] is False
|
|
118
|
+
|
|
119
|
+
# An answer -> each native reply, and each stream -> back to its text.
|
|
120
|
+
text = "Seats agree: ## heading, `code`, ünïcode."
|
|
121
|
+
pieces = ["Seats agree: ", "## heading, ", "`code`, ünïcode."]
|
|
122
|
+
for reply, read in ((F.anthropic_reply, F.anthropic_text), (F.gemini_reply, F.gemini_text),
|
|
123
|
+
(F.responses_reply, F.responses_text), (F.ollama_reply, F.ollama_text),
|
|
124
|
+
(F.ollama_generate_reply, F.ollama_text)):
|
|
125
|
+
assert read(json.loads(json.dumps(reply(text, "bench/core", {"panel": "p"}, (3, 4), 1.5)))) == text, reply
|
|
126
|
+
for stream, delta in ((F.AnthropicStream, F.anthropic_delta), (F.ResponsesStream, F.responses_delta),
|
|
127
|
+
(F.GeminiStream, F.gemini_delta)):
|
|
128
|
+
s = stream("bench/core", 12)
|
|
129
|
+
raw = (s.start() + s.keepalive() + b"".join(s.text(p) for p in pieces) + s.keepalive() + s.end({"panel": "p"}))
|
|
130
|
+
assert "".join(delta(e) for e in sse_json(raw.decode())) == text, stream
|
|
131
|
+
s = F.GeminiStream("bench/core", 12, sse=False) # without ?alt=sse: one JSON array
|
|
132
|
+
raw = s.start() + s.keepalive() + b"".join(s.text(p) for p in pieces) + s.end()
|
|
133
|
+
chunks = json.loads(raw)
|
|
134
|
+
assert "".join(F.gemini_delta(c) for c in chunks) == text and chunks[-1]["candidates"][0]["finishReason"] == "STOP"
|
|
135
|
+
assert chunks[-1]["usageMetadata"]["promptTokenCount"] == 12
|
|
136
|
+
s = F.OllamaStream("bench/core", 3)
|
|
137
|
+
lines = (s.start() + s.keepalive() + b"".join(s.text(p) for p in pieces) + s.end({"panel": "p"})).decode().splitlines()
|
|
138
|
+
rows = [json.loads(line) for line in lines]
|
|
139
|
+
assert "".join(F.ollama_delta(r) for r in rows) == text and rows[-1]["done"] and rows[-1]["bench"] == {"panel": "p"}
|
|
140
|
+
assert rows[-1]["eval_duration"] > 0 and rows[-1]["prompt_eval_count"] == 3
|
|
141
|
+
for s in (F.ChatStream("bench/core"), F.CompletionStream("bench/core")):
|
|
142
|
+
raw = (s.start() + b"".join(s.text(p) for p in pieces) + s.end({"panel": "p"})).decode()
|
|
143
|
+
got = sse_json(raw)
|
|
144
|
+
said = "".join(c["choices"][0].get("text") or c["choices"][0].get("delta", {}).get("content") or "" for c in got)
|
|
145
|
+
assert said == text and got[-1]["bench"] == {"panel": "p"} and raw.endswith("data: [DONE]\n\n"), raw
|
|
146
|
+
# An error in a stream, in each shape, is raised by the outbound reader.
|
|
147
|
+
for stream, delta in ((F.AnthropicStream, F.anthropic_delta), (F.ResponsesStream, F.responses_delta),
|
|
148
|
+
(F.GeminiStream, F.gemini_delta), (F.OllamaStream, F.ollama_delta)):
|
|
149
|
+
raw = stream("m").error(503, "no seat answered").decode()
|
|
150
|
+
events = sse_json(raw) if stream is not F.OllamaStream else [json.loads(raw)]
|
|
151
|
+
try:
|
|
152
|
+
"".join(delta(e) for e in events)
|
|
153
|
+
raise AssertionError(f"{stream.__name__}'s error read as text")
|
|
154
|
+
except F.UpstreamError as exc:
|
|
155
|
+
assert "no seat answered" in str(exc), exc
|
|
156
|
+
assert sse_json(F.ChatStream("m").error(500, "boom").decode())[0]["error"]["message"] == "boom"
|
|
157
|
+
|
|
158
|
+
# Inbound: each format's own request shapes.
|
|
159
|
+
a = F.anthropic_ask({"model": "bench/core", "max_tokens": 9, "stream": True, "system": [{"type": "text", "text": "S"}],
|
|
160
|
+
"messages": [{"role": "user", "content": [
|
|
161
|
+
{"type": "text", "text": "look", "cache_control": {"type": "ephemeral"}},
|
|
162
|
+
{"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": "/9j/AA"}},
|
|
163
|
+
{"type": "image", "source": {"type": "url", "url": PHOTO}}]},
|
|
164
|
+
{"role": "assistant", "content": [{"type": "thinking", "thinking": "hm"},
|
|
165
|
+
{"type": "text", "text": "yes"}]}],
|
|
166
|
+
"bench": {"reply": "merged"}, "with": ["x"]})
|
|
167
|
+
assert a["messages"] == [{"role": "system", "content": "S"},
|
|
168
|
+
{"role": "user", "content": [F.text_part("look"), F.image("data:image/jpeg;base64,/9j/AA"),
|
|
169
|
+
F.image(PHOTO)]},
|
|
170
|
+
{"role": "assistant", "content": "yes"}], a
|
|
171
|
+
assert (a["stream"], a["reply"], a["with"], a["options"]) == (True, "merged", ["x"], {"max_tokens": 9})
|
|
172
|
+
g = F.gemini_ask({"system_instruction": {"parts": [{"text": "S"}]}, "contents": [
|
|
173
|
+
{"role": "user", "parts": [{"text": "look"}, {"inline_data": {"mime_type": "image/png", "data": PIXEL}}]},
|
|
174
|
+
{"role": "model", "parts": [{"text": "thinking...", "thought": True}, {"text": "yes"}]}],
|
|
175
|
+
"generationConfig": {"maxOutputTokens": 9, "stopSequences": ["X"]}}, "bench/core:merged", True)
|
|
176
|
+
assert g["messages"] == [{"role": "system", "content": "S"}, {"role": "user", "content": [F.text_part("look"),
|
|
177
|
+
F.image(DATA_URL)]}, {"role": "assistant", "content": "yes"}], g
|
|
178
|
+
assert (g["model"], g["stream"], g["options"]) == ("bench/core:merged", True, {"max_tokens": 9, "stop": ["X"]})
|
|
179
|
+
r = F.responses_ask({"model": "bench/core", "instructions": "S", "input": [
|
|
180
|
+
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "look"},
|
|
181
|
+
{"type": "input_image", "image_url": DATA_URL, "detail": "low"}]},
|
|
182
|
+
{"type": "reasoning", "id": "rs_1", "summary": []},
|
|
183
|
+
{"role": "assistant", "content": "yes"}], "store": True})
|
|
184
|
+
assert r["messages"] == [{"role": "system", "content": "S"}, {"role": "user", "content": [F.text_part("look"),
|
|
185
|
+
F.image(DATA_URL)]}, {"role": "assistant", "content": "yes"}], r
|
|
186
|
+
assert F.responses_ask({"model": "m", "input": "hi"})["messages"] == [{"role": "user", "content": "hi"}]
|
|
187
|
+
o = F.ollama_chat_ask({"model": "bench/core:latest", "messages": [{"role": "user", "content": "look", "images": [PIXEL]}],
|
|
188
|
+
"options": {"num_predict": 4}})
|
|
189
|
+
assert o["model"] == "bench/core" and o["stream"] is True, o # Ollama streams unless told not to
|
|
190
|
+
assert o["messages"] == [{"role": "user", "content": [F.text_part("look"), F.image(DATA_URL)]}]
|
|
191
|
+
assert F.ollama_model("bench/core:merged") == "bench/core:merged"
|
|
192
|
+
o = F.ollama_generate_ask({"model": "bench/core", "prompt": "hi", "system": "S", "stream": False})
|
|
193
|
+
assert o["messages"] == [{"role": "system", "content": "S"}, {"role": "user", "content": "hi"}] and not o["stream"]
|
|
194
|
+
c = F.completion_ask({"model": "bench/core", "prompt": ["hi"], "stop": "END"})
|
|
195
|
+
assert c["messages"] == [{"role": "user", "content": "hi"}] and c["options"] == {"stop": ["END"]}
|
|
196
|
+
assert F.sniff("/9j/4AAQ") == "image/jpeg" and F.sniff(PIXEL) == "image/png" and F.sniff("R0lGODlh") == "image/gif"
|
|
197
|
+
assert F.gemini_route("/v1beta/models/bench/core:merged:streamGenerateContent") == \
|
|
198
|
+
("bench/core:merged", "streamGenerateContent")
|
|
199
|
+
assert F.gemini_route("/v1beta/models/bench%2Fcore:generateContent") == ("bench/core", "generateContent")
|
|
200
|
+
assert F.gemini_route("/v1beta/models") is None and F.gemini_route("/v1/models") is None
|
|
201
|
+
assert [F.format_of(p) for p in ("/v1/messages", "/api/tags", "/v1beta/models", "/v1/models/x:generateContent",
|
|
202
|
+
"/v1/models", "/v1/responses", "/v1/completions")] == \
|
|
203
|
+
["anthropic", "ollama_chat", "gemini", "gemini", "chat", "responses", "completion"]
|
|
204
|
+
|
|
205
|
+
# Refused, not dropped: tools, tool turns, a stored response, what no seat could read.
|
|
206
|
+
tool = [{"name": "f", "input_schema": {}}]
|
|
207
|
+
user = [{"role": "user", "content": "hi"}]
|
|
208
|
+
for fn, body in ((F.chat_ask, {"messages": user, "tools": [{"type": "function"}]}),
|
|
209
|
+
(F.chat_ask, {"messages": user + [{"role": "tool", "content": "x"}]}),
|
|
210
|
+
(F.anthropic_ask, {"messages": user, "tools": tool}),
|
|
211
|
+
(F.anthropic_ask, {"messages": [{"role": "user", "content": [{"type": "tool_result"}]}]}),
|
|
212
|
+
(F.responses_ask, {"input": "hi", "tools": [{"type": "web_search"}]}),
|
|
213
|
+
(F.responses_ask, {"input": [{"type": "function_call_output", "output": "x"}]}),
|
|
214
|
+
(F.ollama_chat_ask, {"messages": user, "tools": [{}]}),
|
|
215
|
+
(F.ollama_chat_ask, {"messages": [{"role": "assistant", "tool_calls": [{}]}]})):
|
|
216
|
+
refused(fn, body, said="route mode")
|
|
217
|
+
refused(F.gemini_ask, {"contents": [], "tools": [{"functionDeclarations": []}]}, "m", said="route mode")
|
|
218
|
+
refused(F.gemini_ask, {"contents": [{"parts": [{"functionCall": {}}]}]}, "m", said="route mode")
|
|
219
|
+
refused(F.responses_ask, {"input": "hi", "previous_response_id": "resp_1"}, said="never stored")
|
|
220
|
+
refused(F.gemini_ask, {"contents": [{"parts": [{"inlineData": {"mimeType": "application/pdf", "data": "x"}}]}]},
|
|
221
|
+
"m", said="only images")
|
|
222
|
+
refused(F.gemini_ask, {"contents": [{"parts": [{"fileData": {"mimeType": "image/png",
|
|
223
|
+
"fileUri": "files/abc"}}]}]}, "m", said="Files API")
|
|
224
|
+
refused(F.anthropic_ask, {"messages": [{"role": "user", "content": [{"type": "document"}]}]}, said="'document'")
|
|
225
|
+
refused(F.anthropic_ask, {"messages": [{"role": "system", "content": "x"}]}, said="role")
|
|
226
|
+
refused(F.ollama_generate_ask, {"prompt": "hi", "context": [1, 2]}, said="/api/chat")
|
|
227
|
+
refused(F.completion_ask, {"prompt": ["a", "b"]}, said="one string")
|
|
228
|
+
refused(F.chat_ask, {"messages": user, "with": "x"}, said='"with"')
|
|
229
|
+
refused(F.chat_ask, {"messages": user, "bench": {"reply": "both"}}, said='"bench"')
|
|
230
|
+
refused(F.chat_ask, {"messages": user, "model": 5}, said='"model"')
|
|
231
|
+
refused(F.chat_ask, {"messages": "hi"}, said='"messages"')
|
|
232
|
+
assert F.chat_ask({"messages": user, "tools": []})["messages"] == user # an empty list is no tools
|
|
233
|
+
|
|
234
|
+
|
|
235
|
+
class Seats:
|
|
236
|
+
"""Stands in for gateway.ask_seat: an answer per alias, after a delay;
|
|
237
|
+
"boom" raises, as a seat failing outside its own error handling would."""
|
|
238
|
+
|
|
239
|
+
def __init__(self):
|
|
240
|
+
self.asked: list[tuple] = []
|
|
241
|
+
self.delay = {"m1": 0.2}
|
|
242
|
+
|
|
243
|
+
def __call__(self, alias, system, prompt, messages=None):
|
|
244
|
+
self.asked.append((alias, system, prompt, messages))
|
|
245
|
+
time.sleep(self.delay.get(alias, 0.0))
|
|
246
|
+
if alias == "boom":
|
|
247
|
+
raise RuntimeError("the seat's process died")
|
|
248
|
+
answer = "the merged answer" if alias == "mg" else f"answer from {alias}"
|
|
249
|
+
return {"status": "ok", "answer": answer, "error": "", "seconds": 0.0}
|
|
250
|
+
|
|
251
|
+
|
|
252
|
+
def write(name: str, data) -> None:
|
|
253
|
+
(TMP / "config" / name).write_text(data if isinstance(data, str) else json.dumps(data), encoding="utf-8")
|
|
254
|
+
|
|
255
|
+
|
|
256
|
+
def gateway_routes() -> None:
|
|
257
|
+
write("endpoints.json", {"fake": {"base": "http://127.0.0.1:9/v1", "key": "fake", "default": True,
|
|
258
|
+
"kind": "provider"}})
|
|
259
|
+
import endpoints
|
|
260
|
+
write("extra-openai-models.yaml", "".join(endpoints.yaml_block(a, a) for a in ("m1", "m2", "mg", "boom", "tiny")))
|
|
261
|
+
write("model-overrides.json", {"m1": {"vision": True}, "m2": {"vision": True},
|
|
262
|
+
"tiny": {"vision": True, "context_window": 50}})
|
|
263
|
+
write("gateway.json", {"max_calls": 3, "max_explorers": 0}) # no explorers: no router, no shape picker
|
|
264
|
+
write("collections.json", {"core": {"members": ["m1", "m2"], "description": "general"},
|
|
265
|
+
"merge": {"members": ["m1", "m2"], "reply": "merged", "merger": "mg", "description": "x"},
|
|
266
|
+
"broken": {"members": ["m1", "boom"], "description": "x"},
|
|
267
|
+
"small": {"members": ["tiny"], "description": "x"},
|
|
268
|
+
"other": {"members": ["m2"], "description": "x"}})
|
|
269
|
+
import gateway
|
|
270
|
+
import guard
|
|
271
|
+
import serve
|
|
272
|
+
seats = Seats()
|
|
273
|
+
gateway.ask_seat = seats
|
|
274
|
+
gateway.STREAM_GRACE, gateway.KEEPALIVE = 0.1, 0.15
|
|
275
|
+
gw = ThreadingHTTPServer(("127.0.0.1", 0), gateway.Handler)
|
|
276
|
+
dash = ThreadingHTTPServer(("127.0.0.1", 0), serve.Handler)
|
|
277
|
+
serve.PORT = dash.server_port # the dashboard's own Host check uses it
|
|
278
|
+
for s in (gw, dash):
|
|
279
|
+
threading.Thread(target=s.serve_forever, daemon=True).start()
|
|
280
|
+
port = gw.server_port
|
|
281
|
+
|
|
282
|
+
def call(path, body=None, method=None, on=port, **headers):
|
|
283
|
+
"""(status, headers, text, seconds to the status line)."""
|
|
284
|
+
conn = http.client.HTTPConnection("127.0.0.1", on, timeout=30)
|
|
285
|
+
data = json.dumps(body).encode() if body is not None else None
|
|
286
|
+
sent = {"Host": f"127.0.0.1:{on}", **({"Content-Type": "application/json"} if data else {}), **headers}
|
|
287
|
+
started = time.monotonic()
|
|
288
|
+
conn.request(method or ("POST" if data is not None else "GET"), path, data, sent)
|
|
289
|
+
r = conn.getresponse()
|
|
290
|
+
first = time.monotonic() - started
|
|
291
|
+
text = r.read().decode()
|
|
292
|
+
conn.close()
|
|
293
|
+
return r.status, dict(r.getheaders()), text, first
|
|
294
|
+
|
|
295
|
+
def ok(path, body=None, **headers):
|
|
296
|
+
status, h, text, _ = call(path, body, **headers)
|
|
297
|
+
assert status == 200, (path, status, text)
|
|
298
|
+
return text, h
|
|
299
|
+
|
|
300
|
+
user = "say hi"
|
|
301
|
+
saved_err = sys.stderr
|
|
302
|
+
try:
|
|
303
|
+
# NOT STREAMED: each format's reply carries every seat in seat order; the bench where it fits.
|
|
304
|
+
text, _ = ok("/v1/chat/completions", {"model": "bench/core", "messages": [{"role": "user", "content": user}]})
|
|
305
|
+
chat = json.loads(text)
|
|
306
|
+
content, bench = chat["choices"][0]["message"]["content"], chat["bench"]
|
|
307
|
+
assert set(chat) == {"id", "object", "created", "model", "choices", "bench"} and chat["model"] == "bench/core"
|
|
308
|
+
assert content.index("answer from m1") < content.index("answer from m2"), content # seat order
|
|
309
|
+
header, footer = content.split("\n## ")[0], "---\n" + content.split("\n---\n", 1)[1]
|
|
310
|
+
assert content == "\n".join([header] + [gateway.section(s) for s in bench["seats"]] + [footer]), content
|
|
311
|
+
for path, body, read in (
|
|
312
|
+
("/v1/messages", {"model": "bench/core", "max_tokens": 64, "messages": [{"role": "user", "content": user}]},
|
|
313
|
+
F.anthropic_text),
|
|
314
|
+
("/v1/responses", {"model": "bench/core", "input": user}, F.responses_text),
|
|
315
|
+
("/v1/completions", {"model": "bench/core", "prompt": user}, lambda r: r["choices"][0]["text"]),
|
|
316
|
+
("/api/chat", {"model": "bench/core:latest", "stream": False,
|
|
317
|
+
"messages": [{"role": "user", "content": user}]}, F.ollama_text),
|
|
318
|
+
("/api/generate", {"model": "bench/core", "prompt": user, "stream": False}, F.ollama_text),
|
|
319
|
+
("/v1beta/models/bench/core:generateContent", {"contents": [{"parts": [{"text": user}]}]},
|
|
320
|
+
F.gemini_text)):
|
|
321
|
+
text, headers = ok(path, body)
|
|
322
|
+
reply = json.loads(text)
|
|
323
|
+
said = read(reply)
|
|
324
|
+
assert "answer from m1" in said and "answer from m2" in said, (path, said)
|
|
325
|
+
if path.startswith("/v1beta"):
|
|
326
|
+
# Gemini's body has no room for the bench: the call ids are in headers.
|
|
327
|
+
assert "bench" not in reply and headers["X-Bench-Panel"].startswith("p"), headers
|
|
328
|
+
calls = json.loads(headers["X-Bench-Seats"])
|
|
329
|
+
assert [s["alias"] for s in calls] == ["m1", "m2"] and all(s["call"] in said for s in calls), calls
|
|
330
|
+
else:
|
|
331
|
+
assert [s["alias"] for s in reply["bench"]["seats"]] == ["m1", "m2"], (path, reply)
|
|
332
|
+
reply = json.loads(ok("/v1/messages", {"model": "bench/core", "max_tokens": 8,
|
|
333
|
+
"messages": [{"role": "user", "content": user}]})[0])
|
|
334
|
+
assert reply["type"] == "message" and reply["usage"]["input_tokens"] > 0 and reply["usage"]["output_tokens"] > 0
|
|
335
|
+
reply = json.loads(ok("/api/chat", {"model": "bench/core", "stream": False, "messages": [{"role": "user", "content": user}]})[0])
|
|
336
|
+
assert reply["done"] and reply["total_duration"] >= 2e8 and reply["model"] == "bench/core", reply # m1 took 0.2 s
|
|
337
|
+
|
|
338
|
+
# Images in each format's shape reach the seats as OpenAI image_url parts.
|
|
339
|
+
seats.asked.clear()
|
|
340
|
+
ok("/v1/messages", {"model": "bench/core", "max_tokens": 8, "messages": [{"role": "user", "content": [
|
|
341
|
+
{"type": "text", "text": "what is it"},
|
|
342
|
+
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PIXEL}}]}]})
|
|
343
|
+
ok("/api/chat", {"model": "bench/core", "stream": False,
|
|
344
|
+
"messages": [{"role": "user", "content": "what is it", "images": [PIXEL]}]})
|
|
345
|
+
ok("/v1beta/models/bench/core:generateContent", {"contents": [{"parts": [
|
|
346
|
+
{"text": "what is it"}, {"inlineData": {"mimeType": "image/png", "data": PIXEL}}]}]})
|
|
347
|
+
ok("/v1/responses", {"model": "bench/core", "input": [{"role": "user", "content": [
|
|
348
|
+
{"type": "input_text", "text": "what is it"}, {"type": "input_image", "image_url": DATA_URL}]}]})
|
|
349
|
+
sent = [m for _, _, _, m in seats.asked]
|
|
350
|
+
assert len(sent) == 8 and all(m == [{"role": "user", "content": [F.text_part("what is it"), F.image(DATA_URL)]}]
|
|
351
|
+
for m in sent), sent
|
|
352
|
+
# The merged reply, by tag and by mode.
|
|
353
|
+
for path, body in (("/api/chat", {"model": "bench/core:merged", "stream": False,
|
|
354
|
+
"messages": [{"role": "user", "content": user}]}),
|
|
355
|
+
("/v1/messages", {"model": "bench/merge", "max_tokens": 8,
|
|
356
|
+
"messages": [{"role": "user", "content": user}]})):
|
|
357
|
+
reply = json.loads(ok(path, body)[0])
|
|
358
|
+
text = F.ollama_text(reply) if path == "/api/chat" else F.anthropic_text(reply)
|
|
359
|
+
merger = reply["bench"]["merger"] # bench/core has none set: m1
|
|
360
|
+
assert reply["bench"]["reply"] == "merged" and text.startswith(merger["answer"]), reply
|
|
361
|
+
assert merger["call"] in text and merger["alias"] == ("mg" if body["model"] == "bench/merge" else "m1"), merger
|
|
362
|
+
|
|
363
|
+
seats.delay["m1"] = 0.8
|
|
364
|
+
# STREAMED: the 200 at once, every seat as it finishes (m2 before the slower m1), the bench at the end.
|
|
365
|
+
for path, body, parse in (
|
|
366
|
+
("/v1/chat/completions", {"model": "bench/core", "stream": True,
|
|
367
|
+
"messages": [{"role": "user", "content": user}]},
|
|
368
|
+
lambda raw: [c["choices"][0]["delta"].get("content", "") for c in sse_json(raw)]),
|
|
369
|
+
("/v1/completions", {"model": "bench/core", "stream": True, "prompt": user},
|
|
370
|
+
lambda raw: [c["choices"][0]["text"] for c in sse_json(raw)]),
|
|
371
|
+
("/v1/messages", {"model": "bench/core", "stream": True, "max_tokens": 8,
|
|
372
|
+
"messages": [{"role": "user", "content": user}]},
|
|
373
|
+
lambda raw: [F.anthropic_delta(e) for e in sse_json(raw)]),
|
|
374
|
+
("/v1/responses", {"model": "bench/core", "stream": True, "input": user},
|
|
375
|
+
lambda raw: [F.responses_delta(e) for e in sse_json(raw)]),
|
|
376
|
+
("/api/chat", {"model": "bench/core", "messages": [{"role": "user", "content": user}]},
|
|
377
|
+
lambda raw: [F.ollama_delta(json.loads(line)) for line in raw.splitlines()]),
|
|
378
|
+
("/api/generate", {"model": "bench/core", "prompt": user},
|
|
379
|
+
lambda raw: [F.ollama_delta(json.loads(line)) for line in raw.splitlines()]),
|
|
380
|
+
("/v1beta/models/bench/core:streamGenerateContent?alt=sse", {"contents": [{"parts": [{"text": user}]}]},
|
|
381
|
+
lambda raw: [F.gemini_delta(e) for e in sse_json(raw)]),
|
|
382
|
+
("/v1beta/models/bench/core:streamGenerateContent", {"contents": [{"parts": [{"text": user}]}]},
|
|
383
|
+
lambda raw: [F.gemini_delta(c) for c in json.loads(raw)])):
|
|
384
|
+
status, headers, raw, first = call(path, body)
|
|
385
|
+
assert status == 200 and first < 0.6, (path, status, first, raw) # before m1 has answered
|
|
386
|
+
said = "".join(parse(raw))
|
|
387
|
+
assert said.index("answer from m2") < said.index("answer from m1") and said.startswith("Bench bench/core"), (path, said)
|
|
388
|
+
assert "Record a verdict" in said, (path, said)
|
|
389
|
+
if "ndjson" in headers["Content-Type"]:
|
|
390
|
+
rows = [json.loads(line) for line in raw.splitlines()]
|
|
391
|
+
assert any(not r["done"] and not F.ollama_delta(r) for r in rows), raw # an empty line kept it alive
|
|
392
|
+
assert rows[-1]["done"] and rows[-1]["bench"]["seats"], raw
|
|
393
|
+
elif path.startswith("/v1beta"):
|
|
394
|
+
assert headers["Content-Type"] == ("text/event-stream" if "alt=sse" in path else "application/json")
|
|
395
|
+
else:
|
|
396
|
+
assert headers["Content-Type"] == "text/event-stream" and ("comment" in {e for e, _ in sse_events(raw)}
|
|
397
|
+
or "event: ping" in raw), (path, raw)
|
|
398
|
+
if path == "/v1/messages":
|
|
399
|
+
kinds = [e["type"] for e in sse_json(raw)]
|
|
400
|
+
assert kinds[:2] == ["message_start", "content_block_start"] and kinds[-1] == "message_stop", kinds
|
|
401
|
+
assert "ping" in kinds and next(e for e in sse_json(raw) if e["type"] == "message_delta")["bench"]["seats"]
|
|
402
|
+
if path == "/v1/responses":
|
|
403
|
+
done = sse_json(raw)[-1]
|
|
404
|
+
assert done["type"] == "response.completed" and done["response"]["bench"]["seats"], done
|
|
405
|
+
assert F.responses_text(done["response"]) == said
|
|
406
|
+
assert [e["sequence_number"] for e in sse_json(raw)] == list(range(len(sse_json(raw))))
|
|
407
|
+
# Merged and streamed: keep-alives, then the merged answer once, then the end.
|
|
408
|
+
status, _, raw, first = call("/v1/messages", {"model": "bench/merge", "stream": True, "max_tokens": 8,
|
|
409
|
+
"messages": [{"role": "user", "content": user}]})
|
|
410
|
+
texts = [t for t in (F.anthropic_delta(e) for e in sse_json(raw)) if t]
|
|
411
|
+
assert status == 200 and first < 0.6 and len(texts) == 1 and texts[0].startswith("the merged answer"), raw
|
|
412
|
+
assert sse_json(raw)[1]["type"] == "content_block_start" and "event: ping" in raw
|
|
413
|
+
|
|
414
|
+
seats.delay["m1"] = 0.2
|
|
415
|
+
# A refusal before the panel is seated keeps its status, in the format's shape, streamed or not.
|
|
416
|
+
for path, body, shape in (
|
|
417
|
+
("/v1/messages", {"model": "bench/small", "stream": True, "max_tokens": 8,
|
|
418
|
+
"messages": [{"role": "user", "content": "x" * 400}]}, lambda e: e["error"]["type"]),
|
|
419
|
+
("/api/chat", {"model": "bench/small", "messages": [{"role": "user", "content": "x" * 400}]},
|
|
420
|
+
lambda e: e["error"]),
|
|
421
|
+
("/v1beta/models/bench/small:streamGenerateContent?alt=sse",
|
|
422
|
+
{"contents": [{"parts": [{"text": "x" * 400}]}]}, lambda e: e["error"]["status"]),
|
|
423
|
+
("/v1/responses", {"model": "bench/small", "stream": True, "input": "x" * 400},
|
|
424
|
+
lambda e: e["error"]["message"])):
|
|
425
|
+
status, _, text, _ = call(path, body)
|
|
426
|
+
assert status == 413 and shape(json.loads(text)), (path, status, text)
|
|
427
|
+
# An error after the 200 goes into the stream, in the format's error shape: a seat that
|
|
428
|
+
# raises once the heading is out (bench/broken seats m1 and "boom").
|
|
429
|
+
seats.delay["boom"] = 0.3
|
|
430
|
+
for path, body, find in (
|
|
431
|
+
("/v1/chat/completions", {"model": "bench/broken", "stream": True,
|
|
432
|
+
"messages": [{"role": "user", "content": user}]},
|
|
433
|
+
lambda raw: sse_json(raw)[-1]["error"]["message"]),
|
|
434
|
+
("/v1/messages", {"model": "bench/broken", "stream": True, "max_tokens": 8,
|
|
435
|
+
"messages": [{"role": "user", "content": user}]},
|
|
436
|
+
lambda raw: sse_json(raw)[-1]["error"]["message"]),
|
|
437
|
+
("/v1/responses", {"model": "bench/broken", "stream": True, "input": user},
|
|
438
|
+
lambda raw: sse_json(raw)[-1]["response"]["error"]["message"]),
|
|
439
|
+
("/api/chat", {"model": "bench/broken", "messages": [{"role": "user", "content": user}]},
|
|
440
|
+
lambda raw: json.loads(raw.splitlines()[-1])["error"]),
|
|
441
|
+
("/v1beta/models/bench/broken:streamGenerateContent", {"contents": [{"parts": [{"text": user}]}]},
|
|
442
|
+
lambda raw: json.loads(raw)[-1]["error"]["message"])):
|
|
443
|
+
status, _, raw, _ = call(path, body)
|
|
444
|
+
assert status == 200 and "the seat's process died" in find(raw), (path, raw)
|
|
445
|
+
|
|
446
|
+
# Refused before any seat is asked: tools, an unknown collection, in each format's shape.
|
|
447
|
+
seats.asked.clear()
|
|
448
|
+
for path, body, msg in (
|
|
449
|
+
("/v1/chat/completions", {"model": "bench/core", "messages": [{"role": "user", "content": user}],
|
|
450
|
+
"tools": [{"type": "function", "function": {"name": "f"}}]},
|
|
451
|
+
lambda e: e["error"]["message"]),
|
|
452
|
+
("/v1/messages", {"model": "bench/core", "max_tokens": 8, "tools": [{"name": "f"}],
|
|
453
|
+
"messages": [{"role": "user", "content": user}]}, lambda e: e["error"]["message"]),
|
|
454
|
+
("/v1/responses", {"model": "bench/core", "input": user, "tools": [{"type": "function"}]},
|
|
455
|
+
lambda e: e["error"]["message"]),
|
|
456
|
+
("/api/chat", {"model": "bench/core", "tools": [{}], "messages": [{"role": "user", "content": user}]},
|
|
457
|
+
lambda e: e["error"]),
|
|
458
|
+
("/v1beta/models/bench/core:generateContent", {"contents": [{"parts": [{"text": user}]}],
|
|
459
|
+
"tools": [{"functionDeclarations": [{}]}]},
|
|
460
|
+
lambda e: e["error"]["message"])):
|
|
461
|
+
status, _, text, _ = call(path, body)
|
|
462
|
+
assert status == 400 and "route mode" in msg(json.loads(text)), (path, status, text)
|
|
463
|
+
status, _, text, _ = call("/v1/messages", {"model": "bench/nope", "max_tokens": 8,
|
|
464
|
+
"messages": [{"role": "user", "content": user}]})
|
|
465
|
+
assert status == 404 and json.loads(text)["error"]["type"] == "not_found_error", text
|
|
466
|
+
assert "bench/core" in json.loads(text)["error"]["message"] # the collections, to choose from
|
|
467
|
+
status, _, text, _ = call("/v1beta/models/bench/nope:generateContent", {"contents": [{"parts": [{"text": "x"}]}]})
|
|
468
|
+
assert status == 404 and json.loads(text)["error"]["status"] == "NOT_FOUND", text
|
|
469
|
+
status, _, text, _ = call("/v1beta/models/bench/core:embedContent", {"content": {}})
|
|
470
|
+
assert status == 404 and "generateContent" in text, text
|
|
471
|
+
assert not seats.asked, seats.asked
|
|
472
|
+
|
|
473
|
+
# Listing, in each format's shape: every collection, and each in its other reply mode.
|
|
474
|
+
tags = json.loads(ok("/api/tags")[0])["models"]
|
|
475
|
+
assert {"bench/core", "bench/core:merged", "bench/merge", "bench/merge:all"} <= {m["name"] for m in tags}, tags
|
|
476
|
+
assert json.loads(ok("/api/version")[0]) == {"version": F.OLLAMA_VERSION}
|
|
477
|
+
shown = json.loads(ok("/api/show", {"model": "bench/core"})[0])
|
|
478
|
+
assert "vision" in shown["capabilities"] and shown["details"]["family"] == "benchwarmers", shown
|
|
479
|
+
assert json.loads(ok("/api/show", {"name": "core"})[0])["capabilities"]
|
|
480
|
+
status, _, text, _ = call("/api/show", {"model": "bench/nope"})
|
|
481
|
+
assert status == 404 and "error" in json.loads(text), text
|
|
482
|
+
names = [m["name"] for m in json.loads(ok("/v1beta/models")[0])["models"]]
|
|
483
|
+
assert "models/bench/core" in names and "models/bench/core:merged" in names, names
|
|
484
|
+
one = json.loads(ok("/v1beta/models/bench/core")[0])
|
|
485
|
+
assert one["name"] == "models/bench/core" and "generateContent" in one["supportedGenerationMethods"]
|
|
486
|
+
assert call("/v1beta/models/bench/nope")[0] == 404
|
|
487
|
+
# Token counts are estimates, without a model call; an image counts IMAGE_TOKENS.
|
|
488
|
+
counted = {"model": "bench/core", "messages": [
|
|
489
|
+
{"role": "user", "content": [{"type": "text", "text": "x" * 400},
|
|
490
|
+
{"type": "image", "source": {"type": "base64", "media_type": "image/png",
|
|
491
|
+
"data": PIXEL}}]}]}
|
|
492
|
+
n = json.loads(ok("/v1/messages/count_tokens", counted)[0])["input_tokens"]
|
|
493
|
+
assert n == gateway.prompt_tokens(F.anthropic_ask(counted)["messages"]), n
|
|
494
|
+
assert F.estimate("x" * 400) + gateway.IMAGE_TOKENS <= n < F.estimate("x" * 400) + gateway.IMAGE_TOKENS + 20, n
|
|
495
|
+
n = json.loads(ok("/v1beta/models/bench/core:countTokens", {"contents": [{"parts": [{"text": "x" * 400}]}]})[0])
|
|
496
|
+
assert n == {"totalTokens": F.estimate("x" * 400)}, n
|
|
497
|
+
assert not seats.asked
|
|
498
|
+
|
|
499
|
+
# The dashboard's port answers Ollama's /api/ routes as the gateway's, beside its own.
|
|
500
|
+
text, _ = ok("/api/tags", on=dash.server_port)
|
|
501
|
+
assert "bench/core" in text
|
|
502
|
+
assert call("/api/stamp", on=dash.server_port)[0] == 200
|
|
503
|
+
|
|
504
|
+
# KEYS: x-api-key and x-goog-api-key anywhere, ?key= on Gemini routes only; never logged.
|
|
505
|
+
os.environ["BENCHWARMERS_REQUIRE_KEYS"] = "1"
|
|
506
|
+
key = guard.make_key("sdk")
|
|
507
|
+
sys.stderr = log = io.StringIO()
|
|
508
|
+
ask = {"model": "bench/core", "max_tokens": 8, "messages": [{"role": "user", "content": user}]}
|
|
509
|
+
status, h, text, _ = call("/v1/messages", ask)
|
|
510
|
+
assert status == 401 and json.loads(text)["error"]["type"] == "authentication_error", text
|
|
511
|
+
assert h["WWW-Authenticate"].startswith("Bearer")
|
|
512
|
+
assert call("/v1/messages", ask, **{"x-api-key": key})[0] == 200
|
|
513
|
+
assert call("/v1/messages", ask, **{"x-api-key": "sk-ant-dummy"})[0] == 401
|
|
514
|
+
assert call("/api/tags", Authorization=f"Bearer {key}")[0] == 200
|
|
515
|
+
assert call(f"/v1beta/models?key={key}")[0] == 200
|
|
516
|
+
assert call("/v1beta/models", **{"x-goog-api-key": key})[0] == 200
|
|
517
|
+
assert call(f"/v1beta/models/bench/core:generateContent?alt=sse&key={key}",
|
|
518
|
+
{"contents": [{"parts": [{"text": user}]}]})[0] == 200
|
|
519
|
+
status, _, text, _ = call(f"/v1/models?key={key}") # not in an OpenAI route's URL
|
|
520
|
+
assert status == 401, text
|
|
521
|
+
# A request line too broken to parse is quoted by http.server's own error log: the key is not.
|
|
522
|
+
with socket.create_connection(("127.0.0.1", port), timeout=10) as raw:
|
|
523
|
+
raw.sendall(f"GET /v1beta/models?alt=sse&key={key} HTTP/1.1 extra\r\nHost: 127.0.0.1:{port}\r\n\r\n".encode())
|
|
524
|
+
raw.recv(4096)
|
|
525
|
+
time.sleep(0.2)
|
|
526
|
+
sys.stderr = saved_err
|
|
527
|
+
lines = log.getvalue()
|
|
528
|
+
assert "POST /v1/messages 200 sdk" in lines and "GET /v1beta/models 200 sdk" in lines, lines
|
|
529
|
+
assert "Bad request" in lines and key not in lines and key[4:12] not in lines, lines
|
|
530
|
+
os.environ.pop("BENCHWARMERS_REQUIRE_KEYS")
|
|
531
|
+
finally:
|
|
532
|
+
sys.stderr = saved_err
|
|
533
|
+
os.environ.pop("BENCHWARMERS_REQUIRE_KEYS", None)
|
|
534
|
+
gw.shutdown()
|
|
535
|
+
dash.shutdown()
|
|
536
|
+
|
|
537
|
+
|
|
538
|
+
def main() -> int:
|
|
539
|
+
try:
|
|
540
|
+
translators()
|
|
541
|
+
gateway_routes()
|
|
542
|
+
finally:
|
|
543
|
+
shutil.rmtree(TMP, ignore_errors=True)
|
|
544
|
+
print("ok")
|
|
545
|
+
return 0
|
|
546
|
+
|
|
547
|
+
|
|
548
|
+
if __name__ == "__main__":
|
|
549
|
+
sys.exit(main())
|
package/scripts/test_merged.py
CHANGED
|
@@ -157,11 +157,12 @@ def main() -> int:
|
|
|
157
157
|
assert post({"model": "bench/core", "messages": [{"role": "user", "content": "x"}],
|
|
158
158
|
"bench": {"reply": "both"}})[0] == 400
|
|
159
159
|
assert post({"model": "bench/core:both", "messages": [{"role": "user", "content": "x"}]})[0] == 404
|
|
160
|
-
# Streaming: the merged answer
|
|
160
|
+
# Streaming: the merged answer once the merge is done, the seats on the last chunk.
|
|
161
161
|
chunks = [json.loads(line[6:]) for line in post({"model": "bench/core", "stream": True, "messages":
|
|
162
162
|
[{"role": "user", "content": "x"}]})[1].splitlines() if line.startswith("data: {")]
|
|
163
|
-
|
|
164
|
-
assert
|
|
163
|
+
texts = [c["choices"][0]["delta"].get("content") for c in chunks]
|
|
164
|
+
assert texts[0] == "" and texts[1].startswith("answer from merger") and len(chunks) == 3, texts
|
|
165
|
+
assert len(chunks[-1]["bench"]["seats"]) == 2
|
|
165
166
|
|
|
166
167
|
# 6. /v1/models: each collection and its other reply mode.
|
|
167
168
|
with urllib.request.urlopen(url + "/v1/models", timeout=10) as resp:
|