ag2 0.9.6__py3-none-any.whl → 0.9.8.post1__py3-none-any.whl

autogen/messages/client_messages.py

@@ -2,7 +2,8 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Callable, Literal, Optional, Union
+from collections.abc import Callable
+from typing import Any, Literal
 from uuid import UUID
 
 from pydantic import BaseModel
@@ -32,18 +33,18 @@ class ModelUsageSummary(BaseModel):
 class ActualUsageSummary(BaseModel):
     """Actual usage summary."""
 
-    usages: Optional[list[ModelUsageSummary]] = None
+    usages: list[ModelUsageSummary] | None = None
     """List of model usage summaries."""
-    total_cost: Optional[float] = None
+    total_cost: float | None = None
     """Total cost."""
 
 
 class TotalUsageSummary(BaseModel):
     """Total usage summary."""
 
-    usages: Optional[list[ModelUsageSummary]] = None
+    usages: list[ModelUsageSummary] | None = None
     """List of model usage summaries."""
-    total_cost: Optional[float] = None
+    total_cost: float | None = None
     """Total cost."""
 
 
@@ -51,7 +52,7 @@ Mode = Literal["both", "total", "actual"]
 
 
 def _change_usage_summary_format(
-    actual_usage_summary: Optional[dict[str, Any]] = None, total_usage_summary: Optional[dict[str, Any]] = None
+    actual_usage_summary: dict[str, Any] | None = None, total_usage_summary: dict[str, Any] | None = None
 ) -> dict[str, dict[str, Any]]:
     summary: dict[str, Any] = {}
 
@@ -88,9 +89,9 @@ class UsageSummaryMessage(BaseMessage):
     def __init__(
         self,
         *,
-        uuid: Optional[UUID] = None,
-        actual_usage_summary: Optional[dict[str, Any]] = None,
-        total_usage_summary: Optional[dict[str, Any]] = None,
+        uuid: UUID | None = None,
+        actual_usage_summary: dict[str, Any] | None = None,
+        total_usage_summary: dict[str, Any] | None = None,
         mode: Mode = "both",
     ):
         # print(f"{actual_usage_summary=}")
@@ -102,9 +103,9 @@ class UsageSummaryMessage(BaseMessage):
 
     def _print_usage(
         self,
-        usage_summary: Union[ActualUsageSummary, TotalUsageSummary],
+        usage_summary: ActualUsageSummary | TotalUsageSummary,
         usage_type: str = "total",
-        f: Optional[Callable[..., Any]] = None,
+        f: Callable[..., Any] | None = None,
     ) -> None:
         f = f or print
         word_from_type = "including" if usage_type == "total" else "excluding"
@@ -121,7 +122,7 @@ class UsageSummaryMessage(BaseMessage):
                 flush=True,
             )
 
-    def print(self, f: Optional[Callable[..., Any]] = None) -> None:
+    def print(self, f: Callable[..., Any] | None = None) -> None:
         f = f or print
 
         if self.total.usages is None:
@@ -156,10 +157,10 @@ class StreamMessage(BaseMessage):
     content: str
     """Content of the message."""
 
-    def __init__(self, *, uuid: Optional[UUID] = None, content: str) -> None:
+    def __init__(self, *, uuid: UUID | None = None, content: str) -> None:
         super().__init__(uuid=uuid, content=content)
 
-    def print(self, f: Optional[Callable[..., Any]] = None) -> None:
+    def print(self, f: Callable[..., Any] | None = None) -> None:
         f = f or print
 
         # Set the terminal text color to green

autogen/messages/print_message.py

@@ -4,7 +4,8 @@
 
 
 import json
-from typing import Any, Callable, Optional
+from collections.abc import Callable
+from typing import Any
 from uuid import UUID
 
 from ..events import deprecated_by
@@ -24,9 +25,7 @@ class PrintMessage(BaseMessage):
     end: str
     """End of the print"""
 
