aixtools 0.0.0__py3-none-any.whl

aixtools/log_view/filters.py

@@ -0,0 +1,41 @@
+"""
+Functions for filtering nodes based on various criteria.
+"""
+
+import re
+from typing import Any
+
+from aixtools.log_view.node_summary import get_node_type
+
+
+def filter_nodes(nodes: list, filters: dict[str, Any]) -> list:
+    """Filter nodes based on multiple criteria."""
+    if not filters:
+        return nodes
+
+    filtered_nodes = nodes.copy()
+
+    # Apply text filter if provided
+    if "text" in filters and filters["text"]:
+        text_filter = filters["text"].lower()
+        filtered_nodes = [node for node in filtered_nodes if text_filter in str(node).lower()]
+
+    # Apply type filter if provided
+    if "types" in filters and filters["types"]:
+        filtered_nodes = [node for node in filtered_nodes if get_node_type(node) in filters["types"]]
+
+    # Apply attribute filter if provided
+    if "attribute" in filters and filters["attribute"]:
+        attr_filter = filters["attribute"]
+        filtered_nodes = [node for node in filtered_nodes if hasattr(node, "__dict__") and attr_filter in vars(node)]
+
+    # Apply regex filter if provided
+    if "regex" in filters and filters["regex"]:
+        try:
+            pattern = re.compile(filters["regex"], re.IGNORECASE)
+            filtered_nodes = [node for node in filtered_nodes if pattern.search(str(node))]
+        except re.error:
+            # Invalid regex pattern, ignore this filter
+            pass
+
+    return filtered_nodes

aixtools/log_view/log_utils.py

@@ -0,0 +1,26 @@
+"""
+Utility functions for handling log files.
+"""
+
+from datetime import datetime
+from pathlib import Path
+
+
+def get_log_files(log_dir: Path) -> list[Path]:
+    """Get all log files in the specified directory, sorted by modification time (newest first)."""
+    if not log_dir.exists():
+        return []
+    log_files = list(log_dir.glob("agent_run.*.pkl"))
+    log_files.sort(key=lambda x: x.stat().st_mtime, reverse=True)
+    return log_files
+
+
+def format_timestamp_from_filename(filename: str) -> str:
+    """Extract and format the timestamp from a log filename."""
+    try:
+        # Extract timestamp from format "agent_run.YYYYMMDD_HHMMSS.pkl"
+        timestamp_str = filename.split("agent_run.")[1].split(".pkl")[0]
+        timestamp = datetime.strptime(timestamp_str, "%Y%m%d_%H%M%S")
+        return timestamp.strftime("%Y-%m-%d %H:%M:%S")
+    except (IndexError, ValueError):
+        return "Unknown date"

aixtools/log_view/node_summary.py

