Add ruff sort to pre-commit and sort imports in the library (#1259)

* lint

* bump ver

* bump ver

* fixed circular import

---------

Co-authored-by: Jirka Borovec <[email protected]>

flaml/__init__.py

@@ -1,10 +1,10 @@
 import logging
+
 from flaml.automl import AutoML, logger_formatter
-from flaml.tune.searcher import CFO, BlendSearch, FLOW2, BlendSearchTuner, RandomSearch
 from flaml.onlineml.autovw import AutoVW
+from flaml.tune.searcher import CFO, FLOW2, BlendSearch, BlendSearchTuner, RandomSearch
 from flaml.version import __version__
 
-
 # Set the root logger.
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)

flaml/autogen/__init__.py

@@ -1,3 +1,3 @@
-from .oai import *
 from .agentchat import *
 from .code_utils import DEFAULT_MODEL, FAST_MODEL
+from .oai import *

flaml/autogen/agentchat/__init__.py

@@ -1,8 +1,8 @@
 from .agent import Agent
-from .conversable_agent import ConversableAgent
 from .assistant_agent import AssistantAgent
-from .user_proxy_agent import UserProxyAgent
+from .conversable_agent import ConversableAgent
 from .groupchat import GroupChat, GroupChatManager
+from .user_proxy_agent import UserProxyAgent
 
 __all__ = [
     "Agent",

flaml/autogen/agentchat/assistant_agent.py

@@ -1,6 +1,7 @@
-from .conversable_agent import ConversableAgent
 from typing import Callable, Dict, Optional, Union
 
+from .conversable_agent import ConversableAgent
+
 
 class AssistantAgent(ConversableAgent):
     """(In preview) Assistant agent, designed to solve a task with LLM.

flaml/autogen/agentchat/contrib/math_user_proxy_agent.py

@@ -1,14 +1,14 @@
-import re
 import os
-from pydantic import BaseModel, Extra, root_validator
-from typing import Any, Callable, Dict, List, Optional, Union
+import re
 from time import sleep
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Extra, root_validator
 
 from flaml.autogen.agentchat import Agent, UserProxyAgent
-from flaml.autogen.code_utils import UNKNOWN, extract_code, execute_code, infer_lang
+from flaml.autogen.code_utils import UNKNOWN, execute_code, extract_code, infer_lang
 from flaml.autogen.math_utils import get_answer
 
-
 PROMPTS = {
     # default
     "default": """Let's use Python to solve a math problem.

flaml/autogen/agentchat/contrib/retrieve_assistant_agent.py

@@ -1,6 +1,7 @@
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
 from flaml.autogen.agentchat.agent import Agent
 from flaml.autogen.agentchat.assistant_agent import AssistantAgent
-from typing import Callable, Dict, Optional, Union, List, Tuple, Any
 
 
 class RetrieveAssistantAgent(AssistantAgent):

flaml/autogen/agentchat/contrib/retrieve_user_proxy_agent.py

@@ -1,11 +1,12 @@
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
 import chromadb
-from flaml.autogen.agentchat.agent import Agent
+from IPython import get_ipython
+
 from flaml.autogen.agentchat import UserProxyAgent
-from flaml.autogen.retrieve_utils import create_vector_db_from_dir, query_vector_db, num_tokens_from_text
+from flaml.autogen.agentchat.agent import Agent
 from flaml.autogen.code_utils import extract_code
-
-from typing import Callable, Dict, Optional, Union, List, Tuple, Any
-from IPython import get_ipython
+from flaml.autogen.retrieve_utils import create_vector_db_from_dir, num_tokens_from_text, query_vector_db
 
 try:
     from termcolor import colored

flaml/autogen/agentchat/conversable_agent.py

@@ -1,10 +1,10 @@
 import asyncio
-from collections import defaultdict
 import copy
 import json
+from collections import defaultdict
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
+
 from flaml.autogen import oai
-from .agent import Agent
 from flaml.autogen.code_utils import (
     DEFAULT_MODEL,
     UNKNOWN,
@@ -13,6 +13,8 @@
     infer_lang,
 )
 
+from .agent import Agent
+
 try:
     from termcolor import colored
 except ImportError:

flaml/autogen/agentchat/groupchat.py

@@ -1,6 +1,7 @@
-from dataclasses import dataclass
 import sys
+from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
+
 from .agent import Agent
 from .conversable_agent import ConversableAgent
 

flaml/autogen/agentchat/user_proxy_agent.py

@@ -1,6 +1,7 @@
-from .conversable_agent import ConversableAgent
 from typing import Callable, Dict, Optional, Union
 
+from .conversable_agent import ConversableAgent
+
 
 class UserProxyAgent(ConversableAgent):
     """(In preview) A proxy agent for the user, that can execute code and provide feedback to the other agents.

flaml/autogen/code_utils.py

@@ -1,13 +1,14 @@
-import signal
-import subprocess
-import sys
+import logging
 import os
 import pathlib
-from typing import List, Dict, Tuple, Optional, Union, Callable
 import re
+import signal
+import subprocess
+import sys
 import time
 from hashlib import md5
-import logging
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
 from flaml.autogen import oai
 
 try:

flaml/autogen/math_utils.py

@@ -1,5 +1,6 @@
 from typing import Optional
-from flaml.autogen import oai, DEFAULT_MODEL
+
+from flaml.autogen import DEFAULT_MODEL, oai
 
 _MATH_PROMPT = "{problem} Solve the problem carefully. Simplify your answer as much as possible. Put the final answer in \\boxed{{}}."
 _MATH_CONFIG = {

flaml/autogen/oai/__init__.py

@@ -1,10 +1,10 @@
-from flaml.autogen.oai.completion import Completion, ChatCompletion
+from flaml.autogen.oai.completion import ChatCompletion, Completion
 from flaml.autogen.oai.openai_utils import (
-    get_config_list,
+    config_list_from_json,
+    config_list_from_models,
     config_list_gpt4_gpt35,
     config_list_openai_aoai,
-    config_list_from_models,
-    config_list_from_json,
+    get_config_list,
 )
 
 __all__ = [

flaml/autogen/oai/completion.py

@@ -1,28 +1,31 @@
-from time import sleep
 import logging
-import time
-from typing import List, Optional, Dict, Callable, Union
-import sys
 import shutil
+import sys
+import time
+from time import sleep
+from typing import Callable, Dict, List, Optional, Union
+
 import numpy as np
-from flaml import tune, BlendSearch
-from flaml.tune.space import is_constant
+
+from flaml import BlendSearch, tune
 from flaml.automl.logger import logger_formatter
+from flaml.tune.space import is_constant
+
 from .openai_utils import get_key
 
 try:
+    import diskcache
     import openai
+    from openai import Completion as openai_Completion
     from openai.error import (
-        ServiceUnavailableError,
-        RateLimitError,
+        APIConnectionError,
         APIError,
+        AuthenticationError,
         InvalidRequestError,
-        APIConnectionError,
+        RateLimitError,
+        ServiceUnavailableError,
         Timeout,
-        AuthenticationError,
     )
-    from openai import Completion as openai_Completion
-    import diskcache
 
     ERROR = None
 except ImportError:

flaml/autogen/oai/openai_utils.py

@@ -1,7 +1,7 @@
-import os
 import json
-from typing import List, Optional, Dict, Set, Union
 import logging
+import os
+from typing import Dict, List, Optional, Set, Union
 
 NON_CACHE_KEY = ["api_key", "api_base", "api_type", "api_version"]
 

flaml/autogen/retrieve_utils.py

@@ -1,13 +1,14 @@
-from typing import List, Union, Dict, Tuple
+import glob
+import logging
 import os
-import requests
+from typing import Dict, List, Tuple, Union
 from urllib.parse import urlparse
-import glob
-import tiktoken
+
 import chromadb
-from chromadb.api import API
 import chromadb.utils.embedding_functions as ef
-import logging
+import requests
+import tiktoken
+from chromadb.api import API
 
 logger = logging.getLogger(__name__)
 TEXT_FORMATS = ["txt", "json", "csv", "tsv", "md", "html", "htm", "rtf", "rst", "jsonl", "log", "xml", "yaml", "yml"]

flaml/automl/__init__.py

@@ -1,5 +1,5 @@
 from flaml.automl.automl import AutoML, size
 from flaml.automl.logger import logger_formatter
-from flaml.automl.state import SearchState, AutoMLState
+from flaml.automl.state import AutoMLState, SearchState
 
 __all__ = ["AutoML", "AutoMLState", "SearchState", "logger_formatter", "size"]

flaml/automl/automl.py

@@ -3,40 +3,41 @@
 #  * Licensed under the MIT License. See LICENSE file in the
 #  * project root for license information.
 from __future__ import annotations
-import time
+
+import json
+import logging
 import os
 import sys
-from typing import Callable, List, Union, Optional
+import time
 from functools import partial
+from typing import Callable, List, Optional, Union
+
 import numpy as np
-import logging
-import json
 
-from flaml.automl.state import SearchState, AutoMLState
+from flaml import tune
+from flaml.automl.logger import logger, logger_formatter
 from flaml.automl.ml import train_estimator
+from flaml.automl.spark import DataFrame, Series, psDataFrame, psSeries
+from flaml.automl.state import AutoMLState, SearchState
+from flaml.automl.task.factory import task_factory
 
+# TODO check to see when we can remove these
+from flaml.automl.task.task import CLASSIFICATION, Task
 from flaml.automl.time_series import TimeSeriesDataset
+from flaml.automl.training_log import training_log_reader, training_log_writer
 from flaml.config import (
-    MIN_SAMPLE_TRAIN,
+    CV_HOLDOUT_THRESHOLD,
     MEM_THRES,
+    MIN_SAMPLE_TRAIN,
+    N_SPLITS,
     RANDOM_SEED,
+    SAMPLE_MULTIPLY_FACTOR,
     SMALL_LARGE_THRES,
-    CV_HOLDOUT_THRESHOLD,
     SPLIT_RATIO,
-    N_SPLITS,
-    SAMPLE_MULTIPLY_FACTOR,
 )
-
-# TODO check to see when we can remove these
-from flaml.automl.task.task import CLASSIFICATION, Task
-from flaml.automl.task.factory import task_factory
-from flaml import tune
-from flaml.automl.logger import logger, logger_formatter
-from flaml.automl.training_log import training_log_reader, training_log_writer
 from flaml.default import suggest_learner
-from flaml.version import __version__ as flaml_version
-from flaml.automl.spark import psDataFrame, psSeries, DataFrame, Series
 from flaml.tune.spark.utils import check_spark, get_broadcast_data
+from flaml.version import __version__ as flaml_version
 
 ERROR = (
     DataFrame is None and ImportError("please install flaml[automl] option to use the flaml.automl package.") or None

flaml/automl/data.py

@@ -2,15 +2,17 @@
 #  * Copyright (c) Microsoft Corporation. All rights reserved.
 #  * Licensed under the MIT License. See LICENSE file in the
 #  * project root for license information.
-import numpy as np
+import os
 from datetime import datetime
 from typing import TYPE_CHECKING, Union
-import os
+
+import numpy as np
+
+from flaml.automl.spark import DataFrame, Series, pd, ps, psDataFrame, psSeries
 from flaml.automl.training_log import training_log_reader
-from flaml.automl.spark import ps, psDataFrame, psSeries, DataFrame, Series, pd
 
 try:
-    from scipy.sparse import vstack, issparse
+    from scipy.sparse import issparse, vstack
 except ImportError:
     pass
 
@@ -41,8 +43,9 @@ def load_openml_dataset(dataset_id, data_dir=None, random_state=0, dataset_forma
         y_train: A series or array of labels for training data.
         y_test:  A series or array of labels for test data.
     """
-    import openml
     import pickle
+
+    import openml
     from sklearn.model_selection import train_test_split
 
     filename = "openml_ds" + str(dataset_id) + ".pkl"
@@ -93,9 +96,10 @@ def load_openml_task(task_id, data_dir):
         y_train: A series of labels for training data.
         y_test:  A series of labels for test data.
     """
-    import openml
     import pickle
 
+    import openml
+
     task = openml.tasks.get_task(task_id)
     filename = "openml_task" + str(task_id) + ".pkl"
     filepath = os.path.join(data_dir, filename)
@@ -341,8 +345,8 @@ def fit_transform(self, X: Union[DataFrame, np.ndarray], y, task: Union[str, "Ta
                     drop = True
                 else:
                     drop = False
-                from sklearn.impute import SimpleImputer
                 from sklearn.compose import ColumnTransformer
+                from sklearn.impute import SimpleImputer
 
                 self.transformer = ColumnTransformer(
                     [