aixtools 0.2.15__py3-none-any.whl → 0.2.17__py3-none-any.whl

aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.15'
-__version_tuple__ = version_tuple = (0, 2, 15)
+__version__ = version = '0.2.17'
+__version_tuple__ = version_tuple = (0, 2, 17)
 
 __commit_id__ = commit_id = None

aixtools/agents/agent.py

@@ -70,6 +70,7 @@ def _get_model_openai_azure(
     azure_openai_api_key=AZURE_OPENAI_API_KEY,
     azure_openai_endpoint=AZURE_OPENAI_ENDPOINT,
     azure_openai_api_version=AZURE_OPENAI_API_VERSION,
+    http_client=None,
 ):
     assert azure_openai_endpoint, "AZURE_OPENAI_ENDPOINT is not set"
     assert azure_openai_api_key, "AZURE_OPENAI_API_KEY is not set"
@@ -79,6 +80,7 @@ def _get_model_openai_azure(
         azure_endpoint=azure_openai_endpoint,
         api_version=azure_openai_api_version,
         api_key=azure_openai_api_key,
+        http_client=http_client,
     )
     return OpenAIChatModel(model_name=model_name, provider=OpenAIProvider(openai_client=client))
 
@@ -101,9 +103,9 @@ def get_model(model_family=MODEL_FAMILY, model_name=None, http_client=None, **kw
     assert model_family is not None and model_family != "", f"Model family '{model_family}' is not set"
     match model_family:
         case "azure":
-            return _get_model_openai_azure(model_name=model_name or AZURE_MODEL_NAME, **kwargs)
+            return _get_model_openai_azure(model_name=model_name or AZURE_MODEL_NAME, http_client=http_client, **kwargs)
         case "bedrock":
-            return _get_model_bedrock(model_name=model_name or BEDROCK_MODEL_NAME, **kwargs)
+            return _get_model_bedrock(model_name=model_name or BEDROCK_MODEL_NAME, http_client=http_client, **kwargs)
         case "ollama":
             return _get_model_ollama(model_name=model_name or OLLAMA_MODEL_NAME, http_client=http_client, **kwargs)
         case "openai":

aixtools/agents/nodes_to_md.py

@@ -0,0 +1,201 @@
+"""Convert Pydantic-AI Nodes to Markdown format."""
+
+from pydantic_ai import CallToolsNode, ModelRequestNode, UserPromptNode
+from pydantic_ai.messages import (
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+from pydantic_graph.nodes import End
+
+from aixtools.agents.nodes_to_message import NodesToMessage
+from aixtools.utils.utils import is_multiline, is_too_long, to_json_pretty_print
+
+MAX_TITLE_LENGTH = 30
+
+
+class NodesToMarkdown(NodesToMessage):
+    """
+    Convert Pydantic-AI Nodes to Markdown
+    """
+
+    def __init__(
+        self,
+        user_prompt_node: bool = True,
+        user_prompt_part: bool = True,
+        system_prompt_part: bool = True,
+        title_max_length: int = MAX_TITLE_LENGTH,
+    ):
+        """
+        Initialize the NodeToMarkdown converter.
+        Args:
+            user_prompt_node: Whether to include UserPromptNode content
+            user_prompt_part: Whether to include UserPromptPart content
+            system_prompt_part: Whether to include SystemPromptPart content
+        """
+        super().__init__(
+            user_prompt_node=user_prompt_node,
+            user_prompt_part=user_prompt_part,
+            system_prompt_part=system_prompt_part,
+        )
+        self.title_max_length = title_max_length
+
+    def to_str(self, nodes) -> str | None:
+        """Convert a node to its markdown string representation."""
+        _, title, md = self.to_markdown(nodes)
+        return f"# {title}\n\n{md}"
+
+    def to_markdown(self, node) -> tuple[str | None, str | None, str | None]:
+        """
+        Get a name, title, and markdown representation of a node.
+
+        Returns:
+            name: A short name for the node
+            title: A title for the node
+            node_md: A markdown representation of the node or None
+
+        Note: all values can be None if the node should not be converted (e.g., user prompt disabled).
+        """
+        node_md = self.node2md(node)
+        if node_md is None:
+            return None, None, None
+        name = self.node2name(node)
+        title = self.node2title(node)
+        return name, title, node_md
+
+    def _format_prompt_part(self, label: str, content: str, enabled: bool) -> str | None:
+        """Format UserPromptPart or SystemPromptPart with multiline handling."""
+        if not enabled:
+            return None
+        if is_multiline(content):  # type: ignore
+            return f"### {label}\n{content}\n"
+        return f"{label}: {content}\n"
+
+    def _part2md(self, p) -> str | None:  # noqa: PLR0911  # pylint: disable=too-many-return-statements
+        """Convert a Part to a string representation."""
+        match p:
+            case ToolCallPart():
+                return f"### Tool `{p.tool_name}`\n```json\n{to_json_pretty_print(p.args)}\n```\n"
+            case TextPart():
+                return f"### Text\n{p.content}\n" if is_multiline(p.content) else f"{p.content}\n"
+            case ToolReturnPart():
+                if is_multiline(p.content) or is_too_long(p.content):
+                    return f"### Tool return `{p.tool_name}`\n```json\n{to_json_pretty_print(p.content)}\n```\n"
+                return f"Tool return `{p.tool_name}`: `{p.content}`"
+            case UserPromptPart():
+                return self._format_prompt_part("UserPromptPart", p.content, self.user_prompt_part)  # type: ignore
+            case SystemPromptPart():
+                return self._format_prompt_part("SystemPromptPart", p.content, self.system_prompt_part)  # type: ignore
+            case RetryPromptPart():
+                return f"### RetryPromptPart `{p.tool_name}`\n{p.content}\n"
+            case _:
+                return f"### Part {type(p)}\n{p}"
+
+    def _part2title(self, p) -> str:  # noqa: PLR0911  # pylint: disable=too-many-return-statements
+        """Convert a Part to a title representation."""
+        match p:
+            case ToolCallPart():
+                return f"Tool `{p.tool_name}`"
+            case TextPart():
+                return self._to_title_length(p.content)
+            case ToolReturnPart():
+                return f"Tool return `{p.tool_name}`"
+            case UserPromptPart():
+                return f"UserPromptPart: {self._to_title_length(p.content)}\n"
+            case SystemPromptPart():
+                return f"SystemPromptPart: {self._to_title_length(p.content)}\n"
+            case RetryPromptPart():
+                return f"WARNING: Retry {p.tool_name}"
+            case _:
+                return f"{type(p)}: {self._to_title_length(p)}"
+
+    def _parts2md(self, parts) -> str | None:
+        """Convert to string a list of Parts with a given prefix."""
+        if len(parts) == 0:
+            return None
+
+        if len(parts) == 1:
+            s = self._part2md(parts[0])
+            if s is None:
+                return None
+            return s + "\n"
+
+        result = ""
+        for p in parts:
+            s = self._part2md(p)
+            if s is not None:
+                result += s + "\n"
+        return result
+
+    def _parts2title(self, parts) -> str:
+        """Convert to title a list of Parts with a given prefix."""
+        if len(parts) == 0:
+            return ""
+
+        if len(parts) == 1:
+            return self._part2title(parts[0])
+
+        result = ""
+        for p in parts:
+            s = self._part2title(p)
+            # Prefer 'Tool' titles
+            if s.startswith("ERROR"):
+                return s
+            if s.startswith("Tool"):
+                return s
+            if not result:
+                result = s
+
+        return result
+
+    def node2md(self, n) -> str | None:
+        """Convert a node in a markdown format"""
+        match n:
+            case UserPromptNode():
+                return f"# UserPrompt\n{n.user_prompt}\n" if self.user_prompt_node else None
+            case CallToolsNode():
+                return f"# Call tools\n{self._parts2md(n.model_response.parts)}\n"
+            case ModelRequestNode():
+                return f"# Model request\n{self._parts2md(n.request.parts)}\n"
+            case End():
+                return f"# End\n{n.data.output}\n"
+            case RetryPromptPart():
+                return "# Retry tool"
+            case _:
+                return f"{type(n)}: {n}"
+
+    def node2title(self, n) -> str | None:
+        """Get a title for a node."""
+        match n:
+            case UserPromptNode():
+                return self._to_title_length(n.user_prompt)
+            case CallToolsNode():
+                return self._parts2title(n.model_response.parts)
+            case ModelRequestNode():
+                return self._parts2title(n.request.parts)
+            case End():
+                return self._to_title_length(n.data.output)
+            case _:
+                return f"{type(n)}: {n}"
+
+    def node2name(self, n) -> str | None:
+        """Get a short name for a node."""
+        match n:
+            case UserPromptNode():
+                return "UserPrompt"
+            case CallToolsNode():
+                return "Call tools"
+            case ModelRequestNode():
+                return "Model request"
+            case End():
+                return "End"
+            case _:
+                return f"{type(n)}: {n}"
+
+    def _to_title_length(self, s) -> str:
+        """Truncate a string to a maximum length for title display."""
+        s = str(s).replace("\n", " ").strip()
+        return s[: self.title_max_length] + "..." if len(s) > self.title_max_length else s

aixtools/agents/nodes_to_message.py

@@ -0,0 +1,30 @@
+"""Convert Pydantic-AI Nodes to Markdown format."""
+
+from abc import ABC
+
+
+class NodesToMessage(ABC):  # pylint: disable=too-few-public-methods
+    """
+    Convert Pydantic-AI Nodes to Message format
+    """
+
+    def __init__(
+        self,
+        user_prompt_node: bool = True,
+        user_prompt_part: bool = True,
+        system_prompt_part: bool = True,
+    ):
+        """
+        Initialize the NodeToMessage converter.
+        Args:
+            user_prompt_node: Whether to include UserPromptNode content
+            user_prompt_part: Whether to include UserPromptPart content
+            system_prompt_part: Whether to include SystemPromptPart content
+        """
+        self.user_prompt_node = user_prompt_node
+        self.user_prompt_part = user_prompt_part
+        self.system_prompt_part = system_prompt_part
+
+    def to_str(self, nodes) -> str | None:
+        """Convert a node to its string representation."""
+        raise NotImplementedError("to_str method must be implemented by subclasses")

aixtools/agents/nodes_to_str.py

@@ -0,0 +1,109 @@
+"""Convert Pydantic-AI Nodes to String format."""
+
+from collections.abc import Iterable
+
+from pydantic_ai import CallToolsNode, ModelRequestNode, UserPromptNode
+from pydantic_ai.messages import SystemPromptPart, TextPart, ToolCallPart, ToolReturnPart, UserPromptPart
+from pydantic_graph.nodes import End
+
+from aixtools.agents.nodes_to_message import NodesToMessage
+from aixtools.utils.utils import is_multiline, tabit
+
+
+def print_nodes(nodes):
+    """Convert a list of nodes in a readable format."""
+    n2s = NodesToString()
+    print(n2s.to_str(nodes))
+
+
+class NodesToString(NodesToMessage):  # pylint: disable=too-few-public-methods
+    """
+    Convert Pydantic-AI Nodes to String
+    """
+
+    def __init__(self, user_prompt_node: bool = True, user_prompt_part: bool = True, system_prompt_part: bool = True):
+        """
+        Initialize the NodeToMarkdown converter.
+        Args:
+            user_prompt_node: Whether to include UserPromptNode content
+            user_prompt_part: Whether to include UserPromptPart content
+            system_prompt_part: Whether to include SystemPromptPart content
+        """
+        super().__init__(
+            user_prompt_node=user_prompt_node,
+            user_prompt_part=user_prompt_part,
+            system_prompt_part=system_prompt_part,
+        )
+
+    def _format_content(self, label: str, content: str, prefix: str) -> str:
+        """Format content with optional multiline handling."""
+        if is_multiline(content):
+            pre = f"{prefix}\t|"
+            return f"{prefix}{label}:\n{tabit(content, pre)}"
+        return f"{prefix}{label}: {content}"
+
+    def _part2str(self, p, prefix: str = "\t") -> str | None:
+        """Convert a Part to a string representation."""
+        match p:
+            case ToolCallPart():
+                return f"{prefix}Tool: {p.tool_name}, args: {p.args}"
+            case TextPart():
+                return self._format_content("Text", p.content, prefix)
+            case ToolReturnPart():
+                return f"{prefix}Tool return: {p.tool_name}, content: {p.content}"
+            case UserPromptPart():
+                return None if not self.user_prompt_part else self._format_content("UserPromptPart", p.content, prefix)  # type: ignore  # pylint: disable=line-too-long
+            case SystemPromptPart():
+                return (
+                    None if not self.system_prompt_part else self._format_content("SystemPromptPart", p.content, prefix)
+                )  # type: ignore
+            case _:
+                return f"{prefix}Part {type(p)}: {p}"
+
+    def _parts2str(self, parts, prefix: str = "") -> str:
+        """Convert to string a list of Parts with a given prefix."""
+        if len(parts) == 0:
+            return f"{prefix}No parts\n"
+
+        if len(parts) == 1:
+            s = self._part2str(parts[0], prefix=prefix)
+            return s + "\n" if s else ""
+
+        result = ""
+        for i, p in enumerate(parts):
+            s = self._part2str(p, prefix=f"{prefix}{i}: ")
+            result += f"{s}\n" if s else ""
+        return result
+
+    def _node2str(self, n) -> str:
+        """Convert a node in a readable format."""
+        match n:
+            case UserPromptNode():
+                assert n.user_prompt is not None
+
+                return (
+                    f"UserPrompt:\n{tabit(n.user_prompt)}\n"
+                    if n.user_prompt is not None and self.user_prompt_node
+                    else ""
+                )
+            case CallToolsNode():
+                parts_str = self._parts2str(n.model_response.parts, prefix="\t")
+                return f"Call tools:\n{parts_str}\n"
+            case ModelRequestNode():
+                parts_str = self._parts2str(n.request.parts, prefix="\t")
+                return f"Model request:\n{parts_str}\n"
+            case End():
+                return f"End:\n{tabit(n.data.output)}\n"
+            case _:
+                return f"{type(n)}: {n}"
+
+    def to_str(self, nodes) -> str:
+        """Convert a list of nodes in a readable format."""
+        out = ""
+        if isinstance(nodes, Iterable) and not isinstance(nodes, (str, bytes)):
+            for n in nodes:
+                out += self._node2str(n)
+        else:
+            # Assume it's a single node
+            out += self._node2str(nodes)
+        return out

aixtools/logging/log_objects.py

@@ -38,10 +38,13 @@ class ExceptionWrapper:  # pylint: disable=too-few-public-methods
         return f"{self.exc_type}: {self.exc_value}\n{self.exc_traceback}"
 
 
-def is_pickleable(obj):
+def is_pickleable(obj, use_cache: bool = False):
     """
     Check if an object is pickleable.
-    Uses a cache to avoid repeated checks for the same type.
+
+    use_cache: If True, use the cache to store results of previous checks.
+               Why? Some complex objects may have lists of multiple types
+               inside, so whether an object is pickleable may depend on its contents.
     """
     obj_type = type(obj)
     module_name = getattr(obj_type, "__module__", "")
@@ -50,7 +53,7 @@ def is_pickleable(obj):
     if module_name == "fastmcp.utilities.json_schema_type":
         return False
 
-    if obj_type not in _is_pickleable_cache:
+    if not use_cache or obj_type not in _is_pickleable_cache:
         try:
             pickle.loads(pickle.dumps(obj))
             _is_pickleable_cache[obj_type] = True
@@ -76,32 +79,40 @@ def load_from_log(log_file: Path):
     return objects
 
 
-def safe_deepcopy(obj):
+def safe_deepcopy(obj, use_cache: bool = False):
     """
     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.
+
+    Note: If the object is complex (e.g. has a list or dict of varying objects), using the
+    cache may lead to bad results.
+    For example you analyze an object with an empty list first, then another object which is non-pickable
+    is added, but you use the cache result, resulting in the wrong assumption that the object is pickleable
+    So by default the cache is disabled.
     """
     if isinstance(obj, Exception):
         # Wrap exceptions to make them pickleable
         obj = ExceptionWrapper(obj)
 
-    if is_pickleable(obj):
+    if is_pickleable(obj, use_cache=use_cache):
         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)}
