f793a1fd44
Co-authored-by: Kevin Lin <klin5061@gmail.com> Co-authored-by: Matthew Zhou <mattzh1314@gmail.com> Co-authored-by: Kian Jones <11655409+kianjones9@users.noreply.github.com> Co-authored-by: Andy Li <55300002+cliandy@users.noreply.github.com> Co-authored-by: jnjpng <jin@letta.com> Co-authored-by: Jin Peng <jinjpeng@Jins-MacBook-Pro.local> Co-authored-by: Eric Ly <111820150+lyeric2022@users.noreply.github.com> Co-authored-by: Eric Ly <lyyeric@letta.com> Co-authored-by: Shubham Naik <shub@letta.com> Co-authored-by: Shubham Naik <shub@memgpt.ai>
26 lines
859 B
Python
26 lines
859 B
Python
from starlette.middleware.base import BaseHTTPMiddleware
|
|
|
|
|
|
class ProfilerContextMiddleware(BaseHTTPMiddleware):
|
|
"""Middleware to set context if using profiler. Currently just uses google-cloud-profiler."""
|
|
|
|
async def dispatch(self, request, call_next):
|
|
ctx = None
|
|
if request.url.path in {"/v1/health", "/v1/health/"}:
|
|
return await call_next(request)
|
|
try:
|
|
labels = {
|
|
"method": request.method,
|
|
"path": request.url.path,
|
|
"endpoint": request.url.path,
|
|
}
|
|
import googlecloudprofiler
|
|
|
|
ctx = googlecloudprofiler.context.set_labels(**labels)
|
|
except:
|
|
return await call_next(request)
|
|
if ctx:
|
|
with ctx:
|
|
return await call_next(request)
|
|
return await call_next(request)
|