waldiez 0.5.6__py3-none-any.whl → 0.5.8__py3-none-any.whl

waldiez/_version.py

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

waldiez/cli.py

@@ -100,6 +100,19 @@ def run(
         False,
         help="Override the output file if it already exists.",
     ),
+    env_file: Optional[Path] = typer.Option(  # noqa: B008
+        None,
+        "--env-file",
+        "-e",
+        help=(
+            "Path to a .env file containing additional environment variables. "
+            "These variables will be set before running the flow."
+        ),
+        file_okay=True,
+        dir_okay=False,
+        readable=True,
+        resolve_path=True,
+    ),
 ) -> None:
     """Run a Waldiez flow."""
     os.environ["AUTOGEN_USE_DOCKER"] = "0"
@@ -126,6 +139,7 @@ def run(
             structured,  # structured_io
             False,  # skip_mmd
             False,  # skip_timeline
+            env_file,
         )
     else:
         runner.run(
@@ -134,6 +148,7 @@ def run(
             structured_io=structured,
             skip_mmd=False,
             skip_timeline=False,
+            dot_env=env_file,
         )
 
 
@@ -157,7 +172,7 @@ def convert(
             "Path to the output file. "
             "The file extension determines the output format: "
             "`.py` for Python script, `.ipynb` for Jupyter notebook."
-            " If not provided, the output will be saved in the same directory as the input file."
+            " If not provided, the output (.py) will be saved in the same directory as the input file."
             " If the file already exists, it will not be overwritten unless --force is used."
         ),
         file_okay=True,
@@ -230,7 +245,7 @@ def _get_output_path(output: Optional[Path], force: bool) -> Optional[Path]:
     if output is not None and not output.parent.exists():
         output.parent.mkdir(parents=True)
     if output is not None and output.exists():
-        if force is False:
+        if not force:
             LOG.error("Output file already exists.")
             raise typer.Exit(code=1)
         output.unlink()

waldiez/exporter.py

@@ -97,7 +97,7 @@ class WaldiezExporter:
         if path.is_dir():
             raise IsADirectoryError(f"Output is a directory: {path}")
         if path.exists():
-            if force is False:
+            if not force:
                 raise FileExistsError(f"File already exists: {path}")
             path.unlink(missing_ok=True)
         path.parent.mkdir(parents=True, exist_ok=True)

waldiez/exporting/agent/code_execution.py

@@ -5,6 +5,8 @@
 
 import json
 
+from typing_extensions import Literal
+
 from waldiez.models import WaldiezAgent, WaldiezAgentCodeExecutionConfig
 
 from ..core import CodeExecutionConfig, ImportPosition, ImportStatement
@@ -57,6 +59,7 @@ class CodeExecutionProcessor:
             ),
         )
 
+    # noinspection PyMethodMayBeStatic
     def _get_executor_class_name(self, use_docker: bool) -> str:
         """Get the appropriate executor class name."""
         return (
@@ -66,9 +69,13 @@ class CodeExecutionProcessor:
         )
 
     def _build_executor_content(
-        self, config: WaldiezAgentCodeExecutionConfig, use_docker: bool
+        self,
+        config: WaldiezAgentCodeExecutionConfig | Literal[False],
+        use_docker: bool,
     ) -> str:
         """Build the executor content string."""
+        if config is False:
+            return ""
         executor_class = self._get_executor_class_name(use_docker)
         lines = [f"{self.agent_name}_executor = {executor_class}("]
 

waldiez/exporting/agent/exporter.py

@@ -174,7 +174,7 @@ class AgentExporter(Exporter[StandardExtras]):
             model_names=self.model_names,
             initial_chats=self.initial_chats,
             group_chat_members=self.group_chat_members,
-            serializer=(self.context.get_serializer().serialize),
+            serializer=self.context.get_serializer().serialize,
         )
 
         group_manager_extras = group_manager_processor.process(
@@ -376,6 +376,7 @@ class AgentExporter(Exporter[StandardExtras]):
 
 
 # pylint: disable=unused-argument
+# noinspection PyUnusedLocal
 def fallback_args_resolver(agent: WaldiezAgent) -> list[str]:
     """Fallback resolver for agent arguments.
 

waldiez/exporting/agent/extras/captain_agent_extras.py

@@ -248,6 +248,7 @@ class CaptainAgentProcessor:
                 return model
         return self._get_default_model(model_id)
 
+    # noinspection PyMethodMayBeStatic
     def _get_default_model(self, model_id: str) -> WaldiezModel:
         """Get the default model.
 

waldiez/exporting/agent/extras/doc_agent_extras.py

@@ -292,7 +292,6 @@ class DocAgentProcessor:
         if not isinstance(self.agent, WaldiezDocAgent):  # pragma: no cover
             raise TypeError("Agent must be a WaldiezDocAgent instance.")
         query_engine = self.agent.get_query_engine()
-        db_path_str = ""
         if query_engine.get_db_path():
             db_path_str = query_engine.get_db_path()
         else:
@@ -302,7 +301,7 @@ class DocAgentProcessor:
                 db_path_str = str(chroma_dir)
             else:
                 db_path_str = query_engine.get_db_path()
-        if self.agent.data.reset_collection is True:
+        if self.agent.data.reset_collection:
             shutil.rmtree(db_path_str, ignore_errors=True)
             Path(db_path_str).mkdir(parents=True, exist_ok=True)
         return self.path_resolver.resolve(db_path_str)
@@ -333,7 +332,7 @@ class DocAgentProcessor:
             parsed_docs_path = str(Path.cwd() / parsed_docs_path)
         if not Path(parsed_docs_path).is_dir():
             Path(parsed_docs_path).mkdir(parents=True, exist_ok=True)
-        if self.agent.data.reset_collection is True:
+        if self.agent.data.reset_collection:
             shutil.rmtree(parsed_docs_path, ignore_errors=True)
             Path(parsed_docs_path).mkdir(parents=True, exist_ok=True)
         return self.path_resolver.resolve(str(parsed_docs_path))

waldiez/exporting/agent/extras/group_manager_agent_extas.py

@@ -433,6 +433,7 @@ class GroupManagerProcessor:
 
         return lines
 
+    # noinspection PyMethodMayBeStatic
     def _should_check_for_after_work(self, pattern_class: str) -> bool:
         """Check if pattern should have after work configuration."""
         return pattern_class not in ["ManualPattern", "AutoPattern"]

waldiez/exporting/agent/extras/handoffs/after_work.py

@@ -17,10 +17,10 @@ class AfterWorkResult:
 
     Attributes
     ----------
-    registration : str
-        Code to register the after-work transition.
-    after_agent : str
-        Code to be placed after the agent definition.
+    content : str
+        The registration string for the after-work transition.
+    before_content : str
+        Any additional code that should be placed before the main content,
     extra_imports : set[str]
         Additional imports required for the after-work transition.
     """

waldiez/exporting/agent/extras/handoffs/target.py

@@ -87,6 +87,7 @@ class TransitionTargetProcessor:
         what = target.target_type
         result.extra_imports.add(f"from {where} import {what}")
 
+        # noinspection DuplicatedCode
         processors: dict[str, Callable[[WaldiezTransitionTarget], str]] = {
             "AgentTarget": self._process_agent_target,
             "RandomAgentTarget": self._process_random_agent_target,
@@ -152,6 +153,7 @@ class TransitionTargetProcessor:
         target_arg = f'nested_chat_config={{"chat_queue": {chat_name}}}'
         return f"NestedChatTarget({target_arg})"
 
+    # noinspection DuplicatedCode
     def _process_nested_chat_target_full(
         self, _: WaldiezTransitionTarget
     ) -> TargetResult:
@@ -178,6 +180,7 @@ class TransitionTargetProcessor:
         return result
 
     # pylint: disable=no-self-use
+    # noinspection PyMethodMayBeStatic
     def _process_simple_target(self, target: WaldiezTransitionTarget) -> str:
         """Process simple targets that don't need parameters.
 

waldiez/exporting/agent/extras/rag/chroma_extras.py

@@ -15,15 +15,12 @@ def _get_chroma_client_string(agent: WaldiezRagUserProxy) -> tuple[str, str]:
     ----------
     agent : WaldiezRagUserProxy
         The agent.
-    agent_name : str
-        The agent's name.
 
     Returns
     -------
     tuple[str, str]
         The 'client' and what to import.
     """
-    # TODO: also check `connection_url` (chroma in client-server mode)
     to_import = "import chromadb"
     client_str = "chromadb."
     if (
@@ -63,7 +60,6 @@ def _get_chroma_embedding_function_string(
         The 'embedding_function', the import and the custom embedding function.
     """
     to_import = ""
-    embedding_function_arg = ""
     embedding_function_content = ""
     if not agent.retrieve_config.use_custom_embedding:
         to_import = (

waldiez/exporting/agent/extras/rag/mongo_extras.py

@@ -25,7 +25,6 @@ def _get_mongodb_embedding_function_string(
         The 'embedding_function', the import and the custom_embedding_function.
     """
     to_import = ""
-    embedding_function_arg = ""
     embedding_function_content = ""
     if not agent.retrieve_config.use_custom_embedding:
         to_import = "from sentence_transformers import SentenceTransformer"

waldiez/exporting/agent/extras/rag/pgvector_extras.py

@@ -45,10 +45,8 @@ def _get_pgvector_embedding_function_string(
         The 'embedding_function', the import and the custom_embedding_function.
     """
     to_import = ""
-    embedding_function_arg = ""
     embedding_function_content = ""
     if agent.retrieve_config.use_custom_embedding:
-        embedding_function_arg = f"custom_embedding_function_{agent_name}"
         embedding_function_content, embedding_function_arg = (
             agent.retrieve_config.get_custom_embedding_function(
                 name_suffix=agent_name

waldiez/exporting/agent/extras/rag/qdrant_extras.py

@@ -15,8 +15,6 @@ def _get_qdrant_client_string(agent: WaldiezRagUserProxy) -> tuple[str, str]:
     ----------
     agent : WaldiezRagUserProxy
         The agent.
-    agent_name : str
-        The agent's name.
 
     Returns
     -------
@@ -62,7 +60,6 @@ def _get_qdrant_embedding_function_string(
         and the custom_embedding_function if used.
     """
     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:

waldiez/exporting/agent/extras/rag/vector_db_extras.py

@@ -48,8 +48,9 @@ def get_vector_db_extras(
 
     Returns
     -------
-    tuple[str, str, Set[str]]
-        The content before the arg if any, the arg and the related imports.
+    VectorDBExtras
+        The vector db extras containing the before content, vector db arg,
+        and imports.
     """
     before = ""
     imports: set[str] = set()

waldiez/exporting/agent/factory.py

@@ -5,7 +5,7 @@
 """Factory for creating agent exporter."""
 
 from pathlib import Path
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable
 
 from waldiez.models import (
     WaldiezAgent,
@@ -30,11 +30,11 @@ def create_agent_exporter(
     initial_chats: list[WaldiezAgentConnection],
     is_async: bool = False,
     for_notebook: bool = False,
-    cache_seed: Optional[int] = None,
-    group_chat_members: Optional[list[WaldiezAgent]] = None,
-    arguments_resolver: Optional[Callable[[WaldiezAgent], list[str]]] = None,
-    output_dir: Optional[Union[str, Path]] = None,
-    context: Optional[ExporterContext] = None,
+    cache_seed: int | None = None,
+    group_chat_members: list[WaldiezAgent] | None = None,
+    arguments_resolver: Callable[[WaldiezAgent], list[str]] | None = None,
+    output_dir: str | Path | None = None,
+    context: ExporterContext | None = None,
     **kwargs: Any,
 ) -> AgentExporter:
     """Create an agent exporter.
@@ -57,15 +57,15 @@ def create_agent_exporter(
         Whether exporting for notebook, by default False
     cache_seed : Optional[int], optional
         Cache seed if any, by default None
-    initial_chats : Optional[list[WaldiezAgentConnection]], optional
+    initial_chats : list[WaldiezAgentConnection]
         Initial chats for group managers, by default None
-    group_chat_members : Optional[list[WaldiezAgent]], optional
+    group_chat_members : list[WaldiezAgent], optional
         Group chat members if group manager, by default None
-    arguments_resolver : Optional[Callable], optional
+    arguments_resolver : Callable, optional
         Function to resolve additional arguments, by default None
-    output_dir : Optional[Union[str, Path]], optional
+    output_dir : str | Path | None, optional
         Output directory for generated files, by default None
-    context : Optional[ExporterContext], optional
+    context : ExporterContext | None, optional
         Exporter context with dependencies, by default None
     **kwargs : Any
         Additional keyword arguments.

waldiez/exporting/agent/processor.py

@@ -53,8 +53,8 @@ class AgentProcessor:
 
         Returns
         -------
-        str
-            The generated agent definition string.
+        AgentProcessingResult
+            The result containing the agent definition string.
         """
         if (
             isinstance(self.extras, GroupManagerExtras)
@@ -226,6 +226,7 @@ class AgentProcessor:
             )
         return f"    default_auto_reply={default_auto_reply},\n"
 
+    # noinspection PyMethodMayBeStatic
     def add_extra_args(self, agent_str: str, extra_args: str) -> str:
         """Add extra agent args.
 

waldiez/exporting/chats/exporter.py

@@ -43,9 +43,9 @@ class ChatsExporter(Exporter[ChatExtras]):
             Mapping of chat IDs to their names.
         main_chats : list[WaldiezAgentConnection]
             Main chats that are connections between agents.
-        root_group_manager : Optional[WaldiezGroupManager]
+        root_group_manager : WaldiezGroupManager | None
             The root group manager for managing chat groups, if any.
-        context : Optional[ExporterContext], optional
+        context : ExporterContext | None, optional
             Exporter context with dependencies, by default None
         **kwargs : Any
             Additional keyword arguments for the exporter.

waldiez/exporting/chats/factory.py

@@ -2,7 +2,7 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Factory function for creating a ChatsExporter instance."""
 
-from typing import Any, Optional
+from typing import Any
 
 from waldiez.models import (
     WaldiezAgent,
@@ -21,8 +21,8 @@ def create_chats_exporter(
     all_chats: list[WaldiezChat],
     chat_names: dict[str, str],
     main_chats: list[WaldiezAgentConnection],
-    root_group_manager: Optional[WaldiezGroupManager] = None,
-    context: Optional[ExporterContext] = None,
+    root_group_manager: WaldiezGroupManager | None = None,
+    context: ExporterContext | None = None,
     **kwargs: Any,
 ) -> ChatsExporter:
     """Create a chats exporter.
@@ -39,9 +39,9 @@ def create_chats_exporter(
         Mapping of chat IDs to their names.
     main_chats : list[WaldiezAgentConnection]
         Main chats that are connections between agents.
-    root_group_manager : Optional[WaldiezGroupManager], optional
+    root_group_manager : WaldiezGroupManager | None, optional
         The root group manager for managing chat groups, if any.
-    context : Optional[ExporterContext], optional
+    context : ExporterContext | None, optional
         Exporter context with dependencies, by default None
     **kwargs : Any
         Additional keyword arguments for the exporter.

waldiez/exporting/chats/processor.py

@@ -39,6 +39,7 @@ class ChatParams:
 
 
 # pylint: disable=too-many-arguments,too-many-positional-arguments
+# noinspection PyTypeChecker
 class ChatsProcessor:
     """Processor for chats export."""
 
@@ -49,7 +50,7 @@ class ChatsProcessor:
         all_chats: list[WaldiezChat],
         chat_names: dict[str, str],
         main_chats: list[WaldiezAgentConnection],
-        root_group_manager: Optional[WaldiezGroupManager],
+        root_group_manager: WaldiezGroupManager | None,
         for_notebook: bool,
         is_async: bool,
         cache_seed: Optional[int],
@@ -60,6 +61,26 @@ class ChatsProcessor:
 
         Parameters
         ----------
+        all_agents : list[WaldiezAgent]
+            All agents involved in the chats.
+        agent_names : dict[str, str]
+            Mapping of agent IDs to their names.
+        all_chats : list[WaldiezChat]
+            All chats to be exported.
+        chat_names : dict[str, str]
+            Mapping of chat IDs to their names.
+        main_chats : list[WaldiezAgentConnection]
+            Main chats that are connections between agents.
+        root_group_manager : WaldiezGroupManager | None
+            The root group manager for managing chat groups, if any.
+        for_notebook : bool
+            Whether the export is for a notebook.
+        is_async : bool
+            Whether the chat is asynchronous.
+        cache_seed : int | None
+            The cache seed for the export, if any.
+        serializer : Serializer
+            The serializer to use for escaping quotes in strings.
         extras : ChatExtras
             The structured extras for the chats export.
         """

waldiez/exporting/chats/utils/common.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=line-too-long
+# flake8: noqa: E501
 """Common utilities for exporting chats."""
 
 import json
@@ -99,14 +101,56 @@ def get_event_handler_string(
         f"{tab}                should_continue = False\n"
         f"{tab}            if not should_continue:\n"
         f"{tab}                break\n"
+    )
+    content += get_result_dicts_string(tab, is_async)
+    content += (
         f"{tab}else:\n"
         f"{tab}    if not isinstance(results, list):\n"
         f"{tab}        results = [results]\n"
-        f"{tab}    for result in results:\n"
+        f"{tab}    for index, result in enumerate(results):\n"
     )
     if is_async:
         content += f"{tab}        await result.process()\n"
     else:
         content += f"{tab}        result.process()\n"
+    content += get_result_dicts_string(tab, is_async)
 
     return content
+
+
+def get_result_dicts_string(tab: str, is_async: bool) -> str:
+    """Get the result dicts string.
+
+    Parameters
+    ----------
+    tab : str
+        The space string to use for indentation.
+    is_async : bool
+        Whether the function is asynchronous.
+
+    Returns
+    -------
+    str
+        The result dicts string.
+    """
+    space = f"{tab}        "
+    flow_content = f"{space}result_dict = {{\n"
+    flow_content += f"{space}    'index': index,\n"
+    if is_async:
+        flow_content += f"{space}    'messages': await result.messages,\n"
+        flow_content += f"{space}    'summary': await result.summary,\n"
+        flow_content += f"{space}    'cost': (await result.cost).model_dump(mode='json', fallback=str) if await result.cost else None,\n"
+        flow_content += f"{space}    'context_variables': (await result.context_variables).model_dump(mode='json', fallback=str) if await result.context_variables else None,\n"
+        flow_content += (
+            f"{space}    'last_speaker': await result.last_speaker,\n"
+        )
+    else:
+        flow_content += f"{space}    'messages': result.messages,\n"
+        flow_content += f"{space}    'summary': result.summary,\n"
+        flow_content += f"{space}    'cost': result.cost.model_dump(mode='json', fallback=str) if result.cost else None,\n"
+        flow_content += f"{space}    'context_variables': result.context_variables.model_dump(mode='json', fallback=str) if result.context_variables else None,\n"
+        flow_content += f"{space}    'last_speaker': result.last_speaker,\n"
+    flow_content += f"{space}    'uuid': str(result.uuid),\n"
+    flow_content += f"{space}}}\n"
+    flow_content += f"{space}result_dicts.append(result_dict)\n"
+    return flow_content

waldiez/exporting/chats/utils/group.py

@@ -43,7 +43,7 @@ def export_group_chats(
     tab = "    " * tabs
     run_group_chat = "run_group_chat"
     if is_async:
-        run_group_chat = "a_run_group_chat"
+        run_group_chat = "await a_run_group_chat"
     manager_name = agent_names[manager.id]
     pattern_name = f"{manager_name}_pattern"
     content = f"{tab}results = {run_group_chat}(" + "\n"

waldiez/exporting/chats/utils/sequential.py

@@ -22,7 +22,7 @@ def export_sequential_chat(
     is_async: bool,
     skip_cache: bool,
 ) -> tuple[str, str]:
-    r"""Get the chats content, when there are more than one chats in the flow.
+    """Get the chats content, when there are more than one chats in the flow.
 
     Parameters
     ----------
@@ -45,72 +45,6 @@ def export_sequential_chat(
     -------
     tuple[str, str]
         The main chats content and additional methods string if any.
-
-    Example
-    -------
-    ```python
-    >>> from waldiez.models import (
-    ...     WaldiezAgent,
-    ...     WaldiezChat,
-    ...     WaldiezChatData,
-    ...     WaldiezChatMessage,
-    ... )
-    >>> chat1 = WaldiezChat(
-    ...     id="wc-1",
-    ...     name="chat1",
-    ...     description="A chat between two agents.",
-    ...     tags=["chat", "chat1"],
-    ...     requirements=[],
-    ...     data=WaldiezChatData(
-    ...         sender="wa-1",
-    ...         recipient="wa-2",
-    ...         position=0,
-    ...         message=WaldiezChatMessage(
-    ...             type="string",
-    ...             content="Hello, how are you?",
-    ...         ),
-    ...     ),
-    ... )
-    >>> chat2 = WaldiezChat(
-    ...     id="wc-2",
-    ...     name="chat2",
-    ...     description="A chat between two agents.",
-    ...     tags=["chat", "chat2"],
-    ...     requirements=[],
-    ...     data=WaldiezChatData(
-    ...         sender="wa-2",
-    ...         recipient="wa-1",
-    ...         position=1,
-    ...         message=WaldiezChatMessage(
-    ...             type="string",
-    ...             content="I am good, thank you. How about you?",
-    ...         ),
-    ...     ),
-    ... )
-    >>> agent_names = {"wa-1": "agent1", "wa-2": "agent2"}
-    >>> chat_names = {"wc-1": "chat1", "wc-2": "chat2"}
-    >>> serializer = lambda x: x.replace('"', "\"").replace("\n", "\\n")
-    >>>  export_sequential_chat(
-    ...     main_chats=[(chat1, agent1, agent2), (chat2, agent2, agent1)],
-    ...     chat_names=chat_names,
-    ...     agent_names=agent_names,
-    ...     serializer=serializer,
-    ...     tabs=0,
-    ...     is_async=False,
-    ... )
-    results = initiate_chats([
-        {
-            "sender": agent1,
-            "recipient": agent2,
-            "message": "Hello, how are you?",
-        },
-        {
-            "sender": agent2,
-            "recipient": agent1,
-            "message": "I am good, thank you. How about you?",
-        },
-    ])
-    ```
     """
     tab = "    " * tabs if tabs > 0 else ""
     content = "\n"
@@ -138,6 +72,7 @@ def export_sequential_chat(
     return content, additional_methods_string
 
 
+# noinspection PyTypeChecker
 def _get_chat_dict_string(
     connection: WaldiezAgentConnection,
     is_first: bool,
@@ -230,7 +165,7 @@ def _get_chat_string_start(
     if not is_first:
         chat_string += "\n" + f'{tab}    "sender": {agent_names[sender.id]},'
     chat_string += "\n" + f'{tab}    "recipient": {agent_names[recipient.id]},'
-    if skip_cache is False:
+    if not skip_cache:
         chat_string += "\n" + f'{tab}    "cache": cache,'
     # additional_methods_string = ""
     for key, value in chat_args.items():