waldiez 0.5.2__py3-none-any.whl → 0.5.4__py3-none-any.whl

waldiez/running/environment.py

@@ -1,8 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
-# pylint: disable=import-outside-toplevel,reimported
+# pylint: disable=import-outside-toplevel,reimported,line-too-long
+# flake8: noqa: E501, F401
+# pyright: reportUnusedImport=false
 """Environment related utilities."""
 
+import importlib
 import os
 import site
 import sys
@@ -53,12 +56,14 @@ def refresh_environment() -> None:
     # default code execution with docker
     # temp (until we handle/detect docker setup)
     os.environ["AUTOGEN_USE_DOCKER"] = "0"
+    os.environ["ANONYMIZED_TELEMETRY"] = "False"
     try_handle_the_np_thing()
     reload_autogen()
+    reload_chroma_if_needed()
 
 
 # pylint: disable=too-complex,too-many-try-statements,unused-import
-def reload_autogen() -> None:  # noqa: C901
+def reload_autogen() -> None:  # noqa: C901  # pragma: no cover
     """Reload the autogen package.
 
     Try to avoid "please install package x" errors
@@ -86,31 +91,38 @@ def reload_autogen() -> None:  # noqa: C901
         default_io_stream = IOStream.get_default()
     except (ImportError, AttributeError):
         pass
-
+    autogen_modules = sorted(
+        [
+            name
+            for name in sys.modules
+            if name.startswith("autogen.")
+            and not name.startswith("autogen.io")
+            and not name.startswith("autogen.tools")
+        ],
+        key=len,
+        reverse=True,  # Longer names (deeper modules) first
+    )
     try:
         # Remove autogen modules in reverse dependency order
-        autogen_modules = sorted(
-            [
-                name
-                for name in sys.modules
-                if name.startswith("autogen.")
-                and not name.startswith("autogen.io")
-                and not name.startswith("autogen.tools")
-            ],
-            key=len,
-            reverse=True,  # Longer names (deeper modules) first
-        )
         for mod_name in autogen_modules:
             if mod_name in sys.modules:
                 del sys.modules[mod_name]
 
         if "autogen" in sys.modules:
             del sys.modules["autogen"]
-
+        site.main()  # Rebuild the site module cache
         # Re-import autogen
         # pylint: disable=unused-import
         import autogen  # pyright: ignore
 
+        for mod_name in sorted(autogen_modules, key=len, reverse=False):
+            # Re-import each module
+            try:
+                importlib.import_module(mod_name)
+            except ImportError as e:
+                # If a module fails to import, we can log it or handle it
+                print(f"Failed to re-import {mod_name}: {e}", file=sys.stderr)
+
         # Restore IOStream state if we had it
         if default_io_stream is not None:
             try:
@@ -125,11 +137,43 @@ def reload_autogen() -> None:  # noqa: C901
         # If reload fails, at least try to re-import autogen
         try:
             import autogen  # type: ignore  # noqa: F401
-        except ImportError:
+
+            for mod_name in sorted(autogen_modules, key=len, reverse=False):
+                # Re-import each module
+                try:
+                    importlib.import_module(mod_name)
+                except ImportError as err:
+                    # If a module fails to import, we can log it or handle it
+                    print(
+                        f"Failed to re-import {mod_name}: {err}",
+                        file=sys.stderr,
+                    )
+        except Exception:  # pylint: disable=broad-exception-caught
             pass
         raise e
 
 
+def reload_chroma_if_needed() -> None:  # pragma: no cover
+    """Reload the chroma package if it is installed."""
+    cheomadb_modules = [
+        name
+        for name in sys.modules
+        if name.startswith("chromadb.") or name == "chromadb"
+    ]
+    if not cheomadb_modules:
+        return
+
+    for mod_name in sorted(cheomadb_modules, key=len, reverse=True):
+        # Remove chromadb modules in reverse dependency order
+        if mod_name in sys.modules:
+            del sys.modules[mod_name]
+    try:
+        import chromadb  # type: ignore[unused-ignore, import-not-found, import-untyped]
+    except ImportError:
+        # If chromadb is not installed, we can ignore this
+        return
+
+
 def try_handle_the_np_thing() -> None:
     """Try to handle the numpy deprecation warning."""
     # we might get:
@@ -141,7 +185,7 @@ def try_handle_the_np_thing() -> None:
     os.environ["NPY_PROMOTION_STATE"] = "weak"
     import numpy as np
 
-    if not hasattr(np, "_no_pep50_warning"):
+    if not hasattr(np, "_no_pep50_warning"):  # pragma: no branch
         import contextlib
 
         @contextlib.contextmanager

waldiez/running/post_run.py

@@ -5,11 +5,14 @@
 """Utilities for running code."""
 
 import datetime
+import json
 import shutil
 from pathlib import Path
 from typing import Optional, Union
 
 from .gen_seq_diagram import generate_sequence_diagram
+from .patch_io_stream import get_printer
+from .timeline_processor import TimelineProcessor
 
 
 def after_run(
@@ -18,6 +21,7 @@ def after_run(
     flow_name: str,
     uploads_root: Optional[Path] = None,
     skip_mmd: bool = False,
+    skip_timeline: bool = False,
 ) -> None:
     """Actions to perform after running the flow.
 
@@ -34,25 +38,21 @@ def after_run(
     skip_mmd : bool, optional
         Whether to skip the mermaid sequence diagram generation,
         by default, False
+    skip_timeline : bool, optional
+        Whether to skip the timeline processing, by default False
     """
     if isinstance(output_file, str):
         output_file = Path(output_file)
     mmd_dir = output_file.parent if output_file else Path.cwd()
     if skip_mmd is False:
-        events_csv_path = temp_dir / "logs" / "events.csv"
-        if events_csv_path.exists():
-            print("Generating mermaid sequence diagram...")
-            mmd_path = temp_dir / f"{flow_name}.mmd"
-            generate_sequence_diagram(events_csv_path, mmd_path)
-            if (
-                not output_file
-                and mmd_path.exists()
-                and mmd_path != mmd_dir / f"{flow_name}.mmd"
-            ):
-                try:
-                    shutil.copyfile(mmd_path, mmd_dir / f"{flow_name}.mmd")
-                except BaseException:  # pylint: disable=broad-exception-caught
-                    pass
+        _make_mermaid_diagram(
+            temp_dir=temp_dir,
+            output_file=output_file,
+            flow_name=flow_name,
+            mmd_dir=mmd_dir,
+        )
+    if skip_timeline is False:
+        _make_timeline_json(temp_dir=temp_dir)
     if output_file:
         destination_dir = output_file.parent
         destination_dir = (
@@ -71,6 +71,62 @@ def after_run(
     shutil.rmtree(temp_dir)
 
 
+def _make_mermaid_diagram(
+    temp_dir: Path,
+    output_file: Optional[Union[str, Path]],
+    flow_name: str,
+    mmd_dir: Path,
+) -> None:
+    events_csv_path = temp_dir / "logs" / "events.csv"
+    if events_csv_path.exists():
+        print("Generating mermaid sequence diagram...")
+        mmd_path = temp_dir / f"{flow_name}.mmd"
+        generate_sequence_diagram(events_csv_path, mmd_path)
+        if (
+            not output_file
+            and mmd_path.exists()
+            and mmd_path != mmd_dir / f"{flow_name}.mmd"
+        ):
+            try:
+                shutil.copyfile(mmd_path, mmd_dir / f"{flow_name}.mmd")
+            except BaseException:  # pylint: disable=broad-exception-caught
+                pass
+
+
+def _make_timeline_json(
+    temp_dir: Path,
+) -> None:
+    """Make the timeline JSON file."""
+    events_csv_path = temp_dir / "logs" / "events.csv"
+    if events_csv_path.exists():
+        log_files = TimelineProcessor.get_files(temp_dir / "logs")
+        if any(log_files.values()):
+            output_file = temp_dir / "timeline.json"
+            # pylint: disable=too-many-try-statements
+            try:
+                processor = TimelineProcessor()
+                processor.load_csv_files(
+                    agents_file=log_files["agents"],
+                    chat_file=log_files["chat"],
+                    events_file=log_files["events"],
+                    functions_file=log_files["functions"],
+                )
+                results = processor.process_timeline()
+                with open(output_file, "w", encoding="utf-8") as f:
+                    json.dump(results, f, indent=2, default=str)
+                short_results = TimelineProcessor.get_short_results(results)
+                printer = get_printer()
+                printer(
+                    json.dumps(
+                        {"type": "timeline", "content": short_results},
+                        default=str,
+                    ),
+                    flush=True,
+                )
+            except BaseException:  # pylint: disable=broad-exception-caught
+                pass
+
+
 def copy_results(
     temp_dir: Path,
     output_file: Path,

waldiez/running/pre_run.py

@@ -15,6 +15,7 @@ from .utils import strip_ansi
 
 def install_requirements(
     extra_requirements: set[str],
+    upgrade: bool = False,
     printer: Callable[..., None] = print,
 ) -> None:
     """Install the requirements.
@@ -23,6 +24,8 @@ def install_requirements(
     ----------
     extra_requirements : set[str]
         The extra requirements.
+    upgrade : bool, optional
+        Whether to upgrade the requirements, by default False.
     printer : Callable[..., None]
         The printer function to use, defaults to print.
     """
@@ -38,6 +41,8 @@ def install_requirements(
         os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = "1"
         if not is_root():
             pip_install.append("--user")
+    if upgrade:
+        pip_install.append("--upgrade")
     pip_install.extend(extra_requirements)
     # pylint: disable=too-many-try-statements
     try:
@@ -63,6 +68,7 @@ def install_requirements(
 
 async def a_install_requirements(
     extra_requirements: set[str],
+    upgrade: bool = False,
     printer: Callable[..., None] = print,
 ) -> None:
     """Install the requirements asynchronously.
@@ -71,6 +77,8 @@ async def a_install_requirements(
     ----------
     extra_requirements : set[str]
         The extra requirements.
+    upgrade : bool, optional
+        Whether to upgrade the requirements, by default False.
     printer : Callable[..., None]
         The printer function to use, defaults to print.
     """
@@ -83,6 +91,8 @@ async def a_install_requirements(
         os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = "1"
         if not is_root():
             pip_install.extend(["--user"])
+    if upgrade:
+        pip_install.append("--upgrade")
     pip_install.extend(extra_requirements)
     # pylint: disable=too-many-try-statements
     try:
@@ -103,3 +113,35 @@ async def a_install_requirements(
                 os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = break_system_packages
             else:
                 del os.environ["PIP_BREAK_SYSTEM_PACKAGES"]
+
+
+def install_waldiez(
+    upgrade: bool = True,
+    printer: Callable[..., None] = print,
+) -> None:
+    """Install Waldiez.
+
+    Parameters
+    ----------
+    upgrade : bool, optional
+        Whether to upgrade Waldiez, by default True.
+    printer : Callable[..., None]
+        The printer function to use, defaults to print.
+    """
+    install_requirements({"waldiez"}, upgrade, printer)
+
+
+async def a_install_waldiez(
+    upgrade: bool = True,
+    printer: Callable[..., None] = print,
+) -> None:
+    """Install Waldiez asynchronously.
+
+    Parameters
+    ----------
+    upgrade : bool, optional
+        Whether to upgrade Waldiez, by default True.
+    printer : Callable[..., None]
+        The printer function to use, defaults to print.
+    """
+    await a_install_requirements({"waldiez"}, upgrade, printer)

waldiez/running/protocol.py

@@ -7,7 +7,10 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Protocol, Union, runtime_checkable
 
 if TYPE_CHECKING:
-    from autogen import ChatResult  # type: ignore[import-untyped]
+    from autogen.io.run_response import (  # type: ignore[import-untyped]
+        AsyncRunResponseProtocol,
+        RunResponseProtocol,
+    )
 
 
 @runtime_checkable
@@ -59,7 +62,6 @@ class WaldiezRunnerProtocol(Protocol):
         output_path: str | Path | None,
         uploads_root: str | Path | None,
         structured_io: bool | None = None,
-        skip_patch_io: bool | None = None,
         skip_mmd: bool = False,
     ) -> None:
         """Start running the Waldiez flow in a non-blocking way.
@@ -74,9 +76,6 @@ class WaldiezRunnerProtocol(Protocol):
             The runtime uploads root.
         structured_io : bool
             Whether to use structured IO instead of the default 'input/print'.
-        skip_patch_io : bool | None
-            Whether to skip patching I/O, by default None.
-            If None, it will use the value from the context.
         skip_mmd : bool
             Whether to skip generating the mermaid diagram.
 
@@ -91,7 +90,6 @@ class WaldiezRunnerProtocol(Protocol):
         output_path: str | Path | None,
         uploads_root: str | Path | None,
         structured_io: bool | None = None,
-        skip_patch_io: bool | None = None,
         skip_mmd: bool = False,
     ) -> None:
         """Asynchronously start running the Waldiez flow in a non-blocking way.
@@ -106,9 +104,6 @@ class WaldiezRunnerProtocol(Protocol):
             The runtime uploads root.
         structured_io : bool
             Whether to use structured IO instead of the default 'input/print'.
-        skip_patch_io : bool | None
-            Whether to skip patching I/O, by default None.
-            If None, it will use the value from the context.
         skip_mmd : bool
             Whether to skip generating the mermaid diagram.
 
@@ -123,13 +118,10 @@ class WaldiezRunnerProtocol(Protocol):
         output_path: str | Path | None,
         uploads_root: str | Path | None,
         structured_io: bool | None = None,
-        threaded: bool | None = None,
-        skip_patch_io: bool | None = None,
         skip_mmd: bool = False,
     ) -> Union[
-        "ChatResult",
-        list["ChatResult"],
-        dict[int, "ChatResult"],
+        list["RunResponseProtocol"],
+        list["AsyncRunResponseProtocol"],
     ]:  # pyright: ignore
         """Run the Waldiez flow in a blocking way.
 
@@ -141,20 +133,17 @@ class WaldiezRunnerProtocol(Protocol):
             The runtime uploads root.
         structured_io : bool
             Whether to use structured IO instead of the default 'input/print'.
-        threaded : bool | None
-            Whether to run the flow in a separate thread.
-        skip_patch_io : bool
-            Whether to skip patching I/O, by default None.
-            If None, it will use the value from the context.
         skip_mmd : bool
             Whether to skip generating the mermaid diagram.
 
         Returns
         -------
-        Union[ChatResult, list[ChatResult], dict[int, ChatResult]]
-            The result of the run, which can be a single ChatResult,
-            a list of ChatResults,
-            or a dictionary mapping indices to ChatResults.
+        Union[
+            list["RunResponseProtocol"],
+            list["AsyncRunResponseProtocol"],
+        ]
+            The result of the run, which can be a list of RunResponseProtocol
+            or a list of AsyncRunResponseProtocol.
         """
 
     async def a_run(
@@ -162,12 +151,10 @@ class WaldiezRunnerProtocol(Protocol):
         output_path: str | Path | None,
         uploads_root: str | Path | None,
         structured_io: bool | None = None,
-        skip_patch_io: bool | None = None,
         skip_mmd: bool = False,
     ) -> Union[
-        "ChatResult",
-        list["ChatResult"],
-        dict[int, "ChatResult"],
+        list["RunResponseProtocol"],
+        list["AsyncRunResponseProtocol"],
     ]:  # pyright: ignore
         """Run the Waldiez flow.
 
@@ -179,40 +166,41 @@ class WaldiezRunnerProtocol(Protocol):
             The runtime uploads root.
         structured_io : bool
             Whether to use structured IO instead of the default 'input/print'.
-        skip_patch_io : bool
-            Whether to skip patching I/O, by default None.
-            If None, it will use the value from the context.
         skip_mmd : bool
             Whether to skip generating the mermaid diagram.
 
         Returns
         -------
-        Union[ChatResult, list[ChatResult], dict[int, ChatResult]]
-            The result of the run, which can be a single ChatResult,
-            a list of ChatResults,
-            or a dictionary mapping indices to ChatResults.
+        Union[
+            list["RunResponseProtocol"],
+            list["AsyncRunResponseProtocol"],
+        ]
+            The result of the run, which can be a list of RunResponseProtocol
+            or a list of AsyncRunResponseProtocol.
         """
 
     def after_run(
         self,
         results: Union[
-            "ChatResult",
-            list["ChatResult"],
-            dict[int, "ChatResult"],
+            list["RunResponseProtocol"],
+            list["AsyncRunResponseProtocol"],
         ],
         output_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> None:
         """Actions to perform after running the flow.
 
         Parameters
         ----------
-        results : Union[ChatResult, list[ChatResult], dict[int, ChatResult]]
-            The results of the run, which can be a single ChatResult,
-            a list of ChatResults,
-            or a dictionary mapping indices to ChatResults.
+        results : Union[
+            list["RunResponseProtocol"],
+            list["AsyncRunResponseProtocol"],
+        ]
+            The results of the run, which can be a list of RunResponseProtocol
+            or a list of AsyncRunResponseProtocol.
         output_file : Path
             The path to the output file.
         uploads_root : Path | None
@@ -221,28 +209,32 @@ class WaldiezRunnerProtocol(Protocol):
             The path to the temporary directory.
         skip_mmd : bool
             Whether to skip generating the mermaid diagram.
+        skip_timeline : bool
+            Whether to skip generating the timeline JSON.
         """
 
     async def a_after_run(
         self,
         results: Union[
-            "ChatResult",
-            list["ChatResult"],
-            dict[int, "ChatResult"],
+            list["RunResponseProtocol"],
+            list["AsyncRunResponseProtocol"],
         ],
         output_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> None:
         """Asynchronously perform actions after running the flow.
 
         Parameters
         ----------
-        results : Union[ChatResult, list[ChatResult], dict[int, ChatResult]]
-            The results of the run, which can be a single ChatResult,
-            a list of ChatResults,
-            or a dictionary mapping indices to ChatResults.
+        results : Union[
+            list["RunResponseProtocol"],
+            list["AsyncRunResponseProtocol"]
+        ]
+            The results of the run, which can be a list of RunResponseProtocol
+            or a list of AsyncRunResponseProtocol.
         output_file : Path
             The path to the output file.
         uploads_root : Path | None
@@ -251,6 +243,8 @@ class WaldiezRunnerProtocol(Protocol):
             The path to the temporary directory.
         skip_mmd : bool
             Whether to skip generating the mermaid diagram.
+        skip_timeline : bool
+            Whether to skip generating the timeline JSON.
         """
 
     def is_running(self) -> bool:  # pyright: ignore

waldiez/running/run_results.py

@@ -6,17 +6,17 @@
 from typing import TYPE_CHECKING, TypedDict, Union
 
 if TYPE_CHECKING:
-    from autogen import ChatResult  # type: ignore[import-untyped]
+    from autogen.io.run_response import (  # type: ignore[import-untyped]
+        AsyncRunResponseProtocol,
+        RunResponseProtocol,
+    )
 
 
 class WaldiezRunResults(TypedDict):
     """Results of the Waldiez run."""
 
     results: Union[
-        "ChatResult",
-        list["ChatResult"],
-        dict[int, "ChatResult"],
-        None,
+        list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
     ]
     exception: Exception | None
     completed: bool