From 3826815edbb69c7445ee3799a8d4885ef8d29106 Mon Sep 17 00:00:00 2001
From: kingbri <bdashore3@proton.me>
Date: Mon, 22 Jul 2024 21:33:10 -0400
Subject: [PATCH] API: Add request logging

Log all the parts of a request if the config flag is set. The logged
fields are all server side anyways, so nothing is being exposed to
clients.

Signed-off-by: kingbri <bdashore3@proton.me>
---
 common/config.py     |  5 +----
 common/networking.py | 32 ++++++++++++++++++++++++++++-
 config_sample.yml    |  4 ++++
 endpoints/server.py  | 48 +++++++++++++++++++++++---------------------
 4 files changed, 61 insertions(+), 28 deletions(-)

diff --git a/common/config.py b/common/config.py
index 4de7d6b..972b382 100644
--- a/common/config.py
+++ b/common/config.py
@@ -51,10 +51,7 @@ def from_args(args: dict):
         cur_logging_config = logging_config()
         GLOBAL_CONFIG["logging"] = {
             **cur_logging_config,
-            **{
-                k.replace("log_", ""): logging_override[k]
-                for k in logging_override
-            },
+            **{k.replace("log_", ""): logging_override[k] for k in logging_override},
         }
 
     developer_override = args.get("developer")
diff --git a/common/networking.py b/common/networking.py
index 706cb60..7c088a9 100644
--- a/common/networking.py
+++ b/common/networking.py
@@ -1,9 +1,10 @@
 """Common utility functions"""
 
 import asyncio
+import json
 import socket
 import traceback
-from fastapi import HTTPException, Request
+from fastapi import Depends, HTTPException, Request
 from loguru import logger
 from pydantic import BaseModel
 from typing import Optional
@@ -108,3 +109,32 @@ async def add_request_id(request: Request):
 
     request.state.id = uuid4().hex
     return request
+
+
+async def log_request(request: Request):
+    """FastAPI depends to log a request to the user."""
+
+    log_message = [f"Information for {request.method} request {request.state.id}:"]
+
+    log_message.append(f"URL: {request.url}")
+    log_message.append(f"Headers: {dict(request.headers)}")
+
+    if request.method != "GET":
+        body_bytes = await request.body()
+        if body_bytes:
+            body = json.loads(body_bytes.decode("utf-8"))
+
+            log_message.append(f"Body: {dict(body)}")
+
+    logger.info("\n".join(log_message))
+
+
+def get_global_depends():
+    """Returns global dependencies for a FastAPI app."""
+
+    depends = [Depends(add_request_id)]
+
+    if config.logging_config().get("requests"):
+        depends.append(Depends(log_request))
+
+    return depends
diff --git a/config_sample.yml b/config_sample.yml
index 3bde5ee..12458de 100644
--- a/config_sample.yml
+++ b/config_sample.yml
@@ -31,6 +31,10 @@ logging:
   # Enable generation parameter logging (default: False)
   generation_params: False
 
+  # Enable request logging (default: False)
+  # NOTE: Only use this for debugging!
+  requests: False
+
 # Options for sampling
 sampling:
   # Override preset name. Find this in the sampler-overrides folder (default: None)
diff --git a/endpoints/server.py b/endpoints/server.py
index a69ecd1..8c10b63 100644
--- a/endpoints/server.py
+++ b/endpoints/server.py
@@ -1,42 +1,44 @@
 import uvicorn
-from fastapi import Depends, FastAPI
+from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from loguru import logger
 
 from common.logger import UVICORN_LOG_CONFIG
-from common.networking import add_request_id
+from common.networking import get_global_depends
 from endpoints.OAI.router import router as OAIRouter
 
-app = FastAPI(
-    title="TabbyAPI",
-    summary="An OAI compatible exllamav2 API that's both lightweight and fast",
-    description=(
-        "This docs page is not meant to send requests! Please use a service "
-        "like Postman or a frontend UI."
-    ),
-    dependencies=[Depends(add_request_id)],
-)
-
-# ALlow CORS requests
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
 
 def setup_app():
     """Includes the correct routers for startup"""
 
+    app = FastAPI(
+        title="TabbyAPI",
+        summary="An OAI compatible exllamav2 API that's both lightweight and fast",
+        description=(
+            "This docs page is not meant to send requests! Please use a service "
+            "like Postman or a frontend UI."
+        ),
+        dependencies=get_global_depends(),
+    )
+
+    # ALlow CORS requests
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
     app.include_router(OAIRouter)
 
+    return app
+
 
 def export_openapi():
     """Function to return the OpenAPI JSON from the API server"""
 
-    setup_app()
+    app = setup_app()
     return app.openapi()
 
 
@@ -49,7 +51,7 @@ async def start_api(host: str, port: int):
     logger.info(f"Chat completions: http://{host}:{port}/v1/chat/completions")
 
     # Setup app
-    setup_app()
+    app = setup_app()
 
     config = uvicorn.Config(
         app,