aixtools 0.1.0__py3-none-any.whl

aixtools/log_view/display.py

@@ -0,0 +1,285 @@
+"""
+Functions for displaying nodes in the Streamlit interface.
+Provides enhanced display capabilities for various object types,
+including dataclasses, with proper handling of nested structures.
+"""
+
+import inspect
+import json
+from dataclasses import fields as dataclass_fields
+from dataclasses import is_dataclass
+
+import pandas as pd
+import streamlit as st
+from rich.console import Console
+
+from aixtools.utils.utils import prepend_all_lines
+
+# Toggle for using markdown display instead of JSON
+USE_MARKDOWN = True
+
+
+def filter_private_fields(data_dict: dict) -> dict:
+    """Filter out private fields from the data dictionary."""
+    return {k: v for k, v in data_dict.items() if not k.startswith("_")}
+
+
+def filter_private_attributes(obj) -> dict:
+    """
+    Filter out private attributes and methods from an object.
+    Returns a dictionary of public attributes and their values.
+    """
+    if not hasattr(obj, "__dict__"):
+        return {}
+
+    result = {}
+    for attr, value in vars(obj).items():
+        if not attr.startswith("_"):
+            result[attr] = value
+
+    return result
+
+
+def is_method(obj, attr_name: str) -> bool:
+    """Check if an attribute is a method."""
+    try:
+        attr = getattr(obj, attr_name)
+        return inspect.ismethod(attr) or inspect.isfunction(attr)
+    except (AttributeError, TypeError):
+        return False
+
+
+def get_object_type_str(obj) -> str:  # noqa: PLR0911, pylint: disable=too-many-return-statements
+    """Get a string representation of the object's type."""
+    if obj is None:
+        return "null"
+    if isinstance(obj, bool):
+        return "bool"
+    if isinstance(obj, int):
+        return "int"
+    if isinstance(obj, float):
+        return "float"
+    if isinstance(obj, str):
+        return "str"
+    if isinstance(obj, list):
+        return f"list[{len(obj)}]"
+    if isinstance(obj, tuple):
+        return f"tuple[{len(obj)}]"
+    if isinstance(obj, dict):
+        return f"dict[{len(obj)}]"
+    if isinstance(obj, set):
+        return f"set[{len(obj)}]"
+    if is_dataclass(obj):
+        return f"dataclass:{type(obj).__name__}"
+    if hasattr(obj, "__dict__"):
+        return type(obj).__name__
+
+    return type(obj).__name__
+
+
+def object_to_json_with_types(obj, max_depth: int = 5, current_depth: int = 0):  # noqa: PLR0911, PLR0912, pylint: disable=too-many-return-statements,too-many-branches
+    """
+    Convert an object to a JSON-serializable dictionary with type information.
+    Handles nested objects up to max_depth.
+    """
+    # Prevent infinite recursion
+    if current_depth > max_depth:
+        return {"__type": get_object_type_str(obj), "__value": str(obj)}
+
+    # Handle None
+    if obj is None:
+        return None
+
+    # Handle basic types
+    if isinstance(obj, (bool, int, float, str)):
+        return obj
+
+    # Handle lists and tuples
+    if isinstance(obj, (list, tuple)):
+        items = []
+        for item in obj:
+            items.append(object_to_json_with_types(item, max_depth, current_depth + 1))
+        return items
+
+    # Handle dictionaries
+    if isinstance(obj, dict):
+        result = {}
+        for key, value in filter_private_fields(obj).items():
+            result[key] = object_to_json_with_types(value, max_depth, current_depth + 1)
+        return result
+
+    # Handle sets
+    if isinstance(obj, set):
+        items = []
+        for item in obj:
+            items.append(object_to_json_with_types(item, max_depth, current_depth + 1))
+        return {"__type": "set", "__items": items}
+
+    # Handle dataclasses
+    if is_dataclass(obj):
+        result = {"__type": f"dataclass:{type(obj).__name__}"}
+        for field in dataclass_fields(obj):
+            if field.name.startswith("_"):  # Skip private fields
+                continue
+            if not hasattr(obj, field.name):  # Skip not found
+                continue
+            value = getattr(obj, field.name)
+            result[field.name] = object_to_json_with_types(value, max_depth, current_depth + 1)
+        return result
+
+    # Handle objects with __dict__
+    if hasattr(obj, "__dict__"):
+        result = {"__type": type(obj).__name__}
+        for attr, value in filter_private_attributes(obj).items():
+            if not is_method(obj, attr):  # Skip methods
+                result[attr] = object_to_json_with_types(value, max_depth, current_depth + 1)
+        return result
+
+    # Handle other types
+    return {"__type": get_object_type_str(obj), "__value": str(obj)}
+
+
+def object_to_markdown(  # noqa: PLR0911, PLR0912, PLR0915, pylint: disable=too-many-locals,too-many-return-statements,too-many-branches,too-many-statements
+    obj, max_depth: int = 5, current_depth: int = 0, indent: str = ""
+) -> str:
+    """
+    Convert an object to a compact markdown representation.
+    Handles nested objects up to max_depth.
+    """
+    max_display_items = 10  # Show only first MAX_DISPLAY_ITEMS items for large collections, dicts, and sets
+
+    # Prevent infinite recursion
+    if current_depth > max_depth:
+        return f"`{get_object_type_str(obj)}`: {str(obj)}"
+
+    # Handle None
+    if obj is None:
+        return "`None`"
+
+    # Handle basic types
+    if isinstance(obj, bool):
+        return f"`{str(obj).lower()}`"
+
+    if isinstance(obj, (int, float)):
+        return f"`{obj}`"
+
+    if isinstance(obj, str):
+        lines = str(obj).splitlines()
+        if len(lines) > 1:
+            return f"\n{indent}```\n{prepend_all_lines(obj, prepend=indent)}\n{indent}```\n"
+        return obj
+
+    # Handle lists and tuples
+    if isinstance(obj, (list, tuple)):
+        if not obj:  # Empty collection
+            return f"`{get_object_type_str(obj)}`: empty"
+
+        max_inline_length = 3  # For small collections, show inline
+        if (
+            len(obj) <= max_inline_length
+            and current_depth > 0
+            and all(isinstance(x, (bool, int, float, str, type(None))) for x in obj)
+        ):
+            items = [object_to_markdown(item, max_depth, current_depth + 1) for item in obj]
+            return f"`{get_object_type_str(obj)}`: [{', '.join(items)}]"
+
+        # For larger collections, use bullet points
+        result = [f"`{get_object_type_str(obj)}`:"]
+        for i, item in enumerate(obj):
+            if i >= max_display_items and len(obj) > max_display_items + 2:
+                result.append(f"{indent}* ... ({len(obj) - 10} more items)")
+                break
+            item_md = object_to_markdown(item, max_depth, current_depth + 1, indent + "  ")
+            result.append(f"{indent}* {item_md}")
+        return "\n".join(result)
+
+    # Handle dictionaries
+    if isinstance(obj, dict):
+        if not obj:  # Empty dict
+            return "`dict`: empty"
+
+        result = [f"`dict[{len(obj)}]`:"]
+        for i, (key, value) in enumerate(filter_private_fields(obj).items()):
+            if i >= max_display_items and len(obj) > max_display_items + 2:
+                result.append(f"{indent}* ... ({len(obj) - 10} more items)")
+                break
+            value_md = object_to_markdown(value, max_depth, current_depth + 1, indent + "  ")
+            result.append(f"{indent}* **{key}**: {value_md}")
+        return "\n".join(result)
+
+    # Handle sets
+    if isinstance(obj, set):
+        if not obj:  # Empty set
+            return "`set`: empty"
+
+        result = [f"`set[{len(obj)}]`:"]
+        for i, item in enumerate(obj):
+            if i >= max_display_items and len(obj) > max_display_items + 2:
+                result.append(f"{indent}* ... ({len(obj) - 10} more items)")
+                break
+            item_md = object_to_markdown(item, max_depth, current_depth + 1, indent + "  ")
+            result.append(f"{indent}* {item_md}")
+        return "\n".join(result)
+
+    # Handle dataclasses
+    if is_dataclass(obj):
+        result = [f"`{type(obj).__name__}:`"]
+        for field in dataclass_fields(obj):
+            if field.name.startswith("_"):  # Skip private fields
+                continue
+            if not hasattr(obj, field.name):  # Skip not found
+                continue
+            value = getattr(obj, field.name)
+            value_md = object_to_markdown(value, max_depth, current_depth + 1, indent + "  ")
+            result.append(f"{indent}* **{field.name}**: {value_md}")
+        return "\n".join(result)
+
+    # Handle objects with __dict__
+    if hasattr(obj, "__dict__"):
+        attrs = filter_private_attributes(obj)
+        if not attrs:  # No public attributes
+            return f"`{type(obj).__name__}`: (no public attributes)"
+
+        result = [f"`{type(obj).__name__}`:"]
+        for attr, value in attrs.items():
+            if not is_method(obj, attr):  # Skip methods
+                value_md = object_to_markdown(value, max_depth, current_depth + 1, indent + "  ")
+                result.append(f"{indent}* **{attr}**: {value_md}")
+        return "\n".join(result)
+
+    # Handle other types
+    return f"`{get_object_type_str(obj)}`: {str(obj)}"
+
+
+def format_json_for_display(json_obj) -> str:
+    """Format a JSON object for display with proper indentation."""
+    return json.dumps(json_obj, indent=2, default=str)
+
+
+def display_node(node, display_format: str) -> None:
+    """
+    Display node content based on its type, with enhanced formatting.
+    """
+    # Special handling for specific types
+    if isinstance(node, pd.DataFrame):
+        st.dataframe(node)
+        return
+
+    # Toggle between markdown and JSON display
+    match display_format:
+        case "Markdown":
+            st.markdown(object_to_markdown(node))
+        case "JSON":
+            st.json(object_to_json_with_types(node))
+        case "Rich":
+            st.write(rich_print(node))
+        case _:
+            raise ValueError(f"Unsupported display format: {display_format}")
+
+
+def rich_print(node):
+    """Display a node using rich print."""
+    console = Console(color_system=None)
+    with console.capture() as capture:
+        console.print(node)
+    return f"```\n{capture.get()}\n```"

