Skip to content

Commit 65ad2d2

Browse files
noahpodgurskinoahpodgurski
authored
chore: Include litellm v readiness (#180)
* feat: include litellm version in /readiness * feat: add fixture to mock FXA scopes in user signup tests * addr comments * cache litellm_version after first request
1 parent 66cf9ef commit 65ad2d2

7 files changed

Lines changed: 82 additions & 44 deletions

File tree

‎src/mlpa/core/config.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -320,6 +320,7 @@ def __init__(self):
320320
env = Env()
321321

322322
LITELLM_READINESS_URL = f"{env.LITELLM_API_BASE}/health/readiness"
323+
LITELLM_INFO_URL = f"{env.LITELLM_API_BASE}/public/model_hub/info"
323324
LITELLM_COMPLETIONS_URL = f"{env.LITELLM_API_BASE}/v1/chat/completions"
324325
LITELLM_SEARCH_URL = f"{env.LITELLM_API_BASE}/v1/search"
325326
LITELLM_MASTER_AUTH_HEADERS = {

‎src/mlpa/core/routers/health/health.py‎

Lines changed: 29 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,35 @@
22

33
from fastapi import APIRouter
44

5-
from mlpa.core.config import LITELLM_MASTER_AUTH_HEADERS, LITELLM_READINESS_URL
5+
from mlpa.core.config import (
6+
LITELLM_INFO_URL,
7+
LITELLM_MASTER_AUTH_HEADERS,
8+
LITELLM_READINESS_URL,
9+
)
610
from mlpa.core.http_client import get_http_client
711
from mlpa.core.pg_services.services import app_attest_pg, litellm_pg
812

913
mlpa_version = importlib.metadata.version("mlpa")
14+
litellm_version = "N/A"
1015
router = APIRouter()
1116

1217

18+
async def get_litellm_version(client):
19+
global litellm_version
20+
21+
if litellm_version != "N/A":
22+
return litellm_version
23+
24+
try:
25+
response = await client.get(LITELLM_INFO_URL, timeout=3)
26+
litellm_info = response.json()
27+
except Exception:
28+
return litellm_version
29+
30+
litellm_version = litellm_info.get("litellm_version", "N/A")
31+
return litellm_version
32+
33+
1334
@router.get("/liveness", tags=["Health"])
1435
async def liveness_probe():
1536
return {"status": "alive"}
@@ -20,19 +41,22 @@ async def readiness_probe():
2041
# todo add check to PG and LiteLLM status here
2142
pg_status = litellm_pg.check_status()
2243
app_attest_pg_status = app_attest_pg.check_status()
23-
litellm_status = {}
2444
client = get_http_client()
2545
response = await client.get(
2646
LITELLM_READINESS_URL, headers=LITELLM_MASTER_AUTH_HEADERS, timeout=3
2747
)
28-
data = response.json()
29-
litellm_status = data
48+
litellm_status = response.json()
49+
current_litellm_version = await get_litellm_version(client)
50+
3051
return {
3152
"status": "connected",
3253
"mlpa_version": mlpa_version,
3354
"pg_server_dbs": {
3455
"postgres": "connected" if pg_status else "offline",
3556
"app_attest": "connected" if app_attest_pg_status else "offline",
3657
},
37-
"litellm": litellm_status,
58+
"litellm": {
59+
"litellm_version": current_litellm_version,
60+
**litellm_status,
61+
},
3862
}

‎src/tests/conftest.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,12 @@ def mock_request():
1717
def _force_mlpa_debug_false():
1818
monkeypatch = pytest.MonkeyPatch()
1919
monkeypatch.setenv("MLPA_DEBUG", "false")
20+
monkeypatch.setenv("ADDITIONAL_FXA_SCOPE_1", None)
21+
monkeypatch.setenv("ADDITIONAL_FXA_SCOPE_2", None)
22+
monkeypatch.setenv("ADDITIONAL_FXA_SCOPE_3", None)
2023
env.MLPA_DEBUG = False
24+
env.ADDITIONAL_FXA_SCOPE_1 = None
25+
env.ADDITIONAL_FXA_SCOPE_2 = None
26+
env.ADDITIONAL_FXA_SCOPE_3 = None
2127
yield
2228
monkeypatch.undo()

‎src/tests/integration/test_fxa.py‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ def test_ai_missing_purpose_is_allowed_by_default(mocked_client_integration):
3030
"authorization": f"Bearer {TEST_FXA_TOKEN}",
3131
"service-type": "ai",
3232
},
33-
json=SAMPLE_REQUEST.dict(),
33+
json=SAMPLE_REQUEST.model_dump(),
3434
)
3535
assert response.status_code == 200
3636
assert response.json() == SUCCESSFUL_CHAT_RESPONSE
@@ -45,7 +45,7 @@ def test_ai_invalid_purpose_returns_400(mocked_client_integration):
4545
"service-type": "ai",
4646
"purpose": "invalid-purpose",
4747
},
48-
json=SAMPLE_REQUEST.dict(),
48+
json=SAMPLE_REQUEST.model_dump(),
4949
)
5050
assert response.status_code == 400
5151
assert "purpose" in str(response.json().get("detail", "")).lower()
@@ -59,7 +59,7 @@ def test_invalid_fxa_auth(mocked_client_integration):
5959
"service-type": "ai",
6060
"purpose": "chat",
6161
},
62-
json=SAMPLE_REQUEST.dict(),
62+
json=SAMPLE_REQUEST.model_dump(),
6363
)
6464
assert response.status_code == 401
6565

