MindsDB 25.6.2.0__py3-none-any.whl → 25.6.3.0__py3-none-any.whl

mindsdb/integrations/handlers/openai_handler/openai_handler.py

@@ -30,7 +30,7 @@ from mindsdb.integrations.handlers.openai_handler.constants import (
     OPENAI_API_BASE,
     DEFAULT_CHAT_MODEL,
     DEFAULT_EMBEDDING_MODEL,
-    DEFAULT_IMAGE_MODEL
+    DEFAULT_IMAGE_MODEL,
 )
 from mindsdb.integrations.libs.llm.utils import get_completed_prompts
 from mindsdb.integrations.utilities.handler_utils import get_api_key
@@ -43,7 +43,7 @@ class OpenAIHandler(BaseMLEngine):
     This handler handles connection and inference with the OpenAI API.
     """
 
-    name = 'openai'
+    name = "openai"
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -51,15 +51,13 @@ class OpenAIHandler(BaseMLEngine):
         self.default_model = DEFAULT_CHAT_MODEL
         self.default_embedding_model = DEFAULT_EMBEDDING_MODEL
         self.default_image_model = DEFAULT_IMAGE_MODEL
-        self.default_mode = (
-            'default'  # can also be 'conversational' or 'conversational-full'
-        )
+        self.default_mode = "default"  # can also be 'conversational' or 'conversational-full'
         self.supported_modes = [
-            'default',
-            'conversational',
-            'conversational-full',
-            'image',
-            'embedding',
+            "default",
+            "conversational",
+            "conversational-full",
+            "image",
+            "embedding",
         ]
         self.rate_limit = 60  # requests per minute
         self.max_batch_size = 20
@@ -67,8 +65,8 @@ class OpenAIHandler(BaseMLEngine):
         self.chat_completion_models = CHAT_MODELS
         self.supported_ft_models = FINETUNING_MODELS  # base models compatible with finetuning
         # For now this are only used for handlers that inherits OpenAIHandler and don't need to override base methods
-        self.api_key_name = getattr(self, 'api_key_name', self.name)
-        self.api_base = getattr(self, 'api_base', OPENAI_API_BASE)
+        self.api_key_name = getattr(self, "api_key_name", self.name)
+        self.api_base = getattr(self, "api_base", OPENAI_API_BASE)
 
     def create_engine(self, connection_args: Dict) -> None:
         """
@@ -84,10 +82,10 @@ class OpenAIHandler(BaseMLEngine):
             None
         """
         connection_args = {k.lower(): v for k, v in connection_args.items()}
-        api_key = connection_args.get('openai_api_key')
+        api_key = connection_args.get("openai_api_key")
         if api_key is not None:
-            org = connection_args.get('api_organization')
-            api_base = connection_args.get('api_base') or os.environ.get('OPENAI_API_BASE', OPENAI_API_BASE)
+            org = connection_args.get("api_organization")
+            api_base = connection_args.get("api_base") or os.environ.get("OPENAI_API_BASE", OPENAI_API_BASE)
             client = self._get_client(api_key=api_key, base_url=api_base, org=org, args=connection_args)
             OpenAIHandler._check_client_connection(client)
 
@@ -106,13 +104,13 @@ class OpenAIHandler(BaseMLEngine):
             None
         """
         try:
-            client.models.retrieve('test')
+            client.models.retrieve("test")
         except NotFoundError:
             pass
         except AuthenticationError as e:
-            if e.body['code'] == 'invalid_api_key':
-                raise Exception('Invalid api key')
-            raise Exception(f'Something went wrong: {e}')
+            if e.body["code"] == "invalid_api_key":
+                raise Exception("Invalid api key")
+            raise Exception(f"Something went wrong: {e}")
 
     @staticmethod
     def create_validation(target: Text, args: Dict = None, **kwargs: Any) -> None:
@@ -130,41 +128,29 @@ class OpenAIHandler(BaseMLEngine):
         Returns:
             None
         """
-        if 'using' not in args:
-            raise Exception(
-                "OpenAI engine requires a USING clause! Refer to its documentation for more details."
-            )
+        if "using" not in args:
+            raise Exception("OpenAI engine requires a USING clause! Refer to its documentation for more details.")
         else:
-            args = args['using']
+            args = args["using"]
 
-        if (
-            len(
-                set(args.keys())
-                & {'question_column', 'prompt_template', 'prompt'}
-            )
-            == 0
-        ):
-            raise Exception(
-                'One of `question_column`, `prompt_template` or `prompt` is required for this engine.'
-            )
+        if len(set(args.keys()) & {"question_column", "prompt_template", "prompt"}) == 0:
+            raise Exception("One of `question_column`, `prompt_template` or `prompt` is required for this engine.")
 
         keys_collection = [
-            ['prompt_template'],
-            ['question_column', 'context_column'],
-            ['prompt', 'user_column', 'assistant_column'],
+            ["prompt_template"],
+            ["question_column", "context_column"],
+            ["prompt", "user_column", "assistant_column"],
         ]
         for keys in keys_collection:
-            if keys[0] in args and any(
-                x[0] in args for x in keys_collection if x != keys
-            ):
+            if keys[0] in args and any(x[0] in args for x in keys_collection if x != keys):
                 raise Exception(
                     textwrap.dedent(
-                        '''\
+                        """\
                     Please provide one of
                         1) a `prompt_template`
                         2) a `question_column` and an optional `context_column`
                         3) a `prompt', 'user_column' and 'assistant_column`