aixtools/log_view/export.py

@@ -0,0 +1,51 @@
+"""
+Functions for exporting nodes to various formats.
+"""
+
+import json
+
+
+def export_nodes_to_json(nodes: list) -> str:
+    """Export nodes to a JSON string for download."""
+    # Convert nodes to a serializable format
+    serializable_nodes = []
+
+    for node in nodes:
+        if hasattr(node, "__dict__"):
+            # For objects with attributes
+            node_dict = {
+                "type": type(node).__name__,
+                "attributes": {
+                    attr: str(value) if not isinstance(value, (dict, list, int, float, bool, type(None))) else value
+                    for attr, value in vars(node).items()
+                    if not attr.startswith("_")
+                },
+            }
+            serializable_nodes.append(node_dict)
+        elif isinstance(node, dict):
+            # For dictionaries
+            serializable_nodes.append(
+                {
+                    "type": "dict",
+                    "content": {
+                        str(k): str(v) if not isinstance(v, (dict, list, int, float, bool, type(None))) else v
+                        for k, v in node.items()
+                    },
+                }
+            )
+        elif isinstance(node, (list, tuple)):
+            # For lists and tuples
+            serializable_nodes.append(
+                {
+                    "type": "list" if isinstance(node, list) else "tuple",
+                    "content": [
+                        str(item) if not isinstance(item, (dict, list, int, float, bool, type(None))) else item
+                        for item in node
+                    ],
+                }
+            )
+        else:
+            # For primitive types
+            serializable_nodes.append({"type": type(node).__name__, "value": str(node)})
+
+    return json.dumps(serializable_nodes, indent=2)

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"]