Spaces:
Paused
Paused
| from dataclasses import replace | |
| import pytest | |
| from conftest import GOOD_SCENE | |
| from api import budget, routing | |
| from api.llm import ScriptedLLM | |
| from api.logstore import LogStore | |
| from api.simulate import generate_scene | |
| def fenced(xml): | |
| return f"```xml\n{xml}\n```" | |
| def test_code_complexity_picks_the_tier(settings): | |
| simple = routing.plan_code("Publish hello on /chatter every second", "humble", "python", settings) | |
| medium = routing.plan_code( | |
| "A service client that reads a parameter and publishes odometry", "jazzy", "python", settings | |
| ) | |
| complex_ = routing.plan_code( | |
| "A Nav2 behavior tree node that calls MoveIt for pick and place", "jazzy", "cpp", settings | |
| ) | |
| assert (simple.tier, medium.tier, complex_.tier) == ("low", "mid", "high") | |
| assert simple.model == "deepseek/deepseek-v4-flash" and complex_.model == "z-ai/glm-5.2" | |
| assert simple.category == "coding" and "navigation" in complex_.reason | |
| def test_scene_plan_and_simulator_choice(settings): | |
| p = routing.plan_scene("A box on the floor", "unitree_g1", settings) | |
| assert p.tier == "mid" and p.simulator == "mujoco" # scenes never use the cheapest tier | |
| assert routing.plan_scene("A door that opens on a hinge", None, settings).tier == "high" | |
| assert routing.plan_scene("How often does it fall on random slopes?", None, settings).simulator == "mujoco_warp" | |
| assert routing.plan_scene("A photoreal kitchen for camera images", None, settings).simulator == "isaac_sim" | |
| assert not routing.SIMULATORS["gazebo"]["available"] | |
| def test_repairs_escalate_to_a_stronger_model(store, settings): | |
| bad = GOOD_SCENE.replace('mass="3"', 'mass="3" frictoin="1"') | |
| llm = ScriptedLLM([fenced(bad), fenced(GOOD_SCENE)]) | |
| plan = routing.plan_scene("G1 next to a crate", "unitree_g1", settings) | |
| result = generate_scene("G1 next to a crate", "unitree_g1", llm, store, settings, plan=plan) | |
| assert result.valid and llm.models == ["z-ai/glm-5.2", "moonshotai/kimi-k2.6"] # scene models | |
| assert [a.model for a in result.attempts] == llm.models and result.model == "moonshotai/kimi-k2.6" | |
| def test_estimates_grow_with_tier_and_backend(settings): | |
| oc = replace(settings, codegen_backend="opencode") | |
| low = routing.plan_code("Publish hello", "humble", "python", oc).estimate_inr | |
| high = routing.plan_code("Nav2 planner with MoveIt and SLAM", "humble", "python", oc).estimate_inr | |
| assert 0 < low[0] <= low[1] and high[0] > low[0] | |
| direct = routing.plan_code("Publish hello", "humble", "python", replace(settings, codegen_backend="direct")) | |
| assert direct.estimate_inr[0] < low[0] # OpenCode's agent loop reads more context | |
| # --------------------------------------------------------------------------- rupee budget | |
| def spend(logs, settings, tester, usd, cost=None, model="deepseek/deepseek-v4-pro", category="coding"): | |
| spec = routing.model_for_tier("mid") | |
| usage = {"prompt_tokens": 1000, "completion_tokens": 100} | |
| if cost is not None: | |
| usage["cost"] = cost | |
| elif usd: | |
| usage["cost"] = usd | |
| budget.record_call(logs, settings, tester=tester, generation_id="g", category=category, model=spec, usage=usage) | |
| def test_budget_counts_billable_calls_only(tmp_path, settings): | |
| logs = LogStore(tmp_path / "l.db") | |
| spend(logs, settings, "alice", 0.5) # OpenRouter reported $0.50 | |
| spend(logs, settings, "alice", 0, cost=0) # OpenRouter-exempt: reported no charge | |
| spend(logs, settings, "bob", 0.25, category="simulation_ai") | |
| s = budget.summary(logs, settings, "alice") | |
| assert s["used"] == pytest.approx(0.5 * settings.usd_inr) and s["exempt_calls"] == 1 | |
| assert s["by_category"]["coding"] == s["used"] and s["remaining"] == pytest.approx(125 - s["used"]) | |
| assert s["pool"]["used"] == pytest.approx(0.75 * settings.usd_inr) and s["pool"]["total"] == 500 | |
| free = routing.ModelSpec("some/model:free", "low", "x", 0, 0) | |
| assert budget.cost_of_call(free, {"prompt_tokens": 10}).exempt_reason.startswith("free model") | |
| listed = budget.cost_of_call(routing.model_for_tier("low"), {"prompt_tokens": 1_000_000, "completion_tokens": 0}) | |
| assert listed.estimated and listed.cost_usd == pytest.approx(0.089) | |
| def test_budget_check_refuses_and_asks(tmp_path, settings): | |
| logs = LogStore(tmp_path / "l.db") | |
| budget.check(logs, settings, "alice", (1.0, 3.0), confirmed=False) # cheap: no question | |
| with pytest.raises(budget.BudgetError) as e: | |
| budget.check(logs, settings, "alice", (4.0, 12.0), confirmed=False) | |
| assert e.value.status == 409 and e.value.detail["needs_confirmation"] | |
| budget.check(logs, settings, "alice", (4.0, 12.0), confirmed=True) | |
| spend(logs, settings, "alice", 124.5 / settings.usd_inr) | |
| with pytest.raises(budget.BudgetError) as e: | |
| budget.check(logs, settings, "alice", (1.0, 2.0), confirmed=True) | |
| assert e.value.status == 402 and "Your AI budget" in str(e.value) | |
| for name in ("bob", "carol", "dave"): | |
| spend(logs, settings, name, 125 / settings.usd_inr) | |
| with pytest.raises(budget.BudgetError) as e: | |
| budget.check(logs, settings, "erin", (1.0, 2.0), confirmed=True) | |
| assert "team's AI budget" in str(e.value) | |
| # --------------------------------------------------------------------------- through the API | |
| def test_generate_reports_routing_and_cost(run_env): | |
| client, keys_, backend, services, gen_id = run_env | |
| h = {"X-API-Key": keys_["alice"]} | |
| services.llm_factory().responses[:] = [fenced(GOOD_SCENE)] | |
| r = client.post("/simulate/generate", json={"description": "G1 near a crate and a ramp"}, headers=h).json() | |
| assert r["routing"]["tier"] == "mid" and r["routing"]["simulator"] == "mujoco" and not r["routing"]["escalated"] | |
| assert r["cost"]["inr"] >= 0 and r["cost"]["budget"]["allocation"] == 125 | |
| usage = client.get("/me/usage", headers=h).json() | |
| assert usage["ai"]["calls"] >= 2 and usage["ai"]["pool"]["total"] == 500 and not usage["admin"] | |
| g = services.logs.get(r["id"]) | |
| assert g["request"]["routing"]["model"] == "z-ai/glm-5.2" | |
| def test_expensive_task_needs_confirmation(run_env): | |
| client, keys_, backend, services, gen_id = run_env | |
| services.settings = replace(services.settings, confirm_above_inr=0.0001) | |
| h = {"X-API-Key": keys_["alice"]} | |
| r = client.post("/simulate/generate", json={"description": "G1 near a crate and a ramp"}, headers=h) | |
| assert r.status_code == 409 and r.json()["detail"]["needs_confirmation"] | |
| services.llm_factory().responses[:] = [fenced(GOOD_SCENE)] | |
| body = {"description": "G1 near a crate and a ramp", "confirm_cost": True} | |
| assert client.post("/simulate/generate", json=body, headers=h).status_code == 200 | |
| def test_feedback_is_stored_with_context(run_env): | |
| client, keys_, backend, services, gen_id = run_env | |
| h = {"X-API-Key": keys_["alice"]} | |
| body = {"target_type": "generation", "target_id": gen_id, "rating": "not_helpful", "note": "box is floating"} | |
| assert client.post("/feedback", json=body, headers=h).json() == {"ok": True} | |
| (fb,) = services.logs.feedback_rows() | |
| assert fb["rating"] == "not_helpful" and fb["note"] == "box is floating" and fb["context"]["kind"] == "simulate" | |
| assert services.logs.get(gen_id)["tester_rating"] == "broken" | |
| other = {"X-API-Key": keys_["bob"]} | |
| assert client.post("/feedback", json=body, headers=other).status_code == 404 | |
| def test_admin_sees_the_team(run_env): | |
| client, keys_, backend, services, gen_id = run_env | |
| services.settings = replace(services.settings, admin_testers=("alice",)) | |
| team = client.get("/me/usage", headers={"X-API-Key": keys_["alice"]}).json()["team"] | |
| assert {t["tester"] for t in team} >= {"alice", "bob"} and all(t["allocation"] == 125 for t in team) | |
| # --------------------------------------------------------------------------- free debugger | |
| class FlakyFree: | |
| """Stands in for ZenLLM: answers from a script, or fails like a rate-limited free model.""" | |
| model = "opencode/deepseek-v4-flash-free" | |
| def __init__(self, responses, fail=False): | |
| self.responses, self.fail, self.calls = list(responses), fail, 0 | |
| def complete(self, messages, temperature=0.2, model=None): | |
| from api.llm import Completion, LLMError | |
| self.calls += 1 | |
| if self.fail: | |
| raise LLMError("free model: 429 rate limited") | |
| return Completion(self.responses.pop(0), model, 0, {"input_tokens": 900, "output_tokens": 400, "cost": 0.0}) | |
| def test_first_repair_goes_to_the_free_debugger(store, settings): | |
| from api.zen import WithFreeModels | |
| s = replace(settings, free_debug=True) | |
| bad = GOOD_SCENE.replace('mass="3"', 'mass="3" frictoin="1"') | |
| paid, free = ScriptedLLM([fenced(bad)]), FlakyFree([fenced(GOOD_SCENE)]) | |
| plan = routing.plan_scene("G1 next to a crate", "unitree_g1", s) | |
| assert plan.debugger == "opencode/deepseek-v4-flash-free" | |
| result = generate_scene("G1 next to a crate", "unitree_g1", WithFreeModels(paid, free), store, s, plan=plan) | |
| assert result.valid and paid.models == ["z-ai/glm-5.2"] and free.calls == 1 | |
| assert [a.model for a in result.attempts] == ["z-ai/glm-5.2", "opencode/deepseek-v4-flash-free"] | |
| def test_free_debugger_failure_falls_back_to_a_paid_tier(store, settings): | |
| from api.zen import WithFreeModels | |
| s = replace(settings, free_debug=True) | |
| bad = GOOD_SCENE.replace('mass="3"', 'mass="3" frictoin="1"') | |
| paid, free = ScriptedLLM([fenced(bad), fenced(GOOD_SCENE)]), FlakyFree([], fail=True) | |
| plan = routing.plan_scene("G1 next to a crate", "unitree_g1", s) | |
| result = generate_scene("G1 next to a crate", "unitree_g1", WithFreeModels(paid, free), store, s, plan=plan) | |
| assert result.valid and free.calls == 1 # tried, failed, and the same attempt went to the paid model | |
| assert paid.models == ["z-ai/glm-5.2", "moonshotai/kimi-k2.6"] and len(result.attempts) == 2 | |
| def test_attempt_schedule_and_worst_case(settings): | |
| s = replace(settings, free_debug=True) | |
| plan = routing.plan_code("Publish hello on /chatter", "humble", "python", s) | |
| ids = [[m.id for m in routing.candidates_for_attempt(plan, i)] for i in range(3)] | |
| assert ids == [ | |
| ["deepseek/deepseek-v4-flash"], | |
| ["opencode/deepseek-v4-flash-free", "deepseek/deepseek-v4-pro"], | |
| ["z-ai/glm-5.2"], | |
| ] | |
| # the estimate assumes the free call fails, so its paid fallback is in the worst case | |
| tin, tout = routing.TOKENS[("code", "direct")] | |
| paid = [routing.model_for_attempt(plan, i) for i in range(3)] | |
| worst = sum(routing.call_cost_usd(m, tin * (1 + i // 2), tout) for i, m in enumerate(paid)) * s.usd_inr | |
| assert plan.estimate_inr[1] == pytest.approx(worst, abs=0.01) | |
| def test_free_calls_are_recorded_but_never_billed(tmp_path, settings): | |
| free = routing.ModelSpec("opencode/deepseek-v4-flash-free", "free", "OpenCode Zen", 0, 0, free=True) | |
| logs = LogStore(tmp_path / "l.sqlite3") | |
| budget.record_call(logs, settings, tester="a", generation_id="g", category="debugging", model=free, | |
| usage={"input_tokens": 5000, "output_tokens": 900}) # fmt: skip | |
| (row,) = logs.usage_rows() | |
| assert row["billable"] == 0 and row["provider"] == "opencode-zen" and row["exempt_reason"] == "opencode_zen_free" | |
| assert budget.summary(logs, settings, "a")["used"] == 0 | |