feat: add custom request log

Author: CarltonXiang

src/memos/api/context/dependencies.py

@@ -1,4 +1,5 @@
 import logging
+import os
 
 from fastapi import Depends, Header, Request
 
@@ -24,6 +25,13 @@ def get_trace_id_from_header(
     return g_trace_id or x_trace_id or trace_id
 
 
+def generate_trace_id() -> str:
+    """
+    Get a random trace_id.
+    """
+    return os.urandom(16).hex()
+
+
 def get_request_context(
     request: Request, trace_id: str | None = Depends(get_trace_id_from_header)
 ) -> RequestContext:
@@ -57,6 +65,9 @@ def get_g_object(trace_id: str | None = Depends(get_trace_id_from_header)) -> G:
     This creates a RequestContext and sets it globally for access
     throughout the request lifecycle.
     """
+    if trace_id is None:
+        trace_id = generate_trace_id()
+
     g = RequestContext(trace_id=trace_id)
     set_request_context(g)
     logger.info(f"Request g object created with trace_id: {g.trace_id}")

src/memos/api/product_api.py

@@ -1,7 +1,8 @@
 import logging
 
-from fastapi import FastAPI
+from fastapi import Depends, FastAPI
 
+from memos.api.context.dependencies import get_g_object
 from memos.api.exceptions import APIExceptionHandler
 from memos.api.routers.product_router import router as product_router
 
@@ -14,6 +15,7 @@
     title="MemOS Product REST APIs",
     description="A REST API for managing multiple users with MemOS Product.",
     version="1.0.0",
+    dependencies=[Depends(get_g_object)],
 )
 
 # Include routers

src/memos/api/start_api.py

@@ -4,11 +4,12 @@
 from typing import Any, Generic, TypeVar
 
 from dotenv import load_dotenv
-from fastapi import FastAPI
+from fastapi import Depends, FastAPI
 from fastapi.requests import Request
 from fastapi.responses import JSONResponse, RedirectResponse
 from pydantic import BaseModel, Field
 
+from memos.api.context.dependencies import get_g_object
 from memos.configs.mem_os import MOSConfig
 from memos.mem_os.main import MOS
 from memos.mem_user.user_manager import UserManager, UserRole
@@ -76,6 +77,7 @@ def get_mos_instance():
     title="MemOS REST APIs",
     description="A REST API for managing and searching memories using MemOS.",
     version="1.0.0",
+    dependencies=[Depends(get_g_object)],
 )
 
 

src/memos/log.py

@@ -1,12 +1,22 @@
 import logging
+import os
 
+from contextlib import suppress
 from logging.config import dictConfig
 from pathlib import Path
 from sys import stdout
 
+import requests
+
+from dotenv import load_dotenv
+
 from memos import settings
+from memos.api.context.context import get_current_trace_id
 
 
+# Load environment variables
+load_dotenv()
+
 selected_log_level = logging.DEBUG if settings.DEBUG else logging.WARNING
 
 
@@ -21,6 +31,16 @@ def _setup_logfile() -> Path:
     return logfile
 
 
+class CustomLoggerRequestHandler(logging.Handler):
+    def emit(self, record):
+        if os.getenv("CUSTOM_LOGGER_URL") is None:
+            return
+
+        if record.levelno == logging.INFO or record.levelno == logging.ERROR:
+            with suppress(Exception):
+                log_custom_request(record.getMessage())
+
+
 LOGGING_CONFIG = {
     "version": 1,
     "disable_existing_loggers": False,
@@ -48,13 +68,17 @@ def _setup_logfile() -> Path:
             "class": "logging.handlers.RotatingFileHandler",
             "filename": _setup_logfile(),
             "maxBytes": 1024**2 * 10,
-            "backupCount": 10,
+            "backupCount": 3,
             "formatter": "standard",
         },
+        "customRequest": {
+            "level": "INFO",
+            "class": "memos.log.CustomLoggerRequestHandler",
+        },
     },
     "root": {  # Root logger handles all logs
         "level": logging.DEBUG if settings.DEBUG else logging.INFO,
-        "handlers": ["console", "file"],
+        "handlers": ["console", "file", "customRequest"],
     },
     "loggers": {
         "memos": {
@@ -76,3 +100,33 @@ def get_logger(name: str | None = None) -> logging.Logger:
     if name:
         return parent_logger.getChild(name)
     return parent_logger
+
+
+def log_custom_request(message: str):
+    logger_url = os.getenv("CUSTOM_LOGGER_URL")
+    token = os.getenv("CUSTOM_LOGGER_TOKEN")
+
+    trace_id = get_current_trace_id()
+
+    print(f"trace_id: {trace_id}")
+
+    headers = {
+        "Content-Type": "application/json",
+    }
+    post_content = {
+        "message": message,
+    }
+
+    for key, value in os.environ.items():
+        if key.startswith("CUSTOM_LOGGER_ATTRIBUTE_"):
+            attribute_key = key[len("CUSTOM_LOGGER_ATTRIBUTE_") :].lower()
+            post_content[attribute_key] = value
+
+    if token is not None:
+        headers["Authorization"] = token
+
+    if trace_id is not None:
+        headers["traceId"] = trace_id
+        post_content["trace_id"] = trace_id
+
+    requests.post(url=logger_url, headers=headers, json=post_content, timeout=5)