-    def __init__(
-        self, *objects: Any, sep: str = " ", end: str = "\n", flush: bool = False, uuid: Optional[UUID] = None
-    ):
+    def __init__(self, *objects: Any, sep: str = " ", end: str = "\n", flush: bool = False, uuid: UUID | None = None):
         objects_as_string = [self._to_json(x) for x in objects]
 
         super().__init__(uuid=uuid, objects=objects_as_string, sep=sep, end=end)
@@ -43,7 +42,7 @@ class PrintMessage(BaseMessage):
             return str(obj)
             # return repr(obj)
 
-    def print(self, f: Optional[Callable[..., Any]] = None) -> None:
+    def print(self, f: Callable[..., Any] | None = None) -> None:
         f = f or print
 
         f(*self.objects, sep=self.sep, end=self.end, flush=True)

autogen/oai/__init__.py

@@ -8,7 +8,7 @@ from ..cache.cache import Cache
 from .anthropic import AnthropicLLMConfigEntry
 from .bedrock import BedrockLLMConfigEntry
 from .cerebras import CerebrasLLMConfigEntry
-from .client import AzureOpenAILLMConfigEntry, DeepSeekLLMConfigEntry, ModelClient, OpenAILLMConfigEntry, OpenAIWrapper
+from .client import AzureOpenAILLMConfigEntry, DeepSeekLLMConfigEntry, OpenAILLMConfigEntry, OpenAIWrapper
 from .cohere import CohereLLMConfigEntry
 from .gemini import GeminiLLMConfigEntry
 from .groq import GroqLLMConfigEntry
@@ -37,7 +37,6 @@ __all__ = [
     "GeminiLLMConfigEntry",
     "GroqLLMConfigEntry",
     "MistralLLMConfigEntry",
-    "ModelClient",
     "OllamaLLMConfigEntry",
     "OpenAILLMConfigEntry",
     "OpenAIWrapper",

autogen/oai/anthropic.py

@@ -76,12 +76,13 @@ import os
 import re
 import time
 import warnings
-from typing import Any, Literal, Optional, Union
+from typing import Any, Literal
 
 from pydantic import BaseModel, Field
+from typing_extensions import Unpack
 
 from ..import_utils import optional_import_block, require_optional_import
-from ..llm_config import LLMConfigEntry, register_llm_config
+from ..llm_config.entry import LLMConfigEntry, LLMConfigEntryDict
 from .client_utils import FormatterProtocol, validate_parameter
 from .oai_models import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, Choice, CompletionUsage
 
@@ -109,23 +110,39 @@ ANTHROPIC_PRICING_1k = {
 }
 
 
-@register_llm_config
+class AnthropicEntryDict(LLMConfigEntryDict, total=False):
+    api_type: Literal["anthropic"]
+    timeout: int | None
+    stop_sequences: list[str] | None
+    stream: bool
+    price: list[float] | None
+    tool_choice: dict | None
+    thinking: dict | None
+    gcp_project_id: str | None
+    gcp_region: str | None
+    gcp_auth_token: str | None
+
+
 class AnthropicLLMConfigEntry(LLMConfigEntry):
     api_type: Literal["anthropic"] = "anthropic"
-    timeout: Optional[int] = Field(default=None, ge=1)
-    temperature: float = Field(default=1.0, ge=0.0, le=1.0)
-    top_k: Optional[int] = Field(default=None, ge=1)
-    top_p: Optional[float] = Field(default=None, ge=0.0, le=1.0)
-    stop_sequences: Optional[list[str]] = None
-    stream: bool = False
+
+    # Basic options
     max_tokens: int = Field(default=4096, ge=1)
-    price: Optional[list[float]] = Field(default=None, min_length=2, max_length=2)
-    tool_choice: Optional[dict] = None
-    thinking: Optional[dict] = None
+    temperature: float | None = Field(default=None, ge=0.0, le=1.0)
+    top_p: float | None = Field(default=None, ge=0.0, le=1.0)
+
+    # Anthropic-specific options
+    timeout: int | None = Field(default=None, ge=1)
+    top_k: int | None = Field(default=None, ge=1)
+    stop_sequences: list[str] | None = None
+    stream: bool = False
+    price: list[float] | None = Field(default=None, min_length=2, max_length=2)
+    tool_choice: dict | None = None
+    thinking: dict | None = None
 
-    gcp_project_id: Optional[str] = None
-    gcp_region: Optional[str] = None
-    gcp_auth_token: Optional[str] = None
+    gcp_project_id: str | None = None
+    gcp_region: str | None = None
+    gcp_auth_token: str | None = None
 
     def create_client(self):
         raise NotImplementedError("AnthropicLLMConfigEntry.create_client is not implemented.")
@@ -133,37 +150,22 @@ class AnthropicLLMConfigEntry(LLMConfigEntry):
 
 @require_optional_import("anthropic", "anthropic")
 class AnthropicClient:
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Unpack[AnthropicEntryDict]):
         """Initialize the Anthropic API client.
 
         Args:
             **kwargs: The configuration parameters for the client.
         """
-        self._api_key = kwargs.get("api_key")
-        self._aws_access_key = kwargs.get("aws_access_key")
-        self._aws_secret_key = kwargs.get("aws_secret_key")
+        self._api_key = kwargs.get("api_key") or os.getenv("ANTHROPIC_API_KEY")
+        self._aws_access_key = kwargs.get("aws_access_key") or os.getenv("AWS_ACCESS_KEY")
+        self._aws_secret_key = kwargs.get("aws_secret_key") or os.getenv("AWS_SECRET_KEY")
         self._aws_session_token = kwargs.get("aws_session_token")
-        self._aws_region = kwargs.get("aws_region")
+        self._aws_region = kwargs.get("aws_region") or os.getenv("AWS_REGION")
         self._gcp_project_id = kwargs.get("gcp_project_id")
-        self._gcp_region = kwargs.get("gcp_region")
+        self._gcp_region = kwargs.get("gcp_region") or os.getenv("GCP_REGION")
         self._gcp_auth_token = kwargs.get("gcp_auth_token")
         self._base_url = kwargs.get("base_url")
 
-        if not self._api_key:
-            self._api_key = os.getenv("ANTHROPIC_API_KEY")
-
-        if not self._aws_access_key:
-            self._aws_access_key = os.getenv("AWS_ACCESS_KEY")
-
-        if not self._aws_secret_key:
-            self._aws_secret_key = os.getenv("AWS_SECRET_KEY")
-
-        if not self._aws_region:
-            self._aws_region = os.getenv("AWS_REGION")
-
-        if not self._gcp_region:
-            self._gcp_region = os.getenv("GCP_REGION")
-
         if self._api_key is None:
             if self._aws_region:
                 if self._aws_access_key is None or self._aws_secret_key is None:
@@ -181,7 +183,7 @@ class AnthropicClient:
             self._client = Anthropic(**client_kwargs)
         elif self._gcp_region is not None:
             kw = {}
-            for i, p in enumerate(inspect.signature(AnthropicVertex).parameters):
+            for p in inspect.signature(AnthropicVertex).parameters:
                 if hasattr(self, f"_gcp_{p}"):
                     kw[p] = getattr(self, f"_gcp_{p}")
             if self._base_url:
@@ -201,7 +203,7 @@ class AnthropicClient:
         self._last_tooluse_status = {}
 
         # Store the response format, if provided (for structured outputs)
-        self._response_format: Optional[type[BaseModel]] = None
+        self._response_format: type[BaseModel] | None = None
 
     def load_config(self, params: dict[str, Any]):
         """Load the configuration for the Anthropic API client."""
@@ -406,8 +408,7 @@ class AnthropicClient:
 
     @staticmethod
     def convert_tools_to_functions(tools: list) -> list:
-        """
-        Convert tool definitions into Anthropic-compatible functions,
+        """Convert tool definitions into Anthropic-compatible functions,
         updating nested $ref paths in property schemas.
 
         Args:
@@ -550,7 +551,7 @@ def process_image_content(content_item: dict[str, Any]) -> dict[str, Any]:
         return content_item
 
 
-def process_message_content(message: dict[str, Any]) -> Union[str, list[dict[str, Any]]]:
+def process_message_content(message: dict[str, Any]) -> str | list[dict[str, Any]]:
     """Process message content, handling both string and list formats with images."""
     content = message.get("content", "")
 

