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

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,63 @@ 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():
+        print("Processing timeline...")
+        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

@@ -204,6 +204,7 @@ class WaldiezRunnerProtocol(Protocol):
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> None:
         """Actions to perform after running the flow.
 
@@ -221,6 +222,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.
         """
 
     async def a_after_run(
@@ -234,6 +237,7 @@ class WaldiezRunnerProtocol(Protocol):
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> None:
         """Asynchronously perform actions after running the flow.
 
@@ -251,6 +255,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/subprocess_runner.py

@@ -49,6 +49,7 @@ class WaldiezSubprocessRunner(WaldiezBaseRunner):
         output_file: Path,
         uploads_root: Path | None,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> Union[
         "ChatResult",
         list["ChatResult"],
@@ -63,6 +64,7 @@ class WaldiezSubprocessRunner(WaldiezBaseRunner):
         output_file: Path,
         uploads_root: Path | None,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> Union[
         "ChatResult",
         list["ChatResult"],
@@ -77,6 +79,7 @@ class WaldiezSubprocessRunner(WaldiezBaseRunner):
         output_file: Path,
         uploads_root: Path | None,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> None:
         """Start the Waldiez workflow."""
         # This method should be implemented to start the workflow
@@ -88,6 +91,7 @@ class WaldiezSubprocessRunner(WaldiezBaseRunner):
         output_file: Path,
         uploads_root: Path | None,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> None:
         """Start the Waldiez workflow asynchronously."""
         # This method should be implemented to start the workflow asynchronously