Merge branch 'main' into update-generate-reply-function

Author: priyansh4320

autogen/agentchat/group/safeguards/events.py

@@ -5,13 +5,34 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from typing import Any
+from typing import Any, Literal
 from uuid import UUID
 
 from termcolor import colored
 
 from ....events.base_event import BaseEvent, wrap_event
 
+# Type for termcolor colors
+TermColor = Literal[
+    "black",
+    "grey",
+    "red",
+    "green",
+    "yellow",
+    "blue",
+    "magenta",
+    "cyan",
+    "light_grey",
+    "dark_grey",
+    "light_red",
+    "light_green",
+    "light_yellow",
+    "light_blue",
+    "light_magenta",
+    "light_cyan",
+    "white",
+]
+
 
 @wrap_event
 class SafeguardEvent(BaseEvent):
@@ -52,7 +73,7 @@ def print(self, f: Callable[..., Any] | None = None) -> None:
         f = f or print
 
         # Choose color based on event type
-        color = "green"
+        color: TermColor = "green"
         if self.event_type == "load":
             color = "green"
         elif self.event_type == "check":

autogen/interop/pydantic_ai/pydantic_ai.py

@@ -63,7 +63,7 @@ def inject_params(
 
         @wraps(f)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
-            current_retry = 0 if ctx_typed is None else ctx_typed.retries.get(tool_typed.name, 0)
+            current_retry = 0 if ctx_typed is None else ctx_typed.retries.get(tool_typed.name, 0)  # type: ignore[attr-defined]
 
             if current_retry >= max_retries:
                 raise ValueError(f"{tool_typed.name} failed after {max_retries} retries")
@@ -73,12 +73,12 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
                     kwargs.pop("ctx", None)
                     ctx_typed.retry = current_retry
                     result = f(**kwargs, ctx=ctx_typed)  # type: ignore[call-arg]
-                    ctx_typed.retries[tool_typed.name] = 0
+                    ctx_typed.retries[tool_typed.name] = 0  # type: ignore[attr-defined]
                 else:
                     result = f(**kwargs)  # type: ignore[call-arg]
             except Exception as e:
                 if ctx_typed is not None:
-                    ctx_typed.retries[tool_typed.name] = ctx_typed.retries.get(tool_typed.name, 0) + 1
+                    ctx_typed.retries[tool_typed.name] = ctx_typed.retries.get(tool_typed.name, 0) + 1  # type: ignore[attr-defined]
                 raise e
 
             return result
@@ -145,6 +145,10 @@ def convert_tool(cls, tool: Any, deps: Any = None, **kwargs: Any) -> Tool:
             else None
         )
 
+        # Initialize retries dict for tracking retry counts per tool
+        if ctx is not None:
+            ctx.retries = {}  # type: ignore[attr-defined]
+
         func = PydanticAIInteroperability.inject_params(
             ctx=ctx,
             tool=pydantic_ai_tool,
@@ -154,7 +158,7 @@ def convert_tool(cls, tool: Any, deps: Any = None, **kwargs: Any) -> Tool:
             name=pydantic_ai_tool.name,
             description=pydantic_ai_tool.description,
             func_or_tool=func,
-            parameters_json_schema=pydantic_ai_tool.function_schema.json_schema,
+            parameters_json_schema=pydantic_ai_tool.function_schema.json_schema,  # type: ignore[attr-defined]
         )
 
     @classmethod

autogen/llm_config/client.py

@@ -32,7 +32,7 @@ class ModelClient(Protocol):
     class ModelClientResponseProtocol(Protocol):
         class Choice(Protocol):
             class Message(Protocol):
-                content: str | dict[str, Any] | list[dict[str, Any]]
+                content: str | dict[str, Any] | list[dict[str, Any]] | None
 
             message: Message
 

autogen/oai/anthropic.py

@@ -150,6 +150,8 @@ def create_client(self):
 
 @require_optional_import("anthropic", "anthropic")
 class AnthropicClient:
+    RESPONSE_USAGE_KEYS: list[str] = ["prompt_tokens", "completion_tokens", "total_tokens", "cost", "model"]
+
     def __init__(self, **kwargs: Unpack[AnthropicEntryDict]):
         """Initialize the Anthropic API client.
 

autogen/oai/bedrock.py

@@ -97,6 +97,8 @@ def create_client(self):
 class BedrockClient:
     """Client for Amazon's Bedrock Converse API."""
 
+    RESPONSE_USAGE_KEYS: list[str] = ["prompt_tokens", "completion_tokens", "total_tokens", "cost", "model"]
+
     _retries = 5
 
     def __init__(self, **kwargs: Unpack[BedrockEntryDict]):

