(Feat) Add Vertex Model Garden llama 3.1 models (#6763)

* add VertexAIModelGardenModels

* VertexAIModelGardenModels

* test_vertexai_model_garden_model_completion

* docs model garden
This commit is contained in:
Ishaan Jaff 2024-11-15 16:14:06 -08:00 committed by GitHub
parent 0f7ea14992
commit 9ba8f40bd1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 356 additions and 3 deletions

View file

@ -158,6 +158,9 @@ from .llms.vertex_ai_and_google_ai_studio.vertex_ai_partner_models.main import (
from .llms.vertex_ai_and_google_ai_studio.vertex_embeddings.embedding_handler import (
VertexEmbedding,
)
from .llms.vertex_ai_and_google_ai_studio.vertex_model_garden.main import (
VertexAIModelGardenModels,
)
from .llms.watsonx.chat.handler import WatsonXChatHandler
from .llms.watsonx.completion.handler import IBMWatsonXAI
from .types.llms.openai import (
@ -221,6 +224,7 @@ vertex_multimodal_embedding = VertexMultimodalEmbedding()
vertex_image_generation = VertexImageGeneration()
google_batch_embeddings = GoogleBatchEmbeddings()
vertex_partner_models_chat_completion = VertexAIPartnerModels()
vertex_model_garden_chat_completion = VertexAIModelGardenModels()
vertex_text_to_speech = VertexTextToSpeechAPI()
watsonxai = IBMWatsonXAI()
sagemaker_llm = SagemakerLLM()
@ -2355,6 +2359,28 @@ def completion( # type: ignore # noqa: PLR0915
api_base=api_base,
extra_headers=extra_headers,
)
elif "openai" in model:
# Vertex Model Garden - OpenAI compatible models
model_response = vertex_model_garden_chat_completion.completion(
model=model,
messages=messages,
model_response=model_response,
print_verbose=print_verbose,
optional_params=new_params,
litellm_params=litellm_params, # type: ignore
logger_fn=logger_fn,
encoding=encoding,
api_base=api_base,
vertex_location=vertex_ai_location,
vertex_project=vertex_ai_project,
vertex_credentials=vertex_credentials,
logging_obj=logging,
acompletion=acompletion,
headers=headers,
custom_prompt_dict=custom_prompt_dict,
timeout=timeout,
client=client,
)
else:
model_response = vertex_ai_non_gemini.completion(
model=model,