From b43fef0a82080ff9cf40cf126c7e17ce3b36a61c Mon Sep 17 00:00:00 2001
From: Tyler Galske <tgalske@gmail.com>
Date: Sun, 13 Oct 2024 08:08:54 -0400
Subject: [PATCH] feat: add prometheus metrics to /metrics

---
 CHANGELOG.md                       | 4 ++++
 prepline_general/api/app.py        | 6 +++++-
 prepline_general/api/general.py    | 4 ++--
 preprocessing-pipeline-family.yaml | 2 +-
 requirements/base.in               | 1 +
 requirements/base.txt              | 2 ++
 6 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3c2e8fa9..04d02b73 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.0.82
+
+* Expose Prometheus metrics at /metrics
+
 ## 0.0.81
 
 * Update `strategy` parameter to allow `'` and `"` as input surrounding the value.
diff --git a/prepline_general/api/app.py b/prepline_general/api/app.py
index 60247c30..f5567fa6 100644
--- a/prepline_general/api/app.py
+++ b/prepline_general/api/app.py
@@ -2,6 +2,7 @@
 from fastapi.datastructures import FormData
 from fastapi.responses import JSONResponse
 from fastapi.security import APIKeyHeader
+from prometheus_client import make_asgi_app
 import logging
 import os
 
@@ -13,7 +14,7 @@
 app = FastAPI(
     title="Unstructured Pipeline API",
     summary="Partition documents with the Unstructured library",
-    version="0.0.81",
+    version="0.0.82",
     docs_url="/general/docs",
     openapi_url="/general/openapi.json",
     servers=[
@@ -31,6 +32,9 @@
     openapi_tags=[{"name": "general"}],
 )
 
+metrics_app = make_asgi_app()
+app.mount("/metrics", metrics_app)
+
 # Note(austin) - This logger just dumps exceptions
 # We'd rather handle those below, so disable this in deployments
 uvicorn_logger = logging.getLogger("uvicorn.error")
diff --git a/prepline_general/api/general.py b/prepline_general/api/general.py
index 7623702b..e7252a44 100644
--- a/prepline_general/api/general.py
+++ b/prepline_general/api/general.py
@@ -653,7 +653,7 @@ def return_content_type(filename: str):
 
 
 @router.get("/general/v0/general", include_in_schema=False)
-@router.get("/general/v0.0.81/general", include_in_schema=False)
+@router.get("/general/v0.0.82/general", include_in_schema=False)
 async def handle_invalid_get_request():
     raise HTTPException(
         status_code=status.HTTP_405_METHOD_NOT_ALLOWED, detail="Only POST requests are supported."
@@ -668,7 +668,7 @@ async def handle_invalid_get_request():
     description="Description",
     operation_id="partition_parameters",
 )
-@router.post("/general/v0.0.81/general", include_in_schema=False)
+@router.post("/general/v0.0.82/general", include_in_schema=False)
 def general_partition(
     request: Request,
     # cannot use annotated type here because of a bug described here:
diff --git a/preprocessing-pipeline-family.yaml b/preprocessing-pipeline-family.yaml
index 01a64677..116f88e1 100644
--- a/preprocessing-pipeline-family.yaml
+++ b/preprocessing-pipeline-family.yaml
@@ -1,2 +1,2 @@
 name: general
-version: 0.0.81
+version: 0.0.82
diff --git a/requirements/base.in b/requirements/base.in
index fda4c50f..12e0e090 100644
--- a/requirements/base.in
+++ b/requirements/base.in
@@ -16,3 +16,4 @@ backoff
 pypdf
 pycryptodome
 psutil
+prometheus_client
diff --git a/requirements/base.txt b/requirements/base.txt
index 47934cf0..bddb3eca 100644
--- a/requirements/base.txt
+++ b/requirements/base.txt
@@ -244,6 +244,8 @@ pillow==10.4.0
     #   unstructured-pytesseract
 portalocker==2.10.1
     # via iopath
+prometheus-client==0.21.0
+    # via -r requirements/base.in
 proto-plus==1.24.0
     # via
     #   google-api-core