feat: add endpoint metrics (#96)

* feat: add endpoint metrics

* fix: unit test

---------

Co-authored-by: leoguillaume <[email protected]>

app/endpoints/files.py

@@ -5,6 +5,7 @@
 from app.schemas.security import User
 from app.utils.lifespan import clients
 from app.utils.security import check_api_key
+from app.utils.exceptions import FileSizeLimitExceededException
 
 router = APIRouter()
 
@@ -20,10 +21,13 @@ async def upload_file(file: UploadFile = File(...), request: FilesRequest = Body
         For JSON, file structure like a list of documents: [{"text": "hello world", "title": "my document", "metadata": {"autor": "me"}}, ...]} or [{"text": "hello world", "title": "my document"}, ...]}
         Each document must have a "text" and "title" keys and "metadata" key (optional) with dict type value.
     - html: Hypertext Markup Language file.
-
-    Max file size is 10MB.
     """
 
+    file_size = len(file.file.read())
+    if file_size > FileSizeLimitExceededException.MAX_CONTENT_SIZE:
+        raise FileSizeLimitExceededException()
+    file.file.seek(0)  # reset file pointer to the beginning of the file
+
     if request.chunker:
         chunker_args = request.chunker.args.model_dump() if request.chunker.args else ChunkerArgs().model_dump()
         chunker_name = request.chunker.name

app/helpers/__init__.py

@@ -1,9 +1,9 @@
 from ._authenticationclient import AuthenticationClient
 from ._clientsmanager import ClientsManager
-from ._contentsizelimitmiddleware import ContentSizeLimitMiddleware
 from ._fileuploader import FileUploader
 from ._internetclient import InternetClient
 from ._modelclients import ModelClients
 from .searchclients import SearchClient
+from ._metricsmiddleware import MetricsMiddleware
 
-__all__ = ["AuthenticationClient", "ClientsManager", "ContentSizeLimitMiddleware", "FileUploader", "InternetClient", "ModelClients", "SearchClient"]
+__all__ = ["AuthenticationClient", "ClientsManager", "ContentSizeLimitMiddleware", "FileUploader", "InternetClient", "ModelClients", "SearchClient", "MetricsMiddleware"]

app/helpers/_authenticationclient.py

@@ -1,23 +1,21 @@
+import base64
 import datetime as dt
+import hashlib
 import json
 from typing import Optional
 import uuid
+from typing import Any, Callable
 
 from grist_api import GristDocAPI
 from redis import Redis
 
-from app.utils.variables import ROLE_LEVEL_0, ROLE_LEVEL_1, ROLE_LEVEL_2
+from app.schemas.security import Role, User
 
 
 class AuthenticationClient(GristDocAPI):
     CACHE_EXPIRATION = 3600  # 1h
-    ROLE_DICT = {
-        "user": ROLE_LEVEL_0,
-        "client": ROLE_LEVEL_1,
-        "admin": ROLE_LEVEL_2,
-    }
 
-    def __init__(self, cache: Redis, table_id: str, *args, **kwargs):
+    def __init__(self, cache: Redis, table_id: str, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.session_id = str(uuid.uuid4())
         self.table_id = table_id
@@ -35,14 +33,14 @@ def check_api_key(self, key: str) -> Optional[str]:
         """
         keys = self._get_api_keys()
         if key in keys:
-            return keys[key]
+            return User(id=self._api_key_to_user_id(input=key), role=Role[keys[key]["role"]], name=keys[key]["name"])
 
-    def cache(func):
+    def cache(func) -> Callable[..., Any]:
         """
         Decorator to cache the result of a function in Redis.
         """
 
-        def wrapper(self):
+        def wrapper(self) -> Any:
             key = f"auth-{self.session_id}"
             result = self.redis.get(key)
             if result:
@@ -56,18 +54,40 @@ def wrapper(self):
         return wrapper
 
     @cache
-    def _get_api_keys(self):
+    def _get_api_keys(self) -> dict:
         """
         Get all keys from a table in the Grist document.
 
         Returns:
             dict: dictionary of keys and their corresponding access level
         """
-        records = self.fetch_table(self.table_id)
+        records = self.fetch_table(table_name=self.table_id)
 
         keys = dict()
         for record in records:
             if record.EXPIRATION > dt.datetime.now().timestamp():
-                keys[record.KEY] = self.ROLE_DICT.get(record.ROLE, ROLE_LEVEL_0)
+                keys[record.KEY] = {
+                    "id": self._api_key_to_user_id(input=record.KEY),
+                    "role": Role.get(name=record.ROLE.upper(), default=Role.USER)._name_,
+                    "name": record.USER,
+                }
 
         return keys
+
+    @staticmethod
+    def _api_key_to_user_id(input: str) -> str:
+        """
+        Generate a 16 length unique code from an input string using salted SHA-256 hashing.
+
+        Args:
+            input_string (str): The input string to generate the code from.
+
+        Returns:
+            tuple[str, bytes]: A tuple containing the generated code and the salt used.
+        """
+        hash = hashlib.sha256((input).encode()).digest()
+        hash = base64.urlsafe_b64encode(hash).decode()
+        # remove special characters and limit length
+        hash = "".join(c for c in hash if c.isalnum())[:16].lower()
+
+        return hash

app/helpers/_metricsmiddleware.py

@@ -0,0 +1,39 @@
+import json
+
+from prometheus_client import Counter
+
+from app.helpers._authenticationclient import AuthenticationClient
+from fastapi import Request, Response
+from starlette.middleware.base import BaseHTTPMiddleware
+
+
+class MetricsMiddleware(BaseHTTPMiddleware):
+    # TODO: add audio endpoint (support for multipart/form-data)
+    MODELS_ENDPOINTS = ["/v1/chat/completions", "/v1/completions", "/v1/embeddings"]
+    http_requests_by_user = Counter(
+        name="http_requests_by_user_and_endpoint",
+        documentation="Number of HTTP requests by user and endpoint",
+        labelnames=["user", "endpoint", "model"],
+    )
+
+    async def dispatch(self, request: Request, call_next) -> Response:
+        endpoint = request.url.path
+        content_type = request.headers.get("Content-Type", "")
+
+        if endpoint.startswith("/v1"):
+            authorization = request.headers.get("Authorization")
+            model = None
+            if not content_type.startswith("multipart/form-data"):
+                body = await request.body()
+                body = body.decode(encoding="utf-8")
+                model = json.loads(body).get("model") if body else None
+
+            user_id = AuthenticationClient._api_key_to_user_id(input=authorization.split(sep=" ")[1])
+
+            if authorization and authorization.startswith("Bearer "):
+                user_id = AuthenticationClient._api_key_to_user_id(input=authorization.split(sep=" ")[1])
+                self.http_requests_by_user.labels(user=user_id, endpoint=endpoint[3:], model=model).inc()
+
+        response = await call_next(request)
+
+        return response

app/helpers/searchclients/_elasticsearchclient.py

@@ -9,21 +9,21 @@
 from app.schemas.collections import Collection
 from app.schemas.documents import Document
 from app.schemas.chunks import Chunk
+from app.schemas.security import Role
 from app.schemas.security import User
 from app.schemas.search import Filter, Search
 from app.utils.exceptions import (
     DifferentCollectionsModelsException,
-    WrongCollectionTypeException,
     WrongModelTypeException,
     CollectionNotFoundException,
+    InsufficientRightsException,
     SearchMethodNotAvailableException,
 )
 from app.utils.variables import (
     EMBEDDINGS_MODEL_TYPE,
     HYBRID_SEARCH_TYPE,
     LEXICAL_SEARCH_TYPE,
     SEMANTIC_SEARCH_TYPE,
-    ROLE_LEVEL_2,
     PUBLIC_COLLECTION_TYPE,
     PRIVATE_COLLECTION_TYPE,
 )
