waldiez 0.4.8__tar.gz → 0.4.9__tar.gz

{waldiez-0.4.8 → waldiez-0.4.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waldiez
-Version: 0.4.8
+Version: 0.4.9
 Dynamic: Keywords
 Summary: Make AG2 Agents Collaborate: Drag, Drop, and Orchestrate with Waldiez.
 Project-URL: Homepage, https://waldiez.io
@@ -157,15 +157,15 @@ Requires-Dist: natsort==8.4.0; extra == 'docs'
 Provides-Extra: jupyter
 Requires-Dist: jupyter-server==2.16.0; extra == 'jupyter'
 Requires-Dist: jupyterlab<5.0,>=4.3.0; extra == 'jupyter'
-Requires-Dist: waldiez-jupyter==0.4.8; extra == 'jupyter'
+Requires-Dist: waldiez-jupyter==0.4.9; extra == 'jupyter'
 Provides-Extra: mqtt
 Requires-Dist: paho-mqtt<3.0,>=2.1.0; extra == 'mqtt'
 Provides-Extra: redis
 Requires-Dist: ag2[redis]==0.9.2; extra == 'redis'
 Provides-Extra: runner
-Requires-Dist: waldiez-runner==0.4.8; (python_version >= '3.11') and extra == 'runner'
+Requires-Dist: waldiez-runner==0.4.9; (python_version >= '3.11') and extra == 'runner'
 Provides-Extra: studio
-Requires-Dist: waldiez-studio==0.4.8; extra == 'studio'
+Requires-Dist: waldiez-studio==0.4.9; extra == 'studio'
 Provides-Extra: test
 Requires-Dist: ag2[redis]==0.9.2; extra == 'test'
 Requires-Dist: ag2[websockets]==0.9.2; extra == 'test'

{waldiez-0.4.8 → waldiez-0.4.9}/pyproject.toml

@@ -81,15 +81,15 @@ mqtt = [
 ]
 # jupyterlab extension
 jupyter = [
-  "waldiez_jupyter==0.4.8",
+  "waldiez_jupyter==0.4.9",
   "jupyterlab>=4.3.0,<5.0",
   "jupyter_server==2.16.0",
 ]
 runner = [
-  "waldiez_runner==0.4.8; python_version >= '3.11'",
+  "waldiez_runner==0.4.9; python_version >= '3.11'",
 ]
 studio = [
-  "waldiez_studio==0.4.8",
+  "waldiez_studio==0.4.9",
 ]
 ag2_extras = [
   # models

{waldiez-0.4.8 → waldiez-0.4.9}/waldiez/_version.py

@@ -5,4 +5,4 @@
 This file is automatically generated by Hatchling.
 Do not edit this file directly.
 """
-__version__ = VERSION = "0.4.8"
+__version__ = VERSION = "0.4.9"

{waldiez-0.4.8 → waldiez-0.4.9}/waldiez/cli.py

@@ -2,43 +2,25 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 # flake8: noqa: E501
 # pylint: disable=missing-function-docstring,missing-param-doc,missing-raises-doc
+# pylint: disable=line-too-long
 """Command line interface to convert or run a waldiez file."""
 
 import json
 import os
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import Any, Optional
 
 import anyio
 import typer
 from dotenv import load_dotenv
 from typing_extensions import Annotated
 
-# pylint: disable=import-error,line-too-long
-# pyright: reportMissingImports=false
-try:  # pragma: no cover
-    # fmt: off
-    from ._version import __version__ # type: ignore[unused-ignore, unused-import, import-not-found, import-untyped]  # noqa
-    # fmt: on
-except ImportError:  # pragma: no cover
-    import warnings
-
-    warnings.warn(
-        "Importing __version__ failed. Using 'dev' as version.", stacklevel=2
-    )
-    __version__ = "dev"
-
-
 from .exporter import WaldiezExporter
 from .io import StructuredIOStream
 from .logger import get_logger
 from .models import Waldiez
 from .runner import WaldiezRunner
-from .utils import add_cli_extras
-
-if TYPE_CHECKING:
-    from autogen import ChatResult  # type: ignore[import-untyped]
-
+from .utils import add_cli_extras, get_waldiez_version
 
 load_dotenv()
 LOG = get_logger()
@@ -71,7 +53,7 @@ def show_version(
 ) -> None:
     """Show the version of the Waldiez package and exit."""
     if version:
-        typer.echo(f"waldiez version: {__version__}")
+        typer.echo(f"waldiez version: {get_waldiez_version()}")
         raise typer.Exit()
 
 
@@ -130,27 +112,12 @@ def run(
         except json.decoder.JSONDecodeError as error:
             typer.echo("Invalid .waldiez file. Not a valid json?")
             raise typer.Exit(code=1) from error
-    results = _do_run(
+    _do_run(
         data,
         structured=structured,
         uploads_root=uploads_root,
         output_path=output_path,
     )
-    if isinstance(results, list):
-        LOG.info("Results:")
-        for result in results:
-            _log_result(result)
-            sep = "-" * 80
-            print("\n" + f"{sep}" + "\n")
-    elif isinstance(results, dict):
-        LOG.info("Results:")
-        for key, result in results.items():
-            LOG.info("Order: %s", key)
-            _log_result(result)
-            sep = "-" * 80
-            print("\n" + f"{sep}" + "\n")
-    else:
-        _log_result(results)
 
 
 @app.command()
@@ -239,7 +206,7 @@ def _do_run(
     structured: bool,
     uploads_root: Optional[Path],
     output_path: Optional[Path],
-) -> Union["ChatResult", list["ChatResult"], dict[int, "ChatResult"]]:
+) -> None:
     """Run the Waldiez flow and get the results."""
     waldiez = Waldiez.from_dict(data)
     if structured:
@@ -247,30 +214,29 @@ def _do_run(
         with StructuredIOStream.set_default(stream):
             runner = WaldiezRunner(waldiez)
             if waldiez.is_async:
-                results = anyio.run(
+                anyio.run(
                     runner.a_run,
                     output_path,
                     uploads_root,
                 )
             else:
-                results = runner.run(
+                runner.run(
                     output_path=output_path,
                     uploads_root=uploads_root,
                 )
     else:
         runner = WaldiezRunner(waldiez)
         if waldiez.is_async:
-            results = anyio.run(
+            anyio.run(
                 runner.a_run,
                 output_path,
                 uploads_root,
             )
         else:
-            results = runner.run(
+            runner.run(
                 output_path=output_path,
                 uploads_root=uploads_root,
             )
-    return results
 
 
 def _get_output_path(output: Optional[Path], force: bool) -> Optional[Path]:
@@ -286,16 +252,6 @@ def _get_output_path(output: Optional[Path], force: bool) -> Optional[Path]:
     return output
 
 
-def _log_result(result: "ChatResult") -> None:
-    """Log the result of the Waldiez flow."""
-    LOG.info("Chat History:\n")
-    LOG.info(result.chat_history)
-    LOG.info("Summary:\n")
-    LOG.info(result.summary)
-    LOG.info("Cost:\n")
-    LOG.info(result.cost)
-
-
 add_cli_extras(app)
 
 if __name__ == "__main__":

{waldiez-0.4.8 → waldiez-0.4.9}/waldiez/exporting/agent/extras/rag/chroma_extras.py

@@ -65,15 +65,13 @@ def _get_chroma_embedding_function_string(
     to_import = ""
     embedding_function_arg = ""
     embedding_function_content = ""
-    vector_db_model = agent.retrieve_config.db_config.model
     if not agent.retrieve_config.use_custom_embedding:
         to_import = (
             "from chromadb.utils.embedding_functions."
             "sentence_transformer_embedding_function import "
             "SentenceTransformerEmbeddingFunction"
         )
-        embedding_function_arg = "SentenceTransformerEmbeddingFunction("
-        embedding_function_arg += f'model_name="{vector_db_model}")'
+        embedding_function_arg = f"{agent_name}_embedding_function"
     else:
         embedding_function_content, embedding_function_arg = (
             agent.retrieve_config.get_custom_embedding_function(
@@ -124,21 +122,35 @@ def get_chroma_db_args(
     # manually initializing the collection before running the flow,
     # might be a workaround.
     content_before = f"{agent_name}_client = {client_str}" + "\n"
+    vector_db_model = agent.retrieve_config.db_config.model
+    if not embedding_function_body:
+        content_before += (
+            f"{agent_name}_embedding_function = "
+            "SentenceTransformerEmbeddingFunction(\n"
+            f'    model_name="{vector_db_model}",' + "\n"
+            ")\n"
+        )
     collection_name = agent.retrieve_config.collection_name
     get_or_create = agent.retrieve_config.get_or_create
     if collection_name:
         if get_or_create:
             content_before += (
-                f"{agent_name}_client.get_or_create_collection("
-                f'"{collection_name}")' + "\n"
+                f"{agent_name}_client.get_or_create_collection" + "(\n"
+                f'    "{collection_name}",' + "\n"
+                f"    embedding_function={embedding_function_arg}," + "\n"
+                ")" + "\n"
             )
         else:
             content_before += (
                 "try:\n"
-                f'    {agent_name}_client.get_collection("{collection_name}")'
-                + "\n"
+                f"    {agent_name}_client.get_collection(" + "\n"
+                f'        "{collection_name}",' + "\n"
+                f"        embedding_function={embedding_function_arg}," + "\n"
+                "    )\n"
                 "except ValueError:\n"
-                f"    {agent_name}_client.create_collection("
-                f'"{collection_name}")' + "\n"
+                f"    {agent_name}_client.create_collection(" + "\n"
+                f'        "{collection_name}",' + "\n"
+                f"        embedding_function={embedding_function_arg}," + "\n"
+                "    )\n"
             )
     return kwarg_string, to_import, embedding_function_body, content_before

{waldiez-0.4.8 → waldiez-0.4.9}/waldiez/io/structured.py

@@ -36,9 +36,13 @@ class StructuredIOStream(IOStream):
     uploads_root: Path | None = None
 
     def __init__(
-        self, timeout: float = 120, uploads_root: Path | str | None = None
+        self,
+        timeout: float = 120,
+        uploads_root: Path | str | None = None,
+        is_async: bool = False,
     ) -> None:
         self.timeout = timeout
+        self.is_async = is_async
         if uploads_root is not None:
             self.uploads_root = Path(uploads_root).resolve()
             if not self.uploads_root.exists():
@@ -106,6 +110,22 @@ class StructuredIOStream(IOStream):
             uploads_root=self.uploads_root,
             base_name=request_id,
         )
+        # let's keep an eye here:
+        # https://github.com/ag2ai/ag2/blob/main/autogen/agentchat/conversable_agent.py#L2973
+        # reply = await iostream.input(prompt) ???? (await???)
+        if self.is_async:
+            # let's make a coroutine to just return the user response
+            async def async_input() -> str:
+                """Asynchronous input method.
+
+                Returns
+                -------
+                str
+                    The line read from the input stream.
+                """
+                return user_response
+
+            return async_input()  # type: ignore
         return user_response
 
     # noinspection PyMethodMayBeStatic

{waldiez-0.4.8 → waldiez-0.4.9}/waldiez/runner.py

@@ -18,7 +18,6 @@ from pathlib import Path
 from types import ModuleType, TracebackType
 from typing import (
     TYPE_CHECKING,
-    Callable,
     Optional,
     Set,
     Type,
@@ -34,7 +33,6 @@ from .running import (
     after_run,
     before_run,
     chdir,
-    get_printer,
     install_requirements,
     refresh_environment,
     reset_env_vars,
@@ -166,25 +164,22 @@ class WaldiezRunner:
     def install_requirements(self) -> None:
         """Install the requirements for the flow."""
         self._called_install_requirements = True
-        printer = get_printer()
         extra_requirements = self.gather_requirements()
         if extra_requirements:
-            install_requirements(extra_requirements, printer)
+            install_requirements(extra_requirements)
 
     async def a_install_requirements(self) -> None:
         """Install the requirements for the flow asynchronously."""
         self._called_install_requirements = True
-        printer = get_printer()
         extra_requirements = self.gather_requirements()
         if extra_requirements:
-            await a_install_requirements(extra_requirements, printer)
+            await a_install_requirements(extra_requirements)
 
     def _before_run(
         self,
         temp_dir: Path,
         file_name: str,
         module_name: str,
-        printer: Callable[..., None],
     ) -> tuple[ModuleType, dict[str, str]]:
         self._exporter.export(Path(file_name))
         # unique_names = self._exporter.context.get_unique_names()
@@ -197,8 +192,8 @@ class WaldiezRunner:
         old_vars = set_env_vars(self.waldiez.get_flow_env_vars())
         module = importlib.util.module_from_spec(spec)
         spec.loader.exec_module(module)
-        printer("<Waldiez> - Starting workflow...")
-        printer(self.waldiez.info.model_dump_json())
+        print("<Waldiez> - Starting workflow...")
+        print(self.waldiez.info.model_dump_json())
         return module, old_vars
 
     def _run(
@@ -232,8 +227,7 @@ class WaldiezRunner:
         if not self._called_install_requirements:
             self.install_requirements()
         refresh_environment()
-        printer = get_printer()
-        printer(
+        print(
             "Requirements installed.\n"
             "NOTE: If new packages were added and you are using Jupyter, "
             "you might need to restart the kernel."
@@ -246,7 +240,6 @@ class WaldiezRunner:
                 temp_dir=temp_dir,
                 file_name=file_name,
                 module_name=module_name,
-                printer=printer,
             )
             if use_structured_io:
                 stream = StructuredIOStream(
@@ -257,13 +250,12 @@ class WaldiezRunner:
                     results = module.main()
             else:
                 results = module.main()
-            printer("<Waldiez> - Workflow finished")
+            print("<Waldiez> - Workflow finished")
             sys.path.pop(0)
             reset_env_vars(old_vars)
         after_run(
             temp_dir=temp_dir,
             output_path=output_path,
-            printer=printer,
             flow_name=self.waldiez.name,
             skip_mmd=skip_mmd,
         )
@@ -285,8 +277,7 @@ class WaldiezRunner:
         if not self._called_install_requirements:
             await self.a_install_requirements()
         refresh_environment()
-        printer = get_printer()
-        printer(
+        print(
             "Requirements installed.\n"
             "NOTE: If new packages were added and you are using Jupyter, "
             "you might need to restart the kernel."
@@ -299,7 +290,6 @@ class WaldiezRunner:
                 temp_dir=temp_dir,
                 file_name=file_name,
                 module_name=module_name,
-                printer=printer,
             )
             if use_structured_io:
                 stream = StructuredIOStream(
@@ -315,7 +305,6 @@ class WaldiezRunner:
         after_run(
             temp_dir=temp_dir,
             output_path=output_path,
-            printer=printer,
             flow_name=self.waldiez.name,
             skip_mmd=skip_mmd,
         )

{waldiez-0.4.8 → waldiez-0.4.9}/waldiez/running/__init__.py

@@ -9,23 +9,29 @@ from .environment import (
     reset_env_vars,
     set_env_vars,
 )
-from .running import (
-    a_chdir,
+from .post_run import after_run
+from .pre_run import (
     a_install_requirements,
-    after_run,
     before_run,
-    chdir,
-    get_printer,
     install_requirements,
 )
+from .util import (
+    a_chdir,
+    chdir,
+    create_async_subprocess,
+    create_sync_subprocess,
+    strip_ansi,
+)
 
 __all__ = [
     "a_chdir",
     "a_install_requirements",
     "after_run",
     "before_run",
+    "create_async_subprocess",
+    "create_sync_subprocess",
+    "strip_ansi",
     "chdir",
-    "get_printer",
     "in_virtualenv",
     "is_root",
     "install_requirements",

waldiez-0.4.9/waldiez/running/post_run.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Utilities for running code."""
+
+import datetime
+import shutil
+from pathlib import Path
+from typing import Optional, Union
+
+from .gen_seq_diagram import generate_sequence_diagram
+
+
+def after_run(
+    temp_dir: Path,
+    output_path: Optional[Union[str, Path]],
+    flow_name: str,
+    skip_mmd: bool = False,
+) -> None:
+    """Actions to perform after running the flow.
+
+    Parameters
+    ----------
+    temp_dir : Path
+        The temporary directory.
+    output_path : Optional[Union[str, Path]]
+        The output path.
+    flow_name : str
+        The flow name.
+    skip_mmd : bool, optional
+        Whether to skip the mermaid sequence diagram generation,
+        by default, False
+    """
+    if isinstance(output_path, str):
+        output_path = Path(output_path)
+    mmd_dir = output_path.parent if output_path 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_path
+                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
+    if output_path:
+        destination_dir = output_path.parent
+        destination_dir = (
+            destination_dir
+            / "waldiez_out"
+            / datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+        )
+        destination_dir.mkdir(parents=True, exist_ok=True)
+        # copy the contents of the temp dir to the destination dir
+        print(f"Copying the results to {destination_dir}")
+        copy_results(
+            temp_dir=temp_dir,
+            output_path=output_path,
+            output_dir=output_path.parent,
+            destination_dir=destination_dir,
+        )
+    shutil.rmtree(temp_dir)
+
+
+def copy_results(
+    temp_dir: Path,
+    output_path: Path,
+    output_dir: Path,
+    destination_dir: Path,
+) -> None:
+    """Copy the results to the output directory.
+
+    Parameters
+    ----------
+    temp_dir : Path
+        The temporary directory.
+    output_path : Path
+        The output path.
+    output_dir : Path
+        The output directory.
+    destination_dir : Path
+        The destination directory.
+    """
+    temp_dir.mkdir(parents=True, exist_ok=True)
+    for item in temp_dir.iterdir():
+        # skip cache files
+        if (
+            item.name.startswith("__pycache__")
+            or item.name.endswith(".pyc")
+            or item.name.endswith(".pyo")
+            or item.name.endswith(".pyd")
+            or item.name == ".cache"
+        ):
+            continue
+        if item.is_file():
+            # let's also copy the "tree of thoughts" image
+            # to the output directory
+            if item.name.endswith("tree_of_thoughts.png") or item.name.endswith(
+                "reasoning_tree.json"
+            ):
+                shutil.copy(item, output_dir / item.name)
+            shutil.copy(item, destination_dir)
+        else:
+            shutil.copytree(item, destination_dir / item.name)
+    if output_path.is_file():
+        if output_path.suffix == ".waldiez":
+            output_path = output_path.with_suffix(".py")
+        if output_path.suffix == ".py":
+            src = temp_dir / output_path.name
+            if src.exists():
+                dst = destination_dir / output_path.name
+                if dst.exists():
+                    dst.unlink()
+                shutil.copyfile(src, output_dir / output_path.name)