waldiez 0.3.6__tar.gz → 0.3.8__tar.gz

waldiez-0.3.8/NOTICE.md

@@ -0,0 +1,5 @@
+# NOTICE
+
+Copyright 2024 - 2025 Waldiez and contributors.
+
+This Work includes Software developed by Waldiez (<https://waldiez.io/>).

{waldiez-0.3.6 → waldiez-0.3.8}/PKG-INFO

@@ -1,11 +1,12 @@
 Metadata-Version: 2.4
 Name: waldiez
-Version: 0.3.6
+Version: 0.3.8
 Summary: waldiez
 Project-URL: homepage, https://waldiez.github.io/waldiez/python
 Project-URL: repository, https://github.com/waldiez/python.git
 Author-email: Panagiotis Kasnesis <pkasnesis@waldiez.io>, Lazaros Toumanidis <laztoum@waldiez.io>, Stella Ioannidou <stella@humancentered.gr>
 License-File: LICENSE
+License-File: NOTICE.md
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
@@ -17,33 +18,40 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Python: <3.13,>=3.10
+Requires-Dist: aiocsv==1.3.2
+Requires-Dist: aiofiles==24.1.0
+Requires-Dist: aiosqlite==0.20.0
 Requires-Dist: asyncer==0.0.8
+Requires-Dist: graphviz==0.20.3
 Requires-Dist: httpx<0.28.0,>=0.27.0
 Requires-Dist: jupytext
 Requires-Dist: pandas>=2
 Requires-Dist: parso==0.8.4
-Requires-Dist: pyautogen==0.7.1
+Requires-Dist: pyautogen==0.7.2
 Requires-Dist: pydantic<3,>=2.6.1
 Requires-Dist: typer<0.13,>=0.9
 Provides-Extra: ag2-extras
-Requires-Dist: autogen[captainagent]==0.7.1; (platform_system == 'Linux') and extra == 'ag2-extras'
+Requires-Dist: autogen[captainagent]==0.7.2; (platform_system == 'Linux') and extra == 'ag2-extras'
+Requires-Dist: chromadb; (platform_system != 'Linux') and extra == 'ag2-extras'
+Requires-Dist: huggingface-hub>=0.0.17; (platform_system != 'Linux') and extra == 'ag2-extras'
 Requires-Dist: pgvector>=0.3.6; extra == 'ag2-extras'
 Requires-Dist: psycopg[binary]>=3.2.3; extra == 'ag2-extras'
-Requires-Dist: pyautogen[anthropic]==0.7.1; extra == 'ag2-extras'
-Requires-Dist: pyautogen[bedrock]==0.7.1; extra == 'ag2-extras'
-Requires-Dist: pyautogen[gemini]==0.7.1; extra == 'ag2-extras'
-Requires-Dist: pyautogen[groq]==0.7.1; extra == 'ag2-extras'
-Requires-Dist: pyautogen[lmm]==0.7.1; extra == 'ag2-extras'
-Requires-Dist: pyautogen[mistral]==0.7.1; extra == 'ag2-extras'
-Requires-Dist: pyautogen[neo4j]==0.7.1; extra == 'ag2-extras'
-Requires-Dist: pyautogen[retrievechat-mongodb]==0.7.1; extra == 'ag2-extras'
-Requires-Dist: pyautogen[retrievechat-pgvector]==0.7.1; extra == 'ag2-extras'
-Requires-Dist: pyautogen[retrievechat-qdrant]==0.7.1; (python_version < '3.13') and extra == 'ag2-extras'
-Requires-Dist: pyautogen[retrievechat]==0.7.1; extra == 'ag2-extras'
-Requires-Dist: pyautogen[together]==0.7.1; extra == 'ag2-extras'
-Requires-Dist: pyautogen[websurfer]==0.7.1; extra == 'ag2-extras'
+Requires-Dist: pyautogen[anthropic]==0.7.2; extra == 'ag2-extras'
+Requires-Dist: pyautogen[bedrock]==0.7.2; extra == 'ag2-extras'
+Requires-Dist: pyautogen[gemini]==0.7.2; extra == 'ag2-extras'
+Requires-Dist: pyautogen[groq]==0.7.2; extra == 'ag2-extras'
+Requires-Dist: pyautogen[lmm]==0.7.2; extra == 'ag2-extras'
+Requires-Dist: pyautogen[mistral]==0.7.2; extra == 'ag2-extras'
+Requires-Dist: pyautogen[neo4j]==0.7.2; extra == 'ag2-extras'
+Requires-Dist: pyautogen[retrievechat-mongodb]==0.7.2; extra == 'ag2-extras'
+Requires-Dist: pyautogen[retrievechat-pgvector]==0.7.2; extra == 'ag2-extras'
+Requires-Dist: pyautogen[retrievechat-qdrant]==0.7.2; (python_version < '3.13') and extra == 'ag2-extras'
+Requires-Dist: pyautogen[retrievechat]==0.7.2; extra == 'ag2-extras'
+Requires-Dist: pyautogen[together]==0.7.2; extra == 'ag2-extras'
+Requires-Dist: pyautogen[websurfer]==0.7.2; extra == 'ag2-extras'
 Requires-Dist: pymongo==4.10.1; extra == 'ag2-extras'
 Requires-Dist: qdrant-client[fastembed]; (python_version >= '3.13') and extra == 'ag2-extras'
+Requires-Dist: sentence-transformers>=2.0.0; (platform_system != 'Linux') and extra == 'ag2-extras'
 Provides-Extra: dev
 Requires-Dist: autoflake==2.3.1; extra == 'dev'
 Requires-Dist: bandit==1.8.2; extra == 'dev'
@@ -56,7 +64,7 @@ Requires-Dist: pre-commit==4.1.0; extra == 'dev'
 Requires-Dist: pydocstyle==6.3.0; extra == 'dev'
 Requires-Dist: pylint==3.3.3; extra == 'dev'
 Requires-Dist: python-dotenv==1.0.1; extra == 'dev'
-Requires-Dist: ruff==0.9.2; extra == 'dev'
+Requires-Dist: ruff==0.9.3; extra == 'dev'
 Requires-Dist: toml; (python_version <= '3.10') and extra == 'dev'
 Requires-Dist: types-pyyaml==6.0.12.20241230; extra == 'dev'
 Requires-Dist: types-toml==0.10.8.20240310; extra == 'dev'
@@ -73,9 +81,9 @@ Requires-Dist: mkdocstrings-python==1.13.0; extra == 'docs'
 Requires-Dist: mkdocstrings[crystal,python]==0.27.0; extra == 'docs'
 Provides-Extra: jupyter
 Requires-Dist: jupyterlab>=4.3.0; extra == 'jupyter'
-Requires-Dist: waldiez-jupyter==0.3.6; extra == 'jupyter'
+Requires-Dist: waldiez-jupyter==0.3.7; extra == 'jupyter'
 Provides-Extra: studio
-Requires-Dist: waldiez-studio==0.3.6; extra == 'studio'
+Requires-Dist: waldiez-studio==0.3.7; extra == 'studio'
 Provides-Extra: test
 Requires-Dist: pytest-asyncio==0.25.2; extra == 'test'
 Requires-Dist: pytest-cov==6.0.0; extra == 'test'

{waldiez-0.3.6 → waldiez-0.3.8}/pyproject.toml

@@ -22,9 +22,13 @@ classifiers = [
   "Intended Audience :: Developers",
 ]
 dependencies =[
+    "aiocsv==1.3.2",
+    "aiofiles==24.1.0",
+    "aiosqlite==0.20.0",
     "asyncer==0.0.8",
+    "graphviz==0.20.3",
     "pandas>=2",
-    "pyautogen==0.7.1",
+    "pyautogen==0.7.2",
     "jupytext",
     "parso==0.8.4",
     "pydantic>=2.6.1,<3",
@@ -49,10 +53,10 @@ exclude = [ "**/example.py" ]
 
 [project.optional-dependencies]
 studio = [
-    "waldiez_studio==0.3.6",
+    "waldiez_studio==0.3.7",
 ]
 jupyter = [
-    "waldiez_jupyter==0.3.6",
+    "waldiez_jupyter==0.3.7",
     "jupyterlab>=4.3.0",
 ]
 dev = [
@@ -67,7 +71,7 @@ dev = [
     'pydocstyle==6.3.0',
     'pylint==3.3.3',
     'python-dotenv==1.0.1',
-    'ruff==0.9.2',
+    'ruff==0.9.3',
     'types-PyYAML==6.0.12.20241230',
     'yamllint==1.35.1',
     'toml; python_version <= "3.10"',
@@ -99,22 +103,25 @@ dev = [
 ag2_extras =[
     'pgvector>=0.3.6',
     'psycopg[binary]>=3.2.3',
-    'pyautogen[anthropic]==0.7.1',
-    'pyautogen[bedrock]==0.7.1',
-    'autogen[captainagent]==0.7.1; platform_system == "Linux"',
-    'pyautogen[gemini]==0.7.1',
-    'pyautogen[groq]==0.7.1',
-    'pyautogen[lmm]==0.7.1',
-    'pyautogen[mistral]==0.7.1',
-    'pyautogen[neo4j]==0.7.1',
-    'pyautogen[retrievechat-mongodb]==0.7.1',
-    'pyautogen[retrievechat-pgvector]==0.7.1',
-    'pyautogen[retrievechat-qdrant]==0.7.1; python_version < "3.13"',
+    'pyautogen[anthropic]==0.7.2',
+    'pyautogen[bedrock]==0.7.2',
+    'autogen[captainagent]==0.7.2; platform_system == "Linux"',
+    'chromadb; platform_system != "Linux"',
+    'sentence-transformers>=2.0.0; platform_system != "Linux"',
+    'huggingface-hub>=0.0.17; platform_system != "Linux"',
+    'pyautogen[gemini]==0.7.2',
+    'pyautogen[groq]==0.7.2',
+    'pyautogen[lmm]==0.7.2',
+    'pyautogen[mistral]==0.7.2',
+    'pyautogen[neo4j]==0.7.2',
+    'pyautogen[retrievechat-mongodb]==0.7.2',
+    'pyautogen[retrievechat-pgvector]==0.7.2',
+    'pyautogen[retrievechat-qdrant]==0.7.2; python_version < "3.13"',
     'qdrant-client[fastembed]; python_version >= "3.13"',
-    'pyautogen[retrievechat]==0.7.1',
-    'pyautogen[together]==0.7.1',
-    'pyautogen[together]==0.7.1',
-    'pyautogen[websurfer]==0.7.1',
+    'pyautogen[retrievechat]==0.7.2',
+    'pyautogen[together]==0.7.2',
+    'pyautogen[together]==0.7.2',
+    'pyautogen[websurfer]==0.7.2',
     'pymongo==4.10.1',
 ]
 test = [
@@ -347,7 +354,7 @@ addopts = """
 [tool.coverage.run]
 omit = [
   "example.py",
-  "cli_extras.py"
+  "**/cli_extras/*.py",
 ]
 [tool.coverage.report]
 exclude_lines = [

waldiez-0.3.8/waldiez/__init__.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Waldiez package."""
+
+from ._version import __version__
+from .exporter import WaldiezExporter
+from .models import Waldiez
+from .runner import WaldiezRunner
+from .utils import check_conflicts, check_flaml_warnings
+
+__WALDIEZ_INITIALIZED = False
+
+if not __WALDIEZ_INITIALIZED:
+    check_conflicts()
+    check_flaml_warnings()
+    # let's skip the one below
+    # check_pysqlite3()
+    # and use it only if needed:
+    #   captain_agent dependency:
+    #   before calling pip install pyautogen[captainagent]
+    #   we should have pysqlite3 installed (at least on windows)
+    # before running a flow
+    __WALDIEZ_INITIALIZED = True
+
+__all__ = [
+    "Waldiez",
+    "WaldiezExporter",
+    "WaldiezRunner",
+    "__version__",
+]

{waldiez-0.3.6 → waldiez-0.3.8}/waldiez/_version.py

@@ -2,4 +2,4 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Version information for Waldiez."""
 
-__version__ = "0.3.6"
+__version__ = "0.3.8"

{waldiez-0.3.6 → waldiez-0.3.8}/waldiez/cli.py

@@ -7,7 +7,6 @@ import json
 import logging
 import os
 import sys
-import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
@@ -15,14 +14,11 @@ import anyio
 import typer
 from typing_extensions import Annotated
 
-from . import Waldiez, __version__
-from .cli_extras import add_cli_extras  # type: ignore
+from ._version import __version__
 from .exporter import WaldiezExporter
+from .models import Waldiez
 from .runner import WaldiezRunner
-
-warnings.filterwarnings(
-    "ignore", module="flaml", message="^.*flaml.automl is not available.*$"
-)
+from .utils import add_cli_extras
 
 if TYPE_CHECKING:
     from autogen import ChatResult  # type: ignore[import-untyped]
@@ -90,7 +86,7 @@ def run(
     ),
 ) -> None:
     """Run a Waldiez flow."""
-    # swarm without a user,
+    # a swarm chat without a user agent
     # creates a new user (this has a default code execution with docker)
     # temp (until we handle/detect docker setup)
     os.environ["AUTOGEN_USE_DOCKER"] = "0"
@@ -114,6 +110,13 @@ def run(
             _log_result(result, logger)
             sep = "-" * 80
             print("\n" + f"{sep}" + "\n")
+    elif isinstance(results, dict):
+        logger.info("Results:")
+        for key, result in results.items():
+            logger.info("Order: %s", key)
+            _log_result(result, logger)
+            sep = "-" * 80
+            print("\n" + f"{sep}" + "\n")
     else:
         _log_result(results, logger)
 

{waldiez-0.3.6 → waldiez-0.3.8}/waldiez/exporting/__init__.py

@@ -2,11 +2,13 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Tools for exporting agents, models, skills and chats to strings."""
 
+from .agent import AgentExporter
 from .flow import FlowExporter
 from .models import ModelsExporter
 from .skills import SkillsExporter
 
 __all__ = [
+    "AgentExporter",
     "FlowExporter",
     "ModelsExporter",
     "SkillsExporter",

{waldiez-0.3.6 → waldiez-0.3.8}/waldiez/exporting/agent/agent_exporter.py

@@ -24,6 +24,7 @@ from .utils import (
     get_group_manager_extras,
     get_is_termination_message,
     get_rag_user_extras,
+    get_reasoning_agent_extras,
     get_swarm_extras,
 )
 
@@ -114,6 +115,10 @@ class AgentExporter(BaseExporter, ExporterMixin):
         self._termination = get_is_termination_message(
             agent=self.agent, agent_name=self._agent_name
         )
+        self._reasoning = get_reasoning_agent_extras(
+            agent=self.agent,
+            serializer=self.serializer,
+        )
 
     def get_imports(self) -> Optional[List[Tuple[str, ImportPosition]]]:
         """Get the imports.
@@ -135,7 +140,10 @@ class AgentExporter(BaseExporter, ExporterMixin):
         # if the agent has skills, add the register_function import.
         if self.agent.data.skills:
             agent_imports.add("from autogen import register_function")
-        return [(import_string, position) for import_string in agent_imports]
+        return sorted(
+            [(import_string, position) for import_string in agent_imports],
+            key=lambda x: x[0],
+        )
 
     def get_system_message_arg(self) -> str:
         """Get the system message argument.
@@ -223,6 +231,7 @@ class AgentExporter(BaseExporter, ExporterMixin):
             default_auto_reply = (
                 f'"{self.string_escape(agent.data.agent_default_auto_reply)}"'
             )
+        extras = f"{group_chat_arg}{retrieve_arg}{self._reasoning}"
         agent_str = f"""{agent_name} = {agent_class}(
     name="{agent_name}",
     description="{agent.description}"{system_message_arg},
@@ -230,7 +239,7 @@ class AgentExporter(BaseExporter, ExporterMixin):
     max_consecutive_auto_reply={agent.data.max_consecutive_auto_reply},
     default_auto_reply={default_auto_reply},
     code_execution_config={code_execution_arg},
-    is_termination_msg={is_termination},{group_chat_arg}{retrieve_arg}
+    is_termination_msg={is_termination},{extras}
 """
         if self._swarm[1]:
             agent_str += self._swarm[1]

{waldiez-0.3.6 → waldiez-0.3.8}/waldiez/exporting/agent/utils/__init__.py

@@ -7,6 +7,7 @@ from .agent_imports import get_agent_imports
 from .code_execution import get_agent_code_execution_config
 from .group_manager import get_group_manager_extras
 from .rag_user import get_rag_user_extras
+from .reasoning import get_reasoning_agent_extras
 from .swarm_agent import get_swarm_extras
 from .teachability import get_agent_teachability_string
 from .termination_message import get_is_termination_message
@@ -19,5 +20,6 @@ __all__ = [
     "get_group_manager_extras",
     "get_is_termination_message",
     "get_rag_user_extras",
+    "get_reasoning_agent_extras",
     "get_swarm_extras",
 ]

{waldiez-0.3.6 → waldiez-0.3.8}/waldiez/exporting/agent/utils/agent_class_name.py

@@ -31,4 +31,6 @@ def get_agent_class_name(agent: WaldiezAgent) -> str:
         return "RetrieveUserProxyAgent"
     if agent.agent_type == "swarm":
         return "SwarmAgent"
+    if agent.agent_type == "reasoning":
+        return "ReasoningAgent"
     return "ConversableAgent"  # pragma: no cover

{waldiez-0.3.6 → waldiez-0.3.8}/waldiez/exporting/agent/utils/agent_imports.py

@@ -45,6 +45,11 @@ def get_agent_imports(agent_class: str) -> Set[str]:
             "SwarmAgent, "
             "SwarmResult"
         )
+    elif agent_class == "ReasoningAgent":
+        # pylint: disable=line-too-long
+        imports.add(
+            "from autogen.agentchat.contrib.reasoning_agent import ReasoningAgent, visualize_tree"  # noqa: E501
+        )
     else:
         imports.add("from autogen import ConversableAgent")
     return imports

waldiez-0.3.8/waldiez/exporting/agent/utils/reasoning.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Export reasoning agent to string."""
+
+from typing import Callable
+
+from waldiez.models import WaldiezAgent, WaldiezReasoningAgent
+
+
+def get_reasoning_agent_extras(
+    agent: WaldiezAgent,
+    serializer: Callable[..., str],
+) -> str:
+    """Get the reasoning agent extras.
+
+    Parameters
+    ----------
+    agent : WaldiezReasoningAgent
+        The reasoning agent.
+    serializer : Callable[..., str]
+        The serializer to get the string representation of an object.
+
+    Returns
+    -------
+    str
+        The reasoning agent extras.
+    """
+    if agent.agent_type != "reasoning" or not isinstance(
+        agent, WaldiezReasoningAgent
+    ):
+        return ""
+    reasoning_config = agent.get_reasoning_config()
+    serialized = serializer(reasoning_config)
+    content = "\n    verbose=" + f"{agent.data.verbose},"
+    content += "\n    reason_config=" + f"{serialized},"
+    return content

{waldiez-0.3.6 → waldiez-0.3.8}/waldiez/exporting/flow/flow_exporter.py

@@ -49,13 +49,13 @@ from .utils import (
     ensure_unique_names,
     gather_agent_outputs,
     gather_imports,
+    get_after_run_content,
     get_def_main,
     get_ipynb_content_start,
-    get_logging_start_string,
-    get_logging_stop_string,
     get_py_content_start,
     get_sqlite_out,
-    get_sqlite_out_call,
+    get_start_logging,
+    get_stop_logging,
     get_the_imports_string,
 )
 
@@ -201,6 +201,7 @@ class FlowExporter(BaseExporter, ExporterMixin):
         str
             The merged export contents.
         """
+        is_async = self.waldiez.is_async
         content = (
             get_py_content_start(self.waldiez)
             if not self.for_notebook
@@ -209,7 +210,8 @@ class FlowExporter(BaseExporter, ExporterMixin):
         content += self.get_comment("imports", self.for_notebook) + "\n"
         content += imports[0] + "\n"
         content += self.get_comment("logging", self.for_notebook) + "\n"
-        content += get_logging_start_string(tabs=0) + "\n\n"
+        content += get_start_logging(tabs=0) + "\n"
+        content += "start_logging()\n\n"
         if models_output:
             content += self.get_comment("models", self.for_notebook) + "\n"
             content += models_output + "\n"
@@ -221,16 +223,27 @@ class FlowExporter(BaseExporter, ExporterMixin):
             content += agents_content + "\n"
         if before_chats:
             content += before_chats + "\n"
-        content += get_sqlite_out() + "\n"
+        content += get_sqlite_out(is_async=is_async) + "\n"
+        content += get_stop_logging(tabs=0, is_async=is_async) + "\n"
         content += self.get_comment("run", self.for_notebook) + "\n"
+        after_run = get_after_run_content(
+            waldiez=self.waldiez,
+            agent_names=self.agent_names,
+            tabs=0 if self.for_notebook else 1,
+        )
         if self.for_notebook is False:
             content += get_def_main(
-                chats_content, is_async=self.waldiez.is_async
+                chats_content,
+                after_run=after_run,
+                is_async=self.waldiez.is_async,
             )
         else:
             content += "\n" + chats_content + "\n"
-            content += get_logging_stop_string(tabs=0) + "\n"
-            content += get_sqlite_out_call(tabs=0) + "\n"
+            if is_async:
+                content += "await stop_logging()"
+            else:
+                content += "stop_logging()"
+                content += after_run
         content = content.replace("\n\n\n\n", "\n\n\n")
         return content
 

{waldiez-0.3.6 → waldiez-0.3.8}/waldiez/exporting/flow/utils/__init__.py

@@ -11,7 +11,11 @@ from .agent_utils import (
 )
 from .chat_utils import add_after_chat_content, add_before_chat_content
 from .def_main import get_def_main
-from .flow_content import get_ipynb_content_start, get_py_content_start
+from .flow_content import (
+    get_after_run_content,
+    get_ipynb_content_start,
+    get_py_content_start,
+)
 from .flow_names import ensure_unique_names
 from .importing_utils import (
     gather_imports,
@@ -19,10 +23,10 @@ from .importing_utils import (
     get_the_imports_string,
 )
 from .logging_utils import (
-    get_logging_start_string,
-    get_logging_stop_string,
     get_sqlite_out,
     get_sqlite_out_call,
+    get_start_logging,
+    get_stop_logging,
 )
 
 __all__ = [
@@ -35,11 +39,12 @@ __all__ = [
     "ensure_unique_names",
     "gather_agent_outputs",
     "gather_imports",
+    "get_after_run_content",
     "get_def_main",
     "get_py_content_start",
     "get_ipynb_content_start",
-    "get_logging_start_string",
-    "get_logging_stop_string",
+    "get_start_logging",
+    "get_stop_logging",
     "get_sqlite_out",
     "get_sqlite_out_call",
     "get_standard_imports",

{waldiez-0.3.6 → waldiez-0.3.8}/waldiez/exporting/flow/utils/def_main.py

@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
+# flake8: noqa: E501
+# pylint: disable=inconsistent-quotes, line-too-long
 """Get the main function."""
 
-from .logging_utils import get_logging_stop_string, get_sqlite_out_call
 
-
-def get_def_main(flow_chats: str, is_async: bool) -> str:
+def get_def_main(flow_chats: str, after_run: str, is_async: bool) -> str:
     """Get the main function.
 
     When exporting to python, waldiez_chats string will be the
@@ -18,6 +18,8 @@ def get_def_main(flow_chats: str, is_async: bool) -> str:
     ----------
     flow_chats : str
         The content of the main function.
+    after_run : str
+        The content after the run of the flow.
     is_async : bool
         Whether the main function is asynchronous.
     Returns
@@ -29,31 +31,34 @@ def get_def_main(flow_chats: str, is_async: bool) -> str:
     if is_async:
         content += "async "
     content += "def main():\n"
-    content += "    # type: () -> Union[ChatResult, List[ChatResult]]\n"
+    content += "    # type: () -> Union[ChatResult, List[ChatResult], Dict[int, ChatResult]]\n"
     content += '    """Start chatting."""\n'
     content += f"{flow_chats}" + "\n"
-    content += get_logging_stop_string(1) + "\n"
-    content += get_sqlite_out_call(1) + "\n"
-    content += "    return results\n\n\n"
+    if is_async:
+        content += "    await stop_logging()"
+    else:
+        content += "    stop_logging()"
+    content += after_run
+    content += "\n    return results\n\n\n"
     if is_async:
         content += "async def call_main():\n"
     else:
         content += "def call_main() -> None:\n"
     content += '    """Run the main function and print the results."""\n'
-    # fmt: off
+    content += "    results: Union[ChatResult, List[ChatResult], Dict[int, ChatResult]] = "
     if is_async:
-        content += (
-            "    results: Union[ChatResult, List[ChatResult]] = await main()\n"
-        )
-    else:
-        content += (
-            "    results: Union[ChatResult, List[ChatResult]] = main()\n"
-        )
-    # fmt: on
-    content += "    if not isinstance(results, list):\n"
-    content += "        results = [results]\n"
-    content += "    for result in results:\n"
-    content += "        pprint(asdict(result))\n\n\n"
+        content += "await "
+    content += "main()\n"
+    content += "    if isinstance(results, dict):\n"
+    content += "        # order by key\n"
+    content += "        ordered_results = dict(sorted(results.items()))\n"
+    content += "        for _, result in ordered_results.items():\n"
+    content += "            pprint(asdict(result))\n"
+    content += "    elif isinstance(results, list):\n"
+    content += "        for result in results:\n"
+    content += "            pprint(asdict(result))\n"
+    content += "    else:\n"
+    content += "        pprint(asdict(results))\n"
     content += 'if __name__ == "__main__":\n'
     if is_async:
         content += "    anyio.run(call_main)\n"

{waldiez-0.3.6 → waldiez-0.3.8}/waldiez/exporting/flow/utils/flow_content.py

@@ -1,8 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
+# flake8: noqa E501
+# pylint: disable=line-too-long
 """Utils to generate the content of a flow."""
 
-from typing import Callable, List, Optional
+from typing import Callable, Dict, List, Optional
 
 from waldiez.models import Waldiez
 
@@ -55,9 +57,8 @@ def get_ipynb_content_start(
     content += "import sys\n"
     requirements = " ".join(waldiez.requirements)
     if requirements:
-        # pylint: disable=line-too-long
         # fmt: off
-        content += "# " + f"!{{sys.executable}} -m pip install -q {requirements}" + "\n"  # noqa: E501
+        content += "# " + f"!{{sys.executable}} -m pip install -q {requirements}" + "\n"
         # fmt: on
     content += "# flake8: noqa E501"
     content += get_pylint_ignore_comment(True)
@@ -110,3 +111,50 @@ def get_pylint_ignore_comment(
     if notebook is True:
         line = "\n" + line
     return line + "\n"
+
+
+def get_after_run_content(
+    waldiez: Waldiez,
+    agent_names: Dict[str, str],
+    tabs: int,
+) -> str:
+    """Get content to add after the flow is run.
+
+    Parameters
+    ----------
+    waldiez : Waldiez
+        The waldiez object.
+    agent_names : Dict[str, str]
+        The dictionary of agent names and their corresponding ids
+    tabs : int
+        The number of tabs to add before the content.
+    Returns
+    -------
+    str
+        The content to add after the flow is run.
+    """
+    # if th eflow has reasoning agents, we add
+    # visualize_tree(agent._root) for each agent
+    content = ""
+    space = "    " * tabs
+    for agent in waldiez.agents:
+        if agent.agent_type == "reasoning":
+            agent_name = agent_names[agent.id]
+            content += f"""
+{space}# pylint: disable=broad-except,too-many-try-statements
+{space}try:
+{space}{space}visualize_tree({agent_name}._root)  # pylint: disable=protected-access
+{space}{space}if os.path.exists("tree_of_thoughts.png"):
+{space}{space}{space}new_name = "{agent_name}_tree_of_thoughts.png"
+{space}{space}{space}os.rename("tree_of_thoughts.png", new_name)
+{space}except BaseException:
+{space}{space}pass
+{space}# save the tree to json
+{space}try:
+{space}{space}data = {agent_name}._root.to_dict()  # pylint: disable=protected-access
+{space}{space}with open("{agent_name}_reasoning_tree.json", "w", encoding="utf-8") as f:
+{space}{space}{space}json.dump(data, f)
+{space}except BaseException:
+{space}{space}pass
+"""
+    return content