+        return {
+            k: safe_deepcopy(v, use_cache=use_cache) for k, v in obj.items() if is_pickleable(k, use_cache=use_cache)
+        }
 
     if isinstance(obj, Sequence) and not isinstance(obj, str):
-        return [safe_deepcopy(item) for item in obj]
+        return [safe_deepcopy(item, use_cache=use_cache) 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))
+            if is_pickleable(value, use_cache=use_cache):
+                setattr(copy_obj, attr, value)
             else:
-                setattr(copy_obj, attr, None)  # Remove unpickleable field
+                setattr(copy_obj, attr, safe_deepcopy(value, use_cache=use_cache))
         return copy_obj
 
     return None  # fallback for non-serializable, non-introspectable objects

aixtools/logging/mcp_middleware.py

@@ -0,0 +1,54 @@
+"""
+Custom middleware for MCP servers
+"""
+
+import traceback
+
+from fastmcp.server.middleware.error_handling import ErrorHandlingMiddleware
+from fastmcp.server.middleware.middleware import MiddlewareContext
+
+from aixtools.mcp import AixToolError
+
+
+class AixErrorHandlingMiddleware(ErrorHandlingMiddleware):
+    """Custom middleware class for handling errors in MCP servers."""
+
+    def log_as_warn_with_traceback(
+        self, *, error: Exception, original_error: Exception, context: MiddlewareContext
+    ) -> None:
+        """Logs provided error as warning.
+        original_error is an 'unwrapped' error if applicable, otherwise can be the same as error param.
+        """
+        error_type = type(original_error).__name__
+        method = context.method or "unknown"
+        error_key = f"{error_type}:{method}"
+        self.error_counts[error_key] = self.error_counts.get(error_key, 0) + 1
+        base_message = f"{method} resulted in {error_type}: {str(error)}"
+        if self.include_traceback:
+            self.logger.warning(f"{base_message}\n{traceback.format_exc()}")
+        if self.error_callback:
+            try:
+                self.error_callback(error, context)
+            except Exception as callback_error:  # pylint: disable=broad-exception-caught
+                self.logger.warning("Callback failed spectacularly: %s", callback_error)
+
+    def handle_error(self, error: Exception, context: MiddlewareContext) -> bool:
+        """Custom error logging"""
+        if isinstance(error, AixToolError):
+            self.log_as_warn_with_traceback(error=error, original_error=error, context=context)
+            return True
+
+        inner_error = error
+        while hasattr(inner_error, "__cause__") and inner_error.__cause__ is not None:
+            inner_error = inner_error.__cause__
+            if isinstance(inner_error, AixToolError):
+                self.log_as_warn_with_traceback(error=error, original_error=inner_error, context=context)
+                return True
+
+        return False
+
+    def _log_error(self, error: Exception, context: MiddlewareContext) -> None:
+        """Override original _log_error method."""
+        if self.handle_error(error, context):
+            return
+        super()._log_error(error, context)

