mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
initial pass at adding Heroku chat provider
This commit is contained in:
parent
03baf23ad1
commit
4295f3972a
5 changed files with 64 additions and 0 deletions
|
|
@ -216,6 +216,7 @@ nlp_cloud_key: Optional[str] = None
|
|||
novita_api_key: Optional[str] = None
|
||||
snowflake_key: Optional[str] = None
|
||||
nebius_key: Optional[str] = None
|
||||
heroku_key: Optional[str] = None
|
||||
common_cloud_provider_auth_params: dict = {
|
||||
"params": ["project", "region_name", "token"],
|
||||
"providers": ["vertex_ai", "bedrock", "watsonx", "azure", "vertex_ai_beta"],
|
||||
|
|
@ -1151,6 +1152,7 @@ from .llms.azure.azure import (
|
|||
AzureOpenAIAssistantsAPIConfig,
|
||||
)
|
||||
|
||||
from .llms.heroku.chat.transformation import HerokuChatConfig
|
||||
from .llms.azure.chat.gpt_transformation import AzureOpenAIConfig
|
||||
from .llms.azure.completion.transformation import AzureOpenAITextConfig
|
||||
from .llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig
|
||||
|
|
|
|||
|
|
@ -279,6 +279,7 @@ LITELLM_CHAT_PROVIDERS = [
|
|||
"dashscope",
|
||||
"moonshot",
|
||||
"v0",
|
||||
"heroku",
|
||||
]
|
||||
|
||||
LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS = [
|
||||
|
|
|
|||
|
|
@ -350,6 +350,8 @@ def get_llm_provider( # noqa: PLR0915
|
|||
# bytez models
|
||||
elif model.startswith("bytez/"):
|
||||
custom_llm_provider = "bytez"
|
||||
elif model.startswith("heroku/"):
|
||||
custom_llm_provider = "heroku"
|
||||
if not custom_llm_provider:
|
||||
if litellm.suppress_debug_info is False:
|
||||
print() # noqa
|
||||
|
|
|
|||
28
litellm/llms/heroku/chat/transformation.py
Normal file
28
litellm/llms/heroku/chat/transformation.py
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
from typing import Optional, List, Union
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
||||
class HerokuChatConfig(BaseConfig):
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
headers.update({"Authorization": f"Bearer {api_key}"})
|
||||
return headers
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
api_key: Optional[str],
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
stream: Optional[bool] = None,
|
||||
) -> str:
|
||||
return f"{api_base}/v1/chat/completions"
|
||||
|
|
@ -151,6 +151,7 @@ from .llms.custom_llm import CustomLLM, custom_chat_llm_router
|
|||
from .llms.databricks.embed.handler import DatabricksEmbeddingHandler
|
||||
from .llms.deprecated_providers import aleph_alpha, palm
|
||||
from .llms.groq.chat.handler import GroqChatCompletion
|
||||
from .llms.heroku.chat.transformation import HerokuChatConfig
|
||||
from .llms.huggingface.embedding.handler import HuggingFaceEmbedding
|
||||
from .llms.nlp_cloud.chat.handler import completion as nlp_cloud_chat_completion
|
||||
from .llms.ollama.completion import handler as ollama
|
||||
|
|
@ -254,6 +255,7 @@ base_llm_http_handler = BaseLLMHTTPHandler()
|
|||
base_llm_aiohttp_handler = BaseLLMAIOHTTPHandler()
|
||||
sagemaker_chat_completion = SagemakerChatHandler()
|
||||
bytez_transformation = BytezChatConfig()
|
||||
heroku_transformation = HerokuChatConfig()
|
||||
####### COMPLETION ENDPOINTS ################
|
||||
|
||||
|
||||
|
|
@ -1768,6 +1770,35 @@ def completion( # type: ignore # noqa: PLR0915
|
|||
additional_args={"headers": headers},
|
||||
)
|
||||
raise e
|
||||
elif custom_llm_provider == "heroku":
|
||||
try:
|
||||
response = base_llm_http_handler.completion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
headers=headers,
|
||||
model_response=model_response,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
acompletion=acompletion,
|
||||
logging_obj=logging,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
timeout=timeout,
|
||||
client=client,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
encoding=encoding,
|
||||
stream=stream,
|
||||
provider_config=provider_config,
|
||||
)
|
||||
except Exception as e:
|
||||
logging.post_call(
|
||||
input=messages,
|
||||
api_key=api_key,
|
||||
original_response=str(e),
|
||||
additional_args={"headers": headers},
|
||||
)
|
||||
raise e
|
||||
|
||||
elif custom_llm_provider == "xai":
|
||||
## COMPLETION CALL
|
||||
try:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue