okami.mom/albedo

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

nekomimi.petfeat(usage): time provider calls and show tok/s on the turn signoff0a3eddd

main
branchesmainall refs →
12.6 KiB304 linesraw
1"""Provider requests: one durable row per attempt albedo sends, recording what cache
2and quota analysis cannot recover from the transcript.
3
4Catches: usage counts that never reach a request row (cache reads, TTL-split
5writes, reasoning), the account/provider identity of each call, the prefix
6identity that distinguishes a warm append from a compaction rewrite, an HTTP
7failure recorded as a quota reading, and the seq link to the transcript. None
8of this is observable from events or history alone.
9"""
10
11import json
12import unittest
13from datetime import datetime
14
15from harness import Albedo, Provider, Reply, text
16
17FIRST = {
18    "input_tokens": 100,
19    "output_tokens": 20,
20    "input_tokens_details": {"cached_tokens": 40},
21    "output_tokens_details": {"reasoning_tokens": 7},
22}
23TOOL = {
24    "input_tokens": 200,
25    "output_tokens": 30,
26    "input_tokens_details": {"cached_tokens": 80},
27    "output_tokens_details": {"reasoning_tokens": 11},
28}
29AFTER = {
30    "input_tokens": 300,
31    "output_tokens": 40,
32    "input_tokens_details": {"cached_tokens": 120},
33    "output_tokens_details": {"reasoning_tokens": 3},
34}
35THIRD = {
36    "input_tokens": 400,
37    "output_tokens": 50,
38    "input_tokens_details": {"cached_tokens": 160},
39    "output_tokens_details": {"reasoning_tokens": 5},
40}
41SUMMARY = {"input_tokens": 60, "output_tokens": 10}
42POST = {
43    "input_tokens": 70,
44    "output_tokens": 15,
45    "input_tokens_details": {"cached_tokens": 30},
46    "output_tokens_details": {"reasoning_tokens": 2},
47}
48
49
50def reply(request):
51    inputs = request["input"]
52    # The tool result arrives while the newest user message is still the one
53    # that asked for the tool, so it is told apart by input shape.
54    if inputs[-1].get("type") == "function_call_output":
55        return text("tool done", usage=AFTER)
56    user = next(
57        (
58            str(item.get("content", ""))
59            for item in reversed(inputs)
60            if item.get("role") == "user"
61        ),
62        "",
63    )
64    if "<newly-evicted-history>" in user:
65        return text("older conversation summary", usage=SUMMARY)
66    if user == "use the tool":
67        return Reply("python", "2 + 2", usage=TOOL)
68    if user == "hit the limit":
69        return Reply("error", "limit exceeded", status=429)
70    if user == "after compaction":
71        return text("post compaction reply", usage=POST)
72    if user == "third":
73        return text("third reply", usage=THIRD)
74    return text("first reply", usage=FIRST)
75
76
77def instant(timestamp):
78    return datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
79
80
81class ProviderRequestsTest(unittest.TestCase):
82    def rows(self, app, session, query=""):
83        with app.api(f"/sessions/{session}/context?view=requests{query}") as response:
84            return json.load(response)
85
86    def history(self, app, session):
87        items, after = [], 0
88        while True:
89            with app.api(
90                f"/sessions/{session}/history?after={after}&limit=100"
91            ) as response:
92                page = json.load(response)
93            items.extend(page["items"])
94            if page["newer"] is None:
95                return items
96            after = items[-1]["position"]
97
98    def test_rows_record_usage_identity_prefix_and_projection(self):
99        provider = Provider(reply)
100        try:
101            with Albedo(provider, protocol="responses") as app:
102                session = app.session()
103                for prompt in ("first", "use the tool", "third", "hit the limit"):
104                    app.prompt(session, prompt).close()
105                    app.idle(session)
106                with app.api(f"/sessions/{session}/compaction", {}) as response:
107                    self.assertEqual(json.load(response)["state"], "compacted")
108                app.idle(session)
109                app.prompt(session, "after compaction").close()
110                app.idle(session)
111
112                page = self.rows(app, session)
113                rows = page["items"]
114                self.assertEqual(
115                    [row["kind"] for row in rows],
116                    [
117                        "turn",
118                        "turn",
119                        "turn",
120                        "turn",
121                        "turn",
122                        "summarizer",
123                        "summarizer",
124                        "turn",
125                    ],
126                )
127                # Every call carries its own reported counts, split as sent.
128                self.assertEqual(
129                    [
130                        (
131                            row["prompt_tokens"]["observed"],
132                            row["completion_tokens"]["observed"],
133                            row["cached_tokens"]["observed"],
134                            row["reasoning_tokens"]["observed"],
135                        )
136                        for row in rows
137                    ],
138                    [
139                        (100, 20, 40, 7),
140                        (200, 30, 80, 11),
141                        (300, 40, 120, 3),
142                        (400, 50, 160, 5),
143                        (None, None, None, None),
144                        (60, 10, None, None),
145                        (60, 10, None, None),
146                        (70, 15, 30, 2),
147                    ],
148                )
149                # The head (instructions and tools) never changed, so one hash.
150                heads = {row["head_hash"] for row in rows if row["kind"] == "turn"}
151                self.assertEqual(len(heads), 1)
152                self.assertNotIn(rows[5]["head_hash"], heads)
153
154                for row in rows:
155                    self.assertIsNone(row["account_id"])
156                    # The shared harness names the fixture profile per route.
157                    self.assertTrue(row["provider_profile"].startswith("fixture"))
158                    self.assertEqual(row["model"], "fixture-model")
159                    self.assertTrue(row["provider"].startswith("responses:"))
160                    self.assertLessEqual(row["started_at"], row["ended_at"])
161                    # An OpenAI-protocol provider caches on its own; albedo marks nothing.
162                    self.assertEqual(row["cache_marks"], [])
163
164                for row in rows[:4] + [rows[7]]:
165                    self.assertEqual(row["outcome"], "completed")
166                    self.assertIsNone(row["http_status"])
167                    self.assertIsNone(row["failure"])
168
169                for row in rows[:4]:
170                    # Nothing replaced yet: the identity is a plain append.
171                    self.assertEqual(row["replaced_input_count"], 0)
172                    self.assertIsNone(row["projection_hash"])
173
174                # A 429 is a quota reading, with no usage and no transcript row.
175                limited = rows[4]
176                self.assertEqual(limited["outcome"], "failed")
177                self.assertEqual(limited["http_status"], 429)
178                self.assertEqual(set(limited["failure"]), {"code", "detail"})
179                self.assertEqual(limited["failure"]["code"], "provider_request_failed")
180                self.assertIn("limit exceeded", limited["failure"]["detail"])
181                self.assertLessEqual(len(limited["failure"]["detail"]), 4096)
182                self.assertEqual(limited["transcript_positions"], [])
183
184                # Each turn row's seq is the transcript row its response
185                # committed, exactly: the assistant row that turn produced,
186                # found again through the tree route straight from the
187                # transcript.
188                entries = self.history(app, session)
189                assistant = {
190                    part["text"]: entry["position"]
191                    for entry in entries
192                    if entry["kind"] == "assistant"
193                    for part in entry["content"]
194                    if part["kind"] == "text"
195                }
196                for index, message in (
197                    (0, "first reply"),
198                    (2, "tool done"),
199                    (3, "third reply"),
200                    (7, "post compaction reply"),
201                ):
202                    self.assertIn(
203                        assistant[message], rows[index]["transcript_positions"]
204                    )
205                calls = [entry for entry in entries if entry["kind"] == "tool_call"]
206                self.assertEqual(len(calls), 1)
207                self.assertIn(calls[0]["position"], rows[1]["transcript_positions"])
208                tool_output = entries[entries.index(calls[0]) + 1]
209                self.assertEqual(tool_output["kind"], "tool_result")
210                self.assertNotIn(
211                    tool_output["position"], rows[1]["transcript_positions"]
212                )
213                linked = [
214                    position for row in rows for position in row["transcript_positions"]
215                ]
216                self.assertEqual(linked, sorted(linked))
217
218                # The summary and the notes rewrite are their own kind,
219                # outside the session projection.
220                for summary in rows[5:7]:
221                    self.assertEqual(summary["transcript_positions"], [])
222                    self.assertIsNone(summary["replaced_input_count"])
223                    self.assertIsNone(summary["projection_hash"])
224                    self.assertEqual(summary["input_count"], 1)
225
226                # The projection identity changes once compaction rewrites the
227                # head the verbatim tail hangs from.
228                after = rows[7]
229                self.assertGreater(after["replaced_input_count"], 0)
230                self.assertIsNotNone(after["projection_hash"])
231                self.assertTrue(after["strategy"])
232
233                # Paging: continue after the first row, bounded by limit.
234                second = self.rows(
235                    app, session, f"&after={rows[0]['sequence']}&limit=2"
236                )
237                self.assertEqual(
238                    [row["sequence"] for row in second["items"]],
239                    [row["sequence"] for row in rows[1:3]],
240                )
241                self.assertIsNotNone(second["next"])
242        finally:
243            provider.close()
244
245    def test_the_latest_usage_reports_its_calls_span_and_rate(self):
246        provider = Provider(reply)
247        try:
248            with Albedo(provider, protocol="responses") as app:
249                session = app.session()
250                app.prompt(session, "third").close()
251                app.idle(session)
252                row = self.rows(app, session)["items"][-1]
253                with app.api(f"/sessions/{session}?tail=0") as response:
254                    usage = json.load(response)["usage"]
255                started, ended = instant(row["started_at"]), instant(row["ended_at"])
256                span = round((ended - started).total_seconds() * 1000)
257                self.assertEqual(usage["completion_tokens"], 50)
258                self.assertEqual(usage["elapsed_ms"], span)
259                if span:
260                    self.assertAlmostEqual(usage["tokens_per_second"], 50 * 1000 / span)
261                else:
262                    self.assertIsNone(usage["tokens_per_second"])
263        finally:
264            provider.close()
265
266    def test_a_failing_gateway_is_retried_with_doubling_waits(self):
267        failures = []
268
269        def script(request):
270            if len(failures) < 4:
271                failures.append(1)
272                return Reply("error", "gateway down", status=503)
273            return text("recovered")
274
275        provider = Provider(script)
276        try:
277            with Albedo(provider, protocol="responses") as app:
278                session = app.session()
279                app.prompt(session, "hello").close()
280                app.idle(session, timeout=60)
281                rows = self.rows(app, session)["items"]
282                self.assertEqual(
283                    [row["outcome"] for row in rows], ["failed"] * 4 + ["completed"]
284                )
285                self.assertEqual({row["http_status"] for row in rows[:4]}, {503})
286                self.assertTrue(rows[4]["transcript_positions"])
287                # Each wait doubles from 250ms, measured from the failed
288                # attempt's end to the next one's start.
289                gaps = [
290                    (
291                        datetime.fromisoformat(
292                            after["started_at"].replace("Z", "+00:00")
293                        )
294                        - datetime.fromisoformat(
295                            before["ended_at"].replace("Z", "+00:00")
296                        )
297                    ).total_seconds()
298                    * 1000
299                    for before, after in zip(rows, rows[1:])
300                ]
301                for gap, wait in zip(gaps, (250, 500, 1000, 2000)):
302                    self.assertGreaterEqual(gap, wait)
303        finally:
304            provider.close()