From 2fc1b006022fb08581c544f56357ef017476acf2 Mon Sep 17 00:00:00 2001
From: Pavlo Paliychuk <paul@getzep.com>
Date: Mon, 23 Sep 2024 10:12:35 -0400
Subject: [PATCH] feat: add FastAPI lifespan and healthcheck endpoint (#144)

* chore: Add healthcheck endpoint + build indexes and constraints on svc startup

* chore: Bring back driver close call
---
 graphiti_core/graphiti.py            |  2 +-
 server/graph_service/main.py         | 23 +++++++++++++++++++----
 server/graph_service/zep_graphiti.py |  9 +++++++++
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/graphiti_core/graphiti.py b/graphiti_core/graphiti.py
index da66c0c1..bf6f65e9 100644
--- a/graphiti_core/graphiti.py
+++ b/graphiti_core/graphiti.py
@@ -158,8 +158,8 @@ def close(self):
                 # Use graphiti...
             finally:
                 graphiti.close()
-        self.driver.close()
         """
+        self.driver.close()
 
     async def build_indices_and_constraints(self):
         """
diff --git a/server/graph_service/main.py b/server/graph_service/main.py
index 487c1a01..e85638ec 100644
--- a/server/graph_service/main.py
+++ b/server/graph_service/main.py
@@ -1,14 +1,29 @@
+from contextlib import asynccontextmanager
+
 from fastapi import FastAPI
+from fastapi.responses import JSONResponse
 
+from graph_service.config import get_settings
 from graph_service.routers import ingest, retrieve
+from graph_service.zep_graphiti import initialize_graphiti
+
+
+@asynccontextmanager
+async def lifespan(_: FastAPI):
+    settings = get_settings()
+    await initialize_graphiti(settings)
+    yield
+    # Shutdown
+    # No need to close Graphiti here, as it's handled per-request
+
 
-app = FastAPI()
+app = FastAPI(lifespan=lifespan)
 
 
 app.include_router(retrieve.router)
 app.include_router(ingest.router)
 
 
-@app.get('/')
-def read_root():
-    return {'Hello': 'World'}
+@app.get('/healthcheck')
+async def healthcheck():
+    return JSONResponse(content={'status': 'healthy'}, status_code=200)
diff --git a/server/graph_service/zep_graphiti.py b/server/graph_service/zep_graphiti.py
index 6f85079b..13054f37 100644
--- a/server/graph_service/zep_graphiti.py
+++ b/server/graph_service/zep_graphiti.py
@@ -82,6 +82,15 @@ async def get_graphiti(settings: ZepEnvDep):
         client.close()
 
 
+async def initialize_graphiti(settings: ZepEnvDep):
+    client = ZepGraphiti(
+        uri=settings.neo4j_uri,
+        user=settings.neo4j_user,
+        password=settings.neo4j_password,
+    )
+    await client.build_indices_and_constraints()
+
+
 def get_fact_result_from_edge(edge: EntityEdge):
     return FactResult(
         uuid=edge.uuid,