-                '''
+                """
                     )
                 )
 
@@ -202,11 +188,15 @@ class OpenAIHandler(BaseMLEngine):
                 f"Unknown arguments: {', '.join(unknown_args)}.\n Known arguments are: {', '.join(known_args)}"
             )
 
-        engine_storage = kwargs['handler_storage']
+        engine_storage = kwargs["handler_storage"]
         connection_args = engine_storage.get_connection_args()
-        api_key = get_api_key('openai', args, engine_storage=engine_storage)
-        api_base = args.get('api_base') or connection_args.get('api_base') or os.environ.get('OPENAI_API_BASE', OPENAI_API_BASE)
-        org = args.get('api_organization')
+        api_key = get_api_key("openai", args, engine_storage=engine_storage)
+        api_base = (
+            args.get("api_base")
+            or connection_args.get("api_base")
+            or os.environ.get("OPENAI_API_BASE", OPENAI_API_BASE)
+        )
+        org = args.get("api_organization")
         client = OpenAIHandler._get_client(api_key=api_key, base_url=api_base, org=org, args=args)
         OpenAIHandler._check_client_connection(client)
 
@@ -225,33 +215,36 @@ class OpenAIHandler(BaseMLEngine):
         Returns:
             None
         """
-        args = args['using']
-        args['target'] = target
+        args = args["using"]
+        args["target"] = target
         try:
             api_key = get_api_key(self.api_key_name, args, self.engine_storage)
             connection_args = self.engine_storage.get_connection_args()
-            api_base = args.get('api_base') or connection_args.get('api_base') or os.environ.get('OPENAI_API_BASE') or self.api_base
-            client = self._get_client(api_key=api_key, base_url=api_base, org=args.get('api_organization'), args=args)
+            api_base = (
+                args.get("api_base")
+                or connection_args.get("api_base")
+                or os.environ.get("OPENAI_API_BASE")
+                or self.api_base
+            )
+            client = self._get_client(api_key=api_key, base_url=api_base, org=args.get("api_organization"), args=args)
             available_models = get_available_models(client)
 
-            if not args.get('mode'):
-                args['mode'] = self.default_mode
-            elif args['mode'] not in self.supported_modes:
-                raise Exception(
-                    f"Invalid operation mode. Please use one of {self.supported_modes}"
-                )
+            if not args.get("mode"):
+                args["mode"] = self.default_mode
+            elif args["mode"] not in self.supported_modes:
+                raise Exception(f"Invalid operation mode. Please use one of {self.supported_modes}")
 
-            if not args.get('model_name'):
-                if args['mode'] == 'embedding':
-                    args['model_name'] = self.default_embedding_model
-                elif args['mode'] == 'image':
-                    args['model_name'] = self.default_image_model
+            if not args.get("model_name"):
+                if args["mode"] == "embedding":
+                    args["model_name"] = self.default_embedding_model
+                elif args["mode"] == "image":
+                    args["model_name"] = self.default_image_model
                 else:
-                    args['model_name'] = self.default_model
-            elif args['model_name'] not in available_models:
+                    args["model_name"] = self.default_model
+            elif (args["model_name"] not in available_models) and (args["mode"] != "embedding"):
                 raise Exception(f"Invalid model name. Please use one of {available_models}")
         finally:
-            self.model_storage.json_set('args', args)
+            self.model_storage.json_set("args", args)
 
     def predict(self, df: pd.DataFrame, args: Optional[Dict] = None) -> pd.DataFrame:
         """
@@ -269,173 +262,138 @@ class OpenAIHandler(BaseMLEngine):
         """  # noqa
         # TODO: support for edits, embeddings and moderation
 
-        pred_args = args['predict_params'] if args else {}
-        args = self.model_storage.json_get('args')
+        pred_args = args["predict_params"] if args else {}
+        args = self.model_storage.json_get("args")
         connection_args = self.engine_storage.get_connection_args()
 
-        args['api_base'] = (pred_args.get('api_base')
-                            or args.get('api_base')
-                            or connection_args.get('api_base')
-                            or os.environ.get('OPENAI_API_BASE')
-                            or self.api_base)
+        args["api_base"] = (
+            pred_args.get("api_base")
+            or args.get("api_base")
+            or connection_args.get("api_base")
+            or os.environ.get("OPENAI_API_BASE")
+            or self.api_base
+        )
 
-        if pred_args.get('api_organization'):
-            args['api_organization'] = pred_args['api_organization']
+        if pred_args.get("api_organization"):
+            args["api_organization"] = pred_args["api_organization"]
         df = df.reset_index(drop=True)
 
-        if pred_args.get('mode'):
-            if pred_args['mode'] in self.supported_modes:
-                args['mode'] = pred_args['mode']
+        if pred_args.get("mode"):
+            if pred_args["mode"] in self.supported_modes:
+                args["mode"] = pred_args["mode"]
             else:
-                raise Exception(
-                    f"Invalid operation mode. Please use one of {self.supported_modes}."
-                )  # noqa
+                raise Exception(f"Invalid operation mode. Please use one of {self.supported_modes}.")  # noqa
 
         strict_prompt_template = True
-        if pred_args.get('prompt_template', False):
-            base_template = pred_args[
-                'prompt_template'
-            ]  # override with predict-time template if available
+        if pred_args.get("prompt_template", False):
+            base_template = pred_args["prompt_template"]  # override with predict-time template if available
             strict_prompt_template = False
-        elif args.get('prompt_template', False):
-            base_template = args['prompt_template']
+        elif args.get("prompt_template", False):
+            base_template = args["prompt_template"]
         else:
             base_template = None
 
         # Embedding mode
-        if args.get('mode', self.default_mode) == 'embedding':
+        if args.get("mode", self.default_mode) == "embedding":
             api_args = {
-                'question_column': pred_args.get('question_column', None),
-                'model': pred_args.get('model_name') or args.get('model_name'),
+                "question_column": pred_args.get("question_column", None),
+                "model": pred_args.get("model_name") or args.get("model_name"),
             }
-            model_name = 'embedding'
-            if args.get('question_column'):
-                prompts = list(df[args['question_column']].apply(lambda x: str(x)))
-                empty_prompt_ids = np.where(
-                    df[[args['question_column']]].isna().all(axis=1).values
-                )[0]
+            model_name = "embedding"
+            if args.get("question_column"):
+                prompts = list(df[args["question_column"]].apply(lambda x: str(x)))
+                empty_prompt_ids = np.where(df[[args["question_column"]]].isna().all(axis=1).values)[0]
             else:
-                raise Exception('Embedding mode needs a question_column')
+                raise Exception("Embedding mode needs a question_column")
 
         # Image mode
-        elif args.get('mode', self.default_mode) == 'image':
+        elif args.get("mode", self.default_mode) == "image":
             api_args = {
-                'n': pred_args.get('n', None),
-                'size': pred_args.get('size', None),
-                'response_format': pred_args.get('response_format', None),
+                "n": pred_args.get("n", None),
+                "size": pred_args.get("size", None),
+                "response_format": pred_args.get("response_format", None),
             }
-            api_args = {
-                k: v for k, v in api_args.items() if v is not None
-            }  # filter out non-specified api args
-            model_name = pred_args.get('model_name') or args.get('model_name')
+            api_args = {k: v for k, v in api_args.items() if v is not None}  # filter out non-specified api args
+            model_name = pred_args.get("model_name") or args.get("model_name")
 
-            if args.get('question_column'):
-                prompts = list(df[args['question_column']].apply(lambda x: str(x)))
-                empty_prompt_ids = np.where(
-                    df[[args['question_column']]].isna().all(axis=1).values
-                )[0]
-            elif args.get('prompt_template'):
+            if args.get("question_column"):
+                prompts = list(df[args["question_column"]].apply(lambda x: str(x)))
+                empty_prompt_ids = np.where(df[[args["question_column"]]].isna().all(axis=1).values)[0]
+            elif args.get("prompt_template"):
                 prompts, empty_prompt_ids = get_completed_prompts(base_template, df)
             else:
-                raise Exception(
-                    'Image mode needs either `prompt_template` or `question_column`.'
-                )
+                raise Exception("Image mode needs either `prompt_template` or `question_column`.")
 
         # Chat or normal completion mode
         else:
-            if (
-                args.get('question_column', False)
-                and args['question_column'] not in df.columns
-            ):
-                raise Exception(
-                    f"This model expects a question to answer in the '{args['question_column']}' column."
-                )
+            if args.get("question_column", False) and args["question_column"] not in df.columns:
+                raise Exception(f"This model expects a question to answer in the '{args['question_column']}' column.")
 
-            if (
-                args.get('context_column', False)
-                and args['context_column'] not in df.columns
-            ):
-                raise Exception(
-                    f"This model expects context in the '{args['context_column']}' column."
-                )
+            if args.get("context_column", False) and args["context_column"] not in df.columns:
+                raise Exception(f"This model expects context in the '{args['context_column']}' column.")
 
             # API argument validation
-            model_name = args.get('model_name', self.default_model)
+            model_name = args.get("model_name", self.default_model)
             api_args = {
-                'max_tokens': pred_args.get(
-                    'max_tokens', args.get('max_tokens', self.default_max_tokens)
-                ),
-                'temperature': min(
+                "max_tokens": pred_args.get("max_tokens", args.get("max_tokens", self.default_max_tokens)),
+                "temperature": min(
                     1.0,
-                    max(
-                        0.0, pred_args.get('temperature', args.get('temperature', 0.0))
-                    ),
+                    max(0.0, pred_args.get("temperature", args.get("temperature", 0.0))),
                 ),
-                'top_p': pred_args.get('top_p', None),
-                'n': pred_args.get('n', None),
-                'stop': pred_args.get('stop', None),
-                'presence_penalty': pred_args.get('presence_penalty', None),
-                'frequency_penalty': pred_args.get('frequency_penalty', None),
-                'best_of': pred_args.get('best_of', None),
-                'logit_bias': pred_args.get('logit_bias', None),
-                'user': pred_args.get('user', None),
+                "top_p": pred_args.get("top_p", None),
+                "n": pred_args.get("n", None),
+                "stop": pred_args.get("stop", None),
+                "presence_penalty": pred_args.get("presence_penalty", None),
+                "frequency_penalty": pred_args.get("frequency_penalty", None),
+                "best_of": pred_args.get("best_of", None),
+                "logit_bias": pred_args.get("logit_bias", None),
+                "user": pred_args.get("user", None),
             }
 
-            if (
-                args.get('mode', self.default_mode) != 'default'
-                and model_name not in self.chat_completion_models
-            ):
+            if args.get("mode", self.default_mode) != "default" and model_name not in self.chat_completion_models:
                 raise Exception(
                     f"Conversational modes are only available for the following models: {', '.join(self.chat_completion_models)}"
                 )  # noqa
 
-            if args.get('prompt_template', False):
+            if args.get("prompt_template", False):
                 prompts, empty_prompt_ids = get_completed_prompts(base_template, df, strict=strict_prompt_template)
 
-            elif args.get('context_column', False):
+            elif args.get("context_column", False):
                 empty_prompt_ids = np.where(
-                    df[[args['context_column'], args['question_column']]]
-                    .isna()
-                    .all(axis=1)
-                    .values
+                    df[[args["context_column"], args["question_column"]]].isna().all(axis=1).values
                 )[0]
-                contexts = list(df[args['context_column']].apply(lambda x: str(x)))
-                questions = list(df[args['question_column']].apply(lambda x: str(x)))
-                prompts = [
-                    f'Context: {c}\nQuestion: {q}\nAnswer: '
-                    for c, q in zip(contexts, questions)
-                ]
-
-            elif 'prompt' in args:
+                contexts = list(df[args["context_column"]].apply(lambda x: str(x)))
+                questions = list(df[args["question_column"]].apply(lambda x: str(x)))
+                prompts = [f"Context: {c}\nQuestion: {q}\nAnswer: " for c, q in zip(contexts, questions)]
+
+            elif "prompt" in args:
                 empty_prompt_ids = []
-                prompts = list(df[args['user_column']])
+                prompts = list(df[args["user_column"]])
             else:
-                empty_prompt_ids = np.where(
-                    df[[args['question_column']]].isna().all(axis=1).values
-                )[0]
-                prompts = list(df[args['question_column']].apply(lambda x: str(x)))
+                empty_prompt_ids = np.where(df[[args["question_column"]]].isna().all(axis=1).values)[0]
+                prompts = list(df[args["question_column"]].apply(lambda x: str(x)))
 
             # add json struct if available
-            if args.get('json_struct', False):
+            if args.get("json_struct", False):
                 for i, prompt in enumerate(prompts):
-                    json_struct = ''
-                    if 'json_struct' in df.columns and i not in empty_prompt_ids:
+                    json_struct = ""
+                    if "json_struct" in df.columns and i not in empty_prompt_ids:
                         # if row has a specific json, we try to use it instead of the base prompt template
                         try:
-                            if isinstance(df['json_struct'][i], str):
-                                df['json_struct'][i] = json.loads(df['json_struct'][i])
-                            for ind, val in enumerate(df['json_struct'][i].values()):
-                                json_struct = json_struct + f'{ind}. {val}\n'
+                            if isinstance(df["json_struct"][i], str):
+                                df["json_struct"][i] = json.loads(df["json_struct"][i])
+                            for ind, val in enumerate(df["json_struct"][i].values()):
+                                json_struct = json_struct + f"{ind}. {val}\n"
                         except Exception:
                             pass  # if the row's json is invalid, we use the prompt template instead
 
-                    if json_struct == '':
-                        for ind, val in enumerate(args['json_struct'].values()):
-                            json_struct = json_struct + f'{ind + 1}. {val}\n'
+                    if json_struct == "":
+                        for ind, val in enumerate(args["json_struct"].values()):
+                            json_struct = json_struct + f"{ind + 1}. {val}\n"
 
                     p = textwrap.dedent(
-                        f'''\
-                            Based on the text following 'The reference text is:', assign values to the following {len(args['json_struct'])} JSON attributes:
+                        f"""\
+                            Based on the text following 'The reference text is:', assign values to the following {len(args["json_struct"])} JSON attributes:
                             {{{{json_struct}}}}
 
                             Values should follow the same order as the attributes above.
@@ -456,48 +414,51 @@ class OpenAIHandler(BaseMLEngine):
 
                             Now for the real task. The reference text is:
                             {prompt}
-                        '''
+                        """
                     )
 
