clean out AI21 Init Client calls

litellm/llms/ai21.py

@@ -1,4 +1,5 @@
-import os, json
+import os
+import json
 from enum import Enum
 import requests
 import time
@@ -13,43 +14,31 @@ class AI21Error(Exception):
             self.message
         )  # Call the base class constructor with the parameters it needs
 
-
-class AI21LLM:
-    def __init__(
-        self, encoding, logging_obj, api_key=None
-    ):
-        self.encoding = encoding
-        self.completion_url_fragment_1 = "https://api.ai21.com/studio/v1/"
-        self.completion_url_fragment_2 = "/complete"
-        self.api_key = api_key
-        self.logging_obj = logging_obj
-        self.validate_environment(api_key=api_key)
-
-    def validate_environment(
-        self, api_key
-    ):  # set up the environment required to run the model
-        # set the api key
-        if self.api_key == None:
+def validate_environment(api_key):
+    if api_key is None:
         raise ValueError(
             "Missing AI21 API Key - A call is being made to ai21 but no key is set either in the environment variables or via params"
         )
-        self.api_key = api_key
-        self.headers = {
+    headers = {
         "accept": "application/json",
         "content-type": "application/json",
-            "Authorization": "Bearer " + self.api_key,
+        "Authorization": "Bearer " + api_key,
     }
+    return headers
 
 def completion(
-        self,
     model: str,
     messages: list,
     model_response: ModelResponse,
     print_verbose: Callable,
+    encoding,
+    api_key,
+    logging_obj,
     optional_params=None,
     litellm_params=None,
     logger_fn=None,
-    ):  # logic for parsing in - calling - parsing out model completion calls
+):
+    headers = validate_environment(api_key)
     model = model
     prompt = ""
     for message in messages:
@@ -71,22 +60,22 @@ class AI21LLM:
     }
 
     ## LOGGING
-        self.logging_obj.pre_call(
+    logging_obj.pre_call(
             input=prompt,
-            api_key=self.api_key,
+            api_key=api_key,
             additional_args={"complete_input_dict": data},
         )
     ## COMPLETION CALL
     response = requests.post(
-            self.completion_url_fragment_1 + model + self.completion_url_fragment_2, headers=self.headers, data=json.dumps(data)
+        "https://api.ai21.com/studio/v1/" + model + "/complete", headers=headers, data=json.dumps(data)
     )
     if "stream" in optional_params and optional_params["stream"] == True:
         return response.iter_lines()
     else:
         ## LOGGING
-            self.logging_obj.post_call(
+        logging_obj.post_call(
                 input=prompt,
-                api_key=self.api_key,
+                api_key=api_key,
                 original_response=response.text,
                 additional_args={"complete_input_dict": data},
             )
@@ -106,10 +95,10 @@ class AI21LLM:
 
         ## CALCULATING USAGE - baseten charges on time, not tokens - have some mapping of cost here. 
         prompt_tokens = len(
-                self.encoding.encode(prompt)
+            encoding.encode(prompt)
         ) 
         completion_tokens = len(
-                self.encoding.encode(model_response["choices"][0]["message"]["content"])
+            encoding.encode(model_response["choices"][0]["message"]["content"])
         )
 
         model_response["created"] = time.time()
@@ -121,7 +110,6 @@ class AI21LLM:
         }
         return model_response
 
-    def embedding(
-        self,
-    ):  # logic for parsing in - calling - parsing out model embedding calls
+def embedding():
+    # logic for parsing in - calling - parsing out model embedding calls
     pass

litellm/main.py

@@ -21,9 +21,9 @@ from litellm.utils import (
 )
 from .llms import anthropic
 from .llms import together_ai
+from .llms import ai21
 from .llms.huggingface_restapi import HuggingfaceRestAPILLM
 from .llms.baseten import BasetenLLM
-from .llms.ai21 import AI21LLM
 from .llms.aleph_alpha import AlephAlphaLLM
 import tiktoken
 from concurrent.futures import ThreadPoolExecutor
@@ -658,11 +658,7 @@ def completion(
                 or litellm.ai21_key
                 or os.environ.get("AI21_API_KEY")
             )            
-            ai21_client = AI21LLM(
-                encoding=encoding, api_key=ai21_key, logging_obj=logging
-            )
-            
-            model_response = ai21_client.completion(
+            model_response = ai21.completion(
                 model=model,
                 messages=messages,
                 model_response=model_response,
@@ -670,6 +666,9 @@ def completion(
                 optional_params=optional_params,
                 litellm_params=litellm_params,
                 logger_fn=logger_fn,
+                encoding=encoding,
+                api_key=ai21_key,
+                logging_obj=logging
             )
             
             if "stream" in optional_params and optional_params["stream"] == True: