aixtools 0.1.0__py3-none-any.whl

aixtools/logfilters/context_filter.py

@@ -0,0 +1,67 @@
+"""
+A logging filter for injecting contextual information into log records.
+"""
+
+import logging
+
+
+class ContextFilter(logging.Filter):  # pylint: disable=too-few-public-methods
+    """
+    A logging filter that injects a formatted context string (user and session
+    IDs) into the log record. It sources the IDs from the active FastMCP
+    application context and ignores default values.
+    """
+
+    def _extract_from_mcp_context(self) -> tuple[str | None, str | None]:
+        """
+        Retrieve session id (aka conversation id) and user id from the MCP context.
+        Useful in MCP servers.
+        """
+        try:
+            from aixtools.server.utils import (  # noqa: PLC0415 # pylint: disable=import-outside-toplevel
+                get_session_id_tuple,
+            )
+
+            return get_session_id_tuple()
+        except (ImportError, RuntimeError, ValueError):
+            # Context is not available
+            return None, None
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        """
+        Adds a `context` string to the log record.
+
+        The filter attempts to extract user, session (conversation) IDs from
+        context variables. If that fails, it falls back to extracting IDs from
+        the FastMCP context.
+
+        If valid IDs are found, the `context` attribute is formatted as
+        `[conversation:id user:id]`. Otherwise, it is an empty string.
+        """
+        user_id = None
+        session_id = None
+
+        try:
+            # First, try to get context from the global context variables
+            from aixtools.context import (  # noqa: PLC0415 # pylint: disable=import-outside-toplevel
+                session_id_var,
+                user_id_var,
+            )
+
+            user_id = user_id_var.get()
+            session_id = session_id_var.get()
+        except ImportError:
+            pass
+
+        if not user_id and not session_id:
+            user_id, session_id = self._extract_from_mcp_context()
+
+        context = ""
+        if session_id and not str(session_id).startswith("default"):
+            context += f"[{session_id}]"
+        if user_id and not str(user_id).startswith("default"):
+            context += f"[{user_id}]"
+
+        record.context = context
+
+        return True

aixtools/logging/__init__.py

@@ -0,0 +1,30 @@
+"""
+Logging utilities for AI agent operations and model interactions.
+"""
+
+from aixtools.logging.log_objects import ObjectLogger
+from aixtools.logging.mcp_log_models import (
+    BaseLogEntry,
+    CodeLogEntry,
+    CommandLogEntry,
+    Language,
+    LogEntry,
+    LogType,
+    ProcessResult,
+    SystemLogEntry,
+)
+from aixtools.logging.mcp_logger import JSONFileMcpLogger, McpLogger
+
+__all__ = [
+    "ObjectLogger",
+    "LogType",
+    "Language",
+    "ProcessResult",
+    "BaseLogEntry",
+    "CommandLogEntry",
+    "CodeLogEntry",
+    "SystemLogEntry",
+    "LogEntry",
+    "McpLogger",
+    "JSONFileMcpLogger",
+]

aixtools/logging/log_objects.py

@@ -0,0 +1,227 @@
+"""
+This module provides functionality to save objects to a log file using pickle.
+It includes a function to check if an object is pickleable and a function to perform a safe deepcopy of objects.
+It also includes a function to save the objects to a log file with a timestamp.
+"""
+
+import logging
+import pickle
+import traceback
+from copy import copy
+from datetime import datetime
+from pathlib import Path
+from types import NoneType
+from typing import Mapping, Sequence, Union
+
+import rich
+
+from aixtools.logging.logging_config import get_logger
+from aixtools.utils.config import LOG_LEVEL, LOGS_DIR
+
+logger = get_logger(__name__)
+
+_is_pickleable_cache = {}
+
+
+class ExceptionWrapper:  # pylint: disable=too-few-public-methods
+    """
+    A wrapper for exceptions to make them pickleable.
+    It stores the exception type and message.
+    """
+
+    def __init__(self, exception):
+        self.exc_type = str(type(exception))
+        self.exc_value = str(exception)
+        self.exc_traceback = traceback.format_exc()
+
+    def __str__(self):
+        return f"{self.exc_type}: {self.exc_value}\n{self.exc_traceback}"
+
+
+def is_pickleable(obj):
+    """
+    Check if an object is pickleable.
+    Uses a cache to avoid repeated checks for the same type.
+    """
+    obj_type = type(obj)
+    module_name = getattr(obj_type, "__module__", "")
+
+    # FastMCP json_schema_to_type changes __module__ which causes pickle error but for some reason goes to the cache
+    if module_name == "fastmcp.utilities.json_schema_type":
+        return False
+
+    if obj_type not in _is_pickleable_cache:
+        try:
+            pickle.loads(pickle.dumps(obj))
+            _is_pickleable_cache[obj_type] = True
+        except Exception:  # pylint: disable=broad-exception-caught
+            _is_pickleable_cache[obj_type] = False
+    return _is_pickleable_cache[obj_type]
+
+
+def load_from_log(log_file: Path):
+    """
+    Load objects from a log file.
+    It reads the file in binary mode and uses pickle to deserialize the objects.
+    Returns a list of objects.
+    """
+    objects = []
+    with open(log_file, "rb") as f:
+        while True:
+            try:
+                obj = pickle.load(f)
+                objects.append(obj)
+            except EOFError:
+                break
+    return objects
+
+
+def safe_deepcopy(obj):
+    """
+    A safe deepcopy function that handles unpickleable objects.
+    It uses 'is_pickleable' to check if the object is serializable and
+    performs a shallow copy for unpickleable objects.
+    """
+    if isinstance(obj, Exception):
+        # Wrap exceptions to make them pickleable
+        obj = ExceptionWrapper(obj)
+
+    if is_pickleable(obj):
+        return pickle.loads(pickle.dumps(obj))  # Fast path
+
+    if isinstance(obj, Mapping):
+        return {k: safe_deepcopy(v) for k, v in obj.items() if is_pickleable(k)}
+
+    if isinstance(obj, Sequence) and not isinstance(obj, str):
+        return [safe_deepcopy(item) for item in obj]
+
+    if hasattr(obj, "__dict__"):
+        copy_obj = copy(obj)
+        for attr, value in vars(obj).items():
+            if is_pickleable(value):
+                setattr(copy_obj, attr, safe_deepcopy(value))
+            else:
+                setattr(copy_obj, attr, None)  # Remove unpickleable field
+        return copy_obj
+
+    return None  # fallback for non-serializable, non-introspectable objects
+
+
+def save_objects_to_logfile(objects: list, log_dir=LOGS_DIR):
+    """Save the objects to a (pickle) log file"""
+    with ObjectLogger(log_dir=log_dir) as object_logger:
+        for obj in objects:
+            object_logger.log(obj)
+
+
+class ObjectLogger:
+    """
+    A context manager for logging objects to a file.
+    It uses pickle to save the objects and handles exceptions during the save process.
+    """
+
+    def __init__(
+        self,
+        log_dir=LOGS_DIR,
+        verbose: bool = True,
+        debug: bool | None = None,
+        parent_logger: Union["ObjectLogger", NoneType] = None,
+    ):
+        self.verbose = verbose
+        self.debug = (
+            debug if debug is not None else (LOG_LEVEL == logging.DEBUG)
+        )  # Use the debug level from the config if not provided
+        self.log_dir = log_dir
+        self.file = None
+        self.parent_logger = parent_logger
+        self.init_log_file()
+
+    def has_parent(self):
+        """
+        Check if the logger has a parent.
+        If it does, it will not create a new log file.
+        """
+        return self.parent_logger is not None
+
+    def init_log_file(self):
+        """Initialize log file for recording agent operations."""
+        if self.has_parent():
+            # Do nothing: Delegates to the logger
+            return
+        # Create log file name
+        self.log_dir.mkdir(parents=True, exist_ok=True)
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        self.log_file = self.log_dir / f"agent_run.{timestamp}.pkl"
+        logger.info("Logging to %s", self.log_file)
+
+    def __enter__(self):
+        if self.has_parent():
+            # Do nothing: Delegates to the logger
+            return self
+        self.file = open(self.log_file, "ab")  # append in binary mode
+        return self
+
+    def log(self, obj):
+        """
+        Log an object to the file.
+        It uses safe_deepcopy to ensure the object is pickleable.
+        """
+        if self.has_parent():
+            # Delegate to the parent logger
+            self.parent_logger.log(obj)
+        else:
+            try:
+                if self.debug:
+                    rich.print(obj, flush=True)
+                elif self.verbose:
+                    print(obj, flush=True)
+                obj_to_save = safe_deepcopy(obj)
+                pickle.dump(obj_to_save, self.file)
+                self.file.flush()  # ensure it's written immediately
+            except Exception as e:  # pylint: disable=broad-exception-caught
+                logger.error("Failed to log object: %s", e)
+                logger.error(traceback.format_exc())
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.has_parent():
+            # Do nothing: Delegates to the logger
+            pass
+        elif self.file:
+            self.file.close()
+
+
+class NullObjectLogger:
+    """
+    A null logger that does nothing.
+    """
+
+    def __init__(self, **kwargs):
+        pass
+
+    def __enter__(self):
+        pass
+
+    def log(self, obj):
+        """Log an object to the configured destination."""
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+
+class PrintObjectLogger:
+    """
+    Print to stdout
+    """
+
+    def __init__(self, **kwargs):
+        pass
+
+    def __enter__(self):
+        pass
+
+    def log(self, obj):
+        """Log an object using rich print for formatted output."""
+        rich.print(obj, flush=True)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass

aixtools/logging/logging_config.py

@@ -0,0 +1,116 @@
+"""
+Centralized logging configuration for AixTools, based on Python's standard logging.
+"""
+
+import json
+import logging
+import logging.config
+import os
+import time
+from pathlib import Path
+
+# PyYAML is an optional dependency.
+try:
+    import yaml
+except ImportError:
+    yaml = None
+
+# --- Default Configuration ---
+
+logging.Formatter.converter = time.gmtime
+
+DEFAULT_LOGGING_CONFIG = {
+    "version": 1,
+    "disable_existing_loggers": False,
+    "filters": {
+        "context_filter": {
+            "()": "aixtools.logfilters.context_filter.ContextFilter",
+        }
+    },
+    "formatters": {
+        "color": {
+            "()": "colorlog.ColoredFormatter",
+            "format": "%(log_color)s%(asctime)s.%(msecs)03d %(levelname)-8s%(reset)s %(context)s[%(name)s] %(message)s",
+            "datefmt": "%Y-%m-%d %H:%M:%S",
+            "log_colors": {
+                "DEBUG": "cyan",
+                "INFO": "green",
+                "WARNING": "yellow",
+                "ERROR": "red",
+                "CRITICAL": "bold_red",
+            },
+        },
+    },
+    "handlers": {
+        "stream": {
+            "class": "colorlog.StreamHandler",
+            "formatter": "color",
+            "level": "INFO",
+            "filters": ["context_filter"],
+        },
+    },
+    "root": {
+        "handlers": ["stream"],
+        "level": "INFO",
+    },
+}
+
+# --- Public API ---
+
+get_logger = logging.getLogger
+
+
+def configure_logging():
+    """
+    Configure the logging system.
+
+    This function loads the logging configuration from a file or uses the
+    hardcoded default. The configuration source is determined in the
+    following order of precedence:
+
+    1. LOGGING_CONFIG_PATH environment variable.
+    2. logging.yaml in the current working directory.
+    3. logging.json in the current working directory.
+    4. Hardcoded default configuration.
+    """
+    config_path_str = os.environ.get("LOGGING_CONFIG_PATH")
+
+    if config_path_str:
+        config_path = Path(config_path_str)
+        if not config_path.exists():
+            raise FileNotFoundError(f"Logging configuration file not found: {config_path}")
+        _load_config_from_file(config_path)
+        return
+
+    # Check for default config files in the current directory
+    for filename in ["logging.yaml", "logging.json"]:
+        config_path = Path.cwd() / filename
+        if config_path.exists():
+            _load_config_from_file(config_path)
+            return
+
+    # Fallback to the default configuration
+    logging.config.dictConfig(DEFAULT_LOGGING_CONFIG)
+
+
+def _load_config_from_file(path: Path):
+    """Load a logging configuration from a YAML or JSON file."""
+    if path.suffix in [".yaml", ".yml"] and yaml:
+        config = yaml.safe_load(path.read_text(encoding="utf-8"))
+    elif path.suffix == ".json":
+        config = json.loads(path.read_text(encoding="utf-8"))
+    else:
+        raise ValueError(
+            f"Unsupported configuration file format: {path.suffix}. "
+            "Please use .yaml or .json. "
+            "For YAML support, ensure PyYAML is installed (`uv add pyyaml`)."
+        )
+
+    if config:
+        logging.config.dictConfig(config)
+
+
+# --- Initial Configuration ---
+
+# Automatically configure logging when the module is imported.
+configure_logging()

aixtools/logging/mcp_log_models.py

@@ -0,0 +1,102 @@
+"""
+Pydantic models for logging system.
+"""
+
+from datetime import UTC, datetime
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class LogType(str, Enum):
+    """Type of log entry."""
+
+    COMMAND = "command"
+    CODE = "code"
+    SYSTEM = "system"
+    SERVICE = "service"
+
+
+class Language(str, Enum):
+    """Programming language of the code."""
+
+    PYTHON = "python"
+    JAVASCRIPT = "javascript"
+    TYPESCRIPT = "typescript"
+    SHELL = "shell"
+    BASH = "bash"
+    OTHER = "other"
+
+
+class ProcessResult(BaseModel):
+    """
+    Process results from a command or code execution.
+    Includes exit code, stdout, and stderr.
+    """
+
+    exit_code: int = Field(description="Exit code of the command or process")
+    stdout: str = Field(description="Standard output of the command or process")
+    stderr: str = Field(description="Standard error of the command or process")
+
+
+class BaseLogEntry(BaseModel):
+    """Base model for all log entries."""
+
+    id: str = Field(description="Unique identifier for the log entry")
+    user_id: str = Field(description="ID of the user who initiated the action")
+    session_id: str = Field(description="ID of the session")
+    timestamp: datetime = Field(
+        default_factory=lambda: datetime.now(UTC),
+        description="Time when the log entry was created",
+    )
+    log_type: LogType = Field(description="Type of log entry")
+    container_id: str | None = Field(None, description="ID of the container where the action was performed")
+
+
+class CommandLogEntry(BaseLogEntry):
+    """Log entry for shell command execution."""
+
+    log_type: LogType = LogType.COMMAND
+    command: str = Field(description="Shell command that was executed")
+    working_directory: str = Field(description="Working directory where the command was executed")
+    process_result: ProcessResult | None = Field(
+        None,
+        description="Process results: exit status, STDOUT, and STDERR from the command",
+    )
+    duration_ms: int | None = Field(None, description="Duration of command execution in milliseconds")
+
+
+class CodeLogEntry(BaseLogEntry):
+    """Log entry for code execution."""
+
+    log_type: LogType = LogType.CODE
+    language: Language = Field(description="Programming language of the code")
+    code: str = Field(description="Code that was executed")
+    file_path: str | None = Field(None, description="Path to the file where the code was saved")
+    process_result: ProcessResult | None = Field(
+        None,
+        description="Process results: exit status, STDOUT, and STDERR from the command",
+    )
+    duration_ms: int | None = Field(None, description="Duration of code execution in milliseconds")
+
+
+class SystemLogEntry(BaseLogEntry):
+    """Log entry for system events."""
+
+    log_type: LogType = LogType.SYSTEM
+    event: str = Field(description="Description of the system event")
+    details: dict[str, Any] = Field(default_factory=dict, description="Additional details about the event")
+
+
+class ServiceLogEntry(BaseLogEntry):
+    """Log entry for service events."""
+
+    log_type: LogType = LogType.SERVICE
+    service_id: str = Field(description="ID of the service")
+    event: str = Field(description="Description of the service event")
+    details: dict[str, Any] = Field(default_factory=dict, description="Additional details about the event")
+
+
+# Union type for all log entry types
+LogEntry = CommandLogEntry | CodeLogEntry | SystemLogEntry | ServiceLogEntry