aixtools/testing/agent_mock.py

@@ -0,0 +1,143 @@
+"""Mock agent that replays previously recorded nodes for testing purposes."""
+
+import pickle
+from pathlib import Path
+from typing import Any
+
+from aixtools.logging.log_objects import safe_deepcopy
+
+
+class AgentMock:
+    """
+    Mock agent that replays previously recorded nodes.
+    Used for testing or replaying agent runs without executing the actual agent.
+
+    Example:
+
+        # Run an agent and save its nodes
+        agent = get_agent(...)
+        ret, nodes = await run_agent(agent, prompt, ...)
+
+        # Use AgentMock to replay the nodes
+        agent_mock = AgentMock(nodes=nodes, result_output=ret)
+
+        # Now we can use agent_mock in place of the original agent
+        ret, nodes = await run_agent(agent_mock, prompt, ...)
+
+        # Save the mock agent to a file
+        agent_mock.save(Path("agent_mock.pkl"))
+    """
+
+    def __init__(self, nodes: list[Any], result_output: str | None = None):
+        """
+        Initialize the mock agent with pre-recorded nodes.
+
+        Args:
+            nodes: List of nodes from a previous agent run
+            result_output: Optional output to return as the final result
+        """
+        self.nodes = nodes
+        self.result_output = result_output
+
+    def save(self, path: Path) -> None:
+        """
+        Save the mock agent to a file.
+        Uses safe_deepcopy to ensure nodes are serializable.
+
+        Args:
+            path: Path to save the agent mock data
+        """
+        path.parent.mkdir(parents=True, exist_ok=True)
+        data = {"nodes": safe_deepcopy(self.nodes, use_cache=False), "result_output": self.result_output}  # pylint: disable=unexpected-keyword-arg
+        with open(path, "wb") as f:
+            pickle.dump(data, f)
+
+    @classmethod
+    def load(cls, path: Path) -> "AgentMock":
+        """
+        Load a mock agent from a file.
+
+        Args:
+            path: Path to load the agent mock data from
+
+        Returns:
+            AgentMock instance with loaded data
+        """
+        with open(path, "rb") as f:
+            data = pickle.load(f)
+        return cls(nodes=data["nodes"], result_output=data["result_output"])
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+
+    def iter(self, **kwargs) -> "AgentRunMock":  # pylint: disable=unused-argument
+        """
+        Create an async iterator that replays the recorded nodes.
+
+        Args:
+            prompt: The prompt (ignored in mock, kept for interface compatibility)
+            usage_limits: Usage limits (ignored in mock, kept for interface compatibility)
+
+        Returns:
+            An async iterator that yields the recorded nodes
+        """
+        return AgentRunMock(self.nodes, self.result_output)
+
+
+class AgentRunMock:
+    """
+    Mock agent run that yields pre-recorded nodes.
+    """
+
+    def __init__(self, nodes: list[Any], result_output: str | None = None):
+        """
+        Initialize the mock agent run.
+
+        Args:
+            nodes: List of nodes to replay
+            result_output: Optional output to return as the final result
+        """
+        self.nodes = nodes
+        self.result = MockResult(result_output)
+        self._index = 0
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+
+    def __aiter__(self):
+        """Return self as async iterator."""
+        return self
+
+    async def __anext__(self):
+        """Yield the next node from the recorded list."""
+        if self._index >= len(self.nodes):
+            raise StopAsyncIteration
+        node = self.nodes[self._index]
+        self._index += 1
+        return node
+
+
+class MockResult:  # pylint: disable=too-few-public-methods
+    """
+    Mock result object that mimics the agent result structure.
+    """
+
+    def __init__(self, output: str | None = None):
+        """
+        Initialize the mock result.
+
+        Args:
+            output: The output to return
+        """
+        self.output = output
+
+    def __str__(self):
+        return f"MockResult(output={self.output})"

