Spaces:
Sleeping
Sleeping
| import json | |
| import os | |
| from pathlib import Path | |
| import httpx | |
| from fastapi.testclient import TestClient | |
| from service.api import app | |
| from service.policy import Candidate, Limits, select, query_bucket | |
| from fit_policy import fit | |
| def test_closed_form_and_hard_call_cap(): | |
| candidates=[Candidate("small",.7,.002,500,1),Candidate("large",.9,.02,1000,2)] | |
| b=Limits(max_llm_calls=1,deadline_ms=700,max_expected_cost_usd=.1) | |
| result=select(candidates,b) | |
| assert result["selected"]=="small" | |
| assert result["decision_llm_calls"]==0 | |
| assert result["rejected"][0]["reasons"]==["predicted_deadline","call_limit"] | |
| assert select(candidates,Limits(max_llm_calls=0))["selected"]=="STOP" | |
| assert query_bucket("Who invented it?",[.1,.08])=="c0_g0" | |
| def test_calibration_shrinkage_uses_labeled_outcomes(): | |
| rows=[{"bucket":"c0_g0","action":"retrieve_fast","correct":"1","cost_usd":"0.001","latency_ms":"800"}, | |
| {"bucket":"c0_g0","action":"retrieve_fast","correct":"0","cost_usd":"0.003","latency_ms":"1000"}] | |
| obj=fit(rows) | |
| assert obj["actions"]["retrieve_fast"]["global_quality"]==.5 | |
| assert obj["actions"]["retrieve_fast"]["quality_by_bucket"]["c0_g0"]==.5 | |
| assert obj["actions"]["retrieve_fast"]["latency_p95_ms"]==1000 | |
| def test_api_route_and_answer_with_mocked_upstreams(tmp_path,monkeypatch): | |
| monkeypatch.setenv("API_TOKEN","test-only-key") | |
| monkeypatch.setenv("ALLOW_HTTP_ENDPOINTS","1") | |
| monkeypatch.setenv("FAST_LLM_BASE_URL","http://mock.local/v1") | |
| monkeypatch.setenv("FAST_LLM_MODEL","fast-model") | |
| monkeypatch.setenv("FAST_INPUT_USD_PER_1M","0.1") | |
| monkeypatch.setenv("FAST_OUTPUT_USD_PER_1M","0.2") | |
| monkeypatch.setenv("STRONG_LLM_BASE_URL","http://mock.local/v1") | |
| monkeypatch.setenv("STRONG_LLM_MODEL","strong-model") | |
| monkeypatch.setenv("STRONG_INPUT_USD_PER_1M","1") | |
| monkeypatch.setenv("STRONG_OUTPUT_USD_PER_1M","2") | |
| monkeypatch.setenv("RETRIEVAL_URL","http://mock.local/search") | |
| monkeypatch.setenv("RETRIEVAL_USD_PER_CALL","0.0001") | |
| actions=str(Path(__file__).resolve().parents[1]/"config/actions.json") | |
| monkeypatch.setenv("ACTIONS_FILE",actions) | |
| cal={"format_version":1,"gap_threshold":.07,"actions":{ | |
| "retrieve_fast":{"global_quality":.6,"expected_cost_usd":.001,"latency_p95_ms":200}, | |
| "retrieve_strong":{"global_quality":.8,"expected_cost_usd":.003,"latency_p95_ms":300}, | |
| "iterative_strong":{"global_quality":.92,"expected_cost_usd":.006,"latency_p95_ms":600} | |
| }} | |
| f=tmp_path/"policy.json";f.write_text(json.dumps(cal));monkeypatch.setenv("CALIBRATION_FILE",str(f)) | |
| events=[] | |
| def handler(req:httpx.Request): | |
| events.append((req.url.path,req.headers.get("authorization"),req.content.decode())) | |
| if req.url.path=="/search": | |
| return httpx.Response(200,json={"documents":[ | |
| {"id":"1","title":"Tokyo","text":"Tokyo is the capital of Japan","score":.91}, | |
| {"id":"2","title":"Japan","text":"Japan is in East Asia","score":.21}]}) | |
| if req.url.path=="/v1/chat/completions": | |
| data=json.loads(req.content) | |
| text="Tokyo" if "Answer the user's" in data["messages"][0]["content"] else "Tokyo Japan" | |
| return httpx.Response(200,json={"choices":[{"message":{"content":text}}], | |
| "usage":{"prompt_tokens":40,"completion_tokens":2}}) | |
| return httpx.Response(404) | |
| with TestClient(app) as client: | |
| app.state.remote.http=httpx.AsyncClient(transport=httpx.MockTransport(handler)) | |
| assert client.get("/health").status_code==200 | |
| assert client.get("/ready").json()["calibration_loaded"] is True | |
| route_payload={"candidates":[ | |
| {"id":"large","predicted_accuracy":.95,"expected_cost_usd":.003,"latency_p95_ms":400,"llm_calls":2}, | |
| {"id":"small","predicted_accuracy":.7,"expected_cost_usd":.001,"latency_p95_ms":100,"llm_calls":1}], | |
| "budgets":{"max_llm_calls":1}} | |
| assert client.post("/v1/route",json=route_payload).status_code==401 | |
| auth={"Authorization":"Bearer test-only-key"} | |
| assert client.post("/v1/route",json=route_payload,headers=auth).json()["selected"]=="small" | |
| req={"question":"What is the capital of Japan?","budgets":{"deadline_ms":2000,"max_llm_calls":1,"max_retrieval_calls":1,"max_expected_cost_usd":.01}} | |
| r=client.post("/v1/answer",headers=auth,json=req) | |
| assert r.status_code==200,r.text | |
| j=r.json() | |
| assert j["selected_action"]=="retrieve_strong" | |
| assert j["answer"]=="Tokyo" | |
| assert j["resources"]["model_calls"]==1 | |
| assert j["resources"]["retrieval_calls"]==1 | |
| assert j["resources"]["unreported_usage_calls"]==0 | |
| assert j["resources"]["measured_cost_usd"]>0 | |
| assert j["resources"]["retrieval_cost_usd"]==.0001 | |
| assert j["resources"]["total_estimated_cost_usd"]>0 | |
| req["budgets"]["max_llm_calls"]=2 | |
| req["budgets"]["max_retrieval_calls"]=2 | |
| r2=client.post("/v1/answer",headers=auth,json=req) | |
| assert r2.status_code==200,r2.text | |
| j2=r2.json() | |
| assert j2["selected_action"]=="iterative_strong" | |
| assert j2["resources"]["model_calls"]==2 | |
| assert j2["resources"]["retrieval_calls"]==2 | |
| assert all(x[0] in ("/v1/chat/completions","/search") for x in events) | |
| def test_answer_respects_zero_retrieval_limit(tmp_path,monkeypatch): | |
| monkeypatch.setenv("API_TOKEN","t") | |
| monkeypatch.setenv("ALLOW_HTTP_ENDPOINTS","1") | |
| monkeypatch.setenv("FAST_LLM_BASE_URL","http://mock.local/v1") | |
| monkeypatch.setenv("FAST_LLM_MODEL","fast") | |
| monkeypatch.setenv("STRONG_LLM_BASE_URL","http://mock.local/v1") | |
| monkeypatch.setenv("STRONG_LLM_MODEL","strong") | |
| monkeypatch.setenv("RETRIEVAL_URL","http://mock.local/search") | |
| monkeypatch.delenv("CALIBRATION_FILE",raising=False) | |
| monkeypatch.setenv("ACTIONS_FILE",str(Path(__file__).resolve().parents[1]/"config/actions.json")) | |
| with TestClient(app) as client: | |
| response=client.post("/v1/answer",headers={"Authorization":"Bearer t"}, | |
| json={"question":"Hello","budgets":{"max_retrieval_calls":0,"deadline_ms":1000}}) | |
| assert response.status_code==200 | |
| assert response.json()["retrieval_calls"]==0 | |
| assert response.json()["selected_action"]=="STOP" | |