Merge pull request #4279 from tomusher/feature/return-output-vector-s…

…ize-in-modelinfo Return output_vector_size in get_model_info
BerriAI · Jul 5, 2024 · 1807c49 · 1807c49
2 parents ffcf9e6 + 0721e95
commit 1807c49
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/litellm/types/utils.py b/litellm/types/utils.py
@@ -61,6 +61,7 @@ class ModelInfo(TypedDict, total=False):
         float
     ]  # only for vertex ai models
     output_cost_per_image: Optional[float]
+    output_vector_size: Optional[int]
     output_cost_per_video_per_second: Optional[float]  # only for vertex ai models
     output_cost_per_audio_per_second: Optional[float]  # only for vertex ai models
     litellm_provider: Required[str]

diff --git a/litellm/utils.py b/litellm/utils.py
@@ -4685,6 +4685,7 @@ def _get_max_position_embeddings(model_name):
                 output_cost_per_character_above_128k_tokens=_model_info.get(
                     "output_cost_per_character_above_128k_tokens", None
                 ),
+                output_vector_size=_model_info.get("output_vector_size", None),
                 litellm_provider=_model_info.get(
                     "litellm_provider", custom_llm_provider
                 ),