waldiez 0.4.8__tar.gz → 0.4.11__tar.gz

{waldiez-0.4.8 → waldiez-0.4.11}/.gitignore

@@ -29,6 +29,7 @@ dmypy.json
 .coverage
 **/*.coverage*
 **/*_db
+**/*.db
 **/*cache
 **/*.log
 **/*.pid

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waldiez
-Version: 0.4.8
+Version: 0.4.11
 Dynamic: Keywords
 Summary: Make AG2 Agents Collaborate: Drag, Drop, and Orchestrate with Waldiez.
 Project-URL: Homepage, https://waldiez.io
@@ -33,7 +33,7 @@ Requires-Dist: aiofiles==24.1.0
 Requires-Dist: aiosqlite==0.21.0
 Requires-Dist: asyncer==0.0.8
 Requires-Dist: click<8.2
-Requires-Dist: graphviz<=0.20.3
+Requires-Dist: graphviz<=0.21
 Requires-Dist: httpx<1
 Requires-Dist: jupytext
 Requires-Dist: nest-asyncio==1.6.0
@@ -80,7 +80,7 @@ Requires-Dist: ipython; extra == 'ag2-extras'
 Requires-Dist: langchain-community<1,>=0.3.12; extra == 'ag2-extras'
 Requires-Dist: litellm; extra == 'ag2-extras'
 Requires-Dist: markdownify; extra == 'ag2-extras'
-Requires-Dist: mcp<1.6,>=1.4.0; extra == 'ag2-extras'
+Requires-Dist: mcp<2,>=1.4.0; extra == 'ag2-extras'
 Requires-Dist: mistralai>=1.8.1; extra == 'ag2-extras'
 Requires-Dist: networkx; (python_version < '3.11') and extra == 'ag2-extras'
 Requires-Dist: networkx>=3.5; (python_version >= '3.11') and extra == 'ag2-extras'
@@ -106,10 +106,10 @@ Provides-Extra: dev
 Requires-Dist: ag2[redis]==0.9.2; extra == 'dev'
 Requires-Dist: ag2[websockets]==0.9.2; extra == 'dev'
 Requires-Dist: autoflake==2.3.1; extra == 'dev'
-Requires-Dist: bandit==1.8.3; extra == 'dev'
+Requires-Dist: bandit==1.8.5; extra == 'dev'
 Requires-Dist: black[jupyter]==25.1.0; extra == 'dev'
 Requires-Dist: build==1.2.2.post1; extra == 'dev'
-Requires-Dist: fakeredis>=2.28.1; extra == 'dev'
+Requires-Dist: fakeredis<=3.30.0,>=2.28.1; extra == 'dev'
 Requires-Dist: fastjsonschema>=2.21; extra == 'dev'
 Requires-Dist: flake8==7.2.0; extra == 'dev'
 Requires-Dist: hatchling==1.27.0; extra == 'dev'
@@ -117,7 +117,7 @@ Requires-Dist: jsonschema==4.24.0; extra == 'dev'
 Requires-Dist: jupyter-server==2.16.0; extra == 'dev'
 Requires-Dist: jupyterlab<5.0,>=4.4.0; extra == 'dev'
 Requires-Dist: mypy-extensions>=1.1.0; extra == 'dev'
-Requires-Dist: mypy==1.16.0; extra == 'dev'
+Requires-Dist: mypy==1.16.1; extra == 'dev'
 Requires-Dist: nbclient>=0.10.2; extra == 'dev'
 Requires-Dist: nbconvert>=7.16.6; extra == 'dev'
 Requires-Dist: nbformat>=5.10.4; extra == 'dev'
@@ -134,7 +134,7 @@ Requires-Dist: toml==0.10.2; (python_version <= '3.10') and extra == 'dev'
 Requires-Dist: types-jsonschema==4.24.0.20250528; extra == 'dev'
 Requires-Dist: types-pyyaml==6.0.12.20250516; extra == 'dev'
 Requires-Dist: types-redis==4.6.0.20241004; extra == 'dev'
-Requires-Dist: types-requests==2.32.0.20250602; extra == 'dev'
+Requires-Dist: types-requests==2.32.4.20250611; extra == 'dev'
 Requires-Dist: types-toml==0.10.8.20240310; extra == 'dev'
 Requires-Dist: watchdog==6.0.0; extra == 'dev'
 Requires-Dist: yamllint==1.37.1; extra == 'dev'
