feat(predibase.py): add support for predibase provider

Closes https://github.com/BerriAI/litellm/issues/1253

litellm/__init__.py

@@ -71,9 +71,11 @@ maritalk_key: Optional[str] = None
 ai21_key: Optional[str] = None
 ollama_key: Optional[str] = None
 openrouter_key: Optional[str] = None
+predibase_key: Optional[str] = None
 huggingface_key: Optional[str] = None
 vertex_project: Optional[str] = None
 vertex_location: Optional[str] = None
+predibase_tenant_id: Optional[str] = None
 togetherai_api_key: Optional[str] = None
 cloudflare_api_key: Optional[str] = None
 baseten_key: Optional[str] = None
@@ -532,6 +534,7 @@ provider_list: List = [
     "xinference",
     "fireworks_ai",
     "watsonx",
+    "predibase",
     "custom",  # custom apis
 ]
 
@@ -644,6 +647,7 @@ from .utils import (
 )
 from .llms.huggingface_restapi import HuggingfaceConfig
 from .llms.anthropic import AnthropicConfig
+from .llms.predibase import PredibaseConfig
 from .llms.anthropic_text import AnthropicTextConfig
 from .llms.replicate import ReplicateConfig
 from .llms.cohere import CohereConfig

litellm/llms/huggingface_restapi.py

@@ -322,9 +322,9 @@ class Huggingface(BaseLLM):
         encoding,
         api_key,
         logging_obj,
+        optional_params: dict,
         custom_prompt_dict={},
         acompletion: bool = False,
-        optional_params=None,
         litellm_params=None,
         logger_fn=None,
     ):

litellm/llms/predibase.py