-                    p = p.replace('{{json_struct}}', json_struct)
+                    p = p.replace("{{json_struct}}", json_struct)
                     prompts[i] = p
 
         # remove prompts without signal from completion queue
         prompts = [j for i, j in enumerate(prompts) if i not in empty_prompt_ids]
 
         api_key = get_api_key(self.api_key_name, args, self.engine_storage)
-        api_args = {
-            k: v for k, v in api_args.items() if v is not None
-        }  # filter out non-specified api args
+        api_args = {k: v for k, v in api_args.items() if v is not None}  # filter out non-specified api args
         completion = self._completion(model_name, prompts, api_key, api_args, args, df)
 
         # add null completion for empty prompts
         for i in sorted(empty_prompt_ids):
             completion.insert(i, None)
 
-        pred_df = pd.DataFrame(completion, columns=[args['target']])
+        pred_df = pd.DataFrame(completion, columns=[args["target"]])
 
         # restore json struct
-        if args.get('json_struct', False):
+        if args.get("json_struct", False):
             for i in pred_df.index:
                 try:
-                    if 'json_struct' in df.columns:
-                        json_keys = df['json_struct'][i].keys()
+                    if "json_struct" in df.columns:
+                        json_keys = df["json_struct"][i].keys()
                     else:
-                        json_keys = args['json_struct'].keys()
-                    responses = pred_df[args['target']][i].split('\n')
+                        json_keys = args["json_struct"].keys()
+                    responses = pred_df[args["target"]][i].split("\n")
                     responses = [x[3:] for x in responses]  # del question index
 
