Return output_vector_size in get_model_info

This commit is contained in:
Tom Usher 2024-06-19 14:09:20 +01:00
parent 52bcf37079
commit 17482ded74
2 changed files with 4 additions and 0 deletions

View file

@ -53,6 +53,7 @@ class ModelInfo(TypedDict, total=False):
output_cost_per_image: Optional[float] output_cost_per_image: Optional[float]
output_cost_per_video_per_second: Optional[float] output_cost_per_video_per_second: Optional[float]
output_cost_per_audio_per_second: Optional[float] output_cost_per_audio_per_second: Optional[float]
output_vector_size: Optional[int]
litellm_provider: Required[str] litellm_provider: Required[str]
mode: Required[ mode: Required[
Literal[ Literal[

View file

@ -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=_model_info.get(
"output_cost_per_token_above_128k_tokens", None "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=_model_info.get(
"litellm_provider", custom_llm_provider "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=_model_info.get(
"output_cost_per_token_above_128k_tokens", None "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=_model_info.get(
"litellm_provider", custom_llm_provider "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=_model_info.get(
"output_cost_per_token_above_128k_tokens", None "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=_model_info.get(
"litellm_provider", custom_llm_provider "litellm_provider", custom_llm_provider
), ),