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

waldiez/models/waldiez.py

@@ -1,23 +1,33 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=too-many-public-methods
 """Waldiez data class.
 
-A Waldiez class contains all the information that is needed to generate
+A Waldiez class contains all the information needed to generate
 and run an autogen workflow. It has the model/LLM configurations, the agent
-definitions and their optional additional skills to be used.
+definitions and their optional additional tools to be used.
 """
 
 import json
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
-
-from .agents import WaldiezAgent, get_retrievechat_extra_requirements
-from .chat import WaldiezChat
+from typing import Any, Iterator, Optional
+
+from .agents import (
+    WaldiezAgent,
+    WaldiezGroupManager,
+    get_captain_agent_extra_requirements,
+    get_retrievechat_extra_requirements,
+)
 from .common import get_autogen_version
-from .flow import WaldiezFlow, get_flow_data
+from .flow import (
+    WaldiezAgentConnection,
+    WaldiezFlow,
+    WaldiezFlowInfo,
+    get_flow_data,
+)
 from .model import WaldiezModel, get_models_extra_requirements
-from .skill import WaldiezSkill, get_skills_extra_requirements
+from .tool import WaldiezTool, get_tools_extra_requirements
 
 
 @dataclass(frozen=True, slots=True)
@@ -32,18 +42,18 @@ class Waldiez:
     @classmethod
     def from_dict(
         cls,
-        data: Dict[str, Any],
+        data: dict[str, Any],
         flow_id: Optional[str] = None,
         name: Optional[str] = None,
         description: Optional[str] = None,
-        tags: Optional[List[str]] = None,
-        requirements: Optional[List[str]] = None,
+        tags: Optional[list[str]] = None,
+        requirements: Optional[list[str]] = None,
     ) -> "Waldiez":
         """Create a Waldiez from dict.
 
         Parameters
         ----------
-        data : Dict[str, Any]
+        data : dict[str, Any]
             The data.
         flow_id : Optional[str], optional
             The flow id, by default None (retrieved from data or generated).
@@ -51,9 +61,9 @@ class Waldiez:
             The name, by default None (retrieved from data).
         description : Optional[str], optional
             The description, by default None (retrieved from data).
-        tags : Optional[List[str]], optional
+        tags : Optional[list[str]], optional
             The tags, by default None (retrieved from data).
-        requirements : Optional[List[str]], optional
+        requirements : Optional[list[str]], optional
             The requirements, by default None (retrieved from data).
 
         Returns
@@ -69,16 +79,17 @@ class Waldiez:
             tags=tags,
             requirements=requirements,
         )
-        return cls(flow=WaldiezFlow.model_validate(flow))
+        validated = WaldiezFlow.model_validate(flow)
+        return cls(flow=validated)  # pyright: ignore
 
     @classmethod
     def load(
         cls,
-        waldiez_file: Union[str, Path],
+        waldiez_file: str | Path,
         name: Optional[str] = None,
         description: Optional[str] = None,
-        tags: Optional[List[str]] = None,
-        requirements: Optional[List[str]] = None,
+        tags: Optional[list[str]] = None,
+        requirements: Optional[list[str]] = None,
     ) -> "Waldiez":
         """Load a Waldiez from a file.
 
@@ -90,9 +101,9 @@ class Waldiez:
             The name, by default None.
         description : Optional[str], optional
             The description, by default None.
-        tags : Optional[List[str]], optional
+        tags : Optional[list[str]], optional
             The tags, by default None.
-        requirements : Optional[List[str]], optional
+        requirements : Optional[list[str]], optional
             The requirements, by default None.
 
         Returns
@@ -105,7 +116,7 @@ class Waldiez:
         ValueError
             If the file is not found or invalid JSON.
         """
-        data: Dict[str, Any] = {}
+        data: dict[str, Any] = {}
         if not Path(waldiez_file).exists():
             raise ValueError(f"File not found: {waldiez_file}")
         with open(waldiez_file, "r", encoding="utf-8") as file:
@@ -142,23 +153,33 @@ class Waldiez:
         """
         return self.flow.model_dump_json(by_alias=by_alias, indent=indent)
 
+    @property
+    def id(self) -> str:
+        """Get the flow id."""
+        return self.flow.id
+
     @property
     def has_rag_agents(self) -> bool:
         """Check if the flow has RAG agents."""
-        return any(agent.agent_type == "rag_user" for agent in self.agents)
+        return any(agent.is_rag_user for agent in self.agents)
 
     @property
     def has_multimodal_agents(self) -> bool:
         """Check if the flow has multimodal agents."""
-        return any(agent.data.is_multimodal for agent in self.agents)
+        return any(
+            agent.data.is_multimodal
+            for agent in self.flow.data.agents.assistantAgents
+        )
 
     @property
     def has_captain_agents(self) -> bool:
         """Check if the flow has captain agents."""
-        return any(agent.agent_type == "captain" for agent in self.agents)
+        return any(agent.is_captain for agent in self.agents)
 
     @property
-    def chats(self) -> List[Tuple[WaldiezChat, WaldiezAgent, WaldiezAgent]]:
+    def initial_chats(
+        self,
+    ) -> list[WaldiezAgentConnection]:
         """Get the chats."""
         return self.flow.ordered_flow
 
@@ -174,15 +195,15 @@ class Waldiez:
         yield from self.flow.data.agents.members
 
     @property
-    def skills(self) -> Iterator[WaldiezSkill]:
-        """Get the flow skills.
+    def tools(self) -> Iterator[WaldiezTool]:
+        """Get the flow tools.
 
         Yields
         ------
-        WaldiezSkill
-            The skills.
+        WaldiezTool
+            The tools.
         """
-        yield from self.flow.data.skills
+        yield from self.flow.data.tools
 
     @property
     def models(self) -> Iterator[WaldiezModel]:
@@ -195,6 +216,14 @@ class Waldiez:
         """
         yield from self.flow.data.models
 
+    @property
+    def info(self) -> WaldiezFlowInfo:
+        """Get the flow info."""
+        return WaldiezFlowInfo.create(
+            agents=self.agents,
+            agent_names=self.flow.unique_names["agent_names"],
+        )
+
     @property
     def name(self) -> str:
         """Get the flow name."""
@@ -206,7 +235,7 @@ class Waldiez:
         return self.flow.description or "Waldiez Flow description"
 
     @property
-    def tags(self) -> List[str]:
+    def tags(self) -> list[str]:
         """Get the flow tags."""
         return self.flow.tags
 
@@ -226,14 +255,13 @@ class Waldiez:
         return self.flow.is_single_agent_mode
 
     @property
-    def requirements(self) -> List[str]:
+    def requirements(self) -> list[str]:
         """Get the flow requirements."""
         autogen_version = get_autogen_version()
         requirements_list = filter(
             # we use the fixed "ag2=={autogen_version}" below
             lambda requirement: not (
-                # cspell:disable-next-line
-                requirement.startswith("pyautogen")
+                requirement.startswith("ag2")
                 or requirement.startswith("ag2")
                 or requirement.startswith("autogen")
             ),
@@ -241,29 +269,13 @@ class Waldiez:
         )
         requirements = set(requirements_list)
         requirements.add(f"ag2[openai]=={autogen_version}")
-        if self.has_rag_agents:
+        if self.has_rag_agents:  # pragma: no branch
             rag_extras = get_retrievechat_extra_requirements(self.agents)
             requirements.update(rag_extras)
-        if self.has_multimodal_agents:
+        if self.has_multimodal_agents:  # pragma: no branch
             requirements.add(f"ag2[lmm]=={autogen_version}")
-        if self.has_captain_agents:
-            # pysqlite3-binary might not get installed on windows
-            captain_extras = [
-                "chromadb",
-                "sentence-transformers",
-                "huggingface-hub",
-                # tools:
-                "pillow",
-                "markdownify",
-                "arxiv",
-                "pymupdf",
-                "wikipedia-api",
-                "easyocr",
-                "python-pptx",
-                "openai-whisper",
-                "pandas",
-                "scipy",
-            ]
+        if self.has_captain_agents:  # pragma: no branch
+            captain_extras = get_captain_agent_extra_requirements()
             requirements.update(captain_extras)
         requirements.update(
             get_models_extra_requirements(
@@ -272,28 +284,50 @@ class Waldiez:
             )
         )
         requirements.update(
-            get_skills_extra_requirements(
-                self.skills,
+            get_tools_extra_requirements(
+                self.tools,
                 autogen_version=autogen_version,
             )
         )
         return sorted(requirements)
 
-    def get_flow_env_vars(self) -> List[Tuple[str, str]]:
+    # def get_flow_chat_info(self) ->
+
+    def get_flow_env_vars(self) -> list[tuple[str, str]]:
         """Get the flow environment variables.
 
         Returns
         -------
-        List[Tuple[str, str]]
+        list[tuple[str, str]]
             The environment variables for the flow.
         """
-        env_vars: List[Tuple[str, str]] = []
-        for skill in self.skills:
-            for secret_key, secret_value in skill.secrets.items():
+        env_vars: list[tuple[str, str]] = []
+        for tool in self.tools:
+            for secret_key, secret_value in tool.secrets.items():
                 env_vars.append((secret_key, secret_value))
+        for model in self.models:
+            api_env_key = model.api_key_env_key
+            api_key = model.api_key
+            if api_env_key and api_key:  # pragma: no branch
+                env_vars.append((api_env_key, api_key))
         return env_vars
 
-    def get_group_chat_members(self, agent: WaldiezAgent) -> List[WaldiezAgent]:
+    def get_root_group_manager(self) -> WaldiezGroupManager:
+        """Get the root group manager agent.
+
+        Returns
+        -------
+        WaldiezGroupManager
+            The root group manager agent.
+
+        Raises
+        ------
+        ValueError
+            If the root group manager agent is not found.
+        """
+        return self.flow.get_root_group_manager()
+
+    def get_group_chat_members(self, agent: WaldiezAgent) -> list[WaldiezAgent]:
         """Get the chat members that connect to a group chat manager agent.
 
         Parameters
@@ -306,23 +340,6 @@ class Waldiez:
         List[WaldiezAgent]
             The group chat members.
         """
-        if agent.agent_type != "manager":
+        if not agent.is_group_manager:
             return []
         return self.flow.get_group_chat_members(agent.id)
-
-    def get_swarm_members(
-        self, initial_agent: WaldiezAgent
-    ) -> Tuple[List[WaldiezAgent], Optional[WaldiezAgent]]:
-        """Get the chat members that connect to a swarm agent.
-
-        Parameters
-        ----------
-        initial_agent : WaldiezAgent
-            The initial agent.
-
-        Returns
-        -------
-        Tuple[List[WaldiezAgent], Optional[WaldiezAgent]]
-            The swarm agents and the user agent.
-        """
-        return self.flow.get_swarm_chat_members(initial_agent)

waldiez/runner.py

@@ -2,7 +2,7 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Run a waldiez flow.
 
-The flow is first converted to an autogen flow with agents, chats and skills.
+The flow is first converted to an autogen flow with agents, chats, and tools.
 We then chown to temporary directory, call the flow's `main()` and
 return the results. Before running the flow, any additional environment
 variables specified in the waldiez file are set.
@@ -10,14 +10,23 @@ variables specified in the waldiez file are set.
 
 # pylint: disable=import-outside-toplevel,reimported
 
+import gc
 import importlib.util
 import sys
 import tempfile
 from pathlib import Path
-from types import TracebackType
-from typing import TYPE_CHECKING, Dict, List, Optional, Set, Type, Union
+from types import ModuleType, TracebackType
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    Optional,
+    Set,
+    Type,
+    Union,
+)
 
 from .exporter import WaldiezExporter
+from .io import StructuredIOStream
 from .models.waldiez import Waldiez
 from .running import (
     a_chdir,
@@ -31,7 +40,6 @@ from .running import (
     reset_env_vars,
     set_env_vars,
 )
-from .utils import check_pysqlite3
 
 if TYPE_CHECKING:
     from autogen import ChatResult  # type: ignore
@@ -56,8 +64,8 @@ class WaldiezRunner:
         waldiez_file: Union[str, Path],
         name: Optional[str] = None,
         description: Optional[str] = None,
-        tags: Optional[List[str]] = None,
-        requirements: Optional[List[str]] = None,
+        tags: Optional[list[str]] = None,
+        requirements: Optional[list[str]] = None,
     ) -> "WaldiezRunner":
         """Create a WaldiezRunner instance from a file.
 
@@ -69,9 +77,9 @@ class WaldiezRunner:
             The name of the Waldiez, by default None.
         description : Optional[str], optional
             The description of the Waldiez, by default None.
-        tags : Optional[List[str]], optional
+        tags : Optional[list[str]], optional
             The tags of the Waldiez, by default None.
-        requirements : Optional[List[str]], optional
+        requirements : Optional[list[str]], optional
             The requirements of the Waldiez, by default None.
 
         Returns
@@ -132,6 +140,11 @@ class WaldiezRunner:
         """Get the Waldiez instance."""
         return self._waldiez
 
+    @property
+    def is_async(self) -> bool:
+        """Check if the workflow is async."""
+        return self.waldiez.is_async
+
     @property
     def running(self) -> bool:
         """Get the running status."""
@@ -148,8 +161,6 @@ class WaldiezRunner:
         extra_requirements = {
             req for req in self.waldiez.requirements if req not in sys.modules
         }
-        if self.waldiez.has_captain_agents:
-            check_pysqlite3()
         return extra_requirements
 
     def install_requirements(self) -> None:
@@ -159,7 +170,6 @@ class WaldiezRunner:
         extra_requirements = self.gather_requirements()
         if extra_requirements:
             install_requirements(extra_requirements, printer)
-            refresh_environment()
 
     async def a_install_requirements(self) -> None:
         """Install the requirements for the flow asynchronously."""
@@ -168,14 +178,36 @@ class WaldiezRunner:
         extra_requirements = self.gather_requirements()
         if extra_requirements:
             await a_install_requirements(extra_requirements, printer)
-            refresh_environment()
+
+    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()
+        spec = importlib.util.spec_from_file_location(
+            module_name, temp_dir / file_name
+        )
+        if not spec or not spec.loader:
+            raise ImportError("Could not import the flow")
+        sys.path.insert(0, str(temp_dir))
+        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())
+        return module, old_vars
 
     def _run(
         self,
         output_path: Optional[Union[str, Path]],
         uploads_root: Optional[Union[str, Path]],
+        use_structured_io: bool = False,
         skip_mmd: bool = False,
-    ) -> Union["ChatResult", List["ChatResult"], Dict[int, "ChatResult"]]:
+    ) -> Union["ChatResult", list["ChatResult"], dict[int, "ChatResult"]]:
         """Run the Waldiez workflow.
 
         Parameters
@@ -184,12 +216,14 @@ class WaldiezRunner:
             The output path.
         uploads_root : Optional[Union[str, Path]]
             The runtime uploads root.
+        use_structured_io : bool
+            Whether to use structured IO instead of the default 'input/print'.
         skip_mmd : bool
             Whether to skip the Mermaid diagram generation.
 
         Returns
         -------
-        Union[ChatResult, List[ChatResult]]
+        Union[ChatResult, list[ChatResult]]
             The result(s) of the chat(s).
         """
         temp_dir = Path(tempfile.mkdtemp())
@@ -197,8 +231,7 @@ class WaldiezRunner:
         module_name = file_name.replace(".py", "")
         if not self._called_install_requirements:
             self.install_requirements()
-        else:
-            refresh_environment()
+        refresh_environment()
         printer = get_printer()
         printer(
             "Requirements installed.\n"
@@ -206,21 +239,24 @@ class WaldiezRunner:
             "you might need to restart the kernel."
         )
         results: Union[
-            "ChatResult", List["ChatResult"], Dict[int, "ChatResult"]
+            "ChatResult", list["ChatResult"], dict[int, "ChatResult"]
         ] = []
         with chdir(to=temp_dir):
-            self._exporter.export(Path(file_name))
-            spec = importlib.util.spec_from_file_location(
-                module_name, temp_dir / file_name
+            module, old_vars = self._before_run(
+                temp_dir=temp_dir,
+                file_name=file_name,
+                module_name=module_name,
+                printer=printer,
             )
-            if not spec or not spec.loader:
-                raise ImportError("Could not import the flow")
-            sys.path.insert(0, str(temp_dir))
-            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...")
-            results = module.main()
+            if use_structured_io:
+                stream = StructuredIOStream(
+                    timeout=120,  # 2 minutes
+                    uploads_root=uploads_root,
+                )
+                with StructuredIOStream.set_default(stream):
+                    results = module.main()
+            else:
+                results = module.main()
             printer("<Waldiez> - Workflow finished")
             sys.path.pop(0)
             reset_env_vars(old_vars)
@@ -231,22 +267,24 @@ class WaldiezRunner:
             flow_name=self.waldiez.name,
             skip_mmd=skip_mmd,
         )
+        gc.collect()
+        refresh_environment()
         return results
 
     async def _a_run(
         self,
         output_path: Optional[Union[str, Path]],
         uploads_root: Optional[Union[str, Path]],
+        use_structured_io: bool = False,
         skip_mmd: bool = False,
-    ) -> Union["ChatResult", List["ChatResult"], Dict[int, "ChatResult"]]:
+    ) -> Union["ChatResult", list["ChatResult"], dict[int, "ChatResult"]]:
         """Run the Waldiez workflow asynchronously."""
         temp_dir = Path(tempfile.mkdtemp())
         file_name = before_run(output_path, uploads_root)
         module_name = file_name.replace(".py", "")
         if not self._called_install_requirements:
             await self.a_install_requirements()
-        else:
-            refresh_environment()
+        refresh_environment()
         printer = get_printer()
         printer(
             "Requirements installed.\n"
@@ -254,21 +292,24 @@ class WaldiezRunner:
             "you might need to restart the kernel."
         )
         results: Union[
-            "ChatResult", List["ChatResult"], Dict[int, "ChatResult"]
+            "ChatResult", list["ChatResult"], dict[int, "ChatResult"]
         ] = []
         async with a_chdir(to=temp_dir):
-            self._exporter.export(Path(file_name))
-            spec = importlib.util.spec_from_file_location(
-                module_name, temp_dir / file_name
+            module, old_vars = self._before_run(
+                temp_dir=temp_dir,
+                file_name=file_name,
+                module_name=module_name,
+                printer=printer,
             )
-            if not spec or not spec.loader:
-                raise ImportError("Could not import the flow")
-            sys.path.insert(0, str(temp_dir))
-            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...")
-            results = await module.main()
+            if use_structured_io:
+                stream = StructuredIOStream(
+                    timeout=120,  # 2 minutes
+                    uploads_root=uploads_root,
+                )
+                with StructuredIOStream.set_default(stream):
+                    results = await module.main()
+            else:
+                results = await module.main()
             sys.path.pop(0)
             reset_env_vars(old_vars)
         after_run(
@@ -278,14 +319,17 @@ class WaldiezRunner:
             flow_name=self.waldiez.name,
             skip_mmd=skip_mmd,
         )
+        gc.collect()
+        refresh_environment()
         return results
 
     def run(
         self,
         output_path: Optional[Union[str, Path]] = None,
         uploads_root: Optional[Union[str, Path]] = None,
+        use_structured_io: bool = False,
         skip_mmd: bool = False,
-    ) -> Union["ChatResult", List["ChatResult"], Dict[int, "ChatResult"]]:
+    ) -> Union["ChatResult", list["ChatResult"], dict[int, "ChatResult"]]:
         """Run the Waldiez workflow.
 
         Parameters
