(fix) use litellm.vertex_vision_models

This commit is contained in:
ishaan-jaff 2023-12-16 18:39:40 +05:30
parent f3ebfb0517
commit efe8b75200

View file

@ -205,7 +205,7 @@ def completion(
llm_model = GenerativeModel(model)
mode = ""
request_str += f"llm_model = GenerativeModel({model})\n"
elif model in ["gemini-pro-vision"]:
elif model in litellm.vertex_vision_models:
llm_model = GenerativeModel(model)
request_str += f"llm_model = GenerativeModel({model})\n"
mode = "vision"