waldiez 0.1.18__tar.gz → 0.1.20__tar.gz

{waldiez-0.1.18 → waldiez-0.1.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waldiez
-Version: 0.1.18
+Version: 0.1.20
 Summary: waldiez
 Project-URL: homepage, https://waldiez.github.io/waldiez/
 Project-URL: repository, https://github.com/waldiez/waldiez.git

{waldiez-0.1.18 → waldiez-0.1.20}/pyproject.toml

@@ -314,6 +314,7 @@ python_files = [
 ]
 addopts = """
     -d \
+    -s \
     -n auto \
     --exitfirst \
     --durations=10 \

{waldiez-0.1.18 → waldiez-0.1.20}/waldiez/__init__.py

@@ -1,6 +1,7 @@
 """Waldiez package."""
 
 import logging
+import warnings
 
 from ._version import __version__
 from .conflict_checker import check_conflicts
@@ -8,6 +9,8 @@ from .exporter import WaldiezExporter
 from .models import Waldiez
 from .runner import WaldiezRunner
 
+warnings.filterwarnings("ignore", "flaml.automl is not available")
+
 
 # pylint: disable=too-few-public-methods
 class FlamlFilter(logging.Filter):

{waldiez-0.1.18 → waldiez-0.1.20}/waldiez/_version.py

@@ -1,3 +1,3 @@
 """Version information for Waldiez."""
 
-__version__ = "0.1.18"
+__version__ = "0.1.20"

{waldiez-0.1.18 → waldiez-0.1.20}/waldiez/models/model/__init__.py

@@ -1,9 +1,10 @@
 """Waldiez model."""
 
-from .model import WaldiezModel
+from .model import DEFAULT_BASE_URLS, WaldiezModel
 from .model_data import WaldiezModelAPIType, WaldiezModelData, WaldiezModelPrice
 
 __all__ = [
+    "DEFAULT_BASE_URLS",
     "WaldiezModel",
     "WaldiezModelData",
     "WaldiezModelPrice",

{waldiez-0.1.18 → waldiez-0.1.20}/waldiez/models/model/model.py

@@ -7,7 +7,16 @@ from pydantic import Field
 from typing_extensions import Annotated, Literal
 
 from ..common import WaldiezBase, now
-from .model_data import WaldiezModelData
+from .model_data import WaldiezModelAPIType, WaldiezModelData
+
+DEFAULT_BASE_URLS: Dict[WaldiezModelAPIType, str] = {
+    "google": "https://generativelanguage.googleapis.com/v1beta",
+    "anthropic": "https://api.anthropic.com/v1",
+    "mistral": "https://api.mistral.ai/v1",
+    "groq": "https://api.groq.com/openai/v1",
+    "together": "https://api.together.xyz/v1",
+    "nim": "https://integrate.api.nvidia.com/v1",
+}
 
 
 class WaldiezModel(WaldiezBase):
@@ -166,4 +175,30 @@ class WaldiezModel(WaldiezBase):
             value = getattr(self, attr)
             if value:
                 _llm_config[attr] = value
-        return _llm_config
+        return set_default_base_url(_llm_config, self.data.api_type)
+
+
+def set_default_base_url(
+    llm_config: Dict[str, Any], api_type: WaldiezModelAPIType
+) -> Dict[str, Any]:
+    """Set the default base url if not provided.
+
+    Parameters
+    ----------
+    llm_config : Dict[str, Any]
+        The llm config dictionary.
+    api_type : str
+        The api type.
+
+    Returns
+    -------
+    Dict[str, Any]
+        The llm config dictionary with the default base url set.
+    """
+    if api_type in ("openai", "other", "azure"):
+        return llm_config
+    if "base_url" not in llm_config or not llm_config["base_url"]:
+        dict_copy = llm_config.copy()
+        dict_copy["base_url"] = DEFAULT_BASE_URLS.get(api_type, "")
+        return dict_copy
+    return llm_config

{waldiez-0.1.18 → waldiez-0.1.20}/waldiez/runner.py

@@ -7,6 +7,8 @@ return the results. Before running the flow, any additional environment
 variables specified in the waldiez file are set.
 """
 
+# pylint: disable=import-outside-toplevel,reimported
+
 import datetime
 import importlib.util
 import io
@@ -16,6 +18,7 @@ import site
 import subprocess  # nosemgrep # nosec
 import sys
 import tempfile
+import warnings
 from contextlib import contextmanager
 from pathlib import Path
 from types import TracebackType
@@ -26,6 +29,7 @@ from typing import (
     Iterator,
     List,
     Optional,
+    Tuple,
     Type,
     Union,
 )
@@ -59,11 +63,6 @@ def _chdir(to: Union[str, Path]) -> Iterator[None]:
         os.chdir(old_cwd)
 
 
-def refresh_site_packages() -> None:
-    """Refresh the site packages."""
-    site.main()
-
-
 class WaldiezRunner:
     """Waldiez runner class."""
 
@@ -75,6 +74,7 @@ class WaldiezRunner:
         self._running = False
         self._file_path = file_path
         self._exporter = WaldiezExporter(waldiez)
+        self._called_install_requirements = False
 
     @classmethod
     def load(
@@ -147,8 +147,10 @@ class WaldiezRunner:
         """Get the running status."""
         return self._running
 
-    def _install_requirements(self, printer: Callable[..., None]) -> None:
+    def install_requirements(self) -> None:
         """Install the requirements for the flow."""
+        self._called_install_requirements = True
+        printer = get_printer()
         extra_requirements = set(
             req for req in self.waldiez.requirements if req not in sys.modules
         )
@@ -169,13 +171,7 @@ class WaldiezRunner:
                 if proc.stderr:
                     for line in io.TextIOWrapper(proc.stderr, encoding="utf-8"):
                         printer(line.strip())
-            printer("Refreshing site packages...")
-            refresh_site_packages()
-            printer(
-                "Requirements installed.\n"
-                "NOTE: If new packages were added and you are using Jupyter, "
-                "you might need to restart the kernel."
-            )
+            refresh_environment()
 
     @staticmethod
     def _after_run(
@@ -245,11 +241,16 @@ class WaldiezRunner:
         Union[ChatResult, List[ChatResult]]
             The result(s) of the chat(s).
         """
-        # pylint: disable=import-outside-toplevel
-        from autogen.io import IOStream  # type: ignore
-
-        printer = IOStream.get_default().print
-        self._install_requirements(printer)
+        if not self._called_install_requirements:
+            self.install_requirements()
+        else:
+            refresh_environment()
+        printer = get_printer()
+        printer(
+            "Requirements installed.\n"
+            "NOTE: If new packages were added and you are using Jupyter, "
+            "you might need to restart the kernel."
+        )
         results: Union["ChatResult", List["ChatResult"]] = []
         if not uploads_root:
             uploads_root = Path(tempfile.mkdtemp())
@@ -277,7 +278,7 @@ class WaldiezRunner:
             old_vars = self._set_env_vars()
             module = importlib.util.module_from_spec(spec)
             spec.loader.exec_module(module)
-            printer("Starting workflow...")
+            printer("<Waldiez> - Starting workflow...")
             results = module.main()
             sys.path.pop(0)
             self._reset_env_vars(old_vars)
@@ -329,3 +330,91 @@ def in_virtualenv() -> bool:
     return hasattr(sys, "real_prefix") or (
         hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix
     )
+
+
+def refresh_environment() -> None:
+    """Refresh the environment."""
+    # backup the default IOStream
+    from autogen.io import IOStream  # type: ignore
+
+    default_io_stream = IOStream.get_default()
+    site.main()
+    # pylint: disable=import-outside-toplevel
+    modules_to_reload = [
+        mod for mod in sys.modules if mod.startswith("autogen")
+    ]
+    for mod in modules_to_reload:
+        del sys.modules[mod]
+    import autogen
+    from autogen.io import IOStream
+
+    importlib.reload(autogen)
+    # restore the default IOStream
+    IOStream.set_global_default(default_io_stream)
+    warnings.filterwarnings("ignore", "flaml.automl is not available")
+
+
+def get_printer() -> Callable[..., None]:
+    """Get the printer function.
+
+    Returns
+    -------
+    Callable[..., None]
+        The printer function.
+    """
+    from autogen.io import IOStream
+
+    printer = IOStream.get_default().print
+
+    def safe_printer(*args: object, **kwargs: object) -> None:
+        try:
+            printer(*args, **kwargs)
+        except UnicodeEncodeError:
+            # pylint: disable=too-many-try-statements
+            try:
+                msg, flush = get_what_to_print(*args, **kwargs)
+                printer(msg, end="", flush=flush)
+            except UnicodeEncodeError:
+                sys.stdout = io.TextIOWrapper(
+                    sys.stdout.buffer, encoding="utf-8"
+                )
+                sys.stderr = io.TextIOWrapper(
+                    sys.stderr.buffer, encoding="utf-8"
+                )
+                try:
+                    printer(*args, **kwargs)
+                except UnicodeEncodeError:
+                    sys.stderr.write(
+                        "Could not print the message due to encoding issues.\n"
+                    )
+
+    return safe_printer
+
+
+def get_what_to_print(*args: object, **kwargs: object) -> Tuple[str, bool]:
+    """Get what to print.
+
+    Parameters
+    ----------
+    args : object
+        The arguments.
+    kwargs : object
+        The keyword arguments.
+
+    Returns
+    -------
+    Tuple[str, bool]
+        The message and whether to flush.
+    """
+    sep = kwargs.get("sep", " ")
+    if not isinstance(sep, str):
+        sep = " "
+    end = kwargs.get("end", "\n")
+    if not isinstance(end, str):
+        end = "\n"
+    flush = kwargs.get("flush", False)
+    if not isinstance(flush, bool):
+        flush = False
+    msg = sep.join(str(arg) for arg in args) + end
+    utf8_msg = msg.encode("utf-8", errors="replace").decode("utf-8")
+    return utf8_msg, flush