Fixed types related to function calling

autogen/_pydantic.py

@@ -13,7 +13,7 @@
     from pydantic._internal._typing_extra import eval_type_lenient as evaluate_forwardref
     from pydantic.json_schema import JsonSchemaValue
 
-    def type2schema(t: Optional[Type]) -> JsonSchemaValue:
+    def type2schema(t: Optional[Type[Any]]) -> JsonSchemaValue:
         """Convert a type to a JSON schema
 
         Args:
@@ -51,11 +51,11 @@ def model_dump_json(model: BaseModel) -> str:
 # Remove this once we drop support for pydantic 1.x
 else:  # pragma: no cover
     from pydantic import schema_of
-    from pydantic.typing import evaluate_forwardref as evaluate_forwardref
+    from pydantic.typing import evaluate_forwardref as evaluate_forwardref  # type: ignore[no-redef]
 
-    JsonSchemaValue = Dict[str, Any]
+    JsonSchemaValue = Dict[str, Any]  # type: ignore[misc]
 
-    def type2schema(t: Optional[Type]) -> JsonSchemaValue:
+    def type2schema(t: Optional[Type[Any]]) -> JsonSchemaValue:
         """Convert a type to a JSON schema
 
         Args:

autogen/function_utils.py

@@ -73,7 +73,7 @@ def get_typed_return_annotation(call: Callable[..., Any]) -> Any:
     return get_typed_annotation(annotation, globalns)
 
 
-def get_param_annotations(typed_signature: inspect.Signature) -> Dict[int, Union[Annotated[Type[Any], str], Type[Any]]]:
+def get_param_annotations(typed_signature: inspect.Signature) -> Dict[str, Union[Annotated[Type[Any], str], Type[Any]]]:
     """Get the type annotations of the parameters of a function
 
     Args:
@@ -285,7 +285,7 @@ def f(a: Annotated[str, "Parameter a"], b: int = 2, c: Annotated[float, "Paramet
     return model_dump(function)
 
 
-def get_load_param_if_needed_function(t: Any) -> Optional[Callable[[T, Type[Any]], BaseModel]]:
+def get_load_param_if_needed_function(t: Any) -> Optional[Callable[[Dict[str, Any], Type[BaseModel]], BaseModel]]:
     """Get a function to load a parameter if it is a Pydantic model
 
     Args:
@@ -319,10 +319,10 @@ def load_basemodels_if_needed(func: Callable[..., Any]) -> Callable[..., Any]:
     param_annotations = get_param_annotations(typed_signature)
 
     # get functions for loading BaseModels when needed based on the type annotations
-    kwargs_mapping = {k: get_load_param_if_needed_function(t) for k, t in param_annotations.items()}
+    kwargs_mapping_with_nones = {k: get_load_param_if_needed_function(t) for k, t in param_annotations.items()}
 
     # remove the None values
-    kwargs_mapping = {k: f for k, f in kwargs_mapping.items() if f is not None}
+    kwargs_mapping = {k: f for k, f in kwargs_mapping_with_nones.items() if f is not None}
 
     # a function that loads the parameters before calling the original function
     @functools.wraps(func)

pyproject.toml

@@ -68,6 +68,8 @@ files = [
     "autogen/exception_utils.py",
     "autogen/coding",
     "autogen/oai/openai_utils.py",
+    "autogen/_pydantic.py",
+    "autogen/function_utils.py",
     "autogen/io",
     "test/io",
 ]