refactor: Improve code quality based on review feedback

- Add type hints and model version comments in constants.py
- Add API key validation helper function in utils.py
- Improve error handling for file operations
- Add test constants and parameterized tests for Mistral provider

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

src/crewai/cli/constants.py

@@ -1,4 +1,6 @@
-ENV_VARS = {
+from typing import Dict, List
+
+ENV_VARS: Dict[str, List[Dict[str, str]]] = {
     "openai": [
         {
             "prompt": "Enter your OPENAI API key (press Enter to skip)",
@@ -94,7 +96,7 @@ ENV_VARS = {
 }
 
 
-PROVIDERS = [
+PROVIDERS: List[str] = [
     "openai",
     "anthropic",
     "gemini",
@@ -104,16 +106,16 @@ PROVIDERS = [
     "bedrock",
     "azure",
     "cerebras",
-    "mistral",
+    "mistral",  # Added in v0.86.0
 ]
 
-MODELS = {
+MODELS: Dict[str, List[str]] = {
     "openai": ["gpt-4", "gpt-4o", "gpt-4o-mini", "o1-mini", "o1-preview"],
     "mistral": [
-        "mistral-tiny",
-        "mistral-small",
-        "mistral-medium",
-        "mistral-large",
+        "mistral-tiny",    # 7B model optimized for speed
+        "mistral-small",   # 7B model balanced for performance
+        "mistral-medium",  # 8x7B model for enhanced capabilities
+        "mistral-large",   # Latest model with highest performance
     ],
     "anthropic": [
         "claude-3-5-sonnet-20240620",

src/crewai/cli/create_crew.py

@@ -10,7 +10,12 @@ from crewai.cli.provider import (
     select_model,
     select_provider,
 )
-from crewai.cli.utils import copy_template, load_env_vars, write_env_file
+from crewai.cli.utils import (
+    copy_template,
+    load_env_vars,
+    validate_api_keys,
+    write_env_file,
+)
 
 
 def create_folder_structure(name, parent_folder=None):
@@ -162,14 +167,13 @@ def create_crew(name, provider=None, skip_provider=False, parent_folder=None):
                     if api_key_value.strip():
                         env_vars[key_name] = api_key_value
 
-        # Only write env file if we have API keys
-        has_api_keys = any(
-            key in env_vars and env_vars[key].strip()
-            for key in ["MISTRAL_API_KEY", "OPENAI_API_KEY", "ANTHROPIC_API_KEY"]
-        )
-        if has_api_keys:
-            write_env_file(folder_path, env_vars)
-            click.secho("API keys and model saved to .env file", fg="green")
+        if validate_api_keys(env_vars):
+            try:
+                write_env_file(folder_path, env_vars)
+                click.secho("API keys and model saved to .env file", fg="green")
+            except IOError as e:
+                click.secho(f"Error writing .env file: {str(e)}", fg="red")
+                raise
         else:
             click.secho(
                 "No API keys provided. Skipping .env file creation.", fg="yellow"

src/crewai/cli/utils.py

@@ -2,6 +2,7 @@ import os
 import shutil
 import sys
 from functools import reduce
+from pathlib import Path
 from typing import Any, Dict, List
 
 import click
@@ -235,15 +236,38 @@ def update_env_vars(env_vars, provider, model):
     return env_vars
 
 
-def write_env_file(folder_path, env_vars):
+def validate_api_keys(env_vars: Dict[str, str]) -> bool:
+    """
+    Validates that at least one API key is present and non-empty in the environment variables.
+
+    Args:
+        env_vars (Dict[str, str]): Dictionary of environment variables
+
+    Returns:
+        bool: True if at least one API key is present and non-empty
+    """
+    api_keys = ["MISTRAL_API_KEY", "OPENAI_API_KEY", "ANTHROPIC_API_KEY"]
+    return any(
+        key in env_vars and env_vars[key].strip()
+        for key in api_keys
+    )
+
+def write_env_file(folder_path: Path, env_vars: Dict[str, str]) -> None:
     """
     Writes environment variables to a .env file in the specified folder.
 
     Args:
-    - folder_path (Path): The path to the folder where the .env file will be written.
-    - env_vars (dict): A dictionary of environment variables to write.
+        folder_path (Path): The path to the folder where the .env file will be written.
+        env_vars (Dict[str, str]): A dictionary of environment variables to write.
+
+    Raises:
+        IOError: If there is an error writing to the .env file
     """
     env_file_path = folder_path / ".env"
-    with open(env_file_path, "w") as file:
-        for key, value in env_vars.items():
-            file.write(f"{key}={value}\n")
+    try:
+        with open(env_file_path, "w") as file:
+            for key, value in env_vars.items():
+                file.write(f"{key}={value}\n")
+    except IOError as e:
+        click.secho(f"Error writing .env file: {str(e)}", fg="red")
+        raise

tests/cli/cli_test.py

@@ -2,6 +2,7 @@ from pathlib import Path
 from unittest import mock
 
 import pytest
+import click
 from click.testing import CliRunner
 
 from crewai.cli.cli import (
@@ -21,6 +22,13 @@ from crewai.cli.cli import (
 
 
 from crewai.cli.cli import create
+TEST_CONSTANTS = {
+    "CREW_NAME": "test_crew",
+    "MISTRAL_API_KEY": "mistral_api_key_123",
+    "MISTRAL_MODEL": "mistral-tiny",
+    "EMPTY_KEY": "",
+}
+
 @pytest.fixture
 def runner():
     return CliRunner()
@@ -310,6 +318,30 @@ def test_flow_add_crew(mock_path_exists, mock_create_embedded_crew, runner):
     assert isinstance(call_kwargs["parent_folder"], Path)
 
 
+@pytest.mark.parametrize("provider,model,api_key,has_valid_keys,expected_outputs", [
+    (
+        "mistral",
+        TEST_CONSTANTS["MISTRAL_MODEL"],
+        TEST_CONSTANTS["MISTRAL_API_KEY"],
+        True,
+        ["API keys and model saved", f"Selected model: {TEST_CONSTANTS['MISTRAL_MODEL']}"]
+    ),
+    (
+        "mistral",
+        TEST_CONSTANTS["MISTRAL_MODEL"],
+        TEST_CONSTANTS["EMPTY_KEY"],
+        False,
+        ["No API keys provided", f"Selected model: {TEST_CONSTANTS['MISTRAL_MODEL']}"]
+    ),
+    (
+        "mistral",
+        None,
+        TEST_CONSTANTS["EMPTY_KEY"],
+        False,
+        ["No model selected"]
+    ),
+])
+@mock.patch("crewai.cli.create_crew.validate_api_keys")
 @mock.patch("crewai.cli.create_crew.write_env_file")
 @mock.patch("crewai.cli.create_crew.load_env_vars")
 @mock.patch("crewai.cli.create_crew.get_provider_data")
@@ -317,29 +349,46 @@ def test_flow_add_crew(mock_path_exists, mock_create_embedded_crew, runner):
 @mock.patch("crewai.cli.create_crew.select_provider")
 @mock.patch("crewai.cli.create_crew.click.confirm")
 @mock.patch("crewai.cli.create_crew.click.prompt")
-def test_create_crew_with_mistral_provider(
+def test_create_crew_scenarios(
     mock_prompt, mock_confirm, mock_select_provider, mock_select_model,
-    mock_get_provider_data, mock_load_env_vars, mock_write_env_file, runner
+    mock_get_provider_data, mock_load_env_vars, mock_write_env_file, mock_validate_api_keys,
+    runner, provider, model, api_key, has_valid_keys, expected_outputs
 ):
-    # Mock folder override confirmation
+    """Test different scenarios for crew creation with provider configuration.
+    
+    Args:
+        mock_*: Mock objects for various dependencies
+        runner: Click test runner
+        provider: Provider to test (e.g. "mistral")
+        model: Model to select (e.g. "mistral-tiny")
+        api_key: API key to provide
+        has_valid_keys: Whether the API key validation should pass
+        expected_output: Expected message in the output
+    """
     mock_confirm.return_value = True
-    # Mock provider data
-    mock_get_provider_data.return_value = {"mistral": ["mistral-tiny"]}
-    # Mock empty env vars
+    mock_get_provider_data.return_value = {"mistral": [TEST_CONSTANTS["MISTRAL_MODEL"]]}
     mock_load_env_vars.return_value = {}
-    # Mock provider and model selection
-    mock_select_provider.return_value = "mistral"
-    mock_select_model.return_value = "mistral-tiny"
-    # Mock API key input
-    mock_prompt.return_value = "mistral_api_key_123"
-    
-    # Run the command
-    result = runner.invoke(create, ["crew", "test_crew"])
-    
-    assert result.exit_code == 0
-    assert "API keys and model saved to .env file" in result.output
-    assert "Selected model: mistral-tiny" in result.output
+    mock_select_provider.return_value = provider
+    mock_select_model.return_value = model
+    mock_prompt.return_value = api_key
+    mock_validate_api_keys.return_value = has_valid_keys
 
+    # When model is None, simulate model selection being cancelled
+    if model is None:
+        mock_select_model.side_effect = click.UsageError("No model selected")
+
+    result = runner.invoke(create, ["crew", TEST_CONSTANTS["CREW_NAME"]], input="y\n")
+    
+    # For model=None case, we expect error message
+    if model is None:
+        assert result.exit_code == 2  # UsageError exit code
+        assert "No model selected" in result.output
+    else:
+        assert result.exit_code == 0
+        for expected_output in expected_outputs:
+            assert expected_output in result.output
+
+@mock.patch("crewai.cli.create_crew.validate_api_keys")
 @mock.patch("crewai.cli.create_crew.write_env_file")
 @mock.patch("crewai.cli.create_crew.load_env_vars")
 @mock.patch("crewai.cli.create_crew.get_provider_data")
@@ -347,31 +396,32 @@ def test_create_crew_with_mistral_provider(
 @mock.patch("crewai.cli.create_crew.select_provider")
 @mock.patch("crewai.cli.create_crew.click.confirm")
 @mock.patch("crewai.cli.create_crew.click.prompt")
-def test_create_crew_with_mistral_provider_no_key(
+def test_create_crew_with_file_error(
     mock_prompt, mock_confirm, mock_select_provider, mock_select_model,
-    mock_get_provider_data, mock_load_env_vars, mock_write_env_file, runner
+    mock_get_provider_data, mock_load_env_vars, mock_write_env_file, mock_validate_api_keys,
+    runner
 ):
     # Mock folder override confirmation
     mock_confirm.return_value = True
     # Mock provider data
-    mock_get_provider_data.return_value = {"mistral": ["mistral-tiny"]}
+    mock_get_provider_data.return_value = {"mistral": [TEST_CONSTANTS["MISTRAL_MODEL"]]}
     # Mock empty env vars
     mock_load_env_vars.return_value = {}
     # Mock provider and model selection
     mock_select_provider.return_value = "mistral"
-    mock_select_model.return_value = "mistral-tiny"
-    # Mock empty API key input for all prompts
-    mock_prompt.side_effect = [""] * 10  # Enough empty responses for all prompts
+    mock_select_model.return_value = TEST_CONSTANTS["MISTRAL_MODEL"]
+    # Mock API key input
+    mock_prompt.return_value = TEST_CONSTANTS["MISTRAL_API_KEY"]
+    # Mock API key validation
+    mock_validate_api_keys.return_value = True
+    # Mock file write error
+    mock_write_env_file.side_effect = IOError("Permission denied")
     
-    # Run the command
-    result = runner.invoke(create, ["crew", "test_crew"])
+    result = runner.invoke(create, ["crew", TEST_CONSTANTS["CREW_NAME"]], input="y\n")
     
-    assert result.exit_code == 0
-    assert "No API keys provided. Skipping .env file creation." in result.output
-    # Verify env file was not written since no API keys were provided
-    mock_write_env_file.assert_not_called()
-    # Verify model was still selected
-    assert "Selected model: mistral-tiny" in result.output
+    assert result.exit_code == 1
+    assert "Error writing .env file: Permission denied" in result.output
+    assert mock_write_env_file.called
 
 def test_add_crew_to_flow_not_in_root(runner):
     # Simulate not being in the root of a flow project