scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-28 01:32:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 43

fix(ollama): read optional custom prompt template values with a default

Browse source 
initial_prompt_value and final_prompt_value are documented as optional
on custom_prompt_dict, and litellm.completion() only stores them when
truthy. transform_request read them back with direct dict indexing,
so a template that sets roles but omits either key crashed every
request to that model with KeyError, surfaced as an APIConnectionError
even though no connection was ever attempted.

custom_prompt() itself already defaults both to "", and other
providers (e.g. sagemaker) already read them with .get(key, "").
Bring ollama in line with that.
This commit is contained in:
adhavan18 2026-09-06 14:12:00 +05:30
parent 02522a5441
commit e5054d4be7
2 changed files with 57 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/ollama/completion/transformation.py
View file

@ -360,8 +360,8 @@ class OllamaConfig(BaseConfig):
model_prompt_details: Final = custom_prompt_dict[model]
ollama_prompt = custom_prompt(
role_dict=model_prompt_details["roles"],
initial_prompt_value=model_prompt_details["initial_prompt_value"],
final_prompt_value=model_prompt_details["final_prompt_value"],
initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""),
final_prompt_value=model_prompt_details.get("final_prompt_value", ""),
messages=messages,
)
elif text_completion_request: # handle `/completions` requests

55
tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py
View file

@ -410,6 +410,61 @@ class TestOllamaConfig:
assert result.choices[0]["finish_reason"] == "stop"
class TestOllamaConfigTransformRequest:
def test_transform_request_custom_prompt_missing_initial_and_final_value(self):
"""A custom prompt template that only sets `roles` (initial/final prompt
values omitted, as the docs say is allowed) must not raise KeyError."""
config = OllamaConfig()
litellm_params = {
"custom_prompt_dict": {
"llama2": {
"roles": {
"system": {"pre_message": "<<SYS>>\n", "post_message": "\n<</SYS>>\n"},
"user": {"pre_message": "[INST] ", "post_message": " [/INST]"},
}
}
}
}
result = config.transform_request(
model="llama2",
messages=[{"role": "user", "content": "hello"}],
optional_params={},
litellm_params=litellm_params,
headers={},
)
assert "[INST] hello [/INST]" in result["prompt"]
def test_transform_request_custom_prompt_with_initial_and_final_value(self):
"""Existing behavior: explicit initial/final prompt values still apply."""
config = OllamaConfig()
litellm_params = {
"custom_prompt_dict": {
"llama2": {
"roles": {
"user": {"pre_message": "[INST] ", "post_message": " [/INST]"},
},
"initial_prompt_value": "<start>",
"final_prompt_value": "<end>",
}
}
}
result = config.transform_request(
model="llama2",
messages=[{"role": "user", "content": "hello"}],
optional_params={},
litellm_params=litellm_params,
headers={},
)
assert result["prompt"].startswith("<start>")
assert result["prompt"].endswith("<end>")
class TestOllamaTextCompletionResponseIterator:
def test_chunk_parser_with_thinking_field(self):
"""Test that chunks with 'thinking' field and empty 'response' are handled correctly."""