okami.mom/albedo

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

dawnimages: refuse a PNG or JPEG cut short where it is attached, so it never reaches history79b55b3

main
branchesmainall refs →
12.2 KiB292 linesraw
1"""Image tool outputs reach the model, and compaction reclaims the ones it evicts."""
2
3import base64
4import hashlib
5import json
6import sqlite3
7import unittest
8import uuid
9
10from harness import Albedo, Provider, error, exclusive, operation_id, python, text
11
12
13PNG = "iVBORw0KGgoAAAANSUhEUgAAAAIAAAADCAIAAAA2iEnWAAAAAElFTkSuQmCC"
14SHOWN = "iVBORw0KGgoAAAANSUhEUgAAABEAAAADCAIAAAD6EXM9AAAAAElFTkSuQmCC"
15UPLOADED = "iVBORw0KGgoAAAANSUhEUgAAABIAAAADCAIAAAARJsg+AAAAAElFTkSuQmCC"
16MARKER = "image elided and is no longer available"
17
18
19class ImageToolOutputTest(unittest.TestCase):
20    def test_python_image_is_attached_and_invalid_image_is_reported(self):
21        def reply(request):
22            inputs = request["input"]
23            if inputs[-1].get("type") == "function_call_output":
24                return text("done")
25            user = next(
26                item["content"]
27                for item in reversed(inputs)
28                if item.get("role") == "user"
29            )
30            if user == "show a valid image":
31                return python(
32                    "import base64\nshow_image(base64.b64decode('" + PNG + "'))"
33                )
34            return python("show_image(b'\\x89PNG\\r\\n\\x1a\\nnot a header')")
35
36        def tool_output(provider):
37            inputs = provider.requests[-1]["request"]["input"]
38            return next(
39                item for item in inputs if item.get("type") == "function_call_output"
40            )
41
42        provider = Provider(reply)
43        try:
44            with Albedo(provider, protocol="responses") as app:
45                session = app.session()
46                app.prompt(session, "show a valid image").close()
47                app.idle(session)
48                output = tool_output(provider)
49                self.assertIn("attached image/png", json.dumps(output))
50                request = provider.requests[-1]["request"]
51                self.assertIn(PNG, json.dumps(request))
52
53                second = app.session()
54                app.prompt(second, "show an invalid image").close()
55                app.idle(second)
56                output = tool_output(provider)
57                self.assertIn("image_errors", json.dumps(output))
58                self.assertNotIn(PNG, json.dumps(output))
59        finally:
60            provider.close()
61
62    # exclusive: restarts the daemon
63    @exclusive
64    def test_top_level_and_nested_images_survive_restart(self):
65        code = ""
66
67        def reply(request):
68            if request["input"][-1].get("type") == "function_call_output":
69                return text("done")
70            return python(code) if code else text("history restored")
71
72        provider = Provider(reply)
73        try:
74            with Albedo(provider, protocol="responses") as app:
75                session = app.session()
76
77                def run(source):
78                    nonlocal code
79                    code = source
80                    app.prompt(session, "run image scenario").close()
81                    app.idle(session)
82                    result = next(
83                        json.loads(part["value"])
84                        for entry in reversed(app.history(session)["items"])
85                        if entry["kind"] == "tool_result"
86                        and entry["tool"]["name"] == "python"
87                        for part in entry["content"]
88                        if part["kind"] == "json" and part["field"] == "result"
89                    )
90                    self.assertEqual(result["status"], "ok", result)
91                    return result
92
93                first = run(f"show_image(__import__('base64').b64decode({SHOWN!r}))")
94                run(f"await cells.run({first['cell_id']!r}, allow_partial=True)")
95                attachment = "data:image/png;base64," + SHOWN
96                self.assertEqual(
97                    json.dumps(provider.requests[-1]["request"]).count(attachment), 2
98                )
99                code = ""
100                app.restart()
101                app.prompt(session, "recall delivered images").close()
102                app.idle(session)
103                self.assertEqual(
104                    json.dumps(provider.requests[-1]["request"]).count(attachment), 2
105                )
106        finally:
107            provider.close()
108
109    def test_compaction_elides_evicted_tool_images_only(self):
110        # Unique hashes keep global image reclamation assertions fixture-local.
111        suffix = uuid.uuid4().bytes
112        shown_image = base64.b64encode(base64.b64decode(SHOWN) + suffix).decode()
113        uploaded_image = base64.b64encode(base64.b64decode(UPLOADED) + suffix).decode()
114
115        def reply(request):
116            inputs = request["input"]
117            user = next(
118                (
119                    str(item.get("content", ""))
120                    for item in reversed(inputs)
121                    if item.get("role") == "user"
122                ),
123                "",
124            )
125            if "<newly-evicted-history>" in user:
126                return text("older conversation summary")
127            if inputs[-1].get("type") != "function_call_output" and user == "show":
128                return python(
129                    "show_image(__import__('base64').b64decode('" + shown_image + "'))"
130                )
131            return text("done")
132
133        provider = Provider(reply)
134        try:
135            with Albedo(provider, protocol="responses") as app:
136                shown = hashlib.sha256(shown_image.encode()).hexdigest()
137                uploaded = hashlib.sha256(uploaded_image.encode()).hexdigest()
138                database = f"file:{app.home / 'albedo.sqlite'}?mode=ro"
139
140                def converse(sid, prompts):
141                    for prompt in prompts:
142                        app.prompt(sid, prompt).close()
143                        app.idle(sid)
144                    with app.api(f"/sessions/{sid}/compaction", {}) as response:
145                        self.assertEqual(json.load(response)["state"], "compacted")
146                    app.idle(sid)
147
148                def rows(sid, table, needle):
149                    with sqlite3.connect(database, uri=True) as db:
150                        return db.execute(
151                            f"SELECT count(*) FROM {table} WHERE session=? "
152                            "AND instr(payload, CAST(? AS BLOB))>0",
153                            (sid, needle),
154                        ).fetchone()[0]
155
156                old = app.session()
157                app.api(
158                    f"/sessions/{old}/inputs/{operation_id()}",
159                    {
160                        "kind": "message",
161                        "text": "keep this upload",
162                        "images": [
163                            {
164                                "mime_type": "image/png",
165                                "data": uploaded_image,
166                            }
167                        ],
168                    },
169                    method="PUT",
170                ).close()
171                app.idle(old)
172                converse(old, ["show", "second", "third", "fourth"])
173                self.assertEqual(rows(old, "transcript", shown), 0)
174                self.assertEqual(rows(old, "cells", shown), 0)
175                self.assertEqual(rows(old, "transcript", MARKER), 1)
176                self.assertEqual(rows(old, "cells", MARKER), 1)
177                self.assertEqual(rows(old, "transcript", uploaded), 1)
178                with sqlite3.connect(database, uri=True) as db:
179                    stored = {row[0] for row in db.execute("SELECT hash FROM images")}
180                self.assertNotIn(shown, stored)
181                self.assertIn(uploaded, stored)
182                app.prompt(old, "after compaction").close()
183                app.idle(old)
184                self.assertIn(
185                    "older conversation summary",
186                    json.dumps(provider.requests[-1]["request"]["input"]),
187                )
188
189                recent = app.session()
190                converse(recent, ["first", "second", "third", "show"])
191                self.assertEqual(rows(recent, "transcript", shown), 1)
192                self.assertEqual(rows(recent, "cells", shown), 1)
193                self.assertEqual(rows(recent, "cells", shown_image), 0)
194        finally:
195            provider.close()
196
197    # exclusive: changes global rolling settings and restarts the daemon
198    @exclusive
199    def test_automatic_compaction_elides_evicted_tool_images(self):
200        fail_summary = True
201        attempts = []
202        cleaned_before_request = []
203
204        def reply(request):
205            inputs = request["input"]
206            user = next(
207                (
208                    str(item.get("content", ""))
209                    for item in reversed(inputs)
210                    if item.get("role") == "user"
211                ),
212                "",
213            )
214            if "<newly-evicted-history>" in user:
215                attempts.append(fail_summary)
216                if fail_summary:
217                    return error(400, "summary refused")
218                return text("older conversation summary")
219            if "[older conversation summary;" in json.dumps(inputs):
220                cleaned_before_request.append(
221                    (rows("transcript", MARKER), rows("cells", MARKER))
222                )
223                return error(400, "provider refused after committed compaction")
224            if inputs[-1].get("type") != "function_call_output" and user == "show":
225                return python(
226                    "show_image(__import__('base64').b64decode('" + SHOWN + "'))"
227                )
228            return text("done")
229
230        provider = Provider(reply)
231        try:
232            with Albedo(provider, protocol="responses") as app:
233                settings_path = app.home / "extensions.json"
234                original = settings_path.read_text()
235
236                def restore():
237                    settings_path.write_text(original)
238                    app.restart()
239
240                self.addCleanup(restore)
241                settings = json.loads(original)
242                settings["rolling"] = {"contextWindowTokens": 50000}
243                settings_path.write_text(json.dumps(settings))
244                app.restart()
245                session = app.session()
246                shown = hashlib.sha256(SHOWN.encode()).hexdigest()
247                database = f"file:{app.home / 'albedo.sqlite'}?mode=ro"
248
249                def rows(table, needle):
250                    with sqlite3.connect(database, uri=True) as db:
251                        return db.execute(
252                            f"SELECT count(*) FROM {table} WHERE session=? "
253                            "AND instr(payload, CAST(? AS BLOB))>0",
254                            (session, needle),
255                        ).fetchone()[0]
256
257                app.prompt(session, "show").close()
258                app.idle(session)
259                self.assertEqual(rows("transcript", shown), 1)
260                self.assertEqual(rows("cells", shown), 1)
261                pad = "history " * 8000
262                for prompt in (pad, pad, pad):
263                    app.prompt(session, prompt).close()
264                    app.idle(session)
265                self.assertTrue(
266                    attempts, "the automatic trigger must attempt compaction"
267                )
268                self.assertEqual(rows("transcript", shown), 1)
269                self.assertEqual(rows("cells", shown), 1)
270                fail_summary = False
271                app.prompt(session, "retry compaction").close()
272                app.idle(session)
273                self.assertTrue(
274                    cleaned_before_request, "expected a post-compaction request"
275                )
276                self.assertTrue(all(pair == (1, 1) for pair in cleaned_before_request))
277                self.assertEqual(rows("transcript", shown), 0)
278                self.assertEqual(rows("cells", shown), 0)
279                self.assertEqual(rows("transcript", MARKER), 1)
280                self.assertEqual(rows("cells", MARKER), 1)
281                summaries = len(attempts)
282                app.prompt(session, "after compaction").close()
283                app.idle(session)
284                self.assertEqual(
285                    len(attempts), summaries, "reuse the committed summary"
286                )
287                self.assertIn(
288                    "older conversation summary",
289                    json.dumps(provider.requests[-1]["request"]["input"]),
290                )
291        finally:
292            provider.close()