Merge pull request #127 from scaleapi/add-better-logging-for-dd

Add better logging for dd

Author: RoxyFarhad

pyproject.toml

@@ -36,11 +36,14 @@ dependencies = [
     "tzlocal>=5.3.1",
     "tzdata>=2025.2",
     "pytest>=8.4.0",
+    "json_log_formatter>=1.1.1",
     "pytest-asyncio>=1.0.0",
     "scale-gp-beta==0.1.0a20",
     "ipykernel>=6.29.5",
     "openai==1.99.9", # anything higher than 1.99.9 breaks litellm - https://github.com/BerriAI/litellm/issues/13711
     "cloudpickle>=3.1.1",
+    "datadog>=0.52.1",
+    "ddtrace>=3.13.0"
 ]
 requires-python = ">= 3.12,<4"
 classifiers = [
@@ -222,7 +225,7 @@ warn_unused_ignores = false
 warn_redundant_casts = false
 
 disallow_any_generics = true
-disallow_untyped_defs = true
+# disallow_untyped_defs = true
 disallow_untyped_calls = true
 disallow_subclassing_any = true
 disallow_incomplete_defs = true
@@ -238,7 +241,7 @@ cache_fine_grained = true
 # ```
 # Changing this codegen to make mypy happy would increase complexity
 # and would not be worth it.
-disable_error_code = "func-returns-value,overload-cannot-match"
+disable_error_code = "func-returns-value,overload-cannot-match,no-untyped-def"
 
 # https://github.com/python/mypy/issues/12162
 [[tool.mypy.overrides]]

requirements-dev.lock

@@ -40,6 +40,8 @@ attrs==25.3.0
     # via aiohttp
     # via jsonschema
     # via referencing
+bytecode==0.17.0
+    # via ddtrace
 cachetools==5.5.2
     # via google-auth
 certifi==2023.7.22
@@ -61,6 +63,10 @@ colorlog==6.7.0
     # via nox
 comm==0.2.3
     # via ipykernel
+datadog==0.52.1
+    # via agentex-sdk
+ddtrace==3.15.0
+    # via agentex-sdk
 debugpy==1.8.16
     # via ipykernel
 decorator==5.2.1
@@ -73,6 +79,8 @@ distro==1.8.0
     # via openai
     # via scale-gp
     # via scale-gp-beta
+envier==0.6.1
+    # via ddtrace
 execnet==2.1.1
     # via pytest-xdist
 executing==2.2.0
@@ -120,6 +128,7 @@ idna==3.4
     # via yarl
 importlib-metadata==7.0.0
     # via litellm
+    # via opentelemetry-api
 iniconfig==2.0.0
     # via pytest
 ipykernel==6.30.1
@@ -135,6 +144,8 @@ jinja2==3.1.6
     # via litellm
 jiter==0.10.0
     # via openai
+json-log-formatter==1.1.1
+    # via agentex-sdk
 jsonref==1.1.0
     # via agentex-sdk
 jsonschema==4.25.0
@@ -186,6 +197,8 @@ openai==1.99.9
     # via openai-agents
 openai-agents==0.2.7
     # via agentex-sdk
+opentelemetry-api==1.37.0
+    # via ddtrace
 packaging==23.2
     # via huggingface-hub
     # via ipykernel
@@ -207,6 +220,7 @@ propcache==0.3.1
     # via aiohttp
     # via yarl
 protobuf==5.29.5
+    # via ddtrace
     # via temporalio
 psutil==7.0.0
     # via ipykernel
@@ -280,6 +294,7 @@ referencing==0.36.2
 regex==2025.7.34
     # via tiktoken
 requests==2.32.4
+    # via datadog
     # via huggingface-hub
     # via kubernetes
     # via openai-agents
@@ -363,6 +378,7 @@ typing-extensions==4.12.2
     # via nexus-rpc
     # via openai
     # via openai-agents
+    # via opentelemetry-api
     # via pydantic
     # via pydantic-core
     # via pyright
@@ -393,6 +409,8 @@ wcwidth==0.2.13
     # via prompt-toolkit
 websocket-client==1.8.0
     # via kubernetes
+wrapt==1.17.3
+    # via ddtrace
 yarl==1.20.0
     # via aiohttp
 zipp==3.17.0

requirements.lock

@@ -38,6 +38,8 @@ attrs==25.3.0
     # via aiohttp
     # via jsonschema
     # via referencing
+bytecode==0.17.0
+    # via ddtrace
 cachetools==5.5.2
     # via google-auth
 certifi==2023.7.22
@@ -57,6 +59,10 @@ colorama==0.4.6
     # via griffe
 comm==0.2.3
     # via ipykernel
+datadog==0.52.1
+    # via agentex-sdk
+ddtrace==3.15.0
+    # via agentex-sdk
 debugpy==1.8.16
     # via ipykernel
 decorator==5.2.1
@@ -66,6 +72,8 @@ distro==1.8.0
     # via openai
     # via scale-gp
     # via scale-gp-beta
+envier==0.6.1
+    # via ddtrace
 executing==2.2.0
     # via stack-data
 fastapi==0.115.14
@@ -109,6 +117,7 @@ idna==3.4
     # via yarl
 importlib-metadata==8.7.0
     # via litellm
+    # via opentelemetry-api
 iniconfig==2.1.0
     # via pytest
 ipykernel==6.30.1
@@ -124,6 +133,8 @@ jinja2==3.1.6
     # via litellm
 jiter==0.10.0
     # via openai
+json-log-formatter==1.1.1
+    # via agentex-sdk
 jsonref==1.1.0
     # via agentex-sdk
 jsonschema==4.25.0
@@ -169,6 +180,8 @@ openai==1.99.9
     # via openai-agents
 openai-agents==0.2.7
     # via agentex-sdk
+opentelemetry-api==1.37.0
+    # via ddtrace
 packaging==25.0
     # via huggingface-hub
     # via ipykernel
@@ -188,6 +201,7 @@ propcache==0.3.1
     # via aiohttp
     # via yarl
 protobuf==5.29.5
+    # via ddtrace
     # via temporalio
 psutil==7.0.0
     # via ipykernel
@@ -255,6 +269,7 @@ referencing==0.36.2
 regex==2025.7.34
     # via tiktoken
 requests==2.32.4
+    # via datadog
     # via huggingface-hub
     # via kubernetes
     # via openai-agents
@@ -333,6 +348,7 @@ typing-extensions==4.12.2
     # via nexus-rpc
     # via openai
     # via openai-agents
+    # via opentelemetry-api
     # via pydantic
     # via pydantic-core
     # via python-on-whales
@@ -360,6 +376,8 @@ wcwidth==0.2.13
     # via prompt-toolkit
 websocket-client==1.8.0
     # via kubernetes
+wrapt==1.17.3
+    # via ddtrace
 yarl==1.20.0
     # via aiohttp
 zipp==3.23.0

src/agentex/lib/environment_variables.py

@@ -39,6 +39,7 @@ class EnvVarKeys(str, Enum):
 
 
 class Environment(str, Enum):
+    LOCAL = "local"
     DEV = "development"
     STAGING = "staging"
     PROD = "production"

src/agentex/lib/sdk/fastacp/base/base_acp_server.py

@@ -1,5 +1,6 @@
 import asyncio
 import inspect
+import uuid
 from datetime import datetime
 from collections.abc import AsyncGenerator, Awaitable, Callable
 from contextlib import asynccontextmanager
@@ -9,6 +10,7 @@
 from fastapi import FastAPI, Request
 from fastapi.responses import StreamingResponse
 from pydantic import TypeAdapter, ValidationError
+from starlette.middleware.base import BaseHTTPMiddleware
 
 # from agentex.lib.sdk.fastacp.types import BaseACPConfig
 from agentex.lib.environment_variables import EnvironmentVariables, refreshed_environment_variables
@@ -24,7 +26,7 @@
 from agentex.lib.types.json_rpc import JSONRPCError, JSONRPCRequest, JSONRPCResponse
 from agentex.types.task_message_update import StreamTaskMessageFull, TaskMessageUpdate
 from agentex.types.task_message_content import TaskMessageContent
-from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.logging import ctx_var_request_id, make_logger
 from agentex.lib.utils.model_utils import BaseModel
 from agentex.lib.utils.registration import register_agent
 from agentex.lib.sdk.fastacp.base.constants import (
@@ -38,6 +40,20 @@
 task_message_update_adapter = TypeAdapter(TaskMessageUpdate)
 
 
+class RequestIDMiddleware(BaseHTTPMiddleware):
+    """Middleware to extract or generate request IDs and add them to logs and response headers"""
+
+    async def dispatch(self, request: Request, call_next):
+        # Extract request ID from header or generate a new one if there isn't one
+        request_id = request.headers.get("x-request-id") or uuid.uuid4().hex
+        logger.info(f"Request ID: {request_id}")
+        # Store request ID in request state for access in handlers
+        ctx_var_request_id.set(request_id)
+        # Process request
+        response = await call_next(request)
+        return response
+
+
 class BaseACPServer(FastAPI):
     """
     AsyncAgentACP provides RPC-style hooks for agent events and commands asynchronously.
@@ -56,6 +72,8 @@ def __init__(self):
         self.post("/api")(self._handle_jsonrpc)
 
         # Method handlers
+        # this just adds a request ID to the request and response headers
+        self.add_middleware(RequestIDMiddleware)
         self._handlers: dict[RPCMethod, Callable] = {}
 
     @classmethod

src/agentex/lib/utils/logging.py

@@ -1,31 +1,79 @@
 import logging
-
+import contextvars
 from rich.console import Console
 from rich.logging import RichHandler
+import json_log_formatter
+import os
+import ddtrace
+from ddtrace import tracer
+
+_is_datadog_configured = bool(os.environ.get("DD_AGENT_HOST"))
+
+ctx_var_request_id = contextvars.ContextVar[str]("request_id")
+
+
+class CustomJSONFormatter(json_log_formatter.JSONFormatter):
+    def json_record(self, message: str, extra: dict, record: logging.LogRecord) -> dict:
+        extra = super().json_record(message, extra, record)
+        extra["level"] = record.levelname
+        extra["name"] = record.name
+        extra["lineno"] = record.lineno
+        extra["pathname"] = record.pathname
+        extra["request_id"] = ctx_var_request_id.get(None)
+        if _is_datadog_configured:
+            extra["dd.trace_id"] = tracer.get_log_correlation_context().get("dd.trace_id", None) or getattr(
+                record, "dd.trace_id", 0
+            )
+            extra["dd.span_id"] = tracer.get_log_correlation_context().get("dd.span_id", None) or getattr(
+                record, "dd.span_id", 0
+            )
+        # add the env, service, and version configured for the tracer
+        # If tracing is not set up, then this should pull values from DD_ENV, DD_SERVICE, and DD_VERSION.
+        service_override = ddtrace.config.service or os.getenv("DD_SERVICE")
+        if service_override:
+            extra["dd.service"] = service_override
+
+        env_override = ddtrace.config.env or os.getenv("DD_ENV")
+        if env_override:
+            extra["dd.env"] = env_override
 
+        version_override = ddtrace.config.version or os.getenv("DD_VERSION")
+        if version_override:
+            extra["dd.version"] = version_override
 
-def make_logger(name: str):
+        return extra
+
+def make_logger(name: str) -> logging.Logger:
     """
     Creates a logger object with a RichHandler to print colored text.
     :param name: The name of the module to create the logger for.
     :return: A logger object.
     """
     # Create a console object to print colored text
-    console = Console()
-
-    # Create a logger object with the name of the current module
     logger = logging.getLogger(name)
-
-    # Set the global log level to INFO
     logger.setLevel(logging.INFO)
 
-    # Add the RichHandler to the logger to print colored text
-    handler = RichHandler(
-        console=console,
-        show_level=False,
-        show_path=False,
-        show_time=False,
-    )
-    logger.addHandler(handler)
+    environment = os.getenv("ENVIRONMENT")
+    if environment == "local":
+        console = Console()
+        # Add the RichHandler to the logger to print colored text
+        handler = RichHandler(
+            console=console,
+            show_level=False,
+            show_path=False,
+            show_time=False,
+        )
+        logger.addHandler(handler)
+        return logger
+
+    stream_handler = logging.StreamHandler()
+    if _is_datadog_configured:
+        stream_handler.setFormatter(CustomJSONFormatter())
+    else:
+        stream_handler.setFormatter(
+            logging.Formatter("%(asctime)s %(levelname)s [%(name)s] [%(filename)s:%(lineno)d] - %(message)s")
+        )
 
+    logger.addHandler(stream_handler)
+    # Create a logger object with the name of the current module
     return logger