aixtools 0.2.14__tar.gz → 0.2.16__tar.gz

{aixtools-0.2.14 → aixtools-0.2.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.2.14
+Version: 0.2.16
 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.14 → aixtools-0.2.16}/README.md

@@ -30,6 +30,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/`
@@ -491,6 +492,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.
@@ -511,10 +527,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.
@@ -553,7 +565,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.
 
@@ -648,19 +660,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.14 → aixtools-0.2.16}/aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.14'
-__version_tuple__ = version_tuple = (0, 2, 14)
+__version__ = version = '0.2.16'
+__version_tuple__ = version_tuple = (0, 2, 16)
 
-__commit_id__ = commit_id = 'gafbc9d3a4'
+__commit_id__ = commit_id = 'gbf6656c3c'

{aixtools-0.2.14 → aixtools-0.2.16}/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-0.2.16/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-0.2.16/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-0.2.16/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-0.2.14 → aixtools-0.2.16}/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,42 +79,50 @@ 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
 
 
-def save_objects_to_logfile(objects: list, log_dir=LOGS_DIR):
+async 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)
+            await object_logger.log(obj)
 
 
 class BaseLogger:
@@ -144,7 +155,7 @@ class ObjectLogger(BaseLogger):
         log_dir=LOGS_DIR,
         verbose: bool = True,
         debug: bool | None = None,
-        parent_logger: Union["ObjectLogger", NoneType] = None,
+        parent_logger: Union["BaseLogger", NoneType] = None,
     ):
         self.verbose = verbose
         self.debug = (
@@ -195,8 +206,9 @@ class ObjectLogger(BaseLogger):
                 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
+                if self.file is not None:
+                    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())