FIX: fix errors and pyright mypy issues

Author: SongChiYoung

python/packages/autogen-ext/src/autogen_ext/models/openai/_message_transform.py

@@ -30,22 +30,25 @@
 
 from ._utils import assert_valid_name, func_call_to_oai
 
+EMPTY: Dict[str, Any] = {}
+
 
 # ===Mini Transformers===
-def _assert_valid_name(message: LLMMessage, context: Dict[str, Any]):
+def _assert_valid_name(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, Any]:
+    assert isinstance(message, (UserMessage, AssistantMessage))
     assert_valid_name(message.source)
-    result: Dict[str, Any] = {}
-    return result
+    return EMPTY
 
 
-def _set_role(role: str):
+def _set_role(role: str) -> Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]:
     def inner(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, Any]:
         return {"role": role}
 
     return inner
 
 
 def _set_name(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, Any]:
+    assert isinstance(message, (UserMessage, AssistantMessage))
     assert_valid_name(message.source)
     return {"name": message.source}
 
@@ -54,13 +57,16 @@ def _set_content_direct(message: LLMMessage, context: Dict[str, Any]) -> Dict[st
     return {"content": message.content}
 
 
-def _set_prepend_text_content(message: UserMessage, context: Dict[str, Any]) -> Dict[str, Any]:
+def _set_prepend_text_content(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, Any]:
+    assert isinstance(message, (UserMessage, AssistantMessage))
+    assert isinstance(message.content, str)
     prepend = context.get("prepend_name", False)
     prefix = f"{message.source} said:\n" if prepend else ""
     return {"content": prefix + message.content}
 
 
-def _set_multimodal_content(message: UserMessage, context: Dict[str, Any]) -> Dict[str, Any]:
+def _set_multimodal_content(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, Any]:
+    assert isinstance(message, (UserMessage, AssistantMessage))
     prepend = context.get("prepend_name", False)
     parts: List[ChatCompletionContentPartParam] = []
 
@@ -79,13 +85,16 @@ def _set_multimodal_content(message: UserMessage, context: Dict[str, Any]) -> Di
     return {"content": parts}
 
 
-def _set_tool_calls(message: AssistantMessage, context: Dict[str, Any]) -> Dict[str, Any]:
+def _set_tool_calls(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, Any]:
+    assert isinstance(message.content, list)
+    assert isinstance(message, AssistantMessage)
     return {
         "tool_calls": [func_call_to_oai(x) for x in message.content],
     }
 
 
-def _set_thought_as_content(message: AssistantMessage, context: Dict[str, Any]) -> Dict[str, Any]:
+def _set_thought_as_content(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, Any]:
+    assert isinstance(message, AssistantMessage)
     return {"content": message.thought}
 
 
@@ -187,6 +196,7 @@ def user_condition(message: LLMMessage, context: Dict[str, Any]) -> str:
 
 
 def assistant_condition(message: LLMMessage, context: Dict[str, Any]) -> str:
+    assert isinstance(message, AssistantMessage)
     if isinstance(message.content, list):
         if message.thought is not None:
             return "thought"
@@ -208,13 +218,19 @@ def assistant_condition(message: LLMMessage, context: Dict[str, Any]) -> str:
 }
 
 
-def function_execution_result_message(message: LLMMessage, context: Dict[str, Any]) -> List[Dict[str, Any]]:
+def function_execution_result_message(
+    message: LLMMessage, context: Dict[str, Any]
+) -> list[ChatCompletionToolMessageParam]:
+    assert isinstance(message, FunctionExecutionResultMessage)
     return [
         ChatCompletionToolMessageParam(content=x.content, role="tool", tool_call_id=x.call_id) for x in message.content
     ]
 
 
