okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

nekomimi.petgemini: field-reader fast paths for vertex and cloud code assist text chunks, gemini mock routes12f1695

main
branchesmainall refs →
13.1 KiB395 linesraw
1"""A local streaming model server for transport benchmarks.
2
3Run explicitly: python3 test/manual/mock_openai_server.py [port] [cert key]
4The base URL picks the stream: http://127.0.0.1:PORT/<rate>/<tokens>[/<option>...]/v1
5streams <tokens> deltas at <rate> per second (0 for as fast as the socket
6takes them), one SSE event per write, over HTTP/1.1 chunked encoding on
7kept-alive connections. Options: t<percent> streams that share of the
8deltas as one python tool call's arguments after the text (Gemini sends the
9call whole, as one more delta), and s<id> names the stream so parallel
10streams keep separate timings. The route picks the wire format: /responses,
11/chat/completions, Anthropic's /messages, Gemini's /vertex, or Cloud Code
12Assist's /antigravity. With
13a certificate and key it serves https instead. A request to /timings/<id>
14(or /timings for stream 0) returns the send time (ns, CLOCK_REALTIME) of
15every delta of that stream's last response, so a client can compute
16per-token latency.
17"""
18
19import asyncio
20import json
21import random
22import ssl
23import sys
24import time
25from collections.abc import Iterator
26
27WORDS = [
28    "the",
29    " quick",
30    " brown",
31    " fox",
32    " jumps",
33    " over",
34    " lazy",
35    " dog",
36    ".",
37    "\n",
38]
39timings: dict[str, list[int]] = {}
40
41# Each stream yields (is_delta, event); deltas are paced and timed.
42Events = Iterator[tuple[bool, bytes]]
43
44
45def sse(name: str | None, payload: object) -> bytes:
46    data = json.dumps(payload, separators=(",", ":"))
47    head = f"event: {name}\n" if name else ""
48    return f"{head}data: {data}\n\n".encode()
49
50
51def chunked(body: bytes) -> bytes:
52    return b"%x\r\n%s\r\n" % (len(body), body)
53
54
55def split(text: str, parts: int) -> list[str]:
56    """text in at most parts nonempty pieces of about equal size."""
57    if parts <= 0:
58        return []
59    size = -(-len(text) // parts)
60    chunks = [text[i * size : (i + 1) * size] for i in range(parts)]
61    return [chunk for chunk in chunks if chunk] or [text]
62
63
64def plan(count: int, tool_percent: int) -> tuple[list[str], list[str], str]:
65    """count deltas: text words, then tool_percent of them as pieces of one
66    python call's arguments, and those arguments whole."""
67    calls = count * tool_percent // 100
68    words = [WORDS[i % len(WORDS)] for i in range(count - calls)]
69    code = " ".join(f"print({i})" for i in range(max(calls, 1)))
70    arguments = json.dumps({"code": code})
71    return words, split(arguments, calls) if calls else [], arguments
72
73
74def message(text: str, status: str) -> dict:
75    content = [{"type": "output_text", "text": text, "annotations": [], "logprobs": []}]
76    return {
77        "id": "msg_0001",
78        "type": "message",
79        "status": status,
80        "role": "assistant",
81        "content": content,
82    }
83
84
85def function_call(arguments: str, status: str) -> dict:
86    return {
87        "id": "fc_0001",
88        "type": "function_call",
89        "status": status,
90        "call_id": "call_0001",
91        "name": "python",
92        "arguments": arguments,
93    }
94
95
96def responses(words: list[str], calls: list[str], arguments: str) -> Events:
97    sequence = iter(range(1_000_000))
98
99    def event(kind: str, **fields: object) -> bytes:
100        return sse(kind, {"type": kind, "sequence_number": next(sequence), **fields})
101
102    created = {
103        "id": "resp_0001",
104        "object": "response",
105        "status": "in_progress",
106        "model": "mock",
107        "output": [],
108    }
109    yield False, event("response.created", response=created)
110    yield (
111        False,
112        event(
113            "response.output_item.added",
114            output_index=0,
115            item=message("", "in_progress"),
116        ),
117    )
118    for word in words:
119        yield (
120            True,
121            event(
122                "response.output_text.delta",
123                item_id="msg_0001",
124                output_index=0,
125                content_index=0,
126                delta=word,
127                logprobs=[],
128                obfuscation="x7Qp2LmN",
129            ),
130        )
131    output = [message("".join(words), "completed")]
132    yield False, event("response.output_item.done", output_index=0, item=output[0])
133    if calls:
134        yield (
135            False,
136            event(
137                "response.output_item.added",
138                output_index=1,
139                item=function_call("", "in_progress"),
140            ),
141        )
142        for chunk in calls:
143            yield (
144                True,
145                event(
146                    "response.function_call_arguments.delta",
147                    item_id="fc_0001",
148                    output_index=1,
149                    delta=chunk,
150                    obfuscation="x7Qp2LmN",
151                ),
152            )
153        output.append(function_call(arguments, "completed"))
154        yield False, event("response.output_item.done", output_index=1, item=output[1])
155    usage = {
156        "input_tokens": 1200,
157        "input_tokens_details": {"cached_tokens": 1000},
158        "output_tokens": len(words) + len(calls),
159        "output_tokens_details": {"reasoning_tokens": 0},
160    }
161    final = {**created, "status": "completed", "output": output, "usage": usage}
162    yield False, event("response.completed", response=final)
163
164
165def chat(words: list[str], calls: list[str], _arguments: str) -> Events:
166    def chunk(delta: dict, finish: str | None = None) -> bytes:
167        choice = {"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish}
168        return sse(
169            None,
170            {
171                "id": "chatcmpl-0001",
172                "object": "chat.completion.chunk",
173                "created": 1700000000,
174                "model": "mock",
175                "service_tier": "default",
176                "system_fingerprint": "fp_0001",
177                "choices": [choice],
178                "usage": None,
179                "obfuscation": "x7Qp2LmN",
180            },
181        )
182
183    yield False, chunk({"role": "assistant", "content": "", "refusal": None})
184    for word in words:
185        yield True, chunk({"content": word})
186    if calls:
187        opened = {
188            "index": 0,
189            "id": "call_0001",
190            "type": "function",
191            "function": {"name": "python", "arguments": ""},
192        }
193        yield False, chunk({"tool_calls": [opened]})
194        for piece in calls:
195            yield (
196                True,
197                chunk({"tool_calls": [{"index": 0, "function": {"arguments": piece}}]}),
198            )
199    yield False, chunk({}, "tool_calls" if calls else "stop")
200    usage = {
201        "prompt_tokens": 1200,
202        "completion_tokens": len(words) + len(calls),
203        "prompt_tokens_details": {"cached_tokens": 1000},
204    }
205    yield (
206        False,
207        sse(
208            None,
209            {
210                "id": "chatcmpl-0001",
211                "object": "chat.completion.chunk",
212                "choices": [],
213                "usage": usage,
214            },
215        ),
216    )
217    yield False, b"data: [DONE]\n\n"
218
219
220def messages(words: list[str], calls: list[str], _arguments: str) -> Events:
221    def event(kind: str, **fields: object) -> bytes:
222        return sse(kind, {"type": kind, **fields})
223
224    start = {
225        "id": "msg_0001",
226        "type": "message",
227        "role": "assistant",
228        "model": "mock",
229        "content": [],
230        "stop_reason": None,
231        "usage": {
232            "input_tokens": 200,
233            "cache_read_input_tokens": 1000,
234            "output_tokens": 1,
235        },
236    }
237    yield False, event("message_start", message=start)
238    yield (
239        False,
240        event(
241            "content_block_start", index=0, content_block={"type": "text", "text": ""}
242        ),
243    )
244    for word in words:
245        yield (
246            True,
247            event(
248                "content_block_delta",
249                index=0,
250                delta={"type": "text_delta", "text": word},
251            ),
252        )
253    yield False, event("content_block_stop", index=0)
254    if calls:
255        block = {"type": "tool_use", "id": "toolu_0001", "name": "python", "input": {}}
256        yield False, event("content_block_start", index=1, content_block=block)
257        for piece in calls:
258            yield (
259                True,
260                event(
261                    "content_block_delta",
262                    index=1,
263                    delta={"type": "input_json_delta", "partial_json": piece},
264                ),
265            )
266        yield False, event("content_block_stop", index=1)
267    delta = {"stop_reason": "tool_use" if calls else "end_turn", "stop_sequence": None}
268    yield (
269        False,
270        event(
271            "message_delta",
272            delta=delta,
273            usage={"output_tokens": len(words) + len(calls)},
274        ),
275    )
276    yield False, event("message_stop")
277
278
279def gemini(words: list[str], calls: list[str], arguments: str, wrapped: bool) -> Events:
280    def chunk(parts: list[dict], **candidate: object) -> bytes:
281        response = {
282            "candidates": [
283                {"content": {"role": "model", "parts": parts}, "index": 0, **candidate}
284            ],
285            "usageMetadata": {
286                "promptTokenCount": 1200,
287                "totalTokenCount": 1200,
288                "trafficType": "ON_DEMAND",
289            },
290            "modelVersion": "gemini-mock",
291            "responseId": "resp_0001",
292        }
293        return sse(
294            None,
295            {"response": response, "traceId": "trace_0001"} if wrapped else response,
296        )
297
298    for word in words:
299        yield True, chunk([{"text": word}])
300    if calls:
301        call = {"functionCall": {"name": "python", "args": json.loads(arguments)}}
302        yield True, chunk([call])
303    usage = {
304        "promptTokenCount": 1200,
305        "candidatesTokenCount": len(words) + len(calls),
306        "cachedContentTokenCount": 1000,
307        "thoughtsTokenCount": 0,
308    }
309    final = chunk([{"text": ""}], finishReason="STOP")
310    data = json.loads(final.decode().removeprefix("data: "))
311    (data["response"] if wrapped else data)["usageMetadata"] = usage
312    yield False, sse(None, data)
313
314
315def vertex(words: list[str], calls: list[str], arguments: str) -> Events:
316    return gemini(words, calls, arguments, wrapped=False)
317
318
319def antigravity(words: list[str], calls: list[str], arguments: str) -> Events:
320    """Cloud Code Assist wraps each Gemini response in a response field."""
321    return gemini(words, calls, arguments, wrapped=True)
322
323
324async def stream(writer: asyncio.StreamWriter, path: str) -> None:
325    segments = path.split("/")
326    rate, count = float(segments[1]), int(segments[2])
327    options = segments[3 : segments.index("v1")]
328    name = next((option[1:] for option in options if option.startswith("s")), "0")
329    share = int(next((option[1:] for option in options if option.startswith("t")), "0"))
330    routes = {
331        "/chat/completions": chat,
332        "/messages": messages,
333        "/vertex": vertex,
334        "/antigravity": antigravity,
335    }
336    wire = next(
337        (events for route, events in routes.items() if path.endswith(route)), responses
338    )
339    head = "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n"
340    writer.write(head.encode())
341    sent = timings[name] = []
342    interval = 1 / rate if rate > 0 else 0.0
343    # A random phase, so parallel streams do not all send in the same instant.
344    started = time.monotonic() + random.random() * interval
345    for is_delta, event in wire(*plan(count, share)):
346        if is_delta and interval:
347            delay = started + len(sent) * interval - time.monotonic()
348            if delay > 0:
349                await asyncio.sleep(delay)
350        if is_delta:
351            sent.append(time.time_ns())
352        writer.write(chunked(event))
353        await writer.drain()
354    writer.write(b"0\r\n\r\n")
355    await writer.drain()
356
357
358async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
359    try:
360        while request := (await reader.readuntil(b"\r\n")).decode():
361            _, path, _ = request.split(" ", 2)
362            length = 0
363            while (line := await reader.readuntil(b"\r\n")) != b"\r\n":
364                name, _, value = line.decode().partition(":")
365                if name.strip().lower() == "content-length":
366                    length = int(value)
367            await reader.readexactly(length)
368            if path.startswith("/timings"):
369                body = json.dumps(
370                    timings.get(path.removeprefix("/timings").strip("/") or "0", [])
371                ).encode()
372                head = f"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {len(body)}\r\n\r\n"
373                writer.write(head.encode() + body)
374            else:
375                await stream(writer, path)
376            await writer.drain()
377    except (ConnectionError, asyncio.IncompleteReadError):
378        pass
379    finally:
380        writer.close()
381
382
383async def main(port: int, tls: list[str]) -> None:
384    context = None
385    if tls:
386        context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
387        context.load_cert_chain(*tls)
388    server = await asyncio.start_server(handle, "127.0.0.1", port, ssl=context)
389    print(f"listening on {'https' if tls else 'http'}://127.0.0.1:{port}", flush=True)
390    async with server:
391        await server.serve_forever()
392
393
394if __name__ == "__main__":
395    asyncio.run(main(int(sys.argv[1]) if len(sys.argv) > 1 else 8765, sys.argv[2:4]))