autogen/oai/cerebras.py

@@ -75,6 +75,8 @@ def create_client(self):
 class CerebrasClient:
     """Client for Cerebras's API."""
 
+    RESPONSE_USAGE_KEYS: list[str] = ["prompt_tokens", "completion_tokens", "total_tokens", "cost", "model"]
+
     def __init__(self, api_key=None, **kwargs: Unpack[CerebrasEntryDict]):
         """Requires api_key or environment variable to be set
 

autogen/oai/client.py

@@ -19,6 +19,8 @@
 from pydantic import BaseModel, Field, HttpUrl
 from pydantic.type_adapter import TypeAdapter
 
+from autogen.oai.oai_models.chat_completion import ChatCompletionExtended
+
 from ..cache import Cache
 from ..code_utils import content_str
 from ..doc_utils import export_module
@@ -57,10 +59,10 @@
 
     if openai.__version__ >= "1.1.0":
         TOOL_ENABLED = True
-    ERROR = None
+    ERROR: ImportError | None = None
     from openai.lib._pydantic import _ensure_strict_json_schema
 else:
-    ERROR: ImportError | None = ImportError("Please install openai>=1 and diskcache to use autogen.OpenAIWrapper.")
+    ERROR = ImportError("Please install openai>=1 and diskcache to use autogen.OpenAIWrapper.")  # type: ignore[assignment]
 
     # OpenAI = object
     # AzureOpenAI = object
@@ -77,7 +79,7 @@
 if cerebras_result.is_successful:
     cerebras_import_exception: ImportError | None = None
 else:
-    cerebras_AuthenticationError = cerebras_InternalServerError = cerebras_RateLimitError = Exception  # noqa: N816
+    cerebras_AuthenticationError = cerebras_InternalServerError = cerebras_RateLimitError = Exception  # type: ignore[assignment,misc]  # noqa: N816
     cerebras_import_exception = ImportError("cerebras_cloud_sdk not found")
 
 with optional_import_block() as gemini_result:
@@ -91,7 +93,7 @@
 if gemini_result.is_successful:
     gemini_import_exception: ImportError | None = None
 else:
-    gemini_InternalServerError = gemini_ResourceExhausted = Exception  # noqa: N816
+    gemini_InternalServerError = gemini_ResourceExhausted = Exception  # type: ignore[assignment,misc]  # noqa: N816
     gemini_import_exception = ImportError("google-genai not found")
 
 with optional_import_block() as anthropic_result:
@@ -105,7 +107,7 @@
 if anthropic_result.is_successful:
     anthropic_import_exception: ImportError | None = None
 else:
-    anthorpic_InternalServerError = anthorpic_RateLimitError = Exception  # noqa: N816
+    anthorpic_InternalServerError = anthorpic_RateLimitError = Exception  # type: ignore[assignment,misc]  # noqa: N816
     anthropic_import_exception = ImportError("anthropic not found")
 
 with optional_import_block() as mistral_result:
@@ -174,7 +176,7 @@
 if ollama_result.is_successful:
     ollama_import_exception: ImportError | None = None
 else:
-    ollama_RequestError = ollama_ResponseError = Exception  # noqa: N816
+    ollama_RequestError = ollama_ResponseError = Exception  # type: ignore[assignment,misc]  # noqa: N816
     ollama_import_exception = ImportError("ollama not found")
 
 with optional_import_block() as bedrock_result:
@@ -340,6 +342,8 @@ def __init__(self, config):
 class OpenAIClient:
     """Follows the Client protocol and wraps the OpenAI client."""
 
+    RESPONSE_USAGE_KEYS: list[str] = ["prompt_tokens", "completion_tokens", "total_tokens", "cost", "model"]
+
     def __init__(self, client: OpenAI | AzureOpenAI, response_format: BaseModel | dict[str, Any] | None = None):
         self._oai_client = client
         self.response_format = response_format
@@ -712,7 +716,7 @@ def cost(self, response: ChatCompletion | Completion) -> float:
         return tmp_price1K * (n_input_tokens + n_output_tokens) / 1000  # type: ignore [operator]
 
     @staticmethod
-    def get_usage(response: ChatCompletion | Completion) -> dict:
+    def get_usage(response: ChatCompletion | Completion) -> dict[str, Any]:
         return {
             "prompt_tokens": response.usage.prompt_tokens if response.usage is not None else 0,
             "completion_tokens": response.usage.completion_tokens if response.usage is not None else 0,
@@ -900,21 +904,21 @@ def _register_default_client(self, config: dict[str, Any], openai_config: dict[s
                 def create_azure_openai_client() -> AzureOpenAI:
                     self._configure_azure_openai(config, openai_config)
                     client = AzureOpenAI(**openai_config)
-                    self._clients.append(OpenAIClient(client, response_format=response_format))
+                    self._clients.append(OpenAIClient(client, response_format=response_format))  # type: ignore[arg-type]
                     return client
 
                 client = create_azure_openai_client()
             elif api_type is not None and api_type.startswith("cerebras"):
                 if cerebras_import_exception:
                     raise ImportError("Please install `cerebras_cloud_sdk` to use Cerebras OpenAI API.")
                 client = CerebrasClient(response_format=response_format, **openai_config)
-                self._clients.append(client)
+                self._clients.append(client)  # type: ignore[arg-type]
             elif api_type is not None and api_type.startswith("google"):
                 if gemini_import_exception:
                     raise ImportError("Please install `google-genai` and 'vertexai' to use Google's API.")
                 self._configure_openai_config_for_gemini(config, openai_config)
                 client = GeminiClient(response_format=response_format, **openai_config)
-                self._clients.append(client)
+                self._clients.append(client)  # type: ignore[arg-type]
             elif api_type is not None and api_type.startswith("anthropic"):
                 if "api_key" not in config and "aws_region" in config:
                     self._configure_openai_config_for_bedrock(config, openai_config)
@@ -923,44 +927,44 @@ def create_azure_openai_client() -> AzureOpenAI:
                 if anthropic_import_exception:
                     raise ImportError("Please install `anthropic` to use Anthropic API.")
                 client = AnthropicClient(response_format=response_format, **openai_config)
-                self._clients.append(client)
+                self._clients.append(client)  # type: ignore[arg-type]
             elif api_type is not None and api_type.startswith("mistral"):
                 if mistral_import_exception:
                     raise ImportError("Please install `mistralai` to use the Mistral.AI API.")
                 client = MistralAIClient(response_format=response_format, **openai_config)
-                self._clients.append(client)
+                self._clients.append(client)  # type: ignore[arg-type]
             elif api_type is not None and api_type.startswith("together"):
                 if together_import_exception:
                     raise ImportError("Please install `together` to use the Together.AI API.")
                 client = TogetherClient(response_format=response_format, **openai_config)
-                self._clients.append(client)
+                self._clients.append(client)  # type: ignore[arg-type]
             elif api_type is not None and api_type.startswith("groq"):
                 if groq_import_exception:
                     raise ImportError("Please install `groq` to use the Groq API.")
                 client = GroqClient(response_format=response_format, **openai_config)
-                self._clients.append(client)
+                self._clients.append(client)  # type: ignore[arg-type]
             elif api_type is not None and api_type.startswith("cohere"):
                 if cohere_import_exception:
                     raise ImportError("Please install `cohere` to use the Cohere API.")
                 client = CohereClient(response_format=response_format, **openai_config)
-                self._clients.append(client)
+                self._clients.append(client)  # type: ignore[arg-type]
             elif api_type is not None and api_type.startswith("ollama"):
                 if ollama_import_exception:
                     raise ImportError("Please install `ollama` and `fix-busted-json` to use the Ollama API.")
                 client = OllamaClient(response_format=response_format, **openai_config)
-                self._clients.append(client)
+                self._clients.append(client)  # type: ignore[arg-type]
             elif api_type is not None and api_type.startswith("bedrock"):
                 self._configure_openai_config_for_bedrock(config, openai_config)
                 if bedrock_import_exception:
                     raise ImportError("Please install `boto3` to use the Amazon Bedrock API.")
                 client = BedrockClient(response_format=response_format, **openai_config)
-                self._clients.append(client)
+                self._clients.append(client)  # type: ignore[arg-type]
             elif api_type is not None and api_type.startswith("responses"):
                 # OpenAI Responses API (stateful). Reuse the same OpenAI SDK but call the `/responses` endpoint via the new client.
                 @require_optional_import("openai>=1.66.2", "openai")
                 def create_responses_client() -> OpenAI:
                     client = OpenAI(**openai_config)
-                    self._clients.append(OpenAIResponsesClient(client, response_format=response_format))
+                    self._clients.append(OpenAIResponsesClient(client, response_format=response_format))  # type: ignore[arg-type]
                     return client
 
                 client = create_responses_client()
@@ -969,7 +973,7 @@ def create_responses_client() -> OpenAI:
                 @require_optional_import("openai>=1.66.2", "openai")
                 def create_openai_client() -> OpenAI:
                     client = OpenAI(**openai_config)
-                    self._clients.append(OpenAIClient(client, response_format))
+                    self._clients.append(OpenAIClient(client, response_format))  # type: ignore[arg-type]
                     return client
 
                 client = create_openai_client()
@@ -1134,12 +1138,12 @@ def create(self, **config: Any) -> ModelClient.ModelClientResponseProtocol:
                     )
                     request_ts = get_current_ts()
 
-                    response: ModelClient.ModelClientResponseProtocol = cache.get(key, None)
+                    response: ChatCompletionExtended | None = cache.get(key, None)
 
                     if response is not None:
                         response.message_retrieval_function = client.message_retrieval
                         try:
-                            response.cost  # type: ignore [attr-defined]
+                            response.cost
                         except AttributeError:
                             # update attribute if cost is not calculated
                             response.cost = client.cost(response)
@@ -1157,7 +1161,7 @@ def create(self, **config: Any) -> ModelClient.ModelClientResponseProtocol:
                                 request=params,
                                 response=response,
                                 is_cached=1,
-                                cost=response.cost,
+                                cost=response.cost if response.cost is not None else 0.0,
                                 start_time=request_ts,
                             )
 
@@ -1272,12 +1276,10 @@ def create(self, **config: Any) -> ModelClient.ModelClientResponseProtocol:
         raise RuntimeError("Should not reach here.")
 
     @staticmethod
-    def _cost_with_customized_price(
-        response: ModelClient.ModelClientResponseProtocol, price_1k: tuple[float, float]
-    ) -> None:
+    def _cost_with_customized_price(response: ChatCompletion | Completion, price_1k: tuple[float, float]) -> float:
         """If a customized cost is passed, overwrite the cost in the response."""
-        n_input_tokens = response.usage.prompt_tokens if response.usage is not None else 0  # type: ignore [union-attr]
-        n_output_tokens = response.usage.completion_tokens if response.usage is not None else 0  # type: ignore [union-attr]
+        n_input_tokens = response.usage.prompt_tokens if response.usage is not None else 0
+        n_output_tokens = response.usage.completion_tokens if response.usage is not None else 0
         if n_output_tokens is None:
             n_output_tokens = 0
         return (n_input_tokens * price_1k[0] + n_output_tokens * price_1k[1]) / 1000
@@ -1451,17 +1453,17 @@ def clear_usage_summary(self) -> None:
 
     @classmethod
     def extract_text_or_completion_object(
-        cls, response: ModelClient.ModelClientResponseProtocol
-    ) -> list[str] | list[ModelClient.ModelClientResponseProtocol.Choice.Message]:
+        cls, response: ChatCompletionExtended
+    ) -> list[str] | list[ChatCompletionMessage]:
         """Extract the text or ChatCompletion objects from a completion or chat response.
 
         Args:
