baseten client mapping

litellm/llms/baseten.py

@@ -1,11 +1,11 @@
-import os, json
+import os
+import json
 from enum import Enum
 import requests
 import time
 from typing import Callable
 from litellm.utils import ModelResponse
 
-
 class BasetenError(Exception):
     def __init__(self, status_code, message):
         self.status_code = status_code
@@ -14,41 +14,30 @@ class BasetenError(Exception):
             self.message
         )  # Call the base class constructor with the parameters it needs
 
-
-class BasetenLLM:
-    def __init__(self, encoding, logging_obj, api_key=None):
-        self.encoding = encoding
-        self.completion_url_fragment_1 = "https://app.baseten.co/models/"
-        self.completion_url_fragment_2 = "/predict"
-        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:
-            raise ValueError(
-                "Missing Baseten API Key - A call is being made to baseten but no key is set either in the environment variables or via params"
-            )
-        self.api_key = api_key
-        self.headers = {
+def validate_environment(api_key):
+    headers = {
         "accept": "application/json",
         "content-type": "application/json",
-            "Authorization": "Api-Key " + self.api_key,
     }
+    if api_key:
+        headers["Authorization"] = f"Api-Key {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)
+    completion_url_fragment_1 = "https://app.baseten.co/models/"
+    completion_url_fragment_2 = "/predict"
     model = model
     prompt = ""
     for message in messages:
@@ -60,24 +49,22 @@ class BasetenLLM:
         else:
             prompt += f"{message['content']}"
     data = {
-            # "prompt": prompt,
-            "inputs": prompt, # in case it's a TGI deployed model
-            # "instruction": prompt, # some baseten models require the prompt to be passed in via the 'instruction' kwarg
-            # **optional_params,
+        "inputs": prompt,
+        "prompt": prompt,
         "parameters": optional_params,
         "stream": True if "stream" in optional_params and optional_params["stream"] == True else False
     }
 
     ## 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,
+        completion_url_fragment_1 + model + completion_url_fragment_2,
+        headers=headers,
         data=json.dumps(data),
         stream=True if "stream" in optional_params and optional_params["stream"] == True else False
     )
@@ -85,9 +72,9 @@ class BasetenLLM:
         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},
             )
@@ -141,9 +128,9 @@ class BasetenLLM:
                 )
 
         ## CALCULATING USAGE - baseten charges on time, not tokens - have some mapping of cost here.
-            prompt_tokens = len(self.encoding.encode(prompt))
+        prompt_tokens = len(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()
@@ -155,7 +142,6 @@ class BasetenLLM:
         }
         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

@@ -26,7 +26,7 @@ from .llms import sagemaker
 from .llms import bedrock
 from .llms import huggingface_restapi
 from .llms import aleph_alpha
-from .llms.baseten import BasetenLLM
+from .llms import baseten
 import tiktoken
 from concurrent.futures import ThreadPoolExecutor
 
@@ -751,10 +751,8 @@ def completion(
             baseten_key = (
                 api_key or litellm.baseten_key or os.environ.get("BASETEN_API_KEY")
             )
-            baseten_client = BasetenLLM(
-                encoding=encoding, api_key=baseten_key, logging_obj=logging
-            )
-            model_response = baseten_client.completion(
+
+            model_response = baseten.completion(
                 model=model,
                 messages=messages,
                 model_response=model_response,
@@ -762,6 +760,9 @@ def completion(
                 optional_params=optional_params,
                 litellm_params=litellm_params,
                 logger_fn=logger_fn,
+                encoding=encoding, 
+                api_key=baseten_key, 
+                logging_obj=logging
             )
             if inspect.isgenerator(model_response) or ("stream" in optional_params and optional_params["stream"] == True):
                 # don't try to access stream object,