@@ -0,0 +1,417 @@
+# What is this?
+## Controller file for Predibase Integration - https://predibase.com/
+
+
+import os, types
+import json
+from enum import Enum
+import requests, copy  # type: ignore
+import time
+from typing import Callable, Optional, List, Literal, Union
+from litellm.utils import (
+    ModelResponse,
+    Usage,
+    map_finish_reason,
+    CustomStreamWrapper,
+    Message,
+    Choices,
+)
+import litellm
+from .prompt_templates.factory import prompt_factory, custom_prompt
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from .base import BaseLLM
+import httpx  # type: ignore
+
+
+class PredibaseError(Exception):
+    def __init__(
+        self,
+        status_code,
+        message,
+        request: Optional[httpx.Request] = None,
+        response: Optional[httpx.Response] = None,
+    ):
+        self.status_code = status_code
+        self.message = message
+        if request is not None:
+            self.request = request
+        else:
+            self.request = httpx.Request(
+                method="POST",
+                url="https://docs.predibase.com/user-guide/inference/rest_api",
+            )
+        if response is not None:
+            self.response = response
+        else:
+            self.response = httpx.Response(
+                status_code=status_code, request=self.request
+            )
+        super().__init__(
+            self.message
+        )  # Call the base class constructor with the parameters it needs
+
+
+class PredibaseConfig:
+    """
+    Reference:  https://docs.predibase.com/user-guide/inference/rest_api
+
+    """
+
+    adapter_id: Optional[str] = None
+    adapter_source: Optional[Literal["pbase", "hub", "s3"]] = None
+    best_of: Optional[int] = None
+    decoder_input_details: bool = True  # on by default - get the finish reason
+    details: Optional[bool] = True  # enables returning logprobs + best of
+    max_new_tokens: int = (
+        256  # openai default - requests hang if max_new_tokens not given
+    )
+    repetition_penalty: Optional[float] = None
+    return_full_text: Optional[bool] = (
+        False  # by default don't return the input as part of the output
+    )
+    seed: Optional[int] = None
+    stop: Optional[List[str]] = None
+    temperature: Optional[float] = None
+    top_k: Optional[int] = None
+    top_p: Optional[int] = None
+    truncate: Optional[int] = None
+    typical_p: Optional[float] = None
+    watermark: Optional[bool] = None
+
+    def __init__(
+        self,
+        best_of: Optional[int] = None,
+        decoder_input_details: Optional[bool] = None,
+        details: Optional[bool] = None,
+        max_new_tokens: Optional[int] = None,
+        repetition_penalty: Optional[float] = None,
+        return_full_text: Optional[bool] = None,
+        seed: Optional[int] = None,
+        stop: Optional[List[str]] = None,
+        temperature: Optional[float] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[int] = None,
+        truncate: Optional[int] = None,
+        typical_p: Optional[float] = None,
+        watermark: Optional[bool] = None,
+    ) -> None:
+        locals_ = locals()
+        for key, value in locals_.items():
+            if key != "self" and value is not None:
+                setattr(self.__class__, key, value)
+
+    @classmethod
+    def get_config(cls):
+        return {
+            k: v
+            for k, v in cls.__dict__.items()
+            if not k.startswith("__")
+            and not isinstance(
+                v,
+                (
+                    types.FunctionType,
+                    types.BuiltinFunctionType,
+                    classmethod,
+                    staticmethod,
+                ),
+            )
+            and v is not None
+        }
+
+    def get_supported_openai_params(self):
+        return ["stream", "temperature", "max_tokens", "top_p", "stop", "n"]
+
+
+class PredibaseChatCompletion(BaseLLM):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def validate_environment(self, api_key: Optional[str], user_headers: dict) -> dict:
+        if api_key is None:
+            raise ValueError(
+                "Missing Predibase API Key - A call is being made to predibase but no key is set either in the environment variables or via params"
+            )
+        headers = {
+            "content-type": "application/json",
+            "Authorization": "Bearer {}".format(api_key),
+        }
+        if user_headers is not None and isinstance(user_headers, dict):
+            headers = {**headers, **user_headers}
+        return headers
+
+    def output_parser(self, generated_text: str):
+        """
+        Parse the output text to remove any special characters. In our current approach we just check for ChatML tokens.
+
+        Initial issue that prompted this - https://github.com/BerriAI/litellm/issues/763
+        """
+        chat_template_tokens = [
+            "<|assistant|>",
+            "<|system|>",
+            "<|user|>",
+            "<s>",
+            "</s>",
+        ]
+        for token in chat_template_tokens:
+            if generated_text.strip().startswith(token):
+                generated_text = generated_text.replace(token, "", 1)
+            if generated_text.endswith(token):
+                generated_text = generated_text[::-1].replace(token[::-1], "", 1)[::-1]
+        return generated_text
+
+    def process_response(
+        self,
+        model: str,
+        response: requests.Response,
+        model_response: ModelResponse,
+        stream: bool,
+        logging_obj: litellm.utils.Logging,
+        optional_params: dict,
+        api_key: str,
+        data: dict,
+        messages: list,
+        print_verbose,
+        encoding,
+    ) -> ModelResponse:
+        ## LOGGING
+        logging_obj.post_call(
+            input=messages,
+            api_key=api_key,
+            original_response=response.text,
+            additional_args={"complete_input_dict": data},
+        )
+        print_verbose(f"raw model_response: {response.text}")
+        ## RESPONSE OBJECT
+        try:
+            completion_response = response.json()
+        except:
+            raise PredibaseError(
+                message=response.text, status_code=response.status_code
+            )
+        if "error" in completion_response:
+            raise PredibaseError(
+                message=str(completion_response["error"]),
+                status_code=response.status_code,
+            )
+        else:
+            if (
+                not isinstance(completion_response, dict)
+                or "generated_text" not in completion_response
+            ):
+                raise PredibaseError(
+                    status_code=422,
+                    message=f"response is not in expected format - {completion_response}",
+                )
+
+            if len(completion_response["generated_text"]) > 0:
+                model_response["choices"][0]["message"]["content"] = self.output_parser(
+                    completion_response["generated_text"]
+                )
+            ## GETTING LOGPROBS + FINISH REASON
+            if (
+                "details" in completion_response
+                and "tokens" in completion_response["details"]
+            ):
+                model_response.choices[0].finish_reason = completion_response[
+                    "details"
+                ]["finish_reason"]
+                sum_logprob = 0
+                for token in completion_response[0]["details"]["tokens"]:
+                    if token["logprob"] != None:
+                        sum_logprob += token["logprob"]
+                model_response["choices"][0][
+                    "message"
+                ]._logprob = (
+                    sum_logprob  # [TODO] move this to using the actual logprobs
+                )
+            if "best_of" in optional_params and optional_params["best_of"] > 1:
+                if (
+                    "details" in completion_response[0]
+                    and "best_of_sequences" in completion_response[0]["details"]
+                ):
+                    choices_list = []
+                    for idx, item in enumerate(
+                        completion_response[0]["details"]["best_of_sequences"]
+                    ):
+                        sum_logprob = 0
+                        for token in item["tokens"]:
+                            if token["logprob"] != None:
+                                sum_logprob += token["logprob"]
+                        if len(item["generated_text"]) > 0:
+                            message_obj = Message(
+                                content=self.output_parser(item["generated_text"]),
+                                logprobs=sum_logprob,
+                            )
+                        else:
+                            message_obj = Message(content=None)
+                        choice_obj = Choices(
+                            finish_reason=item["finish_reason"],
+                            index=idx + 1,
+                            message=message_obj,
+                        )
+                        choices_list.append(choice_obj)
+                    model_response["choices"].extend(choices_list)
+
+        ## CALCULATING USAGE
+        prompt_tokens = 0
+        try:
+            prompt_tokens = len(
+                encoding.encode(model_response["choices"][0]["message"]["content"])
+            )  ##[TODO] use a model-specific tokenizer here
+        except:
+            # this should remain non blocking we should not block a response returning if calculating usage fails
+            pass
+        output_text = model_response["choices"][0]["message"].get("content", "")
+        if output_text is not None and len(output_text) > 0:
+            completion_tokens = 0
+            try:
+                completion_tokens = len(
+                    encoding.encode(
+                        model_response["choices"][0]["message"].get("content", "")
+                    )
+                )  ##[TODO] use a model-specific tokenizer
+            except:
+                # this should remain non blocking we should not block a response returning if calculating usage fails
+                pass
+        else:
+            completion_tokens = 0
+
+        total_tokens = prompt_tokens + completion_tokens
+
+        model_response["created"] = int(time.time())
+        model_response["model"] = model
+        usage = Usage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=total_tokens,
+        )
+        model_response.usage = usage  # type: ignore
+        return model_response
+
+    def completion(
+        self,
+        model: str,
+        messages: list,
+        api_base: str,
+        custom_prompt_dict: dict,
+        model_response: ModelResponse,
+        print_verbose: Callable,
+        encoding,
+        api_key: str,
+        logging_obj,
+        optional_params: dict,
+        tenant_id: str,
+        acompletion=None,
+        litellm_params=None,
+        logger_fn=None,
+        headers: dict = {},
+    ):
+        headers = self.validate_environment(api_key, headers)
+        completion_url = ""
+        input_text = ""
+        base_url = "https://serving.app.predibase.com"
+        if "https" in model:
+            completion_url = model
+        elif api_base:
+            base_url = api_base
+        elif "PREDIBASE_API_BASE" in os.environ:
+            base_url = os.getenv("PREDIBASE_API_BASE", "")
+
+        completion_url = f"{base_url}/{tenant_id}/deployments/v2/llms/{model}/generate"
+
+        if model in custom_prompt_dict:
+            # check if the model has a registered custom prompt
+            model_prompt_details = custom_prompt_dict[model]
+            prompt = custom_prompt(
+                role_dict=model_prompt_details["roles"],
+                initial_prompt_value=model_prompt_details["initial_prompt_value"],
+                final_prompt_value=model_prompt_details["final_prompt_value"],
+                messages=messages,
+            )
+        else:
+            prompt = prompt_factory(model=model, messages=messages)
+
+        ## Load Config
+        config = litellm.PredibaseConfig.get_config()
+        for k, v in config.items():
+            if (
+                k not in optional_params
+            ):  # completion(top_k=3) > anthropic_config(top_k=3) <- allows for dynamic variables to be passed in
+                optional_params[k] = v
+
+        data = {
+            "inputs": prompt,
+            "parameters": optional_params,
+        }
+        if optional_params.get("stream") and optional_params["stream"] == True:
+            data["stream"] = True
+        input_text = prompt
+        ## LOGGING
+        logging_obj.pre_call(
+            input=input_text,
+            api_key=api_key,
+            additional_args={
+                "complete_input_dict": data,
+                "headers": headers,
+                "api_base": completion_url,
+                "acompletion": acompletion,
+            },
+        )
+        ## COMPLETION CALL
+        if acompletion is True:
+            ### ASYNC STREAMING
+            if optional_params.get("stream", False):
+                return self.async_streaming(logging_obj=logging_obj, api_base=completion_url, data=data, headers=headers, model_response=model_response, model=model, timeout=timeout)  # type: ignore
+            else:
+                ### ASYNC COMPLETION
+                return self.acompletion(api_base=completion_url, data=data, headers=headers, model_response=model_response, task=task, encoding=encoding, input_text=input_text, model=model, optional_params=optional_params, timeout=timeout)  # type: ignore
+
+        ### SYNC STREAMING
+        if "stream" in optional_params and optional_params["stream"] == True:
+            response = requests.post(
+                completion_url,
+                headers=headers,
+                data=json.dumps(data),
+                stream=optional_params["stream"],
+            )
+            return response.iter_lines()
+        ### SYNC COMPLETION
+        else:
+            payload = json.dumps(
+                {
+                    "inputs": "What is your name?",
+                    "parameters": {"max_new_tokens": 20, "temperature": 0.1},
+                }
+                # data
+            )
+            response = requests.post(
+                url="https://serving.app.predibase.com/c4768f95/deployments/v2/llms/llama-3-8b-instruct/generate",
+                headers=headers,
+                data=payload,
+            )
+
+        return self.process_response(
+            model=model,
+            response=response,
+            model_response=model_response,
+            stream=optional_params.get("stream", False),
+            logging_obj=logging_obj,  # type: ignore
+            optional_params=optional_params,
+            api_key=api_key,
+            data=data,
+            messages=messages,
+            print_verbose=print_verbose,
+            encoding=encoding,
+        )
+
+    async def async_completion(self):
+        pass
+
+    async def async_streaming(self):
+        pass
+
+    def streaming(self):
+        pass
+
+    def embedding(self, *args, **kwargs):
+        pass

