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

waldiez/models/model/model.py

@@ -3,15 +3,15 @@
 """Waldiez model model."""
 
 import os
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 from pydantic import Field
 from typing_extensions import Annotated, Literal
 
 from ..common import WaldiezBase, now
-from .model_data import WaldiezModelAPIType, WaldiezModelData
+from .model_data import WaldiezModelAPIType, WaldiezModelAWS, WaldiezModelData
 
-DEFAULT_BASE_URLS: Dict[WaldiezModelAPIType, str] = {
+DEFAULT_BASE_URLS: dict[WaldiezModelAPIType, str] = {
     "deepseek": "https://api.deepseek.com/v1",
     "google": "https://generativelanguage.googleapis.com/v1beta",
     "anthropic": "https://api.anthropic.com/v1",
@@ -24,7 +24,7 @@ DEFAULT_BASE_URLS: Dict[WaldiezModelAPIType, str] = {
 
 
 # we can omit the base_url for these models
-MODEL_NEEDS_BASE_URL: Dict[WaldiezModelAPIType, bool] = {
+MODEL_NEEDS_BASE_URL: dict[WaldiezModelAPIType, bool] = {
     "openai": False,
     "azure": False,
     "google": False,
@@ -50,9 +50,9 @@ class WaldiezModel(WaldiezBase):
         The name of the model.
     description : str
         The description of the model.
-    tags : List[str]
+    tags : list[str]
         The tags of the model.
-    requirements : List[str]
+    requirements : list[str]
         The requirements of the model.
     created_at : str
         The date and time when the model was created.
@@ -86,7 +86,7 @@ class WaldiezModel(WaldiezBase):
         ),
     ]
     tags: Annotated[
-        List[str],
+        list[str],
         Field(
             default_factory=list,
             title="Tags",
@@ -94,7 +94,7 @@ class WaldiezModel(WaldiezBase):
         ),
     ]
     requirements: Annotated[
-        List[str],
+        list[str],
         Field(
             default_factory=list,
             title="Requirements",
@@ -175,7 +175,7 @@ class WaldiezModel(WaldiezBase):
         return api_key or "REPLACE_ME"
 
     @property
-    def price(self) -> Optional[List[float]]:
+    def price(self) -> Optional[list[float]]:
         """Get the model's price."""
         if self.data.price is None:
             return None
@@ -188,29 +188,31 @@ class WaldiezModel(WaldiezBase):
             ]
         return None
 
-    def get_llm_config(self, skip_price: bool = False) -> Dict[str, Any]:
+    def get_llm_config(self, skip_price: bool = False) -> dict[str, Any]:
         """Get the model's llm config.
 
         Parameters
         ----------
         skip_price : bool, optional
-            Whether to skip the price, by default False
+            Whether to skip the price, by default, False
 
         Returns
         -------
-        Dict[str, Any]
+        dict[str, Any]
             The model's llm config dictionary.
         """
-        _llm_config: Dict[str, Any] = {}
+        # noinspection PyDictCreation
+        _llm_config: dict[str, Any] = {}
         _llm_config["model"] = self.name
-        for attr, atr_type in [
+        optionals: list[tuple[str, type]] = [
             ("base_url", str),
             ("max_tokens", int),
             # ("temperature", float),
             ("top_p", float),
             ("api_version", str),
             ("default_headers", dict),
-        ]:
+        ]
+        for attr, atr_type in optionals:
             value = getattr(self.data, attr)
             if value and isinstance(value, atr_type):
                 _llm_config[attr] = value
@@ -221,33 +223,86 @@ class WaldiezModel(WaldiezBase):
             value = getattr(self, attr)
             if value:
                 _llm_config[attr] = value
+        if self.data.api_type == "bedrock":
+            _llm_config.pop("base_url", None)
+            return set_bedrock_aws_config(_llm_config, self.data.aws)
         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]:
+    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]
+    llm_config : dict[str, Any]
         The llm config dictionary.
     api_type : str
         The api type.
 
     Returns
     -------
-    Dict[str, Any]
+    dict[str, Any]
         The llm config dictionary with the default base url set.
     """
     dict_copy = llm_config.copy()
     if "base_url" not in llm_config or not llm_config["base_url"]:
-        if MODEL_NEEDS_BASE_URL.get(api_type, True):
+        if MODEL_NEEDS_BASE_URL.get(api_type, True):  # pragma: no branch
             dict_copy["base_url"] = DEFAULT_BASE_URLS.get(api_type, "")
     if (
         not llm_config.get("base_url", "")
         and MODEL_NEEDS_BASE_URL.get(api_type, True) is False
-    ):
+    ):  # pragma: no cover
         dict_copy.pop("base_url", None)
     return dict_copy
+
+
+def set_bedrock_aws_config(
+    llm_config: dict[str, Any],
+    aws_config: Optional[WaldiezModelAWS],
+) -> dict[str, Any]:
+    """Set the AWS config for Bedrock.
+
+    Parameters
+    ----------
+    llm_config : dict[str, Any]
+        The llm config dictionary.
+    aws_config : Optional[WaldiezModelAWS]
+        The passed aws config if any.
+
+    Returns
+    -------
+    dict[str, Any]
+        The llm config dictionary with the AWS config set.
+    """
+    dict_copy = llm_config.copy()
+    aws_params = [
+        "access_key",
+        "secret_key",
+        "session_token",
+        "profile_name",
+        "region",
+    ]
+
+    extra_args: dict[str, Any] = {}
+    for param in aws_params:
+        config_key = f"aws_{param}"
+        env_var = f"AWS_{param.upper()}"
+
+        # First try to get from aws_config
+        value = getattr(aws_config, param, "") if aws_config else ""
+
+        # If not found, try environment variable
+        if not value:  # pragma: no cover
+            value = os.environ.get(env_var, "")
+
+        # Add to extra_args if value exists
+        if value:  # pragma: no branch
+            extra_args[config_key] = value
+
+    # Update llm_config with extra_args
+    if extra_args:  # pragma: no branch
+        dict_copy.update(extra_args)
+
+    return dict_copy

waldiez/models/model/model_data.py

@@ -3,7 +3,7 @@
 # flake8: noqa: E501
 """Waldiez Model Data."""
 
-from typing import Dict, Optional
+from typing import Optional
 
 from pydantic import Field
 from typing_extensions import Annotated, Literal
@@ -21,11 +21,76 @@ WaldiezModelAPIType = Literal[
     "together",
     "nim",
     "cohere",
+    "bedrock",
     "other",
 ]
 """Possible API types for the model."""
 
 
+class WaldiezModelAWS(WaldiezBase):
+    """AWS related parameters.
+
+    Attributes
+    ----------
+    region : Optional[str]
+        The AWS region, by default None.
+    access_key : Optional[str]
+        The AWS access key, by default None.
+    secret_key : Optional[str]
+        The AWS secret access key, by default None.
+    session_token : Optional[str]
+        The AWS session token, by default None.
+    profile_name : Optional[str]
+        The AWS profile name, by default Nonde.
+    """
+
+    region: Annotated[
+        Optional[str],
+        Field(
+            None,
+            alias="region",
+            title="Region",
+            description="The AWS region",
+        ),
+    ] = None
+    access_key: Annotated[
+        Optional[str],
+        Field(
+            None,
+            alias="accessKey",
+            title="Access Ke",
+            description="The AWS access key",
+        ),
+    ] = None
+    secret_key: Annotated[
+        Optional[str],
+        Field(
+            None,
+            alias="secretKey",
+            title="Secret Key",
+            description="The AWS secret key",
+        ),
+    ] = None
+    session_token: Annotated[
+        Optional[str],
+        Field(
+            None,
+            alias="sessionToken",
+            title="Session Token",
+            description="The AWS session token",
+        ),
+    ] = None
+    profile_name: Annotated[
+        Optional[str],
+        Field(
+            None,
+            alias="profileName",
+            title="Profile Name",
+            description="The AWS Profile name to use",
+        ),
+    ] = None
+
+
 class WaldiezModelPrice(WaldiezBase):
     """Model Price.
 
@@ -65,7 +130,10 @@ class WaldiezModelData(WaldiezBase):
         The top p of the model, by default None.
     max_tokens : Optional[int]
         The max tokens of the model, by default None.
-    default_headers : Dict[str, str]
+    aws : Optional[WaldiezModelAWS]
+    extras: dict[str, str]
+        Any extra attributes to include in the LLM Config.
+    default_headers : dict[str, str]
         The default headers of the model.
     price : Optional[WaldiezModelPrice]
         The price of the model, by default None.
@@ -74,76 +142,96 @@ class WaldiezModelData(WaldiezBase):
     base_url: Annotated[
         Optional[str],
         Field(
-            None,
+            default=None,
             title="Base URL",
             description="The base url of the model",
             alias="baseUrl",
         ),
-    ]
+    ] = None
     api_key: Annotated[
         Optional[str],
         Field(
-            None,
+            default=None,
             alias="apiKey",
             title="API Key",
             description="The api key to use with the model",
         ),
-    ]
+    ] = None
     api_type: Annotated[
         WaldiezModelAPIType,
         Field(
-            "other",
+            default="other",
             alias="apiType",
             title="API Type",
             description="The api type of the model",
         ),
-    ]
+    ] = "other"
     api_version: Annotated[
         Optional[str],
         Field(
-            None,
+            default=None,
             alias="apiVersion",
             title="API Version",
             description="The api version of the model",
         ),
-    ]
+    ] = None
     temperature: Annotated[
         Optional[float],
         Field(
-            None,
+            default=None,
             alias="temperature",
             title="Temperature",
             description="The temperature of the model",
         ),
-    ]
+    ] = None
     top_p: Annotated[
         Optional[float],
         Field(
-            None,
+            default=None,
             alias="topP",
             title="Top P",
             description="The top p of the model",
         ),
-    ]
+    ] = None
     max_tokens: Annotated[
         Optional[int],
         Field(
-            None,
+            default=None,
             alias="maxTokens",
             title="Max Tokens",
             description="The max tokens of the model",
         ),
-    ]
+    ] = None
+    aws: Annotated[
+        Optional[WaldiezModelAWS],
+        Field(
+            default=None,
+            alias="aws",
+            title="AWS",
+            description="The AWS related parameters",
+        ),
+    ] = None
+    extras: Annotated[
+        dict[str, str],
+        Field(
+            alias="extras",
+            default_factory=dict,
+            title="Extras",
+            description="Any extra attributes to include in the LLM Config",
+        ),
+    ] = {}
     default_headers: Annotated[
-        Dict[str, str],
+        dict[str, str],
         Field(
             alias="defaultHeaders",
             default_factory=dict,
             title="Default Headers",
             description="The default headers of the model",
         ),
-    ]
+    ] = {}
     price: Annotated[
         Optional[WaldiezModelPrice],
-        Field(None, title="Price", description="The price of the model"),
-    ]
+        Field(
+            default=None, title="Price", description="The price of the model"
+        ),
+    ] = None

