This commit is contained in:
Udbhav Shrinet 2026-11-06 15:52:33 -05:00 • committed by GitHub
commit a889c8977f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 53 additions and 12 deletions

View file

@ -7,10 +7,17 @@ def get_formatted_prompt(
"acompletion",
"completion",
"embedding",
"aembedding",
"embeddings",
"image_generation",
"aimage_generation",
"audio_transcription",
"transcription",
"atranscription",
"moderation",
"amoderation",
"text_completion",
"atext_completion",
],
) -> str:
"""
@ -19,7 +26,7 @@ def get_formatted_prompt(
Returns a string.
"""
prompt = ""
if call_type == "acompletion" or call_type == "completion":
if call_type in {"acompletion", "completion"}:
for message in data["messages"]:
if message.get("content", None) is not None:
content = message.get("content")
@ -34,17 +41,20 @@ def get_formatted_prompt(
if "function" in tool_call:
function_arguments = tool_call["function"]["arguments"]
prompt += function_arguments
elif call_type == "text_completion":
elif call_type in {"text_completion", "atext_completion"}:
prompt = data["prompt"]
elif call_type == "embedding" or call_type == "moderation":
if isinstance(data["input"], str):
prompt = data["input"]
elif isinstance(data["input"], list):
for m in data["input"]:
prompt += m
elif call_type == "image_generation":
elif call_type in {"embedding", "aembedding", "embeddings", "moderation", "amoderation"}:
if "input" in data:
if isinstance(data["input"], str):
prompt = data["input"]
elif isinstance(data["input"], list):
for m in data["input"]:
if isinstance(m, str):
prompt += m
elif (
call_type in {"image_generation", "aimage_generation"}
or (call_type in {"audio_transcription", "transcription", "atranscription"})
and "prompt" in data
):
prompt = data["prompt"]
elif call_type == "audio_transcription":
if "prompt" in data:
prompt = data["prompt"]
return prompt

View file

@ -6,6 +6,18 @@ from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import (
get_formatted_prompt,
)
CallType = Literal[
"transcription",
"atranscription",
"audio_transcription",
"aembedding",
"embedding",
"embeddings",
"amoderation",
"aimage_generation",
"atext_completion",
]
@pytest.mark.parametrize("call_type", ["acompletion", "completion"])
def test_null_tool_calls_are_skipped(call_type: Literal["acompletion", "completion"]) -> None:
@ -22,3 +34,22 @@ def test_null_tool_calls_are_skipped(call_type: Literal["acompletion", "completi
}
assert get_formatted_prompt(data=data, call_type=call_type) == 'pingpong{"x":1}'
@pytest.mark.parametrize(
("data", "call_type", "expected"),
[
({"prompt": "transcribe speech"}, "transcription", "transcribe speech"),
({"prompt": "async speech"}, "atranscription", "async speech"),
({"prompt": "audio legacy"}, "audio_transcription", "audio legacy"),
({"input": "embed async"}, "aembedding", "embed async"),
({"input": ["embed", "dings"]}, "embeddings", "embeddings"),
({"input": [0.1, 0.2]}, "aembedding", ""),
({"input": [101, 2054]}, "embedding", ""),
({"input": "mod async"}, "amoderation", "mod async"),
({"prompt": "gen image async"}, "aimage_generation", "gen image async"),
({"prompt": "complete async"}, "atext_completion", "complete async"),
],
)
def test_get_formatted_prompt_call_types(data: dict[str, object], call_type: CallType, expected: str) -> None:
assert get_formatted_prompt(data=data, call_type=call_type) == expected