From 186c0ec77bec2900a184d8628ff1457dd25b114e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 9 May 2024 16:39:43 -0700 Subject: [PATCH] feat(predibase.py): add support for predibase provider Closes https://github.com/BerriAI/litellm/issues/1253 --- litellm/__init__.py | 4 + litellm/llms/huggingface_restapi.py | 2 +- litellm/llms/predibase.py | 417 ++++++++++++++++++++++++++++ litellm/main.py | 54 ++++ litellm/tests/test_completion.py | 23 ++ litellm/utils.py | 2 +- 6 files changed, 500 insertions(+), 2 deletions(-) create mode 100644 litellm/llms/predibase.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 4f72504f69..ccf3657fe0 100644 --- a/litellm/__init__.py +++ b/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 diff --git a/litellm/llms/huggingface_restapi.py b/litellm/llms/huggingface_restapi.py index 2937732890..b250f30138 100644 --- a/litellm/llms/huggingface_restapi.py +++ b/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, ): diff --git a/litellm/llms/predibase.py b/litellm/llms/predibase.py new file mode 100644 index 0000000000..728a98b04b --- /dev/null +++ b/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|>", + "", + "", + ] + 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 diff --git a/litellm/main.py b/litellm/main.py index 5ab3fd7c45..f634fd16d9 100644 --- a/litellm/main.py +++ b/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 = ( diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 32b65faea5..7f0977b155 100644 --- a/litellm/tests/test_completion.py +++ b/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 diff --git a/litellm/utils.py b/litellm/utils.py index 6da2960389..7ccb5e8ff9 100644 --- a/litellm/utils.py +++ b/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,