Feature/request id middleware 2 (#4427)

* stubbing out request id

* passthru or create request id's in api and model server

* add onyx request id

* get request id logging into uvicorn

* no logs

* change prefixes

* fix comment

* docker image needs specific shared files

---------

Co-authored-by: Richard Kuo (Onyx) <[email protected]>

Author: rkuo-danswer

backend/Dockerfile.model_server

@@ -46,6 +46,7 @@ WORKDIR /app
 
 # Utils used by model server
 COPY ./onyx/utils/logger.py /app/onyx/utils/logger.py
+COPY ./onyx/utils/middleware.py /app/onyx/utils/middleware.py
 
 # Place to fetch version information
 COPY ./onyx/__init__.py /app/onyx/__init__.py

backend/model_server/main.py

@@ -1,3 +1,4 @@
+import logging
 import os
 import shutil
 from collections.abc import AsyncGenerator
@@ -20,6 +21,8 @@
 from model_server.utils import get_gpu_type
 from onyx import __version__
 from onyx.utils.logger import setup_logger
+from onyx.utils.logger import setup_uvicorn_logger
+from onyx.utils.middleware import add_onyx_request_id_middleware
 from shared_configs.configs import INDEXING_ONLY
 from shared_configs.configs import MIN_THREADS_ML_MODELS
 from shared_configs.configs import MODEL_SERVER_ALLOWED_HOST
@@ -36,6 +39,12 @@
 
 logger = setup_logger()
 
+file_handlers = [
+    h for h in logger.logger.handlers if isinstance(h, logging.FileHandler)
+]
+
+setup_uvicorn_logger(shared_file_handlers=file_handlers)
+
 
 def _move_files_recursively(source: Path, dest: Path, overwrite: bool = False) -> None:
     """
@@ -112,6 +121,12 @@ def get_model_app() -> FastAPI:
     application.include_router(encoders_router)
     application.include_router(custom_models_router)
 
+    request_id_prefix = "INF"
+    if INDEXING_ONLY:
+        request_id_prefix = "IDX"
+
+    add_onyx_request_id_middleware(application, request_id_prefix, logger)
+
     return application
 
 

backend/onyx/main.py

@@ -1,3 +1,4 @@
+import logging
 import sys
 import traceback
 from collections.abc import AsyncGenerator
@@ -102,6 +103,8 @@
 from onyx.setup import setup_multitenant_onyx
 from onyx.setup import setup_onyx
 from onyx.utils.logger import setup_logger
+from onyx.utils.logger import setup_uvicorn_logger
+from onyx.utils.middleware import add_onyx_request_id_middleware
 from onyx.utils.telemetry import get_or_generate_uuid
 from onyx.utils.telemetry import optional_telemetry
 from onyx.utils.telemetry import RecordType
@@ -116,6 +119,12 @@
 
 logger = setup_logger()
 
+file_handlers = [
+    h for h in logger.logger.handlers if isinstance(h, logging.FileHandler)
+]
+
+setup_uvicorn_logger(shared_file_handlers=file_handlers)
+
 
 def validation_exception_handler(request: Request, exc: Exception) -> JSONResponse:
     if not isinstance(exc, RequestValidationError):
@@ -421,6 +430,8 @@ def get_application() -> FastAPI:
     if LOG_ENDPOINT_LATENCY:
         add_latency_logging_middleware(application, logger)
 
+    add_onyx_request_id_middleware(application, "API", logger)
+
     # Ensure all routes have auth enabled or are explicitly marked as public
     check_router_auth(application)
 

backend/onyx/utils/logger.py

@@ -13,6 +13,7 @@
 from shared_configs.configs import SLACK_CHANNEL_ID
 from shared_configs.configs import TENANT_ID_PREFIX
 from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
+from shared_configs.contextvars import ONYX_REQUEST_ID_CONTEXTVAR
 
 
 logging.addLevelName(logging.INFO + 5, "NOTICE")
@@ -71,6 +72,14 @@ def get_log_level_from_str(log_level_str: str = LOG_LEVEL) -> int:
     return log_level_dict.get(log_level_str.upper(), logging.getLevelName("NOTICE"))
 
 
+class OnyxRequestIDFilter(logging.Filter):
+    def filter(self, record: logging.LogRecord) -> bool:
+        from shared_configs.contextvars import ONYX_REQUEST_ID_CONTEXTVAR
+
+        record.request_id = ONYX_REQUEST_ID_CONTEXTVAR.get() or "-"
+        return True
+
+
 class OnyxLoggingAdapter(logging.LoggerAdapter):
     def process(
         self, msg: str, kwargs: MutableMapping[str, Any]
@@ -103,6 +112,7 @@ def process(
                 msg = f"[CC Pair: {cc_pair_id}] {msg}"
 
             break
+
         # Add tenant information if it differs from default
         # This will always be the case for authenticated API requests
         if MULTI_TENANT:
@@ -115,6 +125,11 @@ def process(
                 )
                 msg = f"[t:{short_tenant}] {msg}"
 
+        # request id within a fastapi route
+        fastapi_request_id = ONYX_REQUEST_ID_CONTEXTVAR.get()
+        if fastapi_request_id:
+            msg = f"[{fastapi_request_id}] {msg}"
+
         # For Slack Bot, logs the channel relevant to the request
         channel_id = self.extra.get(SLACK_CHANNEL_ID) if self.extra else None
         if channel_id:
@@ -165,6 +180,14 @@ def format(self, record: logging.LogRecord) -> str:
         return super().format(record)
 
 
+def get_uvicorn_standard_formatter() -> ColoredFormatter:
+    """Returns a standard colored logging formatter."""
+    return ColoredFormatter(
+        "%(asctime)s %(filename)30s %(lineno)4s: [%(request_id)s] %(message)s",
+        datefmt="%m/%d/%Y %I:%M:%S %p",
+    )
+
+
 def get_standard_formatter() -> ColoredFormatter:
     """Returns a standard colored logging formatter."""
     return ColoredFormatter(
@@ -201,12 +224,6 @@ def setup_logger(
 
     logger.addHandler(handler)
 
-    uvicorn_logger = logging.getLogger("uvicorn.access")
-    if uvicorn_logger:
-        uvicorn_logger.handlers = []
-        uvicorn_logger.addHandler(handler)
-        uvicorn_logger.setLevel(log_level)
-
     is_containerized = is_running_in_container()
     if LOG_FILE_NAME and (is_containerized or DEV_LOGGING_ENABLED):
         log_levels = ["debug", "info", "notice"]
@@ -225,14 +242,37 @@ def setup_logger(
             file_handler.setFormatter(formatter)
             logger.addHandler(file_handler)
 
-            if uvicorn_logger:
-                uvicorn_logger.addHandler(file_handler)
-
     logger.notice = lambda msg, *args, **kwargs: logger.log(logging.getLevelName("NOTICE"), msg, *args, **kwargs)  # type: ignore
 
     return OnyxLoggingAdapter(logger, extra=extra)
 
 
+def setup_uvicorn_logger(
+    log_level: int = get_log_level_from_str(),
+    shared_file_handlers: list[logging.FileHandler] | None = None,
+) -> None:
+    uvicorn_logger = logging.getLogger("uvicorn.access")
+    if not uvicorn_logger:
+        return
+
+    formatter = get_uvicorn_standard_formatter()
+
+    handler = logging.StreamHandler()
+    handler.setLevel(log_level)
+    handler.setFormatter(formatter)
+
+    uvicorn_logger.handlers = []
+    uvicorn_logger.addHandler(handler)
+    uvicorn_logger.setLevel(log_level)
+    uvicorn_logger.addFilter(OnyxRequestIDFilter())
+
+    if shared_file_handlers:
+        for fh in shared_file_handlers:
+            uvicorn_logger.addHandler(fh)
+
+    return
+
+
 def print_loggers() -> None:
     """Print information about all loggers. Use to debug logging issues."""
     root_logger = logging.getLogger()

backend/onyx/utils/middleware.py

@@ -0,0 +1,62 @@
+import base64
+import hashlib
+import logging
+import uuid
+from collections.abc import Awaitable
+from collections.abc import Callable
+from datetime import datetime
+from datetime import timezone
+
+from fastapi import FastAPI
+from fastapi import Request
+from fastapi import Response
+
+from shared_configs.contextvars import ONYX_REQUEST_ID_CONTEXTVAR
+
+
+def add_onyx_request_id_middleware(
+    app: FastAPI, prefix: str, logger: logging.LoggerAdapter
+) -> None:
+    @app.middleware("http")
+    async def set_request_id(
+        request: Request, call_next: Callable[[Request], Awaitable[Response]]
+    ) -> Response:
+        """Generate a request hash that can be used to track the lifecycle
+        of a request.  The hash is prefixed to help indicated where the request id
+        originated.
+
+        Format is f"{PREFIX}:{ID}" where PREFIX is 3 chars and ID is 8 chars.
+        Total length is 12 chars.
+        """
+
+        onyx_request_id = request.headers.get("X-Onyx-Request-ID")
+        if not onyx_request_id:
+            onyx_request_id = make_randomized_onyx_request_id(prefix)
+
+        ONYX_REQUEST_ID_CONTEXTVAR.set(onyx_request_id)
+        return await call_next(request)
+
+
+def make_randomized_onyx_request_id(prefix: str) -> str:
+    """generates a randomized request id"""
+
+    hash_input = str(uuid.uuid4())
+    return _make_onyx_request_id(prefix, hash_input)
+
+
+def make_structured_onyx_request_id(prefix: str, request_url: str) -> str:
+    """Not used yet, but could be in the future!"""
+    hash_input = f"{request_url}:{datetime.now(timezone.utc)}"
+    return _make_onyx_request_id(prefix, hash_input)
+
+
+def _make_onyx_request_id(prefix: str, hash_input: str) -> str:
+    """helper function to return an id given a string input"""
+    hash_obj = hashlib.md5(hash_input.encode("utf-8"))
+    hash_bytes = hash_obj.digest()[:6]  # Truncate to 6 bytes
+
+    # 6 bytes becomes 8 bytes. we shouldn't need to strip but just in case
+    # NOTE: possible we'll want more input bytes if id's aren't unique enough
+    hash_str = base64.urlsafe_b64encode(hash_bytes).decode("utf-8").rstrip("=")
+    onyx_request_id = f"{prefix}:{hash_str}"
+    return onyx_request_id

backend/requirements/default.txt

@@ -95,4 +95,4 @@ urllib3==2.2.3
 mistune==0.8.4
 sentry-sdk==2.14.0
 prometheus_client==0.21.0
-fastapi-limiter==0.1.6
+fastapi-limiter==0.1.6

backend/shared_configs/configs.py

@@ -58,6 +58,7 @@
 
 # The process needs to have this for the log file to write to
 # otherwise, it will not create additional log files
+# This should just be the filename base without extension or path.
 LOG_FILE_NAME = os.environ.get("LOG_FILE_NAME") or "onyx"
 
 # Enable generating persistent log files for local dev environments

backend/shared_configs/contextvars.py

@@ -11,6 +11,15 @@
     "current_tenant_id", default=None if MULTI_TENANT else POSTGRES_DEFAULT_SCHEMA
 )
 
+# set by every route in the API server
+INDEXING_REQUEST_ID_CONTEXTVAR: contextvars.ContextVar[
+    str | None
+] = contextvars.ContextVar("indexing_request_id", default=None)
+
+# set by every route in the API server
+ONYX_REQUEST_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "onyx_request_id", default=None
+)
 
 """Utils related to contextvars"""