code: def completion(
      model: str,
      messages: list,
      api_base: str,
      model_response: ModelResponse,
      print_verbose: Callable,
      encoding,
      api_key,
      logging_obj,
      optional_params=None,
      litellm_params=None,
      logger_fn=None,
  ):
      headers = validate_environment(api_key)
      model = model
      prompt = ""
      for message in messages:
          if "role" in message:
              if message["role"] == "user":
                  prompt += f"{message['content']}"
              else:
                  prompt += f"{message['content']}"
          else:
              prompt += f"{message['content']}"
  
      ## Load Config
      config = litellm.AI21Config.get_config()
      for k, v in config.items():
          if (
              k not in optional_params
          ):  # completion(top_k=3) > ai21_config(top_k=3) <- allows for dynamic variables to be passed in
              optional_params[k] = v
  
      data = {
          "prompt": prompt,
          # "instruction": prompt, # some baseten models require the prompt to be passed in via the 'instruction' kwarg
          **optional_params,
      }
  
      ## LOGGING
      logging_obj.pre_call(
          input=prompt,
          api_key=api_key,
          additional_args={"complete_input_dict": data},
      )
      ## COMPLETION CALL
      response = requests.post(
          api_base + model + "/complete",
          headers=headers,
          data=json.dumps(data),
      )
      if response.status_code != 200:
          raise AI21Error(
              status_code=response.status_code, message=response.text
          )
      if (
          "stream" in optional_params
          and optional_params["stream"] == True
      ):
          return response.iter_lines()
      else:
          ## LOGGING
          logging_obj.post_call(
              input=prompt,
              api_key=api_key,
              original_response=response.text,
              additional_args={"complete_input_dict": data},
          )
          ## RESPONSE OBJECT
          completion_response = response.json()
          try:
              choices_list = []
              for idx, item in enumerate(
                  completion_response["completions"]
              ):
                  if len(item["data"]["text"]) > 0:
                      message_obj = Message(
                          content=item["data"]["text"]
                      )
                  else:
                      message_obj = Message(content=None)
                  choice_obj = Choices(
                      finish_reason=item["finishReason"]["reason"],
                      index=idx + 1,
                      message=message_obj,
                  )
                  choices_list.append(choice_obj)
              model_response["choices"] = choices_list
          except Exception as e:
              raise AI21Error(
                  message=traceback.format_exc(),
                  status_code=response.status_code,
              )
  
          ## CALCULATING USAGE - baseten charges on time, not tokens - have some mapping of cost here.
          prompt_tokens = len(encoding.encode(prompt))
          completion_tokens = len(
              encoding.encode(
                  model_response["choices"][0]["message"].get("content")
              )
          )
  
          model_response["created"] = int(time.time())
          model_response["model"] = model
          model_response["usage"] = {
              "prompt_tokens": prompt_tokens,
              "completion_tokens": completion_tokens,
              "total_tokens": prompt_tokens + completion_tokens,
          }
          return model_response
  
file_path: litellm/litellm/llms/ai21.py
imports: import os, types, traceback
  import json
  from enum import Enum
  import requests
  import time, httpx
  from typing import Callable, Optional
  from litellm.utils import ModelResponse, Choices, Message
  import litellm
label_desc: Use image node_id 2 for calling a global function with example usage: completion(model, messages, api_base, model_response, print_verbose, encoding, api_key, logging_obj, optional_params, litellm_params, logger_fn) and returns: response, model_response