From 06c6b54529cfa8c1d1dde812d53dd926d944b8ca Mon Sep 17 00:00:00 2001
From: Yogish Baliga <yogish@together.ai>
Date: Sat, 21 Sep 2024 13:07:55 -0700
Subject: [PATCH] adding support health check to deploy it in k8s

---
 llama_stack/distribution/server/server.py           | 13 ++++++++++++-
 .../providers/adapters/safety/together/safety.py    |  1 +
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/llama_stack/distribution/server/server.py b/llama_stack/distribution/server/server.py
index 38218ab8b..bbd6bfc38 100644
--- a/llama_stack/distribution/server/server.py
+++ b/llama_stack/distribution/server/server.py
@@ -15,6 +15,7 @@ from collections.abc import (
     AsyncIterator as AsyncIteratorABC,
 )
 from contextlib import asynccontextmanager
+from http import HTTPStatus
 from ssl import SSLError
 from typing import (
     Any,
@@ -88,7 +89,7 @@ async def global_exception_handler(request: Request, exc: Exception):
     )
 
 
-def translate_exception(exc: Exception) -> HTTPException:
+def translate_exception(exc: Exception) -> Union[HTTPException, RequestValidationError]:
     if isinstance(exc, ValidationError):
         exc = RequestValidationError(exc.raw_errors)
 
@@ -407,12 +408,22 @@ async def resolve_impls_with_routing(run_config: StackRunConfig) -> Dict[Api, An
     return impls, specs
 
 
+
 def main(yaml_config: str, port: int = 5000, disable_ipv6: bool = False):
     with open(yaml_config, "r") as fp:
         config = StackRunConfig(**yaml.safe_load(fp))
 
     app = FastAPI()
 
+    # Health check is added to enable deploying the docker container image on Kubernetes which require
+    # a health check that can return 200 for readiness and liveness check
+    class HealthCheck(BaseModel):
+        status: str = "OK"
+
+    @app.get("/healthcheck", status_code=HTTPStatus.OK, response_model=HealthCheck)
+    async def healthcheck():
+        return HealthCheck(status="OK")
+
     impls, specs = asyncio.run(resolve_impls_with_routing(config))
     if Api.telemetry in impls:
         setup_logger(impls[Api.telemetry])
diff --git a/llama_stack/providers/adapters/safety/together/safety.py b/llama_stack/providers/adapters/safety/together/safety.py
index b5e28098b..f89d361e9 100644
--- a/llama_stack/providers/adapters/safety/together/safety.py
+++ b/llama_stack/providers/adapters/safety/together/safety.py
@@ -3,6 +3,7 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
+import pydantic
 from together import Together
 
 import asyncio