-def function_execution_result_message_gemini(message: LLMMessage, context: Dict[str, Any]) -> List[Dict[str, Any]]:
+def function_execution_result_message_gemini(
+    message: LLMMessage, context: Dict[str, Any]
+) -> list[ChatCompletionToolMessageParam]:
+    assert isinstance(message, FunctionExecutionResultMessage)
     return [
         ChatCompletionToolMessageParam(content=x.content if x.content else " ", role="tool", tool_call_id=x.call_id)
         for x in message.content
@@ -282,4 +298,4 @@ def function_execution_result_message_gemini(message: LLMMessage, context: Dict[
 for model in __unknown_models:
     register_transformer("openai", model, __BASE_TRANSFORMER_MAP)
 
-register_transformer("openai", "default", __BASE_TRANSFORMER_MAP)
+# register_transformer("openai", "default", __BASE_TRANSFORMER_MAP)

python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py

@@ -11,6 +11,7 @@
 from typing import (
     Any,
     AsyncGenerator,
+    Callable,
     Dict,
     List,
     Mapping,
@@ -58,7 +59,6 @@
     ChatCompletionContentPartParam,
     ChatCompletionContentPartTextParam,
     ChatCompletionMessageParam,
-    ChatCompletionMessageToolCallParam,
     ChatCompletionRole,
     ChatCompletionSystemMessageParam,
     ChatCompletionToolMessageParam,
@@ -195,17 +195,6 @@ def system_message_to_oai(message: SystemMessage) -> ChatCompletionSystemMessage
     )
 
 
-def _old_func_call_to_oai(message: FunctionCall) -> ChatCompletionMessageToolCallParam:
-    return ChatCompletionMessageToolCallParam(
-        id=message.id,
-        function={
-            "arguments": message.arguments,
-            "name": message.name,
-        },
-        type="function",
-    )
-
-
 def tool_message_to_oai(
     message: FunctionExecutionResultMessage,
 ) -> Sequence[ChatCompletionToolMessageParam]:
@@ -248,25 +237,16 @@ def to_oai_type(
     }
     transformers = get_transformer("openai", model_family)
 
-    def raise_value_error(message: LLMMessage, context: Dict[str, Any]) -> Dict[str, Any]:
+    def raise_value_error(message: LLMMessage, context: Dict[str, Any]) -> ChatCompletionMessageParam:
         raise ValueError(f"Unknown message type: {type(message)}")
 
-    transformer = transformers.get(type(message), raise_value_error)
+    transformer: Callable[
+        [LLMMessage, Dict[str, Any]], Union[ChatCompletionMessageParam, Sequence[ChatCompletionMessageParam]]
+    ] = transformers.get(type(message), raise_value_error)
     result = transformer(message, context)
     if isinstance(result, list):
-        return result
-    return [result]
-
-
-def _old_to_oai_type(message: LLMMessage, prepend_name: bool = False) -> Sequence[ChatCompletionMessageParam]:
-    if isinstance(message, SystemMessage):
-        return [system_message_to_oai(message)]
-    elif isinstance(message, UserMessage):
-        return [user_message_to_oai(message, prepend_name)]
-    elif isinstance(message, AssistantMessage):
-        return [assistant_message_to_oai(message)]
-    else:
-        return tool_message_to_oai(message)
+        return cast(List[ChatCompletionMessageParam], result)
+    return cast(List[ChatCompletionMessageParam], [result])
 
 
 def calculate_vision_tokens(image: Image, detail: str = "auto") -> int:
@@ -364,19 +344,6 @@ def normalize_name(name: str) -> str:
     return re.sub(r"[^a-zA-Z0-9_-]", "_", name)[:64]
 
 
-def _old_assert_valid_name(name: str) -> str:
-    """
-    Ensure that configured names are valid, raises ValueError if not.
-
-    For munging LLM responses use _normalize_name to ensure LLM specified names don't break the API.
-    """
-    if not re.match(r"^[a-zA-Z0-9_-]+$", name):
-        raise ValueError(f"Invalid name: {name}. Only letters, numbers, '_' and '-' are allowed.")
-    if len(name) > 64:
-        raise ValueError(f"Invalid name: {name}. Name must be less than 64 characters.")
-    return name
-
-
 @dataclass
 class CreateParams:
     messages: List[ChatCompletionMessageParam]

python/packages/autogen-ext/src/autogen_ext/models/openai/_utils.py

@@ -1,5 +1,4 @@
 import re
-from typing import Any, Dict, List, Union
 
 from autogen_core import FunctionCall
 from openai.types.chat import ChatCompletionMessageToolCallParam

python/packages/autogen-ext/src/autogen_ext/transformation/__init__.py

@@ -6,8 +6,6 @@
     register_transformer,
 )
 from .types import (
-    BuilderFunc,
-    BuilderMap,
     TransformerFunc,
     TransformerMap,
 )
@@ -20,6 +18,4 @@
     "MESSAGE_TRANSFORMERS",
     "TransformerMap",
     "TransformerFunc",
-    "BuilderMap",
-    "BuilderFunc",
 ]

python/packages/autogen-ext/src/autogen_ext/transformation/registry.py

@@ -1,18 +1,16 @@
 from collections import defaultdict
-from typing import Any, Callable, Dict, List, Type
+from typing import Any, Callable, Dict, List
 
-from autogen_core.models import LLMMessage, ModelFamily
+from autogen_core.models import LLMMessage
 
 from .types import (
-    BuilderMap,
     TransformerFunc,
     TransformerMap,
 )
 
 # Global registry of model family → message transformer map
 # Each model family (e.g. "gpt-4o", "gemini-1.5-flash") maps to a dict of LLMMessage type → transformer function
 MESSAGE_TRANSFORMERS: Dict[str, Dict[str, TransformerMap]] = defaultdict(dict)
-MESSAGE_BUILDERS: Dict[str, BuilderMap] = {}
 
 
 def build_transformer_func(
@@ -69,7 +67,7 @@ def transformer(message: LLMMessage, context: Dict[str, Any]) -> Any:
     return transformer
 
 
-def register_transformer(api: str, model_family: str, transformer_map: TransformerMap):
+def register_transformer(api: str, model_family: str, transformer_map: TransformerMap) -> None:
     """
     Registers a transformer map for a given model family.
 
@@ -82,6 +80,17 @@ def register_transformer(api: str, model_family: str, transformer_map: Transform
     MESSAGE_TRANSFORMERS[api][model_family] = transformer_map
 
 
+def _find_model_family(api: str, model: str) -> str:
+    """
+    Finds the best matching model family for the given model.
+    Search via prefix matching (e.g. "gpt-4o" → "gpt-4o-1.0").
+    """
+    for family in MESSAGE_TRANSFORMERS[api].keys():
+        if model.startswith(family):
+            return family
+    return "default"
+
+
 def get_transformer(api: str, model_family: str) -> TransformerMap:
     """
     Returns the registered transformer map for the given model family.
@@ -95,9 +104,10 @@ def get_transformer(api: str, model_family: str) -> TransformerMap:
 
     Keeping this as a function (instead of direct dict access) improves long-term flexibility.
     """
-    transformer = MESSAGE_TRANSFORMERS.get(api, {}).get(model_family, {})
-    if not transformer:
-        transformer = MESSAGE_TRANSFORMERS.get("default", {}).get("default", {})
+
+    model = _find_model_family(api, model_family)
+
+    transformer = MESSAGE_TRANSFORMERS.get(api, {}).get(model, {})
 
     if not transformer:
         raise ValueError(f"No transformer found for model family '{model_family}'")

python/packages/autogen-ext/src/autogen_ext/transformation/types.py

@@ -1,9 +1,6 @@
-from typing import Any, Callable, Dict, List, Type
+from typing import Any, Callable, Dict, Type
 
-from autogen_core.models import LLMMessage, ModelFamily
+from autogen_core.models import LLMMessage
 
 TransformerFunc = Callable[[LLMMessage, Dict[str, Any]], Any]
 TransformerMap = Dict[Type[LLMMessage], TransformerFunc]
-
-BuilderFunc = Callable[[List[Any], Dict[str, Any]], Any]
-BuilderMap = Dict[Type[LLMMessage], BuilderFunc]