-                    pred_df[args['target']][i] = {
-                        key: val for key, val in zip(json_keys, responses)
-                    }
+                    pred_df[args["target"]][i] = {key: val for key, val in zip(json_keys, responses)}
                 except Exception:
-                    pred_df[args['target']][i] = None
+                    pred_df[args["target"]][i] = None
 
         return pred_df
 
     def _completion(
-        self, model_name: Text, prompts: List[Text], api_key: Text, api_args: Dict, args: Dict, df: pd.DataFrame, parallel: bool = True
+        self,
+        model_name: Text,
+        prompts: List[Text],
+        api_key: Text,
+        api_args: Dict,
+        args: Dict,
+        df: pd.DataFrame,
+        parallel: bool = True,
     ) -> List[Any]:
         """
         Handles completion for an arbitrary amount of rows using a model connected to the OpenAI API.
@@ -531,7 +492,9 @@ class OpenAIHandler(BaseMLEngine):
         """
 
         @retry_with_exponential_backoff()
-        def _submit_completion(model_name: Text, prompts: List[Text], api_args: Dict, args: Dict, df: pd.DataFrame) -> List[Text]:
+        def _submit_completion(
+            model_name: Text, prompts: List[Text], api_args: Dict, args: Dict, df: pd.DataFrame
+        ) -> List[Text]:
             """
             Submit a request to the relevant completion endpoint of the OpenAI API based on the type of task.
 
@@ -546,26 +509,22 @@ class OpenAIHandler(BaseMLEngine):
                 List[Text]: List of completions.
             """
             kwargs = {
-                'model': model_name,
+                "model": model_name,
             }
             if model_name in IMAGE_MODELS:
                 return _submit_image_completion(kwargs, prompts, api_args)
-            elif model_name == 'embedding':
+            elif model_name == "embedding":
                 return _submit_embedding_completion(kwargs, prompts, api_args)
             elif model_name in self.chat_completion_models:
                 if model_name == "gpt-3.5-turbo-instruct":
-                    return _submit_normal_completion(
-                        kwargs,
-                        prompts,
-                        api_args
-                    )
+                    return _submit_normal_completion(kwargs, prompts, api_args)
                 else:
                     return _submit_chat_completion(
                         kwargs,
                         prompts,
                         api_args,
                         df,
-                        mode=args.get('mode', 'conversational'),
+                        mode=args.get("mode", "conversational"),
                     )
             else:
                 return _submit_normal_completion(kwargs, prompts, api_args)
@@ -584,9 +543,9 @@ class OpenAIHandler(BaseMLEngine):
             after_openai_query(params, response)
 
             params2 = params.copy()
-            params2.pop('api_key', None)
-            params2.pop('user', None)
-            logger.debug(f'>>>openai call: {params2}:\n{response}')
+            params2.pop("api_key", None)
+            params2.pop("user", None)
+            logger.debug(f">>>openai call: {params2}:\n{response}")
 
         def _submit_normal_completion(kwargs: Dict, prompts: List[Text], api_args: Dict) -> List[Text]:
             """
@@ -616,11 +575,11 @@ class OpenAIHandler(BaseMLEngine):
                 """
                 tidy_comps = []
                 for c in comp.choices:
-                    if hasattr(c, 'text'):
-                        tidy_comps.append(c.text.strip('\n').strip(''))
+                    if hasattr(c, "text"):
+                        tidy_comps.append(c.text.strip("\n").strip(""))
                 return tidy_comps
 
-            kwargs['prompt'] = prompts
+            kwargs["prompt"] = prompts
             kwargs = {**kwargs, **api_args}
 
             before_openai_query(kwargs)
@@ -656,11 +615,11 @@ class OpenAIHandler(BaseMLEngine):
                 """
                 tidy_comps = []
                 for c in comp.data:
-                    if hasattr(c, 'embedding'):
+                    if hasattr(c, "embedding"):
                         tidy_comps.append([c.embedding])
                 return tidy_comps
 
-            kwargs['input'] = prompts
+            kwargs["input"] = prompts
             kwargs = {**kwargs, **api_args}
 
             before_openai_query(kwargs)
@@ -668,7 +627,9 @@ class OpenAIHandler(BaseMLEngine):
             _log_api_call(kwargs, resp)
             return resp
 
-        def _submit_chat_completion(kwargs: Dict, prompts: List[Text], api_args: Dict, df: pd.DataFrame, mode: Text = 'conversational') -> List[Text]:
+        def _submit_chat_completion(
+            kwargs: Dict, prompts: List[Text], api_args: Dict, df: pd.DataFrame, mode: Text = "conversational"
+        ) -> List[Text]:
             """
             Submit a request to the chat completion endpoint of the OpenAI API.
 
@@ -698,48 +659,42 @@ class OpenAIHandler(BaseMLEngine):
                 """
                 tidy_comps = []
                 for c in comp.choices:
-                    if hasattr(c, 'message'):
-                        tidy_comps.append(c.message.content.strip('\n').strip(''))
+                    if hasattr(c, "message"):
+                        tidy_comps.append(c.message.content.strip("\n").strip(""))
                 return tidy_comps
 
             completions = []
-            if mode != 'conversational' or 'prompt' not in args:
+            if mode != "conversational" or "prompt" not in args:
                 initial_prompt = {
                     "role": "system",
                     "content": "You are a helpful assistant. Your task is to continue the chat.",
                 }  # noqa
             else:
                 # get prompt from model
-                initial_prompt = {"role": "system", "content": args['prompt']}  # noqa
+                initial_prompt = {"role": "system", "content": args["prompt"]}  # noqa
 
-            kwargs['messages'] = [initial_prompt]
+            kwargs["messages"] = [initial_prompt]
             last_completion_content = None
 
             for pidx in range(len(prompts)):
-                if mode != 'conversational':
-                    kwargs['messages'].append(
-                        {'role': 'user', 'content': prompts[pidx]}
-                    )
+                if mode != "conversational":
+                    kwargs["messages"].append({"role": "user", "content": prompts[pidx]})
                 else:
                     question = prompts[pidx]
                     if question:
-                        kwargs['messages'].append({'role': 'user', 'content': question})
+                        kwargs["messages"].append({"role": "user", "content": question})
 
-                    assistant_column = args.get('assistant_column')
+                    assistant_column = args.get("assistant_column")
                     if assistant_column in df.columns:
                         answer = df.iloc[pidx][assistant_column]
                     else:
                         answer = None
                     if answer:
-                        kwargs['messages'].append(
-                            {'role': 'assistant', 'content': answer}
-                        )
-
-                if mode == 'conversational-full' or (
-                    mode == 'conversational' and pidx == len(prompts) - 1
-                ):
-                    kwargs['messages'] = truncate_msgs_for_token_limit(
-                        kwargs['messages'], kwargs['model'], api_args['max_tokens']
+                        kwargs["messages"].append({"role": "assistant", "content": answer})
+
+                if mode == "conversational-full" or (mode == "conversational" and pidx == len(prompts) - 1):
+                    kwargs["messages"] = truncate_msgs_for_token_limit(
+                        kwargs["messages"], kwargs["model"], api_args["max_tokens"]
                     )
                     pkwargs = {**kwargs, **api_args}
 
@@ -748,8 +703,8 @@ class OpenAIHandler(BaseMLEngine):
                     _log_api_call(pkwargs, resp)
 
                     completions.extend(resp)
-                elif mode == 'default':
-                    kwargs['messages'] = [initial_prompt] + [kwargs['messages'][-1]]
+                elif mode == "default":
+                    kwargs["messages"] = [initial_prompt] + [kwargs["messages"][-1]]
                     pkwargs = {**kwargs, **api_args}
 
                     before_openai_query(kwargs)
@@ -760,13 +715,11 @@ class OpenAIHandler(BaseMLEngine):
                 else:
                     # in "normal" conversational mode, we request completions only for the last row
                     last_completion_content = None
-                    completions.extend([''])
+                    completions.extend([""])
 
                 if last_completion_content:
                     # interleave assistant responses with user input
-                    kwargs['messages'].append(
-                        {'role': 'assistant', 'content': last_completion_content[0]}
-                    )
+                    kwargs["messages"].append({"role": "assistant", "content": last_completion_content[0]})
 
             return completions
 
@@ -799,46 +752,36 @@ class OpenAIHandler(BaseMLEngine):
                 Returns:
                     List[Text]: List of image completions as URLs or base64 encoded images.
                 """
