okami.mom/albedo

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

dawnwarm: keep a session warm while a remote job it awaits runs8a01b9c

main
branchesmainall refs →
14.1 KiB369 linesraw
1"""A provider's declared image edge holds for every image it is sent.
2
3A profile's `imageEdge` stands in for Claude's 2000px limit, which only the
4real Anthropic upstream declares.
5"""
6
7import base64
8import hashlib
9import json
10import re
11import sqlite3
12import struct
13import unittest
14import urllib.error
15import zlib
16
17from harness import (
18    LOOPBACK_REMOTE,
19    Albedo,
20    Provider,
21    exclusive,
22    operation_id,
23    python,
24    text,
25)
26
27
28def png(width, height, metadata=b""):
29    def chunk(kind, data):
30        body = kind + data
31        return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body))
32
33    rows = b"".join(b"\x00" + b"\x80\x40\x20" * width for _ in range(height))
34    return (
35        b"\x89PNG\r\n\x1a\n"
36        + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0))
37        + (chunk(b"npAD", metadata) if metadata else b"")
38        + chunk(b"IDAT", zlib.compress(rows))
39        + chunk(b"IEND", b"")
40    )
41
42
43WIDE = png(1200, 10)
44WIDE_DATA = base64.b64encode(WIDE).decode()
45NOTE = "An earlier image was scaled from 1200x10 to 1000x"
46
47
48def widths(request):
49    """The pixel width of every PNG the request carries."""
50    found = re.findall(r"data:image/png;base64,([A-Za-z0-9+/=]+)", json.dumps(request))
51    return [struct.unpack(">I", base64.b64decode(data)[16:20])[0] for data in found]
52
53
54def cell_result(output):
55    """The python tool's result: plain text, or the text part beside images."""
56    body = output["output"]
57    if isinstance(body, str):
58        return json.loads(body)
59    [text] = [part["text"] for part in body if part.get("type") == "input_text"]
60    return json.loads(text)
61
62
63class ImageLimitsTest(unittest.TestCase):
64    @exclusive
65    def test_large_valid_image_remains_available_after_durable_reload(self):
66        # A valid ancillary chunk makes encoded length exceed 5 MiB while the
67        # decoded PNG remains below the advertised file-size boundary.
68        image = png(2, 3, b"x" * (4 * 1024 * 1024))
69        data = base64.b64encode(image).decode()
70        self.assertLess(len(image), 5 * 1024 * 1024)
71        self.assertGreater(len(data), 5 * 1024 * 1024)
72        provider = Provider(lambda _: text("image accepted"))
73        self.addCleanup(provider.close)
74        with Albedo(provider, protocol="responses") as app:
75            session = app.session()
76            app.api(
77                f"/sessions/{session}/inputs/{operation_id()}",
78                {
79                    "kind": "message",
80                    "text": "look",
81                    "images": [{"mime_type": "image/png", "data": data}],
82                },
83                method="PUT",
84            ).close()
85            app.idle(session)
86            app.restart()
87            app.prompt(session, "look again").close()
88            app.idle(session)
89            sent = json.dumps(provider.requests[-1]["request"])
90            encoded = re.findall(r"data:image/png;base64,([A-Za-z0-9+/=]+)", sent)
91            self.assertEqual(len(encoded), 1)
92            self.assertEqual(base64.b64decode(encoded[0]), image)
93
94    def test_a_declared_edge_refuses_new_images_and_fits_history(self):
95        def reply(request):
96            inputs = request["input"]
97            if inputs[-1].get("type") == "function_call_output":
98                return text("done")
99            prompts = [
100                item.get("content")
101                for item in inputs
102                if item.get("role") == "user"
103                and "<system-note" not in json.dumps(item.get("content"))
104            ]
105            if prompts[-1] == "show":
106                return python(
107                    "import base64\nshow_image(base64.b64decode('" + WIDE_DATA + "'))"
108                )
109            return text("ok")
110
111        provider = Provider(reply)
112        loose, strict = f"loose-{provider.route}", f"strict-{provider.route}"
113        profile = {
114            "baseUrl": provider.url,
115            "apiKey": "fixture-key",
116            "model": "fixture-model",
117            "protocol": "responses",
118        }
119        try:
120            with Albedo(
121                provider,
122                protocol="responses",
123                providers={loose: profile, strict: dict(profile, imageEdge=1000)},
124            ) as app:
125                session = app.session()
126                attach = {
127                    "kind": "message",
128                    "text": "look",
129                    "images": [{"mime_type": "image/png", "data": WIDE_DATA}],
130                }
131
132                def switch(name):
133                    with app.api(f"/sessions/{session}?view=configuration") as response:
134                        json.load(response)
135                        revision = response.headers["ETag"]
136                    app.api(
137                        f"/sessions/{session}?view=configuration",
138                        {"provider_profile": name, "model": "fixture-model"},
139                        method="PATCH",
140                        headers={"If-Match": revision},
141                    ).close()
142
143                def attach_image():
144                    return app.api(
145                        f"/sessions/{session}/inputs/{operation_id()}",
146                        attach,
147                        method="PUT",
148                    )
149
150                def send(content):
151                    start = len(provider.requests)
152                    app.prompt(session, content).close()
153                    app.idle(session)
154                    return [r["request"] for r in provider.requests[start:]]
155
156                attach_image().close()
157                app.idle(session)
158                self.assertEqual(widths(provider.requests[-1]["request"]), [1200])
159
160                switch(strict)
161                sent = len(provider.requests)
162                with self.assertRaises(urllib.error.HTTPError) as refused:
163                    attach_image().close()
164                self.assertEqual(refused.exception.code, 409)
165                self.assertIn(
166                    "1000px edge limit", json.load(refused.exception)["detail"]
167                )
168                self.assertEqual(len(provider.requests), sent)
169
170                strict_requests = send("show") + send("again")
171                self.assertEqual(len(strict_requests), 3)
172                for request in strict_requests:
173                    self.assertEqual(widths(request), [1000])
174                    self.assertEqual(json.dumps(request).count(NOTE), 1)
175                output = next(
176                    item
177                    for item in strict_requests[1]["input"]
178                    if item.get("type") == "function_call_output"
179                )
180                # show_image refuses it at the call, so the cell fails there.
181                result = cell_result(output)
182                self.assertEqual(result["status"], "error")
183                self.assertIn(
184                    "ValueError: 1200x10 image is over this model's 1000px edge limit",
185                    result["output"],
186                )
187                self.assertNotIn("image_errors", result)
188                entries = app.history(session)["items"]
189                notes = [
190                    part["text"]
191                    for entry in entries
192                    if entry["kind"] == "image_fit"
193                    for part in entry["content"]
194                    if part["kind"] == "text"
195                ]
196                self.assertEqual(len(notes), 1, notes)
197                self.assertIn(NOTE, notes[0])
198                original = [
199                    entry
200                    for entry in entries
201                    if entry["kind"] == "user"
202                    and any(
203                        part["kind"] == "text" and part["text"] == "look"
204                        for part in entry["content"]
205                    )
206                ]
207                self.assertEqual(len(original), 1)
208                [image] = [
209                    part["image"]
210                    for part in original[0]["content"]
211                    if part["kind"] == "image"
212                ]
213                self.assertEqual(image["width"], 1200)
214                self.assertEqual(image["height"], 10)
215                self.assertEqual(image["original_bytes"], len(WIDE))
216
217                source = hashlib.sha256(WIDE_DATA.encode()).hexdigest()
218                database = f"file:{app.home / 'albedo.sqlite'}?mode=ro"
219
220                def rows(session, needle):
221                    with sqlite3.connect(database, uri=True) as db:
222                        return [
223                            seq
224                            for (seq,) in db.execute(
225                                "SELECT seq FROM transcript WHERE session=? "
226                                "AND instr(payload, CAST(? AS BLOB))>0 ORDER BY seq",
227                                (session, needle),
228                            )
229                        ]
230
231                # The original row keeps its image; the fit is a row of its own.
232                self.assertEqual(len(rows(session, source)), 2)
233                self.assertEqual(len(rows(session, "image_fit")), 1)
234
235                # A fit is permanent: a looser provider keeps the copy.
236                switch(loose)
237                [back] = send("back")
238                self.assertEqual(widths(back), [1000])
239                self.assertEqual(json.dumps(back).count(NOTE), 1)
240
241                # A fork from before the fit starts from the original.
242                checkpoint = next(
243                    entry["checkpoint_id"]
244                    for entry in entries
245                    if entry["kind"] == "assistant"
246                )
247                branch = operation_id()
248                app.api(
249                    f"/sessions/{branch}",
250                    {
251                        "kind": "fork",
252                        "source_session_id": session,
253                        "checkpoint_id": checkpoint,
254                    },
255                    method="PUT",
256                    headers={"If-None-Match": "*"},
257                ).close()
258                start = len(provider.requests)
259                app.prompt(branch, "branch").close()
260                app.idle(branch)
261                [forked] = [r["request"] for r in provider.requests[start:]]
262                self.assertEqual(widths(forked), [1200])
263                self.assertNotIn(NOTE, json.dumps(forked))
264        finally:
265            provider.close()
266
267    def test_an_image_cut_short_is_refused_where_it_is_attached(self):
268        # A PNG cut off by a capped read kept a readable header, so it was
269        # attached, and every later request carried an image the provider
270        # could not decode: the session failed until it was forked.
271        cut = png(2, 3, b"x" * 1000)[:500]
272        cut_data = base64.b64encode(cut).decode()
273
274        def reply(request):
275            if request["input"][-1].get("type") == "function_call_output":
276                return text("done")
277            return python(
278                "import base64\nshow_image(base64.b64decode('" + cut_data + "'))"
279            )
280
281        provider = Provider(reply)
282        self.addCleanup(provider.close)
283        with Albedo(provider, protocol="responses") as app:
284            session = app.session()
285            with self.assertRaises(urllib.error.HTTPError) as refused:
286                app.api(
287                    f"/sessions/{session}/inputs/{operation_id()}",
288                    {
289                        "kind": "message",
290                        "text": "look",
291                        "images": [{"mime_type": "image/png", "data": cut_data}],
292                    },
293                    method="PUT",
294                ).close()
295            self.assertEqual(refused.exception.code, 400)
296
297            app.prompt(session, "show").close()
298            app.idle(session)
299            request = provider.requests[-1]["request"]
300            output = next(
301                item
302                for item in request["input"]
303                if item.get("type") == "function_call_output"
304            )
305            result = cell_result(output)
306            self.assertEqual(result["status"], "error")
307            self.assertIn(
308                "ValueError: this 500-byte image/png stops before its IEND chunk",
309                result["output"],
310            )
311            self.assertNotIn("data:image/png", json.dumps(request))
312
313    def test_a_remote_image_is_shown_under_the_same_limits(self):
314        cell = (
315            LOOPBACK_REMOTE
316            + """try:
317    print(await rem.show_image("small.png"))
318    await rem.show_image("wide.png")
319finally:
320    await rem.close()
321"""
322        )
323
324        def reply(request):
325            if request["input"][-1].get("type") == "function_call_output":
326                return text("done")
327            return python(cell)
328
329        provider = Provider(reply)
330        profile = {
331            "baseUrl": provider.url,
332            "apiKey": "fixture-key",
333            "model": "fixture-model",
334            "protocol": "responses",
335            "imageEdge": 1000,
336        }
337        try:
338            with Albedo(
339                provider,
340                protocol="responses",
341                providers={f"strict-{provider.route}": profile},
342            ) as app:
343                (app.workspace / "small.png").write_bytes(png(10, 10))
344                (app.workspace / "wide.png").write_bytes(WIDE)
345                session = app.session()
346                app.prompt(session, "remote").close()
347                app.idle(session)
348                [request] = [
349                    r["request"]
350                    for r in provider.requests
351                    if r["request"]["input"][-1].get("type") == "function_call_output"
352                ]
353                output = next(
354                    item
355                    for item in request["input"]
356                    if item.get("type") == "function_call_output"
357                )
358                result = cell_result(output)
359                # The small image came over the connection and went on to the
360                # model; the wide one failed the cell at its call.
361                self.assertIn("attached image/png", result["output"])
362                self.assertEqual(result["status"], "error")
363                self.assertIn(
364                    "ValueError: 1200x10 image is over this model's 1000px edge limit",
365                    result["output"],
366                )
367                self.assertEqual(widths(request), [10])
368        finally:
369            provider.close()