Spaces:
Running
Running
| """A generation-only Studio cannot advertise or admit the other generation.""" | |
| import asyncio | |
| import os | |
| import unittest | |
| from contextlib import suppress | |
| from unittest.mock import patch | |
| import httpx | |
| from fastapi.testclient import TestClient | |
| from app import create_app | |
| from contract import MODEL | |
| from model_registry import PROFILES, model_registry | |
| from tetris_arena import HTTPDecisionAdapter, LOCAL_MODELS | |
| TOKEN = "t" * 32 | |
| QUESTIONS = {"accept": {"type": "noul", "instructions": "Is this accepted?"}} | |
| def generation_registry(version): | |
| return [ | |
| { | |
| "id": wire, | |
| "label": profile["repo_id"].rsplit("/", 1)[-1].split("-")[2], | |
| "version": version, | |
| "manifest_sha256": f"{index + 1:064x}", | |
| "revision": f"{index + 1:040x}", | |
| } | |
| for index, (wire, profile) in enumerate(PROFILES.items()) | |
| if f"/Decision-{version}-" in profile["repo_id"] | |
| ] | |
| class GenerationSplitTests(unittest.TestCase): | |
| def test_each_generation_is_an_independent_valid_registry(self): | |
| for version in ("1.0", "2.0"): | |
| with self.subTest(version=version): | |
| registry = model_registry(generation_registry(version)) | |
| self.assertEqual(len(registry), 6) | |
| self.assertTrue(all(f"/Decision-{version}-" in row["repo_id"] | |
| for row in registry.values())) | |
| def test_native_rejects_single_models_other_than_decision1_kai(self): | |
| for wire in ("decision2-kai", "decision-lux"): | |
| version = "2.0" if wire.startswith("decision2-") else "1.0" | |
| entry = next(row for row in generation_registry(version) if row["id"] == wire) | |
| with self.subTest(wire=wire), self.assertRaisesRegex(ValueError, "Native mode requires"): | |
| create_app(mode="native", registry=[entry]) | |
| def test_pull_queue_discovery_readiness_and_admission_follow_generation(self): | |
| for version, other in (("1.0", "2.0"), ("2.0", "1.0")): | |
| with self.subTest(version=version), patch.dict(os.environ, { | |
| "DECISION_WORKER_TOKEN": TOKEN, | |
| "TETRIS_LOCAL_API_URL": "http://127.0.0.1:7860/v1/systemone", | |
| "TETRIS_JEV_API_URL": "https://cloud.example.test/v1/systemone", | |
| "TETRIS_JEV_API_KEY": "test-cloud-key", | |
| }, clear=True): | |
| entries = generation_registry(version) | |
| api = create_app(mode="pull_queue", registry=entries) | |
| with TestClient(api) as client: | |
| expected = {PROFILES[row["id"]]["repo_id"] for row in entries} | |
| self.assertEqual(set(api.state.relays), {row["id"] for row in entries}) | |
| self.assertEqual({row["id"] for row in client.get("/v1/models").json()["models"]}, expected) | |
| self.assertEqual(client.get("/api/ready").status_code, 503) | |
| for index, row in enumerate(entries): | |
| heartbeat = client.post("/internal/worker/heartbeat", headers={ | |
| "Authorization": "Bearer " + TOKEN, | |
| }, json={ | |
| "model": row["id"], | |
| "manifest_sha256": row["manifest_sha256"], | |
| "worker_id": f"{index + 1:032x}", | |
| "phase": "ready", | |
| "capabilities": ["context_batch_v1"], | |
| }) | |
| self.assertEqual(heartbeat.status_code, 200, heartbeat.text) | |
| for data in ({"state": "A request"}, {"states": [{"id": "a", "state": "A request"}]}): | |
| job = client.post("/api/jobs", json={ | |
| "model": row["id"], "questions": QUESTIONS, **data, | |
| }) | |
| self.assertEqual(job.status_code, 202, job.text) | |
| self.assertEqual(job.json()["model"], row["id"]) | |
| client.delete(f"/api/jobs/{job.json()['id']}", params={"model": row["id"]}) | |
| ready = client.get("/api/ready") | |
| self.assertEqual(ready.status_code, 200) | |
| self.assertEqual(ready.json()["models"], dict.fromkeys(expected, True)) | |
| catalog = client.get("/api/tetris/config").json()["competitors"] | |
| allowed_competitors = {row.id for row in LOCAL_MODELS if row.request_model in expected} | |
| self.assertEqual({row["id"] for row in catalog if row["family"] == "decision"}, allowed_competitors) | |
| self.assertTrue(all(row["ready"] for row in catalog if row["family"] == "decision")) | |
| self.assertTrue(next(row for row in catalog if row["id"] == "jev-cloud")["ready"]) | |
| excluded = generation_registry(other)[0] | |
| canonical = PROFILES[excluded["id"]]["repo_id"] | |
| for route, data in (("/v1/systemone", {"state": "A request"}), | |
| ("/v1/systemone/batches", {"states": [{"id": "a", "state": "A request"}]})): | |
| self.assertEqual(client.post(route, json={ | |
| "model": canonical, "questions": QUESTIONS, **data, | |
| }).status_code, 422) | |
| self.assertEqual(client.post("/api/jobs", json={ | |
| "model": excluded["id"], "state": "A request", "questions": QUESTIONS, | |
| }).status_code, 422) | |
| excluded_competitor = next(row.id for row in LOCAL_MODELS if row.request_model == canonical) | |
| race = client.post("/api/tetris/races", json={ | |
| "left": excluded_competitor, | |
| "right": next(iter(allowed_competitors)), | |
| "mode": "steps", "max_steps": 1, | |
| }) | |
| self.assertEqual(race.status_code, 422, race.text) | |
| self.assertEqual(api.state.tetris.sessions, {}) | |
| def test_tetris_allowlist_rejects_unknown_model(self): | |
| with self.assertRaisesRegex(ValueError, "Unknown local Decision model allowlist"): | |
| HTTPDecisionAdapter.from_environment({}, allowed_local_models={"unknown/model"}) | |
| class StudioGenerationViewTests(unittest.TestCase): | |
| def test_generation_view_filters_discovery_and_arena_without_changing_queues(self): | |
| entries = generation_registry("2.0") + generation_registry("1.0") | |
| for version in ("", "2.0", "1.0"): | |
| with self.subTest(version=version), patch.dict(os.environ, { | |
| "DECISION_WORKER_TOKEN": TOKEN, | |
| "DECISION_STUDIO_GENERATION": version, | |
| "TETRIS_LOCAL_API_URL": "http://127.0.0.1:7860/v1/systemone", | |
| "TETRIS_JEV_API_URL": "https://jev.example.test/v1/systemone", | |
| "TETRIS_JEV_API_KEY": "test-upstream-secret", | |
| "TETRIS_SYSTEMONE_API_URL": "https://mirror.example.test/v1/systemone", | |
| "TETRIS_SYSTEMONE_API_KEY": "test-cloud-secret", | |
| }, clear=True): | |
| api = create_app(mode="pull_queue", registry=entries) | |
| expected = {row["id"] for row in entries | |
| if not version or row["version"] == version} | |
| with TestClient(api) as client: | |
| document = client.get("/v1/models").json() | |
| self.assertEqual({row["wire_id"] for row in document["models"]}, expected) | |
| self.assertEqual(document["default"], document["models"][0]["id"]) | |
| for row in document["models"]: | |
| configured = api.state.registry[row["wire_id"]] | |
| self.assertEqual(row["revision"], configured["revision"]) | |
| self.assertEqual(row["manifest_sha256"], configured["manifest_sha256"]) | |
| self.assertEqual(len(api.state.registry), 12) | |
| self.assertEqual(len(api.state.relays), 12) | |
| self.assertEqual(len(client.get("/api/ready").json()["models"]), 12) | |
| catalog = client.get("/api/tetris/config").json()["competitors"] | |
| expected_repos = {PROFILES[wire]["repo_id"] for wire in expected} | |
| self.assertEqual({row["id"] for row in catalog if row["family"] == "decision"}, | |
| {row.id for row in LOCAL_MODELS if row.request_model in expected_repos}) | |
| self.assertEqual({row["id"] for row in catalog if row["family"] == "cloud"}, | |
| set() if version == "2.0" else {"jev-cloud", "jev-cloud-mirror"}) | |
| self.assertNotIn("test-upstream-secret", repr(catalog)) | |
| self.assertNotIn("test-cloud-secret", repr(catalog)) | |
| if version == "2.0": | |
| self.assertTrue(all(row["family"] == "decision" for row in catalog)) | |
| self.assertNotIn("jev-cloud", api.state.tetris.adapter._endpoints) | |
| self.assertNotIn("jev-cloud-mirror", api.state.tetris.adapter._endpoints) | |
| for cloud in ("jev-cloud", "jev-cloud-mirror"): | |
| race = client.post("/api/tetris/races", json={ | |
| "left": cloud, "right": cloud, "mode": "steps", "max_steps": 1, | |
| }) | |
| self.assertEqual(race.status_code, 422, race.text) | |
| self.assertEqual(api.state.tetris.sessions, {}) | |
| else: | |
| self.assertTrue(all(row["ready"] for row in catalog if row["family"] == "cloud")) | |
| status = client.get("/api/status", params={"model": MODEL}) | |
| self.assertEqual(status.status_code, 200) | |
| self.assertEqual(status.json()["model"], MODEL) | |
| class SharedGatewayCompatibilityTests(unittest.IsolatedAsyncioTestCase): | |
| async def test_decision1_single_and_batch_still_complete_behind_decision2_view(self): | |
| entries = generation_registry("2.0") + generation_registry("1.0") | |
| entry = next(row for row in entries if row["id"] == MODEL) | |
| canonical = PROFILES[MODEL]["repo_id"] | |
| runtime = PROFILES[MODEL]["runtime_model"] | |
| worker_fields = { | |
| "model": MODEL, | |
| "manifest_sha256": entry["manifest_sha256"], | |
| "worker_id": "1" * 32, | |
| } | |
| headers = {"Authorization": "Bearer " + TOKEN} | |
| questions = { | |
| "route": {"type": "choice", "instructions": "Choose a route.", | |
| "criteria": {"allow": None, "deny": None}}, | |
| "accept": {"type": "noul", "instructions": "Is this accepted?"}, | |
| "urgency": {"type": "score", "instructions": "Rate urgency.", | |
| "criteria": ["Low", "Medium", "High"]}, | |
| } | |
| answers = { | |
| "route": {"type": "choice", "choice": "allow", "confidence": 0.5, | |
| "probabilities": {"allow": 0.75, "deny": 0.25}}, | |
| "accept": {"type": "noul", "noul": 0.8}, | |
| "urgency": {"type": "score", "score": 1.0, "confidence": 0.4, | |
| "probabilities": {"0": 0.2, "1": 0.6, "2": 0.2}, | |
| "legend": {"0": "Low", "1": "Medium", "2": "High"}}, | |
| } | |
| with patch.dict(os.environ, { | |
| "DECISION_WORKER_TOKEN": TOKEN, | |
| "DECISION_STUDIO_GENERATION": "2.0", | |
| }, clear=True): | |
| api = create_app(mode="pull_queue", registry=entries) | |
| try: | |
| async with httpx.AsyncClient(transport=httpx.ASGITransport(app=api), | |
| base_url="http://testserver") as client: | |
| heartbeat = await client.post("/internal/worker/heartbeat", headers=headers, | |
| json={**worker_fields, "phase": "ready", | |
| "capabilities": ["context_batch_v1"]}) | |
| self.assertEqual(heartbeat.status_code, 200) | |
| for batch in (False, True): | |
| with self.subTest(batch=batch): | |
| data = ({"states": [{"id": "a", "state": "One"}, | |
| {"id": "b", "state": "Two"}]} | |
| if batch else {"state": "One"}) | |
| route = "/v1/systemone/batches" if batch else "/v1/systemone" | |
| submitted = asyncio.create_task(client.post(route, json={ | |
| "model": canonical, "questions": questions, **data, | |
| })) | |
| try: | |
| claimed = await client.post("/internal/worker/claim", headers=headers, | |
| json={**worker_fields, "wait_seconds": 1}) | |
| self.assertEqual(claimed.status_code, 200) | |
| job = claimed.json()["job"] | |
| self.assertIsNotNone(job) | |
| self.assertEqual(job["body"]["model"], MODEL) | |
| response = { | |
| "model": runtime, | |
| "usage": {"input_tokens": 120 if batch else 60, "output_tokens": 0}, | |
| } | |
| if batch: | |
| response["results"] = [ | |
| {"id": row["id"], "answers": answers, | |
| "usage": {"input_tokens": 60, "output_tokens": 0}} | |
| for row in data["states"] | |
| ] | |
| else: | |
| response["answers"] = answers | |
| completed = await client.post("/internal/worker/result", headers=headers, | |
| json={**worker_fields, "id": job["id"], | |
| "lease_token": job["lease_token"], | |
| "result": {"kind": "http_runtime_v1", | |
| "response": response}}) | |
| self.assertEqual(completed.status_code, 200, completed.text) | |
| public = await asyncio.wait_for(submitted, timeout=3) | |
| self.assertEqual(public.status_code, 200, public.text) | |
| self.assertEqual(public.json(), dict(response, model=canonical)) | |
| finally: | |
| if not submitted.done(): | |
| submitted.cancel() | |
| with suppress(asyncio.CancelledError): | |
| await submitted | |
| finally: | |
| await api.state.tetris.aclose() | |
| if __name__ == "__main__": | |
| unittest.main() | |