@@ -157,19 +157,19 @@ 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.11; 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.11; (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.11; extra == 'studio'
 Provides-Extra: test
 Requires-Dist: ag2[redis]==0.9.2; extra == 'test'
 Requires-Dist: ag2[websockets]==0.9.2; extra == 'test'
-Requires-Dist: fakeredis>=2.28.1; extra == 'test'
+Requires-Dist: fakeredis<=3.30.0,>=2.28.1; extra == 'test'
 Requires-Dist: paho-mqtt<3.0,>=2.1.0; extra == 'test'
 Requires-Dist: pytest-asyncio==1.0.0; extra == 'test'
 Requires-Dist: pytest-cov==6.1.1; extra == 'test'

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

@@ -34,7 +34,7 @@ dependencies = [
   "asyncer==0.0.8",
   "numpy<=2.3.0",
   "nest_asyncio==1.6.0",
-  "graphviz<=0.20.3",
+  "graphviz<=0.21",
   "pandas>=2",
   "pillow",
   "pip>=25",
@@ -81,15 +81,15 @@ mqtt = [
 ]
 # jupyterlab extension
 jupyter = [
-  "waldiez_jupyter==0.4.8",
+  "waldiez_jupyter==0.4.11",
   "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.11; python_version >= '3.11'",
 ]
 studio = [
-  "waldiez_studio==0.4.8",
+  "waldiez_studio==0.4.11",
 ]
 ag2_extras = [
   # models
@@ -154,16 +154,16 @@ ag2_extras = [
   "google-api-python-client>=2.163.0,<3.0",
   "google-auth-httplib2>=0.2.0,<0.3",
   "google-auth-oauthlib>=1.2.1,<2.0",
-  "mcp>=1.4.0,<1.6",
+  "mcp>=1.4.0,<2",
 ]
 dev = [
   "ag2[redis]==0.9.2",
   "ag2[websockets]==0.9.2",
   "autoflake==2.3.1",
-  "bandit==1.8.3",
+  "bandit==1.8.5",
   "black[jupyter]==25.1.0",
   "build==1.2.2.post1",
-  "fakeredis>=2.28.1",
+  "fakeredis>=2.28.1,<=3.30.0",
   "flake8==7.2.0",
   "jsonschema==4.24.0",
   "jupyterlab>=4.4.0,<5.0",
@@ -174,7 +174,7 @@ dev = [
   "nbformat>=5.10.4",
   "nbconvert>=7.16.6",
   "hatchling==1.27.0",
-  "mypy==1.16.0",
+  "mypy==1.16.1",
   "mypy-extensions>=1.1.0",
   "paho-mqtt>=2.1.0,<3.0",
   "pandas-stubs==2.2.3.250527",
@@ -188,7 +188,7 @@ dev = [
   "types-jsonschema==4.24.0.20250528",
   "types-redis==4.6.0.20241004",
   "types-toml==0.10.8.20240310",
-  "types-requests==2.32.0.20250602",
+  "types-requests==2.32.4.20250611",
   "toml==0.10.2; python_version <= '3.10'",
   "watchdog==6.0.0",
   "yamllint==1.37.1",
@@ -213,9 +213,8 @@ docs = [
 test = [
   "ag2[redis]==0.9.2",
   "ag2[websockets]==0.9.2",
-  "fakeredis>=2.28.1",
+  "fakeredis>=2.28.1,<=3.30.0",
   "paho-mqtt>=2.1.0,<3.0",
-
   "pytest==8.4.0",
   "pytest-asyncio==1.0.0",
   "pytest-cov==6.1.1",
@@ -284,6 +283,10 @@ test-io = """
 hatch run pre-test
 pytest -c pyproject.toml --cov=waldiez/io --cov-branch --cov-report=term-missing:skip-covered -vv tests/io
 """
+test-running = """
+hatch run pre-test
+pytest -c pyproject.toml --cov=waldiez/running --cov-branch --cov-report=term-missing:skip-covered -vv tests/running
+"""
 
 format-black = "black --config pyproject.toml waldiez tests scripts"
 format-autoflake = "autoflake --in-place --remove-all-unused-imports --remove-unused-variables --recursive waldiez tests scripts"

{waldiez-0.4.8 → waldiez-0.4.11}/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-0.4.8 → waldiez-0.4.11}/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-0.4.8 → waldiez-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-0.4.8 → waldiez-0.4.11}/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-0.4.8 → waldiez-0.4.11}/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