Squashed commit of the following:

commit 122f037990da9c008130cde0ec5bd54259b3d0b0
Author: Krrish Dholakia <krrishdholakia@gmail.com>
Date:   Wed May 7 14:19:35 2025 -0700

    fix(ollama/completion): add 'max_completion_token' support for ollama

commit 314b68cfd4da36be8fbbc0c30f86da15e36438b7
Author: Krrish Dholakia <krrishdholakia@gmail.com>
Date:   Wed May 7 13:27:56 2025 -0700

    fix(utils.py): check if ollama model has custom attributes

    Fixes https://github.com/BerriAI/litellm/discussions/9630
This commit is contained in:
Krrish Dholakia
2025-05-07 22:53:28 -07:00
parent 88f5f9b7f8
commit b899e7d44c
4 changed files with 48 additions and 37 deletions
@@ -150,6 +150,7 @@ class OllamaConfig(BaseConfig):
"frequency_penalty",
"stop",
"response_format",
"max_completion_tokens",
]
def map_openai_params(
@@ -160,7 +161,7 @@ class OllamaConfig(BaseConfig):
drop_params: bool,
) -> dict:
for param, value in non_default_params.items():
if param == "max_tokens":
if param == "max_tokens" or param == "max_completion_tokens":
optional_params["num_predict"] = value
if param == "stream":
optional_params["stream"] = value
@@ -257,9 +258,13 @@ class OllamaConfig(BaseConfig):
model_response.choices[0].finish_reason = "stop"
if request_data.get("format", "") == "json":
response_content = json.loads(response_json["response"])
# Check if this is a function call format with name/arguments structure
if isinstance(response_content, dict) and "name" in response_content and "arguments" in response_content:
if (
isinstance(response_content, dict)
and "name" in response_content
and "arguments" in response_content
):
# Handle as function call (original behavior)
function_call = response_content
message = litellm.Message(
File diff suppressed because one or more lines are too long
+6 -13
View File
@@ -60,16 +60,9 @@ model_list:
- model_name: gemini/gemini-2.0-flash
litellm_params:
model: gemini/gemini-2.0-flash
litellm_settings:
num_retries: 0
check_provider_endpoint: true
cache: true
callbacks: ["otel"]
files_settings:
- custom_llm_provider: gemini
api_key: os.environ/GEMINI_API_KEY
general_settings:
store_prompts_in_spend_logs: true
- model_name: llama-qwen
litellm_params:
model: ollama/qwen2:0.5b
model_info:
input_cost_per_token: 0.75
output_cost_per_token: 3
+34 -20
View File
@@ -523,9 +523,9 @@ def function_setup( # noqa: PLR0915
function_id: Optional[str] = kwargs["id"] if "id" in kwargs else None
## DYNAMIC CALLBACKS ##
dynamic_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = (
kwargs.pop("callbacks", None)
)
dynamic_callbacks: Optional[
List[Union[str, Callable, CustomLogger]]
] = kwargs.pop("callbacks", None)
all_callbacks = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks)
if len(all_callbacks) > 0:
@@ -1209,9 +1209,9 @@ def client(original_function): # noqa: PLR0915
exception=e,
retry_policy=kwargs.get("retry_policy"),
)
kwargs["retry_policy"] = (
reset_retry_policy()
) # prevent infinite loops
kwargs[
"retry_policy"
] = reset_retry_policy() # prevent infinite loops
litellm.num_retries = (
None # set retries to None to prevent infinite loops
)
@@ -2754,16 +2754,16 @@ def get_optional_params( # noqa: PLR0915
True # so that main.py adds the function call to the prompt
)
if "tools" in non_default_params:
optional_params["functions_unsupported_model"] = (
non_default_params.pop("tools")
)
optional_params[
"functions_unsupported_model"
] = non_default_params.pop("tools")
non_default_params.pop(
"tool_choice", None
) # causes ollama requests to hang
elif "functions" in non_default_params:
optional_params["functions_unsupported_model"] = (
non_default_params.pop("functions")
)
optional_params[
"functions_unsupported_model"
] = non_default_params.pop("functions")
elif (
litellm.add_function_to_prompt
): # if user opts to add it to prompt instead
@@ -2786,10 +2786,10 @@ def get_optional_params( # noqa: PLR0915
if "response_format" in non_default_params:
if provider_config is not None:
non_default_params["response_format"] = (
provider_config.get_json_schema_from_pydantic_object(
response_format=non_default_params["response_format"]
)
non_default_params[
"response_format"
] = provider_config.get_json_schema_from_pydantic_object(
response_format=non_default_params["response_format"]
)
else:
non_default_params["response_format"] = type_to_response_format_param(
@@ -3805,9 +3805,9 @@ def _count_characters(text: str) -> int:
def get_response_string(response_obj: Union[ModelResponse, ModelResponseStream]) -> str:
_choices: Union[List[Union[Choices, StreamingChoices]], List[StreamingChoices]] = (
response_obj.choices
)
_choices: Union[
List[Union[Choices, StreamingChoices]], List[StreamingChoices]
] = response_obj.choices
response_str = ""
for choice in _choices:
@@ -4127,6 +4127,18 @@ def get_provider_info(
return model_info
def _is_potential_model_name_in_model_cost(
potential_model_names: PotentialModelNamesAndCustomLLMProvider,
) -> bool:
"""
Check if the potential model name is in the model cost.
"""
return any(
potential_model_name in litellm.model_cost
for potential_model_name in potential_model_names.values()
)
def _get_model_info_helper( # noqa: PLR0915
model: str, custom_llm_provider: Optional[str] = None
) -> ModelInfoBase:
@@ -4182,7 +4194,9 @@ def _get_model_info_helper( # noqa: PLR0915
supports_prompt_caching=None,
supports_pdf_input=None,
)
elif custom_llm_provider == "ollama" or custom_llm_provider == "ollama_chat":
elif (
custom_llm_provider == "ollama" or custom_llm_provider == "ollama_chat"
) and not _is_potential_model_name_in_model_cost(potential_model_names):
return litellm.OllamaConfig().get_model_info(model)
else:
"""