diff --git a/autogen/_pydantic.py b/autogen/_pydantic.py index 89dbc4fd2..ced26652f 100644 --- a/autogen/_pydantic.py +++ b/autogen/_pydantic.py @@ -13,7 +13,7 @@ if not PYDANTIC_V1: 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 @@ if not PYDANTIC_V1: # 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: diff --git a/autogen/function_utils.py b/autogen/function_utils.py index 0189836a4..775c7fa0e 100644 --- a/autogen/function_utils.py +++ b/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 get_function_schema(f: Callable[..., Any], *, name: Optional[str] = None, de 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) diff --git a/pyproject.toml b/pyproject.toml index d1b49ee28..115fd5b7d 100644 --- a/pyproject.toml +++ b/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", ]