From 7e4dc1d9355b3801a35f3c6b2f28d9aa81121c60 Mon Sep 17 00:00:00 2001 From: Marcus Date: Fri, 4 Sep 2026 11:30:40 -0400 Subject: [PATCH 1/2] fix(prompt-templates): treat initial_prompt_value and final_prompt_value as optional completion() only stores a custom prompt template key when its value is truthy, and custom_prompt() already defaults both prompt values to "". Seven call sites read them back with direct indexing, so a template that sets only roles raises KeyError instead of rendering, on ollama, petals, vllm, predibase, codestral, anthropic text completion, and response_schema_prompt. Five other providers already read defensively, so the same config works on sagemaker and 500s on ollama. Read both keys with .get(key, ""), matching those five providers and the defaults custom_prompt() already declares. Templates that set all three keys render exactly as before. --- .../prompt_templates/factory.py | 4 +- .../anthropic/completion/transformation.py | 4 +- litellm/llms/codestral/completion/handler.py | 4 +- .../llms/ollama/completion/transformation.py | 4 +- litellm/llms/petals/completion/handler.py | 4 +- litellm/llms/predibase/chat/transformation.py | 4 +- litellm/llms/vllm/completion/handler.py | 8 +- .../test_custom_prompt_optional_keys.py | 198 ++++++++++++++++++ 8 files changed, 214 insertions(+), 16 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/prompt_templates/test_custom_prompt_optional_keys.py diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index ba59e3fa997..3ff4b0cd350 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5165,8 +5165,8 @@ def response_schema_prompt(model: str, response_schema: dict) -> str: if custom_prompt_details is not None: return custom_prompt( role_dict=custom_prompt_details["roles"], - initial_prompt_value=custom_prompt_details["initial_prompt_value"], - final_prompt_value=custom_prompt_details["final_prompt_value"], + initial_prompt_value=custom_prompt_details.get("initial_prompt_value", ""), + final_prompt_value=custom_prompt_details.get("final_prompt_value", ""), messages=response_schema_as_message, ) else: diff --git a/litellm/llms/anthropic/completion/transformation.py b/litellm/llms/anthropic/completion/transformation.py index b15b0159bd9..8e2640e62ea 100644 --- a/litellm/llms/anthropic/completion/transformation.py +++ b/litellm/llms/anthropic/completion/transformation.py @@ -239,8 +239,8 @@ class AnthropicTextConfig(BaseConfig): model_prompt_details: Final = custom_prompt_dict[model] 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, ) else: diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index f8486d3b274..ee2e21b2954 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -263,8 +263,8 @@ class CodestralTextCompletion: model_prompt_details: Final = custom_prompt_dict[model] 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, ) else: 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/litellm/llms/petals/completion/handler.py b/litellm/llms/petals/completion/handler.py index c7cfeb1dd1a..35bb7ced0c5 100644 --- a/litellm/llms/petals/completion/handler.py +++ b/litellm/llms/petals/completion/handler.py @@ -44,8 +44,8 @@ def completion( model_prompt_details: Final = litellm.custom_prompt_dict[model] 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, ) else: diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py index 3265537d1aa..b4b0c621334 100644 --- a/litellm/llms/predibase/chat/transformation.py +++ b/litellm/llms/predibase/chat/transformation.py @@ -261,8 +261,8 @@ class PredibaseConfig(BaseConfig): model_prompt_details: Final = custom_prompt_dict[model] 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, ) else: diff --git a/litellm/llms/vllm/completion/handler.py b/litellm/llms/vllm/completion/handler.py index 78e6c74c2f7..113b2d58981 100644 --- a/litellm/llms/vllm/completion/handler.py +++ b/litellm/llms/vllm/completion/handler.py @@ -58,8 +58,8 @@ def completion( model_prompt_details: Final = custom_prompt_dict[model] 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, ) else: @@ -146,8 +146,8 @@ def batch_completions(model: str, messages: list, optional_params=None, custom_p for message in messages: 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=message, ) prompts.append(prompt) diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_custom_prompt_optional_keys.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_custom_prompt_optional_keys.py new file mode 100644 index 00000000000..29f98be40f8 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_custom_prompt_optional_keys.py @@ -0,0 +1,198 @@ +"""Every consumer of ``custom_prompt_dict`` must treat the two prompt-value keys as optional. + +``litellm.completion()`` only stores ``initial_prompt_value`` / ``final_prompt_value`` when the +caller passes a truthy value, and ``custom_prompt()`` already defaults both to ``""``. A template +that sets ``roles`` alone therefore has to render, not raise ``KeyError``. +""" + +import sys +import types + +import pytest + +import litellm +from litellm.litellm_core_utils.prompt_templates.factory import response_schema_prompt +from litellm.llms.anthropic.completion.transformation import AnthropicTextConfig +from litellm.llms.codestral.completion.handler import CodestralTextCompletion +from litellm.llms.ollama.completion.transformation import OllamaConfig +from litellm.llms.petals.completion import handler as petals_handler +from litellm.llms.predibase.chat.transformation import PredibaseConfig +from litellm.llms.vllm.completion import handler as vllm_handler +from litellm.types.utils import ModelResponse, TextCompletionResponse + +MODEL = "partial-template-model" +ROLES = {"user": {"pre_message": "[INST] ", "post_message": " [/INST]"}} +MESSAGES = [{"role": "user", "content": "hi"}] +RENDERED_MESSAGES = "[INST] hi [/INST]" + +RESPONSE_SCHEMA = {"type": "object"} +RENDERED_SCHEMA = f"[INST] {RESPONSE_SCHEMA} [/INST]" + +INITIAL = "" +FINAL = "" + +PARTIAL_TEMPLATE = {"roles": ROLES} +FULL_TEMPLATE = {"roles": ROLES, "initial_prompt_value": INITIAL, "final_prompt_value": FINAL} + + +class _PromptCaptured(Exception): + """Raised by the test doubles below to stop a provider call once it has built its prompt.""" + + def __init__(self, prompt): + super().__init__(prompt) + self.prompt = prompt + + +class _CapturingLogging: + """Stands in for the logging object providers call with the prompt they are about to send.""" + + def pre_call(self, input, api_key, additional_args=None, **kwargs): + raise _PromptCaptured(input) + + def post_call(self, *args, **kwargs): + pass + + +def _captured_prompt(call): + with pytest.raises(_PromptCaptured) as excinfo: + call() + return excinfo.value.prompt + + +@pytest.fixture(autouse=True) +def stub_vllm_import(monkeypatch): + """``vllm`` is an optional heavyweight dependency; its handler only needs the two names.""" + + class _SamplingParams: + def __init__(self, **kwargs): + self.kwargs = kwargs + + class _LLM: + def __init__(self, model): + self.model = model + + def generate(self, prompts, sampling_params): + raise _PromptCaptured(prompts) + + stub = types.ModuleType("vllm") + stub.LLM = _LLM + stub.SamplingParams = _SamplingParams + + monkeypatch.setitem(sys.modules, "vllm", stub) + monkeypatch.setattr(vllm_handler, "llm", None) + + +def _render_ollama(template, monkeypatch): + request = OllamaConfig().transform_request( + model=MODEL, + messages=MESSAGES, + optional_params={}, + litellm_params={"custom_prompt_dict": {MODEL: template}}, + headers={}, + ) + return request["prompt"] + + +def _render_predibase(template, monkeypatch): + request = PredibaseConfig().transform_request( + model=MODEL, + messages=MESSAGES, + optional_params={}, + litellm_params={"custom_prompt_dict": {MODEL: template}}, + headers={}, + ) + return request["inputs"] + + +def _render_anthropic_text(template, monkeypatch): + monkeypatch.setattr(litellm, "custom_prompt_dict", {MODEL: template}) + return AnthropicTextConfig()._get_anthropic_text_prompt_from_messages(messages=MESSAGES, model=MODEL) + + +def _render_response_schema_prompt(template, monkeypatch): + monkeypatch.setattr(litellm, "custom_prompt_dict", {f"{MODEL}/response_schema_prompt": template}) + return response_schema_prompt(model=MODEL, response_schema=RESPONSE_SCHEMA) + + +def _render_petals(template, monkeypatch): + monkeypatch.setattr(litellm, "custom_prompt_dict", {MODEL: template}) + return _captured_prompt( + lambda: petals_handler.completion( + model=MODEL, + messages=MESSAGES, + api_base="http://petals.invalid", + model_response=ModelResponse(), + print_verbose=lambda *args, **kwargs: None, + encoding=None, + logging_obj=_CapturingLogging(), + optional_params={}, + ) + ) + + +def _render_codestral(template, monkeypatch): + return _captured_prompt( + lambda: CodestralTextCompletion().completion( + model=MODEL, + messages=MESSAGES, + api_base="http://codestral.invalid", + custom_prompt_dict={MODEL: template}, + model_response=TextCompletionResponse(), + print_verbose=lambda *args, **kwargs: None, + encoding=None, + api_key="sk-not-used", + logging_obj=_CapturingLogging(), + optional_params={}, + timeout=1.0, + ) + ) + + +def _render_vllm_completion(template, monkeypatch): + return _captured_prompt( + lambda: vllm_handler.completion( + model=MODEL, + messages=MESSAGES, + model_response=ModelResponse(), + print_verbose=lambda *args, **kwargs: None, + encoding=None, + logging_obj=_CapturingLogging(), + optional_params={}, + custom_prompt_dict={MODEL: template}, + ) + ) + + +def _render_vllm_batch(template, monkeypatch): + prompts = _captured_prompt( + lambda: vllm_handler.batch_completions( + model=MODEL, + messages=[MESSAGES], + optional_params={}, + custom_prompt_dict={MODEL: template}, + ) + ) + assert len(prompts) == 1 + return prompts[0] + + +CALL_SITES = [ + pytest.param(_render_ollama, RENDERED_MESSAGES, id="ollama"), + pytest.param(_render_petals, RENDERED_MESSAGES, id="petals"), + pytest.param(_render_vllm_completion, RENDERED_MESSAGES, id="vllm-completion"), + pytest.param(_render_vllm_batch, RENDERED_MESSAGES, id="vllm-batch-completions"), + pytest.param(_render_predibase, RENDERED_MESSAGES, id="predibase"), + pytest.param(_render_codestral, RENDERED_MESSAGES, id="codestral"), + pytest.param(_render_anthropic_text, RENDERED_MESSAGES, id="anthropic-text"), + pytest.param(_render_response_schema_prompt, RENDERED_SCHEMA, id="response-schema-prompt"), +] + + +@pytest.mark.parametrize(("render", "rendered_body"), CALL_SITES) +def test_template_without_the_optional_prompt_values_renders(render, rendered_body, monkeypatch): + assert render(PARTIAL_TEMPLATE, monkeypatch) == rendered_body + + +@pytest.mark.parametrize(("render", "rendered_body"), CALL_SITES) +def test_template_with_the_optional_prompt_values_still_applies_them(render, rendered_body, monkeypatch): + assert render(FULL_TEMPLATE, monkeypatch) == INITIAL + rendered_body + FINAL From b5dac08b0900f728a1e35f9a319afe3155884c6e Mon Sep 17 00:00:00 2001 From: Marcus Date: Fri, 4 Sep 2026 12:43:38 -0400 Subject: [PATCH 2/2] test(prompt-templates): drop the explanatory docstrings from the optional-keys regression test --- .../test_custom_prompt_optional_keys.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_custom_prompt_optional_keys.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_custom_prompt_optional_keys.py index 29f98be40f8..cead250c580 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_custom_prompt_optional_keys.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_custom_prompt_optional_keys.py @@ -1,10 +1,3 @@ -"""Every consumer of ``custom_prompt_dict`` must treat the two prompt-value keys as optional. - -``litellm.completion()`` only stores ``initial_prompt_value`` / ``final_prompt_value`` when the -caller passes a truthy value, and ``custom_prompt()`` already defaults both to ``""``. A template -that sets ``roles`` alone therefore has to render, not raise ``KeyError``. -""" - import sys import types @@ -36,16 +29,12 @@ FULL_TEMPLATE = {"roles": ROLES, "initial_prompt_value": INITIAL, "final_prompt_ class _PromptCaptured(Exception): - """Raised by the test doubles below to stop a provider call once it has built its prompt.""" - def __init__(self, prompt): super().__init__(prompt) self.prompt = prompt class _CapturingLogging: - """Stands in for the logging object providers call with the prompt they are about to send.""" - def pre_call(self, input, api_key, additional_args=None, **kwargs): raise _PromptCaptured(input) @@ -61,8 +50,6 @@ def _captured_prompt(call): @pytest.fixture(autouse=True) def stub_vllm_import(monkeypatch): - """``vllm`` is an optional heavyweight dependency; its handler only needs the two names.""" - class _SamplingParams: def __init__(self, **kwargs): self.kwargs = kwargs