fix: Improve type annotations across multiple files

- Replace Optional[set[str]] with Union[set[str], None] in json methods
- Fix add_nodes_to_network call parameters in flow_visualizer
- Add __base__=BaseModel to create_model call in structured_tool
- Clean up imports in provider.py

Co-Authored-By: Joe Moura <joao@crewai.com>

src/crewai/cli/provider.py

@@ -5,6 +5,7 @@ from pathlib import Path
 
 import click
 import requests
+from typing import Any
 
 from crewai.cli.constants import JSON_URL, MODELS, PROVIDERS
 
@@ -192,7 +193,7 @@ def download_data(response):
     data_chunks = []
     with click.progressbar(
         length=total_size, label="Downloading", show_pos=True
-    ) as progress_bar:
+    ) as progress_bar:  # type: Any
         for chunk in response.iter_content(block_size):
             if chunk:
                 data_chunks.append(chunk)

src/crewai/crews/crew_output.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any, Dict, Optional
+from typing import Any, Callable, Dict, Optional, Union
 
 from pydantic import BaseModel, Field
 
@@ -23,14 +23,25 @@ class CrewOutput(BaseModel):
     )
     token_usage: UsageMetrics = Field(description="Processed token summary", default={})
 
-    @property
-    def json(self) -> Optional[str]:
+    def json(
+        self,
+        *,
+        include: Union[set[str], None] = None,
+        exclude: Union[set[str], None] = None,
+        by_alias: bool = False,
+        exclude_unset: bool = False,
+        exclude_defaults: bool = False,
+        exclude_none: bool = False,
+        encoder: Optional[Callable[[Any], Any]] = None,
+        models_as_dict: bool = True,
+        **dumps_kwargs: Any,
+    ) -> str:
         if self.tasks_output[-1].output_format != OutputFormat.JSON:
             raise ValueError(
                 "No JSON output found in the final task. Please make sure to set the output_json property in the final task in your crew."
             )
 
-        return json.dumps(self.json_dict)
+        return json.dumps(self.json_dict, default=encoder, **dumps_kwargs)
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert json_output and pydantic_output to a dictionary."""

src/crewai/flow/flow_visualizer.py

@@ -106,7 +106,12 @@ class FlowPlot:
 
             # Add nodes to the network
             try:
-                add_nodes_to_network(net, self.flow, node_positions, self.node_styles)
+                add_nodes_to_network(
+                    net,
+                    flow=self.flow,
+                    pos=node_positions,
+                    node_styles=self.node_styles
+                )
             except Exception as e:
                 raise RuntimeError(f"Failed to add nodes to network: {str(e)}")
 

src/crewai/tasks/task_output.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any, Dict, Optional
+from typing import Any, Callable, Dict, Optional, Union
 
 from pydantic import BaseModel, Field, model_validator
 
@@ -34,8 +34,19 @@ class TaskOutput(BaseModel):
         self.summary = f"{excerpt}..."
         return self
 
-    @property
-    def json(self) -> Optional[str]:
+    def json(
+        self,
+        *,
+        include: Union[set[str], None] = None,
+        exclude: Union[set[str], None] = None,
+        by_alias: bool = False,
+        exclude_unset: bool = False,
+        exclude_defaults: bool = False,
+        exclude_none: bool = False,
+        encoder: Optional[Callable[[Any], Any]] = None,
+        models_as_dict: bool = True,
+        **dumps_kwargs: Any,
+    ) -> str:
         if self.output_format != OutputFormat.JSON:
             raise ValueError(
                 """
@@ -45,7 +56,7 @@ class TaskOutput(BaseModel):
                 """
             )
 
-        return json.dumps(self.json_dict)
+        return json.dumps(self.json_dict, default=encoder, **dumps_kwargs)
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert json_output and pydantic_output to a dictionary."""

src/crewai/tools/structured_tool.py

@@ -142,7 +142,12 @@ class CrewStructuredTool:
 
         # Create model
         schema_name = f"{name.title()}Schema"
-        return create_model(schema_name, **fields)
+        return create_model(
+            schema_name,
+            __base__=BaseModel,
+            __config__=None,
+            **{k: v for k, v in fields.items()}
+        )
 
     def _validate_function_signature(self) -> None:
         """Validate that the function signature matches the args schema."""
@@ -170,7 +175,7 @@ class CrewStructuredTool:
                         f"not found in args_schema"
                     )
 
-    def _parse_args(self, raw_args: Union[str, dict]) -> dict:
+    def _parse_args(self, raw_args: Union[str, dict[str, Any]]) -> dict[str, Any]:
         """Parse and validate the input arguments against the schema.
 
         Args:
@@ -178,6 +183,9 @@ class CrewStructuredTool:
 
         Returns:
             The validated arguments as a dictionary
+
+        Raises:
+            ValueError: If the arguments cannot be parsed or fail validation
         """
         if isinstance(raw_args, str):
             try:
@@ -195,8 +203,8 @@ class CrewStructuredTool:
 
     async def ainvoke(
         self,
-        input: Union[str, dict],
-        config: Optional[dict] = None,
+        input: Union[str, dict[str, Any]],
+        config: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> Any:
         """Asynchronously invoke the tool.
@@ -229,7 +237,10 @@ class CrewStructuredTool:
         return self.invoke(input_dict)
 
     def invoke(
-        self, input: Union[str, dict], config: Optional[dict] = None, **kwargs: Any
+        self,
+        input: Union[str, dict[str, Any]],
+        config: Optional[dict[str, Any]] = None,
+        **kwargs: Any
     ) -> Any:
         """Main method for tool execution."""
         parsed_args = self._parse_args(input)