@@ -0,0 +1,229 @@
+"""
+Utility functions for working with node objects.
+"""
+
+import json
+import traceback
+
+import rich
+from mcp.types import CallToolResult, EmbeddedResource, ImageContent, TextContent
+from pydantic_ai import CallToolsNode, ModelRequestNode, UserPromptNode
+from pydantic_ai.messages import (
+    ModelRequest,
+    ModelResponse,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+from pydantic_ai.models import ModelRequestParameters
+from pydantic_ai.result import FinalResult
+from pydantic_ai.usage import Usage
+from pydantic_graph import End
+
+from aixtools.logging.logging_config import get_logger
+from aixtools.logging.model_patch_logging import ModelRawRequest, ModelRawRequestResult
+from aixtools.utils.utils import escape_newline
+
+logger = get_logger(__name__)
+
+MAX_STR_LEN = 200
+DEBUG = False
+
+
+def has_multiple_lines(s: str) -> bool:
+    """Check if a string has multiple lines."""
+    return s.count("\n") > 1
+
+
+def get_node_type(node):
+    """Return the type name of a node as a string."""
+    return str(type(node).__name__)
+
+
+def extract_node_types(nodes: list) -> set[str]:
+    """Extract all unique node types from a list of nodes."""
+    types = set()
+    for node in nodes:
+        node_type = get_node_type(node)
+        types.add(node_type)
+    return types
+
+
+def to_str(s, max_len=MAX_STR_LEN):
+    """Format string content with appropriate quoting based on content structure."""
+    s = str(s)
+    if has_multiple_lines(s):
+        s = escape_newline(s)
+    if len(s) > max_len:
+        s = s[:max_len] + "..."
+    return s
+
+
+def try_json(s):
+    """Attempt to parse string as JSON, returning parsed object or original string."""
+    # Can it be parsed as a JSON object?
+    try:
+        d = json.loads(s)
+        return d
+    except Exception:  # pylint: disable=broad-exception-caught
+        pass
+    return s
+
+
+class NodeTitle:
+    """Class to create a title for nodes in a human-readable format."""
+
+    def __init__(self):
+        pass
+
+    def summary(self, node):  # noqa: PLR0911, PLR0912, pylint: disable=too-many-return-statements,too-many-branches
+        """Generate a summary string for a node."""
+        if node is None:
+            return "None"
+        _type = str(type(node).__name__)
+        if DEBUG:
+            rich.print(node)
+        try:
+            match node:
+                case str() | bool() | float() | int():
+                    return f"`{_type}`: {to_str(node)}"
+                case list():
+                    return to_str(f"`list` ({len(node)}):\n[" + "\n, ".join([self.summary(n) for n in node]) + "]")
+                case dict():
+                    return to_str(
+                        f"`dict` ({len(node)}): "
+                        + "{"
+                        + "\n, ".join([f"{k}: {self.summary(v)}" for k, v in node.items()])
+                        + "}"
+                    )
+                case tuple():
+                    if len(node) == 0:
+                        return "`tuple`: Empty"
+                    items = [self.summary(n) for n in node]
+                    items_str = "(" + ", ".join([str(item) for item in items]) + ")"
+                    return f"`tuple` ({len(node)}): {to_str(items_str)}"
+                case CallToolsNode():
+                    return f"`{_type}`: {to_str(self.summary(node.model_response))}"
+                case CallToolResult():
+                    return f"`{_type}`: {to_str(self.summary_call_tool_result(node))}"
+                case End():
+                    return f"`{_type}`: {to_str(self.summary(node.data))}"
+                case FinalResult():
+                    if hasattr(node, "data"):
+                        return f"`{_type}`: {to_str(self.summary(node.data))}"
+                    if node.tool_name:
+                        return f"`{_type}`: {to_str(node.tool_name)}"
+                    return f"`{_type}`"
+                case ModelRawRequest():
+                    return f"`{_type}`: {to_str(self.summary_model_raw_request(node))}"
+                case ModelRawRequestResult():
+                    return f"`{_type}`: {to_str(self.summary(node.result))}"
+                case ModelRequest():
+                    return f"`{_type}`: {to_str(self.summary_model_request(node))}"
+                case ModelRequestNode():
+                    return f"`{_type}`: {to_str(self.summary(node.request))}"
+                case ModelRequestParameters():
+                    return f"`{_type}`: {to_str(self.summary_model_request_parameters(node))}"
+                case ModelResponse():
+                    return f"`{_type}`: {to_str(self.summary_model_response(node))}"
+                case TextPart() | SystemPromptPart() | UserPromptPart() | ToolReturnPart() | RetryPromptPart():
+                    return self.summary(node.content)
+                case TextContent():
+                    return self.summary(node.text)
+                case ImageContent():
+                    return f"Image: {node.mimeType}"
+                case EmbeddedResource():
+                    return f"Resource: {node.resource}"
+                case UserPromptNode():
+                    return f"`{_type}`: {to_str(self.summary_user_prompt(node))}"
+                case ToolCallPart():
+                    args = node.args
+                    if isinstance(args, str):
+                        args = try_json(args)
+                    if isinstance(args, dict):
+                        args = ", ".join([f"{k} = {self.summary(v)}" for k, v in args.items()])
+                    return f"{node.tool_name}({to_str(args)})"
+                case Usage():
+                    return f"`{_type}`: {to_str(self.summary_usage(node))}"
+                case _:
+                    logger.debug("NodeSummary.summary(): Unknown node type %s", type(node))
+                    return f"`{type(node)}`: {str(node)}"
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            print(f"Error while summarizing {_type}: {e}")
+            traceback.print_exc()
+            return f"`{_type}`: {to_str(to_str(node))}"
+
+    def summary_call_tool_result(self, node: CallToolResult):
+        """Generate summary for CallToolResult node by joining content summaries."""
+        out = [self.summary(c) for c in node.content]
+        return "\n".join(out)
+
+    def summary_model_raw_request(self, node: ModelRawRequest):
+        """Format ModelRawRequest node showing args and kwargs in readable format."""
+        args = [self.summary(p) for p in node.args]
+        kwargs = [f"{k}={self.summary(v)}" for k, v in node.kwargs.items()]
+        out = ""
+        if len(args) > 0:
+            out += ", ".join(args)
+        if len(kwargs) > 0:
+            if len(out) > 0:
+                out += ", "
+            out += ", ".join([f"{k} = {self.summary(v)}" for k, v in kwargs])
+        return out
+
+    def summary_model_request(self, node: ModelRequest):
+        """Generate summary for ModelRequest by joining part summaries."""
+        out = [self.summary(p) for p in node.parts]
+        return "\n".join(out)
+
+    def summary_model_request_parameters(self, node: ModelRequestParameters):
+        """Format model request parameters with tools and result tools."""
+        out = ""
+
+        if hasattr(node, "function_tools"):
+            tools = [self.tool_description(tool_definition) for tool_definition in node.function_tools]
+            if len(tools) > 0:
+                if len(tools) == 1:
+                    out += f"Tool: {tools[0]}"
+                else:
+                    out += "Tools:\n" + "\n".join(tools)
+
+        if hasattr(node, "output_tools"):
+            result_tools = [self.tool_description(tool_definition) for tool_definition in node.output_tools]
+            if len(result_tools) > 0:
+                if len(out) > 0:
+                    out += "\n"
+                out += "Output Tools:\n" + "\n".join(result_tools)
+
+        return out if len(out) > 0 else ""
+
+    def summary_model_response(self, node: ModelResponse):
+        """Generate summary for ModelResponse by joining part summaries."""
+        out = [self.summary(p) for p in node.parts]
+        return "\n".join(out)
+
+    def summary_usage(self, node: Usage):
+        """Format token usage information showing request and response tokens."""
+        return f"tokens: ({node.request_tokens}, {node.response_tokens}"
+
+    def summary_user_prompt(self, node: UserPromptNode):
+        """Generate summary for UserPromptNode handling both string and list formats."""
+        if isinstance(node.user_prompt, str):
+            return self.summary(node.user_prompt)
+        if node.user_prompt:
+            out = [self.summary(p) for p in node.user_prompt]
+            return "\n".join(out)
+        return "<empty>"
+
+    def tool_description(self, tool_definition):
+        """Format tool definition with name, description and parameters if multi-line."""
+        descr = f"`{tool_definition.name}`: {self.summary(tool_definition.description)}"
+        if has_multiple_lines(descr):
+            args = ""
+            for k, v in tool_definition.parameters_json_schema.items():
+                args += f"- {k}: {v}\n"
+            return f"`{tool_definition.name}`: {self.summary(tool_definition.description)}\n{args}"
+        return descr

aixtools/logfilters/__init__.py

@@ -0,0 +1,7 @@
+"""
+Logging filters for AixTools.
+"""
+
+from aixtools.logfilters.context_filter import ContextFilter
+
+__all__ = ["ContextFilter"]

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