aixtools/utils/utils.py

@@ -32,6 +32,18 @@ def find_file(path: Path, glob="*.pdf"):
     yield from path.rglob(glob)
 
 
+def is_multiline(s: str) -> bool:
+    """Check if a string is multiline."""
+    s = str(s).rstrip()
+    return "\n" in s
+
+
+def is_too_long(s: str, max_length: int = 200) -> bool:
+    """Check if a string is too long."""
+    s = str(s).rstrip()
+    return len(s) > max_length
+
+
 def prepend_all_lines(msg, prepend="\t", skip_first_line: bool = False) -> str:
     """Prepend all lines of a message with a prepend."""
     out = ""
@@ -59,7 +71,7 @@ def remove_quotes(s):
     """
     if s is None:
         return None
-    s = s.strip()
+    s = str(s).strip()
     while (
         (s.startswith('"') and s.endswith('"'))
         or (s.startswith("'") and s.endswith("'"))
@@ -76,9 +88,29 @@ def remove_quotes(s):
     return s
 
 
-def tabit(s: str, prefix="\t|") -> str:
+def tabit(s, prefix="\t|") -> str:
     """Add a prefix to each line of a string for improved readability."""
-    return prefix + str(s).replace("\n", f"\n{prefix}")
+    s = str(s)
+    return prefix + s.replace("\n", f"\n{prefix}")
+
+
+def to_json_pretty_print(obj) -> str:
+    """Convert to a pretty-printed JSON string if possible."""
+    if isinstance(obj, str):
+        # Already a string, try to parse as JSON, then pretty print
+        s = obj
+        try:
+            obj = json.loads(s)
+            return json.dumps(obj, indent=2)
+        except Exception as _:  # pylint: disable=broad-exception-caught
+            # Not a JSON string, return as is
+            return s
+    # Not a string, convert to pretty JSON
+    try:
+        return json.dumps(obj, indent=2)
+    except Exception as _:  # pylint: disable=broad-exception-caught
+        # Fallback to str representation
+        return str(obj)
 
 
 def to_str(data) -> str:

{aixtools-0.2.15.dist-info → aixtools-0.2.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.2.15
+Version: 0.2.17
 Summary: Tools for AI exploration and debugging
 Requires-Python: >=3.11.2
 Description-Content-Type: text/markdown
@@ -63,6 +63,7 @@ Testing Tools & Evals
 - Tool Doctor System - `aixtools/tools/doctor/`
 - Tool Recommendation Engine - `aixtools/tools/doctor/tool_recommendation.py`
 - FaultyMCP - `aixtools/mcp/faulty_mcp.py`
+- Agent Mock - `aixtools/testing/agent_mock.py`
 
 Databases
 - Database Integration - `aixtools/db/`
@@ -524,6 +525,21 @@ The framework includes a custom scoring system with [`average_assertions`](aixto
 
 AIXtools provides comprehensive testing utilities and diagnostic tools for AI agent development and debugging.
 
+### Running Tests
+
+Execute the test suite using the provided scripts:
+
+```bash
+# Run all tests
+./scripts/test.sh
+
+# Run unit tests only
+./scripts/test_unit.sh
+
+# Run integration tests only
+./scripts/test_integration.sh
+```
+
 ### Testing Utilities
 
 The testing module provides mock tools, model patching, and test utilities for comprehensive agent testing.
@@ -544,10 +560,6 @@ cached_response = cache.get_cached_response("test_prompt")
 test_model = AixTestModel()
 ```
 
-### Tool Doctor System
-
-Automated tool analysis and recommendation system for optimizing agent tool usage and analyzing MCP servers.
-
 #### MCP Tool Doctor
 
 Analyze tools from MCP (Model Context Protocol) servers and receive AI-powered recommendations for improvement.
@@ -586,7 +598,7 @@ tool_doctor_mcp --stdio-command fastmcp --stdio-args run my_server.py --debug
 # --debug                Enable debug output
 ```
 
-#### Traditional Tool Doctor
+#### Tool Doctor
 
 Analyze tool usage patterns from agent logs and get optimization recommendations.
 
@@ -681,19 +693,27 @@ python -m aixtools.mcp.faulty_mcp \
     --prob-on-get-crash 0.1
 ```
 
-### Running Tests
+### Agent Mock
 
-Execute the test suite using the provided scripts:
+Replay previously recorded agent runs without executing the actual agent. Useful for testing, debugging, and creating reproducible test cases.
 
-```bash
-# Run all tests
-./scripts/test.sh
+```python
+from aixtools.testing.agent_mock import AgentMock
+from aixtools.agents.agent import get_agent, run_agent
 
-# Run unit tests only
-./scripts/test_unit.sh
+# Run an agent and capture its execution
+agent = get_agent(system_prompt="You are a helpful assistant.")
+result, nodes = await run_agent(agent, "Explain quantum computing")
 
-# Run integration tests only
-./scripts/test_integration.sh
+# Create a mock agent from the recorded nodes
+agent_mock = AgentMock(nodes=nodes, result_output=result)
+
+# Save the mock for later use
+agent_mock.save(Path("test_data/quantum_mock.pkl"))
+
+# Load and replay the mock agent
+loaded_mock = AgentMock.load(Path("test_data/quantum_mock.pkl"))
+result, nodes = await run_agent(loaded_mock, "any prompt")  # Returns recorded nodes
 ```
 
 ## Chainlit & HTTP Server

{aixtools-0.2.15.dist-info → aixtools-0.2.17.dist-info}/RECORD

@@ -1,5 +1,5 @@
 aixtools/__init__.py,sha256=9NGHm7LjsQmsvjTZvw6QFJexSvAU4bCoN_KBk9SCa00,260
-aixtools/_version.py,sha256=fu32ckCvDvmdBESK9QY8685INg5f0hTCLuUhaNOVaJQ,706
+aixtools/_version.py,sha256=sRnPbdnyLakHrE7uBPRC_AQNPiFphtVIa4BPaftkqk4,706
 aixtools/app.py,sha256=JzQ0nrv_bjDQokllIlGHOV0HEb-V8N6k_nGQH-TEsVU,5227
 aixtools/chainlit.md,sha256=yC37Ly57vjKyiIvK4oUvf4DYxZCwH7iocTlx7bLeGLU,761
 aixtools/context.py,sha256=I_MD40ZnvRm5WPKAKqBUAdXIf8YaurkYUUHSVVy-QvU,598
@@ -25,8 +25,11 @@ aixtools/a2a/google_sdk/utils.py,sha256=4VIPV2GtG5IRY-KSZN6iRMS3ntxG2uKgd_d5tQvn
 aixtools/a2a/google_sdk/pydantic_ai_adapter/agent_executor.py,sha256=8VuU2WXeSHUK3_rRm_mjX6elqdC9NA2uz1aELzeC8BU,9784
 aixtools/a2a/google_sdk/pydantic_ai_adapter/storage.py,sha256=nGoVL7MPoZJW7iVR71laqpUYP308yFKZIifJtvUgpiU,878
 aixtools/agents/__init__.py,sha256=MAW196S2_G7uGqv-VNjvlOETRfuV44WlU1leO7SiR0A,282
-aixtools/agents/agent.py,sha256=JRbWWo3lueXyNSd9UgnTztm4bZr7dPxeNOzCFNFG7RE,7620
+aixtools/agents/agent.py,sha256=loBbSrRY-QZH8wec27JjNi_h1jhqkZ9-UrR0lqPNMWI,7725
 aixtools/agents/agent_batch.py,sha256=0Zu9yNCRPAQZPjXQ-dIUAmP1uGTVbxVt7xvnMpoJMjU,2251
+aixtools/agents/nodes_to_md.py,sha256=hAT8dgiZTG4uGoSgeRZkIJ7zgkQUNpdIr8KSFhjWAH0,7515
+aixtools/agents/nodes_to_message.py,sha256=ZqcmxUNf4esiCTRk37wWP1LquhqNsCmydvMr4kjZEjw,1012
+aixtools/agents/nodes_to_str.py,sha256=UkOu5Nry827J4H_ohQU3tPBfJxtr3p6FfCfWoUy5uIs,4325
 aixtools/agents/print_nodes.py,sha256=wVTngNfqM0As845WTRz6G3Rei_Gr3HuBlvu-G_eXuig,1665
 aixtools/agents/prompt.py,sha256=p9OYnyJ4-MyGXwHPrQeJBhZ2a3RV2HqhtdUUCrTMsAQ,3361
 aixtools/auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -52,10 +55,11 @@ aixtools/log_view/node_summary.py,sha256=EJjnBqdBWI-_bI-4nfTxwaost3mtiufb5cK7T54
 aixtools/logfilters/__init__.py,sha256=pTD8ujCqjPWBCeB7yv7lmCtnA2KXOnkIv0HExDagkXs,129
 aixtools/logfilters/context_filter.py,sha256=zR3Bnv3fCqXLeb7bCFTmlnWhC6dFIvUb-u712tOnUPk,2259
 aixtools/logging/__init__.py,sha256=b5oYyGQDUHHxhRtzqKUaQPv8hQeWw54rzDXSV8lDY1w,613
-aixtools/logging/log_objects.py,sha256=K8p2kHTqr9nntHQ0-jyNrDBoBiqH14gFnD7atAiN_do,6964
+aixtools/logging/log_objects.py,sha256=gohsgcfyr8vsY7G_hfmj973-Ek1_PN-bMMLEUA-4u6U,7708
 aixtools/logging/logging_config.py,sha256=LvxV3C75-I0096PpcCIbgM-Cp998LzWXeMM14HYbU20,4985
 aixtools/logging/mcp_log_models.py,sha256=7-H2GJXiiyLhpImuyLLftAGG4skxJal8Swax0ob04MY,3463
 aixtools/logging/mcp_logger.py,sha256=d2I5l4t0d6rQH17w23FpE1IUD8Ax-mSaKfByCH86q4I,6257
+aixtools/logging/mcp_middleware.py,sha256=0kpTAwvz9Fd_mFDP3J9ldH4dPP-bhcUjLJKELGOx0IQ,2257
 aixtools/logging/model_patch_logging.py,sha256=CW5-kKI-zNEgZhNV4vx3EQu6fbrEtX7VjA6fE5loRLQ,2916
 aixtools/logging/open_telemetry.py,sha256=fJjF1ou_8GyfNfbyWDQPGK6JAUrUaPwURYPHhXEtDBE,1121
 aixtools/mcp/__init__.py,sha256=Qp4uD1RtypCYgzWrt_ThBVxa5-CFBgcwMWkHbHFUQ54,232
@@ -71,6 +75,7 @@ aixtools/server/app_mounter.py,sha256=0tJ0tC140ezAjnYdlhpLJQjY-TO8NVw7D8LseYCCVY
 aixtools/server/path.py,sha256=nI4yRQcE6gjKx5GG3PmHf7iT1FelT6Q8Xhw4ol9O1e0,5219
 aixtools/server/utils.py,sha256=tZWITIx6M-luV9yve4j3rPtYGSSA6zWS0JWEAySne_M,2276
 aixtools/testing/__init__.py,sha256=mlmaAR2gmS4SbsYNCxnIprmFpFp-syjgVUkpUszo3mE,166
+aixtools/testing/agent_mock.py,sha256=D4DrRGaeKuWEhq8hdB5s9PdfLAXHMKiC3JA7M8yN13o,4244
 aixtools/testing/aix_test_model.py,sha256=i0xBdmpKoEfJHle6JDmcoJLUENN8Eqt181_WZ7XtDdU,6240
 aixtools/testing/mock_tool.py,sha256=4I0LxxSkLhGIKM2YxCP3cnYI8IYJjdKhfwGZ3dioXsM,2465
 aixtools/testing/model_patch_cache.py,sha256=238gKC_gSpR3BkeejhetObOkpOR1l2Iz3A6B_eUTRNc,10158
@@ -84,13 +89,13 @@ aixtools/utils/config_util.py,sha256=3Ya4Qqhj1RJ1qtTTykQ6iayf5uxlpigPXgEJlTi1wn4
 aixtools/utils/enum_with_description.py,sha256=zjSzWxG74eR4x7dpmb74pLTYCWNSMvauHd7_9LpDYIw,1088
 aixtools/utils/files.py,sha256=8JnxwHJRJcjWCdFpjzWmo0po2fRg8esj4H7sOxElYXU,517
 aixtools/utils/persisted_dict.py,sha256=0jQzV7oF-A6Or-HjcU6V7aMXWQL67SOKpULgmtFwAfg,3110
-aixtools/utils/utils.py,sha256=5911Ej1ES2NU_FKIWA3CWKhKnwgjvi1aDR2aiD6Xv3E,4880
+aixtools/utils/utils.py,sha256=EP4q50wuIwgICzv3unMqDS93HuuySxE7SezqxX4HdKg,5840
 aixtools/utils/chainlit/cl_agent_show.py,sha256=vaRuowp4BRvhxEr5hw0zHEJ7iaSF_5bo_9BH7pGPPpw,4398
 aixtools/utils/chainlit/cl_utils.py,sha256=fxaxdkcZg6uHdM8uztxdPowg3a2f7VR7B26VPY4t-3c,5738
 aixtools/vault/__init__.py,sha256=fsr_NuX3GZ9WZ7dGfe0gp_5-z3URxAfwVRXw7Xyc0dU,141
 aixtools/vault/vault.py,sha256=9dZLWdZQk9qN_Q9Djkofw9LUKnJqnrX5H0fGusVLBhA,6037
-aixtools-0.2.15.dist-info/METADATA,sha256=pMsq-5JO6i7tEE3rfFNwk2zD70H7qLLc8fzo1zObSTI,27230
-aixtools-0.2.15.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aixtools-0.2.15.dist-info/entry_points.txt,sha256=q8412TG4T0S8K0SKeWp2vkVPIDYQs0jNoHqcQ7qxOiA,155
-aixtools-0.2.15.dist-info/top_level.txt,sha256=wBn-rw9bCtxrR4AYEYgjilNCUVmKY0LWby9Zan2PRJM,9
-aixtools-0.2.15.dist-info/RECORD,,
+aixtools-0.2.17.dist-info/METADATA,sha256=dCP4ss1Kz-AZWMQcmaHxR-HfqrtlSEVhIdZgff8ncvU,27958
+aixtools-0.2.17.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aixtools-0.2.17.dist-info/entry_points.txt,sha256=q8412TG4T0S8K0SKeWp2vkVPIDYQs0jNoHqcQ7qxOiA,155
+aixtools-0.2.17.dist-info/top_level.txt,sha256=wBn-rw9bCtxrR4AYEYgjilNCUVmKY0LWby9Zan2PRJM,9
+aixtools-0.2.17.dist-info/RECORD,,