autogen/oai/bedrock.py

@@ -36,13 +36,14 @@ import os
 import re
 import time
 import warnings
-from typing import Any, Literal, Optional
+from typing import Any, Literal
 
 import requests
 from pydantic import Field, SecretStr, field_serializer
+from typing_extensions import Required, Unpack
 
 from ..import_utils import optional_import_block, require_optional_import
-from ..llm_config import LLMConfigEntry, register_llm_config
+from ..llm_config.entry import LLMConfigEntry, LLMConfigEntryDict
 from .client_utils import validate_parameter
 from .oai_models import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, Choice, CompletionUsage
 
@@ -51,26 +52,38 @@ with optional_import_block():
     from botocore.config import Config
 
 
-@register_llm_config
+class BedrockEntryDict(LLMConfigEntryDict, total=False):
+    api_type: Literal["bedrock"]
+    aws_region: Required[str]
+    aws_access_key: SecretStr | None
+    aws_secret_key: SecretStr | None
+    aws_session_token: SecretStr | None
+    aws_profile_name: str | None
+    top_k: int | None
+    k: int | None
+    seed: int | None
+    cache_seed: int | None
+    supports_system_prompts: bool
+    price: list[float] | None
+    timeout: int | None
+
+
 class BedrockLLMConfigEntry(LLMConfigEntry):
     api_type: Literal["bedrock"] = "bedrock"
+
+    # Bedrock-specific options
     aws_region: str
-    aws_access_key: Optional[SecretStr] = None
-    aws_secret_key: Optional[SecretStr] = None
-    aws_session_token: Optional[SecretStr] = None
-    aws_profile_name: Optional[str] = None
-    temperature: Optional[float] = None
-    topP: Optional[float] = None  # noqa: N815
-    maxTokens: Optional[int] = None  # noqa: N815
-    top_p: Optional[float] = None
-    top_k: Optional[int] = None
-    k: Optional[int] = None
-    seed: Optional[int] = None
-    cache_seed: Optional[int] = None
+    aws_access_key: SecretStr | None = None
+    aws_secret_key: SecretStr | None = None
+    aws_session_token: SecretStr | None = None
+    aws_profile_name: str | None = None
+    top_k: int | None = None
+    k: int | None = None
+    seed: int | None = None
+    cache_seed: int | None = None
     supports_system_prompts: bool = True
-    stream: bool = False
-    price: Optional[list[float]] = Field(default=None, min_length=2, max_length=2)
-    timeout: Optional[int] = None
+    price: list[float] | None = Field(default=None, min_length=2, max_length=2)
+    timeout: int | None = None
 
     @field_serializer("aws_access_key", "aws_secret_key", "aws_session_token", when_used="unless-none")
     def serialize_aws_secrets(self, v: SecretStr) -> str:
@@ -86,27 +99,15 @@ class BedrockClient:
 
     _retries = 5
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Unpack[BedrockEntryDict]):
         """Initialises BedrockClient for Amazon's Bedrock Converse API"""
-        self._aws_access_key = kwargs.get("aws_access_key")
-        self._aws_secret_key = kwargs.get("aws_secret_key")
-        self._aws_session_token = kwargs.get("aws_session_token")
-        self._aws_region = kwargs.get("aws_region")
+        self._aws_access_key = kwargs.get("aws_access_key") or os.getenv("AWS_ACCESS_KEY")
+        self._aws_secret_key = kwargs.get("aws_secret_key") or os.getenv("AWS_SECRET_KEY")
+        self._aws_session_token = kwargs.get("aws_session_token") or os.getenv("AWS_SESSION_TOKEN")
+        self._aws_region = kwargs.get("aws_region") or os.getenv("AWS_REGION")
         self._aws_profile_name = kwargs.get("aws_profile_name")
         self._timeout = kwargs.get("timeout")
 
-        if not self._aws_access_key:
-            self._aws_access_key = os.getenv("AWS_ACCESS_KEY")
-
-        if not self._aws_secret_key:
-            self._aws_secret_key = os.getenv("AWS_SECRET_KEY")
-
-        if not self._aws_session_token:
-            self._aws_session_token = os.getenv("AWS_SESSION_TOKEN")
-
-        if not self._aws_region:
-            self._aws_region = os.getenv("AWS_REGION")
-
         if self._aws_region is None:
             raise ValueError("Region is required to use the Amazon Bedrock API.")
 