litellm/main.py

@@ -74,6 +74,7 @@ from .llms.azure_text import AzureTextCompletion
 from .llms.anthropic import AnthropicChatCompletion
 from .llms.anthropic_text import AnthropicTextCompletion
 from .llms.huggingface_restapi import Huggingface
+from .llms.predibase import PredibaseChatCompletion
 from .llms.prompt_templates.factory import (
     prompt_factory,
     custom_prompt,
@@ -110,6 +111,7 @@ anthropic_text_completions = AnthropicTextCompletion()
 azure_chat_completions = AzureChatCompletion()
 azure_text_completions = AzureTextCompletion()
 huggingface = Huggingface()
+predibase_chat_completions = PredibaseChatCompletion()
 ####### COMPLETION ENDPOINTS ################
 
 
@@ -1785,6 +1787,58 @@ def completion(
                 )
                 return response
             response = model_response
+        elif custom_llm_provider == "predibase":
+            tenant_id = (
+                optional_params.pop("tenant_id", None)
+                or optional_params.pop("predibase_tenant_id", None)
+                or litellm.predibase_tenant_id
+                or get_secret("PREDIBASE_TENANT_ID")
+            )
+
+            api_base = (
+                optional_params.pop("api_base", None)
+                or optional_params.pop("base_url", None)
+                or litellm.api_base
+                or get_secret("PREDIBASE_API_BASE")
+            )
+
+            api_key = (
+                api_key
+                or litellm.api_key
+                or litellm.predibase_key
+                or get_secret("PREDIBASE_API_KEY")
+            )
+
+            model_response = predibase_chat_completions.completion(
+                model=model,
+                messages=messages,
+                model_response=model_response,
+                print_verbose=print_verbose,
+                optional_params=optional_params,
+                litellm_params=litellm_params,
+                logger_fn=logger_fn,
+                encoding=encoding,
+                logging_obj=logging,
+                acompletion=acompletion,
+                api_base=api_base,
+                custom_prompt_dict=custom_prompt_dict,
+                api_key=api_key,
+                tenant_id=tenant_id,
+            )
+
+            if (
+                "stream" in optional_params
+                and optional_params["stream"] == True
+                and acompletion == False
+            ):
+                response = CustomStreamWrapper(
+                    model_response,
+                    model,
+                    custom_llm_provider="predibase",
+                    logging_obj=logging,
+                )
+                return response
+            response = model_response
         elif custom_llm_provider == "ai21":
             custom_llm_provider = "ai21"
             ai21_key = (

litellm/tests/test_completion.py

@@ -85,6 +85,29 @@ def test_completion_azure_command_r():
         pytest.fail(f"Error occurred: {e}")
 
 
+@pytest.mark.skip(reason="local test")
+def test_completion_predibase():
+    try:
+        litellm.set_verbose = True
+
+        response = completion(
+            model="predibase/llama-3-8b-instruct",
+            tenant_id="c4768f95",
+            api_base="https://serving.app.predibase.com",
+            api_key=os.getenv("PREDIBASE_API_KEY"),
+            messages=[{"role": "user", "content": "What is the meaning of life?"}],
+        )
+
+        print(response)
+    except litellm.Timeout as e:
+        pass
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+
+# test_completion_predibase()
+
+
 def test_completion_claude():
     litellm.set_verbose = True
     litellm.cache = None

litellm/utils.py

@@ -369,7 +369,7 @@ class ChatCompletionMessageToolCall(OpenAIObject):
 class Message(OpenAIObject):
     def __init__(
         self,
-        content="default",
+        content: Optional[str] = "default",
         role="assistant",
         logprobs=None,
         function_call=None,