-            response (ChatCompletion | Completion): The response from openai.
+            response: The response from openai with message_retrieval_function attached.
 
         Returns:
             A list of text, or a list of ChatCompletion objects if function_call/tool_calls are present.
         """
-        return response.message_retrieval_function(response)
+        return response.message_retrieval_function(response)  # type: ignore [misc]
 
 
 # -----------------------------------------------------------------------------

autogen/oai/cohere.py

@@ -99,6 +99,8 @@ def create_client(self):
 class CohereClient:
     """Client for Cohere's API."""
 
+    RESPONSE_USAGE_KEYS: list[str] = ["prompt_tokens", "completion_tokens", "total_tokens", "cost", "model"]
+
     def __init__(self, **kwargs: Unpack[CohereEntryDict]):
         """Requires api_key or environment variable to be set
 

autogen/oai/gemini.py

@@ -139,6 +139,8 @@ def create_client(self):
 class GeminiClient:
     """Client for Google's Gemini API."""
 
+    RESPONSE_USAGE_KEYS: list[str] = ["prompt_tokens", "completion_tokens", "total_tokens", "cost", "model"]
+
     # Mapping, where Key is a term used by Autogen, and Value is a term used by Gemini
     PARAMS_MAPPING = {
         "max_tokens": "max_output_tokens",

autogen/oai/groq.py

@@ -77,6 +77,8 @@ def create_client(self):
 class GroqClient:
     """Client for Groq's API."""
 
+    RESPONSE_USAGE_KEYS: list[str] = ["prompt_tokens", "completion_tokens", "total_tokens", "cost", "model"]
+
     def __init__(self, **kwargs: Unpack[GroqEntryDict]):
         """Requires api_key or environment variable to be set