@@ -159,13 +160,10 @@ class BedrockClient:
         # This is required because not all models support a system prompt (e.g. Mistral Instruct).
         self._supports_system_prompts = params.get("supports_system_prompts", True)
 
-    def parse_params(self, params: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
+    def parse_params(self, params: BedrockEntryDict | dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
         """Loads the valid parameters required to invoke Bedrock Converse
         Returns a tuple of (base_params, additional_params)
         """
-        base_params = {}
-        additional_params = {}
-
         # Amazon Bedrock  base model IDs are here:
         # https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html
         self._model_id = params.get("model")
@@ -190,38 +188,51 @@ class BedrockClient:
         # https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral-chat-completion.html
 
         # Here are the possible "base" parameters and their suitable types
-        base_parameters = [["temperature", (float, int)], ["topP", (float, int)], ["maxTokens", (int)]]
+        base_params = {}
 
-        for param_name, suitable_types in base_parameters:
-            if param_name in params:
-                base_params[param_name] = validate_parameter(
-                    params, param_name, suitable_types, False, None, None, None
-                )
+        if "temperature" in params:
+            base_params["temperature"] = validate_parameter(
+                params, "temperature", (float, int), False, None, None, None
+            )
+
+        if "top_p" in params:
+            base_params["topP"] = validate_parameter(params, "top_p", (float, int), False, None, None, None)
+
+        if "topP" in params:
+            warnings.warn(
+                ("topP is deprecated, use top_p instead. Scheduled for removal in 0.10.0 version."), DeprecationWarning
+            )
+            base_params["topP"] = validate_parameter(params, "topP", (float, int), False, None, None, None)
+
+        if "max_tokens" in params:
+            base_params["maxTokens"] = validate_parameter(params, "max_tokens", (int,), False, None, None, None)
+
+        if "maxTokens" in params:
+            warnings.warn(
+                ("maxTokens is deprecated, use max_tokens instead. Scheduled for removal in 0.10.0 version."),
+                DeprecationWarning,
+            )
+            base_params["maxTokens"] = validate_parameter(params, "maxTokens", (int,), False, None, None, None)
 
         # Here are the possible "model-specific" parameters and their suitable types, known as additional parameters
-        additional_parameters = [
-            ["top_p", (float, int)],
-            ["top_k", (int)],
-            ["k", (int)],
-            ["seed", (int)],
-        ]
+        additional_params = {}
 
-        for param_name, suitable_types in additional_parameters:
+        for param_name, suitable_types in (
+            ("top_k", (int,)),
+            ("k", (int,)),
+            ("seed", (int,)),
+        ):
             if param_name in params:
                 additional_params[param_name] = validate_parameter(
                     params, param_name, suitable_types, False, None, None, None
                 )
 
-        # Streaming
-        self._streaming = params.get("stream", False)
-
         # For this release we will not support streaming as many models do not support streaming with tool use
-        if self._streaming:
+        if params.get("stream", False):
             warnings.warn(
                 "Streaming is not currently supported, streaming will be disabled.",
                 UserWarning,
             )
-            self._streaming = False
 
         return base_params, additional_params
 

autogen/oai/cerebras.py

@@ -28,12 +28,13 @@ import math
 import os
 import time
 import warnings
-from typing import Any, Literal, Optional
+from typing import Any, Literal
 
-from pydantic import Field, ValidationInfo, field_validator
+from pydantic import Field
+from typing_extensions import Unpack
 
 from ..import_utils import optional_import_block, require_optional_import
-from ..llm_config import LLMConfigEntry, register_llm_config
+from ..llm_config.entry import LLMConfigEntry, LLMConfigEntryDict
 from .client_utils import should_hide_tools, validate_parameter
 from .oai_models import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, Choice, CompletionUsage
 
@@ -47,23 +48,25 @@ CEREBRAS_PRICING_1K = {
 }
 
 
-@register_llm_config
+class CerebrasEntryDict(LLMConfigEntryDict, total=False):
+    api_type: Literal["cerebras"]
+
+    seed: int | None
+    stream: bool
+    hide_tools: Literal["if_all_run", "if_any_run", "never"]
+    tool_choice: Literal["none", "auto", "required"] | None
+
+
 class CerebrasLLMConfigEntry(LLMConfigEntry):
     api_type: Literal["cerebras"] = "cerebras"
-    max_tokens: Optional[int] = None
-    seed: Optional[int] = None
+
+    temperature: float | None = Field(default=None, ge=0.0, le=1.5)
+
+    seed: int | None = None
     stream: bool = False
-    temperature: float = Field(default=1.0, ge=0.0, le=1.5)
-    top_p: Optional[float] = None
     hide_tools: Literal["if_all_run", "if_any_run", "never"] = "never"
-    tool_choice: Optional[Literal["none", "auto", "required"]] = None
-
-    @field_validator("top_p", mode="before")
-    @classmethod
-    def check_top_p(cls, v: Any, info: ValidationInfo) -> Any:
-        if v is not None and info.data.get("temperature") is not None:
-            raise ValueError("temperature and top_p cannot be set at the same time.")
-        return v
+    tool_choice: Literal["none", "auto", "required"] | None = None
+    reasoning_effort: str | None = None
 
     def create_client(self):
         raise NotImplementedError("CerebrasLLMConfigEntry.create_client is not implemented.")
@@ -72,7 +75,7 @@ class CerebrasLLMConfigEntry(LLMConfigEntry):
 class CerebrasClient:
     """Client for Cerebras's API."""
 
-    def __init__(self, api_key=None, **kwargs):
+    def __init__(self, api_key=None, **kwargs: Unpack[CerebrasEntryDict]):
         """Requires api_key or environment variable to be set
 
         Args:
@@ -80,16 +83,14 @@ class CerebrasClient:
             **kwargs: Additional keyword arguments to pass to the Cerebras client
         """
         # Ensure we have the api_key upon instantiation
-        self.api_key = api_key
-        if not self.api_key:
-            self.api_key = os.getenv("CEREBRAS_API_KEY")
+        self.api_key = api_key or os.getenv("CEREBRAS_API_KEY")
 
         assert self.api_key, (
             "Please include the api_key in your config list entry for Cerebras or set the CEREBRAS_API_KEY env variable."
         )
 
         if "response_format" in kwargs and kwargs["response_format"] is not None:
-            warnings.warn("response_format is not supported for Crebras, it will be ignored.", UserWarning)
+            warnings.warn("response_format is not supported for Cerebras, it will be ignored.", UserWarning)
 
     def message_retrieval(self, response: ChatCompletion) -> list:
         """Retrieve and return a list of strings or a list of Choice.Message from the response.
@@ -129,12 +130,12 @@ class CerebrasClient:
         # Validate allowed Cerebras parameters
         # https://inference-docs.cerebras.ai/api-reference/chat-completions
         cerebras_params["max_tokens"] = validate_parameter(params, "max_tokens", int, True, None, (0, None), None)
-        cerebras_params["seed"] = validate_parameter(params, "seed", int, True, None, None, None)
-        cerebras_params["stream"] = validate_parameter(params, "stream", bool, True, False, None, None)
         cerebras_params["temperature"] = validate_parameter(
-            params, "temperature", (int, float), True, 1, (0, 1.5), None
+            params, "temperature", (int, float), True, 1.0, (0, 1.5), None
         )
-        cerebras_params["top_p"] = validate_parameter(params, "top_p", (int, float), True, None, None, None)
+        cerebras_params["top_p"] = validate_parameter(params, "top_p", (int, float), True, None, (0.0, 1.0), None)
+        cerebras_params["seed"] = validate_parameter(params, "seed", int, True, None, None, None)
+        cerebras_params["stream"] = validate_parameter(params, "stream", bool, True, False, None, None)
         cerebras_params["tool_choice"] = validate_parameter(
             params, "tool_choice", str, True, None, None, ["none", "auto", "required"]
         )