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