waldiez 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

waldiez/_version.py

@@ -1,3 +1,3 @@
 """Version information for Waldiez."""
 
-__version__ = "0.1.7"
+__version__ = "0.1.9"

waldiez/exporting/agents/rag_user/chroma_utils.py

@@ -32,9 +32,9 @@ def _get_chroma_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
         # SyntaxError: (unicode error) 'unicodeescape' codec can't decode bytes
         # in position 2-3: truncated \UXXXXXXXX escape
         local_path = Path(agent.retrieve_config.db_config.local_storage_path)
-        client_str += f'PersistentClient(path=r"{local_path}")'
+        client_str += f'PersistentClient(path=r"{local_path}", settings=Settings(anonymized_telemetry=False))'
     else:
-        client_str += "Client()"
+        client_str += "Client(Settings(anonymized_telemetry=False))"
     return client_str, to_import
 
 
@@ -97,11 +97,11 @@ def get_chroma_db_args(
         - The custom embedding function.
         - Any additional content to be used before the `kwargs` string.
     """
-    client_str, to_import_client = _get_chroma_client_string(agent)
+    client_str, client_to_import = _get_chroma_client_string(agent)
     embedding_function_arg, to_import_embedding, embedding_function_body = (
         _get_chroma_embedding_function_string(agent, agent_name)
     )
-    to_import = {to_import_client}
+    to_import = {client_to_import, "from chromadb.config import Settings"}
     if to_import_embedding:
         to_import.add(to_import_embedding)
     kwarg_string = (
@@ -109,8 +109,7 @@ def get_chroma_db_args(
         f"            embedding_function={embedding_function_arg},\n"
     )
     # The RAG example:
-    # https://microsoft.github.io/autogen/docs/\
-    #                                       notebooks/agentchat_groupchat_RAG
+    # https://ag2ai.github.io/ag2/docs/notebooks/agentchat_groupchat_RAG/
     # raises `InvalidCollectionException`: Collection groupchat does not exist.
     # https://github.com/chroma-core/chroma/issues/861
     # https://github.com/microsoft/autogen/issues/3551#issuecomment-2366930994

waldiez/exporting/chats/nested.py

@@ -14,7 +14,6 @@ from .helpers import escape_summary_args_quotes
 
 
 def get_nested_chat_trigger_agent_names(
-    all_chats: List[WaldiezChat],
     nested_chat: WaldiezAgentNestedChat,
     agent_names: Dict[str, str],
 ) -> str:
@@ -22,8 +21,6 @@ def get_nested_chat_trigger_agent_names(
 
     Parameters
     ----------
-    all_chats : List[WaldiezChat]
-        All the chats in the flow.
     nested_chat : WaldiezAgentNestedChat
         The nested chat.
     agent_names : Dict[str, str]
@@ -34,14 +31,7 @@ def get_nested_chat_trigger_agent_names(
     str
         The trigger agent names.
     """
-    trigger_agent_ids: List[str] = []
-    for message in nested_chat.triggered_by:
-        waldiez_chat = next(chat for chat in all_chats if chat.id == message.id)
-        if message.is_reply:
-            trigger_agent_ids.append(waldiez_chat.target)
-        else:
-            trigger_agent_ids.append(waldiez_chat.source)
-    agents = [agent_names[agent_id] for agent_id in trigger_agent_ids]
+    agents = [agent_names[agent_id] for agent_id in nested_chat.triggered_by]
     trigger_string = f'{[", ".join(agents)]}'
     return trigger_string.replace("'", '"')
 
@@ -228,7 +218,7 @@ def export_nested_chat(
     use_suffix = len(agent.data.nested_chats) > 1
     for index, entry in enumerate(agent.data.nested_chats):
         trigger_names = get_nested_chat_trigger_agent_names(
-            all_chats=all_chats, nested_chat=entry, agent_names=agent_names
+            nested_chat=entry, agent_names=agent_names
         )
         chat_queue, extra_methods = get_nested_chat_queue(
             nested_chat=entry,

waldiez/exporting/utils/logging_utils.py

@@ -14,6 +14,7 @@ get_sqlite_to_csv_call_string
 
 
 # Check issue:
+# Also check if in ag2 this still applies
 # https://github.com/microsoft/autogen/issues/2286
 # we cannot log new agents if they have code_execution enabled
 # we get `Path` is not JSON serializable (on code_executor)

waldiez/exporting/utils/path_check.py

@@ -2,11 +2,12 @@
 
 import os
 from pathlib import Path
+from typing import Optional
 
 # pylint: disable=broad-except
 
 
-def _is_local_path(string: str) -> bool:
+def _check_local_path(string: str) -> Optional[Path]:
     """Check if a string is a local path.
 
     Parameters
@@ -21,9 +22,11 @@ def _is_local_path(string: str) -> bool:
     """
     try:
         path = Path(string).resolve()
-        return path.exists()
     except Exception:  # pragma: no cover
-        return False
+        return None
+    if path.exists():
+        return path
+    return None
 
 
 def get_path_string(string: str) -> str:
@@ -41,8 +44,8 @@ def get_path_string(string: str) -> str:
     """
     # On windows, we get paths like "C:\path\to\file"
     # if so, let's try to avoid invalid escape sequences
-    if not _is_local_path(string):
+    if not _check_local_path(string):
         return string
     if os.name == "nt":  # pragma: no cover
         return f"r'{string}'"
-    return f"'{string}'"
+    return f"{Path(string).resolve()}"

waldiez/io/__init__.py

@@ -5,100 +5,48 @@ It is meant to be used when we want to use custom
 is used to trigger a UI element that requires user input.
 and sends back the user's input to the websocket. In the same
 way, we can use it to forward what is meant to be printed.
-
-
-We use:
-
-- one tcp server to handle messaging between the clients
-- one tcp client (provider) to set and forward the user's input
-    that we got elsewhere (e.g. from a websocket connection)
-- one tcp client (consumer) to ask and get the input from the provider
 """
 
-import socket
-from contextlib import closing
+import threading
 from typing import Any, Callable, Optional
 
 from autogen.io import IOStream  # type: ignore[import-untyped]
 
-from .stream import TCPConsumer, TCPProvider, TCPServer
-
 
 class WaldiezIOStream(IOStream):
-    """Custom IOStream class to handle the `print` and `input` functions."""
+    """Custom IOStream class to handle `print` and `input` functions."""
 
     def __init__(
         self,
-        port: int = 0,
-        input_timeout: float = 60,
+        input_timeout: float = 60.0,
         print_function: Optional[Callable[..., None]] = None,
         on_prompt_input: Optional[Callable[[str], None]] = None,
     ) -> None:
-        """Initialize the IOStream.
+        """
+        Initialize the IOStream.
 
         Parameters
         ----------
-        port : int, optional
-            The port to use, by default 0 (auto-assign).
         input_timeout : float, optional
-            The input timeout, by default 60.
+            The input timeout in seconds, by default 60.0.
         print_function : Optional[Callable[..., None]], optional
-            The print function to use, by default None.
+            The function to handle print operations, by default None.
         on_prompt_input : Optional[Callable[[str], None]], optional
-            The function to call for getting an input, by default None.
+            The function to call for processing input prompts, by default None.
+
+        Notes
+        -----
+        - on_prompt_input: It does not return a string (like 'input' does).
+            Instead, it is meant to be used to forward the prompt somewhere else
+            (e.g., a websocket). When we get the input, we can call
+            `waldiez_io_stream.set_input(input_data)` with the input data.
         """
-        self._print_function = print_function
-        if port == 0:
-            port = get_available_port()
-        self._port = port
-        self._input_timeout = input_timeout
-        self._server = TCPServer(port)
-        self._server.start()
-        self._provider = TCPProvider("localhost", port, response=None)
-        self._on_prompt_input = on_prompt_input
-
-    @property
-    def print_function(self) -> Optional[Callable[..., None]]:
-        """Get the print function."""
-        return self._print_function
-
-    def open(self) -> None:
-        """Start the server."""
-        if not self._server.is_running():
-            self._server.start()
-
-    def close(self) -> None:
-        """Stop the server and the provider."""
-        # pylint: disable=broad-except
-        if self._server.is_running():
-            try:
-                self._server.stop()
-            except BaseException:  # pragma: no cover
-                pass
-        try:
-            self._provider.stop()
-        except BaseException:  # pragma: no cover
-            pass
-
-    def __del__(self) -> None:  # pragma: no cover
-        """Delete the instance."""
-        self.close()
-
-    def forward_input(self, input_data: str) -> None:
-        """Forward the user's input to the provider.
-
-        When we have the input data
-            e.g. from 'input(..)' or from a websocket connection,
-            we can forward it to the provider (the tcp client)
-            to make it available to the consumer (the other tcp client).
-        Parameters
-        ----------
-        input_data : str
-            The input data to forward.
-        """
-        if not self._provider.is_running():
-            self._provider.start()
-        self._provider.set_response(input_data)
+        self.input_timeout = input_timeout  # Timeout for input
+        self.print_function = print_function  # Custom print handler
+        self._on_prompt_input = on_prompt_input  # Custom input prompt handler
+        self.current_input: Optional[str] = None  # Store the current input
+        self._input_event = threading.Event()  # Event to signal input readiness
+        self.allow_input = True  # Flag to allow or block input setting
 
     def print(
         self,
@@ -107,7 +55,8 @@ class WaldiezIOStream(IOStream):
         end: str = "\n",
         flush: bool = False,
     ) -> None:
-        """Mock the print function.
+        """
+        Mock the `print` function.
 
         Parameters
         ----------
@@ -116,66 +65,78 @@ class WaldiezIOStream(IOStream):
         sep : str, optional
             The separator, by default " ".
         end : str, optional
-            The end, by default a new line.
+            The ending string, by default a new line.
         flush : bool, optional
-            Whether to flush, by default False.
+            Whether to flush the output, by default False.
         """
         print_function: Callable[..., None] = self.print_function or print
         print_function(*objects, sep=sep, end=end, flush=flush)
 
     def input(self, prompt: str = "", *, password: bool = False) -> str:
-        """Mock the input function.
+        """
+        Mock the `input` function with optional timeout handling.
 
         Parameters
         ----------
         prompt : str, optional
             The prompt to show, by default "".
-        password : bool, optional (not used)
-            Whether to show the input as password, by default False.
+        password : bool, optional
+            Whether to hide the input as password (not used), by default False.
 
         Returns
         -------
         str
-            The user's input.
+            The user's input or '\n' if timeout occurs.
         """
         _prompt = prompt or "Your input:"
-        if _prompt in (">", "> "):
+        if _prompt in (">", "> "):  # pragma: no cover
             _prompt = "Your input:"
         if prompt:
             if self._on_prompt_input:
                 self._on_prompt_input(_prompt)
             self.print(_prompt, end="")
-        if not self._provider.is_running():
-            self._provider.start()
-        # wait for the provider to start
-        self._provider.wait(timeout=self._input_timeout)
-        if not self._provider.is_running():  # pragma: no cover
-            self.print(
-                "WARNING: Provider is not running. Was an input expected?"
+
+        # Only reset if no input is currently set
+        # e.g. handle the case when we call set_input(..) before input(..)
+        already_set = self._input_event.is_set()
+        if not self._input_event.is_set():
+            self.current_input = None  # Reset previous input
+            self._input_event.clear()  # Clear the event before waiting
+
+        # Wait for input or timeout
+        if not self._input_event.wait(self.input_timeout):
+            # Timeout occurred, return what we have so far
+            to_return = (
+                self.current_input if self.current_input is not None else "\n"
             )
-            return "\n"
-        consumer = TCPConsumer(
-            "localhost", self._port, timeout=self._input_timeout
+            self.current_input = None
+            self._input_event.clear()
+            # if we had already set the input, return it
+            return to_return if already_set else "\n"
+
+        # Input is ready, return it
+        to_return = (
+            self.current_input if self.current_input is not None else "\n"
         )
-        consumer.start()
-        # send the prompt and wait for the response
-        consumer.send_prompt(_prompt)
-        response = consumer.get_response()
-        consumer.stop()
-        self._provider.stop()
-        # return the response or a line break (i.e. no input)
-        return response or "\n"
-
-
-def get_available_port() -> int:
-    """Get an available port.
-
-    Returns
-    -------
-    int
-        Available port.
-    """
-    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as soc:
-        soc.bind(("", 0))
-        soc.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        return soc.getsockname()[1]
+        self.current_input = None
+        self._input_event.clear()  # Clear the event after waiting
+        return to_return
+
+    def set_input(self, value: str) -> None:
+        """
+        Set the input value and signal that input is ready.
+
+        Parameters
+        ----------
+        value : str
+            The value to set as input.
+        """
+        if self.allow_input:  # Respect the allow_input flag
+            self.current_input = value
+            self._input_event.set()  # Signal that input is ready
+
+    def reset_state(self) -> None:
+        """Reset the IOStream state for testing."""
+        self.current_input = None
+        self._input_event.clear()
+        self.allow_input = True  # Re-enable input setting

waldiez/models/agents/agent/nested_chat.py

@@ -40,22 +40,18 @@ class WaldiezAgentNestedChat(WaldiezBase):
 
     Attributes
     ----------
-    triggered_by : List[WaldiezAgentNestedChatMessage]
-        A list of chats (id and is_reply) to determine
-        the triggering of the nested chat.
+    triggered_by : List[str]
+        A list of agent ids that trigger the nested chat.
     messages : List[WaldiezAgentNestedChatMessage]
         The list of messages (chat ids and 'is_reply'z)
         to include the in the nested chat registration.
     """
 
     triggered_by: Annotated[
-        List[WaldiezAgentNestedChatMessage],
+        List[str],
         Field(
             title="Triggered By",
-            description=(
-                "A list of chats (id and is_reply) to determine"
-                "the triggering of the nested chat."
-            ),
+            description=("A list of agent ids that trigger the nested chat."),
             alias="triggeredBy",
             default_factory=list,
         ),

waldiez/models/waldiez.py

@@ -205,18 +205,16 @@ class Waldiez:
         requirements_list = filter(
             lambda requirement: not (
                 requirement.startswith("pyautogen")
-                or requirement.startswith("autogen-agentchat")
+                or requirement.startswith("ag2")
             ),
             self.flow.requirements,
         )
         requirements = set(requirements_list)
         if self.has_rag_agents:
-            requirements.add(
-                f"autogen-agentchat[retrievechat]=={autogen_version}"
-            )
+            requirements.add(f"ag2[retrievechat]=={autogen_version}")
         else:
-            requirements.add(f"autogen-agentchat=={autogen_version}")
-        # ref: https://github.com/microsoft/autogen/blob/main/setup.py
+            requirements.add(f"ag2=={autogen_version}")
+        # ref: https://github.com/ag2ai/ag2/blob/main/setup.py
         models_with_additional_requirements = [
             "together",
             "gemini",
@@ -229,8 +227,7 @@ class Waldiez:
         for model in self.models:
             if model.data.api_type in models_with_additional_requirements:
                 requirements.add(
-                    f"autogen-agentchat[{model.data.api_type}]=="
-                    f"{autogen_version}"
+                    f"ag2[{model.data.api_type}]==" f"{autogen_version}"
                 )
         return list(requirements)