From e5054d4be7306a0844c822196b845c1afebed3af Mon Sep 17 00:00:00 2001
From: adhavan18 <sk.tamiladhavan@gmail.com>
Date: Sun, 6 Sep 2026 14:12:00 +0530
Subject: [PATCH] fix(ollama): read optional custom prompt template values with
 a default

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.
---
 .../llms/ollama/completion/transformation.py  |  4 +-
 .../test_ollama_completion_transformation.py  | 55 +++++++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)

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": "<<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."""