@@ -294,12 +338,15 @@ class WaldiezRunner:
             The output path, by default None.
         uploads_root : Optional[Union[str, Path]], optional
             The uploads root, to get user-uploaded files, by default None.
+        use_structured_io : bool, optional
+            Whether to use structured IO instead of the default 'input/print',
+            by default False.
         skip_mmd : bool, optional
             Whether to skip the Mermaid diagram generation, by default False.
 
         Returns
         -------
-        Union["ChatResult", List["ChatResult"], Dict[int, "ChatResult"]]
+        Union["ChatResult", list["ChatResult"], dict[int, "ChatResult"]]
             The result(s) of the chat(s).
 
         Raises
@@ -316,6 +363,7 @@ class WaldiezRunner:
                     self._a_run,
                     output_path,
                     uploads_root,
+                    use_structured_io,
                     skip_mmd,
                 )
         if self._running is True:
@@ -323,7 +371,12 @@ class WaldiezRunner:
         self._running = True
         file_path = output_path or self._file_path
         try:
-            return self._run(file_path, uploads_root, skip_mmd)
+            return self._run(
+                file_path,
+                uploads_root=uploads_root,
+                use_structured_io=use_structured_io,
+                skip_mmd=skip_mmd,
+            )
         finally:
             self._running = False
 
@@ -331,7 +384,9 @@ class WaldiezRunner:
         self,
         output_path: Optional[Union[str, Path]] = None,
         uploads_root: Optional[Union[str, Path]] = None,
-    ) -> Union["ChatResult", List["ChatResult"]]:
+        use_structured_io: bool = False,
+        skip_mmd: bool = False,
+    ) -> Union["ChatResult", list["ChatResult"], dict[int, "ChatResult"]]:
         """Run the Waldiez workflow asynchronously.
 
         Parameters
@@ -340,10 +395,15 @@ class WaldiezRunner:
             The output path, by default None.
         uploads_root : Optional[Union[str, Path]], optional
             The uploads root, to get user-uploaded files, by default None.
+        use_structured_io : bool, optional
+            Whether to use structured IO instead of the default 'input/print',
+            by default False.
+        skip_mmd : bool, optional
+            Whether to skip the Mermaid diagram generation, by default False.
 
         Returns
         -------
-        Union[ChatResult, List[ChatResult]]
+        Union[ChatResult, list[ChatResult]], dict[int, ChatResult]
             The result(s) of the chat(s).
 
         Raises
@@ -356,6 +416,11 @@ class WaldiezRunner:
         self._running = True
         file_path = output_path or self._file_path
         try:
-            return await self._a_run(file_path, uploads_root)
+            return await self._a_run(
+                file_path,
+                uploads_root=uploads_root,
+                use_structured_io=use_structured_io,
+                skip_mmd=skip_mmd,
+            )
         finally:
             self._running = False