diff --git a/pyproject.toml b/pyproject.toml index f81ea801..3694c645 100644 --- a/pyproject.toml +++ b/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]] diff --git a/requirements-dev.lock b/requirements-dev.lock index c2b6e799..3baf1c05 100644 --- a/requirements-dev.lock +++ b/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 diff --git a/requirements.lock b/requirements.lock index 58bc3e38..e055b69f 100644 --- a/requirements.lock +++ b/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 diff --git a/src/agentex/lib/environment_variables.py b/src/agentex/lib/environment_variables.py index a369c867..aff88b68 100644 --- a/src/agentex/lib/environment_variables.py +++ b/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" diff --git a/src/agentex/lib/sdk/fastacp/base/base_acp_server.py b/src/agentex/lib/sdk/fastacp/base/base_acp_server.py index 04ec90a2..3942790a 100644 --- a/src/agentex/lib/sdk/fastacp/base/base_acp_server.py +++ b/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 diff --git a/src/agentex/lib/utils/logging.py b/src/agentex/lib/utils/logging.py index 26d9abcd..8a8b1b20 100644 --- a/src/agentex/lib/utils/logging.py +++ b/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 diff --git a/uv.lock b/uv.lock index e1c22638..c7a35067 100644 --- a/uv.lock +++ b/uv.lock @@ -4,7 +4,7 @@ requires-python = ">=3.12, <4" [[package]] name = "agentex-sdk" -version = "0.4.8" +version = "0.4.18" source = { editable = "." } dependencies = [ { name = "aiohttp" },