Skip to content

Commit 9344cdf

Browse files
committed
feat: add status check
1 parent a755156 commit 9344cdf

File tree

3 files changed

+41
-1
lines changed

3 files changed

+41
-1
lines changed

src/server/api/api.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,11 @@ def custom_openapi():
106106
"/healthcheck", tags=["chore"], openapi_extra=API_X_CODE_DOCS["GET /healthcheck"]
107107
)(api_layer.chore.healthcheck)
108108

109+
router.get(
110+
"/admin/status_check",
111+
tags=["admin"],
112+
# openapi_extra=API_X_CODE_DOCS["GET /admin/status_check"],
113+
)(api_layer.chore.root_running_status_check)
109114

110115
router.post(
111116
"/project/profile_config",

src/server/api/memobase_server/api_layer/chore.py

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,11 @@
1+
import traceback
2+
from fastapi import HTTPException, Request
13
from ..env import LOG
24
from ..models.response import BaseResponse, CODE
5+
from ..models.database import DEFAULT_PROJECT_ID
36
from ..connectors import db_health_check, redis_health_check
4-
from fastapi import HTTPException
7+
from ..llms.embeddings import check_embedding_sanity
8+
from ..llms import llm_sanity_check
59

610

711
async def healthcheck() -> BaseResponse:
@@ -18,3 +22,32 @@ async def healthcheck() -> BaseResponse:
1822
detail="Redis not available",
1923
)
2024
return BaseResponse()
25+
26+
27+
async def root_running_status_check(request: Request) -> BaseResponse:
28+
"""Check if your memobase is set up correctly"""
29+
project_id = request.state.memobase_project_id
30+
if project_id != DEFAULT_PROJECT_ID:
31+
return BaseResponse(
32+
errno=CODE.METHOD_NOT_ALLOWED, errmsg="Only Root can access this"
33+
)
34+
try:
35+
if not db_health_check():
36+
return BaseResponse(
37+
errno=CODE.INTERNAL_SERVER_ERROR,
38+
errmsg="Database not available",
39+
)
40+
if not await redis_health_check():
41+
return BaseResponse(
42+
errno=CODE.INTERNAL_SERVER_ERROR,
43+
errmsg="Redis not available",
44+
)
45+
await check_embedding_sanity()
46+
await llm_sanity_check()
47+
except Exception as e:
48+
return BaseResponse(
49+
errno=CODE.INTERNAL_SERVER_ERROR,
50+
errmsg=f"Occuring Error when status checking: {e}\n{traceback.format_exc()}",
51+
)
52+
53+
return BaseResponse()

src/server/api/memobase_server/api_layer/middleware.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020

2121

2222
PATH_MAPPINGS = [
23+
"/api/v1/admin/status_check",
24+
"/api/v1/users/blobs",
2325
"/api/v1/users/blobs",
2426
"/api/v1/users/profile",
2527
"/api/v1/users/buffer",

0 commit comments

Comments
 (0)