scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

refactor: replace dbrx with 'openai_like'

Browse source
This commit is contained in:
Krrish Dholakia 2024-12-02 20:57:46 -08:00
parent 5438c7d214
commit 2fda90742f
3 changed files with 7 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/sagemaker/sagemaker.py
View file

@ -273,9 +273,9 @@ class SagemakerLLM(BaseAWSLLM):
model_id = optional_params.get("model_id", None)
if use_messages_api is True:
from litellm.llms.databricks.chat.handler import DatabricksChatCompletion
from litellm.llms.openai_like.chat.handler import OpenAILikeChatHandler
openai_like_chat_completions = DatabricksChatCompletion()
openai_like_chat_completions = OpenAILikeChatHandler()
inference_params["stream"] = True if stream is True else False
_data: Dict[str, Any] = {
"model": model,

8
litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_partner_models/main.py
View file

@ -90,14 +90,13 @@ class VertexAIPartnerModels(VertexBase):
from google.cloud import aiplatform
from litellm.llms.anthropic.chat import AnthropicChatCompletion
from litellm.llms.databricks.chat.handler import DatabricksChatCompletion
from litellm.llms.OpenAI.openai import OpenAIChatCompletion
from litellm.llms.openai_like.chat.handler import OpenAILikeChatHandler
from litellm.llms.text_completion_codestral import CodestralTextCompletion
from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import (
VertexLLM,
)
except Exception:
raise VertexAIError(
status_code=400,
message="""vertexai import failed please run `pip install -U "google-cloud-aiplatform>=1.38"`""",
@ -120,7 +119,7 @@ class VertexAIPartnerModels(VertexBase):
custom_llm_provider="vertex_ai",
)
openai_like_chat_completions = DatabricksChatCompletion()
openai_like_chat_completions = OpenAILikeChatHandler()
codestral_fim_completions = CodestralTextCompletion()
anthropic_chat_completions = AnthropicChatCompletion()
@ -133,10 +132,8 @@ class VertexAIPartnerModels(VertexBase):
partner = VertexPartnerProvider.llama
elif "mistral" in model or "codestral" in model:
partner = VertexPartnerProvider.mistralai
optional_params["custom_endpoint"] = True
elif "jamba" in model:
partner = VertexPartnerProvider.ai21
optional_params["custom_endpoint"] = True
elif "claude" in model:
partner = VertexPartnerProvider.claude
@ -233,6 +230,7 @@ class VertexAIPartnerModels(VertexBase):
timeout=timeout,
encoding=encoding,
custom_llm_provider="vertex_ai",
custom_endpoint=True,
)
except Exception as e:

4
litellm/llms/vertex_ai_and_google_ai_studio/vertex_model_garden/main.py
View file

@ -76,8 +76,8 @@ class VertexAIModelGardenModels(VertexBase):
from google.cloud import aiplatform
from litellm.llms.anthropic.chat import AnthropicChatCompletion
from litellm.llms.databricks.chat.handler import DatabricksChatCompletion
from litellm.llms.OpenAI.openai import OpenAIChatCompletion
from litellm.llms.openai_like.chat.handler import OpenAILikeChatHandler
from litellm.llms.text_completion_codestral import CodestralTextCompletion
from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import (
VertexLLM,
@ -106,7 +106,7 @@ class VertexAIModelGardenModels(VertexBase):
custom_llm_provider="vertex_ai",
)
openai_like_chat_completions = DatabricksChatCompletion()
openai_like_chat_completions = OpenAILikeChatHandler()
## CONSTRUCT API BASE
stream: bool = optional_params.get("stream", False) or False