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

Merge e5054d4be7 into 2dccc0dc79

Browse source
This commit is contained in:
Tamil Adhavan S K 2026-09-23 14:41:48 +00:00 • committed by GitHub
commit cd4a5871ba
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
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

@ -383,8 +383,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

@ -415,6 +415,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."""