okami.mom/albedo

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

Aly Raffaufrefactor(text): use shared scalar truncation331aa22

main
branchesmainall refs →
15.3 KiB359 linesraw
1"""Read-only context inspection against actual model requests and catalog data."""
2
3import json
4import unittest
5import urllib.error
6from urllib.parse import urlencode
7
8from harness import Albedo, Provider, exclusive, Reply, text
9
10
11SECRET = "credential-must-never-appear-in-inspector"
12CATALOG = {
13    "fixture-cloud": {
14        "id": "fixture-cloud",
15        "name": "Fixture Cloud",
16        "env": ["FIXTURE_API_KEY"],
17        "api": "http://127.0.0.1/v1",
18        "models": {
19            "fixture": {
20                "id": "fixture",
21                "limit": {"context": 200000, "output": 8000},
22                "modalities": {"input": ["text"]},
23            }
24        },
25    }
26}
27
28
29# exclusive: replaces the global models catalog and asserts catalog resolution
30@exclusive
31class ContextTest(unittest.TestCase):
32    def setUp(self):
33        def reply(request):
34            if "without provider usage" in json.dumps(request):
35                return Reply("text", "done", usage=None)
36            return text("done", usage={"input_tokens": 20, "output_tokens": 1})
37
38        self.provider = Provider(reply, catalog=CATALOG)
39        self.addCleanup(self.provider.close)
40
41        def prepare(app):
42            (app.home / "models.json").unlink(missing_ok=True)
43            app.write_extensions(
44                {
45                    "models": {
46                        "url": self.provider.url + "/models.json",
47                        "refreshHours": 24,
48                    }
49                }
50            )
51
52        self.app = Albedo(
53            self.provider,
54            protocol="responses",
55            prepare=prepare,
56            providers={
57                "fixture": {
58                    "baseUrl": self.provider.url + "/v1",
59                    "apiKey": SECRET,
60                    "model": "fixture",
61                    "protocol": "responses",
62                },
63            },
64        )
65        self.app.__enter__()
66        self.addCleanup(self.app.__exit__, None, None, None)
67        self.session = self.app.session()
68        self.route = f"/sessions/{self.session}/context"
69
70    def read(self, route):
71        with self.app.api(route) as response:
72            return json.load(response)
73
74    def turn(self, message):
75        self.app.prompt(self.session, message).close()
76        self.app.idle(self.session)
77
78    def test_pending_inspection_is_read_only_and_snapshot_matches_request(self):
79        pending = self.read(self.route)
80        self.assertEqual(pending["state"], "pending")
81        self.assertIsNone(pending["snapshot_id"])
82        self.assertEqual(set(pending["reason"]), {"code", "detail"})
83        self.assertEqual(pending["reason"]["code"], "context_pending")
84        self.assertIsInstance(pending["reason"]["detail"], str)
85        self.assertTrue(pending["reason"]["detail"].strip())
86        self.assertLessEqual(len(pending["reason"]["detail"]), 4096)
87        self.assertEqual(self.provider.requests, [])
88        self.turn("inspect exact composition")
89        self.assertEqual(len(self.provider.requests), 1)
90        sent = self.provider.requests[0]["request"]
91        snapshot = self.read(self.route)
92        self.assertEqual(snapshot["state"], "ready")
93        self.assertIsNone(snapshot["reason"])
94        self.assertEqual(
95            (snapshot["provider"], snapshot["model"]), ("fixture", "fixture")
96        )
97        self.assertEqual(snapshot["context_window_tokens"], 200000)
98        compaction = snapshot["compaction"]
99        self.assertEqual(compaction["provider_input_tokens"], 20)
100        self.assertIsNone(compaction["provider_cached_input_tokens"])
101        self.assertEqual(
102            compaction["estimate_method"],
103            "local byte-based estimate; not provider token usage",
104        )
105        self.assertEqual(compaction["status"], "not_needed")
106        self.assertEqual(compaction["input_limit_tokens"], 200000)
107        self.assertEqual(compaction["trigger_free_percent"], 20)
108        self.assertIn("fixture-cloud", compaction["source"])
109        self.assertIn("models.dev catalog", compaction["source"])
110        self.assertIn("<extension-context", sent["instructions"])
111        self.assertNotIn("<extension-context", json.dumps(sent["input"]))
112        self.assertTrue(
113            all(
114                section["preview"]["text"]
115                and len(section["preview"]["text"].encode()) <= 1024
116                for section in snapshot["sections"]
117            )
118        )
119        self.assertNotIn(SECRET, json.dumps(snapshot))
120
121        pages = {}
122        for section in snapshot["sections"]:
123            parts = [
124                self.read(
125                    self.route
126                    + "?"
127                    + urlencode(
128                        {
129                            "view": "section",
130                            "snapshot_id": snapshot["snapshot_id"],
131                            "section_id": section["id"],
132                            "page": page,
133                        }
134                    )
135                )
136                for page in range(section["page_count"])
137            ]
138            self.assertTrue(all(len(part["text"].encode()) <= 32768 for part in parts))
139            self.assertTrue(
140                all(part["snapshot_id"] == snapshot["snapshot_id"] for part in parts)
141            )
142            pages[section["id"]] = "".join(part["text"] for part in parts)
143        self.assertEqual(pages["instructions"], sent["instructions"])
144        self.assertEqual(json.loads(pages["tools"]), sent["tools"])
145        self.assertIn("inspect exact composition", pages["history"])
146        self.assertNotIn(SECRET, json.dumps(pages))
147        self.assertEqual(len(self.provider.requests), 1)
148
149    def test_combining_marks_cannot_make_context_pages_unbounded(self):
150        message = "a" + "\u0301" * 20000
151        self.turn(message)
152        session = self.read(f"/sessions/{self.session}?tail=0")
153        listed = next(
154            item
155            for item in self.read("/sessions")["items"]
156            if item["id"] == self.session
157        )
158        for representation in (session, listed):
159            for field in ("name", "automatic_name"):
160                self.assertTrue(representation[field].startswith("a"))
161                self.assertLessEqual(len(representation[field]), 4096)
162        resource = session["configuration_resource"]
163        configuration = self.read(resource["url"])
164        self.assertEqual(configuration, resource["value"])
165        self.assertIsNone(configuration["name"])
166        self.assertEqual(listed["name"], session["name"])
167        self.assertEqual(listed["automatic_name"], session["automatic_name"])
168        snapshot = self.read(self.route)
169        history = next(
170            section for section in snapshot["sections"] if section["id"] == "history"
171        )
172        parts = []
173        for page in range(history["page_count"]):
174            part = self.read(
175                self.route
176                + "?"
177                + urlencode(
178                    {
179                        "view": "section",
180                        "snapshot_id": snapshot["snapshot_id"],
181                        "section_id": "history",
182                        "page": page,
183                    }
184                )
185            )
186            self.assertLessEqual(len(part["text"]), 8000)
187            self.assertLessEqual(len(part["text"].encode()), 32000)
188            parts.append(part["text"])
189        self.assertGreater(len(parts), 1)
190        history_text = "".join(parts)
191        self.assertIn(message, history_text)
192        self.assertEqual(history["preview"]["text"], history_text[:256])
193        self.assertTrue(history["preview"]["truncated"])
194
195    def test_usage_absence_model_selection_and_unknown_catalog_model(self):
196        self.turn("without provider usage")
197        compaction = self.read(self.route)["compaction"]
198        self.assertIsNone(compaction["provider_input_tokens"])
199        self.assertIsNotNone(compaction["estimated_input_tokens"])
200        resource = self.read(f"/sessions/{self.session}")["configuration_resource"]
201        with self.app.api(
202            resource["url"],
203            {"model": "changed-model"},
204            method="PATCH",
205            headers={"If-Match": resource["etag"]},
206        ) as response:
207            changed = json.load(response)
208        self.assertEqual(changed["session"]["model"], "changed-model")
209        pending = self.read(self.route)
210        self.assertEqual(pending["state"], "pending")
211        self.assertIsNone(pending["snapshot_id"])
212        self.assertEqual(len(self.provider.requests), 1)
213        self.turn("a model outside the catalog")
214        unknown = self.read(self.route)
215        self.assertIsNone(unknown["context_window_tokens"])
216        self.assertEqual(unknown["compaction"]["status"], "unknown")
217
218
219class ContextReplayTest(unittest.TestCase):
220    def test_inspection_preserves_replay_and_replaces_prior_request_for_both_protocols(
221        self,
222    ):
223        for protocol in ("responses", "chat_completions"):
224            with self.subTest(protocol=protocol):
225                opaque = "opaque-replay-must-only-reach-provider"
226                if protocol == "responses":
227                    replay = {
228                        "type": "reasoning",
229                        "id": "fixed-reasoning",
230                        "summary": [],
231                        "encrypted_content": opaque,
232                    }
233                    message = {
234                        "type": "message",
235                        "role": "assistant",
236                        "content": [{"type": "output_text", "text": "first answer"}],
237                    }
238                    initial = Reply(
239                        "raw",
240                        events=[
241                            {
242                                "type": "response.completed",
243                                "response": {
244                                    "id": "fixed-response",
245                                    "status": "completed",
246                                    "output": [replay, message],
247                                    "usage": {
248                                        "input_tokens": 31,
249                                        "output_tokens": 2,
250                                        "input_tokens_details": {"cached_tokens": 9},
251                                    },
252                                },
253                            }
254                        ],
255                    )
256                else:
257                    replay = {
258                        "role": "assistant",
259                        "content": "first answer",
260                        "reasoning_content": opaque,
261                    }
262                    initial = Reply(
263                        "raw",
264                        events=[
265                            {
266                                "id": "fixed-chat",
267                                "choices": [
268                                    {"index": 0, "delta": replay, "finish_reason": None}
269                                ],
270                            },
271                            {
272                                "id": "fixed-chat",
273                                "choices": [
274                                    {"index": 0, "delta": {}, "finish_reason": "stop"}
275                                ],
276                                "usage": {
277                                    "prompt_tokens": 31,
278                                    "completion_tokens": 2,
279                                    "prompt_tokens_details": {"cached_tokens": 9},
280                                },
281                            },
282                        ],
283                    )
284
285                def reply(request):
286                    inputs = request.get("input", request.get("messages", []))
287                    users = [item for item in inputs if item.get("role") == "user"]
288                    return (
289                        initial
290                        if len(users) == 1
291                        else Reply("text", "later answer", usage=None)
292                    )
293
294                provider = Provider(reply)
295                try:
296                    with Albedo(provider, protocol=protocol) as app:
297                        session = app.session()
298                        route = f"/sessions/{session}/context"
299
300                        def inspect():
301                            with app.api(route) as response:
302                                summary = json.load(response)
303                            contents = {}
304                            for section in summary.get("sections", []):
305                                parts = []
306                                for index in range(section["page_count"]):
307                                    with app.api(
308                                        route
309                                        + "?"
310                                        + urlencode(
311                                            {
312                                                "view": "section",
313                                                "snapshot_id": summary["snapshot_id"],
314                                                "section_id": section["id"],
315                                                "page": index,
316                                            }
317                                        )
318                                    ) as response:
319                                        parts.append(json.load(response)["text"])
320                                contents[section["id"]] = "".join(parts)
321                            return summary, contents
322
323                        app.prompt(session, "first request").close()
324                        app.idle(session)
325                        first, pages = inspect()
326                        self.assertNotIn("first answer", pages["history"])
327                        self.assertNotIn(opaque, json.dumps((first, pages)))
328                        self.assertEqual(len(provider.requests), 1)
329                        app.prompt(session, "replacement request").close()
330                        app.idle(session)
331                        second, pages = inspect()
332                        with self.assertRaises(urllib.error.HTTPError) as obsolete:
333                            app.api(
334                                route
335                                + "?"
336                                + urlencode(
337                                    {
338                                        "view": "section",
339                                        "snapshot_id": first["snapshot_id"],
340                                        "section_id": first["sections"][0]["id"],
341                                        "page": 0,
342                                    }
343                                )
344                            )
345                        self.assertEqual(obsolete.exception.code, 410)
346                        self.assertEqual(
347                            json.load(obsolete.exception)["code"], "context_changed"
348                        )
349                        self.assertIn("replacement request", pages["history"])
350                        self.assertNotIn("later answer", pages["history"])
351                        self.assertNotIn(opaque, json.dumps((second, pages)))
352                        sent = provider.requests[-1]["request"]
353                        inputs = sent.get("input", sent.get("messages", []))
354                        self.assertIn(replay, inputs)
355                        self.assertEqual(len(provider.requests), 2)
356                        inspect()
357                        self.assertEqual(len(provider.requests), 2)
358                finally:
359                    provider.close()