diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py
index dccc83efed4..66a73765d10 100644
--- a/litellm/llms/ollama/completion/transformation.py
+++ b/litellm/llms/ollama/completion/transformation.py
@@ -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
diff --git a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py
index eadc2bc9541..cc65b44ccbc 100644
--- a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py
+++ b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py
@@ -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": "<>\n", "post_message": "\n<>\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": "",
+ "final_prompt_value": "",
+ }
+ }
+ }
+
+ result = config.transform_request(
+ model="llama2",
+ messages=[{"role": "user", "content": "hello"}],
+ optional_params={},
+ litellm_params=litellm_params,
+ headers={},
+ )
+
+ assert result["prompt"].startswith("")
+ assert result["prompt"].endswith("")
+
+
class TestOllamaTextCompletionResponseIterator:
def test_chunk_parser_with_thinking_field(self):
"""Test that chunks with 'thinking' field and empty 'response' are handled correctly."""