okami.mom/albedo

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

nekomimi.petlist a generic openai endpoint's own /models ahead of models.devc77d3cf

main
branchesmainall refs →
16.0 KiB398 linesraw
1"""Provider model lists: each provider's own list reaches the picker and reloads.
2
3Alibaba caches its /models answer for a day, so a model it starts offering
4stays out of the picker until that cache ages out; `/reload models` must
5refetch it at once through the configured profile's endpoint and key. A
6catalog that cannot be fetched (Codex, Antigravity, and Claude are never
7signed in here) keeps its previous list: the reload still succeeds and names
8it rather than reporting it refreshed. Claude's picker shows whatever the
9Anthropic Models API listed, newest first, with the API's own limits and
10efforts, so a model released after albedo shipped is offered without an
11update. A generic openai profile lists what its own endpoint's /models serves. Every list is served on loopback or seeded on disk, so no live network
12is reached. The class writes the shared home's list caches and models
13settings, so it is exclusive and removes the caches again.
14"""
15
16import json
17import unittest
18import urllib.parse
19
20from harness import Albedo, Provider, exclusive, operation_id, text
21
22MODELS_CATALOG = {
23    "fixture": {
24        "id": "fixture",
25        "name": "Fixture",
26        "models": {
27            "fixture-model": {"id": "fixture-model", "limit": {"context": 8000}}
28        },
29    },
30    # Claims efforts the Anthropic API says the model does not take.
31    "anthropic": {
32        "id": "anthropic",
33        "name": "Anthropic",
34        "models": {
35            "claude-fixture-small": {
36                "id": "claude-fixture-small",
37                "limit": {"context": 200000, "output": 8000},
38                "reasoning_options": [
39                    {"type": "effort", "values": ["low", "medium", "high"]}
40                ],
41            }
42        },
43    },
44    # Must not hide what a generic openai endpoint's own /models serves.
45    "openai": {
46        "id": "openai",
47        "name": "OpenAI",
48        "models": {"gpt-fixture": {"id": "gpt-fixture", "limit": {"context": 8000}}},
49    },
50}
51
52ALIBABA_MODELS = {"data": [{"id": "qwen-fresh"}, {"id": "qwen-tts-fresh"}]}
53
54GATEWAY_MODELS = {"data": [{"id": "local-b"}, {"id": "local-a"}]}
55
56# As the Anthropic Models API answered, newest first; neither id is one
57# albedo or models.dev knows.
58CLAUDE_MODELS = [
59    {
60        "id": "claude-fixture-6",
61        "context": 500000,
62        "output": 32000,
63        "images": True,
64        "efforts": ["low", "high"],
65    },
66    {
67        "id": "claude-fixture-small",
68        "context": 200000,
69        "output": 8000,
70        "images": False,
71        "efforts": [],
72    },
73]
74
75
76# exclusive: writes global model caches/settings and reloads provider catalogs
77@exclusive
78class ModelCatalogReloadTests(unittest.TestCase):
79    def setUp(self):
80        self.provider = Provider(lambda _request: text("ok"))
81        self.addCleanup(self.provider.close)
82        self.catalog = Provider(lambda _request: text("ok"), catalog=MODELS_CATALOG)
83        self.addCleanup(self.catalog.close)
84        self.alibaba = Provider(lambda _request: text("ok"), catalog=ALIBABA_MODELS)
85        self.addCleanup(self.alibaba.close)
86
87        def prepare(app):
88            # Fresh on disk, so nothing but an explicit reload refetches them.
89            (app.home / "alibaba-models.json").write_text(json.dumps(["qwen-stale"]))
90            (app.home / "claude-models.json").write_text(json.dumps(CLAUDE_MODELS))
91            (app.home / "models.json").unlink(missing_ok=True)
92            app.write_extensions({"models": {"url": self.catalog.url + "/models.json"}})
93
94        self.app = Albedo(
95            self.provider,
96            prepare=prepare,
97            providers={
98                "fixture": {
99                    "baseUrl": self.provider.url,
100                    "apiKey": "fixture-key",
101                    "model": "fixture-model",
102                    "protocol": "chat_completions",
103                },
104                "fixture-alibaba": {
105                    "extension": "alibaba",
106                    "baseUrl": self.alibaba.url,
107                    "apiKey": "fixture-alibaba-key",
108                    "model": "qwen-fresh",
109                    "protocol": "chat_completions",
110                },
111            },
112        )
113        self.app.__enter__()
114        for name in ("alibaba-models.json", "claude-models.json"):
115            self.addCleanup((self.app.home / name).unlink, missing_ok=True)
116        self.addCleanup(self.app.__exit__, None, None, None)
117
118    def listed(self, provider, endpoint=""):
119        query = urllib.parse.urlencode(
120            {"provider": provider, **({"endpoint": endpoint} if endpoint else {})}
121        )
122        with self.app.api(f"/models?{query}") as response:
123            return json.load(response)["items"]
124
125    def reload(self, target):
126        session = self.app.session()
127        with self.app.api(
128            f"/sessions/{session}/reload",
129            {"target": target},
130        ) as response:
131            return json.load(response)
132
133    def alibaba_listed(self):
134        return self.listed("alibaba", self.alibaba.url)
135
136    def test_reload_models_refetches_provider_lists_and_names_failures(self):
137        self.assertEqual([item["id"] for item in self.alibaba_listed()], ["qwen-stale"])
138        outcome = self.reload("models")
139        catalogs = {item["provider"]: item for item in outcome["models"]}
140        self.assertEqual(catalogs["alibaba"]["state"], "refreshed")
141        self.assertEqual(catalogs["models"]["state"], "refreshed")
142        self.assertEqual(
143            {"codex", "antigravity", "claude"},
144            {name for name, item in catalogs.items() if item["state"] == "failed"},
145        )
146        self.assertTrue(catalogs["claude"]["failure"])
147        # The non-chat entitlement stays filtered, as on any other fetch.
148        self.assertEqual([item["id"] for item in self.alibaba_listed()], ["qwen-fresh"])
149        # A failed reload keeps the list it had.
150        self.assertEqual(
151            [item["id"] for item in self.listed("claude")],
152            ["claude-fixture-6", "claude-fixture-small"],
153        )
154
155    def test_generic_openai_endpoint_lists_its_own_models(self):
156        gateway = Provider(lambda _request: text("ok"), catalog=GATEWAY_MODELS)
157        self.addCleanup(gateway.close)
158        self.reload("models")
159        self.assertEqual(
160            [item["id"] for item in self.listed("openai", gateway.url + "/")],
161            ["local-a", "local-b"],
162        )
163
164    def test_claude_picker_offers_the_api_list_with_its_facts(self):
165        self.reload("models")
166        newest, small = self.listed("claude")
167        self.assertEqual(
168            [newest["id"], small["id"]], ["claude-fixture-6", "claude-fixture-small"]
169        )
170        self.assertEqual(
171            (
172                newest["default_context_tokens"],
173                newest["max_output_tokens"],
174                [effort["id"] for effort in newest["efforts"]],
175            ),
176            (500000, 32000, ["low", "high"]),
177        )
178        # No effort listed means the model takes none, not a guessed default.
179        self.assertEqual(small["efforts"], [])
180
181
182# exclusive: seeds global models catalog; several tests restart the daemon
183@exclusive
184class ModelCatalogSelectionTests(unittest.TestCase):
185    """Disputed gateway facts must not borrow one provider's larger limits."""
186
187    def setUp(self):
188        self.provider = Provider(lambda _request: text("ok"))
189        self.addCleanup(self.provider.close)
190        catalog = {
191            "openai": {
192                "api": "https://api.openai.example/v1",
193                "models": {
194                    "shared-model": {
195                        "id": "shared-model",
196                        "limit": {"context": 400000, "output": 8000},
197                        "modalities": {"input": ["text", "image", "pdf"]},
198                        "reasoning_options": [
199                            {"type": "effort", "values": ["low", "medium", "high"]}
200                        ],
201                    }
202                },
203            },
204            "mirror": {
205                "api": "https://mirror.example/v1",
206                "models": {
207                    "shared-model": {
208                        "id": "shared-model",
209                        "limit": {"context": 200000, "output": 16000},
210                        "modalities": {"input": ["image", "text"]},
211                        "reasoning_options": [
212                            {"type": "effort", "values": ["medium", "high"]}
213                        ],
214                    },
215                    "vendor/qualified-model": {
216                        "id": "vendor/qualified-model",
217                        "limit": {"context": 32000},
218                    },
219                },
220            },
221            "unknown": {
222                "models": {
223                    "shared-model": {
224                        "id": "shared-model",
225                        "limit": {"context": None, "output": 0},
226                    }
227                }
228            },
229        }
230
231        def prepare(app):
232            (app.home / "models.json").write_text(json.dumps(catalog))
233            (app.home / "models.json.index").unlink(missing_ok=True)
234            app.write_extensions({"models": {"refreshHours": 0}})
235
236        self.app = Albedo(self.provider, prepare=prepare)
237        self.app.__enter__()
238        self.addCleanup(self.app.__exit__, None, None, None)
239
240    def listed(self, provider, endpoint=""):
241        query = urllib.parse.urlencode(
242            {"provider": provider, **({"endpoint": endpoint} if endpoint else {})}
243        )
244        with self.app.api(f"/models?{query}") as response:
245            return json.load(response)["items"]
246
247    def test_gateway_uses_smallest_limits_and_shared_reported_capabilities(self):
248        gateway = "https://gateway.example/v1"
249        [shared] = self.listed("openai", gateway)
250        self.assertEqual(shared["default_context_tokens"], 200000)
251        self.assertEqual(shared["max_output_tokens"], 8000)
252        self.assertEqual(shared["input_modalities"], ["image", "text"])
253        self.assertEqual(
254            [effort["id"] for effort in shared["efforts"]], ["medium", "high"]
255        )
256        self.assertEqual(self.listed("openai"), [shared])
257        # The reduced index must give the same answer after a daemon restart.
258        self.assertTrue((self.app.home / "models.json.index").exists())
259        self.app.restart()
260        self.assertEqual(self.listed("openai", gateway), [shared])
261
262    def test_endpoint_selects_provider_facts_and_model_list(self):
263        listed = self.listed("openai", "https://MIRROR.example/other-path")
264        self.assertEqual(
265            [item["id"] for item in listed],
266            ["shared-model", "vendor/qualified-model"],
267        )
268        shared, qualified = listed
269        self.assertEqual(shared["default_context_tokens"], 200000)
270        self.assertEqual(shared["max_output_tokens"], 16000)
271        self.assertEqual(shared["input_modalities"], ["image", "text"])
272        self.assertEqual(qualified["default_context_tokens"], 32000)
273        # ChatGPT has no matching API host in models.dev; it uses OpenAI identity.
274        [subscription] = self.listed("openai", "https://chatgpt.com/backend-api")
275        self.assertEqual(subscription["default_context_tokens"], 400000)
276        self.assertEqual(subscription["max_output_tokens"], 8000)
277        self.assertEqual(subscription["input_modalities"], ["text", "image", "pdf"])
278        self.assertEqual(
279            [effort["id"] for effort in subscription["efforts"]],
280            ["low", "medium", "high"],
281        )
282
283    def test_known_reasoning_models_keep_empty_catalog_efforts(self):
284        catalog = {
285            "openai": {
286                "models": {
287                    "o3-disjoint": {
288                        "id": "o3-disjoint",
289                        "reasoning_options": [{"type": "effort", "values": ["low"]}],
290                    },
291                    "o3-empty": {"id": "o3-empty"},
292                }
293            },
294            "mirror": {
295                "models": {
296                    "o3-disjoint": {
297                        "id": "o3-disjoint",
298                        "reasoning_options": [{"type": "effort", "values": ["high"]}],
299                    }
300                }
301            },
302        }
303        (self.app.home / "models.json").write_text(json.dumps(catalog))
304        self.app.restart()
305        listed = self.listed("openai", "https://gateway.example/v1")
306        self.assertEqual([item["id"] for item in listed], ["o3-disjoint", "o3-empty"])
307        self.assertEqual([item["efforts"] for item in listed], [[], []])
308        self.app.restart()
309        self.assertEqual(self.listed("openai"), listed)
310
311    def test_efforts_are_inferred_only_for_a_model_absent_from_a_valid_catalog(self):
312        session = operation_id()
313        self.app.api(
314            f"/sessions/{session}",
315            {
316                "kind": "new",
317                "workspace": str(self.app.workspace),
318                "provider_profile": self.app.profile,
319                "model": "o3-unlisted",
320            },
321            method="PUT",
322            headers={"If-None-Match": "*"},
323        ).close()
324        query = urllib.parse.urlencode({"provider": "openai", "model": "o3-unlisted"})
325        with self.app.api(f"/models?{query}") as response:
326            [model] = json.load(response)["items"]
327        self.assertEqual(
328            [effort["id"] for effort in model["efforts"]], ["low", "medium", "high"]
329        )
330        # A corrupt catalog cannot establish that the model is absent.
331        (self.app.home / "models.json").write_text("{")
332        self.app.restart()
333        with self.app.api(f"/models?{query}") as response:
334            [model] = json.load(response)["items"]
335        self.assertEqual(model["efforts"], [])
336
337
338# exclusive: seeds global Codex catalog and restarts the daemon
339@exclusive
340class CodexCachedCatalogTests(unittest.TestCase):
341    """Account caches keep picker precedence and facts across restart."""
342
343    def test_cache_deduplicates_before_visibility_and_survives_restart(self):
344        provider = Provider(lambda _request: text("ok"))
345        self.addCleanup(provider.close)
346        cache = {
347            "clientVersion": "0.157.1",
348            "versionCheckedAt": 123,
349            "accounts": {
350                "account-a": {
351                    "fetchedAt": 123,
352                    "clientVersion": "0.157.1",
353                    "etag": "fixture-a",
354                    "models": [
355                        {"slug": "hidden-first", "visible": False},
356                        {"slug": "hidden-first", "visible": True},
357                        {
358                            "slug": "codex-fixture",
359                            "name": "Codex fixture",
360                            "context": 200000,
361                            "maxContext": 400000,
362                            "input": ["text", "image"],
363                            "efforts": ["low", "high"],
364                            "priority": 2,
365                        },
366                        {"slug": "codex-fixture", "context": 1000, "priority": 0},
367                    ],
368                },
369                "account-b": {
370                    "models": [{"slug": "other-account", "priority": 1}],
371                },
372            },
373        }
374
375        def prepare(app):
376            (app.home / "codex-models.json").write_text(json.dumps(cache))
377            app.write_extensions({"models": {"refreshHours": 0}})
378
379        with Albedo(provider, prepare=prepare) as app:
380
381            def listed():
382                with app.api("/models?provider=codex") as response:
383                    return json.load(response)["items"]
384
385            expected = ["other-account", "codex-fixture"]
386            self.assertEqual([item["id"] for item in listed()], expected)
387            facts = listed()
388            self.assertEqual([item["id"] for item in facts], expected)
389            model = facts[1]
390            self.assertEqual(model["default_context_tokens"], 200000)
391            self.assertEqual(model["max_context_tokens"], 400000)
392            self.assertEqual(model["input_modalities"], ["text", "image"])
393            self.assertEqual(
394                [effort["id"] for effort in model["efforts"]], ["low", "high"]
395            )
396            app.restart()
397            self.assertEqual([item["id"] for item in listed()], expected)
398            self.assertEqual(listed(), facts)