@@ -72,7 +72,7 @@ def test_successful_request_with_mocked_fxa_auth(mocked_client_integration):
7272
"service-type": "ai",
7373
"purpose": "chat",
7474
},
75-
json=SAMPLE_REQUEST.dict(),
75+
json=SAMPLE_REQUEST.model_dump(),
7676
)
7777
assert response.status_code != 401
7878
assert response.status_code != 400
@@ -93,7 +93,7 @@ def test_x_dev_authorization_success(mocked_client_integration):
9393
"purpose": "chat",
9494
"x-dev-authorization": DEV_TOKEN,
9595
},
96-
json=SAMPLE_REQUEST.dict(),
96+
json=SAMPLE_REQUEST.model_dump(),
9797
)
9898
assert response.status_code == 200
9999
assert response.json() == SUCCESSFUL_CHAT_RESPONSE
@@ -112,7 +112,7 @@ def test_x_dev_authorization_missing_fxa(mocked_client_integration):
112112
"purpose": "chat",
113113
"x-dev-authorization": DEV_TOKEN,
114114
},
115-
json=SAMPLE_REQUEST.dict(),
115+
json=SAMPLE_REQUEST.model_dump(),
116116
)
117117
assert response.status_code == 422
118118

@@ -131,7 +131,7 @@ def test_x_dev_authorization_invalid_token(mocked_client_integration):
131131
"purpose": "chat",
132132
"x-dev-authorization": "wrong-token",
133133
},
134-
json=SAMPLE_REQUEST.dict(),
134+
json=SAMPLE_REQUEST.model_dump(),
135135
)
136136
assert response.status_code == 401
137137

@@ -149,7 +149,7 @@ def test_x_dev_authorization_token_not_configured(mocked_client_integration):
149149
"purpose": "chat",
150150
"x-dev-authorization": "some-token",
151151
},
152-
json=SAMPLE_REQUEST.dict(),
152+
json=SAMPLE_REQUEST.model_dump(),
153153
)
154154
assert response.status_code == 422
155155

@@ -163,7 +163,7 @@ def test_ai_dev_requires_x_dev_authorization(mocked_client_integration):
163163
"service-type": "ai-dev",
164164
"purpose": "chat",
165165
},
166-
json=SAMPLE_REQUEST.dict(),
166+
json=SAMPLE_REQUEST.model_dump(),
167167
)
168168
assert response.status_code == 401
169169
assert "x-dev-authorization required" in str(response.json().get("detail", ""))
@@ -186,7 +186,7 @@ def test_x_dev_authorization_ignored_for_non_dev_service_type(
186186
"purpose": "chat",
187187
"x-dev-authorization": DEV_TOKEN,
188188
},
189-
json=SAMPLE_REQUEST.dict(),
189+
json=SAMPLE_REQUEST.model_dump(),
190190
)
191191
assert response.status_code == 200
192192
assert response.json() == SUCCESSFUL_CHAT_RESPONSE
@@ -206,7 +206,7 @@ def test_x_dev_authorization_token_not_configured_with_fxa(mocked_client_integra
206206
"purpose": "chat",
207207
"x-dev-authorization": "some-token",
208208
},
209-
json=SAMPLE_REQUEST.dict(),
209+
json=SAMPLE_REQUEST.model_dump(),
210210
)
211211
assert response.status_code == 401
212212
assert "Invalid x-dev-authorization" in str(response.json().get("detail", ""))

‎src/tests/integration/test_health.py‎

Lines changed: 17 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import importlib.metadata
2+
13
from mlpa.core.config import env
24

35

@@ -8,39 +10,33 @@ def test_health_liveness(mocked_client_integration, httpx_mock):
810

911

1012
def test_health_readiness(mocked_client_integration, httpx_mock):
13+
mlpa_version = importlib.metadata.version("mlpa")
1114
httpx_mock.add_response(
1215
method="GET",
1316
url=f"{env.LITELLM_API_BASE}/health/readiness",
1417
status_code=200,
1518
json={
16-
"status": "connected",
17-
"pg_server_dbs": {"postgres": "connected", "app_attest": "connected"},
18-
"litellm": {
19-
"status": "connected",
20-
"db": "connected",
21-
"cache": None,
22-
"litellm_version": "1.77.3",
23-
"success_callbacks": [
24-
"sync_deployment_callback_on_success",
25-
"_PROXY_VirtualKeyModelMaxBudgetLimiter",
26-
"_ProxyDBLogger",
27-
"_PROXY_MaxBudgetLimiter",
28-
"_PROXY_MaxParallelRequestsHandler_v3",
29-
"_PROXY_CacheControlCheck",
30-
"_PROXY_LiteLLMManagedFiles",
31-
"ServiceLogging",
32-
],
33-
"use_aiohttp_transport": True,
34-
"last_updated": "2025-10-10T00:00:00",
35-
},
19+
"status": "healthy",
20+
"db": "connected",
3621
},
3722
)
23+
httpx_mock.add_response(
24+
method="GET",
25+
url=f"{env.LITELLM_API_BASE}/public/model_hub/info",
26+
status_code=200,
27+
json={"litellm_version": "1.84.4"},
28+
)
3829