@@ -71,8 +71,8 @@ def __init__(self, models: List[str] = None, hybrid_limit_factor: float = 1.5, *
     def upsert(self, chunks: List[Chunk], collection_id: str, user: Optional[User] = None) -> None:
         collection = self.get_collections(collection_ids=[collection_id], user=user)[0]
 
-        if user.role != ROLE_LEVEL_2 and collection.type == PUBLIC_COLLECTION_TYPE:
-            raise WrongCollectionTypeException()
+        if user.role != Role.ADMIN and collection.type == PUBLIC_COLLECTION_TYPE:
+            raise InsufficientRightsException()
 
         for i in range(0, len(chunks), self.BATCH_SIZE):
             batched_chunks = chunks[i : i + self.BATCH_SIZE]
@@ -168,8 +168,8 @@ def create_collection(
         if self.models[collection_model].type != EMBEDDINGS_MODEL_TYPE:
             raise WrongModelTypeException()
 
-        if user.role != ROLE_LEVEL_2 and collection_type == PUBLIC_COLLECTION_TYPE:
-            raise WrongCollectionTypeException()
+        if user.role != Role.ADMIN and collection_type == PUBLIC_COLLECTION_TYPE:
+            raise InsufficientRightsException()
 
         settings = {
             "similarity": {"default": {"type": "BM25"}},
@@ -224,8 +224,8 @@ def delete_collection(self, collection_id: str, user: User) -> None:
         """
         collection = self.get_collections(collection_ids=[collection_id], user=user)[0]
 
-        if user.role != ROLE_LEVEL_2 and collection.type == PUBLIC_COLLECTION_TYPE:
-            raise WrongCollectionTypeException()
+        if user.role != Role.ADMIN and collection.type == PUBLIC_COLLECTION_TYPE:
+            raise InsufficientRightsException()
 
         self.indices.delete(index=collection_id, ignore_unavailable=True)
 
@@ -279,8 +279,8 @@ def delete_document(self, collection_id: str, document_id: str, user: Optional[U
         """
         collection = self.get_collections(collection_ids=[collection_id], user=user)[0]
 
-        if user.role != ROLE_LEVEL_2 and collection.type == PUBLIC_COLLECTION_TYPE:
-            raise WrongCollectionTypeException()
+        if user.role != Role.ADMIN and collection.type == PUBLIC_COLLECTION_TYPE:
+            raise InsufficientRightsException()
 
         # delete chunks
         body = {"query": {"match": {"metadata.document_id": document_id}}}

app/helpers/searchclients/_qdrantsearchclient.py

@@ -20,20 +20,20 @@
 from app.schemas.collections import Collection
 from app.schemas.documents import Document
 from app.schemas.search import Search
+from app.schemas.security import Role
 from app.schemas.security import User
 from app.utils.exceptions import (
     CollectionNotFoundException,
     DifferentCollectionsModelsException,
     SearchMethodNotAvailableException,
-    WrongCollectionTypeException,
     WrongModelTypeException,
+    InsufficientRightsException,
 )
 from app.utils.variables import (
     EMBEDDINGS_MODEL_TYPE,
     LEXICAL_SEARCH_TYPE,
     HYBRID_SEARCH_TYPE,
     PUBLIC_COLLECTION_TYPE,
-    ROLE_LEVEL_2,
     SEMANTIC_SEARCH_TYPE,
 )
 
@@ -59,8 +59,8 @@ def upsert(self, chunks: List[Chunk], collection_id: str, user: User) -> None:
         """
         collection = self.get_collections(collection_ids=[collection_id], user=user)[0]
 
-        if user.role != ROLE_LEVEL_2 and collection.type == PUBLIC_COLLECTION_TYPE:
-            raise WrongCollectionTypeException()
+        if user.role != Role.ADMIN and collection.type == PUBLIC_COLLECTION_TYPE:
+            raise InsufficientRightsException()
 
         for i in range(0, len(chunks), self.BATCH_SIZE):
             batch = chunks[i : i + self.BATCH_SIZE]
@@ -214,8 +214,8 @@ def create_collection(
         if self.models[collection_model].type != EMBEDDINGS_MODEL_TYPE:
             raise WrongModelTypeException()
 
-        if user.role != ROLE_LEVEL_2 and collection_type == PUBLIC_COLLECTION_TYPE:
-            raise WrongCollectionTypeException()
+        if user.role != Role.ADMIN and collection_type == PUBLIC_COLLECTION_TYPE:
+            raise InsufficientRightsException()
 
         # create metadata
         metadata = {
@@ -241,8 +241,8 @@ def delete_collection(self, collection_id: str, user: User) -> None:
         """
         collection = self.get_collections(collection_ids=[collection_id], user=user)[0]
 
-        if user.role != ROLE_LEVEL_2 and collection.type == PUBLIC_COLLECTION_TYPE:
-            raise WrongCollectionTypeException()
+        if user.role != Role.ADMIN and collection.type == PUBLIC_COLLECTION_TYPE:
+            raise InsufficientRightsException()
 
         super().delete_collection(collection_name=collection.id)
         super().delete(collection_name=self.METADATA_COLLECTION_ID, points_selector=PointIdsList(points=[collection.id]))
@@ -294,8 +294,8 @@ def delete_document(self, collection_id: str, document_id: str, user: User):
         """
         collection = self.get_collections(collection_ids=[collection_id], user=user)[0]
 
-        if user.role != ROLE_LEVEL_2 and collection.type == PUBLIC_COLLECTION_TYPE:
-            raise WrongCollectionTypeException()
+        if user.role != Role.ADMIN and collection.type == PUBLIC_COLLECTION_TYPE:
+            raise InsufficientRightsException()
 
         # delete chunks
         filter = Filter(must=[FieldCondition(key="metadata.document_id", match=MatchAny(any=[document_id]))])

app/main.py

@@ -1,15 +1,13 @@
-from fastapi import FastAPI, Response, Security
-
-
+from fastapi import Depends, FastAPI, Response, Security
+from prometheus_fastapi_instrumentator import Instrumentator
 from slowapi.middleware import SlowAPIASGIMiddleware
 
 from app.endpoints import audio, chat, chunks, collections, completions, documents, embeddings, files, models, search
-from app.helpers import ContentSizeLimitMiddleware
+from app.helpers._metricsmiddleware import MetricsMiddleware
 from app.schemas.security import User
 from app.utils.settings import settings
 from app.utils.lifespan import lifespan
-from app.utils.security import check_api_key
-
+from app.utils.security import check_admin_api_key, check_api_key
 
 app = FastAPI(
     title=settings.app_name,
@@ -22,9 +20,13 @@
     redoc_url="/documentation",
 )
 
+# Prometheus metrics
+# @TODO: env_var_name="ENABLE_METRICS"
+app.instrumentator = Instrumentator().instrument(app=app)
+
 # Middlewares
-app.add_middleware(middleware_class=ContentSizeLimitMiddleware)
 app.add_middleware(middleware_class=SlowAPIASGIMiddleware)
+app.add_middleware(middleware_class=MetricsMiddleware)
 
 
 # Monitoring
@@ -37,6 +39,8 @@ def health(user: User = Security(dependency=check_api_key)) -> Response:
     return Response(status_code=200)
 
 
+app.instrumentator.expose(app=app, should_gzip=True, tags=["Monitoring"], dependencies=[Depends(dependency=check_admin_api_key)])
+
 # Core
 app.include_router(router=models.router, tags=["Core"], prefix="/v1")
 app.include_router(router=chat.router, tags=["Core"], prefix="/v1")