Spaces:
Sleeping
Sleeping
File size: 6,348 Bytes
ea9e5f8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | 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"
|