Files
sglang/python/sglang/test/test_openai_protocol.py
Lianmin Zheng 74b3bfaaf8 format code
2024-01-30 16:36:10 +00:00

52 lines
1.6 KiB
Python

from sglang.srt.managers.openai_protocol import (
ChatCompletionMessageContentImagePart,
ChatCompletionMessageContentImageURL,
ChatCompletionMessageContentTextPart,
ChatCompletionMessageGenericParam,
ChatCompletionMessageUserParam,
ChatCompletionRequest,
)
def test_chat_completion_request_image():
"""Test that Chat Completion Requests with images can be converted."""
image_request = {
"model": "default",
"messages": [
{"role": "system", "content": "You are a helpful AI assistant"},
{
"role": "user",
"content": [
{"type": "text", "text": "Describe this image"},
{"type": "image_url", "image_url": {"url": "https://someurl.com"}},
],
},
],
"temperature": 0,
"max_tokens": 64,
}
request = ChatCompletionRequest(**image_request)
assert len(request.messages) == 2
assert request.messages[0] == ChatCompletionMessageGenericParam(
role="system", content="You are a helpful AI assistant"
)
assert request.messages[1] == ChatCompletionMessageUserParam(
role="user",
content=[
ChatCompletionMessageContentTextPart(
type="text", text="Describe this image"
),
ChatCompletionMessageContentImagePart(
type="image_url",
image_url=ChatCompletionMessageContentImageURL(
url="https://someurl.com"
),
),
],
)
if __name__ == "__main__":
test_chat_completion_request_image()