mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge e4949e9d35 into 1ab6fd89d2
This commit is contained in:
commit
216edf8f0d
5 changed files with 256 additions and 107 deletions
|
|
@ -1,11 +1,31 @@
|
|||
from functools import lru_cache
|
||||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm import model_cost as _model_cost
|
||||
from litellm import verbose_logger
|
||||
from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
||||
|
||||
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
|
||||
def _cached_ollama_show(model: str, api_base: str, headers: tuple[tuple[str, str], ...] = ()) -> dict[str, Any] | None:
|
||||
from litellm import module_level_client
|
||||
|
||||
try:
|
||||
response: Final = module_level_client.post(
|
||||
url=f"{api_base}/api/show",
|
||||
json={"name": model},
|
||||
headers=dict(headers),
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except Exception:
|
||||
verbose_logger.debug("OllamaError: Could not get model info.")
|
||||
return None
|
||||
|
||||
|
||||
class OllamaError(BaseLLMException):
|
||||
def __init__(self, status_code: int, message: str, headers: dict | httpx.Headers):
|
||||
super().__init__(status_code=status_code, message=message, headers=headers)
|
||||
|
|
@ -52,6 +72,8 @@ class OllamaModelInfo(BaseLLMModelInfo):
|
|||
Returns the union of all model names.
|
||||
"""
|
||||
|
||||
_builtin_model_cost_keys: Final = frozenset(key.lower() for key in _model_cost)
|
||||
|
||||
@staticmethod
|
||||
def get_api_key(api_key=None) -> str | None:
|
||||
"""Get API key from environment variables or litellm configuration"""
|
||||
|
|
@ -141,8 +163,8 @@ class OllamaModelInfo(BaseLLMModelInfo):
|
|||
|
||||
@staticmethod
|
||||
def _is_static_ollama_model(model: str) -> bool:
|
||||
from litellm import model_cost
|
||||
|
||||
# Snapshot at class definition time so Router-registered keys
|
||||
# (added via register_model at startup) don't pollute the check
|
||||
stripped_model: Final = OllamaModelInfo._strip_ollama_model_prefix(model)
|
||||
potential_model_names: Final = {
|
||||
model,
|
||||
|
|
@ -150,14 +172,21 @@ class OllamaModelInfo(BaseLLMModelInfo):
|
|||
"ollama/" + stripped_model,
|
||||
"ollama_chat/" + stripped_model,
|
||||
}
|
||||
model_cost_keys: Final = {key.lower() for key in model_cost}
|
||||
return any(name.lower() in model_cost_keys for name in potential_model_names)
|
||||
return any(name.lower() in OllamaModelInfo._builtin_model_cost_keys for name in potential_model_names)
|
||||
|
||||
@staticmethod
|
||||
def _supports_function_calling(ollama_model_info: dict) -> bool:
|
||||
capabilities: Final = ollama_model_info.get("capabilities", [])
|
||||
if isinstance(capabilities, list) and "tools" in capabilities:
|
||||
return True
|
||||
_template: Final[str] = str(ollama_model_info.get("template", "") or "")
|
||||
return "tools" in _template.lower()
|
||||
|
||||
@staticmethod
|
||||
def _supports_vision(ollama_model_info: dict) -> bool:
|
||||
capabilities: Final = ollama_model_info.get("capabilities", [])
|
||||
return isinstance(capabilities, list) and "vision" in capabilities
|
||||
|
||||
@staticmethod
|
||||
def _get_max_tokens(ollama_model_info: dict) -> int | None:
|
||||
_model_info: Final[dict] = ollama_model_info.get("model_info", {})
|
||||
|
|
@ -173,23 +202,15 @@ class OllamaModelInfo(BaseLLMModelInfo):
|
|||
api_base: str | None = None,
|
||||
api_key: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
from litellm import module_level_client
|
||||
|
||||
model = self._strip_ollama_model_prefix(model)
|
||||
passed_api_base: Final = api_base
|
||||
api_base = self.get_server_api_base(api_base)
|
||||
api_key = self.get_api_key(api_key) if passed_api_base is None or api_key else None
|
||||
headers: Final = {"Authorization": f"Bearer {api_key}"} if api_key else {}
|
||||
resolved_api_key: Final = self.get_api_key(api_key) if passed_api_base is None or api_key else None
|
||||
headers: Final = {"Authorization": f"Bearer {resolved_api_key}"} if resolved_api_key else {}
|
||||
|
||||
try:
|
||||
response: Final = module_level_client.post(
|
||||
url=f"{api_base}/api/show",
|
||||
json={"name": model},
|
||||
headers=headers,
|
||||
)
|
||||
response.raise_for_status()
|
||||
except Exception:
|
||||
verbose_logger.debug("OllamaError: Could not get model info.")
|
||||
ollama_model_info = _cached_ollama_show(model, api_base, tuple(sorted(headers.items())))
|
||||
|
||||
if ollama_model_info is None:
|
||||
return {
|
||||
"key": model,
|
||||
"litellm_provider": "ollama",
|
||||
|
|
@ -201,19 +222,19 @@ class OllamaModelInfo(BaseLLMModelInfo):
|
|||
"max_output_tokens": None,
|
||||
}
|
||||
|
||||
model_info: Final = response.json()
|
||||
max_tokens: Final = self._get_max_tokens(model_info)
|
||||
max_tokens: Final = self._get_max_tokens(ollama_model_info)
|
||||
|
||||
return {
|
||||
"key": model,
|
||||
"litellm_provider": "ollama",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": self._supports_function_calling(model_info),
|
||||
"supports_function_calling": self._supports_function_calling(ollama_model_info),
|
||||
"supports_vision": self._supports_vision(ollama_model_info),
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
"max_tokens": max_tokens,
|
||||
"max_input_tokens": max_tokens,
|
||||
"max_output_tokens": max_tokens,
|
||||
"max_output_tokens": None,
|
||||
}
|
||||
|
||||
def get_model_info(
|
||||
|
|
|
|||
|
|
@ -8624,18 +8624,42 @@ def select_data_generator(
|
|||
)
|
||||
|
||||
|
||||
def get_litellm_model_info(model: dict = {}):
|
||||
model_info: Final = model.get("model_info", {})
|
||||
model_to_lookup = model.get("litellm_params", {}).get("model", None)
|
||||
try:
|
||||
if "azure" in model_to_lookup or model_info.get("base_model"):
|
||||
model_to_lookup = model_info.get("base_model", None)
|
||||
litellm_model_info: Final = litellm.get_model_info(model_to_lookup)
|
||||
return litellm_model_info
|
||||
except Exception:
|
||||
# this should not block returning on /model/info
|
||||
# if litellm does not have info on the model it should return {}
|
||||
return {}
|
||||
def get_litellm_model_info(model: dict) -> dict:
|
||||
litellm_params: Final = model.get("litellm_params", {})
|
||||
config_model_info: Final = model.get("model_info", {})
|
||||
api_base: Final = litellm_params.get("api_base")
|
||||
api_key: Final = litellm_params.get("api_key")
|
||||
model_to_lookup: Final = litellm_params.get("model")
|
||||
base_model: Final = config_model_info.get("base_model")
|
||||
|
||||
candidates: Final = tuple(m for m in (base_model, model_to_lookup) if m)
|
||||
|
||||
for candidate in candidates:
|
||||
try:
|
||||
result = litellm.get_model_info(
|
||||
model=candidate,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
)
|
||||
if result:
|
||||
return result
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
if model_to_lookup:
|
||||
split_model: Final = model_to_lookup.split("/")
|
||||
if len(split_model) > 1:
|
||||
try:
|
||||
return litellm.get_model_info(
|
||||
model=split_model[-1],
|
||||
custom_llm_provider=split_model[0],
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
def on_backoff(details):
|
||||
|
|
@ -12665,28 +12689,6 @@ def _enrich_model_info_with_litellm_data(
|
|||
# input_cost_per_token, output_cost_per_token, max_tokens
|
||||
litellm_model_info = get_litellm_model_info(model=model)
|
||||
|
||||
# 2nd pass on the model, try seeing if we can find model in litellm model_cost map
|
||||
if litellm_model_info == {}:
|
||||
# use litellm_param model_name to get model_info
|
||||
litellm_params = model.get("litellm_params", {})
|
||||
litellm_model = litellm_params.get("model", None)
|
||||
try:
|
||||
litellm_model_info = litellm.get_model_info(model=litellm_model)
|
||||
except Exception:
|
||||
litellm_model_info = {}
|
||||
# 3rd pass on the model, try seeing if we can find model but without the "/" in model cost map
|
||||
if litellm_model_info == {}:
|
||||
# use litellm_param model_name to get model_info
|
||||
litellm_params = model.get("litellm_params", {})
|
||||
litellm_model = litellm_params.get("model", None)
|
||||
if litellm_model:
|
||||
split_model: Final = litellm_model.split("/")
|
||||
if len(split_model) > 0:
|
||||
litellm_model = split_model[-1]
|
||||
try:
|
||||
litellm_model_info = litellm.get_model_info(model=litellm_model, custom_llm_provider=split_model[0])
|
||||
except Exception:
|
||||
litellm_model_info = {}
|
||||
for k, v in litellm_model_info.items():
|
||||
if k not in model_info:
|
||||
model_info[k] = v
|
||||
|
|
@ -14088,27 +14090,6 @@ def _get_proxy_model_info(model: dict) -> dict:
|
|||
# input_cost_per_token, output_cost_per_token, max_tokens
|
||||
litellm_model_info = get_litellm_model_info(model=model)
|
||||
|
||||
# 2nd pass on the model, try seeing if we can find model in litellm model_cost map
|
||||
if litellm_model_info == {}:
|
||||
# use litellm_param model_name to get model_info
|
||||
litellm_params = model.get("litellm_params", {})
|
||||
litellm_model = litellm_params.get("model", None)
|
||||
try:
|
||||
litellm_model_info = litellm.get_model_info(model=litellm_model)
|
||||
except Exception:
|
||||
litellm_model_info = {}
|
||||
# 3rd pass on the model, try seeing if we can find model but without the "/" in model cost map
|
||||
if litellm_model_info == {}:
|
||||
# use litellm_param model_name to get model_info
|
||||
litellm_params = model.get("litellm_params", {})
|
||||
litellm_model = litellm_params.get("model", None)
|
||||
split_model: Final = litellm_model.split("/")
|
||||
if len(split_model) > 0:
|
||||
litellm_model = split_model[-1]
|
||||
try:
|
||||
litellm_model_info = litellm.get_model_info(model=litellm_model, custom_llm_provider=split_model[0])
|
||||
except Exception:
|
||||
litellm_model_info = {}
|
||||
for k, v in litellm_model_info.items():
|
||||
if k not in model_info:
|
||||
model_info[k] = v
|
||||
|
|
|
|||
|
|
@ -33835,7 +33835,8 @@
|
|||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": false
|
||||
"supports_function_calling": false,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/codegemma": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33844,7 +33845,9 @@
|
|||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "completion",
|
||||
"output_cost_per_token": 0.0
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_vision": false,
|
||||
"supports_function_calling": false
|
||||
},
|
||||
"ollama/codellama": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33853,7 +33856,9 @@
|
|||
"max_output_tokens": 4096,
|
||||
"max_tokens": 4096,
|
||||
"mode": "completion",
|
||||
"output_cost_per_token": 0.0
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_vision": false,
|
||||
"supports_function_calling": false
|
||||
},
|
||||
"ollama/deepseek-coder-v2-base": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33863,7 +33868,8 @@
|
|||
"max_tokens": 8192,
|
||||
"mode": "completion",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/deepseek-coder-v2-instruct": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33873,7 +33879,8 @@
|
|||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/deepseek-coder-v2-lite-base": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33883,7 +33890,8 @@
|
|||
"max_tokens": 8192,
|
||||
"mode": "completion",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/deepseek-coder-v2-lite-instruct": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33893,7 +33901,8 @@
|
|||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/deepseek-v3.1:671b-cloud": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33903,7 +33912,8 @@
|
|||
"max_tokens": 163840,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/gpt-oss:120b-cloud": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33913,7 +33923,8 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/gpt-oss:20b-cloud": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33923,7 +33934,8 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/internlm2_5-20b-chat": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33933,7 +33945,8 @@
|
|||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/llama2": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33942,7 +33955,9 @@
|
|||
"max_output_tokens": 4096,
|
||||
"max_tokens": 4096,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_vision": false,
|
||||
"supports_function_calling": false
|
||||
},
|
||||
"ollama/llama2-uncensored": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33951,7 +33966,9 @@
|
|||
"max_output_tokens": 4096,
|
||||
"max_tokens": 4096,
|
||||
"mode": "completion",
|
||||
"output_cost_per_token": 0.0
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_vision": false,
|
||||
"supports_function_calling": false
|
||||
},
|
||||
"ollama/llama2:13b": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33960,7 +33977,9 @@
|
|||
"max_output_tokens": 4096,
|
||||
"max_tokens": 4096,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_vision": false,
|
||||
"supports_function_calling": false
|
||||
},
|
||||
"ollama/llama2:70b": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33969,7 +33988,9 @@
|
|||
"max_output_tokens": 4096,
|
||||
"max_tokens": 4096,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_vision": false,
|
||||
"supports_function_calling": false
|
||||
},
|
||||
"ollama/llama2:7b": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33978,7 +33999,9 @@
|
|||
"max_output_tokens": 4096,
|
||||
"max_tokens": 4096,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_vision": false,
|
||||
"supports_function_calling": false
|
||||
},
|
||||
"ollama/llama3": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33987,7 +34010,9 @@
|
|||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_vision": false,
|
||||
"supports_function_calling": false
|
||||
},
|
||||
"ollama/llama3.1": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -33997,7 +34022,8 @@
|
|||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/llama3:70b": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -34006,7 +34032,9 @@
|
|||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_vision": false,
|
||||
"supports_function_calling": false
|
||||
},
|
||||
"ollama/llama3:8b": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -34015,7 +34043,9 @@
|
|||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_vision": false,
|
||||
"supports_function_calling": false
|
||||
},
|
||||
"ollama/mistral": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -34025,7 +34055,8 @@
|
|||
"max_tokens": 8192,
|
||||
"mode": "completion",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/mistral-7B-Instruct-v0.1": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -34035,7 +34066,8 @@
|
|||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/mistral-7B-Instruct-v0.2": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -34045,7 +34077,8 @@
|
|||
"max_tokens": 32768,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/mistral-large-instruct-2407": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -34055,7 +34088,8 @@
|
|||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/mixtral-8x22B-Instruct-v0.1": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -34065,7 +34099,8 @@
|
|||
"max_tokens": 65536,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/mixtral-8x7B-Instruct-v0.1": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -34075,7 +34110,8 @@
|
|||
"max_tokens": 32768,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/orca-mini": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -34084,7 +34120,9 @@
|
|||
"max_output_tokens": 4096,
|
||||
"max_tokens": 4096,
|
||||
"mode": "completion",
|
||||
"output_cost_per_token": 0.0
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_vision": false,
|
||||
"supports_function_calling": false
|
||||
},
|
||||
"ollama/qwen3-coder:480b-cloud": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -34094,7 +34132,8 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_function_calling": true
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"ollama/vicuna": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
@ -34103,7 +34142,9 @@
|
|||
"max_output_tokens": 2048,
|
||||
"max_tokens": 2048,
|
||||
"mode": "completion",
|
||||
"output_cost_per_token": 0.0
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_vision": false,
|
||||
"supports_function_calling": false
|
||||
},
|
||||
"omni-moderation-2024-09-26": {
|
||||
"input_cost_per_token": 0.0,
|
||||
|
|
|
|||
|
|
@ -3146,4 +3146,4 @@ def test_get_litellm_model_info(data):
|
|||
new=get_info_mock,
|
||||
):
|
||||
get_litellm_model_info(model=model)
|
||||
get_info_mock.assert_called_once_with(data["expected"])
|
||||
get_info_mock.assert_called_once_with(model=data["expected"], api_base=None, api_key=None)
|
||||
|
|
|
|||
|
|
@ -20,7 +20,15 @@ if "httpx" not in sys.modules:
|
|||
import httpx
|
||||
import litellm
|
||||
|
||||
from litellm.llms.ollama.common_utils import OllamaModelInfo
|
||||
from litellm.llms.ollama.common_utils import OllamaModelInfo, _cached_ollama_show
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_ollama_show_cache():
|
||||
"""Clear the lru_cache on _cached_ollama_show before each test."""
|
||||
_cached_ollama_show.cache_clear()
|
||||
yield
|
||||
_cached_ollama_show.cache_clear()
|
||||
|
||||
|
||||
class DummyResponse:
|
||||
|
|
@ -513,6 +521,7 @@ class TestOllamaGetModelInfo:
|
|||
)
|
||||
assert captured_json[0]["name"] == "my-custom-model"
|
||||
|
||||
_cached_ollama_show.cache_clear()
|
||||
config.get_model_info(
|
||||
"ollama_chat/my-custom-model", api_base="http://localhost:11434"
|
||||
)
|
||||
|
|
@ -576,6 +585,103 @@ class TestOllamaGetModelInfo:
|
|||
assert model_info["litellm_provider"] == "ollama"
|
||||
|
||||
|
||||
class TestOllamaModelInfoCapabilities:
|
||||
"""Tests for capability detection from Ollama /api/show response."""
|
||||
|
||||
def test_get_runtime_model_info_reports_vision_capability(self, monkeypatch):
|
||||
"""supports_vision should be True when capabilities includes 'vision'."""
|
||||
from litellm.llms.ollama.completion.transformation import OllamaConfig
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
return DummyResponse(
|
||||
{
|
||||
"template": "{{ .System }} tools {{ .Prompt }}",
|
||||
"capabilities": ["completion", "tools", "vision"],
|
||||
"model_info": {"llama.context_length": 131072},
|
||||
},
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
|
||||
config = OllamaConfig()
|
||||
result = config.get_model_info("my-vision-model", api_base="http://localhost:11434")
|
||||
|
||||
assert result["supports_vision"] is True
|
||||
assert result["supports_function_calling"] is True
|
||||
assert result["max_input_tokens"] == 131072
|
||||
|
||||
def test_get_runtime_model_info_no_vision_capability(self, monkeypatch):
|
||||
"""supports_vision should be False when capabilities lacks 'vision'."""
|
||||
from litellm.llms.ollama.completion.transformation import OllamaConfig
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
return DummyResponse(
|
||||
{
|
||||
"template": "{{ .System }} {{ .Prompt }}",
|
||||
"capabilities": ["completion"],
|
||||
"model_info": {"llama.context_length": 8192},
|
||||
},
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
|
||||
config = OllamaConfig()
|
||||
result = config.get_model_info("my-text-model", api_base="http://localhost:11434")
|
||||
|
||||
assert result["supports_vision"] is False
|
||||
|
||||
def test_get_runtime_model_info_no_capabilities_field(self, monkeypatch):
|
||||
"""supports_vision should be False when capabilities field is absent."""
|
||||
from litellm.llms.ollama.completion.transformation import OllamaConfig
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
return DummyResponse(
|
||||
{
|
||||
"template": "{{ .System }} {{ .Prompt }}",
|
||||
"model_info": {"llama.context_length": 8192},
|
||||
},
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
|
||||
config = OllamaConfig()
|
||||
result = config.get_model_info("my-text-model", api_base="http://localhost:11434")
|
||||
|
||||
assert result["supports_vision"] is False
|
||||
|
||||
def test_litellm_get_model_info_threads_api_base_to_ollama(self, monkeypatch):
|
||||
"""litellm.get_model_info should pass api_base through to the Ollama provider hook."""
|
||||
captured_urls = []
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
captured_urls.append(url)
|
||||
return DummyResponse(
|
||||
{
|
||||
"template": "{{ .System }} tools {{ .Prompt }}",
|
||||
"capabilities": ["completion", "tools", "vision"],
|
||||
"model_info": {"llama.context_length": 32768},
|
||||
},
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
litellm.get_model_info.cache_clear()
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
try:
|
||||
model_info = litellm.get_model_info(
|
||||
"ollama_chat/llama3.2-vision:11b",
|
||||
api_base="http://remote-ollama:11434",
|
||||
)
|
||||
finally:
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
assert captured_urls[0] == "http://remote-ollama:11434/api/show"
|
||||
assert model_info["supports_vision"] is True
|
||||
assert model_info["supports_function_calling"] is True
|
||||
assert model_info["max_input_tokens"] == 32768
|
||||
|
||||
class TestOllamaAuthHeaders:
|
||||
"""Tests for Ollama authentication header handling in completion calls."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue