From 17482ded74c9a9a0f97f84b23b23ed64bbee9d49 Mon Sep 17 00:00:00 2001 From: Tom Usher Date: Wed, 19 Jun 2024 14:09:20 +0100 Subject: [PATCH] Return output_vector_size in get_model_info --- litellm/types/utils.py | 1 + litellm/utils.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 142eef300..aeb898ef6 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -53,6 +53,7 @@ class ModelInfo(TypedDict, total=False): output_cost_per_image: Optional[float] output_cost_per_video_per_second: Optional[float] output_cost_per_audio_per_second: Optional[float] + output_vector_size: Optional[int] litellm_provider: Required[str] mode: Required[ Literal[ diff --git a/litellm/utils.py b/litellm/utils.py index ae4c343ba..d31c0a579 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4399,6 +4399,7 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod output_cost_per_token_above_128k_tokens=_model_info.get( "output_cost_per_token_above_128k_tokens", None ), + output_vector_size=_model_info.get("output_vector_size", None), litellm_provider=_model_info.get( "litellm_provider", custom_llm_provider ), @@ -4433,6 +4434,7 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod output_cost_per_token_above_128k_tokens=_model_info.get( "output_cost_per_token_above_128k_tokens", None ), + output_vector_size=_model_info.get("output_vector_size", None), litellm_provider=_model_info.get( "litellm_provider", custom_llm_provider ), @@ -4467,6 +4469,7 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod output_cost_per_token_above_128k_tokens=_model_info.get( "output_cost_per_token_above_128k_tokens", None ), + output_vector_size=_model_info.get("output_vector_size", None), litellm_provider=_model_info.get( "litellm_provider", custom_llm_provider ),