aixtools 0.0.0__py3-none-any.whl

aixtools/logging/logging_config.py

@@ -0,0 +1,161 @@
+"""
+Centralized logging configuration for AixTools, based on Python's standard logging.
+"""
+
+import json
+import logging
+import logging.config
+import os
+import sys
+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.
+
+    Special handling for pytest: If running under pytest without explicit
+    log flags, console logging is suppressed to avoid interfering with
+    pytest's own log capture mechanism.
+    """
+    # Detect if running under pytest and suppress console logging unless explicitly requested
+    is_pytest = "pytest" in sys.modules or "pytest" in sys.argv[0] if sys.argv else False
+
+    # Check for live log flags - handle both separate and combined flags
+    wants_live_logs = False
+    if sys.argv:
+        for arg in sys.argv:
+            # Check for explicit --log-cli
+            if arg == "--log-cli":
+                wants_live_logs = True
+                break
+            # Check for -s either standalone or combined with other flags (like -vsk, -vs, etc.)
+            if arg.startswith("-") and not arg.startswith("--") and "s" in arg:
+                wants_live_logs = True
+                break
+
+    if is_pytest and not wants_live_logs:
+        # Use a minimal configuration that doesn't output to console during pytest
+        pytest_config = {
+            "version": 1,
+            "disable_existing_loggers": False,
+            "formatters": {
+                "simple": {
+                    "format": "%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+                }
+            },
+            "handlers": {
+                # Only use NullHandler to suppress console output but allow pytest log capture
+                "null": {
+                    "class": "logging.NullHandler",
+                }
+            },
+            "root": {
+                "handlers": ["null"],
+                "level": "INFO",
+            },
+        }
+        logging.config.dictConfig(pytest_config)
+        return
+
+    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

aixtools/logging/mcp_logger.py

@@ -0,0 +1,172 @@
+"""
+Logger implementations for MCP server.
+"""
+
+import json
+from abc import ABC, abstractmethod
+from datetime import datetime
+from pathlib import Path
+
+from aixtools.logging.logging_config import get_logger
+
+from .mcp_log_models import CodeLogEntry, CommandLogEntry, LogEntry, ServiceLogEntry, SystemLogEntry
+
+logger = get_logger(__name__)
+
+
+def log_with_default_logger(entry: LogEntry) -> None:
+    """
+    Formats a log entry into a human-readable string and logs it.
+    """
+    if isinstance(entry, SystemLogEntry):
+        logger.info("%s: %s", entry.event, entry.details or "")
+    elif isinstance(entry, ServiceLogEntry):
+        logger.info("%s: %s", entry.event, entry.details or "")
+    elif isinstance(entry, CodeLogEntry):
+        logger.info("%s code: %s", entry.language, entry.code)
+    elif isinstance(entry, CommandLogEntry):
+        logger.info("%s, CWD: %s", entry.command, entry.working_directory)
+    else:
+        logger.debug("Logging entry: %s", entry.model_dump_json(indent=2))
+
+
+class McpLogger(ABC):
+    """Abstract base class for loggers."""
+
+    @abstractmethod
+    def log(self, entry: LogEntry) -> None:
+        """Log an entry."""
+
+    @abstractmethod
+    def get_logs(  # noqa: PLR0913  # pylint: disable=too-many-arguments,too-many-positional-arguments
+        self,
+        user_id: str | None = None,
+        session_id: str | None = None,
+        container_id: str | None = None,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+        limit: int = 100,
+    ) -> list[LogEntry]:
+        """Get logs with optional filters."""
+
+
+class JSONFileMcpLogger(McpLogger):
+    """Logger that stores logs in a single JSON file."""
+
+    def __init__(self, log_dir: str | Path):
+        """Initialize the logger."""
+        self.log_dir = Path(log_dir)
+        self.log_dir.mkdir(parents=True, exist_ok=True)
+        self.log_file_path = self.log_dir / "mcp_logs.jsonl"
+        self.log_file = open(self.log_file_path, "a", encoding="utf-8")  # pylint: disable=consider-using-with
+
+    def __del__(self):
+        """Ensure file is closed when the logger is destroyed."""
+        if hasattr(self, "log_file") and self.log_file and not self.log_file.closed:
+            self.log_file.close()
+
+    def _get_log_file_path(self) -> Path:
+        """Get the path to the log file."""
+        return self.log_file_path
+
+    def log(self, entry: LogEntry) -> None:
+        """Log an entry to the JSON file."""
+        log_with_default_logger(entry)
+        # Convert the entry to a JSON string
+        entry_json = entry.model_dump_json()
+        # Append the entry to the log file and flush to ensure it's written
+        self.log_file.write(entry_json + "\n")
+        self.log_file.flush()
+
+    def get_logs(  # noqa: PLR0913, PLR0912  # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-branches
+        self,
+        user_id: str | None = None,
+        session_id: str | None = None,
+        container_id: str | None = None,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+        limit: int = 100,
+    ) -> list[LogEntry]:
+        """Get logs with optional filters.
+
+        Args:
+            user_id: Filter by user ID.
+            session_id: Filter by session ID.
+            container_id: Filter by container ID.
+            start_time: Filter by start time.
+            end_time: Filter by end time.
+            limit: Maximum number of logs to return.
+
+        Returns:
+            List of log entries.
+        """
+        logs: list[LogEntry] = []
+
+        # Ensure any pending writes are flushed to disk
+        self.log_file.flush()
+
+        # Read from the single log file
+        if self.log_file_path.exists():
+            with open(self.log_file_path, "r", encoding="utf-8") as f:
+                for line in f:
+                    try:
+                        # Parse the JSON entry
+                        entry_dict = json.loads(line)
+
+                        # Convert timestamp string to datetime
+                        entry_dict["timestamp"] = datetime.fromisoformat(entry_dict["timestamp"])
+
+                        # Apply filters
+                        if user_id and entry_dict.get("user_id") != user_id:
+                            continue
+                        if session_id and entry_dict.get("session_id") != session_id:
+                            continue
+                        if container_id and entry_dict.get("container_id") != container_id:
+                            continue
+                        if start_time and entry_dict["timestamp"] < start_time:
+                            continue
+                        if end_time and entry_dict["timestamp"] > end_time:
+                            continue
+
+                        # Create the appropriate log entry object based on log_type
+                        log_type = entry_dict["log_type"]
+                        if log_type == "command":
+                            entry = CommandLogEntry(**entry_dict)
+                        elif log_type == "code":
+                            entry = CodeLogEntry(**entry_dict)
+                        elif log_type == "system":
+                            entry = SystemLogEntry(**entry_dict)
+                        else:
+                            continue
+
+                        logs.append(entry)
+                    except (json.JSONDecodeError, KeyError) as e:
+                        logger.error("Error parsing log entry: %s", e)
+                        continue
+
+                    # Check if we've reached the limit
+                    if len(logs) >= limit:
+                        break
+
+        # Sort logs by timestamp (newest first)
+        logs.sort(key=lambda x: x.timestamp, reverse=True)
+
+        return logs[:limit]
+
+
+# Global logger instance
+_mcp_logger: McpLogger | None = None
+
+
+def initialize_mcp_logger(mcp_logger: McpLogger) -> None:
+    """Initialize the MCP logger"""
+    global _mcp_logger  # noqa: PLW0603, pylint: disable=global-statement
+    _mcp_logger = mcp_logger
+
+
+def get_mcp_logger() -> McpLogger:
+    """Get the global logger for MCP server."""
+    if _mcp_logger is None:
+        raise RuntimeError("Logger not initialized")
+
+    return _mcp_logger

aixtools/logging/model_patch_logging.py

@@ -0,0 +1,87 @@
+"""
+Logging utilities for model patching and request/response tracking.
+"""
+
+import functools
+from contextlib import asynccontextmanager
+from uuid import uuid4
+
+from aixtools.logging.logging_config import get_logger
+from aixtools.model_patch.model_patch import (
+    ModelRawRequest,
+    ModelRawRequestResult,
+    ModelRawRequestYieldItem,
+    get_request_fn,
+    get_request_stream_fn,
+    model_patch,
+)
+
+logger = get_logger(__name__)
+
+
+def log_async_method(fn, agent_logger):
+    """Log async method calls"""
+
+    @functools.wraps(fn)
+    async def model_request_logger_wrapper(*args, **kwargs):
+        # Log request
+        uuid = str(uuid4())  # Create a unique ID for this request
+        log_object = ModelRawRequest(method_name=fn.__name__, request_id=uuid, args=args, kwargs=kwargs)
+        agent_logger.log(log_object)
+        # Invoke the original method
+        try:
+            result = await fn(*args, **kwargs)
+            # Log results
+            log_object = ModelRawRequestResult(method_name=fn.__name__, request_id=uuid, result=result)
+            agent_logger.log(log_object)
+        except Exception as e:
+            # Log exception
+            agent_logger.log(e)
+            raise e
+        return result
+
+    return model_request_logger_wrapper
+
+
+def log_async_stream(fn, agent_logger):
+    """Log async streaming method calls with individual item tracking."""
+
+    @functools.wraps(fn)
+    @asynccontextmanager
+    async def model_request_stream_logger_wrapper(*args, **kwargs):
+        # Log request
+        uuid = str(uuid4())  # Create a unique ID for this request
+        log_object = ModelRawRequest(method_name=fn.__name__, request_id=uuid, args=args, kwargs=kwargs)
+        agent_logger.log(log_object)
+        # Invoke the original method
+        async with fn(*args, **kwargs) as stream:
+
+            async def gen():
+                item_num = 0
+                try:
+                    async for item in stream:
+                        # Log yielded items
+                        log_object = ModelRawRequestYieldItem(
+                            method_name=fn.__name__, request_id=uuid, item=item, item_num=item_num
+                        )
+                        agent_logger.log(log_object)
+                        item_num += 1
+                        yield item
+                except Exception as e:
+                    # Log exception
+                    agent_logger.log(e)
+                    raise e
+
+            yield gen()
+
+    return model_request_stream_logger_wrapper
+
+
+def model_patch_logging(model, agent_logger):
+    """Patch model with logging methods"""
+    logger.debug("Patching model with logging")
+    return model_patch(
+        model,
+        request_method=log_async_method(get_request_fn(model), agent_logger),
+        request_stream_method=log_async_stream(get_request_stream_fn(model), agent_logger),
+    )

aixtools/logging/open_telemetry.py

@@ -0,0 +1,36 @@
+"""
+OpenTelemetry integration for logging and tracing agent operations.
+"""
+
+import os
+
+import logfire  # pylint: disable=import-error
+from pydantic_ai import Agent
+
+from aixtools.utils.config import LOGFIRE_TOKEN, LOGFIRE_TRACES_ENDPOINT
+
+
+def open_telemetry_on():
+    """Configure and enable OpenTelemetry tracing with LogFire integration."""
+    service_name = "agent_poc"
+
+    if LOGFIRE_TRACES_ENDPOINT:
+        os.environ["OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"] = LOGFIRE_TRACES_ENDPOINT
+        logfire.configure(
+            service_name=service_name,
+            # Sending to Logfire is on by default regardless of the OTEL env vars.
+            # Keep this line here if you don't want to send to both Jaeger and Logfire.
+            send_to_logfire=False,
+        )
+        Agent.instrument_all(True)
+        return
+
+    if LOGFIRE_TOKEN:
+        logfire.configure(
+            token=LOGFIRE_TOKEN,
+            service_name=service_name,
+        )
+        Agent.instrument_all(True)
+        return
+
+    print("OpenTelemetry is not enabled. Set the LOGFIRE_TOKEN or LOGFIRE_TRACES_ENDPOINT environment variable.")

aixtools/mcp/__init__.py

@@ -0,0 +1,9 @@
+"""
+Model Context Protocol (MCP) implementation for AI agent communication.
+"""
+
+from aixtools.mcp.fast_mcp_log import FastMcpLog
+
+__all__ = [
+    "FastMcpLog",
+]