agentscope-runtime 0.1.0__py3-none-any.whl

agentscope_runtime/engine/tracing/base.py

@@ -0,0 +1,309 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=protected-access
+import time
+import traceback
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from typing import Any, Dict, List
+
+from .tracing_metric import TraceType
+
+
+# Handler Interface
+class TracerHandler(ABC):
+    """Abstract base class for tracer handlers."""
+
+    @abstractmethod
+    def on_start(
+        self,
+        event_type: TraceType,
+        payload: Dict[str, Any],
+        **kwargs: Any,
+    ) -> None:
+        """Handle the start of a trace event.
+
+        Args:
+            event_type (TraceType): The type of event being traced.
+            payload (Dict[str, Any]): The payload data for the event.
+            **kwargs (Any): Additional keyword arguments.
+        """
+
+    @abstractmethod
+    def on_end(
+        self,
+        event_type: TraceType,
+        start_payload: Dict[str, Any],
+        end_payload: Dict[str, Any],
+        start_time: float,
+        **kwargs: Any,
+    ) -> None:
+        """Handle the end of a trace event.
+
+        Args:
+            event_type (TraceType): The type of event being traced.
+            start_payload (Dict[str, Any]): The payload data from event start.
+            end_payload (Dict[str, Any]): The payload data from event end.
+            start_time (float): The timestamp when the event started.
+            **kwargs (Any): Additional keyword arguments.
+        """
+
+    @abstractmethod
+    def on_log(self, message: str, **kwargs: Any) -> None:
+        """Handle a log message during tracing.
+
+        Args:
+            message (str): The log message.
+            **kwargs (Any): Additional keyword arguments.
+        """
+
+    @abstractmethod
+    def on_error(
+        self,
+        event_type: TraceType,
+        start_payload: Dict[str, Any],
+        error: Exception,
+        start_time: float,
+        traceback_info: str,
+        **kwargs: Any,
+    ) -> None:
+        """Handle an error during tracing.
+
+        Args:
+            event_type (TraceType): The type of event being traced.
+            start_payload (Dict[str, Any]): The payload data from event start.
+            error (Exception): The exception that occurred.
+            start_time (float): The timestamp when the event started.
+            traceback_info (str): The traceback information.
+            **kwargs (Any): Additional keyword arguments.
+        """
+
+
+class BaseLogHandler(TracerHandler):
+    """Basic log handler implementation using Python's logging module."""
+
+    import logging
+
+    logger = logging.getLogger(__name__)
+
+    def on_start(
+        self,
+        event_type: TraceType,
+        payload: Dict[str, Any],
+        **kwargs: Any,
+    ) -> None:
+        """Log the start of a trace event.
+
+        Args:
+            event_type (TraceType): The type of event being traced.
+            payload (Dict[str, Any]): The payload data for the event.
+            **kwargs (Any): Additional keyword arguments.
+        """
+        self.logger.info(f"Event {event_type} started with payload: {payload}")
+
+    def on_end(
+        self,
+        event_type: TraceType,
+        start_payload: Dict[str, Any],
+        end_payload: Dict[str, Any],
+        start_time: float,
+        **kwargs: Any,
+    ) -> None:
+        """Log the end of a trace event.
+
+        Args:
+            event_type (TraceType): The type of event being traced.
+            start_payload (Dict[str, Any]): The payload data from event start.
+            end_payload (Dict[str, Any]): The payload data from event end.
+            start_time (float): The timestamp when the event started.
+            **kwargs (Any): Additional keyword arguments.
+        """
+        self.logger.info(
+            f"Event {event_type} ended with start payload: {start_payload}, "
+            f"end payload: {end_payload}, duration: "
+            f"{time.time() - start_time} seconds, kwargs: {kwargs}",
+        )
+
+    def on_log(self, message: str, **kwargs: Any) -> None:
+        """Log a message during tracing.
+
+        Args:
+            message (str): The log message.
+            **kwargs (Any): Additional keyword arguments.
+        """
+        if message:
+            self.logger.info(f"Log: {message}")
+
+    def on_error(
+        self,
+        event_type: TraceType,
+        start_payload: Dict[str, Any],
+        error: Exception,
+        start_time: float,
+        traceback_info: str,
+        **kwargs: Any,
+    ) -> None:
+        """Log an error during tracing.
+
+        Args:
+            event_type (TraceType): The type of event being traced.
+            start_payload (Dict[str, Any]): The payload data from event start.
+            error (Exception): The exception that occurred.
+            start_time (float): The timestamp when the event started.
+            traceback_info (str): The traceback information.
+            **kwargs (Any): Additional keyword arguments.
+        """
+        self.logger.error(
+            f"Error in event {event_type} with payload: {start_payload}, "
+            f"error: {error}, "
+            f"traceback: {traceback_info}, duration: "
+            f"{time.time() - start_time} seconds, kwargs: {kwargs}",
+        )
+
+
+class Tracer:
+    """
+    Tracer class for logging events
+    usage:
+    with tracer.event(TraceType.LLM, payload) as event:
+        event.log("message")
+        ""...logic here...""
+        end_payload = {xxx}
+        # optional on_end call for additional payload and kwargs
+        event.on_end(end_payload, if_success=True)
+    """
+
+    def __init__(self, handlers: List[TracerHandler]):
+        """Initialize the tracer with a list of handlers.
+
+        Args:
+            handlers (List[TracerHandler]): List of handlers to process trace
+            events.
+        """
+        self.handlers = handlers
+
+    @contextmanager
+    def event(
+        self,
+        event_type: TraceType,
+        payload: Dict[str, Any],
+        **kwargs: Any,
+    ) -> Any:
+        """Create a context manager for tracing an event.
+
+        Args:
+            event_type (TraceType): The type of event being traced.
+            payload (Dict[str, Any]): The payload data for the event.
+            **kwargs (Any): Additional keyword arguments.
+
+        Yields:
+            EventContext: The event context for logging and managing the trace.
+        """
+        start_time = time.time()
+
+        for handle in self.handlers:
+            handle.on_start(event_type, payload, **kwargs)
+
+        event_context = EventContext(
+            self.handlers,
+            event_type,
+            start_time,
+            payload,
+        )
+        try:
+            yield event_context
+        except Exception as e:
+            traceback_info = traceback.format_exc()
+            for handle in self.handlers:
+                handle.on_error(
+                    event_type,
+                    payload,
+                    e,
+                    start_time,
+                    traceback_info=traceback_info,
+                )
+            raise
+        event_context._end(payload)
+
+    def log(self, message: str, **kwargs: Any) -> None:
+        """Log a message using all registered handlers.
+
+        Args:
+            message (str): The log message.
+            **kwargs (Any): Additional keyword arguments.
+        """
+        for handle in self.handlers:
+            handle.on_log(message, **kwargs)
+
+
+class EventContext:
+    """Context manager for individual trace events."""
+
+    def __init__(
+        self,
+        handlers: List[TracerHandler],
+        event_type: TraceType,
+        start_time: float,
+        start_payload: Dict[str, Any],
+    ) -> None:
+        """Initialize the event context.
+
+        Args:
+            handlers (List[TracerHandler]): List of handlers to process
+            trace  events.
+            event_type (TraceType): The type of event being traced.
+            start_time (float): The timestamp when the event started.
+            start_payload (Dict[str, Any]): The payload data from event start.
+        """
+        self.handlers = handlers
+        self.event_type = event_type
+        self.start_time = start_time
+        self.start_payload = start_payload
+        self.end_payload = {}
+        self.kwargs = {}
+
+    def on_end(self, payload: Dict[str, Any], **kwargs: Any) -> None:
+        """Set the end payload and additional kwargs for the event.
+
+        Args:
+            payload (Dict[str, Any]): The payload data for event end.
+            **kwargs (Any): Additional keyword arguments.
+        """
+        self.end_payload = payload
+        self.kwargs = kwargs
+
+    def on_log(self, message: str, **kwargs: Any) -> None:
+        """Log a message within this event context.
+
+        Args:
+            message (str): The log message.
+            **kwargs (Any): Additional keyword arguments.
+        """
+        kwargs["event_type"] = self.event_type
+        kwargs["start_time"] = self.start_time
+        kwargs["start_payload"] = self.start_payload
+        for handle in self.handlers:
+            handle.on_log(message, **kwargs)
+
+    def _end(self, start_payload: Dict[str, Any] = None) -> None:
+        """Finalize the event by calling on_end for all handlers.
+
+        Args:
+            start_payload (Dict[str, Any], optional): The payload data from
+            event start.
+        """
+        for handle in self.handlers:
+            handle.on_end(
+                self.event_type,
+                start_payload,
+                self.end_payload,
+                self.start_time,
+                **self.kwargs,
+            )
+
+    def set_trace_header(self, trace_header: Dict[str, Any]) -> None:
+        """Set trace header for compatible handlers.
+
+        Args:
+            trace_header (Dict[str, Any]): The trace header information.
+        """
+        # TODO: Implement this