waldiez/models/tool/__init__.py

@@ -0,0 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Waldiez Tool related models."""
+
+from .extra_requirements import get_tools_extra_requirements
+from .tool import SHARED_TOOL_NAME, WaldiezTool
+from .tool_data import WaldiezToolData
+from .tool_type import WaldiezToolType
+
+__all__ = [
+    "SHARED_TOOL_NAME",
+    "WaldiezTool",
+    "WaldiezToolData",
+    "WaldiezToolType",
+    "get_tools_extra_requirements",
+]

waldiez/models/tool/extra_requirements.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Waldiez tool extra requirements."""
+
+from typing import Iterator, Set
+
+from .tool import WaldiezTool
+
+
+def get_tools_extra_requirements(
+    tools: Iterator[WaldiezTool],
+    autogen_version: str,
+) -> Set[str]:
+    """Get the tools extra requirements.
+
+    Parameters
+    ----------
+    tools : list[WaldiezTool]
+        The tools.
+    autogen_version : str
+        The ag2 version.
+
+    Returns
+    -------
+    list[str]
+        The tools extra requirements.
+    """
+    tool_requirements: Set[str] = set()
+    for tool in tools:
+        if tool.tool_type == "langchain":
+            tool_requirements.add(f"ag2[interop-langchain]=={autogen_version}")
+        if tool.tool_type == "crewai":
+            tool_requirements.add(f"ag2[interop-crewai]=={autogen_version}")
+        for requirement in tool.requirements:
+            tool_requirements.add(requirement)
+    return tool_requirements