waldiez 0.4.8__py3-none-any.whl → 0.4.11__py3-none-any.whl

waldiez/__init__.py

@@ -5,7 +5,7 @@
 from .exporter import WaldiezExporter
 from .models import Waldiez
 from .runner import WaldiezRunner
-from .utils import check_conflicts, check_flaml_warnings
+from .utils import check_conflicts
 
 # flake8: noqa: F401
 # pylint: disable=import-error,line-too-long
@@ -28,7 +28,6 @@ __waldiez_initialized = False
 if not __waldiez_initialized:
     __waldiez_initialized = True
     check_conflicts()
-    check_flaml_warnings()
 
 __all__ = [
     "Waldiez",

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.11"

waldiez/cli.py

@@ -2,43 +2,24 @@
 # 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 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 +52,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()
 
 
@@ -115,6 +96,20 @@ def run(
             "If set, running the flow will use structured io stream instead of the default 'input/print' "
         ),
     ),
+    threaded: bool = typer.Option(  # noqa: B008
+        False,
+        help=(
+            "If set, the flow will be run in a separate thread. "
+            "This is useful for running flows that require user input or print output."
+        ),
+    ),
+    patch_io: bool = typer.Option(  # noqa: B008
+        False,
+        help=(
+            "If set, the flow will patch ag2's IOStream to safe print and input methods. "
+            "This is useful for running flows that require user input or print output."
+        ),
+    ),
     force: bool = typer.Option(  # noqa: B008
         False,
         help="Override the output file if it already exists.",
@@ -124,33 +119,35 @@ def run(
     os.environ["AUTOGEN_USE_DOCKER"] = "0"
     os.environ["NEP50_DISABLE_WARNING"] = "1"
     output_path = _get_output_path(output, force)
-    with file.open("r", encoding="utf-8") as _file:
-        try:
-            data = json.load(_file)
-        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(
-        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")
+    try:
+        runner = WaldiezRunner.load(file)
+    except FileNotFoundError as error:
+        typer.echo(f"File not found: {file}")
+        raise typer.Exit(code=1) from error
+    except json.decoder.JSONDecodeError as error:
+        typer.echo("Invalid .waldiez file. Not a valid json?")
+        raise typer.Exit(code=1) from error
+    except ValueError as error:
+        typer.echo(f"Invalid .waldiez file: {error}")
+        raise typer.Exit(code=1) from error
+    if runner.is_async:
+        anyio.run(
+            runner.a_run,
+            output_path,
+            uploads_root,
+            structured,
+            not patch_io,  # skip_patch_io
+            False,  # skip_mmd
+        )
     else:
-        _log_result(results)
+        runner.run(
+            output_path=output_path,
+            uploads_root=uploads_root,
+            structured_io=structured,
+            threaded=threaded,
+            skip_patch_io=not patch_io,
+            skip_mmd=False,
+        )
 
 
 @app.command()
@@ -234,43 +231,50 @@ def check(
     LOG.success("Waldiez flow seems valid.")
 
 
-def _do_run(
-    data: dict[str, Any],
-    structured: bool,
-    uploads_root: Optional[Path],
-    output_path: Optional[Path],
-) -> Union["ChatResult", list["ChatResult"], dict[int, "ChatResult"]]:
-    """Run the Waldiez flow and get the results."""
-    waldiez = Waldiez.from_dict(data)
-    if structured:
-        stream = StructuredIOStream(uploads_root=uploads_root)
-        with StructuredIOStream.set_default(stream):
-            runner = WaldiezRunner(waldiez)
-            if waldiez.is_async:
-                results = anyio.run(
-                    runner.a_run,
-                    output_path,
-                    uploads_root,
-                )
-            else:
-                results = runner.run(
-                    output_path=output_path,
-                    uploads_root=uploads_root,
-                )
-    else:
-        runner = WaldiezRunner(waldiez)
-        if waldiez.is_async:
-            results = anyio.run(
-                runner.a_run,
-                output_path,
-                uploads_root,
-            )
-        else:
-            results = runner.run(
-                output_path=output_path,
-                uploads_root=uploads_root,
-            )
-    return results
+# def _do_run(
+#     data: dict[str, Any],
+#     structured: bool,
+#     uploads_root: Optional[Path],
+#     output_path: Optional[Path],
+# ) -> None:
+#     """Run the Waldiez flow and get the results."""
+#     waldiez = Waldiez.from_dict(data)
+#     if structured:
+#         stream = StructuredIOStream(uploads_root=uploads_root)
+#         with StructuredIOStream.set_default(stream):
+#             runner = WaldiezRunner(waldiez)
+#             if waldiez.is_async:
+#                 anyio.run(
+#                     runner.a_run,
+#                     output_path,
+#                     uploads_root,
+#                     True,  # structured_io
+#                     False,  # skip_mmd
+#                 )
+#             else:
+#                 runner.run(
+#                     output_path=output_path,
+#                     uploads_root=uploads_root,
+#                     structured_io=True,
+#                     skip_mmd=False,
+#                 )
+#     else:
+#         runner = WaldiezRunner(waldiez)
+#         if waldiez.is_async:
+#             anyio.run(
+#                 runner.a_run,
+#                 output_path,
+#                 uploads_root,
+#                 False,  # structured_io
+#                 False,  # skip_mmd
+#             )
+#         else:
+#             runner.run(
+#                 output_path=output_path,
+#                 uploads_root=uploads_root,
+#                 structured_io=False,
+#                 skip_mmd=False,
+#             )
 
 
 def _get_output_path(output: Optional[Path], force: bool) -> Optional[Path]:
@@ -286,16 +290,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/exporter.py

@@ -11,7 +11,6 @@ to trigger the chat(s).
 """
 
 from pathlib import Path
-from typing import Union
 
 import jupytext  # type: ignore[import-untyped]
 from jupytext.config import (  # type: ignore[import-untyped]
@@ -61,7 +60,10 @@ class WaldiezExporter:
 
     def export(
         self,
-        path: Union[str, Path],
+        path: str | Path,
+        structured_io: bool = False,
+        uploads_root: Path | None = None,
+        skip_patch_io: bool = True,
         force: bool = False,
         debug: bool = False,
     ) -> None:
@@ -69,8 +71,15 @@ class WaldiezExporter:
 
         Parameters
         ----------
-        path : Union[str, Path]
+        path : str | Path
             The path to export to.
+        structured_io : bool, (optional)
+            Whether to use structured IO instead of the default 'input/print',
+            by default False.
+        uploads_root : str | Path | None, (optional)
+            The uploads root, to get user-uploaded files, by default None.
+        skip_patch_io : bool, (optional)
+            Whether to skip patching I/O, by default True.
         force : bool, (optional)
             Override the output file if it already exists, by default False.
         debug : bool, (optional)
@@ -99,23 +108,45 @@ class WaldiezExporter:
         if extension == ".waldiez":
             self.to_waldiez(path, debug=debug)
         elif extension == ".py":
-            self.to_py(path, debug=debug)
+            self.to_py(
+                path,
+                structured_io=structured_io,
+                uploads_root=uploads_root,
+                skip_patch_io=skip_patch_io,
+                debug=debug,
+            )
         elif extension == ".ipynb":
-            self.to_ipynb(path, debug=debug)
+            self.to_ipynb(
+                path,
+                structured_io=structured_io,
+                uploads_root=uploads_root,
+                skip_patch_io=skip_patch_io,
+                debug=debug,
+            )
         else:
             raise ValueError(f"Invalid extension: {extension}")
 
     def to_ipynb(
         self,
-        path: Path,
+        path: str | Path,
+        structured_io: bool = False,
+        uploads_root: Path | None = None,
+        skip_patch_io: bool = True,
         debug: bool = False,
     ) -> None:
         """Export flow to jupyter notebook.
 
         Parameters
         ----------
-        path : Path
+        path : str | Path
             The path to export to.
+        structured_io : bool, optional
+            Whether to use structured IO instead of the default 'input/print',
+            by default False.
+        uploads_root : Path | None, optional
+            The uploads root, to get user-uploaded files, by default None.
+        skip_patch_io : bool, optional
+            Whether to skip patching I/O, by default True.
         debug : bool, optional
             Whether to enable debug mode, by default False.
 
@@ -126,10 +157,15 @@ class WaldiezExporter:
         """
         # we first create a .py file with the content
         # and then convert it to a notebook using jupytext
+        if not isinstance(path, Path):
+            path = Path(path)
         exporter = create_flow_exporter(
             waldiez=self.waldiez,
             output_dir=path.parent,
+            uploads_root=uploads_root,
+            structured_io=structured_io,
             for_notebook=True,
+            skip_patch_io=skip_patch_io,
             debug=debug,
         )
         self.flow_extras = exporter.extras
@@ -161,13 +197,27 @@ class WaldiezExporter:
         Path(ipynb_path).rename(ipynb_path.replace(".tmp.ipynb", ".ipynb"))
         py_path.unlink(missing_ok=True)
 
-    def to_py(self, path: Path, debug: bool = False) -> None:
+    def to_py(
+        self,
+        path: str | Path,
+        structured_io: bool = False,
+        uploads_root: Path | None = None,
+        skip_patch_io: bool = True,
+        debug: bool = False,
+    ) -> None:
         """Export waldiez flow to a python script.
 
         Parameters
         ----------
-        path : Path
+        path : str | Path
             The path to export to.
+        structured_io : bool, optional
+            Whether to use structured IO instead of the default 'input/print',
+            by default False.
+        uploads_root : Path | None, optional
+            The uploads root, to get user-uploaded files, by default None.
+        skip_patch_io : bool, optional
+            Whether to skip patching I/O, by default True.
         debug : bool, optional
             Whether to enable debug mode, by default False.
 
@@ -176,10 +226,15 @@ class WaldiezExporter:
         RuntimeError
             If the python script could not be generated.
         """
+        if not isinstance(path, Path):
+            path = Path(path)
         exporter = create_flow_exporter(
             waldiez=self.waldiez,
             output_dir=path.parent,
             for_notebook=False,
+            uploads_root=uploads_root,
+            structured_io=structured_io,
+            skip_patch_io=skip_patch_io,
             debug=debug,
         )
         self.flow_extras = exporter.extras

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/exporting/core/context.py

@@ -54,6 +54,8 @@ class ExporterContext:
         is_async: bool = False,
         output_directory: Optional[str] = None,
         cache_seed: Optional[int] = None,
+        structured_io: bool = False,
+        skip_patch_io: bool = True,
     ) -> ExportConfig:
         """Get export config or return default.
 
@@ -75,6 +77,10 @@ class ExporterContext:
             The directory where the output will be saved, by default None
         cache_seed : Optional[int], optional
             The seed for caching, by default None
+        structured_io : bool, optional
+            Whether to use structured I/O, by default False
+        skip_patch_io : bool, optional
+            Whether to skip patching I/O, by default True
 
         Returns
         -------
@@ -85,6 +91,12 @@ class ExporterContext:
             "requirements": requirements or [],
             "tags": tags or [],
             "is_async": self.config.is_async if self.config else is_async,
+            "structured_io": (
+                self.config.structured_io if self.config else structured_io
+            ),
+            "skip_patch_io": (
+                self.config.skip_patch_io if self.config else skip_patch_io
+            ),
         }
         if output_extension is not None:
             kwargs["output_extension"] = output_extension

waldiez/exporting/core/extras/flow_extras.py

@@ -7,19 +7,7 @@
 from dataclasses import dataclass, field
 from typing import Optional
 
-# pylint: disable=import-error,line-too-long
-# pyright: reportMissingImports=false
-try:
-    from waldiez._version import __version__ as waldiez_version  # type: ignore[unused-ignore, import-not-found, import-untyped]  # noqa
-except ImportError:  # pragma: no cover
-    import warnings
-
-    # .gitingored (generated by hatch)
-    warnings.warn(
-        "Importing __version__ failed. Using 'dev' as version.",
-        stacklevel=2,
-    )
-    waldiez_version = "dev"
+from waldiez.utils import get_waldiez_version
 
 from ..result import ExportResult
 from ..types import ExportConfig
@@ -34,7 +22,7 @@ class FlowExtras(BaseExtras):
     flow_name: str = ""
     description: str = ""
     config: ExportConfig = field(default_factory=ExportConfig)
-    version: str = waldiez_version  # pyright: ignore
+    version: str = field(default_factory=get_waldiez_version)
 
     # Sub-exporter results
     tools_result: Optional[ExportResult] = None

waldiez/exporting/core/types.py

@@ -239,10 +239,28 @@ class ExportConfig:
 
     Attributes
     ----------
+    name : str
+        The name of the export.
+    description : str
+        A brief description of the export.
+    requirements : list[str]
+        A list of requirements for the export.
+    tags : list[str]
+        A list of tags associated with the export.
+    output_directory : Optional[str | Path]
+        The directory where the exported content will be saved.
+    uploads_root : Optional[str | Path]
+        The root directory for uploads, if applicable.
+    cache_seed : Optional[int]
+        The seed for caching, if applicable.
+    structured_io : bool
+        Whether the export should use structured I/O.
     output_extension : str
         The file extension for the exported content.
     is_async : bool
         Whether the exported content should be asynchronous.
+    skip_patch_io : bool
+        Whether to skip patching I/O operations.
     """
 
     name: str = "Waldiez Flow"
@@ -254,7 +272,10 @@ class ExportConfig:
     output_extension: str = "py"
     is_async: bool = False
     output_directory: Optional[str | Path] = None
+    uploads_root: Optional[Path] = None
     cache_seed: Optional[int] = None
+    structured_io: bool = False
+    skip_patch_io: bool = True
 
     @property
     def for_notebook(self) -> bool:

waldiez/exporting/flow/exporter.py

@@ -21,6 +21,8 @@ class FlowExporter(Exporter[FlowExtras]):
         waldiez: Waldiez,
         output_dir: Path | None,
         for_notebook: bool,
+        structured_io: bool = False,
+        skip_patch_io: bool = True,
         context: Optional[ExporterContext] = None,
         **kwargs: Any,
     ) -> None:
@@ -52,6 +54,8 @@ class FlowExporter(Exporter[FlowExtras]):
             is_async=waldiez.is_async,
             output_directory=str(self.output_dir) if self.output_dir else None,
             cache_seed=waldiez.cache_seed,
+            structured_io=structured_io,
+            skip_patch_io=skip_patch_io,
         )
         self._extras = self._generate_extras()
 

waldiez/exporting/flow/factory.py

@@ -20,6 +20,7 @@ from .exporter import FlowExporter
 def create_flow_exporter(
     waldiez: Waldiez,
     output_dir: Path | None,
+    uploads_root: Path | None,
     for_notebook: bool,
     context: Optional[ExporterContext] = None,
     **kwargs: Any,
@@ -32,6 +33,8 @@ def create_flow_exporter(
         The Waldiez instance containing the flow data.
     output_dir : Path
         The directory where the exported flow will be saved.
+    uploads_root : Path
+        The root directory for uploads, if applicable.
     for_notebook : bool
         Whether the export is intended for a notebook environment.
     context : Optional[ExporterContext], optional
@@ -44,6 +47,8 @@ def create_flow_exporter(
     ChatsExporter
         The created chats exporter.
     """
+    structured_io = kwargs.pop("structured_io", False)
+    skip_patch_io = kwargs.pop("skip_patch_io", True)
     if context is None:
         config = ExportConfig(
             name=waldiez.name,
@@ -53,7 +58,10 @@ def create_flow_exporter(
             output_extension="ipynb" if for_notebook else "py",
             is_async=waldiez.is_async,
             output_directory=output_dir,
+            uploads_root=uploads_root,
             cache_seed=waldiez.cache_seed,
+            structured_io=structured_io,
+            skip_patch_io=skip_patch_io,
         )
         context = ExporterContext(
             config=config,
@@ -71,7 +79,10 @@ def create_flow_exporter(
                 output_extension="ipynb" if for_notebook else "py",
                 is_async=waldiez.is_async,
                 output_directory=output_dir,
+                uploads_root=uploads_root,
                 cache_seed=waldiez.cache_seed,
+                structured_io=structured_io,
+                skip_patch_io=skip_patch_io,
             )
         else:
             context.config.update(
@@ -82,13 +93,18 @@ def create_flow_exporter(
                 output_extension="ipynb" if for_notebook else "py",
                 is_async=waldiez.is_async,
                 output_directory=output_dir,
+                uploads_root=uploads_root,
                 cache_seed=waldiez.cache_seed,
+                structured_io=structured_io,
+                skip_patch_io=skip_patch_io,
             )
 
     return FlowExporter(
         waldiez=waldiez,
         output_dir=output_dir,
+        uploads_root=uploads_root,
         for_notebook=for_notebook,
+        structured_io=structured_io,
         context=context,
         **kwargs,
     )