initial pass at adding Heroku chat provider

This commit is contained in:
Timothy Lowrimore 2025-07-22 17:17:05 -06:00
parent 03baf23ad1
commit 4295f3972a
5 changed files with 64 additions and 0 deletions

View file

@ -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

View file

@ -279,6 +279,7 @@ LITELLM_CHAT_PROVIDERS = [
"dashscope",
"moonshot",
"v0",
"heroku",
]
LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS = [

View file

@ -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

View 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"

View file

@ -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: