docstr: If there is a prompt template, we use it. Otherwise, we use the concatenation of `context_column` (optional) and `question_column` to ask for a completion.
docstr_tok: ["If","there","is","a","prompt","template",",","we","use","it",".","Otherwise",",","we","use","the","concatenation","of","`","context_column","`","(","optional",")","and","`","question_column","`","to","ask","for","a","completion","."]
code: def predict(
      self, df: pd.DataFrame, args: Optional[Dict] = None
  ) -> pd.DataFrame:
      """
      If there is a prompt template, we use it. Otherwise, we use the concatenation of `context_column` (optional) and `question_column` to ask for a completion.
      """  # noqa
      # TODO: support for edits, embeddings and moderation
  
      pred_args = args["predict_params"] if args else {}
      args = self.model_storage.json_get("args")
      args["api_base"] = pred_args.get(
          "api_base",
          args.get(
              "api_base",
              os.environ.get("OPENAI_API_BASE", OPENAI_API_BASE),
          ),
      )
      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"]
          else:
              raise Exception(
                  f"Invalid operation mode. Please use one of {self.supported_modes}."
              )  # noqa
  
      if pred_args.get("prompt_template", False):
          base_template = pred_args[
              "prompt_template"
          ]  # override with predict-time template if available
      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":
          api_args = {
              "question_column": pred_args.get("question_column", None),
              "model": pred_args.get(
                  "model_name", "text-embedding-ada-002"
              ),
          }
          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")
  
      # Image mode
  
      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),
          }
          api_args = {
              k: v for k, v in api_args.items() if v is not None
          }  # filter out non-specified api args
          model_name = args.get("model_name", "dall-e-2")
  
          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`."
              )
  
      # 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("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)
          api_args = {
              "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)
                      ),
                  ),
              ),
              "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
          ):
              raise Exception(
                  f"Conversational modes are only available for the following models: {', '.join(self.chat_completion_models)}"
              )  # noqa
  
          if args.get("prompt_template", False):
              prompts, empty_prompt_ids = get_completed_prompts(
                  base_template, df
              )
  
          elif args.get("context_column", False):
              empty_prompt_ids = np.where(
                  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 args.get("json_struct", False):
              empty_prompt_ids = np.where(
                  df[[args["input_text"]]].isna().all(axis=1).values
              )[0]
              prompts = []
              for i in df.index:
                  if "json_struct" in df.columns:
                      if isinstance(df["json_struct"][i], str):
                          df["json_struct"][i] = json.loads(
                              df["json_struct"][i]
                          )
                      json_struct = ""
                      for ind, val in enumerate(
                          df["json_struct"][i].values()
                      ):
                          json_struct = json_struct + f"{ind}. {val}\n"
                  else:
                      json_struct = ""
                      for ind, val in enumerate(
                          args["json_struct"].values()
                      ):
                          json_struct = (
                              json_struct + f"{ind + 1}. {val}\n"
                          )
  
                  p = textwrap.dedent(
                      f"""\
                          Using text starting after 'The text is:', give exactly {len(args['json_struct'])} answers to the questions:
                          {{{{json_struct}}}}
  
                          Answers should be in the same order as the questions.
                          Each answer should start with a question number.
                          Each answer must end with new line.
                          If there is no answer to the question in the text, put a -.
                          Answers should be as short as possible, ideally 1-2 words (unless otherwise specified).
  
                          The text is:
                          {{{{{args['input_text']}}}}}
                      """
                  )
                  p = p.replace("{{json_struct}}", json_struct)
                  for column in df.columns:
                      if column == "json_struct":
                          continue
                      p = p.replace(
                          f"{{{{{column}}}}}", str(df[column][i])
                      )
                  prompts.append(p)
          elif "prompt" in args:
              empty_prompt_ids = []
              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))
              )
  
      # 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("openai", 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
      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"]])
  
      # restore json struct
      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()
                  else:
                      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)
                  }
              except Exception:
                  pred_df[args["target"]][i] = None
  
      return pred_df
  
imports: import os
  import math
  import json
  import shutil
  import tempfile
  import datetime
  import textwrap
  import subprocess
  import concurrent.futures
  from typing import Optional, Dict
  import openai
  from openai import OpenAI
  import numpy
  import pandas
  from mindsdb.utilities.hooks import before_openai_query, after_openai_query
  from mindsdb.utilities import log
  from mindsdb.integrations.libs.base import BaseMLEngine
  from mindsdb.integrations.handlers.openai_handler.helpers import retry_with_exponential_backoff, truncate_msgs_for_token_limit, get_available_models, PendingFT
  from mindsdb.integrations.handlers.openai_handler.constants import CHAT_MODELS, IMAGE_MODELS, FINETUNING_MODELS, OPENAI_API_BASE
  from mindsdb.integrations.utilities.handler_utils import get_api_key
  from mindsdb.integrations.libs.llm_utils import get_completed_prompts
label_desc: Use image node_id 4 for calling the OpenAIHandler obj's underlying member method code with example usage: obj.predict(df, args) and returns: pred_df