-                return [
-                    c.url if hasattr(c, 'url') else c.b64_json
-                    for c in comp
-                ]
-
-            completions = [
-                client.images.generate(**{'prompt': p, **kwargs, **api_args}).data[0]
-                for p in prompts
-            ]
+                return [c.url if hasattr(c, "url") else c.b64_json for c in comp]
+
+            completions = [client.images.generate(**{"prompt": p, **kwargs, **api_args}).data[0] for p in prompts]
             return _tidy(completions)
 
         client = self._get_client(
             api_key=api_key,
-            base_url=args.get('api_base'),
-            org=args.pop('api_organization') if 'api_organization' in args else None,
-            args=args
+            base_url=args.get("api_base"),
+            org=args.pop("api_organization") if "api_organization" in args else None,
+            args=args,
         )
 
         try:
             # check if simple completion works
-            completion = _submit_completion(
-                model_name, prompts, api_args, args, df
-            )
+            completion = _submit_completion(model_name, prompts, api_args, args, df)
             return completion
         except Exception as e:
             # else, we get the max batch size
-            if 'you can currently request up to at most a total of' in str(e):
-                pattern = 'a total of'
-                max_batch_size = int(e[e.find(pattern) + len(pattern):].split(').')[0])
+            if "you can currently request up to at most a total of" in str(e):
+                pattern = "a total of"
+                max_batch_size = int(e[e.find(pattern) + len(pattern) :].split(").")[0])
             else:
-                max_batch_size = (
-                    self.max_batch_size
-                )  # guards against changes in the API message
+                max_batch_size = self.max_batch_size  # guards against changes in the API message
 
         if not parallel:
             completion = None
             for i in range(math.ceil(len(prompts) / max_batch_size)):
                 partial = _submit_completion(
                     model_name,
-                    prompts[i * max_batch_size: (i + 1) * max_batch_size],
+                    prompts[i * max_batch_size : (i + 1) * max_batch_size],
                     api_args,
                     args,
                     df,
@@ -846,20 +789,18 @@ class OpenAIHandler(BaseMLEngine):
                 if not completion:
                     completion = partial
                 else:
-                    completion['choices'].extend(partial['choices'])
-                    for field in ('prompt_tokens', 'completion_tokens', 'total_tokens'):
-                        completion['usage'][field] += partial['usage'][field]
+                    completion["choices"].extend(partial["choices"])
+                    for field in ("prompt_tokens", "completion_tokens", "total_tokens"):
+                        completion["usage"][field] += partial["usage"][field]
         else:
             promises = []
             with concurrent.futures.ThreadPoolExecutor() as executor:
                 for i in range(math.ceil(len(prompts) / max_batch_size)):
-                    logger.debug(
-                        f'{i * max_batch_size}:{(i+1) * max_batch_size}/{len(prompts)}'
-                    )
+                    logger.debug(f"{i * max_batch_size}:{(i + 1) * max_batch_size}/{len(prompts)}")
                     future = executor.submit(
                         _submit_completion,
                         model_name,
-                        prompts[i * max_batch_size: (i + 1) * max_batch_size],
+                        prompts[i * max_batch_size : (i + 1) * max_batch_size],
                         api_args,
                         args,
                         df,
@@ -868,9 +809,9 @@ class OpenAIHandler(BaseMLEngine):
             completion = None
             for p in promises:
                 if not completion:
-                    completion = p['choices'].result()
+                    completion = p["choices"].result()
                 else:
-                    completion.extend(p['choices'].result())
+                    completion.extend(p["choices"].result())
 
         return completion
 
@@ -886,26 +827,26 @@ class OpenAIHandler(BaseMLEngine):
         """
         # TODO: Update to use update() artifacts
 
-        args = self.model_storage.json_get('args')
+        args = self.model_storage.json_get("args")
         api_key = get_api_key(self.api_key_name, args, self.engine_storage)
-        if attribute == 'args':
-            return pd.DataFrame(args.items(), columns=['key', 'value'])
-        elif attribute == 'metadata':
-            model_name = args.get('model_name', self.default_model)
+        if attribute == "args":
+            return pd.DataFrame(args.items(), columns=["key", "value"])
+        elif attribute == "metadata":
+            model_name = args.get("model_name", self.default_model)
             try:
                 client = self._get_client(
                     api_key=api_key,
-                    base_url=args.get('api_base'),
-                    org=args.get('api_organization'),
+                    base_url=args.get("api_base"),
+                    org=args.get("api_organization"),
                     args=args,
                 )
                 meta = client.models.retrieve(model_name)
             except Exception as e:
-                meta = {'error': str(e)}
-            return pd.DataFrame(dict(meta).items(), columns=['key', 'value'])
+                meta = {"error": str(e)}
+            return pd.DataFrame(dict(meta).items(), columns=["key", "value"])
         else:
-            tables = ['args', 'metadata']
-            return pd.DataFrame(tables, columns=['tables'])
+            tables = ["args", "metadata"]
+            return pd.DataFrame(tables, columns=["tables"])
 
     def finetune(self, df: Optional[pd.DataFrame] = None, args: Optional[Dict] = None) -> None:
         """
@@ -937,14 +878,14 @@ class OpenAIHandler(BaseMLEngine):
 
         api_key = get_api_key(self.api_key_name, args, self.engine_storage)
 
-        using_args = args.pop('using') if 'using' in args else {}
+        using_args = args.pop("using") if "using" in args else {}
 
-        api_base = using_args.get('api_base', os.environ.get('OPENAI_API_BASE', OPENAI_API_BASE))
-        org = using_args.get('api_organization')
+        api_base = using_args.get("api_base", os.environ.get("OPENAI_API_BASE", OPENAI_API_BASE))
+        org = using_args.get("api_organization")
         client = self._get_client(api_key=api_key, base_url=api_base, org=org, args=args)
 
         args = {**using_args, **args}
-        prev_model_name = self.base_model_storage.json_get('args').get('model_name', '')
+        prev_model_name = self.base_model_storage.json_get("args").get("model_name", "")
 
         if prev_model_name not in self.supported_ft_models:
             # base model may be already FTed, check prefixes
@@ -956,80 +897,73 @@ class OpenAIHandler(BaseMLEngine):
                     f"This model cannot be finetuned. Supported base models are {self.supported_ft_models}."
                 )
 
-        finetune_time = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+        finetune_time = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
 
         temp_storage_path = tempfile.mkdtemp()
         temp_file_name = f"ft_{finetune_time}"
         temp_model_storage_path = f"{temp_storage_path}/{temp_file_name}.jsonl"
 
-        file_names = self._prepare_ft_jsonl(
-            df, temp_storage_path, temp_file_name, temp_model_storage_path
-        )
+        file_names = self._prepare_ft_jsonl(df, temp_storage_path, temp_file_name, temp_model_storage_path)
 
         jsons = {k: None for k in file_names.keys()}
         for split, file_name in file_names.items():
             if os.path.isfile(os.path.join(temp_storage_path, file_name)):
                 jsons[split] = client.files.create(
-                    file=open(f"{temp_storage_path}/{file_name}", "rb"),
-                    purpose='fine-tune'
+                    file=open(f"{temp_storage_path}/{file_name}", "rb"), purpose="fine-tune"
                 )
 
-        if type(jsons['train']) is openai.types.FileObject:
-            train_file_id = jsons['train'].id
+        if type(jsons["train"]) is openai.types.FileObject:
+            train_file_id = jsons["train"].id
         else:
-            train_file_id = jsons['base'].id
+            train_file_id = jsons["base"].id
 
-        if type(jsons['val']) is openai.types.FileObject:
-            val_file_id = jsons['val'].id
+        if type(jsons["val"]) is openai.types.FileObject:
+            val_file_id = jsons["val"].id
         else:
             val_file_id = None
 
         # `None` values are internally imputed by OpenAI to `null` or default values
         ft_params = {
-            'training_file': train_file_id,
-            'validation_file': val_file_id,
-            'model': self._get_ft_model_type(prev_model_name),
+            "training_file": train_file_id,
+            "validation_file": val_file_id,
+            "model": self._get_ft_model_type(prev_model_name),
         }
         ft_params = self._add_extra_ft_params(ft_params, using_args)
 
         start_time = datetime.datetime.now()
 
-        ft_stats, result_file_id = self._ft_call(ft_params, client, args.get('hour_budget', 8))
+        ft_stats, result_file_id = self._ft_call(ft_params, client, args.get("hour_budget", 8))
         ft_model_name = ft_stats.fine_tuned_model
 
         end_time = datetime.datetime.now()
         runtime = end_time - start_time
         name_extension = client.files.retrieve(file_id=result_file_id).filename
-        result_path = f'{temp_storage_path}/ft_{finetune_time}_result_{name_extension}'
+        result_path = f"{temp_storage_path}/ft_{finetune_time}_result_{name_extension}"
 
         try:
             client.files.content(file_id=result_file_id).stream_to_file(result_path)
-            if '.csv' in name_extension:
+            if ".csv" in name_extension:
                 # legacy endpoint
                 train_stats = pd.read_csv(result_path)
-                if 'validation_token_accuracy' in train_stats.columns:
-                    train_stats = train_stats[
-                        train_stats['validation_token_accuracy'].notnull()
-                    ]
-                args['ft_api_info'] = ft_stats.dict()
-                args['ft_result_stats'] = train_stats.to_dict()
-
-            elif '.json' in name_extension:
-                train_stats = pd.read_json(
-                    path_or_buf=result_path, lines=True
-                )  # new endpoint
-                args['ft_api_info'] = args['ft_result_stats'] = train_stats.to_dict()
+                if "validation_token_accuracy" in train_stats.columns:
+                    train_stats = train_stats[train_stats["validation_token_accuracy"].notnull()]
+                args["ft_api_info"] = ft_stats.dict()
+                args["ft_result_stats"] = train_stats.to_dict()
+
+            elif ".json" in name_extension:
+                train_stats = pd.read_json(path_or_buf=result_path, lines=True)  # new endpoint
+                args["ft_api_info"] = args["ft_result_stats"] = train_stats.to_dict()
 
         except Exception:
-            logger.info(f'Error retrieving fine-tuning results. Please check manually for information on job {ft_stats.id} (result file {result_file_id}).')
+            logger.info(
+                f"Error retrieving fine-tuning results. Please check manually for information on job {ft_stats.id} (result file {result_file_id})."
+            )
 
-        args['model_name'] = ft_model_name
-        args['runtime'] = runtime.total_seconds()
-        args['mode'] = self.base_model_storage.json_get('args').get(
-            'mode', self.default_mode
-        )
+        args["model_name"] = ft_model_name
+        args["runtime"] = runtime.total_seconds()
+        args["mode"] = self.base_model_storage.json_get("args").get("mode", self.default_mode)
 
-        self.model_storage.json_set('args', args)
+        self.model_storage.json_set("args", args)
         shutil.rmtree(temp_storage_path)
 
     @staticmethod
@@ -1045,7 +979,7 @@ class OpenAIHandler(BaseMLEngine):
         Returns:
             Dict: File names for the fine-tuning process.
         """
-        df.to_json(temp_model_path, orient='records', lines=True)
+        df.to_json(temp_model_path, orient="records", lines=True)
 
         # TODO avoid subprocess usage once OpenAI enables non-CLI access, or refactor to use our own LLM utils instead
         subprocess.run(
@@ -1055,7 +989,7 @@ class OpenAIHandler(BaseMLEngine):
                 "fine_tunes.prepare_data",
                 "-f",
                 temp_model_path,  # from file
-                '-q',  # quiet mode (accepts all suggestions)
+                "-q",  # quiet mode (accepts all suggestions)
             ],
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
@@ -1063,10 +997,10 @@ class OpenAIHandler(BaseMLEngine):
         )
 
         file_names = {
-            'original': f'{temp_filename}.jsonl',
-            'base': f'{temp_filename}_prepared.jsonl',
-            'train': f'{temp_filename}_prepared_train.jsonl',
-            'val': f'{temp_filename}_prepared_valid.jsonl',
+            "original": f"{temp_filename}.jsonl",
+            "base": f"{temp_filename}_prepared.jsonl",
+            "train": f"{temp_filename}_prepared_train.jsonl",
+            "val": f"{temp_filename}_prepared_valid.jsonl",
         }
         return file_names
 
@@ -1083,7 +1017,7 @@ class OpenAIHandler(BaseMLEngine):
         for model_type in self.supported_ft_models:
             if model_type in model_name.lower():
                 return model_type
-        return 'babbage-002'
+        return "babbage-002"
 
     @staticmethod
     def _add_extra_ft_params(ft_params: Dict, using_args: Dict) -> Dict:
@@ -1098,22 +1032,14 @@ class OpenAIHandler(BaseMLEngine):
             Dict: Fine-tuning parameters with extra parameters.
         """
         extra_params = {
-            'n_epochs': using_args.get('n_epochs', None),
-            'batch_size': using_args.get('batch_size', None),
-            'learning_rate_multiplier': using_args.get(
-                'learning_rate_multiplier', None
-            ),
-            'prompt_loss_weight': using_args.get('prompt_loss_weight', None),
-            'compute_classification_metrics': using_args.get(
-                'compute_classification_metrics', None
-            ),
-            'classification_n_classes': using_args.get(
-                'classification_n_classes', None
-            ),
-            'classification_positive_class': using_args.get(
-                'classification_positive_class', None
-            ),
-            'classification_betas': using_args.get('classification_betas', None),
+            "n_epochs": using_args.get("n_epochs", None),
+            "batch_size": using_args.get("batch_size", None),
+            "learning_rate_multiplier": using_args.get("learning_rate_multiplier", None),
+            "prompt_loss_weight": using_args.get("prompt_loss_weight", None),
+            "compute_classification_metrics": using_args.get("compute_classification_metrics", None),
+            "classification_n_classes": using_args.get("classification_n_classes", None),
+            "classification_positive_class": using_args.get("classification_positive_class", None),
+            "classification_betas": using_args.get("classification_betas", None),
         }
         return {**ft_params, **extra_params}
 
@@ -1137,9 +1063,7 @@ class OpenAIHandler(BaseMLEngine):
         Returns:
             Tuple[FineTuningJob, Text]: Fine-tuning stats and result file ID.
         """
-        ft_result = client.fine_tuning.jobs.create(
-            **{k: v for k, v in ft_params.items() if v is not None}
-        )
+        ft_result = client.fine_tuning.jobs.create(**{k: v for k, v in ft_params.items() if v is not None})
 
         @retry_with_exponential_backoff(
             hour_budget=hour_budget,
@@ -1158,22 +1082,22 @@ class OpenAIHandler(BaseMLEngine):
                 FineTuningJob: Fine-tuning stats.
             """
             ft_retrieved = client.fine_tuning.jobs.retrieve(fine_tuning_job_id=job_id)
-            if ft_retrieved.status in ('succeeded', 'failed', 'cancelled'):
+            if ft_retrieved.status in ("succeeded", "failed", "cancelled"):
                 return ft_retrieved
             else:
-                raise PendingFT('Fine-tuning still pending!')
+                raise PendingFT("Fine-tuning still pending!")
 
         ft_stats = _check_ft_status(ft_result.id)
 
-        if ft_stats.status != 'succeeded':
-            err_message = ft_stats.events[-1].message if hasattr(ft_stats, 'events') else 'could not retrieve!'
-            ft_status = ft_stats.status if hasattr(ft_stats, 'status') else 'N/A'
+        if ft_stats.status != "succeeded":
+            err_message = ft_stats.events[-1].message if hasattr(ft_stats, "events") else "could not retrieve!"
+            ft_status = ft_stats.status if hasattr(ft_stats, "status") else "N/A"
             raise Exception(
                 f"Fine-tuning did not complete successfully (status: {ft_status}). Error message: {err_message}"
             )  # noqa
 
         result_file_id = client.fine_tuning.jobs.retrieve(fine_tuning_job_id=ft_result.id).result_files[0]
-        if hasattr(result_file_id, 'id'):
+        if hasattr(result_file_id, "id"):
             result_file_id = result_file_id.id  # legacy endpoint
 
         return ft_stats, result_file_id
@@ -1191,11 +1115,8 @@ class OpenAIHandler(BaseMLEngine):
         Returns:
             openai.OpenAI: OpenAI client.
         """
-        if args is not None and args.get('provider') == 'azure':
+        if args is not None and args.get("provider") == "azure":
             return AzureOpenAI(
-                api_key=api_key,
-                azure_endpoint=base_url,
-                api_version=args.get('api_version'),
-                organization=org
+                api_key=api_key, azure_endpoint=base_url, api_version=args.get("api_version"), organization=org
             )
         return OpenAI(api_key=api_key, base_url=base_url, organization=org)