3930
readiness_response = mocked_client_integration.get("/health/readiness")
4031
assert readiness_response.status_code == 200
4132
assert readiness_response.json().get("status") == "connected"
33+
assert readiness_response.json().get("mlpa_version") == mlpa_version
4234
assert readiness_response.json().get("pg_server_dbs") is not None
43-
assert readiness_response.json().get("litellm") is not None
35+
assert readiness_response.json().get("litellm") == {
36+
"litellm_version": "1.84.4",
37+
"status": "healthy",
38+
"db": "connected",
39+
}
4440

4541

4642
def test_metrics_endpoint(mocked_client_integration):

‎src/tests/integration/test_security_headers.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,13 @@ def test_security_headers_on_all_endpoints(mocked_client_integration, httpx_mock
5050
method="GET",
5151
url=f"{env.LITELLM_API_BASE}/health/readiness",
5252
status_code=200,
53-
json={"status": "connected"},
53+
json={"status": "healthy", "db": "connected"},
54+
)
55+
httpx_mock.add_response(
56+
method="GET",
57+
url=f"{env.LITELLM_API_BASE}/public/model_hub/info",
58+
status_code=200,
59+
json={"litellm_version": "1.84.4"},
5460
)
5561

5662
endpoints = [

‎src/tests/integration/test_user_signup_cap.py‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,11 @@ def use_real_get_or_create_user():
1212
return True
1313

1414

15+
@pytest.fixture(autouse=True)
16+
def single_fxa_scope(mocker):
17+
mocker.patch("mlpa.core.auth.fxa.FXA_SCOPES", ("profile:uid",))
18+
19+
1520
class _FakeResponse:
1621
def __init__(self, payload: dict):
1722
self._payload = payload
@@ -101,7 +106,7 @@ def verify_token_side_effect(
101106
"authorization": f"Bearer {TEST_FXA_TOKEN}",
102107
"service-type": "ai",
103108
},
104-
json=SAMPLE_REQUEST.dict(),
109+
json=SAMPLE_REQUEST.model_dump(),
105110
)
106111
assert resp1.status_code == 200
107112
assert resp1.json() == SUCCESSFUL_CHAT_RESPONSE
@@ -112,7 +117,7 @@ def verify_token_side_effect(
112117
"authorization": f"Bearer {TEST_FXA_TOKEN}",
113118
"service-type": "s2s",
114119
},
115-
json=SAMPLE_REQUEST.dict(),
120+
json=SAMPLE_REQUEST.model_dump(),
116121
)
117122
assert resp2.status_code == 200
118123
assert resp2.json() == SUCCESSFUL_CHAT_RESPONSE
@@ -123,7 +128,7 @@ def verify_token_side_effect(
123128
"authorization": f"Bearer {TEST_FXA_TOKEN}",
124129
"service-type": "ai",
125130
},
126-
json=SAMPLE_REQUEST.dict(),
131+
json=SAMPLE_REQUEST.model_dump(),
127132
)
128133
assert resp3.status_code == 403
129134
assert resp3.json()["detail"]["error"] == 4
@@ -139,7 +144,7 @@ def verify_token_side_effect(
139144
"authorization": f"Bearer {TEST_FXA_TOKEN}",
140145
"service-type": "memories",
141146
},
142-
json=SAMPLE_REQUEST.dict(),
147+
json=SAMPLE_REQUEST.model_dump(),
143148
)
144149
assert resp4.status_code == 200
145150
assert resp4.json() == SUCCESSFUL_CHAT_RESPONSE
@@ -189,7 +194,7 @@ def verify_token_side_effect(
189194
"authorization": f"Bearer {TEST_FXA_TOKEN}",
190195
"service-type": "ai",
191196
},
192-
json=SAMPLE_REQUEST.dict(),
197+
json=SAMPLE_REQUEST.model_dump(),
193198
)
194199
assert resp1.status_code == 500
195200

@@ -199,7 +204,7 @@ def verify_token_side_effect(
199204
"authorization": f"Bearer {TEST_FXA_TOKEN}",
200205
"service-type": "ai",
201206
},
202-
json=SAMPLE_REQUEST.dict(),
207+
json=SAMPLE_REQUEST.model_dump(),
203208
)
204209
assert resp2.status_code == 200
205210
assert resp2.json() == SUCCESSFUL_CHAT_RESPONSE

0 commit comments

Comments
 (0)