agentscope_runtime/engine/tracing/local_logging_handler.py

@@ -0,0 +1,356 @@
+# -*- coding: utf-8 -*-
+import json
+import logging
+import os
+import time
+from datetime import datetime
+from logging.handlers import RotatingFileHandler
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel
+
+from .base import TracerHandler, TraceType
+
+DEFAULT_LOG_NAME = "agentscope-runtime"
+
+INFO_LOG_FILE_NAME = "info"
+ERROR_LOG_FILE_NAME = "error"
+LOG_EXTENSION = "log"
+DS_SVC_ID = os.getenv("DS_SVC_ID", "test_id")
+DS_SVC_NAME = os.getenv("DS_SVC_NAME", "test_name")
+
+
+class LogContext(BaseModel):
+    """Pydantic model for log context data."""
+
+    time: str = ""
+    step: str = ""
+    model: str = ""
+    user_id: str = ""
+    task_id: str = ""
+    code: str = ""
+    message: str = ""
+    request_id: str = ""
+    context: Dict = {}
+    interval: Dict = {}
+    service_id: str = ""
+    service_name: str = ""
+    ds_service_id: str = ""
+    ds_service_name: str = ""
+
+    class Config:
+        extra = "ignore"  # ignore additional key
+
+
+class JsonFormatter(logging.Formatter):
+    """
+    Custom formatter to output logs in llm chat format.
+    """
+
+    def format(self, record: logging.LogRecord) -> str:
+        """Format a log record as JSON.
+
+        Args:
+            record (logging.LogRecord): The log record to format.
+
+        Returns:
+            str: The formatted log record as a JSON string.
+        """
+        log_record = {
+            "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3],
+            "step": getattr(record, "step", None),
+            "model": getattr(record, "model", None),
+            "user_id": getattr(record, "user_id", None),
+            "code": getattr(record, "code", None),
+            "message": record.getMessage(),
+            "task_id": getattr(record, "task_id", None),
+            "request_id": getattr(record, "request_id", None),
+            "context": getattr(record, "context", None),
+            "interval": getattr(record, "interval", None),
+            "ds_service_id": DS_SVC_ID,
+            "ds_service_name": DS_SVC_NAME,
+        }
+        # Clean up any extra fields that are None (not provided)
+        log_record = {k: v for k, v in log_record.items() if v is not None}
+        if record.exc_info:
+            log_record["exc_info"] = self.formatException(record.exc_info)
+        return json.dumps(log_record, ensure_ascii=False)
+
+
+class LocalLogHandler(TracerHandler):
+    """llm chat log handler for structured JSON logging."""
+
+    def __init__(
+        self,
+        log_level: int = logging.INFO,
+        log_file_name: Optional[str] = None,
+        log_dir: str = f"{os.getcwd()}/logs",
+        max_bytes: int = 1024 * 1024 * 1024,
+        backup_count: int = 7,
+        enable_console: bool = False,
+        propagate: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the llm chat log handler.
+
+        Args:
+            log_level (int): The logging level. Defaults to logging.INFO.
+            log_file_name (Optional[str]): Prefix for log file names.
+                            Defaults to None.
+            log_dir (str): Directory to save log files. Defaults to "./logs".
+            max_bytes (int): Maximum size in bytes for a single log file.
+                    Defaults to 1GB.
+            backup_count (int): Number of log files to keep. Defaults to 7.
+            enable_console (bool): Whether to enable console logging.
+                            Defaults to False.
+            propagate (bool): Whether to propagate log messages to parent
+            loggers.
+                            Defaults to False.
+            **kwargs (Any): Additional keyword arguments.
+        """
+        self.logger = logging.getLogger(DEFAULT_LOG_NAME)
+        # Set propagate attribute to control log message propagation
+        # Note: Python logging propagate defaults to True, we set it to
+        # False by default
+        self.logger.propagate = propagate
+
+        if enable_console:
+            handler = logging.StreamHandler()
+            handler.setFormatter(JsonFormatter())
+            self.logger.addHandler(handler)
+        os.makedirs(log_dir, exist_ok=True)
+        self._set_file_handle(
+            log_dir=log_dir,
+            log_file_name=log_file_name,
+            max_bytes=max_bytes,
+            backup_count=backup_count,
+        )
+
+        self.logger.setLevel(log_level)
+        self.kwargs = kwargs
+
+    def _set_file_handle(
+        self,
+        log_dir: str,
+        log_file_name: Optional[str],
+        max_bytes: int,
+        backup_count: int,
+    ) -> None:
+        """Set up file handlers for logging.
+
+        Args:
+            log_dir (str): Directory to save the log files.
+            log_file_name (Optional[str]): Prefix name of log file name.
+            max_bytes (int): Maximum size in bytes for a single log file.
+            backup_count (int): The number of log files to keep.
+        """
+        log_file_name_prefix = f"{log_file_name}-" if log_file_name else ""
+
+        # Create file handlers with JsonFormatter
+        info_file_path = os.path.join(
+            log_dir,
+            f"{log_file_name_prefix}{INFO_LOG_FILE_NAME}.{LOG_EXTENSION}."
+            f"{os.getpid()}",
+        )
+        info_file_handler = RotatingFileHandler(
+            info_file_path,
+            mode="a",
+            maxBytes=max_bytes,
+            backupCount=backup_count,
+        )
+        info_file_handler.setFormatter(JsonFormatter())
+        info_file_handler.setLevel(logging.INFO)
+
+        # Create error file handler
+        error_file_path = os.path.join(
+            log_dir,
+            f"{log_file_name_prefix}{ERROR_LOG_FILE_NAME}.{LOG_EXTENSION}."
+            f"{os.getpid()}",
+        )
+        error_file_handler = RotatingFileHandler(
+            error_file_path,
+            mode="a",
+            maxBytes=max_bytes,
+            backupCount=backup_count,
+        )
+        error_file_handler.setFormatter(JsonFormatter())
+        error_file_handler.setLevel(logging.ERROR)
+
+        # Add handlers to the logger
+        self.logger.addHandler(info_file_handler)
+        self.logger.addHandler(error_file_handler)
+
+    @staticmethod
+    def _deep_update(original: Dict[str, Any], update: Dict[str, Any]) -> None:
+        """Recursively update a dictionary with another dictionary.
+
+        Args:
+            original (Dict[str, Any]): The original dictionary to update.
+            update (Dict[str, Any]): The dictionary containing updates.
+        """
+        for key, value in update.items():
+            if (
+                isinstance(value, dict)
+                and key in original
+                and isinstance(original[key], dict)
+            ):
+                LocalLogHandler._deep_update(original[key], value)
+            else:
+                original[key] = value
+
+    def on_start(
+        self,
+        event_type: TraceType,
+        payload: Dict[str, Any],
+        **kwargs: Any,
+    ) -> None:
+        """Handle the start of a trace event.
+
+        Args:
+            event_type (TraceType): The type of event being traced.
+            payload (Dict[str, Any]): The payload data for the event.
+            **kwargs (Any): Additional keyword arguments.
+        """
+        step = f"{event_type}_start"
+        request_id = payload.get("request_id", "")
+        context = payload.get("context", payload)
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
+        interval = {"type": step, "cost": 0}
+        runtime_context = LogContext(
+            time=timestamp,
+            step=step,
+            interval=interval,
+            context=context,
+            request_id=request_id,
+        )
+        try:
+            self.logger.info(
+                runtime_context.message,
+                extra=runtime_context.model_dump(exclude={"message"}),
+            )
+        except Exception as e:
+            import traceback
+
+            print(traceback.format_exc())
+            print(e)
+
+    def on_end(
+        self,
+        event_type: TraceType,
+        start_payload: Dict[str, Any],
+        end_payload: Dict[str, Any],
+        start_time: float,
+        **kwargs: Any,
+    ) -> None:
+        """Handle the end of a trace event.
+
+        Args:
+            event_type (TraceType): The type of event being traced.
+            start_payload (Dict[str, Any]): The payload data from event start.
+            end_payload (Dict[str, Any]): The payload data from event end.
+            start_time (float): The timestamp when the event started.
+            **kwargs (Any): Additional keyword arguments.
+        """
+
+        LocalLogHandler._deep_update(end_payload, start_payload)
+
+        step = f"{event_type}_end"
+        request_id = end_payload.get("request_id", "")
+        context = end_payload.get("context", end_payload)
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
+        duration = time.time() - start_time
+        interval = {"type": step, "cost": f"{duration:.3f}"}
+        runtime_context = LogContext(
+            time=timestamp,
+            step=step,
+            interval=interval,
+            context=context,
+            request_id=request_id,
+        )
+        self.logger.info(
+            runtime_context.message,
+            extra=runtime_context.model_dump(exclude={"message"}),
+        )
+
+    def on_log(self, message: str, **kwargs: Any) -> None:
+        """Handle a log message during tracing.
+
+        Args:
+            message (str): The log message.
+            **kwargs (Any): Additional keyword arguments including:
+                - step_suffix: Optional suffix for step identification
+                - event_type: The type of event being traced
+                - payload: The payload data
+                - start_time: The timestamp when the event started
+                - start_payload: The payload data from event start
+        """
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
+        if "step_suffix" in kwargs:
+            step_suffix = kwargs["step_suffix"]
+            event_type = kwargs["event_type"]
+            payload = kwargs["payload"]
+            start_time = kwargs["start_time"]
+            start_payload = kwargs["start_payload"]
+
+            LocalLogHandler._deep_update(payload, start_payload)
+
+            step = f"{event_type}_{step_suffix}"
+            duration = time.time() - start_time
+            interval = {"type": step, "cost": f"{duration:.3f}"}
+        else:
+            step = ""
+            interval = {"type": step, "cost": "0"}
+            payload = {}
+
+        runtime_context = LogContext(
+            time=timestamp,
+            step=step,
+            interval=interval,
+            context=payload,
+            message=message,
+        )
+        self.logger.info(
+            runtime_context.message,
+            extra=runtime_context.model_dump(exclude={"message"}),
+        )
+
+    def on_error(
+        self,
+        event_type: TraceType,
+        start_payload: Dict[str, Any],
+        error: Exception,
+        start_time: float,
+        traceback_info: str,
+        **kwargs: Any,
+    ) -> None:
+        """Handle an error during tracing.
+
+        Args:
+            event_type (TraceType): The type of event being traced.
+            start_payload (Dict[str, Any]): The payload data from event start.
+            error (Exception): The exception that occurred.
+            start_time (float): The timestamp when the event started.
+            traceback_info (str): The traceback information.
+            **kwargs (Any): Additional keyword arguments.
+        """
+        step = f"{event_type}_error"
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
+        duration = time.time() - start_time
+        interval = {"type": step, "cost": f"{duration:.3f}"}
+        if "context" not in start_payload:
+            start_payload["context"] = {}
+        start_payload["context"].update(
+            {"type": error.__class__.__name__, "details": traceback_info},
+        )
+        runtime_context = LogContext(
+            time=timestamp,
+            step=step,
+            interval=interval,
+            code=error.__class__.__name__,
+            message=str(error),
+            **start_payload,
+        )
+        self.logger.error(
+            runtime_context.message,
+            extra=runtime_context.model_dump(exclude={"message"}),
+        )